[AArch64][SVE] Avoid multiple PTRUE values for SETCC.

[sdesmalen]

Summary

Conversion between fixed- and scalable vector types may sometimes
lead to multiple PTRUE values to be used, one 'all active' and one
for the specific fixed-width vector length.

This patch canonicalises the predicate of SETCC_MERGE_ZERO to be
the more specific of the two, when it's inputs come from a conversion
of a fixed -> scalable type.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14106–14121

I thought we have an isel pattern to capture this. Is it the case you need to catch it early? If so then I think you should also remove the isel pattern to prove it's redundant.

17096

This also needs to test V.getValueType().isScalableVector() or have an equivalent assert if you never expect it to be called when such is not true.

17124

This looks misplaced?

17128–17129

This has me wondering if we're in fact missing something else. Do you know where the all active PTRUE is coming from?

You're making the assumption that there exists a fixed length PTRUE, which is probably true but still feels like an arbitrary assumption. For example, when vscale-min==vscale-max we favour all active PTRUEs.

[peterwaller-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

2386

Note: Block of similar forward decls here.

17092–17093

Nit. Is this the right place for this forward declaration? It's not used in the function immediately below. I see a block of forward decls above.

llvm/test/CodeGen/AArch64/sve-setcc.ll

119

Looks like a draft function name slipped in here.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 
 // Forward declarations of SVE fixed length lowering helpers
 static EVT getContainerForFixedLengthVector(SelectionDAG &DAG, EVT VT);
 static SDValue convertToScalableVector(SelectionDAG &DAG, EVT VT, SDValue V);
 static SDValue convertFromScalableVector(SelectionDAG &DAG, EVT VT, SDValue V);
 static SDValue convertFixedMaskToScalableVector(SDValue Mask,
                                                 SelectionDAG &DAG);
 static SDValue getPredicateForScalableVector(SelectionDAG &DAG, SDLoc &DL,
                                              EVT VT);
 
 /// isZerosVector - Check whether SDNode N is a zero-filled vector.
 static bool isZerosVector(const SDNode *N) {
   // Look through a bit convert.
   while (N->getOpcode() == ISD::BITCAST)
     N = N->getOperand(0).getNode();
 
   if (ISD::isConstantSplatVectorAllZeros(N))
Dup = DAG.getNode(AArch64ISD::DUP, DL,
                       DAG.getConstant(Mask.zextOrTrunc(32), DL, MVT::i32));
 
     SDValue And = DAG.getNode(ISD::AND, DL,
                               UnpkOp->getValueType(0), UnpkOp, Dup);
 
     return DAG.getNode(Opc, DL, N->getValueType(0), And);
   }
 
+  // Fold:
+  //    and(
+  //      setcc_merge_zero(pred, ...),
+  //      setcc_merge_zero(pred, ...))
+  // -> setcc_merge_zero(
+  //        setcc_merge_zero(pred, ...),
+  //        ...)
+  if (N->getOperand(0).getOpcode() == AArch64ISD::SETCC_MERGE_ZERO &&
+      N->getOperand(1).getOpcode() == AArch64ISD::SETCC_MERGE_ZERO) {
+    SDValue LHS = N->getOperand(0);
+    SDValue RHS = N->getOperand(1);
+    if (LHS.getOperand(0) == RHS.getOperand(0))
+      return DAG.getNode(AArch64ISD::SETCC_MERGE_ZERO, SDLoc(N),
+                         N->getValueType(0), LHS, RHS.getOperand(1),
+                         RHS.getOperand(2), RHS.getOperand(3));
+  }
+
   if (!EnableCombineMGatherIntrinsics)
     return SDValue();
 
   SDValue Mask = N->getOperand(1);
 
   if (!Src.hasOneUse())
     return SDValue();
 
if (Pred.getOpcode() == AArch64ISD::PTRUE &&
       Pred.getConstantOperandVal(0) == InnerPred.getConstantOperandVal(0) &&
       Pred->getConstantOperandVal(0) >= AArch64SVEPredPattern::vl1 &&
       Pred->getConstantOperandVal(0) <= AArch64SVEPredPattern::vl256)
     return Extract;
 
   return SDValue();
 }
 
