[llvm] [BOLT][NFC] Add sink block to flow CFG in profile inference (PR #95047)

[shaw young via llvm-commits]

https://github.com/shawbyoung updated https://github.com/llvm/llvm-project/pull/95047

>From 1b984ba9747618e277c8ffb2c48467aec7c2a6a3 Mon Sep 17 00:00:00 2001
From: shawbyoung <shawbyoung at gmail.com>
Date: Mon, 10 Jun 2024 12:08:38 -0700
Subject: [PATCH] [BOLT][NFC] Add sink block to flow CFG in profile inference

Test Plan: tbd

Reviewers:
Subscribers:

Tasks:

Tags:

Differential Revision: https://phabricator.intern.facebook.com/D58380996
---
 bolt/lib/Profile/StaleProfileMatching.cpp     | 38 ++++++++++++++++---
 .../Transforms/Utils/SampleProfileInference.h |  3 +-
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/bolt/lib/Profile/StaleProfileMatching.cpp b/bolt/lib/Profile/StaleProfileMatching.cpp
index 365bc5389266d..8ecfb618072ab 100644
--- a/bolt/lib/Profile/StaleProfileMatching.cpp
+++ b/bolt/lib/Profile/StaleProfileMatching.cpp
@@ -309,22 +309,33 @@ createFlowFunction(const BinaryFunction::BasicBlockOrderType &BlockOrder) {
   FlowFunction Func;
 
   // Add a special "dummy" source so that there is always a unique entry point.
-  // Because of the extra source, for all other blocks in FlowFunction it holds
-  // that Block.Index == BB->getIndex() + 1
   FlowBlock EntryBlock;
   EntryBlock.Index = 0;
   Func.Blocks.push_back(EntryBlock);
 
+  auto BinaryBlockIsExit = [&](const BinaryBasicBlock &BB) {
+    if (BB.successors().empty())
+      return true;
+    return false;
+  };
+
   // Create FlowBlock for every basic block in the binary function
   for (const BinaryBasicBlock *BB : BlockOrder) {
     Func.Blocks.emplace_back();
     FlowBlock &Block = Func.Blocks.back();
     Block.Index = Func.Blocks.size() - 1;
+    Block.HasSuccessors = BinaryBlockIsExit(*BB);
     (void)BB;
     assert(Block.Index == BB->getIndex() + 1 &&
            "incorrectly assigned basic block index");
   }
 
+  // Add a special "dummy" sink block so there is always a unique sink
+  FlowBlock SinkBlock;
+  SinkBlock.Index = Func.Blocks.size();
+  Func.Blocks.push_back(SinkBlock);
+  Func.Sink = SinkBlock.Index;
+
   // Create FlowJump for each jump between basic blocks in the binary function
   std::vector<uint64_t> InDegree(Func.Blocks.size(), 0);
   for (const BinaryBasicBlock *SrcBB : BlockOrder) {
@@ -360,18 +371,29 @@ createFlowFunction(const BinaryFunction::BasicBlockOrderType &BlockOrder) {
   // Add dummy edges to the extra sources. If there are multiple entry blocks,
   // add an unlikely edge from 0 to the subsequent ones
   assert(InDegree[0] == 0 && "dummy entry blocks shouldn't have predecessors");
-  for (uint64_t I = 1; I < Func.Blocks.size(); I++) {
+  for (uint64_t I = 1; I < BlockOrder.size() + 1; I++) {
     const BinaryBasicBlock *BB = BlockOrder[I - 1];
     if (BB->isEntryPoint() || InDegree[I] == 0) {
       Func.Jumps.emplace_back();
       FlowJump &Jump = Func.Jumps.back();
-      Jump.Source = 0;
+      Jump.Source = Func.Entry;
       Jump.Target = I;
       if (!BB->isEntryPoint())
         Jump.IsUnlikely = true;
     }
   }
 
+  // Add dummy edges from the exit blocks to the sink block.
+  for (uint64_t I = 1; I < BlockOrder.size() + 1; I++) {
+    FlowBlock &Block = Func.Blocks[I];
+    if (Block.HasSuccessors) {
+      Func.Jumps.emplace_back();
+      FlowJump &Jump = Func.Jumps.back();
+      Jump.Source = I;
+      Jump.Target = Func.Sink;
+    }
+  }
+
   // Create necessary metadata for the flow function
   for (FlowJump &Jump : Func.Jumps) {
     assert(Jump.Source < Func.Blocks.size());
@@ -379,6 +401,7 @@ createFlowFunction(const BinaryFunction::BasicBlockOrderType &BlockOrder) {
     assert(Jump.Target < Func.Blocks.size());
     Func.Blocks[Jump.Target].PredJumps.push_back(&Jump);
   }
+
   return Func;
 }
 
@@ -395,7 +418,7 @@ void matchWeightsByHashes(BinaryContext &BC,
                           const BinaryFunction::BasicBlockOrderType &BlockOrder,
                           const yaml::bolt::BinaryFunctionProfile &YamlBF,
                           FlowFunction &Func) {
-  assert(Func.Blocks.size() == BlockOrder.size() + 1);
+  assert(Func.Blocks.size() == BlockOrder.size() + 2);
 
   std::vector<FlowBlock *> Blocks;
   std::vector<BlendedBlockHash> BlendedHashes;
@@ -618,7 +641,7 @@ void assignProfile(BinaryFunction &BF,
                    FlowFunction &Func) {
   BinaryContext &BC = BF.getBinaryContext();
 
-  assert(Func.Blocks.size() == BlockOrder.size() + 1);
+  assert(Func.Blocks.size() == BlockOrder.size() + 2);
   for (uint64_t I = 0; I < BlockOrder.size(); I++) {
     FlowBlock &Block = Func.Blocks[I + 1];
     BinaryBasicBlock *BB = BlockOrder[I];
@@ -640,6 +663,9 @@ void assignProfile(BinaryFunction &BF,
       if (Jump->Flow == 0)
         continue;
 
+      // Skip the artificial sink block
+      if (Jump->Target == Func.Sink)
+        continue;
       BinaryBasicBlock &SuccBB = *BlockOrder[Jump->Target - 1];
       // Check if the edge corresponds to a regular jump or a landing pad
       if (BB->getSuccessor(SuccBB.getLabel())) {
diff --git a/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h b/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
index b4ea1ad840f9d..b2af05a24c705 100644
--- a/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
+++ b/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
@@ -31,10 +31,10 @@ struct FlowBlock {
   uint64_t Flow{0};
   std::vector<FlowJump *> SuccJumps;
   std::vector<FlowJump *> PredJumps;
+  bool HasSuccessors{false};
 
   /// Check if it is the entry block in the function.
   bool isEntry() const { return PredJumps.empty(); }
-
   /// Check if it is an exit block in the function.
   bool isExit() const { return SuccJumps.empty(); }
 };
@@ -57,6 +57,7 @@ struct FlowFunction {
   std::vector<FlowJump> Jumps;
   /// The index of the entry block.
   uint64_t Entry{0};
+  uint64_t Sink{0};
 };
 
 /// Various thresholds and options controlling the behavior of the profile