Add test for param unboxing + Fixes

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -2235,7 +2235,8 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot
     ( Function_slot.Map.add function_slot approx by_function_slot,
       function_code_ids ) )
 
-let close_functions acc external_env ~current_region function_declarations =
+let close_functions acc external_env ~current_region function_declarations
+    ~ident_map =
   let compilation_unit = Compilation_unit.get_current_exn () in
   let value_slots_from_idents =
     Ident.Set.fold
@@ -2273,12 +2274,31 @@ let close_functions acc external_env ~current_region function_declarations =
     && Flambda_features.classic_mode ()
   in
   let func_decl_list = Function_decls.to_list function_declarations in
+  let idents_for_unboxed_decls =
+    List.fold_left
+      (fun idents_for_unboxed_decls decl ->
+        match Function_decl.calling_convention decl with
+        | Normal_calling_convention -> idents_for_unboxed_decls
+        | Unboxed_calling_convention (_, _, unboxed_function_slot) ->
+          let unboxed_ident = Ident.rename (Function_decl.let_rec_ident decl) in
+          Function_slot.Map.add unboxed_function_slot unboxed_ident
+            idents_for_unboxed_decls)
+      Function_slot.Map.empty func_decl_list
+  in
   let function_slots_from_idents =
     List.fold_left
       (fun map decl ->
         let id = Function_decl.let_rec_ident decl in
         let function_slot = Function_decl.function_slot decl in
-        Ident.Map.add id function_slot map)
+        let map = Ident.Map.add id function_slot map in
+        match Function_decl.calling_convention decl with
+        | Normal_calling_convention -> map
+        | Unboxed_calling_convention (_, _, unboxed_function_slot) ->
+          let unboxed_id =
+            Function_slot.Map.find unboxed_function_slot
+              idents_for_unboxed_decls
+          in
+          Ident.Map.add unboxed_id unboxed_function_slot map)
       Ident.Map.empty func_decl_list
   in
   let function_code_ids =
@@ -2348,8 +2368,6 @@ let close_functions acc external_env ~current_region function_declarations =
             ~loopify:Never_loopify
         in
         let code = Code_or_metadata.create_metadata_only metadata in
-        (* CR ncourant: do we need to add the unboxed function slot to the
-           approx map? *)
         let all_function_slots =
           Ident.Map.data function_slots_from_idents |> Function_slot.Set.of_list
         in
@@ -2366,7 +2384,15 @@ let close_functions acc external_env ~current_region function_declarations =
               symbol = None
             }
         in
-        Function_slot.Map.add function_slot approx approx_map)
+        let approx_map =
+          Function_slot.Map.add function_slot approx approx_map
+        in
+        match Function_decl.calling_convention decl with
+        | Normal_calling_convention -> approx_map
+        | Unboxed_calling_convention (_, _, unboxed_function_slot) ->
+          (* CR mshinwell: add a proper approximation here *)
+          Function_slot.Map.add unboxed_function_slot
+            Value_approximation.Value_unknown approx_map)
       Function_slot.Map.empty func_decl_list
   in
   let acc, external_env, symbol_map =
@@ -2377,8 +2403,9 @@ let close_functions acc external_env ~current_region function_declarations =
           let env, acc, symbol =
             declare_symbol_for_function_slot env acc ident function_slot
           in
+          let approx = Function_slot.Map.find function_slot approx_map in
           let approx =
-            match Function_slot.Map.find function_slot approx_map with
+            match approx with
             | Value_approximation.Closure_approximation
                 { code_id;
                   function_slot;
@@ -2395,6 +2422,7 @@ let close_functions acc external_env ~current_region function_declarations =
                   code;
                   symbol = Some symbol
                 }
+            | Value_unknown -> Value_approximation.Value_unknown
             | _ -> assert false
             (* see above *)
           in
@@ -2404,6 +2432,7 @@ let close_functions acc external_env ~current_region function_declarations =
         (acc, external_env, Function_slot.Map.empty)
     else acc, external_env, Function_slot.Map.empty
   in
+  (* CR mshinwell: rename [approximations] *)
   let acc, (approximations, function_code_ids_in_order) =
     List.fold_left
       (fun (acc, (by_function_slot, function_code_ids_in_order)) function_decl ->
@@ -2446,14 +2475,25 @@ let close_functions acc external_env ~current_region function_declarations =
         Value_slot.Map.add value_slot external_simple map)
       value_slots_from_idents Value_slot.Map.empty
   in
