Allow AGAIN while batching

docs/doxygen/task-batching.md

@@ -1,3 +1,7 @@
+<!--
+Copyright (c) 2026 NVIDIA Corporation.  All rights reserved.
+-->
+
 Task Batching {#task_batching}
 ==============
 
@@ -61,9 +65,13 @@ Recommended collection helper
 
 The preferred interface for GPU submit hooks is
 `parsec_gpu_task_collect_batch()`. The runtime passes the submit hook a
-singleton `parsec_gpu_task_t *gpu_task`. The hook calls the collector with a
-callback that decides, for each task currently pending on the same stream,
-whether that candidate can be added to the batch headed by `gpu_task`.
+singleton `parsec_gpu_task_t *gpu_task` on its initial invocation. The hook
+calls the collector with a callback that decides, for each task currently
+pending on the same stream, whether that candidate can be added to the batch
+headed by `gpu_task`. After a batched hook returns `PARSEC_HOOK_RETURN_AGAIN`,
+the continuation invocation receives the same intact ring. Calling the
+collector again is safe: it returns the existing follower count without
+scanning or modifying the pending FIFO.
 
 The callback has the type `parsec_gpu_task_batch_cb_t` and receives:
 
@@ -73,12 +81,19 @@ The callback has the type `parsec_gpu_task_batch_cb_t` and receives:
 
 The callback return value controls the iterator:
 
-- negative: stop immediately and return that error code;
-- zero: remove `candidate` from the pending FIFO and append it to
-  `batch_head`'s task ring;
-- positive: leave `candidate` pending and continue to the next pending task.
+- `PARSEC_GPU_TASK_BATCH_ACCEPT`: remove `candidate` from the pending FIFO and
+  append it to `batch_head`'s task ring;
+- `PARSEC_GPU_TASK_BATCH_REJECT`: leave `candidate` pending and continue to the
+  next pending task;
+- `PARSEC_GPU_TASK_BATCH_STOP`: end collection successfully, leaving the
+  current candidate and all unvisited candidates pending.
 
-The callback must not modify `gpu_stream->fifo_pending` directly.
+The collector does not impose a scan or batch-size bound. The submit hook owns
+that policy and its callback should return `PARSEC_GPU_TASK_BATCH_STOP` once it
+has accepted enough work. The callback executes while
+`gpu_stream->fifo_pending` is locked, so it must remain short and nonblocking.
+It must not modify or acquire that FIFO, call the collector recursively, or
+otherwise reenter pending-task operations on the same stream.
 
 Example:
 
@@ -93,9 +108,9 @@ gemm_batch_match(parsec_gpu_task_t *candidate,
     if( (batch_head->ec->task_class == candidate->ec->task_class) &&
         (batch_head->ec->selected_chore == candidate->ec->selected_chore) &&
         (batch_head->ec->selected_device == candidate->ec->selected_device) ) {
-        return 0;
+        return PARSEC_GPU_TASK_BATCH_ACCEPT;
     }
-    return 1;
+    return PARSEC_GPU_TASK_BATCH_REJECT;
 }
 
 int
@@ -129,20 +144,183 @@ gemm_kernel_cuda(parsec_device_gpu_module_t *gpu_device,
 }
 ```
 
-`parsec_gpu_task_collect_batch()` returns the number of additional tasks appended
-to the ring on success, or the negative callback error. A return value of 0
-means no task was batched, either because no compatible pending task was found,
-because batching is disabled or unsupported by the head task's selected device,
-or because the head task's selected incarnation is not batch-capable. Tasks
-accepted before an error remain attached to `gpu_task`; tasks not accepted
-remain in `gpu_stream->fifo_pending`.
+`parsec_gpu_task_collect_batch()` returns the number of additional tasks
+appended to the ring on success, including when the callback stops collection.
+A return value of 0 means no task was batched, either because the callback
+stopped before accepting one, no compatible pending task was found, batching
+is disabled or unsupported by the head task's selected device, or the head
+task's selected incarnation is not batch-capable. Any callback value outside
+the three actions above returns `PARSEC_HOOK_RETURN_ERROR`. Tasks accepted
+before that error remain attached to `gpu_task`; tasks not accepted remain in
+`gpu_stream->fifo_pending`.
 
 The submit hook does not need a completion callback merely to return the ring to
 the runtime. If a batched submit hook returns a non-singleton task ring, the GPU
 progress engine automatically chains that ring into the next stream's pending
 FIFO after the recorded device event completes. The normal data retrieval,
 epilog, ownership, pushout, and task completion paths then process the tasks one
-at a time.
+at a time. Every stream insertion merges all ring members according to the
+stream's priority policy; equal-priority members retain their ring order.
+
+The submit result settles ownership of the complete ring:
+
+- `PARSEC_HOOK_RETURN_DONE` commits every member to the submitted device work;
+- `PARSEC_HOOK_RETURN_AGAIN` commits the ring as continuation state, waits for
+  work queued by the hook, and re-enters the hook with the same ring after the
+  stream event completes. The hook must have submitted progress for every ring
+  member before returning `AGAIN`;
+- `PARSEC_HOOK_RETURN_NEXT` immediately restores the followers and applies the
+  existing singleton `NEXT` path to the head;
+- `PARSEC_HOOK_RETURN_ASYNC` transfers every execution context in the ring to
+  the hook, which must eventually complete or reschedule each context. The GPU
+  engine releases all corresponding device-task wrappers;
+- `PARSEC_HOOK_RETURN_ERROR` and `PARSEC_HOOK_RETURN_DISABLE` quiesce the
+  affected execution stream and clean every member of the failed batch before
+  propagating the terminal status.
+
+Device-wide recovery after `PARSEC_HOOK_RETURN_DISABLE` is not implemented.
+The terminal path settles the batch that observed the error, but it does not
+drain other stream FIFOs, event slots, or the device-wide pending queue.
+
+Because `AGAIN` preserves a collected ring, hooks must check predictable
+resource constraints before calling `parsec_gpu_task_collect_batch()`. Returning
+`AGAIN` before collection is a normal singleton retry; returning it afterward
+means that the collected batch has started and must continue as one unit.
+
+The runtime never disbands a ring on the `AGAIN` path. A submit hook that no
+longer wants to continue the batch must split the ring itself before returning:
+detach the followers, restore the head as a valid singleton, and explicitly
+give every detached wrapper a new owner. For example, followers may be returned
+to `gpu_stream->fifo_pending` using that stream's priority-preserving insertion
+policy, or the hook may explicitly retain responsibility for completing the
+execution contexts and releasing their wrappers. Merely breaking the links is
+insufficient because it leaves the followers unreachable. If the hook returns
+`AGAIN` after disbanding, only the ring or singleton it leaves attached to
+`gpu_task` is retained by the event.
+
+Ownership examples
+------------------
+
+### Coroutine-style `AGAIN`
+
+Calling the collector on every coroutine entry is safe. The initial call may
+build a batch; continuation calls only report the followers already attached
+to the same ring. The application-specific progress routine must submit work
+for every member before reporting that another continuation is needed:
+
+```c
+int
+submit_coroutine_batch(parsec_device_gpu_module_t *gpu_device,
+                       parsec_gpu_task_t *gpu_task,
+                       parsec_gpu_exec_stream_t *gpu_stream)
+{
+    int nb_batched;
+
+    (void)gpu_device;
+    nb_batched = parsec_gpu_task_collect_batch(gpu_stream, gpu_task,
+                                               batch_match, NULL);
+    if( nb_batched < 0 ) {
+        return nb_batched;
+    }
+
+    /* Application code: resume and submit one progress step for every member
+     * of gpu_task's ring. The same ring returns here after AGAIN completes.
+     */
+    return submit_coroutine_step(gpu_task, gpu_stream)
+               ? PARSEC_HOOK_RETURN_AGAIN
+               : PARSEC_HOOK_RETURN_DONE;
+}
+```
+
+### Batched `ASYNC`
+
+`ASYNC` transfers every execution context, not the device-task wrappers, to the
+hook. Save or enqueue each `ec` before returning; the GPU engine releases the
+wrappers after the hook returns. The new owner must eventually complete or
+reschedule every saved context:
+
+```c
+parsec_gpu_task_t *current = gpu_task;
+
+do {
+    parsec_gpu_task_t *next =
+        (parsec_gpu_task_t *)current->list_item.list_next;
+    parsec_task_t *ec = current->ec;
+
+    PARSEC_LIST_ITEM_SINGLETON(ec);
+    async_owner_enqueue(ec);  /* Eventually complete or reschedule ec. */
+    current = next;
+} while( current != gpu_task );
+
+return PARSEC_HOOK_RETURN_ASYNC;
+```
+
+The asynchronous owner must retain the execution contexts, not pointers to
+`parsec_gpu_task_t`, because those wrappers become invalid after the return.
+
+### Manually disbanding an `AGAIN` ring
+
+`parsec_list_item_ring_chop()` reconnects the followers but deliberately leaves
+the removed head's links invalid. Restore the head as a singleton, refresh the
+followers' priority snapshots, and return the follower ring to the pending
+stream with its configured ordering policy. Call this helper from the submit
+hook after the collector returns, never from the collector callback while the
+pending FIFO is locked:
+
+```c
+static void
+requeue_batch_followers(parsec_gpu_exec_stream_t *gpu_stream,
+                        parsec_gpu_task_t *gpu_task)
+{
+    parsec_list_item_t *followers;
+
+    followers = parsec_list_item_ring_chop(&gpu_task->list_item);
+    PARSEC_LIST_ITEM_SINGLETON(&gpu_task->list_item);
+    if( NULL == followers ) {
+        return;
+    }
+
+#if PARSEC_GPU_USE_PRIORITIES
+    parsec_gpu_task_t *first = (parsec_gpu_task_t *)followers;
+    parsec_gpu_task_t *current = first;
+
+    do {
+        current->priority = current->ec->priority;
+        current = (parsec_gpu_task_t *)current->list_item.list_next;
+    } while( current != first );
+    parsec_list_chain_sorted(gpu_stream->fifo_pending, followers,
+                             offsetof(parsec_gpu_task_t, priority));
+#else
+    parsec_list_chain_back(gpu_stream->fifo_pending, followers);
+#endif
+}
+```
+
+After this helper, `gpu_task` is the only member retained by an `AGAIN` return.
+The hook must submit progress for that head before returning `AGAIN`; the
+followers have already received a new owner through `fifo_pending`.
+
+Profiling semantics
+-------------------
+
+GPU execution profiling is owned by the common device runtime, not by the
+submit hook. After a hook first returns `PARSEC_HOOK_RETURN_DONE` or
+`PARSEC_HOOK_RETURN_AGAIN`, the runtime emits one task start for every member of
+the finalized ring. It emits one matching task end for every member when the
+submitted execution finally completes. A task therefore contributes exactly
+one start/end pair whether it was submitted alone or as part of a batch.
+
+An `AGAIN` event only resumes the same logical execution. The runtime keeps all
+member intervals open across any number of repeated `AGAIN` continuations and
+does not emit progress markers for those intermediate steps. A committed
+non-singleton ring has already been fully accounted, so repeated continuations
+do not rescan every member looking for new profiling starts.
+
+The open profiling state belongs to each device-task wrapper. If a hook
+manually disbands an `AGAIN` ring and returns followers to the GPU stream's
+pending FIFO, their intervals follow them and close when those wrappers
+eventually complete. A hook must not simply release an open detached wrapper,
+because that would leave its profiling interval without a matching end.
 
 Iterating over the returned ring
 --------------------------------

parsec/class/list_item.h

@@ -185,8 +185,11 @@ parsec_list_item_ring_merge( parsec_list_item_t* ring1,
  *   Removes an item from a ring of items.
  *
  * @details
- *   item must belong to a ring. It is singletoned, and the ring without
- *   item is returned.
+ *   item must belong to a ring. Its neighbors are reconnected and the ring
+ *   without item is returned, but item's own list_next and list_prev are not
+ *   made consistent with that ring. In release builds they retain their old,
+ *   now broken linkage; paranoid builds poison them instead. The caller must
+ *   reinitialize item before using its list links again.
  *  @param[inout] item the item from the ring of items to be removed.
  *  @return the rest of the ring
  * @remark

parsec/include/parsec/execution_stream.h

@@ -48,8 +48,12 @@ struct parsec_execution_stream_s {
 
     void *scheduler_object;
 
-    /* The task to be executed next by this execution_stream. Beware as this bypasses
-     * the scheduler decision.
+    /* One ready task reserved for this execution stream. It bypasses the
+     * scheduler and is invisible to other workers until consumed or flushed.
+     * Once code commits this stream to long-lived work, it must call
+     * __parsec_schedule_flush_private() before entering that work. A path that
+     * merely hands work to an existing manager and returns should retain the
+     * private task for this stream to execute next.
      */
     struct parsec_task_s* next_task;
 

parsec/interfaces/ptg/ptg-compiler/jdf2c.c

@@ -6748,22 +6748,6 @@ static void jdf_generate_code_hook_gpu(const jdf_t *jdf,
     jdf_generate_code_dry_run_before(jdf, f);
     jdf_coutput_prettycomment('-', "%s BODY", f->fname);
 
-    if( profile_on ) {
-        coutput("#if defined(PARSEC_PROF_TRACE)\n"
-                "  if(gpu_stream->prof_event_track_enable) {\n"
-                "    PARSEC_TASK_PROF_TRACE(gpu_stream->profiling,\n"
-                "                           PARSEC_PROF_FUNC_KEY_START(this_task->taskpool,\n"
-                "                                     this_task->task_class->task_class_id),\n"
-                "                           (parsec_task_t*)this_task, 1);\n"
-                "    gpu_task->prof_key_end = PARSEC_PROF_FUNC_KEY_END(this_task->taskpool,\n"
-                "                                   this_task->task_class->task_class_id);\n"
-                "    gpu_task->prof_event_id = this_task->task_class->key_functions->\n"
-                "           key_hash(this_task->task_class->make_key(this_task->taskpool, ((parsec_task_t*)this_task)->locals), NULL);\n"
-                "    gpu_task->prof_tp_id = this_task->taskpool->taskpool_id;\n"
-                "  }\n"
-                "#endif /* PARSEC_PROF_TRACE */\n");
-    }
-
     if ( NULL != dyld ) {
         coutput("  /* Pointer to dynamic gpu function */\n"
                 "  {\n"
@@ -6851,6 +6835,16 @@ static void jdf_generate_code_hook_gpu(const jdf_t *jdf,
             "  gpu_task->task_type = PARSEC_GPU_TASK_TYPE_KERNEL;\n",
             dev_lower, jdf_basename, f->fname);
 
+    /* Execution profiling is emitted by the common GPU runtime after the
+     * submit hook has finalized a possible batch. Preserve the PTG body-level
+     * profiling property on the wrapper so that common path can honor it.
+     */
+    coutput("#if defined(PARSEC_PROF_TRACE)\n"
+            "  gpu_task->prof_exec_state = %s;\n"
+            "#endif\n",
+            profile_on ? "PARSEC_GPU_TASK_PROF_EXEC_PENDING"
+                       : "PARSEC_GPU_TASK_PROF_EXEC_DISABLED");
+
     /* Set up stage in/out callbacks */
     jdf_find_property(body->properties, "stage_in", &stage_in_property);
     coutput("  gpu_task->stage_in  = %s;\n", (NULL == stage_in_property) ? "parsec_default_gpu_stage_in"