[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

 //===- llvm/Analysis/ValueTracking.h - Walk computations --------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
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 pushing freeze into operands of I is okay.
+  /// For example, If I = `add x, y`, rewriting `freeze(add x, y)` into
+  /// `add (freeze x), (freeze y)` is valid, but it isn't okay if the add had
+  /// nsw flag because it can make the result poison.
+  bool canPushFreezeIntoOperands(const Instruction *I);
+
+  /// Return true if V is poison given that ValAssumedPoison is already poison.
+  ///
+  /// If CtxI and DT are specified this method performs flow-sensitive analysis
+  bool isPoisonIf(const Value *V, const Value *ValAssumedPoison,
+                  const Instruction *CtxI = nullptr,
+                  const DominatorTree *DT = nullptr);
+
   /// 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

 //===- ValueTracking.cpp - Walk computations to compute properties --------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
auto AllUsesGuardedByBranch = [&](const BranchInst *BI) {
     }
 
     return true;
   };
 
   return llvm::any_of(GuardingBranches, AllUsesGuardedByBranch);
 }
 
+bool llvm::canPushFreezeIntoOperands(const Instruction *I) {
+  switch (I->getOpcode()) {
+  case Instruction::Add:
+  case Instruction::Sub:
+  case Instruction::Mul:
+    return !I->hasNoUnsignedWrap() && !I->hasNoSignedWrap();
+  case Instruction::UDiv:
+  case Instruction::SDiv:
+  case Instruction::URem:
+  case Instruction::SRem:
+    return !I->isExact();
+  case Instruction::And:
+  case Instruction::Or:
+  case Instruction::Xor:
+  case Instruction::ICmp:
+    return true;
+  case Instruction::FNeg:
+  case Instruction::FAdd:
+  case Instruction::FSub:
+  case Instruction::FMul:
+  case Instruction::FDiv:
+  case Instruction::FRem:
+  case Instruction::FCmp: {
+    auto FMF = I->getFastMathFlags();
+    return !FMF.noNaNs() && !FMF.noInfs();
+  }
+  case Instruction::Shl:
+  case Instruction::AShr:
+  case Instruction::LShr:
+    // For shift operations, it is poison if shiftwidth is larger than the
+    // bitwidth
+    return false;
+  case Instruction::PHI:
+  case Instruction::Select:
+  case Instruction::BitCast:
+  case Instruction::Freeze:
+    return true;
+  case Instruction::GetElementPtr:
+    return !dyn_cast<GetElementPtrInst>(I)->isInBounds();
+  default:
+    // Be conservative here.
+    return false;
+  }
+}
+
+bool llvm::isPoisonIf(const Value *V, const Value *ValAssumedPoison,
+                      const Instruction *CtxI, const DominatorTree *DT) {
+  if (isGuaranteedNotToBeUndefOrPoison(ValAssumedPoison, CtxI, DT))
+    // The assumption is false, so we can say anything
+    return true;
+
+  // Construct a set of values which are known to be poison from
+  // ValAssumedPoison.
+  SmallSet<const Value *, 4> PoisonValues;
+  PoisonValues.insert(ValAssumedPoison);
+  const Instruction *PoisonI = dyn_cast<Instruction>(ValAssumedPoison);
+  unsigned Depth = 0;
+
+  while (PoisonI && Depth < MaxDepth) {
+    // We'd like to know whether an operand of PoisonI is also poison.
+    if (!canPushFreezeIntoOperands(PoisonI))
+      // PoisonI can be a poison-generating instruction, so don't look further
+      break;
+
+    const Value *NextVal = nullptr;
+    bool MoreThanOneCandidate = false;
+    for (auto &Op : PoisonI->operands()) {
+      if (!isGuaranteedNotToBeUndefOrPoison(Op.get(), CtxI, DT)) {
+        // Op can be poison.
+        if (NextVal) {
+          // There is more than one operand that can make PoisonI poison.
+          MoreThanOneCandidate = true;
+          break;
+        }
+        NextVal = Op.get();
+      }
+    }
+
+    if (NextVal == nullptr) {
+      // All operands are non-poison, so PoisonI cannot be poison!
+      return true;
+    } else if (MoreThanOneCandidate)
+      break;
+
+    Depth++;
+    PoisonValues.insert(NextVal);
+    PoisonI = dyn_cast<Instruction>(NextVal);
+  }
+
+  if (PoisonValues.count(V) != 0)
+    // V should be poison.
+    return true;
+  
+  // Let's look one level further, by seeing its arguments if I was an
+  // instruction.
+  auto *I = dyn_cast<Instruction>(V);
+  if (I && propagatesFullPoison(I)) {
+    for (auto &Op : I->operands())
+      if (PoisonValues.count(Op.get()))
+        return true;
+  }
+
+  return false;
+}
+
+
 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

 //===- 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
 //
 //===----------------------------------------------------------------------===//
 
