Refactor: bind platform host callbacks per run instead of per thread - #1685
Refactor: bind platform host callbacks per run instead of per thread#1685ChaoWao wants to merge 1 commit into
Conversation
|
Warning Review limit reached
Next review available in: 3 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (35)
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 |
7e3ecc2 to
6ea6f15
Compare
Split the platform HostApi into an immutable HostApiFunctions table plus a per-run HostApi value object that binds it to one runner and one run's pipeline slot / arena bank. Carry the run's resource selection and trace identity across the C ABI in a required NativeRunDescriptor, and delete the pthread-TLS resource-selection mechanism (NativeRunThreadSelection, both pthread keys, capture/restore) and the four setter exports it made necessary. Launch acceptance moves from a TLS setter into the per-run launch signal. Bump PTO_PIPELINE_CONTRACT_ABI_VERSION 1 -> 2 so a stale host_runtime.so (whose simpler_run/prepare_run still take the old flat argument list) is rejected at load rather than crashing at the first call. PR-A of the hw-native-sys#1650 three-way split; PR-B (direct-L2 async lane) and PR-C (HostTensorAccessScope RAII) follow. Co-Authored-By: Claude <noreply@anthropic.com>
6ea6f15 to
2b6aabe
Compare
CI status — two red checks are self-hosted-runner infra, not this changeEverything that actually executes this change's code is green:
|
Summary
Splits the platform
HostApifrom a process-global table of context-freefunction pointers into an immutable function table (
HostApiFunctions) plus asmall per-run value object (
HostApi) that binds that table to one runner andone run's pipeline slot / arena bank. Carries the run's resource selection and
trace identity across the C ABI in one required
NativeRunDescriptorinsteadof four TLS setter exports, and deletes the pthread-TLS resource-selection
mechanism (
NativeRunThreadSelection, both pthread keys, capture/restore,try_run_selection) that the global table made necessary.This is PR-A of the three-way split of #1650 ("Add: bounded asynchronous
native run lane"). PR-B (the direct-L2 two-slot async lane in
worker.py/the binding) and PR-C (the
HostTensorAccessScopeRAII conversion) follow;both depend only on
HostApiexisting, not on each other. Implementedindependently of #1650's branch; comment there links this PR so it can rebase
the remainder onto it.
Representation change only — no run behaviour moves.
simpler_initkeeps itsflat argument list.
What changes
HostApi→HostApiFunctions(onestatic consttable per backend, everyfunction takes
void *runner_ctxplus an explicitpipeline_slot/arena_bankwhere it indexes per-run storage) +HostApivalue object{runner_ctx, pipeline_slot, arena_bank, functions}whose member functionsforward to the bound table. Constructed per run, stored in
NativeRunState,passed by
const HostApi *into the runtime impls.NativeRunDescriptor {pipeline_slot, arena_bank, run_id, generation, dispatch_id, run_epoch}— copied into prepared state.pipeline_slot/arena_bankare load-bearing (<PTO_PIPELINE_MAX_DEPTH); the rest isdiagnostic identity, zero for a synchronous depth-one run.
simpler_prepare_run/
simpler_runtake a requiredconst NativeRunDescriptor *.NativeRunLaunchSignalgainsconfigure_acceptance/publish_acceptance;simpler_launch_run/simpler_runtakeaccepted_state/accepted_valueand the runner publishes at the real kernel-launch marker.
NativeRunThreadSelection,g_runner_key,g_run_selection_key,try_run_selection,run_selection,capture_/restore_native_run_thread_selection,the two RAII selection guards, and the four setter exports
(
select_pipeline_slot_ctx,select_arena_bank_ctx,set_native_run_identity_ctx,set_task_accepted_state_ctx).host_tensor_access_resettakesconst HostApi *instead of a rawcopy_to_devicepointer (a minimal bridge; the fullHostTensorAccessScopeRAII conversion is PR-C).
PTO_PIPELINE_CONTRACT_ABI_VERSIONbumped 1 → 2 (see ABI note).Per-thread state: why deleting TLS is safe while the per-run executor thread still exists
The plan orders this as B6c ("parameterize selection") after B6a/B6b ("split
run()", "delete the executor thread"), whose stated motivation is "after B6bthere is no second thread". This PR does B6c before B6b, so it removes the
TLS isolation while two threads (the prepare/finalize caller and the per-run
executor) still touch the runner. It is safe because
HostApiis a by-valueobject, one per run, never shared — and the only runner state that crosses the
prepare/execute boundary is already synchronized or disjoint by index:
retained_temp_addrs_[slot],arena_banks_[bank], a2a3run_stream_slots_[slot]): each run'sHostApibinds its own slot/bank, and
try_reserve_native_runrejects any secondreservation that reuses a slot or bank already held, so two concurrent runs
never address the same index. The executor reads its run's
pipeline_slotfrom the
run()parameter and itsHostApifromstate->host_api— neitheris shared across runs.
active_native_run_: a single atomic CAS slot — only one run executes ata time (onboard
std::atomicundernative_run_mu_; simstd::atomic).block_dim_/worker_count_: latched byactivate_launch_shape()onthe executor thread immediately before
run()uses them on that samethread.
prepare_launch_shape()writes only theRuntimeobject(
runtime.set_worker_count), never these runner members — so Add: overlap HBG successor preparation with active execution #1587's fix(two concurrent prepares used to race on
block_dim_) still holds unchanged:prepares do not write them.
device_unusable_:std::atomic<bool>on a2a3 (acquire/release); a5uses a plain
bool(pre-existing, unchanged by this PR).native_launch_signal_: stable for the whole ofrun()—try_acquire_native_runsets it before the executor thread is spawned(happens-before), and
release_native_runclears it only afterfinalizejoins the executor.
So deleting TLS changes how a callback finds its runner/slot/bank, not what
it accesses. If B6a/B6b were found necessary first, it would be because of one
of the items above, none of which this PR weakens.
ABI-version answer
get_pipeline_contract'sabi_versionfield is checked at load time(
ChipWorker::initrunsis_valid_pipeline_contractright afterdlopen+dlsym, before any run; rejection is pinned bytests/ut/cpp/hierarchical/test_pipeline_contract.cpp). Because consumer(
ChipWorker) and producer (host_runtime.so) share thePTO_PIPELINE_CONTRACT_ABI_VERSIONconstant via the same header, a signaturechange does not auto-bump it — so this PR bumps it 1 → 2 in the same commit.
With the bump, a stale
host_runtime.so(compiled with version 1) is rejectedcleanly at load ("host runtime returned a PipelineContract this build cannot
accept") rather than crashing at the first
simpler_runcall. Without thebump, the
simpler_run/prepare_runsignature drift would not be caught atload (
dlsymdoes not validate C signatures). The four deleted setters arecaught in the inverse direction (old
ChipWorkervs new.so) by theload_symbol"dlsym failed for '...'" error; this PR's newChipWorkernolonger asks for them.
Test changes (forced, assertion-neutral)
HostApichanging from an aggregate of function pointers to a value objectforces fixture updates in unit tests that construct it directly — no test
assertion (no expected count, stream provisioning number, or behaviour)
changes:
tests/ut/py/test_host_runtime_abi.py: drops the two deleted setter namesfrom the required-export set (the allowed exception — the symbols are gone).
tests/ut/cpp/a2a3/test_hbg_tensor_access.cpp,tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp: rebuild their fakeHostApi/HostApiFunctionsfixtures against the new shape (fakes gain therunner_ctx/ slot / bank parameters; bodies and all expectations unchanged).tests/st/.../test_l3_launch_acceptance.py: module docstring no longer namesthe deleted setter (comment-only).
The two scene tests that would catch a slot/bank mix-up
(
native_run_lifecycle,run_stream_reuse) pass unchanged.Validation
libhost_runtime.sorebuilt (2 arch × 2 runtime); the box hasno onboard cross-compiler, so onboard DSOs build in CI / on a build host.
nm -Don every builtlibhost_runtime.so: the four setter symbols areabsent.
worker_async_fifo): see PR checks. (a5 onboard is run by CI's st-onboard-a5;
the a5 machine was unavailable locally.)