flambda-backend: Middle and backend support for arrays of unboxed numbers (rebased) (#2185)

Co-authored-by: Mark Shinwell <[email protected]>
Co-authored-by: alanechang <[email protected]>

Author: 3 people

asmcomp/cmm_helpers.ml

@@ -2509,6 +2509,8 @@ let arraylength kind arg dbg =
       Cop(Cor, [addr_array_length_shifted hdr dbg; Cconst_int (1, dbg)], dbg)
   | Pfloatarray ->
       Cop(Cor, [float_array_length_shifted hdr dbg; Cconst_int (1, dbg)], dbg)
+  | Punboxedfloatarray | Punboxedintarray _ ->
+      Misc.fatal_errorf "Unboxed arrays not supported"
 
 let bbswap bi arg dbg =
   let prim, tyarg = match (bi : Primitive.boxed_integer) with
@@ -2699,6 +2701,8 @@ let arrayref_unsafe rkind arg1 arg2 dbg =
       int_array_ref arg1 arg2 dbg
   | Pfloatarray_ref mode ->
       float_array_ref mode arg1 arg2 dbg
+  | Punboxedfloatarray_ref | Punboxedintarray_ref _ ->
+      Misc.fatal_errorf "Unboxed arrays not supported"
 
 let arrayref_safe rkind arg1 arg2 dbg =
   match (rkind : Lambda.array_ref_kind) with
@@ -2752,6 +2756,8 @@ let arrayref_safe rkind arg1 arg2 dbg =
                 (get_header_masked arr dbg) dbg;
               idx],
             unboxed_float_array_ref arr idx dbg))))
+  | Punboxedfloatarray_ref | Punboxedintarray_ref _ ->
+      Misc.fatal_errorf "Unboxed arrays not supported"
 
 type ternary_primitive =
   expression -> expression -> expression -> Debuginfo.t -> expression
@@ -2802,6 +2808,8 @@ let arrayset_unsafe skind arg1 arg2 arg3 dbg =
       int_array_set arg1 arg2 arg3 dbg
   | Pfloatarray_set ->
       float_array_set arg1 arg2 arg3 dbg
+  | Punboxedfloatarray_set | Punboxedintarray_set _ ->
+      Misc.fatal_errorf "Unboxed arrays not supported"
   )
 
 let arrayset_safe skind arg1 arg2 arg3 dbg =
@@ -2865,6 +2873,8 @@ let arrayset_safe skind arg1 arg2 arg3 dbg =
               (get_header_masked arr dbg) dbg;
             idx],
           float_array_set arr idx newval dbg))))
+  | Punboxedfloatarray_set | Punboxedintarray_set _ ->
+      Misc.fatal_errorf "Unboxed arrays not supported"
   )
 
 let bytes_set size unsafe arg1 arg2 arg3 dbg =

asmcomp/cmmgen.ml

@@ -873,6 +873,8 @@ and transl_make_array dbg env kind mode args =
   | Pfloatarray ->
       make_float_alloc ~mode dbg Obj.double_array_tag
                       (List.map (transl_unbox_float dbg env) args)
+  | Punboxedfloatarray | Punboxedintarray _ ->
+      Misc.fatal_errorf "Unboxed arrays not supported"
 
 and transl_ccall env prim args dbg =
   let transl_arg native_repr arg =

bytecomp/bytegen.ml

@@ -492,16 +492,28 @@ let comp_primitive stack_info p sz args =
   | Parrayrefs (Pfloatarray_ref _) -> Kccall("caml_floatarray_get", 2)
   | Parrayrefs (Paddrarray_ref | Pintarray_ref) ->
       Kccall("caml_array_get_addr", 2)
+  | Parrayrefs (Punboxedfloatarray_ref | Punboxedintarray_ref _) ->
+      Misc.fatal_errorf "Cannot use primitive %a for unboxed arrays in bytecode"
+        Printlambda.primitive p
   | Parraysets (Pgenarray_set _) -> Kccall("caml_array_set", 3)
   | Parraysets Pfloatarray_set -> Kccall("caml_floatarray_set", 3)
   | Parraysets (Paddrarray_set _ | Pintarray_set) ->
       Kccall("caml_array_set_addr", 3)
+  | Parraysets (Punboxedfloatarray_set | Punboxedintarray_set _) ->
+      Misc.fatal_errorf "Cannot use primitive %a for unboxed arrays in bytecode"
+        Printlambda.primitive p
   | Parrayrefu (Pgenarray_ref _) -> Kccall("caml_array_unsafe_get", 2)
   | Parrayrefu (Pfloatarray_ref _) -> Kccall("caml_floatarray_unsafe_get", 2)
   | Parrayrefu (Paddrarray_ref | Pintarray_ref) -> Kgetvectitem
+  | Parrayrefu (Punboxedfloatarray_ref | Punboxedintarray_ref _) ->
+      Misc.fatal_errorf "Cannot use primitive %a for unboxed arrays in bytecode"
+        Printlambda.primitive p
   | Parraysetu (Pgenarray_set _) -> Kccall("caml_array_unsafe_set", 3)
   | Parraysetu Pfloatarray_set -> Kccall("caml_floatarray_unsafe_set", 3)
   | Parraysetu (Paddrarray_set _ | Pintarray_set) -> Ksetvectitem
+  | Parraysetu (Punboxedfloatarray_set | Punboxedintarray_set _) ->
+      Misc.fatal_errorf "Cannot use primitive %a for unboxed arrays in bytecode"
+        Printlambda.primitive p
   | Pctconst c ->
      let const_name = match c with
        | Big_endian -> "big_endian"
@@ -826,7 +838,7 @@ let rec comp_expr stack_info env exp sz cont =
       let cont = add_pseudo_event loc !compunit_name cont in
       comp_args stack_info env args sz
         (Kmakefloatblock (List.length args) :: cont)
-  | Lprim(Pmakearray (kind, _, _), args, loc) ->
+  | Lprim((Pmakearray (kind, _, _)) as p, args, loc) ->
       let cont = add_pseudo_event loc !compunit_name cont in
       begin match kind with
         Pintarray | Paddrarray ->
@@ -841,6 +853,10 @@ let rec comp_expr stack_info env exp sz cont =
           else comp_args stack_info env args sz
                  (Kmakeblock(List.length args, 0) ::
                   Kccall("caml_make_array", 1) :: cont)
+      | Punboxedfloatarray | Punboxedintarray _ ->
+          Misc.fatal_errorf
+            "Cannot use Pmakeblock for unboxed arrays in bytecode"
+            Printlambda.primitive p
       end
   | Lprim((Presume|Prunstack), args, _) ->
       let nargs = List.length args - 1 in

lambda/lambda.ml

@@ -305,18 +305,24 @@ and block_shape =
 
 and array_kind =
     Pgenarray | Paddrarray | Pintarray | Pfloatarray
+  | Punboxedfloatarray
+  | Punboxedintarray of unboxed_integer
 
 and array_ref_kind =
   | Pgenarray_ref of alloc_mode
   | Paddrarray_ref
   | Pintarray_ref
   | Pfloatarray_ref of alloc_mode
+  | Punboxedfloatarray_ref
+  | Punboxedintarray_ref of unboxed_integer
 
 and array_set_kind =
   | Pgenarray_set of modify_mode
   | Paddrarray_set of modify_mode
   | Pintarray_set
   | Pfloatarray_set
+  | Punboxedfloatarray_set
+  | Punboxedintarray_set of unboxed_integer
 
 and boxed_integer = Primitive.boxed_integer =
     Pnativeint | Pint32 | Pint64
@@ -1517,8 +1523,10 @@ let primitive_may_allocate : primitive -> alloc_mode option = function
   | Pduparray _ -> Some alloc_heap
   | Parraylength _ -> None
   | Parraysetu _ | Parraysets _
-  | Parrayrefu (Paddrarray_ref | Pintarray_ref)
-  | Parrayrefs (Paddrarray_ref | Pintarray_ref) -> None
+  | Parrayrefu (Paddrarray_ref | Pintarray_ref
+      | Punboxedfloatarray_ref | Punboxedintarray_ref _)
+  | Parrayrefs (Paddrarray_ref | Pintarray_ref
+      | Punboxedfloatarray_ref | Punboxedintarray_ref _) -> None
   | Parrayrefu (Pgenarray_ref m | Pfloatarray_ref m)
   | Parrayrefs (Pgenarray_ref m | Pfloatarray_ref m) -> Some m
   | Pisint _ | Pisout -> None
@@ -1606,7 +1614,11 @@ let layout_of_native_repr : Primitive.native_repr -> _ = function
 let array_ref_kind_result_layout = function
   | Pintarray_ref -> layout_int
   | Pfloatarray_ref _ -> layout_boxed_float
+  | Punboxedfloatarray_ref -> layout_unboxed_float
   | Pgenarray_ref _ | Paddrarray_ref -> layout_field
+  | Punboxedintarray_ref Pint32 -> layout_unboxed_int32
+  | Punboxedintarray_ref Pint64 -> layout_unboxed_int64
+  | Punboxedintarray_ref Pnativeint -> layout_unboxed_nativeint
 
 let primitive_result_layout (p : primitive) =
   assert !Clflags.native_code;
@@ -1741,12 +1753,16 @@ let array_ref_kind mode = function
   | Paddrarray -> Paddrarray_ref
   | Pintarray -> Pintarray_ref
   | Pfloatarray -> Pfloatarray_ref mode
+  | Punboxedintarray int_kind -> Punboxedintarray_ref int_kind
+  | Punboxedfloatarray -> Punboxedfloatarray_ref
 
 let array_set_kind mode = function
   | Pgenarray -> Pgenarray_set mode
   | Paddrarray -> Paddrarray_set mode
   | Pintarray -> Pintarray_set
   | Pfloatarray -> Pfloatarray_set
+  | Punboxedintarray int_kind -> Punboxedintarray_set int_kind
+  | Punboxedfloatarray -> Punboxedfloatarray_set
 
 let is_check_enabled ~opt property =
   match property with

lambda/lambda.mli

@@ -252,6 +252,8 @@ and float_comparison =
 
 and array_kind =
     Pgenarray | Paddrarray | Pintarray | Pfloatarray
+  | Punboxedfloatarray
+  | Punboxedintarray of unboxed_integer
 
 (** When accessing a flat float array, we need to know the mode which we should
     box the resulting float at. *)
@@ -260,6 +262,8 @@ and array_ref_kind =
   | Paddrarray_ref
   | Pintarray_ref
   | Pfloatarray_ref of alloc_mode
+  | Punboxedfloatarray_ref
+  | Punboxedintarray_ref of unboxed_integer
 
 (** When updating an array that might contain pointers, we need to know what
     mode they're at; otherwise, access is uniform. *)
@@ -268,6 +272,8 @@ and array_set_kind =
   | Paddrarray_set of modify_mode
   | Pintarray_set
   | Pfloatarray_set
+  | Punboxedfloatarray_set
+  | Punboxedintarray_set of unboxed_integer
 
 and value_kind =
     Pgenval | Pfloatval | Pboxedintval of boxed_integer | Pintval

lambda/printlambda.ml

@@ -60,6 +60,10 @@ let array_kind = function
   | Paddrarray -> "addr"
   | Pintarray -> "int"
   | Pfloatarray -> "float"
+  | Punboxedfloatarray -> "unboxed_float"
+  | Punboxedintarray Pint32 -> "unboxed_int32"
+  | Punboxedintarray Pint64 -> "unboxed_int64"
+  | Punboxedintarray Pnativeint -> "unboxed_nativeint"
 
 let array_ref_kind ppf k =
   let pp_mode ppf = function
@@ -71,6 +75,10 @@ let array_ref_kind ppf k =
   | Paddrarray_ref -> fprintf ppf "addr"
   | Pintarray_ref -> fprintf ppf "int"
   | Pfloatarray_ref mode -> fprintf ppf "float%a" pp_mode mode
+  | Punboxedfloatarray_ref -> fprintf ppf "unboxed_float"
+  | Punboxedintarray_ref Pint32 -> fprintf ppf "unboxed_int32"
+  | Punboxedintarray_ref Pint64 -> fprintf ppf "unboxed_int64"
+  | Punboxedintarray_ref Pnativeint -> fprintf ppf "unboxed_nativeint"
 
 let array_set_kind ppf k =
   let pp_mode ppf = function
@@ -82,6 +90,10 @@ let array_set_kind ppf k =
   | Paddrarray_set mode -> fprintf ppf "addr%a" pp_mode mode
   | Pintarray_set -> fprintf ppf "int"
   | Pfloatarray_set -> fprintf ppf "float"
+  | Punboxedfloatarray_set -> fprintf ppf "unboxed_float"
+  | Punboxedintarray_set Pint32 -> fprintf ppf "unboxed_int32"
+  | Punboxedintarray_set Pint64 -> fprintf ppf "unboxed_int64"
+  | Punboxedintarray_set Pnativeint -> fprintf ppf "unboxed_nativeint"
 
 let alloc_mode_if_local = function
   | Alloc_heap -> ""

lambda/printlambda.mli

@@ -37,3 +37,4 @@ val print_bigarray :
   Lambda.bigarray_layout -> unit
 val check_attribute : formatter -> check_attribute -> unit
 val alloc_mode : formatter -> alloc_mode -> unit
+val array_kind : array_kind -> string

lambda/transl_array_comprehension.ml

@@ -565,6 +565,7 @@ let clause ~transl_exp ~scopes ~loc = function
     ([Fixed_size]); otherwise, we cannot ([Dynamic_size]), and we have to
     dynamically grow the array as we iterate and shrink it to size at the
     end. *)
+
 type array_sizing =
   | Fixed_size
   | Dynamic_size
@@ -696,13 +697,30 @@ let initial_array ~loc ~array_kind ~array_size ~array_sizing =
     | Fixed_size, (Pintarray | Paddrarray) ->
         Immutable StrictOpt,
         make_vect ~loc ~length:array_size.var ~init:(int 0)
-    | Fixed_size, Pfloatarray ->
+    | Fixed_size, (Pfloatarray | Punboxedfloatarray) ->
+        (* The representations of these two are the same, it's only
+           accesses that differ. *)
         Immutable StrictOpt, make_float_vect ~loc array_size.var
+    | Fixed_size , Punboxedintarray Pint32 ->
+        Immutable StrictOpt, make_unboxed_int32_vect ~loc array_size.var
+    | Fixed_size, Punboxedintarray Pint64 ->
+        Immutable StrictOpt, make_unboxed_int64_vect ~loc array_size.var
+    | Fixed_size, Punboxedintarray Pnativeint ->
+        Immutable StrictOpt, make_unboxed_nativeint_vect ~loc array_size.var
     (* Case 3: Unknown size, known array kind *)
     | Dynamic_size, (Pintarray | Paddrarray) ->
         Mutable, Resizable_array.make ~loc array_kind (int 0)
     | Dynamic_size, Pfloatarray ->
         Mutable, Resizable_array.make ~loc array_kind (float 0.)
+    | Dynamic_size, Punboxedfloatarray ->
+        Mutable, Resizable_array.make ~loc array_kind (unboxed_float 0.)
+    | Dynamic_size, Punboxedintarray Pint32 ->
+        Mutable, Resizable_array.make ~loc array_kind (unboxed_int32 0l)
+    | Dynamic_size, Punboxedintarray Pint64 ->
+        Mutable, Resizable_array.make ~loc array_kind (unboxed_int64 0L)
+    | Dynamic_size, Punboxedintarray Pnativeint ->
+        Mutable, Resizable_array.make ~loc array_kind
+          (unboxed_nativeint Targetint.zero)
   in
   Let_binding.make array_let_kind (Pvalue Pgenval) "array" array_value
 
@@ -790,15 +808,30 @@ let body
                Lassign(array.id, make_array),
                set_element_in_bounds elt.var,
                (Pvalue Pintval) (* [unit] is immediate *)))
-    | Pintarray | Paddrarray | Pfloatarray ->
+    | Pintarray | Paddrarray | Pfloatarray | Punboxedfloatarray
+    | Punboxedintarray _ ->
         set_element_in_bounds body
   in
   Lsequence(
     set_element_known_kind_in_bounds,
     Lassign(index.id, index.var + l1))
 
 let comprehension
