[LoopPredication] Calculate profitability without BPI

[anna]

Summary

Using BPI within loop predication is non-trivial because BPI is only
preserved lossily in loop pass manager (one fix exposed by lossy
preservation is up for review at D111448). However, since loop
predication is only used in downstream pipelines, it is hard to keep BPI
from breaking for incomplete state with upstream changes in BPI.
Also, correctly preserving BPI for all loop passes is a non-trivial
undertaking (D110438 does this lossily), while the benefit of using it
in loop predication isn't clear.

In this patch, we rely on profile metadata to get almost similar benefit as
BPI, without actually using the complete heuristics provided by BPI.
This avoids the compile time explosion we tried to fix with D110438 and
also avoids fragile bugs because BPI can be lossy in loop passes
(D111448).

[anna]

minor changes in naming, clang format. NFC w.r.t. previous version.

[anna]

ping

[asbirlea]

Looks reasonable, assuming you're seeing similar results to using BPI. I'm not very familiar with BPI itself, so if others reviewers are, I'll let them chime in.

Does it make sense to make isValidProfileData and ComputeBranchProbability static utilities in BranchProbability? (nit: consistent case-ing on the names)

[ebrevnov]

llvm/lib/Transforms/Scalar/LoopPredication.cpp

971

This looks like a natural generalization of "bool Instruction::extractProfMetadata(uint64_t &TrueVal, uint64_t &FalseVal) const". Why don't we extend existing implementation of "extractProfMetadata" to handle more than 2 operands and overloaded existing API with "bool Instruction::extractProfMetadata(SmallVectorImpl<uint64_t> &) const"?

PS: Looks tempting to place total weight as a first element in the resulting vector. WDYT?

977

Looks like we already check that there is a "valid" profile at line 968. Why do it again?

[anna]

llvm/lib/Transforms/Scalar/LoopPredication.cpp

971

The SmallVector won't buy us what we need. After retrieving those individual profile data, we still need to calculate the Numerator by checking if (Term->getSuccessor(i) == ExitBlock), i.e. what we do at line 983.

Also, btw, we have Instruction::extractProfTotalWeight, which computes what I have in Denominator, but we anyway need to compute the Numerator so not sure if it's worth it.

977

Because the first check is for Latch terminators, while this one is for all branches exiting the loop. See usage of ComputeBranchProbability below.

[anna]

In D111668#3071306, @asbirlea wrote:

Looks reasonable, assuming you're seeing similar results to using BPI. I'm not very familiar with BPI itself, so if others reviewers are, I'll let them chime in.

Does it make sense to make isValidProfileData and ComputeBranchProbability static utilities in BranchProbability? (nit: consistent case-ing on the names)

I think that's a good idea. Will do as a follow-up after landing patch.

Will fix the naming.

[anna]

addressed review comment about naming. Updated a comment in the testcase to parse better.

[ebrevnov]

llvm/lib/Transforms/Scalar/LoopPredication.cpp

971

This looks like a natural generalization of "bool Instruction::extractProfMetadata(uint64_t &TrueVal, uint64_t &FalseVal) const". Why don't we extend existing implementation of "extractProfMetadata" to handle more than 2 operands and overloaded existing API with "bool Instruction::extractProfMetadata(SmallVectorImpl<uint64_t> &) const"?

PS: Looks tempting to place total weight as a first element in the resulting vector. WDYT?

971

The SmallVector won't buy us what we need. After retrieving those individual profile data, we still need to calculate the Numerator by checking if (Term->getSuccessor(i) == ExitBlock), i.e. what we do at line 983.

Sure, we will still need to iterate other the vector to extract required weight. But all the code responsible for parsing "branch_weight" metadata will move to a single place. If representation of "branch_weight" metadata is changed one day there will be a single place to adjust.

Also, btw, we have Instruction::extractProfTotalWeight, which computes what I have in Denominator, but we anyway need to compute the Numerator so not sure if it's worth it.

