Reland "[ObjCARC][Contract] Optimize bundled RetainRV to ClaimRV"

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/CMakeLists.txt

@@ -22,5 +22,6 @@ add_llvm_component_library(LLVMObjCARCOpts
   Analysis
   Core
   Support
+  TargetParser
   TransformUtils
   )

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

@@ -2423,7 +2423,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)