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 
27 use crate::ir;
28 use souper_ir::ast;
29 use std::collections::{HashMap, HashSet};
30 use std::string::String;
31 use std::sync::mpsc;
32 use 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.
do_souper_harvest(func: &ir::Function, out: &mut mpsc::Sender<String>)37 pub 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)]
60 struct 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 
72 impl Allocs {
73     /// Reset the collections to their empty state (without deallocating their
74     /// backing data).
reset(&mut self)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.
harvest_candidate_lhs( allocs: &mut Allocs, func: &ir::Function, val: ir::Value, out: &mut mpsc::Sender<String>, )83 fn 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::IaddImm
99                 | ir::Opcode::IrsubImm
100                 | ir::Opcode::Imul
101                 | ir::Opcode::ImulImm
102                 | ir::Opcode::Udiv
103                 | ir::Opcode::UdivImm
104                 | ir::Opcode::Sdiv
105                 | ir::Opcode::SdivImm
106                 | ir::Opcode::Urem
107                 | ir::Opcode::UremImm
108                 | ir::Opcode::Srem
109                 | ir::Opcode::SremImm
110                 | ir::Opcode::Band
111                 | ir::Opcode::BandImm
112                 | ir::Opcode::Bor
113                 | ir::Opcode::BorImm
114                 | ir::Opcode::Bxor
115                 | ir::Opcode::BxorImm
116                 | ir::Opcode::Ishl
117                 | ir::Opcode::IshlImm
118                 | ir::Opcode::Sshr
119                 | ir::Opcode::SshrImm
120                 | ir::Opcode::Ushr
121                 | ir::Opcode::UshrImm
122                 | ir::Opcode::Select
123                 | ir::Opcode::Uextend
124                 | ir::Opcode::Sextend
125                 | ir::Opcode::Trunc
126                 | ir::Opcode::Icmp
127                 | ir::Opcode::Popcnt
128                 | ir::Opcode::Bitrev
129                 | ir::Opcode::Clz
130                 | ir::Opcode::Ctz
131                 // TODO: ir::Opcode::IaddCarry
132                 | ir::Opcode::SaddSat
133                 | ir::Opcode::SsubSat
134                 | ir::Opcode::UsubSat => true,
135                 _ => false,
136             },
137         _ => false,
138     };
139 
140     post_order_dfs(allocs, &func.dfg, val, should_trace, |allocs, val| {
141         let souper_assignment_rhs = match func.dfg.value_def(val) {
142             ir::ValueDef::Result(inst, 0) => {
143                 let args = func.dfg.inst_args(inst);
144 
145                 // Get the n^th argument as a souper operand.
146                 let arg = |allocs: &mut Allocs, n| {
147                     let arg = args[n];
148                     if let Some(a) = allocs.ir_to_souper_val.get(&arg).copied() {
149                         a.into()
150                     } else {
151                         // The only arguments we get that we haven't already
152                         // converted into a souper instruction are `iconst`s.
153                         // This is because souper only allows
154                         // constants as operands, and it doesn't allow assigning
155                         // constants to a variable name. So we lazily convert
156                         // `iconst`s into souper operands here,
157                         // when they are actually used.
158                         match func.dfg.value_def(arg) {
159                             ir::ValueDef::Result(inst, 0) => match func.dfg.insts[inst] {
160                                 ir::InstructionData::UnaryImm { opcode, imm } => {
161                                     debug_assert_eq!(opcode, ir::Opcode::Iconst);
162                                     let imm: i64 = imm.into();
163                                     ast::Operand::Constant(ast::Constant {
164                                         value: imm.into(),
165                                         r#type: souper_type_of(&func.dfg, arg),
166                                     })
167                                 }
168                                 _ => unreachable!(
169                                     "only iconst instructions \
170                                      aren't in `ir_to_souper_val`"
171                                 ),
172                             },
173                             _ => unreachable!(
174                                 "only iconst instructions \
175                                  aren't in `ir_to_souper_val`"
176                             ),
177                         }
178                     }
179                 };
180 
181                 match (func.dfg.insts[inst].opcode(), &func.dfg.insts[inst]) {
182                     (ir::Opcode::Iadd, _) => {
183                         let a = arg(allocs, 0);
184                         let b = arg(allocs, 1);
185                         ast::Instruction::Add { a, b }.into()
186                     }
187                     (ir::Opcode::IaddImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
188                         let a = arg(allocs, 0);
189                         let value: i64 = (*imm).into();
190                         let value: i128 = value.into();
191                         let b = ast::Constant {
192                             value,
193                             r#type: souper_type_of(&func.dfg, val),
194                         }
195                         .into();
196                         ast::Instruction::Add { a, b }.into()
197                     }
198                     (ir::Opcode::IrsubImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
199                         let b = arg(allocs, 0);
200                         let value: i64 = (*imm).into();
201                         let value: i128 = value.into();
202                         let a = ast::Constant {
203                             value,
204                             r#type: souper_type_of(&func.dfg, val),
205                         }
206                         .into();
207                         ast::Instruction::Sub { a, b }.into()
208                     }
209                     (ir::Opcode::Imul, _) => {
210                         let a = arg(allocs, 0);
211                         let b = arg(allocs, 1);
212                         ast::Instruction::Mul { a, b }.into()
213                     }
214                     (ir::Opcode::ImulImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
215                         let a = arg(allocs, 0);
216                         let value: i64 = (*imm).into();
217                         let value: i128 = value.into();
218                         let b = ast::Constant {
219                             value,
220                             r#type: souper_type_of(&func.dfg, val),
221                         }
222                         .into();
223                         ast::Instruction::Mul { a, b }.into()
224                     }
225                     (ir::Opcode::Udiv, _) => {
226                         let a = arg(allocs, 0);
227                         let b = arg(allocs, 1);
228                         ast::Instruction::Udiv { a, b }.into()
229                     }
230                     (ir::Opcode::UdivImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
231                         let a = arg(allocs, 0);
232                         let value: i64 = (*imm).into();
233                         let value: i128 = value.into();
234                         let b = ast::Constant {
235                             value,
236                             r#type: souper_type_of(&func.dfg, val),
237                         }
238                         .into();
239                         ast::Instruction::Udiv { a, b }.into()
240                     }
241                     (ir::Opcode::Sdiv, _) => {
242                         let a = arg(allocs, 0);
243                         let b = arg(allocs, 1);
244                         ast::Instruction::Sdiv { a, b }.into()
245                     }
246                     (ir::Opcode::SdivImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
247                         let a = arg(allocs, 0);
248                         let value: i64 = (*imm).into();
249                         let value: i128 = value.into();
250                         let b = ast::Constant {
251                             value,
252                             r#type: souper_type_of(&func.dfg, val),
253                         }
254                         .into();
255                         ast::Instruction::Sdiv { a, b }.into()
256                     }
257                     (ir::Opcode::Urem, _) => {
258                         let a = arg(allocs, 0);
259                         let b = arg(allocs, 1);
260                         ast::Instruction::Urem { a, b }.into()
261                     }
262                     (ir::Opcode::UremImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
263                         let a = arg(allocs, 0);
264                         let value: i64 = (*imm).into();
265                         let value: i128 = value.into();
266                         let b = ast::Constant {
267                             value,
268                             r#type: souper_type_of(&func.dfg, val),
269                         }
270                         .into();
271                         ast::Instruction::Urem { a, b }.into()
272                     }
273                     (ir::Opcode::Srem, _) => {
274                         let a = arg(allocs, 0);
275                         let b = arg(allocs, 1);
276                         ast::Instruction::Srem { a, b }.into()
277                     }
278                     (ir::Opcode::SremImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
279                         let a = arg(allocs, 0);
280                         let value: i64 = (*imm).into();
281                         let value: i128 = value.into();
282                         let b = ast::Constant {
283                             value,
284                             r#type: souper_type_of(&func.dfg, val),
285                         }
286                         .into();
287                         ast::Instruction::Srem { a, b }.into()
288                     }
289                     (ir::Opcode::Band, _) => {
290                         let a = arg(allocs, 0);
291                         let b = arg(allocs, 1);
292                         ast::Instruction::And { a, b }.into()
293                     }
294                     (ir::Opcode::BandImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
295                         let a = arg(allocs, 0);
296                         let value: i64 = (*imm).into();
297                         let value: i128 = value.into();
298                         let b = ast::Constant {
299                             value,
300                             r#type: souper_type_of(&func.dfg, val),
301                         }
302                         .into();
303                         ast::Instruction::And { a, b }.into()
304                     }
305                     (ir::Opcode::Bor, _) => {
306                         let a = arg(allocs, 0);
307                         let b = arg(allocs, 1);
308                         ast::Instruction::Or { a, b }.into()
309                     }
310                     (ir::Opcode::BorImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
311                         let a = arg(allocs, 0);
312                         let value: i64 = (*imm).into();
313                         let value: i128 = value.into();
314                         let b = ast::Constant {
315                             value,
316                             r#type: souper_type_of(&func.dfg, val),
317                         }
318                         .into();
319                         ast::Instruction::Or { a, b }.into()
320                     }
321                     (ir::Opcode::Bxor, _) => {
322                         let a = arg(allocs, 0);
323                         let b = arg(allocs, 1);
324                         ast::Instruction::Xor { a, b }.into()
325                     }
326                     (ir::Opcode::BxorImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
327                         let a = arg(allocs, 0);
328                         let value: i64 = (*imm).into();
329                         let value: i128 = value.into();
330                         let b = ast::Constant {
331                             value,
332                             r#type: souper_type_of(&func.dfg, val),
333                         }
334                         .into();
335                         ast::Instruction::Xor { a, b }.into()
336                     }
337                     (ir::Opcode::Ishl, _) => {
338                         let a = arg(allocs, 0);
339                         let b = arg(allocs, 1);
340                         ast::Instruction::Shl { a, b }.into()
341                     }
342                     (ir::Opcode::IshlImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
343                         let a = arg(allocs, 0);
344                         let value: i64 = (*imm).into();
345                         let value: i128 = value.into();
346                         let b = ast::Constant {
347                             value,
348                             r#type: souper_type_of(&func.dfg, val),
349                         }
350                         .into();
351                         ast::Instruction::Shl { a, b }.into()
352                     }
353                     (ir::Opcode::Sshr, _) => {
354                         let a = arg(allocs, 0);
355                         let b = arg(allocs, 1);
356                         ast::Instruction::Ashr { a, b }.into()
357                     }
358                     (ir::Opcode::SshrImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
359                         let a = arg(allocs, 0);
360                         let value: i64 = (*imm).into();
361                         let value: i128 = value.into();
362                         let b = ast::Constant {
363                             value,
364                             r#type: souper_type_of(&func.dfg, val),
365                         }
366                         .into();
367                         ast::Instruction::Ashr { a, b }.into()
368                     }
369                     (ir::Opcode::Ushr, _) => {
370                         let a = arg(allocs, 0);
371                         let b = arg(allocs, 1);
372                         ast::Instruction::Lshr { a, b }.into()
373                     }
374                     (ir::Opcode::UshrImm, ir::InstructionData::BinaryImm64 { imm, .. }) => {
375                         let a = arg(allocs, 0);
376                         let value: i64 = (*imm).into();
377                         let value: i128 = value.into();
378                         let b = ast::Constant {
379                             value,
380                             r#type: souper_type_of(&func.dfg, val),
381                         }
382                         .into();
383                         ast::Instruction::Lshr { a, b }.into()
384                     }
385                     (ir::Opcode::Select, _) => {
386                         let a = arg(allocs, 0);
387 
388                         // While Cranelift allows any width condition for
389                         // `select` and checks it against `0`, Souper requires
390                         // an `i1`. So insert a `ne %x, 0` as needed.
391                         let a = match a {
392                             ast::Operand::Value(id) => match lhs.get_value(id).r#type {
393                                 Some(ast::Type { width: 1 }) => a,
394                                 _ => lhs
395                                     .assignment(
396                                         None,
397                                         Some(ast::Type { width: 1 }),
398                                         ast::Instruction::Ne {
399                                             a,
400                                             b: ast::Constant {
401                                                 value: 0,
402                                                 r#type: None,
403                                             }
404                                             .into(),
405                                         },
406                                         vec![],
407                                     )
408                                     .into(),
409                             },
410                             ast::Operand::Constant(ast::Constant { value, .. }) => ast::Constant {
411                                 value: (value != 0) as _,
412                                 r#type: Some(ast::Type { width: 1 }),
413                             }
414                             .into(),
415                         };
416 
417                         let b = arg(allocs, 1);
418                         let c = arg(allocs, 2);
419                         ast::Instruction::Select { a, b, c }.into()
420                     }
421                     (ir::Opcode::Uextend, _) => {
422                         let a = arg(allocs, 0);
423                         ast::Instruction::Zext { a }.into()
424                     }
425                     (ir::Opcode::Sextend, _) => {
426                         let a = arg(allocs, 0);
427                         ast::Instruction::Sext { a }.into()
428                     }
429                     (ir::Opcode::Trunc, _) => {
430                         let a = arg(allocs, 0);
431                         ast::Instruction::Trunc { a }.into()
432                     }
433                     (ir::Opcode::Icmp, ir::InstructionData::IntCompare { cond, .. })
434                     | (ir::Opcode::IcmpImm, ir::InstructionData::IntCompare { cond, .. }) => {
435                         let a = arg(allocs, 0);
436                         let b = arg(allocs, 1);
437                         match cond {
438                             ir::condcodes::IntCC::Equal => ast::Instruction::Eq { a, b }.into(),
439                             ir::condcodes::IntCC::NotEqual => ast::Instruction::Ne { a, b }.into(),
440                             ir::condcodes::IntCC::UnsignedLessThan => {
441                                 ast::Instruction::Ult { a, b }.into()
442                             }
443                             ir::condcodes::IntCC::SignedLessThan => {
444                                 ast::Instruction::Slt { a, b }.into()
445                             }
446                             ir::condcodes::IntCC::UnsignedLessThanOrEqual => {
447                                 ast::Instruction::Sle { a, b }.into()
448                             }
449                             ir::condcodes::IntCC::SignedLessThanOrEqual => {
450                                 ast::Instruction::Sle { a, b }.into()
451                             }
452                             _ => ast::AssignmentRhs::Var,
453                         }
454                     }
455                     (ir::Opcode::Popcnt, _) => {
456                         let a = arg(allocs, 0);
457                         ast::Instruction::Ctpop { a }.into()
458                     }
459                     (ir::Opcode::Bitrev, _) => {
460                         let a = arg(allocs, 0);
461                         ast::Instruction::BitReverse { a }.into()
462                     }
463                     (ir::Opcode::Clz, _) => {
464                         let a = arg(allocs, 0);
465                         ast::Instruction::Ctlz { a }.into()
466                     }
467                     (ir::Opcode::Ctz, _) => {
468                         let a = arg(allocs, 0);
469                         ast::Instruction::Cttz { a }.into()
470                     }
471                     // TODO: ir::Opcode::IaddCarry
472                     (ir::Opcode::SaddSat, _) => {
473                         let a = arg(allocs, 0);
474                         let b = arg(allocs, 1);
475                         ast::Instruction::SaddSat { a, b }.into()
476                     }
477                     (ir::Opcode::SsubSat, _) => {
478                         let a = arg(allocs, 0);
479                         let b = arg(allocs, 1);
480                         ast::Instruction::SsubSat { a, b }.into()
481                     }
482                     (ir::Opcode::UsubSat, _) => {
483                         let a = arg(allocs, 0);
484                         let b = arg(allocs, 1);
485                         ast::Instruction::UsubSat { a, b }.into()
486                     }
487                     // Because Souper doesn't allow constants to be on the right
488                     // hand side of an assignment (i.e. `%0:i32 = 1234` is
489                     // disallowed) we have to ignore `iconst`
490                     // instructions until we process them as operands for some
491                     // other instruction. See the `arg` closure above for
492                     // details.
493                     (ir::Opcode::Iconst, _) => return,
494                     _ => ast::AssignmentRhs::Var,
495                 }
496             }
497             _ => ast::AssignmentRhs::Var,
498         };
499 
500         non_var_count += match souper_assignment_rhs {
501             ast::AssignmentRhs::Var => 0,
502             _ => 1,
503         };
504         let souper_ty = souper_type_of(&func.dfg, val);
505         let souper_val = lhs.assignment(None, souper_ty, souper_assignment_rhs, vec![]);
506         let old_value = allocs.ir_to_souper_val.insert(val, souper_val);
507         assert!(old_value.is_none());
508     });
509 
510     // We end up harvesting a lot of candidates like:
511     //
512     //     %0:i32 = var
513     //     infer %0
514     //
515     // and
516     //
517     //     %0:i32 = var
518     //     %1:i32 = var
519     //     %2:i32 = add %0, %1
520     //
521     // Both of these are useless. Only actually harvest the candidate if there
522     // are at least two actual operations.
523     if non_var_count >= 2 {
524         let lhs = lhs.finish(allocs.ir_to_souper_val[&val], None);
525         out.send(format!(
526             ";; Harvested from `{}` in `{}`\n{}\n",
527             val, func.name, lhs
528         ))
529         .unwrap();
530     }
531 }
532 
souper_type_of(dfg: &ir::DataFlowGraph, val: ir::Value) -> Option<ast::Type>533 fn souper_type_of(dfg: &ir::DataFlowGraph, val: ir::Value) -> Option<ast::Type> {
534     let ty = dfg.value_type(val);
535     assert!(ty.is_int());
536     assert_eq!(ty.lane_count(), 1);
537     let width = match dfg.value_def(val).inst() {
538         Some(inst)
539             if dfg.insts[inst].opcode() == ir::Opcode::IcmpImm
540                 || dfg.insts[inst].opcode() == ir::Opcode::Icmp =>
541         {
542             1
543         }
544         _ => ty.bits().try_into().unwrap(),
545     };
546     Some(ast::Type { width })
547 }
548 
549 #[derive(Debug)]
550 enum StackEntry {
551     Visit(ir::Value),
552     Trace(ir::Value),
553 }
554 
post_order_dfs( allocs: &mut Allocs, dfg: &ir::DataFlowGraph, val: ir::Value, should_trace: impl Fn(ir::Value) -> bool, mut visit: impl FnMut(&mut Allocs, ir::Value), )555 fn post_order_dfs(
556     allocs: &mut Allocs,
557     dfg: &ir::DataFlowGraph,
558     val: ir::Value,
559     should_trace: impl Fn(ir::Value) -> bool,
560     mut visit: impl FnMut(&mut Allocs, ir::Value),
561 ) {
562     allocs.dfs_stack.push(StackEntry::Trace(val));
563 
564     while let Some(entry) = allocs.dfs_stack.pop() {
565         match entry {
566             StackEntry::Visit(val) => {
567                 let is_new = allocs.dfs_seen.insert(val);
568                 if is_new {
569                     visit(allocs, val);
570                 }
571             }
572             StackEntry::Trace(val) => {
573                 if allocs.dfs_seen.contains(&val) {
574                     continue;
575                 }
576 
577                 allocs.dfs_stack.push(StackEntry::Visit(val));
578                 if should_trace(val) {
579                     if let ir::ValueDef::Result(inst, 0) = dfg.value_def(val) {
580                         let args = dfg.inst_args(inst);
581                         for v in args.iter().rev().copied() {
582                             allocs.dfs_stack.push(StackEntry::Trace(v));
583                         }
584                     }
585                 }
586             }
587         }
588     }
589 }
590