Support for unboxed products in the middle-end and backend

backend/amd64/proc.ml

@@ -734,6 +734,7 @@ let operation_supported = function
   | Ccheckbound
   | Cvectorcast _ | Cscalarcast _
   | Cprobe _ | Cprobe_is_enabled _ | Copaque | Cbeginregion | Cendregion
+  | Ctuple_field _
     -> true
 
 let trap_size_in_bytes = 16

backend/arm64/proc.ml

@@ -449,7 +449,7 @@ let operation_supported = function
   | Craise _
   | Ccheckbound
   | Cprobe _ | Cprobe_is_enabled _ | Copaque
-  | Cbeginregion | Cendregion
+  | Cbeginregion | Cendregion | Ctuple_field _
     -> true
 
 let trap_size_in_bytes = 16

backend/cmm.ml

@@ -241,6 +241,7 @@ type operation =
   | Cprobe_is_enabled of { name: string }
   | Copaque
   | Cbeginregion | Cendregion
+  | Ctuple_field of int * machtype array
 
 type kind_for_unboxing =
   | Any

backend/cmm.mli

@@ -238,6 +238,8 @@ type operation =
   | Cprobe_is_enabled of { name: string }
   | Copaque (* Sys.opaque_identity *)
   | Cbeginregion | Cendregion
+  | Ctuple_field of int * machtype array
+      (* the [machtype array] refers to the whole tuple *)
 
 (* This is information used exclusively during construction of cmm terms by
    cmmgen, and thus irrelevant for selectgen and flambda2. *)

backend/cmm_helpers.ml

@@ -1126,7 +1126,7 @@ module Extended_machtype = struct
   let change_tagged_int_to_val t =
     Array.map Extended_machtype_component.change_tagged_int_to_val t
 
-  let of_layout (layout : Lambda.layout) =
+  let rec of_layout (layout : Lambda.layout) =
     match layout with
     | Ptop -> Misc.fatal_error "No Extended_machtype for layout [Ptop]"
     | Pbottom ->
@@ -1138,6 +1138,7 @@ module Extended_machtype = struct
       typ_any_int
     | Pvalue Pintval -> typ_tagged_int
     | Pvalue _ -> typ_val
+    | Punboxed_product fields -> Array.concat (List.map of_layout fields)
 end
 
 let machtype_of_layout layout =
@@ -4379,19 +4380,19 @@ let direct_call ~dbg ty pos f_code_sym args =
   Cop (Capply (ty, pos), f_code_sym :: args, dbg)
 
 let indirect_call ~dbg ty pos alloc_mode f args_type args =
-  match args with
-  | [arg] ->
+  match args_type with
+  | [_] ->
     (* Use a variable to avoid duplicating the cmm code of the closure [f]. *)
     let v = Backend_var.create_local "*closure*" in
     let v' = Backend_var.With_provenance.create v in
     letin v' ~defining_expr:f
       ~body:
         (Cop
            ( Capply (Extended_machtype.to_machtype ty, pos),
-             [load ~dbg Word_int Asttypes.Mutable ~addr:(Cvar v); arg; Cvar v],
+             (load ~dbg Word_int Asttypes.Mutable ~addr:(Cvar v) :: args)
+             @ [Cvar v],
              dbg ))
-  | args ->
-    call_caml_apply ty args_type Asttypes.Mutable f args pos alloc_mode dbg
+  | _ -> call_caml_apply ty args_type Asttypes.Mutable f args pos alloc_mode dbg
 
 let indirect_full_call ~dbg ty pos alloc_mode f args_type = function
   (* the single-argument case is already optimized by indirect_call *)
@@ -4528,5 +4529,6 @@ let kind_of_layout (layout : Lambda.layout) =
   | Pvalue (Pboxedintval bi) -> Boxed_integer bi
   | Pvalue (Pboxedvectorval vi) -> Boxed_vector vi
   | Pvalue (Pgenval | Pintval | Pvariant _ | Parrayval _)
-  | Ptop | Pbottom | Punboxed_float | Punboxed_int _ | Punboxed_vector _ ->
+  | Ptop | Pbottom | Punboxed_float | Punboxed_int _ | Punboxed_vector _
+  | Punboxed_product _ ->
     Any

backend/cmmgen.ml

@@ -164,6 +164,8 @@ let get_field env layout ptr n dbg =
     | Pvalue _ -> Word_val
     | Punboxed_float -> Double
     | Punboxed_vector (Pvec128 _) -> Onetwentyeight
+    | Punboxed_product _ ->
+      Misc.fatal_error "Unboxed products cannot be stored as fields for now."
     | Ptop ->
         Misc.fatal_errorf "get_field with Ptop: %a" Debuginfo.print_compact dbg
     | Pbottom ->
