[ValueTracking] Use assume's noundef operand bundle

[aqjune]

Summary

This patch updates isGuaranteedNotToBeUndefOrPoison to use llvm.assume's noundef operand bundle.

[lebedev.ri]

Weak +1 to using assume operand bundles - it's more friendly for llvm-reduce.

[jdoerfert]

llvm/lib/Analysis/ValueTracking.cpp

4843

I thought we had a helper for this, I think it is hasAttributeInAssume in llvm/include/llvm/Analysis/AssumeBundleQueries.h. If that is not enough we should extend it for sure. This kind of query is common and easy to get wrong. For example, this doesn't handle multiple noundef bundles.

[Tyker]

you probably want to add noundef to isUsefullToPreserve such that it gets preserved automatically when knowledge retention is turned on.

llvm/lib/Analysis/ValueTracking.cpp

4843

there is getKnowledgeValidInContext which should be enough for what is being done here. it also has a fallback when there is no AssumptionCache.
also getOperandBundle is not valid since there maybe more than one noundef bundle in this assume.

[aqjune]

Use getKnowledgeValidInContext

[aqjune]

Add NoUndef to isUsefullToPreserve

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

4843

Thanks for the infos!

[jdoerfert]

All of this makes sense to me. Testing coverage is low though.

llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp

63

We should have a test for this I think. I imagine inlining a noundef attribute with knowledge retention enabled should actually cause an assume now.

[aqjune]

Add a test

[aqjune]

llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp

63

It worked, thanks

[jdoerfert]

LGTM

llvm/include/llvm/Analysis/ValueTracking.h

constexpr unsigned MaxAnalysisRecursionDepth = 6;
   /// 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
   /// immediately before the CtxI.
   bool isGuaranteedNotToBeUndefOrPoison(const Value *V,
+                                        AssumptionCache *AC = nullptr,
                                         const Instruction *CtxI = nullptr,
                                         const DominatorTree *DT = nullptr,
                                         unsigned Depth = 0);
-  bool isGuaranteedNotToBePoison(const Value *V,
+  bool isGuaranteedNotToBePoison(const Value *V, AssumptionCache *AC = nullptr,
                                  const Instruction *CtxI = nullptr,
                                  const DominatorTree *DT = nullptr,
                                  unsigned Depth = 0);
 
   /// Specific patterns of select instructions we can match.
   enum SelectPatternFlavor {
     SPF_UNKNOWN = 0,
     SPF_SMIN,                   /// Signed minimum

llvm/lib/Analysis/InstructionSimplify.cpp

static Value *SimplifySelectInst(Value *Cond, Value *TrueVal, Value *FalseVal,
   // select ?, X, X -> X
   if (TrueVal == FalseVal)
     return TrueVal;
 
   // If the true or false value is undef, we can fold to the other value as
   // long as the other value isn't poison.
   // select ?, undef, X -> X
   if (Q.isUndefValue(TrueVal) &&
-      isGuaranteedNotToBeUndefOrPoison(FalseVal, Q.CxtI, Q.DT))
+      isGuaranteedNotToBeUndefOrPoison(FalseVal, Q.AC, Q.CxtI, Q.DT))
     return FalseVal;
   // select ?, X, undef -> X
   if (Q.isUndefValue(FalseVal) &&
-      isGuaranteedNotToBeUndefOrPoison(TrueVal, Q.CxtI, Q.DT))
+      isGuaranteedNotToBeUndefOrPoison(TrueVal, Q.AC, Q.CxtI, Q.DT))
     return TrueVal;
 
   // Deal with partial undef vector constants: select ?, VecC, VecC' --> VecC''
   Constant *TrueC, *FalseC;
   if (TrueVal->getType()->isVectorTy() && match(TrueVal, m_Constant(TrueC)) &&
       match(FalseVal, m_Constant(FalseC))) {
     unsigned NumElts =
         cast<FixedVectorType>(TrueC->getType())->getNumElements();
if (Value *Ret = simplifyIntrinsic(Call, Q))
       return Ret;
 
   return nullptr;
 }
 
 /// Given operands for a Freeze, see if we can fold the result.
 static Value *SimplifyFreezeInst(Value *Op0, const SimplifyQuery &Q) {
   // Use a utility function defined in ValueTracking.
-  if (llvm::isGuaranteedNotToBeUndefOrPoison(Op0, Q.CxtI, Q.DT))
+  if (llvm::isGuaranteedNotToBeUndefOrPoison(Op0, Q.AC, Q.CxtI, Q.DT))
     return Op0;
   // We have room for improvement.
   return nullptr;
 }
 
 Value *llvm::SimplifyFreezeInst(Value *Op0, const SimplifyQuery &Q) {
   return ::SimplifyFreezeInst(Op0, Q);
 }

llvm/lib/Analysis/ValueTracking.cpp

if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I->getOperand(0))) {
           computeKnownBitsMul(II->getArgOperand(0), II->getArgOperand(1), false,
                               DemandedElts, Known, Known2, Depth, Q);
           break;
         }
       }
     }
     break;
   case Instruction::Freeze:
