[AA] Take read-only provenance captures into account (#143097)

Update the AA CaptureAnalysis providers to return CaptureComponents, so
we can distinguish between full provenance and read-only provenance
captures.

Use this to restrict "other" memory effects on call from ModRef to Ref.

Ideally we would also apply the same reasoning for escape sources, but
the current API cannot actually convey the necessary information (we can
only say NoAlias or MayAlias, not MayAlias but only via Ref).

7 files changed, 89 insertions, 72 deletions

diff --git a/llvm/include/llvm/Analysis/AliasAnalysis.h b/llvm/include/llvm/Analysis/AliasAnalysis.h
index 0e736b9..b7d1251 100644
--- a/llvm/include/llvm/Analysis/AliasAnalysis.h
+++ b/llvm/include/llvm/Analysis/AliasAnalysis.h
@@ -149,23 +149,24 @@ LLVM_ABI raw_ostream &operator<<(raw_ostream &OS, AliasResult AR);
 struct LLVM_ABI CaptureAnalysis {
   virtual ~CaptureAnalysis() = 0;
 
-  /// Check whether Object is not captured before instruction I. If OrAt is
-  /// true, captures by instruction I itself are also considered.
+  /// Return how Object may be captured before instruction I, considering only
+  /// provenance captures. If OrAt is true, captures by instruction I itself
+  /// are also considered.
   ///
   /// If I is nullptr, then captures at any point will be considered.
-  virtual bool isNotCapturedBefore(const Value *Object, const Instruction *I,
-                                   bool OrAt) = 0;
+  virtual CaptureComponents
+  getCapturesBefore(const Value *Object, const Instruction *I, bool OrAt) = 0;
 };
 
 /// Context-free CaptureAnalysis provider, which computes and caches whether an
 /// object is captured in the function at all, but does not distinguish whether
 /// it was captured before or after the context instruction.
 class LLVM_ABI SimpleCaptureAnalysis final : public CaptureAnalysis {
-  SmallDenseMap<const Value *, bool, 8> IsCapturedCache;
+  SmallDenseMap<const Value *, CaptureComponents, 8> IsCapturedCache;
 
 public:
-  bool isNotCapturedBefore(const Value *Object, const Instruction *I,
-                           bool OrAt) override;
+  CaptureComponents getCapturesBefore(const Value *Object, const Instruction *I,
+                                      bool OrAt) override;
 };
 
 /// Context-sensitive CaptureAnalysis provider, which computes and caches the
@@ -176,10 +177,12 @@ class LLVM_ABI EarliestEscapeAnalysis final : public CaptureAnalysis {
   const LoopInfo *LI;
 
   /// Map from identified local object to an instruction before which it does
-  /// not escape, or nullptr if it never escapes. The "earliest" instruction
-  /// may be a conservative approximation, e.g. the first instruction in the
-  /// function is always a legal choice.
-  DenseMap<const Value *, Instruction *> EarliestEscapes;
+  /// not escape (or nullptr if it never escapes) and the possible components
+  /// that may be captured (by any instruction, not necessarily the earliest
+  /// one). The "earliest" instruction may be a conservative approximation,
+  /// e.g. the first instruction in the function is always a legal choice.
+  DenseMap<const Value *, std::pair<Instruction *, CaptureComponents>>
+      EarliestEscapes;
 
   /// Reverse map from instruction to the objects it is the earliest escape for.
   /// This is used for cache invalidation purposes.
@@ -189,8 +192,8 @@ public:
   EarliestEscapeAnalysis(DominatorTree &DT, const LoopInfo *LI = nullptr)
       : DT(DT), LI(LI) {}
 
-  bool isNotCapturedBefore(const Value *Object, const Instruction *I,
-                           bool OrAt) override;
+  CaptureComponents getCapturesBefore(const Value *Object, const Instruction *I,
+                                      bool OrAt) override;
 
   void removeInstruction(Instruction *I);
 };
diff --git a/llvm/include/llvm/Analysis/CaptureTracking.h b/llvm/include/llvm/Analysis/CaptureTracking.h
index dd6a7f9..e652bc5 100644
--- a/llvm/include/llvm/Analysis/CaptureTracking.h
+++ b/llvm/include/llvm/Analysis/CaptureTracking.h
@@ -95,21 +95,21 @@ namespace llvm {
       function_ref<bool(CaptureComponents)> StopFn = capturesAnything,
       const LoopInfo *LI = nullptr, unsigned MaxUsesToExplore = 0);
 
-  // Returns the 'earliest' instruction that captures \p V in \F. An instruction
-  // A is considered earlier than instruction B, if A dominates B. If 2 escapes
-  // do not dominate each other, the terminator of the common dominator is
-  // chosen. If not all uses can be analyzed, the earliest escape is set to
-  // the first instruction in the function entry block. If \p V does not escape,
-  // nullptr is returned. Note that the caller of the function has to ensure
-  // that the instruction the result value is compared against is not in a
-  // cycle.
+  // Returns the 'earliest' instruction that captures \p V in \F, and which
+  // components may be captured (by any use, not necessarily the earliest one).
+  // An instruction A is considered earlier than instruction B, if A dominates
+  // B. If 2 escapes do not dominate each other, the terminator of the common
+  // dominator is chosen. If not all uses can be analyzed, the earliest escape
+  // is set to the first instruction in the function entry block. If \p V does
+  // not escape, nullptr is returned. Note that the caller of the function has
+  // to ensure that the instruction the result value is compared against is
+  // not in a cycle.
   //
   // Only consider components that are part of \p Mask.
-  LLVM_ABI Instruction *FindEarliestCapture(const Value *V, Function &F,
-                                            bool ReturnCaptures,
-                                            const DominatorTree &DT,
-                                            CaptureComponents Mask,
-                                            unsigned MaxUsesToExplore = 0);
+  LLVM_ABI std::pair<Instruction *, CaptureComponents>
+  FindEarliestCapture(const Value *V, Function &F, bool ReturnCaptures,
+                      const DominatorTree &DT, CaptureComponents Mask,
+                      unsigned MaxUsesToExplore = 0);
 
   /// Capture information for a specific Use.
   struct UseCaptureInfo {
diff --git a/llvm/lib/Analysis/BasicAliasAnalysis.cpp b/llvm/lib/Analysis/BasicAliasAnalysis.cpp
index f862d69..31611df 100644
--- a/llvm/lib/Analysis/BasicAliasAnalysis.cpp
+++ b/llvm/lib/Analysis/BasicAliasAnalysis.cpp
@@ -192,18 +192,20 @@ static bool areBothVScale(const Value *V1, const Value *V2) {
 
 CaptureAnalysis::~CaptureAnalysis() = default;
 
-bool SimpleCaptureAnalysis::isNotCapturedBefore(const Value *Object,
-                                                const Instruction *I,
-                                                bool OrAt) {
+CaptureComponents SimpleCaptureAnalysis::getCapturesBefore(const Value *Object,
+                                                           const Instruction *I,
+                                                           bool OrAt) {
   if (!isIdentifiedFunctionLocal(Object))
-    return false;
+    return CaptureComponents::Provenance;
 
-  auto [CacheIt, Inserted] = IsCapturedCache.insert({Object, false});
+  auto [CacheIt, Inserted] =
+      IsCapturedCache.insert({Object, CaptureComponents::Provenance});
   if (!Inserted)
     return CacheIt->second;
 
-  bool Ret = !capturesAnything(PointerMayBeCaptured(
-      Object, /*ReturnCaptures=*/false, CaptureComponents::Provenance));
+  CaptureComponents Ret = PointerMayBeCaptured(
+      Object, /*ReturnCaptures=*/false, CaptureComponents::Provenance,
+      [](CaptureComponents CC) { return capturesFullProvenance(CC); });
   CacheIt->second = Ret;
   return Ret;
 }
@@ -216,37 +218,44 @@ static bool isNotInCycle(const Instruction *I, const DominatorTree *DT,
          !isPotentiallyReachableFromMany(Succs, BB, nullptr, DT, LI);
 }
 
-bool EarliestEscapeAnalysis::isNotCapturedBefore(const Value *Object,
-                                                 const Instruction *I,
-                                                 bool OrAt) {
+CaptureComponents
+EarliestEscapeAnalysis::getCapturesBefore(const Value *Object,
+                                          const Instruction *I, bool OrAt) {
   if (!isIdentifiedFunctionLocal(Object))
-    return false;
+    return CaptureComponents::Provenance;
 
   auto Iter = EarliestEscapes.try_emplace(Object);
   if (Iter.second) {
-    Instruction *EarliestCapture = FindEarliestCapture(
-        Object, *const_cast<Function *>(DT.getRoot()->getParent()),
-        /*ReturnCaptures=*/false, DT, CaptureComponents::Provenance);
-    if (EarliestCapture)
-      Inst2Obj[EarliestCapture].push_back(Object);
+    std::pair<Instruction *, CaptureComponents> EarliestCapture =
+        FindEarliestCapture(
+            Object, *const_cast<Function *>(DT.getRoot()->getParent()),
+            /*ReturnCaptures=*/false, DT, CaptureComponents::Provenance);
+    if (EarliestCapture.first)
+      Inst2Obj[EarliestCapture.first].push_back(Object);
     Iter.first->second = EarliestCapture;
   }
 
-  // No capturing instruction.
-  if (!Iter.first->second)
-    return true;
-
-  // No context instruction means any use is capturing.
-  if (!I)
-    return false;
+  auto IsNotCapturedBefore = [&]() {
+    // No capturing instruction.
+    Instruction *CaptureInst = Iter.first->second.first;
+    if (!CaptureInst)
+      return true;
 
-  if (I == Iter.first->second) {
-    if (OrAt)
+    // No context instruction means any use is capturing.
+    if (!I)
       return false;
-    return isNotInCycle(I, &DT, LI);
-  }
 
-  return !isPotentiallyReachable(Iter.first->second, I, nullptr, &DT, LI);
+    if (I == CaptureInst) {
+      if (OrAt)
+        return false;
+      return isNotInCycle(I, &DT, LI);
+    }
+
+    return !isPotentiallyReachable(CaptureInst, I, nullptr, &DT, LI);
+  };
+  if (IsNotCapturedBefore())
+    return CaptureComponents::None;
+  return Iter.first->second.second;
 }
 
 void EarliestEscapeAnalysis::removeInstruction(Instruction *I) {
@@ -946,9 +955,14 @@ ModRefInfo BasicAAResult::getModRefInfo(const CallBase *Call,
   // As an exception, ignore allocas, as setjmp is not required to preserve
   // non-volatile stores for them.
   if (isModOrRefSet(OtherMR) && !isa<Constant>(Object) && Call != Object &&
-      AAQI.CA->isNotCapturedBefore(Object, Call, /*OrAt=*/false) &&
-      (isa<AllocaInst>(Object) || !Call->hasFnAttr(Attribute::ReturnsTwice)))
-    OtherMR = ModRefInfo::NoModRef;
+      (isa<AllocaInst>(Object) || !Call->hasFnAttr(Attribute::ReturnsTwice))) {
+    CaptureComponents CC =
+        AAQI.CA->getCapturesBefore(Object, Call, /*OrAt=*/false);
+    if (capturesNothing(CC))
+      OtherMR = ModRefInfo::NoModRef;
+    else if (capturesReadProvenanceOnly(CC))
+      OtherMR = ModRefInfo::Ref;
+  }
 
   // Refine the modref info for argument memory. We only bother to do this
   // if ArgMR is not a subset of OtherMR, otherwise this won't have an impact
@@ -1614,11 +1628,13 @@ 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))
+    if (isEscapeSource(O1) &&
+        capturesNothing(AAQI.CA->getCapturesBefore(
+            O2, dyn_cast<Instruction>(O1), /*OrAt*/ true)))
       return AliasResult::NoAlias;
-    if (isEscapeSource(O2) && AAQI.CA->isNotCapturedBefore(
-                                  O1, dyn_cast<Instruction>(O2), /*OrAt*/ true))
+    if (isEscapeSource(O2) &&
+        capturesNothing(AAQI.CA->getCapturesBefore(
+            O1, dyn_cast<Instruction>(O2), /*OrAt*/ true)))
       return AliasResult::NoAlias;
   }
 
diff --git a/llvm/lib/Analysis/CaptureTracking.cpp b/llvm/lib/Analysis/CaptureTracking.cpp
index d08ed17..076f417 100644
--- a/llvm/lib/Analysis/CaptureTracking.cpp
+++ b/llvm/lib/Analysis/CaptureTracking.cpp
@@ -249,11 +249,10 @@ bool llvm::PointerMayBeCapturedBefore(const Value *V, bool ReturnCaptures,
       capturesAnything, LI, MaxUsesToExplore));
 }
 
-Instruction *llvm::FindEarliestCapture(const Value *V, Function &F,
-                                       bool ReturnCaptures,
-                                       const DominatorTree &DT,
-                                       CaptureComponents Mask,
-                                       unsigned MaxUsesToExplore) {
+std::pair<Instruction *, CaptureComponents>
+llvm::FindEarliestCapture(const Value *V, Function &F, bool ReturnCaptures,
+                          const DominatorTree &DT, CaptureComponents Mask,
+                          unsigned MaxUsesToExplore) {
   assert(!isa<GlobalValue>(V) &&
          "It doesn't make sense to ask whether a global is captured.");
 
@@ -263,7 +262,7 @@ Instruction *llvm::FindEarliestCapture(const Value *V, Function &F,
     ++NumCapturedBefore;
   else
     ++NumNotCapturedBefore;
-  return CB.EarliestCapture;
+  return {CB.EarliestCapture, CB.CC};
 }
 
 UseCaptureInfo llvm::DetermineUseCaptureKind(const Use &U, const Value *Base) {
diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index 49a0c88..4a2eb92 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -2345,7 +2345,8 @@ bool isFuncLocalAndNotCaptured(Value *Arg, const CallBase *CB,
                                EarliestEscapeAnalysis &EA) {
   const Value *UnderlyingObj = getUnderlyingObject(Arg);
   return isIdentifiedFunctionLocal(UnderlyingObj) &&
-         EA.isNotCapturedBefore(UnderlyingObj, CB, /*OrAt*/ true);
+         capturesNothing(
+             EA.getCapturesBefore(UnderlyingObj, CB, /*OrAt*/ true));
 }
 
 SmallVector<MemoryLocation, 1>
diff --git a/llvm/test/Analysis/BasicAA/captures.ll b/llvm/test/Analysis/BasicAA/captures.ll
index c212a46..c9ed1ea 100644
--- a/llvm/test/Analysis/BasicAA/captures.ll
+++ b/llvm/test/Analysis/BasicAA/captures.ll
@@ -17,8 +17,7 @@ define void @address_capture() {
 
 ; CHECK-LABEL: read_only_capture
 ; CHECK: MayAlias:	i32* %a, i32* %p
-; CHECK: Both ModRef:  Ptr: i32* %a	<->  %p = call ptr @get_ptr()
-; TODO: The ModRef could be just Ref.
+; CHECK: Just Ref:  Ptr: i32* %a	<->  %p = call ptr @get_ptr()
 define void @read_only_capture() {
   %a = alloca i32
   call void @capture(ptr captures(address, read_provenance) %a)
diff --git a/llvm/test/Transforms/GVN/captures.ll b/llvm/test/Transforms/GVN/captures.ll
index ae47e92..96fce43 100644
--- a/llvm/test/Transforms/GVN/captures.ll
+++ b/llvm/test/Transforms/GVN/captures.ll
@@ -43,8 +43,7 @@ define i32 @read_provenance_capture() {
 ; CHECK-NEXT:    call void @capture(ptr captures(address, read_provenance) [[A]])
 ; CHECK-NEXT:    store i32 1, ptr [[A]], align 4
 ; CHECK-NEXT:    call void @unknown_call()
-; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[A]], align 4
-; CHECK-NEXT:    ret i32 [[V]]
+; CHECK-NEXT:    ret i32 1
 ;
   %a = alloca i32
   call void @capture(ptr captures(address, read_provenance) %a)