Add: polling completion scheduler for host_build_graph (a2a3)

[ChaoWao]

Summary

Ports the polling completion scheduler (introduced for tensormap_and_ringbuffer in #1137) into the host_build_graph (hbg) runtime, replacing hbg's wiring dependency subsystem. This is Milestone 1 (core polling; predicated early-dispatch is stubbed, sync_start preserved).

Wiring removed: dep_pool / PTO2FaninPool / PTO2DepListPool, fanin_refcount / fanout_refcount / fanout_head / fanout_lock, the two-phase deferred producer-release, and most of relocate_host_orch_image.

Polling in: per-slot completion_flags byte + monotonic completed_watermark + intrusive wake_list; position-independent inline integer fanin (fanin_local_ids); a device boot initial-classify pass.

Why it's a good fit for hbg

hbg is single-ring (PTO2_MAX_RING_DEPTH==1), whole-graph-resident, and host-orchestrated (host builds the whole graph → relocates → H2Ds; AICPU boots scheduler-only). Polling makes the dependency representation position-independent integers, which collapses the fragile host→device pointer relocation (relocate_host_orch_image) to just task/payload — the single-ring model also removes the cross-ring reclaim gap entirely.

Only the dependency subsystem is swapped; dispatch / core-management / sync_start / async-mailbox are preserved.

Three host-orch-specific correctness points

• task_state is kept — the host still polls it (wait_for_tensor_ready, allocator deadlock detector, cold-path dump); on_mixed_task_complete sets completion_flags (device readiness) and mirrors task_state (host).
• completed_watermark is load-bearing — the host wait_for_consumers now blocks on watermark >= producer.last_consumer_local_id (was fanout_refcount == fanout_count).
• Hidden-alloc producers completed inline on the host pre-set completion_flags in the H2D image, so device consumers don't register_wake on a producer that never runs on device.

Validation

• Builds clean: onboard (host + aicpu + aicore) and sim (host + aicpu).
• Onboard pytest tests/st/a2a3/host_build_graph → 10 passed, 1 skipped (matmul, bgemm, vector_example, dump_args, paged_attention [multi-producer fanin + wake-list drain], prepared_callable [hidden-alloc]).
• Sim st matmul + paged_attention pass.
• No cpput exposure (cpput builds against tensormap_and_ringbuffer).

Notes for review

• The design + the exact wiring→polling cascade map + the 64-byte slot-state layout are in src/a2a3/runtime/host_build_graph/docs/polling-port-plan.md (this PR).
• History is staged one coherent chunk per commit (Stages 1–7); intended for squash-merge — Stages 2–6 don't compile in isolation by design (the 64-byte slot-state change is necessarily big-bang), and Stage 7 (device boot classify) closes the window.
• Milestone 2 (a consumer-pull publish_flags early-dispatch that also fixes tmr's predicated_dispatch) is a follow-up.

🤖 Generated with Claude Code

[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: c9aefe95-0140-4420-a875-22cff8df2f4b

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 runtime now uses inline producer local IDs, shared-memory completion flags, wake lists, and a completed watermark for polling-based readiness and reclamation. Dependency pools, fanout wiring, and deferred producer-release paths are removed across orchestration, scheduling, memory layout, and host runtime preparation.

Changes

Polling scheduler port

Layer / File(s)	Summary
Polling data contracts and runtime layout src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h, runtime/pto_shared_memory.h, runtime/pto_runtime2.h, runtime/shared/*, runtime/orchestrator_core/pto_ring_buffer.*	Task payloads store capped inline fanin IDs; slot state removes fanout and lifecycle fields; shared memory adds completion flags and a completed watermark; dependency-pool arena structures are removed.
Inline fanin submission and reclamation runtime/orchestrator_core/pto_orchestrator.cpp, runtime/orchestrator_core/pto_runtime2.cpp	Fanin construction writes producer local IDs, updates consumer reclaim gates, and uses completed-watermark polling; hidden allocation tasks seed device-visible completion flags.
Boot classification and polling completion runtime/scheduler/pto_scheduler.*, runtime/scheduler/scheduler_cold_path.cpp	Boot classification reads completion flags and routes tasks to ready queues or producer wake lists; task completion publishes flags, drains wake lists, and advances the watermark.
Inline completion and dispatch integration runtime/pto_async_wait.h, runtime/scheduler/scheduler_completion.cpp, runtime/scheduler/scheduler_context.h, runtime/scheduler/scheduler_dispatch.cpp	Deferred producer-release storage, batching, and cleanup paths are removed from async wait, completion, and dispatch flows.
Host wiring and staged implementation plan host/runtime_maker.cpp, docs/polling-port-plan.md	Ring configuration and relocation omit dependency-pool data, while the port plan records the slot layout, staged build sequence, polling cascade, and Milestone-2 hook points.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related PRs

• hw-native-sys/simpler#1415: Both changes modify PTO2TaskSlotState lifecycle and readiness bookkeeping.

Poem

I’m a bunny with flags in a row,
Watching fanins complete as they go.
Wake lists hop, watermarks rise,
Pools vanish beneath polling skies.
The scheduler thumps: “Ready to run!”
Milestone one is begun!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 46.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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 clearly names the main change: adding the polling completion scheduler for host_build_graph.
Description check	✅ Passed	The description matches the PR's main changes, milestone, and validation details.

---

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]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp (1)

359-418: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Stale doc header contradicts the narrowed relocation.

The block comment immediately above relocate_host_orch_image (and the file namespace comment) still describes the pre-polling model — relocating fanin_inline_slot_states[], fanout_head, the dep_pool/dep-entry/ready-queue web across two deltas, and a PTO2_FANIN_INLINE_CAP spill-fatal path. The implementation now relocates only ss->task/ss->payload, and the new inline comment at Line 409 already states this. Please trim the header so the two don't disagree and mislead future maintainers.

🤖 Prompt for 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.

In `@src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp` around lines 359 -
418, Update the documentation immediately above relocate_host_orch_image and the
related file namespace comment to describe only the current relocation behavior:
relocating each slot’s task and payload pointers across the SM or arena windows.
Remove references to fanin_inline_slot_states, fanout_head, dep_pool/dep-entry,
ready queues, dual-delta traversal, and PTO2_FANIN_INLINE_CAP spill handling,
while leaving the implementation unchanged.

🤖 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/scheduler/pto_scheduler.h`:
- Around line 577-586: Update the watermark-advancement loop in the surrounding
completion handler to remove the my_id-based upper bound and continue folding
consecutive set completion_flags entries until the first unset flag. Preserve
the existing atomic compare_exchange_weak progression and rely on the zeroed
unsubmitted slots to terminate traversal without wrapping.

---

Nitpick comments:
In `@src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp`:
- Around line 359-418: Update the documentation immediately above
relocate_host_orch_image and the related file namespace comment to describe only
the current relocation behavior: relocating each slot’s task and payload
pointers across the SM or arena windows. Remove references to
fanin_inline_slot_states, fanout_head, dep_pool/dep-entry, ready queues,
dual-delta traversal, and PTO2_FANIN_INLINE_CAP spill handling, while leaving
the implementation unchanged.

🪄 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

Run ID: 1f142d85-4728-4ea1-bf04-40de02595875

📥 Commits

Reviewing files that changed from the base of the PR and between a7a29ee and 26c8d3a.

📒 Files selected for processing (19)

• src/a2a3/runtime/host_build_graph/docs/polling-port-plan.md
• src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp
• src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp
• src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_ring_buffer.cpp
• src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp
• src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h
• src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h
• src/a2a3/runtime/host_build_graph/runtime/pto_ring_buffer.h
• src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h
• src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h
• src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h
• src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp
• src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp
• src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp
• src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp

💤 Files with no reviewable changes (1)

• src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp