[AssumeBundles] Use assume bundles in isKnownNonZero

[Tyker]

Summary

Use nonnull and dereferenceable from an assume bundle in isKnownNonZero

[jdoerfert]

llvm/lib/Analysis/ValueTracking.cpp

708

I feel we need to move some/all of this complexity into KnowledgeRetention.
Ideally, one (or at least I) would like something along the lines of:

if (Q.DT && Q.CtxI) {
SmallVector<AttributeKind, 2> AttrKinds;
AttrKinds.append(Attribute::NonNull);
if (llvm::NullPointerIsDefined(Q.CtxI->getFunction(), V->getType()->getPointerAddressSpace()))
   AttrKinds.append(Attribute::Dereferenceable);
if (hasKnownAttributeDominatingInst(*V, AttrKinds, *Q.DT, *Q.CtxI))
  return true;
}

The impl of hasKnownAttributeDominatingInst does the required steps, thus looking at the uses. Actually all the steps should be in getKnownNonNullAndDerefBytesForUse which lives in Attributor.cpp. It does actually more, e.g., use in a call site or and recursive queries via the Attributor, but we should somehow share this logic.

My suggestion:

1. Something similar to the above interface in KnowledgeRetention
2. Make Attributor etc. in getKnownNonNullAndDerefBytesForUse optional.
3. Teach getKnownNonNullAndDerefBytesForUse about assume uses.
4. Make the getKnownNonNullAndDerefBytesForUse logic public and invoke it if we have an attribute kind NonNull or Dereferenceable[_or_null].

WDYT?

[jdoerfert]

llvm/lib/Analysis/ValueTracking.cpp

708

Maybe hasKnownAttributeDominatingInst should not need to look for deref at all as we derive it from deref in the Attributor. Otherwise we could make a helper to encapsulate the NullPointerIsDefined stuff.

[Tyker]

Updated based on review comments.

[fhahn]

llvm/lib/Analysis/ValueTracking.cpp

697

nit: is llvm:: needed?

llvm/lib/IR/KnowledgeRetention.cpp

314 ↗

