Vectorize [Specific.Istore_int] (#3450)

Used for array initialization (amd64)

Author: gretay-js

backend/amd64/simd_selection.ml

@@ -855,9 +855,54 @@ let vectorize_operation (width_type : Vectorize_utils.Width_in_bits.t)
       | W32 -> None (* See previous comment *)
       | W16 -> None
       | W8 -> None)
-    | Istore_int _ | Ioffset_loc _ | Ifloatarithmem _ | Ibswap _ | Irdtsc
-    | Irdpmc | Ilfence | Isfence | Imfence | Ipause | Isimd _ | Iprefetch _
-    | Icldemote _ ->
+    | Ifloatarithmem _ -> None
+    | Istore_int (_n, addressing_mode, is_assignment) -> (
+      if not (Vectorize_utils.Width_in_bits.equal width_type W64)
+      then None
+      else
+        let extract_store_int_imm (op : Operation.t) =
+          match op with
+          | Specific (Istore_int (n, _addr, _is_assign)) -> Int64.of_nativeint n
+          | Specific
+              ( Ifloatarithmem _ | Ioffset_loc _ | Iprefetch _ | Icldemote _
+              | Irdtsc | Irdpmc | Ilfence | Isfence | Imfence | Ipause | Isimd _
+              | Ilea _ | Ibswap _ | Isextend32 | Izextend32 )
+          | Intop_imm _ | Move | Load _ | Store _ | Intop _ | Alloc _
+          | Reinterpret_cast _ | Static_cast _ | Spill | Reload | Const_int _
+          | Const_float32 _ | Const_float _ | Const_symbol _ | Const_vec128 _
+          | Stackoffset _ | Intop_atomic _ | Floatop _ | Csel _
+          | Probe_is_enabled _ | Opaque | Begin_region | End_region
+          | Name_for_debugger _ | Dls_get | Poll ->
+            assert false
+        in
+        let consts = List.map extract_store_int_imm cfg_ops in
+        match create_const_vec consts with
+        | None -> None
+        | Some [const_instruction] ->
+          let num_args_addressing = Arch.num_args_addressing addressing_mode in
+          assert (arg_count = num_args_addressing);
+          assert (res_count = 0);
+          assert (Array.length const_instruction.results = 1);
+          let new_reg = Vectorize_utils.Vectorized_instruction.New 0 in
+          const_instruction.results.(0) <- new_reg;
+          let address_args =
+            Array.init num_args_addressing (fun i ->
+                Vectorize_utils.Vectorized_instruction.Original i)
+          in
+          let store_operation =
+            Operation.Store
+              (Onetwentyeight_unaligned, addressing_mode, is_assignment)
+          in
+          let store_instruction : Vectorize_utils.Vectorized_instruction.t =
+            { operation = store_operation;
+              arguments = Array.append [| new_reg |] address_args;
+              results = [||]
+            }
+          in
+          Some [const_instruction; store_instruction]
+        | Some _ -> None)
+    | Ioffset_loc _ | Ibswap _ | Irdtsc | Irdpmc | Ilfence | Isfence | Imfence
+    | Ipause | Isimd _ | Iprefetch _ | Icldemote _ ->
       None)
   | Alloc _ | Reinterpret_cast _ | Static_cast _ | Spill | Reload
   | Const_float32 _ | Const_float _ | Const_symbol _ | Const_vec128 _

backend/amd64/vectorize_specific.ml

@@ -53,3 +53,13 @@ let memory_access : Arch.specific_operation -> Memory_access.t option =
        the moment. *)
     if Simd.is_pure op then None else create Memory_access.Arbitrary
   | Ilea _ | Ibswap _ | Isextend32 | Izextend32 -> None
+
+let is_seed_store :
+    Arch.specific_operation -> Vectorize_utils.Width_in_bits.t option =
+ fun op ->
+  match op with
+  | Istore_int _ -> Some W64
+  | Ifloatarithmem _ | Ioffset_loc _ | Iprefetch _ | Icldemote _ | Irdtsc
+  | Irdpmc | Ilfence | Isfence | Imfence | Ipause | Isimd _ | Ilea _ | Ibswap _
+  | Isextend32 | Izextend32 ->
+    None

backend/arm64/vectorize_specific.ml

@@ -21,3 +21,10 @@ let memory_access : Arch.specific_operation -> Memory_access.t option =
     (* Conservative. we don't have any specific operations with memory
        operations at the moment. *)
     if Arch.operation_is_pure op then None else create Memory_access.Arbitrary
+
+let is_seed_store (op : Arch.specific_operation) =
+  match op with
+  | Ifar_poll _ | Ifar_alloc _ | Ishiftarith _ | Imuladd | Imulsub | Inegmulf
+  | Imuladdf | Inegmuladdf | Imulsubf | Inegmulsubf | Isqrtf | Ibswap _
+  | Imove32 | Isignext _ ->
+    None

backend/cfg/vectorize.ml

@@ -2251,13 +2251,15 @@ end = struct
       | None -> None
       | Some op -> (
         match op with
-        | Store (chunk, _, _) -> Some chunk
+        | Store (chunk, _, _) ->
+          Some (Vectorize_utils.Width_in_bits.of_memory_chunk chunk)
+        | Specific s -> Vectorize_specific.is_seed_store s
         | Alloc _ | Load _ | Move | Reinterpret_cast _ | Static_cast _ | Spill
         | Reload | Const_int _ | Const_float32 _ | Const_float _
         | Const_symbol _ | Const_vec128 _ | Stackoffset _ | Intop _
         | Intop_imm _ | Intop_atomic _ | Floatop _ | Csel _ | Probe_is_enabled _
-        | Opaque | Begin_region | End_region | Specific _ | Name_for_debugger _
-        | Dls_get | Poll ->
+        | Opaque | Begin_region | End_region | Name_for_debugger _ | Dls_get
+        | Poll ->
           None)
 
     let from_block (block : Block.t) deps : t list =
@@ -2280,8 +2282,7 @@ end = struct
         DLL.fold_right body ~init:[] ~f:(fun i acc ->
             let i = Instruction.basic i in
             match is_store i with
-            | Some chunk ->
-              (Vectorize_utils.Width_in_bits.of_memory_chunk chunk, i) :: acc
+            | Some width -> (width, i) :: acc
             | None -> acc)
       in
       Format.(