JIT: Improve modelling of AsyncResumedDef store in value numbering

src/coreclr/jit/async.cpp

@@ -45,6 +45,26 @@
 #include "jitstd/algorithm.h"
 #include "async.h"
 
+//------------------------------------------------------------------------
+// IsKnownResumed:
+//   Check whether a "resumed" indicator value is known to be set.
+//
+// Parameters:
+//   resumed - Value of the indicator
+//
+// Returns:
+//   True if the value is known to be non-zero.
+//
+// Remarks:
+//   The indicator is only ever consumed as a boolean, and the suspension helpers that
+//   take it no-op when it is set, so they can be left out entirely in that case. VN based
+//   constant propagation is what turns the indicator into a constant here.
+//
+static bool IsKnownResumed(GenTree* resumed)
+{
+    return resumed->IsIntegralConst() && !resumed->IsIntegralConst(0);
+}
+
 //------------------------------------------------------------------------
 // ContinuationMember::CustomAwaiterOfLayout:
 //   Create a continuation member that stores a custom awaiter with the
@@ -3486,14 +3506,28 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock*
     // A frame's own indicator is not enough to see that: it is only set when the frame
     // logically returns, which has not happened yet at a suspension. So accumulate the
     // indicators of the frames walked so far and gate on that instead.
