[ValueTracking] Implement impliesPoison

This PR adds impliesPoison(ValAssumedPoison, V) that returns true if V is
poison under the assumption that ValAssumedPoison is poison.

For example, impliesPoison('icmp X, 10', 'icmp X, Y') return true because
'icmp X, Y' is poison if 'icmp X, 10' is poison.

impliesPoison can be used for sound optimization of select, as discussed in
D77868.

Reviewed By: nikic

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

llvm/include/llvm/Analysis/ValueTracking.h

@@ -620,6 +620,11 @@ constexpr unsigned MaxAnalysisRecursionDepth = 6;
   bool canCreateUndefOrPoison(const Operator *Op);
   bool canCreatePoison(const Operator *Op);
 
+  /// Return true if V is poison given that ValAssumedPoison is already poison.
+  /// For example, if ValAssumedPoison is `icmp X, 10` and V is `icmp X, 5`,
+  /// impliesPoison returns true.
+  bool impliesPoison(const Value *ValAssumedPoison, const Value *V);
+
   /// 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.

llvm/lib/Analysis/ValueTracking.cpp

@@ -4809,6 +4809,64 @@ bool llvm::canCreatePoison(const Operator *Op) {
   return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/true);
 }
 
+bool llvm::impliesPoison(const Value *ValAssumedPoison, const Value *V) {
+  // Construct a set of values which are known to be poison from the knowledge
+  // that ValAssumedPoison is poison.
+  SmallPtrSet<const Value *, 4> PoisonValues;
+  PoisonValues.insert(ValAssumedPoison);
+  const Instruction *PoisonI = dyn_cast<Instruction>(ValAssumedPoison);
+  unsigned Depth = 0;
+  const unsigned MaxDepth = 2;
+
+  while (PoisonI && Depth < MaxDepth) {
+    // We'd like to know whether an operand of PoisonI is also poison.
+    if (canCreatePoison(cast<Operator>(PoisonI)))
+      // PoisonI can be a poison-generating instruction, so don't look further
+      break;
+
+    const Value *NextVal = nullptr;
+    bool MoreThanOneCandidate = false;
+    // See which operand can be poison
+    for (const auto &Op : PoisonI->operands()) {
+      if (!isGuaranteedNotToBeUndefOrPoison(Op.get())) {
+        // 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.
+      // Since assumption is false, return true
+      return true;
+    } else if (MoreThanOneCandidate)
+      break;
+
+    Depth++;
+    PoisonValues.insert(NextVal);
+    PoisonI = dyn_cast<Instruction>(NextVal);
+  }
+
+  if (PoisonValues.contains(V))
+    return true;
+
+  // Let's look one level further, by seeing its arguments if I was an
+  // instruction.
+  // This happens when I is e.g. 'icmp X, const' where X is in PoisonValues.
+  const auto *I = dyn_cast<Instruction>(V);
+  if (I && propagatesPoison(cast<Operator>(I))) {
+    for (const auto &Op : I->operands())
+      if (PoisonValues.count(Op.get()))
+        return true;
+  }
+
+  return false;
+}
+
 static bool programUndefinedIfUndefOrPoison(const Value *V,
                                             bool PoisonOnly);
 

llvm/unittests/Analysis/ValueTrackingTest.cpp

@@ -695,6 +695,59 @@ TEST_F(ValueTrackingTest, ComputeNumSignBits_Shuffle2) {
   EXPECT_EQ(ComputeNumSignBits(A, M->getDataLayout()), 1u);
 }
 
+TEST_F(ValueTrackingTest, impliesPoisonTest_Identity) {
+  parseAssembly("define void @test(i32 %x, i32 %y) {\n"
+                "  %A = add i32 %x, %y\n"
+                "  ret void\n"
+                "}");
+  EXPECT_TRUE(impliesPoison(A, A));
+}
+
+TEST_F(ValueTrackingTest, impliesPoisonTest_ICmp) {
+  parseAssembly("define void @test(i32 %x) {\n"
+                "  %A2 = icmp eq i32 %x, 0\n"
+                "  %A = icmp eq i32 %x, 1\n"
+                "  ret void\n"
+                "}");
+  EXPECT_TRUE(impliesPoison(A2, A));
+}
+
+TEST_F(ValueTrackingTest, impliesPoisonTest_ICmpUnknown) {
+  parseAssembly("define void @test(i32 %x, i32 %y) {\n"
+                "  %A2 = icmp eq i32 %x, %y\n"
+                "  %A = icmp eq i32 %x, 1\n"
+                "  ret void\n"
+                "}");
+  EXPECT_FALSE(impliesPoison(A2, A));
+}
+
+TEST_F(ValueTrackingTest, impliesPoisonTest_AddNswOkay) {
+  parseAssembly("define void @test(i32 %x) {\n"
+                "  %A2 = add nsw i32 %x, 1\n"
+                "  %A = add i32 %A2, 1\n"
+                "  ret void\n"
+                "}");
+  EXPECT_TRUE(impliesPoison(A2, A));
+}
+
+TEST_F(ValueTrackingTest, impliesPoisonTest_AddNswOkay2) {
+  parseAssembly("define void @test(i32 %x) {\n"
+                "  %A2 = add i32 %x, 1\n"
+                "  %A = add nsw i32 %A2, 1\n"
+                "  ret void\n"
+                "}");
+  EXPECT_TRUE(impliesPoison(A2, A));
+}
+
+TEST_F(ValueTrackingTest, impliesPoisonTest_AddNsw) {
+  parseAssembly("define void @test(i32 %x) {\n"
+                "  %A2 = add nsw i32 %x, 1\n"
+                "  %A = add i32 %x, 1\n"
+                "  ret void\n"
+                "}");
+  EXPECT_FALSE(impliesPoison(A2, A));
+}
+
 TEST_F(ValueTrackingTest, ComputeNumSignBits_Shuffle_Pointers) {
   parseAssembly(
       "define <2 x i32*> @test(<2 x i32*> %x) {\n"