From 699219ab00389ba9b1100be37a1e55a637467ad7 Mon Sep 17 00:00:00 2001
From: vsadov <8218165+VSadov@users.noreply.github.com>
Date: Sat, 27 Jul 2024 15:50:10 -0700
Subject: [PATCH 1/5] Partial re-revert of #104336. Only JIT fixes are
 included.

---
 src/coreclr/jit/codegencommon.cpp        | 23 -----------------------
 src/coreclr/jit/emit.cpp                 |  6 +++---
 src/coreclr/jit/emitinl.h                |  3 +--
 src/coreclr/jit/gcencode.cpp             |  7 +++----
 src/coreclr/nativeaot/Runtime/thread.cpp |  6 ++++++
 5 files changed, 13 insertions(+), 32 deletions(-)

diff --git a/src/coreclr/jit/codegencommon.cpp b/src/coreclr/jit/codegencommon.cpp
index 3c590feb06edb..70dafbaad1fb8 100644
--- a/src/coreclr/jit/codegencommon.cpp
+++ b/src/coreclr/jit/codegencommon.cpp
@@ -1517,29 +1517,6 @@ void CodeGen::genExitCode(BasicBlock* block)
     if (compiler->getNeedsGSSecurityCookie())
     {
         genEmitGSCookieCheck(jmpEpilog);
-
-        if (jmpEpilog)
-        {
-            // Dev10 642944 -
-            // The GS cookie check created a temp label that has no live
-            // incoming GC registers, we need to fix that
-
-            unsigned   varNum;
-            LclVarDsc* varDsc;
-
-            /* Figure out which register parameters hold pointers */
-
-            for (varNum = 0, varDsc = compiler->lvaTable; varNum < compiler->lvaCount && varDsc->lvIsRegArg;
-                 varNum++, varDsc++)
-            {
-                noway_assert(varDsc->lvIsParam);
-
-                gcInfo.gcMarkRegPtrVal(varDsc->GetArgReg(), varDsc->TypeGet());
-            }
-
-            GetEmitter()->emitThisGCrefRegs = GetEmitter()->emitInitGCrefRegs = gcInfo.gcRegGCrefSetCur;
-            GetEmitter()->emitThisByrefRegs = GetEmitter()->emitInitByrefRegs = gcInfo.gcRegByrefSetCur;
-        }
     }
 
     genReserveEpilog(block);
diff --git a/src/coreclr/jit/emit.cpp b/src/coreclr/jit/emit.cpp
index d91b822afa11a..c82d2df22b9e1 100644
--- a/src/coreclr/jit/emit.cpp
+++ b/src/coreclr/jit/emit.cpp
@@ -10427,9 +10427,9 @@ regMaskTP emitter::emitGetGCRegsKilledByNoGCCall(CorInfoHelpFunc helper)
 //    of the last instruction in the region makes GC safe again.
 //    In particular - once the IP is on the first instruction, but not executed it yet,
 //    it is still safe to do GC.
-//    The only special case is when NoGC region is used for prologs/epilogs.
-//    In such case the GC info could be incorrect until the prolog completes and epilogs
-//    may have unwindability restrictions, so the first instruction cannot have GC.
+//    The only special case is when NoGC region is used for prologs.
+//    In such case the GC info could be incorrect until the prolog completes, so the first
+//    instruction cannot have GC.
 
 void emitter::emitDisableGC()
 {
diff --git a/src/coreclr/jit/emitinl.h b/src/coreclr/jit/emitinl.h
index 022064073d908..a586193dd5b71 100644
--- a/src/coreclr/jit/emitinl.h
+++ b/src/coreclr/jit/emitinl.h
@@ -594,8 +594,7 @@ bool emitter::emitGenNoGCLst(Callback& cb)
             emitter::instrDesc* id = emitFirstInstrDesc(ig->igData);
             assert(id != nullptr);
             assert(id->idCodeSize() > 0);
-            if (!cb(ig->igFuncIdx, ig->igOffs, ig->igSize, id->idCodeSize(),
-                    ig->igFlags & (IGF_FUNCLET_PROLOG | IGF_FUNCLET_EPILOG | IGF_EPILOG)))
+            if (!cb(ig->igFuncIdx, ig->igOffs, ig->igSize, id->idCodeSize(), ig->igFlags & (IGF_FUNCLET_PROLOG)))
             {
                 return false;
             }
diff --git a/src/coreclr/jit/gcencode.cpp b/src/coreclr/jit/gcencode.cpp
index 4863c95a7f59d..ae05bf797f86a 100644
--- a/src/coreclr/jit/gcencode.cpp
+++ b/src/coreclr/jit/gcencode.cpp
@@ -4027,8 +4027,7 @@ class InterruptibleRangeReporter
     // Report everything between the previous region and the current
     // region as interruptible.
 
-    bool operator()(
-        unsigned igFuncIdx, unsigned igOffs, unsigned igSize, unsigned firstInstrSize, bool isInPrologOrEpilog)
+    bool operator()(unsigned igFuncIdx, unsigned igOffs, unsigned igSize, unsigned firstInstrSize, bool isInProlog)
     {
         if (igOffs < m_uninterruptibleEnd)
         {
@@ -4042,9 +4041,9 @@ class InterruptibleRangeReporter
         if (igOffs > m_uninterruptibleEnd)
         {
             // Once the first instruction in IG executes, we cannot have GC.
-            // But it is ok to have GC while the IP is on the first instruction, unless we are in prolog/epilog.
+            // But it is ok to have GC while the IP is on the first instruction, unless we are in prolog.
             unsigned interruptibleEnd = igOffs;
-            if (!isInPrologOrEpilog)
+            if (!isInProlog)
             {
                 interruptibleEnd += firstInstrSize;
             }
diff --git a/src/coreclr/nativeaot/Runtime/thread.cpp b/src/coreclr/nativeaot/Runtime/thread.cpp
index b796b05218226..c2e94a1dc8f38 100644
--- a/src/coreclr/nativeaot/Runtime/thread.cpp
+++ b/src/coreclr/nativeaot/Runtime/thread.cpp
@@ -675,10 +675,16 @@ void Thread::HijackCallback(NATIVE_CONTEXT* pThreadContext, void* pThreadToHijac
     if (runtime->IsConservativeStackReportingEnabled() ||
         codeManager->IsSafePoint(pvAddress))
     {
+        // IsUnwindable is precise on arm64, but can give false negatives on other architectures.
+        // (when IP is on the first instruction of an epilog, we still can unwind,
+        // but we can tell if the instruction is the first only if we can navigate instructions backwards and check)
+        // The preciseness of IsUnwindable is tracked in https://github.com/dotnet/runtime/issues/101932
+#if defined(TARGET_ARM64)
         // we may not be able to unwind in some locations, such as epilogs.
         // such locations should not contain safe points.
         // when scanning conservatively we do not need to unwind
         ASSERT(codeManager->IsUnwindable(pvAddress) || runtime->IsConservativeStackReportingEnabled());
+#endif
 
         // if we are not given a thread to hijack
         // perform in-line wait on the current thread

From 8a5c678f807a7b789d6b285c7d615992d9ffe5e2 Mon Sep 17 00:00:00 2001
From: vsadov <8218165+VSadov@users.noreply.github.com>
Date: Tue, 30 Jul 2024 18:38:15 -0700
Subject: [PATCH 2/5] fix for stress issues

---
 src/coreclr/jit/codegenarmarch.cpp     |  5 +---
 src/coreclr/jit/codegencommon.cpp      | 36 ++++++--------------------
 src/coreclr/jit/codegenloongarch64.cpp |  7 +----
 src/coreclr/jit/codegenriscv64.cpp     |  7 +----
 src/coreclr/jit/codegenxarch.cpp       | 20 +++++++-------
 src/coreclr/jit/emit.cpp               |  5 ++++
 src/coreclr/jit/emit.h                 |  1 +
 7 files changed, 27 insertions(+), 54 deletions(-)

diff --git a/src/coreclr/jit/codegenarmarch.cpp b/src/coreclr/jit/codegenarmarch.cpp
index b0d5d5727fc50..5617c6dfc9135 100644
--- a/src/coreclr/jit/codegenarmarch.cpp
+++ b/src/coreclr/jit/codegenarmarch.cpp
@@ -605,10 +605,7 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
 {
     noway_assert(compiler->gsGlobalSecurityCookieAddr || compiler->gsGlobalSecurityCookieVal);
 
-    // Make sure that the return register is reported as live GC-ref so that any GC that kicks in while
-    // executing GS cookie check will not collect the object pointed to by REG_INTRET (R0).
-    if (!pushReg && (compiler->info.compRetNativeType == TYP_REF))
-        gcInfo.gcRegGCrefSetCur |= RBM_INTRET;
+    assert(GetEmitter()->emitGCDisabled());
 
     // We need two temporary registers, to load the GS cookie values and compare them. We can't use
     // any argument registers if 'pushReg' is true (meaning we have a JMP call). They should be
diff --git a/src/coreclr/jit/codegencommon.cpp b/src/coreclr/jit/codegencommon.cpp
index 70dafbaad1fb8..4b8cb100241f9 100644
--- a/src/coreclr/jit/codegencommon.cpp
+++ b/src/coreclr/jit/codegencommon.cpp
@@ -1516,6 +1516,14 @@ void CodeGen::genExitCode(BasicBlock* block)
     bool jmpEpilog = block->HasFlag(BBF_HAS_JMP);
     if (compiler->getNeedsGSSecurityCookie())
     {
+#ifndef JIT32_GCENCODER
+        // The call to CORINFO_HELP_FAIL_FAST in GS cookie check makes an impression
+        // of trashing all calee-trashed registers.
+        // Instead of ensuring that these registers are live across the check
+        // we just claim that the whole thing is not GC-interruptible.
+        // Effectively this starts the epilog a few instructions earlier.
+        GetEmitter()->emitDisableGC();
+#endif
         genEmitGSCookieCheck(jmpEpilog);
     }
 
@@ -4708,34 +4716,6 @@ void CodeGen::genReserveEpilog(BasicBlock* block)
     regMaskTP gcrefRegsArg = gcInfo.gcRegGCrefSetCur;
     regMaskTP byrefRegsArg = gcInfo.gcRegByrefSetCur;
 
-    /* The return value is special-cased: make sure it goes live for the epilog */
-
-    bool jmpEpilog = block->HasFlag(BBF_HAS_JMP);
-
-    if (IsFullPtrRegMapRequired() && !jmpEpilog)
-    {
-        if (varTypeIsGC(compiler->info.compRetNativeType))
-        {
-            noway_assert(genTypeStSz(compiler->info.compRetNativeType) == genTypeStSz(TYP_I_IMPL));
-
-            gcInfo.gcMarkRegPtrVal(REG_INTRET, compiler->info.compRetNativeType);
-
-            switch (compiler->info.compRetNativeType)
-            {
-                case TYP_REF:
-                    gcrefRegsArg |= RBM_INTRET;
-                    break;
-                case TYP_BYREF:
-                    byrefRegsArg |= RBM_INTRET;
-                    break;
-                default:
-                    break;
-            }
-
-            JITDUMP("Extending return value GC liveness to epilog\n");
-        }
-    }
-
     JITDUMP("Reserving epilog IG for block " FMT_BB "\n", block->bbNum);
 
     assert(block != nullptr);
diff --git a/src/coreclr/jit/codegenloongarch64.cpp b/src/coreclr/jit/codegenloongarch64.cpp
index 91623af10b7cc..5815d5720422d 100644
--- a/src/coreclr/jit/codegenloongarch64.cpp
+++ b/src/coreclr/jit/codegenloongarch64.cpp
@@ -4610,12 +4610,7 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
 {
     noway_assert(compiler->gsGlobalSecurityCookieAddr || compiler->gsGlobalSecurityCookieVal);
 
-    // Make sure that the return register is reported as live GC-ref so that any GC that kicks in while
-    // executing GS cookie check will not collect the object pointed to by REG_INTRET (A0).
-    if (!pushReg && (compiler->info.compRetNativeType == TYP_REF))
-    {
-        gcInfo.gcRegGCrefSetCur |= RBM_INTRET;
-    }
+    assert(GetEmitter()->emitGCDisabled());
 
     // We need two temporary registers, to load the GS cookie values and compare them. We can't use
     // any argument registers if 'pushReg' is true (meaning we have a JMP call). They should be
diff --git a/src/coreclr/jit/codegenriscv64.cpp b/src/coreclr/jit/codegenriscv64.cpp
index 7b06287c742c9..2ed15edb0e3fc 100644
--- a/src/coreclr/jit/codegenriscv64.cpp
+++ b/src/coreclr/jit/codegenriscv64.cpp
@@ -4664,12 +4664,7 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
 {
     noway_assert(compiler->gsGlobalSecurityCookieAddr || compiler->gsGlobalSecurityCookieVal);
 
-    // Make sure that the return register is reported as live GC-ref so that any GC that kicks in while
-    // executing GS cookie check will not collect the object pointed to by REG_INTRET (A0).
-    if (!pushReg && (compiler->info.compRetNativeType == TYP_REF))
-    {
-        gcInfo.gcRegGCrefSetCur |= RBM_INTRET;
-    }
+    assert(GetEmitter()->emitGCDisabled());
 
     // We need two temporary registers, to load the GS cookie values and compare them. We can't use
     // any argument registers if 'pushReg' is true (meaning we have a JMP call). They should be
diff --git a/src/coreclr/jit/codegenxarch.cpp b/src/coreclr/jit/codegenxarch.cpp
index 8504f2fb3ec52..4e66c4d32d3ba 100644
--- a/src/coreclr/jit/codegenxarch.cpp
+++ b/src/coreclr/jit/codegenxarch.cpp
@@ -96,12 +96,9 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
 {
     noway_assert(compiler->gsGlobalSecurityCookieAddr || compiler->gsGlobalSecurityCookieVal);
 
+#ifdef JIT32_GCENCODER
     // Make sure that EAX is reported as live GC-ref so that any GC that kicks in while
     // executing GS cookie check will not collect the object pointed to by EAX.
-    //
-    // For Amd64 System V, a two-register-returned struct could be returned in RAX and RDX
-    // In such case make sure that the correct GC-ness of RDX is reported as well, so
-    // a GC object pointed by RDX will not be collected.
     if (!pushReg)
     {
         if (compiler->compMethodReturnsRetBufAddr())
@@ -116,16 +113,19 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
         }
         else
         {
-            ReturnTypeDesc retTypeDesc = compiler->compRetTypeDesc;
-            const unsigned regCount    = retTypeDesc.GetReturnRegCount();
-
-            for (unsigned i = 0; i < regCount; ++i)
+            if (compiler->info.compRetNativeType == TYP_REF)
             {
-                gcInfo.gcMarkRegPtrVal(retTypeDesc.GetABIReturnReg(i, compiler->info.compCallConv),
-                                       retTypeDesc.GetReturnRegType(i));
+                gcInfo.gcRegGCrefSetCur |= RBM_INTRET;
+            }
+            else if (compiler->info.compRetNativeType == TYP_BYREF)
+            {
+                gcInfo.gcRegByrefSetCur |= RBM_INTRET;
             }
         }
     }
+#else
+    assert(GetEmitter()->emitGCDisabled());
+#endif
 
     regNumber regGSCheck;
     regMaskTP regMaskGSCheck = RBM_NONE;
diff --git a/src/coreclr/jit/emit.cpp b/src/coreclr/jit/emit.cpp
index c82d2df22b9e1..6331adf3b420c 100644
--- a/src/coreclr/jit/emit.cpp
+++ b/src/coreclr/jit/emit.cpp
@@ -10459,6 +10459,11 @@ void emitter::emitDisableGC()
     }
 }
 
+bool emitter::emitGCDisabled()
+{
+    return emitNoGCIG == true;
+}
+
 //------------------------------------------------------------------------
 // emitEnableGC(): Removes a request that the following instruction groups are not GC-interruptible.
 //
diff --git a/src/coreclr/jit/emit.h b/src/coreclr/jit/emit.h
index 83c08395118ef..b227fb182d58b 100644
--- a/src/coreclr/jit/emit.h
+++ b/src/coreclr/jit/emit.h
@@ -2745,6 +2745,7 @@ class emitter
 #if !defined(JIT32_GCENCODER)
     void emitDisableGC();
     void emitEnableGC();
+    bool emitGCDisabled();
 #endif // !defined(JIT32_GCENCODER)
 
 #if defined(TARGET_XARCH)

From 0a96d6650bfeecc3d79d71d85a3c80cd87fd122e Mon Sep 17 00:00:00 2001
From: vsadov <8218165+VSadov@users.noreply.github.com>
Date: Wed, 31 Jul 2024 16:09:35 -0700
Subject: [PATCH 3/5] comments and some cleanup

---
 src/coreclr/jit/codegencommon.cpp | 21 ++++++++++++---------
 src/coreclr/jit/codegenxarch.cpp  | 26 ++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/src/coreclr/jit/codegencommon.cpp b/src/coreclr/jit/codegencommon.cpp
index 4b8cb100241f9..120620e14e4cf 100644
--- a/src/coreclr/jit/codegencommon.cpp
+++ b/src/coreclr/jit/codegencommon.cpp
@@ -1517,11 +1517,15 @@ void CodeGen::genExitCode(BasicBlock* block)
     if (compiler->getNeedsGSSecurityCookie())
     {
 #ifndef JIT32_GCENCODER
-        // The call to CORINFO_HELP_FAIL_FAST in GS cookie check makes an impression
-        // of trashing all calee-trashed registers.
-        // Instead of ensuring that these registers are live across the check
-        // we just claim that the whole thing is not GC-interruptible.
+        // At this point the gc info that we track in codegen is often incorrect,
+        // as it could be missing return registers or arg registers (in a case of tail call).
+        // GS cookie check will emit a call and that will pass our GC info to emit and potentially mess things up.
+        // While we could infer returns/args and force them to be live and it seems to work in JIT32_GCENCODER case,
+        // it appears to be nontrivial in more general case.
+        // So, instead, we just claim that the whole thing is not GC-interruptible.
         // Effectively this starts the epilog a few instructions earlier.
+        //
+        // CONSIDER: is that a good place to be that codegen loses track of returns/args at this point?
         GetEmitter()->emitDisableGC();
 #endif
         genEmitGSCookieCheck(jmpEpilog);
@@ -4713,14 +4717,13 @@ void CodeGen::genReserveProlog(BasicBlock* block)
 
 void CodeGen::genReserveEpilog(BasicBlock* block)
 {
-    regMaskTP gcrefRegsArg = gcInfo.gcRegGCrefSetCur;
-    regMaskTP byrefRegsArg = gcInfo.gcRegByrefSetCur;
-
     JITDUMP("Reserving epilog IG for block " FMT_BB "\n", block->bbNum);
 
     assert(block != nullptr);
-    const VARSET_TP& gcrefVarsArg(GetEmitter()->emitThisGCrefVars);
-    GetEmitter()->emitCreatePlaceholderIG(IGPT_EPILOG, block, gcrefVarsArg, gcrefRegsArg, byrefRegsArg,
+    // We pass empty GC info, because epilog is always an extend IG and will ignore what we pass.
+    // Besides, at this point the GC info that we track in CodeGen is often incorrect.
+    // See comments in genExitCode for more info.
+    GetEmitter()->emitCreatePlaceholderIG(IGPT_EPILOG, block, VarSetOps::MakeEmpty(compiler), 0, 0,
                                           block->IsLast());
 }
 
diff --git a/src/coreclr/jit/codegenxarch.cpp b/src/coreclr/jit/codegenxarch.cpp
index 4e66c4d32d3ba..f6878beaf0974 100644
--- a/src/coreclr/jit/codegenxarch.cpp
+++ b/src/coreclr/jit/codegenxarch.cpp
@@ -97,10 +97,10 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
     noway_assert(compiler->gsGlobalSecurityCookieAddr || compiler->gsGlobalSecurityCookieVal);
 
 #ifdef JIT32_GCENCODER
-    // Make sure that EAX is reported as live GC-ref so that any GC that kicks in while
-    // executing GS cookie check will not collect the object pointed to by EAX.
     if (!pushReg)
     {
+        // Make sure that EAX is reported as live GC-ref so that any GC that kicks in while
+        // executing GS cookie check will not collect the object pointed to by EAX.
         if (compiler->compMethodReturnsRetBufAddr())
         {
             // This is for returning in an implicit RetBuf.
@@ -123,6 +123,28 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
             }
         }
     }
+    else
+    {
+        // Dev10 642944 -
+        // The GS cookie check created a temp label that has no live
+        // incoming GC registers, we need to fix that
+
+        unsigned   varNum;
+        LclVarDsc* varDsc;
+
+        /* Figure out which register parameters hold pointers */
+
+        for (varNum = 0, varDsc = compiler->lvaTable; varNum < compiler->lvaCount && varDsc->lvIsRegArg;
+             varNum++, varDsc++)
+        {
+            noway_assert(varDsc->lvIsParam);
+
+            gcInfo.gcMarkRegPtrVal(varDsc->GetArgReg(), varDsc->TypeGet());
+        }
+
+        GetEmitter()->emitThisGCrefRegs = GetEmitter()->emitInitGCrefRegs = gcInfo.gcRegGCrefSetCur;
+        GetEmitter()->emitThisByrefRegs = GetEmitter()->emitInitByrefRegs = gcInfo.gcRegByrefSetCur;
+    }
 #else
     assert(GetEmitter()->emitGCDisabled());
 #endif

From 7529e47a0468b8ac24f19769d2ed3b35ff83a082 Mon Sep 17 00:00:00 2001
From: vsadov <8218165+VSadov@users.noreply.github.com>
Date: Wed, 31 Jul 2024 17:06:46 -0700
Subject: [PATCH 4/5] JIT format

---
 src/coreclr/jit/codegencommon.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/coreclr/jit/codegencommon.cpp b/src/coreclr/jit/codegencommon.cpp
index 120620e14e4cf..ae233e5583d6f 100644
--- a/src/coreclr/jit/codegencommon.cpp
+++ b/src/coreclr/jit/codegencommon.cpp
@@ -4723,8 +4723,7 @@ void CodeGen::genReserveEpilog(BasicBlock* block)
     // We pass empty GC info, because epilog is always an extend IG and will ignore what we pass.
     // Besides, at this point the GC info that we track in CodeGen is often incorrect.
     // See comments in genExitCode for more info.
-    GetEmitter()->emitCreatePlaceholderIG(IGPT_EPILOG, block, VarSetOps::MakeEmpty(compiler), 0, 0,
-                                          block->IsLast());
+    GetEmitter()->emitCreatePlaceholderIG(IGPT_EPILOG, block, VarSetOps::MakeEmpty(compiler), 0, 0, block->IsLast());
 }
 
 /*****************************************************************************

From 5e06290666a79b199b9b04215d655c51f43fec35 Mon Sep 17 00:00:00 2001
From: vsadov <8218165+VSadov@users.noreply.github.com>
Date: Fri, 2 Aug 2024 17:25:05 -0700
Subject: [PATCH 5/5] keep existing code for x86

---
 src/coreclr/jit/codegencommon.cpp | 26 +++++++++++++++++++++++
 src/coreclr/jit/codegenxarch.cpp  | 34 ++++++-------------------------
 2 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/src/coreclr/jit/codegencommon.cpp b/src/coreclr/jit/codegencommon.cpp
index ae233e5583d6f..0598d968e7388 100644
--- a/src/coreclr/jit/codegencommon.cpp
+++ b/src/coreclr/jit/codegencommon.cpp
@@ -1528,7 +1528,33 @@ void CodeGen::genExitCode(BasicBlock* block)
         // CONSIDER: is that a good place to be that codegen loses track of returns/args at this point?
         GetEmitter()->emitDisableGC();
 #endif
+
         genEmitGSCookieCheck(jmpEpilog);
+
+#ifdef JIT32_GCENCODER
+        if (jmpEpilog)
+        {
+            // Dev10 642944 -
+            // The GS cookie check created a temp label that has no live
+            // incoming GC registers, we need to fix that
+
+            unsigned   varNum;
+            LclVarDsc* varDsc;
+
+            /* Figure out which register parameters hold pointers */
+
+            for (varNum = 0, varDsc = compiler->lvaTable; varNum < compiler->lvaCount && varDsc->lvIsRegArg;
+                 varNum++, varDsc++)
+            {
+                noway_assert(varDsc->lvIsParam);
+
+                gcInfo.gcMarkRegPtrVal(varDsc->GetArgReg(), varDsc->TypeGet());
+            }
+
+            GetEmitter()->emitThisGCrefRegs = GetEmitter()->emitInitGCrefRegs = gcInfo.gcRegGCrefSetCur;
+            GetEmitter()->emitThisByrefRegs = GetEmitter()->emitInitByrefRegs = gcInfo.gcRegByrefSetCur;
+        }
+#endif
     }
 
     genReserveEpilog(block);
diff --git a/src/coreclr/jit/codegenxarch.cpp b/src/coreclr/jit/codegenxarch.cpp
index f6878beaf0974..4102d4c771115 100644
--- a/src/coreclr/jit/codegenxarch.cpp
+++ b/src/coreclr/jit/codegenxarch.cpp
@@ -113,38 +113,16 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
         }
         else
         {
-            if (compiler->info.compRetNativeType == TYP_REF)
-            {
-                gcInfo.gcRegGCrefSetCur |= RBM_INTRET;
-            }
-            else if (compiler->info.compRetNativeType == TYP_BYREF)
+            ReturnTypeDesc retTypeDesc = compiler->compRetTypeDesc;
+            const unsigned regCount    = retTypeDesc.GetReturnRegCount();
+
+            for (unsigned i = 0; i < regCount; ++i)
             {
-                gcInfo.gcRegByrefSetCur |= RBM_INTRET;
+                gcInfo.gcMarkRegPtrVal(retTypeDesc.GetABIReturnReg(i, compiler->info.compCallConv),
+                                       retTypeDesc.GetReturnRegType(i));
             }
         }
     }
-    else
-    {
-        // Dev10 642944 -
-        // The GS cookie check created a temp label that has no live
-        // incoming GC registers, we need to fix that
-
-        unsigned   varNum;
-        LclVarDsc* varDsc;
-
-        /* Figure out which register parameters hold pointers */
-
-        for (varNum = 0, varDsc = compiler->lvaTable; varNum < compiler->lvaCount && varDsc->lvIsRegArg;
-             varNum++, varDsc++)
-        {
-            noway_assert(varDsc->lvIsParam);
-
-            gcInfo.gcMarkRegPtrVal(varDsc->GetArgReg(), varDsc->TypeGet());
-        }
-
-        GetEmitter()->emitThisGCrefRegs = GetEmitter()->emitInitGCrefRegs = gcInfo.gcRegGCrefSetCur;
-        GetEmitter()->emitThisByrefRegs = GetEmitter()->emitInitByrefRegs = gcInfo.gcRegByrefSetCur;
-    }
 #else
     assert(GetEmitter()->emitGCDisabled());
 #endif