JIT: enable instrumentation for inlinees

eng/pipelines/common/templates/runtimes/run-test-job.yml

@@ -474,6 +474,7 @@ jobs:
             - fullpgo_methodprofiling_always_optimized
             - syntheticpgo
             - syntheticpgo_blend
+            - instrument_if_optimizing
         ${{ if in(parameters.testGroup, 'gc-longrunning') }}:
           longRunningGcTests: true
           scenarios:

eng/pipelines/coreclr/libraries-pgo.yml

@@ -71,3 +71,4 @@ extends:
                     - syntheticpgo
                     - syntheticpgo_blend
                     - jitrlcse
+                    - instrument_if_optimizing

src/coreclr/inc/pgo_formatprocessing.h

@@ -364,10 +364,9 @@ bool ReadInstrumentationSchemaWithLayout(const uint8_t *pByte, size_t cbDataMax,
 }
 
 
-// Return true if schemaTable entries are a subset of the schema described by pByte, with matching entries in the same order.
-// Also updates offset of the matching entries in schemaTable to those of the pByte schema.
+// Return number of initial matching schema entries from the pByteSchema and schemaTable.
 //
-inline bool CheckIfPgoSchemaIsCompatibleAndSetOffsets(const uint8_t *pByte, size_t cbDataMax, ICorJitInfo::PgoInstrumentationSchema* schemaTable, size_t cSchemas)
+inline size_t CheckIfPgoSchemaIsCompatibleAndSetOffsets(const uint8_t *pByte, size_t cbDataMax, ICorJitInfo::PgoInstrumentationSchema* schemaTable, size_t cSchemas)
 {
     size_t nMatched = 0;
     size_t initialOffset = cbDataMax;
@@ -383,13 +382,17 @@ inline bool CheckIfPgoSchemaIsCompatibleAndSetOffsets(const uint8_t *pByte, size
             schemaTable[nMatched].Offset = schema.Offset;
             nMatched++;
         }
+        else
+        {
+            return false;
+        }
 
         return true;
     };
 
     ReadInstrumentationSchemaWithLayout(pByte, cbDataMax, initialOffset, handler);
 
-    return (nMatched == cSchemas);
+    return nMatched;
 }
 
 inline bool ReadInstrumentationSchemaWithLayoutIntoSArray(const uint8_t *pByte, size_t cbDataMax, size_t initialOffset, SArray<ICorJitInfo::PgoInstrumentationSchema>* pSchemas)

src/coreclr/jit/compiler.cpp

