Skip to content

JIT: Avoid TLS access when restoring async contexts #127889

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jakobbotsch wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

JIT: Avoid TLS access when restoring async contexts #127889

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
874e2c7
JIT: Avoid TLS access when restoring async contexts
jakobbotsch May 6, 2026
e121c26
Feedback
jakobbotsch May 6, 2026
4c1b874
Run jit-format
jakobbotsch May 6, 2026
fab33bb
Bump JIT-EE GUID
jakobbotsch May 8, 2026
00f22c7
Fix bug
jakobbotsch May 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions ...oreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -982,23 +982,23 @@ private static void RestoreExecutionContext(Thread thread, ExecutionContext? pre
}

[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static void CaptureContexts(out ExecutionContext? execCtx, out SynchronizationContext? syncCtx)
private static void CaptureContexts(out Thread thread, out ExecutionContext? execCtx, out SynchronizationContext? syncCtx)
{
Thread thread = Thread.CurrentThreadAssumedInitialized;
Thread curThread = Thread.CurrentThreadAssumedInitialized;
thread = curThread;
// Here we get the execution context for synchronous restoring,
// not for flowing across suspension to potentially another thread.
// Therefore we do not need to worry about IsFlowSuppressed
execCtx = thread._executionContext;
syncCtx = thread._synchronizationContext;
execCtx = curThread._executionContext;
syncCtx = curThread._synchronizationContext;
}

// Restore contexts onto current Thread. If "resumed" then this is not the first starting call for the async method.
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static void RestoreContexts(bool resumed, ExecutionContext? previousExecCtx, SynchronizationContext? previousSyncCtx)
private static void RestoreContexts(bool resumed, Thread thread, ExecutionContext? previousExecCtx, SynchronizationContext? previousSyncCtx)
{
if (!resumed)
{
Thread thread = Thread.CurrentThreadAssumedInitialized;
if (previousSyncCtx != thread._synchronizationContext)
{
Comment thread
jakobbotsch marked this conversation as resolved.
thread._synchronizationContext = previousSyncCtx;
Expand Down
10 changes: 5 additions & 5 deletions src/coreclr/inc/jiteeversionguid.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,11 @@

#include <minipal/guid.h>

constexpr GUID JITEEVersionIdentifier = { /* e92fbf65-4856-4729-ab9e-f66f7adcecf9 */
0xe92fbf65,
0x4856,
0x4729,
{0xab, 0x9e, 0xf6, 0x6f, 0x7a, 0xdc, 0xec, 0xf9}
constexpr GUID JITEEVersionIdentifier = { /* 744f01cc-80f0-4ff2-a447-52e0417a8149 */
0x744f01cc,
0x80f0,
0x4ff2,
{0xa4, 0x47, 0x52, 0xe0, 0x41, 0x7a, 0x81, 0x49}
};

#endif // JIT_EE_VERSIONING_GUID_H
42 changes: 30 additions & 12 deletions src/coreclr/inc/patchpointinfo.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,15 +39,16 @@ struct PatchpointInfo
// Initialize
void Initialize(uint32_t localCount, int32_t totalFrameSize)
{
m_calleeSaveRegisters = 0;
m_tier0Version = 0;
m_totalFrameSize = totalFrameSize;
m_numberOfLocals = localCount;
m_genericContextArgOffset = -1;
m_keptAliveThisOffset = -1;
m_securityCookieOffset = -1;
m_monitorAcquiredOffset = -1;
m_asyncExecutionContextOffset = -1;
m_calleeSaveRegisters = 0;
m_tier0Version = 0;
m_totalFrameSize = totalFrameSize;
m_numberOfLocals = localCount;
m_genericContextArgOffset = -1;
m_keptAliveThisOffset = -1;
m_securityCookieOffset = -1;
m_monitorAcquiredOffset = -1;
m_asyncThreadObjectOffset = -1;
m_asyncExecutionContextOffset = -1;
m_asyncSynchronizationContextOffset = -1;
}

Expand All @@ -60,6 +61,7 @@ struct PatchpointInfo
m_keptAliveThisOffset = original->m_keptAliveThisOffset;
m_securityCookieOffset = original->m_securityCookieOffset;
m_monitorAcquiredOffset = original->m_monitorAcquiredOffset;
m_asyncThreadObjectOffset = original->m_asyncThreadObjectOffset;
m_asyncExecutionContextOffset = original->m_asyncExecutionContextOffset;
m_asyncSynchronizationContextOffset = original->m_asyncSynchronizationContextOffset;

Expand Down Expand Up @@ -151,13 +153,28 @@ struct PatchpointInfo
m_monitorAcquiredOffset = offset;
}

// Original method FP relative offset for async contexts
// Original method FP relative offset for async thread/contexts
int32_t AsyncThreadOffset() const
{
return m_asyncThreadObjectOffset;
}

bool HasAsyncThread() const
{
return m_asyncThreadObjectOffset != -1;
}

void SetAsyncThreadOffset(int32_t offset)
{
m_asyncThreadObjectOffset = offset;
}

int32_t AsyncExecutionContextOffset() const
{
return m_asyncExecutionContextOffset;
}

bool HasAsyncExecutionContextOffset() const
bool HasAsyncExecutionContext() const
{
return m_asyncExecutionContextOffset != -1;
}
Expand All @@ -172,7 +189,7 @@ struct PatchpointInfo
return m_asyncSynchronizationContextOffset;
}

bool HasAsyncSynchronizationContextOffset() const
bool HasAsyncSynchronizationContext() const
{
return m_asyncSynchronizationContextOffset != -1;
}
Expand Down Expand Up @@ -239,6 +256,7 @@ struct PatchpointInfo
int32_t m_keptAliveThisOffset;
int32_t m_securityCookieOffset;
int32_t m_monitorAcquiredOffset;
int32_t m_asyncThreadObjectOffset;
int32_t m_asyncExecutionContextOffset;
int32_t m_asyncSynchronizationContextOffset;
int32_t m_offsetAndExposureData[];
Expand Down
49 changes: 36 additions & 13 deletions src/coreclr/interpreter/compiler.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2330,7 +2330,7 @@ void InterpCompiler::CreateILVars()
// add some starting extra space for new vars
m_varsCapacity = m_numILVars + getEHcount(m_methodInfo) + 64;
m_pVars = getAllocator(IMK_Var).allocateZeroed<InterpVar>(m_varsCapacity);
m_varsSize = m_numILVars + hasParamArg + hasContinuationArg + (hasThisPointerShadowCopyAsParamIndex ? 1 : 0) + (m_isAsyncMethodWithContextSaveRestore ? 2 : 0) + getEHcount(m_methodInfo);
m_varsSize = m_numILVars + hasParamArg + hasContinuationArg + (hasThisPointerShadowCopyAsParamIndex ? 1 : 0) + (m_isAsyncMethodWithContextSaveRestore ? 3 : 0) + getEHcount(m_methodInfo);

offset = 0;

Expand Down Expand Up @@ -2431,10 +2431,16 @@ void InterpCompiler::CreateILVars()

if (m_isAsyncMethodWithContextSaveRestore)
{
m_threadObjVarIndex = index;
CreateNextLocalVar(index, NULL, InterpTypeO, &offset);
INTERP_DUMP("alloc Async Thread (var %d) to offset %d\n", m_threadObjVarIndex, m_pVars[m_threadObjVarIndex].offset);
index++;

m_execContextVarIndex = index;
CreateNextLocalVar(index, NULL, InterpTypeO, &offset);
INTERP_DUMP("alloc ExecutableContextVar (var %d) to offset %d\n", m_execContextVarIndex, m_pVars[m_execContextVarIndex].offset);
index++;

m_syncContextVarIndex = index;
CreateNextLocalVar(index, NULL, InterpTypeO, &offset);
INTERP_DUMP("alloc SyncContextVar (var %d) to offset %d\n", m_syncContextVarIndex, m_pVars[m_syncContextVarIndex].offset);
Comment thread
jakobbotsch marked this conversation as resolved.
Expand Down Expand Up @@ -4739,7 +4745,14 @@ void InterpCompiler::EmitCall(CORINFO_RESOLVED_TOKEN* pConstrainedToken, bool re
m_pLastNewIns->SetSVar(m_continuationArgIndex);
PushInterpType(InterpTypeI4, NULL);
m_pLastNewIns->SetDVar(m_pStackPointer[-1].var);
int32_t isStartedArg = m_pStackPointer[-1].var;
int32_t isResumedArg = m_pStackPointer[-1].var;
m_pStackPointer--;

AddIns(INTOP_MOV_P);
m_pLastNewIns->SetSVar(m_threadObjVarIndex);
PushInterpType(InterpTypeO, NULL);
m_pLastNewIns->SetDVar(m_pStackPointer[-1].var);
int32_t threadAddressVar = m_pStackPointer[-1].var;
m_pStackPointer--;

AddIns(INTOP_MOV_P);
Expand Down Expand Up @@ -4773,12 +4786,13 @@ void InterpCompiler::EmitCall(CORINFO_RESOLVED_TOKEN* pConstrainedToken, bool re

m_pLastNewIns->flags |= INTERP_INST_FLAG_CALL;
m_pLastNewIns->info.pCallInfo = new (getAllocator(IMK_CallInfo)) InterpCallInfo();
int32_t numArgs = 3;
int32_t numArgs = 4;
int32_t *callArgs = getAllocator(IMK_CallInfo).allocate<int32_t>(numArgs + 1);
callArgs[0] = isStartedArg;
callArgs[1] = execContextAddressVar;
callArgs[2] = syncContextAddressVar;
callArgs[3] = CALL_ARGS_TERMINATOR;
callArgs[0] = isResumedArg;
callArgs[1] = threadAddressVar;
callArgs[2] = execContextAddressVar;
callArgs[3] = syncContextAddressVar;
callArgs[4] = CALL_ARGS_TERMINATOR;
m_pLastNewIns->info.pCallInfo->pCallArgs = callArgs;

m_ip += 5;
Expand Down Expand Up @@ -8040,7 +8054,7 @@ void InterpCompiler::GenerateCode(CORINFO_METHOD_INFO* methodInfo)
INTERP_DUMP("Synchronized method - adding extra IL opcodes for async save/restore\n");

// Async methods are methods implemented by adding a try/finally in the method
// which takes the lock on entry and releases it on exit. To integrate this into the interpreter, we actually
// which captures contexts on entry and restores them on exit. To integrate this into the interpreter, we actually
// add a set of extra "IL" opcodes at the end of the method which do the monitor finally and actual return
// logic.

Expand Down Expand Up @@ -8217,7 +8231,15 @@ void InterpCompiler::GenerateCode(CORINFO_METHOD_INFO* methodInfo)

if (m_isAsyncMethodWithContextSaveRestore)
{
// Load the address of the execution context/sync context locals, and call AsyncHelpers.CaptureContexts
// Load the address of the thread/execution context/sync context locals, and call AsyncHelpers.CaptureContexts
AddIns(INTOP_LDLOCA);
m_pLastNewIns->SetSVar(m_threadObjVarIndex);
PushInterpType(InterpTypeByRef, NULL);
m_pStackPointer[-1].SetAsLocalVariableAddress();
m_pLastNewIns->SetDVar(m_pStackPointer[-1].var);
int32_t threadAddressVar = m_pStackPointer[-1].var;
m_pStackPointer--;

AddIns(INTOP_LDLOCA);
m_pLastNewIns->SetSVar(m_execContextVarIndex);
PushInterpType(InterpTypeByRef, NULL);
Expand Down Expand Up @@ -8252,11 +8274,12 @@ void InterpCompiler::GenerateCode(CORINFO_METHOD_INFO* methodInfo)

m_pLastNewIns->flags |= INTERP_INST_FLAG_CALL;
m_pLastNewIns->info.pCallInfo = new (getAllocator(IMK_CallInfo)) InterpCallInfo();
int32_t numArgs = 2;
int32_t numArgs = 3;
int32_t *callArgs = getAllocator(IMK_CallInfo).allocate<int32_t>(numArgs + 1);
callArgs[0] = execContextAddressVar;
callArgs[1] = syncContextAddressVar;
callArgs[2] = CALL_ARGS_TERMINATOR;
callArgs[0] = threadAddressVar;
callArgs[1] = execContextAddressVar;
callArgs[2] = syncContextAddressVar;
callArgs[3] = CALL_ARGS_TERMINATOR;
m_pLastNewIns->info.pCallInfo->pCallArgs = callArgs;
Comment thread
jakobbotsch marked this conversation as resolved.
}

Expand Down
1 change: 1 addition & 0 deletions src/coreclr/interpreter/compiler.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -889,6 +889,7 @@ class InterpCompiler
int32_t m_synchronizedOrAsyncRetValVarIndex = -1; // If the method is synchronized, ret instructions are replaced with a store to this var and a leave to an epilog instruction.
int32_t m_synchronizedFinallyStartOffset = -1; // If the method is synchronized, this is the offset of the start of the finally epilog

int32_t m_threadObjVarIndex = -1; // If the method is async, this is the var index of the Thread local
int32_t m_execContextVarIndex = -1; // If the method is async, this is the var index of the ExecutionContext local
int32_t m_syncContextVarIndex = -1; // If the method is async, this is the var index of the SynchronizationContext local

Expand Down
18 changes: 15 additions & 3 deletions src/coreclr/jit/async.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,10 @@ PhaseStatus Compiler::SaveAsyncContexts()
return PhaseStatus::MODIFIED_NOTHING;
}

// Create locals for ExecutionContext and SynchronizationContext
// Create locals for Thread, ExecutionContext and SynchronizationContext
lvaAsyncThreadObjectVar = lvaGrabTemp(false DEBUGARG("Async Thread"));
lvaGetDesc(lvaAsyncThreadObjectVar)->lvType = TYP_REF;

lvaAsyncExecutionContextVar = lvaGrabTemp(false DEBUGARG("Async ExecutionContext"));
lvaGetDesc(lvaAsyncExecutionContextVar)->lvType = TYP_REF;

Expand All @@ -97,6 +100,7 @@ PhaseStatus Compiler::SaveAsyncContexts()

if (opts.IsOSR())
{
lvaGetDesc(lvaAsyncThreadObjectVar)->lvIsOSRLocal = true;
lvaGetDesc(lvaAsyncExecutionContextVar)->lvIsOSRLocal = true;
lvaGetDesc(lvaAsyncSynchronizationContextVar)->lvIsOSRLocal = true;
}
Expand Down Expand Up @@ -190,8 +194,10 @@ PhaseStatus Compiler::SaveAsyncContexts()
captureCall->gtArgs.PushFront(this,
NewCallArg::Primitive(gtNewLclAddrNode(lvaAsyncSynchronizationContextVar, 0)));
captureCall->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclAddrNode(lvaAsyncExecutionContextVar, 0)));
lvaGetDesc(lvaAsyncSynchronizationContextVar)->lvHasLdAddrOp = true;
captureCall->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclAddrNode(lvaAsyncThreadObjectVar, 0)));
lvaGetDesc(lvaAsyncThreadObjectVar)->lvHasLdAddrOp = true;
lvaGetDesc(lvaAsyncExecutionContextVar)->lvHasLdAddrOp = true;
lvaGetDesc(lvaAsyncSynchronizationContextVar)->lvHasLdAddrOp = true;

CORINFO_CALL_INFO callInfo = {};
callInfo.hMethod = captureCall->gtCallMethHnd;
Expand Down Expand Up @@ -224,6 +230,7 @@ PhaseStatus Compiler::SaveAsyncContexts()
restoreCall->gtArgs.PushFront(this,
NewCallArg::Primitive(gtNewLclVarNode(lvaAsyncSynchronizationContextVar, TYP_REF)));
restoreCall->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(lvaAsyncExecutionContextVar, TYP_REF)));
restoreCall->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(lvaAsyncThreadObjectVar, TYP_REF)));
restoreCall->gtArgs.PushFront(this, NewCallArg::Primitive(resumed));

