[llvm] [IR][RFC] Add @llvm.mask.beforefirst intrinsic (PR #203874)

[Luke Lau via llvm-commits]

================
@@ -1741,6 +1747,40 @@ SDValue VectorLegalizer::ExpandLOOP_DEPENDENCE_MASK(SDNode *N) {
   return TLI.expandLoopDependenceMask(N, DAG);
 }
 
+SDValue VectorLegalizer::ExpandMASK_BEFOREFIRST(SDNode *N) {
+  // Expand to (setcc ult (step-vector), (cttz_elts x))
+  SDLoc DL(N);
+  EVT VT = N->getValueType(0);
+
+  // Try to expand via get_active_lane_mask if supported.
+  EVT BoolVT = VT.changeVectorElementType(*DAG.getContext(), MVT::i1);
+  EVT VecIdxVT = TLI.getVectorIdxTy(DAG.getDataLayout());
+  if (!TLI.shouldExpandGetActiveLaneMask(BoolVT, VecIdxVT)) {
+    // Perform cttz_elts in bool VT to get the custom target lowering.
+    SDValue CttzElts =
+        DAG.getNode(ISD::CTTZ_ELTS, DL, VecIdxVT,
+                    DAG.getNode(ISD::TRUNCATE, DL, BoolVT, N->getOperand(0)));
+    return DAG.getNode(ISD::GET_ACTIVE_LANE_MASK, DL, VT,
+                       DAG.getConstant(0, DL, VecIdxVT), CttzElts);
+  }
+
+  // Compute the type for the cttz_elts.
+  ConstantRange VScaleRange(1, /*isFullSet=*/true); // Fixed length default.
+  if (VT.isScalableVector())
+    VScaleRange = getVScaleRange(&DAG.getMachineFunction().getFunction(), 64);
+  uint64_t EltWidth = TLI.getBitWidthForCttzElements(
+      EVT(TLI.getVectorIdxTy(DAG.getDataLayout())), VT.getVectorElementCount(),
+      /*ZeroIsPoison=*/false, &VScaleRange);
+
+  EVT EltVT = EVT::getIntegerVT(*DAG.getContext(), EltWidth);
+  EVT CttzEltsVT =
+      EVT::getVectorVT(*DAG.getContext(), EltVT, VT.getVectorElementCount());
+  SDValue CttzElts = DAG.getSplatVector(
+      CttzEltsVT, DL, DAG.getNode(ISD::CTTZ_ELTS, DL, EltVT, N->getOperand(0)));
----------------
lukel97 wrote:

I gave it a stab but ran into some regressions on NEON, that would be great if you could take a look at it. Although this was with removing the SelectionDAGBuilder legalization. I guess you would need to relax some cases in AArch64TargetLowering::shouldExpandGetActiveLaneMask to test this, since I don't think anything will create a GET_ACTIVE_LANE_MASK during legalization?

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