feat: Implement SIMD operations

Author: ospencer

src/dune

@@ -27,6 +27,7 @@
   (flags :standard -O2 -Wall -Wextra))
  (js_of_ocaml
   (javascript_files
+   helpers.js
    array_type.js
    type.js
    op.js

src/expression.c

@@ -351,6 +351,104 @@ caml_binaryen_unreachable(value _module) {
   CAMLreturn(alloc_BinaryenExpressionRef(exp));
 }
 
+CAMLprim value
+caml_binaryen_simd_extract(value _module, value _op, value _vec, value _index) {
+  CAMLparam4(_module, _op, _vec, _index);
+  BinaryenModuleRef module = BinaryenModuleRef_val(_module);
+  BinaryenOp op = BinaryenOp_val(_op);
+  BinaryenExpressionRef vec = BinaryenExpressionRef_val(_vec);
+  int index = Int_val(_index);
+  BinaryenExpressionRef exp = BinaryenSIMDExtract(module, op, vec, index);
+  CAMLreturn(alloc_BinaryenExpressionRef(exp));
+}
+
+CAMLprim value
+caml_binaryen_simd_replace(value _module, value _op, value _vec, value _index, value _val) {
+  CAMLparam5(_module, _op, _vec, _index, _val);
+  BinaryenModuleRef module = BinaryenModuleRef_val(_module);
+  BinaryenOp op = BinaryenOp_val(_op);
+  BinaryenExpressionRef vec = BinaryenExpressionRef_val(_vec);
+  int index = Int_val(_index);
+  BinaryenExpressionRef val = BinaryenExpressionRef_val(_val);
+  BinaryenExpressionRef exp = BinaryenSIMDReplace(module, op, vec, index, val);
+  CAMLreturn(alloc_BinaryenExpressionRef(exp));
+}
+
+CAMLprim value
+caml_binaryen_simd_shuffle(value _module, value _left, value _right, value _mask) {
+  CAMLparam4(_module, _left, _right, _mask);
+  BinaryenModuleRef module = BinaryenModuleRef_val(_module);
+  BinaryenExpressionRef left = BinaryenExpressionRef_val(_left);
+  BinaryenExpressionRef right = BinaryenExpressionRef_val(_right);
+  uint8_t mask[16];
+  for (int i = 0; i < 16; i++) {
+    mask[i] = Int_val(Field(_mask, i));
+  }
+  BinaryenExpressionRef exp = BinaryenSIMDShuffle(module, left, right, mask);
+  CAMLreturn(alloc_BinaryenExpressionRef(exp));
+}
+
+CAMLprim value
+caml_binaryen_simd_ternary(value _module, value _op, value _a, value _b, value _c) {
+  CAMLparam5(_module, _op, _a, _b, _c);
+  BinaryenModuleRef module = BinaryenModuleRef_val(_module);
+  BinaryenOp op = BinaryenOp_val(_op);
+  BinaryenExpressionRef a = BinaryenExpressionRef_val(_a);
+  BinaryenExpressionRef b = BinaryenExpressionRef_val(_b);
+  BinaryenExpressionRef c = BinaryenExpressionRef_val(_c);
+  BinaryenExpressionRef exp = BinaryenSIMDTernary(module, op, a, b, c);
+  CAMLreturn(alloc_BinaryenExpressionRef(exp));
+}
+
+CAMLprim value
+caml_binaryen_simd_shift(value _module, value _op, value _vec, value _shift) {
+  CAMLparam4(_module, _op, _vec, _shift);
+  BinaryenModuleRef module = BinaryenModuleRef_val(_module);
+  BinaryenOp op = BinaryenOp_val(_op);
+  BinaryenExpressionRef vec = BinaryenExpressionRef_val(_vec);
+  BinaryenExpressionRef shift = BinaryenExpressionRef_val(_shift);
+  BinaryenExpressionRef exp = BinaryenSIMDShift(module, op, vec, shift);
+  CAMLreturn(alloc_BinaryenExpressionRef(exp));
+}
+
+CAMLprim value
+caml_binaryen_simd_load(value _module, value _op, value _offset, value _align, value _ptr, value _memoryName) {
+  CAMLparam5(_module, _op, _offset, _align, _ptr);
+  CAMLxparam1(_memoryName);
+  BinaryenModuleRef module = BinaryenModuleRef_val(_module);
+  BinaryenOp op = BinaryenOp_val(_op);
+  int offset = Int_val(_offset);
+  int align = Int_val(_align);
+  BinaryenExpressionRef ptr = BinaryenExpressionRef_val(_ptr);
+  char* memoryName = Safe_String_val(_memoryName);
+  BinaryenExpressionRef exp = BinaryenSIMDLoad(module, op, offset, align, ptr, memoryName);
+  CAMLreturn(alloc_BinaryenExpressionRef(exp));
+}
+CAMLprim value
+caml_binaryen_simd_load__bytecode(value * argv) {
+  return caml_binaryen_simd_load(argv[0], argv[1], argv[2], argv[3], argv[4], argv[5]);
+}
+
+CAMLprim value
+caml_binaryen_simd_load_store_lane(value _module, value _op, value _offset, value _align, value _index, value _ptr, value _vec, value _memoryName) {
+  CAMLparam5(_module, _op, _offset, _align, _index);
+  CAMLxparam3(_ptr, _vec, _memoryName);
+  BinaryenModuleRef module = BinaryenModuleRef_val(_module);
+  BinaryenOp op = BinaryenOp_val(_op);
+  int offset = Int_val(_offset);
+  int align = Int_val(_align);
+  int index = Int_val(_index);
+  BinaryenExpressionRef ptr = BinaryenExpressionRef_val(_ptr);
+  BinaryenExpressionRef vec = BinaryenExpressionRef_val(_vec);
+  char* memoryName = Safe_String_val(_memoryName);
+  BinaryenExpressionRef exp = BinaryenSIMDLoadStoreLane(module, op, offset, align, index, ptr, vec, memoryName);
+  CAMLreturn(alloc_BinaryenExpressionRef(exp));
+}
+CAMLprim value
+caml_binaryen_simd_load_store_lane__bytecode(value * argv) {
+  return caml_binaryen_simd_load_store_lane(argv[0], argv[1], argv[2], argv[3], argv[4], argv[5], argv[6], argv[7]);
+}
+
 CAMLprim value
 caml_binaryen_memory_init(value _module, value _segment, value _dest, value _offset, value _size, value _memoryName) {
   CAMLparam5(_module, _segment, _dest, _offset, _size);

src/expression.js

@@ -387,6 +387,10 @@ function caml_binaryen_const(wasm_mod, lit) {
     );
   }
 
+  if (lit.type === "vec128") {
+    return wasm_mod.v128.const(lit.value);
+  }
+
   caml_failwith("invalid Literal for Binaryen.Const.make");
 }
 
@@ -446,6 +450,107 @@ function caml_binaryen_unreachable(wasm_mod) {
   return wasm_mod.unreachable();
 }
 
+//Provides: caml_binaryen_simd_extract
+//Requires: Binaryen
+function caml_binaryen_simd_extract(wasm_mod, op, vec, index) {
+  return Binaryen._BinaryenSIMDExtract(wasm_mod, op, vec, index);
+}
+
+//Provides: caml_binaryen_simd_replace
+//Requires: Binaryen
+function caml_binaryen_simd_replace(wasm_mod, op, vec, index, value) {
+  return Binaryen._BinaryenSIMDReplace(wasm_mod, op, vec, index, value);
+}
+
+//Provides: caml_binaryen_simd_shuffle
+//Requires: caml_js_from_array
+function caml_binaryen_simd_shuffle(wasm_mod, left, right, mask) {
+  return wasm_mod.i8x16.shuffle(left, right, caml_js_from_array(mask));
+}
+
+//Provides: caml_binaryen_simd_ternary
+//Requires: Binaryen
+function caml_binaryen_simd_ternary(wasm_mod, op, a, b, c) {
+  return Binaryen._BinaryenSIMDTernary(wasm_mod, op, a, b, c);
+}
+
+//Provides: caml_binaryen_simd_shift
+//Requires: Binaryen
+function caml_binaryen_simd_shift(wasm_mod, op, vec, shift) {
+  return Binaryen._BinaryenSIMDShift(wasm_mod, op, vec, shift);
+}
+
+//Provides: caml_binaryen_simd_load
+//Requires: Binaryen
+//Requires: caml_jsstring_of_string, caml_alloc_string_on_heap
+function caml_binaryen_simd_load(wasm_mod, op, offset, align, ptr, memoryName) {
+  var memory = caml_alloc_string_on_heap(caml_jsstring_of_string(memoryName));
+  var exp = Binaryen._BinaryenSIMDLoad(
+    wasm_mod,
+    op,
+    offset,
+    align,
+    ptr,
+    memory
+  );
+  Binaryen._free(memory);
+  return exp;
+}
+//Provides: caml_binaryen_simd_load__bytecode
+//Requires: caml_binaryen_simd_load
+function caml_binaryen_simd_load__bytecode() {
+  return caml_binaryen_simd_load(
+    arguments[0],
+    arguments[1],
+    arguments[2],
+    arguments[3],
+    arguments[4],
+    arguments[5]
+  );
+}
+
+//Provides: caml_binaryen_simd_load_store_lane
+//Requires: Binaryen
+//Requires: caml_jsstring_of_string, caml_alloc_string_on_heap
+function caml_binaryen_simd_load_store_lane(
+  wasm_mod,
+  op,
+  offset,
+  align,
+  index,
+  ptr,
+  vec,
+  memoryName
+) {
+  var memory = caml_alloc_string_on_heap(caml_jsstring_of_string(memoryName));
+  var exp = Binaryen._BinaryenSIMDLoadStoreLane(
+    wasm_mod,
+    op,
+    offset,
+    align,
+    index,
+    ptr,
+    vec,
+    memory
+  );
+  Binaryen._free(memory);
+  return exp;
+}
+//Provides: caml_binaryen_simd_load_store_lane__bytecode
+//Requires: caml_binaryen_simd_load_store_lane
+function caml_binaryen_simd_load_store_lane__bytecode() {
+  return caml_binaryen_simd_load_store_lane(
+    arguments[0],
+    arguments[1],
+    arguments[2],
+    arguments[3],
+    arguments[4],
+    arguments[5],
+    arguments[6],
+    arguments[7]
+  );
+}
+
 //Provides: caml_binaryen_memory_init
 //Requires: caml_jsstring_of_string
 function caml_binaryen_memory_init(

src/expression.ml

@@ -819,6 +819,54 @@ module Pop = struct
   (** Module, type *)
 end
 
+module SIMD_extract = struct
+  external make : Module.t -> Op.t -> t -> int -> t
+    = "caml_binaryen_simd_extract"
+  (** Module, op, vec, index *)
+end
+
+module SIMD_replace = struct
+  external make : Module.t -> Op.t -> t -> int -> t -> t
+    = "caml_binaryen_simd_replace"
+  (** Module, op, vec, index, value *)
+end
+
+module SIMD_shuffle = struct
+  external make : Module.t -> t -> t -> int array -> t
+    = "caml_binaryen_simd_shuffle"
+
+  (** Module, left, right, mask *)
+  let make wasm_mod left right mask =
+    if Array.length mask <> 16 then
+      raise
+        (Invalid_argument
+           "Binaryen.SIMD_shuffle.make: mask must be of length 16");
+    make wasm_mod left right mask
+end
+
+module SIMD_ternary = struct
+  external make : Module.t -> Op.t -> t -> t -> t -> t
+    = "caml_binaryen_simd_ternary"
+  (** Module, op, first, second, third *)
+end
+
+module SIMD_shift = struct
+  external make : Module.t -> Op.t -> t -> t -> t = "caml_binaryen_simd_shift"
+  (** Module, op, left, right *)
+end
+
+module SIMD_load = struct
+  external make : Module.t -> Op.t -> int -> int -> t -> string -> t
+    = "caml_binaryen_simd_load__bytecode" "caml_binaryen_simd_load"
+  (** Module, op, offset, align, ptr, memory *)
+end
+
+module SIMD_load_store_lane = struct
+  external make : Module.t -> Op.t -> int -> int -> int -> t -> t -> string -> t
+    = "caml_binaryen_simd_load_store_lane__bytecode" "caml_binaryen_simd_load_store_lane"
+  (** Module, op, offset, align, index, ptr, vec, memory *)
+end
+
 module I31 = struct
   external make : Module.t -> t -> t = "caml_binaryen_i31_new"
   (** Module, value *)

src/expression.mli

@@ -321,6 +321,34 @@ module Pop : sig
   val make : Module.t -> Type.t -> t
 end
 
+module SIMD_extract : sig
+  val make : Module.t -> Op.t -> t -> int -> t
+end
+
+module SIMD_replace : sig
+  val make : Module.t -> Op.t -> t -> int -> t -> t
+end
+
+module SIMD_shuffle : sig
+  val make : Module.t -> t -> t -> int array -> t
+end
+
+module SIMD_ternary : sig
+  val make : Module.t -> Op.t -> t -> t -> t -> t
+end
+
+module SIMD_shift : sig
+  val make : Module.t -> Op.t -> t -> t -> t
+end
+
+module SIMD_load : sig
+  val make : Module.t -> Op.t -> int -> int -> t -> string -> t
+end
+
+module SIMD_load_store_lane : sig
+  val make : Module.t -> Op.t -> int -> int -> int -> t -> t -> string -> t
+end
+
 module I31 : sig
   val make : Module.t -> t -> t
   (** Module, value *)

src/helpers.js

@@ -0,0 +1,10 @@
+//Requires: Binaryen
+//Provides: caml_alloc_string_on_heap
+function caml_alloc_string_on_heap(str) {
+  var len = str.length;
+  var ptr = Binaryen._malloc(len + 1);
+  var encoder = new TextEncoder();
+  encoder.encodeInto(str, Binaryen.HEAPU8.subarray(ptr, ptr + len));
+  Binaryen.HEAPU8[ptr + len] = 0;
+  return ptr;
+}

src/literal.c

@@ -45,3 +45,14 @@ caml_binaryen_literal_float64_bits(value _i64) {
   struct BinaryenLiteral lit = BinaryenLiteralFloat64Bits(i64);
   CAMLreturn(alloc_BinaryenLiteral(lit));
 }
+
+CAMLprim value
+caml_binaryen_literal_vec128(value _v128) {
+  CAMLparam1(_v128);
+  uint8_t v128[16];
+  for (int i = 0; i < 16; i++) {
+    v128[i] = Int_val(Field(_v128, i));
+  }
+  struct BinaryenLiteral lit = BinaryenLiteralVec128(v128);
+  CAMLreturn(alloc_BinaryenLiteral(lit));
+}

src/literal.js

@@ -27,3 +27,10 @@ function caml_binaryen_literal_float64_bits(i64) {
   // Hack around Binaryen.js
   return { type: 'float64_bits', value: i64 };
 }
+
+//Provides: caml_binaryen_literal_vec128
+//Requires: caml_js_from_array
+function caml_binaryen_literal_vec128(v128) {
+  // Hack around Binaryen.js
+  return { type: "vec128", value: caml_js_from_array(v128) };
+}

src/literal.ml

@@ -8,3 +8,11 @@ external float64_bits : int64 -> t = "caml_binaryen_literal_float64_bits"
 let float32 n = float32_bits @@ Int32.bits_of_float n
 
 external float64 : float -> t = "caml_binaryen_literal_float64"
+external vec128 : int array -> t = "caml_binaryen_literal_vec128"
+
+let vec128 (low, high) =
+  let bytes = Bytes.create 16 in
+  Bytes.set_int64_le bytes 0 low;
+  Bytes.set_int64_le bytes 8 high;
+  let arr = Array.init 16 (Bytes.get_uint8 bytes) in
+  vec128 arr

src/literal.mli

@@ -6,3 +6,4 @@ val float32_bits : int32 -> t
 val float64_bits : int64 -> t
 val float32 : float -> t
 val float64 : float -> t
+val vec128 : int64 * int64 -> t