[ObjC][ARC] Let ARC optimizer bail out if the number of pointer states it keeps track of becomes too large

[ahatanak]

Summary

ARC optimizer does a top-down and a bottom-up traversal of the whole function to pair up retain and release instructions and remove them. This can be expensive if the number of instructions in the function and pointer states it tracks are large since it has to look at each pointer state and determine whether the instruction being visited can potentially use the pointer (see the loop in ObjCARCOpt::VisitInstructionBottomUp that calls HandlePotentialUse).

This patch adds a command line option that sets a limit to the number of pointers it tracks.

rdar://problem/49477063

[pete]

Wow. That must be a huge test case for this to be a problem!

LGTM.

[ahatanak]

The test case I'm looking at has a function with over 500k instructions and lots of ObjC runtime function calls.

[smeenai]

LGTM

In D61100#1477850, @ahatanak wrote:

The test case I'm looking at has a function with over 500k instructions and lots of ObjC runtime function calls.

lib/Transforms/ObjCARC/ObjCARCOpts.cpp

 #include <iterator>
 #include <utility>
 
 using namespace llvm;
 using namespace llvm::objcarc;
 
 #define DEBUG_TYPE "objc-arc-opts"
 
+static cl::opt<unsigned> MaxPtrStates("arc-opt-max-ptr-states",
+    cl::Hidden,
+    cl::desc("Maximum number of ptr states the optimizer keeps track of"),
+    cl::init(4095));
+
 /// \defgroup ARCUtilities Utility declarations/definitions specific to ARC.
 /// @{
 
 /// This is similar to GetRCIdentityRoot but it stops as soon
 /// as it finds a value with multiple uses.
 static const Value *FindSingleUseIdentifiedObject(const Value *Arg) {
   // ConstantData (like ConstantPointerNull and UndefValue) is used across
   // modules.  It's never a single-use value.
public:
     }
     const_top_down_ptr_iterator top_down_ptr_end() const {
       return PerPtrTopDown.end();
     }
     bool hasTopDownPtrs() const {
       return !PerPtrTopDown.empty();
     }
 
+    unsigned top_down_ptr_list_size() const {
+      return std::distance(top_down_ptr_begin(), top_down_ptr_end());
+    }
+
     using bottom_up_ptr_iterator = decltype(PerPtrBottomUp)::iterator;
     using const_bottom_up_ptr_iterator =
         decltype(PerPtrBottomUp)::const_iterator;
 
     bottom_up_ptr_iterator bottom_up_ptr_begin() {
       return PerPtrBottomUp.begin();
     }
     bottom_up_ptr_iterator bottom_up_ptr_end() { return PerPtrBottomUp.end(); }
     const_bottom_up_ptr_iterator bottom_up_ptr_begin() const {
       return PerPtrBottomUp.begin();
     }
     const_bottom_up_ptr_iterator bottom_up_ptr_end() const {
       return PerPtrBottomUp.end();
     }
     bool hasBottomUpPtrs() const {
       return !PerPtrBottomUp.empty();
     }
 
+    unsigned bottom_up_ptr_list_size() const {
+      return std::distance(bottom_up_ptr_begin(), bottom_up_ptr_end());
+    }
+
     /// Mark this block as being an entry block, which has one path from the
     /// entry by definition.
     void SetAsEntry() { TopDownPathCount = 1; }
 
     /// Mark this block as being an exit block, which has one path to an exit by
     /// definition.
     void SetAsExit()  { BottomUpPathCount = 1; }
 
class ObjCARCOpt : public FunctionPass {
 
     /// A cache of MDKinds that can be passed into other functions to propagate
     /// MDKind identifiers.
     ARCMDKindCache MDKindCache;
 
     /// A flag indicating whether this optimization pass should run.
     bool Run;
 
+    /// A flag indicating whether the optimization that removes or moves
+    /// retain/release pairs should be performed.
+    bool DisableRetainReleasePairing = false;
+
     /// Flags which determine whether each of the interesting runtime functions
     /// is in fact used in the current function.
     unsigned UsedInThisFunction;
 
     bool OptimizeRetainRVCall(Function &F, Instruction *RetainRV);
     void OptimizeAutoreleaseRVCall(Function &F, Instruction *AutoreleaseRV,
                                    ARCInstKind &Class);
     void OptimizeIndividualCalls(Function &F);
for (BasicBlock::iterator I = BB->end(), E = BB->begin(); I != E; --I) {
 
     // Invoke instructions are visited as part of their successors (below).
     if (isa<InvokeInst>(Inst))
       continue;
 
     LLVM_DEBUG(dbgs() << "    Visiting " << *Inst << "\n");
 
     NestingDetected |= VisitInstructionBottomUp(Inst, BB, Retains, MyStates);
+
+    // Bail out if the number of pointers being tracked becomes too large so
+    // that this pass can complete in a reasonable amount of time.
+    if (MyStates.bottom_up_ptr_list_size() > MaxPtrStates) {
+      DisableRetainReleasePairing = true;
+      return false;
+    }
   }
 
