Merge pull request #36975 from eeckstein/optimize-async-let

concurrency: make the startAsyncLet closure no-escaping

Author: eeckstein

include/swift/ABI/Task.h

@@ -582,13 +582,13 @@ class FutureAsyncContext : public AsyncContext {
 /// This matches the ABI of a closure `() async throws -> ()`
 using AsyncVoidClosureEntryPoint =
   SWIFT_CC(swiftasync)
-  void (SWIFT_ASYNC_CONTEXT AsyncContext *, SWIFT_CONTEXT HeapObject *);
+  void (SWIFT_ASYNC_CONTEXT AsyncContext *, SWIFT_CONTEXT void *);
 
 /// This matches the ABI of a closure `<T>() async throws -> T`
 using AsyncGenericClosureEntryPoint =
     SWIFT_CC(swiftasync)
     void(OpaqueValue *,
-         SWIFT_ASYNC_CONTEXT AsyncContext *, SWIFT_CONTEXT HeapObject *);
+         SWIFT_ASYNC_CONTEXT AsyncContext *, SWIFT_CONTEXT void *);
 
 /// This matches the ABI of the resume function of a closure
 ///  `() async throws -> ()`.
@@ -602,7 +602,7 @@ class AsyncContextPrefix {
   // passing the closure context instead of via the async context)
   AsyncVoidClosureEntryPoint *__ptrauth_swift_task_resume_function
       asyncEntryPoint;
-   HeapObject *closureContext;
+  void *closureContext;
   SwiftError *errorResult;
 };
 
@@ -615,7 +615,7 @@ class FutureAsyncContextPrefix {
   // passing the closure context instead of via the async context)
   AsyncGenericClosureEntryPoint *__ptrauth_swift_task_resume_function
       asyncEntryPoint;
-  HeapObject *closureContext;
+  void *closureContext;
   SwiftError *errorResult;
 };
 

include/swift/AST/Builtins.def

@@ -793,6 +793,9 @@ BUILTIN_MISC_OPERATION(StartAsyncLet, "startAsyncLet", "", Special)
 /// asyncLetEnd(): (Builtin.RawPointer) -> Void
 ///
 /// Ends and destroys an async-let.
+/// The ClosureLifetimeFixup pass adds a second operand to the builtin to
+/// ensure that optimizations keep the stack-allocated closure arguments alive
+/// until the endAsyncLet.
 BUILTIN_MISC_OPERATION_WITH_SILGEN(EndAsyncLet, "endAsyncLet", "", Special)
 
 /// createAsyncTaskFuture(): (

include/swift/Runtime/Concurrency.h

@@ -281,15 +281,15 @@ bool swift_taskGroup_isEmpty(TaskGroup *group);
 /// \code
 /// func swift_asyncLet_start<T>(
 ///     _ alet: Builtin.RawPointer,
-///     operation: __owned @Sendable @escaping () async throws -> T
+///     operation: __owned @Sendable () async throws -> T
 /// )
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_asyncLet_start(
     AsyncLet *alet,
     const Metadata *futureResultType,
     void *closureEntryPoint,
-    HeapObject * /* +1 */ closureContext);
+    void *closureContext);
 
 /// This matches the ABI of a closure `<T>(Builtin.RawPointer) async -> T`
 using AsyncLetWaitSignature =

include/swift/Runtime/RuntimeFunctions.def

