Merge pull request #85044 from mikeash/emit-into-client-retain-release

[IRGen][Runtime] Add emit-into-client retain/release calls for Darwin ARM64.

Author: mikeash

CMakeLists.txt

@@ -271,6 +271,19 @@ option(SWIFT_BUILD_STDLIB_CXX_MODULE
   "If not building stdlib, controls whether to build the Cxx module"
   TRUE)
 
+# The swiftClientRetainRelease library is currently only available for Darwin
+# platforms.
+if(SWIFT_HOST_VARIANT_SDK IN_LIST SWIFT_DARWIN_PLATFORMS)
+  # Off by default everywhere for now.
+  option(SWIFT_BUILD_CLIENT_RETAIN_RELEASE
+    "Build the swiftClientRetainRelease library"
+    FALSE)
+else()
+  option(SWIFT_BUILD_CLIENT_RETAIN_RELEASE
+    "Build the swiftClientRetainRelease library"
+    FALSE)
+endif()
+
 # In many cases, the CMake build system needs to determine whether to include
 # a directory, or perform other actions, based on whether the stdlib or SDK is
 # being built at all -- statically or dynamically. Please note that these

include/swift/AST/IRGenOptions.h

@@ -556,6 +556,9 @@ class IRGenOptions {
   // Whether to emit mergeable or non-mergeable traps.
   unsigned MergeableTraps : 1;
 
+  /// Enable the use of swift_retain/releaseClient functions.
+  unsigned EnableClientRetainRelease : 1;
+
   /// The number of threads for multi-threaded code generation.
   unsigned NumThreads = 0;
 
@@ -683,6 +686,7 @@ class IRGenOptions {
         EmitAsyncFramePushPopMetadata(true), EmitTypeMallocForCoroFrame(true),
         AsyncFramePointerAll(false), UseProfilingMarkerThunks(false),
         UseCoroCCX8664(false), UseCoroCCArm64(false), MergeableTraps(false),
+        EnableClientRetainRelease(false),
         DebugInfoForProfiling(false), CmdArgs(),
         SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All),

include/swift/Option/FrontendOptions.td

@@ -1477,6 +1477,13 @@ def mergeable_traps :
   Flag<["-"], "mergeable-traps">,
   HelpText<"Emit mergeable traps even in optimized builds">;
 
+def enable_client_retain_release :
+  Flag<["-"], "enable-client-retain-release">,
+  HelpText<"Enable use of swift_retain/releaseClient functions">;
+def disable_client_retain_release :
+  Flag<["-"], "disable-client-retain-release">,
+  HelpText<"Disable use of swift_retain/releaseClient functions">;
+
 def enable_new_llvm_pass_manager :
   Flag<["-"], "enable-new-llvm-pass-manager">,
   HelpText<"Enable the new llvm pass manager">;

include/swift/Runtime/Config.h

@@ -248,6 +248,17 @@ extern uintptr_t __COMPATIBILITY_LIBRARIES_CANNOT_CHECK_THE_IS_SWIFT_BIT_DIRECTL
 // so changing this value is not sufficient.
 #define SWIFT_DEFAULT_LLVM_CC llvm::CallingConv::C
 
+// Define the calling convention for refcounting functions for targets where it
+// differs from the standard calling convention. Currently this is only used for
+// swift_retain, swift_release, and some internal helper functions that they
+// call.
+#if defined(__aarch64__)
+#define SWIFT_REFCOUNT_CC SWIFT_CC_PreserveMost
+#define SWIFT_REFCOUNT_CC_PRESERVEMOST 1
+#else
+#define SWIFT_REFCOUNT_CC
+#endif
+
 /// Should we use absolute function pointers instead of relative ones?
 /// WebAssembly target uses it by default.
 #ifndef SWIFT_COMPACT_ABSOLUTE_FUNCTION_POINTER

include/swift/Runtime/CustomRRABI.h

@@ -59,6 +59,9 @@ namespace swift {
 template <typename Ret, typename Param>
 Param returnTypeHelper(Ret (*)(Param)) {}
 
+template <typename Ret, typename Param>
+Param returnTypeHelper(SWIFT_REFCOUNT_CC Ret (*)(Param)) {}
+
 #if defined(__LP64__) || defined(_LP64)
 #define REGISTER_SUBSTITUTION_PREFIX ""
 #define REGISTER_PREFIX "x"

include/swift/Runtime/HeapObject.h

@@ -137,6 +137,7 @@ HeapObject* swift_allocEmptyBox();
 /// It may also prove worthwhile to have this use a custom CC
 /// which preserves a larger set of registers.
 SWIFT_RUNTIME_EXPORT
+SWIFT_REFCOUNT_CC
 HeapObject *swift_retain(HeapObject *object);
 
 SWIFT_RUNTIME_EXPORT
@@ -173,6 +174,7 @@ bool swift_isDeallocating(HeapObject *object);
 ///      - maybe a variant that can assume a non-null object
 /// It's unlikely that a custom CC would be beneficial here.
 SWIFT_RUNTIME_EXPORT
+SWIFT_REFCOUNT_CC
 void swift_release(HeapObject *object);
 
 SWIFT_RUNTIME_EXPORT

include/swift/Runtime/RuntimeFunctions.def

@@ -220,6 +220,22 @@ FUNCTION(NativeStrongRelease, Swift, swift_release, C_CC, AlwaysAvailable,
          EFFECT(RuntimeEffect::RefCounting, RuntimeEffect::Deallocating),
          UNKNOWN_MEMEFFECTS)
 
+// void *swift_retain(void *ptr);
+FUNCTION(NativeStrongRetainClient, Swift, swift_retainClient, SwiftClientRR_CC, AlwaysAvailable,
+         RETURNS(RefCountedPtrTy),
+         ARGS(RefCountedPtrTy),
+         ATTRS(NoUnwind, FirstParamReturned, WillReturn),
+         EFFECT(RuntimeEffect::RefCounting),
+         UNKNOWN_MEMEFFECTS)
+
+// void swift_release(void *ptr);
+FUNCTION(NativeStrongReleaseClient, Swift, swift_releaseClient, SwiftClientRR_CC, AlwaysAvailable,
+         RETURNS(VoidTy),
+         ARGS(RefCountedPtrTy),
+         ATTRS(NoUnwind),
+         EFFECT(RuntimeEffect::RefCounting, RuntimeEffect::Deallocating),
+         UNKNOWN_MEMEFFECTS)
+
 // void *swift_retain_n(void *ptr, int32_t n);
 FUNCTION(NativeStrongRetainN, Swift, swift_retain_n, C_CC, AlwaysAvailable,
          RETURNS(RefCountedPtrTy),
@@ -420,6 +436,24 @@ FUNCTION(BridgeObjectStrongRelease, Swift, swift_bridgeObjectRelease,
          EFFECT(RuntimeEffect::RefCounting, RuntimeEffect::Deallocating),
          UNKNOWN_MEMEFFECTS)
 
+// void *swift_bridgeObjectRetainClient(void *ptr);
+FUNCTION(BridgeObjectStrongRetainClient, Swift, swift_bridgeObjectRetainClient,
+         SwiftClientRR_CC, AlwaysAvailable,
+         RETURNS(BridgeObjectPtrTy),
+         ARGS(BridgeObjectPtrTy),
+         ATTRS(NoUnwind, FirstParamReturned),
+         EFFECT(RuntimeEffect::RefCounting),
+         UNKNOWN_MEMEFFECTS)
+
+// void *swift_bridgeObjectReleaseClient(void *ptr);
+FUNCTION(BridgeObjectStrongReleaseClient, Swift, swift_bridgeObjectReleaseClient,
+         SwiftClientRR_CC, AlwaysAvailable,
+         RETURNS(VoidTy),
+         ARGS(BridgeObjectPtrTy),
+         ATTRS(NoUnwind),
+         EFFECT(RuntimeEffect::RefCounting, RuntimeEffect::Deallocating),
+         UNKNOWN_MEMEFFECTS)
+
 // void *swift_nonatomic_bridgeObjectRetain(void *ptr);
 FUNCTION(NonAtomicBridgeObjectStrongRetain, Swift, swift_nonatomic_bridgeObjectRetain,
          C_CC, AlwaysAvailable,

lib/Frontend/CompilerInvocation.cpp

@@ -3989,6 +3989,11 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
 
   Opts.MergeableTraps = Args.hasArg(OPT_mergeable_traps);
 
+  Opts.EnableClientRetainRelease =
+    Args.hasFlag(OPT_enable_client_retain_release,
+                 OPT_disable_client_retain_release,
+                 Opts.EnableClientRetainRelease);
+
   Opts.EnableObjectiveCProtocolSymbolicReferences =
     Args.hasFlag(OPT_enable_objective_c_protocol_symbolic_references,
                  OPT_disable_objective_c_protocol_symbolic_references,

lib/IRGen/GenHeap.cpp

@@ -997,11 +997,16 @@ void IRGenFunction::emitNativeStrongRetain(llvm::Value *value,
     value = Builder.CreateBitCast(value, IGM.RefCountedPtrTy);
 
   // Emit the call.
-  llvm::CallInst *call = Builder.CreateCall(
-      (atomicity == Atomicity::Atomic)
-          ? IGM.getNativeStrongRetainFunctionPointer()
-          : IGM.getNativeNonAtomicStrongRetainFunctionPointer(),
-      value);
+  FunctionPointer function;
+  if (atomicity == Atomicity::Atomic &&
+      IGM.TargetInfo.HasSwiftClientRRLibrary &&
+      getOptions().EnableClientRetainRelease)
+    function = IGM.getNativeStrongRetainClientFunctionPointer();
+  else if (atomicity == Atomicity::Atomic)
+    function = IGM.getNativeStrongRetainFunctionPointer();
+  else
+    function = IGM.getNativeNonAtomicStrongRetainFunctionPointer();
+  llvm::CallInst *call = Builder.CreateCall(function, value);
   call->setDoesNotThrow();
   call->addParamAttr(0, llvm::Attribute::Returned);
 }
@@ -1257,10 +1262,16 @@ void IRGenFunction::emitNativeStrongRelease(llvm::Value *value,
                                             Atomicity atomicity) {
   if (doesNotRequireRefCounting(value))
     return;
-  emitUnaryRefCountCall(*this, (atomicity == Atomicity::Atomic)
-                                   ? IGM.getNativeStrongReleaseFn()
-                                   : IGM.getNativeNonAtomicStrongReleaseFn(),
-                        value);
+  llvm::Constant *function;
+  if (atomicity == Atomicity::Atomic &&
+      IGM.TargetInfo.HasSwiftClientRRLibrary &&
+      getOptions().EnableClientRetainRelease)
+    function = IGM.getNativeStrongReleaseClientFn();
+  else if (atomicity == Atomicity::Atomic)
+    function = IGM.getNativeStrongReleaseFn();
+  else
+    function = IGM.getNativeNonAtomicStrongReleaseFn();
+  emitUnaryRefCountCall(*this, function, value);
 }
 
 void IRGenFunction::emitNativeSetDeallocating(llvm::Value *value) {
@@ -1353,20 +1364,30 @@ void IRGenFunction::emitUnknownStrongRelease(llvm::Value *value,
 
 void IRGenFunction::emitBridgeStrongRetain(llvm::Value *value,
                                            Atomicity atomicity) {
-  emitUnaryRefCountCall(*this,
-                        (atomicity == Atomicity::Atomic)
-                            ? IGM.getBridgeObjectStrongRetainFn()
-                            : IGM.getNonAtomicBridgeObjectStrongRetainFn(),
-                        value);
+  llvm::Constant *function;
+  if (atomicity == Atomicity::Atomic &&
+      IGM.TargetInfo.HasSwiftClientRRLibrary &&
+      getOptions().EnableClientRetainRelease)
+    function = IGM.getBridgeObjectStrongRetainClientFn();
+  else if (atomicity == Atomicity::Atomic)
+    function = IGM.getBridgeObjectStrongRetainFn();
+  else
+    function = IGM.getNonAtomicBridgeObjectStrongRetainFn();
+  emitUnaryRefCountCall(*this, function, value);
 }
 
 void IRGenFunction::emitBridgeStrongRelease(llvm::Value *value,
                                             Atomicity atomicity) {
-  emitUnaryRefCountCall(*this,
-                        (atomicity == Atomicity::Atomic)
-                            ? IGM.getBridgeObjectStrongReleaseFn()
-                            : IGM.getNonAtomicBridgeObjectStrongReleaseFn(),
-                        value);
+  llvm::Constant *function;
+  if (atomicity == Atomicity::Atomic &&
+      IGM.TargetInfo.HasSwiftClientRRLibrary &&
+      getOptions().EnableClientRetainRelease)
+    function = IGM.getBridgeObjectStrongReleaseClientFn();
+  else if (atomicity == Atomicity::Atomic)
+    function = IGM.getBridgeObjectStrongReleaseFn();
+  else
+    function = IGM.getNonAtomicBridgeObjectStrongReleaseFn();
+  emitUnaryRefCountCall(*this, function, value);
 }
 
 void IRGenFunction::emitErrorStrongRetain(llvm::Value *value) {

lib/IRGen/IRGenModule.cpp

@@ -564,8 +564,9 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
   InvariantMetadataID = getLLVMContext().getMDKindID("invariant.load");
   InvariantNode = llvm::MDNode::get(getLLVMContext(), {});
   DereferenceableID = getLLVMContext().getMDKindID("dereferenceable");
-  
+
   C_CC = getOptions().PlatformCCallingConvention;
+  SwiftClientRR_CC = llvm::CallingConv::PreserveMost;
   // TODO: use "tinycc" on platforms that support it
   DefaultCC = SWIFT_DEFAULT_LLVM_CC;
 
@@ -1732,6 +1733,11 @@ void IRGenModule::addLinkLibraries() {
     registerLinkLibrary(
         LinkLibrary{"objc", LibraryKind::Library, /*static=*/false});
 
+  if (TargetInfo.HasSwiftClientRRLibrary &&
+      getOptions().EnableClientRetainRelease)
+    registerLinkLibrary(LinkLibrary{"swiftClientRetainRelease",
+                                    LibraryKind::Library, /*static=*/true});
+
   // If C++ interop is enabled, add -lc++ on Darwin and -lstdc++ on linux.
   // Also link with C++ bridging utility module (Cxx) and C++ stdlib overlay
   // (std) if available.