JIT: revise how the jit tracks use of generics context

src/coreclr/src/jit/compiler.h

@@ -2854,9 +2854,7 @@ class Compiler
     // Get the element handle for an array of ref type.
     CORINFO_CLASS_HANDLE gtGetArrayElementClassHandle(GenTree* array);
     // Get a class handle from a helper call argument
-    CORINFO_CLASS_HANDLE gtGetHelperArgClassHandle(GenTree*  array,
-                                                   unsigned* runtimeLookupCount = nullptr,
-                                                   GenTree** handleTree         = nullptr);
+    CORINFO_CLASS_HANDLE gtGetHelperArgClassHandle(GenTree* array, GenTree** handleTree = nullptr);
     // Get the class handle for a field
     CORINFO_CLASS_HANDLE gtGetFieldClassHandle(CORINFO_FIELD_HANDLE fieldHnd, bool* pIsExact, bool* pIsNonNull);
     // Check if this tree is a gc static base helper call
@@ -3129,7 +3127,7 @@ class Compiler
 
 #endif // defined(DEBUG) && defined(TARGET_X86)
 
-    unsigned lvaGenericsContextUseCount;
+    bool lvaGenericsContextInUse;
 
     bool lvaKeepAliveAndReportThis(); // Synchronized instance method of a reference type, or
                                       // CORINFO_GENERICS_CTXT_FROM_THIS?

src/coreclr/src/jit/compiler.hpp

@@ -1938,9 +1938,9 @@ inline bool Compiler::lvaKeepAliveAndReportThis()
         if (opts.compDbgCode)
             return true;
 
-        if (lvaGenericsContextUseCount > 0)
+        if (lvaGenericsContextInUse)
         {
-            JITDUMP("Reporting this as generic context: %u refs\n", lvaGenericsContextUseCount);
+            JITDUMP("Reporting this as generic context\n");
             return true;
         }
     }
@@ -1951,14 +1951,11 @@ inline bool Compiler::lvaKeepAliveAndReportThis()
     // because collectible types need the generics context when gc-ing.
     if (genericsContextIsThis)
     {
-        const bool isUsed   = lvaGenericsContextUseCount > 0;
         const bool mustKeep = (info.compMethodInfo->options & CORINFO_GENERICS_CTXT_KEEP_ALIVE) != 0;
 
-        if (isUsed || mustKeep)
+        if (lvaGenericsContextInUse || mustKeep)
         {
-            JITDUMP("Reporting this as generic context: %u refs%s\n", lvaGenericsContextUseCount,
-                    mustKeep ? ", must keep" : "");
-
+            JITDUMP("Reporting this as generic context: %s\n", mustKeep ? "must keep" : "referenced");
             return true;
         }
     }
@@ -1986,7 +1983,7 @@ inline bool Compiler::lvaReportParamTypeArg()
 
         // Otherwise, if an exact type parameter is needed in the body, report the generics context.
         // We do this because collectible types needs the generics context when gc-ing.
-        if (lvaGenericsContextUseCount > 0)
+        if (lvaGenericsContextInUse)
         {
             return true;
         }

src/coreclr/src/jit/flowgraph.cpp

@@ -7833,7 +7833,7 @@ GenTree* Compiler::fgGetCritSectOfStaticMethod()
         // Collectible types requires that for shared generic code, if we use the generic context paramter
         // that we report it. (This is a conservative approach, we could detect some cases particularly when the
         // context parameter is this that we don't need the eager reporting logic.)
