JIT: enhance escape analysis to understand local struct fields

src/coreclr/jit/compiler.h

@@ -1036,6 +1036,13 @@ class LclVarDsc
         m_layout = layout;
     }
 
+    // Change the layout to one that may not be compatible.
+    void ChangeLayout(ClassLayout* layout)
+    {
+        assert(varTypeIsStruct(lvType));
+        m_layout = layout;
+    }
+
     // Grow the size of a block layout local.
     void GrowBlockLayout(ClassLayout* layout)
     {
@@ -10928,8 +10935,10 @@ class Compiler
     unsigned typGetBlkLayoutNum(unsigned blockSize);
     // Get the layout for the specified array of known length
     ClassLayout* typGetArrayLayout(CORINFO_CLASS_HANDLE classHandle, unsigned length);
-    // Get the number of a layout for the specified array of known length
-    unsigned typGetArrayLayoutNum(CORINFO_CLASS_HANDLE classHandle, unsigned length);
+    // Get a layout like an existing layout, with all gc refs removed
+    ClassLayout* typGetNonGCLayout(ClassLayout* existingLayout);
+    // Get a layout like an existing layout, with all gc refs changed to byrefs
+    ClassLayout* typGetByrefLayout(ClassLayout* existingLayout);
 
     var_types TypeHandleToVarType(CORINFO_CLASS_HANDLE handle, ClassLayout** pLayout = nullptr);
     var_types TypeHandleToVarType(CorInfoType jitType, CORINFO_CLASS_HANDLE handle, ClassLayout** pLayout = nullptr);

src/coreclr/jit/jitconfigvalues.h

@@ -678,6 +678,7 @@ RELEASE_CONFIG_INTEGER(JitObjectStackAllocationConditionalEscape, "JitObjectStac
 CONFIG_STRING(JitObjectStackAllocationConditionalEscapeRange, "JitObjectStackAllocationConditionalEscapeRange")
 RELEASE_CONFIG_INTEGER(JitObjectStackAllocationArray, "JitObjectStackAllocationArray", 1)
 RELEASE_CONFIG_INTEGER(JitObjectStackAllocationSize, "JitObjectStackAllocationSize", 528)
+RELEASE_CONFIG_INTEGER(JitObjectStackAllocationTrackFields, "JitObjectStackAllocationTrackFields", 1)
 
 RELEASE_CONFIG_INTEGER(JitEECallTimingInfo, "JitEECallTimingInfo", 0)
 

src/coreclr/jit/layout.cpp

@@ -414,18 +414,80 @@ ClassLayout* Compiler::typGetBlkLayout(unsigned blockSize)
     return typGetCustomLayout(ClassLayoutBuilder(this, blockSize));
 }
 
-unsigned Compiler::typGetArrayLayoutNum(CORINFO_CLASS_HANDLE classHandle, unsigned length)
+ClassLayout* Compiler::typGetArrayLayout(CORINFO_CLASS_HANDLE classHandle, unsigned length)
 {
     ClassLayoutBuilder b = ClassLayoutBuilder::BuildArray(this, classHandle, length);
-    return typGetCustomLayoutNum(b);
+    return typGetCustomLayout(b);
 }
 
-ClassLayout* Compiler::typGetArrayLayout(CORINFO_CLASS_HANDLE classHandle, unsigned length)
+ClassLayout* Compiler::typGetNonGCLayout(ClassLayout* layout)
 {
-    ClassLayoutBuilder b = ClassLayoutBuilder::BuildArray(this, classHandle, length);
+    assert(layout->HasGCPtr());
+    ClassLayoutBuilder b(this, layout->GetSize());
+    b.CopyPaddingFrom(0, layout);
+
+#ifdef DEBUG
+    b.CopyNameFrom(layout, "[nongc] ");
+#endif
+
+    return typGetCustomLayout(b);
+}
+
+ClassLayout* Compiler::typGetByrefLayout(ClassLayout* layout)
+{
+    assert(layout->HasGCPtr());
+    ClassLayoutBuilder b(this, layout->GetSize());
+    b.CopyPaddingFrom(0, layout);
+    b.CopyGCInfoFromMakeByref(0, layout);
+
+#ifdef DEBUG
+    b.CopyNameFrom(layout, "[byref] ");
+#endif
+
     return typGetCustomLayout(b);
 }
 
+#ifdef DEBUG
+//------------------------------------------------------------------------
+// CopyNameFrom: Copy layout names, with optional prefix.
+//
+// Parameters:
+//   layout      - layout to copy from
+//   prefix      - prefix to add (or nullptr)
+//
+void ClassLayoutBuilder::CopyNameFrom(ClassLayout* layout, const char* prefix)
+{
+    const char* layoutName      = layout->GetClassName();
+    const char* layoutShortName = layout->GetShortClassName();
+
+    if (prefix != nullptr)
+    {
+        char* newName      = nullptr;
+        char* newShortName = nullptr;
+
+        if (layoutName != nullptr)
+        {
+            size_t len = strlen(prefix) + strlen(layoutName) + 1;
+            newName    = m_compiler->getAllocator(CMK_DebugOnly).allocate<char>(len);
+            sprintf_s(newName, len, "%s%s", prefix, layoutShortName);
+        }
+
+        if (layoutShortName != nullptr)
+        {
+            size_t len   = strlen(prefix) + strlen(layoutName) + 1;
+            newShortName = m_compiler->getAllocator(CMK_DebugOnly).allocate<char>(len);
+            sprintf_s(newShortName, len, "%s%s", prefix, layoutShortName);
+        }
+
+        SetName(newName, newShortName);
+    }
+    else
+    {
+        SetName(layoutName, layoutShortName);
+    }
+}
+#endif // DEBUG
+
 //------------------------------------------------------------------------
 // Create: Create a ClassLayout from an EE side class handle.
 //
@@ -646,8 +708,8 @@ const SegmentList& ClassLayout::GetNonPadding(Compiler* comp)
 // AreCompatible: check if 2 layouts are the same for copying.
 //
 // Arguments:
-//    layout1 - the first layout;
-//    layout2 - the second layout.
+//    layout1 - the first layout (copy destination)
+//    layout2 - the second layout (copy source)
 //
 // Return value:
 //    true if compatible, false otherwise.
@@ -658,6 +720,8 @@ const SegmentList& ClassLayout::GetNonPadding(Compiler* comp)
 //
 //    This is an equivalence relation:
 //      AreCompatible(a, b) == AreCompatible(b, a)
+//      AreCompatible(a, a) == true
+//      AreCompatible(a, b) && AreCompatible(b, c) ==> AreCompatible(a, c)
 //
 // static
 bool ClassLayout::AreCompatible(const ClassLayout* layout1, const ClassLayout* layout2)
@@ -746,9 +810,92 @@ bool ClassLayout::AreCompatible(const ClassLayout* layout1, const ClassLayout* l
 //
 bool ClassLayout::CanAssignFrom(const ClassLayout* layout)
 {
-    // Currently this is the same as compatability
+    if (this == layout)
+    {
+        return true;
+    }
+
+    // Do the normal compatibility check first, when possible to do so.
+    //
+    if ((IsCustomLayout() == layout->IsCustomLayout()) || (!HasGCPtr() && !layout->HasGCPtr()))
+    {
+        const bool areCompatible = AreCompatible(this, layout);
+
+        if (areCompatible)
+        {
+            return true;
+        }
+    }
+
+    // Must be same size
+    //
+    if (GetSize() != layout->GetSize())
+    {
+        return false;
+    }
+
+    // Must be same IR type
+    //
+    if (GetType() != layout->GetType())
+    {
+        return false;
+    }
+
+    // Dest is GC, source is GC. Allow, slotwise:
     //
-    return AreCompatible(this, layout);
+    //   byref <- ref, byref, nint
+    //   ref   <- ref
+    //   nint  <- nint
+    //
+    if (HasGCPtr() && layout->HasGCPtr())
+    {
+        const unsigned slotsCount = GetSlotCount();
+        assert(slotsCount == layout->GetSlotCount());
+
+        for (unsigned i = 0; i < slotsCount; ++i)
+        {
+            var_types slotType       = GetGCPtrType(i);
+            var_types layoutSlotType = layout->GetGCPtrType(i);
+
+            if ((slotType != TYP_BYREF) && (slotType != layoutSlotType))
+            {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    // Dest is GC, source is noGC. Allow, slotwise:
+    //
+    //    byref <- nint
+    //    nint  <- nint
+    //
+    if (HasGCPtr() && !layout->HasGCPtr())
+    {
+        const unsigned slotsCount = GetSlotCount();
+
+        for (unsigned i = 0; i < slotsCount; ++i)
+        {
+            var_types slotType = GetGCPtrType(i);
+            if (slotType == TYP_REF)
+            {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    // Dest is noGC, source is GC. Disallow.
+    //
+    if (!HasGCPtr() && layout->HasGCPtr())
+    {
+        assert(!HasGCPtr());
+        return false;
+    }
+
+    // Dest is noGC, source is noGC, and they're not compatible.
+    //
+    return false;
 }
 
 //------------------------------------------------------------------------
@@ -814,7 +961,7 @@ ClassLayoutBuilder ClassLayoutBuilder::BuildArray(Compiler* compiler, CORINFO_CL
             unsigned offset = OFFSETOF__CORINFO_Array__data;
             for (unsigned i = 0; i < length; i++)
             {
-                builder.CopyInfoFrom(offset, elementLayout, /* copy padding */ false);
+                builder.CopyGCInfoFrom(offset, elementLayout);
                 offset += elementSize;
             }
         }
@@ -919,14 +1066,13 @@ void ClassLayoutBuilder::SetGCPtrType(unsigned slot, var_types type)
 }
 
 //------------------------------------------------------------------------
-// CopyInfoFrom: Copy GC pointers and padding information from another layout.
+// CopyInfoGCFrom: Copy GC pointers from another layout.
 //
 // Arguments:
 //   offset      - Offset in this builder to start copy information into.
 //   layout      - Layout to get information from.
-//   copyPadding - Whether padding info should also be copied from the layout.
 //
-void ClassLayoutBuilder::CopyInfoFrom(unsigned offset, ClassLayout* layout, bool copyPadding)
+void ClassLayoutBuilder::CopyGCInfoFrom(unsigned offset, ClassLayout* layout)
 {
     assert(offset + layout->GetSize() <= m_size);
 
@@ -939,18 +1085,53 @@ void ClassLayoutBuilder::CopyInfoFrom(unsigned offset, ClassLayout* layout, bool
             SetGCPtr(startSlot + slot, layout->GetGCPtr(slot));
         }
     }
+}
 
-    if (copyPadding)
-    {
-        AddPadding(SegmentList::Segment(offset, offset + layout->GetSize()));
+//------------------------------------------------------------------------
+// CopyInfoGCFromMakeByref: Copy GC pointers from another layout,and change
+//   all gc references to be TYP_BYREF (TYPE_GC_BYREF)
+//
+// Arguments:
+//   offset      - Offset in this builder to start copy information into.
+//   layout      - Layout to get information from.
+//
+void ClassLayoutBuilder::CopyGCInfoFromMakeByref(unsigned offset, ClassLayout* layout)
+{
+    assert(offset + layout->GetSize() <= m_size);
 
-        for (const SegmentList::Segment& nonPadding : layout->GetNonPadding(m_compiler))
+    if (layout->GetGCPtrCount() > 0)
+    {
+        assert(offset % TARGET_POINTER_SIZE == 0);
+        unsigned startSlot = offset / TARGET_POINTER_SIZE;
+        for (unsigned slot = 0; slot < layout->GetSlotCount(); slot++)
         {
-            RemovePadding(SegmentList::Segment(offset + nonPadding.Start, offset + nonPadding.End));
+            CorInfoGCType gcType = layout->GetGCPtr(slot);
+            if (gcType == TYPE_GC_REF)
+            {
+                gcType = TYPE_GC_BYREF;
+            }
+            SetGCPtr(startSlot + slot, gcType);
         }
     }
 }
 
+//------------------------------------------------------------------------
+// CopyInfoPaddingFrom: Copy padding from another layout.
+//
+// Arguments:
+//   offset      - Offset in this builder to start copy information into.
+//   layout      - Layout to get information from.
+//
+void ClassLayoutBuilder::CopyPaddingFrom(unsigned offset, ClassLayout* layout)
+{
+    AddPadding(SegmentList::Segment(offset, offset + layout->GetSize()));
+
+    for (const SegmentList::Segment& nonPadding : layout->GetNonPadding(m_compiler))
+    {
+        RemovePadding(SegmentList::Segment(offset + nonPadding.Start, offset + nonPadding.End));
+    }
+}
+
 //------------------------------------------------------------------------
 // GetOrCreateNonPadding: Get the non padding segment list, or create it if it
 // does not exist.

src/coreclr/jit/layout.h

@@ -34,12 +34,15 @@ class ClassLayoutBuilder
     ClassLayoutBuilder(Compiler* compiler, unsigned size);
 
     void SetGCPtrType(unsigned slot, var_types type);
-    void CopyInfoFrom(unsigned offset, ClassLayout* layout, bool copyPadding);
+    void CopyGCInfoFrom(unsigned offset, ClassLayout* layout);
+    void CopyGCInfoFromMakeByref(unsigned offset, ClassLayout* layout);
+    void CopyPaddingFrom(unsigned offset, ClassLayout* layout);
     void AddPadding(const SegmentList::Segment& padding);
     void RemovePadding(const SegmentList::Segment& nonPadding);
 
 #ifdef DEBUG
     void SetName(const char* name, const char* shortName);
+    void CopyNameFrom(ClassLayout* layout, const char* prefix);
 #endif
 
     static ClassLayoutBuilder BuildArray(Compiler* compiler, CORINFO_CLASS_HANDLE arrayType, unsigned length);