[LoopUnrollRuntime] Avoid high-cost trip count computation.

[sanjoy]

Summary

Runtime unrolling of loops needs to emit an expression to compute the
loop's runtime trip-count. Avoid runtime unrolling if this computation
will be expensive.

Depends on D8993.

[hfinkel]

You'll need to add a field in the UnrollingPreferences structure to control this. The makes a lot of sense on X86, where we only unroll small loops to fill the LSD's dispatch buffer (and they're still pretty small even after unrolling). On other targets we unroll to much larger sizes and the cost of a division is amortized over many more instructions and iterations.

[sanjoy]

I've added AllowExpensiveTripCount to UnrollPreferences. I don't
have a good understanding for when AllowExpensiveTripCount should be
true, so I've left it false always.

• address Hal's review: add UP.AllowExpensiveTripCount, off by default

[atrick]

LGTM

include/llvm/Analysis/TargetTransformInfo.h

struct UnrollingPreferences {
     unsigned MaxCount;
     /// Allow partial unrolling (unrolling of loops to expand the size of the
     /// loop body, not only to eliminate small constant-trip-count loops).
     bool Partial;
     /// Allow runtime unrolling (unrolling of loops to expand the size of the
     /// loop body even when the number of loop iterations is not known at
     /// compile time).
     bool Runtime;
+    /// Allow emitting expensive instructions (such as divisions) when computing
+    /// the trip count of a loop for runtime unrolling.
+    bool AllowExpensiveTripCount;
   };
 
   /// \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, UnrollingPreferences &UP) const;
 
   /// @}

include/llvm/Transforms/Utils/UnrollLoop.h

 class AssumptionCache;
 class Loop;
 class LoopInfo;
 class LPPassManager;
 class MDNode;
 class Pass;
 
 bool UnrollLoop(Loop *L, unsigned Count, unsigned TripCount, bool AllowRuntime,
-                unsigned TripMultiple, LoopInfo *LI, Pass *PP,
-                LPPassManager *LPM, AssumptionCache *AC);
+                bool AllowExpensiveTripCount, unsigned TripMultiple,
+                LoopInfo *LI, Pass *PP, LPPassManager *LPM,
+                AssumptionCache *AC);
 
-bool UnrollRuntimeLoopProlog(Loop *L, unsigned Count, LoopInfo *LI,
-                             LPPassManager* LPM);
+bool UnrollRuntimeLoopProlog(Loop *L, unsigned Count,
+                             bool AllowExpensiveTripCount, LoopInfo *LI,
+                             LPPassManager *LPM);
 
 MDNode *GetUnrollMetadata(MDNode *LoopID, StringRef Name);
 }
 
 #endif

lib/Transforms/Scalar/LoopUnrollPass.cpp

void getUnrollingPreferences(Loop *L, const TargetTransformInfo &TTI,
       UP.MinPercentOfOptimized = CurrentMinPercentOfOptimized;
       UP.OptSizeThreshold = OptSizeUnrollThreshold;
       UP.PartialThreshold = CurrentThreshold;
       UP.PartialOptSizeThreshold = OptSizeUnrollThreshold;
       UP.Count = CurrentCount;
       UP.MaxCount = UINT_MAX;
       UP.Partial = CurrentAllowPartial;
       UP.Runtime = CurrentRuntime;
+      UP.AllowExpensiveTripCount = false;
       TTI.getUnrollingPreferences(L, UP);
     }
 
     // Select and return an unroll count based on parameters from
     // user, unroll preferences, unroll pragmas, or a heuristic.
     // SetExplicitly is set to true if the unroll count is is set by
     // the user or a pragma rather than selected heuristically.
     unsigned
bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
 
   if (Unrolling != Full && Count < 2) {
     // Partial unrolling by 1 is a nop.  For full unrolling, a factor
     // of 1 makes sense because loop control can be eliminated.
     return false;
   }
 
   // Unroll the loop.
-  if (!UnrollLoop(L, Count, TripCount, AllowRuntime, TripMultiple, LI, this,
-                  &LPM, &AC))
+  if (!UnrollLoop(L, Count, TripCount, AllowRuntime, UP.AllowExpensiveTripCount,
+                  TripMultiple, LI, this, &LPM, &AC))
     return false;
 
   return true;
 }

lib/Transforms/Utils/LoopUnroll.cpp

 /// terminates LatchBlock in order to remove unnecesssary instances of the
 /// test. In other words, control may exit the loop prior to TripCount
 /// iterations via an early branch, but control may not exit the loop from the
 /// LatchBlock's terminator prior to TripCount iterations.
 ///
 /// Similarly, TripMultiple divides the number of times that the LatchBlock may
 /// execute without exiting the loop.
 ///
