[AssumeBundles] Use assume bundles in isKnownNonZero

Summary: Use nonnull and dereferenceable from an assume bundle in isKnownNonZero

Reviewers: jdoerfert, nikic, lebedev.ri, reames, fhahn, sstefan1

Reviewed By: jdoerfert

Subscribers: fhahn, hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D76149

llvm/include/llvm/Analysis/AssumeBundleQueries.h

@@ -20,6 +20,8 @@
 
 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
@@ -98,13 +100,24 @@ using RetainedKnowledgeMap =
 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 four:
+///  - 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.
@@ -129,6 +142,27 @@ inline RetainedKnowledge getKnowledgeFromUseInAssume(const Use *U) {
 /// 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

@@ -7,17 +7,21 @@
 //===----------------------------------------------------------------------===//
 
 #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");
@@ -92,12 +96,8 @@ void llvm::fillMapFromAssume(CallInst &AssumeCI, RetainedKnowledgeMap &Result) {
   }
 }
 
-RetainedKnowledge llvm::getKnowledgeFromOperandInAssume(CallInst &AssumeCI,
+static RetainedKnowledge
-                                                        unsigned Idx) {
+getKnowledgeFromBundle(CallInst &Assume, const CallBase::BundleOpInfo &BOI) {
-  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);
   RetainedKnowledge Result;
   Result.AttrKind = Attribute::getAttrKindFromName(BOI.Tag->getKey());
   Result.WasOn = getValueFromBundleOpInfo(Assume, BOI, ABA_WasOn);
@@ -105,10 +105,18 @@ RetainedKnowledge llvm::getKnowledgeFromOperandInAssume(CallInst &AssumeCI,
     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 &&
@@ -118,3 +126,58 @@ bool llvm::isAssumeWithEmptyBundle(CallInst &CI) {
                    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

@@ -24,6 +24,7 @@
 #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"
@@ -688,6 +689,16 @@ static bool isKnownNonZeroFromAssume(const Value *V, const Query &Q) {
     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;

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

@@ -15,6 +15,7 @@
 
 #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"
@@ -1581,8 +1582,15 @@ static int64_t getKnownNonNullAndDerefBytesForUse(
       F ? llvm::NullPointerIsDefined(F, PtrTy->getPointerAddressSpace()) : true;
   const DataLayout &DL = A.getInfoCache().getDL();
   if (const auto *CB = dyn_cast<CallBase>(I)) {
-    if (CB->isBundleOperand(U))
+    if (CB->isBundleOperand(U)) {
+      if (RetainedKnowledge RK = getKnowledgeFromUse(
+              U, {Attribute::NonNull, Attribute::Dereferenceable})) {
+        IsNonNull |=
+            (RK.AttrKind == Attribute::NonNull || !NullPointerIsDefined);
+        return RK.ArgValue;
+      }
       return 0;
+    }
 
     if (CB->isCallee(U)) {
       IsNonNull |= !NullPointerIsDefined;

llvm/test/Analysis/ValueTracking/assume.ll

@@ -21,9 +21,14 @@ define i32 @assume_add(i32 %a, i32 %b) {
 
 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
@@ -31,3 +36,92 @@ entry-block:
 
 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

@@ -8,6 +8,7 @@
 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() {
@@ -28,6 +29,28 @@ define i8* @test2(i8* nonnull %p) {
   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) {
@@ -181,7 +204,8 @@ define i8* @test9(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) {
 ; CHECK-LABEL: define {{[^@]+}}@test10
 ; CHECK-SAME: (i8* nofree readnone "no-capture-maybe-returned" [[A:%.*]], i64 [[N:%.*]])
@@ -664,7 +688,7 @@ declare i32 @esfp(...)
 
 define i1 @parent8(i8* %a, i8* %bogus1, i8* %b) personality i8* bitcast (i32 (...)* @esfp to i8*){
 ; NOT_CGSCC_OPM-LABEL: define {{[^@]+}}@parent8
-; NOT_CGSCC_OPM-SAME: (i8* nonnull [[A:%.*]], i8* nocapture nofree readnone [[BOGUS1:%.*]], i8* nonnull [[B:%.*]]) #4 personality i8* bitcast (i32 (...)* @esfp to i8*)
+; NOT_CGSCC_OPM-SAME: (i8* nonnull [[A:%.*]], i8* nocapture nofree readnone [[BOGUS1:%.*]], i8* nonnull [[B:%.*]]) {{#[0-9]+}} personality i8* bitcast (i32 (...)* @esfp to i8*)
 ; NOT_CGSCC_OPM-NEXT:  entry:
 ; NOT_CGSCC_OPM-NEXT:    invoke void @use2nonnull(i8* nonnull [[A]], i8* nonnull [[B]])
 ; NOT_CGSCC_OPM-NEXT:    to label [[CONT:%.*]] unwind label [[EXC:%.*]]
@@ -677,7 +701,7 @@ define i1 @parent8(i8* %a, i8* %bogus1, i8* %b) personality i8* bitcast (i32 (..
 ; NOT_CGSCC_OPM-NEXT:    unreachable
 ;
 ; IS__CGSCC_OPM-LABEL: define {{[^@]+}}@parent8
-; IS__CGSCC_OPM-SAME: (i8* nonnull [[A:%.*]], i8* nocapture nofree readnone [[BOGUS1:%.*]], i8* nonnull [[B:%.*]]) #5 personality i8* bitcast (i32 (...)* @esfp to i8*)
+; IS__CGSCC_OPM-SAME: (i8* nonnull [[A:%.*]], i8* nocapture nofree readnone [[BOGUS1:%.*]], i8* nonnull [[B:%.*]]) {{#[0-9]+}} personality i8* bitcast (i32 (...)* @esfp to i8*)
 ; IS__CGSCC_OPM-NEXT:  entry:
 ; IS__CGSCC_OPM-NEXT:    invoke void @use2nonnull(i8* nonnull [[A]], i8* nonnull [[B]])
 ; IS__CGSCC_OPM-NEXT:    to label [[CONT:%.*]] unwind label [[EXC:%.*]]