[BranchProbabilityInfo] Use SmallVector (NFC)

[kazu]

Summary

This patch simplifies BranchProbabilityInfo by changing the type of
Probs.

Without this patch:

DenseMap<Edge, BranchProbability> Probs

maps an ordered pair of a BasicBlock* and a successor index to an edge
probability.

With this patch:

DenseMap<const BasicBlock *, SmallVector<BranchProbability, 2>> Probs

maps a BasicBlock* to a vector of edge probabilities.

BranchProbabilityInfo has a property that for a given basic block, we
either have edge probabilities for all successors or do not have any
edge probability at all. This property combined with the current map
type leads to a somewhat complicated algorithm in eraseBlock to erase
map entries one by one while increasing the successor index.

The new map type allows us to remove the all edge probabilities for a
given basic block in a more intuitive manner, namely:

Probs.erase(BB);

[MaskRay]

It is true that SmallVector is more convenient to use but this change also has potential to increase memory usage. Have you checked whether this can increase memory usage?

[kazu]

In D91017#2381145, @MaskRay wrote:

It is true that SmallVector is more convenient to use but this change also has potential to increase memory usage. Have you checked whether this can increase memory usage?

The differences in memory consumption seem to be really small.

I have invoked clang for three input files, three times each, and looked at Maximum resident set size (kbytes) with /usr/bin/time -v.

JumpThreading.ii
Without this patch: 280792 281016 281392
With this patch:    280736 281484 281568

SROA.ii
Without this patch: 311820 312412 312488
With this patch:    311224 311356 311468

LICM.ii
Without this patch: 346300 346380 348356
With this patch:    345528 348084 348564

Under the hood, the difference should be a wash because:

DenseMap<std::pair<const BasicBlock *, unsigned>, BranchProbability>::value_type

is 24 bytes in size, padded to 32 bytes, and

DenseMap<BasicBlock *, SmallVector<BranchProbability, 2>>::value_type

is 32 bytes in size.

If we assume that we have a function with 100 "if" statements with no switch statements at 50% occupancy of the hash table, we have:

Without this patch: 32 bytes/entry * 200 edges        / 50% occupancy = 12,800 bytes
With this patch:    32 bytes/entry * 100 basic blocks / 50% occupancy =  6,400 bytes

[MaskRay]

In D91017#2381226, @kazu wrote:

In D91017#2381145, @MaskRay wrote:

It is true that SmallVector is more convenient to use but this change also has potential to increase memory usage. Have you checked whether this can increase memory usage?

The differences in memory consumption seem to be really small.

I have invoked clang for three input files, three times each, and looked at Maximum resident set size (kbytes) with /usr/bin/time -v.

JumpThreading.ii
Without this patch: 280792 281016 281392
With this patch:    280736 281484 281568

SROA.ii
Without this patch: 311820 312412 312488
With this patch:    311224 311356 311468

LICM.ii
Without this patch: 346300 346380 348356
With this patch:    345528 348084 348564

Under the hood, the difference should be a wash because:

DenseMap<std::pair<const BasicBlock *, unsigned>, BranchProbability>::value_type

is 24 bytes in size, padded to 32 bytes, and

DenseMap<BasicBlock *, SmallVector<BranchProbability, 2>>::value_type

is 32 bytes in size.

If we assume that we have a function with 100 "if" statements with no switch statements at 50% occupancy of the hash table, we have:

Without this patch: 32 bytes/entry * 200 edges        / 50% occupancy = 12,800 bytes
With this patch:    32 bytes/entry * 100 basic blocks / 50% occupancy =  6,400 bytes

Thanks for the explanation!

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1098

const std::vector<...> & (specify the concrete type if unclear)

1120

const std::vector<...> & (specify the concrete type if unclear)

[kazu]

I've specified for SrcProbs.

[yrouban]

This looks to be a good change towards simplification. I think originally a map over edges was introduced by @atrick back in 2011 in https://github.com/llvm/llvm-project/commit/49371f3f33788 as it could be sparse. But with the commit https://gitlab.azulsystems.com/dev/orca/commit/8138487468e22 we imposed an important restriction on setting outgoing edge probabilities at once for one block. This allowed correctness checks and further improvements like this.

As of performance impact of this patch I would suggest that you check llvm-compile-time-tracker.com out or contact @nikic directly.

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1172–1175

