profi - a flow-based profile inference algorithm: Part III (out of 3)

This is a continuation of D109860 and D109903.

An important challenge for profile inference is caused by the fact that the
sample profile is collected on a fully optimized binary, while the block and
edge frequencies are consumed on an early stage of the compilation that operates
with a non-optimized IR. As a result, some of the basic blocks may not have
associated sample counts, and it is up to the algorithm to deduce missing
frequencies. The problem is illustrated in the figure where three basic
blocks are not present in the optimized binary and hence, receive no samples
during profiling.

We found that it is beneficial to treat all such blocks equally. Otherwise the
compiler may decide that some blocks are “cold” and apply undesirable
optimizations (e.g., hot-cold splitting) regressing the performance. Therefore,
we want to distribute the counts evenly along the blocks with missing samples.
This is achieved by a post-processing step that identifies "dangling" subgraphs
consisting of basic blocks with no sampled counts; once the subgraphs are
found, we rebalance the flow so as every branch probability is 50:50 within the
subgraphs.

Our experiments indicate up to 1% performance win using the optimization on
some binaries and a significant improvement in the quality of profile counts
(when compared to ground-truth instrumentation-based counts)

{F19093045}

Reviewed By: hoy

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

Author: spupyrev

llvm/lib/Transforms/Utils/SampleProfileInference.cpp

@@ -220,7 +220,7 @@ class MinCostMaxFlow {
       Now = Pred;
     }
 
-    assert(PathCapacity > 0 && "found incorrect augmenting path");
+    assert(PathCapacity > 0 && "found an incorrect augmenting path");
 
     // Update the flow along the path
     Now = Target;
@@ -271,7 +271,22 @@ class MinCostMaxFlow {
   uint64_t Target;
 };
 
