[ObjCARC][Contract] Optimize bundled RetainRV to ClaimRV (#138697)

This teaches ObjCARCContract to transform attachedcall bundles
referencing objc_retainAutoreleasedReturnValue to instead reference
objc_claimAutoreleasedReturnValue.

The only distinction between the two is that the latter is required to
be guaranteed to immediately follow the call it's attached to, and, by
construction, the bundles always achieve that by:
- not being separable from the call through IR and the backend
- not getting the marker emitted when claimARV is the attachedcall.

This is enabled only for arm64, arm64e, and arm64_32 on macOS13/iOS16
and related operating systems.

Co-authored-by: Ahmed Bougacha <[email protected]>

Author: citymarina

llvm/lib/Transforms/ObjCARC/ARCRuntimeEntryPoints.h

@@ -42,6 +42,7 @@ enum class ARCRuntimeEntryPointKind {
   Autorelease,
   StoreStrong,
   RetainRV,
+  ClaimRV,
   UnsafeClaimRV,
   RetainAutorelease,
   RetainAutoreleaseRV,
@@ -62,6 +63,7 @@ class ARCRuntimeEntryPoints {
     Autorelease = nullptr;
     StoreStrong = nullptr;
     RetainRV = nullptr;
+    ClaimRV = nullptr;
     UnsafeClaimRV = nullptr;
     RetainAutorelease = nullptr;
     RetainAutoreleaseRV = nullptr;
@@ -87,6 +89,9 @@ class ARCRuntimeEntryPoints {
     case ARCRuntimeEntryPointKind::RetainRV:
       return getIntrinsicEntryPoint(RetainRV,
                                 Intrinsic::objc_retainAutoreleasedReturnValue);
+    case ARCRuntimeEntryPointKind::ClaimRV:
+      return getIntrinsicEntryPoint(
+          ClaimRV, Intrinsic::objc_claimAutoreleasedReturnValue);
     case ARCRuntimeEntryPointKind::UnsafeClaimRV:
       return getIntrinsicEntryPoint(
           UnsafeClaimRV, Intrinsic::objc_unsafeClaimAutoreleasedReturnValue);
@@ -126,6 +131,9 @@ class ARCRuntimeEntryPoints {
   /// Declaration for objc_retainAutoreleasedReturnValue().
   Function *RetainRV = nullptr;
 
+  /// Declaration for objc_claimAutoreleasedReturnValue().
+  Function *ClaimRV = nullptr;
+
   /// Declaration for objc_unsafeClaimAutoreleasedReturnValue().
   Function *UnsafeClaimRV = nullptr;
 

llvm/lib/Transforms/ObjCARC/ObjCARC.cpp

@@ -101,8 +101,37 @@ BundledRetainClaimRVs::~BundledRetainClaimRVs() {
       // can't be tail calls.
       if (auto *CI = dyn_cast<CallInst>(CB))
         CI->setTailCallKind(CallInst::TCK_NoTail);
+
+      // We can also do one final optimization: modify the bundle in the
+      // annotated call, to change the bundle operand from
+      //   objc_retainAutoreleasedReturnValue
+      // to:
+      //   objc_claimAutoreleasedReturnValue
+      // allowing the marker to be omitted from the bundle expansion later.
+      //
+      // Note that, confusingly, ClaimRV is semantically equivalent to RetainRV,
+      // and only differs in that it doesn't require the marker.
+      // The bundle provides the guarantee that we're emitting the ClaimRV call
+      // adjacent to the original call, and providing that guarantee is the
+      // only difference between ClaimRV and RetainRV.
+      //
+      // UnsafeClaimRV has a different RC contract entirely.
+
+      // Find the clang.arc.attachedcall bundle, and rewrite its operand.
+      if (UseClaimRV) {
+        for (auto OBI : CB->bundle_op_infos()) {
+          auto OBU = CB->operandBundleFromBundleOpInfo(OBI);
+          if (OBU.getTagID() == LLVMContext::OB_clang_arc_attachedcall &&
+              OBU.Inputs[0] == EP.get(ARCRuntimeEntryPointKind::RetainRV)) {
+            CB->setOperand(OBI.Begin,
+                           EP.get(ARCRuntimeEntryPointKind::ClaimRV));
+            break;
+          }
+        }
+      }
     }
 
+    // Erase the RV call we emitted earlier: it's already in the bundle.
     EraseInstruction(P.first);
   }
 

llvm/lib/Transforms/ObjCARC/ObjCARC.h

@@ -22,6 +22,7 @@
 #ifndef LLVM_LIB_TRANSFORMS_OBJCARC_OBJCARC_H
 #define LLVM_LIB_TRANSFORMS_OBJCARC_OBJCARC_H
 
+#include "ARCRuntimeEntryPoints.h"
 #include "llvm/Analysis/ObjCARCAnalysisUtils.h"
 #include "llvm/Analysis/ObjCARCUtil.h"
 #include "llvm/IR/EHPersonalities.h"
@@ -104,7 +105,9 @@ CallInst *createCallInstWithColors(
 
 class BundledRetainClaimRVs {
 public:
-  BundledRetainClaimRVs(bool ContractPass) : ContractPass(ContractPass) {}
+  BundledRetainClaimRVs(ARCRuntimeEntryPoints &EP, bool ContractPass,
+                        bool UseClaimRV)
+      : EP(EP), ContractPass(ContractPass), UseClaimRV(UseClaimRV) {}
   ~BundledRetainClaimRVs();
 
   /// Insert a retainRV/claimRV call to the normal destination blocks of invokes
@@ -155,7 +158,9 @@ class BundledRetainClaimRVs {
   /// A map of inserted retainRV/claimRV calls to annotated calls/invokes.
   DenseMap<CallInst *, CallBase *> RVCalls;
 
+  ARCRuntimeEntryPoints &EP;
   bool ContractPass;
+  bool UseClaimRV;
 };
 
 } // end namespace objcarc

llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp

@@ -42,6 +42,7 @@
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/ObjCARC.h"
 
 using namespace llvm;
@@ -52,6 +53,11 @@ using namespace llvm::objcarc;
 STATISTIC(NumPeeps,       "Number of calls peephole-optimized");
 STATISTIC(NumStoreStrongs, "Number objc_storeStrong calls formed");
 
+static cl::opt<cl::boolOrDefault> UseObjCClaimRV(
+    "arc-contract-use-objc-claim-rv",
+    cl::desc(
+        "Enable generation of calls to objc_claimAutoreleasedReturnValue"));
+
 //===----------------------------------------------------------------------===//
 //                                Declarations
 //===----------------------------------------------------------------------===//
@@ -74,6 +80,9 @@ class ObjCARCContract {
   /// A flag indicating whether this optimization pass should run.
   bool Run;
 
+  /// Whether objc_claimAutoreleasedReturnValue is available.
+  bool HasClaimRV = false;
+
   /// The inline asm string to insert between calls and RetainRV calls to make
   /// the optimization work on targets which need it.
   const MDString *RVInstMarker;
@@ -517,6 +526,39 @@ bool ObjCARCContract::tryToPeepholeInstruction(
   }
 }
 
+/// Should we use objc_claimAutoreleasedReturnValue?
+static bool useClaimRuntimeCall(Module &M) {
+  // Let the flag override our OS-based default.
+  if (UseObjCClaimRV != cl::BOU_UNSET)
+    return UseObjCClaimRV == cl::BOU_TRUE;
+
+  Triple TT(M.getTargetTriple());
+
+  // On x86_64, claimARV doesn't make sense, as the marker isn't actually a nop
+  // there (it's needed by the calling convention).
+  if (!TT.isAArch64())
+    return false;
+
+  unsigned Major = TT.getOSMajorVersion();
+  switch (TT.getOS()) {
+  default:
+    return false;
+  case Triple::IOS:
+  case Triple::TvOS:
+    return Major >= 16;
+  case Triple::WatchOS:
+    return Major >= 9;
+  case Triple::BridgeOS:
+    return Major >= 7;
+  case Triple::MacOSX:
+    return Major >= 13;
+  case Triple::Darwin:
+    return Major >= 21;
+  }
+
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 //                              Top Level Driver
 //===----------------------------------------------------------------------===//
@@ -528,6 +570,8 @@ bool ObjCARCContract::init(Module &M) {
 
   EP.init(&M);
 
+  HasClaimRV = useClaimRuntimeCall(M);
+
   // Initialize RVInstMarker.
   RVInstMarker = getRVInstMarker(M);
 
@@ -545,7 +589,7 @@ bool ObjCARCContract::run(Function &F, AAResults *A, DominatorTree *D) {
   AA = A;
   DT = D;
   PA.setAA(A);
-  BundledRetainClaimRVs BRV(/*ContractPass=*/true);
+  BundledRetainClaimRVs BRV(EP, /*ContractPass=*/true, HasClaimRV);
   BundledInsts = &BRV;
 
   std::pair<bool, bool> R = BundledInsts->insertAfterInvokes(F, DT);

llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp

@@ -2426,7 +2426,7 @@ bool ObjCARCOpt::run(Function &F, AAResults &AA) {
     return false;
 
   Changed = CFGChanged = false;
-  BundledRetainClaimRVs BRV(/*ContractPass=*/false);
+  BundledRetainClaimRVs BRV(EP, /*ContractPass=*/false, /*UseClaimRV=*/false);
   BundledInsts = &BRV;
 
   LLVM_DEBUG(dbgs() << "<<< ObjCARCOpt: Visiting Function: " << F.getName()

llvm/test/Transforms/ObjCARC/contract-attached-call-retain-to-claim.ll

@@ -0,0 +1,35 @@
+; RUN: opt -passes=objc-arc-contract -arc-contract-use-objc-claim-rv=1 -S < %s | FileCheck %s --check-prefixes=CHECK,CLAIM
+; RUN: opt -passes=objc-arc-contract -arc-contract-use-objc-claim-rv=0 -S < %s | FileCheck %s --check-prefixes=CHECK,RETAIN
+
+; CHECK-LABEL: define void @test0() {
+; CLAIM: %[[CALL:.*]] = notail call ptr @foo() [ "clang.arc.attachedcall"(ptr @llvm.objc.claimAutoreleasedReturnValue) ]
+; RETAIN: %[[CALL:.*]] = notail call ptr @foo() [ "clang.arc.attachedcall"(ptr @llvm.objc.retainAutoreleasedReturnValue) ]
+; CHECK-NEXT: ret void
+
+define void @test0() {
+  %call1 = call ptr @foo() [ "clang.arc.attachedcall"(ptr @llvm.objc.retainAutoreleasedReturnValue) ]
+  ret void
+}
+
+; CHECK-LABEL: define void @test1() {
+; CHECK: %[[CALL:.*]] = notail call ptr @foo() [ "clang.arc.attachedcall"(ptr @llvm.objc.unsafeClaimAutoreleasedReturnValue) ]
+; CHECK-NEXT: ret void
+
+define void @test1() {
+  %call1 = call ptr @foo() [ "clang.arc.attachedcall"(ptr @llvm.objc.unsafeClaimAutoreleasedReturnValue) ]
+  ret void
+}
+
+; CHECK-LABEL: define void @test2() {
+; CLAIM: %[[CALL:.*]] = notail call ptr @foo() [ "clang.arc.attachedcall"(ptr @llvm.objc.claimAutoreleasedReturnValue), "otherbundle"() ]
+; RETAIN: %[[CALL:.*]] = notail call ptr @foo() [ "clang.arc.attachedcall"(ptr @llvm.objc.retainAutoreleasedReturnValue), "otherbundle"() ]
+; CHECK-NEXT: ret void
+
+define void @test2() {
+  %call1 = call ptr @foo() [ "clang.arc.attachedcall"(ptr @llvm.objc.retainAutoreleasedReturnValue), "otherbundle"() ]
+  ret void
+}
+
+declare ptr @foo()
+declare ptr @llvm.objc.retainAutoreleasedReturnValue(ptr)
+declare ptr @llvm.objc.unsafeClaimAutoreleasedReturnValue(ptr)