[NativeAOT] Emit Align8 flag for thread statics (#105931)

Author: filipnavara

src/coreclr/tools/Common/TypeSystem/Common/TargetDetails.cs

@@ -348,5 +348,17 @@ public int MaxHomogeneousAggregateElementCount
         /// CodeDelta - encapsulate the fact that ARM requires a thumb bit
         /// </summary>
         public int CodeDelta { get => (Architecture == TargetArchitecture.ARM) ? 1 : 0; }
+
+        /// <summary>
+        /// Encapsulates the fact that some architectures require 8-byte (larger than pointer
+        /// size) alignment on some value types and arrays.
+        /// </summary>
+        public bool SupportsAlign8
+        {
+            get
+            {
+                return Architecture is TargetArchitecture.ARM or TargetArchitecture.Wasm32;
+            }
+        }
     }
 }

src/coreclr/tools/Common/TypeSystem/Common/TypeSystemHelpers.cs

@@ -415,7 +415,7 @@ public static bool RequiresSlotUnification(this MethodDesc method)
         /// </summary>
         public static bool RequiresAlign8(this TypeDesc type)
         {
-            if (type.Context.Target.Architecture != TargetArchitecture.ARM && type.Context.Target.Architecture != TargetArchitecture.Wasm32)
+            if (!type.Context.Target.SupportsAlign8)
             {
                 return false;
             }

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/GCStaticEETypeNode.cs

@@ -20,11 +20,13 @@ public class GCStaticEETypeNode : DehydratableObjectNode, ISymbolDefinitionNode
     {
         private GCPointerMap _gcMap;
         private TargetDetails _target;
+        private bool _requiresAlign8;
 
-        public GCStaticEETypeNode(TargetDetails target, GCPointerMap gcMap)
+        public GCStaticEETypeNode(TargetDetails target, GCPointerMap gcMap, bool requiresAlign8)
         {
             _gcMap = gcMap;
             _target = target;
+            _requiresAlign8 = requiresAlign8;
         }
 
         protected override string GetName(NodeFactory factory) => this.GetMangledName(factory.NameMangler);
@@ -83,6 +85,13 @@ protected override ObjectData GetDehydratableData(NodeFactory factory, bool relo
             if (containsPointers)
                 flags |= (uint)EETypeFlags.HasPointersFlag;
 
+            if (_requiresAlign8)
+            {
+                // Mark the method table as non-value type that requires 8-byte alignment
+                flags |= (uint)EETypeFlagsEx.RequiresAlign8Flag;
+                flags |= (uint)EETypeElementType.Class << (byte)EETypeFlags.ElementTypeShift;
+            }
+
             dataBuilder.EmitUInt(flags);
 
             totalSize = Math.Max(totalSize, _target.PointerSize * 3); // minimum GC MethodTable size is 3 pointers

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/GCStaticsNode.cs

@@ -47,7 +47,8 @@ public static string GetMangledName(TypeDesc type, NameMangler nameMangler)
         private ISymbolNode GetGCStaticEETypeNode(NodeFactory factory)
         {
             GCPointerMap map = GCPointerMap.FromStaticLayout(_type);
-            return factory.GCStaticEEType(map);
+            bool requiresAlign8 = _type.GCStaticFieldAlignment.AsInt > factory.Target.PointerSize;
+            return factory.GCStaticEEType(map, requiresAlign8);
         }
 
         protected override DependencyList ComputeNonRelocationBasedDependencies(NodeFactory factory)

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/NativeLayoutVertexNode.cs

@@ -993,7 +993,8 @@ private static TypeDesc GetActualTemplateTypeForType(NodeFactory factory, TypeDe
         private ISymbolNode GetStaticsNode(NodeFactory context, out BagElementKind staticsBagKind)
         {
             MetadataType closestCanonDefType = (MetadataType)_type.GetClosestDefType().ConvertToCanonForm(CanonicalFormKind.Specific);
-            ISymbolNode symbol = context.GCStaticEEType(GCPointerMap.FromStaticLayout(closestCanonDefType));
+            bool requiresAlign8 = closestCanonDefType.GCStaticFieldAlignment.AsInt > context.Target.PointerSize;
+            ISymbolNode symbol = context.GCStaticEEType(GCPointerMap.FromStaticLayout(closestCanonDefType), requiresAlign8);
             staticsBagKind = BagElementKind.GcStaticDesc;
 
             return symbol;
@@ -1002,7 +1003,8 @@ private ISymbolNode GetStaticsNode(NodeFactory context, out BagElementKind stati
         private ISymbolNode GetThreadStaticsNode(NodeFactory context, out BagElementKind staticsBagKind)
         {
             MetadataType closestCanonDefType = (MetadataType)_type.GetClosestDefType().ConvertToCanonForm(CanonicalFormKind.Specific);
-            ISymbolNode symbol = context.GCStaticEEType(GCPointerMap.FromThreadStaticLayout(closestCanonDefType));
+            bool requiresAlign8 = closestCanonDefType.ThreadGcStaticFieldAlignment.AsInt > context.Target.PointerSize;
+            ISymbolNode symbol = context.GCStaticEEType(GCPointerMap.FromThreadStaticLayout(closestCanonDefType), requiresAlign8);
             staticsBagKind = BagElementKind.ThreadStaticDesc;
 
             return symbol;

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/NodeFactory.cs

@@ -241,9 +241,9 @@ private void CreateNodeCaches()
                 return new TypeThreadStaticIndexNode(type, null);
             });
 
-            _GCStaticEETypes = new NodeCache<GCPointerMap, GCStaticEETypeNode>((GCPointerMap gcMap) =>
+            _GCStaticEETypes = new NodeCache<(GCPointerMap, bool), GCStaticEETypeNode>(((GCPointerMap gcMap, bool requiresAlign8) key) =>
             {
-                return new GCStaticEETypeNode(Target, gcMap);
+                return new GCStaticEETypeNode(Target, key.gcMap, key.requiresAlign8);
             });
 
             _readOnlyDataBlobs = new NodeCache<ReadOnlyDataBlobKey, BlobNode>(key =>
@@ -810,11 +810,12 @@ public EmbeddedTrimmingDescriptorNode EmbeddedTrimmingDescriptor(EcmaModule modu
             return _embeddedTrimmingDescriptors.GetOrAdd(module);
         }
 
-        private NodeCache<GCPointerMap, GCStaticEETypeNode> _GCStaticEETypes;
+        private NodeCache<(GCPointerMap, bool), GCStaticEETypeNode> _GCStaticEETypes;
 
-        public ISymbolNode GCStaticEEType(GCPointerMap gcMap)
+        public ISymbolNode GCStaticEEType(GCPointerMap gcMap, bool requiredAlign8)
         {
-            return _GCStaticEETypes.GetOrAdd(gcMap);
+            requiredAlign8 &= Target.SupportsAlign8;
+            return _GCStaticEETypes.GetOrAdd((gcMap, requiredAlign8));
         }
 
         private NodeCache<ReadOnlyDataBlobKey, BlobNode> _readOnlyDataBlobs;

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/ThreadStaticsNode.cs

@@ -61,8 +61,9 @@ private ISymbolNode GetGCStaticEETypeNode(NodeFactory factory)
                     _inlined.GetOffsets(),
                     _inlined.GetSize(),
                     factory.Target.PointerSize);
+            bool requiresAlign8 = _type is not null && _type.ThreadGcStaticFieldAlignment.AsInt > factory.Target.PointerSize;
 
-            return factory.GCStaticEEType(map);
+            return factory.GCStaticEEType(map, requiresAlign8);
         }
 
         public override IEnumerable<DependencyListEntry> GetStaticDependencies(NodeFactory factory)

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/ILScanner.cs

@@ -850,8 +850,9 @@ public ScannedInlinedThreadStatics(NodeFactory factory, ImmutableArray<Dependenc
 
                         types.Add(t);
 
-                        // N.B. for ARM32, we would need to deal with > PointerSize alignments.
-                        //      GCStaticEEType does not currently set RequiresAlign8Flag
+                        // N.B. for ARM32, we would need to deal with > PointerSize alignments. We
+                        // currently don't support inlined thread statics on ARM32, regular GCStaticEEType
+                        // handles this with RequiresAlign8Flag
                         Debug.Assert(t.ThreadGcStaticFieldAlignment.AsInt <= factory.Target.PointerSize);
                         nextDataOffset = nextDataOffset.AlignUp(t.ThreadGcStaticFieldAlignment.AsInt);
 

src/tests/nativeaot/SmokeTests/UnitTests/BasicThreading.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -18,6 +19,9 @@ internal static int Run()
 
         ThreadStaticsTestWithTasks.Run();
 
+        if (ThreadStaticAlignmentTest.Run() != Pass)
+            return Fail;
+
         if (ThreadTest.Run() != Pass)
             return Fail;
 
@@ -187,6 +191,96 @@ public static void Run()
     }
 }
 
+class ThreadStaticAlignmentTest
+{
+    public static int Run()
+    {
+        // Check for 8-byte alignment requirement
+        if (RuntimeInformation.ProcessArchitecture is Architecture.Arm or Architecture.Wasm)
+        {
+            // Assume that these are allocated sequentially, use a padding object of size 12 (mod 8 is not 0)
+            // to move the alignment of the second AddressOfReturnArea in case the first is coincidentally aligned 8.
+            var ts1Addr = ThreadStaticAlignCheck1.returnArea.AddressOfReturnArea();
+            var p = new Padder();
+            var ts2Addr = ThreadStaticAlignCheck2.returnArea.AddressOfReturnArea();
+
+            if (((nint)ts1Addr) % 8 != 0)
+                return BasicThreading.Fail;
+            if (((nint)ts2Addr) % 8 != 0)
+                return BasicThreading.Fail;
+
+            return (int)typeof(ThreadStaticAlignmentTest).GetMethod("RunGeneric").MakeGenericMethod(GetAtom()).Invoke(null, []);
+        }
+
+        return BasicThreading.Pass;
+    }
+
+    public static int RunGeneric<T>()
+    {
+        // Assume that these are allocated sequentially, use a padding object of size 12 (mod 8 is not 0)
+        // to move the alignment of the second AddressOfReturnArea in case the first is coincidentally aligned 8.
+        var ts1Addr = ThreadStaticAlignCheck1<T>.returnArea.AddressOfReturnArea();
+        var p = new Padder();
+        var ts2Addr = ThreadStaticAlignCheck2<T>.returnArea.AddressOfReturnArea();
+
+        if (((nint)ts1Addr) % 8 != 0)
+            return BasicThreading.Fail;
+        if (((nint)ts2Addr) % 8 != 0)
+            return BasicThreading.Fail;
+
+        return BasicThreading.Pass;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static Type GetAtom() => typeof(Atom);
+
+    [InlineArray(3)]
+    private struct ReturnArea
+    {
+        private ulong buffer;
+
+        internal unsafe nint AddressOfReturnArea()
+        {
+            return (nint)Unsafe.AsPointer(ref buffer);
+        }
+    }
+
+    private class ThreadStaticAlignCheck1
+    {
+        [ThreadStatic]
+        [FixedAddressValueType]
+        internal static ReturnArea returnArea = default;
+    }
+
+    private class Padder
+    {
+        private object o1;
+    }
+
+    private class ThreadStaticAlignCheck2
+    {
+        [ThreadStatic]
+        [FixedAddressValueType]
+        internal static ReturnArea returnArea = default;
+    }
+
+    private class ThreadStaticAlignCheck1<T>
+    {
+        [ThreadStatic]
+        [FixedAddressValueType]
+        internal static ReturnArea returnArea = default;
+    }
+
+    private class ThreadStaticAlignCheck2<T>
+    {
+        [ThreadStatic]
+        [FixedAddressValueType]
+        internal static ReturnArea returnArea = default;
+    }
+
+    private class Atom { }
+}
+
 class ThreadTest
 {
     private static readonly List<Thread> s_startedThreads = new List<Thread>();