-/// Post-processing adjustment of the control flow.
+/// A post-processing adjustment of control flow. It applies two steps by
+/// rerouting some flow and making it more realistic:
+///
+/// - First, it removes all isolated components ("islands") with a positive flow
+///   that are unreachable from the entry block. For every such component, we
+///   find the shortest from the entry to an exit passing through the component,
+///   and increase the flow by one unit along the path.
+///
+/// - Second, it identifies all "unknown subgraphs" consisting of basic blocks
+///   with no sampled counts. Then it rebalnces the flow that goes through such
+///   a subgraph so that each branch is taken with probability 50%.
+///   An unknown subgraph is such that for every two nodes u and v:
+///     - u dominates v and u is not unknown;
+///     - v post-dominates u; and
+///     - all inner-nodes of all (u,v)-paths are unknown.
+///
 class FlowAdjuster {
 public:
   FlowAdjuster(FlowFunction &Func) : Func(Func) {
@@ -281,14 +296,16 @@ class FlowAdjuster {
 
   // Run the post-processing
   void run() {
-    /// We adjust the control flow in a function so as to remove all
-    /// "isolated" components with positive flow that are unreachable
-    /// from the entry block. For every such component, we find the shortest
-    /// path from the entry to an exit passing through the component, and
-    /// increase the flow by one unit along the path.
+    /// Adjust the flow to get rid of isolated components.
     joinIsolatedComponents();
+
+    /// Rebalance the flow inside unknown subgraphs.
+    rebalanceUnknownSubgraphs();
   }
 
+  /// The probability for the first successor of a unknown subgraph
+  static constexpr double UnknownFirstSuccProbability = 0.5;
+
 private:
   void joinIsolatedComponents() {
     // Find blocks that are reachable from the source
@@ -315,7 +332,7 @@ class FlowAdjuster {
     }
   }
 
-  /// Run bfs from a given block along the jumps with a positive flow and mark
+  /// Run BFS from a given block along the jumps with a positive flow and mark
   /// all reachable blocks.
   void findReachable(uint64_t Src, std::vector<bool> &Visited) {
     if (Visited[Src])
@@ -435,6 +452,164 @@ class FlowAdjuster {
 
   uint64_t NumBlocks() const { return Func.Blocks.size(); }
 
+  /// Rebalance unknown subgraphs so as each branch splits with probabilities
+  /// UnknownFirstSuccProbability and 1 - UnknownFirstSuccProbability
+  void rebalanceUnknownSubgraphs() {
+    assert(UnknownFirstSuccProbability >= 0.0 &&
+           UnknownFirstSuccProbability <= 1.0 &&
+           "the share of the unknown successor should be between 0 and 1");
+    // Try to find unknown subgraphs from each non-unknown block
+    for (uint64_t I = 0; I < Func.Blocks.size(); I++) {
+      auto SrcBlock = &Func.Blocks[I];
+      // Do not attempt to find unknown successors from a unknown or a
+      // zero-flow block
+      if (SrcBlock->UnknownWeight || SrcBlock->Flow == 0)
+        continue;
+
+      std::vector<FlowBlock *> UnknownSuccs;
+      FlowBlock *DstBlock = nullptr;
+      // Find a unknown subgraphs starting at block SrcBlock
+      if (!findUnknownSubgraph(SrcBlock, DstBlock, UnknownSuccs))
+        continue;
+      // At the moment, we do not rebalance subgraphs containing cycles among
+      // unknown blocks
+      if (!isAcyclicSubgraph(SrcBlock, DstBlock, UnknownSuccs))
+        continue;
+
+      // Rebalance the flow
+      rebalanceUnknownSubgraph(SrcBlock, DstBlock, UnknownSuccs);
+    }
+  }
+
+  /// Find a unknown subgraph starting at block SrcBlock.
+  /// If the search is successful, the method sets DstBlock and UnknownSuccs.
+  bool findUnknownSubgraph(FlowBlock *SrcBlock, FlowBlock *&DstBlock,
+                           std::vector<FlowBlock *> &UnknownSuccs) {
+    // Run BFS from SrcBlock and make sure all paths are going through unknown
+    // blocks and end at a non-unknown DstBlock
+    auto Visited = std::vector<bool>(NumBlocks(), false);
+    std::queue<uint64_t> Queue;
+    DstBlock = nullptr;
+
+    Queue.push(SrcBlock->Index);
+    Visited[SrcBlock->Index] = true;
+    while (!Queue.empty()) {
+      auto &Block = Func.Blocks[Queue.front()];
+      Queue.pop();
+      // Process blocks reachable from Block
+      for (auto Jump : Block.SuccJumps) {
+        uint64_t Dst = Jump->Target;
+        if (Visited[Dst])
+          continue;
+        Visited[Dst] = true;
+        if (!Func.Blocks[Dst].UnknownWeight) {
+          // If we see non-unique non-unknown block reachable from SrcBlock,
+          // stop processing and skip rebalancing
+          FlowBlock *CandidateDstBlock = &Func.Blocks[Dst];
+          if (DstBlock != nullptr && DstBlock != CandidateDstBlock)
+            return false;
+          DstBlock = CandidateDstBlock;
+        } else {
+          Queue.push(Dst);
+          UnknownSuccs.push_back(&Func.Blocks[Dst]);
+        }
+      }
+    }
+
+    // If the list of unknown blocks is empty, we don't need rebalancing
+    if (UnknownSuccs.empty())
+      return false;
+    // If all reachable nodes from SrcBlock are unknown, skip rebalancing
+    if (DstBlock == nullptr)
+      return false;
+    // If any of the unknown blocks is an exit block, skip rebalancing
+    for (auto Block : UnknownSuccs) {
+      if (Block->isExit())
+        return false;
+    }
+
+    return true;
+  }
+
+  /// Verify if the given unknown subgraph is acyclic, and if yes, reorder
+  /// UnknownSuccs in the topological order (so that all jumps are "forward").
+  bool isAcyclicSubgraph(FlowBlock *SrcBlock, FlowBlock *DstBlock,
+                         std::vector<FlowBlock *> &UnknownSuccs) {
+    // Extract local in-degrees in the considered subgraph
+    auto LocalInDegree = std::vector<uint64_t>(NumBlocks(), 0);
+    for (auto Jump : SrcBlock->SuccJumps) {
+      LocalInDegree[Jump->Target]++;
+    }
+    for (uint64_t I = 0; I < UnknownSuccs.size(); I++) {
+      for (auto Jump : UnknownSuccs[I]->SuccJumps) {
+        LocalInDegree[Jump->Target]++;
+      }
+    }
+    // A loop containing SrcBlock
+    if (LocalInDegree[SrcBlock->Index] > 0)
+      return false;
+
+    std::vector<FlowBlock *> AcyclicOrder;
+    std::queue<uint64_t> Queue;
+    Queue.push(SrcBlock->Index);
+    while (!Queue.empty()) {
+      auto &Block = Func.Blocks[Queue.front()];
+      Queue.pop();
+      // Stop propagation once we reach DstBlock
+      if (Block.Index == DstBlock->Index)
+        break;
+
+      AcyclicOrder.push_back(&Block);
+      // Add to the queue all successors with zero local in-degree
+      for (auto Jump : Block.SuccJumps) {
+        uint64_t Dst = Jump->Target;
+        LocalInDegree[Dst]--;
+        if (LocalInDegree[Dst] == 0) {
+          Queue.push(Dst);
+        }
+      }
+    }
+
+    // If there is a cycle in the subgraph, AcyclicOrder contains only a subset
+    // of all blocks
+    if (UnknownSuccs.size() + 1 != AcyclicOrder.size())
+      return false;
+    UnknownSuccs = AcyclicOrder;
+    return true;
+  }
+
+  /// Rebalance a given subgraph.
+  void rebalanceUnknownSubgraph(FlowBlock *SrcBlock, FlowBlock *DstBlock,
+                                std::vector<FlowBlock *> &UnknownSuccs) {
+    assert(SrcBlock->Flow > 0 && "zero-flow block in unknown subgraph");
+    assert(UnknownSuccs.front() == SrcBlock && "incorrect order of unknowns");
+
+    for (auto Block : UnknownSuccs) {
+      // Block's flow is the sum of incoming flows
+      uint64_t TotalFlow = 0;
+      if (Block == SrcBlock) {
+        TotalFlow = Block->Flow;
+      } else {
+        for (auto Jump : Block->PredJumps) {
+          TotalFlow += Jump->Flow;
+        }
+        Block->Flow = TotalFlow;
+      }
+
+      // Process all successor jumps and update corresponding flow values
+      for (uint64_t I = 0; I < Block->SuccJumps.size(); I++) {
+        auto Jump = Block->SuccJumps[I];
+        if (I + 1 == Block->SuccJumps.size()) {
+          Jump->Flow = TotalFlow;
+          continue;
+        }
+        uint64_t Flow = uint64_t(TotalFlow * UnknownFirstSuccProbability);
+        Jump->Flow = Flow;
+        TotalFlow -= Flow;
+      }
+    }
+  }
+
   /// A constant indicating an arbitrary exit block of a function.
   static constexpr uint64_t AnyExitBlock = uint64_t(-1);
 
@@ -622,7 +797,7 @@ void verifyWeights(const FlowFunction &Func) {
     }
   }
 
-  // Run bfs from the source along edges with positive flow
+  // Run BFS from the source along edges with positive flow
   std::queue<uint64_t> Queue;
   auto Visited = std::vector<bool>(NumBlocks, false);
   Queue.push(Func.Entry);

llvm/test/Transforms/SampleProfile/Inputs/profile-inference-noprobes.prof

@@ -0,0 +1,8 @@
+test_4:100:13293
+ 0: 100
+ 1: 100
+ 2: 40
+
+empty:100:13293
+ 0: 7
+ 1: 100

llvm/test/Transforms/SampleProfile/Inputs/profile-inference-rebalance.prof

@@ -0,0 +1,27 @@
+countMultipliers:37078302:0
+ 2: 65536
+ 3: 10
+ 4: 65536
+ 7: 65536
+ 9: 65536
+ 10: 65536
+ !CFGChecksum: 223598586707
+
+countMultipliers2:37078302:0
+ 1: 2100
+ 2: 2000
+ 6: 2100
+ !CFGChecksum: 2235985
+
+countMultipliers3:37078302:0
+ 1: 100
+ 2: 100
+ 3: 100
+ !CFGChecksum: 22985
+
+countMultipliers4:37078302:0
+ 1: 100
+ 2: 50
+ 3: 50
+ 5: 100
+ !CFGChecksum: 2298578

llvm/test/Transforms/SampleProfile/profile-inference-noprobes.ll

@@ -0,0 +1,54 @@
+; RUN: opt < %s -passes=sample-profile -sample-profile-use-profi -sample-profile-file=%S/Inputs/profile-inference-noprobes.prof -S | FileCheck %s
+; RUN: opt < %s -passes=sample-profile -sample-profile-use-profi -sample-profile-file=%S/Inputs/profile-inference-noprobes.prof | opt -analyze -block-freq  -enable-new-pm=0 | FileCheck %s --check-prefix=CHECK2
+
+
+; The test verifies that profile inference can be applied for non-probe-based
+; profiles.
+;
+; +---------+     +----------+
+; | b3 [40] | <-- | b1 [100] |
+; +---------+     +----------+
+;                   |
+;                   |
+;                   v
+;                 +----------+
+;                 | b2 [0]   |
+;                 +----------+
+
+@yydebug = dso_local global i32 0, align 4
+
+define void @test_4() #0 !dbg !4 {
+;entry:
+;  ret void, !dbg !9
+b1:
+  %0 = load i32, i32* @yydebug, align 4
+  %cmp = icmp ne i32 %0, 0, !dbg !9
+  br i1 %cmp, label %b2, label %b3, !dbg !9
+; CHECK2: - b1: float = {{.*}}, int = {{.*}}, count = 100
+
+b2:
+  ret void
+; CHECK2: - b2: float = {{.*}}, int = {{.*}}, count = 60
+
+b3:
+  ret void, !dbg !10
+; CHECK2: - b3: float = {{.*}}, int = {{.*}}, count = 40
+}
+
+; CHECK: {{.*}} = !{!"function_entry_count", i64 100}
+; CHECK: {{.*}} = !{!"branch_weights", i32 60, i32 40}
+
+attributes #0 = { noinline nounwind uwtable "use-sample-profile"}
+
+!llvm.module.flags = !{!6, !7}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, producer: "clang version 3.7.0 (trunk 237249) (llvm/trunk 237261)", isOptimized: false, runtimeVersion: 0, emissionKind: LineTablesOnly, enums: !2, retainedTypes: !2, globals: !2, imports: !2)
+!1 = !DIFile(filename: "entry_counts.c", directory: ".")
+!2 = !{}
+!4 = distinct !DISubprogram(name: "test_4", scope: !1, file: !1, line: 1, type: !5, isLocal: false, isDefinition: true, scopeLine: 1, isOptimized: false, unit: !0, retainedNodes: !2)
+!5 = !DISubroutineType(types: !2)
+!6 = !{i32 2, !"Dwarf Version", i32 4}
+!7 = !{i32 2, !"Debug Info Version", i32 3}
+!8 = !{!"clang version 3.7.0 (trunk 237249) (llvm/trunk 237261)"}
+!9 = !DILocation(line: 1, column: 15, scope: !4)
+!10 = !DILocation(line: 3, column: 15, scope: !4)