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.

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

[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 76 out of 76 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (4)

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

• Same alignment concern as the safe path: Unsafe.As<byte, TAwaiter> + normal load assumes offset yields a naturally aligned address for TAwaiter, but continuation storage only guarantees pointer alignment. Use Unsafe.ReadUnaligned<TAwaiter> to avoid potential unaligned loads for custom awaiters with larger alignment requirements.
  src/coreclr/jit/importercalls.cpp:158
• This optimization introduces GT_CONTINUATION_MEMBER_OFFSET, which is only rewritten to a constant by the async transformation. If this rewrite triggers in a non-async method (compIsAsync() == false), the node will survive past import and there’s no later phase that can lower/codegen it. Guard this optimization on compIsAsync() to ensure it only runs when the async pipeline will replace the symbolic offset.

    CallArg* awaiterArg = call->gtArgs.GetUserArgByIndex(0);
    if (!varTypeIsStruct(awaiterArg->GetSignatureType()))
    {
        return;
    }

src/coreclr/jit/async.cpp:2483

• Same as the FIELD_LIST case: when storing the (possibly wide / highly-aligned) awaiter struct into the continuation, the indirection should be marked unaligned (continuation layout only guarantees pointer alignment). Without GTF_IND_UNALIGNED, some targets may generate aligned loads/stores for e.g. SIMD-containing awaiters.

        GenTree* continuation = m_compiler->gtNewLclvNode(GetNewContinuationVar(), TYP_REF);
        unsigned offset       = OFFSETOF__CORINFO_Continuation__data + layout.ContinuationMemberOffsets[memberIndex];
        GenTree* offsetNode   = m_compiler->gtNewIconNode((ssize_t)offset, TYP_I_IMPL);
        GenTree* address      = m_compiler->gtNewOperNode(GT_ADD, TYP_BYREF, continuation, offsetNode);
        GenTree* store        = m_compiler->gtNewStoreValueNode(awaiterLayout, address, awaiter, GTF_IND_NONFAULTING);
        LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_compiler, store));

src/tests/async/custom-struct-awaiters/custom-struct-awaiters.cs:19

• This test uses sync-over-async (Run().Wait()), which is discouraged in xUnit tests and can introduce deadlock risk if any synchronization context is involved. Prefer making the test method async Task and awaiting Run() directly.

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

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

• The exception message here says getAwaitReturnCall, but the code is in getAwaitAwaiterInContinuationCall. This makes diagnostics confusing when the READYTORUN runtime-JIT fallback is hit.

                        // Leave this method to runtime JIT that will be able to avoid the box
                        throw new RequiresRuntimeJitException($"getAwaitReturnCall: runtime-determined exact instantiation requires runtime JIT ({runtimeDeterminedResult})");
                    }

[jakobbotsch]

@VSadov I do wonder what we would expect to be the most common case here. In the beginning I thought it would be ICriticalNotifyCompletion, but now I am more and more leaning towards IValueTaskSource. What do you think?

[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})");