diff --git a/cranelift/codegen/src/isa/riscv64/inst.isle b/cranelift/codegen/src/isa/riscv64/inst.isle index 10bcbd193b..c7e63d8d96 100644 --- a/cranelift/codegen/src/isa/riscv64/inst.isle +++ b/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 diff --git a/cranelift/codegen/src/isa/riscv64/inst/emit.rs b/cranelift/codegen/src/isa/riscv64/inst/emit.rs index 0bee435712..b85ee23b36 100644 --- a/cranelift/codegen/src/isa/riscv64/inst/emit.rs +++ b/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,15 +1003,18 @@ 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: @@ -1017,6 +1022,19 @@ impl MachInstEmit for Inst { // 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); diff --git a/cranelift/codegen/src/isa/riscv64/inst/mod.rs b/cranelift/codegen/src/isa/riscv64/inst/mod.rs index 350b69b8a1..776af8cf8c 100644 --- a/cranelift/codegen/src/isa/riscv64/inst/mod.rs +++ b/cranelift/codegen/src/isa/riscv64/inst/mod.rs @@ -338,9 +338,12 @@ fn riscv64_get_operands 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 } => { diff --git a/cranelift/codegen/src/isa/riscv64/lower/isle.rs b/cranelift/codegen/src/isa/riscv64/lower/isle.rs index 343992cd85..97852c0c44 100644 --- a/cranelift/codegen/src/isa/riscv64/lower/isle.rs +++ b/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 = targets .into_iter() .copied() @@ -386,6 +387,7 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Riscv64Backend> { self.emit(&MInst::BrTable { index, tmp1, + tmp2, targets, }); } diff --git a/cranelift/filetests/filetests/isa/riscv64/br_table.clif b/cranelift/filetests/filetests/isa/riscv64/br_table.clif new file mode 100644 index 0000000000..aeacae35cf --- /dev/null +++ b/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 + diff --git a/cranelift/filetests/filetests/runtests/br_table.clif b/cranelift/filetests/filetests/runtests/br_table.clif index 7151adcaac..f5ccc43cd3 100644 --- a/cranelift/filetests/filetests/runtests/br_table.clif +++ b/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