(On Diff #251015)

There's AssumptionCache, which is used in various places that inspect assumes, including ValueTracking. Is there a reason to not use it here?

If there is a layering problem, I think that would need a bit more thought, but ideally we would re-use the existing infrastructure to make things available as widely as possible.

318 ↗

(On Diff #251015)

There's isValidAssumeForContext (in ValueTracking.h, https://llvm.org/doxygen/namespacellvm.html#af1d57d70f59588a35ebfe3c4f50fecd2), which already has a similar check (but more general and not necessarily requiring DT. It might be good to use that here instead.

[Tyker]

updated the API so that getKnownNonNullAndDerefBytesForUse can use it as well.

[Tyker]

llvm/lib/Analysis/ValueTracking.cpp

708

i adapted the getKnownNonNullAndDerefBytesForUse to use the same code. but i am not sure this matches what you had in mind.

Maybe hasKnownAttributeDominatingInst should not need to look for deref at all as we derive it from deref in the Attributor. Otherwise we
could make a helper to encapsulate the NullPointerIsDefined stuff.

in ValueTracking this is not needed if the attributor has run before. but i am not sure all places that can add a deref attribute will add a nonnull as well when they can. if you are confident i am fine with not looking for deref.
in the attributor we should definitely look at both nonnull and deref.

llvm/lib/IR/KnowledgeRetention.cpp

314 ↗

(On Diff #251015)

isValidAssumeForContext is now used and i saw the improvement in a test case.

for AssumptionCache. it could be used here but i think it would be costlier to use it than what we currently have.
we are working with knowledge in operand bundles, the assume is always the direct user of the value it knows something about (which is not true for boolean assumes). so we can iterate over all uses (non-recursively) to find all candidate llvm.assume. looking up the knowledge in the assume is cheaper when we have the the operand number which is only available when we lookup from the uses.

[jdoerfert]

I have minor comments but I think this is fine. Please wait on the approval of @fhahn .

llvm/include/llvm/IR/KnowledgeRetention.h

136 ↗

(On Diff #251055)

Return true iff U is a use ... with an attribute in AttrKinds. There is no predicate.

144 ↗

(On Diff #251055)

- and this assume dominate the instruction Ctx.

The above description matches a convenient wrapper we should add (I also mention that below).

llvm/lib/Analysis/ValueTracking.cpp

704

I don't think DT is mandatory here (line 664).

708

This is fine with me. We might want to add a helper that adds this "common" filter (return isValidAssumeForContext(Assume, Q.CxtI, Q.DT);) given an instruction (CtxI) and an optional dom tree.

llvm/test/Analysis/ValueTracking/assume.ll

92

Add test4b, as test4 but with:
define i32 @test4b(i32* readonly %0, i1 %cond) "null-pointer-is-valid"="true" {

[fhahn]

llvm/lib/IR/KnowledgeRetention.cpp

302 ↗

(On Diff #251055)

This could be an early exit (like auto *II = dyn_cast....; if (!II) return none().

314 ↗

(On Diff #251015)

for AssumptionCache. it could be used here but i think it would be costlier to use it than what we currently have.

I am not sure how using assumptionCache would be more costly, assuming it is already populated?

we are working with knowledge in operand bundles, the assume is always the direct user of the value it knows something about (which is not true for boolean assumes).

But you still have to check all uses of each value (of which many may not be assumptions), to find the available assumptions IIUC. By using AssumptionCache, you can skip the traversal of the use list, which should be quite beneficial, assuming most values do not have associated assumptions.

[Tyker]

Fixed review comments.

[Tyker]

llvm/lib/Analysis/ValueTracking.cpp

708

isValidAssumeForContext is in Analysis so we can't access it from KnowledgeRetention.

llvm/lib/IR/KnowledgeRetention.cpp

314 ↗

(On Diff #251015)

when we lookup from use, we can use the operand number so we can use getKnowledgeFromOperandInAssume which is expected to be cheaper than log N for N being the number of element in the assume. if we don't have the operand number, we have to fallback to linear search.

that said i think we can be get the best of both world by improving AssumptionCache to store the operand bundle index with the assume. but this should probably go into an other patch.

[fhahn]

llvm/lib/Analysis/ValueTracking.cpp

708

That's unfortunate in terms of layering :(

I assume the reason this is in IR is because of the AssumeBuilder part, right?

Did you consider splitting the code into an analysis part (that lives in lib/Analysis and can share utilities with ValueTracking & co) & a transform part (that lives in either IR/Transforms)?

I think it would be great to share utilities with the existing code that deals with assumes, as that might lead to general improvements that also benefit the existing code, increases test coverage and reduces duplication.

llvm/lib/IR/KnowledgeRetention.cpp

314 ↗

(On Diff #251015)

Right, I see what you mean now. I guess it all depends on what we expect to be larger, the number of (non-assume) uses of the value or the number of operands in the assume. I would expect the former to be larger in practice, given that most values do not have any assume users.

[Tyker]

rebased

[Tyker]

llvm/lib/Analysis/ValueTracking.cpp

708

That's unfortunate in terms of layering :(

i agree

I assume the reason this is in IR is because of the AssumeBuilder part, right?

Yes this is why.

Did you consider splitting the code into an analysis part (that lives in lib/Analysis and can share utilities with ValueTracking & co) & a transform > part (that lives in either IR/Transforms)?

yeah i think this is going to happen. but there are too many constraints to satisfy them all, although we can satisfy more constraint than we currently do.
a list of constraint is:

• Test for the Query part depend on the AssumeBuilder part.
• We would like the Query part to have access to Analysis for the situation we have here and many more i think.
• In rG57c964aaa76bfaa908398fbd9d8c9d6d19856859#909236, @nikic wrote: "Are you sure IR is the right place for this? I don't think that we have any passes in there, apart from the IR verifier. Possibly this needs splitting up into parts that are needed by analysis and by transforms?"
• When working on improvement on the assume builder i saw some situation where having access to Analysis is beneficial.
• The AssumeBuilder part doesn't logically fit in Analysis

I think that moving the Query part to Analysis is overall a good idea.
however placing the AssumeBuilder part always break some constraint either it is in IR, Analysis or Transform.
if placed in IR we have a pass in IR (for testing) which is subotimal and AssumeBuilder part can't use Analysis.
if placed in Analysis it intuitively doesn't fit there.
if placed in Transform we either need to change the testing method for the Query part or test it from Transform.

[Tyker]

rebased on D77171

[fhahn]

llvm/lib/Analysis/ValueTracking.cpp

708

Test for the Query part depend on the AssumeBuilder part.

Is that the main reason for being in IR/ currently?

I think a better compromise would be to have the AsumeBuilder in Transforms/Utils or something and then test the query part in Transform/ there as well. Also, shouldn't it be also possible to test the query part without the AssumeBuilder, by specifying the expected assume bundles in the IR directly?

In rG57c964aaa76bfaa908398fbd9d8c9d6d19856859#909236, @nikic wrote: "Are you sure IR is the right place for this? I don't think that we have any passes in there, apart from the IR verifier. Possibly this needs splitting up into parts that are needed by analysis and by transforms?"

+1 to moving the pass out of IR/.

I think that moving the Query part to Analysis is overall a good idea.

Sounds good. That should be a relatively straight forward change I think.

however placing the AssumeBuilder part always break some constraint either it is in IR, Analysis or Transform.

As mentioned above, if the main reason for it being in IR/ is the unit tests, I would recommend directly construction instruction with assume bundles for the tests in Analysis/ (if possible. I might miss some details here). Conceptually the pass should be somewhere in Transforms/. I think a comparable case is PredicateInfo, which lives in llvm/include/llvm/Transforms/Utils/PredicateInfo.h.

[Tyker]

llvm/lib/Analysis/ValueTracking.cpp

708

this is going to be fixed by D77171.
and this patch have been rebased on it.
so i added the getKnowledgeValidInContext as requested.

[Tyker]

rebase on D77402 and add support for assumption cache.

[fhahn]

llvm/lib/Analysis/ValueTracking.cpp

708

SGTM, thanks!

[jdoerfert]

I'm fine with this, one test needs an update though.

llvm/test/Transforms/Attributor/nonnull.ll

36

I don't think these test what they should (the change in the Attributor). Since the Attributor knows about assumes (D74888) it can do this deduction already: https://godbolt.org/z/S-PNc-

I think this should work as test though:
https://godbolt.org/z/_6mVi3

[fhahn]

llvm/lib/Analysis/AssumeBundleQueries.cpp

133

This probably could be simplified by using the stuff from PatternMatch.h, e.g. something like match(U->getUser(), m_Intrinsic<Intrinsic::assume>(m_Specific(U))

147

nit: llvm:: not needed there I think

151

Can the cache contain elements that are not assumes? I think ideally it would only contain entries with valid assume calls.

155

llvm:: not needed here I think as there is using namespace llvm.

llvm/lib/Analysis/ValueTracking.cpp

695

nit: llvm:: not needed I think

[Tyker]

addressed comments.
add asserts to make sure that both path of getKnowledgeForValue lead to the same result.

llvm/lib/Analysis/AssumeBundleQueries.cpp

151

the assumption cache cannot hold valid pointers to non-assume instruction.
but it can hold nullptr if an assume was in cache and gets deleted.

[Tyker]

rebased

[fhahn]

llvm/include/llvm/Analysis/AssumeBundleQueries.h

18

Nothing from the header seems to be actually used. Forward declare Instruction?

19

I could not find anything that uses declarations from PassManager.h. Did I miss something?

20

Nothing from the header seems to be actually used. Forward declare DominatorTree?

103

unrelated nit: property that holds?

105–109

unrelated, but it is not clear from the comment what ArgValue is referring to. Might be good to clarify

145

nit: Return

146

nit: period at end of sentence.

150

nit: Return

155

llvm:: not needed?

158

nit: Return

llvm/lib/Analysis/AssumeBundleQueries.cpp

133

It would be slightly better to just use cast<> after the check I think.

llvm/lib/Transforms/IPO/Attributor.cpp

19 ↗

(On Diff #257252)

all new includes unused?

[Tyker]

addressed comments.

[fhahn]

Thanks for all the changes! The patch looks good to me with respect to my comments, but I think it would be good for someone else to have another look as well.

llvm/include/llvm/Analysis/AssumeBundleQueries.h

106

nit: For example ..... of at least for:

• AttrKind will...
• WasOn will be
• argValue will be...

llvm/lib/Analysis/ValueTracking.cpp

703

Just a side note: it would be good to unify the assumption handling for both variants :)

[Tyker]

addressed comments.

[jdoerfert]

LGTM. One nit below and the Attributor test needs to be changed (see my 2 comments) and rebased + updated wrt check lines.

llvm/lib/Analysis/AssumeBundleQueries.cpp

139

Nit: const auto &AttrKind : AttrKinds) or no auto but the type. Especially the name Attr is (by me) associated with llvm::Attribute not the kind enum.

llvm/test/Transforms/Attributor/nonnull.ll

144

We don't check these prefixes anymore. We need to create a call of test10 to see the return attirbute (I think).

llvm/include/llvm/Analysis/AssumeBundleQueries.h

 // built using utilities from Transform/Utils/AssumeBundleBuilder.h
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_ASSUMEBUNDLEQUERIES_H
 #define LLVM_TRANSFORMS_UTILS_ASSUMEBUNDLEQUERIES_H
 
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace llvm {
 class IntrinsicInst;
+class AssumptionCache;
+class DominatorTree;
 
 /// Index of elements in the operand bundle.
 /// If the element exist it is guaranteed to be what is specified in this enum
 /// but it may not exist.
 enum AssumeBundleArg {
   ABA_WasOn = 0,
   ABA_Argument = 1,
 };

 /// Insert into the map all the informations contained in the operand bundles of
 /// the llvm.assume. This should be used instead of hasAttributeInAssume when
 /// many queries are going to be made on the same llvm.assume.
 /// String attributes are not inserted in the map.
 /// If the IR changes the map will be outdated.
 void fillMapFromAssume(CallInst &AssumeCI, RetainedKnowledgeMap &Result);
 
 /// Represent one information held inside an operand bundle of an llvm.assume.
-/// AttrKind is the property that hold.
+/// AttrKind is the property that holds.
 /// WasOn if not null is that Value for which AttrKind holds.
-/// ArgValue is optionally an argument.
+/// ArgValue is optionally an argument of the attribute.
+/// for example if we know that %P has an alignment of at least 4.
+/// AttrKind will be Attribute::Alignment
+/// WasOn will be %P
+/// ArgValue will be 4.
 struct RetainedKnowledge {
   Attribute::AttrKind AttrKind = Attribute::None;
-  Value *WasOn = nullptr;
   unsigned ArgValue = 0;
+  Value *WasOn = nullptr;
+  bool operator==(RetainedKnowledge Other) const {
+    return AttrKind == Other.AttrKind && WasOn == Other.WasOn &&
+           ArgValue == Other.ArgValue;
+  }
+  bool operator!=(RetainedKnowledge Other) const { return !(*this == Other); }
+  operator bool() const { return AttrKind != Attribute::None; }
+  static RetainedKnowledge none() { return RetainedKnowledge{}; }
 };
 
 /// Retreive the information help by Assume on the operand at index Idx.
 /// Assume should be an llvm.assume and Idx should be in the operand bundle.
 RetainedKnowledge getKnowledgeFromOperandInAssume(CallInst &Assume,
                                                   unsigned Idx);
 
 /// Retreive the information help by the Use U of an llvm.assume. the use should
 /// be in the operand bundle.
 inline RetainedKnowledge getKnowledgeFromUseInAssume(const Use *U) {
   return getKnowledgeFromOperandInAssume(*cast<CallInst>(U->getUser()),
                                          U->getOperandNo());
 }
 
 /// Return true iff the operand bundles of the provided llvm.assume doesn't
 /// contain any valuable information. This is true when:
 ///  - The operand bundle is empty
 ///  - The operand bundle only contains information about dropped values or
 ///    constant folded values.
 ///
 /// the argument to the call of llvm.assume may still be useful even if the
 /// function returned true.
 bool isAssumeWithEmptyBundle(CallInst &Assume);
 
+/// Return a valid Knowledge associated to the Use U if its Attribute kind is
+/// in AttrKinds.
+RetainedKnowledge getKnowledgeFromUse(const Use *U,
+                                      ArrayRef<Attribute::AttrKind> AttrKinds);
+
+/// Return a valid Knowledge associated to the Value V if its Attribute kind is
+/// in AttrKinds and it matches the Filter.
+RetainedKnowledge getKnowledgeForValue(
+    const Value *V, ArrayRef<Attribute::AttrKind> AttrKinds,
+    AssumptionCache *AC = nullptr,
+    function_ref<bool(RetainedKnowledge, Instruction *)> Filter =
+        [](RetainedKnowledge, Instruction *) { return true; });
+
+/// Return a valid Knowledge associated to the Value V if its Attribute kind is
+/// in AttrKinds and the knowledge is suitable to be used in the context of
+/// CtxI.
+RetainedKnowledge getKnowledgeValidInContext(
+    const Value *V, ArrayRef<Attribute::AttrKind> AttrKinds,
+    const Instruction *CtxI, const DominatorTree *DT = nullptr,
+    AssumptionCache *AC = nullptr);
+
 } // namespace llvm
 
 #endif

llvm/lib/Analysis/AssumeBundleQueries.cpp

 //===- AssumeBundleQueries.cpp - tool to query assume bundles ---*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/AssumeBundleQueries.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/PatternMatch.h"
 
 using namespace llvm;
+using namespace llvm::PatternMatch;
 
 static bool bundleHasArgument(const CallBase::BundleOpInfo &BOI, unsigned Idx) {
   return BOI.End - BOI.Begin > Idx;
 }
 
-static Value *getValueFromBundleOpInfo(IntrinsicInst &Assume,
+static Value *getValueFromBundleOpInfo(CallInst &Assume,
                                        const CallBase::BundleOpInfo &BOI,
                                        unsigned Idx) {
   assert(bundleHasArgument(BOI, Idx) && "index out of range");
   return (Assume.op_begin() + BOI.Begin + Idx)->get();
 }
 
 bool llvm::hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
                                 StringRef AttrName, uint64_t *ArgVal,
if (Lookup == Result.end() || !Lookup->second.count(&Assume)) {
       Result[Key][&Assume] = {Val, Val};
       continue;
     }
     Lookup->second[&Assume].Min = std::min(Val, Lookup->second[&Assume].Min);
     Lookup->second[&Assume].Max = std::max(Val, Lookup->second[&Assume].Max);
   }
 }
 
-RetainedKnowledge llvm::getKnowledgeFromOperandInAssume(CallInst &AssumeCI,
-                                                        unsigned Idx) {
-  IntrinsicInst &Assume = cast<IntrinsicInst>(AssumeCI);
-  assert(Assume.getIntrinsicID() == Intrinsic::assume &&
-         "this function is intended to be used on llvm.assume");
-  CallBase::BundleOpInfo BOI = Assume.getBundleOpInfoForOperand(Idx);
+static RetainedKnowledge
+getKnowledgeFromBundle(CallInst &Assume, const CallBase::BundleOpInfo &BOI) {
   RetainedKnowledge Result;
   Result.AttrKind = Attribute::getAttrKindFromName(BOI.Tag->getKey());
   Result.WasOn = getValueFromBundleOpInfo(Assume, BOI, ABA_WasOn);
   if (BOI.End - BOI.Begin > ABA_Argument)
     Result.ArgValue =
         cast<ConstantInt>(getValueFromBundleOpInfo(Assume, BOI, ABA_Argument))
             ->getZExtValue();
-
   return Result;
 }
 
+RetainedKnowledge llvm::getKnowledgeFromOperandInAssume(CallInst &AssumeCI,
+                                                        unsigned Idx) {
+  IntrinsicInst &Assume = cast<IntrinsicInst>(AssumeCI);
+  assert(Assume.getIntrinsicID() == Intrinsic::assume &&
+         "this function is intended to be used on llvm.assume");
+  CallBase::BundleOpInfo BOI = Assume.getBundleOpInfoForOperand(Idx);
+  return getKnowledgeFromBundle(AssumeCI, BOI);
+}
+
 bool llvm::isAssumeWithEmptyBundle(CallInst &CI) {
   IntrinsicInst &Assume = cast<IntrinsicInst>(CI);
   assert(Assume.getIntrinsicID() == Intrinsic::assume &&
          "this function is intended to be used on llvm.assume");
   return none_of(Assume.bundle_op_infos(),
                  [](const CallBase::BundleOpInfo &BOI) {
                    return BOI.Tag->getKey() != "ignore";
                  });
 }
+
+RetainedKnowledge
+llvm::getKnowledgeFromUse(const Use *U,
+                          ArrayRef<Attribute::AttrKind> AttrKinds) {
+  if (!match(U->getUser(),
+             m_Intrinsic<Intrinsic::assume>(m_Unless(m_Specific(U->get())))))
+    return RetainedKnowledge::none();
+  auto *Intr = cast<IntrinsicInst>(U->getUser());
+  RetainedKnowledge RK =
+      getKnowledgeFromOperandInAssume(*Intr, U->getOperandNo());
+  for (auto Attr : AttrKinds)
+    if (Attr == RK.AttrKind)
+      return RK;
+  return RetainedKnowledge::none();
+}
+
+RetainedKnowledge llvm::getKnowledgeForValue(
+    const Value *V, ArrayRef<Attribute::AttrKind> AttrKinds,
+    AssumptionCache *AC,
+    function_ref<bool(RetainedKnowledge, Instruction *)> Filter) {
+  if (AC) {
+#ifndef NDEBUG
+    RetainedKnowledge RKCheck =
+        getKnowledgeForValue(V, AttrKinds, nullptr, Filter);
+#endif
+    for (AssumptionCache::ResultElem &Elem : AC->assumptionsFor(V)) {
+      IntrinsicInst *II = cast_or_null<IntrinsicInst>(Elem.Assume);
+      if (!II || Elem.Index == AssumptionCache::ExprResultIdx)
+        continue;
+      if (RetainedKnowledge RK = getKnowledgeFromBundle(
+              *II, II->bundle_op_info_begin()[Elem.Index]))
+        if (is_contained(AttrKinds, RK.AttrKind) && Filter(RK, II)) {
+          assert(!!RKCheck && "invalid Assumption cache");
+          return RK;
+        }
+    }
+    assert(!RKCheck && "invalid Assumption cache");
+    return RetainedKnowledge::none();
+  }
+  for (auto &U : V->uses()) {
+    if (RetainedKnowledge RK = getKnowledgeFromUse(&U, AttrKinds))
+      if (Filter(RK, cast<Instruction>(U.getUser())))
+        return RK;
+  }
+  return RetainedKnowledge::none();
+}
+
+RetainedKnowledge llvm::getKnowledgeValidInContext(
+    const Value *V, ArrayRef<Attribute::AttrKind> AttrKinds,
+    const Instruction *CtxI, const DominatorTree *DT, AssumptionCache *AC) {
+  return getKnowledgeForValue(V, AttrKinds, AC,
+                              [&](RetainedKnowledge, Instruction *I) {
+                                return isValidAssumeForContext(I, CtxI, DT);
+                              });
+}

llvm/lib/Analysis/ValueTracking.cpp

 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/AssumeBundleQueries.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/GuardUtils.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/Argument.h"
auto CmpExcludesZero = [V](ICmpInst *Cmp) {
     if (!match(RHS, m_ConstantInt(CI)))
       return false;
     ConstantRange RHSRange(CI->getValue());
     ConstantRange TrueValues =
         ConstantRange::makeAllowedICmpRegion(Pred, RHSRange);
     return !TrueValues.contains(APInt::getNullValue(CI->getBitWidth()));
   };
 
+  if (Q.CxtI && V->getType()->isPointerTy()) {
+    SmallVector<Attribute::AttrKind, 2> AttrKinds{Attribute::NonNull};
+    if (!NullPointerIsDefined(Q.CxtI->getFunction(),
+                              V->getType()->getPointerAddressSpace()))
+      AttrKinds.push_back(Attribute::Dereferenceable);
+
+    if (getKnowledgeValidInContext(V, AttrKinds, Q.CxtI, Q.DT, Q.AC))
+      return true;
+  }
+
   for (auto &AssumeVH : Q.AC->assumptionsFor(V)) {
     if (!AssumeVH)
       continue;
     CallInst *I = cast<CallInst>(AssumeVH);
     assert(I->getFunction() == Q.CxtI->getFunction() &&
            "Got assumption for the wrong function!");
     if (Q.isExcluded(I))
       continue;
 
     // Warning: This loop can end up being somewhat performance sensitive.
     // We're running this loop for once for each value queried resulting in a
     // runtime of ~O(#assumes * #values).
 
     assert(I->getCalledFunction()->getIntrinsicID() == Intrinsic::assume &&

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

 // more information.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/Attributor.h"
 
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/AssumeBundleQueries.h"
 #include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/LazyValueInfo.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/NoFolder.h"
 #include "llvm/Transforms/IPO/ArgumentPromotion.h"
if (!UseV->getType()->isPointerTy())
     return 0;
 
   Type *PtrTy = UseV->getType();
   const Function *F = I->getFunction();
   bool NullPointerIsDefined =
       F ? llvm::NullPointerIsDefined(F, PtrTy->getPointerAddressSpace()) : true;
   const DataLayout &DL = A.getInfoCache().getDL();
   if (ImmutableCallSite ICS = ImmutableCallSite(I)) {
-    if (ICS.isBundleOperand(U))
+    if (ICS.isBundleOperand(U)) {
+      if (RetainedKnowledge RK = getKnowledgeFromUse(
+              U, {Attribute::NonNull, Attribute::Dereferenceable})) {
+        IsNonNull |=
+            (RK.AttrKind == Attribute::NonNull || !NullPointerIsDefined);
+        return RK.ArgValue;
+      }
       return 0;
+    }
 
     if (ICS.isCallee(U)) {
       IsNonNull |= !NullPointerIsDefined;
       return 0;
     }
 
     unsigned ArgNo = ICS.getArgumentNo(U);
     IRPosition IRP = IRPosition::callsite_argument(ICS, ArgNo);

llvm/test/Analysis/ValueTracking/assume.ll

;
   call void @llvm.assume(i1 %t2)
   %t3 = add i32 %t1, 3
   ret i32 %t3
 }
 
 
 define void @assume_not() {
 ; CHECK-LABEL: @assume_not(
+; CHECK-NEXT:  entry-block:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i1 @get_val()
+; CHECK-NEXT:    [[TMP1:%.*]] = xor i1 [[TMP0]], true
+; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP1]])
+; CHECK-NEXT:    ret void
+;
 entry-block:
   %0 = call i1 @get_val()
-; CHECK: call void @llvm.assume
   %1 = xor i1 %0, true
   call void @llvm.assume(i1 %1)
   ret void
 }
 
 declare i1 @get_val()
 declare void @llvm.assume(i1)
+
+define dso_local i1 @test1(i32* readonly %0) {
+; CHECK-LABEL: @test1(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i32* [[TMP0:%.*]]) ]
+; CHECK-NEXT:    ret i1 false
+;
+  call void @llvm.assume(i1 true) ["nonnull"(i32* %0)]
+  %2 = icmp eq i32* %0, null
+  ret i1 %2
+}
+
+define dso_local i1 @test2(i32* readonly %0) {
+; CHECK-LABEL: @test2(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i32* [[TMP0:%.*]]) ]
+; CHECK-NEXT:    ret i1 false
+;
+  %2 = icmp eq i32* %0, null
+  call void @llvm.assume(i1 true) ["nonnull"(i32* %0)]
+  ret i1 %2
+}
+
+define dso_local i32 @test4(i32* readonly %0, i1 %cond) {
+; CHECK-LABEL: @test4(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[TMP0:%.*]], i32 4) ]
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[A:%.*]], label [[B:%.*]]
+; CHECK:       B:
+; CHECK-NEXT:    br label [[A]]
+; CHECK:       A:
+; CHECK-NEXT:    br i1 false, label [[TMP4:%.*]], label [[TMP2:%.*]]
+; CHECK:       2:
+; CHECK-NEXT:    [[TMP3:%.*]] = load i32, i32* [[TMP0]], align 4
+; CHECK-NEXT:    br label [[TMP4]]
+; CHECK:       4:
+; CHECK-NEXT:    [[TMP5:%.*]] = phi i32 [ [[TMP3]], [[TMP2]] ], [ 0, [[A]] ]
+; CHECK-NEXT:    ret i32 [[TMP5]]
+;
+  call void @llvm.assume(i1 true) ["dereferenceable"(i32* %0, i32 4)]
+  br i1 %cond, label %A, label %B
+
+B:
+  br label %A
+
+A:
+  %2 = icmp eq i32* %0, null
+  br i1 %2, label %5, label %3
+
+3:                                                ; preds = %1
+  %4 = load i32, i32* %0, align 4
+  br label %5
+
+5:                                                ; preds = %1, %3
+  %6 = phi i32 [ %4, %3 ], [ 0, %A ]
+  ret i32 %6
+}
+
+define dso_local i32 @test4b(i32* readonly %0, i1 %cond) "null-pointer-is-valid"="true" {
+; CHECK-LABEL: @test4b(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[TMP0:%.*]], i32 4) ]
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[A:%.*]], label [[B:%.*]]
+; CHECK:       B:
+; CHECK-NEXT:    br label [[A]]
+; CHECK:       A:
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32* [[TMP0]], null
+; CHECK-NEXT:    br i1 [[TMP2]], label [[TMP5:%.*]], label [[TMP3:%.*]]
+; CHECK:       3:
+; CHECK-NEXT:    [[TMP4:%.*]] = load i32, i32* [[TMP0]], align 4
+; CHECK-NEXT:    br label [[TMP5]]
+; CHECK:       5:
+; CHECK-NEXT:    [[TMP6:%.*]] = phi i32 [ [[TMP4]], [[TMP3]] ], [ 0, [[A]] ]
+; CHECK-NEXT:    ret i32 [[TMP6]]
+;
+  call void @llvm.assume(i1 true) ["dereferenceable"(i32* %0, i32 4)]
+  br i1 %cond, label %A, label %B
+
+B:
+  br label %A
+
+A:
+  %2 = icmp eq i32* %0, null
+  br i1 %2, label %5, label %3
+
+3:                                                ; preds = %1
+  %4 = load i32, i32* %0, align 4
+  br label %5
+
+5:                                                ; preds = %1, %3
+  %6 = phi i32 [ %4, %3 ], [ 0, %A ]
+  ret i32 %6
+}

llvm/test/Transforms/Attributor/nonnull.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --scrub-attributes
 ; RUN: opt -attributor --attributor-disable=false -attributor-max-iterations-verify -attributor-annotate-decl-cs -attributor-max-iterations=5 -S < %s | FileCheck %s --check-prefixes=ATTRIBUTOR,ATTRIBUTOR_OPM
 ; RUN: opt -passes=attributor --attributor-disable=false -attributor-max-iterations-verify -attributor-annotate-decl-cs -attributor-max-iterations=5 -S < %s | FileCheck %s --check-prefixes=ATTRIBUTOR,ATTRIBUTOR_NPM
 ; Copied from Transforms/FunctoinAttrs/nonnull.ll
 
 ; UTC_ARGS: --disable
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 
 declare nonnull i8* @ret_nonnull()
+declare void @llvm.assume(i1)
 
 ; Return a pointer trivially nonnull (call return attribute)
 define i8* @test1() {
 ; ATTRIBUTOR: define nonnull i8* @test1
   %ret = call i8* @ret_nonnull()
   ret i8* %ret
 }
 
 ; Return a pointer trivially nonnull (argument attribute)
 define i8* @test2(i8* nonnull %p) {
 ; ATTRIBUTOR: define nonnull i8* @test2
   ret i8* %p
 }
 
+define i8* @test2A(i1 %c, i8* %ret) {
+; ATTRIBUTOR: define nonnull i8* @test2A(i1 %c, i8* nofree nonnull readnone returned %ret)
+  br i1 %c, label %A, label %B
+A:
+  call void @llvm.assume(i1 true) [ "nonnull"(i8* %ret) ]
+  ret i8* %ret
+B:
+  call void @llvm.assume(i1 true) [ "nonnull"(i8* %ret) ]
+  ret i8* %ret
+}
+
+define i8* @test2B(i1 %c, i8* %ret) {
+; ATTRIBUTOR: define nonnull dereferenceable(4) i8* @test2B(i1 %c, i8* nofree nonnull readnone returned dereferenceable(4) %ret)
+  br i1 %c, label %A, label %B
+A:
+  call void @llvm.assume(i1 true) [ "dereferenceable"(i8* %ret, i32 4) ]
+  ret i8* %ret
+B:
+  call void @llvm.assume(i1 true) [ "dereferenceable"(i8* %ret, i32 4) ]
+  ret i8* %ret
+}
+
 ; Given an SCC where one of the functions can not be marked nonnull,
 ; can we still mark the other one which is trivially nonnull
 define i8* @scc_binder(i1 %c) {
 ; ATTRIBUTOR: define noalias align 536870912 i8* @scc_binder
   br i1 %c, label %rec, label %end
 rec:
   call i8* @test3(i1 %c)
   br label %end

 }
 
 ; ATTRIBUTOR: define i8* @test9
 define i8* @test9(i8* %a, i64 %n) {
   %b = getelementptr inbounds i8, i8* %a, i64 %n
   ret i8* %b
 }
 
-declare void @llvm.assume(i1)
 ; ATTRIBUTOR_OPM: define i8* @test10
 ; ATTRIBUTOR_NPM: define nonnull i8* @test10
 define i8* @test10(i8* %a, i64 %n) {
   %cmp = icmp ne i64 %n, 0
   call void @llvm.assume(i1 %cmp)
   %b = getelementptr inbounds i8, i8* %a, i64 %n
   ret i8* %b
 }