Skip to content

Commit 87ecdeb

Browse files
VSadovCopilot
andauthored
[RuntimeAsync] A few fixes for issues found when enabling rt async in Libraries. (#119864)
* Fix AV due to GS cookie check overwriting the continuation return register * runtime async thunks for ValueTask-returning methods. * fix an assert for `(callInfo.SynchronizationContextLclNum != BAD_VAR_NUM)` * JIT format * Update src/coreclr/jit/codegenxarch.cpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 04a748b commit 87ecdeb

File tree

5 files changed

+102
-19
lines changed

5 files changed

+102
-19
lines changed

src/coreclr/jit/async.cpp

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,8 @@ PhaseStatus Compiler::SaveAsyncContexts()
169169
// Await is inside a try, need to insert try-finally around it.
170170
restoreBB = InsertTryFinallyForContextRestore(curBB, stmt, restoreAfterStmt);
171171
restoreAfterStmt = nullptr;
172+
// we have split the block that could have another await.
173+
nextBB = restoreBB->Next();
172174
#endif
173175
}
174176

@@ -1506,8 +1508,9 @@ void AsyncTransformation::FillInGCPointersOnSuspension(GenTreeCall*
15061508
if (layout.ContinuationContextGCDataIndex != UINT_MAX)
15071509
{
15081510
const AsyncCallInfo& callInfo = call->GetAsyncInfo();
1509-
assert(callInfo.SaveAndRestoreSynchronizationContextField &&
1510-
(callInfo.SynchronizationContextLclNum != BAD_VAR_NUM));
1511+
assert(callInfo.SaveAndRestoreSynchronizationContextField);
1512+
assert(callInfo.ExecutionContextHandling == ExecutionContextHandling::SaveAndRestore);
1513+
assert(callInfo.SynchronizationContextLclNum != BAD_VAR_NUM);
15111514

15121515
// Insert call
15131516
// AsyncHelpers.CaptureContinuationContext(

src/coreclr/jit/codegenxarch.cpp

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -106,14 +106,21 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
106106
// we are generating GS cookie check after a GT_RETURN block.
107107
// Note: On Amd64 System V RDX is an arg register - REG_ARG_2 - as well
108108
// as return register for two-register-returned structs.
109+
#ifdef TARGET_X86
110+
// Note: ARG_0 can be REG_ASYNC_CONTINUATION_RET
111+
// we will check for that later if we end up saving/restoring this.
112+
regGSCheck = REG_ARG_0;
113+
regNumber regGSCheckAlternative = REG_ARG_1;
114+
#else
115+
// these cannot be a part of any kind of return
116+
regGSCheck = REG_R8;
117+
regNumber regGSCheckAlternative = REG_R9;
118+
#endif
119+
109120
if (compiler->lvaKeepAliveAndReportThis() && compiler->lvaGetDesc(compiler->info.compThisArg)->lvIsInReg() &&
110-
(compiler->lvaGetDesc(compiler->info.compThisArg)->GetRegNum() == REG_ARG_0))
121+
(compiler->lvaGetDesc(compiler->info.compThisArg)->GetRegNum() == regGSCheck))
111122
{
112-
regGSCheck = REG_ARG_1;
113-
}
114-
else
115-
{
116-
regGSCheck = REG_ARG_0;
123+
regGSCheck = regGSCheckAlternative;
117124
}
118125
}
119126
else
@@ -158,6 +165,14 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
158165
{
159166
// AOT case - GS cookie value needs to be accessed through an indirection.
160167

168+
// if we use the continuation reg, the pop/push requires no-GC
169+
// this can happen only when AOT supports async on x86
170+
if (compiler->compIsAsync() && (regGSCheck == REG_ASYNC_CONTINUATION_RET))
171+
{
172+
regMaskGSCheck = RBM_ASYNC_CONTINUATION_RET;
173+
GetEmitter()->emitDisableGC();
174+
}
175+
161176
pushedRegs = genPushRegs(regMaskGSCheck, &byrefPushedRegs, &norefPushedRegs);
162177

163178
instGen_Set_Reg_To_Imm(EA_HANDLE_CNS_RELOC, regGSCheck, (ssize_t)compiler->gsGlobalSecurityCookieAddr);

src/coreclr/vm/asyncthunks.cpp

Lines changed: 65 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,46 @@ SigPointer MethodDesc::GetAsyncThunkResultTypeSig()
338338
return SigPointer(returnTypeSig, (DWORD)(returnTypeSigEnd - returnTypeSig));
339339
}
340340

