[Attributor] Introduce AAAssumptionInfo to propagate assumptions

[jhuber6]

Summary

This patch introduces a new abstract attributor instance that propagates
assumption information from functions. Conceptually, if a function is
only called by functions that have certain assumptions, then we can
apply the same assumptions to that function. This problem is similar to
calculating the dominator set, but the assumptions are merged instead of
nodes.

[jdoerfert]

llvm/include/llvm/IR/Assumptions.h

60

Consider passing the set as reference, unsure how much work copying it is

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

4750

I don't like the pair approach for long lived data. Make it a struct with properly named members. Costs the same but is much easier to read.

4755

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

15

Style: move down to the block of includes

9671

This should be in the constructor. Initialize can and should already look at the underlying value.

9724

9727

9732

Does this overwrite an existing assumptions attribute? Do we have a test for that?

9734

Adding assumptions is a change. Maybe return if it's empty with unchanged and otherwise add + CHANGED.

9737–9740

Not needed. If you just redirect to the base class, remove it.

9747

No this, also elsewhere.

9753

getIntersection returns if a change was made. CallSitePred should return false only if the traversal of call sites failed and we should give up. I think the return value of the CB here should be:
NotUniversal && AssumptionSet.empty() or something like that.

9759

This is not how checkForAllXXXX work. If they return false it means no all XXXX have been visited and you cannot rely on the information gathered.
Basically if (!checkForAllCXXXX(...)) return indicatePessimisticFixpoint(); is the way to go almost always.

9766

This won't be necessary because checkForAllCallSites returns false if you asked for *all* and not all call be visited.

9770

This is invariant, so it should happen in the initialize.

9809

Same comments as in the other manifest.

llvm/test/Transforms/Attributor/assumes_info.ll

3

Use the same run lines as other tests (4 total)

[jhuber6]

Addressing comments.

[jhuber6]

llvm/include/llvm/IR/Assumptions.h

60

Passing it as a reference would require some external lifetime right? Copy elision should apply in this case, but I'm not sure.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

9732

Made a function that always appends assumptions.

9770

The set at some points will still be in the universal state, we only want to add this attribute once it has some known values.

[jdoerfert]

llvm/lib/Transforms/IPO/Attributor.cpp

2507

We need to initialize this for call sites too, at least if they are not direct call sites.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

9651

9654

I would check the cheap part first, so isvalidstate.

9674

This makes it sound like a push-scheme while the updatImpl uses a pull-scheme.

9711

If the state has "known" and "assumed" assumptions (see also below) you can always fallback to the known ones if you indicate a pessimistic fixpoint. You basically reimplemented parts of the inidicatePessimistic/OptimisticFixpoint stuff here but only partially as part of the state. In line 9742 you do Assumed = Known, which is the definition of indicatePessimisticFixpoint in the integer and boolean state. The special case (line 9737/8) is not necessary at all as Known == Assume == empty in that case anyway.

9723

This is basically the "known assumptions" set while Assumptions is the "assumed assumptions" set. Maybe use those terms to match the rest of the AAs. We might even need to provide two APIs to query a known and an assumed assumption.
It's also not clear why this is only in the Function AA and not the Impl as call sites can have assumptions too which are not part of the caller or callee.

9756

I doubt assumptions.empty() means something changed. If you return changed and nothing changed you just force a timeout after the max iteration count.

[jhuber6]

Rewriting to use standard state encoding semantics. This currently crashes for
the test case provided when using the cgscc pass for some reason.

[jhuber6]

Changing known and assumed usage.

[jdoerfert]

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

2524

Unsure if I'd call this "UniversalSet" as it might not be a universal set ;)

2588

Why would it not be valid if assumed is not universal anymore?
One way is to define Assumed.empty() as invalid.

2624

You should not remove any known assumptions, so remove known from RHS first.

2642

We should not cache "changed" anywhere.

llvm/lib/IR/Assumptions.cpp

129

The two functions above are equal, make it a local template function and call it from the two exposed ones.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

9702

You can/should query the call site AA here. Through inlining and such the call site might have more assumptions than the caller.

9718

The way this is set up it might signal change even if things did not change. That is bad as it will spin until the iteration limit is reached. Either consider known in the intersect or keep track of the number of assumptions you had before and have now.

9730

Technically we can add the caller assumptions and the callee assumptions to the known set here if we want. Not needed right now though.

9732

No need to union something into assumed which is universal at this point.

9747

You can add the Assumed here. We are in a fixpoint state but the AA above might not have been informed of that yet.

9761

We should not look at the callee but caller. The callee is known and asked for "hasAssumption" queries anyway. The caller is what might disappear while the call site sticks around. So propagation direction is "caller -> call site -> callee".

9763

This doesn't actually update the state of this AA at all. You want to add the caller assumptions to the call site and indicate a change if that caused the call site to change.

[jhuber6]

Addressing comments.

[jhuber6]