@@ -1682,7 +1682,7 @@ FUNCTION(DefaultActorDeallocateResilient,
 ///     AsyncLet *alet,
 ///     const Metadata *futureResultType,
 ///     void *closureEntryPoint,
-///     HeapObject *closureContext
+///     OpaqueValue *closureContext
 /// );
 FUNCTION(AsyncLetStart,
          swift_asyncLet_start, SwiftCC,
@@ -1691,7 +1691,7 @@ FUNCTION(AsyncLetStart,
          ARGS(SwiftAsyncLetPtrTy, // AsyncLet*, alias for Int8PtrTy
               TypeMetadataPtrTy,  // futureResultType
               Int8PtrTy,          // closureEntry
-              RefCountedPtrTy,    // closureContext
+              OpaquePtrTy,        // closureContext
          ),
          ATTRS(NoUnwind, ArgMemOnly))
 

lib/AST/Builtins.cpp

@@ -1496,7 +1496,7 @@ static ValueDecl *getStartAsyncLet(ASTContext &ctx, Identifier id) {
   builder.addParameter(makeConcrete(OptionalType::get(ctx.TheRawPointerType)));
 
   // operation async function pointer: () async throws -> T
-  auto extInfo = ASTExtInfoBuilder().withAsync().withThrows().build();
+  auto extInfo = ASTExtInfoBuilder().withAsync().withThrows().withNoEscape().build();
   builder.addParameter(
       makeConcrete(FunctionType::get({ }, genericParam, extInfo)));
 

lib/IRGen/GenBuiltin.cpp

@@ -244,6 +244,9 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
 
   if (Builtin.ID == BuiltinValueKind::EndAsyncLet) {
     emitEndAsyncLet(IGF, args.claimNext());
+    // Ignore a second operand which is inserted by ClosureLifetimeFixup and
+    // only used for dependency tracking.
+    (void)args.claimAll();
     return;
   }
 

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -363,6 +363,61 @@ static SILValue insertMarkDependenceForCapturedArguments(PartialApplyInst *pai,
   return curr;
 }
 
+/// Returns the (single) "endAsyncLet" builtin if \p startAsyncLet is a
+/// "startAsyncLet" builtin.
+static BuiltinInst *getEndAsyncLet(BuiltinInst *startAsyncLet) {
+  if (startAsyncLet->getBuiltinKind() != BuiltinValueKind::StartAsyncLet)
+    return nullptr;
+
+  BuiltinInst *endAsyncLet = nullptr;
+  for (Operand *op : startAsyncLet->getUses()) {
+    auto *endBI = dyn_cast<BuiltinInst>(op->getUser());
+    if (endBI && endBI->getBuiltinKind() == BuiltinValueKind::EndAsyncLet) {
+      // At this stage of the pipeline, it's always the case that a
+      // startAsyncLet has an endAsyncLet: that's how SILGen generates it.
+      // Just to be on the safe side, do this check.
+      if (endAsyncLet)
+        return nullptr;
+      endAsyncLet = endBI;
+    }
+  }
+  return endAsyncLet;
+}
+
+/// Call the \p insertFn with a builder at all insertion points after
+/// a closure is used by \p closureUser.
+static void insertAfterClosureUser(SILInstruction *closureUser,
+                                   function_ref<void(SILBuilder &)> insertFn) {
+  // Don't insert any destroy or deallocation right before an unreachable.
+  // It's not needed an will only add up to code size.
+  auto insertAtNonUnreachable = [&](SILBuilder &builder) {
+    if (isa<UnreachableInst>(builder.getInsertionPoint()))
+      return;
+    insertFn(builder);
+  };
+
+  if (auto *startAsyncLet = dyn_cast<BuiltinInst>(closureUser)) {
+    BuiltinInst *endAsyncLet = getEndAsyncLet(startAsyncLet);
+    assert(endAsyncLet);
+    SILBuilderWithScope builder(std::next(endAsyncLet->getIterator()));
+    insertAtNonUnreachable(builder);
+    return;
+  }
+  FullApplySite fas = FullApplySite::isa(closureUser);
+  assert(fas);
+  fas.insertAfterFullEvaluation(insertAtNonUnreachable);
+}
+
+static SILValue skipConvert(SILValue v) {
+  auto *cvt = dyn_cast<ConvertFunctionInst>(v);
+  if (!cvt)
+    return v;
+  auto *pa = dyn_cast<PartialApplyInst>(cvt->getOperand());
+  if (!pa || !pa->hasOneUse())
+    return v;
+  return pa;
+}
+
 /// Rewrite a partial_apply convert_escape_to_noescape sequence with a single
 /// apply/try_apply user to a partial_apply [stack] terminated with a
 /// dealloc_stack placed after the apply.
@@ -386,12 +441,15 @@ static SILValue insertMarkDependenceForCapturedArguments(PartialApplyInst *pai,
 /// dealloc_stack still needs to be balanced with other dealloc_stacks i.e the
 /// caller needs to use the StackNesting utility to update the dealloc_stack
 /// nesting.
-static bool tryRewriteToPartialApplyStack(
-    SILLocation &loc, PartialApplyInst *origPA,
-    ConvertEscapeToNoEscapeInst *cvt, SILInstruction *singleApplyUser,
-    SILBasicBlock::iterator &advanceIfDelete,
+static SILValue tryRewriteToPartialApplyStack(
+    ConvertEscapeToNoEscapeInst *cvt,
+    SILInstruction *closureUser, SILBasicBlock::iterator &advanceIfDelete,
     llvm::DenseMap<SILInstruction *, SILInstruction *> &memoized) {
-  
+
+  auto *origPA = dyn_cast<PartialApplyInst>(skipConvert(cvt->getOperand()));
+  if (!origPA)
+    return SILValue();
+
   auto *convertOrPartialApply = cast<SingleValueInstruction>(origPA);
   if (cvt->getOperand() != origPA)
     convertOrPartialApply = cast<ConvertFunctionInst>(cvt->getOperand());
@@ -415,7 +473,7 @@ static bool tryRewriteToPartialApplyStack(
       continue;
     }
     if (singleNonDebugNonRefCountUser)
-      return false;
+      return SILValue();
     singleNonDebugNonRefCountUser = user;
   }
 
@@ -480,39 +538,15 @@ static bool tryRewriteToPartialApplyStack(
     }
   }
 
-  // Insert destroys of arguments after the apply and the dealloc_stack.
-  if (auto *apply = dyn_cast<ApplyInst>(singleApplyUser)) {
-    auto insertPt = std::next(SILBasicBlock::iterator(apply));
-    // Don't insert dealloc_stacks at unreachable.
-    if (isa<UnreachableInst>(*insertPt))
-      return true;
-    SILBuilderWithScope b3(insertPt);
-    b3.createDeallocStack(loc, newPA);
-    insertDestroyOfCapturedArguments(newPA, b3);
+  // Insert destroys of arguments after the closure user and the dealloc_stack.
+  insertAfterClosureUser(closureUser, [newPA](SILBuilder &builder) {
+    auto loc = RegularLocation(builder.getInsertionPointLoc());
+    builder.createDeallocStack(loc, newPA);
+    insertDestroyOfCapturedArguments(newPA, builder);
     // dealloc_stack of the in_guaranteed capture is inserted
-    insertDeallocOfCapturedArguments(newPA, b3);
-  } else if (auto *tai = dyn_cast<TryApplyInst>(singleApplyUser)) {
-    for (auto *succBB : tai->getSuccessorBlocks()) {
-      SILBuilderWithScope b3(succBB->begin());
-      b3.createDeallocStack(loc, newPA);
-      insertDestroyOfCapturedArguments(newPA, b3);
-      // dealloc_stack of the in_guaranteed capture is inserted
-      insertDeallocOfCapturedArguments(newPA, b3);
-    }
-  } else {
-    llvm_unreachable("Unknown FullApplySite instruction kind");
-  }
-  return true;
-}
-
-static SILValue skipConvert(SILValue v) {
-  auto *cvt = dyn_cast<ConvertFunctionInst>(v);
-  if (!cvt)
-    return v;
-  auto *pa = dyn_cast<PartialApplyInst>(cvt->getOperand());
-  if (!pa || !pa->hasOneUse())
-    return v;
-  return pa;
+    insertDeallocOfCapturedArguments(newPA, builder);
+  });
+  return closure;
 }
 
 static bool tryExtendLifetimeToLastUse(
@@ -525,45 +559,52 @@ static bool tryExtendLifetimeToLastUse(
   if (!singleUser)
     return false;
 
-  // Handle an apply.
-  if (auto singleApplyUser = FullApplySite::isa(singleUser)) {
-    // FIXME: Don't know how-to handle begin_apply/end_apply yet.
-    if (isa<BeginApplyInst>(singleApplyUser.getInstruction())) {
+  // Handle apply instructions and startAsyncLet.
+  BuiltinInst *endAsyncLet = nullptr;
+  if (FullApplySite::isa(singleUser)) {
+    // TODO: Enable begin_apply/end_apply. It should work, but is not tested yet.
+    if (isa<BeginApplyInst>(singleUser))
       return false;
-    }
-
-    auto loc = RegularLocation::getAutoGeneratedLocation();
-    auto origPA = dyn_cast<PartialApplyInst>(skipConvert(cvt->getOperand()));
-    if (origPA && tryRewriteToPartialApplyStack(
-                      loc, origPA, cvt, singleApplyUser.getInstruction(),
-                      advanceIfDelete, memoized))
-      return true;
+  } else if (auto *bi = dyn_cast<BuiltinInst>(singleUser)) {
+    endAsyncLet = getEndAsyncLet(bi);
+    if (!endAsyncLet)
+      return false;
+  } else {
+    return false;
+  }
 
-    // Insert a copy at the convert_escape_to_noescape [not_guaranteed] and
-    // change the instruction to the guaranteed form.
-    auto escapingClosure = cvt->getOperand();
-    auto *closureCopy =
-        SILBuilderWithScope(cvt).createCopyValue(loc, escapingClosure);
-    cvt->setLifetimeGuaranteed();
-    cvt->setOperand(closureCopy);
-
-    // Insert a destroy after the apply.
-    if (auto *apply = dyn_cast<ApplyInst>(singleApplyUser.getInstruction())) {
-      auto insertPt = std::next(SILBasicBlock::iterator(apply));
-      SILBuilderWithScope(insertPt).createDestroyValue(loc, closureCopy);
-
-    } else if (auto *tai =
-                   dyn_cast<TryApplyInst>(singleApplyUser.getInstruction())) {
-      for (auto *succBB : tai->getSuccessorBlocks()) {
-        SILBuilderWithScope(succBB->begin())
-            .createDestroyValue(loc, closureCopy);
-      }
-    } else {
-      llvm_unreachable("Unknown FullApplySite instruction kind");
+  if (SILValue closure = tryRewriteToPartialApplyStack(cvt, singleUser,
+                                                   advanceIfDelete, memoized)) {
+    if (auto *cfi = dyn_cast<ConvertFunctionInst>(closure))
+      closure = cfi->getOperand();
+    if (endAsyncLet && isa<MarkDependenceInst>(closure)) {
+      // Add the top-level mark_dependence (which keeps the closure arguments
+      // alive) as a second operand to the endAsyncLet builtin.
+      // This ensures that the closure arguments are kept alive until the
+      // endAsyncLet builtin.
+      assert(endAsyncLet->getNumOperands() == 1);
+      SILBuilderWithScope builder(endAsyncLet);
+      builder.createBuiltin(endAsyncLet->getLoc(), endAsyncLet->getName(),
+        endAsyncLet->getType(), endAsyncLet->getSubstitutions(),
+        {endAsyncLet->getOperand(0), closure});
+      endAsyncLet->eraseFromParent();
     }
     return true;
   }
-  return false;
+
+  // Insert a copy at the convert_escape_to_noescape [not_guaranteed] and
+  // change the instruction to the guaranteed form.
+  auto escapingClosure = cvt->getOperand();
+  auto *closureCopy =
+      SILBuilderWithScope(cvt).createCopyValue(cvt->getLoc(), escapingClosure);
+  cvt->setLifetimeGuaranteed();
+  cvt->setOperand(closureCopy);
+
+  insertAfterClosureUser(singleUser, [closureCopy](SILBuilder &builder) {
+    auto loc = RegularLocation(builder.getInsertionPointLoc());
+    builder.createDestroyValue(loc, closureCopy);
+  });
+  return true;
 }
 
 /// Ensure the lifetime of the closure across a two step optional conversion

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -100,9 +100,10 @@ OVERRIDE_TASK(task_create_group_future_common, AsyncTaskAndContext, , , ,
               (JobFlags flags, TaskGroup *group,
                const Metadata *futureResultType,
                FutureAsyncSignature::FunctionType *function,
-               HeapObject */* +1 */ closureContext, size_t initialContextSize),
+               void *closureContext, bool owningClosureContext,
+               size_t initialContextSize),
               (flags, group, futureResultType, function, closureContext,
-               initialContextSize))
+               owningClosureContext, initialContextSize))
 
 OVERRIDE_TASK(task_future_wait, void, SWIFT_EXPORT_FROM(swift_Concurrency),
               SWIFT_CC(swiftasync), swift::,
@@ -155,7 +156,7 @@ OVERRIDE_ASYNC_LET(asyncLet_start, void,
                    (AsyncLet *alet,
                     const Metadata *futureResultType,
                     void *closureEntryPoint,
-                    HeapObject */* +1 */ closureContext
+                    void *closureContext
                    ),
                    (alet, futureResultType,
                     closureEntryPoint, closureContext))

stdlib/public/Concurrency/AsyncLet.cpp

@@ -97,15 +97,15 @@ SWIFT_CC(swift)
 static void swift_asyncLet_startImpl(AsyncLet *alet,
                                      const Metadata *futureResultType,
                                      void *closureEntryPoint,
-                                     HeapObject * /* +1 */ closureContext) {
+                                     void *closureContext) {
   AsyncTask *parent = swift_task_getCurrent();
   assert(parent && "async-let cannot be created without parent task");
 
   auto flags = JobFlags(JobKind::Task, parent->getPriority());
   flags.task_setIsFuture(true);
   flags.task_setIsChildTask(true);
 
-  auto childTaskAndContext = swift_task_create_future(
+  auto childTaskAndContext = swift_task_create_future_no_escaping(
       flags,
       futureResultType,
       closureEntryPoint,

stdlib/public/Concurrency/AsyncLet.swift

@@ -20,7 +20,7 @@ import Swift
 @_silgen_name("swift_asyncLet_start")
 public func _asyncLetStart<T>(
   asyncLet: Builtin.RawPointer,
-  operation: __owned @Sendable @escaping () async throws -> T
+  operation: @Sendable () async throws -> T
 )
 
 /// Similar to _taskFutureGet but for AsyncLet