Skip to content

Refactor primitive handling in Lambda_to_flambda #1382

New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 27 additions & 18 deletions middle_end/flambda2/from_lambda/closure_conversion.ml
Original file line number Diff line number Diff line change
Expand Up @@ -601,16 +601,40 @@ let close_exn_continuation acc env (exn_continuation : IR.exn_continuation) =
Exn_continuation.create ~exn_handler:exn_continuation.exn_handler
~extra_args )

let close_raise acc env ~raise_kind ~arg loc exn_continuation =
let acc, exn_cont = close_exn_continuation acc env exn_continuation in
let exn_handler = Exn_continuation.exn_handler exn_cont in
let acc, arg = find_simple acc env arg in
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since we already do let acc, args = find_simples acc env args in close_primitive, I think this might duplicate the definition of the constant corresponding to the argument if it is one.

let args =
(* CR mshinwell: Share with [Lambda_to_flambda_primitives_helpers] *)
let extra_args =
List.map
(fun (simple, _kind) -> simple)
(Exn_continuation.extra_args exn_cont)
in
arg :: extra_args
in
let raise_kind = Some (Trap_action.Raise_kind.from_lambda raise_kind) in
let trap_action = Trap_action.Pop { exn_handler; raise_kind } in
let dbg = Debuginfo.from_location loc in
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any reason to do that rather than give dbg as an argument (which is already known in one of the use sites) ?

let acc, apply_cont =
Apply_cont_with_acc.create acc ~trap_action exn_handler ~args ~dbg
in
(* Since raising of an exception doesn't terminate, we don't call [k]. *)
Expr_with_acc.create_apply_cont acc apply_cont

let close_primitive acc env ~let_bound_var named (prim : Lambda.primitive) ~args
loc (exn_continuation : IR.exn_continuation option) ~current_region
(k : Acc.t -> Named.t option -> Expr_with_acc.t) : Expr_with_acc.t =
let orig_exn_continuation = exn_continuation in
let acc, exn_continuation =
match exn_continuation with
| None -> acc, None
| Some exn_continuation ->
let acc, cont = close_exn_continuation acc env exn_continuation in
acc, Some cont
in
let orig_args = args in
let acc, args = find_simples acc env args in
let dbg = Debuginfo.from_location loc in
match prim, args with
Expand Down Expand Up @@ -642,29 +666,14 @@ let close_primitive acc env ~let_bound_var named (prim : Lambda.primitive) ~args
k acc (Some named)
| Praise raise_kind, [_] ->
let exn_continuation =
match exn_continuation with
match orig_exn_continuation with
| None ->
Misc.fatal_errorf "Praise is missing exception continuation: %a"
IR.print_named named
| Some exn_continuation -> exn_continuation
in
let exn_handler = Exn_continuation.exn_handler exn_continuation in
let args =
(* CR mshinwell: Share with [Lambda_to_flambda_primitives_helpers] *)
let extra_args =
List.map
(fun (simple, _kind) -> simple)
(Exn_continuation.extra_args exn_continuation)
in
args @ extra_args
in
let raise_kind = Some (Trap_action.Raise_kind.from_lambda raise_kind) in
let trap_action = Trap_action.Pop { exn_handler; raise_kind } in
let acc, apply_cont =
Apply_cont_with_acc.create acc ~trap_action exn_handler ~args ~dbg
in
(* Since raising of an exception doesn't terminate, we don't call [k]. *)
Expr_with_acc.create_apply_cont acc apply_cont
close_raise acc env ~raise_kind ~arg:(List.hd orig_args) loc
exn_continuation
| (Pmakeblock _ | Pmakefloatblock _ | Pmakearray _), [] ->
(* Special case for liftable empty block or array *)
let acc, sym =
Expand Down
9 changes: 9 additions & 0 deletions middle_end/flambda2/from_lambda/closure_conversion.mli
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,15 @@ val close_switch :
IR.switch ->
Expr_with_acc.t

val close_raise :
Acc.t ->
Env.t ->
raise_kind:Lambda.raise_kind ->
arg:IR.simple ->
Lambda.scoped_location ->
IR.exn_continuation ->
Expr_with_acc.t

type 'a close_program_metadata =
| Normal : [`Normal] close_program_metadata
| Classic :
Expand Down
141 changes: 29 additions & 112 deletions middle_end/flambda2/from_lambda/lambda_to_flambda.ml
Original file line number Diff line number Diff line change
Expand Up @@ -985,94 +985,6 @@ let primitive_can_raise (prim : Lambda.primitive) =
| Punbox_int _ | Pbox_int _ ->
false

let primitive_result_kind (prim : Lambda.primitive) :
Flambda_kind.With_subkind.t =
match prim with
| Pccall { prim_native_repr_res = _, Untagged_int; _ } ->
Flambda_kind.With_subkind.tagged_immediate
| Pccall { prim_native_repr_res = _, Unboxed_float; _ }
| Pfloatofint _ | Pnegfloat _ | Pabsfloat _ | Paddfloat _ | Psubfloat _
| Pmulfloat _ | Pdivfloat _ | Pfloatfield _
| Parrayrefs Pfloatarray
| Parrayrefu Pfloatarray
| Pbigarrayref (_, _, (Pbigarray_float32 | Pbigarray_float64), _) ->
Flambda_kind.With_subkind.boxed_float
| Pccall { prim_native_repr_res = _, Unboxed_integer Pnativeint; _ }
| Pbigarrayref (_, _, Pbigarray_native_int, _) ->
Flambda_kind.With_subkind.boxed_nativeint
| Pccall { prim_native_repr_res = _, Unboxed_integer Pint32; _ }
| Pstring_load_32 _ | Pbytes_load_32 _ | Pbigstring_load_32 _
| Pbigarrayref (_, _, Pbigarray_int32, _) ->
Flambda_kind.With_subkind.boxed_int32
| Pccall { prim_native_repr_res = _, Unboxed_integer Pint64; _ }
| Pstring_load_64 _ | Pbytes_load_64 _ | Pbigstring_load_64 _
| Pbigarrayref (_, _, Pbigarray_int64, _) ->
Flambda_kind.With_subkind.boxed_int64
| Pnegint | Paddint | Psubint | Pmulint | Pandint | Porint | Pxorint | Plslint
| Plsrint | Pasrint | Pmodint _ | Pdivint _ | Pignore | Psequand | Psequor
| Pnot | Pbytesrefs | Pstringrefs | Pbytessets | Pstring_load_16 _
| Pbytes_load_16 _ | Pbigstring_load_16 _ | Pbytes_set_16 _ | Pbytes_set_32 _
| Pbytes_set_64 _ | Pbigstring_set_16 _ | Pbigstring_set_32 _
| Pbigstring_set_64 _ | Pintcomp _ | Pcompare_ints | Pcompare_floats
| Pcompare_bints _ | Pintoffloat | Pfloatcomp _ | Parraysets _
| Pbigarrayset _ | Psetfield _ | Psetfield_computed _ | Psetfloatfield _
| Pstringlength | Pstringrefu | Pbyteslength | Pbytesrefu | Pbytessetu
| Parraylength _ | Parraysetu _ | Pisint _ | Pbintcomp _ | Pintofbint _
| Pisout
| Parrayrefs Pintarray
| Parrayrefu Pintarray
| Pprobe_is_enabled _ | Pctconst _ | Pbswap16
| Pbigarrayref
( _,
_,
( Pbigarray_sint8 | Pbigarray_uint8 | Pbigarray_sint16
| Pbigarray_uint16 | Pbigarray_caml_int ),
_ ) ->
Flambda_kind.With_subkind.tagged_immediate
| Pdivbint { size = bi; _ }
| Pmodbint { size = bi; _ }
| Pandbint (bi, _)
| Porbint (bi, _)
| Pxorbint (bi, _)
| Plslbint (bi, _)
| Plsrbint (bi, _)
| Pasrbint (bi, _)
| Pnegbint (bi, _)
| Paddbint (bi, _)
| Psubbint (bi, _)
| Pmulbint (bi, _)
| Pbintofint (bi, _)
| Pcvtbint (_, bi, _)
| Pbbswap (bi, _)
| Pbox_int (bi, _) -> (
match bi with
| Pint32 -> Flambda_kind.With_subkind.boxed_int32
| Pint64 -> Flambda_kind.With_subkind.boxed_int64
| Pnativeint -> Flambda_kind.With_subkind.boxed_nativeint)
| Popaque layout | Pobj_magic layout ->
Flambda_kind.With_subkind.from_lambda layout
| Praise _ ->
(* CR ncourant: this should be bottom, but we don't have it *)
Flambda_kind.With_subkind.any_value
| Pccall { prim_native_repr_res = _, Same_as_ocaml_repr; _ }
| Parrayrefs (Pgenarray | Paddrarray)
| Parrayrefu (Pgenarray | Paddrarray)
| Pbytes_to_string | Pbytes_of_string | Parray_of_iarray | Parray_to_iarray
| Pgetglobal _ | Psetglobal _ | Pgetpredef _ | Pmakeblock _
| Pmakefloatblock _ | Pfield _ | Pfield_computed _ | Pduprecord _
| Poffsetint _ | Poffsetref _ | Pmakearray _ | Pduparray _ | Pbigarraydim _
| Pbigarrayref
(_, _, (Pbigarray_complex32 | Pbigarray_complex64 | Pbigarray_unknown), _)
| Pint_as_pointer | Pobj_dup ->
Flambda_kind.With_subkind.any_value
| Pbox_float _ -> Flambda_kind.With_subkind.boxed_float
| Punbox_float -> Flambda_kind.With_subkind.naked_float
| Punbox_int bi -> (
match bi with
| Pint32 -> Flambda_kind.With_subkind.naked_int32
| Pint64 -> Flambda_kind.With_subkind.naked_int64
| Pnativeint -> Flambda_kind.With_subkind.naked_nativeint)

type cps_continuation =
| Tail of Continuation.t
| Non_tail of (Acc.t -> Env.t -> CCenv.t -> IR.simple -> Expr_with_acc.t)
Expand Down Expand Up @@ -1256,30 +1168,35 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
~current_region:(Env.current_region env)
| Dissected lam -> cps acc env ccenv lam k k_exn)
| Lprim (prim, args, loc) -> (
match transform_primitive env prim args loc with
| Primitive (prim, args, loc) ->
let name = Printlambda.name_of_primitive prim in
let result_var = Ident.create_local name in
let exn_continuation : IR.exn_continuation option =
if primitive_can_raise prim
then
Some
{ exn_handler = k_exn;
extra_args = extra_args_for_exn_continuation env k_exn
}
else None
in
let current_region = Env.current_region env in
let dbg = Debuginfo.from_location loc in
cps_non_tail_list acc env ccenv args
(fun acc env ccenv args ->
let body acc ccenv = apply_cps_cont ~dbg k acc env ccenv result_var in
CC.close_let acc ccenv result_var Not_user_visible
(primitive_result_kind prim)
(Prim { prim; args; loc; exn_continuation; region = current_region })
~body)
k_exn
| Transformed lam -> cps acc env ccenv lam k k_exn)
match[@ocaml.warning "-fragile-match"] prim with
| Praise raise_kind -> (
match args with
| [_] ->
cps_non_tail_list acc env ccenv args
(fun acc _env ccenv args ->
let exn_continuation : IR.exn_continuation =
{ exn_handler = k_exn;
extra_args = extra_args_for_exn_continuation env k_exn
}
in
CC.close_raise acc ccenv ~raise_kind ~arg:(List.hd args) loc
exn_continuation)
k_exn
| [] | _ :: _ ->
Misc.fatal_errorf "Wrong number of arguments for Lraise: %a"
Printlambda.primitive prim)
| _ ->
let id = Ident.create_local "prim" in
let result_layout = L.primitive_result_layout prim in
(match result_layout with
| Pvalue _ | Punboxed_float | Punboxed_int _ -> ()
| Ptop | Pbottom ->
Misc.fatal_errorf "Invalid result layout %a for primitive %a"
Printlambda.layout result_layout Printlambda.primitive prim);
(* CR mshinwell: find a way of making these lets non-user-visible *)
cps acc env ccenv
(L.Llet (Strict, result_layout, id, lam, L.Lvar id))
k k_exn)
| Lswitch (scrutinee, switch, loc, kind) ->
maybe_insert_let_cont "switch_result" kind k acc env ccenv
(fun acc env ccenv k ->
Expand Down
7 changes: 5 additions & 2 deletions ocaml/lambda/lambda.ml
Original file line number Diff line number Diff line change
Expand Up @@ -1430,9 +1430,12 @@ let primitive_result_layout (p : primitive) =
| Paddfloat _ | Psubfloat _ | Pmulfloat _ | Pdivfloat _
| Pbox_float _ -> layout_float
| Punbox_float -> Punboxed_float
| Pccall _p ->
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we should add a case Pccall _ when not !Clflags.native_code -> layout_any_value ?

(* CR ncourant: use native_repr *)
| Pccall { prim_native_repr_res = _, Untagged_int; _} ->layout_int
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
| Pccall { prim_native_repr_res = _, Untagged_int; _} ->layout_int
| Pccall { prim_native_repr_res = _, Untagged_int; _} -> layout_int

| Pccall { prim_native_repr_res = _, Unboxed_float; _} -> layout_float
| Pccall { prim_native_repr_res = _, Same_as_ocaml_repr; _} ->
layout_any_value
| Pccall { prim_native_repr_res = _, Unboxed_integer bi; _} ->
layout_boxedint bi
| Praise _ -> layout_bottom
| Psequor | Psequand | Pnot
| Pnegint | Paddint | Psubint | Pmulint
Expand Down