[AssumeBundles] Use assume bundles in isKnownNonZero

[Tyker]

Summary

Use nonnull and dereferenceable from an assume bundle in isKnownNonZero

[jdoerfert]

llvm/lib/Analysis/ValueTracking.cpp

680

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

680

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

669

nit: is llvm:: needed?

llvm/lib/IR/KnowledgeRetention.cpp

340

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.

344

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

680

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

340

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

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

144

- 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

676

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

680

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

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

340

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

680

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

llvm/lib/IR/KnowledgeRetention.cpp

340

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

680

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

340

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

680

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

680

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

680

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

680

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

131 ↗

(On Diff #254810)

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))

145 ↗

(On Diff #254810)

nit: llvm:: not needed there I think

149 ↗

(On Diff #254810)

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

153 ↗

(On Diff #254810)

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

llvm/lib/Analysis/ValueTracking.cpp

667

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

149 ↗

(On Diff #254810)

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 ↗

(On Diff #257252)

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

19 ↗

(On Diff #257252)

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

20 ↗

(On Diff #257252)

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

104 ↗

(On Diff #257252)

unrelated nit: property that holds?

106 ↗

(On Diff #257252)

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

142 ↗

(On Diff #257252)

nit: Return

143 ↗

(On Diff #257252)

nit: period at end of sentence.

147 ↗

(On Diff #257252)

nit: Return

152 ↗

(On Diff #257252)

llvm:: not needed?

155 ↗

(On Diff #257252)

nit: Return

llvm/lib/Analysis/AssumeBundleQueries.cpp

133 ↗

(On Diff #257252)

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

llvm/lib/Transforms/IPO/Attributor.cpp

23

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 ↗

(On Diff #258259)

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

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

llvm/lib/Analysis/ValueTracking.cpp

678

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 ↗

(On Diff #258371)

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

135

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/IR/KnowledgeRetention.h

 // performing a transformation that may move and delete instructions as those
 // transformation may destroy or worsen information that can be derived from the
 // IR.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_ASSUMEBUILDER_H
 #define LLVM_TRANSFORMS_UTILS_ASSUMEBUILDER_H
 
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace llvm {
 
 /// Build a call to llvm.assume to preserve informations that can be derived
 /// from the given instruction.
 /// If no information derived from \p I, this call returns null.
 /// The returned instruction is not inserted anywhere.

 /// Represent one information held inside an operand bundle of an llvm.assume.
 /// AttrKind is the property that hold.
 /// WasOn if not null is that Value for which AttrKind holds.
 /// ArgValue is optionally an argument.
 struct RetainedKnowledge {
   Attribute::AttrKind AttrKind = Attribute::None;
   Value *WasOn = nullptr;
   unsigned ArgValue = 0;
+  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 true iff U is a Use in an llvm.assume bundle and this use is
+/// associated with any attributes of AttrKinds.
+RetainedKnowledge
+getRetainedKnowledgeFromUse(const Use *U,
+                            ArrayRef<Attribute::AttrKind> AttrKinds);
+
+/// Return true iff the value V has any attribute in AttrKinds in an the operand
+/// bundles of an llvm.assume and matches the Filter predicate.
+RetainedKnowledge getRetainedKnowledgeForValue(
+    const Value *V, ArrayRef<Attribute::AttrKind> AttrKinds,
+    llvm::function_ref<bool(RetainedKnowledge, Instruction *)> Filter =
+        [](RetainedKnowledge, Instruction *) { return true; });
+
 //===----------------------------------------------------------------------===//
 // Utilities for testing
 //===----------------------------------------------------------------------===//
 
 /// This pass will try to build an llvm.assume for every instruction in the
 /// function. Its main purpose is testing.
 struct AssumeBuilderPass : public PassInfoMixin<AssumeBuilderPass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
 } // namespace llvm
 
 #endif

llvm/lib/Analysis/ValueTracking.cpp

 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsAArch64.h"
 #include "llvm/IR/IntrinsicsX86.h"
+#include "llvm/IR/KnowledgeRetention.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.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 (!llvm::NullPointerIsDefined(Q.CxtI->getFunction(),
+                                    V->getType()->getPointerAddressSpace()))
+      AttrKinds.push_back(Attribute::Dereferenceable);
+
+    if (getRetainedKnowledgeForValue(
+            V, AttrKinds, [&](RetainedKnowledge, Instruction *Assume) {
+              return isValidAssumeForContext(Assume, Q.CxtI, Q.DT);
+            }))
+      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

llvm/lib/IR/KnowledgeRetention.cpp

bool llvm::isAssumeWithEmptyBundle(CallInst &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::getRetainedKnowledgeFromUse(const Use *U,
+                                  ArrayRef<Attribute::AttrKind> AttrKinds) {
+  auto *Intr = dyn_cast<IntrinsicInst>(U->getUser());
+  if (!Intr || Intr->getIntrinsicID() != Intrinsic::assume ||
+      &*Intr->args().begin() == U)
+    return RetainedKnowledge::none();
+  RetainedKnowledge RK =
+      getKnowledgeFromOperandInAssume(*Intr, U->getOperandNo());
+  for (auto Attr : AttrKinds)
+    if (Attr == RK.AttrKind)
+      return RK;
+  return RetainedKnowledge::none();
+}
+
+RetainedKnowledge llvm::getRetainedKnowledgeForValue(
+    const Value *V, ArrayRef<Attribute::AttrKind> AttrKinds,
+    llvm::function_ref<bool(RetainedKnowledge, Instruction *)> Filter) {
+  for (auto &U : V->uses()) {
+    if (RetainedKnowledge RK = getRetainedKnowledgeFromUse(&U, AttrKinds))
+      if (Filter(RK, cast<Instruction>(U.getUser())))
+        return RK;
+  }
+  return RetainedKnowledge::none();
+}
+
 PreservedAnalyses AssumeBuilderPass::run(Function &F,
                                          FunctionAnalysisManager &AM) {
   for (Instruction &I : instructions(F))
     if (Instruction *Assume = BuildAssumeFromInst(&I))
       Assume->insertBefore(&I);
   return PreservedAnalyses::all();
 }

llvm/lib/Transforms/IPO/Attributor.cpp

 
 #include "llvm/Transforms/IPO/Attributor.h"
 
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/EHPersonalities.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/LazyValueInfo.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/MustExecute.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/KnowledgeRetention.h"
 #include "llvm/IR/NoFolder.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/IPO/ArgumentPromotion.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.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 = getRetainedKnowledgeFromUse(
+              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 i8* @unknown()
 
 ; 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(i8* %ret) {
+; ATTRIBUTOR: define nonnull i8* @test2A(i8* nonnull returned %ret)
+  call void @llvm.assume(i1 true) [ "nonnull"(i8* %ret) ]
+  ret i8* %ret
+}
+
+define i8* @test2B(i8* %ret) {
+; ATTRIBUTOR: define nonnull dereferenceable(4) i8* @test2B(i8* nonnull returned dereferenceable(4) %ret)
+  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
 }
 

 ; ATTRIBUTOR: define void @test12(i8* nonnull %a)
 ; ATTRIBUTOR-NEXT: tail call void @test12_helper(i8* nonnull %a)
   tail call void @test12_helper(i8* %a)
   ret void
 }
 
 ; TEST 13
 ; Simple Argument Tests
-declare i8* @unknown()
 define void @test13_helper() {
   %nonnullptr = tail call i8* @ret_nonnull()
   %maybenullptr = tail call i8* @unknown()
   tail call void @test13(i8* %nonnullptr, i8* %nonnullptr, i8* %maybenullptr)
   tail call void @test13(i8* %nonnullptr, i8* %maybenullptr, i8* %nonnullptr)
   ret void
 }
 declare void @use_i8_ptr(i8* nofree nocapture readnone) nounwind