cranelift: Move blockparams to critical edge blocks

This lets us stop allocating temporary VRegs for critical edges that
have block parameters. That makes the register allocation problem a
little smaller, and also allows reusing lower_branch_blockparam_args for
all block parameters.

Fixes bytecodealliance#7639, and unblocks bytecodealliance/regalloc2#170

Author: jameysharp

cranelift/codegen/src/machinst/lower.rs

@@ -950,33 +950,46 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
         let loc = self.srcloc(branch);
         self.finish_ir_inst(loc);
         // Add block param outputs for current block.
-        self.lower_branch_blockparam_args(bindex);
+        match self.vcode.block_order().succ_indices(bindex).1 {
+            &[succ] => self.lower_branch_blockparam_args(branch, 0, succ),
+            succs => {
+                // If there are multiple edges, either the destination has
+                // multiple predecessors and there is a split critical edge
+                // block that can hold these block params, or it has one
+                // predecessor and doesn't need block params.
+                debug_assert!(succs
+                    .iter()
+                    .zip(self.f.dfg.insts[branch].branch_destination(&self.f.dfg.jump_tables))
+                    .all(|(succ, block_call)| {
+                        matches!(
+                            self.vcode.block_order().lowered_order()[succ.index()],
+                            LoweredBlock::CriticalEdge { .. }
+                        ) || block_call.args_slice(&self.f.dfg.value_lists).is_empty()
+                    }));
+                let succs: SmallVec<[BlockIndex; 2]> = SmallVec::from_slice(succs);
+                for succ in succs {
+                    self.vcode.add_succ(succ, &[]);
+                }
+            }
+        }
         Ok(())
     }
 
