Skip to content

Commit 4a795cb

Browse files
stedolanstedolan
authored
Flambda support (oxcaml#49)
* Flambda support for local allocations and regions * Optimise away redundant Regions in Flambda
1 parent a39126a commit 4a795cb

40 files changed

Lines changed: 503 additions & 166 deletions

‎lambda/lambda.ml‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1000,3 +1000,21 @@ let max_arity () =
10001000

10011001
let reset () =
10021002
raise_count := 0
1003+
1004+
let join_mode a b =
1005+
match a, b with
1006+
| Alloc_local, _ | _, Alloc_local -> Alloc_local
1007+
| Alloc_heap, Alloc_heap -> Alloc_heap
1008+
1009+
let sub_mode a b =
1010+
match a, b with
1011+
| Alloc_heap, _ -> true
1012+
| _, Alloc_local -> true
1013+
| Alloc_local, Alloc_heap -> false
1014+
1015+
let eq_mode a b =
1016+
match a, b with
1017+
| Alloc_heap, Alloc_heap -> true
1018+
| Alloc_local, Alloc_local -> true
1019+
| Alloc_heap, Alloc_local -> false
1020+
| Alloc_local, Alloc_heap -> false

‎lambda/lambda.mli‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,10 @@ val max_arity : unit -> int
451451
This is unlimited ([max_int]) for bytecode, but limited
452452
(currently to 126) for native code. *)
453453

454+
val join_mode : alloc_mode -> alloc_mode -> alloc_mode
455+
val sub_mode : alloc_mode -> alloc_mode -> bool
456+
val eq_mode : alloc_mode -> alloc_mode -> bool
457+
454458
(***********************)
455459
(* For static failures *)
456460
(***********************)

‎lambda/translcore.ml‎

Lines changed: 2 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -93,17 +93,6 @@ let transl_value_mode mode =
9393
let alloc_mode = Types.Value_mode.regional_to_global_alloc mode in
9494
transl_alloc_mode alloc_mode
9595

96-
let join_mode a b =
97-
match a, b with
98-
| Alloc_local, _ | _, Alloc_local -> Alloc_local
99-
| Alloc_heap, Alloc_heap -> Alloc_heap
100-
101-
let sub_mode a b =
102-
match a, b with
103-
| Alloc_heap, _ -> true
104-
| _, Alloc_local -> true
105-
| Alloc_local, Alloc_heap -> false
106-
10796
let transl_apply_position position =
10897
match position with
10998
| Nontail -> Apply_nontail
@@ -941,10 +930,11 @@ and transl_tupled_function
941930
~scopes ~arity ~mode loc return
942931
repr partial (param:Ident.t) cases =
943932
match cases with
944-
| {c_lhs={pat_desc = Tpat_tuple pl}} :: _
933+
| {c_lhs={pat_desc = Tpat_tuple pl; pat_mode }} :: _
945934
when !Clflags.native_code
946935
&& arity = 1
947936
&& mode = Alloc_heap
937+
&& transl_value_mode pat_mode = Alloc_heap
948938
&& List.length pl <= (Lambda.max_arity ()) ->
949939
begin try
950940
let size = List.length pl in

‎middle_end/flambda/augment_specialised_args.ml‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -468,6 +468,7 @@ module Make (T : S) = struct
468468
spec_args_bound_in_the_wrapper;
469469
kind = Direct (Closure_id.wrap new_fun_var);
470470
dbg = Debuginfo.none;
471+
position = Apply_nontail;
471472
inline = Default_inline;
472473
specialise = Default_specialise;
473474
}
@@ -527,9 +528,15 @@ module Make (T : S) = struct
527528
for_one_function.existing_specialised_args
528529
Variable.Map.empty
529530
in
531+
let alloc_mode =
532+
(* Wrapper closes over no more values than the original function,
533+
so can share the same alloc mode *)
534+
function_decl.alloc_mode
535+
in
530536
let new_function_decl =
531537
Flambda.create_function_declaration
532538
~params:wrapper_params
539+
~alloc_mode
533540
~body:wrapper_body
534541
~stub:true
535542
~dbg:Debuginfo.none
@@ -607,8 +614,12 @@ module Make (T : S) = struct
607614
Variable.Set.elements (Variable.Map.keys
608615
for_one_function.new_inner_to_new_outer_vars)
609616
in
617+
let last_mode =
618+
List.fold_left (fun _mode p -> Parameter.alloc_mode p)
619+
function_decl.alloc_mode function_decl.params
620+
in
610621
let new_params =
611-
List.map Parameter.wrap new_params
622+
List.map (fun p -> Parameter.wrap p last_mode) new_params
612623
in
613624
function_decl.params @ new_params
614625
in
@@ -618,6 +629,7 @@ module Make (T : S) = struct
618629
let rewritten_function_decl =
619630
Flambda.create_function_declaration
620631
~params:all_params
632+
~alloc_mode:function_decl.alloc_mode
621633
~body:function_decl.body
622634
~stub:function_decl.stub
623635
~dbg:function_decl.dbg

‎middle_end/flambda/build_export_info.ml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,10 @@ let rec approx_of_expr (env : Env.t) (flam : Flambda.t) : Export_info.approx =
258258
Closure_id.Map.find closure_id results
259259
| _ -> Value_unknown
260260
end
261+
| Region body ->
262+
approx_of_expr env body
263+
| Tail body ->
264+
approx_of_expr env body
261265
| Assign _ -> Value_id (Env.new_unit_descr env)
262266
| For _ -> Value_id (Env.new_unit_descr env)
263267
| While _ -> Value_id (Env.new_unit_descr env)

‎middle_end/flambda/closure_conversion.ml‎

Lines changed: 32 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ let tupled_function_call_stub original_params unboxed_version ~closure_bound_var
8787
redundancy here (func is also unboxed_version) *)
8888
kind = Direct (Closure_id.wrap unboxed_version);
8989
dbg = Debuginfo.none;
90+
position = Apply_nontail;
9091
inline = Default_inline;
9192
specialise = Default_specialise;
9293
})
@@ -99,8 +100,10 @@ let tupled_function_call_stub original_params unboxed_version ~closure_bound_var
99100
pos + 1, Flambda.create_let param lam body)
100101
(0, call) params
101102
in
102-
let tuple_param = Parameter.wrap tuple_param_var in
103-
Flambda.create_function_declaration ~params:[tuple_param]
103+
(* Tupled functions are always Alloc_heap. See translcore.ml *)
104+
let alloc_mode = Lambda.Alloc_heap in
105+
let tuple_param = Parameter.wrap tuple_param_var alloc_mode in
106+
Flambda.create_function_declaration ~params:[tuple_param] ~alloc_mode
104107
~body ~stub:true ~dbg:Debuginfo.none ~inline:Default_inline
105108
~specialise:Default_specialise ~is_a_functor:false
106109
~closure_origin:(Closure_origin.create (Closure_id.wrap closure_bound_var))
@@ -207,15 +210,15 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
207210
initial_value = var;
208211
body;
209212
contents_kind = block_kind })
210-
| Lfunction { kind; params; body; attr; loc; (* FIXME mode *) } ->
213+
| Lfunction { kind; params; body; attr; loc; mode } ->
211214
let name = Names.anon_fn_with_loc loc in
212215
let closure_bound_var = Variable.create name in
213216
(* CR-soon mshinwell: some of this is now very similar to the let rec case
214217
below *)
215218
let set_of_closures_var = Variable.create Names.set_of_closures in
216219
let set_of_closures =
217220
let decl =
218-
Function_decl.create ~let_rec_ident:None ~closure_bound_var ~kind
221+
Function_decl.create ~let_rec_ident:None ~closure_bound_var ~kind ~mode
219222
~params:(List.map fst params) ~body ~attr ~loc
220223
in
221224
close_functions t env (Function_decls.create [decl])
@@ -228,7 +231,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
228231
Flambda.create_let set_of_closures_var set_of_closures
229232
(name_expr (Project_closure (project_closure)) ~name)
230233
| Lapply { ap_func; ap_args; ap_loc;
231-
ap_tailcall = _; ap_inlined; ap_specialised; } ->
234+
ap_tailcall = _; ap_inlined; ap_specialised; ap_position } ->
232235
Lift_code.lifting_helper (close_list t env ap_args)
233236
~evaluation_order:`Right_to_left
234237
~name:Names.apply_arg
@@ -241,6 +244,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
241244
args;
242245
kind = Indirect;
243246
dbg = Debuginfo.from_location ap_loc;
247+
position = ap_position;
244248
inline = ap_inlined;
245249
specialise = ap_specialised;
246250
})))
@@ -255,13 +259,13 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
255259
will be named after the corresponding identifier in the [let rec]. *)
256260
List.map (function
257261
| (let_rec_ident,
258-
Lambda.Lfunction { kind; params; body; attr; loc (* FIXME mode *) }) ->
262+
Lambda.Lfunction { kind; params; body; attr; loc; mode }) ->
259263
let closure_bound_var =
260264
Variable.create_with_same_name_as_ident let_rec_ident
261265
in
262266
let function_declaration =
263267
Function_decl.create ~let_rec_ident:(Some let_rec_ident)
264-
~closure_bound_var ~kind ~params:(List.map fst params) ~body
268+
~closure_bound_var ~kind ~mode ~params:(List.map fst params) ~body
265269
~attr ~loc
266270
in
267271
Some function_declaration
@@ -312,7 +316,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
312316
in
313317
Let_rec (defs, close t env body)
314318
end
315-
| Lsend (kind, meth, obj, args, _FIXME, loc) ->
319+
| Lsend (kind, meth, obj, args, position, loc) ->
316320
let meth_var = Variable.create Names.meth in
317321
let obj_var = Variable.create Names.obj in
318322
let dbg = Debuginfo.from_location loc in
@@ -322,7 +326,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
322326
~evaluation_order:`Right_to_left
323327
~name:Names.send_arg
324328
~create_body:(fun args ->
325-
Send { kind; meth = meth_var; obj = obj_var; args; dbg; })))
329+
Send { kind; meth = meth_var; obj = obj_var; args; dbg; position })))
326330
| Lprim ((Pdivint Safe | Pmodint Safe
327331
| Pdivbint { is_safe = Safe } | Pmodbint { is_safe = Safe }) as prim,
328332
[arg1; arg2], loc)
@@ -568,8 +572,8 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
568572
or by completely removing it (replacing by unit). *)
569573
Misc.fatal_error "[Lifused] should have been removed by \
570574
[Simplif.simplify_lets]"
571-
| Lregion _ ->
572-
Misc.fatal_error "FIXME: Lregion unimplemented in Flambda"
575+
| Lregion body ->
576+
Region (close t env body)
573577

