[SelectionDAG] Teach SelectionDAG::FoldConstantArithmetic to handle SPLAT_VECTOR

[craig.topper]

Summary

This allows FoldConstantArithmetic to handle SPLAT_VECTOR in
addition to BUILD_VECTOR. This allows it to support scalable
vectors. I'm also allowing fixed length SPLAT_VECTOR which is
used by some targets, but I'm not familiar enough to write tests
for those targets.

I had to block this function from running on CONCAT_VECTORS to
avoid calling getNode for a CONCAT_VECTORS of 2 scalars.
This can happen because the 2 operand getNode calls this
function for any opcode. Previously we were protected because
CONCAT_VECTORs of BUILD_VECTOR is folded to a larger BUILD_VECTOR
before that call. But it's not always possible to fold a CONCAT_VECTORS
of SPLAT_VECTORs, and we don't even try.

This fixes PR49781 where DAG combine thought constant folding
should be possible, but FoldConstantArithmetic couldn't do it.

[nagisa]

Please rebase on top of master to verify this fixes the issue mentioned in D88785.

[frasercrmck]

All looks reasonable to me. I'm not sure whether we should do one iteration for fixed-length SPLAT_VECTORs: it should be sound. I don't know what it would gain us given how infrequently it's used.

I'm also not sure whether there's another way of handling the CONCAT_VECTORs issue so hopefully someone else would be able to provide some insight there.

