fix(scheduler): tag sync_start drain ack/election state by generation

docs/dynamic-linking.md

@@ -223,8 +223,8 @@ SchedulerContext owns its own teardown:
 
 - `SchedulerContext::deinit()` resets every scheduler-owned field —
   per-core states, payloads, sync-start drain coordination
-  (`sync_start_pending` / `drain_worker_elected` / `drain_ack_mask` /
-  `pending_task`), task counters, worker-id lists,
+  (`sync_start_pending` / `drain_worker_elected` / `drain_attempt` /
+  `drain_ack_attempts_` / `pending_task`), task counters, worker-id lists,
   core trackers, `cores_total_num_` / `aic_count_` / `aiv_count_`,
   `regs_`, `sched_`, `func_id_to_addr_`, and the `pto2_init_*` flags.
 - `AicpuExecutor::deinit()` calls `sched_ctx_.deinit()` first, then resets

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

@@ -1003,7 +1003,10 @@ void SchedulerContext::deinit() {
     // would otherwise leave dirty pending/elected/ack state for the next reuse.
     drain_state_.sync_start_pending.store(0, std::memory_order_release);
     drain_state_.drain_worker_elected.store(0, std::memory_order_release);
-    drain_state_.drain_ack_mask.store(0, std::memory_order_release);
+    drain_state_.drain_attempt.store(0, std::memory_order_release);
+    for (int32_t t = 0; t < MAX_AICPU_THREADS; t++) {
+        drain_ack_attempts_[t].store(0, std::memory_order_release);
+    }
     drain_state_.drain_stage_go.store(0, std::memory_order_release);
     drain_state_.drain_stage_done_mask.store(0, std::memory_order_release);
     drain_state_.drain_running_staged.store(0, std::memory_order_release);

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

@@ -472,14 +472,18 @@ void SchedulerContext::check_running_cores_for_completion(
 bool SchedulerContext::enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t block_num) {
     int32_t expected = 0;
     if (!drain_state_.sync_start_pending.compare_exchange_strong(
-            expected, -1, std::memory_order_relaxed, std::memory_order_relaxed
+            expected, -1, std::memory_order_acquire, std::memory_order_relaxed
         )) {
         return false;  // Another thread already holds the drain slot.
     }
-    // We own the drain slot.  Store the task and reset the coordination flags before making
-    // it visible.
+    // We own the drain slot.  Advance the attempt FIRST: a delayed try_elect() from
+    // the previous round races drain_attempt against drain_worker_elected's CAS, so
+    // bumping the attempt before publishing the new task and clearing
+    // drain_worker_elected closes the same ABA window the insufficient-resource
+    // retry path already guards against.
+    uint64_t next_attempt = drain_state_.drain_attempt.load(std::memory_order_relaxed) + 1;
+    drain_state_.drain_attempt.store(next_attempt, std::memory_order_release);
     drain_state_.pending_task.store(slot_state, std::memory_order_release);
-    drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed);
     drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed);
     drain_state_.drain_stage_go.store(0, std::memory_order_relaxed);
     drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed);
@@ -640,7 +644,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block
 //
 // Protocol:
 //   1. Ack barrier: all threads signal they've stopped dispatch, spin until all acked.
-//      If this thread's ack bit gets cleared while waiting, a reset occurred -- return.
+//      Each ack carries the current attempt; an attempt change invalidates the old cohort.
 //   2. Election + availability: one thread wins the CAS. It checks global resources; if
 //      insufficient it resets ack/election so all threads resume completion polling to free
 //      cores, then retry. If sufficient it releases parallel staging (stage_go).
@@ -674,25 +678,23 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
     if (block_num == 0) return;
 
     uint32_t all_acked = (1u << active_sched_threads_) - 1;
+    uint64_t drain_attempt = drain_state_.drain_attempt.load(std::memory_order_acquire);
 
-    // Ack barrier -- signal this thread has stopped dispatch.
-    drain_state_.drain_ack_mask.fetch_or(1u << thread_idx, std::memory_order_release);
+    // Ack barrier -- signal this thread has stopped dispatch for this attempt.
+    drain_ack_attempts_[thread_idx].store(drain_attempt, std::memory_order_release);
 
     // Spin until all threads have acked.
