MLIR DAG rewriter w/ SSA IR, dominance and CFG analysis passes
Merge pull request 'iterate in reverse postorder for faster fixpoint convergence and materialise defs table' (#1) from dev into main
Reviewed-on: https://codeberg.org/50/nemo/pulls/1
author
vm.fail
date
(Jul 24, 2026, 1:15 AM +0200)
commit
54687010
54687010c73a294510af15765bf4b891697e4a00
parent
ee8d504c
ee8d504ca88d5799ebf404057a28e90414ecd178
vm.fail
date
(Jul 24, 2026, 1:15 AM +0200)
commit
54687010
54687010c73a294510af15765bf4b891697e4a00
parent
ee8d504c
ee8d504ca88d5799ebf404057a28e90414ecd178
+91
-38
+15
-1
lib/cfg.ml
+15
-1
lib/cfg.ml
···
6
6
type cfg = {
7
7
blocks : block IntMap.t;
8
8
entry : value;
9
+
mutable defs : inst IntMap.t;
9
10
}
10
11
11
12
let build_cfg (f : func) =
12
13
let blocks = List.fold_left (fun m b -> IntMap.add b.id b m) IntMap.empty f.blocks in
13
-
let cfg = { blocks; entry = f.entry } in
14
+
let cfg = { blocks; entry = f.entry; defs = IntMap.empty } in
14
15
let add_edges b =
15
16
let rec scan acc = function
16
17
| [] -> List.rev acc
···
31
32
with Not_found -> ()
32
33
) b.succs
33
34
) blocks;
35
+
let defs = IntMap.fold (fun _ b acc ->
36
+
List.fold_left (fun d inst ->
37
+
match inst with
38
+
| Alloca a -> IntMap.add a inst acc
39
+
| Load (r, _) -> IntMap.add r inst d
40
+
| BinOp (r, _, _) -> IntMap.add r inst d
41
+
| Move (r, _) -> IntMap.add r inst d
42
+
| Phi (r, _) -> IntMap.add r inst d
43
+
| _ -> d
44
+
) acc b.insts
45
+
) blocks IntMap.empty in
46
+
cfg.defs <- defs;
34
47
cfg
35
48
36
49
let get_block cfg id = IntMap.find id cfg.blocks
37
50
let get_succs cfg id = (get_block cfg id).succs
38
51
let get_preds cfg id = (get_block cfg id).preds
52
+
let get_def cfg v = IntMap.find v cfg.defs
+76
-37
lib/dominance.ml
+76
-37
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 =
7
+
let rev_postorder cfg =
8
+
let visited = ref IntSet.empty in
9
+
let order = ref [] in
10
+
let rec dfs id =
11
+
if IntSet.mem id !visited then ()
12
+
else begin
13
+
visited := IntSet.add id !visited;
14
+
let b = IntMap.find id cfg.blocks in
15
+
List.iter dfs b.succs;
16
+
order := id :: !order
17
+
end
18
+
in
19
+
dfs cfg.entry;
20
+
let remaining = IntMap.fold (fun k _ acc ->
21
+
if IntSet.mem k !visited then acc else k :: acc
22
+
) cfg.blocks [] in
23
+
!order @ remaining
24
+
25
+
let dominators cfg =
26
+
let rpo = rev_postorder cfg in
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
68
+
else
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
75
+
76
+
let dominates doms a b =
77
+
IntSet.mem a (IntMap.find b doms)
78
+
79
+
let idom_from_doms doms id =
13
80
let s = IntMap.find id doms in
14
81
let without = IntSet.diff s (IntSet.singleton id) in
15
82
if IntSet.is_empty without then None
16
83
else
17
84
let strict_doms = IntSet.elements without in
18
-
(* Find the d in strict_doms that every other strict dominator dominates. *)
19
85
let rec find = function
20
86
| [] -> None
21
87
| d :: rest ->
···
25
91
in
26
92
find strict_doms
27
93
28
-
let dominators cfg =
29
-
let all = IntMap.fold (fun k _ s -> IntSet.add k s) cfg.blocks IntSet.empty in
30
-
let init = IntSet.singleton cfg.entry in
31
-
let rec fixpoint doms =
32
-
let new_doms = IntMap.mapi (fun id _ ->
33
-
if id = cfg.entry then init
34
-
else
35
-
let b = IntMap.find id cfg.blocks in
36
-
match b.preds with
37
-
| [] -> IntSet.empty
38
-
| p :: ps ->
39
-
let inter = List.fold_left (fun acc pr ->
40
-
try IntSet.inter acc (IntMap.find pr doms)
41
-
with Not_found -> acc
42
-
) (try IntMap.find p doms with Not_found -> IntSet.empty) ps in
43
-
IntSet.add id inter
44
-
) cfg.blocks in
45
-
if IntMap.equal IntSet.equal doms new_doms then doms
46
-
else fixpoint new_doms
47
-
in
48
-
fixpoint (IntMap.map (fun _ -> all) cfg.blocks)
49
-
50
-
let dominates doms a b =
51
-
IntSet.mem a (IntMap.find b doms)
52
-
53
94
let dominance_frontiers cfg doms =
54
95
let df = ref IntMap.empty in
55
96
IntMap.iter (fun id _ -> df := IntMap.add id IntSet.empty !df) cfg.blocks;
···
57
98
match b.preds with
58
99
| [] | [_] -> ()
59
100
| _ ->
60
-
let idom_y = idom cfg doms id in
101
+
let idom_y = idom_from_doms doms id in
61
102
List.iter (fun p ->
62
-
(* Runner p is a predecessor of Y (p <> Y). Add Y to DF at each node
63
-
from p up to (but not including) Y's idom. p itself is included. *)
64
103
let rec walk cur =
65
-
if Some cur = idom_y then () (* reached idom(Y): stop, don't add *)
104
+
if Some cur = idom_y then ()
66
105
else begin
67
106
df := IntMap.add cur (IntSet.add id (IntMap.find cur !df)) !df;
68
-
(match idom cfg doms cur with
107
+
(match idom_from_doms doms cur with
69
108
| None -> ()
70
109
| Some n -> walk n)
71
110
end