Optimize runtime async OSR resumption performance

[jakobbotsch]

Resuming inside OSR methods is a slow process. It requires first calling the tier0 version which then calls into the runtime to perform an (unwinder-based) OSR transition.

We should not need a call into the runtime to transition to the OSR method from the tier0 method. Instead we can do a non-local jump from the tier 0 version when we realize we have an OSR method continuation.

[jakobbotsch]

The cost of the expensive OSR transition is about an order of magnitude:

using System;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Threading.Tasks;

namespace OSRPerf;

public class Program
{
    static void Main()
    {
        NullAwaiter na = new NullAwaiter();

        Task t = Foo(10_000_000, na);
        while (!t.IsCompleted)
        {
            na.Continue();
        }
    }

    static int s_value;
    static async Task Foo(int n, NullAwaiter na)
    {
        for (int i = 0; i < n; i++)
        {
            s_value += i;
        }

        Stopwatch timer = Stopwatch.StartNew();
        for (int i = 0; i < 10_000_000; i++)
        {
            await na;
        }
        Console.WriteLine("Took {0:F1} ms", timer.Elapsed.TotalMilliseconds);
    }

    private class NullAwaiter : ICriticalNotifyCompletion
    {
        public Action Continue;

        public NullAwaiter GetAwaiter() => this;

        public bool IsCompleted => false;

        public void GetResult()
        {
        }

        public void UnsafeOnCompleted(Action continuation)
        {
            Continue = continuation;
        }

        public void OnCompleted(Action continuation)
        {
            throw new NotImplementedException();
        }
    }
}

DOTNET_TieredCompilation=0: Took 563.0 ms
DOTNET_TieredCompilation=1: Took 6205.5 ms

(Edit: updated numbers with a release coreclr build + using #120411)

[AndyAyersMS]

FWIW I wanted to do something similar for normal OSR transitions once the OSR method was available, and this also would make things like partial compilation more viable (where transitions might also be frequent).

Say the patchpoint helper returns with the address to jump to, or invoke the patchpoint helper via an indirection cell that gets updated with the OSR entry point.

Currently the OSR entry point expects to be called (no args needed, return address ignored), so you would have to emulate that on x64 (push something on the stack then jump, or actually just call it).

[jakobbotsch]

Currently the OSR entry point expects to be called (no args needed, return address ignored), so you would have to emulate that on x64 (push something on the stack then jump, or actually just call it).

Seems like push + jmp would be preferred since otherwise the CPU call-return stack will become unbalanced. Although that worries me for unwinding purposes.
What would it take to get rid of that? Just some updating to the OSR prolog generation to now learn that the stack is actually aligned on entry?

[AndyAyersMS]

What would it take to get rid of that? Just some updating to the OSR prolog generation to now learn that the stack is actually aligned on entry?

Yes, something like that. The OSR prolog could initially (on x64) just push RAX or something harmless, and we could adjust the transition code in the runtime.