From b0f9515c16b643ec0f2f79571d5dbc7558a6dd4f Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Mon, 20 Jul 2026 06:36:54 -0700 Subject: [PATCH] Refactor: name task lifecycle flags accurately Rename the shared slot bitmask and consolidate its flag enums. Keep both A2A3 runtime variants aligned without changing layout, ordering, or atomic behavior. --- .../runtime/pto_runtime2_types.h | 30 ++++++-------- .../runtime/scheduler/pto_scheduler.h | 28 ++++++------- .../docs/RUNTIME_LOGIC.md | 2 +- .../runtime/pto_runtime2_types.h | 39 +++++++------------ .../runtime/scheduler/pto_scheduler.h | 28 ++++++------- tests/ut/cpp/a2a3/test_wiring.cpp | 6 +-- 6 files changed, 59 insertions(+), 74 deletions(-) diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h index d44bbc26fd..6db6a02120 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h @@ -441,7 +441,7 @@ static_assert( * Per-task slot scheduling state (scheduler-private, NOT in shared memory) * * Consolidates all hot-path scheduling fields into a single cache-friendly - * structure (32 bytes = half a cache line). Accessing any field of a task's + * structure (64 bytes = one cache line). Accessing any field of a task's * slot state brings all related fields into the same cache line. * * Concurrency notes: @@ -466,17 +466,11 @@ static_assert( // never reach -> provable deadlock. static constexpr uint32_t PTO2_FANOUT_SCOPE_BIT = 0x80000000u; -enum PTO2ReadyState : uint8_t { - PTO2_READY_UNCLAIMED = 0, - PTO2_READY_CLAIMED = 1, -}; - -enum PTO2CompletionFlag : uint8_t { - PTO2_COMPLETION_DONE = 2, -}; - -enum PTO2DeferredCompletionFlag : uint8_t { - PTO2_SUBTASK_DEFERRED = 4, +enum PTO2TaskLifecycleFlag : uint8_t { + PTO2_LIFECYCLE_FLAGS_NONE = 0, + PTO2_READY_CLAIMED = 1U << 0, + PTO2_COMPLETION_DONE = 1U << 1, + PTO2_SUBTASK_DEFERRED = 1U << 2, }; struct alignas(64) PTO2TaskSlotState { @@ -512,7 +506,7 @@ struct alignas(64) PTO2TaskSlotState { // slot_state (not payload) so fanin walks read the already-hot producer // slot_state cache line. bool allow_early_resolve{false}; - std::atomic ready_state{PTO2_READY_UNCLAIMED}; + std::atomic lifecycle_flags{PTO2_LIFECYCLE_FLAGS_NONE}; int32_t dep_pool_mark{0}; // Dep pool top after Orch-side wiring std::atomic completed_subtasks{0}; // Each core completion increments by 1 @@ -554,11 +548,11 @@ struct alignas(64) PTO2TaskSlotState { void mark_completed() { task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); - ready_state.fetch_or(PTO2_COMPLETION_DONE, std::memory_order_release); + lifecycle_flags.fetch_or(PTO2_COMPLETION_DONE, std::memory_order_release); } bool is_completion_flag_set() const { - return (ready_state.load(std::memory_order_acquire) & PTO2_COMPLETION_DONE) != 0; + return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_COMPLETION_DONE) != 0; } // Set by any subtask FIN that pushed deferred-completion CONDITIONs to the @@ -567,10 +561,10 @@ struct alignas(64) PTO2TaskSlotState { // release write is sequenced before on_subtask_complete's acq_rel fetch_add // and the acquire read after, so all earlier subtasks' writes are visible to // the last subtask. - void mark_any_subtask_deferred() { ready_state.fetch_or(PTO2_SUBTASK_DEFERRED, std::memory_order_release); } + void mark_any_subtask_deferred() { lifecycle_flags.fetch_or(PTO2_SUBTASK_DEFERRED, std::memory_order_release); } bool has_any_subtask_deferred() const { - return (ready_state.load(std::memory_order_acquire) & PTO2_SUBTASK_DEFERRED) != 0; + return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_SUBTASK_DEFERRED) != 0; } void set_allow_early_resolve(bool v) { allow_early_resolve = v; } @@ -594,7 +588,7 @@ struct alignas(64) PTO2TaskSlotState { fanout_refcount.store(0, std::memory_order_relaxed); completed_subtasks.store(0, std::memory_order_relaxed); next_block_idx.store(0, std::memory_order_relaxed); - ready_state.store(PTO2_READY_UNCLAIMED, std::memory_order_relaxed); + lifecycle_flags.store(PTO2_LIFECYCLE_FLAGS_NONE, std::memory_order_relaxed); allow_early_resolve = false; // Note: payload early-dispatch fields (state, masks, fanin, publication count) // are NOT reset here — this method skips the payload by contract. They are diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h index 944a877a8c..feaa87dd3d 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h @@ -512,12 +512,12 @@ struct PTO2SchedulerState { } bool try_claim_ready_once(PTO2TaskSlotState &slot_state) { - uint8_t state = slot_state.ready_state.load(std::memory_order_acquire); + uint8_t flags = slot_state.lifecycle_flags.load(std::memory_order_acquire); for (;;) { - if ((state & PTO2_READY_CLAIMED) != 0) return false; - uint8_t desired = state | PTO2_READY_CLAIMED; - if (slot_state.ready_state.compare_exchange_weak( - state, desired, std::memory_order_acq_rel, std::memory_order_acquire + if ((flags & PTO2_READY_CLAIMED) != 0) return false; + uint8_t desired_flags = flags | PTO2_READY_CLAIMED; + if (slot_state.lifecycle_flags.compare_exchange_weak( + flags, desired_flags, std::memory_order_acq_rel, std::memory_order_acquire )) { return true; } @@ -927,18 +927,18 @@ struct PTO2SchedulerState { PTO2TaskSlotState &slot_state, uint64_t &atomic_count, uint64_t &push_wait, EarlyDispatchReleaseSink *sink = nullptr ) { - uint8_t state = slot_state.ready_state.load(std::memory_order_acquire); - atomic_count += 1; // ready_state load + uint8_t flags = slot_state.lifecycle_flags.load(std::memory_order_acquire); + atomic_count += 1; // lifecycle_flags load for (;;) { - if ((state & PTO2_READY_CLAIMED) != 0) return false; - uint8_t desired = state | PTO2_READY_CLAIMED; - if (slot_state.ready_state.compare_exchange_weak( - state, desired, std::memory_order_acq_rel, std::memory_order_acquire + if ((flags & PTO2_READY_CLAIMED) != 0) return false; + uint8_t desired_flags = flags | PTO2_READY_CLAIMED; + if (slot_state.lifecycle_flags.compare_exchange_weak( + flags, desired_flags, std::memory_order_acq_rel, std::memory_order_acquire )) { - atomic_count += 1; // ready_state CAS + atomic_count += 1; // lifecycle_flags CAS break; } - atomic_count += 1; // failed ready_state CAS + atomic_count += 1; // failed lifecycle_flags CAS } // Early-dispatch: pre-staged tasks are released by doorbell @@ -1080,7 +1080,7 @@ struct PTO2SchedulerState { slot_state.unlock_fanout(); #if SIMPLER_SCHED_PROFILING - lock_atomics += 3; // task_state.store + ready_state.fetch_or + unlock.store + lock_atomics += 3; // task_state.store + lifecycle_flags.fetch_or + unlock.store g_sched_lock_atomic_count[thread_idx] += lock_atomics; g_sched_lock_wait_cycle[thread_idx] += lock_wait; PTO2_SCHED_CYCLE_LAP(g_sched_lock_cycle[thread_idx]); diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index f81f993910..371d75ab45 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -683,7 +683,7 @@ reserved core slot and a launch-visible payload before a consumer can pre-occupy `next_block_idx` records reservation progress; `published_block_count` independently establishes publication and early-candidate readiness. -The producer's slot-local dispatch-propagated bit in `ready_state` and fanout snapshot are +The producer's slot-local dispatch-propagated bit in `lifecycle_flags` and fanout snapshot are serialized under `fanout_lock` with consumer wiring. Wiring already locks and reads the producer's 64-byte `PTO2TaskSlotState`, so testing this bit does not read a producer payload cache line. An edge already in the snapshot is counted by scheduler propagation; wiring seeds diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h index 1a947a9666..ba004cf564 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h @@ -465,21 +465,12 @@ static_assert( // never reach -> provable deadlock. static constexpr uint32_t PTO2_FANOUT_SCOPE_BIT = 0x80000000u; -enum PTO2ReadyState : uint8_t { - PTO2_READY_UNCLAIMED = 0, - PTO2_READY_CLAIMED = 1, -}; - -enum PTO2CompletionFlag : uint8_t { - PTO2_COMPLETION_DONE = 2, -}; - -enum PTO2DeferredCompletionFlag : uint8_t { - PTO2_SUBTASK_DEFERRED = 4, -}; - -enum PTO2DispatchPropagationFlag : uint8_t { - PTO2_DISPATCH_PROPAGATED = 8, +enum PTO2TaskLifecycleFlag : uint8_t { + PTO2_LIFECYCLE_FLAGS_NONE = 0, + PTO2_READY_CLAIMED = 1U << 0, + PTO2_COMPLETION_DONE = 1U << 1, + PTO2_SUBTASK_DEFERRED = 1U << 2, + PTO2_DISPATCH_PROPAGATED = 1U << 3, }; static_assert((PTO2_DISPATCH_PROPAGATED & (PTO2_READY_CLAIMED | PTO2_COMPLETION_DONE | PTO2_SUBTASK_DEFERRED)) == 0); @@ -520,7 +511,7 @@ struct alignas(64) PTO2TaskSlotState { bool allow_early_resolve{false}; // Concurrent lifecycle updates preserve unrelated bits. Slot reuse clears // the byte only after the previous task lifetime is quiescent. - std::atomic ready_state{PTO2_READY_UNCLAIMED}; + std::atomic lifecycle_flags{PTO2_LIFECYCLE_FLAGS_NONE}; int32_t dep_pool_mark{0}; // Dep pool top after Orch-side wiring std::atomic completed_subtasks{0}; // Each core completion increments by 1 @@ -569,23 +560,23 @@ struct alignas(64) PTO2TaskSlotState { // Lock-free callers use this only as a fast-path hint. A false result is // rechecked by try_mark_dispatch_propagated() while holding fanout_lock. bool has_dispatch_propagated() const { - return (ready_state.load(std::memory_order_acquire) & PTO2_DISPATCH_PROPAGATED) != 0; + return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_DISPATCH_PROPAGATED) != 0; } // The propagation owner holds fanout_lock from this claim through its // fanout snapshot so wiring can classify late edges exactly once. bool try_mark_dispatch_propagated() { - return (ready_state.fetch_or(PTO2_DISPATCH_PROPAGATED, std::memory_order_acq_rel) & PTO2_DISPATCH_PROPAGATED) == - 0; + return (lifecycle_flags.fetch_or(PTO2_DISPATCH_PROPAGATED, std::memory_order_acq_rel) & + PTO2_DISPATCH_PROPAGATED) == 0; } void mark_completed() { task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); - ready_state.fetch_or(PTO2_COMPLETION_DONE, std::memory_order_release); + lifecycle_flags.fetch_or(PTO2_COMPLETION_DONE, std::memory_order_release); } bool is_completion_flag_set() const { - return (ready_state.load(std::memory_order_acquire) & PTO2_COMPLETION_DONE) != 0; + return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_COMPLETION_DONE) != 0; } // Set by any subtask FIN that pushed deferred-completion CONDITIONs to the @@ -594,10 +585,10 @@ struct alignas(64) PTO2TaskSlotState { // release write is sequenced before on_subtask_complete's acq_rel fetch_add // and the acquire read after, so all earlier subtasks' writes are visible to // the last subtask. - void mark_any_subtask_deferred() { ready_state.fetch_or(PTO2_SUBTASK_DEFERRED, std::memory_order_release); } + void mark_any_subtask_deferred() { lifecycle_flags.fetch_or(PTO2_SUBTASK_DEFERRED, std::memory_order_release); } bool has_any_subtask_deferred() const { - return (ready_state.load(std::memory_order_acquire) & PTO2_SUBTASK_DEFERRED) != 0; + return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_SUBTASK_DEFERRED) != 0; } void set_allow_early_resolve(bool v) { allow_early_resolve = v; } @@ -621,7 +612,7 @@ struct alignas(64) PTO2TaskSlotState { fanout_refcount.store(0, std::memory_order_relaxed); completed_subtasks.store(0, std::memory_order_relaxed); next_block_idx.store(0, std::memory_order_relaxed); - ready_state.store(PTO2_READY_UNCLAIMED, std::memory_order_relaxed); + lifecycle_flags.store(PTO2_LIFECYCLE_FLAGS_NONE, std::memory_order_relaxed); allow_early_resolve = false; // Note: payload early-dispatch fields (state, masks, fanin, publication count) // are NOT reset here — this method skips the payload by contract. They are diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index c3941aeeb1..87d74878c1 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -544,12 +544,12 @@ struct PTO2SchedulerState { } bool try_claim_ready_once(PTO2TaskSlotState &slot_state) { - uint8_t state = slot_state.ready_state.load(std::memory_order_acquire); + uint8_t flags = slot_state.lifecycle_flags.load(std::memory_order_acquire); for (;;) { - if ((state & PTO2_READY_CLAIMED) != 0) return false; - uint8_t desired = state | PTO2_READY_CLAIMED; - if (slot_state.ready_state.compare_exchange_weak( - state, desired, std::memory_order_acq_rel, std::memory_order_acquire + if ((flags & PTO2_READY_CLAIMED) != 0) return false; + uint8_t desired_flags = flags | PTO2_READY_CLAIMED; + if (slot_state.lifecycle_flags.compare_exchange_weak( + flags, desired_flags, std::memory_order_acq_rel, std::memory_order_acquire )) { return true; } @@ -1053,18 +1053,18 @@ struct PTO2SchedulerState { PTO2TaskSlotState &slot_state, uint64_t &atomic_count, uint64_t &push_wait, EarlyDispatchReleaseSink *sink = nullptr ) { - uint8_t state = slot_state.ready_state.load(std::memory_order_acquire); - atomic_count += 1; // ready_state load + uint8_t flags = slot_state.lifecycle_flags.load(std::memory_order_acquire); + atomic_count += 1; // lifecycle_flags load for (;;) { - if ((state & PTO2_READY_CLAIMED) != 0) return false; - uint8_t desired = state | PTO2_READY_CLAIMED; - if (slot_state.ready_state.compare_exchange_weak( - state, desired, std::memory_order_acq_rel, std::memory_order_acquire + if ((flags & PTO2_READY_CLAIMED) != 0) return false; + uint8_t desired_flags = flags | PTO2_READY_CLAIMED; + if (slot_state.lifecycle_flags.compare_exchange_weak( + flags, desired_flags, std::memory_order_acq_rel, std::memory_order_acquire )) { - atomic_count += 1; // ready_state CAS + atomic_count += 1; // lifecycle_flags CAS break; } - atomic_count += 1; // failed ready_state CAS + atomic_count += 1; // failed lifecycle_flags CAS } // Early-dispatch: pre-staged tasks are released by doorbell @@ -1207,7 +1207,7 @@ struct PTO2SchedulerState { slot_state.unlock_fanout(); #if SIMPLER_SCHED_PROFILING - lock_atomics += 3; // task_state.store + ready_state.fetch_or + unlock.store + lock_atomics += 3; // task_state.store + lifecycle_flags.fetch_or + unlock.store g_sched_lock_atomic_count[thread_idx] += lock_atomics; g_sched_lock_wait_cycle[thread_idx] += lock_wait; PTO2_SCHED_CYCLE_LAP(g_sched_lock_cycle[thread_idx]); diff --git a/tests/ut/cpp/a2a3/test_wiring.cpp b/tests/ut/cpp/a2a3/test_wiring.cpp index 72ef904e39..eeed532c83 100644 --- a/tests/ut/cpp/a2a3/test_wiring.cpp +++ b/tests/ut/cpp/a2a3/test_wiring.cpp @@ -226,7 +226,7 @@ TEST_F(WiringTest, WireTaskProducersPendingTaskNotReady) { EXPECT_EQ(producer_slots[1].fanout_head->slot_state, &task_slot); } -TEST_F(WiringTest, DispatchPropagationMarkerPreservesReadyStateFlagsAndResets) { +TEST_F(WiringTest, DispatchPropagationMarkerPreservesLifecycleFlagsAndResets) { alignas(64) PTO2TaskSlotState producer; init_slot(producer, PTO2_TASK_PENDING, 1, 1); @@ -237,13 +237,13 @@ TEST_F(WiringTest, DispatchPropagationMarkerPreservesReadyStateFlagsAndResets) { producer.mark_any_subtask_deferred(); EXPECT_FALSE(producer.try_mark_dispatch_propagated()); EXPECT_EQ( - producer.ready_state.load(std::memory_order_relaxed), + producer.lifecycle_flags.load(std::memory_order_relaxed), PTO2_READY_CLAIMED | PTO2_COMPLETION_DONE | PTO2_SUBTASK_DEFERRED | PTO2_DISPATCH_PROPAGATED ); producer.reset_for_reuse(); EXPECT_FALSE(producer.has_dispatch_propagated()); - EXPECT_EQ(producer.ready_state.load(std::memory_order_relaxed), PTO2_READY_UNCLAIMED); + EXPECT_EQ(producer.lifecycle_flags.load(std::memory_order_relaxed), PTO2_LIFECYCLE_FLAGS_NONE); } // =============================================================================