armfortas Public

Watch 0 Fork 0 Star 0
Rust · 5748 bytes Raw Blame History
1 use std::collections::BTreeSet;
2 use std::path::PathBuf;
3
4 use armfortas::driver::OptLevel;
5 use armfortas::testing::{capture_from_path, CaptureRequest, CapturedStage, Stage};
6
7 fn fixture(name: &str) -> PathBuf {
8 let path = PathBuf::from("test_programs").join(name);
9 assert!(path.exists(), "missing test fixture {}", path.display());
10 path
11 }
12
13 fn capture_text(request: CaptureRequest, stage: Stage) -> String {
14 let result = capture_from_path(&request).expect("capture should succeed");
15 match result.get(stage) {
16 Some(CapturedStage::Text(text)) => text.clone(),
17 Some(CapturedStage::Run(_)) => panic!("expected text stage for {}", stage.as_str()),
18 None => panic!("missing requested stage {}", stage.as_str()),
19 }
20 }
21
22 fn function_section<'a>(ir: &'a str, name: &str) -> &'a str {
23 let header = format!("func @{}", name);
24 let start = ir
25 .find(&header)
26 .unwrap_or_else(|| panic!("missing function section for {}", name));
27 let rest = &ir[start..];
28 let end = rest
29 .find("\n }\n")
30 .unwrap_or_else(|| panic!("unterminated function section for {}", name));
31 &rest[..end + "\n }".len()]
32 }
33
34 fn block_section<'a>(func_section: &'a str, prefix: &str) -> &'a str {
35 let mut start = None;
36 let mut end = None;
37
38 for (idx, _line) in func_section.match_indices('\n') {
39 let line_start = idx + 1;
40 let tail = &func_section[line_start..];
41 let line_text = tail.split_once('\n').map(|(line, _)| line).unwrap_or(tail);
42
43 if start.is_none() {
44 if line_text.starts_with(" ")
45 && !line_text.starts_with(" ")
46 && line_text[4..].starts_with(prefix)
47 {
48 start = Some(line_start);
49 }
50 continue;
51 }
52
53 if line_text.starts_with(" ") && !line_text.starts_with(" ") {
54 end = Some(idx);
55 break;
56 }
57 if line_text == " }" {
58 end = Some(idx);
59 break;
60 }
61 }
62
63 let start = start.unwrap_or_else(|| panic!("missing block with prefix {}", prefix));
64 let end = end.unwrap_or(func_section.len());
65 &func_section[start..end]
66 }
67
68 fn last_block_section<'a>(func_section: &'a str, prefix: &str) -> &'a str {
69 let mut starts = Vec::new();
70
71 for (idx, _line) in func_section.match_indices('\n') {
72 let line_start = idx + 1;
73 let tail = &func_section[line_start..];
74 let line_text = tail.split_once('\n').map(|(line, _)| line).unwrap_or(tail);
75 if line_text.starts_with(" ")
76 && !line_text.starts_with(" ")
77 && line_text[4..].starts_with(prefix)
78 {
79 starts.push(line_start);
80 }
81 }
82
83 let start = *starts
84 .last()
85 .unwrap_or_else(|| panic!("missing block with prefix {}", prefix));
86 let tail = &func_section[start..];
87 let end_rel = tail
88 .match_indices('\n')
89 .find_map(|(idx, _)| {
90 let line_start = idx + 1;
91 let line_tail = &tail[line_start..];
92 let line_text = line_tail
93 .split_once('\n')
94 .map(|(line, _)| line)
95 .unwrap_or(line_tail);
96 if (line_text.starts_with(" ") && !line_text.starts_with(" "))
97 || line_text == " }"
98 {
99 Some(idx)
100 } else {
101 None
102 }
103 })
104 .unwrap_or(tail.len());
105 &tail[..end_rel]
106 }
107
108 #[test]
109 fn o2_reuses_branch_join_affine_expression() {
110 let source = fixture("realworld_join_bias_sum.f90");
111
112 let raw_ir = capture_text(
113 CaptureRequest {
114 input: source.clone(),
115 requested: BTreeSet::from([Stage::Ir]),
116 opt_level: OptLevel::O0,
117 },
118 Stage::Ir,
119 );
120 let opt_ir = capture_text(
121 CaptureRequest {
122 input: source,
123 requested: BTreeSet::from([Stage::OptIr]),
124 opt_level: OptLevel::O2,
125 },
126 Stage::OptIr,
127 );
128
129 let raw_tally = function_section(&raw_ir, "tally");
130 let raw_join = last_block_section(raw_tally, "if_end_");
131
132 assert!(
133 raw_join.matches("call @offset_value").count() >= 2,
134 "raw join block should still recompute the repeated branch-join PURE helper call:\n{}",
135 raw_join
136 );
137 assert!(
138 opt_ir.matches("call @offset_value").count() < raw_ir.matches("call @offset_value").count(),
139 "O2 should reduce duplicated branch-join PURE helper calls:\n{}",
140 opt_ir
141 );
142 }
143
144 #[test]
145 fn o2_removes_dead_seed_store_across_noalias_call() {
146 let source = fixture("realworld_seed_overwrite.f90");
147
148 let raw_ir = capture_text(
149 CaptureRequest {
150 input: source.clone(),
151 requested: BTreeSet::from([Stage::Ir]),
152 opt_level: OptLevel::O0,
153 },
154 Stage::Ir,
155 );
156 let opt_ir = capture_text(
157 CaptureRequest {
158 input: source,
159 requested: BTreeSet::from([Stage::OptIr]),
160 opt_level: OptLevel::O2,
161 },
162 Stage::OptIr,
163 );
164
165 let raw_fill = function_section(&raw_ir, "seed_and_fill");
166 let opt_fill = function_section(&opt_ir, "seed_and_fill");
167 let raw_body = block_section(raw_fill, "do_body_");
168 let opt_body = block_section(opt_fill, "do_body_");
169
170 assert!(
171 raw_body.matches("store ").count() >= 2,
172 "raw loop body should still contain the seed store and the real fill store:\n{}",
173 raw_body
174 );
175 assert!(
176 opt_body.matches("store ").count() < raw_body.matches("store ").count(),
177 "O2 should remove the dead seed store while keeping the real fill:\n{}",
178 opt_body
179 );
180 }
181