[BOLT] using jump weights in profi

[spupyrev]

Summary

We want to use profile inference (profi) in BOLT for stale profile matching.
This is the second change for existing usages of profi (e.g., CSSPGO):

(i) Added the ability to provide (estimated) jump weights for the algorithm. The
goal of the algorithm is to create a valid control flow for a given function
(that is, one in which incoming counts equal outgoing counts for every basic
block while minimally modifying the original input block and jump weights). The
input jump weights will be provided based on collected LBR profiles in BOLT.

(ii) Added the corresponding options to ProfiParams.

(iii) Slightly modified / simplified the construction of the flow network in profi
so as it utilizes fewer auxiliary nodes. This is done by introducing parallel
edges to the network (which is supported by MMF) and reduces the size of the
network from 3*|V| to 2*|V|, where |V| is the number of basic blocks in the
function.

Inference (profile quality) impact:
The diff is supposed to be a no-op for the inferred counts. However, our
implementation of MCF is not fully deterministic and might return different
results depending on the input network model. Since we changed the model
construction, there are a few differences in comparison to the original
implementation. I checked manually on an internal benchmark and see a minor
difference (+/- 1 count for certain basic blocks) in just a dozen of instances
(out of 10000+ input functions). Hence, the diff is highly unlikely to have an
impact for existing prod workloads.

Runtime impact:
I measure up to 10% speedup for block-only (ie CSSPGO/AutoFDO) inference and up
to 50% speedup for block+jump inference (ie BOLT) in comparison to the original
unoptimized version.

[spupyrev]

upds

[spupyrev]

fixing a test

[spupyrev]

format

[spupyrev]

rebase

[hoy]

llvm/lib/Transforms/Utils/SampleProfileInference.cpp

1088

Does this still hold?

1120

Where is this handled in the new implementation?

1125

Previously the cost it takes to go from Bin to Bout is 2*AuxCostInc. Now without Baux the cost for the same path becomes AuxCostInc. Is it a discrepancy or am I missing anything?

1164

Add a comment to indicate the check is for fall-through branches?

1222

Please add a message for the assert and the one right below.

[spupyrev]

comments

[spupyrev]

llvm/lib/Transforms/Utils/SampleProfileInference.cpp

1088

the assert is moved to verifyInput

1120

This is a good question; we don't have this condition anymore. While in theory there might be a difference, I do not see a single instance (in my benchmark with 10K functions) where this statement yields a different result. So it must be quite rare.
More importantly, I do not remember the original motivation: Why is it "correct" not to penalize blocks with self loops?

1125

Now all the costs are reduced by a factor of 2, that is, instead of 2*AuxCostInc we have AuxCostInc and instead of 2*AuxCostDec we have AuxCostDec. This does not have any impact on the algorithm and the produced solution; only the objective is (uniformly) scaled

[hoy]

llvm/lib/Transforms/Utils/SampleProfileInference.cpp

1120

Oh, it might be related to a Skylake-specific hardware issue where a LBR entry can occur consecutively redundantly. Decreasing the backedge count may make sense. Otherwise I cannot think of why self-loop is special from normal loops.

