[ValueTracking] Check uses of Argument if it is given to isGuaranteedNotToBeUndefOrPoison

[aqjune]

Summary

This is a patch that allows isGuaranteedNotToBeUndefOrPoison to return more precise result
when an argument is given, by looking through its uses at the entry block (and following blocks as well, if it is checking poison only).

This is useful when there is a function call with noundef arguments at the entry block.

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

4889

I would prefer to move the programUndefinedIfUndefOrPoison() call below outside the Instruction check, and just return false from programUndefinedIfUndefOrPoison() on non-Instruction/Argument (aka Constant).

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

4889

Would it make programUndefinedIfUndefOrPoison() differentiate non-undef and undef constants?
We'll want to make programUndefinedIfUndefOrPoison(undef) return false.

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

4889

I don't think it's useful to make a distinction, we can just return false for all constants.

[aqjune]

Address comments

[nikic]

LG with some nits

llvm/lib/Analysis/ValueTracking.cpp

4874

InstOrArg -> V

5165

I think it would be more elegant to write

} else if (auto *Arg = dyn_cast<Argument>(V)) {
   // ...
} else {
   return false;
}

below, rather than having having to keep the checks here and below in sync.

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

5165

Unit tests say V can be MetadataAsValue as well (that is called from AANoUndefImpl::initialize).
But, does it imply that AANoUndefImpl::initialize should be fixed?

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

5165

Unit tests say V can be MetadataAsValue as well (that is called from AANoUndefImpl::initialize).
But, does it imply that AANoUndefImpl::initialize should be fixed?

I don't think there's anything wrong with. It shouldn't cause issues for this code as implemented, right?

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

5165

It happened because there was a function call that had a metadata as its argument, and AANoUndefImpl was somehow trying to work with the metadata argument.
Yes, I think AANoUndefImpl is okay. I feel it might be safe for isGuaranteedNotToBeUndefOrPoison to simply return false when MetadataAsValue is given, to make possible future changes assertion-safe.

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

5165

Made the change at https://github.com/llvm/llvm-project/commit/92106641ae297c24877085e0357e8095aa7b43c9 as a NFC commit, because it's just a couple of lines change

llvm/lib/Analysis/ValueTracking.cpp

 bool llvm::canCreateUndefOrPoison(const Operator *Op) {
   return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/false);
 }
 
 bool llvm::canCreatePoison(const Operator *Op) {
   return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/true);
 }
 
