[AssumeBundles] adapt Assumption cache to assume bundles

[Tyker]

Summary

change assumption cache to store an assume along with an index to the operand bundle containing the knowledge.

[jdoerfert]

How does this work with multiple occurrences of a value in an assume?

---

change assumption cache to store an assume along with an index to the operand bundle containg the knoledge.

Two typos in the commit message

llvm/lib/Analysis/AssumptionCache.cpp

86

Clang format the above please.

llvm/lib/Transforms/Utils/InlineFunction.cpp

1846

Can you make this a variable, hard to read and interpret this way.

[Tyker]

make cache invalidation invalidate no more than needed.
addressed comments.

[jdoerfert]

In D77402#1960415, @jdoerfert wrote:

How does this work with multiple occurrences of a value in an assume?

---

change assumption cache to store an assume along with an index to the operand bundle containg the knoledge.

Two typos in the commit message

These are still open :)

[Tyker]

In D77402#1960415, @jdoerfert wrote:

How does this work with multiple occurrences of a value in an assume?

the cache will contain multiple element with the same assume but a different index.

[jdoerfert]

One more conceptual question and small nits.

llvm/lib/Analysis/AssumptionCache.cpp

60

I usually prefer a struct instead of a pair but this is local only so I guess it's ok. Feel free to use the ResultElem or something similar to improve readability here and at the use site below. (AV.first will then become something meaningful on first glance.)

171

Could you explain the need for the above logic please? Is is strictly related to the index extension or a general improvement?

(Nit: Use a continue to reduce indention and simplify the code)

[Tyker]

replaced std::pair<Value *, unsigned> with ResultElem

[Tyker]

llvm/lib/Analysis/AssumptionCache.cpp

171

the AffectedValues map can associate multiple assumes with a single value this is true before and after the patch although i am not sure it is reachable before the patch.

here is an example of when this change matters:

call void @llvm.assume(i1 true) ["nonnull"(i32* %P)] ; called A
call void @llvm.assume(i1 true) ["align"(i32* %P, i32 4)] ; called B

with both A and B in cache both A and B will be in the AffectedValues for %P.

without this change if unregisterAssumption is called on A, all assumption of P will be invalidate even the one in B, however B will still be in the AssumeHandles (the list of all assumes in the function).
with this change if unregisterAssumption is called on A, A will be removed (set to null) for the assumptions on V and B will not be affected.

this is not strictly related only to AssumeBundles. but it makes a difference in a situtation which is rare if possible without assume bundles.

Note: that having a nullptr in assumption is already possible since Value* are stored in WeakTrackingVH so it can be set to null if it is removed.

[jdoerfert]

LGTM, see below for some additional style changes and documentation you should apply before committing.

llvm/lib/Analysis/AssumptionCache.cpp

171

OK. I expected this to be a improvement not strictly related to assume bundles but I wasn't sure. You explanation makes sense to me.

---

Before you commit maybe replace AVI != ..end() with the opposite condition and a continue.
Add a message to the assert please.

