JIT: Disallow forward substitution of async calls

src/coreclr/jit/compiler.h

@@ -6160,7 +6160,10 @@ class Compiler
     PhaseStatus fgHeadTailMerge(bool early);
     bool fgHeadMerge(BasicBlock* block, bool early);
     bool fgTryOneHeadMerge(BasicBlock* block, bool early);
+    template<typename Predicate>
+    bool gtTreeContainsCall(GenTree* tree, Predicate pred);
     bool gtTreeContainsTailCall(GenTree* tree);
+    bool gtTreeContainsAsyncCall(GenTree* tree);
     bool fgCanMoveFirstStatementIntoPred(bool early, Statement* firstStmt, BasicBlock* pred);
 
     enum FG_RELOCATE_TYPE

src/coreclr/jit/fgopt.cpp

@@ -5605,29 +5605,33 @@ bool Compiler::fgHeadMerge(BasicBlock* block, bool early)
 }
 
 //------------------------------------------------------------------------
-// gtTreeContainsTailCall: Check if a tree contains any tail call or tail call
-// candidate.
+// gtTreeContainsCall:
+//   Check if a tree contains a call node matching the given predicate.
 //
 // Parameters:
 //   tree - The tree
+//   pred - Predicate that the call must match
 //
-// Remarks:
-//   While tail calls are generally expected to be top level nodes we do allow
-//   some other shapes of calls to be tail calls, including some cascading
-//   trivial assignments and casts. This function does a tree walk to check if
-//   any sub tree is a tail call.
+// Returns:
+//   True if a call node matching the predicate was found, false otherwise.
 //
-bool Compiler::gtTreeContainsTailCall(GenTree* tree)
+template <typename Predicate>
+bool Compiler::gtTreeContainsCall(GenTree* tree, Predicate pred)
 {
-    struct HasTailCallCandidateVisitor : GenTreeVisitor<HasTailCallCandidateVisitor>
+    struct HasCallVisitor : GenTreeVisitor<HasCallVisitor>
     {
+    private:
+        Predicate& m_pred;
+
+    public:
         enum
         {
             DoPreOrder = true
         };
 
-        HasTailCallCandidateVisitor(Compiler* comp)
-            : GenTreeVisitor(comp)
+        HasCallVisitor(Compiler* comp, Predicate& pred)
+            : GenTreeVisitor<HasCallVisitor>(comp)
+            , m_pred(pred)
         {
         }
 
@@ -5639,7 +5643,7 @@ bool Compiler::gtTreeContainsTailCall(GenTree* tree)
                 return WALK_SKIP_SUBTREES;
             }
 
-            if (node->IsCall() && (node->AsCall()->CanTailCall() || node->AsCall()->IsTailCall()))
+            if (node->IsCall() && m_pred(node->AsCall()))
             {
                 return WALK_ABORT;
             }
@@ -5648,8 +5652,54 @@ bool Compiler::gtTreeContainsTailCall(GenTree* tree)
         }
     };
 
-    HasTailCallCandidateVisitor visitor(this);
-    return visitor.WalkTree(&tree, nullptr) == WALK_ABORT;
+    if ((tree->gtFlags & GTF_CALL) == 0)
+    {
+        return false;
+    }
+
+    HasCallVisitor hasCall(this, pred);
+    return hasCall.WalkTree(&tree, nullptr) == WALK_ABORT;
+}
+
+//------------------------------------------------------------------------
+// gtTreeContainsTailCall: Check if a tree contains any tail call or tail call
+// candidate.
+//
+// Parameters:
+//   tree - The tree
+//
+// Remarks:
+//   While tail calls are generally expected to be top level nodes we do allow
+//   some other shapes of calls to be tail calls, including some cascading
+//   trivial assignments and casts. This function does a tree walk to check if
+//   any sub tree is a tail call.
+//
+bool Compiler::gtTreeContainsTailCall(GenTree* tree)
+{
+    return gtTreeContainsCall(tree, [](GenTreeCall* call) {
+        return call->CanTailCall() || call->IsTailCall();
+    });
+}
+
+//------------------------------------------------------------------------
+// gtTreeContainsAsyncCall: Check if a tree contains any async call.
+//
+// Parameters:
+//   tree - The tree to check
+//
+// Returns:
+//   True if any node in the tree is an async call, false otherwise.
+//
+bool Compiler::gtTreeContainsAsyncCall(GenTree* tree)
+{
+    if (!compIsAsync())
+    {
+        return false;
+    }
+
+    return gtTreeContainsCall(tree, [](GenTreeCall* call) {
+        return call->IsAsync();
+    });
 }
 
 //------------------------------------------------------------------------

src/coreclr/jit/forwardsub.cpp

@@ -501,19 +501,28 @@ bool Compiler::fgForwardSubStatement(Statement* stmt)
     // Can't substitute GT_CATCH_ARG.
     // Can't substitute GT_LCLHEAP.
     //
-    // Don't substitute a no return call (trips up morph in some cases).
     if (fwdSubNode->OperIs(GT_CATCH_ARG, GT_LCLHEAP, GT_ASYNC_CONTINUATION))
     {
         JITDUMP(" tree to sub is catch arg, or lcl heap\n");
         return false;
     }
 
+    // Don't substitute a no return call (trips up morph in some cases).
+    //
     if (fwdSubNode->IsCall() && fwdSubNode->AsCall()->IsNoReturn())
     {
         JITDUMP(" tree to sub is a 'no return' call\n");
         return false;
     }
 
+    // Do not substitute async calls; if the target node has a temp BYREF node,
+    // that creates illegal IR.
+    if (gtTreeContainsAsyncCall(fwdSubNode))
+    {
+        JITDUMP(" tree has an async call\n");
+        return false;
+    }
+
     // Bail if sub node has embedded stores.
     //
     if ((fwdSubNode->gtFlags & GTF_ASG) != 0)