[LoopUnroll] Enable option to unroll remainder loop

[samparker]

Summary

On some targets, the penalty of executing runtime unrolling checks and then not the unrolled loop can be significantly detrimental to performance. This results in the need to be more conservative with the unroll count, keeping a trip count of 2 reduces the overhead as well as increasing the chance of the unrolled body being executed. But being conservative leaves performance gains on the table.

This patch enables the unrolling of the remainder loop introduced by runtime unrolling. This can help reduce the overhead of mis-unrolled loops because the cost of non-taken branches is much less than the cost of the backedge that would normally be executed in the remainder loop. This allows larger unroll factors to be used without suffering performance loses with smaller iteration counts.

[efriedma]

This makes sense.

lib/Transforms/Utils/LoopUnrollRuntime.cpp

929

I think I would write this using UnrollLoop: full unrolling will result in a slightly simpler CFG because the backedge goes away. (Granted, later optimizations will probably clean it up either way.)

[samparker]

lib/Transforms/Utils/LoopUnrollRuntime.cpp

929

That's not possible here though. As we only know the maximum trip count, we'll need to maintain comparisons for each iteration and convert the backedges to early exits.

[samparker]

Since we know that the maximum number of iterations have been peeled off the loop, I've added a function that removes the remainder loop body. Also now added a test.

[efriedma]

LoopUnroll has been able to unroll using an upper bound since https://reviews.llvm.org/rL284053 .

[samparker]

Ah ok, great. I'll have a play with that then. Thanks

[samparker]

Now unrolling the loop instead of peeling.

[efriedma]

lib/Transforms/Utils/LoopUnrollRuntime.cpp

930

It's difficult to figure out what you're actually passing here; maybe include the parameter names in comments, like UnrollLoop(remainderLoop, /*Count*/Count - 1, /*TripCount*/Count - 1, /*Force*/ false, etc.

test/Transforms/LoopUnroll/runtime-unroll-remainder.ll

1 ↗

(On Diff #110180)

I'd prefer to see the result without the simplifycfg run, to see what the generated IR actually looks like.

[samparker]

Added useful comments for call arguments and modified the test not to use simplifycfg.

[efriedma]

LGTM.

include/llvm/Transforms/Utils/UnrollLoop.h

46

Weird line wrapping.

lib/Transforms/Scalar/LoopUnrollPass.cpp

948

Please commit whitespace change separately.

include/llvm/Analysis/TargetTransformInfo.h

     bool AllowExpensiveTripCount;
     /// Apply loop unroll on any kind of loop
     /// (mainly to loops that fail runtime unrolling).
     bool Force;
     /// Allow using trip count upper bound to unroll loops.
     bool UpperBound;
     /// Allow peeling off loop iterations for loops with low dynamic tripcount.
     bool AllowPeeling;
+    /// Allow peeling off all the iterations of the runtime loop remainder.
+    bool PeelRemainder;
   };
 
   /// \brief Get target-customized preferences for the generic loop unrolling
   /// transformation. The caller will initialize UP with the current
   /// target-independent defaults.
   void getUnrollingPreferences(Loop *L, ScalarEvolution &,
                                UnrollingPreferences &UP) const;
 

include/llvm/Transforms/Utils/UnrollLoop.h

 
 const Loop* addClonedBlockToLoopInfo(BasicBlock *OriginalBB,
                                      BasicBlock *ClonedBB, LoopInfo *LI,
                                      NewLoopsMap &NewLoops);
 
 bool UnrollLoop(Loop *L, unsigned Count, unsigned TripCount, bool Force,
                 bool AllowRuntime, bool AllowExpensiveTripCount,
                 bool PreserveCondBr, bool PreserveOnlyFirst,
-                unsigned TripMultiple, unsigned PeelCount, LoopInfo *LI,
+                unsigned TripMultiple, unsigned PeelCount, bool PeelRemainder,
+                LoopInfo *LI,
                 ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC,
                 OptimizationRemarkEmitter *ORE, bool PreserveLCSSA);
 
 bool UnrollRuntimeLoopRemainder(Loop *L, unsigned Count,
                                 bool AllowExpensiveTripCount,
-                                bool UseEpilogRemainder, LoopInfo *LI,
+                                bool UseEpilogRemainder, bool PeelRemainder,
+                                LoopInfo *LI,
                                 ScalarEvolution *SE, DominatorTree *DT,
+                                AssumptionCache *AC,
                                 bool PreserveLCSSA);
 
 void computePeelCount(Loop *L, unsigned LoopSize,
                       TargetTransformInfo::UnrollingPreferences &UP,
                       unsigned &TripCount);
 
 bool peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI, ScalarEvolution *SE,
               DominatorTree *DT, AssumptionCache *AC, bool PreserveLCSSA);
 
 MDNode *GetUnrollMetadata(MDNode *LoopID, StringRef Name);
 }
 
 #endif