341+
bool MethodDesc::IsValueTaskAsyncThunk()
342+
{
343+
_ASSERTE(IsAsyncThunkMethod());
344+
PCCOR_SIGNATURE pSigRaw;
345+
DWORD cSig;
346+
if (FAILED(GetMDImport()->GetSigOfMethodDef(GetMemberDef(), &cSig, &pSigRaw)))
347+
{
348+
_ASSERTE(!"Loaded MethodDesc should not fail to get signature");
349+
pSigRaw = NULL;
350+
cSig = 0;
351+
}
352+
353+
SigPointer pSig(pSigRaw, cSig);
354+
uint32_t callConvInfo;
355+
IfFailThrow(pSig.GetCallingConvInfo(&callConvInfo));
356+
357+
if ((callConvInfo & IMAGE_CEE_CS_CALLCONV_GENERIC) != 0)
358+
{
359+
// GenParamCount
360+
IfFailThrow(pSig.GetData(NULL));
361+
}
362+
363+
// ParamCount
364+
IfFailThrow(pSig.GetData(NULL));
365+
366+
// ReturnType comes now. Skip the modifiers.
367+
IfFailThrow(pSig.SkipCustomModifiers());
368+
369+
// here we should have something Task, ValueTask, Task<retType> or ValueTask<retType>
370+
BYTE bElementType;
371+
IfFailThrow(pSig.GetByte(&bElementType));
372+
373+
// skip ELEMENT_TYPE_GENERICINST
374+
if (bElementType == ELEMENT_TYPE_GENERICINST)
375+
IfFailThrow(pSig.GetByte(&bElementType));
376+
377+
_ASSERTE(bElementType == ELEMENT_TYPE_VALUETYPE || bElementType == ELEMENT_TYPE_CLASS);
378+
return bElementType == ELEMENT_TYPE_VALUETYPE;
379+
}
380+
341381
// Given a method Foo<T>, return a MethodSpec token for Foo<T> instantiated
342382
// with the result type from the current async method's return type. For
343383
// example, if "this" represents Task<List<T>> Foo<T>(), and "md" is
@@ -435,30 +475,32 @@ void MethodDesc::EmitAsyncMethodThunk(MethodDesc* pAsyncOtherVariant, MetaSig& m
435475
MethodDesc* mdIsCompleted;
436476
MethodDesc* mdGetResult;
437477

478+
bool isValueTask = IsValueTaskAsyncThunk();
479+
438480
if (msig.IsReturnTypeVoid())
439481
{
440-
pMTTask = CoreLibBinder::GetClass(CLASS__TASK);
441-
thTaskAwaiter = CoreLibBinder::GetClass(CLASS__TASK_AWAITER);
442-
mdGetAwaiter = CoreLibBinder::GetMethod(METHOD__TASK__GET_AWAITER);
443-
mdIsCompleted = CoreLibBinder::GetMethod(METHOD__TASK_AWAITER__GET_ISCOMPLETED);
444-
mdGetResult = CoreLibBinder::GetMethod(METHOD__TASK_AWAITER__GET_RESULT);
482+
pMTTask = CoreLibBinder::GetClass(isValueTask ? CLASS__VALUETASK : CLASS__TASK);
483+
thTaskAwaiter = CoreLibBinder::GetClass(isValueTask ? CLASS__VALUETASK_AWAITER : CLASS__TASK_AWAITER);
484+
mdGetAwaiter = CoreLibBinder::GetMethod(isValueTask ? METHOD__VALUETASK__GET_AWAITER : METHOD__TASK__GET_AWAITER);
485+
mdIsCompleted = CoreLibBinder::GetMethod(isValueTask ? METHOD__VALUETASK_AWAITER__GET_ISCOMPLETED : METHOD__TASK_AWAITER__GET_ISCOMPLETED);
486+
mdGetResult = CoreLibBinder::GetMethod(isValueTask ? METHOD__VALUETASK_AWAITER__GET_RESULT : METHOD__TASK_AWAITER__GET_RESULT);
445487
}
446488
else
447489
{
448490
TypeHandle thLogicalRetType = msig.GetRetTypeHandleThrowing();
449-
MethodTable* pMTTaskOpen = CoreLibBinder::GetClass(CLASS__TASK_1);
491+
MethodTable* pMTTaskOpen = CoreLibBinder::GetClass(isValueTask ? CLASS__VALUETASK_1 : CLASS__TASK_1);
450492
pMTTask = ClassLoader::LoadGenericInstantiationThrowing(pMTTaskOpen->GetModule(), pMTTaskOpen->GetCl(), Instantiation(&thLogicalRetType, 1)).GetMethodTable();
451-
MethodTable* pMTTaskAwaiterOpen = CoreLibBinder::GetClass(CLASS__TASK_AWAITER_1);
493+
MethodTable* pMTTaskAwaiterOpen = CoreLibBinder::GetClass(isValueTask ? CLASS__VALUETASK_AWAITER_1 : CLASS__TASK_AWAITER_1);
452494

453495
thTaskAwaiter = ClassLoader::LoadGenericInstantiationThrowing(pMTTaskAwaiterOpen->GetModule(), pMTTaskAwaiterOpen->GetCl(), Instantiation(&thLogicalRetType, 1));
454496

455-
mdGetAwaiter = CoreLibBinder::GetMethod(METHOD__TASK_1__GET_AWAITER);
497+
mdGetAwaiter = CoreLibBinder::GetMethod(isValueTask ? METHOD__VALUETASK_1__GET_AWAITER : METHOD__TASK_1__GET_AWAITER);
456498
mdGetAwaiter = MethodDesc::FindOrCreateAssociatedMethodDesc(mdGetAwaiter, pMTTask, FALSE, Instantiation(), FALSE);
457499

458-
mdIsCompleted = CoreLibBinder::GetMethod(METHOD__TASK_AWAITER_1__GET_ISCOMPLETED);
500+
mdIsCompleted = CoreLibBinder::GetMethod(isValueTask ? METHOD__VALUETASK_AWAITER_1__GET_ISCOMPLETED : METHOD__TASK_AWAITER_1__GET_ISCOMPLETED);
459501
mdIsCompleted = MethodDesc::FindOrCreateAssociatedMethodDesc(mdIsCompleted, thTaskAwaiter.GetMethodTable(), FALSE, Instantiation(), FALSE);
460502

461-
mdGetResult = CoreLibBinder::GetMethod(METHOD__TASK_AWAITER_1__GET_RESULT);
503+
mdGetResult = CoreLibBinder::GetMethod(isValueTask ? METHOD__VALUETASK_AWAITER_1__GET_RESULT : METHOD__TASK_AWAITER_1__GET_RESULT);
462504
mdGetResult = MethodDesc::FindOrCreateAssociatedMethodDesc(mdGetResult, thTaskAwaiter.GetMethodTable(), FALSE, Instantiation(), FALSE);
463505
}
464506

@@ -550,7 +592,19 @@ void MethodDesc::EmitAsyncMethodThunk(MethodDesc* pAsyncOtherVariant, MetaSig& m
550592
getResultToken = pCode->GetToken(mdGetResult);
551593
}
552594

553-
pCode->EmitCALLVIRT(getAwaiterToken, 1, 1);
595+
if (isValueTask)
596+
{
597+
LocalDesc valuetaskLocalDesc(pMTTask);
598+
DWORD valuetaskLocal = pCode->NewLocal(valuetaskLocalDesc);
599+
pCode->EmitSTLOC(valuetaskLocal);
600+
pCode->EmitLDLOCA(valuetaskLocal);
601+
pCode->EmitCALL(getAwaiterToken, 1, 1);
602+
}
603+
else
604+
{
605+
pCode->EmitCALLVIRT(getAwaiterToken, 1, 1);
606+
}
607+
554608
pCode->EmitSTLOC(awaiterLocal);
555609
pCode->EmitLDLOCA(awaiterLocal);
556610
pCode->EmitCALL(getIsCompletedToken, 1, 1);

src/coreclr/vm/corelib.h

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -344,12 +344,14 @@ DEFINE_CLASS(THREAD_START_EXCEPTION,Threading, ThreadStartException
344344
DEFINE_METHOD(THREAD_START_EXCEPTION,EX_CTOR, .ctor, IM_Exception_RetVoid)
345345

346346
DEFINE_CLASS(VALUETASK_1, Tasks, ValueTask`1)
347+
DEFINE_METHOD(VALUETASK_1, GET_AWAITER, GetAwaiter, NoSig)
347348

348349
DEFINE_CLASS(VALUETASK, Tasks, ValueTask)
349350
DEFINE_METHOD(VALUETASK, FROM_EXCEPTION, FromException, SM_Exception_RetValueTask)
350351
DEFINE_METHOD(VALUETASK, FROM_EXCEPTION_1, FromException, GM_Exception_RetValueTaskOfT)
351352
DEFINE_METHOD(VALUETASK, FROM_RESULT_T, FromResult, GM_T_RetValueTaskOfT)
352353
DEFINE_METHOD(VALUETASK, GET_COMPLETED_TASK, get_CompletedTask, SM_RetValueTask)
354+
DEFINE_METHOD(VALUETASK, GET_AWAITER, GetAwaiter, NoSig)
353355

354356
DEFINE_CLASS(TASK_1, Tasks, Task`1)
355357
DEFINE_METHOD(TASK_1, GET_AWAITER, GetAwaiter, NoSig)
@@ -369,6 +371,14 @@ DEFINE_CLASS(TASK_AWAITER, CompilerServices, TaskAwaiter)
369371
DEFINE_METHOD(TASK_AWAITER, GET_ISCOMPLETED, get_IsCompleted, NoSig)
370372
DEFINE_METHOD(TASK_AWAITER, GET_RESULT, GetResult, NoSig)
371373

374+
DEFINE_CLASS(VALUETASK_AWAITER_1, CompilerServices, ValueTaskAwaiter`1)
375+
DEFINE_METHOD(VALUETASK_AWAITER_1, GET_ISCOMPLETED, get_IsCompleted, NoSig)
376+
DEFINE_METHOD(VALUETASK_AWAITER_1, GET_RESULT, GetResult, NoSig)
377+
378+
DEFINE_CLASS(VALUETASK_AWAITER, CompilerServices, ValueTaskAwaiter)
379+
DEFINE_METHOD(VALUETASK_AWAITER, GET_ISCOMPLETED, get_IsCompleted, NoSig)
380+
DEFINE_METHOD(VALUETASK_AWAITER, GET_RESULT, GetResult, NoSig)
381+
372382
DEFINE_CLASS(TYPE_HANDLE, System, RuntimeTypeHandle)
373383
DEFINE_CLASS(RT_TYPE_HANDLE, System, RuntimeTypeHandle)
374384
DEFINE_METHOD(RT_TYPE_HANDLE, PVOID_CTOR, .ctor, IM_RuntimeType_RetVoid)

src/coreclr/vm/method.hpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2099,6 +2099,7 @@ class MethodDesc
20992099
void EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig& thunkMsig, ILStubLinker* pSL);
21002100
void EmitAsyncMethodThunk(MethodDesc* pAsyncOtherVariant, MetaSig& msig, ILStubLinker* pSL);
21012101
SigPointer GetAsyncThunkResultTypeSig();
2102+
bool IsValueTaskAsyncThunk();
21022103
int GetTokenForGenericMethodCallWithAsyncReturnType(ILCodeStream* pCode, MethodDesc* md);
21032104
int GetTokenForGenericTypeMethodCallWithAsyncReturnType(ILCodeStream* pCode, MethodDesc* md);
21042105
int GetTokenForAwaitAwaiterInstantiatedOverTaskAwaiterType(ILCodeStream* pCode, TypeHandle taskAwaiterType);

0 commit comments

Comments
 (0)