JIT: Always track the context for late devirt

src/coreclr/jit/fginline.cpp

@@ -590,14 +590,16 @@ class SubstitutePlaceholdersAndDevirtualizeWalker : public GenTreeVisitor<Substi
 #endif // DEBUG
 
                 CORINFO_CONTEXT_HANDLE context                = nullptr;
+                InlineContext*         inlinersContext        = m_compiler->compInlineContext;
                 CORINFO_METHOD_HANDLE  method                 = call->gtCallMethHnd;
                 unsigned               methodFlags            = 0;
                 const bool             isLateDevirtualization = true;
                 const bool             explicitTailCall       = call->IsTailPrefixedCall();
 
                 if ((call->gtCallMoreFlags & GTF_CALL_M_HAS_LATE_DEVIRT_INFO) != 0)
                 {
-                    context = call->gtLateDevirtualizationInfo->exactContextHnd;
+                    context         = call->gtLateDevirtualizationInfo->exactContextHnd;
+                    inlinersContext = call->gtLateDevirtualizationInfo->inlinersContext;
                     // Note: we might call this multiple times for the same trees.
                     // If the devirtualization below succeeds, the call becomes
                     // non-virtual and we won't get here again. If it does not
@@ -613,9 +615,11 @@ class SubstitutePlaceholdersAndDevirtualizeWalker : public GenTreeVisitor<Substi
                 if (!call->IsVirtual())
                 {
                     assert(context != nullptr);
-                    CORINFO_CALL_INFO callInfo = {};
-                    callInfo.hMethod           = method;
-                    callInfo.methodFlags       = methodFlags;
+                    assert(inlinersContext != nullptr);
+                    m_compiler->compInlineContext = inlinersContext;
+                    CORINFO_CALL_INFO callInfo    = {};
+                    callInfo.hMethod              = method;
+                    callInfo.methodFlags          = methodFlags;
                     m_compiler->impMarkInlineCandidate(call, context, false, &callInfo);
 
                     if (call->IsInlineCandidate())
@@ -652,9 +656,6 @@ class SubstitutePlaceholdersAndDevirtualizeWalker : public GenTreeVisitor<Substi
                             *pTree = retExpr;
                         }
 
-                        call->GetSingleInlineCandidateInfo()->exactContextHandle = context;
-                        INDEBUG(call->GetSingleInlineCandidateInfo()->inlinersContext = call->gtInlineContext);
-
                         JITDUMP("New inline candidate due to late devirtualization:\n");
                         DISPTREE(call);
                     }

src/coreclr/jit/gentree.cpp

@@ -10013,6 +10013,8 @@ GenTreeCall* Compiler::gtCloneExprCallHelper(GenTreeCall* tree)
         copy->gtInlineInfoCount     = tree->gtInlineInfoCount;
     }
 
+    copy->gtLateDevirtualizationInfo = tree->gtLateDevirtualizationInfo;
+
     copy->gtCallType   = tree->gtCallType;
     copy->gtReturnType = tree->gtReturnType;
 

src/coreclr/jit/gentree.h

@@ -5758,11 +5758,13 @@ struct GenTreeCall final : public GenTree
         jitstd::vector<InlineCandidateInfo*>* gtInlineCandidateInfoList;
 
         HandleHistogramProfileCandidateInfo* gtHandleHistogramProfileCandidateInfo;
-        LateDevirtualizationInfo*            gtLateDevirtualizationInfo;
+
         CORINFO_GENERIC_HANDLE compileTimeHelperArgumentHandle; // Used to track type handle argument of dynamic helpers
         void*                  gtDirectCallAddress; // Used to pass direct call address between lower and codegen
     };
 
+    LateDevirtualizationInfo* gtLateDevirtualizationInfo;
+
     // expression evaluated after args are placed which determines the control target
     GenTree* gtControlExpr;
 

src/coreclr/jit/importercalls.cpp

@@ -1241,6 +1241,20 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
 
         // Is it an inline candidate?
         impMarkInlineCandidate(call, exactContextHnd, exactContextNeedsRuntimeLookup, callInfo);
+
+        // If the call is virtual, record the inliner's context for possible use during late devirt inlining.
+        // Also record the generics context if there is any.
+        //
+        if (call->AsCall()->IsVirtual() && (call->AsCall()->gtCallType != CT_INDIRECT))
+        {
+            JITDUMP("\nSaving generic context %p and inline context %p for call [%06u]\n", dspPtr(exactContextHnd),
+                    dspPtr(compInlineContext), dspTreeID(call->AsCall()));
+            call->AsCall()->gtCallMoreFlags |= GTF_CALL_M_HAS_LATE_DEVIRT_INFO;
+            LateDevirtualizationInfo* const info       = new (this, CMK_Inlining) LateDevirtualizationInfo;
+            info->exactContextHnd                      = exactContextHnd;
+            info->inlinersContext                      = compInlineContext;
+            call->AsCall()->gtLateDevirtualizationInfo = info;
+        }
     }
 
     // Extra checks for tail calls and tail recursion.
@@ -1431,22 +1445,6 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
             }
             else
             {
-                // If the call is virtual, and has a generics context, and is not going to have a class probe,
-                // record the context for possible use during late devirt.
-                //
-                // If we ever want to devirt at Tier0, and/or see issues where OSR methods under PGO lose
-                // important devirtualizations, we'll want to allow both a class probe and a captured context.
-                //
-                if (origCall->IsVirtual() && (origCall->gtCallType != CT_INDIRECT) && (exactContextHnd != nullptr) &&
-                    (origCall->gtHandleHistogramProfileCandidateInfo == nullptr))
-                {
-                    JITDUMP("\nSaving context %p for call [%06u]\n", dspPtr(exactContextHnd), dspTreeID(origCall));
-                    origCall->gtCallMoreFlags |= GTF_CALL_M_HAS_LATE_DEVIRT_INFO;
-                    LateDevirtualizationInfo* const info = new (this, CMK_Inlining) LateDevirtualizationInfo;
-                    info->exactContextHnd                = exactContextHnd;
-                    origCall->gtLateDevirtualizationInfo = info;
-                }
-
                 if (isFatPointerCandidate)
                 {
                     // fatPointer candidates should be in statements of the form call() or var = call().

src/coreclr/jit/inline.h

@@ -634,6 +634,7 @@ struct InlineCandidateInfo : public HandleHistogramProfileCandidateInfo
 struct LateDevirtualizationInfo
 {
     CORINFO_CONTEXT_HANDLE exactContextHnd;
+    InlineContext*         inlinersContext;
 };
 
 // InlArgInfo describes inline candidate argument properties.