JIT: revise may/must point to stack analysis (#115080)

Keep track of locals that may point to the heap as well as locals that
may point to the stack.

The set of definitely stack-pointing locals is then the difference in
these two sets. This handles the "cyclic" case where stack pointing locals
are assigned to one another.

Fixes #115017

Author: AndyAyersMS

src/coreclr/jit/assertionprop.cpp

@@ -675,6 +675,8 @@ void Compiler::optAssertionInit(bool isLocalProp)
             {
                 optMaxAssertionCount = (AssertionIndex)min(maxTrackedLocals, ((3 * lvaTrackedCount / 128) + 1) * 64);
             }
+
+            JITDUMP("Cross-block table size %u (for %u tracked locals)\n", optMaxAssertionCount, lvaTrackedCount);
         }
         else
         {

src/coreclr/jit/objectalloc.cpp

@@ -133,7 +133,7 @@ unsigned ObjectAllocator::LocalToIndex(unsigned lclNum)
 {
     unsigned result = BAD_VAR_NUM;
 
-    if (lclNum < comp->lvaCount)
+    if (lclNum < m_firstPseudoLocalNum)
     {
         assert(IsTrackedLocal(lclNum));
         LclVarDsc* const varDsc = comp->lvaGetDesc(lclNum);
@@ -297,6 +297,7 @@ void ObjectAllocator::MarkLclVarAsEscaping(unsigned int lclNum)
 void ObjectAllocator::MarkLclVarAsPossiblyStackPointing(unsigned int lclNum)
 {
     const unsigned bvIndex = LocalToIndex(lclNum);
+    JITDUMP("Marking V%02u (0x%02x) as possibly stack-pointing\n", lclNum, bvIndex);
     BitVecOps::AddElemD(&m_bitVecTraits, m_PossiblyStackPointingPointers, bvIndex);
 }
 
@@ -839,6 +840,11 @@ void ObjectAllocator::ComputeEscapingNodes(BitVecTraits* bitVecTraits, BitVec& e
 
 void ObjectAllocator::ComputeStackObjectPointers(BitVecTraits* bitVecTraits)
 {
+    // Keep track of locals that we know may point at the heap
+    //
+    BitVec possiblyHeapPointingPointers = BitVecOps::MakeEmpty(&m_bitVecTraits);
+    BitVecOps::AddElemD(bitVecTraits, possiblyHeapPointingPointers, m_unknownSourceIndex);
+
     bool     changed = true;
     unsigned pass    = 0;
     while (changed)
@@ -859,23 +865,36 @@ void ObjectAllocator::ComputeStackObjectPointers(BitVecTraits* bitVecTraits)
                                                 m_ConnGraphAdjacencyMatrix[lclIndex]))
             {
                 // We discovered a new pointer that may point to the stack.
+                JITDUMPEXEC(DumpIndex(lclIndex));
+                JITDUMP(" may point to the stack\n");
                 MarkLclVarAsPossiblyStackPointing(lclNum);
+                changed = true;
+            }
 
-                // If all locals assignable to this local are stack pointing, so is this local.
-                //
-                const bool isStackPointing = BitVecOps::IsSubset(bitVecTraits, m_ConnGraphAdjacencyMatrix[lclIndex],
-                                                                 m_DefinitelyStackPointingPointers);
-
-                if (isStackPointing)
-                {
-                    MarkLclVarAsDefinitelyStackPointing(lclNum);
-                    JITDUMP("V%02u is stack pointing\n", lclNum);
-                }
-
+            if (!BitVecOps::IsMember(bitVecTraits, possiblyHeapPointingPointers, lclIndex) &&
+                !BitVecOps::IsEmptyIntersection(bitVecTraits, possiblyHeapPointingPointers,
+                                                m_ConnGraphAdjacencyMatrix[lclIndex]))
+            {
+                // We discovered a new pointer that may point to the heap.
+                JITDUMPEXEC(DumpIndex(lclIndex));
+                JITDUMP(" may point to the heap\n");
+                BitVecOps::AddElemD(bitVecTraits, possiblyHeapPointingPointers, lclIndex);
                 changed = true;
             }
         }
     }
+    JITDUMP("\n---- done computing stack pointing locals\n");
+
+    // If a local is possibly stack pointing and not possibly heap pointing, then it is definitely stack pointing.
+    //
+    BitVec newDefinitelyStackPointingPointers = BitVecOps::UninitVal();
+    BitVecOps::Assign(bitVecTraits, newDefinitelyStackPointingPointers, m_PossiblyStackPointingPointers);
+    BitVecOps::DiffD(bitVecTraits, newDefinitelyStackPointingPointers, possiblyHeapPointingPointers);
+
+    // We should have only added to the set of things that are definitely stack pointing.
+    //
+    assert(BitVecOps::IsSubset(bitVecTraits, m_DefinitelyStackPointingPointers, newDefinitelyStackPointingPointers));
+    BitVecOps::AssignNoCopy(bitVecTraits, m_DefinitelyStackPointingPointers, newDefinitelyStackPointingPointers);
 
 #ifdef DEBUG
     if (comp->verbose)
@@ -1117,10 +1136,10 @@ bool ObjectAllocator::MorphAllocObjNodes()
                 continue;
             }
 
-            bool         canStack     = false;
-            bool         bashCall     = false;
-            const char*  onHeapReason = nullptr;
-            unsigned int lclNum       = stmtExpr->AsLclVar()->GetLclNum();
+            bool               canStack     = false;
+            bool               bashCall     = false;
+            const char*        onHeapReason = nullptr;
+            const unsigned int lclNum       = stmtExpr->AsLclVar()->GetLclNum();
 
             // Don't attempt to do stack allocations inside basic blocks that may be in a loop.
             //
@@ -1306,6 +1325,11 @@ bool ObjectAllocator::MorphAllocObjNodes()
                     stmtExpr->AsLclVar()->Data() = newData;
                     stmtExpr->AddAllEffectsFlags(newData);
                 }
+
+                if (IsTrackedLocal(lclNum))
+                {
+                    AddConnGraphEdge(lclNum, m_unknownSourceLocalNum);
+                }
             }
         }
     }

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

@@ -0,0 +1,48 @@
+// 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.Runtime.CompilerServices;
+using Xunit;
+
+public class X
+{
+    public X() { a = 15; b = 35; }
+    public int a;
+    public int b;
+}
+
+public class ConnectionCycles
+{
+   static bool b;
+
+   [Fact]
+   public static int Problem()
+   {
+       return SubProblem(false) + SubProblem(true);
+   }
+
+   [MethodImpl(MethodImplOptions.NoInlining)]
+   static int SubProblem(bool b)
+   {
+       X x1 = new X();
+       X x2 = new X();
+
+       if (b)
+       {
+           x1 = x2;
+       }
+       else
+       {
+           x2 = x1;
+       }
+
+       SideEffect();
+
+       return x1.a + x2.b;
+   }
+
+
+   [MethodImpl(MethodImplOptions.NoInlining)]
+   static void SideEffect() { }
+}

src/tests/JIT/opt/ObjectStackAllocation/ConnectionCycles.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>