[BOLT] introducing profi params

[spupyrev]

Summary

We want to use profile inference (profi) in BOLT for stale profile matching.
To this end, I am making a few changes modifying the interface of the algorithm.
This is the first change for existing usages of profi (e.g., CSSPGO):

• introducing an object holding the algorithmic parameters;
• some renaming of existing options;
• dropped unused option, SampleProfileInferEntryCount, as we don't plan to change its default value;
• no changes in the output / tests.

[spupyrev]

fix formatting

[spupyrev]

rebase & some updates

[spupyrev]

fixing assert

[spupyrev]

rebase

[hoy]

llvm/lib/Transforms/Utils/SampleProfileInference.cpp

1039

Would it make sense to move this into ProfiParams?

[spupyrev]

llvm/lib/Transforms/Utils/SampleProfileInference.cpp

1039

One earlier version of the diff had this constant in ProfiParams. Later I decided to use ProfiParams as a set of user-configurable options, while MinBaseDistance is more like a constant which shouldn't be changed.

I don't have a strong preference and see pros for both ways. What do you think?

[hoy]

LGTM.

llvm/lib/Transforms/Utils/SampleProfileInference.cpp

1039

I see. I don't have a strong opinion either, was just wondering if it is intentional.

[spupyrev]

rebase

[Amir]

Why was this specific change needed?

introducing an object holding the algorithmic parameters;

Is it in order to specify different BOLT defaults? Or for the theoretical use case where BOLT and Clang are both invoked from the same driver, as in llvm-driver build mode?

[spupyrev]

In D134756#4397850, @Amir wrote:

Why was this specific change needed?

introducing an object holding the algorithmic parameters;

Is it in order to specify different BOLT defaults? Or for the theoretical use case where BOLT and Clang are both invoked from the same driver, as in llvm-driver build mode?

BOLT and CSSPGO are using different profi parameters, as they're based on different profiles. This diff allows to pass the params for the algorithm

[Amir]

In D134756#4399936, @spupyrev wrote:

In D134756#4397850, @Amir wrote:

Why was this specific change needed?

introducing an object holding the algorithmic parameters;

Is it in order to specify different BOLT defaults? Or for the theoretical use case where BOLT and Clang are both invoked from the same driver, as in llvm-driver build mode?

BOLT and CSSPGO are using different profi parameters, as they're based on different profiles. This diff allows to pass the params for the algorithm

In that case, it's possible to override these parameters from BOLT – just like we do for some other knobs in adjustCommandLineOptions.

