vm.fail
MLIR DAG rewriter w/ SSA IR, dominance and CFG analysis passes
1open Ir
2open Cfg
3
4module IntMap = Map.Make(Int)
5module IntSet = Set.Make(Int)
6
7let 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
25let 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
76let dominates doms a b =
77 IntSet.mem a (IntMap.find b doms)
78
79let 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
94let dominance_frontiers cfg doms =
95 let df = ref IntMap.empty in
96 IntMap.iter (fun id _ -> df := IntMap.add id IntSet.empty !df) cfg.blocks;
97 IntMap.iter (fun id b ->
98 match b.preds with
99 | [] | [_] -> ()
100 | _ ->
101 let idom_y = idom_from_doms doms id in
102 List.iter (fun p ->
103 let rec walk cur =
104 if Some cur = idom_y then ()
105 else begin
106 df := IntMap.add cur (IntSet.add id (IntMap.find cur !df)) !df;
107 (match idom_from_doms doms cur with
108 | None -> ()
109 | Some n -> walk n)
110 end
111 in
112 walk p
113 ) b.preds
114 ) cfg.blocks;
115 !df