@@ -617,6 +619,8 @@ let rec transl env e =
   (* Primitives *)
   | Uprim(prim, args, dbg) ->
       begin match (simplif_primitive prim, args) with
+      | (Pmake_unboxed_product layouts, args) ->
+          Ctuple (List.map (transl env) args)
       | (Pread_symbol sym, []) ->
           Cconst_symbol (global_symbol sym, dbg)
       | ((Pmakeblock _ | Pmakeufloatblock _), []) ->
@@ -728,7 +732,7 @@ let rec transl env e =
          | Pasrbint _ | Pbintcomp (_, _) | Pstring_load _ | Pbytes_load _
          | Pbytes_set _ | Pbigstring_load _ | Pbigstring_set _
          | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _
-         | Pbbswap _ | Pget_header _), _)
+         | Punboxed_product_field _ | Pbbswap _ | Pget_header _), _)
         ->
           fatal_error "Cmmgen.transl:prim"
       end
@@ -1044,6 +1048,9 @@ and transl_prim_1 env p arg dbg =
   | Pbswap16 ->
       tag_int (bswap16 (ignore_high_bit_int (untag_int
         (transl env arg) dbg)) dbg) dbg
+  | Punboxed_product_field (field, layouts) ->
+    let layouts = Array.of_list (List.map machtype_of_layout layouts) in
+    Cop (Ctuple_field (field, layouts), [transl env arg], dbg)
   | Pget_header m ->
       box_int dbg Pnativeint m (get_header (transl env arg) dbg)
   | (Pfield_computed | Psequand | Psequor
@@ -1063,7 +1070,9 @@ and transl_prim_1 env p arg dbg =
     | Plslbint _ | Plsrbint _ | Pasrbint _ | Pbintcomp (_, _)
     | Pbigarrayref (_, _, _, _) | Pbigarrayset (_, _, _, _)
     | Pbigarraydim _ | Pstring_load _ | Pbytes_load _ | Pbytes_set _
-    | Pbigstring_load _ | Pbigstring_set _ | Pprobe_is_enabled _)
+    | Pbigstring_load _ | Pbigstring_set _ | Pprobe_is_enabled _
+    | Pmake_unboxed_product _
+    )
     ->
       fatal_errorf "Cmmgen.transl_prim_1: %a"
         Printclambda_primitives.primitive p
@@ -1245,7 +1254,9 @@ and transl_prim_2 env p arg1 arg2 dbg =
   | Pnegbint _ | Pbigarrayref (_, _, _, _) | Pbigarrayset (_, _, _, _)
   | Pbigarraydim _ | Pbytes_set _ | Pbigstring_set _ | Pbbswap _
   | Pprobe_is_enabled _
-  | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _ | Pget_header _
+  | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _
+  | Pmake_unboxed_product _ | Punboxed_product_field _
+  | Pget_header _
     ->
       fatal_errorf "Cmmgen.transl_prim_2: %a"
         Printclambda_primitives.primitive p
@@ -1307,7 +1318,9 @@ and transl_prim_3 env p arg1 arg2 arg3 dbg =
   | Pbigarrayref (_, _, _, _) | Pbigarrayset (_, _, _, _) | Pbigarraydim _
   | Pstring_load _ | Pbytes_load _ | Pbigstring_load _ | Pbbswap _
   | Pprobe_is_enabled _
-  | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _ | Pget_header _
+  | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _
+  | Pmake_unboxed_product _ | Punboxed_product_field _
+  | Pget_header _
     ->
       fatal_errorf "Cmmgen.transl_prim_3: %a"
         Printclambda_primitives.primitive p
@@ -1384,7 +1397,8 @@ and transl_let env str (layout : Lambda.layout) id exp transl_body =
        there may be constant closures inside that need lifting out. *)
     let _cbody : expression = transl_body env in
     cexp
-  | Punboxed_float | Punboxed_int _ | Punboxed_vector _ -> begin
+  | Punboxed_float | Punboxed_int _ | Punboxed_vector _ | Punboxed_product _ ->
+    begin
       let cexp = transl env exp in
       let cbody = transl_body env in
       match str with
@@ -1393,7 +1407,7 @@ and transl_let env str (layout : Lambda.layout) id exp transl_body =
       | Mutable ->
         let typ = machtype_of_layout layout in
         Clet_mut(id, typ, cexp, cbody)
-  end
+    end
   | Pvalue kind ->
     transl_let_value env str kind id exp transl_body
 

backend/printcmm.ml

@@ -247,6 +247,8 @@ let operation d = function
   | Copaque -> "opaque"
   | Cbeginregion -> "beginregion"
   | Cendregion -> "endregion"
