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

This is currently disabled by default. Building the client library can be enabled with the CMake option SWIFT_BUILD_CLIENT_RETAIN_RELEASE, and using the library can be enabled with the flags -Xfrontend -enable-client-retain-release.

To improve retain/release performance, we build a static library containing optimized implementations of the fast paths of swift_retain, swift_release, and the corresponding bridgeObject functions. This avoids going through a stub to make a cross-library call.

IRGen gains awareness of these new functions and emits calls to them when the functionality is enabled and the target supports them. Two options are added to force use of them on or off: -enable-client-retain-release and -disable-client-retain-release. When enabled, the compiler auto-links the static library containing the implementations.

The new calls also use LLVM's preserve_most calling convention. Since retain/release doesn't need a large number of scratch registers, this is mostly harmless for the implementation, while allowing callers to improve code size and performance by spilling fewer registers around refcounting calls. (Experiments with an even more aggressive calling convention preserving x2 and up showed an insignificant savings in code size, so preserve_most seems to be a good middle ground.)

Since the implementations are embedded into client binaries, any change in the runtime's refcounting implementation needs to stay compatible with this new fast path implementation. This is ensured by having the implementation use a runtime-provided mask to check whether it can proceed into its fast path. The mask is provided as the address of the absolute symbol _swift_retainRelease_slowpath_mask_v1. If that mask ANDed with the object's current refcount field is non-zero, then we take the slow path. A future runtime that changes the refcounting implementation can adjust this mask to match, or set the mask to all 1s to disable the old embedded fast path entirely (as long as the new representation never uses 0 as a valid refcount field value).

As part of this work, the overall approach for bridgeObjectRetain is changed slightly. Previously, it would mask off the spare bits from the native pointer and then call through to swift_retain. This either lost the spare bits in the return value (when tail calling swift_retain) which is problematic since it's supposed to return its parameter, or it required pushing a stack frame which is inefficient. Now, swift_retain takes on the responsibility of masking off spare bits from the parameter and preserving them in the return value. This is a trivial addition to the fast path (just a quick mask and an extra register for saving the original value) and makes bridgeObjectRetain quite a bit more efficient when implemented correctly to return the exact value it was passed.

The runtime's implementations of swift_retain/release are now also marked as preserve_most so that they can be tail called from the client library. preserve_most is compatible with callers expecting the standard calling convention so this doesn't break any existing clients. Some ugly tricks were needed to prevent the compiler from creating unnecessary stack frames with the new calling convention. Avert your eyes.

To allow back deployment, the runtime now has aliases for these functions called swift_retain_preservemost and swift_release_preservemost. The client library brings weak definitions of these functions that save the extra registers and call through to swift_retain/release. This allows them to work correctly on older runtimes, with a small performance penalty, while still running at full speed on runtimes that have the new preservemost symbols.

Although this is only supported on Darwin at the moment, it shouldn't be too much work to adapt it to other ARM64 targets. We need to ensure the assembly plays nice with the other platforms' assemblers, and make sure the implementation is correct for the non-ObjC-interop case.

rdar://122595871

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

@@ -548,6 +548,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;
 
@@ -675,6 +678,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

@@ -3983,6 +3983,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;
 
@@ -1730,6 +1731,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.