JIT: escape analysis for delegates

src/coreclr/jit/objectalloc.cpp

@@ -1869,6 +1869,14 @@ bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack<GenTree*>* parent
                             break;
                     }
                 }
+                else if (call->IsDelegateInvoke())
+                {
+                    if (tree == call->gtArgs.GetThisArg()->GetNode())
+                    {
+                        JITDUMP("Delegate invoke this...\n");
+                        canLclVarEscapeViaParentStack = false;
+                    }
+                }
 
                 // Note there is nothing special here about the parent being a call. We could move all this processing
                 // up to the caller and handle any sort of tree that could lead to escapes this way.
@@ -2219,6 +2227,7 @@ void ObjectAllocator::RewriteUses()
                 }
             }
             // Make box accesses explicit for UNBOX_HELPER
+            // Expand delegate invoke for calls where "this" is possibly stack pointing
             //
             else if (tree->IsCall())
             {
@@ -2267,6 +2276,51 @@ void ObjectAllocator::RewriteUses()
                         }
                     }
                 }
+                else if (call->IsDelegateInvoke())
+                {
+                    CallArg* const thisArg      = call->gtArgs.GetThisArg();
+                    GenTree* const delegateThis = thisArg->GetNode();
+
+                    if (delegateThis->OperIs(GT_LCL_VAR))
+                    {
+                        GenTreeLclVarCommon* const lcl = delegateThis->AsLclVarCommon();
+
+                        if (m_allocator->DoesLclVarPointToStack(lcl->GetLclNum()))
+                        {
+                            JITDUMP("Expanding delegate invoke [%06u]\n", m_compiler->dspTreeID(call));
+                            // Expand the delgate invoke early, so that physical promotion has
+                            // a chance to promote the delegate fields.
+                            //
+                            // Note the instance field may also be stack allocatable (someday)
+                            //
+                            GenTree* const cloneThis      = m_compiler->gtClone(lcl);
+                            unsigned const instanceOffset = m_compiler->eeGetEEInfo()->offsetOfDelegateInstance;
+                            GenTree* const newThisAddr =
+                                m_compiler->gtNewOperNode(GT_ADD, TYP_I_IMPL, cloneThis,
+                                                          m_compiler->gtNewIconNode(instanceOffset, TYP_I_IMPL));
+
+                            // For now assume the instance is heap...
+                            //
+                            GenTree* const newThis = m_compiler->gtNewIndir(TYP_REF, newThisAddr);
+                            thisArg->SetEarlyNode(newThis);
+
+                            // the control target is
+                            // [originalThis + firstTgtOffs]
+                            //
+                            unsigned const targetOffset = m_compiler->eeGetEEInfo()->offsetOfDelegateFirstTarget;
+                            GenTree* const targetAddr =
+                                m_compiler->gtNewOperNode(GT_ADD, TYP_I_IMPL, lcl,
+                                                          m_compiler->gtNewIconNode(targetOffset, TYP_I_IMPL));
+                            GenTree* const target = m_compiler->gtNewIndir(TYP_I_IMPL, targetAddr);
+
+                            // Update call state -- now an indirect call to the delegate target
+                            //
+                            call->gtCallAddr = target;
+                            call->gtCallType = CT_INDIRECT;
+                            call->gtCallMoreFlags &= ~(GTF_CALL_M_DELEGATE_INV | GTF_CALL_M_WRAPPER_DELEGATE_INV);
+                        }
+                    }
+                }
             }
             else if (tree->OperIsIndir())
             {

src/tests/JIT/opt/ObjectStackAllocation/Delegates.cs

@@ -0,0 +1,125 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Reflection;
+using System.Diagnostics;
+using System.Runtime.CompilerServices;
+using Xunit;
+
+enum AllocationKind
+{
+    Heap,
+    Stack,
+    Undefined
+}
+
+delegate int Test();
+
+public class Delegates
+{
+    static bool GCStressEnabled()
+    {
+        return Environment.GetEnvironmentVariable("DOTNET_GCStress") != null;
+    }
+
+    static AllocationKind StackAllocation()
+    {
+        AllocationKind expectedAllocationKind = AllocationKind.Stack;
+        if (GCStressEnabled())
+        {
+            Console.WriteLine("GCStress is enabled");
+            expectedAllocationKind = AllocationKind.Undefined;
+        }
+        return expectedAllocationKind;
+    }
+
+    static int CallTestAndVerifyAllocation(Test test, int expectedResult, AllocationKind expectedAllocationsKind, bool throws = false)
+    {
+        string methodName = test.Method.Name;
+        try
+        {
+            long allocatedBytesBefore = GC.GetAllocatedBytesForCurrentThread();
+            int testResult = test();
+            long allocatedBytesAfter = GC.GetAllocatedBytesForCurrentThread();
+
+            if (testResult != expectedResult)
+            {
+                Console.WriteLine($"FAILURE ({methodName}): expected {expectedResult}, got {testResult}");
+                return -1;
+            }
+
+            if ((expectedAllocationsKind == AllocationKind.Stack) && (allocatedBytesBefore != allocatedBytesAfter))
+            {
+                Console.WriteLine($"FAILURE ({methodName}): unexpected allocation of {allocatedBytesAfter - allocatedBytesBefore} bytes");
+                return -1;
+            }
+
+            if ((expectedAllocationsKind == AllocationKind.Heap) && (allocatedBytesBefore == allocatedBytesAfter))
+            {
+                Console.WriteLine($"FAILURE ({methodName}): unexpected stack allocation");
+                return -1;
+            }
+            else
+            {
+                Console.WriteLine($"SUCCESS ({methodName})");
+                return 100;
+            }
+        }
+        catch
+        {
+            if (throws)
+            {
+                Console.WriteLine($"SUCCESS ({methodName})");
+                return 100;
+            }
+            else
+            {
+                return -1;
+            }
+        }
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static int DoTest0(int a)
+    {
+        var f = (int x) => x + 1;
+        return f(a);
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static int RunTest0() => DoTest0(100);
+
+    [Fact]
+    public static int Test0() 
+    {
+        RunTest0();
+        return CallTestAndVerifyAllocation(RunTest0, 101, StackAllocation());
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static int DoTest1(int[] a)
+    {
+        var f = (int x) => x + 1;
+        int sum = 0;
+
+        foreach (int i in a)
+        {
+            sum += f(i);
+        }
+
+        return sum;
+    }
+
+    static int[] s_a = new int[100];
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static int RunTest1() => DoTest1(s_a);
+
+    [Fact]
+    public static int Test1()
+    {
+        RunTest1();
+        return CallTestAndVerifyAllocation(RunTest1, 100, StackAllocation());
+    }
+}

src/tests/JIT/opt/ObjectStackAllocation/Delegates.csproj

@@ -0,0 +1,12 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <!-- Needed for CLRTestEnvironmentVariable -->
+    <RequiresProcessIsolation>true</RequiresProcessIsolation>
+    <DebugType>None</DebugType>
+    <Optimize>True</Optimize>
+    <JitOptimizationSensitive>true</JitOptimizationSensitive>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+</Project>