[llvm-branch-commits] [llvm] [AArch64] Avoid materializing full masks for extractelement users of predicate-as-counter (PR #220961)

[Sander de Smalen via llvm-branch-commits]

================
@@ -229,6 +229,41 @@ static Value *createWhileLO(IRBuilder<> &Builder, unsigned ElementSizeInBits,
       WhileLO, {Start, End, Builder.getInt32(VectorScale)}, "pac.mask");
 }
 
+/// If \p UserI is an extractelement use of the original loop mask, attempt to
+/// rewrite it to `extractelement(pext(count, 0), idx)` if the extract index is
+/// known to be within the first mask section (which means pext index = 0). If
+/// the user of the extract is a branch and the source of the mask is a whilelo,
+/// this form can be optimized into checking the status flags.
+static bool tryRewriteExtractElement(Instruction &UserI,
+                                     const MaskRewriteCandidate &C,
+                                     Value *Count) {
+  auto *EEI = dyn_cast<ExtractElementInst>(&UserI);
+  if (!EEI)
+    return false;
+
+  ElementCount LegalEC = getSVEElementCount(C.ElementSizeInBits);
+  auto *Idx = dyn_cast<ConstantInt>(EEI->getIndexOperand());
+  if (!Idx || Idx->getValue().uge(LegalEC.getKnownMinValue()))
+    return false;
+
+  IRBuilder<> Builder(EEI);
+  Builder.SetCurrentDebugLocation(EEI->getDebugLoc());
+
+  Module *M = EEI->getModule();
+  FunctionCallee PExt = Intrinsic::getOrInsertDeclaration(
+      M, Intrinsic::aarch64_sve_pext,
+      {VectorType::get(Builder.getInt1Ty(), LegalEC)});
+  auto *ExtractMask =
+      Builder.CreateCall(PExt, {Count, Builder.getInt32(0)}, "pac.pext");
----------------
sdesmalen-arm wrote:

nit: can use CreateIntrinsic directly:
```suggestion
  auto *ExtractMask = Builder.CreateIntrinsic(
      Intrinsic::aarch64_sve_pext,                                                                         
      {VectorType::get(Builder.getInt1Ty(), LegalEC)},                                     
      {Count, Builder.getInt32(0)}, /*FMFSource=*/{}, "pn.pext");
```

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