IL emitted for non-capturing lambda declarations can be improved

[neon-sunset]

Given .NET compiler static initialization improvements throughout the years, it appears that Roslyn does not take advantage of them, emitting a JIT-unfriendly pattern for lambda declarations. In theory, this could be improved, yielding performance and codegen size wins.

Consider the following method:

public class Example
{
    public IEnumerable<int> GetRange()
    {
        return Enumerable
            .Range(0, 100)
            .Select(i => i * 2);
    }
}

Currently, it is being lowered to

public class Example
{
    [Serializable]
    [CompilerGenerated]
    private sealed class <>c
    {
        public static readonly <>c <>9 = new <>c();

        public static Func<int, int> <>9__0_0;

        internal int <GetRange>b__0_0(int i)
        {
            return i * 2;
        }
    }

    [System.Runtime.CompilerServices.NullableContext(1)]
    public IEnumerable<int> GetRange()
    {
        return Enumerable.Select(
            Enumerable.Range(0, 100),
            <>c.<>9__0_0 ?? (<>c.<>9__0_0 = new Func<int, int>(<>c.<>9.<GetRange>b__0_0)));
    }
}

Unfortunately, this completely defeats static init optimizations done when GetRange() reaches Tier 1 compilation.
Instead, the Roslyn could possibly emit

public class Example
{
    [CompilerGenerated]
    private static class <>c_GetRange
    {
        public static readonly Func<int, int> <>9__0_0 = new Func<int, int>(<GetRange>b__0_0);

        internal static int <GetRange>b__0_0(int i)
        {
            return i * 2;
        }
    }

    [System.Runtime.CompilerServices.NullableContext(1)]
    public IEnumerable<int> GetRange()
    {
        return Enumerable.Select(Enumerable.Range(0, 100), <>c_GetRange.<>9__0_0);
    }
}

This would allow to completely optimize away the init check in Tier 1 and lower down the cost of referencing/declaring non-capturing lambdas to just a couple of instructions. Additionally, with this form, Native AOT could theoretically start recognizing such pre-init patterns, optimizing them away completely during compilation (when interpreting cctors).

Additionally, it avoids initializing all lambdas in the class by splitting holders into separate static classes per method.

I don't know whether this is specification-compliant, but it would be really nice if it was.

[CyrusNajmabadi]

@stephentoub for guidance.

[neon-sunset]

Example of the codegen for the improved pattern:

C#

static IEnumerable<int> GetRange()
{
    return Enumerable.Range(0, 100).Select(i => i * 2);
}

static IEnumerable<int> GetRangeOptimized()
{
    return Enumerable.Range(0, 100).Select(Holder.Func);
}

static class Holder
{
    public static readonly Func<int, int> Func = new(__GetRange);

    private static int __GetRange(int i)
    {
        return i * 2;
    }
}

Tier 1: GetRange()

; Assembly listing for method Program:<<Main>$>g__GetRange|0_0():System.Collections.Generic.IEnumerable`1[int] (Tier1)
; Emitting BLENDED_CODE for generic ARM64 - MacOS
; Tier1 code
; optimized code
; optimized using Dynamic PGO
; fp based frame
; partially interruptible
; with Dynamic PGO: edge weights are valid, and fgCalledCount is 23378.09
; 1 inlinees with PGO data; 3 single block inlinees; 0 inlinees without PGO data

G_M000_IG01:                ;; offset=0x0000
            stp     fp, lr, [sp, #-0x20]!
            stp     x19, x20, [sp, #0x10]
            mov     fp, sp
 
G_M000_IG02:                ;; offset=0x000C
            movz    x0, #0x14E8
            movk    x0, #0x4AA LSL #16
            movk    x0, #1 LSL #32
            bl      CORINFO_HELP_NEWSFAST
            mov     x19, x0
            bl      System.Environment:get_CurrentManagedThreadId():int
            str     w0, [x19, #0x08]
            str     wzr, [x19, #0x14]
            mov     w0, #100
            str     w0, [x19, #0x18]
            movz    x0, #0x5D50
            movk    x0, #7 LSL #32
            ldr     x1, [x0]
            cbz     x1, G_M000_IG05 ;; <-- manual init check 
 
G_M000_IG03:                ;; offset=0x0044
            mov     x0, x19
            movz    x2, #0x6370
            movk    x2, #0x4AD LSL #16
            movk    x2, #1 LSL #32
            ldr     x2, [x2]
            blr     x2
 
G_M000_IG04:                ;; offset=0x005C
            ldp     x19, x20, [sp, #0x10]
            ldp     fp, lr, [sp], #0x20
            ret     lr
 
 ;; --- inline Func<int, int> initialization logic that cannot be optimized way ---
G_M000_IG05:                ;; offset=0x0068
            movz    x0, #0xFCC8
            movk    x0, #0x4A9 LSL #16
            movk    x0, #1 LSL #32
            bl      CORINFO_HELP_NEWSFAST
            mov     x20, x0
            movz    x1, #0x5D48
            movk    x1, #7 LSL #32
            ldr     x1, [x1]
            mov     x0, x20
            movz    x2, #0x2298
            movk    x2, #0x4AD LSL #16
            movk    x2, #1 LSL #32
            movz    x3, #0x4210
            movk    x3, #0x49F LSL #16
            movk    x3, #1 LSL #32
            ldr     x3, [x3]
            blr     x3
            movz    x14, #0x5D50
            movk    x14, #7 LSL #32
            mov     x15, x20
            bl      CORINFO_HELP_ASSIGN_REF
            mov     x1, x20
            b       G_M000_IG03
 
; Total bytes of code 196

Tier 1: GetRangeOptimized()

; Assembly listing for method Program:<<Main>$>g__GetRangeOptimized|0_1():System.Collections.Generic.IEnumerable`1[int] (Tier1)
; Emitting BLENDED_CODE for generic ARM64 - MacOS
; Tier1 code
; optimized code
; fp based frame
; partially interruptible
; No PGO data
; 1 inlinees with PGO data; 3 single block inlinees; 0 inlinees without PGO data

G_M000_IG01:                ;; offset=0x0000
            stp     fp, lr, [sp, #-0x20]!
            str     x19, [sp, #0x18]
            mov     fp, sp
 
G_M000_IG02:                ;; offset=0x000C
            movz    x0, #0x14E8
            movk    x0, #0x4AA LSL #16
            movk    x0, #1 LSL #32
            bl      CORINFO_HELP_NEWSFAST
            mov     x19, x0
            bl      System.Environment:get_CurrentManagedThreadId():int
            str     w0, [x19, #0x08]
            str     wzr, [x19, #0x14]
            mov     w1, #100
            str     w1, [x19, #0x18]
            movz    x1, #0x5C80 ;; <-- construct the static field address
            movk    x1, #7 LSL #32
            ldr     x1, [x1] ;; <-- load Func<int, int> reference, that's it!
            mov     x0, x19
            movz    x2, #0x6370
            movk    x2, #0x4AD LSL #16
            movk    x2, #1 LSL #32
            ldr     x2, [x2]
            blr     x2
 
G_M000_IG03:                ;; offset=0x0058
            ldr     x19, [sp, #0x18]
            ldp     fp, lr, [sp], #0x20
            ret     lr
 
; Total bytes of code 100

[huoyaoyuan]

What's the codegen on older frameworks? I don't expect it to be significantly worse though.

[neon-sunset]

I'm not entirely sure about pre-.NET 6, but otherwise static init check should be optimized away on older targets too.

[stephentoub]

cc: @VSadov

[MichalPetryka]

This was discussed before in dotnet/csharplang#6746 (reply in thread) and is unlikely to be changed in this way, dotnet/runtime#85014 was created as an alternative fix for it though.

[neon-sunset]

That is true. Part of the reason for submitting this issue is reconsidering it. Alternative options require more complex changes and seem that they are less likely to land.

There was another discussion somewhere that highlighted the reason for current design is reducing the startup cost. I am not sure whether this holds true today with static init improvements in the runtime.

[neon-sunset]

@MichalPetryka I have amended the description so that each method has its own holder static class for lambdas to avoid the issue with forcing initialization of all lambdas that may be unrelated if any of the methods of a particular class is called as suggested by @EgorBo .

[netcorefan1]

I may be wrong, but it seems to be that we have the same problem. I tried to solve it through the use of interceptors, but without success because the compiler still generates the closure (see #71408).
Basically you can avoid closures if you wrap lambda expressions in some helper that injects the argument. The passed data will be stores in custom classes. From there you can cache delegates and do the further optimizations that you need. Problem is that each method will need to be prefixed like LambdaHelper(myvar, myvarinjected => ...) . With an interceptor I was hoping to avoid wrapping and get the call redirected to my class for further processing, but unfortunately the compiler emits the closure before the redirection.

Did you found a workaround?