#include "slang-ir-loop-inversion.h"
#include "slang-ir-clone.h"
#include "slang-ir-dominators.h"
#include "slang-ir-insts.h"
#include "slang-ir-lower-witness-lookup.h"
#include "slang-ir-reachability.h"
#include "slang-ir-ssa-simplification.h"
#include "slang-ir-util.h"
#include "slang-ir.h"
namespace Slang
{
static bool isSameBlockOrTrivialBranch(IRBlock* target, IRBlock* scrutinee)
{
if(target == scrutinee)
return true;
const auto br = as<IRUnconditionalBranch>(scrutinee->getFirstOrdinaryInst());
return br && br->getTargetBlock() == target && br->getArgCount() == 0 && !scrutinee->hasMoreThanOneUse();
};
static bool isSmallBlock(IRBlock* c)
{
// Somewhat arbitrarily, 4 instructions, enough for:
// - Arith
// - Comparison
// - Negation
// - Terminator
Int n = 0;
for([[maybe_unused]] const auto i : c->getOrdinaryInsts())
if(++n > 4)
return false;
return true;
}
static bool hasIrrelevantContinueBlock(IRLoop* loop)
{
const auto c = loop->getContinueBlock();
return c == loop->getTargetBlock() || c->getPredecessors().getCount() <= 1;
}
// Loops are suitable for inversion if:
// - The loop jumps to a conditional branch which has the break block as one of
// its successors (or a trivial break block which we erase) and the other
// successor is empty
// - The conditional block is "small", because we will be duplicating it
// - The loop's continue block is irrelevant, because we'll need to change it,
// either by being the loop header already or by having only a single use
// within the loop body
static bool isSuitableForInversion(IRLoop* loop)
{
const auto nextBlock = loop->getTargetBlock();
const auto breakBlock = loop->getBreakBlock();
// The first thing a loop does must be a conditional
const auto branch = as<IRIfElse>(nextBlock->getTerminator());
if(!branch)
return false;
if(!isSmallBlock(nextBlock))
return false;
if(!hasIrrelevantContinueBlock(loop))
return false;
const auto t = branch->getTrueBlock();
const auto f = branch->getFalseBlock();
const auto a = branch->getAfterBlock();
//
// In principle we could perform this simplification in the cfg simplifier,
// however it relies on slightly more context than is simple to insert
// there, namely that the removed trivial branching block is branching to a
// loop break block.
//
// Do we break on the 'true' side?
if(isSameBlockOrTrivialBranch(breakBlock, t) && f == a)
{
if(t != breakBlock)
{
branch->trueBlock.set(breakBlock);
t->removeAndDeallocate();
}
return true;
}
// ... or the false side
if(isSameBlockOrTrivialBranch(breakBlock, f) && t == a)
{
if(f != breakBlock)
{
branch->falseBlock.set(breakBlock);
f->removeAndDeallocate();
}
return true;
}
return false;
}
static IRParam* duplicateToParamWithDecorations(IRBuilder& builder, IRCloneEnv& cloneEnv, IRInst* i)
{
const auto p = builder.emitParam(i->getFullType());
for(const auto dec : i->getDecorations())
cloneDecoration(&cloneEnv, dec, p, builder.getModule());
return p;
}
// Given
// s: ...1 loop break=b next=c1
// c1: if x then goto b else goto d (merge at d)
// d: goto c1
// b: ...2
//
// Produce:
// s: ...1 goto c2
// c2: if x then goto e2 else goto l (merge at b)
// e2: goto b
// l: loop break=b next=d
// d: goto c1:
// c1: if x then goto e1 else goto e3 (merge at e3)
// e3: goto d
// e1: goto b
// b: ...2
//
// s is the Start block
// c1, c2 are the Condition blocks
// e1, e2, e3 are the critical Edge breakers
// l is the Loop entering block
// d is the loop boDy
// b is the Break block
static void invertLoop(IRBuilder& builder, IRLoop* loop)
{
IRBuilderInsertLocScope builderScope(&builder);
const auto s = as<IRBlock>(loop->getParent());
auto domTree = computeDominatorTree(s->getParent());
SLANG_ASSERT(s);
const auto c1 = loop->getTargetBlock();
const auto c1Terminator = as<IRIfElse>(c1->getTerminator());
SLANG_ASSERT(c1Terminator);
const auto b = loop->getBreakBlock();
auto& c1dUse = c1Terminator->getTrueBlock() == b ? c1Terminator->falseBlock : c1Terminator->trueBlock;
auto& c1bUse = c1Terminator->getTrueBlock() == b ? c1Terminator->trueBlock : c1Terminator->falseBlock;
const auto d = as<IRBlock>(c1dUse.get());
SLANG_ASSERT(d);
IRCloneEnv cloneEnv;
cloneEnv.squashChildrenMapping = true;
// Since we are duplicating the loop break condition block (c1) we must
// introduce phi values for anything in it upon which the rest of the
// program (b onwards) uses. Lift the values fron c1 used in b (and
// onwards) to parameters. To avoid a critical edge, pass these via a new
// block, e1.
builder.setInsertInto(b);
List<IRInst*> c1Params;
for(auto i : IRInstList<IRInst>(c1->getFirstInst(), c1->getLastInst()))
{
IRParam* p = nullptr;
traverseUses(i, [&](IRUse* u){
auto userBlock = u->getUser()->getParent();
if(domTree->dominates(b, userBlock))
{
// A new parameter to replace this 'i'
if(!p)
p = duplicateToParamWithDecorations(builder, cloneEnv, i);
u->set(p);
}
});
if(p)
c1Params.add(i);
}
// Create another break block b2 that will act as the new break block for the
// loop. The original break block b will become the merge point for the outer condition.
//
auto b2 = builder.emitBlock();
b2->insertBefore(b);
// Create a copy of the parameters in b. b2 will simply pass these to b.
List<IRInst*> b2Params;
for(auto p : b->getParams())
{
auto q = duplicateToParamWithDecorations(builder, cloneEnv, p);
b2Params.add(q);
}
builder.setInsertInto(b2);
builder.emitBranch(b, b2Params.getCount(), b2Params.getBuffer());
auto e1 = builder.emitBlock();
e1->insertAfter(c1);
builder.emitBranch(b2, c1Params.getCount(), c1Params.getBuffer());
c1bUse.set(e1);
c1Terminator->afterBlock.set(d);
// We create another block e4 to handle other breaks from inside the loop, and
// rewrite existing breaks to jump to it.
// We keep e4 and e1 distinct since after the inversion step, insts in e4 will
// be re-written to use values from the loop entry, while e1 will use values from
// c1.
//
auto e4 = builder.emitBlock();
e4->insertBefore(c1);
builder.emitBranch(b2, c1Params.getCount(), c1Params.getBuffer());
traverseUses(b, [&](IRUse* u){
auto userBlock = u->getUser()->getParent();
// Restrict this to just those blocks within this loop
if(userBlock != e4 && userBlock != e1 && userBlock != b2 && domTree->dominates(s, userBlock) && !domTree->dominates(b, userBlock))
u->set(e4);
});
// We now have
// s: ...1 loop break=b next=c1
// e4: goto b2
// c1: if x then goto e1 else goto d (merge at d)
// e1: goto b2
// d: goto c1
// b2: goto b
// b: ...2
// Duplicate c1 into c2, and using the same cloneEnv, duplicate e1 into e2
builder.setInsertInto(builder.getFunc());
const auto c2 = as<IRBlock>(cloneInst(&cloneEnv, &builder, c1));
c2->insertBefore(c1);
const auto e2 = as<IRBlock>(cloneInst(&cloneEnv, &builder, e1));
e2->insertAfter(c2);
const auto c2Terminator = as<IRConditionalBranch>(c2->getTerminator());
auto& c2eUse = c2Terminator->getTrueBlock() == e1 ? c2Terminator->trueBlock : c2Terminator->falseBlock;
c2eUse.set(e2);
builder.setInsertAfter(c2Terminator);
const auto newC2Terminator = builder.emitIfElse(c2Terminator->getCondition(), c2Terminator->getTrueBlock(), c2Terminator->getFalseBlock(), b);
c2Terminator->removeAndDeallocate();
// The cloned e2 will branch into b2 by default, rewrite it to branch to b, the correct merge point.
SLANG_ASSERT(cast<IRUnconditionalBranch>(e2->getTerminator())->getTargetBlock() == b2);
cast<IRUnconditionalBranch>(e2->getTerminator())->block.set(b);
// We now have
// s: ...1 loop break=b next=c1
// e4: goto b2
// c2: if x then goto e2 else goto d (merge at b)
// e2: goto b
// c1: if x then goto e1 else goto d (merge at d)
// e1: goto b2
// d: goto c1
// b2: goto b
// b: ...2
// move the loop instruction to its own block, l
const auto l = builder.emitBlock();
l->insertAfter(e2);
loop->insertAtEnd(l);
e4->insertBefore(c1);
// We now have
// s: ...1 no-termiator
// c2: if x then goto e2 else goto d (merge at b)
// e2: goto b
// l: loop break=b next=c1
// e4: goto b2
// c1: if x then goto e1 else goto d (merge at d)
// e1: goto b2
// d: goto c1
// b2: goto b
// b: ...2
// add a new terminator to s. A jump to c2, our outer conditional. retain
// any parameters the loop instruction passed to c1
builder.setInsertInto(s);
List<IRInst*> as;
for(UInt i = 0; i < loop->getArgCount(); ++i)
as.add(loop->getArg(i));
builder.emitBranch(c2, as.getCount(), as.getBuffer());
// We now have
// s: ...1, goto c2
// c2: if x then goto e2 else goto d (merge at b)
// e2: goto b
// l: loop break=b next=c1
// e4: goto b2
// c1: if x then goto e1 else goto d (merge at d)
// e1: goto b2
// d: goto c1
// b2: goto b
// b: ...2
// modify c2 to jump to the new loop
auto& c2dUse = newC2Terminator->getTrueBlock() == e2 ? newC2Terminator->falseBlock : newC2Terminator->trueBlock;
c2dUse.set(l);
// We now have
// s: ...1, goto c2
// c2: if x then goto e2 else goto l (merge at b)
// e2: goto b
// l: loop break=b next=c1
// e4: goto b2
// c1: if x then goto e1 else goto d (merge at d)
// e1: goto b2
// d: goto c1
// b2: goto b
// b: ...2
//
// Now we can modify the loop to jump to the block after the first
// conditional, d, as we know that it won't break out of the loop on the
// first iteration
//
// Since we're only here if the continue block is irrelevant (either the
// target block already or has a single predecessor) we can set it to the
// loop header.
//
// Beyond just retargeting the loop instruction, we need to make sure any
// parameters the loop instruction is passing to c1 are instead passed to
// 'd', and because we've added parameters to 'd' we need to forward them
// from c1 also which we will accomplish using a new block, e3,
//
loop->block.set(d);
loop->breakBlock.set(b2);
// TODO: This really upsets a few later passes, why isn't it ok to do given
// our "irrelevant continue" condition?
// loop->continueBlock.set(loop->getTargetBlock());
SLANG_ASSERT(d->getFirstParam() == nullptr);
c1->insertBefore(b2);
e1->insertAfter(c1);
List<IRInst*> ps;
for(const auto p : c1->getParams())
ps.add(p);
builder.setInsertInto(d);
for(const auto p : ps)
{
const auto q = duplicateToParamWithDecorations(builder, cloneEnv, p);
// Replace all uses, except for those in c1 and e1
List<IRUse*> uses;
traverseUses(p, [&](IRUse* u){if(u->user->getParent() != c1 && u->user->getParent() != e1) uses.add(u);});
for(auto u : uses)
u->set(q);
}
const auto e3 = builder.emitBlock();
e3->insertAfter(c1);
b2->insertBefore(b);
e4->insertAfter(c1);
builder.emitBranch(d, ps.getCount(), ps.getBuffer());
c1dUse.set(e3);
c1Terminator->afterBlock.set(e3);
// We now have the desired output
// s: ...1, goto c2
// c2: if x then goto e2 else goto l (merge at b)
// e2: goto b
// l: loop break=b2 next=d
// d: goto c1
// e4: goto b2
// c1: if x then goto e1 else goto e3 (merge at e3)
// e3: goto d
// e1: goto b2
// b2: goto b
// b: ...2
}
bool invertLoops(IRModule* module)
{
IRBuilder builder(module);
List<IRLoop*> toInvert;
overAllBlocks(module, [&](auto b){
if(auto loop = as<IRLoop>(b->getTerminator()))
{
if(isSuitableForInversion(loop))
toInvert.add(loop);
}
});
for(const auto loop : toInvert)
invertLoop(builder, loop);
return toInvert.getCount() > 0;
}
} // namespace Slang