[ValueTracking] Add canCreateUndefOrPoison & let canCreatePoison use Operator

This patch
- adds `canCreateUndefOrPoison`
- refactors `canCreatePoison` so it can deal with constantexprs

`canCreateUndefOrPoison` will be used at D83926.

Reviewed By: nikic, jdoerfert

Differential Revision: https://reviews.llvm.org/D84007

include/llvm/Analysis/ValueTracking.h

@@ -21,6 +21,7 @@
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/Operator.h"
 #include <cassert>
 #include <cstdint>
 
@@ -591,18 +592,25 @@ class Value;
   /// the parent of I.
   bool programUndefinedIfPoison(const Instruction *PoisonI);
 
-  /// Return true if I can create poison from non-poison operands.
-  /// For vectors, canCreatePoison returns true if there is potential poison in
-  /// any element of the result when vectors without poison are given as
+  /// 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.
-  /// For example, given `I = shl <2 x i32> %x, <0, 32>`, this function returns
-  /// true. If I raises immediate UB but never creates poison (e.g. sdiv I, 0),
-  /// canCreatePoison returns false.
-  bool canCreatePoison(const Instruction *I);
+  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.
-  //
+  /// 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.

lib/Analysis/ValueTracking.cpp

@@ -4665,31 +4665,30 @@ bool llvm::isOverflowIntrinsicNoWrap(const WithOverflowInst *WO,
   return llvm::any_of(GuardingBranches, AllUsesGuardedByBranch);
 }
 
