[RISCV] Fold (sra (add (shl X, 32), C1), 32 - C) -> (shl (sext_inreg (add X, C1), C)

[craig.topper]

Summary

Similar for a subtract with a constant left hand side.

(sra (add (shl X, 32), C1<<32), 32) is the canonical IR from InstCombine
for (sext (add (trunc X to i32), 32) to i32).

For RISCV, we should lower this as addiw which means turning it into
(sext_inreg (add X, C1)).

There is an existing DAG combine to convert back to (sext (add (trunc X
to i32), 32) to i32), but it requires isTruncateFree to return true
and for i32 to be a legal type as it used sign_extend and truncate
nodes. So that doesn't work for RISCV.

If the outer sra happens be used by a shl by constant, it will be
folded and the shift amount of the sra will be changed before we
can do our own DAG combine. This requires us to match the more
general pattern and restore the shl.

I had wanted to do this as a separate (add (shl X, 32), C1<<32) ->
(shl (add X, C1), 32) combine, but that hit an infinite loop for some
values of C1.

[asb]

Really nice improvement - LGTM! (though I think the new tests aren't pre-committed yet or uploaded directly to phab)

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

if (NegAcc) {
     // clang-format on
   }
 
   return Opcode;
 }
 
 // Combine (sra (shl X, 32), 32 - C) -> (shl (sext_inreg X, i32), C)
 // FIXME: Should this be a generic combine? There's a similar combine on X86.
+//
+// Also try these folds where an add or sub is in the middle.
+// (sra (add (shl X, 32), C1), 32 - C) -> (shl (sext_inreg (add X, C1), C)
+// (sra (sub C1, (shl X, 32)), 32 - C) -> (shl (sext_inreg (sub C1, X), C)
 static SDValue performSRACombine(SDNode *N, SelectionDAG &DAG,
                                  const RISCVSubtarget &Subtarget) {
   assert(N->getOpcode() == ISD::SRA && "Unexpected opcode");
 
   if (N->getValueType(0) != MVT::i64 || !Subtarget.is64Bit())
     return SDValue();
 
-  auto *C = dyn_cast<ConstantSDNode>(N->getOperand(1));
-  if (!C || C->getZExtValue() >= 32)
+  auto *ShAmtC = dyn_cast<ConstantSDNode>(N->getOperand(1));
+  if (!ShAmtC || ShAmtC->getZExtValue() > 32)
     return SDValue();
 
   SDValue N0 = N->getOperand(0);
-  if (N0.getOpcode() != ISD::SHL || !N0.hasOneUse() ||
-      !isa<ConstantSDNode>(N0.getOperand(1)) ||
-      N0.getConstantOperandVal(1) != 32)
+
+  SDValue Shl;
+  ConstantSDNode *AddC = nullptr;
+
+  // We might have an ADD or SUB between the SRA and SHL.
+  bool IsAdd = N0.getOpcode() == ISD::ADD;
+  if ((IsAdd || N0.getOpcode() == ISD::SUB)) {
+    if (!N0.hasOneUse())
+      return SDValue();
+    // Other operand needs to be a constant we can modify.
+    AddC = dyn_cast<ConstantSDNode>(N0.getOperand(IsAdd ? 1 : 0));
+    if (!AddC)
+      return SDValue();
+
+    // AddC needs to have at least 32 trailing zeros.
+    if (AddC->getAPIntValue().countTrailingZeros() < 32)
+      return SDValue();
+
+    Shl = N0.getOperand(IsAdd ? 0 : 1);
+  } else {
+    // Not an ADD or SUB.
+    Shl = N0;
+  }
+
+  // Look for a shift left by 32.
+  if (Shl.getOpcode() != ISD::SHL || !Shl.hasOneUse() ||
+      !isa<ConstantSDNode>(Shl.getOperand(1)) ||
+      Shl.getConstantOperandVal(1) != 32)
     return SDValue();
 
   SDLoc DL(N);
-  SDValue SExt = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, MVT::i64,
-                             N0.getOperand(0), DAG.getValueType(MVT::i32));
-  return DAG.getNode(ISD::SHL, DL, MVT::i64, SExt,
-                     DAG.getConstant(32 - C->getZExtValue(), DL, MVT::i64));
+  SDValue In = Shl.getOperand(0);
+
+  // If we looked through an ADD or SUB, we need to rebuild it with the shifted
+  // constant.
+  if (AddC) {
+    SDValue ShiftedAddC =
+        DAG.getConstant(AddC->getAPIntValue().lshr(32), DL, MVT::i64);
+    if (IsAdd)
+      In = DAG.getNode(ISD::ADD, DL, MVT::i64, In, ShiftedAddC);
+    else
+      In = DAG.getNode(ISD::SUB, DL, MVT::i64, ShiftedAddC, In);
+  }
+
+  SDValue SExt = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, MVT::i64, In,
+                             DAG.getValueType(MVT::i32));
+  if (ShAmtC->getZExtValue() == 32)
+    return SExt;
+
+  return DAG.getNode(
+      ISD::SHL, DL, MVT::i64, SExt,
+      DAG.getConstant(32 - ShAmtC->getZExtValue(), DL, MVT::i64));
 }
 
 SDValue RISCVTargetLowering::PerformDAGCombine(SDNode *N,
                                                DAGCombinerInfo &DCI) const {
   SelectionDAG &DAG = DCI.DAG;
 
   // Helper to call SimplifyDemandedBits on an operand of N where only some low
   // bits are demanded. N will be added to the Worklist if it was not deleted.

llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll

 }
 
 ; The ashr+add+shl is canonical IR from InstCombine for
 ; (sext (add (trunc X to i32), 1) to i32).
 ; That can be implemented as addiw make sure we recover it.
 define i64 @test7(i32* %0, i64 %1) {
 ; RV64I-LABEL: test7:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    slli a0, a1, 32
-; RV64I-NEXT:    li a1, 1
-; RV64I-NEXT:    slli a1, a1, 32
-; RV64I-NEXT:    add a0, a0, a1
-; RV64I-NEXT:    srai a0, a0, 32
+; RV64I-NEXT:    addiw a0, a1, 1
 ; RV64I-NEXT:    ret
   %3 = shl i64 %1, 32
   %4 = add i64 %3, 4294967296
   %5 = ashr exact i64 %4, 32
   ret i64 %5
 }
 
 ; The ashr+add+shl is canonical IR from InstCombine for
 ; (sext (sub 1, (trunc X to i32)) to i32).
 ; That can be implemented as (li 1)+subw make sure we recover it.
 define i64 @test8(i32* %0, i64 %1) {
 ; RV64I-LABEL: test8:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    slli a0, a1, 32
-; RV64I-NEXT:    li a1, 1
-; RV64I-NEXT:    slli a1, a1, 32
-; RV64I-NEXT:    sub a0, a1, a0
-; RV64I-NEXT:    srai a0, a0, 32
+; RV64I-NEXT:    li a0, 1
+; RV64I-NEXT:    subw a0, a0, a1
 ; RV64I-NEXT:    ret
   %3 = mul i64 %1, -4294967296
   %4 = add i64 %3, 4294967296
   %5 = ashr exact i64 %4, 32
   ret i64 %5
 }
 
 ; The gep is here to introduce a shl by 2 after the ashr that will get folded
 ; and make this harder to recover.
 define signext i32 @test9(i32* %0, i64 %1) {
 ; RV64I-LABEL: test9:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    slli a1, a1, 32
-; RV64I-NEXT:    lui a2, 4097
-; RV64I-NEXT:    slli a2, a2, 20
-; RV64I-NEXT:    add a1, a1, a2
-; RV64I-NEXT:    srai a1, a1, 30
+; RV64I-NEXT:    lui a2, 1
+; RV64I-NEXT:    addiw a2, a2, 1
+; RV64I-NEXT:    addw a1, a1, a2
+; RV64I-NEXT:    slli a1, a1, 2
 ; RV64I-NEXT:    add a0, a0, a1
 ; RV64I-NEXT:    lw a0, 0(a0)
 ; RV64I-NEXT:    ret
   %3 = shl i64 %1, 32
   %4 = add i64 %3, 17596481011712 ; 4097 << 32
   %5 = ashr exact i64 %4, 32
   %6 = getelementptr inbounds i32, i32* %0, i64 %5
   %7 = load i32, i32* %6, align 4
   ret i32 %7
 }
 
 ; The gep is here to introduce a shl by 2 after the ashr that will get folded
 ; and make this harder to recover.
 define signext i32 @test10(i32* %0, i64 %1) {
 ; RV64I-LABEL: test10:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    slli a1, a1, 32
 ; RV64I-NEXT:    lui a2, 30141
 ; RV64I-NEXT:    addiw a2, a2, -747
-; RV64I-NEXT:    slli a2, a2, 32
-; RV64I-NEXT:    sub a1, a2, a1
-; RV64I-NEXT:    srai a1, a1, 30
+; RV64I-NEXT:    subw a1, a2, a1
+; RV64I-NEXT:    slli a1, a1, 2
 ; RV64I-NEXT:    add a0, a0, a1
 ; RV64I-NEXT:    lw a0, 0(a0)
 ; RV64I-NEXT:    ret
   %3 = mul i64 %1, -4294967296
   %4 = add i64 %3, 530242871224172544 ; 123456789 << 32
   %5 = ashr exact i64 %4, 32
   %6 = getelementptr inbounds i32, i32* %0, i64 %5
   %7 = load i32, i32* %6, align 4
   ret i32 %7
 }
 
 define i64 @test11(i32* %0, i64 %1) {
 ; RV64I-LABEL: test11:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    slli a0, a1, 32
-; RV64I-NEXT:    li a1, -1
-; RV64I-NEXT:    slli a1, a1, 63
-; RV64I-NEXT:    sub a0, a1, a0
-; RV64I-NEXT:    srai a0, a0, 32
+; RV64I-NEXT:    lui a0, 524288
+; RV64I-NEXT:    subw a0, a0, a1
 ; RV64I-NEXT:    ret
   %3 = mul i64 %1, -4294967296
   %4 = add i64 %3, 9223372036854775808 ;0x8000'0000'0000'0000
   %5 = ashr exact i64 %4, 32
   ret i64 %5
 }