Remove boxing for awaited custom awaiters

[jakobbotsch]

This optimizes calls to AsyncHelpers.UnsafeAwaitAwaiter<TAwaiter>(TAwaiter) with struct awaiters to instead call a new function AsyncHelpers.UnsafeAwaitAwaiterInContinuation<TAwaiter>(int offset). The idea is that the JIT ensures that the awaiter will be present in the continuation at the specified offset. The later UnsafeOnCompleted call can then be done without any boxing by extracting it from the continuation.

This introduces a new GT_CONTINUATION_MEMBER_OFFSET which is used to solve the linking problem where we do not know the offset into the continuation until much later. The async transformation is responsible for replacing this node with a constant after it knows the offset.

I have a couple of use cases for GT_CONTINUATION_MEMBER_OFFSET in mind, so I have made the mechanism to represent the type of member easily expandable.

This PR also removes configurable continuation reuse. This is now unconditionally enabled, to simplify the expansion of GT_CONTINUATION_MEMBER_OFFSET nodes.

Fix #119842

Microbenchmark

About 10% improvement, and more importantly, avoids the box that async1 also avoids to gain parity.

using System;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;

public static class Program
{
    static Action s_continuation;
    static long s_value;

    public static void Main()
    {
        for (int i = 0; i < 10; i++)
        {
            for (int j = 0; j < 100; j++)
            {
                Task t = Foo(100);
                while (!t.IsCompleted)
                    s_continuation();
            }

            Thread.Sleep(100);
        }

        for (int i = 0; i < 50; i++)
        {
            Task t = Foo(10_000_000);
            while (!t.IsCompleted)
                s_continuation();
        }
    }

    private static async Task Foo(int n)
    {
        s_value = 0;

        Stopwatch timer = Stopwatch.StartNew();
        for (int i = 0; i < n; i++)
        {
            await new Awaiter(i);
        }

        if (n > 100)
            Console.WriteLine("Took {0} ms", timer.ElapsedMilliseconds);

        Trace.Assert(s_value == ((long)n * (n - 1)) / 2);
    }

    private struct Awaiter : ICriticalNotifyCompletion
    {
        public int X;

        public Awaiter(int x) => X = x;

        public bool IsCompleted => false;
        public Awaiter GetAwaiter() => this;
        public void GetResult() { }

        public void OnCompleted(Action continuation)
        {
        }

        public void UnsafeOnCompleted(Action continuation)
        {
            s_value += X;
            s_continuation = continuation;
        }
    }
}

-Took 323 ms
-Took 318 ms
-Took 320 ms
-Took 323 ms
-Took 322 ms
-Took 318 ms
-Took 320 ms
-Took 318 ms
-Took 317 ms
-Took 319 ms
+Took 290 ms
+Took 291 ms
+Took 292 ms
+Took 288 ms
+Took 292 ms
+Took 291 ms
+Took 290 ms
+Took 290 ms
+Took 287 ms
+Took 292 ms

[azure-pipelines]

Azure Pipelines:
Successfully started running 6 pipeline(s).
10 pipeline(s) were filtered out due to trigger conditions.
There may be pipelines that require an authorized user to comment /azp run to run.

[dotnet-policy-service]

Tagging subscribers to this area: @JulieLeeMSFT, @jakobbotsch
See info in area-owners.md if you want to be subscribed.

[Copilot]

Pull request overview

This PR extends CoreCLR runtime-async to avoid boxing custom struct awaiters on suspension by storing the awaiter value into the continuation object and later invoking OnCompleted/UnsafeOnCompleted by extracting the awaiter from that continuation using a JIT-computed byte offset. It introduces a symbolic “continuation member offset” node so the importer can reference the eventual offset before the async transformation finalizes the continuation layout, and adds a new JIT↔EE query to resolve the specialized helper method.

Changes:

