1//===- MVETailPredication.cpp - MVE Tail Predication ------------*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9/// \file
10/// Armv8.1m introduced MVE, M-Profile Vector Extension, and low-overhead
11/// branches to help accelerate DSP applications. These two extensions,
12/// combined with a new form of predication called tail-predication, can be used
13/// to provide implicit vector predication within a low-overhead loop.
14/// This is implicit because the predicate of active/inactive lanes is
15/// calculated by hardware, and thus does not need to be explicitly passed
16/// to vector instructions. The instructions responsible for this are the
17/// DLSTP and WLSTP instructions, which setup a tail-predicated loop and the
18/// the total number of data elements processed by the loop. The loop-end
19/// LETP instruction is responsible for decrementing and setting the remaining
20/// elements to be processed and generating the mask of active lanes.
21///
22/// The HardwareLoops pass inserts intrinsics identifying loops that the
23/// backend will attempt to convert into a low-overhead loop. The vectorizer is
24/// responsible for generating a vectorized loop in which the lanes are
25/// predicated upon an get.active.lane.mask intrinsic. This pass looks at these
26/// get.active.lane.mask intrinsic and attempts to convert them to VCTP
27/// instructions. This will be picked up by the ARM Low-overhead loop pass later
28/// in the backend, which performs the final transformation to a DLSTP or WLSTP
29/// tail-predicated loop.
30//
31//===----------------------------------------------------------------------===//
32
33#include "ARM.h"
34#include "ARMSubtarget.h"
35#include "ARMTargetTransformInfo.h"
36#include "llvm/Analysis/LoopInfo.h"
37#include "llvm/Analysis/LoopPass.h"
38#include "llvm/Analysis/ScalarEvolution.h"
39#include "llvm/Analysis/ScalarEvolutionExpressions.h"
40#include "llvm/Analysis/TargetLibraryInfo.h"
41#include "llvm/Analysis/TargetTransformInfo.h"
42#include "llvm/Analysis/ValueTracking.h"
43#include "llvm/CodeGen/TargetPassConfig.h"
44#include "llvm/IR/IRBuilder.h"
45#include "llvm/IR/Instructions.h"
46#include "llvm/IR/IntrinsicsARM.h"
47#include "llvm/Support/Debug.h"
48#include "llvm/Transforms/Utils/BasicBlockUtils.h"
49#include "llvm/Transforms/Utils/Local.h"
50#include "llvm/Transforms/Utils/LoopUtils.h"
51#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
52
53using namespace llvm;
54
55#define DEBUG_TYPE "mve-tail-predication"
56#define DESC "Transform predicated vector loops to use MVE tail predication"
57
58cl::opt<TailPredication::Mode> EnableTailPredication(
59 "tail-predication", cl::desc("MVE tail-predication pass options"),
60 cl::init(Val: TailPredication::Enabled),
61 cl::values(clEnumValN(TailPredication::Disabled, "disabled",
62 "Don't tail-predicate loops"),
63 clEnumValN(TailPredication::EnabledNoReductions,
64 "enabled-no-reductions",
65 "Enable tail-predication, but not for reduction loops"),
66 clEnumValN(TailPredication::Enabled,
67 "enabled",
68 "Enable tail-predication, including reduction loops"),
69 clEnumValN(TailPredication::ForceEnabledNoReductions,
70 "force-enabled-no-reductions",
71 "Enable tail-predication, but not for reduction loops, "
72 "and force this which might be unsafe"),
73 clEnumValN(TailPredication::ForceEnabled,
74 "force-enabled",
75 "Enable tail-predication, including reduction loops, "
76 "and force this which might be unsafe")));
77
78
79namespace {
80
81class MVETailPredication : public LoopPass {
82 SmallVector<IntrinsicInst*, 4> MaskedInsts;
83 Loop *L = nullptr;
84 ScalarEvolution *SE = nullptr;
85 TargetTransformInfo *TTI = nullptr;
86 const ARMSubtarget *ST = nullptr;
87
88public:
89 static char ID;
90
91 MVETailPredication() : LoopPass(ID) { }
92
93 void getAnalysisUsage(AnalysisUsage &AU) const override {
94 AU.addRequired<ScalarEvolutionWrapperPass>();
95 AU.addRequired<TargetPassConfig>();
96 AU.addRequired<TargetTransformInfoWrapperPass>();
97 AU.setPreservesCFG();
98 }
99
100 bool runOnLoop(Loop *L, LPPassManager&) override;
101
102private:
103 /// Perform the relevant checks on the loop and convert active lane masks if
104 /// possible.
105 bool TryConvertActiveLaneMask(Value *TripCount);
106
107 /// Perform several checks on the arguments of @llvm.get.active.lane.mask
108 /// intrinsic. E.g., check that the loop induction variable and the element
109 /// count are of the form we expect, and also perform overflow checks for
110 /// the new expressions that are created.
111 const SCEV *IsSafeActiveMask(IntrinsicInst *ActiveLaneMask, Value *TripCount);
112
113 /// Insert the intrinsic to represent the effect of tail predication.
114 void InsertVCTPIntrinsic(IntrinsicInst *ActiveLaneMask, Value *Start);
115};
116
117} // end namespace
118
119bool MVETailPredication::runOnLoop(Loop *L, LPPassManager&) {
120 if (skipLoop(L) || !EnableTailPredication)
121 return false;
122
123 MaskedInsts.clear();
124 Function &F = *L->getHeader()->getParent();
125 auto &TPC = getAnalysis<TargetPassConfig>();
126 auto &TM = TPC.getTM<TargetMachine>();
127 ST = &TM.getSubtarget<ARMSubtarget>(F);
128 TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
129 SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
130 this->L = L;
131
132 // The MVE and LOB extensions are combined to enable tail-predication, but
133 // there's nothing preventing us from generating VCTP instructions for v8.1m.
134 if (!ST->hasMVEIntegerOps() || !ST->hasV8_1MMainlineOps()) {
135 LLVM_DEBUG(dbgs() << "ARM TP: Not a v8.1m.main+mve target.\n");
136 return false;
137 }
138
139 BasicBlock *Preheader = L->getLoopPreheader();
140 if (!Preheader)
141 return false;
142
143 auto FindLoopIterations = [](BasicBlock *BB) -> IntrinsicInst* {
144 for (auto &I : *BB) {
145 auto *Call = dyn_cast<IntrinsicInst>(Val: &I);
146 if (!Call)
147 continue;
148
149 Intrinsic::ID ID = Call->getIntrinsicID();
150 if (ID == Intrinsic::start_loop_iterations ||
151 ID == Intrinsic::test_start_loop_iterations)
152 return cast<IntrinsicInst>(Val: &I);
153 }
154 return nullptr;
155 };
156
157 // Look for the hardware loop intrinsic that sets the iteration count.
158 IntrinsicInst *Setup = FindLoopIterations(Preheader);
159
160 // The test.set iteration could live in the pre-preheader.
161 if (!Setup) {
162 if (!Preheader->getSinglePredecessor())
163 return false;
164 Setup = FindLoopIterations(Preheader->getSinglePredecessor());
165 if (!Setup)
166 return false;
167 }
168
169 LLVM_DEBUG(dbgs() << "ARM TP: Running on Loop: " << *L << *Setup << "\n");
170
171 bool Changed = TryConvertActiveLaneMask(TripCount: Setup->getArgOperand(i: 0));
172
173 return Changed;
174}
175
176// The active lane intrinsic has this form:
177//
178// @llvm.get.active.lane.mask(IV, TC)
179//
180// Here we perform checks that this intrinsic behaves as expected,
181// which means:
182//
183// 1) Check that the TripCount (TC) belongs to this loop (originally).
184// 2) The element count (TC) needs to be sufficiently large that the decrement
185// of element counter doesn't overflow, which means that we need to prove:
186// ceil(ElementCount / VectorWidth) >= TripCount
187// by rounding up ElementCount up:
188// ((ElementCount + (VectorWidth - 1)) / VectorWidth
189// and evaluate if expression isKnownNonNegative:
190// (((ElementCount + (VectorWidth - 1)) / VectorWidth) - TripCount
191// 3) The IV must be an induction phi with an increment equal to the
192// vector width.
193const SCEV *MVETailPredication::IsSafeActiveMask(IntrinsicInst *ActiveLaneMask,
194 Value *TripCount) {
195 bool ForceTailPredication =
196 EnableTailPredication == TailPredication::ForceEnabledNoReductions ||
197 EnableTailPredication == TailPredication::ForceEnabled;
198
199 Value *ElemCount = ActiveLaneMask->getOperand(i_nocapture: 1);
200 bool Changed = false;
201 if (!L->makeLoopInvariant(V: ElemCount, Changed))
202 return nullptr;
203
204 const SCEV *EC = SE->getSCEV(V: ElemCount);
205 const SCEV *TC = SE->getSCEV(V: TripCount);
206 int VectorWidth =
207 cast<FixedVectorType>(Val: ActiveLaneMask->getType())->getNumElements();
208 if (VectorWidth != 2 && VectorWidth != 4 && VectorWidth != 8 &&
209 VectorWidth != 16)
210 return nullptr;
211 ConstantInt *ConstElemCount = nullptr;
212
213 // 1) Smoke tests that the original scalar loop TripCount (TC) belongs to
214 // this loop. The scalar tripcount corresponds the number of elements
215 // processed by the loop, so we will refer to that from this point on.
216 if (!SE->isLoopInvariant(S: EC, L)) {
217 LLVM_DEBUG(dbgs() << "ARM TP: element count must be loop invariant.\n");
218 return nullptr;
219 }
220
221 // 2) Find out if IV is an induction phi. Note that we can't use Loop
222 // helpers here to get the induction variable, because the hardware loop is
223 // no longer in loopsimplify form, and also the hwloop intrinsic uses a
224 // different counter. Using SCEV, we check that the induction is of the
225 // form i = i + 4, where the increment must be equal to the VectorWidth.
226 auto *IV = ActiveLaneMask->getOperand(i_nocapture: 0);
227 const SCEV *IVExpr = SE->getSCEV(V: IV);
228 auto *AddExpr = dyn_cast<SCEVAddRecExpr>(Val: IVExpr);
229
230 if (!AddExpr) {
231 LLVM_DEBUG(dbgs() << "ARM TP: induction not an add expr: "; IVExpr->dump());
232 return nullptr;
233 }
234 // Check that this AddRec is associated with this loop.
235 if (AddExpr->getLoop() != L) {
236 LLVM_DEBUG(dbgs() << "ARM TP: phi not part of this loop\n");
237 return nullptr;
238 }
239 auto *Step = dyn_cast<SCEVConstant>(Val: AddExpr->getOperand(i: 1));
240 if (!Step) {
241 LLVM_DEBUG(dbgs() << "ARM TP: induction step is not a constant: ";
242 AddExpr->getOperand(1)->dump());
243 return nullptr;
244 }
245 auto StepValue = Step->getValue()->getSExtValue();
246 if (VectorWidth != StepValue) {
247 LLVM_DEBUG(dbgs() << "ARM TP: Step value " << StepValue
248 << " doesn't match vector width " << VectorWidth << "\n");
249 return nullptr;
250 }
251
252 if ((ConstElemCount = dyn_cast<ConstantInt>(Val: ElemCount))) {
253 ConstantInt *TC = dyn_cast<ConstantInt>(Val: TripCount);
254 if (!TC) {
255 LLVM_DEBUG(dbgs() << "ARM TP: Constant tripcount expected in "
256 "set.loop.iterations\n");
257 return nullptr;
258 }
259
260 // Calculate 2 tripcount values and check that they are consistent with
261 // each other. The TripCount for a predicated vector loop body is
262 // ceil(ElementCount/Width), or floor((ElementCount+Width-1)/Width) as we
263 // work it out here.
264 uint64_t TC1 = TC->getZExtValue();
265 uint64_t TC2 =
266 (ConstElemCount->getZExtValue() + VectorWidth - 1) / VectorWidth;
267
268 // If the tripcount values are inconsistent, we can't insert the VCTP and
269 // trigger tail-predication; keep the intrinsic as a get.active.lane.mask
270 // and legalize this.
271 if (TC1 != TC2) {
272 LLVM_DEBUG(dbgs() << "ARM TP: inconsistent constant tripcount values: "
273 << TC1 << " from set.loop.iterations, and "
274 << TC2 << " from get.active.lane.mask\n");
275 return nullptr;
276 }
277 } else if (!ForceTailPredication) {
278 // 3) We need to prove that the sub expression that we create in the
279 // tail-predicated loop body, which calculates the remaining elements to be
280 // processed, is non-negative, i.e. it doesn't overflow:
281 //
282 // ((ElementCount + VectorWidth - 1) / VectorWidth) - TripCount >= 0
283 //
284 // This is true if:
285 //
286 // TripCount == (ElementCount + VectorWidth - 1) / VectorWidth
287 //
288 // which what we will be using here.
289 //
290 const SCEV *VW =
291 SE->getSCEV(V: ConstantInt::get(Ty: TripCount->getType(), V: VectorWidth));
292 // ElementCount + (VW-1):
293 const SCEV *Start = AddExpr->getStart();
294 const SCEV *ECPlusVWMinus1 = SE->getAddExpr(
295 LHS: EC,
296 RHS: SE->getSCEV(V: ConstantInt::get(Ty: TripCount->getType(), V: VectorWidth - 1)));
297
298 // Ceil = ElementCount + (VW-1) / VW
299 const SCEV *Ceil = SE->getUDivExpr(LHS: ECPlusVWMinus1, RHS: VW);
300
301 // Prevent unused variable warnings with TC
302 (void)TC;
303 LLVM_DEBUG({
304 dbgs() << "ARM TP: Analysing overflow behaviour for:\n";
305 dbgs() << "ARM TP: - TripCount = " << *TC << "\n";
306 dbgs() << "ARM TP: - ElemCount = " << *EC << "\n";
307 dbgs() << "ARM TP: - Start = " << *Start << "\n";
308 dbgs() << "ARM TP: - BETC = " << *SE->getBackedgeTakenCount(L) << "\n";
309 dbgs() << "ARM TP: - VecWidth = " << VectorWidth << "\n";
310 dbgs() << "ARM TP: - (ElemCount+VW-1) / VW = " << *Ceil << "\n";
311 });
312
313 // As an example, almost all the tripcount expressions (produced by the
314 // vectoriser) look like this:
315 //
316 // TC = ((-4 + (4 * ((3 + %N) /u 4))<nuw> - start) /u 4)
317 //
318 // and "ElementCount + (VW-1) / VW":
319 //
320 // Ceil = ((3 + %N) /u 4)
321 //
322 // Check for equality of TC and Ceil by calculating SCEV expression
323 // TC - Ceil and test it for zero.
324 //
325 const SCEV *Div = SE->getUDivExpr(
326 LHS: SE->getAddExpr(Op0: SE->getMulExpr(LHS: Ceil, RHS: VW), Op1: SE->getNegativeSCEV(V: VW),
327 Op2: SE->getNegativeSCEV(V: Start)),
328 RHS: VW);
329 const SCEV *Sub = SE->getMinusSCEV(LHS: SE->getBackedgeTakenCount(L), RHS: Div);
330 LLVM_DEBUG(dbgs() << "ARM TP: - Sub = "; Sub->dump());
331
332 // Use context sensitive facts about the path to the loop to refine. This
333 // comes up as the backedge taken count can incorporate context sensitive
334 // reasoning, and our RHS just above doesn't.
335 Sub = SE->applyLoopGuards(Expr: Sub, L);
336 LLVM_DEBUG(dbgs() << "ARM TP: - (Guarded) = "; Sub->dump());
337
338 if (!Sub->isZero()) {
339 LLVM_DEBUG(dbgs() << "ARM TP: possible overflow in sub expression.\n");
340 return nullptr;
341 }
342 }
343
344 // Check that the start value is a multiple of the VectorWidth.
345 // TODO: This could do with a method to check if the scev is a multiple of
346 // VectorWidth. For the moment we just check for constants, muls and unknowns
347 // (which use MaskedValueIsZero and seems to be the most common).
348 if (auto *BaseC = dyn_cast<SCEVConstant>(Val: AddExpr->getStart())) {
349 if (BaseC->getAPInt().urem(RHS: VectorWidth) == 0)
350 return SE->getMinusSCEV(LHS: EC, RHS: BaseC);
351 } else if (auto *BaseV = dyn_cast<SCEVUnknown>(Val: AddExpr->getStart())) {
352 Type *Ty = BaseV->getType();
353 APInt Mask = APInt::getLowBitsSet(numBits: Ty->getPrimitiveSizeInBits(),
354 loBitsSet: Log2_64(Value: VectorWidth));
355 if (MaskedValueIsZero(V: BaseV->getValue(), Mask,
356 SQ: L->getHeader()->getDataLayout()))
357 return SE->getMinusSCEV(LHS: EC, RHS: BaseV);
358 } else if (auto *BaseMul = dyn_cast<SCEVMulExpr>(Val: AddExpr->getStart())) {
359 if (auto *BaseC = dyn_cast<SCEVConstant>(Val: BaseMul->getOperand(i: 0)))
360 if (BaseC->getAPInt().urem(RHS: VectorWidth) == 0)
361 return SE->getMinusSCEV(LHS: EC, RHS: BaseC);
362 if (auto *BaseC = dyn_cast<SCEVConstant>(Val: BaseMul->getOperand(i: 1)))
363 if (BaseC->getAPInt().urem(RHS: VectorWidth) == 0)
364 return SE->getMinusSCEV(LHS: EC, RHS: BaseC);
365 }
366
367 LLVM_DEBUG(
368 dbgs() << "ARM TP: induction base is not know to be a multiple of VF: "
369 << *AddExpr->getOperand(0) << "\n");
370 return nullptr;
371}
372
373void MVETailPredication::InsertVCTPIntrinsic(IntrinsicInst *ActiveLaneMask,
374 Value *Start) {
375 IRBuilder<> Builder(L->getLoopPreheader()->getTerminator());
376 Module *M = L->getHeader()->getModule();
377 Type *Ty = IntegerType::get(C&: M->getContext(), NumBits: 32);
378 unsigned VectorWidth =
379 cast<FixedVectorType>(Val: ActiveLaneMask->getType())->getNumElements();
380
381 // Insert a phi to count the number of elements processed by the loop.
382 Builder.SetInsertPoint(TheBB: L->getHeader(), IP: L->getHeader()->getFirstNonPHIIt());
383 PHINode *Processed = Builder.CreatePHI(Ty, NumReservedValues: 2);
384 Processed->addIncoming(V: Start, BB: L->getLoopPreheader());
385
386 // Replace @llvm.get.active.mask() with the ARM specific VCTP intrinsic, and
387 // thus represent the effect of tail predication.
388 Builder.SetInsertPoint(ActiveLaneMask);
389 ConstantInt *Factor = ConstantInt::get(Ty: cast<IntegerType>(Val: Ty), V: VectorWidth);
390
391 Intrinsic::ID VCTPID;
392 switch (VectorWidth) {
393 default:
394 llvm_unreachable("unexpected number of lanes");
395 case 2: VCTPID = Intrinsic::arm_mve_vctp64; break;
396 case 4: VCTPID = Intrinsic::arm_mve_vctp32; break;
397 case 8: VCTPID = Intrinsic::arm_mve_vctp16; break;
398 case 16: VCTPID = Intrinsic::arm_mve_vctp8; break;
399 }
400 Value *VCTPCall = Builder.CreateIntrinsic(ID: VCTPID, Args: Processed);
401 ActiveLaneMask->replaceAllUsesWith(V: VCTPCall);
402
403 // Add the incoming value to the new phi.
404 // TODO: This add likely already exists in the loop.
405 Value *Remaining = Builder.CreateSub(LHS: Processed, RHS: Factor);
406 Processed->addIncoming(V: Remaining, BB: L->getLoopLatch());
407 LLVM_DEBUG(dbgs() << "ARM TP: Insert processed elements phi: "
408 << *Processed << "\n"
409 << "ARM TP: Inserted VCTP: " << *VCTPCall << "\n");
410}
411
412bool MVETailPredication::TryConvertActiveLaneMask(Value *TripCount) {
413 SmallVector<IntrinsicInst *, 4> ActiveLaneMasks;
414 for (auto *BB : L->getBlocks())
415 for (auto &I : *BB)
416 if (auto *Int = dyn_cast<IntrinsicInst>(Val: &I))
417 if (Int->getIntrinsicID() == Intrinsic::get_active_lane_mask)
418 ActiveLaneMasks.push_back(Elt: Int);
419
420 if (ActiveLaneMasks.empty())
421 return false;
422
423 LLVM_DEBUG(dbgs() << "ARM TP: Found predicated vector loop.\n");
424
425 for (auto *ActiveLaneMask : ActiveLaneMasks) {
426 LLVM_DEBUG(dbgs() << "ARM TP: Found active lane mask: "
427 << *ActiveLaneMask << "\n");
428
429 const SCEV *StartSCEV = IsSafeActiveMask(ActiveLaneMask, TripCount);
430 if (!StartSCEV) {
431 LLVM_DEBUG(dbgs() << "ARM TP: Not safe to insert VCTP.\n");
432 return false;
433 }
434 LLVM_DEBUG(dbgs() << "ARM TP: Safe to insert VCTP. Start is " << *StartSCEV
435 << "\n");
436 SCEVExpander Expander(*SE, "start");
437 Instruction *Ins = L->getLoopPreheader()->getTerminator();
438 Value *Start = Expander.expandCodeFor(SH: StartSCEV, Ty: StartSCEV->getType(), I: Ins);
439 LLVM_DEBUG(dbgs() << "ARM TP: Created start value " << *Start << "\n");
440 InsertVCTPIntrinsic(ActiveLaneMask, Start);
441 }
442
443 // Remove dead instructions and now dead phis.
444 for (auto *II : ActiveLaneMasks)
445 RecursivelyDeleteTriviallyDeadInstructions(V: II);
446 for (auto *I : L->blocks())
447 DeleteDeadPHIs(BB: I);
448 return true;
449}
450
451Pass *llvm::createMVETailPredicationPass() {
452 return new MVETailPredication();
453}
454
455char MVETailPredication::ID = 0;
456
457INITIALIZE_PASS(MVETailPredication, DEBUG_TYPE, DESC, false, false)
458