From 76a68f3739ea5b50191a91e997604ef8a3a06a79 Mon Sep 17 00:00:00 2001
From: Jakob Botsch Nielsen <jakob.botsch.nielsen@gmail.com>
Date: Tue, 12 May 2026 13:59:17 +0200
Subject: [PATCH] JIT: Eagerly invalidate GC state for registers trashed by
 linker relaxation sequences

TLS accesses are optimized in NativeAOT with the help of the linker. The
JIT emits a sequence of instructions that is specially recognized, and
the linker then rewrites (relaxes) those instructions to a more
efficient pattern.

Normally, when the JIT emits instructions, it only lazily updates GC
state. For example, if `rax` contains a GC reference, we delay any
actual GC information update until we actually emit an instruction that
clobber `rax`. This is regardless of whether or not the value in `rax`
is dying.

This is a problem for these linker relaxation sequences. The linker's
instructions may end up clobbering registers earlier than we do in our
emitted instructions, and in that case we are reporting a GC register as
live while it had been trashed.

Fix the case for linux-x64 and linux-arm64 by eagerly emitting GC state
updates for these special patterns.
---
 src/coreclr/jit/codegenarmarch.cpp |  8 ++++++++
 src/coreclr/jit/codegenxarch.cpp   |  9 +++++++++
 src/coreclr/jit/emitarm64.cpp      | 11 +++++++++++
 src/coreclr/jit/emitxarch.cpp      |  7 +++++++
 4 files changed, 35 insertions(+)

diff --git a/src/coreclr/jit/codegenarmarch.cpp b/src/coreclr/jit/codegenarmarch.cpp
index 6e177b0a7ad3d9..699e1b103fe1d3 100644
--- a/src/coreclr/jit/codegenarmarch.cpp
+++ b/src/coreclr/jit/codegenarmarch.cpp
@@ -3353,6 +3353,14 @@ void CodeGen::genCallInstruction(GenTreeCall* call)
             // mrs
             emitter->emitIns_R(INS_mrs_tpid0, attr, REG_R1);
 
+            // We remove x0 here since the linker relaxation
+            // sequence will rewrite the instructions we are emitting here with
+            // instructions that may clobber these registers.
+            // (This is more important for the emitter, but we match it here
+            // for symmetry and to avoid confusion about the state of the
+            // registers.)
+            gcInfo.gcMarkRegSetNpt(RBM_R0);
+
             // adrp
             // ldr
             // add
diff --git a/src/coreclr/jit/codegenxarch.cpp b/src/coreclr/jit/codegenxarch.cpp
index b8d5ffb8f9452c..a6b563e97aea01 100644
--- a/src/coreclr/jit/codegenxarch.cpp
+++ b/src/coreclr/jit/codegenxarch.cpp
@@ -484,6 +484,15 @@ void CodeGen::genSetRegToConst(regNumber targetReg, var_types targetType, GenTre
                 else if (con->IsIconHandle(GTF_ICON_TLSGD_OFFSET))
                 {
                     attr = EA_SET_FLG(attr, EA_CNS_TLSGD_RELOC);
+
+                    // This marks the start of a TLS access linker relaxation
+                    // sequence for linux-x64. The linker may rewrite to
+                    // instructions that trash rax at this point, so we update
+                    // GC state eagerly here.
+                    // (This is more important for the emitter, but we match it here
+                    // for symmetry and to avoid confusion about the state of the
+                    // registers.)
+                    gcInfo.gcMarkRegSetNpt(RBM_RAX);
                 }
             }
 
diff --git a/src/coreclr/jit/emitarm64.cpp b/src/coreclr/jit/emitarm64.cpp
index 3bf05d5936ecf4..01699a7d3d0a58 100644
--- a/src/coreclr/jit/emitarm64.cpp
+++ b/src/coreclr/jit/emitarm64.cpp
@@ -11725,6 +11725,17 @@ size_t emitter::emitOutputInstr(insGroup* ig, instrDesc* id, BYTE** dp)
                 emitRecordRelocation(odst, id->idAddr()->iiaAddr,
                                      id->idIsTlsGD() ? CorInfoReloc::ARM64_LIN_TLSDESC_ADR_PAGE21
                                                      : CorInfoReloc::ARM64_PAGEBASE_REL21);
+
+                if (id->idIsTlsGD())
+                {
+                    // This is the beginning of the TLS access linker
+                    // relaxation sequence for linux arm64. The linker may
+                    // replace these with other instructions that may trash x0.
+                    // We thus need to eagerly update GC for x0, in case the
+                    // linker's instructions trashes it earlier than we would
+                    // emit a mutating instruction that trashes it.
+                    emitGCregDeadUpd(REG_R0, dst);
+                }
             }
             else
             {
diff --git a/src/coreclr/jit/emitxarch.cpp b/src/coreclr/jit/emitxarch.cpp
index 06eec05ce90b60..32d20d45385819 100644
--- a/src/coreclr/jit/emitxarch.cpp
+++ b/src/coreclr/jit/emitxarch.cpp
@@ -18451,6 +18451,13 @@ size_t emitter::emitOutputInstr(insGroup* ig, instrDesc* id, BYTE** dp)
             {
                 dst = emitOutputData16(dst);
                 sz  = emitSizeOfInsDsc_NONE(id);
+
+                // This may mark the beginning of a TLS access linker
+                // relaxation sequence. The linker can replace these general
+                // sequences by other instructions that trash rax. We need to
+                // eagerly invalidate GC info in rax before the linker's
+                // instructions would trash it.
+                emitGCregDeadUpd(REG_RAX, dst);
                 break;
             }