diff --git a/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h b/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
--- a/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
+++ b/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
@@ -126,6 +126,9 @@
   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)
@@ -146,6 +149,35 @@
   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);
 
@@ -160,17 +192,15 @@
 };
 
 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)) {
@@ -178,10 +208,12 @@
         (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) {
@@ -190,7 +222,12 @@
       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;
@@ -201,10 +238,34 @@
       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);
@@ -215,14 +276,29 @@
   // 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) {