[Attributor] Use knowledge retained in llvm.assume (operand bundles)

[jdoerfert]

Summary

This patch integrates operand bundle llvm.assumes [0] with the
Attributor. Most IRAttributes will now look at uses of the associated
value and if there are llvm.assume operand bundle uses with the right
tag we will check if they are in the must-be-executed-context (around
the context instruction).

This includes a bug fix for KnowledgeRetention.

[0] http://lists.llvm.org/pipermail/llvm-dev/2019-December/137632.html

[jdoerfert]

Rebase and more tests

[jdoerfert]

Simplify code

[uenoku]

llvm/lib/Transforms/IPO/Attributor.cpp

680

I think using nullptr as a flag to indicate subsuming positions seems a bit difficult to read.
Could you simply add some boolean flag and replace the argument with a reference?

744–749

Can't we share AssumeAttrMap with other AAs whose associated values are the same?

[jdoerfert]

llvm/lib/Transforms/IPO/Attributor.cpp

680

I'm not sure I understand. If A is null, you don't get any assume handling. If A is not null, you get assume handling but only for the current position. not any subsuming one.

744–749

Yes, that is the plan. In KnowledgeRetention.{h/cpp} there is a TODO that says we should provide an API to get all information from an llvm.assume in a map. Once that is in place, the Attributor should be the one doing the querying (once) and then caching the information. Does that make sense?

I hoped @Tyker will look into that API at some point.

[Tyker]

llvm/lib/Transforms/IPO/Attributor.cpp

744–749

I added a revision (D75020) with what i thought about when talking about the map API. I hope it fits with what you had in mind.

but from looking at the code it seems like what could be the most usefull is a way to get from a use in the operand bundles to the knowledge associated with it quickly. and it can be done much quicker than hasAttributeInAssume and may be generally useful.

also sorry for the bug in hasAttributeInAssume.

[uenoku]

llvm/lib/Transforms/IPO/Attributor.cpp

680

I meant that the following implementation is more readable for me. Either way is fine anyway.

bool IRPosition::hasAttr(ArrayRef<Attribute::AttrKind> AKs,
                         bool IgnoreSubsumingPositions, Attributor &A) const {
  ...
  bool IsSubsumingPosition = false;
  for (const IRPosition &EquivIRP : SubsumingPositionIterator(*this)) {
    for (Attribute::AttrKind AK : AKs) {
      ...
      if (!IsSubsumingPosition && EquivIRP.getAttrsFromAssumes(AK, Attrs, A))
        return true;
    }
    ...
    IsSubsumingPosition = true;
  }
  return false;
}

744–749

Yes, that is the plan. In KnowledgeRetention.{h/cpp} there is a TODO that says we should provide an API to get all information from an llvm.assume in a map. Once that is in place, the Attributor should be the one doing the querying (once) and then caching the information. Does that make sense?

Ok.

[jdoerfert]

Use (an extended) map interface to retain (and cache) knowledge

[uenoku]

LGTM

