Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

JIT: Allow stack allocate objects in loops #106526

Closed
wants to merge 3 commits into from
Closed

JIT: Allow stack allocate objects in loops #106526

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
442aefd
Allow stack allocate objects in loop
hez2010 Aug 16, 2024
a2825ac
Format
hez2010 Aug 16, 2024
463f36e
Check preds
hez2010 Aug 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 45 additions & 13 deletions src/coreclr/jit/objectalloc.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,9 @@ void ObjectAllocator::MarkEscapingVarsAndBuildConnGraph()
{
class BuildConnGraphVisitor final : public GenTreeVisitor<BuildConnGraphVisitor>
{
ObjectAllocator* m_allocator;
ObjectAllocator* m_allocator;
bool m_hasBackwardJump;
BitSetShortLongRep m_lclsInPreds;

public:
enum
Expand All @@ -179,9 +181,11 @@ void ObjectAllocator::MarkEscapingVarsAndBuildConnGraph()
ComputeStack = true,
};

BuildConnGraphVisitor(ObjectAllocator* allocator)
BuildConnGraphVisitor(ObjectAllocator* allocator, bool hasBackwardJump, BitSetShortLongRep lclsInPreds)
: GenTreeVisitor<BuildConnGraphVisitor>(allocator->comp)
, m_allocator(allocator)
, m_hasBackwardJump(hasBackwardJump)
, m_lclsInPreds(lclsInPreds)
{
}

Expand All @@ -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;
}
Expand Down Expand Up @@ -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);
}
}
}
}
Comment on lines +264 to +280
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is checking the preds like this sufficient? What prevents the allocated object from being stored in some other block inside the loop that reaches the backedge?

I think doing this analysis is going to require building the proper loop structures and probably also some form of liveness.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What prevents the allocated object from being stored in some other block inside the loop that reaches the backedge?

My thought is that when we see a backward edge, we check the local where the object is being stored to see if the target is a local defined in its preds, otherwise mark it escaping. For blocks that have back edges to the current block, they will be checked later so that we don't miss them.
IMO this should be a conservative approximation without accurate loop and liveness.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree with Jakob here. I think the analysis is more complicated than this, and we need proper loop structure.

If you have a proper loop where the entry block dominates the body, then locals defined in the entry block can still be live on the back edges if they are first used before being defined.

Building loops and running liveness seems to be a necessary prerequisite, and unfortunately, liveness will be impaired at this stage of the jit, because we haven't yet figured out which locals are truly exposed, so we may end up also needing points-to analysis to help refine what can be aliased (and that likely should dovetail with or supplant the existing escape analysis), and proper loopness might also be a challenge given that we might end up creating a new loop after this phase because of a recursive tail call.

Assuming we can indeed soundly prove that an object's references cannot outlive a loop iteration (accounting for any possible liveness from EH), we'd also need to ensure that the object fields are zeroed (as they would be for a heap allocation) and/or that they are non-gc fields that are definitely written before they are read.

None of these problems are impossible, but it's a fair amount of work.


for (Statement* const stmt : block->Statements())
{
BuildConnGraphVisitor buildConnGraphVisitor(this);
BuildConnGraphVisitor buildConnGraphVisitor(this, basicBlockHasBackwardJump, lclsInPreds);
buildConnGraphVisitor.WalkTree(stmt->GetRootNodePointer(), nullptr);
}
}
Expand Down Expand Up @@ -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)
{
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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<GenTree*>* parentStack, unsigned int lclNum)
bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack<GenTree*>* parentStack,
unsigned int lclNum,
bool hasBackwardJump,
BitSetShortLongRep lclsInPreds)
{
assert(parentStack != nullptr);
int parentIndex = 1;
Expand Down Expand Up @@ -675,7 +704,10 @@ bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack<GenTree*>* 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;

Expand Down
5 changes: 4 additions & 1 deletion src/coreclr/jit/objectalloc.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<GenTree*>* parentStack, unsigned int lclNum);
bool CanLclVarEscapeViaParentStack(ArrayStack<GenTree*>* parentStack,
unsigned int lclNum,
bool hasBackwardJump,
BitSetShortLongRep lclsInPreds);
void UpdateAncestorTypes(GenTree* tree, ArrayStack<GenTree*>* parentStack, var_types newType);

static const unsigned int s_StackAllocMaxSize = 0x2000U;
Expand Down
Loading