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

[Paul Walker via llvm-commits]

================
@@ -6022,6 +6022,43 @@ static SDValue optimizeIncrementingWhile(SDNode *N, SelectionDAG &DAG,
   return SDValue();
 }
 
+// Match active.lane.mask(cttz.elts(x)) -> brkb(x)
+// Match active.lane.mask(add(cttz.elts(x), 1)) -> brka(x)
+static SDValue optimizeBrk(SDNode *N, SelectionDAG &DAG) {
+  EVT VT = N->getValueType(0);
+
+  // TODO: Do we need to do anything for fixed types post legalization?
+  if (!VT.isScalableVT())
+    return SDValue();
+
+  SDValue Op = N->getOperand(1);
+
+  // Default to brkb, switch to brka if we find a +1.
+  unsigned BrkID = Intrinsic::aarch64_sve_brkb_z;
+  if (Op->getOpcode() == ISD::ADD && isOneOrOneSplat(Op.getOperand(1))) {
+    Op = Op.getOperand(0);
+    BrkID = Intrinsic::aarch64_sve_brka_z;
+  }
+
+  if (Op.getOpcode() == AArch64ISD::CTTZ_ELTS) {
+    SDValue Mask = Op->getOperand(0);
+    SDLoc DL(N);
+    SDValue PTrue = getPTrue(DAG, DL, VT, AArch64SVEPredPattern::all);
+
+    // brk{a,b} only support .b forms, so reinterpret to make sure all our
+    // p regs will match.
+    PTrue = DAG.getNode(AArch64ISD::REINTERPRET_CAST, DL, MVT::nxv16i1, PTrue);
+    SDValue MaskR =
+        DAG.getNode(AArch64ISD::REINTERPRET_CAST, DL, MVT::nxv16i1, Mask);
----------------
paulwalker-arm wrote:

Presumably that's the point? The PTRUE of the smaller vector when reinterpreted will contain a sequence of 1s and 0s because we're relying on the PTRUE to zero the invisible lanes.

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