[SelectionDAG][AArch64] Restrict matchUnaryPredicate to only handle SPLAT_VECTOR for scalable vectors.

[craig.topper]

Summary

fde24661718c7812a20a10e518cd853e8e060107 added support for
scalable vectors to matchUnaryPredicate by handling SPLAT_VECTOR in
addition to BUILD_VECTOR. This was used to enabled UDIV/SDIV/UREM/SREM
by constant expansion in BuildUDIV/BuildSDIV in TargetLowering.cpp

The caller there expects to call getBuildVector from the match factors.
This leads to a crash right now if there is a SPLAT_VECTOR of
fixed vectors since the number of vectors won't match the number
of elements.

To fix this, this patch updates the callers to check the opcode
instead of whether the type is fixed or scalable. This assumes
that only 3 opcodes are handled by matchUnaryPredicate so
I've added an assertion to the final else to check that opcode.

[RKSimon]

I think for fixed length vectors we have a number of uses of the predicate matchers that assume it will be called once per vector element - some create additional vectors on the fly inside the matchers etc.

[david-arm]

Looking in lib/CodeGen/SelectionDAG/TargetLowering.cpp at the places where we call matchUnaryPredicate it looks like "VT.isFixedVector" and "VT.isScalableVector" calls are being used as proxies for BUILD_VECTOR and SPLAT_VECTOR. I suspect that checking for BUILD_VECTOR and SPLAT_VECTOR instead would give the right behaviour? Not suggesting that has to be fixed in this patch though...

[craig.topper]

In D96174#2548065, @david-arm wrote:

Looking in lib/CodeGen/SelectionDAG/TargetLowering.cpp at the places where we call matchUnaryPredicate it looks like "VT.isFixedVector" and "VT.isScalableVector" calls are being used as proxies for BUILD_VECTOR and SPLAT_VECTOR. I suspect that checking for BUILD_VECTOR and SPLAT_VECTOR instead would give the right behaviour? Not suggesting that has to be fixed in this patch though...

Yeah that would work, but it feels a little like leaking implementation details of matchUnaryPredicate into the caller.

[frasercrmck]

In D96174#2548071, @craig.topper wrote:

In D96174#2548065, @david-arm wrote:

Looking in lib/CodeGen/SelectionDAG/TargetLowering.cpp at the places where we call matchUnaryPredicate it looks like "VT.isFixedVector" and "VT.isScalableVector" calls are being used as proxies for BUILD_VECTOR and SPLAT_VECTOR. I suspect that checking for BUILD_VECTOR and SPLAT_VECTOR instead would give the right behaviour? Not suggesting that has to be fixed in this patch though...

Yeah that would work, but it feels a little like leaking implementation details of matchUnaryPredicate into the caller.

Agreed. I think it's better than checking the vector type, though. Maybe matchUnaryPredicate could return something to indicate the "root" it matched against?

[RKSimon]

In D96174#2548077, @frasercrmck wrote:

In D96174#2548071, @craig.topper wrote:

In D96174#2548065, @david-arm wrote:

Looking in lib/CodeGen/SelectionDAG/TargetLowering.cpp at the places where we call matchUnaryPredicate it looks like "VT.isFixedVector" and "VT.isScalableVector" calls are being used as proxies for BUILD_VECTOR and SPLAT_VECTOR. I suspect that checking for BUILD_VECTOR and SPLAT_VECTOR instead would give the right behaviour? Not suggesting that has to be fixed in this patch though...

Yeah that would work, but it feels a little like leaking implementation details of matchUnaryPredicate into the caller.

Agreed. I think it's better than checking the vector type, though. Maybe matchUnaryPredicate could return something to indicate the "root" it matched against?

Have you got a use case where you actually need to know this?

[craig.topper]

In D96174#2552182, @RKSimon wrote:

In D96174#2548077, @frasercrmck wrote:

In D96174#2548071, @craig.topper wrote:

In D96174#2548065, @david-arm wrote:

Looking in lib/CodeGen/SelectionDAG/TargetLowering.cpp at the places where we call matchUnaryPredicate it looks like "VT.isFixedVector" and "VT.isScalableVector" calls are being used as proxies for BUILD_VECTOR and SPLAT_VECTOR. I suspect that checking for BUILD_VECTOR and SPLAT_VECTOR instead would give the right behaviour? Not suggesting that has to be fixed in this patch though...

Yeah that would work, but it feels a little like leaking implementation details of matchUnaryPredicate into the caller.

Agreed. I think it's better than checking the vector type, though. Maybe matchUnaryPredicate could return something to indicate the "root" it matched against?

Have you got a use case where you actually need to know this?

Depends on what behavior we want for this interface and code like in BuildUDIV/BuildSDIV. Here are some options

1. SPLAT_VECTOR on fixed vectors causes matchUnaryPredicate to visit the splat value NumElements times. BuildUDIV/BuildSDIV pushes to vectors num elements times and calls getBuildVector. DAGCombine will probably create SPLAT_VECTOR from those build_vectors.
2. SPLAT_VECTOR on fixed vectors causes matchUnaryPredicate to visit the splat value 1 time. BuildUDIV/BuildSDDIV push to the vectors 1 times and call getSplatBuildVector. DAGCombine will probably create SPLAT_VECTOR from those build_vectors.
3. SPLAT_VECTOR on fixed vectors causes matchUnaryPredicate to visit the splat value 1 time. BuildUDIV/BuildSDDIV push to the vectors 1 times and call getSplatVector to create SPLAT_VECTOR.

I believe @david-arm is proposing option 3.

[RKSimon]

matchUnaryPredicate was added to try to encourage people to add full (non-uniform) vector support without too much legwork - a lot of combines still just have scalar / uniform support only. I'm just worried that if they have to do extra work to handle different vector cases it'll further discourage vector handling.

[craig.topper]

Maybe a different question we should ask is whether we should be allowing SPLAT_VECTOR with fixed vector types. BuildVectorSDNode already has a query for splat. If I recall correctly we've been adding special scalable vector only shuffle nodes that don't work with fixed vectors. Maybe SPLAT_VECTOR should be the same?

[frasercrmck]

In D96174#2552182, @RKSimon wrote:

Have you got a use case where you actually need to know this?

I don't have a use case, I was just going by the existing users of matchUnaryPredicate which already "need" to know some internal details of what was matched in a roundabout manner, since they're processing the results in different ways. If that makes sense?

[david-arm]

Hi, yeah I agree with @frasercrmck here. The only reason I brought this up originally is because callers of the function do something like this:

if (VT.isFixedLengthVector()) {
... Do something with BUILD_VECTOR ...
} else if (VT.isScalableVector()) {
.. Do something with SPLAT_VECTOR ...
}

It felt like since we're already doing this an alternative solution might be to check the opcode instead of the vector type. The patch you have right now does work and wouldn't want to hold up the patch if you don't have time to fix it right now. It was just a suggestion for another way of dealing with the problem.

[craig.topper]

Check the opcode instead of the type in the callers.

[RKSimon]

LGTM - cheers!

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

auto BuildSDIVPattern = [&](ConstantSDNode *C) {
     return true;
   };
 
   // Collect all magic values from the build vector.
   if (!ISD::matchUnaryPredicate(Op1, BuildSDIVPattern))
     return SDValue();
 
   SDValue Shift, Factor;
-  if (VT.isFixedLengthVector()) {
+  if (Op1.getOpcode() == ISD::BUILD_VECTOR) {
     Shift = DAG.getBuildVector(ShVT, dl, Shifts);
     Factor = DAG.getBuildVector(VT, dl, Factors);
-  } else if (VT.isScalableVector()) {
+  } else if (Op1.getOpcode() == ISD::SPLAT_VECTOR) {
     assert(Shifts.size() == 1 && Factors.size() == 1 &&
            "Expected matchUnaryPredicate to return one element for scalable "
            "vectors");
     Shift = DAG.getSplatVector(ShVT, dl, Shifts[0]);
     Factor = DAG.getSplatVector(VT, dl, Factors[0]);
   } else {
+    assert(isa<ConstantSDNode>(Op1) && "Expected a constant");
     Shift = Shifts[0];
     Factor = Factors[0];
   }
 
   SDValue Res = Op0;
 
   // Shift the value upfront if it is even, so the LSB is one.
   if (UseSRA) {
SDValue TargetLowering::BuildSDIV(SDNode *N, SelectionDAG &DAG,
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
 
   // Collect the shifts / magic values from each element.
   if (!ISD::matchUnaryPredicate(N1, BuildSDIVPattern))
     return SDValue();
 
   SDValue MagicFactor, Factor, Shift, ShiftMask;
-  if (VT.isFixedLengthVector()) {
+  if (N1.getOpcode() == ISD::BUILD_VECTOR) {
     MagicFactor = DAG.getBuildVector(VT, dl, MagicFactors);
     Factor = DAG.getBuildVector(VT, dl, Factors);
     Shift = DAG.getBuildVector(ShVT, dl, Shifts);
     ShiftMask = DAG.getBuildVector(VT, dl, ShiftMasks);
-  } else if (VT.isScalableVector()) {
+  } else if (N1.getOpcode() == ISD::SPLAT_VECTOR) {
     assert(MagicFactors.size() == 1 && Factors.size() == 1 &&
            Shifts.size() == 1 && ShiftMasks.size() == 1 &&
            "Expected matchUnaryPredicate to return one element for scalable "
            "vectors");
     MagicFactor = DAG.getSplatVector(VT, dl, MagicFactors[0]);
     Factor = DAG.getSplatVector(VT, dl, Factors[0]);
     Shift = DAG.getSplatVector(ShVT, dl, Shifts[0]);
     ShiftMask = DAG.getSplatVector(VT, dl, ShiftMasks[0]);
   } else {
+    assert(isa<ConstantSDNode>(N1) && "Expected a constant");
     MagicFactor = MagicFactors[0];
     Factor = Factors[0];
     Shift = Shifts[0];
     ShiftMask = ShiftMasks[0];
   }
 
   // Multiply the numerator (operand 0) by the magic value.
   // FIXME: We should support doing a MUL in a wider type.
SDValue TargetLowering::BuildUDIV(SDNode *N, SelectionDAG &DAG,
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
 
   // Collect the shifts/magic values from each element.
   if (!ISD::matchUnaryPredicate(N1, BuildUDIVPattern))
     return SDValue();
 
   SDValue PreShift, PostShift, MagicFactor, NPQFactor;
-  if (VT.isFixedLengthVector()) {
+  if (N1.getOpcode() == ISD::BUILD_VECTOR) {
     PreShift = DAG.getBuildVector(ShVT, dl, PreShifts);
     MagicFactor = DAG.getBuildVector(VT, dl, MagicFactors);
     NPQFactor = DAG.getBuildVector(VT, dl, NPQFactors);
     PostShift = DAG.getBuildVector(ShVT, dl, PostShifts);
-  } else if (VT.isScalableVector()) {
+  } else if (N1.getOpcode() == ISD::SPLAT_VECTOR) {
     assert(PreShifts.size() == 1 && MagicFactors.size() == 1 &&
            NPQFactors.size() == 1 && PostShifts.size() == 1 &&
            "Expected matchUnaryPredicate to return one for scalable vectors");
     PreShift = DAG.getSplatVector(ShVT, dl, PreShifts[0]);
     MagicFactor = DAG.getSplatVector(VT, dl, MagicFactors[0]);
     NPQFactor = DAG.getSplatVector(VT, dl, NPQFactors[0]);
     PostShift = DAG.getSplatVector(ShVT, dl, PostShifts[0]);
   } else {
+    assert(isa<ConstantSDNode>(N1) && "Expected a constant");
     PreShift = PreShifts[0];
     MagicFactor = MagicFactors[0];
     PostShift = PostShifts[0];
   }
 
   SDValue Q = N0;
   Q = DAG.getNode(ISD::SRL, dl, VT, Q, PreShift);
   Created.push_back(Q.getNode());
if (AllDivisorsAreOnes)
     return SDValue();
 
   // If this is a srem by a powers-of-two (including INT_MIN), avoid the fold
   // since it can be best implemented as a bit test.
   if (AllDivisorsArePowerOfTwo)
     return SDValue();
 
   SDValue PVal, AVal, KVal, QVal;
-  if (VT.isFixedLengthVector()) {
+  if (D.getOpcode() == ISD::BUILD_VECTOR) {
     if (HadOneDivisor) {
       // Try to turn PAmts into a splat, since we don't care about the values
       // that are currently '0'. If we can't, just keep '0'`s.
       turnVectorIntoSplatVector(PAmts, isNullConstant);
       // Try to turn AAmts into a splat, since we don't care about the
       // values that are currently '-1'. If we can't, change them to '0'`s.
       turnVectorIntoSplatVector(AAmts, isAllOnesConstant,
                                 DAG.getConstant(0, DL, SVT));
       // Try to turn KAmts into a splat, since we don't care about the values
       // that are currently '-1'. If we can't, change them to '0'`s.
       turnVectorIntoSplatVector(KAmts, isAllOnesConstant,
                                 DAG.getConstant(0, DL, ShSVT));
     }
 
     PVal = DAG.getBuildVector(VT, DL, PAmts);
     AVal = DAG.getBuildVector(VT, DL, AAmts);
     KVal = DAG.getBuildVector(ShVT, DL, KAmts);
     QVal = DAG.getBuildVector(VT, DL, QAmts);
-  } else if (VT.isScalableVector()) {
+  } else if (D.getOpcode() == ISD::SPLAT_VECTOR) {
     assert(PAmts.size() == 1 && AAmts.size() == 1 && KAmts.size() == 1 &&
            QAmts.size() == 1 &&
            "Expected matchUnaryPredicate to return one element for scalable "
            "vectors");
     PVal = DAG.getSplatVector(VT, DL, PAmts[0]);
     AVal = DAG.getSplatVector(VT, DL, AAmts[0]);
     KVal = DAG.getSplatVector(ShVT, DL, KAmts[0]);
     QVal = DAG.getSplatVector(VT, DL, QAmts[0]);
   } else {
+    assert(isa<ConstantSDNode>(D) && "Expected a constant");
     PVal = PAmts[0];
     AVal = AAmts[0];
     KVal = KAmts[0];
     QVal = QAmts[0];
   }
 
   // (mul N, P)
   SDValue Op0 = DAG.getNode(ISD::MUL, DL, VT, N, PVal);

llvm/test/CodeGen/AArch64/sve-fixed-length-int-div.ll

 ; VBITS_GE_2048-NEXT: ret
   %op1 = load <32 x i64>, <32 x i64>* %a
   %op2 = load <32 x i64>, <32 x i64>* %b
   %res = udiv <32 x i64> %op1, %op2
   store <32 x i64> %res, <32 x i64>* %a
   ret void
 }
 
+; This used to crash because isUnaryPredicate and BuildUDIV don't know how
+; a SPLAT_VECTOR of fixed vector type should be handled.
+define void @udiv_constantsplat_v8i32(<8 x i32>* %a) #0 {
+; CHECK-LABEL: udiv_constantsplat_v8i32:
+; CHECK: ptrue [[PG:p[0-9]+]].s, vl[[#min(div(VBYTES,4),8)]]
+; CHECK-NEXT: ld1w { [[OP1:z[0-9]+]].s }, [[PG]]/z, [x0]
+; CHECK-NEXT: mov [[OP2:z[0-9]+]].s, #95
+; CHECK-NEXT: udiv [[RES:z[0-9]+]].s, [[PG]]/m, [[OP1]].s, [[OP2]].s
+; CHECK-NEXT: st1w { [[RES]].s }, [[PG]], [x0]
+; CHECK-NEXT: ret
+  %op1 = load <8 x i32>, <8 x i32>* %a
+  %res = udiv <8 x i32> %op1, <i32 95, i32 95, i32 95, i32 95, i32 95, i32 95, i32 95, i32 95>
+  store <8 x i32> %res, <8 x i32>* %a
+  ret void
+}
+
 attributes #0 = { "target-features"="+sve" }