[BOLT] Add sink block to flow CFG in profile inference

[shawbyoung]

Summary: Constructing an artificial sink block for the
flow CFG in stale profile inference to allow profile
inference to be run on CFGs with blocks that terminate
and have successors.

Testing Plan: Added infer_no_exits.test to verify that
functions with exit blocks with a landing pad are
covered by stale profile inference.

[llvmbot]

@llvm/pr-subscribers-pgo
@llvm/pr-subscribers-bolt

@llvm/pr-subscribers-llvm-transforms

Author: shaw young (shawbyoung)

Changes

Constructing an artificial sink block for the flow cfg
in stale profile inference.

---

Full diff: https://github.com/llvm/llvm-project/pull/95047.diff

2 Files Affected:

• (modified) bolt/lib/Profile/StaleProfileMatching.cpp (+32-6)
• (modified) llvm/include/llvm/Transforms/Utils/SampleProfileInference.h (+2-1)

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

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[aaupov]

LG with a small nit. Thank you.

[maksfb]

Add a brief description why you are adding the sink block to the summary.

[aaupov]

This state looks good to me. @WenleiHe - do you have any unresolved questions or concerns? We'd like to keep the FlowGraph construction logic in BOLT for now.

[WenleiHe]

lgtm, thanks!