I don't see reasons to call Instruction::extractProfTotalWeight in this case either. My proposal was to include total weight as the first element of the vector populated by "bool Instruction::extractProfMetadata(SmallVectorImpl<uint64_t> &) const". The motivation is simple, the weight itself is mostly useless because it is relative to total weight. Thus we always need total weight as well.

977

Makes sense. Overlooked second call.

llvm/lib/Transforms/Scalar/LoopPredication.cpp

struct LoopICmp {
   }
 };
 
 class LoopPredication {
   AliasAnalysis *AA;
   DominatorTree *DT;
   ScalarEvolution *SE;
   LoopInfo *LI;
-  BranchProbabilityInfo *BPI;
   MemorySSAUpdater *MSSAU;
 
   Loop *L;
   const DataLayout *DL;
   BasicBlock *Preheader;
   LoopICmp LatchCheck;
 
   bool isSupportedStep(const SCEV* Step);
Optional<Value *> widenICmpRangeCheckDecrementingLoop(LoopICmp LatchCheck,
                                                         Instruction *Guard);
   unsigned collectChecks(SmallVectorImpl<Value *> &Checks, Value *Condition,
                          SCEVExpander &Expander, Instruction *Guard);
   bool widenGuardConditions(IntrinsicInst *II, SCEVExpander &Expander);
   bool widenWidenableBranchGuardConditions(BranchInst *Guard, SCEVExpander &Expander);
   // If the loop always exits through another block in the loop, we should not
   // predicate based on the latch check. For example, the latch check can be a
   // very coarse grained check and there can be more fine grained exit checks
-  // within the loop. We identify such unprofitable loops through BPI.
+  // within the loop.
   bool isLoopProfitableToPredicate();
 
   bool predicateLoopExits(Loop *L, SCEVExpander &Rewriter);
 
 public:
   LoopPredication(AliasAnalysis *AA, DominatorTree *DT, ScalarEvolution *SE,
-                  LoopInfo *LI, BranchProbabilityInfo *BPI,
-                  MemorySSAUpdater *MSSAU)
-      : AA(AA), DT(DT), SE(SE), LI(LI), BPI(BPI), MSSAU(MSSAU){};
+                  LoopInfo *LI, MemorySSAUpdater *MSSAU)
+      : AA(AA), DT(DT), SE(SE), LI(LI), MSSAU(MSSAU){};
   bool runOnLoop(Loop *L);
 };
 
 class LoopPredicationLegacyPass : public LoopPass {
 public:
   static char ID;
   LoopPredicationLegacyPass() : LoopPass(ID) {
     initializeLoopPredicationLegacyPassPass(*PassRegistry::getPassRegistry());
if (skipLoop(L))
       return false;
     auto *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
     auto *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     auto *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     auto *MSSAWP = getAnalysisIfAvailable<MemorySSAWrapperPass>();
     std::unique_ptr<MemorySSAUpdater> MSSAU;
     if (MSSAWP)
       MSSAU = std::make_unique<MemorySSAUpdater>(&MSSAWP->getMSSA());
-    BranchProbabilityInfo &BPI =
-        getAnalysis<BranchProbabilityInfoWrapperPass>().getBPI();
     auto *AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
-    LoopPredication LP(AA, DT, SE, LI, &BPI, MSSAU ? MSSAU.get() : nullptr);
+    LoopPredication LP(AA, DT, SE, LI, MSSAU ? MSSAU.get() : nullptr);
     return LP.runOnLoop(L);
   }
 };
 
 char LoopPredicationLegacyPass::ID = 0;
 } // end namespace
 
 INITIALIZE_PASS_BEGIN(LoopPredicationLegacyPass, "loop-predication",
                       "Loop predication", false, false)
 INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopPass)
 INITIALIZE_PASS_END(LoopPredicationLegacyPass, "loop-predication",
                     "Loop predication", false, false)
 
 Pass *llvm::createLoopPredicationPass() {
   return new LoopPredicationLegacyPass();
 }
 
 PreservedAnalyses LoopPredicationPass::run(Loop &L, LoopAnalysisManager &AM,
                                            LoopStandardAnalysisResults &AR,
                                            LPMUpdater &U) {
   std::unique_ptr<MemorySSAUpdater> MSSAU;
   if (AR.MSSA)
     MSSAU = std::make_unique<MemorySSAUpdater>(AR.MSSA);
-  LoopPredication LP(&AR.AA, &AR.DT, &AR.SE, &AR.LI, AR.BPI,
+  LoopPredication LP(&AR.AA, &AR.DT, &AR.SE, &AR.LI,
                      MSSAU ? MSSAU.get() : nullptr);
   if (!LP.runOnLoop(&L))
     return PreservedAnalyses::all();
 
   auto PA = getLoopPassPreservedAnalyses();
   if (AR.MSSA)
     PA.preserve<MemorySSAAnalysis>();
   return PA;
if (IsUnsupportedPredicate(Step, Result->Pred)) {
     return None;
   }
 
   return Result;
 }
 
 
 bool LoopPredication::isLoopProfitableToPredicate() {
-  if (SkipProfitabilityChecks || !BPI)
+  if (SkipProfitabilityChecks)
     return true;
 
   SmallVector<std::pair<BasicBlock *, BasicBlock *>, 8> ExitEdges;
   L->getExitEdges(ExitEdges);
   // If there is only one exiting edge in the loop, it is always profitable to
   // predicate the loop.
   if (ExitEdges.size() == 1)
     return true;
 
   // Calculate the exiting probabilities of all exiting edges from the loop,
   // starting with the LatchExitProbability.
   // Heuristic for profitability: If any of the exiting blocks' probability of
   // exiting the loop is larger than exiting through the latch block, it's not
   // profitable to predicate the loop.
   auto *LatchBlock = L->getLoopLatch();
   assert(LatchBlock && "Should have a single latch at this point!");
   auto *LatchTerm = LatchBlock->getTerminator();
   assert(LatchTerm->getNumSuccessors() == 2 &&
          "expected to be an exiting block with 2 succs!");
   unsigned LatchBrExitIdx =
       LatchTerm->getSuccessor(0) == L->getHeader() ? 1 : 0;
+  // We compute branch probabilities without BPI. We do not rely on BPI since
+  // Loop predication is usually run in an LPM and BPI is only preserved
+  // lossily within loop pass managers, while BPI has an inherent notion of
+  // being complete for an entire function.
+
+  // If the latch exits into a deoptimize or an unreachable block, do not
+  // predicate on that latch check.
+  auto *LatchExitBlock = LatchTerm->getSuccessor(LatchBrExitIdx);
+  if (isa<UnreachableInst>(LatchTerm) ||
+      LatchExitBlock->getTerminatingDeoptimizeCall())
+    return false; 
+
+  auto isValidProfileData = [](MDNode *ProfileData, const Instruction *Term) {
+    if (!ProfileData || !ProfileData->getOperand(0))
+      return false;
+    if (MDString *MDS = dyn_cast<MDString>(ProfileData->getOperand(0)))
+      if (!MDS->getString().equals("branch_weights"))
+        return false;
+    if (ProfileData->getNumOperands() != 1 + Term->getNumSuccessors())
+      return false;
+    return true;
+  };
+  MDNode *LatchProfileData = LatchTerm->getMetadata(LLVMContext::MD_prof);
+  // Latch terminator has no valid profile data, so nothing to check profitability on.
+  if (!isValidProfileData(LatchProfileData, LatchTerm))
+    return true;
+
+  auto ComputeBranchProbability =
+      [&](const BasicBlock *ExitingBlock,
+          const BasicBlock *ExitBlock) -> BranchProbability {
+    auto *Term = ExitingBlock->getTerminator();
+    MDNode *ProfileData = Term->getMetadata(LLVMContext::MD_prof);
+    unsigned succ_num =
+        std::distance(succ_begin(ExitingBlock), succ_end(ExitingBlock));
+    if (isValidProfileData(ProfileData, Term)) {
+      uint64_t numerator = 0, denom = 0, prof_val = 0;
+      for (unsigned i = 0; i < succ_num; i++) {
+        ConstantInt *CI =
+            mdconst::extract<ConstantInt>(ProfileData->getOperand(i + 1));
+        prof_val = CI->getValue().getZExtValue();
+        if (Term->getSuccessor(i) == ExitBlock)
+          numerator += prof_val;
+        denom += prof_val;
+      }
+      return BranchProbability::getBranchProbability(numerator, denom);
+    } else {
+      assert(LatchBlock != ExitingBlock &&
+             "Latch term should always have profile data!");
+      // No profile data, so we choose the weight as 1/num_of_succ(Src)
+      return BranchProbability::getBranchProbability(1, succ_num);
+    }
+  };
+
   BranchProbability LatchExitProbability =
-      BPI->getEdgeProbability(LatchBlock, LatchBrExitIdx);
+      ComputeBranchProbability(LatchBlock, LatchExitBlock);
 
   // Protect against degenerate inputs provided by the user. Providing a value
   // less than one, can invert the definition of profitable loop predication.
   float ScaleFactor = LatchExitProbabilityScale;
   if (ScaleFactor < 1) {
     LLVM_DEBUG(
         dbgs()
         << "Ignored user setting for loop-predication-latch-probability-scale: "
         << LatchExitProbabilityScale << "\n");
     LLVM_DEBUG(dbgs() << "The value is set to 1.0\n");
     ScaleFactor = 1.0;
   }
-  const auto LatchProbabilityThreshold =
-      LatchExitProbability * ScaleFactor;
+  const auto LatchProbabilityThreshold = LatchExitProbability * ScaleFactor;
 
   for (const auto &ExitEdge : ExitEdges) {
     BranchProbability ExitingBlockProbability =
-        BPI->getEdgeProbability(ExitEdge.first, ExitEdge.second);
+        ComputeBranchProbability(ExitEdge.first, ExitEdge.second);
     // Some exiting edge has higher probability than the latch exiting edge.
     // No longer profitable to predicate.
     if (ExitingBlockProbability > LatchProbabilityThreshold)
       return false;
   }
-  // Using BPI, we have concluded that the most probable way to exit from the
+
+  // We have concluded that the most probable way to exit from the
   // loop is through the latch (or there's no profile information and all
   // exits are equally likely).
   return true;
 }
 
 /// If we can (cheaply) find a widenable branch which controls entry into the
 /// loop, return it.
 static BranchInst *FindWidenableTerminatorAboveLoop(Loop *L, LoopInfo &LI) {

llvm/test/Transforms/LoopPredication/profitability.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -S -loop-predication-skip-profitability-checks=false -passes='require<scalar-evolution>,require<branch-prob>,loop-mssa(loop-predication)' -verify-memoryssa < %s 2>&1 | FileCheck %s
+; RUN: opt -S -loop-predication-skip-profitability-checks=false -passes='require<scalar-evolution>,loop-mssa(loop-predication)' -verify-memoryssa < %s 2>&1 | FileCheck %s
 
-; latch block exits to a speculation block. BPI already knows (without prof
+; latch block exits to a speculation block. We account for this (without prof
 ; data) that deopt is very rarely
 ; taken. So we do not predicate this loop using that coarse latch check.
 ; LatchExitProbability: 0x04000000 / 0x80000000 = 3.12%
 ; ExitingBlockProbability: 0x7ffa572a / 0x80000000 = 99.98%
 define i64 @donot_predicate(i64* nocapture readonly %arg, i32 %length, i64* nocapture readonly %arg2, i64* nocapture readonly %n_addr, i64 %i) !prof !21 {
 ; CHECK-LABEL: @donot_predicate(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[LENGTH_EXT:%.*]] = zext i32 [[LENGTH:%.*]] to i64