Cranelift: remove block params on critical-edge blocks. (#10485)

When a block has a terminator branch that targets two or more other
blocks at the CLIF level, and any of these blocks have two or more
precessors, the edge is a "critical edge" and we split it (insert a
new empty block) so that the register allocator has a place to put
moves that happen only on that edge. Otherwise, there is no location
that works: in the predecessor, code runs no matter which outgoing
edge we take; and in the successor, code runs no matter which incoming
edge we came from.

Currently, when we generate these critical-edge blocks, we insert
exactly one instruction: an unconditional branch. We wire up the
blockparam dataflow by (i) adding block parameters to the
critical-edge block with the same signature as the original target,
and (ii) adding all of these arguments to the unconditional branch. In
other words, we maintain the original block signature throughout.

This is fine and correct, but it has two downsides. The first is a
minor loss in compile-time efficiency (more SSA values and
block-params to process). The second, more interesting, is that it
hinders future work with certain kinds of branches that may define
values *on edges*.

In particular, this approach prevents exception-handling support: a
`try_call` instruction that acts as a terminator branch (with
normal-return and exceptional out-edges) defines normal-return values
as block-call arguments that are usable on the normal-return
edge. Some of these normal-return values may be defined by loads from
a return-value area. These loads need somewhere to go; they can't go
"after the terminator" (then it wouldn't be a terminator), so they go
in an edge block; as a result, the block-call for the normal-return
needs to use its arguments only in the unconditional branch out of the
edge block, not in the initial branch to the edge block.

This PR alters the critical-edge blockparam handling to have no
block-call args on the branch into the edge block, and use the
original values (not the newly defined edge-block blockparams) in the
block-call out of the edge block. This will allow these values to be
possibly defined in the edge block rather than in the predecessor (the
block with the original terminator).

This has no functional change today other than some perturbation of
regalloc decisions and a possibly slight compile-time speedup.

Author: cfallin

cranelift/codegen/src/machinst/blockorder.rs

@@ -132,7 +132,6 @@ impl LoweredBlock {
     }
 
     /// The associated out-edge successor, if this is a critical edge.
-    #[cfg(test)]
     pub fn out_edge(&self) -> Option<Block> {
         match self {
             &LoweredBlock::CriticalEdge { succ, .. } => Some(succ),

cranelift/codegen/src/machinst/lower.rs

@@ -872,7 +872,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
         self.vcode.end_bb();
     }
 
-    fn lower_clif_branches<B: LowerBackend<MInst = I>>(
+    fn lower_clif_branch<B: LowerBackend<MInst = I>>(
         &mut self,
         backend: &B,
         // Lowered block index:
@@ -883,7 +883,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
         targets: &[MachLabel],
     ) -> CodegenResult<()> {
         trace!(
-            "lower_clif_branches: block {} branch {:?} targets {:?}",
+            "lower_clif_branch: block {} branch {:?} targets {:?}",
             block,
             branch,
             targets,
@@ -909,31 +909,17 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
     }
 
     fn lower_branch_blockparam_args(&mut self, block: BlockIndex) {
+        let mut branch_arg_vregs: SmallVec<[Reg; 16]> = smallvec![];
+
         // 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 {
-                debug_assert!(self.f.dfg.value_is_real(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[..]);
+            branch_arg_vregs.clear();
+            let (succ, args) = self.collect_block_call(block, succ_idx, &mut branch_arg_vregs);
+            self.vcode.add_succ(succ, args);
         }
     }
 
-    fn collect_branches_and_targets(
+    fn collect_branch_and_targets(
         &self,
         bindex: BlockIndex,
         _bb: Block,
@@ -945,6 +931,56 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
         opt_inst
     }
 
+    /// Collect the outgoing block-call arguments for a given edge out
+    /// of a lowered block.
+    fn collect_block_call<'a>(
+        &mut self,
+        block: BlockIndex,
+        succ_idx: usize,
+        buffer: &'a mut SmallVec<[Reg; 16]>,
+    ) -> (BlockIndex, &'a [Reg]) {
+        let block_order = self.vcode.block_order();
+        let (_, succs) = block_order.succ_indices(block);
+        let succ = succs[succ_idx];
+        let this_lb = block_order.lowered_order()[block.index()];
+        let succ_lb = block_order.lowered_order()[succ.index()];
+
+        let (branch_inst, succ_idx) = match (this_lb, succ_lb) {
+            (_, LoweredBlock::CriticalEdge { .. }) => {
+                // The successor is a split-critical-edge block. In this
+                // case, this block-call has no arguments, and the
+                // arguments go on the critical edge block's unconditional
+                // branch instead.
+                return (succ, &[]);
+            }
+            (LoweredBlock::CriticalEdge { pred, succ_idx, .. }, _) => {
+                // This is a split-critical-edge block. In this case, our
+                // block-call has the arguments that in the CLIF appear in
+                // the predecessor's branch to this edge.
+                let branch_inst = self.f.layout.last_inst(pred).unwrap();
+                (branch_inst, succ_idx as usize)
+            }
+
+            (this, _) => {
+                let block = this.orig_block().unwrap();
+                // Ordinary block, with an ordinary block as
+                // successor. Take the arguments from the branch.
+                let branch_inst = self.f.layout.last_inst(block).unwrap();
+                (branch_inst, succ_idx)
+            }
+        };
+
+        let block_call =
+            self.f.dfg.insts[branch_inst].branch_destination(&self.f.dfg.jump_tables)[succ_idx];
+        let args = block_call.args_slice(&self.f.dfg.value_lists);
+        for &arg in args {
+            debug_assert!(self.f.dfg.value_is_real(arg));
+            let regs = self.put_value_in_regs(arg);
+            buffer.extend_from_slice(regs.regs());
+        }
+        (succ, &buffer[..])
+    }
+
     /// Lower the function.
     pub fn lower<B: LowerBackend<MInst = I>>(
         mut self,
@@ -981,39 +1017,22 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
             // Lower the block body in reverse order (see comment in
             // `lower_clif_block()` for rationale).
 
-            // End branches.
+            // End branch.
             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)?;
+                if let Some(branch) = self.collect_branch_and_targets(bindex, bb, &mut targets) {
+                    self.lower_clif_branch(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);
-                    }
-                }
-                self.vcode.add_succ(succ, &branch_arg_vregs[..]);
-
+                // get the successor and emit a jump. This block has
+                // no block params; and this jump's block-call args
+                // will be filled in by
+                // `lower_branch_blockparam_args`.
+                let succ = self.vcode.block_order().succ_indices(bindex).1[0];
                 self.emit(I::gen_jump(MachLabel::from_block(succ)));
                 self.finish_ir_inst(Default::default());
+                self.lower_branch_blockparam_args(bindex);
             }
 
             // Original block body.

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