Merge pull request swiftlang#35303 from aschwaighofer/introduce_swiftasync_lowering

Add llvm::Attribute::SwiftAsync to the context parameter

Author: aschwaighofer

include/swift/ABI/Executor.h

@@ -118,7 +118,7 @@ using JobInvokeFunction =
 
 using TaskContinuationFunction =
   SWIFT_CC(swiftasync)
-  void (AsyncTask *, ExecutorRef, AsyncContext *);
+  void (AsyncTask *, ExecutorRef, SWIFT_ASYNC_CONTEXT AsyncContext *);
 
 template <class AsyncSignature>
 class AsyncFunctionPointer;

include/swift/ABI/MetadataValues.h

@@ -1172,6 +1172,9 @@ namespace SpecialPointerAuthDiscriminators {
   const uint16_t AsyncContextYield = 0xe207; // = 57863
   const uint16_t CancellationNotificationFunction = 0x1933; // = 6451
   const uint16_t EscalationNotificationFunction = 0x5be4; // = 23524
+
+  /// Swift async context parameter stored in the extended frame info.
+  const uint16_t SwiftAsyncContextExtendedFrameEntry = 0xc31a;
 }
 
 /// The number of arguments that will be passed directly to a generic

include/swift/AST/IRGenOptions.h

@@ -146,6 +146,9 @@ struct PointerAuthOptions : clang::PointerAuthOptions {
 
   /// The resume function stored in AsyncTask.
   PointerAuthSchema TaskResumeFunction;
+
+  /// The swift async context entry in the extended frame info.
+  PointerAuthSchema AsyncContextExtendedFrameEntry;
 };
 
 enum class JITDebugArtifact : unsigned {

include/swift/Runtime/Config.h

@@ -176,6 +176,12 @@ extern uintptr_t __COMPATIBILITY_LIBRARIES_CANNOT_CHECK_THE_IS_SWIFT_BIT_DIRECTL
 #define SWIFT_INDIRECT_RESULT
 #endif
 
+#if __has_attribute(swift_async_context)
+#define SWIFT_ASYNC_CONTEXT __attribute__((swift_async_context))
+#else
+#define SWIFT_ASYNC_CONTEXT
+#endif
+
 // SWIFT_CC(swiftasync) is the Swift async calling convention.
 // We assume that it supports mandatory tail call elimination.
 #if __has_attribute(swiftasynccall)

lib/IRGen/GenCall.cpp

@@ -817,6 +817,8 @@ void SignatureExpansion::addAsyncParameters() {
   //   void (AsyncTask *, ExecutorRef, AsyncContext *);
   ParamIRTypes.push_back(IGM.SwiftTaskPtrTy);
   ParamIRTypes.push_back(IGM.SwiftExecutorPtrTy);
+  Attrs = Attrs.addParamAttribute(IGM.getLLVMContext(), getCurParamIndex(),
+                                  llvm::Attribute::SwiftAsync);
   ParamIRTypes.push_back(IGM.SwiftContextPtrTy);
 }
 

lib/IRGen/GenFunc.cpp

@@ -2385,7 +2385,7 @@ void irgen::emitBlockHeader(IRGenFunction &IGF,
 
 llvm::Function *IRGenFunction::getOrCreateResumePrjFn() {
   auto name = "__swift_async_resume_project_context";
-  return cast<llvm::Function>(IGM.getOrCreateHelperFunction(
+  auto Fn = cast<llvm::Function>(IGM.getOrCreateHelperFunction(
       name, IGM.Int8PtrTy, {IGM.Int8PtrTy},
       [&](IRGenFunction &IGF) {
         auto it = IGF.CurFn->arg_begin();
@@ -2398,9 +2398,29 @@ llvm::Function *IRGenFunction::getOrCreateResumePrjFn() {
               PointerAuthInfo::emit(IGF, schema, addr, PointerAuthEntity());
           callerContext = emitPointerAuthAuth(IGF, callerContext, authInfo);
         }
+        // TODO: remove this once all platforms support lowering the intrinsic.
+        // At the time of this writing only arm64 supports it.
+        if (IGM.TargetInfo.canUseSwiftAsyncContextAddrIntrinsic()) {
+          auto contextLocationInExtendedFrame =
+              Address(Builder.CreateIntrinsicCall(
+                          llvm::Intrinsic::swift_async_context_addr, {}),
+                      IGM.getPointerAlignment());
+          // On arm64e we need to sign this pointer address discriminated
+          // with 0xc31a and process dependent key.
+          if (auto schema = IGF.IGM.getOptions()
+                                .PointerAuth.AsyncContextExtendedFrameEntry) {
+            auto authInfo = PointerAuthInfo::emit(
+                IGF, schema, contextLocationInExtendedFrame.getAddress(),
+                PointerAuthEntity());
+            callerContext = emitPointerAuthSign(IGF, callerContext, authInfo);
+          }
+          Builder.CreateStore(callerContext, contextLocationInExtendedFrame);
+        }
         Builder.CreateRet(callerContext);
       },
       false /*isNoInline*/));
+  Fn->addFnAttr(llvm::Attribute::AlwaysInline);
+  return Fn;
 }
 llvm::Function *
 IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,

lib/IRGen/IRGen.cpp

@@ -702,6 +702,10 @@ static void setPointerAuthOptions(PointerAuthOptions &opts,
   opts.TaskResumeFunction =
       PointerAuthSchema(codeKey, /*address*/ true, Discrimination::Constant,
                         SpecialPointerAuthDiscriminators::TaskResumeFunction);
+
+  opts.AsyncContextExtendedFrameEntry = PointerAuthSchema(
+      dataKey, /*address*/ true, Discrimination::Constant,
+      SpecialPointerAuthDiscriminators::SwiftAsyncContextExtendedFrameEntry);
 }
 
 std::unique_ptr<llvm::TargetMachine>

lib/IRGen/SwiftTargetInfo.cpp

@@ -60,6 +60,8 @@ static void configureARM64(IRGenModule &IGM, const llvm::Triple &triple,
   // arm64 tops out at 56 effective bits of address space and reserves the high
   // half for the kernel.
   target.SwiftRetainIgnoresNegativeValues = true;
+
+  target.UsableSwiftAsyncContextAddrIntrinsic = true;
 }
 
 /// Configures target-specific information for x86-64 platforms.

lib/IRGen/SwiftTargetInfo.h

@@ -51,6 +51,10 @@ class SwiftTargetInfo {
     return ObjCHasOpaqueISAs;
   }
 
+  bool canUseSwiftAsyncContextAddrIntrinsic() const {
+    return UsableSwiftAsyncContextAddrIntrinsic;
+  }
+
   /// The target's object format type.
   llvm::Triple::ObjectFormatType OutputObjectFormat;
 
@@ -104,6 +108,8 @@ class SwiftTargetInfo {
   /// True if `swift_retain` and `swift_release` are no-ops when passed
   /// "negative" pointer values.
   bool SwiftRetainIgnoresNegativeValues = false;
+
+  bool UsableSwiftAsyncContextAddrIntrinsic = false;
 };
 
 }

stdlib/public/Concurrency/Task.cpp

@@ -172,7 +172,7 @@ static FullMetadata<HeapMetadata> taskHeapMetadata = {
 /// to handle the final return.
 SWIFT_CC(swift)
 static void completeTask(AsyncTask *task, ExecutorRef executor,
-                         AsyncContext *context) {
+                         SWIFT_ASYNC_CONTEXT AsyncContext *context) {
   // Tear down the task-local allocator immediately;
   // there's no need to wait for the object to be destroyed.
   _swift_task_alloc_destroy(task);
@@ -305,7 +305,7 @@ AsyncTaskAndContext swift::swift_task_create_future_f(
 
 void swift::swift_task_future_wait(
     AsyncTask *waitingTask, ExecutorRef executor,
-    AsyncContext *rawContext) {
+    SWIFT_ASYNC_CONTEXT AsyncContext *rawContext) {
   // Suspend the waiting task.
   waitingTask->ResumeTask = rawContext->ResumeParent;
   waitingTask->ResumeContext = rawContext;
@@ -402,7 +402,7 @@ using RunAndBlockCalleeContext =
 /// Second half of the runAndBlock async function.
 SWIFT_CC(swiftasync)
 static void runAndBlock_finish(AsyncTask *task, ExecutorRef executor,
-                               AsyncContext *_context) {
+                               SWIFT_ASYNC_CONTEXT AsyncContext *_context) {
   auto calleeContext = static_cast<RunAndBlockCalleeContext*>(_context);
   auto context = popAsyncContext(task, calleeContext);
 
@@ -414,7 +414,7 @@ static void runAndBlock_finish(AsyncTask *task, ExecutorRef executor,
 /// First half of the runAndBlock async function.
 SWIFT_CC(swiftasync)
 static void runAndBlock_start(AsyncTask *task, ExecutorRef executor,
-                              AsyncContext *_context) {
+                              SWIFT_ASYNC_CONTEXT AsyncContext *_context) {
   auto callerContext = static_cast<RunAndBlockContext*>(_context);
 
   size_t calleeContextSize;