[Unroll] Add an option to control complete unrolling

[skatkov]

Summary

Add an ability to specify the max full unroll count for LoopUnrollPass pass
in pass options.

[fedor.sergeev]

I might be in need for the morning coffee, but I cant find the actual use of UP.FullUnrollMaxCount...

[fedor.sergeev]

In D67701#1674862, @fedor.sergeev wrote:

I might be in need for the morning coffee, but I cant find the actual use of UP.FullUnrollMaxCount...

Ugh... I was looking for that use in the new code, but it is there in the old code, sorry.

So this patch just adds a pass-level option to provide to an existing control.

Looked through the tests and it seems that we have no tests to check this existing control (no use of -unroll-full-max-count option).
Can you, please, extend your test to check:

• for actual "max" functionality, where max is non-zero and trip-count is either more or less than the max
• both ways of controlling this max-count (-passes and direct option)

Thanks!

[skatkov]

More tests are added.

[fedor.sergeev]

Great, thanks a lot!

llvm/include/llvm/Transforms/Scalar/LoopUnrollPass.h

 /// additional setters and then pass it to LoopUnrollPass.
 ///
 struct LoopUnrollOptions {
   Optional<bool> AllowPartial;
   Optional<bool> AllowPeeling;
   Optional<bool> AllowRuntime;
   Optional<bool> AllowUpperBound;
   Optional<bool> AllowProfileBasedPeeling;
+  Optional<unsigned> FullUnrollMaxCount;
   int OptLevel;
 
   /// If false, use a cost model to determine whether unrolling of a loop is
   /// profitable. If true, only loops that explicitly request unrolling via
   /// metadata are considered. All other loops are skipped.
   bool OnlyWhenForced;
 
   /// If true, forget all loops when unrolling. If false, forget top-most loop
LoopUnrollOptions &setOptLevel(int O) {
     return *this;
   }
 
   // Enables or disables loop peeling basing on profile.
   LoopUnrollOptions &setProfileBasedPeeling(int O) {
     AllowProfileBasedPeeling = O;
     return *this;
   }
+
+  // Sets the max full unroll count.
+  LoopUnrollOptions &setFullUnrollMaxCount(unsigned O) {
+    FullUnrollMaxCount = O;
+    return *this;
+  }
 };
 
 /// Loop unroll pass that will support both full and partial unrolling.
 /// It is a function pass to have access to function and module analyses.
 /// It will also put loops into canonical form (simplified and LCSSA).
 class LoopUnrollPass : public PassInfoMixin<LoopUnrollPass> {
   LoopUnrollOptions UnrollOpts;
 

llvm/include/llvm/Transforms/Utils/UnrollLoop.h

 MDNode *GetUnrollMetadata(MDNode *LoopID, StringRef Name);
 
 TargetTransformInfo::UnrollingPreferences gatherUnrollingPreferences(
     Loop *L, ScalarEvolution &SE, const TargetTransformInfo &TTI,
     BlockFrequencyInfo *BFI, ProfileSummaryInfo *PSI, int OptLevel,
     Optional<unsigned> UserThreshold, Optional<unsigned> UserCount,
     Optional<bool> UserAllowPartial, Optional<bool> UserRuntime,
     Optional<bool> UserUpperBound, Optional<bool> UserAllowPeeling,
-    Optional<bool> UserAllowProfileBasedPeeling);
+    Optional<bool> UserAllowProfileBasedPeeling,
+    Optional<unsigned> UserFullUnrollMaxCount);
 
 unsigned ApproximateLoopSize(const Loop *L, unsigned &NumCalls,
                              bool &NotDuplicatable, bool &Convergent,
                              const TargetTransformInfo &TTI,
                              const SmallPtrSetImpl<const Value *> &EphValues,
                              unsigned BEInsns);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_UNROLLLOOP_H

llvm/lib/Passes/PassBuilder.cpp

int OptLevel = StringSwitch<int>(ParamName)
                        .Case("O1", 1)
                        .Case("O2", 2)
                        .Case("O3", 3)
                        .Default(-1);
     if (OptLevel >= 0) {
       UnrollOpts.setOptLevel(OptLevel);
       continue;
     }
+    if (ParamName.consume_front("full-unroll-max=")) {
+      int Count;
+      if (ParamName.getAsInteger(0, Count))
+        return make_error<StringError>(
+            formatv("invalid LoopUnrollPass parameter '{0}' ", ParamName).str(),
+            inconvertibleErrorCode());
+      UnrollOpts.setFullUnrollMaxCount(Count);
+      continue;
+    }
 
     bool Enable = !ParamName.consume_front("no-");
     if (ParamName == "partial") {
       UnrollOpts.setPartial(Enable);
     } else if (ParamName == "peeling") {
       UnrollOpts.setPeeling(Enable);
     } else if (ParamName == "profile-peeling") {
       UnrollOpts.setProfileBasedPeeling(Enable);

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

tryToUnrollAndJamLoop(Loop *L, DominatorTree &DT, LoopInfo *LI,
   BasicBlock *SubLoopLatch = SubLoop->getLoopLatch();
   BasicBlock *SubLoopExit = SubLoop->getExitingBlock();
 
   if (Latch != Exit || SubLoopLatch != SubLoopExit)
     return LoopUnrollResult::Unmodified;
 
   TargetTransformInfo::UnrollingPreferences UP =
       gatherUnrollingPreferences(L, SE, TTI, nullptr, nullptr, OptLevel, None,
-                                 None, None, None, None, None, None);
+                                 None, None, None, None, None, None, None);
   if (AllowUnrollAndJam.getNumOccurrences() > 0)
     UP.UnrollAndJam = AllowUnrollAndJam;
   if (UnrollAndJamThreshold.getNumOccurrences() > 0)
     UP.UnrollAndJamInnerLoopThreshold = UnrollAndJamThreshold;
   // Exit early if unrolling is disabled.
   if (!UP.UnrollAndJam || UP.UnrollAndJamInnerLoopThreshold == 0)
     return LoopUnrollResult::Unmodified;
 

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

 /// Gather the various unrolling parameters based on the defaults, compiler
 /// flags, TTI overrides and user specified parameters.
 TargetTransformInfo::UnrollingPreferences llvm::gatherUnrollingPreferences(
     Loop *L, ScalarEvolution &SE, const TargetTransformInfo &TTI,
     BlockFrequencyInfo *BFI, ProfileSummaryInfo *PSI, int OptLevel,
     Optional<unsigned> UserThreshold, Optional<unsigned> UserCount,
     Optional<bool> UserAllowPartial, Optional<bool> UserRuntime,
     Optional<bool> UserUpperBound, Optional<bool> UserAllowPeeling,
-    Optional<bool> UserAllowProfileBasedPeeling) {
+    Optional<bool> UserAllowProfileBasedPeeling,
+    Optional<unsigned> UserFullUnrollMaxCount) {
   TargetTransformInfo::UnrollingPreferences UP;
 
   // Set up the defaults
   UP.Threshold = OptLevel > 2 ? 300 : 150;
   UP.MaxPercentThresholdBoost = 400;
   UP.OptSizeThreshold = 0;
   UP.PartialThreshold = 150;
   UP.PartialOptSizeThreshold = 0;
TargetTransformInfo::UnrollingPreferences llvm::gatherUnrollingPreferences(
   if (UserRuntime.hasValue())
     UP.Runtime = *UserRuntime;
   if (UserUpperBound.hasValue())
     UP.UpperBound = *UserUpperBound;
   if (UserAllowPeeling.hasValue())
     UP.AllowPeeling = *UserAllowPeeling;
   if (UserAllowProfileBasedPeeling.hasValue())
     UP.PeelProfiledIterations = *UserAllowProfileBasedPeeling;
+  if (UserFullUnrollMaxCount.hasValue())
+    UP.FullUnrollMaxCount = *UserFullUnrollMaxCount;
 
   return UP;
 }
 
 namespace {
 
 /// A struct to densely store the state of an instruction after unrolling at
 /// each iteration.
static LoopUnrollResult tryToUnrollLoop(
     Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
     const TargetTransformInfo &TTI, AssumptionCache &AC,
     OptimizationRemarkEmitter &ORE, BlockFrequencyInfo *BFI,
     ProfileSummaryInfo *PSI, bool PreserveLCSSA, int OptLevel,
     bool OnlyWhenForced, bool ForgetAllSCEV, Optional<unsigned> ProvidedCount,
     Optional<unsigned> ProvidedThreshold, Optional<bool> ProvidedAllowPartial,
     Optional<bool> ProvidedRuntime, Optional<bool> ProvidedUpperBound,
     Optional<bool> ProvidedAllowPeeling,
-    Optional<bool> ProvidedAllowProfileBasedPeeling) {
+    Optional<bool> ProvidedAllowProfileBasedPeeling,
+    Optional<unsigned> ProvidedFullUnrollMaxCount) {
   LLVM_DEBUG(dbgs() << "Loop Unroll: F["
                     << L->getHeader()->getParent()->getName() << "] Loop %"
                     << L->getHeader()->getName() << "\n");
   TransformationMode TM = hasUnrollTransformation(L);
   if (TM & TM_Disable)
     return LoopUnrollResult::Unmodified;
   if (!L->isLoopSimplifyForm()) {
     LLVM_DEBUG(
         dbgs() << "  Not unrolling loop which is not in loop-simplify form.\n");
     return LoopUnrollResult::Unmodified;
   }
 
   // When automtatic unrolling is disabled, do not unroll unless overridden for
   // this loop.
   if (OnlyWhenForced && !(TM & TM_Enable))
     return LoopUnrollResult::Unmodified;
 
   bool OptForSize = L->getHeader()->getParent()->hasOptSize();
   unsigned NumInlineCandidates;
   bool NotDuplicatable;
   bool Convergent;
   TargetTransformInfo::UnrollingPreferences UP = gatherUnrollingPreferences(
       L, SE, TTI, BFI, PSI, OptLevel, ProvidedThreshold, ProvidedCount,
       ProvidedAllowPartial, ProvidedRuntime, ProvidedUpperBound,
-      ProvidedAllowPeeling, ProvidedAllowProfileBasedPeeling);
+      ProvidedAllowPeeling, ProvidedAllowProfileBasedPeeling,
+      ProvidedFullUnrollMaxCount);
 
   // Exit early if unrolling is disabled. For OptForSize, we pick the loop size
   // as threshold later on.
   if (UP.Threshold == 0 && (!UP.Partial || UP.PartialThreshold == 0) &&
       !OptForSize)
     return LoopUnrollResult::Unmodified;
 
   SmallPtrSet<const Value *, 32> EphValues;
public:
 
   Optional<unsigned> ProvidedCount;
   Optional<unsigned> ProvidedThreshold;
   Optional<bool> ProvidedAllowPartial;
   Optional<bool> ProvidedRuntime;
   Optional<bool> ProvidedUpperBound;
   Optional<bool> ProvidedAllowPeeling;
   Optional<bool> ProvidedAllowProfileBasedPeeling;
+  Optional<unsigned> ProvidedFullUnrollMaxCount;
 
   LoopUnroll(int OptLevel = 2, bool OnlyWhenForced = false,
              bool ForgetAllSCEV = false, Optional<unsigned> Threshold = None,
              Optional<unsigned> Count = None,
              Optional<bool> AllowPartial = None, Optional<bool> Runtime = None,
              Optional<bool> UpperBound = None,
              Optional<bool> AllowPeeling = None,
-             Optional<bool> AllowProfileBasedPeeling = None)
+             Optional<bool> AllowProfileBasedPeeling = None,
+             Optional<unsigned> ProvidedFullUnrollMaxCount = None)
       : LoopPass(ID), OptLevel(OptLevel), OnlyWhenForced(OnlyWhenForced),
         ForgetAllSCEV(ForgetAllSCEV), ProvidedCount(std::move(Count)),
         ProvidedThreshold(Threshold), ProvidedAllowPartial(AllowPartial),
         ProvidedRuntime(Runtime), ProvidedUpperBound(UpperBound),
         ProvidedAllowPeeling(AllowPeeling),
-        ProvidedAllowProfileBasedPeeling(AllowProfileBasedPeeling) {
+        ProvidedAllowProfileBasedPeeling(AllowProfileBasedPeeling),
+        ProvidedFullUnrollMaxCount(ProvidedFullUnrollMaxCount) {
     initializeLoopUnrollPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnLoop(Loop *L, LPPassManager &LPM) override {
     if (skipLoop(L))
       return false;
 
     Function &F = *L->getHeader()->getParent();
bool runOnLoop(Loop *L, LPPassManager &LPM) override {
     // but ORE cannot be preserved (see comment before the pass definition).
     OptimizationRemarkEmitter ORE(&F);
     bool PreserveLCSSA = mustPreserveAnalysisID(LCSSAID);
 
     LoopUnrollResult Result = tryToUnrollLoop(
         L, DT, LI, SE, TTI, AC, ORE, nullptr, nullptr, PreserveLCSSA, OptLevel,
         OnlyWhenForced, ForgetAllSCEV, ProvidedCount, ProvidedThreshold,
         ProvidedAllowPartial, ProvidedRuntime, ProvidedUpperBound,
-        ProvidedAllowPeeling, ProvidedAllowProfileBasedPeeling);
+        ProvidedAllowPeeling, ProvidedAllowProfileBasedPeeling,
+        ProvidedFullUnrollMaxCount);
 
     if (Result == LoopUnrollResult::FullyUnrolled)
       LPM.markLoopAsDeleted(*L);
 
     return Result != LoopUnrollResult::Unmodified;
   }
 
   /// This transformation requires natural loop information & requires that
PreservedAnalyses LoopFullUnrollPass::run(Loop &L, LoopAnalysisManager &AM,
 
   bool Changed = tryToUnrollLoop(&L, AR.DT, &AR.LI, AR.SE, AR.TTI, AR.AC, *ORE,
                                  /*BFI*/ nullptr, /*PSI*/ nullptr,
                                  /*PreserveLCSSA*/ true, OptLevel,
                                  OnlyWhenForced, ForgetSCEV, /*Count*/ None,
                                  /*Threshold*/ None, /*AllowPartial*/ false,
                                  /*Runtime*/ false, /*UpperBound*/ false,
                                  /*AllowPeeling*/ false,
-                                 /*AllowProfileBasedPeeling*/ false) !=
+                                 /*AllowProfileBasedPeeling*/ false,
+                                 /*FullUnrollMaxCount*/ None) !=
                  LoopUnrollResult::Unmodified;
   if (!Changed)
     return PreservedAnalyses::all();
 
   // The parent must not be damaged by unrolling!
 #ifndef NDEBUG
   if (ParentL)
     ParentL->verifyLoop();
#endif
     // The API here is quite complex to call and we allow to select some
     // flavors of unrolling during construction time (by setting UnrollOpts).
     LoopUnrollResult Result = tryToUnrollLoop(
         &L, DT, &LI, SE, TTI, AC, ORE, BFI, PSI,
         /*PreserveLCSSA*/ true, UnrollOpts.OptLevel, UnrollOpts.OnlyWhenForced,
         UnrollOpts.ForgetSCEV, /*Count*/ None,
         /*Threshold*/ None, UnrollOpts.AllowPartial, UnrollOpts.AllowRuntime,
         UnrollOpts.AllowUpperBound, LocalAllowPeeling,
-        UnrollOpts.AllowProfileBasedPeeling);
+        UnrollOpts.AllowProfileBasedPeeling, UnrollOpts.FullUnrollMaxCount);
     Changed |= Result != LoopUnrollResult::Unmodified;
 
     // The parent must not be damaged by unrolling!
 #ifndef NDEBUG
     if (Result != LoopUnrollResult::Unmodified && ParentL)
       ParentL->verifyLoop();
 #endif
 

llvm/test/Transforms/LoopUnroll/disable-full-unroll-by-opt.ll

+; Default behavior
+; RUN: opt < %s -passes='unroll' -S | FileCheck %s -check-prefixes=ENABLE,COMMON
+
+; Pass option
+; RUN: opt < %s -passes='unroll<full-unroll-max=0>'  -S | FileCheck %s -check-prefixes=DISABLE,COMMON
+; RUN: opt < %s -passes='unroll<full-unroll-max=30>' -S | FileCheck %s -check-prefixes=DISABLE,COMMON
+; RUN: opt < %s -passes='unroll<full-unroll-max=36>' -S | FileCheck %s -check-prefixes=ENABLE,COMMON
+
+; cl::opt option
+; RUN: opt < %s -passes='unroll' -unroll-full-max-count=0 -S | FileCheck %s -check-prefixes=DISABLE,COMMON
+; RUN: opt < %s -passes='unroll' -unroll-full-max-count=30 -S | FileCheck %s -check-prefixes=DISABLE,COMMON
+; RUN: opt < %s -passes='unroll' -unroll-full-max-count=36 -S | FileCheck %s -check-prefixes=ENABLE,COMMON
+
+; Pass option has a priority over cl::opt
+; RUN: opt < %s -passes='unroll<full-unroll-max=30>' -unroll-full-max-count=36 -S | FileCheck %s -check-prefixes=DISABLE,COMMON
+; RUN: opt < %s -passes='unroll<full-unroll-max=36>' -unroll-full-max-count=30 -S | FileCheck %s -check-prefixes=ENABLE,COMMON
+
+define void @test() {
+; COMMON-LABEL: @test(
+ entry:
+  br label %loop
+
+ loop:
+  %idx = phi i32 [ 0, %entry ], [ %idx.inc, %loop ]
+  %idx.inc = add i32 %idx, 1
+  %be = icmp slt i32 %idx, 32
+  br i1 %be, label %loop, label %exit
+
+; COMMON: loop:
+; DISABLE:  %be = icmp slt i32 %idx, 32
+; ENABLE-NOT:  %be = icmp slt i32 %idx, 32
+
+ exit:
+  ret void
+}