JIT: Extend escape analysis to account for arrays with non-gcref elements

src/coreclr/jit/jitmetadatalist.h

@@ -76,6 +76,7 @@ JITMETADATAMETRIC(NewRefClassHelperCalls,                int,              0)
 JITMETADATAMETRIC(StackAllocatedRefClasses,              int,              0)
 JITMETADATAMETRIC(NewBoxedValueClassHelperCalls,         int,              0)
 JITMETADATAMETRIC(StackAllocatedBoxedValueClasses,       int,              0)
+JITMETADATAMETRIC(StackAllocatedArrays,                  int,              0)
 
 #undef JITMETADATA
 #undef JITMETADATAINFO

src/coreclr/jit/objectalloc.cpp

@@ -30,9 +30,9 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 //
 PhaseStatus ObjectAllocator::DoPhase()
 {
-    if ((comp->optMethodFlags & OMF_HAS_NEWOBJ) == 0)
+    if ((comp->optMethodFlags & OMF_HAS_NEWOBJ) == 0 && (comp->optMethodFlags & OMF_HAS_NEWARRAY) == 0)
     {
-        JITDUMP("no newobjs in this method; punting\n");
+        JITDUMP("no newobjs or newarr in this method; punting\n");
         return PhaseStatus::MODIFIED_NOTHING;
     }
 
@@ -89,6 +89,27 @@ void ObjectAllocator::MarkLclVarAsEscaping(unsigned int lclNum)
     BitVecOps::AddElemD(&m_bitVecTraits, m_EscapingPointers, lclNum);
 }
 
+//------------------------------------------------------------------------------
+// MarkLclVarHasLocalStore : Mark local variable having local store.
+//
+//
+// Arguments:
+//    lclNum  - Pointing local variable number
+// Returns:
+//    true if the local variable is first marked as having local store
+//    false if the local variable is already marked as having local store
+
+bool ObjectAllocator::MarkLclVarHasLocalStore(unsigned int lclNum)
+{
+    if (BitVecOps::IsMember(&m_bitVecTraits, m_PointersHasLocalStore, lclNum))
+    {
+        return false;
+    }
+
+    BitVecOps::AddElemD(&m_bitVecTraits, m_PointersHasLocalStore, lclNum);
+    return true;
+}
+
 //------------------------------------------------------------------------------
 // MarkLclVarAsPossiblyStackPointing : Mark local variable as possibly pointing
 //                                     to a stack-allocated object.
@@ -140,6 +161,7 @@ void ObjectAllocator::DoAnalysis()
     if (comp->lvaCount > 0)
     {
         m_EscapingPointers         = BitVecOps::MakeEmpty(&m_bitVecTraits);
+        m_PointersHasLocalStore    = BitVecOps::MakeEmpty(&m_bitVecTraits);
         m_ConnGraphAdjacencyMatrix = new (comp->getAllocator(CMK_ObjectAllocator)) BitSetShortLongRep[comp->lvaCount];
 
         MarkEscapingVarsAndBuildConnGraph();
@@ -201,7 +223,9 @@ void ObjectAllocator::MarkEscapingVarsAndBuildConnGraph()
 
             if (tree->OperIsLocalStore())
             {
-                lclEscapes = false;
+                // conservatively marking locals being reassigned as escaping
+                // this can happen on arrays
+                lclEscapes = !m_allocator->MarkLclVarHasLocalStore(lclNum);
             }
             else if (tree->OperIs(GT_LCL_VAR) && tree->TypeIs(TYP_REF, TYP_BYREF, TYP_I_IMPL))
             {
@@ -391,6 +415,7 @@ bool ObjectAllocator::MorphAllocObjNodes()
             GenTree* data     = nullptr;
 
             bool canonicalAllocObjFound = false;
+            bool canonicalAllocArrFound = false;
 
             if (stmtExpr->OperIs(GT_STORE_LCL_VAR) && stmtExpr->TypeIs(TYP_REF))
             {
@@ -400,6 +425,81 @@ bool ObjectAllocator::MorphAllocObjNodes()
                 {
                     canonicalAllocObjFound = true;
                 }
+                else if (data->IsHelperCall())
+                {
+                    GenTreeCall* call = data->AsCall();
+                    if (call->GetHelperNum() == CORINFO_HELP_NEWARR_1_VC &&
+                        call->gtArgs.GetArgByIndex(0)->GetNode()->IsCnsIntOrI())
+                    {
+                        canonicalAllocArrFound = true;
+                    }
+                }
+            }
+
+            if (canonicalAllocArrFound)
+            {
+                GenTreeCall* asCall = data->AsCall();
+                assert(asCall->GetHelperNum() == CORINFO_HELP_NEWARR_1_VC);
+
+                unsigned int         lclNum       = stmtExpr->AsLclVar()->GetLclNum();
+                CallArg*             arg          = asCall->gtArgs.GetArgByIndex(0);
+                GenTree*             node         = arg->GetNode();
+                CORINFO_CLASS_HANDLE clsHnd       = (CORINFO_CLASS_HANDLE)node->AsIntConCommon()->IntegralValue();
+                GenTree*             len          = arg->GetNext()->GetNode();
+                const char*          onHeapReason = nullptr;
+                unsigned int         blockSize    = 0;
+                bool                 canStack     = false;
+
+                // Don't attempt to do stack allocations inside basic blocks that may be in a loop.
+                //
+                if (!IsObjectStackAllocationEnabled())
+                {
+                    onHeapReason = "[object stack allocation disabled]";
+                    canStack     = false;
+                }
+                else if (basicBlockHasBackwardJump)
+                {
+                    onHeapReason = "[alloc in loop]";
+                    canStack     = false;
+                }
+                else if (!len->IsCnsIntOrI())
+                {
+                    onHeapReason = "[non-constant size]";
+                    canStack     = false;
+                }
+                else if (!CanAllocateLclVarOnStack(lclNum, clsHnd, (unsigned int)len->AsIntCon()->gtIconVal, &blockSize,
+                                                   &onHeapReason))
+                {
+                    // reason set by the call
+                    canStack = false;
+                }
+                else
+                {
+                    JITDUMP("Allocating V%02u on the stack\n", lclNum);
+                    canStack = true;
+                    const unsigned int stackLclNum =
+                        MorphNewArrNodeIntoStackAlloc(asCall, clsHnd, (unsigned int)len->AsIntCon()->gtIconVal,
+                                                      blockSize, block, stmt);
+                    m_HeapLocalToStackLocalMap.AddOrUpdate(lclNum, stackLclNum);
+                    m_LocalArrToLenMap.AddOrUpdate(stackLclNum, (unsigned int)len->AsIntCon()->gtIconVal);
+                    // We keep the set of possibly-stack-pointing pointers as a superset of the set of
+                    // definitely-stack-pointing pointers. All definitely-stack-pointing pointers are in both sets.
+                    MarkLclVarAsDefinitelyStackPointing(lclNum);
+                    MarkLclVarAsPossiblyStackPointing(lclNum);
+                    stmt->GetRootNode()->gtBashToNOP();
+                    comp->optMethodFlags |= OMF_HAS_OBJSTACKALLOC;
+                    didStackAllocate = true;
+                }
+
+                if (canStack)
+                {
+                    comp->Metrics.StackAllocatedArrays++;
+                }
+                else
+                {
+                    assert(onHeapReason != nullptr);
+                    JITDUMP("Allocating V%02u on the heap: %s\n", lclNum, onHeapReason);
+                }
             }
 
             if (canonicalAllocObjFound)
@@ -443,7 +543,7 @@ bool ObjectAllocator::MorphAllocObjNodes()
                     onHeapReason = "[alloc in loop]";
                     canStack     = false;
                 }
-                else if (!CanAllocateLclVarOnStack(lclNum, clsHnd, &onHeapReason))
+                else if (!CanAllocateLclVarOnStack(lclNum, clsHnd, 0, NULL, &onHeapReason))
                 {
                     // reason set by the call
                     canStack = false;
@@ -555,6 +655,78 @@ GenTree* ObjectAllocator::MorphAllocObjNodeIntoHelperCall(GenTreeAllocObj* alloc
     return helperCall;
 }
 
+unsigned int ObjectAllocator::MorphNewArrNodeIntoStackAlloc(GenTreeCall*         newArr,
+                                                            CORINFO_CLASS_HANDLE clsHnd,
+                                                            unsigned int         length,
+                                                            unsigned int         blockSize,
+                                                            BasicBlock*          block,
+                                                            Statement*           stmt)
+{
+    assert(newArr != nullptr);
+    assert(m_AnalysisDone);
+    assert(clsHnd != NO_CLASS_HANDLE);
+
+    const bool         shortLifetime = false;
+    const unsigned int lclNum        = comp->lvaGrabTemp(shortLifetime DEBUGARG("stack allocated array temp"));
+
+    comp->lvaSetStruct(lclNum, comp->typGetBlkLayout(blockSize), /* unsafeValueClsCheck */ false);
+
+    // Initialize the object memory if necessary.
+    bool             bbInALoop  = block->HasFlag(BBF_BACKWARD_JUMP);
+    bool             bbIsReturn = block->KindIs(BBJ_RETURN);
+    LclVarDsc* const lclDsc     = comp->lvaGetDesc(lclNum);
+    lclDsc->lvStackAllocatedBox = false;
+    if (comp->fgVarNeedsExplicitZeroInit(lclNum, bbInALoop, bbIsReturn))
+    {
+        //------------------------------------------------------------------------
+        // STMTx (IL 0x... ???)
+        //   *  STORE_LCL_VAR   struct
+        //   \--*  CNS_INT   int    0
+        //------------------------------------------------------------------------
+
+        GenTree*   init     = comp->gtNewStoreLclVarNode(lclNum, comp->gtNewIconNode(0));
+        Statement* initStmt = comp->gtNewStmt(init);
+
+        comp->fgInsertStmtBefore(block, stmt, initStmt);
+    }
+    else
+    {
+        JITDUMP("\nSuppressing zero-init for V%02u -- expect to zero in prolog\n", lclNum);
+        lclDsc->lvSuppressedZeroInit = 1;
+        comp->compSuppressedZeroInit = true;
+    }
+
+    // Initialize the vtable slot.
+    //
+    //------------------------------------------------------------------------
+    // STMTx (IL 0x... ???)
+    //   * STORE_LCL_FLD    long
+    //   \--*  CNS_INT(h) long
+    //------------------------------------------------------------------------
+
+    // Initialize the method table pointer.
+    GenTree*   init     = comp->gtNewStoreLclFldNode(lclNum, TYP_I_IMPL, 0, newArr->gtArgs.GetArgByIndex(0)->GetNode());
+    Statement* initStmt = comp->gtNewStmt(init);
+
+    comp->fgInsertStmtBefore(block, stmt, initStmt);
+
+    // Initialize the array length.
+    //
+    //------------------------------------------------------------------------
+    // STMTx (IL 0x... ???)
+    //   * STORE_LCL_FLD    int
+    //   \--*  CNS_INT    int
+    //------------------------------------------------------------------------
+
+    // Pass the total length of the array.
+    GenTree*   len     = comp->gtNewStoreLclFldNode(lclNum, TYP_INT, OFFSETOF__CORINFO_Array__length,
+                                                    comp->gtNewIconNode(length, TYP_INT));
+    Statement* lenStmt = comp->gtNewStmt(len);
+    comp->fgInsertStmtBefore(block, stmt, lenStmt);
+
+    return lclNum;
+}
+
 //------------------------------------------------------------------------
 // MorphAllocObjNodeIntoStackAlloc: Morph a GT_ALLOCOBJ node into stack
 //                                  allocation.
@@ -682,6 +854,8 @@ bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack<GenTree*>* parent
             case GT_EQ:
             case GT_NE:
             case GT_NULLCHECK:
+            case GT_ARR_LENGTH:
+            case GT_INDEX_ADDR:
                 canLclVarEscapeViaParentStack = false;
                 break;
 
@@ -778,6 +952,8 @@ void ObjectAllocator::UpdateAncestorTypes(GenTree* tree, ArrayStack<GenTree*>* p
             case GT_EQ:
             case GT_NE:
             case GT_NULLCHECK:
+            case GT_ARR_LENGTH:
+            case GT_INDEX_ADDR:
                 break;
 
             case GT_COMMA:
@@ -1025,6 +1201,50 @@ void ObjectAllocator::RewriteUses()
                     }
                 }
             }
