JIT: Disallow forward substitution of async calls (#115936)

This can be forwarded into an overlapping byref temp, which is invalid
IR.

We could look for that condition specifically, but that's more complex. Do the
simple thing for now and skip forward substituting trees containing async calls.
In the future when we have async contexts in SPMI collections we can try the
more precise check and see if it makes a difference.

The opposite might also be possible: forward substituting a TYP_BYREF
into a tree with an async call. However, I am not 100% sure that can
create overlapping lifetimes, so let's wait and see if Fuzzlyn comes up
with an example.

Author: jakobbotsch

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)