forked from ocaml/ocaml
/
extract_projections.ml
190 lines (184 loc) · 8.29 KB
/
extract_projections.ml
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
(**************************************************************************)
(* *)
(* OCaml *)
(* *)
(* Pierre Chambart, OCamlPro *)
(* Mark Shinwell and Leo White, Jane Street Europe *)
(* *)
(* Copyright 2013--2016 OCamlPro SAS *)
(* Copyright 2014--2016 Jane Street Group LLC *)
(* *)
(* 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. *)
(* *)
(**************************************************************************)
[@@@ocaml.warning "+a-4-9-30-40-41-42"]
open! Int_replace_polymorphic_compare
module A = Simple_value_approx
module E = Inline_and_simplify_aux.Env
(* CR-soon pchambart: should we restrict only to cases
when the field is aliased to a variable outside
of the closure (i.e. when we can certainly remove
the allocation of the block) ?
Note that this may prevent cases with imbricated
closures from benefiting from this transformations.
mshinwell: What word was "imbricated" supposed to be?
(The code this referred to has been deleted, but the same thing is
probably still happening).
*)
let known_valid_projections ~env ~projections ~which_variables =
Projection.Set.filter (fun projection ->
let from = Projection.projecting_from projection in
let outer_var =
match Variable.Map.find from which_variables with
| exception Not_found -> assert false
| (outer_var : Flambda.specialised_to) ->
Freshening.apply_variable (E.freshening env) outer_var.var
in
let approx = E.find_exn env outer_var in
match projection with
| Project_var project_var ->
begin match A.check_approx_for_closure approx with
| Ok (_value_closure, _approx_var, _approx_sym,
value_set_of_closures) ->
Var_within_closure.Map.mem project_var.var
value_set_of_closures.bound_vars
| Wrong -> false
end
| Project_closure project_closure ->
begin match A.strict_check_approx_for_set_of_closures approx with
| Ok (_var, value_set_of_closures) ->
Variable.Set.mem (Closure_id.unwrap project_closure.closure_id)
(Variable.Map.keys value_set_of_closures.function_decls.funs)
| Wrong -> false
end
| Move_within_set_of_closures move ->
begin match A.check_approx_for_closure approx with
| Ok (value_closure, _approx_var, _approx_sym,
_value_set_of_closures) ->
(* We could check that [move.move_to] is in [value_set_of_closures],
but this is unnecessary, since [Closure_id]s are unique. *)
Closure_id.equal value_closure.closure_id move.start_from
| Wrong -> false
end
| Field (field_index, _) ->
match A.check_approx_for_block approx with
| Wrong -> false
| Ok (_tag, fields) ->
field_index >= 0 && field_index < Array.length fields)
projections
let rec analyse_expr ~which_variables expr =
let projections = ref Projection.Set.empty in
let used_which_variables = ref Variable.Set.empty in
let check_free_variable var =
if Variable.Map.mem var which_variables then begin
used_which_variables := Variable.Set.add var !used_which_variables
end
in
let for_expr (expr : Flambda.expr) =
match expr with
| Var var
| Let_mutable { initial_value = var } ->
check_free_variable var
(* CR-soon mshinwell: We don't handle [Apply] for the moment to
avoid disabling unboxing optimizations whenever we see a recursive
call. We should improve this analysis. Leo says this can be
done by a similar thing to the unused argument analysis. *)
| Apply _ -> ()
| Send { meth; obj; args; _ } ->
check_free_variable meth;
check_free_variable obj;
List.iter check_free_variable args
| Assign { new_value; _ } ->
check_free_variable new_value
| If_then_else (var, _, _)
| Switch (var, _)
| String_switch (var, _, _) ->
check_free_variable var
| Static_raise (_, args) ->
List.iter check_free_variable args
| For { from_value; to_value; _ } ->
check_free_variable from_value;
check_free_variable to_value
| Let _ | Let_rec _ | Static_catch _ | While _ | Try_with _
| Proved_unreachable -> ()
in
let for_named (named : Flambda.named) =
match named with
| Project_var project_var
when Variable.Map.mem project_var.closure which_variables ->
projections :=
Projection.Set.add (Project_var project_var) !projections
| Project_closure project_closure
when Variable.Map.mem project_closure.set_of_closures
which_variables ->
projections :=
Projection.Set.add (Project_closure project_closure) !projections
| Move_within_set_of_closures move
when Variable.Map.mem move.closure which_variables ->
projections :=
Projection.Set.add (Move_within_set_of_closures move) !projections
| Prim (Pfield field_index, [var], _dbg)
when Variable.Map.mem var which_variables ->
projections :=
Projection.Set.add (Field (field_index, var)) !projections
| Set_of_closures set_of_closures ->
let aliasing_free_vars =
Variable.Map.filter (fun _ (spec_to : Flambda.specialised_to) ->
Variable.Map.mem spec_to.var which_variables)
set_of_closures.free_vars
in
let aliasing_specialised_args =
Variable.Map.filter (fun _ (spec_to : Flambda.specialised_to) ->
Variable.Map.mem spec_to.var which_variables)
set_of_closures.specialised_args
in
let aliasing_vars =
Variable.Map.disjoint_union
aliasing_free_vars aliasing_specialised_args
in
if not (Variable.Map.is_empty aliasing_vars) then begin
Variable.Map.iter (fun _ (fun_decl : Flambda.function_declaration) ->
(* We ignore projections from within nested sets of closures. *)
let _, used =
analyse_expr fun_decl.body ~which_variables:aliasing_vars
in
Variable.Set.iter (fun var ->
match Variable.Map.find var aliasing_vars with
| exception Not_found -> assert false
| spec_to -> check_free_variable spec_to.var)
used)
set_of_closures.function_decls.funs
end
| Prim (_, vars, _) ->
List.iter check_free_variable vars
| Symbol _ | Const _ | Allocated_const _ | Read_mutable _
| Read_symbol_field _ | Project_var _ | Project_closure _
| Move_within_set_of_closures _
| Expr _ -> ()
in
Flambda_iterators.iter_toplevel for_expr for_named expr;
let projections = !projections in
let used_which_variables = !used_which_variables in
projections, used_which_variables
let from_function_decl ~env ~which_variables
~(function_decl : Flambda.function_declaration) =
let projections, used_which_variables =
analyse_expr ~which_variables function_decl.body
in
(* We must use approximation information to determine which projections
are actually valid in the current environment, other we might lift
expressions too far. *)
let projections =
known_valid_projections ~env ~projections ~which_variables
in
(* Don't extract projections whose [projecting_from] variable is also
used boxed. We could in the future consider being more sophisticated
about this based on the uses in the body, but given we are not doing
that yet, it seems safest in performance terms not to (e.g.) unbox a
specialised argument whose boxed version is used. *)
Projection.Set.filter (fun projection ->
let projecting_from = Projection.projecting_from projection in
not (Variable.Set.mem projecting_from used_which_variables))
projections