From 8b0d7634743965948234b666c77393d4dd8535d7 Mon Sep 17 00:00:00 2001 From: Alex Zhikhartsev Date: Wed, 25 May 2022 19:10:57 -0400 Subject: [PATCH] [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 --- llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp | 130 +++++++++++---------- .../dfa-jump-threading-analysis.ll | 65 +++++++++++ llvm/test/Transforms/DFAJumpThreading/negative.ll | 49 ++++++++ 3 files changed, 181 insertions(+), 63 deletions(-) diff --git a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp index a715df5..d9254f1 100644 --- a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp +++ b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp @@ -103,6 +103,11 @@ static cl::opt MaxPathLength( cl::desc("Max number of blocks searched to find a threading path"), cl::Hidden, cl::init(20)); +static cl::opt MaxNumPaths( + "dfa-max-num-paths", + cl::desc("Max number of paths enumerated around a switch"), + cl::Hidden, cl::init(200)); + static cl::opt 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 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 Q; SmallSet SeenValues; SelectInsts.clear(); - Value *FirstDef = SI->getOperand(0); - auto *Inst = dyn_cast(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(SICond)) return false; - // Require the first definition to be a PHINode - if (!isa(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(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(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(SelI->user_back())) SelectInsts.push_back(SelectInstToUnfold(SelI, SelIUse)); + } else if (isa(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 &SeenValues) { - if (SeenValues.contains(InpVal)) - return true; - - if (isa(InpVal)) - return true; - - // If this is a function argument or another non-instruction, then give up. - if (!isa(InpVal)) - return false; - - return true; - } - - void addInstToQueue(Value *Val, std::deque &Q, - SmallSet &SeenValues) { + void addToQueue(Value *Val, std::deque &Q, + SmallSet &SeenValues) { if (SeenValues.contains(Val)) return; - if (Instruction *I = dyn_cast(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 LoopBBs; + for (const PathType &Path : LoopPaths) { + for (BasicBlock *BB : Path) + LoopBBs.insert(BB); + } + Value *FirstDef = Switch->getOperand(0); - assert(isa(FirstDef) && "After select unfolding, all state " - "definitions are expected to be phi " - "nodes."); + assert(isa(FirstDef) && "The first definition must be a phi."); SmallVector Stack; Stack.push_back(dyn_cast(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(Incoming) || - SeenValues.contains(Incoming)) { + SeenValues.contains(Incoming) || IsOutsideLoops) { continue; } - assert(isa(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(Incoming)) + return StateDefMap(); Stack.push_back(cast(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()) diff --git a/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-analysis.ll b/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-analysis.ll index ccc9b38..537fc9e 100644 --- a/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-analysis.ll +++ b/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-analysis.ll @@ -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 +} diff --git a/llvm/test/Transforms/DFAJumpThreading/negative.ll b/llvm/test/Transforms/DFAJumpThreading/negative.ll index 009ffa3..bd521d7 100644 --- a/llvm/test/Transforms/DFAJumpThreading/negative.ll +++ b/llvm/test/Transforms/DFAJumpThreading/negative.ll @@ -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 +} -- 2.7.4