Refactor stack classes (#3784)

asmcomp/asmgen.ml

@@ -302,7 +302,11 @@ let register_allocator fd : register_allocator =
 let available_regs ~stack_slots ~f x =
   (* Skip DWARF variable range generation for complicated functions to avoid
      high compilation speed penalties *)
-  let total_num_stack_slots = Array.fold_left ( + ) 0 (stack_slots x) in
+  let fun_num_stack_slots = stack_slots x in
+  let total_num_stack_slots =
+    Stack_class.Tbl.fold fun_num_stack_slots ~init:0
+      ~f:(fun _stack_class num acc -> acc + num)
+  in
   if total_num_stack_slots > !Dwarf_flags.dwarf_max_function_complexity
   then x
   else f x

backend/.ocamlformat-enable

@@ -3,12 +3,14 @@ amd64/selection.ml
 amd64/selection_utils.ml
 amd64/simd*.ml
 amd64/stack_check.ml
+amd64/stack_class.ml
 amd64/vectorize_specific.ml
 arm64/cfg_selection.ml
 arm64/selection.ml
 arm64/selection_utils.ml
 arm64/simd*.ml
 arm64/stack_check.ml
+arm64/stack_class.ml
 arm64/vectorize_specific.ml
 arm64_ast.ml
 arm64_ast.mli
@@ -53,6 +55,9 @@ selection.mli
 simple_operation.ml
 simple_operation.mli
 stack_check.mli
+stack_class.mli
+stack_class_utils.ml
+stack_class_utils.mli
 zero_alloc_checker.ml
 zero_alloc_checker.mli
 zero_alloc_info.ml

backend/amd64/emit.ml

@@ -113,7 +113,7 @@ let fastcode_flag = ref true
 (* Layout of the stack frame *)
 let stack_offset = ref 0
 
-let num_stack_slots = Array.make Proc.num_stack_slot_classes 0
+let num_stack_slots = Stack_class.Tbl.make 0
 
 let prologue_required = ref false
 
@@ -317,7 +317,7 @@ let reg = function
       let ofs = n + Domainstate.(idx_of_field Domain_extra_params) * 8 in
       mem64 (x86_data_type_for_stack_slot ty) ofs R14
   | { loc = Stack s; typ = ty } as r ->
-      let ofs = slot_offset s (stack_slot_class r.typ) in
+      let ofs = slot_offset s (Stack_class.of_machtype r.typ) in
       mem64 (x86_data_type_for_stack_slot ty) ofs RSP
   | { loc = Unknown } ->
       assert false
@@ -340,7 +340,7 @@ let reg_low_32_name = Array.map (fun r -> Reg32 r) int_reg_name
 let emit_subreg tbl typ r =
   match r.loc with
   | Reg.Reg r when r < 13 -> tbl.(r)
-  | Stack s -> mem64 typ (slot_offset s (stack_slot_class r.Reg.typ)) RSP
+  | Stack s -> mem64 typ (slot_offset s (Stack_class.of_machtype r.Reg.typ)) RSP
   | _ -> assert false
 
 let arg8 i n = emit_subreg reg_low_8_name BYTE i.arg.(n)
@@ -383,13 +383,13 @@ let record_frame_label live dbg =
           assert (Proc.gc_regs_offset reg = r);
           live_offset := ((r lsl 1) + 1) :: !live_offset
       | {typ = Val; loc = Stack s} as reg ->
-          live_offset := slot_offset s (stack_slot_class reg.typ) :: !live_offset
+          live_offset := slot_offset s (Stack_class.of_machtype reg.typ) :: !live_offset
       | {typ = Valx2; loc = Reg r} as reg ->
           let n = Proc.gc_regs_offset reg in
           let encode n = ((n lsl 1) + 1) in
           live_offset := encode n :: encode (n + 1) :: !live_offset
       | {typ = Valx2; loc = Stack s} as reg ->
-          let n = slot_offset s (stack_slot_class reg.typ)  in
+          let n = slot_offset s (Stack_class.of_machtype reg.typ)  in
           live_offset := n :: n + Arch.size_addr :: !live_offset
       | {typ = Addr} as r ->
           Misc.fatal_error ("bad GC root " ^ Reg.name r)
@@ -867,7 +867,7 @@ let tailrec_entry_point = ref None
 type probe =
   {
     stack_offset: int;
-    num_stack_slots: int array;
+    num_stack_slots: int Stack_class.Tbl.t;
     (* Record frame info held in the corresponding mutable variables. *)
     probe_label: label;
     (* Probe site, recorded in .note.stapsdt section
@@ -2045,7 +2045,7 @@ let emit_instr ~first ~fallthrough i =
       { probe_label;
         probe_insn = i;
         stack_offset = !stack_offset;
-        num_stack_slots = Array.copy num_stack_slots;
+        num_stack_slots = Stack_class.Tbl.copy num_stack_slots;
       }
     in
     probes := probe :: !probes;
@@ -2224,9 +2224,7 @@ let fundecl fundecl =
   local_realloc_sites := [];
   clear_safety_checks ();
   clear_stack_realloc ();
-  for i = 0 to Proc.num_stack_slot_classes - 1 do
-    num_stack_slots.(i) <- fundecl.fun_num_stack_slots.(i);
-  done;
+  Stack_class.Tbl.copy_values ~from:fundecl.fun_num_stack_slots ~to_:num_stack_slots;
   prologue_required := fundecl.fun_prologue_required;
   frame_required := fundecl.fun_frame_required;
   all_functions := fundecl :: !all_functions;
@@ -2473,9 +2471,7 @@ let emit_probe_handler_wrapper p =
      recall that the wrapper does however have its own frame.) *)
   frame_required := true;
   stack_offset := p.stack_offset;
-  for i = 0 to Proc.num_stack_slot_classes - 1 do
-    num_stack_slots.(i) <- p.num_stack_slots.(i);
-  done;
+  Stack_class.Tbl.copy_values ~from:p.num_stack_slots ~to_:num_stack_slots;
   (* Account for the return address that is now pushed on the stack. *)
   stack_offset := !stack_offset + 8;
   (* Emit function entry code *)
@@ -2613,7 +2609,7 @@ let emit_probe_notes0 () =
     let arg_name =
       match arg.loc with
       | Stack s ->
-        Printf.sprintf "%d(%%rsp)" (slot_offset s (stack_slot_class arg.Reg.typ))
+        Printf.sprintf "%d(%%rsp)" (slot_offset s (Stack_class.of_machtype arg.Reg.typ))
       | Reg reg -> Proc.register_name arg.Reg.typ reg
       | Unknown ->
         Misc.fatal_errorf "Cannot create probe: illegal argument: %a"

backend/amd64/proc.ml

@@ -104,13 +104,6 @@ let register_class r =
   | Val | Int | Addr -> 0
   | Float | Float32 | Vec128 | Valx2 -> 1
 
-let num_stack_slot_classes = 3
-
-let stack_slot_class typ =
-  match (typ : machtype_component) with
-  | Val | Addr | Int -> 0
-  | Float | Float32 -> 1
-  | Vec128 | Valx2 -> 2
 
 let types_are_compatible left right =
   match left.typ, right.typ with
@@ -121,13 +114,6 @@ let types_are_compatible left right =
     true
   | (Int | Val | Addr | Float | Float32 | Vec128 | Valx2), _ -> false
 
-let stack_class_tag c =
-  match c with
-  | 0 -> "i"
-  | 1 -> "f"
-  | 2 -> "x"
-  | c -> Misc.fatal_errorf "Unspecified stack slot class %d" c
-
 let num_available_registers = [| 13; 16 |]
 
 let first_available_register = [| 0; 100 |]
@@ -676,9 +662,9 @@ let trap_frame_size_in_bytes = 16
 
 let frame_required ~fun_contains_calls ~fun_num_stack_slots =
   fp || fun_contains_calls ||
-  fun_num_stack_slots.(0) > 0 ||
-  fun_num_stack_slots.(1) > 0 ||
-  fun_num_stack_slots.(2) > 0
+  Stack_class.Tbl.exists
+    fun_num_stack_slots
+    ~f:(fun _stack_class num -> num > 0)
 
 let prologue_required ~fun_contains_calls ~fun_num_stack_slots =
   frame_required ~fun_contains_calls ~fun_num_stack_slots
@@ -691,9 +677,7 @@ let frame_size ~stack_offset ~contains_calls ~num_stack_slots =
     let sz =
       (stack_offset
        + 8
-       + 8 * num_stack_slots.(0)
-       + 8 * num_stack_slots.(1)
-       + 16 * num_stack_slots.(2)
+       + Stack_class.Tbl.total_size_in_bytes num_stack_slots
        + (if fp then 8 else 0))
     in Misc.align sz 16
   end else
@@ -716,14 +700,8 @@ let slot_offset loc ~stack_class ~stack_offset ~fun_contains_calls
         + n)
   | Local n ->
       Bytes_relative_to_stack_pointer (
-        stack_offset +
-          (* Preserves original ordering (int -> float) *)
-          match stack_class with
-          | 2 -> n * 16
-          | 0 -> fun_num_stack_slots.(2) * 16 + n * 8
-          | 1 -> fun_num_stack_slots.(2) * 16
-                 + fun_num_stack_slots.(0) * 8 + n * 8
-          | _ -> Misc.fatal_errorf "Unknown register class %d" stack_class)
+        stack_offset + Stack_class.Tbl.offset_in_bytes fun_num_stack_slots ~stack_class ~slot:n
+)
   | Outgoing n -> Bytes_relative_to_stack_pointer n
   | Domainstate n ->
       Bytes_relative_to_domainstate_pointer (

backend/amd64/stack_class.ml

@@ -0,0 +1,49 @@
+[@@@ocaml.warning "+a-30-40-41-42"]
+
+module T = struct
+  type t =
+    | Int64
+    | Float64
+    | Vector128
+
+  let all = [Int64; Float64; Vector128]
+
+  (* Preserves original ordering (int -> float) *)
+  let frame_order = [| Vector128; Int64; Float64 |]
+
+  let equal : t -> t -> bool =
+   fun left right ->
+    match left, right with
+    | Int64, Int64 -> true
+    | Float64, Float64 -> true
+    | Vector128, Vector128 -> true
+    | (Int64 | Float64 | Vector128), _ -> false
+
+  let hash : t -> int = function Int64 -> 0 | Float64 -> 1 | Vector128 -> 2
+
+  let tag : t -> string = function
+    | Int64 -> "i"
+    | Float64 -> "f"
+    | Vector128 -> "x"
+
+  let print : Format.formatter -> t -> unit =
+   fun ppf stack_class ->
+    Format.fprintf ppf "%s"
+      (match stack_class with
+      | Int64 -> "int64"
+      | Float64 -> "float64"
+      | Vector128 -> "vector128")
+
+  let size_in_bytes : t -> int = function
+    | Int64 -> 8
+    | Float64 -> 8
+    | Vector128 -> 16
+
+  let of_machtype : Cmm.machtype_component -> t = function
+    | Val | Int | Addr -> Int64
+    | Float | Float32 -> Float64
+    | Vec128 | Valx2 -> Vector128
+end
+
+include T
+module Tbl = Stack_class_utils.Make_tbl (T)

backend/arm64/emit.ml

@@ -106,7 +106,7 @@ let femit_wreg out = function
 
 let stack_offset = ref 0
 
-let num_stack_slots = Array.make Proc.num_stack_slot_classes 0
+let num_stack_slots = Stack_class.Tbl.make 0
 
 let prologue_required = ref false
 
@@ -139,7 +139,7 @@ let femit_stack out r =
       let ofs = n + Domainstate.(idx_of_field Domain_extra_params) * 8 in
       Printf.fprintf out "[%a, #%a]" femit_reg reg_domain_state_ptr femit_int ofs
   | Stack ((Local _ | Incoming _ | Outgoing _) as s) ->
-      let ofs = slot_offset s (stack_slot_class r.typ) in
+      let ofs = slot_offset s (Stack_class.of_machtype r.typ) in
       Printf.fprintf out "[sp, #%a]" femit_int ofs
   | Reg _ | Unknown -> fatal_error "Emit.emit_stack"
 
@@ -174,7 +174,7 @@ let record_frame_label live dbg =
       | {typ = Val; loc = Reg r} ->
           live_offset := ((r lsl 1) + 1) :: !live_offset
       | {typ = Val; loc = Stack s} as reg ->
-          live_offset := slot_offset s (stack_slot_class reg.typ) :: !live_offset
+          live_offset := slot_offset s (Stack_class.of_machtype reg.typ) :: !live_offset
       | {typ = Addr} as r ->
           Misc.fatal_error ("bad GC root " ^ Reg.name r)
       | { typ = Valx2; } as r ->
@@ -1680,9 +1680,7 @@ let fundecl fundecl =
   call_gc_sites := [];
   local_realloc_sites := [];
   clear_stack_realloc ();
-  for i = 0 to Proc.num_stack_slot_classes - 1 do
-    num_stack_slots.(i) <- fundecl.fun_num_stack_slots.(i);
-  done;
+  Stack_class.Tbl.copy_values ~from:fundecl.fun_num_stack_slots ~to_:num_stack_slots;
   prologue_required := fundecl.fun_prologue_required;
   contains_calls := fundecl.fun_contains_calls;
   emit_named_text_section !function_name;

backend/arm64/proc.ml

@@ -83,15 +83,6 @@ let register_class_of_machtype_component typ =
 let register_class r =
   register_class_of_machtype_component r.typ
 
-let num_stack_slot_classes = 3
-
-let stack_slot_class typ =
-  match (typ : Cmm.machtype_component) with
-  | Val | Int | Addr  -> 0
-  | Float | Float32 -> 1
-  | Vec128 -> 2
-  | Valx2 -> 2
-
 let types_are_compatible left right =
   match left.typ, right.typ with
   | (Int | Val | Addr), (Int | Val | Addr)
@@ -101,13 +92,6 @@ let types_are_compatible left right =
   | Valx2,Valx2 -> true
   | (Int | Val | Addr | Float | Float32 | Vec128 | Valx2), _ -> false
 
-let stack_class_tag c =
-  match c with
-  | 0 -> "i"
-  | 1 -> "f"
-  | 2 -> "x"
-  | c -> Misc.fatal_errorf "Unspecified stack slot class %d" c
-
 let num_available_registers =
   [| 23; 32 |] (* first 23 int regs allocatable; all float regs allocatable *)
 
@@ -448,9 +432,7 @@ let is_destruction_point ~(more_destruction_points : bool) (terminator : Cfg_int
 (* Layout of the stack *)
 
 let initial_stack_offset ~num_stack_slots ~contains_calls =
-  (8 * num_stack_slots.(0))
-  + (8 * num_stack_slots.(1))
-  + (16 * num_stack_slots.(2))
+  Stack_class.Tbl.total_size_in_bytes num_stack_slots
   + if contains_calls then 8 else 0
 
 let trap_frame_size_in_bytes = 16
@@ -462,9 +444,7 @@ let frame_size ~stack_offset ~contains_calls ~num_stack_slots =
   Misc.align sz 16
 
 let frame_required ~fun_contains_calls ~fun_num_stack_slots =
-  fun_contains_calls
-    || fun_num_stack_slots.(0) > 0
-     || fun_num_stack_slots.(1) > 0
+  fun_contains_calls || Stack_class.Tbl.exists fun_num_stack_slots ~f:(fun _stack_class num -> num > 0)
 
 let prologue_required ~fun_contains_calls ~fun_num_stack_slots =
   frame_required ~fun_contains_calls ~fun_num_stack_slots
@@ -487,13 +467,7 @@ let slot_offset (loc : Reg.stack_location) ~stack_class ~stack_offset
   | Local n ->
       let offset =
         stack_offset +
-        (* Preserves original ordering: int below float. *)
-        (match stack_class with
-        | 2 -> n * 16
-        | 0 -> fun_num_stack_slots.(2) * 16 + n * 8
-        | 1 -> fun_num_stack_slots.(2) * 16 +
-               fun_num_stack_slots.(0) * 8 + n * 8
-        | _ -> Misc.fatal_errorf "Unknown stack class %d" stack_class)
+        Stack_class.Tbl.offset_in_bytes fun_num_stack_slots ~stack_class ~slot:n
       in
       Bytes_relative_to_stack_pointer offset
   | Outgoing n ->