Cranelift: Rewrite conditional branches to unconditional traps into conditional traps during legalization (#10988)

* Cranelift: Legalize via a backwards walk, rather than forwards

Note: the test expectation change in `filetests/egraph/misc.clif` is simply
because we happen to change the order in which we created the legalized
`stack_addr` instructions that get GVN'd together, and therefore also changed
their relative value numbering. We dedupe to the value that occurs first in the
function, which is the one inserted into the DFG *after* the other now because
of the backwards traversal, and it therefore has a different value number from
before. The resulting program is identical, modulo value numbering.

* Cranelift: Rewrite conditional branches to unconditional traps into conditional traps during legalization

Given this instruction:

```clif
brif v0, block1, block2
```

If we know that `block1` does nothing but immediately trap then we can rewrite
that `brif` into the following:

```clif
trapz v0, <trapcode>
jump block2
```

(And we can do the equivalent with `trapz` if `block2` immediately traps).

This transformation allows for the conditional trap instructions to be GVN'd and
for our egraphs mid-end to generally better optimize the program. We
additionally have better codegen in our backends for `trapz` than branches to
unconditional traps.

Fixes #10941

* Update CLIF filetests and Wasmtime disas tests

Author: fitzgen

cranelift/codegen/src/context.rs

@@ -168,13 +168,13 @@ impl Context {
             self.func.display()
         );
 
-        self.compute_cfg();
         if isa.flags().enable_nan_canonicalization() {
             self.canonicalize_nans(isa)?;
         }
 
         self.legalize(isa)?;
 
+        self.compute_cfg();
         self.compute_domtree();
         self.eliminate_unreachable_code(isa)?;
         self.remove_constant_phis(isa)?;
@@ -291,6 +291,7 @@ impl Context {
         // TODO: Avoid doing this when legalization doesn't actually mutate the CFG.
         self.domtree.clear();
         self.loop_analysis.clear();
+        self.cfg.clear();
 
         // Run some specific legalizations only.
         simple_legalize(&mut self.func, isa);

cranelift/codegen/src/ir/layout.rs

@@ -516,6 +516,12 @@ impl Layout {
         }
     }
 
+    /// Does the given block contain exactly one instruction?
+    pub fn block_contains_exactly_one_inst(&self, block: Block) -> bool {
+        let block = &self.blocks[block];
+        block.first_inst.is_some() && block.first_inst == block.last_inst
+    }
+
     /// Split the block containing `before` in two.
     ///
     /// Insert `new_block` after the old block and move `before` and the following instructions to

cranelift/codegen/src/legalizer/branch_to_trap.rs

@@ -0,0 +1,86 @@
+//! Rewrite branch-to-unconditional-trap into conditional trap instructions.
+//!
+//! Given this instruction:
+//!
+//! ```clif
+//! brif v0, block1, block2
+//! ```
+//!
+//! If we know that `block1` does nothing but immediately trap then we can
+//! rewrite that `brif` into the following:
+//!
+//! ```clif
+//! trapz v0, <trapcode>
+//! jump block2
+//! ```
+//!
+//! (And we can do the equivalent with `trapz` if `block2` immediately traps).
+//!
+//! This transformation allows for the conditional trap instructions to be GVN'd
+//! and for our egraphs mid-end to generally better optimize the program. We
+//! additionally have better codegen in our backends for `trapz` than branches
+//! to unconditional traps.
+
+use super::*;
+
+#[derive(Default)]
+pub struct BranchToTrap {
+    /// The set of blocks that contain exactly one unconditional trap
+    /// instruction.
+    just_trap_blocks: EntitySet<ir::Block>,
+}
+
+impl BranchToTrap {
+    /// Analyze the given block.
+    ///
+    /// The `block` must be terminated by a `trap` instruction.
+    pub fn analyze_trapping_block(&mut self, func: &ir::Function, block: ir::Block) {
+        if func.layout.block_contains_exactly_one_inst(block) {
+            self.just_trap_blocks.insert(block);
+        }
+    }
+
+    fn just_trap_block_code(&self, func: &ir::Function, block: ir::Block) -> ir::TrapCode {
+        debug_assert!(self.just_trap_blocks.contains(block));
+        debug_assert!(func.layout.block_contains_exactly_one_inst(block));
+        let inst = func.layout.first_inst(block).unwrap();
+        match func.dfg.insts[inst] {
+            InstructionData::Trap { code, .. } => code,
+            _ => unreachable!(),
+        }
+    }
+
+    /// Process a `brif` instruction, potentially performing our rewrite.
+    ///
+    /// The `inst` must be a `brif` containing the given `arg` and `blocks`.
+    pub fn process_brif(
+        &self,
+        func: &mut ir::Function,
+        inst: ir::Inst,
+        arg: ir::Value,
+        blocks: [ir::BlockCall; 2],
+    ) {
+        let consequent = blocks[0].block(&func.dfg.value_lists);
+        let alternative = blocks[1].block(&func.dfg.value_lists);
+
+        if self.just_trap_blocks.contains(consequent) {
+            let mut pos = FuncCursor::new(func);
+            pos.goto_inst(inst);
+
+            let code = self.just_trap_block_code(pos.func, consequent);
+            pos.ins().trapnz(arg, code);
+
+            let args: SmallVec<[_; 8]> = blocks[1].args(&pos.func.dfg.value_lists).collect();
+            pos.func.dfg.replace(inst).jump(alternative, &args);
+        } else if self.just_trap_blocks.contains(alternative) {
+            let mut pos = FuncCursor::new(func);
+            pos.goto_inst(inst);
+
+            let code = self.just_trap_block_code(pos.func, alternative);
+            pos.ins().trapz(arg, code);
+
+            let args: SmallVec<[_; 8]> = blocks[0].args(&pos.func.dfg.value_lists).collect();
+            pos.func.dfg.replace(inst).jump(consequent, &args);
+        }
+    }
+}

cranelift/codegen/src/legalizer/mod.rs

@@ -19,9 +19,13 @@ use crate::ir::types::{self, I64, I128};
 use crate::ir::{self, InstBuilder, InstructionData, MemFlags, Value};
 use crate::isa::TargetIsa;
 use crate::trace;
+use cranelift_entity::EntitySet;
+use smallvec::SmallVec;
 
+mod branch_to_trap;
 mod globalvalue;
 
+use self::branch_to_trap::BranchToTrap;
 use self::globalvalue::expand_global_value;
 
 fn imm_const(pos: &mut FuncCursor, arg: Value, imm: Imm64, is_signed: bool) -> Value {
@@ -51,34 +55,34 @@ fn imm_const(pos: &mut FuncCursor, arg: Value, imm: Imm64, is_signed: bool) -> V
 
 /// A command describing how the walk over instructions should proceed.
 enum WalkCommand {
-    /// Continue to the next instruction, if any.
+    /// Continue walking to the next instruction, if any.
     Continue,
     /// Revisit the current instruction (presumably because it was legalized
     /// into a new instruction that may also require further legalization).
     Revisit,
 }
 
-/// A simple, naive forwards walk over every instruction in every block in the
+/// A simple, naive backwards walk over every instruction in every block in the
 /// function's layout.
 ///
-/// This does not guarantee any kind of post-order visitation or anything like
-/// that, it is just iterating over blocks in layout order, not any kind of
-/// control-flow graph visitation order.
+/// This does not guarantee any kind of reverse post-order visitation or
+/// anything like that, it is just iterating over blocks in reverse layout
+/// order, not any kind of control-flow graph visitation order.
 ///
 /// The `f` visitor closure controls how the walk proceeds via its `WalkCommand`
 /// result.
-fn forward_walk(
+fn backward_walk(
     func: &mut ir::Function,
-    mut f: impl FnMut(&mut ir::Function, ir::Inst) -> WalkCommand,
+    mut f: impl FnMut(&mut ir::Function, ir::Block, ir::Inst) -> WalkCommand,
 ) {
     let mut pos = FuncCursor::new(func);
-    while let Some(_block) = pos.next_block() {
+    while let Some(block) = pos.prev_block() {
         let mut prev_pos;
         while let Some(inst) = {
             prev_pos = pos.position();
-            pos.next_inst()
+            pos.prev_inst()
         } {
-            match f(pos.func, inst) {
+            match f(pos.func, block, inst) {
                 WalkCommand::Continue => continue,
                 WalkCommand::Revisit => pos.set_position(prev_pos),
             }
@@ -91,7 +95,31 @@ fn forward_walk(
 pub fn simple_legalize(func: &mut ir::Function, isa: &dyn TargetIsa) {
     trace!("Pre-legalization function:\n{}", func.display());
 
-    forward_walk(func, |func, inst| match func.dfg.insts[inst] {
+    let mut branch_to_trap = BranchToTrap::default();
+
+    // We walk the IR backwards because in practice, given the way that
+    // frontends tend to produce CLIF, this means we will visit in roughly
+    // reverse post order, which is helpful for getting the most optimizations
+    // out of the `branch-to-trap` pass that we can (it must analyze trapping
+    // blocks before it can rewrite branches to them) but the order does not
+    // actually affect correctness.
+    backward_walk(func, |func, block, inst| match func.dfg.insts[inst] {
+        InstructionData::Trap {
+            opcode: ir::Opcode::Trap,
+            code: _,
+        } => {
+            branch_to_trap.analyze_trapping_block(func, block);
+            WalkCommand::Continue
+        }
+        InstructionData::Brif {
+            opcode: ir::Opcode::Brif,
+            arg,
+            blocks,
+        } => {
+            branch_to_trap.process_brif(func, inst, arg, blocks);
+            WalkCommand::Continue
+        }
+
         InstructionData::UnaryGlobalValue {
             opcode: ir::Opcode::GlobalValue,
             global_value,

cranelift/filetests/filetests/egraph/misc.clif

@@ -14,7 +14,7 @@ block0(v0: i64):
 ; check: function %stack_load(i64) -> i64 fast {
 ; nextln:    ss0 = explicit_slot 8
 ; check:  block0(v0: i64):
-; nextln:     v2 = stack_addr.i64 ss0
-; nextln:     store notrap v0, v2
+; nextln:     v3 = stack_addr.i64 ss0
+; nextln:     store notrap v0, v3
 ; nextln:     return v0
 ; nextln: }

cranelift/filetests/filetests/isa/x64/branches.clif

@@ -638,15 +638,12 @@ block202:
 ;   ucomiss const(1), %xmm0
 ;   jp,nz   label2; j label1
 ; block1:
-;   jmp     label5
+;   jmp     label3
 ; block2:
 ;   ucomiss const(0), %xmm0
-;   jnp     label4; j label3
+;   jnp #trap=heap_oob
+;   jmp     label3
 ; block3:
-;   jmp     label5
-; block4:
-;   ud2 heap_oob
-; block5:
 ;   movq    %rbp, %rsp
 ;   popq %rbp
 ;   ret
@@ -658,16 +655,15 @@ block202:
 ; block1: ; offset 0x4
 ;   ucomiss 0x25(%rip), %xmm0
 ;   jp 0x17
-;   je 0x26
+;   je 0x24
 ; block2: ; offset 0x17
 ;   ucomiss 0x22(%rip), %xmm0
-;   jp 0x26
+;   jnp 0x29
 ; block3: ; offset 0x24
-;   ud2 ; trap: heap_oob
-; block4: ; offset 0x26
 ;   movq %rbp, %rsp
 ;   popq %rbp
 ;   retq
+;   ud2 ; trap: heap_oob
 ;   addb %al, (%rax)
 ;   addb %al, (%rax)
 ;   addb %al, (%rax)

cranelift/filetests/filetests/legalizer/branch-to-unconditional-trap.clif

@@ -0,0 +1,41 @@
+test legalizer
+target aarch64
+
+function %trapnz(i64, i64) -> i64 {
+block0(v0: i64, v1: i64):
+    brif v0, block1, block2(v1)
+block1:
+    trap user42
+block2(v2: i64):
+    return v2
+}
+
+; check:  block0(v0: i64, v1: i64):
+; nextln:     trapnz v0, user42
+; nextln:     jump block2(v1)
+
+function %trapz(i64, i64) -> i64 {
+block0(v0: i64, v1: i64):
+    brif v0, block1(v1), block2
+block1(v2: i64):
+    return v2
+block2:
+    trap user42
+}
+
+; check:  block0(v0: i64, v1: i64):
+; nextln:     trapz v0, user42
+; nextln:     jump block1(v1)
+
+function %could_be_either_doesnt_matter(i64) -> i64 {
+block0(v0: i64):
+    brif v0, block1, block2
+block1:
+    trap user36
+block2:
+    trap user42
+}
+
+; check:  block0(v0: i64):
+; nextln:     trapnz v0, user36
+; nextln:     jump block2

tests/disas/arith.wat

@@ -25,10 +25,8 @@
 ;; @0021                               v3 = iconst.i32 4
 ;; @0023                               v4 = iconst.i32 4
 ;; @0025                               v5 = isub v3, v4  ; v3 = 4, v4 = 4
-;; @002a                               brif v5, block2, block4
-;;
-;;                                 block2:
-;; @002c                               trap user11
+;;                                     trapnz v5, user11
+;; @002a                               jump block4
 ;;
 ;;                                 block4:
 ;; @002e                               v6 = iconst.i32 6

tests/disas/component-model/enum.wat

@@ -50,34 +50,31 @@
 ;; @0061                               v3 = iconst.i32 0
 ;; @0068                               v8 = icmp eq v7, v3  ; v3 = 0
 ;; @0068                               v9 = uextend.i32 v8
-;; @0069                               brif v9, block2, block3
-;;
-;;                                 block2:
-;; @006b                               trap user11
+;;                                     trapnz v9, user11
+;; @0069                               jump block3
 ;;
 ;;                                 block3:
 ;; @006d                               v10 = load.i64 notrap aligned readonly can_move v0+80
 ;; @006d                               v11 = load.i32 notrap aligned table v10
 ;; @006f                               v12 = iconst.i32 2
 ;; @0071                               v13 = band v11, v12  ; v12 = 2
-;;                                     v82 = iconst.i32 0
-;;                                     v83 = icmp eq v13, v82  ; v82 = 0
-;; @0072                               v15 = uextend.i32 v83
-;; @0073                               brif v15, block4, block5
-;;
-;;                                 block4:
-;; @0075                               trap user11
+;;                                     v80 = iconst.i32 0
+;;                                     v81 = icmp eq v13, v80  ; v80 = 0
+;; @0072                               v15 = uextend.i32 v81
+;;                                     trapnz v15, user11
+;; @0073                               jump block5
 ;;
 ;;                                 block5:
+;; @0077                               v17 = load.i32 notrap aligned table v10
 ;; @0079                               v18 = iconst.i32 -3
-;; @007b                               v19 = band.i32 v11, v18  ; v18 = -3
+;; @007b                               v19 = band v17, v18  ; v18 = -3
 ;; @007c                               store notrap aligned table v19, v10
 ;;                                     v70 = iconst.i32 -4
-;;                                     v76 = band.i32 v11, v70  ; v70 = -4
+;;                                     v76 = band v17, v70  ; v70 = -4
 ;; @0083                               store notrap aligned table v76, v10
-;;                                     v84 = iconst.i32 1
-;;                                     v85 = bor v19, v84  ; v84 = 1
-;; @008a                               store notrap aligned table v85, v10
+;;                                     v82 = iconst.i32 1
+;;                                     v83 = bor v19, v82  ; v82 = 1
+;; @008a                               store notrap aligned table v83, v10
 ;; @008c                               v32 = load.i64 notrap aligned readonly can_move v0+56
 ;; @008c                               v33 = load.i64 notrap aligned readonly can_move v0+72
 ;; @008c                               v34 = call_indirect sig0, v32(v33, v0)
@@ -88,19 +85,18 @@
 ;; @009b                               v40 = iconst.i32 3
 ;; @009d                               v41 = icmp ugt v34, v40  ; v40 = 3
 ;; @009d                               v42 = uextend.i32 v41
-;; @009e                               brif v42, block6, block7
-;;
-;;                                 block6:
-;; @00a0                               trap user11
+;;                                     trapnz v42, user11
+;; @009e                               jump block7
 ;;
 ;;                                 block7:
-;;                                     v86 = iconst.i32 1
-;;                                     v87 = bor.i32 v36, v86  ; v86 = 1
-;; @00a9                               store notrap aligned table v87, v4
+;; @00a4                               v44 = load.i32 notrap aligned table v4
+;;                                     v84 = iconst.i32 1
+;;                                     v85 = bor v44, v84  ; v84 = 1
+;; @00a9                               store notrap aligned table v85, v4
 ;; @00ab                               v49 = load.i32 notrap aligned table v10
-;;                                     v88 = iconst.i32 2
-;;                                     v89 = bor v49, v88  ; v88 = 2
-;; @00b0                               store notrap aligned table v89, v10
+;;                                     v86 = iconst.i32 2
+;;                                     v87 = bor v49, v86  ; v86 = 2
+;; @00b0                               store notrap aligned table v87, v10
 ;; @00b2                               jump block1
 ;;
 ;;                                 block1: