JIT: revise how the jit tracks use of generics context

The generics context is reported by the jit specially whenever it feeds into
runtime lookups, as expansion of those lookups can expose pointers into runtime
data structures, and we don't want those data structures to be collected if
jitted code is still using them.

Sometimes uses of the context are optimized away, and reporting costs
code size and GC space, so we don't want to report the context unless there
is an actual use.

This change revises how the jit keeps track of context use -- instead of trying
to incrementally ref count uses of the generics context, we now just leverage
existing passes which do local accounting.

Initial motivation for this came from dotnet#34641 where the context use was
over-reported, but investigation showed we also some times under-report as
the context var could be cloned without changing the ref count.

So this change fixes both under and over reporting.

Closes dotnet#34641.

src/coreclr/src/jit/compiler.h

@@ -2852,9 +2852,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
@@ -3127,7 +3125,7 @@ class Compiler
 
 #endif // defined(DEBUG) && defined(TARGET_X86)
 
-    unsigned lvaGenericsContextUseCount;
+    unsigned lvaGenericsContextInUse;
 
     bool lvaKeepAliveAndReportThis(); // Synchronized instance method of a reference type, or
                                       // CORINFO_GENERICS_CTXT_FROM_THIS?

src/coreclr/src/jit/compiler.hpp

@@ -1931,9 +1931,8 @@ inline bool Compiler::lvaKeepAliveAndReportThis()
         if (opts.compDbgCode)
             return true;
 
-        if (lvaGenericsContextUseCount > 0)
+        if (lvaGenericsContextInUse)
         {
-            JITDUMP("Reporting this as generic context: %u refs\n", lvaGenericsContextUseCount);
             return true;
         }
     }
@@ -1944,14 +1943,12 @@ inline bool Compiler::lvaKeepAliveAndReportThis()
     // because collectible types need the generics context when gc-ing.
     if (genericsContextIsThis)
     {
-        const bool isUsed   = lvaGenericsContextUseCount > 0;
+        const bool isUsed   = lvaGenericsContextInUse;
         const bool mustKeep = (info.compMethodInfo->options & CORINFO_GENERICS_CTXT_KEEP_ALIVE) != 0;
 
         if (isUsed || 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;
         }
     }
@@ -1979,7 +1976,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

@@ -7849,7 +7849,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)
         {
@@ -7863,13 +7863,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));
@@ -23736,38 +23738,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

@@ -9892,6 +9892,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:
@@ -12421,11 +12428,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;
@@ -12444,11 +12450,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
@@ -12611,15 +12612,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;
 
@@ -12639,11 +12637,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

@@ -770,6 +770,7 @@ struct GenTree
 #define GTF_VAR_USEASG      0x40000000 // GT_LCL_VAR -- this is a partial definition, a use of the previous definition is implied
                                        // A partial definition usually occurs when a struct field is assigned to (s.f = ...) or
                                        // when a scalar typed variable is assigned to via a narrow store (*((byte*)&i) = ...).
+#define GTF_VAR_CONTEXT     0x20000000 // GT_LCL_VAR -- this is a generics context reference
 #define GTF_VAR_CAST        0x10000000 // GT_LCL_VAR -- has been explictly cast (variable node may not be type of local)
 #define GTF_VAR_ITERATOR    0x08000000 // GT_LCL_VAR -- this is a iterator reference in the loop condition
 #define GTF_VAR_CLONED      0x01000000 // GT_LCL_VAR -- this node has been cloned or is a clone

src/coreclr/src/jit/importer.cpp

@@ -1979,12 +1979,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);
@@ -1996,6 +1997,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;
 }
@@ -18595,7 +18597,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();
 
@@ -4008,24 +4021,24 @@ 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
-
+    // 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 = lvaReportParamTypeArg();
     }
     else if (lvaReportParamTypeArg())
     {
-        lvaTable[info.compTypeCtxtArg].lvImplicitlyReferenced = 1;
+        // We should have a context arg.
+        assert(info.compTypeCtxtArg != BAD_VAR_NUM);
+        lvaGetDesc(info.compTypeCtxtArg)->lvImplicitlyReferenced = lvaReportParamTypeArg();
     }
+
+#if ASSERTION_PROP
+    assert(opts.OptimizationEnabled());
+
+    // Note: optAddCopies() depends on lvaRefBlks, which is set in lvaMarkLocalVars(BasicBlock*), called above.
+    optAddCopies();
+#endif
 }
 
 //------------------------------------------------------------------------
@@ -4045,7 +4058,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");
@@ -4140,6 +4156,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
@@ -4175,6 +4196,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;
                     }
 
@@ -4189,6 +4216,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

@@ -16041,7 +16041,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)
         {
@@ -16050,6 +16050,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
@@ -16061,12 +16062,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));
             }