Implement SIMD load splat and load extend (#307)

* Implement SIMD load splat and load extend

Add more pack types for SIMD load splat and load extends.

* Create new pack_simd type and SimdLoad AST

* Reorder functions

* Formatting

* Simplify simd packed loads

* Add SimdStore and make simd_loadop more consistent
New AST nodes SimdLoad and SimdStore.

Author: ngzhian

interpreter/binary/encode.ml

@@ -567,6 +567,8 @@ let encode m =
       | SimdBitmask V128Op.(I32x4 Bitmask) -> simd_op 0xa4l
       | SimdBitmask (_) -> assert false
 
+      | _ -> assert false
+
     let const c =
       list instr c.it; end_ ()
 

interpreter/exec/eval.ml

@@ -220,6 +220,17 @@ let rec step (c : config) : config =
           in v :: vs', []
         with exn -> vs', [Trapping (memory_error e.at exn) @@ e.at])
 
+      | SimdLoad {offset; ty; sz; _}, I32 i :: vs' ->
+        let mem = memory frame.inst (0l @@ e.at) in
+        let addr = I64_convert.extend_i32_u i in
+        (try
+          let v =
+            match sz with
+            | None -> Memory.load_value mem addr offset ty
+            | Some (pack_size, simd_load) -> Memory.load_simd_packed pack_size simd_load mem addr offset ty
+          in v :: vs', []
+        with exn -> vs', [Trapping (memory_error e.at exn) @@ e.at])
+
       | Store {offset; sz; _}, v :: I32 i :: vs' ->
         let mem = memory frame.inst (0l @@ e.at) in
         let addr = I64_convert.extend_i32_u i in
@@ -231,6 +242,14 @@ let rec step (c : config) : config =
           vs', []
         with exn -> vs', [Trapping (memory_error e.at exn) @@ e.at]);
 
+      | SimdStore {offset; sz; _}, v :: I32 i :: vs' ->
+        let mem = memory frame.inst (0l @@ e.at) in
+        let addr = I64_convert.extend_i32_u i in
+        (try
+          Memory.store_value mem addr offset v;
+          vs', []
+        with exn -> vs', [Trapping (memory_error e.at exn) @@ e.at]);
+
       | MemorySize, vs ->
         let mem = memory frame.inst (0l @@ e.at) in
         I32 (Memory.size mem) :: vs, []

interpreter/exec/simd.ml

@@ -180,6 +180,10 @@ sig
     val widen_low_u : t -> t
     val widen_high_u : t -> t
   end
+  module I64x2_convert : sig
+    val widen_low_s : t -> t
+    val widen_low_u : t -> t
+  end
   module F32x4_convert : sig
     val convert_i32x4_s : t -> t
     val convert_i32x4_u : t -> t
@@ -415,6 +419,16 @@ struct
     let widen_high_u = widen Lib.List.drop 0xffffl
   end
 
+  module I64x2_convert = struct
+    let widen mask x =
+      Rep.of_i64x2
+        (List.map
+           (fun i32 -> Int64.(logand mask (of_int32 i32)))
+           (Lib.List.take 2 (Rep.to_i32x4 x)))
+    let widen_low_s = widen 0xffffffffffffffffL
+    let widen_low_u = widen 0xffffffffL
+  end
+
   module F32x4_convert = struct
     let convert f v = Rep.of_f32x4 (List.map f (Rep.to_i32x4 v))
     let convert_i32x4_s = convert F32_convert.convert_i32_s

interpreter/runtime/memory.ml

@@ -131,6 +131,26 @@ let load_packed sz ext mem a o t =
   | I64Type -> I64 x
   | _ -> raise Type
 
+let load_simd_packed pack_size simd_load mem a o t =
+  let n = packed_size pack_size in
+  assert (n <= Types.size t);
+  let x = loadn mem a o n in
+  let b = Bytes.create 16 in
+  Bytes.set_int64_le b 0 x;
+  let v = V128.of_bits (Bytes.to_string b) in
+  match pack_size, simd_load with
+  | Pack64, Pack8x8 SX -> V128 (V128.I16x8_convert.widen_low_s v)
+  | Pack64, Pack8x8 ZX -> V128 (V128.I16x8_convert.widen_low_u v)
+  | Pack64, Pack16x4 SX -> V128 (V128.I32x4_convert.widen_low_s v)
+  | Pack64, Pack16x4 ZX -> V128 (V128.I32x4_convert.widen_low_u v)
+  | Pack64, Pack32x2 SX -> V128 (V128.I64x2_convert.widen_low_s v)
+  | Pack64, Pack32x2 ZX -> V128 (V128.I64x2_convert.widen_low_u v)
+  | Pack8, PackSplat -> V128 (V128.I8x16.splat (I8.of_int_s (Int64.to_int x)))
+  | Pack16, PackSplat -> V128 (V128.I16x8.splat (I16.of_int_s (Int64.to_int x)))
+  | Pack32, PackSplat -> V128 (V128.I32x4.splat (I32.of_int_s (Int64.to_int x)))
+  | Pack64, PackSplat -> V128 (V128.I64x2.splat x)
+  | _ -> assert false
+
 let store_packed sz mem a o v =
   assert (packed_size sz <= Types.size (Values.type_of v));
   let n = packed_size sz in

interpreter/runtime/memory.mli

@@ -35,6 +35,9 @@ val store_value :
 val load_packed :
   pack_size -> extension -> memory -> address -> offset -> value_type -> value
     (* raises Type, Bounds *)
+val load_simd_packed :
+  pack_size -> pack_simd -> memory -> address -> offset -> value_type -> value
+    (* raises Type, Bounds *)
 val store_packed :
   pack_size -> memory -> address -> offset -> value -> unit
     (* raises Type, Bounds *)

interpreter/syntax/ast.ml

@@ -111,6 +111,9 @@ type 'a memop =
 type loadop = (pack_size * extension) memop
 type storeop = pack_size memop
 
+type simd_loadop = (pack_size * pack_simd) memop
+type empty = |
+type simd_storeop = empty memop
 
 (* Expressions *)
 
@@ -142,6 +145,8 @@ and instr' =
   | GlobalSet of var                  (* write global variable *)
   | Load of loadop                    (* read memory at address *)
   | Store of storeop                  (* write memory at address *)
+  | SimdLoad of simd_loadop           (* read memory at address *)
+  | SimdStore of simd_storeop         (* write memory at address *)
   | MemorySize                        (* size of linear memory *)
   | MemoryGrow                        (* grow linear memory *)
   | Const of literal                  (* constant *)

interpreter/syntax/operators.ml

@@ -217,8 +217,29 @@ let memory_size = MemorySize
 let memory_grow = MemoryGrow
 
 (* SIMD *)
-let v128_load align offset = Load {ty = V128Type; align; offset; sz = None}
-let v128_store align offset = Store {ty = V128Type; align; offset; sz = None}
+let v128_load align offset = SimdLoad {ty = V128Type; align; offset; sz = None}
+let i16x8_load8x8_s align offset =
+  SimdLoad {ty = V128Type; align; offset; sz = Some (Pack64, Pack8x8 SX)}
+let i16x8_load8x8_u align offset =
+  SimdLoad {ty = V128Type; align; offset; sz = Some (Pack64, Pack8x8 ZX)}
+let i32x4_load16x4_s align offset =
+  SimdLoad {ty = V128Type; align; offset; sz = Some (Pack64, Pack16x4 SX)}
+let i32x4_load16x4_u align offset =
+  SimdLoad {ty = V128Type; align; offset; sz = Some (Pack64, Pack16x4 ZX)}
+let i64x2_load32x2_s align offset =
+  SimdLoad {ty = V128Type; align; offset; sz = Some (Pack64, Pack32x2 SX)}
+let i64x2_load32x2_u align offset =
+  SimdLoad {ty = V128Type; align; offset; sz = Some (Pack64, Pack32x2 ZX)}
+let v8x16_load_splat align offset =
+  SimdLoad {ty= V128Type; align; offset; sz = Some (Pack8, PackSplat)}
+let v16x8_load_splat align offset =
+  SimdLoad {ty= V128Type; align; offset; sz = Some (Pack16, PackSplat)}
+let v32x4_load_splat align offset =
+  SimdLoad {ty= V128Type; align; offset; sz = Some (Pack32, PackSplat)}
+let v64x2_load_splat align offset =
+  SimdLoad {ty= V128Type; align; offset; sz = Some (Pack64, PackSplat)}
+let v128_store align offset = SimdStore {ty = V128Type; align; offset; sz = None}
+
 let v128_not = Unary (V128 (V128Op.V128 V128Op.Not))
 let v128_and = Binary (V128 (V128Op.V128 V128Op.And))
 let v128_andnot = Binary (V128 (V128Op.V128 V128Op.AndNot))

interpreter/syntax/types.ml

@@ -16,9 +16,13 @@ type extern_type =
   | ExternMemoryType of memory_type
   | ExternGlobalType of global_type
 
-type pack_size = Pack8 | Pack16 | Pack32
+type pack_size = Pack8 | Pack16 | Pack32 | Pack64
 type extension = SX | ZX
-
+type pack_simd =
+  | PackSplat
+  | Pack8x8 of extension
+  | Pack16x4 of extension
+  | Pack32x2 of extension
 
 (* Attributes *)
 
@@ -31,7 +35,7 @@ let packed_size = function
   | Pack8 -> 1
   | Pack16 -> 2
   | Pack32 -> 4
-
+  | Pack64 -> 8
 
 (* Subtyping *)
 

interpreter/text/arrange.ml

@@ -77,6 +77,7 @@ let pack_size = function
   | Pack8 -> "8"
   | Pack16 -> "16"
   | Pack32 -> "32"
+  | Pack64 -> "64"
 
 let extension = function
   | SX -> "_s"
@@ -463,6 +464,8 @@ let rec instr e =
     | GlobalGet x -> "global.get " ^ var x, []
     | GlobalSet x -> "global.set " ^ var x, []
     | Load op -> loadop op, []
+    | SimdLoad op -> failwith "unimplemented SimdLoad arrange"
+    | SimdStore op -> failwith "unimplemented SimdStore arrange"
     | Store op -> storeop op, []
     | MemorySize -> "memory.size", []
     | MemoryGrow -> "memory.grow", []

interpreter/text/lexer.mll

@@ -288,6 +288,20 @@ rule token = parse
             (ext s i64_load8_s i64_load8_u (opt a 0))
             (ext s i64_load16_s i64_load16_u (opt a 1))
             (ext s i64_load32_s i64_load32_u (opt a 2)) o)) }
+  | "i16x8.load8x8_"(sign as s)
+  { LOAD (fun a o -> (ext s i16x8_load8x8_s i16x8_load8x8_u (opt a 3)) o) }
+  | "i32x4.load16x4_"(sign as s)
+  { LOAD (fun a o -> (ext s i32x4_load16x4_s i32x4_load16x4_u (opt a 3)) o) }
+  | "i64x2.load32x2_"(sign as s)
+  { LOAD (fun a o -> (ext s i64x2_load32x2_s i64x2_load32x2_u (opt a 3)) o) }
+  | "v8x16.load_splat"
+  { LOAD (fun a o -> (v8x16_load_splat (opt a 0)) o) }
+  | "v16x8.load_splat"
+  { LOAD (fun a o -> (v16x8_load_splat (opt a 1)) o) }
+  | "v32x4.load_splat"
+  { LOAD (fun a o -> (v32x4_load_splat (opt a 2)) o) }
+  | "v64x2.load_splat"
+  { LOAD (fun a o -> (v64x2_load_splat (opt a 3)) o) }
   | (ixx as t)".store"(mem_size as sz)
     { if t = "i32" && sz = "32" then error lexbuf "unknown operator";
       STORE (fun a o ->