-    if (isGuaranteedNotToBePoison(I->getOperand(0), Q.CxtI, Q.DT, Depth + 1))
+    if (isGuaranteedNotToBePoison(I->getOperand(0), Q.AC, Q.CxtI, Q.DT,
+                                  Depth + 1))
       computeKnownBits(I->getOperand(0), Known, Depth + 1, Q);
     break;
   }
 }
 
 /// Determine which bits of V are known to be either zero or one and return
 /// them.
 KnownBits computeKnownBits(const Value *V, const APInt &DemandedElts,
if (auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType())) {
         DemandedVecElts = APInt::getOneBitSet(NumElts, CIdx->getZExtValue());
       return isKnownNonZero(Vec, DemandedVecElts, Depth, Q);
     }
   }
   // Freeze
   else if (const FreezeInst *FI = dyn_cast<FreezeInst>(V)) {
     auto *Op = FI->getOperand(0);
     if (isKnownNonZero(Op, Depth, Q) &&
-        isGuaranteedNotToBePoison(Op, Q.CxtI, Q.DT, Depth))
+        isGuaranteedNotToBePoison(Op, Q.AC, Q.CxtI, Q.DT, Depth))
       return true;
   }
 
   KnownBits Known(BitWidth);
   computeKnownBits(V, DemandedElts, Known, Depth, Q);
   return Known.One != 0;
 }
 

 bool llvm::canCreatePoison(const Operator *Op) {
   return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/true);
 }
 
 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) {
   if (Depth >= MaxAnalysisRecursionDepth)
     return false;
 
   if (isa<MetadataAsValue>(V))
     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))
       return (PoisonOnly || !C->containsUndefElement()) &&
              !C->containsConstantExpression();
   }
 
   // Strip cast operations from a pointer value.
   // Note that stripPointerCastsSameRepresentation can strip off getelementptr
   // inbounds with zero offset. To guarantee that the result isn't poison, the
   // stripped pointer is checked as it has to be pointing into an allocated
   // object or be null `null` to ensure `inbounds` getelement pointers with a
   // zero offset could not produce poison.
   // It can strip off addrspacecast that do not change bit representation as
   // well. We believe that such addrspacecast is equivalent to no-op.
   auto *StrippedV = V->stripPointerCastsSameRepresentation();
   if (isa<AllocaInst>(StrippedV) || isa<GlobalVariable>(StrippedV) ||
       isa<Function>(StrippedV) || isa<ConstantPointerNull>(StrippedV))
     return true;
 
   auto OpCheck = [&](const Value *V) {
-    return isGuaranteedNotToBeUndefOrPoison(V, CtxI, DT, Depth + 1, PoisonOnly);
+    return isGuaranteedNotToBeUndefOrPoison(V, AC, CtxI, DT, Depth + 1,
+                                            PoisonOnly);
   };
 
   if (auto *Opr = dyn_cast<Operator>(V)) {
     // If the value is a freeze instruction, then it can never
     // be undef or poison.
     if (isa<FreezeInst>(V))
       return true;
 
     if (const auto *CB = dyn_cast<CallBase>(V)) {
       if (CB->hasRetAttr(Attribute::NoUndef))
         return true;
     }
 
     if (const auto *PN = dyn_cast<PHINode>(V)) {
       unsigned Num = PN->getNumIncomingValues();
       bool IsWellDefined = true;
       for (unsigned i = 0; i < Num; ++i) {
         auto *TI = PN->getIncomingBlock(i)->getTerminator();
-        if (!isGuaranteedNotToBeUndefOrPoison(PN->getIncomingValue(i), TI, DT,
-                                              Depth + 1, PoisonOnly)) {
+        if (!isGuaranteedNotToBeUndefOrPoison(PN->getIncomingValue(i), AC, TI,
+                                              DT, Depth + 1, PoisonOnly)) {
           IsWellDefined = false;
           break;
         }
       }
       if (IsWellDefined)
         return true;
     } else if (!canCreateUndefOrPoison(Opr) && all_of(Opr->operands(), OpCheck))
       return true;
if (Cond) {
             any_of(Opr->operand_values(), [&](Value *Op) { return Op == V; }))
           return true;
       }
     }
 
     Dominator = Dominator->getIDom();
   }
 
