1//===- AArch64PredicateAsCounterLoopRewrites.cpp --------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM
4// Exceptions.
5// See https://llvm.org/LICENSE.txt for license information.
6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7//
8//===----------------------------------------------------------------------===//
9//
10// Rewrites IR for loop-carried wide masks that can be represented as
11// predicate-as-counter values. This applies when the mask is used by load/store
12// operations that can be mapped to multi-vector instructions (with +sve2p1).
13//
14// For example, a loop like:
15//
16// entry:
17// %step = vscale x 64
18// %mask.entry = @get.active.lane.mask(0, %n)
19//
20// loop:
21// %iv = phi i64 [0, entry], [%iv.next, loop]
22// %mask = phi <vscale x 64 x i1> [%mask.entry, entry],
23// [%mask.next, loop]
24//
25// %load = load <vscale x 64 x i8> %src[%iv], %mask
26// store <vscale x 64 x i8> %load, %dst[%iv], %mask
27//
28// %iv.next = %iv + %step
29// %mask.next = @get.active.lane.mask(%iv.next, %n)
30// br first.active(%mask.next), loop, exit
31//
32// Could be rewritten to:
33//
34// entry:
35// %step = vscale x 64
36// %mask.entry = @whilelo.c8(0, %n, VLx4)
37//
38// loop:
39// %iv = phi i64 [0, entry], [%iv.next, loop]
40// %mask = phi target("aarch64.svcount") [%mask.entry, entry],
41// [%mask.next, loop]
42//
43// %load = @ld1.pn.x4 <4 x <vscale x 16 x i8>> %src[%iv], %mask
44// @st1.pn.x4 <4 x <vscale x 16 x i8>> %load, %dest[%iv], %mask
45//
46// %iv.next = %iv + %step
47// %mask.next = @whilelo.c8(%iv.next, %n, VLx4)
48// br first.active(@pext(%mask.next, 0)), loop, exit
49//
50// This replaces the `get.active.lane.mask` intrinsics with AArch64
51// predicate-as-counter `whilelo` intrinsics and updates the mask phi to use the
52// `aarch64.svcount` target type. Within the loop, load/store users are mapped
53// to multi-vector load/store intrinsics where possible. Users that cannot be
54// mapped to multi-vector instructions materialize vector masks using the `pext`
55// intrinsic (which extracts vector predicates from a predicate-as-counter).
56//
57// This pass may be a temporary solution that is removed if we gain support
58// for target-specific VPlan transforms in the loop vectorizer.
59//
60//===----------------------------------------------------------------------===//
61
62#include "AArch64.h"
63#include "AArch64Subtarget.h"
64#include "AArch64TargetMachine.h"
65#include "llvm/ADT/DenseMap.h"
66#include "llvm/ADT/SmallVector.h"
67#include "llvm/ADT/Statistic.h"
68#include "llvm/Analysis/LoopInfo.h"
69#include "llvm/Analysis/LoopPass.h"
70#include "llvm/CodeGen/TargetPassConfig.h"
71#include "llvm/IR/Attributes.h"
72#include "llvm/IR/DataLayout.h"
73#include "llvm/IR/IRBuilder.h"
74#include "llvm/IR/IntrinsicInst.h"
75#include "llvm/IR/Intrinsics.h"
76#include "llvm/IR/IntrinsicsAArch64.h"
77#include "llvm/IR/Module.h"
78#include "llvm/InitializePasses.h"
79#include "llvm/Pass.h"
80#include "llvm/Support/Debug.h"
81#include "llvm/Transforms/Utils.h"
82#include "llvm/Transforms/Utils/Local.h"
83#include <optional>
84
85using namespace llvm;
86
87#define DEBUG_TYPE "aarch64-pn-loop-rewrites"
88namespace {
89
90STATISTIC(LoopsRewritten, "Number of loops rewritten");
91
92struct MaskRewriteCandidate {
93 /// The mask phi node (used by masked operations within the loop).
94 PHINode *MaskPhi = nullptr;
95 /// The initial value for the mask (incoming value from the preheader).
96 IntrinsicInst *StartMask = nullptr;
97 /// The updated value for the mask (incoming value from the loop latch).
98 IntrinsicInst *NextMask = nullptr;
99 /// The multi-vector scale for the predicate-as-counter (2 or 4).
100 unsigned VectorScale = 0;
101 /// The element size (in bits) for the predicate-as-counter.
102 unsigned ElementSizeInBits = 0;
103};
104
105static void logLoopBailout(const Loop &L, const Twine &Reason) {
106 LLVM_DEBUG({
107 dbgs() << "PN loop rewrite: skipping loop with header ";
108 L.getHeader()->printAsOperand(dbgs(), /*PrintType=*/false);
109 dbgs() << ": " << Reason << "\n";
110 });
111}
112
113static void logMatchFailure(const PHINode &Phi, const Twine &Reason) {
114 LLVM_DEBUG({
115 dbgs() << "PN loop rewrite: failed to match mask phi ";
116 Phi.printAsOperand(dbgs(), /*PrintType=*/false);
117 dbgs() << ": " << Reason << "\n";
118 });
119}
120
121/// Returns the scalar size in bits for a type according to the data layout.
122static unsigned getScalarSizeInBits(const DataLayout &DL, Type *Ty) {
123 return DL.getTypeSizeInBits(Ty: Ty->getScalarType()).getFixedValue();
124}
125
126/// Returns the SVE element count for \p ElementSizeInBits.
127static ElementCount getSVEElementCount(unsigned ElementSizeInBits) {
128 return ElementCount::getScalable(MinVal: AArch64::SVEBitsPerBlock /
129 ElementSizeInBits);
130}
131
132/// Returns the largest scalar access size of masked loads/stores in loop \p L
133/// where \p MaskPhi is used as the mask. TODO: This heuristic may need
134/// refinement based on the frequency of different access sizes.
135static unsigned getLargestMaskedMemAccessSizeInBits(const Loop &L,
136 PHINode &MaskPhi) {
137 const DataLayout &DL = MaskPhi.getModule()->getDataLayout();
138 unsigned LargestAccessSizeInBits = 0;
139
140 for (User *U : MaskPhi.users()) {
141 auto *II = dyn_cast<IntrinsicInst>(Val: U);
142 if (!II || !L.contains(Inst: II))
143 continue;
144
145 Intrinsic::ID IID = II->getIntrinsicID();
146 if (IID != Intrinsic::masked_load && IID != Intrinsic::masked_store)
147 continue;
148
149 unsigned MaskOpIdx = IID == Intrinsic::masked_load ? 1 : 2;
150 if (II->getArgOperand(i: MaskOpIdx) != &MaskPhi)
151 continue;
152
153 unsigned AccessSizeInBits = getScalarSizeInBits(DL, Ty: II->getAccessType());
154 if (AccessSizeInBits > LargestAccessSizeInBits)
155 LargestAccessSizeInBits = AccessSizeInBits;
156 }
157
158 return LargestAccessSizeInBits;
159}
160
161/// Returns the predicate-as-counter whilelo intrinsic ID for
162/// \p ElementSizeInBits.
163static Intrinsic::ID getWhileLOIntrinsic(unsigned ElementSizeInBits) {
164 switch (ElementSizeInBits) {
165 case 8:
166 return Intrinsic::aarch64_sve_whilelo_c8;
167 case 16:
168 return Intrinsic::aarch64_sve_whilelo_c16;
169 case 32:
170 return Intrinsic::aarch64_sve_whilelo_c32;
171 case 64:
172 return Intrinsic::aarch64_sve_whilelo_c64;
173 default:
174 llvm_unreachable("unsupported predicate-as-counter element size");
175 }
176}
177
178/// Expands the predicate-as-counter mask \p Count into a wide vector mask.
179/// Masks are extracted from the counter using the paired pext intrinsics,
180/// then concatenated to form the wide mask value.
181static Value *buildWideMask(IRBuilder<> &Builder, const MaskRewriteCandidate &C,
182 Value *Count) {
183 ElementCount LegalEC = getSVEElementCount(ElementSizeInBits: C.ElementSizeInBits);
184 Type *LegalMaskTy = VectorType::get(ElementType: Builder.getInt1Ty(), EC: LegalEC);
185
186 Value *WideMask = PoisonValue::get(T: C.MaskPhi->getType());
187 for (unsigned PairOffset = 0; PairOffset != C.VectorScale / 2; ++PairOffset) {
188 auto *Pair =
189 Builder.CreateIntrinsic(ID: Intrinsic::aarch64_sve_pext_x2, OverloadTypes: {LegalMaskTy},
190 Args: {Count, Builder.getInt32(C: PairOffset)},
191 /*FMFSource=*/{}, Name: "pn.pext.pair");
192 for (unsigned SliceInPair = 0; SliceInPair != 2; ++SliceInPair) {
193 Value *Part = Builder.CreateExtractValue(Agg: Pair, Idxs: SliceInPair, Name: "pn.pext");
194 unsigned Slice = PairOffset * 2 + SliceInPair;
195 WideMask = Builder.CreateInsertVector(
196 DstType: C.MaskPhi->getType(), SrcVec: WideMask, SubVec: Part,
197 Idx: Slice * LegalEC.getKnownMinValue(), Name: "pn.mask");
198 }
199 }
200
201 return WideMask;
202}
203
204/// Creates a predicate-as-counter whilelo for \p ElementSizeInBits between
205/// \p Start and \p End with multi-vector \p VectorScale (2 or 4).
206static Value *createWhileLO(IRBuilder<> &Builder, unsigned ElementSizeInBits,
207 Value *Start, Value *End, unsigned VectorScale) {
208 if (Start->getType()->getIntegerBitWidth() < 64) {
209 Start = Builder.CreateZExt(V: Start, DestTy: Builder.getInt64Ty());
210 End = Builder.CreateZExt(V: End, DestTy: Builder.getInt64Ty());
211 }
212
213 Intrinsic::ID WhileLO = getWhileLOIntrinsic(ElementSizeInBits);
214 return Builder.CreateIntrinsic(ID: WhileLO,
215 Args: {Start, End, Builder.getInt32(C: VectorScale)},
216 /*FMFSource=*/{}, Name: "pn.mask");
217}
218
219/// If \p UserI is a masked-load user of the original loop mask, rewrite it to
220/// a masked multi-vector load (using the predicate-as-counter) if the load
221/// access size matches the predicate-as-counter element size.
222static bool tryRewriteMaskedLoad(Instruction &UserI,
223 const MaskRewriteCandidate &C, Value *Count) {
224 auto *II = dyn_cast<IntrinsicInst>(Val: &UserI);
225 if (!II || II->getIntrinsicID() != Intrinsic::masked_load)
226 return false;
227
228 unsigned ElementSizeInBits =
229 getScalarSizeInBits(DL: II->getModule()->getDataLayout(), Ty: II->getType());
230 if (ElementSizeInBits != C.ElementSizeInBits)
231 return false;
232
233 if (!isa<PoisonValue, UndefValue>(Val: II->getArgOperand(i: 2)))
234 return false;
235
236 IRBuilder<> Builder(II);
237
238 ElementCount LegalEC = getSVEElementCount(ElementSizeInBits: C.ElementSizeInBits);
239 Type *ScalarType = II->getType()->getScalarType();
240 auto *WideDataTy = VectorType::get(ElementType: ScalarType, EC: LegalEC * C.VectorScale);
241 auto *LegalDataTy = VectorType::get(ElementType: ScalarType, EC: LegalEC);
242
243 auto *PNLoad = cast<IntrinsicInst>(Val: Builder.CreateIntrinsic(
244 ID: C.VectorScale == 4 ? Intrinsic::aarch64_sve_ld1_pn_x4
245 : Intrinsic::aarch64_sve_ld1_pn_x2,
246 OverloadTypes: {LegalDataTy, II->getArgOperand(i: 0)->getType()},
247 Args: {Count, II->getArgOperand(i: 0)},
248 /*FMFSource=*/{}, Name: "pn.ld1"));
249
250 // Copy pointer parameter attributes.
251 for (Attribute ParamAttr : II->getParamAttributes(ArgNo: 0))
252 PNLoad->addParamAttr(ArgNo: 1, Attr: ParamAttr);
253
254 PNLoad->copyMetadata(SrcInst: *II);
255
256 // Concatenate all results into the original wide value.
257 Value *WideData = PoisonValue::get(T: WideDataTy);
258 for (unsigned Slice = 0; Slice != C.VectorScale; ++Slice) {
259 Value *Part = Builder.CreateExtractValue(Agg: PNLoad, Idxs: Slice, Name: "pn.data");
260 WideData = Builder.CreateInsertVector(DstType: WideDataTy, SrcVec: WideData, SubVec: Part,
261 Idx: Slice * LegalEC.getKnownMinValue(),
262 Name: "pn.vec");
263 }
264
265 II->replaceAllUsesWith(V: WideData);
266 II->eraseFromParent();
267 return true;
268}
269
270/// If \p UserI is a masked-store user of the original loop mask, rewrite it to
271/// a masked multi-vector store (using the predicate-as-counter) if the store
272/// access size matches the predicate-as-counter element size.
273static bool tryRewriteMaskedStore(Instruction &UserI,
274 const MaskRewriteCandidate &C, Value *Count) {
275 auto *II = dyn_cast<IntrinsicInst>(Val: &UserI);
276 if (!II || II->getIntrinsicID() != Intrinsic::masked_store)
277 return false;
278
279 unsigned ElementSizeInBits = getScalarSizeInBits(
280 DL: II->getModule()->getDataLayout(), Ty: II->getArgOperand(i: 0)->getType());
281 if (ElementSizeInBits != C.ElementSizeInBits)
282 return false;
283
284 IRBuilder<> Builder(II);
285
286 Type *ScalarType = II->getArgOperand(i: 0)->getType()->getScalarType();
287 ElementCount LegalEC = getSVEElementCount(ElementSizeInBits: C.ElementSizeInBits);
288 auto *LegalDataTy = VectorType::get(ElementType: ScalarType, EC: LegalEC);
289
290 SmallVector<Value *, 6> StoreArgs;
291 for (unsigned Slice = 0; Slice != C.VectorScale; ++Slice)
292 StoreArgs.push_back(Elt: Builder.CreateExtractVector(
293 DstType: LegalDataTy, SrcVec: II->getArgOperand(i: 0), Idx: Slice * LegalEC.getKnownMinValue(),
294 Name: "pn.data"));
295 StoreArgs.push_back(Elt: Count);
296 StoreArgs.push_back(Elt: II->getArgOperand(i: 1));
297
298 auto *PNStore = cast<IntrinsicInst>(Val: Builder.CreateIntrinsic(
299 ID: C.VectorScale == 4 ? Intrinsic::aarch64_sve_st1_pn_x4
300 : Intrinsic::aarch64_sve_st1_pn_x2,
301 OverloadTypes: {LegalDataTy, II->getArgOperand(i: 1)->getType()}, Args: StoreArgs));
302
303 PNStore->copyMetadata(SrcInst: *II);
304
305 // Copy pointer parameter attributes.
306 for (Attribute ParamAttr : II->getParamAttributes(ArgNo: 1))
307 PNStore->addParamAttr(ArgNo: C.VectorScale + 1, Attr: ParamAttr);
308
309 II->eraseFromParent();
310 return true;
311}
312
313/// If \p UserI is an extractelement use of the original loop mask, attempt to
314/// rewrite it to `extractelement(pext(count, 0), idx)` if the extract index is
315/// known to be within the first mask section (which means pext index = 0). If
316/// the user of the extract is a branch and the source of the mask is a whilelo,
317/// this form can be optimized into checking the status flags.
318static bool tryRewriteExtractElement(Instruction &UserI,
319 const MaskRewriteCandidate &C,
320 Value *Count) {
321 auto *EEI = dyn_cast<ExtractElementInst>(Val: &UserI);
322 if (!EEI)
323 return false;
324
325 ElementCount LegalEC = getSVEElementCount(ElementSizeInBits: C.ElementSizeInBits);
326 auto *Idx = dyn_cast<ConstantInt>(Val: EEI->getIndexOperand());
327 if (!Idx || Idx->getValue().uge(RHS: LegalEC.getKnownMinValue()))
328 return false;
329
330 IRBuilder<> Builder(EEI);
331
332 auto *ExtractMask = Builder.CreateIntrinsic(
333 ID: Intrinsic::aarch64_sve_pext,
334 OverloadTypes: {VectorType::get(ElementType: Builder.getInt1Ty(), EC: LegalEC)},
335 Args: {Count, Builder.getInt32(C: 0)}, /*FMFSource=*/{}, Name: "pn.pext");
336
337 Value *Extracted = Builder.CreateExtractElement(
338 Vec: ExtractMask, Idx: EEI->getIndexOperand(), Name: EEI->getName() + ".pn");
339
340 EEI->replaceAllUsesWith(V: Extracted);
341 EEI->eraseFromParent();
342 return true;
343}
344
345class AArch64PredicateAsCounterLoopRewrites : public LoopPass {
346public:
347 static char ID;
348
349 AArch64PredicateAsCounterLoopRewrites() : LoopPass(ID) {}
350
351 void getAnalysisUsage(AnalysisUsage &AU) const override {
352 AU.addRequired<TargetPassConfig>();
353 // Require loop simplify to ensure loops have a preheader.
354 AU.addRequiredID(ID&: LoopSimplifyID);
355 AU.addPreservedID(ID&: LoopSimplifyID);
356 AU.setPreservesCFG();
357 }
358
359 bool runOnLoop(Loop *L, LPPassManager &) override;
360
361private:
362 std::optional<MaskRewriteCandidate> matchMaskPhi(Loop &L, PHINode &Phi) const;
363 bool rewriteCandidate(const MaskRewriteCandidate &C, Loop &L) const;
364};
365
366} // end anonymous namespace
367
368char AArch64PredicateAsCounterLoopRewrites::ID = 0;
369
370INITIALIZE_PASS_BEGIN(AArch64PredicateAsCounterLoopRewrites, DEBUG_TYPE,
371 "AArch64 Predicate As Counter Loop Rewrites", false,
372 false)
373INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
374INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
375INITIALIZE_PASS_END(AArch64PredicateAsCounterLoopRewrites, DEBUG_TYPE,
376 "AArch64 Predicate As Counter Loop Rewrites", false, false)
377
378Pass *llvm::createAArch64PredicateAsCounterLoopRewritesPass() {
379 return new AArch64PredicateAsCounterLoopRewrites();
380}
381
382bool AArch64PredicateAsCounterLoopRewrites::runOnLoop(Loop *L,
383 LPPassManager &) {
384 if (skipLoop(L)) {
385 logLoopBailout(L: *L, Reason: "skipLoop requested the loop to be skipped");
386 return false;
387 }
388
389 Function &F = *L->getHeader()->getParent();
390 auto &TPC = getAnalysis<TargetPassConfig>();
391 const AArch64Subtarget *ST =
392 TPC.getTM<AArch64TargetMachine>().getSubtargetImpl(F);
393 if (!ST->hasSVE2p1() && !(ST->hasSME2() && ST->isStreaming())) {
394 logLoopBailout(L: *L, Reason: "neither SVE2.1 nor SME2 is available");
395 return false;
396 }
397
398 if (!L->getLoopPreheader()) {
399 logLoopBailout(L: *L, Reason: "loop has no preheader");
400 return false;
401 }
402 if (!L->getLoopLatch()) {
403 logLoopBailout(L: *L, Reason: "loop has no latch");
404 return false;
405 }
406
407 bool Changed = false;
408 BasicBlock *Header = L->getHeader();
409 for (PHINode &Phi : make_early_inc_range(Range: Header->phis())) {
410 if (std::optional<MaskRewriteCandidate> Candidate = matchMaskPhi(L&: *L, Phi))
411 Changed |= rewriteCandidate(C: *Candidate, L&: *L);
412 }
413
414 if (Changed)
415 ++LoopsRewritten;
416
417 return Changed;
418}
419
420static IntrinsicInst *getGetActiveLaneMask(Value *V) {
421 auto *II = dyn_cast<IntrinsicInst>(Val: V);
422 return II && II->getIntrinsicID() == Intrinsic::get_active_lane_mask
423 ? II
424 : nullptr;
425}
426
427std::optional<MaskRewriteCandidate>
428AArch64PredicateAsCounterLoopRewrites::matchMaskPhi(Loop &L,
429 PHINode &Phi) const {
430 auto *PhiTy = dyn_cast<ScalableVectorType>(Val: Phi.getType());
431 if (!PhiTy || !PhiTy->getElementType()->isIntegerTy(BitWidth: 1))
432 return std::nullopt;
433
434 if (Phi.getNumIncomingValues() != 2) {
435 logMatchFailure(Phi, Reason: Twine("phi has ") + Twine(Phi.getNumIncomingValues()) +
436 " incoming values; expected 2");
437 return std::nullopt;
438 }
439
440 Value *StartValue = Phi.getIncomingValueForBlock(BB: L.getLoopPreheader());
441 Value *NextValue = Phi.getIncomingValueForBlock(BB: L.getLoopLatch());
442 IntrinsicInst *StartMask = getGetActiveLaneMask(V: StartValue);
443 IntrinsicInst *NextMask = getGetActiveLaneMask(V: NextValue);
444 if (!StartMask) {
445 logMatchFailure(Phi,
446 Reason: "preheader incoming value is not get_active_lane_mask");
447 return std::nullopt;
448 }
449 if (!NextMask) {
450 logMatchFailure(Phi, Reason: "latch incoming value is not get_active_lane_mask");
451 return std::nullopt;
452 }
453
454 unsigned WideMaskElements = PhiTy->getMinNumElements();
455 if (!isPowerOf2_32(Value: WideMaskElements)) {
456 logMatchFailure(Phi,
457 Reason: Twine("wide mask element count is not a power of 2: ") +
458 Twine(WideMaskElements));
459 return std::nullopt;
460 }
461
462 if (StartMask->getArgOperand(i: 0)->getType()->getIntegerBitWidth() > 64) {
463 logMatchFailure(Phi, Reason: "start mask induction operand is wider than i64");
464 return std::nullopt;
465 }
466 if (NextMask->getArgOperand(i: 0)->getType()->getIntegerBitWidth() > 64) {
467 logMatchFailure(Phi, Reason: "next mask induction operand is wider than i64");
468 return std::nullopt;
469 }
470
471 unsigned PreferredMaskElementSizeInBits =
472 getLargestMaskedMemAccessSizeInBits(L, MaskPhi&: Phi);
473
474 if (!is_contained(Set: {8u, 16u, 32u, 64u}, Element: PreferredMaskElementSizeInBits)) {
475 logMatchFailure(Phi, Reason: Twine("unsupported element size in bits: ") +
476 Twine(PreferredMaskElementSizeInBits));
477 return std::nullopt;
478 }
479
480 unsigned SVEMaskElements =
481 AArch64::SVEBitsPerBlock / PreferredMaskElementSizeInBits;
482 if (WideMaskElements <= SVEMaskElements) {
483 logMatchFailure(Phi, Reason: Twine("wide mask element count ") +
484 Twine(WideMaskElements) +
485 " is not wider than the legal mask width " +
486 Twine(SVEMaskElements));
487 return std::nullopt;
488 }
489
490 unsigned VectorScale = WideMaskElements / SVEMaskElements;
491 if (VectorScale != 2 && VectorScale != 4) {
492 logMatchFailure(Phi, Reason: Twine("unsupported predicate-as-counter scale: ") +
493 Twine(VectorScale));
494 return std::nullopt;
495 }
496
497 return MaskRewriteCandidate{.MaskPhi: &Phi, .StartMask: StartMask, .NextMask: NextMask, .VectorScale: VectorScale,
498 .ElementSizeInBits: PreferredMaskElementSizeInBits};
499}
500
501bool AArch64PredicateAsCounterLoopRewrites::rewriteCandidate(
502 const MaskRewriteCandidate &C, Loop &L) const {
503 IRBuilder<> Builder(C.StartMask);
504 Value *NewStart =
505 createWhileLO(Builder, ElementSizeInBits: C.ElementSizeInBits, Start: C.StartMask->getArgOperand(i: 0),
506 End: C.StartMask->getArgOperand(i: 1), VectorScale: C.VectorScale);
507 Builder.SetInsertPoint(C.NextMask);
508 Value *NewNext =
509 createWhileLO(Builder, ElementSizeInBits: C.ElementSizeInBits, Start: C.NextMask->getArgOperand(i: 0),
510 End: C.NextMask->getArgOperand(i: 1), VectorScale: C.VectorScale);
511
512 Builder.SetInsertPoint(C.MaskPhi);
513 auto *NewPhi =
514 Builder.CreatePHI(Ty: NewStart->getType(), NumReservedValues: 2, Name: C.MaskPhi->getName() + ".pn");
515 NewPhi->addIncoming(V: NewStart, BB: L.getLoopPreheader());
516 NewPhi->addIncoming(V: NewNext, BB: L.getLoopLatch());
517
518 auto RewriteUses = [&](Instruction *OldMask, Value *Count,
519 function_ref<bool(Use &U)> Predicate = nullptr) {
520 SmallVector<Use *, 8> UsesToRewrite;
521 for (Use &U : OldMask->uses()) {
522 if (!Predicate || Predicate(U))
523 UsesToRewrite.push_back(Elt: &U);
524 }
525
526 Value *WideMask = nullptr;
527 for (Use *U : UsesToRewrite) {
528 auto *UserI = cast<Instruction>(Val: U->getUser());
529 if (tryRewriteMaskedLoad(UserI&: *UserI, C, Count) ||
530 tryRewriteMaskedStore(UserI&: *UserI, C, Count) ||
531 tryRewriteExtractElement(UserI&: *UserI, C, Count))
532 continue;
533
534 if (!WideMask) {
535 BasicBlock::iterator InsertPt = OldMask->getIterator();
536 if (isa<PHINode>(Val: OldMask))
537 InsertPt = OldMask->getParent()->getFirstNonPHIIt();
538
539 Builder.SetInsertPoint(InsertPt);
540 Builder.SetCurrentDebugLocation(OldMask->getDebugLoc());
541 WideMask = buildWideMask(Builder, C, Count);
542 }
543
544 U->set(WideMask);
545 }
546 };
547
548 // For the start/next mask, we only replace non-phi in-loop users. Due to CSE,
549 // these masks could be used by other loops, and replacing them with
550 // predicate-as-counter masks could prevent matching those loops.
551 auto IsNonPhiUseInLoop = [&](Use &U) {
552 auto *UseInst = cast<Instruction>(Val: U.getUser());
553 return !isa<PHINode>(Val: UseInst) && L.contains(Inst: UseInst);
554 };
555
556 RewriteUses(C.MaskPhi, NewPhi);
557 RewriteUses(C.StartMask, NewStart, IsNonPhiUseInLoop);
558 RewriteUses(C.NextMask, NewNext, IsNonPhiUseInLoop);
559
560 RecursivelyDeleteTriviallyDeadInstructions(V: C.MaskPhi);
561 return true;
562}
563