Draft: Use alloc_mode with uniqueness for block allocations in lambda

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -1775,6 +1775,7 @@ and cps_function env ~fid ~(recursive : Recursive.t) ?precomputed_free_idents
     List.concat_map
       (fun (({ name; layout; mode; attributes } : L.lparam), kinds) :
            Function_decl.param list ->
+        let mode = L.todo_mode_propagation mode in
         match kinds with
         | [] -> []
         | [kind] -> [{ name; kind; mode; attributes }]
@@ -1819,6 +1820,7 @@ and cps_function env ~fid ~(recursive : Recursive.t) ?precomputed_free_idents
     in
     cps_tail acc new_env ccenv body body_cont body_exn_cont
   in
+  let ret_mode = L.todo_mode_propagation ret_mode in
   Function_decl.create ~let_rec_ident:(Some fid) ~function_slot ~kind ~params
     ~params_arity ~removed_params ~return ~calling_convention
     ~return_continuation:body_cont ~exn_continuation ~my_region ~my_ghost_region

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -968,6 +968,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
   match prim, args with
   | Pmakeblock (tag, mutability, shape, mode), _ ->
     let args = List.flatten args in
+    let mode = L.todo_mode_propagation mode in
     let mode = Alloc_mode.For_allocations.from_lambda mode ~current_region in
     let tag = Tag.Scannable.create_exn tag in
     let shape = convert_block_shape shape ~num_fields:(List.length args) in
@@ -1005,13 +1006,15 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
     List.map (fun arg : H.expr_primitive -> Simple arg) projected_args
   | Pmakefloatblock (mutability, mode), _ ->
     let args = List.flatten args in
+    let mode = L.todo_mode_propagation mode in
     let mode = Alloc_mode.For_allocations.from_lambda mode ~current_region in
     let mutability = Mutability.from_lambda mutability in
     [ Variadic
         (Make_block (Naked_floats, mutability, mode), List.map unbox_float args)
     ]
   | Pmakeufloatblock (mutability, mode), _ ->
     let args = List.flatten args in
+    let mode = L.todo_mode_propagation mode in
     let mode = Alloc_mode.For_allocations.from_lambda mode ~current_region in
     let mutability = Mutability.from_lambda mutability in
     [Variadic (Make_block (Naked_floats, mutability, mode), args)]
@@ -1027,6 +1030,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
           | Flat_suffix Float_boxed -> unbox_float arg)
         args
     in
+    let mode = L.todo_mode_propagation mode in
     let mode = Alloc_mode.For_allocations.from_lambda mode ~current_region in
     let mutability = Mutability.from_lambda mutability in
     let tag = Tag.Scannable.create_exn tag in

ocaml/lambda/lambda.ml

