Implement %array_element_size_in_bytes (#3367)

Co-authored-by: Chris Casinghino <ccasinghino@janestreet.com>

bytecomp/bytegen.ml

@@ -158,6 +158,7 @@ let preserve_tailcall_for_prim = function
   | Psetfield_computed _ | Pfloatfield _ | Psetfloatfield _ | Pduprecord _
   | Pufloatfield _ | Psetufloatfield _ | Pmixedfield _ | Psetmixedfield _
   | Pmake_unboxed_product _ | Punboxed_product_field _
+  | Parray_element_size_in_bytes _
   | Pccall _ | Praise _ | Pnot | Pnegint | Paddint | Psubint | Pmulint
   | Pdivint _ | Pmodint _ | Pandint | Porint | Pxorint | Plslint | Plsrint
   | Pasrint | Pintcomp _ | Poffsetint _ | Poffsetref _ | Pintoffloat _
@@ -415,6 +416,8 @@ let comp_primitive stack_info p sz args =
   | Pcompare_bints bi -> comp_bint_primitive bi "compare" args
   | Pfield (n, _ptr, _sem) -> Kgetfield n
   | Punboxed_product_field (n, _layouts) -> Kgetfield n
+  | Parray_element_size_in_bytes _array_kind ->
+      Kconst (Const_base (Const_int (Sys.word_size / 8)))
   | Pfield_computed _sem -> Kgetvectitem
   | Psetfield(n, _ptr, _init) -> Ksetfield n
   | Psetfield_computed(_ptr, _init) -> Ksetvectitem

lambda/lambda.ml

@@ -153,6 +153,7 @@ type primitive =
   (* Unboxed products *)
   | Pmake_unboxed_product of layout list
   | Punboxed_product_field of int * layout list
+  | Parray_element_size_in_bytes of array_kind
   (* Context switches *)
   | Prunstack
   | Pperform
@@ -1939,7 +1940,8 @@ let primitive_may_allocate : primitive -> locality_mode option = function
   | Patomic_cas
   | Patomic_fetch_add
   | Pdls_get
-  | Preinterpret_unboxed_int64_as_tagged_int63 -> None
+  | Preinterpret_unboxed_int64_as_tagged_int63
+  | Parray_element_size_in_bytes _ -> None
   | Preinterpret_tagged_int63_as_unboxed_int64 ->
     if !Clflags.native_code then None
     else
@@ -2104,7 +2106,8 @@ let primitive_can_raise prim =
   | Patomic_cas | Patomic_fetch_add | Patomic_load _ -> false
   | Prunstack | Pperform | Presume | Preperform -> true (* XXX! *)
   | Pdls_get | Ppoll | Preinterpret_tagged_int63_as_unboxed_int64
-  | Preinterpret_unboxed_int64_as_tagged_int63 ->
+  | Preinterpret_unboxed_int64_as_tagged_int63
+  | Parray_element_size_in_bytes _ ->
     false
 
 let constant_layout: constant -> layout = function
@@ -2216,6 +2219,7 @@ let primitive_result_layout (p : primitive) =
   | Pfield _ | Pfield_computed _ -> layout_value_field
   | Punboxed_product_field (field, layouts) -> (Array.of_list layouts).(field)
   | Pmake_unboxed_product layouts -> layout_unboxed_product layouts
+  | Parray_element_size_in_bytes _ -> layout_int
   | Pfloatfield _ -> layout_boxed_float Boxed_float64
   | Pfloatoffloat32 _ -> layout_boxed_float Boxed_float64
   | Pfloat32offloat _ -> layout_boxed_float Boxed_float32
@@ -2506,6 +2510,11 @@ let rec try_to_find_location lam =
 let try_to_find_debuginfo lam =
   Debuginfo.from_location (try_to_find_location lam)
 
+(* The "count_initializers_*" functions count the number of individual
+   components in an initializer for the corresponding array kind _after_
+   unarization.  These are used to implement the "%array_element_size_in_bytes"
+   primitives for products, as each such component takes a full word in product
+   arrays. *)
 let rec count_initializers_scannable
       (scannable : scannable_product_element_kind) =
   match scannable with

lambda/lambda.mli

@@ -141,6 +141,7 @@ type primitive =
   | Pmake_unboxed_product of layout list
   | Punboxed_product_field of int * (layout list)
       (* the [layout list] is the layout of the whole product *)
+  | Parray_element_size_in_bytes of array_kind
   (* Context switches *)
   | Prunstack
   | Pperform

lambda/printlambda.ml

@@ -609,6 +609,8 @@ let primitive ppf = function
       fprintf ppf "unboxed_product_field %d #(%a)" n
         (pp_print_list ~pp_sep:(fun ppf () -> fprintf ppf ", ") (layout' false))
         layouts
+  | Parray_element_size_in_bytes ak ->
+      fprintf ppf "array_element_size_in_bytes (%s)" (array_kind ak)
   | Pccall p -> fprintf ppf "%s" p.prim_name
   | Praise k -> fprintf ppf "%s" (Lambda.raise_kind k)
   | Psequand -> fprintf ppf "&&"
@@ -953,6 +955,7 @@ let name_of_primitive = function
   | Pduprecord _ -> "Pduprecord"
   | Pmake_unboxed_product _ -> "Pmake_unboxed_product"
   | Punboxed_product_field _ -> "Punboxed_product_field"
+  | Parray_element_size_in_bytes _ -> "Parray_element_size_in_bytes"
   | Pccall _ -> "Pccall"
   | Praise _ -> "Praise"
   | Psequand -> "Psequand"

lambda/tmc.ml

@@ -923,6 +923,7 @@ let rec choice ctx t =
 
     (* nor unboxed products *)
     | Pmake_unboxed_product _ | Punboxed_product_field _
+    | Parray_element_size_in_bytes _
 
     | Pobj_dup
     | Pobj_magic _

lambda/translprim.ml

@@ -561,6 +561,9 @@ let lookup_primitive loc ~poly_mode ~poly_sort pos p =
         src_mutability = Immutable;
         dst_array_set_kind = gen_array_set_kind (get_third_arg_mode ())
       }, 5);
+    | "%array_element_size_in_bytes" ->
+      (* The array kind will be filled in later *)
+      Primitive (Parray_element_size_in_bytes Pgenarray, 1)
     | "%obj_size" -> Primitive ((Parraylength Pgenarray), 1)
     | "%obj_field" -> Primitive ((Parrayrefu (Pgenarray_ref mode, Ptagged_int_index, Mutable)), 2)
     | "%obj_set_field" ->
@@ -1302,6 +1305,12 @@ let specialize_primitive env loc ty ~has_constant_constructor prim =
     if dst_array_set_kind = new_dst_array_set_kind then None
     else Some (Primitive (Parrayblit {
       src_mutability; dst_array_set_kind = new_dst_array_set_kind }, arity))
+  | Primitive (Parray_element_size_in_bytes _, arity), p1 :: _ -> (
+      let array_kind =
+        array_type_kind ~elt_sort:None env (to_location loc) p1
+      in
+      Some (Primitive (Parray_element_size_in_bytes array_kind, arity))
+    )
   | Primitive (Pbigarrayref(unsafe, n, kind, layout), arity), p1 :: _ -> begin
       let (k, l) = bigarray_specialize_kind_and_layout env ~kind ~layout p1 in
       match k, l with
@@ -1815,6 +1824,7 @@ let lambda_primitive_needs_event_after = function
   | Pgetglobal _ | Pgetpredef _ | Pmakeblock _ | Pmakefloatblock _
   | Pmakeufloatblock _ | Pmakemixedblock _
   | Pmake_unboxed_product _ | Punboxed_product_field _
+  | Parray_element_size_in_bytes _
   | Pfield _ | Pfield_computed _ | Psetfield _
   | Psetfield_computed _ | Pfloatfield _ | Psetfloatfield _ | Praise _
   | Pufloatfield _ | Psetufloatfield _ | Pmixedfield _ | Psetmixedfield _

lambda/value_rec_compiler.ml

@@ -368,7 +368,8 @@ let compute_static_size lam =
     | Punboxed_float32_array_set_128 _
     | Punboxed_int32_array_set_128 _
     | Punboxed_int64_array_set_128 _
-    | Punboxed_nativeint_array_set_128 _ ->
+    | Punboxed_nativeint_array_set_128 _
+    | Parray_element_size_in_bytes _ ->
         Constant
 
     | Pmakeufloatblock (_, _)

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -1044,9 +1044,10 @@ let close_primitive acc env ~let_bound_ids_with_kinds named
       | Punbox_vector _
       | Pbox_vector (_, _)
       | Punbox_int _ | Pbox_int _ | Pmake_unboxed_product _
-      | Punboxed_product_field _ | Pget_header _ | Prunstack | Pperform
-      | Presume | Preperform | Patomic_exchange | Patomic_compare_exchange
-      | Patomic_cas | Patomic_fetch_add | Pdls_get | Ppoll | Patomic_load _
+      | Punboxed_product_field _ | Parray_element_size_in_bytes _
+      | Pget_header _ | Prunstack | Pperform | Presume | Preperform
+      | Patomic_exchange | Patomic_compare_exchange | Patomic_cas
+      | Patomic_fetch_add | Pdls_get | Ppoll | Patomic_load _
       | Preinterpret_tagged_int63_as_unboxed_int64
       | Preinterpret_unboxed_int64_as_tagged_int63 ->
         (* Inconsistent with outer match *)

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -1363,6 +1363,8 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
     let mutability = Mutability.from_lambda mutability in
     [Variadic (Make_block (Values (tag, shape), mutability, mode), args)]
   | Pmake_unboxed_product layouts, _ ->
+    (* CR mshinwell: this should check the unarized lengths of [layouts] and
+       [args] (like [Parray_element_size_in_bytes] below) *)
     if List.compare_lengths layouts args <> 0
     then
       Misc.fatal_errorf "Pmake_unboxed_product: expected %d arguments, got %d"
@@ -1392,6 +1394,26 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
       |> Array.to_list
     in
     List.map (fun arg : H.expr_primitive -> Simple arg) projected_args
+  | Parray_element_size_in_bytes array_kind, [_witness] ->
+    (* This is implemented as a unary primitive, but from our point of view it's
+       actually nullary. *)
+    let num_bytes =
+      match array_kind with
+      | Pgenarray | Paddrarray | Pintarray | Pfloatarray -> 8
+      | Punboxedfloatarray Unboxed_float32 ->
+        (* float32# arrays are packed *)
+        4
+      | Punboxedfloatarray Unboxed_float64 -> 8
+      | Punboxedintarray Unboxed_int32 ->
+        (* int32# arrays are packed *)
+        4
+      | Punboxedintarray (Unboxed_int64 | Unboxed_nativeint) -> 8
+      | Punboxedvectorarray Unboxed_vec128 -> 16
+      | Pgcscannableproductarray _ | Pgcignorableproductarray _ ->
+        (* All elements of unboxed product arrays are currently 8 bytes wide. *)
+        L.count_initializers_array_kind array_kind * 8
+    in
+    [Simple (Simple.const_int (Targetint_31_63.of_int num_bytes))]
   | Pmakefloatblock (mutability, mode), _ ->
     let args = List.flatten args in
     let mode = Alloc_mode.For_allocations.from_lambda mode ~current_region in
@@ -2405,7 +2427,8 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
       | Punbox_int _ | Pbox_int _ | Punboxed_product_field _ | Pget_header _
       | Pufloatfield _ | Patomic_load _ | Pmixedfield _
       | Preinterpret_unboxed_int64_as_tagged_int63
-      | Preinterpret_tagged_int63_as_unboxed_int64 ),
+      | Preinterpret_tagged_int63_as_unboxed_int64
+      | Parray_element_size_in_bytes _ ),
       ([] | _ :: _ :: _ | [([] | _ :: _ :: _)]) ) ->
     Misc.fatal_errorf
       "Closure_conversion.convert_primitive: Wrong arity for unary primitive \