+            // Rewrite INDEX_ADDR to ADD for stack-allocated arrays.
+            else if (tree->OperIs(GT_INDEX_ADDR))
+            {
+                GenTreeIndexAddr* const gtIndexAddr = tree->AsIndexAddr();
+                GenTree* const          gtArr       = gtIndexAddr->Arr();
+                GenTree* const          gtInd       = gtIndexAddr->Index();
+                unsigned int            arrLen      = 0;
+
+                if (gtArr->OperIs(GT_LCL_ADDR) && gtInd->IsCnsIntOrI() &&
+                    m_allocator->m_LocalArrToLenMap.TryGetValue(gtArr->AsLclVarCommon()->GetLclNum(), &arrLen))
+                {
+                    if (gtInd->AsIntCon()->gtIconVal < arrLen)
+                    {
+                        JITDUMP("Rewriting INDEX_ADDR to ADD [%06u]\n", m_compiler->dspTreeID(tree));
+                        const ssize_t offset =
+                            OFFSETOF__CORINFO_Array__data + gtInd->AsIntCon()->gtIconVal * gtIndexAddr->gtElemSize;
+                        GenTree* const gtOffset = m_compiler->gtNewIconNode(offset, TYP_I_IMPL);
+                        GenTree* const gtAdd    = m_compiler->gtNewOperNode(GT_ADD, TYP_BYREF, gtArr, gtOffset);
+                        *use                    = gtAdd;
+                    }
+                    else
+                    {
+                        JITDUMP("Rewriting INDEX_ADDR to RNGCHKFAIL helper call [%06u]\n", m_compiler->dspTreeID(tree));
+                        GenTree* const gtRngChkFail =
+                            m_compiler->gtNewMustThrowException(CORINFO_HELP_RNGCHKFAIL, gtIndexAddr->gtType,
+                                                                gtIndexAddr->gtStructElemClass);
+                        *use = gtRngChkFail;
+                    }
+                }
+            }
+            // Rewrite ARR_LENGTH to LCL_FLD for stack-allocated arrays.
+            else if (tree->OperIsArrLength())
+            {
+                GenTreeArrLen* const gtArrLen = tree->AsArrLen();
+                GenTree* const       gtArr    = gtArrLen->ArrRef();
+
+                if (gtArr->OperIs(GT_LCL_ADDR))
+                {
+                    JITDUMP("Rewriting ARR_LENGTH to LCL_FLD [%06u]\n", m_compiler->dspTreeID(tree));
+                    GenTree* const gtLclFld = m_compiler->gtNewLclFldNode(gtArr->AsLclVarCommon()->GetLclNum(), TYP_INT,
+                                                                          OFFSETOF__CORINFO_Array__length);
+                    *use                    = gtLclFld;
+                }
+            }
 
             return Compiler::fgWalkResult::WALK_CONTINUE;
         }