[Unroll] Add an option to control complete unrolling

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

Reviewers: fhahn, fedor.sergeev
Reviewed By: fedor.sergeev
Subscribers: hiraditya, zzheng, dmgreen, llvm-commits
Differential Revision: https://reviews.llvm.org/D67701

llvm-svn: 372305

6 files changed, 72 insertions, 10 deletions

diff --git a/llvm/include/llvm/Transforms/Scalar/LoopUnrollPass.h b/llvm/include/llvm/Transforms/Scalar/LoopUnrollPass.h
index 951dc90d823..afeb1f1da02 100644
--- a/llvm/include/llvm/Transforms/Scalar/LoopUnrollPass.h
+++ b/llvm/include/llvm/Transforms/Scalar/LoopUnrollPass.h
@@ -63,6 +63,7 @@ struct LoopUnrollOptions {
   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
@@ -117,6 +118,12 @@ struct LoopUnrollOptions {
     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.
diff --git a/llvm/include/llvm/Transforms/Utils/UnrollLoop.h b/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
index f3651323c25..9012920745e 100644
--- a/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
+++ b/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
@@ -133,7 +133,8 @@ TargetTransformInfo::UnrollingPreferences gatherUnrollingPreferences(
     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,
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 12c4dcb86ce..08be32b5adb 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1460,6 +1460,15 @@ Expected<LoopUnrollOptions> parseLoopUnrollOptions(StringRef Params) {
       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") {
diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp
index 1b2dd7dacdf..d40ef6be9b2 100644
--- a/llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp
@@ -295,7 +295,7 @@ tryToUnrollAndJamLoop(Loop *L, DominatorTree &DT, LoopInfo *LI,
 
   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)
diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
index 244c7d15192..f801c208903 100644
--- a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -179,7 +179,8 @@ TargetTransformInfo::UnrollingPreferences llvm::gatherUnrollingPreferences(
     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
@@ -261,6 +262,8 @@ TargetTransformInfo::UnrollingPreferences llvm::gatherUnrollingPreferences(
     UP.AllowPeeling = *UserAllowPeeling;
   if (UserAllowProfileBasedPeeling.hasValue())
     UP.PeelProfiledIterations = *UserAllowProfileBasedPeeling;
+  if (UserFullUnrollMaxCount.hasValue())
+    UP.FullUnrollMaxCount = *UserFullUnrollMaxCount;
 
   return UP;
 }
@@ -986,7 +989,8 @@ static LoopUnrollResult tryToUnrollLoop(
     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");
@@ -1011,7 +1015,8 @@ static LoopUnrollResult tryToUnrollLoop(
   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.
@@ -1174,6 +1179,7 @@ public:
   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,
@@ -1181,13 +1187,15 @@ public:
              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());
   }
 
@@ -1213,7 +1221,8 @@ public:
         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);
@@ -1297,7 +1306,8 @@ PreservedAnalyses LoopFullUnrollPass::run(Loop &L, LoopAnalysisManager &AM,
                                  /*Threshold*/ None, /*AllowPartial*/ false,
                                  /*Runtime*/ false, /*UpperBound*/ false,
                                  /*AllowPeeling*/ false,
-                                 /*AllowProfileBasedPeeling*/ false) !=
+                                 /*AllowProfileBasedPeeling*/ false,
+                                 /*FullUnrollMaxCount*/ None) !=
                  LoopUnrollResult::Unmodified;
   if (!Changed)
     return PreservedAnalyses::all();
@@ -1439,7 +1449,7 @@ PreservedAnalyses LoopUnrollPass::run(Function &F,
         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!
diff --git a/llvm/test/Transforms/LoopUnroll/disable-full-unroll-by-opt.ll b/llvm/test/Transforms/LoopUnroll/disable-full-unroll-by-opt.ll
new file mode 100644
index 00000000000..c89948ad52c
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/disable-full-unroll-by-opt.ll
@@ -0,0 +1,35 @@
+; 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
+}