574578
(** Perform closure conversion on a set of function declarations, returning a
575579
set of closures. (The set will often only contain a single function;
@@ -600,22 +604,34 @@ and close_functions t external_env function_declarations : Flambda.named =
600604
not marked as stub but certainly should *)
601605
let stub = Function_decl.stub decl in
602606
let param_vars = List.map (Env.find_var closure_env) params in
603-
let params = List.map Parameter.wrap param_vars in
607+
let nheap =
608+
match Function_decl.mode decl, Function_decl.kind decl with
609+
| _, Curried {nlocal} -> List.length params - nlocal
610+
| Alloc_heap, Tupled -> List.length params
611+
| Alloc_local, Tupled -> 0
612+
in
613+
let params = List.mapi (fun i v ->
614+
let alloc_mode : Lambda.alloc_mode =
615+
if i < nheap then Alloc_heap else Alloc_local in
616+
Parameter.wrap v alloc_mode) param_vars
617+
in
604618
let closure_bound_var = Function_decl.closure_bound_var decl in
605619
let unboxed_version = Variable.rename closure_bound_var in
606620
let body = close t closure_env body in
607621
let closure_origin =
608622
Closure_origin.create (Closure_id.wrap unboxed_version)
609623
in
610624
let fun_decl =
611-
Flambda.create_function_declaration ~params ~body ~stub ~dbg
625+
Flambda.create_function_declaration
626+
~params ~alloc_mode:(Function_decl.mode decl)
627+
~body ~stub ~dbg
612628
~inline:(Function_decl.inline decl)
613629
~specialise:(Function_decl.specialise decl)
614630
~is_a_functor:(Function_decl.is_a_functor decl)
615631
~closure_origin
616632
in
617633
match Function_decl.kind decl with
618-
| Curried _ (* FIXME nlocal *) ->
634+
| Curried _ ->
619635
Variable.Map.add closure_bound_var fun_decl map
620636
| Tupled ->
621637
let unboxed_version = Variable.rename closure_bound_var in
@@ -661,12 +677,12 @@ and close_list t sb l = List.map (close t sb) l
661677
and close_let_bound_expression t ?let_rec_ident let_bound_var env
662678
(lam : Lambda.lambda) : Flambda.named =
663679
match lam with
664-
| Lfunction { kind; params; body; attr; loc; } ->
680+
| Lfunction { kind; params; body; attr; loc; mode } ->
665681
(* Ensure that [let] and [let rec]-bound functions have appropriate
666682
names. *)
667683
let closure_bound_var = Variable.rename let_bound_var in
668684
let decl =
669-
Function_decl.create ~let_rec_ident ~closure_bound_var ~kind
685+
Function_decl.create ~let_rec_ident ~closure_bound_var ~kind ~mode
670686
~params:(List.map fst params) ~body ~attr ~loc
671687
in
672688
let set_of_closures_var = Variable.rename let_bound_var in

‎middle_end/flambda/closure_conversion_aux.ml‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,14 +87,15 @@ module Function_decls = struct
8787
let_rec_ident : Ident.t;
8888
closure_bound_var : Variable.t;
8989
kind : Lambda.function_kind;
90+
mode : Lambda.alloc_mode;
9091
params : Ident.t list;
9192
body : Lambda.lambda;
9293
free_idents_of_body : Ident.Set.t;
9394
attr : Lambda.function_attribute;
9495
loc : Lambda.scoped_location
9596
}
9697

