Stack Allocation Enhancements

AndyAyersMS commented 3 months ago

Stack allocation of non-escaping ref classes and boxed value classes was enabled in #103361, but only works in limited cases. This issue tracks further enhancements (see also #11192).

Abilities:

[ ] stack allocation of constant-sized arrays of non-GC types
- https://github.com/dotnet/runtime/pull/104906
[ ] stack allocation of constant-sized arrays of GC types
[ ] stack allocation of strings
[ ] stack allocation of boxed nullables

[ ] zeroing strategy (zero in prolog vs zero at first use)
[ ] stack allocation of some delegates (and perhaps their closures), eg

    public static int Test()
    {
        int a = 100;
        Func<int> f = () => { return a; };
        return f();
    }

a small tweak to escape analysis gets the delegate on the stack, but the invoke expansion currently happens in lower so we don't get any physical promotion. We would need to move this earlier.

See note below.

[ ] evaluate when/if it makes sense to transform stack objects to localloc instead of fixed allocations (at least for non-gc types; for GC types there's currently no way to do proper GC reporting)
[ ] Allocations in loops:
- https://github.com/dotnet/runtime/pull/106526
- https://devblogs.microsoft.com/java/improving-openjdk-scalar-replacement-part-1-3/

Analysis:

[ ] make the analysis more sophisticated to handle increasingly more complex examples (eg fields of value classes, or fields of unescaped ref classes). See for instance
- 84872
- 61455
- https://github.com/dotnet/runtime/pull/105162
[ ] special case calls to helpers where arguments don't escape (this will require ensuring that the helpers report gc arguments as interior to gc)
[ ] make inliner more aware when it can enhance stack allocation
- https://github.com/dotnet/runtime/issues/104479
[ ] improve address-exposed analysis (an exposed stack allocated ref class probably performs worse than a heap allocated one, in most cases)
- https://github.com/dotnet/runtime/issues/104250

Implementation:

[ ] stop relying on top-level ALLOCOBJ
[ ] stop relying on ALLOCOBJ assigned to single-def temp in importer
[ ] use custom GC layout for boxes; stop fetching placeholder type from the runtime
- https://github.com/dotnet/runtime/issues/103362
[ ] make sure object stack allocation doesn't block fast tail call optimization unnecessarily (currently fast tail call optimization is disabled if there are any exposed local)

NAOT:

[ ] interprocedural escape analysis. May also be viable in jitted contexts, either as a hint for the inliner or as some sort of property we can guarantee on profiler-driven re-jit.

Advanced:

[ ] partial escape analysis. Allocate objects that are unlikely to escape on the stack. Either compute the "escape frontier" of an object and copy it to the stack when that frontier is crossed, or else add capabilities to write barriers to note when a stack allocated object reference is going to be stored on the heap, and "promote" the object at that point (using GC info to rewrite the stack references to heap references). Likely requires PGO, to ensure we're not wrong too often.
[ ] for objects that don't escape but that we don't want to stack allocate, treat them as thread private: we can use more aggressive value numbering for instance, since we don't have to assume the field values can change asynchronously.

Diagnostics:

[ ] if VM/GC/WriteBarriers catch an escaped object-on-stack they should provide a helpful assert instead of a generic GC hole like assert. We can either reserve a debug bit in the sync block or rely on object's address being within the code heap.

FYI @dotnet/jit-contrib

dotnet-policy-service[bot] commented 3 months ago

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

AndyAyersMS commented 3 months ago

For the delegate case, if we add

index 89e28c5978c..a069d98503e 100644
--- a/src/coreclr/jit/objectalloc.cpp
+++ b/src/coreclr/jit/objectalloc.cpp
@@ -719,12 +719,23 @@ bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack<GenTree*>* parent

             case GT_CALL:
             {
-                GenTreeCall* asCall = parent->AsCall();
+                GenTreeCall* const call = parent->AsCall();

-                if (asCall->IsHelperCall())
+                if (call->IsHelperCall())
                 {
                     canLclVarEscapeViaParentStack =
-                        !Compiler::s_helperCallProperties.IsNoEscape(comp->eeGetHelperNum(asCall->gtCallMethHnd));
+                        !Compiler::s_helperCallProperties.IsNoEscape(comp->eeGetHelperNum(call->gtCallMethHnd));
+                }
+                else if (call->gtCallType == CT_USER_FUNC)
+                {
+                    // Delegate invoke won't escape the delegate which is passed as "this"
+                    // And gets expanded inline later.
+                    //
+                    if ((call->gtCallMoreFlags & GTF_CALL_M_DELEGATE_INV) != 0)
+                    {
+                        GenTree* const thisArg = call->gtArgs.GetThisArg()->GetNode();
+                        canLclVarEscapeViaParentStack = thisArg != tree;
+                    }

Then the example above becomes

; Method Y:Test():int (FullOpts)
G_M53607_IG01:  ;; offset=0x0000
       sub      rsp, 88
       vxorps   xmm4, xmm4, xmm4
       vmovdqu  ymmword ptr [rsp+0x20], ymm4
       vmovdqa  xmmword ptr [rsp+0x40], xmm4
       xor      eax, eax
       mov      qword ptr [rsp+0x50], rax
                        ;; size=27 bbWeight=1 PerfScore 5.83

G_M53607_IG02:  ;; offset=0x001B
       mov      rcx, 0x7FFD4BB04580      ; Y+<>c__DisplayClass0_0
       call     CORINFO_HELP_NEWSFAST
       mov      dword ptr [rax+0x08], 100
       mov      rcx, 0x7FFD4BB04B30      ; System.Func`1[int]
       mov      qword ptr [rsp+0x20], rcx
       mov      gword ptr [rsp+0x28], rax
       mov      rcx, 0x7FFD4B8783F0      ; code for Y+<>c__DisplayClass0_0:<Test>b__0():int:this
       mov      qword ptr [rsp+0x38], rcx
       lea      rax, [rsp+0x20]
       mov      rcx, gword ptr [rax+0x08]
       call     [rax+0x18]System.Func`1[int]:Invoke():int:this
       nop      
                        ;; size=70 bbWeight=1 PerfScore 11.50

G_M53607_IG03:  ;; offset=0x0061
       add      rsp, 88
       ret      
                        ;; size=5 bbWeight=1 PerfScore 1.25
; Total bytes of code: 102

where the closure is still on the heap and we're invoking the delegate func "directly" but via a convoluted path where we store the func's (indirection cell) address to the stack allocated delegate and then fetch it back and indirect through it.

Ideally we'd like to be able to inline and perhaps realize the closure doesn't escape either, but that seems far off. Perhaps we can just summarily claim the closure can't escape. I am not sure.

Moving delegate invoke expansion earlier does not look to be simple -- currently there is some prep work in morph and then the actual expansion in lower, and tail calls are a complication.

hez2010 commented 3 months ago

@AndyAyersMS With the array (non-gc elems) support + my field analysis prototype + the above delegate handling (branch at https://github.com/hez2010/runtime/tree/field-stackalloc), the codegen becomes:

G_M30166_IG01:  ;; offset=0x0000
       sub      rsp, 104
       vxorps   xmm4, xmm4, xmm4
       vmovdqu  ymmword ptr [rsp+0x20], ymm4
       vmovdqa  xmmword ptr [rsp+0x40], xmm4
       xor      eax, eax
       mov      qword ptr [rsp+0x50], rax
                                                ;; size=27 bbWeight=1 PerfScore 5.83
G_M30166_IG02:  ;; offset=0x001B
       xor      ecx, ecx
       mov      qword ptr [rsp+0x58], rcx
       mov      dword ptr [rsp+0x60], ecx
       mov      rcx, 0x7FFEBB4211B8      ; Y+<>c__DisplayClass7_0
       mov      qword ptr [rsp+0x58], rcx
       mov      dword ptr [rsp+0x60], 100
       mov      rcx, 0x7FFEBB420EE8      ; System.Func`1[int]
       mov      qword ptr [rsp+0x20], rcx
       lea      rcx, [rsp+0x58]
       mov      qword ptr [rsp+0x28], rcx
       mov      rcx, 0x7FFEBB3F0678      ; code for Y+<>c__DisplayClass7_0:<Test>b__0():int:this
       mov      qword ptr [rsp+0x38], rcx
       lea      rax, [rsp+0x20]
       mov      rcx, gword ptr [rax+0x08]
       call     [rax+0x18]System.Func`1[int]:Invoke():int:this
       nop
                                                ;; size=87 bbWeight=1 PerfScore 14.25
G_M30166_IG03:  ;; offset=0x0072
       add      rsp, 104
       ret

EgorBo commented 2 months ago

Added Diagnostics section

hez2010 commented 2 months ago

Put more useful links:

An overview of the impl of esacpe analysis including interprocedural analysis in JVM: https://cr.openjdk.org/~cslucas/escape-analysis/EscapeAnalysis.html

Some insights on allocating objects in a loop and etc.: https://devblogs.microsoft.com/java/improving-openjdk-scalar-replacement-part-2-3/

dotnet / runtime