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