Fix: Removed false sharing in Completion Watermark and added batched update

[raphael-s-steiner]

Summary

• Extracts the repeated completion_flags[id & task_window_mask] load/store into two helpers on PTO2SharedMemoryRingHeader — is_completion_flag_set() / set_completion_flag() — and switches all 6 call sites (orchestrator, scheduler hot path, scheduler cold path) to use them instead.
• Replaces the single-step completed_watermark CAS loop (advance by one slot, retry) with update_completed_watermark(), which scans the full contiguous completed prefix locally before issuing one CAS to jump the watermark forward by however many slots are ready — cutting the number of atomic RMW ops on completed_watermark from O(advance distance) to O(contention events).
• Adds alignas(64) to task_window_size, pushing the cold, read-mostly layout metadata onto its own cache line, separate from the hot completed_watermark atomic that's CAS'd on every task completion — removes false sharing between the two.
• Updates the PTO2SharedMemoryRingHeader layout static_assert (task_descriptors_offset: 160 → 216) to match; struct size stays 256 bytes.

Why

completed_watermark is CAS'd on the hot completion path for every task and previously shared a cache line with metadata read on nearly every fanin check, so each watermark bump invalidated that line for concurrent readers. Advancing it one slot at a time (instead of scanning the whole completed prefix) also meant up to N CAS attempts when several tasks completed in a burst.

Testing

Ran TestPagedAttention::Case2 (a2a3 / tensormap_and_ringbuffer runtime) on NPU 3, 20 rounds each, from a clean build of both branches:

Metric	up-main	completion_watermark	Δ
Total round-trip, mean	39.70 ms	36.81 ms	-7.27%
Total round-trip, stdev	2.20 ms	1.01 ms	~2.2x tighter
Scheduler phase, mean	23.09 ms	22.73 ms	-1.55%
Scheduler phase, stdev	0.107 ms	0.059 ms	~1.8x tighter

Both branches pass golden comparison. completion_watermark is faster and meaningfully less jittery — consistent with fewer/cheaper atomic ops on the watermark and the removed false sharing.

[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 Plus

Run ID: 5c8eee95-c6ae-4e52-b7d3-ea9cc7e9650f

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 shared-memory ring gains centralized completion-flag and watermark helpers, with an alignment and layout assertion update. Orchestration and scheduler paths now use these helpers for fanin checks, completion publication, diagnostics, and watermark advancement.

Changes

Polling completion state

Layer / File(s)	Summary
Ring completion contract src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h	The ring header aligns task_window_size, adds completion flag accessors and contiguous watermark advancement, and updates the layout assertion.
Completion path adoption src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp, src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h, src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp	Orchestration and scheduler paths use the ring helpers for completion publication, fanin checks, watermark updates, diagnostics, and completed-task filtering.

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

Possibly related PRs

• hw-native-sys/simpler#1435: Introduces the completion flags and completed watermark later centralized by this change.

Poem

A rabbit taps the ring with care,
“Completion flags now flow in air!”
Watermarks hop from slot to slot,
Fanins find the paths they sought.
thump thump—the tasks are through!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main performance and watermark-update changes.
Description check	✅ Passed	The description is directly related to the PR and accurately summarizes the completion-watermark and layout changes.
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.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h`:
- Around line 116-131: Update update_completed_watermark so a failed
completed_watermark.compare_exchange_strong refreshes the scan state: reset next
from the newly loaded curr_watermark and reload submitted from
fc.current_task_index before rescanning completion flags. Preserve the existing
CAS success loop and completion-flag ordering while ensuring concurrent advances
cannot be overwritten or left behind.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: a112fa00-83c4-490b-87d2-d21623888e13

📥 Commits

Reviewing files that changed from the base of the PR and between a53cc16 and e834f03.

📒 Files selected for processing (4)

• src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp
• src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h
• src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp

[ChaoZheng109]

Reviewed against the merge-base — the change is correct and all three optimizations are effective:

• Helper extraction — all 6 completion_flags[id & task_window_mask] call sites are converted consistently; the only raw indexing left is inside the two helpers themselves.
• Batched watermark — update_completed_watermark() is equivalent to the old per-slot loop (same inner scan condition, monotonic, submitted loaded once) but collapses O(advance distance) CAS attempts into one. The failure path next = std::max(next, curr_watermark) is safe: on a failed CAS the observed value is strictly > expected (watermark is monotonic), and the jump target's prefix has already been verified set.
• False-sharing fix — verified the layout: completed_watermark now owns the 128–191 line while the read-mostly metadata (task_window_mask, pointers) packs into 192–255. task_descriptors_offset = 216 and sizeof == 256 both check out; the struct size is unchanged, the padding just moved off the tail. No a5 / tensormap_and_ringbuffer sibling carries this completed_watermark/completion_flags subsystem, so there's nothing to mirror.

One request before merge: please squash the 6 commits into one, with a message describing the final change rather than the working process — the current history ("micro optimisation", "precommit", "padding", "batched watermark update", "modularised completion flags") is the kind of step-by-step narration the repo's codestyle rule asks us to keep out of commit messages.

[raphael-s-steiner]

Single commit version of this PR: #1550.