[llvm] [RISCV] Add combine for shadd family of instructions. (PR #130829)

[Craig Topper via llvm-commits]

================
@@ -14306,6 +14313,84 @@ static SDValue transformAddShlImm(SDNode *N, SelectionDAG &DAG,
   return DAG.getNode(ISD::SHL, DL, VT, SHADD, DAG.getConstant(Bits, DL, VT));
 }
 
+// Check if this SDValue is an add immediate that is fed by a shift of 1, 2,
+// or 3.
+static bool checkAddiForShift(SDValue AddI, int64_t &AddConst,
+                              int64_t &ShlConst) {
+  using namespace llvm::SDPatternMatch;
+  // Based on testing it seems that performance degrades if the ADDI has
+  // more than 2 uses.
+  if (AddI->use_size() > 2)
+    return false;
+
+  APInt AddVal;
+  SDValue SHLVal;
+  assert(sd_match(AddI, m_Add(m_Value(SHLVal), m_ConstInt(AddVal))) &&
+         "Expected an addi with a constant addition.");
+
+  APInt VShift;
+  if (!sd_match(SHLVal, m_BinOp(ISD::SHL, m_Value(), m_ConstInt(VShift))))
+    return false;
+
+  if (VShift.slt(1) || VShift.sgt(3))
+    return false;
+
+  // Set the values at the end when we know that the function will return
+  // true.
+  ShlConst = VShift.getSExtValue();
+  AddConst = AddVal.getSExtValue();
+  return true;
+}
+
+// Optimize (add (add (shl x, c0),  c1), y) ->
+//          (ADDI (SH*ADD y, x), c1), if c0 equals to [1|2|3].
+static SDValue combineShlAddIAdd(SDNode *N, SelectionDAG &DAG,
+                                 const RISCVSubtarget &Subtarget) {
+  using namespace llvm::SDPatternMatch;
+
+  // Perform this optimization only in the zba extension.
+  if (!ReassocShlAddiAdd || !Subtarget.hasStdExtZba())
+    return SDValue();
+
+  // Skip for vector types and larger types.
+  EVT VT = N->getValueType(0);
+  if (VT != Subtarget.getXLenVT())
+    return SDValue();
+
+  // Looking for a reg-reg add and not an addi.
+  if (isa<ConstantSDNode>(N->getOperand(1)))
+    return SDValue();
+
+  SDValue AddI = N->getOperand(0);
+  SDValue Other = N->getOperand(1);
+  bool LHSIsAddI = sd_match(AddI, m_Add(m_Value(), m_ConstInt()));
+  bool RHSIsAddI = sd_match(Other, m_Add(m_Value(), m_ConstInt()));
+  int64_t AddConst = 0;
+  int64_t ShlConst = 0;
+
+  // At least one addi is required.
+  if (!LHSIsAddI && !RHSIsAddI)
+    return SDValue();
+
+  // If the LHS is not the result of an add or both sides are results of an add,
+  // but the LHS does not have the desired structure with a shift, swap the
+  // operands.
+  if (!LHSIsAddI || (RHSIsAddI && !checkAddiForShift(AddI, AddConst, ShlConst)))
+    std::swap(AddI, Other);
+
+  // We simply need to ensure AddI has the desired structure.
+  if (!checkAddiForShift(AddI, AddConst, ShlConst))
+    return SDValue();
+
+  SDValue SHLVal = AddI->getOperand(0);
+  SDLoc DL(N);
+
+  SDValue SHADD = DAG.getNode(RISCVISD::SHL_ADD, DL, VT, SHLVal->getOperand(0),
----------------
topperc wrote:

Can we push this code into the end of checkAddiForShift? If the match is successful just do the transform and return the SDValue.

https://github.com/llvm/llvm-project/pull/130829