More changes, I think it works this time.

[jhuber6]

Fixing isValidState.

[jhuber6]

Fixing issue not initializing isAtFixpoint.

[jhuber6]

The CallBase instance should inherit the attributes from the function it's
calling on initialization.

[jhuber6]

Fix assumption checking and manifest on call sites.

[jdoerfert]

LG, minor comments to address prior to the merge.

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

2521

2592

2620

2626

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

9754

The parent function assumptions should always be in the assumed set, adding them here should not be necessary.

9781

Also add the scope assumptions here, all of them are known.

[jhuber6]

Making suggested changes.

[jhuber6]

Fix bug and unused variable.

llvm/include/llvm/IR/Assumptions.h

 // the KnownAssumptionStrings set. This can be done in various ways, i.a.,
 // via a static KnownAssumptionString object.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_IR_ASSUMPTIONS_H
 #define LLVM_IR_ASSUMPTIONS_H
 
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 
 namespace llvm {
 
 class Function;
 class CallBase;
 
struct KnownAssumptionString {
   }
   operator StringRef() const { return AssumptionStr; }
 
 private:
   StringRef AssumptionStr;
 };
 
 /// Return true if \p F has the assumption \p AssumptionStr attached.
-bool hasAssumption(Function &F, const KnownAssumptionString &AssumptionStr);
+bool hasAssumption(const Function &F,
+                   const KnownAssumptionString &AssumptionStr);
 
 /// Return true if \p CB or the callee has the assumption \p AssumptionStr
 /// attached.
-bool hasAssumption(CallBase &CB, const KnownAssumptionString &AssumptionStr);
+bool hasAssumption(const CallBase &CB,
+                   const KnownAssumptionString &AssumptionStr);
+
+/// Return the set of all assumptions for the function \p F.
+DenseSet<StringRef> getAssumptions(const Function &F);
+
+/// Return the set of all assumptions for the call \p CB.
+DenseSet<StringRef> getAssumptions(const CallBase &CB);
+
+/// Appends the set of assumptions \p Assumptions to \F.
+bool addAssumptions(Function &F, const DenseSet<StringRef> &Assumptions);
+
+/// Appends the set of assumptions \p Assumptions to \CB.
+bool addAssumptions(CallBase &CB, const DenseSet<StringRef> &Assumptions);
 
 } // namespace llvm
 
 #endif

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

 
 #ifndef LLVM_TRANSFORMS_IPO_ATTRIBUTOR_H
 #define LLVM_TRANSFORMS_IPO_ATTRIBUTOR_H
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/Analysis/AssumeBundleQueries.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/LoopInfo.h"
struct IntegerRangeState : public AbstractState {
   IntegerRangeState operator&=(const IntegerRangeState &R) {
     // NOTE: `&=` operator seems like `intersect` but in this case, we need to
     // take `union`.
     unionKnown(R);
     unionAssumed(R);
     return *this;
   }
 };
