[Runtime] Eliminate stack frames in swift_retain and swift_bridgeObjectRetain on ARM64.

Rearrange the slow paths a bit to make them tail calls, which allows the compiler to emit these functions without frames.

Clang is happy to emit frameless functions on ARM64 if no stack space is needed on all execution paths. However, when there's a fast path which doesn't need stack space, and a slow path which does, clang emits code that pushes a stack frame and then decides which path to take. This is fine, but it means we're paying more than we'd like to on the fast path.

We can work around that by manually outlining the slow path, and ensuring that it's invoked with a tail call. Then the original function doesn't need a stack frame on any path and clang omits the stack frame.

We tweak RefCounts::increment to return the object it's being called on, which allows `swift_retain` to tail-call it. We manually outline the objc_retain call in swift_bridgeObjectRetain, which allows the swift_retain path to be frameless.

rdar://101764509

Author: mikeash

stdlib/public/SwiftShims/swift/shims/RefCount.h

@@ -703,7 +703,7 @@ class RefCounts {
   // Out-of-line slow paths.
 
   SWIFT_NOINLINE
-  void incrementSlow(RefCountBits oldbits, uint32_t inc) SWIFT_CC(PreserveMost);
+  HeapObject *incrementSlow(RefCountBits oldbits, uint32_t inc);
 
   SWIFT_NOINLINE
   void incrementNonAtomicSlow(RefCountBits oldbits, uint32_t inc);
@@ -799,14 +799,18 @@ class RefCounts {
   }
 
   // Increment the reference count.
+  //
+  // This returns the enclosing HeapObject so that it the result of this call
+  // can be directly returned from swift_retain. This makes the call to
+  // incrementSlow() a tail call.
   SWIFT_ALWAYS_INLINE
-  void increment(uint32_t inc = 1) {
+  HeapObject *increment(uint32_t inc = 1) {
     auto oldbits = refCounts.load(SWIFT_MEMORY_ORDER_CONSUME);
 
     // constant propagation will remove this in swift_retain, it should only
     // be present in swift_retain_n
     if (inc != 1 && oldbits.isImmortal(true)) {
-      return;
+      return getHeapObject();
     }
 
     RefCountBits newbits;
@@ -815,11 +819,12 @@ class RefCounts {
       bool fast = newbits.incrementStrongExtraRefCount(inc);
       if (SWIFT_UNLIKELY(!fast)) {
         if (oldbits.isImmortal(false))
-          return;
+          return getHeapObject();
         return incrementSlow(oldbits, inc);
       }
     } while (!refCounts.compare_exchange_weak(oldbits, newbits,
                                               std::memory_order_relaxed));
+    return getHeapObject();
   }
 
   SWIFT_ALWAYS_INLINE

stdlib/public/SwiftShims/swift/shims/Visibility.h

@@ -109,6 +109,12 @@
 #define SWIFT_WEAK_IMPORT
 #endif
 
+#if __has_attribute(musttail)
+#define SWIFT_MUSTTAIL [[clang::musttail]]
+#else
+#define SWIFT_MUSTTAIL
+#endif
+
 // Define the appropriate attributes for sharing symbols across
 // image (executable / shared-library) boundaries.
 //

stdlib/public/runtime/HeapObject.cpp

@@ -342,8 +342,12 @@ _swift_release_dealloc(HeapObject *object);
 SWIFT_ALWAYS_INLINE
 static HeapObject *_swift_retain_(HeapObject *object) {
   SWIFT_RT_TRACK_INVOCATION(object, swift_retain);
-  if (isValidPointerForNativeRetain(object))
-    object->refCounts.increment(1);
+  if (isValidPointerForNativeRetain(object)) {
+    // Return the result of increment() to make the eventual call to
+    // incrementSlow a tail call, which avoids pushing a stack frame on the fast
+    // path on ARM64.
+    return object->refCounts.increment(1);
+  }
   return object;
 }
 

stdlib/public/runtime/RefCount.cpp

@@ -15,10 +15,10 @@
 namespace swift {
 
 template <typename RefCountBits>
-void RefCounts<RefCountBits>::incrementSlow(RefCountBits oldbits,
-                                            uint32_t n) {
+HeapObject *RefCounts<RefCountBits>::incrementSlow(RefCountBits oldbits,
+                                                   uint32_t n) {
   if (oldbits.isImmortal(false)) {
-    return;
+    return getHeapObject();
   }
   else if (oldbits.hasSideTable()) {
     // Out-of-line slow path.
@@ -29,9 +29,14 @@ void RefCounts<RefCountBits>::incrementSlow(RefCountBits oldbits,
     // Retain count overflow.
     swift::swift_abortRetainOverflow();
   }
+  return getHeapObject();
 }
-template void RefCounts<InlineRefCountBits>::incrementSlow(InlineRefCountBits oldbits, uint32_t n);
-template void RefCounts<SideTableRefCountBits>::incrementSlow(SideTableRefCountBits oldbits, uint32_t n);
+template HeapObject *
+RefCounts<InlineRefCountBits>::incrementSlow(InlineRefCountBits oldbits,
+                                             uint32_t n);
+template HeapObject *
+RefCounts<SideTableRefCountBits>::incrementSlow(SideTableRefCountBits oldbits,
+                                                uint32_t n);
 
 template <typename RefCountBits>
 void RefCounts<RefCountBits>::incrementNonAtomicSlow(RefCountBits oldbits,

stdlib/public/runtime/SwiftObject.mm

@@ -574,6 +574,20 @@ static bool isBridgeObjectTaggedPointer(void *object) {
   return (void*)(uintptr_t(object) & ~unTaggedNonNativeBridgeObjectBits);
 }
 
+#if SWIFT_OBJC_INTEROP
+#if __arm64__
+// Marking this as noinline allows swift_bridgeObjectRetain to avoid emitting
+// a stack frame for the swift_retain path on ARM64. It makes for worse codegen
+// on x86-64, though, so limit it to ARM64.
+SWIFT_NOINLINE
+#endif
+static void *objcRetainAndReturn(void *object) {
+  auto const objectRef = toPlainObject_unTagged_bridgeObject(object);
+  objc_retain(static_cast<id>(objectRef));
+  return object;
+}
+#endif
+
 void *swift::swift_bridgeObjectRetain(void *object) {
 #if SWIFT_OBJC_INTEROP
   if (isObjCTaggedPointer(object) || isBridgeObjectTaggedPointer(object))
@@ -584,14 +598,18 @@ static bool isBridgeObjectTaggedPointer(void *object) {
 
 #if SWIFT_OBJC_INTEROP
   if (!isNonNative_unTagged_bridgeObject(object)) {
-    swift_retain(static_cast<HeapObject *>(objectRef));
-    return object;
+    return swift_retain(static_cast<HeapObject *>(objectRef));
   }
-  objc_retain(static_cast<id>(objectRef));
-  return object;
+
+  // Put the call to objc_retain in a separate function, tail-called here. This
+  // allows the fast path of swift_bridgeObjectRetain to avoid creating a stack
+  // frame on ARM64. We can't directly tail-call objc_retain, because
+  // swift_bridgeObjectRetain returns the pointer with objectPointerIsObjCBit
+  // set, so we have to make a non-tail call and then return the value with the
+  // bit set.
+  SWIFT_MUSTTAIL return objcRetainAndReturn(object);
 #else
-  swift_retain(static_cast<HeapObject *>(objectRef));
-  return object;
+  return swift_retain(static_cast<HeapObject *>(objectRef));
 #endif
 }