Refactor: isolate A2A3 run stream sets per pipeline slot - #1464
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughA2A3 ChangesRun stream-set reuse
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PythonTest
participant Worker
participant ChipWorker
participant DeviceRunner
participant AICPU
participant AICore
PythonTest->>Worker: run repeated workload
Worker->>ChipWorker: read run_stream_set_create_count
ChipWorker->>DeviceRunner: query creation count
DeviceRunner->>AICPU: launch on slot stream
DeviceRunner->>AICore: launch on slot stream
DeviceRunner-->>ChipWorker: return creation count
ChipWorker-->>Worker: expose read-only property
Worker-->>PythonTest: assert count equals 1
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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. Comment |
aabef49 to
4ef67d6
Compare
4ef67d6 to
e56d50b
Compare
ChaoWao
changed the title
Revision: stream sets are per slot, not per runPushed Blocking: per-run create/destroy cost +22% host time
Measured on a2a3,
The fix is to keep the slot shape and drop the per-run lifetime: a set is created on first use and reused. That is what Falls out of the same change
There was also a narrower hazard: on the Other
Validation61/61 C++, 824 passed / 2 skipped Python, Coverage note: all three files are onboard-only, so |
e56d50b to
47c70da
Compare
Revision: a regression barrier for the reuse invariant, and a question about slot 1Pushed The reuse invariant now has a testI said in the last revision that the create-once invariant had no observable and therefore no test. That was true but not a good place to stop — the +22% regression this PR fixes would have landed green again. Streams are opaque handles, the So the runner reports how many sets it has built, through the same C API →
Slot 1 is unreachable through the whole seriesWorth raising before this merges. I checked the That is consistent with what #1467 actually enables: The second slot only starts paying when a run's So One thing the slot split does not buy, in case it is load-bearing anywhere downstream: separate streams give ordering and independent completion, not AICore exclusivity. Two runs on two slots would still contend for the same cores; serializing that is Gate B's job, not the stream layer's. Validation61/61 C++, 824 passed / 2 skipped Python, |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py (1)
39-65: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate intentional mutable class configuration.
Ruff flags both mutable class attributes with RUF012. Mark them as
ClassVar(or copy them per test if the framework mutates nested values) to make the sharing semantics explicit.Suggested fix
import pytest import torch +from typing import ClassVar from simpler.task_interface import ArgDirection as D @@ - CALLABLE = { + CALLABLE: ClassVar[dict] = { @@ - CASES = [ + CASES: ClassVar[list] = [Also applies to: 67-74
🤖 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 `@tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py` around lines 39 - 65, Annotate the mutable class-level configuration dictionaries containing CALLABLE and the corresponding configuration at the other flagged location with ClassVar, preserving their shared test configuration semantics. Import ClassVar from typing if needed, and use the annotation consistently for both attributes.Source: Linters/SAST tools
🤖 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 `@python/simpler/worker.py`:
- Around line 5960-5965: Update the docstring for run_stream_set_create_count to
state that it reports the number of run stream sets created across pipeline
slots, rather than promising a value of 1; retain the existing 0 behavior for
non-L2 workers and platforms using the persistent bootstrap stream pair.
---
Nitpick comments:
In `@tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py`:
- Around line 39-65: Annotate the mutable class-level configuration dictionaries
containing CALLABLE and the corresponding configuration at the other flagged
location with ClassVar, preserving their shared test configuration semantics.
Import ClassVar from typing if needed, and use the annotation consistently for
both attributes.
🪄 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: e3f22b4f-7a3c-485f-ab81-90655d9d74ea
📒 Files selected for processing (13)
python/bindings/task_interface.cpppython/simpler/task_interface.pypython/simpler/worker.pysrc/a2a3/platform/onboard/host/device_runner.cppsrc/a2a3/platform/onboard/host/device_runner.hsrc/common/platform/onboard/host/c_api_shared.cppsrc/common/platform/onboard/host/device_runner_base.cppsrc/common/platform/onboard/host/device_runner_base.hsrc/common/platform/sim/host/c_api_shared.cppsrc/common/worker/chip_worker.cppsrc/common/worker/chip_worker.hsrc/common/worker/pto_runtime_c_api.htests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py
47c70da to
b70f386
Compare
Give the A2/A3 device runner its own AICPU/AICore stream set per pipeline slot and submit every run on the selected set, so two in-flight runs would never share a stream. The persistent pair created by ensure_device_initialized() stays, now reserved for bootstrap and control work: the Init and RegisterCallable payloads and the block-dim query. A stream set carries no per-run content, so a slot's set is created on first use and reused for every later run on that slot. The pipeline contract classifies both streams as EXEC_HANDLE, which asks for one instance per in-flight run — that is one set per slot, not one per run invocation — so kRunStreamSetCount is PTO_PIPELINE_MAX_DEPTH and only slot 0 is selected while the contract is K=1. destroy_run_stream_sets() in finalize() is the sole release point, running while RTS is live and before the device reset, which is the window finalize_common() uses for the bootstrap pair. It is idempotent, so the destructor's finalize() call is a safe backstop. As there, no pre-destroy sync: rtStreamDestroy is the supported teardown for a stream an op-timeout left in the error state. sync_run_streams() is split so the wait can name a stream pair; DeviceRunnerBase keeps the no-argument form for A5, whose runs still use the persistent pair. Report the number of sets a runner has built through the C API and ChipWorker, alongside the dlopen counters that already serve this purpose, so the reuse invariant is assertable. A new onboard scene test runs one worker four times and requires exactly one set; rebuilding per run makes it report four. Rebuilding the set per run costs two rtStreamCreate plus two rtStreamDestroy, ~300 us each, on the synchronous host path around KernelLaunch. Measured on a2a3 silicon with vector_example at 40 rounds, three alternating repetitions of each binary under one device lock, host time per round (p50 of per-repetition p50, warmup rounds dropped): main 10316.8 us, per-run rebuild 12607.6 us (+22.2%), reuse 10481.1 us (+1.6%, inside main's own 2.7% repetition spread). Device time is unchanged.
…tree An editable install pins the compiled extension at install time (`editable.rebuild = false`) while `python/simpler/*.py` is read live. Switching branches or rebasing therefore moves the Python out from under a fixed binary, nothing rebuilds, and a changed struct layout makes attributes read as 0 with no error anywhere. That is not hypothetical. A worktree installed before hw-native-sys#1309 (which dropped `block_dim` from `CallConfig`) and then branched onto a main containing it read `aicpu_thread_num` as 0, and the whole onboard L2 suite failed with `launch_aicpu_num (0) must be in range [1, 4]` — a plausible-looking runtime rejection that reads as a product bug. It cost a full bisect, and an A/B against main "reproduced" it because both arms shared the same stale binary. The same skew hit again a few hours later as `AttributeError: run_stream_set_create_count` after a rebase across hw-native-sys#1464. `_task_interface` now records the commit it was built from, and `simpler.task_interface` compares it against the working tree at import, raising with the one command that fixes it. Loud beats silent here: the alternative is not an error, it is wrong values. A warning would also have been the wrong choice — pytest relegates import-time warnings to its end-of-run summary, and in the failure above the same warning would have appeared in *both* arms of the A/B and been dismissed as noise. A *missing* stamp raises too, rather than being treated as "cannot tell". The attribute is absent only on an extension compiled before it existed, which in a checkout new enough to run the check is by definition a different revision — and is the state of every already-installed worktree the day this lands. Keyed on git HEAD, matching `RuntimeBuilder._build_cache_stamp` — mtimes do not survive a branch switch, which is why that stamp is git-based too. Deliberately the whole HEAD rather than a subset: excluding paths means maintaining a second "what cannot affect the build" list beside the one CI already keeps, and the guard exists precisely so correctness does not rest on such judgments. It costs little — 120 of the last 200 commits touch the ABI surface anyway, so most of the reinstalls it forces were owed regardless and merely became visible; the rest cost one ~20 s reinstall. Inert outside a source tree: a wheel has no `.git` to compare against, and a build made without git carries an empty stamp. The rebuild table gained the trigger it was missing and no longer contradicts itself. It was framed as "what did you change", so it had no row for the case where you changed nothing and the tree moved — the one that bites, and the one where verifying that `import simpler` resolves into your worktree gives false confidence. Its "no rebuild needed" rows now say what is actually true: nothing to recompile, but the guard keys on HEAD, so a commit still needs a reinstall before the next import. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…tree An editable install pins the compiled extension at install time (`editable.rebuild = false`) while `python/simpler/*.py` is read live. Switching branches or rebasing therefore moves the Python out from under a fixed binary, nothing rebuilds, and a changed struct layout makes attributes read as 0 with no error anywhere. That is not hypothetical. A worktree installed before hw-native-sys#1309 (which dropped `block_dim` from `CallConfig`) and then branched onto a main containing it read `aicpu_thread_num` as 0, and the whole onboard L2 suite failed with `launch_aicpu_num (0) must be in range [1, 4]` — a plausible-looking runtime rejection that reads as a product bug. It cost a full bisect, and an A/B against main "reproduced" it because both arms shared the same stale binary. The same skew hit again a few hours later as `AttributeError: run_stream_set_create_count` after a rebase across hw-native-sys#1464. `_task_interface` now records the commit it was built from, and `simpler.task_interface` compares it against the working tree at import, raising with the one command that fixes it. Loud beats silent here: the alternative is not an error, it is wrong values. A warning would also have been the wrong choice — pytest relegates import-time warnings to its end-of-run summary, and in the failure above the same warning would have appeared in *both* arms of the A/B and been dismissed as noise. A *missing* stamp raises too, rather than being treated as "cannot tell". The attribute is absent only on an extension compiled before it existed, which in a checkout new enough to run the check is by definition a different revision — and is the state of every already-installed worktree the day this lands. Keyed on git HEAD, matching `RuntimeBuilder._build_cache_stamp` — mtimes do not survive a branch switch, which is why that stamp is git-based too. Deliberately the whole HEAD rather than a subset: excluding paths means maintaining a second "what cannot affect the build" list beside the one CI already keeps, and the guard exists precisely so correctness does not rest on such judgments. It costs little — 120 of the last 200 commits touch the ABI surface anyway, so most of the reinstalls it forces were owed regardless and merely became visible; the rest cost one ~20 s reinstall. Inert outside a source tree: a wheel has no `.git` to compare against, and a build made without git carries an empty stamp. The rebuild table gained the trigger it was missing and no longer contradicts itself. It was framed as "what did you change", so it had no row for the case where you changed nothing and the tree moved — the one that bites, and the one where verifying that `import simpler` resolves into your worktree gives false confidence. Its "no rebuild needed" rows now say what is actually true: nothing to recompile, but the guard keys on HEAD, so a commit still needs a reinstall before the next import. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…tree (#1523) An editable install pins the compiled extension at install time (`editable.rebuild = false`) while `python/simpler/*.py` is read live. Switching branches or rebasing therefore moves the Python out from under a fixed binary, nothing rebuilds, and a changed struct layout makes attributes read as 0 with no error anywhere. That is not hypothetical. A worktree installed before #1309 (which dropped `block_dim` from `CallConfig`) and then branched onto a main containing it read `aicpu_thread_num` as 0, and the whole onboard L2 suite failed with `launch_aicpu_num (0) must be in range [1, 4]` — a plausible-looking runtime rejection that reads as a product bug. It cost a full bisect, and an A/B against main "reproduced" it because both arms shared the same stale binary. The same skew hit again a few hours later as `AttributeError: run_stream_set_create_count` after a rebase across #1464. `_task_interface` now records the commit it was built from, and `simpler.task_interface` compares it against the working tree at import, raising with the one command that fixes it. Loud beats silent here: the alternative is not an error, it is wrong values. A warning would also have been the wrong choice — pytest relegates import-time warnings to its end-of-run summary, and in the failure above the same warning would have appeared in *both* arms of the A/B and been dismissed as noise. A *missing* stamp raises too, rather than being treated as "cannot tell". The attribute is absent only on an extension compiled before it existed, which in a checkout new enough to run the check is by definition a different revision — and is the state of every already-installed worktree the day this lands. Keyed on git HEAD, matching `RuntimeBuilder._build_cache_stamp` — mtimes do not survive a branch switch, which is why that stamp is git-based too. Deliberately the whole HEAD rather than a subset: excluding paths means maintaining a second "what cannot affect the build" list beside the one CI already keeps, and the guard exists precisely so correctness does not rest on such judgments. It costs little — 120 of the last 200 commits touch the ABI surface anyway, so most of the reinstalls it forces were owed regardless and merely became visible; the rest cost one ~20 s reinstall. Inert outside a source tree: a wheel has no `.git` to compare against, and a build made without git carries an empty stamp. The rebuild table gained the trigger it was missing and no longer contradicts itself. It was framed as "what did you change", so it had no row for the case where you changed nothing and the tree moved — the one that bites, and the one where verifying that `import simpler` resolves into your worktree gives false confidence. Its "no rebuild needed" rows now say what is actually true: nothing to recompile, but the guard keys on HEAD, so a commit still needs a reinstall before the next import.
2 participants
Depends on #1463 (merged).
Summary
InitandRegisterCallablepayloads and the block-dim queryPTO_PIPELINE_MAX_DEPTH; K=1 selects slot 0 onlydestroy_run_stream_sets()infinalize()is the sole release pointsync_run_streams()so the wait can name a pair; A5 keeps the no-argument form and the persistent pairWhy a set per slot, not per run
The pipeline contract classifies both streams as
EXEC_HANDLE, which asks forone instance per in-flight run — that is one set per slot, i.e.
pipeline_depthsets, not one set perrun()invocation. A stream set carriesno per-run content, so rebuilding it every run buys nothing.
It also costs a great deal.
rtStreamCreate/rtStreamDestroyare ~300 us eachon this silicon, and both sit on the synchronous host path around
KernelLaunch. Measured withvector_exampleat 40 rounds, three alternatingrepetitions of each binary under a single device lock (only
libhost_runtime.soswapped between repetitions), host time per round as thep50 of the per-repetition p50 with warmup rounds dropped:
mainmainmain's own three repetitions span 2.7%, so the remaining +1.6% is inside therun-to-run spread. Device time is unchanged (~250 us) in all three.
Reusing the set also removes the teardown questions the per-run shape raised: a
stream destroy can no longer fail an otherwise successful run, and there is no
window in which a run's streams are destroyed while a launch error left an
AICore kernel spinning on them.
Teardown
run_stream_sets_is this subclass's own RTS-owning member, sodestroy_run_stream_sets()runs infinalize()while RTS is live and beforethe device reset — the window
finalize_common()uses for the bootstrap pair,per the teardown invariant that file documents (#1197). It is idempotent, so
the destructor's
finalize()call is a safe backstop. As there, no pre-destroysync:
rtStreamDestroyis the supported teardown for a stream an op-timeoutleft in the error state.
Observability and the reuse test
Reuse is the invariant this PR turns on, and nothing outside
DeviceRunnercould see it: streams are opaque handles, the
run stream set createdline sitsbelow the default host log threshold, and a leak needs ~990 worker lifetimes to
exhaust a device's 1981 streams. So the runner now reports how many sets it has
built, through the same C API →
ChipWorker→ nanobind path theaicpu_dlopen_count/host_dlopen_countcounters already use for exactly thiskind of assertion. It is diagnostic only and gates nothing; simulation and A5
report 0.
tests/st/a2a3/host_build_graph/run_stream_reuseruns one worker four times andrequires exactly one set. Reinstating the per-run create/destroy makes it report
four, so the regression this PR fixes can no longer land green.
Scope
K=1 throughout.
TASK_ACCEPTED, asynchronous flight, Gate A/B, K=2 and arenabanking are later changes.
run()still callslaunch_runandreap_runback-to-back.
Validation
Rebased onto merged
main; the four superseded ancestor commits are dropped,so this is a single commit of 13 files, +307/-15.
a2a3sim, fullexamples tests/st: 108 passed, 0 faileda2a3onboard, the exact CI command (examples tests/stminus the two SDMAdemos, 4 devices under
task-submit): 115 passed, 1 skipped, 0 failedback it fails with
expected 1 run stream set for 4 runs, got 4a2a3onboard SDMA step (prefetch_async_demo,sdma_async_completion_demo, 2 devices): 2 passedprepared_callablecases inthe onboard suite; the code is onboard-only, so
st-onboard-a2a3is its onlyCI coverage