llvm/include/llvm/Analysis/AssumptionCache.h

 /// A cache of \@llvm.assume calls within a function.
 ///
 /// This cache provides fast lookup of assumptions within a function by caching
 /// them and amortizing the cost of scanning for them across all queries. Passes
 /// that create new assumptions are required to call registerAssumption() to
 /// register any new \@llvm.assume calls that they create. Deletions of
 /// \@llvm.assume calls do not require special handling.
 class AssumptionCache {
+public:
+  /// Value of ResultElem::Index indicating that the argument to the call of the
+  /// llvm.assume.
+  enum : unsigned { ExprResultIdx = std::numeric_limits<unsigned>::max() };
+
+  struct ResultElem {
+    WeakTrackingVH Assume;
+
+    /// contains either ExprResultIdx or the index of the operand bundle
+    /// containing the knowledge.
+    unsigned Index;
+    operator Value *() const { return Assume; }
+  };
+
+private:
   /// The function for which this cache is handling assumptions.
   ///
   /// We track this to lazily populate our assumptions.
   Function &F;
 
   /// Vector of weak value handles to calls of the \@llvm.assume
   /// intrinsic.
-  SmallVector<WeakTrackingVH, 4> AssumeHandles;
+  SmallVector<ResultElem, 4> AssumeHandles;
 
   class AffectedValueCallbackVH final : public CallbackVH {
     AssumptionCache *AC;
 
     void deleted() override;
     void allUsesReplacedWith(Value *) override;
 
   public:
     using DMI = DenseMapInfo<Value *>;
 
     AffectedValueCallbackVH(Value *V, AssumptionCache *AC = nullptr)
         : CallbackVH(V), AC(AC) {}
   };
 
   friend AffectedValueCallbackVH;
 
   /// A map of values about which an assumption might be providing
   /// information to the relevant set of assumptions.
   using AffectedValuesMap =
-      DenseMap<AffectedValueCallbackVH, SmallVector<WeakTrackingVH, 1>,
+      DenseMap<AffectedValueCallbackVH, SmallVector<ResultElem, 1>,
                AffectedValueCallbackVH::DMI>;
   AffectedValuesMap AffectedValues;
 
   /// Get the vector of assumptions which affect a value from the cache.
-  SmallVector<WeakTrackingVH, 1> &getOrInsertAffectedValues(Value *V);
+  SmallVector<ResultElem, 1> &getOrInsertAffectedValues(Value *V);
 
   /// Move affected values in the cache for OV to be affected values for NV.
   void transferAffectedValuesInCache(Value *OV, Value *NV);
 
   /// Flag tracking whether we have scanned the function yet.
   ///
   /// We want to be as lazy about this as possible, and so we scan the function
   /// at the last moment.
public:
   /// Access the list of assumption handles currently tracked for this
   /// function.
   ///
   /// Note that these produce weak handles that may be null. The caller must
   /// handle that case.
   /// FIXME: We should replace this with pointee_iterator<filter_iterator<...>>
   /// when we can write that to filter out the null values. Then caller code
   /// will become simpler.
-  MutableArrayRef<WeakTrackingVH> assumptions() {
+  MutableArrayRef<ResultElem> assumptions() {
     if (!Scanned)
       scanFunction();
     return AssumeHandles;
   }
 
   /// Access the list of assumptions which affect this value.
-  MutableArrayRef<WeakTrackingVH> assumptionsFor(const Value *V) {
+  MutableArrayRef<ResultElem> assumptionsFor(const Value *V) {
     if (!Scanned)
       scanFunction();
 
     auto AVI = AffectedValues.find_as(const_cast<Value *>(V));
     if (AVI == AffectedValues.end())
-      return MutableArrayRef<WeakTrackingVH>();
+      return MutableArrayRef<ResultElem>();
 
     return AVI->second;
   }
 };
 
 /// A function analysis which provides an \c AssumptionCache.
 ///
 /// This analysis is intended for use with the new pass manager and will vend
public:
   bool doFinalization(Module &) override {
     verifyAnalysis();
     return false;
   }
 
   static char ID; // Pass identification, replacement for typeid
 };
 
+template<> struct simplify_type<AssumptionCache::ResultElem> {
+  using SimpleType = Value *;
+
+  static SimpleType getSimplifiedValue(AssumptionCache::ResultElem &Val) {
+    return Val;
+  }
+};
+template<> struct simplify_type<const AssumptionCache::ResultElem> {
+  using SimpleType = /*const*/ Value *;
+
+  static SimpleType getSimplifiedValue(const AssumptionCache::ResultElem &Val) {
+    return Val;
+  }
+};
+
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_ASSUMPTIONCACHE_H

llvm/include/llvm/Transforms/Utils/AssumeBundleBuilder.h

 #define LLVM_TRANSFORMS_UTILS_ASSUMEBUNDLEBUILDER_H
 
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 class IntrinsicInst;
+class AssumptionCache;
 
 /// 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.
 IntrinsicInst *buildAssumeFromInst(Instruction *I);
 
 /// Calls BuildAssumeFromInst and if the resulting llvm.assume is valid insert
 /// if before I. This is usually what need to be done to salvage the knowledge
 /// contained in the instruction I.
-void salvageKnowledge(Instruction *I);
+void salvageKnowledge(Instruction *I, AssumptionCache *AC = nullptr);
 
 /// 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/AssumptionCache.cpp

 //===- AssumptionCache.cpp - Cache finding @llvm.assume calls -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains a pass that keeps track of @llvm.assume intrinsics in
 // the functions of a module.
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/Analysis/AssumeBundleQueries.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"

 using namespace llvm;
 using namespace llvm::PatternMatch;
 
 static cl::opt<bool>
     VerifyAssumptionCache("verify-assumption-cache", cl::Hidden,
                           cl::desc("Enable verification of assumption cache"),
                           cl::init(false));
 
-SmallVector<WeakTrackingVH, 1> &
+SmallVector<AssumptionCache::ResultElem, 1> &
 AssumptionCache::getOrInsertAffectedValues(Value *V) {
   // Try using find_as first to avoid creating extra value handles just for the
   // purpose of doing the lookup.
   auto AVI = AffectedValues.find_as(V);
   if (AVI != AffectedValues.end())
     return AVI->second;
 
   auto AVIP = AffectedValues.insert(
-      {AffectedValueCallbackVH(V, this), SmallVector<WeakTrackingVH, 1>()});
+      {AffectedValueCallbackVH(V, this), SmallVector<ResultElem, 1>()});
   return AVIP.first->second;
 }
 
