From b09f24bdb82f9e91d593d99f3b6e01e4adc9525c Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Wed, 29 Jul 2026 19:21:51 +0800 Subject: [PATCH] Add: overlap HBG successor preparation with active execution - Prepare one HBG successor in a distinct lease-selected slot and arena bank while its predecessor executes. - Keep device launch FIFO-serial and use validation-only staging for diagnostics, TMR, simulation, and unsupported backends. - Carry generation-bound identity through native prepare, launch, tracing, failure, finalization, and teardown. - Require a uniform host-runtime pipeline ABI, with explicit depth-one A5 contracts and sim capability adapters, so mismatched DSOs fail during init. - Resolve expected scheduler admission failures without throwing from the scheduler thread. - Keep per-thread run selection safe when the host runtime DSO unloads. - Cover overlap plus the eight-runtime symbol matrix, cancellation, diagnostics, cleanup, and protocol invariants. --- docs/callable-identity-registration.md | 13 +- docs/dynamic-linking.md | 9 +- docs/task-flow.md | 60 ++- docs/worker-manager.md | 38 +- python/bindings/task_interface.cpp | 17 +- python/simpler/worker.py | 169 ++++++-- .../platform/onboard/host/device_runner.cpp | 41 +- .../platform/onboard/host/device_runner.h | 9 +- .../host_build_graph/host/runtime_maker.cpp | 7 + .../host_build_graph/host/runtime_maker.cpp | 6 + .../host/runtime_maker.cpp | 6 + src/common/hierarchical/scheduler.cpp | 35 +- src/common/hierarchical/worker_manager.cpp | 101 +++-- src/common/hierarchical/worker_manager.h | 20 +- src/common/log/include/common/strace.h | 8 +- .../platform/onboard/host/c_api_shared.cpp | 175 ++++++++- .../onboard/host/device_runner_base.cpp | 182 +++++++-- .../onboard/host/device_runner_base.h | 70 +++- src/common/platform/sim/host/c_api_shared.cpp | 15 + src/common/worker/chip_worker.cpp | 319 ++++++++++----- src/common/worker/chip_worker.h | 33 +- src/common/worker/native_run_state.h | 15 +- src/common/worker/pto_runtime_c_api.h | 26 +- .../test_native_run_lifecycle.py | 119 +++++- .../test_worker_async_endpoint.py | 10 +- .../kernels/aiv/delayed_add.cpp | 66 ++++ .../orchestration/pipelined_vector_orch.cpp | 68 ++++ .../test_worker_async_fifo.py | 149 +++++-- .../hierarchical/test_pipeline_contract.cpp | 2 +- .../hierarchical/test_run_stream_slots.cpp | 26 ++ tests/ut/cpp/hierarchical/test_scheduler.cpp | 138 ++++++- tests/ut/py/test_host_runtime_abi.py | 71 ++++ tests/ut/py/test_worker/test_host_worker.py | 369 +++++++++++++++++- 33 files changed, 2016 insertions(+), 376 deletions(-) create mode 100644 tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/aiv/delayed_add.cpp create mode 100644 tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/orchestration/pipelined_vector_orch.cpp create mode 100644 tests/ut/py/test_host_runtime_abi.py diff --git a/docs/callable-identity-registration.md b/docs/callable-identity-registration.md index e9ef492e00..b2b6065530 100644 --- a/docs/callable-identity-registration.md +++ b/docs/callable-identity-registration.md @@ -525,7 +525,11 @@ The callable digest remains part of the immutable payload, while the trailer lets both sides reject a stale activation, acceptance, or completion from an older use of the same pipeline frame. `FRAME_STAGED` confirms that this exact payload has been validated and retained; it does not resolve to a different -callable identity and does not mean a native run has been prepared. +callable identity and does not by itself prove that a native run is prepared. +An eligible HBG successor normally prepares before this publication when its +predecessor is already active. A successor staged before any active claim +publishes validation-only and may later gain a native token without another +mailbox state transition. The target child loop owns the final execution resolve: @@ -698,8 +702,11 @@ Local mailbox task frames are hashid-based: - Chip and sub child loops resolve `hashid -> local_slot` immediately before execution. - A two-frame chip loop retains the resolved slot and immutable payload at - `FRAME_STAGED`; runtime-native prepare waits for activation and native-run - availability. + `FRAME_STAGED`; an eligible non-diagnostic HBG successor prepares a native + token in the lease-selected inactive bank once its predecessor owns the + active claim. A successor staged before that claim may gain the token without + another state transition. HBG tasks adjacent to diagnostics and all TMR + tasks wait for activation and native-run availability. - Protocol/run/lease-generation/dispatch identity guards task-frame reuse, and a separate sticky word records only real native launch acceptance. - `ChipWorker.run(local_slot)` remains private to the child process. diff --git a/docs/dynamic-linking.md b/docs/dynamic-linking.md index cfc2e96178..517021bb2e 100644 --- a/docs/dynamic-linking.md +++ b/docs/dynamic-linking.md @@ -291,11 +291,16 @@ ChipWorker.init(device_id, bins) # Python wrapper ctypes.CDLL(libcpu_sim_context.so, RTLD_GLOBAL) # sim only, once per process _ChipWorker.init(host_path, aicpu_path, aicore_path, device_id) # C++ dlopen(host_runtime.so, RTLD_LOCAL) - dlsym: create_device_context, destroy_device_context, simpler_init, + dlsym every required export declared in pto_runtime_c_api.h, including: + create_device_context, destroy_device_context, simpler_init, get_runtime_size, get_runtime_alignment, simpler_register_callable, simpler_prepare_run, simpler_launch_run, simpler_poll_run, simpler_wait_run, simpler_finalize_run, simpler_run, - simpler_unregister_callable, finalize_device + simpler_unregister_callable, get_pipeline_contract, + supports_concurrent_native_prepare_ctx, + set_native_run_identity_ctx, set_task_accepted_state_ctx, + get_arena_bank_gm_heap_base_ctx, get_retained_temp_addr_ctx, + finalize_device create_device_context() → DeviceContextHandle allocate zeroed, aligned, stable native-run storage per pipeline slot simpler_init(ctx, device_id, aicpu*, aicpu_size, aicore*, aicore_size) diff --git a/docs/task-flow.md b/docs/task-flow.md index fe19bba30e..27fdde0e38 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -69,9 +69,9 @@ follows `CallConfig`, and the 32-byte digest prefixes the args blob. SUB, nested/remote L3, simulation, A5, and depth-one fallbacks carry that payload in the base compatibility frame. A direct A2/A3 onboard chip child can instead use either of two task frames after the base control frame, with -`PREPARE_READY -> FRAME_STAGED -> ACTIVATE` separating validation from native -launch. The receiving child resolves the digest in its own address space in -both forms. +`PREPARE_READY -> FRAME_STAGED -> ACTIVATE` separating endpoint preparation +from native launch. The receiving child resolves the digest in its own address +space in both forms. The proposed remote L3 path keeps the same callable identity contract, but sends it in a versioned TASK frame. The remote endpoint resolves the digest @@ -348,8 +348,11 @@ finalize`, and the existing `simpler_run` / `ChipWorker.run` surface is the blocking composition of those phases. `prepare` constructs and binds the per-run `Runtime` without crossing the device launch fence; `launch` returns after the backend has actually submitted its execution; `finalize` owns -validation, copy-back, DFX, and Runtime destruction. One runner still permits -only one unfinished native run, including a prepared-but-not-launched run. +validation, copy-back, DFX, and Runtime destruction. A backend that advertises +concurrent native preparation may own one active token and one prepared but +unlaunched and unaccepted successor token in separate lease-selected banks. +Other backends permit only one unfinished native run, including a +prepared-but-not-launched run. #### Two-frame endpoint staging lane @@ -357,7 +360,8 @@ A direct A2/A3 chip endpoint with a negotiated depth of at least two uses two task frames and advertises `supports_frame_staging`. One `WorkerThread` owns both frames and drives them through a non-blocking progress interface; the child process likewise has one loop that services control traffic, both task -frames, and the single native-run lifecycle. There is no thread per frame. +frames, and the bounded active/prepared native lifecycles. There is no thread +per frame. The active and successor paths are: @@ -367,11 +371,28 @@ IDLE -> PREPARE_READY -> FRAME_STAGED -> ACTIVATE -> TASK_LAUNCHED -> TASK_DONE | TASK_FAILED ``` -`FRAME_STAGED` means only that the child validated and owns an immutable frame -snapshot. It is distinct from both an L3 run in `RunPhase::PREPARED` and a -runtime-specific `ChipWorkerNativeRun` in its prepared phase. While an active -native run exists, the successor remains at `FRAME_STAGED`; native prepare is -deferred until the predecessor has polled complete and finalized. +`FRAME_STAGED` means that the child owns an immutable frame snapshot; it does +not by itself distinguish validation-only staging from completed native +preparation. For a `host_build_graph` successor whose own configuration and +active predecessor are both non-diagnostic, the child constructs a +generation-bound native run in the leased inactive arena bank while the +predecessor executes. If the predecessor is already active, that preparation +finishes before `FRAME_STAGED` publication. A successor that reaches the child +before any predecessor owns the active claim publishes validation-only, so the +parent can activate it without deadlock; native prepare follows activation or +a later predecessor claim without another mailbox state transition. HBG tasks +adjacent to diagnostic state and all +`tensormap_and_ringbuffer` tasks also use this as a validation-only state: +native prepare waits for the predecessor's complete device fence because their +shared diagnostic or device-scratch state is not safe to rewrite early. +Backend and per-run capabilities, rather than the mailbox protocol, select +between these meanings. + +An HBG successor's prepared token remains unlaunched and unaccepted until +`ACTIVATE`, and activation still cannot launch it until the predecessor has +polled complete and finalized. The sticky acceptance word therefore remains +zero throughout preparation. Shutdown, stale activation, and pre-launch +failure finalize the token exactly once before the frame becomes terminal. The scheduler stages only the first eligible single NEXT_LEVEL task from the prepared FIFO successor. Tasks from the active run use only the active lane, so @@ -380,6 +401,12 @@ remain on their normal queue and dispatch synchronously after FIFO promotion. Remote, SUB, A5, simulation, nested-worker, and single-frame endpoints retain the blocking compatibility path. +The child validates newly visible metadata from both frames before selecting +the next active `dispatch_id`. It prepares that active token first; preparation +of an ordinary HBG successor starts only after the selected predecessor owns +the native active claim. Physical frame order therefore cannot reorder native +prepare or launch. + Every task frame carries protocol, run, lease slot, generation, dispatch, and callable identity. Its sticky acceptance word is separate from the state word and is set only by the real native launch marker. `FRAME_STAGED` never satisfies @@ -387,7 +414,12 @@ the launch fence. A terminal pre-launch failure may conservatively retire the run-level acceptance waiter, but it does not set the frame's acceptance word. The parent clears that word only immediately before reusing an `IDLE` frame. Control commands continue to use a separate base frame, so they cannot -overwrite either staged task frame. +overwrite either staged task frame. Callable prepare/register/unregister +commands defer while an active or backend-prepared token owns runtime state. +The default unbounded control wait therefore follows child liveness through a +long run; a finite control timeout includes this deferral interval, and expiry +poisons the local endpoint because the pending command's completion is +uncertain. #### TRB temporary buffer @@ -695,8 +727,8 @@ Step-by-step (one chip worker): | 3 | `Orchestrator::submit_next_level` | `slot = ring.alloc()`; move `chip_args` into `slot.task_args`; walk tags → `tensormap.lookup(a.data)`, `tensormap.lookup(b.data)`, `tensormap.insert(c.data, slot)`; push ready | | 4 | Scheduler thread | pop `slot` from worker 0's FIFO; resolve stable worker ID 0 to WT_chip_0; dispatch | | 5 | WT_chip_0 parent side | encode one leased task frame: write `config`, digest prefix, and the args blob; publish `TASK_READY` for the active lane or `PREPARE_READY` for a staged successor | -| 6 | chip_0 child process | validate the frame and resolve its digest, then publish `FRAME_STAGED`; a successor waits for `ACTIVATE`, while the active frame may proceed immediately | -| 7 | chip_0 native-run path | after activation, prepare and launch the native run; poll it to completion and finalize it before another staged frame may launch. Compatibility endpoints perform the equivalent operation through blocking `ChipWorker::run` | +| 6 | chip_0 child process | validate the frame and resolve its digest; ordinary HBG with an active predecessor also prepares the leased inactive arena bank before publishing `FRAME_STAGED`, while a frame with no active predecessor, diagnostic HBG, and TMR publish after validation and defer native prepare | +| 7 | chip_0 native-run path | after activation and the predecessor's finalization fence, launch an already-prepared HBG run or finish deferred native preparation and then launch; poll it to completion and finalize it before another staged frame may launch. Compatibility endpoints perform the equivalent operation through blocking `ChipWorker::run` | | 8 | runtime.so | translate host ptrs → device ptrs; dispatch AICPU / AICore; write output into `c`'s shm | | 9 | chip_0 child | native finalization returns; write `TASK_DONE` | | 10 | WT_chip_0 parent | observe `TASK_DONE`; push success completion | diff --git a/docs/worker-manager.md b/docs/worker-manager.md index d58412b799..eb6126750c 100644 --- a/docs/worker-manager.md +++ b/docs/worker-manager.md @@ -223,7 +223,7 @@ A two-frame `LocalMailboxEndpoint` advertises `supports_frame_staging` and is driven through `submit_progress`, `activate_progress`, and `poll_progress`. The owning `WorkerThread` is the only parent-side progress owner. The child also uses one `run_two_frame_loop`; it services the separate control base, stages -both task frames, and owns the one native-run lifecycle. +both task frames, and owns the bounded active/prepared native lifecycles. The ordinary active dispatch and the staged successor use distinct initial states: @@ -238,24 +238,40 @@ successor: IDLE -> PREPARE_READY -> FRAME_STAGED -> ACTIVATE `FRAME_STAGED` means the child validated the frame identity and arguments, resolved the callable digest, rewrote any mapped host addresses, and retained -an immutable snapshot. It does **not** mean that the runtime-specific native -run is prepared. While one native run is active, the successor remains at -`FRAME_STAGED`. Only after activation and after the predecessor is polled and -finalized does the child call native `prepare`, `launch`, `poll`, and -`finalize` for the successor. This preserves the current one-unfinished-native- -run contract. +an immutable snapshot. The state does not by itself distinguish validation-only +staging from completed native preparation. When both an HBG successor and its +already-active predecessor are non-diagnostic, the child prepares a +generation-bound native token in the successor's leased inactive bank before +publishing `FRAME_STAGED`. A successor that arrives before any active claim +publishes validation-only, allowing the parent to activate it; native prepare +follows that activation or a later predecessor claim without another mailbox +state transition. HBG tasks adjacent to diagnostic state and all TMR tasks keep +the same two-frame protocol but defer native prepare because their shared +diagnostic or device-scratch state cannot be rewritten while another run is +active. + +An HBG token remains unlaunched and unaccepted until activation, and no backend +launches a successor until the predecessor is polled and finalized. Shutdown, +stale activation, and pre-launch failure finalize any unlaunched token exactly +once. Activation is sticky on the parent side: FIFO promotion may be observed before the child reaches `FRAME_STAGED`. The endpoint records that permission and publishes `ACTIVATE` only by a compare/exchange from the matching -`FRAME_STAGED` state. The child chooses activated frames by `dispatch_id`, so a -later frame cannot bypass an earlier eligible dispatch. +`FRAME_STAGED` state. The child validates all newly visible frame metadata +before choosing the next active frame by `dispatch_id`; capable native prepare +starts for that frame first, and a successor prepares only after the selected +predecessor owns the active claim. Frame index therefore cannot let a later +dispatch bypass an earlier eligible dispatch. Task-frame publication briefly shares the base control mutex so it has a defined order relative to a control request. Once published, ordinary controls may run while the device task is active. Registry-mutating controls are -deferred until the active native run is finalized; final unregister also waits -for every published frame using that digest to retire. +deferred until active and backend-prepared native state is finalized; final +unregister also waits for every published frame using that digest to retire. +The default unbounded control wait follows child liveness through this +deferral. A finite timeout includes the deferral interval and poisons the +endpoint if it expires with completion uncertain. Each task frame is bound to its pipeline lease slot and carries a protocol trailer with `{run_id, slot_id, generation, dispatch_id}`. Parent and child diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index bb0dfac049..01fbc7f49b 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -1691,7 +1691,9 @@ NB_MODULE(_task_interface, m) { nb::class_(m, "_ChipWorkerNativeRun") .def_ro("slot_id", &ChipWorkerNativeRun::slot_id) .def_ro("generation", &ChipWorkerNativeRun::generation) - .def_ro("run_epoch", &ChipWorkerNativeRun::run_epoch); + .def_ro("run_epoch", &ChipWorkerNativeRun::run_epoch) + .def_ro("run_id", &ChipWorkerNativeRun::run_id) + .def_ro("dispatch_id", &ChipWorkerNativeRun::dispatch_id); // --- ChipWorker --- nb::class_(m, "_ChipWorker") @@ -1810,12 +1812,15 @@ NB_MODULE(_task_interface, m) { .def( "_prepare_native_run_from_blob", [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity, - const CallConfig &config, uint32_t slot_id, uint64_t generation) { + const CallConfig &config, uint32_t slot_id, uint64_t generation, uint64_t run_id, uint64_t dispatch_id) { TaskArgsView view = read_blob(reinterpret_cast(args_blob_ptr), blob_capacity); - return self.prepare_native_run(callable_id, view, config, PipelineSlotLease{slot_id, 0, generation}); + return self.prepare_native_run( + callable_id, view, config, PipelineSlotLease{slot_id, 0, generation}, run_id, dispatch_id + ); }, nb::arg("callable_id"), nb::arg("args_blob_ptr"), nb::arg("blob_capacity"), nb::arg("config"), - nb::arg("slot_id"), nb::arg("generation"), nb::call_guard(), + nb::arg("slot_id"), nb::arg("generation"), nb::arg("run_id") = 0, nb::arg("dispatch_id") = 0, + nb::call_guard(), "Prepare a generation-bound native run from a raw mailbox TaskArgs blob." ) .def( @@ -1886,6 +1891,10 @@ NB_MODULE(_task_interface, m) { .def_prop_ro("initialized", &ChipWorker::initialized) .def_prop_ro("pipeline_depth", &ChipWorker::pipeline_depth) .def_prop_ro("runtime_slot_count", &ChipWorker::runtime_slot_count) + .def_prop_ro( + "supports_concurrent_native_prepare", &ChipWorker::supports_concurrent_native_prepare, + "Whether non-diagnostic native preparation may overlap one active run in another slot." + ) .def_prop_ro( "runtime_buffer_addrs", &ChipWorker::runtime_buffer_addrs, "Host Runtime staging buffer address of every copy the runtime's " diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 8d367f6d90..9487d92882 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -2394,6 +2394,27 @@ class _StagedFrame: cid: int config: CallConfig activated: bool + native_run: Any = None + published: bool = False + + supports_concurrent_native_prepare = bool(cw._impl.supports_concurrent_native_prepare) + staged_frames: dict[int, _StagedFrame] = {} + active_frame: _StagedFrame | None = None + active_run: Any = None + + def config_has_diagnostics(config: CallConfig) -> bool: + # Mirrors CallConfig::diagnostics_any(); these modes share native + # diagnostic state and therefore use the serial prepare fallback. + return bool( + config.enable_l2_swimlane + or config.enable_dump_args + or config.enable_pmu + or config.enable_dep_gen + or config.enable_scope_stats + ) + + def has_backend_prepared_frame() -> bool: + return any(frame.native_run is not None for frame in staged_frames.values()) def read_identity(frame_buf: memoryview) -> tuple[int, int, int, int, int]: return ( @@ -2417,6 +2438,36 @@ def fail_frame(frame: _StagedFrame, message: str) -> None: _write_error(frame.frame_buf, 1, message) _mailbox_store_i32(frame.frame_addr + _OFF_STATE, _TASK_FAILED) + def publish_frame_staged(frame: _StagedFrame) -> None: + _write_error(frame.frame_buf, 0, "") + _mailbox_store_i32(frame.frame_addr + _OFF_STATE, _FRAME_STAGED) + frame.published = True + + def prepare_frame_native_run(frame: _StagedFrame) -> Any: + if frame.native_run is not None: + return frame.native_run + _protocol, run_id, slot_id, generation, dispatch_id = frame.identity + frame.native_run = cw._impl._prepare_native_run_from_blob( + frame.cid, + frame.frame_addr + _OFF_TASK_ARGS_BLOB, + _MAILBOX_ARGS_CAPACITY, + frame.config, + slot_id, + generation, + run_id, + dispatch_id, + ) + return frame.native_run + + def finalize_frame_native_run(frame: _StagedFrame) -> None: + native_run = frame.native_run + if native_run is None: + return + # Clearing ownership before the native call makes every unwind path + # attempt this token exactly once, including when finalization fails. + frame.native_run = None + cw._impl._finalize_native_run(native_run) + def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: frame_buf = frame_bufs[index] frame_addr = frame_addrs[index] @@ -2466,7 +2517,7 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: _rewrite_blob_host_addrs(frame_buf, _OFF_TASK_ARGS_BLOB, host_buf_ranges) config = _read_config_from_mailbox(frame_buf) activation_required = initial_state != _TASK_READY - staged = _StagedFrame( + return _StagedFrame( index=index, frame_buf=frame_buf, frame_addr=frame_addr, @@ -2475,17 +2526,11 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: config=config, activated=not activation_required or initial_state == _ACTIVATE, ) - _write_error(frame_buf, 0, "") - _mailbox_store_i32(frame_addr + _OFF_STATE, _FRAME_STAGED) - return staged except Exception as e: # noqa: BLE001 _write_error(frame_buf, 1, _format_exc(f"chip_process dev={device_id} frame={index}", e)) _mailbox_store_i32(frame_addr + _OFF_STATE, _TASK_FAILED) return None - staged_frames: dict[int, _StagedFrame] = {} - active_frame: _StagedFrame | None = None - active_run: Any = None parent_pid = os.getppid() liveness_countdown = _PARENT_LIVENESS_POLL_INTERVAL shutdown_message = f"chip_process dev={device_id}: task loop shut down" @@ -2497,7 +2542,8 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: if control_state == _CONTROL_REQUEST: sub_cmd = struct.unpack_from("Q", buf, _OFF_CALLABLE)[0] registry_control = sub_cmd in (_CTRL_PREPARE, _CTRL_REGISTER, _CTRL_UNREGISTER) - defer_control = registry_control and active_frame is not None + backend_prepared = has_backend_prepared_frame() + defer_control = registry_control and (active_frame is not None or backend_prepared) if sub_cmd == _CTRL_UNREGISTER: defer_control = defer_control or task_frame_references_digest(_read_control_digest(buf)) if not defer_control: @@ -2505,6 +2551,7 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: _write_error(buf, code, msg) _mailbox_store_i32(state_addr, _CONTROL_DONE) + stop_after_frame_scan = False for index in range(_TASK_FRAME_COUNT): frame_state = _mailbox_load_i32(frame_addrs[index] + _OFF_STATE) staged = staged_frames.get(index) @@ -2516,19 +2563,81 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: continue if frame_state == _ACTIVATE: if read_identity(staged.frame_buf) != staged.identity: - fail_frame(staged, f"chip_process dev={device_id}: stale activation identity") + stale_message = f"chip_process dev={device_id}: stale activation identity" + finalize_failed = False + try: + finalize_frame_native_run(staged) + except Exception as e: # noqa: BLE001 + finalize_failed = True + stale_message += "; " + _format_exc("native finalize", e) + fail_frame(staged, stale_message) if active_frame is not staged: staged_frames.pop(index, None) + if finalize_failed: + shutdown_message = stale_message + stop_after_frame_scan = True + break continue staged.activated = True + if stop_after_frame_scan: + break + + next_active = None + if active_frame is None: + activated_frames = [frame for frame in staged_frames.values() if frame.activated] + if activated_frames: + next_active = min(activated_frames, key=lambda frame: frame.identity[4]) + + for staged in sorted(staged_frames.values(), key=lambda frame: frame.identity[4]): + # A frame published before any active claim is validation-only. + # Keep considering it so activation or a later predecessor + # claim can add the missing native token. + native_prepare_now = ( + staged.native_run is None + and supports_concurrent_native_prepare + and not config_has_diagnostics(staged.config) + and ( + (active_frame is None and staged is next_active) + or ( + active_frame is not None + and staged is not active_frame + and not config_has_diagnostics(active_frame.config) + ) + ) + ) + if staged.published and not native_prepare_now: + continue + try: + if native_prepare_now: + prepare_frame_native_run(staged) + if not staged.published: + publish_frame_staged(staged) + except Exception as e: # noqa: BLE001 + prepare_message = _format_exc(f"chip_process dev={device_id}: native prepare", e) + finalize_failed = False + try: + finalize_frame_native_run(staged) + except Exception as finalize_error: # noqa: BLE001 + finalize_failed = True + prepare_message += "; " + _format_exc("native finalize", finalize_error) + fail_frame(staged, prepare_message) + staged_frames.pop(staged.index, None) + if finalize_failed: + shutdown_message = prepare_message + stop_after_frame_scan = True + break + + if stop_after_frame_scan: + break + if active_frame is not None: try: run_complete = bool(cw._impl._poll_native_run(active_run)) except Exception as e: # noqa: BLE001 poll_message = _format_exc(f"chip_process dev={device_id}: native poll", e) try: - cw._impl._finalize_native_run(active_run) + finalize_frame_native_run(active_frame) except Exception as finalize_error: # noqa: BLE001 poll_message += "; " + _format_exc("native finalize", finalize_error) fail_frame(active_frame, poll_message) @@ -2548,7 +2657,7 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: msg = "" finalize_failed = False try: - cw._impl._finalize_native_run(active_run) + finalize_frame_native_run(active_frame) except Exception as e: # noqa: BLE001 code = 1 finalize_failed = True @@ -2576,7 +2685,7 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: if active_frame is None: eligible = sorted( - (frame for frame in staged_frames.values() if frame.activated), + (frame for frame in staged_frames.values() if frame.activated and frame.published), key=lambda frame: frame.identity[4], ) if eligible: @@ -2587,23 +2696,22 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: read_identity(next_frame.frame_buf) != next_frame.identity or frame_state not in expected_states ): - fail_frame( - next_frame, - f"chip_process dev={device_id}: staged task frame changed before launch", - ) + launch_message = f"chip_process dev={device_id}: staged task frame changed before launch" + finalize_failed = False + try: + finalize_frame_native_run(next_frame) + except Exception as e: # noqa: BLE001 + finalize_failed = True + launch_message += "; " + _format_exc("native finalize", e) + fail_frame(next_frame, launch_message) staged_frames.pop(next_frame.index, None) + if finalize_failed: + shutdown_message = launch_message + break else: - _protocol, _run_id, slot_id, generation, _dispatch_id = next_frame.identity native_run = None try: - native_run = cw._impl._prepare_native_run_from_blob( - next_frame.cid, - next_frame.frame_addr + _OFF_TASK_ARGS_BLOB, - _MAILBOX_ARGS_CAPACITY, - next_frame.config, - slot_id, - generation, - ) + native_run = prepare_frame_native_run(next_frame) cw._impl._launch_native_run( native_run, next_frame.frame_addr + _OFF_ACCEPTED, @@ -2614,7 +2722,7 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: finalize_failed = False if native_run is not None: try: - cw._impl._finalize_native_run(native_run) + finalize_frame_native_run(next_frame) except Exception as finalize_error: # noqa: BLE001 finalize_failed = True launch_message += "; " + _format_exc("native finalize", finalize_error) @@ -2641,13 +2749,18 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: if active_frame is not None: active_message = shutdown_message try: - cw._impl._finalize_native_run(active_run) + finalize_frame_native_run(active_frame) except Exception as e: # noqa: BLE001 active_message += "; " + _format_exc("native finalize", e) fail_frame(active_frame, active_message) staged_frames.pop(active_frame.index, None) for staged in staged_frames.values(): - fail_frame(staged, shutdown_message) + staged_message = shutdown_message + try: + finalize_frame_native_run(staged) + except Exception as e: # noqa: BLE001 + staged_message += "; " + _format_exc("native finalize", e) + fail_frame(staged, staged_message) for index, frame_buf in enumerate(frame_bufs): frame_state_addr = frame_addrs[index] + _OFF_STATE if _mailbox_load_i32(frame_state_addr) in ( diff --git a/src/a2a3/platform/onboard/host/device_runner.cpp b/src/a2a3/platform/onboard/host/device_runner.cpp index 21a110a6fc..11f9f2658e 100644 --- a/src/a2a3/platform/onboard/host/device_runner.cpp +++ b/src/a2a3/platform/onboard/host/device_runner.cpp @@ -233,8 +233,29 @@ void DeviceRunner::destroy_native_run_thread_state(void *snapshot) noexcept { dep_gen_host_graph_destroy_capture(snapshot); } +int DeviceRunner::provision_native_run_resources(uint32_t pipeline_slot) { + return ensure_run_stream_set(pipeline_slot); +} + +int DeviceRunner::abandon_native_run_resources(uint32_t pipeline_slot) { + return retire_run_aicore_stream(pipeline_slot); +} + int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { const unsigned selected_pipeline_slot = pipeline_slot(); + // The AICore stream is created during native prepare so its provisioning + // can overlap the predecessor's execution. Once run() is entered, every + // exit owns retirement; the success path reports destroy failure, while + // early-error paths keep the original error and leave a failed-destroy + // handle in the slot so it cannot be reused. + bool aicore_stream_retired = false; + auto aicore_stream_retire = RAIIScopeGuard([this, selected_pipeline_slot, &aicore_stream_retired]() { + if (!aicore_stream_retired) (void)retire_run_aicore_stream(selected_pipeline_slot); + }); + if (!run_stream_slots_.ready(selected_pipeline_slot)) { + LOG_ERROR("run stream set %u was not provisioned during native prepare", selected_pipeline_slot); + return -1; + } // Latch this run's diagnostic enables onto the runner before the collector // paths below read them; block_dim/aicpu_thread_num are consumed locally. apply_call_config(config); @@ -267,24 +288,6 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return rc; } - rc = ensure_run_stream_set(selected_pipeline_slot); - if (rc != 0) { - LOG_ERROR("ensure_run_stream_set(%u) failed: %d", selected_pipeline_slot, rc); - return rc; - } - // The AICore stream is this run's alone: creation is the only operation - // this platform offers that is known to leave a core free of a previous - // image's cached instructions, and there is no evidence that selecting an - // already-existing stream does anything to the instruction cache. Retiring - // it on every exit path is what keeps that guarantee from depending on - // which image the next run happens to publish. - // On an early return the guard retires it; the success path below retires - // it explicitly so a failed destroy is reported instead of discarded. - bool aicore_stream_retired = false; - auto aicore_stream_retire = RAIIScopeGuard([this, selected_pipeline_slot, &aicore_stream_retired]() { - if (!aicore_stream_retired) (void)retire_run_aicore_stream(selected_pipeline_slot); - }); - ensure_device_wall_buffer(); if (block_dim < 1) { @@ -507,7 +510,7 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { // The run owns its AICore stream, so a destroy this run cannot complete is // this run's failure: reporting success would leave the caller believing a - // slot is reusable that ensure_run_stream_set will now refuse. + // slot is reusable that the next prepare will now refuse. aicore_stream_retired = true; rc = retire_run_aicore_stream(selected_pipeline_slot); if (rc != 0) return rc; diff --git a/src/a2a3/platform/onboard/host/device_runner.h b/src/a2a3/platform/onboard/host/device_runner.h index 37276deacf..dfd35984cf 100644 --- a/src/a2a3/platform/onboard/host/device_runner.h +++ b/src/a2a3/platform/onboard/host/device_runner.h @@ -24,6 +24,7 @@ #include +#include #include #include #include @@ -115,7 +116,9 @@ class DeviceRunner : public DeviceRunnerBase { * and read off DeviceRunner state / HostLogger here — no per-run args. */ int run(Runtime &runtime, const CallConfig &config) override; - bool can_accept_run() const override { return !device_unusable_; } + bool can_accept_run() const override { return !device_unusable_.load(std::memory_order_acquire); } + int provision_native_run_resources(uint32_t pipeline_slot) override; + int abandon_native_run_resources(uint32_t pipeline_slot) override; // Map/unmap a device buffer into host address space via // halHostRegister(DEV_SVM_MAP_HOST) / halHostUnregister. The returned host @@ -227,7 +230,9 @@ class DeviceRunner : public DeviceRunnerBase { // this path so the next Worker re-inits clean in the same process (see // force_reset_device()). This flag fails run() fast and drives that // recovery. See run() and recover_device_or_mark_unusable(). - bool device_unusable_{false}; + // The prepared-run admission thread reads this while the sole device + // executor may poison the runner after a failed launch. + std::atomic device_unusable_{false}; struct RunStreamSet { rtStream_t aicpu{nullptr}; diff --git a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp index df21926bfd..fa6c6aba1e 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -82,6 +82,13 @@ extern "C" const PipelineContract *get_pipeline_contract(void) { return &contract; } +extern "C" int concurrent_native_prepare_supported_impl(void) { + // HBG can materialize a complete graph into the lease-selected unpublished + // arena bank. The common C API keeps collector-bearing configurations on + // the sequential path until their state is per-epoch. + return 1; +} + // RuntimeEnv (call_config.h) is the cross-runtime ABI for per-ring config and // carries RUNTIME_ENV_RING_COUNT slots, shared with tensormap_and_ringbuffer. // host_build_graph is single-ring (PTO2_MAX_RING_DEPTH == 1) and reads only the diff --git a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp index 2b859dd467..ea917e6542 100644 --- a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp @@ -37,6 +37,7 @@ #include #include +#include "../../../../common/worker/pto_runtime_c_api.h" #include "callable.h" #include "orchestration_api.h" #include "prepare_callable_common.h" @@ -299,6 +300,11 @@ int upload_tensor_allocation_storage(Runtime *runtime, const HostApi *api, const extern "C" { #endif +const PipelineContract *get_pipeline_contract(void) { + static const PipelineContract contract = {PTO_PIPELINE_CONTRACT_ABI_VERSION, 0, 1, {}}; + return &contract; +} + /** * Stage the per-callable resources for the host_build_graph variant: upload * kernel binaries and dlopen the orchestration SO on the host. The dlopen diff --git a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index e855e87fb0..382353333c 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -47,6 +47,7 @@ #include "../runtime/runtime.h" #include "../../../../common/runtime_status/error_log.h" #include "../../../../common/task_interface/call_config.h" +#include "../../../../common/worker/pto_runtime_c_api.h" #include "callable.h" #include "common/platform_config.h" #include "common/strace.h" @@ -61,6 +62,11 @@ static_assert( RUNTIME_ENV_RING_COUNT == PTO2_MAX_RING_DEPTH, "RuntimeEnv ring count must match PTO2 runtime ring depth" ); +extern "C" const PipelineContract *get_pipeline_contract(void) { + static const PipelineContract contract = {PTO_PIPELINE_CONTRACT_ABI_VERSION, 0, 1, {}}; + return &contract; +} + // Helper: return current time in milliseconds static int64_t _now_ms() { struct timeval tv; diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index 78d6c62848..56d1f845a2 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -416,10 +416,10 @@ void Scheduler::try_consume(TaskSlot slot) { // The NEXT_LEVEL worker-id set is fixed by Worker::init() (NextLevelReadyQueues // is reset once, before scheduling starts) and every queued/targeted worker id -// is validated in Orchestrator::submit_impl. The `throw`s in the dispatch -// helpers below are therefore unreachable invariant checks; because they run on -// sched_thread_ with no surrounding handler, any throw is fatal to the whole -// worker tree (std::terminate), not a per-task failure. +// is validated in Orchestrator::submit_impl. WorkerThread resolves expected +// lane/capacity/stopping rejections through exactly one complete_unpublished +// call, under the same non-throwing completion-callback contract as ordinary +// endpoint completion. void Scheduler::dispatch_ready() { dispatch_round_count_.fetch_add(1, std::memory_order_relaxed); std::optional run_snapshot; @@ -448,18 +448,14 @@ bool claim_for_dispatch(TaskSlotState &s) { } void Scheduler::dispatch_claimed(WorkerThread *worker, WorkerDispatch dispatch, bool prepared) { - try { - if (prepared) { - worker->dispatch_prepared(dispatch); - } else { - worker->dispatch(dispatch); - } - } catch (const std::exception &e) { - worker->complete_unpublished( - dispatch, std::string("Scheduler failed to publish a claimed dispatch: ") + e.what() - ); - } catch (...) { - worker->complete_unpublished(dispatch, "Scheduler failed to publish a claimed dispatch"); + // WorkerThread owns publication failure through its one terminal callback. + // Retrying here after that callback starts can duplicate a partially + // published completion. The callback contract is the same non-throwing + // contract used by ordinary endpoint completions. + if (prepared) { + worker->dispatch_prepared(dispatch); + } else { + worker->dispatch(dispatch); } } @@ -483,6 +479,13 @@ void Scheduler::dispatch_preparable_next_level_singles() { cfg_.enqueue_ready_cb(slot); continue; } + // Diagnostic setup mutates runner-global state, so it starts only + // after this run reaches the active FIFO lane. Ordinary tasks may use + // the prepared lane because their backend preparation is run-local. + if (state.config.diagnostics_any()) { + cfg_.enqueue_ready_cb(slot); + continue; + } if (cfg_.before_claim_cb) cfg_.before_claim_cb(slot); if (!claim_for_dispatch(state)) continue; dispatch_claimed(worker, WorkerDispatch{slot, 0}, /*prepared=*/true); diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 619365a992..8fa9c193da 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -298,59 +298,104 @@ void WorkerThread::dispatch(WorkerDispatch d) { d.prepare_only = false; bool expected = false; if (!active_inflight_.compare_exchange_strong(expected, true, std::memory_order_acq_rel)) { - throw std::logic_error("WorkerThread::dispatch: active lane is occupied"); + complete_unpublished(d, "WorkerThread::dispatch: active lane is occupied"); + return; } + EnqueueDispatchResult result; try { - enqueue_dispatch(d); + result = enqueue_dispatch(d); + } catch (const std::exception &e) { + // Restore admission before formatting or publishing the failure: both + // operations may allocate, and neither may strand the active lane. + active_inflight_.store(false, std::memory_order_release); + complete_unpublished(d, std::string("WorkerThread::dispatch: enqueue failed: ") + e.what()); + return; } catch (...) { active_inflight_.store(false, std::memory_order_release); - throw; + complete_unpublished(d, "WorkerThread::dispatch: enqueue failed"); + return; + } + if (result == EnqueueDispatchResult::QUEUED) return; + active_inflight_.store(false, std::memory_order_release); + if (result == EnqueueDispatchResult::STOPPING) { + complete_unpublished(d, "WorkerThread::dispatch: worker is stopping"); + } else { + complete_unpublished(d, "WorkerThread::dispatch: endpoint capacity exceeded"); } } void WorkerThread::dispatch_prepared(WorkerDispatch d) { + d.prepare_only = true; if (!caps().supports_frame_staging) { - throw std::logic_error("WorkerThread::dispatch_prepared: endpoint does not support frame staging"); + complete_unpublished(d, "WorkerThread::dispatch_prepared: endpoint does not support frame staging"); + return; + } + if (ring_ == nullptr) { + complete_unpublished(d, "WorkerThread::dispatch_prepared: null ring"); + return; } - if (ring_ == nullptr) throw std::logic_error("WorkerThread::dispatch_prepared: null ring"); TaskSlotState *slot = ring_->slot_state(d.task_slot); if (slot == nullptr || slot->run_id == INVALID_RUN_ID) { - throw std::logic_error("WorkerThread::dispatch_prepared: dispatch has no run identity"); + complete_unpublished(d, "WorkerThread::dispatch_prepared: dispatch has no run identity"); + return; } + bool staged_lane_occupied = false; { std::lock_guard lane_lk(lane_mu_); if (staged_run_id_.load(std::memory_order_relaxed) != INVALID_RUN_ID) { - throw std::logic_error("WorkerThread::dispatch_prepared: worker already owns a staged run"); + staged_lane_occupied = true; + } else { + staged_run_id_.store(slot->run_id, std::memory_order_relaxed); + staged_dispatch_id_ = 0; } - staged_run_id_.store(slot->run_id, std::memory_order_relaxed); - staged_dispatch_id_ = 0; } - d.prepare_only = true; - try { - enqueue_dispatch(d, slot->run_id); - } catch (...) { + if (staged_lane_occupied) { + complete_unpublished(d, "WorkerThread::dispatch_prepared: worker already owns a staged run"); + return; + } + auto release_staged_lane = [&] { std::lock_guard lane_lk(lane_mu_); if (staged_run_id_.load(std::memory_order_relaxed) == slot->run_id) { staged_run_id_.store(INVALID_RUN_ID, std::memory_order_relaxed); staged_dispatch_id_ = 0; } - throw; + }; + EnqueueDispatchResult result; + try { + result = enqueue_dispatch(d, slot->run_id); + } catch (const std::exception &e) { + release_staged_lane(); + complete_unpublished(d, std::string("WorkerThread::dispatch_prepared: enqueue failed: ") + e.what()); + return; + } catch (...) { + release_staged_lane(); + complete_unpublished(d, "WorkerThread::dispatch_prepared: enqueue failed"); + return; + } + if (result == EnqueueDispatchResult::QUEUED) return; + release_staged_lane(); + if (result == EnqueueDispatchResult::STOPPING) { + complete_unpublished(d, "WorkerThread::dispatch_prepared: worker is stopping"); + } else if (result == EnqueueDispatchResult::CAPACITY_EXCEEDED) { + complete_unpublished(d, "WorkerThread::dispatch_prepared: endpoint capacity exceeded"); + } else { + complete_unpublished(d, "WorkerThread::dispatch_prepared: staged lane identity changed before enqueue"); } } -void WorkerThread::enqueue_dispatch(WorkerDispatch d, RunId staged_run_id) { +WorkerThread::EnqueueDispatchResult WorkerThread::enqueue_dispatch(WorkerDispatch d, RunId staged_run_id) { std::lock_guard lk(mu_); if (shutdown_.load(std::memory_order_acquire)) { - throw std::logic_error("WorkerThread::dispatch: worker is stopping"); + return EnqueueDispatchResult::STOPPING; } if (inflight_.load(std::memory_order_acquire) >= endpoint_->caps().max_inflight_tasks) { - throw std::logic_error("WorkerThread::dispatch: endpoint capacity exceeded"); + return EnqueueDispatchResult::CAPACITY_EXCEEDED; } d.dispatch_id = next_dispatch_id_; if (staged_run_id != INVALID_RUN_ID) { std::lock_guard lane_lk(lane_mu_); if (staged_run_id_.load(std::memory_order_relaxed) != staged_run_id || staged_dispatch_id_ != 0) { - throw std::logic_error("WorkerThread::dispatch_prepared: staged lane identity changed before enqueue"); + return EnqueueDispatchResult::STAGED_IDENTITY_CHANGED; } staged_dispatch_id_ = d.dispatch_id; } @@ -369,6 +414,7 @@ void WorkerThread::enqueue_dispatch(WorkerDispatch d, RunId staged_run_id) { ++next_dispatch_id_; inflight_.fetch_add(1, std::memory_order_release); cv_.notify_one(); + return EnqueueDispatchResult::QUEUED; } bool WorkerThread::activate_prepared(RunId run_id) { @@ -589,7 +635,12 @@ void WorkerThread::loop() { void WorkerThread::finish_progress_dispatch(const WorkerEndpointProgress &progress) { const WorkerDispatch &dispatch = progress.dispatch; - if (progress.kind == WorkerProgressKind::FRAME_STAGED) return; + if (progress.kind == WorkerProgressKind::FRAME_STAGED) { + // The endpoint already owns the prepared frame. This cursor-only event + // keeps the progress poll moving; acceptance, completion, and inflight + // ownership intentionally remain unchanged until activation/terminal. + return; + } if (progress.kind == WorkerProgressKind::ACCEPTED) { if (!accepted_dispatch_ids_.insert(dispatch.dispatch_id).second) return; @@ -1226,10 +1277,14 @@ static void write_control_digest(char *mbox, const uint8_t *digest) { } // Issue a control sub-command and block until the child publishes -// CONTROL_DONE. Caller must hold `mailbox_mu_`. On a non-zero error code -// from the child, throws and leaves the mailbox in IDLE before unwinding -// (so the next claim starts from a clean state). The `op_name` is used -// only for the exception message. +// CONTROL_DONE. Caller must hold `mailbox_mu_`. Registry controls may remain +// in CONTROL_REQUEST while the child owns an active native run or a live frame +// references the callable digest. The default infinite timeout waits for that +// deferral subject to child-liveness checks; an explicit timeout includes the +// deferred interval, and expiry poisons the endpoint. On a non-zero error code +// from the child, throws and leaves the mailbox in IDLE before unwinding (so +// the next claim starts from a clean state). The `op_name` is used only for the +// exception message. void LocalMailboxEndpoint::run_control_command(const char *op_name, double timeout_s) { if (mailbox_control_timed_out_) { throw std::runtime_error(std::string(op_name) + " failed: mailbox has an unresolved timed-out control command"); diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 8ad4f13db4..28e7413bac 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -490,7 +490,7 @@ class WorkerThread { // Enqueue a dispatch for the worker. Non-blocking. void dispatch(WorkerDispatch d); void dispatch_prepared(WorkerDispatch d); - // Complete a slot the scheduler already claimed when publication throws. + // Complete a slot the scheduler already claimed when publication fails. // No endpoint capacity was consumed, but the run-level acceptance waiter // still needs its conservative terminal fallback. void complete_unpublished(WorkerDispatch d, const std::string &error_message); @@ -581,6 +581,13 @@ class WorkerThread { void control_l3_l2_region_release(uint64_t region_id); private: + enum class EnqueueDispatchResult : uint8_t { + QUEUED, + STOPPING, + CAPACITY_EXCEEDED, + STAGED_IDENTITY_CHANGED, + }; + Ring *ring_{nullptr}; std::unique_ptr endpoint_; std::function on_complete_; @@ -605,7 +612,7 @@ class WorkerThread { void loop(); WorkerCompletion dispatch_process(WorkerDispatch d, const std::function &on_accept); - void enqueue_dispatch(WorkerDispatch d, RunId staged_run_id = INVALID_RUN_ID); + EnqueueDispatchResult enqueue_dispatch(WorkerDispatch d, RunId staged_run_id = INVALID_RUN_ID); void finish_progress_dispatch(const WorkerEndpointProgress &progress); void fail_progress_driver(const std::string &reason) noexcept; }; @@ -630,9 +637,12 @@ class WorkerManager { void add_next_level_endpoint(std::unique_ptr endpoint); void add_sub(void *mailbox, int child_pid = -1); - // `on_accept` advances the run's launch fence; pass an empty function only - // when nothing waits on that fence, since an omitted callback leaves - // pending_accepts non-zero forever. No default: the choice is the caller's. + // `on_complete` and `on_accept` are non-throwing contracts: WorkerThread invokes + // `on_complete` after endpoint ownership has ended, so no lower layer can + // reconstruct run accounting if it unwinds. `on_accept` advances the run's + // launch fence; pass an empty function only when nothing waits on that + // fence, since an omitted callback leaves pending_accepts non-zero forever. + // No default: the choice is the caller's. // // `on_idle` fires once per finished dispatch, after that worker's lane // state is published. An edge-triggered scheduler must supply it: a diff --git a/src/common/log/include/common/strace.h b/src/common/log/include/common/strace.h index 9ce4d359d0..5543e5a15e 100644 --- a/src/common/log/include/common/strace.h +++ b/src/common/log/include/common/strace.h @@ -228,8 +228,8 @@ emit_span_at(const char *name, long long ts_ns, long long dur_ns, int depth, con } /** Emit an explicitly timed host-domain span in the active invocation. */ -inline void emit_host_span_at(const char *name, long long ts_ns, long long dur_ns, int depth) { - emit_span_at(name, ts_ns, dur_ns, depth, ""); +inline void emit_host_span_at(const char *name, long long ts_ns, long long dur_ns, int depth, const char *attrs = "") { + emit_span_at(name, ts_ns, dur_ns, depth, attrs); } } // namespace simpler::strace @@ -256,6 +256,9 @@ inline void emit_host_span_at(const char *name, long long ts_ns, long long dur_n /** Emit a host-domain span measured across disjoint API calls. */ #define STRACE_HOST_SPAN_AT(name, ts_ns, dur_ns, depth) \ ::simpler::strace::emit_host_span_at((name), (ts_ns), (dur_ns), (depth)) +/** Emit a disjoint host-domain span with caller-formatted attributes. */ +#define STRACE_HOST_SPAN_AT_A(name, ts_ns, dur_ns, depth, attrs) \ + ::simpler::strace::emit_host_span_at((name), (ts_ns), (dur_ns), (depth), (attrs)) /** Emit a device-domain span (device-clock start `ts_ns` + measured `dur_ns`). */ #define STRACE_DEV_SPAN_AT(name, ts_ns, dur_ns, depth) \ ::simpler::strace::emit_span_at((name), (ts_ns), (dur_ns), (depth)) @@ -270,6 +273,7 @@ inline void emit_host_span_at(const char *name, long long ts_ns, long long dur_n #define STRACE_CONTEXT(inv, hid, depth) ((void)0) #define STRACE_NOW_NS() 0LL #define STRACE_HOST_SPAN_AT(name, ts_ns, dur_ns, depth) ((void)0) +#define STRACE_HOST_SPAN_AT_A(name, ts_ns, dur_ns, depth, attrs) ((void)0) #define STRACE_DEV_SPAN_AT(name, ts_ns, dur_ns, depth) ((void)0) #endif // SIMPLER_HOST_STRACE diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index e19c80b52d..c6322d65b9 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -36,6 +36,7 @@ #include #include +#include #include #include #include @@ -66,6 +67,7 @@ extern "C" { * =========================================================================== */ int register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(const void *), CallableArtifacts *out); int validate_runtime_impl(Runtime *runtime, const HostApi *api, int execution_rc); +__attribute__((weak)) int concurrent_native_prepare_supported_impl(void) { return 0; } /* =========================================================================== * Per-thread DeviceRunnerBase binding (set by simpler_register_callable / simpler_run) @@ -258,7 +260,7 @@ static const HostApi g_host_api = { void destroy_device_context(DeviceContextHandle ctx) { DeviceRunnerBase *runner = static_cast(ctx); - if (runner != nullptr && runner->native_run_active()) { + if (runner != nullptr && runner->native_runs_outstanding()) { LOG_ERROR("destroy_device_context: refusing to destroy a context with an unfinalized native run"); return; } @@ -307,7 +309,7 @@ int finalize_device(DeviceContextHandle ctx) { if (ctx == NULL) return -1; try { DeviceRunnerBase *runner = static_cast(ctx); - if (runner->native_run_active()) { + if (runner->native_runs_outstanding()) { LOG_ERROR("finalize_device: native run must be finalized first"); return -1; } @@ -332,6 +334,9 @@ int simpler_init( auto tsd_guard = RAIIScopeGuard([]() { pthread_setspecific(g_runner_key, nullptr); }); + (void)runner->select_pipeline_slot(0); + (void)runner->select_arena_bank(0); + (void)runner->set_native_run_identity(0, 0, 0, 0); // CANN dlog must be levelled BEFORE the device context is opened // (rtSetDevice inside attach_current_thread): CANN snapshots the @@ -409,7 +414,7 @@ int simpler_init( int simpler_register_callable(DeviceContextHandle ctx, int32_t callable_id, const void *callable) { if (ctx == NULL || callable == NULL) return -1; DeviceRunnerBase *runner = static_cast(ctx); - if (runner->native_run_active()) { + if (runner->native_runs_outstanding()) { LOG_ERROR("simpler_register_callable: native run must be finalized before mutating the callable registry"); return -1; } @@ -584,16 +589,23 @@ static OnboardNativeRunState *native_run_state(DeviceContextHandle ctx, RuntimeH return state; } -static void emit_native_run_host_wall(unsigned trace_inv, uint64_t trace_hid, long long trace_start_ns) { +static void +emit_native_run_host_wall(unsigned trace_inv, uint64_t trace_hid, long long trace_start_ns, const char *trace_attrs) { const long long end_ns = STRACE_NOW_NS(); STRACE_CONTEXT(trace_inv, trace_hid, 0); - STRACE_HOST_SPAN_AT("simpler_run", trace_start_ns, end_ns - trace_start_ns, 0); + STRACE_HOST_SPAN_AT_A("simpler_run", trace_start_ns, end_ns - trace_start_ns, 0, trace_attrs); +} + +int supports_concurrent_native_prepare_ctx(DeviceContextHandle ctx) { + return ctx != nullptr && concurrent_native_prepare_supported_impl() != 0 ? 1 : 0; } static int cleanup_failed_prepare(OnboardNativeRunState *state, int execution_rc, bool clear_gm_sm) { const unsigned trace_inv = state->trace_inv; const uint64_t trace_hid = state->trace_hid; const long long trace_start_ns = state->trace_start_ns; + char trace_attrs[sizeof(state->trace_attrs)]; + std::memcpy(trace_attrs, state->trace_attrs, sizeof(trace_attrs)); if (clear_gm_sm) state->runtime.set_gm_sm_ptr(nullptr); int validation_rc = -1; try { @@ -601,13 +613,28 @@ static int cleanup_failed_prepare(OnboardNativeRunState *state, int execution_rc } catch (...) { validation_rc = -1; } + int resources_rc = 0; + if (state->runner_resources_owned) { + try { + resources_rc = state->runner->abandon_native_run_resources(state->pipeline_slot); + } catch (...) { + resources_rc = -1; + } + state->runner_resources_owned = false; + } if (state->runner_claimed) { state->runner->release_native_run(state); state->runner_claimed = false; } + if (state->runner_reserved) { + state->runner->release_native_run_reservation(state); + state->runner_reserved = false; + } destroy_native_run_state(state); - emit_native_run_host_wall(trace_inv, trace_hid, trace_start_ns); - return validation_rc != 0 ? validation_rc : execution_rc; + emit_native_run_host_wall(trace_inv, trace_hid, trace_start_ns, trace_attrs); + if (validation_rc != 0) return validation_rc; + if (resources_rc != 0) return resources_rc; + return execution_rc; } int simpler_prepare_run( @@ -650,12 +677,30 @@ int simpler_prepare_run( const long long trace_start_ns = STRACE_NOW_NS(); try { state = new (runtime) OnboardNativeRunState(runner, *config, trace_hid); - if (!runner->try_acquire_native_run(state, &state->launch_signal)) { - LOG_ERROR("simpler_prepare_run: another native run is active on this device context"); + const DeviceRunnerBase::NativeRunThreadSelection selection = runner->capture_native_run_thread_selection(); + (void)runner->set_native_run_identity(0, 0, 0, 0); + state->run_id = selection.run_id; + state->generation = selection.generation; + state->dispatch_id = selection.dispatch_id; + state->run_epoch = selection.run_epoch; + state->pipeline_slot = selection.pipeline_slot; + state->arena_bank = selection.arena_bank; + std::snprintf( + state->trace_attrs, sizeof(state->trace_attrs), + "run_id=%llu slot=%u generation=%llu dispatch_id=%llu run_epoch=%llu", + static_cast(state->run_id), state->pipeline_slot, + static_cast(state->generation), static_cast(state->dispatch_id), + static_cast(state->run_epoch) + ); + const bool allow_prepared_successor = + concurrent_native_prepare_supported_impl() != 0 && !config->diagnostics_any(); + if (!runner->try_reserve_native_run(state, state->pipeline_slot, state->arena_bank, allow_prepared_successor)) { + LOG_ERROR("simpler_prepare_run: native-run admission is occupied (%s)", state->trace_attrs); destroy_native_run_state(state); return -1; } - state->runner_claimed = true; + state->runner_reserved = true; + const bool overlaps_active_run = allow_prepared_successor && runner->native_run_active(); state->trace_inv = trace_inv; state->trace_start_ns = trace_start_ns; STRACE_CONTEXT(state->trace_inv, state->trace_hid, 1); @@ -663,10 +708,17 @@ int simpler_prepare_run( int rc = runner->attach_current_thread(runner->device_id()); if (rc != 0) return cleanup_failed_prepare(state, rc, true); + state->runner_resources_owned = true; + rc = runner->provision_native_run_resources(state->pipeline_slot); + if (rc != 0) return cleanup_failed_prepare(state, rc, true); + rc = runner->prepare_launch_shape(state->runtime, state->config); if (rc != 0) return cleanup_failed_prepare(state, rc, true); - runner->apply_call_config(state->config); + // Diagnostic binding reads runner-global collector configuration. It + // is depth-one, while concurrent HBG preparation must leave the active + // run's configuration untouched until launch. + if (!overlaps_active_run) runner->apply_call_config(state->config); { STRACE("simpler_run.bind"); @@ -687,8 +739,34 @@ int simpler_prepare_run( int simpler_launch_run(DeviceContextHandle ctx, RuntimeHandle runtime) { OnboardNativeRunState *state = native_run_state(ctx, runtime, "simpler_launch_run"); if (state == nullptr || state->phase.load(std::memory_order_acquire) != NativeRunPhase::Prepared) return -1; - if (!state->runner->can_accept_run()) return -1; - if (!state->runner_claimed || !state->runner->native_run_owned_by(state)) return -1; + if (!state->runner->can_accept_run() || !state->runner_reserved) return -1; + if (!state->runner->try_acquire_native_run(state, &state->launch_signal)) { + LOG_ERROR("simpler_launch_run: execution claim is occupied (%s)", state->trace_attrs); + return -1; + } + state->runner_claimed = true; + // The active predecessor may poison the device after this successor was + // prepared but before the execution claim becomes available. + if (!state->runner->can_accept_run()) { + state->runner->release_native_run(state); + state->runner_claimed = false; + return -1; + } + + // Phase entry points temporarily install the run's resource selection. + // Preserve the caller's selection so interleaving finalize(A) and + // launch(B) on one progress thread cannot leave that thread bound to A. + const DeviceRunnerBase::NativeRunThreadSelection caller_selection = + state->runner->capture_native_run_thread_selection(); + auto selection_guard = RAIIScopeGuard([runner = state->runner, caller_selection]() { + runner->restore_native_run_thread_selection(caller_selection); + }); + if (state->runner->select_pipeline_slot(state->pipeline_slot) != 0 || + state->runner->select_arena_bank(state->arena_bank) != 0) { + state->runner->release_native_run(state); + state->runner_claimed = false; + return -1; + } state->phase.store(NativeRunPhase::Launching, std::memory_order_release); @@ -700,12 +778,15 @@ int simpler_launch_run(DeviceContextHandle ctx, RuntimeHandle runtime) { pthread_setspecific(g_runner_key, ctx); STRACE_CONTEXT(state->trace_inv, state->trace_hid, 1); int rc = -1; + bool entered_run = false; try { int attach_rc = state->runner->attach_current_thread(state->runner->device_id()); if (attach_rc == 0) { state->adopt_host_thread_state(); + state->runner->activate_launch_shape(state->runtime); { STRACE("simpler_run.runner_run"); + entered_run = true; rc = state->runner->run(state->runtime, state->config); } } else { @@ -714,12 +795,25 @@ int simpler_launch_run(DeviceContextHandle ctx, RuntimeHandle runtime) { } catch (...) { rc = -1; } + if (entered_run) { + // run() owns stream retirement on every exit once entered. + state->runner_resources_owned = false; + } else if (state->runner_resources_owned) { + int resources_rc = -1; + try { + resources_rc = state->runner->abandon_native_run_resources(state->pipeline_slot); + } catch (...) {} + state->runner_resources_owned = false; + if (rc == 0) rc = resources_rc; + } pthread_setspecific(g_runner_key, nullptr); state->execution_rc.store(rc, std::memory_order_relaxed); state->execution_done.store(true, std::memory_order_release); state->launch_signal.notify(); }); } catch (...) { + state->runner->release_native_run(state); + state->runner_claimed = false; state->phase.store(NativeRunPhase::Prepared, std::memory_order_release); return -1; } @@ -763,6 +857,8 @@ int simpler_finalize_run(DeviceContextHandle ctx, RuntimeHandle runtime) { const unsigned trace_inv = state->trace_inv; const uint64_t trace_hid = state->trace_hid; const long long trace_start_ns = state->trace_start_ns; + char trace_attrs[sizeof(state->trace_attrs)]; + std::memcpy(trace_attrs, state->trace_attrs, sizeof(trace_attrs)); pthread_once(&g_runner_key_once, create_runner_key); pthread_setspecific(g_runner_key, ctx); @@ -771,6 +867,19 @@ int simpler_finalize_run(DeviceContextHandle ctx, RuntimeHandle runtime) { }); STRACE_CONTEXT(state->trace_inv, state->trace_hid, 1); + // Finalization can target A after the same caller has already prepared B. + // Resource selection is phase-local; restore B's caller selection on every + // return path, including validation and resource-retirement failures. + const DeviceRunnerBase::NativeRunThreadSelection caller_selection = + state->runner->capture_native_run_thread_selection(); + auto selection_guard = RAIIScopeGuard([runner = state->runner, caller_selection]() { + runner->restore_native_run_thread_selection(caller_selection); + }); + if (state->runner->select_pipeline_slot(state->pipeline_slot) != 0 || + state->runner->select_arena_bank(state->arena_bank) != 0) { + return -1; + } + int execution_rc = -1; const bool launched = phase != NativeRunPhase::Prepared; if (launched) { @@ -795,13 +904,28 @@ int simpler_finalize_run(DeviceContextHandle ctx, RuntimeHandle runtime) { validation_rc = -1; } + int resources_rc = 0; + if (!launched && state->runner_resources_owned) { + try { + resources_rc = state->runner->abandon_native_run_resources(state->pipeline_slot); + } catch (...) { + resources_rc = -1; + } + state->runner_resources_owned = false; + } + if (state->runner_claimed) { state->runner->release_native_run(state); state->runner_claimed = false; } + if (state->runner_reserved) { + state->runner->release_native_run_reservation(state); + state->runner_reserved = false; + } destroy_native_run_state(state); - emit_native_run_host_wall(trace_inv, trace_hid, trace_start_ns); + emit_native_run_host_wall(trace_inv, trace_hid, trace_start_ns, trace_attrs); if (validation_rc != 0) return validation_rc; + if (resources_rc != 0) return resources_rc; return launched ? execution_rc : 0; } @@ -828,9 +952,7 @@ int set_task_accepted_state_ctx(DeviceContextHandle ctx, volatile int32_t *state int select_pipeline_slot_ctx(DeviceContextHandle ctx, uint32_t slot_id) { if (ctx == NULL) return -1; try { - DeviceRunnerBase *runner = static_cast(ctx); - if (runner->native_run_active()) return -1; - return runner->select_pipeline_slot(slot_id); + return static_cast(ctx)->select_pipeline_slot(slot_id); } catch (...) { return -1; } @@ -839,9 +961,20 @@ int select_pipeline_slot_ctx(DeviceContextHandle ctx, uint32_t slot_id) { int select_arena_bank_ctx(DeviceContextHandle ctx, uint32_t bank_id) { if (ctx == NULL) return -1; try { - DeviceRunnerBase *runner = static_cast(ctx); - if (runner->native_run_active()) return -1; - return runner->select_arena_bank(bank_id); + return static_cast(ctx)->select_arena_bank(bank_id); + } catch (...) { + return -1; + } +} + +int set_native_run_identity_ctx( + DeviceContextHandle ctx, uint64_t run_id, uint64_t generation, uint64_t dispatch_id, uint64_t run_epoch +) { + if (ctx == NULL) return -1; + try { + return static_cast(ctx)->set_native_run_identity( + run_id, generation, dispatch_id, run_epoch + ); } catch (...) { return -1; } @@ -869,7 +1002,7 @@ int simpler_unregister_callable(DeviceContextHandle ctx, int32_t callable_id) { if (ctx == NULL) return -1; try { DeviceRunnerBase *runner = static_cast(ctx); - if (runner->native_run_active()) { + if (runner->native_runs_outstanding()) { LOG_ERROR( "simpler_unregister_callable: native run must be finalized before mutating the callable registry" ); diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index fb71abeb84..c03a677c9a 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -27,13 +27,18 @@ #include #include #include +#include #include #include #include #include +#include #include #include +#include +#include +#include #include "callable.h" #include "callable_protocol.h" @@ -64,6 +69,43 @@ extern "C" const char *const *runtime_extra_aicpu_symbols(size_t *count); namespace { +pthread_key_t g_run_selection_key; +pthread_once_t g_run_selection_once = PTHREAD_ONCE_INIT; +int g_run_selection_key_error = 0; + +using NativeRunThreadSelection = DeviceRunnerBase::NativeRunThreadSelection; +static_assert( + std::is_trivially_destructible_v, + "pthread TLS releases native-run selection storage without a DSO-local destructor" +); + +void create_run_selection_key() { + // The pthread key can outlive this dlclosed host-runtime DSO. Point its + // destructor directly at process-lifetime libc instead of a DSO-local + // lambda, which would become an invalid callback when the thread exits. + g_run_selection_key_error = pthread_key_create(&g_run_selection_key, std::free); +} + +DeviceRunnerBase::NativeRunThreadSelection &run_selection() { + int once_rc = pthread_once(&g_run_selection_once, create_run_selection_key); + if (once_rc != 0 || g_run_selection_key_error != 0) { + throw std::runtime_error("failed to create native-run pthread TLS key"); + } + auto *selection = static_cast(pthread_getspecific(g_run_selection_key)); + if (selection == nullptr) { + void *storage = std::malloc(sizeof(NativeRunThreadSelection)); + if (storage == nullptr) throw std::bad_alloc(); + selection = new (storage) NativeRunThreadSelection{}; + int set_rc = pthread_setspecific(g_run_selection_key, selection); + if (set_rc != 0) { + selection->~NativeRunThreadSelection(); + std::free(storage); + throw std::runtime_error("failed to install native-run pthread TLS state"); + } + } + return *selection; +} + HostRuntimeTimeoutConfig resolve_onboard_timeout_config() { RuntimeTimeoutConfig order_defaults{ PLATFORM_OP_EXECUTE_TIMEOUT_US, PLATFORM_STREAM_SYNC_TIMEOUT_MS, PLATFORM_ONBOARD_SCHEDULER_TIMEOUT_MS @@ -132,7 +174,7 @@ int DeviceRunnerBase::select_pipeline_slot(uint32_t slot_id) { LOG_ERROR("pipeline slot %u is outside [0, %u)", slot_id, PTO_PIPELINE_MAX_DEPTH); return -1; } - pipeline_slot_ = slot_id; + run_selection().pipeline_slot = slot_id; return 0; } @@ -141,11 +183,21 @@ int DeviceRunnerBase::select_arena_bank(uint32_t bank_id) { LOG_ERROR("arena bank %u is outside [0, %u)", bank_id, PTO_PIPELINE_MAX_DEPTH); return -1; } - arena_bank_ = bank_id; + run_selection().arena_bank = bank_id; return 0; } -uint32_t DeviceRunnerBase::pipeline_slot() const { return pipeline_slot_; } +uint32_t DeviceRunnerBase::pipeline_slot() const { return run_selection().pipeline_slot; } + +uint32_t DeviceRunnerBase::selected_arena_bank() const { return run_selection().arena_bank; } + +DeviceRunnerBase::NativeRunThreadSelection DeviceRunnerBase::capture_native_run_thread_selection() const { + return run_selection(); +} + +void DeviceRunnerBase::restore_native_run_thread_selection(const NativeRunThreadSelection &selection) noexcept { + run_selection() = selection; +} uint64_t DeviceRunnerBase::arena_bank_gm_heap_base(uint32_t bank_id) const { if (bank_id >= arena_banks_.size()) return 0; @@ -179,13 +231,13 @@ int DeviceRunnerBase::device_memset(void *dev_ptr, int value, std::size_t bytes) } void DeviceRunnerBase::get_retained_temp_buffer(void **addr, size_t *size) { - if (addr != nullptr) *addr = retained_temp_addrs_[pipeline_slot_]; - if (size != nullptr) *size = retained_temp_sizes_[pipeline_slot_]; + if (addr != nullptr) *addr = retained_temp_addrs_[pipeline_slot()]; + if (size != nullptr) *size = retained_temp_sizes_[pipeline_slot()]; } void DeviceRunnerBase::set_retained_temp_buffer(void *addr, size_t size) { - retained_temp_addrs_[pipeline_slot_] = addr; - retained_temp_sizes_[pipeline_slot_] = size; + retained_temp_addrs_[pipeline_slot()] = addr; + retained_temp_sizes_[pipeline_slot()] = size; } void DeviceRunnerBase::clear_temporary_buffer() { @@ -223,7 +275,7 @@ bool DeviceRunnerBase::lookup_prebuilt_runtime_arena_cache( ) const { // The cache holds one entry and its bases point into bank 0, so any other // bank must rebuild rather than be handed a region it does not own. - if (arena_bank_ != 0) return false; + if (selected_arena_bank() != 0) return false; if (!prebuilt_runtime_arena_cache_valid_ || prebuilt_runtime_arena_cache_hash_ != hash || prebuilt_runtime_arena_cache_key_.size() != key_size || key_data == nullptr || gm_heap_base == nullptr || sm_base == nullptr || runtime_arena_base == nullptr || runtime_off == nullptr || image_data == nullptr || @@ -247,7 +299,7 @@ void DeviceRunnerBase::mark_prebuilt_runtime_arena_cached( size_t runtime_off, const void *image_data, size_t image_size ) { // Single-entry cache owned by bank 0; see lookup_prebuilt_runtime_arena_cache. - if (arena_bank_ != 0) return; + if (selected_arena_bank() != 0) return; prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_hash_ = hash; prebuilt_runtime_arena_cache_key_.assign( @@ -323,7 +375,7 @@ int DeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, bank.cached_gm_heap_size = 0; bank.cached_gm_sm_size = 0; bank.cached_runtime_arena_size = 0; - if (arena_bank_ == 0) { + if (selected_arena_bank() == 0) { prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; @@ -333,7 +385,7 @@ int DeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, } return -1; } - if (arena_changed && arena_bank_ == 0) { + if (arena_changed && selected_arena_bank() == 0) { prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; @@ -346,8 +398,10 @@ int DeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, std::thread DeviceRunnerBase::create_thread(std::function fn) { int dev_id = device_id_; - return std::thread([dev_id, fn = std::move(fn)]() { + NativeRunThreadSelection selection = capture_native_run_thread_selection(); + return std::thread([this, dev_id, selection, fn = std::move(fn)]() { rtSetDevice(dev_id); + restore_native_run_thread_selection(selection); fn(); }); } @@ -373,12 +427,14 @@ int DeviceRunnerBase::attach_current_thread(int device_id) { return rc; } + // simpler_init performs the only lifetime write. Prepared-run admission + // and execution subsequently attach different host threads, so repeated + // same-value writes here would still be a C++ data race. if (device_id_ == -1) { timeout_config_ = resolve_onboard_timeout_config(); configure_aicore_op_timeout(); + device_id_ = device_id; } - - device_id_ = device_id; return 0; } @@ -1152,8 +1208,6 @@ int DeviceRunnerBase::finalize_common() { bank->cached_gm_sm_size = 0; bank->cached_runtime_arena_size = 0; } - pipeline_slot_ = 0; - arena_bank_ = 0; return rc; } @@ -1265,9 +1319,8 @@ int DeviceRunnerBase::resolve_block_dim() { ); return -1; } - block_dim_ = max_block_dim_; - LOG_INFO("block_dim resolved to %d (cube=%u, vector=%u)", block_dim_, max_cube_cores_, max_vector_cores_); - return block_dim_; + LOG_INFO("block_dim resolved to %d (cube=%u, vector=%u)", max_block_dim_, max_cube_cores_, max_vector_cores_); + return max_block_dim_; } int DeviceRunnerBase::prepare_launch_shape(Runtime &runtime, const CallConfig &config) { @@ -1286,7 +1339,6 @@ int DeviceRunnerBase::prepare_launch_shape(Runtime &runtime, const CallConfig &c } runtime.set_worker_count(num_aicore); - worker_count_ = num_aicore; // Stored for print_handshake_results in destructor runtime.set_aicpu_thread_num(config.aicpu_thread_num); // First `block_dim` cores are AIC; remaining ~2/3 are AIV. @@ -1301,6 +1353,11 @@ int DeviceRunnerBase::prepare_launch_shape(Runtime &runtime, const CallConfig &c return 0; } +void DeviceRunnerBase::activate_launch_shape(const Runtime &runtime) { + worker_count_ = runtime.get_worker_count(); + block_dim_ = worker_count_ / cores_per_blockdim_; +} + void DeviceRunnerBase::resolve_task_binary_addrs(Runtime &runtime) { // Runtime::func_id_to_addr_[] stores a CoreCallable device address; the // binary code address is one compile-time offset further in. The dispatch @@ -1489,13 +1546,33 @@ void DeviceRunnerBase::teardown_shared_collectors_after_run() { } int DeviceRunnerBase::set_task_accepted_state(volatile int32_t *state, int32_t accepted_value) { - task_accepted_state_ = state; - task_accepted_value_ = accepted_value; + run_selection().accepted_state = state; + run_selection().accepted_value = accepted_value; + return 0; +} + +int DeviceRunnerBase::set_native_run_identity( + uint64_t run_id, uint64_t generation, uint64_t dispatch_id, uint64_t run_epoch +) { + NativeRunThreadSelection &selection = run_selection(); + selection.run_id = run_id; + selection.generation = generation; + selection.dispatch_id = dispatch_id; + selection.run_epoch = run_epoch; return 0; } bool DeviceRunnerBase::try_acquire_native_run(const void *owner, NativeRunLaunchSignal *launch_signal) { if (owner == nullptr || launch_signal == nullptr) return false; + std::lock_guard lk(native_run_mu_); + bool reserved = false; + for (const NativeRunReservation &reservation : native_run_reservations_) { + if (reservation.owner == owner) { + reserved = true; + break; + } + } + if (!reserved) return false; const void *expected = nullptr; if (!active_native_run_.compare_exchange_strong( expected, owner, std::memory_order_acq_rel, std::memory_order_acquire @@ -1507,6 +1584,7 @@ bool DeviceRunnerBase::try_acquire_native_run(const void *owner, NativeRunLaunch } void DeviceRunnerBase::release_native_run(const void *owner) { + std::lock_guard lk(native_run_mu_); if (active_native_run_.load(std::memory_order_acquire) != owner) return; native_launch_signal_ = nullptr; const void *expected = owner; @@ -1523,9 +1601,65 @@ bool DeviceRunnerBase::native_run_owned_by(const void *owner) const { return owner != nullptr && active_native_run_.load(std::memory_order_acquire) == owner; } +bool DeviceRunnerBase::try_reserve_native_run( + const void *owner, uint32_t pipeline_slot, uint32_t arena_bank, bool allow_prepared_successor +) { + if (owner == nullptr || pipeline_slot >= PTO_PIPELINE_MAX_DEPTH || arena_bank >= PTO_PIPELINE_MAX_DEPTH) { + return false; + } + std::lock_guard lk(native_run_mu_); + + size_t occupied = 0; + const NativeRunReservation *existing = nullptr; + for (const NativeRunReservation &reservation : native_run_reservations_) { + if (reservation.owner == nullptr) continue; + if (reservation.owner == owner || reservation.pipeline_slot == pipeline_slot || + reservation.arena_bank == arena_bank) { + return false; + } + ++occupied; + existing = &reservation; + } + if (occupied != 0) { + const void *active = active_native_run_.load(std::memory_order_acquire); + if (!allow_prepared_successor || occupied != 1 || existing == nullptr || + !existing->permits_prepared_successor || active != existing->owner) { + return false; + } + } + + for (NativeRunReservation &reservation : native_run_reservations_) { + if (reservation.owner == nullptr) { + reservation = NativeRunReservation{owner, pipeline_slot, arena_bank, allow_prepared_successor}; + return true; + } + } + return false; +} + +void DeviceRunnerBase::release_native_run_reservation(const void *owner) { + if (owner == nullptr) return; + std::lock_guard lk(native_run_mu_); + for (NativeRunReservation &reservation : native_run_reservations_) { + if (reservation.owner == owner) { + reservation = NativeRunReservation{}; + return; + } + } +} + +bool DeviceRunnerBase::native_runs_outstanding() const { + std::lock_guard lk(native_run_mu_); + for (const NativeRunReservation &reservation : native_run_reservations_) { + if (reservation.owner != nullptr) return true; + } + return false; +} + void DeviceRunnerBase::publish_task_accepted() const { - if (task_accepted_state_ != nullptr) { - __atomic_store_n(task_accepted_state_, task_accepted_value_, __ATOMIC_RELEASE); + NativeRunThreadSelection &selection = run_selection(); + if (selection.accepted_state != nullptr) { + __atomic_store_n(selection.accepted_state, selection.accepted_value, __ATOMIC_RELEASE); } if (native_launch_signal_ != nullptr) native_launch_signal_->notify(); } diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index 2b74cba3a4..ba18d59f74 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -44,6 +44,7 @@ #include #include #include +#include #include #include #include @@ -83,6 +84,17 @@ class NativeRunLaunchSignal; */ class DeviceRunnerBase { public: + struct NativeRunThreadSelection { + uint32_t pipeline_slot{0}; + uint32_t arena_bank{0}; + volatile int32_t *accepted_state{nullptr}; + int32_t accepted_value{0}; + uint64_t run_id{0}; + uint64_t generation{0}; + uint64_t dispatch_id{0}; + uint64_t run_epoch{0}; + }; + // Public virtual dtor so the shared c_api can `delete` a polymorphic // `DeviceRunnerBase *` (the `destroy_device_context` entrypoint). Each // arch's `DeviceRunner` defaults this through the compiler-generated dtor. @@ -94,23 +106,38 @@ class DeviceRunnerBase { /** Bind this runner's launch-acceptance publication target. */ int set_task_accepted_state(volatile int32_t *state, int32_t accepted_value); + int set_native_run_identity(uint64_t run_id, uint64_t generation, uint64_t dispatch_id, uint64_t run_epoch); /** - * Reserve the runner for one native prepared-run execution. The opaque - * owner and runner-owned timing and diagnostic state remain exclusive - * through validation/finalize. + * Claim the runner for one native execution. The opaque owner and + * runner-owned timing and diagnostic state remain exclusive through + * validation/finalize. */ bool try_acquire_native_run(const void *owner, NativeRunLaunchSignal *launch_signal); void release_native_run(const void *owner); bool native_run_active() const; bool native_run_owned_by(const void *owner) const; + /** + * Reserve caller-owned native-run storage before binding starts. A + * concurrent reservation is admitted only while the first reservation + * owns the execution claim and selects distinct per-run resources. + */ + bool try_reserve_native_run( + const void *owner, uint32_t pipeline_slot, uint32_t arena_bank, bool allow_prepared_successor + ); + void release_native_run_reservation(const void *owner); + bool native_runs_outstanding() const; + /** Carry an already-validated run lease into resource selection. */ int select_pipeline_slot(uint32_t slot_id); int select_arena_bank(uint32_t bank_id); /** Slot selected by the current run lease. */ uint32_t pipeline_slot() const; + uint32_t selected_arena_bank() const; + NativeRunThreadSelection capture_native_run_thread_selection() const; + void restore_native_run_thread_selection(const NativeRunThreadSelection &selection) noexcept; /** * Committed GM heap base of one arena bank, or 0 while that bank has never @@ -413,7 +440,7 @@ class DeviceRunnerBase { * Publish this run's core geometry onto `Runtime` before the graph is * built: resolves `block_dim`, derives `num_aicore = block_dim * * cores_per_blockdim_`, range-checks against `RUNTIME_MAX_WORKER`, - * publishes `worker_count` / `worker_count_` / `aicpu_thread_num`, + * publishes the Runtime's `worker_count` / `aicpu_thread_num`, * and zero-initializes the handshake worker array with AIC/AIV core * typing (first `block_dim` cores are AIC, remaining are AIV). * @@ -426,6 +453,9 @@ class DeviceRunnerBase { */ int prepare_launch_shape(Runtime &runtime, const CallConfig &config); + /** Latch a prepared Runtime's geometry immediately before execution. */ + void activate_launch_shape(const Runtime &runtime); + /** * Replay a previously-registered callable's state onto a fresh Runtime and * complete the per-run binding in one step. Writes back kernel addrs and @@ -503,6 +533,10 @@ class DeviceRunnerBase { */ virtual bool can_accept_run() const = 0; + /** Provision/abandon platform resources owned by one prepared native run. */ + virtual int provision_native_run_resources(uint32_t /*pipeline_slot*/) { return 0; } + virtual int abandon_native_run_resources(uint32_t /*pipeline_slot*/) { return 0; } + /** * Execute a Runtime. Each arch implements its own `run()` — the bodies * are too divergent for a shared implementation (FFTS / dep_gen / ACL @@ -719,7 +753,8 @@ class DeviceRunnerBase { * at bind time, before any stream work for the run. * * Returns the resolved block_dim on success, -1 if the ceiling was - * never latched. Updates `block_dim_` on success. + * never latched. The value is latched into runner execution state only + * when the prepared Runtime is launched. */ int resolve_block_dim(); @@ -891,18 +926,21 @@ class DeviceRunnerBase { // Same re-register semantics as `aicpu_dlopen_total_`, but for hbg // variants. size_t host_dlopen_total_{0}; - volatile int32_t *task_accepted_state_{nullptr}; - int32_t task_accepted_value_{0}; + struct NativeRunReservation { + const void *owner{nullptr}; + uint32_t pipeline_slot{0}; + uint32_t arena_bank{0}; + bool permits_prepared_successor{false}; + }; + mutable std::mutex native_run_mu_; + std::array native_run_reservations_{}; std::atomic active_native_run_{nullptr}; NativeRunLaunchSignal *native_launch_signal_{nullptr}; // ---- State shared by both a2a3 and a5 --------------------------------- // - // `device_id_` is set once in `attach_current_thread()` (called from - // simpler_init during ChipWorker::init) and read on every subsequent - // op. All ChipWorker callers run on the same thread that called - // init, so plain int + the init→user happens-before edge is - // sufficient. + // `device_id_` is written once by simpler_init and is immutable while + // native prepare, execution, and collector threads attach to the runner. int device_id_{-1}; int block_dim_{0}; int cores_per_blockdim_{PLATFORM_CORES_PER_BLOCKDIM}; @@ -980,13 +1018,7 @@ class DeviceRunnerBase { // Held by pointer because DeviceArena is non-copyable and non-movable, so // the array cannot be brace-initialised without naming every bank. std::array, PTO_PIPELINE_MAX_DEPTH> arena_banks_; - ArenaBank &arena_bank() { return *arena_banks_[arena_bank_]; } - - // Selection carried by the lease of the run currently being served. Both - // default to zero, which is what an unleased run and every depth-one - // runtime use. - uint32_t pipeline_slot_{0}; - uint32_t arena_bank_{0}; + ArenaBank &arena_bank() { return *arena_banks_[selected_arena_bank()]; } bool prebuilt_runtime_arena_cache_valid_{false}; uint64_t prebuilt_runtime_arena_cache_hash_{0}; diff --git a/src/common/platform/sim/host/c_api_shared.cpp b/src/common/platform/sim/host/c_api_shared.cpp index 5a95715378..d9e510b135 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -774,6 +774,21 @@ int simpler_run( return finalize_rc != 0 ? finalize_rc : rc; } +int supports_concurrent_native_prepare_ctx(DeviceContextHandle) { return 0; } + +int set_task_accepted_state_ctx(DeviceContextHandle ctx, volatile int32_t *state, int32_t accepted_value) { + if (ctx == NULL) return -1; + try { + return static_cast(ctx)->set_task_accepted_state(state, accepted_value); + } catch (...) { + return -1; + } +} + +int set_native_run_identity_ctx(DeviceContextHandle ctx, uint64_t, uint64_t, uint64_t, uint64_t) { + return ctx == NULL ? -1 : 0; +} + int select_pipeline_slot_ctx(DeviceContextHandle ctx, uint32_t slot_id) { if (ctx == NULL) return -1; SimDeviceRunnerBase *runner = static_cast(ctx); diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 4430489012..1e2866ed3b 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -43,14 +43,6 @@ T load_symbol(void *handle, const char *name) { return reinterpret_cast(sym); } -template -T load_optional_symbol(void *handle, const char *name) { - dlerror(); - void *sym = dlsym(handle, name); - if (dlerror() != nullptr) return nullptr; - return reinterpret_cast(sym); -} - uint64_t next_native_run_epoch() { static std::atomic epoch{0}; uint64_t current = epoch.load(std::memory_order_relaxed); @@ -62,6 +54,12 @@ uint64_t next_native_run_epoch() { throw std::overflow_error("native-run epoch space is exhausted"); } +std::string format_native_run_identity(const ChipWorkerNativeRun &run) { + return "(run_id=" + std::to_string(run.run_id) + " slot=" + std::to_string(run.slot_id) + + " generation=" + std::to_string(run.generation) + " dispatch_id=" + std::to_string(run.dispatch_id) + + " run_epoch=" + std::to_string(run.run_epoch) + ")"; +} + std::vector read_binary_file(const std::string &path) { std::ifstream f(path, std::ios::binary | std::ios::ate); if (!f) { @@ -163,14 +161,16 @@ void ChipWorker::init( poll_run_fn_ = load_symbol(handle, "simpler_poll_run"); wait_run_fn_ = load_symbol(handle, "simpler_wait_run"); finalize_run_fn_ = load_symbol(handle, "simpler_finalize_run"); + supports_concurrent_native_prepare_fn_ = + load_symbol(handle, "supports_concurrent_native_prepare_ctx"); select_pipeline_slot_fn_ = load_symbol(handle, "select_pipeline_slot_ctx"); select_arena_bank_fn_ = load_symbol(handle, "select_arena_bank_ctx"); + set_native_run_identity_fn_ = load_symbol(handle, "set_native_run_identity_ctx"); get_arena_bank_gm_heap_base_fn_ = - load_optional_symbol(handle, "get_arena_bank_gm_heap_base_ctx"); - get_retained_temp_addr_fn_ = load_optional_symbol(handle, "get_retained_temp_addr_ctx"); - set_task_accepted_state_fn_ = - load_optional_symbol(handle, "set_task_accepted_state_ctx"); - get_pipeline_contract_fn = load_optional_symbol(handle, "get_pipeline_contract"); + load_symbol(handle, "get_arena_bank_gm_heap_base_ctx"); + get_retained_temp_addr_fn_ = load_symbol(handle, "get_retained_temp_addr_ctx"); + set_task_accepted_state_fn_ = load_symbol(handle, "set_task_accepted_state_ctx"); + get_pipeline_contract_fn = load_symbol(handle, "get_pipeline_contract"); unregister_callable_fn_ = load_symbol(handle, "simpler_unregister_callable"); get_aicpu_dlopen_count_fn_ = load_symbol(handle, "get_aicpu_dlopen_count"); get_host_dlopen_count_fn_ = load_symbol(handle, "get_host_dlopen_count"); @@ -202,15 +202,12 @@ void ChipWorker::init( throw; } - PipelineContract resolved_contract{PTO_PIPELINE_CONTRACT_ABI_VERSION, 0, 1, {}}; - if (get_pipeline_contract_fn != nullptr) { - const PipelineContract *contract = get_pipeline_contract_fn(); - if (!is_valid_pipeline_contract(contract) || !has_serviceable_arena_topology(*contract)) { - dlclose(handle); - throw std::runtime_error("host runtime returned a PipelineContract this build cannot accept"); - } - resolved_contract = *contract; + const PipelineContract *contract = get_pipeline_contract_fn(); + if (!is_valid_pipeline_contract(contract) || !has_serviceable_arena_topology(*contract)) { + dlclose(handle); + throw std::runtime_error("host runtime returned a PipelineContract this build cannot accept"); } + const PipelineContract resolved_contract = *contract; lib_handle_ = handle; @@ -301,8 +298,10 @@ void ChipWorker::init( poll_run_fn_ = nullptr; wait_run_fn_ = nullptr; finalize_run_fn_ = nullptr; + supports_concurrent_native_prepare_fn_ = nullptr; select_pipeline_slot_fn_ = nullptr; select_arena_bank_fn_ = nullptr; + set_native_run_identity_fn_ = nullptr; get_arena_bank_gm_heap_base_fn_ = nullptr; get_retained_temp_addr_fn_ = nullptr; set_task_accepted_state_fn_ = nullptr; @@ -353,8 +352,10 @@ void ChipWorker::init( poll_run_fn_ = nullptr; wait_run_fn_ = nullptr; finalize_run_fn_ = nullptr; + supports_concurrent_native_prepare_fn_ = nullptr; select_pipeline_slot_fn_ = nullptr; select_arena_bank_fn_ = nullptr; + set_native_run_identity_fn_ = nullptr; get_arena_bank_gm_heap_base_fn_ = nullptr; get_retained_temp_addr_fn_ = nullptr; set_task_accepted_state_fn_ = nullptr; @@ -436,8 +437,10 @@ void ChipWorker::finalize() { poll_run_fn_ = nullptr; wait_run_fn_ = nullptr; finalize_run_fn_ = nullptr; + supports_concurrent_native_prepare_fn_ = nullptr; select_pipeline_slot_fn_ = nullptr; select_arena_bank_fn_ = nullptr; + set_native_run_identity_fn_ = nullptr; get_arena_bank_gm_heap_base_fn_ = nullptr; get_retained_temp_addr_fn_ = nullptr; set_task_accepted_state_fn_ = nullptr; @@ -533,12 +536,12 @@ std::vector ChipWorker::runtime_buffer_addrs() const { } uint64_t ChipWorker::arena_bank_gm_heap_base(uint32_t bank_id) const { - if (!initialized_ || get_arena_bank_gm_heap_base_fn_ == nullptr) return 0; + if (!initialized_) return 0; return get_arena_bank_gm_heap_base_fn_(device_ctx_, bank_id); } uint64_t ChipWorker::retained_temp_addr(uint32_t slot_id) const { - if (!initialized_ || get_retained_temp_addr_fn_ == nullptr) return 0; + if (!initialized_) return 0; return get_retained_temp_addr_fn_(device_ctx_, slot_id); } @@ -573,15 +576,18 @@ void ChipWorker::run_on_slot( } (void)select_slot_resources(slot_id); + if (set_native_run_identity_fn_(device_ctx_, 0, 0, 0, 0) != 0) { + throw std::runtime_error("set_native_run_identity_ctx failed"); + } void *rt = runtime_bufs_[slot_id].data(); - if (accepted_state != nullptr && set_task_accepted_state_fn_ != nullptr) { + if (accepted_state != nullptr) { int bind_rc = set_task_accepted_state_fn_(device_ctx_, accepted_state, accepted_value); if (bind_rc != 0) { throw std::runtime_error("set_task_accepted_state_ctx failed with code " + std::to_string(bind_rc)); } } auto clear_accepted_state = [&]() { - if (accepted_state != nullptr && set_task_accepted_state_fn_ != nullptr) { + if (accepted_state != nullptr) { (void)set_task_accepted_state_fn_(device_ctx_, nullptr, 0); } }; @@ -602,27 +608,31 @@ void ChipWorker::run_on_slot( } ChipWorkerNativeRun ChipWorker::prepare_native_run( - int32_t callable_id, TaskArgsView args, const CallConfig &config, const PipelineSlotLease &lease + int32_t callable_id, TaskArgsView args, const CallConfig &config, const PipelineSlotLease &lease, uint64_t run_id, + uint64_t dispatch_id ) { ChipStorageTaskArgs chip_storage = view_to_chip_storage(args); - return prepare_native_run(callable_id, &chip_storage, config, lease); + return prepare_native_run(callable_id, &chip_storage, config, lease, run_id, dispatch_id); } ChipWorkerNativeRun ChipWorker::prepare_native_run( - int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, const PipelineSlotLease &lease + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, const PipelineSlotLease &lease, + uint64_t run_id, uint64_t dispatch_id ) { if (lease.reserved != 0 || lease.generation == 0 || lease.slot_id >= pipeline_contract_.pipeline_depth) { throw std::runtime_error("native-run pipeline lease is outside the runtime PipelineContract"); } - if (!pipeline_generations_.admit(lease)) { - throw std::runtime_error("native-run pipeline lease generation is stale"); - } - return prepare_native_run_on_slot(callable_id, args, config, lease.slot_id, lease.generation); + return prepare_native_run_on_slot(callable_id, args, config, lease.slot_id, lease.generation, run_id, dispatch_id); +} + +bool ChipWorker::supports_concurrent_native_prepare() const { + return initialized_ && pipeline_contract_.pipeline_depth > 1 && + supports_concurrent_native_prepare_fn_(device_ctx_) > 0; } ChipWorkerNativeRun ChipWorker::prepare_native_run_on_slot( int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, uint32_t slot_id, - uint64_t generation + uint64_t generation, uint64_t run_id, uint64_t dispatch_id ) { config.validate(); if (!initialized_) { @@ -634,119 +644,231 @@ ChipWorkerNativeRun ChipWorker::prepare_native_run_on_slot( if (slot_id >= runtime_bufs_.size()) { throw std::runtime_error("prepare_native_run slot is outside the runtime PipelineContract"); } - NativeRunSlotState &state = native_run_states_[slot_id]; - if (state.phase != NativeRunPhase::EMPTY) { - throw std::runtime_error("prepare_native_run slot already owns an unfinished native run"); - } - for (const NativeRunSlotState &candidate : native_run_states_) { - if (candidate.phase != NativeRunPhase::EMPTY) { + const uint64_t run_epoch = next_native_run_epoch(); + const ChipWorkerNativeRun run_identity{slot_id, generation, run_epoch, run_id, dispatch_id}; + const bool allow_prepared_successor = supports_concurrent_native_prepare() && !config.diagnostics_any(); + { + std::lock_guard lk(native_run_mu_); + NativeRunSlotState &state = native_run_states_[slot_id]; + if (state.phase != NativeRunPhase::EMPTY) { throw std::runtime_error( - "prepare_native_run cannot select another slot while an unfinished native run owns the runner" + "prepare_native_run slot already owns an unfinished native run " + + format_native_run_identity(run_identity) ); } + size_t occupied = 0; + for (const NativeRunSlotState &candidate : native_run_states_) { + if (candidate.phase == NativeRunPhase::EMPTY) continue; + ++occupied; + if (!allow_prepared_successor || !candidate.permits_prepared_successor || + (candidate.phase != NativeRunPhase::LAUNCHED && candidate.phase != NativeRunPhase::REAPED)) { + throw std::runtime_error( + "prepare_native_run requires an active predecessor before staging a successor " + + format_native_run_identity(run_identity) + ); + } + } + if (occupied != 0 && occupied != 1) { + throw std::runtime_error( + "prepare_native_run already owns a prepared successor " + format_native_run_identity(run_identity) + ); + } + if (!pipeline_generations_.admit(PipelineSlotLease{slot_id, 0, generation})) { + throw std::runtime_error( + "native-run pipeline lease generation is stale " + format_native_run_identity(run_identity) + ); + } + state.lease_generation = generation; + state.run_epoch = run_epoch; + state.phase = NativeRunPhase::PREPARING; + state.wait_rc = 0; + state.permits_prepared_successor = allow_prepared_successor; } - (void)select_slot_resources(slot_id); - const uint64_t run_epoch = next_native_run_epoch(); - int rc = prepare_run_fn_(device_ctx_, runtime_bufs_[slot_id].data(), callable_id, args, &config); + int rc = -1; + try { + (void)select_slot_resources(slot_id); + int identity_rc = set_native_run_identity_fn_(device_ctx_, run_id, generation, dispatch_id, run_epoch); + if (identity_rc != 0) { + throw std::runtime_error("set_native_run_identity_ctx failed with code " + std::to_string(identity_rc)); + } + rc = prepare_run_fn_(device_ctx_, runtime_bufs_[slot_id].data(), callable_id, args, &config); + } catch (...) { + std::lock_guard lk(native_run_mu_); + NativeRunSlotState &state = native_run_states_[slot_id]; + if (state.run_epoch == run_epoch && state.phase == NativeRunPhase::PREPARING) { + state = NativeRunSlotState{}; + } + throw; + } if (rc != 0) { - throw std::runtime_error("prepare_native_run failed with code " + std::to_string(rc)); + std::lock_guard lk(native_run_mu_); + NativeRunSlotState &state = native_run_states_[slot_id]; + if (state.run_epoch == run_epoch && state.phase == NativeRunPhase::PREPARING) { + state = NativeRunSlotState{}; + } + throw std::runtime_error( + "prepare_native_run failed with code " + std::to_string(rc) + " " + format_native_run_identity(run_identity) + ); } - state.lease_generation = generation; - state.run_epoch = run_epoch; - state.phase = NativeRunPhase::PREPARED; - state.wait_rc = 0; - return ChipWorkerNativeRun{slot_id, generation, run_epoch}; -} -ChipWorker::NativeRunSlotState & -ChipWorker::require_native_run(const ChipWorkerNativeRun &run, NativeRunPhase first, NativeRunPhase second) { - if (run.slot_id >= runtime_bufs_.size()) { - throw std::runtime_error("native-run token slot is outside the runtime PipelineContract"); - } - NativeRunSlotState &state = native_run_states_[run.slot_id]; - if (state.lease_generation != run.generation || state.run_epoch != run.run_epoch || - (state.phase != first && state.phase != second)) { - throw std::runtime_error("native-run token is stale or used in the wrong phase"); + { + std::lock_guard lk(native_run_mu_); + NativeRunSlotState &state = native_run_states_[slot_id]; + if (state.lease_generation != generation || state.run_epoch != run_epoch || + state.phase != NativeRunPhase::PREPARING) { + (void)finalize_run_fn_(device_ctx_, runtime_bufs_[slot_id].data()); + state = NativeRunSlotState{}; + throw std::runtime_error("native-run identity changed while prepare was in progress"); + } + state.phase = NativeRunPhase::PREPARED; } - return state; + return run_identity; } void ChipWorker::launch_native_run( const ChipWorkerNativeRun &run, volatile int32_t *accepted_state, int32_t accepted_value ) { - NativeRunSlotState &state = require_native_run(run, NativeRunPhase::PREPARED, NativeRunPhase::PREPARED); - if (accepted_state != nullptr && set_task_accepted_state_fn_ != nullptr) { - int bind_rc = set_task_accepted_state_fn_(device_ctx_, accepted_state, accepted_value); - if (bind_rc != 0) { - throw std::runtime_error("set_task_accepted_state_ctx failed with code " + std::to_string(bind_rc)); + { + std::lock_guard lk(native_run_mu_); + if (run.slot_id >= runtime_bufs_.size()) { + throw std::runtime_error("native-run token slot is outside the runtime PipelineContract"); + } + NativeRunSlotState &state = native_run_states_[run.slot_id]; + if (state.lease_generation != run.generation || state.run_epoch != run.run_epoch || + state.phase != NativeRunPhase::PREPARED) { + throw std::runtime_error("native-run token is stale or used in the wrong phase"); } + state.phase = NativeRunPhase::LAUNCHING; } + auto clear_accepted_state = [&]() { - if (accepted_state != nullptr && set_task_accepted_state_fn_ != nullptr) { + if (accepted_state != nullptr) { (void)set_task_accepted_state_fn_(device_ctx_, nullptr, 0); } }; int rc = -1; try { + if (accepted_state != nullptr) { + int bind_rc = set_task_accepted_state_fn_(device_ctx_, accepted_state, accepted_value); + if (bind_rc != 0) { + throw std::runtime_error("set_task_accepted_state_ctx failed with code " + std::to_string(bind_rc)); + } + } rc = launch_run_fn_(device_ctx_, runtime_bufs_[run.slot_id].data()); } catch (...) { clear_accepted_state(); + std::lock_guard lk(native_run_mu_); + NativeRunSlotState &state = native_run_states_[run.slot_id]; + if (state.lease_generation == run.generation && state.run_epoch == run.run_epoch && + state.phase == NativeRunPhase::LAUNCHING) { + state.phase = NativeRunPhase::PREPARED; + } throw; } clear_accepted_state(); if (rc != 0) { - state.phase = NativeRunPhase::REAPED; - state.wait_rc = rc; - throw std::runtime_error("launch_native_run failed with code " + std::to_string(rc)); + int poll_rc = poll_run_fn_(device_ctx_, runtime_bufs_[run.slot_id].data()); + std::lock_guard lk(native_run_mu_); + NativeRunSlotState &state = native_run_states_[run.slot_id]; + state.phase = poll_rc == SIMPLER_NATIVE_RUN_POLL_COMPLETE ? NativeRunPhase::REAPED : NativeRunPhase::PREPARED; + state.wait_rc = poll_rc == SIMPLER_NATIVE_RUN_POLL_COMPLETE ? rc : 0; + throw std::runtime_error( + "launch_native_run failed with code " + std::to_string(rc) + " " + format_native_run_identity(run) + ); } + std::lock_guard lk(native_run_mu_); + NativeRunSlotState &state = native_run_states_[run.slot_id]; state.phase = NativeRunPhase::LAUNCHED; } bool ChipWorker::poll_native_run(const ChipWorkerNativeRun &run) { - NativeRunSlotState &state = require_native_run(run, NativeRunPhase::LAUNCHED, NativeRunPhase::REAPED); - if (state.phase == NativeRunPhase::REAPED) { - return true; + { + std::lock_guard lk(native_run_mu_); + if (run.slot_id >= runtime_bufs_.size()) { + throw std::runtime_error("native-run token slot is outside the runtime PipelineContract"); + } + const NativeRunSlotState &state = native_run_states_[run.slot_id]; + if (state.lease_generation != run.generation || state.run_epoch != run.run_epoch || + (state.phase != NativeRunPhase::LAUNCHED && state.phase != NativeRunPhase::REAPED)) { + throw std::runtime_error("native-run token is stale or used in the wrong phase"); + } + if (state.phase == NativeRunPhase::REAPED) return true; } int rc = poll_run_fn_(device_ctx_, runtime_bufs_[run.slot_id].data()); if (rc == SIMPLER_NATIVE_RUN_POLL_NOT_READY) { return false; } if (rc == SIMPLER_NATIVE_RUN_POLL_COMPLETE) { + std::lock_guard lk(native_run_mu_); + NativeRunSlotState &state = native_run_states_[run.slot_id]; state.phase = NativeRunPhase::REAPED; return true; } - throw std::runtime_error("poll_native_run failed with code " + std::to_string(rc)); + throw std::runtime_error( + "poll_native_run failed with code " + std::to_string(rc) + " " + format_native_run_identity(run) + ); } void ChipWorker::wait_native_run(const ChipWorkerNativeRun &run) { - NativeRunSlotState &state = require_native_run(run, NativeRunPhase::LAUNCHED, NativeRunPhase::REAPED); - if (state.phase == NativeRunPhase::REAPED) { - return; + { + std::lock_guard lk(native_run_mu_); + if (run.slot_id >= runtime_bufs_.size()) { + throw std::runtime_error("native-run token slot is outside the runtime PipelineContract"); + } + const NativeRunSlotState &state = native_run_states_[run.slot_id]; + if (state.lease_generation != run.generation || state.run_epoch != run.run_epoch || + (state.phase != NativeRunPhase::LAUNCHED && state.phase != NativeRunPhase::REAPED)) { + throw std::runtime_error("native-run token is stale or used in the wrong phase"); + } + if (state.phase == NativeRunPhase::REAPED) return; } - state.wait_rc = wait_run_fn_(device_ctx_, runtime_bufs_[run.slot_id].data()); + int wait_rc = wait_run_fn_(device_ctx_, runtime_bufs_[run.slot_id].data()); + std::lock_guard lk(native_run_mu_); + NativeRunSlotState &state = native_run_states_[run.slot_id]; + state.wait_rc = wait_rc; state.phase = NativeRunPhase::REAPED; } void ChipWorker::finalize_native_run(const ChipWorkerNativeRun &run) { - if (run.slot_id >= runtime_bufs_.size()) { - throw std::runtime_error("native-run token slot is outside the runtime PipelineContract"); - } - NativeRunSlotState &state = native_run_states_[run.slot_id]; - if (state.lease_generation != run.generation || state.run_epoch != run.run_epoch || - state.phase == NativeRunPhase::EMPTY) { - throw std::runtime_error("native-run token is stale or already finalized"); + NativeRunPhase phase; + int wait_rc; + { + std::lock_guard lk(native_run_mu_); + if (run.slot_id >= runtime_bufs_.size()) { + throw std::runtime_error("native-run token slot is outside the runtime PipelineContract"); + } + NativeRunSlotState &state = native_run_states_[run.slot_id]; + if (state.lease_generation != run.generation || state.run_epoch != run.run_epoch || + (state.phase != NativeRunPhase::PREPARED && state.phase != NativeRunPhase::LAUNCHED && + state.phase != NativeRunPhase::REAPED)) { + throw std::runtime_error("native-run token is stale or already finalized"); + } + phase = state.phase; + wait_rc = state.wait_rc; + state.phase = NativeRunPhase::FINALIZING; } - if (state.phase == NativeRunPhase::LAUNCHED) { - wait_native_run(run); + + try { + if (phase == NativeRunPhase::LAUNCHED) { + wait_rc = wait_run_fn_(device_ctx_, runtime_bufs_[run.slot_id].data()); + } + } catch (...) { + std::lock_guard lk(native_run_mu_); + native_run_states_[run.slot_id].phase = phase; + throw; } - int wait_rc = state.wait_rc; int finalize_rc = finalize_run_fn_(device_ctx_, runtime_bufs_[run.slot_id].data()); - state = NativeRunSlotState{}; + { + std::lock_guard lk(native_run_mu_); + native_run_states_[run.slot_id] = NativeRunSlotState{}; + } int rc = finalize_rc != 0 ? finalize_rc : wait_rc; if (rc != 0) { - throw std::runtime_error("finalize_native_run failed with code " + std::to_string(rc)); + throw std::runtime_error( + "finalize_native_run failed with code " + std::to_string(rc) + " " + format_native_run_identity(run) + ); } } @@ -755,15 +877,24 @@ void ChipWorker::cleanup_native_runs_noexcept() noexcept { return; } for (size_t slot_id = 0; slot_id < runtime_bufs_.size(); ++slot_id) { - NativeRunSlotState &state = native_run_states_[slot_id]; - if (state.phase == NativeRunPhase::EMPTY) { - continue; + NativeRunPhase phase; + { + std::lock_guard lk(native_run_mu_); + NativeRunSlotState &state = native_run_states_[slot_id]; + phase = state.phase; + if (phase == NativeRunPhase::EMPTY) continue; + if (phase == NativeRunPhase::PREPARING || phase == NativeRunPhase::LAUNCHING || + phase == NativeRunPhase::FINALIZING) { + continue; + } + state.phase = NativeRunPhase::FINALIZING; } - if (state.phase == NativeRunPhase::LAUNCHED && wait_run_fn_ != nullptr) { + if (phase == NativeRunPhase::LAUNCHED && wait_run_fn_ != nullptr) { (void)wait_run_fn_(device_ctx_, runtime_bufs_[slot_id].data()); } (void)finalize_run_fn_(device_ctx_, runtime_bufs_[slot_id].data()); - state = NativeRunSlotState{}; + std::lock_guard lk(native_run_mu_); + native_run_states_[slot_id] = NativeRunSlotState{}; } } diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 5582237ce0..36ab118e38 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -33,6 +34,9 @@ struct ChipWorkerNativeRun { uint64_t generation{0}; // Process-unique identity for exactly one prepare attempt. uint64_t run_epoch{0}; + // Optional L3 mailbox identity. Direct L2 callers leave these as zero. + uint64_t run_id{0}; + uint64_t dispatch_id{0}; }; class ChipWorker { @@ -114,17 +118,20 @@ class ChipWorker { * the token, and the caller must still finalize it. The blocking composition * performs that cleanup internally on every exit. * - * The current backend keeps execution-only state on DeviceRunner, so it - * admits one prepared/active native run at a time. The - * slot/lease-generation/process-unique-run-epoch token prevents a delayed - * phase call from touching reused storage, including another run under the - * same pipeline lease or on another ChipWorker. + * Onboard HBG may prepare one distinct-slot successor while another run + * owns the execution claim. Diagnostics and backends without the explicit + * capability remain depth-one. The slot/lease-generation/process-unique- + * run-epoch token prevents a delayed phase call from touching reused + * storage, including another run under the same pipeline lease or on + * another ChipWorker. */ ChipWorkerNativeRun prepare_native_run( - int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, const PipelineSlotLease &lease + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, const PipelineSlotLease &lease, + uint64_t run_id = 0, uint64_t dispatch_id = 0 ); ChipWorkerNativeRun prepare_native_run( - int32_t callable_id, TaskArgsView args, const CallConfig &config, const PipelineSlotLease &lease + int32_t callable_id, TaskArgsView args, const CallConfig &config, const PipelineSlotLease &lease, + uint64_t run_id = 0, uint64_t dispatch_id = 0 ); void launch_native_run( const ChipWorkerNativeRun &run, volatile int32_t *accepted_state = nullptr, int32_t accepted_value = 0 @@ -207,6 +214,7 @@ class ChipWorker { bool initialized() const { return initialized_; } unsigned pipeline_depth() const { return pipeline_contract_.pipeline_depth; } size_t runtime_slot_count() const { return runtime_bufs_.size(); } + bool supports_concurrent_native_prepare() const; /// Opaque host native-run storage address for every slot the contract /// asked for. Two slots hold distinct storage; tests read this to prove @@ -242,9 +250,11 @@ class ChipWorker { using SimplerRunFn = int (*)(void *, void *, int32_t, const void *, const CallConfig *); using SimplerPrepareRunFn = int (*)(void *, void *, int32_t, const void *, const CallConfig *); using SimplerNativeRunFn = int (*)(void *, void *); + using SupportsConcurrentNativePrepareFn = int (*)(void *); using SetTaskAcceptedStateFn = int (*)(void *, volatile int32_t *, int32_t); using SelectPipelineSlotFn = int (*)(void *, uint32_t); using SelectArenaBankFn = int (*)(void *, uint32_t); + using SetNativeRunIdentityFn = int (*)(void *, uint64_t, uint64_t, uint64_t, uint64_t); using GetArenaBankGmHeapBaseFn = uint64_t (*)(void *, uint32_t); using GetRetainedTempAddrFn = uint64_t (*)(void *, uint32_t); using GetPipelineContractFn = const PipelineContract *(*)(); @@ -301,9 +311,11 @@ class ChipWorker { SimplerNativeRunFn poll_run_fn_ = nullptr; SimplerNativeRunFn wait_run_fn_ = nullptr; SimplerNativeRunFn finalize_run_fn_ = nullptr; + SupportsConcurrentNativePrepareFn supports_concurrent_native_prepare_fn_ = nullptr; SetTaskAcceptedStateFn set_task_accepted_state_fn_ = nullptr; SelectPipelineSlotFn select_pipeline_slot_fn_ = nullptr; SelectArenaBankFn select_arena_bank_fn_ = nullptr; + SetNativeRunIdentityFn set_native_run_identity_fn_ = nullptr; GetArenaBankGmHeapBaseFn get_arena_bank_gm_heap_base_fn_ = nullptr; GetRetainedTempAddrFn get_retained_temp_addr_fn_ = nullptr; SimplerUnregisterCallableFn unregister_callable_fn_ = nullptr; @@ -339,18 +351,18 @@ class ChipWorker { ); uint32_t select_slot_resources(uint32_t slot_id); - enum class NativeRunPhase : uint8_t { EMPTY, PREPARED, LAUNCHED, REAPED }; + enum class NativeRunPhase : uint8_t { EMPTY, PREPARING, PREPARED, LAUNCHING, LAUNCHED, REAPED, FINALIZING }; struct NativeRunSlotState { uint64_t lease_generation{0}; uint64_t run_epoch{0}; NativeRunPhase phase{NativeRunPhase::EMPTY}; int wait_rc{0}; + bool permits_prepared_successor{false}; }; ChipWorkerNativeRun prepare_native_run_on_slot( int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, uint32_t slot_id, - uint64_t generation + uint64_t generation, uint64_t run_id, uint64_t dispatch_id ); - NativeRunSlotState &require_native_run(const ChipWorkerNativeRun &run, NativeRunPhase first, NativeRunPhase second); void cleanup_native_runs_noexcept() noexcept; class RuntimeStorage { @@ -375,6 +387,7 @@ class ChipWorker { // prepare/finalize even if the owning vector itself is moved. std::vector runtime_bufs_; std::array native_run_states_{}; + mutable std::mutex native_run_mu_; PipelineSlotGenerationFilter pipeline_generations_; PipelineContract pipeline_contract_{PTO_PIPELINE_CONTRACT_ABI_VERSION, 0, 1, {}}; // device_id_ is set once in init() and never modified afterward. All diff --git a/src/common/worker/native_run_state.h b/src/common/worker/native_run_state.h index f21283dd9f..831c4f9e75 100644 --- a/src/common/worker/native_run_state.h +++ b/src/common/worker/native_run_state.h @@ -30,9 +30,9 @@ enum class NativeRunPhase : uint8_t { }; /** - * Caller-owned state for one progressable native lifecycle. Runtime and - * CallConfig are per-run; Runner-owned streams, diagnostics, and timing require - * exclusive ownership from prepare through finalize. + * Caller-owned state for one progressable native lifecycle. Runtime, + * CallConfig, and resource selection are per-run. Runner-owned execution, + * diagnostics, and timing remain exclusive from launch through finalize. */ template struct NativeRunState { @@ -70,6 +70,15 @@ struct NativeRunState { std::atomic phase{NativeRunPhase::Prepared}; NativeRunLaunchSignal launch_signal{}; void *host_thread_state{nullptr}; + uint64_t run_id{0}; + uint64_t generation{0}; + uint64_t dispatch_id{0}; + uint64_t run_epoch{0}; + uint32_t pipeline_slot{0}; + uint32_t arena_bank{0}; + char trace_attrs[192]{}; + bool runner_resources_owned{false}; + bool runner_reserved{false}; bool runner_claimed{false}; }; diff --git a/src/common/worker/pto_runtime_c_api.h b/src/common/worker/pto_runtime_c_api.h index 1e6757618d..179d4f3222 100644 --- a/src/common/worker/pto_runtime_c_api.h +++ b/src/common/worker/pto_runtime_c_api.h @@ -31,14 +31,18 @@ * get_aicpu_dlopen_count, get_host_dlopen_count, * get_run_stream_set_create_count, * simpler_provision_dma_workspace + * - pipeline: get_pipeline_contract, select_pipeline_slot_ctx, + * select_arena_bank_ctx, + * supports_concurrent_native_prepare_ctx, + * set_native_run_identity_ctx, + * get_arena_bank_gm_heap_base_ctx, + * get_retained_temp_addr_ctx, + * set_task_accepted_state_ctx * - ACL/stream: ensure_acl_ready_ctx, create_comm_stream_ctx, * destroy_comm_stream_ctx * - comm: comm_init, comm_alloc_windows, comm_get_local_window_base, * comm_get_window_size, comm_barrier, comm_destroy * - * Optional metadata: - * - pipeline: get_pipeline_contract - * * Native-run storage: caller allocates at least get_runtime_size() bytes with * get_runtime_alignment() alignment, zero-initializes it before first use, and * keeps its address stable from prepare through finalize. After finalize the @@ -163,7 +167,7 @@ typedef struct PipelineSlotLease { * Public API (resolved by ChipWorker via dlsym) * =========================================================================== */ -/** Return this runtime's immutable pipeline resource declaration. Optional. */ +/** Return this runtime's immutable pipeline resource declaration. */ const PipelineContract *get_pipeline_contract(void); /** @@ -333,6 +337,12 @@ int simpler_prepare_run( DeviceContextHandle ctx, RuntimeHandle runtime, int32_t callable_id, const void *args, const CallConfig *config ); +/** + * Return nonzero when non-diagnostic preparation may overlap the execution + * claim held by a run in another pipeline slot. + */ +int supports_concurrent_native_prepare_ctx(DeviceContextHandle ctx); + /** * Launch a prepared run. Returns only after the platform has published its * real kernel-launch marker, or after execution terminates before that marker. @@ -361,6 +371,11 @@ int select_pipeline_slot_ctx(DeviceContextHandle ctx, uint32_t slot_id); /** Select the HOST_PER_RUN arena bank used by the next synchronous run. */ int select_arena_bank_ctx(DeviceContextHandle ctx, uint32_t bank_id); +/** Attach optional L3 identity metadata to the current thread's next native prepare. */ +int set_native_run_identity_ctx( + DeviceContextHandle ctx, uint64_t run_id, uint64_t generation, uint64_t dispatch_id, uint64_t run_epoch +); + /** * Committed GM heap base of one arena bank, or 0 when that bank has never been * committed or the platform keeps a single shared arena set. Reports which @@ -377,8 +392,7 @@ uint64_t get_retained_temp_addr_ctx(DeviceContextHandle ctx, uint32_t slot_id); /** * Bind an optional host state word that the runner publishes after both device - * kernels have been enqueued. Onboard runtimes may export this symbol; callers - * must fall back to completion when it is absent. + * kernels have been enqueued. Passing NULL clears the binding. */ int set_task_accepted_state_ctx(DeviceContextHandle ctx, volatile int32_t *state, int32_t accepted_value); diff --git a/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py b/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py index ea2a35fc9f..022f697623 100644 --- a/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py +++ b/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py @@ -9,6 +9,8 @@ # ----------------------------------------------------------------------------------------------------------- """End-to-end validation of the B3a prepare/launch/poll/wait/finalize seam.""" +import tempfile + import pytest import torch from simpler.task_interface import ArgDirection as D @@ -72,7 +74,8 @@ def test_run(self, st_platform, st_worker, request, capfd): spans = list(parse_spans(capfd.readouterr().err.splitlines())) invocations = [inv for inv in group_invocations(spans) if "simpler_run" in inv.by_name()] - assert len(invocations) == 3, "abandoned, direct, and blocking runs must each emit one trace invocation" + expected_invocations = 3 if st_platform.endswith("sim") else 7 + assert len(invocations) == expected_invocations common_depths = { "simpler_run": 0, @@ -100,9 +103,18 @@ def test_run(self, st_platform, st_worker, request, capfd): for name in expected_depths.keys() - {"simpler_run", "simpler_run.runner_run.device_wall"}: stage = by_name[name] assert root.ts <= stage.ts <= stage.ts + stage.dur <= root_end - assert launched_count == 2 + expected_launched = 2 if st_platform.endswith("sim") else 6 + assert launched_count == expected_launched + if not st_platform.endswith("sim"): + root_attrs = [inv.by_name()["simpler_run"].attrs for inv in invocations] + assert all( + key in attrs + for attrs in root_attrs + for key in ("run_id=", "slot=", "generation=", "dispatch_id=", "run_epoch=") + ) + assert any("slot=1" in attrs and "generation=1" in attrs for attrs in root_attrs) - def _run_and_validate_l2( # noqa: PLR0913 + def _run_and_validate_l2( # noqa: PLR0913, PLR0915 -- lifecycle contract is intentionally sequential self, worker, callable_obj, @@ -121,8 +133,10 @@ def _run_and_validate_l2( # noqa: PLR0913 config = self._build_config(case["config"]) chip_worker = worker._chip_worker assert chip_worker is not None + supports_concurrent_prepare = bool(chip_worker._impl.supports_concurrent_native_prepare) chip_worker._register_callable_at_slot(_SLOT, callable_obj) native_run = None + successor_run = None try: test_args = self.generate_args(case["params"]) chip_args, output_names = _build_chip_task_args(test_args, self.CALLABLE["orchestration"]["signature"]) @@ -134,9 +148,10 @@ def _run_and_validate_l2( # noqa: PLR0913 _SLOT, chip_args, _SLOT, _GENERATION, config=config ) first_run = native_run - assert chip_worker.run_stream_set_create_count == stream_count_before_prepare + expected_stream_count = stream_count_before_prepare + int(supports_concurrent_prepare) + assert chip_worker.run_stream_set_create_count == expected_stream_count assert torch.count_nonzero(test_args.out) == 0, "prepare crossed the device launch fence" - with pytest.raises(RuntimeError, match="unfinished native run|owns the runner"): + with pytest.raises(RuntimeError, match="unfinished native run|owns the runner|active predecessor"): chip_worker._prepare_native_run_with_pipeline_lease(_SLOT, chip_args, 1, _GENERATION, config=config) with pytest.raises(RuntimeError, match="unregister_callable failed"): chip_worker._unregister_slot(_SLOT) @@ -181,10 +196,100 @@ def _run_and_validate_l2( # noqa: PLR0913 self.compute_golden(second_golden, case["params"]) chip_worker._run_slot(_SLOT, second_chip_args, config=config) _compare_outputs(second_args, second_golden, second_output_names, self.RTOL, self.ATOL) + + if supports_concurrent_prepare: + + def build_run_args(): + run_args = self.generate_args(case["params"]) + run_chip_args, run_output_names = _build_chip_task_args( + run_args, self.CALLABLE["orchestration"]["signature"] + ) + run_golden = run_args.clone() + self.compute_golden(run_golden, case["params"]) + return run_args, run_chip_args, run_output_names, run_golden + + # The successor owns a distinct bank and fresh stream while A + # still owns the execution claim. A failed early launch must + # leave B prepared so the same token can launch after A's + # complete fence and finalization. + active_args, active_chip_args, active_outputs, active_golden = build_run_args() + successor_args, successor_chip_args, successor_outputs, successor_golden = build_run_args() + stream_count = chip_worker.run_stream_set_create_count + native_run = chip_worker._prepare_native_run_with_pipeline_lease( + _SLOT, active_chip_args, 0, _GENERATION + 1, config=config + ) + assert chip_worker.run_stream_set_create_count == stream_count + 1 + chip_worker._launch_native_run(native_run) + successor_run = chip_worker._prepare_native_run_with_pipeline_lease( + _SLOT, successor_chip_args, 1, _GENERATION, config=config + ) + assert chip_worker.run_stream_set_create_count == stream_count + 2 + bank0 = chip_worker.arena_bank_gm_heap_base(0) + bank1 = chip_worker.arena_bank_gm_heap_base(1) + assert bank0 != 0 + assert bank1 != 0 + assert bank0 != bank1 + assert torch.count_nonzero(successor_args.out) == 0 + + with pytest.raises(RuntimeError, match="launch_native_run failed") as claim_error: + chip_worker._launch_native_run(successor_run) + assert "slot=1" in str(claim_error.value) + assert "generation=1" in str(claim_error.value) + assert "run_epoch=" in str(claim_error.value) + + chip_worker._wait_native_run(native_run) + chip_worker._finalize_native_run(native_run) + native_run = None + _compare_outputs(active_args, active_golden, active_outputs, self.RTOL, self.ATOL) + + chip_worker._launch_native_run(successor_run) + chip_worker._wait_native_run(successor_run) + chip_worker._finalize_native_run(successor_run) + successor_run = None + _compare_outputs(successor_args, successor_golden, successor_outputs, self.RTOL, self.ATOL) + + with tempfile.TemporaryDirectory(prefix="simpler-native-diagnostics-") as output_dir: + diagnostic_config = self._build_config(case["config"]) + diagnostic_config.enable_dep_gen = True + diagnostic_config.output_prefix = output_dir + + # A diagnostic successor cannot overlap an ordinary active + # run, even though the predecessor otherwise permits one. + active_args, active_chip_args, active_outputs, active_golden = build_run_args() + native_run = chip_worker._prepare_native_run_with_pipeline_lease( + _SLOT, active_chip_args, 0, _GENERATION + 2, config=config + ) + chip_worker._launch_native_run(native_run) + with pytest.raises(RuntimeError, match="active predecessor"): + chip_worker._prepare_native_run_with_pipeline_lease( + _SLOT, successor_chip_args, 1, _GENERATION + 1, config=diagnostic_config + ) + chip_worker._wait_native_run(native_run) + chip_worker._finalize_native_run(native_run) + native_run = None + _compare_outputs(active_args, active_golden, active_outputs, self.RTOL, self.ATOL) + + # A diagnostic predecessor also cannot admit an ordinary + # successor while it owns the execution claim. + active_args, active_chip_args, active_outputs, active_golden = build_run_args() + native_run = chip_worker._prepare_native_run_with_pipeline_lease( + _SLOT, active_chip_args, 0, _GENERATION + 3, config=diagnostic_config + ) + chip_worker._launch_native_run(native_run) + with pytest.raises(RuntimeError, match="active predecessor"): + chip_worker._prepare_native_run_with_pipeline_lease( + _SLOT, successor_chip_args, 1, _GENERATION + 1, config=config + ) + chip_worker._wait_native_run(native_run) + chip_worker._finalize_native_run(native_run) + native_run = None + _compare_outputs(active_args, active_golden, active_outputs, self.RTOL, self.ATOL) finally: - if native_run is not None: + for unfinished_run in (successor_run, native_run): + if unfinished_run is None: + continue try: - chip_worker._finalize_native_run(native_run) + chip_worker._finalize_native_run(unfinished_run) except Exception: pass chip_worker._unregister_slot(_SLOT) diff --git a/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py b/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py index 3d5d4d3ea5..b3fc24d060 100644 --- a/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py +++ b/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py @@ -24,7 +24,6 @@ _FRAME_STAGED, _OFF_ACCEPTED, _OFF_STATE, - _TASK_LAUNCHED, MAILBOX_FRAME_SIZE, _mailbox_load_i32, ) @@ -198,8 +197,6 @@ def second_graph(orch, _args, _cfg): shm_buf = st_worker._chip_shms[0].buf assert shm_buf is not None mailbox_addr = ctypes.addressof(ctypes.c_char.from_buffer(shm_buf)) - predecessor_frame_addr = mailbox_addr + MAILBOX_FRAME_SIZE - predecessor_state_addr = predecessor_frame_addr + _OFF_STATE successor_frame_addr = mailbox_addr + 2 * MAILBOX_FRAME_SIZE successor_state_addr = successor_frame_addr + _OFF_STATE successor_accepted_addr = successor_frame_addr + _OFF_ACCEPTED @@ -211,15 +208,12 @@ def second_graph(orch, _args, _cfg): accepted = _mailbox_load_i32(successor_accepted_addr) if state == _FRAME_STAGED: assert accepted == 0, "the successor crossed its launch fence before activation" - assert _mailbox_load_i32(predecessor_state_addr) == _TASK_LAUNCHED, ( - "the successor staged only after the predecessor native run had already terminalized" - ) saw_staged = True break - assert accepted == 0, "the successor launched while its predecessor was still active" + assert accepted == 0, "the successor crossed FIFO activation while its predecessor run remained open" time.sleep(0.001) - assert saw_staged, "the successor did not reach FRAME_STAGED behind its predecessor" + assert saw_staged, "the successor did not reach FRAME_STAGED while its predecessor run remained open" assert not first.done, "the predecessor escaped its SubTask fence" _SUB_RELEASE.set() diff --git a/tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/aiv/delayed_add.cpp b/tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/aiv/delayed_add.cpp new file mode 100644 index 0000000000..aefd623626 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/aiv/delayed_add.cpp @@ -0,0 +1,66 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include +#include + +#include "tensor.h" + +#include "pipe_sync.h" + +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *src0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *src1_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + + // Match the repository's bounded slow-producer test shape. The caller uses + // a loop count calibrated to keep this task on-core long enough to observe + // successor preparation while remaining below the scheduler timeout, + // without an unbounded device-side handshake that teardown could not release. + const int32_t spin_iters = static_cast(args[3]); + volatile int32_t spin_accumulator = 0; + for (int32_t i = 0; i < spin_iters; ++i) { + ++spin_accumulator; + } + (void)spin_accumulator; + + __gm__ float *src0 = reinterpret_cast<__gm__ float *>(src0_tensor->buffer.addr) + src0_tensor->start_offset; + __gm__ float *src1 = reinterpret_cast<__gm__ float *>(src1_tensor->buffer.addr) + src1_tensor->start_offset; + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + constexpr int kTRows = 128; + constexpr int kTCols = 128; + using DynShapeDim5 = pto::Shape<1, 1, 1, kTRows, kTCols>; + using DynStridDim5 = pto::Stride<1, 1, 1, kTCols, 1>; + using GlobalData = pto::GlobalTensor; + using TileData = pto::Tile; + + TileData src0_tile(kTRows, kTCols); + TileData src1_tile(kTRows, kTCols); + TileData dst_tile(kTRows, kTCols); + TASSIGN(src0_tile, 0x0); + TASSIGN(src1_tile, 0x10000); + TASSIGN(dst_tile, 0x20000); + + GlobalData src0_global(src0); + GlobalData src1_global(src1); + GlobalData dst_global(out); + + TLOAD(src0_tile, src0_global); + TLOAD(src1_tile, src1_global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(dst_tile, src0_tile, src1_tile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(dst_global, dst_tile); + pipe_sync(); +} diff --git a/tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/orchestration/pipelined_vector_orch.cpp b/tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/orchestration/pipelined_vector_orch.cpp new file mode 100644 index 0000000000..85c7aa66fc --- /dev/null +++ b/tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/orchestration/pipelined_vector_orch.cpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +namespace { + +constexpr uint64_t kDelayedAdd = 0; +constexpr uint64_t kAddScalar = 1; +constexpr int kChainLength = 512; + +} // namespace + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &args) { + (void)args; + return PTO2OrchestrationConfig{.expected_arg_count = 4}; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &args) { + const Tensor &a = args.tensor(0).ref(); + const Tensor &b = args.tensor(1).ref(); + const Tensor &out = args.tensor(2).ref(); + const uint64_t spin_iters = args.scalar(0); + uint32_t shape[1] = {a.shapes[0]}; + TensorCreateInfo temporary(shape, 1, DataType::FLOAT32); + + L0TaskArgs add_args; + add_args.add_input(a); + add_args.add_input(b); + add_args.add_output(temporary); + add_args.add_scalar(spin_iters); + TaskOutputTensors add_outputs = rt_submit_aiv_task(kDelayedAdd, add_args); + Tensor current = add_outputs.get_ref(0); + + union { + float f32; + uint64_t u64; + } scalar{}; + scalar.f32 = 1.0F; + for (int i = 0; i < kChainLength; ++i) { + L0TaskArgs step_args; + step_args.add_input(current); + if (i + 1 == kChainLength) { + step_args.add_output(out); + step_args.add_scalar(scalar.u64); + rt_submit_aiv_task(kAddScalar, step_args); + } else { + step_args.add_output(temporary); + step_args.add_scalar(scalar.u64); + TaskOutputTensors step_outputs = rt_submit_aiv_task(kAddScalar, step_args); + current = step_outputs.get_ref(0); + } + } +} + +} // extern "C" diff --git a/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py b/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py index 5c22e0dd69..ad3b8f5c3e 100644 --- a/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py +++ b/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py @@ -9,13 +9,14 @@ # ----------------------------------------------------------------------------------------------------------- """Onboard validation for bounded whole-run FIFO admission. -The first run completes real NPU work but remains active behind a SubTask -fence. The second run builds its graph into the other pipeline slot and must -not dispatch until the first run becomes terminal. A third submission must -block before its graph callback while both slots are admitted. +The second run builds and backend-prepares in the other pipeline slot while the +first run executes a bounded delayed NPU chain, but cannot launch until the first run's +SubTask fence releases. A third submission blocks before its graph callback +while both slots are admitted. """ import atexit +import ctypes import tempfile import threading import time @@ -26,11 +27,22 @@ import pytest import torch from simpler.task_interface import ArgDirection as D - -from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler.worker import ( + _FRAME_STAGED, + _OFF_ACCEPTED, + _OFF_STATE, + _TASK_LAUNCHED, + MAILBOX_FRAME_SIZE, + _mailbox_load_i32, +) + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test from simpler_setup.scene_test import _build_l3_task_args _VECTOR_KERNELS = "../vector_example/kernels" +_PIPELINED_VECTOR_ORCH = "kernels/orchestration/pipelined_vector_orch.cpp" +_CHAIN_LENGTH = 512 +_DEVICE_SPIN_ITERS = 200_000_000 _SIZE = 128 * 128 @@ -75,6 +87,36 @@ def _wait_for_release(_args): raise RuntimeError("whole-run FIFO test timed out waiting for the release fence") +def _wait_for_backend_prepared_successor(worker, timeout: float) -> None: + shm_buf = worker._chip_shms[0].buf # noqa: SLF001 -- white-box backend-prepare observation + assert shm_buf is not None + mailbox_addr = ctypes.addressof(ctypes.c_char.from_buffer(shm_buf)) + state_addrs = [mailbox_addr + (1 + index) * MAILBOX_FRAME_SIZE + _OFF_STATE for index in range(2)] + accepted_addrs = [mailbox_addr + (1 + index) * MAILBOX_FRAME_SIZE + _OFF_ACCEPTED for index in range(2)] + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + states = [_mailbox_load_i32(addr) for addr in state_addrs] + for index, state in enumerate(states): + if state == _FRAME_STAGED and states[1 - index] == _TASK_LAUNCHED: + assert _mailbox_load_i32(accepted_addrs[index]) == 0 + return + time.sleep(0.001) + raise AssertionError("the successor did not finish backend preparation while its predecessor was launched") + + +def _wait_for_active_device_run(worker, timeout: float) -> None: + shm_buf = worker._chip_shms[0].buf # noqa: SLF001 -- white-box device-launch observation + assert shm_buf is not None + mailbox_addr = ctypes.addressof(ctypes.c_char.from_buffer(shm_buf)) + state_addrs = [mailbox_addr + (1 + index) * MAILBOX_FRAME_SIZE + _OFF_STATE for index in range(2)] + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if any(_mailbox_load_i32(addr) == _TASK_LAUNCHED for addr in state_addrs): + return + time.sleep(0.001) + raise AssertionError("the predecessor did not reach its device launch fence") + + @scene_test(level=3, runtime="host_build_graph") class TestWorkerAsyncWholeRunFifo(SceneTestCase): """A prepared run may build ahead but cannot dispatch ahead.""" @@ -84,14 +126,14 @@ class TestWorkerAsyncWholeRunFifo(SceneTestCase): { "name": "vector", "orchestration": { - "source": f"{_VECTOR_KERNELS}/orchestration/example_orch.cpp", + "source": _PIPELINED_VECTOR_ORCH, "function_name": "aicpu_orchestration_entry", "signature": [D.IN, D.IN, D.OUT], }, "incores": [ { "func_id": 0, - "source": f"{_VECTOR_KERNELS}/aiv/kernel_add.cpp", + "source": "kernels/aiv/delayed_add.cpp", "core_type": "aiv", "signature": [D.IN, D.IN, D.OUT], }, @@ -101,12 +143,6 @@ class TestWorkerAsyncWholeRunFifo(SceneTestCase): "core_type": "aiv", "signature": [D.IN, D.OUT], }, - { - "func_id": 2, - "source": f"{_VECTOR_KERNELS}/aiv/kernel_mul.cpp", - "core_type": "aiv", - "signature": [D.IN, D.IN, D.OUT], - }, ], }, {"name": "wait_for_release", "callable": _wait_for_release}, @@ -203,7 +239,12 @@ def test_prepared_run_device_control_waits_for_the_active_run(self, st_platform, sub_handle = type(self)._st_sub_handles["wait_for_release"] def first_graph(orch, _args, _cfg): - builder = TaskArgsBuilder(Tensor("a", first_a), Tensor("b", first_b), Tensor("f", first_out)) + builder = TaskArgsBuilder( + Tensor("a", first_a), + Tensor("b", first_b), + Tensor("f", first_out), + Scalar("spin_iters", 0), + ) chip_args, _ = _build_l3_task_args(builder, vector_signature) orch.submit_next_level(vector_handle, chip_args, self._build_config(self.CASES[0]["config"]), worker=0) orch.submit_sub(sub_handle) @@ -220,9 +261,13 @@ def second_graph(orch, _args, _cfg): control_returned.set() orch.free(0, ptr) - submitter = threading.Thread( - target=lambda: result.setdefault("handle", st_worker.submit(second_graph)), daemon=True - ) + def submit_second(): + try: + result["handle"] = st_worker.submit(second_graph) + except BaseException as error: # noqa: BLE001 -- re-raised on the test thread + result["error"] = error + + submitter = threading.Thread(target=submit_second, daemon=True) submitter.start() assert entered_callback.wait(10.0), "the prepared run's graph callback never entered" @@ -232,9 +277,11 @@ def second_graph(orch, _args, _cfg): _SUB_RELEASE.set() first.wait(30.0) - assert control_returned.wait(30.0), "device control stayed blocked after the active run became terminal" submitter.join(30.0) assert not submitter.is_alive() + if "error" in result: + raise result["error"] + assert control_returned.is_set(), "device control stayed blocked after the active run became terminal" result["handle"].wait(30.0) finally: _SUB_RELEASE.set() @@ -291,9 +338,13 @@ def first_graph(orch, _args, _cfg): def second_graph(_orch, _args, _cfg): entered_callback.set() - submitter = threading.Thread( - target=lambda: result.setdefault("handle", st_worker.submit(second_graph)), daemon=True - ) + def submit_second(): + try: + result["handle"] = st_worker.submit(second_graph) + except BaseException as error: # noqa: BLE001 -- re-raised on the test thread + result["error"] = error + + submitter = threading.Thread(target=submit_second, daemon=True) submitter.start() assert not entered_callback.wait(2.0), ( @@ -302,9 +353,11 @@ def second_graph(_orch, _args, _cfg): _SUB_RELEASE.set() first.wait(30.0) - assert entered_callback.wait(30.0), "the successor stayed blocked after its predecessor's cleanup ran" submitter.join(30.0) assert not submitter.is_alive() + if "error" in result: + raise result["error"] + assert entered_callback.is_set(), "the successor stayed blocked after its predecessor's cleanup ran" result["handle"].wait(30.0) finally: _SUB_RELEASE.set() @@ -334,29 +387,31 @@ def test_run(self, st_platform, st_worker): buffers.append(buffer) tensors.append(tensor) first_a, first_b, first_out, second_a, second_b, second_out = tensors - vector_handle = type(self)._st_chip_handles["vector"] vector_signature = type(self)._st_chip_handles["vector_sig"] sub_handle = type(self)._st_sub_handles["wait_for_release"] - def submit_vector(orch, a, b, out, *, hold_open=False): - builder = TaskArgsBuilder(Tensor("a", a), Tensor("b", b), Tensor("f", out)) + def submit_vector(orch, a, b, out, *, spin_iters=0, hold_open=False): + builder = TaskArgsBuilder( + Tensor("a", a), Tensor("b", b), Tensor("f", out), Scalar("spin_iters", spin_iters) + ) chip_args, _ = _build_l3_task_args(builder, vector_signature) orch.submit_next_level(vector_handle, chip_args, self._build_config(self.CASES[0]["config"]), worker=0) if hold_open: orch.submit_sub(sub_handle) first = st_worker.submit( - lambda orch, _args, _cfg: submit_vector(orch, first_a, first_b, first_out, hold_open=True) + lambda orch, _args, _cfg: submit_vector( + orch, first_a, first_b, first_out, spin_iters=_DEVICE_SPIN_ITERS, hold_open=True + ) ) + first_expected = first_a + first_b + _CHAIN_LENGTH + # Run-level acceptance also includes the intentionally blocked SUB + # task, whose compatibility endpoint acknowledges only on return. + # Observe the chip frame directly so that fence cannot postpone the + # successor submission until its 30-second timeout. + _wait_for_active_device_run(st_worker, 10.0) assert _SUB_ENTERED.wait(10.0), "the first run's SubTask did not start" - - first_expected = (first_a + first_b + 1) * (first_a + first_b + 2) - deadline = time.monotonic() + 10.0 - while not torch.allclose(first_out, first_expected) and time.monotonic() < deadline: - time.sleep(0.001) - assert torch.allclose(first_out, first_expected), "the first run's NPU task did not complete" - second_graph_done = threading.Event() def second_graph(orch, _args, _cfg): @@ -365,6 +420,10 @@ def second_graph(orch, _args, _cfg): second = st_worker.submit(second_graph) assert second_graph_done.is_set(), "the second run did not build ahead" + _wait_for_backend_prepared_successor(st_worker, 10.0) + assert torch.count_nonzero(first_out).item() == 0, ( + "successor backend preparation did not finish during predecessor device execution" + ) assert torch.count_nonzero(second_out).item() == 0, ( "the prepared run dispatched before the active run ended" ) @@ -372,9 +431,13 @@ def second_graph(orch, _args, _cfg): def third_graph(_orch, _args, _cfg): third_callback.set() - submitter = threading.Thread( - target=lambda: third_result.setdefault("handle", st_worker.submit(third_graph)), daemon=True - ) + def submit_third(): + try: + third_result["handle"] = st_worker.submit(third_graph) + except BaseException as error: # noqa: BLE001 -- re-raised on the test thread + third_result["error"] = error + + submitter = threading.Thread(target=submit_third, daemon=True) submitter.start() assert not third_callback.wait(0.1), "the third graph callback entered before admission capacity was free" @@ -389,13 +452,15 @@ def third_graph(_orch, _args, _cfg): _SUB_RELEASE.set() first.wait(10.0) - assert third_callback.wait(10.0), "the third submission did not enter after the first run freed its slot" - second.wait(10.0) - second_expected = (second_a + second_b + 1) * (second_a + second_b + 2) - assert torch.allclose(second_out, second_expected), "the prepared run did not execute correctly on the NPU" - submitter.join(10.0) assert not submitter.is_alive() + if "error" in third_result: + raise third_result["error"] + assert third_callback.is_set(), "the third submission did not enter after the first run freed its slot" + second.wait(10.0) + second_expected = second_a + second_b + _CHAIN_LENGTH + assert torch.allclose(first_out, first_expected), "the first run did not execute correctly on the NPU" + assert torch.allclose(second_out, second_expected), "the prepared run did not execute correctly on the NPU" third_result["handle"].wait(10.0) finally: _SUB_RELEASE.set() diff --git a/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp b/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp index b40aae4853..1a01c91b90 100644 --- a/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp +++ b/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp @@ -35,7 +35,7 @@ TEST(PipelineContract, AcceptsADeclarationThisBuildCanHonor) { EXPECT_TRUE(is_valid_pipeline_contract(&c)); } -// A runtime that exports no contract is handled by the caller, not here. +// Runtime loading requires a contract; this helper still rejects malformed null values. TEST(PipelineContract, RejectsNull) { EXPECT_FALSE(is_valid_pipeline_contract(nullptr)); } TEST(PipelineContract, AcceptsAnEmptyResourceList) { diff --git a/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp b/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp index 250b645663..85358178b5 100644 --- a/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp +++ b/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp @@ -131,6 +131,32 @@ TEST(RunStreamSlots, AStrandedSlotDoesNotBlockItsPeer) { EXPECT_EQ(slots.retire_aicore(1), 0); } +TEST(RunStreamSlots, SuccessorProvisionAndAbandonDoNotTouchTheActiveSlot) { + FakeStreams fake; + RunStreamSlots slots = make_slots(fake); + + ASSERT_EQ(slots.acquire(0), 0); + void *active_aicore = slots.aicore(0); + + // Native prepare provisions the successor before the predecessor retires. + ASSERT_EQ(slots.acquire(1), 0); + EXPECT_TRUE(slots.ready(0)); + EXPECT_TRUE(slots.ready(1)); + EXPECT_NE(slots.aicore(1), active_aicore); + EXPECT_EQ(slots.created_count(), 2u); + + // Abandoning the unpublished successor retires only its fresh AICore + // stream. Its slot-persistent AICPU stream remains reusable. + ASSERT_EQ(slots.retire_aicore(1), 0); + EXPECT_TRUE(slots.ready(0)); + EXPECT_EQ(slots.aicore(0), active_aicore); + EXPECT_EQ(slots.aicore(1), nullptr); + EXPECT_NE(slots.aicpu(1), nullptr); + + EXPECT_EQ(slots.retire_aicore(0), 0); + EXPECT_EQ(slots.destroy_all(), 0); +} + // destroy_all reports the first failure and keeps only what it could not free, // so a second teardown attempt is still meaningful. TEST(RunStreamSlots, DestroyAllReportsFailureAndRetriesWhatSurvived) { diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 044641ed11..6f9bb1afd1 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -1051,6 +1052,88 @@ TEST(WorkerManagerTest, WorkerThreadUsesOneProgressOwnerForActiveAndStagedLanes) allocator.shutdown(); } +TEST(WorkerManagerTest, AdmissionRejectionsCompleteClaimedDispatchesWithoutThrowing) { + Ring allocator; + allocator.init(/*heap_bytes=*/0); + TaskSlot active_slot = make_progress_slot(allocator, /*run_id=*/13, /*pipeline_slot=*/0, /*generation=*/1); + TaskSlot capacity_rejected = make_progress_slot(allocator, /*run_id=*/14, /*pipeline_slot=*/1, /*generation=*/1); + TaskSlot lane_rejected = make_progress_slot(allocator, /*run_id=*/15, /*pipeline_slot=*/1, /*generation=*/2); + TaskSlot stopping_rejected = make_progress_slot(allocator, /*run_id=*/16, /*pipeline_slot=*/0, /*generation=*/2); + ASSERT_NE(active_slot, INVALID_SLOT); + ASSERT_NE(capacity_rejected, INVALID_SLOT); + ASSERT_NE(lane_rejected, INVALID_SLOT); + ASSERT_NE(stopping_rejected, INVALID_SLOT); + + WorkerThread worker; + auto endpoint = std::make_unique(/*worker_id=*/0, /*max_inflight_tasks=*/1); + DeterministicProgressEndpoint *endpoint_ptr = endpoint.get(); + std::mutex callback_mu; + std::condition_variable callback_cv; + std::vector accepted; + std::vector completed; + worker.start( + &allocator, + [&](WorkerCompletion completion) { + std::lock_guard lk(callback_mu); + completed.push_back(std::move(completion)); + callback_cv.notify_all(); + }, + [&](WorkerDispatch dispatch) { + std::lock_guard lk(callback_mu); + accepted.push_back(dispatch); + callback_cv.notify_all(); + }, + {}, std::move(endpoint) + ); + + worker.dispatch(WorkerDispatch{active_slot, 0}); + ASSERT_TRUE(endpoint_ptr->wait_submitted(1)); + + EXPECT_NO_THROW(worker.dispatch_prepared(WorkerDispatch{capacity_rejected, 0})); + EXPECT_TRUE(worker.can_stage()); + EXPECT_NO_THROW(worker.dispatch(WorkerDispatch{lane_rejected, 0})); + { + std::unique_lock lk(callback_mu); + ASSERT_TRUE(callback_cv.wait_for(lk, std::chrono::seconds(3), [&] { + return accepted.size() == 2 && completed.size() == 2; + })); + std::set accepted_slots; + std::set completed_slots; + for (const WorkerDispatch &dispatch : accepted) { + accepted_slots.insert(dispatch.task_slot); + EXPECT_EQ(dispatch.prepare_only, dispatch.task_slot == capacity_rejected); + } + for (const WorkerCompletion &completion : completed) { + EXPECT_EQ(completion.outcome, EndpointOutcome::ENDPOINT_FAILURE); + completed_slots.insert(completion.task_slot); + } + EXPECT_EQ(accepted_slots, (std::set{capacity_rejected, lane_rejected})); + EXPECT_EQ(completed_slots, accepted_slots); + } + EXPECT_EQ(endpoint_ptr->submitted().size(), 1u); + + WorkerDispatch active = endpoint_ptr->submitted().front(); + endpoint_ptr->emit(WorkerProgressKind::ACCEPTED, active); + endpoint_ptr->emit(WorkerProgressKind::COMPLETED, active); + { + std::unique_lock lk(callback_mu); + ASSERT_TRUE(callback_cv.wait_for(lk, std::chrono::seconds(3), [&] { + return accepted.size() == 3 && completed.size() == 3; + })); + } + worker.stop(); + + EXPECT_NO_THROW(worker.dispatch(WorkerDispatch{stopping_rejected, 0})); + { + std::lock_guard lk(callback_mu); + ASSERT_EQ(accepted.size(), 4u); + ASSERT_EQ(completed.size(), 4u); + EXPECT_EQ(completed.back().outcome, EndpointOutcome::ENDPOINT_FAILURE); + EXPECT_EQ(completed.back().task_slot, stopping_rejected); + } + allocator.shutdown(); +} + TEST(WorkerManagerTest, TwoFrameLeaseSlotsDoNotDefineFifoOrAcceptance) { alignas(8) std::array mailbox{}; Ring allocator; @@ -1111,9 +1194,11 @@ TEST(WorkerManagerTest, ThirdDispatchCannotMutateTwoOccupiedFrames) { TaskSlot active_slot = make_progress_slot(allocator, /*run_id=*/31, /*pipeline_slot=*/0, /*generation=*/1); TaskSlot staged_slot = make_progress_slot(allocator, /*run_id=*/32, /*pipeline_slot=*/1, /*generation=*/1); TaskSlot third_slot = make_progress_slot(allocator, /*run_id=*/33, /*pipeline_slot=*/0, /*generation=*/1); + TaskSlot fourth_slot = make_progress_slot(allocator, /*run_id=*/34, /*pipeline_slot=*/1, /*generation=*/2); ASSERT_NE(active_slot, INVALID_SLOT); ASSERT_NE(staged_slot, INVALID_SLOT); ASSERT_NE(third_slot, INVALID_SLOT); + ASSERT_NE(fourth_slot, INVALID_SLOT); WorkerThread worker; std::mutex completion_mu; @@ -1147,12 +1232,16 @@ TEST(WorkerManagerTest, ThirdDispatchCannotMutateTwoOccupiedFrames) { const uint64_t first_dispatch_id = test_frame_dispatch_id(frame0); const uint64_t second_dispatch_id = test_frame_dispatch_id(frame1); - EXPECT_THROW(worker.dispatch(WorkerDispatch{third_slot, 0}), std::logic_error); - EXPECT_THROW(worker.dispatch_prepared(WorkerDispatch{third_slot, 0}), std::logic_error); + EXPECT_NO_THROW(worker.dispatch(WorkerDispatch{third_slot, 0})); + EXPECT_NO_THROW(worker.dispatch_prepared(WorkerDispatch{fourth_slot, 0})); EXPECT_EQ(test_frame_state(frame0), MailboxState::TASK_READY); EXPECT_EQ(test_frame_state(frame1), MailboxState::PREPARE_READY); EXPECT_EQ(test_frame_dispatch_id(frame0), first_dispatch_id); EXPECT_EQ(test_frame_dispatch_id(frame1), second_dispatch_id); + { + std::lock_guard lk(completion_mu); + EXPECT_EQ(completion_count, 2); + } set_test_frame_accepted(frame0); set_test_frame_accepted(frame1); @@ -1161,7 +1250,7 @@ TEST(WorkerManagerTest, ThirdDispatchCannotMutateTwoOccupiedFrames) { { std::unique_lock lk(completion_mu); EXPECT_TRUE(completion_cv.wait_for(lk, std::chrono::seconds(3), [&] { - return completion_count == 2; + return completion_count == 4; })); } worker.stop(); @@ -1639,6 +1728,9 @@ TEST_F(ProgressSchedulerFixture, SuccessorStagesButActivatesOnlyAfterFifoPromoti SubmitResult first = orchestrator.submit_next_level(C(1), single_tensor_args(0x1000, TensorArgType::OUTPUT), config, 0); orchestrator.close_run_submission(first_run); + // Pin the production ordering: the predecessor is already active before + // closing the successor publishes the edge that makes it preparable. + ASSERT_TRUE(endpoint0->wait_submitted(1)); RunId second_run = orchestrator.begin_run(); SubmitResult second = orchestrator.submit_next_level(C(2), single_tensor_args(0x2000, TensorArgType::OUTPUT), config, 0); @@ -1674,6 +1766,46 @@ TEST_F(ProgressSchedulerFixture, SuccessorStagesButActivatesOnlyAfterFifoPromoti if (orchestrator.run_done(second_run)) orchestrator.release_run(second_run); } +TEST_F(ProgressSchedulerFixture, DiagnosticSuccessorWaitsForActiveLaneInsteadOfStaging) { + RunId first_run = orchestrator.begin_run(); + SubmitResult first = + orchestrator.submit_next_level(C(1), single_tensor_args(0x1100, TensorArgType::OUTPUT), config, 0); + orchestrator.close_run_submission(first_run); + ASSERT_TRUE(endpoint0->wait_submitted(1)); + + CallConfig diagnostic_config; + diagnostic_config.enable_dump_args = 1; + std::snprintf( + diagnostic_config.output_prefix, sizeof(diagnostic_config.output_prefix), "%s", + "/tmp/simpler-diagnostic-successor" + ); + RunId second_run = orchestrator.begin_run(); + SubmitResult second = + orchestrator.submit_next_level(C(2), single_tensor_args(0x2200, TensorArgType::OUTPUT), diagnostic_config, 0); + orchestrator.close_run_submission(second_run); + + EXPECT_FALSE(endpoint0->wait_submitted(2, std::chrono::milliseconds(50))); + std::vector submitted = endpoint0->submitted(); + ASSERT_EQ(submitted.size(), 1u); + EXPECT_EQ(submitted[0].task_slot, first.task_slot); + EXPECT_FALSE(submitted[0].prepare_only); + + endpoint0->emit(WorkerProgressKind::ACCEPTED, submitted[0]); + endpoint0->emit(WorkerProgressKind::COMPLETED, submitted[0]); + ASSERT_TRUE(endpoint0->wait_submitted(2)); + submitted = endpoint0->submitted(); + ASSERT_EQ(submitted.size(), 2u); + EXPECT_EQ(submitted[1].task_slot, second.task_slot); + EXPECT_FALSE(submitted[1].prepare_only); + + endpoint0->emit(WorkerProgressKind::ACCEPTED, submitted[1]); + endpoint0->emit(WorkerProgressKind::COMPLETED, submitted[1]); + EXPECT_TRUE(orchestrator.wait_run_for(first_run, 3.0)); + EXPECT_TRUE(orchestrator.wait_run_for(second_run, 3.0)); + if (orchestrator.run_done(first_run)) orchestrator.release_run(first_run); + if (orchestrator.run_done(second_run)) orchestrator.release_run(second_run); +} + TEST_F(ProgressSchedulerFixture, ActivatedRunReleasesTheStagingLaneForItsSuccessor) { RunId first_run = orchestrator.begin_run(); SubmitResult first = diff --git a/tests/ut/py/test_host_runtime_abi.py b/tests/ut/py/test_host_runtime_abi.py new file mode 100644 index 0000000000..670fe95e2b --- /dev/null +++ b/tests/ut/py/test_host_runtime_abi.py @@ -0,0 +1,71 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Uniform host-runtime pipeline-symbol contract tests.""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +import pytest + +_PROJECT_ROOT = Path(__file__).resolve().parents[3] +_NEWLY_REQUIRED_PIPELINE_SYMBOLS = { + "get_arena_bank_gm_heap_base_ctx", + "get_pipeline_contract", + "get_retained_temp_addr_ctx", + "set_native_run_identity_ctx", + "set_task_accepted_state_ctx", + "supports_concurrent_native_prepare_ctx", +} + +_SIM_CASES = [ + pytest.param(arch, "sim", runtime, id=f"{arch}-sim-{runtime}") + for arch in ("a2a3", "a5") + for runtime in ("host_build_graph", "tensormap_and_ringbuffer") +] +_ONBOARD_CASES = [ + pytest.param( + arch, + "onboard", + runtime, + id=f"{arch}-onboard-{runtime}", + marks=[pytest.mark.requires_hardware, pytest.mark.platforms([arch])], + ) + for arch in ("a2a3", "a5") + for runtime in ("host_build_graph", "tensormap_and_ringbuffer") +] + + +def _defined_external_symbols(path: Path) -> set[str]: + if sys.platform == "darwin": + command = ["nm", "-gU", str(path)] + else: + command = ["nm", "-D", "--defined-only", str(path)] + result = subprocess.run(command, check=True, capture_output=True, text=True) + symbols = {line.split()[-1] for line in result.stdout.splitlines() if line.split()} + if sys.platform == "darwin": + return {symbol.removeprefix("_") for symbol in symbols} + return symbols + + +@pytest.mark.parametrize( + ("arch", "variant", "runtime"), + _SIM_CASES + _ONBOARD_CASES, +) +def test_host_runtime_exports_required_pipeline_symbols(arch: str, variant: str, runtime: str): + runtime_dir = _PROJECT_ROOT / "build" / "lib" / arch / variant / runtime + runtime_libraries = tuple(runtime_dir.glob("libhost_runtime.*")) + assert len(runtime_libraries) == 1, runtime_dir + + runtime_path = runtime_libraries[0] + symbols = _defined_external_symbols(runtime_path) + + assert _NEWLY_REQUIRED_PIPELINE_SYMBOLS <= symbols, sorted(_NEWLY_REQUIRED_PIPELINE_SYMBOLS - symbols) diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index 7096a79cbb..97b156df1c 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -252,15 +252,18 @@ def fake_await_children_ready(shms, _pids, kind: str, _deadline: float) -> None: class _FakeNativeRunImpl: - def __init__(self) -> None: + def __init__(self, *, supports_concurrent_native_prepare: bool = False) -> None: + self.supports_concurrent_native_prepare = supports_concurrent_native_prepare self.events: list[tuple] = [] self.completed = [threading.Event(), threading.Event()] self.prepared = [threading.Event(), threading.Event()] self.launched = [threading.Event(), threading.Event()] self.finalized = [threading.Event(), threading.Event()] self.launch_errors: dict[tuple[int, int], BaseException] = {} + self.prepare_errors: dict[tuple[int, int], BaseException] = {} self.poll_errors: dict[tuple[int, int], BaseException] = {} self.finalize_errors: dict[tuple[int, int], BaseException] = {} + self.prepare_identities: list[tuple[int, int, int, int]] = [] self.register_calls: list[tuple[int, int]] = [] self.register_called = threading.Event() self._finalized_runs: set[tuple[int, int]] = set() @@ -270,8 +273,22 @@ def register_callable_from_blob(self, cid: int, blob_addr: int) -> None: self.register_calls.append((int(cid), int(blob_addr))) self.register_called.set() - def _prepare_native_run_from_blob(self, _cid, _blob_addr, _capacity, _cfg, slot_id, generation): + def _prepare_native_run_from_blob( + self, + _cid, + _blob_addr, + _capacity, + _cfg, + slot_id, + generation, + _run_id=0, + _dispatch_id=0, + ): slot = int(slot_id) + prepare_error = self.prepare_errors.get((slot, int(generation))) + if prepare_error is not None: + raise prepare_error + self.prepare_identities.append((slot, int(generation), int(_run_id), int(_dispatch_id))) token = SimpleNamespace(slot_id=slot, generation=int(generation), run_epoch=slot + 1) self.events.append(("prepare", slot)) self.prepared[slot].set() @@ -320,8 +337,8 @@ def _finalize_native_run(self, token) -> None: class _FakeTwoFrameChipWorker: pipeline_depth = 2 - def __init__(self) -> None: - self._impl = _FakeNativeRunImpl() + def __init__(self, *, supports_concurrent_native_prepare: bool = False) -> None: + self._impl = _FakeNativeRunImpl(supports_concurrent_native_prepare=supports_concurrent_native_prepare) self.malloc_called = threading.Event() self.unregister_calls: list[int] = [] self.unregister_called = threading.Event() @@ -340,13 +357,18 @@ def _unregister_slot(self, cid: int) -> None: class _TwoFrameLoopHarness: - def __init__(self) -> None: + def __init__( + self, + *, + supports_concurrent_native_prepare: bool = False, + chip_runtime: str = "", + ) -> None: self.shm = SharedMemory(create=True, size=MAILBOX_SIZE) self.buf = cast(memoryview, self.shm.buf) assert self.buf is not None self.mailbox_addr = _mailbox_addr(self.shm) self.digest = bytes([0x42]) * worker_mod.CALLABLE_HASH_DIGEST_BYTES - self.cw = _FakeTwoFrameChipWorker() + self.cw = _FakeTwoFrameChipWorker(supports_concurrent_native_prepare=supports_concurrent_native_prepare) self.registry = {7: object()} self.identity_table = {self.digest: 7} self.identity_refs = {self.digest: 1} @@ -363,7 +385,12 @@ def __init__(self) -> None: self.identity_table, self.identity_refs, ), - kwargs={"chip_platform": "a2a3", "prepared": self.prepared, "task_frame_count": 2}, + kwargs={ + "chip_platform": "a2a3", + "chip_runtime": chip_runtime, + "prepared": self.prepared, + "task_frame_count": 2, + }, ) def _frame_offset(self, index: int) -> int: @@ -382,6 +409,7 @@ def publish( *, state: int = worker_mod._TASK_READY, generation: int = 11, + diagnostics: bool = False, ) -> None: offset = self._frame_offset(index) frame = self.buf[offset : offset + worker_mod.MAILBOX_FRAME_SIZE] @@ -389,7 +417,9 @@ def publish( frame[worker_mod._OFF_TASK_CALLABLE_HASH : worker_mod._OFF_TASK_ARGS_BLOB] = self.digest struct.pack_into("=ii", frame, worker_mod._OFF_TASK_ARGS_BLOB, 0, 0) cfg_values = [0] * (6 + 3 * worker_mod.RUNTIME_ENV_RING_COUNT) - worker_mod._CFG_FMT.pack_into(frame, worker_mod._OFF_CONFIG, *cfg_values, b"") + cfg_values[3] = int(diagnostics) + output_prefix = b"/tmp/simpler-test" if diagnostics else b"" + worker_mod._CFG_FMT.pack_into(frame, worker_mod._OFF_CONFIG, *cfg_values, output_prefix) worker_mod._PIPELINE_LEASE_FMT.pack_into(frame, worker_mod._OFF_PIPELINE_LEASE, index, 0, generation) struct.pack_into("=Q", frame, worker_mod._OFF_FRAME_PROTOCOL, worker_mod._TASK_PROTOCOL_VERSION) struct.pack_into("=Q", frame, worker_mod._OFF_FRAME_RUN_ID, 5) @@ -453,7 +483,7 @@ def close(self) -> None: def test_two_frame_stages_b_without_native_prepare_until_a_finalizes(): - harness = _TwoFrameLoopHarness() + harness = _TwoFrameLoopHarness(chip_runtime="tensormap_and_ringbuffer") try: harness.publish(0, 1) harness.start() @@ -489,6 +519,193 @@ def test_two_frame_stages_b_without_native_prepare_until_a_finalizes(): harness.close() +def test_two_frame_hbg_prepares_b_while_a_runs_but_accepts_only_after_launch(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + try: + harness.publish(0, 1) + harness.start() + assert harness.cw._impl.launched[0].wait(5.0) + + harness.publish(1, 2, state=worker_mod._PREPARE_READY) + harness.wait_state(1, worker_mod._FRAME_STAGED) + assert harness.cw._impl.prepared[1].is_set() + assert not harness.cw._impl.finalized[0].is_set() + assert not harness.cw._impl.launched[1].is_set() + assert _mailbox_load_i32(harness.accepted_addr(1)) == 0 + + _mailbox_store_i32(harness.state_addr(1), worker_mod._ACTIVATE) + assert not harness.cw._impl.launched[1].wait(0.05) + assert _mailbox_load_i32(harness.accepted_addr(1)) == 0 + + harness.cw._impl.completed[0].set() + assert harness.cw._impl.finalized[0].wait(5.0) + assert harness.cw._impl.launched[1].wait(5.0) + assert _mailbox_load_i32(harness.accepted_addr(1)) == worker_mod._TASK_ACCEPTED + harness.cw._impl.completed[1].set() + harness.wait_state(1, worker_mod._TASK_DONE) + + lifecycle = [event[:2] for event in harness.cw._impl.events if event[0] in {"prepare", "launch", "finalize"}] + assert lifecycle == [ + ("prepare", 0), + ("launch", 0), + ("prepare", 1), + ("finalize", 0), + ("launch", 1), + ("finalize", 1), + ] + finally: + harness.close() + + +def test_two_frame_hbg_publishes_failure_instead_of_staged_when_prepare_fails(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + try: + harness.cw._impl.prepare_errors[(0, 11)] = RuntimeError("prepare failed") + harness.publish(0, 1) + harness.start() + harness.wait_state(0, worker_mod._TASK_FAILED) + assert not harness.cw._impl.prepared[0].is_set() + assert not harness.cw._impl.launched[0].is_set() + assert _mailbox_load_i32(harness.accepted_addr(0)) == 0 + finally: + harness.close() + + +def test_two_frame_hbg_waits_for_first_token_to_launch_before_preparing_second(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + try: + harness.publish(0, 1) + harness.publish(1, 2, state=worker_mod._PREPARE_READY) + harness.start() + + assert harness.cw._impl.launched[0].wait(5.0) + harness.wait_state(1, worker_mod._FRAME_STAGED) + assert harness.cw._impl.prepared[1].is_set() + assert not harness.cw._impl.finalized[0].is_set() + assert [event[:2] for event in harness.cw._impl.events if event[0] in {"prepare", "launch"}] == [ + ("prepare", 0), + ("launch", 0), + ("prepare", 1), + ] + + _mailbox_store_i32(harness.state_addr(1), worker_mod._ACTIVATE) + harness.cw._impl.completed[0].set() + assert harness.cw._impl.launched[1].wait(5.0) + harness.cw._impl.completed[1].set() + harness.wait_state(1, worker_mod._TASK_DONE) + finally: + harness.close() + + +def test_two_frame_hbg_prepares_and_launches_reverse_ready_frames_by_dispatch_id(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + try: + harness.publish(0, 2) + harness.publish(1, 1) + harness.start() + + deadline = time.monotonic() + 5.0 + while not any(event.is_set() for event in harness.cw._impl.launched): + assert time.monotonic() < deadline + time.sleep(0.001) + assert harness.cw._impl.launched[1].is_set() + assert not harness.cw._impl.launched[0].is_set() + assert [event[:2] for event in harness.cw._impl.events if event[0] in {"prepare", "launch"}] == [ + ("prepare", 1), + ("launch", 1), + ("prepare", 0), + ] + assert harness.cw._impl.prepare_identities == [(1, 11, 5, 1), (0, 11, 5, 2)] + + harness.cw._impl.completed[1].set() + assert harness.cw._impl.launched[0].wait(5.0) + harness.cw._impl.completed[0].set() + harness.wait_state(0, worker_mod._TASK_DONE) + finally: + harness.close() + + +def test_two_frame_hbg_does_not_prepare_high_dispatch_successor_before_active_frame(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + try: + harness.publish(0, 2, state=worker_mod._PREPARE_READY) + harness.publish(1, 1) + harness.start() + + assert harness.cw._impl.launched[1].wait(5.0) + harness.wait_state(0, worker_mod._FRAME_STAGED) + assert harness.cw._impl.prepared[0].is_set() + assert not harness.cw._impl.launched[0].is_set() + assert [event[:2] for event in harness.cw._impl.events if event[0] in {"prepare", "launch"}] == [ + ("prepare", 1), + ("launch", 1), + ("prepare", 0), + ] + + _mailbox_store_i32(harness.state_addr(0), worker_mod._ACTIVATE) + harness.cw._impl.completed[1].set() + assert harness.cw._impl.launched[0].wait(5.0) + harness.cw._impl.completed[0].set() + harness.wait_state(0, worker_mod._TASK_DONE) + finally: + harness.close() + + +@pytest.mark.parametrize("diagnostic_frame", ["active", "successor"]) +def test_two_frame_hbg_defers_diagnostic_native_prepare_until_predecessor_finalizes(diagnostic_frame): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + try: + harness.publish(0, 1, diagnostics=diagnostic_frame == "active") + harness.start() + assert harness.cw._impl.launched[0].wait(5.0) + + harness.publish( + 1, + 2, + state=worker_mod._PREPARE_READY, + diagnostics=diagnostic_frame == "successor", + ) + harness.wait_state(1, worker_mod._FRAME_STAGED) + assert not harness.cw._impl.prepared[1].is_set() + + _mailbox_store_i32(harness.state_addr(1), worker_mod._ACTIVATE) + harness.cw._impl.completed[0].set() + assert harness.cw._impl.finalized[0].wait(5.0) + assert harness.cw._impl.launched[1].wait(5.0) + harness.cw._impl.completed[1].set() + harness.wait_state(1, worker_mod._TASK_DONE) + + lifecycle = [event[:2] for event in harness.cw._impl.events if event[0] in {"prepare", "launch", "finalize"}] + assert lifecycle == [ + ("prepare", 0), + ("launch", 0), + ("finalize", 0), + ("prepare", 1), + ("launch", 1), + ("finalize", 1), + ] + finally: + harness.close() + + def test_two_frame_launches_by_dispatch_id_when_frames_are_ready_in_reverse_order(): harness = _TwoFrameLoopHarness() try: @@ -523,6 +740,55 @@ def test_two_frame_prepare_ready_waits_for_sticky_activation(): harness.close() +def test_two_frame_hbg_lone_prepare_ready_stages_before_native_prepare(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + try: + harness.publish(0, 1, state=worker_mod._PREPARE_READY) + harness.start() + harness.wait_state(0, worker_mod._FRAME_STAGED) + assert not harness.cw._impl.prepared[0].is_set() + + _mailbox_store_i32(harness.state_addr(0), worker_mod._ACTIVATE) + assert harness.cw._impl.launched[0].wait(5.0) + assert harness.cw._impl.prepared[0].is_set() + harness.cw._impl.completed[0].set() + harness.wait_state(0, worker_mod._TASK_DONE) + finally: + harness.close() + + +def test_two_frame_hbg_prepares_already_staged_successor_after_active_claim(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + try: + harness.publish(1, 2, state=worker_mod._PREPARE_READY) + harness.start() + harness.wait_state(1, worker_mod._FRAME_STAGED) + assert not harness.cw._impl.prepared[1].is_set() + + harness.publish(0, 1) + assert harness.cw._impl.launched[0].wait(5.0) + assert harness.cw._impl.prepared[1].wait(5.0) + assert [event[:2] for event in harness.cw._impl.events if event[0] in {"prepare", "launch"}] == [ + ("prepare", 0), + ("launch", 0), + ("prepare", 1), + ] + + _mailbox_store_i32(harness.state_addr(1), worker_mod._ACTIVATE) + harness.cw._impl.completed[0].set() + assert harness.cw._impl.launched[1].wait(5.0) + harness.cw._impl.completed[1].set() + harness.wait_state(1, worker_mod._TASK_DONE) + finally: + harness.close() + + def test_two_frame_processes_control_while_native_run_is_active(): harness = _TwoFrameLoopHarness() try: @@ -596,6 +862,40 @@ def test_two_frame_defers_register_until_active_native_run_finalizes(): payload_shm.unlink() +def test_two_frame_defers_register_until_backend_prepared_frame_finalizes(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + callable_obj = _unique_chip_callable(29) + digest = _chip_digest(callable_obj, platform="a2a3", runtime="host_build_graph") + payload_shm = _chip_payload_shm(callable_obj) + try: + harness.publish(0, 1) + harness.start() + assert harness.cw._impl.launched[0].wait(5.0) + harness.publish(1, 2, state=worker_mod._PREPARE_READY) + harness.wait_state(1, worker_mod._FRAME_STAGED) + assert harness.cw._impl.prepared[1].is_set() + + harness.publish_register(callable_obj, payload_shm, digest) + harness.assert_control_stays_pending() + assert harness.cw._impl.register_calls == [] + + _mailbox_store_i32(harness.state_addr(1), worker_mod._ACTIVATE) + harness.cw._impl.completed[0].set() + assert harness.cw._impl.launched[1].wait(5.0) + harness.assert_control_stays_pending() + harness.cw._impl.completed[1].set() + harness.wait_state(1, worker_mod._TASK_DONE) + harness.wait_control_state(worker_mod._CONTROL_DONE) + assert len(harness.cw._impl.register_calls) == 1 + finally: + harness.close() + payload_shm.close() + payload_shm.unlink() + + def test_two_frame_defers_final_unregister_while_matching_frame_is_staged(): harness = _TwoFrameLoopHarness() try: @@ -667,6 +967,57 @@ def test_two_frame_shutdown_finalizes_active_and_fails_staged_without_launch(): harness.close() +def test_two_frame_shutdown_finalizes_backend_prepared_successor_once(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + try: + harness.publish(0, 1) + harness.start() + assert harness.cw._impl.launched[0].wait(5.0) + harness.publish(1, 2, state=worker_mod._PREPARE_READY) + harness.wait_state(1, worker_mod._FRAME_STAGED) + assert harness.cw._impl.prepared[1].is_set() + + _mailbox_store_i32(harness.mailbox_addr + _OFF_STATE, worker_mod._SHUTDOWN) + harness.thread.join(5.0) + assert not harness.thread.is_alive() + assert _mailbox_load_i32(harness.state_addr(1)) == worker_mod._TASK_FAILED + assert not harness.cw._impl.launched[1].is_set() + assert sum(event == ("finalize", 0) for event in harness.cw._impl.events) == 1 + assert sum(event == ("finalize", 1) for event in harness.cw._impl.events) == 1 + finally: + harness.close() + + +def test_two_frame_stale_activation_finalizes_backend_prepared_successor_once(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="host_build_graph", + ) + try: + harness.publish(0, 1) + harness.start() + assert harness.cw._impl.launched[0].wait(5.0) + harness.publish(1, 2, state=worker_mod._PREPARE_READY) + harness.wait_state(1, worker_mod._FRAME_STAGED) + assert harness.cw._impl.prepared[1].is_set() + + frame_offset = harness._frame_offset(1) + struct.pack_into("=Q", harness.buf, frame_offset + worker_mod._OFF_FRAME_DISPATCH_ID, 99) + _mailbox_store_i32(harness.state_addr(1), worker_mod._ACTIVATE) + harness.wait_state(1, worker_mod._TASK_FAILED) + assert not harness.cw._impl.launched[1].is_set() + assert sum(event == ("finalize", 1) for event in harness.cw._impl.events) == 1 + + harness.cw._impl.completed[0].set() + harness.wait_state(0, worker_mod._TASK_DONE) + assert sum(event == ("finalize", 1) for event in harness.cw._impl.events) == 1 + finally: + harness.close() + + def test_two_frame_reuses_completed_frame_for_next_staged_successor(): harness = _TwoFrameLoopHarness() try: