From d7c7affacaaea3aabfe21a77443a6d82fce374f7 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin@google.com>
Date: Tue, 28 Jul 2020 13:04:11 -0700
Subject: [PATCH] [llvm][NFC] refactor setBlockFrequency for clarity.

The refactoring encapsulates frequency calculation in MachineBlockFrequencyInfo,
and renames the API to clarify its motivation. It should clarify
frequencies may not be reset 'freely' by users of the analysis, as the
API serves as a partial update to avoid a full analysis recomputation.

Differential Revision: https://reviews.llvm.org/D84427
---
 include/llvm/CodeGen/MachineBlockFrequencyInfo.h |  6 +++++-
 lib/CodeGen/MachineBlockFrequencyInfo.cpp        | 11 ++++++++---
 lib/CodeGen/MachineSink.cpp                      |  8 +++-----
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/include/llvm/CodeGen/MachineBlockFrequencyInfo.h b/include/llvm/CodeGen/MachineBlockFrequencyInfo.h
index 7ce11c784b0..6c442d3d07b 100644
--- a/include/llvm/CodeGen/MachineBlockFrequencyInfo.h
+++ b/include/llvm/CodeGen/MachineBlockFrequencyInfo.h
@@ -74,7 +74,11 @@ public:
 
   bool isIrrLoopHeader(const MachineBasicBlock *MBB) const;
 
-  void setBlockFreq(const MachineBasicBlock *MBB, uint64_t Freq);
+  /// incrementally calculate block frequencies when we split edges, to avoid
+  /// full CFG traversal.
+  void onEdgeSplit(const MachineBasicBlock &NewPredecessor,
+                   const MachineBasicBlock &NewSuccessor,
+                   const MachineBranchProbabilityInfo &MBPI);
 
   const MachineFunction *getFunction() const;
   const MachineBranchProbabilityInfo *getMBPI() const;
diff --git a/lib/CodeGen/MachineBlockFrequencyInfo.cpp b/lib/CodeGen/MachineBlockFrequencyInfo.cpp
index c8369970601..54e0a14e055 100644
--- a/lib/CodeGen/MachineBlockFrequencyInfo.cpp
+++ b/lib/CodeGen/MachineBlockFrequencyInfo.cpp
@@ -247,10 +247,15 @@ bool MachineBlockFrequencyInfo::isIrrLoopHeader(
   return MBFI->isIrrLoopHeader(MBB);
 }
 
-void MachineBlockFrequencyInfo::setBlockFreq(const MachineBasicBlock *MBB,
-                                             uint64_t Freq) {
+void MachineBlockFrequencyInfo::onEdgeSplit(
+    const MachineBasicBlock &NewPredecessor,
+    const MachineBasicBlock &NewSuccessor,
+    const MachineBranchProbabilityInfo &MBPI) {
   assert(MBFI && "Expected analysis to be available");
-  MBFI->setBlockFreq(MBB, Freq);
+  auto NewSuccFreq = MBFI->getBlockFreq(&NewPredecessor) *
+                     MBPI.getEdgeProbability(&NewPredecessor, &NewSuccessor);
+
+  MBFI->setBlockFreq(&NewSuccessor, NewSuccFreq.getFrequency());
 }
 
 const MachineFunction *MachineBlockFrequencyInfo::getFunction() const {
diff --git a/lib/CodeGen/MachineSink.cpp b/lib/CodeGen/MachineSink.cpp
index 5f958bbc31b..dfb88f6bf51 100644
--- a/lib/CodeGen/MachineSink.cpp
+++ b/lib/CodeGen/MachineSink.cpp
@@ -347,11 +347,9 @@ bool MachineSinking::runOnMachineFunction(MachineFunction &MF) {
                           << printMBBReference(*Pair.first) << " -- "
                           << printMBBReference(*NewSucc) << " -- "
                           << printMBBReference(*Pair.second) << '\n');
-        if (MBFI) {
-          auto NewSuccFreq = MBFI->getBlockFreq(Pair.first) *
-                             MBPI->getEdgeProbability(Pair.first, NewSucc);
-          MBFI->setBlockFreq(NewSucc, NewSuccFreq.getFrequency());
-        }
+        if (MBFI)
+          MBFI->onEdgeSplit(*Pair.first, *NewSucc, *MBPI);
+
         MadeChange = true;
         ++NumSplit;
       } else