JIT: fix issue in conditional escape analysis (#118214)

Conditional escape analysis is triggered by a GDV guarded call to
`GetEnumerator`. On the fast path this call is devirtualized and inlined
and exposes one or more enumerator object allocations. The "more" allows
for collections to do special-case optimizations like returning a static
enumerator object for an empty collection.

When there is more than one enumerator object allocation under the GDV guard,
ensure that only the allocation that was analyzed for conditional escape is stack
allocated; the remaining enumerators will end up on the "slow path" where enumeration
happens via interface calls, and so must be heap objects.

If we clone an allocation site the allocation block won't be in
the DFS tree. Such sites will not be in an enumerator GDV hammock
so can be exempted from checking.

Fixes #111922.

Author: AndyAyersMS

src/coreclr/jit/objectalloc.cpp

@@ -1324,7 +1324,8 @@ void ObjectAllocator::MorphAllocObjNode(AllocationCandidate& candidate)
     else
     {
         assert(candidate.m_onHeapReason != nullptr);
-        JITDUMP("Allocating V%02u on the heap: %s\n", lclNum, candidate.m_onHeapReason);
+        JITDUMP("Allocating V%02u / [%06u] on the heap: %s\n", lclNum, comp->dspTreeID(candidate.m_tree),
+                candidate.m_onHeapReason);
         if ((candidate.m_allocType == OAT_NEWOBJ) || (candidate.m_allocType == OAT_NEWOBJ_HEAP))
         {
             GenTree* const stmtExpr      = candidate.m_tree;
@@ -1367,6 +1368,58 @@ bool ObjectAllocator::MorphAllocObjNodeHelper(AllocationCandidate& candidate)
         return false;
     }
 
+    // If we have done any cloning for conditional escape, verify this allocation
+    // is not on one of the "slow paths" that are now forcibly making interface calls.
+    //
+    for (CloneInfo* const c : CloneMap::ValueIteration(&m_CloneMap))
+    {
+        if (!c->m_willClone)
+        {
+            // We didn't clone so there is no check needed.
+            //
+            continue;
+        }
+
+        if (c->m_guardBlock == nullptr)
+        {
+            // This clone wasn't guarded by a GDV, so we should
+            // only have a fast path.
+            continue;
+        }
+
+        GenTree* const allocTree = candidate.m_tree->AsLclVar()->Data();
+
+        if (c->m_allocTree == allocTree)
+        {
+            // This is the conditionally escaping allocation, so it's
+            // on the fast path.
+            //
+            continue;
+        }
+
+        if (!comp->m_dfsTree->Contains(candidate.m_block))
+        {
+            // If the def block is not in the DFS tree then the allocation
+            // is not within the hammock, so it is not on the slow path.
+            //
+            continue;
+        }
+
+        // This allocation candidate might be on the "slow path". Check.
+        //
+        // c->m_guardBlock and c->m_defBlock form a GDV hammock.
+        //
+        // So if an allocation site is dominated by c->m_guardBlock and
+        // not dominated by c->m_defBlock, it is within the hammock.
+        //
+        if (comp->m_domTree->Dominates(c->m_guardBlock, candidate.m_block) &&
+            !comp->m_domTree->Dominates(c->m_defBlock, candidate.m_block))
+        {
+            candidate.m_onHeapReason = "[on slow path of conditional escape clone]";
+            return false;
+        }
+    }
+
     switch (candidate.m_allocType)
     {
         case OAT_NEWARR:
@@ -3772,6 +3825,8 @@ bool ObjectAllocator::CheckCanClone(CloneInfo* info)
     JITDUMP("V%02u has single def in " FMT_BB " at [%06u]\n", info->m_local, defBlock->bbNum,
             comp->dspTreeID(defStmt->GetRootNode()));
 
+    info->m_defBlock = defBlock;
+
     // We expect to be able to follow all paths from alloc block to defBlock
     // without reaching "beyond" defBlock.
     //
@@ -3861,6 +3916,35 @@ bool ObjectAllocator::CheckCanClone(CloneInfo* info)
         JITDUMP("Unexpected, alloc site " FMT_BB " dominates def block " FMT_BB ". We will clone anyways.\n",
                 allocBlock->bbNum, defBlock->bbNum);
     }
+    else
+    {
+        // The allocation is conditional. See if we can find the GDV guard
+        // so we can check what other possible stack allocations might reach
+        // the def block.
+        //
+        BasicBlock* possibleGuardBlock = defBlock->bbIDom;
+        GuardInfo   gi;
+        if (IsGuard(possibleGuardBlock, &gi))
+        {
+            // Todo: validate guard is the right guard.
+            //
+            JITDUMP("Conditional allocation is guarded by a GDV\n");
+            info->m_guardBlock = possibleGuardBlock;
+        }
+        else
+        {
+            // The allocation is conditional but the enumerator var def is
+            // not dominated by a GDV. Perhaps we resolved the GDV already.
+            //
+            // We want to allow cloning to proceed in this case as alternative
+            // enumerators might be non-allocating (say the static empty array
+            // enumerator).
+            //
+            // But consider adding more checking in this case.
+            //
+            JITDUMP("Conditional allocation is not guarded by a GDV\n");
+        }
+    }
 
     // Classify the other local appearances
     // as Ts (allocTemps) or Us (useTemps), and look for guard appearances.

src/coreclr/jit/objectalloc.h

@@ -89,6 +89,12 @@ struct CloneInfo : public GuardInfo
     unsigned                  m_appearanceCount = 0;
     jitstd::vector<unsigned>* m_allocTemps      = nullptr;
 
+    // The initial GDV guard block, if any.
+    BasicBlock* m_guardBlock = nullptr;
+
+    // The block where the enumerator var is assigned.
+    BasicBlock* m_defBlock = nullptr;
+
     // Where the enumerator allocation happens
     GenTree*    m_allocTree  = nullptr;
     Statement*  m_allocStmt  = nullptr;

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

@@ -0,0 +1,40 @@
+// 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.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Runtime.CompilerServices;
+using Xunit;
+
+public class Runtime_111922v2
+{
+    [Fact]
+    public static int Test()
+    {
+        LinkedList<string> e = new LinkedList<string>();
+        LinkedList<string> e1 = new LinkedList<string>();
+        e1.AddLast("b");
+        e1.AddLast("a");
+
+        int sum = -80;
+
+        for (int i = 0; i < 200; i++)
+        {
+            sum += Problem(i % 10 == 0 ? e : e1) ? 1 : 0;
+            Thread.Sleep(5);
+        }
+
+        Console.WriteLine(sum);
+        return sum;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static bool Problem(IEnumerable<string> e)
+    {
+        return e.Contains("a", StringComparer.OrdinalIgnoreCase);
+    }
+}
+
+

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

@@ -0,0 +1,9 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <DebugType>None</DebugType>
+    <Optimize>True</Optimize>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+</Project>