+
+/// Simple state for a set.
+///
+/// This represents a state containing a set of values. The interface supports
+/// modelling sets that contain all possible elements. The state's internal
+/// value is modified using union or intersection operations.
+template <typename BaseTy> struct SetState : public AbstractState {
+  /// A wrapper around a set that has semantics for handling unions and
+  /// intersections with a "universal" set that contains all elements.
+  struct SetContents {
+    /// Creates a universal set with no concrete elements or an empty set.
+    SetContents(bool Universal) : Universal(Universal) {}
+
+    /// Creates a non-universal set with concrete values.
+    SetContents(const DenseSet<BaseTy> &Assumptions)
+        : Universal(false), Set(Assumptions) {}
+
+    SetContents(bool Universal, const DenseSet<BaseTy> &Assumptions)
+        : Universal(Universal), Set(Assumptions) {}
+
+    const DenseSet<BaseTy> &getSet() const { return Set; }
+
+    bool isUniversal() const { return Universal; }
+
+    bool empty() const { return Set.empty() && !Universal; }
+
+    /// Finds A := A ^ B where A or B could be the "Universal" set which
+    /// contains every possible attribute. Returns true if changes were made.
+    bool getIntersection(const SetContents &RHS) {
+      bool IsUniversal = Universal;
+      unsigned Size = Set.size();
+
+      // A := A ^ U = A
+      if (RHS.isUniversal())
+        return false;
+
+      // A := U ^ B = B
+      if (Universal)
+        Set = RHS.getSet();
+      else
+        set_intersect(Set, RHS.getSet());
+
+      Universal &= RHS.isUniversal();
+      return IsUniversal != Universal || Size != Set.size();
+    }
+
+    /// Finds A := A u B where A or B could be the "Universal" set which
+    /// contains every possible attribute. returns true if changes were made.
+    bool getUnion(const SetContents &RHS) {
+      bool IsUniversal = Universal;
+      unsigned Size = Set.size();
+
+      // A := A u U = U = U u B
+      if (!RHS.isUniversal() && !Universal)
+        set_union(Set, RHS.getSet());
+
+      Universal |= RHS.isUniversal();
+      return IsUniversal != Universal || Size != Set.size();
+    }
+
+  private:
+    /// Indicates if this set is "universal", containing every possible element.
+    bool Universal;
+
+    /// The set of currently active assumptions.
+    DenseSet<BaseTy> Set;
+  };
+
+  SetState() : Known(false), Assumed(true), IsAtFixedpoint(false) {}
+
+  /// Initializes the known state with an initial set and initializes the
+  /// assumed state as universal.
+  SetState(const DenseSet<BaseTy> &Known)
+      : Known(Known), Assumed(true), IsAtFixedpoint(false) {}
+
+  /// See AbstractState::isValidState()
+  bool isValidState() const override { return !Assumed.empty(); }
+
+  /// See AbstractState::isAtFixpoint()
+  bool isAtFixpoint() const override { return IsAtFixedpoint; }
+
+  /// See AbstractState::indicateOptimisticFixpoint(...)
+  ChangeStatus indicateOptimisticFixpoint() override {
+    IsAtFixedpoint = true;
+    Known = Assumed;
+    return ChangeStatus::UNCHANGED;
+  }
+
+  /// See AbstractState::indicatePessimisticFixpoint(...)
+  ChangeStatus indicatePessimisticFixpoint() override {
+    IsAtFixedpoint = true;
+    Assumed = Known;
+    return ChangeStatus::CHANGED;
+  }
+
+  /// Return the known state encoding.
+  const SetContents &getKnown() const { return Known; }
+
+  /// Return the assumed state encoding.
+  const SetContents &getAssumed() const { return Assumed; }
+
+  /// Returns if the set state contains the element.
+  bool setContains(const BaseTy &Elem) const {
+    return Assumed.getSet().contains(Elem) || Known.getSet().contains(Elem);
+  }
+
+  /// Performs the set intersection between this set and \p RHS. Returns true if
+  /// changes were made.
+  bool getIntersection(const SetContents &RHS) {
+    unsigned SizeBefore = Assumed.getSet().size();
+
+    // Get intersection and make sure that the known set is still a proper
+    // subset of the assumed set. A := K u (A ^ R).
+    Assumed.getIntersection(RHS);
+    Assumed.getUnion(Known);
+
+    return SizeBefore != Assumed.getSet().size();
+  }
+
+  /// Performs the set union between this set and \p RHS. Returns true if
+  /// changes were made.
+  bool getUnion(const SetContents &RHS) { return Assumed.getUnion(RHS); }
+
+private:
+  /// The set of values known for this state.
+  SetContents Known;
+
+  /// The set of assumed values for this state.
+  SetContents Assumed;
+
+  bool IsAtFixedpoint;
+};
+
 /// Helper struct necessary as the modular build fails if the virtual method
 /// IRAttribute::manifest is defined in the Attributor.cpp.
 struct IRAttributeManifest {
   static ChangeStatus manifestAttrs(Attributor &A, const IRPosition &IRP,
                                     const ArrayRef<Attribute> &DeducedAttrs,
                                     bool ForceReplace = false);
 };
 
struct AAPointerInfo : public AbstractAttribute {
   static bool classof(const AbstractAttribute *AA) {
     return (AA->getIdAddr() == &ID);
   }
 
   /// Unique ID (due to the unique address)
   static const char ID;
 };
 
