[ARM] Custom lower i1 vector truncates

[dmgreen]

Summary

The ISel patterns we have for truncating to i1's under MVE do seem to be correct. Instead custom lower to icmp(ne, and(x, 1), 0).

[samtebbs]

LGTM. Have you noticed any performance changes with this?

[dmgreen]

Thanks. I made a typo in the summary, it should have said "do not seem to be correct", not "do seem...".

Have you noticed any performance changes with this?

No. I think it should be pretty rare - which is why we did not see this problem before.

llvm/lib/Target/ARM/ARMISelLowering.cpp

for (auto VT : pTypes) {
     setOperationAction(ISD::EXTRACT_SUBVECTOR, VT, Custom);
     setOperationAction(ISD::CONCAT_VECTORS, VT, Custom);
     setOperationAction(ISD::INSERT_VECTOR_ELT, VT, Custom);
     setOperationAction(ISD::EXTRACT_VECTOR_ELT, VT, Custom);
     setOperationAction(ISD::SETCC, VT, Custom);
     setOperationAction(ISD::SCALAR_TO_VECTOR, VT, Expand);
     setOperationAction(ISD::LOAD, VT, Custom);
     setOperationAction(ISD::STORE, VT, Custom);
+    setOperationAction(ISD::TRUNCATE, VT, Custom);
   }
 }
 
 ARMTargetLowering::ARMTargetLowering(const TargetMachine &TM,
                                      const ARMSubtarget &STI)
     : TargetLowering(TM), Subtarget(&STI) {
   RegInfo = Subtarget->getRegisterInfo();
   Itins = Subtarget->getInstrItineraryData();
static SDValue LowerEXTRACT_SUBVECTOR(SDValue Op, SelectionDAG &DAG,
   }
 
   // Now return the result of comparing the subvector with zero,
   // which will generate a real predicate, i.e. v4i1, v8i1 or v16i1.
   return DAG.getNode(ARMISD::VCMPZ, dl, VT, SubVec,
                      DAG.getConstant(ARMCC::NE, dl, MVT::i32));
 }
 