• Add new CoreLib helpers to await via “awaiter-in-continuation” with an offset, and update suspension handling to call OnCompleted/UnsafeOnCompleted by reading the awaiter out of the continuation.
• Add GT_CONTINUATION_MEMBER_OFFSET plus continuation-member layout plumbing in the JIT async transformation; update importer to rewrite struct-await helper calls to the new helper and record the member offset.
• Extend the JIT↔EE interface and supporting tooling (VM, SuperPMI, ILC/ReadyToRun) to resolve the new helper method; add a regression test covering safe/unsafe custom awaiters.

Reviewed changes

Copilot reviewed 32 out of 32 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/tests/async/struct/struct.cs	Adds a regression test for safe and unsafe custom struct awaiters under runtime async.
src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.cs	Adds new private awaiter-in-continuation helpers and extraction-based OnCompleted/UnsafeOnCompleted paths.
src/coreclr/vm/jitinterface.h	Adds VM declaration for the new JIT↔EE query to resolve the awaiter-in-continuation helper.
src/coreclr/vm/jitinterface.cpp	Implements helper resolution + runtime lookup computation for the new awaiter-in-continuation call.
src/coreclr/vm/corelib.h	Adds CoreLibBinder method IDs for the new AsyncHelpers helper methods.
src/coreclr/tools/superpmi/superpmi/icorjitinfo.cpp	Plumbs the new ICorJitInfo method through SuperPMI.
src/coreclr/tools/superpmi/superpmi-shim-simple/icorjitinfo_generated.cpp	Updates generated shim to forward the new ICorJitInfo entrypoint.
src/coreclr/tools/superpmi/superpmi-shim-counter/icorjitinfo_generated.cpp	Updates generated counter shim to track/forward the new call.
src/coreclr/tools/superpmi/superpmi-shim-collector/icorjitinfo.cpp	Updates collector shim to record/replay the new query.
src/coreclr/tools/superpmi/superpmi-shared/methodcontext.h	Adds recording/replay declarations and a new packet kind for the query.
src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp	Implements record/dump/replay for the new query.
src/coreclr/tools/superpmi/superpmi-shared/lwmlist.h	Adds LWM map entry for the query.
src/coreclr/tools/superpmi/superpmi-shared/agnostic.h	Adds agnostic key struct for recording the query inputs.
src/coreclr/tools/Common/JitInterface/ThunkGenerator/ThunkInput.txt	Adds the new API to the JIT interface thunk generator input.
src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs	Adds ILC implementation to resolve the new helper method handle.
src/coreclr/tools/Common/JitInterface/CorInfoImpl_generated.cs	Adds generated unmanaged callback plumbing for the new API.
src/coreclr/tools/aot/jitinterface/jitinterface_generated.h	Updates NativeAOT jitinterface wrapper with the new callback.
src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/ReadyToRunCodegenCompilation.cs	Ensures R2R token availability for the new helper methods.
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs	Adds runtime stack-state fields + dispatcher logic to call awaiter continuation callbacks.
src/coreclr/jit/wellknownargs.h	Introduces a new well-known arg kind for the async awaiter pseudo-arg.
src/coreclr/jit/importercalls.cpp	Adds importer rewrite to swap struct awaiter calls to the new helper and attach symbolic continuation-member offset.
src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp	Adds wrapper method for the new JIT↔EE query.
src/coreclr/jit/ICorJitInfo_names_generated.h	Adds the new API name for logging/tracing.
src/coreclr/jit/gtstructs.h	Extends GenTreeVal struct mapping to include GT_CONTINUATION_MEMBER_OFFSET.
src/coreclr/jit/gtlist.h	Adds the GT_CONTINUATION_MEMBER_OFFSET node kind and documentation.
src/coreclr/jit/gentree.cpp	Updates pseudo-arg handling and debug printing for the new node.
src/coreclr/jit/compiler.h	Adds compiler state for tracking continuation members and the importer helper prototype.
src/coreclr/jit/async.h	Adds continuation-member abstractions and extends layout builder to allocate member offsets.
src/coreclr/jit/async.cpp	Implements continuation-member tracking, reserves storage in continuation layouts, and bashes symbolic offsets to constants.
src/coreclr/inc/jiteeversionguid.h	Bumps JIT↔EE version GUID due to interface change.
src/coreclr/inc/icorjitinfoimpl_generated.h	Adds the new virtual method to the generated ICorJitInfo impl surface.
src/coreclr/inc/corinfo.h	Adds the new ICorStaticInfo virtual method for helper resolution.