+  SmallVector<Attribute::AttrKind, 2> AttrKinds{Attribute::NoUndef};
+  if (getKnowledgeValidInContext(V, AttrKinds, CtxI, DT, AC))
+    return true;
+
   return false;
 }
 
-bool llvm::isGuaranteedNotToBeUndefOrPoison(const Value *V,
+bool llvm::isGuaranteedNotToBeUndefOrPoison(const Value *V, AssumptionCache *AC,
                                             const Instruction *CtxI,
                                             const DominatorTree *DT,
                                             unsigned Depth) {
-  return ::isGuaranteedNotToBeUndefOrPoison(V, CtxI, DT, Depth, false);
+  return ::isGuaranteedNotToBeUndefOrPoison(V, AC, CtxI, DT, Depth, false);
 }
 
-bool llvm::isGuaranteedNotToBePoison(const Value *V, const Instruction *CtxI,
+bool llvm::isGuaranteedNotToBePoison(const Value *V, AssumptionCache *AC,
+                                     const Instruction *CtxI,
                                      const DominatorTree *DT, unsigned Depth) {
-  return ::isGuaranteedNotToBeUndefOrPoison(V, CtxI, DT, Depth, true);
+  return ::isGuaranteedNotToBeUndefOrPoison(V, AC, CtxI, DT, Depth, true);
 }
 
 OverflowResult llvm::computeOverflowForSignedAdd(const AddOperator *Add,
                                                  const DataLayout &DL,
                                                  AssumptionCache *AC,
                                                  const Instruction *CxtI,
                                                  const DominatorTree *DT) {
   return ::computeOverflowForSignedAdd(Add->getOperand(0), Add->getOperand(1),

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

else
       AANoUndef::initialize(A);
   }
 
   /// See followUsesInMBEC
   bool followUseInMBEC(Attributor &A, const Use *U, const Instruction *I,
                        AANoUndef::StateType &State) {
     const Value *UseV = U->get();
     const DominatorTree *DT = nullptr;
-    if (Function *F = getAnchorScope())
-      DT = A.getInfoCache().getAnalysisResultForFunction<DominatorTreeAnalysis>(
-          *F);
-    State.setKnown(isGuaranteedNotToBeUndefOrPoison(UseV, I, DT));
+    AssumptionCache *AC = nullptr;
+    InformationCache &InfoCache = A.getInfoCache();
+    if (Function *F = getAnchorScope()) {
+      DT = InfoCache.getAnalysisResultForFunction<DominatorTreeAnalysis>(*F);
+      AC = InfoCache.getAnalysisResultForFunction<AssumptionAnalysis>(*F);
+    }
+    State.setKnown(isGuaranteedNotToBeUndefOrPoison(UseV, AC, I, DT));
     bool TrackUse = false;
     // Track use for instructions which must produce undef or poison bits when
     // at least one operand contains such bits.
     if (isa<CastInst>(*I) || isa<GetElementPtrInst>(*I))
       TrackUse = true;
     return TrackUse;
   }
 

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

Instruction *InstCombinerImpl::foldSelectValueEquivalence(SelectInst &Sel,
   }
 
   // In X == Y ? f(X) : Z, try to evaluate f(Y) and replace the operand.
   // Make sure Y cannot be undef though, as we might pick different values for
   // undef in the icmp and in f(Y). Additionally, take care to avoid replacing
   // X == Y ? X : Z with X == Y ? Y : Z, as that would lead to an infinite
   // replacement cycle.
   Value *CmpLHS = Cmp.getOperand(0), *CmpRHS = Cmp.getOperand(1);
-  if (TrueVal != CmpLHS && isGuaranteedNotToBeUndefOrPoison(CmpRHS, &Sel, &DT))
+  if (TrueVal != CmpLHS &&
+      isGuaranteedNotToBeUndefOrPoison(CmpRHS, SQ.AC, &Sel, &DT))
     if (Value *V = SimplifyWithOpReplaced(TrueVal, CmpLHS, CmpRHS, SQ,
                                           /* AllowRefinement */ true))
       return replaceOperand(Sel, Swapped ? 2 : 1, V);
-  if (TrueVal != CmpRHS && isGuaranteedNotToBeUndefOrPoison(CmpLHS, &Sel, &DT))
+  if (TrueVal != CmpRHS &&
+      isGuaranteedNotToBeUndefOrPoison(CmpLHS, SQ.AC, &Sel, &DT))
     if (Value *V = SimplifyWithOpReplaced(TrueVal, CmpRHS, CmpLHS, SQ,
                                           /* AllowRefinement */ true))
       return replaceOperand(Sel, Swapped ? 2 : 1, V);
 
   auto *FalseInst = dyn_cast<Instruction>(FalseVal);
   if (!FalseInst)
     return nullptr;
 

llvm/lib/Transforms/Scalar/DivRemPairs.cpp

if (HasDivRemOp) {
       //   %div = sdiv undef, 1 // %div = undef
       //   %rem = srem undef, 1 // %rem = 0
       // =>
       //   %div = sdiv undef, 1 // %div = undef
       //   %mul = mul %div, 1   // %mul = undef
       //   %rem = sub %x, %mul  // %rem = undef - undef = undef
       // If X is not frozen, %rem becomes undef after transformation.
       // TODO: We need a undef-specific checking function in ValueTracking
-      if (!isGuaranteedNotToBeUndefOrPoison(X, DivInst, &DT)) {
+      if (!isGuaranteedNotToBeUndefOrPoison(X, nullptr, DivInst, &DT)) {
         auto *FrX = new FreezeInst(X, X->getName() + ".frozen", DivInst);
         DivInst->setOperand(0, FrX);
         Sub->setOperand(0, FrX);
       }
       // Same for Y. If X = 1 and Y = (undef | 1), %rem in src is either 1 or 0,
       // but %rem in tgt can be one of many integer values.
-      if (!isGuaranteedNotToBeUndefOrPoison(Y, DivInst, &DT)) {
+      if (!isGuaranteedNotToBeUndefOrPoison(Y, nullptr, DivInst, &DT)) {
         auto *FrY = new FreezeInst(Y, Y->getName() + ".frozen", DivInst);
         DivInst->setOperand(1, FrY);
         Mul->setOperand(1, FrY);
       }
 
       // Now kill the explicit remainder. We have replaced it with:
       // (sub X, (mul (div X, Y), Y)
       Sub->setName(RemInst->getName() + ".decomposed");

llvm/lib/Transforms/Scalar/JumpThreading.cpp

for (Use &U : PN->uses()) {
       }
     }
 
     if (!SI)
       continue;
     // Expand the select.
     Value *Cond = SI->getCondition();
     if (InsertFreezeWhenUnfoldingSelect &&
-        !isGuaranteedNotToBeUndefOrPoison(Cond, SI, &DTU->getDomTree()))
+        !isGuaranteedNotToBeUndefOrPoison(Cond, nullptr, SI,
+                                          &DTU->getDomTree()))
       Cond = new FreezeInst(Cond, "cond.fr", SI);
     Instruction *Term = SplitBlockAndInsertIfThen(Cond, SI, false);
     BasicBlock *SplitBB = SI->getParent();
     BasicBlock *NewBB = Term->getParent();
     PHINode *NewPN = PHINode::Create(SI->getType(), 2, "", SI);
     NewPN->addIncoming(SI->getTrueValue(), Term->getParent());
     NewPN->addIncoming(SI->getFalseValue(), BB);
     SI->replaceAllUsesWith(NewPN);

llvm/lib/Transforms/Scalar/SCCP.cpp

if (ReturnValue.isConstantRange() &&
         auto *CB = dyn_cast<CallBase>(User);
         if (!CB || CB->getCalledFunction() != F)
           continue;
 
         // Limit to cases where the return value is guaranteed to be neither
         // poison nor undef. Poison will be outside any range and currently
         // values outside of the specified range cause immediate undefined
         // behavior.
-        if (!isGuaranteedNotToBeUndefOrPoison(CB, CB))
+        if (!isGuaranteedNotToBeUndefOrPoison(CB, nullptr, CB))
           continue;
 
         // Do not touch existing metadata for now.
         // TODO: We should be able to take the intersection of the existing
         // metadata and the inferred range.
         if (CB->getMetadata(LLVMContext::MD_range))
           continue;
 

llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp

 DEBUG_COUNTER(BuildAssumeCounter, "assume-builder-counter",
               "Controls which assumes gets created");
 
 namespace {
 
 bool isUsefullToPreserve(Attribute::AttrKind Kind) {
   switch (Kind) {
     case Attribute::NonNull:
+    case Attribute::NoUndef:
     case Attribute::Alignment:
     case Attribute::Dereferenceable:
     case Attribute::DereferenceableOrNull:
     case Attribute::Cold:
       return true;
     default:
       return false;
   }
 }
 
 /// This function will try to transform the given knowledge into a more
 /// canonical one. the canonical knowledge maybe the given one.
 RetainedKnowledge canonicalizedKnowledge(RetainedKnowledge RK, Module *M) {
   switch (RK.AttrKind) {
   default:
     return RK;

llvm/lib/Transforms/Utils/CanonicalizeFreezeInLoops.cpp

 // SCEV forget user. The inserted freeze is placed in the preheader.
 void CanonicalizeFreezeInLoopsImpl::InsertFreezeAndForgetFromSCEV(Use &U) {
   auto *PH = L->getLoopPreheader();
 
   auto *UserI = cast<Instruction>(U.getUser());
   auto *ValueToFr = U.get();
   assert(L->contains(UserI->getParent()) &&
          "Should not process an instruction that isn't inside the loop");
-  if (isGuaranteedNotToBeUndefOrPoison(ValueToFr, UserI, &DT))
+  if (isGuaranteedNotToBeUndefOrPoison(ValueToFr, nullptr, UserI, &DT))
     return;
 
   LLVM_DEBUG(dbgs() << "canonfr: inserting freeze:\n");
   LLVM_DEBUG(dbgs() << "\tUser: " << *U.getUser() << "\n");
   LLVM_DEBUG(dbgs() << "\tOperand: " << *U.get() << "\n");
 
   U.set(new FreezeInst(ValueToFr, ValueToFr->getName() + ".frozen",
                        PH->getTerminator()));
for (const auto &Info : Candidates) {
     PHINode *PHI = Info.PHI;
     if (!ProcessedPHIs.insert(Info.PHI).second)
       continue;
 
     BinaryOperator *StepI = Info.StepInst;
     assert(StepI && "Step instruction should have been found");
 
     // Drop flags from the step instruction.
-    if (!isGuaranteedNotToBeUndefOrPoison(StepI, StepI, &DT)) {
+    if (!isGuaranteedNotToBeUndefOrPoison(StepI, nullptr, StepI, &DT)) {
       LLVM_DEBUG(dbgs() << "canonfr: drop flags: " << *StepI << "\n");
       StepI->dropPoisonGeneratingFlags();
       SE.forgetValue(StepI);
     }
 
     InsertFreezeAndForgetFromSCEV(StepI->getOperandUse(Info.StepValIdx));
 
     unsigned OperandIdx =

llvm/test/Transforms/Inline/arg-attr-propagation.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
 ; RUN: opt -inline -S < %s | FileCheck --check-prefixes=CHECK,NO_ASSUME %s
 ; RUN: opt -inline -S --enable-knowledge-retention < %s | FileCheck %s --check-prefixes=CHECK,USE_ASSUME
 
 ; The callee guarantees that the pointer argument is nonnull and dereferenceable.
 ; That information should transfer to the caller.
 
 define i32 @callee(i32* dereferenceable(32) %t1) {
 ; CHECK-LABEL: define {{[^@]+}}@callee
-; CHECK-SAME: (i32* dereferenceable(32) [[T1:%.*]])
-; CHECK-NEXT:    [[T2:%.*]] = load i32, i32* [[T1]]
+; CHECK-SAME: (i32* dereferenceable(32) [[T1:%.*]]) {
+; CHECK-NEXT:    [[T2:%.*]] = load i32, i32* [[T1]], align 4
 ; CHECK-NEXT:    ret i32 [[T2]]
 ;
   %t2 = load i32, i32* %t1
   ret i32 %t2
 }
 
+define i32 @callee2(i32* dereferenceable(32) %t1, i32 noundef %t2) {
+; CHECK-LABEL: define {{[^@]+}}@callee2
+; CHECK-SAME: (i32* dereferenceable(32) [[T1:%.*]], i32 noundef [[T2:%.*]]) {
+; CHECK-NEXT:    [[V:%.*]] = load i32, i32* [[T1]], align 4
+; CHECK-NEXT:    ret i32 [[V]]
+;
+  %v = load i32, i32* %t1
+  ret i32 %v
+}
+
 ; FIXME: All dereferenceability information is lost.
 ; The caller argument could be known nonnull and dereferenceable(32).
 
-define i32 @caller1(i32* %t1) {
+define i32 @caller1(i32* %t1, i32 %t2) {
 ; NO_ASSUME-LABEL: define {{[^@]+}}@caller1
-; NO_ASSUME-SAME: (i32* [[T1:%.*]])
-; NO_ASSUME-NEXT:    [[T2_I:%.*]] = load i32, i32* [[T1]]
-; NO_ASSUME-NEXT:    ret i32 [[T2_I]]
+; NO_ASSUME-SAME: (i32* [[T1:%.*]], i32 [[T2:%.*]]) {
+; NO_ASSUME-NEXT:    [[V_I:%.*]] = load i32, i32* [[T1]], align 4
+; NO_ASSUME-NEXT:    ret i32 [[V_I]]
 ;
 ; USE_ASSUME-LABEL: define {{[^@]+}}@caller1
-; USE_ASSUME-SAME: (i32* [[T1:%.*]])
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[T1]], i64 32) ]
-; USE_ASSUME-NEXT:    [[T2_I:%.*]] = load i32, i32* [[T1]]
-; USE_ASSUME-NEXT:    ret i32 [[T2_I]]
+; USE_ASSUME-SAME: (i32* [[T1:%.*]], i32 [[T2:%.*]]) {
+; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[T1]], i64 32), "noundef"(i32 [[T2]]) ]
+; USE_ASSUME-NEXT:    [[V_I:%.*]] = load i32, i32* [[T1]], align 4
+; USE_ASSUME-NEXT:    ret i32 [[V_I]]
 ;
-  %t2 = tail call i32 @callee(i32* dereferenceable(32) %t1)
-  ret i32 %t2
+  %v = tail call i32 @callee2(i32* dereferenceable(32) %t1, i32 noundef %t2)
+  ret i32 %v
 }
 
 ; The caller argument is nonnull, but that can be explicit.
 ; The dereferenceable amount could be increased.
 
 define i32 @caller2(i32* dereferenceable(31) %t1) {
 ; NO_ASSUME-LABEL: define {{[^@]+}}@caller2
 ; NO_ASSUME-SAME: (i32* dereferenceable(31) [[T1:%.*]])

llvm/unittests/Analysis/ValueTrackingTest.cpp

 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 namespace {
 
-static Instruction &findInstructionByName(Function *F, StringRef Name) {
+static Instruction *findInstructionByNameOrNull(Function *F, StringRef Name) {
   for (Instruction &I : instructions(F))
     if (I.getName() == Name)
-      return I;
+      return &I;
+
+  return nullptr;
+}
+
+static Instruction &findInstructionByName(Function *F, StringRef Name) {
+  auto *I = findInstructionByNameOrNull(F, Name);
+  if (I)
+    return *I;
 
   llvm_unreachable("Expected value not found");
 }
 
 class ValueTrackingTest : public testing::Test {
 protected:
   std::unique_ptr<Module> parseModule(StringRef Assembly) {
     SMDiagnostic Error;
void parseAssembly(StringRef Assembly) {
     M = parseModule(Assembly);
     ASSERT_TRUE(M);
 
     F = M->getFunction("test");
     ASSERT_TRUE(F) << "Test must have a function @test";
     if (!F)
       return;
 
-    A = &findInstructionByName(F, "A");
+    A = findInstructionByNameOrNull(F, "A");
     ASSERT_TRUE(A) << "@test must have an instruction %A";
+
+    CxtI = findInstructionByNameOrNull(F, "CxtI");
+    CxtI2 = findInstructionByNameOrNull(F, "CxtI2");
+    CxtI3 = findInstructionByNameOrNull(F, "CxtI3");
   }
 
   LLVMContext Context;
   std::unique_ptr<Module> M;
   Function *F = nullptr;
   Instruction *A = nullptr;
+
+  // Context instructions (optional)
+  Instruction *CxtI = nullptr, *CxtI2 = nullptr, *CxtI3 = 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("declare i1 @any_bool()"
                 "BB2:\n"
                 "  ret void\n"
                 "}\n");
   DominatorTree DT(*F);
   for (auto &BB : *F) {
     if (&BB == &F->getEntryBlock())
       continue;
 
-    EXPECT_EQ(isGuaranteedNotToBePoison(A, BB.getTerminator(), &DT), true)
+    EXPECT_EQ(isGuaranteedNotToBePoison(A, nullptr, BB.getTerminator(), &DT),
+              true)
         << "isGuaranteedNotToBePoison does not hold at " << *BB.getTerminator();
   }
 }
 
 TEST_F(ValueTrackingTest, isGuaranteedNotToBePoison_phi) {
   parseAssembly("declare i32 @any_i32(i32)"
                 "define void @test() {\n"
                 "ENTRY:\n"
                 "  br label %LOOP\n"
                 "LOOP:\n"
                 "  %A = phi i32 [0, %ENTRY], [%A.next, %NEXT]\n"
                 "  %A.next = call i32 @any_i32(i32 %A)\n"
                 "  %cond = icmp eq i32 %A.next, 0\n"
                 "  br i1 %cond, label %NEXT, label %EXIT\n"
                 "NEXT:\n"
                 "  br label %LOOP\n"
                 "EXIT:\n"
                 "  ret void\n"
                 "}\n");
   DominatorTree DT(*F);
   for (auto &BB : *F) {
     if (BB.getName() == "LOOP") {
-      EXPECT_EQ(isGuaranteedNotToBePoison(A, A, &DT), true)
+      EXPECT_EQ(isGuaranteedNotToBePoison(A, nullptr, A, &DT), true)
           << "isGuaranteedNotToBePoison does not hold";
     }
   }
 }
 
 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_F(ValueTrackingTest, isGuaranteedNotToBeUndefOrPoison_assume) {
+  parseAssembly("declare i1 @f_i1()\n"
+                "declare i32 @f_i32()\n"
+                "declare void @llvm.assume(i1)\n"
+                "define void @test() {\n"
+                "  %A = call i32 @f_i32()\n"
+                "  %cond = call i1 @f_i1()\n"
+                "  %CxtI = add i32 0, 0\n"
+                "  br i1 %cond, label %BB1, label %EXIT\n"
+                "BB1:\n"
+                "  %CxtI2 = add i32 0, 0\n"
+                "  %cond2 = call i1 @f_i1()\n"
+                "  call void @llvm.assume(i1 true) [ \"noundef\"(i32 %A) ]\n"
+                "  br i1 %cond2, label %BB2, label %EXIT\n"
+                "BB2:\n"
+                "  %CxtI3 = add i32 0, 0\n"
+                "  ret void\n"
+                "EXIT:\n"
+                "  ret void\n"
+                "}");
+  AssumptionCache AC(*F);
+  DominatorTree DT(*F);
+  EXPECT_FALSE(isGuaranteedNotToBeUndefOrPoison(A, &AC, CxtI, &DT));
+  EXPECT_FALSE(isGuaranteedNotToBeUndefOrPoison(A, &AC, CxtI2, &DT));
+  EXPECT_TRUE(isGuaranteedNotToBeUndefOrPoison(A, &AC, CxtI3, &DT));
+}
+
 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 = {