llvm/include/llvm/Transforms/Utils/SampleProfileInference.h

 } // end namespace afdo_detail
 
 struct FlowJump;
 
 /// A wrapper of a binary basic block.
 struct FlowBlock {
   uint64_t Index;
   uint64_t Weight{0};
-  bool UnknownWeight{false};
+  bool HasUnknownWeight{false};
   uint64_t Flow{0};
   bool HasSelfEdge{false};
   std::vector<FlowJump *> SuccJumps;
   std::vector<FlowJump *> PredJumps;
 
   /// Check if it is the entry block in the function.
   bool isEntry() const { return PredJumps.empty(); }
 
struct FlowJump {
   bool IsUnlikely{false};
 };
 
 /// A wrapper of binary function with basic blocks and jumps.
 struct FlowFunction {
   std::vector<FlowBlock> Blocks;
   std::vector<FlowJump> Jumps;
   /// The index of the entry block.
-  uint64_t Entry;
+  uint64_t Entry{0};
 };
 
+/// Various thresholds and options controlling the behavior of the profile
+/// inference algorithm. Default values are tuned for several large-scale
+/// applications, and can be modified via corresponding command-line flags.
+struct ProfiParams {
+  /// Evenly distribute flow when there are multiple equally likely options.
+  bool EvenFlowDistribution{false};
+
+  /// Evenly re-distribute flow among unknown subgraphs.
+  bool RebalanceUnknown{false};
+
+  /// Join isolated components having positive flow.
+  bool JoinIslands{false};
+
+  /// The cost of increasing a block's count by one.
+  unsigned CostBlockInc{0};
+
+  /// The cost of decreasing a block's count by one.
+  unsigned CostBlockDec{0};
+
+  /// The cost of increasing a count of zero-weight block by one.
+  unsigned CostBlockZeroInc{0};
+
+  /// The cost of increasing the entry block's count by one.
+  unsigned CostBlockEntryInc{0};
+
+  /// The cost of decreasing the entry block's count by one.
+  unsigned CostBlockEntryDec{0};
+
+  /// The cost of increasing an unknown block's count by one.
+  unsigned CostBlockUnknownInc{0};
+
+  /// The cost of taking an unlikely block/jump.
+  const int64_t CostUnlikely = ((int64_t)1) << 30;
+};
+
+void applyFlowInference(const ProfiParams &Params, FlowFunction &Func);
 void applyFlowInference(FlowFunction &Func);
 
 /// Sample profile inference pass.
 template <typename BT> class SampleProfileInference {
 public:
   using BasicBlockT = typename afdo_detail::TypeMap<BT>::BasicBlockT;
   using FunctionT = typename afdo_detail::TypeMap<BT>::FunctionT;
   using Edge = std::pair<const BasicBlockT *, const BasicBlockT *>;
void SampleProfileInference<BT>::apply(BlockWeightMap &BlockWeights,
 
   // Create necessary objects
   FlowFunction Func;
   Func.Blocks.reserve(BasicBlocks.size());
   // Create FlowBlocks
   for (const auto *BB : BasicBlocks) {
     FlowBlock Block;
     if (SampleBlockWeights.find(BB) != SampleBlockWeights.end()) {
-      Block.UnknownWeight = false;
+      Block.HasUnknownWeight = false;
       Block.Weight = SampleBlockWeights[BB];
     } else {
-      Block.UnknownWeight = true;
+      Block.HasUnknownWeight = true;
       Block.Weight = 0;
     }
     Block.Index = Func.Blocks.size();
     Func.Blocks.push_back(Block);
   }
   // Create FlowEdges
   for (const auto *BB : BasicBlocks) {
     for (auto *Succ : Successors[BB]) {

llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h

template <> struct IRTraits<BasicBlock> {
   }
   static pred_range getPredecessors(BasicBlock *BB) { return predecessors(BB); }
   static succ_range getSuccessors(BasicBlock *BB) { return successors(BB); }
 };
 
 } // end namespace afdo_detail
 
 extern cl::opt<bool> SampleProfileUseProfi;
-extern cl::opt<bool> SampleProfileInferEntryCount;
 
 template <typename BT> class SampleProfileLoaderBaseImpl {
 public:
   SampleProfileLoaderBaseImpl(std::string Name, std::string RemapName)
       : Filename(Name), RemappingFilename(RemapName) {}
   void dump() { Reader->dump(); }
 
   using InstructionT = typename afdo_detail::IRTraits<BT>::InstructionT;
void SampleProfileLoaderBaseImpl<BT>::finalizeWeightPropagation(
   // counts and set the entry count as computed by the algorithm. This is
   // primarily done to sync the counts produced by profi and BFI inference,
   // which uses the entry count for mass propagation.
   // If profi produces a zero-value for the entry count, we fallback to
   // Samples->getHeadSamples() + 1 to avoid functions with zero count.
   if (SampleProfileUseProfi) {
     const BasicBlockT *EntryBB = getEntryBB(&F);
     ErrorOr<uint64_t> EntryWeight = getBlockWeight(EntryBB);
-    if (BlockWeights[EntryBB] > 0 &&
-        (SampleProfileInferEntryCount || !EntryWeight)) {
+    if (BlockWeights[EntryBB] > 0) {
       getFunction(F).setEntryCount(
           ProfileCount(BlockWeights[EntryBB], Function::PCT_Real),
           &InlinedGUIDs);
     }
   }
 }
 
 template <typename BT>

llvm/lib/Transforms/Utils/SampleProfileInference.cpp

 #include <set>
 #include <stack>
 
 using namespace llvm;
 #define DEBUG_TYPE "sample-profile-inference"
 
 namespace {
 
-static cl::opt<bool> SampleProfileEvenCountDistribution(
-    "sample-profile-even-count-distribution", cl::init(true), cl::Hidden,
-    cl::desc("Try to evenly distribute counts when there are multiple equally "
+static cl::opt<bool> SampleProfileEvenFlowDistribution(
+    "sample-profile-even-flow-distribution", cl::init(true), cl::Hidden,
+    cl::desc("Try to evenly distribute flow when there are multiple equally "
              "likely options."));
 
-static cl::opt<unsigned> SampleProfileMaxDfsCalls(
-    "sample-profile-max-dfs-calls", cl::init(10), cl::Hidden,
-    cl::desc("Maximum number of dfs iterations for even count distribution."));
-
-static cl::opt<unsigned> SampleProfileProfiCostInc(
-    "sample-profile-profi-cost-inc", cl::init(10), cl::Hidden,
-    cl::desc("A cost of increasing a block's count by one."));
-
-static cl::opt<unsigned> SampleProfileProfiCostDec(
-    "sample-profile-profi-cost-dec", cl::init(20), cl::Hidden,
-    cl::desc("A cost of decreasing a block's count by one."));
-
-static cl::opt<unsigned> SampleProfileProfiCostIncZero(
-    "sample-profile-profi-cost-inc-zero", cl::init(11), cl::Hidden,
-    cl::desc("A cost of increasing a count of zero-weight block by one."));
-
-static cl::opt<unsigned> SampleProfileProfiCostIncEntry(
-    "sample-profile-profi-cost-inc-entry", cl::init(40), cl::Hidden,
-    cl::desc("A cost of increasing the entry block's count by one."));
-
-static cl::opt<unsigned> SampleProfileProfiCostDecEntry(
-    "sample-profile-profi-cost-dec-entry", cl::init(10), cl::Hidden,
-    cl::desc("A cost of decreasing the entry block's count by one."));
+static cl::opt<bool> SampleProfileRebalanceUnknown(
+    "sample-profile-rebalance-unknown", cl::init(true), cl::Hidden,
+    cl::desc("Evenly re-distribute flow among unknown subgraphs."));
+
+static cl::opt<bool> SampleProfileJoinIslands(
+    "sample-profile-join-islands", cl::init(true), cl::Hidden,
+    cl::desc("Join isolated components having positive flow."));
+
+static cl::opt<unsigned> SampleProfileProfiCostBlockInc(
+    "sample-profile-profi-cost-block-inc", cl::init(10), cl::Hidden,
+    cl::desc("The cost of increasing a block's count by one."));
+
+static cl::opt<unsigned> SampleProfileProfiCostBlockDec(
+    "sample-profile-profi-cost-block-dec", cl::init(20), cl::Hidden,
+    cl::desc("The cost of decreasing a block's count by one."));
+
+static cl::opt<unsigned> SampleProfileProfiCostBlockEntryInc(
+    "sample-profile-profi-cost-block-entry-inc", cl::init(40), cl::Hidden,
+    cl::desc("The cost of increasing the entry block's count by one."));
+
+static cl::opt<unsigned> SampleProfileProfiCostBlockEntryDec(
+    "sample-profile-profi-cost-block-entry-dec", cl::init(10), cl::Hidden,
+    cl::desc("The cost of decreasing the entry block's count by one."));
+
+static cl::opt<unsigned> SampleProfileProfiCostBlockZeroInc(
+    "sample-profile-profi-cost-block-zero-inc", cl::init(11), cl::Hidden,
+    cl::desc("The cost of increasing a count of zero-weight block by one."));
+
+static cl::opt<unsigned> SampleProfileProfiCostBlockUnknownInc(
+    "sample-profile-profi-cost-block-unknown-inc", cl::init(0), cl::Hidden,
+    cl::desc("The cost of increasing an unknown block's count by one."));
 
 /// A value indicating an infinite flow/capacity/weight of a block/edge.
 /// Not using numeric_limits<int64_t>::max(), as the values can be summed up
 /// during the execution.
 static constexpr int64_t INF = ((int64_t)1) << 50;
 
 /// The minimum-cost maximum flow algorithm.
 ///
 /// The algorithm finds the maximum flow of minimum cost on a given (directed)
 /// network using a modified version of the classical Moore-Bellman-Ford
 /// approach. The algorithm applies a number of augmentation iterations in which
 /// flow is sent along paths of positive capacity from the source to the sink.
 /// The worst-case time complexity of the implementation is O(v(f)*m*n), where
 /// where m is the number of edges, n is the number of vertices, and v(f) is the
 /// value of the maximum flow. However, the observed running time on typical
 /// instances is sub-quadratic, that is, o(n^2).
 ///
 /// The input is a set of edges with specified costs and capacities, and a pair
 /// of nodes (source and sink). The output is the flow along each edge of the
 /// minimum total cost respecting the given edge capacities.
 class MinCostMaxFlow {
 public:
+  MinCostMaxFlow(const ProfiParams &Params) : Params(Params) {}
+
   // Initialize algorithm's data structures for a network of a given size.
   void initialize(uint64_t NodeCount, uint64_t SourceNode, uint64_t SinkNode) {
     Source = SourceNode;
     Target = SinkNode;
 
     Nodes = std::vector<Node>(NodeCount);
     Edges = std::vector<std::vector<Edge>>(NodeCount, std::vector<Edge>());
-    if (SampleProfileEvenCountDistribution)
+    if (Params.EvenFlowDistribution)
       AugmentingEdges =
           std::vector<std::vector<Edge *>>(NodeCount, std::vector<Edge *>());
   }
 
   // Run the algorithm.
   int64_t run() {
     // Iteratively find an augmentation path/dag in the network and send the
     // flow along its edges
int64_t getFlow(uint64_t Src, uint64_t Dst) const {
     for (const auto &Edge : Edges[Src]) {
       if (Edge.Dst == Dst) {
         Flow += Edge.Flow;
       }
     }
     return Flow;
   }
 
-  /// A cost of taking an unlikely jump.
-  static constexpr int64_t AuxCostUnlikely = ((int64_t)1) << 30;
-  /// Minimum BaseDistance for the jump distance values in island joining.
-  static constexpr uint64_t MinBaseDistance = 10000;
-
 private:
   /// Iteratively find an augmentation path/dag in the network and send the
   /// flow along its edges. The method returns the number of applied iterations.
   size_t applyFlowAugmentation() {
     size_t AugmentationIters = 0;
     while (findAugmentingPath()) {
       uint64_t PathCapacity = computeAugmentingPathCapacity();
       while (PathCapacity > 0) {
         bool Progress = false;
-        if (SampleProfileEvenCountDistribution) {
+        if (Params.EvenFlowDistribution) {
           // Identify node/edge candidates for augmentation
           identifyShortestEdges(PathCapacity);
 
           // Find an augmenting DAG
           auto AugmentingOrder = findAugmentingDAG();
 
           // Apply the DAG augmentation
           Progress = augmentFlowAlongDAG(AugmentingOrder);
while (!Queue.empty()) {
       // the following early-stop criteria:
       // -- If we find a path with zero-distance from Source to Target, stop the
       //    search, as the path is the shortest since Dist[Source, Target] >= 0;
       // -- If we have Dist[Source, V] > Dist[Source, Target], then do not
       //    process node V, as it is guaranteed _not_ to be on a shortest path
       //    from Source to Target; it follows from inequalities
       //    Dist[Source, Target] >= Dist[Source, V] + Dist[V, Target]
       //                         >= Dist[Source, V]
-      if (!SampleProfileEvenCountDistribution && Nodes[Target].Distance == 0)
+      if (!Params.EvenFlowDistribution && Nodes[Target].Distance == 0)
         break;
       if (Nodes[Src].Distance > Nodes[Target].Distance)
         continue;
 
       // Process adjacent edges
       for (uint64_t EdgeIdx = 0; EdgeIdx < Edges[Src].size(); EdgeIdx++) {
         auto &Edge = Edges[Src][EdgeIdx];
         if (Edge.Flow < Edge.Capacity) {
while (!Stack.empty()) {
       // If we haven't scanned all edges out of NodeIdx, continue scanning
       if (EdgeIdx < Edges[NodeIdx].size()) {
         auto &Edge = Edges[NodeIdx][EdgeIdx];
         auto &Dst = Nodes[Edge.Dst];
         Stack.top().second++;
 
         if (Edge.OnShortestPath) {
           // If we haven't seen Edge.Dst so far, continue DFS search there
-          if (Dst.Discovery == 0 && Dst.NumCalls < SampleProfileMaxDfsCalls) {
+          if (Dst.Discovery == 0 && Dst.NumCalls < MaxDfsCalls) {
             Dst.Discovery = ++Time;
             Stack.emplace(Edge.Dst, 0);
             Dst.NumCalls++;
           } else if (Dst.Taken && Dst.Finish != 0) {
             // Else, if Edge.Dst already have a path to Target, so that NodeIdx
             Nodes[NodeIdx].Taken = true;
           }
         }
for (size_t Src = 0; Src < Nodes.size(); Src++) {
             Nodes[Dst].Distance <= Nodes[Target].Distance &&
             Nodes[Dst].Distance == Nodes[Src].Distance + Edge.Cost &&
             Edge.Capacity > Edge.Flow &&
             uint64_t(Edge.Capacity - Edge.Flow) >= MinCapacity;
       }
     }
   }
 
+  /// Maximum number of DFS iterations for DAG finding.
+  static constexpr uint64_t MaxDfsCalls = 10;
+
   /// A node in a flow network.
   struct Node {
     /// The cost of the cheapest path from the source to the current node.
     int64_t Distance;
     /// The node preceding the current one in the path.
     uint64_t ParentNode;
     /// The index of the edge between ParentNode and the current node.
     uint64_t ParentEdgeIndex;
private:
   /// The set of network edges.
   std::vector<std::vector<Edge>> Edges;
   /// Source node of the flow.
   uint64_t Source;
   /// Target (sink) node of the flow.
   uint64_t Target;
   /// Augmenting edges.
   std::vector<std::vector<Edge *>> AugmentingEdges;
+  /// Params for flow computation.
+  const ProfiParams &Params;
 };
 
 /// 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) {
+  FlowAdjuster(const ProfiParams &Params, FlowFunction &Func)
+      : Params(Params), Func(Func) {
     assert(Func.Blocks[Func.Entry].isEntry() &&
            "incorrect index of the entry block");
   }
 
   // Run the post-processing
   void run() {
+    if (Params.JoinIslands) {
       /// Adjust the flow to get rid of isolated components.
       joinIsolatedComponents();
+    }
 
+    if (Params.RebalanceUnknown) {
       /// Rebalance the flow inside unknown subgraphs.
       rebalanceUnknownSubgraphs();
     }
+  }
 
 private:
   void joinIsolatedComponents() {
     // Find blocks that are reachable from the source
     auto Visited = BitVector(NumBlocks(), false);
     findReachable(Func.Entry, Visited);
 
     // Iterate over all non-reachable blocks and adjust their weights
private:
   /// and avoid changing branch probability of outgoing edges drastically,
   /// set the jump distance so as:
   ///   - to minimize the number of unlikely jumps used and subject to that,
   ///   - to minimize the number of Flow == 0 jumps used and subject to that,
   ///   - minimizes total multiplicative Flow increase for the remaining edges.
   /// To capture this objective with integer distances, we round off fractional
   /// parts to a multiple of 1 / BaseDistance.
   int64_t jumpDistance(FlowJump *Jump) const {
+    if (Jump->IsUnlikely)
+      return Params.CostUnlikely;
+
     uint64_t BaseDistance =
-        std::max(MinCostMaxFlow::MinBaseDistance,
+        std::max(FlowAdjuster::MinBaseDistance,
                  std::min(Func.Blocks[Func.Entry].Flow,
-                          MinCostMaxFlow::AuxCostUnlikely / NumBlocks()));
-    if (Jump->IsUnlikely)
-      return MinCostMaxFlow::AuxCostUnlikely;
+                          Params.CostUnlikely / NumBlocks()));
     if (Jump->Flow > 0)
       return BaseDistance + BaseDistance / Jump->Flow;
     return BaseDistance * NumBlocks();
   };
 
   uint64_t NumBlocks() const { return Func.Blocks.size(); }
 
   /// Rebalance unknown subgraphs so that the flow is split evenly across the
for (const FlowBlock &SrcBlock : Func.Blocks) {
       rebalanceUnknownSubgraph(&SrcBlock, DstBlock, UnknownBlocks);
     }
   }
 
   /// Verify if rebalancing rooted at a given block is possible.
   bool canRebalanceAtRoot(const FlowBlock *SrcBlock) {
     // Do not attempt to find unknown subgraphs from an unknown or a
     // zero-flow block
-    if (SrcBlock->UnknownWeight || SrcBlock->Flow == 0)
+    if (SrcBlock->HasUnknownWeight || SrcBlock->Flow == 0)
       return false;
 
     // Do not attempt to process subgraphs from a block w/o unknown sucessors
     bool HasUnknownSuccs = false;
     for (auto *Jump : SrcBlock->SuccJumps) {
-      if (Func.Blocks[Jump->Target].UnknownWeight) {
+      if (Func.Blocks[Jump->Target].HasUnknownWeight) {
         HasUnknownSuccs = true;
         break;
       }
     }
     if (!HasUnknownSuccs)
       return false;
 
     return true;
while (!Queue.empty()) {
           continue;
 
         uint64_t Dst = Jump->Target;
         // If Dst has been visited, skip Jump
         if (Visited[Dst])
           continue;
         // Process block Dst
         Visited[Dst] = true;
-        if (!Func.Blocks[Dst].UnknownWeight) {
+        if (!Func.Blocks[Dst].HasUnknownWeight) {
           KnownDstBlocks.push_back(&Func.Blocks[Dst]);
         } else {
           Queue.push(Dst);
           UnknownBlocks.push_back(&Func.Blocks[Dst]);
         }
       }
     }
   }
bool ignoreJump(const FlowBlock *SrcBlock, const FlowBlock *DstBlock,
     auto JumpSource = &Func.Blocks[Jump->Source];
     auto JumpTarget = &Func.Blocks[Jump->Target];
 
     // Do not ignore jumps coming into DstBlock
     if (DstBlock != nullptr && JumpTarget == DstBlock)
       return false;
 
     // Ignore jumps out of SrcBlock to known blocks
-    if (!JumpTarget->UnknownWeight && JumpSource == SrcBlock)
+    if (!JumpTarget->HasUnknownWeight && JumpSource == SrcBlock)
       return true;
 
     // Ignore jumps to known blocks with zero flow
-    if (!JumpTarget->UnknownWeight && JumpTarget->Flow == 0)
+    if (!JumpTarget->HasUnknownWeight && JumpTarget->Flow == 0)
       return true;
 
     return false;
   }
 
   /// Verify if the given unknown subgraph is acyclic, and if yes, reorder
   /// UnknownBlocks in the topological order (so that all jumps are "forward").
   bool isAcyclicSubgraph(const FlowBlock *SrcBlock, const FlowBlock *DstBlock,
bool isAcyclicSubgraph(const FlowBlock *SrcBlock, const FlowBlock *DstBlock,
     while (!Queue.empty()) {
       FlowBlock *Block = &Func.Blocks[Queue.front()];
       Queue.pop();
       // Stop propagation once we reach DstBlock, if any
       if (DstBlock != nullptr && Block == DstBlock)
         break;
 
       // Keep an acyclic order of unknown blocks
-      if (Block->UnknownWeight && Block != SrcBlock)
+      if (Block->HasUnknownWeight && Block != SrcBlock)
         AcyclicOrder.push_back(Block);
 
       // Add to the queue all successors with zero local in-degree
       for (auto *Jump : Block->SuccJumps) {
         if (ignoreJump(SrcBlock, DstBlock, Jump))
           continue;
         uint64_t Dst = Jump->Target;
         LocalInDegree[Dst]--;
for (auto *Jump : SrcBlock->SuccJumps) {
       if (ignoreJump(SrcBlock, DstBlock, Jump))
         continue;
       BlockFlow += Jump->Flow;
     }
     rebalanceBlock(SrcBlock, DstBlock, SrcBlock, BlockFlow);
 
     // Ditribute flow from the remaining blocks
     for (auto *Block : UnknownBlocks) {
-      assert(Block->UnknownWeight && "incorrect unknown subgraph");
+      assert(Block->HasUnknownWeight && "incorrect unknown subgraph");
       uint64_t BlockFlow = 0;
       // Block's flow is the sum of incoming flows
       for (auto *Jump : Block->PredJumps) {
         BlockFlow += Jump->Flow;
       }
       Block->Flow = BlockFlow;
       rebalanceBlock(SrcBlock, DstBlock, Block, BlockFlow);
     }
for (auto *Jump : Block->SuccJumps) {
       Jump->Flow = Flow;
       BlockFlow -= Flow;
     }
     assert(BlockFlow == 0 && "not all flow is propagated");
   }
 
   /// A constant indicating an arbitrary exit block of a function.
   static constexpr uint64_t AnyExitBlock = uint64_t(-1);
+  /// Minimum BaseDistance for the jump distance values in island joining.
+  static constexpr uint64_t MinBaseDistance = 10000;
 
+  /// Params for flow computation.
+  const ProfiParams &Params;
   /// The function.
   FlowFunction &Func;
 };
 
 /// Initializing flow network for a given function.
 ///
 /// Every block is split into three nodes that are responsible for (i) an
 /// incoming flow, (ii) an outgoing flow, and (iii) penalizing an increase or
 /// reduction of the block weight.
-void initializeNetwork(MinCostMaxFlow &Network, FlowFunction &Func) {
+void initializeNetwork(const ProfiParams &Params, MinCostMaxFlow &Network,
+                       FlowFunction &Func) {
   uint64_t NumBlocks = Func.Blocks.size();
   assert(NumBlocks > 1 && "Too few blocks in a function");
   LLVM_DEBUG(dbgs() << "Initializing profi for " << NumBlocks << " blocks\n");
 
   // Pre-process data: make sure the entry weight is at least 1
   if (Func.Blocks[Func.Entry].Weight == 0) {
     Func.Blocks[Func.Entry].Weight = 1;
   }
   // Introducing dummy source/sink pairs to allow flow circulation.
   // The nodes corresponding to blocks of Func have indicies in the range
   // [0..3 * NumBlocks); the dummy nodes are indexed by the next four values.
   uint64_t S = 3 * NumBlocks;
   uint64_t T = S + 1;
   uint64_t S1 = S + 2;
   uint64_t T1 = S + 3;
 
   Network.initialize(3 * NumBlocks + 4, S1, T1);
 
   // Create three nodes for every block of the function
   for (uint64_t B = 0; B < NumBlocks; B++) {
     auto &Block = Func.Blocks[B];
-    assert((!Block.UnknownWeight || Block.Weight == 0 || Block.isEntry()) &&
+    assert((!Block.HasUnknownWeight || Block.Weight == 0 || Block.isEntry()) &&
            "non-zero weight of a block w/o weight except for an entry");
 
     // Split every block into two nodes
     uint64_t Bin = 3 * B;
     uint64_t Bout = 3 * B + 1;
     uint64_t Baux = 3 * B + 2;
     if (Block.Weight > 0) {
       Network.addEdge(S1, Bout, Block.Weight, 0);
       Network.addEdge(Bin, T1, Block.Weight, 0);
     }
 
     // Edges from S and to T
     assert((!Block.isEntry() || !Block.isExit()) &&
            "a block cannot be an entry and an exit");
     if (Block.isEntry()) {
       Network.addEdge(S, Bin, 0);
     } else if (Block.isExit()) {
       Network.addEdge(Bout, T, 0);
     }
 
     // An auxiliary node to allow increase/reduction of block counts:
     // We assume that decreasing block counts is more expensive than increasing,
     // and thus, setting separate costs here. In the future we may want to tune
     // the relative costs so as to maximize the quality of generated profiles.
-    int64_t AuxCostInc = SampleProfileProfiCostInc;
-    int64_t AuxCostDec = SampleProfileProfiCostDec;
-    if (Block.UnknownWeight) {
+    int64_t AuxCostInc = Params.CostBlockInc;
+    int64_t AuxCostDec = Params.CostBlockDec;
+    if (Block.HasUnknownWeight) {
       // Do not penalize changing weights of blocks w/o known profile count
-      AuxCostInc = 0;
+      AuxCostInc = Params.CostBlockUnknownInc;
       AuxCostDec = 0;
     } else {
       // Increasing the count for "cold" blocks with zero initial count is more
       // expensive than for "hot" ones
       if (Block.Weight == 0) {
-        AuxCostInc = SampleProfileProfiCostIncZero;
+        AuxCostInc = Params.CostBlockZeroInc;
       }
       // Modifying the count of the entry block is expensive
       if (Block.isEntry()) {
-        AuxCostInc = SampleProfileProfiCostIncEntry;
-        AuxCostDec = SampleProfileProfiCostDecEntry;
+        AuxCostInc = Params.CostBlockEntryInc;
+        AuxCostDec = Params.CostBlockEntryDec;
       }
     }
     // For blocks with self-edges, do not penalize a reduction of the count,
     // as all of the increase can be attributed to the self-edge
     if (Block.HasSelfEdge) {
       AuxCostDec = 0;
     }
 
     Network.addEdge(Bin, Baux, AuxCostInc);
     Network.addEdge(Baux, Bout, AuxCostInc);
     if (Block.Weight > 0) {
       Network.addEdge(Bout, Baux, AuxCostDec);
       Network.addEdge(Baux, Bin, AuxCostDec);
     }
   }
 
   // Creating edges for every jump
   for (auto &Jump : Func.Jumps) {
     uint64_t Src = Jump.Source;
     uint64_t Dst = Jump.Target;
     if (Src != Dst) {
       uint64_t SrcOut = 3 * Src + 1;
       uint64_t DstIn = 3 * Dst;
-      uint64_t Cost = Jump.IsUnlikely ? MinCostMaxFlow::AuxCostUnlikely : 0;
+      uint64_t Cost = Jump.IsUnlikely ? Params.CostUnlikely : 0;
       Network.addEdge(SrcOut, DstIn, Cost);
     }
   }
 
   // Make sure we have a valid flow circulation
   Network.addEdge(T, S, 0);
 }
 
for (uint64_t I = 0; I < NumBlocks; I++) {
     assert((Visited[I] || Block.Flow == 0) && "an isolated flow component");
   }
 }
 #endif
 
 } // end of anonymous namespace
 
 /// Apply the profile inference algorithm for a given flow function
-void llvm::applyFlowInference(FlowFunction &Func) {
+void llvm::applyFlowInference(const ProfiParams &Params, FlowFunction &Func) {
   // Create and apply an inference network model
-  auto InferenceNetwork = MinCostMaxFlow();
-  initializeNetwork(InferenceNetwork, Func);
+  auto InferenceNetwork = MinCostMaxFlow(Params);
+  initializeNetwork(Params, InferenceNetwork, Func);
   InferenceNetwork.run();
 
   // Extract flow values for every block and every edge
   extractWeights(InferenceNetwork, Func);
 
   // Post-processing adjustments to the flow
-  auto Adjuster = FlowAdjuster(Func);
+  auto Adjuster = FlowAdjuster(Params, Func);
   Adjuster.run();
 
 #ifndef NDEBUG
   // Verify the result
   verifyWeights(Func);
 #endif
 }
+
+/// Apply the profile inference algorithm for a given flow function
+void llvm::applyFlowInference(FlowFunction &Func) {
+  ProfiParams Params;
+  // Set the params from the command-line flags.
+  Params.EvenFlowDistribution = SampleProfileEvenFlowDistribution;
+  Params.RebalanceUnknown = SampleProfileRebalanceUnknown;
+  Params.JoinIslands = SampleProfileJoinIslands;
+  Params.CostBlockInc = SampleProfileProfiCostBlockInc;
+  Params.CostBlockDec = SampleProfileProfiCostBlockDec;
+  Params.CostBlockEntryInc = SampleProfileProfiCostBlockEntryInc;
+  Params.CostBlockEntryDec = SampleProfileProfiCostBlockEntryDec;
+  Params.CostBlockZeroInc = SampleProfileProfiCostBlockZeroInc;
+  Params.CostBlockUnknownInc = SampleProfileProfiCostBlockUnknownInc;
+
+  applyFlowInference(Params, Func);
+}

llvm/lib/Transforms/Utils/SampleProfileLoaderBaseUtil.cpp

cl::opt<bool> NoWarnSampleUnused(
     "no-warn-sample-unused", cl::init(false), cl::Hidden,
     cl::desc("Use this option to turn off/on warnings about function with "
              "samples but without debug information to use those samples. "));
 
 cl::opt<bool> SampleProfileUseProfi(
     "sample-profile-use-profi", cl::Hidden,
     cl::desc("Use profi to infer block and edge counts."));
 
-cl::opt<bool> SampleProfileInferEntryCount(
-    "sample-profile-infer-entry-count", cl::init(true), cl::Hidden,
-    cl::desc("Use profi to infer function entry count."));
-
 namespace sampleprofutil {
 
 /// Return true if the given callsite is hot wrt to hot cutoff threshold.
 ///
 /// Functions that were inlined in the original binary will be represented
 /// in the inline stack in the sample profile. If the profile shows that
 /// the original inline decision was "good" (i.e., the callsite is executed
 /// frequently), then we will recreate the inline decision and apply the

llvm/test/Transforms/SampleProfile/csspgo-import-list.ll

 ; Make sure Import GUID list for ThinLTO properly set for CSSPGO
-; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list.prof -sample-profile-even-count-distribution=0 -S | FileCheck %s
+; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list.prof --sample-profile-even-flow-distribution=0 -S | FileCheck %s
 ; RUN: llvm-profdata merge --sample --extbinary %S/Inputs/csspgo-import-list.prof -o %t.prof
-; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%t.prof -sample-profile-even-count-distribution=0 -S | FileCheck %s
+; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%t.prof --sample-profile-even-flow-distribution=0 -S | FileCheck %s
 ; RUN: llvm-profdata show --sample -show-sec-info-only %t.prof | FileCheck %s --check-prefix=CHECK-ORDERED
 ; RUN: llvm-profdata merge --sample --extbinary --use-md5 %S/Inputs/csspgo-import-list.prof -o %t.md5
-; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%t.md5 -sample-profile-even-count-distribution=0 -S | FileCheck %s
+; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%t.md5 --sample-profile-even-flow-distribution=0 -S | FileCheck %s
 ; RUN: llvm-profdata show --sample -show-sec-info-only %t.md5 | FileCheck %s --check-prefix=CHECK-ORDERED
 
 ;; Validate that with replay in effect, we import call sites even if they are below the threshold
 ;; Baseline import decisions
-; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list.prof -profile-summary-hot-count=10000 -sample-profile-even-count-distribution=0 -S | FileCheck %s --check-prefix=THRESHOLD
+; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list.prof -profile-summary-hot-count=10000 --sample-profile-even-flow-distribution=0 -S | FileCheck %s --check-prefix=THRESHOLD
 ;; With replay
-; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list.prof -sample-profile-inline-replay=%S/Inputs/csspgo-import-list-replay.txt -sample-profile-inline-replay-scope=Module -profile-summary-hot-count=10000 -sample-profile-even-count-distribution=0 -S | FileCheck %s --check-prefix=THRESHOLD-REPLAY
+; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list.prof -sample-profile-inline-replay=%S/Inputs/csspgo-import-list-replay.txt -sample-profile-inline-replay-scope=Module -profile-summary-hot-count=10000 --sample-profile-even-flow-distribution=0 -S | FileCheck %s --check-prefix=THRESHOLD-REPLAY
 ;; With replay but no profile information for call to _Z5funcAi. We import _Z5funcAi because it's explicitly in the replay but don't go further to its callee (_Z3fibi) because we lack samples
-; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list-no-funca.prof -sample-profile-inline-replay=%S/Inputs/csspgo-import-list-replay.txt -sample-profile-inline-replay-scope=Module -profile-summary-hot-count=10000 -sample-profile-even-count-distribution=0 -S | FileCheck %s --check-prefix=THRESHOLD-REPLAY-NO-FUNCA
+; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list-no-funca.prof -sample-profile-inline-replay=%S/Inputs/csspgo-import-list-replay.txt -sample-profile-inline-replay-scope=Module -profile-summary-hot-count=10000 --sample-profile-even-flow-distribution=0 -S | FileCheck %s --check-prefix=THRESHOLD-REPLAY-NO-FUNCA
 
 declare i32 @_Z5funcBi(i32 %x)
 declare i32 @_Z5funcAi(i32 %x)
 
 define dso_local i32 @main() local_unnamed_addr #0 !dbg !18 {
 entry:
   br label %for.body, !dbg !25
 

llvm/test/Transforms/SampleProfile/profile-context-tracker.ll

 ; Test for CSSPGO's SampleContextTracker to make sure context profile tree is promoted and merged properly
 ; based on inline decision, so post inline counts are accurate.
 
 ; RUN: llvm-profdata merge --sample --extbinary %S/Inputs/profile-context-tracker.prof -o %t
 
 ; Note that we need new pass manager to enable top-down processing for sample profile loader
 ; Test we inlined the following in top-down order and entry counts accurate reflects post-inline base profile
 ;   main:3 @ _Z5funcAi
 ;   main:3 @ _Z5funcAi:1 @ _Z8funcLeafi
 ;   _Z5funcBi:1 @ _Z8funcLeafi
-; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/profile-context-tracker.prof -sample-profile-inline-size -sample-profile-prioritized-inline=0 -profile-sample-accurate -sample-profile-even-count-distribution=0 -S | FileCheck %s --check-prefix=INLINE-ALL
-; RUN: opt < %s -passes=sample-profile -sample-profile-file=%t -sample-profile-inline-size -sample-profile-prioritized-inline=0 -profile-sample-accurate -sample-profile-even-count-distribution=0 -S | FileCheck %s --check-prefix=INLINE-ALL
-; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/profile-context-tracker.prof -sample-profile-inline-size -sample-profile-cold-inline-threshold=200 -profile-sample-accurate -sample-profile-even-count-distribution=0 -S | FileCheck %s --check-prefix=INLINE-ALL
-; RUN: opt < %s -passes=sample-profile -sample-profile-file=%t -sample-profile-inline-size -sample-profile-cold-inline-threshold=200 -profile-sample-accurate -sample-profile-even-count-distribution=0 -S | FileCheck %s --check-prefix=INLINE-ALL
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/profile-context-tracker.prof -sample-profile-inline-size -sample-profile-prioritized-inline=0 -profile-sample-accurate -sample-profile-even-flow-distribution=0 -S | FileCheck %s --check-prefix=INLINE-ALL
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%t -sample-profile-inline-size -sample-profile-prioritized-inline=0 -profile-sample-accurate -sample-profile-even-flow-distribution=0 -S | FileCheck %s --check-prefix=INLINE-ALL
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/profile-context-tracker.prof -sample-profile-inline-size -sample-profile-cold-inline-threshold=200 -profile-sample-accurate -sample-profile-even-flow-distribution=0 -S | FileCheck %s --check-prefix=INLINE-ALL
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%t -sample-profile-inline-size -sample-profile-cold-inline-threshold=200 -profile-sample-accurate -sample-profile-even-flow-distribution=0 -S | FileCheck %s --check-prefix=INLINE-ALL
 ;
 ; Test we inlined the following in top-down order and entry counts accurate reflects post-inline base profile
 ;   _Z5funcAi:1 @ _Z8funcLeafi
 ;   _Z5funcBi:1 @ _Z8funcLeafi
 
 ; Test the functions won't be inlined as a result of sampled profile if `disable-sample-loader-inlining` is true.
 ;
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/profile-context-tracker.prof -sample-profile-inline-size -sample-profile-prioritized-inline=0 -profile-sample-accurate -disable-sample-loader-inlining -S | FileCheck %s --check-prefix=INLINE-NONE