   // If there's a predecessor with an invoke, visit the invoke as if it were
   // part of this block, since we can't insert code after an invoke in its own
   // block, and we don't want to split critical edges.
   for (BBState::edge_iterator PI(MyStates.pred_begin()),
        PE(MyStates.pred_end()); PI != PE; ++PI) {
     BasicBlock *Pred = *PI;
LLVM_DEBUG(dbgs() << "Before:\n"
                     << BBStates[BB] << "\n"
                     << "Performing Dataflow:\n");
 
   // Visit all the instructions, top-down.
   for (Instruction &Inst : *BB) {
     LLVM_DEBUG(dbgs() << "    Visiting " << Inst << "\n");
 
     NestingDetected |= VisitInstructionTopDown(&Inst, Releases, MyStates);
+
+    // Bail out if the number of pointers being tracked becomes too large so
+    // that this pass can complete in a reasonable amount of time.
+    if (MyStates.top_down_ptr_list_size() > MaxPtrStates) {
+      DisableRetainReleasePairing = true;
+      return false;
+    }
   }
 
   LLVM_DEBUG(dbgs() << "\nState Before Checking for CFG Hazards:\n"
                     << BBStates[BB] << "\n\n");
   CheckForCFGHazards(BB, BBStates, MyStates);
   LLVM_DEBUG(dbgs() << "Final State:\n" << BBStates[BB] << "\n");
   return NestingDetected;
 }
bool ObjCARCOpt::Visit(Function &F,
   SmallVector<BasicBlock *, 16> PostOrder;
   SmallVector<BasicBlock *, 16> ReverseCFGPostOrder;
   ComputePostOrders(F, PostOrder, ReverseCFGPostOrder,
                     MDKindCache.get(ARCMDKindID::NoObjCARCExceptions),
                     BBStates);
 
   // Use reverse-postorder on the reverse CFG for bottom-up.
   bool BottomUpNestingDetected = false;
-  for (BasicBlock *BB : llvm::reverse(ReverseCFGPostOrder))
+  for (BasicBlock *BB : llvm::reverse(ReverseCFGPostOrder)) {
     BottomUpNestingDetected |= VisitBottomUp(BB, BBStates, Retains);
+    if (DisableRetainReleasePairing)
+      return false;
+  }
 
   // Use reverse-postorder for top-down.
   bool TopDownNestingDetected = false;
-  for (BasicBlock *BB : llvm::reverse(PostOrder))
+  for (BasicBlock *BB : llvm::reverse(PostOrder)) {
     TopDownNestingDetected |= VisitTopDown(BB, BBStates, Releases);
+    if (DisableRetainReleasePairing)
+      return false;
+  }
 
   return TopDownNestingDetected && BottomUpNestingDetected;
 }
 
 /// Move the calls in RetainsToMove and ReleasesToMove.
 void ObjCARCOpt::MoveCalls(Value *Arg, RRInfo &RetainsToMove,
                            RRInfo &ReleasesToMove,
                            BlotMapVector<Value *, RRInfo> &Retains,
bool ObjCARCOpt::OptimizeSequences(Function &F) {
 
   // This is used during the traversal of the function to track the
   // states for each identified object at each block.
   DenseMap<const BasicBlock *, BBState> BBStates;
 
   // Analyze the CFG of the function, and all instructions.
   bool NestingDetected = Visit(F, BBStates, Retains, Releases);
 
+  if (DisableRetainReleasePairing)
+    return false;
+
   // Transform.
   bool AnyPairsCompletelyEliminated = PerformCodePlacement(BBStates, Retains,
                                                            Releases,
                                                            F.getParent());
 
   return AnyPairsCompletelyEliminated && NestingDetected;
 }
 

test/Transforms/ObjCARC/opt-max-ptr-states.ll

+; RUN: opt -objc-arc -S < %s | FileCheck -check-prefix=ENABLE -check-prefix=CHECK %s
+; RUN: opt -objc-arc -arc-opt-max-ptr-states=1 -S < %s | FileCheck -check-prefix=DISABLE -check-prefix=CHECK %s
+
+@g0 = common global i8* null, align 8
+
+; CHECK: call i8* @llvm.objc.retain
+; ENABLE-NOT: call i8* @llvm.objc.retain
+; DISABLE: call i8* @llvm.objc.retain
+; CHECK: call void @llvm.objc.release
+; ENABLE-NOT: call void @llvm.objc.release
+; DISABLE: call void @llvm.objc.release
+
+define void @foo0(i8* %a) {
+  %1 = tail call i8* @llvm.objc.retain(i8* %a)
+  %2 = tail call i8* @llvm.objc.retain(i8* %a)
+  %3 = load i8*, i8** @g0, align 8
+  store i8* %a, i8** @g0, align 8
+  tail call void @llvm.objc.release(i8* %3)
+  tail call void @llvm.objc.release(i8* %a), !clang.imprecise_release !0
+  ret void
+}
+
+declare i8* @llvm.objc.retain(i8*)
+declare void @llvm.objc.release(i8*)
+
+!0 = !{}