JIT: change how we block gc refs from callee saves for inline pinvokes (#22477)

Add a new marker instruction that we emit once we've enabled preepmtive gc in
the inline pinvoke method prolog. Use that to kill off callee saves registers
with GC references, instead of waiting until the call.

This closes a window of vulnerability we see in GC stress where if a stress
interrupt happens between the point at which we enable preeemptive GC and
the point at which we make the call, we may report callee saves as GC live
when they're actually dead.

Closes #19211.

13 files changed, 80 insertions, 13 deletions

diff --git a/src/jit/codegenarmarch.cpp b/src/jit/codegenarmarch.cpp
index 0682f609e6..e8f38b7b81 100644
--- a/src/jit/codegenarmarch.cpp
+++ b/src/jit/codegenarmarch.cpp
@@ -71,6 +71,13 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
             getEmitter()->emitDisableGC();
             break;
 
+        case GT_START_PREEMPTGC:
+            // Kill callee saves GC registers, and create a label
+            // so that information gets propagated to the emitter.
+            gcInfo.gcMarkRegSetNpt(RBM_INT_CALLEE_SAVED);
+            genDefineTempLabel(genCreateTempLabel());
+            break;
+
         case GT_PROF_HOOK:
             // We should be seeing this only if profiler hook is needed
             noway_assert(compiler->compIsProfilerHookNeeded());
diff --git a/src/jit/codegenxarch.cpp b/src/jit/codegenxarch.cpp
index 2327988079..f80feeb088 100644
--- a/src/jit/codegenxarch.cpp
+++ b/src/jit/codegenxarch.cpp
@@ -1560,6 +1560,13 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
             break;
 #endif // !defined(JIT32_GCENCODER)
 
+        case GT_START_PREEMPTGC:
+            // Kill callee saves GC registers, and create a label
+            // so that information gets propagated to the emitter.
+            gcInfo.gcMarkRegSetNpt(RBM_INT_CALLEE_SAVED);
+            genDefineTempLabel(genCreateTempLabel());
+            break;
+
         case GT_PROF_HOOK:
 #ifdef PROFILING_SUPPORTED
             // We should be seeing this only if profiler hook is needed
diff --git a/src/jit/compiler.cpp b/src/jit/compiler.cpp
index 6123fb126a..691f35b800 100644
--- a/src/jit/compiler.cpp
+++ b/src/jit/compiler.cpp
@@ -10156,6 +10156,7 @@ int cLeafIR(Compiler* comp, GenTree* tree)
 
         case GT_NO_OP:
         case GT_START_NONGC:
+        case GT_START_PREEMPTGC:
         case GT_PROF_HOOK:
         case GT_CATCH_ARG:
         case GT_MEMORYBARRIER:
@@ -11083,5 +11084,10 @@ bool Compiler::killGCRefs(GenTree* tree)
             return true;
         }
     }
+    else if (tree->OperIs(GT_START_PREEMPTGC))
+    {
+        return true;
+    }
+
     return false;
 }
diff --git a/src/jit/compiler.h b/src/jit/compiler.h
index 2940bc9d77..48ce14109d 100644
--- a/src/jit/compiler.h
+++ b/src/jit/compiler.h
@@ -9822,6 +9822,7 @@ public:
             case GT_SETCC:
             case GT_NO_OP:
             case GT_START_NONGC:
+            case GT_START_PREEMPTGC:
             case GT_PROF_HOOK:
 #if !FEATURE_EH_FUNCLETS
             case GT_END_LFIN:
diff --git a/src/jit/compiler.hpp b/src/jit/compiler.hpp
index 56f4248ee4..a9dc3bbc5f 100644
--- a/src/jit/compiler.hpp
+++ b/src/jit/compiler.hpp
@@ -4196,6 +4196,7 @@ void GenTree::VisitOperands(TVisitor visitor)
         case GT_SETCC:
         case GT_NO_OP:
         case GT_START_NONGC:
+        case GT_START_PREEMPTGC:
         case GT_PROF_HOOK:
 #if !FEATURE_EH_FUNCLETS
         case GT_END_LFIN:
diff --git a/src/jit/gentree.cpp b/src/jit/gentree.cpp
index 724a6555c7..f197d249b6 100644
--- a/src/jit/gentree.cpp
+++ b/src/jit/gentree.cpp
@@ -4700,6 +4700,7 @@ bool GenTree::TryGetUse(GenTree* def, GenTree*** use)
         case GT_SETCC:
         case GT_NO_OP:
         case GT_START_NONGC:
+        case GT_START_PREEMPTGC:
         case GT_PROF_HOOK:
 #if !FEATURE_EH_FUNCLETS
         case GT_END_LFIN:
@@ -8400,6 +8401,7 @@ GenTreeUseEdgeIterator::GenTreeUseEdgeIterator(GenTree* node)
         case GT_SETCC:
         case GT_NO_OP:
         case GT_START_NONGC:
+        case GT_START_PREEMPTGC:
         case GT_PROF_HOOK:
 #if !FEATURE_EH_FUNCLETS
         case GT_END_LFIN:
@@ -10386,6 +10388,7 @@ void Compiler::gtDispLeaf(GenTree* tree, IndentStack* indentStack)
 
         case GT_NO_OP:
         case GT_START_NONGC:
+        case GT_START_PREEMPTGC:
         case GT_PROF_HOOK:
         case GT_CATCH_ARG:
         case GT_MEMORYBARRIER:
diff --git a/src/jit/gtlist.h b/src/jit/gtlist.h
index eb8b0a883a..dd23db486b 100644
--- a/src/jit/gtlist.h
+++ b/src/jit/gtlist.h
@@ -255,6 +255,8 @@ GTNODE(NO_OP            , GenTree            ,0,GTK_LEAF|GTK_NOVALUE)   // nop!
 
 GTNODE(START_NONGC      , GenTree            ,0,GTK_LEAF|GTK_NOVALUE)   // starts a new instruction group that will be non-gc interruptible
 
+GTNODE(START_PREEMPTGC  , GenTree            ,0,GTK_LEAF|GTK_NOVALUE)   // starts a new instruction group where preemptive GC is enabled
+
 GTNODE(PROF_HOOK        , GenTree            ,0,GTK_LEAF|GTK_NOVALUE)   // profiler Enter/Leave/TailCall hook
 
 GTNODE(RETFILT          , GenTreeOp          ,0,GTK_UNOP|GTK_NOVALUE)   // end filter with TYP_I_IMPL return value
diff --git a/src/jit/liveness.cpp b/src/jit/liveness.cpp
index 0f24d825d2..ded6c05ce1 100644
--- a/src/jit/liveness.cpp
+++ b/src/jit/liveness.cpp
@@ -1947,6 +1947,7 @@ void Compiler::fgComputeLifeLIR(VARSET_TP& life, BasicBlock* block, VARSET_VALAR
             case GT_SWITCH:
             case GT_RETFILT:
             case GT_START_NONGC:
+            case GT_START_PREEMPTGC:
             case GT_PROF_HOOK:
 #if !FEATURE_EH_FUNCLETS
             case GT_END_LFIN:
diff --git a/src/jit/lower.cpp b/src/jit/lower.cpp
index 55f44a71e0..2c4c071a35 100644
--- a/src/jit/lower.cpp
+++ b/src/jit/lower.cpp
@@ -3584,6 +3584,7 @@ void Lowering::InsertPInvokeCallProlog(GenTreeCall* call)
     // InlinedCallFrame.callTarget = methodHandle   // stored in m_Datum
     // InlinedCallFrame.m_pCallSiteSP = SP          // x86 only
     // InlinedCallFrame.m_pCallerReturnAddress = return address
+    // GT_START_PREEEMPTC
     // Thread.gcState = 0
     // (non-stub) - update top Frame on TCB         // 64-bit targets only
 
@@ -3688,14 +3689,19 @@ void Lowering::InsertPInvokeCallProlog(GenTreeCall* call)
     }
 #endif // _TARGET_64BIT_
 
