[CaptureTracking][AA] Only consider provenance captures (llvm#130777)

For the purposes of alias analysis, we should only consider provenance
captures, not address captures. To support this, change (or add)
CaptureTracking APIs to accept a Mask and StopFn argument. The Mask
determines which components we are interested in (for AA that would be
Provenance).

The StopFn determines when we can abort the walk early. Currently, we
want to do this as soon as any of the components in the Mask is
captured. The purpose of making this a separate predicate is that in the
future we will also want to distinguish between capturing full
provenance and read-only provenance. In that case, we can only stop
early once full provenance is captured. The earliest escape analysis
does not get a StopFn, because it must always inspect all captures.

Author: nikic

llvm/include/llvm/Analysis/CaptureTracking.h

@@ -44,6 +44,15 @@ namespace llvm {
   bool PointerMayBeCaptured(const Value *V, bool ReturnCaptures,
                             unsigned MaxUsesToExplore = 0);
 
+  /// Return which components of the pointer may be captured. Only consider
+  /// components that are part of \p Mask. Once \p StopFn on the accumulated
+  /// components returns true, the traversal is aborted early. By default, this
+  /// happens when *any* of the components in \p Mask are captured.
+  CaptureComponents PointerMayBeCaptured(
+      const Value *V, bool ReturnCaptures, CaptureComponents Mask,
+      function_ref<bool(CaptureComponents)> StopFn = capturesAnything,
+      unsigned MaxUsesToExplore = 0);
+
   /// PointerMayBeCapturedBefore - Return true if this pointer value may be
   /// captured by the enclosing function (which is required to exist). If a
   /// DominatorTree is provided, only captures which happen before the given
@@ -61,6 +70,17 @@ namespace llvm {
                                   unsigned MaxUsesToExplore = 0,
                                   const LoopInfo *LI = nullptr);
 
+  /// Return which components of the pointer may be captured on the path to
+  /// \p I. Only consider components that are part of \p Mask. Once \p StopFn
+  /// on the accumulated components returns true, the traversal is aborted
+  /// early. By default, this happens when *any* of the components in \p Mask
+  /// are captured.
+  CaptureComponents PointerMayBeCapturedBefore(
+      const Value *V, bool ReturnCaptures, const Instruction *I,
+      const DominatorTree *DT, bool IncludeI, CaptureComponents Mask,
+      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
@@ -69,8 +89,11 @@ namespace llvm {
   // 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.
   Instruction *FindEarliestCapture(const Value *V, Function &F,
                                    bool ReturnCaptures, const DominatorTree &DT,
+                                   CaptureComponents Mask,
                                    unsigned MaxUsesToExplore = 0);
 
   /// Capture information for a specific Use.

llvm/include/llvm/Support/ModRef.h

@@ -326,6 +326,10 @@ inline bool capturesFullProvenance(CaptureComponents CC) {
   return (CC & CaptureComponents::Provenance) == CaptureComponents::Provenance;
 }
 
+inline bool capturesAnyProvenance(CaptureComponents CC) {
+  return (CC & CaptureComponents::Provenance) != CaptureComponents::None;
+}
+
 inline bool capturesAll(CaptureComponents CC) {
   return CC == CaptureComponents::All;
 }

llvm/lib/Analysis/AliasAnalysis.cpp

@@ -622,8 +622,9 @@ ModRefInfo AAResults::callCapturesBefore(const Instruction *I,
   if (!Call || Call == Object)
     return ModRefInfo::ModRef;
 
-  if (PointerMayBeCapturedBefore(Object, /* ReturnCaptures */ true, I, DT,
-                                 /* include Object */ true))
+  if (capturesAnything(PointerMayBeCapturedBefore(
+          Object, /* ReturnCaptures */ true, I, DT,
+          /* include Object */ true, CaptureComponents::Provenance)))
     return ModRefInfo::ModRef;
 
   unsigned ArgNo = 0;
@@ -637,10 +638,11 @@ ModRefInfo AAResults::callCapturesBefore(const Instruction *I,
     if (!(*CI)->getType()->isPointerTy())
       continue;
 
-    // Make sure we still check captures(ret: address, provenance) arguments,
-    // as these wouldn't be treated as a capture at the call-site.
+    // Make sure we still check captures(ret: address, provenance) and
+    // captures(address) arguments, as these wouldn't be treated as a capture
+    // at the call-site.
     CaptureInfo Captures = Call->getCaptureInfo(ArgNo);
-    if (!capturesNothing(Captures.getOtherComponents()))
+    if (capturesAnyProvenance(Captures.getOtherComponents()))
       continue;
 
     AliasResult AR =

llvm/lib/Analysis/BasicAliasAnalysis.cpp

@@ -220,7 +220,7 @@ bool EarliestEscapeAnalysis::isNotCapturedBefore(const Value *Object,
   if (Iter.second) {
     Instruction *EarliestCapture = FindEarliestCapture(
         Object, *const_cast<Function *>(DT.getRoot()->getParent()),
-        /*ReturnCaptures=*/false, DT);
+        /*ReturnCaptures=*/false, DT, CaptureComponents::Provenance);
     if (EarliestCapture)
       Inst2Obj[EarliestCapture].push_back(Object);
     Iter.first->second = EarliestCapture;

llvm/lib/Analysis/CaptureTracking.cpp

@@ -73,28 +73,32 @@ bool CaptureTracker::isDereferenceableOrNull(Value *O, const DataLayout &DL) {
 
 namespace {
 struct SimpleCaptureTracker : public CaptureTracker {
-  explicit SimpleCaptureTracker(bool ReturnCaptures)
-      : ReturnCaptures(ReturnCaptures) {}
+  explicit SimpleCaptureTracker(bool ReturnCaptures, CaptureComponents Mask,
+                                function_ref<bool(CaptureComponents)> StopFn)
+      : ReturnCaptures(ReturnCaptures), Mask(Mask), StopFn(StopFn) {}
 
   void tooManyUses() override {
     LLVM_DEBUG(dbgs() << "Captured due to too many uses\n");
-    Captured = true;
+    CC = Mask;
   }
 
   Action captured(const Use *U, UseCaptureInfo CI) override {
-    // TODO(captures): Use UseCaptureInfo.
     if (isa<ReturnInst>(U->getUser()) && !ReturnCaptures)
       return ContinueIgnoringReturn;
 
-    LLVM_DEBUG(dbgs() << "Captured by: " << *U->getUser() << "\n");
+    if (capturesNothing(CI.UseCC & Mask))
+      return Continue;
 
-    Captured = true;
-    return Stop;
+    LLVM_DEBUG(dbgs() << "Captured by: " << *U->getUser() << "\n");
+    CC |= CI.UseCC & Mask;
+    return StopFn(CC) ? Stop : Continue;
   }
 
   bool ReturnCaptures;
+  CaptureComponents Mask;
+  function_ref<bool(CaptureComponents)> StopFn;
 
-  bool Captured = false;
+  CaptureComponents CC = CaptureComponents::None;
 };
 
 /// Only find pointer captures which happen before the given instruction. Uses
@@ -104,11 +108,13 @@ struct SimpleCaptureTracker : public CaptureTracker {
 struct CapturesBefore : public CaptureTracker {
 
   CapturesBefore(bool ReturnCaptures, const Instruction *I,
-                 const DominatorTree *DT, bool IncludeI, const LoopInfo *LI)
+                 const DominatorTree *DT, bool IncludeI, const LoopInfo *LI,
+                 CaptureComponents Mask,
+                 function_ref<bool(CaptureComponents)> StopFn)
       : BeforeHere(I), DT(DT), ReturnCaptures(ReturnCaptures),
-        IncludeI(IncludeI), LI(LI) {}
+        IncludeI(IncludeI), LI(LI), Mask(Mask), StopFn(StopFn) {}
 
-  void tooManyUses() override { Captured = true; }
+  void tooManyUses() override { CC = Mask; }
 
   bool isSafeToPrune(Instruction *I) {
     if (BeforeHere == I)
@@ -124,7 +130,6 @@ struct CapturesBefore : public CaptureTracker {
   }
 
   Action captured(const Use *U, UseCaptureInfo CI) override {
-    // TODO(captures): Use UseCaptureInfo.
     Instruction *I = cast<Instruction>(U->getUser());
     if (isa<ReturnInst>(I) && !ReturnCaptures)
       return ContinueIgnoringReturn;
@@ -136,8 +141,11 @@ struct CapturesBefore : public CaptureTracker {
       // If the use is not reachable, the instruction result isn't either.
       return ContinueIgnoringReturn;
 
-    Captured = true;
-    return Stop;
+    if (capturesNothing(CI.UseCC & Mask))
+      return Continue;
+
+    CC |= CI.UseCC & Mask;
+    return StopFn(CC) ? Stop : Continue;
   }
 
   const Instruction *BeforeHere;
@@ -146,9 +154,11 @@ struct CapturesBefore : public CaptureTracker {
   bool ReturnCaptures;
   bool IncludeI;
 
-  bool Captured = false;
+  CaptureComponents CC = CaptureComponents::None;
 
   const LoopInfo *LI;
+  CaptureComponents Mask;
+  function_ref<bool(CaptureComponents)> StopFn;
 };
 
 /// Find the 'earliest' instruction before which the pointer is known not to
@@ -161,104 +171,110 @@ struct CapturesBefore : public CaptureTracker {
 // escape are not in a cycle.
 struct EarliestCaptures : public CaptureTracker {
 
-  EarliestCaptures(bool ReturnCaptures, Function &F, const DominatorTree &DT)
-      : DT(DT), ReturnCaptures(ReturnCaptures), F(F) {}
+  EarliestCaptures(bool ReturnCaptures, Function &F, const DominatorTree &DT,
+                   CaptureComponents Mask)
+      : DT(DT), ReturnCaptures(ReturnCaptures), F(F), Mask(Mask) {}
 
   void tooManyUses() override {
-    Captured = true;
+    CC = Mask;
     EarliestCapture = &*F.getEntryBlock().begin();
   }
 
   Action captured(const Use *U, UseCaptureInfo CI) override {
-    // TODO(captures): Use UseCaptureInfo.
     Instruction *I = cast<Instruction>(U->getUser());
     if (isa<ReturnInst>(I) && !ReturnCaptures)
       return ContinueIgnoringReturn;
 
-    if (!EarliestCapture)
-      EarliestCapture = I;
-    else
-      EarliestCapture = DT.findNearestCommonDominator(EarliestCapture, I);
-    Captured = true;
+    if (capturesAnything(CI.UseCC & Mask)) {
+      if (!EarliestCapture)
+        EarliestCapture = I;
+      else
+        EarliestCapture = DT.findNearestCommonDominator(EarliestCapture, I);
+      CC |= CI.UseCC & Mask;
+    }
 
-    // Continue analysis, as we need to see all potential captures. However,
-    // we do not need to follow the instruction result, as this use will
-    // dominate any captures made through the instruction result.
-    return ContinueIgnoringReturn;
+    // Continue analysis, as we need to see all potential captures.
+    return Continue;
   }
 
-  Instruction *EarliestCapture = nullptr;
-
   const DominatorTree &DT;
-
   bool ReturnCaptures;
-
-  bool Captured = false;
-
   Function &F;
+  CaptureComponents Mask;
+
+  Instruction *EarliestCapture = nullptr;
+  CaptureComponents CC = CaptureComponents::None;
 };
 } // namespace
 
-/// PointerMayBeCaptured - Return true if this pointer value may be captured
-/// by the enclosing function (which is required to exist).  This routine can
-/// be expensive, so consider caching the results.  The boolean ReturnCaptures
-/// specifies whether returning the value (or part of it) from the function
-/// counts as capturing it or not.
-bool llvm::PointerMayBeCaptured(const Value *V, bool ReturnCaptures,
-                                unsigned MaxUsesToExplore) {
+CaptureComponents llvm::PointerMayBeCaptured(
+    const Value *V, bool ReturnCaptures, CaptureComponents Mask,
+    function_ref<bool(CaptureComponents)> StopFn, unsigned MaxUsesToExplore) {
   assert(!isa<GlobalValue>(V) &&
          "It doesn't make sense to ask whether a global is captured.");
 
   LLVM_DEBUG(dbgs() << "Captured?: " << *V << " = ");
 
-  SimpleCaptureTracker SCT(ReturnCaptures);
+  SimpleCaptureTracker SCT(ReturnCaptures, Mask, StopFn);
   PointerMayBeCaptured(V, &SCT, MaxUsesToExplore);
-  if (SCT.Captured)
+  if (capturesAnything(SCT.CC))
     ++NumCaptured;
   else {
     ++NumNotCaptured;
     LLVM_DEBUG(dbgs() << "not captured\n");
   }
-  return SCT.Captured;
+  return SCT.CC;
 }
 
-/// PointerMayBeCapturedBefore - Return true if this pointer value may be
-/// captured by the enclosing function (which is required to exist). If a
-/// DominatorTree is provided, only captures which happen before the given
-/// instruction are considered. This routine can be expensive, so consider
-/// caching the results.  The boolean ReturnCaptures specifies whether
-/// returning the value (or part of it) from the function counts as capturing
-/// it or not.
-bool llvm::PointerMayBeCapturedBefore(const Value *V, bool ReturnCaptures,
-                                      const Instruction *I,
-                                      const DominatorTree *DT, bool IncludeI,
-                                      unsigned MaxUsesToExplore,
-                                      const LoopInfo *LI) {
+bool llvm::PointerMayBeCaptured(const Value *V, bool ReturnCaptures,
+                                unsigned MaxUsesToExplore) {
+  return capturesAnything(
+      PointerMayBeCaptured(V, ReturnCaptures, CaptureComponents::All,
+                           capturesAnything, MaxUsesToExplore));
+}
+
+CaptureComponents llvm::PointerMayBeCapturedBefore(
+    const Value *V, bool ReturnCaptures, const Instruction *I,
+    const DominatorTree *DT, bool IncludeI, CaptureComponents Mask,
+    function_ref<bool(CaptureComponents)> StopFn, const LoopInfo *LI,
+    unsigned MaxUsesToExplore) {
   assert(!isa<GlobalValue>(V) &&
          "It doesn't make sense to ask whether a global is captured.");
 
   if (!DT)
-    return PointerMayBeCaptured(V, ReturnCaptures, MaxUsesToExplore);
+    return PointerMayBeCaptured(V, ReturnCaptures, Mask, StopFn,
+                                MaxUsesToExplore);
 
-  CapturesBefore CB(ReturnCaptures, I, DT, IncludeI, LI);
+  CapturesBefore CB(ReturnCaptures, I, DT, IncludeI, LI, Mask, StopFn);
   PointerMayBeCaptured(V, &CB, MaxUsesToExplore);
-  if (CB.Captured)
+  if (capturesAnything(CB.CC))
     ++NumCapturedBefore;
   else
     ++NumNotCapturedBefore;
-  return CB.Captured;
+  return CB.CC;
+}
+
+bool llvm::PointerMayBeCapturedBefore(const Value *V, bool ReturnCaptures,
+                                      const Instruction *I,
+                                      const DominatorTree *DT, bool IncludeI,
+                                      unsigned MaxUsesToExplore,
+                                      const LoopInfo *LI) {
+  return capturesAnything(PointerMayBeCapturedBefore(
+      V, ReturnCaptures, I, DT, IncludeI, CaptureComponents::All,
+      capturesAnything, LI, MaxUsesToExplore));
 }
 
 Instruction *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.");
 
-  EarliestCaptures CB(ReturnCaptures, F, DT);
+  EarliestCaptures CB(ReturnCaptures, F, DT, Mask);
   PointerMayBeCaptured(V, &CB, MaxUsesToExplore);
-  if (CB.Captured)
+  if (capturesAnything(CB.CC))
     ++NumCapturedBefore;
   else
     ++NumNotCapturedBefore;
@@ -473,8 +489,10 @@ bool llvm::isNonEscapingLocalObject(
   }
 
   // If this is an identified function-local object, check to see if it escapes.
+  // We only care about provenance here, not address capture.
   if (isIdentifiedFunctionLocal(V)) {
-    auto Ret = !PointerMayBeCaptured(V, /*ReturnCaptures=*/false);
+    bool Ret = !capturesAnything(PointerMayBeCaptured(
+        V, /*ReturnCaptures=*/false, CaptureComponents::Provenance));
     if (IsCapturedCache)
       CacheIt->second = Ret;
     return Ret;

llvm/test/Analysis/BasicAA/captures.ll

@@ -0,0 +1,42 @@
+; RUN: opt < %s -passes=aa-eval -print-all-alias-modref-info -disable-output 2>&1 | FileCheck %s
+
+declare void @capture(ptr)
+declare ptr @get_ptr()
+
+; CHECK-LABEL: address_capture
+; CHECK: NoAlias:	i32* %a, i32* %p
+; CHECK: NoModRef:  Ptr: i32* %a	<->  %p = call ptr @get_ptr()
+define void @address_capture() {
+  %a = alloca i32
+  call void @capture(ptr captures(address) %a)
+  %p = call ptr @get_ptr()
+  store i32 0, ptr %p
+  load i32, ptr %a
+  ret void
+}
+
+; 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.
+define void @read_only_capture() {
+  %a = alloca i32
+  call void @capture(ptr captures(address, read_provenance) %a)
+  %p = call ptr @get_ptr()
+  store i32 0, ptr %p
+  load i32, ptr %a
+  ret void
+}
+
+; CHECK-LABEL: address_capture_and_full_capture
+; CHECK: MayAlias:	i32* %a, i32* %p
+; CHECK: Both ModRef:  Ptr: i32* %a	<->  %p = call ptr @get_ptr()
+define void @address_capture_and_full_capture() {
+  %a = alloca i32
+  call void @capture(ptr captures(address) %a)
+  call void @capture(ptr %a)
+  %p = call ptr @get_ptr()
+  store i32 0, ptr %p
+  load i32, ptr %a
+  ret void
+}