[RISCV] Optimize (add (shl x, c0), (shl y, c1)) with SH*ADD

[benshi001]

Summary

Optimize (add (shl x, c0), (shl y, c1)) ->

(SLLI (SH*ADD x, y), c1), if c0-c1 == 1/2/3.

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5790

You can't depend on this. I believe the add node will be visited before the shift nodes are visited to replace any negative shift amounts with undef. So if the IR contains a negative shift amount I think you will fail this assertion.

6233

You can just check

if (NA)

SDValue operator bool() checks getNode() != nullptr.

[benshi001]

ping ...

[benshi001]

ping...

Can this patch be revriewed and accepted? It is safe and does improve code quality.

[luismarques]

LGTM.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5771

Nit: vector

5793–5794

Is this just abs(C1 - C0)?

[luismarques]

I had started reviewing this before D109729, now I'm very confused about the point of the two patches...

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5773–5774

See craig's comment in D109729: The VT.getSizeInBits() > Subtarget->getXLen() should never fail. If the VT is scalar it must be XLenVT.

[benshi001]

In D108916#3003531, @luismarques wrote:

I had started reviewing this before D109729, now I'm very confused about the point of the two patches...

The reason I created https://reviews.llvm.org/D109729 is that current patch has been pending for weeks and I thought you disagree with it.

So I created https://reviews.llvm.org/D109729, which initial version was the same optimization but implemented in TD rules. However TD rules were in bad form so I changed it to DAG2GAG.

However I think both patches are OK. If current one can be accepted, then I will close https://reviews.llvm.org/D109729.

[benshi001]

The build failure x64 debian > libomp.api::omp_get_wtime.c seems not related to my patch.

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5773–5774

This comment was not correct for this patch. The XLen check is needed here.

6232

Move this into performADDCombine

[luismarques]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5773–5774

This comment was not correct for this patch. The XLen check is needed here.

Yup, my bad. Sorry!

[benshi001]

Also I have moved transformAddShlImm into performADDCombine

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5784

NC0 -> N0C as I think that's the style DAGCombiner.cpp uses.

5917

Don't N0 for the name here. N0 usually means N->getOperand(0)

Use

if (SDValue V = transformAddShlImm(N, DAG, Subtarget))
  return V;

V is pretty generic and putting it in the if limits the scope.

[craig.topper]

LGTM

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

static SDValue combineORToSHFL(SDValue Op, SelectionDAG &DAG,
   if (Mask != ExpMask)
     return SDValue();
 
   SDLoc DL(Op);
   return DAG.getNode(RISCVISD::SHFL, DL, VT, Match1->Op,
                      DAG.getConstant(Match1->ShAmt, DL, VT));
 }
 
