Backend changes for local allocations with Flambda 2

backend/amd64/proc.ml

@@ -465,5 +465,5 @@ let operation_supported = function
   | Cfloatofint | Cintoffloat | Ccmpf _
   | Craise _
   | Ccheckbound
-  | Cprobe _ | Cprobe_is_enabled _ | Copaque
+  | Cprobe _ | Cprobe_is_enabled _ | Copaque | Cbeginregion | Cendregion
     -> true

backend/arm64/proc.ml

@@ -325,4 +325,5 @@ let operation_supported = function
   | Craise _
   | Ccheckbound
   | Cprobe _ | Cprobe_is_enabled _ | Copaque
+  | Cbeginregion | Cendregion
     -> true

backend/cmm.ml

@@ -202,6 +202,7 @@ and operation =
   | Cprobe of { name: string; handler_code_sym: string; }
   | Cprobe_is_enabled of { name: string }
   | Copaque
+  | Cbeginregion | Cendregion
 
 type expression =
     Cconst_int of int * Debuginfo.t

backend/cmm.mli

@@ -207,6 +207,7 @@ and operation =
   | Cprobe of { name: string; handler_code_sym: string; }
   | Cprobe_is_enabled of { name: string }
   | Copaque (* Sys.opaque_identity *)
+  | Cbeginregion | Cendregion
 
 (** Every basic block should have a corresponding [Debuginfo.t] for its
     beginning. *)

backend/cmm_helpers.ml

@@ -783,7 +783,6 @@ let unboxed_float_array_ref arr ofs dbg =
 let float_array_ref arr ofs dbg =
   box_float dbg Alloc_heap (unboxed_float_array_ref arr ofs dbg)
 
-(* TODO support mutation of local arrays *)
 let addr_array_set arr ofs newval dbg =
   Cop(Cextcall { func = "caml_modify"; ty = typ_void; alloc = false;
                  builtin = false;
@@ -793,6 +792,15 @@ let addr_array_set arr ofs newval dbg =
                  ty_args = []},
       [array_indexing log2_size_addr arr ofs dbg; newval], dbg)
 