lib/Transforms/Scalar/LoopUnrollPass.cpp

              "threshold, the loop is considered as flat and will be less "
              "aggressively unrolled."));
 
 static cl::opt<bool>
     UnrollAllowPeeling("unroll-allow-peeling", cl::init(true), cl::Hidden,
                        cl::desc("Allows loops to be peeled when the dynamic "
                                 "trip count is known to be low."));
 
+static cl::opt<bool> UnrollPeelRemainder(
+  "peel-remainder", cl::Hidden,
+  cl::desc("Allow the loop remainder to be peeled."));
+
 // This option isn't ever intended to be enabled, it serves to allow
 // experiments to check the assumptions about when this kind of revisit is
 // necessary.
 static cl::opt<bool> UnrollRevisitChildLoops(
     "unroll-revisit-child-loops", cl::Hidden,
     cl::desc("Enqueue and re-visit child loops in the loop PM after unrolling. "
              "This shouldn't typically be needed as child loops (or their "
              "clones) were already visited."));

   UP.PeelCount = 0;
   UP.DefaultUnrollRuntimeCount = 8;
   UP.MaxCount = UINT_MAX;
   UP.FullUnrollMaxCount = UINT_MAX;
   UP.BEInsns = 2;
   UP.Partial = false;
   UP.Runtime = false;
   UP.AllowRemainder = true;
+  UP.PeelRemainder = false;
   UP.AllowExpensiveTripCount = false;
   UP.Force = false;
   UP.UpperBound = false;
   UP.AllowPeeling = true;
 
   // Override with any target specific settings
   TTI.getUnrollingPreferences(L, SE, UP);
 

   if (UnrollAllowRemainder.getNumOccurrences() > 0)
     UP.AllowRemainder = UnrollAllowRemainder;
   if (UnrollRuntime.getNumOccurrences() > 0)
     UP.Runtime = UnrollRuntime;
   if (UnrollMaxUpperBound == 0)
     UP.UpperBound = false;
   if (UnrollAllowPeeling.getNumOccurrences() > 0)
     UP.AllowPeeling = UnrollAllowPeeling;
