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 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
50let dominates doms a b =
51 IntSet.mem a (IntMap.find b doms)
52
53let dominance_frontiers cfg doms =
54 let df = ref IntMap.empty in
55 IntMap.iter (fun id _ -> df := IntMap.add id IntSet.empty !df) cfg.blocks;
56 IntMap.iter (fun id b ->
57 match b.preds with
58 | [] | [_] -> ()
59 | _ ->
60 let idom_y = idom cfg doms id in
61 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 let rec walk cur =
65 if Some cur = idom_y then () (* reached idom(Y): stop, don't add *)
66 else begin
67 df := IntMap.add cur (IntSet.add id (IntMap.find cur !df)) !df;
68 (match idom cfg doms cur with
69 | None -> ()
70 | Some n -> walk n)
71 end
72 in
73 walk p
74 ) b.preds
75 ) cfg.blocks;
76 !df