[JIT] Add support to inline the field access of primitive types marked with TLS (#82973)

* Add CORINFO_THREAD_LOCAL_FIELD_INFO and getThreadLocalFieldInfo()

* Implementation of getThreadLocalFieldInfo()

* Change the field types from CORINFO_CONST_LOOKUP to uint32_t

* Introduce CORINFO_FIELD_STATIC_TLS_MANAGED

* Introduce GTF_FLD_TLS_MANAGED

* Introduce eeGetThreadLocalFieldInfo()

* Fix the offsetOfThreadStaticBlocks, hardcode threadStaticBlockIndex for now

* Add impThreadLocalFieldAccess()

* Switch to GS segment register with a TODO

* Add reverse map `g_threadStaticBlockTypeIDMap` for type -> ID

* Fix GS segment register encoding

* Add comment for g_threadStaticBlockTypeIDMap

* Add extra parameter for typeIndex in JIT_GetSharedNonGCThreadStaticBase

* Update JIT_GetSharedNonGCThreadStaticBase to add logic for storing typeIndex in t_threadStaticBlocks

* Fix encoding of gs:[0x58]

* fix some bugs in importer to change size from 4 to 8

* add some printf in jithelpers

* To revert

* Update the helper to skip making re-entry for staticBlock

* Add impThreadLocalFieldWrite() for store

* Fix a issue for helper

* Use JITDUMP

* Just enable the optimization for primitive types

* Update the JITEE guid

* Fix an assert for volatile variables

* Mark helper block as cold block

* wip

* Fix the size of maxThreadStaticBlock

* fix a bug to insert the entry in cache

* Add #ifdef for HOST_WINDOWS / TARGET_WINDOWS

* Dynamic memory allocation for static block array

* Call InitTypeMap only if host == windows

* Add CORINFO_HELP_GETSHARED_NONGCTHREADSTATIC_BASE_NOCTOR_OPTIMIZED

* Expand the TLS field access in late phase

* Move rarely ran block to the end

* wip

* Fix the bug to access the value from fastPath

* Remove code from importer

* Cleanup the code

* Moved CreateBlockFromTree to fgbasic.cpp

* Move fgExpandThreadLocalAccess() to flowgraph.cpp

* Add getThreadLocalStaticBlocksInfo() method

* Consume getThreadLocalStaticBlocksInfo() method

* Updated definition of getThreadLocalFieldInfo()

* Rename CORINFO_THREAD_LOCAL_FIELD_INFO to CORINFO_THREAD_STATIC_BLOCKS_INFO

* Add TARGET_WINDOWS

* Make TARGET_WINDOWS at some definitions

* Add dummy definition for unix

* jit format

* remote unnecessary code from morph

* Fix an issue on linux, add assert for tls_index != 0

* Convert from TARGET_WINDOWS to HOST_WINDOWS

* jit format

* fix the if-check

* fix gcc build error

* Added MethodHaveTlsFieldAccess()

* Handle the case for '_tls_index == 0'

* fix a typo

* Fix windows/arm64 issue

* fix encoding of teb loading for arm64

* Fix the condition in jitinterface to select correct helper

* fix the condition for other helper too

* fix the windows/x86 case

* jit format

* Use CORINFO_CONST_LOOKUP for _tls_index

* Move typeIdMap to BaseDomain class

* Introduce useFatPointerDispatch parameter for GetTypeID()

* jit format

* Do not pass classID and moduleID to the new helper

* fix build for non-windows

* fix superpmi methods

* review feedback

* fix linux build errors

* review feedback

* fix weight inherit

* inline JIT_GetNonGCThreadStaticBase_Helper to fix the contract error

* Remove dead code from helper

* Add GTF_CALL_M_EXP_TLS_ACCESS to check if we do not revisit the already expanded call node

* address feedback

* regenerate the files

* Introduce fgExpandHelper and fgExpandHelperForBlock and reuse it for various helper expansion

* fix the typo

* update the assert for runtimelookup

* check if we should skip rarely run blocks

* review feedback

* Remove GTF_IND_INVARIANT from `typeIndex` access because it produces nullptr 1st time and valid value 2nd time onwards

Author: kunalspathak

src/coreclr/inc/corinfo.h

@@ -546,6 +546,7 @@ enum CorInfoHelpFunc
     CORINFO_HELP_GETSHARED_NONGCTHREADSTATIC_BASE_NOCTOR,
     CORINFO_HELP_GETSHARED_GCTHREADSTATIC_BASE_DYNAMICCLASS,
     CORINFO_HELP_GETSHARED_NONGCTHREADSTATIC_BASE_DYNAMICCLASS,
+    CORINFO_HELP_GETSHARED_NONGCTHREADSTATIC_BASE_NOCTOR_OPTIMIZED,
 
     /* Debugger */
 
@@ -1686,6 +1687,7 @@ enum CORINFO_FIELD_ACCESSOR
     CORINFO_FIELD_STATIC_GENERICS_STATIC_HELPER, // static field access using the "generic static" helper (argument is MethodTable *)
     CORINFO_FIELD_STATIC_ADDR_HELPER,       // static field accessed using address-of helper (argument is FieldDesc *)
     CORINFO_FIELD_STATIC_TLS,               // unmanaged TLS access
+    CORINFO_FIELD_STATIC_TLS_MANAGED,       // managed TLS access
     CORINFO_FIELD_STATIC_READYTORUN_HELPER, // static field access using a runtime lookup helper
     CORINFO_FIELD_STATIC_RELOCATABLE,       // static field access using relocation (used in AOT)
     CORINFO_FIELD_INTRINSIC_ZERO,           // intrinsic zero (IntPtr.Zero, UIntPtr.Zero)
@@ -1725,6 +1727,17 @@ struct CORINFO_FIELD_INFO
     CORINFO_CONST_LOOKUP    fieldLookup;        // Used by Ready-to-Run
 };
 
+//----------------------------------------------------------------------------
+// getThreadLocalStaticBlocksInfo and CORINFO_THREAD_STATIC_BLOCKS_INFO: The EE instructs the JIT about how to access a thread local field
+
+struct CORINFO_THREAD_STATIC_BLOCKS_INFO
+{
+    CORINFO_CONST_LOOKUP tlsIndex;
+    uint32_t offsetOfThreadLocalStoragePointer;
+    uint32_t offsetOfMaxThreadStaticBlocks;
+    uint32_t offsetOfThreadStaticBlocks;
+};
+
 //----------------------------------------------------------------------------
 // Exception handling
 
@@ -2743,6 +2756,12 @@ class ICorStaticInfo
                                CORINFO_FIELD_INFO    *pResult
                               ) = 0;
 
+    virtual uint32_t getThreadLocalFieldInfo (
+                        CORINFO_FIELD_HANDLE  field) = 0;
+
+    virtual void getThreadLocalStaticBlocksInfo (
+                        CORINFO_THREAD_STATIC_BLOCKS_INFO* pInfo) = 0;
+
     // Returns true iff "fldHnd" represents a static field.
     virtual bool isFieldStatic(CORINFO_FIELD_HANDLE fldHnd) = 0;
 

src/coreclr/inc/icorjitinfoimpl_generated.h

@@ -414,6 +414,12 @@ void getFieldInfo(
           CORINFO_ACCESS_FLAGS flags,
           CORINFO_FIELD_INFO* pResult) override;
 
+uint32_t getThreadLocalFieldInfo(
+          CORINFO_FIELD_HANDLE field) override;
+
+void getThreadLocalStaticBlocksInfo(
+          CORINFO_THREAD_STATIC_BLOCKS_INFO* pInfo) override;
+
 bool isFieldStatic(
           CORINFO_FIELD_HANDLE fldHnd) override;
 

src/coreclr/inc/jiteeversionguid.h

@@ -43,11 +43,11 @@ typedef const GUID *LPCGUID;
 #define GUID_DEFINED
 #endif // !GUID_DEFINED
 
-constexpr GUID JITEEVersionIdentifier =  { /* 3054e9ba-bcfe-417c-9043-92ccc8738b80 */
-    0x3054e9ba,
-    0xbcfe,
-    0x417c,
-    {0x90, 0x43, 0x92, 0xcc, 0xc8, 0x73, 0x8b, 0x80}
+constexpr GUID JITEEVersionIdentifier = { /* 236d7997-3d71-45f9-b7d7-5241ad89a56f */
+    0x236d7997,
+    0x3d71,
+    0x45f9,
+    { 0xb7, 0xd7, 0x52, 0x41, 0xad, 0x89, 0xa5, 0x6f }
   };
 
 //////////////////////////////////////////////////////////////////////////////////////////////////////////

src/coreclr/inc/jithelpers.h

@@ -203,6 +203,7 @@
     JITHELPER(CORINFO_HELP_GETSHARED_NONGCTHREADSTATIC_BASE_NOCTOR,       JIT_GetSharedNonGCThreadStaticBase, CORINFO_HELP_SIG_REG_ONLY)
     JITHELPER(CORINFO_HELP_GETSHARED_GCTHREADSTATIC_BASE_DYNAMICCLASS,    JIT_GetSharedGCThreadStaticBaseDynamicClass, CORINFO_HELP_SIG_REG_ONLY)
     JITHELPER(CORINFO_HELP_GETSHARED_NONGCTHREADSTATIC_BASE_DYNAMICCLASS, JIT_GetSharedNonGCThreadStaticBaseDynamicClass, CORINFO_HELP_SIG_REG_ONLY)
+    JITHELPER(CORINFO_HELP_GETSHARED_NONGCTHREADSTATIC_BASE_NOCTOR_OPTIMIZED,       JIT_GetSharedNonGCThreadStaticBaseOptimized, CORINFO_HELP_SIG_REG_ONLY)
 
     // Debugger
     JITHELPER(CORINFO_HELP_DBG_IS_JUST_MY_CODE, JIT_DbgIsJustMyCode,CORINFO_HELP_SIG_REG_ONLY)

src/coreclr/jit/CMakeLists.txt

@@ -129,7 +129,7 @@ set( JIT_SOURCES
   hwintrinsic.cpp
   hostallocator.cpp
   ifconversion.cpp
-  runtimelookup.cpp
+  helperexpansion.cpp
   indirectcalltransformer.cpp
   importercalls.cpp
   importer.cpp

src/coreclr/jit/ICorJitInfo_names_generated.h

@@ -103,6 +103,8 @@ DEF_CLR_API(getFieldClass)
 DEF_CLR_API(getFieldType)
 DEF_CLR_API(getFieldOffset)
 DEF_CLR_API(getFieldInfo)
+DEF_CLR_API(getThreadLocalFieldInfo)
+DEF_CLR_API(getThreadLocalStaticBlocksInfo)
 DEF_CLR_API(isFieldStatic)
 DEF_CLR_API(getArrayOrStringLength)
 DEF_CLR_API(getBoundaries)

src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp

@@ -981,6 +981,23 @@ void WrapICorJitInfo::getFieldInfo(
     API_LEAVE(getFieldInfo);
 }
 
+uint32_t WrapICorJitInfo::getThreadLocalFieldInfo(
+          CORINFO_FIELD_HANDLE field)
+{
+    API_ENTER(getThreadLocalFieldInfo);
+    uint32_t temp = wrapHnd->getThreadLocalFieldInfo(field);
+    API_LEAVE(getThreadLocalFieldInfo);
+    return temp;
+}
+
+void WrapICorJitInfo::getThreadLocalStaticBlocksInfo(
+          CORINFO_THREAD_STATIC_BLOCKS_INFO* pInfo)
+{
+    API_ENTER(getThreadLocalStaticBlocksInfo);
+    wrapHnd->getThreadLocalStaticBlocksInfo(pInfo);
+    API_LEAVE(getThreadLocalStaticBlocksInfo);
+}
+
 bool WrapICorJitInfo::isFieldStatic(
           CORINFO_FIELD_HANDLE fldHnd)
 {

src/coreclr/jit/codegenxarch.cpp

@@ -5420,8 +5420,13 @@ void CodeGen::genCodeForIndir(GenTreeIndir* tree)
     if (addr->IsIconHandle(GTF_ICON_TLS_HDL))
     {
         noway_assert(EA_ATTR(genTypeSize(targetType)) == EA_PTRSIZE);
+#if TARGET_64BIT
+        emit->emitIns_R_C(ins_Load(TYP_I_IMPL), EA_PTRSIZE, tree->GetRegNum(), FLD_GLOBAL_GS,
+                          (int)addr->AsIntCon()->gtIconVal);
+#else
         emit->emitIns_R_C(ins_Load(TYP_I_IMPL), EA_PTRSIZE, tree->GetRegNum(), FLD_GLOBAL_FS,
                           (int)addr->AsIntCon()->gtIconVal);
+#endif
     }
     else
     {

src/coreclr/jit/compiler.cpp

@@ -5044,6 +5044,12 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     // Partially inline static initializations
     DoPhase(this, PHASE_EXPAND_STATIC_INIT, &Compiler::fgExpandStaticInit);
 
+    if (TargetOS::IsWindows)
+    {
+        // Currently this is only applicable for Windows
+        DoPhase(this, PHASE_EXPAND_TLS, &Compiler::fgExpandThreadLocalAccess);
+    }
+
     // Insert GC Polls
     DoPhase(this, PHASE_INSERT_GC_POLLS, &Compiler::fgInsertGCPolls);
 

src/coreclr/jit/compiler.h

@@ -4834,7 +4834,7 @@ class Compiler
 
     GenTree* fgInitThisClass();
 
-    GenTreeCall* fgGetStaticsCCtorHelper(CORINFO_CLASS_HANDLE cls, CorInfoHelpFunc helper);
+    GenTreeCall* fgGetStaticsCCtorHelper(CORINFO_CLASS_HANDLE cls, CorInfoHelpFunc helper, uint32_t typeIndex = 0);
 
     GenTreeCall* fgGetSharedCCtor(CORINFO_CLASS_HANDLE cls);
 
@@ -5287,11 +5287,21 @@ class Compiler
     PhaseStatus StressSplitTree();
     void SplitTreesRandomly();
     void SplitTreesRemoveCommas();
+
+    template <bool (Compiler::*ExpansionFunction)(BasicBlock*, Statement*, GenTreeCall*)>
+    PhaseStatus fgExpandHelper(bool skipRarelyRunBlocks);
+
+    template <bool (Compiler::*ExpansionFunction)(BasicBlock*, Statement*, GenTreeCall*)>
+    bool fgExpandHelperForBlock(BasicBlock* block);
+
     PhaseStatus fgExpandRuntimeLookups();
+    bool fgExpandRuntimeLookupsForCall(BasicBlock* block, Statement* stmt, GenTreeCall* call);
+
+    PhaseStatus fgExpandThreadLocalAccess();
+    bool fgExpandThreadLocalAccessForCall(BasicBlock* block, Statement* stmt, GenTreeCall* call);
 
-    bool fgExpandStaticInitForBlock(BasicBlock* block);
-    bool fgExpandStaticInitForCall(BasicBlock* block, Statement* stmt, GenTreeCall* call);
     PhaseStatus fgExpandStaticInit();
+    bool fgExpandStaticInitForCall(BasicBlock* block, Statement* stmt, GenTreeCall* call);
 
     PhaseStatus fgInsertGCPolls();
     BasicBlock* fgCreateGCPoll(GCPollType pollType, BasicBlock* block);
@@ -7018,6 +7028,7 @@ class Compiler
 #define OMF_HAS_MDNEWARRAY                     0x00002000 // Method contains 'new' of an MD array
 #define OMF_HAS_MDARRAYREF                     0x00004000 // Method contains multi-dimensional intrinsic array element loads or stores.
 #define OMF_HAS_STATIC_INIT                    0x00008000 // Method has static initializations we might want to partially inline
+#define OMF_HAS_TLS_FIELD                      0x00010000 // Method contains TLS field access
 
     // clang-format on
 
@@ -7068,6 +7079,16 @@ class Compiler
         optMethodFlags |= OMF_HAS_GUARDEDDEVIRT;
     }
 
+    bool doesMethodHasTlsFieldAccess()
+    {
+        return (optMethodFlags & OMF_HAS_TLS_FIELD) != 0;
+    }
+
+    void setMethodHasTlsFieldAccess()
+    {
+        optMethodFlags |= OMF_HAS_TLS_FIELD;
+    }
+
     void pickGDV(GenTreeCall*           call,
                  IL_OFFSET              ilOffset,
                  bool                   isInterface,
@@ -7769,6 +7790,7 @@ class Compiler
     void eeGetFieldInfo(CORINFO_RESOLVED_TOKEN* pResolvedToken,
                         CORINFO_ACCESS_FLAGS    flags,
                         CORINFO_FIELD_INFO*     pResult);
+    uint32_t eeGetThreadLocalFieldInfo(CORINFO_FIELD_HANDLE field);
 
     // Get the flags