llvm/include/llvm/IR/KnowledgeRetention.h

 #define LLVM_TRANSFORMS_UTILS_ASSUMEBUILDER_H
 
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace llvm {
+class IntrinsicInst;
 
 /// 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.
 CallInst *BuildAssumeFromInst(const Instruction *I, Module *M);
 inline CallInst *BuildAssumeFromInst(Instruction *I) {
   return BuildAssumeFromInst(I, I->getModule());

 /// If the Attribute is not on any value, the Value is nullptr.
 using RetainedKnowledgeKey = std::pair<Value *, Attribute::AttrKind>;
 
 struct MinMax {
   unsigned Min;
   unsigned Max;
 };
 
-using RetainedKnowledgeMap = DenseMap<RetainedKnowledgeKey, MinMax>;
+/// A mapping from intrinsics (=`llvm.assume` calls) to a value range
+/// (=knowledge) that is encoded in them. How the value range is interpreted
+/// depends on the RetainedKnowledgeKey that was used to get this out of the
+/// RetainedKnowledgeMap.
+using Assume2KnowledgeMap = DenseMap<IntrinsicInst *, MinMax>;
+
+using RetainedKnowledgeMap =
+    DenseMap<RetainedKnowledgeKey, Assume2KnowledgeMap>;
 
 /// 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);
 

llvm/include/llvm/Transforms/IPO/Attributor.h

 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/MustExecute.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/ConstantRange.h"
+#include "llvm/IR/KnowledgeRetention.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
 
 namespace llvm {
 
+struct Attributor;
 struct AbstractAttribute;
 struct InformationCache;
 struct AAIsDead;
 
 class Function;
 
 /// Simple enum classes that forces properties to be spelled out explicitly.
 ///
struct IRPosition {
   ///       here or somewhere else.
 
   /// Return true if any kind in \p AKs existing in the IR at a position that
   /// will affect this one. See also getAttrs(...).
   /// \param IgnoreSubsumingPositions Flag to determine if subsuming positions,
   ///                                 e.g., the function position if this is an
   ///                                 argument position, should be ignored.
   bool hasAttr(ArrayRef<Attribute::AttrKind> AKs,
-               bool IgnoreSubsumingPositions = false) const;
+               bool IgnoreSubsumingPositions = false,
+               Attributor *A = nullptr) const;
 
   /// Return the attributes of any kind in \p AKs existing in the IR at a
   /// position that will affect this one. While each position can only have a
   /// single attribute of any kind in \p AKs, there are "subsuming" positions
   /// that could have an attribute as well. This method returns all attributes
   /// found in \p Attrs.
   /// \param IgnoreSubsumingPositions Flag to determine if subsuming positions,
   ///                                 e.g., the function position if this is an
   ///                                 argument position, should be ignored.
   void getAttrs(ArrayRef<Attribute::AttrKind> AKs,
                 SmallVectorImpl<Attribute> &Attrs,
-                bool IgnoreSubsumingPositions = false) const;
+                bool IgnoreSubsumingPositions = false,
+                Attributor *A = nullptr) const;
 
   /// Remove the attribute of kind \p AKs existing in the IR at this position.
   void removeAttrs(ArrayRef<Attribute::AttrKind> AKs) const {
     if (getPositionKind() == IRP_INVALID || getPositionKind() == IRP_FLOAT)
       return;
 
     AttributeList AttrList;
     CallSite CS = CallSite(&getAnchorValue());
private:
 
   /// Verify internal invariants.
   void verify();
 
   /// Return the attributes of kind \p AK existing in the IR as attribute.
   bool getAttrsFromIRAttr(Attribute::AttrKind AK,
                           SmallVectorImpl<Attribute> &Attrs) const;
 
+  /// Return the attributes of kind \p AK existing in the IR as operand bundles
+  /// of an llvm.assume.
+  bool getAttrsFromAssumes(Attribute::AttrKind AK,
+                           SmallVectorImpl<Attribute> &Attrs,
+                           Attributor &A) const;
+
 protected:
   /// The value this position is anchored at.
   Value *AnchorVal;
 
   /// If AnchorVal is Argument or CallBase then this number should be
   /// non-negative and it denotes the argument or call site argument index
   /// respectively. Otherwise, it denotes the kind of this IRPosition according
   /// to Kind above.
unsigned getSccSize(const Function &F) {
     if (CGSCC && CGSCC->count(const_cast<Function *>(&F)))
       return CGSCC->size();
     return 0;
   }
 
   /// Return datalayout used in the module.
   const DataLayout &getDL() { return DL; }
 
+  /// Return the map conaining all the knowledge we have from `llvm.assume`s.
+  const RetainedKnowledgeMap &getKnowledgeMap() const { return KnowledgeMap; }
+
 private:
   /// A map type from functions to opcode to instruction maps.
   using FuncInstOpcodeMapTy = DenseMap<const Function *, OpcodeInstMapTy>;
 
   /// A map type from functions to their read or write instructions.
   using FuncRWInstsMapTy = DenseMap<const Function *, InstructionVectorTy>;
 
   /// A nested map that remembers all instructions in a function with a certain
   /// instruction opcode (Instruction::getOpcode()).
   FuncInstOpcodeMapTy FuncInstOpcodeMap;
 
   /// A map from functions to their instructions that may read or write memory.
   FuncRWInstsMapTy FuncRWInstsMap;
 
   /// The datalayout used in the module.
   const DataLayout &DL;
 
   /// MustBeExecutedContextExplorer
   MustBeExecutedContextExplorer Explorer;
 
+  /// A map with knowledge retained in `llvm.assume` instructions.
+  RetainedKnowledgeMap KnowledgeMap;
+
   /// Getters for analysis.
   AnalysisGetter &AG;
 
   /// The underlying CGSCC, or null if not available.
   SetVector<Function *> *CGSCC;
 
   /// Set of inlineable functions
   SmallPtrSet<const Function *, 8> InlineableFunctions;

 template <Attribute::AttrKind AK, typename Base>
 struct IRAttribute : public IRPosition, public Base {
   IRAttribute(const IRPosition &IRP) : IRPosition(IRP) {}
   ~IRAttribute() {}
 
   /// See AbstractAttribute::initialize(...).
   virtual void initialize(Attributor &A) override {
     const IRPosition &IRP = this->getIRPosition();
-    if (isa<UndefValue>(IRP.getAssociatedValue()) || hasAttr(getAttrKind())) {
+    if (isa<UndefValue>(IRP.getAssociatedValue()) ||
+        hasAttr(getAttrKind(), /* IgnoreSubsumingPositions */ false, &A)) {
       this->getState().indicateOptimisticFixpoint();
       return;
     }
 
     bool IsFnInterface = IRP.isFnInterfaceKind();
     const Function *FnScope = IRP.getAnchorScope();
     // TODO: Not all attributes require an exact definition. Find a way to
     //       enable deduction for some but not all attributes in case the

llvm/lib/IR/KnowledgeRetention.cpp

static Value *getValueFromBundleOpInfo(IntrinsicInst &Assume,
                                 unsigned Idx) {
   assert(BundleHasArguement(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,
                                 AssumeQuery AQR) {
+  assert(isa<IntrinsicInst>(AssumeCI) &&
+         "this function is intended to be used on llvm.assume");
   IntrinsicInst &Assume = cast<IntrinsicInst>(AssumeCI);
   assert(Assume.getIntrinsicID() == Intrinsic::assume &&
          "this function is intended to be used on llvm.assume");
   assert(Attribute::isExistingAttribute(AttrName) &&
          "this attribute doesn't exist");
   assert((ArgVal == nullptr || Attribute::doesAttrKindHaveArgument(
                                    Attribute::getAttrKindFromName(AttrName))) &&
          "requested value for an attribute that has no argument");
for (auto &Bundles : Assume.bundle_op_infos()) {
     std::pair<Value *, Attribute::AttrKind> Key{
         nullptr, Attribute::getAttrKindFromName(Bundles.Tag->getKey())};
     if (BundleHasArguement(Bundles, BOIE_WasOn))
       Key.first = getValueFromBundleOpInfo(Assume, Bundles, BOIE_WasOn);
 
     if (Key.first == nullptr && Key.second == Attribute::None)
       continue;
     if (!BundleHasArguement(Bundles, BOIE_Argument)) {
-      Result[Key] = {0, 0};
+      Result[Key][&Assume] = {0, 0};
       continue;
     }
     unsigned Val = cast<ConstantInt>(
                        getValueFromBundleOpInfo(Assume, Bundles, BOIE_Argument))
                        ->getZExtValue();
     auto Lookup = Result.find(Key);
-    if (Lookup == Result.end()) {
-      Result[Key] = {Val, Val};
+    if (Lookup == Result.end() || !Lookup->second.count(&Assume)) {
+      Result[Key][&Assume] = {Val, Val};
       continue;
     }
-    Lookup->second.Min = std::min(Val, Lookup->second.Min);
-    Lookup->second.Max = std::max(Val, Lookup->second.Max);
+    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");

llvm/lib/Transforms/IPO/Attributor.cpp

case IRPosition::IRP_CALL_SITE_ARGUMENT: {
     }
     IRPositions.emplace_back(IRPosition::value(IRP.getAssociatedValue()));
     return;
   }
   }
 }
 
 bool IRPosition::hasAttr(ArrayRef<Attribute::AttrKind> AKs,
-                         bool IgnoreSubsumingPositions) const {
+                         bool IgnoreSubsumingPositions, Attributor *A) const {
   SmallVector<Attribute, 4> Attrs;
   for (const IRPosition &EquivIRP : SubsumingPositionIterator(*this)) {
     for (Attribute::AttrKind AK : AKs)
       if (EquivIRP.getAttrsFromIRAttr(AK, Attrs))
         return true;
     // The first position returned by the SubsumingPositionIterator is
     // always the position itself. If we ignore subsuming positions we
     // are done after the first iteration.
     if (IgnoreSubsumingPositions)
       break;
   }
+  if (A)
+    for (Attribute::AttrKind AK : AKs)
+      if (getAttrsFromAssumes(AK, Attrs, *A))
+        return true;
   return false;
 }
 
 void IRPosition::getAttrs(ArrayRef<Attribute::AttrKind> AKs,
                           SmallVectorImpl<Attribute> &Attrs,
-                          bool IgnoreSubsumingPositions) const {
+                          bool IgnoreSubsumingPositions, Attributor *A) const {
   for (const IRPosition &EquivIRP : SubsumingPositionIterator(*this)) {
     for (Attribute::AttrKind AK : AKs)
       EquivIRP.getAttrsFromIRAttr(AK, Attrs);
     // The first position returned by the SubsumingPositionIterator is
     // always the position itself. If we ignore subsuming positions we
     // are done after the first iteration.
     if (IgnoreSubsumingPositions)
       break;
   }
+  if (A)
+    for (Attribute::AttrKind AK : AKs)
+     getAttrsFromAssumes(AK, Attrs, *A);
 }
 
 bool IRPosition::getAttrsFromIRAttr(Attribute::AttrKind AK,
                                     SmallVectorImpl<Attribute> &Attrs) const {
   if (getPositionKind() == IRP_INVALID || getPositionKind() == IRP_FLOAT)
     return false;
 
   AttributeList AttrList;
   if (ImmutableCallSite ICS = ImmutableCallSite(&getAnchorValue()))
     AttrList = ICS.getAttributes();
   else
     AttrList = getAssociatedFunction()->getAttributes();
 
   bool HasAttr = AttrList.hasAttribute(getAttrIdx(), AK);
   if (HasAttr)
     Attrs.push_back(AttrList.getAttribute(getAttrIdx(), AK));
   return HasAttr;
 }
 
+bool IRPosition::getAttrsFromAssumes(Attribute::AttrKind AK,
+                                     SmallVectorImpl<Attribute> &Attrs,
+                                     Attributor &A) const {
+  assert(getPositionKind() != IRP_INVALID && "Did expect a valid position!");
+  Value &AssociatedValue = getAssociatedValue();
+
+  const Assume2KnowledgeMap &A2K =
+      A.getInfoCache().getKnowledgeMap().lookup({&AssociatedValue, AK});
+
+  // Check if we found any potential assume use, if not we don't need to create
+  // explorer iterators.
+  if (A2K.empty())
+    return false;
+
+  LLVMContext &Ctx = AssociatedValue.getContext();
+  unsigned AttrsSize = Attrs.size();
+  MustBeExecutedContextExplorer &Explorer =
+      A.getInfoCache().getMustBeExecutedContextExplorer();
+  auto EIt = Explorer.begin(getCtxI()), EEnd = Explorer.end(getCtxI());
+  for (auto &It : A2K)
+    if (Explorer.findInContextOf(It.first, EIt, EEnd))
+      Attrs.push_back(Attribute::get(Ctx, AK, It.second.Max));
+  return AttrsSize != Attrs.size();
+}
+
 void IRPosition::verify() {
   switch (KindOrArgNo) {
   default:
     assert(KindOrArgNo >= 0 && "Expected argument or call site argument!");
     assert((isa<CallBase>(AnchorVal) || isa<Argument>(AnchorVal)) &&
            "Expected call base or argument for positive attribute index!");
     if (isa<Argument>(AnchorVal)) {
       assert(cast<Argument>(AnchorVal)->getArgNo() == unsigned(getArgNo()) &&
AANonNullImpl(const IRPosition &IRP)
       : AANonNull(IRP),
         NullIsDefined(NullPointerIsDefined(
             getAnchorScope(),
             getAssociatedValue().getType()->getPointerAddressSpace())) {}
 
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     if (!NullIsDefined &&
-        hasAttr({Attribute::NonNull, Attribute::Dereferenceable}))
+        hasAttr({Attribute::NonNull, Attribute::Dereferenceable},
+                /* IgnoreSubsumingPositions */ false, &A))
       indicateOptimisticFixpoint();
     else if (isa<ConstantPointerNull>(getAssociatedValue()))
       indicatePessimisticFixpoint();
     else
       AANonNull::initialize(A);
   }
 
   /// See AAFromMustBeExecutedContext

 
 struct AADereferenceableImpl : AADereferenceable {
   AADereferenceableImpl(const IRPosition &IRP) : AADereferenceable(IRP) {}
   using StateType = DerefState;
 
   void initialize(Attributor &A) override {
     SmallVector<Attribute, 4> Attrs;
     getAttrs({Attribute::Dereferenceable, Attribute::DereferenceableOrNull},
-             Attrs);
+             Attrs, /* IgnoreSubsumingPositions */ false, &A);
     for (const Attribute &Attr : Attrs)
       takeKnownDerefBytesMaximum(Attr.getValueAsInt());
 
     NonNullAA = &A.getAAFor<AANonNull>(*this, getIRPosition(),
                                        /* TrackDependence */ false);
 
     const IRPosition &IRP = this->getIRPosition();
     bool IsFnInterface = IRP.isFnInterfaceKind();
struct AACaptureUseTracker final : public CaptureTracker {
 
   /// Register \p CS as potential copy of the value we are checking.
   void addPotentialCopy(CallSite CS) {
     PotentialCopies.push_back(CS.getInstruction());
   }
 
   /// See CaptureTracker::shouldExplore(...).
   bool shouldExplore(const Use *U) override {
-    // Check liveness.
-    return !A.isAssumedDead(*U, &NoCaptureAA, &IsDeadAA);
+    // Check liveness and ignore droppable users.
+    return !U->getUser()->isDroppable() &&
+           !A.isAssumedDead(*U, &NoCaptureAA, &IsDeadAA);
   }
 
   /// Update the state according to \p CapturedInMem, \p CapturedInInt, and
   /// \p CapturedInRet, then return the appropriate value for use in the
   /// CaptureTracker::captured() interface.
   bool isCapturedIn(bool CapturedInMem, bool CapturedInInt,
                     bool CapturedInRet) {
     LLVM_DEBUG(dbgs() << " - captures [Mem " << CapturedInMem << "|Int "
for (unsigned i = 0; i < Uses.size() && !isAtFixpoint(); i++) {
     const Use *U = Uses[i];
     Instruction *UserI = cast<Instruction>(U->getUser());
     LLVM_DEBUG(dbgs() << "[AAMemoryBehavior] Use: " << **U << " in " << *UserI
                       << " [Dead: " << (A.isAssumedDead(*U, this, &LivenessAA))
                       << "]\n");
     if (A.isAssumedDead(*U, this, &LivenessAA))
       continue;
 
+    // Droppable users, e.g., llvm::assume does not actually perform any action.
+    if (UserI->isDroppable())
+      continue;
+
     // Check if the users of UserI should also be visited.
     if (followUsersOfUseIn(A, U, UserI))
       for (const Use &UserIUse : UserI->uses())
         Uses.insert(&UserIUse);
 
     // If UserI might touch memory we analyze the use in detail.
     if (UserI->mayReadOrWriteMemory())
       analyzeUseIn(A, U, UserI);
if (!Visited.insert(U).second)
       continue;
     LLVM_DEBUG(dbgs() << "[Attributor] Check use: " << **U << " in "
                       << *U->getUser() << "\n");
     if (isAssumedDead(*U, &QueryingAA, LivenessAA,
                       /* CheckBBLivenessOnly */ false, LivenessDepClass)) {
       LLVM_DEBUG(dbgs() << "[Attributor] Dead use, skip!\n");
       continue;
     }
+    if (U->getUser()->isDroppable()) {
+      LLVM_DEBUG(dbgs() << "[Attributor] Droppable user, skip!\n");
+      continue;
+    }
 
     bool Follow = false;
     if (!Pred(*U, Follow))
       return false;
     if (!Follow)
       continue;
     for (const Use &UU : U->getUser()->uses())
       Worklist.push_back(&UU);
for (Instruction &I : instructions(&F)) {
     // the following switch.
     // Note: There are no concrete attributes now so this is initially empty.
     switch (I.getOpcode()) {
     default:
       assert((!ImmutableCallSite(&I)) && (!isa<CallBase>(&I)) &&
              "New call site/base instruction type needs to be known in the "
              "Attributor.");
       break;
+    case Instruction::Call:
+      // Calls are interesting but for `llvm.assume` calls we also fill the
+      // KnowledgeMap as we find them.
+      if (IntrinsicInst *Assume = dyn_cast<IntrinsicInst>(&I)) {
+        if (Assume->getIntrinsicID() == Intrinsic::assume)
+          fillMapFromAssume(*Assume, InfoCache.KnowledgeMap);
+      }
+      LLVM_FALLTHROUGH;
     case Instruction::Load:
       // The alignment of a pointer is interesting for loads.
     case Instruction::Store:
       // The alignment of a pointer is interesting for stores.
-    case Instruction::Call:
     case Instruction::CallBr:
     case Instruction::Invoke:
     case Instruction::CleanupRet:
     case Instruction::CatchSwitch:
     case Instruction::AtomicRMW:
     case Instruction::AtomicCmpXchg:
     case Instruction::Br:
     case Instruction::Resume:

llvm/test/Transforms/Attributor/dereferenceable-1.ll

 if.else6:                                         ; preds = %if.else3
   tail call void @rec-branch-2(i32 1, i32 1, i32 1, i32* %ptr)
   br label %if.end8
 
 if.end8:                                          ; preds = %if.then5, %if.else6, %if.then2, %if.else
   ret void
 }
 
+declare void @unknown()
+define void @nonnull_assume_pos(i8* %arg1, i8* %arg2, i8* %arg3, i8* %arg4) {
+; ATTRIBUTOR-LABEL: define {{[^@]+}}@nonnull_assume_pos
+; ATTRIBUTOR-SAME: (i8* nocapture nofree nonnull readnone dereferenceable(101) [[ARG1:%.*]], i8* nocapture nofree readnone dereferenceable_or_null(31) [[ARG2:%.*]], i8* nocapture nofree nonnull readnone [[ARG3:%.*]], i8* nocapture nofree readnone dereferenceable_or_null(42) [[ARG4:%.*]])
+; ATTRIBUTOR-NEXT:    call void @llvm.assume(i1 true) #6 [ "nonnull"(i8* undef), "dereferenceable"(i8* undef, i64 1), "dereferenceable"(i8* undef, i64 2), "dereferenceable"(i8* undef, i64 101), "dereferenceable_or_null"(i8* undef, i64 31), "dereferenceable_or_null"(i8* undef, i64 42) ]
+; ATTRIBUTOR-NEXT:    call void @unknown()
+; ATTRIBUTOR-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) [ "nonnull"(i8* %arg3), "dereferenceable"(i8* %arg1, i64 1), "dereferenceable"(i8* %arg1, i64 2), "dereferenceable"(i8* %arg1, i64 101), "dereferenceable_or_null"(i8* %arg2, i64 31), "dereferenceable_or_null"(i8* %arg4, i64 42)]
+  call void @unknown()
+  ret void
+}
+define void @nonnull_assume_neg(i8* %arg1, i8* %arg2, i8* %arg3) {
+; ATTRIBUTOR-LABEL: define {{[^@]+}}@nonnull_assume_neg
+; ATTRIBUTOR-SAME: (i8* nocapture nofree readnone [[ARG1:%.*]], i8* nocapture nofree readnone [[ARG2:%.*]], i8* nocapture nofree readnone [[ARG3:%.*]])
+; ATTRIBUTOR-NEXT:    call void @unknown()
+; ATTRIBUTOR-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i8* undef, i64 101), "dereferenceable"(i8* undef, i64 -2), "dereferenceable_or_null"(i8* undef, i64 31) ]
+; ATTRIBUTOR-NEXT:    ret void
+;
+  call void @unknown()
+  call void @llvm.assume(i1 true) ["dereferenceable"(i8* %arg1, i64 101), "dereferenceable"(i8* %arg2, i64 -2), "dereferenceable_or_null"(i8* %arg3, i64 31)]
+  ret void
+}
+define void @nonnull_assume_call(i8* %arg1, i8* %arg2, i8* %arg3, i8* %arg4) {
+; ATTRIBUTOR-LABEL: define {{[^@]+}}@nonnull_assume_call
+; ATTRIBUTOR-SAME: (i8* [[ARG1:%.*]], i8* [[ARG2:%.*]], i8* [[ARG3:%.*]], i8* [[ARG4:%.*]])
+; ATTRIBUTOR-NEXT:    call void @unknown()
+; ATTRIBUTOR-NEXT:    [[P:%.*]] = call nonnull dereferenceable(101) i32* @unkown_ptr()
+; ATTRIBUTOR-NEXT:    call void @unknown_use32(i32* nonnull dereferenceable(101) [[P]])
+; ATTRIBUTOR-NEXT:    call void @unknown_use8(i8* nonnull dereferenceable(42) [[ARG4]])
+; ATTRIBUTOR-NEXT:    call void @unknown_use8(i8* nonnull [[ARG3]])
+; ATTRIBUTOR-NEXT:    call void @unknown_use8(i8* nonnull dereferenceable(31) [[ARG2]])
+; ATTRIBUTOR-NEXT:    call void @unknown_use8(i8* nonnull dereferenceable(2) [[ARG1]])
+; ATTRIBUTOR-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i8* [[ARG3]]), "dereferenceable"(i8* [[ARG1]], i64 1), "dereferenceable"(i8* [[ARG1]], i64 2), "dereferenceable"(i32* [[P]], i64 101), "dereferenceable_or_null"(i8* [[ARG2]], i64 31), "dereferenceable_or_null"(i8* [[ARG4]], i64 42) ]
+; ATTRIBUTOR-NEXT:    call void @unknown_use8(i8* nonnull dereferenceable(2) [[ARG1]])
+; ATTRIBUTOR-NEXT:    call void @unknown_use8(i8* nonnull dereferenceable(31) [[ARG2]])
+; ATTRIBUTOR-NEXT:    call void @unknown_use8(i8* nonnull [[ARG3]])
+; ATTRIBUTOR-NEXT:    call void @unknown_use8(i8* nonnull dereferenceable(42) [[ARG4]])
+; ATTRIBUTOR-NEXT:    call void @unknown_use32(i32* nonnull dereferenceable(101) [[P]])
+; ATTRIBUTOR-NEXT:    call void @unknown()
+; ATTRIBUTOR-NEXT:    ret void
+;
+  call void @unknown()
+  %p = call i32* @unkown_ptr()
+  call void @unknown_use32(i32* %p)
+  call void @unknown_use8(i8* %arg4)
+  call void @unknown_use8(i8* %arg3)
+  call void @unknown_use8(i8* %arg2)
+  call void @unknown_use8(i8* %arg1)
+  call void @llvm.assume(i1 true) [ "nonnull"(i8* %arg3), "dereferenceable"(i8* %arg1, i64 1), "dereferenceable"(i8* %arg1, i64 2), "dereferenceable"(i32* %p, i64 101), "dereferenceable_or_null"(i8* %arg2, i64 31), "dereferenceable_or_null"(i8* %arg4, i64 42)]
+  call void @unknown_use8(i8* %arg1)
+  call void @unknown_use8(i8* %arg2)
+  call void @unknown_use8(i8* %arg3)
+  call void @unknown_use8(i8* %arg4)
+  call void @unknown_use32(i32* %p)
+  call void @unknown()
+  ret void
+}
+declare void @unknown_use8(i8*) willreturn nounwind
+declare void @unknown_use32(i32*) willreturn nounwind
+declare void @llvm.assume(i1)
+
 !0 = !{i64 10, i64 100}
 

llvm/test/Transforms/Attributor/nofree.ll

 }
 
 ; ATTRIBUTOR: define void @test14(i8* nocapture %0, i8* nocapture nofree readnone %1)
 define void @test14(i8* nocapture %0, i8* nocapture %1) {
   tail call void @free(i8* %0) #1
   ret void
 }
 
+; UTC_ARGS: --enable
+
+define void @nonnull_assume_pos(i8* %arg1, i8* %arg2, i8* %arg3, i8* %arg4) {
+; ATTRIBUTOR-LABEL: define {{[^@]+}}@nonnull_assume_pos
+; ATTRIBUTOR-SAME: (i8* nofree [[ARG1:%.*]], i8* [[ARG2:%.*]], i8* nofree [[ARG3:%.*]], i8* [[ARG4:%.*]])
+; ATTRIBUTOR-NEXT:    call void @llvm.assume(i1 true) #11 [ "nofree"(i8* [[ARG1]]), "nofree"(i8* [[ARG3]]) ]
+; ATTRIBUTOR-NEXT:    call void @unknown(i8* nofree [[ARG1]], i8* [[ARG2]], i8* nofree [[ARG3]], i8* [[ARG4]])
+; ATTRIBUTOR-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) ["nofree"(i8* %arg1), "nofree"(i8* %arg3)]
+  call void @unknown(i8* %arg1, i8* %arg2, i8* %arg3, i8* %arg4)
+  ret void
+}
+define void @nonnull_assume_neg(i8* %arg1, i8* %arg2, i8* %arg3, i8* %arg4) {
+; ATTRIBUTOR-LABEL: define {{[^@]+}}@nonnull_assume_neg
+; ATTRIBUTOR-SAME: (i8* [[ARG1:%.*]], i8* [[ARG2:%.*]], i8* [[ARG3:%.*]], i8* [[ARG4:%.*]])
+; ATTRIBUTOR-NEXT:    call void @unknown(i8* [[ARG1]], i8* [[ARG2]], i8* [[ARG3]], i8* [[ARG4]])
+; ATTRIBUTOR-NEXT:    call void @llvm.assume(i1 true) [ "nofree"(i8* [[ARG1]]), "nofree"(i8* [[ARG3]]) ]
+; ATTRIBUTOR-NEXT:    ret void
+;
+  call void @unknown(i8* %arg1, i8* %arg2, i8* %arg3, i8* %arg4)
+  call void @llvm.assume(i1 true) ["nofree"(i8* %arg1), "nofree"(i8* %arg3)]
+  ret void
+}
+define void @nonnull_assume_call(i8* %arg1, i8* %arg2, i8* %arg3, i8* %arg4) {
+; ATTRIBUTOR-LABEL: define {{[^@]+}}@nonnull_assume_call
+; ATTRIBUTOR-SAME: (i8* [[ARG1:%.*]], i8* [[ARG2:%.*]], i8* [[ARG3:%.*]], i8* [[ARG4:%.*]])
+; ATTRIBUTOR-NEXT:    call void @unknown(i8* [[ARG1]], i8* [[ARG2]], i8* [[ARG3]], i8* [[ARG4]])
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr(i8* noalias readnone [[ARG1]])
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr(i8* noalias readnone [[ARG2]])
+; ATTRIBUTOR-NEXT:    call void @llvm.assume(i1 true) [ "nofree"(i8* [[ARG1]]), "nofree"(i8* [[ARG3]]) ]
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr(i8* noalias nofree readnone [[ARG3]])
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr(i8* noalias readnone [[ARG4]])
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr_ret(i8* noalias nofree readnone [[ARG1]])
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr_ret(i8* noalias readnone [[ARG2]])
+; ATTRIBUTOR-NEXT:    call void @llvm.assume(i1 true) [ "nofree"(i8* [[ARG1]]), "nofree"(i8* [[ARG4]]) ]
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr_ret(i8* noalias nofree readnone [[ARG3]])
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr_ret(i8* noalias nofree readnone [[ARG4]])
+; ATTRIBUTOR-NEXT:    ret void
+;
+  call void @unknown(i8* %arg1, i8* %arg2, i8* %arg3, i8* %arg4)
+  call void @use_i8_ptr(i8* %arg1)
+  call void @use_i8_ptr(i8* %arg2)
+  call void @llvm.assume(i1 true) ["nofree"(i8* %arg1), "nofree"(i8* %arg3)]
+  call void @use_i8_ptr(i8* %arg3)
+  call void @use_i8_ptr(i8* %arg4)
+  call void @use_i8_ptr_ret(i8* %arg1)
+  call void @use_i8_ptr_ret(i8* %arg2)
+  call void @llvm.assume(i1 true) ["nofree"(i8* %arg1), "nofree"(i8* %arg4)]
+  call void @use_i8_ptr_ret(i8* %arg3)
+  call void @use_i8_ptr_ret(i8* %arg4)
+  ret void
+}
+declare void @llvm.assume(i1)
+declare void @unknown(i8*, i8*, i8*, i8*)
+declare void @use_i8_ptr(i8* nocapture readnone) nounwind
+declare void @use_i8_ptr_ret(i8* nocapture readnone) nounwind willreturn
+
 declare noalias i8* @malloc(i64)
 
 attributes #0 = { nounwind uwtable noinline }
 attributes #1 = { nounwind }
 attributes #2 = { nobuiltin nounwind }

llvm/test/Transforms/Attributor/nonnull.ll

 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
 define internal void @test13(i8* %a, i8* %b, i8* %c) {
 ; ATTRIBUTOR: define internal void @test13(i8* noalias nocapture nofree nonnull readnone %a, i8* noalias nocapture nofree readnone %b, i8* noalias nocapture nofree readnone %c)
   call void @use_i8_ptr(i8* %a)
   call void @use_i8_ptr(i8* %b)
   call void @use_i8_ptr(i8* %c)
   ret void
 }
 

 ;
   %call = call i8* @strrchr(i8* %str, i32 47)
   %tobool = icmp ne i8* %call, null
   %add.ptr = getelementptr inbounds i8, i8* %call, i64 1
   %cond = select i1 %tobool, i8* %add.ptr, i8* %str
   ret i8* %cond
 }
 
+define void @nonnull_assume_pos(i8* %arg) {
+; ATTRIBUTOR-LABEL: define {{[^@]+}}@nonnull_assume_pos
+; ATTRIBUTOR-SAME: (i8* nocapture nofree nonnull readnone [[ARG:%.*]])
+; ATTRIBUTOR-NEXT:    call void @llvm.assume(i1 true) #11 [ "nonnull"(i8* [[ARG]]) ]
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr(i8* noalias nocapture nofree nonnull readnone [[ARG]])
+; ATTRIBUTOR-NEXT:    [[TMP1:%.*]] = call i8* @unknown()
+; ATTRIBUTOR-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) ["nonnull"(i8* %arg)]
+  call void @use_i8_ptr(i8* %arg)
+  call i8* @unknown()
+  ret void
+}
+define void @nonnull_assume_neg(i8* %arg) {
+; ATTRIBUTOR-LABEL: define {{[^@]+}}@nonnull_assume_neg
+; ATTRIBUTOR-SAME: (i8* nocapture nofree readnone [[ARG:%.*]])
+; ATTRIBUTOR-NEXT:    [[TMP1:%.*]] = call i8* @unknown()
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr(i8* noalias nocapture nofree readnone [[ARG]])
+; ATTRIBUTOR-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i8* [[ARG]]) ]
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr(i8* noalias nocapture nofree nonnull readnone [[ARG]])
+; ATTRIBUTOR-NEXT:    [[TMP2:%.*]] = call i8* @unknown()
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr_ret(i8* noalias nocapture nofree nonnull readnone [[ARG]])
+; ATTRIBUTOR-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i8* [[ARG]]) ]
+; ATTRIBUTOR-NEXT:    call void @use_i8_ptr_ret(i8* noalias nocapture nofree nonnull readnone [[ARG]])
+; ATTRIBUTOR-NEXT:    ret void
+;
+  call i8* @unknown()
+  call void @use_i8_ptr(i8* %arg)
+  call void @llvm.assume(i1 true) ["nonnull"(i8* %arg)]
+  call void @use_i8_ptr(i8* %arg)
+  call i8* @unknown()
+  call void @use_i8_ptr_ret(i8* %arg)
+  call void @llvm.assume(i1 true) ["nonnull"(i8* %arg)]
+  call void @use_i8_ptr_ret(i8* %arg)
+  ret void
+}
+declare void @use_i8_ptr(i8* nofree nocapture readnone) nounwind
+declare void @use_i8_ptr_ret(i8* nofree nocapture readnone) nounwind willreturn
+
 attributes #0 = { "null-pointer-is-valid"="true" }
 attributes #1 = { nounwind willreturn}

llvm/unittests/IR/KnowledgeRetentionTest.cpp

 #include "gtest/gtest.h"
 #include <random>
 
 using namespace llvm;
 
 extern cl::opt<bool> ShouldPreserveAllAttributes;
 extern cl::opt<bool> EnableKnowledgeRetention;
 
+static IntrinsicInst *buildAssumeFromInst(Instruction *I) {
+  return cast_or_null<IntrinsicInst>(BuildAssumeFromInst(I));
+}
+
 static void RunTest(
     StringRef Head, StringRef Tail,
     std::vector<std::pair<StringRef, llvm::function_ref<void(Instruction *)>>>
         &Tests) {
   for (auto &Elem : Tests) {
     std::string IR;
     IR.append(Head.begin(), Head.end());
     IR.append(Elem.first.begin(), Elem.first.end());
     IR.append(Tail.begin(), Tail.end());
     LLVMContext C;
     SMDiagnostic Err;
     std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
     if (!Mod)
       Err.print("AssumeQueryAPI", errs());
     Elem.second(&*(Mod->getFunction("test")->begin()->begin()));
   }
 }
 
-bool hasMatchesExactlyAttributes(CallInst *Assume, Value *WasOn,
+bool hasMatchesExactlyAttributes(IntrinsicInst *Assume, Value *WasOn,
                                     StringRef AttrToMatch) {
   Regex Reg(AttrToMatch);
   SmallVector<StringRef, 1> Matches;
   for (StringRef Attr : {
 #define GET_ATTR_NAMES
 #define ATTRIBUTE_ALL(ENUM_NAME, DISPLAY_NAME) StringRef(#DISPLAY_NAME),
 #include "llvm/IR/Attributes.inc"
        }) {
     bool ShouldHaveAttr = Reg.match(Attr, &Matches) && Matches[0] == Attr;
     if (ShouldHaveAttr != hasAttributeInAssume(*Assume, WasOn, Attr))
       return false;
   }
   return true;
 }
 
-bool hasTheRightValue(CallInst *Assume, Value *WasOn,
+bool hasTheRightValue(IntrinsicInst *Assume, Value *WasOn,
                             Attribute::AttrKind Kind, unsigned Value, bool Both,
                             AssumeQuery AQ = AssumeQuery::Highest) {
   if (!Both) {
     uint64_t ArgVal = 0;
     if (!hasAttributeInAssume(*Assume, WasOn, Kind, &ArgVal, AQ))
       return false;
     if (ArgVal != Value)
       return false;
TEST(AssumeQueryAPI, hasAttributeInAssume) {
   StringRef Tail = "ret void\n"
                    "}";
   std::vector<std::pair<StringRef, llvm::function_ref<void(Instruction *)>>>
       Tests;
   Tests.push_back(std::make_pair(
       "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
       "8 noalias %P1)\n",
       [](Instruction *I) {
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(0),
                                        "(nonnull|align|dereferenceable)"));
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(1),
                                        "(align)"));
         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
                                Attribute::AttrKind::Dereferenceable, 16, true));
         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
                                Attribute::AttrKind::Alignment, 4, true));
         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
                                Attribute::AttrKind::Alignment, 4, true));
       }));
   Tests.push_back(std::make_pair(
       "call void @func1(i32* nonnull align 32 dereferenceable(48) %P, i32* "
       "nonnull "
       "align 8 dereferenceable(28) %P, i32* nonnull align 64 "
       "dereferenceable(4) "
       "%P, i32* nonnull align 16 dereferenceable(12) %P)\n",
       [](Instruction *I) {
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(0),
                                        "(nonnull|align|dereferenceable)"));
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(1),
                                        "(nonnull|align|dereferenceable)"));
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(2),
                                        "(nonnull|align|dereferenceable)"));
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(3),
Tests.push_back(std::make_pair(
                                AssumeQuery::Lowest));
         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(1),
                                Attribute::AttrKind::Alignment, 8, false,
                                AssumeQuery::Lowest));
       }));
   Tests.push_back(std::make_pair(
       "call void @func_many(i32* align 8 %P1) cold\n", [](Instruction *I) {
         ShouldPreserveAllAttributes.setValue(true);
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
         ASSERT_TRUE(hasMatchesExactlyAttributes(
             Assume, nullptr,
             "(align|no-jump-tables|less-precise-fpmad|"
             "nounwind|norecurse|willreturn|cold)"));
         ShouldPreserveAllAttributes.setValue(false);
       }));
   Tests.push_back(
       std::make_pair("call void @llvm.assume(i1 true)\n", [](Instruction *I) {
-        CallInst *Assume = cast<CallInst>(I);
+        IntrinsicInst *Assume = cast<IntrinsicInst>(I);
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, nullptr, ""));
       }));
   Tests.push_back(std::make_pair(
       "call void @func1(i32* readnone align 32 "
       "dereferenceable(48) noalias %P, i32* "
       "align 8 dereferenceable(28) %P1, i32* align 64 "
       "dereferenceable(4) "
       "%P2, i32* nonnull align 16 dereferenceable(12) %P3)\n",
       [](Instruction *I) {
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
         ASSERT_TRUE(hasMatchesExactlyAttributes(
             Assume, I->getOperand(0),
             "(align|dereferenceable)"));
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(1),
                                        "(align|dereferenceable)"));
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(2),
                                        "(align|dereferenceable)"));
