Finish flambda2 wiring + add caml_modify_local_prototype in memory.h

Author: mshinwell

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -993,8 +993,8 @@ let primitive_result_kind (prim : Lambda.primitive) :
   | Pccall { prim_native_repr_res = _, Unboxed_float; _ }
   | Pfloatofint _ | Pnegfloat _ | Pabsfloat _ | Paddfloat _ | Psubfloat _
   | Pmulfloat _ | Pdivfloat _ | Pfloatfield _
-  | Parrayrefs Pfloatarray
-  | Parrayrefu Pfloatarray
+  | Parrayrefs (Pfloatarray_ref _)
+  | Parrayrefu (Pfloatarray_ref _)
   | Pbigarrayref (_, _, (Pbigarray_float32 | Pbigarray_float64), _) ->
     Flambda_kind.With_subkind.boxed_float
   | Pccall { prim_native_repr_res = _, Unboxed_integer Pnativeint; _ }
@@ -1019,8 +1019,8 @@ let primitive_result_kind (prim : Lambda.primitive) :
   | Pstringlength | Pstringrefu | Pbyteslength | Pbytesrefu | Pbytessetu
   | Parraylength _ | Parraysetu _ | Pisint _ | Pbintcomp _ | Pintofbint _
   | Pisout
-  | Parrayrefs Pintarray
-  | Parrayrefu Pintarray
+  | Parrayrefs Pintarray_ref
+  | Parrayrefu Pintarray_ref
   | Pprobe_is_enabled _ | Pctconst _ | Pbswap16
   | Pbigarrayref
       ( _,
@@ -1055,8 +1055,8 @@ let primitive_result_kind (prim : Lambda.primitive) :
     (* CR ncourant: this should be bottom, but we don't have it *)
     Flambda_kind.With_subkind.any_value
   | Pccall { prim_native_repr_res = _, Same_as_ocaml_repr; _ }
-  | Parrayrefs (Pgenarray | Paddrarray)
-  | Parrayrefu (Pgenarray | Paddrarray)
+  | Parrayrefs (Pgenarray_ref _ | Paddrarray_ref)
+  | Parrayrefu (Pgenarray_ref _ | Paddrarray_ref)
   | Pbytes_to_string | Pbytes_of_string | Parray_of_iarray | Parray_to_iarray
   | Pgetglobal _ | Psetglobal _ | Pgetpredef _ | Pmakeblock _
   | Pmakefloatblock _ | Pfield _ | Pfield_computed _ | Pduprecord _

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -131,34 +131,40 @@ let convert_array_kind (kind : L.array_kind) : converted_array_kind =
   | Pintarray -> Array_kind Immediates
   | Pfloatarray -> Array_kind Naked_floats
 
+module Array_ref_kind = struct
+  type t =
+    | Immediates
+    | Values
+    | Naked_floats of L.alloc_mode
+end
+
 type converted_array_ref_kind =
-  | Array_ref_kind of P.Array_ref_kind.t
-  | Float_array_opt_dynamic_ref of Alloc_mode.For_allocations.t
+  | Array_ref_kind of Array_ref_kind.t
+  | Float_array_opt_dynamic_ref of L.alloc_mode
 
-let convert_array_ref_kind
-      ~current_region (kind : L.array_ref_kind) : converted_array_ref_kind =
+let convert_array_ref_kind (kind : L.array_ref_kind) : converted_array_ref_kind
+    =
   match kind with
   | Pgenarray_ref mode ->
     check_float_array_optimisation_enabled ();
-    Float_array_opt_dynamic_ref
-      (Alloc_mode.For_allocations.from_lambda ~current_region mode)
+    Float_array_opt_dynamic_ref mode
   | Paddrarray_ref -> Array_ref_kind Values
   | Pintarray_ref -> Array_ref_kind Immediates
-  | Pfloatarray_ref mode ->
-    Array_ref_kind
-      (Naked_floats (Alloc_mode.For_allocations.from_lambda ~current_region mode))
+  | Pfloatarray_ref mode -> Array_ref_kind (Naked_floats mode)
 
 type converted_array_set_kind =
   | Array_set_kind of P.Array_set_kind.t
   | Float_array_opt_dynamic_set of Alloc_mode.For_assignments.t
 
-let convert_array_set_kind (kind : L.array_set_kind) : converted_array_set_kind =
+let convert_array_set_kind (kind : L.array_set_kind) : converted_array_set_kind
+    =
   match kind with
   | Pgenarray_set mode ->
     check_float_array_optimisation_enabled ();
     Float_array_opt_dynamic_set (Alloc_mode.For_assignments.from_lambda mode)
   | Paddrarray_set mode ->
-    Array_set_kind (Values (Alloc_mode.For_assignments.from_lambda mode))
+    Array_set_kind
+      (Values (Assignment (Alloc_mode.For_assignments.from_lambda mode)))
   | Pintarray_set -> Array_set_kind Immediates
   | Pfloatarray_set -> Array_set_kind Naked_floats
 
@@ -522,53 +528,39 @@ let check_array_access ~dbg ~array ~index primitive : H.expr_primitive =
     ~conditions:(array_access_validity_condition array index)
     ~dbg
 
-let array_load_unsafe ~array ~index (array_ref_kind : P.Array_ref_kind.t)
-      ~current_region : H.expr_primitive =
+let array_load_unsafe ~array ~index (array_ref_kind : Array_ref_kind.t)
+    ~current_region : H.expr_primitive =
   match array_ref_kind with
-  | Immediates | Values ->
-    Binary (Array_load (array_ref_kind, Mutable), array, index)
+  | Immediates -> Binary (Array_load (Immediates, Mutable), array, index)
+  | Values -> Binary (Array_load (Values, Mutable), array, index)
   | Naked_floats mode ->
     box_float mode
       (Binary (Array_load (Naked_floats, Mutable), array, index))
       ~current_region
 
-let array_set_unsafe ~array ~index ~new_value (array_set_kind : P.Array_set_kind.t) :
-    H.expr_primitive =
-  match array_set_kind with
-  | Immediates
-    Ternary
-      ( Array_set (array_kind, Assignment Alloc_mode.For_assignments.heap),
-        array,
-        index,
-        new_value )
-  | Values ->
-    Ternary
-      ( Array_set (array_kind, Assignment Alloc_mode.For_assignments.heap),
-        array,
-        index,
-        new_value )
-  | Naked_floats ->
-    Ternary
-      ( Array_set
-          (Naked_floats, Assignment (Alloc_mode.For_assignments.local ())),
-        array,
-        index,
-        unbox_float new_value )
+let array_set_unsafe ~array ~index ~new_value
+    (array_set_kind : P.Array_set_kind.t) : H.expr_primitive =
+  let new_value =
+    match array_set_kind with
+    | Immediates | Values _ -> new_value
+    | Naked_floats -> unbox_float new_value
+  in
+  Ternary (Array_set array_set_kind, array, index, new_value)
 
-let[@inline always] match_on_array_ref_kind ~current_region ~array array_ref_kind f
-  : H.expr_primitive =
-  match convert_array_ref_kind ~current_region array_ref_kind with
+let[@inline always] match_on_array_ref_kind ~array array_ref_kind f :
+    H.expr_primitive =
+  match convert_array_ref_kind array_ref_kind with
   | Array_ref_kind array_ref_kind -> f array_ref_kind
   | Float_array_opt_dynamic_ref mode ->
     (* CR keryan: we should push the ITE as low as possible to avoid duplicating
        too much *)
     If_then_else
       ( Unary (Is_flat_float_array, array),
-        f (P.Array_ref_kind.Naked_floats mode),
-        f P.Array_ref_kind.Values )
+        f (Array_ref_kind.Naked_floats mode),
+        f Array_ref_kind.Values )
 
-let[@inline always] match_on_array_set_kind ~array array_ref_kind f
-  : H.expr_primitive =
+let[@inline always] match_on_array_set_kind ~array array_ref_kind f :
+    H.expr_primitive =
   match convert_array_set_kind array_ref_kind with
   | Array_set_kind array_set_kind -> f array_set_kind
   | Float_array_opt_dynamic_set mode ->
@@ -577,7 +569,7 @@ let[@inline always] match_on_array_set_kind ~array array_ref_kind f
     If_then_else
       ( Unary (Is_flat_float_array, array),
         f P.Array_set_kind.Naked_floats,
-        f (P.Array_set_kind.Values mode) )
+        f (P.Array_set_kind.Values (Assignment mode)) )
 
 (* Safe arith (div/mod by zero) *)
 let checked_arith_op ~dbg (bi : Lambda.boxed_integer option) op mode arg1 arg2