+  let unboxed_function_slots =
+    List.concat_map
+      (fun decl ->
+        match Function_decl.calling_convention decl with
+        | Normal_calling_convention -> []
+        | Unboxed_calling_convention (_, _, function_slot) -> [function_slot])
+      func_decl_list
+    |> Function_slot.Set.of_list
+  in
   let approximations =
     Function_slot.Map.mapi
       (fun function_slot code ->
         let code_id =
           Code_metadata.code_id (Code_or_metadata.code_metadata code)
         in
         let all_function_slots =
-          Function_slot.Lmap.keys funs |> Function_slot.Set.of_list
+          Function_slot.Set.union
+            (Function_slot.Lmap.keys funs |> Function_slot.Set.of_list)
+            unboxed_function_slots
         in
         Value_approximation.Closure_approximation
           { code_id;
@@ -2475,6 +2515,10 @@ let close_functions acc external_env ~current_region function_declarations =
   let acc =
     Acc.add_set_of_closures_offsets ~is_phantom:false acc set_of_closures
   in
+  let ident_map =
+    Function_slot.Map.disjoint_union ~eq:Ident.same ident_map
+      idents_for_unboxed_decls
+  in
   if can_be_lifted
   then
     let symbols_with_approx =
@@ -2507,8 +2551,8 @@ let close_functions acc external_env ~current_region function_declarations =
     in
     let symbols = Function_slot.Lmap.map fst symbols_with_approx in
     let acc = Acc.add_lifted_set_of_closures ~symbols ~set_of_closures acc in
-    acc, Lifted symbols_with_approx
-  else acc, Dynamic (set_of_closures, approximations)
+    acc, Lifted symbols_with_approx, ident_map
+  else acc, Dynamic (set_of_closures, approximations), ident_map
 
 let close_let_rec acc env ~function_declarations
     ~(body : Acc.t -> Env.t -> Expr_with_acc.t) ~current_region =
@@ -2560,10 +2604,10 @@ let close_let_rec acc env ~function_declarations
     | None ->
       Misc.fatal_error "let-rec group of [lfunction] declarations is empty"
   in
-  let acc, closed_functions =
+  let acc, closed_functions, ident_map =
     close_functions acc env
       (Function_decls.create function_declarations alloc_mode)
-      ~current_region
+      ~current_region ~ident_map
   in
   match closed_functions with
   | Lifted symbols ->

middle_end/flambda2/from_lambda/closure_conversion_aux.ml

@@ -746,6 +746,11 @@ module Function_decls = struct
       | Unboxed_number of Flambda_kind.Boxable_number.t
       | Unboxed_float_record of int
 
+    let num_params_for_unboxing_kind = function
+      | Fields_of_block_with_tag_zero l -> List.length l
+      | Unboxed_number _ -> 1
+      | Unboxed_float_record n -> n
+
     type calling_convention =
       | Normal_calling_convention
       | Unboxed_calling_convention of

middle_end/flambda2/from_lambda/closure_conversion_aux.mli

@@ -321,6 +321,8 @@ module Function_decls : sig
       | Unboxed_number of Flambda_kind.Boxable_number.t
       | Unboxed_float_record of int
 
+    val num_params_for_unboxing_kind : unboxing_kind -> int
+
     type calling_convention =
       | Normal_calling_convention
       | Unboxed_calling_convention of

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -1566,6 +1566,13 @@ and cps_function env ~fid ~(recursive : Recursive.t) ?precomputed_free_idents
     in
     if attr.stub || ((not attr.unbox_return) && not is_a_param_unboxed)
     then Normal_calling_convention
+    else if Flambda_features.classic_mode ()
+    then
+      (* Unboxing of arguments/returns in classic mode does not yield the
+         expected benefits, because that would require to inline and simplify
+         stubs or regenerate them on the fly at callsites. Therefore we disable
+         it in classic mode to avoid generating worse code. *)
+      Normal_calling_convention
     else
       let unboxed_function_slot =
         Function_slot.create
@@ -1578,7 +1585,9 @@ and cps_function env ~fid ~(recursive : Recursive.t) ?precomputed_free_idents
       in
       let unboxed_param (param : Lambda.lparam) =
         if param.attributes.unbox_param
-        then unboxing_kind param.layout
+        then
+          (* No limit is currently placed on the number of return values. *)
+          unboxing_kind param.layout
         else None
       in
       let unboxed_params =
@@ -1592,8 +1601,23 @@ and cps_function env ~fid ~(recursive : Recursive.t) ?precomputed_free_idents
                  Misc.fatal_error "Trying to unbox an unboxed product.")
              params unarized_per_param)
       in
-      Unboxed_calling_convention
-        (unboxed_params, unboxed_return, unboxed_function_slot)
+      let non_unboxed_params_length =
+        List.fold_left (fun acc l -> acc + List.length l) 0 unarized_per_param
+      in
+      let unboxed_params_length =
+        List.fold_left
+          (fun length unboxed_param ->
+            match unboxed_param with
+            | None -> length
+            | Some unboxing_kind ->
+              length + Function_decl.num_params_for_unboxing_kind unboxing_kind)
+          0 unboxed_params
+      in
+      if non_unboxed_params_length + unboxed_params_length > 127
+      then Normal_calling_convention
+      else
+        Unboxed_calling_convention
+          (unboxed_params, unboxed_return, unboxed_function_slot)
   in
   let body_cont =
     match calling_convention with

