Merge pull request #61794 from mikeash/retain-stack-frames

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

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
 }