Merge pull request #41281 from apple/rokhinip/88600541-task-priority-escalation-arm64_32

Task Escalation support on arm64_32 Apple platforms

Author: rokhinip

include/swift/ABI/Task.h

@@ -317,8 +317,9 @@ class AsyncTask : public Job {
   /// Flag that this task is now suspended.
   void flagAsSuspended();
 
-  /// Flag that the task is to be enqueued on the provided executor
-  void flagAsEnqueuedOnExecutor(ExecutorRef newExecutor);
+  /// Flag that the task is to be enqueued on the provided executor and actually
+  /// enqueue it
+  void flagAsAndEnqueueOnExecutor(ExecutorRef newExecutor);
 
   /// Flag that this task is now completed. This normally does not do anything
   /// but can be used to locally insert logging.

include/swift/Reflection/RuntimeInternals.h

@@ -108,11 +108,12 @@ struct ActiveTaskStatusWithoutEscalation {
 
 template <typename Runtime, typename ActiveTaskStatus>
 struct AsyncTaskPrivateStorage {
+  typename Runtime::StoredPointer ExclusivityAccessSet[2];
   ActiveTaskStatus Status;
   StackAllocator<Runtime> Allocator;
   typename Runtime::StoredPointer Local;
-  typename Runtime::StoredPointer ExclusivityAccessSet[2];
   uint32_t Id;
+  uint32_t BasePriority;
 };
 
 template <typename Runtime, typename ActiveTaskStatus>

include/swift/Runtime/Concurrency.h

@@ -45,7 +45,7 @@
 #ifndef SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
 #if SWIFT_CONCURRENCY_ENABLE_DISPATCH && \
     __has_include(<dispatch/swift_concurrency_private.h>) && __APPLE__ && \
-    SWIFT_POINTER_IS_8_BYTES
+    (defined(__arm64__) || defined(__x86_64__))
 #define SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION 1
 #else
 #define SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION 0

stdlib/public/Concurrency/Actor.cpp

@@ -1616,7 +1616,7 @@ static void swift_task_switchImpl(SWIFT_ASYNC_CONTEXT AsyncContext *resumeContex
   SWIFT_TASK_DEBUG_LOG("switch failed, task %p enqueued on executor %p", task,
                        newExecutor.getIdentity());
 
-  task->flagAsEnqueuedOnExecutor(newExecutor);
+  task->flagAsAndEnqueueOnExecutor(newExecutor);
   _swift_task_clearCurrent();
 }
 

stdlib/public/Concurrency/Task.cpp

@@ -148,7 +148,7 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
 
       // Escalate the priority of this task based on the priority
       // of the waiting task.
-      auto status = waitingTask->_private().Status.load(std::memory_order_relaxed);
+      auto status = waitingTask->_private()._status().load(std::memory_order_relaxed);
       swift_task_escalate(this, status.getStoredPriority());
 
       _swift_task_clearCurrent();
@@ -234,7 +234,7 @@ void AsyncTask::completeFuture(AsyncContext *context) {
 
     // Enqueue the waiter on the global executor.
     // TODO: allow waiters to fill in a suggested executor
-    waitingTask->flagAsEnqueuedOnExecutor(ExecutorRef::generic());
+    waitingTask->flagAsAndEnqueueOnExecutor(ExecutorRef::generic());
 
     // Move to the next task.
     waitingTask = nextWaitingTask;
@@ -491,7 +491,7 @@ const void *AsyncTask::getResumeFunctionForLogging() {
 JobPriority swift::swift_task_currentPriority(AsyncTask *task)
 {
   // This is racey but this is to be used in an API is inherently racey anyways.
-  auto oldStatus = task->_private().Status.load(std::memory_order_relaxed);
+  auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
   return oldStatus.getStoredPriority();
 }
 
@@ -847,7 +847,7 @@ static AsyncTaskAndContext swift_task_create_commonImpl(
   // If we're supposed to enqueue the task, do so now.
   if (taskCreateFlags.enqueueJob()) {
     swift_retain(task);
-    task->flagAsEnqueuedOnExecutor(executor);
+    task->flagAsAndEnqueueOnExecutor(executor);
   }
 
   return {task, initialContext};
@@ -1020,7 +1020,7 @@ SWIFT_CC(swift)
 static void
 swift_task_enqueueTaskOnExecutorImpl(AsyncTask *task, ExecutorRef executor)
 {
-  task->flagAsEnqueuedOnExecutor(executor);
+  task->flagAsAndEnqueueOnExecutor(executor);
 }
 
 SWIFT_CC(swift)
@@ -1151,7 +1151,7 @@ static void resumeTaskAfterContinuation(AsyncTask *task,
   // to make a stronger best-effort attempt to catch racing attempts to
   // resume the continuation?
 
-  task->flagAsEnqueuedOnExecutor(context->ResumeToExecutor);
+  task->flagAsAndEnqueueOnExecutor(context->ResumeToExecutor);
 }
 
 SWIFT_CC(swift)

stdlib/public/Concurrency/TaskGroup.cpp

@@ -622,7 +622,7 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
         _swift_tsan_acquire(static_cast<Job *>(waitingTask));
 
         // TODO: allow the caller to suggest an executor
-        waitingTask->flagAsEnqueuedOnExecutor(ExecutorRef::generic());
+        waitingTask->flagAsAndEnqueueOnExecutor(ExecutorRef::generic());
         return;
       } // else, try again
     }

stdlib/public/Concurrency/TaskPrivate.h

@@ -231,7 +231,28 @@ class TaskFutureWaitAsyncContext : public AsyncContext {
 /// only need to do double wide atomics if we need to reach for the
 /// StatusRecord pointers and therefore have to update the flags at the same
 /// time.
-class alignas(sizeof(void*) * 2) ActiveTaskStatus {
+///
+/// Size requirements:
+///     On 64 bit systems or if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION=1,
+///     the field is 16 bytes long.
+///
+///     Otherwise, it is 8 bytes long.
+///
+/// Alignment requirements:
+///     On 64 bit systems or if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION=1,
+///     this 16-byte field needs to be 16 byte aligned to be able to do aligned
+///     atomic stores field.
+///
+///     On all other systems, it needs to be 8 byte aligned for the atomic
+///     stores.
+///
+///     As a result of varying alignment needs, we've marked the class as
+///     needing 2-word alignment but on arm64_32 with
+///     SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION=1, 16 byte alignment is
+///     achieved through careful arrangement of the storage for this in the
+///     AsyncTask::PrivateStorage. The additional alignment requirements are
+///     enforced by static asserts below.
+class alignas(2 * sizeof(void*)) ActiveTaskStatus {
   enum : uint32_t {
     /// The max priority of the task. This is always >= basePriority in the task
     PriorityMask = 0xFF,
@@ -465,12 +486,12 @@ class alignas(sizeof(void*) * 2) ActiveTaskStatus {
 };
 
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION && SWIFT_POINTER_IS_4_BYTES
-static_assert(sizeof(ActiveTaskStatus) == 4 * sizeof(uintptr_t),
-  "ActiveTaskStatus is 4 words large");
+#define ACTIVE_TASK_STATUS_SIZE (4 * (sizeof(uintptr_t)))
 #else
-static_assert(sizeof(ActiveTaskStatus) == 2 * sizeof(uintptr_t),
-  "ActiveTaskStatus is 2 words large");
+#define ACTIVE_TASK_STATUS_SIZE (2 * (sizeof(uintptr_t)))
 #endif
+static_assert(sizeof(ActiveTaskStatus) == ACTIVE_TASK_STATUS_SIZE,
+  "ActiveTaskStatus is of incorrect size");
 
 /// The size of an allocator slab. We want the full allocation to fit into a
 /// 1024-byte malloc quantum. We subtract off the slab header size, plus a
@@ -483,9 +504,16 @@ using TaskAllocator = StackAllocator<SlabCapacity, &TaskAllocatorSlabMetadata>;
 
 /// Private storage in an AsyncTask object.
 struct AsyncTask::PrivateStorage {
-  /// The currently-active information about cancellation.
-  /// Currently two words.
-  swift::atomic<ActiveTaskStatus> Status;
+  /// State inside the AsyncTask whose state is only managed by the exclusivity
+  /// runtime in stdlibCore. We zero initialize to provide a safe initial value,
+  /// but actually initialize its bit state to a const global provided by
+  /// libswiftCore so that libswiftCore can control the layout of our initial
+  /// state.
+  uintptr_t ExclusivityAccessSet[2] = {0, 0};
+
+  /// Storage for the ActiveTaskStatus. See doc for ActiveTaskStatus for size
+  /// and alignment requirements.
+  alignas(2 * sizeof(void *)) char StatusStorage[sizeof(ActiveTaskStatus)];
 
   /// The allocator for the task stack.
   /// Currently 2 words + 8 bytes.
@@ -495,13 +523,6 @@ struct AsyncTask::PrivateStorage {
   /// Currently one word.
   TaskLocal::Storage Local;
 
-  /// State inside the AsyncTask whose state is only managed by the exclusivity
-  /// runtime in stdlibCore. We zero initialize to provide a safe initial value,
-  /// but actually initialize its bit state to a const global provided by
-  /// libswiftCore so that libswiftCore can control the layout of our initial
-  /// state.
-  uintptr_t ExclusivityAccessSet[2] = {0, 0};
-
   /// The top 32 bits of the task ID. The bottom 32 bits are in Job::Id.
   uint32_t Id;
 
@@ -516,19 +537,22 @@ struct AsyncTask::PrivateStorage {
   // Always create an async task with max priority in ActiveTaskStatus = base
   // priority. It will be updated later if needed.
   PrivateStorage(JobPriority basePri)
-      : Status(ActiveTaskStatus(basePri)), Local(TaskLocal::Storage()),
-        BasePriority(basePri) {}
+      : Local(TaskLocal::Storage()), BasePriority(basePri) {
+    _status().store(ActiveTaskStatus(basePri), std::memory_order_relaxed);
+  }
 
   PrivateStorage(JobPriority basePri, void *slab, size_t slabCapacity)
-      : Status(ActiveTaskStatus(basePri)), Allocator(slab, slabCapacity),
-        Local(TaskLocal::Storage()), BasePriority(basePri) {}
+      : Allocator(slab, slabCapacity), Local(TaskLocal::Storage()),
+          BasePriority(basePri) {
+    _status().store(ActiveTaskStatus(basePri), std::memory_order_relaxed);
+  }
 
   /// Called on the thread that was previously executing the task that we are
   /// now trying to complete.
   void complete(AsyncTask *task) {
     // Drain unlock the task and remove any overrides on thread as a
     // result of the task
-    auto oldStatus = task->_private().Status.load(std::memory_order_relaxed);
+    auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
     while (true) {
       // Task is completing, it shouldn't have any records and therefore
       // cannot be status record locked.
@@ -544,7 +568,7 @@ struct AsyncTask::PrivateStorage {
 
       // This can fail since the task can still get concurrently cancelled or
       // escalated.
-      if (task->_private().Status.compare_exchange_weak(oldStatus, newStatus,
+      if (task->_private()._status().compare_exchange_weak(oldStatus, newStatus,
               /* success */ std::memory_order_relaxed,
               /* failure */ std::memory_order_relaxed)) {
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
@@ -563,12 +587,21 @@ struct AsyncTask::PrivateStorage {
 
     this->~PrivateStorage();
   }
+
+  swift::atomic<ActiveTaskStatus> &_status() {
+    return reinterpret_cast<swift::atomic<ActiveTaskStatus>&> (this->StatusStorage);
+  }
+
+  const swift::atomic<ActiveTaskStatus> &_status() const {
+    return reinterpret_cast<const swift::atomic<ActiveTaskStatus>&> (this->StatusStorage);
+  }
 };
 
-static_assert(sizeof(AsyncTask::PrivateStorage)
-                <= sizeof(AsyncTask::OpaquePrivateStorage) &&
-              alignof(AsyncTask::PrivateStorage)
-                <= alignof(AsyncTask::OpaquePrivateStorage),
+// It will be aligned to 2 words on all platforms. On arm64_32, we have an
+// additional requirement where it is aligned to 4 words.
+static_assert(((offsetof(AsyncTask, Private) + offsetof(AsyncTask::PrivateStorage, StatusStorage)) % ACTIVE_TASK_STATUS_SIZE == 0),
+   "StatusStorage is not aligned in the AsyncTask");
+static_assert(sizeof(AsyncTask::PrivateStorage) <= sizeof(AsyncTask::OpaquePrivateStorage),
               "Task-private storage doesn't fit in reserved space");
 
 inline AsyncTask::PrivateStorage &
@@ -601,7 +634,7 @@ inline const AsyncTask::PrivateStorage &AsyncTask::_private() const {
 }
 
 inline bool AsyncTask::isCancelled() const {
-  return _private().Status.load(std::memory_order_relaxed)
+  return _private()._status().load(std::memory_order_relaxed)
                           .isCancelled();
 }
 
@@ -614,7 +647,7 @@ inline void AsyncTask::flagAsRunning() {
   qos_class_t overrideFloor = threadOverrideInfo.override_qos_floor;
 retry:;
 #endif
-  auto oldStatus = _private().Status.load(std::memory_order_relaxed);
+  auto oldStatus = _private()._status().load(std::memory_order_relaxed);
   while (true) {
     // We can get here from being suspended or being enqueued
     assert(!oldStatus.isRunning());
@@ -638,7 +671,7 @@ retry:;
     newStatus = newStatus.withoutStoredPriorityEscalation();
     newStatus = newStatus.withoutEnqueued();
 
-    if (_private().Status.compare_exchange_weak(oldStatus, newStatus,
+    if (_private()._status().compare_exchange_weak(oldStatus, newStatus,
              /* success */ std::memory_order_relaxed,
              /* failure */ std::memory_order_relaxed)) {
       newStatus.traceStatusChanged(this);
@@ -666,10 +699,10 @@ retry:;
 /// original enqueueing thread.
 ///
 /// rdar://88366470 (Direct handoff behaviour when tasks switch executors)
-inline void AsyncTask::flagAsEnqueuedOnExecutor(ExecutorRef newExecutor) {
+inline void AsyncTask::flagAsAndEnqueueOnExecutor(ExecutorRef newExecutor) {
 
-  SWIFT_TASK_DEBUG_LOG("%p->flagAsEnqueuedOnExecutor()", this);
-  auto oldStatus = _private().Status.load(std::memory_order_relaxed);
+  SWIFT_TASK_DEBUG_LOG("%p->flagAsAndEnqueueOnExecutor()", this);
+  auto oldStatus = _private()._status().load(std::memory_order_relaxed);
   auto newStatus = oldStatus;
 
   while (true) {
@@ -684,7 +717,7 @@ inline void AsyncTask::flagAsEnqueuedOnExecutor(ExecutorRef newExecutor) {
     newStatus = newStatus.withoutStoredPriorityEscalation();
     newStatus = newStatus.withEnqueued();
 
-    if (_private().Status.compare_exchange_weak(oldStatus, newStatus,
+    if (_private()._status().compare_exchange_weak(oldStatus, newStatus,
             /* success */std::memory_order_relaxed,
             /* failure */std::memory_order_relaxed)) {
       break;
@@ -715,7 +748,7 @@ inline void AsyncTask::flagAsEnqueuedOnExecutor(ExecutorRef newExecutor) {
 
 inline void AsyncTask::flagAsSuspended() {
   SWIFT_TASK_DEBUG_LOG("%p->flagAsSuspended()", this);
-  auto oldStatus = _private().Status.load(std::memory_order_relaxed);
+  auto oldStatus = _private()._status().load(std::memory_order_relaxed);
   auto newStatus = oldStatus;
   while (true) {
     // We can only be suspended if we were previously running. See state
@@ -725,7 +758,7 @@ inline void AsyncTask::flagAsSuspended() {
     newStatus = oldStatus.withRunning(false);
     newStatus = newStatus.withoutStoredPriorityEscalation();
 
-    if (_private().Status.compare_exchange_weak(oldStatus, newStatus,
+    if (_private()._status().compare_exchange_weak(oldStatus, newStatus,
             /* success */std::memory_order_relaxed,
             /* failure */std::memory_order_relaxed)) {
       break;