Merge pull request #34600 from DougGregor/async-task-builtins

[Concurrency] Add cancelAsyncTask() builtin.

Author: DougGregor

include/swift/ABI/Task.h

@@ -182,7 +182,7 @@ class AsyncTask : public HeapObject, public Job {
   /// prevent it from being corrupted in flight.
   AsyncContext * __ptrauth_swift_task_resume_context ResumeContext;
 
-  /// The currntly-active information about cancellation.
+  /// The currently-active information about cancellation.
   std::atomic<ActiveTaskStatus> Status;
 
   /// Reserved for the use of the task-local stack allocator.

include/swift/AST/Builtins.def

@@ -723,6 +723,12 @@ BUILTIN_MISC_OPERATION(IntInstrprofIncrement, "int_instrprof_increment", "", Spe
 // function is executing.
 BUILTIN_MISC_OPERATION_WITH_SILGEN(GetCurrentAsyncTask, "getCurrentAsyncTask", "n", Special)
 
+/// cancelAsyncTask(): (Builtin.NativeObject) -> Void
+///
+/// Cancel the given asynchronous task.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(CancelAsyncTask, "cancelAsyncTask", "", Special)
+
+
 /// globalStringTablePointer has type String -> Builtin.RawPointer.
 /// It returns an immortal, global string table pointer for strings constructed
 /// from string literals. We consider it effects as readnone meaning that it

include/swift/Runtime/RuntimeFunctions.def

@@ -1489,6 +1489,14 @@ FUNCTION(TaskDealloc,
          ARGS(SwiftTaskPtrTy, Int8PtrTy),
          ATTRS(NoUnwind, ArgMemOnly))
 
+// void swift_task_dealloc(AsyncTask *task, void *ptr);
+FUNCTION(TaskCancel,
+         swift_task_cancel, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(VoidTy),
+         ARGS(SwiftTaskPtrTy),
+         ATTRS(NoUnwind, ArgMemOnly))
+
 #undef RETURNS
 #undef ARGS
 #undef ATTRS

lib/AST/Builtins.cpp

@@ -1345,6 +1345,11 @@ static ValueDecl *getGetCurrentAsyncTask(ASTContext &ctx, Identifier id) {
   return getBuiltinFunction(id, { }, ctx.TheNativeObjectType);
 }
 
+static ValueDecl *getCancelAsyncTask(ASTContext &ctx, Identifier id) {
+  return getBuiltinFunction(
+      id, { ctx.TheNativeObjectType }, ctx.TheEmptyTupleType);
+}
+
 static ValueDecl *getPoundAssert(ASTContext &Context, Identifier Id) {
   auto int1Type = BuiltinIntegerType::get(1, Context);
   auto optionalRawPointerType = BoundGenericEnumType::get(
@@ -2474,6 +2479,9 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
   case BuiltinValueKind::GetCurrentAsyncTask:
     return getGetCurrentAsyncTask(Context, Id);
 
+  case BuiltinValueKind::CancelAsyncTask:
+    return getCancelAsyncTask(Context, Id);
+
   case BuiltinValueKind::PoundAssert:
     return getPoundAssert(Context, Id);
 

lib/IRGen/GenBuiltin.cpp

@@ -220,6 +220,11 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
 
   // Everything else cares about the (rvalue) argument.
 
+  if (Builtin.ID == BuiltinValueKind::CancelAsyncTask) {
+    emitTaskCancel(IGF, args.claimNext());
+    return;
+  }
+
   // If this is an LLVM IR intrinsic, lower it to an intrinsic call.
   const IntrinsicInfo &IInfo = IGF.getSILModule().getIntrinsicInfo(FnId);
   llvm::Intrinsic::ID IID = IInfo.ID;

lib/IRGen/GenCall.cpp

@@ -3458,6 +3458,18 @@ void irgen::emitTaskDealloc(IRGenFunction &IGF, Address address,
                      llvm::Attribute::ReadNone);
 }
 
+void irgen::emitTaskCancel(IRGenFunction &IGF, llvm::Value *task) {
+  if (task->getType() != IGF.IGM.SwiftTaskPtrTy) {
+    task = IGF.Builder.CreateBitCast(task, IGF.IGM.SwiftTaskPtrTy);
+  }
+
+  auto *call = IGF.Builder.CreateCall(IGF.IGM.getTaskCancelFn(), {task});
+  call->setDoesNotThrow();
+  call->setCallingConv(IGF.IGM.SwiftCC);
+  call->addAttribute(llvm::AttributeList::FunctionIndex,
+                     llvm::Attribute::ReadNone);
+}
+
 std::pair<Address, Size> irgen::emitAllocAsyncContext(IRGenFunction &IGF,
                                                       AsyncContextLayout layout,
                                                       llvm::Value *sizeValue,

lib/IRGen/GenCall.h

@@ -388,6 +388,8 @@ namespace irgen {
   Address emitTaskAlloc(IRGenFunction &IGF, llvm::Value *size,
                         Alignment alignment);
   void emitTaskDealloc(IRGenFunction &IGF, Address address, llvm::Value *size);
+  void emitTaskCancel(IRGenFunction &IGF, llvm::Value *task);
+
   /// Allocate task local storage for the specified layout but using the
   /// provided dynamic size.  Allowing the size to be specified dynamically is
   /// necessary for applies of thick functions the sizes of whose async contexts

lib/SIL/IR/OperandOwnership.cpp

@@ -231,6 +231,7 @@ CONSTANT_OWNERSHIP_INST(None, MustBeLive, UncheckedTakeEnumDataAddr)
 CONSTANT_OWNERSHIP_INST(None, MustBeLive, UnconditionalCheckedCastAddr)
 CONSTANT_OWNERSHIP_INST(None, MustBeLive, AllocValueBuffer)
 CONSTANT_OWNERSHIP_INST(None, MustBeLive, DeallocValueBuffer)
+
 #define NEVER_LOADABLE_CHECKED_REF_STORAGE(Name, ...)                          \
   CONSTANT_OWNERSHIP_INST(None, MustBeLive, Load##Name)
 #define ALWAYS_LOADABLE_CHECKED_REF_STORAGE(Name, ...)                         \
@@ -1033,6 +1034,8 @@ ANY_OWNERSHIP_BUILTIN(IntInstrprofIncrement)
   }
 CONSTANT_OWNERSHIP_BUILTIN(Owned, MustBeInvalidated, COWBufferForReading)
 CONSTANT_OWNERSHIP_BUILTIN(Owned, MustBeInvalidated, UnsafeGuaranteed)
+CONSTANT_OWNERSHIP_BUILTIN(Guaranteed, MustBeLive, CancelAsyncTask)
+
 #undef CONSTANT_OWNERSHIP_BUILTIN
 
 #define SHOULD_NEVER_VISIT_BUILTIN(ID)                              \

lib/SIL/IR/ValueOwnership.cpp

@@ -542,6 +542,7 @@ CONSTANT_OWNERSHIP_BUILTIN(None, TypePtrAuthDiscriminator)
 CONSTANT_OWNERSHIP_BUILTIN(None, IntInstrprofIncrement)
 CONSTANT_OWNERSHIP_BUILTIN(None, GlobalStringTablePointer)
 CONSTANT_OWNERSHIP_BUILTIN(Owned, GetCurrentAsyncTask)
+CONSTANT_OWNERSHIP_BUILTIN(None, CancelAsyncTask)
 
 #undef CONSTANT_OWNERSHIP_BUILTIN
 

lib/SIL/Utils/MemAccessUtils.cpp

@@ -1799,6 +1799,7 @@ static void visitBuiltinAddress(BuiltinInst *builtin,
     case BuiltinValueKind::PoundAssert:
     case BuiltinValueKind::IntInstrprofIncrement:
     case BuiltinValueKind::TSanInoutAccess:
+    case BuiltinValueKind::CancelAsyncTask:
       return;
 
     // General memory access to a pointer in first operand position.

lib/SILGen/SILGenBuiltin.cpp

@@ -1405,6 +1405,20 @@ static ManagedValue emitBuiltinGetCurrentAsyncTask(
   return SGF.emitManagedRValueWithEndLifetimeCleanup(apply);
 }
 
+// Emit SIL for the named builtin: getCurrentAsyncTask.
+static ManagedValue emitBuiltinCancelAsyncTask(
+    SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
+    ArrayRef<ManagedValue> args, SGFContext C) {
+  ASTContext &ctx = SGF.getASTContext();
+  auto argument = args[0].borrow(SGF, loc).forward(SGF);
+  auto apply = SGF.B.createBuiltin(
+      loc,
+      ctx.getIdentifier(getBuiltinName(BuiltinValueKind::CancelAsyncTask)),
+      SGF.getLoweredType(ctx.TheEmptyTupleType), SubstitutionMap(),
+      { argument });
+  return ManagedValue::forUnmanaged(apply);
+}
+
 Optional<SpecializedEmitter>
 SpecializedEmitter::forDecl(SILGenModule &SGM, SILDeclRef function) {
   // Only consider standalone declarations in the Builtin module.

lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp

@@ -165,6 +165,7 @@ static bool isBarrier(SILInstruction *inst) {
     case BuiltinValueKind::AssignTakeArray:
     case BuiltinValueKind::UnsafeGuaranteed:
     case BuiltinValueKind::UnsafeGuaranteedEnd:
+    case BuiltinValueKind::CancelAsyncTask:
       return true;
     }
   }

stdlib/public/Concurrency/Task.swift

@@ -102,7 +102,10 @@ extension Task {
   /// i.e. the task will run regardless of the handle still being present or not.
   /// Dropping a handle however means losing the ability to await on the task's result
   /// and losing the ability to cancel it.
-  public final class Handle<Success> {
+  @_frozen
+  public struct Handle<Success> {
+    private let task: Builtin.NativeObject
+
     /// Wait for the task to complete, returning (or throwing) its result.
     ///
     /// ### Priority
@@ -130,7 +133,7 @@ extension Task {
     /// their "actual work", however this is not a requirement nor is it guaranteed
     /// how and when tasks check for cancellation in general.
     public func cancel() {
-      fatalError("\(#function) not implemented yet.")
+      Builtin.cancelAsyncTask(task)
     }
   }
 }

stdlib/public/Concurrency/TaskStatus.cpp

@@ -505,7 +505,7 @@ JobPriority
 swift::swift_task_escalate(AsyncTask *task, JobPriority newPriority) {
   Optional<StatusRecordLockRecord> recordLockRecord;
 
-  // Fast path: check that the task's priority is not already at laest
+  // Fast path: check that the task's priority is not already at least
   // as high as the target.  The task's priority can only be modified
   // under the status record lock; it's possible that the priority could
   // be getting simultaneously escalated, but it's okay for us to return

test/IRGen/async/builtins.sil

@@ -1,5 +1,5 @@
-// RUN: %target-swift-frontend -enable-experimental-concurrency -enable-objc-interop  -primary-file %s -emit-ir | %FileCheck %s -DINT=i%target-ptrsize --check-prefixes=CHECK,CHECK-objc
-// RUN: %target-swift-frontend -enable-experimental-concurrency -disable-objc-interop -primary-file %s -emit-ir | %FileCheck %s -DINT=i%target-ptrsize --check-prefixes=CHECK,CHECK-native
+// RUN: %target-swift-frontend -enable-experimental-concurrency -enable-objc-interop  -primary-file %s -emit-ir -sil-verify-all | %FileCheck %s -DINT=i%target-ptrsize --check-prefixes=CHECK,CHECK-objc
+// RUN: %target-swift-frontend -enable-experimental-concurrency -disable-objc-interop -primary-file %s -emit-ir -sil-verify-all | %FileCheck %s -DINT=i%target-ptrsize --check-prefixes=CHECK,CHECK-native
 
 // REQUIRES: concurrency
 
@@ -17,3 +17,14 @@ bb0:
   end_lifetime %0 : $Builtin.NativeObject
   return %1 : $Builtin.NativeObject
 }
+
+// CHECK-LABEL: define hidden swiftcc void @cancel_task(%swift.refcounted* %0)
+// CHECK: [[TASK:%.*]] = bitcast %swift.refcounted* %0 to %swift.task*
+// CHECK-NEXT:  call swiftcc void @swift_task_cancel(%swift.task* [[TASK]])
+
+sil hidden [ossa] @cancel_task : $@convention(method) (@guaranteed Builtin.NativeObject) -> () {
+bb0(%0 : @guaranteed $Builtin.NativeObject):
+  %4 = builtin "cancelAsyncTask"(%0 : $Builtin.NativeObject) : $()
+  %5 = tuple ()
+  return %5 : $()
+}

test/SILGen/async_builtins.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -emit-silgen %s -module-name test -swift-version 5 -enable-experimental-concurrency -parse-stdlib | %FileCheck %s
+// RUN: %target-swift-frontend -emit-silgen %s -module-name test -swift-version 5 -enable-experimental-concurrency -parse-stdlib -sil-verify-all | %FileCheck %s
 // REQUIRES: concurrency
 
 struct X {
@@ -7,4 +7,10 @@ struct X {
     // CHECK: builtin "getCurrentAsyncTask"() : $Builtin.NativeObject
     return Builtin.getCurrentAsyncTask()
   }
+
+  // CHECK-LABEL: sil hidden [ossa] @$s4test1XV8doCancel4taskyBo_tF : $@convention(method) (@guaranteed Builtin.NativeObject, X) -> ()
+  func doCancel(task: Builtin.NativeObject) {
+    // CHECK: builtin "cancelAsyncTask"(%0 : $Builtin.NativeObject) : $()
+    Builtin.cancelAsyncTask(task)
+  }
 }