NativeAOT: Partially expand static initialization #83911

EgorBo · 2023-03-24T23:03:41Z

Closes #80954 and contributes to #64242

In the JIT world we successfully get rid of static class initializations (either beforefieldinit or normal cctors) naturally or via tiered compilation (Tier1 means that we already initialized eveyrything we needed on hot path in the previous tier). However, it's not the case for AOT. The plan is to help NativeAOT to avoid hitting call overhead in this case by partially inline "is class already initialized?" check with a fast path, e.g.:

// non-GC static
static int s_Int = int.Parse("42");
int GetInt() => s_Int;

// GC static
static string s_Str = (42).ToString();
string GetStr() => s_Str;

Current codegen on NativeAOT:

; Method Prog:GetStr():System.String:this
G_M45136_IG01:              
       4883EC28             sub      rsp, 40
G_M45136_IG02:              
       E800000000           call     CORINFO_HELP_READYTORUN_GCSTATIC_BASE
       488B4008             mov      rax, gword ptr [rax+08H]
G_M45136_IG03:              
       4883C428             add      rsp, 40
       C3                   ret      
; Total bytes of code: 18


; Method Prog:GetInt():int:this
G_M29019_IG01:              
       4883EC28             sub      rsp, 40
G_M29019_IG02:              
       E800000000           call     CORINFO_HELP_READYTORUN_NONGCSTATIC_BASE
       8B00                 mov      eax, dword ptr [rax]
G_M29019_IG03:              
       4883C428             add      rsp, 40
       C3                   ret      
; Total bytes of code: 16

New codegen (NativeAOT):

; Method Prog:GetStr():System.String:this
       4883EC28             sub      rsp, 40
       488D0500000000       lea      rax, [(reloc 0x4000000000421128)]
       8378F801             cmp      dword ptr [rax-08H], 1
       7510                 jne      SHORT G_M25482_IG05
G_M25482_IG03:              
       488B0500000000       mov      rax, qword ptr [(reloc 0x4000000000421130)]
       488B4008             mov      rax, gword ptr [rax+08H]
       4883C428             add      rsp, 40
       C3                   ret      
G_M25482_IG05:              
       E800000000           call     CORINFO_HELP_READYTORUN_GCSTATIC_BASE
       EBE9                 jmp      SHORT G_M25482_IG03
; Total bytes of code: 40


; Method Prog:GetInt():int:this
       56                   push     rsi
       4883EC20             sub      rsp, 32
       488D3500000000       lea      rsi, [(reloc 0x4000000000421008)]
       837EF801             cmp      dword ptr [rsi-08H], 1
       7508                 jne      SHORT G_M56961_IG05
G_M56961_IG03:             
       8B06                 mov      eax, dword ptr [rsi]
       4883C420             add      rsp, 32
       5E                   pop      rsi
       C3                   ret      
G_M56961_IG05:              
       E800000000           call     CORINFO_HELP_READYTORUN_NONGCSTATIC_BASE
       EBF1                 jmp      SHORT G_M56961_IG03
; Total bytes of code: 33

It also works for JIT (mainly, to test it, but could be useful for TC=0).

Benchmarks

I was using JIT with TieredCompilation=0

// case 1: normal cctor. JIT doesn't hoist such initializations
// from loops (needs loop peeling)

public class TestClassWithCctor
{
    public static int field;

    static TestClassWithCctor()
    {
        field = 42;
    }
}

[Benchmark]
public int NonHoistableStaticInit()
{
    int sum = 0;
    for (int i = 0; i < 10000; i++)
    {
        sum += TestClassWithCctor.field;
    }
    return sum;
}

// case 2: beforefieldinit

public class TestClassWithBeforefieldinit
{
    public static int field = 42;
}

[Benchmark]
public int SimpleStaticInit()
{
    return TestClassWithBeforefieldinit.field;
}

Method	Toolchain	Mean
NonHoistableStaticInit	\runtime-base\corerun.exe	11,123.0890 ns
NonHoistableStaticInit	\runtime\corerun.exe	1,893.7071 ns

