[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:

Thanks for the review, I'm not too fussy on the expansion so happy to default to GET_ACTIVE_LANE_MASK. However this has the same problem where get_active_lane_mask isn't expanded at all during DAG legalization, only during initial DAG construction in SelectionDAGBuilder. So currently it will always crash for `!TLI.shouldExpandGetActiveLaneMask`, e.g. on NEON. 

I can try and move that legalization out of SelectionDAGBuilder and into LegalizeVectorOps but that would need to be a separate PR before this. Do you want me to give that a stab before this PR or after this PR?

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