TEST(AssumeQueryAPI, hasAttributeInAssume) {
 
   Tests.push_back(std::make_pair(
       "call void @func1(i32* readnone align 32 "
       "dereferenceable(48) noalias %P, i32* "
       "align 8 dereferenceable(28) %P1, i32* align 64 "
       "dereferenceable(4) "
       "%P2, i32* nonnull align 16 dereferenceable(12) %P3)\n",
       [](Instruction *I) {
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
         I->getOperand(1)->dropDroppableUses();
         I->getOperand(2)->dropDroppableUses();
         I->getOperand(3)->dropDroppableUses();
         ASSERT_TRUE(hasMatchesExactlyAttributes(
             Assume, I->getOperand(0),
             "(align|dereferenceable)"));
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(1),
                                        ""));
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(2),
                                        ""));
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(3),
                                        ""));
         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
                                Attribute::AttrKind::Alignment, 32, true));
         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
                                Attribute::AttrKind::Dereferenceable, 48, true));
       }));
   Tests.push_back(std::make_pair(
       "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
       "8 noalias %P1)\n",
       [](Instruction *I) {
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
         Value *New = I->getFunction()->getArg(3);
         Value *Old = I->getOperand(0);
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, New, ""));
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, Old,
                                        "(nonnull|align|dereferenceable)"));
         Old->replaceAllUsesWith(New);
         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, New,
