[Performance] A2/A3/A5 host_build_graph: skip redundant orchestrator init in arena build

[SergioMartin86]

Human Summary

The run_host_orchestration() function was being called twice. This PR removes the redundant call.

Applies to A2/A3/A5.

AI Summary

PR notes — host_build_graph: skip redundant orchestrator init in arena build

• Branch: hbg-skip-redundant-orch-init (commit a2a42446, off upstream/main cf0fbc06)
• Diff: 6 files, +264 / −12 — the one-function change + its doc comment
  mirrored across a2a3 and a5, plus a concurrent-prepare white-box stress test
  (concurrent_prepare_stress/ + isolated-worker conftest).
• Arch/runtime: host_build_graph on both a2a3 and a5 (see a5 parity)
• Independent of hbg-ready-queue-8192 (different file; either merge order works).

TL;DR

The host bind path initialized the orchestrator twice per run: once against
the device SM while building the prebuilt arena, then again against the host SM
in run_host_orchestration. The first init is dead work — its arena content is
overwritten by the second, and the orchestrator arena block is never uploaded to
the device. Skipping it cuts the arena_build phase of bind by ~46% (the
redundant tensormap memset), trimming host bind wall −13.6% on paged_attention
(8677 → 7494 µs), device time unchanged, golden-clean.

The double-init

Two calls to PTO2OrchestratorState::init_data_from_layout happen per bind:

1. arena_build — runtime_init_data_from_layout (pto_runtime2_init.cpp) calls
   rt->orchestrator.init_data_from_layout(..., sm_dev_base, ...) against the
   device SM while assembling the prebuilt arena image.
2. host_orch — run_host_orchestration (runtime_maker.cpp) calls
   rt->orchestrator.init_data_from_layout(..., host_sm, ...) against the host
   SM, because the orchestrator's submit loop runs on the host and writes task
   descriptors into a host-side SM buffer that is allocated only at that point.
   After the submit loop the image is relocated (host→device pointers) and the
   live prefix is uploaded.

init_data_from_layout does two kinds of work: arena work (the
fanin_seen_epoch memset + tensor_map.init_data_from_layout — the expensive,
SM-independent part) and SM-pointer setup (task_allocator.init, sm_header
— cheap, SM-specific). Both inits memset the same arena regions; only the SM
pointers differ. So the first init's expensive arena work is entirely redone by
the second.

Why the first init is safe to drop

