Skip to content

Commit c1da9f4

Browse files
hhugohhugo
committed
Compiler: simplify branch
1 parent 9de6314 commit c1da9f4

File tree

4 files changed

+614
-721
lines changed

4 files changed

+614
-721
lines changed

‎CHANGES.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
* Compiler: support for OCaml 4.14.3+trunk (#1844)
88
* Compiler: add the `--empty-sourcemap` flag
99
* Compiler: improve debug/sourcemap location of closures (#1947)
10-
* Compiler: optimize compilation of switches
10+
* Compiler: optimize compilation of switches (#1921, #2057)
1111
* Compiler: evaluate statically more primitives (#1912, #1915, #1965, #1969)
1212
* Compiler: rewrote inlining pass (#1935, #2018, #2027)
1313
* Compiler: improve tailcall optimization (#1943)

‎compiler/lib/specialize.ml‎

Lines changed: 176 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -130,11 +130,107 @@ let f ~shape ~update_def p =
130130

131131
(***)
132132

133+
module Simple_block : sig
134+
type t
135+
136+
val hash : t -> int
137+
138+
val equal : t -> t -> bool
139+
140+
val make : block -> t
141+
end = struct
142+
type t = block
143+
144+
let subst_cont s (pc, arg) = pc, List.map arg ~f:s
145+
146+
let expr s e =
147+
match e with
148+
| Constant _ -> e
149+
| Apply { f; args; exact } -> Apply { f = s f; args = List.map args ~f:s; exact }
150+
| Block (n, a, k, mut) -> Block (n, Array.map a ~f:s, k, mut)
151+
| Field (x, n, typ) -> Field (s x, n, typ)
152+
| Closure (l, pc, loc) -> Closure (l, subst_cont s pc, loc)
153+
| Special _ -> e
154+
| Prim (p, l) ->
155+
Prim
156+
( p
157+
, List.map l ~f:(fun x ->
158+
match x with
159+
| Pv x -> Pv (s x)
160+
| Pc _ -> x) )
161+
162+
let instr s d i =
163+
match i with
164+
| Let (x, e) ->
165+
let x = d x in
166+
Let (x, expr s e)
167+
| Assign (x, y) -> Assign (s x, s y)
168+
| Set_field (x, n, typ, y) -> Set_field (s x, n, typ, s y)
169+
| Offset_ref (x, n) -> Offset_ref (s x, n)
170+
| Array_set (x, y, z) -> Array_set (s x, s y, s z)
171+
| Event _ -> Event Parse_info.zero
172+
173+
let instrs s d l = List.map l ~f:(fun i -> instr s d i)
174+
175+
let last s l =
176+
match l with
177+
| Stop -> l
178+
| Branch cont -> Branch (subst_cont s cont)
179+
| Pushtrap (cont1, x, cont2) -> Pushtrap (subst_cont s cont1, s x, subst_cont s cont2)
180+
| Return x -> Return (s x)
181+
| Raise (x, k) -> Raise (s x, k)
182+
| Cond (x, cont1, cont2) -> Cond (s x, subst_cont s cont1, subst_cont s cont2)
183+
| Switch (x, conts) -> Switch (s x, Array.map conts ~f:(fun cont -> subst_cont s cont))
184+
| Poptrap cont -> Poptrap (subst_cont s cont)
185+
186+
let block s d block =
187+
let params = List.map block.params ~f:s in
188+
let body = instrs s d block.body in
189+
let branch = last s block.branch in
190+
{ params; body; branch }
191+
192+
let make blk =
193+
let t = Var.Hashtbl.create 17 in
194+
let s x =
195+
match Var.Hashtbl.find_opt t x with
196+
| None -> x
197+
| Some x -> x
198+
in
199+
let d x =
200+
let v = Var.of_idx (-Var.Hashtbl.length t) in
201+
Var.Hashtbl.add t x v;
202+
v
203+
in
204+
block s d blk
205+
206+
let instr_equal a b =
207+
match a, b with
208+
| Event _, Event _ -> true
209+
| Event _, _ | _, Event _ -> false
210+
| a, b -> Poly.equal a b
211+
212+
let equal a b =
213+
List.equal ~eq:Var.equal a.params b.params
214+
&& List.equal ~eq:instr_equal a.body b.body
215+
&& Poly.equal a.branch b.branch
216+
217+
let hash (x : block) = Hashtbl.hash x
218+
end
219+
220+
module SBT = Hashtbl.Make (Simple_block)
221+
133222
(* For switches, at this point, we know that this it is sufficient to
134223
check the [pc]. *)
135224
let equal (pc, _) (pc', _) = pc = pc'
136225

137-
let find_outlier_index arr =
226+
type switch_to_cond =
227+
[ `All_equals
228+
| `Distinguished of int
229+
| `Splitted of int
230+
| `Splitted_shifted of int * int
231+
]
232+
233+
let find_outlier_index arr : [ switch_to_cond | `Many_cases ] =
138234
let len = Array.length arr in
139235
let rec find w i =
140236
if i >= len
@@ -159,6 +255,37 @@ let find_outlier_index arr =
159255
| `All_equals -> if j = i + 1 then `Distinguished i else `Splitted_shifted (i, j)
160256
| `Distinguished _ -> `Many_cases))
161257

258+
let optimize_switch_to_cond block x l (opt : switch_to_cond) =
259+
match opt with
260+
| `All_equals -> { block with branch = Branch l.(0) }
261+
| `Distinguished i ->
262+
let c = Var.fresh () in
263+
{ block with
264+
body =
265+
block.body @ [ Let (c, Prim (Eq, [ Pc (Int (Targetint.of_int_exn i)); Pv x ])) ]
266+
; branch = Cond (c, l.(i), l.((i + 1) mod Array.length l))
267+
}
268+
| `Splitted i ->
269+
let c = Var.fresh () in
270+
{ block with
271+
body =
272+
block.body @ [ Let (c, Prim (Lt, [ Pv x; Pc (Int (Targetint.of_int_exn i)) ])) ]
273+
; branch = Cond (c, l.(i - 1), l.(i))
274+
}
275+
| `Splitted_shifted (i, j) ->
276+
let shifted = Var.fresh () in
277+
let c = Var.fresh () in
278+
{ block with
279+
body =
280+
block.body
281+
@ [ Let
282+
( shifted
283+
, Prim (Extern "%int_sub", [ Pv x; Pc (Int (Targetint.of_int_exn i)) ]) )
284+
; Let (c, Prim (Ult, [ Pv shifted; Pc (Int (Targetint.of_int_exn (j - i))) ]))
285+
]
286+
; branch = Cond (c, l.(i), l.(j))
287+
}
288+
162289
let switches p =
163290
let previous_p = p in
164291
let t = Timer.make () in
@@ -171,63 +298,56 @@ let switches p =
171298
match block.branch with
172299
| Switch (x, l) -> (
173300
match find_outlier_index l with
174-
| `All_equals ->
175-
incr opt_count;
176-
Addr.Map.add pc { block with branch = Branch l.(0) } blocks
177-
| `Distinguished i ->
301+
| #switch_to_cond as opt ->
178302
incr opt_count;
179-
let block =
180-
let c = Var.fresh () in
181-
{ block with
182-
body =
183-
block.body
184-
@ [ Let
185-
(c, Prim (Eq, [ Pc (Int (Targetint.of_int_exn i)); Pv x ]))
186-
]
187-
; branch = Cond (c, l.(i), l.((i + 1) mod Array.length l))
188-
}
189-
in
303+
let block = optimize_switch_to_cond block x l opt in
190304
Addr.Map.add pc block blocks
191-
| `Splitted i ->
192-
incr opt_count;
193-
let block =
194-
let c = Var.fresh () in
195-
{ block with
196-
body =
197-
block.body
198-
@ [ Let
199-
(c, Prim (Lt, [ Pv x; Pc (Int (Targetint.of_int_exn i)) ]))
200-
]
201-
; branch = Cond (c, l.(i - 1), l.(i))
202-
}
203-
in
204-
Addr.Map.add pc block blocks
205-
| `Splitted_shifted (i, j) ->
206-
incr opt_count;
207-
let block =
208-
let shifted = Var.fresh () in
209-
let c = Var.fresh () in
210-
{ block with
211-
body =
212-
block.body
213-
@ [ Let
214-
( shifted
215-
, Prim
216-
( Extern "%int_sub"
217-
, [ Pv x; Pc (Int (Targetint.of_int_exn i)) ] ) )
218-
; Let
219-
( c
220-
, Prim
221-
( Ult
222-
, [ Pv shifted
223-
; Pc (Int (Targetint.of_int_exn (j - i)))
224-
] ) )
225-
]
226-
; branch = Cond (c, l.(i), l.(j))
227-
}
305+
| `Many_cases ->
306+
let t = SBT.create 0 in
307+
let rewrite = ref Addr.Set.empty in
308+
let l =
309+
Array.map l ~f:(fun ((pc, _) as cont) ->
310+
let block = Code.Addr.Map.find pc blocks in
311+
if List.compare_length_with block.body ~len:7 <= 0
312+
then (
313+
let sb = Simple_block.make block in
314+
match SBT.find_opt t sb with
315+
| Some cont' when not (equal cont' cont) ->
316+
rewrite := Addr.Set.add (fst cont') !rewrite;
317+
cont'
318+
| Some _ | None ->
319+
SBT.add t sb cont;
320+
cont)
321+
else cont)
228322
in
229-
Addr.Map.add pc block blocks
230-
| `Many_cases -> blocks)
323+
if not (Addr.Set.is_empty !rewrite)
324+
then (
325+
incr opt_count;
326+
let blocks =
327+
Addr.Set.fold
328+
(fun pc blocks ->
329+
let block = Code.Addr.Map.find pc blocks in
330+
Addr.Map.add
331+
pc
332+
{ block with
333+
body =
334+
List.filter
335+
~f:(function
336+
| Event _ -> false
337+
| _ -> true)
338+
block.body
339+
}
340+
blocks)
341+
!rewrite
342+
blocks
343+
in
344+
match find_outlier_index l with
345+
| #switch_to_cond as opt ->
346+
let block = optimize_switch_to_cond block x l opt in
347+
Addr.Map.add pc block blocks
348+
| `Many_cases ->
349+
Addr.Map.add pc { block with branch = Switch (x, l) } blocks)
350+
else blocks)
231351
| _ -> blocks)
232352
p.blocks
233353
p.blocks
@@ -237,4 +357,4 @@ let switches p =
237357
if stats () then Format.eprintf "Stats - switches: %d@." !opt_count;
238358
if debug_stats ()
239359
then Code.check_updates ~name:"switches" previous_p p ~updates:!opt_count;
240-
p
360+
Deadcode.remove_unused_blocks p

‎compiler/tests-compiler/cond.ml‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,3 +67,97 @@ let%expect_test "conditional" =
6767
}
6868
//end
6969
|}]
70+
71+
let%expect_test "conditional" =
72+
let program =
73+
compile_and_parse
74+
{|
75+
type rip_relative_kind =
76+
| Explicitly_rip_relative
77+
| Implicitly_rip_relative
78+
| Not_rip_relative
79+
80+
(** val rip_relative_kind_beq :
81+
rip_relative_kind -> rip_relative_kind -> bool **)
82+
83+
let rip_relative_kind_beq x y =
84+
match x with
85+
| Explicitly_rip_relative ->
86+
(match y with
87+
| Explicitly_rip_relative -> true
88+
| Implicitly_rip_relative -> false
89+
| Not_rip_relative -> false)
90+
| Implicitly_rip_relative ->
91+
(match y with
92+
| Explicitly_rip_relative -> false
93+
| Implicitly_rip_relative -> true
94+
| Not_rip_relative -> false)
95+
| Not_rip_relative ->
96+
(match y with
97+
| Explicitly_rip_relative -> false
98+
| Implicitly_rip_relative -> false
99+
| Not_rip_relative -> true)
100+
|}
101+
in
102+
print_fun_decl program (Some "rip_relative_kind_beq");
103+
[%expect
104+
{|
105+
function rip_relative_kind_beq(x, y){
106+
switch(x){
107+
case 0:
108+
return 0 === y ? 1 : 0;
109+
case 1:
110+
return 1 === y ? 1 : 0;
111+
default: return 2 === y ? 1 : 0;
112+
}
113+
}
114+
//end
115+
|}]
116+
117+
let%expect_test "conditional" =
118+
let program =
119+
compile_and_parse
120+
{|
121+
type rip_relative_kind =
122+
| Explicitly_rip_relative
123+
| Implicitly_rip_relative
124+
| Not_rip_relative
125+
126+
(** val rip_relative_kind_beq :
127+
rip_relative_kind -> rip_relative_kind -> bool **)
128+
129+
let rip_relative_kind_beq x y =
130+
let i = match x with
131+
| Explicitly_rip_relative ->
132+
(match y with
133+
| Explicitly_rip_relative -> 1
134+
| Implicitly_rip_relative -> 2
135+
| Not_rip_relative -> 2)
136+
| Implicitly_rip_relative ->
137+
(match y with
138+
| Explicitly_rip_relative -> 2
139+
| Implicitly_rip_relative -> 1
140+
| Not_rip_relative -> 2)
141+
| Not_rip_relative ->
142+
(match y with
143+
| Explicitly_rip_relative -> 2
144+
| Implicitly_rip_relative -> 2
145+
| Not_rip_relative -> 1)
146+
in print_int i
147+
|}
148+
in
149+
print_fun_decl program (Some "rip_relative_kind_beq");
150+
[%expect
151+
{|
152+
function rip_relative_kind_beq(x, y){
153+
switch(x){
154+
case 0:
155+
var i = 0 === y ? 1 : 2; break;
156+
case 1:
157+
var i = 1 === y ? 1 : 2; break;
158+
default: var i = 2 === y ? 1 : 2;
159+
}
160+
return caml_call1(Stdlib[44], i);
161+
}
162+
//end
163+
|}]

0 commit comments

Comments
 (0)