[ValueTracking] Add UndefOrPoison/Poison-only version of relevant functions

[aqjune]

Summary

This patch adds isGuaranteedNotToBePoison and programUndefinedIfUndefOrPoison.

isGuaranteedNotToBePoison will be used at D75808. The latter function is used at isGuaranteedNotToBePoison.

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

4918–4919

This is a fix to a bug that caused returning false more often. Since it was a small change, embedded here

[nikic]

llvm/include/llvm/Analysis/ValueTracking.h

589–592

As these functions are the same, I would just rename getGuaranteedNonPoisonOp to getGuaranteedNonUndefOrPoisonOp.

llvm/lib/Analysis/ValueTracking.cpp

4879

If PoisonOnly=true, does this actually end up returning true for UndefValue in all cases? It's not completely obvious to me. Maybe if (isa<UndefValue>(C)) return PoisonOnly;?

4887

As a followup, we may want to recursively check the elements instead.

4918–4919

I'd prefer to land this separately with the corresponding test case, otherwise it gets lost in the rest of the patch.

4926

To make sure I get the logic here: Undef is a bit-wise concept and programUndefinedIfUndefOrPoison() tells us whether we get UB for a full undef operand, while isGuaranteedNotToBeUndefOrPoison() tells us whether there can be any undef bits, so we can apply this only for the i1 case. Is that right?

[aqjune]

llvm/include/llvm/Analysis/ValueTracking.h

589–592

I think splitting these two is good (even if their contents are actually the same) because users like PoisonChecking want to only consider poison values. What do you think?

llvm/lib/Analysis/ValueTracking.cpp

4879

You're right, it shouldn't fall through. Will fix this

4887

Sounds good!

4918–4919

Okay, will make a separate patch for this

4926

Yes, exactly.

[aqjune]

Fix undef case

[aqjune]

clang-format

[aqjune]

Rebase,
Remove getGuaranteedNonUndefOrPoisonOp since getGuaranteedNonUndefOrPoisonOp and getGuaranteedNonPoisonOp are equivalent

[aqjune]

llvm/include/llvm/Analysis/ValueTracking.h

589–592

Leave getGuaranteedNonPoisonOp only & left a comment at previous getGuaranteedNonUndefOrPoisonOp's uses saying that non-poison op and non-undef op are equivalent

[aqjune]

ping

[jdoerfert]

I think I'm fine with this, though it's a lot. @nikic ?

llvm/lib/Transforms/Instrumentation/PoisonChecking.cpp

298

Interesting that we need these.

[aqjune]

Rebase, and have a few changes to maximally utilize noundef attributes.
programUndefinedIfUndefOrPoison(I) is updated in a way that it returns true
if I is partially undef.
For example, the updated programUndefinedIfUndefOrPoison(I) returns true in this case:

%I = or i32 x, 1 ; if x = undef, %I is partially undef
call void @f(i32 noundef %I)

[nikic]

LGTM

llvm/unittests/Analysis/ValueTrackingTest.cpp

734

avlue -> value

[aqjune]

Typos

llvm/include/llvm/Analysis/ValueTracking.h

bool isGuaranteedToExecuteForEveryIteration(const Instruction *I,
                                               const Loop *L);
 
   /// Return true if I yields poison or raises UB if any of its operands is
   /// poison.
   /// Formally, given I = `r = op v1 v2 .. vN`, propagatesPoison returns true
   /// 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);
 
   /// Insert operands of I into Ops such that I will trigger undefined behavior
   /// if I is executed and that operand has a poison value.
   void getGuaranteedNonPoisonOps(const Instruction *I,
                                  SmallPtrSetImpl<const Value *> &Ops);
 
-  /// Return true if the given instruction must trigger undefined behavior.
+  /// 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.
   bool mustTriggerUB(const Instruction *I,
                      const SmallSet<const Value *, 16>& KnownPoison);
 
-  /// 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 value or undef bits, 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.
   /// 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.
   /// For example, given `Op = shl <2 x i32> %x, <0, 32>`, this function returns
   /// true. If Op raises immediate UB but never creates poison or undef
   /// (e.g. sdiv I, 0), canCreatePoison returns false.
   ///
   /// canCreatePoison returns true if Op can create poison from non-poison
   /// operands.
   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
