[profi][NFC] Refactor SampleProfileInference::apply

[Amir]

Summary

Split up apply into the following steps:

• findUnreachable
• createStableOrder
• preAssignWeights
• canApplyInference
• extractProfile
• verifyProfileConsistency

With some specializations, apply should apply to BOLT stale profile matching
(D144500), or potentially a custom apply reusing some of the interfaces.

[spupyrev]

Do you have a follow-up diff to adjust BOLT's part?

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

261–264

"stable order" confuses me. How about "fixed order" or simply "an order of blocks"?

266–268

I'd separate the computation of HasSamples and preAssignWeights, as they seem to be unrelated:

277–278

this comment doesn't make sense (perhaps, outdated). remove maybe

302

should this marker be moved to the function call?

[Amir]

In D152217#4415501, @spupyrev wrote:

Do you have a follow-up diff to adjust BOLT's part?

Not at the moment, it's WIP.

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

public:
   using NodeRef = typename GraphTraits<BT *>::NodeRef;
   using BasicBlockT = typename std::remove_pointer<NodeRef>::type;
   using FunctionT = BT;
   using Edge = std::pair<const BasicBlockT *, const BasicBlockT *>;
   using BlockWeightMap = DenseMap<const BasicBlockT *, uint64_t>;
   using EdgeWeightMap = DenseMap<Edge, uint64_t>;
   using BlockEdgeMap =
       DenseMap<const BasicBlockT *, SmallVector<const BasicBlockT *, 8>>;
+  using BlockIndexMap = DenseMap<const BasicBlockT *, uint64_t>;
+  using BlockList = std::vector<const BasicBlockT *>;
+  using BlockVisitedSet = df_iterator_default_set<const BasicBlockT *>;
 
   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.
   FlowFunction
   createFlowFunction(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);
 
+  /// Collect all forwards and backwards reachable blocks which the inference
+  /// algorithm will be applied on.
+  void findReachable(BlockVisitedSet &Reachable,
+                     BlockVisitedSet &InverseReachable);
+
+  /// Produce a stable order for all reachable blocks.
+  void createStableOrder(BlockIndexMap &BlockIndex, BlockList &BasicBlocks,
+                         const BlockVisitedSet &Reachable,
+                         const BlockVisitedSet &InverseReachable);
+
+  /// Pre-assign weights from SampleBlockWeights. Returns whether the function
+  /// has any samples.
+  bool preAssignWeights(BlockWeightMap &BlockWeights,
+                        EdgeWeightMap &EdgeWeights,
+                        const BlockList &BasicBlocks);
+
+  /// Pre-flight check whether the inference could/should be applied.
+  bool canApplyInference(const BlockList &BasicBlocks, const bool HasSamples);
+
+  /// Extract profile from FlowFunction to \p BlockWeights and \p EdgeWeights.
+  void extractProfile(BlockWeightMap &BlockWeights, EdgeWeightMap &EdgeWeights,
+                      const FlowFunction &Func, const BlockList &BasicBlocks,
+                      const BlockIndexMap &BlockIndex);
+
+  void verifyProfileConsistency(const BlockWeightMap &BlockWeights,
+                                const EdgeWeightMap &EdgeWeights,
+                                const BlockVisitedSet &Reachable,
+                                const BlockVisitedSet &InverseReachable);
+
   /// Determine whether the block is an exit in the CFG.
   bool isExit(const BasicBlockT *BB);
 
   /// Function.
   const FunctionT &F;
 
   /// Successors for each basic block in the CFG.
   BlockEdgeMap &Successors;
 
   /// Map basic blocks to their sampled weights.
   BlockWeightMap &SampleBlockWeights;
 };
 
 template <typename BT>
