[LoopPeeling] Add flag to disable support for peeling loops with non-latch exits

[asbirlea]

Summary

Add a flag to allow disabling the changes in
https://reviews.llvm.org/D134803.

[tra]

Thank you, Alina!

llvm/lib/Transforms/Utils/LoopPeel.cpp

977

Style nit: no need for {} around if/else bodies here.

[bgraur]

Thanks a lot Alina!

[nikic]

llvm/lib/Transforms/Utils/LoopPeel.cpp

112

Is it possible to reduce these conditions to the ones relevant to the affected test case? Like, I doubt that the check for BranchInst is relevant, and it would be good to know whether for the problematic case having a non-latch exit is the issue, or the structure of other exits.

528

Why is all this code needed? I would have expected the change to be just in canPeel(). For the cases where canPeel() previously returned true, the new implementation should behave the same apart from minor differences in branch weights.

Does that mean that you see issues if you limit peeling, but use the new branch weight updating code? That would be quite surprising, and would point towards something there being broken.

[nikic]

Generally fine with a temporary option, as long as the code duplication introduced here is not actually necessary.

[asbirlea]

Simplified to only the changes in canPeel.

[asbirlea]

Simplify conditions in canPeel.

llvm/lib/Transforms/Utils/LoopPeel.cpp

 static cl::opt<unsigned> UnrollPeelMaxCount(
     "unroll-peel-max-count", cl::init(7), cl::Hidden,
     cl::desc("Max average trip count which will cause loop peeling."));
 
 static cl::opt<unsigned> UnrollForcePeelCount(
     "unroll-force-peel-count", cl::init(0), cl::Hidden,
     cl::desc("Force a peel count regardless of profiling information."));
 
+static cl::opt<bool> DisableAdvancedPeeling(
+    "disable-advanced-peeling", cl::init(false), cl::Hidden,
+    cl::desc(
+        "Disable advance peeling. Issues for convergent targets (D134803)."));
+
 static const char *PeeledCountMetaData = "llvm.loop.peeled.count";
 
 // Check whether we are capable of peeling this loop.
 bool llvm::canPeel(Loop *L) {
   // Make sure the loop is in simplified form
-  return L->isLoopSimplifyForm();
+  if (!L->isLoopSimplifyForm())
+    return false;
+  if (!DisableAdvancedPeeling)
+    return true;
+
+  // Don't try to peel loops where the latch is not the exiting block.
+  // This can be an indication of two different things:
+  // 1) The loop is not rotated.
+  // 2) The loop contains irreducible control flow that involves the latch.
+  const BasicBlock *Latch = L->getLoopLatch();
+  if (!L->isLoopExiting(Latch))
+    return false;
+
+  // Peeling is only supported if the latch is a branch.
+  if (!isa<BranchInst>(Latch->getTerminator()))
+    return false;
+
+  SmallVector<BasicBlock *, 4> Exits;
+  L->getUniqueNonLatchExitBlocks(Exits);
+  // The latch must either be the only exiting block or all non-latch exit
+  // blocks have either a deopt or unreachable terminator or compose a chain of
+  // blocks where the last one is either deopt or unreachable terminated. Both
+  // deopt and unreachable terminators are a strong indication they are not
+  // taken. Note that this is a profitability check, not a legality check. Also
+  // note that LoopPeeling currently can only update the branch weights of latch
+  // blocks and branch weights to blocks with deopt or unreachable do not need
+  // updating.
+  return llvm::all_of(Exits, IsBlockFollowedByDeoptOrUnreachable);
 }
 
 // This function calculates the number of iterations after which the given Phi
 // becomes an invariant. The pre-calculated values are memorized in the map. The
 // function (shortcut is I) is calculated according to the following definition:
 // Given %x = phi <Inputs from above the loop>, ..., [%y, %back.edge].
 //   If %y is a loop invariant, then I(%x) = 1.
 //   If %y is a Phi from the loop header, I(%x) = I(%y) + 1.
Term->setMetadata(LLVMContext::MD_prof,
                     MDB.createBranchWeights(Info.Weights));
   for (auto [Idx, SubWeight] : enumerate(Info.SubWeights))
     if (SubWeight != 0)
       Info.Weights[Idx] = Info.Weights[Idx] > SubWeight
                               ? Info.Weights[Idx] - SubWeight
                               : 1;
 }
 