+  if (UnrollPeelRemainder.getNumOccurrences() > 0)
+    UP.PeelRemainder = UnrollPeelRemainder;
 
   // Apply user values provided by argument
   if (UserThreshold.hasValue()) {
     UP.Threshold = *UserThreshold;
     UP.PartialThreshold = *UserThreshold;
   }
   if (UserCount.hasValue())
     UP.Count = *UserCount;

                             bool PreserveLCSSA, int OptLevel,
                             Optional<unsigned> ProvidedCount,
                             Optional<unsigned> ProvidedThreshold,
                             Optional<bool> ProvidedAllowPartial,
                             Optional<bool> ProvidedRuntime,
                             Optional<bool> ProvidedUpperBound) {
   DEBUG(dbgs() << "Loop Unroll: F[" << L->getHeader()->getParent()->getName()
                << "] Loop %" << L->getHeader()->getName() << "\n");
-  if (HasUnrollDisablePragma(L)) 
+  if (HasUnrollDisablePragma(L))
     return false;
   if (!L->isLoopSimplifyForm()) { 
     DEBUG(
         dbgs() << "  Not unrolling loop which is not in loop-simplify form.\n");
     return false;
   }
 
   unsigned NumInlineCandidates;

     return false;
   // Unroll factor (Count) must be less or equal to TripCount.
   if (TripCount && UP.Count > TripCount)
     UP.Count = TripCount;
 
   // Unroll the loop.
   if (!UnrollLoop(L, UP.Count, TripCount, UP.Force, UP.Runtime,
                   UP.AllowExpensiveTripCount, UseUpperBound, MaxOrZero,
-                  TripMultiple, UP.PeelCount, LI, &SE, &DT, &AC, &ORE,
+                  TripMultiple, UP.PeelCount, UP.PeelRemainder,
+                  LI, &SE, &DT, &AC, &ORE,
                   PreserveLCSSA))
     return false;
 
   // If loop has an unroll count pragma or unrolled by explicitly set count
   // mark loop as unrolled to prevent unrolling beyond that requested.
   // If the loop was peeled, we already "used up" the profile information
   // we had, so we don't want to unroll or peel again.
   if (IsCountSetExplicitly || UP.PeelCount)

lib/Transforms/Utils/LoopUnroll.cpp

 ///
 /// The LoopInfo Analysis that is passed will be kept consistent.
 ///
 /// This utility preserves LoopInfo. It will also preserve ScalarEvolution and
 /// DominatorTree if they are non-null.
 bool llvm::UnrollLoop(Loop *L, unsigned Count, unsigned TripCount, bool Force,
                       bool AllowRuntime, bool AllowExpensiveTripCount,
                       bool PreserveCondBr, bool PreserveOnlyFirst,
-                      unsigned TripMultiple, unsigned PeelCount, LoopInfo *LI,
+                      unsigned TripMultiple, unsigned PeelCount,
+                      bool PeelRemainder, LoopInfo *LI,
                       ScalarEvolution *SE, DominatorTree *DT,
                       AssumptionCache *AC, OptimizationRemarkEmitter *ORE,
                       bool PreserveLCSSA) {
 
   BasicBlock *Preheader = L->getLoopPreheader();
   if (!Preheader) {
     DEBUG(dbgs() << "  Can't unroll; loop preheader-insertion failed.\n");
     return false;

       });
 
   bool EpilogProfitability =
       UnrollRuntimeEpilog.getNumOccurrences() ? UnrollRuntimeEpilog
                                               : isEpilogProfitable(L);
 
   if (RuntimeTripCount && TripMultiple % Count != 0 &&
       !UnrollRuntimeLoopRemainder(L, Count, AllowExpensiveTripCount,
-                                  EpilogProfitability, LI, SE, DT,
+                                  EpilogProfitability, PeelRemainder,
+                                  LI, SE, DT, AC,
                                   PreserveLCSSA)) {
     if (Force)
       RuntimeTripCount = false;
     else {
       DEBUG(
           dbgs() << "Wont unroll; remainder loop could not be generated"
                     "when assuming runtime trip count\n");
       return false;

lib/Transforms/Utils/LoopUnrollRuntime.cpp

 /// otherwise a new loop will be created including all cloned blocks, and the
 /// iterator of it switches to count NewIter down to 0.
 /// The cloned blocks should be inserted between InsertTop and InsertBot.
 /// If loop structure is cloned InsertTop should be new preheader, InsertBot
 /// new loop exit.
 /// Return the new cloned loop that is created when CreateRemainderLoop is true.
 static Loop *
 CloneLoopBlocks(Loop *L, Value *NewIter, const bool CreateRemainderLoop,
-                const bool UseEpilogRemainder, BasicBlock *InsertTop,
+                const bool UseEpilogRemainder, const bool PeelRemainder,
+                BasicBlock *InsertTop,
                 BasicBlock *InsertBot, BasicBlock *Preheader,
                 std::vector<BasicBlock *> &NewBlocks, LoopBlocksDFS &LoopBlocks,
                 ValueToValueMapTy &VMap, DominatorTree *DT, LoopInfo *LI) {
   StringRef suffix = UseEpilogRemainder ? "epil" : "prol";
   BasicBlock *Header = L->getHeader();
   BasicBlock *Latch = L->getLoopLatch();
   Function *F = Header->getParent();
   LoopBlocksDFS::RPOIterator BlockBegin = LoopBlocks.beginRPO();

           IsUnrollMetadata = S && S->getString().startswith("llvm.loop.unroll.");
         }
         if (!IsUnrollMetadata)
           MDs.push_back(LoopID->getOperand(i));
       }
     }
 
     LLVMContext &Context = NewLoop->getHeader()->getContext();
-    SmallVector<Metadata *, 1> DisableOperands;
-    DisableOperands.push_back(MDString::get(Context, "llvm.loop.unroll.disable"));
-    MDNode *DisableNode = MDNode::get(Context, DisableOperands);
-    MDs.push_back(DisableNode);
+    if (!PeelRemainder) {
+      SmallVector<Metadata *, 1> DisableOperands;
+      DisableOperands.push_back(MDString::get(Context,
+                                              "llvm.loop.unroll.disable"));
+      MDNode *DisableNode = MDNode::get(Context, DisableOperands);
+      MDs.push_back(DisableNode);
+    }
 
     MDNode *NewLoopID = MDNode::get(Context, MDs);
     // Set operand 0 to refer to the loop id itself.
     NewLoopID->replaceOperandWith(0, NewLoopID);
     NewLoop->setLoopID(NewLoopID);
     return NewLoop;
   }
   else

                                       PreserveLCSSA, UseEpilogRemainder) &&
          "Should be safe to unroll before checking profitability!");
 #endif
   // Priority goes to UnrollRuntimeMultiExit if it's supplied.
   return UnrollRuntimeMultiExit.getNumOccurrences() ? UnrollRuntimeMultiExit
                                                     : false;
 }
 