• Single caller. The 8-arg runtime_init_data_from_layout is called from
  exactly one place — the host bind path (runtime_maker.cpp). The AICPU boots
  scheduler-only (per [Performance] HBG: up to 99% host-side overhead reduction #1659) and never calls it, so no device path depends on it.
• Always re-initialized. That sole caller unconditionally calls
  run_host_orchestration right after, which re-inits the orchestrator against
  the host SM. The first init's result never survives.
• Never uploaded. [Performance] HBG: up to 99% host-side overhead reduction #1659 already skips the orchestrator arena block from the
  H2D upload, so whatever the first init wrote there does not reach the device.
• Nothing reads it in between. The intervening runtime_wire_arena_pointers
  only sets pointers (it does not read orchestrator state), and
  rt->orchestrator is already zeroed by the memset(rt, 0, ...) at the top of
  runtime_init_data_from_layout. So leaving it zeroed until
  run_host_orchestration is well-defined.

The change initializes only the scheduler in runtime_init_data_from_layout and
lets run_host_orchestration own the single orchestrator init.

a5 parity

The change is mirrored into the a5 host_build_graph sibling in the same
commit — a2a3 and a5 are near-duplicate trees with no architectural divergence
here. All four safety pillars above were re-verified on a5:

• Single caller — the 8-arg runtime_init_data_from_layout is called from
  exactly one place, the a5 host bind path (src/a5/.../host/runtime_maker.cpp);
  the AICPU never calls it.
• Always re-initialized — a5's run_host_orchestration re-inits the
  orchestrator against the host SM (src/a5/.../host/runtime_maker.cpp), same as
  a2a3.
• Never uploaded — a5 boots the scheduler without the orchestrator arena
  block.
• Nothing reads it in between — a5's runtime_wire_arena_pointers only sets
  pointers.

Files: src/a5/.../runtime/shared/pto_runtime2_init.cpp (the orchestrator init
removed from the 8-arg runtime_init_data_from_layout, replaced by the same
explanatory comment) and src/a5/.../runtime/pto_runtime2.h (matching doc
comment). Verification: a5 (and a2a3) host libs compile + link clean; all
pre-commit hooks pass on all six files. a5 is compile-verified only — not
perf-measured, as the dev box is a2a3 silicon (the numbers below are a2a3). The
concurrent_prepare_stress test lives under tests/st/a2a3/ and exercises the
shared bind/prepare logic; it is not duplicated per-arch.

Going forward, HBG optimizations land in both arches in one commit unless there
is a genuine arch-level difference.

Measurement

a2a3 onboard, paged_attention, median of 50 rounds, from the [STRACE] bind span
tree (default capacity 65536; this PR does not change the queue size).

Phase	before	after	Δ
arena_build	2564.8	1372.6	−1192 µs (−46.5%)
host_orch (incl. orch_reinit)	~1610	~1708	unchanged (the surviving init)
arena_upload	2676.1	2693.3	unchanged
host wall (simpler_run)	8677.1	7493.6	−1184 µs (−13.6%)
device_wall	69.8	71.1	unchanged

orch_reinit (the host-SM init the submit loop actually needs) is untouched, as
expected — this PR removes only the duplicate. Device time is flat.

Stacked on hbg-ready-queue-8192 the two levers compound: measured together on
the 8192 build, arena_build fell to ~186 µs and host wall to ~3790 µs (−56% vs
upstream). Standalone off main the effect is the −13.6% above.

Correctness

Full host_build_graph golden suite (pytest tests/st/a2a3/host_build_graph/,
a2a3 onboard): 22 passed, 2 skipped, 2 failed. Both failures are
pre-existing and unrelated, verified by rebuilding the runtime source at
cf0fbc06 baseline (orch init present) with a matching binding — both reproduce
identically:

• run_stream_reuse::test_depth_two_slots_own_separate_resources — "a served
  bank is uncommitted: …, 0x0" (arena-bank-commit path; already tracked).
• worker_async_fifo::TestWorkerAsyncWholeRunFifo::test_run — 507901
  finalize/hdc-disconnect (same native-run pipeline-slot / arena-bank family).

The diff touches neither path (ChipWorker::arena_bank_gm_heap_base / native-run
finalize are untouched), and the pass/fail set is identical with and without the
change.

Concurrent-prepare stress (added by this PR)

Because the orchestrator init is per-run scratch, the one place worth exercising
is concurrent prepare — a successor prepared (bound) into bank B while a
predecessor is still executing in bank A (overlaps_active_run). Ordinary
blocking run() never hits it, and the two existing tests that do
(run_stream_reuse, worker_async_fifo) are the pre-existing-broken ones, so
they cannot confirm safety.

concurrent_prepare_stress/ adds a white-box test that drives a 2-deep pipeline
over both banks through the canonical direct-chip lane
(_submit_chip_run_direct, the exact path production L2 Worker.run uses):
40 iterations, distinct per-iteration inputs, each result golden-checked, so any
cross-bank orchestrator interference would fail. It passes identically on this
change and a baseline build (orch init present) — confirming the test is sound
and the change is safe under concurrent prepare. This is safe by construction:
each run has its own arena bank → its own PTO2Runtime/orchestrator; the
orchestrator block is never uploaded to the device; and nothing reads it between
the (removed) init and the surviving host-SM init.

Scope & follow-ups

• a2a3 + a5 host_build_graph. The same double-init existed in a5 HBG and is
  fixed in this same commit (see a5 parity); a5 is compile-verified,
  perf-measured on a2a3.
• Next lever: orch_reinit (~1.3 ms) is now the largest single bind
  sub-phase — the surviving orchestrator init (tensormap memset against host
  SM). Reducing it (bounded tensormap, or a persistent-zeroed pool that the
  submit loop cleans up) is the logical successor, and harder.

Reproduce

python tests/st/a2a3/host_build_graph/paged_attention/test_paged_attention.py \
    -p a2a3 -d <dev> --rounds 50 --skip-golden
python -m simpler_setup.tools.strace_timing <captured-stderr> --tree   # Host wall + bind sub-phases

cmake --build build/cache/a2a3/onboard/host_build_graph/host  -j"$(nproc)"
cmake --build build/cache/a2a3/onboard/host_build_graph/aicpu -j"$(nproc)"

# a5 compile-verify (no a5 silicon on this box, so build-only):
cmake --build build/cache/a5/onboard/host_build_graph/host -j"$(nproc)"

# Concurrent-prepare overlap stress:
python -m pytest tests/st/a2a3/host_build_graph/concurrent_prepare_stress/ \
    --platform a2a3 --device <dev>

[coderabbitai]

📝 Walkthrough

Walkthrough

The runtime layout setup no longer initializes the orchestrator or handles that initialization failure. The orchestrator remains zero-initialized until run_host_orchestration, while scheduler initialization and later failure handling remain unchanged.

Changes

Runtime initialization

Layer / File(s)	Summary
Defer orchestrator initialization src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp	runtime_init_data_from_layout removes the orchestrator initialization call and its failure return path. Host orchestration initializes the orchestrator later.

Estimated code review effort: 2 (Simple) | ~10 minutes

Poem

I’m a rabbit with a runtime plan,
The orchestrator waits where it can.
Zeroed at the start,
Then orchestration takes part,
While the scheduler follows its span.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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 identifies the performance optimization and the affected host_build_graph targets.
Description check	✅ Passed	The description directly explains the redundant initialization removal, performance impact, scope, and validation results.

---

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

[ChaoZheng109]

The optimization itself is correct and well-covered — I verified the four safety pillars (single caller of the 8-arg runtime_init_data_from_layout; unconditional orchestrator re-init in run_host_orchestration against the host SM; the orchestrator arena block is never uploaded to the device; and runtime_wire_arena_pointers between the two only sets pointers and does not read orchestrator state). The concurrent_prepare_stress test is sound too.

One ask: please apply the same change to a5 host_build_graph in this PR. a2a3 and a5 are near-duplicate trees, and we keep optimizations in sync across both architectures unless there is a genuine architectural divergence.

There is no divergence here — I checked, a5 has the identical double-init pattern:

• src/a5/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp still calls rt->orchestrator.init_data_from_layout(...) inside runtime_init_data_from_layout (the redundant init).
• src/a5/runtime/host_build_graph/host/runtime_maker.cpp:463 re-initializes the orchestrator against the host SM inside run_host_orchestration (the surviving init), exactly as a2a3 does.

So the same removal + the matching doc-comment update in src/a5/runtime/host_build_graph/runtime/pto_runtime2.h should land here too. Going forward, unless there's a real arch-level difference, please bundle the a5 sibling in the same change.

[SergioMartin86]

The optimization itself is correct and well-covered — I verified the four safety pillars (single caller of the 8-arg runtime_init_data_from_layout; unconditional orchestrator re-init in run_host_orchestration against the host SM; the orchestrator arena block is never uploaded to the device; and runtime_wire_arena_pointers between the two only sets pointers and does not read orchestrator state). The concurrent_prepare_stress test is sound too.

One ask: please apply the same change to a5 host_build_graph in this PR. a2a3 and a5 are near-duplicate trees, and we keep optimizations in sync across both architectures unless there is a genuine architectural divergence.

There is no divergence here — I checked, a5 has the identical double-init pattern:

• src/a5/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp still calls rt->orchestrator.init_data_from_layout(...) inside runtime_init_data_from_layout (the redundant init).
• src/a5/runtime/host_build_graph/host/runtime_maker.cpp:463 re-initializes the orchestrator against the host SM inside run_host_orchestration (the surviving init), exactly as a2a3 does.

So the same removal + the matching doc-comment update in src/a5/runtime/host_build_graph/runtime/pto_runtime2.h should land here too. Going forward, unless there's a real arch-level difference, please bundle the a5 sibling in the same change.

Done — applied the identical change to a5 host_build_graph in this PR (same commit, pto_runtime2_init.cpp + pto_runtime2.h). Agreed there's no divergence; I re-verified all four safety pillars on a5:

• Single caller — the 8-arg runtime_init_data_from_layout is called only from the a5 host bind path (src/a5/.../host/runtime_maker.cpp); the AICPU never calls it.
• Surviving init — a5's run_host_orchestration re-inits the orchestrator against the host SM (runtime_maker.cpp:463), exactly as a2a3.
• Never uploaded — a5 boots scheduler-only; the orchestrator arena block isn't shipped.
• runtime_wire_arena_pointers between the two only sets pointers, doesn't read orchestrator state.

The removed init and the two doc-comment updates match the a2a3 change verbatim. a5 is compile-verified only (host lib links clean; the dev box is a2a3 silicon, so the perf numbers stay a2a3); all pre-commit hooks pass on all six files. The concurrent_prepare_stress test stays under tests/st/a2a3/ since it exercises the shared bind/prepare logic rather than anything arch-specific.

Kept it a single commit; force-pushed. Will bundle the a5 sibling in the same change for HBG work going forward.