[AArch64] Improve codegen for extract.last.active

When SVE support is present, we can use the 'clastb' instruction to
perform the work of extract.last.active.

Author: huntergr-arm

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -490,6 +490,11 @@ class TargetLoweringBase {
     return true;
   }
 
+  virtual bool
+  shouldExpandVectorExtractLastActive(const IntrinsicInst *I) const {
+    return true;
+  }
+
   // Return true if op(vecreduce(x), vecreduce(y)) should be reassociated to
   // vecreduce(op(x, y)) for the reduction opcode RedOpc.
   virtual bool shouldReassociateReduction(unsigned RedOpc, EVT VT) const {

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -6417,6 +6417,13 @@ void SelectionDAGBuilder::visitVectorHistogram(const CallInst &I,
 
 void SelectionDAGBuilder::visitVectorExtractLastActive(const CallInst &I,
                                                        unsigned Intrinsic) {
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+
+  if (!TLI.shouldExpandVectorExtractLastActive(cast<IntrinsicInst>(&I))) {
+    visitTargetIntrinsic(I, Intrinsic);
+    return;
+  }
+
   assert(Intrinsic == Intrinsic::experimental_vector_extract_last_active &&
          "Tried lowering invalid vector extract last");
   SDLoc sdl = getCurSDLoc();
@@ -6432,7 +6439,6 @@ void SelectionDAGBuilder::visitVectorExtractLastActive(const CallInst &I,
   ConstantRange VScaleRange(1, /*isFullSet=*/true); // Dummy value.
   if (DataVT.isScalableVector())
     VScaleRange = getVScaleRange(I.getCaller(), 64);
-  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   unsigned EltWidth = TLI.getBitWidthForCttzElements(
       I.getType(), DataVT.getVectorElementCount(), /*ZeroIsPoison=*/true,
       &VScaleRange);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -2072,6 +2072,23 @@ bool AArch64TargetLowering::shouldExpandVectorMatch(EVT VT,
   return true;
 }
 
+bool AArch64TargetLowering::shouldExpandVectorExtractLastActive(
+    const IntrinsicInst *I) const {
+  // 'clastb' requires SVE support.
+  if (!Subtarget->hasSVE())
+    return true;
+
+  // Check if the input data vector is a legal supported type.
+  EVT VT = EVT::getEVT(I->getArgOperand(0)->getType());
+  EVT ScalarVT = VT.getScalarType();
+
+  if (ScalarVT != MVT::i8 && ScalarVT != MVT::i16 && ScalarVT != MVT::i32 &&
+      ScalarVT != MVT::i64 && ScalarVT != MVT::f32 && ScalarVT != MVT::f64)
+    return true;
+
+  return VT.getStoreSizeInBits().getKnownMinValue() != 128;
+}
+
 void AArch64TargetLowering::addTypeForFixedLengthSVE(MVT VT) {
   assert(VT.isFixedLengthVector() && "Expected fixed length vector type!");
 
@@ -6405,6 +6422,22 @@ SDValue AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
   case Intrinsic::experimental_vector_match: {
     return LowerVectorMatch(Op, DAG);
   }
+  case Intrinsic::experimental_vector_extract_last_active: {
+    SDValue Data = Op.getOperand(1);
+    SDValue Mask = Op.getOperand(2);
+    SDValue PassThru = Op.getOperand(3);
+    EVT VT = Op.getValueType();
+    EVT DataVT = Data.getValueType();
+
+    if (DataVT.isFixedLengthVector()) {
+      EVT ContainerVT = getContainerForFixedLengthVector(DAG, DataVT);
+      EVT MaskVT = ContainerVT.changeElementType(MVT::i1);
+      Data = convertToScalableVector(DAG, ContainerVT, Data);
+      Mask = convertToScalableVector(DAG, MaskVT, Mask);
+    }
+
+    return DAG.getNode(AArch64ISD::CLASTB_N, dl, VT, Mask, PassThru, Data);
+  }
   }
 }
 
@@ -27192,6 +27225,18 @@ void AArch64TargetLowering::ReplaceNodeResults(
       Results.push_back(DAG.getNode(ISD::TRUNCATE, DL, VT, V));
       return;
     }
+    case Intrinsic::experimental_vector_extract_last_active: {
+      assert((VT == MVT::i8 || VT == MVT::i16) &&
+             "custom lowering for unexpected type");
+      SDLoc DL(N);
+      auto PassThru =
+          DAG.getNode(ISD::ANY_EXTEND, DL, MVT::i32, N->getOperand(3));
+      auto Extract =
+          DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, MVT::i32, N->getOperand(0),
+                      N->getOperand(1), N->getOperand(2), PassThru);
+      Results.push_back(DAG.getNode(ISD::TRUNCATE, DL, VT, Extract));
+      return;
+    }
     }
   }
   case ISD::READ_REGISTER: {

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -987,6 +987,9 @@ class AArch64TargetLowering : public TargetLowering {
 
   bool shouldExpandVectorMatch(EVT VT, unsigned SearchSize) const override;
 
+  bool
+  shouldExpandVectorExtractLastActive(const IntrinsicInst *) const override;
+
   /// If a change in streaming mode is required on entry to/return from a
   /// function call it emits and returns the corresponding SMSTART or SMSTOP
   /// node. \p Condition should be one of the enum values from