[NativeAOT] Use 8.1 atomics, if available, in RhpCheckedXchg/RhpCheckedLockCmpXchg (#85283)

* use atomics when available

* exch win

* conditional atomics

* moving ARM64_ATOMICS_FEATURE_FLAG_BIT to AsmOffsets.h and adding a static assert.

* enable lse locally in the asm helpers

Author: VSadov

src/coreclr/nativeaot/Runtime/AsmOffsets.h

@@ -34,6 +34,13 @@ ASM_CONST(     2,     2, STRING_COMPONENT_SIZE)
 ASM_CONST(     E,    16, STRING_BASE_SIZE)
 ASM_CONST(3FFFFFDF,3FFFFFDF,MAX_STRING_LENGTH)
 
+
+#if defined(HOST_ARM64)
+// Bit position for the ARM64IntrinsicConstants_Atomics flags, to be used with tbz / tbnz instructions
+// ARM64IntrinsicConstants_Atomics = 0x0080
+ASM_CONST(     7,     7, ARM64_ATOMICS_FEATURE_FLAG_BIT)
+#endif
+
 ASM_OFFSET(    0,     0, MethodTable, m_usComponentSize)
 ASM_OFFSET(    0,     0, MethodTable, m_uFlags)
 ASM_OFFSET(    4,     4, MethodTable, m_uBaseSize)

src/coreclr/nativeaot/Runtime/EHHelpers.cpp

@@ -286,6 +286,11 @@ EXTERN_C void * RhpCheckedXchgAVLocation;
 EXTERN_C void * RhpLockCmpXchg32AVLocation;
 EXTERN_C void * RhpLockCmpXchg64AVLocation;
 
+#if defined(HOST_ARM64) && !defined(LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT)
+EXTERN_C void* RhpCheckedLockCmpXchgAVLocation2;
+EXTERN_C void* RhpCheckedXchgAVLocation2;
+#endif
+
 static bool InWriteBarrierHelper(uintptr_t faultingIP)
 {
 #ifndef USE_PORTABLE_HELPERS
@@ -297,6 +302,10 @@ static bool InWriteBarrierHelper(uintptr_t faultingIP)
         (uintptr_t)&RhpCheckedXchgAVLocation,
         (uintptr_t)&RhpLockCmpXchg32AVLocation,
         (uintptr_t)&RhpLockCmpXchg64AVLocation,
+#if defined(HOST_ARM64) && !defined(LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT)
+        (uintptr_t)&RhpCheckedLockCmpXchgAVLocation2,
+        (uintptr_t)&RhpCheckedXchgAVLocation2,
+#endif
     };
 
     // compare the IP against the list of known possible AV locations in the write barrier helpers

src/coreclr/nativeaot/Runtime/IntrinsicConstants.h

@@ -50,6 +50,11 @@ enum ARM64IntrinsicConstants
     ARM64IntrinsicConstants_Atomics = 0x0080,
     ARM64IntrinsicConstants_Rcpc = 0x0100,
 };
+
+// Bit position for the ARM64IntrinsicConstants_Atomics flags, to be used with tbz / tbnz instructions
+static const int ARM64_ATOMICS_FEATURE_FLAG_BIT = 7;
+static_assert((1 << ARM64_ATOMICS_FEATURE_FLAG_BIT) == ARM64IntrinsicConstants_Atomics, "ARM64_ATOMICS_FEATURE_FLAG_BIT must match with ARM64IntrinsicConstants_Atomics");
+
 #endif //HOST_ARM64
 
 #endif //!INTRINSICCONSTANTS_INCLUDED

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

@@ -420,4 +420,3 @@ endif
 ifdef FEATURE_USE_SOFTWARE_WRITE_WATCH_FOR_GC_HEAP
 EXTERN g_write_watch_table  : QWORD
 endif
-

src/coreclr/nativeaot/Runtime/arm64/AsmMacros.h

@@ -116,6 +116,7 @@ OFFSETOF__Thread__m_alloc_context__alloc_limit      equ OFFSETOF__Thread__m_rgbA
     EXTERN g_write_watch_table
 #endif
 
+    EXTERN g_cpuFeatures
 
 ;; -----------------------------------------------------------------------------
 ;; Macro used to assign an alternate name to a symbol containing characters normally disallowed in a symbol
@@ -163,6 +164,16 @@ MovInstr SETS "movk"
         ldr  $Reg, [$Reg, $Name]
     MEND
 
+;; ---------------------------------------------------------------------------- -
+;; Macro for loading a 32bit value of a global variable into a register
+    MACRO
+        PREPARE_EXTERNAL_VAR_INDIRECT_W $Name, $RegNum
+
+        adrp x$RegNum, $Name
+        ldr  w$RegNum, [x$RegNum, $Name]
+    MEND
+
+
 ;; -----------------------------------------------------------------------------
 ;;
 ;; Macro to export a pointer to an address inside a stub as a 64-bit variable

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

@@ -268,6 +268,10 @@ LEAF_END RhpAssignRef, _TEXT
 // - Function "InWriteBarrierHelper" assumes an AV due to passed in null pointer will happen at RhpCheckedLockCmpXchgAVLocation
 // - Function "UnwindSimpleHelperToCaller" assumes no registers were pushed and LR contains the return address
 
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+.arch_extension lse
+#endif
+
 // RhpCheckedLockCmpXchg(Object** dest, Object* value, Object* comparand)
 //
 // Interlocked compare exchange on objectref.
@@ -279,21 +283,36 @@ LEAF_END RhpAssignRef, _TEXT
 //
 // On exit:
 //  x0: original value of objectref
-//  x10, x12, x17: trashed
+//  x10, x12, x16, x17: trashed
 //
     LEAF_ENTRY RhpCheckedLockCmpXchg
 
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        PREPARE_EXTERNAL_VAR_INDIRECT_W g_cpuFeatures, 16
+        tbz    w16, #ARM64_ATOMICS_FEATURE_FLAG_BIT, CmpXchgRetry
+#endif
+
+        mov    x10, x2
+    ALTERNATE_ENTRY RhpCheckedLockCmpXchgAVLocation
+        casal  x10, x1, [x0]                  // exchange
+        cmp    x2, x10
+        bne    CmpXchgNoUpdate
+
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        b      DoCardsCmpXchg
 CmpXchgRetry:
         // Check location value is what we expect.
-    ALTERNATE_ENTRY  RhpCheckedLockCmpXchgAVLocation
+    ALTERNATE_ENTRY  RhpCheckedLockCmpXchgAVLocation2
         ldaxr   x10, [x0]
         cmp     x10, x2
         bne     CmpXchgNoUpdate
 
         // Current value matches comparand, attempt to update with the new value.
         stlxr   w12, x1, [x0]
         cbnz    w12, CmpXchgRetry
+#endif
 
+DoCardsCmpXchg:
         // We have successfully updated the value of the objectref so now we need a GC write barrier.
         // The following barrier code takes the destination in x0 and the value in x1 so the arguments are
         // already correctly set up.
@@ -303,7 +322,12 @@ CmpXchgRetry:
 CmpXchgNoUpdate:
         // x10 still contains the original value.
         mov     x0, x10
+
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        tbnz    w16, #ARM64_ATOMICS_FEATURE_FLAG_BIT, NoBarrierCmpXchg
         InterlockedOperationBarrier
+NoBarrierCmpXchg:
+#endif
         ret     lr
 
     LEAF_END RhpCheckedLockCmpXchg, _TEXT
@@ -323,19 +347,31 @@ CmpXchgNoUpdate:
 // On exit:
 //  x0: original value of objectref
 //  x10: trashed
-//  x12, x17: trashed
+//  x12, x16, x17: trashed
 //
     LEAF_ENTRY RhpCheckedXchg, _TEXT
 
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        PREPARE_EXTERNAL_VAR_INDIRECT_W g_cpuFeatures, 16
+        tbz    w16, #ARM64_ATOMICS_FEATURE_FLAG_BIT, ExchangeRetry
+#endif
+
+    ALTERNATE_ENTRY  RhpCheckedXchgAVLocation
+        swpal  x1, x10, [x0]                   // exchange
+
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        b      DoCardsXchg
 ExchangeRetry:
         // Read the existing memory location.
-    ALTERNATE_ENTRY  RhpCheckedXchgAVLocation
+    ALTERNATE_ENTRY  RhpCheckedXchgAVLocation2
         ldaxr   x10,  [x0]
 
         // Attempt to update with the new value.
         stlxr   w12, x1, [x0]
         cbnz    w12, ExchangeRetry
+#endif
 
+DoCardsXchg:
         // We have successfully updated the value of the objectref so now we need a GC write barrier.
         // The following barrier code takes the destination in x0 and the value in x1 so the arguments are
         // already correctly set up.
@@ -344,7 +380,16 @@ ExchangeRetry:
 
         // x10 still contains the original value.
         mov     x0, x10
