[llvm] [AArch64][SVE] Use ADD/ADR instead of MUL/MLA for x*N (PR #198566)

[Tomas Matheson via llvm-commits]

================
@@ -20579,10 +20618,169 @@ static bool isOnlyUsedAsMemoryAddress(SDNode *N) {
   return true;
 }
 
+/// Build an all-active SVE left shift by a constant element value.
+static SDValue buildSVEShiftLeftImmediate(SDValue Op, unsigned Shift,
+                                          SelectionDAG &DAG, SDLoc DL) {
+  EVT VT = Op.getValueType();
+  assert(VT.isScalableVector() && "Expected scalable vector type");
+
+  EVT ScalarVT = VT.getScalarType();
+  SDValue ShiftSplat = DAG.getNode(ISD::SPLAT_VECTOR, DL, VT,
+                                   DAG.getConstant(Shift, DL, ScalarVT));
+  return DAG.getNode(AArch64ISD::SHL_PRED, DL, VT,
+                     getPredicateForVector(DAG, DL, VT), Op, ShiftSplat);
+}
+
+static bool isSupportedSVEMulByConstantVT(EVT VT) {
+  return VT.isScalableVector() &&
+         (VT.getScalarType() == MVT::i32 || VT.getScalarType() == MVT::i64);
+}
+
+/// Combine an SVE multiply by a constant that does not need MUL:
+///   x * 2     -> x + x
+///   x * 3/5/9 -> x + (x << 1/2/3), which can select ADR.
+static SDValue performSVEMulByConstantCombine(SDValue Op, uint64_t Imm,
+                                              SelectionDAG &DAG, SDLoc DL) {
+  switch (Imm) {
+  default:
+    return SDValue();
+  case 2:
+    return DAG.getNode(ISD::ADD, DL, Op.getValueType(), Op, Op);
+  case 3:
+    return DAG.getNode(ISD::ADD, DL, Op.getValueType(), Op,
+                       buildSVEShiftLeftImmediate(Op, 1, DAG, DL));
+  case 5:
+    return DAG.getNode(ISD::ADD, DL, Op.getValueType(), Op,
+                       buildSVEShiftLeftImmediate(Op, 2, DAG, DL));
+  case 9:
+    return DAG.getNode(ISD::ADD, DL, Op.getValueType(), Op,
+                       buildSVEShiftLeftImmediate(Op, 3, DAG, DL));
+  }
+}
+
+/// Build the scaled multiply operand for:
+///   acc + x * 1     -> acc + x
+///   acc + x * 2/4/8 -> acc + (x << 1/2/3)
+/// The final ADD can then select ADR.
+static SDValue buildSVEScaledMulOperand(SDValue Op, uint64_t Imm,
+                                        SelectionDAG &DAG, SDLoc DL) {
+  switch (Imm) {
+  default:
+    return SDValue();
+  case 1:
+    return Op;
+  case 2:
+    return buildSVEShiftLeftImmediate(Op, 1, DAG, DL);
+  case 4:
+    return buildSVEShiftLeftImmediate(Op, 2, DAG, DL);
+  case 8:
+    return buildSVEShiftLeftImmediate(Op, 3, DAG, DL);
+  }
+}
+
+/// Combine scalable i32/i64 SVE multiply by small constants to ADD/SHL DAG
+/// nodes that avoid MUL and expose ADR selection opportunities.
+static SDValue performSVEMulByConstantNodeCombine(SDNode *N,
+                                                  SelectionDAG &DAG) {
+  EVT VT = N->getValueType(0);
+  bool IsSupportedSVEInt = isSupportedSVEMulByConstantVT(VT);
+  if (!IsSupportedSVEInt)
+    return SDValue();
+
+  SDValue Op = N->getOperand(0);
+  SDValue ImmOp = N->getOperand(1);
+  ConstantSDNode *ImmC = isConstOrConstSplat(ImmOp);
+  if (!ImmC) {
+    // MUL is commutative, so accept the splat on either side.
+    Op = N->getOperand(1);
+    ImmOp = N->getOperand(0);
+    ImmC = isConstOrConstSplat(ImmOp);
+    if (!ImmC)
+      return SDValue();
+  }
+
+  const APInt &Imm = ImmC->getAPIntValue();
+
+  // acc + (x * 2) -> acc + (x << 1), which can select ADR.
+  if (Imm == 2 && N->hasOneUse() && N->user_begin()->getOpcode() == ISD::ADD)
+    return buildSVEShiftLeftImmediate(Op, 1, DAG, SDLoc(N));
+
+  return performSVEMulByConstantCombine(Op, Imm.getZExtValue(), DAG, SDLoc(N));
+}
+
+/// Rewrite all-active SVE MUL intrinsics to the same ADD/SHL forms as plain
+/// arithmetic. Partial predicates are not equivalent because inactive lanes
+/// must keep their predicated intrinsic semantics.
+static SDValue performAllActiveSVEMulIntrinsicCombine(SDNode *N,
+                                                      SelectionDAG &DAG) {
+  if (!isSupportedSVEMulByConstantVT(N->getValueType(0)) ||
+      !isAllActivePredicate(DAG, N->getOperand(1)))
+    return SDValue();
+
+  SDValue Op = N->getOperand(2);
+  SDValue ImmOp = N->getOperand(3);
+  ConstantSDNode *ImmC = isConstOrConstSplat(ImmOp);
+  if (!ImmC) {
+    // MUL is commutative, so accept the splat on either side.
+    Op = N->getOperand(3);
+    ImmOp = N->getOperand(2);
+    ImmC = isConstOrConstSplat(ImmOp);
+    if (!ImmC)
+      return DAG.getNode(ISD::MUL, SDLoc(N), N->getValueType(0),
+                         N->getOperand(2), N->getOperand(3));
+  }
+
+  const APInt &Imm = ImmC->getAPIntValue();
+  if (SDValue Mul =
+          performSVEMulByConstantCombine(Op, Imm.getZExtValue(), DAG, SDLoc(N)))
+    return Mul;
+
+  return DAG.getNode(ISD::MUL, SDLoc(N), N->getValueType(0), N->getOperand(2),
+                     N->getOperand(3));
+}
+
+/// Rewrite all-active SVE MLA intrinsics as ADD plus a scaled multiply operand:
+///   acc + x * 1     -> acc + x
+///   acc + x * 2/4/8 -> acc + (x << 1/2/3)
----------------
tommat01 wrote:

Reimplemented as isel patterns

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