Backend support for local allocations

backend/CSEgen.ml

@@ -249,6 +249,7 @@ method class_of_operation op =
   | Ispecific _ -> Op_other
   | Iname_for_debugger _ -> Op_pure
   | Iprobe_is_enabled _ -> Op_other
+  | Ibeginregion | Iendregion -> Op_other
 
 (* Operations that are so cheap that it isn't worth factoring them. *)
 

backend/afl_instrument.ml

@@ -88,6 +88,8 @@ and instrument = function
      in
      Ccatch (isrec, cases, instrument body)
   | Cexit (ex, args, traps) -> Cexit (ex, List.map instrument args, traps)
+  | Cregion e -> Cregion (instrument e)
+  | Ctail e -> Ctail (instrument e)
 
   (* these are base cases and have no logging *)
   | Cconst_int _ | Cconst_natint _ | Cconst_float _

backend/amd64/CSE.ml

@@ -47,6 +47,7 @@ method! class_of_operation op =
   | Istackoffset _ | Iload _ | Istore _ | Ialloc _
   | Iintop _ | Iintop_imm _
   | Iname_for_debugger _ | Iprobe _ | Iprobe_is_enabled _ | Iopaque
+  | Ibeginregion | Iendregion
     -> super#class_of_operation op
 
 end

backend/amd64/emit.mlp

@@ -305,6 +305,19 @@ let emit_call_gc gc =
   def_label gc.gc_frame;
   I.jmp (label gc.gc_return_lbl)
 
+(* Record calls to local stack reallocation *)
+
+type local_realloc_call =
+  { lr_lbl: label;
+    lr_return_lbl: label; }
+
+let local_realloc_sites = ref ([] : local_realloc_call list)
+
+let emit_local_realloc lr =
+  def_label lr.lr_lbl;
+  emit_call "caml_call_local_realloc";
+  I.jmp (label lr.lr_return_lbl)
+
 (* Record calls to caml_ml_array_bound_error.
    In -g mode we maintain one call to
    caml_ml_array_bound_error per bound check site.  Without -g, we can share
@@ -760,7 +773,7 @@ let emit_instr fallthrough i =
       | Double ->
           I.movsd (arg i 0) (addressing addr REAL8 i 1)
       end
-  | Lop(Ialloc { bytes = n; dbginfo }) ->
+  | Lop(Ialloc { bytes = n; dbginfo; mode = Alloc_heap }) ->
       assert (n <= (Config.max_young_wosize + 1) * Arch.size_addr);
       if !fastcode_flag then begin
         I.sub (int n) r15;
@@ -790,6 +803,21 @@ let emit_instr fallthrough i =
         def_label label;
         I.lea (mem64 NONE 8 R15) (res i 0)
       end
+  | Lop(Ialloc { bytes = n; dbginfo=_; mode = Alloc_local }) ->
+      let r = res i 0 in
+      I.mov (domain_field Domainstate.Domain_local_sp) r;
+      I.sub (int n) r;
+      I.mov r (domain_field Domainstate.Domain_local_sp);
+      I.cmp (domain_field Domainstate.Domain_local_limit) r;
+      let lbl_call = new_label () in
+      I.j L (label lbl_call);
+      let lbl_after_alloc = new_label () in
+      def_label lbl_after_alloc;
+      I.add (domain_field Domainstate.Domain_local_top) r;
+      I.add (int 8) r;
+      local_realloc_sites :=
+        { lr_lbl = lbl_call;
+          lr_return_lbl = lbl_after_alloc } :: !local_realloc_sites
   | Lop(Iintop(Icomp cmp)) ->
       I.cmp (arg i 1) (arg i 0);
       I.set (cond cmp) al;
@@ -963,6 +991,10 @@ let emit_instr fallthrough i =
       | High -> T0
     in
     I.prefetch is_write locality (addressing addr QWORD i 0)
+  | Lop(Ibeginregion) ->
+      I.mov (domain_field Domainstate.Domain_local_sp) (res i 0)
+  | Lop(Iendregion) ->
+      I.mov (arg i 0) (domain_field Domainstate.Domain_local_sp)
   | Lop (Iname_for_debugger _) -> ()
   | Lop (Iprobe _) ->
     let probe_label = new_label () in
@@ -1149,6 +1181,7 @@ let fundecl fundecl =
   tailrec_entry_point := fundecl.fun_tailrec_entry_point_label;
   stack_offset := 0;
   call_gc_sites := [];
+  local_realloc_sites := [];
   bound_error_sites := [];
   bound_error_call := 0;
   for i = 0 to Proc.num_register_classes - 1 do
@@ -1172,6 +1205,7 @@ let fundecl fundecl =
   cfi_startproc ();
   emit_all true fundecl.fun_body;
   List.iter emit_call_gc !call_gc_sites;
+  List.iter emit_local_realloc !local_realloc_sites;
   emit_call_bound_errors ();
   if !frame_required then begin
     let n = frame_size() - 8 - (if fp then 8 else 0) in

backend/amd64/proc.ml

@@ -337,6 +337,7 @@ let destroyed_at_oper = function
        | Iname_for_debugger _ | Iprobe _| Iprobe_is_enabled _ | Iopaque)
   | Iend | Ireturn _ | Iifthenelse (_, _, _) | Icatch (_, _, _, _)
   | Iexit _ | Iraise _
+  | Iop(Ibeginregion | Iendregion)
     ->
     if fp then
 (* prevent any use of the frame pointer ! *)
@@ -362,6 +363,7 @@ let safe_register_pressure = function
   | Istackoffset _ | Iload (_, _) | Istore (_, _, _)
   | Iintop _ | Iintop_imm (_, _) | Ispecific _ | Iname_for_debugger _
   | Iprobe _ | Iprobe_is_enabled _ | Iopaque
+  | Ibeginregion | Iendregion
     -> if fp then 10 else 11
 
 let max_register_pressure =
@@ -401,6 +403,7 @@ let max_register_pressure =
              | Ioffset_loc (_, _) | Ifloatarithmem (_, _)
              | Ibswap _ | Ifloatsqrtf _ | Isqrtf)
   | Iname_for_debugger _ | Iprobe _ | Iprobe_is_enabled _ | Iopaque
