JIT: Eagerly invalidate GC state for registers trashed by linker relaxation sequences (#128073)

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
clobbers `rax`. This is regardless of whether or not the value in `rax`
died a long time ago.

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.

Author: jakobbotsch

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

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);
                 }
             }
 

src/coreclr/jit/emitarm64.cpp

@@ -11749,6 +11749,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
             {

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;
             }