Adding tests for the various iree_task_t types.

Author: benvanik

iree/hal/local/task_queue.c

@@ -427,11 +427,18 @@ static iree_status_t iree_hal_task_queue_submit_batch(
   // NOTE: if we fail from here on we must drop the retire_cmd arena.
   iree_status_t status = iree_ok_status();
 
+  // A fence we'll use to detect when the entire submission has completed.
+  // TODO(benvanik): fold into the retire command.
+  iree_task_fence_t* fence = NULL;
+  status =
+      iree_task_executor_acquire_fence(queue->executor, &queue->scope, &fence);
+  iree_task_set_completion_task(&retire_cmd->task.header, &fence->header);
+
   // Task to fork and wait for unsatisfied semaphore dependencies.
   // This is optional and only required if we have previous submissions still
   // in-flight - if the queue is empty then we can directly schedule the waits.
   iree_hal_task_queue_wait_cmd_t* wait_cmd = NULL;
-  if (batch->wait_semaphores.count > 0) {
+  if (iree_status_is_ok(status) && batch->wait_semaphores.count > 0) {
     status = iree_hal_task_queue_wait_cmd_allocate(
         &queue->scope, &batch->wait_semaphores, &retire_cmd->arena, &wait_cmd);
   }

iree/task/BUILD

@@ -124,6 +124,25 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "task_tests",
+    srcs = [
+        "task_test_barrier.cc",
+        "task_test_call.cc",
+        "task_test_dispatch.cc",
+        "task_test_fence.cc",
+        "task_test_nop.cc",
+        "task_test_wait.cc",
+    ],
+    deps = [
+        ":task",
+        "//iree/base:api",
+        "//iree/task/testing:task_test",
+        "//iree/testing:gtest",
+        "//iree/testing:gtest_main",
+    ],
+)
+
 cc_test(
     name = "topology_test",
     srcs = ["topology_test.cc"],

iree/task/CMakeLists.txt

@@ -122,6 +122,24 @@ iree_cc_test(
     iree::testing::gtest_main
 )
 
+iree_cc_test(
+  NAME
+    task_tests
+  SRCS
+    "task_test_barrier.cc"
+    "task_test_call.cc"
+    "task_test_dispatch.cc"
+    "task_test_fence.cc"
+    "task_test_nop.cc"
+    "task_test_wait.cc"
+  DEPS
+    ::task
+    iree::base::api
+    iree::task::testing::task_test
+    iree::testing::gtest
+    iree::testing::gtest_main
+)
+
 iree_cc_test(
   NAME
     topology_test

iree/task/executor.c

@@ -86,16 +86,18 @@ iree_status_t iree_task_executor_create(
   // executor and since we know the precise lifetime of them we can keep them
   // entirely within the system here.
   if (iree_status_is_ok(status)) {
-    status = iree_task_pool_initialize(
-        allocator, sizeof(iree_task_dispatch_slice_t),
-        worker_count * IREE_TASK_EXECUTOR_INITIAL_SLICE_RESERVATION_PER_WORKER,
-        &executor->slice_task_pool);
+    status = iree_task_pool_initialize(allocator, sizeof(iree_task_fence_t), 8,
+                                       &executor->fence_task_pool);
   }
   if (iree_status_is_ok(status)) {
     status = iree_task_pool_initialize(
-        allocator, sizeof(iree_task_dispatch_shard_t),
-        worker_count * IREE_TASK_EXECUTOR_INITIAL_SHARD_RESERVATION_PER_WORKER,
-        &executor->shard_task_pool);
+        allocator,
+        iree_max(sizeof(iree_task_dispatch_shard_t),
+                 sizeof(iree_task_dispatch_slice_t)),
+        worker_count *
+            iree_max(IREE_TASK_EXECUTOR_INITIAL_SHARD_RESERVATION_PER_WORKER,
+                     IREE_TASK_EXECUTOR_INITIAL_SLICE_RESERVATION_PER_WORKER),
+        &executor->dispatch_task_pool);
   }
 
   // Bring up the workers; the threads will be created here but be suspended
@@ -169,8 +171,8 @@ static void iree_task_executor_destroy(iree_task_executor_t* executor) {
   iree_slim_mutex_deinitialize(&executor->coordinator_mutex);
   iree_atomic_task_slist_deinitialize(&executor->incoming_ready_slist);
   iree_atomic_task_slist_deinitialize(&executor->incoming_waiting_slist);
-  iree_task_pool_deinitialize(&executor->slice_task_pool);
-  iree_task_pool_deinitialize(&executor->shard_task_pool);
+  iree_task_pool_deinitialize(&executor->fence_task_pool);
+  iree_task_pool_deinitialize(&executor->dispatch_task_pool);
   iree_allocator_free(executor->allocator, executor);
 
   IREE_TRACE_ZONE_END(z0);
@@ -188,6 +190,19 @@ void iree_task_executor_release(iree_task_executor_t* executor) {
   }
 }
 
+iree_status_t iree_task_executor_acquire_fence(iree_task_executor_t* executor,
+                                               iree_task_scope_t* scope,
+                                               iree_task_fence_t** out_fence) {
+  *out_fence = NULL;
+  iree_task_fence_t* fence = NULL;
+  IREE_RETURN_IF_ERROR(iree_task_pool_acquire(&executor->fence_task_pool,
+                                              (iree_task_t**)&fence));
+  iree_task_fence_initialize(scope, fence);
+  fence->header.pool = &executor->fence_task_pool;
+  *out_fence = fence;
+  return iree_ok_status();
+}
+
 // Schedules a generic task to a worker matching its affinity.
 // The task will be posted to the worker mailbox and available for the worker to
 // begin processing as soon as the |post_batch| is submitted.
@@ -262,11 +277,11 @@ void iree_task_executor_schedule_ready_tasks(
         } else {
           if (task->flags & IREE_TASK_FLAG_DISPATCH_SLICED) {
             iree_task_dispatch_issue_sliced((iree_task_dispatch_t*)task,
-                                            &executor->slice_task_pool,
+                                            &executor->dispatch_task_pool,
                                             pending_submission, post_batch);
           } else {
             iree_task_dispatch_issue_sharded((iree_task_dispatch_t*)task,
-                                             &executor->shard_task_pool,
+                                             &executor->dispatch_task_pool,
                                              pending_submission, post_batch);
           }
         }
@@ -520,17 +535,7 @@ static void iree_task_executor_wait_any_task(
 void iree_task_executor_coordinate(iree_task_executor_t* executor,
                                    iree_task_worker_t* current_worker,
                                    bool speculative) {
-  if (speculative) {
-    if (!iree_slim_mutex_try_lock(&executor->coordinator_mutex)) {
-      // Another thread is already holding the coordination lock.
-      // Return to the caller to wait for it to finish.
-      // TODO(benvanik): spin here if it's likely we'll have work after the
-      // other coordinator finishes - that way we don't enter the wait.
-      return;
-    }
-  } else {
-    iree_slim_mutex_lock(&executor->coordinator_mutex);
-  }
+  iree_slim_mutex_lock(&executor->coordinator_mutex);
   IREE_TRACE_ZONE_BEGIN(z0);
 
   // We may be adding tasks/waiting/etc on each pass through coordination - to

iree/task/executor.h

@@ -319,6 +319,11 @@ void iree_task_executor_retain(iree_task_executor_t* executor);
 // Releases the given |executor| from the caller.
 void iree_task_executor_release(iree_task_executor_t* executor);
 
+// Acquires a fence for the given |scope| from the executor fence pool.
+iree_status_t iree_task_executor_acquire_fence(iree_task_executor_t* executor,
+                                               iree_task_scope_t* scope,
+                                               iree_task_fence_t** out_fence);
+
 // TODO(benvanik): scheduling mode mutation, compute quota control, etc.
 
 // Submits a batch of tasks for execution.

iree/task/executor_impl.h

@@ -50,8 +50,8 @@ struct iree_task_executor_s {
   // Pools of transient dispatch tasks shared across all workers.
   // Depending on configuration the task pool may allocate after creation using
   // the allocator provided upon executor creation.
-  iree_task_pool_t slice_task_pool;
-  iree_task_pool_t shard_task_pool;
+  iree_task_pool_t fence_task_pool;
+  iree_task_pool_t dispatch_task_pool;
 
   // A list of incoming tasks that are ready to execute immediately.
   // The list is LIFO and we require that task lists are reversed by the

iree/task/executor_test.cc

@@ -123,7 +123,7 @@ TEST(ExecutorTest, Any) {
           },
           0),
       workgroup_size_1, workgroup_count_1, &dispatch1);
-  // dispatch1.header.flags |= IREE_TASK_FLAG_DISPATCH_SLICED;
+  dispatch1.header.flags |= IREE_TASK_FLAG_DISPATCH_SLICED;
 
   //
   iree_task_call_t call1;
@@ -155,9 +155,9 @@ TEST(ExecutorTest, Any) {
 #endif
 
   // fence
-  iree_task_fence_t fence0;
-  iree_task_fence_initialize(&scope_a, &fence0);
-  iree_task_set_completion_task(&call1.header, &fence0.header);
+  iree_task_fence_t* fence0 = NULL;
+  IREE_CHECK_OK(iree_task_executor_acquire_fence(executor, &scope_a, &fence0));
+  iree_task_set_completion_task(&call1.header, &fence0->header);
 
   //
   iree_task_submission_t sub0;

iree/task/scope.c

@@ -116,7 +116,7 @@ void iree_task_scope_fail(iree_task_scope_t* scope, iree_task_t* task,
 
 bool iree_task_scope_is_idle(iree_task_scope_t* scope) {
   return iree_atomic_load_int32(&scope->pending_submissions,
-                                iree_memory_order_relaxed) == 0;
+                                iree_memory_order_acquire) == 0;
 }
 
 iree_status_t iree_task_scope_wait_idle(iree_task_scope_t* scope,

iree/task/task.c

@@ -100,7 +100,7 @@ void iree_task_discard(iree_task_t* task, iree_task_list_t* discard_worklist) {
       // TODO(benvanik): signal as error.
       // iree_task_fence_t* fence_task = (iree_task_fence_t*)task;
       iree_atomic_fetch_sub_int32(&task->scope->pending_submissions, 1,
-                                  iree_memory_order_relaxed);
+                                  iree_memory_order_release);
       break;
     }
     case IREE_TASK_TYPE_WAIT:
@@ -232,7 +232,7 @@ void iree_task_fence_initialize(iree_task_scope_t* scope,
                                 iree_task_fence_t* out_task) {
   iree_task_initialize(IREE_TASK_TYPE_FENCE, scope, &out_task->header);
   iree_atomic_fetch_add_int32(&scope->pending_submissions, 1,
-                              iree_memory_order_relaxed);
+                              iree_memory_order_release);
 }
 
 void iree_task_fence_retire(iree_task_fence_t* task,
@@ -408,6 +408,14 @@ void iree_task_dispatch_issue_sliced(iree_task_dispatch_t* dispatch_task,
     memcpy(workgroup_count, dispatch_task->workgroup_count.value,
            sizeof(workgroup_count));
   }
+  uint32_t total_workgroup_count =
+      workgroup_count[0] * workgroup_count[1] * workgroup_count[2];
+  if (total_workgroup_count == 0) {
+    // No workgroups to execute - bail early.
+    iree_task_dispatch_retire(dispatch_task, pending_submission);
+    IREE_TRACE_ZONE_END(z0);
+    return;
+  }
 
 #if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
   char xyz_string[32];
@@ -449,12 +457,15 @@ void iree_task_dispatch_issue_sliced(iree_task_dispatch_t* dispatch_task,
         workgroup_base[1] = slice_y * tiles_per_slice_y;
         workgroup_base[2] = slice_z * tiles_per_slice_z;
         uint32_t workgroup_range[3];
-        workgroup_range[0] = iree_min(
-            workgroup_count[0], workgroup_base[0] + tiles_per_slice_x - 1);
-        workgroup_range[1] = iree_min(
-            workgroup_count[1], workgroup_base[1] + tiles_per_slice_y - 1);
-        workgroup_range[2] = iree_min(
-            workgroup_count[2], workgroup_base[2] + tiles_per_slice_z - 1);
+        workgroup_range[0] = iree_min(workgroup_count[0],
+                                      workgroup_base[0] + tiles_per_slice_x) -
+                             1;
+        workgroup_range[1] = iree_min(workgroup_count[1],
+                                      workgroup_base[1] + tiles_per_slice_y) -
+                             1;
+        workgroup_range[2] = iree_min(workgroup_count[2],
+                                      workgroup_base[2] + tiles_per_slice_z) -
+                             1;
 
         // Allocate and initialize the slice.
         iree_task_dispatch_slice_t* slice_task =
@@ -789,10 +800,10 @@ iree_status_t iree_task_dispatch_shard_execute(
       // TODO(benvanik): faster math here, especially knowing we pull off N
       // sequential indices per reservation.
       uint32_t tile_i = tile_index;
-      tile_context.workgroup_xyz[0] = tile_i % (workgroup_count_x + 1);
-      tile_i /= (workgroup_count_x + 1);
-      tile_context.workgroup_xyz[1] = tile_i % (workgroup_count_y + 1);
-      tile_i /= (workgroup_count_y + 1);
+      tile_context.workgroup_xyz[0] = tile_i % workgroup_count_x;
+      tile_i /= workgroup_count_x;
+      tile_context.workgroup_xyz[1] = tile_i % workgroup_count_y;
+      tile_i /= workgroup_count_y;
       tile_context.workgroup_xyz[2] = tile_i;
 
       IREE_TRACE_ZONE_BEGIN_NAMED(z_tile,