MLIR DAG rewriter w/ SSA IR, dominance and CFG analysis passes
1open Ir
2open Nemo
3open Printer
4open Parser
5
6(* ---- Inline unit tests (no file I/O) ---- *)
7
8let test_simple () =
9 Printf.printf "Test Simple\n";
10 let b1 = mk_block 1 [Alloca 100; Store (1, 100); Br 2] in
11 let b2 = mk_block 2 [Load (200, 100); Ret (Some 200)] in
12 let f = { blocks = [b1; b2]; entry = 1 } in
13 Printf.printf "Before:\n";
14 print_func f;
15 let f' = mem2reg f in
16 Printf.printf "\nAfter:\n";
17 print_func f'
18
19let test_phi () =
20 Printf.printf "\nTest Phi\n";
21 let b1 = mk_block 1 [Alloca 100; Store (1, 100); Br 2] in
22 let b2 = mk_block 2 [Alloca 101; Store (2, 101); Br 3] in
23 let b3 = mk_block 3 [Load (200, 100); Load (201, 101); Ret (Some 200)] in
24 let f = { blocks = [b1; b2; b3]; entry = 1 } in
25 Printf.printf "Before:\n";
26 print_func f;
27 let f' = mem2reg f in
28 Printf.printf "\nAfter:\n";
29 print_func f'
30
31let test_branch () =
32 Printf.printf "\nTest Branch with Phi\n";
33 let b1 = mk_block 1 [Alloca 100; CondBr (1, 2, 3)] in
34 let b2 = mk_block 2 [Store (10, 100); Br 4] in
35 let b3 = mk_block 3 [Store (20, 100); Br 4] in
36 let b4 = mk_block 4 [Load (200, 100); Ret (Some 200)] in
37 let f = { blocks = [b1; b2; b3; b4]; entry = 1 } in
38 Printf.printf "Before:\n";
39 print_func f;
40 let f' = mem2reg f in
41 Printf.printf "\nAfter:\n";
42 print_func f'
43
44(* ---- File-based regression test ---- *)
45
46(* Parse a .ir file, run mem2reg, and assert the printed SSA output equals the
47 expected string. Exits with status 1 on mismatch so `make test` fails. *)
48let string_of_func f =
49 let buf = Buffer.create 256 in
50 List.iter (fun b ->
51 Buffer.add_string buf (Printf.sprintf "block %d:\n" b.id);
52 List.iter (fun i ->
53 Buffer.add_string buf (Printf.sprintf " %s\n" (Printer.string_of_inst i))
54 ) b.insts
55 ) f.blocks;
56 Buffer.contents buf
57
58let test_file name expected () =
59 Printf.printf "\nTest file: %s\n" name;
60 let f = parse_func_from_file ("examples/" ^ name) in
61 let f' = mem2reg f in
62 let got = string_of_func f' in
63 if got <> expected then begin
64 Printf.eprintf "FAIL %s\n--- expected ---\n%s\n--- got ---\n%s\n" name expected got;
65 exit 1
66 end;
67 print_string got
68
69let expected_simple =
70"block 1:\n br 2\nblock 2:\n move 200, 1\n ret 200\n"
71
72let expected_basic =
73"block 1:\n br 2\nblock 2:\n move 200, 42\n ret 200\n"
74
75(* Example exercising the phi-handling path: an input phi node is filled in
76 from the renamed value at each predecessor. *)
77let expected_phi =
78"block 1:\n condbr 1, 2, 3\nblock 2:\n br 4\nblock 3:\n br 4\nblock 4:\n phi 100 [(20, 3), (10, 2)]\n move 200, 100\n ret 200\n"
79
80let () =
81 test_simple ();
82 test_phi ();
83 test_branch ();
84 test_file "simple.ir" expected_simple ();
85 test_file "basic.ir" expected_basic ();
86 test_file "phi.ir" expected_phi ()