IRGen: Add metadata for async funclets denoting frame entry and frame exists

aschwaighofer · aschwaighofer · commit eaf90dff384f · 2024-09-03T08:44:16.000-07:00
Adds sections `__TEXT,__swift_as_entry`, and `__TEXT,__swift_as_ret` that
contain relative pointers to async functlets modelling async function entries,
and function returns, respectively.

Emission of the sections can be trigger with the frontend option
`-Xfrontend -enable-async-frame-push-pop-metadata`.

This is done by:

* IRGen adding a `async_entry` function attribute to async functions.
* LLVM's coroutine splitting identifying continuation funclets that
  model the return from an async function call by adding the function
  attribute `async_ret`.  (see #llvm-project/pull/9204)
* An LLVM pass that keys off these two function attribute and emits the
  metadata into the above mention sections.

rdar://134460666
diff --git a/include/swift/AST/IRGenOptions.h b/include/swift/AST/IRGenOptions.h
@@ -478,6 +478,8 @@ class IRGenOptions {
   // Whether to run the HotColdSplitting pass when optimizing.
   unsigned EnableHotColdSplit : 1;
 
+  unsigned EmitAsyncFramePushPopMetadata : 1;
+
   /// The number of threads for multi-threaded code generation.
   unsigned NumThreads = 0;
 
@@ -567,7 +569,7 @@ class IRGenOptions {
         DisableReadonlyStaticObjects(false), CollocatedMetadataFunctions(false),
         ColocateTypeDescriptors(true), UseRelativeProtocolWitnessTables(false),
         UseFragileResilientProtocolWitnesses(false),
-        EnableHotColdSplit(false),
+        EnableHotColdSplit(false), EmitAsyncFramePushPopMetadata(false),
         CmdArgs(), SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All),
         PlatformCCallingConvention(llvm::CallingConv::C), UseCASBackend(false),
diff --git a/include/swift/LLVMPasses/Passes.h b/include/swift/LLVMPasses/Passes.h
@@ -149,6 +149,12 @@ namespace swift {
     llvm::PreservedAnalyses run(llvm::Module &M,
                                 llvm::ModuleAnalysisManager &AM);
   };
+
+  struct AsyncEntryReturnMetadataPass
+      : public llvm::PassInfoMixin<AsyncEntryReturnMetadataPass> {
+    llvm::PreservedAnalyses run(llvm::Module &M,
+                                llvm::ModuleAnalysisManager &AM);
+  };
 } // end namespace swift
 
 #endif
diff --git a/include/swift/Option/FrontendOptions.td b/include/swift/Option/FrontendOptions.td
@@ -1310,6 +1310,13 @@ def disable_fragile_resilient_protocol_witnesses :
   Flag<["-"], "disable-fragile-relative-protocol-tables">,
   HelpText<"Disable relative protocol witness tables">;
 
+def enable_async_frame_push_pop_metadata :
+  Flag<["-"], "enable-async-frame-push-pop-metadata">,
+  HelpText<"Enable async frame push pop metadata">;
+def disable_async_frame_push_pop_metadata :
+  Flag<["-"], "disable-async-frame-push-pop-metadata">,
+  HelpText<"Disable async frame push pop metadata">;
+
 def enable_split_cold_code :
   Flag<["-"], "enable-split-cold-code">,
   HelpText<"Enable splitting of cold code when optimizing">;
diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
@@ -3435,6 +3435,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
       Args.hasFlag(OPT_enable_split_cold_code,
                    OPT_disable_split_cold_code,
                    Opts.EnableHotColdSplit);
+  Opts.EmitAsyncFramePushPopMetadata =
+    Args.hasFlag(OPT_enable_async_frame_push_pop_metadata,
+                 OPT_disable_async_frame_push_pop_metadata,
+                 Opts.EmitAsyncFramePushPopMetadata);
   Opts.EnableLargeLoadableTypesReg2Mem =
       Args.hasFlag(OPT_enable_large_loadable_types_reg2mem,
                    OPT_disable_large_loadable_types_reg2mem,
diff --git a/lib/IRGen/GenCall.cpp b/lib/IRGen/GenCall.cpp
@@ -6150,6 +6150,7 @@ void irgen::emitAsyncReturn(
   // Setup the coro.end.async intrinsic call.
   auto &Builder = IGF.Builder;
   auto mustTailCallFn = IGF.createAsyncDispatchFn(fnPtr,Args);
+
   auto handle = IGF.getCoroutineHandle();
   auto rawFnPtr =
       Builder.CreateBitOrPointerCast(fnPtr.getRawPointer(), IGF.IGM.Int8PtrTy);
diff --git a/lib/IRGen/IRGen.cpp b/lib/IRGen/IRGen.cpp
@@ -270,12 +270,14 @@ void swift::performLLVMOptimizations(const IRGenOptions &Opts,
           if (Level != OptimizationLevel::O0)
             FPM.addPass(SwiftARCOptPass());
         });
-    PB.registerOptimizerLastEPCallback([](ModulePassManager &MPM,
+    PB.registerOptimizerLastEPCallback([&](ModulePassManager &MPM,
                                           OptimizationLevel Level) {
       if (Level != OptimizationLevel::O0)
         MPM.addPass(createModuleToFunctionPassAdaptor(SwiftARCContractPass()));
       if (Level == OptimizationLevel::O0)
         MPM.addPass(AlwaysInlinerPass());
+      if (Opts.EmitAsyncFramePushPopMetadata)
+        MPM.addPass(AsyncEntryReturnMetadataPass());
     });
   }
 
diff --git a/lib/IRGen/IRGenSIL.cpp b/lib/IRGen/IRGenSIL.cpp
@@ -2522,6 +2522,10 @@ void IRGenSILFunction::emitSILFunction() {
                              CurFn,
                              LinkEntity::forSILFunction(CurSILFn),
                              getAsyncContextLayout(*this).getSize());
+
+    if (IGM.getOptions().EmitAsyncFramePushPopMetadata) {
+      CurFn->addFnAttr("async_entry");
+    }
   }
   if (isAsyncFn) {
     IGM.noteSwiftAsyncFunctionDef();
diff --git a/lib/LLVMPasses/CMakeLists.txt b/lib/LLVMPasses/CMakeLists.txt
@@ -6,6 +6,7 @@ add_swift_host_library(swiftLLVMPasses STATIC
   LLVMARCContract.cpp
   LLVMInlineTree.cpp
   LLVMMergeFunctions.cpp
+  LLVMEmitAsyncEntryReturnMetadata.cpp
 
   LLVM_LINK_COMPONENTS
   analysis
diff --git a/lib/LLVMPasses/LLVMEmitAsyncEntryReturnMetadata.cpp b/lib/LLVMPasses/LLVMEmitAsyncEntryReturnMetadata.cpp
@@ -0,0 +1,90 @@
+//===--- LLVMEmitAsyncEntryReturnMetadata.cpp - Async function metadata ---===//
+//
+
+#include "swift/LLVMPasses/Passes.h"
+#include "llvm/Pass.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+
+using namespace llvm;
+using namespace swift;
+
+#define DEBUG_TYPE "swift-async-return"
+
+PreservedAnalyses AsyncEntryReturnMetadataPass::run(Module &M,
+                                                    ModuleAnalysisManager &AM) {
+  bool changed = false;
+
+  SmallVector<llvm::Function *, 16> asyncEntries;
+  SmallVector<llvm::Function *, 16> asyncReturns;
+  for (auto &F : M) {
+    if (F.isDeclaration())
+      continue;
+
+    if (F.hasFnAttribute("async_entry"))
+      asyncEntries.push_back(&F);
+    if (F.hasFnAttribute("async_ret"))
+      asyncReturns.push_back(&F);
+  }
+
+  auto &ctxt = M.getContext();
+  auto int32Ty = llvm::Type::getInt32Ty(ctxt);
+  auto sizeTy = M.getDataLayout().getIntPtrType(ctxt, /*addrspace*/ 0);
+
+  auto addSection = [&] (const char * sectionName, const char *globalName,
+                         SmallVectorImpl<llvm::Function *> & entries) {
+    if (entries.empty())
+      return;
+
+    auto intArrayTy = llvm::ArrayType::get(int32Ty, entries.size());
+    auto global =
+      new llvm::GlobalVariable(M, intArrayTy, true,
+                               llvm::GlobalValue::InternalLinkage,
+                               nullptr, /*init*/ globalName,
+                               nullptr, /*insertBefore*/
+                               llvm::GlobalValue::NotThreadLocal,
+                               0/*address space*/);
+    global->setAlignment(Align(4));
+    global->setSection(sectionName);
+    size_t index = 0;
+    SmallVector<llvm::Constant*, 16> offsets;
+    for (auto *fn : entries) {
+      llvm::Constant *indices[] = { llvm::ConstantInt::get(int32Ty, 0),
+        llvm::ConstantInt::get(int32Ty, index)};
+      ++index;
+
+      llvm::Constant *base = llvm::ConstantExpr::getInBoundsGetElementPtr(
+       intArrayTy, global, indices);
+      base = llvm::ConstantExpr::getPtrToInt(base, sizeTy);
+      auto *target = llvm::ConstantExpr::getPtrToInt(fn, sizeTy);
+      llvm::Constant *offset = llvm::ConstantExpr::getSub(target, base);
+
+      if (sizeTy != int32Ty) {
+        offset = llvm::ConstantExpr::getTrunc(offset, int32Ty);
+      }
+      offsets.push_back(offset);
+    }
+    auto constant = llvm::ConstantArray::get(intArrayTy, offsets);
+    global->setInitializer(constant);
+    appendToUsed(M, global);
+
+    llvm::GlobalVariable::SanitizerMetadata Meta;
+    Meta.IsDynInit = false;
+    Meta.NoAddress = true;
+    global->setSanitizerMetadata(Meta);
+
+    changed = true;
+  };
+
+  addSection("__TEXT,__swift_as_entry, coalesced, no_dead_strip",
+             "__swift_async_entry_functlets",
+             asyncEntries);
+  addSection("__TEXT,__swift_as_ret, coalesced, no_dead_strip",
+             "__swift_async_ret_functlets",
+             asyncReturns);
+
+  if (!changed)
+    return PreservedAnalyses::all();
+
+  return PreservedAnalyses::none();
+}
diff --git a/test/IRGen/async_frame_entry_return_metadata.swift b/test/IRGen/async_frame_entry_return_metadata.swift
@@ -0,0 +1,40 @@
+// RUN: %target-swift-frontend -primary-file %s -emit-ir  -module-name async -disable-availability-checking -enable-async-frame-push-pop-metadata | %FileCheck %s --check-prefix=ENABLED
+// RUN: %target-swift-frontend -primary-file %s -emit-ir  -module-name async -disable-availability-checking -O -enable-async-frame-push-pop-metadata | %FileCheck %s --check-prefix=ENABLED
+// RUN: %target-swift-frontend -primary-file %s -emit-ir  -module-name async -disable-availability-checking -disable-async-frame-push-pop-metadata | %FileCheck %s --check-prefix=DISABLED
+// RUN: %target-swift-frontend -primary-file %s -emit-ir  -module-name async -disable-availability-checking | %FileCheck %s --check-prefix=DISABLED
+
+// REQUIRES: OS=macosx || OS=iphoneos
+// REQUIRES: PTRSIZE=64
+
+// ENABLED: @__swift_async_entry_functlets = internal constant [2 x i32] [i32 trunc (i64 sub (i64 ptrtoint (ptr @"$s5async6calleeyyYaF" to i64), i64 ptrtoint (ptr @__swift_async_entry_functlets to i64)) to i32), i32 trunc (i64 sub (i64 ptrtoint (ptr @"$s5async6callerySiSbYaF" to i64), i64 ptrtoint (ptr getelementptr inbounds ([2 x i32], ptr @__swift_async_entry_functlets, i32 0, i32 1) to i64)) to i32)], section "__TEXT,__swift_as_entry, coalesced, no_dead_strip", no_sanitize_address, align 4
+// ENABLED: @__swift_async_ret_functlets = internal constant [1 x i32] [i32 trunc (i64 sub (i64 ptrtoint (ptr @"$s5async6callerySiSbYaFTQ1_" to i64), i64 ptrtoint (ptr @__swift_async_ret_functlets to i64)) to i32)], section "__TEXT,__swift_as_ret, coalesced, no_dead_strip", no_sanitize_address, align 4
+
+// DISABLED-NOT: @__swift_async_entry_functlets
+// DISABLED-NOT: @__swift_async_ret_functlets
+// DISABLED-NOT: s5async6calleeyyYaF.0
+// DISABLED-NOT: s5async6callerySiSbYaF.0
+
+@inline(never)
+public func plusOne() {
+    print("+1")
+}
+
+@inline(never)
+public func minusOne() {
+}
+
+@inline(never)
+public func callee() async {
+  print("callee")
+}
+
+public func caller(_ b: Bool) async -> Int {
+  plusOne()
+
+  if b {
+      await callee()
+  }
+
+  minusOne()
+  return 1
+}