+  | Ctuple_field (field, _ty) ->
+    to_string "tuple_field %i" field
 
 let rec expr ppf = function
   | Cconst_int (n, _dbg) -> fprintf ppf "%i" n

backend/selectgen.ml

@@ -236,6 +236,7 @@ let oper_result_type = function
        naked pointer into the local allocation stack. *)
     typ_int
   | Cendregion -> typ_void
+  | Ctuple_field (field, fields_ty) -> fields_ty.(field)
 
 (* Infer the size in bytes of the result of an expression whose evaluation
    may be deferred (cf. [emit_parts]). *)
@@ -537,6 +538,7 @@ method is_simple_expr = function
       | Cabsf | Caddf | Csubf | Cmulf | Cdivf | Cfloatofint | Cintoffloat
       | Cvectorcast _ | Cscalarcast _
       | Cvalueofint | Cintofvalue
+      | Ctuple_field _
       | Ccmpf _ -> List.for_all self#is_simple_expr args
       end
   | Cassign _ | Cifthenelse _ | Cswitch _ | Ccatch _ | Cexit _
@@ -584,6 +586,7 @@ method effects_of exp =
       | Cload (_, Asttypes.Immutable) -> EC.none
       | Cload (_, Asttypes.Mutable) -> EC.coeffect_only Coeffect.Read_mutable
       | Cprobe_is_enabled _ -> EC.coeffect_only Coeffect.Arbitrary
+      | Ctuple_field _
       | Caddi | Csubi | Cmuli | Cmulhi _ | Cdivi | Cmodi | Cand | Cor | Cxor
       | Cbswap _
       | Ccsel _
@@ -979,6 +982,21 @@ method emit_expr_aux (env:environment) exp ~bound_name :
          let rs = self#emit_tuple env simple_args in
          ret (self#insert_op_debug env Iopaque dbg rs rs)
       end
+  | Cop(Ctuple_field(field, fields_layout), [arg], dbg) ->
+      begin match self#emit_expr env arg ~bound_name:None with
+        None -> None
+      | Some loc_exp ->
+        let flat_size a =
+          Array.fold_left (fun acc t -> acc + Array.length t) 0 a
+        in
+        assert(Array.length loc_exp = flat_size fields_layout);
+        let before = Array.sub fields_layout 0 field in
+        let size_before = flat_size before in
+        let field_slice =
+          Array.sub loc_exp size_before (Array.length fields_layout.(field))
+        in
+        ret field_slice
+      end
   | Cop(op, args, dbg) ->
       begin match self#emit_parts_list env args with
         None -> None

dune

@@ -87,6 +87,7 @@
   backend_var
   clambda
   clambda_primitives
+  clambda_layout
   compilenv
   mangling
   convert_primitives

middle_end/.ocamlformat-enable

@@ -1,2 +1,4 @@
+clambda_layout.ml
+clambda_layout.mli
 mangling.ml
 mangling.mli

middle_end/clambda_layout.ml

@@ -0,0 +1,136 @@
+(**************************************************************************)
+(*                                                                        *)
+(*                                 OCaml                                  *)
+(*                                                                        *)
+(*                        Pierre Chambart, OCamlPro                       *)
+(*                                                                        *)
+(*   Copyright 2023 OCamlPro SAS                                          *)
+(*                                                                        *)
+(*   All rights reserved.  This file is distributed under the terms of    *)
+(*   the GNU Lesser General Public License version 2.1, with the          *)
+(*   special exception on linking described in the file LICENSE.          *)
+(*                                                                        *)
+(**************************************************************************)
+
+type atom =
+  | Value
+  | Value_int
+  | Unboxed_float
+  | Unboxed_int of Lambda.boxed_integer
+  | Unboxed_vector of Lambda.boxed_vector
+
+let rec fold_left_layout (f : 'acc -> 'e -> atom -> 'acc) (acc : 'acc)
+    (expr : Clambda.ulambda) (layout : Clambda_primitives.layout) : 'acc =
+  match layout with
+  | Ptop -> Misc.fatal_error "[Ptop] can't be stored in a closure."
+  | Pbottom ->
+    Misc.fatal_error
+      "[Pbottom] should have been eliminated as dead code and not stored in a \
+       closure."
+  | Punboxed_float -> f acc expr Unboxed_float
+  | Punboxed_int bi -> f acc expr (Unboxed_int bi)
+  | Punboxed_vector bv -> f acc expr (Unboxed_vector bv)
+  | Pvalue Pintval -> f acc expr Value_int
+  | Pvalue _ -> f acc expr Value
+  | Punboxed_product layouts ->
+    List.fold_left
+      (fun acc (field, layout) ->
+        let expr : Clambda.ulambda =
+          Uprim (Punboxed_product_field (field, layouts), [expr], Debuginfo.none)
+        in
+        fold_left_layout f acc expr layout)
+      acc
+      (List.mapi (fun i v -> i, v) layouts)
+
+type ('visible, 'invisible) decomposition' =
+  | Gc_visible of ('visible * atom)
+  | Gc_invisible of ('invisible * atom)
+  | Product of ('visible, 'invisible) decomposition' array
+
+type decomposition =
+  | Atom of
+      { offset : int;
+        layout : atom
+      }
+  | Product of decomposition array
+
+let print_atom ppf = function
+  | Value -> Format.fprintf ppf "val"
+  | Value_int -> Format.fprintf ppf "int"
+  | Unboxed_float -> Format.fprintf ppf "#float"
+  | Unboxed_int Pint32 -> Format.fprintf ppf "unboxed_int32"
+  | Unboxed_int Pint64 -> Format.fprintf ppf "unboxed_int64"
+  | Unboxed_int Pnativeint -> Format.fprintf ppf "unboxed_nativeint"
+  | Unboxed_vector (Pvec128 _) -> Format.fprintf ppf "unboxed_vec128"
+
+let equal_decomposition = ( = )
+
+let rec print_decomposition ppf dec =
+  match dec with
+  | Atom { offset; layout } ->
+    Format.fprintf ppf "(%d: %a)" offset print_atom layout
+  | Product a ->
+    Format.fprintf ppf "@[<hov 2>[%a]@]"
+      (Format.pp_print_list ~pp_sep:Format.pp_print_space print_decomposition)
+      (Array.to_list a)
+
+let rec decompose (layout : Lambda.layout) : _ decomposition' =
+  match layout with
+  | Ptop -> Misc.fatal_error "[Ptop] can't be stored in a closure."
+  | Pbottom ->
+    Misc.fatal_error
+      "[Pbottom] should have been eliminated as dead code and not stored in a \
+       closure."
+  | Punboxed_float -> Gc_invisible ((), Unboxed_float)
+  | Punboxed_int bi -> Gc_invisible ((), Unboxed_int bi)
+  | Punboxed_vector bv -> Gc_invisible ((), Unboxed_vector bv)
+  | Pvalue Pintval -> Gc_invisible ((), Value_int)
+  | Pvalue _ -> Gc_visible ((), Value)
+  | Punboxed_product l -> Product (Array.of_list (List.map decompose l))
+
+let rec solidify (dec : (int, int) decomposition') : decomposition =
+  match dec with
+  | Gc_visible (offset, layout) -> Atom { offset; layout }
+  | Gc_invisible (offset, layout) -> Atom { offset; layout }
+  | Product a -> Product (Array.map solidify a)
+
+let rec fold_decompose (f1 : 'acc -> 'a -> atom -> 'acc * 'b)
+    (f2 : 'acc -> 'c -> atom -> 'acc * 'd) (acc : 'acc)
+    (d : ('a, 'c) decomposition') : 'acc * ('b, 'd) decomposition' =
+  match d with
+  | Gc_visible (v, layout) ->
+    let acc, v = f1 acc v layout in
+    acc, Gc_visible (v, layout)
+  | Gc_invisible (v, layout) ->
+    let acc, v = f2 acc v layout in
+    acc, Gc_invisible (v, layout)
+  | Product elts ->
+    let acc, elts = Array.fold_left_map (fold_decompose f1 f2) acc elts in
+    acc, Product elts
+
+let atom_size (layout : atom) =
+  match layout with
+  | Value | Value_int | Unboxed_float | Unboxed_int _ -> 1
+  | Unboxed_vector (Pvec128 _) -> 2
+
+let assign_invisible_offsets init_pos (var, dec) =
+  let f_visible acc () _layout = acc, () in
+  let f_invisible acc () layout = acc + atom_size layout, acc in
+  let acc, dec = fold_decompose f_visible f_invisible init_pos dec in
+  acc, (var, dec)
+
+let assign_visible_offsets init_pos (var, dec) =
+  let f_visible acc () layout = acc + atom_size layout, acc in
+  let f_invisible acc off _layout = acc, off in
+  let acc, dec = fold_decompose f_visible f_invisible init_pos dec in
+  acc, (var, solidify dec)
+
+let decompose_free_vars ~base_offset ~free_vars =
+  let free_vars = List.map (fun (var, kind) -> var, decompose kind) free_vars in
+  let base_offset, free_vars =
+    List.fold_left_map assign_invisible_offsets base_offset free_vars
+  in
+  let _base_offset, free_vars =
+    List.fold_left_map assign_visible_offsets base_offset free_vars
+  in
+  free_vars