#include "llvm/IR/Attributes.inc"
     bool ShouldHaveAttr = Reg.match(Attr, &Matches) && Matches[0] == Attr;
 
     if (ShouldHaveAttr != (Map.find(RetainedKnowledgeKey{WasOn, Attribute::getAttrKindFromName(Attr)}) != Map.end()))
       return false;
   }
   return true;
 }
 
-static bool MapHasRightValue(RetainedKnowledgeMap &Map,
-                                   RetainedKnowledgeKey Key, MinMax MM) {
+static bool MapHasRightValue(RetainedKnowledgeMap &Map, IntrinsicInst *II,
   auto LookupIt = Map.find(Key);
-  return (LookupIt != Map.end()) && (LookupIt->second.Min == MM.Min) &&
-         (LookupIt->second.Max == MM.Max);
+  return (LookupIt != Map.end()) && (LookupIt->second[II].Min == MM.Min) &&
+         (LookupIt->second[II].Max == MM.Max);
 }
 
 TEST(AssumeQueryAPI, fillMapFromAssume) {
   EnableKnowledgeRetention.setValue(true);
   StringRef Head =
       "declare void @llvm.assume(i1)\n"
       "declare void @func(i32*, i32*)\n"
       "declare void @func1(i32*, i32*, i32*, i32*)\n"
       "declare void @func_many(i32*) \"no-jump-tables\" nounwind "
       "\"less-precise-fpmad\" willreturn norecurse\n"
       "define void @test(i32* %P, i32* %P1, i32* %P2, i32* %P3) {\n";
   StringRef Tail = "ret void\n"
                    "}";
   std::vector<std::pair<StringRef, llvm::function_ref<void(Instruction *)>>>
       Tests;
   Tests.push_back(std::make_pair(
       "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
       "8 noalias %P1)\n",
       [](Instruction *I) {
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
 
         RetainedKnowledgeMap Map;
         fillMapFromAssume(*Assume, Map);
         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(0),
                                        "(nonnull|align|dereferenceable)"));
         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(1),
                                        "(align)"));
         ASSERT_TRUE(MapHasRightValue(
-            Map, {I->getOperand(0), Attribute::Dereferenceable}, {16, 16}));
-        ASSERT_TRUE(MapHasRightValue(Map, {I->getOperand(0), Attribute::Alignment},
+            Map, Assume, {I->getOperand(0), Attribute::Dereferenceable}, {16, 16}));
+        ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(0), Attribute::Alignment},
                                {4, 4}));
-        ASSERT_TRUE(MapHasRightValue(Map, {I->getOperand(0), Attribute::Alignment},
+        ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(0), Attribute::Alignment},
                                {4, 4}));
       }));
   Tests.push_back(std::make_pair(
       "call void @func1(i32* nonnull align 32 dereferenceable(48) %P, i32* "
       "nonnull "
       "align 8 dereferenceable(28) %P, i32* nonnull align 64 "
       "dereferenceable(4) "
       "%P, i32* nonnull align 16 dereferenceable(12) %P)\n",
       [](Instruction *I) {
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
 
         RetainedKnowledgeMap Map;
         fillMapFromAssume(*Assume, Map);
 
         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(0),
                                        "(nonnull|align|dereferenceable)"));
         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(1),
                                        "(nonnull|align|dereferenceable)"));
         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(2),
                                        "(nonnull|align|dereferenceable)"));
         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(3),
                                        "(nonnull|align|dereferenceable)"));
         ASSERT_TRUE(MapHasRightValue(
-            Map, {I->getOperand(0), Attribute::Dereferenceable}, {4, 48}));
-        ASSERT_TRUE(MapHasRightValue(Map, {I->getOperand(0), Attribute::Alignment},
+            Map, Assume, {I->getOperand(0), Attribute::Dereferenceable}, {4, 48}));
+        ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(0), Attribute::Alignment},
                                {8, 64}));
       }));
   Tests.push_back(std::make_pair(
       "call void @func_many(i32* align 8 %P1) cold\n", [](Instruction *I) {
         ShouldPreserveAllAttributes.setValue(true);
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
 
         RetainedKnowledgeMap Map;
         fillMapFromAssume(*Assume, Map);
 
         ASSERT_TRUE(FindExactlyAttributes(
             Map, nullptr, "(nounwind|norecurse|willreturn|cold)"));
         ShouldPreserveAllAttributes.setValue(false);
       }));
   Tests.push_back(
       std::make_pair("call void @llvm.assume(i1 true)\n", [](Instruction *I) {
         RetainedKnowledgeMap Map;
-        fillMapFromAssume(*cast<CallInst>(I), Map);
+        fillMapFromAssume(*cast<IntrinsicInst>(I), Map);
 
         ASSERT_TRUE(FindExactlyAttributes(Map, nullptr, ""));
         ASSERT_TRUE(Map.empty());
       }));
   Tests.push_back(std::make_pair(
       "call void @func1(i32* readnone align 32 "
       "dereferenceable(48) noalias %P, i32* "
       "align 8 dereferenceable(28) %P1, i32* align 64 "
       "dereferenceable(4) "
       "%P2, i32* nonnull align 16 dereferenceable(12) %P3)\n",
       [](Instruction *I) {
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
 
         RetainedKnowledgeMap Map;
         fillMapFromAssume(*Assume, Map);
 
         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(0),
                                     "(align|dereferenceable)"));
         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(1),
                                     "(align|dereferenceable)"));
         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(2),
                                        "(align|dereferenceable)"));
         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(3),
                                        "(nonnull|align|dereferenceable)"));
