[llvm] [RISCV] Use RISCVISD::SHL_ADD in transformAddShlImm [NFC] (PR #89832)

[Philip Reames via llvm-commits]

https://github.com/preames created https://github.com/llvm/llvm-project/pull/89832

Doing so avoids negative interactions with other combines which don't know the shl_add is a single instruction.  From the commit log, we've had several combine loops already.

This was originally posted as part of #88791, where a bug was pointed out. That bug was fixed by #89789 which hits the same issue from another angle. To confirm the fix, I included the reduced test case here.

>From 4c0b7d815b508cc48ed4573df343cf74b4fa7d30 Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Thu, 18 Apr 2024 13:33:38 -0700
Subject: [PATCH] [RISCV] Use RISCVISD::SHL_ADD in transformAddShlImm [NFC]

Doing so avoids negative interactions with other combines which don't
know the shl_add is a single instruction.  From the commit log, we've had
several combine loops already.

This was originally posted as part of #88791, where a bug was pointed out.
That bug was fixed by #89789 which hits the same issue from another angle.
To confirm the fix, I included the reduce test case here.
---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp | 18 +++++++-------
 llvm/test/CodeGen/RISCV/addimm-mulimm.ll    | 26 +++++++++++++++++++++
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index 9c66f09a0cbc85..e6c765d5fc1d51 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -12817,10 +12817,9 @@ static SDValue transformAddShlImm(SDNode *N, SelectionDAG &DAG,
   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));
+  SDValue SHADD =
+      DAG.getNode(RISCVISD::SHL_ADD, DL, VT, NL, DAG.getConstant(Diff, DL, VT), NS);
+  return DAG.getNode(ISD::SHL, DL, VT, SHADD, DAG.getConstant(Bits, DL, VT));
 }
 
 // Combine a constant select operand into its use:
@@ -13056,14 +13055,17 @@ static SDValue combineAddOfBooleanXor(SDNode *N, SelectionDAG &DAG) {
                      N0.getOperand(0));
 }
 
-static SDValue performADDCombine(SDNode *N, SelectionDAG &DAG,
+static SDValue performADDCombine(SDNode *N,
+                                 TargetLowering::DAGCombinerInfo &DCI,
                                  const RISCVSubtarget &Subtarget) {
+  SelectionDAG &DAG = DCI.DAG;
   if (SDValue V = combineAddOfBooleanXor(N, DAG))
     return V;
   if (SDValue V = transformAddImmMulImm(N, DAG, Subtarget))
     return V;
-  if (SDValue V = transformAddShlImm(N, DAG, Subtarget))
-    return V;
+  if (!DCI.isBeforeLegalize() && !DCI.isCalledByLegalizer())
+    if (SDValue V = transformAddShlImm(N, DAG, Subtarget))
+      return V;
   if (SDValue V = combineBinOpToReduce(N, DAG, Subtarget))
     return V;
   if (SDValue V = combineBinOpOfExtractToReduceTree(N, DAG, Subtarget))
@@ -16027,7 +16029,7 @@ SDValue RISCVTargetLowering::PerformDAGCombine(SDNode *N,
       return V;
     if (SDValue V = combineToVWMACC(N, DAG, Subtarget))
       return V;
-    return performADDCombine(N, DAG, Subtarget);
+    return performADDCombine(N, DCI, Subtarget);
   }
   case ISD::SUB: {
     if (SDValue V = combineBinOp_VLToVWBinOp_VL(N, DCI, Subtarget))
diff --git a/llvm/test/CodeGen/RISCV/addimm-mulimm.ll b/llvm/test/CodeGen/RISCV/addimm-mulimm.ll
index 736c8e7d55c75b..190b38cb577355 100644
--- a/llvm/test/CodeGen/RISCV/addimm-mulimm.ll
+++ b/llvm/test/CodeGen/RISCV/addimm-mulimm.ll
@@ -898,3 +898,29 @@ define i1 @pr53831(i32 %x) {
   %tmp5 = icmp eq i32 %tmp4, %tmp2
   ret i1 %tmp5
 }
+
+define i64 @sh2add_uw(i64 signext %0, i32 signext %1) {
+; RV32IMB-LABEL: sh2add_uw:
+; RV32IMB:       # %bb.0: # %entry
+; RV32IMB-NEXT:    srli a3, a2, 27
+; RV32IMB-NEXT:    slli a2, a2, 5
+; RV32IMB-NEXT:    srli a4, a0, 29
+; RV32IMB-NEXT:    sh3add a1, a1, a4
+; RV32IMB-NEXT:    sh3add a0, a0, a2
+; RV32IMB-NEXT:    sltu a2, a0, a2
+; RV32IMB-NEXT:    add a1, a3, a1
+; RV32IMB-NEXT:    add a1, a1, a2
+; RV32IMB-NEXT:    ret
+;
+; RV64IMB-LABEL: sh2add_uw:
+; RV64IMB:       # %bb.0: # %entry
+; RV64IMB-NEXT:    sh2add.uw a0, a1, a0
+; RV64IMB-NEXT:    slli a0, a0, 3
+; RV64IMB-NEXT:    ret
+entry:
+  %2 = zext i32 %1 to i64
+  %3 = shl i64 %2, 5
+  %4 = shl i64 %0, 3
+  %5 = add i64 %3, %4
+  ret i64 %5
+}