[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.

This is a resubmission of 408349a with
fixed compiler warning and MSVC compilation error.

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

Reviewed-By: asbirlea, davidxl

Tags: #llvm

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

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,19 +548,24 @@ 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;
   using LoopInfoT = LoopInfo;
 };
 template <> struct TypeMap<MachineBasicBlock> {
   using BlockT = MachineBasicBlock;
+  using BlockKeyT = const MachineBasicBlock *;
   using FunctionT = MachineFunction;
   using BranchProbabilityInfoT = MachineBranchProbabilityInfo;
   using LoopT = MachineLoop;
   using LoopInfoT = MachineLoopInfo;
 };
 
+template <class BlockT, class BFIImplT>
+class BFICallbackVH;
+
 /// Get the name of a MachineBasicBlock.
 ///
 /// Get the name of a MachineBasicBlock.  It's templated so that including from
@@ -845,21 +851,24 @@ 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;
   using LoopT = typename bfi_detail::TypeMap<BT>::LoopT;
   using LoopInfoT = typename bfi_detail::TypeMap<BT>::LoopInfoT;
   using Successor = GraphTraits<const BlockT *>;
   using Predecessor = GraphTraits<Inverse<const BlockT *>>;
+  using BFICallbackVH =
+      bfi_detail::BFICallbackVH<BlockT, BlockFrequencyInfoImpl>;
 
   const BranchProbabilityInfoT *BPI = nullptr;
   const LoopInfoT *LI = nullptr;
   const FunctionT *F = nullptr;
 
   // 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 +880,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 +1002,13 @@ template <class BT> class BlockFrequencyInfoImpl : BlockFrequencyInfoImplBase {
 
   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 +1036,36 @@ template <class BT> class BlockFrequencyInfoImpl : BlockFrequencyInfoImplBase {
   }
 };
 
+namespace bfi_detail {
+
+template <class BFIImplT>
+class BFICallbackVH<BasicBlock, BFIImplT> : public CallbackVH {
+  BFIImplT *BFIImpl;
+
+public:
+  BFICallbackVH() = default;
+
+  BFICallbackVH(const BasicBlock *BB, BFIImplT *BFIImpl)
+      : CallbackVH(BB), BFIImpl(BFIImpl) {}
+
+  virtual ~BFICallbackVH() = default;
+
+  void deleted() override {
+    BFIImpl->forgetBlock(cast<BasicBlock>(getValPtr()));
+  }
+};
+
+/// Dummy implementation since MachineBasicBlocks aren't Values, so ValueHandles
+/// don't apply to them.
+template <class BFIImplT>
+class BFICallbackVH<MachineBasicBlock, BFIImplT> {
+public:
+  BFICallbackVH() = default;
+  BFICallbackVH(const MachineBasicBlock *, BFIImplT *) {}
+};
+
+} // end namespace bfi_detail
+
 template <class BT>
 void BlockFrequencyInfoImpl<BT>::calculate(const FunctionT &F,
                                            const BranchProbabilityInfoT &BPI,
@@ -1066,7 +1113,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 +1133,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());

llvm/include/llvm/IR/ValueHandle.h

@@ -414,6 +414,7 @@ class CallbackVH : public ValueHandleBase {
 public:
   CallbackVH() : ValueHandleBase(Callback) {}
   CallbackVH(Value *P) : ValueHandleBase(Callback, P) {}
+  CallbackVH(const Value *P) : CallbackVH(const_cast<Value *>(P)) {}
 
   operator Value*() const {
     return getValPtr();