+  /// 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.
   ///
   /// If CtxI and DT are specified this method performs flow-sensitive analysis
   /// and returns true if it is guaranteed to be never undef or poison
   /// immediately before the CtxI.
   bool isGuaranteedNotToBeUndefOrPoison(const Value *V,
                                         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 {
     SPF_UNKNOWN = 0,
     SPF_SMIN,                   /// Signed minimum
     SPF_UMIN,                   /// Unsigned minimum
     SPF_SMAX,                   /// Signed maximum
     SPF_UMAX,                   /// Unsigned maximum

llvm/lib/Analysis/ScalarEvolution.cpp

bool ScalarEvolution::isAddRecNeverPoison(const Instruction *I, const Loop *L) {
   Pushed.insert(I);
   PoisonStack.push_back(I);
 
   bool LatchControlDependentOnPoison = false;
   while (!PoisonStack.empty() && !LatchControlDependentOnPoison) {
     const Instruction *Poison = PoisonStack.pop_back_val();
 
     for (auto *PoisonUser : Poison->users()) {
-      if (propagatesPoison(cast<Instruction>(PoisonUser))) {
+      if (propagatesPoison(cast<Operator>(PoisonUser))) {
         if (Pushed.insert(cast<Instruction>(PoisonUser)).second)
           PoisonStack.push_back(cast<Instruction>(PoisonUser));
       } else if (auto *BI = dyn_cast<BranchInst>(PoisonUser)) {
         assert(BI->isConditional() && "Only possibility!");
         if (BI->getParent() == LatchBB) {
           LatchControlDependentOnPoison = true;
           break;
         }

llvm/lib/Analysis/ValueTracking.cpp

 bool llvm::canCreateUndefOrPoison(const Operator *Op) {
   return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/false);
 }
 
 bool llvm::canCreatePoison(const Operator *Op) {
   return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/true);
 }
 
-bool llvm::isGuaranteedNotToBeUndefOrPoison(const Value *V,
+static bool programUndefinedIfUndefOrPoison(const Instruction *Inst,
+                                            bool PoisonOnly);
+
+static bool isGuaranteedNotToBeUndefOrPoison(const Value *V,
                                              const Instruction *CtxI,
                                              const DominatorTree *DT,
-                                            unsigned Depth) {
+                                             unsigned Depth, bool PoisonOnly) {
   if (Depth >= MaxAnalysisRecursionDepth)
     return false;
 
   if (const auto *A = dyn_cast<Argument>(V)) {
     if (A->hasAttribute(Attribute::NoUndef))
       return true;
   }
 
   if (auto *C = dyn_cast<Constant>(V)) {
     if (isa<UndefValue>(C))
-      return false;
+      return PoisonOnly;
 
     if (isa<ConstantInt>(C) || isa<GlobalVariable>(C) || isa<ConstantFP>(V) ||
         isa<ConstantPointerNull>(C) || isa<Function>(C))
       return true;
 
     if (C->getType()->isVectorTy() && !isa<ConstantExpr>(C))
-      return !C->containsConstantExpression() && !C->containsUndefElement();
+      return (PoisonOnly || !C->containsUndefElement()) &&
+             !C->containsConstantExpression();
   }
 
   // Strip cast operations from a pointer value.
   // Note that stripPointerCastsSameRepresentation can strip off getelementptr
   // inbounds with zero offset. To guarantee that the result isn't poison, the
   // stripped pointer is checked as it has to be pointing into an allocated
   // object or be null `null` to ensure `inbounds` getelement pointers with a
   // zero offset could not produce poison.
   // It can strip off addrspacecast that do not change bit representation as
   // well. We believe that such addrspacecast is equivalent to no-op.
   auto *StrippedV = V->stripPointerCastsSameRepresentation();
   if (isa<AllocaInst>(StrippedV) || isa<GlobalVariable>(StrippedV) ||
       isa<Function>(StrippedV) || isa<ConstantPointerNull>(StrippedV))
     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<Operator>(V)) {
     // If the value is a freeze instruction, then it can never
     // be undef or poison.
     if (isa<FreezeInst>(V))
       return true;
 
     if (const auto *CB = dyn_cast<CallBase>(V)) {
       if (CB->hasRetAttr(Attribute::NoUndef))
         return true;
     }
 
     if (!canCreateUndefOrPoison(Opr) && all_of(Opr->operands(), OpCheck))
       return true;
   }
 
   if (auto *I = dyn_cast<Instruction>(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 (programUndefinedIfUndefOrPoison(I, PoisonOnly))
       return true;
   }
 
   // CxtI may be null or a cloned instruction.
   if (!CtxI || !CtxI->getParent() || !DT)
     return false;
 
   auto *DNode = DT->getNode(CtxI->getParent());
   if (!DNode)
     // Unreachable block
     return false;
 
   // If V is used as a branch condition before reaching CtxI, V cannot be
   // undef or poison.
   //   br V, BB1, BB2
   // BB1:
   //   CtxI ; V cannot be undef or poison here
   auto *Dominator = DNode->getIDom();
   while (Dominator) {
     auto *TI = Dominator->getBlock()->getTerminator();
 
+    Value *Cond = nullptr;
     if (auto BI = dyn_cast<BranchInst>(TI)) {
-      if (BI->isConditional() && BI->getCondition() == V)
-        return true;
+      if (BI->isConditional())
+        Cond = BI->getCondition();
     } else if (auto SI = dyn_cast<SwitchInst>(TI)) {
-      if (SI->getCondition() == V)
+      Cond = SI->getCondition();
+    }
+
+    if (Cond) {
+      if (Cond == V)
         return true;
+      else if (PoisonOnly && isa<Operator>(Cond)) {
+        // For poison, we can analyze further
+        auto *Opr = cast<Operator>(Cond);
+        if (propagatesPoison(Opr) &&
+            any_of(Opr->operand_values(), [&](Value *Op) { return Op == V; }))
+          return true;
+      }
     }
 
     Dominator = Dominator->getIDom();
   }
 
   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,
                                                  const Instruction *CxtI,
                                                  const DominatorTree *DT) {
   return ::computeOverflowForSignedAdd(Add->getOperand(0), Add->getOperand(1),
                                        Add, DL, AC, CxtI, DT);
 }
bool llvm::isGuaranteedToExecuteForEveryIteration(const Instruction *I,
 
   for (const Instruction &LI : *L->getHeader()) {
     if (&LI == I) return true;
     if (!isGuaranteedToTransferExecutionToSuccessor(&LI)) return false;
   }
   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:
   case Instruction::PHI:
   case Instruction::Call:
   case Instruction::Invoke:
     return false;
   case Instruction::ICmp:
bool llvm::mustTriggerUB(const Instruction *I,
 
   for (const auto *V : NonPoisonOps)
     if (KnownPoison.count(V))
       return true;
 
   return false;
 }
 
-
-bool llvm::programUndefinedIfPoison(const Instruction *PoisonI) {
-  // We currently only look for uses of poison values within the same basic
+static bool programUndefinedIfUndefOrPoison(const Instruction *Inst,
+                                            bool PoisonOnly) {
+  // We currently only look for uses of 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
+  BasicBlock::const_iterator Begin = Inst->getIterator(), End = BB->end();
+
+  if (!PoisonOnly) {
+    // Be conservative & just check whether a value is passed to a noundef
+    // argument.
+    // Instructions that raise UB with a poison operand are well-defined
+    // or have unclear semantics when the input is partially undef.
+    // For example, 'udiv x, (undef | 1)' isn't UB.
+
+    for (auto &I : make_range(Begin, End)) {
+      if (const auto *CB = dyn_cast<CallBase>(&I)) {
+        for (unsigned i = 0; i < CB->arg_size(); ++i) {
+          if (CB->paramHasAttr(i, Attribute::NoUndef) &&
+              CB->getArgOperand(i) == Inst)
+            return true;
+        }
+      }
+      if (!isGuaranteedToTransferExecutionToSuccessor(&I))
+        break;
+    }
+    return false;
+  }
+
+  // Set of instructions that we have proved will yield poison if Inst
   // does.
   SmallSet<const Value *, 16> YieldsPoison;
   SmallSet<const BasicBlock *, 4> Visited;
-  YieldsPoison.insert(PoisonI);
-  Visited.insert(PoisonI->getParent());
-
-  BasicBlock::const_iterator Begin = PoisonI->getIterator(), End = BB->end();
+  YieldsPoison.insert(Inst);
+  Visited.insert(Inst->getParent());
 
   unsigned Iter = 0;
   while (Iter++ < MaxAnalysisRecursionDepth) {
     for (auto &I : make_range(Begin, End)) {
-      if (&I != PoisonI) {
+      if (&I != Inst) {
         if (mustTriggerUB(&I, YieldsPoison))
           return true;
         if (!isGuaranteedToTransferExecutionToSuccessor(&I))
           return false;
       }
 
       // Mark poison that propagates from I through uses of I.
       if (YieldsPoison.count(&I)) {
         for (const User *User : I.users()) {
           const Instruction *UserI = cast<Instruction>(User);
-          if (propagatesPoison(UserI))
+          if (propagatesPoison(cast<Operator>(UserI)))
             YieldsPoison.insert(User);
         }
       }
     }
 
     if (auto *NextBB = BB->getSingleSuccessor()) {
       if (Visited.insert(NextBB).second) {
         BB = NextBB;
         Begin = BB->getFirstNonPHI()->getIterator();
         End = BB->end();
         continue;
       }
     }
 
     break;
   }
   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;
 
   if (auto *C = dyn_cast<ConstantFP>(V))
     return !C->isNaN();
 
   if (auto *C = dyn_cast<ConstantDataVector>(V)) {

llvm/lib/Transforms/Instrumentation/PoisonChecking.cpp

for (Instruction &I : BB) {
       if (LocalCheck)
         if (auto *RI = dyn_cast<ReturnInst>(&I))
           if (RI->getNumOperands() != 0) {
             Value *Op = RI->getOperand(0);
             CreateAssertNot(B, getPoisonFor(ValToPoison, Op));
           }
 
       SmallVector<Value*, 4> Checks;
-      if (propagatesPoison(&I))
+      if (propagatesPoison(cast<Operator>(&I)))
         for (Value *V : I.operands())
           Checks.push_back(getPoisonFor(ValToPoison, V));
 
       if (canCreatePoison(cast<Operator>(&I)))
         generateCreationChecks(I, Checks);
       ValToPoison[&I] = buildOrChain(B, Checks);
     }
 

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

while (!Worklist.empty()) {
     const Instruction *I = Worklist.pop_back_val();
 
     // If we know this must trigger UB on a path leading our target.
     if (mustTriggerUB(I, KnownPoison) && DT->dominates(I, OnPathTo))
       return true;
 
     // 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<Operator>(I)) && I != Root)
       continue;
 
     if (KnownPoison.insert(I).second)
       for (const User *User : I->users())
         Worklist.push_back(cast<Instruction>(User));
   }
 
   // Might be non-UB, or might have a path we couldn't prove must execute on

llvm/unittests/Analysis/ValueTrackingTest.cpp

 //===- ValueTrackingTest.cpp - ValueTracking tests ------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/ValueTracking.h"
 #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"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/SourceMgr.h"
TEST(ValueTracking, propagatesPoison) {
   assert(F && "Bad assembly?");
 
   auto &BB = F->getEntryBlock();
 
   int Index = 0;
   for (auto &I : BB) {
     if (isa<ReturnInst>(&I))
       break;
-    EXPECT_EQ(propagatesPoison(&I), Data[Index].first)
+    EXPECT_EQ(propagatesPoison(cast<Operator>(&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 value
+  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"
       "define void @f(i32 %x, i32 %y, float %fx, float %fy, i1 %cond, "
       "<4 x i32> %vx, <4 x i32> %vx2, <vscale x 4 x i32> %svx, i8* %p) {\n";
   std::string AsmTail = "  ret void\n}";
   // (can create poison?, can create undef?, IR instruction)
   SmallVector<std::pair<std::pair<bool, bool>, std::string>, 32> Data = {