[ValueTracking] Let getGuaranteedNonPoisonOp find multiple non-poison operands

This patch helps getGuaranteedNonPoisonOp find multiple non-poison operands.

Instead of special-casing llvm.assume, I think it is also a viable option to
add noundef to Intrinsics.td. If it makes sense, I'll make a patch for that.

Reviewed By: jdoerfert

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

Author: aqjune

llvm/include/llvm/Analysis/ValueTracking.h

@@ -584,10 +584,10 @@ constexpr unsigned MaxAnalysisRecursionDepth = 6;
   /// getGuaranteedNonPoisonOp.
   bool propagatesPoison(const Instruction *I);
 
-  /// Return either nullptr or an operand of I such that I will trigger
-  /// undefined behavior if I is executed and that operand has a poison
-  /// value.
-  const Value *getGuaranteedNonPoisonOp(const Instruction *I);
+  /// Insert operands of I into Ops such that I will trigger undefined behavior
+  /// if I is executed and that operand has a poison value.
+  void getGuaranteedNonPoisonOps(const Instruction *I,
+                                 SmallPtrSetImpl<const Value *> &Ops);
 
   /// Return true if the given instruction must trigger undefined behavior.
   /// when I is executed with any operands which appear in KnownPoison holding

llvm/lib/Analysis/ValueTracking.cpp

@@ -5065,46 +5065,69 @@ bool llvm::propagatesPoison(const Instruction *I) {
   }
 }
 
-const Value *llvm::getGuaranteedNonPoisonOp(const Instruction *I) {
+void llvm::getGuaranteedNonPoisonOps(const Instruction *I,
+                                     SmallPtrSetImpl<const Value *> &Operands) {
   switch (I->getOpcode()) {
     case Instruction::Store:
-      return cast<StoreInst>(I)->getPointerOperand();
+      Operands.insert(cast<StoreInst>(I)->getPointerOperand());
+      break;
 
     case Instruction::Load:
-      return cast<LoadInst>(I)->getPointerOperand();
+      Operands.insert(cast<LoadInst>(I)->getPointerOperand());
+      break;
 
     case Instruction::AtomicCmpXchg:
-      return cast<AtomicCmpXchgInst>(I)->getPointerOperand();
+      Operands.insert(cast<AtomicCmpXchgInst>(I)->getPointerOperand());
+      break;
 
     case Instruction::AtomicRMW:
-      return cast<AtomicRMWInst>(I)->getPointerOperand();
+      Operands.insert(cast<AtomicRMWInst>(I)->getPointerOperand());
+      break;
 
     case Instruction::UDiv:
     case Instruction::SDiv:
     case Instruction::URem:
     case Instruction::SRem:
-      return I->getOperand(1);
+      Operands.insert(I->getOperand(1));
+      break;
 
     case Instruction::Call:
+    case Instruction::Invoke: {
       if (auto *II = dyn_cast<IntrinsicInst>(I)) {
         switch (II->getIntrinsicID()) {
         case Intrinsic::assume:
-          return II->getArgOperand(0);
+          Operands.insert(II->getArgOperand(0));
+          break;
         default:
-          return nullptr;
+          break;
         }
       }
-      return nullptr;
+
+      const CallBase *CB = cast<CallBase>(I);
+      if (CB->isIndirectCall())
+        Operands.insert(CB->getCalledOperand());
+      for (unsigned i = 0; i < CB->arg_size(); ++i) {
+        if (CB->paramHasAttr(i, Attribute::NoUndef))
+          Operands.insert(CB->getArgOperand(i));
+      }
+      break;
+    }
 
     default:
-      return nullptr;
+      break;
   }
 }
 
 bool llvm::mustTriggerUB(const Instruction *I,
                          const SmallSet<const Value *, 16>& KnownPoison) {
-  auto *NotPoison = getGuaranteedNonPoisonOp(I);
-  return (NotPoison && KnownPoison.count(NotPoison));
+  SmallPtrSet<const Value *, 4> NonPoisonOps;
+  getGuaranteedNonPoisonOps(I, NonPoisonOps);
+
+  for (const auto *V : NonPoisonOps)
+    if (KnownPoison.count(V))
+      return true;
+
+  return false;
 }
 
 

llvm/lib/Transforms/Instrumentation/PoisonChecking.cpp

@@ -282,8 +282,10 @@ static bool rewrite(Function &F) {
 
       // Note: There are many more sources of documented UB, but this pass only
       // attempts to find UB triggered by propagation of poison.
-      if (Value *Op = const_cast<Value*>(getGuaranteedNonPoisonOp(&I)))
-        CreateAssertNot(B, getPoisonFor(ValToPoison, Op));
+      SmallPtrSet<const Value *, 4> NonPoisonOps;
+      getGuaranteedNonPoisonOps(&I, NonPoisonOps);
+      for (const Value *Op : NonPoisonOps)
+        CreateAssertNot(B, getPoisonFor(ValToPoison, const_cast<Value *>(Op)));
 
       if (LocalCheck)
         if (auto *RI = dyn_cast<ReturnInst>(&I))

llvm/test/Transforms/InstSimplify/freeze-noundef.ll

@@ -97,8 +97,7 @@ define i1 @used_by_fncall(i1 %x) {
 ; CHECK-LABEL: @used_by_fncall(
 ; CHECK-NEXT:    [[Y:%.*]] = add nsw i1 [[X:%.*]], true
 ; CHECK-NEXT:    call void @use_i1(i1 [[Y]])
-; CHECK-NEXT:    [[F:%.*]] = freeze i1 [[Y]]
-; CHECK-NEXT:    ret i1 [[F]]
+; CHECK-NEXT:    ret i1 [[Y]]
 ;
   %y = add nsw i1 %x, 1
   call void @use_i1(i1 %y)