+  | Ibeginregion | Iendregion
     -> consumes ~int:0 ~float:0
 
 (* Pure operations (without any side effect besides updating their result
@@ -410,6 +413,7 @@ let op_is_pure = function
   | Icall_ind | Icall_imm _ | Itailcall_ind | Itailcall_imm _
   | Iextcall _ | Istackoffset _ | Istore _ | Ialloc _
   | Iintop(Icheckbound) | Iintop_imm(Icheckbound, _) | Iopaque -> false
+  | Ibeginregion | Iendregion -> false
   | Ispecific(Ipause)
   | Ispecific(Iprefetch _) -> false
   | Ispecific(Ilea _ | Isextend32 | Izextend32 | Ifloat_iround | Ifloat_round _
@@ -451,7 +455,7 @@ let init () =
 let operation_supported = function
   | Cpopcnt -> !popcnt_support
   | Cprefetch _
-  | Capply _ | Cextcall _ | Cload _ | Calloc | Cstore _
+  | Capply _ | Cextcall _ | Cload _ | Calloc _ | Cstore _
   | Caddi | Csubi | Cmuli | Cmulhi _ | Cdivi | Cmodi
   | Cand | Cor | Cxor | Clsl | Clsr | Casr
   | Cclz _ | Cctz _

backend/amd64/reload.ml

@@ -158,6 +158,7 @@ method! reload_operation op arg res =
   | Itailcall_ind|Itailcall_imm _|Iextcall _|Istackoffset _|Iload (_, _)
   | Istore (_, _, _)|Ialloc _|Iname_for_debugger _|Iprobe _|Iprobe_is_enabled _
   | Iopaque
+  | Ibeginregion | Iendregion
     -> (* Other operations: all args and results in registers *)
       super#reload_operation op arg res
 

backend/amd64/selection.ml

@@ -155,6 +155,7 @@ let pseudoregs_for_operation op arg res =
   | Iconst_symbol _|Icall_ind|Icall_imm _|Itailcall_ind|Itailcall_imm _
   | Iextcall _|Istackoffset _|Iload (_, _)|Istore (_, _, _)|Ialloc _
   | Iname_for_debugger _|Iprobe _|Iprobe_is_enabled _ | Iopaque
+  | Ibeginregion | Iendregion
     -> raise Use_default
 
 let select_locality (l : Cmm.prefetch_temporal_locality_hint)
@@ -242,6 +243,7 @@ method! select_store is_assign addr exp =
   | Cassign (_, _) | Ctuple _ | Cop (_, _, _) | Csequence (_, _)
   | Cifthenelse (_, _, _, _, _, _) | Cswitch (_, _, _, _) | Ccatch (_, _, _)
   | Cexit (_, _, _) | Ctrywith (_, _, _, _, _)
+  | Cregion _ | Ctail _
     ->
       super#select_store is_assign addr exp
 

backend/arm64/emit.mlp

@@ -511,6 +511,8 @@ module BR = Branch_relaxation.Make (struct
       | 16 | 24 | 32 -> 1
       | _ -> 1 + num_instructions_for_intconst (Nativeint.of_int num_bytes)
       end
+    | Lop (Ibeginregion | Iendregion) ->
+       Misc.fatal_error "Local allocations not supported on this architecture"
     | Lop (Iintop (Icomp _)) -> 2
     | Lop (Icompf _) -> 2
     | Lop (Iintop_imm (Icomp _, _)) -> 2
@@ -787,10 +789,12 @@ let emit_instr i =
         | Word_int | Word_val | Double ->
             `	str	{emit_reg src}, {emit_addressing addr base}\n`
         end
-    | Lop(Ialloc { bytes = n; dbginfo }) ->
+    | Lop(Ialloc { bytes = n; dbginfo; mode = Alloc_heap }) ->
         assembly_code_for_allocation i ~n ~far:false ~dbginfo
     | Lop(Ispecific (Ifar_alloc { bytes = n; dbginfo })) ->
         assembly_code_for_allocation i ~n ~far:true ~dbginfo
+    | Lop(Ialloc { mode = Alloc_local } | Ibeginregion | Iendregion) ->
+       Misc.fatal_error "Local allocations not supported on this architecture"
     | Lop(Iintop_imm(Iadd, n)) ->
         emit_addimm i.res.(0) i.arg.(0) n
     | Lop(Iintop_imm(Isub, n)) ->

backend/arm64/proc.ml

@@ -288,6 +288,7 @@ let op_is_pure = function
   | Icall_ind | Icall_imm _ | Itailcall_ind | Itailcall_imm _
   | Iextcall _ | Istackoffset _ | Istore _ | Ialloc _
   | Iintop(Icheckbound) | Iintop_imm(Icheckbound, _) | Iopaque
+  | Ibeginregion | Iendregion
   | Ispecific(Ishiftcheckbound _) -> false
   | _ -> true
 
@@ -314,7 +315,7 @@ let operation_supported = function
   | Cclz _ | Cctz _ | Cpopcnt
   | Cprefetch _
     -> false   (* Not implemented *)
-  | Capply _ | Cextcall _ | Cload _ | Calloc | Cstore _
+  | Capply _ | Cextcall _ | Cload _ | Calloc _ | Cstore _
   | Caddi | Csubi | Cmuli | Cmulhi _ | Cdivi | Cmodi
   | Cand | Cor | Cxor | Clsl | Clsr | Casr
   | Ccmpi _ | Caddv | Cadda | Ccmpa _

backend/cfg/cfg.ml

@@ -221,6 +221,8 @@ let dump_op ppf = function
     Format.fprintf ppf "probe %s %s" name handler_code_sym
   | Probe_is_enabled { name } -> Format.fprintf ppf "probe_is_enabled %s" name
   | Opaque -> Format.fprintf ppf "opaque"
+  | Begin_region -> Format.fprintf ppf "beginregion"
+  | End_region -> Format.fprintf ppf "endregion"
   | Name_for_debugger _ -> Format.fprintf ppf "name_for_debugger"
 
 let dump_call ppf = function
@@ -319,6 +321,8 @@ let can_raise_operation : operation -> bool = function
   | Specific _ -> false (* CR xclerc for xclerc: double check *)
   | Opaque -> false
   | Name_for_debugger _ -> false
+  | Begin_region -> false
+  | End_region -> false
 
 let can_raise_basic : basic -> bool = function
   | Op op -> can_raise_operation op

backend/cfg/cfg_equivalence.ml

@@ -250,6 +250,8 @@ let check_operation : location -> Cfg.operation -> Cfg.operation -> unit =
     when Arch.equal_specific_operation expected_spec result_spec ->
     ()
   | Opaque, Opaque -> ()
+  | Begin_region, Begin_region -> ()
+  | End_region, End_region -> ()
   | ( Name_for_debugger
         { ident = left_ident;
           which_parameter = left_which_parameter;
@@ -276,9 +278,16 @@ let check_prim_call_operation :
   match expected, result with
   | External expected, External result ->
     check_external_call_operation location expected result
-  | ( Alloc { bytes = expected_bytes; dbginfo = _expected_dbginfo },
-      Alloc { bytes = result_bytes; dbginfo = _result_dbginfo } )
-    when Int.equal expected_bytes result_bytes ->
+  | ( Alloc
+        { bytes = expected_bytes;
+          dbginfo = _expected_dbginfo;
+          mode = expected_mode
+        },
+      Alloc
+        { bytes = result_bytes; dbginfo = _result_dbginfo; mode = result_mode }
+    )
+    when Int.equal expected_bytes result_bytes
+         && Lambda.eq_mode expected_mode result_mode ->
     (* CR xclerc for xclerc: also check debug info *)
     ()
   | ( Checkbound { immediate = expected_immediate },

backend/cfg/cfg_intf.ml

@@ -49,7 +49,8 @@ module S = struct
     | External of external_call_operation
     | Alloc of
         { bytes : int;
-          dbginfo : Debuginfo.alloc_dbginfo
+          dbginfo : Debuginfo.alloc_dbginfo;
+          mode : Lambda.alloc_mode
         }
     | Checkbound of { immediate : int option }
 
@@ -80,6 +81,8 @@ module S = struct
         }
     | Probe_is_enabled of { name : string }
     | Opaque
+    | Begin_region
+    | End_region
     | Specific of Arch.specific_operation
     | Name_for_debugger of
         { ident : Ident.t;

backend/cfg/cfg_to_linear.ml

@@ -51,7 +51,8 @@ let from_basic (basic : Cfg.basic) : L.instruction_desc =
   | Call (P (Checkbound { immediate = None })) -> Lop (Iintop Icheckbound)
   | Call (P (Checkbound { immediate = Some i })) ->
     Lop (Iintop_imm (Icheckbound, i))
-  | Call (P (Alloc { bytes; dbginfo })) -> Lop (Ialloc { bytes; dbginfo })
+  | Call (P (Alloc { bytes; dbginfo; mode })) ->
+    Lop (Ialloc { bytes; dbginfo; mode })
   | Op op ->
     let op : Mach.operation =
       match op with
@@ -79,6 +80,8 @@ let from_basic (basic : Cfg.basic) : L.instruction_desc =
       | Probe_is_enabled { name } -> Iprobe_is_enabled { name }
       | Opaque -> Iopaque
       | Specific op -> Ispecific op
+      | Begin_region -> Ibeginregion
+      | End_region -> Iendregion
       | Name_for_debugger { ident; which_parameter; provenance; is_assignment }
         ->
         Iname_for_debugger { ident; which_parameter; provenance; is_assignment }

backend/cfg/cfgize.ml

@@ -158,7 +158,8 @@ let basic_or_terminator_of_operation :
   | Istackoffset ofs -> Basic (Op (Stackoffset ofs))
   | Iload (mem, mode) -> Basic (Op (Load (mem, mode)))
   | Istore (mem, mode, assignment) -> Basic (Op (Store (mem, mode, assignment)))
-  | Ialloc { bytes; dbginfo } -> Basic (Call (P (Alloc { bytes; dbginfo })))
+  | Ialloc { bytes; dbginfo; mode } ->
+    Basic (Call (P (Alloc { bytes; dbginfo; mode })))
   | Iintop Icheckbound -> Basic (Call (P (Checkbound { immediate = None })))
   | Iintop_imm (Icheckbound, i) ->
     Basic (Call (P (Checkbound { immediate = Some i })))
@@ -189,6 +190,8 @@ let basic_or_terminator_of_operation :
   | Iprobe { name; handler_code_sym } ->
     Basic (Op (Probe { name; handler_code_sym }))
   | Iprobe_is_enabled { name } -> Basic (Op (Probe_is_enabled { name }))
+  | Ibeginregion -> Basic (Op Begin_region)
+  | Iendregion -> Basic (Op End_region)
 
 let float_test_of_float_comparison :
     Cmm.float_comparison ->
@@ -316,7 +319,8 @@ let is_noop_move (instr : Cfg.basic Cfg.instruction) : bool =
       ( Const_int _ | Const_float _ | Const_symbol _ | Stackoffset _ | Load _
       | Store _ | Intop _ | Intop_imm _ | Negf | Absf | Addf | Subf | Mulf
       | Divf | Compf _ | Floatofint | Intoffloat | Probe _ | Opaque
-      | Probe_is_enabled _ | Specific _ | Name_for_debugger _ )
+      | Probe_is_enabled _ | Specific _ | Name_for_debugger _ | Begin_region
+      | End_region )
   | Call _ | Reloadretaddr | Pushtrap _ | Poptrap | Prologue ->
     false
 

backend/cfg/linear_to_cfg.ml

@@ -393,7 +393,7 @@ let to_basic (mop : Mach.operation) : C.basic =
          | Ipopcnt | Iclz _ | Ictz _ | Ilsl | Ilsr | Iasr | Icomp _ ) as op),
         i ) ->
     Op (Intop_imm (op, i))
-  | Ialloc { bytes; dbginfo } -> Call (P (Alloc { bytes; dbginfo }))
+  | Ialloc { bytes; dbginfo; mode } -> Call (P (Alloc { bytes; dbginfo; mode }))
   | Iprobe { name; handler_code_sym } -> Op (Probe { name; handler_code_sym })
   | Iprobe_is_enabled { name } -> Op (Probe_is_enabled { name })
   | Istackoffset i -> Op (Stackoffset i)
@@ -415,6 +415,8 @@ let to_basic (mop : Mach.operation) : C.basic =
   | Ifloatofint -> Op Floatofint
   | Iintoffloat -> Op Intoffloat
   | Iopaque -> Op Opaque
+  | Ibeginregion -> Op Begin_region
+  | Iendregion -> Op End_region
   | Ispecific op -> Op (Specific op)
   | Iname_for_debugger { ident; which_parameter; provenance; is_assignment } ->
     Op (Name_for_debugger { ident; which_parameter; provenance; is_assignment })
@@ -611,8 +613,8 @@ let rec create_blocks (t : t) (i : L.instruction) (block : C.basic_block)
     | Istore (_, _, _)
     | Ialloc _ | Iintop _
     | Iintop_imm (_, _)
-    | Iopaque | Iprobe _ | Iprobe_is_enabled _ | Ispecific _
-    | Iname_for_debugger _ ->
+    | Iopaque | Iprobe _ | Iprobe_is_enabled _ | Ispecific _ | Ibeginregion
+    | Iendregion | Iname_for_debugger _ ->
       let desc = to_basic mop in
       block.body <- create_instruction t desc i ~trap_depth :: block.body;
       if Mach.operation_can_raise mop