[Attributor][FIX] Do treat byval arguments special

When we reason about the pointer argument that is byval we actually
reason about a local copy of the value passed at the call site. This was
not the case before and we wrongly introduced attributes based on the
surrounding function.

AAMemoryBehaviorArgument, AAMemoryBehaviorCallSiteArgument and
AANoCaptureCallSiteArgument are made aware of byval now. The code
to skip "subsuming positions" for reasoning follows a common pattern and
we should refactor it. A TODO was added.

Discovered by @efriedma as part of D69748.

3 files changed, 117 insertions, 13 deletions

diff --git a/llvm/include/llvm/Transforms/IPO/Attributor.h b/llvm/include/llvm/Transforms/IPO/Attributor.h
index bacbccda02c..b931d0f5746 100644
--- a/llvm/include/llvm/Transforms/IPO/Attributor.h
+++ b/llvm/include/llvm/Transforms/IPO/Attributor.h
@@ -398,8 +398,12 @@ struct IRPosition {
   /// 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) const;
+                SmallVectorImpl<Attribute> &Attrs,
+                bool IgnoreSubsumingPositions = false) const;
 
   /// Return the attribute of kind \p AK existing in the IR at this position.
   Attribute getAttr(Attribute::AttrKind AK) const {
diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index 305cc289823..07033c03fd6 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -465,13 +465,20 @@ bool IRPosition::hasAttr(ArrayRef<Attribute::AttrKind> AKs,
 }
 
 void IRPosition::getAttrs(ArrayRef<Attribute::AttrKind> AKs,
-                          SmallVectorImpl<Attribute> &Attrs) const {
-  for (const IRPosition &EquivIRP : SubsumingPositionIterator(*this))
+                          SmallVectorImpl<Attribute> &Attrs,
+                          bool IgnoreSubsumingPositions) const {
+  for (const IRPosition &EquivIRP : SubsumingPositionIterator(*this)) {
     for (Attribute::AttrKind AK : AKs) {
       const Attribute &Attr = EquivIRP.getAttr(AK);
       if (Attr.getKindAsEnum() == AK)
         Attrs.push_back(Attr);
     }
+    // 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;
+  }
 }
 
 void IRPosition::verify() {
@@ -3773,6 +3780,14 @@ struct AANoCaptureArgument final : AANoCaptureImpl {
 struct AANoCaptureCallSiteArgument final : AANoCaptureImpl {
   AANoCaptureCallSiteArgument(const IRPosition &IRP) : AANoCaptureImpl(IRP) {}
 
+  /// See AbstractAttribute::initialize(...).
+  void initialize(Attributor &A) override {
+    if (Argument *Arg = getAssociatedArgument())
+      if (Arg->hasByValAttr())
+        indicateOptimisticFixpoint();
+    AANoCaptureImpl::initialize(A);
+  }
+
   /// See AbstractAttribute::updateImpl(...).
   ChangeStatus updateImpl(Attributor &A) override {
     // TODO: Once we have call site specific value information we can provide
@@ -4337,9 +4352,10 @@ struct AAMemoryBehaviorImpl : public AAMemoryBehavior {
 
   /// Return the memory behavior information encoded in the IR for \p IRP.
   static void getKnownStateFromValue(const IRPosition &IRP,
-                                     BitIntegerState &State) {
+                                     BitIntegerState &State,
+                                     bool IgnoreSubsumingPositions = false) {
     SmallVector<Attribute, 2> Attrs;
-    IRP.getAttrs(AttrKinds, Attrs);
+    IRP.getAttrs(AttrKinds, Attrs, IgnoreSubsumingPositions);
     for (const Attribute &Attr : Attrs) {
       switch (Attr.getKindAsEnum()) {
       case Attribute::ReadNone:
@@ -4461,12 +4477,25 @@ struct AAMemoryBehaviorArgument : AAMemoryBehaviorFloating {
 
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
-    AAMemoryBehaviorFloating::initialize(A);
+    intersectAssumedBits(BEST_STATE);
+    const IRPosition &IRP = getIRPosition();
+    // TODO: Make IgnoreSubsumingPositions a property of an IRAttribute so we
+    // can query it when we use has/getAttr. That would allow us to reuse the
+    // initialize of the base class here.
+    bool HasByVal =
+        IRP.hasAttr({Attribute::ByVal}, /* IgnoreSubsumingPositions */ true);
+    getKnownStateFromValue(IRP, getState(),
+                           /* IgnoreSubsumingPositions */ HasByVal);
 
     // Initialize the use vector with all direct uses of the associated value.
     Argument *Arg = getAssociatedArgument();
-    if (!Arg || !Arg->getParent()->hasExactDefinition())
+    if (!Arg || !Arg->getParent()->hasExactDefinition()) {
       indicatePessimisticFixpoint();
+    } else {
+      // Initialize the use vector with all direct uses of the associated value.
+      for (const Use &U : Arg->uses())
+        Uses.insert(&U);
+    }
   }
 
   ChangeStatus manifest(Attributor &A) override {
@@ -4494,6 +4523,19 @@ struct AAMemoryBehaviorCallSiteArgument final : AAMemoryBehaviorArgument {
   AAMemoryBehaviorCallSiteArgument(const IRPosition &IRP)
       : AAMemoryBehaviorArgument(IRP) {}
 
+  /// See AbstractAttribute::initialize(...).
+  void initialize(Attributor &A) override {
+    if (Argument *Arg = getAssociatedArgument()) {
+      if (Arg->hasByValAttr()) {
+        addKnownBits(NO_WRITES);
+        removeKnownBits(NO_READS);
+        removeAssumedBits(NO_READS);
+      }
+    } else {
+    }
+    AAMemoryBehaviorArgument::initialize(A);
+  }
+
   /// See AbstractAttribute::updateImpl(...).
   ChangeStatus updateImpl(Attributor &A) override {
     // TODO: Once we have call site specific value information we can provide
@@ -4640,11 +4682,17 @@ ChangeStatus AAMemoryBehaviorFloating::updateImpl(Attributor &A) {
 
   // First, check the function scope. We take the known information and we avoid
   // work if the assumed information implies the current assumed information for
-  // this attribute.
-  const auto &FnMemAA = A.getAAFor<AAMemoryBehavior>(*this, FnPos);
-  S.addKnownBits(FnMemAA.getKnown());
-  if ((S.getAssumed() & FnMemAA.getAssumed()) == S.getAssumed())
-    return ChangeStatus::UNCHANGED;
+  // this attribute. This is a valid for all but byval arguments.
+  Argument *Arg = IRP.getAssociatedArgument();
+  AAMemoryBehavior::base_t FnMemAssumedState =
+      AAMemoryBehavior::StateType::getWorstState();
+  if (!Arg || !Arg->hasByValAttr()) {
+    const auto &FnMemAA = A.getAAFor<AAMemoryBehavior>(*this, FnPos);
+    FnMemAssumedState = FnMemAA.getAssumed();
+    S.addKnownBits(FnMemAA.getKnown());
+    if ((S.getAssumed() & FnMemAA.getAssumed()) == S.getAssumed())
+      return ChangeStatus::UNCHANGED;
+  }
 
   // Make sure the value is not captured (except through "return"), if
   // it is, any information derived would be irrelevant anyway as we cannot
@@ -4652,7 +4700,7 @@ ChangeStatus AAMemoryBehaviorFloating::updateImpl(Attributor &A) {
   // to fall back to anythign less optimistic than the function state.
   const auto &ArgNoCaptureAA = A.getAAFor<AANoCapture>(*this, IRP);
   if (!ArgNoCaptureAA.isAssumedNoCaptureMaybeReturned()) {
-    S.intersectAssumedBits(FnMemAA.getAssumed());
+    S.intersectAssumedBits(FnMemAssumedState);
     return ChangeStatus::CHANGED;
   }
 
diff --git a/llvm/test/Transforms/Attributor/readattrs.ll b/llvm/test/Transforms/Attributor/readattrs.ll
index 9c148ef160b..cfb4f71ce0e 100644
--- a/llvm/test/Transforms/Attributor/readattrs.ll
+++ b/llvm/test/Transforms/Attributor/readattrs.ll
@@ -143,3 +143,55 @@ define void @unsound_readonly(i8* %ignored, i8* %escaped_then_written) {
   store i8 0, i8* %addr.ld
   ret void
 }
+
+; Byval but not readonly/none tests
+;
+;{
+declare void @escape_i8(i8* %ptr)
+
+; ATTRIBUTOR:      @byval_not_readonly_1
+; ATTRIBUTOR-SAME: i8* byval %written
+define void @byval_not_readonly_1(i8* byval %written) readonly {
+  call void @escape_i8(i8* %written)
+  ret void
+}
+
+; ATTRIBUTOR:      @byval_not_readonly_2
+; ATTRIBUTOR-SAME: i8* nocapture nofree nonnull writeonly byval dereferenceable(1) %written
+define void @byval_not_readonly_2(i8* byval %written) readonly {
+  store i8 0, i8* %written
+  ret void
+}
+
+; ATTRIBUTOR:      @byval_not_readnone_1
+; ATTRIBUTOR-SAME: i8* byval %written
+define void @byval_not_readnone_1(i8* byval %written) readnone {
+  call void @escape_i8(i8* %written)
+  ret void
+}
+
+; ATTRIBUTOR:      @byval_not_readnone_2
+; ATTRIBUTOR-SAME: i8* nocapture nofree nonnull writeonly byval dereferenceable(1) %written
+define void @byval_not_readnone_2(i8* byval %written) readnone {
+  store i8 0, i8* %written
+  ret void
+}
+
+; ATTRIBUTOR:      @byval_no_fnarg
+; ATTRIBUTOR-SAME: i8* nocapture nofree nonnull writeonly byval dereferenceable(1) %written
+define void @byval_no_fnarg(i8* byval %written) {
+  store i8 0, i8* %written
+  ret void
+}
+
+; ATTRIBUTOR: @testbyval
+; ATTRIBUTOR-SAME: i8* nocapture readonly %read_only
+define void @testbyval(i8* %read_only) {
+  call void @byval_not_readonly_1(i8* %read_only)
+  call void @byval_not_readonly_2(i8* %read_only)
+  call void @byval_not_readnone_1(i8* %read_only)
+  call void @byval_not_readnone_2(i8* %read_only)
+  call void @byval_no_fnarg(i8* %read_only)
+  ret void
+}
+;}