Lay out resumption stubs next to async method variants (#128380)

In large enough composite compilations targeting ARM32, some async
methods will be placed too far apart from their resumption stubs, and
the offset doesn't fit into a reloc. To keep the offsets small (and
perhaps a small perf win for icache locality), layout the resumption
stubs next to the async methods.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: jtschuster

src/coreclr/tools/aot/ILCompiler.ReadyToRun.Tests/TestCases/R2RTestSuites.cs

@@ -244,6 +244,7 @@ static void Validate(ReadyToRunReader reader)
             Assert.True(R2RAssert.HasContinuationLayout(reader, "CaptureObjectAcrossAwait", out diag), diag);
             Assert.True(R2RAssert.HasContinuationLayout(reader, "CaptureMultipleRefsAcrossAwait", out diag), diag);
             Assert.True(R2RAssert.HasResumptionStubFixup(reader, "CaptureObjectAcrossAwait", out diag), diag);
+            Assert.True(R2RAssert.AsyncMethodsWithResumptionStubsAreAdjacent(reader, out diag), diag);
         }
     }
 
@@ -358,6 +359,7 @@ static void Validate(ReadyToRunReader reader)
             Assert.True(R2RAssert.HasResumptionStubFixup(reader, ".MultipleAwaitsWithRefs(", out diag), diag);
             Assert.True(R2RAssert.HasFixupKindCountOnMethod(reader, ReadyToRunFixupKind.ResumptionStubEntryPoint, ".MultipleAwaits(", 1, out diag), diag);
             Assert.True(R2RAssert.HasFixupKindCountOnMethod(reader, ReadyToRunFixupKind.ResumptionStubEntryPoint, ".MultipleAwaitsWithRefs(", 1, out diag), diag);
+            Assert.True(R2RAssert.AsyncMethodsWithResumptionStubsAreAdjacent(reader, out diag), diag);
         }
     }
 

src/coreclr/tools/aot/ILCompiler.ReadyToRun.Tests/TestCasesRunner/R2RResultChecker.cs

@@ -334,6 +334,64 @@ public static bool HasResumptionStub(ReadyToRunReader reader, string methodName,
         return found;
     }
 
+    /// <summary>
+    /// Returns true if every [ASYNC] method with a ResumptionStubEntryPoint fixup is followed
+    /// immediately in emitted code by its [RESUME] stub.
+    /// </summary>
+    public static bool AsyncMethodsWithResumptionStubsAreAdjacent(ReadyToRunReader reader, out string diagnostic)
+    {
+        var methodsByRva = GetAllMethods(reader)
+            .Select(method => (Method: method, EntryPointRva: GetEntryPointRva(method)))
+            .Where(entry => entry.EntryPointRva >= 0)
+            .OrderBy(entry => entry.EntryPointRva)
+            .ToList();
+
+        var failures = new List<string>();
+        int checkedMethodCount = 0;
+        for (int i = 0; i < methodsByRva.Count; i++)
+        {
+            ReadyToRunMethod method = methodsByRva[i].Method;
+            if (!HasSignaturePrefix(method, "[ASYNC]") ||
+                !HasFixupKind(method, ReadyToRunFixupKind.ResumptionStubEntryPoint))
+            {
+                continue;
+            }
+
+            checkedMethodCount++;
+            if (i + 1 == methodsByRva.Count)
+            {
+                failures.Add($"'{method.SignatureString}' at RVA 0x{methodsByRva[i].EntryPointRva:X} is the final method.");
+                continue;
+            }
+
+            ReadyToRunMethod nextMethod = methodsByRva[i + 1].Method;
+            if (!HasSignaturePrefix(nextMethod, "[RESUME]") ||
+                StripAsyncMethodPrefix(nextMethod.SignatureString) != StripAsyncMethodPrefix(method.SignatureString))
+            {
+                failures.Add(
+                    $"'{method.SignatureString}' at RVA 0x{methodsByRva[i].EntryPointRva:X} " +
+                    $"is followed by '{nextMethod.SignatureString}' at RVA 0x{methodsByRva[i + 1].EntryPointRva:X}.");
+            }
+        }
+
+        if (checkedMethodCount == 0)
+        {
+            diagnostic = "No [ASYNC] methods with ResumptionStubEntryPoint fixups were found.";
+            return false;
+        }
+
+        if (failures.Count != 0)
+        {
+            diagnostic =
+                $"Expected each [ASYNC] method with a ResumptionStubEntryPoint fixup to be followed by its [RESUME] stub, " +
+                $"but found {failures.Count} mismatch(es):\n  {string.Join("\n  ", failures)}";
+            return false;
+        }
+
+        diagnostic = $"Found {checkedMethodCount} [ASYNC] method(s), each followed by its [RESUME] stub.";
+        return true;
+    }
+
     /// <summary>
     /// Returns true if the R2R image contains at least one ContinuationLayout fixup.
     /// </summary>