+/// If AllowRuntime is true then UnrollLoop will consider unrolling loops that
+/// have a runtime (i.e. not compile time constant) trip count.  Unrolling these
+/// loops require a unroll "prologue" that runs "RuntimeTripCount % Count"
+/// iterations before branching into the unrolled loop.  UnrollLoop will not
+/// runtime-unroll the loop if computing RuntimeTripCount will be expensive and
+/// AllowExpensiveTripCount is false.
+///
 /// The LoopInfo Analysis that is passed will be kept consistent.
 ///
 /// If a LoopPassManager is passed in, and the loop is fully removed, it will be
 /// removed from the LoopPassManager as well. LPM can also be NULL.
 ///
 /// This utility preserves LoopInfo. If DominatorTree or ScalarEvolution are
 /// available from the Pass it must also preserve those analyses.
 bool llvm::UnrollLoop(Loop *L, unsigned Count, unsigned TripCount,
-                      bool AllowRuntime, unsigned TripMultiple, LoopInfo *LI,
-                      Pass *PP, LPPassManager *LPM, AssumptionCache *AC) {
+                      bool AllowRuntime, bool AllowExpensiveTripCount,
+                      unsigned TripMultiple, LoopInfo *LI, Pass *PP,
+                      LPPassManager *LPM, AssumptionCache *AC) {
   BasicBlock *Preheader = L->getLoopPreheader();
   if (!Preheader) {
     DEBUG(dbgs() << "  Can't unroll; loop preheader-insertion failed.\n");
     return false;
   }
 
   BasicBlock *LatchBlock = L->getLoopLatch();
   if (!LatchBlock) {
bool llvm::UnrollLoop(Loop *L, unsigned Count, unsigned TripCount,
   // Are we eliminating the loop control altogether?
   bool CompletelyUnroll = Count == TripCount;
 
   // We assume a run-time trip count if the compiler cannot
   // figure out the loop trip count and the unroll-runtime
   // flag is specified.
   bool RuntimeTripCount = (TripCount == 0 && Count > 0 && AllowRuntime);
 
-  if (RuntimeTripCount && !UnrollRuntimeLoopProlog(L, Count, LI, LPM))
+  if (RuntimeTripCount &&
+      !UnrollRuntimeLoopProlog(L, Count, AllowExpensiveTripCount, LI, LPM))
     return false;
 
   // Notify ScalarEvolution that the loop will be substantially changed,
   // if not outright eliminated.
   ScalarEvolution *SE =
       PP ? PP->getAnalysisIfAvailable<ScalarEvolution>() : nullptr;
   if (SE)
     SE->forgetLoop(L);

lib/Transforms/Utils/LoopUnrollRuntime.cpp

 /// Prol:  LoopBody;
 ///        extraiters -= 1                 // Omitted if unroll factor is 2.
 ///        if (extraiters != 0) jump Prol: // Omitted if unroll factor is 2.
 ///        if (tripcount < loopfactor) jump End
 /// Loop:
 /// ...
 /// End:
 ///
-bool llvm::UnrollRuntimeLoopProlog(Loop *L, unsigned Count, LoopInfo *LI,
+bool llvm::UnrollRuntimeLoopProlog(Loop *L, unsigned Count,
+                                   bool AllowExpensiveTripCount, LoopInfo *LI,
                                    LPPassManager *LPM) {
   // for now, only unroll loops that contain a single exit
   if (!L->getExitingBlock())
     return false;
 
   // Make sure the loop is in canonical form, and there is a single
   // exit block only.
   if (!L->isLoopSimplifyForm() || !L->getUniqueExitBlock())
bool llvm::UnrollRuntimeLoopProlog(Loop *L, unsigned Count,
   unsigned BEWidth = cast<IntegerType>(BECountSC->getType())->getBitWidth();
 
   // Add 1 since the backedge count doesn't include the first loop iteration
   const SCEV *TripCountSC =
     SE->getAddExpr(BECountSC, SE->getConstant(BECountSC->getType(), 1));
   if (isa<SCEVCouldNotCompute>(TripCountSC))
     return false;
 
+  BasicBlock *Header = L->getHeader();
+  const DataLayout &DL = Header->getModule()->getDataLayout();
+  SCEVExpander Expander(*SE, DL, "loop-unroll");
+  if (!AllowExpensiveTripCount && Expander.isHighCostExpansion(TripCountSC, L))
+    return false;
+
   // We only handle cases when the unroll factor is a power of 2.
   // Count is the loop unroll factor, the number of extra copies added + 1.
   if (!isPowerOf2_32(Count))
     return false;
 
   // This constraint lets us deal with an overflowing trip count easily; see the
   // comment on ModVal below.
   if (Log2_32(Count) > BEWidth)
     return false;
 
   // If this loop is nested, then the loop unroller changes the code in
   // parent loop, so the Scalar Evolution pass needs to be run again
   if (Loop *ParentLoop = L->getParentLoop())
     SE->forgetLoop(ParentLoop);
 
   // Grab analyses that we preserve.
   auto *DTWP = LPM->getAnalysisIfAvailable<DominatorTreeWrapperPass>();
   auto *DT = DTWP ? &DTWP->getDomTree() : nullptr;
 
   BasicBlock *PH = L->getLoopPreheader();
-  BasicBlock *Header = L->getHeader();
   BasicBlock *Latch = L->getLoopLatch();
   // It helps to splits the original preheader twice, one for the end of the
   // prolog code and one for a new loop preheader
   BasicBlock *PEnd = SplitEdge(PH, Header, DT, LI);
   BasicBlock *NewPH = SplitBlock(PEnd, PEnd->getTerminator(), DT, LI);
   BranchInst *PreHeaderBR = cast<BranchInst>(PH->getTerminator());
-  const DataLayout &DL = Header->getModule()->getDataLayout();
 
   // Compute the number of extra iterations required, which is:
   //  extra iterations = run-time trip count % (loop unroll factor + 1)
-  SCEVExpander Expander(*SE, DL, "loop-unroll");
   Value *TripCount = Expander.expandCodeFor(TripCountSC, TripCountSC->getType(),
                                             PreHeaderBR);
   Value *BECount = Expander.expandCodeFor(BECountSC, BECountSC->getType(),
                                           PreHeaderBR);
 
   IRBuilder<> B(PreHeaderBR);
   Value *ModVal = B.CreateAnd(TripCount, Count - 1, "xtraiter");
 

test/Transforms/LoopUnroll/high-cost-trip-count-computation.ll

+; RUN: opt -S -unroll-runtime -loop-unroll < %s | FileCheck %s
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+
+;; Check that we don't emit expensive instructions to compute trip
+;; counts when unrolling loops.
+
+define i32 @test(i64 %v12, i8* %array, i64* %loc) {
+; CHECK-LABEL: @test(
+; CHECK-NOT: udiv
+entry:
+  %step = load i64, i64* %loc, !range !0
+  br label %loop
+
+loop:                                           ; preds = %entry, %loop
+  %k.015 = phi i64 [ %v15, %loop ], [ %v12, %entry ]
+  %v14 = getelementptr inbounds i8, i8* %array, i64 %k.015
+  store i8 0, i8* %v14
+  %v15 = add nuw nsw i64 %k.015, %step
+  %v16 = icmp slt i64 %v15, 8193
+  br i1 %v16, label %loop, label %loopexit
+
+loopexit:                             ; preds = %loop
+  ret i32 0
+}
+
+!0 = !{i64 1, i64 100}

test/Transforms/LoopUnroll/runtime-loop.ll

 ; RUN: opt < %s -S -loop-unroll -unroll-runtime=true | FileCheck %s
 
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+
 ; Tests for unrolling loops with run-time trip counts
 
 ; CHECK: %xtraiter = and i32 %n
 ; CHECK:  %lcmp.mod = icmp ne i32 %xtraiter, 0
 ; CHECK:  br i1 %lcmp.mod, label %for.body.prol, label %for.body.preheader.split
 
 ; CHECK: for.body.prol:
 ; CHECK: %indvars.iv.prol = phi i64 [ %indvars.iv.next.prol, %for.body.prol ], [ 0, %for.body.preheader ]

test/Transforms/LoopUnroll/runtime-loop4.ll

 
 loop1:
   %iv1 = phi i32 [ 0, %entry ], [ %inc1, %loop1.latch ]
   %offset1 = getelementptr i32, i32* %addr1, i32 %iv1
   store i32 %iv1, i32* %offset1, align 4
   br label %loop2.header
 
 loop2.header:
-  br label %loop2
+  %e = icmp uge i32 %iter, 1
+  br i1 %e, label %loop2, label %exit2
 
 loop2:
   %iv2 = phi i32 [ 0, %loop2.header ], [ %inc2, %loop2 ]
   %offset2 = getelementptr i32, i32* %addr2, i32 %iv2
   store i32 %iv2, i32* %offset2, align 4
   %inc2 = add i32 %iv2, 1
   %exitcnd2 = icmp uge i32 %inc2, %iter
   br i1 %exitcnd2, label %exit2, label %loop2