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

[benshi001]

Summary

This patch does the following optimization.

Rx + Ry * 6 => (SH1ADD (SH2ADD Rx, Ry), Ry)
Rx + Ry * 10 => (SH1ADD (SH3ADD Rx, Ry), Ry)
Rx + Ry * 12 => (SH2ADD (SH3ADD Rx, Ry), Ry)

[benshi001]

The SH1ADDUW/SH2ADDUW/SH3ADDUW are different, and my previous test cases do not patch.

I will try SH1ADDUW/SH2ADDUW/SH3ADDUW next week, and you are appreciated to review current optimization with SH1ADD/SH2ADD/SH3ADD.

[benshi001]

In D104436#2823697, @benshi001 wrote:

The SH1ADDUW/SH2ADDUW/SH3ADDUW are different, and my previous test cases do not patch.

I will try SH1ADDUW/SH2ADDUW/SH3ADDUW next week, and you are appreciated to review current optimization with SH1ADD/SH2ADD/SH3ADD.

It seems SH1/2/3ADDUW are quite different and are not simply 32-bit SH1/2/3ADD on RV64, they can not be used to optimize mul-add.

[craig.topper]

There's a more generic optimization hiding here. Could we teach decomposeMulByConstant to emit (shl (sh1add X, X), C) to handle any constant of the form (3 << C). Similar for (shl (sh2add X, X)) to handle (5 << C), and (shl (sh3add X, X)) to handle (9 << C). If the multiply happens to be used by an add the existing patterns would combine the ADD and the SHL when possible.

If you want to try that as a followup. I'd suggest using the sequence you'd get from that instead. So your isel patterns would be

def : Pat<(add (mul GPR:$rs1, (XLenVT 6)), GPR:$rs2),
          (SH1ADD (SH1ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 10)), GPR:$rs2),
          (SH1ADD (SH2ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 12)), GPR:$rs2),
          (SH2ADD (SH1ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;

And you can add these additional cases.

def : Pat<(add (mul GPR:$rs1, (XLenVT 24)), GPR:$rs2),
          (SH3ADD (SH1ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 20)), GPR:$rs2),
          (SH2ADD (SH2ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 40)), GPR:$rs2),
          (SH3ADD (SH2ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 18)), GPR:$rs2),
          (SH1ADD (SH3ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 36)), GPR:$rs2),
          (SH2ADD (SH3ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 72)), GPR:$rs2),
          (SH3ADD (SH3ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;

X86 does basically the same optimization using LEA which is like our SHNADD. https://godbolt.org/z/e8PTT3oTo

llvm/lib/Target/RISCV/RISCVInstrInfoB.td

974

What if the multiply has an additional user that isn't the add?

[benshi001]

In D104436#2825904, @craig.topper wrote:

There's a more generic optimization hiding here. Could we teach decomposeMulByConstant to emit (shl (sh1add X, X), C) to handle any constant of the form (3 << C). Similar for (shl (sh2add X, X)) to handle (5 << C), and (shl (sh3add X, X)) to handle (9 << C). If the multiply happens to be used by an add the existing patterns would combine the ADD and the SHL when possible.

If you want to try that as a followup. I'd suggest using the sequence you'd get from that instead. So your isel patterns would be

def : Pat<(add (mul GPR:$rs1, (XLenVT 6)), GPR:$rs2),
          (SH1ADD (SH1ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 10)), GPR:$rs2),
          (SH1ADD (SH2ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 12)), GPR:$rs2),
          (SH2ADD (SH1ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;

And you can add these additional cases.

def : Pat<(add (mul GPR:$rs1, (XLenVT 24)), GPR:$rs2),
          (SH3ADD (SH1ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 20)), GPR:$rs2),
          (SH2ADD (SH2ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 40)), GPR:$rs2),
          (SH3ADD (SH2ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 18)), GPR:$rs2),
          (SH1ADD (SH3ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 36)), GPR:$rs2),
          (SH2ADD (SH3ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 72)), GPR:$rs2),
          (SH3ADD (SH3ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;

X86 does basically the same optimization using LEA which is like our SHNADD. https://godbolt.org/z/e8PTT3oTo

Thanks for teaching me so much skills. It seems I should submit another patch first, which contains tests for all above ways.

[benshi001]

In D104436#2826096, @benshi001 wrote:

In D104436#2825904, @craig.topper wrote:

There's a more generic optimization hiding here. Could we teach decomposeMulByConstant to emit (shl (sh1add X, X), C) to handle any constant of the form (3 << C). Similar for (shl (sh2add X, X)) to handle (5 << C), and (shl (sh3add X, X)) to handle (9 << C). If the multiply happens to be used by an add the existing patterns would combine the ADD and the SHL when possible.

If you want to try that as a followup. I'd suggest using the sequence you'd get from that instead. So your isel patterns would be

def : Pat<(add (mul GPR:$rs1, (XLenVT 6)), GPR:$rs2),
          (SH1ADD (SH1ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 10)), GPR:$rs2),
          (SH1ADD (SH2ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 12)), GPR:$rs2),
          (SH2ADD (SH1ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;

And you can add these additional cases.

def : Pat<(add (mul GPR:$rs1, (XLenVT 24)), GPR:$rs2),
          (SH3ADD (SH1ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 20)), GPR:$rs2),
          (SH2ADD (SH2ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 40)), GPR:$rs2),
          (SH3ADD (SH2ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 18)), GPR:$rs2),
          (SH1ADD (SH3ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 36)), GPR:$rs2),
          (SH2ADD (SH3ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
def : Pat<(add (mul GPR:$rs1, (XLenVT 72)), GPR:$rs2),
          (SH3ADD (SH3ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;

X86 does basically the same optimization using LEA which is like our SHNADD. https://godbolt.org/z/e8PTT3oTo

Thanks for teaching me so much skills. It seems I should submit another patch first, which contains tests for all above ways.

I have added new tests in https://reviews.llvm.org/D104507

llvm/lib/Target/RISCV/RISCVInstrInfoB.td

 
 let Predicates = [HasStdExtZba] in {
 def : Pat<(add (shl GPR:$rs1, (XLenVT 1)), GPR:$rs2),
           (SH1ADD GPR:$rs1, GPR:$rs2)>;
 def : Pat<(add (shl GPR:$rs1, (XLenVT 2)), GPR:$rs2),
           (SH2ADD GPR:$rs1, GPR:$rs2)>;
 def : Pat<(add (shl GPR:$rs1, (XLenVT 3)), GPR:$rs2),
           (SH3ADD GPR:$rs1, GPR:$rs2)>;
+
+def : Pat<(add (mul GPR:$rs1, (XLenVT 6)), GPR:$rs2),
+          (SH1ADD GPR:$rs1, (SH2ADD GPR:$rs1, GPR:$rs2))>;
+def : Pat<(add (mul GPR:$rs1, (XLenVT 10)), GPR:$rs2),
+          (SH1ADD GPR:$rs1, (SH3ADD GPR:$rs1, GPR:$rs2))>;
+def : Pat<(add (mul GPR:$rs1, (XLenVT 12)), GPR:$rs2),
+          (SH2ADD GPR:$rs1, (SH3ADD GPR:$rs1, GPR:$rs2))>;
 } // Predicates = [HasStdExtZba]
 
 let Predicates = [HasStdExtZba, IsRV64] in {
 def : Pat<(i64 (SLLIUWPat GPR:$rs1, uimm5:$shamt)),
           (SLLIUW GPR:$rs1, uimm5:$shamt)>;
 def : Pat<(i64 (shl (and GPR:$rs1, 0xFFFFFFFF), uimm5:$shamt)),
           (SLLIUW GPR:$rs1, uimm5:$shamt)>;
 def : Pat<(i64 (add (and GPR:$rs1, 0xFFFFFFFF), GPR:$rs2)),

llvm/test/CodeGen/RISCV/rv32zba.ll

 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    addi a2, zero, 6
 ; RV32I-NEXT:    mul a0, a0, a2
 ; RV32I-NEXT:    add a0, a0, a1
 ; RV32I-NEXT:    ret
 ;
 ; RV32IB-LABEL: addmul6:
 ; RV32IB:       # %bb.0:
-; RV32IB-NEXT:    addi a2, zero, 6
-; RV32IB-NEXT:    mul a0, a0, a2
-; RV32IB-NEXT:    add a0, a0, a1
+; RV32IB-NEXT:    sh2add a1, a0, a1
+; RV32IB-NEXT:    sh1add a0, a0, a1
 ; RV32IB-NEXT:    ret
 ;
 ; RV32IBA-LABEL: addmul6:
 ; RV32IBA:       # %bb.0:
-; RV32IBA-NEXT:    addi a2, zero, 6
-; RV32IBA-NEXT:    mul a0, a0, a2
-; RV32IBA-NEXT:    add a0, a0, a1
+; RV32IBA-NEXT:    sh2add a1, a0, a1
+; RV32IBA-NEXT:    sh1add a0, a0, a1
 ; RV32IBA-NEXT:    ret
   %c = mul i32 %a, 6
   %d = add i32 %c, %b
   ret i32 %d
 }
 
 define i32 @addmul10(i32 %a, i32 %b) {
 ; RV32I-LABEL: addmul10:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    addi a2, zero, 10
 ; RV32I-NEXT:    mul a0, a0, a2
 ; RV32I-NEXT:    add a0, a0, a1
 ; RV32I-NEXT:    ret
 ;
 ; RV32IB-LABEL: addmul10:
 ; RV32IB:       # %bb.0:
-; RV32IB-NEXT:    addi a2, zero, 10
-; RV32IB-NEXT:    mul a0, a0, a2
-; RV32IB-NEXT:    add a0, a0, a1
+; RV32IB-NEXT:    sh3add a1, a0, a1
+; RV32IB-NEXT:    sh1add a0, a0, a1
 ; RV32IB-NEXT:    ret
 ;
 ; RV32IBA-LABEL: addmul10:
 ; RV32IBA:       # %bb.0:
-; RV32IBA-NEXT:    addi a2, zero, 10
-; RV32IBA-NEXT:    mul a0, a0, a2
-; RV32IBA-NEXT:    add a0, a0, a1
+; RV32IBA-NEXT:    sh3add a1, a0, a1
+; RV32IBA-NEXT:    sh1add a0, a0, a1
 ; RV32IBA-NEXT:    ret
   %c = mul i32 %a, 10
   %d = add i32 %c, %b
   ret i32 %d
 }
 
 define i32 @addmul12(i32 %a, i32 %b) {
 ; RV32I-LABEL: addmul12:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    addi a2, zero, 12
 ; RV32I-NEXT:    mul a0, a0, a2
 ; RV32I-NEXT:    add a0, a0, a1
 ; RV32I-NEXT:    ret
 ;
 ; RV32IB-LABEL: addmul12:
 ; RV32IB:       # %bb.0:
-; RV32IB-NEXT:    addi a2, zero, 12
-; RV32IB-NEXT:    mul a0, a0, a2
-; RV32IB-NEXT:    add a0, a0, a1
+; RV32IB-NEXT:    sh3add a1, a0, a1
+; RV32IB-NEXT:    sh2add a0, a0, a1
 ; RV32IB-NEXT:    ret
 ;
 ; RV32IBA-LABEL: addmul12:
 ; RV32IBA:       # %bb.0:
-; RV32IBA-NEXT:    addi a2, zero, 12
-; RV32IBA-NEXT:    mul a0, a0, a2
-; RV32IBA-NEXT:    add a0, a0, a1
+; RV32IBA-NEXT:    sh3add a1, a0, a1
+; RV32IBA-NEXT:    sh2add a0, a0, a1
 ; RV32IBA-NEXT:    ret
   %c = mul i32 %a, 12
   %d = add i32 %c, %b
   ret i32 %d
 }

llvm/test/CodeGen/RISCV/rv64zba.ll

 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    addi a2, zero, 6
 ; RV64I-NEXT:    mul a0, a0, a2
 ; RV64I-NEXT:    add a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: addmul6:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    addi a2, zero, 6
-; RV64IB-NEXT:    mul a0, a0, a2
-; RV64IB-NEXT:    add a0, a0, a1
+; RV64IB-NEXT:    sh2add a1, a0, a1
+; RV64IB-NEXT:    sh1add a0, a0, a1
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBA-LABEL: addmul6:
 ; RV64IBA:       # %bb.0:
-; RV64IBA-NEXT:    addi a2, zero, 6
-; RV64IBA-NEXT:    mul a0, a0, a2
-; RV64IBA-NEXT:    add a0, a0, a1
+; RV64IBA-NEXT:    sh2add a1, a0, a1
+; RV64IBA-NEXT:    sh1add a0, a0, a1
 ; RV64IBA-NEXT:    ret
   %c = mul i64 %a, 6
   %d = add i64 %c, %b
   ret i64 %d
 }
 
 define i64 @addmul10(i64 %a, i64 %b) {
 ; RV64I-LABEL: addmul10:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    addi a2, zero, 10
 ; RV64I-NEXT:    mul a0, a0, a2
 ; RV64I-NEXT:    add a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: addmul10:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    addi a2, zero, 10
-; RV64IB-NEXT:    mul a0, a0, a2
-; RV64IB-NEXT:    add a0, a0, a1
+; RV64IB-NEXT:    sh3add a1, a0, a1
+; RV64IB-NEXT:    sh1add a0, a0, a1
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBA-LABEL: addmul10:
 ; RV64IBA:       # %bb.0:
-; RV64IBA-NEXT:    addi a2, zero, 10
-; RV64IBA-NEXT:    mul a0, a0, a2
-; RV64IBA-NEXT:    add a0, a0, a1
+; RV64IBA-NEXT:    sh3add a1, a0, a1
+; RV64IBA-NEXT:    sh1add a0, a0, a1
 ; RV64IBA-NEXT:    ret
   %c = mul i64 %a, 10
   %d = add i64 %c, %b
   ret i64 %d
 }
 
 define i64 @addmul12(i64 %a, i64 %b) {
 ; RV64I-LABEL: addmul12:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    addi a2, zero, 12
 ; RV64I-NEXT:    mul a0, a0, a2
 ; RV64I-NEXT:    add a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: addmul12:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    addi a2, zero, 12
-; RV64IB-NEXT:    mul a0, a0, a2
-; RV64IB-NEXT:    add a0, a0, a1
+; RV64IB-NEXT:    sh3add a1, a0, a1
+; RV64IB-NEXT:    sh2add a0, a0, a1
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBA-LABEL: addmul12:
 ; RV64IBA:       # %bb.0:
-; RV64IBA-NEXT:    addi a2, zero, 12
-; RV64IBA-NEXT:    mul a0, a0, a2
-; RV64IBA-NEXT:    add a0, a0, a1
+; RV64IBA-NEXT:    sh3add a1, a0, a1
+; RV64IBA-NEXT:    sh2add a0, a0, a1
 ; RV64IBA-NEXT:    ret
   %c = mul i64 %a, 12
   %d = add i64 %c, %b
   ret i64 %d
 }
-
-define i32 @addmulw6(i32 signext %a, i32 signext %b) {
-; RV64I-LABEL: addmulw6:
-; RV64I:       # %bb.0:
-; RV64I-NEXT:    addi a2, zero, 6
-; RV64I-NEXT:    mul a0, a0, a2
-; RV64I-NEXT:    addw a0, a0, a1
-; RV64I-NEXT:    ret
-;
-; RV64IB-LABEL: addmulw6:
-; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    addi a2, zero, 6
-; RV64IB-NEXT:    mul a0, a0, a2
-; RV64IB-NEXT:    addw a0, a0, a1
-; RV64IB-NEXT:    ret
-;
-; RV64IBA-LABEL: addmulw6:
-; RV64IBA:       # %bb.0:
-; RV64IBA-NEXT:    addi a2, zero, 6
-; RV64IBA-NEXT:    mul a0, a0, a2
-; RV64IBA-NEXT:    addw a0, a0, a1
-; RV64IBA-NEXT:    ret
-  %c = mul i32 %a, 6
-  %d = add i32 %c, %b
-  ret i32 %d
-}
-
-define i32 @addmulw10(i32 signext %a, i32 signext %b) {
-; RV64I-LABEL: addmulw10:
-; RV64I:       # %bb.0:
-; RV64I-NEXT:    addi a2, zero, 10
-; RV64I-NEXT:    mul a0, a0, a2
-; RV64I-NEXT:    addw a0, a0, a1
-; RV64I-NEXT:    ret
-;
-; RV64IB-LABEL: addmulw10:
-; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    addi a2, zero, 10
-; RV64IB-NEXT:    mul a0, a0, a2
-; RV64IB-NEXT:    addw a0, a0, a1
-; RV64IB-NEXT:    ret
-;
-; RV64IBA-LABEL: addmulw10:
-; RV64IBA:       # %bb.0:
-; RV64IBA-NEXT:    addi a2, zero, 10
-; RV64IBA-NEXT:    mul a0, a0, a2
-; RV64IBA-NEXT:    addw a0, a0, a1
-; RV64IBA-NEXT:    ret
-  %c = mul i32 %a, 10
-  %d = add i32 %c, %b
-  ret i32 %d
-}
-
-define i32 @addmulw12(i32 signext %a, i32 signext %b) {
-; RV64I-LABEL: addmulw12:
-; RV64I:       # %bb.0:
-; RV64I-NEXT:    addi a2, zero, 12
-; RV64I-NEXT:    mul a0, a0, a2
-; RV64I-NEXT:    addw a0, a0, a1
-; RV64I-NEXT:    ret
-;
-; RV64IB-LABEL: addmulw12:
-; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    addi a2, zero, 12
-; RV64IB-NEXT:    mul a0, a0, a2
-; RV64IB-NEXT:    addw a0, a0, a1
-; RV64IB-NEXT:    ret
-;
-; RV64IBA-LABEL: addmulw12:
-; RV64IBA:       # %bb.0:
-; RV64IBA-NEXT:    addi a2, zero, 12
-; RV64IBA-NEXT:    mul a0, a0, a2
-; RV64IBA-NEXT:    addw a0, a0, a1
-; RV64IBA-NEXT:    ret
-  %c = mul i32 %a, 12
-  %d = add i32 %c, %b
-  ret i32 %d
-}