[BFI] Use CallbackVH to notify BFI about deletion of basic blocks

With AssertingVHs instead of bare pointers in
BlockFrequencyInfoImpl::Nodes (but without CallbackVHs) ~1/36 of all
tests ran by make check fail. It means that there are users of BFI that
delete basic blocks while keeping BFI. Some of those transformations add
new basic blocks, so if a new basic block happens to be allocated at
address where an already deleted block was and we don't explicitly set
block frequency for that new block, BFI will report some non-default
frequency for the block even though frequency for the block was never
set. Inliner is an example of a transformation that adds and removes BBs
while querying and updating BFI.
With this patch, thanks to updates via CallbackVH, BFI won't keep stale
pointers in its Nodes map.

Reviewers: davidxl, yamauchi, asbirlea, fhahn, fedor.sergeev

Reviewed-By: asbirlea, davidxl
Tags: #llvm

Differential Revision: https://reviews.llvm.org/D75341

diff --git a/llvm/include/llvm/Analysis/BlockFrequencyInfoImpl.h b/llvm/include/llvm/Analysis/BlockFrequencyInfoImpl.h
index 4d80591..ae62f80 100644 (file)
--- a/llvm/include/llvm/Analysis/BlockFrequencyInfoImpl.h
+++ b/llvm/include/llvm/Analysis/BlockFrequencyInfoImpl.h
@@ -24,6 +24,7 @@
 #include "llvm/ADT/Twine.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/ValueHandle.h"
 #include "llvm/Support/BlockFrequency.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/CommandLine.h"
@@ -547,6 +548,7 @@ namespace bfi_detail {
 template <class BlockT> struct TypeMap {};
 template <> struct TypeMap<BasicBlock> {
   using BlockT = BasicBlock;
+  using BlockKeyT = AssertingVH<const BasicBlock>;
   using FunctionT = Function;
   using BranchProbabilityInfoT = BranchProbabilityInfo;
   using LoopT = Loop;
@@ -554,6 +556,7 @@ template <> struct TypeMap<BasicBlock> {
 };
 template <> struct TypeMap<MachineBasicBlock> {
   using BlockT = MachineBasicBlock;
+  using BlockKeyT = const MachineBasicBlock *;
   using FunctionT = MachineFunction;
   using BranchProbabilityInfoT = MachineBranchProbabilityInfo;
   using LoopT = MachineLoop;
@@ -845,6 +848,7 @@ template <class BT> class BlockFrequencyInfoImpl : BlockFrequencyInfoImplBase {
   friend struct bfi_detail::BlockEdgesAdder<BT>;
 
   using BlockT = typename bfi_detail::TypeMap<BT>::BlockT;
+  using BlockKeyT = typename bfi_detail::TypeMap<BT>::BlockKeyT;
   using FunctionT = typename bfi_detail::TypeMap<BT>::FunctionT;
   using BranchProbabilityInfoT =
       typename bfi_detail::TypeMap<BT>::BranchProbabilityInfoT;
@@ -857,9 +861,11 @@ template <class BT> class BlockFrequencyInfoImpl : BlockFrequencyInfoImplBase {
   const LoopInfoT *LI = nullptr;
   const FunctionT *F = nullptr;
 
+  class BFICallbackVH;
+
   // All blocks in reverse postorder.
   std::vector<const BlockT *> RPOT;
-  DenseMap<const BlockT *, BlockNode> Nodes;
+  DenseMap<BlockKeyT, std::pair<BlockNode, BFICallbackVH>> Nodes;
 
   using rpot_iterator = typename std::vector<const BlockT *>::const_iterator;
 
@@ -871,7 +877,8 @@ template <class BT> class BlockFrequencyInfoImpl : BlockFrequencyInfoImplBase {
   BlockNode getNode(const rpot_iterator &I) const {
     return BlockNode(getIndex(I));
   }
-  BlockNode getNode(const BlockT *BB) const { return Nodes.lookup(BB); }
+
+  BlockNode getNode(const BlockT *BB) const { return Nodes.lookup(BB).first; }
 
   const BlockT *getBlock(const BlockNode &Node) const {
     assert(Node.Index < RPOT.size());
@@ -992,6 +999,13 @@ public:
 
   void setBlockFreq(const BlockT *BB, uint64_t Freq);
 
+  void forgetBlock(const BlockT *BB) {
+    // We don't erase corresponding items from `Freqs`, `RPOT` and other to
+    // avoid invalidating indices. Doing so would have saved some memory, but
+    // it's not worth it.
+    Nodes.erase(BB);
+  }
+
   Scaled64 getFloatingBlockFreq(const BlockT *BB) const {
     return BlockFrequencyInfoImplBase::getFloatingBlockFreq(getNode(BB));
   }
@@ -1019,6 +1033,30 @@ public:
   }
 };
 
+template <>
+class BlockFrequencyInfoImpl<BasicBlock>::BFICallbackVH : public CallbackVH {
+  BlockFrequencyInfoImpl *BFIImpl;
+
+public:
+  BFICallbackVH() = default;
+
+  BFICallbackVH(const BasicBlock *BB, BlockFrequencyInfoImpl *BFIImpl)
+    : CallbackVH(BB), BFIImpl(BFIImpl) {}
+
+  void deleted() override {
+    BFIImpl->forgetBlock(cast<BasicBlock>(getValPtr()));
+  }
+};
+
+/// Dummy implementation since MachineBasicBlocks aren't Values, so ValueHandles
+/// don't apply to them.
+template <>
+class BlockFrequencyInfoImpl<MachineBasicBlock>::BFICallbackVH {
+public:
+  BFICallbackVH() = default;
+  BFICallbackVH(const MachineBasicBlock *, BlockFrequencyInfoImpl *) {}
+};
+
 template <class BT>
 void BlockFrequencyInfoImpl<BT>::calculate(const FunctionT &F,
                                            const BranchProbabilityInfoT &BPI,
@@ -1066,7 +1104,7 @@ void BlockFrequencyInfoImpl<BT>::setBlockFreq(const BlockT *BB, uint64_t Freq) {
     // BlockNode for it assigned with a new index. The index can be determined
     // by the size of Freqs.
     BlockNode NewNode(Freqs.size());
-    Nodes[BB] = NewNode;
+    Nodes[BB] = {NewNode, BFICallbackVH(BB, this)};
     Freqs.emplace_back();
     BlockFrequencyInfoImplBase::setBlockFreq(NewNode, Freq);
   }
@@ -1086,7 +1124,7 @@ template <class BT> void BlockFrequencyInfoImpl<BT>::initializeRPOT() {
     BlockNode Node = getNode(I);
     LLVM_DEBUG(dbgs() << " - " << getIndex(I) << ": " << getBlockName(Node)
                       << "\n");
-    Nodes[*I] = Node;
+    Nodes[*I] = {Node, BFICallbackVH(*I, this)};
   }
 
   Working.reserve(RPOT.size());

diff --git a/llvm/include/llvm/IR/ValueHandle.h b/llvm/include/llvm/IR/ValueHandle.h
index a21d017..81438c6 100644 (file)
--- a/llvm/include/llvm/IR/ValueHandle.h
+++ b/llvm/include/llvm/IR/ValueHandle.h
@@ -414,6 +414,7 @@ protected:
 public:
   CallbackVH() : ValueHandleBase(Callback) {}
   CallbackVH(Value *P) : ValueHandleBase(Callback, P) {}
+  CallbackVH(const Value *P) : CallbackVH(const_cast<Value *>(P)) {}
 
   operator Value*() const {
     return getValPtr();