diff --git a/src/coreclr/jit/objectalloc.cpp b/src/coreclr/jit/objectalloc.cpp index 89e28c5978c6d4..7ca0bf4791adb3 100644 --- a/src/coreclr/jit/objectalloc.cpp +++ b/src/coreclr/jit/objectalloc.cpp @@ -169,7 +169,9 @@ void ObjectAllocator::MarkEscapingVarsAndBuildConnGraph() { class BuildConnGraphVisitor final : public GenTreeVisitor { - ObjectAllocator* m_allocator; + ObjectAllocator* m_allocator; + bool m_hasBackwardJump; + BitSetShortLongRep m_lclsInPreds; public: enum @@ -179,9 +181,11 @@ void ObjectAllocator::MarkEscapingVarsAndBuildConnGraph() ComputeStack = true, }; - BuildConnGraphVisitor(ObjectAllocator* allocator) + BuildConnGraphVisitor(ObjectAllocator* allocator, bool hasBackwardJump, BitSetShortLongRep lclsInPreds) : GenTreeVisitor(allocator->comp) , m_allocator(allocator) + , m_hasBackwardJump(hasBackwardJump) + , m_lclsInPreds(lclsInPreds) { } @@ -207,7 +211,12 @@ void ObjectAllocator::MarkEscapingVarsAndBuildConnGraph() { assert(tree == m_ancestors.Top()); - if (!m_allocator->CanLclVarEscapeViaParentStack(&m_ancestors, lclNum)) + if (m_hasBackwardJump) + { + BitVecOps::AddElemD(&m_allocator->m_bitVecTraits, m_lclsInPreds, lclNum); + } + + if (!m_allocator->CanLclVarEscapeViaParentStack(&m_ancestors, lclNum, m_hasBackwardJump, m_lclsInPreds)) { lclEscapes = false; } @@ -249,9 +258,30 @@ void ObjectAllocator::MarkEscapingVarsAndBuildConnGraph() for (BasicBlock* const block : comp->Blocks()) { + const bool basicBlockHasBackwardJump = block->HasFlag(BBF_BACKWARD_JUMP); + BitSetShortLongRep lclsInPreds = BitVecOps::UninitVal(); + + if (basicBlockHasBackwardJump) + { + lclsInPreds = BitVecOps::MakeEmpty(&m_bitVecTraits); + for (BasicBlock* const predBlock : block->PredBlocks()) + { + for (Statement* const stmt : predBlock->Statements()) + { + GenTree* stmtExpr = stmt->GetRootNode(); + if (stmtExpr->OperIsLocalStore()) + { + GenTreeLclVarCommon* lclVar = stmtExpr->AsLclVarCommon(); + unsigned int lclNum = lclVar->GetLclNum(); + BitVecOps::AddElemD(&m_bitVecTraits, lclsInPreds, lclNum); + } + } + } + } + for (Statement* const stmt : block->Statements()) { - BuildConnGraphVisitor buildConnGraphVisitor(this); + BuildConnGraphVisitor buildConnGraphVisitor(this, basicBlockHasBackwardJump, lclsInPreds); buildConnGraphVisitor.WalkTree(stmt->GetRootNodePointer(), nullptr); } } @@ -376,8 +406,7 @@ bool ObjectAllocator::MorphAllocObjNodes() for (BasicBlock* const block : comp->Blocks()) { - const bool basicBlockHasNewObj = block->HasFlag(BBF_HAS_NEWOBJ); - const bool basicBlockHasBackwardJump = block->HasFlag(BBF_BACKWARD_JUMP); + const bool basicBlockHasNewObj = block->HasFlag(BBF_HAS_NEWOBJ); #ifndef DEBUG if (!basicBlockHasNewObj) { @@ -438,11 +467,6 @@ bool ObjectAllocator::MorphAllocObjNodes() onHeapReason = "[object stack allocation disabled]"; canStack = false; } - else if (basicBlockHasBackwardJump) - { - onHeapReason = "[alloc in loop]"; - canStack = false; - } else if (!CanAllocateLclVarOnStack(lclNum, clsHnd, &onHeapReason)) { // reason set by the call @@ -633,6 +657,8 @@ unsigned int ObjectAllocator::MorphAllocObjNodeIntoStackAlloc( // Arguments: // parentStack - Parent stack of the current visit // lclNum - Local variable number +// hasBackwardJump - Indicates if the block containing the parent stack has a backward jump +// lclsInPreds - Local variables stored in the predecessors // // Return Value: // true if the local can escape via the parent stack; false otherwise @@ -641,7 +667,10 @@ unsigned int ObjectAllocator::MorphAllocObjNodeIntoStackAlloc( // The method currently treats all locals assigned to a field as escaping. // The can potentially be tracked by special field edges in the connection graph. // -bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack* parentStack, unsigned int lclNum) +bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack* parentStack, + unsigned int lclNum, + bool hasBackwardJump, + BitSetShortLongRep lclsInPreds) { assert(parentStack != nullptr); int parentIndex = 1; @@ -675,7 +704,10 @@ bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack* parent const unsigned int srcLclNum = lclNum; AddConnGraphEdge(dstLclNum, srcLclNum); - canLclVarEscapeViaParentStack = false; + + canLclVarEscapeViaParentStack = hasBackwardJump && + parent->AsLclVar()->TypeIs(TYP_REF, TYP_BYREF, TYP_I_IMPL) && + !BitVecOps::IsMember(&m_bitVecTraits, lclsInPreds, dstLclNum); } break; diff --git a/src/coreclr/jit/objectalloc.h b/src/coreclr/jit/objectalloc.h index 8a873be8b34ad7..a2bdc7f6df18b0 100644 --- a/src/coreclr/jit/objectalloc.h +++ b/src/coreclr/jit/objectalloc.h @@ -65,7 +65,10 @@ class ObjectAllocator final : public Phase unsigned int MorphAllocObjNodeIntoStackAlloc( GenTreeAllocObj* allocObj, CORINFO_CLASS_HANDLE clsHnd, bool isValueClass, BasicBlock* block, Statement* stmt); struct BuildConnGraphVisitorCallbackData; - bool CanLclVarEscapeViaParentStack(ArrayStack* parentStack, unsigned int lclNum); + bool CanLclVarEscapeViaParentStack(ArrayStack* parentStack, + unsigned int lclNum, + bool hasBackwardJump, + BitSetShortLongRep lclsInPreds); void UpdateAncestorTypes(GenTree* tree, ArrayStack* parentStack, var_types newType); static const unsigned int s_StackAllocMaxSize = 0x2000U;