riscv64: Clear the top 32bits in the br_table index (bytecodealliance#5831)

We were unintentionally relying on these to be zeroed when jumping.

Author: afonso360

cranelift/codegen/src/isa/riscv64/inst.isle

@@ -213,6 +213,7 @@
     (BrTable
       (index Reg)
       (tmp1 WritableReg)
+      (tmp2 WritableReg)
       (targets VecBranchTarget))
 
     ;; atomic compare and set operation

cranelift/codegen/src/isa/riscv64/inst/emit.rs

@@ -984,11 +984,13 @@ impl MachInstEmit for Inst {
             &Inst::BrTable {
                 index,
                 tmp1,
+                tmp2,
                 ref targets,
             } => {
                 let index = allocs.next(index);
                 let tmp1 = allocs.next_writable(tmp1);
-                let tmp2 = writable_spilltmp_reg();
+                let tmp2 = allocs.next_writable(tmp2);
+                let ext_index = writable_spilltmp_reg();
 
                 // The default target is passed in as the 0th element of `targets`
                 // separate it here for clarity.
@@ -1001,22 +1003,38 @@ impl MachInstEmit for Inst {
                 // that offset. Each jump table entry is a regular auipc+jalr which we emit sequentially.
                 //
                 // Build the following sequence:
+                //
+                // extend_index:
+                //     zext.w  ext_index, index
                 // bounds_check:
                 //     li      tmp, n_labels
-                //     bltu    index, tmp, compute_target
+                //     bltu    ext_index, tmp, compute_target
                 // jump_to_default_block:
                 //     auipc   pc, 0
                 //     jalr    zero, pc, default_block
                 // compute_target:
                 //     auipc   pc, 0
-                //     slli    tmp, index, 3
+                //     slli    tmp, ext_index, 3
                 //     add     pc, pc, tmp
                 //     jalr    zero, pc, 0x10
                 // jump_table:
                 //     ; This repeats for each entry in the jumptable
                 //     auipc   pc, 0
                 //     jalr    zero, pc, block_target
 
+                // Extend the index to 64 bits.
+                //
+                // This prevents us branching on the top 32 bits of the index, which
+                // are undefined.
+                Inst::Extend {
+                    rd: ext_index,
+                    rn: index,
+                    signed: false,
+                    from_bits: 32,
+                    to_bits: 64,
+                }
+                .emit(&[], sink, emit_info, state);
+
                 // Bounds check.
                 //
                 // Check if the index passed in is larger than the number of jumptable
@@ -1030,7 +1048,7 @@ impl MachInstEmit for Inst {
                     not_taken: BranchTarget::zero(),
                     kind: IntegerCompare {
                         kind: IntCC::UnsignedLessThan,
-                        rs1: index,
+                        rs1: ext_index.to_reg(),
                         rs2: tmp2.to_reg(),
                     },
                 }
@@ -1059,7 +1077,7 @@ impl MachInstEmit for Inst {
                 Inst::AluRRImm12 {
                     alu_op: AluOPRRI::Slli,
                     rd: tmp2,
-                    rs: index,
+                    rs: ext_index.to_reg(),
                     imm12: Imm12::from_bits(3),
                 }
                 .emit(&[], sink, emit_info, state);

cranelift/codegen/src/isa/riscv64/inst/mod.rs

@@ -338,9 +338,12 @@ fn riscv64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
     match inst {
         &Inst::Nop0 => {}
         &Inst::Nop4 => {}
-        &Inst::BrTable { index, tmp1, .. } => {
+        &Inst::BrTable {
+            index, tmp1, tmp2, ..
+        } => {
             collector.reg_use(index);
             collector.reg_early_def(tmp1);
+            collector.reg_early_def(tmp2);
         }
         &Inst::Auipc { rd, .. } => collector.reg_def(rd),
         &Inst::Lui { rd, .. } => collector.reg_def(rd),
@@ -1171,15 +1174,17 @@ impl Inst {
             &Inst::BrTable {
                 index,
                 tmp1,
+                tmp2,
                 ref targets,
             } => {
                 let targets: Vec<_> = targets.iter().map(|x| x.as_label().unwrap()).collect();
                 format!(
-                    "{} {},{}##tmp1={}",
+                    "{} {},{}##tmp1={},tmp2={}",
                     "br_table",
                     format_reg(index, allocs),
                     format_labels(&targets[..]),
                     format_reg(tmp1.to_reg(), allocs),
+                    format_reg(tmp2.to_reg(), allocs),
                 )
             }
             &Inst::Auipc { rd, imm } => {

cranelift/codegen/src/isa/riscv64/lower/isle.rs

@@ -378,6 +378,7 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Riscv64Backend> {
 
     fn lower_br_table(&mut self, index: Reg, targets: &VecMachLabel) -> Unit {
         let tmp1 = self.temp_writable_reg(I64);
+        let tmp2 = self.temp_writable_reg(I64);
         let targets: Vec<BranchTarget> = targets
             .into_iter()
             .copied()
@@ -386,6 +387,7 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Riscv64Backend> {
         self.emit(&MInst::BrTable {
             index,
             tmp1,
+            tmp2,
             targets,
         });
     }

cranelift/filetests/filetests/isa/riscv64/br_table.clif

@@ -0,0 +1,95 @@
+test compile precise-output
+set unwind_info=false
+target riscv64
+
+function %br_table(i32) -> i32 {
+block0(v0: i32):
+  br_table v0, block4, [block1, block2, block2, block3]
+
+block1:
+  v1 = iconst.i32 1
+  jump block5(v1)
+
+block2:
+  v2 = iconst.i32 2
+  jump block5(v2)
+
+block3:
+  v3 = iconst.i32 3
+  jump block5(v3)
+
+block4:
+  v4 = iconst.i32 4
+  jump block5(v4)
+
+block5(v5: i32):
+  v6 = iadd.i32 v0, v5
+  return v6
+}
+
+; VCode:
+; block0:
+;   br_table a0,[MachLabel(1),MachLabel(3),MachLabel(5),MachLabel(6),MachLabel(8)]##tmp1=t3,tmp2=t4
+; block1:
+;   li a1,4
+;   j label2
+; block2:
+;   j label10
+; block3:
+;   li a1,1
+;   j label4
+; block4:
+;   j label10
+; block5:
+;   j label7
+; block6:
+;   j label7
+; block7:
+;   li a1,2
+;   j label10
+; block8:
+;   li a1,3
+;   j label9
+; block9:
+;   j label10
+; block10:
+;   addw a0,a0,a1
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   slli t6, a0, 0x20
+;   srli t6, t6, 0x20
+;   addi t4, zero, 4
+;   bltu t6, t4, 0xc
+;   auipc t4, 0
+;   jalr zero, t4, 0x38
+;   auipc t3, 0
+;   slli t4, t6, 3
+;   add t3, t3, t4
+;   jalr zero, t3, 0x10
+;   auipc t4, 0
+;   jalr zero, t4, 0x28
+;   auipc t4, 0
+;   jalr zero, t4, 0x28
+;   auipc t4, 0
+;   jalr zero, t4, 0x20
+;   auipc t4, 0
+;   jalr zero, t4, 0x20
+; block1: ; offset 0x48
+;   addi a1, zero, 4
+; block2: ; offset 0x4c
+;   j 0x18
+; block3: ; offset 0x50
+;   addi a1, zero, 1
+; block4: ; offset 0x54
+;   j 0x10
+; block5: ; offset 0x58
+;   addi a1, zero, 2
+;   j 8
+; block6: ; offset 0x60
+;   addi a1, zero, 3
+; block7: ; offset 0x64
+;   addw a0, a0, a1
+;   ret
+

cranelift/filetests/filetests/runtests/br_table.clif

@@ -100,3 +100,32 @@ block1(v5: i32):
 ; run: %br_table_i32_inline_varied(4) == 4
 ; run: %br_table_i32_inline_varied(297) == 4
 ; run: %br_table_i32_inline_varied(65535) == 4
+
+
+
+; This is a regression test for #5831.
+; The riscv64 backend was failing to clear the upper half of the
+; index register on a br_table, which caused it to jump to the wrong
+; block.
+function %br_table_upper_reg() -> i32 {
+block0:
+    v0 = iconst.i32 -555163938
+    v1 = iconst.i8 -34
+    jump block1(v0, v1)
+
+block1(v2: i32, v3: i8):
+    v4 = ishl.i32 v2, v2
+    v5 = rotr v4, v3
+    br_table v5, block2, [block2, block2, block3]
+
+block2 cold:
+    v100 = iconst.i32 100
+    return v100
+
+block3 cold:
+    v200 = iconst.i32 200
+    return v200
+}
+
+
+; run: %br_table_upper_reg() == 200