void parseAssembly(StringRef Assembly) {
     M = parseModule(Assembly);
     ASSERT_TRUE(M);
 
     Function *F = M->getFunction("test");
     ASSERT_TRUE(F) << "Test must have a function @test";
     if (!F)
       return;
 
-    A = nullptr;
+    A = B = nullptr;
     for (inst_iterator I = inst_begin(F), E = inst_end(F); I != E; ++I) {
       if (I->hasName()) {
         if (I->getName() == "A")
           A = &*I;
+        else if (I->getName() == "B")
+          B = &*I;
       }
     }
     ASSERT_TRUE(A) << "@test must have an instruction %A";
   }
 
   LLVMContext Context;
   std::unique_ptr<Module> M;
   Instruction *A = nullptr;
+  Instruction *B = nullptr;
 };
 
 class MatchSelectPatternTest : public ValueTrackingTest {
 protected:
   void expectPattern(const SelectPatternResult &P) {
     Value *LHS, *RHS;
     Instruction::CastOps CastOp;
     SelectPatternResult R = matchSelectPattern(A, LHS, RHS, &CastOp);
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_F(ValueTrackingTest, IsPoisonIfTest_Identity) {
+  parseAssembly(
+      "define void @test(i32 %x, i32 %y) {\n"
+      "  %A = add i32 %x, %y\n"
+      "  ret void\n"
+      "}");
+  EXPECT_EQ(isPoisonIf(A, A), 1u);
+}
+
+TEST_F(ValueTrackingTest, IsPoisonIfTest_ICmp) {
+  parseAssembly(
+      "define void @test(i32 %x) {\n"
+      "  %B = icmp eq i32 %x, 0\n"
+      "  %A = icmp eq i32 %x, 1\n"
+      "  ret void\n"
+      "}");
+  EXPECT_EQ(isPoisonIf(A, B), 1u);
+}
+
+TEST_F(ValueTrackingTest, IsPoisonIfTest_ICmpUnknown) {
+  parseAssembly(
+      "define void @test(i32 %x, i32 %y) {\n"
+      "  %B = icmp eq i32 %x, %y\n"
+      "  %A = icmp eq i32 %x, 1\n"
+      "  ret void\n"
+      "}");
+  EXPECT_EQ(isPoisonIf(A, B), 0u);
+}
+
+TEST_F(ValueTrackingTest, IsPoisonIfTest_AddNswOkay) {
+  parseAssembly(
+      "define void @test(i32 %x) {\n"
+      "  %B = add nsw i32 %x, 1\n"
+      "  %A = add i32 %B, 1\n"
+      "  ret void\n"
+      "}");
+  EXPECT_EQ(isPoisonIf(A, B), 1u);
+}
+
+TEST_F(ValueTrackingTest, IsPoisonIfTest_AddNswOkay2) {
+  parseAssembly(
+      "define void @test(i32 %x) {\n"
+      "  %B = add i32 %x, 1\n"
+      "  %A = add nsw i32 %B, 1\n"
+      "  ret void\n"
+      "}");
+  EXPECT_EQ(isPoisonIf(A, B), 1u);
+}
+
+TEST_F(ValueTrackingTest, IsPoisonIfTest_AddNsw) {
+  parseAssembly(
+      "define void @test(i32 %x) {\n"
+      "  %B = add nsw i32 %x, 1\n"
+      "  %A = add i32 %x, 1\n"
+      "  ret void\n"
+      "}");
+  EXPECT_EQ(isPoisonIf(A, B), 0u);
+}
+
+TEST(ValueTracking, CanPushFreezeIntoOperands) {
+  StringRef Assembly =
+      "define void @f(i32 %x, i32 %y, float %fx, float %fy, i1 %cond) { "
+      "  add i32 %x, %y\n"
+      "  add nsw nuw i32 %x, %y\n"
+      "  shl i32 %x, %y\n"
+      "  udiv i32 %x, %y\n"
+      "  udiv exact i32 %x, %y\n"
+      "  fadd float %fx, %fy\n"
+      "  fadd nnan float %fx, %fy\n"
+      "  select i1 %cond, i32 %x, i32 %y\n"
+      "  ret void "
+      "} ";
+
+  LLVMContext Context;
+  SMDiagnostic Error;
+  auto M = parseAssemblyString(Assembly, Error, Context);
+  assert(M && "Bad assembly?");
+
+  auto *F = M->getFunction("f");
+  assert(F && "Bad assembly?");
+
+  auto &BB = F->getEntryBlock();
+  bool ExpectedAnswers[] = {
+      true,  // add
+      false, // add nsw nuw
+      false, // shl
+      true,  // udiv
+      false, // udiv exact
+      true,  // fadd
+      false, // fadd nnan
+      true,  // select
+  };
+
+  int Index = 0;
+  for (auto &I : BB) {
+    if (isa<ReturnInst>(&I))
+      break;
+    EXPECT_EQ(canPushFreezeIntoOperands(&I), ExpectedAnswers[Index])
+        << "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"