[NativeAOT] Simplifying access to thread static variables (#84566)

Fixes: #84373

- [x] separate "fast" inlinable case . (used for singlemodule, not dynamic cases, when optimizing)
- [x] make the storage for fast threadstatics a single "combo" instance instead of array of instances.

Author: VSadov

src/coreclr/nativeaot/Runtime/AsmOffsets.h

@@ -58,8 +58,6 @@ ASM_OFFSET(    0,    78, Thread, m_uHijackedReturnValueFlags)
 ASM_OFFSET(   48,    80, Thread, m_pExInfoStackHead)
 ASM_OFFSET(   4c,    88, Thread, m_threadAbortException)
 
-ASM_OFFSET(   50,    90, Thread, m_pThreadLocalModuleStatics)
-
 ASM_SIZEOF(   14,    20, EHEnum)
 
 ASM_OFFSET(    0,     0, gc_alloc_context, alloc_ptr)

src/coreclr/nativeaot/Runtime/RuntimeInstance.cpp

@@ -275,6 +275,25 @@ RuntimeInstance::TypeManagerList& RuntimeInstance::GetTypeManagerList()
     return m_TypeManagerList;
 }
 
+TypeManager* RuntimeInstance::GetSingleTypeManager()
+{
+    auto head = m_TypeManagerList.GetHead();
+    if (head != NULL && head->m_pNext == NULL)
+    {
+        return head->m_pTypeManager;
+    }
+
+    return NULL;
+}
+
+COOP_PINVOKE_HELPER(TypeManagerHandle, RhGetSingleTypeManager, ())
+{
+    TypeManager* typeManager = GetRuntimeInstance()->GetSingleTypeManager();
+    ASSERT(typeManager != NULL);
+
+    return TypeManagerHandle::Create(typeManager);
+}
+
 // static
 bool RuntimeInstance::Initialize(HANDLE hPalInstance)
 {

src/coreclr/nativeaot/Runtime/RuntimeInstance.h

@@ -99,6 +99,7 @@ class RuntimeInstance
 
     bool RegisterTypeManager(TypeManager * pTypeManager);
     TypeManagerList& GetTypeManagerList();
+    TypeManager* GetSingleTypeManager();
     OsModuleList* GetOsModuleList();
 
     bool RegisterUnboxingStubs(PTR_VOID pvStartRange, uint32_t cbRange);

src/coreclr/nativeaot/Runtime/amd64/AsmMacros.inc

@@ -235,11 +235,11 @@ Name    dq      offset AddressToExport
 _tls_array     equ 58h     ;; offsetof(TEB, ThreadLocalStoragePointer)
 
 ;;
-;; __declspec(thread) version
+;; __declspec(thread) variable
 ;;
-INLINE_GETTHREAD macro destReg, trashReg
+INLINE_GET_TLS_VAR macro destReg, trashReg, variable
     EXTERN _tls_index : DWORD
-    EXTERN tls_CurrentThread:DWORD
+    EXTERN variable:DWORD
 
 ;;
 ;; construct 'eax' from 'rax' so that the register size and data size match
@@ -255,11 +255,18 @@ endif
     mov         destRegDWORD, [_tls_index]
     mov         trashReg, gs:[_tls_array]
     mov         trashReg, [trashReg + destReg * 8]
-    mov         destRegDWORD, SECTIONREL tls_CurrentThread
+    mov         destRegDWORD, SECTIONREL variable
     add         destReg, trashReg
 
 endm
 
+;;
+;; __declspec(thread) tls_CurrentThread
+;;
+INLINE_GETTHREAD macro destReg, trashReg
+    INLINE_GET_TLS_VAR destReg, trashReg, tls_CurrentThread
+endm
+
 INLINE_THREAD_UNHIJACK macro threadReg, trashReg1, trashReg2
         ;;
         ;; Thread::Unhijack()

src/coreclr/nativeaot/Runtime/amd64/MiscStubs.S

@@ -45,57 +45,20 @@ LOCAL_LABEL(ProbeLoop):
         ret
 NESTED_END RhpStackProbe, _TEXT
 
-NESTED_ENTRY RhpGetThreadStaticBaseForType, _TEXT, NoHandler
-        // On entry:
-        //   rdi - TypeManagerSlot*
-        //   rsi - type index
+NESTED_ENTRY RhpGetInlinedThreadStaticBase, _TEXT, NoHandler
         // On exit:
         //   rax - the thread static base for the given type
 
-        push_nonvol_reg rbx
-        push_nonvol_reg r12
-
-        mov     rbx, rdi               // Save TypeManagerSlot*
-        mov     r12, rsi               // Save type index
-
-        // rax = GetThread()
-        INLINE_GETTHREAD
-
-        mov     r8d, [rbx + 8]         // Get ModuleIndex out of the TypeManagerSlot
+        // rdi = &tls_InlinedThreadStatics
+        INLINE_GET_TLS_VAR tls_InlinedThreadStatics
+        mov     rdi, rax
 
         // get per-thread storage
-        mov     rax, [rax + OFFSETOF__Thread__m_pThreadLocalModuleStatics]
-
-        // get per-module storage
+        mov     rax, [rdi]
         test    rax, rax
-        jz      LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
-        cmp     r8d, [rax + OFFSETOF__Array__m_Length]
-        jae     LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
-        mov     rax, [rax + r8 * 8 + 0x10]
+        jz      C_FUNC(RhpGetInlinedThreadStaticBaseSlow)    // rdi contains the storage ref
 
-        // get the actual per-type storage
-        test    rax, rax
-        jz      LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
-        cmp     r12d, [rax + OFFSETOF__Array__m_Length]
-        jae     LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
-        mov     rax, [rax + r12 * 8 + 0x10]
-
-        // if have storage, return it
-        test    rax, rax
-        jz      LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
-
-        .cfi_remember_state
-        pop_nonvol_reg r12
-        pop_nonvol_reg rbx
+        // return it
         ret
+NESTED_END RhpGetInlinedThreadStaticBase, _TEXT
 
-        .cfi_restore_state
-        .cfi_def_cfa_offset 24  // workaround cfi_restore_state bug
-LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath):
-        mov         rdi, rbx    // restore TypeManagerSlot*
-        mov         rsi, r12    // restore type index
-
-        pop_nonvol_reg r12
-        pop_nonvol_reg rbx
-        jmp C_FUNC(RhpGetThreadStaticBaseForTypeSlow)
-NESTED_END RhpGetThreadStaticBaseForType, _TEXT