Statement* restoreStmt = fgNewStmtFromTree(restoreCall);
Expand Down Expand Up @@ -391,6 +398,7 @@ BasicBlock* Compiler::CreateReturnBB(unsigned* mergedReturnLcl)
restoreCall->gtArgs.PushFront(this,
NewCallArg::Primitive(gtNewLclVarNode(lvaAsyncSynchronizationContextVar, TYP_REF)));
restoreCall->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(lvaAsyncExecutionContextVar, TYP_REF)));
restoreCall->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewLclVarNode(lvaAsyncThreadObjectVar, TYP_REF)));
restoreCall->gtArgs.PushFront(this, NewCallArg::Primitive(resumed));

// This restore is an inline candidate (unlike the fault one)
Expand Down Expand Up @@ -993,6 +1001,10 @@ void AsyncTransformation::CreateLiveSetForSuspension(BasicBlock*
call->VisitLocalDefs(m_compiler, visitDef);

// Exclude method-level context locals (only live on synchronous path)
if (m_compiler->lvaAsyncThreadObjectVar != BAD_VAR_NUM)
{
excludedLocals.AddOrUpdate(m_compiler->lvaAsyncThreadObjectVar, true);
}
if (m_compiler->lvaAsyncSynchronizationContextVar != BAD_VAR_NUM)
{
excludedLocals.AddOrUpdate(m_compiler->lvaAsyncSynchronizationContextVar, true);
Expand Down Expand Up @@ -2307,7 +2319,7 @@ void AsyncTransformation::RestoreContexts(BasicBlock* block, GenTreeCall* call,
JITDUMP(" Call [%06u] has async contexts; will restore on suspension\n", Compiler::dspTreeID(call));

// Insert call
// AsyncHelpers.RestoreContexts(resumed, execContext, syncContext);
// AsyncHelpers.RestoreContextsOnSuspension(resumed, execContext, syncContext);

GenTree* resumedPlaceholder = m_compiler->gtNewIconNode(0);
GenTree* execContextPlaceholder = m_compiler->gtNewNull();
Expand Down
4 changes: 4 additions & 0 deletions src/coreclr/jit/codegenarm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2491,6 +2491,10 @@ void CodeGen::genCaptureFuncletPrologEpilogInfo()
{
saveSizeWithPSP += TARGET_POINTER_SIZE;
}
if (m_compiler->lvaAsyncThreadObjectVar != BAD_VAR_NUM)
{
saveSizeWithPSP += TARGET_POINTER_SIZE;
}
if (m_compiler->lvaAsyncExecutionContextVar != BAD_VAR_NUM)
{
saveSizeWithPSP += TARGET_POINTER_SIZE;
Expand Down
5 changes: 5 additions & 0 deletions src/coreclr/jit/codegenarm64.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1716,6 +1716,11 @@ void CodeGen::genCaptureFuncletPrologEpilogInfo()
saveRegsSize += m_compiler->lvaLclStackHomeSize(m_compiler->lvaMonAcquired);
}

if ((m_compiler->lvaAsyncThreadObjectVar != BAD_VAR_NUM) && !m_compiler->opts.IsOSR())
{
saveRegsSize += m_compiler->lvaLclStackHomeSize(m_compiler->lvaAsyncThreadObjectVar);
}

if ((m_compiler->lvaAsyncExecutionContextVar != BAD_VAR_NUM) && !m_compiler->opts.IsOSR())
{
saveRegsSize += m_compiler->lvaLclStackHomeSize(m_compiler->lvaAsyncExecutionContextVar);
Expand Down
7 changes: 7 additions & 0 deletions src/coreclr/jit/codegenarmarch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3843,6 +3843,13 @@ void CodeGen::genCreateAndStoreGCInfo(unsigned codeSize,
assert(m_compiler->lvaGetCallerSPRelativeOffset(m_compiler->lvaMonAcquired) == -preservedAreaSize);
}

if (m_compiler->lvaAsyncThreadObjectVar != BAD_VAR_NUM)
{
preservedAreaSize += TARGET_POINTER_SIZE;

assert(m_compiler->lvaGetCallerSPRelativeOffset(m_compiler->lvaAsyncThreadObjectVar) == -preservedAreaSize);
}

if (m_compiler->lvaAsyncExecutionContextVar != BAD_VAR_NUM)
{
preservedAreaSize += TARGET_POINTER_SIZE;
Expand Down
3 changes: 2 additions & 1 deletion src/coreclr/jit/codegencommon.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2322,7 +2322,8 @@ void CodeGen::genEmitMachineCode()
trackedStackPtrsContig = false;
#else
// We try to allocate GC pointers contiguously on the stack except for some special cases.
trackedStackPtrsContig = !m_compiler->opts.compDbgEnC && (m_compiler->lvaAsyncExecutionContextVar == BAD_VAR_NUM) &&
trackedStackPtrsContig = !m_compiler->opts.compDbgEnC && (m_compiler->lvaAsyncThreadObjectVar == BAD_VAR_NUM) &&
(m_compiler->lvaAsyncExecutionContextVar == BAD_VAR_NUM) &&
(m_compiler->lvaAsyncSynchronizationContextVar == BAD_VAR_NUM);

#ifdef TARGET_ARM
Expand Down
7 changes: 7 additions & 0 deletions src/coreclr/jit/codegenxarch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8552,6 +8552,13 @@ void CodeGen::genCreateAndStoreGCInfoX64(unsigned codeSize, unsigned prologSize
assert(m_compiler->lvaGetCallerSPRelativeOffset(m_compiler->lvaMonAcquired) == -preservedAreaSize);
}

if (m_compiler->lvaAsyncThreadObjectVar != BAD_VAR_NUM)
{
preservedAreaSize += TARGET_POINTER_SIZE;

assert(m_compiler->lvaGetCallerSPRelativeOffset(m_compiler->lvaAsyncThreadObjectVar) == -preservedAreaSize);
}

if (m_compiler->lvaAsyncExecutionContextVar != BAD_VAR_NUM)
{
preservedAreaSize += TARGET_POINTER_SIZE;
Expand Down
Loading
Loading