[ValueTracking] Implement impliesPoison

[aqjune]

Summary

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.

[aqjune]

Remove DT / CtxI as passing them to isGuaranteedNotToBeUndefOrPoison increases compile time

[aqjune]

Make the function lightweight

[nikic]

LGTM with some nits. Please also fix the clang-tidy warnings wrt const.

There's various ways this could be improved, but I guess we should start simple and see which useful cases slip by.

llvm/lib/Analysis/ValueTracking.cpp

4806

SmallPtrSet

4809

You probably want to leave these two functions next to the main canCreateUndefOrPoison() definition.

4845

count -> contains

llvm/unittests/Analysis/ValueTrackingTest.cpp

703

Why 1u and not true? Or for that matter, EXPECT_TRUE()?

[aqjune]

update

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
 //
 //===----------------------------------------------------------------------===//
 //
constexpr unsigned MaxAnalysisRecursionDepth = 6;
   /// 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.
   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.
   /// 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

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
 //
 //===----------------------------------------------------------------------===//
 //
default: {
     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::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);
 
 static bool isGuaranteedNotToBeUndefOrPoison(const Value *V,
                                              AssumptionCache *AC,
                                              const Instruction *CtxI,
                                              const DominatorTree *DT,
                                              unsigned Depth, bool PoisonOnly) {

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
 //
 //===----------------------------------------------------------------------===//
 
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, 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"
       "  %A = shufflevector <2 x i32*> zeroinitializer, <2 x i32*> undef, <2 x i32> zeroinitializer\n"
       "  ret <2 x i32*> %A\n"
       "}\n");
   EXPECT_EQ(ComputeNumSignBits(A, M->getDataLayout()), 64u);
 }