[MBFIWrapper] Add a new function getBlockProfileCount

[Carrot]

Summary

MBFIWrapper keeps track of block frequencies of newly created blocks and modified blocks, modified block frequencies should also impact block profile count. This class doesn't provide interface getBlockProfileCount, users can only use the underlying MBFI to query profile count, the underlying MBFI doesn't know the modifications made in MBFIWrapper, so it either provides stale profile count for modified block or simply crashes on new blocks.

So this patch add function getBlockProfileCount to class MBFIWrapper to handle new blocks or modified blocks.

[davidxl]

Is it possible to add a test case?

[Carrot]

It can be triggered by a test case in my next basic block duplication patch, because newly created basic block can cause MBFI.getBlockProfileCount crash.

It is very difficult to write a test case in current code base.

[davidxl]

llvm/lib/CodeGen/MBFIWrapper.cpp

37

Can you add some comments in the file here ?

[davidxl]

llvm/lib/CodeGen/MBFIWrapper.cpp

42

If the code reaches here, will the call simply return 0 (according to its implementation)?

[Carrot]

llvm/lib/CodeGen/MBFIWrapper.cpp

42

See the function getBlockFreq and setBlockFreq, when code reaches here, it means MBB frequency is not changed, so its profile count is not impacted, so we can simply query the underlying MBFI for profile count.

[davidxl]

lgtm

llvm/include/llvm/CodeGen/MBFIWrapper.h

 class MachineBlockFrequencyInfo;
 
 class MBFIWrapper {
  public:
   MBFIWrapper(const MachineBlockFrequencyInfo &I) : MBFI(I) {}
 
   BlockFrequency getBlockFreq(const MachineBasicBlock *MBB) const;
   void setBlockFreq(const MachineBasicBlock *MBB, BlockFrequency F);
+  Optional<uint64_t> getBlockProfileCount(const MachineBasicBlock *MBB) const;
+
   raw_ostream &printBlockFreq(raw_ostream &OS,
                               const MachineBasicBlock *MBB) const;
   raw_ostream &printBlockFreq(raw_ostream &OS,
                               const BlockFrequency Freq) const;
   void view(const Twine &Name, bool isSimple = true);
   uint64_t getEntryFreq() const;
   const MachineBlockFrequencyInfo &getMBFI() { return MBFI; }
 
  private:
   const MachineBlockFrequencyInfo &MBFI;
   DenseMap<const MachineBasicBlock *, BlockFrequency> MergedBBFreq;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_MBFIWRAPPER_H

llvm/lib/CodeGen/MBFIWrapper.cpp

BlockFrequency MBFIWrapper::getBlockFreq(const MachineBasicBlock *MBB) const {
   return MBFI.getBlockFreq(MBB);
 }
 
 void MBFIWrapper::setBlockFreq(const MachineBasicBlock *MBB,
                                BlockFrequency F) {
   MergedBBFreq[MBB] = F;
 }
 
+Optional<uint64_t>
+MBFIWrapper::getBlockProfileCount(const MachineBasicBlock *MBB) const {
+  auto I = MergedBBFreq.find(MBB);
+
+  // Modified block frequency also impacts profile count. So we should compute
+  // profile count from new block frequency if it has been changed.
+  if (I != MergedBBFreq.end())
+    return MBFI.getProfileCountFromFreq(I->second.getFrequency());
+
+  return MBFI.getBlockProfileCount(MBB);
+}
+
 raw_ostream & MBFIWrapper::printBlockFreq(raw_ostream &OS,
                                           const MachineBasicBlock *MBB) const {
   return MBFI.printBlockFreq(OS, getBlockFreq(MBB));
 }
 
 raw_ostream & MBFIWrapper::printBlockFreq(raw_ostream &OS,
                                           const BlockFrequency Freq) const {
   return MBFI.printBlockFreq(OS, Freq);

llvm/lib/CodeGen/MachineBlockPlacement.cpp

#ifndef NDEBUG
   /// blocks and their successors through the pass.
   SmallPtrSet<MachineBasicBlock *, 4> BlocksWithUnanalyzableExits;
 #endif
 
   /// Get block profile count or frequency according to UseProfileCount.
   /// The return value is used to model tail duplication cost.
   BlockFrequency getBlockCountOrFrequency(const MachineBasicBlock *BB) {
     if (UseProfileCount) {
-      auto Count = MBFI->getMBFI().getBlockProfileCount(BB);
+      auto Count = MBFI->getBlockProfileCount(BB);
       if (Count)
         return *Count;
       else
         return 0;
     } else
       return MBFI->getBlockFreq(BB);
   }