@@ -49,6 +49,12 @@ include (struct
     | Alloc_heap
     | Alloc_local
 
+  type uniqueness_mode =
+    | Alloc_unique
+    | Alloc_aliased
+
+  type alloc_mode = locality_mode * uniqueness_mode
+
   type modify_mode =
     | Modify_heap
     | Modify_maybe_stack
@@ -59,6 +65,13 @@ include (struct
     if Config.stack_allocation then Alloc_local
     else Alloc_heap
 
+  let alloc_unique = Alloc_unique
+  let alloc_aliased = Alloc_aliased
+  let alloc_heap_unique = alloc_heap, alloc_unique
+  let alloc_heap_aliased = alloc_heap, Alloc_aliased
+  let alloc_local_unique = alloc_local, alloc_unique
+  let alloc_local_aliased = alloc_local, Alloc_aliased
+
   let modify_heap = Modify_heap
 
   let modify_maybe_stack : modify_mode =
@@ -69,26 +82,52 @@ include (struct
     match a, b with
     | Alloc_local, _ | _, Alloc_local -> Alloc_local
     | Alloc_heap, Alloc_heap -> Alloc_heap
+
+  let join_uniqueness_mode a b =
+    match a, b with
+    | Alloc_aliased, _ | _, Alloc_aliased -> Alloc_aliased
+    | Alloc_unique, Alloc_unique -> Alloc_unique
+
+  let join_mode (al, au) (bl, bu) =
+    (join_locality_mode al bl, join_uniqueness_mode au bu)
 end : sig
 
   type locality_mode = private
     | Alloc_heap
     | Alloc_local
 
+  type uniqueness_mode = private
+    | Alloc_unique
+    | Alloc_aliased
+
+  type alloc_mode = locality_mode * uniqueness_mode
+
   type modify_mode = private
     | Modify_heap
     | Modify_maybe_stack
 
   val alloc_heap : locality_mode
   val alloc_local : locality_mode
 
+  val alloc_unique : uniqueness_mode
+  val alloc_aliased : uniqueness_mode
+
+  val alloc_heap_unique : alloc_mode
+  val alloc_heap_aliased : alloc_mode
+  val alloc_local_unique : alloc_mode
+  val alloc_local_aliased : alloc_mode
+
   val modify_heap : modify_mode
 
   val modify_maybe_stack : modify_mode
 
   val join_locality_mode : locality_mode -> locality_mode -> locality_mode
+  val join_mode : alloc_mode -> alloc_mode -> alloc_mode
 end)
 
+let todo_mode_propagation mode =
+  fst mode
+
 let is_local_mode = function
   | Alloc_heap -> false
   | Alloc_local -> true
@@ -103,13 +142,32 @@ let sub_locality_mode a b =
   | _, Alloc_local -> true
   | Alloc_local, Alloc_heap -> false
 
+let sub_uniqueness_mode a b =
+  match a, b with
+  | Alloc_unique, _ -> true
+  | _, Alloc_aliased -> true
+  | Alloc_aliased, Alloc_unique -> false
+
+let sub_mode (al, au) (bl, bu) =
+  sub_locality_mode al bl && sub_uniqueness_mode au bu
+
 let eq_locality_mode a b =
   match a, b with
   | Alloc_heap, Alloc_heap -> true
   | Alloc_local, Alloc_local -> true
   | Alloc_heap, Alloc_local -> false
   | Alloc_local, Alloc_heap -> false
 
+let eq_uniqueness_mode a b =
+  match a, b with
+  | Alloc_unique, Alloc_unique -> true
+  | Alloc_aliased, Alloc_aliased -> true
+  | Alloc_unique, Alloc_aliased -> false
+  | Alloc_aliased, Alloc_unique -> false
+
+let eq_mode (al, au) (bl, bu) =
+  eq_locality_mode al bl && eq_uniqueness_mode au bu
+
 type initialization_or_assignment =
   | Assignment of modify_mode
   | Heap_initialization
@@ -129,10 +187,10 @@ type primitive =
   | Psetglobal of Compilation_unit.t
   | Pgetpredef of Ident.t
   (* Operations on heap blocks *)
-  | Pmakeblock of int * mutable_flag * block_shape * locality_mode
-  | Pmakefloatblock of mutable_flag * locality_mode
-  | Pmakeufloatblock of mutable_flag * locality_mode
-  | Pmakemixedblock of int * mutable_flag * mixed_block_shape * locality_mode
+  | Pmakeblock of int * mutable_flag * block_shape * alloc_mode
+  | Pmakefloatblock of mutable_flag * alloc_mode
+  | Pmakeufloatblock of mutable_flag * alloc_mode
+  | Pmakemixedblock of int * mutable_flag * mixed_block_shape * alloc_mode
   | Pfield of int * immediate_or_pointer * field_read_semantics
   | Pfield_computed of field_read_semantics
   | Psetfield of int * immediate_or_pointer * initialization_or_assignment
@@ -714,7 +772,7 @@ type lparam = {
   name : Ident.t;
   layout : layout;
   attributes : parameter_attribute;
-  mode : locality_mode
+  mode : alloc_mode
 }
 
 type pop_region =
@@ -765,7 +823,7 @@ and lfunction =
     attr: function_attribute; (* specified with [@inline] attribute *)
     loc: scoped_location;
     mode: locality_mode;
-    ret_mode: locality_mode;
+    ret_mode: alloc_mode;
     region: bool; }
 
 and lambda_while =
@@ -1700,10 +1758,10 @@ let primitive_may_allocate : primitive -> locality_mode option = function
   | Parray_to_iarray | Parray_of_iarray
   | Pignore -> None
   | Pgetglobal _ | Psetglobal _ | Pgetpredef _ -> None
-  | Pmakeblock (_, _, _, m) -> Some m
-  | Pmakefloatblock (_, m) -> Some m
-  | Pmakeufloatblock (_, m) -> Some m
-  | Pmakemixedblock (_, _, _, m) -> Some m
+  | Pmakeblock (_, _, _, m) -> Some (fst m)
+  | Pmakefloatblock (_, m) -> Some (fst m)
+  | Pmakeufloatblock (_, m) -> Some (fst m)
+  | Pmakemixedblock (_, _, _, m) -> Some (fst m)
   | Pfield _ | Pfield_computed _ | Psetfield _ | Psetfield_computed _ -> None
   | Pfloatfield (_, _, m) -> Some m
   | Pufloatfield _ -> None
@@ -1785,7 +1843,7 @@ let primitive_may_allocate : primitive -> locality_mode option = function
   | Punboxed_float32_array_load_128 { mode = m; _ }
   | Punboxed_int32_array_load_128 { mode = m; _ }
   | Punboxed_int64_array_load_128 { mode = m; _ }
-  | Punboxed_nativeint_array_load_128 { mode = m; _ }
+  | Punboxed_nativeint_array_load_128 { mode = m; _ } -> Some m
   | Pget_header m -> Some m
   | Pstring_load_32 { boxed = false; _ }
   | Pstring_load_f32 { boxed = false; _ }

ocaml/lambda/lambda.mli

@@ -41,6 +41,12 @@ type locality_mode = private
   | Alloc_heap
   | Alloc_local
 
+type uniqueness_mode = private
+  | Alloc_unique (* Might be reused *)
+  | Alloc_aliased (* No reuse *)
+
+type alloc_mode = locality_mode * uniqueness_mode
+
 type modify_mode = private
   | Modify_heap
   | Modify_maybe_stack
@@ -50,10 +56,20 @@ val alloc_heap : locality_mode
 (* Actually [Alloc_heap] if [Config.stack_allocation] is [false] *)
 val alloc_local : locality_mode
 
+val alloc_aliased : uniqueness_mode
+val alloc_unique : uniqueness_mode
+
+val alloc_heap_unique : alloc_mode
+val alloc_heap_aliased : alloc_mode
+val alloc_local_unique : alloc_mode
+val alloc_local_aliased : alloc_mode
+
 val modify_heap : modify_mode
 
 val modify_maybe_stack : modify_mode
 
+val todo_mode_propagation : alloc_mode -> locality_mode
+
 type initialization_or_assignment =
   (* [Assignment Alloc_local] is a mutation of a block that may be heap or local.
      [Assignment Alloc_heap] is a mutation of a block that's definitely heap. *)
@@ -113,10 +129,10 @@ type primitive =
   | Psetglobal of Compilation_unit.t
   | Pgetpredef of Ident.t
   (* Operations on heap blocks *)
-  | Pmakeblock of int * mutable_flag * block_shape * locality_mode
-  | Pmakefloatblock of mutable_flag * locality_mode
-  | Pmakeufloatblock of mutable_flag * locality_mode
-  | Pmakemixedblock of int * mutable_flag * mixed_block_shape * locality_mode
+  | Pmakeblock of int * mutable_flag * block_shape * alloc_mode
+  | Pmakefloatblock of mutable_flag * alloc_mode
+  | Pmakeufloatblock of mutable_flag * alloc_mode
+  | Pmakemixedblock of int * mutable_flag * mixed_block_shape * alloc_mode
   | Pfield of int * immediate_or_pointer * field_read_semantics
   | Pfield_computed of field_read_semantics
   | Psetfield of int * immediate_or_pointer * initialization_or_assignment
@@ -641,7 +657,7 @@ type lparam = {
   name : Ident.t;
   layout : layout;
   attributes : parameter_attribute;
-  mode : locality_mode
+  mode : alloc_mode
 }
 
 type scoped_location = Debuginfo.Scoped_location.t
@@ -714,7 +730,7 @@ and lfunction = private
     attr: function_attribute; (* specified with [@inline] attribute *)
     loc : scoped_location;
     mode : locality_mode;     (* locality of the closure itself *)
-    ret_mode: locality_mode;
+    ret_mode: alloc_mode;
     region : bool;         (* false if this function may locally
                               allocate in the caller's region *)
   }
@@ -848,7 +864,7 @@ val lfunction :
   attr:function_attribute -> (* specified with [@inline] attribute *)
   loc:scoped_location ->
   mode:locality_mode ->
-  ret_mode:locality_mode ->
+  ret_mode:alloc_mode ->
   region:bool ->
   lambda
 
@@ -860,7 +876,7 @@ val lfunction' :
   attr:function_attribute -> (* specified with [@inline] attribute *)
   loc:scoped_location ->
   mode:locality_mode ->
-  ret_mode:locality_mode ->
+  ret_mode:alloc_mode ->
   region:bool ->
   lfunction
 
@@ -967,6 +983,10 @@ val max_arity : unit -> int
       This is unlimited ([max_int]) for bytecode, but limited
       (currently to 126) for native code. *)
 
+val join_mode : alloc_mode -> alloc_mode -> alloc_mode
+val sub_mode : alloc_mode -> alloc_mode -> bool
+val eq_mode : alloc_mode -> alloc_mode -> bool
+
 val join_locality_mode : locality_mode -> locality_mode -> locality_mode
 val sub_locality_mode : locality_mode -> locality_mode -> bool
 val eq_locality_mode : locality_mode -> locality_mode -> bool

ocaml/lambda/matching.ml

@@ -3907,7 +3907,7 @@ let failure_handler ~scopes loc ~failer () =
     Lprim
       ( Praise Raise_regular,
         [ Lprim
-            ( Pmakeblock (0, Immutable, None, alloc_heap),
+            ( Pmakeblock (0, Immutable, None, alloc_heap_aliased),
               [ slot;
                 Lconst
                   (Const_block

ocaml/lambda/matching.mli

@@ -36,7 +36,7 @@ val for_let:
         lambda
 val for_multiple_match:
         scopes:scopes -> return_layout:layout -> Location.t ->
-        (lambda * Jkind.sort * layout) list -> locality_mode ->
+        (lambda * Jkind.sort * layout) list -> alloc_mode ->
         (pattern * lambda) list -> partial ->
         lambda