The Skylake issue is worked around in the profile generation time so we should be free from that now.

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 HasUnknownWeight{false};
+  bool HasUnknownWeight{true};
+  bool IsUnlikely{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(); }
 
   /// Check if it is an exit block in the function.
   bool isExit() const { return SuccJumps.empty(); }
 };
 
 /// A wrapper of a jump between two basic blocks.
 struct FlowJump {
   uint64_t Source;
   uint64_t Target;
-  uint64_t Flow{0};
+  uint64_t Weight{0};
+  bool HasUnknownWeight{true};
   bool IsUnlikely{false};
+  uint64_t Flow{0};
 };
 
 /// A wrapper of binary function with basic blocks and jumps.
 struct FlowFunction {
+  /// Basic blocks in the function.
   std::vector<FlowBlock> Blocks;
+  /// Jumps between the basic blocks.
   std::vector<FlowJump> Jumps;
   /// The index of the entry block.
   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 {
   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 increasing a jump's count by one.
+  unsigned CostJumpInc{0};
+
+  /// The cost of increasing a fall-through jump's count by one.
+  unsigned CostJumpFTInc{0};
+
+  /// The cost of decreasing a jump's count by one.
+  unsigned CostJumpDec{0};
+
+  /// The cost of decreasing a fall-through jump's count by one.
+  unsigned CostJumpFTDec{0};
+
+  /// The cost of increasing an unknown jump's count by one.
+  unsigned CostJumpUnknownInc{0};
+
+  /// The cost of increasing an unknown fall-through jump's count by one.
+  unsigned CostJumpUnknownFTInc{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.
public:
   SampleProfileInference(FunctionT &F, BlockEdgeMap &Successors,
                          BlockWeightMap &SampleBlockWeights)
       : F(F), Successors(Successors), SampleBlockWeights(SampleBlockWeights) {}
 
   /// Apply the profile inference algorithm for a given function
   void apply(BlockWeightMap &BlockWeights, EdgeWeightMap &EdgeWeights);
 
 private:
+  /// Initialize flow function blocks, jumps and misc metadata.
+  void initFunction(FlowFunction &Func,
+                    const std::vector<const BasicBlockT *> &BasicBlocks,
+                    DenseMap<const BasicBlockT *, uint64_t> &BlockIndex);
+
   /// Try to infer branch probabilities mimicking implementation of
   /// BranchProbabilityInfo. Unlikely taken branches are marked so that the
   /// inference algorithm can avoid sending flow along corresponding edges.
   void findUnlikelyJumps(const std::vector<const BasicBlockT *> &BasicBlocks,
                          BlockEdgeMap &Successors, FlowFunction &Func);
 
   /// Determine whether the block is an exit in the CFG.
   bool isExit(const BasicBlockT *BB);
void SampleProfileInference<BT>::apply(BlockWeightMap &BlockWeights,
   }
   // Quit early for functions with a single block or ones w/o samples
   if (BasicBlocks.size() <= 1 || !HasSamples) {
     return;
   }
 
   // Create necessary objects
   FlowFunction Func;
+  initFunction(Func, BasicBlocks, BlockIndex);
+
+  // Create and apply the inference network model.
+  applyFlowInference(Func);
+
+  // Extract the resulting weights from the control flow
+  // All weights are increased by one to avoid propagation errors introduced by
+  // zero weights.
+  for (const auto *BB : BasicBlocks) {
+    BlockWeights[BB] = Func.Blocks[BlockIndex[BB]].Flow;
+  }
+  for (auto &Jump : Func.Jumps) {
+    Edge E = std::make_pair(BasicBlocks[Jump.Source], BasicBlocks[Jump.Target]);
+    EdgeWeights[E] = Jump.Flow;
+  }
+
+#ifndef NDEBUG
+  // Unreachable blocks and edges should not have a weight.
+  for (auto &I : BlockWeights) {
+    assert(Reachable.contains(I.first));
+    assert(InverseReachable.contains(I.first));
+  }
+  for (auto &I : EdgeWeights) {
+    assert(Reachable.contains(I.first.first) &&
+           Reachable.contains(I.first.second));
+    assert(InverseReachable.contains(I.first.first) &&
+           InverseReachable.contains(I.first.second));
+  }
+#endif
+}
+
+template <typename BT>
+void SampleProfileInference<BT>::initFunction(
+    FlowFunction &Func, const std::vector<const BasicBlockT *> &BasicBlocks,
+    DenseMap<const BasicBlockT *, uint64_t> &BlockIndex) {
   Func.Blocks.reserve(BasicBlocks.size());
   // Create FlowBlocks
   for (const auto *BB : BasicBlocks) {
     FlowBlock Block;
     if (SampleBlockWeights.find(BB) != SampleBlockWeights.end()) {
       Block.HasUnknownWeight = false;
       Block.Weight = SampleBlockWeights[BB];
     } else {
       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]) {
       if (!BlockIndex.count(Succ))
         continue;
       FlowJump Jump;
       Jump.Source = BlockIndex[BB];
       Jump.Target = BlockIndex[Succ];
       Func.Jumps.push_back(Jump);
-      if (BB == Succ) {
-        Func.Blocks[BlockIndex[BB]].HasSelfEdge = true;
-      }
     }
   }
   for (auto &Jump : Func.Jumps) {
-    Func.Blocks[Jump.Source].SuccJumps.push_back(&Jump);
-    Func.Blocks[Jump.Target].PredJumps.push_back(&Jump);
+    uint64_t Src = Jump.Source;
+    uint64_t Dst = Jump.Target;
+    Func.Blocks[Src].SuccJumps.push_back(&Jump);
+    Func.Blocks[Dst].PredJumps.push_back(&Jump);
   }
 
   // Try to infer probabilities of jumps based on the content of basic block
   findUnlikelyJumps(BasicBlocks, Successors, Func);
 
   // Find the entry block
   for (size_t I = 0; I < Func.Blocks.size(); I++) {
     if (Func.Blocks[I].isEntry()) {
       Func.Entry = I;
       break;
     }
   }
+  assert(Func.Entry == 0 && "incorrect index of the entry block");
 
-  // Create and apply the inference network model.
-  applyFlowInference(Func);
-
-  // Extract the resulting weights from the control flow
-  // All weights are increased by one to avoid propagation errors introduced by
-  // zero weights.
-  for (const auto *BB : BasicBlocks) {
-    BlockWeights[BB] = Func.Blocks[BlockIndex[BB]].Flow;
+  // Pre-process data: make sure the entry weight is at least 1
+  auto &EntryBlock = Func.Blocks[Func.Entry];
+  if (EntryBlock.Weight == 0 && !EntryBlock.HasUnknownWeight) {
+    EntryBlock.Weight = 1;
+    EntryBlock.HasUnknownWeight = false;
   }
-  for (auto &Jump : Func.Jumps) {
-    Edge E = std::make_pair(BasicBlocks[Jump.Source], BasicBlocks[Jump.Target]);
-    EdgeWeights[E] = Jump.Flow;
-  }
-
-#ifndef NDEBUG
-  // Unreachable blocks and edges should not have a weight.
-  for (auto &I : BlockWeights) {
-    assert(Reachable.contains(I.first));
-    assert(InverseReachable.contains(I.first));
-  }
-  for (auto &I : EdgeWeights) {
-    assert(Reachable.contains(I.first.first) &&
-           Reachable.contains(I.first.second));
-    assert(InverseReachable.contains(I.first.first) &&
-           InverseReachable.contains(I.first.second));
-  }
-#endif
 }
 
 template <typename BT>
 inline void SampleProfileInference<BT>::findUnlikelyJumps(
     const std::vector<const BasicBlockT *> &BasicBlocks,
     BlockEdgeMap &Successors, FlowFunction &Func) {}
 
 template <>

llvm/lib/Transforms/Utils/SampleProfileInference.cpp

void initialize(uint64_t NodeCount, uint64_t SourceNode, uint64_t SinkNode) {
     Edges = std::vector<std::vector<Edge>>(NodeCount, std::vector<Edge>());
     if (Params.EvenFlowDistribution)
       AugmentingEdges =
           std::vector<std::vector<Edge *>>(NodeCount, std::vector<Edge *>());
   }
 
   // Run the algorithm.
   int64_t run() {
+    LLVM_DEBUG(dbgs() << "Starting profi for " << Nodes.size() << " nodes\n");
+
     // Iteratively find an augmentation path/dag in the network and send the
     // flow along its edges
     size_t AugmentationIters = applyFlowAugmentation();
 
     // Compute the total flow and its cost
     int64_t TotalCost = 0;
     int64_t TotalFlow = 0;
     for (uint64_t Src = 0; Src < Nodes.size(); Src++) {
private:
   /// 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
+/// A post-processing adjustment of the 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(const ProfiParams &Params, FlowFunction &Func)
-      : Params(Params), Func(Func) {
-    assert(Func.Blocks[Func.Entry].isEntry() &&
-           "incorrect index of the entry block");
-  }
+      : Params(Params), Func(Func) {}
 
-  // Run the post-processing
+  /// Apply the post-processing.
   void run() {
     if (Params.JoinIslands) {
-      /// Adjust the flow to get rid of isolated components.
+      // Adjust the flow to get rid of isolated components
       joinIsolatedComponents();
     }
 
     if (Params.RebalanceUnknown) {
-      /// Rebalance the flow inside unknown subgraphs.
+      // Rebalance the flow inside unknown subgraphs
       rebalanceUnknownSubgraphs();
     }
   }
 
 private:
   void joinIsolatedComponents() {
     // Find blocks that are reachable from the source
     auto Visited = BitVector(NumBlocks(), false);
private:
   ///   - 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(FlowAdjuster::MinBaseDistance,
                  std::min(Func.Blocks[Func.Entry].Flow,
-                          Params.CostUnlikely / NumBlocks()));
+                          Params.CostUnlikely / (2 * (NumBlocks() + 1))));
     if (Jump->Flow > 0)
       return BaseDistance + BaseDistance / Jump->Flow;
-    return BaseDistance * NumBlocks();
+    return 2 * BaseDistance * (NumBlocks() + 1);
   };
 
   uint64_t NumBlocks() const { return Func.Blocks.size(); }
 
   /// Rebalance unknown subgraphs so that the flow is split evenly across the
   /// outgoing branches of every block of the subgraph. The method iterates over
   /// blocks with known weight and identifies unknown subgraphs rooted at the
   /// blocks. Then it verifies if flow rebalancing is feasible and applies it.
private:
   static constexpr uint64_t MinBaseDistance = 10000;
 
   /// Params for flow computation.
   const ProfiParams &Params;
   /// The function.
   FlowFunction &Func;
 };
 
