From 782d0fb128866919074ed1ba2f0fc578e7d84e66 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Thu, 28 May 2026 19:13:08 +0200 Subject: [PATCH 1/3] Revert "JIT: Tail merge async suspensions (#128559)" This reverts commit 269c0befa6ec5b6b5af1031a8533ba54bdd4e70b. --- src/coreclr/jit/async.cpp | 374 +----------------------------- src/coreclr/jit/async.h | 35 +-- src/coreclr/jit/gentree.cpp | 41 ---- src/coreclr/jit/gentree.h | 2 - src/coreclr/jit/jitmetadatalist.h | 1 - 5 files changed, 21 insertions(+), 432 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index b43572225035e5..e87309c7c03227 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -513,79 +513,6 @@ bool ContinuationLayoutBuilder::ContainsLocal(unsigned lclNum) const return BinarySearch(m_locals.data(), (int)m_locals.size(), lclNum) != nullptr; } -//------------------------------------------------------------------------ -// ContinuationLayoutBuilder::Equals: -// Check if two builders produce the same layout. -// -// Parameters: -// a - The first builder to compare. -// b - The second builder to compare. -// -// Returns: -// True if the builders produce the same layout. -// -bool ContinuationLayoutBuilder::Equals(const ContinuationLayoutBuilder& a, const ContinuationLayoutBuilder& b) -{ - if (a.m_needsOSRAddress != b.m_needsOSRAddress) - { - return false; - } - - if (a.m_needsException != b.m_needsException) - { - return false; - } - - if (a.m_needsContinuationContext != b.m_needsContinuationContext) - { - return false; - } - - if (a.m_needsKeepAlive != b.m_needsKeepAlive) - { - return false; - } - - if (a.m_needsExecutionContext != b.m_needsExecutionContext) - { - return false; - } - - if (a.m_returns.size() != b.m_returns.size()) - { - return false; - } - - for (size_t i = 0; i < a.m_returns.size(); i++) - { - if (a.m_returns[i].ReturnType != b.m_returns[i].ReturnType) - { - return false; - } - - if ((a.m_returns[i].ReturnType == TYP_STRUCT) && - !ClassLayout::AreCompatible(a.m_returns[i].ReturnLayout, b.m_returns[i].ReturnLayout)) - { - return false; - } - } - - if (a.m_locals.size() != b.m_locals.size()) - { - return false; - } - - for (size_t i = 0; i < a.m_locals.size(); i++) - { - if (a.m_locals[i] != b.m_locals[i]) - { - return false; - } - } - - return true; -} - //------------------------------------------------------------------------ // TransformAsync: Run async transformation. // @@ -1017,31 +944,17 @@ void AsyncTransformation::Transform(BasicBlock* block, CallDefinitionInfo callDefInfo = CanonicalizeCallDefinition(block, call, &analyses); - const AsyncState* reusedState = - FindReusableSuspension(block, call, callDefInfo, layoutBuilder, resumeReachable, mutatedSinceResumption); - - if (reusedState != nullptr) - { - JITDUMP(" Reused state %u\n", reusedState->Number); - CreateCheckAndSuspendAfterCall(block, call, callDefInfo, reusedState->SuspensionBB, remainder); - - HandleReusedSuspension(block, call); - m_compiler->Metrics.SuspensionPointsMerged++; - } - else - { - unsigned stateNum = (unsigned)m_states.size(); - JITDUMP(" Assigned state %u\n", stateNum); + unsigned stateNum = (unsigned)m_states.size(); + JITDUMP(" Assigned state %u\n", stateNum); - BasicBlock* suspendBB = CreateSuspensionBlock(block, stateNum); + BasicBlock* suspendBB = CreateSuspensionBlock(block, stateNum); - CreateCheckAndSuspendAfterCall(block, call, callDefInfo, suspendBB, remainder); + CreateCheckAndSuspendAfterCall(block, call, callDefInfo, suspendBB, remainder); - BasicBlock* resumeBB = CreateResumptionBlock(*remainder, stateNum); + BasicBlock* resumeBB = CreateResumptionBlock(*remainder, stateNum); - m_states.push_back(AsyncState(stateNum, layoutBuilder, block, call, callDefInfo, suspendBB, resumeBB, - resumeReachable, mutatedSinceResumption)); - } + m_states.push_back(AsyncState(stateNum, layoutBuilder, block, call, callDefInfo, suspendBB, resumeBB, + resumeReachable, mutatedSinceResumption)); JITDUMP("\n"); } @@ -1743,269 +1656,6 @@ CallDefinitionInfo AsyncTransformation::CanonicalizeCallDefinition(BasicBlock* return callDefInfo; } -//------------------------------------------------------------------------ -// AsyncTransformation::FindReusableSuspension: -// Try to find a suspension/resumption state that we can merge with. -// -// Parameters: -// block - The block containing the async call. -// call - The async call -// defInfo - Async call def info -// layoutBuilder - Layout for this async call -// resumeReachable - Whether 'call' is reachable after a previous resumption -// mutatedSinceResumption - Set of variables mutated since the last resumption -// -// Returns: -// Reusable existing state, or nullptr if no such state. -// -const AsyncState* AsyncTransformation::FindReusableSuspension(BasicBlock* block, - GenTreeCall* call, - const CallDefinitionInfo& defInfo, - ContinuationLayoutBuilder* layoutBuilder, - bool resumeReachable, - VARSET_VALARG_TP mutatedSinceResumption) -{ - if (m_compiler->opts.OptimizationDisabled()) - { - // Tail merging breaks debugging. - return nullptr; - } - - // Call must be at the tail of its block. - GenTree* thisLastNode = block->lastNode(); - if ((thisLastNode != call) && (thisLastNode != defInfo.DefinitionNode)) - { - return nullptr; - } - - // We can only merge if we logically could move the suspension point into a - // successor of two blocks. - BasicBlock* target = block->GetUniqueSucc(); - if (target == nullptr) - { - return nullptr; - } - - for (BasicBlock* pred : target->PredBlocks()) - { - if (pred == block) - { - continue; - } - - if (!pred->KindIs(BBJ_ALWAYS) || !pred->isEmpty()) - { - // We are looking for the "remainder" block that was created for a previous async transformation. - // If they are in a usable tail position they will always be empty BBJ_ALWAYS blocks. - continue; - } - - int limit = 128; - for (size_t i = m_states.size(); i != 0; i--) - { - if (limit-- == 0) - { - break; - } - - const AsyncState& state = m_states[i - 1]; - - assert(state.ResumptionBB->KindIs(BBJ_ALWAYS)); - BasicBlock* resumptionJoin = state.ResumptionBB->GetUniqueSucc(); - - // This existing async state resumes at a sibling join block. We - // can reuse that resumption if it has the same state as us. - if ((resumptionJoin == pred) && IsReusableSuspension(&state, block, call, defInfo, layoutBuilder, - resumeReachable, mutatedSinceResumption)) - { - return &state; - } - } - } - - return nullptr; -} - -//------------------------------------------------------------------------ -// AsyncTransformation::IsReusableSuspension: -// Check if the suspension for a specified state is reusable for an async call. -// -// Parameters: -// state - State for existing suspension. -// block - The block containing the async call. -// call - The async call -// defInfo - Async call def info -// layoutBuilder - Layout for this async call -// resumeReachable - Whether 'call' is reachable after a previous resumption -// mutatedSinceResumption - Set of variables mutated since the last resumption -// -// Returns: -// True if the suspension is reusable. -// -bool AsyncTransformation::IsReusableSuspension(const AsyncState* state, - BasicBlock* block, - GenTreeCall* call, - const CallDefinitionInfo& defInfo, - ContinuationLayoutBuilder* layoutBuilder, - bool resumeReachable, - VARSET_VALARG_TP mutatedSinceResumption) -{ - GenTreeCall* predAsyncCall = state->Call; - - JITDUMP( - " Sibling to the join is resumption for async call [%06u]; checking for possible tail merging of suspension points\n", - Compiler::dspTreeID(predAsyncCall)); - - if (!BasicBlock::sameEHRegion(block, state->CallBlock)) - { - JITDUMP(" Not same EH region\n"); - return false; - } - - // We have a sibling of the join that ends with an async call. Check if it is compatible. - if ((defInfo.DefinitionNode == nullptr) != (state->CallDefInfo.DefinitionNode == nullptr)) - { - JITDUMP(" No; disagreement on presence of return value\n"); - return false; - } - - if ((defInfo.DefinitionNode != nullptr) && - !GenTreeLclVarCommon::EqualsLocal(defInfo.DefinitionNode, state->CallDefInfo.DefinitionNode)) - { - JITDUMP(" No; disagreement on return value destination ([%06u] does not equal [%06u])\n", - Compiler::dspTreeID(defInfo.DefinitionNode), Compiler::dspTreeID(state->CallDefInfo.DefinitionNode)); - return false; - } - - if (call->gtReturnType != predAsyncCall->AsCall()->gtReturnType) - { - JITDUMP(" No; disagreement on return type (%s vs %s)\n", varTypeName(call->gtReturnType), - varTypeName(predAsyncCall->AsCall()->gtReturnType)); - return false; - } - - if (call->gtReturnType == TYP_STRUCT) - { - ClassLayout* thisLayout = m_compiler->typGetObjLayout(call->gtRetClsHnd); - ClassLayout* otherLayout = m_compiler->typGetObjLayout(predAsyncCall->AsCall()->gtRetClsHnd); - if (!ClassLayout::AreCompatible(thisLayout, otherLayout)) - { - JITDUMP(" No; disagreement on return type (%s vs %s)\n", thisLayout->GetClassName(), - otherLayout->GetClassName()); - return false; - } - } - - const AsyncCallInfo& asyncInfoThis = call->GetAsyncInfo(); - const AsyncCallInfo& asyncInfoOther = predAsyncCall->AsCall()->GetAsyncInfo(); - - if (asyncInfoThis.ContinuationContextHandling != asyncInfoOther.ContinuationContextHandling) - { - JITDUMP(" No; disagreement on continuation context handling (%u vs %u)\n", - asyncInfoThis.ContinuationContextHandling, asyncInfoOther.ContinuationContextHandling); - return false; - } - - if (asyncInfoThis.NeedsToSaveAndRestoreExecutionContext() != asyncInfoOther.NeedsToSaveAndRestoreExecutionContext()) - { - JITDUMP(" No; disagreement on whether execution context needs to be saved and restored (%s vs %s)\n", - asyncInfoThis.NeedsToSaveAndRestoreExecutionContext() ? "yes" : "no", - asyncInfoOther.NeedsToSaveAndRestoreExecutionContext() ? "yes" : "no"); - return false; - } - - if (state->ResumeReachable != resumeReachable) - { - // Not being resume reachable means we can skip checking for a - // reusable continuation. We do not want to give up that - // optimization. - JITDUMP(" No; disagreement on resume reachability (%s vs %s)\n", state->ResumeReachable ? "yes" : "no", - resumeReachable ? "yes" : "no"); - return false; - } - - if (resumeReachable) - { - // If both are reachable after resuming then we need to take care that - // we update the right set of locals in the continuation on suspension. - // In one path we may have mutated a local that we didn't in another - // path, and that results in a difference when we reuse a continuation - // and need to decide if that local needs to be saved again or not. - for (unsigned lclNum : layoutBuilder->Locals()) - { - LclVarDsc* dsc = m_compiler->lvaGetDesc(lclNum); - if (GetLocalSaveSet(dsc, mutatedSinceResumption) != - GetLocalSaveSet(dsc, state->MutatedSincePreviousResumption)) - { - JITDUMP(" No; disagreement on save set for V%02u (%d vs %d)\n", lclNum, - (unsigned)GetLocalSaveSet(dsc, mutatedSinceResumption), - (unsigned)GetLocalSaveSet(dsc, state->MutatedSincePreviousResumption)); - return false; - } - } - } - - static const WellKnownArg validateArgs[] = {WellKnownArg::AsyncExecutionContext, - WellKnownArg::AsyncSynchronizationContext}; - for (WellKnownArg arg : validateArgs) - { - CallArg* thisArg = call->gtArgs.FindWellKnownArg(arg); - CallArg* otherArg = predAsyncCall->gtArgs.FindWellKnownArg(arg); - - if ((thisArg == nullptr) != (otherArg == nullptr)) - { - JITDUMP(" No; disagreement on presence of %s argument\n", getWellKnownArgName(arg)); - return false; - } - - if (thisArg != nullptr) - { - if (!thisArg->GetNode()->OperIsLocal()) - { - JITDUMP(" No; %s argument is too complex\n", getWellKnownArgName(arg)); - return false; - } - - if (!GenTree::Compare(thisArg->GetNode(), otherArg->GetNode())) - { - JITDUMP(" No; disagreement on value of %s argument ([%06u] vs [%06u])\n", getWellKnownArgName(arg), - Compiler::dspTreeID(thisArg->GetNode()), Compiler::dspTreeID(otherArg->GetNode())); - return false; - } - } - } - - // At this point we expect both suspensions to agree on the live locals - // since they both join at the same point. So we only need to assert that - // the layouts match. - assert(ContinuationLayoutBuilder::Equals(*layoutBuilder, *state->Layout)); - return true; -} - -//------------------------------------------------------------------------ -// AsyncTransformation::HandleReusedSuspension: -// Handle the case where we were able to reuse a previously created suspension state. -// -// Parameters: -// callBlock - Block containing async call -// call - The async call -// -void AsyncTransformation::HandleReusedSuspension(BasicBlock* callBlock, GenTreeCall* call) -{ - static const WellKnownArg argsToRemove[] = {WellKnownArg::AsyncExecutionContext, - WellKnownArg::AsyncSynchronizationContext}; - for (WellKnownArg wka : argsToRemove) - { - CallArg* arg = call->gtArgs.FindWellKnownArg(wka); - if (arg != nullptr) - { - assert(arg->GetNode()->OperIsLocal()); - LIR::AsRange(callBlock).Remove(arg->GetNode()); - call->gtArgs.RemoveUnsafe(arg); - } - } -} - //------------------------------------------------------------------------ // AsyncTransformation::CreateSuspensionBlock: // Create an empty basic block that will hold the suspension IR for the @@ -2140,7 +1790,7 @@ void AsyncTransformation::CreateSuspension(BasicBlock* call // In the path where we allocated a new continuation we save only locals that we know to be unmutated since the // last resumption. - FillInDataOnSuspension(layout, subLayout, allocNewBB, mutatedSinceResumption, SaveSet::UnmutatedLocals); + FillInDataOnSuspension(call, layout, subLayout, allocNewBB, mutatedSinceResumption, SaveSet::UnmutatedLocals); // We can skip saving unmutated locals in the shared path -- we only need to save locals that may have been // mutated since the last resumption. @@ -2218,7 +1868,7 @@ void AsyncTransformation::CreateSuspension(BasicBlock* call GenTree* storeFlags = StoreAtOffset(newContinuation, flagsOffset, flagsNode, TYP_INT); LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_compiler, storeFlags)); - FillInDataOnSuspension(layout, subLayout, suspendBB, mutatedSinceResumption, tailSaveSet); + FillInDataOnSuspension(call, layout, subLayout, suspendBB, mutatedSinceResumption, tailSaveSet); FinishContextHandlingAndSuspension(callBlock, call, suspendBB, layout, subLayout); } @@ -2268,13 +1918,13 @@ GenTreeCall* AsyncTransformation::CreateAllocContinuationCall(bool // context. // // Parameters: +// call - The async call. // layout - Information about the continuation layout. // subLayout - Per-call layout builder indicating which fields are needed. // suspendBB - Basic block to add IR to. -// mutatedSinceResumption - Set of locals mutated since the last resumption. -// saveSet - Set of locals to save // -void AsyncTransformation::FillInDataOnSuspension(const ContinuationLayout& layout, +void AsyncTransformation::FillInDataOnSuspension(GenTreeCall* call, + const ContinuationLayout& layout, const ContinuationLayoutBuilder& subLayout, BasicBlock* suspendBB, VARSET_VALARG_TP mutatedSinceResumption, diff --git a/src/coreclr/jit/async.h b/src/coreclr/jit/async.h index 6e194fef03019b..ce25cd991362cf 100644 --- a/src/coreclr/jit/async.h +++ b/src/coreclr/jit/async.h @@ -119,8 +119,6 @@ struct ContinuationLayoutBuilder return m_locals; } - static bool Equals(const ContinuationLayoutBuilder& a, const ContinuationLayoutBuilder& b); - struct ContinuationLayout* Create(); static ContinuationLayoutBuilder* CreateSharedLayout(Compiler* comp, @@ -406,21 +404,6 @@ class AsyncTransformation CallDefinitionInfo CanonicalizeCallDefinition(BasicBlock* block, GenTreeCall* call, AsyncAnalysis* analyses); - const AsyncState* FindReusableSuspension(BasicBlock* block, - GenTreeCall* call, - const CallDefinitionInfo& defInfo, - ContinuationLayoutBuilder* layoutBuilder, - bool resumeReachable, - VARSET_VALARG_TP mutatedSinceResumption); - bool IsReusableSuspension(const AsyncState* state, - BasicBlock* block, - GenTreeCall* call, - const CallDefinitionInfo& defInfo, - ContinuationLayoutBuilder* layoutBuilder, - bool resumeReachable, - VARSET_VALARG_TP mutatedSinceResumption); - void HandleReusedSuspension(BasicBlock* callBlock, GenTreeCall* call); - BasicBlock* CreateSuspensionBlock(BasicBlock* block, unsigned stateNum); void CreateSuspension(BasicBlock* callBlock, GenTreeCall* call, @@ -435,7 +418,8 @@ class AsyncTransformation GenTree* prevContinuation, const ContinuationLayout& layout); - void FillInDataOnSuspension(const ContinuationLayout& layout, + void FillInDataOnSuspension(GenTreeCall* call, + const ContinuationLayout& layout, const ContinuationLayoutBuilder& subLayout, BasicBlock* suspendBB, VARSET_VALARG_TP mutatedSinceResumption, @@ -459,14 +443,13 @@ class AsyncTransformation const CallDefinitionInfo& callDefInfo, BasicBlock* suspendBB, BasicBlock** remainder); - - BasicBlock* CreateResumptionBlock(BasicBlock* remainder, unsigned stateNum); - void CreateResumption(BasicBlock* callBlock, - GenTreeCall* call, - BasicBlock* resumeBB, - const CallDefinitionInfo& callDefInfo, - const ContinuationLayout& layout, - const ContinuationLayoutBuilder& subLayout); + BasicBlock* CreateResumptionBlock(BasicBlock* remainder, unsigned stateNum); + void CreateResumption(BasicBlock* callBlock, + GenTreeCall* call, + BasicBlock* resumeBB, + const CallDefinitionInfo& callDefInfo, + const ContinuationLayout& layout, + const ContinuationLayoutBuilder& subLayout); void RestoreFromDataOnResumption(const ContinuationLayout& layout, const ContinuationLayoutBuilder& subLayout, diff --git a/src/coreclr/jit/gentree.cpp b/src/coreclr/jit/gentree.cpp index d5d2fb86a90960..74fb96430251ce 100644 --- a/src/coreclr/jit/gentree.cpp +++ b/src/coreclr/jit/gentree.cpp @@ -33107,47 +33107,6 @@ ClassLayout* GenTreeLclVarCommon::GetLayout(Compiler* compiler) const return AsLclFld()->GetLayout(); } -//------------------------------------------------------------------------ -// EqualsLocal: -// Check if the locals information of two locals is equal. -// -// Arguments: -// lcl1 - The first local -// lcl2 - The second local -// -// Returns: -// True if equal. -// -// Remarks: -// For LCL_VAR, LCL_FLD and LCL_ADDR this is equivalent to GenTree::Compare. -// For STORE_LCL_VAR and STORE_LCL_FLD this only checks the local -// information; it does not check the data node for equality. -// -bool GenTreeLclVarCommon::EqualsLocal(GenTreeLclVarCommon* lcl1, GenTreeLclVarCommon* lcl2) -{ - if (lcl1->OperGet() != lcl2->OperGet()) - { - return false; - } - - if (lcl1->GetLclNum() != lcl2->GetLclNum()) - { - return false; - } - - if (lcl1->GetLclOffs() != lcl2->GetLclOffs()) - { - return false; - } - - if (lcl1->OperIs(GT_LCL_FLD, GT_STORE_LCL_FLD) && lcl1->AsLclFld()->GetLayout() != lcl2->AsLclFld()->GetLayout()) - { - return false; - } - - return true; -} - //------------------------------------------------------------------------ // GenTreeLclVar::IsNeverNegative: Gets true if the lcl var is never negative; otherwise false. // diff --git a/src/coreclr/jit/gentree.h b/src/coreclr/jit/gentree.h index 3c51fc9541e201..525c088b1db2c6 100644 --- a/src/coreclr/jit/gentree.h +++ b/src/coreclr/jit/gentree.h @@ -3851,8 +3851,6 @@ struct GenTreeLclVarCommon : public GenTreeUnOp { } #endif - - static bool EqualsLocal(GenTreeLclVarCommon* lcl1, GenTreeLclVarCommon* lcl2); }; //------------------------------------------------------------------------ diff --git a/src/coreclr/jit/jitmetadatalist.h b/src/coreclr/jit/jitmetadatalist.h index 53c363ef5f24f2..e750ec79dcab03 100644 --- a/src/coreclr/jit/jitmetadatalist.h +++ b/src/coreclr/jit/jitmetadatalist.h @@ -94,7 +94,6 @@ JITMETADATAMETRIC(MorphTrackedLocals, int, 0) JITMETADATAMETRIC(MorphLocals, int, 0) JITMETADATAMETRIC(EnumeratorGDVProvisionalNoEscape, int, 0) JITMETADATAMETRIC(EnumeratorGDVCanCloneToEnsureNoEscape, int, 0) -JITMETADATAMETRIC(SuspensionPointsMerged, int, 0) #undef JITMETADATA #undef JITMETADATAINFO From edb9739ed2887a19b682d60aa9e78bb64b628539 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Thu, 28 May 2026 19:14:31 +0200 Subject: [PATCH 2/3] Reapply "JIT: Tail merge async suspensions (#128559)" This reverts commit 782d0fb128866919074ed1ba2f0fc578e7d84e66. --- src/coreclr/jit/async.cpp | 374 +++++++++++++++++++++++++++++- src/coreclr/jit/async.h | 35 ++- src/coreclr/jit/gentree.cpp | 41 ++++ src/coreclr/jit/gentree.h | 2 + src/coreclr/jit/jitmetadatalist.h | 1 + 5 files changed, 432 insertions(+), 21 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index e87309c7c03227..b43572225035e5 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -513,6 +513,79 @@ bool ContinuationLayoutBuilder::ContainsLocal(unsigned lclNum) const return BinarySearch(m_locals.data(), (int)m_locals.size(), lclNum) != nullptr; } +//------------------------------------------------------------------------ +// ContinuationLayoutBuilder::Equals: +// Check if two builders produce the same layout. +// +// Parameters: +// a - The first builder to compare. +// b - The second builder to compare. +// +// Returns: +// True if the builders produce the same layout. +// +bool ContinuationLayoutBuilder::Equals(const ContinuationLayoutBuilder& a, const ContinuationLayoutBuilder& b) +{ + if (a.m_needsOSRAddress != b.m_needsOSRAddress) + { + return false; + } + + if (a.m_needsException != b.m_needsException) + { + return false; + } + + if (a.m_needsContinuationContext != b.m_needsContinuationContext) + { + return false; + } + + if (a.m_needsKeepAlive != b.m_needsKeepAlive) + { + return false; + } + + if (a.m_needsExecutionContext != b.m_needsExecutionContext) + { + return false; + } + + if (a.m_returns.size() != b.m_returns.size()) + { + return false; + } + + for (size_t i = 0; i < a.m_returns.size(); i++) + { + if (a.m_returns[i].ReturnType != b.m_returns[i].ReturnType) + { + return false; + } + + if ((a.m_returns[i].ReturnType == TYP_STRUCT) && + !ClassLayout::AreCompatible(a.m_returns[i].ReturnLayout, b.m_returns[i].ReturnLayout)) + { + return false; + } + } + + if (a.m_locals.size() != b.m_locals.size()) + { + return false; + } + + for (size_t i = 0; i < a.m_locals.size(); i++) + { + if (a.m_locals[i] != b.m_locals[i]) + { + return false; + } + } + + return true; +} + //------------------------------------------------------------------------ // TransformAsync: Run async transformation. // @@ -944,17 +1017,31 @@ void AsyncTransformation::Transform(BasicBlock* block, CallDefinitionInfo callDefInfo = CanonicalizeCallDefinition(block, call, &analyses); - unsigned stateNum = (unsigned)m_states.size(); - JITDUMP(" Assigned state %u\n", stateNum); + const AsyncState* reusedState = + FindReusableSuspension(block, call, callDefInfo, layoutBuilder, resumeReachable, mutatedSinceResumption); + + if (reusedState != nullptr) + { + JITDUMP(" Reused state %u\n", reusedState->Number); + CreateCheckAndSuspendAfterCall(block, call, callDefInfo, reusedState->SuspensionBB, remainder); + + HandleReusedSuspension(block, call); + m_compiler->Metrics.SuspensionPointsMerged++; + } + else + { + unsigned stateNum = (unsigned)m_states.size(); + JITDUMP(" Assigned state %u\n", stateNum); - BasicBlock* suspendBB = CreateSuspensionBlock(block, stateNum); + BasicBlock* suspendBB = CreateSuspensionBlock(block, stateNum); - CreateCheckAndSuspendAfterCall(block, call, callDefInfo, suspendBB, remainder); + CreateCheckAndSuspendAfterCall(block, call, callDefInfo, suspendBB, remainder); - BasicBlock* resumeBB = CreateResumptionBlock(*remainder, stateNum); + BasicBlock* resumeBB = CreateResumptionBlock(*remainder, stateNum); - m_states.push_back(AsyncState(stateNum, layoutBuilder, block, call, callDefInfo, suspendBB, resumeBB, - resumeReachable, mutatedSinceResumption)); + m_states.push_back(AsyncState(stateNum, layoutBuilder, block, call, callDefInfo, suspendBB, resumeBB, + resumeReachable, mutatedSinceResumption)); + } JITDUMP("\n"); } @@ -1656,6 +1743,269 @@ CallDefinitionInfo AsyncTransformation::CanonicalizeCallDefinition(BasicBlock* return callDefInfo; } +//------------------------------------------------------------------------ +// AsyncTransformation::FindReusableSuspension: +// Try to find a suspension/resumption state that we can merge with. +// +// Parameters: +// block - The block containing the async call. +// call - The async call +// defInfo - Async call def info +// layoutBuilder - Layout for this async call +// resumeReachable - Whether 'call' is reachable after a previous resumption +// mutatedSinceResumption - Set of variables mutated since the last resumption +// +// Returns: +// Reusable existing state, or nullptr if no such state. +// +const AsyncState* AsyncTransformation::FindReusableSuspension(BasicBlock* block, + GenTreeCall* call, + const CallDefinitionInfo& defInfo, + ContinuationLayoutBuilder* layoutBuilder, + bool resumeReachable, + VARSET_VALARG_TP mutatedSinceResumption) +{ + if (m_compiler->opts.OptimizationDisabled()) + { + // Tail merging breaks debugging. + return nullptr; + } + + // Call must be at the tail of its block. + GenTree* thisLastNode = block->lastNode(); + if ((thisLastNode != call) && (thisLastNode != defInfo.DefinitionNode)) + { + return nullptr; + } + + // We can only merge if we logically could move the suspension point into a + // successor of two blocks. + BasicBlock* target = block->GetUniqueSucc(); + if (target == nullptr) + { + return nullptr; + } + + for (BasicBlock* pred : target->PredBlocks()) + { + if (pred == block) + { + continue; + } + + if (!pred->KindIs(BBJ_ALWAYS) || !pred->isEmpty()) + { + // We are looking for the "remainder" block that was created for a previous async transformation. + // If they are in a usable tail position they will always be empty BBJ_ALWAYS blocks. + continue; + } + + int limit = 128; + for (size_t i = m_states.size(); i != 0; i--) + { + if (limit-- == 0) + { + break; + } + + const AsyncState& state = m_states[i - 1]; + + assert(state.ResumptionBB->KindIs(BBJ_ALWAYS)); + BasicBlock* resumptionJoin = state.ResumptionBB->GetUniqueSucc(); + + // This existing async state resumes at a sibling join block. We + // can reuse that resumption if it has the same state as us. + if ((resumptionJoin == pred) && IsReusableSuspension(&state, block, call, defInfo, layoutBuilder, + resumeReachable, mutatedSinceResumption)) + { + return &state; + } + } + } + + return nullptr; +} + +//------------------------------------------------------------------------ +// AsyncTransformation::IsReusableSuspension: +// Check if the suspension for a specified state is reusable for an async call. +// +// Parameters: +// state - State for existing suspension. +// block - The block containing the async call. +// call - The async call +// defInfo - Async call def info +// layoutBuilder - Layout for this async call +// resumeReachable - Whether 'call' is reachable after a previous resumption +// mutatedSinceResumption - Set of variables mutated since the last resumption +// +// Returns: +// True if the suspension is reusable. +// +bool AsyncTransformation::IsReusableSuspension(const AsyncState* state, + BasicBlock* block, + GenTreeCall* call, + const CallDefinitionInfo& defInfo, + ContinuationLayoutBuilder* layoutBuilder, + bool resumeReachable, + VARSET_VALARG_TP mutatedSinceResumption) +{ + GenTreeCall* predAsyncCall = state->Call; + + JITDUMP( + " Sibling to the join is resumption for async call [%06u]; checking for possible tail merging of suspension points\n", + Compiler::dspTreeID(predAsyncCall)); + + if (!BasicBlock::sameEHRegion(block, state->CallBlock)) + { + JITDUMP(" Not same EH region\n"); + return false; + } + + // We have a sibling of the join that ends with an async call. Check if it is compatible. + if ((defInfo.DefinitionNode == nullptr) != (state->CallDefInfo.DefinitionNode == nullptr)) + { + JITDUMP(" No; disagreement on presence of return value\n"); + return false; + } + + if ((defInfo.DefinitionNode != nullptr) && + !GenTreeLclVarCommon::EqualsLocal(defInfo.DefinitionNode, state->CallDefInfo.DefinitionNode)) + { + JITDUMP(" No; disagreement on return value destination ([%06u] does not equal [%06u])\n", + Compiler::dspTreeID(defInfo.DefinitionNode), Compiler::dspTreeID(state->CallDefInfo.DefinitionNode)); + return false; + } + + if (call->gtReturnType != predAsyncCall->AsCall()->gtReturnType) + { + JITDUMP(" No; disagreement on return type (%s vs %s)\n", varTypeName(call->gtReturnType), + varTypeName(predAsyncCall->AsCall()->gtReturnType)); + return false; + } + + if (call->gtReturnType == TYP_STRUCT) + { + ClassLayout* thisLayout = m_compiler->typGetObjLayout(call->gtRetClsHnd); + ClassLayout* otherLayout = m_compiler->typGetObjLayout(predAsyncCall->AsCall()->gtRetClsHnd); + if (!ClassLayout::AreCompatible(thisLayout, otherLayout)) + { + JITDUMP(" No; disagreement on return type (%s vs %s)\n", thisLayout->GetClassName(), + otherLayout->GetClassName()); + return false; + } + } + + const AsyncCallInfo& asyncInfoThis = call->GetAsyncInfo(); + const AsyncCallInfo& asyncInfoOther = predAsyncCall->AsCall()->GetAsyncInfo(); + + if (asyncInfoThis.ContinuationContextHandling != asyncInfoOther.ContinuationContextHandling) + { + JITDUMP(" No; disagreement on continuation context handling (%u vs %u)\n", + asyncInfoThis.ContinuationContextHandling, asyncInfoOther.ContinuationContextHandling); + return false; + } + + if (asyncInfoThis.NeedsToSaveAndRestoreExecutionContext() != asyncInfoOther.NeedsToSaveAndRestoreExecutionContext()) + { + JITDUMP(" No; disagreement on whether execution context needs to be saved and restored (%s vs %s)\n", + asyncInfoThis.NeedsToSaveAndRestoreExecutionContext() ? "yes" : "no", + asyncInfoOther.NeedsToSaveAndRestoreExecutionContext() ? "yes" : "no"); + return false; + } + + if (state->ResumeReachable != resumeReachable) + { + // Not being resume reachable means we can skip checking for a + // reusable continuation. We do not want to give up that + // optimization. + JITDUMP(" No; disagreement on resume reachability (%s vs %s)\n", state->ResumeReachable ? "yes" : "no", + resumeReachable ? "yes" : "no"); + return false; + } + + if (resumeReachable) + { + // If both are reachable after resuming then we need to take care that + // we update the right set of locals in the continuation on suspension. + // In one path we may have mutated a local that we didn't in another + // path, and that results in a difference when we reuse a continuation + // and need to decide if that local needs to be saved again or not. + for (unsigned lclNum : layoutBuilder->Locals()) + { + LclVarDsc* dsc = m_compiler->lvaGetDesc(lclNum); + if (GetLocalSaveSet(dsc, mutatedSinceResumption) != + GetLocalSaveSet(dsc, state->MutatedSincePreviousResumption)) + { + JITDUMP(" No; disagreement on save set for V%02u (%d vs %d)\n", lclNum, + (unsigned)GetLocalSaveSet(dsc, mutatedSinceResumption), + (unsigned)GetLocalSaveSet(dsc, state->MutatedSincePreviousResumption)); + return false; + } + } + } + + static const WellKnownArg validateArgs[] = {WellKnownArg::AsyncExecutionContext, + WellKnownArg::AsyncSynchronizationContext}; + for (WellKnownArg arg : validateArgs) + { + CallArg* thisArg = call->gtArgs.FindWellKnownArg(arg); + CallArg* otherArg = predAsyncCall->gtArgs.FindWellKnownArg(arg); + + if ((thisArg == nullptr) != (otherArg == nullptr)) + { + JITDUMP(" No; disagreement on presence of %s argument\n", getWellKnownArgName(arg)); + return false; + } + + if (thisArg != nullptr) + { + if (!thisArg->GetNode()->OperIsLocal()) + { + JITDUMP(" No; %s argument is too complex\n", getWellKnownArgName(arg)); + return false; + } + + if (!GenTree::Compare(thisArg->GetNode(), otherArg->GetNode())) + { + JITDUMP(" No; disagreement on value of %s argument ([%06u] vs [%06u])\n", getWellKnownArgName(arg), + Compiler::dspTreeID(thisArg->GetNode()), Compiler::dspTreeID(otherArg->GetNode())); + return false; + } + } + } + + // At this point we expect both suspensions to agree on the live locals + // since they both join at the same point. So we only need to assert that + // the layouts match. + assert(ContinuationLayoutBuilder::Equals(*layoutBuilder, *state->Layout)); + return true; +} + +//------------------------------------------------------------------------ +// AsyncTransformation::HandleReusedSuspension: +// Handle the case where we were able to reuse a previously created suspension state. +// +// Parameters: +// callBlock - Block containing async call +// call - The async call +// +void AsyncTransformation::HandleReusedSuspension(BasicBlock* callBlock, GenTreeCall* call) +{ + static const WellKnownArg argsToRemove[] = {WellKnownArg::AsyncExecutionContext, + WellKnownArg::AsyncSynchronizationContext}; + for (WellKnownArg wka : argsToRemove) + { + CallArg* arg = call->gtArgs.FindWellKnownArg(wka); + if (arg != nullptr) + { + assert(arg->GetNode()->OperIsLocal()); + LIR::AsRange(callBlock).Remove(arg->GetNode()); + call->gtArgs.RemoveUnsafe(arg); + } + } +} + //------------------------------------------------------------------------ // AsyncTransformation::CreateSuspensionBlock: // Create an empty basic block that will hold the suspension IR for the @@ -1790,7 +2140,7 @@ void AsyncTransformation::CreateSuspension(BasicBlock* call // In the path where we allocated a new continuation we save only locals that we know to be unmutated since the // last resumption. - FillInDataOnSuspension(call, layout, subLayout, allocNewBB, mutatedSinceResumption, SaveSet::UnmutatedLocals); + FillInDataOnSuspension(layout, subLayout, allocNewBB, mutatedSinceResumption, SaveSet::UnmutatedLocals); // We can skip saving unmutated locals in the shared path -- we only need to save locals that may have been // mutated since the last resumption. @@ -1868,7 +2218,7 @@ void AsyncTransformation::CreateSuspension(BasicBlock* call GenTree* storeFlags = StoreAtOffset(newContinuation, flagsOffset, flagsNode, TYP_INT); LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_compiler, storeFlags)); - FillInDataOnSuspension(call, layout, subLayout, suspendBB, mutatedSinceResumption, tailSaveSet); + FillInDataOnSuspension(layout, subLayout, suspendBB, mutatedSinceResumption, tailSaveSet); FinishContextHandlingAndSuspension(callBlock, call, suspendBB, layout, subLayout); } @@ -1918,13 +2268,13 @@ GenTreeCall* AsyncTransformation::CreateAllocContinuationCall(bool // context. // // Parameters: -// call - The async call. // layout - Information about the continuation layout. // subLayout - Per-call layout builder indicating which fields are needed. // suspendBB - Basic block to add IR to. +// mutatedSinceResumption - Set of locals mutated since the last resumption. +// saveSet - Set of locals to save // -void AsyncTransformation::FillInDataOnSuspension(GenTreeCall* call, - const ContinuationLayout& layout, +void AsyncTransformation::FillInDataOnSuspension(const ContinuationLayout& layout, const ContinuationLayoutBuilder& subLayout, BasicBlock* suspendBB, VARSET_VALARG_TP mutatedSinceResumption, diff --git a/src/coreclr/jit/async.h b/src/coreclr/jit/async.h index ce25cd991362cf..6e194fef03019b 100644 --- a/src/coreclr/jit/async.h +++ b/src/coreclr/jit/async.h @@ -119,6 +119,8 @@ struct ContinuationLayoutBuilder return m_locals; } + static bool Equals(const ContinuationLayoutBuilder& a, const ContinuationLayoutBuilder& b); + struct ContinuationLayout* Create(); static ContinuationLayoutBuilder* CreateSharedLayout(Compiler* comp, @@ -404,6 +406,21 @@ class AsyncTransformation CallDefinitionInfo CanonicalizeCallDefinition(BasicBlock* block, GenTreeCall* call, AsyncAnalysis* analyses); + const AsyncState* FindReusableSuspension(BasicBlock* block, + GenTreeCall* call, + const CallDefinitionInfo& defInfo, + ContinuationLayoutBuilder* layoutBuilder, + bool resumeReachable, + VARSET_VALARG_TP mutatedSinceResumption); + bool IsReusableSuspension(const AsyncState* state, + BasicBlock* block, + GenTreeCall* call, + const CallDefinitionInfo& defInfo, + ContinuationLayoutBuilder* layoutBuilder, + bool resumeReachable, + VARSET_VALARG_TP mutatedSinceResumption); + void HandleReusedSuspension(BasicBlock* callBlock, GenTreeCall* call); + BasicBlock* CreateSuspensionBlock(BasicBlock* block, unsigned stateNum); void CreateSuspension(BasicBlock* callBlock, GenTreeCall* call, @@ -418,8 +435,7 @@ class AsyncTransformation GenTree* prevContinuation, const ContinuationLayout& layout); - void FillInDataOnSuspension(GenTreeCall* call, - const ContinuationLayout& layout, + void FillInDataOnSuspension(const ContinuationLayout& layout, const ContinuationLayoutBuilder& subLayout, BasicBlock* suspendBB, VARSET_VALARG_TP mutatedSinceResumption, @@ -443,13 +459,14 @@ class AsyncTransformation const CallDefinitionInfo& callDefInfo, BasicBlock* suspendBB, BasicBlock** remainder); - BasicBlock* CreateResumptionBlock(BasicBlock* remainder, unsigned stateNum); - void CreateResumption(BasicBlock* callBlock, - GenTreeCall* call, - BasicBlock* resumeBB, - const CallDefinitionInfo& callDefInfo, - const ContinuationLayout& layout, - const ContinuationLayoutBuilder& subLayout); + + BasicBlock* CreateResumptionBlock(BasicBlock* remainder, unsigned stateNum); + void CreateResumption(BasicBlock* callBlock, + GenTreeCall* call, + BasicBlock* resumeBB, + const CallDefinitionInfo& callDefInfo, + const ContinuationLayout& layout, + const ContinuationLayoutBuilder& subLayout); void RestoreFromDataOnResumption(const ContinuationLayout& layout, const ContinuationLayoutBuilder& subLayout, diff --git a/src/coreclr/jit/gentree.cpp b/src/coreclr/jit/gentree.cpp index 74fb96430251ce..d5d2fb86a90960 100644 --- a/src/coreclr/jit/gentree.cpp +++ b/src/coreclr/jit/gentree.cpp @@ -33107,6 +33107,47 @@ ClassLayout* GenTreeLclVarCommon::GetLayout(Compiler* compiler) const return AsLclFld()->GetLayout(); } +//------------------------------------------------------------------------ +// EqualsLocal: +// Check if the locals information of two locals is equal. +// +// Arguments: +// lcl1 - The first local +// lcl2 - The second local +// +// Returns: +// True if equal. +// +// Remarks: +// For LCL_VAR, LCL_FLD and LCL_ADDR this is equivalent to GenTree::Compare. +// For STORE_LCL_VAR and STORE_LCL_FLD this only checks the local +// information; it does not check the data node for equality. +// +bool GenTreeLclVarCommon::EqualsLocal(GenTreeLclVarCommon* lcl1, GenTreeLclVarCommon* lcl2) +{ + if (lcl1->OperGet() != lcl2->OperGet()) + { + return false; + } + + if (lcl1->GetLclNum() != lcl2->GetLclNum()) + { + return false; + } + + if (lcl1->GetLclOffs() != lcl2->GetLclOffs()) + { + return false; + } + + if (lcl1->OperIs(GT_LCL_FLD, GT_STORE_LCL_FLD) && lcl1->AsLclFld()->GetLayout() != lcl2->AsLclFld()->GetLayout()) + { + return false; + } + + return true; +} + //------------------------------------------------------------------------ // GenTreeLclVar::IsNeverNegative: Gets true if the lcl var is never negative; otherwise false. // diff --git a/src/coreclr/jit/gentree.h b/src/coreclr/jit/gentree.h index 525c088b1db2c6..3c51fc9541e201 100644 --- a/src/coreclr/jit/gentree.h +++ b/src/coreclr/jit/gentree.h @@ -3851,6 +3851,8 @@ struct GenTreeLclVarCommon : public GenTreeUnOp { } #endif + + static bool EqualsLocal(GenTreeLclVarCommon* lcl1, GenTreeLclVarCommon* lcl2); }; //------------------------------------------------------------------------ diff --git a/src/coreclr/jit/jitmetadatalist.h b/src/coreclr/jit/jitmetadatalist.h index e750ec79dcab03..53c363ef5f24f2 100644 --- a/src/coreclr/jit/jitmetadatalist.h +++ b/src/coreclr/jit/jitmetadatalist.h @@ -94,6 +94,7 @@ JITMETADATAMETRIC(MorphTrackedLocals, int, 0) JITMETADATAMETRIC(MorphLocals, int, 0) JITMETADATAMETRIC(EnumeratorGDVProvisionalNoEscape, int, 0) JITMETADATAMETRIC(EnumeratorGDVCanCloneToEnsureNoEscape, int, 0) +JITMETADATAMETRIC(SuspensionPointsMerged, int, 0) #undef JITMETADATA #undef JITMETADATAINFO From 6aa2fcef7c081f4284e614c5407571d4f5dc3e25 Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Thu, 28 May 2026 19:15:29 +0200 Subject: [PATCH 3/3] Just make it a check instead --- src/coreclr/jit/async.cpp | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/coreclr/jit/async.cpp b/src/coreclr/jit/async.cpp index b43572225035e5..a535d7a0a741f1 100644 --- a/src/coreclr/jit/async.cpp +++ b/src/coreclr/jit/async.cpp @@ -1975,10 +1975,13 @@ bool AsyncTransformation::IsReusableSuspension(const AsyncState* state, } } - // At this point we expect both suspensions to agree on the live locals - // since they both join at the same point. So we only need to assert that - // the layouts match. - assert(ContinuationLayoutBuilder::Equals(*layoutBuilder, *state->Layout)); + if (!ContinuationLayoutBuilder::Equals(*layoutBuilder, *state->Layout)) + { + // TODO: We would expect liveness to match in tail positions here. But + // sometimes that isn't the case -- why not? + return false; + } + return true; }