SimpleStaticInit	\runtime-base\corerun.exe	0.9206 ns
SimpleStaticInit	\runtime\corerun.exe	0.0302 ns

codegen diff: https://www.diffchecker.com/87bWA1er/

Can't say yet the total size overhead for a hello world, in the worst case we can leave this optimization only for "prefer speed" mode (-Os)

ghost · 2023-03-24T23:08:19Z

Tagging subscribers to this area: @JulieLeeMSFT, @jakobbotsch, @kunalspathak
See info in area-owners.md if you want to be subscribed.

Issue Details

Closes #80954 and contributes to #64242

In the JIT world we successfully get rid of static class initializations (either beforefieldinit or normal cctors) naturally or via tiered compilation (Tier1 means that we already initialized eveyrything we needed on hot path in the previous tier). However, it's not the case for AOT. The plan is to help NativeAOT to avoid hitting call overhead in this case by partially inline "is class already initialized?" check with a fast path, e.g.:

static int field = int.Parse("42"); // mimic a complex cctor

int Test()
{
    return field;
}

Current codegen on NativeAOT:

; Method Prog:Test():int:this
       sub      rsp, 40
       call     CORINFO_HELP_READYTORUN_NONGCSTATIC_BASE
       mov      eax, dword ptr [rax]
       add      rsp, 40
       ret

Expected codegen:

       sub      rsp, 40
       test     byte  ptr [(reloc)], 1 ;; is already initialized?
       jne      SHORT G_M3272_IG04
       call     CORINFO_HELP_READYTORUN_NONGCSTATIC_BASE ;; it's not -- fallback
SHORT G_M3272_IG04:
       mov      eax, dword ptr [(reloc)] ;; just access field's value
       add      rsp, 40
       ret

The implementation in this PR only handles JIT for now just to test it, while I'm trying to figure out how to implement getIsClassInitedFieldAddress API on NativeAOT side, might need @MichalStrehovsky help

Author:	EgorBo
Assignees:	-
Labels:	`area-CodeGen-coreclr`
Milestone:	-

src/coreclr/jit/flowgraph.cpp

jkotas · 2023-03-24T23:59:22Z

getIsClassInitedFieldAddress API on NativeAOT side,

factory.TypeNonGCStaticsSymbol(type) and subtract NonGCStaticsNode.GetClassConstructorContextSize(factory.Target) offset from it.

It should be compared with 1 (CoreCLR does bit test, native aot uses comparison).

…SE in JitDisasm

EgorBo · 2023-03-25T14:47:09Z

Managed to get NAOT working locally (not yet pushed), so for

static int field = int.Parse("42"); // mimic a complex cctor

int Test()
{
    return field;
}

Was:

; Method Prog:Test():int:this
       sub      rsp, 40
       call     CORINFO_HELP_READYTORUN_NONGCSTATIC_BASE
       mov      eax, dword ptr [rax]
       add      rsp, 40
       ret
; Total bytes of code 16

Now I get:

; Assembly listing for method Prog:Test():int:this
       sub      rsp, 40
       lea      rax, [(reloc)]
       cmp      dword ptr [rax-08H], 0
       je       SHORT G_M3272_IG05
G_M3272_IG03:
       mov      eax, dword ptr [(reloc)]
       add      rsp, 40
       ret
G_M3272_IG05:
       call     CORINFO_HELP_READYTORUN_NONGCSTATIC_BASE
       jmp      SHORT G_M3272_IG03
; Total bytes of code 35

(nongc statics)

EgorBo · 2023-03-25T15:28:15Z

Same but for JIT:

Was:

; Assembly listing for method Prog:Test():int:this
G_M3272_IG01:
       sub      rsp, 40
       mov      rcx, 0xD1FFAB1E
       mov      edx, 3
       call     CORINFO_HELP_GETSHARED_NONGCSTATIC_BASE
       mov      eax, dword ptr [(reloc)]
       add      rsp, 40
       ret      
; Total bytes of code 35

New:

; Assembly listing for method Prog:Test():int:this
       sub      rsp, 40
       test     byte  ptr [(reloc)], 1      
       je       SHORT G_M3272_IG05
G_M3272_IG03:
       mov      eax, dword ptr [(reloc)]
       add      rsp, 40
       ret      
G_M3272_IG05:
       mov      rcx, 0xD1FFAB1E
       mov      edx, 3
       call     CORINFO_HELP_GETSHARED_NONGCSTATIC_BASE
       jmp      SHORT G_M3272_IG03
; Total bytes of code 46

Presumably, JIT version can be made 5 bytes smaller if we replace the cold part with "call INITCLASS(cls)".

EgorBo · 2023-03-25T16:10:04Z

GC statics come with a bigger size increase on NAOT:

static string field = (42).ToString();

[MethodImpl(MethodImplOptions.NoInlining)]
string Test()
{
    return field;
}

Was:

; Method Prog:Test():System.String:this
       sub      rsp, 40
       call     CORINFO_HELP_READYTORUN_GCSTATIC_BASE
       mov      rax, gword ptr [rax+08H]
       add      rsp, 40
       ret      
; Total bytes of code: 18

Now:

; Method Prog:Test():System.String:this
       sub      rsp, 40
       lea      rax, [(reloc)] ;; NonGCStaticsBase
       cmp      dword ptr [rax-08H], 0
       je       SHORT G_M48517_IG05
G_M48517_IG03:
       mov      rax, qword ptr [(reloc)] ;; this is a different reloc (GCStaticBase)
       mov      rax, gword ptr [rax+08H]
       add      rsp, 40
       ret      
G_M48517_IG05:
       call     CORINFO_HELP_READYTORUN_GCSTATIC_BASE
       jmp      SHORT G_M48517_IG03
; Total bytes of code: 40

due to double-indirect for field access

EgorBo · 2023-03-25T16:46:01Z

Benchmarks (JIT, TieredCompilation=0):

// case 1: normal cctor. JIT doesn't hoist such initializations
// from loops (needs loop peeling)

public class TestClassWithCctor
{
    public static int field;

    static TestClassWithCctor()
    {
        field = 42;
    }
}

[Benchmark]
public int NonHoistableStaticInit()
{
    int sum = 0;
    for (int i = 0; i < 10000; i++)
    {
        sum += TestClassWithCctor.field;
    }
    return sum;
}

// case 2: beforefieldinit

public class TestClassWithBeforefieldinit
{
    public static int field = 42;
}

[Benchmark]
public int SimpleStaticInit()
{
    return TestClassWithBeforefieldinit.field;
}

Method	Toolchain	Mean
NonHoistableStaticInit	\runtime-base\corerun.exe	11,123.0890 ns
NonHoistableStaticInit	\runtime\corerun.exe	1,893.7071 ns

SimpleStaticInit	\runtime-base\corerun.exe	0.9206 ns
SimpleStaticInit	\runtime\corerun.exe	0.0302 ns

codegen diff: https://www.diffchecker.com/87bWA1er/

EgorBo · 2023-03-25T20:40:21Z

/azp list

EgorBo · 2023-03-25T20:41:24Z

/azp run runtime-coreclr outerloop, runtime-extra-platforms, runtime-coreclr jitstress

azure-pipelines · 2023-03-25T20:41:55Z

Azure Pipelines successfully started running 3 pipeline(s).

Co-authored-by: Jan Kotas <[email protected]>

… expand-static-init

EgorBo · 2023-04-06T11:17:34Z

@kunalspathak @AndyAyersMS @SingleAccretion I've addressed feedback, anything else?

src/coreclr/jit/flowgraph.cpp

src/coreclr/jit/importer.cpp

Co-authored-by: SingleAccretion <[email protected]>

EgorBo · 2023-04-06T16:45:59Z

@SingleAccretion Thanks!

@kunalspathak @AndyAyersMS waiting for a green approve now 🙂

src/coreclr/vm/jitinterface.cpp

src/coreclr/jit/flowgraph.cpp

src/coreclr/tools/aot/ILCompiler.RyuJit/JitInterface/CorInfoImpl.RyuJit.cs

src/coreclr/jit/flowgraph.cpp

AndyAyersMS

Would it make sense to have some kind of "post phase" check to make sure each possible candidate was either expanded or intentionally not expanded?

src/coreclr/jit/flowgraph.cpp

EgorBo · 2023-04-06T20:26:27Z

Would it make sense to have some kind of "post phase" check to make sure each possible candidate was either expanded or intentionally not expanded?

E.g. for runtime lookups I put an assert in Lower.cpp to make sure all of them are expanded (because it was required). In this case I didn't do that because it's just an optimization and we skip some type of static initializations on JIT already + we skip cold blocks so I didn't do that

kunalspathak

LGTM

Expand static initialization in JIT

7100058

dotnet-issue-labeler bot added the area-CodeGen-coreclr CLR JIT compiler in src/coreclr/src/jit and related components such as SuperPMI label Mar 24, 2023

ghost assigned EgorBo Mar 24, 2023

Clean up

023e2c8

jkotas reviewed Mar 24, 2023

View reviewed changes

src/coreclr/jit/flowgraph.cpp Outdated Show resolved Hide resolved

EgorBo added 3 commits March 25, 2023 01:28

Address feedback, clean up

6019949

Address feedback, baby steps towards NAOT impl

791244b

Remove test that always expects CORINFO_HELP_GETSHARED_NONGCSTATIC_BA…

4f7c857

…SE in JitDisasm

This was referenced Mar 25, 2023

Arm32: System.Threading.Tasks.Dataflow.Tests failing with NRE #80857

Closed

WasmTestOnBrowser-System.* test failures in CI #83655

Closed

NativeAOT non-GC impl

620bbd1

Fix NAOT

4279269

Enable GC statics for NativeAOT

7b9a259

Merge branch 'main' of github.com:dotnet/runtime into expand-static-init

a79eba7

This was referenced Mar 25, 2023

IOException running NuGet-Migrations during tests in dotnet CLI first run #80619

Closed

[release/6.0] Doublelinklist GC failures on Mono #83245

Closed

Fix assert

c1744bc

This comment was marked as outdated.

Sign in to view

build-analysis bot mentioned this pull request Mar 25, 2023

Tracking issue for CI build timeouts #76454

Closed

EgorBo added 2 commits March 25, 2023 23:34

code clean up

141e99e

Merge branch 'main' of github.com:dotnet/runtime into expand-static-init

9ab6dc1

EgorBo and others added 6 commits April 6, 2023 00:35

Update src/coreclr/jit/flowgraph.cpp

deb4eb6

Co-authored-by: Jan Kotas <[email protected]>

Update src/coreclr/jit/flowgraph.cpp

f0a44ba

Co-authored-by: Jan Kotas <[email protected]>

Merge branch 'expand-static-init' of github.com:EgorBo/runtime-1 into…

9b11b91

… expand-static-init

formatting

2570f84

Address feedback (1st round)

16c75b2

Address feedback (round 2)

2bbf5bf

EgorBo mentioned this pull request Apr 6, 2023

Fix ThunkGenerator #84390

Merged

EgorBo added 2 commits April 6, 2023 03:17

Address feedback

ddc1731

Fix assert on NAOT

f65d080

MichalStrehovsky mentioned this pull request Apr 6, 2023

CctorForWrite failing with NativeAOT #84400

Closed

SingleAccretion approved these changes Apr 6, 2023

View reviewed changes

EgorBo and others added 2 commits April 6, 2023 18:44

Apply suggestions from code review

c7bc7d8

Co-authored-by: SingleAccretion <[email protected]>

Update src/coreclr/jit/flowgraph.cpp

6d1d18a

Co-authored-by: SingleAccretion <[email protected]>