@@ -441,6 +499,47 @@ public static bool HasFixupKind(ReadyToRunReader reader, ReadyToRunFixupKind kin
         return found;
     }
 
+    private static bool HasFixupKind(ReadyToRunMethod method, ReadyToRunFixupKind kind)
+    {
+        if (method.Fixups is null)
+            return false;
+
+        foreach (var cell in method.Fixups)
+        {
+            if (cell.Signature is not null && cell.Signature.FixupKind == kind)
+                return true;
+        }
+
+        return false;
+    }
+
+    private static bool HasSignaturePrefix(ReadyToRunMethod method, string prefix)
+        => method.SignaturePrefixes is not null && method.SignaturePrefixes.Contains(prefix);
+
+    private static int GetEntryPointRva(ReadyToRunMethod method)
+    {
+        foreach (RuntimeFunction runtimeFunction in method.RuntimeFunctions)
+        {
+            if (runtimeFunction.Id == method.EntryPointRuntimeFunctionId)
+                return runtimeFunction.StartAddress;
+        }
+
+        return -1;
+    }
+
+    private static string StripAsyncMethodPrefix(string signature)
+    {
+        const string AsyncPrefix = "[ASYNC] ";
+        const string ResumePrefix = "[RESUME] ";
+        if (signature.StartsWith(AsyncPrefix, StringComparison.Ordinal))
+            return signature.Substring(AsyncPrefix.Length);
+
+        if (signature.StartsWith(ResumePrefix, StringComparison.Ordinal))
+            return signature.Substring(ResumePrefix.Length);
+
+        return signature;
+    }
+
     /// <summary>
     /// Returns true if exactly one method whose signature contains <paramref name="methodName"/>
     /// has at least one fixup of the given kind.

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

@@ -384,7 +384,35 @@ public void InitializeDebugEHClauseInfos(DebugEHClauseInfo[] debugEHClauseInfos)
         public override int CompareToImpl(ISortableNode other, CompilerComparer comparer)
         {
             MethodWithGCInfo otherNode = (MethodWithGCInfo)other;
+            MethodDesc methodLayoutAnchor = GetMethodLayoutAnchor(_method, out int methodLayoutRank);
+            MethodDesc otherLayoutAnchor = GetMethodLayoutAnchor(otherNode._method, out int otherLayoutRank);
+
+            int result = comparer.Compare(methodLayoutAnchor, otherLayoutAnchor);
+            if (result != 0)
+            {
+                return result;
+            }
+
+            result = methodLayoutRank.CompareTo(otherLayoutRank);
+            if (result != 0)
+            {
+                return result;
+            }
+
             return comparer.Compare(_method, otherNode._method);
+
+            static MethodDesc GetMethodLayoutAnchor(MethodDesc method, out int layoutRank)
+            {
+                if (method is AsyncResumptionStub resumptionStub)
+                {
+                    Debug.Assert(resumptionStub.TargetMethod is not AsyncResumptionStub);
+                    layoutRank = 1;
+                    return resumptionStub.TargetMethod;
+                }
+
+                layoutRank = 0;
+                return method;
+            }
         }
 
         public void InitializeInliningInfo(MethodDesc[] inlinedMethods, NodeFactory factory)

src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/ReadyToRunTableManager.cs

@@ -112,7 +112,10 @@ public IEnumerable<IMethodNode> GetCompiledMethods(EcmaModule moduleToEnumerate,
                         int methodOnlyResult = comparer.Compare(x.Method, y.Method);
 
                         // Assert the two sorting techniques produce the same result unless there is a CustomSort applied
+                        // or MethodWithGCInfo ordering is keeping async resumption stubs next to their target methods.
                         Debug.Assert((nodeComparerResult == methodOnlyResult) ||
+                            x.Method is AsyncResumptionStub ||
+                            y.Method is AsyncResumptionStub ||
                             ((x is SortableDependencyNode sortableX && sortableX.CustomSort != Int32.MaxValue) ||
                              (y is SortableDependencyNode sortableY && sortableY.CustomSort != Int32.MaxValue)));
 #endif

src/coreclr/tools/aot/ILCompiler.Reflection.ReadyToRun/ReadyToRunReader.cs

@@ -1216,7 +1216,10 @@ private void MarkResumptionStubEntryPoints(bool[] isEntryPoint, ReadyToRunSectio
                             constrainedType: null,
                             instanceArgs: method.InstanceArgs,
                             signaturePrefixes: ["[RESUME]"],
-                            fixupOffset: null);
+                            fixupOffset: null)
+                        {
+                            RuntimeFunctionCount = 1,
+                        };
                         _instanceMethods.Add(new InstanceMethod(0, stubMethod));
                     }
                 }