-static void findAffectedValues(CallInst *CI,
-                               SmallVectorImpl<Value *> &Affected) {
+static void
+findAffectedValues(CallInst *CI,
+                   SmallVectorImpl<AssumptionCache::ResultElem> &Affected) {
   // Note: This code must be kept in-sync with the code in
   // computeKnownBitsFromAssume in ValueTracking.
 
-  auto AddAffected = [&Affected](Value *V) {
+  auto AddAffected = [&Affected](Value *V, unsigned Idx =
+                                               AssumptionCache::ExprResultIdx) {
     if (isa<Argument>(V)) {
-      Affected.push_back(V);
+      Affected.push_back({V, Idx});
     } else if (auto *I = dyn_cast<Instruction>(V)) {
-      Affected.push_back(I);
+      Affected.push_back({I, Idx});
 
       // Peek through unary operators to find the source of the condition.
       Value *Op;
       if (match(I, m_BitCast(m_Value(Op))) ||
-          match(I, m_PtrToInt(m_Value(Op))) ||
-          match(I, m_Not(m_Value(Op)))) {
+          match(I, m_PtrToInt(m_Value(Op))) || match(I, m_Not(m_Value(Op)))) {
         if (isa<Instruction>(Op) || isa<Argument>(Op))
-          Affected.push_back(Op);
+          Affected.push_back({Op, Idx});
       }
     }
   };
 
+  for (unsigned Idx = 0; Idx != CI->getNumOperandBundles(); Idx++) {
+    if (CI->getOperandBundleAt(Idx).Inputs.size() > ABA_WasOn &&
+        CI->getOperandBundleAt(Idx).getTagName() != "ignore")
+      AddAffected(CI->getOperandBundleAt(Idx).Inputs[ABA_WasOn], Idx);
+  }
+
   Value *Cond = CI->getArgOperand(0), *A, *B;
   AddAffected(Cond);
 
   CmpInst::Predicate Pred;
   if (match(Cond, m_ICmp(Pred, m_Value(A), m_Value(B)))) {
     AddAffected(A);
     AddAffected(B);
 
if (Pred == ICmpInst::ICMP_EQ) {
 
       AddAffectedFromEq(A);
       AddAffectedFromEq(B);
     }
   }
 }
 
 void AssumptionCache::updateAffectedValues(CallInst *CI) {
-  SmallVector<Value *, 16> Affected;
+  SmallVector<AssumptionCache::ResultElem, 16> Affected;
   findAffectedValues(CI, Affected);
 
   for (auto &AV : Affected) {
-    auto &AVV = getOrInsertAffectedValues(AV);
-    if (std::find(AVV.begin(), AVV.end(), CI) == AVV.end())
-      AVV.push_back(CI);
+    auto &AVV = getOrInsertAffectedValues(AV.Assume);
+    if (std::find_if(AVV.begin(), AVV.end(), [&](ResultElem &Elem) {
+          return Elem.Assume == CI && Elem.Index == AV.Index;
+        }) == AVV.end())
+      AVV.push_back({CI, AV.Index});
   }
 }
 
 void AssumptionCache::unregisterAssumption(CallInst *CI) {
-  SmallVector<Value *, 16> Affected;
+  SmallVector<AssumptionCache::ResultElem, 16> Affected;
   findAffectedValues(CI, Affected);
 
   for (auto &AV : Affected) {
-    auto AVI = AffectedValues.find_as(AV);
-    if (AVI != AffectedValues.end())
+    auto AVI = AffectedValues.find_as(AV.Assume);
+    if (AVI == AffectedValues.end())
+      continue;
+    bool Found = false;
+    bool HasNonnull = false;
+    for (ResultElem &Elem : AVI->second) {
+      if (Elem.Assume == CI) {
+        Found = true;
+        Elem.Assume = nullptr;
+      }
+      HasNonnull |= !!Elem.Assume;
+      if (HasNonnull && Found)
+        break;
+    }
+    assert(Found && "already unregistered or incorrect cache state");
+    if (!HasNonnull)
       AffectedValues.erase(AVI);
   }
 
   AssumeHandles.erase(
-      remove_if(AssumeHandles, [CI](WeakTrackingVH &VH) { return CI == VH; }),
+      remove_if(AssumeHandles, [CI](ResultElem &RE) { return CI == RE; }),
       AssumeHandles.end());
 }
 
 void AssumptionCache::AffectedValueCallbackVH::deleted() {
   auto AVI = AC->AffectedValues.find(getValPtr());
   if (AVI != AC->AffectedValues.end())
     AC->AffectedValues.erase(AVI);
   // 'this' now dangles!
 }
 
 void AssumptionCache::transferAffectedValuesInCache(Value *OV, Value *NV) {
   auto &NAVV = getOrInsertAffectedValues(NV);
   auto AVI = AffectedValues.find(OV);
   if (AVI == AffectedValues.end())
     return;
 
   for (auto &A : AVI->second)
     if (std::find(NAVV.begin(), NAVV.end(), A) == NAVV.end())
       NAVV.push_back(A);
void AssumptionCache::scanFunction() {
   assert(!Scanned && "Tried to scan the function twice!");
   assert(AssumeHandles.empty() && "Already have assumes when scanning!");
 
   // Go through all instructions in all blocks, add all calls to @llvm.assume
   // to this cache.
   for (BasicBlock &B : F)
     for (Instruction &II : B)
       if (match(&II, m_Intrinsic<Intrinsic::assume>()))
-        AssumeHandles.push_back(&II);
+        AssumeHandles.push_back({&II, ExprResultIdx});
 
   // Mark the scan as complete.
   Scanned = true;
 
   // Update affected values.
   for (auto &A : AssumeHandles)
     updateAffectedValues(cast<CallInst>(A));
 }
 
 void AssumptionCache::registerAssumption(CallInst *CI) {
   assert(match(CI, m_Intrinsic<Intrinsic::assume>()) &&
          "Registered call does not call @llvm.assume");
 
   // If we haven't scanned the function yet, just drop this assumption. It will
   // be found when we scan later.
   if (!Scanned)
     return;
 
-  AssumeHandles.push_back(CI);
+  AssumeHandles.push_back({CI, ExprResultIdx});
 
 #ifndef NDEBUG
   assert(CI->getParent() &&
          "Cannot register @llvm.assume call not in a basic block");
   assert(&F == CI->getParent()->getParent() &&
          "Cannot register @llvm.assume call not in this function");
 
   // We expect the number of assumptions to be small, so in an asserts build

llvm/lib/Transforms/Scalar/EarlyCSE.cpp

for (Instruction &Inst : make_early_inc_range(BB->getInstList())) {
     // Dead instructions should just be removed.
     if (isInstructionTriviallyDead(&Inst, &TLI)) {
       LLVM_DEBUG(dbgs() << "EarlyCSE DCE: " << Inst << '\n');
       if (!DebugCounter::shouldExecute(CSECounter)) {
         LLVM_DEBUG(dbgs() << "Skipping due to debug counter\n");
         continue;
       }
 
-      salvageKnowledge(&Inst);
+      salvageKnowledge(&Inst, &AC);
       salvageDebugInfoOrMarkUndef(Inst);
       removeMSSA(Inst);
       Inst.eraseFromParent();
       Changed = true;
       ++NumSimplify;
       continue;
     }
 
if (isGuard(&Inst)) {
         if (SimpleValue::canHandle(CondI)) {
           // Do we already know the actual value of this condition?
           if (auto *KnownCond = AvailableValues.lookup(CondI)) {
             // Is the condition known to be true?
             if (isa<ConstantInt>(KnownCond) &&
                 cast<ConstantInt>(KnownCond)->isOne()) {
               LLVM_DEBUG(dbgs()
                          << "EarlyCSE removing guard: " << Inst << '\n');
-              salvageKnowledge(&Inst);
+              salvageKnowledge(&Inst, &AC);
               removeMSSA(Inst);
               Inst.eraseFromParent();
               Changed = true;
               continue;
             } else
               // Use the known value if it wasn't true.
               cast<CallInst>(Inst).setArgOperand(0, KnownCond);
           }
if (Value *V = SimplifyInstruction(&Inst, SQ)) {
         LLVM_DEBUG(dbgs() << "Skipping due to debug counter\n");
       } else {
         bool Killed = false;
         if (!Inst.use_empty()) {
           Inst.replaceAllUsesWith(V);
           Changed = true;
         }
         if (isInstructionTriviallyDead(&Inst, &TLI)) {
-          salvageKnowledge(&Inst);
+          salvageKnowledge(&Inst, &AC);
           removeMSSA(Inst);
           Inst.eraseFromParent();
           Changed = true;
           Killed = true;
         }
         if (Changed)
           ++NumSimplify;
         if (Killed)
if (SimpleValue::canHandle(&Inst)) {
                           << '\n');
         if (!DebugCounter::shouldExecute(CSECounter)) {
           LLVM_DEBUG(dbgs() << "Skipping due to debug counter\n");
           continue;
         }
         if (auto *I = dyn_cast<Instruction>(V))
           I->andIRFlags(&Inst);
         Inst.replaceAllUsesWith(V);
-        salvageKnowledge(&Inst);
+        salvageKnowledge(&Inst, &AC);
         removeMSSA(Inst);
         Inst.eraseFromParent();
         Changed = true;
         ++NumCSE;
         continue;
       }
 
       // Otherwise, just remember that this value is available.
if (MemInst.isValid() && MemInst.isLoad()) {
           LLVM_DEBUG(dbgs() << "EarlyCSE CSE LOAD: " << Inst
                             << "  to: " << *InVal.DefInst << '\n');
           if (!DebugCounter::shouldExecute(CSECounter)) {
             LLVM_DEBUG(dbgs() << "Skipping due to debug counter\n");
             continue;
           }
           if (!Inst.use_empty())
             Inst.replaceAllUsesWith(Op);
-          salvageKnowledge(&Inst);
+          salvageKnowledge(&Inst, &AC);
           removeMSSA(Inst);
           Inst.eraseFromParent();
           Changed = true;
           ++NumCSELoad;
           continue;
         }
       }
 
if (CallValue::canHandle(&Inst)) {
         LLVM_DEBUG(dbgs() << "EarlyCSE CSE CALL: " << Inst
                           << "  to: " << *InVal.first << '\n');
         if (!DebugCounter::shouldExecute(CSECounter)) {
           LLVM_DEBUG(dbgs() << "Skipping due to debug counter\n");
           continue;
         }
         if (!Inst.use_empty())
           Inst.replaceAllUsesWith(InVal.first);
-        salvageKnowledge(&Inst);
+        salvageKnowledge(&Inst, &AC);
         removeMSSA(Inst);
         Inst.eraseFromParent();
         Changed = true;
         ++NumCSECall;
         continue;
       }
 
       // Otherwise, remember that we have this instruction.
if (MemInst.isValid() && MemInst.isStore()) {
                     MemInst.getPointerOperand() ||
                 MSSA) &&
                "can't have an intervening store if not using MemorySSA!");
         LLVM_DEBUG(dbgs() << "EarlyCSE DSE (writeback): " << Inst << '\n');
         if (!DebugCounter::shouldExecute(CSECounter)) {
           LLVM_DEBUG(dbgs() << "Skipping due to debug counter\n");
           continue;
         }
-        salvageKnowledge(&Inst);
+        salvageKnowledge(&Inst, &AC);
         removeMSSA(Inst);
         Inst.eraseFromParent();
         Changed = true;
         ++NumDSE;
         // We can avoid incrementing the generation count since we were able
         // to eliminate this store.
         continue;
       }
if (Inst.mayWriteToMemory()) {
                  !LastStoreMemInst.isVolatile() &&
                  "Violated invariant");
           if (LastStoreMemInst.isMatchingMemLoc(MemInst)) {
             LLVM_DEBUG(dbgs() << "EarlyCSE DEAD STORE: " << *LastStore
                               << "  due to: " << Inst << '\n');
             if (!DebugCounter::shouldExecute(CSECounter)) {
               LLVM_DEBUG(dbgs() << "Skipping due to debug counter\n");
             } else {
-              salvageKnowledge(&Inst);
+              salvageKnowledge(&Inst, &AC);
               removeMSSA(*LastStore);
               LastStore->eraseFromParent();
               Changed = true;
               ++NumDSE;
               LastStore = nullptr;
             }
           }
           // fallthrough - we can exploit information about this store

llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp

 //===- AssumeBundleBuilder.cpp - tools to preserve informations -*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/AssumeBundleBuilder.h"
 #include "llvm/Analysis/AssumeBundleQueries.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace llvm;

 IntrinsicInst *llvm::buildAssumeFromInst(Instruction *I) {
   if (!EnableKnowledgeRetention)
     return nullptr;
   AssumeBuilderState Builder(I->getModule());
   Builder.addInstruction(I);
   return Builder.build();
 }
 
-void llvm::salvageKnowledge(Instruction *I) {
-  if (Instruction *Intr = buildAssumeFromInst(I))
+void llvm::salvageKnowledge(Instruction *I, AssumptionCache *AC) {
+  if (IntrinsicInst *Intr = buildAssumeFromInst(I)) {
     Intr->insertBefore(I);
+    if (AC)
+      AC->registerAssumption(Intr);
+  }
 }
 
 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/Utils/InlineFunction.cpp

Function::iterator FirstNewBlock;
     }
 
     // TODO: Remove this when users have been updated to the assume bundles.
     // Add alignment assumptions if necessary. We do this before the inlined
     // instructions are actually cloned into the caller so that we can easily
     // check what will be known at the start of the inlined code.
     AddAlignmentAssumptions(CS, IFI);
 
+    AssumptionCache *AC =
+        IFI.GetAssumptionCache ? &(*IFI.GetAssumptionCache)(*Caller) : nullptr;
+
     /// Preserve all attributes on of the call and its parameters.
-    if (Instruction *Assume = buildAssumeFromInst(CS.getInstruction()))
-      Assume->insertBefore(CS.getInstruction());
+    salvageKnowledge(CS.getInstruction(), AC);
 
     // We want the inliner to prune the code as it copies.  We would LOVE to
     // have no dead or constant instructions leftover after inlining occurs
     // (which can happen, e.g., because an argument was constant), but we'll be
     // happy with whatever the cloner can do.
     CloneAndPruneFunctionInto(Caller, CalledFunc, VMap,
                               /*ModuleLevelChanges=*/false, Returns, ".i",
                               &InlinedFunctionInfo, TheCall);
     // Remember the first block that is newly cloned over.
     FirstNewBlock = LastBlock; ++FirstNewBlock;

llvm/unittests/Analysis/AssumeBundleQueriesTest.cpp

 //===- AssumeBundleQueriesTest.cpp ------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/AssumeBundleQueries.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/CommandLine.h"
TEST(AssumeQueryAPI, getKnowledgeFromUseInAssume) {
   // }
   RunRandTest(23456, 4, 0, 2, 100);
   RunRandTest(560987, 25, -3, 2, 100);
 
   // Large bundles can lead to special cases. this is why this test is soo
   // large.
   RunRandTest(9876789, 100000, -0, 7, 100);
 }
+
+TEST(AssumeQueryAPI, AssumptionCache) {
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> Mod = parseAssemblyString(
+      "declare void @llvm.assume(i1)\n"
+      "define void @test(i32* %P, i32* %P1, i32* %P2, i32* %P3, i1 %B) {\n"
+      "call void @llvm.assume(i1 true) [\"nonnull\"(i32* %P), \"align\"(i32* "
+      "%P2, i32 4), \"align\"(i32* %P, i32 8)]\n"
+      "call void @llvm.assume(i1 %B) [\"test\"(i32* %P1), "
+      "\"dereferenceable\"(i32* %P, i32 4)]\n"
+      "ret void\n}\n",
+      Err, C);
+  if (!Mod)
+    Err.print("AssumeQueryAPI", errs());
+  Function *F = Mod->getFunction("test");
+  BasicBlock::iterator First = F->begin()->begin();
+  BasicBlock::iterator Second = F->begin()->begin();
+  Second++;
+  AssumptionCacheTracker ACT;
+  AssumptionCache &AC = ACT.getAssumptionCache(*F);
+  auto AR = AC.assumptionsFor(F->getArg(3));
+  ASSERT_EQ(AR.size(), 0u);
+  AR = AC.assumptionsFor(F->getArg(1));
+  ASSERT_EQ(AR.size(), 1u);
+  ASSERT_EQ(AR[0].Index, 0u);
+  ASSERT_EQ(AR[0].Assume, &*Second);
+  AR = AC.assumptionsFor(F->getArg(2));
+  ASSERT_EQ(AR.size(), 1u);
+  ASSERT_EQ(AR[0].Index, 1u);
+  ASSERT_EQ(AR[0].Assume, &*First);
+  AR = AC.assumptionsFor(F->getArg(0));
+  ASSERT_EQ(AR.size(), 3u);
+  llvm::sort(AR,
+             [](const auto &L, const auto &R) { return L.Index < R.Index; });
+  ASSERT_EQ(AR[0].Assume, &*First);
+  ASSERT_EQ(AR[0].Index, 0u);
+  ASSERT_EQ(AR[1].Assume, &*Second);
+  ASSERT_EQ(AR[1].Index, 1u);
+  ASSERT_EQ(AR[2].Assume, &*First);
+  ASSERT_EQ(AR[2].Index, 2u);
+  AR = AC.assumptionsFor(F->getArg(4));
+  ASSERT_EQ(AR.size(), 1u);
+  ASSERT_EQ(AR[0].Assume, &*Second);
+  ASSERT_EQ(AR[0].Index, AssumptionCache::ExprResultIdx);
+  AC.unregisterAssumption(cast<CallInst>(&*Second));
+  AR = AC.assumptionsFor(F->getArg(1));
+  ASSERT_EQ(AR.size(), 0u);
+  AR = AC.assumptionsFor(F->getArg(0));
+  ASSERT_EQ(AR.size(), 3u);
+  llvm::sort(AR,
+             [](const auto &L, const auto &R) { return L.Index < R.Index; });
+  ASSERT_EQ(AR[0].Assume, &*First);
+  ASSERT_EQ(AR[0].Index, 0u);
+  ASSERT_EQ(AR[1].Assume, nullptr);
+  ASSERT_EQ(AR[1].Index, 1u);
+  ASSERT_EQ(AR[2].Assume, &*First);
+  ASSERT_EQ(AR[2].Index, 2u);
+  AR = AC.assumptionsFor(F->getArg(2));
+  ASSERT_EQ(AR.size(), 1u);
+  ASSERT_EQ(AR[0].Index, 1u);
+  ASSERT_EQ(AR[0].Assume, &*First);
+}