[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

4683

SmallPtrSet

4722

count -> contains

4750–4754

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

llvm/unittests/Analysis/ValueTrackingTest.cpp

678

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
 //
 //===----------------------------------------------------------------------===//
 //
class Value;
   /// If I is returning a vector value, it checks the existence of poison
   /// element. To be specific, canCreatePoison returns true if the result can
   /// contain poison when vectors without poison are given as operands.
   /// For example, given `I = shl <2 x i32> %x, <0, 32>`, this returns true.
   /// If I raises immediate UB but never creates poison (e.g. sdiv I, 0), this
   /// returns false.
   bool canCreatePoison(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 impliesPoison(const Value *ValAssumedPoison, const Value *V,
+                     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
 //
 //===----------------------------------------------------------------------===//
 //
if (isa<CastInst>(I))
       return false;
     else if (isa<BinaryOperator>(I))
       return false;
     // Be conservative and return true.
     return true;
   }
 }
 
+bool llvm::impliesPoison(const Value *ValAssumedPoison, const Value *V,
+                         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 (canCreatePoison(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
   // nor poison if their arguments are not poison/undef.
 
   if (auto *C = dyn_cast<Constant>(V)) {
     // TODO: We can analyze ConstExpr by opcode to determine if there is any
     //       possibility of poison.
     if (isa<UndefValue>(C) || isa<ConstantExpr>(C))
       return false;
 
     // TODO: Add ConstantFP and pointers.
     if (isa<ConstantInt>(C) || isa<GlobalVariable>(C) )
       return true;
 
     if (C->getType()->isVectorTy())
       return !C->containsUndefElement() && !C->containsConstantExpression();
 

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, impliesPoisonTest_Identity) {
+  parseAssembly("define void @test(i32 %x, i32 %y) {\n"
+                "  %A = add i32 %x, %y\n"
+                "  ret void\n"
+                "}");
+  EXPECT_EQ(impliesPoison(A, A), 1u);
+}
+
+TEST_F(ValueTrackingTest, impliesPoisonTest_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(impliesPoison(B, A), 1u);
+}
+
+TEST_F(ValueTrackingTest, impliesPoisonTest_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(impliesPoison(B, A), 0u);
+}
+
+TEST_F(ValueTrackingTest, impliesPoisonTest_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(impliesPoison(B, A), 1u);
+}
+
+TEST_F(ValueTrackingTest, impliesPoisonTest_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(impliesPoison(B, A), 1u);
+}
+
+TEST_F(ValueTrackingTest, impliesPoisonTest_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(impliesPoison(B, A), 0u);
+}
+
 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 = {