[ValueTracking] Implement canCreatePoison

[aqjune]

Summary

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

[spatel]

llvm/lib/Analysis/ValueTracking.cpp

4603–4604

URem/SRem are not PossiblyExactOperator, so I think this will assert. Please add at least 1 test to confirm.

Would it be better to create a preliminary "canCreatePoison()" helper function that is structured like Instruction::andIRFlags() and return 'false' on any of those instructions?

[aqjune]

Renamed canPushFreezeIntoOperands to canCreatePoison (merged two rather than splitting a helper function, because the two seems to have no difference at this moment).
Used the idiom from Instruction::andIRFlags().
Also renamed isPoisonIf to impliesPoison because the latter seems more formal.

[aqjune]

clang-format

[spatel]

I have enough questions about canCreatePoison() alone that I have not looked at impliesPoison() yet.
That might be a sign that we should split the review into 2 patches, but I think it's ok either way.

llvm/include/llvm/Analysis/ValueTracking.h

596

typo: returnd

597

Add an example of the vector case in this code comment? Is it possible with the current code?

llvm/lib/Analysis/ValueTracking.cpp

4600–4602

We could refine this, right? Ie, if the shift amount is a constant that is less than bitwidth, then shifts can not produce poison (unless they have wrapping/exact flags).
I don't know if that's the only case like this, so it's fine to make a TODO comment for a future patch if that's correct.

4619

plain "cast" rather than "dyn_cast"?

4643

I find it confusing to switch on opcode, check flags, then switch on opcode again. Can we check flags first, and then switch on opcode just once?

4646–4647

Immediate undefined behavior (eg, sdiv %x, 0) is explicitly not poison?
If correct, please note that in the documentation comment for this function, so there is no confusion.
(There's existing code called "MightCreatePoisonOrUB" in instcombine, but I'm not sure if it can be changed to use this new analysis function.)

[aqjune]

• Leave canCreatePoison only
• Leave the latter switch only (merged)
• Let shl/ashr/lshr deal with constant shift amount
• Address other comments, add more tests

[aqjune]

• Minor update to a comment

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

4646–4647

For MightCreatePoisonOrUB, we can use this function, but it can be broken if canCreatePoison uses deeper analysis.
You mean this optimization when M is given such that shufflevector can work as a vector select:

// shuffle (op X, C0), (op Y, C1), M --> op (shuffle X, Y, M), C'
// shuffle (op C0, X), (op C1, Y), M --> op C', (shuffle X, Y, M)

For example:

X = 0 & (freeze X0)
Y = 0 & (freeze Y0) // X and Y are 0
// canCreatePoison concluded that shl C0, X and shl C1, Y are safe
shuffle (shl C0, X), (shl C1, Y), undef
=>
shl C', (shufflevector X, Y, undef)
=>
shl C', undef

To make using canCreatePoison safe, M should not contain undef elements.

[spatel]

LGTM

llvm/include/llvm/Analysis/ValueTracking.h

596–597

This still doesn't read quite right to me. Should it be:
"For vectors, return true if there is potential poison in any element."

[aqjune]

For the impliesPoison one, I made a separate patch here: https://reviews.llvm.org/D78152
But I guess the only place where impliesPoison would be used right now will be the select optimization, then having it in ValueTracking might be too much, so I'm fine with putting it into InstCombine or somewhere else.

llvm/include/llvm/Analysis/ValueTracking.h

class Value;
   /// Return true if this function can prove that if PoisonI is executed
   /// and yields a full-poison value (all bits poison), then that will
   /// trigger undefined behavior.
   ///
   /// Note that this currently only considers the basic block that is
   /// 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.
   //
   /// 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,

llvm/lib/Analysis/ValueTracking.cpp

auto AllUsesGuardedByBranch = [&](const BranchInst *BI) {
     }
 
     return true;
   };
 
   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) {
   // If the value is a freeze instruction, then it can never
   // be undef or poison.
   if (isa<FreezeInst>(V))
     return true;
   // TODO: Some instructions are guaranteed to return neither undef

llvm/unittests/Analysis/ValueTrackingTest.cpp

TEST_F(ValueTrackingTest, ComputeNumSignBits_PR32045) {
   parseAssembly(
       "define i32 @test(i32 %a) {\n"
       "  %A = ashr i32 %a, -1\n"
       "  ret i32 %A\n"
       "}\n");
   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"
       "  %A = shufflevector <2 x i32> undef, <2 x i32> undef, <2 x i32> <i32 0, i32 0>\n"
       "  ret <2 x i32> %A\n"
       "}\n");
   EXPECT_EQ(ComputeNumSignBits(A, M->getDataLayout()), 1u);
 }
 
 // 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"
       "  %sext = sext <2 x i1> %x to <2 x i32>\n"
       "  %A = shufflevector <2 x i32> %sext, <2 x i32> undef, <2 x i32> <i32 0, i32 2>\n"
       "  ret <2 x i32> %A\n"
       "}\n");
   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"
       "  %ash = mul i32 %a, 8\n"
       "  %aad = add i32 %ash, 7\n"
       "  %aan = and i32 %aad, 4095\n"
       "  %bsh = shl i32 %b, 4\n"
       "  %bad = or i32 %bsh, 6\n"