[DFAJumpThreading] Relax analysis to handle unpredictable initial values

Responding to a feature request from the Rust community:

https://github.com/rust-lang/rust/issues/80630

    void foo(X) {
      for (...)
	switch (X)
	  case A
	    X = B
	  case B
	    X = C
    }

Even though the initial switch value is non-constant, the switch
statement can still be threaded: the initial value will hit the switch
statement but the rest of the state changes will proceed by jumping
unconditionally.

The early predictability check is relaxed to allow unpredictable values
anywhere, but later, after the paths through the switch statement have
been enumerated, no non-constant state values are allowed along the
paths. Any state value not along a path will be an initial switch value,
which can be safely ignored.

Differential Revision: https://reviews.llvm.org/D124394
This commit is contained in:
Alex Zhikhartsev 2022-05-25 19:10:57 -04:00
parent 9c66ed9b73
commit 8b0d763474
3 changed files with 181 additions and 63 deletions

View File

@ -103,6 +103,11 @@ static cl::opt<unsigned> MaxPathLength(
cl::desc("Max number of blocks searched to find a threading path"),
cl::Hidden, cl::init(20));
static cl::opt<unsigned> MaxNumPaths(
"dfa-max-num-paths",
cl::desc("Max number of paths enumerated around a switch"),
cl::Hidden, cl::init(200));
static cl::opt<unsigned>
CostThreshold("dfa-cost-threshold",
cl::desc("Maximum cost accepted for the transformation"),
@ -415,7 +420,7 @@ inline raw_ostream &operator<<(raw_ostream &OS, const ThreadingPath &TPath) {
struct MainSwitch {
MainSwitch(SwitchInst *SI, OptimizationRemarkEmitter *ORE) {
if (isPredictable(SI)) {
if (isCandidate(SI)) {
Instr = SI;
} else {
ORE->emit([&]() {
@ -433,83 +438,60 @@ struct MainSwitch {
}
private:
/// Do a use-def chain traversal. Make sure the value of the switch variable
/// is always a known constant. This means that all conditional jumps based on
/// switch variable can be converted to unconditional jumps.
bool isPredictable(const SwitchInst *SI) {
std::deque<Instruction *> Q;
/// Do a use-def chain traversal starting from the switch condition to see if
/// \p SI is a potential condidate.
///
/// Also, collect select instructions to unfold.
bool isCandidate(const SwitchInst *SI) {
std::deque<Value *> Q;
SmallSet<Value *, 16> SeenValues;
SelectInsts.clear();
Value *FirstDef = SI->getOperand(0);
auto *Inst = dyn_cast<Instruction>(FirstDef);
// If this is a function argument or another non-instruction, then give up.
// We are interested in loop local variables.
if (!Inst)
Value *SICond = SI->getCondition();
LLVM_DEBUG(dbgs() << "\tSICond: " << *SICond << "\n");
if (!isa<PHINode>(SICond))
return false;
// Require the first definition to be a PHINode
if (!isa<PHINode>(Inst))
return false;
LLVM_DEBUG(dbgs() << "\tisPredictable() FirstDef: " << *Inst << "\n");
Q.push_back(Inst);
SeenValues.insert(FirstDef);
addToQueue(SICond, Q, SeenValues);
while (!Q.empty()) {
Instruction *Current = Q.front();
Value *Current = Q.front();
Q.pop_front();
if (auto *Phi = dyn_cast<PHINode>(Current)) {
for (Value *Incoming : Phi->incoming_values()) {
if (!isPredictableValue(Incoming, SeenValues))
return false;
addInstToQueue(Incoming, Q, SeenValues);
addToQueue(Incoming, Q, SeenValues);
}
LLVM_DEBUG(dbgs() << "\tisPredictable() phi: " << *Phi << "\n");
LLVM_DEBUG(dbgs() << "\tphi: " << *Phi << "\n");
} else if (SelectInst *SelI = dyn_cast<SelectInst>(Current)) {
if (!isValidSelectInst(SelI))
return false;
if (!isPredictableValue(SelI->getTrueValue(), SeenValues) ||
!isPredictableValue(SelI->getFalseValue(), SeenValues)) {
return false;
}
addInstToQueue(SelI->getTrueValue(), Q, SeenValues);
addInstToQueue(SelI->getFalseValue(), Q, SeenValues);
LLVM_DEBUG(dbgs() << "\tisPredictable() select: " << *SelI << "\n");
addToQueue(SelI->getTrueValue(), Q, SeenValues);
addToQueue(SelI->getFalseValue(), Q, SeenValues);
LLVM_DEBUG(dbgs() << "\tselect: " << *SelI << "\n");
if (auto *SelIUse = dyn_cast<PHINode>(SelI->user_back()))
SelectInsts.push_back(SelectInstToUnfold(SelI, SelIUse));
} else if (isa<Constant>(Current)) {
LLVM_DEBUG(dbgs() << "\tconst: " << *Current << "\n");
continue;
} else {
// If it is neither a phi nor a select, then we give up.
return false;
LLVM_DEBUG(dbgs() << "\tother: " << *Current << "\n");
// Allow unpredictable values. The hope is that those will be the
// initial switch values that can be ignored (they will hit the
// unthreaded switch) but this assumption will get checked later after
// paths have been enumerated (in function getStateDefMap).
continue;
}
}
return true;
}
bool isPredictableValue(Value *InpVal, SmallSet<Value *, 16> &SeenValues) {
if (SeenValues.contains(InpVal))
return true;
if (isa<ConstantInt>(InpVal))
return true;
// If this is a function argument or another non-instruction, then give up.
if (!isa<Instruction>(InpVal))
return false;
return true;
}
void addInstToQueue(Value *Val, std::deque<Instruction *> &Q,
SmallSet<Value *, 16> &SeenValues) {
void addToQueue(Value *Val, std::deque<Value *> &Q,
SmallSet<Value *, 16> &SeenValues) {
if (SeenValues.contains(Val))
return;
if (Instruction *I = dyn_cast<Instruction>(Val))
Q.push_back(I);
Q.push_back(Val);
SeenValues.insert(Val);
}
@ -563,7 +545,16 @@ struct AllSwitchPaths {
void run() {
VisitedBlocks Visited;
PathsType LoopPaths = paths(SwitchBlock, Visited, /* PathDepth = */ 1);
StateDefMap StateDef = getStateDefMap();
StateDefMap StateDef = getStateDefMap(LoopPaths);
if (StateDef.empty()) {
ORE->emit([&]() {
return OptimizationRemarkMissed(DEBUG_TYPE, "SwitchNotPredictable",
Switch)
<< "Switch instruction is not predictable.";
});
return;
}
for (PathType Path : LoopPaths) {
ThreadingPath TPath;
@ -638,6 +629,9 @@ private:
PathType NewPath(Path);
NewPath.push_front(BB);
Res.push_back(NewPath);
if (Res.size() >= MaxNumPaths) {
return Res;
}
}
}
// This block could now be visited again from a different predecessor. Note
@ -648,14 +642,22 @@ private:
}
/// Walk the use-def chain and collect all the state-defining instructions.
StateDefMap getStateDefMap() const {
///
/// Return an empty map if unpredictable values encountered inside the basic
/// blocks of \p LoopPaths.
StateDefMap getStateDefMap(const PathsType &LoopPaths) const {
StateDefMap Res;
// Basic blocks belonging to any of the loops around the switch statement.
SmallPtrSet<BasicBlock *, 16> LoopBBs;
for (const PathType &Path : LoopPaths) {
for (BasicBlock *BB : Path)
LoopBBs.insert(BB);
}
Value *FirstDef = Switch->getOperand(0);
assert(isa<PHINode>(FirstDef) && "After select unfolding, all state "
"definitions are expected to be phi "
"nodes.");
assert(isa<PHINode>(FirstDef) && "The first definition must be a phi.");
SmallVector<PHINode *, 8> Stack;
Stack.push_back(dyn_cast<PHINode>(FirstDef));
@ -667,15 +669,17 @@ private:
Res[CurPhi->getParent()] = CurPhi;
SeenValues.insert(CurPhi);
for (Value *Incoming : CurPhi->incoming_values()) {
for (BasicBlock *IncomingBB : CurPhi->blocks()) {
Value *Incoming = CurPhi->getIncomingValueForBlock(IncomingBB);
bool IsOutsideLoops = LoopBBs.count(IncomingBB) == 0;
if (Incoming == FirstDef || isa<ConstantInt>(Incoming) ||
SeenValues.contains(Incoming)) {
SeenValues.contains(Incoming) || IsOutsideLoops) {
continue;
}
assert(isa<PHINode>(Incoming) && "After select unfolding, all state "
"definitions are expected to be phi "
"nodes.");
// Any unpredictable value inside the loops means we must bail out.
if (!isa<PHINode>(Incoming))
return StateDefMap();
Stack.push_back(cast<PHINode>(Incoming));
}
@ -1279,7 +1283,7 @@ bool DFAJumpThreading::run(Function &F) {
continue;
LLVM_DEBUG(dbgs() << "\nCheck if SwitchInst in BB " << BB.getName()
<< " is predictable\n");
<< " is a candidate\n");
MainSwitch Switch(SI, ORE);
if (!Switch.getInstr())

View File

@ -184,3 +184,68 @@ bb49: ; preds = %bb43, %bb43
bb66: ; preds = %bb59
ret i32 0
}
; Value %init is not predictable but it's okay since it is the value initial to the switch.
define i32 @initial.value.positive1(i32 %init) {
; CHECK: < loop.3 case2 > [ 3, loop.3 ]
; CHECK-NEXT: < loop.3 case2 loop.1.backedge loop.1 loop.2 > [ 1, loop.1 ]
; CHECK-NEXT: < loop.3 case2 loop.1.backedge si.unfold.false loop.1 loop.2 > [ 4, loop.1.backedge ]
; CHECK-NEXT: < loop.3 case3 loop.2.backedge loop.2 > [ 0, loop.2.backedge ]
; CHECK-NEXT: < loop.3 case3 case4 loop.2.backedge loop.2 > [ 3, loop.2.backedge ]
; CHECK-NEXT: < loop.3 case3 case4 loop.1.backedge loop.1 loop.2 > [ 1, loop.1 ]
; CHECK-NEXT: < loop.3 case3 case4 loop.1.backedge si.unfold.false loop.1 loop.2 > [ 2, loop.1.backedge ]
; CHECK-NEXT: < loop.3 case4 loop.2.backedge loop.2 > [ 3, loop.2.backedge ]
; CHECK-NEXT: < loop.3 case4 loop.1.backedge loop.1 loop.2 > [ 1, loop.1 ]
; CHECK-NEXT: < loop.3 case4 loop.1.backedge si.unfold.false loop.1 loop.2 > [ 2, loop.1.backedge ]
entry:
%cmp = icmp eq i32 %init, 0
br label %loop.1
loop.1:
%state.1 = phi i32 [ %init, %entry ], [ %state.1.be2, %loop.1.backedge ]
br label %loop.2
loop.2:
%state.2 = phi i32 [ %state.1, %loop.1 ], [ %state.2.be, %loop.2.backedge ]
br label %loop.3
loop.3:
%state = phi i32 [ %state.2, %loop.2 ], [ 3, %case2 ]
switch i32 %state, label %infloop.i [
i32 2, label %case2
i32 3, label %case3
i32 4, label %case4
i32 0, label %case0
i32 1, label %case1
]
case2:
br i1 %cmp, label %loop.3, label %loop.1.backedge
case3:
br i1 %cmp, label %loop.2.backedge, label %case4
case4:
br i1 %cmp, label %loop.2.backedge, label %loop.1.backedge
loop.1.backedge:
%state.1.be = phi i32 [ 2, %case4 ], [ 4, %case2 ]
%state.1.be2 = select i1 %cmp, i32 1, i32 %state.1.be
br label %loop.1
loop.2.backedge:
%state.2.be = phi i32 [ 3, %case4 ], [ 0, %case3 ]
br label %loop.2
case0:
br label %exit
case1:
br label %exit
infloop.i:
br label %infloop.i
exit:
ret i32 0
}

View File

@ -214,3 +214,52 @@ for.inc:
for.end:
ret i32 0
}
declare i32 @arbitrary_function()
; Don't confuse %state.2 for the initial switch value.
define i32 @negative6(i32 %init) {
; REMARK: SwitchNotPredictable
; REMARK-NEXT: negative6
entry:
%cmp = icmp eq i32 %init, 0
br label %loop.2
loop.2:
%state.2 = call i32 @arbitrary_function()
br label %loop.3
loop.3:
%state = phi i32 [ %state.2, %loop.2 ], [ 3, %case2 ]
switch i32 %state, label %infloop.i [
i32 2, label %case2
i32 3, label %case3
i32 4, label %case4
i32 0, label %case0
i32 1, label %case1
]
case2:
br label %loop.3
case3:
br i1 %cmp, label %loop.2.backedge, label %case4
case4:
br label %loop.2.backedge
loop.2.backedge:
br label %loop.2
case0:
br label %exit
case1:
br label %exit
infloop.i:
br label %infloop.i
exit:
ret i32 0
}