remove braces

[kazu]

Removed unnecessary curly braces.

[kazu]

@nickc kindly measured the compilation time for me:

https://llvm-compile-time-tracker.com/compare.php?from=dd5b51f4fabbb44c2f6629f0dafefa7b7f3f0b23&to=2d731aa30cd9f27daedae206d4cb0de8a4fe3ad8&stat=instructions

shows a very slight improvement.

Qualitatively speaking, we are making the data structure more contiguous while maintaining the initial size of Probs. Once we start inserting edge probabilities, then the SmallVector approach should come out positive (that is, less memory) because a basic block ending with an "if" instruction takes up a single hash entry as opposed to two.

[MaskRay]

Thanks!

[atrick]

I won't weigh in decisively, but it does look like a positive change to me.

When it comes to designing for compile time and memory usage, it's more important to protect against the worst-case (pathological input) than streamlining the average case--there's typically one straggler function that takes excessive time while the rest can be easily parallelized. That was the motivation behind avoiding allocations that may scale with the function size.

On the other hand, the worst-case with this change--a very large number of switches in the same function---seems truly pathological, and the consequence is not so bad--it's just a linear increase in allocations and memory.

When in doubt, the simpler implementation should be preferred.

[dmajor]

Our bots are failing on the new assertion "The number of edge probabilities must match the number of successors." during PGO builds of clang.

The backtrace and reproducers are available from https://treeherder.mozilla.org/logviewer?job_id=321399701&repo=try&lineNumber=50587 (you may need to click "Show Job Info" in the top bar to get the files).

[yrouban]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1121

could you please extract this assertion and the one at line 1099 to a separate patch that the rest NFC should depend on? This new extracted patch would be a functional change that impose a stronger restriction on BPI correctness.

[kazu]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1121

I'll do that. Again, let me do additional testing on my side and fix whatever issues come up before I'll post the assertion-only patch.

llvm/include/llvm/Analysis/BranchProbabilityInfo.h

private:
   // Default weight value. Used when we don't have information about the edge.
   // TODO: DEFAULT_WEIGHT makes sense during static predication, when none of
   // the successors have a weight yet. But it doesn't make sense when providing
   // weight to an edge that may have siblings with non-zero weights. This can
   // be handled various ways, but it's probably fine for an edge with unknown
   // weight to just "inherit" the non-zero weight of an adjacent successor.
   static const uint32_t DEFAULT_WEIGHT = 16;
 
-  DenseMap<Edge, BranchProbability> Probs;
+  DenseMap<const BasicBlock *, SmallVector<BranchProbability, 2>> Probs;
 
   /// Track the last function we run over for printing.
   const Function *LastF = nullptr;
 
   /// Keeps information about all SCCs in a function.
   std::unique_ptr<const SccInfo> SccI;
 
   /// Track the set of blocks directly succeeded by a returning block.

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");
-
-  if (I != Probs.end())
-    return I->second;
-
+  auto I = Probs.find(Src);
+  if (I == Probs.end())
     return {1, static_cast<uint32_t>(succ_size(Src))};
+
+  const SmallVector<BranchProbability, 2> &SrcProbs = I->second;
+  assert(SrcProbs.size() == Src->getTerminator()->getNumSuccessors() &&
+         "The number of edge probabilities must match the number of "
+         "successors.");
+  return SrcProbs[IndexInSuccessors];
 }
 
 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)))
+  auto SrcProbsIt = Probs.find(Src);
+  if (SrcProbsIt == Probs.end())
     return BranchProbability(llvm::count(successors(Src), Dst), succ_size(Src));
 
+  const SmallVector<BranchProbability, 2> &SrcProbs = SrcProbsIt->second;
+  assert(SrcProbs.size() == Src->getTerminator()->getNumSuccessors() &&
+         "The number of edge probabilities must match the number of "
+         "successors.");
   auto Prob = BranchProbability::getZero();
   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 += SrcProbs[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());
   eraseBlock(Src); // Erase stale data if any.
   if (Probs.size() == 0)
     return; // Nothing to set.
 
   Handles.insert(BasicBlockCallbackVH(Src, this));
   uint64_t TotalNumerator = 0;
