[llvm] [AArch64] Add predicate-as-counter loop rewrite pass (PR #218690)
Benjamin Maxwell via llvm-commits
llvm-commits at lists.llvm.org
Tue Aug 25 07:10:12 PDT 2026
https://github.com/MacDue updated https://github.com/llvm/llvm-project/pull/218690
>From 8af82acd17b0bd0a4ba2e7b82f61226ed3b080b7 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Tue, 25 Aug 2026 12:18:21 +0000
Subject: [PATCH] [AArch64] Add predicate-as-counter loop rewrite pass
Add an AArch64 IR loop pass that rewrites wide loop-carried
`llvm.get.active.lane.mask` phis to predicate-as-counter `whilelo` phis
(for SVE2.1 or streaming SME2 targets).
Masked load/store users of the loop mask are rewritten to multi-vector
`ld1/st1 ... pn x2/x4` operations. Remaining users of the original mask
are preserved by materializing vector predicates with `aarch64.sve.pext`
and `aarch64.sve.pext.x2`.
For example, a loop like:
```
entry:
%step = vscale x 64
%mask.entry = @get.active.lane.mask(0, %n)
loop:
%iv = phi i64 [0, entry], [%iv.next, loop]
%mask = phi <vscale x 64 x i1> [%mask.entry, entry],
[%mask.next, loop]
%load = load <vscale x 64 x i8> %src[%iv], %mask
store <vscale x 64 x i8> %load, %dst[%iv], %mask
%iv.next = %iv + %step
%mask.next = @get.active.lane.mask(%iv.next, %n)
br first.active(%mask.next), loop, exit
```
Could be rewritten to:
```
entry:
%step = vscale x 64
%mask.entry = @whilelo.c8(0, %n, VLx4)
loop:
%iv = phi i64 [0, entry], [%iv.next, loop]
%mask = phi target("aarch64.svcount") [%mask.entry, entry],
[%mask.next, loop]
%load = @ld1.pn.x4 <4 x <vscale x 16 x i8>> %src[%iv], %mask
@st1.pn.x4 <4 x <vscale x 16 x i8>> %load, %dest[%iv], %mask
%iv.next = %iv + %step
%mask.next = @whilelo.c8(%iv.next, %n, VLx4)
br first.active(@pext(%mask.next, 0)), loop, exit
```
Assisted-by: Codex
---
llvm/lib/Target/AArch64/AArch64.h | 2 +
.../AArch64PredicateAsCounterLoopRewrites.cpp | 607 ++++++++++++++++++
.../Target/AArch64/AArch64TargetMachine.cpp | 10 +
llvm/lib/Target/AArch64/CMakeLists.txt | 1 +
.../predicate-as-counter-loop-rewrites.ll | 566 ++++++++++++++++
5 files changed, 1186 insertions(+)
create mode 100644 llvm/lib/Target/AArch64/AArch64PredicateAsCounterLoopRewrites.cpp
create mode 100644 llvm/test/CodeGen/AArch64/predicate-as-counter-loop-rewrites.ll
diff --git a/llvm/lib/Target/AArch64/AArch64.h b/llvm/lib/Target/AArch64/AArch64.h
index bcf40a02e3cf5..f2e49d6578f9c 100644
--- a/llvm/lib/Target/AArch64/AArch64.h
+++ b/llvm/lib/Target/AArch64/AArch64.h
@@ -76,6 +76,7 @@ FunctionPass *createAArch64PTrueCoalescingLegacyPass();
FunctionPass *createAArch64CleanupLocalDynamicTLSPass();
FunctionPass *createAArch64CollectLOHPass();
+Pass *createAArch64PredicateAsCounterLoopRewritesPass();
FunctionPass *createSMEPeepholeOptPass();
FunctionPass *createMachineSMEABIPass(CodeGenOptLevel);
FunctionPass *createAArch64SRLTDefineSuperRegsLegacyPass();
@@ -167,6 +168,7 @@ void initializeAArch64A57FPLoadBalancingLegacyPass(PassRegistry &);
void initializeAArch64AdvSIMDScalarLegacyPass(PassRegistry &);
void initializeAArch64AsmPrinterPass(PassRegistry &);
void initializeAArch64PointerAuthLegacyPass(PassRegistry &);
+void initializeAArch64PredicateAsCounterLoopRewritesPass(PassRegistry &);
void initializeAArch64BranchTargetsLegacyPass(PassRegistry &);
void initializeAArch64CFIFixupPass(PassRegistry&);
void initializeAArch64CollectLOHLegacyPass(PassRegistry &);
diff --git a/llvm/lib/Target/AArch64/AArch64PredicateAsCounterLoopRewrites.cpp b/llvm/lib/Target/AArch64/AArch64PredicateAsCounterLoopRewrites.cpp
new file mode 100644
index 0000000000000..7881751137425
--- /dev/null
+++ b/llvm/lib/Target/AArch64/AArch64PredicateAsCounterLoopRewrites.cpp
@@ -0,0 +1,607 @@
+//===- AArch64PredicateAsCounterLoopRewrites.cpp --------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Rewrites IR for loop-carried wide masks that can be represented as
+// predicate-as-counter values. This applies when the mask is used by load/store
+// operations that can be mapped to multi-vector instructions (with +sve2p1).
+//
+// For example, a loop like:
+//
+// entry:
+// %step = vscale x 64
+// %mask.entry = @get.active.lane.mask(0, %n)
+//
+// loop:
+// %iv = phi i64 [0, entry], [%iv.next, loop]
+// %mask = phi <vscale x 64 x i1> [%mask.entry, entry],
+// [%mask.next, loop]
+//
+// %load = load <vscale x 64 x i8> %src[%iv], %mask
+// store <vscale x 64 x i8> %load, %dst[%iv], %mask
+//
+// %iv.next = %iv + %step
+// %mask.next = @get.active.lane.mask(%iv.next, %n)
+// br first.active(%mask.next), loop, exit
+//
+// Could be rewritten to:
+//
+// entry:
+// %step = vscale x 64
+// %mask.entry = @whilelo.c8(0, %n, VLx4)
+//
+// loop:
+// %iv = phi i64 [0, entry], [%iv.next, loop]
+// %mask = phi target("aarch64.svcount") [%mask.entry, entry],
+// [%mask.next, loop]
+//
+// %load = @ld1.pn.x4 <4 x <vscale x 16 x i8>> %src[%iv], %mask
+// @st1.pn.x4 <4 x <vscale x 16 x i8>> %load, %dest[%iv], %mask
+//
+// %iv.next = %iv + %step
+// %mask.next = @whilelo.c8(%iv.next, %n, VLx4)
+// br first.active(@pext(%mask.next, 0)), loop, exit
+//
+// This replaces the `get.active.lane.mask` intrinsics with AArch64
+// predicate-as-counter `whilelo` intrinsics and updates the mask phi to use the
+// `aarch64.svcount` target type. Within the loop, load/store users are mapped
+// to multi-vector load/store intrinsics where possible. Users that cannot be
+// mapped to multi-vector instructions materialize vector masks using the `pext`
+// intrinsic (which extracts vector predicates from a predicate-as-counter).
+//
+// This pass may be a temporary solution that is removed if we gain support
+// for target-specific VPlan transforms in the loop vectorizer.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AArch64.h"
+#include "AArch64Subtarget.h"
+#include "AArch64TargetMachine.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/LoopPass.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicsAArch64.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils.h"
+#include "llvm/Transforms/Utils/Local.h"
+#include <optional>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "aarch64-predicate-as-counter-loop-rewrites"
+namespace {
+
+STATISTIC(LoopsRewritten, "Number of loops rewritten");
+
+struct MaskRewriteCandidate {
+ /// The preheader block for the loop.
+ BasicBlock *Preheader = nullptr;
+ /// The latch block for the loop.
+ BasicBlock *Latch = nullptr;
+ /// The mask phi node (used by masked operations within the loop).
+ PHINode *MaskPhi = nullptr;
+ /// The initial value for the mask (incoming value from the preheader).
+ IntrinsicInst *StartMask = nullptr;
+ /// The updated value for the mask (incoming value from the loop latch).
+ IntrinsicInst *NextMask = nullptr;
+ /// The multi-vector scale for the predicate-as-counter (2 or 4).
+ unsigned VectorScale = 0;
+ /// The element size (in bits) for the predicate-as-counter.
+ unsigned ElementSizeInBits = 0;
+};
+
+static void logLoopBailout(const Loop &L, const Twine &Reason) {
+ LLVM_DEBUG({
+ dbgs() << "PAC loop rewrite: skipping loop with header ";
+ L.getHeader()->printAsOperand(dbgs(), /*PrintType=*/false);
+ dbgs() << ": " << Reason << "\n";
+ });
+}
+
+static void logMatchFailure(const PHINode &Phi, const Twine &Reason) {
+ LLVM_DEBUG({
+ dbgs() << "PAC loop rewrite: failed to match mask phi ";
+ Phi.printAsOperand(dbgs(), /*PrintType=*/false);
+ dbgs() << ": " << Reason << "\n";
+ });
+}
+
+/// Returns the scalar size in bits for a type according to the data layout.
+static unsigned getScalarSizeInBits(const DataLayout &DL, Type *Ty) {
+ return DL.getTypeSizeInBits(Ty->getScalarType()).getFixedValue();
+}
+
+/// Returns the SVE element count for \p ElementSizeInBits.
+static ElementCount getSVEElementCount(unsigned ElementSizeInBits) {
+ return ElementCount::getScalable(AArch64::SVEBitsPerBlock /
+ ElementSizeInBits);
+}
+
+/// Returns the most common scalar access size of masked loads/stores in loop
+/// \p L where \p MaskPhi is used as the mask. Ties are broken in favor of
+/// larger access sizes.
+static std::optional<unsigned>
+getMostCommonMaskedMemAccessSizeInBits(const Loop &L, PHINode &MaskPhi) {
+ const DataLayout &DL = MaskPhi.getModule()->getDataLayout();
+ DenseMap<unsigned, unsigned> AccessSizeCounts;
+ std::optional<unsigned> BestAccessSizeInBits;
+ unsigned BestAccessSizeCount = 0;
+
+ for (User *U : MaskPhi.users()) {
+ auto *II = dyn_cast<IntrinsicInst>(U);
+ if (!II || !L.contains(II))
+ continue;
+
+ Intrinsic::ID IID = II->getIntrinsicID();
+ if (IID != Intrinsic::masked_load && IID != Intrinsic::masked_store)
+ continue;
+
+ unsigned MaskOpIdx = IID == Intrinsic::masked_load ? 1 : 2;
+ if (II->getArgOperand(MaskOpIdx) != &MaskPhi)
+ continue;
+
+ unsigned AccessSizeInBits = getScalarSizeInBits(DL, II->getAccessType());
+ unsigned AccessSizeCount = ++AccessSizeCounts[AccessSizeInBits];
+
+ if (!BestAccessSizeInBits || AccessSizeCount > BestAccessSizeCount ||
+ (AccessSizeCount == BestAccessSizeCount &&
+ AccessSizeInBits > *BestAccessSizeInBits)) {
+ BestAccessSizeInBits = AccessSizeInBits;
+ BestAccessSizeCount = AccessSizeCount;
+ }
+ }
+
+ return BestAccessSizeInBits;
+}
+
+/// Returns the predicate-as-counter whilelo intrinsic ID for
+/// \p ElementSizeInBits.
+static Intrinsic::ID getWhileLOIntrinsic(unsigned ElementSizeInBits) {
+ switch (ElementSizeInBits) {
+ case 8:
+ return Intrinsic::aarch64_sve_whilelo_c8;
+ case 16:
+ return Intrinsic::aarch64_sve_whilelo_c16;
+ case 32:
+ return Intrinsic::aarch64_sve_whilelo_c32;
+ case 64:
+ return Intrinsic::aarch64_sve_whilelo_c64;
+ default:
+ llvm_unreachable("unsupported predicate-as-counter element size");
+ }
+}
+
+/// Expands the predicate-as-counter mask \p Count into a wide vector mask.
+/// Masks are extracted from the counter using the paired pext intrinsics,
+/// then concatenated to form the wide mask value.
+static Value *buildWideMask(IRBuilder<> &Builder, const MaskRewriteCandidate &C,
+ Value *Count) {
+ ElementCount LegalEC = getSVEElementCount(C.ElementSizeInBits);
+ Module *M = Builder.GetInsertBlock()->getModule();
+ Type *LegalMaskTy = VectorType::get(Builder.getInt1Ty(), LegalEC);
+ FunctionCallee PExtX2 = Intrinsic::getOrInsertDeclaration(
+ M, Intrinsic::aarch64_sve_pext_x2, {LegalMaskTy});
+
+ Value *WideMask = PoisonValue::get(C.MaskPhi->getType());
+ for (unsigned PairOffset = 0; PairOffset != C.VectorScale / 2; ++PairOffset) {
+ auto *Pair = Builder.CreateCall(
+ PExtX2, {Count, Builder.getInt32(PairOffset)}, "pac.pext.pair");
+ for (unsigned SliceInPair = 0; SliceInPair != 2; ++SliceInPair) {
+ Value *Part = Builder.CreateExtractValue(Pair, SliceInPair, "pac.pext");
+ unsigned Slice = PairOffset * 2 + SliceInPair;
+ WideMask = Builder.CreateInsertVector(
+ C.MaskPhi->getType(), WideMask, Part,
+ Slice * LegalEC.getKnownMinValue(), "pac.mask");
+ }
+ }
+
+ return WideMask;
+}
+
+/// Creates a predicate-as-counter whilelo for \p ElementSizeInBits between
+/// \p Start and \p End with multi-vector \p VectorScale (2 or 4).
+static Value *createWhileLO(IRBuilder<> &Builder, unsigned ElementSizeInBits,
+ Value *Start, Value *End, unsigned VectorScale) {
+ if (Start->getType()->getIntegerBitWidth() < 64) {
+ Start = Builder.CreateZExt(Start, Builder.getInt64Ty());
+ End = Builder.CreateZExt(End, Builder.getInt64Ty());
+ }
+ Module *M = Builder.GetInsertBlock()->getModule();
+ auto ID = getWhileLOIntrinsic(ElementSizeInBits);
+ FunctionCallee WhileLO = Intrinsic::getOrInsertDeclaration(M, ID);
+ return Builder.CreateCall(
+ WhileLO, {Start, End, Builder.getInt32(VectorScale)}, "pac.mask");
+}
+
+/// Returns the multi-vector load/store intrinsic ID for \p VectorScale.
+static Intrinsic::ID getPNLoadStoreIntrinsic(unsigned VectorScale,
+ bool IsLoad) {
+ if (VectorScale == 2)
+ return IsLoad ? Intrinsic::aarch64_sve_ld1_pn_x2
+ : Intrinsic::aarch64_sve_st1_pn_x2;
+ if (VectorScale == 4)
+ return IsLoad ? Intrinsic::aarch64_sve_ld1_pn_x4
+ : Intrinsic::aarch64_sve_st1_pn_x4;
+ llvm_unreachable("unsupported predicate-as-counter scale");
+}
+
+/// If \p UserI is a masked-load user of the original loop mask, rewrite it to
+/// a masked multi-vector load (using the predicate-as-counter) if the load
+/// access size matches the predicate-as-counter element size.
+static bool tryRewriteMaskedLoadUser(Instruction &UserI,
+ const MaskRewriteCandidate &C,
+ Value *Count) {
+ auto *II = dyn_cast<IntrinsicInst>(&UserI);
+ if (!II || II->getIntrinsicID() != Intrinsic::masked_load)
+ return false;
+
+ unsigned ElementSizeInBits =
+ getScalarSizeInBits(II->getModule()->getDataLayout(), II->getType());
+ if (ElementSizeInBits != C.ElementSizeInBits)
+ return false;
+
+ if (!isa<PoisonValue, UndefValue>(II->getArgOperand(2)))
+ return false;
+
+ IRBuilder<> Builder(II);
+ Builder.SetCurrentDebugLocation(II->getDebugLoc());
+
+ ElementCount LegalEC = getSVEElementCount(C.ElementSizeInBits);
+ Type *ScalarType = II->getType()->getScalarType();
+ auto *WideDataTy = VectorType::get(ScalarType, LegalEC * C.VectorScale);
+ auto *LegalDataTy = VectorType::get(ScalarType, LegalEC);
+
+ Module *M = II->getModule();
+ FunctionCallee LD1 = Intrinsic::getOrInsertDeclaration(
+ M, getPNLoadStoreIntrinsic(C.VectorScale, /*IsLoad=*/true),
+ {LegalDataTy, II->getArgOperand(0)->getType()});
+ auto *PNLoad =
+ Builder.CreateCall(LD1, {Count, II->getArgOperand(0)}, "pac.ld1");
+
+ // Copy pointer parameter attributes.
+ for (Attribute ParamAttr : II->getParamAttributes(0))
+ PNLoad->addParamAttr(1, ParamAttr);
+
+ // Concatenate all results into the original wide value.
+ Value *WideData = PoisonValue::get(WideDataTy);
+ for (unsigned Slice = 0; Slice != C.VectorScale; ++Slice) {
+ Value *Part = Builder.CreateExtractValue(PNLoad, Slice, "pac.data");
+ WideData = Builder.CreateInsertVector(WideDataTy, WideData, Part,
+ Slice * LegalEC.getKnownMinValue(),
+ "pac.vec");
+ }
+
+ II->replaceAllUsesWith(WideData);
+ II->eraseFromParent();
+ return true;
+}
+
+/// If \p UserI is a masked-store user of the original loop mask, rewrite it to
+/// a masked multi-vector store (using the predicate-as-counter) if the store
+/// access size matches the predicate-as-counter element size.
+static bool tryRewriteMaskedStoreUser(Instruction &UserI,
+ const MaskRewriteCandidate &C,
+ Value *Count) {
+ auto *II = dyn_cast<IntrinsicInst>(&UserI);
+ if (!II || II->getIntrinsicID() != Intrinsic::masked_store)
+ return false;
+
+ unsigned ElementSizeInBits = getScalarSizeInBits(
+ II->getModule()->getDataLayout(), II->getArgOperand(0)->getType());
+ if (ElementSizeInBits != C.ElementSizeInBits)
+ return false;
+
+ IRBuilder<> Builder(II);
+ Builder.SetCurrentDebugLocation(II->getDebugLoc());
+
+ Type *ScalarType = II->getArgOperand(0)->getType()->getScalarType();
+ ElementCount LegalEC = getSVEElementCount(C.ElementSizeInBits);
+ auto *LegalDataTy = VectorType::get(ScalarType, LegalEC);
+
+ SmallVector<Value *, 6> StoreArgs;
+ for (unsigned Slice = 0; Slice != C.VectorScale; ++Slice)
+ StoreArgs.push_back(Builder.CreateExtractVector(
+ LegalDataTy, II->getArgOperand(0), Slice * LegalEC.getKnownMinValue(),
+ "pac.data"));
+ StoreArgs.push_back(Count);
+ StoreArgs.push_back(II->getArgOperand(1));
+
+ Module *M = II->getModule();
+ FunctionCallee ST1 = Intrinsic::getOrInsertDeclaration(
+ M, getPNLoadStoreIntrinsic(C.VectorScale, /*IsLoad=*/false),
+ {LegalDataTy, II->getArgOperand(1)->getType()});
+ auto *PNStore = Builder.CreateCall(ST1, StoreArgs);
+
+ // Copy pointer parameter attributes.
+ for (Attribute ParamAttr : II->getParamAttributes(1))
+ PNStore->addParamAttr(C.VectorScale + 1, ParamAttr);
+
+ II->eraseFromParent();
+ return true;
+}
+
+/// 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");
+
+ Value *Extracted = Builder.CreateExtractElement(
+ ExtractMask, EEI->getIndexOperand(), EEI->getName() + ".pac");
+
+ EEI->replaceAllUsesWith(Extracted);
+ EEI->eraseFromParent();
+ return true;
+}
+
+class AArch64PredicateAsCounterLoopRewrites : public LoopPass {
+public:
+ static char ID;
+
+ AArch64PredicateAsCounterLoopRewrites() : LoopPass(ID) {}
+
+ void getAnalysisUsage(AnalysisUsage &AU) const override {
+ AU.addRequired<TargetPassConfig>();
+ // Require loop simplify to ensure loops have a preheader.
+ AU.addRequiredID(LoopSimplifyID);
+ AU.addPreservedID(LoopSimplifyID);
+ AU.setPreservesCFG();
+ }
+
+ bool runOnLoop(Loop *L, LPPassManager &) override;
+
+private:
+ std::optional<MaskRewriteCandidate> matchMaskPhi(Loop &L, PHINode &Phi) const;
+ bool rewriteCandidate(const MaskRewriteCandidate &C, Loop &L) const;
+};
+
+} // end anonymous namespace
+
+char AArch64PredicateAsCounterLoopRewrites::ID = 0;
+
+INITIALIZE_PASS_BEGIN(AArch64PredicateAsCounterLoopRewrites, DEBUG_TYPE,
+ "AArch64 Predicate As Counter Loop Rewrites", false,
+ false)
+INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
+INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
+INITIALIZE_PASS_END(AArch64PredicateAsCounterLoopRewrites, DEBUG_TYPE,
+ "AArch64 Predicate As Counter Loop Rewrites", false, false)
+
+Pass *llvm::createAArch64PredicateAsCounterLoopRewritesPass() {
+ return new AArch64PredicateAsCounterLoopRewrites();
+}
+
+bool AArch64PredicateAsCounterLoopRewrites::runOnLoop(Loop *L,
+ LPPassManager &) {
+ if (skipLoop(L)) {
+ logLoopBailout(*L, "skipLoop requested the loop to be skipped");
+ return false;
+ }
+
+ Function &F = *L->getHeader()->getParent();
+ auto &TPC = getAnalysis<TargetPassConfig>();
+ const AArch64Subtarget *ST =
+ TPC.getTM<AArch64TargetMachine>().getSubtargetImpl(F);
+ if (!ST->isSVEorStreamingSVEAvailable()) {
+ logLoopBailout(*L, "SVE or streaming SVE is unavailable");
+ return false;
+ }
+ if (!ST->hasSVE2p1() && !(ST->hasSME2() && ST->isStreaming())) {
+ logLoopBailout(*L, "neither SVE2.1 nor SME2 is available");
+ return false;
+ }
+
+ bool Changed = false;
+ BasicBlock *Header = L->getHeader();
+ for (PHINode &Phi : make_early_inc_range(Header->phis())) {
+ auto Candidate = matchMaskPhi(*L, Phi);
+ if (Candidate)
+ Changed |= rewriteCandidate(*Candidate, *L);
+ }
+
+ if (Changed)
+ ++LoopsRewritten;
+
+ return Changed;
+}
+
+static IntrinsicInst *getGetActiveLaneMask(Value *V) {
+ auto *II = dyn_cast<IntrinsicInst>(V);
+ return II && II->getIntrinsicID() == Intrinsic::get_active_lane_mask
+ ? II
+ : nullptr;
+}
+
+std::optional<MaskRewriteCandidate>
+AArch64PredicateAsCounterLoopRewrites::matchMaskPhi(Loop &L,
+ PHINode &Phi) const {
+ BasicBlock *Preheader = L.getLoopPreheader();
+ BasicBlock *Latch = L.getLoopLatch();
+ if (!Preheader) {
+ logMatchFailure(Phi, "loop has no preheader");
+ return std::nullopt;
+ }
+ if (!Latch) {
+ logMatchFailure(Phi, "loop has no latch");
+ return std::nullopt;
+ }
+
+ auto *PhiTy = dyn_cast<ScalableVectorType>(Phi.getType());
+ if (!PhiTy || !PhiTy->getElementType()->isIntegerTy(1)) {
+ logMatchFailure(Phi, "phi type is not a scalable i1 vector mask");
+ return std::nullopt;
+ }
+ if (Phi.getNumIncomingValues() != 2) {
+ logMatchFailure(Phi, Twine("phi has ")
+ .concat(Twine(Phi.getNumIncomingValues()))
+ .concat(" incoming values; expected 2"));
+ return std::nullopt;
+ }
+
+ auto *StartMask =
+ getGetActiveLaneMask(Phi.getIncomingValueForBlock(Preheader));
+ auto *NextMask = getGetActiveLaneMask(Phi.getIncomingValueForBlock(Latch));
+ if (!StartMask) {
+ logMatchFailure(Phi,
+ "preheader incoming value is not get_active_lane_mask");
+ return std::nullopt;
+ }
+ if (!NextMask) {
+ logMatchFailure(Phi, "latch incoming value is not get_active_lane_mask");
+ return std::nullopt;
+ }
+
+ unsigned WideMaskElements = PhiTy->getMinNumElements();
+ if (!isPowerOf2_32(WideMaskElements)) {
+ logMatchFailure(Phi, Twine("wide mask element count is not a power of 2: ")
+ .concat(Twine(WideMaskElements)));
+ return std::nullopt;
+ }
+
+ if (StartMask->getArgOperand(0)->getType()->getIntegerBitWidth() > 64) {
+ logMatchFailure(Phi, "start mask induction operand is wider than i64");
+ return std::nullopt;
+ }
+ if (NextMask->getArgOperand(0)->getType()->getIntegerBitWidth() > 64) {
+ logMatchFailure(Phi, "next mask induction operand is wider than i64");
+ return std::nullopt;
+ }
+
+ std::optional<unsigned> PreferredMaskElementSizeInBits =
+ getMostCommonMaskedMemAccessSizeInBits(L, Phi);
+ if (!PreferredMaskElementSizeInBits) {
+ logMatchFailure(Phi, "mask phi has no masked load/store users in the loop");
+ return std::nullopt;
+ }
+
+ if (!is_contained({8u, 16u, 32u, 64u}, *PreferredMaskElementSizeInBits)) {
+ logMatchFailure(Phi, Twine("unsupported element size in bits: ")
+ .concat(Twine(*PreferredMaskElementSizeInBits)));
+ return std::nullopt;
+ }
+
+ unsigned SVEMaskElements =
+ getSVEElementCount(*PreferredMaskElementSizeInBits).getKnownMinValue();
+ if (WideMaskElements <= SVEMaskElements) {
+ logMatchFailure(Phi, Twine("wide mask element count ")
+ .concat(Twine(WideMaskElements))
+ .concat(" is not wider than the legal mask width ")
+ .concat(Twine(SVEMaskElements)));
+ return std::nullopt;
+ }
+
+ unsigned VectorScale = WideMaskElements / SVEMaskElements;
+ if (VectorScale != 2 && VectorScale != 4) {
+ logMatchFailure(Phi, Twine("unsupported predicate-as-counter scale: ")
+ .concat(Twine(VectorScale)));
+ return std::nullopt;
+ }
+
+ return MaskRewriteCandidate{Preheader,
+ Latch,
+ &Phi,
+ StartMask,
+ NextMask,
+ VectorScale,
+ *PreferredMaskElementSizeInBits};
+}
+
+bool AArch64PredicateAsCounterLoopRewrites::rewriteCandidate(
+ const MaskRewriteCandidate &C, Loop &L) const {
+ IRBuilder<> Builder(C.StartMask);
+ Value *NewStart =
+ createWhileLO(Builder, C.ElementSizeInBits, C.StartMask->getArgOperand(0),
+ C.StartMask->getArgOperand(1), C.VectorScale);
+ Builder.SetInsertPoint(C.NextMask);
+ Value *NewNext =
+ createWhileLO(Builder, C.ElementSizeInBits, C.NextMask->getArgOperand(0),
+ C.NextMask->getArgOperand(1), C.VectorScale);
+
+ Builder.SetInsertPoint(C.MaskPhi);
+ auto *NewPhi =
+ Builder.CreatePHI(NewStart->getType(), 2, C.MaskPhi->getName() + ".pn");
+ NewPhi->addIncoming(NewStart, C.Preheader);
+ NewPhi->addIncoming(NewNext, C.Latch);
+
+ auto RewriteUses = [&](Instruction *OldMask, Value *Count,
+ function_ref<bool(Use & U)> Predicate = nullptr) {
+ SmallVector<Use *, 8> UsesToRewrite;
+ for (Use &U : OldMask->uses()) {
+ if (!Predicate || Predicate(U))
+ UsesToRewrite.push_back(&U);
+ }
+
+ Value *WideMask = nullptr;
+ for (Use *U : UsesToRewrite) {
+ auto *UserI = cast<Instruction>(U->getUser());
+ if (tryRewriteMaskedLoadUser(*UserI, C, Count) ||
+ tryRewriteMaskedStoreUser(*UserI, C, Count) ||
+ tryRewriteExtractElement(*UserI, C, Count))
+ continue;
+
+ if (!WideMask) {
+ BasicBlock::iterator InsertPt = OldMask->getIterator();
+ if (isa<PHINode>(OldMask))
+ InsertPt = OldMask->getParent()->getFirstNonPHIIt();
+
+ Builder.SetInsertPoint(InsertPt);
+ Builder.SetCurrentDebugLocation(OldMask->getDebugLoc());
+ WideMask = buildWideMask(Builder, C, Count);
+ }
+
+ U->set(WideMask);
+ }
+ };
+
+ // For the start/next mask, we only replace non-phi in-loop users. Due to CSE,
+ // these masks could be used by other loops, and replacing them with
+ // predicate-as-counter masks could prevent matching those loops.
+ auto IsNonPhiUseInLoop = [&](Use &U) {
+ auto *UseInst = cast<Instruction>(U.getUser());
+ return !isa<PHINode>(UseInst) && L.contains(UseInst);
+ };
+
+ RewriteUses(C.MaskPhi, NewPhi);
+ RewriteUses(C.StartMask, NewStart, IsNonPhiUseInLoop);
+ RewriteUses(C.NextMask, NewNext, IsNonPhiUseInLoop);
+
+ RecursivelyDeleteTriviallyDeadInstructions(C.MaskPhi);
+ return true;
+}
diff --git a/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp b/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
index bad3ce59bd4a1..883fd68786289 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
@@ -230,6 +230,12 @@ static cl::opt<bool> EnableSVEShuffleOpt(
"instructions like tbl or the bottom/top variants"),
cl::init(true), cl::Hidden);
+static cl::opt<bool> EnablePredicateAsCounterLoopRewrites(
+ "aarch64-enable-predicate-as-counter-loop-rewrites",
+ cl::desc("Enable rewriting loops with wide loop-carried masks to use "
+ "predicate-as-counter"),
+ cl::init(false), cl::Hidden);
+
extern "C" LLVM_ABI LLVM_EXTERNAL_VISIBILITY void
LLVMInitializeAArch64Target() {
// Register the target.
@@ -257,6 +263,7 @@ LLVMInitializeAArch64Target() {
initializeAArch64PTrueCoalescingLegacyPass(PR);
initializeAArch64SIMDInstrOptLegacyPass(PR);
initializeAArch64O0PreLegalizerCombinerLegacyPass(PR);
+ initializeAArch64PredicateAsCounterLoopRewritesPass(PR);
initializeAArch64PreLegalizerCombinerLegacyPass(PR);
initializeAArch64PointerAuthLegacyPass(PR);
initializeAArch64PostCoalescerLegacyPass(PR);
@@ -649,6 +656,9 @@ void AArch64PassConfig::addIRPasses() {
// ourselves.
addPass(createAtomicExpandLegacyPass());
+ if (EnablePredicateAsCounterLoopRewrites)
+ addPass(createAArch64PredicateAsCounterLoopRewritesPass());
+
// Cmpxchg instructions are often used with a subsequent comparison to
// determine whether it succeeded. We can exploit existing control-flow in
// ldrex/strex loops to simplify this, but it needs tidying up.
diff --git a/llvm/lib/Target/AArch64/CMakeLists.txt b/llvm/lib/Target/AArch64/CMakeLists.txt
index 12a2214f8e58e..5c4740baf79de 100644
--- a/llvm/lib/Target/AArch64/CMakeLists.txt
+++ b/llvm/lib/Target/AArch64/CMakeLists.txt
@@ -79,6 +79,7 @@ add_llvm_target(AArch64CodeGen
AArch64PromoteConstant.cpp
AArch64PTrueCoalescing.cpp
AArch64PBQPRegAlloc.cpp
+ AArch64PredicateAsCounterLoopRewrites.cpp
AArch64RegisterInfo.cpp
AArch64SMEAttributes.cpp
AArch64SLSHardening.cpp
diff --git a/llvm/test/CodeGen/AArch64/predicate-as-counter-loop-rewrites.ll b/llvm/test/CodeGen/AArch64/predicate-as-counter-loop-rewrites.ll
new file mode 100644
index 0000000000000..f887048eb8ad5
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/predicate-as-counter-loop-rewrites.ll
@@ -0,0 +1,566 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc -mtriple=aarch64-linux-gnu -aarch64-enable-predicate-as-counter-loop-rewrites -mattr=+sve2p1 -aarch64-enable-subreg-liveness-tracking < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+; Tests the baseline i64 masked load/store rewrite to predicate-as-counter with vlx4.
+define void @rewrite_masked_load_store_i64_vlx4(ptr %x, i64 %n) #0 {
+; CHECK-LABEL: rewrite_masked_load_store_i64_vlx4:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: cnth x8
+; CHECK-NEXT: whilelo pn8.d, xzr, x1, vlx4
+; CHECK-NEXT: .LBB0_1: // %loop
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: ld1d { z0.d - z3.d }, pn8/z, [x0]
+; CHECK-NEXT: add z3.d, z3.d, #1 // =0x1
+; CHECK-NEXT: add z2.d, z2.d, #1 // =0x1
+; CHECK-NEXT: add z1.d, z1.d, #1 // =0x1
+; CHECK-NEXT: add z0.d, z0.d, #1 // =0x1
+; CHECK-NEXT: st1d { z0.d - z3.d }, pn8, [x0]
+; CHECK-NEXT: whilelo pn8.d, x8, x1, vlx4
+; CHECK-NEXT: inch x8
+; CHECK-NEXT: incb x0, all, mul #4
+; CHECK-NEXT: b.mi .LBB0_1
+; CHECK-NEXT: // %bb.2: // %exit
+; CHECK-NEXT: ret
+entry:
+ br label %preheader
+
+preheader:
+ %start.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 0, i64 %n)
+ %vscale = call i64 @llvm.vscale.i64()
+ %step = shl i64 %vscale, 3
+ br label %loop
+
+loop:
+ %base = phi i64 [ 0, %preheader ], [ %next.base, %loop ]
+ %update.mask = phi <vscale x 8 x i1> [ %start.mask, %preheader ], [ %new.mask, %loop ]
+ %ptr = getelementptr i64, ptr %x, i64 %base
+ %v = call <vscale x 8 x i64> @llvm.masked.load.nxv8i64(ptr %ptr, i32 8, <vscale x 8 x i1> %update.mask, <vscale x 8 x i64> poison)
+ %add = add <vscale x 8 x i64> %v, splat (i64 1)
+ call void @llvm.masked.store.nxv8i64(<vscale x 8 x i64> %add, ptr %ptr, i32 8, <vscale x 8 x i1> %update.mask)
+ %next.base = add i64 %base, %step
+ %new.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %next.base, i64 %n)
+ %more = extractelement <vscale x 8 x i1> %new.mask, i64 0
+ br i1 %more, label %loop, label %exit
+
+exit:
+ ret void
+}
+
+; Tests rewriting a loop whose header has multiple outside predecessors before LoopSimplify creates a preheader.
+define void @rewrite_masked_load_store_i64_multi_pred_header(ptr %x, i64 %n, i1 %c) #0 {
+; CHECK-LABEL: rewrite_masked_load_store_i64_multi_pred_header:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: cnth x8
+; CHECK-NEXT: whilelo pn8.d, xzr, x1, vlx4
+; CHECK-NEXT: .LBB1_1: // %loop
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: ld1d { z0.d - z3.d }, pn8/z, [x0]
+; CHECK-NEXT: add z3.d, z3.d, #1 // =0x1
+; CHECK-NEXT: add z2.d, z2.d, #1 // =0x1
+; CHECK-NEXT: add z1.d, z1.d, #1 // =0x1
+; CHECK-NEXT: add z0.d, z0.d, #1 // =0x1
+; CHECK-NEXT: st1d { z0.d - z3.d }, pn8, [x0]
+; CHECK-NEXT: whilelo pn8.d, x8, x1, vlx4
+; CHECK-NEXT: inch x8
+; CHECK-NEXT: incb x0, all, mul #4
+; CHECK-NEXT: b.mi .LBB1_1
+; CHECK-NEXT: // %bb.2: // %exit
+; CHECK-NEXT: ret
+entry:
+ %start.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 0, i64 %n)
+ %vscale = call i64 @llvm.vscale.i64()
+ %step = shl i64 %vscale, 3
+ br i1 %c, label %left, label %right
+
+left:
+ br label %loop
+
+right:
+ br label %loop
+
+loop:
+ %base = phi i64 [ 0, %left ], [ 0, %right ], [ %next.base, %loop ]
+ %update.mask = phi <vscale x 8 x i1> [ %start.mask, %left ], [ %start.mask, %right ], [ %new.mask, %loop ]
+ %ptr = getelementptr i64, ptr %x, i64 %base
+ %v = call <vscale x 8 x i64> @llvm.masked.load.nxv8i64(ptr %ptr, i32 8, <vscale x 8 x i1> %update.mask, <vscale x 8 x i64> poison)
+ %add = add <vscale x 8 x i64> %v, splat (i64 1)
+ call void @llvm.masked.store.nxv8i64(<vscale x 8 x i64> %add, ptr %ptr, i32 8, <vscale x 8 x i1> %update.mask)
+ %next.base = add i64 %base, %step
+ %new.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %next.base, i64 %n)
+ %more = extractelement <vscale x 8 x i1> %new.mask, i64 0
+ br i1 %more, label %loop, label %exit
+
+exit:
+ ret void
+}
+
+; Tests that a shared start mask remains available so a later loop can also be rewritten.
+define void @shared_start_mask_between_loops(ptr %x, ptr %y, i64 %n) #0 {
+; CHECK-LABEL: shared_start_mask_between_loops:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: whilelo pn8.d, xzr, x2, vlx4
+; CHECK-NEXT: cnth x8
+; CHECK-NEXT: mov p9.b, p8.b
+; CHECK-NEXT: mov x9, x8
+; CHECK-NEXT: .LBB2_1: // %loop1
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: ld1d { z0.d - z3.d }, pn9/z, [x0]
+; CHECK-NEXT: add z3.d, z3.d, #1 // =0x1
+; CHECK-NEXT: add z2.d, z2.d, #1 // =0x1
+; CHECK-NEXT: add z1.d, z1.d, #1 // =0x1
+; CHECK-NEXT: add z0.d, z0.d, #1 // =0x1
+; CHECK-NEXT: st1d { z0.d - z3.d }, pn9, [x0]
+; CHECK-NEXT: whilelo pn9.d, x9, x2, vlx4
+; CHECK-NEXT: inch x9
+; CHECK-NEXT: incb x0, all, mul #4
+; CHECK-NEXT: b.mi .LBB2_1
+; CHECK-NEXT: .LBB2_2: // %loop2
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: ld1d { z0.d - z3.d }, pn8/z, [x1]
+; CHECK-NEXT: st1d { z0.d - z3.d }, pn8, [x1]
+; CHECK-NEXT: whilelo pn8.d, x8, x2, vlx4
+; CHECK-NEXT: inch x8
+; CHECK-NEXT: incb x1, all, mul #4
+; CHECK-NEXT: b.mi .LBB2_2
+; CHECK-NEXT: // %bb.3: // %exit
+; CHECK-NEXT: ret
+entry:
+ %start.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 0, i64 %n)
+ %vscale = call i64 @llvm.vscale.i64()
+ %step = shl i64 %vscale, 3
+ br label %loop1
+
+loop1:
+ %base1 = phi i64 [ 0, %entry ], [ %next.base1, %loop1 ]
+ %mask1 = phi <vscale x 8 x i1> [ %start.mask, %entry ], [ %new.mask1, %loop1 ]
+ %ptr1 = getelementptr i64, ptr %x, i64 %base1
+ %v1 = call <vscale x 8 x i64> @llvm.masked.load.nxv8i64(ptr %ptr1, i32 8, <vscale x 8 x i1> %mask1, <vscale x 8 x i64> poison)
+ %add1 = add <vscale x 8 x i64> %v1, splat (i64 1)
+ call void @llvm.masked.store.nxv8i64(<vscale x 8 x i64> %add1, ptr %ptr1, i32 8, <vscale x 8 x i1> %mask1)
+ %next.base1 = add i64 %base1, %step
+ %new.mask1 = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %next.base1, i64 %n)
+ %more1 = extractelement <vscale x 8 x i1> %new.mask1, i64 0
+ br i1 %more1, label %loop1, label %between
+
+between:
+ br label %loop2
+
+loop2:
+ %base2 = phi i64 [ 0, %between ], [ %next.base2, %loop2 ]
+ %mask2 = phi <vscale x 8 x i1> [ %start.mask, %between ], [ %new.mask2, %loop2 ]
+ %ptr2 = getelementptr ptr, ptr %y, i64 %base2
+ %v2 = call <vscale x 8 x ptr> @llvm.masked.load.nxv8p0.p0(ptr %ptr2, i32 8, <vscale x 8 x i1> %mask2, <vscale x 8 x ptr> poison)
+ call void @llvm.masked.store.nxv8p0.p0(<vscale x 8 x ptr> %v2, ptr %ptr2, i32 8, <vscale x 8 x i1> %mask2)
+ %next.base2 = add i64 %base2, %step
+ %new.mask2 = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %next.base2, i64 %n)
+ %more2 = extractelement <vscale x 8 x i1> %new.mask2, i64 0
+ br i1 %more2, label %loop2, label %exit
+
+exit:
+ ret void
+}
+
+; Tests the i8 masked load/store rewrite when the induction operands are i32 and must be extended.
+define void @rewrite_masked_load_store_i8_i32_induction(ptr %x, i32 %n) #0 {
+; CHECK-LABEL: rewrite_masked_load_store_i8_i32_induction:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: mov w8, w1
+; CHECK-NEXT: rdvl x9, #2
+; CHECK-NEXT: whilelo pn8.b, xzr, x8, vlx2
+; CHECK-NEXT: .LBB3_1: // %loop
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: mov x10, x9
+; CHECK-NEXT: mov w11, w9
+; CHECK-NEXT: incb x9, all, mul #2
+; CHECK-NEXT: decb x10, all, mul #2
+; CHECK-NEXT: sxtw x10, w10
+; CHECK-NEXT: ld1b { z0.b, z1.b }, pn8/z, [x0, x10]
+; CHECK-NEXT: add x10, x0, x10
+; CHECK-NEXT: add z1.b, z1.b, #1 // =0x1
+; CHECK-NEXT: add z0.b, z0.b, #1 // =0x1
+; CHECK-NEXT: st1b { z0.b, z1.b }, pn8, [x10]
+; CHECK-NEXT: whilelo pn8.b, x11, x8, vlx2
+; CHECK-NEXT: b.mi .LBB3_1
+; CHECK-NEXT: // %bb.2: // %exit
+; CHECK-NEXT: ret
+entry:
+ br label %preheader
+
+preheader:
+ %start.mask = call <vscale x 32 x i1> @llvm.get.active.lane.mask.nxv32i1.i32(i32 0, i32 %n)
+ %vscale = call i32 @llvm.vscale.i32()
+ %step = shl i32 %vscale, 5
+ br label %loop
+
+loop:
+ %base = phi i32 [ 0, %preheader ], [ %next.base, %loop ]
+ %update.mask = phi <vscale x 32 x i1> [ %start.mask, %preheader ], [ %new.mask, %loop ]
+ %ptr = getelementptr i8, ptr %x, i32 %base
+ %v = call <vscale x 32 x i8> @llvm.masked.load.nxv32i8(ptr %ptr, i32 1, <vscale x 32 x i1> %update.mask, <vscale x 32 x i8> poison)
+ %add = add <vscale x 32 x i8> %v, splat (i8 1)
+ call void @llvm.masked.store.nxv32i8(<vscale x 32 x i8> %add, ptr %ptr, i32 1, <vscale x 32 x i1> %update.mask)
+ %next.base = add i32 %base, %step
+ %new.mask = call <vscale x 32 x i1> @llvm.get.active.lane.mask.nxv32i1.i32(i32 %next.base, i32 %n)
+ %more = extractelement <vscale x 32 x i1> %new.mask, i64 0
+ br i1 %more, label %loop, label %exit
+
+exit:
+ ret void
+}
+
+; Tests the i32 masked load/store rewrite to predicate-as-counter with vlx2.
+define void @rewrite_masked_load_store_i32_vlx2(ptr %x, i64 %n) #0 {
+; CHECK-LABEL: rewrite_masked_load_store_i32_vlx2:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: cnth x8
+; CHECK-NEXT: whilelo pn8.s, xzr, x1, vlx2
+; CHECK-NEXT: .LBB4_1: // %loop
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: ld1w { z0.s, z1.s }, pn8/z, [x0]
+; CHECK-NEXT: add z1.s, z1.s, #1 // =0x1
+; CHECK-NEXT: add z0.s, z0.s, #1 // =0x1
+; CHECK-NEXT: st1w { z0.s, z1.s }, pn8, [x0]
+; CHECK-NEXT: whilelo pn8.s, x8, x1, vlx2
+; CHECK-NEXT: inch x8
+; CHECK-NEXT: incb x0, all, mul #2
+; CHECK-NEXT: b.mi .LBB4_1
+; CHECK-NEXT: // %bb.2: // %exit
+; CHECK-NEXT: ret
+entry:
+ br label %preheader
+
+preheader:
+ %start.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 0, i64 %n)
+ %vscale = call i64 @llvm.vscale.i64()
+ %step = shl i64 %vscale, 3
+ br label %loop
+
+loop:
+ %base = phi i64 [ 0, %preheader ], [ %next.base, %loop ]
+ %update.mask = phi <vscale x 8 x i1> [ %start.mask, %preheader ], [ %new.mask, %loop ]
+ %ptr = getelementptr i32, ptr %x, i64 %base
+ %v = call <vscale x 8 x i32> @llvm.masked.load.nxv8i32(ptr %ptr, i32 4, <vscale x 8 x i1> %update.mask, <vscale x 8 x i32> poison)
+ %add = add <vscale x 8 x i32> %v, splat (i32 1)
+ call void @llvm.masked.store.nxv8i32(<vscale x 8 x i32> %add, ptr %ptr, i32 4, <vscale x 8 x i1> %update.mask)
+ %next.base = add i64 %base, %step
+ %new.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %next.base, i64 %n)
+ %more = extractelement <vscale x 8 x i1> %new.mask, i64 0
+ br i1 %more, label %loop, label %exit
+
+exit:
+ ret void
+}
+
+; Tests choosing an i64 predicate-as-counter rewrite while materializing a wide mask for mismatched i32 users.
+define void @mixed_data_vector_types(ptr %x, ptr %y, i64 %n) #0 {
+; CHECK-LABEL: mixed_data_vector_types:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: whilelo pn8.d, xzr, x2, vlx4
+; CHECK-NEXT: cnth x8
+; CHECK-NEXT: .LBB5_1: // %loop
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: pext { p0.d, p1.d }, pn8[0]
+; CHECK-NEXT: pext { p2.d, p3.d }, pn8[1]
+; CHECK-NEXT: ld1d { z0.d - z3.d }, pn8/z, [x0]
+; CHECK-NEXT: uzp1 p0.s, p0.s, p1.s
+; CHECK-NEXT: whilelo pn8.d, x8, x2, vlx4
+; CHECK-NEXT: inch x8
+; CHECK-NEXT: uzp1 p1.s, p2.s, p3.s
+; CHECK-NEXT: incb x0, all, mul #4
+; CHECK-NEXT: ld1w { z5.s }, p0/z, [x1]
+; CHECK-NEXT: ld1w { z4.s }, p1/z, [x1, #1, mul vl]
+; CHECK-NEXT: incb x1, all, mul #2
+; CHECK-NEXT: // fake_use: $z0
+; CHECK-NEXT: // fake_use: $z1
+; CHECK-NEXT: // fake_use: $z2
+; CHECK-NEXT: // fake_use: $z3
+; CHECK-NEXT: // fake_use: $z5
+; CHECK-NEXT: // fake_use: $z4
+; CHECK-NEXT: b.mi .LBB5_1
+; CHECK-NEXT: // %bb.2: // %exit
+; CHECK-NEXT: ret
+entry:
+ br label %preheader
+
+preheader:
+ %start.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 0, i64 %n)
+ %vscale = call i64 @llvm.vscale.i64()
+ %step = shl i64 %vscale, 3
+ br label %loop
+
+loop:
+ %base = phi i64 [ 0, %preheader ], [ %next.base, %loop ]
+ %update.mask = phi <vscale x 8 x i1> [ %start.mask, %preheader ], [ %new.mask, %loop ]
+ %ptr64 = getelementptr i64, ptr %x, i64 %base
+ %ptr32 = getelementptr i32, ptr %y, i64 %base
+ %v64 = call <vscale x 8 x i64> @llvm.masked.load.nxv8i64(ptr %ptr64, i32 8, <vscale x 8 x i1> %update.mask, <vscale x 8 x i64> poison)
+ %v32 = call <vscale x 8 x i32> @llvm.masked.load.nxv8i32(ptr %ptr32, i32 4, <vscale x 8 x i1> %update.mask, <vscale x 8 x i32> poison)
+ call void (...) @llvm.fake.use(<vscale x 8 x i64> %v64)
+ call void (...) @llvm.fake.use(<vscale x 8 x i32> %v32)
+ %next.base = add i64 %base, %step
+ %new.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %next.base, i64 %n)
+ %more = extractelement <vscale x 8 x i1> %new.mask, i64 0
+ br i1 %more, label %loop, label %exit
+
+exit:
+ ret void
+}
+
+; Tests preferring the most common masked memory access size when selecting the predicate-as-counter element size.
+define void @prefer_more_common_masked_access_size(ptr %x16, ptr %y32, i64 %n) #0 {
+; CHECK-LABEL: prefer_more_common_masked_access_size:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v0.2d, #0000000000000000
+; CHECK-NEXT: whilelo pn8.h, xzr, x2, vlx2
+; CHECK-NEXT: rdvl x8, #1
+; CHECK-NEXT: mov z1.d, z0.d
+; CHECK-NEXT: .LBB6_1: // %loop
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: pext { p0.h, p1.h }, pn8[0]
+; CHECK-NEXT: st1h { z0.h, z1.h }, pn8, [x0]
+; CHECK-NEXT: st1h { z0.h, z1.h }, pn8, [x0]
+; CHECK-NEXT: punpkhi p2.h, p1.b
+; CHECK-NEXT: whilelo pn8.h, x8, x2, vlx2
+; CHECK-NEXT: incb x8
+; CHECK-NEXT: punpklo p1.h, p1.b
+; CHECK-NEXT: incb x0, all, mul #2
+; CHECK-NEXT: punpkhi p3.h, p0.b
+; CHECK-NEXT: ld1w { z2.s }, p2/z, [x1, #3, mul vl]
+; CHECK-NEXT: punpklo p0.h, p0.b
+; CHECK-NEXT: ld1w { z3.s }, p1/z, [x1, #2, mul vl]
+; CHECK-NEXT: ld1w { z4.s }, p3/z, [x1, #1, mul vl]
+; CHECK-NEXT: ld1w { z5.s }, p0/z, [x1]
+; CHECK-NEXT: incb x1, all, mul #4
+; CHECK-NEXT: // fake_use: $z5
+; CHECK-NEXT: // fake_use: $z4
+; CHECK-NEXT: // fake_use: $z3
+; CHECK-NEXT: // fake_use: $z2
+; CHECK-NEXT: b.mi .LBB6_1
+; CHECK-NEXT: // %bb.2: // %exit
+; CHECK-NEXT: ret
+entry:
+ br label %preheader
+
+preheader:
+ %start.mask = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 0, i64 %n)
+ %vscale = call i64 @llvm.vscale.i64()
+ %step = shl i64 %vscale, 4
+ br label %loop
+
+loop:
+ %base = phi i64 [ 0, %preheader ], [ %next.base, %loop ]
+ %update.mask = phi <vscale x 16 x i1> [ %start.mask, %preheader ], [ %new.mask, %loop ]
+ %ptr16 = getelementptr i16, ptr %x16, i64 %base
+ %ptr32 = getelementptr i32, ptr %y32, i64 %base
+ call void @llvm.masked.store.nxv16i16(<vscale x 16 x i16> zeroinitializer, ptr %ptr16, i32 2, <vscale x 16 x i1> %update.mask)
+ call void @llvm.masked.store.nxv16i16(<vscale x 16 x i16> zeroinitializer, ptr %ptr16, i32 2, <vscale x 16 x i1> %update.mask)
+ %v32 = call <vscale x 16 x i32> @llvm.masked.load.nxv16i32(ptr %ptr32, i32 4, <vscale x 16 x i1> %update.mask, <vscale x 16 x i32> poison)
+ call void (...) @llvm.fake.use(<vscale x 16 x i32> %v32)
+ %next.base = add i64 %base, %step
+ %new.mask = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 %next.base, i64 %n)
+ %more = extractelement <vscale x 16 x i1> %new.mask, i64 0
+ br i1 %more, label %loop, label %exit
+
+exit:
+ ret void
+}
+
+; Tests breaking equal access-size counts in favor of the larger access size.
+define void @prefer_larger_access_size_on_tie(ptr %x16, ptr %y32, i64 %n) #0 {
+; CHECK-LABEL: prefer_larger_access_size_on_tie:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v0.2d, #0000000000000000
+; CHECK-NEXT: whilelo pn8.s, xzr, x2, vlx4
+; CHECK-NEXT: rdvl x8, #1
+; CHECK-NEXT: .LBB7_1: // %loop
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: pext { p0.s, p1.s }, pn8[1]
+; CHECK-NEXT: pext { p2.s, p3.s }, pn8[0]
+; CHECK-NEXT: uzp1 p0.h, p0.h, p1.h
+; CHECK-NEXT: uzp1 p1.h, p2.h, p3.h
+; CHECK-NEXT: st1h { z0.h }, p0, [x0, #1, mul vl]
+; CHECK-NEXT: st1h { z0.h }, p1, [x0]
+; CHECK-NEXT: incb x0, all, mul #2
+; CHECK-NEXT: ld1w { z4.s - z7.s }, pn8/z, [x1]
+; CHECK-NEXT: whilelo pn8.s, x8, x2, vlx4
+; CHECK-NEXT: incb x8
+; CHECK-NEXT: incb x1, all, mul #4
+; CHECK-NEXT: // fake_use: $z4
+; CHECK-NEXT: // fake_use: $z5
+; CHECK-NEXT: // fake_use: $z6
+; CHECK-NEXT: // fake_use: $z7
+; CHECK-NEXT: b.mi .LBB7_1
+; CHECK-NEXT: // %bb.2: // %exit
+; CHECK-NEXT: ret
+entry:
+ br label %preheader
+
+preheader:
+ %start.mask = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 0, i64 %n)
+ %vscale = call i64 @llvm.vscale.i64()
+ %step = shl i64 %vscale, 4
+ br label %loop
+
+loop:
+ %base = phi i64 [ 0, %preheader ], [ %next.base, %loop ]
+ %update.mask = phi <vscale x 16 x i1> [ %start.mask, %preheader ], [ %new.mask, %loop ]
+ %ptr16 = getelementptr i16, ptr %x16, i64 %base
+ %ptr32 = getelementptr i32, ptr %y32, i64 %base
+ call void @llvm.masked.store.nxv16i16(<vscale x 16 x i16> zeroinitializer, ptr %ptr16, i32 2, <vscale x 16 x i1> %update.mask)
+ %v32 = call <vscale x 16 x i32> @llvm.masked.load.nxv16i32(ptr %ptr32, i32 4, <vscale x 16 x i1> %update.mask, <vscale x 16 x i32> poison)
+ call void (...) @llvm.fake.use(<vscale x 16 x i32> %v32)
+ %next.base = add i64 %base, %step
+ %new.mask = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 %next.base, i64 %n)
+ %more = extractelement <vscale x 16 x i1> %new.mask, i64 0
+ br i1 %more, label %loop, label %exit
+
+exit:
+ ret void
+}
+
+; Tests rewriting extractelement when the extracted lane is within the first legal predicate section.
+define void @rewrite_extractelement_within_first_section_i32(ptr %x, i64 %n) #0 {
+; CHECK-LABEL: rewrite_extractelement_within_first_section_i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: cnth x8
+; CHECK-NEXT: whilelo pn8.s, xzr, x1, vlx2
+; CHECK-NEXT: .LBB8_1: // %loop
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: ld1w { z0.s, z1.s }, pn8/z, [x0]
+; CHECK-NEXT: whilelo pn8.s, x8, x1, vlx2
+; CHECK-NEXT: inch x8
+; CHECK-NEXT: pext p0.s, pn8[0]
+; CHECK-NEXT: incb x0, all, mul #2
+; CHECK-NEXT: mov z2.s, p0/z, #1 // =0x1
+; CHECK-NEXT: mov w9, v2.s[3]
+; CHECK-NEXT: // fake_use: $z0
+; CHECK-NEXT: // fake_use: $z1
+; CHECK-NEXT: tbnz w9, #0, .LBB8_1
+; CHECK-NEXT: // %bb.2: // %exit
+; CHECK-NEXT: ret
+entry:
+ br label %preheader
+
+preheader:
+ %start.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 0, i64 %n)
+ %vscale = call i64 @llvm.vscale.i64()
+ %step = shl i64 %vscale, 3
+ br label %loop
+
+loop:
+ %base = phi i64 [ 0, %preheader ], [ %next.base, %loop ]
+ %update.mask = phi <vscale x 8 x i1> [ %start.mask, %preheader ], [ %new.mask, %loop ]
+ %ptr = getelementptr i32, ptr %x, i64 %base
+ %v = call <vscale x 8 x i32> @llvm.masked.load.nxv8i32(ptr %ptr, i32 4, <vscale x 8 x i1> %update.mask, <vscale x 8 x i32> poison)
+ call void (...) @llvm.fake.use(<vscale x 8 x i32> %v)
+ %next.base = add i64 %base, %step
+ %new.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %next.base, i64 %n)
+ %more = extractelement <vscale x 8 x i1> %new.mask, i64 3
+ br i1 %more, label %loop, label %exit
+
+exit:
+ ret void
+}
+
+; Negative test: masked loads with a non-poison passthru value are not rewritten to pn loads.
+define void @masked_load_passthru_non_poison(ptr %x, i64 %n) #0 {
+; CHECK-LABEL: masked_load_passthru_non_poison:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: cnth x8
+; CHECK-NEXT: whilelo pn8.d, xzr, x1, vlx4
+; CHECK-NEXT: .LBB9_1: // %loop
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: pext { p0.d, p1.d }, pn8[1]
+; CHECK-NEXT: ld1d { z0.d }, p1/z, [x0, #3, mul vl]
+; CHECK-NEXT: pext { p1.d, p2.d }, pn8[0]
+; CHECK-NEXT: ld1d { z1.d }, p0/z, [x0, #2, mul vl]
+; CHECK-NEXT: ld1d { z2.d }, p2/z, [x0, #1, mul vl]
+; CHECK-NEXT: ld1d { z3.d }, p1/z, [x0]
+; CHECK-NEXT: movprfx z7, z0
+; CHECK-NEXT: add z7.d, z7.d, #1 // =0x1
+; CHECK-NEXT: movprfx z6, z1
+; CHECK-NEXT: add z6.d, z6.d, #1 // =0x1
+; CHECK-NEXT: movprfx z5, z2
+; CHECK-NEXT: add z5.d, z5.d, #1 // =0x1
+; CHECK-NEXT: movprfx z4, z3
+; CHECK-NEXT: add z4.d, z4.d, #1 // =0x1
+; CHECK-NEXT: st1d { z4.d - z7.d }, pn8, [x0]
+; CHECK-NEXT: whilelo pn8.d, x8, x1, vlx4
+; CHECK-NEXT: inch x8
+; CHECK-NEXT: incb x0, all, mul #4
+; CHECK-NEXT: b.mi .LBB9_1
+; CHECK-NEXT: // %bb.2: // %exit
+; CHECK-NEXT: ret
+entry:
+ br label %preheader
+
+preheader:
+ %start.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 0, i64 %n)
+ %vscale = call i64 @llvm.vscale.i64()
+ %step = shl i64 %vscale, 3
+ br label %loop
+
+loop:
+ %base = phi i64 [ 0, %preheader ], [ %next.base, %loop ]
+ %update.mask = phi <vscale x 8 x i1> [ %start.mask, %preheader ], [ %new.mask, %loop ]
+ %ptr = getelementptr i64, ptr %x, i64 %base
+ %v = call <vscale x 8 x i64> @llvm.masked.load.nxv8i64(ptr %ptr, i32 8, <vscale x 8 x i1> %update.mask, <vscale x 8 x i64> zeroinitializer)
+ %add = add <vscale x 8 x i64> %v, splat (i64 1)
+ call void @llvm.masked.store.nxv8i64(<vscale x 8 x i64> %add, ptr %ptr, i32 8, <vscale x 8 x i1> %update.mask)
+ %next.base = add i64 %base, %step
+ %new.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %next.base, i64 %n)
+ %more = extractelement <vscale x 8 x i1> %new.mask, i64 0
+ br i1 %more, label %loop, label %exit
+
+exit:
+ ret void
+}
+
+; Negative test: extractelement is not rewritten when the extracted lane is outside the first predicate section.
+define void @extractelement_outside_first_section_i32(ptr %x, i64 %n) #0 {
+; CHECK-LABEL: extractelement_outside_first_section_i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: cnth x8
+; CHECK-NEXT: whilelo pn8.s, xzr, x1, vlx2
+; CHECK-NEXT: .LBB10_1: // %loop
+; CHECK-NEXT: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT: ld1w { z0.s, z1.s }, pn8/z, [x0]
+; CHECK-NEXT: whilelo pn8.s, x8, x1, vlx2
+; CHECK-NEXT: inch x8
+; CHECK-NEXT: pext { p0.s, p1.s }, pn8[0]
+; CHECK-NEXT: incb x0, all, mul #2
+; CHECK-NEXT: uzp1 p0.h, p0.h, p1.h
+; CHECK-NEXT: mov z2.h, p0/z, #1 // =0x1
+; CHECK-NEXT: umov w9, v2.h[6]
+; CHECK-NEXT: // fake_use: $z0
+; CHECK-NEXT: // fake_use: $z1
+; CHECK-NEXT: tbnz w9, #0, .LBB10_1
+; CHECK-NEXT: // %bb.2: // %exit
+; CHECK-NEXT: ret
+entry:
+ br label %preheader
+
+preheader:
+ %start.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 0, i64 %n)
+ %vscale = call i64 @llvm.vscale.i64()
+ %step = shl i64 %vscale, 3
+ br label %loop
+
+loop:
+ %base = phi i64 [ 0, %preheader ], [ %next.base, %loop ]
+ %update.mask = phi <vscale x 8 x i1> [ %start.mask, %preheader ], [ %new.mask, %loop ]
+ %ptr = getelementptr i32, ptr %x, i64 %base
+ %v = call <vscale x 8 x i32> @llvm.masked.load.nxv8i32(ptr %ptr, i32 4, <vscale x 8 x i1> %update.mask, <vscale x 8 x i32> poison)
+ call void (...) @llvm.fake.use(<vscale x 8 x i32> %v)
+ %next.base = add i64 %base, %step
+ %new.mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %next.base, i64 %n)
+ %more = extractelement <vscale x 8 x i1> %new.mask, i64 6
+ br i1 %more, label %loop, label %exit
+
+exit:
+ ret void
+}
+
+attributes #0 = { "target-features"="+sve2p1" }
More information about the llvm-commits
mailing list