ocaml/testsuite/tests/unboxed-params/classic.ml

@@ -0,0 +1,93 @@
+(* TEST_BELOW *)
+
+(* Unboxing of function params/returns is not available in classic mode,
+   but we should at least check that it does not make the compiler crash. *)
+[@@@ocaml.flambda_oclassic]
+
+(* About testing for allocations.
+
+   The call to `Gc.counters` itself allocates a bit (about 10 words: a 3-tuple
+   plus some floats), so we will never have zero allocation. It is a bit
+   fragile to check the exact number of words allocated, so instead we check
+   that the number of allocated minor words is "not a lot more" than what
+   `Gc.counters` allocates by using a reasonable threshold.
+
+   To clearly differentiate the allocations from `Gc.counters` from those that
+   are (potentially) in a function `f`, we run `f` in a loop enough times so
+   that if it had any allocation, the number of minor words allocated would be
+   a lot bigger than the threshold we check for. *)
+let test_allocs test_name f arg1 arg2 =
+  let (minor, promoted, major) = Gc.counters () in
+  for i = 1 to 1_000 do
+    if not (f arg1 arg2) then failwith "incorrect result"
+  done;
+  let (minor', promoted', major') = Gc.counters () in
+  if minor' <= minor +. 20. && promoted = promoted' && major = major' then
+    Format.printf "%s: allocs ok.@." test_name
+  else
+    Format.printf "%s: allocation check failed.@." test_name
+[@@inline never]
+
+(* Check unboxability of floats *)
+module Floats = struct
+  let[@unboxable] f (x[@unboxable]) y = x +. y [@@inline never]
+  let g t y = let x = t +. 1. in f x y = 0. [@@inline never]
+end
+
+(* Check unboxability of Int32s *)
+module Int32s = struct
+  let[@unboxable] f (x[@unboxable]) y = Int32.add x y [@@inline never]
+  let g t y = let x = Int32.add t Int32.one in f x y = Int32.zero [@@inline never]
+end
+
+(* Check unboxability of Int32s *)
+module Int64s = struct
+  let[@unboxable] f (x[@unboxable]) y = Int64.add x y [@@inline never]
+  let g t y = let x = Int64.add t Int64.one in f x y = Int64.zero [@@inline never]
+end
+
+(* Check unboxability of Int32s *)
+module Nativeints = struct
+  let[@unboxable] f (x[@unboxable]) y = Nativeint.add x y [@@inline never]
+  let g t y = let x = Nativeint.add t Nativeint.one in f x y = Nativeint.zero [@@inline never]
+end
+
+(* Check unboability of tuples *)
+module Tuples = struct
+  let[@unboxable] f ((x, y)[@unboxable]) = (y, x) [@@inline never]
+  let g (a : int) (b : int) = let x, y = f (a, b) in x = b && y = a [@@inline never]
+end
+
+(* This does not work **as of now** *)
+module Variants = struct
+  type t = A of int | B of int
+  let[@unboxable] f (t[@unboxable]) =
+    match t with
+    | A i -> B i
+    | B i -> A i
+  [@@inline never]
+
+  let g i j =
+    (match f (A i) with
+    | B i' -> i = i'
+    | A _ -> false) &&
+    (match f (B j) with
+    | A j' -> j = j'
+    | B _ -> false)
+  [@@inline never]
+end
+
+(* Actual tests *)
+let () =
+  test_allocs "floats" Floats.g 0. (- 1.);
+  test_allocs "int32s" Int32s.g 0l (- 1l);
+  test_allocs "int64s" Int64s.g 0L (- 1L);
+  test_allocs "nativeints" Nativeints.g 0n (- 1n);
+  test_allocs "tuples" Tuples.g 13 42;
+  test_allocs "variants" Variants.g 13 42;
+  ()
+
+(* TEST
+   * flambda2
+   ** native
+*)

ocaml/testsuite/tests/unboxed-params/classic.reference

@@ -0,0 +1,6 @@
+floats: allocation check failed.
+int32s: allocation check failed.
+int64s: allocation check failed.
+nativeints: allocation check failed.
+tuples: allocation check failed.
+variants: allocation check failed.

ocaml/testsuite/tests/unboxed-params/test.compilers.reference

@@ -0,0 +1,15 @@
+File "test.ml", lines 63-66, characters 20-16:
+63 | ....................(t[@unboxable]) =
+64 |     match t with
+65 |     | A i -> B i
+66 |     | B i -> A i
+Warning 210 [unboxing-impossible]: This [@unboxable] attribute cannot be used.
+The type of this value does not allow unboxing.
+
+File "test.ml", lines 63-66, characters 20-16:
+63 | ....................(t[@unboxable]) =
+64 |     match t with
+65 |     | A i -> B i
+66 |     | B i -> A i
+Warning 210 [unboxing-impossible]: This [@unboxable] attribute cannot be used.
+The type of this value does not allow unboxing.