+static void updateBranchWeightsLimited(BasicBlock *Header, BranchInst *LatchBR,
+                                       uint64_t ExitWeight,
+                                       uint64_t &FallThroughWeight) {
+  // FallThroughWeight is 0 means that there is no branch weights on original
+  // latch block or estimated trip count is zero.
+  if (!FallThroughWeight)
+    return;
+
+  unsigned HeaderIdx = (LatchBR->getSuccessor(0) == Header ? 0 : 1);
+  MDBuilder MDB(LatchBR->getContext());
+  MDNode *WeightNode =
+      HeaderIdx ? MDB.createBranchWeights(ExitWeight, FallThroughWeight)
+                : MDB.createBranchWeights(FallThroughWeight, ExitWeight);
+  LatchBR->setMetadata(LLVMContext::MD_prof, WeightNode);
+  FallThroughWeight =
+      FallThroughWeight > ExitWeight ? FallThroughWeight - ExitWeight : 1;
+}
+
 /// Initialize the weights for all exiting blocks.
 static void initBranchWeights(DenseMap<Instruction *, WeightInfo> &WeightInfos,
                               Loop *L) {
   SmallVector<BasicBlock *> ExitingBlocks;
   L->getExitingBlocks(ExitingBlocks);
   for (BasicBlock *ExitingBlock : ExitingBlocks) {
     Instruction *Term = ExitingBlock->getTerminator();
     SmallVector<uint32_t> Weights;
for (auto [Succ, Weight] : zip(successors(Term), Weights)) {
       double W = (double)Weight / (double)FallThroughWeights;
       SubWeights.push_back((uint32_t)(ExitWeights * W));
     }
 
     WeightInfos.insert({Term, {std::move(Weights), std::move(SubWeights)}});
   }
 }
 
+static void initBranchWeightsLimited(BasicBlock *Header, BranchInst *LatchBR,
+                                     uint64_t &ExitWeight,
+                                     uint64_t &FallThroughWeight) {
+  uint64_t TrueWeight, FalseWeight;
+  if (!extractBranchWeights(*LatchBR, TrueWeight, FalseWeight))
+    return;
+  unsigned HeaderIdx = LatchBR->getSuccessor(0) == Header ? 0 : 1;
+  ExitWeight = HeaderIdx ? TrueWeight : FalseWeight;
+  FallThroughWeight = HeaderIdx ? FalseWeight : TrueWeight;
+}
+
 /// Update the weights of original exiting block after peeling off all
 /// iterations.
 static void fixupBranchWeights(Instruction *Term, const WeightInfo &Info) {
   MDBuilder MDB(Term->getContext());
   Term->setMetadata(LLVMContext::MD_prof,
                     MDB.createBranchWeights(Info.Weights));
 }
 
+static void fixupBranchWeightsLimited(BasicBlock *Header, BranchInst *LatchBR,
+                                      uint64_t ExitWeight,
+                                      uint64_t FallThroughWeight) {
+  // FallThroughWeight is 0 means that there is no branch weights on original
+  // latch block or estimated trip count is zero.
+  if (!FallThroughWeight)
+    return;
+
+  // Sets the branch weights on the loop exit.
+  MDBuilder MDB(LatchBR->getContext());
+  unsigned HeaderIdx = LatchBR->getSuccessor(0) == Header ? 0 : 1;
+  MDNode *WeightNode =
+      HeaderIdx ? MDB.createBranchWeights(ExitWeight, FallThroughWeight)
+                : MDB.createBranchWeights(FallThroughWeight, ExitWeight);
+  LatchBR->setMetadata(LLVMContext::MD_prof, WeightNode);
+}
+
 /// Clones the body of the loop L, putting it between \p InsertTop and \p
 /// InsertBot.
 /// \param IterNumber The serial number of the iteration currently being
 /// peeled off.
 /// \param ExitEdges The exit edges of the original loop.
 /// \param[out] NewBlocks A list of the blocks in the newly created clone
 /// \param[out] VMap The value map between the loop and the new clone.
 /// \param LoopBlocks A helper for DFS-traversal of the loop.
bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
   ValueToValueMapTy LVMap;
 
   Instruction *LatchTerm =
       cast<Instruction>(cast<BasicBlock>(Latch)->getTerminator());
 
   // If we have branch weight information, we'll want to update it for the
   // newly created branches.
   DenseMap<Instruction *, WeightInfo> Weights;