-    fn lower_branch_blockparam_args(&mut self, block: BlockIndex) {
-        // TODO: why not make `block_order` public?
-        for succ_idx in 0..self.vcode.block_order().succ_indices(block).1.len() {
-            // Avoid immutable borrow by explicitly indexing.
-            let (opt_inst, succs) = self.vcode.block_order().succ_indices(block);
-            let inst = opt_inst.expect("lower_branch_blockparam_args called on a critical edge!");
-            let succ = succs[succ_idx];
-
-            // The use of `succ_idx` to index `branch_destination` is valid on the assumption that
-            // the traversal order defined in `visit_block_succs` mirrors the order returned by
-            // `branch_destination`. If that assumption is violated, the branch targets returned
-            // here will not match the clif.
-            let branches = self.f.dfg.insts[inst].branch_destination(&self.f.dfg.jump_tables);
-            let branch_args = branches[succ_idx].args_slice(&self.f.dfg.value_lists);
-
-            let mut branch_arg_vregs: SmallVec<[Reg; 16]> = smallvec![];
-            for &arg in branch_args {
-                let arg = self.f.dfg.resolve_aliases(arg);
-                let regs = self.put_value_in_regs(arg);
-                branch_arg_vregs.extend_from_slice(regs.regs());
-            }
-            self.vcode.add_succ(succ, &branch_arg_vregs[..]);
+    fn lower_branch_blockparam_args(&mut self, inst: Inst, succ_idx: usize, succ: BlockIndex) {
+        // The use of `succ_idx` to index `branch_destination` is valid on the assumption that
+        // the traversal order defined in `visit_block_succs` mirrors the order returned by
+        // `branch_destination`. If that assumption is violated, the branch targets returned
+        // here will not match the clif.
+        let branches = self.f.dfg.insts[inst].branch_destination(&self.f.dfg.jump_tables);
+        let branch_args = branches[succ_idx].args_slice(&self.f.dfg.value_lists);
+
+        let mut branch_arg_vregs: SmallVec<[Reg; 16]> = smallvec![];
+        for &arg in branch_args {
+            let arg = self.f.dfg.resolve_aliases(arg);
+            let regs = self.put_value_in_regs(arg);
+            branch_arg_vregs.extend_from_slice(regs.regs());
         }
+        self.vcode.add_succ(succ, &branch_arg_vregs[..]);
     }
 
     fn collect_branches_and_targets(
@@ -1028,44 +1041,29 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
             // `lower_clif_block()` for rationale).
 
             // End branches.
-            if let Some(bb) = lb.orig_block() {
-                if let Some(branch) = self.collect_branches_and_targets(bindex, bb, &mut targets) {
-                    self.lower_clif_branches(backend, bindex, bb, branch, &targets)?;
-                    self.finish_ir_inst(self.srcloc(branch));
-                }
-            } else {
-                // If no orig block, this must be a pure edge block;
-                // get the successor and emit a jump. Add block params
-                // according to the one successor, and pass them
-                // through; note that the successor must have an
-                // original block.
-                let (_, succs) = self.vcode.block_order().succ_indices(bindex);
-                let succ = succs[0];
-
-                let orig_succ = lowered_order[succ.index()];
-                let orig_succ = orig_succ
-                    .orig_block()
-                    .expect("Edge block succ must be body block");
-
-                let mut branch_arg_vregs: SmallVec<[Reg; 16]> = smallvec![];
-                for ty in self.f.dfg.block_param_types(orig_succ) {
-                    let regs = self.vregs.alloc(ty)?;
-                    for &reg in regs.regs() {
-                        branch_arg_vregs.push(reg);
-                        let vreg = reg.to_virtual_reg().unwrap();
-                        self.vcode.add_block_param(vreg);
+            match lb {
+                &LoweredBlock::Orig { block: bb } => {
+                    if let Some(branch) =
+                        self.collect_branches_and_targets(bindex, bb, &mut targets)
+                    {
+                        self.lower_clif_branches(backend, bindex, bb, branch, &targets)?;
+                        self.finish_ir_inst(self.srcloc(branch));
                     }
-                }
-                self.vcode.add_succ(succ, &branch_arg_vregs[..]);
 
-                self.emit(I::gen_jump(MachLabel::from_block(succ)));
-                self.finish_ir_inst(Default::default());
-            }
-
-            // Original block body.
-            if let Some(bb) = lb.orig_block() {
-                self.lower_clif_block(backend, bb, ctrl_plane)?;
-                self.emit_value_label_markers_for_block_args(bb);
+                    // Original block body.
+                    self.lower_clif_block(backend, bb, ctrl_plane)?;
+                    self.emit_value_label_markers_for_block_args(bb);
+                }
+                &LoweredBlock::CriticalEdge { pred, succ_idx, .. } => {
+                    // Emit a jump to the successor, placing the block params
+                    // that the predecessor was going to pass along here.
+                    let (_, succs) = self.vcode.block_order().succ_indices(bindex);
+                    let succ = succs[0];
+                    let branch = self.f.layout.last_inst(pred).unwrap();
+                    self.lower_branch_blockparam_args(branch, succ_idx as usize, succ);
+                    self.emit(I::gen_jump(MachLabel::from_block(succ)));
+                    self.finish_ir_inst(Default::default());
+                }
             }
 
             if bindex.index() == 0 {

cranelift/filetests/filetests/isa/aarch64/cold-blocks.clif

@@ -16,8 +16,8 @@ block2:
 
 ; VCode:
 ; block0:
-;   mov w5, w0
-;   cbnz x5, label1 ; b label2
+;   mov w4, w0
+;   cbnz x4, label1 ; b label2
 ; block1:
 ;   b label3
 ; block2:
@@ -28,8 +28,8 @@ block2:
 ;
 ; Disassembled:
 ; block0: ; offset 0x0
-;   mov w5, w0
-;   cbnz x5, #0xc
+;   mov w4, w0
+;   cbnz x4, #0xc
 ; block1: ; offset 0x8
 ;   mov w0, #0x61
 ; block2: ; offset 0xc
@@ -49,8 +49,8 @@ block2 cold:
 
 ; VCode:
 ; block0:
-;   mov w5, w0
-;   cbnz x5, label1 ; b label2
+;   mov w4, w0
+;   cbnz x4, label1 ; b label2
 ; block1:
 ;   b label3
 ; block3:
@@ -61,8 +61,8 @@ block2 cold:
 ;
 ; Disassembled:
 ; block0: ; offset 0x0
-;   mov w5, w0
-;   cbz x5, #0xc
+;   mov w4, w0
+;   cbz x4, #0xc
 ; block1: ; offset 0x8
 ;   ret
 ; block2: ; offset 0xc

cranelift/filetests/filetests/isa/riscv64/bitops-float.clif

@@ -22,52 +22,52 @@ block1(v4: f32):
 ; VCode:
 ; block0:
 ;   li a0,0
-;   li a1,0
+;   li a5,0
+;   fmv.w.x fa1,a5
+;   fmv.x.w a4,fa1
+;   not a1,a4
+;   fmv.w.x fa2,a1
+;   fmv.x.w a3,fa2
+;   fmv.x.w a5,fa2
+;   or a1,a3,a5
 ;   fmv.w.x fa3,a1
-;   fmv.x.w a1,fa3
-;   not a2,a1
-;   fmv.w.x fa4,a2
-;   fmv.x.w a5,fa4
-;   fmv.x.w a1,fa4
-;   or a3,a5,a1
-;   fmv.w.x fa4,a3
-;   br_table a0,[MachLabel(1),MachLabel(2)]##tmp1=a1,tmp2=a2
+;   br_table a0,[MachLabel(1),MachLabel(2)]##tmp1=a5,tmp2=a1
 ; block1:
 ;   j label3
 ; block2:
-;   fmv.d fa4,fa3
+;   fmv.d fa3,fa1
 ;   j label3
 ; block3:
 ;   ret
 ;
 ; Disassembled:
 ; block0: ; offset 0x0
 ;   mv a0, zero
-;   mv a1, zero
+;   mv a5, zero
+;   fmv.w.x fa1, a5
+;   fmv.x.w a4, fa1
+;   not a1, a4
+;   fmv.w.x fa2, a1
+;   fmv.x.w a3, fa2
+;   fmv.x.w a5, fa2
+;   or a1, a3, a5
 ;   fmv.w.x fa3, a1
-;   fmv.x.w a1, fa3
-;   not a2, a1
-;   fmv.w.x fa4, a2
-;   fmv.x.w a5, fa4
-;   fmv.x.w a1, fa4
-;   or a3, a5, a1
-;   fmv.w.x fa4, a3
 ;   slli t6, a0, 0x20
 ;   srli t6, t6, 0x20
-;   addi a2, zero, 1
-;   bltu t6, a2, 0xc
-;   auipc a2, 0
-;   jalr zero, a2, 0x28
+;   addi a1, zero, 1
+;   bltu t6, a1, 0xc
 ;   auipc a1, 0
-;   slli a2, t6, 3
-;   add a1, a1, a2
-;   jalr zero, a1, 0x10
-;   auipc a2, 0
-;   jalr zero, a2, 0xc
+;   jalr zero, a1, 0x28
+;   auipc a5, 0
+;   slli a1, t6, 3
+;   add a5, a5, a1
+;   jalr zero, a5, 0x10
+;   auipc a1, 0
+;   jalr zero, a1, 0xc
 ; block1: ; offset 0x58
 ;   j 8
 ; block2: ; offset 0x5c
-;   fmv.d fa4, fa3
+;   fmv.d fa3, fa1
 ; block3: ; offset 0x60
 ;   ret
 

cranelift/filetests/filetests/isa/riscv64/cold-blocks.clif

@@ -16,8 +16,8 @@ block2:
 
 ; VCode:
 ; block0:
-;   sext.w a5,a0
-;   bne a5,zero,taken(label1),not_taken(label2)
+;   sext.w a4,a0
+;   bne a4,zero,taken(label1),not_taken(label2)
 ; block1:
 ;   j label3
 ; block2:
@@ -28,8 +28,8 @@ block2:
 ;
 ; Disassembled:
 ; block0: ; offset 0x0
-;   sext.w a5, a0
-;   bnez a5, 8
+;   sext.w a4, a0
+;   bnez a4, 8
 ; block1: ; offset 0x8
 ;   addi a0, zero, 0x61
 ; block2: ; offset 0xc
@@ -49,8 +49,8 @@ block2 cold:
 
 ; VCode:
 ; block0:
-;   sext.w a5,a0
-;   bne a5,zero,taken(label1),not_taken(label2)
+;   sext.w a4,a0
+;   bne a4,zero,taken(label1),not_taken(label2)
 ; block1:
 ;   j label3
 ; block3:
@@ -61,8 +61,8 @@ block2 cold:
 ;
 ; Disassembled:
 ; block0: ; offset 0x0
-;   sext.w a5, a0
-;   beqz a5, 8
+;   sext.w a4, a0
+;   beqz a4, 8
 ; block1: ; offset 0x8
 ;   ret
 ; block2: ; offset 0xc