- Add isLane1KnownActive helper

- Extend canRemovePTestInstr for PTEST_PP_FIRST
- Remove getVectorElementType() != MVT::i1 check for Mask

Author: kmclaughlin-arm

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -20370,7 +20370,7 @@ static bool hasPairwiseAdd(unsigned Opcode, EVT VT, bool FullFP16) {
 }
 
 static SDValue getPTest(SelectionDAG &DAG, EVT VT, SDValue Pg, SDValue Op,
-                        AArch64CC::CondCode Cond, bool EmitCSel = true);
+                        AArch64CC::CondCode Cond);
 
 static bool isPredicateCCSettingOp(SDValue N) {
   if ((N.getOpcode() == ISD::SETCC) ||
@@ -20495,7 +20495,6 @@ static SDValue
 performExtractVectorEltCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
                                const AArch64Subtarget *Subtarget) {
   assert(N->getOpcode() == ISD::EXTRACT_VECTOR_ELT);
-
   if (SDValue Res = performFirstTrueTestVectorCombine(N, DCI, Subtarget))
     return Res;
   if (SDValue Res = performLastTrueTestVectorCombine(N, DCI, Subtarget))
@@ -22536,7 +22535,7 @@ static SDValue tryConvertSVEWideCompare(SDNode *N, ISD::CondCode CC,
 }
 
 static SDValue getPTest(SelectionDAG &DAG, EVT VT, SDValue Pg, SDValue Op,
-                        AArch64CC::CondCode Cond, bool EmitCSel) {
+                        AArch64CC::CondCode Cond) {
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
 
   SDLoc DL(Op);
@@ -22569,8 +22568,6 @@ static SDValue getPTest(SelectionDAG &DAG, EVT VT, SDValue Pg, SDValue Op,
 
   // Set condition code (CC) flags.
   SDValue Test = DAG.getNode(PTest, DL, MVT::i32, Pg, Op);
-  if (!EmitCSel)
-    return Test;
 
   // Convert CC to integer based on requested condition.
   // NOTE: Cond is inverted to promote CSEL's removal when it feeds a compare.
@@ -27237,6 +27234,21 @@ static bool isLanes1toNKnownZero(SDValue Op) {
   }
 }
 
+// Return true if the vector operation can guarantee that the first lane of its
+// result is active.
+static bool isLane1KnownActive(SDValue Op) {
+  switch (Op.getOpcode()) {
+  default:
+    return false;
+  case AArch64ISD::REINTERPRET_CAST:
+    return isLane1KnownActive(Op->getOperand(0));
+  case ISD::SPLAT_VECTOR:
+    return isOneConstant(Op.getOperand(0));
+  case AArch64ISD::PTRUE:
+    return Op.getConstantOperandVal(0) == AArch64SVEPredPattern::all;
+  };
+}
+
 static SDValue removeRedundantInsertVectorElt(SDNode *N) {
   assert(N->getOpcode() == ISD::INSERT_VECTOR_ELT && "Unexpected node!");
   SDValue InsertVec = N->getOperand(0);
@@ -27532,22 +27544,17 @@ static SDValue performPTestFirstCombine(SDNode *N,
   auto Mask = N->getOperand(0);
   auto Pred = N->getOperand(1);
 
-  if (Mask->getOpcode() == AArch64ISD::REINTERPRET_CAST)
-    Mask = Mask->getOperand(0);
-
   if (Pred->getOpcode() == AArch64ISD::REINTERPRET_CAST)
     Pred = Pred->getOperand(0);
 
-  if (Pred->getValueType(0).getVectorElementType() != MVT::i1 ||
-      !isAllActivePredicate(DAG, Mask))
+  if (!isLane1KnownActive(Mask))
     return SDValue();
 
   if (Pred->getOpcode() == ISD::CONCAT_VECTORS) {
     Pred = Pred->getOperand(0);
-    SDValue Mask = DAG.getSplatVector(Pred->getValueType(0), DL,
-                                      DAG.getAllOnesConstant(DL, MVT::i64));
-    return getPTest(DAG, N->getValueType(0), Mask, Pred,
-                    AArch64CC::FIRST_ACTIVE, /* EmitCSel */ false);
+    Pred = DAG.getNode(AArch64ISD::REINTERPRET_CAST, DL, MVT::nxv16i1, Pred);
+    return DAG.getNode(AArch64ISD::PTEST_FIRST, DL, N->getValueType(0), Mask,
+                       Pred);
   }
 
   return SDValue();

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -1503,6 +1503,10 @@ AArch64InstrInfo::canRemovePTestInstr(MachineInstr *PTest, MachineInstr *Mask,
             getElementSizeForOpcode(PredOpcode))
       return PredOpcode;
 
+    if (PTest->getOpcode() == AArch64::PTEST_PP_FIRST &&
+        isPTrueOpcode(MaskOpcode) && Mask->getOperand(1).getImm() == 31)
+      return PredOpcode;
+
     return {};
   }