1//===- LoopUnrollAndJam.cpp - Loop unroll and jam pass --------------------===//
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// This pass implements an unroll and jam pass. Most of the work is done by
10// Utils/UnrollLoopAndJam.cpp.
11//===----------------------------------------------------------------------===//
12
13#include "llvm/Transforms/Scalar/LoopUnrollAndJamPass.h"
14#include "llvm/ADT/ArrayRef.h"
15#include "llvm/ADT/PriorityWorklist.h"
16#include "llvm/ADT/SmallPtrSet.h"
17#include "llvm/ADT/StringRef.h"
18#include "llvm/Analysis/AssumptionCache.h"
19#include "llvm/Analysis/CodeMetrics.h"
20#include "llvm/Analysis/DependenceAnalysis.h"
21#include "llvm/Analysis/LoopAnalysisManager.h"
22#include "llvm/Analysis/LoopInfo.h"
23#include "llvm/Analysis/LoopNestAnalysis.h"
24#include "llvm/Analysis/LoopPass.h"
25#include "llvm/Analysis/OptimizationRemarkEmitter.h"
26#include "llvm/Analysis/ScalarEvolution.h"
27#include "llvm/Analysis/TargetTransformInfo.h"
28#include "llvm/IR/BasicBlock.h"
29#include "llvm/IR/Constants.h"
30#include "llvm/IR/Dominators.h"
31#include "llvm/IR/Function.h"
32#include "llvm/IR/Instructions.h"
33#include "llvm/IR/Metadata.h"
34#include "llvm/IR/PassManager.h"
35#include "llvm/Support/Casting.h"
36#include "llvm/Support/CommandLine.h"
37#include "llvm/Support/Compiler.h"
38#include "llvm/Support/Debug.h"
39#include "llvm/Support/raw_ostream.h"
40#include "llvm/Transforms/Scalar/LoopPassManager.h"
41#include "llvm/Transforms/Utils/LoopPeel.h"
42#include "llvm/Transforms/Utils/LoopUtils.h"
43#include "llvm/Transforms/Utils/UnrollLoop.h"
44#include <cassert>
45#include <cstdint>
46
47namespace llvm {
48class Instruction;
49class Value;
50} // namespace llvm
51
52using namespace llvm;
53
54#define DEBUG_TYPE "loop-unroll-and-jam"
55
56/// @{
57/// Metadata attribute names
58static const char *const LLVMLoopUnrollAndJamFollowupAll =
59 "llvm.loop.unroll_and_jam.followup_all";
60static const char *const LLVMLoopUnrollAndJamFollowupInner =
61 "llvm.loop.unroll_and_jam.followup_inner";
62static const char *const LLVMLoopUnrollAndJamFollowupOuter =
63 "llvm.loop.unroll_and_jam.followup_outer";
64static const char *const LLVMLoopUnrollAndJamFollowupRemainderInner =
65 "llvm.loop.unroll_and_jam.followup_remainder_inner";
66static const char *const LLVMLoopUnrollAndJamFollowupRemainderOuter =
67 "llvm.loop.unroll_and_jam.followup_remainder_outer";
68/// @}
69
70static cl::opt<bool>
71 AllowUnrollAndJam("allow-unroll-and-jam", cl::Hidden,
72 cl::desc("Allows loops to be unroll-and-jammed."));
73
74static cl::opt<unsigned> UnrollAndJamCount(
75 "unroll-and-jam-count", cl::Hidden,
76 cl::desc("Use this unroll count for all loops including those with "
77 "unroll_and_jam_count pragma values, for testing purposes"));
78
79static cl::opt<unsigned> UnrollAndJamThreshold(
80 "unroll-and-jam-threshold", cl::init(Val: 60), cl::Hidden,
81 cl::desc("Threshold to use for inner loop when doing unroll and jam."));
82
83static cl::opt<unsigned> PragmaUnrollAndJamThreshold(
84 "pragma-unroll-and-jam-threshold", cl::init(Val: 1024), cl::Hidden,
85 cl::desc("Unrolled size limit for loops with an unroll_and_jam(full) or "
86 "unroll_count pragma."));
87
88// Returns the loop hint metadata node with the given name (for example,
89// "llvm.loop.unroll.count"). If no such metadata node exists, then nullptr is
90// returned.
91static MDNode *getUnrollMetadataForLoop(const Loop *L, StringRef Name) {
92 if (MDNode *LoopID = L->getLoopID())
93 return GetUnrollMetadata(LoopID, Name);
94 return nullptr;
95}
96
97// Returns true if the loop has any metadata starting with Prefix. For example a
98// Prefix of "llvm.loop.unroll." returns true if we have any unroll metadata.
99static bool hasAnyUnrollPragma(const Loop *L, StringRef Prefix) {
100 if (MDNode *LoopID = L->getLoopID()) {
101 // First operand should refer to the loop id itself.
102 assert(LoopID->getNumOperands() > 0 && "requires at least one operand");
103 assert(LoopID->getOperand(0) == LoopID && "invalid loop id");
104
105 for (unsigned I = 1, E = LoopID->getNumOperands(); I < E; ++I) {
106 MDNode *MD = dyn_cast<MDNode>(Val: LoopID->getOperand(I));
107 if (!MD)
108 continue;
109
110 MDString *S = dyn_cast<MDString>(Val: MD->getOperand(I: 0));
111 if (!S)
112 continue;
113
114 if (S->getString().starts_with(Prefix))
115 return true;
116 }
117 }
118 return false;
119}
120
121// Returns true if the loop has an unroll_and_jam(enable) pragma.
122static bool hasUnrollAndJamEnablePragma(const Loop *L) {
123 return getUnrollMetadataForLoop(L, Name: "llvm.loop.unroll_and_jam.enable");
124}
125
126// If loop has an unroll_and_jam_count pragma return the (necessarily
127// positive) value from the pragma. Otherwise return 0.
128static unsigned unrollAndJamCountPragmaValue(const Loop *L) {
129 MDNode *MD = getUnrollMetadataForLoop(L, Name: "llvm.loop.unroll_and_jam.count");
130 if (MD) {
131 assert(MD->getNumOperands() == 2 &&
132 "Unroll count hint metadata should have two operands.");
133 unsigned Count =
134 mdconst::extract<ConstantInt>(MD: MD->getOperand(I: 1))->getZExtValue();
135 assert(Count >= 1 && "Unroll count must be positive.");
136 return Count;
137 }
138 return 0;
139}
140
141// Returns loop size estimation for unrolled loop.
142static uint64_t
143getUnrollAndJammedLoopSize(unsigned LoopSize,
144 TargetTransformInfo::UnrollingPreferences &UP) {
145 assert(LoopSize >= UP.BEInsns && "LoopSize should not be less than BEInsns!");
146 return static_cast<uint64_t>(LoopSize - UP.BEInsns) * UP.Count + UP.BEInsns;
147}
148
149// Calculates unroll and jam count and writes it to UP.Count. Returns true if
150// unroll count was set explicitly.
151static bool computeUnrollAndJamCount(
152 Loop *L, Loop *SubLoop, const TargetTransformInfo &TTI, DominatorTree &DT,
153 LoopInfo *LI, AssumptionCache *AC, ScalarEvolution &SE,
154 const SmallPtrSetImpl<const Value *> &EphValues,
155 OptimizationRemarkEmitter *ORE, unsigned OuterTripCount,
156 unsigned OuterTripMultiple, const UnrollCostEstimator &OuterUCE,
157 unsigned InnerTripCount, unsigned InnerLoopSize,
158 TargetTransformInfo::UnrollingPreferences &UP,
159 TargetTransformInfo::PeelingPreferences &PP) {
160 unsigned OuterLoopSize = OuterUCE.getRolledLoopSize();
161 // First up use computeUnrollCount from the loop unroller to get a count
162 // for unrolling the outer loop, plus any loops requiring explicit
163 // unrolling we leave to the unroller. This uses UP.Threshold /
164 // UP.PartialThreshold / UP.MaxCount to come up with sensible loop values.
165 // We have already checked that the loop has no unroll.* pragmas.
166 unsigned MaxTripCount = 0;
167 bool UseUpperBound = false;
168 bool ExplicitUnroll = computeUnrollCount(
169 L, TTI, DT, LI, AC, SE, EphValues, ORE, TripCount: OuterTripCount, MaxTripCount,
170 /*MaxOrZero*/ false, TripMultiple: OuterTripMultiple, UCE: OuterUCE, UP, PP,
171 UseUpperBound);
172 if (ExplicitUnroll || UseUpperBound) {
173 // If the user explicitly set the loop as unrolled, dont UnJ it. Leave it
174 // for the unroller instead.
175 LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; explicit count set by "
176 "computeUnrollCount\n");
177 UP.Count = 0;
178 return false;
179 }
180
181 // Override with any explicit Count from the "unroll-and-jam-count" option.
182 bool UserUnrollCount = UnrollAndJamCount.getNumOccurrences() > 0;
183 if (UserUnrollCount) {
184 UP.Count = UnrollAndJamCount;
185 UP.Force = true;
186 if (UP.AllowRemainder &&
187 getUnrollAndJammedLoopSize(LoopSize: OuterLoopSize, UP) < UP.Threshold &&
188 getUnrollAndJammedLoopSize(LoopSize: InnerLoopSize, UP) <
189 UP.UnrollAndJamInnerLoopThreshold)
190 return true;
191 }
192
193 // Check for unroll_and_jam pragmas
194 unsigned PragmaCount = unrollAndJamCountPragmaValue(L);
195 if (PragmaCount > 0) {
196 UP.Count = PragmaCount;
197 UP.Runtime = true;
198 UP.Force = true;
199 if ((UP.AllowRemainder || (OuterTripMultiple % PragmaCount == 0)) &&
200 getUnrollAndJammedLoopSize(LoopSize: OuterLoopSize, UP) < UP.Threshold &&
201 getUnrollAndJammedLoopSize(LoopSize: InnerLoopSize, UP) <
202 UP.UnrollAndJamInnerLoopThreshold)
203 return true;
204 }
205
206 bool PragmaEnableUnroll = hasUnrollAndJamEnablePragma(L);
207 bool ExplicitUnrollAndJamCount = PragmaCount > 0 || UserUnrollCount;
208 bool ExplicitUnrollAndJam = PragmaEnableUnroll || ExplicitUnrollAndJamCount;
209
210 // If the loop has an unrolling pragma, we want to be more aggressive with
211 // unrolling limits.
212 if (ExplicitUnrollAndJam)
213 UP.UnrollAndJamInnerLoopThreshold = PragmaUnrollAndJamThreshold;
214
215 if (!UP.AllowRemainder && getUnrollAndJammedLoopSize(LoopSize: InnerLoopSize, UP) >=
216 UP.UnrollAndJamInnerLoopThreshold) {
217 LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; can't create remainder and "
218 "inner loop too large\n");
219 UP.Count = 0;
220 return false;
221 }
222
223 // We have a sensible limit for the outer loop, now adjust it for the inner
224 // loop and UP.UnrollAndJamInnerLoopThreshold. If the outer limit was set
225 // explicitly, we want to stick to it.
226 if (!ExplicitUnrollAndJamCount && UP.AllowRemainder) {
227 while (UP.Count != 0 && getUnrollAndJammedLoopSize(LoopSize: InnerLoopSize, UP) >=
228 UP.UnrollAndJamInnerLoopThreshold)
229 UP.Count--;
230 }
231
232 // If we are explicitly unroll and jamming, we are done. Otherwise there are a
233 // number of extra performance heuristics to check.
234 if (ExplicitUnrollAndJam)
235 return true;
236
237 // If the inner loop count is known and small, leave the entire loop nest to
238 // be the unroller
239 if (InnerTripCount && InnerLoopSize * InnerTripCount < UP.Threshold) {
240 LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; small inner loop count is "
241 "being left for the unroller\n");
242 UP.Count = 0;
243 return false;
244 }
245
246 // Check for situations where UnJ is likely to be unprofitable. Including
247 // subloops with more than 1 block.
248 if (SubLoop->getBlocks().size() != 1) {
249 LLVM_DEBUG(
250 dbgs() << "Won't unroll-and-jam; More than one inner loop block\n");
251 UP.Count = 0;
252 return false;
253 }
254
255 // Limit to loops where there is something to gain from unrolling and
256 // jamming the loop. In this case, look for loads that are invariant in the
257 // outer loop and can become shared.
258 unsigned NumInvariant = 0;
259 for (BasicBlock *BB : SubLoop->getBlocks()) {
260 for (Instruction &I : *BB) {
261 if (auto *Ld = dyn_cast<LoadInst>(Val: &I)) {
262 Value *V = Ld->getPointerOperand();
263 const SCEV *LSCEV = SE.getSCEVAtScope(V, L);
264 if (SE.isLoopInvariant(S: LSCEV, L))
265 NumInvariant++;
266 }
267 }
268 }
269 if (NumInvariant == 0) {
270 LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; No loop invariant loads\n");
271 UP.Count = 0;
272 return false;
273 }
274
275 return false;
276}
277
278static LoopUnrollResult
279tryToUnrollAndJamLoop(Loop *L, DominatorTree &DT, LoopInfo *LI,
280 ScalarEvolution &SE, const TargetTransformInfo &TTI,
281 AssumptionCache &AC, DependenceInfo &DI,
282 OptimizationRemarkEmitter &ORE, int OptLevel) {
283 TargetTransformInfo::UnrollingPreferences UP = gatherUnrollingPreferences(
284 L, SE, TTI, BFI: nullptr, PSI: nullptr, ORE, OptLevel, UserThreshold: std::nullopt, UserCount: std::nullopt,
285 UserAllowPartial: std::nullopt, UserRuntime: std::nullopt, UserUpperBound: std::nullopt, UserFullUnrollMaxCount: std::nullopt);
286 TargetTransformInfo::PeelingPreferences PP =
287 gatherPeelingPreferences(L, SE, TTI, UserAllowPeeling: std::nullopt, UserAllowProfileBasedPeeling: std::nullopt);
288
289 TransformationMode EnableMode = hasUnrollAndJamTransformation(L);
290 if (EnableMode & TM_Disable)
291 return LoopUnrollResult::Unmodified;
292 if (EnableMode & TM_ForcedByUser)
293 UP.UnrollAndJam = true;
294
295 if (AllowUnrollAndJam.getNumOccurrences() > 0)
296 UP.UnrollAndJam = AllowUnrollAndJam;
297 if (UnrollAndJamThreshold.getNumOccurrences() > 0)
298 UP.UnrollAndJamInnerLoopThreshold = UnrollAndJamThreshold;
299 // Exit early if unrolling is disabled.
300 if (!UP.UnrollAndJam || UP.UnrollAndJamInnerLoopThreshold == 0)
301 return LoopUnrollResult::Unmodified;
302
303 LLVM_DEBUG(dbgs() << "Loop Unroll and Jam: F["
304 << L->getHeader()->getParent()->getName() << "] Loop %"
305 << L->getHeader()->getName() << "\n");
306
307 // A loop with any unroll pragma (enabling/disabling/count/etc) is left for
308 // the unroller, so long as it does not explicitly have unroll_and_jam
309 // metadata. This means #pragma nounroll will disable unroll and jam as well
310 // as unrolling
311 if (hasAnyUnrollPragma(L, Prefix: "llvm.loop.unroll.") &&
312 !hasAnyUnrollPragma(L, Prefix: "llvm.loop.unroll_and_jam.")) {
313 LLVM_DEBUG(dbgs() << " Disabled due to pragma.\n");
314 return LoopUnrollResult::Unmodified;
315 }
316
317 if (!isSafeToUnrollAndJam(L, SE, DT, DI, LI&: *LI)) {
318 LLVM_DEBUG(dbgs() << " Disabled due to not being safe.\n");
319 return LoopUnrollResult::Unmodified;
320 }
321
322 // Approximate the loop size and collect useful info
323 SmallPtrSet<const Value *, 32> EphValues;
324 CodeMetrics::collectEphemeralValues(L, AC: &AC, EphValues);
325 Loop *SubLoop = L->getSubLoops()[0];
326 UnrollCostEstimator InnerUCE(SubLoop, TTI, EphValues, UP.BEInsns);
327 UnrollCostEstimator OuterUCE(L, TTI, EphValues, UP.BEInsns);
328
329 if (!InnerUCE.canUnroll() || !OuterUCE.canUnroll()) {
330 LLVM_DEBUG(dbgs() << " Loop not considered unrollable\n");
331 return LoopUnrollResult::Unmodified;
332 }
333
334 unsigned InnerLoopSize = InnerUCE.getRolledLoopSize();
335 LLVM_DEBUG(dbgs() << " Outer Loop Size: " << OuterUCE.getRolledLoopSize()
336 << "\n");
337 LLVM_DEBUG(dbgs() << " Inner Loop Size: " << InnerLoopSize << "\n");
338
339 if (InnerUCE.NumInlineCandidates != 0 || OuterUCE.NumInlineCandidates != 0) {
340 LLVM_DEBUG(dbgs() << " Not unrolling loop with inlinable calls.\n");
341 return LoopUnrollResult::Unmodified;
342 }
343 // FIXME: The call to canUnroll() allows some controlled convergent
344 // operations, but we block them here for future changes.
345 if (InnerUCE.Convergence != ConvergenceKind::None ||
346 OuterUCE.Convergence != ConvergenceKind::None) {
347 LLVM_DEBUG(
348 dbgs() << " Not unrolling loop with convergent instructions.\n");
349 return LoopUnrollResult::Unmodified;
350 }
351
352 // Save original loop IDs for after the transformation.
353 MDNode *OrigOuterLoopID = L->getLoopID();
354 MDNode *OrigSubLoopID = SubLoop->getLoopID();
355
356 // To assign the loop id of the epilogue, assign it before unrolling it so it
357 // is applied to every inner loop of the epilogue. We later apply the loop ID
358 // for the jammed inner loop.
359 std::optional<MDNode *> NewInnerEpilogueLoopID = makeFollowupLoopID(
360 OrigLoopID: OrigOuterLoopID, FollowupAttrs: {LLVMLoopUnrollAndJamFollowupAll,
361 LLVMLoopUnrollAndJamFollowupRemainderInner});
362 if (NewInnerEpilogueLoopID)
363 SubLoop->setLoopID(*NewInnerEpilogueLoopID);
364
365 // Find trip count and trip multiple
366 BasicBlock *Latch = L->getLoopLatch();
367 BasicBlock *SubLoopLatch = SubLoop->getLoopLatch();
368 unsigned OuterTripCount = SE.getSmallConstantTripCount(L, ExitingBlock: Latch);
369 unsigned OuterTripMultiple = SE.getSmallConstantTripMultiple(L, ExitingBlock: Latch);
370 unsigned InnerTripCount = SE.getSmallConstantTripCount(L: SubLoop, ExitingBlock: SubLoopLatch);
371
372 // Decide if, and by how much, to unroll
373 bool IsCountSetExplicitly = computeUnrollAndJamCount(
374 L, SubLoop, TTI, DT, LI, AC: &AC, SE, EphValues, ORE: &ORE, OuterTripCount,
375 OuterTripMultiple, OuterUCE, InnerTripCount, InnerLoopSize, UP, PP);
376 if (UP.Count <= 1)
377 return LoopUnrollResult::Unmodified;
378 // Unroll factor (Count) must be less or equal to TripCount.
379 if (OuterTripCount && UP.Count > OuterTripCount)
380 UP.Count = OuterTripCount;
381
382 Loop *EpilogueOuterLoop = nullptr;
383 LoopUnrollResult UnrollResult = UnrollAndJamLoop(
384 L, Count: UP.Count, TripCount: OuterTripCount, TripMultiple: OuterTripMultiple, UnrollRemainder: UP.UnrollRemainder, LI,
385 SE: &SE, DT: &DT, AC: &AC, TTI: &TTI, ORE: &ORE, EpilogueLoop: &EpilogueOuterLoop);
386
387 // Assign new loop attributes.
388 if (EpilogueOuterLoop) {
389 std::optional<MDNode *> NewOuterEpilogueLoopID = makeFollowupLoopID(
390 OrigLoopID: OrigOuterLoopID, FollowupAttrs: {LLVMLoopUnrollAndJamFollowupAll,
391 LLVMLoopUnrollAndJamFollowupRemainderOuter});
392 if (NewOuterEpilogueLoopID)
393 EpilogueOuterLoop->setLoopID(*NewOuterEpilogueLoopID);
394 }
395
396 std::optional<MDNode *> NewInnerLoopID =
397 makeFollowupLoopID(OrigLoopID: OrigOuterLoopID, FollowupAttrs: {LLVMLoopUnrollAndJamFollowupAll,
398 LLVMLoopUnrollAndJamFollowupInner});
399 if (NewInnerLoopID)
400 SubLoop->setLoopID(*NewInnerLoopID);
401 else
402 SubLoop->setLoopID(OrigSubLoopID);
403
404 if (UnrollResult == LoopUnrollResult::PartiallyUnrolled) {
405 std::optional<MDNode *> NewOuterLoopID = makeFollowupLoopID(
406 OrigLoopID: OrigOuterLoopID,
407 FollowupAttrs: {LLVMLoopUnrollAndJamFollowupAll, LLVMLoopUnrollAndJamFollowupOuter});
408 if (NewOuterLoopID) {
409 L->setLoopID(*NewOuterLoopID);
410
411 // Do not setLoopAlreadyUnrolled if a followup was given.
412 return UnrollResult;
413 }
414 }
415
416 // If loop has an unroll count pragma or unrolled by explicitly set count
417 // mark loop as unrolled to prevent unrolling beyond that requested.
418 if (UnrollResult != LoopUnrollResult::FullyUnrolled && IsCountSetExplicitly)
419 L->setLoopAlreadyUnrolled();
420
421 return UnrollResult;
422}
423
424static bool tryToUnrollAndJamLoop(LoopNest &LN, DominatorTree &DT, LoopInfo &LI,
425 ScalarEvolution &SE,
426 const TargetTransformInfo &TTI,
427 AssumptionCache &AC, DependenceInfo &DI,
428 OptimizationRemarkEmitter &ORE, int OptLevel,
429 LPMUpdater &U) {
430 bool DidSomething = false;
431 ArrayRef<Loop *> Loops = LN.getLoops();
432 Loop *OutmostLoop = &LN.getOutermostLoop();
433
434 // Add the loop nests in the reverse order of LN. See method
435 // declaration.
436 SmallPriorityWorklist<Loop *, 4> Worklist;
437 appendLoopsToWorklist(Loops, Worklist);
438 while (!Worklist.empty()) {
439 Loop *L = Worklist.pop_back_val();
440 std::string LoopName = std::string(L->getName());
441 LoopUnrollResult Result =
442 tryToUnrollAndJamLoop(L, DT, LI: &LI, SE, TTI, AC, DI, ORE, OptLevel);
443 if (Result != LoopUnrollResult::Unmodified)
444 DidSomething = true;
445 if (L == OutmostLoop && Result == LoopUnrollResult::FullyUnrolled)
446 U.markLoopAsDeleted(L&: *L, Name: LoopName);
447 }
448
449 return DidSomething;
450}
451
452PreservedAnalyses LoopUnrollAndJamPass::run(LoopNest &LN,
453 LoopAnalysisManager &AM,
454 LoopStandardAnalysisResults &AR,
455 LPMUpdater &U) {
456 Function &F = *LN.getParent();
457
458 DependenceInfo DI(&F, &AR.AA, &AR.SE, &AR.LI);
459 OptimizationRemarkEmitter ORE(&F);
460
461 if (!tryToUnrollAndJamLoop(LN, DT&: AR.DT, LI&: AR.LI, SE&: AR.SE, TTI: AR.TTI, AC&: AR.AC, DI, ORE,
462 OptLevel, U))
463 return PreservedAnalyses::all();
464
465 auto PA = getLoopPassPreservedAnalyses();
466 PA.preserve<LoopNestAnalysis>();
467 return PA;
468}
469