flambda-backend: Functions with unboxed parameters and returns (#1271)

Co-authored-by: Vincent Laviron <[email protected]>
Co-authored-by: Guillaume Bury <[email protected]>

Author: 3 people

lambda/lambda.ml

@@ -602,11 +602,14 @@ type function_attribute = {
   stub: bool;
   tmc_candidate: bool;
   may_fuse_arity: bool;
+  unbox_return: bool;
 }
 
 type scoped_location = Debuginfo.Scoped_location.t
 
-type parameter_attribute = No_attributes
+type parameter_attribute = {
+  unbox_param: bool;
+}
 
 type lparam = {
   name : Ident.t;
@@ -803,12 +806,13 @@ let default_function_attribute = {
      them multi-argument. So, we keep arity fusion turned on by default for now.
   *)
   may_fuse_arity = true;
+  unbox_return = false;
 }
 
 let default_stub_attribute =
   { default_function_attribute with stub = true; check = Ignore_assert_all Zero_alloc }
 
-let default_param_attribute = No_attributes
+let default_param_attribute = { unbox_param = false }
 
 (* Build sharing keys *)
 (*

lambda/lambda.mli

@@ -499,9 +499,12 @@ type function_attribute = {
   (* [may_fuse_arity] is true if [simplif.ml] is permitted to fuse arity, i.e.,
      to perform the rewrite [fun x -> fun y -> e] to [fun x y -> e] *)
   may_fuse_arity: bool;
+  unbox_return: bool;
 }
 
-type parameter_attribute = No_attributes
+type parameter_attribute = {
+  unbox_param: bool;
+}
 
 type lparam = {
   name : Ident.t;

lambda/printlambda.ml

@@ -885,21 +885,23 @@ let rec lam ppf = function
         | Curried {nlocal} ->
             fprintf ppf "@ {nlocal = %d}" nlocal;
             List.iter (fun (p : Lambda.lparam) ->
-                (* Make sure we change this once there are attributes *)
-                let No_attributes = p.attributes in
-                fprintf ppf "@ %a%s%a"
-                  Ident.print p.name (alloc_kind p.mode) layout p.layout) params
+                let { unbox_param } = p.attributes in
+                fprintf ppf "@ %a%s%a%s"
+                  Ident.print p.name (alloc_kind p.mode) layout p.layout
+                  (if unbox_param then "[@unboxed]" else "")
+              ) params
         | Tupled ->
             fprintf ppf " (";
             let first = ref true in
             List.iter
               (fun (p : Lambda.lparam) ->
-                 (* Make sure we change this once there are attributes *)
-                 let No_attributes = p.attributes in
+                 let { unbox_param } = p.attributes in
                  if !first then first := false else fprintf ppf ",@ ";
                  Ident.print ppf p.name;
                  Format.fprintf ppf "%s" (alloc_kind p.mode);
-                 layout ppf p.layout)
+                 layout ppf p.layout;
+                 if unbox_param then Format.fprintf ppf "[@unboxed]"
+              )
               params;
             fprintf ppf ")" in
       fprintf ppf "@[<2>(function%s%a@ %a%a%a)@]"

lambda/translattribute.ml

@@ -53,6 +53,8 @@ let is_loop_attribute =
 let is_opaque_attribute =
   [ ["opaque"; "ocaml.opaque"], true ]
 
+let is_unboxed_attribute =
+  [ ["unboxed"; "ocaml.unboxed"], true ]
 
 let find_attribute ?mark_used p attributes =
   let inline_attribute =
@@ -586,6 +588,21 @@ let add_opaque_attribute expr loc attributes =
       end
   | _ -> expr
 
+let add_unbox_return_attribute expr loc attributes =
+  match expr with
+  | Lfunction funct ->
+      let attr = find_attribute is_unboxed_attribute attributes in
+      begin match attr with
+      | None -> expr
+      | Some _ ->
+          if funct.attr.unbox_return then
+            Location.prerr_warning loc
+              (Warnings.Duplicated_attribute "unboxed");
+          let attr = { funct.attr with unbox_return = true } in
+          lfunction_with_attr ~attr funct
+      end
+  | _ -> expr
+
 
 (* Get the [@inlined] attribute payload (or default if not present). *)
 let get_inlined_attribute e =
@@ -646,6 +663,9 @@ let add_function_attributes lam loc attr =
   let lam =
     add_tmc_attribute lam loc attr
   in
+  let lam =
+    add_unbox_return_attribute lam loc attr
+  in
   (* last because poll and opaque overrides inline and local *)
   let lam =
     add_poll_attribute lam loc attr
@@ -654,3 +674,10 @@ let add_function_attributes lam loc attr =
     add_opaque_attribute lam loc attr
   in
   lam
+
+let transl_param_attributes pat =
+      let attrs = pat.pat_attributes in
+      let unbox_param =
+        Option.is_some (find_attribute is_unboxed_attribute attrs)
+      in
+      { unbox_param }

lambda/translattribute.mli

@@ -66,3 +66,8 @@ val add_function_attributes
   -> Lambda.lambda
 
 val get_assume_zero_alloc : with_warnings:bool ->  Parsetree.attributes -> bool
+
+val transl_param_attributes
+  : Typedtree.pattern
+  -> Lambda.parameter_attribute
+

lambda/translcore.ml

@@ -963,6 +963,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
           stub = false;
           poll = Default_poll;
           tmc_candidate = false;
+          unbox_return = false;
           may_fuse_arity = false;
         } in
       let funcid = Ident.create_local ("probe_handler_" ^ name) in
@@ -1332,10 +1333,15 @@ and transl_curried_function ~scopes loc repr params body
               layout_of_sort fc_loc fc_arg_sort
         in
         let arg_mode = transl_alloc_mode fc_arg_mode in
+        let attributes =
+          match fc_cases with
+          | [ { c_lhs }] -> Translattribute.transl_param_attributes c_lhs
+          | [] | _ :: _ :: _ -> Lambda.default_param_attribute
+        in
         let param =
           { name = fc_param;
             layout = arg_layout;
-            attributes = Lambda.default_param_attribute;
+            attributes;
             mode = arg_mode;
           }
         in
@@ -1350,18 +1356,19 @@ and transl_curried_function ~scopes loc repr params body
     List.fold_right
       (fun fp (body, params) ->
         let { fp_param; fp_kind; fp_mode; fp_sort; fp_partial; fp_loc } = fp in
-        let arg_env, arg_type =
+        let arg_env, arg_type, attributes =
           match fp_kind with
-          | Tparam_pat pat -> pat.pat_env, pat.pat_type
-          | Tparam_optional_default (_, expr, _) ->
-              expr.exp_env, Predef.type_option expr.exp_type
+          | Tparam_pat pat ->
+              pat.pat_env, pat.pat_type, Translattribute.transl_param_attributes pat
+          | Tparam_optional_default (pat, expr, _) ->
+              expr.exp_env, Predef.type_option expr.exp_type, Translattribute.transl_param_attributes pat
         in
         let arg_layout = layout arg_env fp_loc fp_sort arg_type in
         let arg_mode = transl_alloc_mode fp_mode in
         let param =
           { name = fp_param;
             layout = arg_layout;
-            attributes = Lambda.default_param_attribute;
+            attributes;
             mode = arg_mode;
           }
         in

lambda/translmod.ml

@@ -603,6 +603,7 @@ let rec compile_functor ~scopes mexp coercion root_path loc =
       stub = false;
       tmc_candidate = false;
       may_fuse_arity = true;
+      unbox_return = false;
     }
     ~loc
     ~mode:alloc_heap

middle_end/closure/closure.ml

@@ -1546,7 +1546,7 @@ and close_functions { backend; fenv; cenv; mutable_vars; kinds; catch_env } fun_
                fun_region = region;
                fun_poll = attr.poll } in
             let dbg = Debuginfo.from_location loc in
-            (id, List.map (fun (p : Lambda.lparam) -> let No_attributes = p.attributes in (p.name, p.layout, p.mode)) params,
+            (id, List.map (fun (p : Lambda.lparam) -> (p.name, p.layout, p.mode)) params,
              return, body, mode, fundesc, dbg)
         | (_, _) -> fatal_error "Closure.close_functions")
       fun_defs in