+  auto &SrcProbs = this->Probs[Src];
+  SrcProbs.reserve(Probs.size());
   for (unsigned SuccIdx = 0; SuccIdx < Probs.size(); ++SuccIdx) {
-    this->Probs[std::make_pair(Src, SuccIdx)] = Probs[SuccIdx];
+    SrcProbs.push_back(Probs[SuccIdx]);
     LLVM_DEBUG(dbgs() << "set edge " << Src->getName() << " -> " << SuccIdx
                       << " successor probability to " << Probs[SuccIdx]
                       << "\n");
     TotalNumerator += Probs[SuccIdx].getNumerator();
   }
 
   // Because of rounding errors the total probability cannot be checked to be
   // 1.0 exactly. That is TotalNumerator == BranchProbability::getDenominator.
   // Instead, every single probability in Probs must be as accurate as possible.
   // This results in error 1/denominator at most, thus the total absolute error
   // should be within Probs.size / BranchProbability::getDenominator.
   assert(TotalNumerator <= BranchProbability::getDenominator() + Probs.size());
   assert(TotalNumerator >= BranchProbability::getDenominator() - Probs.size());
 }
 
 void BranchProbabilityInfo::copyEdgeProbabilities(BasicBlock *Src,
                                                   BasicBlock *Dst) {
   eraseBlock(Dst); // Erase stale data if any.
   unsigned NumSuccessors = Src->getTerminator()->getNumSuccessors();
   assert(NumSuccessors == Dst->getTerminator()->getNumSuccessors());
   if (NumSuccessors == 0)
     return; // Nothing to set.
-  if (this->Probs.find(std::make_pair(Src, 0)) == this->Probs.end())
+  auto SrcProbsIt = this->Probs.find(Src);
+  if (SrcProbsIt == this->Probs.end())
     return; // No probability is set for edges from Src. Keep the same for Dst.
 
   Handles.insert(BasicBlockCallbackVH(Dst, this));
-  for (unsigned SuccIdx = 0; SuccIdx < NumSuccessors; ++SuccIdx) {
-    auto Prob = this->Probs[std::make_pair(Src, SuccIdx)];
-    this->Probs[std::make_pair(Dst, SuccIdx)] = Prob;
+  auto &DstProbs = this->Probs[Dst];
+  DstProbs = SrcProbsIt->second;
+  for (unsigned SuccIdx = 0; SuccIdx < NumSuccessors; ++SuccIdx)
     LLVM_DEBUG(dbgs() << "set edge " << Dst->getName() << " -> " << SuccIdx
-                      << " successor probability to " << Prob << "\n");
-  }
+                      << " successor probability to " << DstProbs[SuccIdx]
+                      << "\n");
 }
 
 raw_ostream &
 BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
                                             const BasicBlock *Src,
                                             const BasicBlock *Dst) const {
   const BranchProbability Prob = getEdgeProbability(Src, Dst);
   OS << "edge " << Src->getName() << " -> " << Dst->getName()
      << " probability is " << Prob
      << (isEdgeHot(Src, Dst) ? " [HOT edge]\n" : "\n");
 
   return OS;
 }
 
 void BranchProbabilityInfo::eraseBlock(const BasicBlock *BB) {
-  // Note that we cannot use successors of BB because the terminator of BB may
-  // have changed when eraseBlock is called as a BasicBlockCallbackVH callback.
-  // Instead we remove prob data for the block by iterating successors by their
-  // indices from 0 till the last which exists. There could not be prob data for
-  // a pair (BB, N) if there is no data for (BB, N-1) because the data is always
-  // set for all successors from 0 to M at once by the method
-  // setEdgeProbability().
   Handles.erase(BasicBlockCallbackVH(BB, this));
-  for (unsigned I = 0;; ++I) {
-    auto MapI = Probs.find(std::make_pair(BB, I));
-    if (MapI == Probs.end()) {
-      assert(Probs.count(std::make_pair(BB, I + 1)) == 0 &&
-             "Must be no more successors");
-      return;
-    }
-    Probs.erase(MapI);
-  }
+  Probs.erase(BB);
 }
 
 void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LI,
                                       const TargetLibraryInfo *TLI,
                                       PostDominatorTree *PDT) {
   LLVM_DEBUG(dbgs() << "---- Branch Probability Info : " << F.getName()
                     << " ----\n\n");
   LastF = &F; // Store the last function we ran on for printing.