Skip to content

Feature: CAS-based watermark advancement and slot-reuse validation#143

Merged
ChaoWao merged 1 commit into
mainfrom
fix/cas-watermark-heap-tail-and-thread-guard
Feb 28, 2026
Merged

Feature: CAS-based watermark advancement and slot-reuse validation#143
ChaoWao merged 1 commit into
mainfrom
fix/cas-watermark-heap-tail-and-thread-guard

Conversation

@ChaoWao
Copy link
Copy Markdown
Collaborator

@ChaoWao ChaoWao commented Feb 28, 2026

Add ring buffer flow control so scheduler threads advance last_task_alive and heap_tail after consuming tasks, enabling slot reuse for long-running orchestrations.

aicpu_executor.cpp:

  • orch_pointers_ready_ atomic flag: scheduler threads wait for last thread to finish pointer setup before entering main loop
  • s_pto2_completed_by_task[] array: records which task_id set each slot's completion state to prevent stale state from recycled slots
  • CAS watermark loop: after task completion, mark slot as state=3, then CAS-advance last_task_alive and heap_tail
  • Fix heap_tail: convert absolute pointer to offset from gm_heap_base (same semantic as cda8448 fixed on scheduler side)
  • Fix thread_num_ == 4 hardcoded guard to thread_num_ > 1 for general multi-thread device orchestration scenarios
  • Add orch ready queue pointer cleanup in deinit() to prevent stale pointer access on deinit/re-init cycles

pto_orchestrator.h/cpp:

  • Add aicpu_completed_by_task pointer to PTO2OrchestratorState
  • Early-return path: verify completed_by_task matches producer_id to guard against stale state from recycled slots
  • Slot reset at allocation: clear task_completed and completed_by_task
  • Add RELAXED load rationale comment for completed_by_task

@gemini-code-assist
Copy link
Copy Markdown

Summary of Changes

Hello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the ring buffer flow control mechanism within the AICPU executor, particularly for long-running orchestrations in multi-threaded scenarios. By implementing a CAS-based watermark advancement for last_task_alive and heap_tail, it enables more efficient slot reuse. Additionally, a robust validation system using completed_by_task is introduced to prevent issues arising from stale state when slots are recycled, improving the overall stability and performance of the orchestration system.

Highlights

  • CAS-based Watermark Advancement: Implemented a Compare-And-Swap (CAS) mechanism to advance last_task_alive and heap_tail watermarks, enabling efficient slot reuse in ring buffers after tasks are fully consumed.
  • Slot Reuse Validation: Introduced s_pto2_completed_by_task array to record the task_id that set a slot's completion state, preventing issues from stale state when slots are recycled in multi-threaded scenarios.
  • Scheduler Thread Synchronization: Added an orch_pointers_ready_ atomic flag to ensure scheduler threads wait for the main thread to complete pointer setup before entering the main scheduling loop.
  • Heap Tail Calculation Correction: Fixed the heap_tail calculation to correctly represent an offset from gm_heap_base, aligning with existing semantic conventions for memory management.
  • Generalized Multi-threading Guard: Broadened the multi-threading condition from a hardcoded thread_num_ == 4 to thread_num_ > 1, making the logic applicable to general multi-thread device orchestration scenarios.
  • Orchestrator Pointer Cleanup: Ensured proper cleanup and reset of orchestrator ready queue pointers during the deinit() process to prevent stale pointer access upon re-initialization cycles.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • src/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp
    • Added orch_pointers_ready_ atomic flag to AicpuExecutor struct.
    • Declared s_pto2_completed_by_task static array for slot-reuse validation.
    • Updated thread synchronization condition from thread_num_ == 4 to thread_num_ > 1.
    • Initialized s_pto2_completed_by_task with -1 during one-time initialization.
    • Introduced a wait loop for orch_pointers_ready_ before the scheduling loop.
    • Stored task_id in s_pto2_completed_by_task upon task completion.
    • Implemented a CAS loop to advance last_task_alive and heap_tail after tasks are fully consumed, including resetting s_pto2_fanin_refcount.
    • Retrieved gm_heap_base for correct heap_tail offset calculation.
    • Assigned s_pto2_completed_by_task to rt->orchestrator.aicpu_completed_by_task.
    • Set orch_pointers_ready_ to true after pointer setup.
    • Reset orch_pointers_ready_ and orchestrator ready queue pointers in deinit().
  • src/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp
    • Modified pto2_add_consumer_to_producer to include a check against aicpu_completed_by_task to validate producer completion against the correct task ID.
    • Added logic to reset aicpu_task_completed and aicpu_completed_by_task when a slot is allocated for a new task.
  • src/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h
    • Added aicpu_completed_by_task member (pointer to int32_t) to PTO2OrchestratorState struct.
Activity
  • No human activity (comments, reviews, etc.) has been recorded for this pull request yet.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces significant enhancements for flow control and slot reuse in the ring buffer, crucial for long-running orchestrations, using CAS-based watermark advancement and completed_by_task validation. However, the implementation introduces security risks related to handling untrusted data from shared memory, specifically lacking validation for offsets and indices in aicpu_executor.cpp which could lead to arbitrary memory reads or system-wide stalls (DoS). Additionally, the use of raw pointers in shared memory structures for heap reclamation is fragile and potentially exploitable. Furthermore, a critical race condition in the update logic for heap_tail could lead to performance degradation.

Comment thread src/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp Outdated
Comment thread src/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp
@ChaoWao ChaoWao force-pushed the fix/cas-watermark-heap-tail-and-thread-guard branch 5 times, most recently from 97d20d2 to 9325d17 Compare February 28, 2026 14:38
…simulation thread starvation fix

Add lock-free heap_tail/last_task_alive advancement in AICPU parallel mode
with ticket-based generation counter (heap_tail_gen) to serialize concurrent
heap_tail writes. Add completed_by_task array for slot-reuse validation in
the orchestrator's early-return path, preventing stale state from recycled
slots.

Fix simulation thread starvation on resource-constrained CI runners (e.g.,
2 CPU cores running 13+ threads) by introducing a platform-level
SPIN_WAIT_HINT() macro for both AICore and AICPU idle polling loops. On
simulation (a2a3sim), expands to CPU pause/yield hint + sched_yield() to
let the OS scheduler give time slices to threads doing real work; on
hardware (a2a3), expands to no-op. Without this, idle threads in tight
spin loops starve the executing core, causing AICPU iteration-based
timeouts (MAX_IDLE_ITERATIONS) before kernels can complete.

Also fix race where scheduler threads could access orch_ready_queue pointers
before Thread 3 finished initializing them (orch_pointers_ready_ guard).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@ChaoWao ChaoWao force-pushed the fix/cas-watermark-heap-tail-and-thread-guard branch from 9325d17 to 6ec38cb Compare February 28, 2026 15:09
@ChaoWao ChaoWao merged commit 00a4ea4 into main Feb 28, 2026
3 checks passed
@ChaoWao ChaoWao deleted the fix/cas-watermark-heap-tail-and-thread-guard branch March 1, 2026 06:58
PKUZHOU pushed a commit to PKUZHOU/simpler that referenced this pull request Mar 31, 2026
…simulation thread starvation fix (hw-native-sys#143)

Add lock-free heap_tail/last_task_alive advancement in AICPU parallel mode
with ticket-based generation counter (heap_tail_gen) to serialize concurrent
heap_tail writes. Add completed_by_task array for slot-reuse validation in
the orchestrator's early-return path, preventing stale state from recycled
slots.

Fix simulation thread starvation on resource-constrained CI runners (e.g.,
2 CPU cores running 13+ threads) by introducing a platform-level
SPIN_WAIT_HINT() macro for both AICore and AICPU idle polling loops. On
simulation (a2a3sim), expands to CPU pause/yield hint + sched_yield() to
let the OS scheduler give time slices to threads doing real work; on
hardware (a2a3), expands to no-op. Without this, idle threads in tight
spin loops starve the executing core, causing AICPU iteration-based
timeouts (MAX_IDLE_ITERATIONS) before kernels can complete.

Also fix race where scheduler threads could access orch_ready_queue pointers
before Thread 3 finished initializing them (orch_pointers_ready_ guard).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant