JIT: Use a "resumed?" indicator variable in runtime async instead of lvaAsyncContinuationArg != nullptr

[jakobbotsch]

We will need this to enable inlining, since lvaAsyncContinuationArg != nullptr is not a correct way to compute this value for inlinees.

Diffs are pretty mixed and not generally improvements since we have to generate stores to this local now. One possibility is that we use lvaAsyncContinuationArg != null for the root function and only use indicator locals in inlinees instead. This would also make things simpler for OSR.

[jakobbotsch]

cc @AndyAyersMS, this is the PR with the indicator variables we discussed

[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

Replaces the implicit "have we resumed?" check (lvaAsyncContinuationArg != nullptr) with an explicit boolean indicator local (lvaResumedIndicator). The root motivation is to enable async inlining, where the existing check is not valid for inlinees. Two new pseudo-args (AsyncResumedUse / AsyncResumedDef) are threaded through every async call so each callsite can read the indicator and the resumption path can write 1 back into it; the JIT machinery (lcl morph, address-of escape analysis, fgInline, local def visitors, diagnostic checks) is updated to recognize the new defined-via-address local.

Changes:

• Introduce lvaResumedIndicator (TYP_UBYTE) and initialize it in the prolog / OSR entry / loops, replacing the inline-special-cased continuation != null recipe at all uses.
• Add WellKnownArg::AsyncResumedUse / AsyncResumedDef plus gtCallGetDefinedAsyncResumedLclAddr, and propagate them through arg setup, IR sequencing, escape analysis, local-def visiting, and diagnostic checks.
• In AsyncTransformation, store 1 to the indicator on the resumption path, add common-def head-merging into the resumption switch, and add a shared "resumed" temp for shared finish-context handling.

Reviewed changes

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

Show a summary per file

File	Description
src/coreclr/jit/async.h	Declares new shared resumed var, helper signatures, and updated parameter lists for finish-context BB construction.
src/coreclr/jit/async.cpp	Core change: introduces indicator local, initializes it, threads new well-known args through calls, writes 1 on resumption, and implements common-def head merging.
src/coreclr/jit/compiler.h / compiler.cpp	Adds lvaResumedIndicator field and includes it in the OSR-local sanity check.
src/coreclr/jit/compiler.hpp	Extends VisitLocalDefs / VisitLocalDefNodes to also visit the async-resumed def on calls.
src/coreclr/jit/gentree.h / gentree.cpp	Adds the two new WellKnownArg values, their names, and gtCallGetDefinedAsyncResumedLclAddr.
src/coreclr/jit/morph.cpp	Wires the new well-known arg names into getWellKnownArgName.
src/coreclr/jit/importer.cpp / fginline.cpp	Skips the new pseudo-args in inline arg-info construction and prepend-statements.
src/coreclr/jit/lclmorph.cpp	Generalizes the callsite address-escape logic to recognize the resumed def (size 1) alongside retbuf, and sequences the def address.
src/coreclr/jit/lclvars.cpp	Minor cleanup: uses INDEBUG(...) macro for SetDefinedViaAddress.
src/coreclr/jit/fgdiagnostic.cpp	Extends linked-locals debug check to also push/expect the async-resumed def.

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

src/coreclr/jit/async.cpp:3618

• CreateSharedFinishContextHandlingBB creates the shared "resumed" local as TYP_UBYTE and also loads it as TYP_UBYTE. This value is later substituted into an already-morphed helper call whose placeholder arg was an int, so passing a 1-byte typed node here can lead to arg type / ABI mismatches. Load the local as TYP_INT (keeping the underlying local as TYP_UBYTE if desired) to match the rest of the async resumed flow, which consistently uses TYP_INT nodes.

    GenTree* resumed;
    if (invariantResumed == nullptr)
    {
        if (m_sharedFinishContextHandlingResumedVar == BAD_VAR_NUM)
        {
            m_sharedFinishContextHandlingResumedVar =
                m_compiler->lvaGrabTemp(false DEBUGARG("'resumed' for shared finish context handling"));
            m_compiler->lvaGetDesc(m_sharedFinishContextHandlingResumedVar)->lvType = TYP_UBYTE;
        }

        resumed = m_compiler->gtNewLclVarNode(m_sharedFinishContextHandlingResumedVar, TYP_UBYTE);
    }

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (2)

src/coreclr/jit/gentree.cpp:21662

• gtCallGetDefinedAsyncResumedLclAddr assumes the AsyncResumedDef argument node is always a raw GT_LCL_ADDR. However CallArg::GetNode() can be a GT_PUTARG_* after lowering (same reason gtCallGetDefinedRetBufLclAddr unwraps PUTARG_* and calls gtSkipReloadOrCopy). This can trigger the assert(node->OperIs(GT_LCL_ADDR) ...) in debug builds and mis-identify the def when the argument is wrapped.

Consider unwrapping GT_PUTARG_REG/GT_PUTARG_STK and skipping reload/copy, mirroring the retbuf helper.
src/coreclr/jit/compiler.hpp:4630

• VisitLocalDefs computes isEntire for the AsyncResumedDef call definition via lvaLclExactSize(...) == 1. That ignores the def's offset and doesn't use the same IsEntireAccess(...) logic as the retbuf call definition and LocalAddressVisitor (which sets GTF_VAR_USEASG based on IsEntireAccess).

To keep the call-def contract consistent (esp. if the resumed indicator ever becomes a field/offset def), compute isEntire using IsEntireAccess(lclNum, offs, ValueSize(1)).

        GenTreeLclVarCommon* asyncResumedLclAddr = comp->gtCallGetDefinedAsyncResumedLclAddr(call);
        if (asyncResumedLclAddr != nullptr)
        {
            bool isEntire = comp->lvaLclExactSize(asyncResumedLclAddr->GetLclNum()) == 1;

            RETURN_IF_ABORT(
                visitor(LocalDef(asyncResumedLclAddr, isEntire, asyncResumedLclAddr->GetLclOffs(), ValueSize(1))));
        }

[azure-pipelines]

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

[jakobbotsch]

cc @dotnet/jit-contrib PTAL @AndyAyersMS

Diffs. TP regressions are a bit larger than I'd like. Looking at the detailed TP diffs it looks like we no longer inline a bunch of the visitor functions passed into GenTree::VisitLocalDefs and GenTree::VisitLocalDefNodes with the new case added there. Not sure that there is much we can do about that here, hopefully PGO will decide to selectively inline the important cases.

[AndyAyersMS]

This is more complicated than I would have imagined.

[jakobbotsch]

This is more complicated than I would have imagined.

I think it's more or less the cost of making the dependencies of the async transformation visible in the IR. Before this PR the uses of this state just weren't visible -- e.g. on suspension we need to query information about "have we resumed before?". We did that by checking if there is a non-null continuation argument, which is not correct once inlining is in the picture. And then we furthermore need to model the fact that async resumptions conditionally define that state.

If we expanded the async transformation earlier it would be simpler since we could just introduce the IR there.

Now that we model this in IR I think it does subsume some of the previous stuff. The async transformation does
an analysis to be able to answer at every async call whether or not that call is reachable from a previous async call. I think it should be equivalent to checking whether the resumed indicator is known to be false, and that should fall out from normal constant propagation now.