@@ -2086,8 +2086,6 @@ void Compiler::compInitOptions(JitFlags* jitFlags)
     {
         // The following flags are lost when inlining. (They are removed in
         // Compiler::fgInvokeInlineeCompiler().)
-        assert(!jitFlags->IsSet(JitFlags::JIT_FLAG_BBINSTR));
-        assert(!jitFlags->IsSet(JitFlags::JIT_FLAG_BBINSTR_IF_LOOPS));
         assert(!jitFlags->IsSet(JitFlags::JIT_FLAG_PROF_ENTERLEAVE));
         assert(!jitFlags->IsSet(JitFlags::JIT_FLAG_DEBUG_EnC));
         assert(!jitFlags->IsSet(JitFlags::JIT_FLAG_REVERSE_PINVOKE));
@@ -7024,6 +7022,14 @@ int Compiler::compCompileHelper(CORINFO_MODULE_HANDLE classPtr,
 
     compSetOptimizationLevel();
 
+#ifdef DEBUG
+    if (JitConfig.JitInstrumentIfOptimizing() && opts.OptimizationEnabled() && !IsReadyToRun())
+    {
+        JITDUMP("\nForcibly enabling instrumentation\n");
+        opts.jitFlags->Set(JitFlags::JIT_FLAG_BBINSTR);
+    }
+#endif
+
     if ((JitConfig.JitDisasmOnlyOptimized() != 0) && (!opts.OptimizationEnabled()))
     {
         // Disable JitDisasm for non-optimized code.

src/coreclr/jit/compiler.h

@@ -6431,6 +6431,7 @@ class Compiler
 
     PhaseStatus fgPrepareToInstrumentMethod();
     PhaseStatus fgInstrumentMethod();
+    PhaseStatus fgInstrumentMethodCore();
     PhaseStatus fgIncorporateProfileData();
     bool        fgIncorporateBlockCounts();
     bool        fgIncorporateEdgeCounts();

src/coreclr/jit/fginline.cpp

@@ -1415,8 +1415,6 @@ void Compiler::fgInvokeInlineeCompiler(GenTreeCall* call, InlineResult* inlineRe
 
             // The following flags are lost when inlining.
             // (This is checked in Compiler::compInitOptions().)
-            compileFlagsForInlinee.Clear(JitFlags::JIT_FLAG_BBINSTR);
-            compileFlagsForInlinee.Clear(JitFlags::JIT_FLAG_BBINSTR_IF_LOOPS);
             compileFlagsForInlinee.Clear(JitFlags::JIT_FLAG_PROF_ENTERLEAVE);
             compileFlagsForInlinee.Clear(JitFlags::JIT_FLAG_DEBUG_EnC);
             compileFlagsForInlinee.Clear(JitFlags::JIT_FLAG_REVERSE_PINVOKE);

src/coreclr/jit/fgprofile.cpp

@@ -484,14 +484,6 @@ void BlockCountInstrumentor::RelocateProbes()
 
     JITDUMP("Optimized + instrumented + potential tail calls --- preparing to relocate edge probes\n");
 
-    // We should be in a root method compiler instance. We currently do not instrument inlinees.
-    //
-    // Relaxing this will require changes below because inlinee compilers
-    // share the root compiler flow graph (and hence bb epoch), and flow
-    // from inlinee tail calls to returns can be more complex.
-    //
-    assert(!m_comp->compIsForInlining());
-
     // Keep track of return blocks needing special treatment.
     //
     ArrayStack<BasicBlock*> criticalPreds(m_comp->getAllocator(CMK_Pgo));
@@ -850,26 +842,16 @@ void Compiler::WalkSpanningTree(SpanningTreeVisitor* visitor)
     // Push the method entry and all EH handler region entries on the stack.
     // (push method entry last so it's visited first).
     //
-    // Note inlinees are "contaminated" with root method EH structures.
-    // We know the inlinee itself doesn't have EH, so we only look at
-    // handlers for root methods.
-    //
-    // If we ever want to support inlining methods with EH, we'll
-    // have to revisit this.
-    //
-    if (!compIsForInlining())
+    for (EHblkDsc* const HBtab : EHClauses(this))
     {
-        for (EHblkDsc* const HBtab : EHClauses(this))
+        BasicBlock* hndBegBB = HBtab->ebdHndBeg;
+        stack.Push(hndBegBB);
+        BitVecOps::AddElemD(&traits, marked, hndBegBB->bbID);
+        if (HBtab->HasFilter())
         {
-            BasicBlock* hndBegBB = HBtab->ebdHndBeg;
-            stack.Push(hndBegBB);
-            BitVecOps::AddElemD(&traits, marked, hndBegBB->bbID);
-            if (HBtab->HasFilter())
-            {
-                BasicBlock* filterBB = HBtab->ebdFilter;
-                stack.Push(filterBB);
-                BitVecOps::AddElemD(&traits, marked, filterBB->bbID);
-            }
+            BasicBlock* filterBB = HBtab->ebdFilter;
+            stack.Push(filterBB);
+            BitVecOps::AddElemD(&traits, marked, filterBB->bbID);
         }
     }
 
@@ -1559,14 +1541,6 @@ void EfficientEdgeCountInstrumentor::RelocateProbes()
 
     JITDUMP("Optimized + instrumented + potential tail calls --- preparing to relocate edge probes\n");
 
-    // We should be in a root method compiler instance. We currently do not instrument inlinees.
-    //
-    // Relaxing this will require changes below because inlinee compilers
-    // share the root compiler flow graph (and hence bb epoch), and flow
-    // from inlinee tail calls to returns can be more complex.
-    //
-    assert(!m_comp->compIsForInlining());
-
     // We may need to track the critical predecessors of some blocks.
     //
     ArrayStack<BasicBlock*> criticalPreds(m_comp->getAllocator(CMK_Pgo));
@@ -2508,7 +2482,11 @@ void ValueInstrumentor::Instrument(BasicBlock* block, Schema& schema, uint8_t* p
 //
 PhaseStatus Compiler::fgPrepareToInstrumentMethod()
 {
-    noway_assert(!compIsForInlining());
+    if (compIsForInlining() && JitConfig.JitInstrumentInlinees() == 0)
+    {
+        JITDUMP("Inlinee instrumentation disabled by config\n");
+        return PhaseStatus::MODIFIED_NOTHING;
+    }
 
     // Choose instrumentation technology.
     //
@@ -2626,6 +2604,59 @@ PhaseStatus Compiler::fgPrepareToInstrumentMethod()
 //   appropriate phase status
 //
 // Note:
+//   Wrapper around fgInstrumentMethodCore, which handles
+//   special cases when instrumenting inlinees.
+//
+PhaseStatus Compiler::fgInstrumentMethod()
+{
+    // If this is an inlinee that returns a value, and we don't have a return
+    // value temp, the return value tree may not be linked into the return block.
+    //
+    // Temporarily link it in so the passes below can operate on it as needed.
+    //
+    BasicBlock* retBB                 = nullptr;
+    Statement*  tempInlineeReturnStmt = nullptr;
+
+    if (compIsForInlining())
+    {
+        if (JitConfig.JitInstrumentInlinees() == 0)
+        {
+            JITDUMP("Inlinee instrumentation disabled by config\n");
+            return PhaseStatus::MODIFIED_NOTHING;
+        }
+
+        GenTreeRetExpr* const retExpr = impInlineInfo->inlineCandidateInfo->retExpr;
+
+        // If there's a retExpr but no gtSubstBB, we assume the retExpr is a temp
+        // and so not interesting to instrumentation.
+        //
+        if ((retExpr != nullptr) && (retExpr->gtSubstBB != nullptr))
+        {
+            assert(retExpr->gtSubstExpr != nullptr);
+            retBB                 = retExpr->gtSubstBB;
+            tempInlineeReturnStmt = fgNewStmtAtEnd(retBB, retExpr->gtSubstExpr);
+            JITDUMP("Temporarily adding ret expr [%06u] to " FMT_BB "\n", dspTreeID(retExpr->gtSubstExpr),
+                    retBB->bbNum);
+        }
+    }
+
+    PhaseStatus status = fgInstrumentMethodCore();
+
+    if (tempInlineeReturnStmt != nullptr)
+    {
+        fgRemoveStmt(retBB, tempInlineeReturnStmt);
+    }
+
+    return status;
+}
+
+//------------------------------------------------------------------------
+// fgInstrumentMethodCore: add instrumentation probes to the method
+//
+// Returns:
+//   appropriate phase status
+//
+// Note:
 //
 //   By default this instruments each non-internal block with
 //   a counter probe.
@@ -2635,10 +2666,9 @@ PhaseStatus Compiler::fgPrepareToInstrumentMethod()
 //   Probe structure is described by a schema array, which is created
 //   here based on flowgraph and IR structure.
 //
-PhaseStatus Compiler::fgInstrumentMethod()
-{
-    noway_assert(!compIsForInlining());
 
+PhaseStatus Compiler::fgInstrumentMethodCore()
+{
     // Make post-import preparations.
     //
     const bool isPreImport = false;
@@ -2718,7 +2748,7 @@ PhaseStatus Compiler::fgInstrumentMethod()
     //
     // If this is an OSR method, we should use the same buffer that the Tier0 method used.
     //
-    // This is supported by allocPgoInsrumentationDataBySchema, which will verify the schema
+    // This is supported by allocPgoInstrumentationBySchema, which will verify the schema
     // we provide here matches the one from Tier0, and will fill in the data offsets in
     // our schema properly.
     //

src/coreclr/jit/importercalls.cpp

@@ -1294,13 +1294,10 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
             compCurBB->SetFlags(BBF_RECURSIVE_TAILCALL);
         }
 
-        // We only do these OSR checks in the root method because:
-        // * If we fail to import the root method entry when importing the root method, we can't go back
-        //    and import it during inlining. So instead of checking just for recursive tail calls we also
-        //    have to check for anything that might introduce a recursive tail call.
-        // * We only instrument root method blocks in OSR methods,
+        // If we might be instrumenting, flag blocks that might be tail call successors
+        // so we can relocate probes before the calls.
         //
-        if ((opts.IsInstrumentedAndOptimized() || opts.IsOSR()) && !compIsForInlining())
+        if (opts.IsInstrumentedAndOptimized() || opts.IsOSR())
         {
             // If a root method tail call candidate block is not a BBJ_RETURN, it should have a unique
             // BBJ_RETURN successor. Mark that successor so we can handle it specially during profile
@@ -9298,16 +9295,12 @@ bool Compiler::impConsiderCallProbe(GenTreeCall* call, IL_OFFSET ilOffset)
     // to revisit -- if we can devirtualize we should be able to
     // suppress the probe.
     //
-    // We strip BBINSTR from inlinees currently, so we'll only
-    // do this for the root method calls.
-    //
     if (!opts.jitFlags->IsSet(JitFlags::JIT_FLAG_BBINSTR))
     {
         return false;
     }
 
     assert(opts.OptimizationDisabled() || opts.IsInstrumentedAndOptimized());
-    assert(!compIsForInlining());
 
     // During importation, optionally flag this block as one that
     // contains calls requiring class profiling. Ideally perhaps

src/coreclr/jit/jitconfigvalues.h

@@ -754,6 +754,11 @@ RELEASE_CONFIG_INTEGER(JitVTableProfiling, "JitVTableProfiling", 0)       // Pro
 RELEASE_CONFIG_INTEGER(JitEdgeProfiling, "JitEdgeProfiling", 1)           // Profile edges instead of blocks
 RELEASE_CONFIG_INTEGER(JitCollect64BitCounts, "JitCollect64BitCounts", 0) // Collect counts as 64-bit values.
 
+CONFIG_INTEGER(JitInstrumentIfOptimizing, "JitInstrumentIfOptimizing", 0) // 1: Always add instrumentation if optimizing
+                                                                          // and not prejitting
+
+RELEASE_CONFIG_INTEGER(JitInstrumentInlinees, "JitInstrumentInlinees", 1) // Add instrumentation to inlined methods
+
 // Profile consumption options
 RELEASE_CONFIG_INTEGER(JitDisablePGO, "JitDisablePGO", 0)     // Ignore PGO data for all methods
 CONFIG_STRING(JitEnablePGORange, "JitEnablePGORange")         // Enable PGO data for only some methods

src/coreclr/vm/jitinterface.cpp

@@ -12460,7 +12460,8 @@ HRESULT CEEJitInfo::allocPgoInstrumentationBySchema(
     JIT_TO_EE_TRANSITION();
 
 #ifdef FEATURE_PGO
-    hr = PgoManager::allocPgoInstrumentationBySchema(m_pMethodBeingCompiled, pSchema, countSchemaItems, pInstrumentationData);
+    MethodDesc* pMD = (MethodDesc*)ftnHnd;
+    hr = PgoManager::allocPgoInstrumentationBySchema(pMD, pSchema, countSchemaItems, pInstrumentationData);
 #else
     _ASSERTE(!"allocMethodBlockCounts not implemented on CEEJitInfo!");
     hr = E_NOTIMPL;

src/coreclr/vm/pgo.cpp

@@ -749,8 +749,10 @@ HRESULT PgoManager::allocPgoInstrumentationBySchemaInstance(MethodDesc* pMD,
         HeaderList *currentHeaderList = m_pgoHeaders;
         if (currentHeaderList != NULL)
         {
-            if (!CheckIfPgoSchemaIsCompatibleAndSetOffsets(currentHeaderList->header.GetData(), currentHeaderList->header.countsOffset, pSchema, countSchemaItems))
+            size_t matchedCount = CheckIfPgoSchemaIsCompatibleAndSetOffsets(currentHeaderList->header.GetData(), currentHeaderList->header.countsOffset, pSchema, countSchemaItems);
+            if (matchedCount != (size_t) countSchemaItems)
             {
+                // We don't expect to see schema mismatches for dynamic methods
                 return E_NOTIMPL;
             }
             _ASSERTE(currentHeaderList->header.method == pMD);
@@ -779,12 +781,29 @@ HRESULT PgoManager::allocPgoInstrumentationBySchemaInstance(MethodDesc* pMD,
         HeaderList* existingData = laPgoManagerThis->m_pgoDataLookup.Lookup(pMD);
         if (existingData != NULL)
         {
-            if (!CheckIfPgoSchemaIsCompatibleAndSetOffsets(existingData->header.GetData(), existingData->header.countsOffset, pSchema, countSchemaItems))
+            size_t matchedCount = CheckIfPgoSchemaIsCompatibleAndSetOffsets(existingData->header.GetData(), existingData->header.countsOffset, pSchema, countSchemaItems);
+            if (matchedCount != (size_t) countSchemaItems)
             {
-                return E_NOTIMPL;
+                // Assume existing data is stale static PGO. Remove it add add the new schema below.
+                //
+                laPgoManagerThis->m_pgoDataLookup.Remove(pMD);
+
+                // We may have altered some of the offsets if we had a partial match. Redo the layout.
+                //
+                ICorJitInfo::PgoInstrumentationSchema prevSchema;
+                memset(&prevSchema, 0, sizeof(ICorJitInfo::PgoInstrumentationSchema));
+                prevSchema.Offset = offsetOfInstrumentationDataFromStartOfDataRegion;
+                for (UINT32 iSchema = 0; iSchema < countSchemaItems; iSchema++)
+                {
+                    LayoutPgoInstrumentationSchema(prevSchema, &pSchema[iSchema]);
+                    prevSchema = pSchema[iSchema];
+                }
+            }
+            else
+            {
+                *pInstrumentationData = existingData->header.GetData();
+                return S_OK;
             }
-            *pInstrumentationData = existingData->header.GetData();
-            return S_OK;
         }
 
         AllocMemTracker loaderHeapAllocation;

src/tests/Common/testenvironment.proj

@@ -72,6 +72,7 @@
       DOTNET_JitOptRepeat;
       DOTNET_JitOptRepeatCount;
       DOTNET_JitDoReversePostOrderLayout;
+      DOTNET_JitInstrumentIfOptimizing;
     </DOTNETVariables>
   </PropertyGroup>
   <ItemGroup>
@@ -199,6 +200,7 @@
     <TestEnvironment Include="gcstandalone" Condition="'$(TargetsWindows)' != 'true'" GCName="libclrgc.so"/>
     <TestEnvironment Include="gcstandaloneserver" Condition="'$(TargetsWindows)' == 'true'" gcServer="1" GCName="clrgc.dll"/>
     <TestEnvironment Include="gcstandaloneserver" Condition="'$(TargetsWindows)' != 'true'" gcServer="1" GCName="libclrgc.so"/>
+    <TestEnvironment Include="instrument_if_optimizing" TieredPGO="1" TieredCompilation="1" TC_QuickJitForLoops="1" JitInstrumentIfOptimizing="1" />
   </ItemGroup>
 
   <!-- Create a random number for use by random JitStress. Note that we want to define a random value here,