Ensure proper InterpThreadContext initialization (#117479)

* Ensure proper InterpThreadContext initialization

In case an interpreted method was compiled on one thread and then later
executed on another thread that never executed any interpreted code
before, the InterpThreadContext was not initialized.
This change fixes that by calling Thread::GetInterpThreadContext from
the InterpreterStub in case it was not initialized for the current
thread yet.

* Fix Unix Amd64

Author: janvorli

src/coreclr/pal/inc/unixasmmacrosarm64.inc

@@ -204,9 +204,9 @@ C_FUNC(\Name\()_End):
 
         __PWTB_TransitionBlock = __PWTB_FloatArgumentRegisters
 
-        .if (__PWTB_SaveFPArgs == 1)
-            __PWTB_TransitionBlock = __PWTB_TransitionBlock + SIZEOF__FloatArgumentRegisters
-        .endif
+        // Always reserve space for the FP arg regs, even if they are not saved in the PROLOG_WITH_TRANSITION_BLOCK
+        // Some code using this macro saves them optionally
+        __PWTB_TransitionBlock = __PWTB_TransitionBlock + SIZEOF__FloatArgumentRegisters
 
         __PWTB_StackAlloc = __PWTB_TransitionBlock
         __PWTB_ArgumentRegisters = __PWTB_StackAlloc + 104

src/coreclr/vm/amd64/AsmHelpers.asm

@@ -4,6 +4,8 @@
 include AsmMacros.inc
 include asmconstants.inc
 
+Thread_GetInterpThreadContext  TEXTEQU <?GetInterpThreadContext@Thread@@QEAAPEAUInterpThreadContext@@XZ>
+
 extern NDirectImportWorker:proc
 extern ThePreStub:proc
 extern  ProfileEnter:proc
@@ -13,6 +15,7 @@ extern OnHijackWorker:proc
 extern JIT_RareDisableHelperWorker:proc
 ifdef FEATURE_INTERPRETER
 extern ExecuteInterpretedMethod:proc
+extern Thread_GetInterpThreadContext:proc
 endif
 
 extern g_pPollGC:QWORD
@@ -565,12 +568,21 @@ NESTED_ENTRY InterpreterStub, _TEXT
 
         INLINE_GETTHREAD r10; thrashes rax and r11
 
+        mov             rax, qword ptr [r10 + OFFSETOF__Thread__m_pInterpThreadContext]
+        test            rax, rax
+        jnz             HaveInterpThreadContext
+
+        mov             rcx, r10
+        call            Thread_GetInterpThreadContext
+        RESTORE_ARGUMENT_REGISTERS __PWTB_ArgumentRegisters
+        RESTORE_FLOAT_ARGUMENT_REGISTERS __PWTB_FloatArgumentRegisters
+
+HaveInterpThreadContext:
+        mov             r10, qword ptr [rax + OFFSETOF__InterpThreadContext__pStackPointer]
         ; Load the InterpMethod pointer from the IR bytecode
         mov             rax, qword ptr [rbx]
         mov             rax, qword ptr [rax + OFFSETOF__InterpMethod__pCallStub]
         lea             r11, qword ptr [rax + OFFSETOF__CallStubHeader__Routines]
-        mov             r10, qword ptr [r10 + OFFSETOF__Thread__m_pInterpThreadContext]
-        mov             r10, qword ptr [r10 + OFFSETOF__InterpThreadContext__pStackPointer]
         lea             rax, [rsp + __PWTB_TransitionBlock]
         ; Copy the arguments to the interpreter stack, invoke the InterpExecMethod and load the return value
         call            qword ptr [r11]

src/coreclr/vm/amd64/asmhelpers.S

@@ -444,6 +444,15 @@ NESTED_ENTRY InterpreterStub, _TEXT, NoHandler
         INLINE_GETTHREAD // result in rax, it can thrash all argument registers as it can call a helper
         mov             r10, rax
 
+        mov             rax, qword ptr [r10 + OFFSETOF__Thread__m_pInterpThreadContext]
+        test            rax, rax
+        jnz             LOCAL_LABEL(HaveInterpThreadContext)
+
+        mov             rcx, r10
+        call            C_FUNC(_ZN6Thread22GetInterpThreadContextEv) // Thread::GetInterpThreadContext
+
+LOCAL_LABEL(HaveInterpThreadContext):
+        mov             r10, qword ptr [rax + OFFSETOF__InterpThreadContext__pStackPointer]
         // Load the InterpMethod pointer from the IR bytecode
         mov             rax, qword ptr [rbx]
         mov             rax, qword ptr [rax + OFFSETOF__InterpMethod__pCallStub]
@@ -463,8 +472,6 @@ NESTED_ENTRY InterpreterStub, _TEXT, NoHandler
         movsd           xmm6, real8 ptr [rsp + __PWTB_FloatArgumentRegisters + 0x60]
         movsd           xmm7, real8 ptr [rsp + __PWTB_FloatArgumentRegisters + 0x70]
         lea             r11, qword ptr [rax + OFFSETOF__CallStubHeader__Routines]
-        mov             r10, qword ptr [r10 + OFFSETOF__Thread__m_pInterpThreadContext]
-        mov             r10, qword ptr [r10 + OFFSETOF__InterpThreadContext__pStackPointer]
         lea             rax, [rsp + __PWTB_TransitionBlock]
         // rbx contains IR bytecode address
         // Copy the arguments to the interpreter stack, invoke the InterpExecMethod and load the return value

src/coreclr/vm/arm64/asmhelpers.S

@@ -695,27 +695,40 @@ NESTED_ENTRY InterpreterStub, _TEXT, NoHandler
 #endif
     INLINE_GETTHREAD x20 // thrashes x0 on Apple OSes (and possibly other arg registers on other Unixes)
 
-    // On Apple, the INLINE_GETTHREAD is guaranteed to not to thrash argument registers other than x0
-    // On other Unixes, there is no such guarantee, so we need to always restore the argument registers
+    ldr x11, [x20, #OFFSETOF__Thread__m_pInterpThreadContext] 
+    cbnz x11, LOCAL_LABEL(HaveInterpThreadContext)
+
+#ifdef TARGET_APPLE
+    // There Thread::GetInterpThreadContext can destroy all argument registers, so we 
+    // need to save them. For non-Apple, they have been already saved in the PROLOG_WITH_TRANSITION_BLOCK
+    // Restore x0 thrashed by the INLINE_GETTHREAD
+    mov x0, x21
+    SAVE_ARGUMENT_REGISTERS sp, __PWTB_ArgumentRegisters
+    SAVE_FLOAT_ARGUMENT_REGISTERS sp, __PWTB_FloatArgumentRegisters
+#endif
+
+    mov x0, x20
+    bl C_FUNC(_ZN6Thread22GetInterpThreadContextEv) // Thread::GetInterpThreadContext
+    mov x11, x0
+
 #ifndef TARGET_APPLE
-    ldp x0, x1, [sp, #__PWTB_ArgumentRegisters + 8]
-    ldp x2, x3, [sp, #__PWTB_ArgumentRegisters + 0x18]
-    ldp x4, x5, [sp, #__PWTB_ArgumentRegisters + 0x28]
-    ldp x6, x7, [sp, #__PWTB_ArgumentRegisters + 0x38]
-    ldp q0, q1, [sp, #__PWTB_FloatArgumentRegisters]
-    ldp q2, q3, [sp, #__PWTB_FloatArgumentRegisters + 0x20]
-    ldp q4, q5, [sp, #__PWTB_FloatArgumentRegisters + 0x40]
-    ldp q6, q7, [sp, #__PWTB_FloatArgumentRegisters + 0x60]
-#else // !TARGET_APPLE
-    // Restore the thrashed x0
+LOCAL_LABEL(HaveInterpThreadContext):
+#endif
+
+    RESTORE_ARGUMENT_REGISTERS sp, __PWTB_ArgumentRegisters
+    RESTORE_FLOAT_ARGUMENT_REGISTERS sp, __PWTB_FloatArgumentRegisters
+
+#ifdef TARGET_APPLE
+LOCAL_LABEL(HaveInterpThreadContext):
+    // On Apple, the INLINE_GETTHREAD is guaranteed to not to thrash argument registers other than x0
+    // So we restore just the x0
     mov x0, x21
-#endif // !TARGET_APPLE
+#endif // TARGET_APPLE
 
     ldr x9, [x19] // InterpMethod*
     ldr x9, [x9, #OFFSETOF__InterpMethod__pCallStub]
     add x10, x9, #OFFSETOF__CallStubHeader__Routines
-    ldr x9, [x20, #OFFSETOF__Thread__m_pInterpThreadContext] 
-    ldr x9, [x9, #OFFSETOF__InterpThreadContext__pStackPointer]
+    ldr x9, [x11, #OFFSETOF__InterpThreadContext__pStackPointer]
     // x19 contains IR bytecode address
     // Copy the arguments to the interpreter stack, invoke the InterpExecMethod and load the return value
     ldr x11, [x10], #8

src/coreclr/vm/arm64/asmhelpers.asm

@@ -24,6 +24,8 @@
     IMPORT HijackHandler
     IMPORT ThrowControlForThread
 #ifdef FEATURE_INTERPRETER
+    SETALIAS Thread_GetInterpThreadContext, ?GetInterpThreadContext@Thread@@QEAAPEAUInterpThreadContext@@XZ
+    IMPORT $Thread_GetInterpThreadContext
     IMPORT ExecuteInterpretedMethod
 #endif
 
@@ -1064,13 +1066,22 @@ JIT_PollGCRarePath
 
         INLINE_GETTHREAD x20, x19
 
+        ldr x11, [x20, #OFFSETOF__Thread__m_pInterpThreadContext]
+        cbnz x11, HaveInterpThreadContext
+
+        mov x0, x20
+        bl $Thread_GetInterpThreadContext
+        mov x11, x0
+        RESTORE_ARGUMENT_REGISTERS sp, __PWTB_ArgumentRegisters
+        RESTORE_FLOAT_ARGUMENT_REGISTERS sp, __PWTB_FloatArgumentRegisters
+
+HaveInterpThreadContext
         ; IR bytecode address
         mov x19, METHODDESC_REGISTER
         ldr x9, [METHODDESC_REGISTER]
         ldr x9, [x9, #OFFSETOF__InterpMethod__pCallStub]
         add x10, x9, #OFFSETOF__CallStubHeader__Routines
-        ldr x9, [x20, #OFFSETOF__Thread__m_pInterpThreadContext] 
-        ldr x9, [x9, #OFFSETOF__InterpThreadContext__pStackPointer]
+        ldr x9, [x11, #OFFSETOF__InterpThreadContext__pStackPointer]
         ; x19 contains IR bytecode address
         ; Copy the arguments to the interpreter stack, invoke the InterpExecMethod and load the return value
         ldr x11, [x10], #8