+/// An abstract attribute for getting assumption information.
+struct AAAssumptionInfo
+    : public StateWrapper<SetState<StringRef>, AbstractAttribute,
+                          DenseSet<StringRef>> {
+  using Base =
+      StateWrapper<SetState<StringRef>, AbstractAttribute, DenseSet<StringRef>>;
+
+  AAAssumptionInfo(const IRPosition &IRP, Attributor &A,
+                   const DenseSet<StringRef> &Known)
+      : Base(IRP, Known) {}
+
+  /// Returns true if the assumption set contains the assumption \p Assumption.
+  virtual bool hasAssumption(const StringRef Assumption) const = 0;
+
+  /// Create an abstract attribute view for the position \p IRP.
+  static AAAssumptionInfo &createForPosition(const IRPosition &IRP,
+                                             Attributor &A);
+
+  /// See AbstractAttribute::getName()
+  const std::string getName() const override { return "AAAssumptionInfo"; }
+
+  /// See AbstractAttribute::getIdAddr()
+  const char *getIdAddr() const override { return &ID; }
+
+  /// This function should return true if the type of the \p AA is
+  /// AAAssumptionInfo
+  static bool classof(const AbstractAttribute *AA) {
+    return (AA->getIdAddr() == &ID);
+  }
+
+  /// Unique ID (due to the unique address)
+  static const char ID;
+};
+
 raw_ostream &operator<<(raw_ostream &, const AAPointerInfo::Access &);
 
 /// Run options, used by the pass manager.
 enum AttributorRunOption {
   NONE = 0,
   MODULE = 1 << 0,
   CGSCC = 1 << 1,
   ALL = MODULE | CGSCC
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_IPO_ATTRIBUTOR_H

llvm/lib/IR/Assumptions.cpp

 //===- Assumptions.cpp ------ Collection of helpers for assumptions -------===//
 //
 // 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 implements helper functions for accessing assumption infomration
+//  inside of the "llvm.assume" metadata.
+//
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Assumptions.h"
+#include "llvm/ADT/SetOperations.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
 
 using namespace llvm;
 
 namespace {
 bool hasAssumption(const Attribute &A,
                    const KnownAssumptionString &AssumptionStr) {
   if (!A.isValid())
     return false;
   assert(A.isStringAttribute() && "Expected a string attribute!");
 
   SmallVector<StringRef, 8> Strings;
   A.getValueAsString().split(Strings, ",");
 
   return llvm::is_contained(Strings, AssumptionStr);
 }
+
+DenseSet<StringRef> getAssumptions(const Attribute &A) {
+  if (!A.isValid())
+    return DenseSet<StringRef>();
+  assert(A.isStringAttribute() && "Expected a string attribute!");
+
+  DenseSet<StringRef> Assumptions;
+  SmallVector<StringRef, 8> Strings;
+  A.getValueAsString().split(Strings, ",");
+
+  for (StringRef Str : Strings)
+    Assumptions.insert(Str);
+  return Assumptions;
+}
+
+template <typename AttrSite>
+bool addAssumptionsImpl(AttrSite &Site,
+                        const DenseSet<StringRef> &Assumptions) {
+  if (Assumptions.empty())
+    return false;
+
+  DenseSet<StringRef> CurAssumptions = getAssumptions(Site);
+
+  if (!set_union(CurAssumptions, Assumptions))
+    return false;
+
+  LLVMContext &Ctx = Site.getContext();
+  Site.addFnAttr(llvm::Attribute::get(
+      Ctx, llvm::AssumptionAttrKey,
+      llvm::join(CurAssumptions.begin(), CurAssumptions.end(), ",")));
+
+  return true;
+}
 } // namespace
 
-bool llvm::hasAssumption(Function &F,
+bool llvm::hasAssumption(const Function &F,
                          const KnownAssumptionString &AssumptionStr) {
   const Attribute &A = F.getFnAttribute(AssumptionAttrKey);
   return ::hasAssumption(A, AssumptionStr);
 }
 
-bool llvm::hasAssumption(CallBase &CB,
+bool llvm::hasAssumption(const CallBase &CB,
                          const KnownAssumptionString &AssumptionStr) {
   if (Function *F = CB.getCalledFunction())
     if (hasAssumption(*F, AssumptionStr))
       return true;
 
   const Attribute &A = CB.getFnAttr(AssumptionAttrKey);
   return ::hasAssumption(A, AssumptionStr);
 }
 
+DenseSet<StringRef> llvm::getAssumptions(const Function &F) {
+  const Attribute &A = F.getFnAttribute(AssumptionAttrKey);
+  return ::getAssumptions(A);
+}
+
+DenseSet<StringRef> llvm::getAssumptions(const CallBase &CB) {
+  const Attribute &A = CB.getFnAttr(AssumptionAttrKey);
+  return ::getAssumptions(A);
+}
+
+bool llvm::addAssumptions(Function &F, const DenseSet<StringRef> &Assumptions) {
+  return ::addAssumptionsImpl(F, Assumptions);
+}
+
+bool llvm::addAssumptions(CallBase &CB,
+                          const DenseSet<StringRef> &Assumptions) {
+  return ::addAssumptionsImpl(CB, Assumptions);
+}
+
 StringSet<> llvm::KnownAssumptionStrings({
     "omp_no_openmp",          // OpenMP 5.1
     "omp_no_openmp_routines", // OpenMP 5.1
     "omp_no_parallelism",     // OpenMP 5.1
     "ompx_spmd_amenable",     // OpenMPOpt extension
 });

llvm/lib/Transforms/IPO/Attributor.cpp

void Attributor::identifyDefaultAbstractAttributes(Function &F) {
   getOrCreateAAFor<AANoRecurse>(FPos);
 
   // Every function might be "readnone/readonly/writeonly/...".
   getOrCreateAAFor<AAMemoryBehavior>(FPos);
 
   // Every function can be "readnone/argmemonly/inaccessiblememonly/...".
   getOrCreateAAFor<AAMemoryLocation>(FPos);
 
+  // Every function can track active assumptions.
+  getOrCreateAAFor<AAAssumptionInfo>(FPos);
+
   // Every function might be applicable for Heap-To-Stack conversion.
   if (EnableHeapToStack)
     getOrCreateAAFor<AAHeapToStack>(FPos);
 
   // Return attributes are only appropriate if the return type is non void.
   Type *ReturnType = F.getReturnType();
   if (!ReturnType->isVoidTy()) {
     // Argument attribute "returned" --- Create only one per function even
if (Arg.getType()->isPointerTy()) {
       // Every argument with pointer type might be privatizable (or promotable)
       getOrCreateAAFor<AAPrivatizablePtr>(ArgPos);
     }
   }
 
   auto CallSitePred = [&](Instruction &I) -> bool {
     auto &CB = cast<CallBase>(I);
     IRPosition CBRetPos = IRPosition::callsite_returned(CB);
+    IRPosition CBFnPos = IRPosition::callsite_function(CB);
 
     // Call sites might be dead if they do not have side effects and no live
     // users. The return value might be dead if there are no live users.
     getOrCreateAAFor<AAIsDead>(CBRetPos);
 
     Function *Callee = CB.getCalledFunction();
     // TODO: Even if the callee is not known now we might be able to simplify
     //       the call/callee.
     if (!Callee)
       return true;
 
+    // Every call site can track active assumptions.
+    getOrCreateAAFor<AAAssumptionInfo>(CBFnPos);
+
     // Skip declarations except if annotations on their call sites were
     // explicitly requested.
     if (!AnnotateDeclarationCallSites && Callee->isDeclaration() &&
         !Callee->hasMetadata(LLVMContext::MD_callback))
       return true;
 
     if (!Callee->getReturnType()->isVoidTy() && !CB.use_empty()) {
 

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

 //===- AttributorAttributes.cpp - Attributes for Attributor deduction -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // See the Attributor.h file comment and the class descriptions in that file for
 // more information.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/Attributor.h"
 
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/SCCIterator.h"
+#include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/AssumeBundleQueries.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LazyValueInfo.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/Assumptions.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/NoFolder.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/Casting.h"

 PIPE_OPERATOR(AAValueConstantRange)
 PIPE_OPERATOR(AAPrivatizablePtr)
 PIPE_OPERATOR(AAUndefinedBehavior)
 PIPE_OPERATOR(AAPotentialValues)
 PIPE_OPERATOR(AANoUndef)
 PIPE_OPERATOR(AACallEdges)
 PIPE_OPERATOR(AAFunctionReachability)
 PIPE_OPERATOR(AAPointerInfo)
+PIPE_OPERATOR(AAAssumptionInfo)
 
 #undef PIPE_OPERATOR
 
 template <>
 ChangeStatus clampStateAndIndicateChange<DerefState>(DerefState &S,
                                                      const DerefState &R) {
   ChangeStatus CS0 =
       clampStateAndIndicateChange(S.DerefBytesState, R.DerefBytesState);
size_t QueryCount =
         WholeFunction.Reachable.size() + WholeFunction.Unreachable.size();
 
     return "FunctionReachability [" +
            std::to_string(WholeFunction.Reachable.size()) + "," +
            std::to_string(QueryCount) + "]";
   }
 
   void trackStatistics() const override {}
+
 private:
   bool canReachUnknownCallee() const override {
     return WholeFunction.CanReachUnknownCallee;
   }
 
   /// Used to answer if a the whole function can reacha a specific function.
   QuerySet WholeFunction;
 
   /// Used to answer if a call base inside this function can reach a specific
   /// function.
   DenseMap<CallBase *, QuerySet> CBQueries;
 };
 
+/// ---------------------- Assumption Propagation ------------------------------
+struct AAAssumptionInfoImpl : public AAAssumptionInfo {
+  AAAssumptionInfoImpl(const IRPosition &IRP, Attributor &A,
+                       const DenseSet<StringRef> &Known)
+      : AAAssumptionInfo(IRP, A, Known) {}
+
+  bool hasAssumption(const StringRef Assumption) const override {
+    return isValidState() && setContains(Assumption);
+  }
+
+  /// See AbstractAttribute::getAsStr()
+  const std::string getAsStr() const override {
+    const SetContents &Known = getKnown();
+    const SetContents &Assumed = getAssumed();
+
+    const std::string KnownStr =
+        llvm::join(Known.getSet().begin(), Known.getSet().end(), ",");
+    const std::string AssumedStr =
+        (Assumed.isUniversal())
+            ? "Universal"
+            : llvm::join(Assumed.getSet().begin(), Assumed.getSet().end(), ",");
+
+    return "Known [" + KnownStr + "]," + " Assumed [" + AssumedStr + "]";
+  }
+};
+
+/// Propagates assumption information from parent functions to all of their
+/// successors. An assumption can be propagated if the containing function
+/// dominates the called function.
+///
+/// We start with a "known" set of assumptions already valid for the associated
+/// function and an "assumed" set that initially contains all possible
+/// assumptions. The assumed set is inter-procedurally updated by narrowing its
+/// contents as concrete values are known. The concrete values are seeded by the
+/// first nodes that are either entries into the call graph, or contains no
+/// assumptions. Each node is updated as the intersection of the assumed state
+/// with all of its predecessors.
+struct AAAssumptionInfoFunction final : AAAssumptionInfoImpl {
+  AAAssumptionInfoFunction(const IRPosition &IRP, Attributor &A)
+      : AAAssumptionInfoImpl(IRP, A,
+                             getAssumptions(*IRP.getAssociatedFunction())) {}
+
+  /// See AbstractAttribute::manifest(...).
+  ChangeStatus manifest(Attributor &A) override {
+    const auto &Assumptions = getKnown();
+
+    // Don't manifest a universal set if it somehow made it here.
+    if (Assumptions.isUniversal())
+      return ChangeStatus::UNCHANGED;
+
+    Function *AssociatedFunction = getAssociatedFunction();
+
+    bool Changed = addAssumptions(*AssociatedFunction, Assumptions.getSet());
+
+    return Changed ? ChangeStatus::CHANGED : ChangeStatus::UNCHANGED;
+  }
+
+  /// See AbstractAttribute::updateImpl(...).
+  ChangeStatus updateImpl(Attributor &A) override {
+    bool Changed = false;
+
+    auto CallSitePred = [&](AbstractCallSite ACS) {
+      const auto &AssumptionAA = A.getAAFor<AAAssumptionInfo>(
+          *this, IRPosition::callsite_function(*ACS.getInstruction()),
+          DepClassTy::REQUIRED);
+      // Get the set of assumptions shared by all of this function's callers.
+      Changed |= getIntersection(AssumptionAA.getAssumed());
+      return !getAssumed().empty() || !getKnown().empty();
+    };
+
+    bool AllCallSitesKnown;
+    // Get the intersection of all assumptions held by this node's predecessors.
+    // If we don't know all the call sites then this is either an entry into the
+    // call graph or an empty node. This node is known to only contain its own
+    // assumptions and can be propagated to its successors.
+    if (!A.checkForAllCallSites(CallSitePred, *this, true, AllCallSitesKnown))
+      return indicatePessimisticFixpoint();
+
+    return Changed ? ChangeStatus::CHANGED : ChangeStatus::UNCHANGED;
+  }
+
+  void trackStatistics() const override {}
+};
+
+/// Assumption Info defined for call sites.
+struct AAAssumptionInfoCallSite final : AAAssumptionInfoImpl {
+
+  AAAssumptionInfoCallSite(const IRPosition &IRP, Attributor &A)
+      : AAAssumptionInfoImpl(IRP, A, getInitialAssumptions(IRP)) {}
+
+  /// See AbstractAttribute::initialize(...).
+  void initialize(Attributor &A) override {
+    const IRPosition &FnPos = IRPosition::function(*getAnchorScope());
+    A.getAAFor<AAAssumptionInfo>(*this, FnPos, DepClassTy::REQUIRED);
+  }
+
+  /// See AbstractAttribute::manifest(...).
+  ChangeStatus manifest(Attributor &A) override {
+    // Don't manifest a universal set if it somehow made it here.
+    if (getKnown().isUniversal())
+      return ChangeStatus::UNCHANGED;
+
+    CallBase &AssociatedCall = cast<CallBase>(getAssociatedValue());
+    bool Changed = addAssumptions(AssociatedCall, getAssumed().getSet());
+
+    return Changed ? ChangeStatus::CHANGED : ChangeStatus::UNCHANGED;
+  }
+
+  /// See AbstractAttribute::updateImpl(...).
+  ChangeStatus updateImpl(Attributor &A) override {
+    const IRPosition &FnPos = IRPosition::function(*getAnchorScope());
+    auto &AssumptionAA =
+        A.getAAFor<AAAssumptionInfo>(*this, FnPos, DepClassTy::REQUIRED);
+    bool Changed = getIntersection(AssumptionAA.getAssumed());
+    return Changed ? ChangeStatus::CHANGED : ChangeStatus::UNCHANGED;
+  }
+
+  /// See AbstractAttribute::trackStatistics()
+  void trackStatistics() const override {}
+
+private:
+  /// Helper to initialized the known set as all the assumptions this call and
+  /// the callee contain.
+  DenseSet<StringRef> getInitialAssumptions(const IRPosition &IRP) {
+    const CallBase &CB = cast<CallBase>(IRP.getAssociatedValue());
+    auto Assumptions = getAssumptions(CB);
+    if (Function *F = IRP.getAssociatedFunction())
+      set_union(Assumptions, getAssumptions(*F));
+    if (Function *F = IRP.getAssociatedFunction())
+      set_union(Assumptions, getAssumptions(*F));
+    return Assumptions;
+  }
+};
+
 } // namespace
 
 AACallGraphNode *AACallEdgeIterator::operator*() const {
   return static_cast<AACallGraphNode *>(const_cast<AACallEdges *>(
       &A.getOrCreateAAFor<AACallEdges>(IRPosition::function(**I))));
 }
 
 void AttributorCallGraph::print() { llvm::WriteGraph(outs(), this); }
 
 const char AAReturnedValues::ID = 0;

 const char AADereferenceable::ID = 0;
 const char AAAlign::ID = 0;
 const char AANoCapture::ID = 0;
 const char AAValueSimplify::ID = 0;
 const char AAHeapToStack::ID = 0;
 const char AAPrivatizablePtr::ID = 0;
 const char AAMemoryBehavior::ID = 0;
 const char AAMemoryLocation::ID = 0;
 const char AAValueConstantRange::ID = 0;
 const char AAPotentialValues::ID = 0;
 const char AANoUndef::ID = 0;
 const char AACallEdges::ID = 0;
 const char AAFunctionReachability::ID = 0;
 const char AAPointerInfo::ID = 0;
+const char AAAssumptionInfo::ID = 0;
 
 // Macro magic to create the static generator function for attributes that
 // follow the naming scheme.
 
 #define SWITCH_PK_INV(CLASS, PK, POS_NAME)                                     \
   case IRPosition::PK:                                                         \
     llvm_unreachable("Cannot create " #CLASS " for a " POS_NAME " position!");
 

 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoUnwind)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoSync)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoRecurse)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAWillReturn)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoReturn)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAReturnedValues)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAMemoryLocation)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AACallEdges)
+CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAAssumptionInfo)
 
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANonNull)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoAlias)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAPrivatizablePtr)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AADereferenceable)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAAlign)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoCapture)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAValueConstantRange)

llvm/test/Transforms/Attributor/assumes_info.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-attributes --check-globals
+; RUN: opt -attributor -enable-new-pm=0 -attributor-manifest-internal  -attributor-max-iterations-verify -attributor-annotate-decl-cs -attributor-max-iterations=3 -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_CGSCC_NPM,NOT_CGSCC_OPM,NOT_TUNIT_NPM,IS__TUNIT____,IS________OPM,IS__TUNIT_OPM
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal  -attributor-max-iterations-verify -attributor-annotate-decl-cs -attributor-max-iterations=3 -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_CGSCC_OPM,NOT_CGSCC_NPM,NOT_TUNIT_OPM,IS__TUNIT____,IS________NPM,IS__TUNIT_NPM
+; RUN: opt -attributor-cgscc -enable-new-pm=0 -attributor-manifest-internal  -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_TUNIT_NPM,NOT_TUNIT_OPM,NOT_CGSCC_NPM,IS__CGSCC____,IS________OPM,IS__CGSCC_OPM
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor-cgscc -attributor-manifest-internal  -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_TUNIT_NPM,NOT_TUNIT_OPM,NOT_CGSCC_OPM,IS__CGSCC____,IS________NPM,IS__CGSCC_NPM
+
+define dso_local void @entry(i1 %cond) #0 {
+; CHECK-LABEL: define {{[^@]+}}@entry
+; CHECK-SAME: (i1 [[COND:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @foo(i1 [[COND]]) #[[ATTR1:[0-9]+]]
+; CHECK-NEXT:    call void @bar() #[[ATTR2:[0-9]+]]
+; CHECK-NEXT:    call void @qux() #[[ATTR1]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  call void @foo(i1 %cond)
+  call void @bar()
+  call void @qux() #1
+  ret void
+}
+
+define internal void @foo(i1 %cond) #1 {
+; CHECK-LABEL: define {{[^@]+}}@foo
+; CHECK-SAME: (i1 [[COND:%.*]]) #[[ATTR1]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @baz(i1 [[COND]]) #[[ATTR1]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  call void @baz(i1 %cond)
+  ret void
+}
+
+define internal void @bar() #2 {
+; CHECK-LABEL: define {{[^@]+}}@bar
+; CHECK-SAME: () #[[ATTR2]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @baz(i1 noundef false) #[[ATTR2]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  call void @baz(i1 0)
+  ret void
+}
+
+define internal void @baz(i1 %Cond) {
+; CHECK-LABEL: define {{[^@]+}}@baz
+; CHECK-SAME: (i1 [[COND:%.*]]) #[[ATTR1]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp ne i1 [[COND]], false
+; CHECK-NEXT:    br i1 [[TOBOOL]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    call void @baz(i1 noundef false) #[[ATTR1]]
+; CHECK-NEXT:    br label [[IF_END]]
+; CHECK:       if.end:
+; CHECK-NEXT:    call void @qux() #[[ATTR1]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %tobool = icmp ne i1 %Cond, 0
+  br i1 %tobool, label %if.then, label %if.end
+
+if.then:
+  call void @baz(i1 0)
+  br label %if.end
+
+if.end:
+  call void @qux()
+  ret void
+}
+
+define internal void @qux() {
+; CHECK-LABEL: define {{[^@]+}}@qux
+; CHECK-SAME: () #[[ATTR1]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @call() #[[ATTR2]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  call void @call()
+  ret void
+}
+
+declare void @call() #3
+
+attributes #0 = { "llvm.assume"="A" }
+attributes #1 = { "llvm.assume"="B" }
+attributes #2 = { "llvm.assume"="B,C" }
+attributes #3 = { "llvm.assume"="B,C,A" }
+;.
+; CHECK: attributes #[[ATTR0]] = { "llvm.assume"="A" }
+; CHECK: attributes #[[ATTR1]] = { "llvm.assume"="B,A" }
+; CHECK: attributes #[[ATTR2]] = { "llvm.assume"="B,C,A" }
+;.

llvm/test/Transforms/Attributor/depgraph.ll

 ; GRAPH-NEXT:   updates [AANoCapture] for CtxI '  %2 = load i32, i32* %0, align 4' at position {arg: [@0]} with state assumed not-captured-maybe-returned
 ; GRAPH-NEXT:   updates [AANoCapture] for CtxI '  %2 = load i32, i32* %0, align 4' at position {arg: [@0]} with state assumed not-captured-maybe-returned
 ; GRAPH-NEXT:   updates [AANoCapture] for CtxI '  %2 = load i32, i32* %0, align 4' at position {arg: [@0]} with state assumed not-captured-maybe-returned
 ; GRAPH-EMPTY:
 ; GRAPH-NEXT: [AAMemoryLocation] for CtxI '  %2 = load i32, i32* %0, align 4' at position {fn:checkAndAdvance [checkAndAdvance@-1]} with state memory:argument
 ; GRAPH-NEXT:   updates [AAMemoryLocation] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs: [@-1]} with state memory:argument
 ; GRAPH-NEXT:   updates [AAMemoryLocation] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs: [@-1]} with state memory:argument
 ; GRAPH-EMPTY:
