[BasicAA] Improve escape source analysis for return-only captures

This fixes an issue where improving inference from a general capture
to a return-only capture can make alias analysis results worse.

The problem is that if there are return-only captures, the call
can no longer treated as a plain escape source, as the call result
may alias one of the return-only-capture arguments, even if they
don't escape.

Fix this by having isEscapeSource() return the relevant arguments and
explicitly check that these aren't aliased.

Fixes #131168.

Author: nikic

llvm/include/llvm/Analysis/AliasAnalysis.h

@@ -888,9 +888,9 @@ bool isIdentifiedFunctionLocal(const Value *V);
 /// can be inbounds w.r.t the actual underlying object.
 bool isBaseOfObject(const Value *V);
 
-/// Returns true if the pointer is one which would have been considered an
-/// escape by isNonEscapingLocalObject.
-bool isEscapeSource(const Value *V);
+/// Returns true if the pointer cannot alias a non-escaping local object,
+/// except via the values returned in the MayAlias argument.
+bool isEscapeSource(const Value *V, SmallVectorImpl<Value *> &MayAlias);
 
 /// Return true if Object memory is not visible after an unwind, in the sense
 /// that program semantics cannot depend on Object containing any particular

llvm/include/llvm/IR/InstrTypes.h

@@ -1692,10 +1692,11 @@ class CallBase : public Instruction {
     return capturesNothing(getCaptureInfo(OpNo));
   }
 
-  /// Returns whether the call has an argument that has an attribute like
+  /// Collects call arguments that have an attribute like
   /// captures(ret: address, provenance), where the return capture components
   /// are not a subset of the other capture components.
-  bool hasArgumentWithAdditionalReturnCaptureComponents() const;
+  void getArgumentsWithAdditionalReturnCaptureComponents(
+      SmallVectorImpl<Value *> &Args) const;
 
   /// Determine whether this argument is passed by value.
   bool isByValArgument(unsigned ArgNo) const {

llvm/lib/Analysis/AliasAnalysis.cpp

@@ -840,15 +840,16 @@ bool llvm::isBaseOfObject(const Value *V) {
   return (isa<AllocaInst>(V) || isa<GlobalVariable>(V));
 }
 
-bool llvm::isEscapeSource(const Value *V) {
+bool llvm::isEscapeSource(const Value *V, SmallVectorImpl<Value *> &MayAlias) {
   if (auto *CB = dyn_cast<CallBase>(V)) {
     if (isIntrinsicReturningPointerAliasingArgumentWithoutCapturing(CB, true))
       return false;
 
     // The return value of a function with a captures(ret: address, provenance)
-    // attribute is not necessarily an escape source. The return value may
-    // alias with a non-escaping object.
-    return !CB->hasArgumentWithAdditionalReturnCaptureComponents();
+    // attribute may alias with the corresponding argument. If it doesn't, then
+    // it is an escape source.
+    CB->getArgumentsWithAdditionalReturnCaptureComponents(MayAlias);
+    return true;
   }
 
   // The load case works because isNonEscapingLocalObject considers all

llvm/lib/Analysis/BasicAliasAnalysis.cpp

@@ -1606,11 +1606,21 @@ AliasResult BasicAAResult::aliasCheck(const Value *V1, LocationSize V1Size,
     // temporary store the nocapture argument's value in a temporary memory
     // location if that memory location doesn't escape. Or it may pass a
     // nocapture value to other functions as long as they don't capture it.
-    if (isEscapeSource(O1) && AAQI.CA->isNotCapturedBefore(
-                                  O2, dyn_cast<Instruction>(O1), /*OrAt*/ true))
-      return AliasResult::NoAlias;
-    if (isEscapeSource(O2) && AAQI.CA->isNotCapturedBefore(
-                                  O1, dyn_cast<Instruction>(O2), /*OrAt*/ true))
+    auto EscapeSourceDoesNotAliasNonEscapingLocalObject =
+        [&](const Value *O1, const Value *O2) {
+          SmallVector<Value *> MayAlias;
+          if (!isEscapeSource(O1, MayAlias) ||
+              !AAQI.CA->isNotCapturedBefore(O2, dyn_cast<Instruction>(O1),
+                                            /*OrAt=*/true))
+            return false;
+          for (Value *V : MayAlias)
+            if (AAQI.AAR.alias(MemoryLocation::getBeforeOrAfter(O2),
+                               MemoryLocation::getBeforeOrAfter(V), AAQI))
+              return false;
+          return true;
+        };
+    if (EscapeSourceDoesNotAliasNonEscapingLocalObject(O1, O2) ||
+        EscapeSourceDoesNotAliasNonEscapingLocalObject(O2, O1))
       return AliasResult::NoAlias;
   }
 

llvm/lib/IR/Instructions.cpp

@@ -713,7 +713,8 @@ CaptureInfo CallBase::getCaptureInfo(unsigned OpNo) const {
   return OBU.isDeoptOperandBundle() ? CaptureInfo::none() : CaptureInfo::all();
 }
 
-bool CallBase::hasArgumentWithAdditionalReturnCaptureComponents() const {
+void CallBase::getArgumentsWithAdditionalReturnCaptureComponents(
+    SmallVectorImpl<Value *> &Args) const {
   for (unsigned I = 0, E = arg_size(); I < E; ++I) {
     if (!getArgOperand(I)->getType()->isPointerTy())
       continue;
@@ -722,9 +723,8 @@ bool CallBase::hasArgumentWithAdditionalReturnCaptureComponents() const {
     if (auto *Fn = dyn_cast<Function>(getCalledOperand()))
       CI &= Fn->getAttributes().getParamAttrs(I).getCaptureInfo();
     if (capturesAnything(CI.getRetComponents() & ~CI.getOtherComponents()))
-      return true;
+      Args.push_back(getArgOperand(I));
   }
-  return false;
 }
 
 //===----------------------------------------------------------------------===//

llvm/lib/Transforms/Utils/InlineFunction.cpp

@@ -1273,7 +1273,8 @@ static void AddAliasScopeMetadata(CallBase &CB, ValueToValueMapTy &VMap,
           UsesAliasingPtr = true;
         }
 
-        if (isEscapeSource(V)) {
+        SmallVector<Value *> MayAlias;
+        if (isEscapeSource(V, MayAlias) && MayAlias.empty()) {
           // An escape source can only alias with a noalias argument if it has
           // been captured beforehand.
           RequiresNoCaptureBefore = true;

llvm/test/Analysis/BasicAA/captures.ll

@@ -45,7 +45,7 @@ declare ptr @capture_ret(ptr, ptr)
 
 ; CHECK-LABEL: capture_ret_only
 ; CHECK: MayAlias:	i8* %a, i8* %ret
-; CHECK: MayAlias:	i8* %b, i8* %ret
+; CHECK: NoAlias:	i8* %b, i8* %ret
 define void @capture_ret_only(ptr noalias %a, ptr noalias %b) {
   %ret = call ptr @capture_ret(ptr captures(ret: address, provenance) %a, ptr captures(none) %b)
   load i8, ptr %ret