https://github.com/shader-slang/slang
Tip revision: 5902acdabc4445a65741a7a6a3a95f223e301059 authored by Yong He on 23 January 2024, 07:19:40 UTC
[LSP] Fetch configs directly from didConfigurationChanged message. (#3478)
[LSP] Fetch configs directly from didConfigurationChanged message. (#3478)
Tip revision: 5902acd
slang-ir-loop-inversion.cpp
#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=b2 next=d
// d: goto c1:
// c1: if x then goto e1 else goto e3 (merge at e3)
// e3: goto d
// e1: goto b2
// b2: 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 merge point for the outer condition
// b2 is the new break block for the loop
//
static void invertLoop(IRBuilder& builder, IRLoop* loop)
{
IRBuilderInsertLocScope builderScope(&builder);
builder.setInsertInto(loop->getParent());
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;
// We don't expect 'd' to have any parameters, because it used to be the
// target of a conditional branch
SLANG_ASSERT(d->getFirstParam() == nullptr);
// 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 (inside the loop, and b onwards) uses. Lift the values from c1
// used in b (and onwards) to parameters, the same for those used before b.
// To avoid a critical edge, pass these via a new block, e1.
// For any such values used within the loop we can pass directly to d.
//
// c1PostLoopParams are values form c1 used after the loop
List<IRInst*> c1PostLoopParams;
// c1LoopParams are values from c1 used within the loop itself
List<IRInst*> c1LoopParams;
for(auto i : IRInstList<IRInst>(c1->getFirstInst(), c1->getLastInst()))
{
IRParam* postLoopParam = nullptr;
IRParam* loopParam = nullptr;
traverseUses(i, [&](IRUse* u){
auto userBlock = u->getUser()->getParent();
if(domTree->dominates(b, userBlock))
{
// A new parameter to replace this 'i'
if(!postLoopParam)
{
postLoopParam = duplicateToParamWithDecorations(builder, cloneEnv, i);
b->addParam(postLoopParam);
}
u->set(postLoopParam);
}
else if(userBlock != c1)
{
// A new parameter to replace this 'i'
if(!loopParam)
{
loopParam = duplicateToParamWithDecorations(builder, cloneEnv, i);
d->addParam(loopParam);
}
u->set(loopParam);
}
});
if(postLoopParam)
c1PostLoopParams.add(i);
if(loopParam)
c1LoopParams.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.
//
const 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, c1PostLoopParams.getCount(), c1PostLoopParams.getBuffer());
c1bUse.set(e1);
c1Terminator->afterBlock.set(d);
// We'll rewrite existing breaks to jump to b2, but via an intermediate jump block to
// avoid a critical edge.
//
traverseUses(b, [&](IRUse* u){
auto userBlock = u->getUser()->getParent();
// Restrict this to just those blocks within this loop
if(userBlock != e1
&& userBlock != b2
&& userBlock != s
&& domTree->dominates(s, userBlock)
&& !domTree->dominates(b, userBlock))
{
const auto jumpToB2Block = builder.emitBlock();
jumpToB2Block->insertAfter(userBlock);
builder.emitBranch(b2, c1PostLoopParams.getCount(), c1PostLoopParams.getBuffer());
u->set(jumpToB2Block);
}
});
// We now have
// s: ...1 loop break=b next=c1
// 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
// 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);
// 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
// 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
// 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
// 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 during the back-edge, which we will accomplish using a new
// block, e3,
//
// Utilize the cloneenv to make sure that when entering the loop we use
// c1's instructions as cloned into c2
builder.setInsertAfter(loop);
List<IRInst*> loopEntryArgs;
for(const auto p : c1LoopParams)
loopEntryArgs.add(cloneInst(&cloneEnv, &builder, p));
for(UInt i = 0; i < loop->getArgCount(); ++i)
loopEntryArgs.add(loop->getArg(i));
const auto newLoop = builder.emitLoop(d, b2, loop->getContinueBlock(), loopEntryArgs.getCount(), loopEntryArgs.getBuffer());
newLoop->sourceLoc = loop->sourceLoc;
loop->transferDecorationsTo(newLoop);
loop->removeAndDeallocate();
// 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());
c1->insertBefore(b2);
e1->insertAfter(c1);
builder.setInsertInto(d);
// Add the necessary parameters for the loop state to d, first the
// paramters for instructions in the duplicated conditional block, then the
// ones from the loop body.
List<IRInst*> ps = c1LoopParams;
for(const auto p : c1->getParams())
{
ps.add(p);
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);
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
// 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