Refactor: hbg gives in-graph tasks their own compact state array

src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler.h

@@ -16,10 +16,10 @@
  * 1. Maintaining per-resource-shape ready queues
  * 2. Polling-completion dependency resolution: a GLOBAL task is ready when every
  *    producer named in its inline fanin has a COMPLETED task_states byte, an
- *    IN_GRAPH one when every producer in its Graph's fanin wire has reached
- *    task_state == COMPLETED (that table has no flag bytes); a producer publishes
- *    completion + drains its wake list on finish
- * 3. Publishing completion (task_state PENDING -> COMPLETED, task_states byte)
+ *    IN_GRAPH one when every producer in its Graph's fanin wire does — the same
+ *    array shape, held by the GraphExecution instead of the task header; a
+ *    producer publishes completion + drains its wake list on finish
+ * 3. Publishing completion (task_states byte PENDING -> COMPLETED)
  * 4. Two-stage mixed-task completion (subtask done bits -> mixed-task complete)
  *
  * The Scheduler runs on Device AI_CPU. host_build_graph is scheduler-only (the
@@ -659,15 +659,14 @@ struct SchedulerState {
         }
     }
 
-    // Producer completion under polling: publish the task_state completion
-    // mirror + the device-visible task_states byte, then drain the wake list
+    // Producer completion under polling: publish the device-visible task_states
+    // byte, then drain the wake list
     // (route/re-register each waiter). Whole-graph-resident hbg
     // has no device slot reclaim, so nothing advances a reclaim cursor here.
     void on_mixed_task_complete(ChipTaskSlotState &slot_state) {
         const int32_t task_id = slot_state.to_descriptor().task_id.local_id();
         SharedMemoryTaskHeader &tasks = *task_view.tasks;
 
-        slot_state.mark_completed();  // completion mirror (task_state = COMPLETED)
         tasks.store_completed(task_id);
         // COMPLETED >= PUBLISHED, so a tracked producer that never published
         // (DUMMY, predicate-retired) still releases its publish-list waiters
@@ -1156,8 +1155,7 @@ struct SchedulerState {
         const int32_t count = execution.fanin_offsets[task_index + 1] - begin;
         const int32_t start = consumer.wake_scan_cursor < count ? consumer.wake_scan_cursor : count - 1;
         for (int32_t row = start; row >= 0; --row) {
-            const ChipTaskSlotState &producer = execution.task_at(execution.fanin_indices[begin + row]).slot;
-            if (producer.task_state.load(std::memory_order_acquire) != CHIP_TASK_COMPLETED) {
+            if (!execution.is_completed(execution.fanin_indices[begin + row])) {
                 consumer.wake_scan_cursor = static_cast<uint16_t>(row);
                 return row;
             }
@@ -1355,7 +1353,7 @@ struct SchedulerState {
         // Publish completion before closing the wake list. A consumer that
         // loses registration to the sentinel acquires this state when it
         // rescans the Orch-built fanin wire, so no wakeup can be lost.
-        slot_state.mark_completed();
+        execution->store_completed(task_index);
         outcome.fanout_edges = drain_graph_wake_list(*execution, slot_state);
 
         const bool graph_completed = graph_execution_complete_in_graph_task(*execution);
@@ -1410,8 +1408,8 @@ struct SchedulerState {
         int thread_idx
 #endif
     ) {
-        // Polling completion: publish the task_state completion mirror + the
-        // device-visible task_states byte and drain the wake list (route or
+        // Polling completion: publish the device-visible task_states
+        // byte and drain the wake list (route or
         // re-register each waiter). Replaces the
         // fanout-list walk + fanin_refcount decrements of the wiring model.
         on_mixed_task_complete(slot_state);

src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp

@@ -219,7 +219,7 @@ void SchedulerContext::log_stall_diagnostics(
         // below it was claimed by the host orchestrator, and no slot above it was.
         for (int32_t si = 0; si < task_count; si++) {
             ChipTaskSlotState &slot_state = tasks.get_slot_state_by_task_id(si);
-            ChipTaskState st = slot_state.task_state.load(std::memory_order_relaxed);
+            const bool completed = tasks.is_completed(si, std::memory_order_relaxed);
             // Polling: no fanin_refcount. Recompute met/total from the inline
             // fanin ids vs the task_states array (rc = satisfied producers,
             // fi = raw producer count) so the stall dump still shows readiness.
@@ -235,10 +235,10 @@ void SchedulerContext::log_stall_diagnostics(
             int32_t kid_aiv0 = slot_state.to_descriptor().kernel_id[1];
             int32_t kid_aiv1 = slot_state.to_descriptor().kernel_id[2];
             int64_t task_id = static_cast<int64_t>(slot_state.to_descriptor().task_id.raw);
-            if (st >= CHIP_TASK_COMPLETED) continue;
-            // The slot mirror has no intermediate ready/running value — it
-            // stays PENDING until the worker stores COMPLETED (PUBLISHED
-            // lives in the task_states array, not here). Classify
+            if (completed) continue;
+            // The state byte has no intermediate ready/running value — a task
+            // stays PENDING until it publishes PUBLISHED or COMPLETED, neither
+            // of which distinguishes queued from running. Classify
             // by the ground truth instead: a slot is RUNNING iff some
             // core has it as running_slot_state. A task occupies at most
             // 3 cores (one cluster), all under the same owner thread by

src/a5/runtime/host_build_graph/host/runtime_maker.cpp

@@ -975,25 +975,29 @@ bool create_scheduler_state(
     uint64_t aic_worker_demand = 0;
     uint64_t aiv_worker_demand = 0;
     int64_t legacy_shape_task_id = -1;
+    // The progress byte is PENDING or COMPLETED throughout this walk, never
+    // PUBLISHED: that value is stored only by the device dispatch path, which
+    // has not run yet, and orch::prepare_task resets every claimed slot this
+    // bind. So `is_completed` and its negation partition the tasks here, and a
+    // future host-side publication would silently change what these tests mean.
     for (int64_t task_id = 0; task_id < total_tasks; ++task_id) {
         ChipTaskSlotState &slot = host_sm_handle.header->tasks.get_slot_state_by_task_id(task_id);
         SchedulerTaskShape shape{};
         SchedulerGraphResult status = scheduler_classify_task_shape(host_graph, task_id, &shape);
         bool inline_dispatch_task = false;
         if (status != SchedulerGraphResult::OK) {
-            bool inline_completed_task = status == SchedulerGraphResult::UNSUPPORTED_SHAPE &&
-                                         slot.active_mask.raw() == 0 && slot.logical_block_num == 1 &&
-                                         slot.total_required_subtasks == 0 &&
-                                         slot.task_state.load(std::memory_order_acquire) == CHIP_TASK_COMPLETED &&
-                                         slot.task_attrs.allow_early_resolve() &&
-                                         !slot.task_attrs.requires_sync_start() && !slot.task_attrs.has_predicate();
+            bool inline_completed_task =
+                status == SchedulerGraphResult::UNSUPPORTED_SHAPE && slot.active_mask.raw() == 0 &&
+                slot.logical_block_num == 1 && slot.total_required_subtasks == 0 &&
+                host_sm_handle.header->tasks.is_completed(task_id) && slot.task_attrs.allow_early_resolve() &&
+                !slot.task_attrs.requires_sync_start() && !slot.task_attrs.has_predicate();
             if (inline_completed_task) {
                 inline_completed_task_ids.push_back(task_id);
                 continue;
             }
             inline_dispatch_task = status == SchedulerGraphResult::UNSUPPORTED_SHAPE && slot.active_mask.raw() == 0 &&
                                    slot.logical_block_num == 1 && slot.total_required_subtasks == 0 &&
-                                   slot.task_state.load(std::memory_order_acquire) == CHIP_TASK_PENDING &&
+                                   !host_sm_handle.header->tasks.is_completed(task_id) &&
                                    !slot.task_attrs.requires_sync_start() && !slot.task_attrs.has_predicate();
             if (!inline_dispatch_task) {
                 LOG_ERROR(
@@ -1030,7 +1034,7 @@ bool create_scheduler_state(
         const uint32_t expected_subtasks = logical_block_num * active_subtasks;
         if ((!inline_dispatch_task && (slot.active_mask.raw() != classified_active_mask ||
                                        static_cast<uint32_t>(slot.total_required_subtasks) != expected_subtasks)) ||
-            expected_subtasks > UINT16_MAX || slot.task_state.load(std::memory_order_acquire) != CHIP_TASK_PENDING ||
+            expected_subtasks > UINT16_MAX || host_sm_handle.header->tasks.is_completed(task_id) ||
             (slot.task_attrs.has_predicate() && (active_subtasks != 1 || logical_block_num != 1))) {
             LOG_ERROR(
                 "A5 HBG AICore scheduler: task id=%" PRId64

src/a5/runtime/host_build_graph/runtime/scheduler/scheduler.h

@@ -16,10 +16,10 @@
  * 1. Maintaining per-resource-shape ready queues
  * 2. Polling-completion dependency resolution: a GLOBAL task is ready when every
  *    producer named in its inline fanin has a COMPLETED task_states byte, an
- *    IN_GRAPH one when every producer in its Graph's fanin wire has reached
- *    task_state == COMPLETED (that table has no flag bytes); a producer publishes
- *    completion + drains its wake list on finish
- * 3. Publishing completion (task_state PENDING -> COMPLETED, task_states byte)
+ *    IN_GRAPH one when every producer in its Graph's fanin wire does — the same
+ *    array shape, held by the GraphExecution instead of the task header; a
+ *    producer publishes completion + drains its wake list on finish
+ * 3. Publishing completion (task_states byte PENDING -> COMPLETED)
  * 4. Two-stage mixed-task completion (subtask done bits -> mixed-task complete)
  *
  * The Scheduler runs on Device AI_CPU. host_build_graph is scheduler-only (the
@@ -670,15 +670,14 @@ struct SchedulerState {
         }
     }
 
-    // Producer completion under polling: publish the task_state completion
-    // mirror + the device-visible task_states byte, then drain the wake list
+    // Producer completion under polling: publish the device-visible task_states
+    // byte, then drain the wake list
     // (route/re-register each waiter). Whole-graph-resident hbg
     // has no device slot reclaim, so nothing advances a reclaim cursor here.
     void on_mixed_task_complete(ChipTaskSlotState &slot_state) {
         const int32_t task_id = slot_state.to_descriptor().task_id.local_id();
         SharedMemoryTaskHeader &tasks = *task_view.tasks;
 
-        slot_state.mark_completed();  // completion mirror (task_state = COMPLETED)
         tasks.store_completed(task_id);
         // COMPLETED >= PUBLISHED, so a tracked producer that never published
         // (DUMMY, predicate-retired) still releases its publish-list waiters
@@ -1167,8 +1166,7 @@ struct SchedulerState {
         const int32_t count = execution.fanin_offsets[task_index + 1] - begin;
         const int32_t start = consumer.wake_scan_cursor < count ? consumer.wake_scan_cursor : count - 1;
         for (int32_t row = start; row >= 0; --row) {
-            const ChipTaskSlotState &producer = execution.task_at(execution.fanin_indices[begin + row]).slot;
-            if (producer.task_state.load(std::memory_order_acquire) != CHIP_TASK_COMPLETED) {
+            if (!execution.is_completed(execution.fanin_indices[begin + row])) {
                 consumer.wake_scan_cursor = static_cast<uint16_t>(row);
                 return row;
             }
@@ -1366,7 +1364,7 @@ struct SchedulerState {
         // Publish completion before closing the wake list. A consumer that
         // loses registration to the sentinel acquires this state when it
         // rescans the Orch-built fanin wire, so no wakeup can be lost.
-        slot_state.mark_completed();
+        execution->store_completed(task_index);
         outcome.fanout_edges = drain_graph_wake_list(*execution, slot_state);
 
         const bool graph_completed = graph_execution_complete_in_graph_task(*execution);
@@ -1421,8 +1419,8 @@ struct SchedulerState {
         int thread_idx
 #endif
     ) {
-        // Polling completion: publish the task_state completion mirror + the
-        // device-visible task_states byte and drain the wake list (route or
+        // Polling completion: publish the device-visible task_states
+        // byte and drain the wake list (route or
         // re-register each waiter). Replaces the
         // fanout-list walk + fanin_refcount decrements of the wiring model.
         on_mixed_task_complete(slot_state);

src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp

@@ -219,7 +219,7 @@ void SchedulerContext::log_stall_diagnostics(
         // below it was claimed by the host orchestrator, and no slot above it was.
         for (int32_t si = 0; si < task_count; si++) {
             ChipTaskSlotState &slot_state = tasks.get_slot_state_by_task_id(si);
-            ChipTaskState st = slot_state.task_state.load(std::memory_order_relaxed);
+            const bool completed = tasks.is_completed(si, std::memory_order_relaxed);
             // Polling: no fanin_refcount. Recompute met/total from the inline
             // fanin ids vs the task_states array (rc = satisfied producers,
             // fi = raw producer count) so the stall dump still shows readiness.
@@ -235,10 +235,10 @@ void SchedulerContext::log_stall_diagnostics(
             int32_t kid_aiv0 = slot_state.to_descriptor().kernel_id[1];
             int32_t kid_aiv1 = slot_state.to_descriptor().kernel_id[2];
             int64_t task_id = static_cast<int64_t>(slot_state.to_descriptor().task_id.raw);
-            if (st >= CHIP_TASK_COMPLETED) continue;
-            // The slot mirror has no intermediate ready/running value — it
-            // stays PENDING until the worker stores COMPLETED (PUBLISHED
-            // lives in the task_states array, not here). Classify
+            if (completed) continue;
+            // The state byte has no intermediate ready/running value — a task
+            // stays PENDING until it publishes PUBLISHED or COMPLETED, neither
+            // of which distinguishes queued from running. Classify
             // by the ground truth instead: a slot is RUNNING iff some
             // core has it as running_slot_state. A task occupies at most
             // 3 cores (one cluster), all under the same owner thread by

src/common/host_build_graph/device/graph_execution.cpp

@@ -37,6 +37,7 @@ GraphExecution *acquire_execution_storage(
     execution->task_storage = reinterpret_cast<ChipTaskStorage *>(base + layout.tasks_offset);
     execution->task_tensor_pool = reinterpret_cast<simpler::hbg::Tensor *>(base + layout.tensors_offset);
     execution->task_scalar_pool = reinterpret_cast<uint64_t *>(base + layout.scalars_offset);
+    execution->task_states = reinterpret_cast<std::atomic<ChipTaskState> *>(base + layout.states_offset);
     return execution;
 }
 
@@ -412,7 +413,10 @@ GraphMaterializeResult graph_execution_materialize_slice(
         task.packed_buffer_end = reinterpret_cast<void *>(outer_base + task_offset + output_bytes);
 
         slot.reset_for_reuse();
-        slot.task_state.store(CHIP_TASK_PENDING, std::memory_order_relaxed);
+        // The readiness a consumer polls, cleared before this task becomes
+        // visible to one: materialization publishes tasks incrementally, so a
+        // peer may scan this index as soon as published_tasks passes it.
+        execution.reset_task_state(i);
         slot.active_mask = ActiveMask(source.active_mask);
         slot.task_attrs = TaskAttrs(source.task_attrs);
         slot.total_required_subtasks = source.total_required_subtasks;

src/common/host_build_graph/docs/GRAPH_EXECUTION.md

@@ -380,8 +380,10 @@ directly, so boundary metadata does not need to expand to full `ChipTensor`
 records on either side of H2D.
 
 In-graph tasks consume no task-table slots. Their descriptor, payload, slot
-state, and argument pools live in the tail of the outer `GRAPH` task's own heap block,
-past `required_heap`: `[GraphExecution][ChipTaskStorage...][tensor pool][scalar pool]`. One
+state, argument pools, and completion states live in the tail of the outer `GRAPH` task's
+own heap block, past `required_heap`:
+`[GraphExecution][ChipTaskStorage...][tensor pool][scalar pool][task_states]`. The state
+array is last because a byte needs no alignment, so appending it moves no other region. One
 `TaskAllocator::alloc` covers both the packed outputs and this execution storage,
 so they are reclaimed together without a separate device allocation or release path.
 
@@ -448,8 +450,9 @@ dependency wiring remains an Orchestrator responsibility:
   operand order, so the tail is exactly the deepest producer only on an
   early-dispatch candidate's row, which recording sorts by producer index;
   elsewhere the direction is a heuristic;
-- an in-graph task's release/acquire `task_state` is its Graph-local completion truth, so
-  such tasks need neither a shared-memory `task_states` byte nor a task-table slot;
+- an in-graph task's completion truth is its execution's own `task_states` byte,
+  the same shape the task header gives a GLOBAL task, so such tasks need neither
+  a task-table slot nor a byte in the shared-memory array;
 - producer completion closes and drains only its current wake-list rather than
   traversing the saved fanout CSR;
 - a woken consumer with a single producer enters its shape queue directly; any