src/coreclr/nativeaot/Runtime/amd64/MiscStubs.asm

@@ -3,7 +3,7 @@
 
 include AsmMacros.inc
 
-EXTERN RhpGetThreadStaticBaseForTypeSlow : PROC
+EXTERN RhpGetInlinedThreadStaticBaseSlow : PROC
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ; The following helper will access ("probe") a word on each page of the stack
@@ -39,41 +39,20 @@ ProbeLoop:
 
 LEAF_END RhpStackProbe, _TEXT
 
-LEAF_ENTRY RhpGetThreadStaticBaseForType, _TEXT
-        ; On entry and thorough the procedure:
-        ;   rcx - TypeManagerSlot*
-        ;   rdx - type index
+LEAF_ENTRY RhpGetInlinedThreadStaticBase, _TEXT
         ; On exit:
         ;   rax - the thread static base for the given type
 
-        ;; rax = GetThread(), TRASHES r8
-        INLINE_GETTHREAD rax, r8
-
-        mov     r8d, [rcx + 8]         ; Get ModuleIndex out of the TypeManagerSlot
+        ;; rcx = &tls_InlinedThreadStatics, TRASHES r8
+        INLINE_GET_TLS_VAR rcx, r8, tls_InlinedThreadStatics
 
         ;; get per-thread storage
-        mov     rax, [rax + OFFSETOF__Thread__m_pThreadLocalModuleStatics]
-
-        ;; get per-module storage
-        test    rax, rax
-        jz      RhpGetThreadStaticBaseForTypeSlow
-        cmp     r8d, [rax + OFFSETOF__Array__m_Length]
-        jae     RhpGetThreadStaticBaseForTypeSlow
-        mov     rax, [rax + r8 * 8 + 10h]
-
-        ;; get the actual per-type storage
+        mov     rax, [rcx]
         test    rax, rax
-        jz      RhpGetThreadStaticBaseForTypeSlow
-        cmp     edx, [rax + OFFSETOF__Array__m_Length]
-        jae     RhpGetThreadStaticBaseForTypeSlow
-        mov     rax, [rax + rdx * 8 + 10h]
-
-        ;; if have storage, return it
-        test    rax, rax
-        jz      RhpGetThreadStaticBaseForTypeSlow
+        jz      RhpGetInlinedThreadStaticBaseSlow   ;; rcx contains the storage ref
 
+        ;; return it
         ret
-
-LEAF_END RhpGetThreadStaticBaseForType, _TEXT
+LEAF_END RhpGetInlinedThreadStaticBase, _TEXT
 
 end

src/coreclr/nativeaot/Runtime/arm64/MiscStubs.S

@@ -4,3 +4,21 @@
 #include <unixasmmacros.inc>
 #include "AsmOffsets.inc"
 
