vm.fail
MLIR DAG rewriter w/ SSA IR, dominance and CFG analysis passes
1open Ir
2open Cfg
3open Dominance
4
5module PairMap = Map.Make(struct
6 type t = int * int
7 let compare (a1, b1) (a2, b2) =
8 match Stdlib.compare a1 a2 with
9 | 0 -> Stdlib.compare b1 b2
10 | c -> c
11end)
12
13type alloca_info = {
14 aid : value;
15 mutable stores : (value * value) list;
16 mutable loads : (value * value) list;
17}
18
19let find_allocas cfg =
20 let info = ref IntMap.empty in
21 IntMap.iter (fun _ b ->
22 List.iter (function
23 | Alloca a -> info := IntMap.add a { aid = a; stores = []; loads = [] } !info
24 | _ -> ()
25 ) b.insts
26 ) cfg.blocks;
27 IntMap.iter (fun _ b ->
28 List.iter (function
29 | Store (v, p) ->
30 (try let ai = IntMap.find p !info in ai.stores <- (v, b.id) :: ai.stores with Not_found -> ())
31 | Load (r, p) ->
32 (try let ai = IntMap.find p !info in ai.loads <- (r, b.id) :: ai.loads with Not_found -> ())
33 | _ -> ()
34 ) b.insts
35 ) cfg.blocks;
36 IntMap.fold (fun _ ai acc -> ai :: acc) !info []
37
38let place_phi cfg df allocas =
39 let phi_needed = ref PairMap.empty in
40 List.iter (fun ai ->
41 let blocks_with_store = ref (List.fold_left (fun s (_, b) -> IntSet.add b s) IntSet.empty ai.stores) in
42 let work = ref (IntSet.elements !blocks_with_store) in
43 while !work <> [] do
44 let b = List.hd !work in
45 work := List.tl !work;
46 (try
47 let frontiers = IntMap.find b df in
48 IntSet.iter (fun f ->
49 let key = (f, ai.aid) in
50 if not (PairMap.mem key !phi_needed) then begin
51 phi_needed := PairMap.add key true !phi_needed;
52 if not (IntSet.mem f !blocks_with_store) then begin
53 blocks_with_store := IntSet.add f !blocks_with_store;
54 work := f :: !work
55 end
56 end
57 ) frontiers
58 with Not_found -> ())
59 done
60 ) allocas;
61 !phi_needed
62
63let rename cfg phi_needed allocas =
64 let stacks = ref IntMap.empty in
65 let push aid v =
66 let s = try IntMap.find aid !stacks with Not_found -> [] in
67 stacks := IntMap.add aid (v :: s) !stacks
68 in
69 let pop aid =
70 let s = IntMap.find aid !stacks in
71 if List.tl s = [] then stacks := IntMap.remove aid !stacks
72 else stacks := IntMap.add aid (List.tl s) !stacks
73 in
74 let top aid =
75 try Some (List.hd (IntMap.find aid !stacks))
76 with Not_found -> None
77 in
78 (* For each (block, alloca) needing a phi, accumulate operands as the DFS
79 visits each predecessor. Each entry maps a predecessor id to the SSA
80 version of the variable on that edge. These are read only AFTER the full
81 DFS completes, so every predecessor has been visited. *)
82 let phi_ops = ref (PairMap.map (fun _ -> ref []) phi_needed) in
83 let body_insts = ref IntMap.empty in
84 let process_block id b =
85 (* A block that needs a phi defines a new version of the variable (using
86 the alloca id as the SSA name). Push it so subsequent loads in this
87 block read the phi's result rather than a stale predecessor version. *)
88 let pushed = ref [] in
89 List.iter (fun ai ->
90 if PairMap.mem (id, ai.aid) phi_needed then (push ai.aid ai.aid; pushed := ai.aid :: !pushed)
91 ) allocas;
92 let insts = ref [] in
93 List.iter (fun i ->
94 match i with
95 | Alloca a -> () (* promoted away *)
96 | Store (v, p) ->
97 (* A store defines a new version of the variable. Push it once. *)
98 push p v; pushed := p :: !pushed
99 | Load (r, p) ->
100 (match top p with
101 | Some v -> insts := Move (r, v) :: !insts
102 | None -> ())
103 | Phi _ -> () (* input phis are replaced by the inserted ones *)
104 | _ -> insts := i :: !insts
105 ) b.insts;
106 (* Rev to restore original order (insts was built by consing). *)
107 body_insts := IntMap.add id (List.rev !insts) !body_insts;
108 !pushed
109 in
110 let add_phi_operand block aid pred v =
111 if PairMap.mem (block, aid) phi_needed then
112 let ops_ref = PairMap.find (block, aid) !phi_ops in
113 (* Store as (pred, value) so IntMap.of_list below builds a pred->value map. *)
114 ops_ref := (pred, v) :: !ops_ref
115 in
116 let rec dfs id visited =
117 if IntSet.mem id visited then visited
118 else
119 let visited = IntSet.add id visited in
120 let b = IntMap.find id cfg.blocks in
121 let pushed = process_block id b in
122 let visited = List.fold_left (fun vis s ->
123 (* Edge id -> s: the version of each variable live on entry to s is the
124 current top of its stack. Record it as s's phi operand from id,
125 BEFORE descending into s (so s's own phi push doesn't interfere). *)
126 List.iter (fun ai ->
127 match top ai.aid with
128 | Some v -> add_phi_operand s ai.aid id v
129 | None -> ()
130 ) allocas;
131 dfs s vis
132 ) visited b.succs in
133 (* Pop every version this block pushed before backtracking. *)
134 List.iter pop pushed;
135 visited
136 in
137 ignore (dfs cfg.entry IntSet.empty);
138 (* Now build phi instructions from the fully-collected operands. *)
139 let phi_insts = ref IntMap.empty in
140 PairMap.iter (fun (block, aid) ops_ref ->
141 let ops_map = IntMap.of_list !ops_ref in
142 let b = IntMap.find block cfg.blocks in
143 let phi_ops_list = List.map (fun p ->
144 match IntMap.find_opt p ops_map with
145 | Some v -> (v, p)
146 | None -> (0, p)
147 ) b.preds in
148 let existing = try IntMap.find block !phi_insts with Not_found -> [] in
149 phi_insts := IntMap.add block (Phi (aid, phi_ops_list) :: existing) !phi_insts
150 ) !phi_ops;
151 (* Final per-block instruction list: phis first, then the rewritten body. *)
152 let new_insts = ref IntMap.empty in
153 IntMap.iter (fun id body ->
154 let phis = try IntMap.find id !phi_insts with Not_found -> [] in
155 new_insts := IntMap.add id (phis @ body) !new_insts
156 ) !body_insts;
157 !new_insts
158
159let mem2reg f =
160 let cfg = build_cfg f in
161 let doms = dominators cfg in
162 let df = dominance_frontiers cfg doms in
163 let allocas = find_allocas cfg in
164 let phi_needed = place_phi cfg df allocas in
165 let new_insts = rename cfg phi_needed allocas in
166 let new_blocks = IntMap.mapi (fun id b ->
167 { b with insts = try IntMap.find id new_insts with Not_found -> b.insts }
168 ) cfg.blocks in
169 (* Preserve the original input block order for deterministic output. *)
170 let ordered = List.map (fun b -> IntMap.find b.id new_blocks) f.blocks in
171 { f with blocks = ordered }