[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

681

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?

748–753

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

[jdoerfert]

llvm/lib/Transforms/IPO/Attributor.cpp

681

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.

748–753

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

748–753

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

681

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;
}

748–753

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/Transforms/IPO/Attributor.h

 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/ConstantRange.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;
 
   /// The argument number, if non-negative, or the position "kind".
   int KindOrArgNo;
 };
 

 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/Transforms/IPO/Attributor.cpp

 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/NoFolder.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/InitializePasses.h"
-#include "llvm/IR/NoFolder.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/IPO/ArgumentPromotion.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/KnowledgeRetention.h"
 #include "llvm/Transforms/Utils/Local.h"
 
 #include <cassert>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "attributor"
 
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)
+    for (Attribute::AttrKind AK : AKs) {
       if (EquivIRP.getAttrsFromIRAttr(AK, Attrs))
         return true;
+      if (A && EquivIRP.getAttrsFromAssumes(AK, Attrs, *A))
+        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;
+    // Do not perform an search for assumes on subsuming positions.
+    A = nullptr;
   }
   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)
+    for (Attribute::AttrKind AK : AKs) {
       EquivIRP.getAttrsFromIRAttr(AK, Attrs);
+      if (A)
+        EquivIRP.getAttrsFromAssumes(AK, Attrs, *A);
+    }
     // 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;
+    // Do not perform an search for assumes on subsuming positions.
+    A = nullptr;
   }
 }
 
 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();
+  LLVMContext &Ctx = AssociatedValue.getContext();
+  uint64_t ArgVal = 0;
+  uint64_t *ArgValPtr =
+      Attribute::doesAttrKindHaveArgument(AK) ? &ArgVal : nullptr;
+  StringRef AttrName = Attribute::getNameFromAttrKind(AK);
+  DenseMap<Instruction *, Attribute> AssumeAttrMap;
+  for (const Use &U : AssociatedValue.uses())
+    if (IntrinsicInst *II = dyn_cast_or_null<IntrinsicInst>(U.getUser()))
+      if (II->getIntrinsicID() == Intrinsic::assume && !II->isArgOperand(&U))
+        if (hasAttributeInAssume(*II, &AssociatedValue, AttrName, ArgValPtr))
+          if (!ArgValPtr || ArgVal > 0)
+            AssumeAttrMap[II] = Attribute::get(Ctx, AK, ArgVal);
+
+  // Check if we found any potential assume use, if not we don't need to create
+  // explorer iterators.
+  if (AssumeAttrMap.empty())
+    return false;
+
+  unsigned AttrsSize = Attrs.size();
+  MustBeExecutedContextExplorer &Explorer =
+      A.getInfoCache().getMustBeExecutedContextExplorer();
+  auto EIt = Explorer.begin(getCtxI()), EEnd = Explorer.end(getCtxI());
+  for (auto &It : AssumeAttrMap)
+    if (Explorer.findInContextOf(It.first, EIt, EEnd))
+      Attrs.push_back(It.second);
+  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)) {
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();

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

#include "llvm/IR/Attributes.inc"
       .Default(false);
 }
 
 #endif
 
 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(isExistingAttribute(AttrName) && "this attribute doesn't exist");
   assert((ArgVal == nullptr || Attribute::doesAttrKindHaveArgument(
                                    Attribute::getAttrKindFromName(AttrName))) &&
          "requested value for an attribute that has no argument");
   if (Assume.bundle_op_infos().empty())
     return false;
 
   CallInst::bundle_op_iterator Lookup;
 
   /// The right attribute can be found by binary search. After this finding the
   /// right WasOn needs to be done via linear search.
   /// Element have been ordered by argument value so the first we find is the
   /// one we need.
-  if (AQR == AssumeQuery::Lowest)
+  if (AQR == AssumeQuery::Lowest) {
     Lookup =
         llvm::lower_bound(Assume.bundle_op_infos(), AttrName,
                           [](const CallBase::BundleOpInfo &BOI, StringRef RHS) {
                             assert(isExistingAttribute(BOI.Tag->getKey()) &&
                                    "this attribute doesn't exist");
                             return BOI.Tag->getKey() < RHS;
                           });
-  else
-    Lookup = std::prev(
+  } else {
+    Lookup =
         llvm::upper_bound(Assume.bundle_op_infos(), AttrName,
                           [](StringRef LHS, const CallBase::BundleOpInfo &BOI) {
                             assert(isExistingAttribute(BOI.Tag->getKey()) &&
                                    "this attribute doesn't exist");
                             return LHS < BOI.Tag->getKey();
-                          }));
+                          });
+    // If we are at the very beginning the tag was not found. If we are not at
+    // the very beginning we don't know if the tag was found but we can go back
+    // to find out.
+    if (Lookup == Assume.bundle_op_info_begin())
+      Lookup = Assume.bundle_op_info_end();
+    else
+      Lookup = std::prev(Lookup);
+  }
 
   auto getValueFromBundleOpInfo = [&Assume](const CallBase::BundleOpInfo &BOI,
                                             unsigned Idx) {
     assert(BOI.End - BOI.Begin > Idx && "index out of range");
     return (Assume.op_begin() + BOI.Begin + Idx)->get();
   };
 
   if (Lookup == Assume.bundle_op_info_end() ||

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

 ;
 entry:
   %0 = load i64, i64* %range, align 8, !range !0
   tail call void @fill_range_inbounds(i32* %p, i64 %0)
   tail call void @fill_range_not_inbounds(i32* %p, i64 %0)
   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 readnone [[ARG3:%.*]], i8* nocapture nofree readnone dereferenceable_or_null(42) [[ARG4:%.*]])
+; ATTRIBUTOR-NEXT:    call void @unknown()
+; ATTRIBUTOR-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) ["dereferenceable"(i8* %arg1, i64 1), "dereferenceable"(i8* %arg1, i64 2), "dereferenceable"(i8* %arg1, i64 101), "dereferenceable"(i8* %arg3, i64 0), "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:    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
+}
+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
+
+declare void @unknown(i8*, i8*, i8*, i8*)
+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 @unknown(i8* [[ARG1]], i8* [[ARG2]], i8* [[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:    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
+}
+declare void @llvm.assume(i1)
+
 declare noalias i8* @malloc(i64)
 
 attributes #0 = { nounwind uwtable noinline }
 attributes #1 = { nounwind }
 attributes #2 = { nobuiltin nounwind }

llvm/test/Transforms/Attributor/nonnull.ll

 ;
   %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:    [[TMP1:%.*]] = call i8* @unknown()
+; ATTRIBUTOR-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) ["nonnull"(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:    ret void
+;
+  call i8* @unknown()
+  call void @llvm.assume(i1 true) ["nonnull"(i8* %arg)]
+  ret void
+}
+
 attributes #0 = { "null-pointer-is-valid"="true" }
 attributes #1 = { nounwind willreturn}