[AArch64] Enhance last active true vector combine

[Allen]

Summary

Last active extracting will output LASTB + WHILELS, and the WHILELS itself
is a flag-setting operation, so perform it unconditionally.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14444–14445

nit: this comment only makes sense compared to the previous code which had an explicit check for SETCC, so I think you can just remove this comment altogether.

Can you merge this condition with the condition above, e.g.:

if (Idx.getOpcode() != ISD::ADD || !isAllOnesConstant(Idx.getOperand(1)))

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 performLastTrueTestVectorCombine(SDNode *N,
                                  TargetLowering::DAGCombinerInfo &DCI,
                                  const AArch64Subtarget *Subtarget) {
   assert(N->getOpcode() == ISD::EXTRACT_VECTOR_ELT);
   // Make sure PTEST is legal types.
   if (!Subtarget->hasSVE() || DCI.isBeforeLegalize())
     return SDValue();
 
-  SDValue SetCC = N->getOperand(0);
-  EVT OpVT = SetCC.getValueType();
+  SDValue N0 = N->getOperand(0);
+  EVT OpVT = N0.getValueType();
 
   if (!OpVT.isScalableVector() || OpVT.getVectorElementType() != MVT::i1)
     return SDValue();
 
   // Idx == (add (mul vscale, NumEls), -1)
   SDValue Idx = N->getOperand(1);
   if (Idx.getOpcode() != ISD::ADD)
     return SDValue();
 
   SDValue VS = Idx.getOperand(0);
   if (VS.getOpcode() != ISD::VSCALE)
     return SDValue();
 
   unsigned NumEls = OpVT.getVectorElementCount().getKnownMinValue();
   if (VS.getConstantOperandVal(0) != NumEls)
     return SDValue();
 
-  // Restricted the DAG combine to only cases where we're extracting from a
-  // flag-setting operation
-  auto *CI = dyn_cast<ConstantSDNode>(Idx.getOperand(1));
-  if (!CI || !CI->isAllOnes() || SetCC.getOpcode() != ISD::SETCC)
+  // Last active extracting will output LASTB + WHILELS, and the WHILELS itself
+  // is a flag-setting operation, so perform it unconditionally.
+  if (!isAllOnesConstant(Idx.getOperand(1)))
     return SDValue();
 
   // Extracts of lane EC-1 for SVE can be expressed as PTEST(Op, LAST) ? 1 : 0
   SelectionDAG &DAG = DCI.DAG;
   SDValue Pg = getPTrue(DAG, SDLoc(N), OpVT, AArch64SVEPredPattern::all);
-  return getPTest(DAG, N->getValueType(0), Pg, SetCC, AArch64CC::LAST_ACTIVE);
+  return getPTest(DAG, N->getValueType(0), Pg, N0, AArch64CC::LAST_ACTIVE);
 }
 
 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;

llvm/test/CodeGen/AArch64/sve-extract-element.ll

 ; CHECK-NEXT:    ret
   %b = extractelement <vscale x 8 x i1> %a, i32 9
   ret i1 %b
 }
 
 define i1 @test_last_8xi1(<vscale x 8 x i1> %a) #0 {
 ; CHECK-LABEL: test_last_8xi1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    cnth x8
-; CHECK-NEXT:    mov z0.h, p0/z, #1 // =0x1
-; CHECK-NEXT:    sub x8, x8, #1
-; CHECK-NEXT:    whilels p0.h, xzr, x8
-; CHECK-NEXT:    lastb w8, p0, z0.h
-; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ptrue p1.h
+; CHECK-NEXT:    ptest p1, p0.b
+; CHECK-NEXT:    cset w0, lo
 ; CHECK-NEXT:    ret
   %vscale = call i64 @llvm.vscale.i64()
   %shl = shl nuw nsw i64 %vscale, 3
   %idx = add nuw nsw i64 %shl, -1
   %bit = extractelement <vscale x 8 x i1> %a, i64 %idx
   ret i1 %bit
 }