Refactor: consolidate per-task markers into one TaskAttrs byte (a2a3/hbg)

src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp

@@ -385,7 +385,7 @@ static bool check_scope_can_accept_task(PTO2OrchestratorState *orch, PTO2TaskAll
 
 static bool prepare_task(
     PTO2OrchestratorState *orch, const L0TaskArgs &args, int32_t total_output_size, ActiveMask active_mask,
-    PTO2PreparedTask *out
+    TaskAttrs task_attrs, PTO2PreparedTask *out
 ) {
     uint8_t ring_id = 0;
     auto &allocator = orch->ring.task_allocator;
@@ -435,6 +435,7 @@ static bool prepare_task(
         static_cast<int16_t>(block_num * __builtin_popcount(active_mask.core_mask()));
     out->slot_state->logical_block_num = block_num;
     out->slot_state->active_mask = active_mask;
+    out->slot_state->task_attrs = task_attrs;
     // Reclaim gate: seed last_consumer to self, so a producer with no consumers
     // is retirable once completed_watermark >= its own id. Each fanin edge bumps
     // it in append_fanin_or_fail. completion_flags for this slot are already 0
@@ -655,20 +656,19 @@ static bool ensure_tensormap_capacity(PTO2OrchestratorState *orch, int32_t neede
 // computation (explicit_deps + auto), output registration, slot init, and
 // Orch-side wiring/ready publication.
 static TaskOutputTensors submit_task_common(
-    PTO2OrchestratorState *orch, const L0TaskArgs &args, ActiveMask active_mask, int32_t aic_kernel_id,
-    int32_t aiv0_kernel_id, int32_t aiv1_kernel_id
+    PTO2OrchestratorState *orch, const L0TaskArgs &args, ActiveMask active_mask, TaskAttrs task_attrs,
+    int32_t aic_kernel_id, int32_t aiv0_kernel_id, int32_t aiv1_kernel_id
 ) {
     CYCLE_COUNT_START();
     TaskOutputTensors result;
     PTO2OutputLayout layout = calculate_output_layout(args);
     PTO2PreparedTask prepared;
-    if (!prepare_task(orch, args, layout.total_output_size, active_mask, &prepared)) {
+    if (!prepare_task(orch, args, layout.total_output_size, active_mask, task_attrs, &prepared)) {
         return result;
     }
     PTO2SchedulerState *sched = orch->scheduler;
     PTO2RingFlowControl &fc = orch->sm_header->ring.fc;
     PTO2TaskId task_id = prepared.task_id;
-    PTO2TaskSlotState &cur_slot_state = *prepared.slot_state;
     PTO2TaskDescriptor &task = *prepared.task;
     PTO2TaskPayload &payload = *prepared.payload;
     result.set_task_id(task_id);
@@ -789,15 +789,13 @@ static TaskOutputTensors submit_task_common(
     task.kernel_id[static_cast<int>(PTO2SubtaskSlot::AIC)] = aic_kernel_id;
     task.kernel_id[static_cast<int>(PTO2SubtaskSlot::AIV0)] = aiv0_kernel_id;
     task.kernel_id[static_cast<int>(PTO2SubtaskSlot::AIV1)] = aiv1_kernel_id;
-    task.task_timing_slot = args.task_timing_slot();
     task.packed_buffer_base = prepared.alloc_result.packed_base;
     task.packed_buffer_end = prepared.alloc_result.packed_end;
 
     // append_fanin_or_fail wrote each producer's local id straight into
     // payload.fanin_local_ids and bumped its last_consumer_local_id; the count is
     // published in STEP 6 below. payload.init does not touch the fanin region.
     payload.init(args, result, prepared.alloc_result, layout);
-    cur_slot_state.set_allow_early_resolve(args.allow_early_resolve());
 
     // Dispatch predicate: resolve the (tensor, indices) to an absolute GM address
     // now so the scheduler can read it at the dispatch point with a single load,
@@ -905,7 +903,11 @@ TaskOutputTensors PTO2OrchestratorState::submit_task(const MixedKernels &mixed_k
         active_mask = normalized.to_active_mask();
     }
 
-    // Encode require_sync_start into active_mask bit 3 (only meaningful for tasks with block_num > 1)
+    TaskAttrs task_attrs;
+    task_attrs.set_early_resolve(args.allow_early_resolve());
+    task_attrs.set_timing_slot(args.task_timing_slot());
+
+    // sync_start is only meaningful for tasks with block_num > 1.
     if (block_num > 1 && args.launch_spec.require_sync_start()) {
         // Deadlock check: block_num >= total available slots of the required type.
         // For MIX/AIC: limit is total_cluster_count (one AIC per cluster).
@@ -919,15 +921,16 @@ TaskOutputTensors PTO2OrchestratorState::submit_task(const MixedKernels &mixed_k
             );
             return TaskOutputTensors{};
         }
-        active_mask.set_sync_start();
+        task_attrs.set_sync_start();
     }
 
     if (args.predicate().op != PredicateOp::NONE) {
-        active_mask.set_has_predicate();
+        task_attrs.set_predicate();
     }
 
     return submit_task_common(
-        orch, args, active_mask, normalized.aic_kernel_id, normalized.aiv0_kernel_id, normalized.aiv1_kernel_id
+        orch, args, active_mask, task_attrs, normalized.aic_kernel_id, normalized.aiv0_kernel_id,
+        normalized.aiv1_kernel_id
     );
 }
 
@@ -955,7 +958,15 @@ TaskOutputTensors PTO2OrchestratorState::submit_dummy_task(const L0TaskArgs &arg
     }
     always_assert(orch->scheduler != nullptr);
 
-    return submit_task_common(orch, args, ActiveMask{}, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID);
+    // Dummy tasks never dispatch to an AICore, so sync_start / has_predicate do
+    // not apply; only the early-dispatch hint and timing tag carry over.
+    TaskAttrs task_attrs;
+    task_attrs.set_early_resolve(args.allow_early_resolve());
+    task_attrs.set_timing_slot(args.task_timing_slot());
+
+    return submit_task_common(
+        orch, args, ActiveMask{}, task_attrs, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID
+    );
 }
 
 TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) {
@@ -995,7 +1006,9 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) {
 
     PTO2OutputLayout layout = calculate_output_layout(args);
     PTO2PreparedTask prepared;
-    if (!prepare_task(orch, args, layout.total_output_size, ActiveMask{}, &prepared)) {
+    // Kernel-less alloc task: no active subtasks, no dispatch-time attributes. The
+    // early-dispatch hint is force-set below (see the flag-the-creator note).
+    if (!prepare_task(orch, args, layout.total_output_size, ActiveMask{}, TaskAttrs{}, &prepared)) {
         return TaskOutputTensors{};
     }
 
@@ -1015,9 +1028,6 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) {
     task.kernel_id[static_cast<int>(PTO2SubtaskSlot::AIC)] = INVALID_KERNEL_ID;
     task.kernel_id[static_cast<int>(PTO2SubtaskSlot::AIV0)] = INVALID_KERNEL_ID;
     task.kernel_id[static_cast<int>(PTO2SubtaskSlot::AIV1)] = INVALID_KERNEL_ID;
-    // alloc_tensors builds a kernel-less descriptor that never dispatches; keep
-    // the slot untagged so a recycled ring slot cannot leak a stale tag.
-    task.task_timing_slot = TASK_TIMING_SLOT_NONE;
     task.packed_buffer_base = prepared.alloc_result.packed_base;
     task.packed_buffer_end = prepared.alloc_result.packed_end;
 
@@ -1044,7 +1054,7 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) {
         // creation — it should always be transparent, never a barrier. Unlike a
         // codegen task there is no Arg-driven hint to honor here, so mark it
         // unconditionally.
-        prepared.slot_state->allow_early_resolve = true;
+        prepared.slot_state->task_attrs.set_early_resolve(true);
         prepared.slot_state->mark_completed();  // host-visible task_state mirror
         // Polling: pre-set the device-visible completion_flags byte in the H2D
         // image. Consumers poll completion_flags (not task_state), so a hidden-alloc

src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h

@@ -197,26 +197,14 @@ struct PTO2TaskDescriptor {
     // Per-slot kernel IDs (INVALID_KERNEL_ID = inactive)
     int32_t kernel_id[PTO2_SUBTASK_SLOT_COUNT];
 
-    // Selective task-timing slot (TASK_TIMING_SLOT_NONE = untagged; 0..15 valid).
-    // Occupies the 4-byte alignment pad that already followed kernel_id[3], so
-    // the descriptor does not grow; the scheduler folds this task's dispatch/
-    // finish cycles into the tagged slot (see common/device_phase.h).
-    int32_t task_timing_slot;
-
     // Packed output buffer (all outputs packed into single contiguous buffer)
     void *packed_buffer_base;  // Start of packed buffer in GM Heap
     void *packed_buffer_end;   // End of packed buffer (for heap reclamation)
 };
 
-// task_timing_slot must occupy the former padding after kernel_id[3] without
-// growing the descriptor or shifting packed_buffer_base — the scheduler and
-// shared-memory ABI depend on the existing size and pointer offset.
-static_assert(sizeof(PTO2TaskDescriptor) == 40, "PTO2TaskDescriptor must not grow: slot uses the existing pad");
-static_assert(
-    offsetof(PTO2TaskDescriptor, task_timing_slot) ==
-        offsetof(PTO2TaskDescriptor, kernel_id) + sizeof(int32_t) * PTO2_SUBTASK_SLOT_COUNT,
-    "task_timing_slot must sit immediately after kernel_id in the former pad"
-);
+// A 4-byte alignment pad follows kernel_id[3]; the scheduler and shared-memory
+// ABI depend on the descriptor size and packed_buffer_base offset staying fixed.
+static_assert(sizeof(PTO2TaskDescriptor) == 40, "PTO2TaskDescriptor size is part of the shared-memory ABI");
 static_assert(offsetof(PTO2TaskDescriptor, packed_buffer_base) == 24, "packed_buffer_base offset must be unchanged");
 
 // =============================================================================
@@ -480,9 +468,12 @@ struct alignas(64) PTO2TaskSlotState {
 
     // --- Set per-submit (depend on task inputs) ---
     ActiveMask active_mask;  // Bitmask of active subtask slots (set once)
-    // Codegen early-dispatch hint, copied from Arg at submit. Stored for the
-    // Milestone-2 publish-status early-dispatch; no Milestone-1 hot-path reader.
-    bool allow_early_resolve{false};
+    // Single per-task attributes byte (early-dispatch hint, sync_start,
+    // has_predicate, selective timing tag). Lives on slot_state (not payload) so
+    // fanin walks and the completion path read them off the already-hot producer
+    // slot_state cache line. Plain-write (set once at submit, before the slot is
+    // scheduler-visible).
+    TaskAttrs task_attrs{};
     // Set by any subtask FIN that pushed a deferred-completion CONDITION to the
     // runtime mailbox; read by the last subtask FIN to decide inline vs
     // MPSC-deferred completion. The release write is sequenced before
@@ -530,8 +521,6 @@ struct alignas(64) PTO2TaskSlotState {
 
     bool has_any_subtask_deferred() const { return any_subtask_deferred.load(std::memory_order_acquire); }
 
-    void set_allow_early_resolve(bool v) { allow_early_resolve = v; }
-
     /**
      * Reset dynamic scheduling fields to their pristine values. Runs once per
      * slot at init (pto_shared_memory.cpp) — whole-graph-resident hbg has no
@@ -545,7 +534,8 @@ struct alignas(64) PTO2TaskSlotState {
         any_subtask_deferred.store(false, std::memory_order_relaxed);
         completed_subtasks.store(0, std::memory_order_relaxed);
         next_block_idx.store(0, std::memory_order_relaxed);
-        allow_early_resolve = false;
+        // Note: active_mask and task_attrs are per-submit-constant fields
+        // rewritten in prepare_task on every reuse, so they are not reset here.
         // last_consumer_local_id is seeded in prepare_task once the id is known.
         // Payload early-dispatch/fanin fields are (re)initialized in
         // PTO2TaskPayload::init on every submit, before the slot is visible.

src/a2a3/runtime/host_build_graph/runtime/pto_submit_types.h

@@ -39,11 +39,9 @@ enum class PTO2SubtaskSlot : uint8_t {
 /**
  * Subtask mask bits (for ActiveMask)
  */
-inline constexpr uint8_t PTO2_SUBTASK_MASK_AIC = (1u << 0);            // 0x1
-inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV0 = (1u << 1);           // 0x2
-inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV1 = (1u << 2);           // 0x4
-inline constexpr uint8_t PTO2_SUBTASK_FLAG_SYNC_START = (1u << 3);     // 0x8: all blocks must launch atomically
-inline constexpr uint8_t PTO2_SUBTASK_FLAG_HAS_PREDICATE = (1u << 4);  // 0x10: task carries a dispatch predicate
+inline constexpr uint8_t PTO2_SUBTASK_MASK_AIC = (1u << 0);   // 0x1
+inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV0 = (1u << 1);  // 0x2
+inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV1 = (1u << 2);  // 0x4
 
 // Dispatch-predicate comparison operator. The scheduler evaluates the predicate
 // at the dispatch point — the task is ready (fanin satisfied), so the predicate
@@ -120,7 +118,11 @@ enum class PTO2ResourceShape : uint8_t {
 inline constexpr int32_t PTO2_NUM_RESOURCE_SHAPES = 3;
 
 /**
- * Bitmask of active subtask slots + flags, sizeof == 1.
+ * Bitmask of active subtask slots (AIC/AIV0/AIV1), sizeof == 1.
+ *
+ * Pure subtask-slot mask: only the low 3 bits (core_mask) are meaningful, so it
+ * can be OR/==-combined when building MIX clusters without dragging unrelated
+ * flag bits along. Per-task scheduling flags live on TaskAttrs instead.
  */
 class ActiveMask {
 public:
@@ -134,11 +136,6 @@ class ActiveMask {
 
     uint8_t core_mask() const { return raw_ & 0x07u; }
 
-    bool requires_sync_start() const { return (raw_ & PTO2_SUBTASK_FLAG_SYNC_START) != 0; }
-
-    // True when the task carries a dispatch predicate.
-    bool has_predicate() const { return (raw_ & PTO2_SUBTASK_FLAG_HAS_PREDICATE) != 0; }
-
     PTO2ResourceShape to_shape() const {
         uint8_t cmask = core_mask();
         if (cmask == 0) return PTO2ResourceShape::DUMMY;
@@ -148,10 +145,6 @@ class ActiveMask {
         return PTO2ResourceShape::AIV;
     }
 
-    void set_sync_start() { raw_ |= PTO2_SUBTASK_FLAG_SYNC_START; }
-
-    void set_has_predicate() { raw_ |= PTO2_SUBTASK_FLAG_HAS_PREDICATE; }
-
     bool operator==(ActiveMask other) const { return raw_ == other.raw_; }
     bool operator!=(ActiveMask other) const { return raw_ != other.raw_; }
 
@@ -173,6 +166,67 @@ class ActiveMask {
 
 static_assert(sizeof(ActiveMask) == 1, "ActiveMask must be exactly 1 byte");
 
+/**
+ * Per-task scheduling attributes, packed into one byte on PTO2TaskSlotState.
+ *
+ * Single home for the independent per-task flags: an early-dispatch hint, the
+ * two dispatch-time predicates (sync_start / has_predicate), and the selective
+ * timing tag. Consolidating them here keeps active_mask a pure subtask-slot mask
+ * and lands the timing tag on the scheduler's hot slot_state cache line.
+ *
+ *   bit 0     allow_early_resolve
+ *   bit 1     sync_start
+ *   bit 2     has_predicate
+ *   bit 3     is_timed          (0 => untagged; timing_tag ignored)
+ *   bits 4-7  timing_tag (0..15)
+ */
+class TaskAttrs {
+public:
+    constexpr TaskAttrs() = default;
+
+    bool allow_early_resolve() const { return (raw_ & BIT_EARLY_RESOLVE) != 0; }
+    void set_early_resolve(bool v) {
+        if (v) {
+            raw_ |= BIT_EARLY_RESOLVE;
+        } else {
+            raw_ &= static_cast<uint8_t>(~BIT_EARLY_RESOLVE);
+        }
+    }
+
+    bool requires_sync_start() const { return (raw_ & BIT_SYNC_START) != 0; }
+    void set_sync_start() { raw_ |= BIT_SYNC_START; }
+
+    bool has_predicate() const { return (raw_ & BIT_HAS_PREDICATE) != 0; }
+    void set_predicate() { raw_ |= BIT_HAS_PREDICATE; }
+
+    bool is_timed() const { return (raw_ & BIT_IS_TIMED) != 0; }
+
+    // 0..15 timing tag when tagged, else -1 (matches TASK_TIMING_SLOT_NONE; the
+    // equality is guarded by a static_assert in pto_types.h).
+    int32_t timing_slot() const { return is_timed() ? static_cast<int32_t>(raw_ >> TIMING_TAG_SHIFT) : -1; }
+
+    // slot < 0 => untagged; 0..15 => tagged. Arg::set_task_timing_slot already
+    // range-checks 0..NUM_TASK_TIMING_SLOTS-1, so only the low 4 bits are stored.
+    void set_timing_slot(int32_t slot) {
+        raw_ &= static_cast<uint8_t>(~(BIT_IS_TIMED | TIMING_TAG_MASK));
+        if (slot >= 0) {
+            raw_ |= static_cast<uint8_t>(BIT_IS_TIMED | ((slot & 0x0F) << TIMING_TAG_SHIFT));
+        }
+    }
+
+private:
+    static constexpr uint8_t BIT_EARLY_RESOLVE = 1u << 0;
+    static constexpr uint8_t BIT_SYNC_START = 1u << 1;
+    static constexpr uint8_t BIT_HAS_PREDICATE = 1u << 2;
+    static constexpr uint8_t BIT_IS_TIMED = 1u << 3;
+    static constexpr uint8_t TIMING_TAG_SHIFT = 4;
+    static constexpr uint8_t TIMING_TAG_MASK = 0xF0u;
+
+    uint8_t raw_{0};
+};
+
+static_assert(sizeof(TaskAttrs) == 1, "TaskAttrs must be exactly 1 byte");
+
 /**
  * Mixed-task submit contract.
  *

src/a2a3/runtime/host_build_graph/runtime/pto_types.h

@@ -45,6 +45,11 @@
 #include "tensor.h"
 #include "tensor_create_info.h"  // runtime-only TensorCreateInfo + materialization helpers
 
+// TaskAttrs packs the timing tag into a 4-bit field and reports "untagged" as
+// -1, so the tag domain must fit 0..15 and the untagged sentinel must be -1.
+static_assert(NUM_TASK_TIMING_SLOTS <= 16, "timing tag must fit TaskAttrs' 4-bit field");
+static_assert(TASK_TIMING_SLOT_NONE == -1, "TaskAttrs::timing_slot() reports untagged as -1");
+
 typedef enum {
     ASYNC_ENGINE_SDMA = 0,
     ASYNC_ENGINE_ROCE = 1,

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

@@ -482,9 +482,9 @@ struct PTO2SchedulerState {
     void push_ready_routed(PTO2TaskSlotState *slot_state) {
         PTO2ResourceShape shape = slot_state->active_mask.to_shape();
         if (shape == PTO2ResourceShape::DUMMY ||
-            (slot_state->active_mask.has_predicate() && !slot_state->payload->predicate.pass())) {
+            (slot_state->task_attrs.has_predicate() && !slot_state->payload->predicate.pass())) {
             dummy_ready_queue.push(slot_state);
-        } else if (slot_state->active_mask.requires_sync_start()) {
+        } else if (slot_state->task_attrs.requires_sync_start()) {
             ready_sync_queues[static_cast<int32_t>(shape)].push(slot_state);
         } else {
             ready_queues[static_cast<int32_t>(shape)].push(slot_state);
@@ -703,7 +703,7 @@ struct PTO2SchedulerState {
     }
 
     inline void record_published_blocks(PTO2TaskSlotState &slot_state, int32_t count) {
-        if (count <= 0 || !slot_state.allow_early_resolve) return;
+        if (count <= 0 || !slot_state.task_attrs.allow_early_resolve()) return;
         slot_state.payload->published_block_count.fetch_add(static_cast<int16_t>(count), std::memory_order_seq_cst);
     }