[Attributor] Teach AANoCapture to use information in-flight more aggressively

AAReturnedValues, AAMemoryBehavior, and AANoUnwind, can provide
information that helps during the tracking or even justifies no-capture.
We now use this information and enable no-capture in some test cases
designed a long while a ago for these cases.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@375382 91177308-0d34-0410-b5e6-96231b3b80d8

6 files changed, 103 insertions, 21 deletions

diff --git a/lib/Transforms/IPO/Attributor.cpp b/lib/Transforms/IPO/Attributor.cpp
index 11a54bdb0a7..e7b7404b171 100644
--- a/lib/Transforms/IPO/Attributor.cpp
+++ b/lib/Transforms/IPO/Attributor.cpp
@@ -2939,7 +2939,7 @@ struct AANoCaptureImpl : public AANoCapture {
 
     // Check what state the associated function can actually capture.
     if (F)
-      determineFunctionCaptureCapabilities(*F, *this);
+      determineFunctionCaptureCapabilities(IRP, *F, *this);
     else
       indicatePessimisticFixpoint();
   }
@@ -2965,7 +2965,8 @@ struct AANoCaptureImpl : public AANoCapture {
   /// Set the NOT_CAPTURED_IN_MEM and NOT_CAPTURED_IN_RET bits in \p Known
   /// depending on the ability of the function associated with \p IRP to capture
   /// state in memory and through "returning/throwing", respectively.
-  static void determineFunctionCaptureCapabilities(const Function &F,
+  static void determineFunctionCaptureCapabilities(const IRPosition &IRP,
+                                                   const Function &F,
                                                    IntegerState &State) {
     // TODO: Once we have memory behavior attributes we should use them here.
 
@@ -2987,6 +2988,21 @@ struct AANoCaptureImpl : public AANoCapture {
     // exceptions and doesn not return values.
     if (F.doesNotThrow() && F.getReturnType()->isVoidTy())
       State.addKnownBits(NOT_CAPTURED_IN_RET);
+
+    // Check existing "returned" attributes.
+    int ArgNo = IRP.getArgNo();
+    if (F.doesNotThrow() && ArgNo >= 0) {
+      for (unsigned u = 0, e = F.arg_size(); u< e; ++u)
+        if (F.hasParamAttribute(u, Attribute::Returned)) {
+          if (u == ArgNo)
+            State.removeAssumedBits(NOT_CAPTURED_IN_RET);
+          else if (F.onlyReadsMemory())
+            State.addKnownBits(NO_CAPTURE);
+          else
+            State.addKnownBits(NOT_CAPTURED_IN_RET);
+          break;
+        }
+    }
   }
 
   /// See AbstractState::getAsStr().
@@ -3158,15 +3174,54 @@ ChangeStatus AANoCaptureImpl::updateImpl(Attributor &A) {
   const Function *F =
       getArgNo() >= 0 ? IRP.getAssociatedFunction() : IRP.getAnchorScope();
   assert(F && "Expected a function!");
-  const auto &IsDeadAA = A.getAAFor<AAIsDead>(*this, IRPosition::function(*F));
+  const IRPosition &FnPos = IRPosition::function(*F);
+  const auto &IsDeadAA = A.getAAFor<AAIsDead>(*this, FnPos);
 
   AANoCapture::StateType T;
-  // TODO: Once we have memory behavior attributes we should use them here
-  // similar to the reasoning in
-  // AANoCaptureImpl::determineFunctionCaptureCapabilities(...).
 
-  // TODO: Use the AAReturnedValues to learn if the argument can return or
-  // not.
+  // Readonly means we cannot capture through memory.
+  const auto &FnMemAA = A.getAAFor<AAMemoryBehavior>(*this, FnPos);
+  if (FnMemAA.isAssumedReadOnly()) {
+    T.addKnownBits(NOT_CAPTURED_IN_MEM);
+    if (FnMemAA.isKnownReadOnly())
+      addKnownBits(NOT_CAPTURED_IN_MEM);
+  }
+
+  // Make sure all returned values are different than the underlying value.
+  // TODO: we could do this in a more sophisticated way inside
+  //       AAReturnedValues, e.g., track all values that escape through returns
+  //       directly somehow.
+  auto CheckReturnedArgs = [&](const AAReturnedValues &RVAA) {
+    bool SeenConstant = false;
+    for (auto &It : RVAA.returned_values()) {
+      if (isa<Constant>(It.first)) {
+        if (SeenConstant)
+          return false;
+        SeenConstant = true;
+      } else if (!isa<Argument>(It.first) ||
+                 It.first == getAssociatedArgument())
+        return false;
+    }
+    return true;
+  };
+
+  const auto &NoUnwindAA = A.getAAFor<AANoUnwind>(*this, FnPos);
+  if (NoUnwindAA.isAssumedNoUnwind()) {
+    bool IsVoidTy = F->getReturnType()->isVoidTy();
+    const AAReturnedValues *RVAA =
+        IsVoidTy ? nullptr : &A.getAAFor<AAReturnedValues>(*this, FnPos);
+    if (IsVoidTy || CheckReturnedArgs(*RVAA)) {
+      T.addKnownBits(NOT_CAPTURED_IN_RET);
+      if (T.isKnown(NOT_CAPTURED_IN_MEM))
+        return ChangeStatus::UNCHANGED;
+      if (NoUnwindAA.isKnownNoUnwind() &&
+          (IsVoidTy || RVAA->getState().isAtFixpoint())) {
+        addKnownBits(NOT_CAPTURED_IN_RET);
+        if (isKnown(NOT_CAPTURED_IN_MEM))
+          return indicateOptimisticFixpoint();
+      }
+    }
+  }
 
   // Use the CaptureTracker interface and logic with the specialized tracker,
   // defined in AACaptureUseTracker, that can look at in-flight abstract
diff --git a/test/Transforms/FunctionAttrs/arg_nocapture.ll b/test/Transforms/FunctionAttrs/arg_nocapture.ll
index 18050dba201..7afaab6637e 100644
--- a/test/Transforms/FunctionAttrs/arg_nocapture.ll
+++ b/test/Transforms/FunctionAttrs/arg_nocapture.ll
@@ -405,35 +405,33 @@ entry:
 ;
 ; Make sure the returned flag on %r is strong enough to justify nocapture on %b but **not** on %r.
 ;
-; FIXME: The "returned" information is not propagated to the fullest extend causing us to miss "nocapture" on %b in the following:
-; CHECK: define i32* @not_captured_by_readonly_call_not_returned_either1(i32* readonly %b, i32* readonly returned %r)
+; CHECK: define i32* @not_captured_by_readonly_call_not_returned_either1(i32* nocapture readonly %b, i32* readonly returned %r)
 ;
-; CHECK: define i32* @not_captured_by_readonly_call_not_returned_either2(i32* readonly %b, i32* readonly returned %r)
-; CHECK: define i32* @not_captured_by_readonly_call_not_returned_either3(i32* readonly %b, i32* readonly returned %r)
+; CHECK: define i32* @not_captured_by_readonly_call_not_returned_either2(i32* nocapture readonly %b, i32* readonly returned %r)
+; CHECK: define i32* @not_captured_by_readonly_call_not_returned_either3(i32* nocapture readonly %b, i32* readonly returned %r)
 ;
-; FIXME: The "nounwind" information is not derived to the fullest extend causing us to miss "nocapture" on %b in the following:
-; CHECK: define i32* @not_captured_by_readonly_call_not_returned_either4(i32* readonly %b, i32* readonly returned %r)
-define i32* @not_captured_by_readonly_call_not_returned_either1(i32* %b, i32* returned %r) #0 {
+; CHECK: define i32* @not_captured_by_readonly_call_not_returned_either4(i32* nocapture readonly %b, i32* readonly returned %r)
+define i32* @not_captured_by_readonly_call_not_returned_either1(i32* %b, i32* returned %r) {
 entry:
   %call = call i32* @readonly_unknown(i32* %b, i32* %r) nounwind
   ret i32* %call
 }
 
 declare i32* @readonly_unknown_r1a(i32*, i32* returned) readonly
-define i32* @not_captured_by_readonly_call_not_returned_either2(i32* %b, i32* %r) #0 {
+define i32* @not_captured_by_readonly_call_not_returned_either2(i32* %b, i32* %r) {
 entry:
   %call = call i32* @readonly_unknown_r1a(i32* %b, i32* %r) nounwind
   ret i32* %call
 }
 
 declare i32* @readonly_unknown_r1b(i32*, i32* returned) readonly nounwind
-define i32* @not_captured_by_readonly_call_not_returned_either3(i32* %b, i32* %r) #0 {
+define i32* @not_captured_by_readonly_call_not_returned_either3(i32* %b, i32* %r) {
 entry:
   %call = call i32* @readonly_unknown_r1b(i32* %b, i32* %r)
   ret i32* %call
 }
 
-define i32* @not_captured_by_readonly_call_not_returned_either4(i32* %b, i32* %r) #0 {
+define i32* @not_captured_by_readonly_call_not_returned_either4(i32* %b, i32* %r) nounwind {
 entry:
   %call = call i32* @readonly_unknown_r1a(i32* %b, i32* %r)
   ret i32* %call
diff --git a/test/Transforms/FunctionAttrs/arg_returned.ll b/test/Transforms/FunctionAttrs/arg_returned.ll
index 3fe960c87e6..d927cdf7927 100644
--- a/test/Transforms/FunctionAttrs/arg_returned.ll
+++ b/test/Transforms/FunctionAttrs/arg_returned.ll
@@ -1,5 +1,5 @@
 ; RUN: opt -functionattrs -S < %s | FileCheck %s --check-prefix=FNATTR
-; RUN: opt -attributor -attributor-manifest-internal -attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=6 -S < %s | FileCheck %s --check-prefix=ATTRIBUTOR
+; RUN: opt -attributor -attributor-manifest-internal -attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=7 -S < %s | FileCheck %s --check-prefix=ATTRIBUTOR
 ; RUN: opt -attributor -attributor-manifest-internal -attributor-disable=false -functionattrs -S < %s | FileCheck %s --check-prefix=BOTH
 ;
 ; Test cases specifically designed for the "returned" argument attribute.
diff --git a/test/Transforms/FunctionAttrs/nocapture.ll b/test/Transforms/FunctionAttrs/nocapture.ll
index e0580849983..f8927a60cb1 100644
--- a/test/Transforms/FunctionAttrs/nocapture.ll
+++ b/test/Transforms/FunctionAttrs/nocapture.ll
@@ -39,6 +39,20 @@ l1:
 	ret i1 1 ; escaping value not caught by def-use chaining.
 }
 
+; c4b is c4 but without the escaping part
+; FNATTR: define i1 @c4b(i32* %q, i32 %bitno)
+; ATTRIBUTOR: define i1 @c4b(i32* nocapture readnone %q, i32 %bitno)
+define i1 @c4b(i32* %q, i32 %bitno) {
+	%tmp = ptrtoint i32* %q to i32
+	%tmp2 = lshr i32 %tmp, %bitno
+	%bit = trunc i32 %tmp2 to i1
+	br i1 %bit, label %l1, label %l0
+l0:
+	ret i1 0 ; not escaping!
+l1:
+	ret i1 0 ; not escaping!
+}
+
 @lookup_table = global [2 x i1] [ i1 0, i1 1 ]
 
 ; FNATTR: define i1 @c5(i32* %q, i32 %bitno)
@@ -331,5 +345,20 @@ entry:
   ret void
 }
 
+declare i8* @unknownpi8pi8(i8*,i8* returned)
+define i8* @test_returned1(i8* %A, i8* returned %B) nounwind readonly {
+; ATTRIBUTOR: define i8* @test_returned1(i8* nocapture readonly %A, i8* readonly returned %B)
+entry:
+  %p = call i8* @unknownpi8pi8(i8* %A, i8* %B)
+  ret i8* %p
+}
+
+define i8* @test_returned2(i8* %A, i8* %B) {
+; ATTRIBUTOR: define i8* @test_returned2(i8* nocapture readonly %A, i8* readonly returned %B)
+entry:
+  %p = call i8* @unknownpi8pi8(i8* %A, i8* %B) nounwind readonly
+  ret i8* %p
+}
+
 declare i8* @llvm.launder.invariant.group.p0i8(i8*)
 declare i8* @llvm.strip.invariant.group.p0i8(i8*)
diff --git a/test/Transforms/FunctionAttrs/nonnull.ll b/test/Transforms/FunctionAttrs/nonnull.ll
index 84d6720032f..9a7eb114eae 100644
--- a/test/Transforms/FunctionAttrs/nonnull.ll
+++ b/test/Transforms/FunctionAttrs/nonnull.ll
@@ -220,7 +220,7 @@ bb:
 
 define dso_local noalias i32* @f3(i32* %arg) {
 ; FIXME: missing nonnull. It should be nonnull @f3(i32* nonnull readonly %arg) 
-; ATTRIBUTOR: define dso_local noalias i32* @f3(i32* readonly %arg)
+; ATTRIBUTOR: define dso_local noalias i32* @f3(i32* nocapture readonly %arg)
 bb:
 ; FIXME: missing nonnull. It should be @f1(i32* nonnull readonly %arg) 
 ; ATTRIBUTOR:   %tmp = call i32* @f1(i32* readonly %arg)
diff --git a/test/Transforms/FunctionAttrs/read_write_returned_arguments_scc.ll b/test/Transforms/FunctionAttrs/read_write_returned_arguments_scc.ll
index 5f9e477679c..6fbc54502b1 100644
--- a/test/Transforms/FunctionAttrs/read_write_returned_arguments_scc.ll
+++ b/test/Transforms/FunctionAttrs/read_write_returned_arguments_scc.ll
@@ -1,4 +1,4 @@
-; RUN: opt -functionattrs -enable-nonnull-arg-prop -attributor -attributor-manifest-internal -attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=7 -S < %s | FileCheck %s
+; RUN: opt -functionattrs -enable-nonnull-arg-prop -attributor -attributor-manifest-internal -attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=8 -S < %s | FileCheck %s
 ;
 ; This is an evolved example to stress test SCC parameter attribute propagation.
 ; The SCC in this test is made up of the following six function, three of which