JIT: prepare for instrumentation before incorporating profile counts

src/coreclr/jit/block.h

@@ -1089,15 +1089,13 @@ struct BasicBlock : private LIR::Range
     BlockSet bbReach; // Set of all blocks that can reach this one
 
     union {
-        BasicBlock* bbIDom;   // Represent the closest dominator to this block (called the Immediate
-                              // Dominator) used to compute the dominance tree.
-        FlowEdge* bbLastPred; // Used early on by fgLinkBasicBlock/fgAddRefPred
+        BasicBlock* bbIDom;          // Represent the closest dominator to this block (called the Immediate
+                                     // Dominator) used to compute the dominance tree.
+        FlowEdge* bbLastPred;        // Used early on by fgLinkBasicBlock/fgAddRefPred
+        void*     bbSparseProbeList; // Used early on by fgInstrument
     };
 
-    union {
-        void* bbSparseCountInfo; // Used early on by fgIncorporateEdgeCounts
-        void* bbSparseProbeList; // Used early on by fgInstrument
-    };
+    void* bbSparseCountInfo; // Used early on by fgIncorporateEdgeCounts
 
     unsigned bbPreorderNum;  // the block's  preorder number in the graph (1...fgMaxBBNum]
     unsigned bbPostorderNum; // the block's postorder number in the graph (1...fgMaxBBNum]

src/coreclr/jit/compiler.cpp

@@ -4496,6 +4496,14 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     };
     DoPhase(this, PHASE_PRE_IMPORT, preImportPhase);
 
+    // If we're going to instrument code, we may need to prepare before
+    // we import. Also do this before we read in any profile data.
+    //
+    if (compileFlags->IsSet(JitFlags::JIT_FLAG_BBINSTR))
+    {
+        DoPhase(this, PHASE_IBCPREP, &Compiler::fgPrepareToInstrumentMethod);
+    }
+
     // Incorporate profile data.
     //
     // Note: the importer is sensitive to block weights, so this has
@@ -4505,14 +4513,6 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     DoPhase(this, PHASE_INCPROFILE, &Compiler::fgIncorporateProfileData);
     activePhaseChecks &= ~PhaseChecks::CHECK_PROFILE;
 
-    // If we're going to instrument code, we may need to prepare before
-    // we import.
-    //
-    if (compileFlags->IsSet(JitFlags::JIT_FLAG_BBINSTR))
-    {
-        DoPhase(this, PHASE_IBCPREP, &Compiler::fgPrepareToInstrumentMethod);
-    }
-
     // If we are doing OSR, update flow to initially reach the appropriate IL offset.
     //
     if (opts.IsOSR())

src/coreclr/jit/fgprofile.cpp

@@ -1220,6 +1220,16 @@ void Compiler::WalkSpanningTree(SpanningTreeVisitor* visitor)
             break;
         }
     }
+
+    // Notify visitor of remaining blocks
+    //
+    for (BasicBlock* const block : Blocks())
+    {
+        if (!BlockSetOps::IsMember(this, marked, block->bbNum))
+        {
+            visitor->VisitBlock(block);
+        }
+    }
 }
 
 // Map a block into its schema key we will use for storing basic blocks.