Fix: serialize block range claims

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

@@ -491,12 +491,28 @@ struct alignas(64) PTO2TaskSlotState {
     std::atomic<int16_t> completed_subtasks{0};  // Each core completion increments by 1
     int16_t total_required_subtasks{0};          // = logical_block_num * popcount(active_mask)
     int16_t logical_block_num{1};                // Total logical blocks (set by orchestrator)
-    // Next block to dispatch. Atomic so concurrent early-dispatch stagers can each
-    // claim a distinct block via CAS; normal dispatch (ready-queue serialized)
-    // uses plain relaxed load/store. The two phases never overlap in time (staging
-    // happens before release; normal dispatch of the remainder happens after).
+    // Next block to dispatch. Normal dispatch and late early-dispatch stagers
+    // can run concurrently after a partial staged release. All paths claim
+    // ranges through claim_block_range().
     std::atomic<int16_t> next_block_idx{0};
 
+    int32_t claim_block_range(int32_t block_limit, int32_t max_count, int32_t &start) {
+        int16_t current = next_block_idx.load(std::memory_order_relaxed);
+        while (current < block_limit && max_count > 0) {
+            int32_t count = block_limit - current;
+            if (count > max_count) count = max_count;
+            int16_t desired = static_cast<int16_t>(current + count);
+            if (next_block_idx.compare_exchange_weak(
+                    current, desired, std::memory_order_seq_cst, std::memory_order_relaxed
+                )) {
+                start = current;
+                return count;
+            }
+        }
+        start = current;
+        return 0;
+    }
+
     /**
      * Bind the slot-invariant ring id. Called once per slot during
      * RingSchedState::init(); ring_id never changes across reuses.

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

@@ -557,20 +557,8 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block
     auto stage = [&](CoreTracker::BitStates valid, bool to_pending) {
         while (valid.has_value()) {
             int32_t avail = valid.count();
-            int32_t start = 0, claim = 0;
-            while (true) {
-                int16_t cur = slot_state->next_block_idx.load(std::memory_order_relaxed);
-                if (cur >= block_num) return;  // all blocks claimed
-                int32_t cnt = block_num - cur;
-                if (cnt > avail) cnt = avail;
-                if (slot_state->next_block_idx.compare_exchange_weak(
-                        cur, static_cast<int16_t>(cur + cnt), std::memory_order_seq_cst, std::memory_order_relaxed
-                    )) {
-                    start = cur;
-                    claim = cnt;
-                    break;
-                }
-            }
+            int32_t start = 0;
+            int32_t claim = slot_state->claim_block_range(block_num, avail, start);
             if (claim == 0) return;
 #if SIMPLER_DFX
             bool sub_prof = l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES;

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

@@ -320,7 +320,7 @@ class SchedulerContext {
 
     // Stage the already-claimed range [start, start+count) of consumer `c` onto
     // thread_idx's idle (RUNNING slot) then pending (gated-pending, promote-on-FIN)
-    // cores from the provided free-core sets. The caller advances next_block_idx and
+    // cores from the provided free-core sets. The caller claims next_block_idx and
     // re-pushes `c` BEFORE calling, so this expensive prepare+publish runs
     // concurrently with peers (mirrors the normal SPMD dispatch path). Returns the
     // number of blocks staged.

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

@@ -437,20 +437,19 @@ void SchedulerContext::dispatch_shape(
                 break;
             }
 
-            dispatched_any = true;
-            try_pushed = true;
             // Claim a contiguous range of blocks, hand the slot back to the
             // ready queue immediately, then perform the expensive dispatches.
             // This lets other schedulers concurrently claim and dispatch the
             // remaining blocks of the same SPMD task instead of spinning while
             // this thread fills all its own cores. Only local `start + b` is
             // read after the push — `next_block_idx` may already be advanced
             // by another scheduler that popped the slot.
-            int32_t start = slot_state->next_block_idx.load(std::memory_order_relaxed);
-            int32_t remaining = slot_state->logical_block_num - start;
             int32_t available = is_mix ? selected_mix_clusters.count() : cores.count();
-            int32_t claim = std::min(available, remaining);
-            slot_state->next_block_idx.store(static_cast<int16_t>(start + claim), std::memory_order_relaxed);
+            int32_t start = 0;
+            int32_t claim = slot_state->claim_block_range(slot_state->logical_block_num, available, start);
+            if (claim == 0) continue;
+            dispatched_any = true;
+            try_pushed = true;
 
             published_list[published_n] = slot_state;
             published_counts[published_n] = static_cast<int16_t>(claim);
@@ -598,11 +597,11 @@ void SchedulerContext::dispatch_ready_tasks(
 }
 
 // Stage the ALREADY-CLAIMED range [start, start+count) of consumer `c` onto
-// thread_idx's idle then pending cores. The caller (the queue drain) has advanced
-// next_block_idx by `count` under pop-exclusivity AND re-pushed `c` for peers
+// thread_idx's idle then pending cores. The caller has atomically advanced
+// next_block_idx by `count` AND re-pushed `c` for peers
 // BEFORE calling this — so this, the expensive prepare+publish, runs CONCURRENTLY
 // with peers staging other ranges of the same consumer. This mirrors the normal
-// SPMD dispatch path (claim range -> store next_block_idx -> re-push -> dispatch).
+// SPMD dispatch path (claim range -> re-push -> dispatch).
 // `idle`/`pend` are this thread's free-core sets, sized so idle.count+pend.count >=
 // count (the caller clamped the claim to them), so all `count` blocks get a core.
 //
@@ -760,22 +759,8 @@ SchedulerContext::early_dispatch_shape(int32_t thread_idx, PTO2ResourceShape sha
             sched_->early_dispatch_queues[s].push_batch(&batch[bi], got - bi);
             break;
         }
-        // CAS-claim a contiguous range [start, start+claim) sized to this thread's
-        // free cores; CAS keeps it atomic against peers AND normal dispatch.
-        int32_t start = 0, claim = 0;
-        while (true) {
-            int16_t cur = c->next_block_idx.load(std::memory_order_relaxed);
-            if (cur >= c->logical_block_num) break;  // fully claimed
-            int32_t cnt = c->logical_block_num - cur;
-            if (cnt > freecores) cnt = freecores;
-            if (c->next_block_idx.compare_exchange_weak(
-                    cur, static_cast<int16_t>(cur + cnt), std::memory_order_seq_cst, std::memory_order_relaxed
-                )) {
-                start = cur;
-                claim = cnt;
-                break;
-            }
-        }
+        int32_t start = 0;
+        int32_t claim = c->claim_block_range(c->logical_block_num, freecores, start);
         if (claim == 0) continue;  // nothing left to claim -> drop (no re-push)
         // Re-push for concurrent peers BEFORE the expensive staging.
         if (start + claim < c->logical_block_num) {

tests/ut/cpp/a2a3/test_wiring.cpp

@@ -382,6 +382,61 @@ TEST_F(WiringTest, EarlyDispatchWaitsForAllProducerBlocksPublished) {
     EXPECT_EQ(payload.dispatch_fanin.load(), payload.fanin_actual_count);
 }
 
+TEST_F(WiringTest, ConcurrentBlockRangeClaimsDoNotOverlap) {
+    alignas(64) PTO2TaskSlotState task_slot;
+    init_slot(task_slot, PTO2_TASK_PENDING, 1, 1);
+    task_slot.logical_block_num = 8;
+
+    struct ClaimedRange {
+        int32_t start = -1;
+        int32_t count = 0;
+    } ranges[2];
+    std::atomic<int32_t> ready{0};
+    std::atomic<bool> start{false};
+
+    auto claim = [&](int32_t index) {
+        ready.fetch_add(1, std::memory_order_release);
+        while (!start.load(std::memory_order_acquire)) {}
+        ranges[index].count = task_slot.claim_block_range(task_slot.logical_block_num, 5, ranges[index].start);
+    };
+
+    std::thread first(claim, 0);
+    std::thread second(claim, 1);
+    while (ready.load(std::memory_order_acquire) != 2) {}
+    start.store(true, std::memory_order_release);
+    first.join();
+    second.join();
+
+    ClaimedRange *lower = ranges[0].start < ranges[1].start ? &ranges[0] : &ranges[1];
+    ClaimedRange *upper = lower == &ranges[0] ? &ranges[1] : &ranges[0];
+    EXPECT_EQ(lower->start, 0);
+    EXPECT_EQ(lower->count, 5);
+    EXPECT_EQ(upper->start, 5);
+    EXPECT_EQ(upper->count, 3);
+    EXPECT_EQ(lower->start + lower->count, upper->start);
+    EXPECT_EQ(task_slot.next_block_idx.load(std::memory_order_relaxed), task_slot.logical_block_num);
+}
+
+TEST_F(WiringTest, PartialStagedReleaseRoutesRemainderToReadyQueue) {
+    alignas(64) PTO2TaskSlotState consumer;
+    init_slot(consumer, PTO2_TASK_PENDING, 1, 1);
+    consumer.logical_block_num = 5;
+    consumer.next_block_idx.store(2, std::memory_order_relaxed);
+    consumer.payload->early_dispatch_state.store(PTO2_EARLY_DISPATCH_STAGING, std::memory_order_relaxed);
+    consumer.payload->staged_core_mask[0].store(1, std::memory_order_relaxed);
+
+    EXPECT_TRUE(sched.route_ready_once(consumer));
+    EXPECT_EQ(consumer.payload->early_dispatch_state.load(), PTO2_EARLY_DISPATCH_DISPATCHED);
+    EXPECT_EQ(consumer.next_block_idx.load(), 2);
+
+    PTO2ResourceShape shape = consumer.active_mask.to_shape();
+    EXPECT_EQ(sched.ready_queues[static_cast<int32_t>(shape)].pop(), &consumer);
+
+    int32_t remaining_start = -1;
+    EXPECT_EQ(consumer.claim_block_range(consumer.logical_block_num, 5, remaining_start), 3);
+    EXPECT_EQ(remaining_start, 2);
+}
+
 TEST_F(WiringTest, EarlyDispatchDoorbellBitsHaveOneOwner) {
     constexpr uint64_t all_bits = 0b1111;
     constexpr uint64_t late_bits = 0b1010;