diff --git a/llvm/include/llvm/Analysis/ValueTracking.h b/llvm/include/llvm/Analysis/ValueTracking.h --- a/llvm/include/llvm/Analysis/ValueTracking.h +++ b/llvm/include/llvm/Analysis/ValueTracking.h @@ -572,28 +572,31 @@ /// if, for all i, r is evaluated to poison or op raises UB if vi = poison. /// To filter out operands that raise UB on poison, you can use /// getGuaranteedNonPoisonOp. - bool propagatesPoison(const Instruction *I); + bool propagatesPoison(const Operator *I); /// Return either nullptr or an operand of I such that I will trigger - /// undefined behavior if I is executed and that operand has a poison - /// value. + /// undefined behavior if I is executed and that operand has either a poison + /// or a fully undef value. + const Value *getGuaranteedNonUndefOrPoisonOp(const Instruction *I); const Value *getGuaranteedNonPoisonOp(const Instruction *I); /// Return true if the given instruction must trigger undefined behavior. - /// when I is executed with any operands which appear in KnownPoison holding - /// a poison value at the point of execution. + /// when I is executed with any operands which appear in KnownUndefOrPoison + /// holding a poison or fully undef value at the point of execution. bool mustTriggerUB(const Instruction *I, - const SmallSet& KnownPoison); + const SmallSet& KnownUndefOrPoison); - /// Return true if this function can prove that if PoisonI is executed - /// and yields a poison value, then that will trigger undefined behavior. + /// Return true if this function can prove that if Inst is executed + /// and yields a poison or a fully undef value, then that will trigger + /// undefined behavior. /// /// Note that this currently only considers the basic block that is - /// the parent of I. - bool programUndefinedIfPoison(const Instruction *PoisonI); + /// the parent of Inst. + bool programUndefinedIfUndefOrPoison(const Instruction *Inst); + bool programUndefinedIfPoison(const Instruction *Inst); - /// canCreateUndefOrPoison returns true if Op can create undef or poison from - /// non-undef & non-poison operands. + /// canCreateUndefOrPoison returns true if Op can create undef bits or poison + /// from non-undef & non-poison operands. /// For vectors, canCreateUndefOrPoison returns true if there is potential /// poison or undef in any element of the result when vectors without /// undef/poison poison are given as operands. @@ -606,9 +609,9 @@ bool canCreateUndefOrPoison(const Operator *Op); bool canCreatePoison(const Operator *Op); - /// Return true if this function can prove that V is never undef value - /// or poison value. If V is an aggregate value or vector, check whether all - /// elements (except padding) are not undef or poison. + /// Return true if this function can prove that V does not have undef bits or + /// and is never poison. If V is an aggregate value or vector, check whether + /// all elements (except padding) are not undef or poison. /// Note that this is different from canCreateUndefOrPoison because the /// function assumes Op's operands are not poison/undef. /// @@ -619,6 +622,10 @@ const Instruction *CtxI = nullptr, const DominatorTree *DT = nullptr, unsigned Depth = 0); + bool isGuaranteedNotToBePoison(const Value *V, + const Instruction *CtxI = nullptr, + const DominatorTree *DT = nullptr, + unsigned Depth = 0); /// Specific patterns of select instructions we can match. enum SelectPatternFlavor { diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp --- a/llvm/lib/Analysis/ScalarEvolution.cpp +++ b/llvm/lib/Analysis/ScalarEvolution.cpp @@ -5913,7 +5913,7 @@ const Instruction *Poison = PoisonStack.pop_back_val(); for (auto *PoisonUser : Poison->users()) { - if (propagatesPoison(cast(PoisonUser))) { + if (propagatesPoison(cast(PoisonUser))) { if (Pushed.insert(cast(PoisonUser)).second) PoisonStack.push_back(cast(PoisonUser)); } else if (auto *BI = dyn_cast(PoisonUser)) { diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp --- a/llvm/lib/Analysis/ValueTracking.cpp +++ b/llvm/lib/Analysis/ValueTracking.cpp @@ -4772,10 +4772,14 @@ return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/true); } -bool llvm::isGuaranteedNotToBeUndefOrPoison(const Value *V, - const Instruction *CtxI, - const DominatorTree *DT, - unsigned Depth) { +static bool programUndefinedIfUndefOrPoison(const Instruction *Inst, + bool PoisonOnly); + +static bool isGuaranteedNotToBeUndefOrPoison(const Value *V, + const Instruction *CtxI, + const DominatorTree *DT, + unsigned Depth, + bool PoisonOnly) { if (Depth >= MaxDepth) return false; @@ -4786,14 +4790,15 @@ if (auto *C = dyn_cast(V)) { if (isa(C)) - return false; + return PoisonOnly; if (isa(C) || isa(C) || isa(V) || isa(C) || isa(C)) return true; if (C->getType()->isVectorTy()) - return !C->containsUndefElement() && !C->containsConstantExpression(); + return (PoisonOnly || !C->containsUndefElement()) && + !C->containsConstantExpression(); } // Strip cast operations from a pointer value. @@ -4810,7 +4815,7 @@ return true; auto OpCheck = [&](const Value *V) { - return isGuaranteedNotToBeUndefOrPoison(V, CtxI, DT, Depth + 1); + return isGuaranteedNotToBeUndefOrPoison(V, CtxI, DT, Depth + 1, PoisonOnly); }; if (auto *Opr = dyn_cast(V)) { @@ -4829,9 +4834,9 @@ } if (auto *I = dyn_cast(V)) { - if (programUndefinedIfPoison(I) && I->getType()->isIntegerTy(1)) - // Note: once we have an agreement that poison is a value-wise concept, - // we can remove the isIntegerTy(1) constraint. + // If V is undef, it should be i1 type, otherwise it can be partially undef + if ((PoisonOnly || I->getType()->isIntegerTy(1)) && + programUndefinedIfUndefOrPoison(I, PoisonOnly)) return true; } @@ -4853,12 +4858,23 @@ while (Dominator) { auto *TI = Dominator->getBlock()->getTerminator(); + Value *Cond = nullptr; if (auto BI = dyn_cast(TI)) { - if (BI->isConditional() && BI->getCondition() == V) - return true; + if (BI->isConditional()) + Cond = BI->getCondition(); } else if (auto SI = dyn_cast(TI)) { - if (SI->getCondition() == V) + Cond = SI->getCondition(); + } + + if (Cond) { + if (Cond == V) return true; + else if (PoisonOnly && isa(Cond)) { + auto *Opr = cast(Cond); + if (propagatesPoison(Opr) && + any_of(Opr->operand_values(), [&](Value *Op) { return Op == V; })) + return true; + } } Dominator = Dominator->getIDom(); @@ -4867,6 +4883,18 @@ return false; } +bool llvm::isGuaranteedNotToBeUndefOrPoison(const Value *V, + const Instruction *CtxI, + const DominatorTree *DT, + unsigned Depth) { + return ::isGuaranteedNotToBeUndefOrPoison(V, CtxI, DT, Depth, false); +} + +bool llvm::isGuaranteedNotToBePoison(const Value *V, const Instruction *CtxI, + const DominatorTree *DT, unsigned Depth) { + return ::isGuaranteedNotToBeUndefOrPoison(V, CtxI, DT, Depth, true); +} + OverflowResult llvm::computeOverflowForSignedAdd(const AddOperator *Add, const DataLayout &DL, AssumptionCache *AC, @@ -4960,7 +4988,7 @@ llvm_unreachable("Instruction not contained in its own parent basic block."); } -bool llvm::propagatesPoison(const Instruction *I) { +bool llvm::propagatesPoison(const Operator *I) { switch (I->getOpcode()) { case Instruction::Freeze: case Instruction::Select: @@ -4981,7 +5009,7 @@ } } -const Value *llvm::getGuaranteedNonPoisonOp(const Instruction *I) { +const Value *llvm::getGuaranteedNonUndefOrPoisonOp(const Instruction *I) { switch (I->getOpcode()) { case Instruction::Store: return cast(I)->getPointerOperand(); @@ -5017,48 +5045,57 @@ } } +const Value *llvm::getGuaranteedNonPoisonOp(const Instruction *I) { + // If passing undef to a specific operand of an instruction is UB, passing + // poison is also UB. + return getGuaranteedNonUndefOrPoisonOp(I); +} + bool llvm::mustTriggerUB(const Instruction *I, - const SmallSet& KnownPoison) { - auto *NotPoison = getGuaranteedNonPoisonOp(I); - return (NotPoison && KnownPoison.count(NotPoison)); + const SmallSet& KnownUndefOrPoison) { + auto *NotUndefOrPoison = getGuaranteedNonUndefOrPoisonOp(I); + return (NotUndefOrPoison && KnownUndefOrPoison.count(NotUndefOrPoison)); } -bool llvm::programUndefinedIfPoison(const Instruction *PoisonI) { +// Checks whether a program raises UB if Inst returned undef or poison. +// If PoisonOnly is true, it is assumed that Inst returned poison. +static bool programUndefinedIfUndefOrPoison(const Instruction *Inst, + bool PoisonOnly) { // We currently only look for uses of poison values within the same basic // block, as that makes it easier to guarantee that the uses will be - // executed given that PoisonI is executed. + // executed given that Inst is executed. // // FIXME: Expand this to consider uses beyond the same basic block. To do // this, look out for the distinction between post-dominance and strong // post-dominance. - const BasicBlock *BB = PoisonI->getParent(); + const BasicBlock *BB = Inst->getParent(); - // Set of instructions that we have proved will yield poison if PoisonI + // Set of instructions that we have proved will yield undef or poison if Inst // does. - SmallSet YieldsPoison; + SmallSet YieldsUndefOrPoison; SmallSet Visited; - YieldsPoison.insert(PoisonI); - Visited.insert(PoisonI->getParent()); + YieldsUndefOrPoison.insert(Inst); + Visited.insert(Inst->getParent()); - BasicBlock::const_iterator Begin = PoisonI->getIterator(), End = BB->end(); + BasicBlock::const_iterator Begin = Inst->getIterator(), End = BB->end(); unsigned Iter = 0; while (Iter++ < MaxDepth) { for (auto &I : make_range(Begin, End)) { - if (&I != PoisonI) { - if (mustTriggerUB(&I, YieldsPoison)) + if (&I != Inst) { + if (mustTriggerUB(&I, YieldsUndefOrPoison)) return true; if (!isGuaranteedToTransferExecutionToSuccessor(&I)) return false; } // Mark poison that propagates from I through uses of I. - if (YieldsPoison.count(&I)) { + if (PoisonOnly && YieldsUndefOrPoison.count(&I)) { for (const User *User : I.users()) { const Instruction *UserI = cast(User); - if (propagatesPoison(UserI)) - YieldsPoison.insert(User); + if (propagatesPoison(cast(UserI))) + YieldsUndefOrPoison.insert(User); } } } @@ -5077,6 +5114,14 @@ return false; } +bool llvm::programUndefinedIfUndefOrPoison(const Instruction *Inst) { + return ::programUndefinedIfUndefOrPoison(Inst, false); +} + +bool llvm::programUndefinedIfPoison(const Instruction *Inst) { + return ::programUndefinedIfUndefOrPoison(Inst, true); +} + static bool isKnownNonNaN(const Value *V, FastMathFlags FMF) { if (FMF.noNaNs()) return true; diff --git a/llvm/lib/Transforms/Instrumentation/PoisonChecking.cpp b/llvm/lib/Transforms/Instrumentation/PoisonChecking.cpp --- a/llvm/lib/Transforms/Instrumentation/PoisonChecking.cpp +++ b/llvm/lib/Transforms/Instrumentation/PoisonChecking.cpp @@ -293,7 +293,7 @@ } SmallVector Checks; - if (propagatesPoison(&I)) + if (propagatesPoison(cast(&I))) for (Value *V : I.operands()) Checks.push_back(getPoisonFor(ValToPoison, V)); diff --git a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp --- a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp +++ b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp @@ -1824,7 +1824,7 @@ // If we can't analyze propagation through this instruction, just skip it // and transitive users. Safe as false is a conservative result. - if (!propagatesPoison(I) && I != Root) + if (!propagatesPoison(cast(I)) && I != Root) continue; if (KnownPoison.insert(I).second) diff --git a/llvm/unittests/Analysis/ValueTrackingTest.cpp b/llvm/unittests/Analysis/ValueTrackingTest.cpp --- a/llvm/unittests/Analysis/ValueTrackingTest.cpp +++ b/llvm/unittests/Analysis/ValueTrackingTest.cpp @@ -10,6 +10,7 @@ #include "llvm/Analysis/AssumptionCache.h" #include "llvm/AsmParser/Parser.h" #include "llvm/IR/ConstantRange.h" +#include "llvm/IR/Dominators.h" #include "llvm/IR/Function.h" #include "llvm/IR/InstIterator.h" #include "llvm/IR/Instructions.h" @@ -716,12 +717,60 @@ for (auto &I : BB) { if (isa(&I)) break; - EXPECT_EQ(propagatesPoison(&I), Data[Index].first) + EXPECT_EQ(propagatesPoison(cast(&I)), Data[Index].first) << "Incorrect answer at instruction " << Index << " = " << I; Index++; } } +TEST_F(ValueTrackingTest, programUndefinedIfPoison) { + parseAssembly( + "declare i32 @any_num()" + "define void @test(i32 %mask) {\n" + " %A = call i32 @any_num()\n" + " %B = or i32 %A, %mask\n" + " udiv i32 1, %B" + " ret void\n" + "}\n"); + // If %A was poison, udiv raises UB regardless of %mask's avlue + EXPECT_EQ(programUndefinedIfPoison(A), true); +} + +TEST_F(ValueTrackingTest, programUndefinedIfUndefOrPoison) { + parseAssembly( + "declare i32 @any_num()" + "define void @test(i32 %mask) {\n" + " %A = call i32 @any_num()\n" + " %B = or i32 %A, %mask\n" + " udiv i32 1, %B" + " ret void\n" + "}\n"); + // If %A was undef and %mask was 1, udiv does not raise UB + EXPECT_EQ(programUndefinedIfUndefOrPoison(A), false); +} + +TEST_F(ValueTrackingTest, isGuaranteedNotToBePoison_exploitBranchCond) { + parseAssembly( + "declare i1 @any_bool()" + "define void @test(i1 %y) {\n" + " %A = call i1 @any_bool()\n" + " %cond = and i1 %A, %y\n" + " br i1 %cond, label %BB1, label %BB2\n" + "BB1:\n" + " ret void\n" + "BB2:\n" + " ret void\n" + "}\n"); + DominatorTree DT(*F); + for (auto &BB : *F) { + if (&BB == &F->getEntryBlock()) + continue; + + EXPECT_EQ(isGuaranteedNotToBePoison(A, BB.getTerminator(), &DT), true) + << "isGuaranteedNotToBePoison does not hold at " << *BB.getTerminator(); + } +} + TEST(ValueTracking, canCreatePoisonOrUndef) { std::string AsmHead = "declare i32 @g(i32)\n"