Wasm: encode struct arg alignment in wasm signature strings

docs/design/coreclr/botr/readytorun-format.md

@@ -1077,6 +1077,7 @@ The string format is:
 | `d` | returns `f64` |
 | `V` | returns `v128` (a `Vector128<T>`, or a 16-byte `Vector<T>`) |
 | `S<N>` | struct return via hidden buffer, `N` is the struct size in bytes |
+| `A<N>` | as `S<N>`, but the struct's argument slot requires 16-byte alignment |
 
 **This pointer** (if the method has a `this` parameter):
 
@@ -1091,8 +1092,8 @@ The string format is:
 2. **Async continuation** (`a`): present for async calls.
 
 Note: the hidden return buffer pointer is **not** encoded in the signature string. Its
-presence is implied by the return type being `S<N>` — when the caller sees a struct return,
-it knows a hidden retbuf pointer argument is present in the Wasm parameter list.
+presence is implied by the return type being `S<N>` or `A<N>` — when the caller sees a struct
+return, it knows a hidden retbuf pointer argument is present in the Wasm parameter list.
 
 **Explicit parameters** (one token per parameter, in declaration order):
 
@@ -1104,6 +1105,7 @@ it knows a hidden retbuf pointer argument is present in the Wasm parameter list.
 | `d` | `f64` parameter |
 | `V` | `v128` parameter (a `Vector128<T>`, or a 16-byte `Vector<T>`, passed by value) |
 | `S<N>` | struct parameter passed by reference, `<N>` is the struct size in bytes |
+| `A<N>` | as `S<N>`, but the struct's argument slot requires 16-byte alignment |
 | `e` | empty struct parameter — elided from Wasm args but present in the string |
 | `<slot><E>` | multi-slot parameter passed by value, see below |
 
@@ -1127,15 +1129,23 @@ reads it. So `ll2VV4` is `i64`, `Int128`, `Vector128<T>`, `Vector512<T>`. The gr
 unambiguous because no other token places a digit after a slot character; `S<N>` consumes
 its own digits.
 
-These types are still *returned* through a hidden buffer, encoded as `S<N>` like any other
-aggregate. Limitation: a single digit carries both the slot count and the elevation factor,
-so an aggregate whose elevation differs from its slot count has no spelling. That includes
-one whose alignment is merely natural for its slot type, which would need count `N` with
-elevation 1. No such type exists in the Wasm ABI today.
+These types are still *returned* through a hidden buffer, encoded like any other aggregate:
+`A<N>`, since each of them clamps to 16-byte alignment. Limitation: a single digit carries
+both the slot count and the elevation factor, so an aggregate whose elevation differs from
+its slot count has no spelling. That includes one whose alignment is merely natural for its
+slot type, which would need count `N` with elevation 1. No such type exists in the Wasm ABI
+today.
 
 WasmAppBuilder does not emit or consume multi-slot tokens: they do not appear in
 `InternalCall` or `PInvoke` signatures.
 
+A struct's argument slot in the transition block is aligned to the struct's own alignment,
+clamped to `[8, 16]`. Because thunks are keyed by (and shared across modules by) the signature
+string alone, that alignment must be part of the encoding: a 32-byte struct of `long`s and a
+32-byte struct of `Int128`s have different argument layouts and cannot share a thunk. Structs
+whose alignment clamps to 16 (those containing `Int128`/`UInt128` or a 128-bit vector) use the
+`A<N>` token; all others use `S<N>`.
+
 **Suffix**:
 
 | Encoding | Meaning |

src/coreclr/inc/jiteeversionguid.h

@@ -37,11 +37,11 @@
 
 #include <minipal/guid.h>
 