-        lvaGenericsContextUseCount++;
+        lvaGenericsContextInUse = true;
 
         switch (kind.runtimeLookupKind)
         {
@@ -7847,13 +7847,15 @@ GenTree* Compiler::fgGetCritSectOfStaticMethod()
             {
                 // In this case, the hidden param is the class handle.
                 tree = gtNewLclvNode(info.compTypeCtxtArg, TYP_I_IMPL);
+                tree->gtFlags |= GTF_VAR_CONTEXT;
                 break;
             }
 
             case CORINFO_LOOKUP_METHODPARAM:
             {
                 // In this case, the hidden param is the method handle.
                 tree = gtNewLclvNode(info.compTypeCtxtArg, TYP_I_IMPL);
+                tree->gtFlags |= GTF_VAR_CONTEXT;
                 // Call helper CORINFO_HELP_GETCLASSFROMMETHODPARAM to get the class handle
                 // from the method handle.
                 tree = gtNewHelperCallNode(CORINFO_HELP_GETCLASSFROMMETHODPARAM, TYP_I_IMPL, gtNewCallArgs(tree));
@@ -23720,38 +23722,6 @@ Statement* Compiler::fgInlinePrependStatements(InlineInfo* inlineInfo)
 
 void Compiler::fgInlineAppendStatements(InlineInfo* inlineInfo, BasicBlock* block, Statement* stmtAfter)
 {
-    // If this inlinee was passed a runtime lookup generic context and
-    // ignores it, we can decrement the "generic context was used" ref
-    // count, because we created a new lookup tree and incremented the
-    // count when we imported the type parameter argument to pass to
-    // the inlinee. See corresponding logic in impImportCall that
-    // checks the sig for CORINFO_CALLCONV_PARAMTYPE.
-    //
-    // Does this method require a context (type) parameter?
-    if ((inlineInfo->inlineCandidateInfo->methInfo.args.callConv & CORINFO_CALLCONV_PARAMTYPE) != 0)
-    {
-        // Did the computation of that parameter require the
-        // caller to perform a runtime lookup?
-        if (inlineInfo->inlineCandidateInfo->exactContextNeedsRuntimeLookup)
-        {
-            // Fetch the temp for the generic context as it would
-            // appear in the inlinee's body.
-            const unsigned typeCtxtArg = inlineInfo->typeContextArg;
-            const unsigned tmpNum      = inlineInfo->lclTmpNum[typeCtxtArg];
-
-            // Was it used in the inline body?
-            if (tmpNum == BAD_VAR_NUM)
-            {
-                // No -- so the associated runtime lookup is not needed
-                // and also no longer provides evidence that the generic
-                // context should be kept alive.
-                JITDUMP("Inlinee ignores runtime lookup generics context\n");
-                assert(lvaGenericsContextUseCount > 0);
-                lvaGenericsContextUseCount--;
-            }
-        }
-    }
-
     // Null out any gc ref locals
     if (!inlineInfo->HasGcRefLocals())
     {

src/coreclr/src/jit/gentree.cpp

@@ -9947,6 +9947,13 @@ void Compiler::gtDispNode(GenTree* tree, IndentStack* indentStack, __in __in_z _
                     --msgLength;
                     break;
                 }
+                if (tree->gtFlags & GTF_VAR_CONTEXT)
+                {
+                    printf("!");
+                    --msgLength;
+                    break;
+                }
+
                 goto DASH;
 
             case GT_EQ:
@@ -12478,11 +12485,10 @@ GenTree* Compiler::gtFoldTypeCompare(GenTree* tree)
         GenTree*             op2TunneledHandle  = nullptr;
         CORINFO_CLASS_HANDLE cls1Hnd            = NO_CLASS_HANDLE;
         CORINFO_CLASS_HANDLE cls2Hnd            = NO_CLASS_HANDLE;
-        unsigned             runtimeLookupCount = 0;
 
         // Try and find class handles from op1 and op2
-        cls1Hnd = gtGetHelperArgClassHandle(op1ClassFromHandle, &runtimeLookupCount, &op1TunneledHandle);
-        cls2Hnd = gtGetHelperArgClassHandle(op2ClassFromHandle, &runtimeLookupCount, &op2TunneledHandle);
+        cls1Hnd = gtGetHelperArgClassHandle(op1ClassFromHandle, &op1TunneledHandle);
+        cls2Hnd = gtGetHelperArgClassHandle(op2ClassFromHandle, &op2TunneledHandle);
 
         // If we have both class handles, try and resolve the type equality test completely.
         bool resolveFailed = false;
@@ -12501,11 +12507,6 @@ GenTree* Compiler::gtFoldTypeCompare(GenTree* tree)
                 const int  compareResult = operatorIsEQ ^ typesAreEqual ? 0 : 1;
                 JITDUMP("Runtime reports comparison is known at jit time: %u\n", compareResult);
                 GenTree* result = gtNewIconNode(compareResult);
-
-                // Any runtime lookups that fed into this compare are
-                // now dead code, so they no longer require the runtime context.
-                assert(lvaGenericsContextUseCount >= runtimeLookupCount);
-                lvaGenericsContextUseCount -= runtimeLookupCount;
                 return result;
             }
             else
@@ -12668,15 +12669,12 @@ GenTree* Compiler::gtFoldTypeCompare(GenTree* tree)
 //
 // Arguments:
 //    tree - tree that passes the handle to the helper
-//    runtimeLookupCount [optional, in/out] - incremented if tree was a runtime lookup
 //    handleTree [optional, out] - set to the literal operand tree for indirect handles
 //
 // Returns:
 //    The compile time class handle if known.
 //
-CORINFO_CLASS_HANDLE Compiler::gtGetHelperArgClassHandle(GenTree*  tree,
-                                                         unsigned* runtimeLookupCount,
-                                                         GenTree** handleTree)
+CORINFO_CLASS_HANDLE Compiler::gtGetHelperArgClassHandle(GenTree* tree, GenTree** handleTree)
 {
     CORINFO_CLASS_HANDLE result = NO_CLASS_HANDLE;
 
@@ -12696,11 +12694,6 @@ CORINFO_CLASS_HANDLE Compiler::gtGetHelperArgClassHandle(GenTree*  tree,
     else if (tree->OperGet() == GT_RUNTIMELOOKUP)
     {
         result = tree->AsRuntimeLookup()->GetClassHandle();
-
-        if (runtimeLookupCount != nullptr)
-        {
-            *runtimeLookupCount = *runtimeLookupCount + 1;
-        }
     }
     // Or something reached indirectly
     else if (tree->gtOper == GT_IND)

src/coreclr/src/jit/gentree.h

@@ -792,6 +792,8 @@ struct GenTree
 #define GTF_VAR_CAST        0x01000000 // GT_LCL_VAR -- has been explictly cast (variable node may not be type of local)
 #define GTF_VAR_ITERATOR    0x00800000 // GT_LCL_VAR -- this is a iterator reference in the loop condition
 #define GTF_VAR_CLONED      0x00400000 // GT_LCL_VAR -- this node has been cloned or is a clone
+#define GTF_VAR_CONTEXT     0x00200000 // GT_LCL_VAR -- this node is part of a runtime lookup
+
                                        // Relevant for inlining optimizations (see fgInlinePrependStatements)
 
 #define GTF_VAR_ARR_INDEX   0x00000020 // The variable is part of (the index portion of) an array index expression.

src/coreclr/src/jit/importer.cpp

@@ -1983,12 +1983,13 @@ GenTree* Compiler::getRuntimeContextTree(CORINFO_RUNTIME_LOOKUP_KIND kind)
     // Collectible types requires that for shared generic code, if we use the generic context parameter
     // that we report it. (This is a conservative approach, we could detect some cases particularly when the
     // context parameter is this that we don't need the eager reporting logic.)
-    lvaGenericsContextUseCount++;
+    lvaGenericsContextInUse = true;
 
     if (kind == CORINFO_LOOKUP_THISOBJ)
     {
         // this Object
         ctxTree = gtNewLclvNode(info.compThisArg, TYP_REF);
+        ctxTree->gtFlags |= GTF_VAR_CONTEXT;
 
         // Vtable pointer of this object
         ctxTree = gtNewOperNode(GT_IND, TYP_I_IMPL, ctxTree);
@@ -2000,6 +2001,7 @@ GenTree* Compiler::getRuntimeContextTree(CORINFO_RUNTIME_LOOKUP_KIND kind)
         assert(kind == CORINFO_LOOKUP_METHODPARAM || kind == CORINFO_LOOKUP_CLASSPARAM);
 
         ctxTree = gtNewLclvNode(info.compTypeCtxtArg, TYP_I_IMPL); // Exact method descriptor as passed in as last arg
+        ctxTree->gtFlags |= GTF_VAR_CONTEXT;
     }
     return ctxTree;
 }
@@ -18613,7 +18615,7 @@ void Compiler::impCheckCanInline(GenTreeCall*           call,
 //
 //   Checks for various inline blocking conditions and makes notes in
 //   the inline info arg table about the properties of the actual. These
-//   properties are used later by impFetchArg to determine how best to
+//   properties are used later by impInlineFetchArg to determine how best to
 //   pass the argument into the inlinee.
 
 void Compiler::impInlineRecordArgInfo(InlineInfo*   pInlineInfo,

src/coreclr/src/jit/lclvars.cpp

@@ -39,7 +39,7 @@ void Compiler::lvaInit()
     /* We haven't allocated stack variables yet */
     lvaRefCountState = RCS_INVALID;
 
-    lvaGenericsContextUseCount = 0;
+    lvaGenericsContextInUse = false;
 
     lvaTrackedToVarNumSize = 0;
     lvaTrackedToVarNum     = nullptr;
@@ -3606,6 +3606,8 @@ var_types LclVarDsc::lvaArgType()
 //     eligible for assertion prop, single defs, and tracking which blocks
 //     hold uses.
 //
+//     Looks for uses of generic context and sets lvaGenericsContextInUse.
+//
 //     In checked builds:
 //
 //     Verifies that local accesses are consistenly typed.
@@ -3711,6 +3713,17 @@ void Compiler::lvaMarkLclRefs(GenTree* tree, BasicBlock* block, Statement* stmt,
 
     /* This must be a local variable reference */
 
+    // See if this is a generics context use.
+    if ((tree->gtFlags & GTF_VAR_CONTEXT) != 0)
+    {
+        assert(tree->OperIs(GT_LCL_VAR));
+        if (!lvaGenericsContextInUse)
+        {
+            JITDUMP("-- generic context in use at [%06u]\n", dspTreeID(tree));
+            lvaGenericsContextInUse = true;
+        }
+    }
+
     assert((tree->gtOper == GT_LCL_VAR) || (tree->gtOper == GT_LCL_FLD));
     unsigned lclNum = tree->AsLclVarCommon()->GetLclNum();
 
@@ -4009,24 +4022,26 @@ void Compiler::lvaMarkLocalVars()
         return;
     }
 
-#if ASSERTION_PROP
-    assert(opts.OptimizationEnabled());
-
-    // Note: optAddCopies() depends on lvaRefBlks, which is set in lvaMarkLocalVars(BasicBlock*), called above.
-    optAddCopies();
-#endif
+    const bool reportParamTypeArg = lvaReportParamTypeArg();
 
+    // Update bookkeeping on the generic context.
     if (lvaKeepAliveAndReportThis())
     {
-        lvaTable[0].lvImplicitlyReferenced = 1;
-        // This isn't strictly needed as we will make a copy of the param-type-arg
-        // in the prolog. However, this ensures that the LclVarDsc corresponding to
-        // info.compTypeCtxtArg is valid.
+        lvaGetDesc(0u)->lvImplicitlyReferenced = reportParamTypeArg;
     }
     else if (lvaReportParamTypeArg())
     {
-        lvaTable[info.compTypeCtxtArg].lvImplicitlyReferenced = 1;
+        // We should have a context arg.
+        assert(info.compTypeCtxtArg != BAD_VAR_NUM);
+        lvaGetDesc(info.compTypeCtxtArg)->lvImplicitlyReferenced = reportParamTypeArg;
     }
+
+#if ASSERTION_PROP
+    assert(opts.OptimizationEnabled());
+
+    // Note: optAddCopies() depends on lvaRefBlks, which is set in lvaMarkLocalVars(BasicBlock*), called above.
+    optAddCopies();
+#endif
 }
 
 //------------------------------------------------------------------------
@@ -4046,7 +4061,10 @@ void Compiler::lvaMarkLocalVars()
 //    In fast-jitting modes where we don't ref count locals, this bypasses
 //    actual counting, and makes all locals implicitly referenced on first
 //    compute. It asserts all locals are implicitly referenced on recompute.
-
+//
+//    When optimizing we also recompute lvaGenericsContextInUse based
+//    on specially flagged LCL_VAR appearances.
+//
 void Compiler::lvaComputeRefCounts(bool isRecompute, bool setSlotNumbers)
 {
     JITDUMP("\n*** lvaComputeRefCounts ***\n");
@@ -4141,6 +4159,11 @@ void Compiler::lvaComputeRefCounts(bool isRecompute, bool setSlotNumbers)
         varDsc->lvSingleDef = varDsc->lvIsParam;
     }
 
+    // Remember current state of generic context use, and prepare
+    // to compute new state.
+    const bool oldLvaGenericsContextInUse = lvaGenericsContextInUse;
+    lvaGenericsContextInUse               = false;
+
     JITDUMP("\n*** lvaComputeRefCounts -- explicit counts ***\n");
 
     // Second, account for all explicit local variable references
@@ -4176,6 +4199,12 @@ void Compiler::lvaComputeRefCounts(bool isRecompute, bool setSlotNumbers)
                         {
                             varDsc->incRefCnts(weight, this);
                         }
+
+                        if ((node->gtFlags & GTF_VAR_CONTEXT) != 0)
+                        {
+                            assert(node->OperIs(GT_LCL_VAR));
+                            lvaGenericsContextInUse = true;
+                        }
                         break;
                     }
 
@@ -4190,6 +4219,21 @@ void Compiler::lvaComputeRefCounts(bool isRecompute, bool setSlotNumbers)
         }
     }
 
+    if (oldLvaGenericsContextInUse && !lvaGenericsContextInUse)
+    {
+        // Context was in use but no longer is. This can happen
+        // if we're able to optimize, so just leave a note.
+        JITDUMP("\n** Generics context no longer in use\n");
+    }
+    else if (lvaGenericsContextInUse && !oldLvaGenericsContextInUse)
+    {
+        // Context was not in use but now is.
+        //
+        // Changing from unused->used should never happen; creation of any new IR
+        // for context use should also be settting lvaGenericsContextInUse.
+        assert(!"unexpected new use of generics context");
+    }
+
     JITDUMP("\n*** lvaComputeRefCounts -- implicit counts ***\n");
 
     // Third, bump ref counts for some implicit prolog references

src/coreclr/src/jit/morph.cpp

@@ -16048,7 +16048,7 @@ GenTree* Compiler::fgInitThisClass()
         // Collectible types requires that for shared generic code, if we use the generic context paramter
         // that we report it. (This is a conservative approach, we could detect some cases particularly when the
         // context parameter is this that we don't need the eager reporting logic.)
-        lvaGenericsContextUseCount++;
+        lvaGenericsContextInUse = true;
 
         switch (kind.runtimeLookupKind)
         {
@@ -16057,6 +16057,7 @@ GenTree* Compiler::fgInitThisClass()
                 // the hierarchy
                 {
                     GenTree* vtTree = gtNewLclvNode(info.compThisArg, TYP_REF);
+                    vtTree->gtFlags |= GTF_VAR_CONTEXT;
                     // Vtable pointer of this object
                     vtTree = gtNewOperNode(GT_IND, TYP_I_IMPL, vtTree);
                     vtTree->gtFlags |= GTF_EXCEPT; // Null-pointer exception
@@ -16068,12 +16069,14 @@ GenTree* Compiler::fgInitThisClass()
             case CORINFO_LOOKUP_CLASSPARAM:
             {
                 GenTree* vtTree = gtNewLclvNode(info.compTypeCtxtArg, TYP_I_IMPL);
+                vtTree->gtFlags |= GTF_VAR_CONTEXT;
                 return gtNewHelperCallNode(CORINFO_HELP_INITCLASS, TYP_VOID, gtNewCallArgs(vtTree));
             }
 
             case CORINFO_LOOKUP_METHODPARAM:
             {
                 GenTree* methHndTree = gtNewLclvNode(info.compTypeCtxtArg, TYP_I_IMPL);
+                methHndTree->gtFlags |= GTF_VAR_CONTEXT;
                 return gtNewHelperCallNode(CORINFO_HELP_INITINSTCLASS, TYP_VOID,
                                            gtNewCallArgs(gtNewIconNode(0), methHndTree));
             }