cranelift_codegen/souper_harvest.rs
1//! Harvest left-hand side superoptimization candidates.
2//!
3//! Given a clif function, harvest all its integer subexpressions, so that they
4//! can be fed into [Souper](https://github.com/google/souper) as candidates for
5//! superoptimization. For some of these candidates, Souper will successfully
6//! synthesize a right-hand side that is equivalent but has lower cost than the
7//! left-hand side. Then, we can combine these left- and right-hand sides into a
8//! complete optimization, and add it to our peephole passes.
9//!
10//! To harvest the expression that produced a given value `x`, we do a
11//! post-order traversal of the dataflow graph starting from `x`. As we do this
12//! traversal, we maintain a map from clif values to their translated Souper
13//! values. We stop traversing when we reach anything that can't be translated
14//! into Souper IR: a memory load, a float-to-int conversion, a block parameter,
15//! etc. For values produced by these instructions, we create a Souper `var`,
16//! which is an input variable to the optimization. For instructions that have a
17//! direct mapping into Souper IR, we get the Souper version of each of its
18//! operands and then create the Souper version of the instruction itself. It
19//! should now be clear why we do a post-order traversal: we need an
20//! instruction's translated operands in order to translate the instruction
21//! itself. Once this instruction is translated, we update the clif-to-souper
22//! map with this new translation so that any other instruction that uses this
23//! result as an operand has access to the translated value. When the traversal
24//! is complete we return the translation of `x` as the root of left-hand side
25//! candidate.
26
27use crate::ir;
28use souper_ir::ast;
29use std::collections::{HashMap, HashSet};
30use std::string::String;
31use std::sync::mpsc;
32use std::vec::Vec;
33
34/// Harvest Souper left-hand side candidates from the given function.
35///
36/// Candidates are reported through the given MPSC sender.
37pub fn do_souper_harvest(func: &ir::Function, out: &mut mpsc::Sender<String>) {
38 let mut allocs = Allocs::default();
39
40 // Iterate over each instruction in each block and try and harvest a
41 // left-hand side from its result.
42 for block in func.layout.blocks() {
43 let mut option_inst = func.layout.first_inst(block);
44 while let Some(inst) = option_inst {
45 let results = func.dfg.inst_results(inst);
46 if results.len() == 1 {
47 let val = results[0];
48 let ty = func.dfg.value_type(val);
49 if ty.is_int() && ty.lane_count() == 1 {
50 harvest_candidate_lhs(&mut allocs, func, val, out);
51 }
52 }
53 option_inst = func.layout.next_inst(inst);
54 }
55 }
56}
57
58/// Allocations that we reuse across many LHS candidate harvests.
59#[derive(Default)]
60struct Allocs {
61 /// A map from cranelift IR to souper IR for values that we've already
62 /// translated into souper IR.
63 ir_to_souper_val: HashMap<ir::Value, ast::ValueId>,
64
65 /// Stack of to-visit and to-trace values for the post-order DFS.
66 dfs_stack: Vec<StackEntry>,
67
68 /// Set of values we've already seen in our post-order DFS.
69 dfs_seen: HashSet<ir::Value>,
70}
71
72impl Allocs {
73 /// Reset the collections to their empty state (without deallocating their
74 /// backing data).
75 fn reset(&mut self) {
76 self.ir_to_souper_val.clear();
77 self.dfs_stack.clear();
78 self.dfs_seen.clear();
79 }
80}
81
82/// Harvest a candidate LHS for `val` from the dataflow graph.
83fn harvest_candidate_lhs(
84 allocs: &mut Allocs,
85 func: &ir::Function,
86 val: ir::Value,
87 out: &mut mpsc::Sender<String>,
88) {
89 allocs.reset();
90 let mut lhs = ast::LeftHandSideBuilder::default();
91 let mut non_var_count = 0;
92
93 // Should we keep tracing through the given `val`? Only if it is defined
94 // by an instruction that we can translate to Souper IR.
95 let should_trace = |val| match func.dfg.value_def(val) {
96 ir::ValueDef::Result(inst, 0) => match func.dfg.insts[inst].opcode() {
97 ir::Opcode::Iadd
98 | ir::Opcode::Imul
99 | ir::Opcode::Udiv
100 | ir::Opcode::Sdiv
101 | ir::Opcode::Urem
102 | ir::Opcode::Srem
103 | ir::Opcode::Band
104 | ir::Opcode::Bor
105 | ir::Opcode::Bxor
106 | ir::Opcode::Bnot
107 | ir::Opcode::Ishl
108 | ir::Opcode::Sshr
109 | ir::Opcode::Ushr
110 | ir::Opcode::Select
111 | ir::Opcode::Uextend
112 | ir::Opcode::Sextend
113 | ir::Opcode::Trunc
114 | ir::Opcode::Icmp
115 | ir::Opcode::Popcnt
116 | ir::Opcode::Bitrev
117 | ir::Opcode::Clz
118 | ir::Opcode::Ctz
119 // TODO: ir::Opcode::IaddCarry
120 | ir::Opcode::SaddSat
121 | ir::Opcode::SsubSat
122 | ir::Opcode::UsubSat => true,
123 _ => false,
124 },
125 _ => false,
126 };
127
128 post_order_dfs(allocs, &func.dfg, val, should_trace, |allocs, val| {
129 let souper_assignment_rhs = match func.dfg.value_def(val) {
130 ir::ValueDef::Result(inst, 0) => {
131 let args = func.dfg.inst_args(inst);
132
133 // Get the n^th argument as a souper operand.
134 let arg = |allocs: &mut Allocs, n| {
135 let arg = args[n];
136 if let Some(a) = allocs.ir_to_souper_val.get(&arg).copied() {
137 a.into()
138 } else {
139 // The only arguments we get that we haven't already
140 // converted into a souper instruction are `iconst`s.
141 // This is because souper only allows
142 // constants as operands, and it doesn't allow assigning
143 // constants to a variable name. So we lazily convert
144 // `iconst`s into souper operands here,
145 // when they are actually used.
146 match func.dfg.value_def(arg) {
147 ir::ValueDef::Result(inst, 0) => match func.dfg.insts[inst] {
148 ir::InstructionData::UnaryImm { opcode, imm } => {
149 debug_assert_eq!(opcode, ir::Opcode::Iconst);
150 let imm: i64 = imm.into();
151 ast::Operand::Constant(ast::Constant {
152 value: imm.into(),
153 r#type: souper_type_of(&func.dfg, arg),
154 })
155 }
156 _ => unreachable!(
157 "only iconst instructions \
158 aren't in `ir_to_souper_val`"
159 ),
160 },
161 _ => unreachable!(
162 "only iconst instructions \
163 aren't in `ir_to_souper_val`"
164 ),
165 }
166 }
167 };
168
169 match (func.dfg.insts[inst].opcode(), &func.dfg.insts[inst]) {
170 (ir::Opcode::Iadd, _) => {
171 let a = arg(allocs, 0);
172 let b = arg(allocs, 1);
173 ast::Instruction::Add { a, b }.into()
174 }
175 (ir::Opcode::Imul, _) => {
176 let a = arg(allocs, 0);
177 let b = arg(allocs, 1);
178 ast::Instruction::Mul { a, b }.into()
179 }
180 (ir::Opcode::Udiv, _) => {
181 let a = arg(allocs, 0);
182 let b = arg(allocs, 1);
183 ast::Instruction::Udiv { a, b }.into()
184 }
185 (ir::Opcode::Sdiv, _) => {
186 let a = arg(allocs, 0);
187 let b = arg(allocs, 1);
188 ast::Instruction::Sdiv { a, b }.into()
189 }
190 (ir::Opcode::Urem, _) => {
191 let a = arg(allocs, 0);
192 let b = arg(allocs, 1);
193 ast::Instruction::Urem { a, b }.into()
194 }
195 (ir::Opcode::Srem, _) => {
196 let a = arg(allocs, 0);
197 let b = arg(allocs, 1);
198 ast::Instruction::Srem { a, b }.into()
199 }
200 (ir::Opcode::Band, _) => {
201 let a = arg(allocs, 0);
202 let b = arg(allocs, 1);
203 ast::Instruction::And { a, b }.into()
204 }
205 (ir::Opcode::Bor, _) => {
206 let a = arg(allocs, 0);
207 let b = arg(allocs, 1);
208 ast::Instruction::Or { a, b }.into()
209 }
210 (ir::Opcode::Bxor, _) => {
211 let a = arg(allocs, 0);
212 let b = arg(allocs, 1);
213 ast::Instruction::Xor { a, b }.into()
214 }
215 (ir::Opcode::Bnot, _) => {
216 let a = arg(allocs, 0);
217 let b = ast::Constant {
218 value: -1,
219 r#type: souper_type_of(&func.dfg, val),
220 }
221 .into();
222 ast::Instruction::Xor { a, b }.into()
223 }
224 (ir::Opcode::Ishl, _) => {
225 let a = arg(allocs, 0);
226 let b = arg(allocs, 1);
227 ast::Instruction::Shl { a, b }.into()
228 }
229 (ir::Opcode::Sshr, _) => {
230 let a = arg(allocs, 0);
231 let b = arg(allocs, 1);
232 ast::Instruction::Ashr { a, b }.into()
233 }
234 (ir::Opcode::Ushr, _) => {
235 let a = arg(allocs, 0);
236 let b = arg(allocs, 1);
237 ast::Instruction::Lshr { a, b }.into()
238 }
239 (ir::Opcode::Select, _) => {
240 let a = arg(allocs, 0);
241
242 // While Cranelift allows any width condition for
243 // `select` and checks it against `0`, Souper requires
244 // an `i1`. So insert a `ne %x, 0` as needed.
245 let a = match a {
246 ast::Operand::Value(id) => match lhs.get_value(id).r#type {
247 Some(ast::Type { width: 1 }) => a,
248 _ => lhs
249 .assignment(
250 None,
251 Some(ast::Type { width: 1 }),
252 ast::Instruction::Ne {
253 a,
254 b: ast::Constant {
255 value: 0,
256 r#type: None,
257 }
258 .into(),
259 },
260 vec![],
261 )
262 .into(),
263 },
264 ast::Operand::Constant(ast::Constant { value, .. }) => ast::Constant {
265 value: (value != 0) as _,
266 r#type: Some(ast::Type { width: 1 }),
267 }
268 .into(),
269 };
270
271 let b = arg(allocs, 1);
272 let c = arg(allocs, 2);
273 ast::Instruction::Select { a, b, c }.into()
274 }
275 (ir::Opcode::Uextend, _) => {
276 let a = arg(allocs, 0);
277 ast::Instruction::Zext { a }.into()
278 }
279 (ir::Opcode::Sextend, _) => {
280 let a = arg(allocs, 0);
281 ast::Instruction::Sext { a }.into()
282 }
283 (ir::Opcode::Trunc, _) => {
284 let a = arg(allocs, 0);
285 ast::Instruction::Trunc { a }.into()
286 }
287 (ir::Opcode::Icmp, ir::InstructionData::IntCompare { cond, .. }) => {
288 let a = arg(allocs, 0);
289 let b = arg(allocs, 1);
290 let cmp = match cond {
291 ir::condcodes::IntCC::Equal => ast::Instruction::Eq { a, b }.into(),
292 ir::condcodes::IntCC::NotEqual => ast::Instruction::Ne { a, b }.into(),
293 ir::condcodes::IntCC::UnsignedLessThan => {
294 ast::Instruction::Ult { a, b }.into()
295 }
296 ir::condcodes::IntCC::SignedLessThan => {
297 ast::Instruction::Slt { a, b }.into()
298 }
299 ir::condcodes::IntCC::UnsignedLessThanOrEqual => {
300 ast::Instruction::Ule { a, b }.into()
301 }
302 ir::condcodes::IntCC::SignedLessThanOrEqual => {
303 ast::Instruction::Sle { a, b }.into()
304 }
305 _ => ast::AssignmentRhs::Var,
306 };
307
308 match cmp {
309 ast::AssignmentRhs::Var => ast::AssignmentRhs::Var,
310 cmp => {
311 let cmp = lhs
312 .assignment(None, Some(ast::Type { width: 1 }), cmp, vec![])
313 .into();
314 ast::Instruction::Zext { a: cmp }.into()
315 }
316 }
317 }
318 (ir::Opcode::Popcnt, _) => {
319 let a = arg(allocs, 0);
320 ast::Instruction::Ctpop { a }.into()
321 }
322 (ir::Opcode::Bitrev, _) => {
323 let a = arg(allocs, 0);
324 ast::Instruction::BitReverse { a }.into()
325 }
326 (ir::Opcode::Clz, _) => {
327 let a = arg(allocs, 0);
328 ast::Instruction::Ctlz { a }.into()
329 }
330 (ir::Opcode::Ctz, _) => {
331 let a = arg(allocs, 0);
332 ast::Instruction::Cttz { a }.into()
333 }
334 // TODO: ir::Opcode::IaddCarry
335 (ir::Opcode::SaddSat, _) => {
336 let a = arg(allocs, 0);
337 let b = arg(allocs, 1);
338 ast::Instruction::SaddSat { a, b }.into()
339 }
340 (ir::Opcode::SsubSat, _) => {
341 let a = arg(allocs, 0);
342 let b = arg(allocs, 1);
343 ast::Instruction::SsubSat { a, b }.into()
344 }
345 (ir::Opcode::UsubSat, _) => {
346 let a = arg(allocs, 0);
347 let b = arg(allocs, 1);
348 ast::Instruction::UsubSat { a, b }.into()
349 }
350 // Because Souper doesn't allow constants to be on the right
351 // hand side of an assignment (i.e. `%0:i32 = 1234` is
352 // disallowed) we have to ignore `iconst`
353 // instructions until we process them as operands for some
354 // other instruction. See the `arg` closure above for
355 // details.
356 (ir::Opcode::Iconst, _) => return,
357 _ => ast::AssignmentRhs::Var,
358 }
359 }
360 _ => ast::AssignmentRhs::Var,
361 };
362
363 non_var_count += match souper_assignment_rhs {
364 ast::AssignmentRhs::Var => 0,
365 _ => 1,
366 };
367 let souper_ty = souper_type_of(&func.dfg, val);
368 let souper_val = lhs.assignment(None, souper_ty, souper_assignment_rhs, vec![]);
369 let old_value = allocs.ir_to_souper_val.insert(val, souper_val);
370 assert!(old_value.is_none());
371 });
372
373 // We end up harvesting a lot of candidates like:
374 //
375 // %0:i32 = var
376 // infer %0
377 //
378 // and
379 //
380 // %0:i32 = var
381 // %1:i32 = var
382 // %2:i32 = add %0, %1
383 //
384 // Both of these are useless. Only actually harvest the candidate if there
385 // are at least two actual operations.
386 if non_var_count >= 2 {
387 let lhs = lhs.finish(allocs.ir_to_souper_val[&val], None);
388 out.send(format!(
389 ";; Harvested from `{}` in `{}`\n{}\n",
390 val, func.name, lhs
391 ))
392 .unwrap();
393 }
394}
395
396fn souper_type_of(dfg: &ir::DataFlowGraph, val: ir::Value) -> Option<ast::Type> {
397 let ty = dfg.value_type(val);
398 assert!(ty.is_int());
399 assert_eq!(ty.lane_count(), 1);
400 let width = ty.bits().try_into().unwrap();
401 Some(ast::Type { width })
402}
403
404#[derive(Debug)]
405enum StackEntry {
406 Visit(ir::Value),
407 Trace(ir::Value),
408}
409
410fn post_order_dfs(
411 allocs: &mut Allocs,
412 dfg: &ir::DataFlowGraph,
413 val: ir::Value,
414 should_trace: impl Fn(ir::Value) -> bool,
415 mut visit: impl FnMut(&mut Allocs, ir::Value),
416) {
417 allocs.dfs_stack.push(StackEntry::Trace(val));
418
419 while let Some(entry) = allocs.dfs_stack.pop() {
420 match entry {
421 StackEntry::Visit(val) => {
422 let is_new = allocs.dfs_seen.insert(val);
423 if is_new {
424 visit(allocs, val);
425 }
426 }
427 StackEntry::Trace(val) => {
428 if allocs.dfs_seen.contains(&val) {
429 continue;
430 }
431
432 allocs.dfs_stack.push(StackEntry::Visit(val));
433 if should_trace(val) {
434 if let ir::ValueDef::Result(inst, 0) = dfg.value_def(val) {
435 let args = dfg.inst_args(inst);
436 for v in args.iter().rev().copied() {
437 allocs.dfs_stack.push(StackEntry::Trace(v));
438 }
439 }
440 }
441 }
442 }
443 }
444}