[RISCV] Make use of SHXADD instructions in RVV spill/reload code.

[craig.topper]

Summary

We can use SH1ADD, SH2ADD, SH3ADD to multipy by 3, 5, and 9 respectively.

We could extend this to 3, 5, or 9 multiplied by a power 2 by also
emitting a SLLI.

While there remove the code that looks up TII. We're already a
member of RISCV's TII.

[reames]

LGTM w/minor comments.

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

1764–1765

Can you commit this separately as a cleanup please?

1822

Aside: Shouldn't this simply be one of options inside movImm?

[craig.topper]

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

1822

That's a good point. movImm probably handles this correctly.

[craig.topper]

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

1760

I'm considering merging this with it's one caller. There might be some more optimization opportunities if we include the ADD/SUB to SP.

For example, vlenb * 7 + sp could become

csrrd a0, vlenb
sub sp, sp, a0
slli a0, a0, 3
add sp, sp, a0

instead of

csrrd a0, vlenb
slli a1, a0, 3
sub a0, a1, a0
add sp, sp, a0

That would save a temporary register.

[reames]

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

1760

Separate commit please. :)

Any chance this could be done in a post-process combine like thing? Duplicating what appears to be generic multiple improvements seems like poor code structure.

[craig.topper]

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

1760

Yes separate commit.

I'm not sure if we could do it in a generic combine. This code runs after register allocation and needs to scavenge a free register. An emergency spill will happen if it can't find one. Avoiding emergency spills would be one of the benefits of saving the use of a temporary register. If did a generic combine later we'd also have to remove the emergency spill.

1764–1765

Will do

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

 #undef CASE_WIDEOP_OPCODE_COMMON
 
 Register RISCVInstrInfo::getVLENFactoredAmount(MachineFunction &MF,
                                                MachineBasicBlock &MBB,
                                                MachineBasicBlock::iterator II,
                                                const DebugLoc &DL,
                                                int64_t Amount,
                                                MachineInstr::MIFlag Flag) const {
   assert(Amount > 0 && "There is no need to get VLEN scaled value.");
   assert(Amount % 8 == 0 &&
          "Reserve the stack by the multiple of one vector size.");
 
   MachineRegisterInfo &MRI = MF.getRegInfo();
   int64_t NumOfVReg = Amount / 8;
 
   Register VL = MRI.createVirtualRegister(&RISCV::GPRRegClass);
   BuildMI(MBB, II, DL, get(RISCV::PseudoReadVLENB), VL)
     .setMIFlag(Flag);
   assert(isInt<32>(NumOfVReg) &&
          "Expect the number of vector registers within 32-bits.");
   if (isPowerOf2_32(NumOfVReg)) {
     uint32_t ShiftAmount = Log2_32(NumOfVReg);
     if (ShiftAmount == 0)
       return VL;
     BuildMI(MBB, II, DL, get(RISCV::SLLI), VL)
         .addReg(VL, RegState::Kill)
         .addImm(ShiftAmount)
         .setMIFlag(Flag);
+  } else if ((NumOfVReg == 3 || NumOfVReg == 5 || NumOfVReg == 9) &&
+             STI.hasStdExtZba()) {
+    // We can use Zba SHXADD instructions for multiply in some cases.
+    // TODO: Generalize to SHXADD+SLLI.
+    unsigned Opc;
+    switch (NumOfVReg) {
+    default: llvm_unreachable("Unexpected number of vregs");
+    case 3: Opc = RISCV::SH1ADD; break;
+    case 5: Opc = RISCV::SH2ADD; break;
+    case 9: Opc = RISCV::SH3ADD; break;
+    }
+    BuildMI(MBB, II, DL, get(Opc), VL)
+        .addReg(VL, RegState::Kill)
+        .addReg(VL)
+        .setMIFlag(Flag);
   } else if (isPowerOf2_32(NumOfVReg - 1)) {
     Register ScaledRegister = MRI.createVirtualRegister(&RISCV::GPRRegClass);
     uint32_t ShiftAmount = Log2_32(NumOfVReg - 1);
     BuildMI(MBB, II, DL, get(RISCV::SLLI), ScaledRegister)
         .addReg(VL)
         .addImm(ShiftAmount)
         .setMIFlag(Flag);
     BuildMI(MBB, II, DL, get(RISCV::ADD), VL)
BuildMI(MBB, II, DL, get(RISCV::SUB), VL)
         .addReg(ScaledRegister, RegState::Kill)
         .addReg(VL, RegState::Kill)
         .setMIFlag(Flag);
   } else {
     Register N = MRI.createVirtualRegister(&RISCV::GPRRegClass);
     if (!isInt<12>(NumOfVReg))
       movImm(MBB, II, DL, N, NumOfVReg);
     else {
       BuildMI(MBB, II, DL, get(RISCV::ADDI), N)
           .addReg(RISCV::X0)
           .addImm(NumOfVReg)
           .setMIFlag(Flag);
     }
     if (!STI.hasStdExtM())
       MF.getFunction().getContext().diagnose(DiagnosticInfoUnsupported{
           MF.getFunction(),
           "M-extension must be enabled to calculate the vscaled size/offset."});

llvm/test/CodeGen/RISCV/rvv/allocate-lmul-2-4-8.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv64 -mattr=+m,+v -verify-machineinstrs < %s \
-; RUN:    | FileCheck %s
+; RUN:    | FileCheck %s --check-prefixes=CHECK,NOZBA
+; RUN: llc -mtriple=riscv64 -mattr=+m,+v,+zba -verify-machineinstrs < %s \
+; RUN:    | FileCheck %s --check-prefixes=CHECK,ZBA
 
 define void @lmul1() nounwind {
 ; CHECK-LABEL: lmul1:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    csrr a0, vlenb
 ; CHECK-NEXT:    sub sp, sp, a0
 ; CHECK-NEXT:    csrr a0, vlenb
 ; CHECK-NEXT:    add sp, sp, a0

 ; CHECK-NEXT:    ld s0, 48(sp) # 8-byte Folded Reload
 ; CHECK-NEXT:    addi sp, sp, 64
 ; CHECK-NEXT:    ret
   %v = alloca <vscale x 8 x i64>
   ret void
 }
 
 define void @lmul1_and_2() nounwind {
-; CHECK-LABEL: lmul1_and_2:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    slli a1, a0, 1
-; CHECK-NEXT:    add a0, a1, a0
-; CHECK-NEXT:    sub sp, sp, a0
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    slli a1, a0, 1
-; CHECK-NEXT:    add a0, a1, a0
-; CHECK-NEXT:    add sp, sp, a0
-; CHECK-NEXT:    ret
+; NOZBA-LABEL: lmul1_and_2:
+; NOZBA:       # %bb.0:
+; NOZBA-NEXT:    csrr a0, vlenb
+; NOZBA-NEXT:    slli a1, a0, 1
+; NOZBA-NEXT:    add a0, a1, a0
+; NOZBA-NEXT:    sub sp, sp, a0
+; NOZBA-NEXT:    csrr a0, vlenb
+; NOZBA-NEXT:    slli a1, a0, 1
+; NOZBA-NEXT:    add a0, a1, a0
+; NOZBA-NEXT:    add sp, sp, a0
+; NOZBA-NEXT:    ret
+;
+; ZBA-LABEL: lmul1_and_2:
+; ZBA:       # %bb.0:
+; ZBA-NEXT:    csrr a0, vlenb
+; ZBA-NEXT:    sh1add a0, a0, a0
+; ZBA-NEXT:    sub sp, sp, a0
+; ZBA-NEXT:    csrr a0, vlenb
+; ZBA-NEXT:    sh1add a0, a0, a0
+; ZBA-NEXT:    add sp, sp, a0
+; ZBA-NEXT:    ret
   %v1 = alloca <vscale x 1 x i64>
   %v2 = alloca <vscale x 2 x i64>
   ret void
 }
 
 define void @lmul2_and_4() nounwind {
 ; CHECK-LABEL: lmul2_and_4:
 ; CHECK:       # %bb.0:

 ; CHECK-NEXT:    addi sp, sp, 32
 ; CHECK-NEXT:    ret
   %v1 = alloca <vscale x 2 x i64>
   %v2 = alloca <vscale x 4 x i64>
   ret void
 }
 
 define void @lmul1_and_4() nounwind {
-; CHECK-LABEL: lmul1_and_4:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    addi sp, sp, -32
-; CHECK-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
-; CHECK-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
-; CHECK-NEXT:    addi s0, sp, 32
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    slli a1, a0, 2
-; CHECK-NEXT:    add a0, a1, a0
-; CHECK-NEXT:    sub sp, sp, a0
-; CHECK-NEXT:    andi sp, sp, -32
-; CHECK-NEXT:    addi sp, s0, -32
-; CHECK-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
-; CHECK-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
-; CHECK-NEXT:    addi sp, sp, 32
-; CHECK-NEXT:    ret
+; NOZBA-LABEL: lmul1_and_4:
+; NOZBA:       # %bb.0:
+; NOZBA-NEXT:    addi sp, sp, -32
+; NOZBA-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
+; NOZBA-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
+; NOZBA-NEXT:    addi s0, sp, 32
+; NOZBA-NEXT:    csrr a0, vlenb
+; NOZBA-NEXT:    slli a1, a0, 2
+; NOZBA-NEXT:    add a0, a1, a0
+; NOZBA-NEXT:    sub sp, sp, a0
+; NOZBA-NEXT:    andi sp, sp, -32
+; NOZBA-NEXT:    addi sp, s0, -32
+; NOZBA-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
+; NOZBA-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
+; NOZBA-NEXT:    addi sp, sp, 32
+; NOZBA-NEXT:    ret
+;
+; ZBA-LABEL: lmul1_and_4:
+; ZBA:       # %bb.0:
+; ZBA-NEXT:    addi sp, sp, -32
+; ZBA-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
+; ZBA-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
+; ZBA-NEXT:    addi s0, sp, 32
+; ZBA-NEXT:    csrr a0, vlenb
+; ZBA-NEXT:    sh2add a0, a0, a0
+; ZBA-NEXT:    sub sp, sp, a0
+; ZBA-NEXT:    andi sp, sp, -32
+; ZBA-NEXT:    addi sp, s0, -32
+; ZBA-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
+; ZBA-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
+; ZBA-NEXT:    addi sp, sp, 32
+; ZBA-NEXT:    ret
   %v1 = alloca <vscale x 1 x i64>
   %v2 = alloca <vscale x 4 x i64>
   ret void
 }
 
 define void @lmul2_and_1() nounwind {
-; CHECK-LABEL: lmul2_and_1:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    slli a1, a0, 1
-; CHECK-NEXT:    add a0, a1, a0
-; CHECK-NEXT:    sub sp, sp, a0
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    slli a1, a0, 1
-; CHECK-NEXT:    add a0, a1, a0
-; CHECK-NEXT:    add sp, sp, a0
-; CHECK-NEXT:    ret
+; NOZBA-LABEL: lmul2_and_1:
+; NOZBA:       # %bb.0:
+; NOZBA-NEXT:    csrr a0, vlenb
+; NOZBA-NEXT:    slli a1, a0, 1
+; NOZBA-NEXT:    add a0, a1, a0
+; NOZBA-NEXT:    sub sp, sp, a0
+; NOZBA-NEXT:    csrr a0, vlenb
+; NOZBA-NEXT:    slli a1, a0, 1
+; NOZBA-NEXT:    add a0, a1, a0
+; NOZBA-NEXT:    add sp, sp, a0
+; NOZBA-NEXT:    ret
+;
+; ZBA-LABEL: lmul2_and_1:
+; ZBA:       # %bb.0:
+; ZBA-NEXT:    csrr a0, vlenb
+; ZBA-NEXT:    sh1add a0, a0, a0
+; ZBA-NEXT:    sub sp, sp, a0
+; ZBA-NEXT:    csrr a0, vlenb
+; ZBA-NEXT:    sh1add a0, a0, a0
+; ZBA-NEXT:    add sp, sp, a0
+; ZBA-NEXT:    ret
   %v1 = alloca <vscale x 2 x i64>
   %v2 = alloca <vscale x 1 x i64>
   ret void
 }
 
 define void @lmul4_and_1() nounwind {
-; CHECK-LABEL: lmul4_and_1:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    addi sp, sp, -32
-; CHECK-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
-; CHECK-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
-; CHECK-NEXT:    addi s0, sp, 32
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    slli a1, a0, 2
-; CHECK-NEXT:    add a0, a1, a0
-; CHECK-NEXT:    sub sp, sp, a0
-; CHECK-NEXT:    andi sp, sp, -32
-; CHECK-NEXT:    addi sp, s0, -32
-; CHECK-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
-; CHECK-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
-; CHECK-NEXT:    addi sp, sp, 32
-; CHECK-NEXT:    ret
+; NOZBA-LABEL: lmul4_and_1:
+; NOZBA:       # %bb.0:
+; NOZBA-NEXT:    addi sp, sp, -32
+; NOZBA-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
+; NOZBA-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
+; NOZBA-NEXT:    addi s0, sp, 32
+; NOZBA-NEXT:    csrr a0, vlenb
+; NOZBA-NEXT:    slli a1, a0, 2
+; NOZBA-NEXT:    add a0, a1, a0
+; NOZBA-NEXT:    sub sp, sp, a0
+; NOZBA-NEXT:    andi sp, sp, -32
+; NOZBA-NEXT:    addi sp, s0, -32
+; NOZBA-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
+; NOZBA-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
+; NOZBA-NEXT:    addi sp, sp, 32
+; NOZBA-NEXT:    ret
+;
+; ZBA-LABEL: lmul4_and_1:
+; ZBA:       # %bb.0:
+; ZBA-NEXT:    addi sp, sp, -32
+; ZBA-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
+; ZBA-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
+; ZBA-NEXT:    addi s0, sp, 32
+; ZBA-NEXT:    csrr a0, vlenb
+; ZBA-NEXT:    sh2add a0, a0, a0
+; ZBA-NEXT:    sub sp, sp, a0
+; ZBA-NEXT:    andi sp, sp, -32
+; ZBA-NEXT:    addi sp, s0, -32
+; ZBA-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
+; ZBA-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
+; ZBA-NEXT:    addi sp, sp, 32
+; ZBA-NEXT:    ret
   %v1 = alloca <vscale x 4 x i64>
   %v2 = alloca <vscale x 1 x i64>
   ret void
 }
 
 define void @lmul4_and_2() nounwind {
 ; CHECK-LABEL: lmul4_and_2:
 ; CHECK:       # %bb.0:
; CHECK-NEXT:    ret
   %v3 = alloca <vscale x 4 x i64>
   %v2 = alloca <vscale x 2 x i64>
   %v4 = alloca <vscale x 2 x i64>
   ret void
 }
 
 
 define void @gpr_and_lmul1_and_2() nounwind {
-; CHECK-LABEL: gpr_and_lmul1_and_2:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    addi sp, sp, -16
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    slli a1, a0, 1
-; CHECK-NEXT:    add a0, a1, a0
-; CHECK-NEXT:    sub sp, sp, a0
-; CHECK-NEXT:    li a0, 3
-; CHECK-NEXT:    sd a0, 8(sp)
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    slli a1, a0, 1
-; CHECK-NEXT:    add a0, a1, a0
-; CHECK-NEXT:    add sp, sp, a0
-; CHECK-NEXT:    addi sp, sp, 16
-; CHECK-NEXT:    ret
+; NOZBA-LABEL: gpr_and_lmul1_and_2:
+; NOZBA:       # %bb.0:
+; NOZBA-NEXT:    addi sp, sp, -16
+; NOZBA-NEXT:    csrr a0, vlenb
+; NOZBA-NEXT:    slli a1, a0, 1
+; NOZBA-NEXT:    add a0, a1, a0
+; NOZBA-NEXT:    sub sp, sp, a0
+; NOZBA-NEXT:    li a0, 3
+; NOZBA-NEXT:    sd a0, 8(sp)
+; NOZBA-NEXT:    csrr a0, vlenb
+; NOZBA-NEXT:    slli a1, a0, 1
+; NOZBA-NEXT:    add a0, a1, a0
+; NOZBA-NEXT:    add sp, sp, a0
+; NOZBA-NEXT:    addi sp, sp, 16
+; NOZBA-NEXT:    ret
+;
+; ZBA-LABEL: gpr_and_lmul1_and_2:
+; ZBA:       # %bb.0:
+; ZBA-NEXT:    addi sp, sp, -16
+; ZBA-NEXT:    csrr a0, vlenb
+; ZBA-NEXT:    sh1add a0, a0, a0
+; ZBA-NEXT:    sub sp, sp, a0
+; ZBA-NEXT:    li a0, 3
+; ZBA-NEXT:    sd a0, 8(sp)
+; ZBA-NEXT:    csrr a0, vlenb
+; ZBA-NEXT:    sh1add a0, a0, a0
+; ZBA-NEXT:    add sp, sp, a0
+; ZBA-NEXT:    addi sp, sp, 16
+; ZBA-NEXT:    ret
   %x1 = alloca i64
   %v1 = alloca <vscale x 1 x i64>
   %v2 = alloca <vscale x 2 x i64>
   store volatile i64 3, i64* %x1
   ret void
 }
 
 define void @gpr_and_lmul1_and_4() nounwind {
-; CHECK-LABEL: gpr_and_lmul1_and_4:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    addi sp, sp, -32
-; CHECK-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
-; CHECK-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
-; CHECK-NEXT:    addi s0, sp, 32
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    slli a1, a0, 2
-; CHECK-NEXT:    add a0, a1, a0
-; CHECK-NEXT:    sub sp, sp, a0
-; CHECK-NEXT:    andi sp, sp, -32
-; CHECK-NEXT:    li a0, 3
-; CHECK-NEXT:    sd a0, 8(sp)
-; CHECK-NEXT:    addi sp, s0, -32
-; CHECK-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
-; CHECK-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
-; CHECK-NEXT:    addi sp, sp, 32
-; CHECK-NEXT:    ret
+; NOZBA-LABEL: gpr_and_lmul1_and_4:
+; NOZBA:       # %bb.0:
+; NOZBA-NEXT:    addi sp, sp, -32
+; NOZBA-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
+; NOZBA-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
+; NOZBA-NEXT:    addi s0, sp, 32
+; NOZBA-NEXT:    csrr a0, vlenb
+; NOZBA-NEXT:    slli a1, a0, 2
+; NOZBA-NEXT:    add a0, a1, a0
+; NOZBA-NEXT:    sub sp, sp, a0
+; NOZBA-NEXT:    andi sp, sp, -32
+; NOZBA-NEXT:    li a0, 3
+; NOZBA-NEXT:    sd a0, 8(sp)
+; NOZBA-NEXT:    addi sp, s0, -32
+; NOZBA-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
+; NOZBA-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
+; NOZBA-NEXT:    addi sp, sp, 32
+; NOZBA-NEXT:    ret
+;
+; ZBA-LABEL: gpr_and_lmul1_and_4:
+; ZBA:       # %bb.0:
+; ZBA-NEXT:    addi sp, sp, -32
+; ZBA-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
+; ZBA-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
+; ZBA-NEXT:    addi s0, sp, 32
+; ZBA-NEXT:    csrr a0, vlenb
+; ZBA-NEXT:    sh2add a0, a0, a0
+; ZBA-NEXT:    sub sp, sp, a0
+; ZBA-NEXT:    andi sp, sp, -32
+; ZBA-NEXT:    li a0, 3
+; ZBA-NEXT:    sd a0, 8(sp)
+; ZBA-NEXT:    addi sp, s0, -32
+; ZBA-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
+; ZBA-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
+; ZBA-NEXT:    addi sp, sp, 32
+; ZBA-NEXT:    ret
   %x1 = alloca i64
   %v1 = alloca <vscale x 1 x i64>
   %v2 = alloca <vscale x 4 x i64>
   store volatile i64 3, i64* %x1
   ret void
 }
 
 define void @lmul_1_2_4_8() nounwind {