+
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        tbnz    w16, #ARM64_ATOMICS_FEATURE_FLAG_BIT, NoBarrierXchg
         InterlockedOperationBarrier
+NoBarrierXchg:
+#endif
         ret
 
     LEAF_END RhpCheckedXchg, _TEXT
+
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+.arch_extension nolse
+#endif

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

@@ -294,22 +294,37 @@ NotInHeap
 ;;
 ;; On exit:
 ;;  x0: original value of objectref
-;;  x10, x12, x17: trashed
+;;  x10, x12, x16, x17: trashed
 ;;
     LEAF_ENTRY RhpCheckedLockCmpXchg
 
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        PREPARE_EXTERNAL_VAR_INDIRECT_W g_cpuFeatures, 16
+        tbz    x16, #ARM64_ATOMICS_FEATURE_FLAG_BIT, CmpXchgRetry
+#endif
+
+        mov    x10, x2
+    ALTERNATE_ENTRY RhpCheckedLockCmpXchgAVLocation
+        casal  x10, x1, [x0]                  ;; exchange
+        cmp    x2, x10
+        bne    CmpXchgNoUpdate
+
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        b      DoCardsCmpXchg
 CmpXchgRetry
         ;; Check location value is what we expect.
-    ALTERNATE_ENTRY RhpCheckedLockCmpXchgAVLocation
+    ALTERNATE_ENTRY  RhpCheckedLockCmpXchgAVLocation2
         ldaxr   x10, [x0]
         cmp     x10, x2
         bne     CmpXchgNoUpdate
 
         ;; Current value matches comparand, attempt to update with the new value.
         stlxr   w12, x1, [x0]
         cbnz    w12, CmpXchgRetry
+#endif
 
-        ;; We've successfully updated the value of the objectref so now we need a GC write barrier.
+DoCardsCmpXchg
+        ;; We have successfully updated the value of the objectref so now we need a GC write barrier.
         ;; The following barrier code takes the destination in x0 and the value in x1 so the arguments are
         ;; already correctly set up.
 
@@ -318,7 +333,12 @@ CmpXchgRetry
 CmpXchgNoUpdate
         ;; x10 still contains the original value.
         mov     x0, x10
+
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        tbnz    x16, #ARM64_ATOMICS_FEATURE_FLAG_BIT, NoBarrierCmpXchg
         InterlockedOperationBarrier
+NoBarrierCmpXchg
+#endif
         ret     lr
 
     LEAF_END RhpCheckedLockCmpXchg
@@ -338,28 +358,45 @@ CmpXchgNoUpdate
 ;; On exit:
 ;;  x0: original value of objectref
 ;;  x10: trashed
-;;  x12, x17: trashed
+;;  x12, x16, x17: trashed
 ;;
     LEAF_ENTRY RhpCheckedXchg
 
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        PREPARE_EXTERNAL_VAR_INDIRECT_W g_cpuFeatures, 16
+        tbz    x16, #ARM64_ATOMICS_FEATURE_FLAG_BIT, ExchangeRetry
+#endif
+
+    ALTERNATE_ENTRY  RhpCheckedXchgAVLocation
+        swpal  x1, x10, [x0]                   ;; exchange
+
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        b      DoCardsXchg
 ExchangeRetry
         ;; Read the existing memory location.
-    ALTERNATE_ENTRY RhpCheckedXchgAVLocation
+    ALTERNATE_ENTRY  RhpCheckedXchgAVLocation2
         ldaxr   x10,  [x0]
 
         ;; Attempt to update with the new value.
         stlxr   w12, x1, [x0]
         cbnz    w12, ExchangeRetry
+#endif
 
-        ;; We've successfully updated the value of the objectref so now we need a GC write barrier.
+DoCardsXchg
+        ;; We have successfully updated the value of the objectref so now we need a GC write barrier.
         ;; The following barrier code takes the destination in x0 and the value in x1 so the arguments are
         ;; already correctly set up.
 
         INSERT_CHECKED_WRITE_BARRIER_CORE x0, x1
 
         ;; x10 still contains the original value.
         mov     x0, x10
+
+#ifndef LSE_INSTRUCTIONS_ENABLED_BY_DEFAULT
+        tbnz    x16, #ARM64_ATOMICS_FEATURE_FLAG_BIT, NoBarrierXchg
         InterlockedOperationBarrier
+NoBarrierXchg
+#endif
         ret
 
     LEAF_END RhpCheckedXchg