From 41bfc514f820e148038785ff615df03724814e23 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 13 May 2026 12:10:57 +0200 Subject: [PATCH 01/41] JIT: Use an indicator variable "did resume?" in runtime async We will need this to enable inlining. --- src/coreclr/jit/async.cpp | 320 ++++++++++++++++++++++++++----- src/coreclr/jit/async.h | 13 +- src/coreclr/jit/compiler.cpp | 2 +- src/coreclr/jit/compiler.h | 4 + src/coreclr/jit/compiler.hpp | 35 +++- src/coreclr/jit/fgdiagnostic.cpp | 12 +- src/coreclr/jit/fginline.cpp | 2 + src/coreclr/jit/gentree.cpp | 34 ++++ src/coreclr/jit/gentree.h | 2 + src/coreclr/jit/importer.cpp | 2 + src/coreclr/jit/lclmorph.cpp | 61 ++++-- src/coreclr/jit/lclvars.cpp | 4 +- src/coreclr/jit/morph.cpp | 4 + 13 files changed, 411 insertions(+), 84 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index cd4bb35514149e..c6badac1ea89ea 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -95,6 +95,9 @@ PhaseStatus Compiler::SaveAsyncContexts() lvaAsyncSynchronizationContextVar = lvaGrabTemp(false DEBUGARG("Async SynchronizationContext")); lvaGetDesc(lvaAsyncSynchronizationContextVar)->lvType = TYP_REF; + lvaResumedIndicator = lvaGrabTemp(false DEBUGARG("Async Resumed")); + lvaGetDesc(lvaResumedIndicator)->lvType = TYP_UBYTE; + if (opts.IsOSR()) { lvaGetDesc(lvaAsyncExecutionContextVar)->lvIsOSRLocal = true; @@ -183,7 +186,17 @@ PhaseStatus Compiler::SaveAsyncContexts() // Insert CaptureContexts call before the try (keep it before so the // try/finally can be removed if there is no exception side effects). // For OSR, we did this in the tier0 method. - if (!opts.IsOSR()) + if (opts.IsOSR()) + { + // In the OSR method we compute the initial value of the resumption indicator based on the continuation arg. + GenTree* continuation = gtNewLclVarNode(lvaAsyncContinuationArg, TYP_REF); + GenTree* null = gtNewNull(); + GenTree* contNeNull = gtNewOperNode(GT_NE, TYP_INT, continuation, null); + GenTree* storeIndicator = gtNewStoreLclVarNode(lvaResumedIndicator, contNeNull); + Statement* storeIndicatorStmt = fgNewStmtFromTree(storeIndicator); + fgInsertStmtAtBeg(fgFirstBB, storeIndicatorStmt); + } + else { GenTreeCall* captureCall = gtNewCallNode(CT_USER_FUNC, asyncInfo->captureContextsMethHnd, TYP_VOID); SetCallEntrypointForR2R(captureCall, this, asyncInfo->captureContextsMethHnd); @@ -203,21 +216,30 @@ PhaseStatus Compiler::SaveAsyncContexts() JITDUMP("Inserted capture\n"); DISPSTMT(captureStmt); + + // Also initialize resumed indicator var if it will not be initialized by the prolog. + BasicBlock* containingBlock = compIsForInlining() ? impInlineInfo->iciBlock : fgFirstBB; + bool inALoop = containingBlock->HasFlag(BBF_BACKWARD_JUMP); + bool isReturn = containingBlock->KindIs(BBJ_RETURN); + + if ((inALoop && !isReturn) || !impInlineRoot()->info.compInitMem) + { + GenTree* storeIndicator = gtNewStoreLclVarNode(lvaResumedIndicator, gtNewIconNode(0)); + Statement* storeIndicatorStmt = fgNewStmtFromTree(storeIndicator); + fgInsertStmtAtBeg(fgFirstBB, storeIndicatorStmt); + + JITDUMP("Inserted resumed indicator initialization\n"); + DISPSTMT(storeIndicatorStmt); + } + else + { + JITDUMP("Skipping zero init of resumed indicator due to compInitMem\n"); + } } // Insert RestoreContexts call in fault (exceptional case) // First argument: resumed = (continuation != null) - GenTree* resumed; - if (compIsForInlining()) - { - resumed = gtNewFalse(); - } - else - { - GenTree* continuation = gtNewLclvNode(lvaAsyncContinuationArg, TYP_REF); - GenTree* null = gtNewNull(); - resumed = gtNewOperNode(GT_NE, TYP_INT, continuation, null); - } + GenTree* resumed = gtNewLclvNode(lvaResumedIndicator, TYP_INT); GenTreeCall* restoreCall = gtNewCallNode(CT_USER_FUNC, asyncInfo->restoreContextsMethHnd, TYP_VOID); SetCallEntrypointForR2R(restoreCall, this, asyncInfo->restoreContextsMethHnd); @@ -332,15 +354,23 @@ void Compiler::AddContextArgsToAsyncCalls(BasicBlock* block) return WALK_CONTINUE; } - GenTreeCall* call = tree->AsCall(); - GenTree* execCtx = m_compiler->gtNewLclVarNode(m_compiler->lvaAsyncExecutionContextVar, TYP_REF); + GenTreeCall* call = tree->AsCall(); + GenTree* resumed = m_compiler->gtNewLclVarNode(m_compiler->lvaResumedIndicator, TYP_INT); + GenTree* resumedAddr = m_compiler->gtNewLclAddrNode(m_compiler->lvaResumedIndicator, 0); + GenTree* execCtx = m_compiler->gtNewLclVarNode(m_compiler->lvaAsyncExecutionContextVar, TYP_REF); GenTree* syncCtx = m_compiler->gtNewLclVarNode(m_compiler->lvaAsyncSynchronizationContextVar, TYP_REF); - JITDUMP("Adding exec context [%06u], sync context [%06u] to async call [%06u]\n", dspTreeID(execCtx), - dspTreeID(syncCtx), dspTreeID(call)); + JITDUMP( + "Adding resumed use [%06u], resumed def [%06u] exec context [%06u], sync context [%06u] to async call [%06u]\n", + dspTreeID(resumed), dspTreeID(resumedAddr), dspTreeID(execCtx), dspTreeID(syncCtx), dspTreeID(call)); call->gtArgs.PushFront(m_compiler, NewCallArg::Primitive(syncCtx).WellKnown(WellKnownArg::AsyncSynchronizationContext)); call->gtArgs.PushFront(m_compiler, NewCallArg::Primitive(execCtx).WellKnown(WellKnownArg::AsyncExecutionContext)); + call->gtArgs.PushFront(m_compiler, + NewCallArg::Primitive(resumedAddr).WellKnown(WellKnownArg::AsyncResumedDef)); + call->gtArgs.PushFront(m_compiler, NewCallArg::Primitive(resumed).WellKnown(WellKnownArg::AsyncResumedUse)); + + m_compiler->lvaGetDesc(m_compiler->lvaResumedIndicator)->lvHasLdAddrOp = true; return WALK_CONTINUE; } }; @@ -374,17 +404,7 @@ BasicBlock* Compiler::CreateReturnBB(unsigned* mergedReturnLcl) // Insert "restore" call CORINFO_ASYNC_INFO* asyncInfo = eeGetAsyncInfo(); - GenTree* resumed; - if (compIsForInlining()) - { - resumed = gtNewFalse(); - } - else - { - GenTree* continuation = gtNewLclvNode(lvaAsyncContinuationArg, TYP_REF); - GenTree* null = gtNewNull(); - resumed = gtNewOperNode(GT_NE, TYP_INT, continuation, null); - } + GenTree* resumed = gtNewLclvNode(lvaResumedIndicator, TYP_INT); GenTreeCall* restoreCall = gtNewCallNode(CT_USER_FUNC, asyncInfo->restoreContextsMethHnd, TYP_VOID); SetCallEntrypointForR2R(restoreCall, this, asyncInfo->restoreContextsMethHnd); @@ -688,11 +708,13 @@ PhaseStatus AsyncTransformation::Run() } } + GenTreeLclVarCommon* commonAsyncResumedDef = FindAndRemoveCommonAsyncResumedDef(); + CreateResumptionsAndSuspensions(); // After transforming all async calls we have created resumption blocks; // create the resumption switch. - CreateResumptionSwitch(); + CreateResumptionSwitch(commonAsyncResumedDef); m_compiler->fgInvalidateDfsTree(); @@ -2213,8 +2235,10 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc ? m_sharedFinishContextHandlingWithContinuationContextBB : m_sharedFinishContextHandlingWithoutContinuationContextBB; + CallArg* resumedArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse); CallArg* execContextArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncExecutionContext); CallArg* syncContextArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncSynchronizationContext); + assert((resumedArg != nullptr) && (execContextArg != nullptr)); assert((execContextArg != nullptr) && (syncContextArg != nullptr)); // Get the contexts from the call node: @@ -2222,6 +2246,17 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc // 2. For non-shared finish, just make sure it is a GT_LCL_VAR since we need to create // a use in a different block. // Also remove the nodes from the original block and the call args. + GenTree* resumed = resumedArg->GetNode(); + if (!resumed->IsInvariant() && !resumed->OperIs(GT_LCL_VAR)) + { + // We are moving resumed into a different BB so create a temp for it. + LIR::Use use(LIR::AsRange(callBlock), &resumedArg->NodeRef(), call); + use.ReplaceWithLclVar(m_compiler); + resumed = use.Def(); + } + LIR::AsRange(callBlock).Remove(resumed); + call->gtArgs.RemoveUnsafe(resumedArg); + GenTree* execContext = execContextArg->GetNode(); if (!execContext->OperIs(GT_LCL_VAR)) { @@ -2246,7 +2281,13 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc if (sharedFinish != nullptr) { - // Store the contexts to the shared locals that the shared finish block will take them from. + // Store the vars to the shared locals that the shared finish block will take them from. + if (m_sharedFinishContextHandlingResumedVar != BAD_VAR_NUM) + { + GenTree* storeResumed = m_compiler->gtNewStoreLclVarNode(m_sharedFinishContextHandlingResumedVar, resumed); + LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_compiler, storeResumed)); + } + if (m_sharedFinishContextHandlingExecContextVar != BAD_VAR_NUM) { GenTree* storeExecContext = @@ -2267,7 +2308,7 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc else { // Otherwise insert a new call - InsertFinishContextHandlingCall(suspendBB, layout, helper, execContext, syncContext); + InsertFinishContextHandlingCall(suspendBB, layout, helper, resumed, execContext, syncContext); // And return either via a new GT_RETURN_SUSPEND or via the shared return BB. if (m_sharedReturnBB != nullptr) @@ -2294,10 +2335,12 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc // void AsyncTransformation::RestoreContexts(BasicBlock* block, GenTreeCall* call, BasicBlock* suspendBB) { + CallArg* resumedArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse); CallArg* execContextArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncExecutionContext); CallArg* syncContextArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncSynchronizationContext); + assert((resumedArg != nullptr) == (execContextArg != nullptr)); assert((execContextArg != nullptr) == (syncContextArg != nullptr)); - if (execContextArg == nullptr) + if (resumedArg == nullptr) { JITDUMP(" Call [%06u] does not have async contexts; skipping restore on suspension\n", Compiler::dspTreeID(call)); @@ -2325,19 +2368,27 @@ void AsyncTransformation::RestoreContexts(BasicBlock* block, GenTreeCall* call, LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_compiler, restoreCall)); - // Replace resumedPlaceholder with actual "continuationParameter != null" arg + // Replace resumedPlaceholder with actual resumed arg + GenTree* resumed = resumedArg->GetNode(); + if (!resumed->IsInvariant() && !resumed->OperIs(GT_LCL_VAR)) + { + // We are moving resumed into a different BB so create a temp for it. + LIR::Use use(LIR::AsRange(block), &resumedArg->NodeRef(), call); + use.ReplaceWithLclVar(m_compiler); + resumed = use.Def(); + } + LIR::Use use; bool gotUse = LIR::AsRange(suspendBB).TryGetUse(resumedPlaceholder, &use); assert(gotUse); - GenTree* continuation = m_compiler->gtNewLclvNode(m_compiler->lvaAsyncContinuationArg, TYP_REF); - GenTree* null = m_compiler->gtNewNull(); - GenTree* resumed = m_compiler->gtNewOperNode(GT_NE, TYP_INT, continuation, null); - - LIR::AsRange(suspendBB).InsertBefore(resumedPlaceholder, LIR::SeqTree(m_compiler, resumed)); + LIR::AsRange(block).Remove(resumed); + LIR::AsRange(suspendBB).InsertBefore(resumedPlaceholder, resumed); use.ReplaceWith(resumed); LIR::AsRange(suspendBB).Remove(resumedPlaceholder); + call->gtArgs.RemoveUnsafe(resumedArg); + // Replace execContextPlaceholder with actual value GenTree* execContext = execContextArg->GetNode(); if (!execContext->OperIs(GT_LCL_VAR)) @@ -2502,6 +2553,8 @@ void AsyncTransformation::CreateResumption(BasicBlock* call RestoreFromDataOnResumption(layout, subLayout, resumeBB); } + StoreResumedDef(callBlock, call, resumeBB); + BasicBlock* storeResultBB = resumeBB; if (subLayout.NeedsException()) @@ -2580,6 +2633,70 @@ void AsyncTransformation::RestoreFromDataOnResumption(const ContinuationLayout& } } +//------------------------------------------------------------------------ +// AsyncTransformation::StoreResumedDef: +// Assign the resumed def to 1 from the resumption path. +// +// Parameters: +// callBlock - The basic block containing the async call +// call - The async call node +// resumeBB - The basic block to append IR to +// +void AsyncTransformation::StoreResumedDef(BasicBlock* callBlock, GenTreeCall* call, BasicBlock* resumeBB) +{ + CallArg* resumedDefArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedDef); + + if (resumedDefArg == nullptr) + { + return; + } + + GenTreeLclVarCommon* resumedDef = m_compiler->gtCallGetDefinedAsyncResumedLclAddr(call); + assert((resumedDef != nullptr) && (resumedDefArg->GetNode() == resumedDef)); + + StoreResumedDef(resumedDef, resumeBB); + + LIR::AsRange(callBlock).Remove(resumedDef); + call->gtArgs.RemoveUnsafe(resumedDefArg); +} + +//------------------------------------------------------------------------ +// AsyncTransformation::StoreResumedDef: +// Assign the resumed def to 1 in the specified block. +// +// Parameters: +// resumedDef - The local variable representing the resumed def +// block - The basic block to append IR to +// +void AsyncTransformation::StoreResumedDef(GenTreeLclVarCommon* resumedDef, BasicBlock* block) +{ + JITDUMP(" Have resume def [%06u] to store to\n", Compiler::dspTreeID(resumedDef)); + + LclVarDsc* varDsc = m_compiler->lvaGetDesc(resumedDef); + GenTree* store; + if ((resumedDef->GetLclOffs() == 0) && varDsc->TypeIs(TYP_UBYTE)) + { + store = m_compiler->gtNewStoreLclVarNode(resumedDef->GetLclNum(), m_compiler->gtNewIconNode(1)); + } + else + { + store = m_compiler->gtNewStoreLclFldNode(resumedDef->GetLclNum(), TYP_UBYTE, resumedDef->GetLclOffs(), + m_compiler->gtNewIconNode(1)); + m_compiler->lvaSetVarDoNotEnregister(resumedDef->GetLclNum() DEBUGARG(DoNotEnregisterReason::LocalField)); + } + + if (block->HasTerminator()) + { + LIR::AsRange(block).InsertBefore(block->lastNode(), LIR::SeqTree(m_compiler, store)); + } + else + { + LIR::AsRange(block).InsertAtEnd(LIR::SeqTree(m_compiler, store)); + } + + JITDUMP(" Created store [%06u] to set resumed def to 1\n", Compiler::dspTreeID(store)); +} + //------------------------------------------------------------------------ // AsyncTransformation::RethrowExceptionOnResumption: // Create IR that checks for an exception and rethrows it at the original @@ -2980,6 +3097,7 @@ void AsyncTransformation::CreateSharedReturnBB() // Parameters: // helper - The type of helper to call // layout - The continuation layout +// invariantResumed - Tree node to clone and use for "resumed" computation // execContextMayVary - If true, callers may use different execution // contexts, and thus we need a local to allow it to vary. // syncContextMayVary - If true, callers may use different synchronization @@ -2990,6 +3108,7 @@ void AsyncTransformation::CreateSharedReturnBB() // BasicBlock* AsyncTransformation::CreateSharedFinishContextHandlingBB(SuspensionContextHelper helper, const ContinuationLayout& layout, + GenTree* invariantResumed, bool execContextMayVary, bool syncContextMayVary) { @@ -3007,6 +3126,23 @@ BasicBlock* AsyncTransformation::CreateSharedFinishContextHandlingBB(SuspensionC block->SetFlags(BBF_PROF_WEIGHT); } + 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_REF; + } + + resumed = m_compiler->gtNewLclVarNode(m_sharedFinishContextHandlingResumedVar, TYP_INT); + } + else + { + resumed = m_compiler->gtCloneExpr(invariantResumed); + } + unsigned execContextLclNum; if (execContextMayVary) { @@ -3041,7 +3177,8 @@ BasicBlock* AsyncTransformation::CreateSharedFinishContextHandlingBB(SuspensionC syncContextLclNum = m_compiler->lvaAsyncSynchronizationContextVar; } - InsertFinishContextHandlingCall(block, layout, helper, m_compiler->gtNewLclvNode(execContextLclNum, TYP_REF), + InsertFinishContextHandlingCall(block, layout, helper, resumed, + m_compiler->gtNewLclvNode(execContextLclNum, TYP_REF), m_compiler->gtNewLclvNode(syncContextLclNum, TYP_REF)); return block; @@ -3055,12 +3192,14 @@ BasicBlock* AsyncTransformation::CreateSharedFinishContextHandlingBB(SuspensionC // block - Block that should contain the call (inserted at the end) // layout - The continuation layout // helper - The type of helper +// resumed - The resumed tree to pass to the helper // execContext - The execution context tree to pass to the helper // syncContext - The synchronization context tree to pass to the helper // void AsyncTransformation::InsertFinishContextHandlingCall(BasicBlock* block, const ContinuationLayout& layout, SuspensionContextHelper helper, + GenTree* resumed, GenTree* execContext, GenTree* syncContext) { @@ -3157,11 +3296,7 @@ void AsyncTransformation::InsertFinishContextHandlingCall(BasicBlock* gotUse = LIR::AsRange(block).TryGetUse(resumedPlaceholder, &use); assert(gotUse); - GenTree* continuation = m_compiler->gtNewLclvNode(m_compiler->lvaAsyncContinuationArg, TYP_REF); - GenTree* null = m_compiler->gtNewNull(); - GenTree* resumed = m_compiler->gtNewOperNode(GT_NE, TYP_INT, continuation, null); - - LIR::AsRange(block).InsertBefore(resumedPlaceholder, LIR::SeqTree(m_compiler, resumed)); + LIR::AsRange(block).InsertBefore(resumedPlaceholder, resumed); use.ReplaceWith(resumed); LIR::AsRange(block).Remove(resumedPlaceholder); @@ -3185,6 +3320,65 @@ void AsyncTransformation::InsertFinishContextHandlingCall(BasicBlock* DISPTREERANGE(LIR::AsRange(block), finishCall); } +//------------------------------------------------------------------------ +// AsyncTransformation::FindAndRemoveCommonAsyncResumedDef: +// If all async calls define the same async resumption indicator variable, +// then remove the def from all calls and return it. +// +// Returns: +// The common def, or null if there is no common def. +// +GenTreeLclVarCommon* AsyncTransformation::FindAndRemoveCommonAsyncResumedDef() +{ + if (m_states.size() <= 1) + { + return nullptr; + } + + bool hasCommonDef = true; + GenTreeLclVarCommon* commonDef = nullptr; + unsigned numWithCommonDef = 0; + + for (const AsyncState& state : m_states) + { + GenTreeLclVarCommon* def = m_compiler->gtCallGetDefinedAsyncResumedLclAddr(state.Call); + if (def == nullptr) + { + continue; + } + + if ((commonDef == nullptr) || GenTree::Compare(def, commonDef)) + { + commonDef = def; + numWithCommonDef++; + } + else + { + hasCommonDef = false; + } + } + + if (!hasCommonDef || (numWithCommonDef <= 1)) + { + return nullptr; + } + + JITDUMP(" Found common async resumed def node:\n"); + DISPTREE(commonDef); + + for (const AsyncState& state : m_states) + { + CallArg* arg = state.Call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedDef); + if (arg != nullptr) + { + LIR::AsRange(state.CallBlock).Remove(arg->GetNode()); + state.Call->gtArgs.RemoveUnsafe(arg); + } + } + + return commonDef; +} + //------------------------------------------------------------------------ // AsyncTransformation::CreateResumptionsAndSuspensions: // Walk all recorded async states and create the suspension and resumption @@ -3205,8 +3399,10 @@ void AsyncTransformation::CreateResumptionsAndSuspensions() unsigned numSharedSuspensionsWithContinuationContext = 0; unsigned numSharedSuspensionsWithoutContinuationContext = 0; - bool execContextMayVary = false; - bool syncContextMayVary = false; + bool resumedMayVary = false; + bool execContextMayVary = false; + bool syncContextMayVary = false; + GenTree* invariantResumed = nullptr; for (const AsyncState& state : m_states) { @@ -3228,6 +3424,26 @@ void AsyncTransformation::CreateResumptionsAndSuspensions() // unnecessary additional register moves. This is a common case. if (helper != SuspensionContextHelper::None) { + CallArg* resumedArg = state.Call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse); + assert(resumedArg != nullptr); + GenTree* resumed = resumedArg->GetNode(); + + if (resumed->IsInvariant() || resumed->OperIs(GT_LCL_VAR)) + { + if ((invariantResumed == nullptr) || GenTree::Compare(invariantResumed, resumed)) + { + invariantResumed = resumed; + } + else + { + resumedMayVary = true; + } + } + else + { + resumedMayVary = true; + } + CallArg* execContextArg = state.Call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncExecutionContext); CallArg* syncContextArg = state.Call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncSynchronizationContext); @@ -3248,7 +3464,8 @@ void AsyncTransformation::CreateResumptionsAndSuspensions() numSharedSuspensionsWithContinuationContext); m_sharedFinishContextHandlingWithContinuationContextBB = CreateSharedFinishContextHandlingBB(SuspensionContextHelper::WithContinuationContext, *sharedLayout, - execContextMayVary, syncContextMayVary); + resumedMayVary ? nullptr : invariantResumed, execContextMayVary, + syncContextMayVary); } if (numSharedSuspensionsWithoutContinuationContext > 1) @@ -3258,7 +3475,8 @@ void AsyncTransformation::CreateResumptionsAndSuspensions() numSharedSuspensionsWithoutContinuationContext); m_sharedFinishContextHandlingWithoutContinuationContextBB = CreateSharedFinishContextHandlingBB(SuspensionContextHelper::WithoutContinuationContext, *sharedLayout, - execContextMayVary, syncContextMayVary); + resumedMayVary ? nullptr : invariantResumed, execContextMayVary, + syncContextMayVary); } } @@ -3363,7 +3581,7 @@ ContinuationLayoutBuilder* ContinuationLayoutBuilder::CreateSharedLayout(Compile // Create the IR for the entry of the function that checks the continuation // and dispatches on its state number. // -void AsyncTransformation::CreateResumptionSwitch() +void AsyncTransformation::CreateResumptionSwitch(GenTreeLclVarCommon* commonAsyncResumedDef) { m_compiler->fgCreateNewInitBB(); BasicBlock* newEntryBB = m_compiler->fgFirstBB; @@ -3461,6 +3679,12 @@ void AsyncTransformation::CreateResumptionSwitch() resumingEdge->setLikelihood(0); newEntryBB->GetFalseEdge()->setLikelihood(1); + if (commonAsyncResumedDef != nullptr) + { + // If we have a common async resumption def, then we do a manual head merge to move it into the switch block + StoreResumedDef(commonAsyncResumedDef, resumingEdge->getDestinationBlock()); + } + if (m_compiler->doesMethodHavePatchpoints()) { JITDUMP(" Method has patch points...\n"); diff --git a/src/coreclr/jit/async.h b/src/coreclr/jit/async.h index ce25cd991362cf..cc360e2bfedddb 100644 --- a/src/coreclr/jit/async.h +++ b/src/coreclr/jit/async.h @@ -366,7 +366,8 @@ class AsyncTransformation // saves/restores and then suspend. BasicBlock* m_sharedFinishContextHandlingWithContinuationContextBB = nullptr; BasicBlock* m_sharedFinishContextHandlingWithoutContinuationContextBB = nullptr; - // Variables that shared suspension finishing BBs take the exec/sync contexts in + // Variables that shared suspension finishing BBs take the resumed/exec/sync contexts in + unsigned m_sharedFinishContextHandlingResumedVar = BAD_VAR_NUM; unsigned m_sharedFinishContextHandlingExecContextVar = BAD_VAR_NUM; unsigned m_sharedFinishContextHandlingSyncContextVar = BAD_VAR_NUM; @@ -454,6 +455,8 @@ class AsyncTransformation void RestoreFromDataOnResumption(const ContinuationLayout& layout, const ContinuationLayoutBuilder& subLayout, BasicBlock* resumeBB); + void StoreResumedDef(BasicBlock* callBlock, GenTreeCall* call, BasicBlock* resumeBB); + void StoreResumedDef(GenTreeLclVarCommon* resumedDef, BasicBlock* block); BasicBlock* RethrowExceptionOnResumption(BasicBlock* block, const ContinuationLayout& layout, BasicBlock* resumeBB); void CopyReturnValueOnResumption(GenTreeCall* call, const CallDefinitionInfo& callDefInfo, @@ -479,16 +482,20 @@ class AsyncTransformation void CreateSharedReturnBB(); BasicBlock* CreateSharedFinishContextHandlingBB(SuspensionContextHelper helper, const ContinuationLayout& layout, + GenTree* invariantResumed, bool execContextMayVary, bool syncContextMayVary); void InsertFinishContextHandlingCall(BasicBlock* block, const ContinuationLayout& layout, SuspensionContextHelper helper, + GenTree* resumed, GenTree* execContext, GenTree* syncContext); bool ReuseContinuations(); - void CreateResumptionsAndSuspensions(); - void CreateResumptionSwitch(); + + GenTreeLclVarCommon* FindAndRemoveCommonAsyncResumedDef(); + void CreateResumptionsAndSuspensions(); + void CreateResumptionSwitch(GenTreeLclVarCommon* commonAsyncResumedDef); public: AsyncTransformation(Compiler* comp) diff --git a/src/coreclr/jit/compiler.cpp b/src/coreclr/jit/compiler.cpp index 24b980ccc8f755..308e6d465f7a6b 100644 --- a/src/coreclr/jit/compiler.cpp +++ b/src/coreclr/jit/compiler.cpp @@ -10432,7 +10432,7 @@ bool Compiler::lvaIsOSRLocal(unsigned varNum) { // Sanity check for promoted fields of OSR locals. // - if ((varNum >= info.compLocalsCount) && (varNum != lvaMonAcquired) && + if ((varNum >= info.compLocalsCount) && (varNum != lvaMonAcquired) && (varNum != lvaResumedIndicator) && (varNum != lvaAsyncExecutionContextVar) && (varNum != lvaAsyncSynchronizationContextVar)) { assert(varDsc->lvIsStructField); diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h index 6d83b3ff0030e5..05ea91e7fed5f0 100644 --- a/src/coreclr/jit/compiler.h +++ b/src/coreclr/jit/compiler.h @@ -3987,6 +3987,7 @@ class Compiler bool gtIsTypeof(GenTree* tree, CORINFO_CLASS_HANDLE* handle = nullptr); GenTreeLclVarCommon* gtCallGetDefinedRetBufLclAddr(GenTreeCall* call); + GenTreeLclVarCommon* gtCallGetDefinedAsyncResumedLclAddr(GenTreeCall* call); //------------------------------------------------------------------------- // Functions to display the trees @@ -4259,6 +4260,9 @@ class Compiler // Variable representing async continuation argument passed. unsigned lvaAsyncContinuationArg = BAD_VAR_NUM; + // Variable representing "have we resumed?" for async methods + unsigned lvaResumedIndicator = BAD_VAR_NUM; + #if defined(DEBUG) && defined(TARGET_XARCH) unsigned lvaReturnSpCheck = BAD_VAR_NUM; // Stores SP to confirm it is not corrupted on return. diff --git a/src/coreclr/jit/compiler.hpp b/src/coreclr/jit/compiler.hpp index 03b5b33273ebaa..3fcd3901445be5 100644 --- a/src/coreclr/jit/compiler.hpp +++ b/src/coreclr/jit/compiler.hpp @@ -4568,15 +4568,25 @@ GenTree::VisitResult GenTree::VisitLocalDefs(Compiler* comp, TVisitor visitor) } if (OperIs(GT_CALL)) { - GenTreeCall* call = AsCall(); - GenTreeLclVarCommon* lclAddr = comp->gtCallGetDefinedRetBufLclAddr(call); - if (lclAddr != nullptr) + GenTreeCall* call = AsCall(); + + 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)))); + } + + GenTreeLclVarCommon* retBufLclAddr = comp->gtCallGetDefinedRetBufLclAddr(call); + if (retBufLclAddr != nullptr) { unsigned storeSize = comp->typGetObjLayout(AsCall()->gtRetClsHnd)->GetSize(); - bool isEntire = storeSize == comp->lvaLclExactSize(lclAddr->GetLclNum()); + bool isEntire = storeSize == comp->lvaLclExactSize(retBufLclAddr->GetLclNum()); - return visitor(LocalDef(lclAddr, isEntire, lclAddr->GetLclOffs(), ValueSize(storeSize))); + return visitor(LocalDef(retBufLclAddr, isEntire, retBufLclAddr->GetLclOffs(), ValueSize(storeSize))); } } @@ -4611,11 +4621,18 @@ GenTree::VisitResult GenTree::VisitLocalDefNodes(Compiler* comp, TVisitor visito } if (OperIs(GT_CALL)) { - GenTreeCall* call = AsCall(); - GenTreeLclVarCommon* lclAddr = comp->gtCallGetDefinedRetBufLclAddr(call); - if (lclAddr != nullptr) + GenTreeCall* call = AsCall(); + + GenTreeLclVarCommon* asyncResumedLclAddr = comp->gtCallGetDefinedAsyncResumedLclAddr(call); + if (asyncResumedLclAddr != nullptr) + { + RETURN_IF_ABORT(visitor(asyncResumedLclAddr)); + } + + GenTreeLclVarCommon* retBufLclAddr = comp->gtCallGetDefinedRetBufLclAddr(call); + if (retBufLclAddr != nullptr) { - return visitor(lclAddr); + return visitor(retBufLclAddr); } } diff --git a/src/coreclr/jit/fgdiagnostic.cpp b/src/coreclr/jit/fgdiagnostic.cpp index 172b1ccc3b1609..86be4639ad8771 100644 --- a/src/coreclr/jit/fgdiagnostic.cpp +++ b/src/coreclr/jit/fgdiagnostic.cpp @@ -3815,7 +3815,8 @@ void Compiler::fgDebugCheckLinkedLocals() if (ShouldLink(node)) { if ((user != nullptr) && user->IsCall() && - (node == m_compiler->gtCallGetDefinedRetBufLclAddr(user->AsCall()))) + ((node == m_compiler->gtCallGetDefinedRetBufLclAddr(user->AsCall())) || + (node == m_compiler->gtCallGetDefinedAsyncResumedLclAddr(user->AsCall())))) { } else @@ -3826,7 +3827,14 @@ void Compiler::fgDebugCheckLinkedLocals() if (node->IsCall()) { - GenTree* defined = m_compiler->gtCallGetDefinedRetBufLclAddr(node->AsCall()); + GenTree* defined = m_compiler->gtCallGetDefinedAsyncResumedLclAddr(node->AsCall()); + if (defined != nullptr) + { + assert(ShouldLink(defined)); + m_locals.Push(defined); + } + + defined = m_compiler->gtCallGetDefinedRetBufLclAddr(node->AsCall()); if (defined != nullptr) { assert(ShouldLink(defined)); diff --git a/src/coreclr/jit/fginline.cpp b/src/coreclr/jit/fginline.cpp index dda8a3582d7a2c..56d35f7575e026 100644 --- a/src/coreclr/jit/fginline.cpp +++ b/src/coreclr/jit/fginline.cpp @@ -2320,6 +2320,8 @@ Statement* Compiler::fgInlinePrependStatements(InlineInfo* inlineInfo) case WellKnownArg::AsyncContinuation: case WellKnownArg::AsyncExecutionContext: case WellKnownArg::AsyncSynchronizationContext: + case WellKnownArg::AsyncResumedUse: + case WellKnownArg::AsyncResumedDef: continue; case WellKnownArg::InstParam: argInfo = inlineInfo->inlInstParamArgInfo; diff --git a/src/coreclr/jit/gentree.cpp b/src/coreclr/jit/gentree.cpp index aeaf0182c43130..267720afba11ff 100644 --- a/src/coreclr/jit/gentree.cpp +++ b/src/coreclr/jit/gentree.cpp @@ -1621,6 +1621,8 @@ bool CallArgs::GetCustomRegister(Compiler* comp, CorInfoCallConvExtension cc, We case WellKnownArg::StackArrayLocal: case WellKnownArg::AsyncExecutionContext: case WellKnownArg::AsyncSynchronizationContext: + case WellKnownArg::AsyncResumedUse: + case WellKnownArg::AsyncResumedDef: // These are pseudo-args; they are not actual arguments, but we // reuse the argument mechanism to represent them as arbitrary uses // that are later expanded out. @@ -14521,6 +14523,10 @@ const char* Compiler::gtGetWellKnownArgNameForArgMsg(WellKnownArg arg) return "exec ctx"; case WellKnownArg::AsyncSynchronizationContext: return "sync ctx"; + case WellKnownArg::AsyncResumedUse: + return "resumed"; + case WellKnownArg::AsyncResumedDef: + return "resumed def"; case WellKnownArg::WasmShadowStackPointer: return "wasm sp"; case WellKnownArg::WasmPortableEntryPoint: @@ -20941,6 +20947,34 @@ GenTreeLclVarCommon* Compiler::gtCallGetDefinedRetBufLclAddr(GenTreeCall* call) return node->AsLclVarCommon(); } +//------------------------------------------------------------------------ +// gtCallGetDefinedAsyncResumedLclAddr: +// Get the tree corresponding to the address of the async resumed indicator that this call defines. +// +// Parameters: +// call - The call node +// +// Returns: +// A tree representing the address of a local. +// +GenTreeLclVarCommon* Compiler::gtCallGetDefinedAsyncResumedLclAddr(GenTreeCall* call) +{ + if (!call->IsAsync()) + { + return nullptr; + } + + CallArg* arg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedDef); + if (arg == nullptr) + { + return nullptr; + } + + GenTree* node = arg->GetNode(); + assert(node->OperIs(GT_LCL_ADDR) && lvaGetDesc(node->AsLclVarCommon())->IsDefinedViaAddress()); + return node->AsLclVarCommon(); +} + //------------------------------------------------------------------------ // ParseArrayAddress: Rehydrate the array and index expression from ARR_ADDR. // diff --git a/src/coreclr/jit/gentree.h b/src/coreclr/jit/gentree.h index 1897b43014f17d..5bb371e3817eda 100644 --- a/src/coreclr/jit/gentree.h +++ b/src/coreclr/jit/gentree.h @@ -4779,6 +4779,8 @@ enum class WellKnownArg : unsigned RuntimeMethodHandle, AsyncExecutionContext, AsyncSynchronizationContext, + AsyncResumedUse, + AsyncResumedDef, WasmShadowStackPointer, WasmPortableEntryPoint }; diff --git a/src/coreclr/jit/importer.cpp b/src/coreclr/jit/importer.cpp index a83b16e1591cb9..d53e3eb4859276 100644 --- a/src/coreclr/jit/importer.cpp +++ b/src/coreclr/jit/importer.cpp @@ -13474,6 +13474,8 @@ void Compiler::impInlineInitVars(InlineInfo* pInlineInfo) case WellKnownArg::AsyncContinuation: case WellKnownArg::AsyncExecutionContext: case WellKnownArg::AsyncSynchronizationContext: + case WellKnownArg::AsyncResumedUse: + case WellKnownArg::AsyncResumedDef: // These do not appear in the table of inline arg info; do not include them continue; case WellKnownArg::InstParam: diff --git a/src/coreclr/jit/lclmorph.cpp b/src/coreclr/jit/lclmorph.cpp index 8f86759cdd284d..bc6a319802f660 100644 --- a/src/coreclr/jit/lclmorph.cpp +++ b/src/coreclr/jit/lclmorph.cpp @@ -113,6 +113,15 @@ class LocalSequencer final : public GenTreeVisitor // void SequenceCall(GenTreeCall* call) { + if (call->IsAsync()) + { + GenTreeLclVarCommon* asyncResumedDef = m_compiler->gtCallGetDefinedAsyncResumedLclAddr(call); + if (asyncResumedDef != nullptr) + { + MoveNodeToEnd(asyncResumedDef); + } + } + if (call->IsOptimizingRetBufAsLocal()) { // Correct the point at which the definition of the retbuf local appears. @@ -1503,31 +1512,47 @@ class LocalAddressVisitor final : public GenTreeVisitor GenTreeFlags defFlag = GTF_EMPTY; GenTreeCall* callUser = (user != nullptr) && user->IsCall() ? user->AsCall() : nullptr; bool escapeAddr = true; - if (m_compiler->opts.compJitOptimizeStructHiddenBuffer && (callUser != nullptr) && - m_compiler->IsValidLclAddr(lclNum, val.Offset())) - { - // We will only attempt this optimization for locals that do not - // later turn into indirections. - bool isSuitableLocal = - varTypeIsStruct(varDsc) && !m_compiler->lvaIsImplicitByRefLocal(lclNum) && - (!varDsc->lvIsStructField || !m_compiler->lvaIsImplicitByRefLocal(varDsc->lvParentLcl)); + if ((callUser != nullptr) && m_compiler->IsValidLclAddr(lclNum, val.Offset())) + { + unsigned defSize = UINT_MAX; + if (callUser->gtArgs.HasRetBuffer() && (val.Node() == callUser->gtArgs.GetRetBufferArg()->GetNode())) + { + // We will only attempt this optimization for locals that do not + // later turn into indirections. + bool isSuitableLocal = + m_compiler->opts.compJitOptimizeStructHiddenBuffer && varTypeIsStruct(varDsc) && + !m_compiler->lvaIsImplicitByRefLocal(lclNum) && + (!varDsc->lvIsStructField || !m_compiler->lvaIsImplicitByRefLocal(varDsc->lvParentLcl)); #ifdef TARGET_X86 - if (m_compiler->lvaIsArgAccessedViaVarArgsCookie(lclNum)) + if (m_compiler->lvaIsArgAccessedViaVarArgsCookie(lclNum)) + { + isSuitableLocal = false; + } +#endif // TARGET_X86 + + if (isSuitableLocal) + { + m_compiler->lvaSetHiddenBufferStructArg(lclNum); + callUser->gtCallMoreFlags |= GTF_CALL_M_RETBUFFARG_LCLOPT; + defSize = m_compiler->typGetObjLayout(callUser->gtRetClsHnd)->GetSize(); + } + } + else if (callUser->IsAsync()) { - isSuitableLocal = false; + CallArg* asyncResumedDef = callUser->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedDef); + if ((asyncResumedDef != nullptr) && (val.Node() == asyncResumedDef->GetNode())) + { + defSize = 1; + } } -#endif // TARGET_X86 - if (isSuitableLocal && callUser->gtArgs.HasRetBuffer() && - (val.Node() == callUser->gtArgs.GetRetBufferArg()->GetNode())) + if (defSize != UINT_MAX) { - m_compiler->lvaSetHiddenBufferStructArg(lclNum); + INDEBUG(varDsc->SetDefinedViaAddress(true)); escapeAddr = false; - callUser->gtCallMoreFlags |= GTF_CALL_M_RETBUFFARG_LCLOPT; - defFlag = GTF_VAR_DEF; + defFlag = GTF_VAR_DEF; - if ((val.Offset() != 0) || - (varDsc->lvExactSize() != m_compiler->typGetObjLayout(callUser->gtRetClsHnd)->GetSize())) + if ((val.Offset() != 0) || (varDsc->lvExactSize() != defSize)) { defFlag |= GTF_VAR_USEASG; } diff --git a/src/coreclr/jit/lclvars.cpp b/src/coreclr/jit/lclvars.cpp index adf573deda4674..443e756e4ad879 100644 --- a/src/coreclr/jit/lclvars.cpp +++ b/src/coreclr/jit/lclvars.cpp @@ -2243,9 +2243,7 @@ void Compiler::lvaSetHiddenBufferStructArg(unsigned varNum) { LclVarDsc* varDsc = lvaGetDesc(varNum); -#ifdef DEBUG - varDsc->SetDefinedViaAddress(true); -#endif + INDEBUG(varDsc->SetDefinedViaAddress(true)); if (varDsc->lvPromoted) { diff --git a/src/coreclr/jit/morph.cpp b/src/coreclr/jit/morph.cpp index 771732a23624b0..fab998a1790e14 100644 --- a/src/coreclr/jit/morph.cpp +++ b/src/coreclr/jit/morph.cpp @@ -633,6 +633,10 @@ const char* getWellKnownArgName(WellKnownArg arg) return "AsyncExecutionContext"; case WellKnownArg::AsyncSynchronizationContext: return "AsyncSynchronizationContext"; + case WellKnownArg::AsyncResumedUse: + return "AsyncResumedUse"; + case WellKnownArg::AsyncResumedDef: + return "AsyncResumedDef"; case WellKnownArg::WasmShadowStackPointer: return "WasmShadowStackPointer"; case WellKnownArg::WasmPortableEntryPoint: From 80c8b105d5e99bc314e028bcd36ad6195917516d Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 13 May 2026 18:31:03 +0200 Subject: [PATCH 02/41] Fix --- src/coreclr/jit/async.cpp | 27 +++++++-------------------- 1 file changed, 7 insertions(+), 20 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index c6badac1ea89ea..b63b17054189cf 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -3335,32 +3335,22 @@ GenTreeLclVarCommon* AsyncTransformation::FindAndRemoveCommonAsyncResumedDef() return nullptr; } - bool hasCommonDef = true; - GenTreeLclVarCommon* commonDef = nullptr; - unsigned numWithCommonDef = 0; + GenTreeLclVarCommon* commonDef = nullptr; for (const AsyncState& state : m_states) { GenTreeLclVarCommon* def = m_compiler->gtCallGetDefinedAsyncResumedLclAddr(state.Call); if (def == nullptr) { - continue; + return nullptr; } - if ((commonDef == nullptr) || GenTree::Compare(def, commonDef)) + if ((commonDef != nullptr) && !GenTree::Compare(def, commonDef)) { - commonDef = def; - numWithCommonDef++; + return nullptr; } - else - { - hasCommonDef = false; - } - } - if (!hasCommonDef || (numWithCommonDef <= 1)) - { - return nullptr; + commonDef = def; } JITDUMP(" Found common async resumed def node:\n"); @@ -3369,11 +3359,8 @@ GenTreeLclVarCommon* AsyncTransformation::FindAndRemoveCommonAsyncResumedDef() for (const AsyncState& state : m_states) { CallArg* arg = state.Call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedDef); - if (arg != nullptr) - { - LIR::AsRange(state.CallBlock).Remove(arg->GetNode()); - state.Call->gtArgs.RemoveUnsafe(arg); - } + LIR::AsRange(state.CallBlock).Remove(arg->GetNode()); + state.Call->gtArgs.RemoveUnsafe(arg); } return commonDef; From ca5ef77e0f00b63640037b521452d13b340190b9 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 13 May 2026 18:31:55 +0200 Subject: [PATCH 03/41] Comment --- src/coreclr/jit/async.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index b63b17054189cf..10421c8370c08e 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -3668,7 +3668,9 @@ void AsyncTransformation::CreateResumptionSwitch(GenTreeLclVarCommon* commonAsyn if (commonAsyncResumedDef != nullptr) { - // If we have a common async resumption def, then we do a manual head merge to move it into the switch block + // If we have a common async resumption def (common), then we do a + // manual head merge to move it into the switch block to avoid storing + // it in every resumption. StoreResumedDef(commonAsyncResumedDef, resumingEdge->getDestinationBlock()); } From cca3fc554030efb93100ac21c51a5b848c13cc09 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 13 May 2026 18:37:21 +0200 Subject: [PATCH 04/41] Copilot feedback --- src/coreclr/jit/async.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 10421c8370c08e..30e1a6495aa6ac 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -3133,10 +3133,10 @@ BasicBlock* AsyncTransformation::CreateSharedFinishContextHandlingBB(SuspensionC { m_sharedFinishContextHandlingResumedVar = m_compiler->lvaGrabTemp(false DEBUGARG("'resumed' for shared finish context handling")); - m_compiler->lvaGetDesc(m_sharedFinishContextHandlingResumedVar)->lvType = TYP_REF; + m_compiler->lvaGetDesc(m_sharedFinishContextHandlingResumedVar)->lvType = TYP_UBYTE; } - resumed = m_compiler->gtNewLclVarNode(m_sharedFinishContextHandlingResumedVar, TYP_INT); + resumed = m_compiler->gtNewLclVarNode(m_sharedFinishContextHandlingResumedVar, TYP_UBYTE); } else { From 6564c125bc9235e4220ee6889de1aec7853445e1 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 15 May 2026 13:32:11 +0200 Subject: [PATCH 05/41] Generalize --- src/coreclr/jit/fgdiagnostic.cpp | 33 ++++++++++++--------------- src/coreclr/jit/lclmorph.cpp | 24 +++++++------------ src/coreclr/jit/liveness.cpp | 26 ++++++++++++--------- src/coreclr/jit/morph.cpp | 8 +++---- src/coreclr/jit/promotionliveness.cpp | 23 +++++++++++++++---- 5 files changed, 60 insertions(+), 54 deletions(-) diff --git a/src/coreclr/jit/fgdiagnostic.cpp b/src/coreclr/jit/fgdiagnostic.cpp index 86be4639ad8771..112720bbe14406 100644 --- a/src/coreclr/jit/fgdiagnostic.cpp +++ b/src/coreclr/jit/fgdiagnostic.cpp @@ -3814,12 +3814,7 @@ void Compiler::fgDebugCheckLinkedLocals() GenTree* node = *use; if (ShouldLink(node)) { - if ((user != nullptr) && user->IsCall() && - ((node == m_compiler->gtCallGetDefinedRetBufLclAddr(user->AsCall())) || - (node == m_compiler->gtCallGetDefinedAsyncResumedLclAddr(user->AsCall())))) - { - } - else + if ((user == nullptr) || !user->IsCall() || !IsDefinedByCall(user->AsCall(), node)) { m_locals.Push(node); } @@ -3827,23 +3822,25 @@ void Compiler::fgDebugCheckLinkedLocals() if (node->IsCall()) { - GenTree* defined = m_compiler->gtCallGetDefinedAsyncResumedLclAddr(node->AsCall()); - if (defined != nullptr) - { - assert(ShouldLink(defined)); - m_locals.Push(defined); - } + auto linkDefs = [&](GenTree* def) { + assert(ShouldLink(def)); + m_locals.Push(def); + return GenTree::VisitResult::Continue; + }; - defined = m_compiler->gtCallGetDefinedRetBufLclAddr(node->AsCall()); - if (defined != nullptr) - { - assert(ShouldLink(defined)); - m_locals.Push(defined); - } + node->VisitLocalDefNodes(m_compiler, linkDefs); } return WALK_CONTINUE; } + + bool IsDefinedByCall(GenTreeCall* call, GenTree* node) + { + auto defIsNode = [=](GenTree* def) { + return node == def ? GenTree::VisitResult::Abort : GenTree::VisitResult::Continue; + }; + return call->VisitLocalDefNodes(m_compiler, defIsNode) == GenTree::VisitResult::Abort; + } }; DebugLocalSequencer seq(this); diff --git a/src/coreclr/jit/lclmorph.cpp b/src/coreclr/jit/lclmorph.cpp index bc6a319802f660..be403a634b8100 100644 --- a/src/coreclr/jit/lclmorph.cpp +++ b/src/coreclr/jit/lclmorph.cpp @@ -102,31 +102,23 @@ class LocalSequencer final : public GenTreeVisitor } //------------------------------------------------------------------- - // SequenceCall: Post-process a call that may define a local. + // SequenceCall: Post-process a call that may define locals. // // Arguments: // call - the call // // Remarks: - // calls may also define a local that we would like to see - // after all other operands of the call have been evaluated. + // calls may also define locals that we would like to see after all + // other operands of the call have been evaluated. // void SequenceCall(GenTreeCall* call) { - if (call->IsAsync()) - { - GenTreeLclVarCommon* asyncResumedDef = m_compiler->gtCallGetDefinedAsyncResumedLclAddr(call); - if (asyncResumedDef != nullptr) - { - MoveNodeToEnd(asyncResumedDef); - } - } + auto moveToEnd = [&](GenTree* def) { + MoveNodeToEnd(def); + return GenTree::VisitResult::Continue; + }; - if (call->IsOptimizingRetBufAsLocal()) - { - // Correct the point at which the definition of the retbuf local appears. - MoveNodeToEnd(m_compiler->gtCallGetDefinedRetBufLclAddr(call)); - } + call->VisitLocalDefNodes(m_compiler, moveToEnd); } //------------------------------------------------------------------- diff --git a/src/coreclr/jit/liveness.cpp b/src/coreclr/jit/liveness.cpp index 8449490c068feb..a1d89821ec3944 100644 --- a/src/coreclr/jit/liveness.cpp +++ b/src/coreclr/jit/liveness.cpp @@ -107,7 +107,7 @@ class Liveness bool* pStoreRemoved DEBUGARG(bool* treeModf)); void ComputeLifeLIR(VARSET_TP& life, BasicBlock* block, VARSET_VALARG_TP keepAliveVars); - bool IsTrackedRetBufferAddress(LIR::Range& range, GenTree* node); + bool IsTrackedCallDefinition(LIR::Range& range, GenTree* node); bool TryRemoveDeadStoreLIR(GenTree* store, GenTreeLclVarCommon* lclNode, BasicBlock* block); bool TryRemoveNonLocalLIR(GenTree* node, LIR::Range* blockRange); bool CanUncontainOrRemoveOperands(GenTree* node); @@ -2432,12 +2432,11 @@ void Liveness::ComputeLifeLIR(VARSET_TP& life, BasicBlock* block, VAR } else { - // For LCL_ADDRs that are defined by being passed as a - // retbuf we will handle them when we get to the call. We - // cannot consider them to be defined at the point of the - // LCL_ADDR since there may be uses between the LCL_ADDR - // and call. - if (IsTrackedRetBufferAddress(blockRange, node)) + // For LCL_ADDRs that are definitions for the call we will + // handle them when we get to the call. We cannot consider + // them to be defined at the point of the LCL_ADDR since + // there may be uses between the LCL_ADDR and call. + if (IsTrackedCallDefinition(blockRange, node)) { break; } @@ -2639,15 +2638,15 @@ void Liveness::ComputeLifeLIR(VARSET_TP& life, BasicBlock* block, VAR } //--------------------------------------------------------------------- -// IsTrackedRetBufferAddress - given a LCL_ADDR node, check if it is the -// return buffer definition of a call. +// IsTrackedCallDefinition - given a LCL_ADDR node, check if it is an +// extra definition of a call. // // Arguments // range - the block range containing the LCL_ADDR // node - the LCL_ADDR // template -bool Liveness::IsTrackedRetBufferAddress(LIR::Range& range, GenTree* node) +bool Liveness::IsTrackedCallDefinition(LIR::Range& range, GenTree* node) { assert(node->OperIs(GT_LCL_ADDR)); if ((node->gtFlags & GTF_VAR_DEF) == 0) @@ -2674,7 +2673,12 @@ bool Liveness::IsTrackedRetBufferAddress(LIR::Range& range, GenTree* if (curNode->IsCall()) { - return m_compiler->gtCallGetDefinedRetBufLclAddr(curNode->AsCall()) == node; + auto visit = [=](GenTree* callDef) { + return node == callDef ? GenTree::VisitResult::Abort : GenTree::VisitResult::Continue; + }; + + return + curNode->VisitLocalDefNodes(m_compiler, visit) == GenTree::VisitResult::Abort; } } while (curNode->OperIs(GT_FIELD_LIST) || curNode->OperIsPutArg()); diff --git a/src/coreclr/jit/morph.cpp b/src/coreclr/jit/morph.cpp index fab998a1790e14..94d61531877b9e 100644 --- a/src/coreclr/jit/morph.cpp +++ b/src/coreclr/jit/morph.cpp @@ -1494,10 +1494,10 @@ void CallArgs::EvalArgsToTemps(Compiler* comp, GenTreeCall* call) arg.SetEarlyNode(setupArg); call->gtFlags |= setupArg->gtFlags & GTF_SIDE_EFFECT; - // Make sure we do not break recognition of retbuf-as-local - // optimization here. If this is hit it indicates that we are - // unnecessarily creating temps for some ret buf addresses, and - // gtCallGetDefinedRetBufLclAddr relies on this not to happen. + // Make sure we do not break recognition of defs optimization here. + // If this is hit it indicates that we are unnecessarily creating + // temps for some ret buf addresses, and call defs rely on this not + // to happen. noway_assert((arg.GetWellKnownArg() != WellKnownArg::RetBuffer) || !call->IsOptimizingRetBufAsLocal()); } diff --git a/src/coreclr/jit/promotionliveness.cpp b/src/coreclr/jit/promotionliveness.cpp index e137caab827500..5eae2295fb8d75 100644 --- a/src/coreclr/jit/promotionliveness.cpp +++ b/src/coreclr/jit/promotionliveness.cpp @@ -257,12 +257,25 @@ unsigned PromotionLiveness::GetSizeOfStructLocal(Statement* stmt, GenTreeLclVarC { if (lcl->OperIs(GT_LCL_ADDR)) { - // Retbuf definition. Find the definition size from the - // containing call. + // LCL_ADDR definition. Currently we only have calls that define via + // LCL_ADDRs. Find the definition size from the containing call. Compiler::FindLinkData data = m_compiler->gtFindLink(stmt, lcl); - assert((data.parent != nullptr) && data.parent->IsCall() && - (m_compiler->gtCallGetDefinedRetBufLclAddr(data.parent->AsCall()) == lcl)); - return m_compiler->typGetObjLayout(data.parent->AsCall()->gtRetClsHnd)->GetSize(); + assert((data.parent != nullptr) && data.parent->IsCall()); + //(m_compiler->gtCallGetDefinedRetBufLclAddr(data.parent->AsCall()) == lcl)); + unsigned defSize = UINT_MAX; + auto findDef = [&](const LocalDef& def) { + if (def.Def == lcl) + { + defSize = def.Size.GetExact(); + return GenTree::VisitResult::Abort; + } + + return GenTree::VisitResult::Continue; + }; + + GenTree::VisitResult result = data.parent->VisitLocalDefs(m_compiler, findDef); + assert(result == GenTree::VisitResult::Abort); + return defSize; } return lcl->GetLayout(m_compiler)->GetSize(); From 73187afd5f01730a52a64e5d10932e7c19d4dbce Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 15 May 2026 13:41:32 +0200 Subject: [PATCH 06/41] Fixes --- src/coreclr/jit/lclmorph.cpp | 2 +- src/coreclr/jit/liveness.cpp | 6 +++--- src/coreclr/jit/promotionliveness.cpp | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/coreclr/jit/lclmorph.cpp b/src/coreclr/jit/lclmorph.cpp index be403a634b8100..a66cc246ae7859 100644 --- a/src/coreclr/jit/lclmorph.cpp +++ b/src/coreclr/jit/lclmorph.cpp @@ -113,7 +113,7 @@ class LocalSequencer final : public GenTreeVisitor // void SequenceCall(GenTreeCall* call) { - auto moveToEnd = [&](GenTree* def) { + auto moveToEnd = [&](GenTreeLclVarCommon* def) { MoveNodeToEnd(def); return GenTree::VisitResult::Continue; }; diff --git a/src/coreclr/jit/liveness.cpp b/src/coreclr/jit/liveness.cpp index a1d89821ec3944..db8f215c3acd1b 100644 --- a/src/coreclr/jit/liveness.cpp +++ b/src/coreclr/jit/liveness.cpp @@ -845,9 +845,9 @@ void Liveness::PerNodeLocalVarLiveness(GenTree* tree) case GT_LCL_ADDR: if (TLiveness::IsLIR) { - // If this is a definition of a retbuf then we process it as - // part of the GT_CALL node. - if (IsTrackedRetBufferAddress(LIR::AsRange(m_compiler->compCurBB), tree)) + // If this is a call definition then we process it as part of + // the GT_CALL node. + if (IsTrackedCallDefinition(LIR::AsRange(m_compiler->compCurBB), tree)) { break; } diff --git a/src/coreclr/jit/promotionliveness.cpp b/src/coreclr/jit/promotionliveness.cpp index 5eae2295fb8d75..1b5e695e09915a 100644 --- a/src/coreclr/jit/promotionliveness.cpp +++ b/src/coreclr/jit/promotionliveness.cpp @@ -261,7 +261,7 @@ unsigned PromotionLiveness::GetSizeOfStructLocal(Statement* stmt, GenTreeLclVarC // LCL_ADDRs. Find the definition size from the containing call. Compiler::FindLinkData data = m_compiler->gtFindLink(stmt, lcl); assert((data.parent != nullptr) && data.parent->IsCall()); - //(m_compiler->gtCallGetDefinedRetBufLclAddr(data.parent->AsCall()) == lcl)); + unsigned defSize = UINT_MAX; auto findDef = [&](const LocalDef& def) { if (def.Def == lcl) From d13f065356d6bf2b1bd9cf3af07c9f42c73218e5 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 15 May 2026 13:06:10 +0200 Subject: [PATCH 07/41] JIT: Move invariant nodes and LCL_VARs in LiftLIREdges Invariant nodes and LCL_VARs do not need to be lifted across async calls. --- src/coreclr/jit/async.cpp | 33 +++++++++++++++++++-------------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 30e1a6495aa6ac..ab5bd051ed91a4 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -1093,24 +1093,29 @@ void AsyncTransformation::LiftLIREdges(BasicBlock* block, for (GenTree* tree : defs) { - // TODO-CQ: Enable this. It currently breaks our recognition of how the - // call is stored. - // if (tree->OperIs(GT_LCL_VAR)) - //{ - // LclVarDsc* dsc = m_compiler->lvaGetDesc(tree->AsLclVarCommon()); - // if (!dsc->IsAddressExposed()) - // { - // // No interference by IR invariants. - // LIR::AsRange(block).Remove(tree); - // LIR::AsRange(block).InsertAfter(beyond, tree); - // continue; - // } - //} - LIR::Use use; bool gotUse = LIR::AsRange(block).TryGetUse(tree, &use); assert(gotUse); // Defs list should not contain unused values. + if (tree->IsInvariant()) + { + LIR::AsRange(block).Remove(tree); + LIR::AsRange(block).InsertBefore(use.User(), tree); + continue; + } + + if (tree->OperIs(GT_LCL_VAR)) + { + LclVarDsc* dsc = m_compiler->lvaGetDesc(tree->AsLclVarCommon()); + if (!dsc->IsAddressExposed()) + { + // No interference by IR invariants + LIR::AsRange(block).Remove(tree); + LIR::AsRange(block).InsertAfter(use.User(), tree); + continue; + } + } + unsigned newLclNum = use.ReplaceWithLclVar(m_compiler); layoutBuilder->AddLocal(newLclNum); GenTree* newUse = use.Def(); From e5b707ec3a91cfac9da76b43ee1ba783ab4474fa Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 15 May 2026 13:33:09 +0200 Subject: [PATCH 08/41] Fix --- src/coreclr/jit/async.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index ab5bd051ed91a4..2081bb85521a85 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -1111,7 +1111,7 @@ void AsyncTransformation::LiftLIREdges(BasicBlock* block, { // No interference by IR invariants LIR::AsRange(block).Remove(tree); - LIR::AsRange(block).InsertAfter(use.User(), tree); + LIR::AsRange(block).InsertBefore(use.User(), tree); continue; } } From 6d2556d669ad466f8f4f8a29bafda96745a4ad27 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 15 May 2026 14:27:05 +0200 Subject: [PATCH 09/41] Run jit-format --- src/coreclr/jit/fgdiagnostic.cpp | 4 ++-- src/coreclr/jit/lclmorph.cpp | 2 +- src/coreclr/jit/liveness.cpp | 5 ++--- src/coreclr/jit/promotionliveness.cpp | 4 ++-- 4 files changed, 7 insertions(+), 8 deletions(-) diff --git a/src/coreclr/jit/fgdiagnostic.cpp b/src/coreclr/jit/fgdiagnostic.cpp index 112720bbe14406..75c42931cdfef1 100644 --- a/src/coreclr/jit/fgdiagnostic.cpp +++ b/src/coreclr/jit/fgdiagnostic.cpp @@ -3826,7 +3826,7 @@ void Compiler::fgDebugCheckLinkedLocals() assert(ShouldLink(def)); m_locals.Push(def); return GenTree::VisitResult::Continue; - }; + }; node->VisitLocalDefNodes(m_compiler, linkDefs); } @@ -3838,7 +3838,7 @@ void Compiler::fgDebugCheckLinkedLocals() { auto defIsNode = [=](GenTree* def) { return node == def ? GenTree::VisitResult::Abort : GenTree::VisitResult::Continue; - }; + }; return call->VisitLocalDefNodes(m_compiler, defIsNode) == GenTree::VisitResult::Abort; } }; diff --git a/src/coreclr/jit/lclmorph.cpp b/src/coreclr/jit/lclmorph.cpp index a66cc246ae7859..f681866626b95a 100644 --- a/src/coreclr/jit/lclmorph.cpp +++ b/src/coreclr/jit/lclmorph.cpp @@ -116,7 +116,7 @@ class LocalSequencer final : public GenTreeVisitor auto moveToEnd = [&](GenTreeLclVarCommon* def) { MoveNodeToEnd(def); return GenTree::VisitResult::Continue; - }; + }; call->VisitLocalDefNodes(m_compiler, moveToEnd); } diff --git a/src/coreclr/jit/liveness.cpp b/src/coreclr/jit/liveness.cpp index db8f215c3acd1b..ea21161b4c92d5 100644 --- a/src/coreclr/jit/liveness.cpp +++ b/src/coreclr/jit/liveness.cpp @@ -2675,10 +2675,9 @@ bool Liveness::IsTrackedCallDefinition(LIR::Range& range, GenTree* no { auto visit = [=](GenTree* callDef) { return node == callDef ? GenTree::VisitResult::Abort : GenTree::VisitResult::Continue; - }; + }; - return - curNode->VisitLocalDefNodes(m_compiler, visit) == GenTree::VisitResult::Abort; + return curNode->VisitLocalDefNodes(m_compiler, visit) == GenTree::VisitResult::Abort; } } while (curNode->OperIs(GT_FIELD_LIST) || curNode->OperIsPutArg()); diff --git a/src/coreclr/jit/promotionliveness.cpp b/src/coreclr/jit/promotionliveness.cpp index 1b5e695e09915a..8ba68f7711c5ca 100644 --- a/src/coreclr/jit/promotionliveness.cpp +++ b/src/coreclr/jit/promotionliveness.cpp @@ -263,7 +263,7 @@ unsigned PromotionLiveness::GetSizeOfStructLocal(Statement* stmt, GenTreeLclVarC assert((data.parent != nullptr) && data.parent->IsCall()); unsigned defSize = UINT_MAX; - auto findDef = [&](const LocalDef& def) { + auto findDef = [&](const LocalDef& def) { if (def.Def == lcl) { defSize = def.Size.GetExact(); @@ -271,7 +271,7 @@ unsigned PromotionLiveness::GetSizeOfStructLocal(Statement* stmt, GenTreeLclVarC } return GenTree::VisitResult::Continue; - }; + }; GenTree::VisitResult result = data.parent->VisitLocalDefs(m_compiler, findDef); assert(result == GenTree::VisitResult::Abort); From f448161665fe42a70844423ddb46e1a4987f0bbd Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Mon, 27 Jul 2026 16:16:28 +0200 Subject: [PATCH 10/41] Squash change with GT_CONTINUATION_MEMBER_OFFSET --- .../CompilerServices/AsyncHelpers.CoreCLR.cs | 14 +- src/coreclr/inc/corinfo.h | 8 + src/coreclr/inc/icorjitinfoimpl_generated.h | 7 + src/coreclr/inc/jiteeversionguid.h | 10 +- src/coreclr/jit/ICorJitInfo_names_generated.h | 1 + .../jit/ICorJitInfo_wrapper_generated.hpp | 13 + src/coreclr/jit/async.cpp | 314 ++++++++++++++---- src/coreclr/jit/async.h | 67 ++-- src/coreclr/jit/compiler.h | 15 + src/coreclr/jit/gentree.cpp | 12 + src/coreclr/jit/gtlist.h | 6 + src/coreclr/jit/gtstructs.h | 2 +- src/coreclr/jit/importercalls.cpp | 81 +++++ src/coreclr/jit/jitconfigvalues.h | 6 - src/coreclr/jit/morph.cpp | 13 + src/coreclr/jit/valuenum.cpp | 1 + src/coreclr/jit/wellknownargs.h | 1 + .../tools/Common/JitInterface/CorInfoImpl.cs | 55 +++ .../JitInterface/CorInfoImpl_generated.cs | 17 + .../ThunkGenerator/ThunkInput.txt | 1 + .../Compiler/ReadyToRunCodegenCompilation.cs | 2 + .../aot/jitinterface/jitinterface_generated.h | 14 + .../tools/superpmi/superpmi-shared/agnostic.h | 7 + .../tools/superpmi/superpmi-shared/lwmlist.h | 1 + .../superpmi-shared/methodcontext.cpp | 67 ++++ .../superpmi/superpmi-shared/methodcontext.h | 15 + .../superpmi-shim-collector/icorjitinfo.cpp | 14 + .../icorjitinfo_generated.cpp | 11 + .../icorjitinfo_generated.cpp | 10 + .../tools/superpmi/superpmi/icorjitinfo.cpp | 12 + src/coreclr/vm/corelib.h | 2 + src/coreclr/vm/jitinterface.cpp | 113 +++++++ src/coreclr/vm/jitinterface.h | 5 + .../Runtime/CompilerServices/AsyncHelpers.cs | 52 +++ .../GitHub_130437/GitHub_130437.csproj | 2 +- .../awaitingnoasyncgeneric.csproj | 2 +- .../awaitingnotasync/awaitingnotasync.csproj | 2 +- .../async/byref-param/byref-param.csproj | 2 +- .../async/cancellation/cancellation.csproj | 2 +- .../collectible-alc/collectible-alc.csproj | 2 +- .../custom-struct-awaiters.cs | 85 +++++ .../custom-struct-awaiters.csproj | 5 + .../async/eh-microbench/eh-microbench.csproj | 2 +- .../execution-context.csproj | 2 +- .../async/gc-roots-scan/gc-roots-scan.csproj | 2 +- src/tests/async/implement/implement.csproj | 2 +- .../inst-unbox-thunks.csproj | 2 +- .../mincallcost-microbench.csproj | 2 +- src/tests/async/object/object.csproj | 2 +- .../objects-captured/objects-captured.csproj | 2 +- src/tests/async/override/override.csproj | 2 +- src/tests/async/pgo/pgo.csproj | 2 +- src/tests/async/pinvoke/pinvoke.csproj | 2 +- src/tests/async/regression/125042.csproj | 2 +- src/tests/async/regression/125615.csproj | 2 +- src/tests/async/regression/125805.csproj | 2 +- src/tests/async/regression/128566.csproj | 2 +- .../async/regression/managed-thread-id.csproj | 2 +- src/tests/async/returns/returns.csproj | 2 +- .../shared-generic/shared-generic.csproj | 2 +- .../async/staticvirtual/staticvirtual.csproj | 2 +- src/tests/async/struct/struct.cs | 1 - src/tests/async/struct/struct.csproj | 2 +- .../synchronization-context.csproj | 2 +- .../valuetask-source-stub.csproj | 2 +- .../valuetask-source/valuetask-source.csproj | 2 +- src/tests/async/valuetask/valuetask.csproj | 2 +- .../varying-yields/varying-yields-task.csproj | 2 +- .../varying-yields-valuetask.csproj | 2 +- .../varying-yields/varying-yields.csproj | 2 +- src/tests/async/void/void.csproj | 2 +- .../async/with-yields/with-yields.csproj | 2 +- .../without-yields/without-yields.csproj | 2 +- 73 files changed, 993 insertions(+), 139 deletions(-) create mode 100644 src/tests/async/custom-struct-awaiters/custom-struct-awaiters.cs create mode 100644 src/tests/async/custom-struct-awaiters/custom-struct-awaiters.csproj diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index fd3e75881619ff..2fd4bb51f3b564 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -223,6 +223,8 @@ private ref struct RuntimeAsyncStackState public INotifyCompletion? Notifier; public ValueTaskSourceContinuation? ValueTaskSourceContinuation; public RuntimeAsyncTaskContinuation? TaskContinuation; + public delegate* AwaiterContinuation; + public int AwaiterOffset; // When we suspend in the leaf, the contexts are captured into these fields. public ExecutionContext? LeafExecutionContext; @@ -265,6 +267,8 @@ public void CaptureContexts() [NonVersionable] public void Push(RuntimeAsyncStackState* stackState) { + stackState->AwaiterContinuation = null; + stackState->AwaiterOffset = 0; stackState->Next = StackState; StackState = stackState; CurrentThread ??= Thread.CurrentThread; @@ -805,7 +809,15 @@ internal unsafe bool HandleSuspended(ref RuntimeAsyncAwaitState state) try { - if (stackState->CriticalNotifier is { } critNotifier) + if (stackState->AwaiterContinuation != null) + { + // The awaiter is stored in the continuation for the caller of + // AwaitAwaiterInContinuation or UnsafeAwaitAwaiterInContinuation. + Debug.Assert((headContinuation.Flags & ContinuationFlags.AllContinuationFlags) == 0); + stackState->AwaiterContinuation( + headContinuation, stackState->AwaiterOffset, GetContinuationAction()); + } + else if (stackState->CriticalNotifier is { } critNotifier) { // Result of async call to AwaitAwaiter or UnsafeAwaitAwaiter. // These never have special continuation context handling. diff --git a/src/coreclr/inc/corinfo.h b/src/coreclr/inc/corinfo.h index cd6b359c9f5d76..f35be06c2cf136 100644 --- a/src/coreclr/inc/corinfo.h +++ b/src/coreclr/inc/corinfo.h @@ -3154,6 +3154,14 @@ class ICorStaticInfo // instantiation argument that must be passed to the await call. virtual CORINFO_METHOD_HANDLE getAwaitReturnCall(CORINFO_METHOD_HANDLE callerHandle, CORINFO_CONTEXT_HANDLE* contextHandle, CORINFO_LOOKUP* instArg) = 0; + virtual CORINFO_METHOD_HANDLE getAwaitAwaiterInContinuationCall( + CORINFO_METHOD_HANDLE callerHandle, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg + ) = 0; + /*********************************************************************************/ // // Diagnostic methods diff --git a/src/coreclr/inc/icorjitinfoimpl_generated.h b/src/coreclr/inc/icorjitinfoimpl_generated.h index 6c0c39f71f0293..829e1073cd8d25 100644 --- a/src/coreclr/inc/icorjitinfoimpl_generated.h +++ b/src/coreclr/inc/icorjitinfoimpl_generated.h @@ -503,6 +503,13 @@ CORINFO_METHOD_HANDLE getAwaitReturnCall( CORINFO_CONTEXT_HANDLE* contextHandle, CORINFO_LOOKUP* instArg) override; +CORINFO_METHOD_HANDLE getAwaitAwaiterInContinuationCall( + CORINFO_METHOD_HANDLE callerHandle, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg) override; + mdMethodDef getMethodDefFromMethod( CORINFO_METHOD_HANDLE hMethod) override; diff --git a/src/coreclr/inc/jiteeversionguid.h b/src/coreclr/inc/jiteeversionguid.h index 08a3bcde5744d6..cb3210e624b6a2 100644 --- a/src/coreclr/inc/jiteeversionguid.h +++ b/src/coreclr/inc/jiteeversionguid.h @@ -37,11 +37,11 @@ #include -constexpr GUID JITEEVersionIdentifier = { /* 5bb9b1e5-9941-4762-a195-351b6a736588 */ - 0x5bb9b1e5, - 0x9941, - 0x4762, - {0xa1, 0x95, 0x35, 0x1b, 0x6a, 0x73, 0x65, 0x88} +constexpr GUID JITEEVersionIdentifier = { /* 364f7f40-1feb-4b18-bc70-a3cb6481665b */ + 0x364f7f40, + 0x1feb, + 0x4b18, + {0xbc, 0x70, 0xa3, 0xcb, 0x64, 0x81, 0x66, 0x5b} }; #endif // JIT_EE_VERSIONING_GUID_H diff --git a/src/coreclr/jit/ICorJitInfo_names_generated.h b/src/coreclr/jit/ICorJitInfo_names_generated.h index b3ccde2acebc8f..bd0c0dd01da10f 100644 --- a/src/coreclr/jit/ICorJitInfo_names_generated.h +++ b/src/coreclr/jit/ICorJitInfo_names_generated.h @@ -125,6 +125,7 @@ DEF_CLR_API(runWithSPMIErrorTrap) DEF_CLR_API(getEEInfo) DEF_CLR_API(getAsyncInfo) DEF_CLR_API(getAwaitReturnCall) +DEF_CLR_API(getAwaitAwaiterInContinuationCall) DEF_CLR_API(getMethodDefFromMethod) DEF_CLR_API(printMethodName) DEF_CLR_API(getMethodNameFromMetadata) diff --git a/src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp b/src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp index 3f40e71e6bbb26..bec4c91776ee52 100644 --- a/src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp +++ b/src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp @@ -1190,6 +1190,19 @@ CORINFO_METHOD_HANDLE WrapICorJitInfo::getAwaitReturnCall( return temp; } +CORINFO_METHOD_HANDLE WrapICorJitInfo::getAwaitAwaiterInContinuationCall( + CORINFO_METHOD_HANDLE callerHandle, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg) +{ + API_ENTER(getAwaitAwaiterInContinuationCall); + CORINFO_METHOD_HANDLE temp = wrapHnd->getAwaitAwaiterInContinuationCall(callerHandle, callSig, isUnsafe, contextHandle, instArg); + API_LEAVE(getAwaitAwaiterInContinuationCall); + return temp; +} + mdMethodDef WrapICorJitInfo::getMethodDefFromMethod( CORINFO_METHOD_HANDLE hMethod) { diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 7b3b895f54c41e..8b9de8fffdc229 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -45,6 +45,84 @@ #include "jitstd/algorithm.h" #include "async.h" +ContinuationMember ContinuationMember::CustomAwaiterOfLayout(ClassLayout* layout) +{ + ContinuationMember member; + member.Type = ContinuationMemberType::CustomAwaiterOfLayout; + member.m_customAwaiterLayout = layout; + return member; +} + +ClassLayout* ContinuationMember::GetCustomAwaiterLayout() const +{ + assert(Type == ContinuationMemberType::CustomAwaiterOfLayout); + return m_customAwaiterLayout; +} + +bool ContinuationMember::AreCompatible(const ContinuationMember& a, const ContinuationMember& b) +{ + if (a.Type != b.Type) + { + return false; + } + + switch (a.Type) + { + case ContinuationMemberType::CustomAwaiterOfLayout: + return ClassLayout::AreCompatible(a.m_customAwaiterLayout, b.m_customAwaiterLayout); + default: + unreached(); + } +} + +#ifdef DEBUG +void ContinuationMember::Print() const +{ + switch (Type) + { + case ContinuationMemberType::CustomAwaiterOfLayout: + printf("CustomAwaiter<%s>", m_customAwaiterLayout->GetClassName()); + break; + default: + unreached(); + } +} +#endif + +size_t Compiler::GetContinuationMemberIndex(const ContinuationMember& member) +{ + if (m_asyncContinuationMembers == nullptr) + { + m_asyncContinuationMembers = new (this, CMK_Async) jitstd::vector(getAllocator(CMK_Async)); + } + else + { + for (size_t i = 0; i < m_asyncContinuationMembers->size(); i++) + { + const ContinuationMember& existingMember = m_asyncContinuationMembers->at(i); + + if (ContinuationMember::AreCompatible(member, existingMember)) + { + return i; + } + } + } + + m_asyncContinuationMembers->push_back(member); + return m_asyncContinuationMembers->size() - 1; +} + +size_t Compiler::GetContinuationMemberCount() const +{ + return m_asyncContinuationMembers == nullptr ? 0 : m_asyncContinuationMembers->size(); +} + +const ContinuationMember& Compiler::GetContinuationMember(size_t index) +{ + assert(index < m_asyncContinuationMembers->size()); + return m_asyncContinuationMembers->at(index); +} + //------------------------------------------------------------------------ // Compiler::SaveAsyncContexts: // Insert code in async methods that saves and restores contexts. @@ -600,7 +678,8 @@ PhaseStatus AsyncTransformation::Run() PhaseStatus result = PhaseStatus::MODIFIED_NOTHING; ArrayStack blocksWithNormalAwaits(m_compiler->getAllocator(CMK_Async)); ArrayStack blocksWithTailAwaits(m_compiler->getAllocator(CMK_Async)); - AggregatedAwaitInfo awaits = FindAwaits(blocksWithNormalAwaits, blocksWithTailAwaits); + ArrayStack continuationMemberOffsets(m_compiler->getAllocator(CMK_Async)); + AggregatedAwaitInfo awaits = FindAwaits(blocksWithNormalAwaits, blocksWithTailAwaits, continuationMemberOffsets); if (awaits.NumNormalAwaits + awaits.NumTailAwaits > 1) { @@ -620,7 +699,8 @@ PhaseStatus AsyncTransformation::Run() // This may have changed blocks, so refind the normal awaits. blocksWithNormalAwaits.Reset(); blocksWithTailAwaits.Reset(); - awaits = FindAwaits(blocksWithNormalAwaits, blocksWithTailAwaits); + continuationMemberOffsets.Reset(); + awaits = FindAwaits(blocksWithNormalAwaits, blocksWithTailAwaits, continuationMemberOffsets); } result = PhaseStatus::MODIFIED_EVERYTHING; @@ -630,6 +710,7 @@ PhaseStatus AsyncTransformation::Run() if (awaits.NumNormalAwaits <= 0) { + assert(continuationMemberOffsets.Empty()); return result; } @@ -727,29 +808,37 @@ PhaseStatus AsyncTransformation::Run() } while (any); } - if (ReuseContinuations()) + // Set up the local containing the continuation we can reuse. For OSR + // things are special: we can transition to the OSR method after having + // resumed in the tier0 method. In that case we end up with the tier0 + // continuation in the OSR method, but we cannot reuse it. + if (m_compiler->opts.IsOSR()) { - // Set up the local containing the continuation we can reuse. For OSR - // things are special: we can transition to the OSR method after having - // resumed in the tier0 method. In that case we end up with the tier0 - // continuation in the OSR method, but we cannot reuse it. - if (m_compiler->opts.IsOSR()) - { - m_reuseContinuationVar = m_compiler->lvaGrabTemp(false DEBUGARG("OSR reusable continuation")); - m_compiler->lvaGetDesc(m_reuseContinuationVar)->lvType = TYP_REF; - } - else - { - m_reuseContinuationVar = m_compiler->lvaAsyncContinuationArg; - } + m_reuseContinuationVar = m_compiler->lvaGrabTemp(false DEBUGARG("OSR reusable continuation")); + m_compiler->lvaGetDesc(m_reuseContinuationVar)->lvType = TYP_REF; + } + else + { + m_reuseContinuationVar = m_compiler->lvaAsyncContinuationArg; } - CreateResumptionsAndSuspensions(); + const ContinuationLayout* continuationLayout = CreateResumptionsAndSuspensions(continuationMemberOffsets); // After transforming all async calls we have created resumption blocks; // create the resumption switch. CreateResumptionSwitch(); + // Now bash all GT_CONTINUATION_MEMBER_OFFSET into appropriate constants. + for (GenTree* node : continuationMemberOffsets.BottomUpOrder()) + { + size_t memberIndex = node->AsVal()->gtVal1; + assert(memberIndex < continuationLayout->ContinuationMemberOffsets.size()); + assert(continuationLayout->ContinuationMemberOffsets[memberIndex] != UINT_MAX); + ssize_t offset = (OFFSETOF__CORINFO_Continuation__data - SIZEOF__CORINFO_Object) + + continuationLayout->ContinuationMemberOffsets[memberIndex]; + node->BashToConst(offset, TYP_INT); + } + m_compiler->fgInvalidateDfsTree(); if (m_compiler->opts.OptimizationDisabled()) @@ -785,12 +874,14 @@ PhaseStatus AsyncTransformation::Run() // Parameters: // blocksWithNormalAwaits - [out] Blocks with normal awaits are pushed onto this stack // blocksWithTailAwaits - [out] Blocks with tail awaits are pushed onto this stack +// continuationMemberOffsets - [out] Symbolic continuation member offset nodes // // Returns: // Information about awaits in the function. // AggregatedAwaitInfo AsyncTransformation::FindAwaits(ArrayStack& blocksWithNormalAwaits, - ArrayStack& blocksWithTailAwaits) + ArrayStack& blocksWithTailAwaits, + ArrayStack& continuationMemberOffsets) { AggregatedAwaitInfo awaits; for (BasicBlock* block : m_compiler->Blocks()) @@ -799,6 +890,12 @@ AggregatedAwaitInfo AsyncTransformation::FindAwaits(ArrayStack& blo bool hasTailAwait = false; for (GenTree* tree : LIR::AsRange(block)) { + if (tree->OperIs(GT_CONTINUATION_MEMBER_OFFSET)) + { + continuationMemberOffsets.Push(tree); + continue; + } + if (!tree->IsCall()) { continue; @@ -1363,6 +1460,16 @@ void ContinuationLayout::Dump(int indent) printf("%*s +%03u Keep alive object\n", indent, "", KeepAliveOffset); } + for (size_t i = 0; i < ContinuationMemberOffsets.size(); i++) + { + if (ContinuationMemberOffsets[i] != UINT_MAX) + { + printf("%*s +%03u ", indent, "", ContinuationMemberOffsets[i]); + JitTls::GetCompiler()->GetContinuationMember(i).Print(); + printf("\n"); + } + } + for (const LiveLocalInfo& inf : Locals) { printf("%*s +%03u V%02u: %u bytes\n", indent, "", inf.Offset, inf.LclNum, inf.Size); @@ -1413,10 +1520,12 @@ const ReturnInfo* ContinuationLayout::FindReturn(Compiler* comp, GenTreeCall* ca // The finalized ContinuationLayout with computed offsets and a class // handle for the continuation type. // -ContinuationLayout* ContinuationLayoutBuilder::Create() +ContinuationLayout* ContinuationLayoutBuilder::Create(ArrayStack& continuationMemberOffsets) { ContinuationLayout* layout = new (m_compiler, CMK_Async) ContinuationLayout(m_compiler); layout->Locals.reserve(m_locals.size()); + size_t continuationMemberCount = m_compiler->GetContinuationMemberCount(); + layout->ContinuationMemberOffsets.resize(continuationMemberCount, UINT_MAX); for (unsigned lclNum : m_locals) { @@ -1530,6 +1639,24 @@ ContinuationLayout* ContinuationLayoutBuilder::Create() ret.Offset = allocLayout(ret.HeapAlignment(), ret.Size); } + for (GenTree* memberOffsetNode : continuationMemberOffsets.BottomUpOrder()) + { + size_t memberIndex = memberOffsetNode->AsVal()->gtVal1; + assert(memberIndex < continuationMemberCount); + if (layout->ContinuationMemberOffsets[memberIndex] != UINT_MAX) + { + continue; + } + + ClassLayout* memberLayout = m_compiler->GetContinuationMember(memberIndex).GetCustomAwaiterLayout(); + unsigned alignment = + memberLayout->IsCustomLayout() + ? (memberLayout->HasGCPtr() ? TARGET_POINTER_SIZE : 1) + : m_compiler->info.compCompHnd->getClassAlignmentRequirement(memberLayout->GetClassHandle()); + layout->ContinuationMemberOffsets[memberIndex] = + allocLayout(std::min(alignment, (unsigned)TARGET_POINTER_SIZE), memberLayout->GetSize()); + } + if (m_needsKeepAlive) { layout->KeepAliveOffset = allocLayout(TARGET_POINTER_SIZE, TARGET_POINTER_SIZE); @@ -1578,6 +1705,15 @@ ContinuationLayout* ContinuationLayoutBuilder::Create() bitmapBuilder.SetType(ret.Offset, ret.Type.ReturnType, ret.Type.ReturnLayout); } + for (size_t i = 0; i < layout->ContinuationMemberOffsets.size(); i++) + { + if (layout->ContinuationMemberOffsets[i] != UINT_MAX) + { + bitmapBuilder.SetType(layout->ContinuationMemberOffsets[i], TYP_STRUCT, + m_compiler->GetContinuationMember(i).GetCustomAwaiterLayout()); + } + } + #ifdef DEBUG if (m_compiler->verbose) { @@ -1920,7 +2056,7 @@ bool AsyncTransformation::IsReusableSuspension(const AsyncState* state, } } - static const WellKnownArg validateArgs[] = {WellKnownArg::AsyncExecutionContext, + static const WellKnownArg validateArgs[] = {WellKnownArg::AsyncAwaiter, WellKnownArg::AsyncExecutionContext, WellKnownArg::AsyncSynchronizationContext}; for (WellKnownArg arg : validateArgs) { @@ -1972,7 +2108,7 @@ bool AsyncTransformation::IsReusableSuspension(const AsyncState* state, // void AsyncTransformation::HandleReusedSuspension(BasicBlock* callBlock, GenTreeCall* call) { - static const WellKnownArg argsToRemove[] = {WellKnownArg::AsyncExecutionContext, + static const WellKnownArg argsToRemove[] = {WellKnownArg::AsyncAwaiter, WellKnownArg::AsyncExecutionContext, WellKnownArg::AsyncSynchronizationContext}; for (WellKnownArg wka : argsToRemove) { @@ -2078,7 +2214,7 @@ void AsyncTransformation::CreateSuspension(BasicBlock* call LIR::AsRange(suspendBB).InsertAtEnd(storeNewContinuation); SaveSet tailSaveSet = SaveSet::All; - if (ReuseContinuations() && resumeReachable) + if (resumeReachable) { // Split suspendBB into suspendBB -> [reuse continuation with Next store] -> [allocNewBlock with allocation // call] -> suspendBBTail [empty] @@ -2204,10 +2340,78 @@ void AsyncTransformation::CreateSuspension(BasicBlock* call LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_compiler, storeFlags)); FillInDataOnSuspension(layout, subLayout, suspendBB, mutatedSinceResumption, tailSaveSet); + StoreAsyncAwaiter(callBlock, call, suspendBB, layout); FinishContextHandlingAndSuspension(callBlock, call, suspendBB, layout, subLayout); } +//------------------------------------------------------------------------ +// AsyncTransformation::StoreAsyncAwaiter: +// Move the pseudo awaiter argument into its reserved continuation member. +// +void AsyncTransformation::StoreAsyncAwaiter(BasicBlock* callBlock, + GenTreeCall* call, + BasicBlock* suspendBB, + const ContinuationLayout& layout) +{ + CallArg* awaiterArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncAwaiter); + if (awaiterArg == nullptr) + { + return; + } + + ClassLayout* awaiterLayout = awaiterArg->GetSignatureLayout(); + assert(awaiterLayout != nullptr); + size_t memberIndex = + m_compiler->GetContinuationMemberIndex(ContinuationMember::CustomAwaiterOfLayout(awaiterLayout)); + assert(memberIndex < layout.ContinuationMemberOffsets.size()); + assert(layout.ContinuationMemberOffsets[memberIndex] != UINT_MAX); + + GenTree* awaiter = awaiterArg->GetNode(); + assert(varTypeIsStruct(awaiter)); + + if (awaiter->OperIs(GT_FIELD_LIST)) + { + GenTreeFieldList* fieldList = awaiter->AsFieldList(); + for (GenTreeFieldList::Use& use : fieldList->Uses()) + { + if (!use.GetNode()->IsInvariant() && !use.GetNode()->OperIs(GT_LCL_VAR)) + { + LIR::Use lirUse(LIR::AsRange(callBlock), &use.NodeRef(), fieldList); + lirUse.ReplaceWithLclVar(m_compiler); + } + + GenTree* field = use.GetNode(); + LIR::AsRange(callBlock).Remove(field); + + GenTree* continuation = m_compiler->gtNewLclvNode(GetNewContinuationVar(), TYP_REF); + unsigned offset = + OFFSETOF__CORINFO_Continuation__data + layout.ContinuationMemberOffsets[memberIndex] + use.GetOffset(); + GenTree* store = StoreAtOffset(continuation, offset, field, use.GetType()); + LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_compiler, store)); + } + } + else + { + if (!awaiter->OperIs(GT_LCL_VAR)) + { + LIR::Use use(LIR::AsRange(callBlock), &awaiterArg->NodeRef(), call); + use.ReplaceWithLclVar(m_compiler); + awaiter = use.Def(); + } + + 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)); + } + + LIR::AsRange(callBlock).Remove(awaiter); + call->gtArgs.RemoveUnsafe(awaiterArg); +} + //------------------------------------------------------------------------ // AsyncTransformation::CreateAllocContinuationCall: // Create a call to the JIT helper that allocates a continuation. @@ -3025,18 +3229,15 @@ BasicBlock* AsyncTransformation::RethrowExceptionOnResumption(BasicBlock* GenTree* storeException = m_compiler->gtNewStoreLclVarNode(exceptionLclNum, exceptionInd); LIR::AsRange(resumeBB).InsertAtEnd(LIR::SeqTree(m_compiler, storeException)); - if (ReuseContinuations()) - { - // If we may reuse this continuation later then make sure we don't see the same exception again. - GenTree* continuation = m_compiler->gtNewLclvNode(m_compiler->lvaAsyncContinuationArg, TYP_REF); - unsigned exceptionOffset = OFFSETOF__CORINFO_Continuation__data + layout.ExceptionOffset; - GenTree* null = m_compiler->gtNewNull(); - GenTree* nullException = StoreAtOffset(continuation, exceptionOffset, null, TYP_REF); - LIR::AsRange(resumeBB).InsertAtEnd(LIR::SeqTree(m_compiler, nullException)); - } + // Since we may reuse this continuation later we make sure we don't see the same exception again. + continuation = m_compiler->gtNewLclvNode(m_compiler->lvaAsyncContinuationArg, TYP_REF); + exceptionOffset = OFFSETOF__CORINFO_Continuation__data + layout.ExceptionOffset; + GenTree* null = m_compiler->gtNewNull(); + GenTree* nullException = StoreAtOffset(continuation, exceptionOffset, null, TYP_REF); + LIR::AsRange(resumeBB).InsertAtEnd(LIR::SeqTree(m_compiler, nullException)); GenTree* exception = m_compiler->gtNewLclVarNode(exceptionLclNum, TYP_REF); - GenTree* null = m_compiler->gtNewNull(); + null = m_compiler->gtNewNull(); GenTree* neNull = m_compiler->gtNewOperNode(GT_NE, TYP_INT, exception, null); GenTree* jtrue = m_compiler->gtNewOperNode(GT_JTRUE, TYP_VOID, neNull); LIR::AsRange(resumeBB).InsertAtEnd(exception, null, neNull, jtrue); @@ -3160,10 +3361,7 @@ void AsyncTransformation::CopyReturnValueOnResumption(GenTreeCall* LIR::AsRange(storeResultBB).InsertAtEnd(LIR::SeqTree(m_compiler, storeResult)); } - if (ReuseContinuations()) - { - ClearReturnValueOnResumption(retInfo, resultOffset, storeResultBB); - } + ClearReturnValueOnResumption(retInfo, resultOffset, storeResultBB); } //------------------------------------------------------------------------ @@ -3661,9 +3859,10 @@ void AsyncTransformation::InsertFinishContextHandlingCall(BasicBlock* // Walk all recorded async states and create the suspension and resumption // IR, continuation layouts, and debug info for each one. // -void AsyncTransformation::CreateResumptionsAndSuspensions() +const ContinuationLayout* AsyncTransformation::CreateResumptionsAndSuspensions( + ArrayStack& continuationMemberOffsets) { - bool useSharedLayout = (m_states.size() > 1) && ReuseContinuations(); + bool useSharedLayout = m_states.size() > 1; ContinuationLayout* sharedLayout = nullptr; if (useSharedLayout) @@ -3671,7 +3870,7 @@ void AsyncTransformation::CreateResumptionsAndSuspensions() JITDUMP("Creating shared layout:\n"); ContinuationLayoutBuilder* sharedLayoutBuilder = ContinuationLayoutBuilder::CreateSharedLayout(m_compiler, m_states); - sharedLayout = sharedLayoutBuilder->Create(); + sharedLayout = sharedLayoutBuilder->Create(continuationMemberOffsets); unsigned numSharedSuspensionsWithContinuationContext = 0; unsigned numSharedSuspensionsWithoutContinuationContext = 0; @@ -3733,12 +3932,14 @@ void AsyncTransformation::CreateResumptionsAndSuspensions() } } + ContinuationLayout* layout = nullptr; JITDUMP("Creating suspensions and resumptions for %zu states\n", m_states.size()); for (const AsyncState& state : m_states) { JITDUMP("State %u suspend @ " FMT_BB ", resume @ " FMT_BB "\n", state.Number, state.SuspensionBB->bbNum, state.ResumptionBB->bbNum); - ContinuationLayout* layout = sharedLayout == nullptr ? state.Layout->Create() : sharedLayout; + layout = sharedLayout == nullptr ? state.Layout->Create(continuationMemberOffsets) : sharedLayout; + CreateSuspension(state.CallBlock, state.Call, state.SuspensionBB, state.Number, *layout, *state.Layout, state.ResumeReachable, state.MutatedSincePreviousResumption); CreateResumption(state.CallBlock, state.Call, state.ResumptionBB, state.CallDefInfo, *layout, *state.Layout); @@ -3746,28 +3947,8 @@ void AsyncTransformation::CreateResumptionsAndSuspensions() JITDUMP("\n"); } -} -//------------------------------------------------------------------------ -// AsyncTransformation::ReuseContinuations: -// Returns true if continuation reuse is enabled. -// -// Returns: -// True if so. -// -bool AsyncTransformation::ReuseContinuations() -{ -#ifdef DEBUG - static ConfigMethodRange s_range; - s_range.EnsureInit(JitConfig.JitAsyncReuseContinuationsRange()); - - if (!s_range.Contains(m_compiler->info.compMethodHash())) - { - return false; - } -#endif - - return JitConfig.JitAsyncReuseContinuations() != 0; + return layout; } //------------------------------------------------------------------------ @@ -4016,12 +4197,9 @@ void AsyncTransformation::CreateResumptionSwitch() GenTree* jtrue = m_compiler->gtNewOperNode(GT_JTRUE, TYP_VOID, eqZero); LIR::AsRange(checkOSRAddressOffsetBB).InsertAtEnd(LIR::SeqTree(m_compiler, jtrue)); - if (ReuseContinuations()) - { - // Also, save the fact that we have a reusable continuation - continuationArg = m_compiler->gtNewLclvNode(m_compiler->lvaAsyncContinuationArg, TYP_REF); - GenTree* storeReusable = m_compiler->gtNewStoreLclVarNode(m_reuseContinuationVar, continuationArg); - LIR::AsRange(onContinuationBB).InsertAtBeginning(continuationArg, storeReusable); - } + // Also, save the fact that we have a reusable continuation + continuationArg = m_compiler->gtNewLclvNode(m_compiler->lvaAsyncContinuationArg, TYP_REF); + GenTree* storeReusable = m_compiler->gtNewStoreLclVarNode(m_reuseContinuationVar, continuationArg); + LIR::AsRange(onContinuationBB).InsertAtBeginning(continuationArg, storeReusable); } } diff --git a/src/coreclr/jit/async.h b/src/coreclr/jit/async.h index fb9abf2d8be779..97c1a96ed3b2be 100644 --- a/src/coreclr/jit/async.h +++ b/src/coreclr/jit/async.h @@ -1,6 +1,25 @@ // Licensed to the .NET Foundation under one or more agreements. // The .NET Foundation licenses this file to you under the MIT license. +enum class ContinuationMemberType +{ + CustomAwaiterOfLayout, +}; + +struct ContinuationMember +{ + ContinuationMemberType Type; + + void Print() const; +private: + ClassLayout* m_customAwaiterLayout; + +public: + ClassLayout* GetCustomAwaiterLayout() const; + static ContinuationMember CustomAwaiterOfLayout(ClassLayout* layout); + static bool AreCompatible(const ContinuationMember& a, const ContinuationMember& b); +}; + struct ReturnTypeInfo { var_types ReturnType = TYP_UNDEF; @@ -121,7 +140,7 @@ struct ContinuationLayoutBuilder static bool Equals(const ContinuationLayoutBuilder& a, const ContinuationLayoutBuilder& b); - struct ContinuationLayout* Create(); + struct ContinuationLayout* Create(ArrayStack& continuationMemberOffsets); static ContinuationLayoutBuilder* CreateSharedLayout(Compiler* comp, const jitstd::vector& states); @@ -137,11 +156,13 @@ struct ContinuationLayout unsigned ExecutionContextOffset = UINT_MAX; jitstd::vector Locals; jitstd::vector Returns; + jitstd::vector ContinuationMemberOffsets; CORINFO_CLASS_HANDLE ClassHnd = NO_CLASS_HANDLE; ContinuationLayout(Compiler* comp) : Locals(comp->getAllocator(CMK_Async)) , Returns(comp->getAllocator(CMK_Async)) + , ContinuationMemberOffsets(comp->getAllocator(CMK_Async)) { } @@ -373,7 +394,8 @@ class AsyncTransformation unsigned m_sharedFinishContextHandlingSyncContextVar = BAD_VAR_NUM; AggregatedAwaitInfo FindAwaits(ArrayStack& blocksWithNormalAwaits, - ArrayStack& blocksWithTailAwaits); + ArrayStack& blocksWithTailAwaits, + ArrayStack& continuationMemberOffsets); void TransformTailAwaits(ArrayStack& blocksWithTailAwaits); void TransformTailAwait(BasicBlock* block, GenTreeCall* call, BasicBlock** remainder); @@ -454,6 +476,10 @@ class AsyncTransformation const ContinuationLayoutBuilder& subLayout, SuspensionContextHelper helper); void RestoreContexts(BasicBlock* block, GenTreeCall* call, BasicBlock* insertionBB); + void StoreAsyncAwaiter(BasicBlock* callBlock, + GenTreeCall* call, + BasicBlock* suspendBB, + const ContinuationLayout& layout); void CreateCheckAndSuspendAfterCall(BasicBlock* block, GenTreeCall* call, const CallDefinitionInfo& callDefInfo, @@ -488,25 +514,24 @@ class AsyncTransformation var_types storeType, GenTreeFlags indirFlags = GTF_IND_NONFAULTING); - void CreateDebugInfoForSuspensionPoint(const ContinuationLayout& layout, - const ContinuationLayoutBuilder& subLayout); - unsigned GetReturnedContinuationVar(); - unsigned GetNewContinuationVar(); - unsigned GetResultBaseVar(); - unsigned GetExceptionVar(); - void CreateSharedReturnBB(); - BasicBlock* CreateSharedFinishContextHandlingBB(SuspensionContextHelper helper, - const ContinuationLayout& layout, - bool execContextMayVary, - bool syncContextMayVary); - void InsertFinishContextHandlingCall(BasicBlock* block, - const ContinuationLayout& layout, - SuspensionContextHelper helper, - GenTree* execContext, - GenTree* syncContext); - bool ReuseContinuations(); - void CreateResumptionsAndSuspensions(); - void CreateResumptionSwitch(); + void CreateDebugInfoForSuspensionPoint(const ContinuationLayout& layout, + const ContinuationLayoutBuilder& subLayout); + unsigned GetReturnedContinuationVar(); + unsigned GetNewContinuationVar(); + unsigned GetResultBaseVar(); + unsigned GetExceptionVar(); + void CreateSharedReturnBB(); + BasicBlock* CreateSharedFinishContextHandlingBB(SuspensionContextHelper helper, + const ContinuationLayout& layout, + bool execContextMayVary, + bool syncContextMayVary); + void InsertFinishContextHandlingCall(BasicBlock* block, + const ContinuationLayout& layout, + SuspensionContextHelper helper, + GenTree* execContext, + GenTree* syncContext); + const ContinuationLayout* CreateResumptionsAndSuspensions(ArrayStack& continuationMemberOffsets); + void CreateResumptionSwitch(); public: AsyncTransformation(Compiler* comp) diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h index 56f79d07641ec3..b54967c66114c3 100644 --- a/src/coreclr/jit/compiler.h +++ b/src/coreclr/jit/compiler.h @@ -95,6 +95,7 @@ enum class WasmValueType : unsigned; #ifdef DEBUG struct IndentStack; #endif +struct ContinuationMember; class Lowering; // defined in lower.h @@ -5246,6 +5247,15 @@ class Compiler unsigned clsFlags, bool isReadonlyCall); + void impTryOptimizeAwaitAwaiter(GenTreeCall* call, + CORINFO_RESOLVED_TOKEN* pResolvedToken, + CORINFO_SIG_INFO* sig, + CORINFO_CALL_INFO* callInfo, + CORINFO_METHOD_HANDLE* methHnd, + CORINFO_CONTEXT_HANDLE* exactContextHnd, + GenTree** instParam, + NamedIntrinsic ni); + void impSetupAsyncCall(GenTreeCall* call, CORINFO_METHOD_HANDLE methHnd, OPCODE opcode, unsigned prefixFlags, NamedIntrinsic ni, const DebugInfo& callDI); void impInsertAsyncArgsForLdvirtftnCall(GenTreeCall* call); @@ -6237,6 +6247,11 @@ class Compiler PhaseStatus placeLoopAlignInstructions(); #endif + jitstd::vector* m_asyncContinuationMembers = nullptr; + size_t GetContinuationMemberIndex(const ContinuationMember& member); + size_t GetContinuationMemberCount() const; + const ContinuationMember& GetContinuationMember(size_t index); + PhaseStatus SaveAsyncContexts(); void AddContextArgsToAsyncCalls(BasicBlock* block); BasicBlock* CreateReturnBB(unsigned* mergedReturnLcl); diff --git a/src/coreclr/jit/gentree.cpp b/src/coreclr/jit/gentree.cpp index f050c5eeeee885..f2ef3f32885b03 100644 --- a/src/coreclr/jit/gentree.cpp +++ b/src/coreclr/jit/gentree.cpp @@ -13,6 +13,7 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX #include "jitpch.h" #include "hwintrinsic.h" #include "simd.h" +#include "async.h" #ifdef _MSC_VER #pragma hdrstop @@ -1648,6 +1649,7 @@ bool CallArgs::GetCustomRegister(Compiler* comp, CorInfoCallConvExtension cc, We #endif // HAS_FIXED_REGISTER_SET case WellKnownArg::StackArrayLocal: + case WellKnownArg::AsyncAwaiter: case WellKnownArg::AsyncExecutionContext: case WellKnownArg::AsyncSynchronizationContext: // These are pseudo-args; they are not actual arguments, but we @@ -11917,6 +11919,7 @@ GenTreeUseEdgeIterator::GenTreeUseEdgeIterator(GenTree* node) case GT_CATCH_ARG: case GT_ASYNC_CONTINUATION: case GT_ASYNC_RESUME_INFO: + case GT_CONTINUATION_MEMBER_OFFSET: case GT_LABEL: case GT_FTN_ADDR: case GT_FTN_ENTRY: @@ -14008,6 +14011,15 @@ void Compiler::gtDispLeaf(GenTree* tree, IndentStack* indentStack) case GT_WASM_JEXCEPT: break; + case GT_CONTINUATION_MEMBER_OFFSET: + { + printf(" index=%zu ", tree->AsVal()->gtVal1); + + const ContinuationMember& member = GetContinuationMember(tree->AsVal()->gtVal1); + member.Print(); + break; + } + case GT_RET_EXPR: { GenTreeCall* inlineCand = tree->AsRetExpr()->gtInlineCandidate; diff --git a/src/coreclr/jit/gtlist.h b/src/coreclr/jit/gtlist.h index be81ca9f22052d..42ff6a03a84018 100644 --- a/src/coreclr/jit/gtlist.h +++ b/src/coreclr/jit/gtlist.h @@ -41,6 +41,12 @@ GTNODE(GCPOLL , GenTree ,0,0,GTK_LEAF|GTK_NOVALUE|DBK_NOTLI GTNODE(ASYNC_RESUME_INFO, GenTreeVal ,0,0,GTK_LEAF) // Address of async resume info for a state GTNODE(FTN_ENTRY , GenTree ,0,0,GTK_LEAF) // Address of this function's entry point +// Offset into continuation, starting after the MethodTable*, of a special +// member. GenTreeVal::gtVal1 is an index into Compiler::m_asyncContinuationMembers. +// The async transformation allocates and lays out the continuation and +// substitutes these by constants at that time. +GTNODE(CONTINUATION_MEMBER_OFFSET, GenTreeVal ,0,0,GTK_LEAF) + //----------------------------------------------------------------------------- // Constant nodes: //----------------------------------------------------------------------------- diff --git a/src/coreclr/jit/gtstructs.h b/src/coreclr/jit/gtstructs.h index 59cd282c9544e1..bec1093d34d0b9 100644 --- a/src/coreclr/jit/gtstructs.h +++ b/src/coreclr/jit/gtstructs.h @@ -50,7 +50,7 @@ GTSTRUCT_0(UnOp , GT_OP) GTSTRUCT_0(Op , GT_OP) -GTSTRUCT_N(Val , GT_JMP, GT_RECORD_ASYNC_RESUME, GT_ASYNC_RESUME_INFO) +GTSTRUCT_N(Val , GT_JMP, GT_RECORD_ASYNC_RESUME, GT_ASYNC_RESUME_INFO, GT_CONTINUATION_MEMBER_OFFSET) GTSTRUCT_2_SPECIAL(IntConCommon, GT_CNS_INT, GT_CNS_LNG) GTSTRUCT_1(IntCon , GT_CNS_INT) GTSTRUCT_1(LngCon , GT_CNS_LNG) diff --git a/src/coreclr/jit/importercalls.cpp b/src/coreclr/jit/importercalls.cpp index f9a583a850e2c8..1c49f4540a1fd0 100644 --- a/src/coreclr/jit/importercalls.cpp +++ b/src/coreclr/jit/importercalls.cpp @@ -2,6 +2,7 @@ // The .NET Foundation licenses this file to you under the MIT license. #include "jitpch.h" +#include "async.h" //------------------------------------------------------------------------ // impGetInstParamArg: compute the hidden instantiation / generic-context argument @@ -127,6 +128,79 @@ GenTree* Compiler::impGetInstParamArg(CORINFO_RESOLVED_TOKEN* pResolvedToken, return instParam; } +//------------------------------------------------------------------------ +// impTryOptimizeAwaitAwaiter: +// Rewrite a struct AwaitAwaiter or UnsafeAwaitAwaiter call to use the +// corresponding helper that reads the awaiter from the continuation. +// +void Compiler::impTryOptimizeAwaitAwaiter(GenTreeCall* call, + CORINFO_RESOLVED_TOKEN* pResolvedToken, + CORINFO_SIG_INFO* sig, + CORINFO_CALL_INFO* callInfo, + CORINFO_METHOD_HANDLE* methHnd, + CORINFO_CONTEXT_HANDLE* exactContextHnd, + GenTree** instParam, + NamedIntrinsic ni) +{ + CallArg* awaiterArg = call->gtArgs.GetUserArgByIndex(0); + if (!varTypeIsStruct(awaiterArg->GetSignatureType())) + { + return; + } + + CORINFO_LOOKUP newInstArgLookup; + CORINFO_SIG_INFO lookupSig = *sig; + if (pResolvedToken->pMethodSpec != nullptr) + { + lookupSig.pSig = pResolvedToken->pMethodSpec; + lookupSig.cbSig = pResolvedToken->cbMethodSpec; + lookupSig.scope = pResolvedToken->tokenScope; + } + + bool isUnsafe = ni == NI_System_Runtime_CompilerServices_AsyncHelpers_UnsafeAwaitAwaiter; + CORINFO_CONTEXT_HANDLE newExactContextHnd; + CORINFO_METHOD_HANDLE newMethod = + info.compCompHnd->getAwaitAwaiterInContinuationCall(info.compMethodHnd, &lookupSig, isUnsafe, + &newExactContextHnd, &newInstArgLookup); + + CORINFO_SIG_INFO newSig; + info.compCompHnd->getMethodSig(newMethod, &newSig); + + GenTree* newInstParam = nullptr; + if (newSig.hasTypeArg()) + { + newInstParam = impLookupToTree(&newInstArgLookup, GTF_ICON_METHOD_HDL, newMethod); + if (newInstParam == nullptr) + { + JITDUMP("Failed to optimize awaiter call [%06u] because its replacement lookup could not be created\n", + dspTreeID(call)); + return; + } + } + + JITDUMP("Optimizing awaiter call [%06u] to read its struct awaiter from the continuation\n", dspTreeID(call)); + + *methHnd = newMethod; + *exactContextHnd = newExactContextHnd; + call->gtCallMethHnd = newMethod; + callInfo->hMethod = newMethod; + callInfo->methodFlags = info.compCompHnd->getMethodAttribs(newMethod); + callInfo->sig = newSig; + *instParam = newInstParam; + + GenTree* awaiter = awaiterArg->GetNode(); + var_types awaiterType = awaiterArg->GetSignatureType(); + ClassLayout* awaiterLayout = awaiterArg->GetSignatureLayout(); + call->gtArgs.Remove(awaiterArg); + call->gtArgs + .PushFront(this, NewCallArg::Struct(awaiter, awaiterType, awaiterLayout).WellKnown(WellKnownArg::AsyncAwaiter)); + + size_t memberIndex = GetContinuationMemberIndex(ContinuationMember::CustomAwaiterOfLayout(awaiterLayout)); + GenTree* offset = + new (this, GT_CONTINUATION_MEMBER_OFFSET) GenTreeVal(GT_CONTINUATION_MEMBER_OFFSET, TYP_INT, memberIndex); + call->gtArgs.PushBack(this, NewCallArg::Primitive(offset)); +} + //------------------------------------------------------------------------ // impImportCall: import a call-inspiring opcode // @@ -952,6 +1026,13 @@ var_types Compiler::impImportCall(OPCODE opcode, impPopCallArgs(sig, call->AsCall()); + if (opts.OptimizationEnabled() && ((ni == NI_System_Runtime_CompilerServices_AsyncHelpers_AwaitAwaiter) || + (ni == NI_System_Runtime_CompilerServices_AsyncHelpers_UnsafeAwaitAwaiter))) + { + impTryOptimizeAwaitAwaiter(call->AsCall(), pResolvedToken, sig, callInfo, &methHnd, &exactContextHnd, + &instParam, ni); + } + // Extra args if ((instParam != nullptr) || (asyncContinuation != nullptr) || (varArgsCookie != nullptr)) { diff --git a/src/coreclr/jit/jitconfigvalues.h b/src/coreclr/jit/jitconfigvalues.h index 5e3095e94dd87c..bfb21eaf40c916 100644 --- a/src/coreclr/jit/jitconfigvalues.h +++ b/src/coreclr/jit/jitconfigvalues.h @@ -620,12 +620,6 @@ OPT_CONFIG_STRING(JitAsyncDefaultValueAnalysisRange, // a continuation is being reused. OPT_CONFIG_STRING(JitAsyncPreservedValueAnalysisRange, "JitAsyncPreservedValueAnalysisRange") -// Enable continuation reuse based on method hash range -OPT_CONFIG_STRING(JitAsyncReuseContinuationsRange, "JitAsyncReuseContinuationsRange") -// Save and reuse continuation instances in runtime async functions. Also -// implies use of shared continuation layouts for all suspension points. -RELEASE_CONFIG_INTEGER(JitAsyncReuseContinuations, "JitAsyncReuseContinuations", 1) - RELEASE_CONFIG_INTEGER(JitEnableOptRepeat, "JitEnableOptRepeat", 1) // If zero, do not allow JitOptRepeat RELEASE_CONFIG_METHODSET(JitOptRepeat, "JitOptRepeat") // Runs optimizer multiple times on specified methods RELEASE_CONFIG_INTEGER(JitOptRepeatCount, "JitOptRepeatCount", 2) // Number of times to repeat opts when repeating diff --git a/src/coreclr/jit/morph.cpp b/src/coreclr/jit/morph.cpp index 39e2c9f0eea43e..b7ef9b9185585f 100644 --- a/src/coreclr/jit/morph.cpp +++ b/src/coreclr/jit/morph.cpp @@ -2214,6 +2214,19 @@ bool Compiler::fgTryMorphStructArg(CallArg* arg) GenTree* argNode = *use; assert(varTypeIsStruct(argNode)); + if (arg->AbiInfo.NumSegments == 0) + { + // Pseudo arg. One case is WellKnownArg::AsyncAwaiter. We just handle + // these as arbitrary struct operands that can be expanded into + // FIELD_LIST. The async transformation will later store the value into + // the continuation, so FIELD_LIST allows using decomposed stores. + if (fgTryReplaceStructLocalWithFields(&arg->NodeRef())) + { + arg->GetNode()->SetMorphed(this, true); + } + return true; + } + bool isSplit = arg->AbiInfo.IsSplitAcrossRegistersAndStack(); #ifdef TARGET_ARM if ((isSplit && (arg->AbiInfo.CountRegsAndStackSlots() > 4)) || (!isSplit && arg->AbiInfo.HasAnyStackSegment())) diff --git a/src/coreclr/jit/valuenum.cpp b/src/coreclr/jit/valuenum.cpp index 5b08ab7fea7142..6ab8c7ba469efb 100644 --- a/src/coreclr/jit/valuenum.cpp +++ b/src/coreclr/jit/valuenum.cpp @@ -13153,6 +13153,7 @@ void Compiler::fgValueNumberTree(GenTree* tree) case GT_CATCH_ARG: case GT_ASYNC_CONTINUATION: + case GT_CONTINUATION_MEMBER_OFFSET: case GT_SWIFT_ERROR: // We know nothing about the value of these. tree->gtVNPair.SetBoth(vnStore->VNForExpr(compCurBB, tree->TypeGet())); diff --git a/src/coreclr/jit/wellknownargs.h b/src/coreclr/jit/wellknownargs.h index a963226292cbbc..fd7318171bc4b4 100644 --- a/src/coreclr/jit/wellknownargs.h +++ b/src/coreclr/jit/wellknownargs.h @@ -36,6 +36,7 @@ WELL_KNOWN_ARG(SwiftSelf, "swift self", true, false) WELL_KNOWN_ARG(X86TailCallSpecialArg, "tail call", false, false) WELL_KNOWN_ARG(StackArrayLocal, "&lcl arr", false, false) WELL_KNOWN_ARG(RuntimeMethodHandle, "meth hnd", false, false) +WELL_KNOWN_ARG(AsyncAwaiter, "awaiter", false, false) WELL_KNOWN_ARG(AsyncExecutionContext, "exec ctx", false, false) WELL_KNOWN_ARG(AsyncSynchronizationContext, "sync ctx", false, false) WELL_KNOWN_ARG(WasmShadowStackPointer, "wasm sp", false, false) diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs index 329abbf03b7f78..2cc004d9c64ab7 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs @@ -3674,6 +3674,61 @@ private void getWasmWellKnownGlobals(ref CORINFO_WASM_WELLKNOWN_GLOBALS pWellKno return ObjectToHandle(targetMethod); } + private CORINFO_METHOD_STRUCT_* getAwaitAwaiterInContinuationCall( + CORINFO_METHOD_STRUCT_* callerHandle, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_STRUCT** contextHandle, + ref CORINFO_LOOKUP instArg) + { + Debug.Assert(callSig->sigInst.methInstCount == 1); + + instArg.lookupKind.needsRuntimeLookup = false; + instArg.constLookup.accessType = InfoAccessType.IAT_VALUE; + instArg.constLookup.addr = null; + + MethodDesc caller = HandleToObject(callerHandle); + TypeDesc awaiterType = HandleToObject(callSig->sigInst.methInst[0]); + CompilerTypeSystemContext context = _compilation.TypeSystemContext; + DefType asyncHelpers = + context.SystemModule.GetKnownType("System.Runtime.CompilerServices"u8, "AsyncHelpers"u8); + MethodSignature signature = new MethodSignature( + MethodSignatureFlags.Static, + 1, + context.GetWellKnownType(WellKnownType.Void), + [context.GetWellKnownType(WellKnownType.Int32)]); + MethodDesc result = asyncHelpers + .GetKnownMethod( + isUnsafe ? "UnsafeAwaitAwaiterInContinuation"u8 : "AwaitAwaiterInContinuation"u8, + signature) + .MakeInstantiatedMethod(awaiterType); + MethodDesc targetMethod = result.GetCanonMethodTarget(CanonicalFormKind.Specific); + *contextHandle = contextFromMethod(result); + + if (targetMethod.RequiresInstArg()) + { +#if READYTORUN + instArg.constLookup = CreateConstLookupToSymbol( + _compilation.SymbolNodeFactory.CreateReadyToRunHelper( + ReadyToRunHelperId.MethodDictionary, + new MethodWithToken( + result, + _compilation.NodeFactory.Resolver.GetModuleTokenForMethod( + result, + allowDynamicallyCreatedReference: true, + throwIfNotFound: true), + constrainedType: null, + unboxing: false, + genericContextObject: caller))); +#else + ComputeLookup(caller != MethodBeingCompiled, result, ReadyToRunHelperId.MethodDictionary, caller, + ref instArg); +#endif + } + + return ObjectToHandle(targetMethod); + } + private CORINFO_CLASS_STRUCT_* getContinuationType(nuint dataSize, ref bool objRefs, nuint objRefsSize) { Debug.Assert(objRefsSize == (dataSize + (nuint)(PointerSize - 1)) / (nuint)PointerSize); diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoImpl_generated.cs b/src/coreclr/tools/Common/JitInterface/CorInfoImpl_generated.cs index ef00dcf0188b1b..f3389802e9ad5b 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoImpl_generated.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoImpl_generated.cs @@ -141,6 +141,7 @@ static ICorJitInfoCallbacks() s_callbacks.getEEInfo = &_getEEInfo; s_callbacks.getAsyncInfo = &_getAsyncInfo; s_callbacks.getAwaitReturnCall = &_getAwaitReturnCall; + s_callbacks.getAwaitAwaiterInContinuationCall = &_getAwaitAwaiterInContinuationCall; s_callbacks.getMethodDefFromMethod = &_getMethodDefFromMethod; s_callbacks.printMethodName = &_printMethodName; s_callbacks.getMethodNameFromMetadata = &_getMethodNameFromMetadata; @@ -325,6 +326,7 @@ static ICorJitInfoCallbacks() public delegate* unmanaged getEEInfo; public delegate* unmanaged getAsyncInfo; public delegate* unmanaged getAwaitReturnCall; + public delegate* unmanaged getAwaitAwaiterInContinuationCall; public delegate* unmanaged getMethodDefFromMethod; public delegate* unmanaged printMethodName; public delegate* unmanaged getMethodNameFromMetadata; @@ -2178,6 +2180,21 @@ private static void _getAsyncInfo(IntPtr thisHandle, IntPtr* ppException, CORINF } } + [UnmanagedCallersOnly] + private static CORINFO_METHOD_STRUCT_* _getAwaitAwaiterInContinuationCall(IntPtr thisHandle, IntPtr* ppException, CORINFO_METHOD_STRUCT_* callerHandle, CORINFO_SIG_INFO* callSig, byte isUnsafe, CORINFO_CONTEXT_STRUCT** contextHandle, CORINFO_LOOKUP* instArg) + { + var _this = GetThis(thisHandle); + try + { + return _this.getAwaitAwaiterInContinuationCall(callerHandle, callSig, isUnsafe != 0, contextHandle, ref *instArg); + } + catch (Exception ex) + { + *ppException = _this.AllocException(ex); + return default; + } + } + [UnmanagedCallersOnly] private static mdToken _getMethodDefFromMethod(IntPtr thisHandle, IntPtr* ppException, CORINFO_METHOD_STRUCT_* hMethod) { diff --git a/src/coreclr/tools/Common/JitInterface/ThunkGenerator/ThunkInput.txt b/src/coreclr/tools/Common/JitInterface/ThunkGenerator/ThunkInput.txt index 8a6ab34f947c4b..8368a69891a2d9 100644 --- a/src/coreclr/tools/Common/JitInterface/ThunkGenerator/ThunkInput.txt +++ b/src/coreclr/tools/Common/JitInterface/ThunkGenerator/ThunkInput.txt @@ -296,6 +296,7 @@ FUNCTIONS void getEEInfo(CORINFO_EE_INFO* pEEInfoOut); void getAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut); CORINFO_METHOD_HANDLE getAwaitReturnCall(CORINFO_METHOD_HANDLE callerHandle, CORINFO_CONTEXT_HANDLE* contextHandle, CORINFO_LOOKUP* instArg); + CORINFO_METHOD_HANDLE getAwaitAwaiterInContinuationCall(CORINFO_METHOD_HANDLE callerHandle, CORINFO_SIG_INFO* callSig, bool isUnsafe, CORINFO_CONTEXT_HANDLE* contextHandle, CORINFO_LOOKUP* instArg); mdMethodDef getMethodDefFromMethod(CORINFO_METHOD_HANDLE hMethod); size_t printMethodName(CORINFO_METHOD_HANDLE ftn, char* buffer, size_t bufferSize, size_t* pRequiredBufferSize) const char* getMethodNameFromMetadata(CORINFO_METHOD_HANDLE ftn, const char **className, const char **namespaceName, const char **enclosingClassNames, size_t maxEnclosingClassNames); diff --git a/src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/ReadyToRunCodegenCompilation.cs b/src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/ReadyToRunCodegenCompilation.cs index 5834498f43ea3e..1a7a9012ecd2c8 100644 --- a/src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/ReadyToRunCodegenCompilation.cs +++ b/src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/ReadyToRunCodegenCompilation.cs @@ -1064,6 +1064,8 @@ private void AddNecessaryAsyncReferences(MethodDesc method) asyncHelpers.GetKnownMethod("TransparentAwait"u8, new MethodSignature(MethodSignatureFlags.Static, 0, voidType, [valueTaskType])), asyncHelpers.GetKnownMethod("TransparentAwait"u8, new MethodSignature(MethodSignatureFlags.Static, 1, methodVar, [taskOfTType.MakeInstantiatedType(methodVar)])), asyncHelpers.GetKnownMethod("TransparentAwait"u8, new MethodSignature(MethodSignatureFlags.Static, 1, methodVar, [valueTaskOfTType.MakeInstantiatedType(methodVar)])), + asyncHelpers.GetKnownMethod("AwaitAwaiterInContinuation"u8, null), + asyncHelpers.GetKnownMethod("UnsafeAwaitAwaiterInContinuation"u8, null), ]; var moduleForNewReferences = ((EcmaMethod)method.GetPrimaryMethodDesc().GetTypicalMethodDefinition()).Module; _tokenManager.EnsureDefTokensAreAvailable([..requiredMethods, ..requiredTypes, ..requiredFields], moduleForNewReferences, true); diff --git a/src/coreclr/tools/aot/jitinterface/jitinterface_generated.h b/src/coreclr/tools/aot/jitinterface/jitinterface_generated.h index 3a0514e49ad2a7..9f036e818e547d 100644 --- a/src/coreclr/tools/aot/jitinterface/jitinterface_generated.h +++ b/src/coreclr/tools/aot/jitinterface/jitinterface_generated.h @@ -132,6 +132,7 @@ struct JitInterfaceCallbacks void (* getEEInfo)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_EE_INFO* pEEInfoOut); void (* getAsyncInfo)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_ASYNC_INFO* pAsyncInfoOut); CORINFO_METHOD_HANDLE (* getAwaitReturnCall)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_METHOD_HANDLE callerHandle, CORINFO_CONTEXT_HANDLE* contextHandle, CORINFO_LOOKUP* instArg); + CORINFO_METHOD_HANDLE (* getAwaitAwaiterInContinuationCall)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_METHOD_HANDLE callerHandle, CORINFO_SIG_INFO* callSig, bool isUnsafe, CORINFO_CONTEXT_HANDLE* contextHandle, CORINFO_LOOKUP* instArg); mdMethodDef (* getMethodDefFromMethod)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_METHOD_HANDLE hMethod); size_t (* printMethodName)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_METHOD_HANDLE ftn, char* buffer, size_t bufferSize, size_t* pRequiredBufferSize); const char* (* getMethodNameFromMetadata)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_METHOD_HANDLE ftn, const char** className, const char** namespaceName, const char** enclosingClassNames, size_t maxEnclosingClassNames); @@ -1374,6 +1375,19 @@ class JitInterfaceWrapper : public ICorJitInfo return temp; } + virtual CORINFO_METHOD_HANDLE getAwaitAwaiterInContinuationCall( + CORINFO_METHOD_HANDLE callerHandle, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg) +{ + CorInfoExceptionClass* pException = nullptr; + CORINFO_METHOD_HANDLE temp = _callbacks->getAwaitAwaiterInContinuationCall(_thisHandle, &pException, callerHandle, callSig, isUnsafe, contextHandle, instArg); + if (pException != nullptr) throw pException; + return temp; +} + virtual mdMethodDef getMethodDefFromMethod( CORINFO_METHOD_HANDLE hMethod) { diff --git a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h index 2c36e6500c75ac..0f64d7436d779f 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h +++ b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h @@ -255,6 +255,13 @@ struct Agnostic_GetAwaitReturnCallResult Agnostic_CORINFO_LOOKUP instArg; }; +struct Agnostic_GetAwaitAwaiterInContinuationCall +{ + DWORDLONG callerHnd; + Agnostic_CORINFO_SIG_INFO callSig; + DWORD isUnsafe; +}; + struct Agnostic_GetOSRInfo { DWORD index; diff --git a/src/coreclr/tools/superpmi/superpmi-shared/lwmlist.h b/src/coreclr/tools/superpmi/superpmi-shared/lwmlist.h index e44cf227fbd412..c3f540189da97c 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/lwmlist.h +++ b/src/coreclr/tools/superpmi/superpmi-shared/lwmlist.h @@ -80,6 +80,7 @@ LWM(GetDelegateCtor, Agnostic_GetDelegateCtorIn, Agnostic_GetDelegateCtorOut) LWM(GetEEInfo, DWORD, Agnostic_CORINFO_EE_INFO) LWM(GetAsyncInfo, DWORD, Agnostic_CORINFO_ASYNC_INFO) LWM(GetAwaitReturnCall, DWORDLONG, Agnostic_GetAwaitReturnCallResult) +LWM(GetAwaitAwaiterInContinuationCall, Agnostic_GetAwaitAwaiterInContinuationCall, Agnostic_GetAwaitReturnCallResult) LWM(GetEHinfo, DLD, Agnostic_CORINFO_EH_CLAUSE) LWM(GetStaticFieldContent, DLDDD, DD) LWM(GetObjectContent, DLDD, DD) diff --git a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp index db89aaed33f8a7..2bf7973f50d754 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp +++ b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp @@ -4483,6 +4483,73 @@ CORINFO_METHOD_HANDLE MethodContext::repGetAwaitReturnCall(CORINFO_METHOD_HANDLE return (CORINFO_METHOD_HANDLE)result.methodHnd; } +void MethodContext::recGetAwaitAwaiterInContinuationCall(CORINFO_METHOD_HANDLE callerHnd, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg, + CORINFO_METHOD_HANDLE methHnd) +{ + if (GetAwaitAwaiterInContinuationCall == nullptr) + { + GetAwaitAwaiterInContinuationCall = + new LightWeightMap(); + } + + Agnostic_GetAwaitAwaiterInContinuationCall key; + ZeroMemory(&key, sizeof(key)); + key.callerHnd = CastHandle(callerHnd); + key.callSig = + SpmiRecordsHelper::StoreAgnostic_CORINFO_SIG_INFO(*callSig, GetAwaitAwaiterInContinuationCall, SigInstHandleMap); + key.isUnsafe = isUnsafe; + + Agnostic_GetAwaitReturnCallResult value; + ZeroMemory(&value, sizeof(value)); + value.methodHnd = CastHandle(methHnd); + value.contextHandle = CastHandle(*contextHandle); + value.instArg = SpmiRecordsHelper::StoreAgnostic_CORINFO_LOOKUP(instArg); + + GetAwaitAwaiterInContinuationCall->Add(key, value); + DEBUG_REC(dmpGetAwaitAwaiterInContinuationCall(key, value)); +} + +void MethodContext::dmpGetAwaitAwaiterInContinuationCall( + const Agnostic_GetAwaitAwaiterInContinuationCall& key, + Agnostic_GetAwaitReturnCallResult& value) +{ + printf("GetAwaitAwaiterInContinuationCall caller-%016" PRIX64 " sig-%s unsafe-%u " + "methodHnd-%016" PRIX64 " contextHandle-%016" PRIX64 " instArg %s", + key.callerHnd, + SpmiDumpHelper::DumpAgnostic_CORINFO_SIG_INFO( + key.callSig, GetAwaitAwaiterInContinuationCall, SigInstHandleMap).c_str(), + key.isUnsafe, + value.methodHnd, + value.contextHandle, + SpmiDumpHelper::DumpAgnostic_CORINFO_LOOKUP(value.instArg).c_str()); +} + +CORINFO_METHOD_HANDLE MethodContext::repGetAwaitAwaiterInContinuationCall( + CORINFO_METHOD_HANDLE callerHnd, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg) +{ + Agnostic_GetAwaitAwaiterInContinuationCall key; + ZeroMemory(&key, sizeof(key)); + key.callerHnd = CastHandle(callerHnd); + key.callSig = SpmiRecordsHelper::RestoreAgnostic_CORINFO_SIG_INFO( + *callSig, GetAwaitAwaiterInContinuationCall, SigInstHandleMap); + key.isUnsafe = isUnsafe; + + Agnostic_GetAwaitReturnCallResult value = + LookupByKeyOrMissNoMessage(GetAwaitAwaiterInContinuationCall, key); + DEBUG_REP(dmpGetAwaitAwaiterInContinuationCall(key, value)); + *contextHandle = (CORINFO_CONTEXT_HANDLE)value.contextHandle; + *instArg = SpmiRecordsHelper::RestoreCORINFO_LOOKUP(value.instArg); + return (CORINFO_METHOD_HANDLE)value.methodHnd; +} + void MethodContext::recGetGSCookie(GSCookie* pCookieVal, GSCookie** ppCookieVal) { if (GetGSCookie == nullptr) diff --git a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.h b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.h index 5ca689f36ce384..0fa494f306e8b8 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.h +++ b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.h @@ -563,6 +563,20 @@ class MethodContext void dmpGetAwaitReturnCall(DWORDLONG key, Agnostic_GetAwaitReturnCallResult& value); CORINFO_METHOD_HANDLE repGetAwaitReturnCall(CORINFO_METHOD_HANDLE callerHnd, CORINFO_CONTEXT_HANDLE* contextHandle, CORINFO_LOOKUP* instArg); + void recGetAwaitAwaiterInContinuationCall(CORINFO_METHOD_HANDLE callerHnd, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg, + CORINFO_METHOD_HANDLE methHnd); + void dmpGetAwaitAwaiterInContinuationCall(const Agnostic_GetAwaitAwaiterInContinuationCall& key, + Agnostic_GetAwaitReturnCallResult& value); + CORINFO_METHOD_HANDLE repGetAwaitAwaiterInContinuationCall(CORINFO_METHOD_HANDLE callerHnd, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg); + void recGetWasmWellKnownGlobals(const CORINFO_WASM_WELLKNOWN_GLOBALS* pBaseGlobals); void dmpGetWasmWellKnownGlobals(DWORD key, const Agnostic_CORINFO_WASM_WELLKNOWN_GLOBALS& value); void repGetWasmWellKnownGlobals(CORINFO_WASM_WELLKNOWN_GLOBALS* pWellKnownGlobalsOut); @@ -1224,6 +1238,7 @@ enum mcPackets Packet_GetAwaitReturnCall = 238, Packet_GetAddressAlignment = 239, Packet_GetWasmWellKnownGlobals = 240, + Packet_GetAwaitAwaiterInContinuationCall = 241, }; void SetDebugDumpVariables(); diff --git a/src/coreclr/tools/superpmi/superpmi-shim-collector/icorjitinfo.cpp b/src/coreclr/tools/superpmi/superpmi-shim-collector/icorjitinfo.cpp index 9744fc70aaa470..e9ade6d19548a2 100644 --- a/src/coreclr/tools/superpmi/superpmi-shim-collector/icorjitinfo.cpp +++ b/src/coreclr/tools/superpmi/superpmi-shim-collector/icorjitinfo.cpp @@ -1396,6 +1396,20 @@ CORINFO_METHOD_HANDLE interceptor_ICJI::getAwaitReturnCall(CORINFO_METHOD_HANDLE return result; } +CORINFO_METHOD_HANDLE interceptor_ICJI::getAwaitAwaiterInContinuationCall( + CORINFO_METHOD_HANDLE callerHandle, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg) +{ + mc->cr->AddCall("getAwaitAwaiterInContinuationCall"); + CORINFO_METHOD_HANDLE result = original_ICorJitInfo->getAwaitAwaiterInContinuationCall( + callerHandle, callSig, isUnsafe, contextHandle, instArg); + mc->recGetAwaitAwaiterInContinuationCall(callerHandle, callSig, isUnsafe, contextHandle, instArg, result); + return result; +} + /*********************************************************************************/ // // Diagnostic methods diff --git a/src/coreclr/tools/superpmi/superpmi-shim-counter/icorjitinfo_generated.cpp b/src/coreclr/tools/superpmi/superpmi-shim-counter/icorjitinfo_generated.cpp index 53166c4d1538d6..1dc724f10fa38a 100644 --- a/src/coreclr/tools/superpmi/superpmi-shim-counter/icorjitinfo_generated.cpp +++ b/src/coreclr/tools/superpmi/superpmi-shim-counter/icorjitinfo_generated.cpp @@ -978,6 +978,17 @@ CORINFO_METHOD_HANDLE interceptor_ICJI::getAwaitReturnCall( return original_ICorJitInfo->getAwaitReturnCall(callerHandle, contextHandle, instArg); } +CORINFO_METHOD_HANDLE interceptor_ICJI::getAwaitAwaiterInContinuationCall( + CORINFO_METHOD_HANDLE callerHandle, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg) +{ + mcs->AddCall("getAwaitAwaiterInContinuationCall"); + return original_ICorJitInfo->getAwaitAwaiterInContinuationCall(callerHandle, callSig, isUnsafe, contextHandle, instArg); +} + mdMethodDef interceptor_ICJI::getMethodDefFromMethod( CORINFO_METHOD_HANDLE hMethod) { diff --git a/src/coreclr/tools/superpmi/superpmi-shim-simple/icorjitinfo_generated.cpp b/src/coreclr/tools/superpmi/superpmi-shim-simple/icorjitinfo_generated.cpp index 458287f1d32ac2..199f0b6fce2631 100644 --- a/src/coreclr/tools/superpmi/superpmi-shim-simple/icorjitinfo_generated.cpp +++ b/src/coreclr/tools/superpmi/superpmi-shim-simple/icorjitinfo_generated.cpp @@ -857,6 +857,16 @@ CORINFO_METHOD_HANDLE interceptor_ICJI::getAwaitReturnCall( return original_ICorJitInfo->getAwaitReturnCall(callerHandle, contextHandle, instArg); } +CORINFO_METHOD_HANDLE interceptor_ICJI::getAwaitAwaiterInContinuationCall( + CORINFO_METHOD_HANDLE callerHandle, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg) +{ + return original_ICorJitInfo->getAwaitAwaiterInContinuationCall(callerHandle, callSig, isUnsafe, contextHandle, instArg); +} + mdMethodDef interceptor_ICJI::getMethodDefFromMethod( CORINFO_METHOD_HANDLE hMethod) { diff --git a/src/coreclr/tools/superpmi/superpmi/icorjitinfo.cpp b/src/coreclr/tools/superpmi/superpmi/icorjitinfo.cpp index e71c1dd6071c4a..e5cff9a0c71581 100644 --- a/src/coreclr/tools/superpmi/superpmi/icorjitinfo.cpp +++ b/src/coreclr/tools/superpmi/superpmi/icorjitinfo.cpp @@ -1208,6 +1208,18 @@ CORINFO_METHOD_HANDLE MyICJI::getAwaitReturnCall(CORINFO_METHOD_HANDLE callerHan return jitInstance->mc->repGetAwaitReturnCall(callerHandle, contextHandle, instArg); } +CORINFO_METHOD_HANDLE MyICJI::getAwaitAwaiterInContinuationCall( + CORINFO_METHOD_HANDLE callerHandle, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg) +{ + jitInstance->mc->cr->AddCall("getAwaitAwaiterInContinuationCall"); + return jitInstance->mc->repGetAwaitAwaiterInContinuationCall( + callerHandle, callSig, isUnsafe, contextHandle, instArg); +} + /*********************************************************************************/ // // Diagnostic methods diff --git a/src/coreclr/vm/corelib.h b/src/coreclr/vm/corelib.h index 1584e25df9f8f1..60747f18d01d72 100644 --- a/src/coreclr/vm/corelib.h +++ b/src/coreclr/vm/corelib.h @@ -723,6 +723,8 @@ DEFINE_METHOD(ASYNC_HELPERS, TRANSPARENT_AWAIT_TASK, Transpa DEFINE_METHOD(ASYNC_HELPERS, TRANSPARENT_AWAIT_VALUETASK, TransparentAwait, SM_ValueTask_RetVoid) DEFINE_METHOD(ASYNC_HELPERS, TRANSPARENT_AWAIT_TASK_OF_T, TransparentAwait, GM_TaskOfT_RetT) DEFINE_METHOD(ASYNC_HELPERS, TRANSPARENT_AWAIT_VALUETASK_OF_T, TransparentAwait, GM_ValueTaskOfT_RetT) +DEFINE_METHOD(ASYNC_HELPERS, AWAIT_AWAITER_IN_CONTINUATION, AwaitAwaiterInContinuation, NoSig) +DEFINE_METHOD(ASYNC_HELPERS, UNSAFE_AWAIT_AWAITER_IN_CONTINUATION, UnsafeAwaitAwaiterInContinuation, NoSig) DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_EXECUTION_CONTEXT, CaptureExecutionContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_CONTINUATION_CONTEXT, CaptureContinuationContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_CONTEXTS, CaptureContexts, NoSig) diff --git a/src/coreclr/vm/jitinterface.cpp b/src/coreclr/vm/jitinterface.cpp index 5bb833b15d8527..6d132f96b4688a 100644 --- a/src/coreclr/vm/jitinterface.cpp +++ b/src/coreclr/vm/jitinterface.cpp @@ -10445,6 +10445,60 @@ CORINFO_METHOD_HANDLE CEEInfo::getAwaitReturnCall(CORINFO_METHOD_HANDLE callerHa return CORINFO_METHOD_HANDLE(pMD); } +CORINFO_METHOD_HANDLE CEEInfo::getAwaitAwaiterInContinuationCall( + CORINFO_METHOD_HANDLE callerHandle, + CORINFO_SIG_INFO* callSig, + bool isUnsafe, + CORINFO_CONTEXT_HANDLE* contextHandle, + CORINFO_LOOKUP* instArg) +{ + CONTRACTL { + THROWS; + GC_TRIGGERS; + MODE_PREEMPTIVE; + } CONTRACTL_END; + + MethodDesc* pMD = NULL; + + JIT_TO_EE_TRANSITION(); + + _ASSERTE(callSig->sigInst.methInstCount == 1); + TypeHandle awaiterType(callSig->sigInst.methInst[0]); + + instArg->lookupKind.needsRuntimeLookup = false; + instArg->constLookup.accessType = IAT_VALUE; + instArg->constLookup.addr = NULL; + + MethodDesc* pCallerMD = GetMethod(callerHandle); + MethodDesc* pTypicalAwaitMD = CoreLibBinder::GetMethod( + isUnsafe ? METHOD__ASYNC_HELPERS__UNSAFE_AWAIT_AWAITER_IN_CONTINUATION + : METHOD__ASYNC_HELPERS__AWAIT_AWAITER_IN_CONTINUATION); + pMD = MethodDesc::FindOrCreateAssociatedMethodDesc(pTypicalAwaitMD, pTypicalAwaitMD->GetMethodTable(), FALSE, + Instantiation(&awaiterType, 1), TRUE); + + MethodDesc* pInliningContext = pMD; + if (pMD->RequiresInstArg()) + { + if (awaiterType.IsCanonicalSubtype()) + { + ComputeRuntimeLookupForAwaitAwaiterInContinuationCall(pCallerMD, pTypicalAwaitMD, callSig, instArg); + } + else + { + MethodDesc* pContext = MethodDesc::FindOrCreateAssociatedMethodDesc( + pTypicalAwaitMD, pTypicalAwaitMD->GetMethodTable(), FALSE, Instantiation(&awaiterType, 1), FALSE); + instArg->constLookup.addr = pContext; + pInliningContext = pContext; + } + } + + *contextHandle = MAKE_METHODCONTEXT(pInliningContext); + + EE_TO_JIT_TRANSITION(); + + return CORINFO_METHOD_HANDLE(pMD); +} + // Compute the runtime lookup for the instantiation argument for an // AsyncHelpers.TransparentAwait call, to be used for wrapping a return value // from pCallerMD for its runtime async version. @@ -10509,6 +10563,65 @@ void CEEInfo::ComputeRuntimeLookupForAwaitCall(MethodDesc* pCallerMD, MethodDesc FinishComputeRuntimeLookup(sigBuilder, pCallerMD, lookup); } +void CEEInfo::ComputeRuntimeLookupForAwaitAwaiterInContinuationCall( + MethodDesc* pCallerMD, + MethodDesc* pTypicalAwaitMD, + CORINFO_SIG_INFO* callSig, + CORINFO_LOOKUP* lookup) +{ + lookup->lookupKind.needsRuntimeLookup = true; + + CORINFO_RUNTIME_LOOKUP* rlookup = &lookup->runtimeLookup; + rlookup->signature = NULL; + rlookup->indirectFirstOffset = 0; + rlookup->indirectSecondOffset = 0; + rlookup->sizeOffset = CORINFO_NO_SIZE_CHECK; + rlookup->indirections = CORINFO_USEHELPER; + + if (pCallerMD->RequiresInstMethodDescArg()) + { + lookup->lookupKind.runtimeLookupKind = CORINFO_LOOKUP_METHODPARAM; + rlookup->helper = CORINFO_HELP_RUNTIMEHANDLE_METHOD; + } + else if (pCallerMD->RequiresInstMethodTableArg()) + { + lookup->lookupKind.runtimeLookupKind = CORINFO_LOOKUP_CLASSPARAM; + rlookup->helper = CORINFO_HELP_RUNTIMEHANDLE_CLASS; + } + else + { + lookup->lookupKind.runtimeLookupKind = CORINFO_LOOKUP_THISOBJ; + rlookup->helper = CORINFO_HELP_RUNTIMEHANDLE_CLASS; + } + + SigBuilder sigBuilder; + sigBuilder.AppendData(MethodDescSlot); + if (lookup->lookupKind.runtimeLookupKind != CORINFO_LOOKUP_METHODPARAM) + { + sigBuilder.AppendData(pCallerMD->GetMethodTable()->GetNumDicts() - 1); + } + + sigBuilder.AppendElementType(ELEMENT_TYPE_INTERNAL); + sigBuilder.AppendPointer(pTypicalAwaitMD->GetMethodTable()); + sigBuilder.AppendData(ENCODE_METHOD_SIG_MethodInstantiation | ENCODE_METHOD_SIG_InstantiatingStub); + sigBuilder.AppendElementType(ELEMENT_TYPE_INTERNAL); + sigBuilder.AppendPointer(pTypicalAwaitMD->GetMethodTable()); + sigBuilder.AppendData(RidFromToken(pTypicalAwaitMD->GetMemberDef())); + sigBuilder.AppendData(1); + + _ASSERTE(callSig->pSig != NULL); + SigPointer instantiation(callSig->pSig, callSig->cbSig); + BYTE callingConvention; + IfFailThrow(instantiation.GetByte(&callingConvention)); + _ASSERTE(callingConvention == IMAGE_CEE_CS_CALLCONV_GENERICINST); + uint32_t numArgs; + IfFailThrow(instantiation.GetData(&numArgs)); + _ASSERTE(numArgs == 1); + instantiation.ConvertToInternalExactlyOne(GetModule(callSig->scope), NULL, &sigBuilder); + + FinishComputeRuntimeLookup(sigBuilder, pCallerMD, lookup); +} + static MethodTable* getContinuationType( size_t dataSize, bool* objRefs, diff --git a/src/coreclr/vm/jitinterface.h b/src/coreclr/vm/jitinterface.h index 62280728dc1cba..5072ab4c854889 100644 --- a/src/coreclr/vm/jitinterface.h +++ b/src/coreclr/vm/jitinterface.h @@ -413,6 +413,11 @@ class CEEInfo : public ICorJitInfo CORINFO_LOOKUP* pResultLookup); void ComputeRuntimeLookupForAwaitCall(MethodDesc* pCallerMD, MethodDesc* pTypicalAwaitMD, CORINFO_LOOKUP* lookup); + void ComputeRuntimeLookupForAwaitAwaiterInContinuationCall( + MethodDesc* pCallerMD, + MethodDesc* pTypicalAwaitMD, + CORINFO_SIG_INFO* callSig, + CORINFO_LOOKUP* lookup); #if defined(FEATURE_GDBJIT) CalledMethod * GetCalledMethods() { return m_pCalledMethods; } diff --git a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.cs b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.cs index 7e11855f9b82d0..9e638aa9204141 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.cs @@ -36,6 +36,29 @@ public static unsafe void AwaitAwaiter(TAwaiter awaiter) where TAwaite AsyncSuspend(sentinelContinuation); } + [BypassReadyToRun] + [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)] + [StackTraceHidden] + private static unsafe void AwaitAwaiterInContinuation(int offset) + where TAwaiter : INotifyCompletion + { + ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState; + Continuation? sentinelContinuation = state.SentinelContinuation ??= new Continuation(); + state.StackState->AwaiterContinuation = &AwaiterOnCompletedFromContinuation; + state.StackState->AwaiterOffset = offset; + state.CaptureContexts(); + AsyncSuspend(sentinelContinuation); + } + + private static void AwaiterOnCompletedFromContinuation( + Continuation headContinuation, int offset, Action continuation) + where TAwaiter : INotifyCompletion + { + ref byte data = ref RuntimeHelpers.GetRawData(headContinuation); + TAwaiter awaiter = Unsafe.As(ref Unsafe.Add(ref data, offset)); + awaiter.OnCompleted(continuation); + } + // Must be NoInlining because we use AsyncSuspend to manufacture an explicit suspension point. // It will not capture/restore any local state that is live across it. @@ -70,6 +93,35 @@ public static unsafe void UnsafeAwaitAwaiter(TAwaiter awaiter) where T AsyncSuspend(sentinelContinuation); } + [BypassReadyToRun] + [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)] + [StackTraceHidden] + private static unsafe void UnsafeAwaitAwaiterInContinuation(int offset) + where TAwaiter : ICriticalNotifyCompletion + { + ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState; + Continuation? sentinelContinuation = state.SentinelContinuation ??= new Continuation(); + state.StackState->AwaiterContinuation = &UnsafeAwaiterOnCompletedFromContinuation; + state.StackState->AwaiterOffset = offset; + state.CaptureContexts(); + AsyncSuspend(sentinelContinuation); + } + + private static void UnsafeAwaiterOnCompletedFromContinuation( + Continuation headContinuation, int offset, Action continuation) + where TAwaiter : ICriticalNotifyCompletion + { + if (typeof(TAwaiter) == typeof(YieldAwaitable.YieldAwaiter)) + { + RuntimeAsyncYielder.Instance.UnsafeOnCompleted(continuation); + return; + } + + ref byte data = ref RuntimeHelpers.GetRawData(headContinuation); + TAwaiter awaiter = Unsafe.As(ref Unsafe.Add(ref data, offset)); + awaiter.UnsafeOnCompleted(continuation); + } + /// /// Awaits the specified and returns its result, throwing any exception produced by the task. /// diff --git a/src/tests/JIT/Regression/JitBlue/GitHub_130437/GitHub_130437.csproj b/src/tests/JIT/Regression/JitBlue/GitHub_130437/GitHub_130437.csproj index 8f66634434cd1a..f07af3a622402e 100644 --- a/src/tests/JIT/Regression/JitBlue/GitHub_130437/GitHub_130437.csproj +++ b/src/tests/JIT/Regression/JitBlue/GitHub_130437/GitHub_130437.csproj @@ -1,4 +1,4 @@ - + $(RepoRoot)src/coreclr/.nuget/Microsoft.NET.Sdk.IL/Microsoft.NET.Sdk.IL.targets.template True diff --git a/src/tests/async/awaitingnoasyncgeneric/awaitingnoasyncgeneric.csproj b/src/tests/async/awaitingnoasyncgeneric/awaitingnoasyncgeneric.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/awaitingnoasyncgeneric/awaitingnoasyncgeneric.csproj +++ b/src/tests/async/awaitingnoasyncgeneric/awaitingnoasyncgeneric.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/awaitingnotasync/awaitingnotasync.csproj b/src/tests/async/awaitingnotasync/awaitingnotasync.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/awaitingnotasync/awaitingnotasync.csproj +++ b/src/tests/async/awaitingnotasync/awaitingnotasync.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/byref-param/byref-param.csproj b/src/tests/async/byref-param/byref-param.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/byref-param/byref-param.csproj +++ b/src/tests/async/byref-param/byref-param.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/cancellation/cancellation.csproj b/src/tests/async/cancellation/cancellation.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/cancellation/cancellation.csproj +++ b/src/tests/async/cancellation/cancellation.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/collectible-alc/collectible-alc.csproj b/src/tests/async/collectible-alc/collectible-alc.csproj index 5bc9c9b10779eb..e65e440150c631 100644 --- a/src/tests/async/collectible-alc/collectible-alc.csproj +++ b/src/tests/async/collectible-alc/collectible-alc.csproj @@ -1,4 +1,4 @@ - + True diff --git a/src/tests/async/custom-struct-awaiters/custom-struct-awaiters.cs b/src/tests/async/custom-struct-awaiters/custom-struct-awaiters.cs new file mode 100644 index 00000000000000..6946cb6975745b --- /dev/null +++ b/src/tests/async/custom-struct-awaiters/custom-struct-awaiters.cs @@ -0,0 +1,85 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +public class CustomStructAwaiters +{ + private static int s_safeAwaiterValue; + private static int s_unsafeAwaiterValue; + + [Fact] + public static void TestEntryPoint() + { + Run().Wait(); + } + + private static async Task Run() + { + Volatile.Write(ref s_safeAwaiterValue, 0); + await new SafeAwaitable(42); + Assert.Equal(42, Volatile.Read(ref s_safeAwaiterValue)); + + Volatile.Write(ref s_unsafeAwaiterValue, 0); + await new UnsafeAwaitable(43); + Assert.Equal(43, Volatile.Read(ref s_unsafeAwaiterValue)); + } + + private readonly struct SafeAwaitable + { + private readonly int _value; + + public SafeAwaitable(int value) => _value = value; + + public SafeAwaiter GetAwaiter() => new SafeAwaiter(_value); + + public readonly struct SafeAwaiter : INotifyCompletion + { + private readonly int _value; + + public SafeAwaiter(int value) => _value = value; + + public bool IsCompleted => false; + + public void OnCompleted(Action continuation) + { + Volatile.Write(ref s_safeAwaiterValue, _value); + ThreadPool.QueueUserWorkItem(static state => ((Action)state!).Invoke(), continuation); + } + + public void GetResult() => Assert.Equal(42, _value); + } + } + + private readonly struct UnsafeAwaitable + { + private readonly int _value; + + public UnsafeAwaitable(int value) => _value = value; + + public UnsafeAwaiter GetAwaiter() => new UnsafeAwaiter(_value); + + public readonly struct UnsafeAwaiter : ICriticalNotifyCompletion + { + private readonly int _value; + + public UnsafeAwaiter(int value) => _value = value; + + public bool IsCompleted => false; + + public void OnCompleted(Action continuation) => throw new InvalidOperationException(); + + public void UnsafeOnCompleted(Action continuation) + { + Volatile.Write(ref s_unsafeAwaiterValue, _value); + ThreadPool.UnsafeQueueUserWorkItem(static state => ((Action)state!).Invoke(), continuation); + } + + public void GetResult() => Assert.Equal(43, _value); + } + } +} diff --git a/src/tests/async/custom-struct-awaiters/custom-struct-awaiters.csproj b/src/tests/async/custom-struct-awaiters/custom-struct-awaiters.csproj new file mode 100644 index 00000000000000..197767e2c4e249 --- /dev/null +++ b/src/tests/async/custom-struct-awaiters/custom-struct-awaiters.csproj @@ -0,0 +1,5 @@ + + + + + diff --git a/src/tests/async/eh-microbench/eh-microbench.csproj b/src/tests/async/eh-microbench/eh-microbench.csproj index e0c0b671ef1829..6650af6b67607c 100644 --- a/src/tests/async/eh-microbench/eh-microbench.csproj +++ b/src/tests/async/eh-microbench/eh-microbench.csproj @@ -1,4 +1,4 @@ - + BuildOnly diff --git a/src/tests/async/execution-context/execution-context.csproj b/src/tests/async/execution-context/execution-context.csproj index 644025b102673e..3a8b32c0ae764d 100644 --- a/src/tests/async/execution-context/execution-context.csproj +++ b/src/tests/async/execution-context/execution-context.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/gc-roots-scan/gc-roots-scan.csproj b/src/tests/async/gc-roots-scan/gc-roots-scan.csproj index 638c237d91983d..21cc479969f131 100644 --- a/src/tests/async/gc-roots-scan/gc-roots-scan.csproj +++ b/src/tests/async/gc-roots-scan/gc-roots-scan.csproj @@ -1,4 +1,4 @@ - + BuildOnly diff --git a/src/tests/async/implement/implement.csproj b/src/tests/async/implement/implement.csproj index 644025b102673e..3a8b32c0ae764d 100644 --- a/src/tests/async/implement/implement.csproj +++ b/src/tests/async/implement/implement.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/inst-unbox-thunks/inst-unbox-thunks.csproj b/src/tests/async/inst-unbox-thunks/inst-unbox-thunks.csproj index 8e9fed2b02ad47..d06e5df35dd4a2 100644 --- a/src/tests/async/inst-unbox-thunks/inst-unbox-thunks.csproj +++ b/src/tests/async/inst-unbox-thunks/inst-unbox-thunks.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/mincallcost-microbench/mincallcost-microbench.csproj b/src/tests/async/mincallcost-microbench/mincallcost-microbench.csproj index e0c0b671ef1829..6650af6b67607c 100644 --- a/src/tests/async/mincallcost-microbench/mincallcost-microbench.csproj +++ b/src/tests/async/mincallcost-microbench/mincallcost-microbench.csproj @@ -1,4 +1,4 @@ - + BuildOnly diff --git a/src/tests/async/object/object.csproj b/src/tests/async/object/object.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/object/object.csproj +++ b/src/tests/async/object/object.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/objects-captured/objects-captured.csproj b/src/tests/async/objects-captured/objects-captured.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/objects-captured/objects-captured.csproj +++ b/src/tests/async/objects-captured/objects-captured.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/override/override.csproj b/src/tests/async/override/override.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/override/override.csproj +++ b/src/tests/async/override/override.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/pgo/pgo.csproj b/src/tests/async/pgo/pgo.csproj index 8d460faebada27..1c5f5d7e02eb7d 100644 --- a/src/tests/async/pgo/pgo.csproj +++ b/src/tests/async/pgo/pgo.csproj @@ -1,4 +1,4 @@ - + true diff --git a/src/tests/async/pinvoke/pinvoke.csproj b/src/tests/async/pinvoke/pinvoke.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/pinvoke/pinvoke.csproj +++ b/src/tests/async/pinvoke/pinvoke.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/regression/125042.csproj b/src/tests/async/regression/125042.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/regression/125042.csproj +++ b/src/tests/async/regression/125042.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/regression/125615.csproj b/src/tests/async/regression/125615.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/regression/125615.csproj +++ b/src/tests/async/regression/125615.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/regression/125805.csproj b/src/tests/async/regression/125805.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/regression/125805.csproj +++ b/src/tests/async/regression/125805.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/regression/128566.csproj b/src/tests/async/regression/128566.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/regression/128566.csproj +++ b/src/tests/async/regression/128566.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/regression/managed-thread-id.csproj b/src/tests/async/regression/managed-thread-id.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/regression/managed-thread-id.csproj +++ b/src/tests/async/regression/managed-thread-id.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/returns/returns.csproj b/src/tests/async/returns/returns.csproj index ff94a80928aef7..7548a677889c52 100644 --- a/src/tests/async/returns/returns.csproj +++ b/src/tests/async/returns/returns.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/shared-generic/shared-generic.csproj b/src/tests/async/shared-generic/shared-generic.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/shared-generic/shared-generic.csproj +++ b/src/tests/async/shared-generic/shared-generic.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/staticvirtual/staticvirtual.csproj b/src/tests/async/staticvirtual/staticvirtual.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/staticvirtual/staticvirtual.csproj +++ b/src/tests/async/staticvirtual/staticvirtual.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/struct/struct.cs b/src/tests/async/struct/struct.cs index b8d8ab53b14077..df393687e5216e 100644 --- a/src/tests/async/struct/struct.cs +++ b/src/tests/async/struct/struct.cs @@ -32,7 +32,6 @@ private static async Task Async2() AssertEqual(100, s.Value); } - [MethodImpl(MethodImplOptions.NoInlining)] private static void AssertEqual(int expected, int val) { diff --git a/src/tests/async/struct/struct.csproj b/src/tests/async/struct/struct.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/struct/struct.csproj +++ b/src/tests/async/struct/struct.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/synchronization-context/synchronization-context.csproj b/src/tests/async/synchronization-context/synchronization-context.csproj index 644025b102673e..3a8b32c0ae764d 100644 --- a/src/tests/async/synchronization-context/synchronization-context.csproj +++ b/src/tests/async/synchronization-context/synchronization-context.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/valuetask-source/valuetask-source-stub.csproj b/src/tests/async/valuetask-source/valuetask-source-stub.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/valuetask-source/valuetask-source-stub.csproj +++ b/src/tests/async/valuetask-source/valuetask-source-stub.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/valuetask-source/valuetask-source.csproj b/src/tests/async/valuetask-source/valuetask-source.csproj index 7e637235662d05..c2c78be9d50708 100644 --- a/src/tests/async/valuetask-source/valuetask-source.csproj +++ b/src/tests/async/valuetask-source/valuetask-source.csproj @@ -1,4 +1,4 @@ - + true diff --git a/src/tests/async/valuetask/valuetask.csproj b/src/tests/async/valuetask/valuetask.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/valuetask/valuetask.csproj +++ b/src/tests/async/valuetask/valuetask.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/varying-yields/varying-yields-task.csproj b/src/tests/async/varying-yields/varying-yields-task.csproj index a063851b97b392..b80a06405ee432 100644 --- a/src/tests/async/varying-yields/varying-yields-task.csproj +++ b/src/tests/async/varying-yields/varying-yields-task.csproj @@ -1,4 +1,4 @@ - + ASYNC1_TASK;$(DefineConstants) diff --git a/src/tests/async/varying-yields/varying-yields-valuetask.csproj b/src/tests/async/varying-yields/varying-yields-valuetask.csproj index 7679b339a8019b..9e82183ae9f70b 100644 --- a/src/tests/async/varying-yields/varying-yields-valuetask.csproj +++ b/src/tests/async/varying-yields/varying-yields-valuetask.csproj @@ -1,4 +1,4 @@ - + ASYNC1_VALUETASK;$(DefineConstants) diff --git a/src/tests/async/varying-yields/varying-yields.csproj b/src/tests/async/varying-yields/varying-yields.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/varying-yields/varying-yields.csproj +++ b/src/tests/async/varying-yields/varying-yields.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/void/void.csproj b/src/tests/async/void/void.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/void/void.csproj +++ b/src/tests/async/void/void.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/with-yields/with-yields.csproj b/src/tests/async/with-yields/with-yields.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/with-yields/with-yields.csproj +++ b/src/tests/async/with-yields/with-yields.csproj @@ -1,4 +1,4 @@ - + diff --git a/src/tests/async/without-yields/without-yields.csproj b/src/tests/async/without-yields/without-yields.csproj index 3fc50cde4b3443..197767e2c4e249 100644 --- a/src/tests/async/without-yields/without-yields.csproj +++ b/src/tests/async/without-yields/without-yields.csproj @@ -1,4 +1,4 @@ - + From 0330364ea53ed5045ff400df692f8f0229b30813 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Mon, 27 Jul 2026 19:45:09 +0200 Subject: [PATCH 11/41] Run jit-format --- src/coreclr/jit/async.cpp | 12 ++++++------ src/coreclr/jit/compiler.cpp | 4 ++-- src/coreclr/jit/importercalls.cpp | 2 +- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 300fbefce973f4..4a6994e2e34b8f 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -1946,9 +1946,9 @@ bool AsyncTransformation::IsReusableSuspension(const AsyncState* state, } } - static const WellKnownArg validateArgs[] = { - WellKnownArg::AsyncResumedUse, WellKnownArg::AsyncResumedDef, WellKnownArg::AsyncExecutionContext, - WellKnownArg::AsyncSynchronizationContext}; + static const WellKnownArg validateArgs[] = {WellKnownArg::AsyncResumedUse, WellKnownArg::AsyncResumedDef, + WellKnownArg::AsyncExecutionContext, + WellKnownArg::AsyncSynchronizationContext}; for (WellKnownArg arg : validateArgs) { CallArg* thisArg = call->gtArgs.FindWellKnownArg(arg); @@ -1999,9 +1999,9 @@ bool AsyncTransformation::IsReusableSuspension(const AsyncState* state, // void AsyncTransformation::HandleReusedSuspension(BasicBlock* callBlock, GenTreeCall* call) { - static const WellKnownArg argsToRemove[] = { - WellKnownArg::AsyncResumedUse, WellKnownArg::AsyncResumedDef, WellKnownArg::AsyncExecutionContext, - WellKnownArg::AsyncSynchronizationContext}; + static const WellKnownArg argsToRemove[] = {WellKnownArg::AsyncResumedUse, WellKnownArg::AsyncResumedDef, + WellKnownArg::AsyncExecutionContext, + WellKnownArg::AsyncSynchronizationContext}; for (WellKnownArg wka : argsToRemove) { CallArg* arg = call->gtArgs.FindWellKnownArg(wka); diff --git a/src/coreclr/jit/compiler.cpp b/src/coreclr/jit/compiler.cpp index 5f880fd6806071..0e5dae4e7ac605 100644 --- a/src/coreclr/jit/compiler.cpp +++ b/src/coreclr/jit/compiler.cpp @@ -10383,8 +10383,8 @@ bool Compiler::lvaIsOSRLocal(unsigned varNum) // Sanity check for promoted fields of OSR locals. // if ((varNum >= info.compLocalsCount) && (varNum != lvaMonAcquired) && (varNum != lvaAsyncThreadObjectVar) && - (varNum != lvaResumedIndicator) && - (varNum != lvaAsyncExecutionContextVar) && (varNum != lvaAsyncSynchronizationContextVar)) + (varNum != lvaResumedIndicator) && (varNum != lvaAsyncExecutionContextVar) && + (varNum != lvaAsyncSynchronizationContextVar)) { assert(varDsc->lvIsStructField); assert(varDsc->lvParentLcl < info.compLocalsCount); diff --git a/src/coreclr/jit/importercalls.cpp b/src/coreclr/jit/importercalls.cpp index 206ccbf9fe7d12..775339048a6af1 100644 --- a/src/coreclr/jit/importercalls.cpp +++ b/src/coreclr/jit/importercalls.cpp @@ -7429,7 +7429,7 @@ void Compiler::impInheritAsyncContextsFromInliner(GenTreeCall* call) return; } - GenTreeCall* inlCall = impInlineInfo->iciCall; + GenTreeCall* inlCall = impInlineInfo->iciCall; CallArg* resumedUseArg = inlCall->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse); CallArg* resumedDefArg = inlCall->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedDef); CallArg* execArg = inlCall->gtArgs.FindWellKnownArg(WellKnownArg::AsyncExecutionContext); From dc1941d9c92263b831f9c1f64f76c55dfb7b9bdc Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Tue, 28 Jul 2026 16:57:21 +0200 Subject: [PATCH 12/41] JIT: Address async transformation review feedback Consolidate the duplicated non-null assert in FinishContextHandlingAndSuspensionWithHelper into a single assert covering all three well-known args. Update the stale comment in InsertFinishContextHandlingCall that still described the resumed placeholder as being replaced by a 'continuation != null' computation; it is now replaced by the already-computed resumed value. No functional change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d86cad55-abc0-42d6-8369-bc6bc619326f --- src/coreclr/jit/async.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 4a6994e2e34b8f..810ca2d37a75cb 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -2594,8 +2594,7 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc CallArg* resumedArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse); CallArg* execContextArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncExecutionContext); CallArg* syncContextArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncSynchronizationContext); - assert((resumedArg != nullptr) && (execContextArg != nullptr)); - assert((execContextArg != nullptr) && (syncContextArg != nullptr)); + assert((resumedArg != nullptr) && (execContextArg != nullptr) && (syncContextArg != nullptr)); // Get the contexts from the call node: // 1. For shared finish, store it directly to the shared locals in the same block @@ -3769,7 +3768,7 @@ void AsyncTransformation::InsertFinishContextHandlingCall(BasicBlock* use.ReplaceWith(execContextAddrOffset); LIR::AsRange(block).Remove(execContextAddrPlaceholder); - // Replace resumedPlaceholder with actual "continuationParameter != null" arg + // Replace resumedPlaceholder with the resumed value gotUse = LIR::AsRange(block).TryGetUse(resumedPlaceholder, &use); assert(gotUse); From 3a37708bef59c34b0617553e89e5bcab246447db Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Tue, 28 Jul 2026 17:21:24 +0200 Subject: [PATCH 13/41] JIT: Allow invariant async pseudo-args when reusing suspensions Value numbering can fold the AsyncResumedUse argument to a constant when the resumed indicator is provably zero at the call, for example when two awaits sit on mutually exclusive paths. IsReusableSuspension only accepted locals, so it bailed out with 'argument is too complex' and gave up on merging the suspension points. That produced a duplicated CORINFO_HELP_ALLOC_CONTINUATION sequence, an extra resumption block and an extra readonly data slot. In smoke_tests.nativeaot this showed up as +59 bytes (17.15%) on Generics+TestAsyncGVMScenarios:AsyncGvm1[System.__Canon] and +57 bytes (16.52%) on RunAsync. Accept invariant nodes as well, matching the predicate already used for invariantResumed in CreateResumptionsAndSuspensions. Invariant nodes are constants, LCL_ADDR or FTN_ADDR, so comparing them and removing them from the call is safe. Both methods now improve by 15 bytes relative to the base instead of regressing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d86cad55-abc0-42d6-8369-bc6bc619326f --- src/coreclr/jit/async.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 810ca2d37a75cb..d8e501d4b10514 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -1962,7 +1962,10 @@ bool AsyncTransformation::IsReusableSuspension(const AsyncState* state, if (thisArg != nullptr) { - if (!thisArg->GetNode()->OperIsAnyLocal()) + // The value may have been folded to a constant (e.g. when the + // indicator is provably zero at this point), which is still fine to + // compare and to remove from the call. + if (!thisArg->GetNode()->OperIsAnyLocal() && !thisArg->GetNode()->IsInvariant()) { JITDUMP(" No; %s argument is too complex\n", getWellKnownArgName(arg)); return false; @@ -2007,7 +2010,7 @@ void AsyncTransformation::HandleReusedSuspension(BasicBlock* callBlock, GenTreeC CallArg* arg = call->gtArgs.FindWellKnownArg(wka); if (arg != nullptr) { - assert(arg->GetNode()->OperIsAnyLocal()); + assert(arg->GetNode()->OperIsAnyLocal() || arg->GetNode()->IsInvariant()); LIR::AsRange(callBlock).Remove(arg->GetNode()); call->gtArgs.RemoveUnsafe(arg); } From 545d8bd35fca5400eed38b7fbd76372e04193130 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 12:51:05 +0200 Subject: [PATCH 14/41] JIT: Fix LIR corruption when storing async awaiter into continuation StoreAsyncAwaiter removed the awaiter node from the call block after it had already been spliced into the suspension block's LIR range, corrupting the range and tripping 'prev == m_lastNode'. Remove the node from the call block before building the store instead, and remove the FIELD_LIST node itself in the multi-field case. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index a406afc20ab806..e6c576f2fc10eb 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -2421,6 +2421,8 @@ void AsyncTransformation::StoreAsyncAwaiter(BasicBlock* callBlock, GenTree* store = StoreAtOffset(continuation, offset, field, use.GetType()); LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_compiler, store)); } + + LIR::AsRange(callBlock).Remove(fieldList); } else { @@ -2431,6 +2433,8 @@ void AsyncTransformation::StoreAsyncAwaiter(BasicBlock* callBlock, awaiter = use.Def(); } + LIR::AsRange(callBlock).Remove(awaiter); + 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); @@ -2439,7 +2443,6 @@ void AsyncTransformation::StoreAsyncAwaiter(BasicBlock* callBlock, LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_compiler, store)); } - LIR::AsRange(callBlock).Remove(awaiter); call->gtArgs.RemoveUnsafe(awaiterArg); } From 673390a0158fa1e5e2c09a24537c42eb20e8aedd Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 13:21:57 +0200 Subject: [PATCH 15/41] JIT: Add design document for general runtime async inlining Folds in dotnet/runtime#127800. Renames GT_CONTINUATION_FIELD_OFFSET to GT_CONTINUATION_MEMBER_OFFSET to match the node actually implemented. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../coreclr/jit/runtime-async-inlining.md | 342 ++++++++++++++++++ 1 file changed, 342 insertions(+) create mode 100644 docs/design/coreclr/jit/runtime-async-inlining.md diff --git a/docs/design/coreclr/jit/runtime-async-inlining.md b/docs/design/coreclr/jit/runtime-async-inlining.md new file mode 100644 index 00000000000000..dcc1293d2ec6fb --- /dev/null +++ b/docs/design/coreclr/jit/runtime-async-inlining.md @@ -0,0 +1,342 @@ +# General runtime async inlining + +This document outlines some challenges and thoughts behind general inlining of runtime async calls. + +## Background + +Runtime async comes with several user-visible behaviors that must be preserved when inlining. +These can be categorized in behaviors attached to the call site and behaviors attached to the callee. + +### Await/call site behaviors + +Custom awaiters and synchronous Task/ValueTask-returning functions can be inlined today, so below we consider only runtime async backed implementations. + +- Every suspending await saves/restores the `ExecutionContext`. +Concretely this is what saves and restores the values of `AsyncLocal` if suspension/resumption happens around an await. +The behavior of this cannot be configured per-await, but it can be configured globally via `ExecutionContext.SuppressFlow()`. +For our purposes we can consider this to always require saving/restoration. +- For suspending `Task` and `ValueTask` awaits the resumption behavior is configurable: + + By default, the current `SynchronizationContext` or `TaskScheduler` is captured and the continuation is posted to it. + + If `ConfigureAwait(false)` is applied the continuation is run on the thread pool. + + The "location to continue", either a `SynchronizationContext`, `TaskScheduler` or the thread pool, is collectively called the _continuation context_ in this document. + +For awaits that finish synchronously (i.e. do not suspend) no additional work happens at the call site. + +To accomplish the above the JIT works together with async infrastructure: +- Suspension generates code to obtain and save the `ExecutionContext` into the `Continuation` object. +On resumption, async infrastructure restores the `ExecutionContext` before calling back into JIT'ed code. +- Suspension generates code to capture continuation context information into the `Continuation` object. +On resumption, async infrastructure ensures that the continuation runs in the captured continuation context. + +The above are accomplished with several helpers returned via the `getAsyncInfo` JIT-EE helper: +- `AsyncHelpers.CaptureExecutionContext` +- `AsyncHelpers.CaptureContinuationContext` +- `AsyncHelpers.FinishSuspensionNoContinuationContext` (optimized case) +- `AsyncHelpers.FinishSuspensionWithContinuationContext` (optimized case) + +### Function/callee behaviors +Every runtime async function body is wrapped with save and restore of the async contexts on the current Thread object. +This restore happens only when the function finishes synchronously. +Concretely a runtime async function ends up looking like: +```csharp +ExecutionContext execContext; +SynchronizationContext syncContext; +AsyncHelpers.CaptureContexts(out execContext, out syncContext); +try +{ + // actual user code +} +finally +{ + AsyncHelpers.RestoreContexts(resumed, execContext, syncContext); +} +``` + +Note the `resumed` argument, used to check if the runtime async function finished synchronously or not. +If the runtime async function finished after resuming then contexts are not restored. + +Additionally, every suspension _also_ restores the contexts. +This is done by inserting a similar call +```csharp +AsyncHelpers.RestoreContextsOnSuspension(resumed, execContext, syncContext) +``` +in suspensions. +This is folded into the `AsyncHelpers.FinishSuspensionNoContinuationContext` and `AsyncHelpers.FinishSuspensionWithContinuationContext` helpers. +The JIT uses those helpers whenever possible. + +## Inlining runtime async calls +When inlining a runtime async call we remove both a call site and the callee. +This removes the user-visible behaviors discussed above. +To inline correctly we need to ensure the JIT preserves these behaviors. + +### Suspending inside an inlinee + +Consider a scenario like: + +```csharp +async Task A() +{ + ... + await B(); + ... +} + +async Task B() +{ + ... + await C(); + ... +} + +async Task C() +{ + ... + await Task.Yield(); + ... +} +``` +Now assume that we have available bools `resumed_A`, `resumed_B`, `resumed_C` that represent whether `A`, `B` and `C` are currently running after being resumed (true), or whether they are running because they were started and have never suspended before (false). +Note that `resumed_C` and `resumed_B` can switch between true/false multiple times during the execution if `B` and `C` are in loops, while `resumed_A` will switch to `true` and then stay `true`. + +Assume we inlined `B()` and `C()` into `A`, and that we are logically suspending on `Task.Yield()` in `C`. +What does this suspension capture? That depends on `resumed_B` and `resumed_C`. +- Let's first consider `!resumed_C`. +This means that we just started running `C`, and on suspension we would restore `ExecutionContext` and `SynchronizationContext` from the beginning of `C` into the `Thread` object. +If we hadn't inlined then `B`'s frame would be physically present on the stack due to `!resumed_C`, so after returning it would create and link its own continuation. +This continuation would capture the `ExecutionContext` restored by `C` and also capture continuation context information based on the `SynchronizationContext` restored by `C`. +- If `resumed_C` then `B`'s frame would not be physically present on the stack. +Rather, the existing `B` continuation would be linked to the newly created `C` continuation by async infrastructure. +It would maintain its values from the time we suspended in `C` with `!resumed_C`. + +The take away is that the suspension for the inlined `Task.Yield()` needs to roughly accomplish the following, in addition to storing its normal state: +```csharp +continuation = AllocOrReuseContinuation(); +CaptureContinuationContext(continuation, ref continuation.ContinuationContext, ref continuation.Flags); // In the general case, but not for Task.Yield() +CaptureExecutionContext(continuation, ref continuation.ExecutionContext); +if (!resumed_C) +{ + RestoreContextsOnSuspension(false, execContext_C, syncContext_C); + // Logical return to B and linking in B's continuation + CaptureContinuationContext(continuation, ref continuation.ContinuationContextForB, ref continuation.FlagsForB) + CaptureExecutionContext(continuation, ref continuation.ExecutionContextForB); + if (!resumed_B) + { + RestoreContextsOnSuspension(false, execContext_B, syncContext_B); + // Logical return to A and linking in A's continuation + CaptureContinuationContext(continuation, ref continuation.ContinuationContextForA, ref continuation.FlagsForA) + CaptureExecutionContext(continuation, ref continuation.ExecutionContextForA); + + RestoreContextsOnSuspension(resumed_A, execContext_A, syncContext_A); + } + else + { + // keep continuation.ContinuationContextForA, FlagsForA, ExecutionContextForA + } +} +else +{ + // keep continuation.ContinuationContextForB, FlagsForB, ExecutionContextForB +} +``` +Note also that we have the implications `resumed_C` implies `resumed_B` implies `resumed_A` simplifying the `else` cases here. +The else cases also rely on `continuation.ContinuationContextForB/A` staying assigned in the continuations, which is the case only because the JIT allocates a single continuation and reuses it throughout the execution of `A`. + +This may look like a large amount of code for every suspension inside `C`, but sharing of tails is possible since all the inlined suspensions end with the same code that does not change per suspension site. + +### Resuming inside an inlinee + +Let's imagine we resume at the inlined `Task.Yield()` in `C`. +Initially such a resumption progresses as normal: async infrastructure restores the `ExecutionContext` that was saved. +If this was a task await (which it isn't, but often will be), the async infrastructure ensures it is running in the right continuation context. +Then it resumes in the code. + +When `C` logically returns to `B` we need to handle what the async infrastructure would have handled for us. +Call this the _post-inline IR_. +- We need to restore the `ExecutionContext` that we captured for the suspension at `C()`. +Logically, we need to do something like +```csharp +if (resumed_C) + AsyncHelpers.RestoreExecutionContext(Thread.CurrentThreadAssumedInitialized, continuation.ExecutionContextForB); +``` +- We need to ensure we're running on the proper continuation context. +The JIT cannot accomplish this on its own; there is no way to directly set up the environment without posting a callback. +Hence, in the case where we are not running in the right context we need help from async infrastructure to proceed with a suspension+switch+resumption: +```csharp +if (resumed_C && !AsyncHelpers.IsOnRightContext(continuation.ContinuationContextForB, continuation.FlagsForB)) + await AsyncHelpers.SwitchContext(continuation.ContinuationContextForB, continuation.FlagsForB); +``` + +We expect `IsOnRightContext` to be true almost always, since it is very rare that the current synchronization context is modified, especially during the execution of async methods. +If this was a common case then it is unlikely that inlining would ever be profitable. + +### Handling synchronous saves and restores of contexts + +Recall that every runtime async function's body was wrapped with a save/restore of the contexts. +Given `resumed_A`, `resumed_B` and `resumed_C` it is not hard to insert these for every inlinee. +For example, when inlining `C` we insert code similar to: + +```csharp +ExecutionContext execContext; +SynchronizationContext syncContext; +AsyncHelpers.CaptureContexts(out execContext, out syncContext); +try +{ + // actual user code +} +finally +{ + AsyncHelpers.RestoreContexts(resumed_C, execContext, syncContext); +} +``` + +### Special cases + +The only case supported today is when we know the inlinee never suspends, by virtue of having no awaits at all. +This simplifies the cases above since `resumed_B` and `resumed_C` are always false, and all suspension points generated belong to `A`. + +### Computing `resumed_A`, `resumed_B`, `resumed_C` ... + +Today we compute the `resumed` value only for the top-level async function, since it is computed based on the async continuation parameter. +With inlinees this strategy no longer works. +Instead insert IR to compute these values at several points: +1. When starting a function `F` (whether it be the top function or start of an inlined function), assign `resumed_F = false` +2. When resuming `F` at one of its resumption points assign `resumed_F = true` +3. When logically returning from `F` to its caller, inherit the resumption status: `resumed_caller |= resumed_F`. + +Note that these are introduced at different points. (1) and (3) will likely be inserted during `PHASE_SAVE_ASYNC_CONTEXTS` while (2) will be inserted in `PHASE_ASYNC`. + +(2) is a store to `resumed_F` inserted very late. +Until the async transformation we will model this as a `LCL_ADDR` well-known argument added to all async calls. + +It is important that the JIT can reason about `resumed` booleans to be able to eliminate the cruft for inlinees that are proven to never suspend. +We would like constant propagation to realize `resumed_C` is always false if `C` has no suspension point, for example. + +### Representation of post-inline IR + +The post-inline IR we need to insert when `C` returns to `B` looks like: + +```csharp +if (resumed_C) +{ + AsyncHelpers.RestoreExecutionContext(Thread.CurrentThreadAssumedInitialized, continuation.ExecutionContextForB); + + if (!AsyncHelpers.IsOnRightContext(continuation.ContinuationContextForB, continuation.FlagsForB)) + { + await AsyncHelpers.SwitchContext(continuation.ContinuationContextForB, continuation.FlagsForB); + } + + resumed_B = true; +} +``` + +It is not clear how to represent `continuation.ExecutionContextForB`, `continuation.ContinuationContextForB` and `continuation.FlagsForB`. +This IR is likely something to insert as part of inlining, yet at that point we have not laid out the continuation yet. +To solve this we will introduce a `GT_CONTINUATION_MEMBER_OFFSET` node that represents the member offset and that is replaced by a constant as part of the async transformation. + +### Exceptions + +Inlined functions can finish by throwing exceptions too. +Under normal circumstances the async infrastructures catches the exception and looks for the next continuation that may handle the exception. +Once located, the same resumption mechanism applies: the `ExecutionContext` is restored and the proper continuation context is ensured. + +Correct handling requires catching the possible exception and then rethrowing it in the post-inline IR. +The expansion for calls in try clauses would look roughly like: +```csharp +try +{ + // user code for C +} +catch (throwable ex) when (resumed_C) +{ + throwable_C = ex; +} + +if (resumed_C) +{ + AsyncHelpers.RestoreExecutionContext(Thread.CurrentThreadAssumedInitialized, continuation.ExecutionContextForB); + + if (!AsyncHelpers.IsOnRightContext(continuation.ContinuationContextForB, continuation.FlagsForB)) + { + await AsyncHelpers.SwitchContext(continuation.ContinuationContextForB, continuation.FlagsForB); + } + + resumed_B = true; + + if (throwable_C != null) + ThrowExact(throwable_C); +} +``` + +Initially we will skip inlining async calls that may suspend in try clauses. + +## Some examples + +In the following assume the presence of two synchronization contexts `_syncContext1` and `_syncContext2` that set up the global environment in some distinguishable way. + +```csharp +public static async Task Foo() +{ + SynchronizationContext.SetSynchronizationContext(_syncContext1); + await Bar(); +} + +private static async Task Bar() +{ + SynchronizationContext.SetSynchronizationContext(_syncContext2); + await Baz(); +} + +private static async Task Baz() +{ + await new AlwaysThreadPoolAwaitable(); +} + +private struct AlwaysThreadPoolAwaitable : INotifyCompletion +{ + public bool IsCompleted => false; + public void OnCompleted(Action continuation) + { + ThreadPool.QueueUserWorkItem(_ => { Thread.Sleep(100); continuation(); }); + } + public void GetResult() { } + + public AlwaysThreadPoolAwaitable GetAwaiter() => this; +} +``` + +- `Foo()` runs in the ambient synchronization context but switches its own synchronization context before awaiting `Bar()`. Its continuation captures `_syncContext1` as the continuation context. +- `Bar()` also switches its own synchronization context before awaiting `Baz()`. Its continuation captures `_syncContext2` as the continuation context. +- `Baz()` always suspends, switching onto a thread pool thread. No continuation context is captured for custom awaitables. +- On resumption, async infrastructure runs `Baz` directly in the context of `AlwaysThreadPoolAwaitable` +- Next, async infrastructure switches to `_syncContext2` to run `Bar`'s continuation +- Next, async infrastructure switches to `_syncContext1` to run `Foo`'s continuation + +Inlining in this case would need to suspend+switch+resume in both the `IsOnRightContext` checks. + +```csharp +private static AsyncLocal s_local = new(); +public static async Task Foo() +{ + s_local.Value = 1; + await Bar(); + Console.WriteLine(s_local.Value); +} + +private static async Task Bar() +{ + s_local.Value = 2; + await Baz(); + Console.WriteLine(s_local.Value); +} + +private static async Task Baz() +{ + s_local.Value = 3; + await new AlwaysThreadPoolAwaitable(); + Console.WriteLine(s_local.Value); +} +``` + +Similar to above, except this time we do not need a full bail out; rather, the post-inline handling will restore the `ExecutionContext` to make sure the value of `s_local` is properly restored. +Proper output is 3,2,1 and relies on `ExecutionContext`s having been restored by the post-inline handling. From dd130c3ddb536a98eabbde4f5b01f102f602c276 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 13:34:19 +0200 Subject: [PATCH 16/41] Add AsyncHelpers.IsOnRightContext and expose it and RestoreExecutionContext to the JIT General runtime async inlining needs to check, when an inlined callee logically returns to its caller after having been resumed, whether it is already running in the continuation context the caller's continuation captured, and to restore that continuation's ExecutionContext. IsOnRightContext mirrors the 'can inline' conditions in RuntimeAsyncTaskContinuation.QueueIfNecessary; the two must be kept in sync. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../CompilerServices/AsyncHelpers.CoreCLR.cs | 37 +++++++++++++++++++ src/coreclr/inc/corinfo.h | 6 +++ src/coreclr/inc/jiteeversionguid.h | 10 ++--- .../tools/Common/JitInterface/CorInfoImpl.cs | 2 + .../tools/Common/JitInterface/CorInfoTypes.cs | 2 + .../tools/superpmi/superpmi-shared/agnostic.h | 2 + .../superpmi-shared/methodcontext.cpp | 4 ++ src/coreclr/vm/corelib.h | 2 + src/coreclr/vm/jitinterface.cpp | 2 + 9 files changed, 62 insertions(+), 5 deletions(-) diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index 2fd4bb51f3b564..5a9738bfbc8928 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -1489,6 +1489,43 @@ private static void CaptureContinuationContextFlags(ref ContinuationFlags flags, flags |= ContinuationFlags.ContinueOnThreadPool; } + // Check whether the current thread already satisfies the continuation context captured in a + // continuation, i.e. whether resuming that continuation here would be dispatched inline. + // + // Used by the JIT when inlining runtime async calls: when an inlined callee logically returns + // to its caller after having been resumed, the JIT must ensure it is running in the continuation + // context the caller's continuation captured. This mirrors the "can inline" conditions in + // RuntimeAsyncTaskContinuation.QueueIfNecessary; the two must be kept in sync. + private static bool IsOnRightContext(object? continuationContext, ContinuationFlags flags) + { + if ((flags & ContinuationFlags.ContinueOnThreadPool) != 0) + { + SynchronizationContext? syncCtx = Thread.CurrentThreadAssumedInitialized._synchronizationContext; + if (syncCtx != null && syncCtx.GetType() != typeof(SynchronizationContext)) + { + return false; + } + + TaskScheduler? sched = TaskScheduler.InternalCurrent; + return sched is null || sched == TaskScheduler.Default; + } + + if ((flags & ContinuationFlags.ContinueOnCapturedSynchronizationContext) != 0) + { + Debug.Assert(continuationContext is SynchronizationContext); + return (SynchronizationContext)continuationContext! == Thread.CurrentThreadAssumedInitialized._synchronizationContext; + } + + if ((flags & ContinuationFlags.ContinueOnCapturedTaskScheduler) != 0) + { + Debug.Assert(continuationContext is TaskScheduler); + return (TaskScheduler)continuationContext! == TaskScheduler.InternalCurrent; + } + + // No continuation context was captured, so there is nothing to switch to. + return true; + } + // Finish suspension in the common case of a custom await or for a ConfigureAwait(false) task await: // - Capture current ExecutionContext into the continuation // - Restore ExecutionContext and SynchronizationContext to the current Thread object diff --git a/src/coreclr/inc/corinfo.h b/src/coreclr/inc/corinfo.h index f4f63bc3e41739..3400f61e270a0a 100644 --- a/src/coreclr/inc/corinfo.h +++ b/src/coreclr/inc/corinfo.h @@ -1834,6 +1834,12 @@ struct CORINFO_ASYNC_INFO CORINFO_METHOD_HANDLE restoreContextsMethHnd; // Method handle for AsyncHelpers.RestoreContextsOnSuspension, used before suspending in async methods CORINFO_METHOD_HANDLE restoreContextsOnSuspensionMethHnd; + // Method handle for AsyncHelpers.RestoreExecutionContext, used when an inlined async callee + // logically returns to its caller after having been resumed + CORINFO_METHOD_HANDLE restoreExecutionContextMethHnd; + // Method handle for AsyncHelpers.IsOnRightContext, used to check whether an inlined async callee + // can logically return to its caller without switching continuation context + CORINFO_METHOD_HANDLE isOnRightContextMethHnd; // Finish suspension without saving continuation context (i.e. custom awaiter or ConfigureAwait(false)) CORINFO_METHOD_HANDLE finishSuspensionNoContinuationContextMethHnd; // Finish suspension with saving continuation context (i.e. normal task await) diff --git a/src/coreclr/inc/jiteeversionguid.h b/src/coreclr/inc/jiteeversionguid.h index c9cc3ba2e7c957..2afeabaeb50afd 100644 --- a/src/coreclr/inc/jiteeversionguid.h +++ b/src/coreclr/inc/jiteeversionguid.h @@ -37,11 +37,11 @@ #include -constexpr GUID JITEEVersionIdentifier = { /* 5eb17041-1abf-4d57-80a9-a597869ddab4 */ - 0x5eb17041, - 0x1abf, - 0x4d57, - {0x80, 0xa9, 0xa5, 0x97, 0x86, 0x9d, 0xda, 0xb4} +constexpr GUID JITEEVersionIdentifier = { /* 923c9eda-2e4e-4d36-b52b-07170ec36748 */ + 0x923c9eda, + 0x2e4e, + 0x4d36, + {0xb5, 0x2b, 0x07, 0x17, 0x0e, 0xc3, 0x67, 0x48} }; #endif // JIT_EE_VERSIONING_GUID_H diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs index 41f1e9b61a587b..9497667313d65c 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs @@ -3621,6 +3621,8 @@ private void getAsyncInfo(ref CORINFO_ASYNC_INFO pAsyncInfoOut) pAsyncInfoOut.captureContextsMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("CaptureContexts"u8, null)); pAsyncInfoOut.restoreContextsMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreContexts"u8, null)); pAsyncInfoOut.restoreContextsOnSuspensionMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreContextsOnSuspension"u8, null)); + pAsyncInfoOut.restoreExecutionContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreExecutionContext"u8, null)); + pAsyncInfoOut.isOnRightContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("IsOnRightContext"u8, null)); pAsyncInfoOut.finishSuspensionNoContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionNoContinuationContext"u8, null)); pAsyncInfoOut.finishSuspensionWithContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionWithContinuationContext"u8, null)); } diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs index 2402eda54f0d59..e2e508b1e63d5a 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs @@ -970,6 +970,8 @@ public unsafe struct CORINFO_ASYNC_INFO public CORINFO_METHOD_STRUCT_* captureContextsMethHnd; public CORINFO_METHOD_STRUCT_* restoreContextsMethHnd; public CORINFO_METHOD_STRUCT_* restoreContextsOnSuspensionMethHnd; + public CORINFO_METHOD_STRUCT_* restoreExecutionContextMethHnd; + public CORINFO_METHOD_STRUCT_* isOnRightContextMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionNoContinuationContextMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionWithContinuationContextMethHnd; } diff --git a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h index 0f64d7436d779f..d992b078bf957a 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h +++ b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h @@ -237,6 +237,8 @@ struct Agnostic_CORINFO_ASYNC_INFO DWORDLONG captureContextsMethHnd; DWORDLONG restoreContextsMethHnd; DWORDLONG restoreContextsOnSuspensionMethHnd; + DWORDLONG restoreExecutionContextMethHnd; + DWORDLONG isOnRightContextMethHnd; DWORDLONG finishSuspensionNoContinuationContextMethHnd; DWORDLONG finishSuspensionWithContinuationContextMethHnd; }; diff --git a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp index 778254d4c7e867..293c60c696181d 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp +++ b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp @@ -4394,6 +4394,8 @@ void MethodContext::recGetAsyncInfo(const CORINFO_ASYNC_INFO* pAsyncInfo) value.captureContextsMethHnd = CastHandle(pAsyncInfo->captureContextsMethHnd); value.restoreContextsMethHnd = CastHandle(pAsyncInfo->restoreContextsMethHnd); value.restoreContextsOnSuspensionMethHnd = CastHandle(pAsyncInfo->restoreContextsOnSuspensionMethHnd); + value.restoreExecutionContextMethHnd = CastHandle(pAsyncInfo->restoreExecutionContextMethHnd); + value.isOnRightContextMethHnd = CastHandle(pAsyncInfo->isOnRightContextMethHnd); value.finishSuspensionNoContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionNoContinuationContextMethHnd); value.finishSuspensionWithContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionWithContinuationContextMethHnd); @@ -4420,6 +4422,8 @@ void MethodContext::repGetAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->captureContextsMethHnd = (CORINFO_METHOD_HANDLE)value.captureContextsMethHnd; pAsyncInfoOut->restoreContextsMethHnd = (CORINFO_METHOD_HANDLE)value.restoreContextsMethHnd; pAsyncInfoOut->restoreContextsOnSuspensionMethHnd = (CORINFO_METHOD_HANDLE)value.restoreContextsOnSuspensionMethHnd; + pAsyncInfoOut->restoreExecutionContextMethHnd = (CORINFO_METHOD_HANDLE)value.restoreExecutionContextMethHnd; + pAsyncInfoOut->isOnRightContextMethHnd = (CORINFO_METHOD_HANDLE)value.isOnRightContextMethHnd; pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionNoContinuationContextMethHnd; pAsyncInfoOut->finishSuspensionWithContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionWithContinuationContextMethHnd; DEBUG_REP(dmpGetAsyncInfo(0, value)); diff --git a/src/coreclr/vm/corelib.h b/src/coreclr/vm/corelib.h index 60747f18d01d72..22557e639f0d56 100644 --- a/src/coreclr/vm/corelib.h +++ b/src/coreclr/vm/corelib.h @@ -730,6 +730,8 @@ DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_CONTINUATION_CONTEXT, CaptureContinuat DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_CONTEXTS, CaptureContexts, NoSig) DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS, RestoreContexts, NoSig) DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS_ON_SUSPENSION, RestoreContextsOnSuspension, NoSig) +DEFINE_METHOD(ASYNC_HELPERS, RESTORE_EXECUTION_CONTEXT, RestoreExecutionContext, NoSig) +DEFINE_METHOD(ASYNC_HELPERS, IS_ON_RIGHT_CONTEXT, IsOnRightContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT, FinishSuspensionNoContinuationContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_WITH_CONTINUATION_CONTEXT, FinishSuspensionWithContinuationContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, ASYNC_CALL_CONTINUATION, AsyncCallContinuation, NoSig) diff --git a/src/coreclr/vm/jitinterface.cpp b/src/coreclr/vm/jitinterface.cpp index 784da9288c1ff1..17281405dff813 100644 --- a/src/coreclr/vm/jitinterface.cpp +++ b/src/coreclr/vm/jitinterface.cpp @@ -10389,6 +10389,8 @@ void CEEInfo::getAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->captureContextsMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__CAPTURE_CONTEXTS)); pAsyncInfoOut->restoreContextsMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_CONTEXTS)); pAsyncInfoOut->restoreContextsOnSuspensionMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_CONTEXTS_ON_SUSPENSION)); + pAsyncInfoOut->restoreExecutionContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_EXECUTION_CONTEXT)); + pAsyncInfoOut->isOnRightContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__IS_ON_RIGHT_CONTEXT)); pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT)); pAsyncInfoOut->finishSuspensionWithContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_WITH_CONTINUATION_CONTEXT)); From 76237421a02f6bf0a636ee7a1000b07c692597d2 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 13:44:10 +0200 Subject: [PATCH 17/41] Add AsyncHelpers.SwitchContext and expose it to the JIT Suspends and resumes in a specified continuation context, rather than one captured at the suspension point. General runtime async inlining uses this when an inlined callee logically returns to its caller after having been resumed and IsOnRightContext reports the current context does not match the one the caller's continuation captured. It suspends on an already completed task so that the dispatcher re-dispatches immediately. That dispatch runs with canInline: false, so it always posts or schedules onto the requested context instead of running inline. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../CompilerServices/AsyncHelpers.CoreCLR.cs | 43 +++++++++++++++++++ src/coreclr/inc/corinfo.h | 3 ++ src/coreclr/inc/jiteeversionguid.h | 10 ++--- .../tools/Common/JitInterface/CorInfoImpl.cs | 1 + .../tools/Common/JitInterface/CorInfoTypes.cs | 1 + .../tools/superpmi/superpmi-shared/agnostic.h | 1 + .../superpmi-shared/methodcontext.cpp | 2 + src/coreclr/vm/corelib.h | 1 + src/coreclr/vm/jitinterface.cpp | 1 + 9 files changed, 58 insertions(+), 5 deletions(-) diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index 5a9738bfbc8928..a377c8a84a9ded 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -1526,6 +1526,49 @@ private static bool IsOnRightContext(object? continuationContext, ContinuationFl return true; } + // Suspend and resume in the specified continuation context. + // + // Used by the JIT when inlining runtime async calls: when an inlined callee logically + // returns to its caller after having been resumed, and IsOnRightContext reports that the + // current context does not match the one the caller's continuation captured, we must get + // back onto that context before continuing. + // + // 'flags' must contain only ContinuationFlags.AllContinuationFlags bits. + // + // Suspending on an already completed task makes the dispatcher re-dispatch the continuation + // immediately. Because that dispatch happens with canInline: false, it always posts or + // schedules onto the requested context rather than running inline here, which is what we + // want -- the JIT only calls this when we are known to be on the wrong context. + [BypassReadyToRun] + [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)] + private static unsafe void SwitchContext(object? continuationContext, ContinuationFlags flags) + { + Debug.Assert((flags & ~ContinuationFlags.AllContinuationFlags) == 0); + + ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState; + Continuation? sentinelContinuation = state.SentinelContinuation ??= new Continuation(); + + RuntimeAsyncTaskContinuation? taskCont = state.CachedTaskContinuation; + if (taskCont != null) + { + state.CachedTaskContinuation = null; + } + else + { + taskCont = new RuntimeAsyncTaskContinuation(); + } + + taskCont.Initialize(Task.CompletedTask); + taskCont.ContinuationContext = continuationContext; + taskCont.Flags |= flags; + + sentinelContinuation.Next = taskCont; + state.StackState->TaskContinuation = taskCont; + + state.CaptureContexts(); + AsyncSuspend(taskCont); + } + // Finish suspension in the common case of a custom await or for a ConfigureAwait(false) task await: // - Capture current ExecutionContext into the continuation // - Restore ExecutionContext and SynchronizationContext to the current Thread object diff --git a/src/coreclr/inc/corinfo.h b/src/coreclr/inc/corinfo.h index 3400f61e270a0a..e7a621111f37ce 100644 --- a/src/coreclr/inc/corinfo.h +++ b/src/coreclr/inc/corinfo.h @@ -1840,6 +1840,9 @@ struct CORINFO_ASYNC_INFO // Method handle for AsyncHelpers.IsOnRightContext, used to check whether an inlined async callee // can logically return to its caller without switching continuation context CORINFO_METHOD_HANDLE isOnRightContextMethHnd; + // Method handle for AsyncHelpers.SwitchContext, used to switch to the continuation context of an + // inlined async callee's caller when IsOnRightContext reports a mismatch + CORINFO_METHOD_HANDLE switchContextMethHnd; // Finish suspension without saving continuation context (i.e. custom awaiter or ConfigureAwait(false)) CORINFO_METHOD_HANDLE finishSuspensionNoContinuationContextMethHnd; // Finish suspension with saving continuation context (i.e. normal task await) diff --git a/src/coreclr/inc/jiteeversionguid.h b/src/coreclr/inc/jiteeversionguid.h index 2afeabaeb50afd..79f35210db1882 100644 --- a/src/coreclr/inc/jiteeversionguid.h +++ b/src/coreclr/inc/jiteeversionguid.h @@ -37,11 +37,11 @@ #include -constexpr GUID JITEEVersionIdentifier = { /* 923c9eda-2e4e-4d36-b52b-07170ec36748 */ - 0x923c9eda, - 0x2e4e, - 0x4d36, - {0xb5, 0x2b, 0x07, 0x17, 0x0e, 0xc3, 0x67, 0x48} +constexpr GUID JITEEVersionIdentifier = { /* 0ab9b4f4-b822-4e35-9aea-4c5696b6b26d */ + 0x0ab9b4f4, + 0xb822, + 0x4e35, + {0x9a, 0xea, 0x4c, 0x56, 0x96, 0xb6, 0xb2, 0x6d} }; #endif // JIT_EE_VERSIONING_GUID_H diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs index 9497667313d65c..fc2c19d922c63e 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs @@ -3623,6 +3623,7 @@ private void getAsyncInfo(ref CORINFO_ASYNC_INFO pAsyncInfoOut) pAsyncInfoOut.restoreContextsOnSuspensionMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreContextsOnSuspension"u8, null)); pAsyncInfoOut.restoreExecutionContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreExecutionContext"u8, null)); pAsyncInfoOut.isOnRightContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("IsOnRightContext"u8, null)); + pAsyncInfoOut.switchContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("SwitchContext"u8, null)); pAsyncInfoOut.finishSuspensionNoContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionNoContinuationContext"u8, null)); pAsyncInfoOut.finishSuspensionWithContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionWithContinuationContext"u8, null)); } diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs index e2e508b1e63d5a..32105aa5a396d0 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs @@ -972,6 +972,7 @@ public unsafe struct CORINFO_ASYNC_INFO public CORINFO_METHOD_STRUCT_* restoreContextsOnSuspensionMethHnd; public CORINFO_METHOD_STRUCT_* restoreExecutionContextMethHnd; public CORINFO_METHOD_STRUCT_* isOnRightContextMethHnd; + public CORINFO_METHOD_STRUCT_* switchContextMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionNoContinuationContextMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionWithContinuationContextMethHnd; } diff --git a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h index d992b078bf957a..c4280fd420175a 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h +++ b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h @@ -239,6 +239,7 @@ struct Agnostic_CORINFO_ASYNC_INFO DWORDLONG restoreContextsOnSuspensionMethHnd; DWORDLONG restoreExecutionContextMethHnd; DWORDLONG isOnRightContextMethHnd; + DWORDLONG switchContextMethHnd; DWORDLONG finishSuspensionNoContinuationContextMethHnd; DWORDLONG finishSuspensionWithContinuationContextMethHnd; }; diff --git a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp index 293c60c696181d..089832bb12b781 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp +++ b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp @@ -4396,6 +4396,7 @@ void MethodContext::recGetAsyncInfo(const CORINFO_ASYNC_INFO* pAsyncInfo) value.restoreContextsOnSuspensionMethHnd = CastHandle(pAsyncInfo->restoreContextsOnSuspensionMethHnd); value.restoreExecutionContextMethHnd = CastHandle(pAsyncInfo->restoreExecutionContextMethHnd); value.isOnRightContextMethHnd = CastHandle(pAsyncInfo->isOnRightContextMethHnd); + value.switchContextMethHnd = CastHandle(pAsyncInfo->switchContextMethHnd); value.finishSuspensionNoContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionNoContinuationContextMethHnd); value.finishSuspensionWithContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionWithContinuationContextMethHnd); @@ -4424,6 +4425,7 @@ void MethodContext::repGetAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->restoreContextsOnSuspensionMethHnd = (CORINFO_METHOD_HANDLE)value.restoreContextsOnSuspensionMethHnd; pAsyncInfoOut->restoreExecutionContextMethHnd = (CORINFO_METHOD_HANDLE)value.restoreExecutionContextMethHnd; pAsyncInfoOut->isOnRightContextMethHnd = (CORINFO_METHOD_HANDLE)value.isOnRightContextMethHnd; + pAsyncInfoOut->switchContextMethHnd = (CORINFO_METHOD_HANDLE)value.switchContextMethHnd; pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionNoContinuationContextMethHnd; pAsyncInfoOut->finishSuspensionWithContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionWithContinuationContextMethHnd; DEBUG_REP(dmpGetAsyncInfo(0, value)); diff --git a/src/coreclr/vm/corelib.h b/src/coreclr/vm/corelib.h index 22557e639f0d56..498af055bc34b7 100644 --- a/src/coreclr/vm/corelib.h +++ b/src/coreclr/vm/corelib.h @@ -732,6 +732,7 @@ DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS, RestoreContexts, No DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS_ON_SUSPENSION, RestoreContextsOnSuspension, NoSig) DEFINE_METHOD(ASYNC_HELPERS, RESTORE_EXECUTION_CONTEXT, RestoreExecutionContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, IS_ON_RIGHT_CONTEXT, IsOnRightContext, NoSig) +DEFINE_METHOD(ASYNC_HELPERS, SWITCH_CONTEXT, SwitchContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT, FinishSuspensionNoContinuationContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_WITH_CONTINUATION_CONTEXT, FinishSuspensionWithContinuationContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, ASYNC_CALL_CONTINUATION, AsyncCallContinuation, NoSig) diff --git a/src/coreclr/vm/jitinterface.cpp b/src/coreclr/vm/jitinterface.cpp index 17281405dff813..31bc523a307d8c 100644 --- a/src/coreclr/vm/jitinterface.cpp +++ b/src/coreclr/vm/jitinterface.cpp @@ -10391,6 +10391,7 @@ void CEEInfo::getAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->restoreContextsOnSuspensionMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_CONTEXTS_ON_SUSPENSION)); pAsyncInfoOut->restoreExecutionContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_EXECUTION_CONTEXT)); pAsyncInfoOut->isOnRightContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__IS_ON_RIGHT_CONTEXT)); + pAsyncInfoOut->switchContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__SWITCH_CONTEXT)); pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT)); pAsyncInfoOut->finishSuspensionWithContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_WITH_CONTINUATION_CONTEXT)); From 27edc8f8d6df038e39b56fbcdd967fbb3bae4454 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 13:54:31 +0200 Subject: [PATCH 18/41] JIT: Add continuation members for inlined async frames Extends ContinuationMember with slots for the ExecutionContext, continuation context and continuation flags that an inlined async frame captures when it logically returns to its caller. These are keyed by inline depth rather than by individual inline frame. Two frames at the same depth can share storage because their live ranges cannot overlap: the value is written at a frame's first suspension and consumed by that same frame's post-inline IR, and one frame at a given depth must be left before another can be entered. This bounds the extra continuation members by the maximum inline depth instead of by the number of inlined async frames. No functional change yet; nothing requests these members. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 104 +++++++++++++++++++++++++++++++++++--- src/coreclr/jit/async.h | 23 ++++++++- src/coreclr/jit/inline.h | 12 +++++ 3 files changed, 132 insertions(+), 7 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index e6c576f2fc10eb..539528b5ed3167 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -53,12 +53,68 @@ ContinuationMember ContinuationMember::CustomAwaiterOfLayout(ClassLayout* layout return member; } +ContinuationMember ContinuationMember::InlineFrameExecutionContext(unsigned inlineDepth) +{ + ContinuationMember member; + member.Type = ContinuationMemberType::InlineFrameExecutionContext; + member.m_inlineDepth = inlineDepth; + return member; +} + +ContinuationMember ContinuationMember::InlineFrameContinuationContext(unsigned inlineDepth) +{ + ContinuationMember member; + member.Type = ContinuationMemberType::InlineFrameContinuationContext; + member.m_inlineDepth = inlineDepth; + return member; +} + +ContinuationMember ContinuationMember::InlineFrameFlags(unsigned inlineDepth) +{ + ContinuationMember member; + member.Type = ContinuationMemberType::InlineFrameFlags; + member.m_inlineDepth = inlineDepth; + return member; +} + ClassLayout* ContinuationMember::GetCustomAwaiterLayout() const { assert(Type == ContinuationMemberType::CustomAwaiterOfLayout); return m_customAwaiterLayout; } +unsigned ContinuationMember::GetInlineDepth() const +{ + assert((Type == ContinuationMemberType::InlineFrameExecutionContext) || + (Type == ContinuationMemberType::InlineFrameContinuationContext) || + (Type == ContinuationMemberType::InlineFrameFlags)); + return m_inlineDepth; +} + +//------------------------------------------------------------------------ +// ContinuationMember::GetStorageType: +// Get the type of the storage this member occupies in the continuation. +// +// Returns: +// TYP_STRUCT for custom awaiters (described by GetCustomAwaiterLayout), otherwise +// the primitive type stored. +// +var_types ContinuationMember::GetStorageType() const +{ + switch (Type) + { + case ContinuationMemberType::CustomAwaiterOfLayout: + return TYP_STRUCT; + case ContinuationMemberType::InlineFrameExecutionContext: + case ContinuationMemberType::InlineFrameContinuationContext: + return TYP_REF; + case ContinuationMemberType::InlineFrameFlags: + return TYP_INT; + default: + unreached(); + } +} + bool ContinuationMember::AreCompatible(const ContinuationMember& a, const ContinuationMember& b) { if (a.Type != b.Type) @@ -70,6 +126,15 @@ bool ContinuationMember::AreCompatible(const ContinuationMember& a, const Contin { case ContinuationMemberType::CustomAwaiterOfLayout: return ClassLayout::AreCompatible(a.m_customAwaiterLayout, b.m_customAwaiterLayout); + case ContinuationMemberType::InlineFrameExecutionContext: + case ContinuationMemberType::InlineFrameContinuationContext: + case ContinuationMemberType::InlineFrameFlags: + // Keyed by inline depth rather than by individual inline frame. Two frames at + // the same depth can share storage because their live ranges cannot overlap: + // the value is written at a frame's first suspension and consumed by that same + // frame's post-inline IR, and one frame at a given depth must be left before + // another can be entered. + return a.m_inlineDepth == b.m_inlineDepth; default: unreached(); } @@ -83,6 +148,15 @@ void ContinuationMember::Print() const case ContinuationMemberType::CustomAwaiterOfLayout: printf("CustomAwaiter<%s>", m_customAwaiterLayout->GetClassName()); break; + case ContinuationMemberType::InlineFrameExecutionContext: + printf("ExecutionContext for inline depth %u", m_inlineDepth); + break; + case ContinuationMemberType::InlineFrameContinuationContext: + printf("Continuation context for inline depth %u", m_inlineDepth); + break; + case ContinuationMemberType::InlineFrameFlags: + printf("Continuation flags for inline depth %u", m_inlineDepth); + break; default: unreached(); } @@ -1674,13 +1748,28 @@ ContinuationLayout* ContinuationLayoutBuilder::Create(ArrayStack& cont continue; } - ClassLayout* memberLayout = m_compiler->GetContinuationMember(memberIndex).GetCustomAwaiterLayout(); - unsigned alignment = - memberLayout->IsCustomLayout() + const ContinuationMember& member = m_compiler->GetContinuationMember(memberIndex); + + unsigned alignment; + unsigned size; + if (member.Type == ContinuationMemberType::CustomAwaiterOfLayout) + { + ClassLayout* memberLayout = member.GetCustomAwaiterLayout(); + alignment = + memberLayout->IsCustomLayout() ? (memberLayout->HasGCPtr() ? TARGET_POINTER_SIZE : 1) : m_compiler->info.compCompHnd->getClassAlignmentRequirement(memberLayout->GetClassHandle()); + size = memberLayout->GetSize(); + } + else + { + var_types storageType = member.GetStorageType(); + alignment = genTypeAlignments[storageType]; + size = genTypeSize(storageType); + } + layout->ContinuationMemberOffsets[memberIndex] = - allocLayout(std::min(alignment, (unsigned)TARGET_POINTER_SIZE), memberLayout->GetSize()); + allocLayout(std::min(alignment, (unsigned)TARGET_POINTER_SIZE), size); } if (m_needsKeepAlive) @@ -1735,8 +1824,11 @@ ContinuationLayout* ContinuationLayoutBuilder::Create(ArrayStack& cont { if (layout->ContinuationMemberOffsets[i] != UINT_MAX) { - bitmapBuilder.SetType(layout->ContinuationMemberOffsets[i], TYP_STRUCT, - m_compiler->GetContinuationMember(i).GetCustomAwaiterLayout()); + const ContinuationMember& member = m_compiler->GetContinuationMember(i); + ClassLayout* memberLayout = member.Type == ContinuationMemberType::CustomAwaiterOfLayout + ? member.GetCustomAwaiterLayout() + : nullptr; + bitmapBuilder.SetType(layout->ContinuationMemberOffsets[i], member.GetStorageType(), memberLayout); } } diff --git a/src/coreclr/jit/async.h b/src/coreclr/jit/async.h index d26c6edb826f87..f001af4cb8238a 100644 --- a/src/coreclr/jit/async.h +++ b/src/coreclr/jit/async.h @@ -3,7 +3,18 @@ enum class ContinuationMemberType { + // A custom awaiter, stored so that it can be awaited from the continuation. + // Deduplicated by layout: all awaits of compatible awaiter types share one slot. CustomAwaiterOfLayout, + // The ExecutionContext an inlined async frame captured when it logically returned + // to its caller. Deduplicated by inline depth. + InlineFrameExecutionContext, + // The continuation context an inlined async frame captured when it logically + // returned to its caller. Deduplicated by inline depth. + InlineFrameContinuationContext, + // The ContinuationFlags describing the above continuation context. Deduplicated by + // inline depth. + InlineFrameFlags, }; struct ContinuationMember @@ -11,12 +22,22 @@ struct ContinuationMember ContinuationMemberType Type; void Print() const; + private: ClassLayout* m_customAwaiterLayout; + unsigned m_inlineDepth; public: - ClassLayout* GetCustomAwaiterLayout() const; + ClassLayout* GetCustomAwaiterLayout() const; + unsigned GetInlineDepth() const; + + // The type of the storage this member requires in the continuation. + var_types GetStorageType() const; + static ContinuationMember CustomAwaiterOfLayout(ClassLayout* layout); + static ContinuationMember InlineFrameExecutionContext(unsigned inlineDepth); + static ContinuationMember InlineFrameContinuationContext(unsigned inlineDepth); + static ContinuationMember InlineFrameFlags(unsigned inlineDepth); static bool AreCompatible(const ContinuationMember& a, const ContinuationMember& b); }; diff --git a/src/coreclr/jit/inline.h b/src/coreclr/jit/inline.h index 939a7b767fd81c..68d39e37258c60 100644 --- a/src/coreclr/jit/inline.h +++ b/src/coreclr/jit/inline.h @@ -793,6 +793,18 @@ class InlineContext return m_Parent; } + // Get the inline depth of this context. The root context has depth 0. + unsigned GetDepth() const + { + unsigned depth = 0; + for (InlineContext* parent = m_Parent; parent != nullptr; parent = parent->m_Parent) + { + depth++; + } + + return depth; + } + // Get the sibling context. InlineContext* GetSibling() const { From 80bbc3d7a4e2b66c2e8f04f0433c8e738369584a Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 14:10:46 +0200 Subject: [PATCH 19/41] JIT: Add JitAsyncInlining knob and give inlinee awaits their own contexts Adds the JitAsyncInlining config knob, off by default, which lifts the CALLEE_AWAIT restriction and allows inlining async callees that may suspend. When an await in an inlinee may suspend it no longer inherits context handling from the inlining call. Instead it is treated exactly like an await in a non-inlined method and picks up the inlinee's own resumed/ExecutionContext/ SynchronizationContext locals, which the inlinee's own SaveAsyncContexts phase already creates. The cheap paths (async versions of synchronous methods and tail awaits) keep inheriting from the caller and are unchanged. With the knob off behavior is unchanged: the restructured gate reports the same CALLEE_AWAIT observation, and inlinees do not get their own context args. The knob is not yet functional when enabled; the post-inline IR and the nested capture chain at suspensions are still missing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 25 +++++++++--- src/coreclr/jit/compiler.h | 14 +++++++ src/coreclr/jit/importercalls.cpp | 67 +++++++++++++++++++++---------- src/coreclr/jit/jitconfigvalues.h | 6 +++ 4 files changed, 85 insertions(+), 27 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 539528b5ed3167..252a6d3c9c89aa 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -394,7 +394,10 @@ PhaseStatus Compiler::SaveAsyncContexts() for (BasicBlock* block : Blocks()) { - if (!compIsForInlining()) + // Async calls that inherited their context handling from the inlining call + // already have these args; AddContextArgsToAsyncCalls skips those. Only general + // async inlining introduces inlinee async calls that need their own contexts. + if (!compIsForInlining() || generalAsyncInliningEnabled()) { AddContextArgsToAsyncCalls(block); } @@ -489,11 +492,21 @@ void Compiler::AddContextArgsToAsyncCalls(BasicBlock* block) return WALK_CONTINUE; } - GenTreeCall* call = tree->AsCall(); - GenTree* resumed = m_compiler->gtNewLclVarNode(m_compiler->lvaResumedIndicator, TYP_INT); - GenTree* resumedAddr = m_compiler->gtNewLclAddrNode(m_compiler->lvaResumedIndicator, 0); - GenTree* execCtx = m_compiler->gtNewLclVarNode(m_compiler->lvaAsyncExecutionContextVar, TYP_REF); - GenTree* syncCtx = m_compiler->gtNewLclVarNode(m_compiler->lvaAsyncSynchronizationContextVar, TYP_REF); + GenTreeCall* call = tree->AsCall(); + + if (call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse) != nullptr) + { + // Already has context args: this is an async call in an inlinee that + // inherited them from the inlining call (see + // impInheritAsyncContextsFromInliner). + assert(m_compiler->compIsForInlining()); + return WALK_CONTINUE; + } + + GenTree* resumed = m_compiler->gtNewLclVarNode(m_compiler->lvaResumedIndicator, TYP_INT); + GenTree* resumedAddr = m_compiler->gtNewLclAddrNode(m_compiler->lvaResumedIndicator, 0); + GenTree* execCtx = m_compiler->gtNewLclVarNode(m_compiler->lvaAsyncExecutionContextVar, TYP_REF); + GenTree* syncCtx = m_compiler->gtNewLclVarNode(m_compiler->lvaAsyncSynchronizationContextVar, TYP_REF); JITDUMP( "Adding resumed use [%06u], resumed def [%06u] exec context [%06u], sync context [%06u] to async call [%06u]\n", dspTreeID(resumed), dspTreeID(resumedAddr), dspTreeID(execCtx), dspTreeID(syncCtx), dspTreeID(call)); diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h index 1b2e81c68b0d36..bce1460e58b9fb 100644 --- a/src/coreclr/jit/compiler.h +++ b/src/coreclr/jit/compiler.h @@ -5632,6 +5632,12 @@ class Compiler bool m_nextAwaitIsTail = false; + // Set while importing an inlinee for an async call that gets its own context + // handling instead of inheriting it from the inlining call, i.e. an await that may + // suspend inside a generally-inlined async callee. Consumed by + // impInheritAsyncContextsFromInliner. + bool m_nextAsyncCallUsesOwnContexts = false; + bool impSpillStackEntry(unsigned level, unsigned varNum #ifdef DEBUG @@ -11919,6 +11925,14 @@ class Compiler return (info.compMethodInfo->options & CORINFO_ASYNC_VERSION) != 0; } + // Is general inlining of runtime async calls enabled, i.e. inlining of async callees + // that may suspend? When disabled only the restricted cases are inlined: callees + // without any awaits, async versions of synchronous methods, and tail awaits. + static bool generalAsyncInliningEnabled() + { + return JitConfig.JitAsyncInlining() != 0; + } + //------------------------------------------------------------------------ // compMethodReturnsMultiRegRetType: Does this method return a multi-reg value? // diff --git a/src/coreclr/jit/importercalls.cpp b/src/coreclr/jit/importercalls.cpp index 288abbcc05bf17..7504ed1bc26be2 100644 --- a/src/coreclr/jit/importercalls.cpp +++ b/src/coreclr/jit/importercalls.cpp @@ -7394,38 +7394,55 @@ void Compiler::impSetupAsyncCall(GenTreeCall* call, if (compIsForInlining()) { - if (!m_nextAwaitIsTail && !compIsAsyncVersion()) + // Two cases are inlined cheaply: async versions of synchronous methods, where + // all async calls are in tail position, and explicit tail awaits. In both the + // inlinee's tail can run in the caller's context, so we can inherit all context + // handling from the inlining call and no logical frame transition is needed when + // the inlinee returns. + bool inheritsCallerContexts = m_nextAwaitIsTail || compIsAsyncVersion(); + m_nextAsyncCallUsesOwnContexts = false; + + // We cannot inline if the callee returns valueTask.AsTask(). We need to preserve + // the continuation in this case to be able to mark it with + // CORINFO_CONTINUATION_VALUETASK_ADAPTED_TO_TASK. + if ((prefixFlags & PREFIX_IS_ADAPTED_FROM_VALUETASK) != 0) { compInlineResult->NoteFatal(InlineObservation::CALLEE_AWAIT); return; } - // We cannot inline if the callee returns valueTask.AsTask() in an - // async version. We need to preserve the continuation in this case to - // be able to mark it with CORINFO_CONTINUATION_VALUETASK_ADAPTED_TO_TASK. - if ((prefixFlags & PREFIX_IS_ADAPTED_FROM_VALUETASK) != 0) + if (!inheritsCallerContexts) { - compInlineResult->NoteFatal(InlineObservation::CALLEE_AWAIT); - return; - } + if (!generalAsyncInliningEnabled()) + { + compInlineResult->NoteFatal(InlineObservation::CALLEE_AWAIT); + return; + } - // For async versions of synchronous methods all async calls are in - // tail position. Inlining is simple for these cases: we can just - // inherit all context handling from the inlining call. - assert(!compIsAsyncVersion() || ((prefixFlags & PREFIX_IS_ASYNC_VERSION_TAIL_AWAIT) != 0)); + // General case: this is a real await inside the inlinee that may suspend. It + // gets its own context handling below, exactly like an await in a non-inlined + // method, and the inlinee's own contexts (created by its SaveAsyncContexts) + // are used rather than the caller's. + JITDUMP("Call [%06u] is an await in an inlinee that may suspend\n", dspTreeID(call)); + m_nextAsyncCallUsesOwnContexts = true; + } + else + { + assert(!compIsAsyncVersion() || ((prefixFlags & PREFIX_IS_ASYNC_VERSION_TAIL_AWAIT) != 0)); - GenTreeCall* inlCall = impInlineInfo->iciCall; - JITDUMP("Call [%06u] is to function with a tail async call [%06u]\n", dspTreeID(inlCall), dspTreeID(call)); + GenTreeCall* inlCall = impInlineInfo->iciCall; + JITDUMP("Call [%06u] is to function with a tail async call [%06u]\n", dspTreeID(inlCall), dspTreeID(call)); - assert(inlCall->IsAsync()); + assert(inlCall->IsAsync()); - asyncInfo.ContinuationContextHandling = inlCall->GetAsyncInfo().ContinuationContextHandling; - // Validate that below code won't override the handling - assert((prefixFlags & PREFIX_IS_TASK_AWAIT) == 0); + asyncInfo.ContinuationContextHandling = inlCall->GetAsyncInfo().ContinuationContextHandling; + // Validate that below code won't override the handling + assert((prefixFlags & PREFIX_IS_TASK_AWAIT) == 0); - asyncInfo.IsTailAwait = - inlCall->GetAsyncInfo().IsTailAwait && (m_nextAwaitIsTail || (call->gtReturnType == info.compRetType)); - m_nextAwaitIsTail = false; + asyncInfo.IsTailAwait = + inlCall->GetAsyncInfo().IsTailAwait && (m_nextAwaitIsTail || (call->gtReturnType == info.compRetType)); + m_nextAwaitIsTail = false; + } } else { @@ -7510,6 +7527,14 @@ void Compiler::impInheritAsyncContextsFromInliner(GenTreeCall* call) return; } + if (m_nextAsyncCallUsesOwnContexts) + { + // General async inlining: this await may suspend, so it uses the inlinee's own + // contexts, which are added by the inlinee's SaveAsyncContexts phase. + m_nextAsyncCallUsesOwnContexts = false; + return; + } + GenTreeCall* inlCall = impInlineInfo->iciCall; CallArg* resumedUseArg = inlCall->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse); CallArg* resumedDefArg = inlCall->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedDef); diff --git a/src/coreclr/jit/jitconfigvalues.h b/src/coreclr/jit/jitconfigvalues.h index bfb21eaf40c916..12a5c44f6cbdc3 100644 --- a/src/coreclr/jit/jitconfigvalues.h +++ b/src/coreclr/jit/jitconfigvalues.h @@ -620,6 +620,12 @@ OPT_CONFIG_STRING(JitAsyncDefaultValueAnalysisRange, // a continuation is being reused. OPT_CONFIG_STRING(JitAsyncPreservedValueAnalysisRange, "JitAsyncPreservedValueAnalysisRange") +// Enable general inlining of runtime async calls, i.e. inlining of async +// callees that may suspend. When zero, only the restricted cases are inlined: +// callees without any awaits, async versions of synchronous methods, and tail +// awaits. +RELEASE_CONFIG_INTEGER(JitAsyncInlining, "JitAsyncInlining", 0) + RELEASE_CONFIG_INTEGER(JitEnableOptRepeat, "JitEnableOptRepeat", 1) // If zero, do not allow JitOptRepeat RELEASE_CONFIG_METHODSET(JitOptRepeat, "JitOptRepeat") // Runs optimizer multiple times on specified methods RELEASE_CONFIG_INTEGER(JitOptRepeatCount, "JitOptRepeatCount", 2) // Number of times to repeat opts when repeating From fe0ae8b3d918af1e9edb3e1ca1abd53936f8d17e Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 14:19:35 +0200 Subject: [PATCH 20/41] JIT: Detect whether an async body may suspend, and skip OSR for general async inlining SaveAsyncContexts now records whether the body contains any async call. For an inlinee without one the resumed indicator is provably always false, so the post-inline async frame IR can be skipped entirely, keeping the cheap shape for what is by far the common case. Also bails out of general async inlining in OSR methods. An OSR method is entered with the tier0 method's continuation, whose layout differs from the one the OSR method computes, so an inlined frame's members cannot be read from it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 19 ++++++++++++------- src/coreclr/jit/compiler.h | 6 ++++++ src/coreclr/jit/importercalls.cpp | 5 ++++- 3 files changed, 22 insertions(+), 8 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 252a6d3c9c89aa..48fb9d3d0df13e 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -394,13 +394,7 @@ PhaseStatus Compiler::SaveAsyncContexts() for (BasicBlock* block : Blocks()) { - // Async calls that inherited their context handling from the inlining call - // already have these args; AddContextArgsToAsyncCalls skips those. Only general - // async inlining introduces inlinee async calls that need their own contexts. - if (!compIsForInlining() || generalAsyncInliningEnabled()) - { - AddContextArgsToAsyncCalls(block); - } + AddContextArgsToAsyncCalls(block); if (!block->KindIs(BBJ_RETURN) || (block == newReturnBB)) { @@ -494,6 +488,10 @@ void Compiler::AddContextArgsToAsyncCalls(BasicBlock* block) GenTreeCall* call = tree->AsCall(); + // Record that this body has a suspension point. Inlinees without one keep the + // cheap shape: no post-inline IR and no capture chain is emitted for them. + m_compiler->compAsyncBodyMaySuspend = true; + if (call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse) != nullptr) { // Already has context args: this is an async call in an inlinee that @@ -503,6 +501,13 @@ void Compiler::AddContextArgsToAsyncCalls(BasicBlock* block) return WALK_CONTINUE; } + if (m_compiler->compIsForInlining() && !generalAsyncInliningEnabled()) + { + // Only general async inlining introduces inlinee async calls that need + // their own contexts. + return WALK_CONTINUE; + } + GenTree* resumed = m_compiler->gtNewLclVarNode(m_compiler->lvaResumedIndicator, TYP_INT); GenTree* resumedAddr = m_compiler->gtNewLclAddrNode(m_compiler->lvaResumedIndicator, 0); GenTree* execCtx = m_compiler->gtNewLclVarNode(m_compiler->lvaAsyncExecutionContextVar, TYP_REF); diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h index bce1460e58b9fb..66fa71dc1dad13 100644 --- a/src/coreclr/jit/compiler.h +++ b/src/coreclr/jit/compiler.h @@ -4391,6 +4391,12 @@ class Compiler // Variable representing "have we resumed?" for async methods unsigned lvaResumedIndicator = BAD_VAR_NUM; + // For async methods with CORINFO_ASYNC_SAVE_CONTEXTS: whether the body contains any + // async call, i.e. any point at which this method may suspend. Computed by + // SaveAsyncContexts. When false for an inlinee, its resumed indicator is provably + // always false and no post-inline async frame IR needs to be emitted for it. + bool compAsyncBodyMaySuspend = false; + #if defined(DEBUG) && defined(TARGET_XARCH) unsigned lvaReturnSpCheck = BAD_VAR_NUM; // Stores SP to confirm it is not corrupted on return. diff --git a/src/coreclr/jit/importercalls.cpp b/src/coreclr/jit/importercalls.cpp index 7504ed1bc26be2..2fe0f7d983f2be 100644 --- a/src/coreclr/jit/importercalls.cpp +++ b/src/coreclr/jit/importercalls.cpp @@ -7413,7 +7413,10 @@ void Compiler::impSetupAsyncCall(GenTreeCall* call, if (!inheritsCallerContexts) { - if (!generalAsyncInliningEnabled()) + // OSR methods are entered with the tier0 method's continuation, whose layout + // differs from the one this method computes, so the inlined frame's members + // cannot be read from it. Not supported for now. + if (!generalAsyncInliningEnabled() || impInlineRoot()->opts.IsOSR()) { compInlineResult->NoteFatal(InlineObservation::CALLEE_AWAIT); return; From 2ed700c929d48fb595c032d6a0dbb9bb04b93c3c Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 14:27:24 +0200 Subject: [PATCH 21/41] Use a purpose-built helper for restoring an inlined frame's ExecutionContext RestoreExecutionContext takes the Thread whose contexts were captured on method entry. The restore an inlined async frame needs runs only after a resumption, so it must target the thread we were resumed on instead. Adds RestoreInlinedFrameExecutionContext, which reads the current thread itself, and points the JIT-EE handle at it. This also keeps the JIT side to a single argument. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs | 11 +++++++++++ src/coreclr/inc/corinfo.h | 6 +++--- src/coreclr/inc/jiteeversionguid.h | 10 +++++----- src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs | 2 +- src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs | 2 +- src/coreclr/tools/superpmi/superpmi-shared/agnostic.h | 2 +- .../tools/superpmi/superpmi-shared/methodcontext.cpp | 4 ++-- src/coreclr/vm/corelib.h | 2 +- src/coreclr/vm/jitinterface.cpp | 2 +- 9 files changed, 26 insertions(+), 15 deletions(-) diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index a377c8a84a9ded..c13d40edf1dc19 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -1526,6 +1526,17 @@ private static bool IsOnRightContext(object? continuationContext, ContinuationFl return true; } + // Restore the ExecutionContext that an inlined async frame captured when it logically + // returned to its caller. + // + // Unlike the synchronous restore at the end of a method, this runs only after a + // resumption, so it must target the thread we were resumed on rather than the one + // whose contexts were captured on entry. + private static void RestoreInlinedFrameExecutionContext(ExecutionContext? previousExecCtx) + { + RestoreExecutionContext(Thread.CurrentThreadAssumedInitialized, previousExecCtx); + } + // Suspend and resume in the specified continuation context. // // Used by the JIT when inlining runtime async calls: when an inlined callee logically diff --git a/src/coreclr/inc/corinfo.h b/src/coreclr/inc/corinfo.h index e7a621111f37ce..04c30fa2a97c3e 100644 --- a/src/coreclr/inc/corinfo.h +++ b/src/coreclr/inc/corinfo.h @@ -1834,9 +1834,9 @@ struct CORINFO_ASYNC_INFO CORINFO_METHOD_HANDLE restoreContextsMethHnd; // Method handle for AsyncHelpers.RestoreContextsOnSuspension, used before suspending in async methods CORINFO_METHOD_HANDLE restoreContextsOnSuspensionMethHnd; - // Method handle for AsyncHelpers.RestoreExecutionContext, used when an inlined async callee - // logically returns to its caller after having been resumed - CORINFO_METHOD_HANDLE restoreExecutionContextMethHnd; + // Method handle for AsyncHelpers.RestoreInlinedFrameExecutionContext, used when an inlined + // async callee logically returns to its caller after having been resumed + CORINFO_METHOD_HANDLE restoreInlinedFrameExecutionContextMethHnd; // Method handle for AsyncHelpers.IsOnRightContext, used to check whether an inlined async callee // can logically return to its caller without switching continuation context CORINFO_METHOD_HANDLE isOnRightContextMethHnd; diff --git a/src/coreclr/inc/jiteeversionguid.h b/src/coreclr/inc/jiteeversionguid.h index 79f35210db1882..7aa7cbf94d576a 100644 --- a/src/coreclr/inc/jiteeversionguid.h +++ b/src/coreclr/inc/jiteeversionguid.h @@ -37,11 +37,11 @@ #include -constexpr GUID JITEEVersionIdentifier = { /* 0ab9b4f4-b822-4e35-9aea-4c5696b6b26d */ - 0x0ab9b4f4, - 0xb822, - 0x4e35, - {0x9a, 0xea, 0x4c, 0x56, 0x96, 0xb6, 0xb2, 0x6d} +constexpr GUID JITEEVersionIdentifier = { /* 2f34f847-03c5-43b5-9261-e526a9bcedc9 */ + 0x2f34f847, + 0x03c5, + 0x43b5, + {0x92, 0x61, 0xe5, 0x26, 0xa9, 0xbc, 0xed, 0xc9} }; #endif // JIT_EE_VERSIONING_GUID_H diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs index fc2c19d922c63e..a7aba5a6d1f302 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs @@ -3621,7 +3621,7 @@ private void getAsyncInfo(ref CORINFO_ASYNC_INFO pAsyncInfoOut) pAsyncInfoOut.captureContextsMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("CaptureContexts"u8, null)); pAsyncInfoOut.restoreContextsMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreContexts"u8, null)); pAsyncInfoOut.restoreContextsOnSuspensionMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreContextsOnSuspension"u8, null)); - pAsyncInfoOut.restoreExecutionContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreExecutionContext"u8, null)); + pAsyncInfoOut.restoreInlinedFrameExecutionContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreInlinedFrameExecutionContext"u8, null)); pAsyncInfoOut.isOnRightContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("IsOnRightContext"u8, null)); pAsyncInfoOut.switchContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("SwitchContext"u8, null)); pAsyncInfoOut.finishSuspensionNoContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionNoContinuationContext"u8, null)); diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs index 32105aa5a396d0..013a5c6638ac1d 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs @@ -970,7 +970,7 @@ public unsafe struct CORINFO_ASYNC_INFO public CORINFO_METHOD_STRUCT_* captureContextsMethHnd; public CORINFO_METHOD_STRUCT_* restoreContextsMethHnd; public CORINFO_METHOD_STRUCT_* restoreContextsOnSuspensionMethHnd; - public CORINFO_METHOD_STRUCT_* restoreExecutionContextMethHnd; + public CORINFO_METHOD_STRUCT_* restoreInlinedFrameExecutionContextMethHnd; public CORINFO_METHOD_STRUCT_* isOnRightContextMethHnd; public CORINFO_METHOD_STRUCT_* switchContextMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionNoContinuationContextMethHnd; diff --git a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h index c4280fd420175a..e062a03b1bd02c 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h +++ b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h @@ -237,7 +237,7 @@ struct Agnostic_CORINFO_ASYNC_INFO DWORDLONG captureContextsMethHnd; DWORDLONG restoreContextsMethHnd; DWORDLONG restoreContextsOnSuspensionMethHnd; - DWORDLONG restoreExecutionContextMethHnd; + DWORDLONG restoreInlinedFrameExecutionContextMethHnd; DWORDLONG isOnRightContextMethHnd; DWORDLONG switchContextMethHnd; DWORDLONG finishSuspensionNoContinuationContextMethHnd; diff --git a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp index 089832bb12b781..efeb1ab0ac2b62 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp +++ b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp @@ -4394,7 +4394,7 @@ void MethodContext::recGetAsyncInfo(const CORINFO_ASYNC_INFO* pAsyncInfo) value.captureContextsMethHnd = CastHandle(pAsyncInfo->captureContextsMethHnd); value.restoreContextsMethHnd = CastHandle(pAsyncInfo->restoreContextsMethHnd); value.restoreContextsOnSuspensionMethHnd = CastHandle(pAsyncInfo->restoreContextsOnSuspensionMethHnd); - value.restoreExecutionContextMethHnd = CastHandle(pAsyncInfo->restoreExecutionContextMethHnd); + value.restoreInlinedFrameExecutionContextMethHnd = CastHandle(pAsyncInfo->restoreInlinedFrameExecutionContextMethHnd); value.isOnRightContextMethHnd = CastHandle(pAsyncInfo->isOnRightContextMethHnd); value.switchContextMethHnd = CastHandle(pAsyncInfo->switchContextMethHnd); value.finishSuspensionNoContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionNoContinuationContextMethHnd); @@ -4423,7 +4423,7 @@ void MethodContext::repGetAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->captureContextsMethHnd = (CORINFO_METHOD_HANDLE)value.captureContextsMethHnd; pAsyncInfoOut->restoreContextsMethHnd = (CORINFO_METHOD_HANDLE)value.restoreContextsMethHnd; pAsyncInfoOut->restoreContextsOnSuspensionMethHnd = (CORINFO_METHOD_HANDLE)value.restoreContextsOnSuspensionMethHnd; - pAsyncInfoOut->restoreExecutionContextMethHnd = (CORINFO_METHOD_HANDLE)value.restoreExecutionContextMethHnd; + pAsyncInfoOut->restoreInlinedFrameExecutionContextMethHnd = (CORINFO_METHOD_HANDLE)value.restoreInlinedFrameExecutionContextMethHnd; pAsyncInfoOut->isOnRightContextMethHnd = (CORINFO_METHOD_HANDLE)value.isOnRightContextMethHnd; pAsyncInfoOut->switchContextMethHnd = (CORINFO_METHOD_HANDLE)value.switchContextMethHnd; pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionNoContinuationContextMethHnd; diff --git a/src/coreclr/vm/corelib.h b/src/coreclr/vm/corelib.h index 498af055bc34b7..dbec6e83f5bbd0 100644 --- a/src/coreclr/vm/corelib.h +++ b/src/coreclr/vm/corelib.h @@ -730,7 +730,7 @@ DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_CONTINUATION_CONTEXT, CaptureContinuat DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_CONTEXTS, CaptureContexts, NoSig) DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS, RestoreContexts, NoSig) DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS_ON_SUSPENSION, RestoreContextsOnSuspension, NoSig) -DEFINE_METHOD(ASYNC_HELPERS, RESTORE_EXECUTION_CONTEXT, RestoreExecutionContext, NoSig) +DEFINE_METHOD(ASYNC_HELPERS, RESTORE_INLINED_FRAME_EXECUTION_CONTEXT, RestoreInlinedFrameExecutionContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, IS_ON_RIGHT_CONTEXT, IsOnRightContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, SWITCH_CONTEXT, SwitchContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT, FinishSuspensionNoContinuationContext, NoSig) diff --git a/src/coreclr/vm/jitinterface.cpp b/src/coreclr/vm/jitinterface.cpp index 31bc523a307d8c..051e9c0eed19d0 100644 --- a/src/coreclr/vm/jitinterface.cpp +++ b/src/coreclr/vm/jitinterface.cpp @@ -10389,7 +10389,7 @@ void CEEInfo::getAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->captureContextsMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__CAPTURE_CONTEXTS)); pAsyncInfoOut->restoreContextsMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_CONTEXTS)); pAsyncInfoOut->restoreContextsOnSuspensionMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_CONTEXTS_ON_SUSPENSION)); - pAsyncInfoOut->restoreExecutionContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_EXECUTION_CONTEXT)); + pAsyncInfoOut->restoreInlinedFrameExecutionContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_INLINED_FRAME_EXECUTION_CONTEXT)); pAsyncInfoOut->isOnRightContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__IS_ON_RIGHT_CONTEXT)); pAsyncInfoOut->switchContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__SWITCH_CONTEXT)); pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT)); From 5ca846e4248b99819ffa61dbe1a6854439b76e8e Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 14:43:39 +0200 Subject: [PATCH 22/41] JIT: Apply OSR async context handling to the root frame only SaveAsyncContexts runs for inlinees too, and JIT_FLAG_OSR is not cleared for an inlinee, so opts.IsOSR() is true while importing an inlinee of an OSR method. The OSR-specific handling was therefore applied to inlinees, where all three parts of it are wrong: - the resumed indicator was initialized from the continuation argument, but an inlinee has not resumed just because the OSR method was entered via a resumption - CaptureContexts was skipped, leaving the inlinee's context locals unassigned for its own RestoreContexts - the inlinee's fresh temps were marked lvIsOSRLocal, though they have no home in the tier0 frame An inlinee inside an OSR method starts a fresh logical frame, so it needs the same treatment as in a non-OSR method. With this, general async inlining works in OSR methods: an inlinee's resumed indicator can only become true at a resumption point belonging to that inlinee's own async calls, which are resumption points of the OSR method's continuation layout, so the continuation is never the tier0 one where it is read. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 10 ++++++++-- src/coreclr/jit/importercalls.cpp | 5 +---- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 48fb9d3d0df13e..ac6e885a3280e3 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -218,6 +218,12 @@ PhaseStatus Compiler::SaveAsyncContexts() return PhaseStatus::MODIFIED_NOTHING; } + // OSR context handling concerns the root frame only. An inlinee inside an OSR method + // still starts a fresh logical frame: it must capture its own contexts, its resumed + // indicator starts out false regardless of how the OSR method was entered, and its + // locals have no home in the tier0 frame. + bool isOSRRootFrame = opts.IsOSR() && !compIsForInlining(); + // Create locals for Thread, ExecutionContext and SynchronizationContext lvaAsyncThreadObjectVar = lvaGrabTemp(false DEBUGARG("Async Thread")); lvaGetDesc(lvaAsyncThreadObjectVar)->lvType = TYP_REF; @@ -231,7 +237,7 @@ PhaseStatus Compiler::SaveAsyncContexts() lvaResumedIndicator = lvaGrabTemp(false DEBUGARG("Async Resumed")); lvaGetDesc(lvaResumedIndicator)->lvType = TYP_UBYTE; - if (opts.IsOSR()) + if (isOSRRootFrame) { lvaGetDesc(lvaAsyncThreadObjectVar)->lvIsOSRLocal = true; lvaGetDesc(lvaAsyncExecutionContextVar)->lvIsOSRLocal = true; @@ -320,7 +326,7 @@ PhaseStatus Compiler::SaveAsyncContexts() // Insert CaptureContexts call before the try (keep it before so the // try/finally can be removed if there is no exception side effects). // For OSR, we did this in the tier0 method. - if (opts.IsOSR()) + if (isOSRRootFrame) { // In the OSR method we compute the initial value of the resumption indicator based on the continuation arg. GenTree* continuation = gtNewLclVarNode(lvaAsyncContinuationArg, TYP_REF); diff --git a/src/coreclr/jit/importercalls.cpp b/src/coreclr/jit/importercalls.cpp index 2fe0f7d983f2be..7504ed1bc26be2 100644 --- a/src/coreclr/jit/importercalls.cpp +++ b/src/coreclr/jit/importercalls.cpp @@ -7413,10 +7413,7 @@ void Compiler::impSetupAsyncCall(GenTreeCall* call, if (!inheritsCallerContexts) { - // OSR methods are entered with the tier0 method's continuation, whose layout - // differs from the one this method computes, so the inlined frame's members - // cannot be read from it. Not supported for now. - if (!generalAsyncInliningEnabled() || impInlineRoot()->opts.IsOSR()) + if (!generalAsyncInliningEnabled()) { compInlineResult->NoteFatal(InlineObservation::CALLEE_AWAIT); return; From 110df7e89cab5766b998ea1809b17fe3bcc1a529 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 15:01:10 +0200 Subject: [PATCH 23/41] JIT: Do not generally inline async calls that may suspend inside a try region Suspending inside a protected region additionally requires catching the exception and rethrowing it around the post-inline handling, per the design document. That is not implemented, so bail out for now. Inlinees with EH are only inlined when compInlineMethodsWithEH is enabled, so this is reachable but not the common case. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/importercalls.cpp | 9 +++++++++ src/coreclr/jit/inline.def | 1 + 2 files changed, 10 insertions(+) diff --git a/src/coreclr/jit/importercalls.cpp b/src/coreclr/jit/importercalls.cpp index 7504ed1bc26be2..e10c32831267fe 100644 --- a/src/coreclr/jit/importercalls.cpp +++ b/src/coreclr/jit/importercalls.cpp @@ -7423,6 +7423,15 @@ void Compiler::impSetupAsyncCall(GenTreeCall* call, // gets its own context handling below, exactly like an await in a non-inlined // method, and the inlinee's own contexts (created by its SaveAsyncContexts) // are used rather than the caller's. + + // Suspending inside a protected region additionally requires catching and + // rethrowing around the post-inline handling, which is not implemented yet. + if ((compCurBB != nullptr) && (compCurBB->hasTryIndex() || compCurBB->hasHndIndex())) + { + compInlineResult->NoteFatal(InlineObservation::CALLEE_AWAIT_IN_TRY); + return; + } + JITDUMP("Call [%06u] is an await in an inlinee that may suspend\n", dspTreeID(call)); m_nextAsyncCallUsesOwnContexts = true; } diff --git a/src/coreclr/jit/inline.def b/src/coreclr/jit/inline.def index aaa7ce52dc82d1..e2a63f2e783616 100644 --- a/src/coreclr/jit/inline.def +++ b/src/coreclr/jit/inline.def @@ -57,6 +57,7 @@ INLINE_OBSERVATION(STFLD_NEEDS_HELPER, bool, "stfld needs helper", INLINE_OBSERVATION(TOO_MANY_ARGUMENTS, bool, "too many arguments", FATAL, CALLEE) INLINE_OBSERVATION(TOO_MANY_LOCALS, bool, "too many locals", FATAL, CALLEE) INLINE_OBSERVATION(AWAIT, bool, "has await", FATAL, CALLEE) +INLINE_OBSERVATION(AWAIT_IN_TRY, bool, "has await inside try region", FATAL, CALLEE) INLINE_OBSERVATION(ASYNC_SUSPEND, bool, "has async suspend", FATAL, CALLEE) // ------ Callee Performance ------- From 65fdb576552aa6c978598d7921b9a2ab6eb6d47d Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 15:26:59 +0200 Subject: [PATCH 24/41] JIT: Emit post-inline IR for an inlined async frame returning to its caller When an inlined async callee is resumed inside its own body, the async infrastructure only restores the context of the suspension point itself. The work it would otherwise have done when the callee's frame returned to its caller now has to be performed by the JIT: if (resumed_F) { RestoreInlinedFrameExecutionContext(continuation.ExecutionContextFor); if (!IsOnRightContext(continuation.ContinuationContextFor, continuation.FlagsFor)) { await SwitchContext(...); } resumed_caller = true; } This is emitted as a real control flow diamond at the join point all of the inlinee's returns converge on, alongside the other post-inline IR. IsOnRightContext stays visible in the IR and is an inline candidate, so it needs to be a statement root with its value consumed through a GT_RET_EXPR. Two bugs this exposed are fixed here as well: - Continuation members were registered on whichever Compiler instance was importing, so an awaiter inside an inlinee got an index into the inlinee's vector while its offset node was resolved against the root's layout. Members describe the root method's continuation, so they are now always registered there. - GT_CONTINUATION_MEMBER_OFFSET was always bashed to a TYP_INT constant. Offsets used in address arithmetic are TYP_I_IMPL, so the node's type is now preserved. Also records each inlined frame's async locals on its InlineContext, since the inlinee's Compiler is discarded once it has been spliced in. The nested capture chain at suspensions is still missing, so enabling JitAsyncInlining is not yet correct. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 32 ++-- src/coreclr/jit/compiler.h | 6 +- src/coreclr/jit/fginline.cpp | 305 +++++++++++++++++++++++++++++++++++ src/coreclr/jit/inline.h | 34 ++++ 4 files changed, 365 insertions(+), 12 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index ac6e885a3280e3..a64714d0107681 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -165,15 +165,21 @@ void ContinuationMember::Print() const size_t Compiler::GetContinuationMemberIndex(const ContinuationMember& member) { - if (m_asyncContinuationMembers == nullptr) + // Members describe the root method's continuation, so they are always registered + // there: an inlinee's IR ends up in the root and its member offset nodes are + // resolved against the root's layout. + Compiler* const root = impInlineRoot(); + + if (root->m_asyncContinuationMembers == nullptr) { - m_asyncContinuationMembers = new (this, CMK_Async) jitstd::vector(getAllocator(CMK_Async)); + root->m_asyncContinuationMembers = + new (root, CMK_Async) jitstd::vector(root->getAllocator(CMK_Async)); } else { - for (size_t i = 0; i < m_asyncContinuationMembers->size(); i++) + for (size_t i = 0; i < root->m_asyncContinuationMembers->size(); i++) { - const ContinuationMember& existingMember = m_asyncContinuationMembers->at(i); + const ContinuationMember& existingMember = root->m_asyncContinuationMembers->at(i); if (ContinuationMember::AreCompatible(member, existingMember)) { @@ -182,19 +188,21 @@ size_t Compiler::GetContinuationMemberIndex(const ContinuationMember& member) } } - m_asyncContinuationMembers->push_back(member); - return m_asyncContinuationMembers->size() - 1; + root->m_asyncContinuationMembers->push_back(member); + return root->m_asyncContinuationMembers->size() - 1; } -size_t Compiler::GetContinuationMemberCount() const +size_t Compiler::GetContinuationMemberCount() { - return m_asyncContinuationMembers == nullptr ? 0 : m_asyncContinuationMembers->size(); + Compiler* const root = impInlineRoot(); + return root->m_asyncContinuationMembers == nullptr ? 0 : root->m_asyncContinuationMembers->size(); } const ContinuationMember& Compiler::GetContinuationMember(size_t index) { - assert(index < m_asyncContinuationMembers->size()); - return m_asyncContinuationMembers->at(index); + Compiler* const root = impInlineRoot(); + assert(index < root->m_asyncContinuationMembers->size()); + return root->m_asyncContinuationMembers->at(index); } //------------------------------------------------------------------------ @@ -956,7 +964,9 @@ PhaseStatus AsyncTransformation::Run() assert(continuationLayout->ContinuationMemberOffsets[memberIndex] != UINT_MAX); ssize_t offset = (OFFSETOF__CORINFO_Continuation__data - SIZEOF__CORINFO_Object) + continuationLayout->ContinuationMemberOffsets[memberIndex]; - node->BashToConst(offset, TYP_INT); + // Preserve the node's type: offsets used in address arithmetic are TYP_I_IMPL, + // while those passed to helpers are TYP_INT. + node->BashToConst(offset, node->TypeGet()); } m_compiler->fgInvalidateDfsTree(); diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h index 66fa71dc1dad13..02aac6c033e8e6 100644 --- a/src/coreclr/jit/compiler.h +++ b/src/coreclr/jit/compiler.h @@ -6266,7 +6266,7 @@ class Compiler jitstd::vector* m_asyncContinuationMembers = nullptr; size_t GetContinuationMemberIndex(const ContinuationMember& member); - size_t GetContinuationMemberCount() const; + size_t GetContinuationMemberCount(); const ContinuationMember& GetContinuationMember(size_t index); PhaseStatus SaveAsyncContexts(); @@ -7517,6 +7517,10 @@ class Compiler void fgInsertInlineeArgument(const InlArgInfo& argInfo, BasicBlock* block, Statement** afterStmt, Statement** newStmt, const DebugInfo& callDI); Statement* fgInlinePrependStatements(InlineInfo* inlineInfo); void fgInlineAppendStatements(InlineInfo* inlineInfo, BasicBlock* block, Statement* stmt); + void fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicBlock* joinBlock); + void fgSetupAsyncFrameSwitchCall(GenTreeCall* call, InlineContext* inlineeContext, const DebugInfo& di); + void fgMarkAsyncHelperInlineCandidate(GenTreeCall* call); + GenTree* gtNewContinuationMemberIndir(const struct ContinuationMember& member, var_types type); #ifdef DEBUG static fgWalkPreFn fgDebugCheckInlineCandidates; diff --git a/src/coreclr/jit/fginline.cpp b/src/coreclr/jit/fginline.cpp index b11ab31fea2d16..d3362fbc8e8a8d 100644 --- a/src/coreclr/jit/fginline.cpp +++ b/src/coreclr/jit/fginline.cpp @@ -2,6 +2,7 @@ // The .NET Foundation licenses this file to you under the MIT license. #include "jitpch.h" +#include "async.h" #ifdef _MSC_VER #pragma hdrstop @@ -1644,6 +1645,10 @@ void Compiler::fgInsertInlineeBlocks(InlineInfo* pInlineInfo) // Append statements to null out gc ref locals, if necessary. fgInlineAppendStatements(pInlineInfo, iciBlock, stmtAfter); + + // An inlinee that does async context handling always has multiple blocks: its + // SaveAsyncContexts creates a try/fault region and a merged return block. + assert(InlineeCompiler->lvaResumedIndicator == BAD_VAR_NUM); insertInlineeBlocks = false; } else @@ -1890,6 +1895,9 @@ void Compiler::fgInsertInlineeBlocks(InlineInfo* pInlineInfo) // Append statements to null out gc ref locals, if necessary. fgInlineAppendStatements(pInlineInfo, bottomBlock, nullptr); JITDUMPEXEC(fgDispBasicBlocks(InlineeCompiler->fgFirstBB, InlineeCompiler->fgLastBB, true)); + + // Handle the inlined async frame logically returning to its caller. + fgInlineAppendAsyncFrameStatements(pInlineInfo, bottomBlock); } // @@ -2422,6 +2430,303 @@ Statement* Compiler::fgInlinePrependStatements(InlineInfo* inlineInfo) return afterStmt; } +//------------------------------------------------------------------------ +// fgSetupAsyncFrameSwitchCall: Turn a call to AsyncHelpers.SwitchContext into a proper +// async call. +// +// Arguments: +// call - the call, with its user arguments already added +// inlineeContext - inline context of the frame that is logically returning +// di - debug info to use +// +// Notes: +// The switch is an await: it always suspends, and gets resumed on the requested +// context. It runs after the inlinee's frame has logically returned, so the frame it +// belongs to is the inlinee's caller. +// +void Compiler::fgSetupAsyncFrameSwitchCall(GenTreeCall* call, InlineContext* inlineeContext, const DebugInfo& di) +{ + AsyncCallInfo asyncInfo; + // SwitchContext establishes the continuation context itself, so no further handling + // is attached to this await. + asyncInfo.ContinuationContextHandling = ContinuationContextHandling::None; + asyncInfo.AlwaysSuspends = true; + asyncInfo.CallAsyncDebugInfo = di; + call->SetIsAsync(new (this, CMK_Async) AsyncCallInfo(asyncInfo)); + + if (Target::g_tgtArgOrder == Target::ARG_ORDER_R2L) + { + call->gtArgs.PushFront(this, + NewCallArg::Primitive(gtNewNull(), TYP_REF).WellKnown(WellKnownArg::AsyncContinuation)); + } + else + { + call->gtArgs.PushBack(this, + NewCallArg::Primitive(gtNewNull(), TYP_REF).WellKnown(WellKnownArg::AsyncContinuation)); + } + + // Contexts of the frame this call belongs to, i.e. the inlinee's caller. + InlineContext* const parentContext = inlineeContext->GetParent(); + unsigned resumed; + unsigned execCtx; + unsigned syncCtx; + if ((parentContext == nullptr) || !parentContext->HasAsyncFrameLocals()) + { + resumed = lvaResumedIndicator; + execCtx = lvaAsyncExecutionContextVar; + syncCtx = lvaAsyncSynchronizationContextVar; + } + else + { + resumed = parentContext->GetAsyncResumedIndicator(); + execCtx = parentContext->GetAsyncExecutionContextVar(); + syncCtx = parentContext->GetAsyncSynchronizationContextVar(); + } + + assert((resumed != BAD_VAR_NUM) && (execCtx != BAD_VAR_NUM) && (syncCtx != BAD_VAR_NUM)); + + // TODO-Async: when this call suspends, the frames enclosing the caller need their + // context handling run too. Those come from the enclosing frame chain, which is added + // to async calls as additional pseudo-args. + call->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(syncCtx, TYP_REF)) + .WellKnown(WellKnownArg::AsyncSynchronizationContext)); + call->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(execCtx, TYP_REF)) + .WellKnown(WellKnownArg::AsyncExecutionContext)); + call->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(resumed, TYP_INT)) + .WellKnown(WellKnownArg::AsyncResumedUse)); +} + +//------------------------------------------------------------------------ +// gtNewContinuationMemberIndir: Create a load of a continuation member from the +// continuation this method was resumed with. +// +// Arguments: +// member - the continuation member to load +// type - type of the member +// +// Returns: +// A load of the member. +// +// Notes: +// The member offset is not known until the async transformation has laid out the +// continuation, so it is represented symbolically until then. The offset is relative +// to the instance data, hence the additional object header adjustment. +// +GenTree* Compiler::gtNewContinuationMemberIndir(const ContinuationMember& member, var_types type) +{ + size_t const memberIndex = GetContinuationMemberIndex(member); + + GenTree* const memberOffset = + new (this, GT_CONTINUATION_MEMBER_OFFSET) GenTreeVal(GT_CONTINUATION_MEMBER_OFFSET, TYP_I_IMPL, memberIndex); + GenTree* const headerSize = gtNewIconNode(SIZEOF__CORINFO_Object, TYP_I_IMPL); + GenTree* const offset = gtNewOperNode(GT_ADD, TYP_I_IMPL, memberOffset, headerSize); + + GenTree* const continuation = gtNewLclvNode(lvaAsyncContinuationArg, TYP_REF); + GenTree* const addr = gtNewOperNode(GT_ADD, TYP_BYREF, continuation, offset); + + GenTree* const indir = gtNewIndir(type, addr, GTF_IND_NONFAULTING); + return indir; +} + +//------------------------------------------------------------------------ +// fgMarkAsyncHelperInlineCandidate: Mark a call to an async helper as an inline +// candidate. +// +// Arguments: +// call - the call +// +void Compiler::fgMarkAsyncHelperInlineCandidate(GenTreeCall* call) +{ + CORINFO_CALL_INFO callInfo = {}; + callInfo.hMethod = call->gtCallMethHnd; + callInfo.methodFlags = info.compCompHnd->getMethodAttribs(callInfo.hMethod); + impMarkInlineCandidate(call, MAKE_METHODCONTEXT(callInfo.hMethod), &callInfo, compInlineContext); +} + +//------------------------------------------------------------------------ +// fgInlineAppendAsyncFrameStatements: Emit the IR that handles an inlined async +// frame logically returning to its caller. +// +// Arguments: +// inlineInfo - information about the inline +// joinBlock - the block all of the inlinee's returns converge on +// +// Notes: +// When an inlined async callee is resumed inside its own body, the async +// infrastructure only restored the context of the suspension point itself. The work +// it would otherwise have done when the callee's frame returned to its caller has to +// be performed here instead: +// +// if (resumed_F) +// { +// RestoreInlinedFrameExecutionContext(continuation.ExecutionContextFor); +// if (!IsOnRightContext(continuation.ContinuationContextFor, +// continuation.FlagsFor)) +// { +// await SwitchContext(continuation.ContinuationContextFor, +// continuation.FlagsFor); +// } +// resumed_caller = true; +// } +// +// The members are read off the continuation this method was resumed with. That is +// only valid when resumed_F is true, which is exactly the guard: resumed_F can only +// become true at a resumption point belonging to this inlinee's own async calls, and +// those belong to this method's continuation layout. +// +void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicBlock* joinBlock) +{ + if (!generalAsyncInliningEnabled()) + { + return; + } + + if (InlineeCompiler->lvaResumedIndicator == BAD_VAR_NUM) + { + // Inlinee does not do context handling, so it inherited everything from this + // call and there is no logical frame transition to handle. + return; + } + + InlineContext* const inlineContext = inlineInfo->inlineContext; + inlineContext->SetAsyncFrameLocals(InlineeCompiler->lvaResumedIndicator, + InlineeCompiler->lvaAsyncExecutionContextVar, + InlineeCompiler->lvaAsyncSynchronizationContextVar); + + if (!InlineeCompiler->compAsyncBodyMaySuspend) + { + // The inlinee has no suspension point, so its resumed indicator is provably + // always false and all of the below would be dead code. + JITDUMP("Inlinee cannot suspend; no async frame transition IR needed\n"); + return; + } + + // Locals of the frame we are logically returning into. + InlineContext* const parentContext = inlineContext->GetParent(); + unsigned const resumedCaller = ((parentContext == nullptr) || !parentContext->HasAsyncFrameLocals()) + ? lvaResumedIndicator + : parentContext->GetAsyncResumedIndicator(); + assert(resumedCaller != BAD_VAR_NUM); + + unsigned const resumedInlinee = InlineeCompiler->lvaResumedIndicator; + unsigned const inlineDepth = inlineContext->GetDepth(); + + JITDUMP("Adding async frame transition IR for inline depth %u: resumed V%02u -> V%02u\n", inlineDepth, + resumedInlinee, resumedCaller); + + const DebugInfo& di = inlineInfo->iciStmt->GetDebugInfo(); + + // Split so that the transition IR precedes the remainder of the caller's code. + BasicBlock* const restBlock = fgSplitBlockAtBeginning(joinBlock); + + // joinBlock is now empty and jumps to restBlock. Turn it into the "did we resume + // inside the inlinee?" test. + BasicBlock* const restoreBlock = fgNewBBafter(BBJ_ALWAYS, joinBlock, /* extendRegion */ true); + BasicBlock* const switchBlock = fgNewBBafter(BBJ_ALWAYS, restoreBlock, /* extendRegion */ true); + BasicBlock* const doneBlock = fgNewBBafter(BBJ_ALWAYS, switchBlock, /* extendRegion */ true); + + restoreBlock->inheritWeightPercentage(joinBlock, 0); + switchBlock->inheritWeightPercentage(joinBlock, 0); + doneBlock->inheritWeightPercentage(joinBlock, 0); + + // joinBlock: if (resumed_F == 0) goto restBlock; else goto restoreBlock + { + GenTree* const resumed = gtNewLclvNode(resumedInlinee, TYP_INT); + GenTree* const isZero = gtNewOperNode(GT_EQ, TYP_INT, resumed, gtNewIconNode(0)); + GenTree* const jtrue = gtNewOperNode(GT_JTRUE, TYP_VOID, isZero); + fgInsertStmtAtEnd(joinBlock, gtNewStmt(jtrue)); + + fgRemoveRefPred(joinBlock->GetTargetEdge()); + FlowEdge* const toRest = fgAddRefPred(restBlock, joinBlock); + FlowEdge* const toRestore = fgAddRefPred(restoreBlock, joinBlock); + joinBlock->SetCond(toRest, toRestore); + // Resuming inside the inlinee is the cold path. + toRest->setLikelihood(1.0); + toRestore->setLikelihood(0.0); + } + + CORINFO_ASYNC_INFO* const asyncInfo = eeGetAsyncInfo(); + + // restoreBlock: restore the ExecutionContext, then check the continuation context. + { + GenTreeCall* const restoreCall = + gtNewUserCallNode(asyncInfo->restoreInlinedFrameExecutionContextMethHnd, TYP_VOID); + restoreCall->gtArgs.PushFront(this, + NewCallArg::Primitive( + gtNewContinuationMemberIndir(ContinuationMember::InlineFrameExecutionContext( + inlineDepth), + TYP_REF))); + fgMarkAsyncHelperInlineCandidate(restoreCall); + fgInsertStmtAtEnd(restoreBlock, gtNewStmt(restoreCall)); + + GenTreeCall* const isOnRightContextCall = gtNewUserCallNode(asyncInfo->isOnRightContextMethHnd, TYP_INT); + isOnRightContextCall->gtArgs + .PushFront(this, + NewCallArg::Primitive( + gtNewContinuationMemberIndir(ContinuationMember::InlineFrameFlags(inlineDepth), TYP_INT))); + isOnRightContextCall->gtArgs + .PushFront(this, + NewCallArg::Primitive( + gtNewContinuationMemberIndir(ContinuationMember::InlineFrameContinuationContext(inlineDepth), + TYP_REF))); + fgMarkAsyncHelperInlineCandidate(isOnRightContextCall); + + // Inline candidates must be statement roots; the value is consumed through a + // GT_RET_EXPR placeholder that the inliner substitutes. + GenTree* isOnRightContext; + if (isOnRightContextCall->IsInlineCandidate()) + { + fgInsertStmtAtEnd(restoreBlock, gtNewStmt(isOnRightContextCall)); + isOnRightContext = gtNewInlineCandidateReturnExpr(isOnRightContextCall, TYP_INT); + } + else + { + isOnRightContext = isOnRightContextCall; + } + + GenTree* const isWrongContext = gtNewOperNode(GT_EQ, TYP_INT, isOnRightContext, gtNewIconNode(0)); + GenTree* const jtrue = gtNewOperNode(GT_JTRUE, TYP_VOID, isWrongContext); + fgInsertStmtAtEnd(restoreBlock, gtNewStmt(jtrue)); + + FlowEdge* const toSwitch = fgAddRefPred(switchBlock, restoreBlock); + FlowEdge* const toDone = fgAddRefPred(doneBlock, restoreBlock); + restoreBlock->SetCond(toSwitch, toDone); + // We expect to already be on the right context in almost all cases; if we were + // not, inlining would be unlikely to be profitable in the first place. + toSwitch->setLikelihood(0.0); + toDone->setLikelihood(1.0); + } + + // switchBlock: get back onto the continuation context the caller's continuation + // captured. This is an await, so it may suspend. + { + GenTreeCall* const switchCall = gtNewUserCallNode(asyncInfo->switchContextMethHnd, TYP_VOID); + switchCall->gtArgs.PushFront(this, + NewCallArg::Primitive( + gtNewContinuationMemberIndir(ContinuationMember::InlineFrameFlags(inlineDepth), + TYP_INT))); + switchCall->gtArgs + .PushFront(this, + NewCallArg::Primitive( + gtNewContinuationMemberIndir(ContinuationMember::InlineFrameContinuationContext(inlineDepth), + TYP_REF))); + fgSetupAsyncFrameSwitchCall(switchCall, inlineContext, di); + fgInsertStmtAtEnd(switchBlock, gtNewStmt(switchCall)); + + switchBlock->SetKindAndTargetEdge(BBJ_ALWAYS, fgAddRefPred(doneBlock, switchBlock)); + } + + // doneBlock: the caller's frame has now observed a resumption too. + { + GenTree* const store = gtNewStoreLclVarNode(resumedCaller, gtNewIconNode(1)); + fgInsertStmtAtEnd(doneBlock, gtNewStmt(store)); + + doneBlock->SetKindAndTargetEdge(BBJ_ALWAYS, fgAddRefPred(restBlock, doneBlock)); + } + + JITDUMPEXEC(fgDispBasicBlocks(joinBlock, restBlock, true)); +} + //------------------------------------------------------------------------ // fgInlineAppendStatements: Append statements that are needed // after the inlined call. diff --git a/src/coreclr/jit/inline.h b/src/coreclr/jit/inline.h index 68d39e37258c60..f2a0df1b828613 100644 --- a/src/coreclr/jit/inline.h +++ b/src/coreclr/jit/inline.h @@ -911,6 +911,36 @@ class InlineContext return (m_PgoInfo.PgoSchema != nullptr) && (m_PgoInfo.PgoSchemaCount > 0) && (m_PgoInfo.PgoData != nullptr); } + // The async context locals of the inlined frame this context represents. Recorded when + // the inlinee is spliced in, since the inlinee's Compiler is discarded afterwards. + // Used to build the logical frame transitions for general runtime async inlining. + void SetAsyncFrameLocals(unsigned resumedIndicator, unsigned executionContextVar, unsigned syncContextVar) + { + m_asyncResumedIndicator = resumedIndicator; + m_asyncExecutionContextVar = executionContextVar; + m_asyncSynchronizationContextVar = syncContextVar; + } + + bool HasAsyncFrameLocals() const + { + return m_asyncResumedIndicator != BAD_VAR_NUM; + } + + unsigned GetAsyncResumedIndicator() const + { + return m_asyncResumedIndicator; + } + + unsigned GetAsyncExecutionContextVar() const + { + return m_asyncExecutionContextVar; + } + + unsigned GetAsyncSynchronizationContextVar() const + { + return m_asyncSynchronizationContextVar; + } + private: InlineContext(InlineStrategy* strategy); @@ -931,6 +961,10 @@ class InlineContext unsigned m_Ordinal; // Ordinal number of this inline bool m_Success : 1; // true if this was a successful inline + unsigned m_asyncResumedIndicator = BAD_VAR_NUM; + unsigned m_asyncExecutionContextVar = BAD_VAR_NUM; + unsigned m_asyncSynchronizationContextVar = BAD_VAR_NUM; + #if defined(DEBUG) InlinePolicy* m_Policy; // policy that evaluated this inline From 117822b7801e30e4d48ccbbb1456aecea3e7652b Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 15:40:18 +0200 Subject: [PATCH 25/41] JIT: Carry enclosing inlined frames' context values on async calls When an async call inside an inlined frame suspends, every enclosing frame that has not yet resumed needs its contexts captured into the continuation and restored onto the thread, as if its physical frame had returned. The IR for that is created by the async transformation, long after the optimizer has run, so those values have to be kept live and GC reported at the call until then. That is exactly what the existing per-frame context args do, so the enclosing frames just add more of them: one set of AsyncResumedUse/AsyncExecutionContext/ AsyncSynchronizationContext per frame, appended outward when the inlinee is spliced in. These are pseudo-args, so duplicates take no registers and are expanded out later. The innermost frame stays first. Only uses are needed for the enclosing frames. AsyncResumedDef stays innermost-only, since the caller's indicator is set by the post-inline IR rather than by a resumption point. Adjusts the two places that assumed these args were unique: IsReusableSuspension now requires all of them to agree, and HandleReusedSuspension removes all of them. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 58 ++++++++++++--- src/coreclr/jit/compiler.h | 1 + src/coreclr/jit/fginline.cpp | 133 +++++++++++++++++++++++++++++++++++ 3 files changed, 182 insertions(+), 10 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index a64714d0107681..8f33ebdbbe328e 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -163,6 +163,31 @@ void ContinuationMember::Print() const } #endif +//------------------------------------------------------------------------ +// CollectWellKnownArgs: +// Collect the nodes of all arguments of a call with a specific well-known kind, in +// argument order. +// +// Parameters: +// call - The call +// wka - The well-known argument kind +// nodes - [out] Stack to push the argument nodes onto +// +// Remarks: +// Async context pseudo-args can appear multiple times on a call: general async +// inlining adds one set per enclosing inlined frame. +// +static void CollectWellKnownArgs(GenTreeCall* call, WellKnownArg wka, ArrayStack& nodes) +{ + for (CallArg& arg : call->gtArgs.Args()) + { + if (arg.GetWellKnownArg() == wka) + { + nodes.Push(arg.GetNode()); + } + } +} + size_t Compiler::GetContinuationMemberIndex(const ContinuationMember& member) { // Members describe the root method's continuation, so they are always registered @@ -2208,35 +2233,46 @@ bool AsyncTransformation::IsReusableSuspension(const AsyncState* state, } } + // These can be duplicated: general async inlining adds one set of context args per + // enclosing inlined frame, describing the chain of frames whose handling a suspension + // has to run. All of them must agree for a suspension to be reusable. static const WellKnownArg validateArgs[] = {WellKnownArg::AsyncAwaiter, WellKnownArg::AsyncResumedUse, WellKnownArg::AsyncResumedDef, WellKnownArg::AsyncExecutionContext, WellKnownArg::AsyncSynchronizationContext}; + ArrayStack thisNodes(m_compiler->getAllocator(CMK_Async)); + ArrayStack otherNodes(m_compiler->getAllocator(CMK_Async)); for (WellKnownArg arg : validateArgs) { - CallArg* thisArg = call->gtArgs.FindWellKnownArg(arg); - CallArg* otherArg = predAsyncCall->gtArgs.FindWellKnownArg(arg); + thisNodes.Reset(); + otherNodes.Reset(); + CollectWellKnownArgs(call, arg, thisNodes); + CollectWellKnownArgs(predAsyncCall, arg, otherNodes); - if ((thisArg == nullptr) != (otherArg == nullptr)) + if (thisNodes.Height() != otherNodes.Height()) { - JITDUMP(" No; disagreement on presence of %s argument\n", getWellKnownArgName(arg)); + JITDUMP(" No; disagreement on number of %s arguments (%d vs %d)\n", getWellKnownArgName(arg), + thisNodes.Height(), otherNodes.Height()); return false; } - if (thisArg != nullptr) + for (int i = 0; i < thisNodes.Height(); i++) { + GenTree* thisNode = thisNodes.Bottom(i); + GenTree* otherNode = otherNodes.Bottom(i); + // The value may have been folded to a constant (e.g. when the // indicator is provably zero at this point), which is still fine to // compare and to remove from the call. - if (!thisArg->GetNode()->OperIsAnyLocal() && !thisArg->GetNode()->IsInvariant()) + if (!thisNode->OperIsAnyLocal() && !thisNode->IsInvariant()) { JITDUMP(" No; %s argument is too complex\n", getWellKnownArgName(arg)); return false; } - if (!GenTree::Compare(thisArg->GetNode(), otherArg->GetNode())) + if (!GenTree::Compare(thisNode, otherNode)) { JITDUMP(" No; disagreement on value of %s argument ([%06u] vs [%06u])\n", getWellKnownArgName(arg), - Compiler::dspTreeID(thisArg->GetNode()), Compiler::dspTreeID(otherArg->GetNode())); + Compiler::dspTreeID(thisNode), Compiler::dspTreeID(otherNode)); return false; } } @@ -2269,8 +2305,10 @@ void AsyncTransformation::HandleReusedSuspension(BasicBlock* callBlock, GenTreeC WellKnownArg::AsyncSynchronizationContext}; for (WellKnownArg wka : argsToRemove) { - CallArg* arg = call->gtArgs.FindWellKnownArg(wka); - if (arg != nullptr) + // These can be duplicated: general async inlining adds one set of context args + // per enclosing inlined frame. All of them must go. + for (CallArg* arg = call->gtArgs.FindWellKnownArg(wka); arg != nullptr; + arg = call->gtArgs.FindWellKnownArg(wka)) { assert(arg->GetNode()->OperIsAnyLocal() || arg->GetNode()->IsInvariant()); LIR::AsRange(callBlock).Remove(arg->GetNode()); diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h index 02aac6c033e8e6..752b1540814cd6 100644 --- a/src/coreclr/jit/compiler.h +++ b/src/coreclr/jit/compiler.h @@ -7518,6 +7518,7 @@ class Compiler Statement* fgInlinePrependStatements(InlineInfo* inlineInfo); void fgInlineAppendStatements(InlineInfo* inlineInfo, BasicBlock* block, Statement* stmt); void fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicBlock* joinBlock); + void fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo); void fgSetupAsyncFrameSwitchCall(GenTreeCall* call, InlineContext* inlineeContext, const DebugInfo& di); void fgMarkAsyncHelperInlineCandidate(GenTreeCall* call); GenTree* gtNewContinuationMemberIndir(const struct ContinuationMember& member, var_types type); diff --git a/src/coreclr/jit/fginline.cpp b/src/coreclr/jit/fginline.cpp index d3362fbc8e8a8d..1b03a8e1a4589c 100644 --- a/src/coreclr/jit/fginline.cpp +++ b/src/coreclr/jit/fginline.cpp @@ -1897,6 +1897,7 @@ void Compiler::fgInsertInlineeBlocks(InlineInfo* pInlineInfo) JITDUMPEXEC(fgDispBasicBlocks(InlineeCompiler->fgFirstBB, InlineeCompiler->fgLastBB, true)); // Handle the inlined async frame logically returning to its caller. + fgAppendEnclosingAsyncFrameContextArgs(pInlineInfo); fgInlineAppendAsyncFrameStatements(pInlineInfo, bottomBlock); } @@ -2430,6 +2431,138 @@ Statement* Compiler::fgInlinePrependStatements(InlineInfo* inlineInfo) return afterStmt; } +//------------------------------------------------------------------------ +// fgAppendEnclosingAsyncFrameContextArgs: Add the context values of the frames +// enclosing an inlined async frame to the async calls inside it. +// +// Arguments: +// inlineInfo - information about the inline +// +// Notes: +// When one of these calls suspends, every enclosing frame that has not yet resumed +// must have its contexts captured into the continuation and restored onto the thread, +// as if its physical frame had returned. The IR for that is created by the async +// transformation, long after the optimizer has run, so the values have to be kept +// live and GC reported at the call until then. That is what the existing per-frame +// context args do, so the enclosing frames simply add more of them. +// +// These are pseudo-args, so duplicates are fine: they take no registers and are +// expanded out by the async transformation. The innermost frame stays first, since +// the inlinee pushed its own args to the front. +// +void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) +{ + if (!generalAsyncInliningEnabled()) + { + return; + } + + if ((InlineeCompiler->lvaResumedIndicator == BAD_VAR_NUM) || !InlineeCompiler->compAsyncBodyMaySuspend) + { + return; + } + + struct AsyncFrameLocals + { + unsigned Resumed; + unsigned ExecutionContext; + unsigned SynchronizationContext; + }; + + ArrayStack frames(getAllocator(CMK_Async)); + for (InlineContext* ctx = inlineInfo->inlineContext->GetParent(); ctx != nullptr; ctx = ctx->GetParent()) + { + if (ctx->HasAsyncFrameLocals()) + { + frames.Push({ctx->GetAsyncResumedIndicator(), ctx->GetAsyncExecutionContextVar(), + ctx->GetAsyncSynchronizationContextVar()}); + } + } + + // The root method's frame is always the outermost one. + frames.Push({lvaResumedIndicator, lvaAsyncExecutionContextVar, lvaAsyncSynchronizationContextVar}); + + unsigned const inlineeResumed = InlineeCompiler->lvaResumedIndicator; + + struct Visitor : GenTreeVisitor + { + enum + { + DoPreOrder = true, + }; + + ArrayStack& m_frames; + unsigned m_inlineeResumed; + + Visitor(Compiler* comp, ArrayStack& frames, unsigned inlineeResumed) + : GenTreeVisitor(comp) + , m_frames(frames) + , m_inlineeResumed(inlineeResumed) + { + } + + fgWalkResult PreOrderVisit(GenTree** use, GenTree* user) + { + GenTree* tree = *use; + if ((tree->gtFlags & GTF_CALL) == 0) + { + return WALK_SKIP_SUBTREES; + } + + if (!tree->IsCall() || !tree->AsCall()->IsAsync()) + { + return WALK_CONTINUE; + } + + GenTreeCall* const call = tree->AsCall(); + CallArg* const ownArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse); + if ((ownArg == nullptr) || !ownArg->GetNode()->OperIs(GT_LCL_VAR) || + (ownArg->GetNode()->AsLclVar()->GetLclNum() != m_inlineeResumed)) + { + // Either no context handling, or the call inherited its contexts from the + // inlining call, in which case there is no frame transition. + return WALK_CONTINUE; + } + + for (int i = 0; i < m_frames.Height(); i++) + { + const AsyncFrameLocals& frame = m_frames.Bottom(i); + assert((frame.Resumed != BAD_VAR_NUM) && (frame.ExecutionContext != BAD_VAR_NUM) && + (frame.SynchronizationContext != BAD_VAR_NUM)); + + call->gtArgs.PushBack(m_compiler, + NewCallArg::Primitive(m_compiler->gtNewLclVarNode(frame.Resumed, TYP_INT)) + .WellKnown(WellKnownArg::AsyncResumedUse)); + call->gtArgs.PushBack(m_compiler, NewCallArg::Primitive( + m_compiler->gtNewLclVarNode(frame.ExecutionContext, TYP_REF)) + .WellKnown(WellKnownArg::AsyncExecutionContext)); + call->gtArgs.PushBack(m_compiler, + NewCallArg::Primitive( + m_compiler->gtNewLclVarNode(frame.SynchronizationContext, TYP_REF)) + .WellKnown(WellKnownArg::AsyncSynchronizationContext)); + } + + JITDUMP("Added %d enclosing frame context arg sets to async call [%06u]\n", m_frames.Height(), + Compiler::dspTreeID(call)); + return WALK_CONTINUE; + } + }; + + Visitor visitor(this, frames, inlineeResumed); + for (BasicBlock* block = InlineeCompiler->fgFirstBB; block != nullptr; block = block->Next()) + { + for (Statement* const stmt : block->Statements()) + { + visitor.WalkTree(stmt->GetRootNodePointer(), nullptr); + } + + if (block == InlineeCompiler->fgLastBB) + { + break; + } + } +} + //------------------------------------------------------------------------ // fgSetupAsyncFrameSwitchCall: Turn a call to AsyncHelpers.SwitchContext into a proper // async call. From 810302550b4d137e292cc70b01e96253ff5859f3 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 16:53:09 +0200 Subject: [PATCH 26/41] JIT: Run enclosing inlined frames' context handling on suspension Completes general runtime async inlining. When an async call inside an inlined frame suspends, every enclosing frame that has not yet resumed must capture the contexts it would hand to its caller and restore its caller's contexts onto the thread, as if its physical frame had returned. A frame having resumed implies its caller has too, so the frames can be walked outward as a straight line rather than as the nested conditionals of the design document: once one frame has resumed, the helpers for it and every frame outside it no-op. Each suspension jumps to a tail block that does this, shared by all suspensions in the same frame. Adds AsyncHelpers.CaptureInlinedFrameTransition for the capture, which no-ops when the frame has already resumed. The chain of frames is recorded on AsyncCallInfo rather than being recovered from the call's context args, since optimizations may rewrite the arg nodes; the args remain purely to keep the values live and GC reported until here. Also fixes impInheritAsyncContextsFromInliner, which propagated only the innermost frame's contexts to an inlined callee's async calls. An async call carrying a frame chain can itself be an inline candidate, and dropping the rest of the chain silently lost the handling for every frame outside the immediate one. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../CompilerServices/AsyncHelpers.CoreCLR.cs | 24 +++ src/coreclr/inc/corinfo.h | 3 + src/coreclr/inc/jiteeversionguid.h | 10 +- src/coreclr/jit/async.cpp | 201 +++++++++++++++++- src/coreclr/jit/async.h | 10 + src/coreclr/jit/fginline.cpp | 50 +++-- src/coreclr/jit/gentree.h | 32 +++ src/coreclr/jit/importer.cpp | 9 +- src/coreclr/jit/importercalls.cpp | 36 ++++ src/coreclr/jit/inline.h | 17 +- .../tools/Common/JitInterface/CorInfoImpl.cs | 1 + .../tools/Common/JitInterface/CorInfoTypes.cs | 1 + .../tools/superpmi/superpmi-shared/agnostic.h | 1 + .../superpmi-shared/methodcontext.cpp | 2 + src/coreclr/vm/corelib.h | 1 + src/coreclr/vm/jitinterface.cpp | 1 + 16 files changed, 371 insertions(+), 28 deletions(-) diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index c13d40edf1dc19..835880f39325b8 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -1580,6 +1580,30 @@ private static unsafe void SwitchContext(object? continuationContext, Continuati AsyncSuspend(taskCont); } + // Capture the contexts an inlined async frame hands to its caller when it logically + // returns during a suspension, i.e. what the caller's continuation would have captured + // had the callee's frame been physically present. + // + // No-ops when the frame has already resumed: in that case the caller's continuation + // already exists and keeps the values it captured when the frame first suspended. + // + // The suspension walks the inlined frames outward, and a frame having resumed implies + // its caller has too, so these can be emitted as a straight line: once one frame has + // resumed, this and every subsequent restore for the frames outside it no-op. + private static void CaptureInlinedFrameTransition(bool resumed, + ref object? continuationContext, + ref ContinuationFlags flags, + ref ExecutionContext? execContext) + { + if (resumed) + { + return; + } + + CaptureContinuationContext(ref continuationContext, ref flags); + execContext = CaptureExecutionContext(); + } + // Finish suspension in the common case of a custom await or for a ConfigureAwait(false) task await: // - Capture current ExecutionContext into the continuation // - Restore ExecutionContext and SynchronizationContext to the current Thread object diff --git a/src/coreclr/inc/corinfo.h b/src/coreclr/inc/corinfo.h index 04c30fa2a97c3e..c21513b1d1b076 100644 --- a/src/coreclr/inc/corinfo.h +++ b/src/coreclr/inc/corinfo.h @@ -1843,6 +1843,9 @@ struct CORINFO_ASYNC_INFO // Method handle for AsyncHelpers.SwitchContext, used to switch to the continuation context of an // inlined async callee's caller when IsOnRightContext reports a mismatch CORINFO_METHOD_HANDLE switchContextMethHnd; + // Method handle for AsyncHelpers.CaptureInlinedFrameTransition, used on suspension to capture + // the contexts each inlined async frame hands to its caller + CORINFO_METHOD_HANDLE captureInlinedFrameTransitionMethHnd; // Finish suspension without saving continuation context (i.e. custom awaiter or ConfigureAwait(false)) CORINFO_METHOD_HANDLE finishSuspensionNoContinuationContextMethHnd; // Finish suspension with saving continuation context (i.e. normal task await) diff --git a/src/coreclr/inc/jiteeversionguid.h b/src/coreclr/inc/jiteeversionguid.h index 7aa7cbf94d576a..f1cc0272b69ff3 100644 --- a/src/coreclr/inc/jiteeversionguid.h +++ b/src/coreclr/inc/jiteeversionguid.h @@ -37,11 +37,11 @@ #include -constexpr GUID JITEEVersionIdentifier = { /* 2f34f847-03c5-43b5-9261-e526a9bcedc9 */ - 0x2f34f847, - 0x03c5, - 0x43b5, - {0x92, 0x61, 0xe5, 0x26, 0xa9, 0xbc, 0xed, 0xc9} +constexpr GUID JITEEVersionIdentifier = { /* ec3a33a7-a176-461b-9520-3b64a122bb39 */ + 0xec3a33a7, + 0xa176, + 0x461b, + {0x95, 0x20, 0x3b, 0x64, 0xa1, 0x22, 0xbb, 0x39} }; #endif // JIT_EE_VERSIONING_GUID_H diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 8f33ebdbbe328e..590976d4b36483 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -2923,7 +2923,12 @@ void AsyncTransformation::FinishContextHandlingAndSuspension(BasicBlock* assert(suspendBB->KindIs(BBJ_RETURN)); - if (m_sharedReturnBB != nullptr) + BasicBlock* const frameTail = CreateInlinedFrameSuspensionTail(callBlock, call, layout); + if (frameTail != nullptr) + { + suspendBB->SetKindAndTargetEdge(BBJ_ALWAYS, m_compiler->fgAddRefPred(frameTail, suspendBB)); + } + else if (m_sharedReturnBB != nullptr) { suspendBB->SetKindAndTargetEdge(BBJ_ALWAYS, m_compiler->fgAddRefPred(m_sharedReturnBB, suspendBB)); } @@ -3008,7 +3013,12 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc LIR::AsRange(callBlock).Remove(syncContext); call->gtArgs.RemoveUnsafe(syncContextArg); - if (sharedFinish != nullptr) + // A call inside an inlined async frame needs the enclosing frames' handling to run + // after this one, so it cannot use a finish block shared with suspensions in other + // frames. + BasicBlock* const frameTail = CreateInlinedFrameSuspensionTail(callBlock, call, layout); + + if ((sharedFinish != nullptr) && (frameTail == nullptr)) { // Store the vars to the shared locals that the shared finish block will take them from. if (m_sharedFinishContextHandlingResumedVar != BAD_VAR_NUM) @@ -3039,8 +3049,13 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc // Otherwise insert a new call InsertFinishContextHandlingCall(suspendBB, layout, helper, resumed, execContext, syncContext); - // And return either via a new GT_RETURN_SUSPEND or via the shared return BB. - if (m_sharedReturnBB != nullptr) + // And continue with the enclosing frames, or return either via a new + // GT_RETURN_SUSPEND or via the shared return BB. + if (frameTail != nullptr) + { + suspendBB->SetKindAndTargetEdge(BBJ_ALWAYS, m_compiler->fgAddRefPred(frameTail, suspendBB)); + } + else if (m_sharedReturnBB != nullptr) { suspendBB->SetKindAndTargetEdge(BBJ_ALWAYS, m_compiler->fgAddRefPred(m_sharedReturnBB, suspendBB)); } @@ -3053,6 +3068,184 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc } } +//------------------------------------------------------------------------ +// AsyncTransformation::CreateInlinedFrameSuspensionTail: +// Create, or reuse, the block that runs the context handling of the frames enclosing +// an inlined async frame on suspension. +// +// Parameters: +// callBlock - The block containing the async call +// call - The async call +// layout - Information about the continuation layout +// +// Returns: +// The block to continue suspension in, or nullptr if this call is not inside an +// inlined async frame. +// +// Remarks: +// Each frame that has not yet resumed must capture the contexts it would hand to its +// caller, and restore its caller's contexts onto the thread, as if its physical frame +// had returned. A frame having resumed implies its caller has too, so the frames can be +// walked outward as a straight line: once one frame has resumed, the helpers for it and +// every frame outside it no-op. +// +// The result depends only on the chain of frames, so it is shared by all suspensions in +// the same frame. +// +BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* callBlock, + GenTreeCall* call, + const ContinuationLayout& layout) +{ + // Remove the context args of the enclosing frames. They exist only to keep the values + // live and reported up to here; which locals they are is recorded in the call info, + // since optimizations may have rewritten the arg nodes themselves. + while (true) + { + CallArg* resumedArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse); + CallArg* execContextArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncExecutionContext); + CallArg* syncContextArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncSynchronizationContext); + if (resumedArg == nullptr) + { + assert((execContextArg == nullptr) && (syncContextArg == nullptr)); + break; + } + + assert((execContextArg != nullptr) && (syncContextArg != nullptr)); + + for (CallArg* arg : {resumedArg, execContextArg, syncContextArg}) + { + GenTree* node = arg->GetNode(); + if (LIR::AsRange(callBlock).Contains(node)) + { + LIR::AsRange(callBlock).Remove(node); + } + + call->gtArgs.RemoveUnsafe(arg); + } + } + + const AsyncFrameLocals* const frameLocals = call->GetAsyncInfo().InlineFrameLocals; + if (frameLocals == nullptr) + { + JITDUMP(" Call [%06u] is not inside an inlined async frame\n", Compiler::dspTreeID(call)); + return nullptr; + } + + unsigned const innermostDepth = call->GetAsyncInfo().InlineFrameDepth; + unsigned const numFrames = innermostDepth + 1; + assert(numFrames >= 2); + + JITDUMP(" Call [%06u] is inside an inlined async frame; %u frames in chain\n", Compiler::dspTreeID(call), + numFrames); + + unsigned const key = frameLocals[0].Resumed; + if (m_inlinedFrameTails == nullptr) + { + m_inlinedFrameTails = new (m_compiler, CMK_Async) InlinedFrameTailMap(m_compiler->getAllocator(CMK_Async)); + } + else + { + BasicBlock* existing; + if (m_inlinedFrameTails->Lookup(key, &existing)) + { + JITDUMP(" Reusing inlined frame suspension tail " FMT_BB "\n", existing->bbNum); + return existing; + } + } + + if (m_sharedReturnBB == nullptr) + { + CreateSharedReturnBB(); + } + + BasicBlock* tailBB = m_compiler->fgNewBBbefore(BBJ_ALWAYS, m_sharedReturnBB, false); + tailBB->bbSetRunRarely(); + tailBB->clearTryIndex(); + tailBB->clearHndIndex(); + tailBB->SetKindAndTargetEdge(BBJ_ALWAYS, m_compiler->fgAddRefPred(m_sharedReturnBB, tailBB)); + + if (m_compiler->fgIsUsingProfileWeights()) + { + tailBB->SetFlags(BBF_PROF_WEIGHT); + } + + JITDUMP(" Created inlined frame suspension tail " FMT_BB " for %u frames\n", tailBB->bbNum, numFrames); + + for (unsigned i = 0; i + 1 < numFrames; i++) + { + const AsyncFrameLocals& frame = frameLocals[i]; + const AsyncFrameLocals& outer = frameLocals[i + 1]; + unsigned const depth = innermostDepth - i; + + // Capture what this frame hands to its caller. + GenTreeCall* captureCall = + m_compiler->gtNewUserCallNode(m_asyncInfo->captureInlinedFrameTransitionMethHnd, TYP_VOID); + captureCall->gtArgs + .PushFront(m_compiler, + NewCallArg::Primitive( + ContinuationMemberAddress(layout, ContinuationMember::InlineFrameExecutionContext(depth)))); + captureCall->gtArgs.PushFront(m_compiler, + NewCallArg::Primitive( + ContinuationMemberAddress(layout, + ContinuationMember::InlineFrameFlags(depth)))); + captureCall->gtArgs.PushFront(m_compiler, + NewCallArg::Primitive( + ContinuationMemberAddress(layout, + ContinuationMember::InlineFrameContinuationContext( + depth)))); + captureCall->gtArgs.PushFront(m_compiler, + NewCallArg::Primitive(m_compiler->gtNewLclvNode(frame.Resumed, TYP_INT))); + + m_compiler->compCurBB = tailBB; + m_compiler->fgMorphTree(captureCall); + LIR::AsRange(tailBB).InsertAtEnd(LIR::SeqTree(m_compiler, captureCall)); + + // Then restore the caller's contexts, as its physical frame's return would have. + GenTreeCall* restoreCall = + m_compiler->gtNewUserCallNode(m_asyncInfo->restoreContextsOnSuspensionMethHnd, TYP_VOID); + restoreCall->gtArgs.PushFront(m_compiler, + NewCallArg::Primitive( + m_compiler->gtNewLclvNode(outer.SynchronizationContext, TYP_REF))); + restoreCall->gtArgs.PushFront(m_compiler, NewCallArg::Primitive( + m_compiler->gtNewLclvNode(outer.ExecutionContext, TYP_REF))); + restoreCall->gtArgs.PushFront(m_compiler, + NewCallArg::Primitive(m_compiler->gtNewLclvNode(outer.Resumed, TYP_INT))); + + m_compiler->compCurBB = tailBB; + m_compiler->fgMorphTree(restoreCall); + LIR::AsRange(tailBB).InsertAtEnd(LIR::SeqTree(m_compiler, restoreCall)); + } + + DISPRANGE(LIR::AsRange(tailBB)); + + m_inlinedFrameTails->Set(key, tailBB); + return tailBB; +} + +//------------------------------------------------------------------------ +// AsyncTransformation::ContinuationMemberAddress: +// Create the address of a continuation member in the newly created continuation. +// +// Parameters: +// layout - Information about the continuation layout +// member - The member +// +// Returns: +// IR node computing the address. +// +GenTree* AsyncTransformation::ContinuationMemberAddress(const ContinuationLayout& layout, + const ContinuationMember& member) +{ + size_t const memberIndex = m_compiler->GetContinuationMemberIndex(member); + assert(memberIndex < layout.ContinuationMemberOffsets.size()); + assert(layout.ContinuationMemberOffsets[memberIndex] != UINT_MAX); + + unsigned const offset = OFFSETOF__CORINFO_Continuation__data + layout.ContinuationMemberOffsets[memberIndex]; + GenTree* const continuation = m_compiler->gtNewLclvNode(GetNewContinuationVar(), TYP_REF); + return m_compiler->gtNewOperNode(GT_ADD, TYP_BYREF, continuation, + m_compiler->gtNewIconNode((ssize_t)offset, TYP_I_IMPL)); +} + //------------------------------------------------------------------------ // AsyncTransformation::RestoreContexts: // Create IR to restore contexts on suspension. diff --git a/src/coreclr/jit/async.h b/src/coreclr/jit/async.h index f001af4cb8238a..45a9dd5648d286 100644 --- a/src/coreclr/jit/async.h +++ b/src/coreclr/jit/async.h @@ -560,6 +560,16 @@ class AsyncTransformation const ContinuationLayout* CreateResumptionsAndSuspensions(ArrayStack& continuationMemberOffsets); void CreateResumptionSwitch(GenTreeLclVarCommon* commonAsyncResumedDef); + typedef JitHashTable, BasicBlock*> InlinedFrameTailMap; + // Suspension tails handling the frames enclosing an inlined async frame, keyed by the + // innermost frame's resumed indicator local. Shared by all suspensions in that frame. + InlinedFrameTailMap* m_inlinedFrameTails = nullptr; + + BasicBlock* CreateInlinedFrameSuspensionTail(BasicBlock* callBlock, + GenTreeCall* call, + const ContinuationLayout& layout); + GenTree* ContinuationMemberAddress(const ContinuationLayout& layout, const ContinuationMember& member); + public: AsyncTransformation(Compiler* comp) : m_compiler(comp) diff --git a/src/coreclr/jit/fginline.cpp b/src/coreclr/jit/fginline.cpp index 1b03a8e1a4589c..b8a51caee04193 100644 --- a/src/coreclr/jit/fginline.cpp +++ b/src/coreclr/jit/fginline.cpp @@ -2462,14 +2462,16 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) return; } - struct AsyncFrameLocals - { - unsigned Resumed; - unsigned ExecutionContext; - unsigned SynchronizationContext; - }; - + // AsyncFrameLocals is declared alongside AsyncCallInfo, which is where the chain is + // recorded for the async transformation. ArrayStack frames(getAllocator(CMK_Async)); + + // The inlinee's own frame comes first: the suspension needs its resumed indicator to + // decide whether to capture the contexts it hands to its caller. Its other context + // args were consumed by the innermost handling by then. + frames.Push({InlineeCompiler->lvaResumedIndicator, InlineeCompiler->lvaAsyncExecutionContextVar, + InlineeCompiler->lvaAsyncSynchronizationContextVar}); + for (InlineContext* ctx = inlineInfo->inlineContext->GetParent(); ctx != nullptr; ctx = ctx->GetParent()) { if (ctx->HasAsyncFrameLocals()) @@ -2482,6 +2484,12 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) // The root method's frame is always the outermost one. frames.Push({lvaResumedIndicator, lvaAsyncExecutionContextVar, lvaAsyncSynchronizationContextVar}); + AsyncFrameLocals* const frameLocals = new (this, CMK_Async) AsyncFrameLocals[frames.Height()]; + for (int i = 0; i < frames.Height(); i++) + { + frameLocals[i] = frames.Bottom(i); + } + unsigned const inlineeResumed = InlineeCompiler->lvaResumedIndicator; struct Visitor : GenTreeVisitor @@ -2492,11 +2500,16 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) }; ArrayStack& m_frames; + AsyncFrameLocals* m_frameLocals; unsigned m_inlineeResumed; - Visitor(Compiler* comp, ArrayStack& frames, unsigned inlineeResumed) + Visitor(Compiler* comp, + ArrayStack& frames, + AsyncFrameLocals* frameLocals, + unsigned inlineeResumed) : GenTreeVisitor(comp) , m_frames(frames) + , m_frameLocals(frameLocals) , m_inlineeResumed(inlineeResumed) { } @@ -2544,11 +2557,17 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) JITDUMP("Added %d enclosing frame context arg sets to async call [%06u]\n", m_frames.Height(), Compiler::dspTreeID(call)); + + // The chain holds this call's own frame followed by its enclosing frames, + // ending with the root, so its depth in that numbering is one less than the + // number of entries. + call->GetAsyncInfo().InlineFrameDepth = (unsigned)m_frames.Height() - 1; + call->GetAsyncInfo().InlineFrameLocals = m_frameLocals; return WALK_CONTINUE; } }; - Visitor visitor(this, frames, inlineeResumed); + Visitor visitor(this, frames, frameLocals, inlineeResumed); for (BasicBlock* block = InlineeCompiler->fgFirstBB; block != nullptr; block = block->Next()) { for (Statement* const stmt : block->Statements()) @@ -2742,9 +2761,9 @@ void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicB assert(resumedCaller != BAD_VAR_NUM); unsigned const resumedInlinee = InlineeCompiler->lvaResumedIndicator; - unsigned const inlineDepth = inlineContext->GetDepth(); + unsigned const inlineDepth = inlineContext->GetAsyncFrameDepth(); - JITDUMP("Adding async frame transition IR for inline depth %u: resumed V%02u -> V%02u\n", inlineDepth, + JITDUMP("Adding async frame transition IR for async frame depth %u: resumed V%02u -> V%02u\n", inlineDepth, resumedInlinee, resumedCaller); const DebugInfo& di = inlineInfo->iciStmt->GetDebugInfo(); @@ -2792,7 +2811,7 @@ void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicB fgMarkAsyncHelperInlineCandidate(restoreCall); fgInsertStmtAtEnd(restoreBlock, gtNewStmt(restoreCall)); - GenTreeCall* const isOnRightContextCall = gtNewUserCallNode(asyncInfo->isOnRightContextMethHnd, TYP_INT); + GenTreeCall* const isOnRightContextCall = gtNewUserCallNode(asyncInfo->isOnRightContextMethHnd, TYP_UBYTE); isOnRightContextCall->gtArgs .PushFront(this, NewCallArg::Primitive( @@ -2805,12 +2824,15 @@ void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicB fgMarkAsyncHelperInlineCandidate(isOnRightContextCall); // Inline candidates must be statement roots; the value is consumed through a - // GT_RET_EXPR placeholder that the inliner substitutes. + // GT_RET_EXPR placeholder that the inliner substitutes. The candidate info has to + // point back at it so that a failed inline can put the call back. GenTree* isOnRightContext; if (isOnRightContextCall->IsInlineCandidate()) { fgInsertStmtAtEnd(restoreBlock, gtNewStmt(isOnRightContextCall)); - isOnRightContext = gtNewInlineCandidateReturnExpr(isOnRightContextCall, TYP_INT); + GenTreeRetExpr* const retExpr = gtNewInlineCandidateReturnExpr(isOnRightContextCall, TYP_UBYTE); + isOnRightContextCall->GetSingleInlineCandidateInfo()->retExpr = retExpr; + isOnRightContext = retExpr; } else { diff --git a/src/coreclr/jit/gentree.h b/src/coreclr/jit/gentree.h index b451be41097493..6d2e608181d17a 100644 --- a/src/coreclr/jit/gentree.h +++ b/src/coreclr/jit/gentree.h @@ -4529,6 +4529,14 @@ enum class ContinuationContextHandling }; // Additional async call info. +// Async context locals of one inlined frame. +struct AsyncFrameLocals +{ + unsigned Resumed; + unsigned ExecutionContext; + unsigned SynchronizationContext; +}; + struct AsyncCallInfo { // DebugInfo with SOURCE_TYPE_ASYNC pointing at the await call IL instruction @@ -4554,6 +4562,24 @@ struct AsyncCallInfo // and suspend unconditionally. bool AlwaysSuspends = false; + // Depth of the inlined frame this call belongs to, counting only frames that have + // async context handling. 0 is the root method's frame. + // + // On suspension every enclosing frame that has not yet resumed must capture the + // contexts it would hand to its caller. Those live in continuation members keyed by + // this depth, and the enclosing frames' depths follow by decrementing while draining + // the call's chain of context args. + unsigned InlineFrameDepth = 0; + + // Async context locals of this call's own inlined frame followed by its enclosing + // frames, ending with the root method's. Null when the call is not inside an inlined + // async frame; otherwise it holds InlineFrameDepth + 1 entries. + // + // The matching context args on the call keep these values live and reported up to the + // async transformation. This records which locals they are, since optimizations may + // rewrite the arg nodes themselves. + AsyncFrameLocals* InlineFrameLocals = nullptr; + bool NeedsToSaveAndRestoreExecutionContext() const { return true; @@ -5273,6 +5299,12 @@ struct GenTreeCall final : public GenTree return *asyncInfo; } + AsyncCallInfo& GetAsyncInfo() + { + assert(IsAsync()); + return *asyncInfo; + } + //--------------------------------------------------------------------------- // GetRegNumByIdx: get i'th return register allocated to this call node. // diff --git a/src/coreclr/jit/importer.cpp b/src/coreclr/jit/importer.cpp index 90396e6a294bc5..925b76b0c39e58 100644 --- a/src/coreclr/jit/importer.cpp +++ b/src/coreclr/jit/importer.cpp @@ -12013,10 +12013,17 @@ bool Compiler::impWrapTopOfStackInAwait() JITDUMP("Inheriting continuation handling %d from caller [%06u]\n", (unsigned)inlCall->GetAsyncInfo().ContinuationContextHandling, dspTreeID(inlCall)); asyncInfo->ContinuationContextHandling = inlCall->GetAsyncInfo().ContinuationContextHandling; + + // Mark the call async first: inheriting the contexts also inherits the inlined + // frame depth, which lives in the async call info. + awaitCall->SetIsAsync(asyncInfo); impInheritAsyncContextsFromInliner(awaitCall); } - awaitCall->SetIsAsync(asyncInfo); + if (!awaitCall->IsAsync()) + { + awaitCall->SetIsAsync(asyncInfo); + } if (awaitCall->IsInlineCandidate()) { diff --git a/src/coreclr/jit/importercalls.cpp b/src/coreclr/jit/importercalls.cpp index e10c32831267fe..3884e157db7f5c 100644 --- a/src/coreclr/jit/importercalls.cpp +++ b/src/coreclr/jit/importercalls.cpp @@ -7578,6 +7578,42 @@ void Compiler::impInheritAsyncContextsFromInliner(GenTreeCall* call) call->gtArgs.PushFront(this, NewCallArg::Primitive(execNode).WellKnown(WellKnownArg::AsyncExecutionContext)); call->gtArgs.PushFront(this, NewCallArg::Primitive(resumedDefNode).WellKnown(WellKnownArg::AsyncResumedDef)); call->gtArgs.PushFront(this, NewCallArg::Primitive(resumedUseNode).WellKnown(WellKnownArg::AsyncResumedUse)); + + // The inlining call may carry further sets describing the frames enclosing it, which + // this call inherits as well: it ends up in the same frame, so a suspension in it has + // to run the same chain of frame transitions. Dropping them would silently lose the + // handling for every frame outside the immediate one. + bool skippedFirst = false; + for (CallArg& arg : inlCall->gtArgs.Args()) + { + WellKnownArg wka = arg.GetWellKnownArg(); + if ((wka != WellKnownArg::AsyncResumedUse) && (wka != WellKnownArg::AsyncExecutionContext) && + (wka != WellKnownArg::AsyncSynchronizationContext)) + { + continue; + } + + if ((wka == WellKnownArg::AsyncResumedUse) && !skippedFirst) + { + // Already inherited above, along with the resumed def that only the innermost + // frame has. + skippedFirst = true; + continue; + } + + if (&arg == execArg || &arg == syncArg) + { + continue; + } + + call->gtArgs.PushBack(this, NewCallArg::Primitive(gtCloneExpr(arg.GetNode())).WellKnown(wka)); + } + + if (call->IsAsync()) + { + call->GetAsyncInfo().InlineFrameDepth = inlCall->GetAsyncInfo().InlineFrameDepth; + call->GetAsyncInfo().InlineFrameLocals = inlCall->GetAsyncInfo().InlineFrameLocals; + } } //------------------------------------------------------------------------ diff --git a/src/coreclr/jit/inline.h b/src/coreclr/jit/inline.h index f2a0df1b828613..de7fdedd4ab065 100644 --- a/src/coreclr/jit/inline.h +++ b/src/coreclr/jit/inline.h @@ -793,13 +793,22 @@ class InlineContext return m_Parent; } - // Get the inline depth of this context. The root context has depth 0. - unsigned GetDepth() const + // Depth of the inlined frame this context represents, counting only frames that have + // async context handling. The root method's frame is depth 0, so the shallowest + // inlined frame with context handling is depth 1. + // + // Continuation members holding the contexts captured at frame transitions are keyed + // by this, so frames at the same depth share storage. That is safe because their live + // ranges cannot overlap. + unsigned GetAsyncFrameDepth() const { - unsigned depth = 0; + unsigned depth = 1; for (InlineContext* parent = m_Parent; parent != nullptr; parent = parent->m_Parent) { - depth++; + if (parent->HasAsyncFrameLocals()) + { + depth++; + } } return depth; diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs index a7aba5a6d1f302..4743519a85850b 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs @@ -3624,6 +3624,7 @@ private void getAsyncInfo(ref CORINFO_ASYNC_INFO pAsyncInfoOut) pAsyncInfoOut.restoreInlinedFrameExecutionContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreInlinedFrameExecutionContext"u8, null)); pAsyncInfoOut.isOnRightContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("IsOnRightContext"u8, null)); pAsyncInfoOut.switchContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("SwitchContext"u8, null)); + pAsyncInfoOut.captureInlinedFrameTransitionMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("CaptureInlinedFrameTransition"u8, null)); pAsyncInfoOut.finishSuspensionNoContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionNoContinuationContext"u8, null)); pAsyncInfoOut.finishSuspensionWithContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionWithContinuationContext"u8, null)); } diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs index 013a5c6638ac1d..1ce537052d2450 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs @@ -973,6 +973,7 @@ public unsafe struct CORINFO_ASYNC_INFO public CORINFO_METHOD_STRUCT_* restoreInlinedFrameExecutionContextMethHnd; public CORINFO_METHOD_STRUCT_* isOnRightContextMethHnd; public CORINFO_METHOD_STRUCT_* switchContextMethHnd; + public CORINFO_METHOD_STRUCT_* captureInlinedFrameTransitionMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionNoContinuationContextMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionWithContinuationContextMethHnd; } diff --git a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h index e062a03b1bd02c..97b3162ee6afa3 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h +++ b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h @@ -240,6 +240,7 @@ struct Agnostic_CORINFO_ASYNC_INFO DWORDLONG restoreInlinedFrameExecutionContextMethHnd; DWORDLONG isOnRightContextMethHnd; DWORDLONG switchContextMethHnd; + DWORDLONG captureInlinedFrameTransitionMethHnd; DWORDLONG finishSuspensionNoContinuationContextMethHnd; DWORDLONG finishSuspensionWithContinuationContextMethHnd; }; diff --git a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp index efeb1ab0ac2b62..89bc70c963b682 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp +++ b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp @@ -4397,6 +4397,7 @@ void MethodContext::recGetAsyncInfo(const CORINFO_ASYNC_INFO* pAsyncInfo) value.restoreInlinedFrameExecutionContextMethHnd = CastHandle(pAsyncInfo->restoreInlinedFrameExecutionContextMethHnd); value.isOnRightContextMethHnd = CastHandle(pAsyncInfo->isOnRightContextMethHnd); value.switchContextMethHnd = CastHandle(pAsyncInfo->switchContextMethHnd); + value.captureInlinedFrameTransitionMethHnd = CastHandle(pAsyncInfo->captureInlinedFrameTransitionMethHnd); value.finishSuspensionNoContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionNoContinuationContextMethHnd); value.finishSuspensionWithContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionWithContinuationContextMethHnd); @@ -4426,6 +4427,7 @@ void MethodContext::repGetAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->restoreInlinedFrameExecutionContextMethHnd = (CORINFO_METHOD_HANDLE)value.restoreInlinedFrameExecutionContextMethHnd; pAsyncInfoOut->isOnRightContextMethHnd = (CORINFO_METHOD_HANDLE)value.isOnRightContextMethHnd; pAsyncInfoOut->switchContextMethHnd = (CORINFO_METHOD_HANDLE)value.switchContextMethHnd; + pAsyncInfoOut->captureInlinedFrameTransitionMethHnd = (CORINFO_METHOD_HANDLE)value.captureInlinedFrameTransitionMethHnd; pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionNoContinuationContextMethHnd; pAsyncInfoOut->finishSuspensionWithContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionWithContinuationContextMethHnd; DEBUG_REP(dmpGetAsyncInfo(0, value)); diff --git a/src/coreclr/vm/corelib.h b/src/coreclr/vm/corelib.h index dbec6e83f5bbd0..c079ed86db6398 100644 --- a/src/coreclr/vm/corelib.h +++ b/src/coreclr/vm/corelib.h @@ -733,6 +733,7 @@ DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS_ON_SUSPENSION, RestoreContext DEFINE_METHOD(ASYNC_HELPERS, RESTORE_INLINED_FRAME_EXECUTION_CONTEXT, RestoreInlinedFrameExecutionContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, IS_ON_RIGHT_CONTEXT, IsOnRightContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, SWITCH_CONTEXT, SwitchContext, NoSig) +DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_INLINED_FRAME_TRANSITION, CaptureInlinedFrameTransition, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT, FinishSuspensionNoContinuationContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_WITH_CONTINUATION_CONTEXT, FinishSuspensionWithContinuationContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, ASYNC_CALL_CONTINUATION, AsyncCallContinuation, NoSig) diff --git a/src/coreclr/vm/jitinterface.cpp b/src/coreclr/vm/jitinterface.cpp index 051e9c0eed19d0..4719f93c11a05a 100644 --- a/src/coreclr/vm/jitinterface.cpp +++ b/src/coreclr/vm/jitinterface.cpp @@ -10392,6 +10392,7 @@ void CEEInfo::getAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->restoreInlinedFrameExecutionContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_INLINED_FRAME_EXECUTION_CONTEXT)); pAsyncInfoOut->isOnRightContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__IS_ON_RIGHT_CONTEXT)); pAsyncInfoOut->switchContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__SWITCH_CONTEXT)); + pAsyncInfoOut->captureInlinedFrameTransitionMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__CAPTURE_INLINED_FRAME_TRANSITION)); pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT)); pAsyncInfoOut->finishSuspensionWithContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_WITH_CONTINUATION_CONTEXT)); From d6c80dd635a704a9de0040d315bb3d7c31e00800 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 17:12:24 +0200 Subject: [PATCH 27/41] Add tests for context handling across inlined async frames The worked examples from the design document. Inlining removes a call site and a callee, and with them the context handling the async infrastructure would otherwise perform at each frame boundary, so these pin down the observable results: each frame resumes on the synchronization context its own continuation captured, and sees the AsyncLocal values from the ExecutionContext it captured. The callees are marked AggressiveInlining because the profitability heuristic otherwise rejects them, and the test runs with tiering disabled, since at tier 0 nothing is inlined and the tests would not exercise the feature at all. Verified to exercise a three frame chain, and to fail if the frame transition IR is removed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../inlined-frame-contexts.cs | 196 ++++++++++++++++++ .../inlined-frame-contexts.csproj | 14 ++ 2 files changed, 210 insertions(+) create mode 100644 src/tests/async/inlined-frame-contexts/inlined-frame-contexts.cs create mode 100644 src/tests/async/inlined-frame-contexts/inlined-frame-contexts.csproj diff --git a/src/tests/async/inlined-frame-contexts/inlined-frame-contexts.cs b/src/tests/async/inlined-frame-contexts/inlined-frame-contexts.cs new file mode 100644 index 00000000000000..0e078d81193d99 --- /dev/null +++ b/src/tests/async/inlined-frame-contexts/inlined-frame-contexts.cs @@ -0,0 +1,196 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +// Behaviors that must be preserved when runtime async calls are inlined. These are the +// worked examples from docs/design/coreclr/jit/runtime-async-inlining.md. +// +// Inlining removes both a call site and a callee, and with them the context handling the +// async infrastructure would otherwise have performed at each frame boundary. These tests +// pin down the observable results of that handling, so they must produce the same answers +// whether or not the calls are inlined. +// +// The inner frames are marked AggressiveInlining so that the inlining actually happens +// rather than being left to the profitability heuristic, which rejects these callees. +public class Async2InlinedFrameContexts +{ + private sealed class NamedContext : SynchronizationContext + { + public readonly string Name; + + public NamedContext(string name) => Name = name; + + public override void Post(SendOrPostCallback d, object? state) + { + SynchronizationContext currentCtx = Current; + SetSynchronizationContext(this); + try + { + d(state); + } + finally + { + SetSynchronizationContext(currentCtx); + } + } + } + + // Always suspends and resumes on a thread pool thread, so resumption never happens on + // whatever context the awaiting frames were running in. + private struct AlwaysThreadPoolAwaitable : INotifyCompletion + { + public bool IsCompleted => false; + + public void OnCompleted(Action continuation) + { + ThreadPool.QueueUserWorkItem(_ => continuation()); + } + + public void GetResult() + { + } + + public AlwaysThreadPoolAwaitable GetAwaiter() => this; + } + + private static readonly NamedContext s_syncContext1 = new NamedContext("ctx1"); + private static readonly NamedContext s_syncContext2 = new NamedContext("ctx2"); + + private static string CurrentContextName => (SynchronizationContext.Current as NamedContext)?.Name ?? "none"; + + // Each frame sets its own synchronization context before awaiting, so each frame's + // continuation captures a different one. After the innermost await resumes on the + // thread pool, every frame must be back on the context it captured by the time it + // observes it again. + private static string s_fooContextAfterAwait = ""; + private static string s_barContextAfterAwait = ""; + + private static async Task ContextFoo() + { + SynchronizationContext.SetSynchronizationContext(s_syncContext1); + await ContextBar(); + s_fooContextAfterAwait = CurrentContextName; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static async Task ContextBar() + { + SynchronizationContext.SetSynchronizationContext(s_syncContext2); + await ContextBaz(); + s_barContextAfterAwait = CurrentContextName; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static async Task ContextBaz() + { + await new AlwaysThreadPoolAwaitable(); + } + + [Fact] + public static void SynchronizationContextIsRestoredPerFrame() + { + SynchronizationContext original = SynchronizationContext.Current; + try + { + ContextFoo().GetAwaiter().GetResult(); + } + finally + { + SynchronizationContext.SetSynchronizationContext(original); + } + + // Bar awaited Baz having set ctx2, so Bar resumes on ctx2. Foo awaited Bar having + // set ctx1, so Foo resumes on ctx1. + Assert.Equal("ctx2", s_barContextAfterAwait); + Assert.Equal("ctx1", s_fooContextAfterAwait); + } + + private static readonly AsyncLocal s_local = new AsyncLocal(); + + private static int s_bazLocalAfterAwait; + private static int s_barLocalAfterAwait; + private static int s_fooLocalAfterAwait; + + private static async Task LocalFoo() + { + s_local.Value = 1; + await LocalBar(); + s_fooLocalAfterAwait = s_local.Value; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static async Task LocalBar() + { + s_local.Value = 2; + await LocalBaz(); + s_barLocalAfterAwait = s_local.Value; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static async Task LocalBaz() + { + s_local.Value = 3; + await new AlwaysThreadPoolAwaitable(); + s_bazLocalAfterAwait = s_local.Value; + } + + [Fact] + public static void ExecutionContextIsRestoredPerFrame() + { + LocalFoo().GetAwaiter().GetResult(); + + // Each frame sees the AsyncLocal value it set, because every frame boundary + // restores the ExecutionContext that frame captured. + Assert.Equal(3, s_bazLocalAfterAwait); + Assert.Equal(2, s_barLocalAfterAwait); + Assert.Equal(1, s_fooLocalAfterAwait); + } + + // The same, but with the awaits inside loops so the frames suspend and resume + // repeatedly. A frame's captured contexts must survive later suspensions of the + // frames nested inside it. + private static async Task LoopFoo(int iterations) + { + SynchronizationContext.SetSynchronizationContext(s_syncContext1); + s_local.Value = 1; + + for (int i = 0; i < iterations; i++) + { + await LoopBar(iterations); + Assert.Equal("ctx1", CurrentContextName); + Assert.Equal(1, s_local.Value); + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static async Task LoopBar(int iterations) + { + SynchronizationContext.SetSynchronizationContext(s_syncContext2); + s_local.Value = 2; + + for (int i = 0; i < iterations; i++) + { + await new AlwaysThreadPoolAwaitable(); + Assert.Equal(2, s_local.Value); + } + } + + [Fact] + public static void ContextsSurviveRepeatedSuspensions() + { + SynchronizationContext original = SynchronizationContext.Current; + try + { + LoopFoo(4).GetAwaiter().GetResult(); + } + finally + { + SynchronizationContext.SetSynchronizationContext(original); + } + } +} diff --git a/src/tests/async/inlined-frame-contexts/inlined-frame-contexts.csproj b/src/tests/async/inlined-frame-contexts/inlined-frame-contexts.csproj new file mode 100644 index 00000000000000..5ca3f694332957 --- /dev/null +++ b/src/tests/async/inlined-frame-contexts/inlined-frame-contexts.csproj @@ -0,0 +1,14 @@ + + + + true + + + + + + + + + From 9911625560797e2b2341903c71efa35b7bff48ac Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 17:25:34 +0200 Subject: [PATCH 28/41] Add OSR coverage for async inlining, and correct an earlier OSR claim Adds a test covering runtime async inlining in methods rejitted via OSR, including resuming inside an inlined frame there. Verified to exercise an inlined frame with a suspension tail inside an OSR variant. Also corrects the change from 'Apply OSR async context handling to the root frame only'. That commit claimed opts.IsOSR() was true while importing an inlinee, so the OSR-specific handling was being wrongly applied to inlinees. It is not: compInitOptions clears JIT_FLAG_OSR for inlinees (compiler.cpp), so the added condition was redundant and the bugs it described did not exist. The condition is removed again and replaced with a comment recording why the handling only ever applies to the root frame. Removing the earlier bail on inlining into OSR methods was still correct, which the new test demonstrates. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 10 +- .../osr-inlined-contexts.cs | 125 ++++++++++++++++++ .../osr-inlined-contexts.csproj | 18 +++ 3 files changed, 148 insertions(+), 5 deletions(-) create mode 100644 src/tests/async/osr-inlined-contexts/osr-inlined-contexts.cs create mode 100644 src/tests/async/osr-inlined-contexts/osr-inlined-contexts.csproj diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 590976d4b36483..90804246b807af 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -251,11 +251,11 @@ PhaseStatus Compiler::SaveAsyncContexts() return PhaseStatus::MODIFIED_NOTHING; } - // OSR context handling concerns the root frame only. An inlinee inside an OSR method - // still starts a fresh logical frame: it must capture its own contexts, its resumed - // indicator starts out false regardless of how the OSR method was entered, and its - // locals have no home in the tier0 frame. - bool isOSRRootFrame = opts.IsOSR() && !compIsForInlining(); + // Note this only ever applies to the root frame: JIT_FLAG_OSR is cleared for inlinees + // in compInitOptions. An inlinee inside an OSR method starts a fresh logical frame and + // so captures its own contexts, with a resumed indicator that starts out false however + // the OSR method was entered. + bool isOSRRootFrame = opts.IsOSR(); // Create locals for Thread, ExecutionContext and SynchronizationContext lvaAsyncThreadObjectVar = lvaGrabTemp(false DEBUGARG("Async Thread")); diff --git a/src/tests/async/osr-inlined-contexts/osr-inlined-contexts.cs b/src/tests/async/osr-inlined-contexts/osr-inlined-contexts.cs new file mode 100644 index 00000000000000..c6af045972fd4b --- /dev/null +++ b/src/tests/async/osr-inlined-contexts/osr-inlined-contexts.cs @@ -0,0 +1,125 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +// Coverage for runtime async inlining in methods that are rejitted via OSR. +// +// An OSR method is entered with the tier 0 method's continuation, whose layout differs +// from the one the OSR method computes. Reading an inlined frame's state off it would be +// wrong, so this pins down that inlining async callees into an OSR method still behaves: +// an inlined callee's contexts are its own, and resuming inside one keeps every frame's +// contexts intact. +public class Async2OsrInlinedContexts +{ + private sealed class MarkerContext : SynchronizationContext + { + // Re-establish this context around the callback, as a real UI-style context would. + // The default implementation queues to the thread pool without doing so, which + // would leave the continuation with no current context at all. + public override void Post(SendOrPostCallback d, object? state) + { + SynchronizationContext currentCtx = Current; + SetSynchronizationContext(this); + try + { + d(state); + } + finally + { + SetSynchronizationContext(currentCtx); + } + } + } + + // Enough iterations to get the enclosing method rejitted via OSR. + private const int Iterations = 100_000; + + private static int s_sideEffect; + + // Async, and so has its own context save and restore, but never suspends, so it is + // inlinable into the loop below. + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static async Task BumpAsync(int value) + { + s_sideEffect += value; + } + + private static async Task LoopWithInlinedCallee() + { + int total = 0; + for (int i = 0; i < Iterations; i++) + { + await BumpAsync(1); + total += i; + } + + return total; + } + + [Fact] + public static void InlinedCalleeDoesNotClobberContextsInOsrMethod() + { + SynchronizationContext original = SynchronizationContext.Current; + MarkerContext marker = new MarkerContext(); + try + { + SynchronizationContext.SetSynchronizationContext(marker); + + LoopWithInlinedCallee().GetAwaiter().GetResult(); + + // The inlined callee has its own context save and restore, which must not + // disturb the caller's. + Assert.Same(marker, SynchronizationContext.Current); + } + finally + { + SynchronizationContext.SetSynchronizationContext(original); + } + } + + private static readonly AsyncLocal s_local = new AsyncLocal(); + + // Same, but the callee suspends, so the loop also exercises resuming inside an inlined + // frame of an OSR method. + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static async Task YieldAsync() + { + await Task.Yield(); + } + + private static async Task LoopWithSuspendingInlinedCallee() + { + SynchronizationContext.SetSynchronizationContext(new MarkerContext()); + s_local.Value = 7; + + int total = 0; + for (int i = 0; i < 2_000; i++) + { + await YieldAsync(); + Assert.IsType(SynchronizationContext.Current); + Assert.Equal(7, s_local.Value); + total += i; + } + + return total; + } + + [Fact] + public static void ResumingInsideInlinedFrameOfOsrMethodKeepsContexts() + { + SynchronizationContext original = SynchronizationContext.Current; + try + { + LoopWithSuspendingInlinedCallee().GetAwaiter().GetResult(); + } + finally + { + SynchronizationContext.SetSynchronizationContext(original); + } + } +} diff --git a/src/tests/async/osr-inlined-contexts/osr-inlined-contexts.csproj b/src/tests/async/osr-inlined-contexts/osr-inlined-contexts.csproj new file mode 100644 index 00000000000000..8525b085aae062 --- /dev/null +++ b/src/tests/async/osr-inlined-contexts/osr-inlined-contexts.csproj @@ -0,0 +1,18 @@ + + + + true + True + + + + + + + + + + + + From 34a87848d13ed48ae9470f50c98880a1c44caecd Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 17:35:28 +0200 Subject: [PATCH 29/41] JIT: Support cloning GT_CONTINUATION_MEMBER_OFFSET It is a GenTreeVal like GT_RECORD_ASYNC_RESUME and GT_ASYNC_RESUME_INFO, but was missing from gtCloneExpr, so cloning one hit 'Cloning of node not supported'. Reachable whenever such a node ends up somewhere that gets cloned, for example a loop body that is unrolled. Found by Fuzzlyn. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/gentree.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/coreclr/jit/gentree.cpp b/src/coreclr/jit/gentree.cpp index 1a48d4a15488d6..484ca0097d4e29 100644 --- a/src/coreclr/jit/gentree.cpp +++ b/src/coreclr/jit/gentree.cpp @@ -11165,6 +11165,7 @@ GenTree* Compiler::gtCloneExpr(GenTree* tree) case GT_JMP: case GT_RECORD_ASYNC_RESUME: case GT_ASYNC_RESUME_INFO: + case GT_CONTINUATION_MEMBER_OFFSET: copy = new (this, oper) GenTreeVal(oper, tree->gtType, tree->AsVal()->gtVal1); goto DONE; From 30078a635b83f419269b0de4aa95110b7618a3f8 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 18:59:32 +0200 Subject: [PATCH 30/41] Enable by default for testing --- src/coreclr/jit/jitconfigvalues.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coreclr/jit/jitconfigvalues.h b/src/coreclr/jit/jitconfigvalues.h index 12a5c44f6cbdc3..b497598377ea40 100644 --- a/src/coreclr/jit/jitconfigvalues.h +++ b/src/coreclr/jit/jitconfigvalues.h @@ -624,7 +624,7 @@ OPT_CONFIG_STRING(JitAsyncPreservedValueAnalysisRange, "JitAsyncPreservedValueAn // callees that may suspend. When zero, only the restricted cases are inlined: // callees without any awaits, async versions of synchronous methods, and tail // awaits. -RELEASE_CONFIG_INTEGER(JitAsyncInlining, "JitAsyncInlining", 0) +RELEASE_CONFIG_INTEGER(JitAsyncInlining, "JitAsyncInlining", 1) RELEASE_CONFIG_INTEGER(JitEnableOptRepeat, "JitEnableOptRepeat", 1) // If zero, do not allow JitOptRepeat RELEASE_CONFIG_METHODSET(JitOptRepeat, "JitOptRepeat") // Runs optimizer multiple times on specified methods From 096e3683c9b234fc91c1ad5d57c9b345c652b749 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Wed, 29 Jul 2026 19:16:55 +0200 Subject: [PATCH 31/41] Fix --- src/coreclr/jit/async.cpp | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 628600cf005c74..042cbf13fbfdf8 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -3124,10 +3124,7 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* for (CallArg* arg : {resumedArg, execContextArg, syncContextArg}) { GenTree* node = arg->GetNode(); - if (LIR::AsRange(callBlock).Contains(node)) - { - LIR::AsRange(callBlock).Remove(node); - } + LIR::AsRange(callBlock).Remove(node); call->gtArgs.RemoveUnsafe(arg); } From 0d86ed624bf4c0fc34bdde21ff9335a0b129f11a Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Thu, 30 Jul 2026 17:19:27 +0200 Subject: [PATCH 32/41] JIT: Take inlined frame suspension values from the IR, not a side table The context pseudo-args on an async call are the source of truth for what the suspension stores. The whole optimizer runs between adding them and the async transformation, so those nodes can be rewritten: 'resumed' indicators in particular get const folded, and contexts can be copy propagated to other locals. Reconstructing the values from the locals recorded at inline time ignored all of that and stored whatever the original locals happened to hold. Take the actual nodes off the call instead, spilling to a temp only when a node is neither invariant nor a local, exactly as the existing shared finish-context handling does when it needs to reference an arg from another block. The nodes also become the identity used for sharing a tail, so suspensions share one only when they agree on every value. That is now observable: with two awaits in the same inlined frame, the first has its resumed indicator folded to the constant zero while the second still reads the local, and they correctly get separate tails. Sharing them would have stored the wrong value at one of them. This removes the need for the side table, so AsyncCallInfo::InlineFrameLocals and the AsyncFrameLocals type go away. InlineContext keeps its frame locals, which are still the right source when creating the args during inlining. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 119 ++++++++++++++++++------------ src/coreclr/jit/async.h | 17 ++++- src/coreclr/jit/fginline.cpp | 28 +++---- src/coreclr/jit/gentree.h | 17 ----- src/coreclr/jit/importercalls.cpp | 3 +- 5 files changed, 95 insertions(+), 89 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index 1591798111dec5..a5fb7fefc9475b 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -3104,57 +3104,81 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* GenTreeCall* call, const ContinuationLayout& layout) { - // Remove the context args of the enclosing frames. They exist only to keep the values - // live and reported up to here; which locals they are is recorded in the call info, - // since optimizations may have rewritten the arg nodes themselves. - while (true) - { - CallArg* resumedArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse); - CallArg* execContextArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncExecutionContext); - CallArg* syncContextArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncSynchronizationContext); - if (resumedArg == nullptr) - { - assert((execContextArg == nullptr) && (syncContextArg == nullptr)); - break; - } + unsigned const numFrames = call->GetAsyncInfo().InlineFrameDepth + 1; + if (numFrames < 2) + { + // Not inside an inlined async frame, so there are no logical frame transitions to + // run and no context args beyond the ones the suspension itself consumed. + assert(call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse) == nullptr); + return nullptr; + } + + JITDUMP(" Call [%06u] is inside an inlined async frame; %u frames in chain\n", Compiler::dspTreeID(call), + numFrames); - assert((execContextArg != nullptr) && (syncContextArg != nullptr)); + // Take the values of the enclosing frames off the call. These args are the source of + // truth for what to store: the optimizer may have rewritten them since they were + // added, for example folding a "resumed" indicator to a constant, and the stores must + // reflect that rather than re-reading the locals the args originally named. + unsigned const numValues = numFrames * 3; + GenTree** const values = new (m_compiler, CMK_Async) GenTree*[numValues]; - for (CallArg* arg : {resumedArg, execContextArg, syncContextArg}) + for (unsigned i = 0; i < numFrames; i++) + { + CallArg* frameArgs[3] = {call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse), + call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncExecutionContext), + call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncSynchronizationContext)}; + + for (unsigned j = 0; j < 3; j++) { + CallArg* arg = frameArgs[j]; + assert(arg != nullptr); + GenTree* node = arg->GetNode(); - LIR::AsRange(callBlock).Remove(node); + if (!node->IsInvariant() && !node->OperIs(GT_LCL_VAR)) + { + // We are going to reference this from a different block, so spill it. + LIR::Use use(LIR::AsRange(callBlock), &arg->NodeRef(), call); + use.ReplaceWithLclVar(m_compiler); + node = use.Def(); + } + LIR::AsRange(callBlock).Remove(node); call->gtArgs.RemoveUnsafe(arg); + values[i * 3 + j] = node; } } - const AsyncFrameLocals* const frameLocals = call->GetAsyncInfo().InlineFrameLocals; - if (frameLocals == nullptr) + // The tail depends only on these values, so suspensions that agree on all of them can + // share one. That is the common case: every suspension in the same inlined frame sees + // the same locals unless the optimizer rewrote one of them differently. + if (m_inlinedFrameTails == nullptr) { - JITDUMP(" Call [%06u] is not inside an inlined async frame\n", Compiler::dspTreeID(call)); - return nullptr; + m_inlinedFrameTails = + new (m_compiler, CMK_Async) jitstd::vector(m_compiler->getAllocator(CMK_Async)); } - unsigned const innermostDepth = call->GetAsyncInfo().InlineFrameDepth; - unsigned const numFrames = innermostDepth + 1; - assert(numFrames >= 2); + for (const InlinedFrameTail& tail : *m_inlinedFrameTails) + { + if (tail.NumValues != numValues) + { + continue; + } - JITDUMP(" Call [%06u] is inside an inlined async frame; %u frames in chain\n", Compiler::dspTreeID(call), - numFrames); + bool match = true; + for (unsigned i = 0; i < numValues; i++) + { + if (!GenTree::Compare(tail.Values[i], values[i])) + { + match = false; + break; + } + } - unsigned const key = frameLocals[0].Resumed; - if (m_inlinedFrameTails == nullptr) - { - m_inlinedFrameTails = new (m_compiler, CMK_Async) InlinedFrameTailMap(m_compiler->getAllocator(CMK_Async)); - } - else - { - BasicBlock* existing; - if (m_inlinedFrameTails->Lookup(key, &existing)) + if (match) { - JITDUMP(" Reusing inlined frame suspension tail " FMT_BB "\n", existing->bbNum); - return existing; + JITDUMP(" Reusing inlined frame suspension tail " FMT_BB "\n", tail.Block->bbNum); + return tail.Block; } } @@ -3178,9 +3202,11 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* for (unsigned i = 0; i + 1 < numFrames; i++) { - const AsyncFrameLocals& frame = frameLocals[i]; - const AsyncFrameLocals& outer = frameLocals[i + 1]; - unsigned const depth = innermostDepth - i; + GenTree* const frameResumed = values[i * 3]; + GenTree* const outerResumed = values[(i + 1) * 3]; + GenTree* const outerExec = values[(i + 1) * 3 + 1]; + GenTree* const outerSync = values[(i + 1) * 3 + 2]; + unsigned const depth = numFrames - 1 - i; // Capture what this frame hands to its caller. GenTreeCall* captureCall = @@ -3198,8 +3224,7 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* ContinuationMemberAddress(layout, ContinuationMember::InlineFrameContinuationContext( depth)))); - captureCall->gtArgs.PushFront(m_compiler, - NewCallArg::Primitive(m_compiler->gtNewLclvNode(frame.Resumed, TYP_INT))); + captureCall->gtArgs.PushFront(m_compiler, NewCallArg::Primitive(m_compiler->gtCloneExpr(frameResumed))); m_compiler->compCurBB = tailBB; m_compiler->fgMorphTree(captureCall); @@ -3208,13 +3233,9 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* // Then restore the caller's contexts, as its physical frame's return would have. GenTreeCall* restoreCall = m_compiler->gtNewUserCallNode(m_asyncInfo->restoreContextsOnSuspensionMethHnd, TYP_VOID); - restoreCall->gtArgs.PushFront(m_compiler, - NewCallArg::Primitive( - m_compiler->gtNewLclvNode(outer.SynchronizationContext, TYP_REF))); - restoreCall->gtArgs.PushFront(m_compiler, NewCallArg::Primitive( - m_compiler->gtNewLclvNode(outer.ExecutionContext, TYP_REF))); - restoreCall->gtArgs.PushFront(m_compiler, - NewCallArg::Primitive(m_compiler->gtNewLclvNode(outer.Resumed, TYP_INT))); + restoreCall->gtArgs.PushFront(m_compiler, NewCallArg::Primitive(m_compiler->gtCloneExpr(outerSync))); + restoreCall->gtArgs.PushFront(m_compiler, NewCallArg::Primitive(m_compiler->gtCloneExpr(outerExec))); + restoreCall->gtArgs.PushFront(m_compiler, NewCallArg::Primitive(m_compiler->gtCloneExpr(outerResumed))); m_compiler->compCurBB = tailBB; m_compiler->fgMorphTree(restoreCall); @@ -3223,7 +3244,7 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* DISPRANGE(LIR::AsRange(tailBB)); - m_inlinedFrameTails->Set(key, tailBB); + m_inlinedFrameTails->push_back({numValues, values, tailBB}); return tailBB; } diff --git a/src/coreclr/jit/async.h b/src/coreclr/jit/async.h index 45a9dd5648d286..f99e9ae61d6220 100644 --- a/src/coreclr/jit/async.h +++ b/src/coreclr/jit/async.h @@ -560,10 +560,19 @@ class AsyncTransformation const ContinuationLayout* CreateResumptionsAndSuspensions(ArrayStack& continuationMemberOffsets); void CreateResumptionSwitch(GenTreeLclVarCommon* commonAsyncResumedDef); - typedef JitHashTable, BasicBlock*> InlinedFrameTailMap; - // Suspension tails handling the frames enclosing an inlined async frame, keyed by the - // innermost frame's resumed indicator local. Shared by all suspensions in that frame. - InlinedFrameTailMap* m_inlinedFrameTails = nullptr; + struct InlinedFrameTail + { + // Values the tail stores, in the order they were taken off the call: one + // resumed/exec/sync triple per frame, innermost first. + unsigned NumValues; + GenTree** Values; + BasicBlock* Block; + }; + + // Suspension tails running the context handling of the frames enclosing an inlined + // async frame. Keyed by the values they store, so suspensions that agree on all of + // them share one. + jitstd::vector* m_inlinedFrameTails = nullptr; BasicBlock* CreateInlinedFrameSuspensionTail(BasicBlock* callBlock, GenTreeCall* call, diff --git a/src/coreclr/jit/fginline.cpp b/src/coreclr/jit/fginline.cpp index b8a51caee04193..55644e709f889c 100644 --- a/src/coreclr/jit/fginline.cpp +++ b/src/coreclr/jit/fginline.cpp @@ -2462,8 +2462,14 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) return; } - // AsyncFrameLocals is declared alongside AsyncCallInfo, which is where the chain is - // recorded for the async transformation. + // Async context locals of one inlined frame. + struct AsyncFrameLocals + { + unsigned Resumed; + unsigned ExecutionContext; + unsigned SynchronizationContext; + }; + ArrayStack frames(getAllocator(CMK_Async)); // The inlinee's own frame comes first: the suspension needs its resumed indicator to @@ -2484,12 +2490,6 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) // The root method's frame is always the outermost one. frames.Push({lvaResumedIndicator, lvaAsyncExecutionContextVar, lvaAsyncSynchronizationContextVar}); - AsyncFrameLocals* const frameLocals = new (this, CMK_Async) AsyncFrameLocals[frames.Height()]; - for (int i = 0; i < frames.Height(); i++) - { - frameLocals[i] = frames.Bottom(i); - } - unsigned const inlineeResumed = InlineeCompiler->lvaResumedIndicator; struct Visitor : GenTreeVisitor @@ -2500,16 +2500,11 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) }; ArrayStack& m_frames; - AsyncFrameLocals* m_frameLocals; unsigned m_inlineeResumed; - Visitor(Compiler* comp, - ArrayStack& frames, - AsyncFrameLocals* frameLocals, - unsigned inlineeResumed) + Visitor(Compiler* comp, ArrayStack& frames, unsigned inlineeResumed) : GenTreeVisitor(comp) , m_frames(frames) - , m_frameLocals(frameLocals) , m_inlineeResumed(inlineeResumed) { } @@ -2561,13 +2556,12 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) // The chain holds this call's own frame followed by its enclosing frames, // ending with the root, so its depth in that numbering is one less than the // number of entries. - call->GetAsyncInfo().InlineFrameDepth = (unsigned)m_frames.Height() - 1; - call->GetAsyncInfo().InlineFrameLocals = m_frameLocals; + call->GetAsyncInfo().InlineFrameDepth = (unsigned)m_frames.Height() - 1; return WALK_CONTINUE; } }; - Visitor visitor(this, frames, frameLocals, inlineeResumed); + Visitor visitor(this, frames, inlineeResumed); for (BasicBlock* block = InlineeCompiler->fgFirstBB; block != nullptr; block = block->Next()) { for (Statement* const stmt : block->Statements()) diff --git a/src/coreclr/jit/gentree.h b/src/coreclr/jit/gentree.h index cfeeef6a7c4901..14764c4fdc74b5 100644 --- a/src/coreclr/jit/gentree.h +++ b/src/coreclr/jit/gentree.h @@ -4529,14 +4529,6 @@ enum class ContinuationContextHandling }; // Additional async call info. -// Async context locals of one inlined frame. -struct AsyncFrameLocals -{ - unsigned Resumed; - unsigned ExecutionContext; - unsigned SynchronizationContext; -}; - struct AsyncCallInfo { // DebugInfo with SOURCE_TYPE_ASYNC pointing at the await call IL instruction @@ -4571,15 +4563,6 @@ struct AsyncCallInfo // the call's chain of context args. unsigned InlineFrameDepth = 0; - // Async context locals of this call's own inlined frame followed by its enclosing - // frames, ending with the root method's. Null when the call is not inside an inlined - // async frame; otherwise it holds InlineFrameDepth + 1 entries. - // - // The matching context args on the call keep these values live and reported up to the - // async transformation. This records which locals they are, since optimizations may - // rewrite the arg nodes themselves. - AsyncFrameLocals* InlineFrameLocals = nullptr; - bool NeedsToSaveAndRestoreExecutionContext() const { return true; diff --git a/src/coreclr/jit/importercalls.cpp b/src/coreclr/jit/importercalls.cpp index 3884e157db7f5c..bfff818d6992c3 100644 --- a/src/coreclr/jit/importercalls.cpp +++ b/src/coreclr/jit/importercalls.cpp @@ -7611,8 +7611,7 @@ void Compiler::impInheritAsyncContextsFromInliner(GenTreeCall* call) if (call->IsAsync()) { - call->GetAsyncInfo().InlineFrameDepth = inlCall->GetAsyncInfo().InlineFrameDepth; - call->GetAsyncInfo().InlineFrameLocals = inlCall->GetAsyncInfo().InlineFrameLocals; + call->GetAsyncInfo().InlineFrameDepth = inlCall->GetAsyncInfo().InlineFrameDepth; } } From 3447a52ede6964a1fa8380fde77d4c5e79c9b699 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Thu, 30 Jul 2026 17:28:26 +0200 Subject: [PATCH 33/41] JIT: Remove sharing of inlined frame suspension tails Each suspension now gets its own tail. Sharing required the suspensions to agree on every value the tail stores, which in practice they do not: the first await in a frame typically has its resumed indicator folded to a constant while later ones still read the local, so the tails were never actually shared. Left as a TODO; it is an optimization that saves cold code and can come back separately. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 55 +++++++-------------------------------- src/coreclr/jit/async.h | 14 ---------- 2 files changed, 10 insertions(+), 59 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index a5fb7fefc9475b..c50e7576a56431 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -3078,8 +3078,8 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc //------------------------------------------------------------------------ // AsyncTransformation::CreateInlinedFrameSuspensionTail: -// Create, or reuse, the block that runs the context handling of the frames enclosing -// an inlined async frame on suspension. +// Create the block that runs the context handling of the frames enclosing an inlined +// async frame on suspension. // // Parameters: // callBlock - The block containing the async call @@ -3097,8 +3097,8 @@ void AsyncTransformation::FinishContextHandlingAndSuspensionWithHelper(BasicBloc // walked outward as a straight line: once one frame has resumed, the helpers for it and // every frame outside it no-op. // -// The result depends only on the chain of frames, so it is shared by all suspensions in -// the same frame. +// TODO-Async: suspensions in the same frame could share a tail when they agree on all +// of the values below, which would save some cold code. // BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* callBlock, GenTreeCall* call, @@ -3120,8 +3120,7 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* // truth for what to store: the optimizer may have rewritten them since they were // added, for example folding a "resumed" indicator to a constant, and the stores must // reflect that rather than re-reading the locals the args originally named. - unsigned const numValues = numFrames * 3; - GenTree** const values = new (m_compiler, CMK_Async) GenTree*[numValues]; + ArrayStack values(m_compiler->getAllocator(CMK_Async)); for (unsigned i = 0; i < numFrames; i++) { @@ -3145,40 +3144,7 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* LIR::AsRange(callBlock).Remove(node); call->gtArgs.RemoveUnsafe(arg); - values[i * 3 + j] = node; - } - } - - // The tail depends only on these values, so suspensions that agree on all of them can - // share one. That is the common case: every suspension in the same inlined frame sees - // the same locals unless the optimizer rewrote one of them differently. - if (m_inlinedFrameTails == nullptr) - { - m_inlinedFrameTails = - new (m_compiler, CMK_Async) jitstd::vector(m_compiler->getAllocator(CMK_Async)); - } - - for (const InlinedFrameTail& tail : *m_inlinedFrameTails) - { - if (tail.NumValues != numValues) - { - continue; - } - - bool match = true; - for (unsigned i = 0; i < numValues; i++) - { - if (!GenTree::Compare(tail.Values[i], values[i])) - { - match = false; - break; - } - } - - if (match) - { - JITDUMP(" Reusing inlined frame suspension tail " FMT_BB "\n", tail.Block->bbNum); - return tail.Block; + values.Push(node); } } @@ -3202,10 +3168,10 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* for (unsigned i = 0; i + 1 < numFrames; i++) { - GenTree* const frameResumed = values[i * 3]; - GenTree* const outerResumed = values[(i + 1) * 3]; - GenTree* const outerExec = values[(i + 1) * 3 + 1]; - GenTree* const outerSync = values[(i + 1) * 3 + 2]; + GenTree* const frameResumed = values.Bottom((int)(i * 3)); + GenTree* const outerResumed = values.Bottom((int)((i + 1) * 3)); + GenTree* const outerExec = values.Bottom((int)((i + 1) * 3 + 1)); + GenTree* const outerSync = values.Bottom((int)((i + 1) * 3 + 2)); unsigned const depth = numFrames - 1 - i; // Capture what this frame hands to its caller. @@ -3244,7 +3210,6 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock* DISPRANGE(LIR::AsRange(tailBB)); - m_inlinedFrameTails->push_back({numValues, values, tailBB}); return tailBB; } diff --git a/src/coreclr/jit/async.h b/src/coreclr/jit/async.h index f99e9ae61d6220..0871050776bb95 100644 --- a/src/coreclr/jit/async.h +++ b/src/coreclr/jit/async.h @@ -560,20 +560,6 @@ class AsyncTransformation const ContinuationLayout* CreateResumptionsAndSuspensions(ArrayStack& continuationMemberOffsets); void CreateResumptionSwitch(GenTreeLclVarCommon* commonAsyncResumedDef); - struct InlinedFrameTail - { - // Values the tail stores, in the order they were taken off the call: one - // resumed/exec/sync triple per frame, innermost first. - unsigned NumValues; - GenTree** Values; - BasicBlock* Block; - }; - - // Suspension tails running the context handling of the frames enclosing an inlined - // async frame. Keyed by the values they store, so suspensions that agree on all of - // them share one. - jitstd::vector* m_inlinedFrameTails = nullptr; - BasicBlock* CreateInlinedFrameSuspensionTail(BasicBlock* callBlock, GenTreeCall* call, const ContinuationLayout& layout); From 036133c234d6703683d66a8e8ecde30a76bb1b3e Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 31 Jul 2026 14:04:32 +0200 Subject: [PATCH 34/41] JIT: Emit the inlined async frame transition as calls rather than IR The post-inline IR for an inlined async frame returning to its caller was a full diamond: an ExecutionContext restore, an IsOnRightContext call with its own conditional, and a SwitchContext await, spread over four blocks and relying on inlining the two helpers back in. Only the resumed check is worth having as IR. It is what guards the synchronous path, which is the one to optimize for; anything past it has already suspended and resumed at least once, so a call is cheap there. Fold IsOnRightContext and SwitchContext into a single AsyncHelpers.RestoreInlinedFrameContinuationContext await that does the check internally, leaving two blocks and no inline candidates. The ExecutionContext restore stays a separate non-async call. It cannot move into the same helper: a runtime async method restores the contexts it captured on entry when it returns, which would undo it on the non-suspending path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../coreclr/jit/runtime-async-inlining.md | 8 + .../CompilerServices/AsyncHelpers.CoreCLR.cs | 26 +++- src/coreclr/inc/corinfo.h | 9 +- src/coreclr/inc/jiteeversionguid.h | 10 +- src/coreclr/jit/compiler.h | 3 +- src/coreclr/jit/fginline.cpp | 143 +++++------------- .../tools/Common/JitInterface/CorInfoImpl.cs | 3 +- .../tools/Common/JitInterface/CorInfoTypes.cs | 3 +- .../tools/superpmi/superpmi-shared/agnostic.h | 3 +- .../superpmi-shared/methodcontext.cpp | 6 +- src/coreclr/vm/corelib.h | 3 +- src/coreclr/vm/jitinterface.cpp | 3 +- 12 files changed, 82 insertions(+), 138 deletions(-) diff --git a/docs/design/coreclr/jit/runtime-async-inlining.md b/docs/design/coreclr/jit/runtime-async-inlining.md index dcc1293d2ec6fb..cb69d6495105f2 100644 --- a/docs/design/coreclr/jit/runtime-async-inlining.md +++ b/docs/design/coreclr/jit/runtime-async-inlining.md @@ -170,6 +170,14 @@ if (resumed_C && !AsyncHelpers.IsOnRightContext(continuation.ContinuationContext We expect `IsOnRightContext` to be true almost always, since it is very rare that the current synchronization context is modified, especially during the execution of async methods. If this was a common case then it is unlikely that inlining would ever be profitable. +The implementation expands only the `resumed_C` check as JIT IR and leaves the rest as two calls +(`AsyncHelpers.RestoreInlinedFrameExecutionContext` and `AsyncHelpers.RestoreInlinedFrameContinuationContext`, +the latter doing the `IsOnRightContext` check and the suspension itself internally). +The synchronous case is the one worth optimizing for, and anything past the `resumed_C` check has +already paid for at least one suspension and resumption. +The `ExecutionContext` restore has to be a separate, non-async call: a runtime async method restores +the contexts it captured on entry when it returns, which would undo it. + ### Handling synchronous saves and restores of contexts Recall that every runtime async function's body was wrapped with a save/restore of the contexts. diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index 835880f39325b8..15905b48f9575a 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -1492,8 +1492,8 @@ private static void CaptureContinuationContextFlags(ref ContinuationFlags flags, // Check whether the current thread already satisfies the continuation context captured in a // continuation, i.e. whether resuming that continuation here would be dispatched inline. // - // Used by the JIT when inlining runtime async calls: when an inlined callee logically returns - // to its caller after having been resumed, the JIT must ensure it is running in the continuation + // Used when inlining runtime async calls: when an inlined callee logically returns + // to its caller after having been resumed, we must ensure we are running in the continuation // context the caller's continuation captured. This mirrors the "can inline" conditions in // RuntimeAsyncTaskContinuation.QueueIfNecessary; the two must be kept in sync. private static bool IsOnRightContext(object? continuationContext, ContinuationFlags flags) @@ -1532,30 +1532,40 @@ private static bool IsOnRightContext(object? continuationContext, ContinuationFl // Unlike the synchronous restore at the end of a method, this runs only after a // resumption, so it must target the thread we were resumed on rather than the one // whose contexts were captured on entry. + // + // This deliberately is not a runtime async method: such a method restores the contexts + // it captured on entry when it returns, which would undo the restore done here. private static void RestoreInlinedFrameExecutionContext(ExecutionContext? previousExecCtx) { RestoreExecutionContext(Thread.CurrentThreadAssumedInitialized, previousExecCtx); } - // Suspend and resume in the specified continuation context. + // Get back onto the continuation context that an inlined async frame's caller captured, + // if we are not already on it. // // Used by the JIT when inlining runtime async calls: when an inlined callee logically - // returns to its caller after having been resumed, and IsOnRightContext reports that the - // current context does not match the one the caller's continuation captured, we must get - // back onto that context before continuing. + // returns to its caller after having been resumed, we must continue on the context the + // caller's continuation captured. The JIT emits the check of whether the frame was + // resumed at all and calls this when it was; the far less likely check of whether the + // context actually differs is left to this method. // // 'flags' must contain only ContinuationFlags.AllContinuationFlags bits. // // Suspending on an already completed task makes the dispatcher re-dispatch the continuation // immediately. Because that dispatch happens with canInline: false, it always posts or // schedules onto the requested context rather than running inline here, which is what we - // want -- the JIT only calls this when we are known to be on the wrong context. + // want -- we only suspend when we are known to be on the wrong context. [BypassReadyToRun] [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)] - private static unsafe void SwitchContext(object? continuationContext, ContinuationFlags flags) + private static unsafe void RestoreInlinedFrameContinuationContext(object? continuationContext, ContinuationFlags flags) { Debug.Assert((flags & ~ContinuationFlags.AllContinuationFlags) == 0); + if (IsOnRightContext(continuationContext, flags)) + { + return; + } + ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState; Continuation? sentinelContinuation = state.SentinelContinuation ??= new Continuation(); diff --git a/src/coreclr/inc/corinfo.h b/src/coreclr/inc/corinfo.h index 7f65421eb61ffd..93d69a2bfd6156 100644 --- a/src/coreclr/inc/corinfo.h +++ b/src/coreclr/inc/corinfo.h @@ -1839,12 +1839,9 @@ struct CORINFO_ASYNC_INFO // Method handle for AsyncHelpers.RestoreInlinedFrameExecutionContext, used when an inlined // async callee logically returns to its caller after having been resumed CORINFO_METHOD_HANDLE restoreInlinedFrameExecutionContextMethHnd; - // Method handle for AsyncHelpers.IsOnRightContext, used to check whether an inlined async callee - // can logically return to its caller without switching continuation context - CORINFO_METHOD_HANDLE isOnRightContextMethHnd; - // Method handle for AsyncHelpers.SwitchContext, used to switch to the continuation context of an - // inlined async callee's caller when IsOnRightContext reports a mismatch - CORINFO_METHOD_HANDLE switchContextMethHnd; + // Method handle for AsyncHelpers.RestoreInlinedFrameContinuationContext, used to get back onto + // the continuation context of an inlined async callee's caller + CORINFO_METHOD_HANDLE restoreInlinedFrameContinuationContextMethHnd; // Method handle for AsyncHelpers.CaptureInlinedFrameTransition, used on suspension to capture // the contexts each inlined async frame hands to its caller CORINFO_METHOD_HANDLE captureInlinedFrameTransitionMethHnd; diff --git a/src/coreclr/inc/jiteeversionguid.h b/src/coreclr/inc/jiteeversionguid.h index 302c29cd9b26bf..57f6efdbb05087 100644 --- a/src/coreclr/inc/jiteeversionguid.h +++ b/src/coreclr/inc/jiteeversionguid.h @@ -37,11 +37,11 @@ #include -constexpr GUID JITEEVersionIdentifier = { /* e5c01024-ef09-46e0-8222-62f13997551e */ - 0xe5c01024, - 0xef09, - 0x46e0, - {0x82, 0x22, 0x62, 0xf1, 0x39, 0x97, 0x55, 0x1e} +constexpr GUID JITEEVersionIdentifier = { /* 0e083167-f348-426d-b207-6de9f15b8914 */ + 0x0e083167, + 0xf348, + 0x426d, + {0xb2, 0x07, 0x6d, 0xe9, 0xf1, 0x5b, 0x89, 0x14} }; #endif // JIT_EE_VERSIONING_GUID_H diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h index d583583b59e974..fc8a8fbb625bfb 100644 --- a/src/coreclr/jit/compiler.h +++ b/src/coreclr/jit/compiler.h @@ -7535,8 +7535,7 @@ class Compiler void fgInlineAppendStatements(InlineInfo* inlineInfo, BasicBlock* block, Statement* stmt); void fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicBlock* joinBlock); void fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo); - void fgSetupAsyncFrameSwitchCall(GenTreeCall* call, InlineContext* inlineeContext, const DebugInfo& di); - void fgMarkAsyncHelperInlineCandidate(GenTreeCall* call); + void fgSetupAsyncFrameTransitionCall(GenTreeCall* call, InlineContext* inlineeContext, const DebugInfo& di); GenTree* gtNewContinuationMemberIndir(const struct ContinuationMember& member, var_types type); #ifdef DEBUG diff --git a/src/coreclr/jit/fginline.cpp b/src/coreclr/jit/fginline.cpp index 55644e709f889c..07dad23f8c5664 100644 --- a/src/coreclr/jit/fginline.cpp +++ b/src/coreclr/jit/fginline.cpp @@ -2577,8 +2577,8 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) } //------------------------------------------------------------------------ -// fgSetupAsyncFrameSwitchCall: Turn a call to AsyncHelpers.SwitchContext into a proper -// async call. +// fgSetupAsyncFrameTransitionCall: Turn a call to +// AsyncHelpers.RestoreInlinedFrameContinuationContext into a proper async call. // // Arguments: // call - the call, with its user arguments already added @@ -2586,17 +2586,16 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) // di - debug info to use // // Notes: -// The switch is an await: it always suspends, and gets resumed on the requested -// context. It runs after the inlinee's frame has logically returned, so the frame it -// belongs to is the inlinee's caller. +// The call is an await: it suspends when it has to switch continuation context. It runs +// after the inlinee's frame has logically returned, so the frame it belongs to is the +// inlinee's caller. // -void Compiler::fgSetupAsyncFrameSwitchCall(GenTreeCall* call, InlineContext* inlineeContext, const DebugInfo& di) +void Compiler::fgSetupAsyncFrameTransitionCall(GenTreeCall* call, InlineContext* inlineeContext, const DebugInfo& di) { AsyncCallInfo asyncInfo; - // SwitchContext establishes the continuation context itself, so no further handling - // is attached to this await. + // The helper resumes on the requested context, so no further handling is attached to + // this await. asyncInfo.ContinuationContextHandling = ContinuationContextHandling::None; - asyncInfo.AlwaysSuspends = true; asyncInfo.CallAsyncDebugInfo = di; call->SetIsAsync(new (this, CMK_Async) AsyncCallInfo(asyncInfo)); @@ -2674,21 +2673,6 @@ GenTree* Compiler::gtNewContinuationMemberIndir(const ContinuationMember& member return indir; } -//------------------------------------------------------------------------ -// fgMarkAsyncHelperInlineCandidate: Mark a call to an async helper as an inline -// candidate. -// -// Arguments: -// call - the call -// -void Compiler::fgMarkAsyncHelperInlineCandidate(GenTreeCall* call) -{ - CORINFO_CALL_INFO callInfo = {}; - callInfo.hMethod = call->gtCallMethHnd; - callInfo.methodFlags = info.compCompHnd->getMethodAttribs(callInfo.hMethod); - impMarkInlineCandidate(call, MAKE_METHODCONTEXT(callInfo.hMethod), &callInfo, compInlineContext); -} - //------------------------------------------------------------------------ // fgInlineAppendAsyncFrameStatements: Emit the IR that handles an inlined async // frame logically returning to its caller. @@ -2705,16 +2689,22 @@ void Compiler::fgMarkAsyncHelperInlineCandidate(GenTreeCall* call) // // if (resumed_F) // { -// RestoreInlinedFrameExecutionContext(continuation.ExecutionContextFor); -// if (!IsOnRightContext(continuation.ContinuationContextFor, -// continuation.FlagsFor)) -// { -// await SwitchContext(continuation.ContinuationContextFor, -// continuation.FlagsFor); -// } +// AsyncHelpers.RestoreInlinedFrameExecutionContext(continuation.ExecutionContextFor); +// await AsyncHelpers.RestoreInlinedFrameContinuationContext(continuation.ContinuationContextFor, +// continuation.FlagsFor); // resumed_caller = true; // } // +// Only the check is expanded as IR; the restores are left as plain calls. The +// synchronous case is the one worth optimizing, and by the time we get past the check +// we have already suspended and resumed at least once, so the calls are cheap. In +// particular the "are we already on the right context?" test, which is what actually +// decides whether we suspend, stays inside the helper. +// +// The ExecutionContext restore cannot be folded into the same helper: a runtime async +// method restores the contexts it captured on entry when it returns, which would undo +// it. It is a separate, non-async call for that reason. +// // The members are read off the continuation this method was resumed with. That is // only valid when resumed_F is true, which is exactly the guard: resumed_F can only // become true at a resumption point belonging to this inlinee's own async calls, and @@ -2768,12 +2758,7 @@ void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicB // joinBlock is now empty and jumps to restBlock. Turn it into the "did we resume // inside the inlinee?" test. BasicBlock* const restoreBlock = fgNewBBafter(BBJ_ALWAYS, joinBlock, /* extendRegion */ true); - BasicBlock* const switchBlock = fgNewBBafter(BBJ_ALWAYS, restoreBlock, /* extendRegion */ true); - BasicBlock* const doneBlock = fgNewBBafter(BBJ_ALWAYS, switchBlock, /* extendRegion */ true); - restoreBlock->inheritWeightPercentage(joinBlock, 0); - switchBlock->inheritWeightPercentage(joinBlock, 0); - doneBlock->inheritWeightPercentage(joinBlock, 0); // joinBlock: if (resumed_F == 0) goto restBlock; else goto restoreBlock { @@ -2791,86 +2776,38 @@ void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicB toRestore->setLikelihood(0.0); } - CORINFO_ASYNC_INFO* const asyncInfo = eeGetAsyncInfo(); - - // restoreBlock: restore the ExecutionContext, then check the continuation context. + // restoreBlock: restore the contexts the caller's continuation captured, then record + // that the caller's frame has observed a resumption too. { - GenTreeCall* const restoreCall = + CORINFO_ASYNC_INFO* const asyncInfo = eeGetAsyncInfo(); + + GenTreeCall* const restoreExecCtxCall = gtNewUserCallNode(asyncInfo->restoreInlinedFrameExecutionContextMethHnd, TYP_VOID); - restoreCall->gtArgs.PushFront(this, - NewCallArg::Primitive( - gtNewContinuationMemberIndir(ContinuationMember::InlineFrameExecutionContext( - inlineDepth), - TYP_REF))); - fgMarkAsyncHelperInlineCandidate(restoreCall); - fgInsertStmtAtEnd(restoreBlock, gtNewStmt(restoreCall)); - - GenTreeCall* const isOnRightContextCall = gtNewUserCallNode(asyncInfo->isOnRightContextMethHnd, TYP_UBYTE); - isOnRightContextCall->gtArgs + restoreExecCtxCall->gtArgs .PushFront(this, NewCallArg::Primitive( - gtNewContinuationMemberIndir(ContinuationMember::InlineFrameFlags(inlineDepth), TYP_INT))); - isOnRightContextCall->gtArgs - .PushFront(this, - NewCallArg::Primitive( - gtNewContinuationMemberIndir(ContinuationMember::InlineFrameContinuationContext(inlineDepth), + gtNewContinuationMemberIndir(ContinuationMember::InlineFrameExecutionContext(inlineDepth), TYP_REF))); - fgMarkAsyncHelperInlineCandidate(isOnRightContextCall); + fgInsertStmtAtEnd(restoreBlock, gtNewStmt(restoreExecCtxCall)); - // Inline candidates must be statement roots; the value is consumed through a - // GT_RET_EXPR placeholder that the inliner substitutes. The candidate info has to - // point back at it so that a failed inline can put the call back. - GenTree* isOnRightContext; - if (isOnRightContextCall->IsInlineCandidate()) - { - fgInsertStmtAtEnd(restoreBlock, gtNewStmt(isOnRightContextCall)); - GenTreeRetExpr* const retExpr = gtNewInlineCandidateReturnExpr(isOnRightContextCall, TYP_UBYTE); - isOnRightContextCall->GetSingleInlineCandidateInfo()->retExpr = retExpr; - isOnRightContext = retExpr; - } - else - { - isOnRightContext = isOnRightContextCall; - } - - GenTree* const isWrongContext = gtNewOperNode(GT_EQ, TYP_INT, isOnRightContext, gtNewIconNode(0)); - GenTree* const jtrue = gtNewOperNode(GT_JTRUE, TYP_VOID, isWrongContext); - fgInsertStmtAtEnd(restoreBlock, gtNewStmt(jtrue)); - - FlowEdge* const toSwitch = fgAddRefPred(switchBlock, restoreBlock); - FlowEdge* const toDone = fgAddRefPred(doneBlock, restoreBlock); - restoreBlock->SetCond(toSwitch, toDone); - // We expect to already be on the right context in almost all cases; if we were - // not, inlining would be unlikely to be profitable in the first place. - toSwitch->setLikelihood(0.0); - toDone->setLikelihood(1.0); - } - - // switchBlock: get back onto the continuation context the caller's continuation - // captured. This is an await, so it may suspend. - { - GenTreeCall* const switchCall = gtNewUserCallNode(asyncInfo->switchContextMethHnd, TYP_VOID); - switchCall->gtArgs.PushFront(this, - NewCallArg::Primitive( - gtNewContinuationMemberIndir(ContinuationMember::InlineFrameFlags(inlineDepth), - TYP_INT))); - switchCall->gtArgs + GenTreeCall* const restoreContinuationCtxCall = + gtNewUserCallNode(asyncInfo->restoreInlinedFrameContinuationContextMethHnd, TYP_VOID); + restoreContinuationCtxCall->gtArgs + .PushFront(this, + NewCallArg::Primitive( + gtNewContinuationMemberIndir(ContinuationMember::InlineFrameFlags(inlineDepth), TYP_INT))); + restoreContinuationCtxCall->gtArgs .PushFront(this, NewCallArg::Primitive( gtNewContinuationMemberIndir(ContinuationMember::InlineFrameContinuationContext(inlineDepth), TYP_REF))); - fgSetupAsyncFrameSwitchCall(switchCall, inlineContext, di); - fgInsertStmtAtEnd(switchBlock, gtNewStmt(switchCall)); - - switchBlock->SetKindAndTargetEdge(BBJ_ALWAYS, fgAddRefPred(doneBlock, switchBlock)); - } + fgSetupAsyncFrameTransitionCall(restoreContinuationCtxCall, inlineContext, di); + fgInsertStmtAtEnd(restoreBlock, gtNewStmt(restoreContinuationCtxCall)); - // doneBlock: the caller's frame has now observed a resumption too. - { GenTree* const store = gtNewStoreLclVarNode(resumedCaller, gtNewIconNode(1)); - fgInsertStmtAtEnd(doneBlock, gtNewStmt(store)); + fgInsertStmtAtEnd(restoreBlock, gtNewStmt(store)); - doneBlock->SetKindAndTargetEdge(BBJ_ALWAYS, fgAddRefPred(restBlock, doneBlock)); + restoreBlock->SetKindAndTargetEdge(BBJ_ALWAYS, fgAddRefPred(restBlock, restoreBlock)); } JITDUMPEXEC(fgDispBasicBlocks(joinBlock, restBlock, true)); diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs index c8d476a0716ff0..9ee9b167680251 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs @@ -3651,8 +3651,7 @@ private void getAsyncInfo(ref CORINFO_ASYNC_INFO pAsyncInfoOut) pAsyncInfoOut.restoreContextsMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreContexts"u8, null)); pAsyncInfoOut.restoreContextsOnSuspensionMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreContextsOnSuspension"u8, null)); pAsyncInfoOut.restoreInlinedFrameExecutionContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreInlinedFrameExecutionContext"u8, null)); - pAsyncInfoOut.isOnRightContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("IsOnRightContext"u8, null)); - pAsyncInfoOut.switchContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("SwitchContext"u8, null)); + pAsyncInfoOut.restoreInlinedFrameContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreInlinedFrameContinuationContext"u8, null)); pAsyncInfoOut.captureInlinedFrameTransitionMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("CaptureInlinedFrameTransition"u8, null)); pAsyncInfoOut.finishSuspensionNoContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionNoContinuationContext"u8, null)); pAsyncInfoOut.finishSuspensionWithContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionWithContinuationContext"u8, null)); diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs index 4370bbcb4d6757..723dbcde23723c 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs @@ -973,8 +973,7 @@ public unsafe struct CORINFO_ASYNC_INFO public CORINFO_METHOD_STRUCT_* restoreContextsMethHnd; public CORINFO_METHOD_STRUCT_* restoreContextsOnSuspensionMethHnd; public CORINFO_METHOD_STRUCT_* restoreInlinedFrameExecutionContextMethHnd; - public CORINFO_METHOD_STRUCT_* isOnRightContextMethHnd; - public CORINFO_METHOD_STRUCT_* switchContextMethHnd; + public CORINFO_METHOD_STRUCT_* restoreInlinedFrameContinuationContextMethHnd; public CORINFO_METHOD_STRUCT_* captureInlinedFrameTransitionMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionNoContinuationContextMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionWithContinuationContextMethHnd; diff --git a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h index 9f16723fd9a3ff..b17ea30f91ff8e 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h +++ b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h @@ -238,8 +238,7 @@ struct Agnostic_CORINFO_ASYNC_INFO DWORDLONG restoreContextsMethHnd; DWORDLONG restoreContextsOnSuspensionMethHnd; DWORDLONG restoreInlinedFrameExecutionContextMethHnd; - DWORDLONG isOnRightContextMethHnd; - DWORDLONG switchContextMethHnd; + DWORDLONG restoreInlinedFrameContinuationContextMethHnd; DWORDLONG captureInlinedFrameTransitionMethHnd; DWORDLONG finishSuspensionNoContinuationContextMethHnd; DWORDLONG finishSuspensionWithContinuationContextMethHnd; diff --git a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp index 7647921354c0da..85adcb43da5c52 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp +++ b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp @@ -4395,8 +4395,7 @@ void MethodContext::recGetAsyncInfo(const CORINFO_ASYNC_INFO* pAsyncInfo) value.restoreContextsMethHnd = CastHandle(pAsyncInfo->restoreContextsMethHnd); value.restoreContextsOnSuspensionMethHnd = CastHandle(pAsyncInfo->restoreContextsOnSuspensionMethHnd); value.restoreInlinedFrameExecutionContextMethHnd = CastHandle(pAsyncInfo->restoreInlinedFrameExecutionContextMethHnd); - value.isOnRightContextMethHnd = CastHandle(pAsyncInfo->isOnRightContextMethHnd); - value.switchContextMethHnd = CastHandle(pAsyncInfo->switchContextMethHnd); + value.restoreInlinedFrameContinuationContextMethHnd = CastHandle(pAsyncInfo->restoreInlinedFrameContinuationContextMethHnd); value.captureInlinedFrameTransitionMethHnd = CastHandle(pAsyncInfo->captureInlinedFrameTransitionMethHnd); value.finishSuspensionNoContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionNoContinuationContextMethHnd); value.finishSuspensionWithContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionWithContinuationContextMethHnd); @@ -4425,8 +4424,7 @@ void MethodContext::repGetAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->restoreContextsMethHnd = (CORINFO_METHOD_HANDLE)value.restoreContextsMethHnd; pAsyncInfoOut->restoreContextsOnSuspensionMethHnd = (CORINFO_METHOD_HANDLE)value.restoreContextsOnSuspensionMethHnd; pAsyncInfoOut->restoreInlinedFrameExecutionContextMethHnd = (CORINFO_METHOD_HANDLE)value.restoreInlinedFrameExecutionContextMethHnd; - pAsyncInfoOut->isOnRightContextMethHnd = (CORINFO_METHOD_HANDLE)value.isOnRightContextMethHnd; - pAsyncInfoOut->switchContextMethHnd = (CORINFO_METHOD_HANDLE)value.switchContextMethHnd; + pAsyncInfoOut->restoreInlinedFrameContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.restoreInlinedFrameContinuationContextMethHnd; pAsyncInfoOut->captureInlinedFrameTransitionMethHnd = (CORINFO_METHOD_HANDLE)value.captureInlinedFrameTransitionMethHnd; pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionNoContinuationContextMethHnd; pAsyncInfoOut->finishSuspensionWithContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionWithContinuationContextMethHnd; diff --git a/src/coreclr/vm/corelib.h b/src/coreclr/vm/corelib.h index c3062ae2c9cf57..2534d74a019bc7 100644 --- a/src/coreclr/vm/corelib.h +++ b/src/coreclr/vm/corelib.h @@ -733,8 +733,7 @@ DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_CONTEXTS, CaptureContexts, No DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS, RestoreContexts, NoSig) DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS_ON_SUSPENSION, RestoreContextsOnSuspension, NoSig) DEFINE_METHOD(ASYNC_HELPERS, RESTORE_INLINED_FRAME_EXECUTION_CONTEXT, RestoreInlinedFrameExecutionContext, NoSig) -DEFINE_METHOD(ASYNC_HELPERS, IS_ON_RIGHT_CONTEXT, IsOnRightContext, NoSig) -DEFINE_METHOD(ASYNC_HELPERS, SWITCH_CONTEXT, SwitchContext, NoSig) +DEFINE_METHOD(ASYNC_HELPERS, RESTORE_INLINED_FRAME_CONTINUATION_CONTEXT, RestoreInlinedFrameContinuationContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_INLINED_FRAME_TRANSITION, CaptureInlinedFrameTransition, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT, FinishSuspensionNoContinuationContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_WITH_CONTINUATION_CONTEXT, FinishSuspensionWithContinuationContext, NoSig) diff --git a/src/coreclr/vm/jitinterface.cpp b/src/coreclr/vm/jitinterface.cpp index 20dc48471e71dc..3959bb223baba1 100644 --- a/src/coreclr/vm/jitinterface.cpp +++ b/src/coreclr/vm/jitinterface.cpp @@ -10446,8 +10446,7 @@ void CEEInfo::getAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->restoreContextsMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_CONTEXTS)); pAsyncInfoOut->restoreContextsOnSuspensionMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_CONTEXTS_ON_SUSPENSION)); pAsyncInfoOut->restoreInlinedFrameExecutionContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_INLINED_FRAME_EXECUTION_CONTEXT)); - pAsyncInfoOut->isOnRightContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__IS_ON_RIGHT_CONTEXT)); - pAsyncInfoOut->switchContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__SWITCH_CONTEXT)); + pAsyncInfoOut->restoreInlinedFrameContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_INLINED_FRAME_CONTINUATION_CONTEXT)); pAsyncInfoOut->captureInlinedFrameTransitionMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__CAPTURE_INLINED_FRAME_TRANSITION)); pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT)); pAsyncInfoOut->finishSuspensionWithContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_WITH_CONTINUATION_CONTEXT)); From ce2aec6e4b6a64faff5a388741362c8a71137867 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 31 Jul 2026 14:20:44 +0200 Subject: [PATCH 35/41] JIT: Merge the inlined async frame transition into a single helper call RestoreInlinedFrameExecutionContext and RestoreInlinedFrameContinuationContext become one RestoreInlinedFrameContexts await. The ExecutionContext restore can live in the async helper after all: it is a manually marked async method rather than the async variant of a task returning one, so CORINFO_ASYNC_SAVE_CONTEXTS is not set for it and there is no save/restore around its body to undo the restore on the way out. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../coreclr/jit/runtime-async-inlining.md | 8 ++- .../CompilerServices/AsyncHelpers.CoreCLR.cs | 45 ++++++++--------- src/coreclr/inc/corinfo.h | 7 +-- src/coreclr/inc/jiteeversionguid.h | 10 ++-- src/coreclr/jit/fginline.cpp | 49 ++++++++----------- .../tools/Common/JitInterface/CorInfoImpl.cs | 3 +- .../tools/Common/JitInterface/CorInfoTypes.cs | 3 +- .../tools/superpmi/superpmi-shared/agnostic.h | 3 +- .../superpmi-shared/methodcontext.cpp | 6 +-- src/coreclr/vm/corelib.h | 3 +- src/coreclr/vm/jitinterface.cpp | 3 +- 11 files changed, 57 insertions(+), 83 deletions(-) diff --git a/docs/design/coreclr/jit/runtime-async-inlining.md b/docs/design/coreclr/jit/runtime-async-inlining.md index cb69d6495105f2..9d8c84ea6fc6a0 100644 --- a/docs/design/coreclr/jit/runtime-async-inlining.md +++ b/docs/design/coreclr/jit/runtime-async-inlining.md @@ -170,13 +170,11 @@ if (resumed_C && !AsyncHelpers.IsOnRightContext(continuation.ContinuationContext We expect `IsOnRightContext` to be true almost always, since it is very rare that the current synchronization context is modified, especially during the execution of async methods. If this was a common case then it is unlikely that inlining would ever be profitable. -The implementation expands only the `resumed_C` check as JIT IR and leaves the rest as two calls -(`AsyncHelpers.RestoreInlinedFrameExecutionContext` and `AsyncHelpers.RestoreInlinedFrameContinuationContext`, -the latter doing the `IsOnRightContext` check and the suspension itself internally). +The implementation expands only the `resumed_C` check as JIT IR and leaves the rest as a single +`AsyncHelpers.RestoreInlinedFrameContexts` call, which does the `ExecutionContext` restore, the +`IsOnRightContext` check and the suspension itself internally. The synchronous case is the one worth optimizing for, and anything past the `resumed_C` check has already paid for at least one suspension and resumption. -The `ExecutionContext` restore has to be a separate, non-async call: a runtime async method restores -the contexts it captured on entry when it returns, which would undo it. ### Handling synchronous saves and restores of contexts diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index 15905b48f9575a..f9b8aac5e71e17 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -1526,41 +1526,36 @@ private static bool IsOnRightContext(object? continuationContext, ContinuationFl return true; } - // Restore the ExecutionContext that an inlined async frame captured when it logically - // returned to its caller. + // Restore the contexts that an inlined async frame captured when it logically returned to + // its caller, after that frame was resumed inside its own body. // - // Unlike the synchronous restore at the end of a method, this runs only after a - // resumption, so it must target the thread we were resumed on rather than the one - // whose contexts were captured on entry. + // Used by the JIT when inlining runtime async calls. The JIT emits the check of whether + // the frame was resumed at all and calls this when it was; everything the async + // infrastructure would otherwise have done at that frame boundary happens here. // - // This deliberately is not a runtime async method: such a method restores the contexts - // it captured on entry when it returns, which would undo the restore done here. - private static void RestoreInlinedFrameExecutionContext(ExecutionContext? previousExecCtx) - { - RestoreExecutionContext(Thread.CurrentThreadAssumedInitialized, previousExecCtx); - } - - // Get back onto the continuation context that an inlined async frame's caller captured, - // if we are not already on it. + // Unlike the synchronous restore at the end of a method, the ExecutionContext restore runs + // only after a resumption, so it must target the thread we were resumed on rather than the + // one whose contexts were captured on entry. Note that this relies on no context save and + // restore being emitted around this method itself; that would undo the restore on the way + // out. It is not, since this is a manually marked async method and not an async variant of + // a task returning one. // - // Used by the JIT when inlining runtime async calls: when an inlined callee logically - // returns to its caller after having been resumed, we must continue on the context the - // caller's continuation captured. The JIT emits the check of whether the frame was - // resumed at all and calls this when it was; the far less likely check of whether the - // context actually differs is left to this method. + // If we are not already on the continuation context the caller's continuation captured we + // suspend to get back onto it. Suspending on an already completed task makes the dispatcher + // re-dispatch the continuation immediately. Because that dispatch happens with + // canInline: false, it always posts or schedules onto the requested context rather than + // running inline here, which is what we want -- we only suspend when we are known to be on + // the wrong context. // // 'flags' must contain only ContinuationFlags.AllContinuationFlags bits. - // - // Suspending on an already completed task makes the dispatcher re-dispatch the continuation - // immediately. Because that dispatch happens with canInline: false, it always posts or - // schedules onto the requested context rather than running inline here, which is what we - // want -- we only suspend when we are known to be on the wrong context. [BypassReadyToRun] [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)] - private static unsafe void RestoreInlinedFrameContinuationContext(object? continuationContext, ContinuationFlags flags) + private static unsafe void RestoreInlinedFrameContexts(ExecutionContext? previousExecCtx, object? continuationContext, ContinuationFlags flags) { Debug.Assert((flags & ~ContinuationFlags.AllContinuationFlags) == 0); + RestoreExecutionContext(Thread.CurrentThreadAssumedInitialized, previousExecCtx); + if (IsOnRightContext(continuationContext, flags)) { return; diff --git a/src/coreclr/inc/corinfo.h b/src/coreclr/inc/corinfo.h index 93d69a2bfd6156..c92f38cf269b10 100644 --- a/src/coreclr/inc/corinfo.h +++ b/src/coreclr/inc/corinfo.h @@ -1836,12 +1836,9 @@ struct CORINFO_ASYNC_INFO CORINFO_METHOD_HANDLE restoreContextsMethHnd; // Method handle for AsyncHelpers.RestoreContextsOnSuspension, used before suspending in async methods CORINFO_METHOD_HANDLE restoreContextsOnSuspensionMethHnd; - // Method handle for AsyncHelpers.RestoreInlinedFrameExecutionContext, used when an inlined + // Method handle for AsyncHelpers.RestoreInlinedFrameContexts, used when an inlined // async callee logically returns to its caller after having been resumed - CORINFO_METHOD_HANDLE restoreInlinedFrameExecutionContextMethHnd; - // Method handle for AsyncHelpers.RestoreInlinedFrameContinuationContext, used to get back onto - // the continuation context of an inlined async callee's caller - CORINFO_METHOD_HANDLE restoreInlinedFrameContinuationContextMethHnd; + CORINFO_METHOD_HANDLE restoreInlinedFrameContextsMethHnd; // Method handle for AsyncHelpers.CaptureInlinedFrameTransition, used on suspension to capture // the contexts each inlined async frame hands to its caller CORINFO_METHOD_HANDLE captureInlinedFrameTransitionMethHnd; diff --git a/src/coreclr/inc/jiteeversionguid.h b/src/coreclr/inc/jiteeversionguid.h index 57f6efdbb05087..916805184634bb 100644 --- a/src/coreclr/inc/jiteeversionguid.h +++ b/src/coreclr/inc/jiteeversionguid.h @@ -37,11 +37,11 @@ #include -constexpr GUID JITEEVersionIdentifier = { /* 0e083167-f348-426d-b207-6de9f15b8914 */ - 0x0e083167, - 0xf348, - 0x426d, - {0xb2, 0x07, 0x6d, 0xe9, 0xf1, 0x5b, 0x89, 0x14} +constexpr GUID JITEEVersionIdentifier = { /* a41245a2-74c0-47ac-87d7-9375cb833125 */ + 0xa41245a2, + 0x74c0, + 0x47ac, + {0x87, 0xd7, 0x93, 0x75, 0xcb, 0x83, 0x31, 0x25} }; #endif // JIT_EE_VERSIONING_GUID_H diff --git a/src/coreclr/jit/fginline.cpp b/src/coreclr/jit/fginline.cpp index 07dad23f8c5664..ca6acfa7b0795e 100644 --- a/src/coreclr/jit/fginline.cpp +++ b/src/coreclr/jit/fginline.cpp @@ -2577,8 +2577,8 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) } //------------------------------------------------------------------------ -// fgSetupAsyncFrameTransitionCall: Turn a call to -// AsyncHelpers.RestoreInlinedFrameContinuationContext into a proper async call. +// fgSetupAsyncFrameTransitionCall: Turn a call to AsyncHelpers.RestoreInlinedFrameContexts +// into a proper async call. // // Arguments: // call - the call, with its user arguments already added @@ -2689,22 +2689,18 @@ GenTree* Compiler::gtNewContinuationMemberIndir(const ContinuationMember& member // // if (resumed_F) // { -// AsyncHelpers.RestoreInlinedFrameExecutionContext(continuation.ExecutionContextFor); -// await AsyncHelpers.RestoreInlinedFrameContinuationContext(continuation.ContinuationContextFor, -// continuation.FlagsFor); +// await AsyncHelpers.RestoreInlinedFrameContexts(continuation.ExecutionContextFor, +// continuation.ContinuationContextFor, +// continuation.FlagsFor); // resumed_caller = true; // } // -// Only the check is expanded as IR; the restores are left as plain calls. The +// Only the check is expanded as IR; the restore itself is left as a single call. The // synchronous case is the one worth optimizing, and by the time we get past the check -// we have already suspended and resumed at least once, so the calls are cheap. In +// we have already suspended and resumed at least once, so one more call is cheap. In // particular the "are we already on the right context?" test, which is what actually // decides whether we suspend, stays inside the helper. // -// The ExecutionContext restore cannot be folded into the same helper: a runtime async -// method restores the contexts it captured on entry when it returns, which would undo -// it. It is a separate, non-async call for that reason. -// // The members are read off the continuation this method was resumed with. That is // only valid when resumed_F is true, which is exactly the guard: resumed_F can only // become true at a resumption point belonging to this inlinee's own async calls, and @@ -2781,28 +2777,23 @@ void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicB { CORINFO_ASYNC_INFO* const asyncInfo = eeGetAsyncInfo(); - GenTreeCall* const restoreExecCtxCall = - gtNewUserCallNode(asyncInfo->restoreInlinedFrameExecutionContextMethHnd, TYP_VOID); - restoreExecCtxCall->gtArgs - .PushFront(this, - NewCallArg::Primitive( - gtNewContinuationMemberIndir(ContinuationMember::InlineFrameExecutionContext(inlineDepth), - TYP_REF))); - fgInsertStmtAtEnd(restoreBlock, gtNewStmt(restoreExecCtxCall)); - - GenTreeCall* const restoreContinuationCtxCall = - gtNewUserCallNode(asyncInfo->restoreInlinedFrameContinuationContextMethHnd, TYP_VOID); - restoreContinuationCtxCall->gtArgs - .PushFront(this, - NewCallArg::Primitive( - gtNewContinuationMemberIndir(ContinuationMember::InlineFrameFlags(inlineDepth), TYP_INT))); - restoreContinuationCtxCall->gtArgs + GenTreeCall* const restoreCall = gtNewUserCallNode(asyncInfo->restoreInlinedFrameContextsMethHnd, TYP_VOID); + restoreCall->gtArgs.PushFront(this, NewCallArg::Primitive( + gtNewContinuationMemberIndir(ContinuationMember::InlineFrameFlags( + inlineDepth), + TYP_INT))); + restoreCall->gtArgs .PushFront(this, NewCallArg::Primitive( gtNewContinuationMemberIndir(ContinuationMember::InlineFrameContinuationContext(inlineDepth), TYP_REF))); - fgSetupAsyncFrameTransitionCall(restoreContinuationCtxCall, inlineContext, di); - fgInsertStmtAtEnd(restoreBlock, gtNewStmt(restoreContinuationCtxCall)); + restoreCall->gtArgs.PushFront(this, + NewCallArg::Primitive( + gtNewContinuationMemberIndir(ContinuationMember::InlineFrameExecutionContext( + inlineDepth), + TYP_REF))); + fgSetupAsyncFrameTransitionCall(restoreCall, inlineContext, di); + fgInsertStmtAtEnd(restoreBlock, gtNewStmt(restoreCall)); GenTree* const store = gtNewStoreLclVarNode(resumedCaller, gtNewIconNode(1)); fgInsertStmtAtEnd(restoreBlock, gtNewStmt(store)); diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs index 9ee9b167680251..68175209d25a61 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs @@ -3650,8 +3650,7 @@ private void getAsyncInfo(ref CORINFO_ASYNC_INFO pAsyncInfoOut) pAsyncInfoOut.captureContextsMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("CaptureContexts"u8, null)); pAsyncInfoOut.restoreContextsMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreContexts"u8, null)); pAsyncInfoOut.restoreContextsOnSuspensionMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreContextsOnSuspension"u8, null)); - pAsyncInfoOut.restoreInlinedFrameExecutionContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreInlinedFrameExecutionContext"u8, null)); - pAsyncInfoOut.restoreInlinedFrameContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreInlinedFrameContinuationContext"u8, null)); + pAsyncInfoOut.restoreInlinedFrameContextsMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("RestoreInlinedFrameContexts"u8, null)); pAsyncInfoOut.captureInlinedFrameTransitionMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("CaptureInlinedFrameTransition"u8, null)); pAsyncInfoOut.finishSuspensionNoContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionNoContinuationContext"u8, null)); pAsyncInfoOut.finishSuspensionWithContinuationContextMethHnd = ObjectToHandle(asyncHelpers.GetKnownMethod("FinishSuspensionWithContinuationContext"u8, null)); diff --git a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs index 723dbcde23723c..d421e9b5905e58 100644 --- a/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs +++ b/src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs @@ -972,8 +972,7 @@ public unsafe struct CORINFO_ASYNC_INFO public CORINFO_METHOD_STRUCT_* captureContextsMethHnd; public CORINFO_METHOD_STRUCT_* restoreContextsMethHnd; public CORINFO_METHOD_STRUCT_* restoreContextsOnSuspensionMethHnd; - public CORINFO_METHOD_STRUCT_* restoreInlinedFrameExecutionContextMethHnd; - public CORINFO_METHOD_STRUCT_* restoreInlinedFrameContinuationContextMethHnd; + public CORINFO_METHOD_STRUCT_* restoreInlinedFrameContextsMethHnd; public CORINFO_METHOD_STRUCT_* captureInlinedFrameTransitionMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionNoContinuationContextMethHnd; public CORINFO_METHOD_STRUCT_* finishSuspensionWithContinuationContextMethHnd; diff --git a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h index b17ea30f91ff8e..8b8e5e5f39b8aa 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h +++ b/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h @@ -237,8 +237,7 @@ struct Agnostic_CORINFO_ASYNC_INFO DWORDLONG captureContextsMethHnd; DWORDLONG restoreContextsMethHnd; DWORDLONG restoreContextsOnSuspensionMethHnd; - DWORDLONG restoreInlinedFrameExecutionContextMethHnd; - DWORDLONG restoreInlinedFrameContinuationContextMethHnd; + DWORDLONG restoreInlinedFrameContextsMethHnd; DWORDLONG captureInlinedFrameTransitionMethHnd; DWORDLONG finishSuspensionNoContinuationContextMethHnd; DWORDLONG finishSuspensionWithContinuationContextMethHnd; diff --git a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp index 85adcb43da5c52..1fa1e4c1e3204f 100644 --- a/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp +++ b/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp @@ -4394,8 +4394,7 @@ void MethodContext::recGetAsyncInfo(const CORINFO_ASYNC_INFO* pAsyncInfo) value.captureContextsMethHnd = CastHandle(pAsyncInfo->captureContextsMethHnd); value.restoreContextsMethHnd = CastHandle(pAsyncInfo->restoreContextsMethHnd); value.restoreContextsOnSuspensionMethHnd = CastHandle(pAsyncInfo->restoreContextsOnSuspensionMethHnd); - value.restoreInlinedFrameExecutionContextMethHnd = CastHandle(pAsyncInfo->restoreInlinedFrameExecutionContextMethHnd); - value.restoreInlinedFrameContinuationContextMethHnd = CastHandle(pAsyncInfo->restoreInlinedFrameContinuationContextMethHnd); + value.restoreInlinedFrameContextsMethHnd = CastHandle(pAsyncInfo->restoreInlinedFrameContextsMethHnd); value.captureInlinedFrameTransitionMethHnd = CastHandle(pAsyncInfo->captureInlinedFrameTransitionMethHnd); value.finishSuspensionNoContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionNoContinuationContextMethHnd); value.finishSuspensionWithContinuationContextMethHnd = CastHandle(pAsyncInfo->finishSuspensionWithContinuationContextMethHnd); @@ -4423,8 +4422,7 @@ void MethodContext::repGetAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->captureContextsMethHnd = (CORINFO_METHOD_HANDLE)value.captureContextsMethHnd; pAsyncInfoOut->restoreContextsMethHnd = (CORINFO_METHOD_HANDLE)value.restoreContextsMethHnd; pAsyncInfoOut->restoreContextsOnSuspensionMethHnd = (CORINFO_METHOD_HANDLE)value.restoreContextsOnSuspensionMethHnd; - pAsyncInfoOut->restoreInlinedFrameExecutionContextMethHnd = (CORINFO_METHOD_HANDLE)value.restoreInlinedFrameExecutionContextMethHnd; - pAsyncInfoOut->restoreInlinedFrameContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.restoreInlinedFrameContinuationContextMethHnd; + pAsyncInfoOut->restoreInlinedFrameContextsMethHnd = (CORINFO_METHOD_HANDLE)value.restoreInlinedFrameContextsMethHnd; pAsyncInfoOut->captureInlinedFrameTransitionMethHnd = (CORINFO_METHOD_HANDLE)value.captureInlinedFrameTransitionMethHnd; pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionNoContinuationContextMethHnd; pAsyncInfoOut->finishSuspensionWithContinuationContextMethHnd = (CORINFO_METHOD_HANDLE)value.finishSuspensionWithContinuationContextMethHnd; diff --git a/src/coreclr/vm/corelib.h b/src/coreclr/vm/corelib.h index 2534d74a019bc7..774cd84cd32d97 100644 --- a/src/coreclr/vm/corelib.h +++ b/src/coreclr/vm/corelib.h @@ -732,8 +732,7 @@ DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_CONTINUATION_CONTEXT, CaptureContinuat DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_CONTEXTS, CaptureContexts, NoSig) DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS, RestoreContexts, NoSig) DEFINE_METHOD(ASYNC_HELPERS, RESTORE_CONTEXTS_ON_SUSPENSION, RestoreContextsOnSuspension, NoSig) -DEFINE_METHOD(ASYNC_HELPERS, RESTORE_INLINED_FRAME_EXECUTION_CONTEXT, RestoreInlinedFrameExecutionContext, NoSig) -DEFINE_METHOD(ASYNC_HELPERS, RESTORE_INLINED_FRAME_CONTINUATION_CONTEXT, RestoreInlinedFrameContinuationContext, NoSig) +DEFINE_METHOD(ASYNC_HELPERS, RESTORE_INLINED_FRAME_CONTEXTS, RestoreInlinedFrameContexts, NoSig) DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_INLINED_FRAME_TRANSITION, CaptureInlinedFrameTransition, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT, FinishSuspensionNoContinuationContext, NoSig) DEFINE_METHOD(ASYNC_HELPERS, FINISH_SUSPENSION_WITH_CONTINUATION_CONTEXT, FinishSuspensionWithContinuationContext, NoSig) diff --git a/src/coreclr/vm/jitinterface.cpp b/src/coreclr/vm/jitinterface.cpp index 3959bb223baba1..1a1f9fbb290e0a 100644 --- a/src/coreclr/vm/jitinterface.cpp +++ b/src/coreclr/vm/jitinterface.cpp @@ -10445,8 +10445,7 @@ void CEEInfo::getAsyncInfo(CORINFO_ASYNC_INFO* pAsyncInfoOut) pAsyncInfoOut->captureContextsMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__CAPTURE_CONTEXTS)); pAsyncInfoOut->restoreContextsMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_CONTEXTS)); pAsyncInfoOut->restoreContextsOnSuspensionMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_CONTEXTS_ON_SUSPENSION)); - pAsyncInfoOut->restoreInlinedFrameExecutionContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_INLINED_FRAME_EXECUTION_CONTEXT)); - pAsyncInfoOut->restoreInlinedFrameContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_INLINED_FRAME_CONTINUATION_CONTEXT)); + pAsyncInfoOut->restoreInlinedFrameContextsMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__RESTORE_INLINED_FRAME_CONTEXTS)); pAsyncInfoOut->captureInlinedFrameTransitionMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__CAPTURE_INLINED_FRAME_TRANSITION)); pAsyncInfoOut->finishSuspensionNoContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_NO_CONTINUATION_CONTEXT)); pAsyncInfoOut->finishSuspensionWithContinuationContextMethHnd = CORINFO_METHOD_HANDLE(CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINISH_SUSPENSION_WITH_CONTINUATION_CONTEXT)); From 6134f4ebbba460e43a75b738449afd27a9584b1e Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 31 Jul 2026 14:40:41 +0200 Subject: [PATCH 36/41] Make RestoreInlinedFrameContexts inlinable Split the suspension out into a tail awaited SwitchToContinuationContext so that what is left is just the ExecutionContext restore and the continuation context check. Take the thread from the cached one in the await state rather than a second TLS lookup, and inline IsOnRightContext into its only caller. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../CompilerServices/AsyncHelpers.CoreCLR.cs | 109 +++++++++--------- 1 file changed, 57 insertions(+), 52 deletions(-) diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index f9b8aac5e71e17..743dd815fc259b 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -1489,43 +1489,6 @@ private static void CaptureContinuationContextFlags(ref ContinuationFlags flags, flags |= ContinuationFlags.ContinueOnThreadPool; } - // Check whether the current thread already satisfies the continuation context captured in a - // continuation, i.e. whether resuming that continuation here would be dispatched inline. - // - // Used when inlining runtime async calls: when an inlined callee logically returns - // to its caller after having been resumed, we must ensure we are running in the continuation - // context the caller's continuation captured. This mirrors the "can inline" conditions in - // RuntimeAsyncTaskContinuation.QueueIfNecessary; the two must be kept in sync. - private static bool IsOnRightContext(object? continuationContext, ContinuationFlags flags) - { - if ((flags & ContinuationFlags.ContinueOnThreadPool) != 0) - { - SynchronizationContext? syncCtx = Thread.CurrentThreadAssumedInitialized._synchronizationContext; - if (syncCtx != null && syncCtx.GetType() != typeof(SynchronizationContext)) - { - return false; - } - - TaskScheduler? sched = TaskScheduler.InternalCurrent; - return sched is null || sched == TaskScheduler.Default; - } - - if ((flags & ContinuationFlags.ContinueOnCapturedSynchronizationContext) != 0) - { - Debug.Assert(continuationContext is SynchronizationContext); - return (SynchronizationContext)continuationContext! == Thread.CurrentThreadAssumedInitialized._synchronizationContext; - } - - if ((flags & ContinuationFlags.ContinueOnCapturedTaskScheduler) != 0) - { - Debug.Assert(continuationContext is TaskScheduler); - return (TaskScheduler)continuationContext! == TaskScheduler.InternalCurrent; - } - - // No continuation context was captured, so there is nothing to switch to. - return true; - } - // Restore the contexts that an inlined async frame captured when it logically returned to // its caller, after that frame was resumed inside its own body. // @@ -1535,33 +1498,75 @@ private static bool IsOnRightContext(object? continuationContext, ContinuationFl // // Unlike the synchronous restore at the end of a method, the ExecutionContext restore runs // only after a resumption, so it must target the thread we were resumed on rather than the - // one whose contexts were captured on entry. Note that this relies on no context save and - // restore being emitted around this method itself; that would undo the restore on the way - // out. It is not, since this is a manually marked async method and not an async variant of - // a task returning one. + // one whose contexts were captured on entry. // - // If we are not already on the continuation context the caller's continuation captured we - // suspend to get back onto it. Suspending on an already completed task makes the dispatcher - // re-dispatch the continuation immediately. Because that dispatch happens with - // canInline: false, it always posts or schedules onto the requested context rather than - // running inline here, which is what we want -- we only suspend when we are known to be on - // the wrong context. + // The continuation context check determines whether resuming the caller's continuation here + // would be dispatched inline. It mirrors the "can inline" conditions in + // RuntimeAsyncTaskContinuation.QueueIfNecessary; the two must be kept in sync. // // 'flags' must contain only ContinuationFlags.AllContinuationFlags bits. [BypassReadyToRun] - [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)] - private static unsafe void RestoreInlinedFrameContexts(ExecutionContext? previousExecCtx, object? continuationContext, ContinuationFlags flags) + [MethodImpl(MethodImplOptions.Async)] + private static void RestoreInlinedFrameContexts(ExecutionContext? previousExecCtx, object? continuationContext, ContinuationFlags flags) { Debug.Assert((flags & ~ContinuationFlags.AllContinuationFlags) == 0); - RestoreExecutionContext(Thread.CurrentThreadAssumedInitialized, previousExecCtx); + // We are inside a runtime async chain, so the thread has already been cached. Use it + // instead of Thread.CurrentThreadAssumedInitialized to keep this to one TLS lookup. + ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState; + Thread? currentThread = state.CurrentThread; + Debug.Assert(currentThread != null); + + RestoreExecutionContext(currentThread, previousExecCtx); - if (IsOnRightContext(continuationContext, flags)) + if ((flags & ContinuationFlags.ContinueOnThreadPool) != 0) { + SynchronizationContext? syncCtx = currentThread._synchronizationContext; + if (syncCtx is null || syncCtx.GetType() == typeof(SynchronizationContext)) + { + TaskScheduler? sched = TaskScheduler.InternalCurrent; + if (sched is null || sched == TaskScheduler.Default) + { + return; + } + } + } + else if ((flags & ContinuationFlags.ContinueOnCapturedSynchronizationContext) != 0) + { + Debug.Assert(continuationContext is SynchronizationContext); + if ((SynchronizationContext)continuationContext! == currentThread._synchronizationContext) + { + return; + } + } + else if ((flags & ContinuationFlags.ContinueOnCapturedTaskScheduler) != 0) + { + Debug.Assert(continuationContext is TaskScheduler); + if ((TaskScheduler)continuationContext! == TaskScheduler.InternalCurrent) + { + return; + } + } + else + { + // No continuation context was captured, so there is nothing to switch to. return; } - ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState; + TailAwait(); + SwitchToContinuationContext(ref state, continuationContext, flags); + } + + // Suspend and resume in the specified continuation context. + // + // Suspending on an already completed task makes the dispatcher re-dispatch the continuation + // immediately. Because that dispatch happens with canInline: false, it always posts or + // schedules onto the requested context rather than running inline here, which is what we + // want -- we only get here when we are known to be on the wrong context. + [BypassReadyToRun] + [MethodImpl(MethodImplOptions.Async)] + private static unsafe void SwitchToContinuationContext(ref RuntimeAsyncAwaitState state, object? continuationContext, ContinuationFlags flags) + { Continuation? sentinelContinuation = state.SentinelContinuation ??= new Continuation(); RuntimeAsyncTaskContinuation? taskCont = state.CachedTaskContinuation; From 38e42b5ef7f0eca9d466770c49142d809ab6f8d2 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 31 Jul 2026 14:51:13 +0200 Subject: [PATCH 37/41] JIT: Make the inlined async frame transition call an inline candidate The check for whether we are already on the right continuation context is what guards the suspension, so it is worth having in the caller. Mark the call so the inliner can pull it in, and force it with AggressiveInlining since the helper is over the IL size budget. This also needs the call to carry an AsyncResumedDef like every other async call. Without it, inlining a body containing an async call trips the use/def parity assert in impInheritAsyncContextsFromInliner, which propagates both to the inlinee's own async calls. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs | 2 +- src/coreclr/jit/fginline.cpp | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index 743dd815fc259b..84c2150001a1c9 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -1506,7 +1506,7 @@ private static void CaptureContinuationContextFlags(ref ContinuationFlags flags, // // 'flags' must contain only ContinuationFlags.AllContinuationFlags bits. [BypassReadyToRun] - [MethodImpl(MethodImplOptions.Async)] + [MethodImpl(MethodImplOptions.AggressiveInlining | MethodImplOptions.Async)] private static void RestoreInlinedFrameContexts(ExecutionContext? previousExecCtx, object? continuationContext, ContinuationFlags flags) { Debug.Assert((flags & ~ContinuationFlags.AllContinuationFlags) == 0); diff --git a/src/coreclr/jit/fginline.cpp b/src/coreclr/jit/fginline.cpp index ca6acfa7b0795e..363a8e54cede99 100644 --- a/src/coreclr/jit/fginline.cpp +++ b/src/coreclr/jit/fginline.cpp @@ -2637,8 +2637,12 @@ void Compiler::fgSetupAsyncFrameTransitionCall(GenTreeCall* call, InlineContext* .WellKnown(WellKnownArg::AsyncSynchronizationContext)); call->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(execCtx, TYP_REF)) .WellKnown(WellKnownArg::AsyncExecutionContext)); + call->gtArgs + .PushFront(this, NewCallArg::Primitive(gtNewLclAddrNode(resumed, 0)).WellKnown(WellKnownArg::AsyncResumedDef)); call->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(resumed, TYP_INT)) .WellKnown(WellKnownArg::AsyncResumedUse)); + + lvaGetDesc(resumed)->lvHasLdAddrOp = true; } //------------------------------------------------------------------------ @@ -2793,6 +2797,14 @@ void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicB inlineDepth), TYP_REF))); fgSetupAsyncFrameTransitionCall(restoreCall, inlineContext, di); + + // The helper is small on the path that does not suspend, which is the one we care + // about, so let the inliner have a look at it. + CORINFO_CALL_INFO callInfo = {}; + callInfo.hMethod = restoreCall->gtCallMethHnd; + callInfo.methodFlags = info.compCompHnd->getMethodAttribs(callInfo.hMethod); + impMarkInlineCandidate(restoreCall, MAKE_METHODCONTEXT(callInfo.hMethod), &callInfo, compInlineContext); + fgInsertStmtAtEnd(restoreBlock, gtNewStmt(restoreCall)); GenTree* const store = gtNewStoreLclVarNode(resumedCaller, gtNewIconNode(1)); From 6962878ffc9c62db38e11d24cb07fa6af0e428b2 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 31 Jul 2026 15:12:11 +0200 Subject: [PATCH 38/41] Remove AggressiveInlining from RestoreInlinedFrameContexts Whether to inline it should be the JIT's call based on PGO, not forced here. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index 84c2150001a1c9..743dd815fc259b 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -1506,7 +1506,7 @@ private static void CaptureContinuationContextFlags(ref ContinuationFlags flags, // // 'flags' must contain only ContinuationFlags.AllContinuationFlags bits. [BypassReadyToRun] - [MethodImpl(MethodImplOptions.AggressiveInlining | MethodImplOptions.Async)] + [MethodImpl(MethodImplOptions.Async)] private static void RestoreInlinedFrameContexts(ExecutionContext? previousExecCtx, object? continuationContext, ContinuationFlags flags) { Debug.Assert((flags & ~ContinuationFlags.AllContinuationFlags) == 0); From 9263050b8deeccc2adf400dc9cf67f06cdc3c42c Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 31 Jul 2026 15:12:11 +0200 Subject: [PATCH 39/41] JIT: Do not add context pseudo-args to the inlined frame transition call Those args model what a frame captures and restores when a suspension unwinds through it. None of that applies here: we only run at all because the frame we are returning out of was resumed, so every frame in the chain already has a continuation holding what it captured when it first suspended, and a suspension here goes straight back to the dispatcher, which restores the thread's contexts itself. This also drops the AsyncResumedDef added along with the inline candidate marking; with no args at all the use/def parity in impInheritAsyncContextsFromInliner holds trivially and the inlinee's own async calls correctly get no context args either. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/compiler.h | 2 +- src/coreclr/jit/fginline.cpp | 54 ++++++++---------------------------- 2 files changed, 13 insertions(+), 43 deletions(-) diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h index fc8a8fbb625bfb..10657288fc5f6a 100644 --- a/src/coreclr/jit/compiler.h +++ b/src/coreclr/jit/compiler.h @@ -7535,7 +7535,7 @@ class Compiler void fgInlineAppendStatements(InlineInfo* inlineInfo, BasicBlock* block, Statement* stmt); void fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicBlock* joinBlock); void fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo); - void fgSetupAsyncFrameTransitionCall(GenTreeCall* call, InlineContext* inlineeContext, const DebugInfo& di); + void fgSetupAsyncFrameTransitionCall(GenTreeCall* call, const DebugInfo& di); GenTree* gtNewContinuationMemberIndir(const struct ContinuationMember& member, var_types type); #ifdef DEBUG diff --git a/src/coreclr/jit/fginline.cpp b/src/coreclr/jit/fginline.cpp index 363a8e54cede99..6ee67e5856cd4b 100644 --- a/src/coreclr/jit/fginline.cpp +++ b/src/coreclr/jit/fginline.cpp @@ -2581,16 +2581,20 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) // into a proper async call. // // Arguments: -// call - the call, with its user arguments already added -// inlineeContext - inline context of the frame that is logically returning -// di - debug info to use +// call - the call, with its user arguments already added +// di - debug info to use // // Notes: -// The call is an await: it suspends when it has to switch continuation context. It runs -// after the inlinee's frame has logically returned, so the frame it belongs to is the -// inlinee's caller. +// The call is an await: it suspends when it has to switch continuation context. // -void Compiler::fgSetupAsyncFrameTransitionCall(GenTreeCall* call, InlineContext* inlineeContext, const DebugInfo& di) +// Unlike normal awaits it gets no context pseudo-args. Those model what a frame +// captures and restores when a suspension unwinds through it, and none of that applies +// here: we only run at all because the frame we are returning out of was resumed, so +// every frame in the chain already has a continuation holding the values it captured +// when it first suspended, and a suspension here returns straight back to the +// dispatcher, which restores the thread's contexts itself. +// +void Compiler::fgSetupAsyncFrameTransitionCall(GenTreeCall* call, const DebugInfo& di) { AsyncCallInfo asyncInfo; // The helper resumes on the requested context, so no further handling is attached to @@ -2609,40 +2613,6 @@ void Compiler::fgSetupAsyncFrameTransitionCall(GenTreeCall* call, InlineContext* call->gtArgs.PushBack(this, NewCallArg::Primitive(gtNewNull(), TYP_REF).WellKnown(WellKnownArg::AsyncContinuation)); } - - // Contexts of the frame this call belongs to, i.e. the inlinee's caller. - InlineContext* const parentContext = inlineeContext->GetParent(); - unsigned resumed; - unsigned execCtx; - unsigned syncCtx; - if ((parentContext == nullptr) || !parentContext->HasAsyncFrameLocals()) - { - resumed = lvaResumedIndicator; - execCtx = lvaAsyncExecutionContextVar; - syncCtx = lvaAsyncSynchronizationContextVar; - } - else - { - resumed = parentContext->GetAsyncResumedIndicator(); - execCtx = parentContext->GetAsyncExecutionContextVar(); - syncCtx = parentContext->GetAsyncSynchronizationContextVar(); - } - - assert((resumed != BAD_VAR_NUM) && (execCtx != BAD_VAR_NUM) && (syncCtx != BAD_VAR_NUM)); - - // TODO-Async: when this call suspends, the frames enclosing the caller need their - // context handling run too. Those come from the enclosing frame chain, which is added - // to async calls as additional pseudo-args. - call->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(syncCtx, TYP_REF)) - .WellKnown(WellKnownArg::AsyncSynchronizationContext)); - call->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(execCtx, TYP_REF)) - .WellKnown(WellKnownArg::AsyncExecutionContext)); - call->gtArgs - .PushFront(this, NewCallArg::Primitive(gtNewLclAddrNode(resumed, 0)).WellKnown(WellKnownArg::AsyncResumedDef)); - call->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(resumed, TYP_INT)) - .WellKnown(WellKnownArg::AsyncResumedUse)); - - lvaGetDesc(resumed)->lvHasLdAddrOp = true; } //------------------------------------------------------------------------ @@ -2796,7 +2766,7 @@ void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicB gtNewContinuationMemberIndir(ContinuationMember::InlineFrameExecutionContext( inlineDepth), TYP_REF))); - fgSetupAsyncFrameTransitionCall(restoreCall, inlineContext, di); + fgSetupAsyncFrameTransitionCall(restoreCall, di); // The helper is small on the path that does not suspend, which is the one we care // about, so let the inliner have a look at it. From 285d6d92c7bd8075ba65ef8d61c4b845305c65df Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 31 Jul 2026 15:24:28 +0200 Subject: [PATCH 40/41] Drop redundant casts in RestoreInlinedFrameContexts The comparisons are reference comparisons either way, so the castclass just costs IL. That matters here: it brings the method from 130 to 120 IL bytes, under the 128 byte inlining limit, so the JIT can now actually inline it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- .../System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index 743dd815fc259b..ad89e7be28fe2e 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -1534,7 +1534,7 @@ private static void RestoreInlinedFrameContexts(ExecutionContext? previousExecCt else if ((flags & ContinuationFlags.ContinueOnCapturedSynchronizationContext) != 0) { Debug.Assert(continuationContext is SynchronizationContext); - if ((SynchronizationContext)continuationContext! == currentThread._synchronizationContext) + if (continuationContext == currentThread._synchronizationContext) { return; } @@ -1542,7 +1542,7 @@ private static void RestoreInlinedFrameContexts(ExecutionContext? previousExecCt else if ((flags & ContinuationFlags.ContinueOnCapturedTaskScheduler) != 0) { Debug.Assert(continuationContext is TaskScheduler); - if ((TaskScheduler)continuationContext! == TaskScheduler.InternalCurrent) + if (continuationContext == TaskScheduler.InternalCurrent) { return; } From 27aac9fa42faa9d437baa622a7f11a36a2874d57 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 31 Jul 2026 16:13:39 +0200 Subject: [PATCH 41/41] JIT: Derive the enclosing async frame chain from the inlining call The chain was built from a side table of context locals recorded on InlineContext. That table is populated when a frame is spliced into its caller, which happens after the frame's own inlinees have already been processed, so an inner frame never saw the frames between it and the root and silently dropped their transitions. Copy the context args off the call being inlined instead. They describe exactly the chain the inlinee ends up in, and they compose: a frame picks up its caller's chain when it is spliced in, and every call in it, including ones that came from its own inlinees, is extended again when that caller is inlined. The first set on the call belongs to the frame it suspends in and is consumed by that frame's own handling, so it is skipped unless it is all there is. This also drops the assumption that the root method is always a logical async frame. It is not when it does no context handling itself, for example when it is the async version of a synchronous method; there is then nothing to hand back and no transitions to run. InlineContext keeps only a bit saying whether the frame does context handling, set while the frame is compiled rather than when it is spliced in, so that depths are right for its own inlinees. Fixes an assert compiling System.IO.ConnectedStreams.BidirectionalStreamBufferStream.ReadAsync. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ed3e552-b907-4815-9683-56a7fed2cdf3 --- src/coreclr/jit/async.cpp | 5 + src/coreclr/jit/fginline.cpp | 175 +++++++++++++++++++++++------------ src/coreclr/jit/inline.h | 46 +++------ 3 files changed, 134 insertions(+), 92 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index c50e7576a56431..c6c21d06f83f60 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -269,6 +269,11 @@ PhaseStatus Compiler::SaveAsyncContexts() lvaResumedIndicator = lvaGrabTemp(false DEBUGARG("Async Resumed")); lvaGetDesc(lvaResumedIndicator)->lvType = TYP_UBYTE; + // Mark this frame as a logical async frame. This has to happen here rather than when + // an inlinee is spliced into its caller: the frame's own inlinees are processed while + // it is still being compiled, and they need to see it in the enclosing frame chain. + compInlineContext->SetIsAsyncFrame(); + if (opts.IsOSR()) { lvaGetDesc(lvaAsyncThreadObjectVar)->lvIsOSRLocal = true; diff --git a/src/coreclr/jit/fginline.cpp b/src/coreclr/jit/fginline.cpp index 6ee67e5856cd4b..a78f66ae1ea632 100644 --- a/src/coreclr/jit/fginline.cpp +++ b/src/coreclr/jit/fginline.cpp @@ -2446,6 +2446,12 @@ Statement* Compiler::fgInlinePrependStatements(InlineInfo* inlineInfo) // live and GC reported at the call until then. That is what the existing per-frame // context args do, so the enclosing frames simply add more of them. // +// The enclosing frames are exactly the ones the call being inlined already describes, +// so its context args are simply copied over. That composes for nested inlines: a +// frame picks up its immediate caller's chain when it is spliced in, and every call in +// it, including ones that came from its own inlinees, is extended again when that +// caller is itself inlined. +// // These are pseudo-args, so duplicates are fine: they take no registers and are // expanded out by the async transformation. The innermost frame stays first, since // the inlinee pushed its own args to the front. @@ -2462,35 +2468,44 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) return; } - // Async context locals of one inlined frame. - struct AsyncFrameLocals + // Context args of the inlining call. It carries one set for the frame it suspends in, + // which that frame's own handling consumes, followed by that frame's chain of frames + // out to the root. + ArrayStack callerArgs(getAllocator(CMK_Async)); + for (CallArg& arg : inlineInfo->iciCall->gtArgs.Args()) { - unsigned Resumed; - unsigned ExecutionContext; - unsigned SynchronizationContext; - }; - - ArrayStack frames(getAllocator(CMK_Async)); - - // The inlinee's own frame comes first: the suspension needs its resumed indicator to - // decide whether to capture the contexts it hands to its caller. Its other context - // args were consumed by the innermost handling by then. - frames.Push({InlineeCompiler->lvaResumedIndicator, InlineeCompiler->lvaAsyncExecutionContextVar, - InlineeCompiler->lvaAsyncSynchronizationContextVar}); - - for (InlineContext* ctx = inlineInfo->inlineContext->GetParent(); ctx != nullptr; ctx = ctx->GetParent()) - { - if (ctx->HasAsyncFrameLocals()) + switch (arg.GetWellKnownArg()) { - frames.Push({ctx->GetAsyncResumedIndicator(), ctx->GetAsyncExecutionContextVar(), - ctx->GetAsyncSynchronizationContextVar()}); + case WellKnownArg::AsyncResumedUse: + case WellKnownArg::AsyncExecutionContext: + case WellKnownArg::AsyncSynchronizationContext: + callerArgs.Push(&arg); + break; + default: + break; } } - // The root method's frame is always the outermost one. - frames.Push({lvaResumedIndicator, lvaAsyncExecutionContextVar, lvaAsyncSynchronizationContextVar}); + if (callerArgs.Height() == 0) + { + // The inlining call does no context handling, so no enclosing frame is a logical + // async frame. The inlinee's frame has nothing to hand back and there are no frame + // transitions to run on suspension. + JITDUMP("Inlining call [%06u] has no context args; inlinee has no enclosing async frame\n", + dspTreeID(inlineInfo->iciCall)); + return; + } + + assert((callerArgs.Height() % 3) == 0); + + // The frames enclosing the inlinee are the inlining call's chain. When the call has + // only the one set it sits in the root frame, whose chain is never materialized + // because it has no transitions of its own; that frame is then the only enclosing one. + int const firstEnclosing = (callerArgs.Height() == 3) ? 0 : 3; unsigned const inlineeResumed = InlineeCompiler->lvaResumedIndicator; + unsigned const inlineeExecCtx = InlineeCompiler->lvaAsyncExecutionContextVar; + unsigned const inlineeSyncCtx = InlineeCompiler->lvaAsyncSynchronizationContextVar; struct Visitor : GenTreeVisitor { @@ -2499,13 +2514,24 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) DoPreOrder = true, }; - ArrayStack& m_frames; - unsigned m_inlineeResumed; - - Visitor(Compiler* comp, ArrayStack& frames, unsigned inlineeResumed) + ArrayStack& m_callerArgs; + int m_firstEnclosing; + unsigned m_inlineeResumed; + unsigned m_inlineeExecCtx; + unsigned m_inlineeSyncCtx; + + Visitor(Compiler* comp, + ArrayStack& callerArgs, + int firstEnclosing, + unsigned inlineeResumed, + unsigned inlineeExecCtx, + unsigned inlineeSyncCtx) : GenTreeVisitor(comp) - , m_frames(frames) + , m_callerArgs(callerArgs) + , m_firstEnclosing(firstEnclosing) , m_inlineeResumed(inlineeResumed) + , m_inlineeExecCtx(inlineeExecCtx) + , m_inlineeSyncCtx(inlineeSyncCtx) { } @@ -2522,46 +2548,70 @@ void Compiler::fgAppendEnclosingAsyncFrameContextArgs(InlineInfo* inlineInfo) return WALK_CONTINUE; } - GenTreeCall* const call = tree->AsCall(); - CallArg* const ownArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse); - if ((ownArg == nullptr) || !ownArg->GetNode()->OperIs(GT_LCL_VAR) || - (ownArg->GetNode()->AsLclVar()->GetLclNum() != m_inlineeResumed)) + GenTreeCall* const call = tree->AsCall(); + + unsigned numOwn = 0; + for (CallArg& arg : call->gtArgs.Args()) { - // Either no context handling, or the call inherited its contexts from the - // inlining call, in which case there is no frame transition. - return WALK_CONTINUE; + if (arg.GetWellKnownArg() == WellKnownArg::AsyncResumedUse) + { + numOwn++; + } } - for (int i = 0; i < m_frames.Height(); i++) + if (numOwn == 0) { - const AsyncFrameLocals& frame = m_frames.Bottom(i); - assert((frame.Resumed != BAD_VAR_NUM) && (frame.ExecutionContext != BAD_VAR_NUM) && - (frame.SynchronizationContext != BAD_VAR_NUM)); + // Call does no context handling, so it has no frame chain to extend. + return WALK_CONTINUE; + } + if (numOwn == 1) + { + // The call only has the set describing the frame it suspends in, so that + // frame is the inlinee's. The suspension needs the inlinee's resumed + // indicator again to decide whether to capture the contexts it hands to + // its caller, since the first set is consumed by the innermost handling. call->gtArgs.PushBack(m_compiler, - NewCallArg::Primitive(m_compiler->gtNewLclVarNode(frame.Resumed, TYP_INT)) + NewCallArg::Primitive(m_compiler->gtNewLclVarNode(m_inlineeResumed, TYP_INT)) .WellKnown(WellKnownArg::AsyncResumedUse)); - call->gtArgs.PushBack(m_compiler, NewCallArg::Primitive( - m_compiler->gtNewLclVarNode(frame.ExecutionContext, TYP_REF)) - .WellKnown(WellKnownArg::AsyncExecutionContext)); call->gtArgs.PushBack(m_compiler, - NewCallArg::Primitive( - m_compiler->gtNewLclVarNode(frame.SynchronizationContext, TYP_REF)) + NewCallArg::Primitive(m_compiler->gtNewLclVarNode(m_inlineeExecCtx, TYP_REF)) + .WellKnown(WellKnownArg::AsyncExecutionContext)); + call->gtArgs.PushBack(m_compiler, + NewCallArg::Primitive(m_compiler->gtNewLclVarNode(m_inlineeSyncCtx, TYP_REF)) .WellKnown(WellKnownArg::AsyncSynchronizationContext)); } - JITDUMP("Added %d enclosing frame context arg sets to async call [%06u]\n", m_frames.Height(), - Compiler::dspTreeID(call)); + // Then the frames enclosing the inlinee, taken from the inlining call. These + // are the values as they appear in the caller's IR, which is what the + // suspension has to store. + for (int i = m_firstEnclosing; i < m_callerArgs.Height(); i++) + { + CallArg* const arg = m_callerArgs.Bottom(i); + call->gtArgs.PushBack(m_compiler, NewCallArg::Primitive(m_compiler->gtCloneExpr(arg->GetNode())) + .WellKnown(arg->GetWellKnownArg())); + } - // The chain holds this call's own frame followed by its enclosing frames, - // ending with the root, so its depth in that numbering is one less than the - // number of entries. - call->GetAsyncInfo().InlineFrameDepth = (unsigned)m_frames.Height() - 1; + unsigned numFrames = 0; + for (CallArg& arg : call->gtArgs.Args()) + { + if (arg.GetWellKnownArg() == WellKnownArg::AsyncResumedUse) + { + numFrames++; + } + } + + // One set is consumed by the handling of the frame the call suspends in, and + // the rest describe that frame and the ones enclosing it. + assert(numFrames >= 2); + call->GetAsyncInfo().InlineFrameDepth = numFrames - 2; + + JITDUMP("Extended async call [%06u] to %u frames in chain\n", Compiler::dspTreeID(call), numFrames - 1); return WALK_CONTINUE; } }; - Visitor visitor(this, frames, inlineeResumed); + Visitor visitor(this, callerArgs, firstEnclosing, inlineeResumed, inlineeExecCtx, inlineeSyncCtx); for (BasicBlock* block = InlineeCompiler->fgFirstBB; block != nullptr; block = block->Next()) { for (Statement* const stmt : block->Statements()) @@ -2695,9 +2745,6 @@ void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicB } InlineContext* const inlineContext = inlineInfo->inlineContext; - inlineContext->SetAsyncFrameLocals(InlineeCompiler->lvaResumedIndicator, - InlineeCompiler->lvaAsyncExecutionContextVar, - InlineeCompiler->lvaAsyncSynchronizationContextVar); if (!InlineeCompiler->compAsyncBodyMaySuspend) { @@ -2707,12 +2754,20 @@ void Compiler::fgInlineAppendAsyncFrameStatements(InlineInfo* inlineInfo, BasicB return; } - // Locals of the frame we are logically returning into. - InlineContext* const parentContext = inlineContext->GetParent(); - unsigned const resumedCaller = ((parentContext == nullptr) || !parentContext->HasAsyncFrameLocals()) - ? lvaResumedIndicator - : parentContext->GetAsyncResumedIndicator(); - assert(resumedCaller != BAD_VAR_NUM); + // The resumed indicator of the frame we are logically returning into. The inlining + // call carries it as the definition of its own frame's indicator, which is a local + // address and so something we can store to. + CallArg* const resumedDefArg = inlineInfo->iciCall->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedDef); + if (resumedDefArg == nullptr) + { + // The inlining call does no context handling, so there is no enclosing logical + // async frame to hand anything back to. + JITDUMP("Inlining call does no context handling; no async frame transition IR needed\n"); + return; + } + + assert(resumedDefArg->GetNode()->OperIs(GT_LCL_ADDR)); + unsigned const resumedCaller = resumedDefArg->GetNode()->AsLclVarCommon()->GetLclNum(); unsigned const resumedInlinee = InlineeCompiler->lvaResumedIndicator; unsigned const inlineDepth = inlineContext->GetAsyncFrameDepth(); diff --git a/src/coreclr/jit/inline.h b/src/coreclr/jit/inline.h index de7fdedd4ab065..4318cf603ff9b8 100644 --- a/src/coreclr/jit/inline.h +++ b/src/coreclr/jit/inline.h @@ -793,19 +793,20 @@ class InlineContext return m_Parent; } - // Depth of the inlined frame this context represents, counting only frames that have - // async context handling. The root method's frame is depth 0, so the shallowest - // inlined frame with context handling is depth 1. + // Depth of the inlined frame this context represents, counting only frames that do + // async context handling. The shallowest such frame is depth 1 when the root method + // does context handling itself, and depth 0 otherwise (for example when the root is + // the async version of a synchronous method). // // Continuation members holding the contexts captured at frame transitions are keyed // by this, so frames at the same depth share storage. That is safe because their live // ranges cannot overlap. unsigned GetAsyncFrameDepth() const { - unsigned depth = 1; + unsigned depth = 0; for (InlineContext* parent = m_Parent; parent != nullptr; parent = parent->m_Parent) { - if (parent->HasAsyncFrameLocals()) + if (parent->IsAsyncFrame()) { depth++; } @@ -920,34 +921,17 @@ class InlineContext return (m_PgoInfo.PgoSchema != nullptr) && (m_PgoInfo.PgoSchemaCount > 0) && (m_PgoInfo.PgoData != nullptr); } - // The async context locals of the inlined frame this context represents. Recorded when - // the inlinee is spliced in, since the inlinee's Compiler is discarded afterwards. - // Used to build the logical frame transitions for general runtime async inlining. - void SetAsyncFrameLocals(unsigned resumedIndicator, unsigned executionContextVar, unsigned syncContextVar) + // Whether the frame this context represents does async context handling, i.e. whether + // it is a logical async frame. Set when the corresponding Compiler creates its context + // locals, so that it is already known while the frame's own inlinees are processed. + void SetIsAsyncFrame() { - m_asyncResumedIndicator = resumedIndicator; - m_asyncExecutionContextVar = executionContextVar; - m_asyncSynchronizationContextVar = syncContextVar; + m_isAsyncFrame = true; } - bool HasAsyncFrameLocals() const + bool IsAsyncFrame() const { - return m_asyncResumedIndicator != BAD_VAR_NUM; - } - - unsigned GetAsyncResumedIndicator() const - { - return m_asyncResumedIndicator; - } - - unsigned GetAsyncExecutionContextVar() const - { - return m_asyncExecutionContextVar; - } - - unsigned GetAsyncSynchronizationContextVar() const - { - return m_asyncSynchronizationContextVar; + return m_isAsyncFrame; } private: @@ -970,9 +954,7 @@ class InlineContext unsigned m_Ordinal; // Ordinal number of this inline bool m_Success : 1; // true if this was a successful inline - unsigned m_asyncResumedIndicator = BAD_VAR_NUM; - unsigned m_asyncExecutionContextVar = BAD_VAR_NUM; - unsigned m_asyncSynchronizationContextVar = BAD_VAR_NUM; + bool m_isAsyncFrame = false; #if defined(DEBUG)