[david-arm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5129

I wonder if it makes more sense to only set NumOps to VT.getVectorNumElements() if at least one opcode is BUILD_VECTOR. Then you can have an assert that for scalable vectors neither of the opcodes is BUILD_VECTOR.

5177

This is just a suggestion, but instead of checking for scalable vectors, if you change NumOps above to be always be 1 if no BUILD_VECTORs are involved, then maybe you could write:

if (Outputs.size() == 1) {
  // Must be a combination of either SPLAT_VECTOR or UNDEF.
  return getSplatVector(VT, SDLoc(), Outputs[0]);
}

However, I'm not sure about the case when both inputs are UNDEF and maybe still want to return UNDEF?

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5177

Both inputs being undef doesn't go through this code, it's handled earlier. One of the operands must be a build_vector or splat_vector.

I think using Outputs.size() == 1 would be incorrect if the fixed vector type has 1 element and the input is build_vector. We shouldn't generate splat_vector in that case.

[david-arm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5177

Ah ok, fair enough. I was just thinking that it might be nice to avoid tying scalable vectors to SPLAT_VECTOR that's all, given that your changes above should work for fixed-length vectors too. For example, if we've already gone to the trouble of creating inputs for fixed-length vectors, then it might make sense to generate one here too?

[craig.topper]

Rebase

Will work on generating splat vectors for fixed vectors when possible next.

[craig.topper]

Create SPLAT_VECTOR result for SPLAT_VECTOR inputs.

[david-arm]

LGTM! Thanks a lot for making the changes.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5183

nit: Perhaps this should be something like:

If one of the inputs is a fixed length SPLAT_VECTOR ...

since at least one input is BUILD_VECTOR?

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5183

I think I just forgot to delete that TODO when I refactored.

[craig.topper]

Remove stale TODO

[craig.topper]

@david-arm, you wrote LGTM, but didn't accept the revision. Do I need someone else to approve?

[david-arm]

Ah, sorry, that was just a simple mistake and forget to "Accept Revision"!

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

bool SelectionDAG::isUndef(unsigned Opcode, ArrayRef<SDValue> Ops) {
   }
 }
 
 SDValue SelectionDAG::FoldConstantArithmetic(unsigned Opcode, const SDLoc &DL,
                                              EVT VT, ArrayRef<SDValue> Ops) {
   // If the opcode is a target-specific ISD node, there's nothing we can
   // do here and the operand rules may not line up with the below, so
   // bail early.
-  if (Opcode >= ISD::BUILTIN_OP_END)
+  // We can't create a scalar CONCAT_VECTORS so skip it. It will break
+  // for concats involving SPLAT_VECTOR. Concats of BUILD_VECTORS are handled by
+  // foldCONCAT_VECTORS in getNode before this is called.
+  if (Opcode >= ISD::BUILTIN_OP_END || Opcode == ISD::CONCAT_VECTORS)
     return SDValue();
 
   // For now, the array Ops should only contain two values.
   // This enforcement will be removed once this function is merged with
   // FoldConstantVectorArithmetic
   if (Ops.size() != 2)
     return SDValue();
 
SDValue SelectionDAG::FoldConstantArithmetic(unsigned Opcode, const SDLoc &DL,
 
   // fold (add Sym, c) -> Sym+c
   if (GlobalAddressSDNode *GA = dyn_cast<GlobalAddressSDNode>(N1))
     return FoldSymbolOffset(Opcode, VT, GA, N2);
   if (TLI->isCommutativeBinOp(Opcode))
     if (GlobalAddressSDNode *GA = dyn_cast<GlobalAddressSDNode>(N2))
       return FoldSymbolOffset(Opcode, VT, GA, N1);
 
-  // TODO: All the folds below are performed lane-by-lane and assume a fixed
-  // vector width, however we should be able to do constant folds involving
-  // splat vector nodes too.
-  if (VT.isScalableVector())
-    return SDValue();
-
   // For fixed width vectors, extract each constant element and fold them
   // individually. Either input may be an undef value.
-  auto *BV1 = dyn_cast<BuildVectorSDNode>(N1);
-  if (!BV1 && !N1->isUndef())
-    return SDValue();
-  auto *BV2 = dyn_cast<BuildVectorSDNode>(N2);
-  if (!BV2 && !N2->isUndef())
+  bool IsBVOrSV1 = N1->getOpcode() == ISD::BUILD_VECTOR ||
+                   N1->getOpcode() == ISD::SPLAT_VECTOR;
+  if (!IsBVOrSV1 && !N1->isUndef())
+    return SDValue();
+  bool IsBVOrSV2 = N2->getOpcode() == ISD::BUILD_VECTOR ||
+                   N2->getOpcode() == ISD::SPLAT_VECTOR;
+  if (!IsBVOrSV2 && !N2->isUndef())
     return SDValue();
   // If both operands are undef, that's handled the same way as scalars.
-  if (!BV1 && !BV2)
+  if (!IsBVOrSV1 && !IsBVOrSV2)
     return SDValue();
 
-  assert((!BV1 || !BV2 || BV1->getNumOperands() == BV2->getNumOperands()) &&
-         "Vector binop with different number of elements in operands?");
-
   EVT SVT = VT.getScalarType();
   EVT LegalSVT = SVT;
   if (NewNodesMustHaveLegalTypes && LegalSVT.isInteger()) {
     LegalSVT = TLI->getTypeToTransformTo(*getContext(), LegalSVT);
     if (LegalSVT.bitsLT(SVT))
       return SDValue();
   }
+
   SmallVector<SDValue, 4> Outputs;
-  unsigned NumOps = BV1 ? BV1->getNumOperands() : BV2->getNumOperands();
+  unsigned NumOps = 0;
+  if (IsBVOrSV1)
+    NumOps = std::max(NumOps, N1->getNumOperands());
+  if (IsBVOrSV2)
+    NumOps = std::max(NumOps, N2->getNumOperands());
+  assert(NumOps != 0 && "Expected non-zero operands");
+  // Scalable vectors should only be SPLAT_VECTOR or UNDEF here. We only need
+  // one iteration for that.
+  assert((!VT.isScalableVector() || NumOps == 1) &&
+         "Scalar vector should only have one scalar");
+
   for (unsigned I = 0; I != NumOps; ++I) {
-    SDValue V1 = BV1 ? BV1->getOperand(I) : getUNDEF(SVT);
-    SDValue V2 = BV2 ? BV2->getOperand(I) : getUNDEF(SVT);
+    // We can have a fixed length SPLAT_VECTOR and a BUILD_VECTOR so we need
+    // to use operand 0 of the SPLAT_VECTOR for each fixed element.
+    SDValue V1;
+    if (N1->getOpcode() == ISD::BUILD_VECTOR)
+      V1 = N1->getOperand(I);
+    else if (N1->getOpcode() == ISD::SPLAT_VECTOR)
+      V1 = N1->getOperand(0);
+    else
+      V1 = getUNDEF(SVT);
+
+    SDValue V2;
+    if (N2->getOpcode() == ISD::BUILD_VECTOR)
+      V2 = N2->getOperand(I);
+    else if (N2->getOpcode() == ISD::SPLAT_VECTOR)
+      V2 = N2->getOperand(0);
+    else
+      V2 = getUNDEF(SVT);
+
     if (SVT.isInteger()) {
-      if (V1->getValueType(0).bitsGT(SVT))
+      if (V1.getValueType().bitsGT(SVT))
         V1 = getNode(ISD::TRUNCATE, DL, SVT, V1);
-      if (V2->getValueType(0).bitsGT(SVT))
+      if (V2.getValueType().bitsGT(SVT))
         V2 = getNode(ISD::TRUNCATE, DL, SVT, V2);
     }
 
-    if (V1->getValueType(0) != SVT || V2->getValueType(0) != SVT)
+    if (V1.getValueType() != SVT || V2.getValueType() != SVT)
       return SDValue();
 
     // Fold one vector element.
     SDValue ScalarResult = getNode(Opcode, DL, SVT, V1, V2);
     if (LegalSVT != SVT)
       ScalarResult = getNode(ISD::SIGN_EXTEND, DL, LegalSVT, ScalarResult);
 
     // Scalar folding only succeeded if the result is a constant or UNDEF.
     if (!ScalarResult.isUndef() && ScalarResult.getOpcode() != ISD::Constant &&
         ScalarResult.getOpcode() != ISD::ConstantFP)
       return SDValue();
     Outputs.push_back(ScalarResult);
   }
 
+  if (N1->getOpcode() == ISD::BUILD_VECTOR ||
+      N2->getOpcode() == ISD::BUILD_VECTOR) {
     assert(VT.getVectorNumElements() == Outputs.size() &&
            "Vector size mismatch!");
 
     // Build a big vector out of the scalar elements we generated.
     return getBuildVector(VT, SDLoc(), Outputs);
   }
 
+  assert((N1->getOpcode() == ISD::SPLAT_VECTOR ||
+          N2->getOpcode() == ISD::SPLAT_VECTOR) &&
+         "One operand should be a splat vector");
+
+  assert(Outputs.size() == 1 && "Vector size mismatch!");
+  return getSplatVector(VT, SDLoc(), Outputs[0]);
+}
+
 // TODO: Merge with FoldConstantArithmetic
 SDValue SelectionDAG::FoldConstantVectorArithmetic(unsigned Opcode,
                                                    const SDLoc &DL, EVT VT,
                                                    ArrayRef<SDValue> Ops,
                                                    const SDNodeFlags Flags) {
   // If the opcode is a target-specific ISD node, there's nothing we can
   // do here and the operand rules may not line up with the below, so
   // bail early.

llvm/test/CodeGen/AArch64/pr49781.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=aarch64 -mattr=+sve | FileCheck %s
+
+define <vscale x 2 x i64> @foo(<vscale x 2 x i64> %a) {
+; CHECK-LABEL: foo:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    sub z0.d, z0.d, #2 // =0x2
+; CHECK-NEXT:    ret
+ %idx = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 1, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
+ %b = sub <vscale x 2 x i64> %a, %idx
+ %c = sub <vscale x 2 x i64> %b, %idx
+ ret <vscale x 2 x i64> %c
+}

llvm/test/CodeGen/RISCV/rvv/vadd-sdnode-rv32.ll

 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 1 x i8> undef, i8 2, i32 0
   %splat = shufflevector <vscale x 1 x i8> %head, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
   %vc = add <vscale x 1 x i8> %va, %splat
   ret <vscale x 1 x i8> %vc
 }
 
 ; Test constant adds to see if we can optimize them away for scalable vectors.
-; FIXME: We can't.
 define <vscale x 1 x i8> @vadd_ii_nxv1i8_1() {
 ; CHECK-LABEL: vadd_ii_nxv1i8_1:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e8,mf8,ta,mu
-; CHECK-NEXT:    vmv.v.i v25, 2
-; CHECK-NEXT:    vadd.vi v8, v25, 3
+; CHECK-NEXT:    vmv.v.i v8, 5
 ; CHECK-NEXT:    ret
   %heada = insertelement <vscale x 1 x i8> undef, i8 2, i32 0
   %splata = shufflevector <vscale x 1 x i8> %heada, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
   %headb = insertelement <vscale x 1 x i8> undef, i8 3, i32 0
   %splatb = shufflevector <vscale x 1 x i8> %headb, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
   %vc = add <vscale x 1 x i8> %splata, %splatb
   ret <vscale x 1 x i8> %vc
 }

llvm/test/CodeGen/RISCV/rvv/vadd-sdnode-rv64.ll

 ; CHECK-NEXT:    vadd.vi v8, v8, 2
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 1 x i8> undef, i8 2, i32 0
   %splat = shufflevector <vscale x 1 x i8> %head, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
   %vc = add <vscale x 1 x i8> %va, %splat
   ret <vscale x 1 x i8> %vc
 }
 
+; Test constant adds to see if we can optimize them away for scalable vectors.
+define <vscale x 1 x i8> @vadd_ii_nxv1i8_1() {
+; CHECK-LABEL: vadd_ii_nxv1i8_1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a0, zero, e8,mf8,ta,mu
+; CHECK-NEXT:    vmv.v.i v8, 5
+; CHECK-NEXT:    ret
+  %heada = insertelement <vscale x 1 x i8> undef, i8 2, i32 0
+  %splata = shufflevector <vscale x 1 x i8> %heada, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
+  %headb = insertelement <vscale x 1 x i8> undef, i8 3, i32 0
+  %splatb = shufflevector <vscale x 1 x i8> %headb, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
+  %vc = add <vscale x 1 x i8> %splata, %splatb
+  ret <vscale x 1 x i8> %vc
+}
+
 define <vscale x 2 x i8> @vadd_vx_nxv2i8(<vscale x 2 x i8> %va, i8 signext %b) {
 ; CHECK-LABEL: vadd_vx_nxv2i8:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a1, zero, e8,mf4,ta,mu
 ; CHECK-NEXT:    vadd.vx v8, v8, a0
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 2 x i8> undef, i8 %b, i32 0
   %splat = shufflevector <vscale x 2 x i8> %head, <vscale x 2 x i8> undef, <vscale x 2 x i32> zeroinitializer

llvm/test/CodeGen/RISCV/rvv/vmulh-sdnode-rv32.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv32 -mattr=+experimental-v -verify-machineinstrs < %s | FileCheck %s
 
 ; Test that the prepareSREMEqFold optimization doesn't crash on scalable
 ; vector types.
 define <vscale x 4 x i1> @srem_eq_fold_nxv4i8(<vscale x 4 x i8> %va) {
 ; CHECK-LABEL: srem_eq_fold_nxv4i8:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    addi a0, zero, -85
 ; CHECK-NEXT:    vsetvli a1, zero, e8,mf2,ta,mu
 ; CHECK-NEXT:    vmul.vx v25, v8, a0
 ; CHECK-NEXT:    addi a0, zero, 42
 ; CHECK-NEXT:    vadd.vx v25, v25, a0
-; CHECK-NEXT:    vmv.v.i v26, 1
-; CHECK-NEXT:    vrsub.vi v27, v26, 0
-; CHECK-NEXT:    vand.vi v27, v27, 7
-; CHECK-NEXT:    vsll.vv v27, v25, v27
-; CHECK-NEXT:    vand.vi v26, v26, 7
-; CHECK-NEXT:    vsrl.vv v25, v25, v26
-; CHECK-NEXT:    vor.vv v25, v25, v27
+; CHECK-NEXT:    vsll.vi v26, v25, 7
+; CHECK-NEXT:    vsrl.vi v25, v25, 1
+; CHECK-NEXT:    vor.vv v25, v25, v26
 ; CHECK-NEXT:    vmsleu.vx v0, v25, a0
 ; CHECK-NEXT:    ret
   %head_six = insertelement <vscale x 4 x i8> undef, i8 6, i32 0
   %splat_six = shufflevector <vscale x 4 x i8> %head_six, <vscale x 4 x i8> undef, <vscale x 4 x i32> zeroinitializer
   %rem = srem <vscale x 4 x i8> %va, %splat_six
 
   %cc = icmp eq <vscale x 4 x i8> %rem, zeroinitializer
   ret <vscale x 4 x i1> %cc
 }

llvm/test/CodeGen/RISCV/rvv/vsub-sdnode-rv32.ll

 ; CHECK-NEXT:    vsub.vx v8, v8, a0
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 1 x i8> undef, i8 1, i32 0
   %splat = shufflevector <vscale x 1 x i8> %head, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
   %vc = sub <vscale x 1 x i8> %va, %splat
   ret <vscale x 1 x i8> %vc
 }
 
+; Test constant subs to see if we can optimize them away for scalable vectors.
+define <vscale x 1 x i8> @vsub_ii_nxv1i8_1() {
+; CHECK-LABEL: vsub_ii_nxv1i8_1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a0, zero, e8,mf8,ta,mu
+; CHECK-NEXT:    vmv.v.i v8, -1
+; CHECK-NEXT:    ret
+  %heada = insertelement <vscale x 1 x i8> undef, i8 2, i32 0
+  %splata = shufflevector <vscale x 1 x i8> %heada, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
+  %headb = insertelement <vscale x 1 x i8> undef, i8 3, i32 0
+  %splatb = shufflevector <vscale x 1 x i8> %headb, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
+  %vc = sub <vscale x 1 x i8> %splata, %splatb
+  ret <vscale x 1 x i8> %vc
+}
+
 define <vscale x 2 x i8> @vsub_vv_nxv2i8(<vscale x 2 x i8> %va, <vscale x 2 x i8> %vb) {
 ; CHECK-LABEL: vsub_vv_nxv2i8:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e8,mf4,ta,mu
 ; CHECK-NEXT:    vsub.vv v8, v8, v9
 ; CHECK-NEXT:    ret
   %vc = sub <vscale x 2 x i8> %va, %vb
   ret <vscale x 2 x i8> %vc

llvm/test/CodeGen/RISCV/rvv/vsub-sdnode-rv64.ll

 ; CHECK-NEXT:    vsub.vx v8, v8, a0
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 1 x i8> undef, i8 1, i32 0
   %splat = shufflevector <vscale x 1 x i8> %head, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
   %vc = sub <vscale x 1 x i8> %va, %splat
   ret <vscale x 1 x i8> %vc
 }
 
+; Test constant subs to see if we can optimize them away for scalable vectors.
+define <vscale x 1 x i8> @vsub_ii_nxv1i8_1() {
+; CHECK-LABEL: vsub_ii_nxv1i8_1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a0, zero, e8,mf8,ta,mu
+; CHECK-NEXT:    vmv.v.i v8, -1
+; CHECK-NEXT:    ret
+  %heada = insertelement <vscale x 1 x i8> undef, i8 2, i32 0
+  %splata = shufflevector <vscale x 1 x i8> %heada, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
+  %headb = insertelement <vscale x 1 x i8> undef, i8 3, i32 0
+  %splatb = shufflevector <vscale x 1 x i8> %headb, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
+  %vc = sub <vscale x 1 x i8> %splata, %splatb
+  ret <vscale x 1 x i8> %vc
+}
+
 define <vscale x 2 x i8> @vsub_vv_nxv2i8(<vscale x 2 x i8> %va, <vscale x 2 x i8> %vb) {
 ; CHECK-LABEL: vsub_vv_nxv2i8:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e8,mf4,ta,mu
 ; CHECK-NEXT:    vsub.vv v8, v8, v9
 ; CHECK-NEXT:    ret
   %vc = sub <vscale x 2 x i8> %va, %vb
   ret <vscale x 2 x i8> %vc