+  uint64_t ExitWeight = 0, FallThroughWeight = 0;
+  if (DisableAdvancedPeeling)
+    initBranchWeightsLimited(Header, cast<BranchInst>(LatchTerm), ExitWeight,
+                             FallThroughWeight);
+  else
     initBranchWeights(Weights, L);
 
   // Identify what noalias metadata is inside the loop: if it is inside the
   // loop, the associated metadata must be cloned for each iteration.
   SmallVector<MDNode *, 6> LoopLocalNoAliasDeclScopes;
   identifyNoAliasScopesToClone(L->getBlocks(), LoopLocalNoAliasDeclScopes);
 
   // For each peeled-off iteration, make a copy of the loop.
   for (unsigned Iter = 0; Iter < PeelCount; ++Iter) {
for (unsigned Iter = 0; Iter < PeelCount; ++Iter) {
     if (Iter == 0)
       for (auto BBIDom : NonLoopBlocksIDom)
         DT.changeImmediateDominator(BBIDom.first,
                                      cast<BasicBlock>(LVMap[BBIDom.second]));
 #ifdef EXPENSIVE_CHECKS
     assert(DT.verify(DominatorTree::VerificationLevel::Fast));
 #endif
 
+    if (DisableAdvancedPeeling) {
+      auto *LatchBRCopy = cast<BranchInst>(VMap[cast<BranchInst>(LatchTerm)]);
+      updateBranchWeightsLimited(InsertBot, LatchBRCopy, ExitWeight,
+                                 FallThroughWeight);
+
+      // Remove Loop metadata from the latch branch instruction
+      // because it is not the Loop's latch branch anymore.
+      LatchBRCopy->setMetadata(LLVMContext::MD_loop, nullptr);
+    } else {
       for (auto &[Term, Info] : Weights) {
         auto *TermCopy = cast<Instruction>(VMap[Term]);
         updateBranchWeights(TermCopy, Info);
       }
 
       // Remove Loop metadata from the latch branch instruction
       // because it is not the Loop's latch branch anymore.
       auto *LatchTermCopy = cast<Instruction>(VMap[LatchTerm]);
       LatchTermCopy->setMetadata(LLVMContext::MD_loop, nullptr);
+    }
 
     InsertTop = InsertBot;
     InsertBot = SplitBlock(InsertBot, InsertBot->getTerminator(), &DT, LI);
     InsertBot->setName(Header->getName() + ".peel.next");
 
     F->getBasicBlockList().splice(InsertTop->getIterator(),
                                   F->getBasicBlockList(),
                                   NewBlocks[0]->getIterator(), F->end());
   }
 
   // Now adjust the phi nodes in the loop header to get their initial values
   // from the last peeled-off iteration instead of the preheader.
   for (BasicBlock::iterator I = Header->begin(); isa<PHINode>(I); ++I) {
     PHINode *PHI = cast<PHINode>(I);
     Value *NewVal = PHI->getIncomingValueForBlock(Latch);
     Instruction *LatchInst = dyn_cast<Instruction>(NewVal);
     if (LatchInst && L->contains(LatchInst))
       NewVal = LVMap[LatchInst];
 
     PHI->setIncomingValueForBlock(NewPreHeader, NewVal);
   }
 
+  if (DisableAdvancedPeeling) {
+    fixupBranchWeightsLimited(Header, cast<BranchInst>(LatchTerm), ExitWeight,
+                              FallThroughWeight);
+  } else {
     for (const auto &[Term, Info] : Weights)
       fixupBranchWeights(Term, Info);
+  }
 
   // Update Metadata for count of peeled off iterations.
   unsigned AlreadyPeeled = 0;
   if (auto Peeled = getOptionalIntLoopAttribute(L, PeeledCountMetaData))
     AlreadyPeeled = *Peeled;
   addStringMetadataToLoop(L, PeeledCountMetaData, AlreadyPeeled + PeelCount);
 
   if (Loop *ParentLoop = L->getParentLoop())

llvm/test/Transforms/LoopUnroll/peel-branch-weights.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals
 ; RUN: opt < %s -S -loop-unroll -unroll-force-peel-count=2 2>&1 | FileCheck %s
+; RUN: opt < %s -S -loop-unroll -unroll-force-peel-count=2 -disable-advanced-peeling 2>&1 | FileCheck %s --check-prefix=DISABLEADV
 
 declare i32 @get.x()
 
 ; Test branch weight update for terminator with multiple fallthrough and
 ; multiple exit edges.
 define void @test() {
 ; CHECK-LABEL: @test(
 ; CHECK-NEXT:  entry:

 ; CHECK-NEXT:    i32 2, label [[LOOP_EXIT_LOOPEXIT]]
 ; CHECK-NEXT:    ], !prof [[PROF2:![0-9]+]]
 ; CHECK:       loop.latch:
 ; CHECK-NEXT:    br label [[LOOP]], !llvm.loop [[LOOP3:![0-9]+]]
 ; CHECK:       loop.exit.loopexit:
 ; CHECK-NEXT:    br label [[LOOP_EXIT]]
 ; CHECK:       loop.exit:
 ; CHECK-NEXT:    ret void
-;
+
+; DISABLEADV-LABEL: @test()
+; DISABLEADV-NEXT: entry:
+; DISABLEADV-NEXT:  br label %loop
+; DISABLEADV: loop
+; DISABLEADV-NEXT:  %x = call i32 @get.x()
+; DISABLEADV-NEXT:  switch i32 %x, label %loop.latch [
+; DISABLEADV-NEXT:    i32 0, label %loop.latch
+; DISABLEADV-NEXT:    i32 1, label %loop.exit
+; DISABLEADV-NEXT:    i32 2, label %loop.exit
+; DISABLEADV-NEXT:  ], !prof !0
+; DISABLEADV: loop.latch:
+; DISABLEADV-NEXT:  br label %loop
+; DISABLEADV: loop.exit:
+; DISABLEADV-NEXT:  ret void
+
 entry:
   br label %loop
 
 loop:
   %x = call i32 @get.x()
   switch i32 %x, label %loop.latch [
   i32 0, label %loop.latch
   i32 1, label %loop.exit