@@ -1022,11 +1014,11 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list)
     (* For this and the following cases we will end up relying on the backend to
        CSE the two accesses to the array's header word in the [Pgenarray]
        case. *)
-    match_on_array_ref_kind ~current_region ~array array_ref_kind
+    match_on_array_ref_kind ~array array_ref_kind
       (array_load_unsafe ~array ~index ~current_region)
   | Parrayrefs array_ref_kind, [array; index] ->
     check_array_access ~dbg ~array ~index
-      (match_on_array_ref_kind ~current_region ~array array_ref_kind
+      (match_on_array_ref_kind ~array array_ref_kind
          (array_load_unsafe ~array ~index ~current_region))
   | Parraysetu array_set_kind, [array; index; new_value] ->
     match_on_array_set_kind ~array array_set_kind

middle_end/flambda2/parser/fexpr_to_flambda.ml

@@ -441,7 +441,14 @@ let binop (binop : Fexpr.binop) : Flambda_primitive.binary_primitive =
 
 let ternop env (ternop : Fexpr.ternop) : Flambda_primitive.ternary_primitive =
   match ternop with
-  | Array_set (ak, ia) -> Array_set (ak, init_or_assign env ia)
+  | Array_set (ak, ia) ->
+    let ask : Flambda_primitive.Array_set_kind.t =
+      match ak, ia with
+      | Immediates, _ -> Immediates
+      | Naked_floats, _ -> Naked_floats
+      | Values, ia -> Values (init_or_assign env ia)
+    in
+    Array_set ask
   | Block_set (bk, ia) -> Block_set (block_access_kind bk, init_or_assign env ia)
   | Bytes_or_bigstring_set (blv, saw) -> Bytes_or_bigstring_set (blv, saw)
 

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -582,7 +582,19 @@ let binop (op : Flambda_primitive.binary_primitive) : Fexpr.binop =
 
 let ternop env (op : Flambda_primitive.ternary_primitive) : Fexpr.ternop =
   match op with
-  | Array_set (ak, ia) -> Array_set (ak, init_or_assign env ia)
+  | Array_set ak ->
+    let ia : Flambda_primitive.Init_or_assign.t =
+      match ak with
+      | Values ia -> ia
+      | Immediates | Naked_floats -> Assignment Alloc_mode.For_assignments.heap
+    in
+    let ak : Flambda_primitive.Array_kind.t =
+      match ak with
+      | Immediates -> Immediates
+      | Values _ -> Values
+      | Naked_floats -> Naked_floats
+    in
+    Array_set (ak, init_or_assign env ia)
   | Block_set (bk, ia) -> Block_set (block_access_kind bk, init_or_assign env ia)
   | Bytes_or_bigstring_set (blv, saw) -> Bytes_or_bigstring_set (blv, saw)
   | Bigarray_set _ ->

middle_end/flambda2/simplify/simplify_ternary_primitive.ml

@@ -16,12 +16,16 @@
 
 open! Simplify_import
 
-let simplify_array_set (array_kind : P.Array_kind.t) init_or_assign dacc
-    ~original_term:_ dbg ~arg1:array ~arg1_ty:array_ty ~arg2:index ~arg2_ty:_
-    ~arg3:new_value ~arg3_ty:_ ~result_var =
-  let elt_kind = P.Array_kind.element_kind array_kind |> K.With_subkind.kind in
+let simplify_array_set (array_set_kind : P.Array_set_kind.t) dacc ~original_term
+    dbg ~arg1:array ~arg1_ty:array_ty ~arg2:index ~arg2_ty:_ ~arg3:new_value
+    ~arg3_ty:_ ~result_var =
+  let elt_kind =
+    P.Array_set_kind.element_kind array_set_kind |> K.With_subkind.kind
+  in
   let array_kind =
-    Simplify_common.specialise_array_kind dacc array_kind ~array_ty
+    Simplify_common.specialise_array_kind dacc
+      (P.Array_set_kind.array_kind array_set_kind)
+      ~array_ty
   in
   match array_kind with
   | Bottom -> SPR.create_invalid dacc
@@ -30,10 +34,26 @@ let simplify_array_set (array_kind : P.Array_kind.t) init_or_assign dacc
       P.Array_kind.element_kind array_kind |> K.With_subkind.kind
     in
     assert (K.equal elt_kind elt_kind');
+    let array_set_kind : P.Array_set_kind.t =
+      match array_kind with
+      | Immediates -> Immediates
+      | Values -> (
+        match array_set_kind with
+        | Values init_or_assign -> Values init_or_assign
+        | Immediates
+        (* We don't expect specialisation regressions from Immediates to
+           Values. *)
+        | Naked_floats ->
+          Misc.fatal_errorf
+            "Didn't expect array specialisation to yield array kind %a from \
+             array set kind %a:@ %a"
+            P.Array_kind.print array_kind P.Array_set_kind.print array_set_kind
+            Named.print original_term)
+      | Naked_floats -> Naked_floats
+    in
     let named =
       Named.create_prim
-        (Ternary
-           (Array_set (array_kind, init_or_assign), array, index, new_value))
+        (Ternary (Array_set array_set_kind, array, index, new_value))
         dbg
     in
     let unit_ty = Flambda2_types.this_tagged_immediate Targetint_31_63.zero in
@@ -59,8 +79,7 @@ let simplify_ternary_primitive dacc original_prim (prim : P.ternary_primitive)
   let original_term = Named.create_prim original_prim dbg in
   let simplifier =
     match prim with
-    | Array_set (array_kind, init_or_assign) ->
-      simplify_array_set array_kind init_or_assign
+    | Array_set array_kind -> simplify_array_set array_kind
     | Block_set (block_access_kind, init_or_assign) ->
       simplify_block_set block_access_kind init_or_assign
     | Bytes_or_bigstring_set (bytes_like_value, string_accessor_width) ->

middle_end/flambda2/terms/code_size.ml

@@ -148,12 +148,11 @@ let block_set (kind : Flambda_primitive.Block_access_kind.t)
   | Values _, (Assignment Local | Initialization) -> 1 (* cadda + store *)
   | Naked_floats _, (Assignment _ | Initialization) -> 1
 
-let array_set (kind : Flambda_primitive.Array_kind.t)
-    (init : Flambda_primitive.Init_or_assign.t) =
-  match kind, init with
-  | Values, Assignment Heap -> nonalloc_extcall_size
-  | Values, (Assignment Local | Initialization) -> 1
-  | (Immediates | Naked_floats), (Assignment _ | Initialization) -> 1
+let array_set (kind : Flambda_primitive.Array_set_kind.t) =
+  match kind with
+  | Values (Assignment Heap) -> nonalloc_extcall_size
+  | Values (Assignment Local | Initialization) -> 1
+  | Immediates | Naked_floats -> 1
 
 let string_or_bigstring_load kind width =
   let start_address_load =
@@ -360,7 +359,7 @@ let binary_prim_size prim =
 let ternary_prim_size prim =
   match (prim : Flambda_primitive.ternary_primitive) with
   | Block_set (block_access, init) -> block_set block_access init
-  | Array_set (kind, init) -> array_set kind init
+  | Array_set kind -> array_set kind
   | Bytes_or_bigstring_set (kind, width) -> bytes_like_set kind width
   | Bigarray_set (_dims, (Complex32 | Complex64), _layout) ->
     5 (* ~ 3 block_load + 2 block_set *)