-    // If our bit is cleared while waiting, elected reset due to insufficient resources.
+    // An attempt change invalidates every ack from the prior cohort.
     while (true) {
         if (is_completed()) return;
-        uint32_t ack = drain_state_.drain_ack_mask.load(std::memory_order_acquire);
+        if (drain_state_.drain_attempt.load(std::memory_order_acquire) != drain_attempt) return;
+        uint32_t ack = sync_start_drain_ack_mask_for_attempt(drain_ack_attempts_, active_sched_threads_, drain_attempt);
         if ((ack & all_acked) == all_acked) break;
-        if ((ack & (1u << thread_idx)) == 0) return;
         SPIN_WAIT_HINT();
     }
     // Election -- exactly one thread wins the CAS.
-    int32_t expected = 0;
-    drain_state_.drain_worker_elected.compare_exchange_strong(
-        expected, thread_idx + 1, std::memory_order_acquire, std::memory_order_relaxed
-    );
-    bool elected = drain_state_.drain_worker_elected.load(std::memory_order_relaxed) == thread_idx + 1;
+    bool elected = sync_start_drain_try_elect(drain_state_, thread_idx, drain_attempt);
+    if (!elected && drain_state_.drain_attempt.load(std::memory_order_acquire) != drain_attempt) return;
 
     PTO2TaskSlotState *slot_state = drain_state_.pending_task.load(std::memory_order_acquire);
     // OWNER is acquired before the drain is published and persists through
@@ -704,9 +706,12 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
     if (elected) {
         if (slot_state == nullptr) {
             // pending_task observed null only when a concurrent drain completion already cleared
-            // it. Stale-elected: release the election lock and return. Do NOT clear drain_ack_mask
-            // / sync_start_pending -- a *new* drain run may already be accumulating acks.
-            drain_state_.drain_worker_elected.store(0, std::memory_order_release);
+            // it. Stale-elected: release our election lock and return. Do NOT clear
+            // sync_start_pending -- a new drain run may already be accumulating acks.
+            int32_t owner = thread_idx + 1;
+            drain_state_.drain_worker_elected.compare_exchange_strong(
+                owner, 0, std::memory_order_release, std::memory_order_relaxed
+            );
             return;
         }
         PTO2ResourceShape shape = slot_state->active_mask.to_shape();
@@ -716,7 +721,7 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
             count_global_available(shape, slot_state->active_mask.core_mask(), /*include_pending=*/gated);
         if (available < block_num) {
             // Insufficient -- reset so all threads resume completion polling to free cores, then retry.
-            drain_state_.drain_ack_mask.store(0, std::memory_order_release);
+            drain_state_.drain_attempt.store(drain_attempt + 1, std::memory_order_release);
             drain_state_.drain_worker_elected.store(0, std::memory_order_release);
             return;
         }
@@ -729,6 +734,8 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
         // insufficient resources).
         while (drain_state_.drain_stage_go.load(std::memory_order_acquire) == 0) {
             if (is_completed()) return;
+            if (drain_state_.sync_start_pending.load(std::memory_order_acquire) == 0) return;
+            if (drain_state_.drain_attempt.load(std::memory_order_acquire) != drain_attempt) return;
             if (drain_state_.drain_worker_elected.load(std::memory_order_acquire) == 0) return;
             SPIN_WAIT_HINT();
         }
@@ -751,13 +758,12 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
     drain_state_.drain_stage_done_mask.fetch_or(1u << thread_idx, std::memory_order_release);
 
     if (!elected) {
-        // Non-elected: staging done; wait for the elected thread to reopen the gate. Exiting via
-        // sync_start_pending==0 (release/acquire) or drain_worker_elected==0 both synchronize
-        // with the elected's finalize (its release fence sequences the seed before both stores),
-        // so the running_slot_count seed is visible before this thread resumes completions.
+        // Non-elected: staging done; wait for the elected thread to reopen the gate.
+        // sync_start_pending==0 synchronizes with finalize, so the running_slot_count seed is
+        // visible before this thread resumes completions.
         while (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) {
             if (is_completed()) return;
-            if (drain_state_.drain_worker_elected.load(std::memory_order_acquire) == 0) return;
+            if (drain_state_.drain_attempt.load(std::memory_order_acquire) != drain_attempt) return;
             SPIN_WAIT_HINT();
         }
         return;
@@ -778,16 +784,15 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
         );
     }
     // Clear drain state and reopen the gate FIRST, so the other threads resume immediately.
-    // Release fence sequences the seed + tracker mutations before every clear, so any thread
-    // that acquire-observes one of them (sync_start_pending==0 / drain_worker_elected==0) sees
-    // the seed. `slot_state` is a local holding the fa_fused slot (not drain_state_), so it stays
-    // valid for the propagate below even if a new drain reuses pending_task after reopen.
+    // Release fence sequences the seed + tracker mutations before sync_start_pending. The
+    // elected value remains published while the gate is open; the next drain owner resets it
+    // behind the -1 sentinel. Followers identify reopen by gate-open or attempt change.
+    // `slot_state` is a local holding the fa_fused slot (not drain_state_), so it stays valid for
+    // the propagate below even if a new drain reuses pending_task after reopen.
     std::atomic_thread_fence(std::memory_order_release);
     drain_state_.pending_task.store(nullptr, std::memory_order_release);
     drain_state_.drain_stage_go.store(0, std::memory_order_relaxed);
     drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed);
-    drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed);
-    drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed);
     drain_state_.sync_start_pending.store(0, std::memory_order_release);
 
     // Recheck after publishing the drain seed. The producer-side rendezvous check can race

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

@@ -139,6 +139,7 @@ class SchedulerContext {
 
     // sync_start drain coordination
     SyncStartDrainState drain_state_;
+    std::atomic<uint64_t> drain_ack_attempts_[MAX_AICPU_THREADS]{};
 
 #if SIMPLER_DFX
     SchedL2SwimlaneCounters sched_l2_swimlane_[MAX_AICPU_THREADS];

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

@@ -532,10 +532,10 @@ struct alignas(64) SchedL2SwimlaneCounters {
 // When sync_start_pending != 0, all scheduler threads skip dispatch
 // (only process completions) until the drain worker finishes launching all blocks.
 struct alignas(64) SyncStartDrainState {
-    std::atomic<int32_t> sync_start_pending{0};    // 0=normal; -1=initializing; >0=active (value=block_num)
-    std::atomic<int32_t> drain_worker_elected{0};  // 0=none; >0: elected thread's (thread_idx+1)
-    std::atomic<uint32_t> drain_ack_mask{0};       // bit per thread; all-set = all threads reached ack barrier
+    std::atomic<int32_t> sync_start_pending{0};              // 0=normal; -1=initializing; >0=active (value=block_num)
+    std::atomic<int32_t> drain_worker_elected{0};            // 0=none; >0: elected thread's (thread_idx+1)
     std::atomic<PTO2TaskSlotState *> pending_task{nullptr};  // held task (not re-queued)
+    std::atomic<uint64_t> drain_attempt{0};                  // incremented whenever an ack/election round is reset
     // Parallel staging: after the elected thread confirms global availability it sets
     // stage_go, releasing every thread to stage its OWN cores concurrently (vs the old
     // single-thread serial fill). Each thread ORs its bit into stage_done_mask when it
@@ -548,4 +548,33 @@ struct alignas(64) SyncStartDrainState {
 };
 static_assert(sizeof(SyncStartDrainState) == 64);
 
+inline uint32_t sync_start_drain_ack_mask_for_attempt(
+    const std::atomic<uint64_t> *ack_attempts, int32_t thread_count, uint64_t attempt
+) {
+    uint32_t mask = 0;
+    for (int32_t t = 0; t < thread_count; t++) {
+        if (ack_attempts[t].load(std::memory_order_acquire) == attempt) {
+            mask |= 1u << t;
+        }
+    }
+    return mask;
+}
+
+inline bool sync_start_drain_try_elect(SyncStartDrainState &state, int32_t thread_idx, uint64_t attempt) {
+    int32_t expected = 0;
+    bool won = state.drain_worker_elected.compare_exchange_strong(
+        expected, thread_idx + 1, std::memory_order_acquire, std::memory_order_relaxed
+    );
+    if (state.drain_attempt.load(std::memory_order_acquire) == attempt) {
+        return won;
+    }
+    if (won) {
+        int32_t owner = thread_idx + 1;
+        state.drain_worker_elected.compare_exchange_strong(
+            owner, 0, std::memory_order_release, std::memory_order_relaxed
+        );
+    }
+    return false;
+}
+
 #endif  // SCHEDULER_TYPES_H

src/a2a3/runtime/tensormap_and_ringbuffer/common/pto_runtime_status.h

@@ -39,6 +39,8 @@
 #define PTO2_ERROR_ASYNC_COMPLETION_INVALID 101
 #define PTO2_ERROR_ASYNC_WAIT_OVERFLOW 102
 #define PTO2_ERROR_ASYNC_REGISTRATION_FAILED 103
+#define PTO2_ERROR_DRAIN_ABA_TEST_NOT_EXERCISED \
+    104  // SIMPLER_DRAIN_ABA_TEST: hook never observed the forced attempt change
 
 // Sub-classification of a PTO2_ERROR_SCHEDULER_TIMEOUT (code 100). The top-level
 // sched_error_code stays 100 for backward compatibility; this detail value tells

src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp

@@ -608,9 +608,13 @@ static void apply_orch_sched_env_flags(Runtime *runtime) {
     const char *serial_env = std::getenv("SIMPLER_TMR_SERIAL_ORCH_SCHED_ENABLE");
     runtime->dev.serial_orch_sched =
         serial_env && (serial_env[0] == '1' || serial_env[0] == 't' || serial_env[0] == 'T');
+    const char *drain_aba_env = std::getenv("SIMPLER_DRAIN_ABA_TEST");
+    runtime->dev.drain_aba_test_mode =
+        drain_aba_env && (drain_aba_env[0] == '1' || drain_aba_env[0] == 't' || drain_aba_env[0] == 'T');
     LOG_INFO_V0(
         "Serial orchestrator-to-scheduler start gate: %s", runtime->dev.serial_orch_sched ? "enabled" : "disabled"
     );
+    LOG_INFO_V0("Drain ABA test hook: %s", runtime->dev.drain_aba_test_mode ? "enabled" : "disabled");
 }
 
 // per-(cid,config): reserve and acquire the static device pools. GM heap, PTO2

src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h

@@ -187,6 +187,12 @@ struct alignas(64) DeviceRuntimeLaunchDesc {
     // Controlled via SIMPLER_TMR_SERIAL_ORCH_SCHED_ENABLE environment variable.
     bool serial_orch_sched;
 
+    // Test-only drain interleaving hook. When enabled, scheduler threads create a
+    // deterministic stale-attempt window in sync_start drain to prove whether the
+    // reusable ack/election fields are generation-less.
+    // Controlled via SIMPLER_DRAIN_ABA_TEST environment variable.
+    bool drain_aba_test_mode;
+
     void *gm_sm_ptr_;                        // GM pointer to PTO2 shared memory (device)
     ChipStorageTaskArgs orch_args_storage_;  // Copy of args for device
 

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

@@ -1130,6 +1130,7 @@ int32_t SchedulerContext::pre_handshake_init(
     aic_count_ = cores_total_num_ / 3;
     aiv_count_ = (cores_total_num_ * 2) / 3;
     active_sched_threads_ = (sched_thread_num_ > 0) ? sched_thread_num_ : aicpu_thread_num_;
+    drain_aba_test_mode_ = runtime->dev.drain_aba_test_mode;
     handshake_failed_.store(false, std::memory_order_release);
 
     // State the barrier-free per-thread init path no longer reaches via
@@ -1307,7 +1308,10 @@ void SchedulerContext::deinit() {
     // would otherwise leave dirty pending/elected/ack state for the next reuse.
     drain_state_.sync_start_pending.store(0, std::memory_order_release);
     drain_state_.drain_worker_elected.store(0, std::memory_order_release);
-    drain_state_.drain_ack_mask.store(0, std::memory_order_release);
+    drain_state_.drain_attempt.store(0, std::memory_order_release);
+    for (int32_t t = 0; t < MAX_AICPU_THREADS; t++) {
+        drain_ack_attempts_[t].store(0, std::memory_order_release);
+    }
     drain_state_.drain_stage_go.store(0, std::memory_order_release);
     drain_state_.drain_stage_done_mask.store(0, std::memory_order_release);
     drain_state_.drain_running_staged.store(0, std::memory_order_release);
@@ -1324,6 +1328,8 @@ void SchedulerContext::deinit() {
     aiv_count_ = 0;
     cores_total_num_ = 0;
     aicpu_thread_num_ = 0;
+    drain_aba_test_mode_ = false;
+    drain_test_victim_armed_.store(0, std::memory_order_relaxed);
     sched_thread_num_ = 0;
     active_sched_threads_ = 0;
     for (int32_t t = 0; t < MAX_AICPU_THREADS; t++) {