-      ~transl_exp ~scopes ~loc ~array_kind { comp_body; comp_clauses } =
+      ~transl_exp ~scopes ~loc ~(array_kind : Lambda.array_kind)
+      { comp_body; comp_clauses } =
+  (match array_kind with
+  | Pgenarray | Paddrarray | Pintarray | Pfloatarray -> ()
+  | Punboxedfloatarray | Punboxedintarray _ ->
+    if not !Clflags.native_code then
+      Misc.fatal_errorf
+        "Array comprehensions for kind %s are not allowed in bytecode"
+        (Printlambda.array_kind array_kind);
+    if Targetint.size <> 64 then
+      Misc.fatal_errorf
+        "Array comprehensions for kind %s can only be compiled for \
+         64-bit native targets"
+        (Printlambda.array_kind array_kind)
+  );
   let { array_sizing_info; array_size; make_comprehension } =
     clauses ~transl_exp ~scopes ~loc comp_clauses
   in

lambda/transl_comprehension_utils.ml

@@ -36,6 +36,16 @@ module Lambda_utils = struct
 
     let float f = Lconst (Const_base (Const_float (Float.to_string f)))
 
+    let unboxed_float f =
+      Lconst (Const_base (Const_unboxed_float (Float.to_string f)))
+
+    let unboxed_int32 i = Lconst (Const_base (Const_unboxed_int32 i))
+    let unboxed_int64 i = Lconst (Const_base (Const_unboxed_int64 i))
+    let unboxed_nativeint i =
+      (* See CR in typedtree.mli *)
+      let i = i |> Targetint.to_int64 |> Int64.to_nativeint in
+      Lconst (Const_base (Const_unboxed_nativeint i))
+
     let string ~loc s = Lconst (Const_base (Const_string(s, loc, None)))
   end
 
@@ -140,6 +150,12 @@ module Lambda_utils = struct
 
     let make_float_vect = unary "caml_make_float_vect"
 
+    let make_unboxed_int32_vect = unary "caml_make_unboxed_int32_vect"
+
+    let make_unboxed_int64_vect = unary "caml_make_unboxed_int64_vect"
+
+    let make_unboxed_nativeint_vect = unary "caml_make_unboxed_nativeint_vect"
+
     let array_append = binary "caml_array_append"
 
     let array_sub =

lambda/transl_comprehension_utils.mli

@@ -57,6 +57,12 @@ module Lambda_utils : sig
         [Float.to_string] *)
     val float  : float -> lambda
 
+    (** Unboxed floats and ints *)
+    val unboxed_float : float -> lambda
+    val unboxed_int32 : Int32.t -> lambda
+    val unboxed_int64 : Int64.t -> lambda
+    val unboxed_nativeint : Targetint.t -> lambda
+
     (** Lambda string literals; these require a location, and are constructed as
         "quoted strings", not {fancy|delimited strings|fancy}. *)
     val string : loc:Location.t -> string -> lambda
@@ -118,6 +124,15 @@ module Lambda_utils : sig
         uninitialized *)
     val make_float_vect : loc:scoped_location -> lambda -> lambda
 
+    (** Like [make_float_vect] but for unboxed int32 arrays. *)
+    val make_unboxed_int32_vect : loc:scoped_location -> lambda -> lambda
+
+    (** Like [make_float_vect] but for unboxed int64 arrays. *)
+    val make_unboxed_int64_vect : loc:scoped_location -> lambda -> lambda
+
+    (** Like [make_float_vect] but for unboxed nativeint arrays. *)
+    val make_unboxed_nativeint_vect : loc:scoped_location -> lambda -> lambda
+
     (** [array_append a1 a2] calls the [caml_array_append] C primitive, which
         creates a new array by appending [a1] and [a2] *)
     val array_append : loc:scoped_location -> lambda -> lambda -> lambda

lambda/translcore.ml

@@ -643,6 +643,8 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
                   Lconst(Const_float_array(List.map extract_float cl))
                 | Pgenarray ->
                   raise Not_constant    (* can this really happen? *)
+                | Punboxedfloatarray | Punboxedintarray _ ->
+                  Misc.fatal_error "Use flambda2 for unboxed arrays"
             in
             match amut with
             | Mutable   -> duparray_to_mutable const