+static void RemovePeeledLoop(Loop *L, LoopInfo *LI, DominatorTree *DT) {
+  // Loops that can be peeled are in their simple form, with a single block
+  // and exit.
+  BasicBlock *Preheader = L->getLoopPreheader();
+  auto *ExitBB = L->getExitingBlock();
+  auto *Term = ExitBB->getTerminator();
+  BasicBlock *Dest = nullptr;
+
+  // Find the successor block.
+  for (unsigned i = 0; i < 2; ++i) {
+    BasicBlock *Succ = Term->getSuccessor(i);
+    if (Succ != ExitBB) {
+      Dest = Succ;
+      break;
+    }
+  }
+
+  // Remove the loop body from the CFG and dominator tree.
+  Dest->removePredecessor(ExitBB);
+  ExitBB->getTerminator()->eraseFromParent();
+  DT->deleteEdge(ExitBB, Dest);
+
+  ExitBB->removePredecessor(Preheader);
+  Preheader->getTerminator()->eraseFromParent();
+  DT->deleteEdge(Preheader, ExitBB);
+
+  // Have the preheader branch directly to the loop's successor.
+  BranchInst::Create(Dest, Preheader);
+  DT->insertEdge(Preheader, Dest);
+
+  ExitBB->dropAllReferences();
+  ExitBB->removeFromParent();
+  LI->markAsRemoved(L);
+}
+
 /// Insert code in the prolog/epilog code when unrolling a loop with a
 /// run-time trip-count.
 ///
 /// This method assumes that the loop unroll factor is total number
 /// of loop bodies in the loop after unrolling. (Some folks refer
 /// to the unroll factor as the number of *extra* copies added).
 /// We assume also that the loop unroll factor is a power-of-two. So, after
 /// unrolling the loop, the number of loop bodies executed is 2,

 /// Epil:  LoopBody; (executes extraiters times)
 ///        extraiters -= 1                 // Omitted if unroll factor is 2.
 ///        if (extraiters != 0) jump Epil: // Omitted if unroll factor is 2.
 /// EpilExit:
 
 bool llvm::UnrollRuntimeLoopRemainder(Loop *L, unsigned Count,
                                       bool AllowExpensiveTripCount,
                                       bool UseEpilogRemainder,
+                                      bool PeelRemainder,
                                       LoopInfo *LI, ScalarEvolution *SE,
-                                      DominatorTree *DT, bool PreserveLCSSA) {
+                                      DominatorTree *DT, AssumptionCache *AC,
+                                      bool PreserveLCSSA) {
   DEBUG(dbgs() << "Trying runtime unrolling on Loop: \n");
   DEBUG(L->dump());
 
   // Make sure the loop is in canonical form.
   if (!L->isLoopSimplifyForm()) {
     DEBUG(dbgs() << "Not in simplify form!\n");
     return false;
   }

   bool CreateRemainderLoop = (Count != 2);
 
   // Clone all the basic blocks in the loop. If Count is 2, we don't clone
   // the loop, otherwise we create a cloned loop to execute the extra
   // iterations. This function adds the appropriate CFG connections.
   BasicBlock *InsertBot = UseEpilogRemainder ? LatchExit : PrologExit;
   BasicBlock *InsertTop = UseEpilogRemainder ? EpilogPreHeader : PrologPreHeader;
   Loop *remainderLoop = CloneLoopBlocks(
-      L, ModVal, CreateRemainderLoop, UseEpilogRemainder, InsertTop, InsertBot,
+      L, ModVal, CreateRemainderLoop, UseEpilogRemainder, PeelRemainder,
+      InsertTop, InsertBot,
       NewPreHeader, NewBlocks, LoopBlocks, VMap, DT, LI);
 
   // Insert the cloned blocks into the function.
   F->getBasicBlockList().splice(InsertBot->getIterator(),
                                 F->getBasicBlockList(),
                                 NewBlocks[0]->getIterator(),
                                 F->end());
 

     // LoopSimplifyForm.
     formDedicatedExitBlocks(L, DT, LI, PreserveLCSSA);
     // Generate dedicated exit blocks for the remainder loop if one exists, to
     // preserve LoopSimplifyForm.
     if (remainderLoop)
       formDedicatedExitBlocks(remainderLoop, DT, LI, PreserveLCSSA);
   }
 