-    // IMPORTANT **** This instruction must come last!!! ****
+    // IMPORTANT **** This instruction must be the last real instruction ****
     // It changes the thread's state to Preemptive mode
     // ----------------------------------------------------------------------------------
     //  [tcb + offsetOfGcState] = 0
-
     GenTree* storeGCState = SetGCState(0);
     BlockRange().InsertBefore(insertBefore, LIR::SeqTree(comp, storeGCState));
     ContainCheckStoreIndir(storeGCState->AsIndir());
+
+    // Indicate that codegen has switched this thread to preemptive GC.
+    // This tree node doesn't generate any code, but impacts LSRA and gc reporting.
+    // This tree node is simple so doesn't require sequencing.
+    GenTree* preemptiveGCNode = new (comp, GT_START_PREEMPTGC) GenTree(GT_START_PREEMPTGC, TYP_VOID);
+    BlockRange().InsertBefore(insertBefore, preemptiveGCNode);
 }
 
 //------------------------------------------------------------------------
diff --git a/src/jit/lsraarm.cpp b/src/jit/lsraarm.cpp
index f64b7fb3f0..27e5e99111 100644
--- a/src/jit/lsraarm.cpp
+++ b/src/jit/lsraarm.cpp
@@ -411,6 +411,13 @@ int LinearScan::BuildNode(GenTree* tree)
             assert(dstCount == 0);
             break;
 
+        case GT_START_PREEMPTGC:
+            // This kills GC refs in callee save regs
+            srcCount = 0;
+            assert(dstCount == 0);
+            BuildDefsWithKills(tree, 0, RBM_NONE, RBM_NONE);
+            break;
+
         case GT_LONG:
             assert(tree->IsUnusedValue()); // Contained nodes are already processed, only unused GT_LONG can reach here.
                                            // An unused GT_LONG doesn't produce any registers.
diff --git a/src/jit/lsraarm64.cpp b/src/jit/lsraarm64.cpp
index d5446362fe..13907a7b24 100644
--- a/src/jit/lsraarm64.cpp
+++ b/src/jit/lsraarm64.cpp
@@ -135,6 +135,13 @@ int LinearScan::BuildNode(GenTree* tree)
             assert(dstCount == 0);
             break;
 