+static SDValue getPredicateForFixedLengthVector(SelectionDAG &DAG, SDLoc &DL,
+                                                EVT VT);
+// Pattern match utility function to return if V is a conversion of a
+// fixed-width vector -> scalable vector.
+static bool isConvertToScalableVector(SDValue V) {
+  return V.getOpcode() == ISD::INSERT_SUBVECTOR && V.getOperand(0).isUndef() &&
+         V.getOperand(1).getValueType().isFixedLengthVector() &&
+         V.getConstantOperandVal(2) == 0;
+}
+
 static SDValue
 performSetccMergeZeroCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI) {
   assert(N->getOpcode() == AArch64ISD::SETCC_MERGE_ZERO &&
          "Unexpected opcode!");
 
   SelectionDAG &DAG = DCI.DAG;
   SDValue Pred = N->getOperand(0);
   SDValue LHS = N->getOperand(1);
   SDValue RHS = N->getOperand(2);
   ISD::CondCode Cond = cast<CondCodeSDNode>(N->getOperand(3))->get();
 
   if (SDValue V = performSetCCPunpkCombine(N, DAG))
     return V;
 
+  // Make the predicate more specific, because there is likely already a
+  // predicate available for the given VL.
+  //    setcc(cc, <ptrue all>, insert_subvector(undef, fixed_width_vec, 0), vec)
+  // -> setcc(cc, <ptrue VL>, insert_subvector(undef, fixed_width_vec, 0), vec)
+  //
+  // This avoids ending up with multiple PTRUEs and also helps some of the other
+  // folds for setcc_merge_zero which try to match multiple sequences of setcc's
+  // when the predicate is equivalent.
+  EVT FixedLengthVT;
+  if (isConvertToScalableVector(LHS) && isAllActivePredicate(DAG, Pred)) {
+    FixedLengthVT = LHS.getOperand(1).getValueType();
+    SDLoc DL(N);
+    SDValue FixedLengthPred =
+        getPredicateForFixedLengthVector(DAG, DL, FixedLengthVT);
+    return DAG.getNode(AArch64ISD::SETCC_MERGE_ZERO, DL, N->getValueType(0),
+                       FixedLengthPred, LHS, RHS, N->getOperand(3));
+  }
+
   if (Cond == ISD::SETNE && isZerosVector(RHS.getNode()) &&
       LHS->getOpcode() == ISD::SIGN_EXTEND &&
       LHS->getOperand(0)->getValueType(0) == N->getValueType(0)) {
     //    setcc_merge_zero(
     //       pred, extend(setcc_merge_zero(CC, pred, ...)), != splat(0))
     // => setcc_merge_zero(CC, pred, ...)
     if (LHS->getOperand(0)->getOpcode() == AArch64ISD::SETCC_MERGE_ZERO &&
         LHS->getOperand(0)->getOperand(0) == Pred)

llvm/test/CodeGen/AArch64/sve-setcc.ll

; CHECK-NEXT:    ret
   %cmp1 = call <vscale x 16 x i1> @llvm.aarch64.sve.cmpne.nxv16i8(<vscale x 16 x i1> %pg1, <vscale x 16 x i8> %vec, <vscale x 16 x i8> zeroinitializer)
   %cmp1.sext = sext <vscale x 16 x i1> %cmp1 to <vscale x 16 x i8>
   %cmp2 = call <vscale x 16 x i1> @llvm.aarch64.sve.cmpne.nxv16i8(<vscale x 16 x i1> %pg2, <vscale x 16 x i8> %cmp1.sext, <vscale x 16 x i8> zeroinitializer)
   ret <vscale x 16 x i1> %cmp2
 }
 
 declare <vscale x 16 x i1> @llvm.aarch64.sve.cmpne.nxv16i8(<vscale x 16 x i1>, <vscale x 16 x i8>, <vscale x 16 x i8>)
 
+; Make sure that only a single PTRUE instruction is generated (to avoid having both a `ptrue p0.s` and a `ptrue p0.s, vl16` for the different operations).
+define <vscale x 16 x i1> @foo(<vscale x 16 x i8> %other, <16 x i8> %subvec) {
+; CHECK-LABEL: foo:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.b, vl16
+; CHECK-NEXT:    // kill: def $q1 killed $q1 def $z1
+; CHECK-NEXT:    cmpne p1.b, p0/z, z0.b, #0
+; CHECK-NEXT:    cmpne p0.b, p0/z, z1.b, z0.b
+; CHECK-NEXT:    mov p0.b, p1/m, p1.b
+; CHECK-NEXT:    ret
+  %ptrue_vl16 = call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 9)
+  %cmp1 = call <vscale x 16 x i1> @llvm.aarch64.sve.cmpne.nxv16i8(<vscale x 16 x i1> %ptrue_vl16, <vscale x 16 x i8> %other, <vscale x 16 x i8> zeroinitializer)
+  %vec = call <vscale x 16 x i8> @llvm.experimental.vector.insert.nxv16i8.v16i8(<vscale x 16 x i8> poison, <16 x i8> %subvec, i64 0)
+  %cmp2 = icmp ne <vscale x 16 x i8> %vec, %other
+  %retval = or <vscale x 16 x i1> %cmp1, %cmp2
+  ret <vscale x 16 x i1> %retval
+}
+
+; Same as the above test, but make sure we still fold the AND operation into the SETCC.
+;
+; Fold:
+;    and(setcc_merge_zero(pred, ...),setcc_merge_zero(pred, ...))
+; -> setcc_merge_zero(setcc_merge_zero(pred, ...), ...)
+define <vscale x 16 x i1> @setcc_reuse_ptrue_vl16_fold_and(<vscale x 16 x i8> %other, <16 x i8> %subvec) {
+; CHECK-LABEL: setcc_reuse_ptrue_vl16_fold_and:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.b, vl16
+; CHECK-NEXT:    // kill: def $q1 killed $q1 def $z1
+; CHECK-NEXT:    cmpne p0.b, p0/z, z0.b, #0
+; CHECK-NEXT:    cmpne p0.b, p0/z, z1.b, z0.b
+; CHECK-NEXT:    ret
+  %ptrue_vl16 = call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 9)
+  %cmp1 = call <vscale x 16 x i1> @llvm.aarch64.sve.cmpne.nxv16i8(<vscale x 16 x i1> %ptrue_vl16, <vscale x 16 x i8> %other, <vscale x 16 x i8> zeroinitializer)
+  %vec = call <vscale x 16 x i8> @llvm.experimental.vector.insert.nxv16i8.v16i8(<vscale x 16 x i8> poison, <16 x i8> %subvec, i64 0)
+  %cmp2 = icmp ne <vscale x 16 x i8> %vec, %other
+  %retval = and <vscale x 16 x i1> %cmp1, %cmp2
+  ret <vscale x 16 x i1> %retval
+}
+
+declare <vscale x 16 x i8> @llvm.experimental.vector.insert.nxv16i8.v16i8(<vscale x 16 x i8>, <16 x i8>, i64)
+declare <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32)
+
 declare i1 @llvm.aarch64.sve.ptest.any.nxv8i1(<vscale x 8 x i1>, <vscale x 8 x i1>)
 declare i1 @llvm.aarch64.sve.ptest.last.nxv8i1(<vscale x 8 x i1>, <vscale x 8 x i1>)
 
 declare <vscale x 8 x i1> @llvm.aarch64.sve.cmplt.wide.nxv8i16(<vscale x 8 x i1>, <vscale x 8 x i16>, <vscale x 2 x i64>)
 
 declare void @llvm.masked.store.nxv8i16.p0nxv8i16(<vscale x 8 x i16>, <vscale x 8 x i16>*, i32, <vscale x 8 x i1>)