Stack allocate unescaped boxes (#103361)

Enable object stack allocation for ref classes and extend the support to include boxed value classes. Use a specialized unbox helper for stack allocated boxes, both to avoid apparent escape of the box by the helper, and to ensure all box field accesses are visible to the JIT. Update the local address visitor to rewrite trees involving address of stack allocated boxes in some cases to avoid address exposure. Disable old promotion for stack allocated boxes (since we have no field handles) and allow physical promotion to enregister the box method table and/or payload as appropriate. In OSR methods handle the fact that the stack allocation may actually have been a heap allocation by the Tier0 method.

The analysis TP cost is around 0.4-0.7% (notes below). Boxes are much less likely to escape than ref classes (roughly ~90% of boxes escape, ~99.8% of ref classes escape). Codegen impact is diminished somewhat because many of the boxes are dead and were already getting optimized away.
 
Fixes #4584, #9118, #10195, #11192, #53585, #58554, #85570

---------

Co-authored-by: Jakob Botsch Nielsen <jakob.botsch.nielsen@gmail.com>
Co-authored-by: Jan Kotas <jkotas@microsoft.com>

src/coreclr/System.Private.CoreLib/System.Private.CoreLib.csproj

@@ -197,6 +197,7 @@
     <Compile Include="$(BclSourcesRoot)\System\Runtime\CompilerServices\CastHelpers.cs" />
     <Compile Include="$(BclSourcesRoot)\System\Runtime\CompilerServices\ICastableHelpers.cs" />
     <Compile Include="$(BclSourcesRoot)\System\Runtime\CompilerServices\RuntimeHelpers.CoreCLR.cs" />
+    <Compile Include="$(BclSourcesRoot)\System\Runtime\CompilerServices\StackAllocatedBox.cs" />
     <Compile Include="$(BclSourcesRoot)\System\Runtime\ControlledExecution.CoreCLR.cs" />
     <Compile Include="$(BclSourcesRoot)\System\Runtime\DependentHandle.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)..\nativeaot\Common\src\System\Runtime\RhFailFastReason.cs" />

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/StackAllocatedBox.cs

@@ -0,0 +1,15 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Runtime.InteropServices;
+
+namespace System.Runtime.CompilerServices
+{
+    [StructLayout(LayoutKind.Sequential)]
+    internal unsafe struct StackAllocatedBox<T>
+    {
+        // These fields are only accessed from jitted code
+        private MethodTable* _pMethodTable;
+        private T _value;
+    }
+}

src/coreclr/inc/corinfo.h

@@ -445,6 +445,7 @@ enum CorInfoHelpFunc
     CORINFO_HELP_BOX,               // Fast box helper. Only possible exception is OutOfMemory
     CORINFO_HELP_BOX_NULLABLE,      // special form of boxing for Nullable<T>
     CORINFO_HELP_UNBOX,
+    CORINFO_HELP_UNBOX_TYPETEST,    // Verify unbox type, throws if incompatible
     CORINFO_HELP_UNBOX_NULLABLE,    // special form of unboxing for Nullable<T>
     CORINFO_HELP_GETREFANY,         // Extract the byref from a TypedReference, checking that it is the expected type
 
@@ -2552,6 +2553,14 @@ class ICorStaticInfo
             CORINFO_CLASS_HANDLE        cls
             ) = 0;
 
+    // Get a representation for a stack-allocated boxed value type.
+    //
+    // This differs from getTypeForBox in that it includes an explicit field
+    // for the method table pointer.
+    virtual CORINFO_CLASS_HANDLE getTypeForBoxOnStack(
+            CORINFO_CLASS_HANDLE        cls
+            ) = 0;
+
     // returns the correct box helper for a particular class.  Note
     // that if this returns CORINFO_HELP_BOX, the JIT can assume
     // 'standard' boxing (allocate object and copy), and optimize

src/coreclr/inc/icorjitinfoimpl_generated.h