-        ASSERT_TRUE(MapHasRightValue(Map, {I->getOperand(0), Attribute::Alignment},
+        ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(0), Attribute::Alignment},
                                {32, 32}));
         ASSERT_TRUE(MapHasRightValue(
-            Map, {I->getOperand(0), Attribute::Dereferenceable}, {48, 48}));
+            Map, Assume, {I->getOperand(0), Attribute::Dereferenceable}, {48, 48}));
         ASSERT_TRUE(MapHasRightValue(
-            Map, {I->getOperand(1), Attribute::Dereferenceable}, {28, 28}));
-        ASSERT_TRUE(MapHasRightValue(Map, {I->getOperand(1), Attribute::Alignment},
+            Map, Assume, {I->getOperand(1), Attribute::Dereferenceable}, {28, 28}));
+        ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(1), Attribute::Alignment},
                                {8, 8}));
-        ASSERT_TRUE(MapHasRightValue(Map, {I->getOperand(2), Attribute::Alignment},
+        ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(2), Attribute::Alignment},
                                {64, 64}));
         ASSERT_TRUE(MapHasRightValue(
-            Map, {I->getOperand(2), Attribute::Dereferenceable}, {4, 4}));
-        ASSERT_TRUE(MapHasRightValue(Map, {I->getOperand(3), Attribute::Alignment},
+            Map, Assume, {I->getOperand(2), Attribute::Dereferenceable}, {4, 4}));
+        ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(3), Attribute::Alignment},
                                {16, 16}));
         ASSERT_TRUE(MapHasRightValue(
-            Map, {I->getOperand(3), Attribute::Dereferenceable}, {12, 12}));
+            Map, Assume, {I->getOperand(3), Attribute::Dereferenceable}, {12, 12}));
       }));
 
   /// Keep this test last as it modifies the function.
   Tests.push_back(std::make_pair(
       "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
       "8 noalias %P1)\n",
       [](Instruction *I) {
-        CallInst *Assume = BuildAssumeFromInst(I);
+        IntrinsicInst *Assume = buildAssumeFromInst(I);
         Assume->insertBefore(I);
 
         RetainedKnowledgeMap Map;
         fillMapFromAssume(*Assume, Map);
 
         Value *New = I->getFunction()->getArg(3);
         Value *Old = I->getOperand(0);
         ASSERT_TRUE(FindExactlyAttributes(Map, New, ""));
if (count > 0) {
       HasArg[i] = true;
     }
     if (count > 1)
       Args.push_back(ConstantInt::get(Type::getInt32Ty(C), value));
 
     OpBundle.push_back(OperandBundleDef{ss.str().c_str(), std::move(Args)});
   }
 
-  Instruction *Assume =
-      CallInst::Create(FnAssume, ArrayRef<Value *>({ConstantInt::getTrue(C)}),
-                       std::move(OpBundle));
+  auto *Assume = cast<IntrinsicInst>(IntrinsicInst::Create(
+      FnAssume, ArrayRef<Value *>({ConstantInt::getTrue(C)}),
+      std::move(OpBundle)));
   Assume->insertBefore(&F->begin()->front());
   RetainedKnowledgeMap Map;
-  fillMapFromAssume(*cast<CallInst>(Assume), Map);
+  fillMapFromAssume(*Assume, Map);
   for (int i = 0; i < (Size * 2); i++) {
     if (!HasArg[i])
       continue;
     RetainedKnowledge K =
         getKnowledgeFromUseInAssume(&*ShuffledArgs[i]->use_begin());
     auto LookupIt = Map.find(RetainedKnowledgeKey{K.WasOn, K.AttrKind});
     ASSERT_TRUE(LookupIt != Map.end());
-    MinMax MM = LookupIt->second;
+    MinMax MM = LookupIt->second[Assume];
     ASSERT_TRUE(MM.Min == MM.Max);
     ASSERT_TRUE(MM.Min == K.ArgValue);
   }
 }
 
 TEST(AssumeQueryAPI, getKnowledgeFromUseInAssume) {
   // // For Fuzzing
   // std::random_device dev;