[BFI][CGP] Add limited support for detecting missed BFI updates and fix one in CodeGenPrepare.

Summary:
This helps detect some missed BFI updates during CodeGenPrepare.

This is debug build only and disabled behind a flag.

Fix a missed update in CodeGenPrepare::dupRetToEnableTailCallOpts().

Reviewers: davidxl

Subscribers: hiraditya, llvm-commits

Tags: #llvm

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

include/llvm/Analysis/BlockFrequencyInfo.h

@@ -103,6 +103,9 @@ public:
   uint64_t getEntryFreq() const;
   void releaseMemory();
   void print(raw_ostream &OS) const;
+
+  // Compare to the other BFI and verify they match.
+  void verifyMatch(BlockFrequencyInfo &Other) const;
 };
 
 /// Analysis pass which computes \c BlockFrequencyInfo.

include/llvm/Analysis/BlockFrequencyInfoImpl.h

@@ -1034,6 +1034,8 @@ public:
   raw_ostream &printBlockFreq(raw_ostream &OS, const BlockT *BB) const {
     return BlockFrequencyInfoImplBase::printBlockFreq(OS, getNode(BB));
   }
+
+  void verifyMatch(BlockFrequencyInfoImpl<BT> &Other) const;
 };
 
 namespace bfi_detail {
@@ -1417,6 +1419,61 @@ raw_ostream &BlockFrequencyInfoImpl<BT>::print(raw_ostream &OS) const {
   return OS;
 }
 
+template <class BT>
+void BlockFrequencyInfoImpl<BT>::verifyMatch(
+    BlockFrequencyInfoImpl<BT> &Other) const {
+  bool Match = true;
+  DenseMap<const BlockT *, BlockNode> ValidNodes;
+  DenseMap<const BlockT *, BlockNode> OtherValidNodes;
+  for (auto &Entry : Nodes) {
+    const BlockT *BB = Entry.first;
+    if (BB) {
+      ValidNodes[BB] = Entry.second.first;
+    }
+  }
+  for (auto &Entry : Other.Nodes) {
+    const BlockT *BB = Entry.first;
+    if (BB) {
+      OtherValidNodes[BB] = Entry.second.first;
+    }
+  }
+  unsigned NumValidNodes = ValidNodes.size();
+  unsigned NumOtherValidNodes = OtherValidNodes.size();
+  if (NumValidNodes != NumOtherValidNodes) {
+    Match = false;
+    dbgs() << "Number of blocks mismatch: " << NumValidNodes << " vs "
+           << NumOtherValidNodes << "\n";
+  } else {
+    for (auto &Entry : ValidNodes) {
+      const BlockT *BB = Entry.first;
+      BlockNode Node = Entry.second;
+      if (OtherValidNodes.count(BB)) {
+        BlockNode OtherNode = OtherValidNodes[BB];
+        auto Freq = Freqs[Node.Index];
+        auto OtherFreq = Other.Freqs[OtherNode.Index];
+        if (Freq.Integer != OtherFreq.Integer) {
+          Match = false;
+          dbgs() << "Freq mismatch: " << bfi_detail::getBlockName(BB) << " "
+                 << Freq.Integer << " vs " << OtherFreq.Integer << "\n";
+        }
+      } else {
+        Match = false;
+        dbgs() << "Block " << bfi_detail::getBlockName(BB) << " index "
+               << Node.Index << " does not exist in Other.\n";
+      }
+    }
+    // If there's a valid node in OtherValidNodes that's not in ValidNodes,
+    // either the above num check or the check on OtherValidNodes will fail.
+  }
+  if (!Match) {
+    dbgs() << "This\n";
+    print(dbgs());
+    dbgs() << "Other\n";
+    Other.print(dbgs());
+  }
+  assert(Match && "BFI mismatch");
+}
+
 // Graph trait base class for block frequency information graph
 // viewer.
 

lib/Analysis/BlockFrequencyInfo.cpp

@@ -287,6 +287,11 @@ void BlockFrequencyInfo::print(raw_ostream &OS) const {
     BFI->print(OS);
 }
 
+void BlockFrequencyInfo::verifyMatch(BlockFrequencyInfo &Other) const {
+  if (BFI)
+    BFI->verifyMatch(*Other.BFI);
+}
+
 INITIALIZE_PASS_BEGIN(BlockFrequencyInfoWrapperPass, "block-freq",
                       "Block Frequency Analysis", true, true)
 INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)