+        case GT_START_PREEMPTGC:
+            // This kills GC refs in callee save regs
+            srcCount = 0;
+            assert(dstCount == 0);
+            BuildDefsWithKills(tree, 0, RBM_NONE, RBM_NONE);
+            break;
+
         case GT_CNS_DBL:
         {
             GenTreeDblCon* dblConst   = tree->AsDblCon();
diff --git a/src/jit/lsrabuild.cpp b/src/jit/lsrabuild.cpp
index 4f30182926..7e32e5c614 100644
--- a/src/jit/lsrabuild.cpp
+++ b/src/jit/lsrabuild.cpp
@@ -1113,6 +1113,7 @@ regMaskTP LinearScan::getKillSetForNode(GenTree* tree)
 // Arguments:
 //    tree       - the tree for which kill positions should be generated
 //    currentLoc - the location at which the kills should be added
+//    killMask   - The mask of registers killed by this node
 //
 // Return Value:
 //    true       - kills were inserted
@@ -1124,10 +1125,14 @@ regMaskTP LinearScan::getKillSetForNode(GenTree* tree)
 //    the multiple defs for a regPair are in different locations.
 //    If we generate any kills, we will mark all currentLiveVars as being preferenced
 //    to avoid the killed registers.  This is somewhat conservative.
+//
+//    This method can add kills even if killMask is RBM_NONE, if this tree is one of the
+//    special cases that signals that we can't permit callee save registers to hold GC refs.
 
 bool LinearScan::buildKillPositionsForNode(GenTree* tree, LsraLocation currentLoc, regMaskTP killMask)
 {
-    bool isCallKill = ((killMask == RBM_INT_CALLEE_TRASH) || (killMask == RBM_CALLEE_TRASH));
+    bool insertedKills = false;
+
     if (killMask != RBM_NONE)
     {
         // The killMask identifies a set of registers that will be used during codegen.
@@ -1143,7 +1148,7 @@ bool LinearScan::buildKillPositionsForNode(GenTree* tree, LsraLocation currentLo
         addRefsForPhysRegMask(killMask, currentLoc, RefTypeKill, true);
 
         // TODO-CQ: It appears to be valuable for both fp and int registers to avoid killing the callee
-        // save regs on infrequently exectued paths.  However, it results in a large number of asmDiffs,
+        // save regs on infrequently executed paths.  However, it results in a large number of asmDiffs,
         // many of which appear to be regressions (because there is more spill on the infrequently path),
         // but are not really because the frequent path becomes smaller.  Validating these diffs will need
         // to be done before making this change.
@@ -1171,7 +1176,9 @@ bool LinearScan::buildKillPositionsForNode(GenTree* tree, LsraLocation currentLo
                 {
                     continue;
                 }
-                Interval* interval = getIntervalForLocalVar(varIndex);
+                Interval*  interval   = getIntervalForLocalVar(varIndex);
+                const bool isCallKill = ((killMask == RBM_INT_CALLEE_TRASH) || (killMask == RBM_CALLEE_TRASH));
+
                 if (isCallKill)
                 {
                     interval->preferCalleeSave = true;
@@ -1192,15 +1199,17 @@ bool LinearScan::buildKillPositionsForNode(GenTree* tree, LsraLocation currentLo
             }
         }
 
-        if (compiler->killGCRefs(tree))
-        {
-            RefPosition* pos = newRefPosition((Interval*)nullptr, currentLoc, RefTypeKillGCRefs, tree,
-                                              (allRegs(TYP_REF) & ~RBM_ARG_REGS));
-        }
-        return true;
+        insertedKills = true;
     }
 
-    return false;
+    if (compiler->killGCRefs(tree))
+    {
+        RefPosition* pos =
+            newRefPosition((Interval*)nullptr, currentLoc, RefTypeKillGCRefs, tree, (allRegs(TYP_REF) & ~RBM_ARG_REGS));
+        insertedKills = true;
+    }
+
+    return insertedKills;
 }
 
 //----------------------------------------------------------------------------
@@ -2556,9 +2565,12 @@ void LinearScan::BuildDefsWithKills(GenTree* tree, int dstCount, regMaskTP dstCa
     VARSET_TP liveLargeVectors(VarSetOps::UninitVal());
     bool      doLargeVectorRestore = false;
 #endif // FEATURE_PARTIAL_SIMD_CALLEE_SAVE
+
+    // Call this even when killMask is RBM_NONE, as we have to check for some special cases
+    buildKillPositionsForNode(tree, currentLoc + 1, killMask);
+
     if (killMask != RBM_NONE)
     {
-        buildKillPositionsForNode(tree, currentLoc + 1, killMask);
 #if FEATURE_PARTIAL_SIMD_CALLEE_SAVE
         if (enregisterLocalVars && ((killMask & RBM_FLT_CALLEE_TRASH) != RBM_NONE))
         {
diff --git a/src/jit/lsraxarch.cpp b/src/jit/lsraxarch.cpp
index cc6419ae0d..67fe539a2f 100644
--- a/src/jit/lsraxarch.cpp
+++ b/src/jit/lsraxarch.cpp
@@ -156,6 +156,13 @@ int LinearScan::BuildNode(GenTree* tree)
             assert(dstCount == 0);
             break;
 
+        case GT_START_PREEMPTGC:
+            // This kills GC refs in callee save regs
+            srcCount = 0;
+            assert(dstCount == 0);
+            BuildDefsWithKills(tree, 0, RBM_NONE, RBM_NONE);
+            break;
+
         case GT_PROF_HOOK:
             srcCount = 0;
             assert(dstCount == 0);