Arm64: opportunistically light up LSE atomics for NativeAOT

src/coreclr/System.Private.CoreLib/src/System/Threading/Interlocked.CoreCLR.cs

@@ -5,6 +5,9 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
+#if TARGET_ARM64
+using System.Runtime.Intrinsics.Arm;
+#endif
 
 namespace System.Threading
 {
@@ -52,6 +55,13 @@ public static long Decrement(ref long location) =>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static int Exchange(ref int location1, int value)
         {
+#if TARGET_ARM64
+            // Outlined on AOT, where LSE may not be in the baseline instruction set.
+            if (Lse.IsSupported)
+            {
+                return Lse.Swap(ref location1, value);
+            }
+#endif
 #if TARGET_X86 || TARGET_AMD64 || TARGET_ARM64 || TARGET_RISCV64
             return Exchange(ref location1, value); // Must expand intrinsic
 #else
@@ -73,6 +83,13 @@ public static int Exchange(ref int location1, int value)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static long Exchange(ref long location1, long value)
         {
+#if TARGET_ARM64
+            // Outlined on AOT, where LSE may not be in the baseline instruction set.
+            if (Lse.IsSupported)
+            {
+                return Lse.Swap(ref location1, value);
+            }
+#endif
 #if TARGET_AMD64 || TARGET_ARM64 || TARGET_RISCV64
             return Exchange(ref location1, value); // Must expand intrinsic
 #else
@@ -116,6 +133,13 @@ public static long Exchange(ref long location1, long value)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static int CompareExchange(ref int location1, int value, int comparand)
         {
+#if TARGET_ARM64
+            // Outlined on AOT, where LSE may not be in the baseline instruction set.
+            if (Lse.IsSupported)
+            {
+                return Lse.CompareAndSwap(ref location1, value, comparand);
+            }
+#endif
 #if TARGET_X86 || TARGET_AMD64 || TARGET_ARM64 || TARGET_RISCV64
             return CompareExchange(ref location1, value, comparand); // Must expand intrinsic
 #else
@@ -158,6 +182,13 @@ internal static unsafe int CompareExchange(int* location1, int value, int compar
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static long CompareExchange(ref long location1, long value, long comparand)
         {
+#if TARGET_ARM64
+            // Outlined on AOT, where LSE may not be in the baseline instruction set.
+            if (Lse.IsSupported)
+            {
+                return Lse.CompareAndSwap(ref location1, value, comparand);
+            }
+#endif
 #if TARGET_AMD64 || TARGET_ARM64 || TARGET_RISCV64
             return CompareExchange(ref location1, value, comparand); // Must expand intrinsic
 #else
@@ -212,6 +243,13 @@ public static long Add(ref long location1, long value) =>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private static int ExchangeAdd(ref int location1, int value)
         {
+#if TARGET_ARM64
+            // Outlined on AOT, where LSE may not be in the baseline instruction set.
+            if (Lse.IsSupported)
+            {
+                return Lse.LoadAdd(ref location1, value);
+            }
+#endif
 #if TARGET_X86 || TARGET_AMD64 || TARGET_ARM64 || TARGET_RISCV64
             return ExchangeAdd(ref location1, value); // Must expand intrinsic
 #else
@@ -228,6 +266,13 @@ private static int ExchangeAdd(ref int location1, int value)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private static long ExchangeAdd(ref long location1, long value)
         {
+#if TARGET_ARM64
+            // Outlined on AOT, where LSE may not be in the baseline instruction set.
+            if (Lse.IsSupported)
+            {
+                return Lse.LoadAdd(ref location1, value);
+            }
+#endif
 #if TARGET_AMD64 || TARGET_ARM64 || TARGET_RISCV64
             return ExchangeAdd(ref location1, value); // Must expand intrinsic
 #else

src/coreclr/jit/codegenarm64.cpp

@@ -3810,7 +3810,7 @@ void CodeGen::genLockedInstructions(GenTreeOp* treeNode)
 
     emitAttr dataSize = emitActualTypeSize(data);
 
-    if (m_compiler->compOpportunisticallyDependsOn(InstructionSet_Atomics))
+    if (m_compiler->compGetAtomicsImplForNode(treeNode) == Compiler::AtomicsImpl::Lse)
     {
         assert(!data->isContainedIntOrIImmed());
 
@@ -3988,8 +3988,11 @@ void CodeGen::genCodeForCmpXchg(GenTreeCmpXchg* treeNode)
 
     emitAttr dataSize = emitActualTypeSize(data);
 
-    if (m_compiler->compOpportunisticallyDependsOn(InstructionSet_Atomics))
+    if (m_compiler->compGetAtomicsImplForNode(treeNode) == Compiler::AtomicsImpl::Lse)
     {
+        // 'casal' has no immediate form, so lowering must not have contained the comparand.
+        assert(!comparand->isContained());
+
         // casal use the comparand as the target reg
         GetEmitter()->emitIns_Mov(INS_mov, dataSize, targetReg, comparandReg, /* canSkip */ true);
 

src/coreclr/jit/compiler.cpp

@@ -2017,6 +2017,68 @@ bool Compiler::notifyInstructionSetUsage(CORINFO_InstructionSet isa, bool suppor
     return info.compCompHnd->notifyInstructionSetUsage(isa, supported);
 }
 
+#ifdef TARGET_ARM64
+//------------------------------------------------------------------------
+// compGetAtomicsImpl: Decide how atomic operations are to be expanded in this method.
+//
+// Return Value:
+//    AtomicsImpl::Lse     - the Armv8.1 atomics are known to be available, use them directly.
+//    AtomicsImpl::LlSc    - they are known not to be available, the ldaxr/stlxr retry loops
+//                           have to be used.
+//    AtomicsImpl::Dynamic - they are not part of the baseline instruction set, but they may
+//                           still be available on the machine that ends up running this code.
+//                           The Interlocked APIs are left unexpanded so that the managed
+//                           `if (Lse.IsSupported)` check in their bodies gets inlined instead.
+//
+// Notes:
+//    The result is computed on demand and cached: asking the EE about instruction set support has
+//    side effects (for ReadyToRun it records a hard requirement on the method), so we must not do
+//    it for methods that contain no atomic operations at all.
+//
+//    'Dynamic' is what makes NativeAOT interesting here: its baseline is typically armv8-a, so
+//    compExactlyDependsOn(InstructionSet_Atomics) is false even though the vast majority of the
+//    hardware in the wild does implement LSE. Note that this state cannot occur for the runtime
+//    JIT, where the opportunistic and the exact answer always agree - see JitStressAtomicsLightUp
+//    for a way to force it for testing.
+//
+Compiler::AtomicsImpl Compiler::compGetAtomicsImpl()
+{
+    if (m_atomicsImpl != AtomicsImpl::Uninitialized)
+    {
+        return m_atomicsImpl;
+    }
+
+    AtomicsImpl impl;
+    if (!opts.compSupportsISA.HasInstructionSet(InstructionSet_Atomics))
+    {
+        impl = AtomicsImpl::LlSc;
+    }
+#ifdef DEBUG
+    else if (JitConfig.JitStressAtomicsLightUp() != 0)
+    {
+        // Pretend the atomics are only opportunistically available so that we exercise the
+        // managed light-up path. Note we deliberately do not report the ISA as used to the EE.
+        impl = AtomicsImpl::Dynamic;
+    }
+#endif
+    else if (compExactlyDependsOn(InstructionSet_Atomics))
+    {
+        impl = AtomicsImpl::Lse;
+    }
+    else
+    {
+        impl = AtomicsImpl::Dynamic;
+    }
+
+    JITDUMP("Atomics will be expanded as %s\n", (impl == AtomicsImpl::Lse)    ? "LSE"
+                                                : (impl == AtomicsImpl::LlSc) ? "ldaxr/stlxr loops"
+                                                                              : "managed light-up");
+
+    m_atomicsImpl = impl;
+    return impl;
+}
+#endif // TARGET_ARM64
+
 #ifdef PROFILING_SUPPORTED
 // A Dummy routine to receive Enter/Leave/Tailcall profiler callbacks.
 // These are used when DOTNET_JitEltHookEnabled=1

src/coreclr/jit/compiler.h

@@ -5418,6 +5418,7 @@ class Compiler
                           R2RARG(CORINFO_CONST_LOOKUP* entryPoint),
                           NamedIntrinsic*         pIntrinsicName,
                           bool*                   isSpecialIntrinsic = nullptr);
+    bool     impIsAtomicLightUpCandidate(NamedIntrinsic ni, CORINFO_SIG_INFO* sig, bool mustExpand);
     GenTree* impEstimateIntrinsic(CORINFO_METHOD_HANDLE method,
                                   CORINFO_SIG_INFO*     sig,
                                   CorInfoType           callJitType,
@@ -10936,6 +10937,46 @@ class Compiler
         return opts.compSupportsISA.HasInstructionSet(isa);
     }
 
+#ifdef TARGET_ARM64
+    // How the Interlocked/atomic operations (GT_CMPXCHG, GT_XADD, GT_XCHG, GT_XORR, GT_XAND)
+    // are to be expanded by codegen.
+    enum class AtomicsImpl : uint8_t
+    {
+        Uninitialized,
+        LlSc,    // ldaxr/stlxr retry loop
+        Lse,     // Armv8.1 single instruction atomics
+        Dynamic, // not in the baseline, but may still be available at run time
+    };
+
+    AtomicsImpl compGetAtomicsImpl();
+
+    //------------------------------------------------------------------------
+    // compGetAtomicsImplForNode: How a specific atomic node is to be expanded.
+    //
+    // The importer only ever creates an atomic node when it knows which sequence to use: in the
+    // AtomicsImpl::Dynamic case the Interlocked APIs are left unexpanded and the managed
+    // `if (Lse.IsSupported)` check in their bodies picks the arm, marking the node it creates
+    // with GTF_ATOMIC_LSE. So lowering, LSRA and codegen only ever see a concrete answer.
+    //
+    AtomicsImpl compGetAtomicsImplForNode(GenTree* node)
+    {
+        assert(node->OperIsAtomicOp());
+
+        if ((node->gtFlags & GTF_ATOMIC_LSE) != 0)
+        {
+            return AtomicsImpl::Lse;
+        }
+
+        AtomicsImpl impl = compGetAtomicsImpl();
+        return (impl == AtomicsImpl::Lse) ? AtomicsImpl::Lse : AtomicsImpl::LlSc;
+    }
+
+private:
+    AtomicsImpl m_atomicsImpl = AtomicsImpl::Uninitialized;
+
+public:
+#endif // TARGET_ARM64
+
 private:
 #ifdef DEBUG
     //------------------------------------------------------------------------

src/coreclr/jit/gentree.cpp

@@ -2770,6 +2770,15 @@ bool GenTree::Compare(GenTree* op1, GenTree* op2, bool swapOK)
         return false;
     }
 
+#ifdef TARGET_ARM64
+    // The two arms of an opportunistic atomics check hold otherwise identical nodes that must not
+    // be merged back together - the flag is what tells codegen which sequence to emit.
+    if (OperIsAtomicOp(oper) && (((op1->gtFlags ^ op2->gtFlags) & GTF_ATOMIC_LSE) != 0))
+    {
+        return false;
+    }
+#endif // TARGET_ARM64
+
     /* Sensible flags must be equal */
     if (op1->IsUnsigned() != op2->IsUnsigned())
     {

src/coreclr/jit/gentree.h

@@ -487,6 +487,10 @@ enum GenTreeFlags : unsigned
     GTF_IND_INITCLASS           = 0x00200000, // OperIsIndir() -- the indirection requires preceding static cctor
     GTF_IND_ALLOW_NON_ATOMIC    = 0x00100000, // GT_IND -- this memory access does not need to be atomic
 
+    // Set by the importer on arm64 for the atomic nodes created from the Lse intrinsics. Note this
+    // bit is not used by any of the GTF_IND_* flags, which the atomic opers otherwise share.
+    GTF_ATOMIC_LSE              = 0x04000000, // GT_XADD/GT_XAND/GT_XORR/GT_XCHG/GT_CMPXCHG -- use the Armv8.1 atomics
+
     // Represents flags that an indirection based on another indirection must preserve
     GTF_IND_MUST_PRESERVE_FLAGS = GTF_IND_VOLATILE | GTF_IND_UNALIGNED | GTF_IND_INITCLASS,
 

src/coreclr/jit/hwintrinsic.cpp

@@ -983,7 +983,7 @@ static const HWIntrinsicIsaRange hwintrinsicIsaRangeArray[] = {
     { FIRST_NI_Fp16, LAST_NI_Fp16 },                            // Fp16
     { FIRST_NI_Sha1, LAST_NI_Sha1 },                            // Sha1
     { FIRST_NI_Sha256, LAST_NI_Sha256 },                        // Sha256
-    { NI_Illegal, NI_Illegal },                                 //      Atomics
+    { FIRST_NI_Atomics, LAST_NI_Atomics },                      // Atomics
     { FIRST_NI_Vector, LAST_NI_Vector },                        // Vector64
     { FIRST_NI_Vector, LAST_NI_Vector },                        // Vector128
     { NI_Illegal, NI_Illegal },                                 //      VectorT

src/coreclr/jit/hwintrinsicarm64.cpp

@@ -94,6 +94,13 @@ CORINFO_InstructionSet Compiler::lookupInstructionSet(const char* className)
             return InstructionSet_Dp;
         }
     }
+    else if (className[0] == 'L')
+    {
+        if (strcmp(className, "Lse") == 0)
+        {
+            return InstructionSet_Atomics;
+        }
+    }
     else if (className[0] == 'R')
     {
         if (strcmp(className, "Rdm") == 0)
@@ -703,7 +710,7 @@ GenTree* Compiler::impSpecialIntrinsic(NamedIntrinsic        intrinsic,
     const int                 numArgs  = sig->numArgs;
 
     // The vast majority of "special" intrinsics are Vector64/Vector128 methods.
-    // The only exception is ArmBase.Yield which should be treated differently.
+    // The only exceptions are ArmBase.Yield and the Lse atomics, which are treated differently.
     if (intrinsic == NI_ArmBase_Yield)
     {
         assert(sig->numArgs == 0);
@@ -713,6 +720,80 @@ GenTree* Compiler::impSpecialIntrinsic(NamedIntrinsic        intrinsic,
         return gtNewScalarHWIntrinsicNode(TYP_VOID, intrinsic);
     }
 
+    if (isa == InstructionSet_Atomics)
+    {
+        // The Armv8.1 atomics reuse the nodes that back the Interlocked APIs; GTF_ATOMIC_LSE is
+        // what tells lowering and codegen to use the single instruction forms unconditionally.
+        //
+        // Note we must use simdBaseType rather than retType here: the latter has been widened to
+        // TYP_INT for the byte and halfword forms, which would give us an access of the wrong size.
+        assert(simdSize == 0);
+
+        genTreeOps oper;
+        switch (intrinsic)
+        {
+            case NI_Atomics_CompareAndSwap:
+                oper = GT_CMPXCHG;
+                break;
+            case NI_Atomics_LoadAdd:
+                oper = GT_XADD;
+                break;
+            case NI_Atomics_LoadClear:
+                oper = GT_XAND;
+                break;
+            case NI_Atomics_LoadSet:
+                oper = GT_XORR;
+                break;
+            case NI_Atomics_Swap:
+                oper = GT_XCHG;
+                break;
+            default:
+                unreached();
+        }
+
+        // Only "cas" and "swp" have byte and halfword forms.
+        const bool hasSmallForms = (oper == GT_CMPXCHG) || (oper == GT_XCHG);
+
+        // These are generic, so the type argument has to be checked rather than assumed: only an
+        // integer no wider than a pointer can be encoded. Anything else (floating point, an object
+        // reference, a struct) throws instead of silently accessing memory at the wrong width.
+        //
+        // Note the call is made before any argument is popped, and asks for a throwing expansion
+        // unconditionally: the managed bodies are recursive stubs, so falling back to a real call
+        // would simply recurse forever.
+        if (!varTypeIsIntegral(simdBaseType) || (genTypeSize(simdBaseType) > TARGET_POINTER_SIZE) ||
+            (varTypeIsSmall(simdBaseType) && !hasSmallForms))
+        {
+            return impUnsupportedNamedIntrinsic(CORINFO_HELP_THROW_TYPE_NOT_SUPPORTED, method, sig,
+                                                /* mustExpand */ true);
+        }
+
+        GenTree* comparand = nullptr;
+        if (oper == GT_CMPXCHG)
+        {
+            assert(sig->numArgs == 3);
+            comparand = impPopStack().val;
+
+            if (varTypeIsSmall(simdBaseType))
+            {
+                // Small types need the comparand to have its upper bits zeroed.
+                comparand = gtNewCastNode(genActualType(simdBaseType), comparand, /* uns */ false,
+                                          varTypeToUnsigned(simdBaseType));
+            }
+        }
+        else
+        {
+            assert(sig->numArgs == 2);
+        }
+
+        GenTree* value = impPopStack().val;
+        GenTree* addr  = impPopStack().val;
+
+        GenTree* node = gtNewAtomicNode(oper, simdBaseType, addr, value, comparand);
+        node->gtFlags |= GTF_ATOMIC_LSE;
+        return node;
+    }
+
     bool isScalar = (category == HW_Category_Scalar);
     assert(numArgs >= 0);