-constexpr GUID JITEEVersionIdentifier = { /* aa3cece7-a5f9-4b4e-9309-852c8bd4bdb1 */
-    0xaa3cece7,
-    0xa5f9,
-    0x4b4e,
-    {0x93, 0x09, 0x85, 0x2c, 0x8b, 0xd4, 0xbd, 0xb1}
+constexpr GUID JITEEVersionIdentifier = { /* 5203056d-9af9-4dc8-a134-56a53f27e274 */
+    0x5203056d,
+    0x9af9,
+    0x4dc8,
+    {0xa1, 0x34, 0x56, 0xa5, 0x3f, 0x27, 0xe2, 0x74}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H

src/coreclr/tools/Common/Compiler/CompilerTypeSystemContext.Wasm.cs

@@ -13,8 +13,7 @@ namespace ILCompiler
     public partial class CompilerTypeSystemContext
     {
         private readonly object _structCacheLock = new object();
-        private readonly Dictionary<int, TypeDesc> _structsBySize = new Dictionary<int, TypeDesc>();
-        private readonly Dictionary<int, TypeDesc> _returnStructsBySize = new Dictionary<int, TypeDesc>();
+        private readonly Dictionary<(int Size, int Alignment), TypeDesc> _structsBySizeAndAlignment = new Dictionary<(int, int), TypeDesc>();
         private volatile TypeDesc _cachedEmptyStruct;
         private volatile TypeDesc _wasmV128Type;
         private volatile TypeDesc _wasmInt128Type;
@@ -90,61 +89,32 @@ public void CacheEmptyStruct(TypeDesc type)
         }
 
         /// <summary>
-        /// Caches a struct type by its element size, so RaiseSignature can retrieve a real
-        /// type of that size. Only the first struct encountered for a given size is retained.
+        /// Caches a struct type by its element size and argument alignment, so RaiseSignature can
+        /// retrieve a real type with the same Wasm argument layout. Only the first struct
+        /// encountered for a given (size, alignment) pair is retained.
         /// </summary>
-        public void CacheStructBySize(TypeDesc type)
+        public void CacheStructBySize(TypeDesc type, int alignment)
         {
             int size = type.GetElementSize().AsInt;
             if (size <= 0)
                 return;
 
             lock (_structCacheLock)
             {
-                _structsBySize.TryAdd(size, type);
+                _structsBySizeAndAlignment.TryAdd((size, alignment), type);
             }
         }
 
         /// <summary>
-        /// Caches a struct return type by size. Kept apart from the parameter cache because the two
-        /// classes are not interchangeable: a multi-slot type spells <c>S&lt;N&gt;</c> as a return but
-        /// re-lowers to its slot form as a parameter, so letting one answer for the other would give
-        /// an ordinary same-sized struct that type's larger alignment.
+        /// Gets a previously cached struct type of the specified byte size and argument alignment.
+        /// Returns null if no such struct has been cached.
+        /// Used by RaiseSignature to produce a roundtrippable type for the 'S&lt;N&gt;'/'A&lt;N&gt;' encodings.
         /// </summary>
-        public void CacheReturnStructBySize(TypeDesc type)
-        {
-            int size = type.GetElementSize().AsInt;
-            if (size <= 0)
-                return;
-
-            lock (_structCacheLock)
-            {
-                _returnStructsBySize.TryAdd(size, type);
-            }
-        }
-
-        /// <summary>
-        /// Gets a previously cached struct return type of the specified byte size, falling back to a
-        /// parameter of that size. Returns null if neither has been cached.
-        /// </summary>
-        public TypeDesc GetCachedReturnStructOfSize(int size)
-        {
-            lock (_structCacheLock)
-            {
-                if (_returnStructsBySize.TryGetValue(size, out TypeDesc result))
-                {
-                    return result;
-                }
-            }
-
-            return GetCachedStructOfSize(size);
-        }
-
-        public TypeDesc GetCachedStructOfSize(int size)
+        public TypeDesc GetCachedStructOfSize(int size, int alignment)
         {
             lock (_structCacheLock)
             {
-                if (_structsBySize.TryGetValue(size, out TypeDesc result))
+                if (_structsBySizeAndAlignment.TryGetValue((size, alignment), out TypeDesc result))
                     return result;
             }
 

src/coreclr/tools/Common/JitInterface/WasmLowering.cs

@@ -409,10 +409,11 @@ public static WasmValueType LowerType(TypeDesc type)
             _ => throw new InvalidOperationException($"Unknown signature char: {c}")
         };
 
-        private static int ParseStructSize(string sig, ref int pos)
+        private static int ParseStructSize(string sig, ref int pos, out int align)
         {
-            Debug.Assert(sig[pos] == 'S');
-            pos++; // skip 'S'
+            Debug.Assert(sig[pos] is 'S' or 'A');
+            align = (sig[pos] == 'A') ? WideStructArgAlignment : DefaultStructArgAlignment;
+            pos++; // skip 'S'/'A'
             int start = pos;
             while (pos < sig.Length && char.IsDigit(sig[pos]))
             {
@@ -421,6 +422,38 @@ private static int ParseStructSize(string sig, ref int pos)
             return int.Parse(sig.AsSpan(start, pos - start));
         }
 
+        // Wasm passes struct arguments in transition block slots aligned to the struct's own
+        // alignment, clamped to [8, 16]. This must match ArgIterator's Wasm32 case.
+        //
+        private const int DefaultStructArgAlignment = 8;
+        private const int WideStructArgAlignment = 16;
+
+        /// <summary>
+        /// Gets the alignment a struct argument or return type requires in the transition block.
+        /// Mirrors the Wasm32 case of <c>ArgIterator</c> and of the native <c>ArgIteratorTemplate</c>.
+        /// </summary>
+        public static int GetStructArgAlignment(TypeDesc type)
+        {
+            int alignment = ((DefType)type).InstanceFieldAlignment.AsInt;
+            return Math.Clamp(alignment, DefaultStructArgAlignment, WideStructArgAlignment);
+        }
+
+        /// <summary>
+        /// Returns true if <paramref name="c"/> starts a struct token ('S&lt;N&gt;' or 'A&lt;N&gt;').
+        /// </summary>
+        public static bool IsStructToken(char c) => c is 'S' or 'A';
+
+        // Appends the 'S<N>'/'A<N>' token for a struct passed by reference, and caches the type so
+        // RaiseSignature can recover a type with the same argument layout.
+        //
+        private static void AppendStructToken(StringBuilder sigBuilder, TypeDesc type)
+        {
+            int alignment = GetStructArgAlignment(type);
+            sigBuilder.Append(alignment == WideStructArgAlignment ? 'A' : 'S');
+            sigBuilder.Append(type.GetElementSize().AsInt);
+            ((CompilerTypeSystemContext)type.Context).CacheStructBySize(type, alignment);
+        }
+
         public static MethodSignature RaiseSignature(WasmSignature wasmSignature, TypeSystemContext context)
         {
             string sig = wasmSignature.SignatureString;
@@ -433,11 +466,11 @@ public static MethodSignature RaiseSignature(WasmSignature wasmSignature, TypeSy
                 returnType = context.GetWellKnownType(WellKnownType.Void);
                 pos++;
             }
-            else if (sig[pos] == 'S')
+            else if (sig[pos] is 'S' or 'A')
             {
-                int structSize = ParseStructSize(sig, ref pos);
-                returnType = ((CompilerTypeSystemContext)context).GetCachedReturnStructOfSize(structSize);
-                Debug.Assert(returnType is not null, $"No cached struct of size {structSize} for return type in signature '{sig}'");
+                int structSize = ParseStructSize(sig, ref pos, out int structAlign);
+                returnType = ((CompilerTypeSystemContext)context).GetCachedStructOfSize(structSize, structAlign);
+                Debug.Assert(returnType is not null, $"No cached struct of size {structSize} and alignment {structAlign} for return type in signature '{sig}'");
             }
             else
             {
@@ -496,11 +529,11 @@ public static MethodSignature RaiseSignature(WasmSignature wasmSignature, TypeSy
                     parameters.Add(((CompilerTypeSystemContext)context).GetWasmElevatedType(c, elevation));
                     pos += 2;
                 }
-                else if (c == 'S')
+                else if (IsStructToken(c))
                 {
-                    int structSize = ParseStructSize(sig, ref pos);
-                    TypeDesc cachedStruct = ((CompilerTypeSystemContext)context).GetCachedStructOfSize(structSize);
-                    Debug.Assert(cachedStruct is not null, $"No cached struct of size {structSize} for parameter in signature '{sig}'");
+                    int structSize = ParseStructSize(sig, ref pos, out int structAlign);
+                    TypeDesc cachedStruct = ((CompilerTypeSystemContext)context).GetCachedStructOfSize(structSize, structAlign);
+                    Debug.Assert(cachedStruct is not null, $"No cached struct of size {structSize} and alignment {structAlign} for parameter in signature '{sig}'");
                     parameters.Add(cachedStruct);
                 }
                 else
@@ -611,19 +644,7 @@ public static WasmSignature GetSignature(MethodSignature signature, LoweringFlag
                 {
                     hasReturnBuffer = true;
                     returnIsVoid = true;
-                    int returnSize = returnType.GetElementSize().AsInt;
-                    sigBuilder.Append('S');
-                    sigBuilder.Append(returnSize);
-
-                    // A multi-slot type spells 'S<N>' only as a return; as a parameter it re-lowers
-                    // to its slot form. Keep it in the return cache alone, so an ordinary same-sized
-                    // struct parameter does not raise with this type's larger alignment.
-                    CompilerTypeSystemContext returnContext = (CompilerTypeSystemContext)returnType.Context;
-                    returnContext.CacheReturnStructBySize(returnType);
-                    if (!TryGetMultiSegmentLayout(returnType, out _, out _))
-                    {
-                        returnContext.CacheStructBySize(returnType);
-                    }
+                    AppendStructToken(sigBuilder, returnType);
                 }
             }
             else if (loweredReturnType.IsVoid)
@@ -701,7 +722,6 @@ public static WasmSignature GetSignature(MethodSignature signature, LoweringFlag
                     }
 
                     // Struct that cannot be lowered to a single primitive — passed by reference
-                    int paramSize = paramType.GetElementSize().AsInt;
                     if (TryGetMultiSegmentLayout(paramType, out WasmValueType slotType, out int slotCount))
                     {
                         // Passed by value across several wasm parameters, matching the wasm C ABI.
@@ -719,9 +739,7 @@ public static WasmSignature GetSignature(MethodSignature signature, LoweringFlag
                     }
                     else
                     {
-                        sigBuilder.Append('S');
-                        sigBuilder.Append(paramSize);
-                        ((CompilerTypeSystemContext)paramType.Context).CacheStructBySize(paramType);
+                        AppendStructToken(sigBuilder, paramType);
                         result.Add(pointerType);
                     }
                 }

src/coreclr/tools/aot/ILCompiler.ReadyToRun.Tests/WasmArgumentLayoutTests.cs

@@ -283,9 +283,10 @@ public void NarrowVectorsAreNotMultiSlot(string vectorType, int expectedSize, st
     }
 
     /// <summary>
-    /// A multi-slot type is returned through a hidden buffer and so spells <c>S&lt;N&gt;</c>, but it
-    /// must not be remembered as the type that encoding raises to: it re-lowers to its multi-slot
-    /// form, so an ordinary same-sized struct parameter would raise with the wrong alignment.
+    /// A multi-slot type is returned through a hidden buffer, so it takes a struct token rather
+    /// than its slot form. Its argument alignment is part of that token, so an ordinary struct of
+    /// the same size — which is only 8-byte aligned — spells a different one and raises with its
+    /// own layout rather than the multi-slot type's.
     /// </summary>
     [Theory]
     [InlineData(Int128Type, 2)]
@@ -297,10 +298,13 @@ public void MultiSlotReturnDoesNotPoisonTheStructSizeCache(string typeName, int
 
         // Lower a method whose return is the multi-slot type; this is what caches by size.
         DefType multiSlot = InstantiateMultiSlotType(context, typeName);
-        WasmLowering.GetSignature(
+        WasmSignature returnSignature = WasmLowering.GetSignature(
             new MethodSignature(MethodSignatureFlags.Static, 0, multiSlot, Array.Empty<TypeDesc>()),
             WasmLowering.LoweringFlags.None);
 
+        // Every multi-slot type clamps to 16-byte alignment, so the return takes the wide token.
+        Assert.Equal($"A{multiSlot.InstanceFieldSize.AsInt}p", returnSignature.SignatureString);
+
         // Now an ordinary struct of the same size, which is only 8-byte aligned.
         MetadataType tupleType = fieldCount == 2
             ? context.SystemModule.GetType("System"u8, "ValueTuple`2"u8)
@@ -462,9 +466,10 @@ public void SingleFieldStructIsPassedAsTheTypeItWraps(string typeName, string ex
         Assert.Equal(expectedSignature, SignatureOf(context, wrapper));
         Assert.Equal(OffsetsOf(context, wrapped), OffsetsOf(context, wrapper));
 
-        // A second field makes it an ordinary aggregate, so it goes back to the struct ABI.
+        // A second field makes it an ordinary aggregate, so it goes back to the struct ABI. Two
+        // vector fields still leave it over-aligned, so it takes the wide-alignment token.
         DefType twoFields = MakeValueTuple(context, wrapped, wrapped);
-        Assert.Equal($"vlS{twoFields.InstanceFieldSize.AsInt}ip", SignatureOf(context, twoFields));
+        Assert.Equal($"vlA{twoFields.InstanceFieldSize.AsInt}ip", SignatureOf(context, twoFields));
     }
 
     /// <summary>
@@ -482,6 +487,29 @@ public void SingleFieldStructWrappingAMultiSlotScalarKeepsItsSlots()
         Assert.Equal(OffsetsOf(context, int128), OffsetsOf(context, wrapper));
     }
 
+    /// <summary>
+    /// A thunk is keyed by its signature string, so two structs of the same size must not spell the
+    /// same token unless they also agree on argument alignment. The over-aligned one is padded to 16
+    /// in the transition block and the other is not, so one thunk cannot describe both.
+    /// </summary>
+    [Fact]
+    public void SameSizedStructsOfDifferentAlignmentGetDifferentTokens()
+    {
+        ReadyToRunCompilerContext context = CreateWasmContext();
+        TypeDesc int64 = context.GetWellKnownType(WellKnownType.Int64);
+        DefType vector = InstantiateVector(context, Vector128OfT, WellKnownType.Int32);
+
+        DefType wideAligned = MakeValueTuple(context, vector, vector);
+        DefType slotAligned = MakeValueTuple(context, int64, int64, int64, int64);
+
+        Assert.Equal(wideAligned.InstanceFieldSize.AsInt, slotAligned.InstanceFieldSize.AsInt);
+        Assert.NotEqual(wideAligned.InstanceFieldAlignment.AsInt, slotAligned.InstanceFieldAlignment.AsInt);
+
+        Assert.Equal("vlA32ip", SignatureOf(context, wideAligned));
+        Assert.Equal("vlS32ip", SignatureOf(context, slotAligned));
+        Assert.NotEqual(OffsetsOf(context, wideAligned), OffsetsOf(context, slotAligned));
+    }
+
     private static DefType MakeValueTuple(ReadyToRunCompilerContext context, params TypeDesc[] fields) =>
         ((MetadataType)context.SystemModule.GetType(
             "System"u8, System.Text.Encoding.UTF8.GetBytes($"ValueTuple`{fields.Length}")))

src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/DependencyAnalysis/ReadyToRun/WasmImportThunk.cs

@@ -105,7 +105,7 @@ private bool HasGenericContextBeforeAsync
                 }
 
                 int pos = 1;
-                if (_wasmSignature.SignatureString[0] == 'S')
+                if (WasmLowering.IsStructToken(_wasmSignature.SignatureString[0]))
                 {
                     while ((pos < _wasmSignature.SignatureString.Length) && char.IsDigit(_wasmSignature.SignatureString[pos]))
                     {
@@ -161,7 +161,7 @@ protected override void EmitCode(NodeFactory factory, ref Wasm.WasmEmitter instr
 
             int[] offsets = new int[methodSignature.Length];
             bool[] isIndirectStructArg = new bool[methodSignature.Length];
-            bool hasRetBuffArg = _wasmSignature.SignatureString[0] == 'S';
+            bool hasRetBuffArg = WasmLowering.IsStructToken(_wasmSignature.SignatureString[0]);
 
             int argIndex = 0;
             int argOffset;