-static bool programUndefinedIfUndefOrPoison(const Instruction *Inst,
+static bool programUndefinedIfUndefOrPoison(const Value *V,
                                             bool PoisonOnly);
 
 static bool isGuaranteedNotToBeUndefOrPoison(const Value *V,
                                              const Instruction *CtxI,
                                              const DominatorTree *DT,
                                              unsigned Depth, bool PoisonOnly) {
   if (Depth >= MaxAnalysisRecursionDepth)
     return false;
 
   if (const auto *A = dyn_cast<Argument>(V)) {
     if (A->hasAttribute(Attribute::NoUndef))
       return true;
   }
 
   if (auto *C = dyn_cast<Constant>(V)) {
     if (isa<UndefValue>(C))
       return PoisonOnly;
 
     if (isa<ConstantInt>(C) || isa<GlobalVariable>(C) || isa<ConstantFP>(V) ||
         isa<ConstantPointerNull>(C) || isa<Function>(C))
       return true;
 
     if (C->getType()->isVectorTy() && !isa<ConstantExpr>(C))
if (const auto *CB = dyn_cast<CallBase>(V)) {
       if (CB->hasRetAttr(Attribute::NoUndef))
         return true;
     }
 
     if (!canCreateUndefOrPoison(Opr) && all_of(Opr->operands(), OpCheck))
       return true;
   }
 
-  if (auto *I = dyn_cast<Instruction>(V)) {
-    if (programUndefinedIfUndefOrPoison(I, PoisonOnly))
-      return true;
-  }
+  if (programUndefinedIfUndefOrPoison(V, PoisonOnly))
 
   // CxtI may be null or a cloned instruction.
   if (!CtxI || !CtxI->getParent() || !DT)
     return false;
 
   auto *DNode = DT->getNode(CtxI->getParent());
   if (!DNode)
     // Unreachable block
bool llvm::mustTriggerUB(const Instruction *I,
 
   for (const auto *V : NonPoisonOps)
     if (KnownPoison.count(V))
       return true;
 
   return false;
 }
 
-static bool programUndefinedIfUndefOrPoison(const Instruction *Inst,
+static bool programUndefinedIfUndefOrPoison(const Value *V,
                                             bool PoisonOnly) {
   // We currently only look for uses of values within the same basic
   // block, as that makes it easier to guarantee that the uses will be
   // executed given that Inst is executed.
   //
   // FIXME: Expand this to consider uses beyond the same basic block. To do
   // this, look out for the distinction between post-dominance and strong
   // post-dominance.
-  const BasicBlock *BB = Inst->getParent();
+  const BasicBlock *BB = nullptr;
+  BasicBlock::const_iterator Begin;
+  if (const auto *Inst = dyn_cast<Instruction>(V)) {
+    BB = Inst->getParent();
+    Begin = Inst->getIterator();
+    Begin++;
+  } else if (const auto *Arg = dyn_cast<Argument>(V)) {
+    BB = &Arg->getParent()->getEntryBlock();
+    Begin = BB->begin();
+  } else {
+    return false;
+  }
 
-  BasicBlock::const_iterator Begin = Inst->getIterator(), End = BB->end();
+  BasicBlock::const_iterator End = BB->end();
 
   if (!PoisonOnly) {
     // Be conservative & just check whether a value is passed to a noundef
     // argument.
     // Instructions that raise UB with a poison operand are well-defined
     // or have unclear semantics when the input is partially undef.
     // For example, 'udiv x, (undef | 1)' isn't UB.
 
     for (auto &I : make_range(Begin, End)) {
       if (const auto *CB = dyn_cast<CallBase>(&I)) {
         for (unsigned i = 0; i < CB->arg_size(); ++i) {
           if (CB->paramHasAttr(i, Attribute::NoUndef) &&
-              CB->getArgOperand(i) == Inst)
+              CB->getArgOperand(i) == V)
             return true;
         }
       }
       if (!isGuaranteedToTransferExecutionToSuccessor(&I))
         break;
     }
     return false;
   }
 
   // Set of instructions that we have proved will yield poison if Inst
   // does.
   SmallSet<const Value *, 16> YieldsPoison;
   SmallSet<const BasicBlock *, 4> Visited;
-  YieldsPoison.insert(Inst);
-  Visited.insert(Inst->getParent());
+
+  YieldsPoison.insert(V);
+  auto Propagate = [&](const User *User) {
+    if (propagatesPoison(cast<Operator>(User)))
+      YieldsPoison.insert(User);
+  };
+  for_each(V->users(), Propagate);
+  Visited.insert(BB);
 
   unsigned Iter = 0;
   while (Iter++ < MaxAnalysisRecursionDepth) {
     for (auto &I : make_range(Begin, End)) {
-      if (&I != Inst) {
-        if (mustTriggerUB(&I, YieldsPoison))
-          return true;
-        if (!isGuaranteedToTransferExecutionToSuccessor(&I))
-          return false;
-      }
 
       // Mark poison that propagates from I through uses of I.
-      if (YieldsPoison.count(&I)) {
-        for (const User *User : I.users()) {
-          const Instruction *UserI = cast<Instruction>(User);
-          if (propagatesPoison(cast<Operator>(UserI)))
-            YieldsPoison.insert(User);
-        }
-      }
+      if (YieldsPoison.count(&I))
+        for_each(I.users(), Propagate);
     }
 
     if (auto *NextBB = BB->getSingleSuccessor()) {
       if (Visited.insert(NextBB).second) {
         BB = NextBB;
         Begin = BB->getFirstNonPHI()->getIterator();
         End = BB->end();
         continue;

llvm/unittests/Analysis/ValueTrackingTest.cpp

for (auto &BB : *F) {
     if (&BB == &F->getEntryBlock())
       continue;
 
     EXPECT_EQ(isGuaranteedNotToBePoison(A, BB.getTerminator(), &DT), true)
         << "isGuaranteedNotToBePoison does not hold at " << *BB.getTerminator();
   }
 }
 
+TEST_F(ValueTrackingTest, isGuaranteedNotToBeUndefOrPoison) {
+  parseAssembly("declare void @f(i32 noundef)"
+                "define void @test(i32 %x) {\n"
+                "  %A = bitcast i32 %x to i32\n"
+                "  call void @f(i32 noundef %x)\n"
+                "  ret void\n"
+                "}\n");
+  EXPECT_EQ(isGuaranteedNotToBeUndefOrPoison(A), true);
+}
+
 TEST(ValueTracking, canCreatePoisonOrUndef) {
   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?, can create undef?, IR instruction)
   SmallVector<std::pair<std::pair<bool, bool>, std::string>, 32> Data = {