Introduce MaxUsesToExplore argument to capture tracking

[apilipenko]

Summary

Currently CaptureTracker gives up if it encounters a value with more than 20 uses. The motivation for this cap is to keep it relatively cheap for BasicAliasAnalysis use case, where the results can't be cached. Although, other clients of CaptureTracker might be ok with higher cost. This patch introduces an argument for PointerMayBeCaptured functions to specify the max number of uses to explore. The motivation for this change is a downstream user of CaptureTracker, but I believe upstream clients of CaptureTracker might also benefit from more fine grained cap.

[hfinkel]

Can you please upload the patch with full context?

I think this makes sense, and is consistent with how we expose the depth threshold in GetUnderlyingObject (and similar).

[apilipenko]

Add diff with context.

[hfinkel]

LGTM

include/llvm/Analysis/CaptureTracking.h

 namespace llvm {
 
   class Value;
   class Use;
   class Instruction;
   class DominatorTree;
   class OrderedBasicBlock;
 
+  /// The default value for MaxUsesToExplore argument. It's relatively small to
+  /// keep the cost of analysis reasonable for clients like BasicAliasAnalysis,
+  /// where the results can't be cached.
+  /// TODO: we should probably introduce a caching CaptureTracking analysis and
+  /// use it where possible. The caching version can use much higher limit or
+  /// don't have this cap at all.
+  unsigned constexpr DefaultMaxUsesToExplore = 20;
+
   /// 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.  The boolean StoreCaptures specified
   /// whether storing the value (or part of it) into memory anywhere
   /// automatically counts as capturing it or not.
+  /// MaxUsesToExplore specifies how many uses should the analysis explore for
+  /// one value before giving up due too "too many uses".
   bool PointerMayBeCaptured(const Value *V,
                             bool ReturnCaptures,
-                            bool StoreCaptures);
+                            bool StoreCaptures,
+                            unsigned MaxUsesToExplore = DefaultMaxUsesToExplore);
 
   /// 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.  The boolean StoreCaptures specified whether storing the value
   /// (or part of it) into memory anywhere automatically counts as capturing it
   /// or not. Captures by the provided instruction are considered if the
   /// final parameter is true. An ordered basic block in \p OBB could be used
   /// to speed up capture-tracker queries.
+  /// MaxUsesToExplore specifies how many uses should the analysis explore for
+  /// one value before giving up due too "too many uses".
   bool PointerMayBeCapturedBefore(const Value *V, bool ReturnCaptures,
                                   bool StoreCaptures, const Instruction *I,
                                   const DominatorTree *DT, bool IncludeI = false,
-                                  OrderedBasicBlock *OBB = nullptr);
+                                  OrderedBasicBlock *OBB = nullptr,
+                                  unsigned MaxUsesToExplore = DefaultMaxUsesToExplore);
 
   /// This callback is used in conjunction with PointerMayBeCaptured. In
   /// addition to the interface here, you'll need to provide your own getters
   /// to see whether anything was captured.
   struct CaptureTracker {
     virtual ~CaptureTracker();
 
     /// tooManyUses - The depth of traversal has breached a limit. There may be
struct CaptureTracker {
     /// use U. Return true to stop the traversal or false to continue looking
     /// for more capturing instructions.
     virtual bool captured(const Use *U) = 0;
   };
 
   /// PointerMayBeCaptured - Visit the value and the values derived from it and
   /// find values which appear to be capturing the pointer value. This feeds
   /// results into and is controlled by the CaptureTracker object.
-  void PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker);
+  /// MaxUsesToExplore specifies how many uses should the analysis explore for
+  /// one value before giving up due too "too many uses".
+  void PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker,
+                            unsigned MaxUsesToExplore = DefaultMaxUsesToExplore);
 } // end namespace llvm
 
 #endif

lib/Analysis/CaptureTracking.cpp

 /// 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.  The boolean StoreCaptures specified whether
 /// storing the value (or part of it) into memory anywhere automatically
 /// counts as capturing it or not.
 bool llvm::PointerMayBeCaptured(const Value *V,
-                                bool ReturnCaptures, bool StoreCaptures) {
+                                bool ReturnCaptures, bool StoreCaptures,
+                                unsigned MaxUsesToExplore) {
   assert(!isa<GlobalValue>(V) &&
          "It doesn't make sense to ask whether a global is captured.");
 
   // TODO: If StoreCaptures is not true, we could do Fancy analysis
   // to determine whether this store is not actually an escape point.
   // In that case, BasicAliasAnalysis should be updated as well to
   // take advantage of this.
   (void)StoreCaptures;

 /// returning the value (or part of it) from the function counts as capturing
 /// it or not.  The boolean StoreCaptures specified whether storing the value
 /// (or part of it) into memory anywhere automatically counts as capturing it
 /// or not. A ordered basic block \p OBB can be used in order to speed up
 /// queries about relative order among instructions in the same basic block.
 bool llvm::PointerMayBeCapturedBefore(const Value *V, bool ReturnCaptures,
                                       bool StoreCaptures, const Instruction *I,
                                       const DominatorTree *DT, bool IncludeI,
-                                      OrderedBasicBlock *OBB) {
+                                      OrderedBasicBlock *OBB,
+                                      unsigned MaxUsesToExplore) {
   assert(!isa<GlobalValue>(V) &&
          "It doesn't make sense to ask whether a global is captured.");
   bool UseNewOBB = OBB == nullptr;
 
   if (!DT)
     return PointerMayBeCaptured(V, ReturnCaptures, StoreCaptures);
   if (UseNewOBB)
     OBB = new OrderedBasicBlock(I->getParent());
 
   // TODO: See comment in PointerMayBeCaptured regarding what could be done
   // with StoreCaptures.
 
   CapturesBefore CB(ReturnCaptures, I, DT, IncludeI, OBB);
   PointerMayBeCaptured(V, &CB);
 
   if (UseNewOBB)
     delete OBB;
   return CB.Captured;
 }
 
-/// TODO: Write a new FunctionPass AliasAnalysis so that it can keep
-/// a cache. Then we can move the code from BasicAliasAnalysis into
-/// that path, and remove this threshold.
-static unsigned const Threshold = 20;
-
-void llvm::PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker) {
+void llvm::PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker,
+                                unsigned MaxUsesToExplore) {
   assert(V->getType()->isPointerTy() && "Capture is for pointers only!");
-  SmallVector<const Use *, Threshold> Worklist;
-  SmallSet<const Use *, Threshold> Visited;
+  SmallVector<const Use *, DefaultMaxUsesToExplore> Worklist;
+  SmallSet<const Use *, DefaultMaxUsesToExplore> Visited;
 
   auto AddUses = [&](const Value *V) {
     unsigned Count = 0;
     for (const Use &U : V->uses()) {
       // If there are lots of uses, conservatively say that the value
       // is captured to avoid taking too much compile time.
-      if (Count++ >= Threshold)
+      if (Count++ >= MaxUsesToExplore)
         return Tracker->tooManyUses();
       if (!Visited.insert(&U).second)
         continue;
       if (!Tracker->shouldExplore(&U))
         continue;
       Worklist.push_back(&U);
     }
   };