Skip to content

Commit

Permalink
JIT: Generalize FlowGraphNaturalLoop::AnalyzeIteration (#97327)
Browse files Browse the repository at this point in the history
Consider all exit edges as potential tests on a IV variable, and try to
recognize induction from those. Previously we looked only at the
lexically bottom most block.

Overall this unlocks only a pretty minor number of new cases which then
leads to a bit more cloning. Size regressions are expected because of
that.
  • Loading branch information
jakobbotsch authored Jan 23, 2024
1 parent 5e1ad18 commit c5944d4
Show file tree
Hide file tree
Showing 3 changed files with 95 additions and 76 deletions.
2 changes: 1 addition & 1 deletion src/coreclr/jit/compiler.h
Original file line number Diff line number Diff line change
Expand Up @@ -2089,7 +2089,7 @@ class FlowGraphNaturalLoop
GenTreeLclVarCommon* FindDef(unsigned lclNum);

void MatchInit(NaturalLoopIterInfo* info, BasicBlock* initBlock, GenTree* init);
bool MatchLimit(NaturalLoopIterInfo* info, GenTree* test);
bool MatchLimit(unsigned iterVar, GenTree* test, NaturalLoopIterInfo* info);
bool CheckLoopConditionBaseCase(BasicBlock* initBlock, NaturalLoopIterInfo* info);
template<typename T>
static bool EvaluateRelop(T op1, T op2, genTreeOps oper);
Expand Down
135 changes: 82 additions & 53 deletions src/coreclr/jit/flowgraph.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -4349,7 +4349,7 @@ FlowGraphNaturalLoops* FlowGraphNaturalLoops::Find(const FlowGraphDfsTree* dfsTr

// Find the exit edges
//
loop->VisitLoopBlocks([=](BasicBlock* loopBlock) {
loop->VisitLoopBlocksReversePostOrder([=](BasicBlock* loopBlock) {
loopBlock->VisitRegularSuccs(comp, [=](BasicBlock* succBlock) {
if (!loop->ContainsBlock(succBlock))
{
Expand Down Expand Up @@ -4929,51 +4929,79 @@ bool FlowGraphNaturalLoop::AnalyzeIteration(NaturalLoopIterInfo* info)

JITDUMP(" Preheader = " FMT_BB "\n", preheader->bbNum);

// TODO-Quirk: For backwards compatibility always try the lexically
// bottom-most block for the loop variable.
BasicBlock* bottom = GetLexicallyBottomMostBlock();
JITDUMP(" Bottom = " FMT_BB "\n", bottom->bbNum);
BasicBlock* initBlock = nullptr;
GenTree* init = nullptr;
GenTree* test = nullptr;

BasicBlock* cond = bottom;
BasicBlock* initBlock = preheader;
GenTree* init;
GenTree* test;
if (!cond->KindIs(BBJ_COND) ||
!comp->optExtractInitTestIncr(&initBlock, bottom, m_header, &init, &test, &info->IterTree))
info->IterVar = BAD_VAR_NUM;

for (FlowEdge* exitEdge : ExitEdges())
{
// TODO-CQ: Try all exit edges here to see if we can find an induction variable.
JITDUMP(" Could not extract induction variable from bottom\n");
return false;
}
BasicBlock* cond = exitEdge->getSourceBlock();
JITDUMP(" Checking exiting block " FMT_BB "\n", cond->bbNum);
if (!cond->KindIs(BBJ_COND))
{
JITDUMP(" Not a BBJ_COND\n");
continue;
}

assert(ContainsBlock(cond->GetTrueTarget()) != ContainsBlock(cond->GetFalseTarget()));
GenTree* iterTree = nullptr;
initBlock = preheader;
if (!comp->optExtractInitTestIncr(&initBlock, cond, m_header, &init, &test, &iterTree))
{
JITDUMP(" Could not extract an IV\n");
continue;
}

info->TestBlock = cond;
info->IterVar = comp->optIsLoopIncrTree(info->IterTree);
info->ExitedOnTrue = !ContainsBlock(cond->GetTrueTarget());
unsigned iterVar = comp->optIsLoopIncrTree(iterTree);
assert(iterVar != BAD_VAR_NUM);
LclVarDsc* const iterVarDsc = comp->lvaGetDesc(iterVar);
// Bail on promoted case, otherwise we'd have to search the loop
// for both iterVar and its parent.
// TODO-CQ: Fix this
//
if (iterVarDsc->lvIsStructField)
{
JITDUMP(" iterVar V%02u is a promoted field\n", iterVar);
continue;
}

// TODO-CQ: Currently, since we only look at the lexically bottom most block all loops have
// ExitedOnTrue == false. Once we recognize more structures this can be true.
assert(!info->ExitedOnTrue);
// Bail on the potentially aliased case.
//
if (iterVarDsc->IsAddressExposed())
{
JITDUMP(" iterVar V%02u is address exposed\n", iterVar);
continue;
}

assert(info->IterVar != BAD_VAR_NUM);
LclVarDsc* const iterVarDsc = comp->lvaGetDesc(info->IterVar);
if (!MatchLimit(iterVar, test, info))
{
continue;
}

// Bail on promoted case, otherwise we'd have to search the loop
// for both iterVar and its parent.
// TODO-CQ: Fix this
//
if (iterVarDsc->lvIsStructField)
{
JITDUMP(" iterVar V%02u is a promoted field\n", info->IterVar);
return false;
bool result = VisitDefs([=](GenTreeLclVarCommon* def) {
if ((def->GetLclNum() != iterVar) || (def == iterTree))
return true;

JITDUMP(" Loop has extraneous def [%06u]\n", Compiler::dspTreeID(def));
return false;
});

if (!result)
{
continue;
}

info->TestBlock = cond;
info->IterVar = iterVar;
info->IterTree = iterTree;
info->ExitedOnTrue = !ContainsBlock(cond->GetTrueTarget());
break;
}

// Bail on the potentially aliased case.
//
if (iterVarDsc->IsAddressExposed())
if (info->IterVar == BAD_VAR_NUM)
{
JITDUMP(" iterVar V%02u is address exposed\n", info->IterVar);
JITDUMP(" Could not find any IV\n");
return false;
}

Expand All @@ -4988,11 +5016,6 @@ bool FlowGraphNaturalLoop::AnalyzeIteration(NaturalLoopIterInfo* info)
Compiler::dspTreeID(info->IterTree));
}

if (!MatchLimit(info, test))
{
return false;
}

MatchInit(info, initBlock, init);

bool result = VisitDefs([=](GenTreeLclVarCommon* def) {
Expand Down Expand Up @@ -5036,7 +5059,7 @@ bool FlowGraphNaturalLoop::AnalyzeIteration(NaturalLoopIterInfo* info)
}
#endif

return result;
return true;
}

//------------------------------------------------------------------------
Expand Down Expand Up @@ -5071,8 +5094,9 @@ void FlowGraphNaturalLoop::MatchInit(NaturalLoopIterInfo* info, BasicBlock* init
// induction variable.
//
// Parameters:
// info - [in, out] Info structure to query and fill out
// test - Loop condition test
// iterVar - Local number of potential IV
// test - Loop condition test
// info - [out] Info structure to fill out with information about the limit
//
// Returns:
// True if the loop condition was recognized and "info" was filled out.
Expand All @@ -5081,8 +5105,13 @@ void FlowGraphNaturalLoop::MatchInit(NaturalLoopIterInfo* info, BasicBlock* init
// Unlike the initialization, we do require that we are able to match the
// loop condition.
//
bool FlowGraphNaturalLoop::MatchLimit(NaturalLoopIterInfo* info, GenTree* test)
bool FlowGraphNaturalLoop::MatchLimit(unsigned iterVar, GenTree* test, NaturalLoopIterInfo* info)
{
info->HasConstLimit = false;
info->HasSimdLimit = false;
info->HasArrayLengthLimit = false;
info->HasInvariantLocalLimit = false;

Compiler* comp = m_dfsTree->GetCompiler();

// Obtain the relop from the "test" tree.
Expand All @@ -5106,12 +5135,12 @@ bool FlowGraphNaturalLoop::MatchLimit(NaturalLoopIterInfo* info, GenTree* test)
GenTree* limitOp;

// Make sure op1 or op2 is the iterVar.
if (opr1->OperIsScalarLocal() && (opr1->AsLclVarCommon()->GetLclNum() == info->IterVar))
if (opr1->OperIsScalarLocal() && (opr1->AsLclVarCommon()->GetLclNum() == iterVar))
{
iterOp = opr1;
limitOp = opr2;
}
else if (opr2->OperIsScalarLocal() && (opr2->AsLclVarCommon()->GetLclNum() == info->IterVar))
else if (opr2->OperIsScalarLocal() && (opr2->AsLclVarCommon()->GetLclNum() == iterVar))
{
iterOp = opr2;
limitOp = opr1;
Expand Down Expand Up @@ -5141,14 +5170,14 @@ bool FlowGraphNaturalLoop::MatchLimit(NaturalLoopIterInfo* info, GenTree* test)
//
if (comp->lvaGetDesc(limitOp->AsLclVarCommon())->IsAddressExposed())
{
JITDUMP(" Limit var V%02u is address exposed\n", limitOp->AsLclVarCommon()->GetLclNum());
JITDUMP(" Limit var V%02u is address exposed\n", limitOp->AsLclVarCommon()->GetLclNum());
return false;
}

GenTreeLclVarCommon* def = FindDef(limitOp->AsLclVarCommon()->GetLclNum());
if (def != nullptr)
{
JITDUMP(" Limit var V%02u modified by [%06u]\n", limitOp->AsLclVarCommon()->GetLclNum(),
JITDUMP(" Limit var V%02u modified by [%06u]\n", limitOp->AsLclVarCommon()->GetLclNum(),
Compiler::dspTreeID(def));
return false;
}
Expand All @@ -5163,20 +5192,20 @@ bool FlowGraphNaturalLoop::MatchLimit(NaturalLoopIterInfo* info, GenTree* test)

if (!array->OperIs(GT_LCL_VAR))
{
JITDUMP(" Array limit tree [%06u] not analyzable\n", Compiler::dspTreeID(limitOp));
JITDUMP(" Array limit tree [%06u] not analyzable\n", Compiler::dspTreeID(limitOp));
return false;
}

if (comp->lvaGetDesc(array->AsLclVarCommon())->IsAddressExposed())
{
JITDUMP(" Array base local V%02u is address exposed\n", array->AsLclVarCommon()->GetLclNum());
JITDUMP(" Array base local V%02u is address exposed\n", array->AsLclVarCommon()->GetLclNum());
return false;
}

GenTreeLclVarCommon* def = FindDef(array->AsLclVarCommon()->GetLclNum());
if (def != nullptr)
{
JITDUMP(" Array limit var V%02u modified by [%06u]\n", array->AsLclVarCommon()->GetLclNum(),
JITDUMP(" Array limit var V%02u modified by [%06u]\n", array->AsLclVarCommon()->GetLclNum(),
Compiler::dspTreeID(def));
return false;
}
Expand All @@ -5185,7 +5214,7 @@ bool FlowGraphNaturalLoop::MatchLimit(NaturalLoopIterInfo* info, GenTree* test)
}
else
{
JITDUMP(" Loop limit tree [%06u] not analyzable\n", Compiler::dspTreeID(limitOp));
JITDUMP(" Loop limit tree [%06u] not analyzable\n", Compiler::dspTreeID(limitOp));
return false;
}

Expand Down
34 changes: 12 additions & 22 deletions src/coreclr/jit/optimizer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -393,11 +393,11 @@ bool Compiler::optIsLoopTestEvalIntoTemp(Statement* testStmt, Statement** newTes
// Arguments:
// pInitBlock - [IN/OUT] *pInitBlock is the loop head block on entry, and is set to the initBlock on exit,
// if `**ppInit` is non-null.
// bottom - Loop bottom block
// top - Loop top block
// ppInit - The init stmt of the loop if found.
// ppTest - The test stmt of the loop if found.
// ppIncr - The incr stmt of the loop if found.
// cond - A BBJ_COND block that exits the loop
// header - Loop header block
// ppInit - [out] The init stmt of the loop if found.
// ppTest - [out] The test stmt of the loop if found.
// ppIncr - [out] The incr stmt of the loop if found.
//
// Return Value:
// The results are put in "ppInit", "ppTest" and "ppIncr" if the method
Expand All @@ -406,25 +406,15 @@ bool Compiler::optIsLoopTestEvalIntoTemp(Statement* testStmt, Statement** newTes
// to nullptr. Return value will never be false if `init` is not found.
//
// Operation:
// Check if the "test" stmt is last stmt in the loop "bottom". Try to find the "incr" stmt.
// Check previous stmt of "test" to get the "incr" stmt. If it is not found it could be a loop of the
// below form.
//
// +-------<-----------------<-----------+
// | |
// v |
// BBinit(head) -> BBcond(top) -> BBLoopBody(bottom) ---^
//
// Check if the "incr" tree is present in the loop "top" node as the last stmt.
// Also check if the "test" tree is assigned to a tmp node and the tmp is used
// in the jtrue condition.
// Check if the "test" stmt is last stmt in an exiting BBJ_COND block of the loop. Try to find the "incr" stmt.
// Check previous stmt of "test" to get the "incr" stmt.
//
// Note:
// This method just retrieves what it thinks is the "test" node,
// the callers are expected to verify that "iterVar" is used in the test.
//
bool Compiler::optExtractInitTestIncr(
BasicBlock** pInitBlock, BasicBlock* bottom, BasicBlock* top, GenTree** ppInit, GenTree** ppTest, GenTree** ppIncr)
BasicBlock** pInitBlock, BasicBlock* cond, BasicBlock* header, GenTree** ppInit, GenTree** ppTest, GenTree** ppIncr)
{
assert(pInitBlock != nullptr);
assert(ppInit != nullptr);
Expand All @@ -433,8 +423,8 @@ bool Compiler::optExtractInitTestIncr(

// Check if last two statements in the loop body are the increment of the iterator
// and the loop termination test.
noway_assert(bottom->bbStmtList != nullptr);
Statement* testStmt = bottom->lastStmt();
noway_assert(cond->bbStmtList != nullptr);
Statement* testStmt = cond->lastStmt();
noway_assert(testStmt != nullptr && testStmt->GetNextStmt() == nullptr);

Statement* newTestStmt;
Expand All @@ -444,7 +434,7 @@ bool Compiler::optExtractInitTestIncr(
}

// Check if we have the incr stmt before the test stmt, if we don't,
// check if incr is part of the loop "top".
// check if incr is part of the loop "header".
Statement* incrStmt = testStmt->GetPrevStmt();

// If we've added profile instrumentation, we may need to skip past a BB counter update.
Expand Down Expand Up @@ -472,7 +462,7 @@ bool Compiler::optExtractInitTestIncr(
// If we are rebuilding the loops, we would already have the pre-header block introduced
// the first time, which might be empty if no hoisting has yet occurred. In this case, look a
// little harder for the possible loop initialization statement.
if (initBlock->KindIs(BBJ_ALWAYS) && initBlock->TargetIs(top) && (initBlock->countOfInEdges() == 1) &&
if (initBlock->KindIs(BBJ_ALWAYS) && initBlock->TargetIs(header) && (initBlock->countOfInEdges() == 1) &&
!initBlock->IsFirst() && initBlock->Prev()->bbFallsThrough())
{
initBlock = initBlock->Prev();
Expand Down

0 comments on commit c5944d4

Please sign in to comment.