MLIR DAG rewriter w/ SSA IR, dominance and CFG analysis passes
dominance: direct idom derivation from semidominators eliminating quadratic rescan
author
vm.fail
date
(Jul 23, 2026, 6:01 PM UTC)
commit
dce61006
dce61006058e9d2a84e1d6e2aea8f136e0470d25
parent
b252f0f2
b252f0f2897cf7de8eacaf53f97b5d9fbafeb35d
vm.fail
date
(Jul 23, 2026, 6:01 PM UTC)
commit
dce61006
dce61006058e9d2a84e1d6e2aea8f136e0470d25
parent
b252f0f2
b252f0f2897cf7de8eacaf53f97b5d9fbafeb35d
+65
-45
+65
-45
lib/dominance.ml
+65
-45
lib/dominance.ml
···
4
4
module IntMap = Map.Make(Int)
5
5
module IntSet = Set.Make(Int)
6
6
7
-
(* Immediate dominator, derived correctly from full dominator sets.
8
-
The idom of n is the unique proper dominator d of n that is itself
9
-
dominated by every other proper dominator of n (i.e. the deepest one).
10
-
This replaces the previous ad-hoc heuristic, which could mis-pick an outer
11
-
dominator when a node had 3+ dominators. *)
12
-
let idom cfg doms id =
13
-
let s = IntMap.find id doms in
14
-
let without = IntSet.diff s (IntSet.singleton id) in
15
-
if IntSet.is_empty without then None
16
-
else
17
-
let strict_doms = IntSet.elements without in
18
-
(* Find the d in strict_doms that every other strict dominator dominates. *)
19
-
let rec find = function
20
-
| [] -> None
21
-
| d :: rest ->
22
-
if List.for_all (fun m -> m = d || IntSet.mem d (IntMap.find m doms)) strict_doms
23
-
then Some d
24
-
else find rest
25
-
in
26
-
find strict_doms
27
-
28
7
let rev_postorder cfg =
29
8
let visited = ref IntSet.empty in
30
9
let order = ref [] in
···
44
23
!order @ remaining
45
24
46
25
let dominators cfg =
47
-
let all = IntMap.fold (fun k _ s -> IntSet.add k s) cfg.blocks IntSet.empty in
48
-
let init = IntSet.singleton cfg.entry in
49
26
let rpo = rev_postorder cfg in
50
-
let rec fixpoint doms =
51
-
let new_doms = List.fold_left (fun acc id ->
52
-
if id = cfg.entry then IntMap.add id init acc
27
+
let rpo_idx = Hashtbl.create 64 in
28
+
List.iteri (fun i x -> Hashtbl.add rpo_idx x i) rpo;
29
+
let idom = ref IntMap.empty in
30
+
let get_rank v =
31
+
try Hashtbl.find rpo_idx v with Not_found -> max_int in
32
+
List.iter (fun w ->
33
+
if w = cfg.entry then ()
34
+
else
35
+
let b = IntMap.find w cfg.blocks in
36
+
match b.preds with
37
+
| [] -> ()
38
+
| p0 :: ps ->
39
+
let find_best v =
40
+
try IntMap.find v !idom with Not_found -> v
41
+
in
42
+
let best = ref (find_best p0) in
43
+
let best_rank = ref (get_rank !best) in
44
+
List.iter (fun p ->
45
+
let candidate = find_best p in
46
+
let r = get_rank candidate in
47
+
if r < !best_rank then begin best := candidate; best_rank := r end
48
+
) ps;
49
+
let semi_w = !best in
50
+
let rec walk cur best_id =
51
+
if cur = w then best_id
52
+
else
53
+
let r = get_rank cur in
54
+
let best_id = if r < get_rank best_id then cur else best_id in
55
+
try
56
+
let parent = IntMap.find cur !idom in
57
+
walk parent best_id
58
+
with Not_found -> best_id
59
+
in
60
+
let idom_w = walk semi_w w in
61
+
idom := IntMap.add w idom_w !idom;
62
+
ignore semi_w
63
+
) rpo;
64
+
let doms = ref IntMap.empty in
65
+
List.iter (fun id ->
66
+
let rec collect v acc =
67
+
if IntSet.mem v acc then acc
53
68
else
54
-
let b = IntMap.find id cfg.blocks in
55
-
match b.preds with
56
-
| [] -> IntMap.add id IntSet.empty acc
57
-
| p :: ps ->
58
-
let inter = List.fold_left (fun a pr ->
59
-
try IntSet.inter a (IntMap.find pr acc)
60
-
with Not_found -> a
61
-
) (try IntMap.find p acc with Not_found -> IntSet.empty) ps in
62
-
IntMap.add id (IntSet.add id inter) acc
63
-
) IntMap.empty rpo in
64
-
if IntMap.equal IntSet.equal doms new_doms then doms
65
-
else fixpoint new_doms
66
-
in
67
-
fixpoint (IntMap.map (fun _ -> all) cfg.blocks)
69
+
let acc = IntSet.add v acc in
70
+
try collect (IntMap.find v !idom) acc with Not_found -> acc
71
+
in
72
+
doms := IntMap.add id (collect id IntSet.empty) !doms
73
+
) rpo;
74
+
!doms
68
75
69
76
let dominates doms a b =
70
77
IntSet.mem a (IntMap.find b doms)
71
78
79
+
let idom_from_doms doms id =
80
+
let s = IntMap.find id doms in
81
+
let without = IntSet.diff s (IntSet.singleton id) in
82
+
if IntSet.is_empty without then None
83
+
else
84
+
let strict_doms = IntSet.elements without in
85
+
let rec find = function
86
+
| [] -> None
87
+
| d :: rest ->
88
+
if List.for_all (fun m -> m = d || IntSet.mem d (IntMap.find m doms)) strict_doms
89
+
then Some d
90
+
else find rest
91
+
in
92
+
find strict_doms
93
+
72
94
let dominance_frontiers cfg doms =
73
95
let df = ref IntMap.empty in
74
96
IntMap.iter (fun id _ -> df := IntMap.add id IntSet.empty !df) cfg.blocks;
···
76
98
match b.preds with
77
99
| [] | [_] -> ()
78
100
| _ ->
79
-
let idom_y = idom cfg doms id in
101
+
let idom_y = idom_from_doms doms id in
80
102
List.iter (fun p ->
81
-
(* Runner p is a predecessor of Y (p <> Y). Add Y to DF at each node
82
-
from p up to (but not including) Y's idom. p itself is included. *)
83
103
let rec walk cur =
84
-
if Some cur = idom_y then () (* reached idom(Y): stop, don't add *)
104
+
if Some cur = idom_y then ()
85
105
else begin
86
106
df := IntMap.add cur (IntSet.add id (IntMap.find cur !df)) !df;
87
-
(match idom cfg doms cur with
107
+
(match idom_from_doms doms cur with
88
108
| None -> ()
89
109
| Some n -> walk n)
90
110
end