Update: modularisation of completion watermark together with a false-sharing alignment fix and batch updates.

[raphael-s-steiner]

This PR is the squashed version of #1536.

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]

📝 Walkthrough

Walkthrough

Completion flag reads and writes are centralized in PTO2SharedMemoryRingHeader, which also owns contiguous watermark advancement. Scheduler, orchestrator, diagnostics, and classification paths now use these helpers, alongside an updated shared-memory layout assertion.

Changes

Completion flag and watermark flow

Layer / File(s)	Summary
Ring header completion API src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h	Aligns task_window_size, adds completion flag read/write helpers and watermark advancement, and updates the layout offset assertion.
Scheduler completion integration src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h	Uses ring-header helpers for fan-in checks, completion publication, and contiguous completed-watermark updates.
Orchestration completion integration src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp, src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp	Routes hidden-allocation completion, stall diagnostics, and orchestration classification through the completion flag helper.
Estimated code review effort: 3 (Moderate)	~20 minutes

Possibly related PRs

• hw-native-sys/simpler#1435: Introduces overlapping polling-completion flags, watermark advancement, and scheduler/orchestrator wiring.
• hw-native-sys/simpler#1536: Refactors completion flag publication and watermark handling around the same ring-header helpers.

Poem

A bunny hops where completion flags glow,
Ring helpers guide the tasks in a row.
Watermarks advance through the night,
Fan-ins find every signal right.
“Hop hooray!” says the rabbit, “all aligned!”

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title accurately summarizes the main changes: watermark modularization, cache-line alignment, and batched updates.
Description check	✅ Passed	The description is clearly related to the code changes and matches the implemented completion-flag and watermark refactor.

---

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.}