+; GRAPH-NEXT:  [AAAssumptionInfo] for CtxI ' %2 = load i32, i32* %0, align 4' at position {fn:checkAndAdvance [checkAndAdvance@-1]} with state Known [], Assumed []
+; GRAPH-EMPTY:
 ; GRAPH-NEXT: [AAHeapToStack] for CtxI '  %2 = load i32, i32* %0, align 4' at position {fn:checkAndAdvance [checkAndAdvance@-1]} with state [H2S] Mallocs Good/Bad: 0/0
 ; GRAPH-EMPTY:
 ; GRAPH-NEXT: [AAValueSimplify] for CtxI '  %2 = load i32, i32* %0, align 4' at position {fn_ret:checkAndAdvance [checkAndAdvance@-1]} with state not-simple
 ; GRAPH-EMPTY:
 ; GRAPH-NEXT: [AAAlign] for CtxI '  %2 = load i32, i32* %0, align 4' at position {fn_ret:checkAndAdvance [checkAndAdvance@-1]} with state align<1-16>
 ; GRAPH-NEXT:   updates [AAAlign] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_ret: [@-1]} with state align<1-16>
 ; GRAPH-NEXT:   updates [AAAlign] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_ret: [@-1]} with state align<1-16>
 ; GRAPH-EMPTY:

 ; GRAPH-NEXT:   updates [AAIsDead] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_ret: [@-1]} with state assumed-live
 ; GRAPH-NEXT:   updates [AAIsDead] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_ret: [@-1]} with state assumed-live
 ; GRAPH-NEXT:   updates [AANoUnwind] for CtxI '  %2 = load i32, i32* %0, align 4' at position {fn:checkAndAdvance [checkAndAdvance@-1]} with state nounwind
 ; GRAPH-EMPTY:
 ; GRAPH-NEXT: [AAMemoryBehavior] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs: [@-1]} with state readonly
 ; GRAPH-NEXT:   updates [AAIsDead] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_ret: [@-1]} with state assumed-live
 ; GRAPH-NEXT:   updates [AAMemoryBehavior] for CtxI '  %2 = load i32, i32* %0, align 4' at position {fn:checkAndAdvance [checkAndAdvance@-1]} with state readonly
 ; GRAPH-EMPTY:
+; GRAPH-NEXT: [AAAssumptionInfo] for CtxI ' %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs: [@-1]} with state Known [], Assumed []
+; GRAPH-EMPTY:
 ; GRAPH-NEXT: [AAValueSimplify] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_ret: [@-1]} with state not-simple
 ; GRAPH-EMPTY:
 ; GRAPH-NEXT: [AAIsDead] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_arg: [@0]} with state assumed-live
 ; GRAPH-EMPTY:
 ; GRAPH-NEXT: [AAValueSimplify] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_arg: [@0]} with state simplified
 ; GRAPH-EMPTY:
 ; GRAPH-NEXT: [AANoUndef] for CtxI '  %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_arg: [@0]} with state may-undef-or-poison
 ; GRAPH-EMPTY: