JIT: Enable inlining of runtime async methods without awaits

src/coreclr/jit/async.cpp

@@ -183,10 +183,18 @@ PhaseStatus Compiler::SaveAsyncContexts()
     }
 
     // Insert RestoreContexts call in fault (exceptional case)
-    // First argument: started = (continuation == null)
-    GenTree* continuation = gtNewLclvNode(lvaAsyncContinuationArg, TYP_REF);
-    GenTree* null         = gtNewNull();
-    GenTree* resumed      = gtNewOperNode(GT_NE, TYP_INT, continuation, null);
+    // 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);
+    }
 
     GenTreeCall* restoreCall = gtNewCallNode(CT_USER_FUNC, asyncInfo->restoreContextsMethHnd, TYP_VOID);
     restoreCall->gtArgs.PushFront(this,
@@ -205,7 +213,10 @@ PhaseStatus Compiler::SaveAsyncContexts()
 
     for (BasicBlock* block : Blocks())
     {
-        AddContextArgsToAsyncCalls(block);
+        if (!compIsForInlining())
+        {
+            AddContextArgsToAsyncCalls(block);
+        }
 
         if (!block->KindIs(BBJ_RETURN) || (block == newReturnBB))
         {
@@ -220,23 +231,28 @@ PhaseStatus Compiler::SaveAsyncContexts()
             newReturnBB->inheritWeightPercentage(block, 0);
         }
 
-        // Store return value to common local
-        Statement* retStmt = block->lastStmt();
-        assert((retStmt != nullptr) && retStmt->GetRootNode()->OperIs(GT_RETURN));
-
-        if (mergedReturnLcl != BAD_VAR_NUM)
+        // When inlining we do merging during import, so we do not need to do
+        // any storing there.
+        if (!compIsForInlining())
         {
-            GenTree*   retVal      = retStmt->GetRootNode()->AsOp()->GetReturnValue();
-            Statement* insertAfter = retStmt;
-            GenTree*   storeRetVal =
-                gtNewTempStore(mergedReturnLcl, retVal, CHECK_SPILL_NONE, &insertAfter, retStmt->GetDebugInfo(), block);
-            Statement* storeStmt = fgNewStmtFromTree(storeRetVal);
-            fgInsertStmtAtEnd(block, storeStmt);
-            JITDUMP("Inserted store to common return local\n");
-            DISPSTMT(storeStmt);
-        }
+            // Store return value to common local
+            Statement* retStmt = block->lastStmt();
+            assert((retStmt != nullptr) && retStmt->GetRootNode()->OperIs(GT_RETURN));
 
-        retStmt->GetRootNode()->gtBashToNOP();
+            if (mergedReturnLcl != BAD_VAR_NUM)
+            {
+                GenTree*   retVal      = retStmt->GetRootNode()->AsOp()->GetReturnValue();
+                Statement* insertAfter = retStmt;
+                GenTree*   storeRetVal = gtNewTempStore(mergedReturnLcl, retVal, CHECK_SPILL_NONE, &insertAfter,
+                                                        retStmt->GetDebugInfo(), block);
+                Statement* storeStmt   = fgNewStmtFromTree(storeRetVal);
+                fgInsertStmtAtEnd(block, storeStmt);
+                JITDUMP("Inserted store to common return local\n");
+                DISPSTMT(storeStmt);
+            }
+
+            retStmt->GetRootNode()->gtBashToNOP();
+        }
 
         // Jump to new shared restore + return block
         block->SetKindAndTargetEdge(BBJ_ALWAYS, fgAddRefPred(newReturnBB, block));
@@ -334,9 +350,17 @@ BasicBlock* Compiler::CreateReturnBB(unsigned* mergedReturnLcl)
     // Insert "restore" call
     CORINFO_ASYNC_INFO* asyncInfo = eeGetAsyncInfo();
 
-    GenTree* continuation = gtNewLclvNode(lvaAsyncContinuationArg, TYP_REF);
-    GenTree* null         = gtNewNull();
-    GenTree* resumed      = gtNewOperNode(GT_NE, TYP_INT, 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);
+    }
 
     GenTreeCall* restoreCall = gtNewCallNode(CT_USER_FUNC, asyncInfo->restoreContextsMethHnd, TYP_VOID);
     restoreCall->gtArgs.PushFront(this,
@@ -355,42 +379,46 @@ BasicBlock* Compiler::CreateReturnBB(unsigned* mergedReturnLcl)
     JITDUMP("Inserted restore statement in return block\n");
     DISPSTMT(restoreStmt);
 
-    *mergedReturnLcl = BAD_VAR_NUM;
-
-    GenTree* ret;
-    if (compMethodHasRetVal())
+    if (!compIsForInlining())
     {
-        *mergedReturnLcl = lvaGrabTemp(false DEBUGARG("Async merged return local"));
+        *mergedReturnLcl = BAD_VAR_NUM;
 
-        var_types retLclType = compMethodReturnsRetBufAddr() ? TYP_BYREF : genActualType(info.compRetType);
-
-        if (varTypeIsStruct(retLclType))
+        GenTree* ret;
+        if (compMethodHasRetVal())
         {
-            lvaSetStruct(*mergedReturnLcl, info.compMethodInfo->args.retTypeClass, false);
+            *mergedReturnLcl = lvaGrabTemp(false DEBUGARG("Async merged return local"));
+
+            var_types retLclType = compMethodReturnsRetBufAddr() ? TYP_BYREF : genActualType(info.compRetType);
 
-            if (compMethodReturnsMultiRegRetType())
+            if (varTypeIsStruct(retLclType))
+            {
+                lvaSetStruct(*mergedReturnLcl, info.compMethodInfo->args.retTypeClass, false);
+
+                if (compMethodReturnsMultiRegRetType())
+                {
+                    lvaGetDesc(*mergedReturnLcl)->lvIsMultiRegRet = true;
+                }
+            }
+            else
             {
-                lvaGetDesc(*mergedReturnLcl)->lvIsMultiRegRet = true;
+                lvaGetDesc(*mergedReturnLcl)->lvType = retLclType;
             }
+
+            GenTree* retTemp = gtNewLclVarNode(*mergedReturnLcl);
+            ret              = gtNewOperNode(GT_RETURN, retTemp->TypeGet(), retTemp);
         }
         else
         {
-            lvaGetDesc(*mergedReturnLcl)->lvType = retLclType;
+            ret = new (this, GT_RETURN) GenTreeOp(GT_RETURN, TYP_VOID);
         }
 
-        GenTree* retTemp = gtNewLclVarNode(*mergedReturnLcl);
-        ret              = gtNewOperNode(GT_RETURN, retTemp->TypeGet(), retTemp);
-    }
-    else
-    {
-        ret = new (this, GT_RETURN) GenTreeOp(GT_RETURN, TYP_VOID);
-    }
+        Statement* retStmt = fgNewStmtFromTree(ret);
 
-    Statement* retStmt = fgNewStmtFromTree(ret);
+        fgInsertStmtAtEnd(newReturnBB, retStmt);
+        JITDUMP("Inserted return statement in return block\n");
+        DISPSTMT(retStmt);
+    }
 
-    fgInsertStmtAtEnd(newReturnBB, retStmt);
-    JITDUMP("Inserted return statement in return block\n");
-    DISPSTMT(retStmt);
     return newReturnBB;
 }
 

src/coreclr/jit/fgbasic.cpp

@@ -3455,7 +3455,7 @@ void Compiler::fgFindBasicBlocks()
 
     // Are there any exception handlers?
     //
-    if (info.compXcptnsCount > 0)
+    if (info.compXcptnsCount > 0 || ((info.compMethodInfo->options & CORINFO_ASYNC_SAVE_CONTEXTS) != 0))
     {
         assert(!compIsForInlining() || opts.compInlineMethodsWithEH);
 
@@ -3464,7 +3464,14 @@ void Compiler::fgFindBasicBlocks()
             // Verify we can expand the EH table as needed to incorporate the callee's EH clauses.
             // Failing here should be extremely rare.
             //
-            EHblkDsc* const dsc = fgTryAddEHTableEntries(0, info.compXcptnsCount, /* deferAdding */ true);
+            unsigned numEHEntries = info.compXcptnsCount;
+            // We will introduce another EH clause before inlining finishes to restore async contexts
+            if ((info.compMethodInfo->options & CORINFO_ASYNC_SAVE_CONTEXTS) != 0)
+            {
+                numEHEntries++;
+            }
+
+            EHblkDsc* const dsc = fgTryAddEHTableEntries(0, numEHEntries, /* deferAdding */ true);
             if (dsc == nullptr)
             {
                 compInlineResult->NoteFatal(InlineObservation::CALLSITE_EH_TABLE_FULL);

src/coreclr/jit/importercalls.cpp

@@ -404,6 +404,11 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
                 if (sig->isAsyncCall())
                 {
                     impSetupAsyncCall(call->AsCall(), opcode, prefixFlags, di);
+
+                    if (compDonotInline())
+                    {
+                        return TYP_UNDEF;
+                    }
                 }
 
                 impPopCallArgs(sig, call->AsCall());
@@ -717,6 +722,11 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
     {
         impSetupAsyncCall(call->AsCall(), opcode, prefixFlags, di);
 
+        if (compDonotInline())
+        {
+            return TYP_UNDEF;
+        }
+
         if (lvaNextCallAsyncContinuation != BAD_VAR_NUM)
         {
             asyncContinuation            = gtNewLclVarNode(lvaNextCallAsyncContinuation);
@@ -6868,6 +6878,12 @@ void Compiler::impCheckForPInvokeCall(
 //
 void Compiler::impSetupAsyncCall(GenTreeCall* call, OPCODE opcode, unsigned prefixFlags, const DebugInfo& callDI)
 {
+    if (compIsForInlining())
+    {
+        compInlineResult->NoteFatal(InlineObservation::CALLEE_AWAIT);
+        return;
+    }
+
     AsyncCallInfo asyncInfo;
 
     unsigned newSourceTypes = ICorDebugInfo::ASYNC;
@@ -8123,14 +8139,6 @@ void Compiler::impMarkInlineCandidateHelper(GenTreeCall*           call,
         return;
     }
 
-    if (call->IsAsync() && (call->GetAsyncInfo().ContinuationContextHandling != ContinuationContextHandling::None))
-    {
-        // Cannot currently handle moving to captured context/thread pool when logically returning from inlinee.
-        //
-        inlineResult->NoteFatal(InlineObservation::CALLSITE_CONTINUATION_HANDLING);
-        return;
-    }
-
     // Ignore indirect calls, unless they are indirect virtual stub calls with profile info.
     //
     if (call->gtCallType == CT_INDIRECT)

src/coreclr/jit/inline.def

@@ -56,6 +56,7 @@ INLINE_OBSERVATION(STACK_CRAWL_MARK,          bool,   "uses stack crawl mark",
 INLINE_OBSERVATION(STFLD_NEEDS_HELPER,        bool,   "stfld needs helper",                   FATAL,       CALLEE)
 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)
 
 // ------ Callee Performance -------
 
@@ -161,7 +162,6 @@ INLINE_OBSERVATION(RETURN_TYPE_MISMATCH,      bool,   "return type mismatch",
 INLINE_OBSERVATION(STFLD_NEEDS_HELPER,        bool,   "stfld needs helper",                   FATAL,       CALLSITE)
 INLINE_OBSERVATION(TOO_MANY_LOCALS,           bool,   "too many locals",                      FATAL,       CALLSITE)
 INLINE_OBSERVATION(PINVOKE_EH,                bool,   "PInvoke call site with EH",            FATAL,       CALLSITE)
-INLINE_OBSERVATION(CONTINUATION_HANDLING,     bool,   "Callsite needs continuation handling", FATAL,       CALLSITE)
 
 // ------ Call Site Performance -------