-    unsigned const   anyResumedLcl = m_compiler->lvaGrabTemp(false DEBUGARG("Async any inlined frame resumed"));
-    LclVarDsc* const anyResumedDsc = m_compiler->lvaGetDesc(anyResumedLcl);
-    anyResumedDsc->lvType          = TYP_INT;
-    anyResumedDsc->lvOnlyUsedOnSynchronousPath = true;
+    //
+    // If the optimizer proved one of the indicators to be set then the accumulation is
+    // known to be set from that point on, and all the helpers gated on it are no-ops. In
+    // that case emit none of it.
+    bool anyResumedIsSet = IsKnownResumed(frameResumed);
+
+    unsigned anyResumedLcl = BAD_VAR_NUM;
+    if (!anyResumedIsSet)
+    {
+        anyResumedLcl                  = m_compiler->lvaGrabTemp(false DEBUGARG("Async any inlined frame resumed"));
+        LclVarDsc* const anyResumedDsc = m_compiler->lvaGetDesc(anyResumedLcl);
+        anyResumedDsc->lvType          = TYP_INT;
+        anyResumedDsc->lvOnlyUsedOnSynchronousPath = true;
 
-    GenTree* const initAnyResumed =
-        m_compiler->gtNewStoreLclVarNode(anyResumedLcl, m_compiler->gtCloneExpr(frameResumed));
-    LIR::AsRange(tailBB).InsertAtEnd(LIR::SeqTree(m_compiler, initAnyResumed));
+        GenTree* const initAnyResumed =
+            m_compiler->gtNewStoreLclVarNode(anyResumedLcl, m_compiler->gtCloneExpr(frameResumed));
+        LIR::AsRange(tailBB).InsertAtEnd(LIR::SeqTree(m_compiler, initAnyResumed));
+    }
+    else
+    {
+        JITDUMP("    Frame is known to have resumed; skipping all frame transition handling\n");
+    }
 
     for (unsigned i = 0; i + 1 < numFrames; i++)
     {
@@ -3514,6 +3548,13 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock*
         assert(membersAreLive && "suspension tail needs a frame whose members were never registered");
         membersAreLive = membersAreLive && (layout.ContinuationMemberOffsets[flagsIndex] != UINT_MAX);
 
+        if (anyResumedIsSet)
+        {
+            // Both the capture and the restore below no-op on an already resumed frame.
+            JITDUMP("    Skipping no-op frame transition handling for inline frame depth %u\n", depth);
+            continue;
+        }
+
         if (membersAreLive)
         {
             // The caller's await decides whether this frame has to come back to the
@@ -3569,6 +3610,15 @@ BasicBlock* AsyncTransformation::CreateInlinedFrameSuspensionTail(BasicBlock*
 
         // Fold in the frame we are about to hand off to: from here on out its handling,
         // and that of every frame outside it, is likewise only needed the first time.
+        if (IsKnownResumed(outerResumed))
+        {
+            // The accumulation is set from here on, so this restore and everything
+            // after it no-ops. Leave the local alone; nothing reads it anymore.
+            JITDUMP("    Inline frame depth %u is known to have resumed; skipping the rest of the tail\n", depth);
+            anyResumedIsSet = true;
+            continue;
+        }
+
         GenTree* const merged =
             m_compiler->gtNewOperNode(GT_OR, TYP_INT, m_compiler->gtNewLclvNode(anyResumedLcl, TYP_INT),
                                       m_compiler->gtCloneExpr(outerResumed));
@@ -3652,6 +3702,36 @@ GenTree* AsyncTransformation::RestoreContexts(BasicBlock* block, GenTreeCall* ca
 
     JITDUMP("    Call [%06u] has async contexts; will restore on suspension\n", Compiler::dspTreeID(call));
 
+    // Take the values off the call. They are used from a different block, so anything that
+    // is not invariant or a local use is spilled to a temp first, which leaves its side
+    // effects behind in the block with the call.
+    auto takeValue = [=](CallArg* arg) {
+        GenTree* node = arg->GetNode();
+        if (!node->IsInvariant() && !node->OperIs(GT_LCL_VAR))
+        {
+            LIR::Use use(LIR::AsRange(block), &arg->NodeRef(), call);
+            use.ReplaceWithLclVar(m_compiler);
+            node = use.Def();
+        }
+
+        LIR::AsRange(block).Remove(node);
+        call->gtArgs.RemoveUnsafe(arg);
+        return node;
+    };
+
+    GenTree* const resumed     = takeValue(resumedArg);
+    GenTree* const execContext = takeValue(execContextArg);
+    GenTree* const syncContext = takeValue(syncContextArg);
+
+    if (IsKnownResumed(resumed))
+    {
+        // RestoreContextsOnSuspension no-ops on an already resumed frame, so do not emit
+        // it at all. The values it would have taken are simply left out; any side effects
+        // they had stayed behind in the block with the call.
+        JITDUMP("    Frame is known to have resumed; skipping no-op restore on suspension\n");
+        return resumed;
+    }
+
     // Insert call
     //   AsyncHelpers.RestoreContextsOnSuspension(resumed, execContext, syncContext);
 
@@ -3669,66 +3749,20 @@ GenTree* AsyncTransformation::RestoreContexts(BasicBlock* block, GenTreeCall* ca
 
     LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_compiler, restoreCall));
 
-    // Replace resumedPlaceholder with actual resumed arg
-    GenTree* resumed = resumedArg->GetNode();
-    if (!resumed->IsInvariant() && !resumed->OperIs(GT_LCL_VAR))
-    {
-        // We are moving resumed into a different BB so create a temp for it.
-        LIR::Use use(LIR::AsRange(block), &resumedArg->NodeRef(), call);
-        use.ReplaceWithLclVar(m_compiler);
-        resumed = use.Def();
-    }
-
-    LIR::Use use;
-    bool     gotUse = LIR::AsRange(suspendBB).TryGetUse(resumedPlaceholder, &use);
-    assert(gotUse);
-
-    LIR::AsRange(block).Remove(resumed);
-    LIR::AsRange(suspendBB).InsertBefore(resumedPlaceholder, resumed);
-    use.ReplaceWith(resumed);
-    LIR::AsRange(suspendBB).Remove(resumedPlaceholder);
-
-    call->gtArgs.RemoveUnsafe(resumedArg);
-
-    // Replace execContextPlaceholder with actual value
-    GenTree* execContext = execContextArg->GetNode();
-    if (!execContext->OperIs(GT_LCL_VAR))
-    {
-        // We are moving execContext into a different BB so create a temp for it.
-        LIR::Use use(LIR::AsRange(block), &execContextArg->NodeRef(), call);
-        use.ReplaceWithLclVar(m_compiler);
-        execContext = use.Def();
-    }
-
-    gotUse = LIR::AsRange(suspendBB).TryGetUse(execContextPlaceholder, &use);
-    assert(gotUse);
-
-    LIR::AsRange(block).Remove(execContext);
-    LIR::AsRange(suspendBB).InsertBefore(execContextPlaceholder, execContext);
-    use.ReplaceWith(execContext);
-    LIR::AsRange(suspendBB).Remove(execContextPlaceholder);
-
-    call->gtArgs.RemoveUnsafe(execContextArg);
-
-    // Replace syncContextPlaceholder with actual value
-    GenTree* syncContext = syncContextArg->GetNode();
-    if (!syncContext->OperIs(GT_LCL_VAR))
-    {
-        // We are moving syncContext into a different BB so create a temp for it.
-        LIR::Use use(LIR::AsRange(block), &syncContextArg->NodeRef(), call);
-        use.ReplaceWithLclVar(m_compiler);
-        syncContext = use.Def();
-    }
-
-    gotUse = LIR::AsRange(suspendBB).TryGetUse(syncContextPlaceholder, &use);
-    assert(gotUse);
+    // Replace the placeholders with the actual values.
+    auto replacePlaceholder = [=](GenTree* placeholder, GenTree* value) {
+        LIR::Use use;
+        bool     gotUse = LIR::AsRange(suspendBB).TryGetUse(placeholder, &use);
+        assert(gotUse);
 
-    LIR::AsRange(block).Remove(syncContext);
-    LIR::AsRange(suspendBB).InsertBefore(syncContextPlaceholder, syncContext);
-    use.ReplaceWith(syncContext);
-    LIR::AsRange(suspendBB).Remove(syncContextPlaceholder);
+        LIR::AsRange(suspendBB).InsertBefore(placeholder, value);
+        use.ReplaceWith(value);
+        LIR::AsRange(suspendBB).Remove(placeholder);
+    };
 
-    call->gtArgs.RemoveUnsafe(syncContextArg);
+    replacePlaceholder(resumedPlaceholder, resumed);
+    replacePlaceholder(execContextPlaceholder, execContext);
+    replacePlaceholder(syncContextPlaceholder, syncContext);
 
     JITDUMP("    Created RestoreContexts call on suspension:\n");
     DISPTREERANGE(LIR::AsRange(suspendBB), restoreCall);

src/coreclr/jit/valuenum.cpp

@@ -14968,11 +14968,67 @@ void Compiler::fgValueNumberCall(GenTreeCall* call)
         }
     }
 
+    // Async calls define the "resumed" indicator: it is set to 1 when the call
+    // suspends and we are later resumed in this frame, and is left alone
+    // otherwise. Thus we know the value after the call is 1 if the call always
+    // suspends, or if the indicator was already known to be 1 before it.
+    GenTreeLclVarCommon* asyncResumedLclAddr = gtCallGetDefinedAsyncResumedLclAddr(call);
+    ValueNumPair         asyncResumedVNP;
+    if (asyncResumedLclAddr != nullptr)
+    {
+        if (call->GetAsyncInfo().AlwaysSuspends)
+        {
+            JITDUMP("Call [%06u] always suspends; establishing always resumed state for V%02u\n", dspTreeID(call),
+                    asyncResumedLclAddr->GetLclNum());
+            asyncResumedVNP.SetBoth(vnStore->VNOneForType(TYP_I_IMPL));
+        }
+        else
+        {
+            CallArg* resumedUseArg = call->gtArgs.FindWellKnownArg(WellKnownArg::AsyncResumedUse);
+            if (resumedUseArg != nullptr)
+            {
+                // The arg is usually a TYP_INT use of the TYP_I_IMPL indicator var.
+                assert(genActualTypeIsInt(resumedUseArg->GetNode()));
+
+                // The indicator is monotonic, so it stays 1 for whichever kinds
+                // it was already known to be 1 for.
+                ValueNumKind vnKinds[2] = {VNK_Liberal, VNK_Conservative};
+                for (ValueNumKind vnKind : vnKinds)
+                {
+                    ValueNum resumedUseVN = vnStore->VNNormalValue(resumedUseArg->GetNode()->GetVN(vnKind));
+                    if (vnStore->IsVNConstant(resumedUseVN) && (vnStore->CoercedConstantValue<int>(resumedUseVN) == 1))
+                    {
+                        JITDUMP("Call [%06u] propagates always resumed state for V%02u (%s VN)\n", dspTreeID(call),
+                                asyncResumedLclAddr->GetLclNum(), vnKind == VNK_Liberal ? "liberal" : "conservative");
+                        asyncResumedVNP.Set(vnKind, vnStore->VNOneForType(TYP_I_IMPL));
+                    }
+                }
+            }
+        }
+    }
+
     // If the call generates any definitions, for example because it uses "return buffer", then VN the local
     // as well.
     auto visitDef = [=](const LocalDef& def) {
         ValueNumPair storeValue;
-        storeValue.SetBoth(vnStore->VNForExpr(compCurBB, lvaGetDesc(def.Def->AsLclVarCommon())->TypeGet()));
+        if (def.Def == asyncResumedLclAddr)
+        {
+            storeValue = asyncResumedVNP;
+        }
+
+        if (!storeValue.BothDefined())
+        {
+            // Use the same new unique VN for any kind we do not know the value of.
+            ValueNum newVN = vnStore->VNForExpr(compCurBB, lvaGetDesc(def.Def->AsLclVarCommon())->TypeGet());
+            if (storeValue.GetLiberal() == ValueNumStore::NoVN)
+            {
+                storeValue.SetLiberal(newVN);
+            }
+            if (storeValue.GetConservative() == ValueNumStore::NoVN)
+            {
+                storeValue.SetConservative(newVN);
+            }
+        }
 
         fgValueNumberLocalStore(call, def.Def, def.Offset, def.Size, storeValue);
         return GenTree::VisitResult::Continue;

src/tests/async/resumed-known-true/resumed-known-true.cs

@@ -0,0 +1,85 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
+using Xunit;
+
+// Shapes where the JIT can prove that the "resumed" indicator is 1 at a
+// suspension point: an await that always suspends dominates a later await, with
+// no merge in between. The helpers emitted on the later await's suspension path
+// (RestoreContextsOnSuspension, the inlined frame transition captures) no-op
+// when the indicator is set, so they could be folded away entirely.
+public class Async2ResumedKnownTrue
+{
+    [Fact]
+    public static void TestTwoYields()
+    {
+        TwoYields().GetAwaiter().GetResult();
+    }
+
+    [Fact]
+    public static void TestYieldThenAwaitTask()
+    {
+        YieldThenAwaitTask().GetAwaiter().GetResult();
+    }
+
+    [Fact]
+    public static void TestYieldThenAwaitInlinedFrame()
+    {
+        YieldThenAwaitInlinedFrame().GetAwaiter().GetResult();
+    }
+
+    [Fact]
+    public static void TestYieldThenAwaitResumedInlinedFrame()
+    {
+        YieldThenAwaitResumedInlinedFrame().GetAwaiter().GetResult();
+    }
+
+    // The second Yield's suspension path sees the indicator defined by the
+    // first, which always suspends.
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static async Task TwoYields()
+    {
+        await Task.Yield();
+        await Task.Yield();
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static async Task YieldThenAwaitTask()
+    {
+        await Task.Yield();
+        await Task.Delay(1);
+    }
+
+    // Same, but the later await sits in an inlined async frame, so its
+    // suspension tail runs the enclosing frames' transition captures too.
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static async Task YieldThenAwaitInlinedFrame()
+    {
+        await Task.Yield();
+        await Inner();
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private static async Task Inner()
+    {
+        await Task.Delay(1);
+    }
+
+    // Here the inlined frame's own indicator is set as well by the time it suspends, so
+    // none of the frame transition handling is needed.
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static async Task YieldThenAwaitResumedInlinedFrame()
+    {
+        await Task.Yield();
+        await InnerYieldThenAwait();
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private static async Task InnerYieldThenAwait()
+    {
+        await Task.Yield();
+        await Task.Delay(1);
+    }
+}

src/tests/async/resumed-known-true/resumed-known-true.csproj

@@ -0,0 +1,5 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <ItemGroup>
+   <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+</Project>