[Copilot]

Pull request overview

Copilot reviewed 33 out of 33 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 73 out of 73 changed files in this pull request and generated no new comments.

[jakobbotsch]

/azp run runtime-coreclr jitstress, runtime-coreclr libraries-jitstress

[azure-pipelines]

Azure Pipelines:
Successfully started running 2 pipeline(s).

[Copilot]

Pull request overview

Copilot reviewed 76 out of 76 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs:3759

• The exception message thrown from getAwaitAwaiterInContinuationCall still says "getAwaitReturnCall", which is misleading for diagnostics and log triage (this path is specifically for awaiter-in-continuation).

                        throw new RequiresRuntimeJitException($"getAwaitReturnCall: runtime-determined exact instantiation requires runtime JIT ({runtimeDeterminedResult})");

[Copilot]

Pull request overview

Copilot reviewed 76 out of 76 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

src/coreclr/inc/corinfo.h:3191

• New ICorStaticInfo virtual method should be appended at the end of the interface to preserve vtable slot ordering. Even with the JIT/EE GUID bump, keeping existing slots stable reduces risk for out-of-tree JITs/tools and matches the runtime convention for these interfaces. Consider moving getAwaitAwaiterInContinuationCall to the end of ICorStaticInfo (just before the class closes), rather than inserting it before the diagnostic-method block.

    // Get the method to use to await a struct awaiter that is stored inside the
    // continuation instead of being boxed. 'pResolvedToken' is the resolved
    // token of the AsyncHelpers.AwaitAwaiter/UnsafeAwaitAwaiter call site that
    // is being replaced, and 'isUnsafe' indicates whether the unsafe variant is
    // being replaced.
    //
    // Returns the method handle of the call to insert, or NULL if the
    // transformation cannot be performed. 'contextHandle' is set to the context
    // to use when inlining the call, exactly as getCallInfo would report it for
    // a direct call to it (it may be an approximate/shared instantiation when
    // 'instArg' requires a runtime lookup). 'instArg' is filled with the
    // (potentially runtime-looked-up) instantiation argument that must be
    // passed to the call.
    virtual CORINFO_METHOD_HANDLE getAwaitAwaiterInContinuationCall(
        CORINFO_METHOD_HANDLE callerHandle,
        CORINFO_RESOLVED_TOKEN* pResolvedToken,
        bool isUnsafe,
        CORINFO_CONTEXT_HANDLE* contextHandle,
        CORINFO_LOOKUP* instArg
    ) = 0;

[jakobbotsch]

Had some random failure here (went away on a clean rebuild), but I didn't get any information about what the problem was so I expanded this a bit.

[azure-pipelines]

Azure Pipelines:
Successfully started running 6 pipeline(s).
10 pipeline(s) were filtered out due to trigger conditions.
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

Copilot reviewed 75 out of 75 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 75 out of 75 changed files in this pull request and generated no new comments.

Suppressed comments (2)

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.cs:58

• Reading a struct awaiter out of the continuation via Unsafe.As<byte, TAwaiter> creates a typed ref from a byte ref, which can introduce alignment assumptions for TAwaiter (and is generally riskier for structs with stricter alignment). Prefer an explicit unaligned read when materializing the copy.
  src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.cs:115
• Same as above: extracting the struct awaiter via Unsafe.As<byte, TAwaiter> can rely on alignment that isn't guaranteed by the byte-based offset computation. Use an explicit unaligned read when copying the awaiter out of the continuation buffer.