[ValueTracking] Implement canCreatePoison

Summary:
This PR adds `canCreatePoison(Instruction *I)` which returns true if `I` can generate poison from non-poison
operands.

Reviewers: spatel, nikic, lebedev.ri

Reviewed By: spatel

Subscribers: hiraditya, llvm-commits, regehr, nlopes

Tags: #llvm

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

llvm/include/llvm/Analysis/ValueTracking.h

@@ -592,6 +592,15 @@ class Value;
   /// the parent of I.
   bool programUndefinedIfFullPoison(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
+  /// 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);
+
   /// Return true if this function can prove that V is never undef value
   /// or poison value.
   //

llvm/lib/Analysis/ValueTracking.cpp

@@ -4592,6 +4592,91 @@ bool llvm::isOverflowIntrinsicNoWrap(const WithOverflowInst *WO,
   return llvm::any_of(GuardingBranches, AllUsesGuardedByBranch);
 }
 
+bool llvm::canCreatePoison(const Instruction *I) {
+  // 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)) {
+    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();
+
+  // Check whether opcode is a poison-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))) {
+      SmallVector<Constant *, 4> ShiftAmounts;
+      if (C->getType()->isVectorTy()) {
+        unsigned NumElts = cast<VectorType>(C->getType())->getNumElements();
+        for (unsigned i = 0; i < NumElts; ++i)
+          ShiftAmounts.push_back(C->getAggregateElement(i));
+      } else
+        ShiftAmounts.push_back(C);
+
+      bool Safe = llvm::all_of(ShiftAmounts, [](Constant *C) {
+        auto *CI = dyn_cast<ConstantInt>(C);
+        return CI && CI->getZExtValue() < C->getType()->getIntegerBitWidth();
+      });
+      return !Safe;
+    }
+    return true;
+  }
+  case Instruction::FPToSI:
+  case Instruction::FPToUI:
+    // fptosi/ui yields poison if the resulting value does not fit in the
+    // destination type.
+    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. CallBr returns poison when jumping to indirect labels.
+    return true;
+  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));
+    if (!Idx || Idx->getZExtValue() >= VTy->getElementCount().Min)
+      return true;
+    return false;
+  }
+  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))
+      return false;
+    else if (isa<BinaryOperator>(I))
+      return false;
+    // Be conservative and return true.
+    return true;
+  }
+}
+
 bool llvm::isGuaranteedNotToBeUndefOrPoison(const Value *V,
                                             const Instruction *CtxI,
                                             const DominatorTree *DT) {

llvm/unittests/Analysis/ValueTrackingTest.cpp

@@ -645,7 +645,7 @@ TEST_F(ValueTrackingTest, ComputeNumSignBits_PR32045) {
   EXPECT_EQ(ComputeNumSignBits(A, M->getDataLayout()), 1u);
 }
 
-// No guarantees for canonical IR in this analysis, so this just bails out. 
+// No guarantees for canonical IR in this analysis, so this just bails out.
 TEST_F(ValueTrackingTest, ComputeNumSignBits_Shuffle) {
   parseAssembly(
       "define <2 x i32> @test() {\n"
@@ -656,7 +656,7 @@ TEST_F(ValueTrackingTest, ComputeNumSignBits_Shuffle) {
 }
 
 // No guarantees for canonical IR in this analysis, so a shuffle element that
-// references an undef value means this can't return any extra information. 
+// references an undef value means this can't return any extra information.
 TEST_F(ValueTrackingTest, ComputeNumSignBits_Shuffle2) {
   parseAssembly(
       "define <2 x i32> @test(<2 x i1> %x) {\n"
@@ -667,6 +667,83 @@ TEST_F(ValueTrackingTest, ComputeNumSignBits_Shuffle2) {
   EXPECT_EQ(ComputeNumSignBits(A, M->getDataLayout()), 1u);
 }
 
+TEST(ValueTracking, canCreatePoison) {
+  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"}};
+
+  std::string AssemblyStr = AsmHead;
+  for (auto &Itm : Data)
+    AssemblyStr += Itm.second + "\n";
+  AssemblyStr += AsmTail;
+
+  LLVMContext Context;
+  SMDiagnostic Error;
+  auto M = parseAssemblyString(AssemblyStr, Error, Context);
+  assert(M && "Bad assembly?");
+
+  auto *F = M->getFunction("f");
+  assert(F && "Bad assembly?");
+
+  auto &BB = F->getEntryBlock();
+
+  int Index = 0;
+  for (auto &I : BB) {
+    if (isa<ReturnInst>(&I))
+      break;
+    EXPECT_EQ(canCreatePoison(&I), Data[Index].first)
+        << "Incorrect answer at instruction " << Index << " = " << I;
+    Index++;
+  }
+}
+
 TEST_F(ComputeKnownBitsTest, ComputeKnownBits) {
   parseAssembly(
       "define i32 @test(i32 %a, i32 %b) {\n"