[RISCV] Optimize mul in the zba extension with SH*ADD

[benshi001]

Summary

This patch makes the following optimization, if the
immediate multiplier is not a simm12.

(mul x, (power_of_2 + 2)) => (SH1ADD x, (SLLI x, bits))
(mul x, (power_of_2 + 4)) => (SH2ADD x, (SLLI x, bits))
(mul x, (power_of_2 + 8)) => (SH3ADD x, (SLLI x, bits))

[benshi001]

It is obvious a win if the immediate is non-simm12, but not sure for simm12, so I leave that unchanged.

[benshi001]

Ping ...

I saw the similar optimization on x86, and think it will also benefit riscv.

[craig.topper]

LGTM

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

if (VT.isScalarInteger()) {
     if (Subtarget.hasStdExtM() && VT.getSizeInBits() > Subtarget.getXLen())
       return false;
     if (auto *ConstNode = dyn_cast<ConstantSDNode>(C.getNode())) {
       // Break the MUL to a SLLI and an ADD/SUB.
       const APInt &Imm = ConstNode->getAPIntValue();
       if ((Imm + 1).isPowerOf2() || (Imm - 1).isPowerOf2() ||
           (1 - Imm).isPowerOf2() || (-1 - Imm).isPowerOf2())
         return true;
+      // Optimize the MUL to (SH*ADD x, (SLLI x, bits)) if Imm is not simm12.
+      if (Subtarget.hasStdExtZba() && !Imm.isSignedIntN(12) &&
+          ((Imm - 2).isPowerOf2() || (Imm - 4).isPowerOf2() ||
+           (Imm - 8).isPowerOf2()))
+        return true;
       // Omit the following optimization if the sub target has the M extension
       // and the data size >= XLen.
       if (Subtarget.hasStdExtM() && VT.getSizeInBits() >= Subtarget.getXLen())
         return false;
       // Break the MUL to two SLLI instructions and an ADD/SUB, if Imm needs
       // a pair of LUI/ADDI.
       if (!Imm.isSignedIntN(12) && Imm.countTrailingZeros() < 12) {
         APInt ImmS = Imm.ashr(Imm.countTrailingZeros());

llvm/test/CodeGen/RISCV/rv32zba.ll

 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    lui a1, 1
 ; RV32I-NEXT:    addi a1, a1, 2
 ; RV32I-NEXT:    mul a0, a0, a1
 ; RV32I-NEXT:    ret
 ;
 ; RV32IB-LABEL: mul4098:
 ; RV32IB:       # %bb.0:
-; RV32IB-NEXT:    lui a1, 1
-; RV32IB-NEXT:    addi a1, a1, 2
-; RV32IB-NEXT:    mul a0, a0, a1
+; RV32IB-NEXT:    slli a1, a0, 12
+; RV32IB-NEXT:    sh1add a0, a0, a1
 ; RV32IB-NEXT:    ret
 ;
 ; RV32IBA-LABEL: mul4098:
 ; RV32IBA:       # %bb.0:
-; RV32IBA-NEXT:    lui a1, 1
-; RV32IBA-NEXT:    addi a1, a1, 2
-; RV32IBA-NEXT:    mul a0, a0, a1
+; RV32IBA-NEXT:    slli a1, a0, 12
+; RV32IBA-NEXT:    sh1add a0, a0, a1
 ; RV32IBA-NEXT:    ret
   %c = mul i32 %a, 4098
   ret i32 %c
 }
 
 define i32 @mul4100(i32 %a) {
 ; RV32I-LABEL: mul4100:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    lui a1, 1
 ; RV32I-NEXT:    addi a1, a1, 4
 ; RV32I-NEXT:    mul a0, a0, a1
 ; RV32I-NEXT:    ret
 ;
 ; RV32IB-LABEL: mul4100:
 ; RV32IB:       # %bb.0:
-; RV32IB-NEXT:    lui a1, 1
-; RV32IB-NEXT:    addi a1, a1, 4
-; RV32IB-NEXT:    mul a0, a0, a1
+; RV32IB-NEXT:    slli a1, a0, 12
+; RV32IB-NEXT:    sh2add a0, a0, a1
 ; RV32IB-NEXT:    ret
 ;
 ; RV32IBA-LABEL: mul4100:
 ; RV32IBA:       # %bb.0:
-; RV32IBA-NEXT:    lui a1, 1
-; RV32IBA-NEXT:    addi a1, a1, 4
-; RV32IBA-NEXT:    mul a0, a0, a1
+; RV32IBA-NEXT:    slli a1, a0, 12
+; RV32IBA-NEXT:    sh2add a0, a0, a1
 ; RV32IBA-NEXT:    ret
   %c = mul i32 %a, 4100
   ret i32 %c
 }
 
 define i32 @mul4104(i32 %a) {
 ; RV32I-LABEL: mul4104:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    lui a1, 1
 ; RV32I-NEXT:    addi a1, a1, 8
 ; RV32I-NEXT:    mul a0, a0, a1
 ; RV32I-NEXT:    ret
 ;
 ; RV32IB-LABEL: mul4104:
 ; RV32IB:       # %bb.0:
-; RV32IB-NEXT:    lui a1, 1
-; RV32IB-NEXT:    addi a1, a1, 8
-; RV32IB-NEXT:    mul a0, a0, a1
+; RV32IB-NEXT:    slli a1, a0, 12
+; RV32IB-NEXT:    sh3add a0, a0, a1
 ; RV32IB-NEXT:    ret
 ;
 ; RV32IBA-LABEL: mul4104:
 ; RV32IBA:       # %bb.0:
-; RV32IBA-NEXT:    lui a1, 1
-; RV32IBA-NEXT:    addi a1, a1, 8
-; RV32IBA-NEXT:    mul a0, a0, a1
+; RV32IBA-NEXT:    slli a1, a0, 12
+; RV32IBA-NEXT:    sh3add a0, a0, a1
 ; RV32IBA-NEXT:    ret
   %c = mul i32 %a, 4104
   ret i32 %c
 }

llvm/test/CodeGen/RISCV/rv64zba.ll

 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    lui a1, 1
 ; RV64I-NEXT:    addiw a1, a1, 2
 ; RV64I-NEXT:    mul a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: mul4098:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    lui a1, 1
-; RV64IB-NEXT:    addiw a1, a1, 2
-; RV64IB-NEXT:    mul a0, a0, a1
+; RV64IB-NEXT:    slli a1, a0, 12
+; RV64IB-NEXT:    sh1add a0, a0, a1
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBA-LABEL: mul4098:
 ; RV64IBA:       # %bb.0:
-; RV64IBA-NEXT:    lui a1, 1
-; RV64IBA-NEXT:    addiw a1, a1, 2
-; RV64IBA-NEXT:    mul a0, a0, a1
+; RV64IBA-NEXT:    slli a1, a0, 12
+; RV64IBA-NEXT:    sh1add a0, a0, a1
 ; RV64IBA-NEXT:    ret
   %c = mul i64 %a, 4098
   ret i64 %c
 }
 
 define i64 @mul4100(i64 %a) {
 ; RV64I-LABEL: mul4100:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    lui a1, 1
 ; RV64I-NEXT:    addiw a1, a1, 4
 ; RV64I-NEXT:    mul a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: mul4100:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    lui a1, 1
-; RV64IB-NEXT:    addiw a1, a1, 4
-; RV64IB-NEXT:    mul a0, a0, a1
+; RV64IB-NEXT:    slli a1, a0, 12
+; RV64IB-NEXT:    sh2add a0, a0, a1
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBA-LABEL: mul4100:
 ; RV64IBA:       # %bb.0:
-; RV64IBA-NEXT:    lui a1, 1
-; RV64IBA-NEXT:    addiw a1, a1, 4
-; RV64IBA-NEXT:    mul a0, a0, a1
+; RV64IBA-NEXT:    slli a1, a0, 12
+; RV64IBA-NEXT:    sh2add a0, a0, a1
 ; RV64IBA-NEXT:    ret
   %c = mul i64 %a, 4100
   ret i64 %c
 }
 
 define i64 @mul4104(i64 %a) {
 ; RV64I-LABEL: mul4104:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    lui a1, 1
 ; RV64I-NEXT:    addiw a1, a1, 8
 ; RV64I-NEXT:    mul a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: mul4104:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    lui a1, 1
-; RV64IB-NEXT:    addiw a1, a1, 8
-; RV64IB-NEXT:    mul a0, a0, a1
+; RV64IB-NEXT:    slli a1, a0, 12
+; RV64IB-NEXT:    sh3add a0, a0, a1
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBA-LABEL: mul4104:
 ; RV64IBA:       # %bb.0:
-; RV64IBA-NEXT:    lui a1, 1
-; RV64IBA-NEXT:    addiw a1, a1, 8
-; RV64IBA-NEXT:    mul a0, a0, a1
+; RV64IBA-NEXT:    slli a1, a0, 12
+; RV64IBA-NEXT:    sh3add a0, a0, a1
 ; RV64IBA-NEXT:    ret
   %c = mul i64 %a, 4104
   ret i64 %c
 }