MLIR DAG rewriter w/ SSA IR, dominance and CFG analysis passes
1open Ir
2
3module IntMap = Map.Make(Int)
4module IntSet = Set.Make(Int)
5
6type cfg = {
7 blocks : block IntMap.t;
8 entry : value;
9 mutable defs : inst IntMap.t;
10}
11
12let build_cfg (f : func) =
13 let blocks = List.fold_left (fun m b -> IntMap.add b.id b m) IntMap.empty f.blocks in
14 let cfg = { blocks; entry = f.entry; defs = IntMap.empty } in
15 let add_edges b =
16 let rec scan acc = function
17 | [] -> List.rev acc
18 | Ret _ :: _ -> List.rev acc
19 | BinOp (_, _, t) :: rest -> scan (t :: acc) rest
20 | Br t :: rest -> scan (t :: acc) rest
21 | CondBr (_, t1, t2) :: rest -> scan (t2 :: t1 :: acc) rest
22 | _ :: rest -> scan acc rest
23 in
24 b.succs <- scan [] b.insts
25 in
26 IntMap.iter (fun _ b -> add_edges b) blocks;
27 IntMap.iter (fun _ b ->
28 List.iter (fun s ->
29 try
30 let sb = IntMap.find s cfg.blocks in
31 sb.preds <- b.id :: sb.preds
32 with Not_found -> ()
33 ) b.succs
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;
47 cfg
48
49let get_block cfg id = IntMap.find id cfg.blocks
50let get_succs cfg id = (get_block cfg id).succs
51let get_preds cfg id = (get_block cfg id).preds
52let get_def cfg v = IntMap.find v cfg.defs