JIT: Extend escape analysis to account for arrays with non-gcref elements #104906

hez2010 · 2024-07-15T17:09:51Z

Positive case:

var chs = new char[42];
chs[1] = 'a';
Console.WriteLine((int)chs[1] + chs.Length);

Codegen:

; Assembly listing for method ArrayAllocator.Program:Main() (FullOpts)
; Emitting BLENDED_CODE for X64 with AVX - Windows
; FullOpts code
; optimized code
; rsp based frame
; partially interruptible
; No PGO data
; Final local variable assignments
;
;* V00 loc0         [V00    ] (  0,  0   )    long  ->  zero-ref    class-hnd exact <short[]>
;  V01 OutArgs      [V01    ] (  1,  1   )  struct (32) [rsp+0x00]  do-not-enreg[XS] addr-exposed "OutgoingArgSpace"
;* V02 tmp1         [V02    ] (  0,  0   )  struct (104) zero-ref    do-not-enreg[SF] "stack allocated array temp"
;* V03 tmp2         [V03    ] (  0,  0   )    long  ->  zero-ref    single-def "V02.[000..008)"
;* V04 tmp3         [V04    ] (  0,  0   )     int  ->  zero-ref    single-def "V02.[008..012)"
;* V05 tmp4         [V05    ] (  0,  0   )   short  ->  zero-ref    "V02.[018..020)"
;
; Lcl frame size = 40

G_M25548_IG01:  ;; offset=0x0000
       sub      rsp, 40
                                                ;; size=4 bbWeight=1 PerfScore 0.25
G_M25548_IG02:  ;; offset=0x0004
       mov      ecx, 84
       call     [System.Console:WriteLine(int)]
       nop
                                                ;; size=12 bbWeight=1 PerfScore 3.50
G_M25548_IG03:  ;; offset=0x0010
       add      rsp, 40
       ret
                                                ;; size=5 bbWeight=1 PerfScore 1.25

; Total bytes of code 21, prolog size 4, PerfScore 5.00, instruction count 6, allocated bytes for code 21 (MethodHash=5b0b9c33) for method ArrayAllocator.Program:Main() (FullOpts)

Negative case:

var chs = new char[42];
chs[1] = 'a';
Console.WriteLine((int)chs[42] + chs.Length);

Codegen:

; Assembly listing for method ArrayAllocator.Program:Main() (FullOpts)
; Emitting BLENDED_CODE for X64 with AVX - Windows
; FullOpts code
; optimized code
; rsp based frame
; partially interruptible
; No PGO data
; Final local variable assignments
;
;* V00 loc0         [V00    ] (  0,  0   )    long  ->  zero-ref    class-hnd exact <short[]>
;  V01 OutArgs      [V01    ] (  1,  1   )  struct (32) [rsp+0x00]  do-not-enreg[XS] addr-exposed "OutgoingArgSpace"
;* V02 tmp1         [V02    ] (  0,  0   )  struct (104) zero-ref    do-not-enreg[SF] "stack allocated array temp"
;  V03 tmp2         [V03,T00] (  1,  0   )   byref  ->  rbx         must-init "dummy temp of must thrown exception"
;* V04 tmp3         [V04    ] (  0,  0   )    long  ->  zero-ref    single-def "V02.[000..008)"
;* V05 tmp4         [V05    ] (  0,  0   )     int  ->  zero-ref    single-def "V02.[008..012)"
;* V06 tmp5         [V06    ] (  0,  0   )   short  ->  zero-ref    single-def "V02.[018..020)"
;
; Lcl frame size = 32

G_M25548_IG01:  ;; offset=0x0000
       push     rbx
       sub      rsp, 32
       xor      ebx, ebx
                                                ;; size=7 bbWeight=0 PerfScore 0.00
G_M25548_IG02:  ;; offset=0x0007
       call     CORINFO_HELP_RNGCHKFAIL
       movsx    rcx, word  ptr [rbx]
       call     [System.Console:WriteLine(int)]
       int3
                                                ;; size=16 bbWeight=0 PerfScore 0.00

; Total bytes of code 23, prolog size 5, PerfScore 0.00, instruction count 7, allocated bytes for code 23 (MethodHash=5b0b9c33) for method ArrayAllocator.Program:Main() (FullOpts)
; ============================================================

