[ObjCARC][Contract] Optimize bundled RetainRV to ClaimRV (#138697)users/citymarina/objc-claim-3a

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 <ahmed@bougacha.org>

6 files changed, 124 insertions, 3 deletions

diff --git a/llvm/lib/Transforms/ObjCARC/ARCRuntimeEntryPoints.h b/llvm/lib/Transforms/ObjCARC/ARCRuntimeEntryPoints.h
index 0dedd02..3fa844e 100644
--- a/llvm/lib/Transforms/ObjCARC/ARCRuntimeEntryPoints.h
+++ b/llvm/lib/Transforms/ObjCARC/ARCRuntimeEntryPoints.h
@@ -42,6 +42,7 @@ enum class ARCRuntimeEntryPointKind {
   Autorelease,
   StoreStrong,
   RetainRV,
+  ClaimRV,
   UnsafeClaimRV,
   RetainAutorelease,
   RetainAutoreleaseRV,
@@ -62,6 +63,7 @@ public:
     Autorelease = nullptr;
     StoreStrong = nullptr;
     RetainRV = nullptr;
+    ClaimRV = nullptr;
     UnsafeClaimRV = nullptr;
     RetainAutorelease = nullptr;
     RetainAutoreleaseRV = nullptr;
@@ -87,6 +89,9 @@ public:
     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 @@ private:
   /// Declaration for objc_retainAutoreleasedReturnValue().
   Function *RetainRV = nullptr;
 
+  /// Declaration for objc_claimAutoreleasedReturnValue().
+  Function *ClaimRV = nullptr;
+
   /// Declaration for objc_unsafeClaimAutoreleasedReturnValue().
   Function *UnsafeClaimRV = nullptr;
 
diff --git a/llvm/lib/Transforms/ObjCARC/ObjCARC.cpp b/llvm/lib/Transforms/ObjCARC/ObjCARC.cpp
index b6ade1c..32e7092 100644
--- a/llvm/lib/Transforms/ObjCARC/ObjCARC.cpp
+++ b/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);
   }
 
diff --git a/llvm/lib/Transforms/ObjCARC/ObjCARC.h b/llvm/lib/Transforms/ObjCARC/ObjCARC.h
index f4d7c92..d0bff00 100644
--- a/llvm/lib/Transforms/ObjCARC/ObjCARC.h
+++ b/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 @@ private:
   /// A map of inserted retainRV/claimRV calls to annotated calls/invokes.
   DenseMap<CallInst *, CallBase *> RVCalls;
 
+  ARCRuntimeEntryPoints &EP;
   bool ContractPass;
+  bool UseClaimRV;
 };
 
 } // end namespace objcarc
diff --git a/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp b/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
index e11748b..86d7e2f 100644
--- a/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
+++ b/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);
diff --git a/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp b/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp
index 9d7f5e6..65e2339 100644
--- a/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp
+++ b/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()
diff --git a/llvm/test/Transforms/ObjCARC/contract-attached-call-retain-to-claim.ll b/llvm/test/Transforms/ObjCARC/contract-attached-call-retain-to-claim.ll
new file mode 100644
index 0000000..d0b8ce9
--- /dev/null
+++ b/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)