-bool llvm::canCreatePoison(const Instruction *I) {
+static bool canCreateUndefOrPoison(const Operator *Op, bool PoisonOnly) {
   // See whether I has flags that may create poison
-  if (isa<OverflowingBinaryOperator>(I) &&
-      (I->hasNoSignedWrap() || I->hasNoUnsignedWrap()))
-    return true;
-  if (isa<PossiblyExactOperator>(I) && I->isExact())
-    return true;
-  if (auto *FP = dyn_cast<FPMathOperator>(I)) {
+  if (const auto *OvOp = dyn_cast<OverflowingBinaryOperator>(Op)) {
+    if (OvOp->hasNoSignedWrap() || OvOp->hasNoUnsignedWrap())
+      return true;
+  }
+  if (const auto *ExactOp = dyn_cast<PossiblyExactOperator>(Op))
+    if (ExactOp->isExact())
+      return true;
+  if (const auto *FP = dyn_cast<FPMathOperator>(Op)) {
     auto FMF = FP->getFastMathFlags();
     if (FMF.noNaNs() || FMF.noInfs())
       return true;
   }
-  if (auto *GEP = dyn_cast<GetElementPtrInst>(I))
-    if (GEP->isInBounds())
-      return true;
 
-  unsigned Opcode = I->getOpcode();
+  unsigned Opcode = Op->getOpcode();
 
-  // Check whether opcode is a poison-generating operation
+  // Check whether opcode is a poison/undef-generating operation
   switch (Opcode) {
   case Instruction::Shl:
   case Instruction::AShr:
   case Instruction::LShr: {
     // Shifts return poison if shiftwidth is larger than the bitwidth.
-    if (auto *C = dyn_cast<Constant>(I->getOperand(1))) {
+    if (auto *C = dyn_cast<Constant>(Op->getOperand(1))) {
       SmallVector<Constant *, 4> ShiftAmounts;
       if (auto *FVTy = dyn_cast<FixedVectorType>(C->getType())) {
         unsigned NumElts = FVTy->getNumElements();
@@ -4715,41 +4714,62 @@ bool llvm::canCreatePoison(const Instruction *I) {
     return true;
   case Instruction::Call:
   case Instruction::CallBr:
-  case Instruction::Invoke:
-    // Function calls can return a poison value even if args are non-poison
-    // values.
-    return true;
+  case Instruction::Invoke: {
+    const auto *CB = cast<CallBase>(Op);
+    return !CB->hasRetAttr(Attribute::NoUndef);
+  }
   case Instruction::InsertElement:
   case Instruction::ExtractElement: {
     // If index exceeds the length of the vector, it returns poison
-    auto *VTy = cast<VectorType>(I->getOperand(0)->getType());
-    unsigned IdxOp = I->getOpcode() == Instruction::InsertElement ? 2 : 1;
-    auto *Idx = dyn_cast<ConstantInt>(I->getOperand(IdxOp));
+    auto *VTy = cast<VectorType>(Op->getOperand(0)->getType());
+    unsigned IdxOp = Op->getOpcode() == Instruction::InsertElement ? 2 : 1;
+    auto *Idx = dyn_cast<ConstantInt>(Op->getOperand(IdxOp));
     if (!Idx || Idx->getZExtValue() >= VTy->getElementCount().Min)
       return true;
     return false;
   }
+  case Instruction::ShuffleVector: {
+    // shufflevector may return undef.
+    if (PoisonOnly)
+      return false;
+    ArrayRef<int> Mask = isa<ConstantExpr>(Op)
+                             ? cast<ConstantExpr>(Op)->getShuffleMask()
+                             : cast<ShuffleVectorInst>(Op)->getShuffleMask();
+    return any_of(Mask, [](int Elt) { return Elt == UndefMaskElem; });
+  }
   case Instruction::FNeg:
   case Instruction::PHI:
   case Instruction::Select:
   case Instruction::URem:
   case Instruction::SRem:
-  case Instruction::ShuffleVector:
   case Instruction::ExtractValue:
   case Instruction::InsertValue:
   case Instruction::Freeze:
   case Instruction::ICmp:
   case Instruction::FCmp:
-  case Instruction::GetElementPtr:
     return false;
-  default:
-    if (isa<CastInst>(I))
+  case Instruction::GetElementPtr: {
+    const auto *GEP = cast<GEPOperator>(Op);
+    return GEP->isInBounds();
+  }
+  default: {
+    const auto *CE = dyn_cast<ConstantExpr>(Op);
+    if (isa<CastInst>(Op) || (CE && CE->isCast()))
       return false;
-    else if (isa<BinaryOperator>(I))
+    else if (isa<BinaryOperator>(Op))
       return false;
     // Be conservative and return true.
     return true;
   }
+  }
+}
+
+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,

lib/Transforms/Instrumentation/PoisonChecking.cpp

@@ -297,7 +297,7 @@ static bool rewrite(Function &F) {
         for (Value *V : I.operands())
           Checks.push_back(getPoisonFor(ValToPoison, V));
 
-      if (canCreatePoison(&I))
+      if (canCreatePoison(cast<Operator>(&I)))
         generateCreationChecks(I, Checks);
       ValToPoison[&I] = buildOrChain(B, Checks);
     }

unittests/Analysis/ValueTrackingTest.cpp

@@ -722,61 +722,71 @@ TEST(ValueTracking, propagatesPoison) {
   }
 }
 
-TEST(ValueTracking, canCreatePoison) {
+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?, IR instruction)
-  SmallVector<std::pair<bool, std::string>, 32> Data = {
-      {false, "add i32 %x, %y"},
-      {true, "add nsw nuw i32 %x, %y"},
-      {true, "shl i32 %x, %y"},
-      {true, "shl <4 x i32> %vx, %vx2"},
-      {true, "shl nsw i32 %x, %y"},
-      {true, "shl nsw <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 3>"},
-      {false, "shl i32 %x, 31"},
-      {true, "shl i32 %x, 32"},
-      {false, "shl <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 3>"},
-      {true, "shl <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 32>"},
-      {true, "ashr i32 %x, %y"},
-      {true, "ashr exact i32 %x, %y"},
-      {false, "ashr i32 %x, 31"},
-      {true, "ashr exact i32 %x, 31"},
-      {false, "ashr <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 3>"},
-      {true, "ashr <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 32>"},
-      {true, "ashr exact <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 3>"},
-      {true, "lshr i32 %x, %y"},
-      {true, "lshr exact i32 %x, 31"},
-      {false, "udiv i32 %x, %y"},
-      {true, "udiv exact i32 %x, %y"},
-      {false, "getelementptr i8, i8* %p, i32 %x"},
-      {true, "getelementptr inbounds i8, i8* %p, i32 %x"},
-      {true, "fneg nnan float %fx"},
-      {false, "fneg float %fx"},
-      {false, "fadd float %fx, %fy"},
-      {true, "fadd nnan float %fx, %fy"},
-      {false, "urem i32 %x, %y"},
-      {true, "fptoui float %fx to i32"},
-      {true, "fptosi float %fx to i32"},
-      {false, "bitcast float %fx to i32"},
-      {false, "select i1 %cond, i32 %x, i32 %y"},
-      {true, "select nnan i1 %cond, float %fx, float %fy"},
-      {true, "extractelement <4 x i32> %vx, i32 %x"},
-      {false, "extractelement <4 x i32> %vx, i32 3"},
-      {true, "extractelement <vscale x 4 x i32> %svx, i32 4"},
-      {true, "insertelement <4 x i32> %vx, i32 %x, i32 %y"},
-      {false, "insertelement <4 x i32> %vx, i32 %x, i32 3"},
-      {true, "insertelement <vscale x 4 x i32> %svx, i32 %x, i32 4"},
-      {false, "freeze i32 %x"},
-      {true, "call i32 @g(i32 %x)"},
-      {true, "fcmp nnan oeq float %fx, %fy"},
-      {false, "fcmp oeq float %fx, %fy"}};
+  // (can create poison?, can create undef?, IR instruction)
+  SmallVector<std::tuple<bool, bool, std::string>, 32> Data = {
+      {false, false, "add i32 %x, %y"},
+      {true, false, "add nsw nuw i32 %x, %y"},
+      {true, false, "shl i32 %x, %y"},
+      {true, false, "shl <4 x i32> %vx, %vx2"},
+      {true, false, "shl nsw i32 %x, %y"},
+      {true, false, "shl nsw <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 3>"},
+      {false, false, "shl i32 %x, 31"},
+      {true, false, "shl i32 %x, 32"},
+      {false, false, "shl <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 3>"},
+      {true, false, "shl <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 32>"},
+      {true, false, "ashr i32 %x, %y"},
+      {true, false, "ashr exact i32 %x, %y"},
+      {false, false, "ashr i32 %x, 31"},
+      {true, false, "ashr exact i32 %x, 31"},
+      {false, false, "ashr <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 3>"},
+      {true, false, "ashr <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 32>"},
+      {true, false, "ashr exact <4 x i32> %vx, <i32 0, i32 1, i32 2, i32 3>"},
+      {true, false, "lshr i32 %x, %y"},
+      {true, false, "lshr exact i32 %x, 31"},
+      {false, false, "udiv i32 %x, %y"},
+      {true, false, "udiv exact i32 %x, %y"},
+      {false, false, "getelementptr i8, i8* %p, i32 %x"},
+      {true, false, "getelementptr inbounds i8, i8* %p, i32 %x"},
+      {true, false, "fneg nnan float %fx"},
+      {false, false, "fneg float %fx"},
+      {false, false, "fadd float %fx, %fy"},
+      {true, false, "fadd nnan float %fx, %fy"},
+      {false, false, "urem i32 %x, %y"},
+      {true, false, "fptoui float %fx to i32"},
+      {true, false, "fptosi float %fx to i32"},
+      {false, false, "bitcast float %fx to i32"},
+      {false, false, "select i1 %cond, i32 %x, i32 %y"},
+      {true, false, "select nnan i1 %cond, float %fx, float %fy"},
+      {true, false, "extractelement <4 x i32> %vx, i32 %x"},
+      {false, false, "extractelement <4 x i32> %vx, i32 3"},
+      {true, false, "extractelement <vscale x 4 x i32> %svx, i32 4"},
+      {true, false, "insertelement <4 x i32> %vx, i32 %x, i32 %y"},
+      {false, false, "insertelement <4 x i32> %vx, i32 %x, i32 3"},
+      {true, false, "insertelement <vscale x 4 x i32> %svx, i32 %x, i32 4"},
+      {false, false, "freeze i32 %x"},
+      {false, false,
+       "shufflevector <4 x i32> %vx, <4 x i32> %vx2, "
+       "<4 x i32> <i32 0, i32 1, i32 2, i32 3>"},
+      {false, true,
+       "shufflevector <4 x i32> %vx, <4 x i32> %vx2, "
+       "<4 x i32> <i32 0, i32 1, i32 2, i32 undef>"},
+      {false, true,
+       "shufflevector <vscale x 4 x i32> %svx, "
+       "<vscale x 4 x i32> %svx, <vscale x 4 x i32> undef"},
+      {true, false, "call i32 @g(i32 %x)"},
+      {false, false, "call noundef i32 @g(i32 %x)"},
+      {true, false, "fcmp nnan oeq float %fx, %fy"},
+      {false, false, "fcmp oeq float %fx, %fy"}};
 
   std::string AssemblyStr = AsmHead;
   for (auto &Itm : Data)
-    AssemblyStr += Itm.second + "\n";
+    AssemblyStr += std::get<2>(Itm) + "\n";
   AssemblyStr += AsmTail;
 
   LLVMContext Context;
@@ -793,8 +803,14 @@ TEST(ValueTracking, canCreatePoison) {
   for (auto &I : BB) {
     if (isa<ReturnInst>(&I))
       break;
-    EXPECT_EQ(canCreatePoison(&I), Data[Index].first)
-        << "Incorrect answer at instruction " << Index << " = " << I;
+    bool Poison = std::get<0>(Data[Index]);
+    bool Undef = std::get<1>(Data[Index]);
+    EXPECT_EQ(canCreatePoison(cast<Operator>(&I)), Poison)
+        << "Incorrect answer of canCreatePoison at instruction " << Index
+        << " = " << I;
+    EXPECT_EQ(canCreateUndefOrPoison(cast<Operator>(&I)), Undef || Poison)
+        << "Incorrect answer of canCreateUndef at instruction " << Index
+        << " = " << I;
     Index++;
   }
 }