1 files changed, 13 insertions, 6 deletions
diff --git a/llvm/lib/CodeGen/MIRSampleProfile.cpp b/llvm/lib/CodeGen/MIRSampleProfile.cpp
index ce82f280c1c5..23db09b89599 100644
--- a/llvm/lib/CodeGen/MIRSampleProfile.cpp
+++ b/llvm/lib/CodeGen/MIRSampleProfile.cpp
@@ -126,8 +126,10 @@ template <> struct IRTraits<MachineBasicBlock> {
   using PostDominatorTreeT = MachinePostDominatorTree;
   using OptRemarkEmitterT = MachineOptimizationRemarkEmitter;
   using OptRemarkAnalysisT = MachineOptimizationRemarkAnalysis;
-  using PredRangeT = iterator_range<std::vector<MachineBasicBlock *>::iterator>;
-  using SuccRangeT = iterator_range<std::vector<MachineBasicBlock *>::iterator>;
+  using PredRangeT =
+      iterator_range<SmallVectorImpl<MachineBasicBlock *>::iterator>;
+  using SuccRangeT =
+      iterator_range<SmallVectorImpl<MachineBasicBlock *>::iterator>;
   static Function &getFunction(MachineFunction &F) { return F.getFunction(); }
   static const MachineBasicBlock *getEntryBB(const MachineFunction *F) {
     return GraphTraits<const MachineFunction *>::getEntryNode(F);
@@ -364,13 +366,18 @@ bool MIRProfileLoaderPass::runOnMachineFunction(MachineFunction &MF) {
   LLVM_DEBUG(dbgs() << "MIRProfileLoader pass working on Func: "
                     << MF.getFunction().getName() << "\n");
   MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+  auto *MDT = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
+  auto *MPDT =
+      &getAnalysis<MachinePostDominatorTreeWrapperPass>().getPostDomTree();
+
+  MF.RenumberBlocks();
+  MDT->updateBlockNumbers();
+  MPDT->updateBlockNumbers();
+
   MIRSampleLoader->setInitVals(
-      &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree(),
-      &getAnalysis<MachinePostDominatorTreeWrapperPass>().getPostDomTree(),
-      &getAnalysis<MachineLoopInfoWrapperPass>().getLI(), MBFI,
+      MDT, MPDT, &getAnalysis<MachineLoopInfoWrapperPass>().getLI(), MBFI,
       &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE());
 
-  MF.RenumberBlocks();
   if (ViewBFIBefore && ViewBlockLayoutWithBFI != GVDT_None &&
       (ViewBlockFreqFuncName.empty() ||
        MF.getFunction().getName() == ViewBlockFreqFuncName)) {