@@ -275,6 +275,9 @@ CorInfoHelpFunc getSharedCCtorHelper(
 CORINFO_CLASS_HANDLE getTypeForBox(
           CORINFO_CLASS_HANDLE cls) override;
 
+CORINFO_CLASS_HANDLE getTypeForBoxOnStack(
+          CORINFO_CLASS_HANDLE cls) override;
+
 CorInfoHelpFunc getBoxHelper(
           CORINFO_CLASS_HANDLE cls) override;
 

src/coreclr/inc/jiteeversionguid.h

@@ -43,11 +43,11 @@ typedef const GUID *LPCGUID;
 #define GUID_DEFINED
 #endif // !GUID_DEFINED
 
-constexpr GUID JITEEVersionIdentifier = { /* e428e66d-5e0e-4320-ad8a-fa5a50f6da07 */
-    0xe428e66d,
-    0x5e0e,
-    0x4320,
-    {0xad, 0x8a, 0xfa, 0x5a, 0x50, 0xf6, 0xda, 0x07}
+constexpr GUID JITEEVersionIdentifier = { /* 748cd07e-c686-4085-8f5f-54c1b9cff44c */
+    0x748cd07e,
+    0xc686,
+    0x4085,
+    {0x8f, 0x5f, 0x54, 0xc1, 0xb9, 0xcf, 0xf4, 0x4c}
   };
 
 //////////////////////////////////////////////////////////////////////////////////////////////////////////

src/coreclr/inc/jithelpers.h

@@ -106,6 +106,7 @@
     DYNAMICJITHELPER(CORINFO_HELP_BOX,          JIT_Box,            CORINFO_HELP_SIG_REG_ONLY)
     JITHELPER(CORINFO_HELP_BOX_NULLABLE,        JIT_Box,            CORINFO_HELP_SIG_REG_ONLY)
     DYNAMICJITHELPER(CORINFO_HELP_UNBOX,        NULL,               CORINFO_HELP_SIG_REG_ONLY)
+    JITHELPER(CORINFO_HELP_UNBOX_TYPETEST,      JIT_Unbox_TypeTest, CORINFO_HELP_SIG_REG_ONLY)
     JITHELPER(CORINFO_HELP_UNBOX_NULLABLE,      JIT_Unbox_Nullable, CORINFO_HELP_SIG_4_STACK)
 
     JITHELPER(CORINFO_HELP_GETREFANY,           JIT_GetRefAny,      CORINFO_HELP_SIG_8_STACK)

src/coreclr/jit/ICorJitInfo_names_generated.h

@@ -68,6 +68,7 @@ DEF_CLR_API(getNewArrHelper)
 DEF_CLR_API(getCastingHelper)
 DEF_CLR_API(getSharedCCtorHelper)
 DEF_CLR_API(getTypeForBox)
+DEF_CLR_API(getTypeForBoxOnStack)
 DEF_CLR_API(getBoxHelper)
 DEF_CLR_API(getUnBoxHelper)
 DEF_CLR_API(getRuntimeTypePointer)

src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp

@@ -635,6 +635,15 @@ CORINFO_CLASS_HANDLE WrapICorJitInfo::getTypeForBox(
     return temp;
 }
 
+CORINFO_CLASS_HANDLE WrapICorJitInfo::getTypeForBoxOnStack(
+          CORINFO_CLASS_HANDLE cls)
+{
+    API_ENTER(getTypeForBoxOnStack);
+    CORINFO_CLASS_HANDLE temp = wrapHnd->getTypeForBoxOnStack(cls);
+    API_LEAVE(getTypeForBoxOnStack);
+    return temp;
+}
+
 CorInfoHelpFunc WrapICorJitInfo::getBoxHelper(
           CORINFO_CLASS_HANDLE cls)
 {

src/coreclr/jit/compiler.h

@@ -690,7 +690,8 @@ class LclVarDsc
     unsigned char lvSingleDefDisqualifyReason = 'H';
 #endif
 
-    unsigned char lvAllDefsAreNoGc : 1; // For pinned locals: true if all defs of this local are no-gc
+    unsigned char lvAllDefsAreNoGc    : 1; // For pinned locals: true if all defs of this local are no-gc
+    unsigned char lvStackAllocatedBox : 1; // Local is a stack allocated box
 
 #if FEATURE_MULTIREG_ARGS
     regNumber lvRegNumForSlot(unsigned slotNum)
@@ -805,6 +806,11 @@ class LclVarDsc
         return lvIsMultiRegArg || lvIsMultiRegRet;
     }
 
+    bool IsStackAllocatedBox() const
+    {
+        return lvStackAllocatedBox;
+    }
+
 #if defined(DEBUG)
 private:
     DoNotEnregisterReason m_doNotEnregReason;
@@ -3523,7 +3529,7 @@ class Compiler
 
     GenTree* gtNewRuntimeLookup(CORINFO_GENERIC_HANDLE hnd, CorInfoGenericHandleType hndTyp, GenTree* lookupTree);
 
-    GenTreeIndir* gtNewMethodTableLookup(GenTree* obj);
+    GenTreeIndir* gtNewMethodTableLookup(GenTree* obj, bool onStack = false);
 
     //------------------------------------------------------------------------
     // Other GenTree functions

src/coreclr/jit/compiler.hpp

@@ -1872,9 +1872,9 @@ inline GenTreeCast* Compiler::gtNewCastNodeL(var_types typ, GenTree* op1, bool f
     return cast;
 }
 
-inline GenTreeIndir* Compiler::gtNewMethodTableLookup(GenTree* object)
+inline GenTreeIndir* Compiler::gtNewMethodTableLookup(GenTree* object, bool onStack)
 {
-    assert(object->TypeIs(TYP_REF));
+    assert(onStack || object->TypeIs(TYP_REF));
     GenTreeIndir* result = gtNewIndir(TYP_I_IMPL, object, GTF_IND_INVARIANT);
     return result;
 }

src/coreclr/jit/inlinepolicy.cpp

@@ -1628,7 +1628,7 @@ double ExtendedDefaultPolicy::DetermineMultiplier()
         //
         //  void Caller() => DoNothing(42); // 42 is going to be boxed at the call site.
         //
-        multiplier += 0.5;
+        multiplier += 0.5 * m_ArgIsBoxedAtCallsite;
         JITDUMP("\nCallsite is going to box %d arguments.  Multiplier increased to %g.", m_ArgIsBoxedAtCallsite,
                 multiplier);
     }

src/coreclr/jit/jitconfigvalues.h

@@ -646,7 +646,10 @@ RELEASE_CONFIG_INTEGER(JitExtDefaultPolicyProfScale, W("JitExtDefaultPolicyProfS
 RELEASE_CONFIG_INTEGER(JitInlinePolicyModel, W("JitInlinePolicyModel"), 0)
 RELEASE_CONFIG_INTEGER(JitInlinePolicyProfile, W("JitInlinePolicyProfile"), 0)
 RELEASE_CONFIG_INTEGER(JitInlinePolicyProfileThreshold, W("JitInlinePolicyProfileThreshold"), 40)
-RELEASE_CONFIG_INTEGER(JitObjectStackAllocation, W("JitObjectStackAllocation"), 0)
+CONFIG_STRING(JitObjectStackAllocationRange, W("JitObjectStackAllocationRange"))
+RELEASE_CONFIG_INTEGER(JitObjectStackAllocation, W("JitObjectStackAllocation"), 1)
+RELEASE_CONFIG_INTEGER(JitObjectStackAllocationRefClass, W("JitObjectStackAllocationRefClass"), 1)
+RELEASE_CONFIG_INTEGER(JitObjectStackAllocationBoxedValueClass, W("JitObjectStackAllocationBoxedValueClass"), 1)
 
 RELEASE_CONFIG_INTEGER(JitEECallTimingInfo, W("JitEECallTimingInfo"), 0)
 

src/coreclr/jit/jitmetadatalist.h

@@ -71,6 +71,10 @@ JITMETADATAMETRIC(ProfileInconsistentInlineeScale,       int,              0)
 JITMETADATAMETRIC(ProfileInconsistentInlinee,            int,              0)
 JITMETADATAMETRIC(ProfileInconsistentNoReturnInlinee,    int,              0)
 JITMETADATAMETRIC(ProfileInconsistentMayThrowInlinee,    int,              0)
+JITMETADATAMETRIC(NewRefClassHelperCalls,                int,              0)
+JITMETADATAMETRIC(StackAllocatedRefClasses,              int,              0)
+JITMETADATAMETRIC(NewBoxedValueClassHelperCalls,         int,              0)
+JITMETADATAMETRIC(StackAllocatedBoxedValueClasses,       int,              0)
 
 #undef JITMETADATA
 #undef JITMETADATAINFO

src/coreclr/jit/lclmorph.cpp

@@ -340,6 +340,7 @@ class LocalEqualsLocalAddrAssertions
     //
     void OnExposed(unsigned lclNum)
     {
+        JITDUMP("On exposed: V%02u\n", lclNum);
         BitVecTraits localsTraits(m_comp->lvaCount, m_comp);
         BitVecOps::AddElemD(&localsTraits, m_localsToExpose, lclNum);
     }
@@ -515,6 +516,13 @@ class LocalAddressVisitor final : public GenTreeVisitor<LocalAddressVisitor>
             return m_offset;
         }
 
+        bool IsSameAddress(const Value& other) const
+        {
+            assert(IsAddress() && other.IsAddress());
+
+            return ((LclNum() == other.LclNum()) && (Offset() == other.Offset()));
+        }
+
         //------------------------------------------------------------------------
         // Address: Produce an address value from a LCL_ADDR node.
         //
@@ -806,7 +814,6 @@ class LocalAddressVisitor final : public GenTreeVisitor<LocalAddressVisitor>
         }
 
         PushValue(use);
-
         return Compiler::WALK_CONTINUE;
     }
 