97-
let create ~let_rec_ident ~closure_bound_var ~kind ~params ~body
98+
let create ~let_rec_ident ~closure_bound_var ~kind ~mode ~params ~body
9899
~attr ~loc =
99100
let let_rec_ident =
100101
match let_rec_ident with
@@ -104,6 +105,7 @@ module Function_decls = struct
104105
{ let_rec_ident;
105106
closure_bound_var;
106107
kind;
108+
mode;
107109
params;
108110
body;
109111
free_idents_of_body = Lambda.free_variables body;
@@ -114,6 +116,7 @@ module Function_decls = struct
114116
let let_rec_ident t = t.let_rec_ident
115117
let closure_bound_var t = t.closure_bound_var
116118
let kind t = t.kind
119+
let mode t = t.mode
117120
let params t = t.params
118121
let body t = t.body
119122
let free_idents t = t.free_idents_of_body
@@ -160,7 +163,7 @@ module Function_decls = struct
160163
(all_params function_decls))
161164
(let_rec_idents function_decls)
162165

163-
let create function_decls =
166+
let create (function_decls : Function_decl.t list) =
164167
{ function_decls;
165168
all_free_idents = all_free_idents function_decls;
166169
}

‎middle_end/flambda/closure_conversion_aux.mli‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ module Function_decls : sig
5656
: let_rec_ident:Ident.t option
5757
-> closure_bound_var:Variable.t
5858
-> kind:Lambda.function_kind
59+
-> mode:Lambda.alloc_mode
5960
-> params:Ident.t list
6061
-> body:Lambda.lambda
6162
-> attr:Lambda.function_attribute
@@ -65,6 +66,7 @@ module Function_decls : sig
6566
val let_rec_ident : t -> Ident.t
6667
val closure_bound_var : t -> Variable.t
6768
val kind : t -> Lambda.function_kind
69+
val mode : t -> Lambda.alloc_mode
6870
val params : t -> Ident.t list
6971
val body : t -> Lambda.lambda
7072
val inline : t -> Lambda.inline_attribute

‎middle_end/flambda/effect_analysis.ml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,10 @@ let rec no_effects (flam : Flambda.t) =
4545
(* If there is a [raise] in [body], the whole [Try_with] may have an
4646
effect, so there is no need to test the handler. *)
4747
no_effects body
48+
| Region body ->
49+
no_effects body
50+
| Tail body ->
51+
no_effects body
4852
| While _ | For _ | Apply _ | Send _ | Assign _ | Static_raise _ -> false
4953
| Proved_unreachable -> true
5054

‎middle_end/flambda/export_info_for_pack.ml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,8 @@ and import_function_declarations_for_pack_aux units pack
139139
let funs =
140140
Variable.Map.map
141141
(fun (function_decl : Flambda.function_declaration) ->
142-
Flambda.create_function_declaration ~params:function_decl.params
142+
Flambda.create_function_declaration
143+
~params:function_decl.params ~alloc_mode:function_decl.alloc_mode
143144
~body:(import_code_for_pack units pack function_decl.body)
144145
~stub:function_decl.stub ~dbg:function_decl.dbg
145146
~inline:function_decl.inline

0 commit comments

Comments
 (0)