+let addr_array_set_local arr ofs newval dbg =
+  Cop(Cextcall { func = "caml_modify_local"; ty = typ_void; alloc = false;
+                 builtin = false;
+                 returns = true;
+                 effects = Arbitrary_effects;
+                 coeffects = Has_coeffects;
+                 ty_args = []},
+      [arr; untag_int ofs dbg; newval], dbg)
+
 let addr_array_initialize arr ofs newval dbg =
   Cop(Cextcall { func = "caml_initialize";
                  builtin = false;
@@ -903,7 +911,7 @@ let make_alloc_generic ~mode set_fn dbg tag wordsize args =
          fill_fields 1 args)
   end
 
-let make_alloc ?(mode=Lambda.Alloc_heap) dbg tag args =
+let make_alloc ~mode dbg tag args =
   let addr_array_init arr ofs newval dbg =
     Cop(Cextcall { func = "caml_initialize"; ty = typ_void; alloc = false;
                    builtin = false;
@@ -915,7 +923,7 @@ let make_alloc ?(mode=Lambda.Alloc_heap) dbg tag args =
   in
   make_alloc_generic ~mode addr_array_init dbg tag (List.length args) args
 
-let make_float_alloc ?(mode=Lambda.Alloc_heap) dbg tag args =
+let make_float_alloc ~mode dbg tag args =
   make_alloc_generic ~mode float_array_set dbg tag
                      (List.length args * size_float / size_addr) args
 
@@ -2665,9 +2673,7 @@ let setfield_computed ptr init arg1 arg2 arg3 dbg =
   | Caml_modify ->
       return_unit dbg (addr_array_set arg1 arg2 arg3 dbg)
   | Caml_modify_local ->
-      (* TODO: support this, if there are any uses.
-         (Currently, setfield_computed is only used by classes) *)
-      Misc.fatal_error "setfield_computed: local"
+      return_unit dbg (addr_array_set_local arg1 arg2 arg3 dbg)
   | Simple ->
       return_unit dbg (int_array_set arg1 arg2 arg3 dbg)
 

backend/cmm_helpers.mli

@@ -267,6 +267,8 @@ val unboxed_float_array_ref :
 val float_array_ref : expression -> expression -> Debuginfo.t -> expression
 val addr_array_set :
   expression -> expression -> expression -> Debuginfo.t -> expression
+val addr_array_set_local :
+  expression -> expression -> expression -> Debuginfo.t -> expression
 val addr_array_initialize :
   expression -> expression -> expression -> Debuginfo.t -> expression
 val int_array_set :
@@ -309,11 +311,11 @@ val call_cached_method :
 
 (** Allocate a block of regular values with the given tag *)
 val make_alloc :
-  ?mode:Lambda.alloc_mode -> Debuginfo.t -> int -> expression list -> expression
+  mode:Lambda.alloc_mode -> Debuginfo.t -> int -> expression list -> expression
 
 (** Allocate a block of unboxed floats with the given tag *)
 val make_float_alloc :
-  ?mode:Lambda.alloc_mode -> Debuginfo.t -> int -> expression list -> expression
+  mode:Lambda.alloc_mode -> Debuginfo.t -> int -> expression list -> expression
 
 (** Bounds checking *)
 

backend/cmmgen.ml

@@ -437,7 +437,7 @@ let rec transl env e =
         | [] -> Debuginfo.none
         | fundecl::_ -> fundecl.dbg
       in
-      make_alloc dbg Obj.closure_tag (transl_fundecls 0 fundecls)
+      make_alloc ~mode:Alloc_heap dbg Obj.closure_tag (transl_fundecls 0 fundecls)
   | Uoffset(arg, offset) ->
       (* produces a valid Caml value, pointing just after an infix header *)
       let ptr = transl env arg in

backend/printcmm.ml

@@ -207,6 +207,8 @@ let operation d = function
     Printf.sprintf "prefetch is_write=%b prefetch_temporal_locality_hint=%s"
       is_write (temporal_locality locality)
   | Copaque -> "opaque"
+  | Cbeginregion -> "beginregion"
+  | Cendregion -> "endregion"
 
 let rec expr ppf = function
   | Cconst_int (n, _dbg) -> fprintf ppf "%i" n

backend/selectgen.ml

@@ -185,6 +185,11 @@ let oper_result_type = function
   | Cprobe _ -> typ_void
   | Cprobe_is_enabled _ -> typ_int
   | Copaque -> typ_val
+  | Cbeginregion ->
+    (* This must not be typ_val; the begin-region operation returns a
+       naked pointer into the local allocation stack. *)
+    typ_int
+  | Cendregion -> typ_void
 
 (* Infer the size in bytes of the result of an expression whose evaluation
    may be deferred (cf. [emit_parts]). *)
@@ -441,7 +446,7 @@ method is_simple_expr = function
       | Capply _ | Cextcall _ | Calloc _ | Cstore _
       | Craise _ | Ccheckbound
       | Cprobe _ | Cprobe_is_enabled _ | Copaque -> false
-      | Cprefetch _ -> false (* avoid reordering *)
+      | Cprefetch _ | Cbeginregion | Cendregion -> false (* avoid reordering *)
         (* The remaining operations are simple if their args are *)
       | Cload _ | Caddi | Csubi | Cmuli | Cmulhi _ | Cdivi | Cmodi | Cand | Cor
       | Cxor | Clsl | Clsr | Casr | Ccmpi _ | Caddv | Cadda | Ccmpa _ | Cnegf
@@ -488,6 +493,7 @@ method effects_of exp =
       | Calloc Alloc_heap -> EC.none
       | Calloc Alloc_local -> EC.coeffect_only Coeffect.Arbitrary
       | Cstore _ -> EC.effect_only Effect.Arbitrary
+      | Cbeginregion | Cendregion -> EC.arbitrary
       | Cprefetch _ -> EC.arbitrary
       | Craise _ | Ccheckbound -> EC.effect_only Effect.Raise
       | Cload (_, Asttypes.Immutable) -> EC.none
@@ -623,6 +629,8 @@ method select_operation op args _dbg =
   | (Cprobe { name; handler_code_sym; }, _) ->
     Iprobe { name; handler_code_sym; }, args
   | (Cprobe_is_enabled {name}, _) -> Iprobe_is_enabled {name}, []
+  | (Cbeginregion, _) -> Ibeginregion, []
+  | (Cendregion, _) -> Iendregion, args
   | _ -> Misc.fatal_error "Selection.select_oper"
 
 method private select_arith_comm op = function

middle_end/flambda2/to_cmm/to_cmm_helper.ml

@@ -319,19 +319,22 @@ let check_alloc_fields = function
 let make_array ?(dbg = Debuginfo.none) kind args =
   check_alloc_fields args;
   match (kind : Flambda_primitive.Array_kind.t) with
-  | Immediates | Values -> make_alloc dbg 0 args
-  | Naked_floats -> make_float_alloc dbg (Tag.to_int Tag.double_array_tag) args
+  | Immediates | Values -> make_alloc ~mode:Alloc_heap dbg 0 args
+  | Naked_floats ->
+    make_float_alloc ~mode:Alloc_heap dbg (Tag.to_int Tag.double_array_tag) args
 
 let make_block ?(dbg = Debuginfo.none) kind args =
   check_alloc_fields args;
   match (kind : Flambda_primitive.Block_kind.t) with
-  | Values (tag, _) -> make_alloc dbg (Tag.Scannable.to_int tag) args
-  | Naked_floats -> make_float_alloc dbg (Tag.to_int Tag.double_array_tag) args
+  | Values (tag, _) ->
+    make_alloc ~mode:Alloc_heap dbg (Tag.Scannable.to_int tag) args
+  | Naked_floats ->
+    make_float_alloc ~mode:Alloc_heap dbg (Tag.to_int Tag.double_array_tag) args
 
 let make_closure_block ?(dbg = Debuginfo.none) l =
   assert (List.compare_length_with l 0 > 0);
   let tag = Tag.(to_int closure_tag) in
-  make_alloc dbg tag l
+  make_alloc ~mode:Alloc_heap dbg tag l
 
 (* Block access *)