+std::pair<int64_t, int64_t> assignBlockCosts(const ProfiParams &Params,
+                                             const FlowBlock &Block);
+std::pair<int64_t, int64_t> assignJumpCosts(const ProfiParams &Params,
+                                            const FlowJump &Jump);
+
 /// 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
+/// Every block is split into two nodes that are responsible for (i) an
+/// incoming flow, (ii) an outgoing flow; they penalize an increase or a
 /// reduction of the block weight.
 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");
+  uint64_t NumJumps = Func.Jumps.size();
+  assert(NumJumps > 0 && "Too few jumps in a function");
 
-  // 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;
+  // The nodes corresponding to blocks of the function have indicies in
+  // the range [0 .. 2 * NumBlocks); the dummy sources/sinks are indexed by the
+  // next four values.
+  uint64_t S = 2 * NumBlocks;
   uint64_t T = S + 1;
   uint64_t S1 = S + 2;
   uint64_t T1 = S + 3;
 
-  Network.initialize(3 * NumBlocks + 4, S1, T1);
+  Network.initialize(2 * NumBlocks + 4, S1, T1);
 
-  // Create three nodes for every block of the function
+  // Initialize nodes of the flow network
   for (uint64_t B = 0; B < NumBlocks; B++) {
     auto &Block = Func.Blocks[B];
-    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);
-    }
+    // Split every block into two auxiliary nodes to allow
+    // increase/reduction of the block count.
+    uint64_t Bin = 2 * B;
+    uint64_t Bout = 2 * B + 1;
 
     // 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 = Params.CostBlockInc;