+NESTED_ENTRY RhpGetInlinedThreadStaticBase, _TEXT, NoHandler
+        // On exit:
+        //   x0 - the thread static base for the given type
+
+        // x1 = GetThread()
+        INLINE_GET_TLS_VAR x1, C_FUNC(tls_InlinedThreadStatics)
+
+        // get per-thread storage
+        ldr     x0, [x1]
+        cbnz    x0, HaveValue
+        mov     x0, x1
+        b       C_FUNC(RhpGetInlinedThreadStaticBaseSlow)
+
+HaveValue:
+        // return it
+        ret
+
+NESTED_END RhpGetInlinedThreadStaticBase, _TEXT

src/coreclr/nativeaot/Runtime/arm64/MiscStubs.asm

@@ -3,6 +3,25 @@
 
 #include "AsmMacros.h"
 
+    EXTERN RhpGetInlinedThreadStaticBaseSlow
+
     TEXTAREA
 
+;; On exit:
+;;   x0 - the thread static base for the given type
+    LEAF_ENTRY RhpGetInlinedThreadStaticBase
+        ;; x1 = &tls_InlinedThreadStatics, TRASHES x2
+        INLINE_GET_TLS_VAR x1, x2, tls_InlinedThreadStatics
+
+        ;; get per-thread storage
+        ldr     x0, [x1]
+        cbnz    x0, HaveValue
+        mov     x0, x1
+        b       RhpGetInlinedThreadStaticBaseSlow
+
+HaveValue
+        ;; return it
+        ret
+    LEAF_END RhpGetInlinedThreadStaticBase
+
     end

src/coreclr/nativeaot/Runtime/gcrhscan.cpp

@@ -54,6 +54,14 @@ void GCToEEInterface::GcScanRoots(EnumGcRefCallbackFunc * fn,  int condemned, in
         else
 #endif
         {
+            InlinedThreadStaticRoot* pRoot = pThread->GetInlinedThreadStaticList();
+            while (pRoot != NULL)
+            {
+                STRESS_LOG2(LF_GC | LF_GCROOTS, LL_INFO100, "{ Scanning Thread's %p inline thread statics root %p. \n", pThread, pRoot);
+                GcEnumObject(&pRoot->m_threadStaticsBase, 0 /*flags*/, fn, sc);
+                pRoot = pRoot->m_next;
+            }
+
             STRESS_LOG1(LF_GC | LF_GCROOTS, LL_INFO100, "{ Scanning Thread's %p thread statics root. \n", pThread);
             GcEnumObject(pThread->GetThreadStaticStorage(), 0 /*flags*/, fn, sc);
 

src/coreclr/nativeaot/Runtime/thread.cpp

@@ -284,7 +284,8 @@ void Thread::Construct()
 
     // Everything else should be initialized to 0 via the static initialization of tls_CurrentThread.
 
-    ASSERT(m_pThreadLocalModuleStatics == NULL);
+    ASSERT(m_pThreadLocalStatics == NULL);
+    ASSERT(m_pInlinedThreadLocalStatics == NULL);
 
     ASSERT(m_pGCFrameRegistrations == NULL);
 
@@ -1266,15 +1267,33 @@ COOP_PINVOKE_HELPER(Object *, RhpGetThreadAbortException, ())
 
 Object** Thread::GetThreadStaticStorage()
 {
-    return &m_pThreadLocalModuleStatics;
+    return &m_pThreadLocalStatics;
 }
 
 COOP_PINVOKE_HELPER(Object**, RhGetThreadStaticStorage, ())
 {
-    Thread * pCurrentThread = ThreadStore::RawGetCurrentThread();
+    Thread* pCurrentThread = ThreadStore::RawGetCurrentThread();
     return pCurrentThread->GetThreadStaticStorage();
 }
 
+InlinedThreadStaticRoot* Thread::GetInlinedThreadStaticList()
+{
+    return m_pInlinedThreadLocalStatics;
+}
+
+void Thread::RegisterInlinedThreadStaticRoot(InlinedThreadStaticRoot* newRoot)
+{
+    ASSERT(newRoot->m_next == NULL);
+    newRoot->m_next = m_pInlinedThreadLocalStatics;
+    m_pInlinedThreadLocalStatics = newRoot;
+}
+
+COOP_PINVOKE_HELPER(void, RhRegisterInlinedThreadStaticRoot, (Object** root))
+{
+    Thread* pCurrentThread = ThreadStore::RawGetCurrentThread();
+    pCurrentThread->RegisterInlinedThreadStaticRoot((InlinedThreadStaticRoot*)root);
+}
+
 // This is function is used to quickly query a value that can uniquely identify a thread
 COOP_PINVOKE_HELPER(uint8_t*, RhCurrentNativeThreadId, ())
 {