[BPI] Handle incomplete state in BPI

[anna]

Summary

BPI has this inherent property that once it is computed for a block's
successor, Probs data exists for all successors of the block.

However, this is not true during lossy BPI preservation in loop passes,
where BPI is updated when a block is removed, but not when a block is
added (D110438).

This patch makes BPI more relaxed in the face of incomplete information:
removing an assertion that is no longer true and not relying on Probs
data to unconditionally have the information for a successor.

The unfortunate bit is that I cannot come up with any unit tests that
trigger the failure when we have lossy BPI (even downstream, the
failures are intemittent within the loop predication pass which uses
BPI). Upstream pipeline does not have a loop pass manager invocation
that uses lossy BPI preservation, but when it does, we will definitely
trip on these issues.

[asbirlea]

I'm confused how the intermittent assert occurs. Is the compilation non-deterministic?

[anna]

In D111448#3052089, @asbirlea wrote:

I'm confused how the intermittent assert occurs. Is the compilation non-deterministic?

Compilation is deterministic.

There can be couple of reasons (I couldn't find the reason for the intermittent failure) :

1. Underlying data structure for Probs is DenseMap<Edge, BranchProbability>. Iteration through Probs is non-determistic, but I don't see us relying on that in BPI.
2. It could also be that some pass within our LPM invocation has non-determinism in iteration through use-list or the way the successors are updated (this LPM is made up of all upstream passes: LoopPredication, LICM, LoopSimplifyCFG, LoopUnswitch<non-trivial>, GuardWidening).

[asbirlea]

The question I was trying to get answer was: can there be the case where the probability is queried such that you can get different values based on when (what order) the blocks are removed or added? I think the answer is no, and I *think* I know why you're sometimes *not* triggering the assert: it's when IndexInSuccessors == 0. I suspect the successor order at the callsite in LoopPredication.cpp:947 is non-deterministic, so sometimes it's comparing against index 1, which does not have prob, and sometimes against itself, for which it does.

llvm/lib/Analysis/BranchProbabilityInfo.cpp

 /// Get the raw edge probability for the edge. If can't find it, return a
 /// default probability 1/N where N is the number of successors. Here an edge is
 /// specified using PredBlock and an
 /// index to the successors.
 BranchProbability
 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
                                           unsigned IndexInSuccessors) const {
   auto I = Probs.find(std::make_pair(Src, IndexInSuccessors));
-  assert((Probs.end() == Probs.find(std::make_pair(Src, 0))) ==
-             (Probs.end() == I) &&
-         "Probability for I-th successor must always be defined along with the "
-         "probability for the first successor");
 
+  // We can be having cases where we have probability defined for first
+  // successor but not for the I-th successor. This happens when BPI is updated
+  // lossily through loop passes.
   if (I != Probs.end())
     return I->second;
 
   return {1, static_cast<uint32_t>(succ_size(Src))};
 }
 
 BranchProbability
 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
                                           const_succ_iterator Dst) const {
   return getEdgeProbability(Src, Dst.getSuccessorIndex());
 }
 
 /// Get the raw edge probability calculated for the block pair. This returns the
 /// sum of all raw edge probabilities from Src to Dst.
 BranchProbability
 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
                                           const BasicBlock *Dst) const {
   if (!Probs.count(std::make_pair(Src, 0)))
     return BranchProbability(llvm::count(successors(Src), Dst), succ_size(Src));
 
   auto Prob = BranchProbability::getZero();
+  // We cannot rely on retrieving info in `Probs` since the successor may
+  // not have been updated into Probs if BPI is not updated when new blocks are
+  // added through loop passes. 
   for (const_succ_iterator I = succ_begin(Src), E = succ_end(Src); I != E; ++I)
     if (*I == Dst)
-      Prob += Probs.find(std::make_pair(Src, I.getSuccessorIndex()))->second;
+      Prob += getEdgeProbability(Src, I.getSuccessorIndex());
 
   return Prob;
 }
 
 /// Set the edge probability for all edges at once.
 void BranchProbabilityInfo::setEdgeProbability(
     const BasicBlock *Src, const SmallVectorImpl<BranchProbability> &Probs) {
   assert(Src->getTerminator()->getNumSuccessors() == Probs.size());

llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp

 //===- BranchProbabilityInfoTest.cpp - BranchProbabilityInfo unit tests ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/DataTypes.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/BreakCriticalEdges.h"
 #include "gtest/gtest.h"
 
 namespace llvm {
 namespace {
 
 struct BranchProbabilityInfoTest : public testing::Test {
   std::unique_ptr<BranchProbabilityInfo> BPI;
   std::unique_ptr<DominatorTree> DT;
TEST_F(BranchProbabilityInfoTest, StressUnreachableHeuristic) {
 
   // FIXME: This doesn't seem optimal. Since all of the cases handled by the
   // switch have the *same* destination block ("preexit"), shouldn't it be the
   // hot one? I'd expect the results to be reversed here...
   EXPECT_FALSE(BPI.isEdgeHot(EntryBB, PreExitBB));
   EXPECT_TRUE(BPI.isEdgeHot(EntryBB, ExitBB));
 }
 
+static std::unique_ptr<Module>
+makeLLVMModuleWithCustomFunc(LLVMContext &Context, const char *ModuleStr) {
+  SMDiagnostic Err;
+  return parseAssemblyString(ModuleStr, Err, Context);
+}
+
+static BasicBlock *getBasicBlockByName(Function &F, StringRef Name) {
+  for (BasicBlock &BB : F)
+    if (BB.getName() == Name)
+      return &BB;
+  llvm_unreachable("Expected to find basic block!");
+}
+
+// We may update BPI lossily in loop passes, which means newly created blocks do
+// not have BPI.
+TEST_F(BranchProbabilityInfoTest, TestLossyBPIInLoop) {
+  const char *ModuleStr =
+      "target datalayout = \"e-m:o-i64:64-f80:128-n8:16:32:64-S128\"\n"
+      "define void @foo(i32 %n, i32 %m, i1 %inv_cond) {\n"
+      "entry:\n"
+      "  br i1 undef, label %for.end, label %for.cond\n"
+      "for.cond:\n"
+      "  %i.0 = phi i32 [ 0, %entry ], [ %inc, %latch.1 ], [ %inc, %latch.2 ]\n"
+      "  %inc = add nsw i32 %i.0, 1\n"
+      "  %cmp = icmp slt i32 %i.0, %n\n"
+      "  br i1 %inv_cond, label %deadbb, label %latch.1\n"
+      "deadbb:\n"
+      "  br label %latch.1\n"
+      "latch.1:\n"
+      "  br i1 %cmp, label %for.cond, label %latch.2\n"
+      "latch.2:\n"
+      "  %cmp2 = icmp slt i32 %i.0, %m\n"
+      "  br i1 %cmp2, label %for.cond, label %for.end\n"
+      "for.end:\n"
+      "  %non.lcssa.phi = phi i1 [ %cmp2, %latch.2 ], [ %inv_cond, %entry ]\n"
+      "  ret void\n"
+      "}\n";
+
+  // Parse the module.
+  LLVMContext Context;
+  std::unique_ptr<Module> M = makeLLVMModuleWithCustomFunc(Context, ModuleStr);
+  Function &F = *(M->getFunction("foo"));
+  auto *ExitBB = getBasicBlockByName(F, "for.end");
+  auto *Header = getBasicBlockByName(F, "for.cond");
+  auto *Latch2 = getBasicBlockByName(F, "latch.2");
+  auto *Latch1 = getBasicBlockByName(F, "latch.1");
+  auto *DeadBB = getBasicBlockByName(F, "deadbb");
+
+  BranchProbabilityInfo &BPI = buildBPI(F);
+
+  ASSERT_NE(BPI.getEdgeProbability(Latch2, ExitBB),
+            BranchProbability::getZero());
+  ASSERT_NE(BPI.getEdgeProbability(Header, Latch1),
+            BranchProbability::getZero());
+  // Maybe some loop pass decides to preserveLCSSA form and split the critical
+  // edges to the exit block.
+  auto *CritExit = SplitCriticalEdge(
+      Latch2, ExitBB,
+      CriticalEdgeSplittingOptions().unsetPreserveLoopSimplify());
+  // Remove dead block from loop.
+  auto *OldBr = Header->getTerminator();
+  IRBuilder<> Builder(OldBr);
+  auto *NewBI = Builder.CreateBr(Latch1);
+  OldBr->eraseFromParent();
+  DeadBB->eraseFromParent();
+
+  // We should not fail here when retrieving the edge probability of newly
+  // created blocks or blocks whose successors have changed.
+  EXPECT_NE(BPI.getEdgeProbability(Latch2, CritExit),
+            BranchProbability::getZero());
+  EXPECT_NE(BPI.getEdgeProbability(Header, Latch1),
+            BranchProbability::getZero());
+  EXPECT_NE(BPI.getEdgeProbability(Latch1, Latch2),
+            BranchProbability::getZero());
+  EXPECT_NE(BPI.getEdgeProbability(Latch1, Header),
+            BranchProbability::getZero());
+  EXPECT_NE(BPI.getEdgeProbability(CritExit, ExitBB),
+            BranchProbability::getZero());
+
+  BasicBlock *Latch2Pred =
+      splitBlockBefore(Latch2, Latch2->getTerminator(), nullptr, nullptr,
+                       nullptr, "latch2.pred");
+  EXPECT_NE(BPI.getEdgeProbability(Latch2Pred, Latch2),
+            BranchProbability::getZero());
+  EXPECT_NE(BPI.getEdgeProbability(Latch1, Latch2Pred),
+            BranchProbability::getZero());
+}
 } // end anonymous namespace
 } // end namespace llvm