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 47 namespace llvm { 48 class Instruction; 49 class Value; 50 } // namespace llvm 51 52 using namespace llvm; 53 54 #define DEBUG_TYPE "loop-unroll-and-jam" 55 56 /// @{ 57 /// Metadata attribute names 58 static const char *const LLVMLoopUnrollAndJamFollowupAll = 59 "llvm.loop.unroll_and_jam.followup_all"; 60 static const char *const LLVMLoopUnrollAndJamFollowupInner = 61 "llvm.loop.unroll_and_jam.followup_inner"; 62 static const char *const LLVMLoopUnrollAndJamFollowupOuter = 63 "llvm.loop.unroll_and_jam.followup_outer"; 64 static const char *const LLVMLoopUnrollAndJamFollowupRemainderInner = 65 "llvm.loop.unroll_and_jam.followup_remainder_inner"; 66 static const char *const LLVMLoopUnrollAndJamFollowupRemainderOuter = 67 "llvm.loop.unroll_and_jam.followup_remainder_outer"; 68 /// @} 69 70 static cl::opt<bool> 71 AllowUnrollAndJam("allow-unroll-and-jam", cl::Hidden, 72 cl::desc("Allows loops to be unroll-and-jammed.")); 73 74 static 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 79 static cl::opt<unsigned> UnrollAndJamThreshold( 80 "unroll-and-jam-threshold", cl::init(60), cl::Hidden, 81 cl::desc("Threshold to use for inner loop when doing unroll and jam.")); 82 83 static cl::opt<unsigned> PragmaUnrollAndJamThreshold( 84 "pragma-unroll-and-jam-threshold", cl::init(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. 91 static 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. 99 static 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>(LoopID->getOperand(I)); 107 if (!MD) 108 continue; 109 110 MDString *S = dyn_cast<MDString>(MD->getOperand(0)); 111 if (!S) 112 continue; 113 114 if (S->getString().startswith(Prefix)) 115 return true; 116 } 117 } 118 return false; 119 } 120 121 // Returns true if the loop has an unroll_and_jam(enable) pragma. 122 static bool hasUnrollAndJamEnablePragma(const Loop *L) { 123 return getUnrollMetadataForLoop(L, "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. 128 static unsigned unrollAndJamCountPragmaValue(const Loop *L) { 129 MDNode *MD = getUnrollMetadataForLoop(L, "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->getOperand(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. 142 static uint64_t 143 getUnrollAndJammedLoopSize(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. 151 static 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, unsigned OuterLoopSize, unsigned InnerTripCount, 157 unsigned InnerLoopSize, TargetTransformInfo::UnrollingPreferences &UP, 158 TargetTransformInfo::PeelingPreferences &PP) { 159 // First up use computeUnrollCount from the loop unroller to get a count 160 // for unrolling the outer loop, plus any loops requiring explicit 161 // unrolling we leave to the unroller. This uses UP.Threshold / 162 // UP.PartialThreshold / UP.MaxCount to come up with sensible loop values. 163 // We have already checked that the loop has no unroll.* pragmas. 164 unsigned MaxTripCount = 0; 165 bool UseUpperBound = false; 166 bool ExplicitUnroll = computeUnrollCount( 167 L, TTI, DT, LI, AC, SE, EphValues, ORE, OuterTripCount, MaxTripCount, 168 /*MaxOrZero*/ false, OuterTripMultiple, OuterLoopSize, UP, PP, 169 UseUpperBound); 170 if (ExplicitUnroll || UseUpperBound) { 171 // If the user explicitly set the loop as unrolled, dont UnJ it. Leave it 172 // for the unroller instead. 173 LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; explicit count set by " 174 "computeUnrollCount\n"); 175 UP.Count = 0; 176 return false; 177 } 178 179 // Override with any explicit Count from the "unroll-and-jam-count" option. 180 bool UserUnrollCount = UnrollAndJamCount.getNumOccurrences() > 0; 181 if (UserUnrollCount) { 182 UP.Count = UnrollAndJamCount; 183 UP.Force = true; 184 if (UP.AllowRemainder && 185 getUnrollAndJammedLoopSize(OuterLoopSize, UP) < UP.Threshold && 186 getUnrollAndJammedLoopSize(InnerLoopSize, UP) < 187 UP.UnrollAndJamInnerLoopThreshold) 188 return true; 189 } 190 191 // Check for unroll_and_jam pragmas 192 unsigned PragmaCount = unrollAndJamCountPragmaValue(L); 193 if (PragmaCount > 0) { 194 UP.Count = PragmaCount; 195 UP.Runtime = true; 196 UP.Force = true; 197 if ((UP.AllowRemainder || (OuterTripMultiple % PragmaCount == 0)) && 198 getUnrollAndJammedLoopSize(OuterLoopSize, UP) < UP.Threshold && 199 getUnrollAndJammedLoopSize(InnerLoopSize, UP) < 200 UP.UnrollAndJamInnerLoopThreshold) 201 return true; 202 } 203 204 bool PragmaEnableUnroll = hasUnrollAndJamEnablePragma(L); 205 bool ExplicitUnrollAndJamCount = PragmaCount > 0 || UserUnrollCount; 206 bool ExplicitUnrollAndJam = PragmaEnableUnroll || ExplicitUnrollAndJamCount; 207 208 // If the loop has an unrolling pragma, we want to be more aggressive with 209 // unrolling limits. 210 if (ExplicitUnrollAndJam) 211 UP.UnrollAndJamInnerLoopThreshold = PragmaUnrollAndJamThreshold; 212 213 if (!UP.AllowRemainder && getUnrollAndJammedLoopSize(InnerLoopSize, UP) >= 214 UP.UnrollAndJamInnerLoopThreshold) { 215 LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; can't create remainder and " 216 "inner loop too large\n"); 217 UP.Count = 0; 218 return false; 219 } 220 221 // We have a sensible limit for the outer loop, now adjust it for the inner 222 // loop and UP.UnrollAndJamInnerLoopThreshold. If the outer limit was set 223 // explicitly, we want to stick to it. 224 if (!ExplicitUnrollAndJamCount && UP.AllowRemainder) { 225 while (UP.Count != 0 && getUnrollAndJammedLoopSize(InnerLoopSize, UP) >= 226 UP.UnrollAndJamInnerLoopThreshold) 227 UP.Count--; 228 } 229 230 // If we are explicitly unroll and jamming, we are done. Otherwise there are a 231 // number of extra performance heuristics to check. 232 if (ExplicitUnrollAndJam) 233 return true; 234 235 // If the inner loop count is known and small, leave the entire loop nest to 236 // be the unroller 237 if (InnerTripCount && InnerLoopSize * InnerTripCount < UP.Threshold) { 238 LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; small inner loop count is " 239 "being left for the unroller\n"); 240 UP.Count = 0; 241 return false; 242 } 243 244 // Check for situations where UnJ is likely to be unprofitable. Including 245 // subloops with more than 1 block. 246 if (SubLoop->getBlocks().size() != 1) { 247 LLVM_DEBUG( 248 dbgs() << "Won't unroll-and-jam; More than one inner loop block\n"); 249 UP.Count = 0; 250 return false; 251 } 252 253 // Limit to loops where there is something to gain from unrolling and 254 // jamming the loop. In this case, look for loads that are invariant in the 255 // outer loop and can become shared. 256 unsigned NumInvariant = 0; 257 for (BasicBlock *BB : SubLoop->getBlocks()) { 258 for (Instruction &I : *BB) { 259 if (auto *Ld = dyn_cast<LoadInst>(&I)) { 260 Value *V = Ld->getPointerOperand(); 261 const SCEV *LSCEV = SE.getSCEVAtScope(V, L); 262 if (SE.isLoopInvariant(LSCEV, L)) 263 NumInvariant++; 264 } 265 } 266 } 267 if (NumInvariant == 0) { 268 LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; No loop invariant loads\n"); 269 UP.Count = 0; 270 return false; 271 } 272 273 return false; 274 } 275 276 static LoopUnrollResult 277 tryToUnrollAndJamLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, 278 ScalarEvolution &SE, const TargetTransformInfo &TTI, 279 AssumptionCache &AC, DependenceInfo &DI, 280 OptimizationRemarkEmitter &ORE, int OptLevel) { 281 TargetTransformInfo::UnrollingPreferences UP = gatherUnrollingPreferences( 282 L, SE, TTI, nullptr, nullptr, ORE, OptLevel, std::nullopt, std::nullopt, 283 std::nullopt, std::nullopt, std::nullopt, std::nullopt); 284 TargetTransformInfo::PeelingPreferences PP = 285 gatherPeelingPreferences(L, SE, TTI, std::nullopt, std::nullopt); 286 287 TransformationMode EnableMode = hasUnrollAndJamTransformation(L); 288 if (EnableMode & TM_Disable) 289 return LoopUnrollResult::Unmodified; 290 if (EnableMode & TM_ForcedByUser) 291 UP.UnrollAndJam = true; 292 293 if (AllowUnrollAndJam.getNumOccurrences() > 0) 294 UP.UnrollAndJam = AllowUnrollAndJam; 295 if (UnrollAndJamThreshold.getNumOccurrences() > 0) 296 UP.UnrollAndJamInnerLoopThreshold = UnrollAndJamThreshold; 297 // Exit early if unrolling is disabled. 298 if (!UP.UnrollAndJam || UP.UnrollAndJamInnerLoopThreshold == 0) 299 return LoopUnrollResult::Unmodified; 300 301 LLVM_DEBUG(dbgs() << "Loop Unroll and Jam: F[" 302 << L->getHeader()->getParent()->getName() << "] Loop %" 303 << L->getHeader()->getName() << "\n"); 304 305 // A loop with any unroll pragma (enabling/disabling/count/etc) is left for 306 // the unroller, so long as it does not explicitly have unroll_and_jam 307 // metadata. This means #pragma nounroll will disable unroll and jam as well 308 // as unrolling 309 if (hasAnyUnrollPragma(L, "llvm.loop.unroll.") && 310 !hasAnyUnrollPragma(L, "llvm.loop.unroll_and_jam.")) { 311 LLVM_DEBUG(dbgs() << " Disabled due to pragma.\n"); 312 return LoopUnrollResult::Unmodified; 313 } 314 315 if (!isSafeToUnrollAndJam(L, SE, DT, DI, *LI)) { 316 LLVM_DEBUG(dbgs() << " Disabled due to not being safe.\n"); 317 return LoopUnrollResult::Unmodified; 318 } 319 320 // Approximate the loop size and collect useful info 321 unsigned NumInlineCandidates; 322 bool NotDuplicatable; 323 bool Convergent; 324 SmallPtrSet<const Value *, 32> EphValues; 325 CodeMetrics::collectEphemeralValues(L, &AC, EphValues); 326 Loop *SubLoop = L->getSubLoops()[0]; 327 InstructionCost InnerLoopSizeIC = 328 ApproximateLoopSize(SubLoop, NumInlineCandidates, NotDuplicatable, 329 Convergent, TTI, EphValues, UP.BEInsns); 330 InstructionCost OuterLoopSizeIC = 331 ApproximateLoopSize(L, NumInlineCandidates, NotDuplicatable, Convergent, 332 TTI, EphValues, UP.BEInsns); 333 LLVM_DEBUG(dbgs() << " Outer Loop Size: " << OuterLoopSizeIC << "\n"); 334 LLVM_DEBUG(dbgs() << " Inner Loop Size: " << InnerLoopSizeIC << "\n"); 335 336 if (!InnerLoopSizeIC.isValid() || !OuterLoopSizeIC.isValid()) { 337 LLVM_DEBUG(dbgs() << " Not unrolling loop which contains instructions" 338 << " with invalid cost.\n"); 339 return LoopUnrollResult::Unmodified; 340 } 341 unsigned InnerLoopSize = *InnerLoopSizeIC.getValue(); 342 unsigned OuterLoopSize = *OuterLoopSizeIC.getValue(); 343 344 if (NotDuplicatable) { 345 LLVM_DEBUG(dbgs() << " Not unrolling loop which contains non-duplicatable " 346 "instructions.\n"); 347 return LoopUnrollResult::Unmodified; 348 } 349 if (NumInlineCandidates != 0) { 350 LLVM_DEBUG(dbgs() << " Not unrolling loop with inlinable calls.\n"); 351 return LoopUnrollResult::Unmodified; 352 } 353 if (Convergent) { 354 LLVM_DEBUG( 355 dbgs() << " Not unrolling loop with convergent instructions.\n"); 356 return LoopUnrollResult::Unmodified; 357 } 358 359 // Save original loop IDs for after the transformation. 360 MDNode *OrigOuterLoopID = L->getLoopID(); 361 MDNode *OrigSubLoopID = SubLoop->getLoopID(); 362 363 // To assign the loop id of the epilogue, assign it before unrolling it so it 364 // is applied to every inner loop of the epilogue. We later apply the loop ID 365 // for the jammed inner loop. 366 std::optional<MDNode *> NewInnerEpilogueLoopID = makeFollowupLoopID( 367 OrigOuterLoopID, {LLVMLoopUnrollAndJamFollowupAll, 368 LLVMLoopUnrollAndJamFollowupRemainderInner}); 369 if (NewInnerEpilogueLoopID) 370 SubLoop->setLoopID(*NewInnerEpilogueLoopID); 371 372 // Find trip count and trip multiple 373 BasicBlock *Latch = L->getLoopLatch(); 374 BasicBlock *SubLoopLatch = SubLoop->getLoopLatch(); 375 unsigned OuterTripCount = SE.getSmallConstantTripCount(L, Latch); 376 unsigned OuterTripMultiple = SE.getSmallConstantTripMultiple(L, Latch); 377 unsigned InnerTripCount = SE.getSmallConstantTripCount(SubLoop, SubLoopLatch); 378 379 // Decide if, and by how much, to unroll 380 bool IsCountSetExplicitly = computeUnrollAndJamCount( 381 L, SubLoop, TTI, DT, LI, &AC, SE, EphValues, &ORE, OuterTripCount, 382 OuterTripMultiple, OuterLoopSize, InnerTripCount, InnerLoopSize, UP, PP); 383 if (UP.Count <= 1) 384 return LoopUnrollResult::Unmodified; 385 // Unroll factor (Count) must be less or equal to TripCount. 386 if (OuterTripCount && UP.Count > OuterTripCount) 387 UP.Count = OuterTripCount; 388 389 Loop *EpilogueOuterLoop = nullptr; 390 LoopUnrollResult UnrollResult = UnrollAndJamLoop( 391 L, UP.Count, OuterTripCount, OuterTripMultiple, UP.UnrollRemainder, LI, 392 &SE, &DT, &AC, &TTI, &ORE, &EpilogueOuterLoop); 393 394 // Assign new loop attributes. 395 if (EpilogueOuterLoop) { 396 std::optional<MDNode *> NewOuterEpilogueLoopID = makeFollowupLoopID( 397 OrigOuterLoopID, {LLVMLoopUnrollAndJamFollowupAll, 398 LLVMLoopUnrollAndJamFollowupRemainderOuter}); 399 if (NewOuterEpilogueLoopID) 400 EpilogueOuterLoop->setLoopID(*NewOuterEpilogueLoopID); 401 } 402 403 std::optional<MDNode *> NewInnerLoopID = 404 makeFollowupLoopID(OrigOuterLoopID, {LLVMLoopUnrollAndJamFollowupAll, 405 LLVMLoopUnrollAndJamFollowupInner}); 406 if (NewInnerLoopID) 407 SubLoop->setLoopID(*NewInnerLoopID); 408 else 409 SubLoop->setLoopID(OrigSubLoopID); 410 411 if (UnrollResult == LoopUnrollResult::PartiallyUnrolled) { 412 std::optional<MDNode *> NewOuterLoopID = makeFollowupLoopID( 413 OrigOuterLoopID, 414 {LLVMLoopUnrollAndJamFollowupAll, LLVMLoopUnrollAndJamFollowupOuter}); 415 if (NewOuterLoopID) { 416 L->setLoopID(*NewOuterLoopID); 417 418 // Do not setLoopAlreadyUnrolled if a followup was given. 419 return UnrollResult; 420 } 421 } 422 423 // If loop has an unroll count pragma or unrolled by explicitly set count 424 // mark loop as unrolled to prevent unrolling beyond that requested. 425 if (UnrollResult != LoopUnrollResult::FullyUnrolled && IsCountSetExplicitly) 426 L->setLoopAlreadyUnrolled(); 427 428 return UnrollResult; 429 } 430 431 static bool tryToUnrollAndJamLoop(LoopNest &LN, DominatorTree &DT, LoopInfo &LI, 432 ScalarEvolution &SE, 433 const TargetTransformInfo &TTI, 434 AssumptionCache &AC, DependenceInfo &DI, 435 OptimizationRemarkEmitter &ORE, int OptLevel, 436 LPMUpdater &U) { 437 bool DidSomething = false; 438 ArrayRef<Loop *> Loops = LN.getLoops(); 439 Loop *OutmostLoop = &LN.getOutermostLoop(); 440 441 // Add the loop nests in the reverse order of LN. See method 442 // declaration. 443 SmallPriorityWorklist<Loop *, 4> Worklist; 444 appendLoopsToWorklist(Loops, Worklist); 445 while (!Worklist.empty()) { 446 Loop *L = Worklist.pop_back_val(); 447 std::string LoopName = std::string(L->getName()); 448 LoopUnrollResult Result = 449 tryToUnrollAndJamLoop(L, DT, &LI, SE, TTI, AC, DI, ORE, OptLevel); 450 if (Result != LoopUnrollResult::Unmodified) 451 DidSomething = true; 452 if (L == OutmostLoop && Result == LoopUnrollResult::FullyUnrolled) 453 U.markLoopAsDeleted(*L, LoopName); 454 } 455 456 return DidSomething; 457 } 458 459 PreservedAnalyses LoopUnrollAndJamPass::run(LoopNest &LN, 460 LoopAnalysisManager &AM, 461 LoopStandardAnalysisResults &AR, 462 LPMUpdater &U) { 463 Function &F = *LN.getParent(); 464 465 DependenceInfo DI(&F, &AR.AA, &AR.SE, &AR.LI); 466 OptimizationRemarkEmitter ORE(&F); 467 468 if (!tryToUnrollAndJamLoop(LN, AR.DT, AR.LI, AR.SE, AR.TTI, AR.AC, DI, ORE, 469 OptLevel, U)) 470 return PreservedAnalyses::all(); 471 472 auto PA = getLoopPassPreservedAnalyses(); 473 PA.preserve<LoopNestAnalysis>(); 474 return PA; 475 } 476