[llvm] [AArch64] Use brk{a, b} for a lane mask from cttz.elts (PR #178674)

[Sander de Smalen via llvm-commits]

================
@@ -6027,6 +6027,44 @@ static SDValue optimizeIncrementingWhile(SDNode *N, SelectionDAG &DAG,
   return SDValue();
 }
 
+// Match get.active.lane.mask(0, cttz.elts(x)) -> brkb(x)
+// Match get.active.lane.mask(0, add(cttz.elts(x), 1)) -> brka(x)
+static SDValue optimizeBrk(SDNode *N, SelectionDAG &DAG) {
+  SDLoc DL(N);
+  EVT VT = N->getValueType(0);
+  // Lower bound must be 0.
+  if (!isZeroOrZeroSplat(N->getOperand(0)))
+    return SDValue();
+
+  SDValue Upper = N->getOperand(1);
+
+  // Default to brkb, switch to brka if we find a +1.
+  unsigned BrkID = Intrinsic::aarch64_sve_brkb_z;
+  if (Upper->getOpcode() == ISD::ADD && isOneOrOneSplat(Upper.getOperand(1))) {
+    Upper = Upper.getOperand(0);
+    BrkID = Intrinsic::aarch64_sve_brka_z;
+  }
+
+  // We're looking for an upper bound based on CTTZ_ELTS; this would be selected
+  // as a cntp(brk(Pg, Mask)), but if we're just going to make a whilelo based
+  // on that then we just need the brk.
+  if (Upper.getOpcode() != AArch64ISD::CTTZ_ELTS || !VT.isScalableVector())
+    return SDValue();
+
+  SDValue Mask = Upper->getOperand(0);
+  const APInt &PgPattern = Upper.getConstantOperandAPInt(1);
+  SDValue Pg = getPTrue(DAG, DL, VT, PgPattern.getZExtValue());
----------------
sdesmalen-arm wrote:

I should probably have mentioned this on the previous PR, but in retrospect I think it makes more sense to just pass a predicate vector to the `AArch64ISD::CTTZ_ELTS` rather than the predicate pattern. That avoids having to write PTRUE in all the patterns, and is also more intuitive, and allows for passing predicates that are not necessarily a pattern (if that's needed for whatever reason).

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