Benchmark on Mandelbrot:

Method	Job	Mean	Error	StdDev	Code Size	Allocated
MandelBrot	NoStackAllocationArray	199.7 us	1.30 us	1.22 us	1,996 B	2.49 KB
MandelBrot	StackAllocationArray	195.8 us	1.16 us	1.08 us	2,414 B	1.14 KB

Diff: https://www.diffchecker.com/bNP4qHdF/

src/coreclr/jit/objectalloc.h

src/coreclr/jit/objectalloc.cpp

AndyAyersMS

For arrays (and also perhaps boxes and ref classes) we ought to have some kind of size limit... possibly similar to the one we use for stackallocs.

We need to be careful we don't allocate a lot of stack for an object that might not be heavily used, as we'll pay per-call prolog zeroing costs.

src/coreclr/jit/lclmorph.cpp

hez2010 · 2025-01-18T16:26:34Z

@AndyAyersMS Just experimented a bit with late dead stores removal by unexposing locals in the liveness and repeating to convergence: commit 3450580

Before adding late dead stores removal: diffs
After adding late dead stores removal: diffs

Relative diffs:

linux arm64: -26.588 bytes
linux x64: -27,113 bytes
osx arm64: -14,360 bytes
windows arm64: -25,528 bytes
windows x64: -26,628 bytes

Diffs look interesting but TP impacts are too high. Now I'm reverting the experiment commit.

AndyAyersMS · 2025-01-20T01:53:32Z

@jakobbotsch can you look this one over again?

src/coreclr/jit/gentree.h

src/coreclr/jit/helperexpansion.cpp

src/coreclr/jit/gentree.cpp

src/coreclr/jit/objectalloc.cpp

jakobbotsch · 2025-01-20T08:25:43Z

src/coreclr/jit/objectalloc.cpp

+#ifdef FEATURE_READYTORUN
+                        if (comp->opts.IsReadyToRun() && data->IsHelperCall(comp, CORINFO_HELP_READYTORUN_NEWARR_1))
+                        {
+                            len = data->AsCall()->gtArgs.GetArgByIndex(0)->GetNode();


I removed R2R support more completely

src/coreclr/jit/helperexpansion.cpp

jakobbotsch · 2025-01-20T08:49:38Z

src/coreclr/jit/helperexpansion.cpp

+    GenTree* const   mtStore = gtNewStoreLclFldNode(lclNum, TYP_I_IMPL, 0, mt);
+    Statement* const mtStmt  = gtNewStmt(mtStore);
+
+    fgInsertStmtBefore(block, newStmt, mtStmt);
+    gtSetStmtInfo(mtStmt);
+    fgSetStmtSeq(mtStmt);


Suggested change

GenTree* const mtStore = gtNewStoreLclFldNode(lclNum, TYP_I_IMPL, 0, mt);

Statement* const mtStmt = gtNewStmt(mtStore);

fgInsertStmtBefore(block, newStmt, mtStmt);

gtSetStmtInfo(mtStmt);

fgSetStmtSeq(mtStmt);

GenTree* const mtStore = gtNewStoreLclFldNode(lclNum, TYP_I_IMPL, 0, mt);

Statement* const mtStmt = fgNewStmtFromTree(mtStore);

fgInsertStmtBefore(block, newStmt, mtStmt);

Ditto below, more things can switch to use fgNewStmtFromTree

src/coreclr/jit/helperexpansion.cpp

jakobbotsch · 2025-01-20T09:03:47Z

src/coreclr/jit/objectalloc.h

@@ -181,6 +254,11 @@ inline bool ObjectAllocator::CanAllocateLclVarOnStack(unsigned int         lclNu
        return false;
    }


It seems like we should have a limit on the aggregate stack allocated size. That's somewhat preexisting, but probably more important now.

Agree but will do that in a future PR

jakobbotsch · 2025-01-20T09:19:47Z

src/coreclr/jit/valuenum.cpp

+    // If this is a local array, there are no asyncronous modifications, so we can set the
+    // conservative VN to the liberal VN.
+    //
+    VNFuncApp arrFn;
+    if (vnStore->IsVNNewLocalArr(arrVN, &arrFn))
+    {
+        loadTree->gtVNPair.SetConservative(loadValueVN);
+    }
+    else
+    {
+        loadTree->gtVNPair.SetConservative(vnStore->VNForExpr(compCurBB, loadType));
+    }


Does this actually show up as benefits?

In some limited cases, yes... eg we can const propagate through a[2] = 1; y = a[2];

AndyAyersMS · 2025-01-21T15:15:42Z

@jakobbotsch think I've addressed most of the key points. Overall size limit will come in a future PR.

jakobbotsch · 2025-01-22T12:29:41Z

@AndyAyersMS Did you push those changes?

AndyAyersMS · 2025-01-22T15:10:11Z

Ah, I pushed to my fork, but ... this PR is not from my fork.

AndyAyersMS · 2025-01-22T15:35:26Z

@jakobbotsch changes are there now

jakobbotsch · 2025-01-22T19:24:39Z

src/coreclr/jit/gentree.h

@@ -4227,6 +4227,7 @@ enum GenTreeCallFlags : unsigned int
    GTF_CALL_M_CAST_CAN_BE_EXPANDED    = 0x04000000, // this cast (helper call) can be expanded if it's profitable. To be removed.
    GTF_CALL_M_CAST_OBJ_NONNULL        = 0x08000000, // if we expand this specific cast we don't need to check the input object for null
                                                     // NOTE: if needed, this flag can be removed, and we can introduce new _NONNUL cast helpers
+    GTF_CALL_M_STACK_ARRAY             = 0x10000000, // this call is a new array helper for a stack allocated array.


I suppose we weren't able to get rid of this since it's also used in VN.

jakobbotsch · 2025-01-22T19:27:33Z

src/coreclr/jit/valuenum.cpp

+                    // modHeap = false;
+                }
+                else if (vnf == VNF_JitReadyToRunNewArr)
+                {
+                    vnf = VNF_JitReadyToRunNewLclArr;
+                    // modHeap = false;


Nit: commented code (feel free to remove as part of a follow up)

jakobbotsch · 2025-01-22T19:28:21Z

src/coreclr/jit/valuenum.cpp

@@ -14338,7 +14367,22 @@ bool Compiler::fgValueNumberHelperCall(GenTreeCall* call)
                }
            }

+            if (isAlloc && ((call->gtCallMoreFlags & GTF_CALL_M_STACK_ARRAY) != 0))


I suppose this could be changed to FindWellKnownArg(StackArrayLocal) != nullptr to get rid of the flag

jakobbotsch

LGTM

AndyAyersMS · 2025-01-22T19:30:44Z

Thanks. I have more changes in this area so I can handle the last few bits in a subsequent PR.

AndyAyersMS · 2025-01-22T19:35:03Z

@hez2010 thanks for all the work you did here.

hez2010 added 7 commits July 15, 2024 20:23

initial prototype

1b0e3d3

Morph ARR_LENGTH and INDEX_ADDR

57b7e42

Fix incorrect array length storage

1b5b25e

Use offset and correct type

395b735

handle reassignment

17de70b

range check

5443c42

throw range check failure

b2d07da

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

dotnet-policy-service bot added the community-contribution Indicates that the PR has been added by a community member label Jul 15, 2024

hez2010 added 2 commits July 16, 2024 02:13

update comments

b5ae9e7

add metrics

87b29de

jakobbotsch reviewed Jul 15, 2024

View reviewed changes

src/coreclr/jit/objectalloc.h Outdated Show resolved Hide resolved

jakobbotsch reviewed Jul 15, 2024

View reviewed changes

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

jakobbotsch reviewed Jul 15, 2024

View reviewed changes

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

minor cleanup

eeb681d

AndyAyersMS reviewed Jul 15, 2024

View reviewed changes

AndyAyersMS mentioned this pull request Jul 16, 2024

Stack Allocation Enhancements #104936

Open

20 tasks

hez2010 added 5 commits July 16, 2024 14:13

Introduce new temp and implement local address morphing

dee9f38

handle index out-of-range

94c103b

Refactor to remove duplicates

12b297b

Remove invalid asserts

e0fa91e

make compiler happy

9e0a04f

jakobbotsch reviewed Jul 16, 2024

View reviewed changes

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