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
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. *)
12let 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
28let rev_postorder cfg =
29 let visited = ref IntSet.empty in
30 let order = ref [] in
31 let rec dfs id =
32 if IntSet.mem id !visited then ()
33 else begin
34 visited := IntSet.add id !visited;
35 let b = IntMap.find id cfg.blocks in
36 List.iter dfs b.succs;
37 order := id :: !order
38 end
39 in
40 dfs cfg.entry;
41 let remaining = IntMap.fold (fun k _ acc ->
42 if IntSet.mem k !visited then acc else k :: acc
43 ) cfg.blocks [] in
44 !order @ remaining
45
46let 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 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
53 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)
68
69let dominates doms a b =
70 IntSet.mem a (IntMap.find b doms)
71
72let dominance_frontiers cfg doms =
73 let df = ref IntMap.empty in
74 IntMap.iter (fun id _ -> df := IntMap.add id IntSet.empty !df) cfg.blocks;
75 IntMap.iter (fun id b ->
76 match b.preds with
77 | [] | [_] -> ()
78 | _ ->
79 let idom_y = idom cfg doms id in
80 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 let rec walk cur =
84 if Some cur = idom_y then () (* reached idom(Y): stop, don't add *)
85 else begin
86 df := IntMap.add cur (IntSet.add id (IntMap.find cur !df)) !df;
87 (match idom cfg doms cur with
88 | None -> ()
89 | Some n -> walk n)
90 end
91 in
92 walk p
93 ) b.preds
94 ) cfg.blocks;
95 !df