Refactor: give each run its own host tensor accessor - #1695
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:
📝 WalkthroughWalkthroughThe change replaces global host tensor mappings with per-run ChangesHost tensor access
Estimated code review effort: 4 (Complex) | ~45 minutes 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 |
ac437d2 to
ba69d43
Compare
Crane-Liu
changed the title
de9d197 to
1434cb3
Compare
Status of this branchI rebased this onto current
Everything is squashed into one commit with CI so far: green. A scoping question worth settling before mergeThe change bundles two independent things, and I think only one of them has to be in it:
(a) is unambiguously right and I'd merge it as-is. (b) I could not find a buyer for. Three things I checked: 1. The host-side bind was dead code. Every reader of So 2. The 3. Pipelining does not produce concurrent orchestration. In So the concurrency property (b) buys has no user I could find, today or under the pipelined design. What's left is readability — real, but the price is 119 files of author-visible churn. The part that I think blocks (b) regardless
plus five hand-written collectives templates ( Those templates still use And the entry ABI break is the silent one. ProposalSplit (b) out and merge (a) alone:
That's ~400 lines in If (b) is wanted later — the concrete trigger would be W1c "direct depth two" genuinely putting two runs in @Crane-Liu happy to do the split myself, or hand the branch back — your call. Either way I'd like to record the three findings above in |
The host views a run stages lived in two file-scope globals in libhost_runtime.so — a region list and the HostApi pointer serving mirror-mode writes. One instance per process, shared by every worker, so two runs staging tensors overwrite each other's regions. The window was bounded by a hand-paired host_tensor_access_reset(api) / reset(nullptr) plus an RAIIScopeGuard, and the mapping each region needed was registered by the caller but released far away in cleanup. HostTensorAccessor replaces both globals with one object per run: - It owns every mapping it registers and releases them in close(), which the destructor also calls, so no exit path can leak one. Registration reserves both tables before it registers, so a throwing push_back cannot strand a mapping it has no record of. - register_device_memory_to_host moves inside add(), which takes the staging buffer as the fallback view. Register and unregister are now balanced per run; previously a run that bound more than once per validate leaked the extra mappings. - A null HostApi admits no region, so a mirrored write can no longer reach a null copy_to_device and write() needs no second check. The runtime carries the accessor in a field past its first two, which are the only ones the orchestration .so's partial PTO2Runtime definition can see. get_tensor_data / set_tensor_data read it from there, so the ops table and every author-facing signature are unchanged and no kernel needs editing. run_host_orchestration no longer binds the host library's own copy of framework_current_runtime. Nothing outside the orchestration .so includes pto_orchestration_api.h, so nothing read it; rt_scope_* and rt_orchestration_done take the runtime as an argument. The .so's own copy is still resolved and bound — its inline rt_submit_* read that one. Tests cover null-api rejection, two regions resolving independently within one accessor, a span starting before a region base, an unregistered write issuing no device copy, and two concurrent runs each staging, reading and closing their own accessor over overlapping device addresses. Co-authored-by: Crane-Liu <c.wliu@outlook.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1434cb3 to
ae36953
Compare
ChaoWao
changed the title
CI triage on
|
2 participants
Summary
The host views a run stages lived in two file-scope globals in
libhost_runtime.so— a region list and theHostApipointer serving mirror-mode writes. One instance per process, shared by every worker, so two runs staging tensors overwrite each other's regions. The window was bounded by a hand-pairedhost_tensor_access_reset(api)/reset(nullptr)plus anRAIIScopeGuard, and the mapping each region needed was registered by the caller but released far away in cleanup.HostTensorAccessorreplaces both globals with one object per run:close(), which the destructor also calls, so no exit path can leak one. Registration reserves both tables before it registers, so a throwingpush_backcannot strand a mapping it has no record of.register_device_memory_to_hostmoves insideadd(), which takes the staging buffer as the fallback view. Register and unregister are now balanced per run — previously a run that bound more than once per validate leaked the extra mappings.HostApiadmits no region, so a mirrored write can no longer reach a nullcopy_to_deviceandwrite()needs no second check.The runtime carries the accessor in a field past its first two — the only ones the orchestration
.so's partialPTO2Runtimedefinition can see.get_tensor_data/set_tensor_dataread it from there, so the ops table and every author-facing signature are unchanged and no kernel needs editing.run_host_orchestrationno longer binds the host library's own copy offramework_current_runtime. Nothing outside the orchestration.soincludespto_orchestration_api.h, so nothing read it;rt_scope_*andrt_orchestration_donetake the runtime as an argument. The.so's own copy is still resolved and bound — its inlinert_submit_*read that one.Scope change from the previous revision
This PR previously also threaded an explicit
OrchestrationContextparameter through every orchestration API (194 files, +2676 −2217). That half is not in this revision. Three findings, each checked against the code rather than reasoned about:The host-side bind was dead. Every reader of
current_runtime()is inpto_orchestration_api.h, and no source underhost/,aicpu/,runtime/,runtime/orchestrator_core/orruntime/shared/#includes it — the only two greps that hit are comments:The
.so-side global is per-callable, not per-process.register_callable_implcreates a unique temp file per callable anddlopens itRTLD_LOCAL; TMR keysorch_so_table_[callable_id]. Two workers get two.somappings and two independent globals.Pipelining does not produce concurrent orchestration. In
[orch N] → [device N]overlapped with[orch N+1],orch Nhas already returned beforeorch N+1starts — orch being shorter than the device phase is exactly what keeps the write/read serialized. TMR runs several orchestrator threads, but the comment on that code records that "All orchestrator threads bind the same rt value".So the global region table was the only state a concurrent run could actually corrupt, and it is what this PR fixes. Separately,
pto_orchestration_api.his a Tier C external contract undercodestyle.mdrule 10 —pypto's orchestration codegen (src/codegen/orchestration/orchestration_codegen.cpp) emitsrt_submit_aic_task/PTO2_SCOPE()/alloc_tensors(against it, and five collectives templates#includeit directly — so changing it needs a coordinated cross-repo change, not a unilateral sweep. If explicit context is wanted later (the concrete trigger would be W1c "direct depth two" genuinely putting two runs inprepareat once), rule 10's incremental doctrine applies: additive overloads plus the pypto change, not one sweep.Deliberately left out
framework_bind_runtimecalls (HBG teardown, TMR bind + teardown) are dead by the same argument, but they live in a different binary and a different file from this change. Separate cleanup.SUBMIT_BY_CLUSTER.mddocumentsrt_submit_task(PTO2Runtime*, Arg*, int32_t), and the TMRRUNTIME_LOGIC.mdexport section documentsaicpu_orchestration_entry(uint64_t*, int). Both are wrong onmaintoday and unrelated to this change, so they are not fixed here.Testing
g++ -fsyntax-onlyagainst each runtime'sbuild_config.pyinclude set) — direct evidence the author-facing API is untouched.clang-formatclean;examples/andtests/st/have zero changed files.New coverage
test_hbg_tensor_access.cppgains the cases the old global-based tests could not express:apiadmits no region — the invariant that letswrite()skip re-checking the copy hook;