[ObjCARC] Delete empty autoreleasepools with no autoreleases in them

Erase empty autorelease pools that have no autorelease in them

Author: AZero13

llvm/lib/Transforms/ObjCARC/ARCRuntimeEntryPoints.h

@@ -46,6 +46,8 @@ enum class ARCRuntimeEntryPointKind {
   UnsafeClaimRV,
   RetainAutorelease,
   RetainAutoreleaseRV,
+  AutoreleasePoolPush,
+  AutoreleasePoolPop,
 };
 
 /// Declarations for ObjC runtime functions and constants. These are initialized
@@ -67,6 +69,8 @@ class ARCRuntimeEntryPoints {
     UnsafeClaimRV = nullptr;
     RetainAutorelease = nullptr;
     RetainAutoreleaseRV = nullptr;
+    AutoreleasePoolPush = nullptr;
+    AutoreleasePoolPop = nullptr;
   }
 
   Function *get(ARCRuntimeEntryPointKind kind) {
@@ -101,6 +105,12 @@ class ARCRuntimeEntryPoints {
     case ARCRuntimeEntryPointKind::RetainAutoreleaseRV:
       return getIntrinsicEntryPoint(RetainAutoreleaseRV,
                                 Intrinsic::objc_retainAutoreleaseReturnValue);
+    case ARCRuntimeEntryPointKind::AutoreleasePoolPush:
+      return getIntrinsicEntryPoint(AutoreleasePoolPush,
+                                    Intrinsic::objc_autoreleasePoolPush);
+    case ARCRuntimeEntryPointKind::AutoreleasePoolPop:
+      return getIntrinsicEntryPoint(AutoreleasePoolPop,
+                                    Intrinsic::objc_autoreleasePoolPop);
     }
 
     llvm_unreachable("Switch should be a covered switch.");
@@ -143,6 +153,12 @@ class ARCRuntimeEntryPoints {
   /// Declaration for objc_retainAutoreleaseReturnValue().
   Function *RetainAutoreleaseRV = nullptr;
 
+  /// Declaration for objc_autoreleasePoolPush().
+  Function *AutoreleasePoolPush = nullptr;
+
+  /// Declaration for objc_autoreleasePoolPop().
+  Function *AutoreleasePoolPop = nullptr;
+
   Function *getIntrinsicEntryPoint(Function *&Decl, Intrinsic::ID IntID) {
     if (Decl)
       return Decl;

llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp

@@ -132,11 +132,8 @@ static const Value *FindSingleUseIdentifiedObject(const Value *Arg) {
 //
 // The second retain and autorelease can be deleted.
 
-// TODO: It should be possible to delete
-// objc_autoreleasePoolPush and objc_autoreleasePoolPop
-// pairs if nothing is actually autoreleased between them. Also, autorelease
-// calls followed by objc_autoreleasePoolPop calls (perhaps in ObjC++ code
-// after inlining) can be turned into plain release calls.
+// TODO: Autorelease calls followed by objc_autoreleasePoolPop calls (perhaps in
+// ObjC++ code after inlining) can be turned into plain release calls.
 
 // TODO: Critical-edge splitting. If the optimial insertion point is
 // a critical edge, the current algorithm has to fail, because it doesn't
@@ -566,6 +563,8 @@ class ObjCARCOpt {
 
   void OptimizeReturns(Function &F);
 
+  void OptimizeAutoreleasePools(Function &F);
+
   template <typename PredicateT>
   static void cloneOpBundlesIf(CallBase *CI,
                                SmallVectorImpl<OperandBundleDef> &OpBundles,
@@ -2473,6 +2472,11 @@ bool ObjCARCOpt::run(Function &F, AAResults &AA) {
                             (1 << unsigned(ARCInstKind::AutoreleaseRV))))
     OptimizeReturns(F);
 
+  // Optimizations for autorelease pools.
+  if (UsedInThisFunction & ((1 << unsigned(ARCInstKind::AutoreleasepoolPush)) |
+                            (1 << unsigned(ARCInstKind::AutoreleasepoolPop))))
+    OptimizeAutoreleasePools(F);
+
   // Gather statistics after optimization.
 #ifndef NDEBUG
   if (AreStatisticsEnabled()) {
@@ -2485,6 +2489,107 @@ bool ObjCARCOpt::run(Function &F, AAResults &AA) {
   return Changed;
 }
 
+/// Optimize autorelease pools by eliminating empty push/pop pairs.
+void ObjCARCOpt::OptimizeAutoreleasePools(Function &F) {
+  LLVM_DEBUG(dbgs() << "\n== ObjCARCOpt::OptimizeAutoreleasePools ==\n");
+
+  // Track empty autorelease pool push/pop pairs
+  SmallVector<std::pair<CallInst *, CallInst *>, 4> EmptyPoolPairs;
+
+  // Process each basic block independently.
+  // TODO: Can we optimize inter-block autorelease pool pairs?
+  // This would involve tracking autorelease pool state across blocks.
+  for (BasicBlock &BB : F) {
+    CallInst *PendingPush = nullptr;
+    bool HasAutoreleaseInScope = false;
+
+    for (Instruction &Inst : BB) {
+      ARCInstKind Class = GetBasicARCInstKind(&Inst);
+
+      switch (Class) {
+      case ARCInstKind::AutoreleasepoolPush: {
+        // Start tracking a new autorelease pool scope
+        PendingPush = cast<CallInst>(&Inst);
+        HasAutoreleaseInScope = false;
+        LLVM_DEBUG(dbgs() << "Found autorelease pool push: " << *PendingPush
+                          << "\n");
+        break;
+      }
+
+      case ARCInstKind::AutoreleasepoolPop: {
+        CallInst *Pop = cast<CallInst>(&Inst);
+
+        if (PendingPush) {
+          // Check if this pop matches the pending push by comparing the token
+          Value *PopArg = Pop->getArgOperand(0);
+          bool IsMatchingPop = (PopArg == PendingPush);
+
+          // Also handle bitcast case
+          if (!IsMatchingPop && isa<BitCastInst>(PopArg)) {
+            Value *BitcastSrc = cast<BitCastInst>(PopArg)->getOperand(0);
+            IsMatchingPop = (BitcastSrc == PendingPush);
+          }
+
+          if (IsMatchingPop && !HasAutoreleaseInScope) {
+            LLVM_DEBUG(dbgs() << "Eliminating empty autorelease pool pair: "
+                              << *PendingPush << " and " << *Pop << "\n");
+
+            // Store the pair for careful deletion later
+            EmptyPoolPairs.push_back({PendingPush, Pop});
+
+            Changed = true;
+            ++NumNoops;
+          }
+        }
+
+        PendingPush = nullptr;
+        HasAutoreleaseInScope = false;
+        break;
+      }
+      case ARCInstKind::CallOrUser:
+      case ARCInstKind::Call:
+      case ARCInstKind::Autorelease:
+      case ARCInstKind::AutoreleaseRV: {
+        // Track that we have autorelease calls in the current pool scope
+        if (PendingPush) {
+          HasAutoreleaseInScope = true;
+          LLVM_DEBUG(
+              dbgs()
+              << "Found autorelease or potiential autorelease in pool scope: "
+              << Inst << "\n");
+        }
+        break;
+      }
+
+      default:
+        break;
+      }
+    }
+  }
+
+  // Handle empty pool pairs carefully to avoid use-after-delete
+  SmallVector<CallInst *, 8> DeadInsts;
+  for (auto &Pair : EmptyPoolPairs) {
+    CallInst *Push = Pair.first;
+    CallInst *Pop = Pair.second;
+
+    // Replace the pop's argument with poison to break dependencies
+    Value *PoisonToken = PoisonValue::get(Push->getType());
+    Pop->setArgOperand(0, PoisonToken);
+
+    LLVM_DEBUG(dbgs() << "Erasing empty pool pair: " << *Push << " and " << *Pop
+                      << "\n");
+    DeadInsts.push_back(Pop);
+    DeadInsts.push_back(Push);
+  }
+
+  // Remove the pairs
+  for (CallInst *DeadInst : DeadInsts) {
+    LLVM_DEBUG(dbgs() << "Erasing dead instruction: " << *DeadInst << "\n");
+    DeadInst->eraseFromParent();
+  }
+}
+
 /// @}
 ///
 

llvm/test/Transforms/ObjCARC/test_autorelease_pool.ll

@@ -0,0 +1,79 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; Test for autorelease pool optimizations
+; RUN: opt -passes=objc-arc < %s -S | FileCheck %s
+
+declare ptr @llvm.objc.autoreleasePoolPush()
+declare void @llvm.objc.autoreleasePoolPop(ptr)
+declare ptr @llvm.objc.autorelease(ptr)
+declare ptr @llvm.objc.retain(ptr)
+declare ptr @create_object()
+declare void @use_object(ptr)
+declare ptr @object_with_thing()
+
+; Test 1: Empty autorelease pool should be eliminated
+define void @test_empty_pool() {
+; CHECK-LABEL: define void @test_empty_pool() {
+; CHECK-NEXT:    ret void
+;
+  %pool = call ptr @llvm.objc.autoreleasePoolPush()
+  call void @llvm.objc.autoreleasePoolPop(ptr %pool)
+  ret void
+}
+
+; Test 2: Pool with only release should be removed
+define void @test_autorelease_to_release() {
+; CHECK-LABEL: define void @test_autorelease_to_release() {
+; CHECK-NEXT:    [[OBJ:%.*]] = call ptr @create_object()
+; CHECK-NEXT:    call void @llvm.objc.release(ptr [[OBJ]]) #[[ATTR0:[0-9]+]], !clang.imprecise_release [[META0:![0-9]+]]
+; CHECK-NEXT:    ret void
+;
+  %obj = call ptr @create_object()
+  %pool = call ptr @llvm.objc.autoreleasePoolPush()
+  call ptr @llvm.objc.autorelease(ptr %obj)
+  call void @llvm.objc.autoreleasePoolPop(ptr %pool)
+  ret void
+}
+
+; Test 3: Pool with autoreleases should not be optimized
+define void @test_multiple_autoreleases() {
+; CHECK-LABEL: define void @test_multiple_autoreleases() {
+; CHECK-NEXT:    [[OBJ1:%.*]] = call ptr @create_object()
+; CHECK-NEXT:    [[OBJ2:%.*]] = call ptr @create_object()
+; CHECK-NEXT:    [[POOL:%.*]] = call ptr @llvm.objc.autoreleasePoolPush() #[[ATTR0]]
+; CHECK-NEXT:    call void @use_object(ptr [[OBJ1]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call ptr @llvm.objc.autorelease(ptr [[OBJ1]]) #[[ATTR0]]
+; CHECK-NEXT:    call void @use_object(ptr [[OBJ2]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call ptr @llvm.objc.autorelease(ptr [[OBJ2]]) #[[ATTR0]]
+; CHECK-NEXT:    call void @llvm.objc.autoreleasePoolPop(ptr [[POOL]]) #[[ATTR0]]
+; CHECK-NEXT:    ret void
+;
+  %obj1 = call ptr @create_object()
+  %obj2 = call ptr @create_object()
+  %pool = call ptr @llvm.objc.autoreleasePoolPush()
+  call void @use_object(ptr %obj1)
+  call ptr @llvm.objc.autorelease(ptr %obj1)
+  call void @use_object(ptr %obj2)
+  call ptr @llvm.objc.autorelease(ptr %obj2)
+  call void @llvm.objc.autoreleasePoolPop(ptr %pool)
+  ret void
+}
+
+; Test 4: Pool with calls should not be optimized
+define void @test_calls() {
+; CHECK-LABEL: define void @test_calls() {
+; CHECK-NEXT:    [[POOL:%.*]] = call ptr @llvm.objc.autoreleasePoolPush() #[[ATTR0]]
+; CHECK-NEXT:    [[OBJ1:%.*]] = call ptr @object_with_thing()
+; CHECK-NEXT:    call void @use_object(ptr [[OBJ1]])
+; CHECK-NEXT:    call void @llvm.objc.autoreleasePoolPop(ptr [[POOL]]) #[[ATTR0]]
+; CHECK-NEXT:    ret void
+;
+  %pool = call ptr @llvm.objc.autoreleasePoolPush()
+  %obj1 = call ptr @object_with_thing()
+  call void @use_object(ptr %obj1)
+  call void @llvm.objc.autoreleasePoolPop(ptr %pool)
+  ret void
+}
+
+;.
+; CHECK: [[META0]] = !{}
+;.