+  if (remainderLoop && PeelRemainder) {
+    if (peelLoop(remainderLoop, Count - 1, LI, SE, DT, AC, PreserveLCSSA))
+      RemovePeeledLoop(remainderLoop, LI, DT);
+  }
+
   NumRuntimeUnrolled++;
   return true;
 }

test/Transforms/LoopUnroll/runtime-peel-remainder.ll

+; RUN: opt < %s -S -loop-unroll -unroll-runtime=true -unroll-count=4 -peel-remainder -simplifycfg -instcombine | FileCheck %s
+
+; CHECK-LABEL: unroll
+define i32 @unroll(i32* nocapture readonly %a, i32* nocapture readonly %b, i32 %N) local_unnamed_addr #0 {
+entry:
+  %cmp9 = icmp eq i32 %N, 0
+  br i1 %cmp9, label %for.cond.cleanup, label %for.body.lr.ph
+
+for.body.lr.ph:
+  %wide.trip.count = zext i32 %N to i64
+  br label %for.body
+
+for.cond.cleanup:
+  %c.0.lcssa = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  ret i32 %c.0.lcssa
+
+; CHECK-LABEL: for.body.lr.ph
+; CHECK: [[COUNT:%[a-z.0-9]+]] = add nsw i64 %wide.trip.count, -1
+; CHECK: %xtraiter = and i64 %wide.trip.count, 3
+; CHECK: [[CMP:%[a-z.0-9]+]] = icmp ult i64 [[COUNT]], 3
+; CHECK: br i1 [[CMP]], label %[[CLEANUP:.*]], label %for.body.lr.ph.new
+
+; CHECK-LABEL: for.body.lr.ph.new:
+; CHECK: sub nsw i64 %wide.trip.count, %xtraiter
+; CHECK: br label %for.body
+
+; CHECK: [[CLEANUP]]:
+; CHECK: [[MOD:%[a-z.0-9]+]] = icmp eq i64 %xtraiter, 0
+; CHECK: br i1 [[MOD]], label %[[EXIT:.*]], label %[[EPIL_PEEL0:.*]]
+
+; CHECK: [[EPIL_PEEL0]]:
+; CHECK: [[PEEL_CMP0:%[a-z.0-9]+]] = icmp eq i64 %xtraiter, 1
+; CHECK: br i1 [[PEEL_CMP0]], label %[[EXIT]], label %[[EPIL_PEEL1:.*]],
+
+; CHECK: [[EPIL_PEEL1]]:
+; CHECK: [[PEEL_CMP1:%[a-z.0-9]+]] = icmp eq i64 %xtraiter, 2
+; CHECK: br i1 [[PEEL_CMP1]], label %[[EXIT]], label %[[EPIL_PEEL2:.*]],
+
+; CHECK: [[EPIL_PEEL2]]:
+; CHECK: ret i32
+
+; CHECK: [[EXIT]]:
+; CHECK: ret i32
+
+; CHECK-LABEL: for.body:
+; CHECK: [[INDVAR0:%[a-z.0-9]+]] = phi i64 [ 0, %for.body.lr.ph
+; CHECK: [[ITER:%[a-z.0-9]+]] = phi i64 [ %unroll_iter
+; CHECK: or i64 [[INDVAR0]], 1
+; CHECK: or i64 [[INDVAR0]], 2
+; CHECK: or i64 [[INDVAR0]], 3
+; CHECK: add nsw i64 [[INDVAR0]], 4
+; CHECK: [[SUB:%[a-z.0-9]+]] = add i64 [[ITER]], -4
+; CHECK: [[ITER_CMP:%[a-z.0-9]+]] = icmp eq i64 [[SUB]]
+; CHECK: br i1 [[ITER_CMP]], label %[[CLEANUP]], label %for.body
+
+for.body:
+  %indvars.iv = phi i64 [ 0, %for.body.lr.ph ], [ %indvars.iv.next, %for.body ]
+  %c.010 = phi i32 [ 0, %for.body.lr.ph ], [ %add, %for.body ]
+  %arrayidx = getelementptr inbounds i32, i32* %a, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4
+  %arrayidx2 = getelementptr inbounds i32, i32* %b, i64 %indvars.iv
+  %1 = load i32, i32* %arrayidx2, align 4
+  %mul = mul nsw i32 %1, %0
+  %add = add nsw i32 %mul, %c.010
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond = icmp eq i64 %indvars.iv.next, %wide.trip.count
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+}