[ValueTracking] Let getGuaranteedNonPoisonOp find multiple non-poison operands

[aqjune]

Summary

This patch helps getGuaranteedNonPoisonOp find multiple non-poison operands.

Another patch that adds noundef to Intrinsics.td will be uploaded.

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

5123–5130

We don't really care about duplicates in this context, so I'd suggest a more efficient SmallVector.

5128

This should be the other way around, i.e. iterate over NonPoisonOps and check if they are in KnownPoison. The NonPoisonOps set is generally expected to be smaller than the KnownPoison one.

[aqjune]

Iterate over NonPoisonOps

[aqjune]

Use SmallPtrSet

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

5123–5130

I think set is beneficial for its use at PoisonChecking - it allows introducing one check for one value.
Instead I changed its type to SmallPtrSet as it seems to be more efficient.

[jdoerfert]

LGTM.

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.

I think that makes sense.

llvm/include/llvm/Analysis/ValueTracking.h

590

Nit: SmallPtrSetImpl

llvm/lib/Analysis/ValueTracking.cpp

5111

The function pointer too, but unclear if that helps much ;)

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

5111

I think that's a valid point, I'll add it.

[aqjune]

Use SmallPtrSetImpl, consider function pointers

llvm/include/llvm/Analysis/ValueTracking.h

constexpr unsigned MaxAnalysisRecursionDepth = 6;
   /// Return true if I yields poison or raises UB if any of its operands is
   /// poison.
   /// Formally, given I = `r = op v1 v2 .. vN`, propagatesPoison returns true
   /// if, for all i, r is evaluated to poison or op raises UB if vi = poison.
   /// To filter out operands that raise UB on poison, you can use
   /// 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
   /// a poison value at the point of execution.
   bool mustTriggerUB(const Instruction *I,
                      const SmallSet<const Value *, 16>& KnownPoison);
 
   /// Return true if this function can prove that if PoisonI is executed

llvm/lib/Analysis/ValueTracking.cpp

default:
     if (isa<BinaryOperator>(I) || isa<UnaryOperator>(I) || isa<CastInst>(I))
       return true;
 
     // Be conservative and return false.
     return false;
   }
 }
 
-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;
 }
 
 
 bool llvm::programUndefinedIfPoison(const Instruction *PoisonI) {
   // We currently only look for uses of poison values within the same basic
   // block, as that makes it easier to guarantee that the uses will be
   // executed given that PoisonI is executed.
   //

llvm/lib/Transforms/Instrumentation/PoisonChecking.cpp

static bool rewrite(Function &F) {
   for (BasicBlock &BB : F)
     for (Instruction &I : BB) {
       if (isa<PHINode>(I)) continue;
 
       IRBuilder<> B(cast<Instruction>(&I));
 
       // 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))
           if (RI->getNumOperands() != 0) {
             Value *Op = RI->getOperand(0);
             CreateAssertNot(B, getPoisonFor(ValToPoison, Op));
           }
 

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

 }
 
 declare void @use_i1(i1 noundef)
 
 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)
   %f = freeze i1 %y
   ret i1 %f
 }