+// Optimize (add (shl x, c0), (shl y, c1)) ->
+//          (SLLI (SH*ADD x, y), c0), if c1-c0 equals to [1|2|3].
+static SDValue transformAddShlImm(SDNode *N, SelectionDAG &DAG,
+                                  const RISCVSubtarget &Subtarget) {
+  // Perform this optimization only in the zba extension.
+  if (!Subtarget.hasStdExtZba())
+    return SDValue();
+
+  // Skip for vector types and larger types.
+  EVT VT = N->getValueType(0);
+  if (VT.isVector() || VT.getSizeInBits() > Subtarget.getXLen())
+    return SDValue();
+
+  // The two operand nodes must be SHL and have no other use.
+  SDValue N0 = N->getOperand(0);
+  SDValue N1 = N->getOperand(1);
+  if (N0->getOpcode() != ISD::SHL || N1->getOpcode() != ISD::SHL ||
+      !N0->hasOneUse() || !N1->hasOneUse())
+    return SDValue();
+
+  // Check c0 and c1.
+  auto *NC0 = dyn_cast<ConstantSDNode>(N0->getOperand(1));
+  auto *NC1 = dyn_cast<ConstantSDNode>(N1->getOperand(1));
+  if (!NC0 || !NC1)
+    return SDValue();
+  int64_t C0 = NC0->getSExtValue();
+  int64_t C1 = NC1->getSExtValue();
+  if (C0 <= 0 || C1 <= 0)
+    return SDValue();
+
+  // Skip if SH1ADD/SH2ADD/SH3ADD are not applicable.
+  int64_t Bits = std::min(C0, C1);
+  int64_t Diff = std::abs(C0 - C1);
+  if (Diff != 1 && Diff != 2 && Diff != 3)
+    return SDValue();
+
+  // Build nodes.
+  SDLoc DL(N);
+  SDValue NS = (C0 < C1) ? N0->getOperand(0) : N1->getOperand(0);
+  SDValue NL = (C0 > C1) ? N0->getOperand(0) : N1->getOperand(0);
+  SDValue NA0 =
+      DAG.getNode(ISD::SHL, DL, VT, NL, DAG.getConstant(Diff, DL, VT));
+  SDValue NA1 = DAG.getNode(ISD::ADD, DL, VT, NA0, NS);
+  return DAG.getNode(ISD::SHL, DL, VT, NA1, DAG.getConstant(Bits, DL, VT));
+}
+
 // Combine (GREVI (GREVI x, C2), C1) -> (GREVI x, C1^C2) when C1^C2 is
 // non-zero, and to x when it is. Any repeated GREVI stage undoes itself.
 // Combine (GORCI (GORCI x, C2), C1) -> (GORCI x, C1|C2). Repeated stage does
 // not undo itself, but they are redundant.
 static SDValue combineGREVI_GORCI(SDNode *N, SelectionDAG &DAG) {
   SDValue Src = N->getOperand(0);
 
   if (Src.getOpcode() != N->getOpcode())
static SDValue combineSelectAndUseCommutative(SDNode *N, SelectionDAG &DAG,
   SDValue N1 = N->getOperand(1);
   if (SDValue Result = combineSelectAndUse(N, N0, N1, DAG, AllOnes))
     return Result;
   if (SDValue Result = combineSelectAndUse(N, N1, N0, DAG, AllOnes))
     return Result;
   return SDValue();
 }
 
-static SDValue performADDCombine(SDNode *N, SelectionDAG &DAG) {
+static SDValue performADDCombine(SDNode *N, SelectionDAG &DAG,
+                                 const RISCVSubtarget &Subtarget) {
+  // Fold (add (shl x, c0), (shl y, c1)) ->
+  //      (SLLI (SH*ADD x, y), c0), if c1-c0 equals to [1|2|3].
+  SDValue N0 = transformAddShlImm(N, DAG, Subtarget);
+  if (N0)
+    return N0;
   // fold (add (select lhs, rhs, cc, 0, y), x) ->
   //      (select lhs, rhs, cc, x, (add x, y))
   return combineSelectAndUseCommutative(N, DAG, /*AllOnes*/ false);
 }
 
 static SDValue performSUBCombine(SDNode *N, SelectionDAG &DAG) {
   // fold (sub x, (select lhs, rhs, cc, 0, y)) ->
   //      (select lhs, rhs, cc, x, (sub x, y))
if (Op0.getOpcode() == ISD::FNEG)
       return DAG.getNode(ISD::XOR, DL, VT, NewFMV,
                          DAG.getConstant(SignBit, DL, VT));
 
     assert(Op0.getOpcode() == ISD::FABS);
     return DAG.getNode(ISD::AND, DL, VT, NewFMV,
                        DAG.getConstant(~SignBit, DL, VT));
   }
   case ISD::ADD:
-    return performADDCombine(N, DAG);
+    return performADDCombine(N, DAG, Subtarget);
   case ISD::SUB:
     return performSUBCombine(N, DAG);
   case ISD::AND:
     return performANDCombine(N, DAG);
   case ISD::OR:
     return performORCombine(N, DAG, Subtarget);
   case ISD::XOR:
     return performXORCombine(N, DAG);
   case ISD::ANY_EXTEND:

llvm/test/CodeGen/RISCV/rv32zba.ll

 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    slli a0, a0, 5
 ; RV32I-NEXT:    slli a1, a1, 6
 ; RV32I-NEXT:    add a0, a0, a1
 ; RV32I-NEXT:    ret
 ;
 ; RV32B-LABEL: addshl_5_6:
 ; RV32B:       # %bb.0:
+; RV32B-NEXT:    sh1add a0, a1, a0
 ; RV32B-NEXT:    slli a0, a0, 5
-; RV32B-NEXT:    slli a1, a1, 6
-; RV32B-NEXT:    add a0, a0, a1
 ; RV32B-NEXT:    ret
 ;
 ; RV32ZBA-LABEL: addshl_5_6:
 ; RV32ZBA:       # %bb.0:
+; RV32ZBA-NEXT:    sh1add a0, a1, a0
 ; RV32ZBA-NEXT:    slli a0, a0, 5
-; RV32ZBA-NEXT:    slli a1, a1, 6
-; RV32ZBA-NEXT:    add a0, a0, a1
 ; RV32ZBA-NEXT:    ret
   %c = shl i32 %a, 5
   %d = shl i32 %b, 6
   %e = add i32 %c, %d
   ret i32 %e
 }
 
 define i32 @addshl_5_7(i32 %a, i32 %b) {
 ; RV32I-LABEL: addshl_5_7:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    slli a0, a0, 5
 ; RV32I-NEXT:    slli a1, a1, 7
 ; RV32I-NEXT:    add a0, a0, a1
 ; RV32I-NEXT:    ret
 ;
 ; RV32B-LABEL: addshl_5_7:
 ; RV32B:       # %bb.0:
+; RV32B-NEXT:    sh2add a0, a1, a0
 ; RV32B-NEXT:    slli a0, a0, 5
-; RV32B-NEXT:    slli a1, a1, 7
-; RV32B-NEXT:    add a0, a0, a1
 ; RV32B-NEXT:    ret
 ;
 ; RV32ZBA-LABEL: addshl_5_7:
 ; RV32ZBA:       # %bb.0:
+; RV32ZBA-NEXT:    sh2add a0, a1, a0
 ; RV32ZBA-NEXT:    slli a0, a0, 5
-; RV32ZBA-NEXT:    slli a1, a1, 7
-; RV32ZBA-NEXT:    add a0, a0, a1
 ; RV32ZBA-NEXT:    ret
   %c = shl i32 %a, 5
   %d = shl i32 %b, 7
   %e = add i32 %c, %d
   ret i32 %e
 }
 
 define i32 @addshl_5_8(i32 %a, i32 %b) {
 ; RV32I-LABEL: addshl_5_8:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    slli a0, a0, 5
 ; RV32I-NEXT:    slli a1, a1, 8
 ; RV32I-NEXT:    add a0, a0, a1
 ; RV32I-NEXT:    ret
 ;
 ; RV32B-LABEL: addshl_5_8:
 ; RV32B:       # %bb.0:
+; RV32B-NEXT:    sh3add a0, a1, a0
 ; RV32B-NEXT:    slli a0, a0, 5
-; RV32B-NEXT:    slli a1, a1, 8
-; RV32B-NEXT:    add a0, a0, a1
 ; RV32B-NEXT:    ret
 ;
 ; RV32ZBA-LABEL: addshl_5_8:
 ; RV32ZBA:       # %bb.0:
+; RV32ZBA-NEXT:    sh3add a0, a1, a0
 ; RV32ZBA-NEXT:    slli a0, a0, 5
-; RV32ZBA-NEXT:    slli a1, a1, 8
-; RV32ZBA-NEXT:    add a0, a0, a1
 ; RV32ZBA-NEXT:    ret
   %c = shl i32 %a, 5
   %d = shl i32 %b, 8
   %e = add i32 %c, %d
   ret i32 %e
 }

llvm/test/CodeGen/RISCV/rv64zba.ll

 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    slliw a0, a0, 5
 ; RV64I-NEXT:    slliw a1, a1, 6
 ; RV64I-NEXT:    addw a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64B-LABEL: addshl32_5_6:
 ; RV64B:       # %bb.0:
+; RV64B-NEXT:    sh1add a0, a1, a0
 ; RV64B-NEXT:    slliw a0, a0, 5
-; RV64B-NEXT:    slliw a1, a1, 6
-; RV64B-NEXT:    addw a0, a0, a1
 ; RV64B-NEXT:    ret
 ;
 ; RV64ZBA-LABEL: addshl32_5_6:
 ; RV64ZBA:       # %bb.0:
+; RV64ZBA-NEXT:    sh1add a0, a1, a0
 ; RV64ZBA-NEXT:    slliw a0, a0, 5
-; RV64ZBA-NEXT:    slliw a1, a1, 6
-; RV64ZBA-NEXT:    addw a0, a0, a1
 ; RV64ZBA-NEXT:    ret
   %c = shl i32 %a, 5
   %d = shl i32 %b, 6
   %e = add i32 %c, %d
   ret i32 %e
 }
 
 define i64 @addshl64_5_6(i64 %a, i64 %b) {
 ; RV64I-LABEL: addshl64_5_6:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    slli a0, a0, 5
 ; RV64I-NEXT:    slli a1, a1, 6
 ; RV64I-NEXT:    add a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64B-LABEL: addshl64_5_6:
 ; RV64B:       # %bb.0:
+; RV64B-NEXT:    sh1add a0, a1, a0
 ; RV64B-NEXT:    slli a0, a0, 5
-; RV64B-NEXT:    slli a1, a1, 6
-; RV64B-NEXT:    add a0, a0, a1
 ; RV64B-NEXT:    ret
 ;
 ; RV64ZBA-LABEL: addshl64_5_6:
 ; RV64ZBA:       # %bb.0:
+; RV64ZBA-NEXT:    sh1add a0, a1, a0
 ; RV64ZBA-NEXT:    slli a0, a0, 5
-; RV64ZBA-NEXT:    slli a1, a1, 6
-; RV64ZBA-NEXT:    add a0, a0, a1
 ; RV64ZBA-NEXT:    ret
   %c = shl i64 %a, 5
   %d = shl i64 %b, 6
   %e = add i64 %c, %d
   ret i64 %e
 }
 
 define signext i32 @addshl32_5_7(i32 signext %a, i32 signext %b) {
 ; RV64I-LABEL: addshl32_5_7:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    slliw a0, a0, 5
 ; RV64I-NEXT:    slliw a1, a1, 7
 ; RV64I-NEXT:    addw a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64B-LABEL: addshl32_5_7:
 ; RV64B:       # %bb.0:
+; RV64B-NEXT:    sh2add a0, a1, a0
 ; RV64B-NEXT:    slliw a0, a0, 5
-; RV64B-NEXT:    slliw a1, a1, 7
-; RV64B-NEXT:    addw a0, a0, a1
 ; RV64B-NEXT:    ret
 ;
 ; RV64ZBA-LABEL: addshl32_5_7:
 ; RV64ZBA:       # %bb.0:
+; RV64ZBA-NEXT:    sh2add a0, a1, a0
 ; RV64ZBA-NEXT:    slliw a0, a0, 5
-; RV64ZBA-NEXT:    slliw a1, a1, 7
-; RV64ZBA-NEXT:    addw a0, a0, a1
 ; RV64ZBA-NEXT:    ret
   %c = shl i32 %a, 5
   %d = shl i32 %b, 7
   %e = add i32 %c, %d
   ret i32 %e
 }
 
 define i64 @addshl64_5_7(i64 %a, i64 %b) {
 ; RV64I-LABEL: addshl64_5_7:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    slli a0, a0, 5
 ; RV64I-NEXT:    slli a1, a1, 7
 ; RV64I-NEXT:    add a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64B-LABEL: addshl64_5_7:
 ; RV64B:       # %bb.0:
+; RV64B-NEXT:    sh2add a0, a1, a0
 ; RV64B-NEXT:    slli a0, a0, 5
-; RV64B-NEXT:    slli a1, a1, 7
-; RV64B-NEXT:    add a0, a0, a1
 ; RV64B-NEXT:    ret
 ;
 ; RV64ZBA-LABEL: addshl64_5_7:
 ; RV64ZBA:       # %bb.0:
+; RV64ZBA-NEXT:    sh2add a0, a1, a0
 ; RV64ZBA-NEXT:    slli a0, a0, 5
-; RV64ZBA-NEXT:    slli a1, a1, 7
-; RV64ZBA-NEXT:    add a0, a0, a1
 ; RV64ZBA-NEXT:    ret
   %c = shl i64 %a, 5
   %d = shl i64 %b, 7
   %e = add i64 %c, %d
   ret i64 %e
 }
 
 define signext i32 @addshl32_5_8(i32 signext %a, i32 signext %b) {
 ; RV64I-LABEL: addshl32_5_8:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    slliw a0, a0, 5
 ; RV64I-NEXT:    slliw a1, a1, 8
 ; RV64I-NEXT:    addw a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64B-LABEL: addshl32_5_8:
 ; RV64B:       # %bb.0:
+; RV64B-NEXT:    sh3add a0, a1, a0
 ; RV64B-NEXT:    slliw a0, a0, 5
-; RV64B-NEXT:    slliw a1, a1, 8
-; RV64B-NEXT:    addw a0, a0, a1
 ; RV64B-NEXT:    ret
 ;
 ; RV64ZBA-LABEL: addshl32_5_8:
 ; RV64ZBA:       # %bb.0:
+; RV64ZBA-NEXT:    sh3add a0, a1, a0
 ; RV64ZBA-NEXT:    slliw a0, a0, 5
-; RV64ZBA-NEXT:    slliw a1, a1, 8
-; RV64ZBA-NEXT:    addw a0, a0, a1
 ; RV64ZBA-NEXT:    ret
   %c = shl i32 %a, 5
   %d = shl i32 %b, 8
   %e = add i32 %c, %d
   ret i32 %e
 }
 
 define i64 @addshl64_5_8(i64 %a, i64 %b) {
 ; RV64I-LABEL: addshl64_5_8:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    slli a0, a0, 5
 ; RV64I-NEXT:    slli a1, a1, 8
 ; RV64I-NEXT:    add a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64B-LABEL: addshl64_5_8:
 ; RV64B:       # %bb.0:
+; RV64B-NEXT:    sh3add a0, a1, a0
 ; RV64B-NEXT:    slli a0, a0, 5
-; RV64B-NEXT:    slli a1, a1, 8
-; RV64B-NEXT:    add a0, a0, a1
 ; RV64B-NEXT:    ret
 ;
 ; RV64ZBA-LABEL: addshl64_5_8:
 ; RV64ZBA:       # %bb.0:
+; RV64ZBA-NEXT:    sh3add a0, a1, a0
 ; RV64ZBA-NEXT:    slli a0, a0, 5
-; RV64ZBA-NEXT:    slli a1, a1, 8
-; RV64ZBA-NEXT:    add a0, a0, a1
 ; RV64ZBA-NEXT:    ret
   %c = shl i64 %a, 5
   %d = shl i64 %b, 8
   %e = add i64 %c, %d
   ret i64 %e
 }