@@ -1028,6 +1035,77 @@ class LocalAddressVisitor final : public GenTreeVisitor<LocalAddressVisitor>
                 SequenceCall(node->AsCall());
                 break;
 
+            case GT_EQ:
+            case GT_NE:
+            {
+                // If we see &lcl EQ/NE null, rewrite to 0/1 comparison
+                // to reduce overall address exposure.
+                //
+                assert(TopValue(2).Node() == node);
+                assert(TopValue(1).Node() == node->AsOp()->gtOp1);
+                assert(TopValue(0).Node() == node->AsOp()->gtOp2);
+
+                Value& lhs = TopValue(1);
+                Value& rhs = TopValue(0);
+
+                if ((lhs.IsAddress() && rhs.Node()->IsIntegralConst(0)) ||
+                    (rhs.IsAddress() && lhs.Node()->IsIntegralConst(0)))
+                {
+                    JITDUMP("Rewriting known address vs null comparison [%06u]\n", m_compiler->dspTreeID(node));
+                    *lhs.Use()     = m_compiler->gtNewIconNode(0);
+                    *rhs.Use()     = m_compiler->gtNewIconNode(1);
+                    m_stmtModified = true;
+
+                    INDEBUG(TopValue(0).Consume());
+                    INDEBUG(TopValue(1).Consume());
+                    PopValue();
+                    PopValue();
+                }
+                else if (lhs.IsAddress() && rhs.IsAddress())
+                {
+                    JITDUMP("Rewriting known address vs address comparison [%06u]\n", m_compiler->dspTreeID(node));
+                    bool isSameAddress = lhs.IsSameAddress(rhs);
+                    *lhs.Use()         = m_compiler->gtNewIconNode(0);
+                    *rhs.Use()         = m_compiler->gtNewIconNode(isSameAddress ? 0 : 1);
+                    m_stmtModified     = true;
+
+                    INDEBUG(TopValue(0).Consume());
+                    INDEBUG(TopValue(1).Consume());
+                    PopValue();
+                    PopValue();
+                }
+                else
+                {
+                    EscapeValue(TopValue(0), node);
+                    PopValue();
+                    EscapeValue(TopValue(0), node);
+                    PopValue();
+                }
+
+                break;
+            }
+
+            case GT_NULLCHECK:
+            {
+                assert(TopValue(1).Node() == node);
+                assert(TopValue(0).Node() == node->AsOp()->gtOp1);
+                Value& op = TopValue(0);
+                if (op.IsAddress())
+                {
+                    JITDUMP("Bashing nullcheck of local [%06u] to NOP\n", m_compiler->dspTreeID(node));
+                    node->gtBashToNOP();
+                    INDEBUG(TopValue(0).Consume());
+                    PopValue();
+                    m_stmtModified = true;
+                }
+                else
+                {
+                    EscapeValue(TopValue(0), node);
+                    PopValue();
+                }
+                break;
+            }
+
             default:
                 while (TopValue(0).Node() != node)
                 {

src/coreclr/jit/lclvars.cpp

@@ -2572,6 +2572,13 @@ bool Compiler::StructPromotionHelper::CanPromoteStructVar(unsigned lclNum)
 
     if (varDsc->GetLayout()->IsBlockLayout())
     {
+        JITDUMP("  struct promotion of V%02u is disabled because it has block layout\n", lclNum);
+        return false;
+    }
+
+    if (varDsc->lvStackAllocatedBox)
+    {
+        JITDUMP("  struct promotion of V%02u is disabled because it is a stack allocated box\n", lclNum);
         return false;
     }
 
@@ -3359,6 +3366,8 @@ bool Compiler::lvaIsLocalImplicitlyAccessedByRef(unsigned lclNum) const
 // this information is already available on the CallArgABIInformation, and shouldn't need to be
 // recomputed.
 //
+// Also seems like this info could be cached in the layout.
+//
 bool Compiler::lvaIsMultiregStruct(LclVarDsc* varDsc, bool isVarArg)
 {
     if (varTypeIsStruct(varDsc->TypeGet()))