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::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
533fn 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)]
550enum StackEntry {
551    Visit(ir::Value),
552    Trace(ir::Value),
553}
554
555fn 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}
cranelift_codegen/souper_harvest.rs

cranelift_codegen/
souper_harvest.rs
