[NativeAOT] Inline TLS access for windows/x64 (#89472)

* wip

* working model

* wip

* wip

* working

* Add helper for tlsIndex

* add methods in superpmi

* revert some local changes

* misc fixes

* Stop emitting TLS access code for windows/x64

* fix linux build errors

* Do not throw not implemented for windows/x64

* fix the problem where ThreadStaticBase helper was still getting invoked

* Revert certain changes from JIT method

* Introduce getThreadLocalStaticInfo_ReadyToRun()

* Consume getThreadLocalStaticInfo_ReadyToRun()

* Remove getTlsRootInfo() and other methods

* Revert unneeded changes

* missing gtInitCldHnd initialization

* save target address

* jit format

* run thunkgenerator

* resolve merge conflicts

* fix issues so the TLS is inlined

* Rename data structures from *_ReadyToRun to *_NativeAOT

* jit format

* fix some unit test

* fix a bug

* fix the weird jump problem

* use enclosing type cls handle for VN of static gc/non-gc helper

* fix a bug of resetting the flag

* useEnclosingTypeOnly from runtime to determine if VN should optimize it

* do not use vnf, but only use useEnclosingTypeAsArg0

* Use GT_COMMA to add GCStaticBase call next to TLS call

* optimize the cctor call

* Remove lazy ctor generation from tls

* Update jitinterface to not fetch data for lazy ctor

* fix errors after merge

* fix test build errors

* fix bug in CSE

* Use CORINFO_FLG_FIELD_INITCLASS instead of separate flag

* Use the INITCLASS flag

* Remove useEnclosingTypeOnly

* Add NoCtor

* Use CORINFO_HELP_READYTORUN_THREADSTATIC_BASE_NOCTOR

* Minor cleanup

* Renegenrate thunk

* Add the SetFalseTarget

* fix merge conflict resolution

* better handling of GTF_ICON_SECREL_OFFSET

better handling of GTF_ICON_SECREL_OFFSET

* review feedback

* Disable optimization for minopts

* Add comments around iiaSecRel

* jit format

* create emitNewInstrCns()

* Expand TLS even if optimization is disabled

* Track t_inlinedThreadStaticBase

Better tracking `t_inlinedThreadStaticBase` as TYP_REF

* jit format

Author: kunalspathak

src/coreclr/inc/corinfo.h

@@ -596,6 +596,7 @@ enum CorInfoHelpFunc
     CORINFO_HELP_READYTORUN_GCSTATIC_BASE,           // static gc field access
     CORINFO_HELP_READYTORUN_NONGCSTATIC_BASE,        // static non gc field access
     CORINFO_HELP_READYTORUN_THREADSTATIC_BASE,
+    CORINFO_HELP_READYTORUN_THREADSTATIC_BASE_NOCTOR,
     CORINFO_HELP_READYTORUN_NONGCTHREADSTATIC_BASE,
     CORINFO_HELP_READYTORUN_VIRTUAL_FUNC_PTR,
     CORINFO_HELP_READYTORUN_GENERIC_HANDLE,
@@ -1739,6 +1740,17 @@ struct CORINFO_THREAD_STATIC_BLOCKS_INFO
     uint32_t offsetOfGCDataPointer;
 };
 
+//----------------------------------------------------------------------------
+// getThreadLocalStaticInfo_NativeAOT and CORINFO_THREAD_STATIC_INFO_NATIVEAOT: The EE instructs the JIT about how to access a thread local field
+
+struct CORINFO_THREAD_STATIC_INFO_NATIVEAOT
+{
+    uint32_t offsetOfThreadLocalStoragePointer;
+    CORINFO_CONST_LOOKUP tlsRootObject;
+    CORINFO_CONST_LOOKUP tlsIndexObject;
+    CORINFO_CONST_LOOKUP threadStaticBaseSlow;
+};
+
 //----------------------------------------------------------------------------
 // Exception handling
 
@@ -2832,6 +2844,10 @@ class ICorStaticInfo
             bool                                isGCType
             ) = 0;
 
+    virtual void getThreadLocalStaticInfo_NativeAOT(
+            CORINFO_THREAD_STATIC_INFO_NATIVEAOT* pInfo
+            ) = 0;
+
     // Returns true iff "fldHnd" represents a static field.
     virtual bool isFieldStatic(CORINFO_FIELD_HANDLE fldHnd) = 0;
 
@@ -3366,6 +3382,7 @@ class ICorDynamicInfo : public ICorStaticInfo
 // It would be nicer to use existing IMAGE_REL_XXX constants instead of defining our own here...
 #define IMAGE_REL_BASED_REL32           0x10
 #define IMAGE_REL_BASED_THUMB_BRANCH24  0x13
+#define IMAGE_REL_SECREL                0x104
 
 // The identifier for ARM32-specific PC-relative address
 // computation corresponds to the following instruction

src/coreclr/inc/icorjitinfoimpl_generated.h

@@ -402,6 +402,9 @@ void getThreadLocalStaticBlocksInfo(
           CORINFO_THREAD_STATIC_BLOCKS_INFO* pInfo,
           bool isGCType) override;
 
+void getThreadLocalStaticInfo_NativeAOT(
+          CORINFO_THREAD_STATIC_INFO_NATIVEAOT* 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 = { /* 0cb8113f-52e5-444f-b713-dcb749b5d211 */
-    0x0cb8113f,
-    0x52e5,
-    0x444f,
-    {0xb7, 0x13, 0xdc, 0xb7, 0x49, 0xb5, 0xd2, 0x11}
+constexpr GUID JITEEVersionIdentifier = { /* d7bbeb5a-aa7d-43ec-b29e-6f24dd3bca9c */
+    0xd7bbeb5a,
+    0xaa7d,
+    0x43ec,
+    {0xb2, 0x9e, 0x6f, 0x24, 0xdd, 0x3b, 0xca, 0x9c}
   };
 
 //////////////////////////////////////////////////////////////////////////////////////////////////////////

src/coreclr/inc/jithelpers.h

@@ -264,6 +264,7 @@
     JITHELPER(CORINFO_HELP_READYTORUN_GCSTATIC_BASE,       NULL,   CORINFO_HELP_SIG_NO_ALIGN_STUB)
     JITHELPER(CORINFO_HELP_READYTORUN_NONGCSTATIC_BASE,    NULL,   CORINFO_HELP_SIG_NO_ALIGN_STUB)
     JITHELPER(CORINFO_HELP_READYTORUN_THREADSTATIC_BASE,   NULL,   CORINFO_HELP_SIG_NO_ALIGN_STUB)
+    JITHELPER(CORINFO_HELP_READYTORUN_THREADSTATIC_BASE_NOCTOR,   NULL,   CORINFO_HELP_SIG_NO_ALIGN_STUB)
     JITHELPER(CORINFO_HELP_READYTORUN_NONGCTHREADSTATIC_BASE, NULL,CORINFO_HELP_SIG_NO_ALIGN_STUB)
     JITHELPER(CORINFO_HELP_READYTORUN_VIRTUAL_FUNC_PTR,    NULL,   CORINFO_HELP_SIG_NO_ALIGN_STUB)
     JITHELPER(CORINFO_HELP_READYTORUN_GENERIC_HANDLE,      NULL,   CORINFO_HELP_SIG_NO_ALIGN_STUB)

src/coreclr/jit/ICorJitInfo_names_generated.h

@@ -100,6 +100,7 @@ DEF_CLR_API(getFieldOffset)
 DEF_CLR_API(getFieldInfo)
 DEF_CLR_API(getThreadLocalFieldInfo)
 DEF_CLR_API(getThreadLocalStaticBlocksInfo)
+DEF_CLR_API(getThreadLocalStaticInfo_NativeAOT)
 DEF_CLR_API(isFieldStatic)
 DEF_CLR_API(getArrayOrStringLength)
 DEF_CLR_API(getBoundaries)

src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp

@@ -950,6 +950,14 @@ void WrapICorJitInfo::getThreadLocalStaticBlocksInfo(
     API_LEAVE(getThreadLocalStaticBlocksInfo);
 }
 
+void WrapICorJitInfo::getThreadLocalStaticInfo_NativeAOT(
+          CORINFO_THREAD_STATIC_INFO_NATIVEAOT* pInfo)
+{
+    API_ENTER(getThreadLocalStaticInfo_NativeAOT);
+    wrapHnd->getThreadLocalStaticInfo_NativeAOT(pInfo);
+    API_LEAVE(getThreadLocalStaticInfo_NativeAOT);
+}
+
 bool WrapICorJitInfo::isFieldStatic(
           CORINFO_FIELD_HANDLE fldHnd)
 {

src/coreclr/jit/codegenxarch.cpp

@@ -411,9 +411,17 @@ void CodeGen::instGen_Set_Reg_To_Imm(emitAttr  size,
         // instruction selection due to different memory placement at runtime.
         if (EA_IS_RELOC(origAttr) && genDataIndirAddrCanBeEncodedAsPCRelOffset(imm))
         {
-            // We will use lea so displacement and not immediate will be relocatable
-            size = EA_SET_FLG(EA_REMOVE_FLG(size, EA_CNS_RELOC_FLG), EA_DSP_RELOC_FLG);
-            GetEmitter()->emitIns_R_AI(INS_lea, size, reg, imm DEBUGARG(targetHandle) DEBUGARG(gtFlags));
+            if (!EA_IS_CNS_SEC_RELOC(origAttr))
+            {
+                // We will use lea so displacement and not immediate will be relocatable
+                size = EA_SET_FLG(EA_REMOVE_FLG(size, EA_CNS_RELOC_FLG), EA_DSP_RELOC_FLG);
+                GetEmitter()->emitIns_R_AI(INS_lea, size, reg, imm DEBUGARG(targetHandle) DEBUGARG(gtFlags));
+            }
+            else
+            {
+                // For section constant, the immediate will be relocatable
+                GetEmitter()->emitIns_R_I(INS_mov, size, reg, imm DEBUGARG(targetHandle) DEBUGARG(gtFlags));
+            }
         }
         else
         {
@@ -616,6 +624,11 @@ void CodeGen::genSetRegToConst(regNumber targetReg, var_types targetType, GenTre
                 attr = EA_SET_FLG(attr, EA_BYREF_FLG);
             }
 
+            if (con->IsIconHandle(GTF_ICON_SECREL_OFFSET))
+            {
+                attr = EA_SET_FLG(attr, EA_CNS_SEC_RELOC);
+            }
+
             instGen_Set_Reg_To_Imm(attr, targetReg, cnsVal,
                                    INS_FLAGS_DONT_CARE DEBUGARG(con->gtTargetHandle) DEBUGARG(con->gtFlags));
             regSet.verifyRegUsed(targetReg);

src/coreclr/jit/compiler.h

@@ -5807,6 +5807,7 @@ class Compiler
 
     PhaseStatus fgExpandThreadLocalAccess();
     bool fgExpandThreadLocalAccessForCall(BasicBlock** pBlock, Statement* stmt, GenTreeCall* call);
+    bool fgExpandThreadLocalAccessForCallNativeAOT(BasicBlock** pBlock, Statement* stmt, GenTreeCall* call);
 
     PhaseStatus fgExpandStaticInit();
     bool fgExpandStaticInitForCall(BasicBlock** pBlock, Statement* stmt, GenTreeCall* call);

src/coreclr/jit/compiler.hpp

@@ -3941,6 +3941,7 @@ inline bool Compiler::IsSharedStaticHelper(GenTree* tree)
 #ifdef FEATURE_READYTORUN
         helper == CORINFO_HELP_READYTORUN_GENERIC_STATIC_BASE || helper == CORINFO_HELP_READYTORUN_GCSTATIC_BASE ||
         helper == CORINFO_HELP_READYTORUN_NONGCSTATIC_BASE || helper == CORINFO_HELP_READYTORUN_THREADSTATIC_BASE ||
+        helper == CORINFO_HELP_READYTORUN_THREADSTATIC_BASE_NOCTOR ||
         helper == CORINFO_HELP_READYTORUN_NONGCTHREADSTATIC_BASE ||
 #endif
         helper == CORINFO_HELP_CLASSINIT_SHARED_DYNAMICCLASS;

src/coreclr/jit/emit.h

@@ -1021,6 +1021,7 @@ class emitter
             {
                 return iiaJmpOffset;
             }
+
 #elif defined(TARGET_RISCV64)
             struct
             {
@@ -1039,6 +1040,9 @@ class emitter
             }
 #endif // defined(TARGET_RISCV64)
 
+            // Used for instrDesc that has relocatable immediate offset
+            bool iiaSecRel;
+
         } _idAddrUnion;
 
         /* Trivial wrappers to return properly typed enums */

src/coreclr/jit/emitxarch.cpp

@@ -5702,7 +5702,15 @@ void emitter::emitIns_R_I(instruction ins,
             break;
     }
 
-    id = emitNewInstrSC(attr, val);
+    if (emitComp->IsTargetAbi(CORINFO_NATIVEAOT_ABI) && EA_IS_CNS_SEC_RELOC(attr))
+    {
+        id                      = emitNewInstrCns(attr, val);
+        id->idAddr()->iiaSecRel = true;
+    }
+    else
+    {
+        id = emitNewInstrSC(attr, val);
+    }
     id->idIns(ins);
     id->idInsFmt(fmt);
     id->idReg1(reg);
@@ -15098,7 +15106,16 @@ BYTE* emitter::emitOutputRI(BYTE* dst, instrDesc* id)
 
         if (id->idIsCnsReloc())
         {
-            emitRecordRelocation((void*)(dst - (unsigned)EA_SIZE(size)), (void*)(size_t)val, IMAGE_REL_BASED_MOFFSET);
+            if (emitComp->IsTargetAbi(CORINFO_NATIVEAOT_ABI) && id->idAddr()->iiaSecRel)
+            {
+                // For section relative, the immediate offset is relocatable and hence need IMAGE_REL_SECREL
+                emitRecordRelocation((void*)(dst - (unsigned)EA_SIZE(size)), (void*)(size_t)val, IMAGE_REL_SECREL);
+            }
+            else
+            {
+                emitRecordRelocation((void*)(dst - (unsigned)EA_SIZE(size)), (void*)(size_t)val,
+                                     IMAGE_REL_BASED_MOFFSET);
+            }
         }
 
         goto DONE;

src/coreclr/jit/gentree.cpp

@@ -10771,6 +10771,8 @@ const char* GenTree::gtGetHandleKindString(GenTreeFlags flags)
             return "GTF_ICON_FIELD_SEQ";
         case GTF_ICON_STATIC_ADDR_PTR:
             return "GTF_ICON_STATIC_ADDR_PTR";
+        case GTF_ICON_SECREL_OFFSET:
+            return "GTF_ICON_SECREL_OFFSET";
         default:
             return "ILLEGAL!";
     }
@@ -12000,6 +12002,9 @@ void Compiler::gtDispConst(GenTree* tree)
                         case GTF_ICON_STATIC_ADDR_PTR:
                             printf(" static base addr cell");
                             break;
+                        case GTF_ICON_SECREL_OFFSET:
+                            printf(" relative offset in section");
+                            break;
                         default:
                             printf(" UNKNOWN");
                             break;

src/coreclr/jit/gentree.h

@@ -536,6 +536,7 @@ enum GenTreeFlags : unsigned int
     GTF_ICON_STATIC_BOX_PTR     = 0x10000000, // GT_CNS_INT -- constant is an address of the box for a STATIC_IN_HEAP field
     GTF_ICON_FIELD_SEQ          = 0x11000000, // <--------> -- constant is a FieldSeq* (used only as VNHandle)
     GTF_ICON_STATIC_ADDR_PTR    = 0x13000000, // GT_CNS_INT -- constant is a pointer to a static base address
+    GTF_ICON_SECREL_OFFSET      = 0x14000000, // GT_CNS_INT -- constant is an offset in a certain section.
 
  // GTF_ICON_REUSE_REG_VAL      = 0x00800000  // GT_CNS_INT -- GTF_REUSE_REG_VAL, defined above
     GTF_ICON_SIMD_COUNT         = 0x00200000, // GT_CNS_INT -- constant is Vector<T>.Count