lib/CodeGen/CodeGenPrepare.cpp

@@ -229,6 +229,11 @@ static cl::opt<bool> EnableICMP_EQToICMP_ST(
     "cgp-icmp-eq2icmp-st", cl::Hidden, cl::init(false),
     cl::desc("Enable ICMP_EQ to ICMP_S(L|G)T conversion."));
 
+static cl::opt<bool>
+    VerifyBFIUpdates("cgp-verify-bfi-updates", cl::Hidden, cl::init(false),
+                     cl::desc("Enable BFI update verification for "
+                              "CodeGenPrepare."));
+
 namespace {
 
 enum ExtType {
@@ -405,6 +410,7 @@ class TypePromotionTransaction;
     bool optimizeCmp(CmpInst *Cmp, bool &ModifiedDT);
     bool combineToUSubWithOverflow(CmpInst *Cmp, bool &ModifiedDT);
     bool combineToUAddWithOverflow(CmpInst *Cmp, bool &ModifiedDT);
+    void verifyBFIUpdates(Function &F);
   };
 
 } // end anonymous namespace
@@ -565,9 +571,23 @@ bool CodeGenPrepare::runOnFunction(Function &F) {
   // preparatory transforms.
   EverMadeChange |= placeDbgValues(F);
 
+#ifndef NDEBUG
+  if (VerifyBFIUpdates)
+    verifyBFIUpdates(F);
+#endif
+
   return EverMadeChange;
 }
 
+// Verify BFI has been updated correctly by recomputing BFI and comparing them.
+void CodeGenPrepare::verifyBFIUpdates(Function &F) {
+  DominatorTree NewDT(F);
+  LoopInfo NewLI(NewDT);
+  BranchProbabilityInfo NewBPI(F, NewLI, TLInfo);
+  BlockFrequencyInfo NewBFI(F, NewBPI, NewLI);
+  NewBFI.verifyMatch(*BFI);
+}
+
 /// Merge basic blocks which are connected by a single edge, where one of the
 /// basic blocks has a single successor pointing to the other basic block,
 /// which has a single predecessor.
@@ -2204,6 +2224,11 @@ bool CodeGenPrepare::dupRetToEnableTailCallOpts(BasicBlock *BB, bool &ModifiedDT
 
     // Duplicate the return into TailCallBB.
     (void)FoldReturnIntoUncondBranch(RetI, BB, TailCallBB);
+    assert(!VerifyBFIUpdates ||
+           BFI->getBlockFreq(BB) >= BFI->getBlockFreq(TailCallBB));
+    BFI->setBlockFreq(
+        BB,
+        (BFI->getBlockFreq(BB) - BFI->getBlockFreq(TailCallBB)).getFrequency());
     ModifiedDT = Changed = true;
     ++NumRetsDup;
   }

test/CodeGen/AArch64/aarch64-tbz.ll

@@ -1,4 +1,5 @@
 ; RUN: llc -verify-machineinstrs -mtriple=aarch64-linux-gnueabi < %s | FileCheck %s
+; RUN: llc -verify-machineinstrs -mtriple=aarch64-linux-gnueabi -cgp-verify-bfi-updates=true < %s | FileCheck %s
 
 ; CHECK-LABEL: test1
 ; CHECK: tbz {{w[0-9]}}, #3, {{.LBB0_3}}