-void SampleProfileInference<BT>::apply(BlockWeightMap &BlockWeights,
-                                       EdgeWeightMap &EdgeWeights) {
+void SampleProfileInference<BT>::findReachable(
+    BlockVisitedSet &Reachable, BlockVisitedSet &InverseReachable) {
   // Find all forwards reachable blocks which the inference algorithm will be
   // applied on.
-  df_iterator_default_set<const BasicBlockT *> Reachable;
   for (auto *BB : depth_first_ext(&F, Reachable))
     (void)BB /* Mark all reachable blocks */;
 
   // Find all backwards reachable blocks which the inference algorithm will be
   // applied on.
-  df_iterator_default_set<const BasicBlockT *> InverseReachable;
   for (const auto &BB : F) {
     // An exit block is a block without any successors.
     if (isExit(&BB)) {
       for (auto *RBB : inverse_depth_first_ext(&BB, InverseReachable))
         (void)RBB;
     }
   }
+}
 
-  // Keep a stable order for reachable blocks
-  DenseMap<const BasicBlockT *, uint64_t> BlockIndex;
-  std::vector<const BasicBlockT *> BasicBlocks;
+template <typename BT>
+void SampleProfileInference<BT>::createStableOrder(
+    BlockIndexMap &BlockIndex, BlockList &BasicBlocks,
+    const BlockVisitedSet &Reachable, const BlockVisitedSet &InverseReachable) {
   BlockIndex.reserve(Reachable.size());
   BasicBlocks.reserve(Reachable.size());
   for (const auto &BB : F) {
     if (Reachable.count(&BB) && InverseReachable.count(&BB)) {
       BlockIndex[&BB] = BasicBlocks.size();
       BasicBlocks.push_back(&BB);
     }
   }
+}
 
+template <typename BT>
+bool SampleProfileInference<BT>::preAssignWeights(
+    BlockWeightMap &BlockWeights, EdgeWeightMap &EdgeWeights,
+    const BlockList &BasicBlocks) {
   BlockWeights.clear();
   EdgeWeights.clear();
   bool HasSamples = false;
   for (const auto *BB : BasicBlocks) {
     auto It = SampleBlockWeights.find(BB);
     if (It != SampleBlockWeights.end() && It->second > 0) {
       HasSamples = true;
       BlockWeights[BB] = It->second;
     }
   }
+  return HasSamples;
+}
+
+template <typename BT>
+bool SampleProfileInference<BT>::canApplyInference(const BlockList &BasicBlocks,
+                                                   const bool HasSamples) {
   // Quit early for functions with a single block or ones w/o samples
   if (BasicBlocks.size() <= 1 || !HasSamples) {
-    return;
-  }
+    return false;
+  }
+  return true;
 
+template <typename BT>
+void SampleProfileInference<BT>::apply(BlockWeightMap &BlockWeights,
+                                       EdgeWeightMap &EdgeWeights) {
+  BlockVisitedSet Reachable;
+  BlockVisitedSet InverseReachable;
+  findReachable(Reachable, InverseReachable);
+
+  // Keep a stable order for reachable blocks
+  BlockIndexMap BlockIndex;
+  BlockList BasicBlocks;
+  createStableOrder(BlockIndex, BasicBlocks, Reachable, InverseReachable);
+
+  bool HasSamples = preAssignWeights(BlockWeights, EdgeWeights, BasicBlocks);
+  if (!canApplyInference(BasicBlocks, HasSamples))
+    return;
+
   // Create necessary objects
   FlowFunction Func = createFlowFunction(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;
-  }
+  extractProfile(BlockWeights, EdgeWeights, Func, BasicBlocks, BlockIndex);
+
+  verifyProfileConsistency(BlockWeights, EdgeWeights, Reachable,
+                           InverseReachable);
+
+template <typename BT>
+void SampleProfileInference<BT>::extractProfile(
+    BlockWeightMap &BlockWeights, EdgeWeightMap &EdgeWeights,
+    const FlowFunction &Func, const BlockList &BasicBlocks,
+    const BlockIndexMap &BlockIndex) {
+  for (const auto &[BB, Idx] : BlockIndex)
+    BlockWeights[BB] = Func.Blocks[Idx].Flow;
   for (auto &Jump : Func.Jumps) {
     Edge E = std::make_pair(BasicBlocks[Jump.Source], BasicBlocks[Jump.Target]);
     EdgeWeights[E] = Jump.Flow;
   }
+}
 
+template <typename BT>
+void SampleProfileInference<BT>::verifyProfileConsistency(
+    const BlockWeightMap &BlockWeights, const EdgeWeightMap &EdgeWeights,
+    const BlockVisitedSet &Reachable, const BlockVisitedSet &InverseReachable) {
 #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));