[ObjC][ARC] Check the basic block size before calling DominatorTree::dominate

[ahatanak]

Summary

ARC contract pass has an optimization that replaces the uses of the argument of an ObjC runtime function calls with the call result. For example:

; Before optimization
%1 = tail call i8* @foo1()
%2 = tail call i8* @llvm.objc.retainAutoreleasedReturnValue(i8* %1)
store i8* %1, i8** @g0, align 8

; After optimization
%1 = tail call i8* @foo1()
%2 = tail call i8* @llvm.objc.retainAutoreleasedReturnValue(i8* %1)
store i8* %2, i8** @g0, align 8 // %1 is replaced with %2

Before replacing the argument, DominatorTree::dominate is called to determine whether the user instruction is dominated by the ObjC runtime function call instruction. The call to DominatorTree::dominate can be expensive if the two instructions belong to the same basic block and the size of the basic block is large. This patch checks the basic block size and just bails out if the size exceeds the limit.

rdar://problem/49477063

[pete]

LGTM.

[ahatanak]

Fix the comment on why we have to check the size of the basic block before calling DominatorTree::dominate to compute the dominance relation between two instructions in the same basic block.

lib/Transforms/ObjCARC/ObjCARCContract.cpp

 using namespace llvm;
 using namespace llvm::objcarc;
 
 #define DEBUG_TYPE "objc-arc-contract"
 
 STATISTIC(NumPeeps,       "Number of calls peephole-optimized");
 STATISTIC(NumStoreStrongs, "Number objc_storeStrong calls formed");
 
+static cl::opt<unsigned> MaxBBSize("arc-contract-max-bb-size", cl::Hidden,
+    cl::desc("Maximum basic block size to discover the dominance relation of "
+             "two instructions in the same basic block"), cl::init(65535));
+
 //===----------------------------------------------------------------------===//
 //                                Declarations
 //===----------------------------------------------------------------------===//
 
 namespace {
   /// Late ARC optimizations
   ///
   /// These change the IR in a way that makes it difficult to be analyzed by
bool ObjCARCContract::runOnFunction(Function &F) {
   bool TailOkForStoreStrongs =
       !F.isVarArg() && !F.callsFunctionThatReturnsTwice();
 
   // For ObjC library calls which return their argument, replace uses of the
   // argument with uses of the call return value, if it dominates the use. This
   // reduces register pressure.
   SmallPtrSet<Instruction *, 4> DependingInstructions;
   SmallPtrSet<const BasicBlock *, 4> Visited;
+
+  // Cache the basic block size.
+  DenseMap<const BasicBlock *, unsigned> BBSizeMap;
+
+  // A lambda that lazily computes the size of a basic block and determines
+  // whether the size exceeds MaxBBSize.
+  auto IsLargeBB = [&](const BasicBlock *BB) {
+    unsigned BBSize;
+    auto I = BBSizeMap.find(BB);
+
+    if (I != BBSizeMap.end())
+      BBSize = I->second;
+    else
+      BBSize = BBSizeMap[BB] = BB->size();
+
+    return BBSize > MaxBBSize;
+  };
+
   for (inst_iterator I = inst_begin(&F), E = inst_end(&F); I != E;) {
     Instruction *Inst = &*I++;
 
     LLVM_DEBUG(dbgs() << "Visiting: " << *Inst << "\n");
 
     // First try to peephole Inst. If there is nothing further we can do in
     // terms of undoing objc-arc-expand, process the next inst.
     if (tryToPeepholeInstruction(F, Inst, I, DependingInstructions, Visited,
                                  TailOkForStoreStrongs, BlockColors))
       continue;
 
     // Otherwise, try to undo objc-arc-expand.
 
     // Don't use GetArgRCIdentityRoot because we don't want to look through bitcasts
     // and such; to do the replacement, the argument must have type i8*.
 
     // Function for replacing uses of Arg dominated by Inst.
-    auto ReplaceArgUses = [Inst, this](Value *Arg) {
+    auto ReplaceArgUses = [Inst, IsLargeBB, this](Value *Arg) {
       // If we're compiling bugpointed code, don't get in trouble.
       if (!isa<Instruction>(Arg) && !isa<Argument>(Arg))
         return;
 
       // Look through the uses of the pointer.
       for (Value::use_iterator UI = Arg->use_begin(), UE = Arg->use_end();
            UI != UE; ) {
         // Increment UI now, because we may unlink its element.
         Use &U = *UI++;
         unsigned OperandNo = U.getOperandNo();
 
+        // Don't replace the uses if Inst and the user belong to the same basic
+        // block and the size of the basic block is large. We don't want to call
+        // DominatorTree::dominate in that case. We can remove this check when
+        // we have a way to compute the dominance relation between two
+        // instructions that belong to the same basic block in constant time.
+        if (Inst->getParent() == cast<Instruction>(U.getUser())->getParent() &&
+            IsLargeBB(Inst->getParent()))
+          continue;
+
         // If the call's return value dominates a use of the call's argument
         // value, rewrite the use to use the return value. We check for
         // reachability here because an unreachable call is considered to
         // trivially dominate itself, which would lead us to rewriting its
         // argument in terms of its return value, which would lead to
         // infinite loops in GetArgRCIdentityRoot.
         if (!DT->isReachableFromEntry(U) || !DT->dominates(Inst, U))
           continue;

test/Transforms/ObjCARC/contract-max-bb-size.ll

+; RUN: opt -objc-arc-contract -S < %s | FileCheck -check-prefix=ENABLE %s
+; RUN: opt -objc-arc-contract -arc-contract-max-bb-size=3 -S < %s | FileCheck -check-prefix=DISABLE %s
+
+@g0 = common global i8* null, align 8
+
+; ENABLE: store i8* %2, i8** @g0
+; DISABLE: store i8* %1, i8** @g0
+
+define void @foo0() {
+  %1 = tail call i8* @foo1()
+  %2 = tail call i8* @llvm.objc.retainAutoreleasedReturnValue(i8* %1)
+  store i8* %1, i8** @g0, align 8
+  ret void
+}
+
+declare i8* @foo1()
+declare i8* @llvm.objc.retainAutoreleasedReturnValue(i8*)