+// Turn a truncate into a predicate (an i1 vector) into icmp(and(x, 1), 0).
+static SDValue LowerTruncatei1(SDValue N, SelectionDAG &DAG,
+                               const ARMSubtarget *ST) {
+  assert(ST->hasMVEIntegerOps() && "Expected MVE!");
+  EVT VT = N.getValueType();
+  assert((VT == MVT::v16i1 || VT == MVT::v8i1 || VT == MVT::v4i1) &&
+         "Expected a vector i1 type!");
+  SDValue Op = N.getOperand(0);
+  EVT FromVT = Op.getValueType();
+  SDLoc DL(N);
+
+  SDValue And =
+      DAG.getNode(ISD::AND, DL, FromVT, Op, DAG.getConstant(1, DL, FromVT));
+  return DAG.getNode(ISD::SETCC, DL, VT, And, DAG.getConstant(0, DL, FromVT),
+                     DAG.getCondCode(ISD::SETNE));
+}
+
 /// isExtendedBUILD_VECTOR - Check if N is a constant BUILD_VECTOR where each
 /// element has been zero/sign-extended, depending on the isSigned parameter,
 /// from an integer type half its size.
 static bool isExtendedBUILD_VECTOR(SDNode *N, SelectionDAG &DAG,
                                    bool isSigned) {
   // A v2i64 BUILD_VECTOR will have been legalized to a BITCAST from v4i32.
   EVT VT = N->getValueType(0);
   if (VT == MVT::v2i64 && N->getOpcode() == ISD::BITCAST) {
SDValue ARMTargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
   case ISD::SETCCCARRY:    return LowerSETCCCARRY(Op, DAG);
   case ISD::ConstantFP:    return LowerConstantFP(Op, DAG, Subtarget);
   case ISD::BUILD_VECTOR:  return LowerBUILD_VECTOR(Op, DAG, Subtarget);
   case ISD::VECTOR_SHUFFLE: return LowerVECTOR_SHUFFLE(Op, DAG, Subtarget);
   case ISD::EXTRACT_SUBVECTOR: return LowerEXTRACT_SUBVECTOR(Op, DAG, Subtarget);
   case ISD::INSERT_VECTOR_ELT: return LowerINSERT_VECTOR_ELT(Op, DAG);
   case ISD::EXTRACT_VECTOR_ELT: return LowerEXTRACT_VECTOR_ELT(Op, DAG, Subtarget);
   case ISD::CONCAT_VECTORS: return LowerCONCAT_VECTORS(Op, DAG, Subtarget);
+  case ISD::TRUNCATE:      return LowerTruncatei1(Op, DAG, Subtarget);
   case ISD::FLT_ROUNDS_:   return LowerFLT_ROUNDS_(Op, DAG);
   case ISD::MUL:           return LowerMUL(Op, DAG);
   case ISD::SDIV:
     if (Subtarget->isTargetWindows() && !Op.getValueType().isVector())
       return LowerDIV_Windows(Op, DAG, /* Signed */ true);
     return LowerSDIV(Op, DAG, Subtarget);
   case ISD::UDIV:
     if (Subtarget->isTargetWindows() && !Op.getValueType().isVector())

llvm/lib/Target/ARM/ARMInstrMVE.td

def : Pat<(v4i32 (sext  (v4i1  VCCR:$pred))),
             (v4i32 (MVE_VPSEL (MVE_VMOVimmi8 255), (MVE_VMOVimmi32 0), ARMVCCNone, VCCR:$pred))>;
 
   def : Pat<(v16i8 (anyext  (v16i1 VCCR:$pred))),
             (v16i8 (MVE_VPSEL (MVE_VMOVimmi8 1), (MVE_VMOVimmi8 0), ARMVCCNone, VCCR:$pred))>;
   def : Pat<(v8i16 (anyext  (v8i1  VCCR:$pred))),
             (v8i16 (MVE_VPSEL (MVE_VMOVimmi16 1), (MVE_VMOVimmi16 0), ARMVCCNone, VCCR:$pred))>;
   def : Pat<(v4i32 (anyext  (v4i1  VCCR:$pred))),
             (v4i32 (MVE_VPSEL (MVE_VMOVimmi32 1), (MVE_VMOVimmi32 0), ARMVCCNone, VCCR:$pred))>;
-
-  def : Pat<(v16i1 (trunc (v16i8 MQPR:$v1))),
-            (v16i1 (MVE_VCMPi32r (v16i8 MQPR:$v1), ZR, ARMCCne))>;
-  def : Pat<(v8i1 (trunc (v8i16  MQPR:$v1))),
-            (v8i1 (MVE_VCMPi32r (v8i16 MQPR:$v1), ZR, ARMCCne))>;
-  def : Pat<(v4i1 (trunc (v4i32  MQPR:$v1))),
-            (v4i1 (MVE_VCMPi32r (v4i32 MQPR:$v1), ZR, ARMCCne))>;
 }
 
 let Predicates = [HasMVEFloat] in {
   // Pred <-> Float
   // 112 is 1.0 in float
   def : Pat<(v4f32 (uint_to_fp (v4i1 VCCR:$pred))),
             (v4f32 (MVE_VPSEL (v4f32 (MVE_VMOVimmf32 112)), (v4f32 (MVE_VMOVimmi32 0)), ARMVCCNone, VCCR:$pred))>;
   // 2620 in 1.0 in half

llvm/test/CodeGen/Thumb2/mve-pred-ext.ll

entry:
   %0 = zext <2 x i1> %c to <2 x i64>
   ret <2 x i64> %0
 }
 
 
 define arm_aapcs_vfpcc <4 x i32> @trunc_v4i1_v4i32(<4 x i32> %src) {
 ; CHECK-LABEL: trunc_v4i1_v4i32:
 ; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    vmov.i32 q2, #0x1
 ; CHECK-NEXT:    vmov.i32 q1, #0x0
-; CHECK-NEXT:    vcmp.i32 ne, q0, zr
+; CHECK-NEXT:    vand q2, q0, q2
+; CHECK-NEXT:    vcmp.i32 ne, q2, zr
 ; CHECK-NEXT:    vpsel q0, q0, q1
 ; CHECK-NEXT:    bx lr
 entry:
   %0 = trunc <4 x i32> %src to <4 x i1>
   %1 = select <4 x i1> %0, <4 x i32> %src, <4 x i32> zeroinitializer
   ret <4 x i32> %1
 }
 
 define arm_aapcs_vfpcc <8 x i16> @trunc_v8i1_v8i16(<8 x i16> %src) {
 ; CHECK-LABEL: trunc_v8i1_v8i16:
 ; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    vmov.i16 q2, #0x1
 ; CHECK-NEXT:    vmov.i32 q1, #0x0
-; CHECK-NEXT:    vcmp.i32 ne, q0, zr
+; CHECK-NEXT:    vand q2, q0, q2
+; CHECK-NEXT:    vcmp.i16 ne, q2, zr
 ; CHECK-NEXT:    vpsel q0, q0, q1
 ; CHECK-NEXT:    bx lr
 entry:
   %0 = trunc <8 x i16> %src to <8 x i1>
   %1 = select <8 x i1> %0, <8 x i16> %src, <8 x i16> zeroinitializer
   ret <8 x i16> %1
 }
 
 define arm_aapcs_vfpcc <16 x i8> @trunc_v16i1_v16i8(<16 x i8> %src) {
 ; CHECK-LABEL: trunc_v16i1_v16i8:
 ; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    vmov.i8 q2, #0x1
 ; CHECK-NEXT:    vmov.i32 q1, #0x0
-; CHECK-NEXT:    vcmp.i32 ne, q0, zr
+; CHECK-NEXT:    vand q2, q0, q2
+; CHECK-NEXT:    vcmp.i8 ne, q2, zr
 ; CHECK-NEXT:    vpsel q0, q0, q1
 ; CHECK-NEXT:    bx lr
 entry:
   %0 = trunc <16 x i8> %src to <16 x i1>
   %1 = select <16 x i1> %0, <16 x i8> %src, <16 x i8> zeroinitializer
   ret <16 x i8> %1
 }