-    int64_t AuxCostDec = Params.CostBlockDec;
-    if (Block.HasUnknownWeight) {
-      // Do not penalize changing weights of blocks w/o known profile count
-      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 = Params.CostBlockZeroInc;
-      }
-      // Modifying the count of the entry block is expensive
-      if (Block.isEntry()) {
-        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;
-    }
+    // Assign costs for increasing/decreasing the block counts
+    auto [AuxCostInc, AuxCostDec] = assignBlockCosts(Params, Block);
+
+    // Add the corresponding edges to the network
+    Network.addEdge(Bin, Bout, AuxCostInc);
+    if (Block.Weight > 0) {
+      Network.addEdge(Bout, Bin, Block.Weight, AuxCostDec);
+      Network.addEdge(S1, Bout, Block.Weight, 0);
+      Network.addEdge(Bin, T1, Block.Weight, 0);
+    }
+  }
+
+  // Initialize edges of the flow network
+  for (uint64_t J = 0; J < NumJumps; J++) {
+    auto &Jump = Func.Jumps[J];
+
+    // Get the endpoints corresponding to the jump
+    uint64_t Jin = 2 * Jump.Source + 1;
+    uint64_t Jout = 2 * Jump.Target;
+
+    // Assign costs for increasing/decreasing the jump counts
+    auto [AuxCostInc, AuxCostDec] = assignJumpCosts(Params, Jump);
+
+    // Add the corresponding edges to the network
+    Network.addEdge(Jin, Jout, AuxCostInc);
+    if (Jump.Weight > 0) {
+      Network.addEdge(Jout, Jin, Jump.Weight, AuxCostDec);
+      Network.addEdge(S1, Jout, Jump.Weight, 0);
+      Network.addEdge(Jin, T1, Jump.Weight, 0);
+    }
+  }
+
+  // Make sure we have a valid flow circulation
+  Network.addEdge(T, S, 0);
+}
+
+/// Assign costs for increasing/decreasing the block counts.
+std::pair<int64_t, int64_t> assignBlockCosts(const ProfiParams &Params,
+                                             const FlowBlock &Block) {
+  // Modifying the weight of an unlikely block is expensive
+  if (Block.IsUnlikely)
+    return std::make_pair(Params.CostUnlikely, Params.CostUnlikely);
+
+  // Assign default values for the costs
+  int64_t CostInc = Params.CostBlockInc;
+  int64_t CostDec = Params.CostBlockDec;
+  // Update the costs depending on the block metadata
+    CostInc = Params.CostBlockUnknownInc;
+    CostDec = 0;
+    if (Block.Weight == 0)
+      CostInc = Params.CostBlockZeroInc;
+      CostInc = Params.CostBlockEntryInc;
+      CostDec = Params.CostBlockEntryDec;
+  return std::make_pair(CostInc, CostDec);
 
-    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 ? Params.CostUnlikely : 0;
-      Network.addEdge(SrcOut, DstIn, Cost);
-    }
-  }
-
-  // Make sure we have a valid flow circulation
-  Network.addEdge(T, S, 0);
+/// Assign costs for increasing/decreasing the jump counts.
+std::pair<int64_t, int64_t> assignJumpCosts(const ProfiParams &Params,
+                                            const FlowJump &Jump) {
+  // Modifying the weight of an unlikely jump is expensive
+  if (Jump.IsUnlikely)
+    return std::make_pair(Params.CostUnlikely, Params.CostUnlikely);
+
+  // Assign default values for the costs
+  int64_t CostInc = Params.CostJumpInc;
+  int64_t CostDec = Params.CostJumpDec;
+  // Update the costs depending on the block metadata
+  if (Jump.Source + 1 == Jump.Target) {
+    CostInc = Params.CostJumpFTInc;
+    CostDec = Params.CostJumpFTDec;
+  }
+  if (Jump.HasUnknownWeight) {
+    if (Jump.Source + 1 == Jump.Target)
+      CostInc = Params.CostJumpUnknownFTInc;
+    else
+      CostInc = Params.CostJumpUnknownInc;
+    CostDec = 0;
+  } else {
+    assert(Jump.Weight > 0 && "found zero-weight jump with a positive weight");
+  return std::make_pair(CostInc, CostDec);
 }
 
 /// Extract resulting block and edge counts from the flow network.
-void extractWeights(MinCostMaxFlow &Network, FlowFunction &Func) {
+void extractWeights(const ProfiParams &Params, MinCostMaxFlow &Network,
+                    FlowFunction &Func) {
   uint64_t NumBlocks = Func.Blocks.size();
+  uint64_t NumJumps = Func.Jumps.size();
+
+  // Extract resulting jump counts
+  for (uint64_t J = 0; J < NumJumps; J++) {
+    auto &Jump = Func.Jumps[J];
+    uint64_t SrcOut = 2 * Jump.Source + 1;
+    uint64_t DstIn = 2 * Jump.Target;
 
-  // Extract resulting block counts
-  for (uint64_t Src = 0; Src < NumBlocks; Src++) {
-    auto &Block = Func.Blocks[Src];
-    uint64_t SrcOut = 3 * Src + 1;
     int64_t Flow = 0;
-    for (const auto &Adj : Network.getFlow(SrcOut)) {
-      uint64_t DstIn = Adj.first;
-      int64_t DstFlow = Adj.second;
-      bool IsAuxNode = (DstIn < 3 * NumBlocks && DstIn % 3 == 2);
-      if (!IsAuxNode || Block.HasSelfEdge) {
-        Flow += DstFlow;
-      }
-    }
-    Block.Flow = Flow;
-    assert(Flow >= 0 && "negative block flow");
+    int64_t AuxFlow = Network.getFlow(SrcOut, DstIn);
+    if (Jump.Source != Jump.Target)
+      Flow = int64_t(Jump.Weight) + AuxFlow;
+    else
+      Flow = int64_t(Jump.Weight) + (AuxFlow > 0 ? AuxFlow : 0);
+
+    Jump.Flow = Flow;
+    assert(Flow >= 0 && "negative jump flow");
   }
 
-  // Extract resulting jump counts
+  // Extract resulting block counts
+  auto InFlow = std::vector<uint64_t>(NumBlocks, 0);
+  auto OutFlow = std::vector<uint64_t>(NumBlocks, 0);
   for (auto &Jump : Func.Jumps) {
-    uint64_t Src = Jump.Source;
-    uint64_t Dst = Jump.Target;
-    int64_t Flow = 0;
-    if (Src != Dst) {
-      uint64_t SrcOut = 3 * Src + 1;
-      uint64_t DstIn = 3 * Dst;
-      Flow = Network.getFlow(SrcOut, DstIn);
-    } else {
-      uint64_t SrcOut = 3 * Src + 1;
-      uint64_t SrcAux = 3 * Src + 2;
-      int64_t AuxFlow = Network.getFlow(SrcOut, SrcAux);
-      if (AuxFlow > 0)
-        Flow = AuxFlow;
-    }
-    Jump.Flow = Flow;
-    assert(Flow >= 0 && "negative jump flow");
+    InFlow[Jump.Target] += Jump.Flow;
+    OutFlow[Jump.Source] += Jump.Flow;
+  for (uint64_t B = 0; B < NumBlocks; B++) {
+    auto &Block = Func.Blocks[B];
+    Block.Flow = std::max(OutFlow[B], InFlow[B]);
   }
 }
 
 #ifndef NDEBUG
-/// Verify that the computed flow values satisfy flow conservation rules
-void verifyWeights(const FlowFunction &Func) {
+/// Verify that the provided block/jump weights are as expected.
+void verifyInput(const FlowFunction &Func) {
+  // Verify the entry block
+  assert(Func.Entry == 0 && Func.Blocks[0].isEntry());
+  for (size_t I = 1; I < Func.Blocks.size(); I++) {
+    assert(!Func.Blocks[I].isEntry());
+  }
+  // Verify CFG jumps
+  for (auto &Block : Func.Blocks) {
+    assert(!Block.PredJumps.empty() || !Block.SuccJumps.empty());
+  }
+  // Verify input block weights
+  for (auto &Block : Func.Blocks) {
+    assert((!Block.HasUnknownWeight || Block.Weight == 0 || Block.isEntry()) &&
+           "non-zero weight of a block w/o weight except for an entry");
+  }
+  // Verify input jump weights
+  for (auto &Jump : Func.Jumps) {
+    assert((!Jump.HasUnknownWeight || Jump.Weight == 0) &&
+           "non-zero weight of a jump w/o weight");
+  }
+}
+
+/// Verify that the computed flow values satisfy flow conservation rules.
+void verifyOutput(const FlowFunction &Func) {
   const uint64_t NumBlocks = Func.Blocks.size();
   auto InFlow = std::vector<uint64_t>(NumBlocks, 0);
   auto OutFlow = std::vector<uint64_t>(NumBlocks, 0);
   for (const auto &Jump : Func.Jumps) {
     InFlow[Jump.Target] += Jump.Flow;
     OutFlow[Jump.Source] += Jump.Flow;
   }
 
for (uint64_t I = 0; I < NumBlocks; I++) {
     auto &Block = Func.Blocks[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
+/// Apply the profile inference algorithm for a given function
 void llvm::applyFlowInference(const ProfiParams &Params, FlowFunction &Func) {
+  // TODO: uncomment!
+  //#ifndef NDEBUG
+  // Verify the input data
+  verifyInput(Func);
+  //#endif
+
   // Create and apply an inference network model
   auto InferenceNetwork = MinCostMaxFlow(Params);
   initializeNetwork(Params, InferenceNetwork, Func);
   InferenceNetwork.run();
 
   // Extract flow values for every block and every edge
-  extractWeights(InferenceNetwork, Func);
+  extractWeights(Params, InferenceNetwork, Func);
 
   // Post-processing adjustments to the flow
   auto Adjuster = FlowAdjuster(Params, Func);
   Adjuster.run();
 
-#ifndef NDEBUG
+  //#ifndef NDEBUG
   // Verify the result
-  verifyWeights(Func);
-#endif
+  verifyOutput(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;

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

 ; INLINE-ALL-DAG-SAME: [[LEAF_PROF]] = !{!"function_entry_count", i64 0}
 ; INLINE-ALL-DAG: [[FUNCB_PROF]] = !{!"function_entry_count", i64 13}
 
 ; INLINE-HOT-DAG: [[MAIN_PROF]] = !{!"function_entry_count", i64 1}
 ; INLINE-HOT-DAG: [[FUNCA_PROF]] = !{!"function_entry_count", i64 12}
 ; INLINE-HOT-DAG-SAME: [[LEAF_PROF]] = !{!"function_entry_count", i64 0}
 ; INLINE-HOT-DAG: [[FUNCB_PROF]] = !{!"function_entry_count", i64 13}
 
-; INLINE-NONE: [[MAIN_PROF]] = !{!"function_entry_count", i64 13}
+; INLINE-NONE: [[MAIN_PROF]] = !{!"function_entry_count", i64 14}
 ; INLINE-NONE: [[FUNCA_PROF]] = !{!"function_entry_count", i64 24}
 ; INLINE-NONE-DAG-SAME: [[LEAF_PROF]] = !{!"function_entry_count", i64 21}
 ; INLINE-NONE-DAG: [[FUNCB_PROF]] = !{!"function_entry_count", i64 32}
 
 declare i32 @_Z3fibi(i32)
 
 attributes #0 = { nofree noinline norecurse nounwind uwtable "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" "use-sample-profile" }
 attributes #1 = { nofree norecurse nounwind uwtable "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" "use-sample-profile" }