[AArch64][SVE] Add support for trunc to <vscale x N x i1>.

[efriedma]

Summary

This isn't a natively supported operation, so convert it to a mask+compare.

In addition to the operation itself, fix up some surrounding stuff to make the testcase work: we need concat_vectors on i1 vectors, we need legalization of i1 vector truncates, and we need to fix up all the relevant uses of getVectorNumElements().

[sdesmalen]

LGTM, thanks!

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

69

It seems our downstream compiler uses lsl with #(elementwidth-1) instead of and #0x1, but using and seems a bit more intuitive.

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19070

Could use getVectorMinNumElements() here.

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

2567

Since this is now an ElementCount this could be

assert(ElementCount / 2 && ....

which already asserts the division is lossless.

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

77–144

Based on the assumption that these tests are really testing the CONCAT_VECTORS functionality, is it worth naming them as such (perhaps also move into sve-shuffles.ll or something)?

[efriedma]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

2567

We do that division a bit later in the code.

I could drop the assertion, I guess.

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

77–144

Well, it's testing the truncate itself, and the legalization of scalable truncates, and yes, also the concat_vectors functionality. I can rename them, sure; maybe trunc_i64toi1_split_op_and_concatenate_result. I think the location is fine, since the way it's expressed in IR is just a plain truncate.

[paulwalker-arm]

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

77–144

I guess I was just wondering if you would have still written these tests if there was a route to test CONCAT_VECTORS in isolation (of which I understand there currently isn't). If you're happy with the location then I'm happy with their current names.

[efriedma]

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

77–144

We need these tests to test the legalization aspect, which can't be tested any other way.

[paulwalker-arm]

LGTM as long as we can replace the two instances of ElementCount.Min as suggested.

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

77–144

Yep, understood.

[efriedma]

b8f765a

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

for (unsigned i = 0, e = N0.getNumOperands(); i != e; ++i) {
       if (!X.isUndef()) {
         V = X;
         Idx = i;
         NumDefs++;
       }
       // Stop if more than one members are non-undef.
       if (NumDefs > 1)
         break;
+
       VTs.push_back(EVT::getVectorVT(*DAG.getContext(),
                                      VT.getVectorElementType(),
-                                     X.getValueType().getVectorNumElements()));
+                                     X.getValueType().getVectorElementCount()));
     }
 
     if (NumDefs == 0)
       return DAG.getUNDEF(VT);
 
     if (NumDefs == 1) {
       assert(V.getNode() && "The single defined operand is empty!");
       SmallVector<SDValue, 8> Opnds;

 
 // Check to see if this is a CONCAT_VECTORS of a bunch of EXTRACT_SUBVECTOR
 // operations. If so, and if the EXTRACT_SUBVECTOR vector inputs come from at
 // most two distinct vectors the same size as the result, attempt to turn this
 // into a legal shuffle.
 static SDValue combineConcatVectorOfExtracts(SDNode *N, SelectionDAG &DAG) {
   EVT VT = N->getValueType(0);
   EVT OpVT = N->getOperand(0).getValueType();
+
+  // We currently can't generate an appropriate shuffle for a scalable vector.
+  if (VT.isScalableVector())
+    return SDValue();
+
   int NumElts = VT.getVectorNumElements();
   int NumOpElts = OpVT.getVectorNumElements();
 
   SDValue SV0 = DAG.getUNDEF(VT), SV1 = DAG.getUNDEF(VT);
   SmallVector<int, 8> Mask;
 
   for (SDValue Op : N->ops()) {
     Op = peekThroughBitcasts(Op);
SDValue DAGCombiner::visitCONCAT_VECTORS(SDNode *N) {
   if (Level < AfterLegalizeVectorOps && TLI.isTypeLegal(VT))
     if (SDValue V = combineConcatVectorOfExtracts(N, DAG))
       return V;
 
   if (SDValue V = combineConcatVectorOfCasts(N, DAG))
     return V;
 
   // Type legalization of vectors and DAG canonicalization of SHUFFLE_VECTOR
-  // nodes often generate nop CONCAT_VECTOR nodes.
-  // Scan the CONCAT_VECTOR operands and look for a CONCAT operations that
-  // place the incoming vectors at the exact same location.
+  // nodes often generate nop CONCAT_VECTOR nodes. Scan the CONCAT_VECTOR
+  // operands and look for a CONCAT operations that place the incoming vectors
+  // at the exact same location.
+  //
+  // For scalable vectors, EXTRACT_SUBVECTOR indexes are implicitly scaled.
   SDValue SingleSource = SDValue();
-  unsigned PartNumElem = N->getOperand(0).getValueType().getVectorNumElements();
+  unsigned PartNumElem =
+      N->getOperand(0).getValueType().getVectorElementCount().Min;
 
   for (unsigned i = 0, e = N->getNumOperands(); i != e; ++i) {
     SDValue Op = N->getOperand(i);
 
     if (Op.isUndef())
       continue;
 
     // Check if this is the identity extract:

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

SDValue DAGTypeLegalizer::SplitVecOp_UnaryOp(SDNode *N) {
   // The result has a legal vector type, but the input needs splitting.
   EVT ResVT = N->getValueType(0);
   SDValue Lo, Hi;
   SDLoc dl(N);
   GetSplitVector(N->getOperand(N->isStrictFPOpcode() ? 1 : 0), Lo, Hi);
   EVT InVT = Lo.getValueType();
 
   EVT OutVT = EVT::getVectorVT(*DAG.getContext(), ResVT.getVectorElementType(),
-                               InVT.getVectorNumElements());
+                               InVT.getVectorElementCount());
 
   if (N->isStrictFPOpcode()) {
     Lo = DAG.getNode(N->getOpcode(), dl, { OutVT, MVT::Other }, 
                      { N->getOperand(0), Lo });
     Hi = DAG.getNode(N->getOpcode(), dl, { OutVT, MVT::Other }, 
                      { N->getOperand(0), Hi });
 
     // Build a factor node to remember that this operation is independent
SDValue DAGTypeLegalizer::SplitVecOp_TruncateHelper(SDNode *N) {
   //   %res = v8i8 trunc v8i16 %in16
   //
   // Without this transform, the original truncate would end up being
   // scalarized, which is pretty much always a last resort.
   unsigned OpNo = N->isStrictFPOpcode() ? 1 : 0;
   SDValue InVec = N->getOperand(OpNo);
   EVT InVT = InVec->getValueType(0);
   EVT OutVT = N->getValueType(0);
-  unsigned NumElements = OutVT.getVectorNumElements();
+  ElementCount NumElements = OutVT.getVectorElementCount();
   bool IsFloat = OutVT.isFloatingPoint();
 
   // Widening should have already made sure this is a power-two vector
   // if we're trying to split it at all. assert() that's true, just in case.
-  assert(!(NumElements & 1) && "Splitting vector, but not in half!");
+  assert(!(NumElements.Min & 1) && "Splitting vector, but not in half!");
 
   unsigned InElementSize = InVT.getScalarSizeInBits();
   unsigned OutElementSize = OutVT.getScalarSizeInBits();
 
   // Determine the split output VT. If its legal we can just split dirctly.
   EVT LoOutVT, HiOutVT;
   std::tie(LoOutVT, HiOutVT) = DAG.GetSplitDestVTs(OutVT);
   assert(LoOutVT == HiOutVT && "Unequal split?");

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

for (MVT VT : MVT::integer_scalable_vector_valuetypes()) {
         setOperationAction(ISD::UDIV, VT, Custom);
         setOperationAction(ISD::SMIN, VT, Custom);
         setOperationAction(ISD::UMIN, VT, Custom);
         setOperationAction(ISD::SMAX, VT, Custom);
         setOperationAction(ISD::UMAX, VT, Custom);
         setOperationAction(ISD::SHL, VT, Custom);
         setOperationAction(ISD::SRL, VT, Custom);
         setOperationAction(ISD::SRA, VT, Custom);
-        if (VT.getScalarType() == MVT::i1)
+        if (VT.getScalarType() == MVT::i1) {
           setOperationAction(ISD::SETCC, VT, Custom);
+          setOperationAction(ISD::TRUNCATE, VT, Custom);
+          setOperationAction(ISD::CONCAT_VECTORS, VT, Legal);
+        }
       }
     }
 
     for (auto VT : {MVT::nxv8i8, MVT::nxv4i16, MVT::nxv2i32})
       setOperationAction(ISD::EXTRACT_SUBVECTOR, VT, Custom);
 
     setOperationAction(ISD::INTRINSIC_WO_CHAIN, MVT::i8, Custom);
     setOperationAction(ISD::INTRINSIC_WO_CHAIN, MVT::i16, Custom);

 // is called the srl will have been lowered to AArch64ISD::VLSHR and the
 // ((OpA + OpB + 1) >> 1) expression will have been changed to (OpB - (~OpA)).
 // This pass can also recognize a variant of this pattern that uses sign
 // extension instead of zero extension and form a srhadd(OpA, OpB) from it.
 SDValue AArch64TargetLowering::LowerTRUNCATE(SDValue Op,
                                              SelectionDAG &DAG) const {
   EVT VT = Op.getValueType();
 
+  if (VT.getScalarType() == MVT::i1) {
+    // Lower i1 truncate to `(x & 1) != 0`.
+    SDLoc dl(Op);
+    EVT OpVT = Op.getOperand(0).getValueType();
+    SDValue Zero = DAG.getConstant(0, dl, OpVT);
+    SDValue One = DAG.getConstant(1, dl, OpVT);
+    SDValue And = DAG.getNode(ISD::AND, dl, OpVT, Op.getOperand(0), One);
+    return DAG.getSetCC(dl, VT, And, Zero, ISD::SETNE);
+  }
+
   if (!VT.isVector() || VT.isScalableVector())
     return Op;
 
   if (useSVEForFixedLengthVectorVT(Op.getOperand(0).getValueType()))
     return LowerFixedLengthVectorTruncateToSVE(Op, DAG);
 
   // Since we are looking for a right shift by a constant value of 1 and we are
   // operating on types at least 16 bits in length (sign/zero extended OpA and

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

multiclass sve_prefetch<SDPatternOperator prefetch, ValueType PredTy, Instruction RegImmInst, Instruction RegRegInst, int scale, ComplexPattern AddrCP> {
 
   defm ZIP1_PPP : sve_int_perm_bin_perm_pp<0b000, "zip1", AArch64zip1>;
   defm ZIP2_PPP : sve_int_perm_bin_perm_pp<0b001, "zip2", AArch64zip2>;
   defm UZP1_PPP : sve_int_perm_bin_perm_pp<0b010, "uzp1", AArch64uzp1>;
   defm UZP2_PPP : sve_int_perm_bin_perm_pp<0b011, "uzp2", AArch64uzp2>;
   defm TRN1_PPP : sve_int_perm_bin_perm_pp<0b100, "trn1", AArch64trn1>;
   defm TRN2_PPP : sve_int_perm_bin_perm_pp<0b101, "trn2", AArch64trn2>;
 
+  def : Pat<(nxv4i1 (concat_vectors nxv2i1:$p1, nxv2i1:$p2)),
+            (UZP1_PPP_S $p1, $p2)>;
+  def : Pat<(nxv8i1 (concat_vectors nxv4i1:$p1, nxv4i1:$p2)),
+            (UZP1_PPP_H $p1, $p2)>;
+  def : Pat<(nxv16i1 (concat_vectors nxv8i1:$p1, nxv8i1:$p2)),
+            (UZP1_PPP_B $p1, $p2)>;
+
   defm CMPHS_PPzZZ : sve_int_cmp_0<0b000, "cmphs", SETUGE, SETULE>;
   defm CMPHI_PPzZZ : sve_int_cmp_0<0b001, "cmphi", SETUGT, SETULT>;
   defm CMPGE_PPzZZ : sve_int_cmp_0<0b100, "cmpge", SETGE, SETLE>;
   defm CMPGT_PPzZZ : sve_int_cmp_0<0b101, "cmpgt", SETGT, SETLT>;
   defm CMPEQ_PPzZZ : sve_int_cmp_0<0b110, "cmpeq", SETEQ, SETEQ>;
   defm CMPNE_PPzZZ : sve_int_cmp_0<0b111, "cmpne", SETNE, SETNE>;
 
   defm CMPEQ_WIDE_PPzZZ : sve_int_cmp_0_wide<0b010, "cmpeq", int_aarch64_sve_cmpeq_wide>;

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

 define <vscale x 2 x i32> @trunc_i64toi32(<vscale x 2 x i64> %in) {
 ; CHECK-LABEL: trunc_i64toi32:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    ret
 entry:
   %out = trunc <vscale x 2 x i64> %in to <vscale x 2 x i32>
   ret <vscale x 2 x i32> %out
 }
+
+; Truncating to i1 requires convert it to a cmp
+
+define <vscale x 2 x i1> @trunc_i64toi1(<vscale x 2 x i64> %in) {
+; CHECK-LABEL: trunc_i64toi1:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    and z0.d, z0.d, #0x1
+; CHECK-NEXT:    cmpne p0.d, p0/z, z0.d, #0
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 2 x i64> %in to <vscale x 2 x i1>
+  ret <vscale x 2 x i1> %out
+}
+
+define <vscale x 4 x i1> @trunc_i64toi1_split(<vscale x 4 x i64> %in) {
+; CHECK-LABEL: trunc_i64toi1_split:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    and z1.d, z1.d, #0x1
+; CHECK-NEXT:    and z0.d, z0.d, #0x1
+; CHECK-NEXT:    cmpne p1.d, p0/z, z1.d, #0
+; CHECK-NEXT:    cmpne p0.d, p0/z, z0.d, #0
+; CHECK-NEXT:    uzp1 p0.s, p0.s, p1.s
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 4 x i64> %in to <vscale x 4 x i1>
+  ret <vscale x 4 x i1> %out
+}
+
+define <vscale x 8 x i1> @trunc_i64toi1_split2(<vscale x 8 x i64> %in) {
+; CHECK-LABEL: trunc_i64toi1_split2:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    and z3.d, z3.d, #0x1
+; CHECK-NEXT:    and z2.d, z2.d, #0x1
+; CHECK-NEXT:    and z1.d, z1.d, #0x1
+; CHECK-NEXT:    and z0.d, z0.d, #0x1
+; CHECK-NEXT:    cmpne p1.d, p0/z, z3.d, #0
+; CHECK-NEXT:    cmpne p2.d, p0/z, z2.d, #0
+; CHECK-NEXT:    uzp1 p1.s, p2.s, p1.s
+; CHECK-NEXT:    cmpne p2.d, p0/z, z1.d, #0
+; CHECK-NEXT:    cmpne p0.d, p0/z, z0.d, #0
+; CHECK-NEXT:    uzp1 p0.s, p0.s, p2.s
+; CHECK-NEXT:    uzp1 p0.h, p0.h, p1.h
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 8 x i64> %in to <vscale x 8 x i1>
+  ret <vscale x 8 x i1> %out
+}
+
+define <vscale x 16 x i1> @trunc_i64toi1_split3(<vscale x 16 x i64> %in) {
+; CHECK-LABEL: trunc_i64toi1_split3:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    and z7.d, z7.d, #0x1
+; CHECK-NEXT:    and z6.d, z6.d, #0x1
+; CHECK-NEXT:    and z5.d, z5.d, #0x1
+; CHECK-NEXT:    and z4.d, z4.d, #0x1
+; CHECK-NEXT:    and z3.d, z3.d, #0x1
+; CHECK-NEXT:    and z2.d, z2.d, #0x1
+; CHECK-NEXT:    cmpne p1.d, p0/z, z7.d, #0
+; CHECK-NEXT:    cmpne p2.d, p0/z, z6.d, #0
+; CHECK-NEXT:    cmpne p3.d, p0/z, z5.d, #0
+; CHECK-NEXT:    cmpne p4.d, p0/z, z4.d, #0
+; CHECK-NEXT:    and z1.d, z1.d, #0x1
+; CHECK-NEXT:    and z0.d, z0.d, #0x1
+; CHECK-NEXT:    uzp1 p1.s, p2.s, p1.s
+; CHECK-NEXT:    cmpne p2.d, p0/z, z3.d, #0
+; CHECK-NEXT:    uzp1 p3.s, p4.s, p3.s
+; CHECK-NEXT:    cmpne p4.d, p0/z, z2.d, #0
+; CHECK-NEXT:    uzp1 p2.s, p4.s, p2.s
+; CHECK-NEXT:    cmpne p4.d, p0/z, z1.d, #0
+; CHECK-NEXT:    cmpne p0.d, p0/z, z0.d, #0
+; CHECK-NEXT:    uzp1 p0.s, p0.s, p4.s
+; CHECK-NEXT:    uzp1 p1.h, p3.h, p1.h
+; CHECK-NEXT:    uzp1 p0.h, p0.h, p2.h
+; CHECK-NEXT:    uzp1 p0.b, p0.b, p1.b
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 16 x i64> %in to <vscale x 16 x i1>
+  ret <vscale x 16 x i1> %out
+}
+
+
+define <vscale x 4 x i1> @trunc_i32toi1(<vscale x 4 x i32> %in) {
+; CHECK-LABEL: trunc_i32toi1:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    and z0.s, z0.s, #0x1
+; CHECK-NEXT:    cmpne p0.s, p0/z, z0.s, #0
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 4 x i32> %in to <vscale x 4 x i1>
+  ret <vscale x 4 x i1> %out
+}
+
+define <vscale x 8 x i1> @trunc_i16toi1(<vscale x 8 x i16> %in) {
+; CHECK-LABEL: trunc_i16toi1:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    and z0.h, z0.h, #0x1
+; CHECK-NEXT:    cmpne p0.h, p0/z, z0.h, #0
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 8 x i16> %in to <vscale x 8 x i1>
+  ret <vscale x 8 x i1> %out
+}
+
+define <vscale x 16 x i1> @trunc_i8toi1(<vscale x 16 x i8> %in) {
+; CHECK-LABEL: trunc_i8toi1:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.b
+; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    cmpne p0.b, p0/z, z0.b, #0
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 16 x i8> %in to <vscale x 16 x i1>
+  ret <vscale x 16 x i1> %out
+}