middle_end/flambda/closure_conversion.ml

@@ -226,7 +226,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
     (* CR-soon mshinwell: some of this is now very similar to the let rec case
        below *)
     let set_of_closures_var = Variable.create Names.set_of_closures in
-    let params = List.map (fun (p : Lambda.lparam) -> let No_attributes = p.attributes in (p.name, p.layout)) params in
+    let params = List.map (fun (p : Lambda.lparam) -> (p.name, p.layout)) params in
     let set_of_closures =
       let decl =
         Function_decl.create ~let_rec_ident:None ~closure_bound_var ~kind ~mode
@@ -279,7 +279,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
             let closure_bound_var =
               Variable.create_with_same_name_as_ident let_rec_ident
             in
-            let params = List.map (fun (p : Lambda.lparam) -> let No_attributes = p.attributes in (p.name, p.layout)) params in
+            let params = List.map (fun (p : Lambda.lparam) -> (p.name, p.layout)) params in
             let function_declaration =
               Function_decl.create ~let_rec_ident:(Some let_rec_ident)
                 ~closure_bound_var ~kind ~mode ~region
@@ -699,7 +699,7 @@ and close_let_bound_expression t ?let_rec_ident let_bound_var env
     (* Ensure that [let] and [let rec]-bound functions have appropriate
        names. *)
     let closure_bound_var = Variable.rename let_bound_var in
-    let params = List.map (fun (p : Lambda.lparam) -> let No_attributes = p.attributes in (p.name, p.layout)) params in
+    let params = List.map (fun (p : Lambda.lparam) -> (p.name, p.layout)) params in
     let decl =
       Function_decl.create ~let_rec_ident ~closure_bound_var ~kind ~mode ~region
         ~params ~body ~attr ~loc ~return_layout:return

utils/warnings.ml

@@ -112,6 +112,7 @@ type t =
   | Unnecessarily_partial_tuple_pattern     (* 189 *)
   | Probe_name_too_long of string           (* 190 *)
   | Unchecked_property_attribute of string  (* 199 *)
+  | Unboxing_impossible                     (* 210 *)
 
 (* If you remove a warning, leave a hole in the numbering.  NEVER change
    the numbers of existing warnings.
@@ -196,12 +197,13 @@ let number = function
   | Unnecessarily_partial_tuple_pattern -> 189
   | Probe_name_too_long _ -> 190
   | Unchecked_property_attribute _ -> 199
+  | Unboxing_impossible -> 210
 ;;
 (* DO NOT REMOVE the ;; above: it is used by
    the testsuite/ests/warnings/mnemonics.mll test to determine where
    the  definition of the number function above ends *)
 
-let last_warning_number = 199
+let last_warning_number = 210
 ;;
 
 type description =
@@ -555,6 +557,10 @@ let descriptions = [
     description = "A property of a function that was \
                    optimized away cannot be checked.";
     since = since 4 14 };
+  { number = 210;
+    names = ["unboxing-impossible"];
+    description = "The parameter or return value corresponding @unboxed attribute cannot be unboxed.";
+    since = since 4 14 };
 ]
 
 let name_to_number =
@@ -1170,6 +1176,10 @@ let message = function
       You can try to mark this function as [@inline never] \n\
       or move the attribute to the relevant callers of this function."
       property
+  | Unboxing_impossible ->
+      Printf.sprintf
+        "This [@unboxed] attribute cannot be used.\n\
+         The type of this value does not allow unboxing."
 ;;
 
 let nerrors = ref 0

utils/warnings.mli

@@ -116,6 +116,7 @@ type t =
   | Unnecessarily_partial_tuple_pattern     (* 189 *)
   | Probe_name_too_long of string           (* 190 *)
   | Unchecked_property_attribute of string  (* 199 *)
+  | Unboxing_impossible                     (* 210 *)
 
 type alert = {kind:string; message:string; def:loc; use:loc}