arch extensions

Author: TheNumbat

backend/amd64/arch.ml

@@ -30,6 +30,9 @@ module Extension = struct
       | LZCNT
       | BMI
       | BMI2
+      | AVX
+      | AVX2
+      | AVX512F
 
     let rank = function
       | POPCNT -> 0
@@ -43,6 +46,9 @@ module Extension = struct
       | LZCNT -> 8
       | BMI -> 9
       | BMI2 -> 10
+      | AVX -> 11
+      | AVX2 -> 12
+      | AVX512F -> 13
 
     let compare left right = Int.compare (rank left) (rank right)
   end
@@ -62,6 +68,9 @@ module Extension = struct
     | LZCNT -> "LZCNT"
     | BMI -> "BMI"
     | BMI2 -> "BMI2"
+    | AVX -> "AVX"
+    | AVX2 -> "AVX2"
+    | AVX512F -> "AVX512F"
 
   let generation = function
     | POPCNT -> "Nehalem+"
@@ -75,18 +84,29 @@ module Extension = struct
     | LZCNT -> "Haswell+"
     | BMI -> "Haswell+"
     | BMI2 -> "Haswell+"
+    | AVX -> "Sandybridge+"
+    | AVX2 -> "Haswell+"
+    | AVX512F -> "SkylakeXeon+"
 
   let enabled_by_default = function
     | SSE3 | SSSE3 | SSE4_1 | SSE4_2
-    | POPCNT | CLMUL | LZCNT | BMI | BMI2 -> true
-    | PREFETCHW | PREFETCHWT1 -> false
+    | POPCNT | CLMUL | LZCNT | BMI | BMI2 | AVX | AVX2 -> true
+    | PREFETCHW | PREFETCHWT1 | AVX512F -> false
 
-  let all = Set.of_list [ POPCNT; PREFETCHW; PREFETCHWT1; SSE3; SSSE3; SSE4_1; SSE4_2; CLMUL; LZCNT; BMI; BMI2 ]
+  let all = Set.of_list [ POPCNT; PREFETCHW; PREFETCHWT1; SSE3; SSSE3; SSE4_1; SSE4_2; CLMUL; LZCNT; BMI; BMI2; AVX; AVX2; AVX512F ]
   let config = ref (Set.filter enabled_by_default all)
 
   let enabled t = Set.mem t !config
   let disabled t = not (enabled t)
 
+  let allow_vec256 () = List.exists (fun t -> enabled t) [AVX; AVX2; AVX512F]
+  let allow_vec512 () = List.exists (fun t -> enabled t) [AVX512F]
+
+  let require_vec256 () =
+    if not (allow_vec256 ()) then Misc.fatal_error "AVX or AVX512 is required for 256-bit vectors"
+  let require_vec512 () =
+    if not (allow_vec512 ()) then Misc.fatal_error "AVX512 is required for 512-bit vectors"
+
   let args =
     let y t = "-f" ^ (name t |> String.lowercase_ascii) in
     let n t = "-fno-" ^ (name t |> String.lowercase_ascii) in

backend/amd64/arch.mli

@@ -31,11 +31,19 @@ module Extension : sig
              to Haswell, i.e. they do not cause an illegal instruction fault.
              That means code using LZCNT/TZCNT will silently produce wrong results. *)
     | BMI2
+    | AVX
+    | AVX2
+    | AVX512F
 
   val name : t -> string
 
   val enabled : t -> bool
   val available : unit -> t list
+
+  val allow_vec256 : unit -> bool
+  val allow_vec512 : unit -> bool
+  val require_vec256 : unit -> unit
+  val require_vec512 : unit -> unit
 end
 
 val trap_notes : bool ref

backend/amd64/emit.ml

@@ -144,8 +144,12 @@ let register_name typ r : X86_ast.arg =
   match (typ : Cmm.machtype_component) with
   | Int | Val | Addr -> Reg64 int_reg_name.(r)
   | Float | Float32 | Vec128 | Valx2 -> Regf xmm_reg_name.(r - 100)
-  | Vec256 -> Regf ymm_reg_name.(r - 100)
-  | Vec512 -> Regf zmm_reg_name.(r - 100)
+  | Vec256 ->
+    Arch.Extension.require_vec256 ();
+    Regf ymm_reg_name.(r - 100)
+  | Vec512 ->
+    Arch.Extension.require_vec512 ();
+    Regf zmm_reg_name.(r - 100)
 
 let phys_rax = phys_reg Int 0
 
@@ -375,8 +379,12 @@ let x86_data_type_for_stack_slot : Cmm.machtype_component -> X86_ast.data_type =
   function
   | Float -> REAL8
   | Vec128 -> VEC128
-  | Vec256 -> VEC256
-  | Vec512 -> VEC512
+  | Vec256 ->
+    Arch.Extension.require_vec256 ();
+    VEC256
+  | Vec512 ->
+    Arch.Extension.require_vec512 ();
+    VEC512
   | Valx2 -> VEC128
   | Int | Addr | Val -> QWORD
   | Float32 -> REAL4

backend/amd64/proc.ml

@@ -107,17 +107,32 @@ let hard_vec256_reg = Array.map (fun r -> {r with Reg.typ = Vec256}) hard_float_
 let hard_vec512_reg = Array.map (fun r -> {r with Reg.typ = Vec512}) hard_float_reg
 let hard_float32_reg = Array.map (fun r -> {r with Reg.typ = Float32}) hard_float_reg
 
+let add_hard_vec256_regs list ~f =
+  if Arch.Extension.allow_vec256 ()
+  then f hard_vec256_reg :: list else list
+
+let add_hard_vec512_regs list ~f =
+  if Arch.Extension.allow_vec512 ()
+  then f hard_vec512_reg :: list else list
+
 let all_phys_regs =
-  Array.concat [hard_int_reg; hard_float_reg; hard_float32_reg; hard_vec128_reg; hard_vec256_reg; hard_vec512_reg]
+  [hard_int_reg; hard_float_reg; hard_float32_reg; hard_vec128_reg]
+  |> add_hard_vec256_regs ~f:(fun regs -> regs)
+  |> add_hard_vec512_regs ~f:(fun regs -> regs)
+  |> Array.concat
 
 let phys_reg ty n =
   match (ty : machtype_component) with
   | Int | Addr | Val -> hard_int_reg.(n)
   | Float -> hard_float_reg.(n - 100)
   | Float32 -> hard_float32_reg.(n - 100)
   | Vec128 | Valx2 -> hard_vec128_reg.(n - 100)
-  | Vec256 -> hard_vec256_reg.(n - 100)
-  | Vec512 -> hard_vec512_reg.(n - 100)
+  | Vec256 ->
+    Arch.Extension.require_vec256 ();
+    hard_vec256_reg.(n - 100)
+  | Vec512 ->
+    Arch.Extension.require_vec512 ();
+    hard_vec512_reg.(n - 100)
 
 let rax = phys_reg Int 0
 let rdi = phys_reg Int 2
@@ -128,9 +143,14 @@ let r11 = phys_reg Int 11
 let rbp = phys_reg Int 12
 
 (* CSE needs to know that all versions of xmm15 are destroyed. *)
-let destroy_xmm n =
-  [| phys_reg Float (100 + n); phys_reg Float32 (100 + n);
-     phys_reg Vec128 (100 + n); phys_reg Vec256 (100 + n); phys_reg Vec512 (100 + n) |]
+let destroy_xmm =
+  let types =
+    ([ Float; Float32; Vec128 ] : machtype_component list)
+    |> add_hard_vec256_regs ~f:(fun _ -> Vec256)
+    |> add_hard_vec512_regs ~f:(fun _ -> Vec512)
+    |> Array.of_list
+  in
+  fun n -> Array.map (fun t -> phys_reg t (100 + n)) types
 
 let destroyed_by_plt_stub =
   if not X86_proc.use_plt then [| |] else [| r10; r11 |]
@@ -189,6 +209,7 @@ let calling_conventions
         ofs := !ofs + size_vec128
       end
     | Vec256 ->
+      Arch.Extension.require_vec256 ();
       if !float <= last_float then begin
         loc.(i) <- phys_reg Vec256 !float;
         incr float
@@ -198,6 +219,7 @@ let calling_conventions
         ofs := !ofs + size_vec256
       end
     | Vec512 ->
+      Arch.Extension.require_vec512 ();
       if !float <= last_float then begin
         loc.(i) <- phys_reg Vec512 !float;
         incr float
@@ -390,21 +412,23 @@ let int_regs_destroyed_at_c_call =
 
 let destroyed_at_c_call_win64 =
   (* Win64: rbx, rbp, rsi, rdi, r12-r15, xmm6-xmm15 preserved *)
-  Array.concat [
-    Array.map (phys_reg Int) int_regs_destroyed_at_c_call_win64;
+  [ Array.map (phys_reg Int) int_regs_destroyed_at_c_call_win64;
     Array.sub hard_float_reg 0 6;
     Array.sub hard_float32_reg 0 6;
-    Array.sub hard_vec128_reg 0 6
-  ]
+    Array.sub hard_vec128_reg 0 6 ]
+  |> add_hard_vec256_regs ~f:(fun regs -> Array.sub regs 0 6)
+  |> add_hard_vec512_regs ~f:(fun regs -> Array.sub regs 0 6)
+  |> Array.concat
 
 let destroyed_at_c_call_unix =
   (* Unix: rbx, rbp, r12-r15 preserved *)
-  Array.concat [
-      Array.map (phys_reg Int) int_regs_destroyed_at_c_call;
-      hard_float_reg;
-      hard_float32_reg;
-      hard_vec128_reg
-  ]
+  [ Array.map (phys_reg Int) int_regs_destroyed_at_c_call;
+    hard_float_reg;
+    hard_float32_reg;
+    hard_vec128_reg ]
+  |> add_hard_vec256_regs ~f:(fun regs -> regs)
+  |> add_hard_vec512_regs ~f:(fun regs -> regs)
+  |> Array.concat
 
 let destroyed_at_c_call =
   (* C calling conventions preserve rbx, but it is clobbered

flambda-backend/tests/simd/dune

@@ -55,7 +55,7 @@
  (libraries simd_test_helpers stdlib_stable stdlib_upstream_compatible)
  (foreign_archives stubs)
  (ocamlopt_flags
-  (:standard -extension simd_alpha)))
+  (:standard -extension simd_beta)))
 
 (rule
  (enabled_if
@@ -219,7 +219,7 @@
  (libraries simd_test_helpers stdlib_stable stdlib_upstream_compatible)
  (foreign_archives stubs)
  (ocamlopt_flags
-  (:standard -nodynlink -extension simd_alpha)))
+  (:standard -nodynlink -extension simd_beta)))
 
 (rule
  (enabled_if
@@ -391,7 +391,7 @@
   (<> %{system} macosx))
  (foreign_archives stubs)
  (ocamlopt_flags
-  (:standard -internal-assembler -extension simd_alpha)))
+  (:standard -internal-assembler -extension simd_beta)))
 
 (rule
  (enabled_if