Fix: serialize block range claims

[ChaoWao]

Summary

• Add one CAS-based claim_block_range() helper for task block ranges.
• Use it in normal ready dispatch, early-dispatch staging, and sync-start
  drain dispatch.
• Cover concurrent claims and partial early-staged release routing.

Root cause

After a partially pre-staged task is released, normal ready dispatch can run
at the same time as a late early-dispatch stager. The stager already claims
through CAS, but normal dispatch used a relaxed load followed by a relaxed
store:

normal: load start=2
stager: CAS 2 -> 7, owns [2,7)
normal: store 5, owns [2,5)

That both overlaps a block range and moves next_block_idx backwards.
Routing every claimant through the same CAS loop makes ranges disjoint.

This race exists on main independently of #1304, so it is intentionally
submitted separately.

Testing

• C++ non-hardware unit suite: 54/54 passed
• Targeted concurrent claim and partial-release tests: 1000 repeats passed
• Full pre-commit run --from-ref main --to-ref HEAD: passed
• Combined latest-main + Add: early-dispatch for require_sync_start SPMD cohorts (a2a3 tmr) #1304 Qwen3-14B reference-kernel stress:
  10/10 passed

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 133b67d9-6dc0-4f40-b5be-4b29e9df85fa

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The change adds an atomic block-range claim helper to task slots, updates scheduler dispatch paths to use it, revises staging comments, and adds tests for concurrent claims and partially staged early-dispatch routing.

Changes

Concurrent block claiming

Layer / File(s)	Summary
Atomic block-range claim helper src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h	Adds CAS-based bounded range reservation through claim_block_range() and documents concurrent cursor advancement.
Scheduler dispatch integration src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp, src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp, src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h	Routes worker, normal, and early dispatch through the shared claim helper while updating staging-order comments.
Concurrent and partial-release validation tests/ut/cpp/a2a3/test_wiring.cpp	Tests non-overlapping concurrent claims and remainder routing after partial staging.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant WorkerDispatch
  participant SchedulerDispatch
  participant PTO2TaskSlotState
  WorkerDispatch->>PTO2TaskSlotState: claim_block_range(block_num, valid.count(), start)
  SchedulerDispatch->>PTO2TaskSlotState: claim_block_range(logical_block_num, available, start)
  PTO2TaskSlotState-->>SchedulerDispatch: claimed contiguous range
  SchedulerDispatch->>SchedulerDispatch: publish or stage claimed blocks

Loading

Possibly related PRs

• hw-native-sys/simpler#989: Updates related scheduler block-claim and dispatch paths.
• hw-native-sys/simpler#1326: Modifies the related early-dispatch scheduling flow.

Poem

I’m a rabbit claiming blocks in a row,
With CAS hops quick as the winds that blow.
No overlap, no muddle, the ranges align,
Partial dispatch leaves a tail just fine.
Hop, scheduler—your queues now rhyme!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly matches the main change: serializing block range claims via a shared CAS-based helper.
Description check	✅ Passed	The description accurately describes the CAS helper, affected dispatch paths, race condition, and added tests.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request introduces a thread-safe 'claim_block_range' helper method in 'PTO2TaskSlotState' using a CAS loop to manage concurrent block claims. It refactors the scheduler dispatch and completion logic to use this helper, ensuring safe concurrent execution, and adds corresponding unit tests. The feedback suggests a minor optimization in 'claim_block_range' to directly assign 'start' to 'current' upon a successful CAS, simplifying the logic and avoiding unnecessary arithmetic.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.