Merge pull request #85224 from gottesmm/pr-6cb303a9f5f15489fa44c26d00c70155a6d7cc97

[concurrency] Create builtins for invoking specific concurrency runtime functions.

Author: gottesmm

include/swift/AST/ASTSynthesis.h

@@ -52,13 +52,14 @@ enum SingletonTypeSynthesizer {
   _rawUnsafeContinuation,
   _void,
   _word,
-  _swiftInt,       // Swift.Int
-  _serialExecutor, // the '_Concurrency.SerialExecutor' protocol
-  _taskExecutor,   // the '_Concurrency.TaskExecutor' protocol
-  _actor,          // the '_Concurrency.Actor' protocol
-  _distributedActor,  // the 'Distributed.DistributedActor' protocol
+  _swiftInt,               // Swift.Int
+  _serialExecutor,         // the '_Concurrency.SerialExecutor' protocol
+  _taskExecutor,           // the '_Concurrency.TaskExecutor' protocol
+  _actor,                  // the '_Concurrency.Actor' protocol
+  _distributedActor,       // the 'Distributed.DistributedActor' protocol
   _unsafeRawBufferPointer, // UnsafeRawBufferPointer
-  _unconstrainedAny, // any ~Copyable & ~Escapable
+  _unconstrainedAny,       // any ~Copyable & ~Escapable
+  _unsafeRawPointer,       // UnsafeRawPointer
 };
 inline Type synthesizeType(SynthesisContext &SC,
                            SingletonTypeSynthesizer kind) {
@@ -93,6 +94,8 @@ inline Type synthesizeType(SynthesisContext &SC,
       ->getDeclaredInterfaceType();
   case _unsafeRawBufferPointer:
     return SC.Context.getUnsafeRawBufferPointerType();
+  case _unsafeRawPointer:
+    return SC.Context.getUnsafeRawPointerType();
   case _copyable:
     return SC.Context.getProtocol(KnownProtocolKind::Copyable)
         ->getDeclaredInterfaceType();

include/swift/AST/Builtins.def

@@ -1174,6 +1174,42 @@ BUILTIN_TYPE_CHECKER_OPERATION(TriggerFallbackDiagnostic, trigger_fallback_diagn
 /// Maybe.
 BUILTIN_TYPE_TRAIT_OPERATION(CanBeObjCClass, canBeClass)
 
+/// Equivalent to calling swift_task_addCancellationHandler.
+///
+/// Signature: (handler: () -> Void) -> UnsafeRawPointer
+BUILTIN_MISC_OPERATION_WITH_SILGEN(TaskAddCancellationHandler,
+                                  "taskAddCancellationHandler", "",
+                                   Special)
+
+/// Equivalent to calling swift_task_removeCancellationHandler.
+///
+/// Signature: (record: UnsafeRawPointer) -> ()
+BUILTIN_MISC_OPERATION(TaskRemoveCancellationHandler,
+                       "taskRemoveCancellationHandler", "", Special)
+
+/// Equivalent to calling swift_task_addPriorityEscalationHandler.
+///
+/// Signature: (handler: (UInt8, UInt8) -> Void) -> UnsafeRawPointer
+BUILTIN_MISC_OPERATION_WITH_SILGEN(TaskAddPriorityEscalationHandler,
+                                   "taskAddPriorityEscalationHandler", "",
+                                   Special)
+
+/// Equivalent to calling swift_task_removePriorityEscalationHandler.
+///
+/// Signature: (record: UnsafeRawPointer) -> ()
+BUILTIN_MISC_OPERATION(TaskRemovePriorityEscalationHandler,
+                       "taskRemovePriorityEscalationHandler", "", Special)
+
+/// Equivalent to calling swift_task_localValuePush.
+///
+/// Signature: <Value> (key: Builtin.RawPointer, value: __owned Value) -> ()
+BUILTIN_MISC_OPERATION(TaskLocalValuePush, "taskLocalValuePush", "", Special)
+
+/// Equivalent to calling swift_task_localValuePop.
+///
+/// Signature: () -> ()
+BUILTIN_MISC_OPERATION(TaskLocalValuePop, "taskLocalValuePop", "", Special)
+
 #undef BUILTIN_TYPE_TRAIT_OPERATION
 #undef BUILTIN_UNARY_OPERATION
 #undef BUILTIN_BINARY_PREDICATE

include/swift/AST/Builtins.h

@@ -107,10 +107,21 @@ getLLVMIntrinsicIDForBuiltinWithOverflow(BuiltinValueKind ID);
 ///
 /// Returns null if the name does not identifier a known builtin value.
 ValueDecl *getBuiltinValueDecl(ASTContext &Context, Identifier Name);
-  
+
+/// Overload that takes a StringRef instead of an Identifier. We convert it to
+/// an identifier internally.
+ValueDecl *getBuiltinValueDecl(ASTContext &Context, StringRef Name);
+
 /// Returns the name of a builtin declaration given a builtin ID.
 StringRef getBuiltinName(BuiltinValueKind ID);
-  
+
+/// Namespace containing constexpr StringLiterals of BuiltinNames so that
+/// builtin names can be referred to without using raw string literals.
+namespace BuiltinNames {
+#define BUILTIN(Id, Name, Attrs) constexpr StringLiteral Id = Name;
+#include "swift/AST/Builtins.def"
+} // namespace BuiltinNames
+
 /// The information identifying the builtin - its kind and types.
 class BuiltinInfo {
 public:

include/swift/Runtime/RuntimeFunctions.def

@@ -2660,6 +2660,72 @@ FUNCTION(TaskGroupDestroy,
          EFFECT(RuntimeEffect::Concurrency),
          UNKNOWN_MEMEFFECTS)
 
+// CancellationNotificationStatusRecord*
+// swift_task_addCancellationHandler(
+//     CancellationNotificationStatusRecord::FunctionType handler,
+//     void *context);
+FUNCTION(TaskAddCancellationHandler,
+        _Concurrency, swift_task_addCancellationHandler, SwiftCC,
+        ConcurrencyAvailability,
+        RETURNS(Int8PtrTy),
+        ARGS(Int8PtrTy, Int8PtrTy),
+        ATTRS(NoUnwind),
+        EFFECT(RuntimeEffect::Concurrency),
+        UNKNOWN_MEMEFFECTS)
+
+// void swift_task_removeCancellationHandler(
+//   CancellationNotificationStatusRecord *record);
+FUNCTION(TaskRemoveCancellationHandler,
+         _Concurrency, swift_task_removeCancellationHandler, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(VoidTy),
+         ARGS(Int8PtrTy),
+         ATTRS(NoUnwind),
+         EFFECT(RuntimeEffect::Concurrency),
+         UNKNOWN_MEMEFFECTS)
+
+// EscalationNotificationStatusRecord*
+// swift_task_addPriorityEscalationHandler(
+//     EscalationNotificationStatusRecord::FunctionType handler,
+//     void *context);
+FUNCTION(TaskAddPriorityEscalationHandler,
+        _Concurrency, swift_task_addPriorityEscalationHandler, SwiftCC,
+        ConcurrencyAvailability,
+        RETURNS(Int8PtrTy),
+        ARGS(Int8PtrTy, Int8PtrTy),
+        ATTRS(NoUnwind),
+        EFFECT(RuntimeEffect::Concurrency),
+        UNKNOWN_MEMEFFECTS)
+
+// void swift_task_removePriorityEscalationHandler(
+//   EscalationNotificationStatusRecord *record);
+FUNCTION(TaskRemovePriorityEscalationHandler,
+         _Concurrency, swift_task_removePriorityEscalationHandler, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(VoidTy),
+         ARGS(Int8PtrTy),
+         ATTRS(NoUnwind),
+         EFFECT(RuntimeEffect::Concurrency),
+         UNKNOWN_MEMEFFECTS)
+
+// void swift_task_localValuePush(const HeapObject *key,
+//                                /* +1 */ OpaqueValue *value,
+//                                const Metadata *valueType);
+FUNCTION(TaskLocalValuePush, _Concurrency, swift_task_localValuePush, SwiftCC,
+         ConcurrencyAvailability, RETURNS(),
+         ARGS(RefCountedPtrTy, OpaquePtrTy, TypeMetadataPtrTy), ATTRS(NoUnwind),
+         EFFECT(RuntimeEffect::Concurrency), UNKNOWN_MEMEFFECTS)
+
+// void swift_task_localValuePop();
+FUNCTION(TaskLocalValuePop,
+         _Concurrency, swift_task_localValuePop, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(),
+         ARGS(),
+         ATTRS(NoUnwind),
+         EFFECT(RuntimeEffect::Concurrency),
+         UNKNOWN_MEMEFFECTS)
+
 // AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContextWithType(const Metadata *);
 FUNCTION(AutoDiffCreateLinearMapContextWithType,
          Swift, swift_autoDiffCreateLinearMapContextWithType, SwiftCC,

include/swift/SIL/AddressWalker.h

@@ -295,6 +295,7 @@ TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) {
         case BuiltinValueKind::AddressOfRawLayout:
         case BuiltinValueKind::FlowSensitiveSelfIsolation:
         case BuiltinValueKind::FlowSensitiveDistributedSelfIsolation:
+        case BuiltinValueKind::TaskLocalValuePush:
           callVisitUse(op);
           continue;
         default:

include/swift/SIL/SILBuilder.h

@@ -611,6 +611,12 @@ class SILBuilder {
                                       Subs, Args, getInstructionContext()));
   }
 
+  BuiltinInst *createBuiltin(SILLocation Loc, StringRef Name, SILType ResultTy,
+                             SubstitutionMap Subs, ArrayRef<SILValue> Args) {
+    return createBuiltin(Loc, getASTContext().getIdentifier(Name), ResultTy,
+                         Subs, Args);
+  }
+
   /// Create a binary function with the signature: OpdTy, OpdTy -> ResultTy.
   BuiltinInst *createBuiltinBinaryFunction(SILLocation Loc, StringRef Name,
                                            SILType OpdTy, SILType ResultTy,

include/swift/SIL/SILType.h

@@ -1067,6 +1067,9 @@ class SILType {
   /// Return Builtin.ImplicitActor.
   static SILType getBuiltinImplicitActorType(const ASTContext &ctx);
 
+  /// Return UnsafeRawPointer.
+  static SILType getUnsafeRawPointer(const ASTContext &ctx);
+
   //
   // Utilities for treating SILType as a pointer-like type.
   //

lib/AST/Builtins.cpp

@@ -2364,6 +2364,52 @@ static ValueDecl *getEmplace(ASTContext &ctx, Identifier id) {
   return builder.build(id);
 }
 
+static ValueDecl *getTaskAddCancellationHandler(ASTContext &ctx,
+                                                Identifier id) {
+  auto extInfo = ASTExtInfoBuilder().withNoEscape().build();
+  auto fnType = FunctionType::get({}, ctx.TheEmptyTupleType, extInfo);
+  return getBuiltinFunction(ctx, id, _thin,
+                            _parameters(_label("handler", fnType)),
+                            _unsafeRawPointer);
+}
+
+static ValueDecl *getTaskRemoveCancellationHandler(ASTContext &ctx,
+                                                   Identifier id) {
+  return getBuiltinFunction(
+      ctx, id, _thin, _parameters(_label("record", _unsafeRawPointer)), _void);
+}
+
+static ValueDecl *getTaskAddPriorityEscalationHandler(ASTContext &ctx,
+                                                      Identifier id) {
+  std::array<AnyFunctionType::Param, 2> params = {
+      AnyFunctionType::Param(ctx.getUInt8Type()),
+      AnyFunctionType::Param(ctx.getUInt8Type()),
+  };
+  // (UInt8, UInt8) -> ()
+  auto extInfo = ASTExtInfoBuilder().withNoEscape().build();
+  auto *functionType =
+      FunctionType::get(params, ctx.TheEmptyTupleType, extInfo);
+  return getBuiltinFunction(ctx, id, _thin,
+                            _parameters(_label("handler", functionType)),
+                            _unsafeRawPointer);
+}
+
+static ValueDecl *getTaskRemovePriorityEscalationHandler(ASTContext &ctx,
+                                                         Identifier id) {
+  return getBuiltinFunction(
+      ctx, id, _thin, _parameters(_label("record", _unsafeRawPointer)), _void);
+}
+
+static ValueDecl *getTaskLocalValuePush(ASTContext &ctx, Identifier id) {
+  return getBuiltinFunction(ctx, id, _thin, _generics(_unrestricted),
+                            _parameters(_rawPointer, _consuming(_typeparam(0))),
+                            _void);
+}
+
+static ValueDecl *getTaskLocalValuePop(ASTContext &ctx, Identifier id) {
+  return getBuiltinFunction(ctx, id, _thin, _parameters(), _void);
+}
+
 /// An array of the overloaded builtin kinds.
 static const OverloadedBuiltinKind OverloadedBuiltinKinds[] = {
   OverloadedBuiltinKind::None,
@@ -3452,6 +3498,24 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
 
   case BuiltinValueKind::Emplace:
     return getEmplace(Context, Id);
+
+  case BuiltinValueKind::TaskAddCancellationHandler:
+    return getTaskAddCancellationHandler(Context, Id);
+
+  case BuiltinValueKind::TaskRemoveCancellationHandler:
+    return getTaskRemoveCancellationHandler(Context, Id);
+
+  case BuiltinValueKind::TaskAddPriorityEscalationHandler:
+    return getTaskAddPriorityEscalationHandler(Context, Id);
+
+  case BuiltinValueKind::TaskRemovePriorityEscalationHandler:
+    return getTaskRemovePriorityEscalationHandler(Context, Id);
+
+  case BuiltinValueKind::TaskLocalValuePush:
+    return getTaskLocalValuePush(Context, Id);
+
+  case BuiltinValueKind::TaskLocalValuePop:
+    return getTaskLocalValuePop(Context, Id);
   }
 
   llvm_unreachable("bad builtin value!");
@@ -3815,3 +3879,7 @@ BuiltinFixedArrayType::isFixedNegativeSize() const {
   }
   return false;
 }
+
+ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, StringRef Name) {
+  return getBuiltinValueDecl(Context, Context.getIdentifier(Name));
+}

lib/IRGen/GenBuiltin.cpp

@@ -1565,6 +1565,29 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
     return;
   }
 
+  case BuiltinValueKind::TaskRemovePriorityEscalationHandler:
+  case BuiltinValueKind::TaskRemoveCancellationHandler: {
+    auto rawPointer = args.claimNext();
+    emitBuiltinTaskRemoveHandler(IGF, Builtin.ID, rawPointer);
+    return;
+  }
+  case BuiltinValueKind::TaskAddCancellationHandler:
+  case BuiltinValueKind::TaskAddPriorityEscalationHandler: {
+    auto func = args.claimNext();
+    auto context = args.claimNext();
+    out.add(emitBuiltinTaskAddHandler(IGF, Builtin.ID, func, context));
+    return;
+  }
+  case BuiltinValueKind::TaskLocalValuePop:
+    return emitBuiltinTaskLocalValuePop(IGF);
+  case BuiltinValueKind::TaskLocalValuePush: {
+    auto *key = args.claimNext();
+    auto *value = args.claimNext();
+    // Grab T from the builtin.
+    auto *valueMetatype = IGF.emitTypeMetadataRef(argTypes[1].getASTType());
+    return emitBuiltinTaskLocalValuePush(IGF, key, value, valueMetatype);
+  }
+
   // Builtins without IRGen implementations.
   case BuiltinValueKind::None:
   case BuiltinValueKind::CondFailMessage:

lib/IRGen/GenConcurrency.cpp

@@ -313,6 +313,61 @@ llvm::Value *irgen::emitBuiltinStartAsyncLet(IRGenFunction &IGF,
   return alet;
 }
 
+llvm::Value *irgen::emitBuiltinTaskAddHandler(IRGenFunction &IGF,
+                                              BuiltinValueKind kind,
+                                              llvm::Value *func,
+                                              llvm::Value *context) {
+  auto callee = [&]() -> FunctionPointer {
+    if (kind == BuiltinValueKind::TaskAddCancellationHandler) {
+      return IGF.IGM.getTaskAddCancellationHandlerFunctionPointer();
+    }
+    if (kind == BuiltinValueKind::TaskAddPriorityEscalationHandler) {
+      return IGF.IGM.getTaskAddPriorityEscalationHandlerFunctionPointer();
+    }
+    llvm::report_fatal_error("Unhandled builtin");
+  }();
+  auto *call = IGF.Builder.CreateCall(callee, {func, context});
+  call->setDoesNotThrow();
+  call->setCallingConv(IGF.IGM.SwiftCC);
+  return call;
+}
+
+void irgen::emitBuiltinTaskRemoveHandler(IRGenFunction &IGF,
+                                         BuiltinValueKind kind,
+                                         llvm::Value *record) {
+  auto callee = [&]() -> FunctionPointer {
+    if (kind == BuiltinValueKind::TaskRemoveCancellationHandler) {
+      return IGF.IGM.getTaskRemoveCancellationHandlerFunctionPointer();
+    }
+    if (kind == BuiltinValueKind::TaskRemovePriorityEscalationHandler) {
+      return IGF.IGM.getTaskRemovePriorityEscalationHandlerFunctionPointer();
+    }
+    llvm::report_fatal_error("Unhandled builtin");
+  }();
+  auto *call = IGF.Builder.CreateCall(callee, {record});
+  call->setDoesNotThrow();
+  call->setCallingConv(IGF.IGM.SwiftCC);
+}
+
+void irgen::emitBuiltinTaskLocalValuePush(IRGenFunction &IGF, llvm::Value *key,
+                                          llvm::Value *value,
+                                          llvm::Value *valueMetatype) {
+  auto callee = IGF.IGM.getTaskLocalValuePushFunctionPointer();
+
+  // We pass in Value at +1, but we are luckily given the value already at +1,
+  // so the end lifetime is performed for us.
+  auto *call = IGF.Builder.CreateCall(callee, {key, value, valueMetatype});
+  call->setDoesNotThrow();
+  call->setCallingConv(IGF.IGM.SwiftCC);
+}
+
+void irgen::emitBuiltinTaskLocalValuePop(IRGenFunction &IGF) {
+  auto *call =
+      IGF.Builder.CreateCall(IGF.IGM.getTaskLocalValuePopFunctionPointer(), {});
+  call->setDoesNotThrow();
+  call->setCallingConv(IGF.IGM.SwiftCC);
+}
+
 void irgen::emitFinishAsyncLet(IRGenFunction &IGF,
                                llvm::Value *asyncLet,
                                llvm::Value *resultBuffer) {