[llvm] [LLVM][CodeGen][AArch64] Implement expansion for fixed-length get.active.lane.mask. (PR #225459)

[Paul Walker via llvm-commits]

================
@@ -1737,6 +1743,47 @@ SDValue VectorLegalizer::ExpandVP_REM(SDNode *Node) {
   return DAG.getNode(ISD::SUB, DL, VT, Dividend, Mul);
 }
 
+SDValue VectorLegalizer::ExpandGET_ACTIVE_LANE_MASK(SDNode *N) {
+  SDLoc DL(N);
+
+  SDValue Start = N->getOperand(0);
+  SDValue End = N->getOperand(1);
+  EVT VT = N->getValueType(0);
+  EVT OpVT = Start.getValueType();
+
+  if (VT.isScalableVector())
+    return SDValue();
+
+  // Try a promoted comparison type to simplify saturation.
+  EVT PromoteVT = VT.changeVectorElementType(*DAG.getContext(), OpVT);
+  if (TLI.isTypeLegal(PromoteVT) &&
+      isUIntN(OpVT.getScalarSizeInBits(), VT.getVectorNumElements())) {
+    SDValue StartV = DAG.getSplat(PromoteVT, DL, Start);
+    SDValue Seq = DAG.getStepVector(DL, PromoteVT);
+    Seq = DAG.getNode(ISD::UADDSAT, DL, PromoteVT, Seq, StartV);
+
+    EVT MaskVT = TLI.getSetCCResultType(DAG.getDataLayout(), *DAG.getContext(),
+                                        PromoteVT);
+    SDValue EndV = DAG.getSplat(PromoteVT, DL, End);
+    SDValue Mask = DAG.getSetCC(DL, MaskVT, Seq, EndV, ISD::SETULT);
+    return DAG.getBoolExtOrTrunc(Mask, DL, VT, PromoteVT);
+  }
+
+  // Is VT's element type big enough to hold all rebased indices?
+  if (!isUIntN(VT.getScalarSizeInBits(), VT.getVectorNumElements()))
+    return SDValue();
+
+  // Rebase and saturate the termination value.
+  SDValue Max = DAG.getConstant(maxUIntN(VT.getScalarSizeInBits()), DL, OpVT);
+  End = DAG.getNode(ISD::USUBSAT, DL, OpVT, End, Start);
+  End = DAG.getNode(ISD::UMIN, DL, OpVT, End, Max);
----------------
paulwalker-arm wrote:

Yes.

https://github.com/llvm/llvm-project/pull/225459