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

Specifically:

1. swift_task_addCancellationHandler
2. swift_task_removeCancellationHandler
3. swift_task_addPriorityEscalationHandler
4. swift_task_removePriorityEscalationHandler
5. swift_task_localValuePush
6. swift_task_localValuePop

rdar://109850951

Author: gottesmm

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/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:

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,
@@ -3450,6 +3496,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!");

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) {

lib/IRGen/GenConcurrency.h

@@ -115,6 +115,29 @@ emitTaskCreate(IRGenFunction &IGF, llvm::Value *flags,
 llvm::Value *clearImplicitIsolatedActorBits(IRGenFunction &IGF,
                                             llvm::Value *value);
 
+/// Emit IR for a builtin that adds a handler to the Task's task record.
+///
+/// \returns the record that can be used to refer to and cancel the handler.
+///
+/// Currently supports TaskAddCancellationHandler and
+/// TaskAddPriorityEscalationHandler.
+llvm::Value *emitBuiltinTaskAddHandler(IRGenFunction &IGF,
+                                       BuiltinValueKind kind, llvm::Value *func,
+                                       llvm::Value *context);
+
+/// Emit IR for a builtin that cancels some sort of handler by calling an ABI
+/// entry point. Record is a value that was returned by the handler creator.
+///
+/// E.x.: TaskRemoveCancellationHandler, TaskRemovePriorityEscalationHandler.
+void emitBuiltinTaskRemoveHandler(IRGenFunction &IGF, BuiltinValueKind kind,
+                                  llvm::Value *record);
+
+void emitBuiltinTaskLocalValuePush(IRGenFunction &IGF, llvm::Value *key,
+                                   llvm::Value *value,
+                                   llvm::Value *valueMetatype);
+
+void emitBuiltinTaskLocalValuePop(IRGenFunction &IGF);
+
 } // end namespace irgen
 } // end namespace swift
 

lib/SIL/IR/OperandOwnership.cpp

@@ -1052,6 +1052,19 @@ BUILTIN_OPERAND_OWNERSHIP(BitwiseEscape, BuildDefaultActorExecutorRef)
 BUILTIN_OPERAND_OWNERSHIP(BitwiseEscape, BuildMainActorExecutorRef)
 
 BUILTIN_OPERAND_OWNERSHIP(TrivialUse, AutoDiffCreateLinearMapContextWithType)
+
+// InstantaneousUse since we take in a closure at +0.
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, TaskAddCancellationHandler)
+// Trivial use since our operand is just an UnsafeRawPointer.
+BUILTIN_OPERAND_OWNERSHIP(TrivialUse, TaskRemoveCancellationHandler)
+// InstantaneousUse since we take in a closure at +0.
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, TaskAddPriorityEscalationHandler)
+// Trivial use since our operand is just an UnsafeRawPointer.
+BUILTIN_OPERAND_OWNERSHIP(TrivialUse, TaskRemovePriorityEscalationHandler)
+// This is a trivial use since our first operand is a Builtin.RawPointer and our
+// second is an address to our generic Value.
+BUILTIN_OPERAND_OWNERSHIP(TrivialUse, TaskLocalValuePush)
+
 #undef BUILTIN_OPERAND_OWNERSHIP
 
 #define SHOULD_NEVER_VISIT_BUILTIN(ID)                                         \
@@ -1062,6 +1075,7 @@ BUILTIN_OPERAND_OWNERSHIP(TrivialUse, AutoDiffCreateLinearMapContextWithType)
   }
 SHOULD_NEVER_VISIT_BUILTIN(GetCurrentAsyncTask)
 SHOULD_NEVER_VISIT_BUILTIN(GetCurrentExecutor)
+SHOULD_NEVER_VISIT_BUILTIN(TaskLocalValuePop)
 #undef SHOULD_NEVER_VISIT_BUILTIN
 
 // Builtins that should be lowered to SIL instructions so we should never see

lib/SIL/IR/ValueOwnership.cpp

@@ -664,6 +664,13 @@ CONSTANT_OWNERSHIP_BUILTIN(Owned, DistributedActorAsAnyActor)
 CONSTANT_OWNERSHIP_BUILTIN(Guaranteed, ExtractFunctionIsolation) // unreachable
 CONSTANT_OWNERSHIP_BUILTIN(None, AddressOfRawLayout)
 
+CONSTANT_OWNERSHIP_BUILTIN(None, TaskAddCancellationHandler)
+CONSTANT_OWNERSHIP_BUILTIN(None, TaskRemoveCancellationHandler)
+CONSTANT_OWNERSHIP_BUILTIN(None, TaskAddPriorityEscalationHandler)
+CONSTANT_OWNERSHIP_BUILTIN(None, TaskRemovePriorityEscalationHandler)
+CONSTANT_OWNERSHIP_BUILTIN(None, TaskLocalValuePush)
+CONSTANT_OWNERSHIP_BUILTIN(None, TaskLocalValuePop)
+
 #undef CONSTANT_OWNERSHIP_BUILTIN
 
 // Check all of these...

lib/SILGen/SILGenBuiltin.cpp

@@ -2182,6 +2182,26 @@ static ManagedValue emitBuiltinEmplace(SILGenFunction &SGF,
   return SGF.B.createLoadTake(loc, result);
 }
 
+static ManagedValue emitBuiltinTaskAddCancellationHandler(
+    SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
+    ArrayRef<ManagedValue> args, SGFContext C) {
+  auto *b =
+      SGF.B.createBuiltin(loc, BuiltinNames::TaskAddCancellationHandler,
+                          SILType::getUnsafeRawPointer(SGF.getASTContext()),
+                          subs, {args[0].getValue()});
+  return ManagedValue::forRValueWithoutOwnership(b);
+}
+
+static ManagedValue emitBuiltinTaskAddPriorityEscalationHandler(
+    SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
+    ArrayRef<ManagedValue> args, SGFContext C) {
+  auto *b =
+      SGF.B.createBuiltin(loc, BuiltinNames::TaskAddPriorityEscalationHandler,
+                          SILType::getUnsafeRawPointer(SGF.getASTContext()),
+                          subs, {args[0].getValue()});
+  return ManagedValue::forRValueWithoutOwnership(b);
+}
+
 std::optional<SpecializedEmitter>
 SpecializedEmitter::forDecl(SILGenModule &SGM, SILDeclRef function) {
   // Only consider standalone declarations in the Builtin module.