From af454c7548940fad7f5de08bcea3c0d708c45323 Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Mon, 27 Jul 2026 21:55:29 +0800 Subject: [PATCH] Feat: add generation-safe depth-two pipeline slots MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Permit A2/A3 host runtimes to declare `pipeline_depth = 2` and give every resource class the copy count its declaration implies, without admitting a second run into device execution. `HOST_PER_RUN` and `EXEC_HANDLE` gain one copy per slot, `DEVICE_SCRATCH` keeps one, and `pipeline_resource_copy_count` / `pipeline_resource_slot` derive both from the contract so no caller hardcodes a replication rule. GM heap, GM shared memory, and the runtime image are committed together into one arena bank, so a contract that classifies them inconsistently — or repeats one of them, where only the first entry is ever read — describes a layout no runner can build. `has_serviceable_arena_topology` rejects it when the runtime loads rather than at the first launch that would need the second bank. Repeated stream kinds select no bank and stay legal. A run reaches its resources through a `PipelineSlotLease {slot_id, generation}`. `PipelineSlotPool` is the only mint and the only authority on ownership: it hands out one lease per free slot, bumps the generation on reuse, and makes release idempotent for the current generation while rejecting a stale one. Consumers downstream of the pool cannot re-derive ownership, so `ChipWorker` carries a `PipelineSlotGenerationFilter` that only refuses generations older than the newest that slot has presented. The filter mints nothing: an unleased run selects slot 0 and advances no generation, so a pool's first lease is still admitted after any number of ordinary runs. Each slot owns a host `Runtime` staging buffer. That buffer is not the `RUNTIME_IMAGE` resource — the contract classifies the device-resident image, while this holds the host-side object a run constructs in place, with its tensor leases, launch arguments, and validate/finalize state. That is per-run whatever the device image sharing is, so a runtime whose image is `DEVICE_SCRATCH` still needs its own buffer per slot. The A2/A3 runner owns its slot and arena-bank selection as member state, so two contexts on one thread cannot see each other's choice — a fresh worker's prewarm lays out its arenas in the bank it will actually serve runs from. The three pooled device regions become one bank per slot, and the single-entry prebuilt-arena cache stays owned by bank 0. AICore streams are outside that lease and outside the slot: each run creates its own and retires it on every exit path. The instruction cache belongs to the cores, every slot publishes its image to the same GM code address, and the platform offers no cache invalidation for code replaced there. Creating a stream is the only operation known to leave a core free of the previous image's instructions — selecting an existing one is not — so no record of which image a stream last ran can make reuse safe, and none is kept. The AICPU stream carries no such state and stays with its slot for the worker's lifetime. Simulation implements the same depth, so the contract means the same thing on both platforms: its runner owns one arena bank and one retained temp buffer per slot, and a leased slot-1 run there exercises the second copy rather than being refused. The generation filter is a replay filter, not an ownership check, and says so: it rejects a lease whose successor has already presented itself, but one that was released while its successor has not yet reached the consumer still passes. Closing that window needs dispatch gated on `PipelineSlotPool::owns`, which belongs to whoever admits runs. The ordinary synchronous path stays on slot 0 and the chip child's mailbox loop passes no lease, so every production run is unleased. No prepared FIFO, second mailbox frame, or backend publication overlap is enabled here. A run owns its AICore stream, so a destroy it cannot complete is that run's failure: the retire path returns the driver's error and keeps the handle, which both leaves the slot unusable for the next run — the stream may still hold the previous image's instructions — and leaves teardown something to retry. Reporting success there would have told the caller a slot is reusable that the next acquire now refuses. That state machine lives in `RunStreamSlots`, whose stream creation and destruction are injected, so the failure paths are covered without a device. `get_arena_bank_gm_heap_base_ctx`, `get_retained_temp_addr_ctx`, and `ChipWorker::runtime_buffer_addrs` report which storage a bank and a slot actually own. Sequential runs re-stage their arguments every round, so correct output alone survives two slots sharing one buffer; the tests compare addresses instead. --- docs/task-flow.md | 49 ++++- python/bindings/task_interface.cpp | 64 ++++++- python/simpler/task_interface.py | 30 ++- python/simpler/worker.py | 10 +- .../platform/onboard/host/device_runner.cpp | 119 ++++-------- .../platform/onboard/host/device_runner.h | 34 +++- src/a2a3/platform/sim/host/device_runner.cpp | 18 +- .../host_build_graph/host/runtime_maker.cpp | 2 +- .../docs/RUNTIME_LOGIC.md | 12 +- .../host/runtime_maker.cpp | 2 +- src/a5/platform/sim/host/device_runner.cpp | 18 +- .../docs/RUNTIME_LOGIC.md | 2 +- src/common/platform/include/common/host_api.h | 14 +- .../platform/include/host/run_stream_slots.h | 117 ++++++++++++ .../platform/onboard/host/c_api_shared.cpp | 36 ++++ .../onboard/host/device_runner_base.cpp | 133 +++++++++---- .../onboard/host/device_runner_base.h | 89 +++++++-- src/common/platform/sim/host/c_api_shared.cpp | 24 +++ .../platform/sim/host/device_runner_base.cpp | 79 +++++--- .../platform/sim/host/device_runner_base.h | 76 ++++++-- src/common/worker/chip_worker.cpp | 116 ++++++++++- src/common/worker/chip_worker.h | 51 ++++- src/common/worker/pipeline_contract.h | 61 +++++- src/common/worker/pipeline_slot_pool.h | 130 +++++++++++++ src/common/worker/pto_runtime_c_api.h | 42 +++- .../run_stream_reuse/test_run_stream_reuse.py | 168 +++++++++++++--- .../pipeline_slots/test_pipeline_slots.py | 141 ++++++++++++++ tests/ut/cpp/CMakeLists.txt | 13 ++ .../hierarchical/test_pipeline_contract.cpp | 145 +++++++++++--- .../hierarchical/test_run_stream_slots.cpp | 180 ++++++++++++++++++ 30 files changed, 1678 insertions(+), 297 deletions(-) create mode 100644 src/common/platform/include/host/run_stream_slots.h create mode 100644 src/common/worker/pipeline_slot_pool.h create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/pipeline_slots/test_pipeline_slots.py create mode 100644 tests/ut/cpp/hierarchical/test_run_stream_slots.cpp diff --git a/docs/task-flow.md b/docs/task-flow.md index 944f0f0028..f9b109ea9d 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -257,10 +257,57 @@ void ChipWorker::run(int32_t local_slot, TaskArgsView view, const CallConfig &co One memcpy of a few KB per task; negligible. +#### Pipeline resource leases + +A2/A3 host runtimes declare `pipeline_depth = 2` as resource capacity. The +contract determines the concrete copy count rather than permitting concurrent +device execution by itself: + +| Resource class | Copies | Selection | +| -------------- | -----: | --------- | +| `HOST_PER_RUN` | `pipeline_depth` | lease `slot_id` | +| `DEVICE_SCRATCH` | 1 | slot 0 | +| `EXEC_HANDLE` | `pipeline_depth` | lease `slot_id`, plus any hardware-generation key | + +A run-owned lease is `{slot_id, generation}`. `PipelineSlotPool` mints it and +is the authority on ownership: releasing the current lease is idempotent, while +releasing it after the slot has been re-leased is rejected. + +`ChipWorker` is downstream of that pool and cannot re-derive ownership — it +never sees an acquire or a release. It keeps a per-slot high-water mark and +rejects any generation below it, which stops a *superseded* lease from +selecting resources the slot's newer owner holds. That is strictly weaker than +an ownership check: a lease that was released but whose successor has not yet +reached `ChipWorker` still passes, because nothing has raised the mark. Closing +that window needs the admission layer to gate dispatch on `pool.owns(lease)` +before handing work down, which is whole-run admission's job, not this +layer's. + +AICore streams are outside that lease. The instruction cache belongs to the +cores, every slot publishes its image to the same GM code address, and the +platform offers no cache invalidation for code replaced there. Creating a +stream is the only operation known to leave a core free of the previous image's +instructions; selecting an already-existing one is not. So each run creates its +own AICore stream and retires it on every exit path, and no record of which +image a stream last ran is load-bearing. The AICPU stream carries no such state +and stays with its slot. + +This is dormant capacity at this layer. The ordinary synchronous entry point +continues to use slot 0, and the chip child's mailbox loop passes no lease, so +every production run is unleased. This contract does not enable a second +mailbox frame, a second device execution, or cross-run publication overlap. +Carrying a lease across the mailbox, and deciding when slot 1 may be leased at +all, belong to whole-run admission. + +Simulation implements the same depth, so the contract means the same thing on +both platforms: its runner owns one arena bank and one retained temporary +buffer per slot, and its single-entry prebuilt-arena cache stays owned by +bank 0 exactly as onboard's does. + #### TRB temporary buffer `tensormap_and_ringbuffer` stages ordinary non-child tensor arguments through a -runner-scoped retained temporary buffer instead of a per-run `device_malloc()` / +retained temporary buffer owned per pipeline slot, instead of a per-run `device_malloc()` / `device_free()` pair. This is always on for TRB — an internal allocation optimization with no user-facing switch. It is not serialized in task mailboxes and does not change `TaskArgs`, `CallConfig`, child-memory tensors, or public diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 1aac74a091..56f862a9cd 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -1430,10 +1430,29 @@ NB_MODULE(_task_interface, m) { "Launch a callable_id from a TaskArgs (used for in-process callers). " "Returns None; timing is emitted as `[STRACE]` log markers." ) + .def( + "_run_with_pipeline_lease", + [](ChipWorker &self, int32_t callable_id, TaskArgs &args, const CallConfig &config, uint32_t slot_id, + uint64_t generation) { + self.run_with_lease(callable_id, make_view(args), config, PipelineSlotLease{slot_id, 0, generation}); + }, + nb::arg("callable_id"), nb::arg("args"), nb::arg("config"), nb::arg("slot_id"), nb::arg("generation"), + "Internal generation-safe pipeline-slot launch used by hierarchical admission." + ) + .def( + "_run_with_pipeline_lease", + [](ChipWorker &self, int32_t callable_id, ChipStorageTaskArgs &args, const CallConfig &config, + uint32_t slot_id, uint64_t generation) { + self.run_with_lease(callable_id, &args, config, PipelineSlotLease{slot_id, 0, generation}); + }, + nb::arg("callable_id"), nb::arg("args"), nb::arg("config"), nb::arg("slot_id"), nb::arg("generation"), + "Internal generation-safe pipeline-slot launch for pre-encoded task args." + ) .def( "run_from_blob", [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity, - const CallConfig &config, uint64_t accepted_state_addr, int32_t accepted_value) { + const CallConfig &config, uint64_t accepted_state_addr, int32_t accepted_value, uint32_t pipeline_slot, + uint64_t pipeline_generation) { // The mailbox region is the on-wire format `write_blob` produced; // `read_blob` is the matching reader that returns a zero-copy // TaskArgsView into the caller-owned bytes. Forwards to the @@ -1441,12 +1460,21 @@ NB_MODULE(_task_interface, m) { // loops never re-implement the tensor/scalar layout in Python // (where it has historically dropped fields like child_memory). TaskArgsView view = read_blob(reinterpret_cast(args_blob_ptr), blob_capacity); - self.run( - callable_id, view, config, reinterpret_cast(accepted_state_addr), accepted_value - ); + if (pipeline_generation == 0) { + self.run( + callable_id, view, config, reinterpret_cast(accepted_state_addr), + accepted_value + ); + } else { + self.run_with_lease( + callable_id, view, config, PipelineSlotLease{pipeline_slot, 0, pipeline_generation}, + reinterpret_cast(accepted_state_addr), accepted_value + ); + } }, nb::arg("callable_id"), nb::arg("args_blob_ptr"), nb::arg("blob_capacity"), nb::arg("config"), - nb::arg("accepted_state_addr") = 0, nb::arg("accepted_value") = 0, + nb::arg("accepted_state_addr") = 0, nb::arg("accepted_value") = 0, nb::arg("pipeline_slot") = 0, + nb::arg("pipeline_generation") = 0, "Launch a callable_id from a raw mailbox-blob pointer + capacity " "(used by chip-child mailbox loops to avoid Python-side re-deserialisation " "of the per-task tensor/scalar layout). The blob must be in the format " @@ -1464,6 +1492,23 @@ NB_MODULE(_task_interface, m) { ) .def_prop_ro("device_id", &ChipWorker::device_id) .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( + "runtime_buffer_addrs", &ChipWorker::runtime_buffer_addrs, + "Host Runtime staging buffer address of every copy the runtime's " + "PipelineContract asked for, in slot order." + ) + .def( + "retained_temp_addr", &ChipWorker::retained_temp_addr, nb::arg("slot_id"), + "Retained temporary-buffer address held for one pipeline slot, or 0 " + "while that slot holds none." + ) + .def( + "arena_bank_gm_heap_base", &ChipWorker::arena_bank_gm_heap_base, nb::arg("bank_id"), + "Committed GM heap base of one arena bank on the bound runner, or 0 " + "when that bank has never been committed." + ) .def_prop_ro( "aicpu_dlopen_count", &ChipWorker::aicpu_dlopen_count, "Number of distinct callable entries the AICPU has dlopened for on the " @@ -1479,11 +1524,10 @@ NB_MODULE(_task_interface, m) { ) .def_prop_ro( "run_stream_set_create_count", &ChipWorker::run_stream_set_create_count, - "Number of run stream sets the bound runner has created. A set " - "belongs to a pipeline slot and is reused for every run on that " - "slot, so a worker that has served any number of runs reports 1; " - "platforms whose runs use the persistent bootstrap pair report 0. " - "Tests assert this to verify repeated runs do not rebuild the set." + "Number of AICore run streams the bound runner has created. The AICPU " + "stream belongs to a pipeline slot for the worker's lifetime, while each " + "run creates and retires its own AICore stream, so this advances once per " + "run; platforms whose runs use the persistent bootstrap pair report 0." ) .def("malloc", &ChipWorker::malloc, nb::arg("size")) .def("free", &ChipWorker::free, nb::arg("ptr")) diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index e29bec4c58..73e453c631 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -1320,6 +1320,13 @@ def _run_slot(self, callable_id, args, config=None, **kwargs): # Returns None; per-stage timing is emitted as `[STRACE]` log markers. self._impl.run(int(callable_id), args, config) + def _run_slot_with_pipeline_lease(self, callable_id, args, slot_id, generation, config=None, **kwargs): + if config is None: + config = CallConfig() + for k, v in kwargs.items(): + setattr(config, k, v) + self._impl._run_with_pipeline_lease(int(callable_id), args, config, int(slot_id), int(generation)) + def _unregister_slot(self, callable_id): self._impl.unregister_callable(int(callable_id)) @@ -1335,9 +1342,30 @@ def host_dlopen_count(self): @property def run_stream_set_create_count(self): - """Number of run stream generations the bound runner has created.""" + """Number of AICore run streams the bound runner has created.""" return self._impl.run_stream_set_create_count + @property + def pipeline_depth(self): + return self._impl.pipeline_depth + + @property + def runtime_slot_count(self): + return self._impl.runtime_slot_count + + @property + def runtime_buffer_addrs(self): + """Address of each host Runtime staging buffer, in slot order.""" + return list(self._impl.runtime_buffer_addrs) + + def arena_bank_gm_heap_base(self, bank_id): + """Committed GM heap base of one arena bank, or 0 when uncommitted.""" + return int(self._impl.arena_bank_gm_heap_base(int(bank_id))) + + def retained_temp_addr(self, slot_id): + """Retained temporary-buffer address for one slot, or 0 when unheld.""" + return int(self._impl.retained_temp_addr(int(slot_id))) + def malloc(self, size): """Allocate memory. Returns a pointer (uint64).""" return int(self._impl.malloc(int(size))) diff --git a/python/simpler/worker.py b/python/simpler/worker.py index ed9a99e42b..c8e71b2398 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -6275,12 +6275,12 @@ def host_dlopen_count(self) -> int: @property def run_stream_set_create_count(self) -> int: - """L2 only: number of run stream generations the runner has created. + """L2 only: number of AICore run streams the runner has created. - AICPU streams belong to pipeline slots. AICore streams are reused only - while the loaded AICore image is unchanged, so each code transition - advances this count. Returns 0 on non-L2 workers and on platforms whose - runs use the persistent bootstrap stream pair (simulation, a5). + AICPU streams belong to pipeline slots for the worker's lifetime, while + each run creates and retires its own AICore stream, so this advances + once per run. Returns 0 on non-L2 workers and on platforms whose runs + use the persistent bootstrap stream pair (simulation, a5). """ if self.level != 2 or self._chip_worker is None: return 0 diff --git a/src/a2a3/platform/onboard/host/device_runner.cpp b/src/a2a3/platform/onboard/host/device_runner.cpp index ae35b65156..838b65828e 100644 --- a/src/a2a3/platform/onboard/host/device_runner.cpp +++ b/src/a2a3/platform/onboard/host/device_runner.cpp @@ -215,7 +215,7 @@ void DeviceRunner::set_dep_gen_enabled(bool enable) { } int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { - constexpr unsigned kPipelineSlot = 0; + const unsigned selected_pipeline_slot = pipeline_slot(); // 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); @@ -248,18 +248,23 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return rc; } - const int32_t callable_id = runtime.get_active_callable_id(); - auto callable_it = callables_.find(callable_id); - if (callable_it == callables_.end()) { - LOG_ERROR("run() has no registered state for callable_id=%d", callable_id); - return -1; - } - const uint64_t aicore_image_hash = callable_it->second.aicore_image_hash; - rc = ensure_run_stream_set(kPipelineSlot, aicore_image_hash); + rc = ensure_run_stream_set(selected_pipeline_slot); if (rc != 0) { - LOG_ERROR("ensure_run_stream_set(%u, image=0x%lx) failed: %d", kPipelineSlot, aicore_image_hash, rc); + 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(); @@ -475,10 +480,17 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { l2_swimlane_collector_.set_core_types(core_types.data(), num_aicore); } - rc = launch_run(runtime, num_aicore, launch_aicpu_num, kPipelineSlot); + rc = launch_run(runtime, num_aicore, launch_aicpu_num, selected_pipeline_slot); + if (rc != 0) return rc; + + rc = reap_run(selected_pipeline_slot); if (rc != 0) return rc; - rc = reap_run(kPipelineSlot); + // 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. + aicore_stream_retired = true; + rc = retire_run_aicore_stream(selected_pipeline_slot); if (rc != 0) return rc; // Print handshake results (reads from device memory, must be before free) @@ -487,75 +499,30 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return 0; } -int DeviceRunner::ensure_run_stream_set(unsigned slot, uint64_t aicore_image_hash) { - if (slot >= kRunStreamSetCount) { - LOG_ERROR("ensure_run_stream_set: invalid slot %u", slot); - return -1; - } - RunStreamSet &streams = run_stream_sets_[slot]; - if (streams.aicpu == nullptr) { - int rc = rtStreamCreate(&streams.aicpu, 0); - if (rc != 0) { - LOG_ERROR("rtStreamCreate (run AICPU slot %u) failed: %d", slot, rc); - ACL_LOG_ERROR_DETAIL(rc); - return rc; - } - } - if (streams.aicore != nullptr && streams.has_aicore_image && streams.aicore_image_hash == aicore_image_hash) { - return 0; - } - - if (streams.aicore != nullptr) { - int rc = rtStreamDestroy(streams.aicore); - if (rc != 0) { - LOG_ERROR("rtStreamDestroy (retired AICore slot %u) failed: %d", slot, rc); - return rc; - } - streams.aicore = nullptr; - streams.has_aicore_image = false; +int DeviceRunner::ensure_run_stream_set(unsigned slot) { + int rc = run_stream_slots_.acquire(slot); + if (rc != 0) { + LOG_ERROR("ensure_run_stream_set(%u) failed: %d", slot, rc); + ACL_LOG_ERROR_DETAIL(rc); } + return rc; +} - int rc = rtStreamCreate(&streams.aicore, 0); +int DeviceRunner::retire_run_aicore_stream(unsigned slot) { + int rc = run_stream_slots_.retire_aicore(slot); if (rc != 0) { - LOG_ERROR("rtStreamCreate (run AICore slot %u) failed: %d", slot, rc); - ACL_LOG_ERROR_DETAIL(rc); - streams.aicore = nullptr; - return rc; + LOG_ERROR("rtStreamDestroy (run AICore slot %u) failed: %d, slot is now unusable", slot, rc); } - streams.aicore_image_hash = aicore_image_hash; - streams.has_aicore_image = true; - ++run_stream_sets_created_; - LOG_INFO_V0("DeviceRunner: run stream generation %zu created for slot %u", run_stream_sets_created_, slot); - return 0; + return rc; } int DeviceRunner::destroy_run_stream_sets() { - int rc = 0; - auto capture = [&rc](int err) { - if (err != 0 && rc == 0) rc = err; - }; // No pre-destroy sync, for the reason finalize_common() documents for the // bootstrap pair: rtStreamDestroy is the supported teardown for a stream // left in the error state by an op-timeout. - for (unsigned slot = 0; slot < kRunStreamSetCount; ++slot) { - RunStreamSet &streams = run_stream_sets_[slot]; - if (streams.aicpu != nullptr) { - int destroy_rc = rtStreamDestroy(streams.aicpu); - if (destroy_rc != 0) { - LOG_ERROR("rtStreamDestroy (run AICPU slot %u) failed: %d", slot, destroy_rc); - } - capture(destroy_rc); - streams.aicpu = nullptr; - } - if (streams.aicore != nullptr) { - int destroy_rc = rtStreamDestroy(streams.aicore); - if (destroy_rc != 0) { - LOG_ERROR("rtStreamDestroy (run AICore slot %u) failed: %d", slot, destroy_rc); - } - capture(destroy_rc); - streams.aicore = nullptr; - streams.has_aicore_image = false; - } + int rc = run_stream_slots_.destroy_all(); + if (rc != 0) { + LOG_ERROR("destroy_run_stream_sets: a stream survived teardown: %d", rc); } return rc; } @@ -564,12 +531,11 @@ int DeviceRunner::launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_ // KernelLaunch is the pipeline boundary: this method clears the handshake // consumed by the launch and submits exactly the AICore and AICPU kernels. // It intentionally performs no stream synchronization or per-run cleanup. - if (slot >= kRunStreamSetCount || run_stream_sets_[slot].aicpu == nullptr || - run_stream_sets_[slot].aicore == nullptr) { + if (!run_stream_slots_.ready(slot)) { LOG_ERROR("launch_run: stream set %u is not ready", slot); return -1; } - RunStreamSet &streams = run_stream_sets_[slot]; + RunStreamSet streams{run_stream_slots_.aicpu(slot), run_stream_slots_.aicore(slot)}; int rc = 0; // Launch the AICore worker BEFORE the AICPU Run task — mirrors the a5 path @@ -626,12 +592,11 @@ int DeviceRunner::launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_ } int DeviceRunner::reap_run(unsigned slot) { - if (slot >= kRunStreamSetCount) { + if (!run_stream_slots_.ready(slot)) { LOG_ERROR("reap_run: invalid stream set %u", slot); return -1; } - RunStreamSet &streams = run_stream_sets_[slot]; - int rc = sync_stream_pair(streams.aicpu, streams.aicore); + int rc = sync_stream_pair(run_stream_slots_.aicpu(slot), run_stream_slots_.aicore(slot)); if (rc != 0) { // The pair wait surfaces the AICore op-timeout (STARS-reaped op -> // 507000/507018/507046 at AICPU/AICore stream sync). The op-timeout diff --git a/src/a2a3/platform/onboard/host/device_runner.h b/src/a2a3/platform/onboard/host/device_runner.h index cc5544415f..62f8183754 100644 --- a/src/a2a3/platform/onboard/host/device_runner.h +++ b/src/a2a3/platform/onboard/host/device_runner.h @@ -38,6 +38,7 @@ #include "callable.h" #include "prepare_callable_common.h" #include "pto_runtime_c_api.h" // PTO_PIPELINE_MAX_DEPTH +#include "host/run_stream_slots.h" #include "common/kernel_args.h" #include "common/memory_barrier.h" #include "common/l2_swimlane_profiling.h" @@ -191,7 +192,7 @@ class DeviceRunner : public DeviceRunnerBase { // `aicpu_dlopen_count`, and `host_dlopen_count` are inherited from // `DeviceRunnerBase`. - size_t run_stream_set_create_count() const override { return run_stream_sets_created_; } + size_t run_stream_set_create_count() const override { return run_stream_slots_.created_count(); } private: // Most lifecycle state (device_id_, block_dim_, cores_per_blockdim_, @@ -228,18 +229,33 @@ class DeviceRunner : public DeviceRunnerBase { struct RunStreamSet { rtStream_t aicpu{nullptr}; rtStream_t aicore{nullptr}; - uint64_t aicore_image_hash{0}; - bool has_aicore_image{false}; }; + // One slot per in-flight run the pipeline contract can declare. static constexpr unsigned kRunStreamSetCount = PTO_PIPELINE_MAX_DEPTH; - // AICPU streams belong to slots. AICore streams additionally belong to the - // loaded code image because the platform has no explicit instruction-cache - // invalidation operation for code replaced in GM. - RunStreamSet run_stream_sets_[kRunStreamSetCount]{}; - size_t run_stream_sets_created_{0}; - int ensure_run_stream_set(unsigned slot, uint64_t aicore_image_hash); + // AICPU streams belong to slots for the worker's lifetime. AICore streams + // belong to a single run: the platform offers no instruction-cache + // invalidation for code replaced at a reused GM address, and creating the + // stream is the only operation known to leave a core free of the previous + // image's instructions. Selecting an existing stream is not. + // Stream lifetimes live in RunStreamSlots so the slot state machine — + // fresh AICore stream per run, handle kept when a destroy fails — is + // testable without a device. + RunStreamSlots run_stream_slots_{ + [](void **out) { + return rtStreamCreate(reinterpret_cast(out), 0); + }, + [](void *stream) { + return rtStreamDestroy(static_cast(stream)); + } + }; + int ensure_run_stream_set(unsigned slot); + // Destroys this run's AICore stream. Returns the driver's error and KEEPS + // the handle when the destroy fails: the stream may still hold the previous + // image's instructions, so the slot must refuse the next run until finalize + // reclaims it. + int retire_run_aicore_stream(unsigned slot); int destroy_run_stream_sets(); // The kernel submission boundary is separate from the stream wait and the diff --git a/src/a2a3/platform/sim/host/device_runner.cpp b/src/a2a3/platform/sim/host/device_runner.cpp index 6a4333cdcf..df8a232534 100644 --- a/src/a2a3/platform/sim/host/device_runner.cpp +++ b/src/a2a3/platform/sim/host/device_runner.cpp @@ -694,13 +694,19 @@ int DeviceRunner::finalize() { // Release the three per-Worker pooled arenas. Must precede mem_alloc_.finalize() // so the arenas free through the still-live allocator, not after it. - gm_heap_arena_.release(); - gm_sm_arena_.release(); - runtime_arena_pool_.release(); + for (auto &bank : arena_banks_) { + bank->gm_heap.release(); + bank->gm_sm.release(); + bank->runtime_pool.release(); + } clear_temporary_buffer(); - cached_gm_heap_size_ = 0; - cached_gm_sm_size_ = 0; - cached_runtime_arena_size_ = 0; + for (auto &bank : arena_banks_) { + bank->cached_gm_heap_size = 0; + bank->cached_gm_sm_size = 0; + bank->cached_runtime_arena_size = 0; + } + pipeline_slot_ = 0; + arena_bank_ = 0; prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = 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 80cc399a0c..025095976d 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -70,7 +70,7 @@ extern "C" const PipelineContract *get_pipeline_contract(void) { static const PipelineContract contract = { PTO_PIPELINE_CONTRACT_ABI_VERSION, 5, - 1, + 2, { {PTO_PIPELINE_GM_HEAP, PTO_PIPELINE_HOST_PER_RUN, 0}, {PTO_PIPELINE_GM_SM, PTO_PIPELINE_HOST_PER_RUN, 0}, diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index 77f388b0f3..469d3a8a27 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -96,12 +96,16 @@ Two platform implementations exist under `src/platform/`, sharing a common inter TRB bind normally allocates one device buffer per ordinary non-child tensor during host-side argument staging, copies input bytes as needed, records copy-back metadata, and frees those temporary buffers during runtime -validation. TRB replaces those per-run malloc/free pairs with a single -runner-scoped retained buffer that is reused across runs. This is always on for -TRB — an internal allocation optimization, not user-facing configuration. +validation. TRB replaces those per-run malloc/free pairs with a retained buffer +that its runs reuse. This is always on for TRB — an internal allocation +optimization, not user-facing configuration. + +The buffer is owned per pipeline slot, not per runner: TRB's task args are +`HOST_PER_RUN`, so two runs holding different slot leases stage through +different buffers even though they share arena bank 0 for device scratch. The platform side is deliberately thin: `DeviceRunnerBase` only remembers a -`{addr, size}` slot across runs, exposed through two HostApi callbacks — +`{addr, size}` slot per pipeline slot, exposed through two HostApi callbacks — `get_retained_temp_buffer` and `set_retained_temp_buffer`. It is not an allocator; all grow/pack/slice logic lives in `runtime_maker.cpp` (`RetainedTempBump`). diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index 9bb1b611cd..52603949ff 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -67,7 +67,7 @@ extern "C" const PipelineContract *get_pipeline_contract(void) { static const PipelineContract contract = { PTO_PIPELINE_CONTRACT_ABI_VERSION, 6, - 1, + 2, { {PTO_PIPELINE_TASK_ARGS, PTO_PIPELINE_HOST_PER_RUN, 0}, {PTO_PIPELINE_GM_HEAP, PTO_PIPELINE_DEVICE_SCRATCH, 0}, diff --git a/src/a5/platform/sim/host/device_runner.cpp b/src/a5/platform/sim/host/device_runner.cpp index 1ace457e95..cd3c6c1500 100644 --- a/src/a5/platform/sim/host/device_runner.cpp +++ b/src/a5/platform/sim/host/device_runner.cpp @@ -623,13 +623,19 @@ int DeviceRunner::finalize() { unload_executor_binaries(); - gm_heap_arena_.release(); - gm_sm_arena_.release(); - runtime_arena_pool_.release(); + for (auto &bank : arena_banks_) { + bank->gm_heap.release(); + bank->gm_sm.release(); + bank->runtime_pool.release(); + } clear_temporary_buffer(); - cached_gm_heap_size_ = 0; - cached_gm_sm_size_ = 0; - cached_runtime_arena_size_ = 0; + for (auto &bank : arena_banks_) { + bank->cached_gm_heap_size = 0; + bank->cached_gm_sm_size = 0; + bank->cached_runtime_arena_size = 0; + } + pipeline_slot_ = 0; + arena_bank_ = 0; prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; diff --git a/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index f6d3cab788..475d417693 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -98,7 +98,7 @@ TRB bind normally allocates one device buffer per ordinary non-child tensor during host-side argument staging, copies input bytes as needed, records copy-back metadata, and frees those temporary buffers during runtime validation. TRB replaces those per-run malloc/free pairs with a single -runner-scoped retained buffer that is reused across runs. This is always on for +retained buffer, owned per pipeline slot, that its runs reuse. This is always on for TRB — an internal allocation optimization, not user-facing configuration. The platform side is deliberately thin: `DeviceRunnerBase` only remembers a diff --git a/src/common/platform/include/common/host_api.h b/src/common/platform/include/common/host_api.h index e4ff984849..9cbd54cf0a 100644 --- a/src/common/platform/include/common/host_api.h +++ b/src/common/platform/include/common/host_api.h @@ -54,14 +54,16 @@ struct HostApi { // device_free's the old one, device_malloc's a bigger one, and writes the // new {addr, size} back. The grow/pack/slice logic lives in trb bind // (runtime_maker); the platform only remembers the slot so it can be reused - // next run and freed at finalize. hbg leaves these null and stages tensors - // through per-tensor device_malloc. `get` returns {nullptr, 0} when nothing - // is retained yet. + // by later runs and freed at finalize. The slot is per pipeline slot, so + // the pair reads and writes whichever one the current run's lease selected + // — two runs in different slots never share a staging buffer. hbg leaves + // these null and stages tensors through per-tensor device_malloc. `get` + // returns {nullptr, 0} when nothing is retained yet. void (*get_retained_temp_buffer)(void **addr, size_t *size); void (*set_retained_temp_buffer)(void *addr, size_t size); - // Commit the three per-Worker pooled regions (PTO2 GM heap, PTO2 shared - // memory, trb prebuilt runtime arena) as three independent device - // allocations. `runtime_arena_size == 0` skips the third region (hbg + // Commit the three pooled regions (PTO2 GM heap, PTO2 shared memory, trb + // prebuilt runtime arena) of the arena bank the current run's lease + // selected, as three independent device allocations. `runtime_arena_size == 0` skips the third region (hbg // path: hbg has no prebuilt runtime arena). Idempotent on identical // sizes; returns 0 on success, -1 on allocation failure. int (*setup_static_arena)(size_t gm_heap_size, size_t gm_sm_size, size_t runtime_arena_size); diff --git a/src/common/platform/include/host/run_stream_slots.h b/src/common/platform/include/host/run_stream_slots.h new file mode 100644 index 0000000000..194a164747 --- /dev/null +++ b/src/common/platform/include/host/run_stream_slots.h @@ -0,0 +1,117 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#ifndef SRC_COMMON_PLATFORM_INCLUDE_HOST_RUN_STREAM_SLOTS_H_ +#define SRC_COMMON_PLATFORM_INCLUDE_HOST_RUN_STREAM_SLOTS_H_ + +#include +#include +#include + +#include "pto_runtime_c_api.h" + +/** + * Per-slot ownership of the two streams a run submits on. + * + * The AICPU stream carries no instruction-cache state, so it belongs to its + * slot for the runner's lifetime. The AICore stream belongs to a single run: + * the platform offers no instruction-cache invalidation for code replaced at a + * reused GM address, and creating the stream is the only operation known to + * leave a core free of the previous image's instructions — selecting an + * existing one is not. + * + * A destroy that fails keeps its handle. Such a stream may still hold the + * previous image's instructions, so the slot must refuse the next run rather + * than hand it a fresh stream beside a live one, and teardown needs the handle + * to retry. Stream creation and destruction are injected so this state machine + * is exercisable without a device. + */ +class RunStreamSlots { +public: + using CreateFn = std::function; + using DestroyFn = std::function; + + RunStreamSlots(CreateFn create, DestroyFn destroy) : + create_(std::move(create)), + destroy_(std::move(destroy)) {} + + /** + * Ready `slot` for a run: its AICPU stream on first use, and always a fresh + * AICore stream. Fails when the slot still holds an AICore stream a prior + * run could not retire. + */ + int acquire(unsigned slot) { + if (slot >= slots_.size()) return -1; + Slot &s = slots_[slot]; + if (s.aicpu == nullptr) { + int rc = create_(&s.aicpu); + if (rc != 0) { + s.aicpu = nullptr; + return rc; + } + } + if (s.aicore != nullptr) return -1; + int rc = create_(&s.aicore); + if (rc != 0) { + s.aicore = nullptr; + return rc; + } + ++created_count_; + return 0; + } + + /** Retire `slot`'s AICore stream. The handle survives a failed destroy. */ + int retire_aicore(unsigned slot) { + if (slot >= slots_.size()) return -1; + Slot &s = slots_[slot]; + if (s.aicore == nullptr) return 0; + int rc = destroy_(s.aicore); + if (rc != 0) return rc; + s.aicore = nullptr; + return 0; + } + + /** Destroy every stream, keeping handles whose destroy failed. */ + int destroy_all() { + int first_error = 0; + for (Slot &s : slots_) { + for (void **stream : {&s.aicpu, &s.aicore}) { + if (*stream == nullptr) continue; + int rc = destroy_(*stream); + if (rc != 0) { + if (first_error == 0) first_error = rc; + continue; + } + *stream = nullptr; + } + } + return first_error; + } + + void *aicpu(unsigned slot) const { return slot < slots_.size() ? slots_[slot].aicpu : nullptr; } + void *aicore(unsigned slot) const { return slot < slots_.size() ? slots_[slot].aicore : nullptr; } + bool ready(unsigned slot) const { return aicpu(slot) != nullptr && aicore(slot) != nullptr; } + size_t created_count() const { return created_count_; } + static constexpr size_t capacity() { return PTO_PIPELINE_MAX_DEPTH; } + +private: + struct Slot { + void *aicpu{nullptr}; + void *aicore{nullptr}; + }; + + CreateFn create_; + DestroyFn destroy_; + std::array slots_{}; + size_t created_count_{0}; +}; + +#endif // SRC_COMMON_PLATFORM_INCLUDE_HOST_RUN_STREAM_SLOTS_H_ diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index 101118d6e8..70275d91ca 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -659,6 +659,42 @@ 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 { + return static_cast(ctx)->select_pipeline_slot(slot_id); + } catch (...) { + return -1; + } +} + +int select_arena_bank_ctx(DeviceContextHandle ctx, uint32_t bank_id) { + if (ctx == NULL) return -1; + try { + return static_cast(ctx)->select_arena_bank(bank_id); + } catch (...) { + return -1; + } +} + +uint64_t get_arena_bank_gm_heap_base_ctx(DeviceContextHandle ctx, uint32_t bank_id) { + if (ctx == NULL) return 0; + try { + return static_cast(ctx)->arena_bank_gm_heap_base(bank_id); + } catch (...) { + return 0; + } +} + +uint64_t get_retained_temp_addr_ctx(DeviceContextHandle ctx, uint32_t slot_id) { + if (ctx == NULL) return 0; + try { + return static_cast(ctx)->retained_temp_addr(slot_id); + } catch (...) { + return 0; + } +} + int simpler_unregister_callable(DeviceContextHandle ctx, int32_t callable_id) { if (ctx == NULL) return -1; try { diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index e1ef5115c9..9944a575f0 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -120,10 +120,42 @@ HostRuntimeTimeoutConfig resolve_onboard_timeout_config() { } // namespace -DeviceRunnerBase::DeviceRunnerBase() : - gm_heap_arena_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_), - gm_sm_arena_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_), - runtime_arena_pool_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_) {} +DeviceRunnerBase::DeviceRunnerBase() { + for (auto &bank : arena_banks_) { + bank = std::make_unique(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_); + } +} + +int DeviceRunnerBase::select_pipeline_slot(uint32_t slot_id) { + if (slot_id >= PTO_PIPELINE_MAX_DEPTH) { + LOG_ERROR("pipeline slot %u is outside [0, %u)", slot_id, PTO_PIPELINE_MAX_DEPTH); + return -1; + } + pipeline_slot_ = slot_id; + return 0; +} + +int DeviceRunnerBase::select_arena_bank(uint32_t bank_id) { + if (bank_id >= PTO_PIPELINE_MAX_DEPTH) { + LOG_ERROR("arena bank %u is outside [0, %u)", bank_id, PTO_PIPELINE_MAX_DEPTH); + return -1; + } + arena_bank_ = bank_id; + return 0; +} + +uint32_t DeviceRunnerBase::pipeline_slot() const { return pipeline_slot_; } + +uint64_t DeviceRunnerBase::arena_bank_gm_heap_base(uint32_t bank_id) const { + if (bank_id >= arena_banks_.size()) return 0; + const ArenaBank &bank = *arena_banks_[bank_id]; + return bank.gm_heap.is_committed() ? reinterpret_cast(bank.gm_heap.base()) : 0; +} + +uint64_t DeviceRunnerBase::retained_temp_addr(uint32_t slot_id) const { + if (slot_id >= retained_temp_addrs_.size()) return 0; + return reinterpret_cast(retained_temp_addrs_[slot_id]); +} void *DeviceRunnerBase::allocate_tensor(std::size_t bytes) { return mem_alloc_.alloc(bytes); } @@ -146,44 +178,51 @@ 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_addr_; - if (size != nullptr) *size = retained_temp_size_; + 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_addr_ = addr; - retained_temp_size_ = size; + retained_temp_addrs_[pipeline_slot_] = addr; + retained_temp_sizes_[pipeline_slot_] = size; } void DeviceRunnerBase::clear_temporary_buffer() { - if (retained_temp_addr_ != nullptr) { - mem_alloc_.free(retained_temp_addr_); - retained_temp_addr_ = nullptr; - retained_temp_size_ = 0; + for (size_t slot = 0; slot < retained_temp_addrs_.size(); ++slot) { + if (retained_temp_addrs_[slot] == nullptr) continue; + mem_alloc_.free(retained_temp_addrs_[slot]); + retained_temp_addrs_[slot] = nullptr; + retained_temp_sizes_[slot] = 0; } } void *DeviceRunnerBase::acquire_pooled_gm_heap() { - if (!gm_heap_arena_.is_committed()) return nullptr; - return gm_heap_arena_.base(); + DeviceArena &arena = arena_bank().gm_heap; + if (!arena.is_committed()) return nullptr; + return arena.base(); } void *DeviceRunnerBase::acquire_pooled_gm_sm() { - if (!gm_sm_arena_.is_committed()) return nullptr; - return gm_sm_arena_.base(); + DeviceArena &arena = arena_bank().gm_sm; + if (!arena.is_committed()) return nullptr; + return arena.base(); } void *DeviceRunnerBase::acquire_pooled_runtime_arena() { - // hbg calls setup_static_arena(...,0) and leaves runtime_arena_pool_ + // hbg calls setup_static_arena(...,0) and leaves the runtime pool // uncommitted — fail loudly if a caller asks for it anyway. - if (!runtime_arena_pool_.is_committed()) return nullptr; - return runtime_arena_pool_.base(); + DeviceArena &arena = arena_bank().runtime_pool; + if (!arena.is_committed()) return nullptr; + return arena.base(); } bool DeviceRunnerBase::lookup_prebuilt_runtime_arena_cache( uint64_t hash, const void *key_data, size_t key_size, void **gm_heap_base, void **sm_base, void **runtime_arena_base, size_t *runtime_off, const void **image_data, size_t *image_size ) 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 (!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 || @@ -206,6 +245,8 @@ void DeviceRunnerBase::mark_prebuilt_runtime_arena_cached( uint64_t hash, const void *key_data, size_t key_size, void *gm_heap_base, void *sm_base, void *runtime_arena_base, 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; prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_hash_ = hash; prebuilt_runtime_arena_cache_key_.assign( @@ -232,6 +273,8 @@ int DeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, // worker's lifetime). If a caller asks for a larger layout on any // region, redo just that region — already-committed peers stay alive // so their callers don't have to re-acquire. + ArenaBank &bank = arena_bank(); + bool arena_changed = false; auto commit_region = [&arena_changed](DeviceArena &arena, size_t &cached_size, size_t requested_size) -> int { if (requested_size == 0) { @@ -269,25 +312,27 @@ int DeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, // asking for a larger layout) fails midway, defeating the // "failure means failure" guarantee. Reset everything to the // post-construction state so the caller can retry with a new layout. - bool ok = commit_region(gm_heap_arena_, cached_gm_heap_size_, gm_heap_size) == 0; - ok = ok && commit_region(gm_sm_arena_, cached_gm_sm_size_, gm_sm_size) == 0; - ok = ok && commit_region(runtime_arena_pool_, cached_runtime_arena_size_, runtime_arena_size) == 0; + bool ok = commit_region(bank.gm_heap, bank.cached_gm_heap_size, gm_heap_size) == 0; + ok = ok && commit_region(bank.gm_sm, bank.cached_gm_sm_size, gm_sm_size) == 0; + ok = ok && commit_region(bank.runtime_pool, bank.cached_runtime_arena_size, runtime_arena_size) == 0; if (!ok) { - gm_heap_arena_.release(); - gm_sm_arena_.release(); - runtime_arena_pool_.release(); - cached_gm_heap_size_ = 0; - cached_gm_sm_size_ = 0; - cached_runtime_arena_size_ = 0; - prebuilt_runtime_arena_cache_valid_ = false; - prebuilt_runtime_arena_cache_key_.clear(); - prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; - prebuilt_runtime_arena_cache_sm_base_ = nullptr; - prebuilt_runtime_arena_cache_runtime_arena_base_ = nullptr; - prebuilt_runtime_arena_cache_image_.clear(); + bank.gm_heap.release(); + bank.gm_sm.release(); + bank.runtime_pool.release(); + bank.cached_gm_heap_size = 0; + bank.cached_gm_sm_size = 0; + bank.cached_runtime_arena_size = 0; + if (arena_bank_ == 0) { + prebuilt_runtime_arena_cache_valid_ = false; + prebuilt_runtime_arena_cache_key_.clear(); + prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; + prebuilt_runtime_arena_cache_sm_base_ = nullptr; + prebuilt_runtime_arena_cache_runtime_arena_base_ = nullptr; + prebuilt_runtime_arena_cache_image_.clear(); + } return -1; } - if (arena_changed) { + if (arena_changed && arena_bank_ == 0) { prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; @@ -1062,9 +1107,11 @@ int DeviceRunnerBase::finalize_common() { // trb prebuilt runtime arena — each its own device_malloc). Must precede // mem_alloc_.finalize() so the arenas free through the still-live // allocator, not after it. - gm_heap_arena_.release(); - gm_sm_arena_.release(); - runtime_arena_pool_.release(); + for (auto &bank : arena_banks_) { + bank->gm_heap.release(); + bank->gm_sm.release(); + bank->runtime_pool.release(); + } prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; @@ -1094,9 +1141,13 @@ int DeviceRunnerBase::finalize_common() { max_cube_cores_ = 0; max_vector_cores_ = 0; aicore_kernel_binary_.clear(); - cached_gm_heap_size_ = 0; - cached_gm_sm_size_ = 0; - cached_runtime_arena_size_ = 0; + for (auto &bank : arena_banks_) { + bank->cached_gm_heap_size = 0; + bank->cached_gm_sm_size = 0; + bank->cached_runtime_arena_size = 0; + } + pipeline_slot_ = 0; + arena_bank_ = 0; return rc; } diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index 96b9c87ab9..d34b975933 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -38,9 +38,11 @@ #include +#include #include #include #include +#include #include #include #include @@ -63,6 +65,7 @@ #include "host/scope_stats_collector.h" #include "host/args_dump_collector.h" #include "prepare_callable_common.h" +#include "pto_runtime_c_api.h" struct HostApi; // common/host_api.h — fwd-declared to keep task_interface headers out struct CallConfig; // task_interface/call_config.h — per-run config threaded into run() @@ -90,6 +93,28 @@ class DeviceRunnerBase { /** Bind this runner's launch-acceptance publication target. */ int set_task_accepted_state(volatile int32_t *state, int32_t accepted_value); + /** 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; + + /** + * Committed GM heap base of one arena bank, or 0 while that bank has never + * been committed. Two banks that have both served a run hold distinct + * device allocations; tests read this to prove the depth-two split is real + * rather than two names for one region. + */ + uint64_t arena_bank_gm_heap_base(uint32_t bank_id) const; + + /** + * Retained temporary-buffer address held for one pipeline slot, or 0 while + * that slot holds none. Two slots that have both staged arguments hold + * distinct buffers; tests read this to prove the split is real. + */ + uint64_t retained_temp_addr(uint32_t slot_id) const; + /** Allocate / free / copy on the per-Worker `MemoryAllocator` + CANN runtime. */ void *allocate_tensor(std::size_t bytes); void free_tensor(void *dev_ptr); @@ -138,11 +163,12 @@ class DeviceRunnerBase { int setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, size_t runtime_arena_size); /** - * Return the pooled GM heap / PTO2 SM / runtime arena base pointer. - * `setup_static_arena` (arch subclass) must have already committed - * the relevant region; otherwise returns nullptr. The runtime arena - * accessor is trb-only — hbg's `setup_static_arena(...,0)` leaves - * `runtime_arena_pool_` uncommitted and this returns nullptr. + * Return the pooled GM heap / PTO2 SM / runtime arena base pointer of the + * selected arena bank. `setup_static_arena` (arch subclass) must have + * already committed the relevant region on that bank; otherwise returns + * nullptr. The runtime arena accessor is trb-only — hbg's + * `setup_static_arena(...,0)` leaves the runtime pool uncommitted and this + * returns nullptr. */ void *acquire_pooled_gm_heap(); void *acquire_pooled_gm_sm(); @@ -903,22 +929,43 @@ class DeviceRunnerBase { host::LoadAicpuOp load_aicpu_op_; MemoryAllocator mem_alloc_; - // Retained temporary buffer slot for TRB device-arg staging (see HostApi - // get/set_retained_temp_buffer). Just a remembered {addr, size} reused - // across runs and freed in finalize; the grow/pack logic lives in trb bind. - void *retained_temp_addr_ = nullptr; - std::size_t retained_temp_size_ = 0; - DeviceArena gm_heap_arena_; - DeviceArena gm_sm_arena_; - DeviceArena runtime_arena_pool_; - - // Cached arena sizes for `setup_static_arena`'s "fits" check — avoids - // re-allocating the same buffer when a later worker init asks for an - // equal-or-smaller layout on an already-committed arena. Reset by - // the subclass's `finalize()` alongside the other identity state. - size_t cached_gm_heap_size_{0}; - size_t cached_gm_sm_size_{0}; - size_t cached_runtime_arena_size_{0}; + // Retained temporary buffer for TRB device-arg staging, one per pipeline + // slot (see HostApi get/set_retained_temp_buffer). Just a remembered + // {addr, size} that the slot reuses across its runs and finalize frees; + // the grow/pack logic lives in trb bind. + std::array retained_temp_addrs_{}; + std::array retained_temp_sizes_{}; + + // One independently committed set of the three pooled device regions. A + // run reaches its set through the arena bank its lease selects, so + // preparing one bank never mutates a region the active run is executing + // out of. `cached_*` back `setup_static_arena`'s "fits" check: a later + // init asking for an equal-or-smaller layout on an already-committed + // arena reuses it instead of re-allocating. + struct ArenaBank { + ArenaBank(DeviceArena::AllocFn alloc, DeviceArena::FreeFn free_fn, void *ctx) : + gm_heap(alloc, free_fn, ctx), + gm_sm(alloc, free_fn, ctx), + runtime_pool(alloc, free_fn, ctx) {} + + DeviceArena gm_heap; + DeviceArena gm_sm; + DeviceArena runtime_pool; + size_t cached_gm_heap_size{0}; + size_t cached_gm_sm_size{0}; + size_t cached_runtime_arena_size{0}; + }; + // 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}; + bool prebuilt_runtime_arena_cache_valid_{false}; uint64_t prebuilt_runtime_arena_cache_hash_{0}; std::vector prebuilt_runtime_arena_cache_key_; diff --git a/src/common/platform/sim/host/c_api_shared.cpp b/src/common/platform/sim/host/c_api_shared.cpp index 0a8b9af98d..42a41ba6cc 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -597,6 +597,30 @@ int simpler_run( } } +int select_pipeline_slot_ctx(DeviceContextHandle ctx, uint32_t slot_id) { + if (ctx == NULL) return -1; + return static_cast(ctx)->select_pipeline_slot(slot_id); +} + +int select_arena_bank_ctx(DeviceContextHandle ctx, uint32_t bank_id) { + if (ctx == NULL) return -1; + return static_cast(ctx)->select_arena_bank(bank_id); +} + +uint64_t get_arena_bank_gm_heap_base_ctx(DeviceContextHandle ctx, uint32_t bank_id) { + if (ctx == NULL) return 0; + return static_cast(ctx)->arena_bank_gm_heap_base(bank_id); +} + +uint64_t get_retained_temp_addr_ctx(DeviceContextHandle ctx, uint32_t slot_id) { + if (ctx == NULL) return 0; + try { + return static_cast(ctx)->retained_temp_addr(slot_id); + } catch (...) { + return 0; + } +} + int simpler_unregister_callable(DeviceContextHandle ctx, int32_t callable_id) { if (ctx == NULL) return -1; try { diff --git a/src/common/platform/sim/host/device_runner_base.cpp b/src/common/platform/sim/host/device_runner_base.cpp index fc7d1a9964..4b781cdff9 100644 --- a/src/common/platform/sim/host/device_runner_base.cpp +++ b/src/common/platform/sim/host/device_runner_base.cpp @@ -82,7 +82,31 @@ bool create_temp_so_file(const std::string &path_template, const uint8_t *data, // SimDeviceRunnerBase Implementation // ============================================================================= +int SimDeviceRunnerBase::select_pipeline_slot(uint32_t slot_id) { + if (slot_id >= PTO_PIPELINE_MAX_DEPTH) return -1; + pipeline_slot_ = slot_id; + return 0; +} + +int SimDeviceRunnerBase::select_arena_bank(uint32_t bank_id) { + if (bank_id >= PTO_PIPELINE_MAX_DEPTH) return -1; + arena_bank_ = bank_id; + return 0; +} + +uint64_t SimDeviceRunnerBase::arena_bank_gm_heap_base(uint32_t bank_id) const { + if (bank_id >= arena_banks_.size()) return 0; + const ArenaBank &bank = *arena_banks_[bank_id]; + return bank.gm_heap.is_committed() ? reinterpret_cast(bank.gm_heap.base()) : 0; +} + +uint64_t SimDeviceRunnerBase::retained_temp_addr(uint32_t slot_id) const { + if (slot_id >= retained_temp_addrs_.size()) return 0; + return reinterpret_cast(retained_temp_addrs_[slot_id]); +} + int SimDeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, size_t runtime_arena_size) { + ArenaBank &bank = arena_bank(); // Three independent device_malloc'd buffers: GM heap, PTO2 SM, prebuilt // runtime arena. Split out from a single large allocation because the // combined size can exceed the device allocator's largest contiguous @@ -120,16 +144,16 @@ int SimDeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_si // all on any failure" semantic (PR #922). Pooled pointers from a prior // successful call stay valid; a failed resize attempt does not leave a // partial layout behind. - bool ok = commit_region(gm_heap_arena_, cached_gm_heap_size_, gm_heap_size) == 0; - ok = ok && commit_region(gm_sm_arena_, cached_gm_sm_size_, gm_sm_size) == 0; - ok = ok && commit_region(runtime_arena_pool_, cached_runtime_arena_size_, runtime_arena_size) == 0; + bool ok = commit_region(bank.gm_heap, bank.cached_gm_heap_size, gm_heap_size) == 0; + ok = ok && commit_region(bank.gm_sm, bank.cached_gm_sm_size, gm_sm_size) == 0; + ok = ok && commit_region(bank.runtime_pool, bank.cached_runtime_arena_size, runtime_arena_size) == 0; if (!ok) { - gm_heap_arena_.release(); - gm_sm_arena_.release(); - runtime_arena_pool_.release(); - cached_gm_heap_size_ = 0; - cached_gm_sm_size_ = 0; - cached_runtime_arena_size_ = 0; + bank.gm_heap.release(); + bank.gm_sm.release(); + bank.runtime_pool.release(); + bank.cached_gm_heap_size = 0; + bank.cached_gm_sm_size = 0; + bank.cached_runtime_arena_size = 0; prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; @@ -153,6 +177,9 @@ bool SimDeviceRunnerBase::lookup_prebuilt_runtime_arena_cache( uint64_t hash, const void *key_data, size_t key_size, void **gm_heap_base, void **sm_base, void **runtime_arena_base, size_t *runtime_off, const void **image_data, size_t *image_size ) 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 (!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 || @@ -175,6 +202,8 @@ void SimDeviceRunnerBase::mark_prebuilt_runtime_arena_cached( uint64_t hash, const void *key_data, size_t key_size, void *gm_heap_base, void *sm_base, void *runtime_arena_base, 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; prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_hash_ = hash; prebuilt_runtime_arena_cache_key_.assign( @@ -191,18 +220,21 @@ void SimDeviceRunnerBase::mark_prebuilt_runtime_arena_cached( } void *SimDeviceRunnerBase::acquire_pooled_gm_heap() { - if (!gm_heap_arena_.is_committed()) return nullptr; - return gm_heap_arena_.base(); + DeviceArena &arena = arena_bank().gm_heap; + if (!arena.is_committed()) return nullptr; + return arena.base(); } void *SimDeviceRunnerBase::acquire_pooled_gm_sm() { - if (!gm_sm_arena_.is_committed()) return nullptr; - return gm_sm_arena_.base(); + DeviceArena &arena = arena_bank().gm_sm; + if (!arena.is_committed()) return nullptr; + return arena.base(); } void *SimDeviceRunnerBase::acquire_pooled_runtime_arena() { - if (!runtime_arena_pool_.is_committed()) return nullptr; - return runtime_arena_pool_.base(); + DeviceArena &arena = arena_bank().runtime_pool; + if (!arena.is_committed()) return nullptr; + return arena.base(); } std::thread SimDeviceRunnerBase::create_thread(std::function fn) { @@ -301,20 +333,21 @@ int SimDeviceRunnerBase::device_memset(void *dev_ptr, int value, size_t bytes) { } void SimDeviceRunnerBase::get_retained_temp_buffer(void **addr, size_t *size) { - if (addr != nullptr) *addr = retained_temp_addr_; - if (size != nullptr) *size = retained_temp_size_; + if (addr != nullptr) *addr = retained_temp_addrs_[pipeline_slot_]; + if (size != nullptr) *size = retained_temp_sizes_[pipeline_slot_]; } void SimDeviceRunnerBase::set_retained_temp_buffer(void *addr, size_t size) { - retained_temp_addr_ = addr; - retained_temp_size_ = size; + retained_temp_addrs_[pipeline_slot_] = addr; + retained_temp_sizes_[pipeline_slot_] = size; } void SimDeviceRunnerBase::clear_temporary_buffer() { - if (retained_temp_addr_ != nullptr) { - mem_alloc_.free(retained_temp_addr_); - retained_temp_addr_ = nullptr; - retained_temp_size_ = 0; + for (size_t slot = 0; slot < retained_temp_addrs_.size(); ++slot) { + if (retained_temp_addrs_[slot] == nullptr) continue; + mem_alloc_.free(retained_temp_addrs_[slot]); + retained_temp_addrs_[slot] = nullptr; + retained_temp_sizes_[slot] = 0; } } diff --git a/src/common/platform/sim/host/device_runner_base.h b/src/common/platform/sim/host/device_runner_base.h index fa94f8cc2e..0f37d9f6d5 100644 --- a/src/common/platform/sim/host/device_runner_base.h +++ b/src/common/platform/sim/host/device_runner_base.h @@ -28,15 +28,19 @@ #include #include +#include #include #include #include +#include #include #include #include #include #include +#include "pto_runtime_c_api.h" + #include "callable.h" #include "prepare_callable_common.h" #include "utils/device_arena.h" @@ -66,10 +70,24 @@ constexpr int SIM_AUTO_BLOCKDIM = 8; class SimDeviceRunnerBase { public: - SimDeviceRunnerBase() : - gm_heap_arena_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_), - gm_sm_arena_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_), - runtime_arena_pool_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_) {} + SimDeviceRunnerBase() { + for (auto &bank : arena_banks_) { + bank = std::make_unique(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_); + } + } + + /** 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); + uint32_t pipeline_slot() const { return pipeline_slot_; } + uint64_t arena_bank_gm_heap_base(uint32_t bank_id) const; + + /** + * Retained temporary-buffer address held for one pipeline slot, or 0 while + * that slot holds none. Two slots that have both staged arguments hold + * distinct buffers; tests read this to prove the split is real. + */ + uint64_t retained_temp_addr(uint32_t slot_id) const; // Public virtual dtor so c_api_shared can `delete` a SimDeviceRunnerBase * // (destroy_device_context entrypoint). @@ -234,29 +252,45 @@ class SimDeviceRunnerBase { std::vector aicore_kernel_binary_; MemoryAllocator mem_alloc_; - void *retained_temp_addr_ = nullptr; - size_t retained_temp_size_ = 0; - - // Three independent per-Worker arenas, each backing a single pooled - // region (PTO2 GM heap / PTO2 shared memory / trb prebuilt runtime - // arena). Split out from a single backing allocation because the - // combined size can exceed the device allocator's largest contiguous - // block. Released explicitly in finalize() before mem_alloc_.finalize(). + std::array retained_temp_addrs_{}; + std::array retained_temp_sizes_{}; + + // Each arena bank backs the three pooled regions (PTO2 GM heap / PTO2 + // shared memory / trb prebuilt runtime arena) for one pipeline slot. They + // are separate allocations because the combined size can exceed the device + // allocator's largest contiguous block. Released explicitly in finalize() + // before mem_alloc_.finalize(). // - // runtime_arena_pool_ stays unreserved when setup_static_arena was + // A bank's runtime pool stays unreserved when setup_static_arena was // invoked with runtime_arena_size == 0 (hbg path). static void *arena_alloc_trampoline(void *ctx, size_t size) { return static_cast(ctx)->alloc(size); } static void arena_free_trampoline(void *ctx, void *p) { static_cast(ctx)->free(p); } - DeviceArena gm_heap_arena_; - DeviceArena gm_sm_arena_; - DeviceArena runtime_arena_pool_; - // Cached sizes for setup_static_arena's "fits" check — avoids re-allocating - // a buffer when a later worker init asks for an equal-or-smaller layout. - size_t cached_gm_heap_size_{0}; - size_t cached_gm_sm_size_{0}; - size_t cached_runtime_arena_size_{0}; + // One independently committed set of the three pooled regions per pipeline + // slot, so preparing one bank never mutates a region the active run is + // executing out of. `cached_*` back setup_static_arena's "fits" check — + // avoids re-allocating when a later worker init asks for an equal-or- + // smaller layout. Held by pointer because DeviceArena is neither copyable + // nor movable, so the array cannot be brace-initialised without naming + // every bank. + struct ArenaBank { + ArenaBank(DeviceArena::AllocFn alloc, DeviceArena::FreeFn free_fn, void *ctx) : + gm_heap(alloc, free_fn, ctx), + gm_sm(alloc, free_fn, ctx), + runtime_pool(alloc, free_fn, ctx) {} + + DeviceArena gm_heap; + DeviceArena gm_sm; + DeviceArena runtime_pool; + size_t cached_gm_heap_size{0}; + size_t cached_gm_sm_size{0}; + size_t cached_runtime_arena_size{0}; + }; + std::array, PTO_PIPELINE_MAX_DEPTH> arena_banks_; + ArenaBank &arena_bank() { return *arena_banks_[arena_bank_]; } + uint32_t pipeline_slot_{0}; + uint32_t arena_bank_{0}; bool prebuilt_runtime_arena_cache_valid_{false}; uint64_t prebuilt_runtime_arena_cache_hash_{0}; std::vector prebuilt_runtime_arena_cache_key_; diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 68e48e2ce6..1136deaf8b 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -117,6 +117,11 @@ void ChipWorker::init( simpler_init_fn_ = load_symbol(handle, "simpler_init"); register_callable_fn_ = load_symbol(handle, "simpler_register_callable"); run_fn_ = load_symbol(handle, "simpler_run"); + select_pipeline_slot_fn_ = load_symbol(handle, "select_pipeline_slot_ctx"); + select_arena_bank_fn_ = load_symbol(handle, "select_arena_bank_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"); @@ -154,7 +159,7 @@ void ChipWorker::init( 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_depth1_pipeline_contract(contract)) { + 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"); } @@ -170,7 +175,25 @@ void ChipWorker::init( throw std::runtime_error("create_device_context returned null"); } - runtime_buf_.resize(get_runtime_size_fn_()); + try { + // One host Runtime per slot, always. This buffer is not the + // RUNTIME_IMAGE resource: the contract classifies the device-resident + // image, while this holds the host-side Runtime object a run + // constructs in place — its tensor leases, launch arguments, and + // validate/finalize state. That is per-run whatever the device image + // sharing is, so a runtime whose image is DEVICE_SCRATCH still needs + // its own buffer per slot or preparing N+1 overwrites live state of N. + std::vector> runtime_bufs( + resolved_contract.pipeline_depth, std::vector(get_runtime_size_fn_()) + ); + runtime_bufs_.swap(runtime_bufs); + } catch (...) { + destroy_device_context_fn_(device_ctx_); + device_ctx_ = nullptr; + dlclose(handle); + lib_handle_ = nullptr; + throw; + } // One-shot platform-side init: attach the calling thread to `device_id` // (rtSetDevice on onboard, sim bind+acquire on sim), transfer ownership @@ -218,6 +241,10 @@ void ChipWorker::init( simpler_init_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + select_pipeline_slot_fn_ = nullptr; + select_arena_bank_fn_ = nullptr; + get_arena_bank_gm_heap_base_fn_ = nullptr; + get_retained_temp_addr_fn_ = nullptr; set_task_accepted_state_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; @@ -236,7 +263,7 @@ void ChipWorker::init( comm_release_domain_windows_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; - runtime_buf_.clear(); + runtime_bufs_.clear(); throw; } if (init_rc != 0) { @@ -259,6 +286,10 @@ void ChipWorker::init( simpler_init_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + select_pipeline_slot_fn_ = nullptr; + select_arena_bank_fn_ = nullptr; + get_arena_bank_gm_heap_base_fn_ = nullptr; + get_retained_temp_addr_fn_ = nullptr; set_task_accepted_state_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; @@ -278,7 +309,7 @@ void ChipWorker::init( comm_release_domain_windows_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; - runtime_buf_.clear(); + runtime_bufs_.clear(); throw std::runtime_error("simpler_init failed with code " + std::to_string(init_rc)); } @@ -330,6 +361,10 @@ void ChipWorker::finalize() { get_runtime_size_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + select_pipeline_slot_fn_ = nullptr; + select_arena_bank_fn_ = nullptr; + get_arena_bank_gm_heap_base_fn_ = nullptr; + get_retained_temp_addr_fn_ = nullptr; set_task_accepted_state_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; @@ -349,7 +384,8 @@ void ChipWorker::finalize() { comm_release_domain_windows_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; - runtime_buf_.clear(); + runtime_bufs_.clear(); + pipeline_generations_.reset(); pipeline_contract_ = {PTO_PIPELINE_CONTRACT_ABI_VERSION, 0, 1, {}}; initialized_ = false; device_id_ = -1; @@ -388,13 +424,81 @@ void ChipWorker::run(int32_t callable_id, const ChipStorageTaskArgs *args, const void ChipWorker::run( int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, volatile int32_t *accepted_state, int32_t accepted_value +) { + run_on_slot(callable_id, args, config, UNLEASED_SLOT, accepted_state, accepted_value); +} + +uint32_t ChipWorker::select_slot_resources(uint32_t slot_id) { + // init() rejected any contract whose three arena kinds disagree, so + // whichever of them the runtime declares yields the same bank. + const PipelineSlotLease selector{slot_id, 0, 0}; + uint32_t arena_bank = 0; + for (uint32_t kind : {PTO_PIPELINE_GM_HEAP, PTO_PIPELINE_GM_SM, PTO_PIPELINE_RUNTIME_IMAGE}) { + const PipelineResource *resource = find_pipeline_resource(pipeline_contract_, kind); + if (resource == nullptr) continue; + arena_bank = pipeline_resource_slot(pipeline_contract_, *resource, selector); + break; + } + if (select_pipeline_slot_fn_(device_ctx_, slot_id) != 0) { + throw std::runtime_error("select_pipeline_slot_ctx failed"); + } + if (select_arena_bank_fn_(device_ctx_, arena_bank) != 0) { + throw std::runtime_error("select_arena_bank_ctx failed"); + } + return arena_bank; +} + +std::vector ChipWorker::runtime_buffer_addrs() const { + std::vector addrs; + addrs.reserve(runtime_bufs_.size()); + for (const auto &buf : runtime_bufs_) { + addrs.push_back(reinterpret_cast(buf.data())); + } + return addrs; +} + +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; + 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; + return get_retained_temp_addr_fn_(device_ctx_, slot_id); +} + +void ChipWorker::run_with_lease( + int32_t callable_id, TaskArgsView args, const CallConfig &config, const PipelineSlotLease &lease, + volatile int32_t *accepted_state, int32_t accepted_value +) { + ChipStorageTaskArgs chip_storage = view_to_chip_storage(args); + run_with_lease(callable_id, &chip_storage, config, lease, accepted_state, accepted_value); +} + +void ChipWorker::run_with_lease( + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, const PipelineSlotLease &lease, + volatile int32_t *accepted_state, int32_t accepted_value +) { + if (lease.reserved != 0 || lease.generation == 0 || lease.slot_id >= pipeline_contract_.pipeline_depth) { + throw std::runtime_error("run pipeline lease is outside the runtime PipelineContract"); + } + if (!pipeline_generations_.admit(lease)) { + throw std::runtime_error("run pipeline lease generation is stale"); + } + run_on_slot(callable_id, args, config, lease.slot_id, accepted_state, accepted_value); +} + +void ChipWorker::run_on_slot( + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, uint32_t slot_id, + volatile int32_t *accepted_state, int32_t accepted_value ) { config.validate(); if (!initialized_) { throw std::runtime_error("ChipWorker not initialized; call init() first"); } - void *rt = runtime_buf_.data(); + (void)select_slot_resources(slot_id); + void *rt = runtime_bufs_[slot_id].data(); 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) { diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 7e0dbb84fc..f248fb25f9 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -19,6 +19,7 @@ #include "../task_interface/call_config.h" #include "../task_interface/task_args.h" +#include "pipeline_slot_pool.h" #include "pto_runtime_c_api.h" #include "types.h" @@ -79,6 +80,14 @@ class ChipWorker { void run(int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, volatile int32_t *accepted_state, int32_t accepted_value); + void run_with_lease( + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, const PipelineSlotLease &lease, + volatile int32_t *accepted_state = nullptr, int32_t accepted_value = 0 + ); + void run_with_lease( + int32_t callable_id, TaskArgsView args, const CallConfig &config, const PipelineSlotLease &lease, + volatile int32_t *accepted_state = nullptr, int32_t accepted_value = 0 + ); // Per-callable_id preparation. Requires init() first and a callable_id // in [0, MAX_REGISTERED_CALLABLE_IDS) (cap 64). @@ -96,9 +105,10 @@ class ChipWorker { /// `aicpu_dlopen_count` for the trb path; returns 0 on device-orch variants. size_t host_dlopen_count() const; - /// Number of run stream generations the bound runner has created. AICPU - /// streams belong to slots; AICore streams are reused only while the loaded - /// code image is unchanged. Platforms using the persistent pair report 0. + /// Number of AICore run streams the bound runner has created. AICPU streams + /// belong to slots for the worker's lifetime; each run gets a freshly + /// created AICore stream, so this advances once per run. Platforms using + /// the persistent bootstrap pair report 0. size_t run_stream_set_create_count() const; uint64_t malloc(size_t size); @@ -152,6 +162,20 @@ class ChipWorker { int device_id() const { return device_id_; } bool initialized() const { return initialized_; } unsigned pipeline_depth() const { return pipeline_contract_.pipeline_depth; } + size_t runtime_slot_count() const { return runtime_bufs_.size(); } + + /// Host Runtime staging buffer address of every copy the contract asked + /// for, in slot order. Two copies hold distinct storage; tests read this to + /// prove per-run buffers are not one buffer under two slot ids. + std::vector runtime_buffer_addrs() const; + + /// Committed GM heap base of one arena bank on the bound runner, or 0 when + /// that bank has never been committed or the platform shares one arena set. + uint64_t arena_bank_gm_heap_base(uint32_t bank_id) const; + + /// Retained temporary-buffer address the bound runner holds for one + /// pipeline slot, or 0 while that slot holds none. + uint64_t retained_temp_addr(uint32_t slot_id) const; private: using CreateDeviceContextFn = void *(*)(); @@ -170,6 +194,10 @@ class ChipWorker { using SimplerRegisterCallableFn = int (*)(void *, int32_t, const void *); using SimplerRunFn = int (*)(void *, void *, int32_t, const void *, const CallConfig *); using SetTaskAcceptedStateFn = int (*)(void *, volatile int32_t *, int32_t); + using SelectPipelineSlotFn = int (*)(void *, uint32_t); + using SelectArenaBankFn = int (*)(void *, uint32_t); + using GetArenaBankGmHeapBaseFn = uint64_t (*)(void *, uint32_t); + using GetRetainedTempAddrFn = uint64_t (*)(void *, uint32_t); using GetPipelineContractFn = const PipelineContract *(*)(); using SimplerUnregisterCallableFn = int (*)(void *, int32_t); using GetAicpuDlopenCountFn = size_t (*)(void *); @@ -218,6 +246,10 @@ class ChipWorker { SimplerRegisterCallableFn register_callable_fn_ = nullptr; SimplerRunFn run_fn_ = nullptr; SetTaskAcceptedStateFn set_task_accepted_state_fn_ = nullptr; + SelectPipelineSlotFn select_pipeline_slot_fn_ = nullptr; + SelectArenaBankFn select_arena_bank_fn_ = nullptr; + GetArenaBankGmHeapBaseFn get_arena_bank_gm_heap_base_fn_ = nullptr; + GetRetainedTempAddrFn get_retained_temp_addr_fn_ = nullptr; SimplerUnregisterCallableFn unregister_callable_fn_ = nullptr; GetAicpuDlopenCountFn get_aicpu_dlopen_count_fn_ = nullptr; GetAicpuDlopenCountFn get_host_dlopen_count_fn_ = nullptr; @@ -241,7 +273,18 @@ class ChipWorker { std::unordered_map comm_session_index_; uint64_t base_comm_handle_ = 0; - std::vector runtime_buf_; + // Slot 0 with no generation bookkeeping: an unleased run is a caller that + // is not using the pipeline, and minting a generation for it here would + // advance the filter past the leases a real pool later presents. + static constexpr uint32_t UNLEASED_SLOT = 0; + void run_on_slot( + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, uint32_t slot_id, + volatile int32_t *accepted_state, int32_t accepted_value + ); + uint32_t select_slot_resources(uint32_t slot_id); + + std::vector> runtime_bufs_; + 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 // ChipWorker callers run on the thread that called init() (the same diff --git a/src/common/worker/pipeline_contract.h b/src/common/worker/pipeline_contract.h index 7b7dd9d035..9df7ba4355 100644 --- a/src/common/worker/pipeline_contract.h +++ b/src/common/worker/pipeline_contract.h @@ -27,8 +27,8 @@ * Whether this build can honor `contract`. * * Rejects a contract whose ABI version is not the one compiled in, that - * declares more resources than fit, that asks for a pipeline depth other than - * 1, or whose resources carry an unspecified or out-of-range kind, an + * declares more resources than fit, that asks for a pipeline depth outside + * the supported range, or whose resources carry an unspecified or out-of-range kind, an * out-of-range class, or a non-zero reserved `bytes_per_copy`. * * The unspecified-kind rule is what catches a `resource_count` larger than the @@ -36,9 +36,10 @@ * and zero is not a resource. Repeated kinds stay legal, so a runtime may * declare several resources of one kind. */ -inline bool is_valid_depth1_pipeline_contract(const PipelineContract *contract) { +inline bool is_valid_pipeline_contract(const PipelineContract *contract) { if (contract == nullptr || contract->abi_version != PTO_PIPELINE_CONTRACT_ABI_VERSION || - contract->resource_count > PTO_PIPELINE_MAX_RESOURCES || contract->pipeline_depth != 1) { + contract->resource_count > PTO_PIPELINE_MAX_RESOURCES || contract->pipeline_depth == 0 || + contract->pipeline_depth > PTO_PIPELINE_MAX_DEPTH) { return false; } for (uint32_t i = 0; i < contract->resource_count; ++i) { @@ -51,4 +52,56 @@ inline bool is_valid_depth1_pipeline_contract(const PipelineContract *contract) return true; } +/** Return the number of concrete copies required for one resource. */ +inline uint32_t pipeline_resource_copy_count(const PipelineContract &contract, const PipelineResource &resource) { + return resource.resource_class == PTO_PIPELINE_DEVICE_SCRATCH ? 1u : contract.pipeline_depth; +} + +/** Select the concrete copy of `resource` owned by `lease`. */ +inline uint32_t pipeline_resource_slot( + const PipelineContract &contract, const PipelineResource &resource, const PipelineSlotLease &lease +) { + return pipeline_resource_copy_count(contract, resource) == 1 ? 0u : lease.slot_id; +} + +/** Find a declared resource kind, or return nullptr when the runtime does not use it. */ +inline const PipelineResource *find_pipeline_resource(const PipelineContract &contract, uint32_t kind) { + for (uint32_t i = 0; i < contract.resource_count; ++i) { + if (contract.resources[i].kind == kind) return &contract.resources[i]; + } + return nullptr; +} + +/** + * Whether the three pooled device regions can be served by one bank selector. + * + * GM heap, GM shared memory, and the runtime image are committed together by + * one `setup_static_arena` call into one arena bank, so they cannot be given + * different copy counts. A contract that classifies them inconsistently — or + * declares one of them twice, where only the first entry would ever be read — + * describes a layout the executor cannot produce, and is rejected at load + * rather than at the first launch that would need the second bank. + */ +inline bool has_serviceable_arena_topology(const PipelineContract &contract) { + constexpr uint32_t ARENA_KINDS[] = {PTO_PIPELINE_GM_HEAP, PTO_PIPELINE_GM_SM, PTO_PIPELINE_RUNTIME_IMAGE}; + bool seen = false; + uint32_t shared_copies = 0; + for (uint32_t kind : ARENA_KINDS) { + uint32_t declared = 0; + const PipelineResource *resource = nullptr; + for (uint32_t i = 0; i < contract.resource_count; ++i) { + if (contract.resources[i].kind != kind) continue; + ++declared; + resource = &contract.resources[i]; + } + if (declared > 1) return false; + if (resource == nullptr) continue; + uint32_t copies = pipeline_resource_copy_count(contract, *resource); + if (seen && copies != shared_copies) return false; + shared_copies = copies; + seen = true; + } + return true; +} + #endif // SRC_COMMON_WORKER_PIPELINE_CONTRACT_H_ diff --git a/src/common/worker/pipeline_slot_pool.h b/src/common/worker/pipeline_slot_pool.h new file mode 100644 index 0000000000..66f18a1762 --- /dev/null +++ b/src/common/worker/pipeline_slot_pool.h @@ -0,0 +1,130 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#ifndef SRC_COMMON_WORKER_PIPELINE_SLOT_POOL_H_ +#define SRC_COMMON_WORKER_PIPELINE_SLOT_POOL_H_ + +#include +#include +#include +#include +#include +#include + +#include "pto_runtime_c_api.h" + +/** + * Fixed-capacity owner of generation-safe pipeline slot leases. + * + * This is capability only: callers decide when admission is allowed. A + * released lease remains idempotently releasable until the slot is acquired + * again; after reuse, the old generation can neither access nor release it. + */ +class PipelineSlotPool { +public: + explicit PipelineSlotPool(uint32_t depth) : + depth_(depth) { + if (depth == 0 || depth > PTO_PIPELINE_MAX_DEPTH) { + throw std::invalid_argument("pipeline slot depth is outside the supported range"); + } + } + + std::optional try_acquire() { + std::lock_guard lock(mu_); + for (uint32_t slot = 0; slot < depth_; ++slot) { + SlotState &state = slots_[slot]; + if (state.in_use) continue; + if (state.generation == std::numeric_limits::max()) { + throw std::overflow_error("pipeline slot generation space exhausted"); + } + ++state.generation; + state.in_use = true; + return PipelineSlotLease{slot, 0, state.generation}; + } + return std::nullopt; + } + + bool owns(const PipelineSlotLease &lease) const { + std::lock_guard lock(mu_); + return owns_locked(lease); + } + + bool release(const PipelineSlotLease &lease) { + std::lock_guard lock(mu_); + if (!matches_generation_locked(lease)) return false; + // Releasing the current generation twice is deliberately idempotent. + slots_[lease.slot_id].in_use = false; + return true; + } + + uint32_t depth() const { return depth_; } + +private: + struct SlotState { + uint64_t generation{0}; + bool in_use{false}; + }; + + bool matches_generation_locked(const PipelineSlotLease &lease) const { + return lease.reserved == 0 && lease.generation != 0 && lease.slot_id < depth_ && + slots_[lease.slot_id].generation == lease.generation; + } + + bool owns_locked(const PipelineSlotLease &lease) const { + return matches_generation_locked(lease) && slots_[lease.slot_id].in_use; + } + + const uint32_t depth_; + mutable std::mutex mu_; + std::array slots_{}; +}; + +/** + * Replay filter for a consumer that executes leases minted elsewhere. + * + * `PipelineSlotPool` is the sole mint and the sole authority on who owns a + * slot. A consumer downstream of it never sees an acquire or a release, so it + * cannot re-derive ownership; it records the newest generation each slot has + * presented and refuses anything older. Repeating the current generation stays + * admissible, because one run may dispatch several times under one lease. + * + * This is strictly weaker than an ownership check, and deliberately so. It + * rejects a *superseded* lease — one whose successor has already presented + * itself — but a lease that was released while its successor has not yet + * reached this consumer still passes, because nothing has raised the mark. + * Closing that window requires gating dispatch on `PipelineSlotPool::owns` + * before work is handed down, which belongs to whoever admits runs. + * + * The filter mints nothing. A consumer that also served unleased work must not + * invent generations for it, or a later real lease is rejected as stale. + */ +class PipelineSlotGenerationFilter { +public: + bool admit(const PipelineSlotLease &lease) { + if (lease.slot_id >= newest_.size()) return false; + std::lock_guard lock(mu_); + uint64_t &newest = newest_[lease.slot_id]; + if (lease.generation < newest) return false; + newest = lease.generation; + return true; + } + + void reset() { + std::lock_guard lock(mu_); + newest_.fill(0); + } + +private: + std::mutex mu_; + std::array newest_{}; +}; + +#endif // SRC_COMMON_WORKER_PIPELINE_SLOT_POOL_H_ diff --git a/src/common/worker/pto_runtime_c_api.h b/src/common/worker/pto_runtime_c_api.h index 0c77507fbe..c4092eab19 100644 --- a/src/common/worker/pto_runtime_c_api.h +++ b/src/common/worker/pto_runtime_c_api.h @@ -123,6 +123,20 @@ typedef struct PipelineContract { PipelineResource resources[PTO_PIPELINE_MAX_RESOURCES]; } PipelineContract; +/** + * Ownership token for one pipeline resource slot. + * + * `slot_id` selects the per-run/exec-handle copies. `generation` changes every + * time that slot is leased, so delayed completion or cleanup from an older run + * cannot mutate resources now owned by its successor. Generation zero is + * reserved for an invalid/uninitialised lease. + */ +typedef struct PipelineSlotLease { + uint32_t slot_id; + uint32_t reserved; + uint64_t generation; +} PipelineSlotLease; + /* Per-stage run timing is no longer returned. The platform emits it as * `[STRACE]` log markers (host stages + the AICPU device-phase breakdown, * gated on SIMPLER_HOST_STRACE) — parse with simpler_setup.tools.strace_timing. @@ -271,6 +285,26 @@ int simpler_run( DeviceContextHandle ctx, RuntimeHandle runtime, int32_t callable_id, const void *args, const CallConfig *config ); +/** Select the per-run/exec-handle slot used by the next synchronous run. */ +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); + +/** + * 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 + * device allocation a bank actually owns; changes nothing. + */ +uint64_t get_arena_bank_gm_heap_base_ctx(DeviceContextHandle ctx, uint32_t bank_id); + +/** + * Retained temporary-buffer address held for one pipeline slot, or 0 while that + * slot holds none. Reports which staging buffer a slot actually owns; changes + * nothing. + */ +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 @@ -314,10 +348,10 @@ size_t get_aicpu_dlopen_count(DeviceContextHandle ctx); size_t get_host_dlopen_count(DeviceContextHandle ctx); /** - * Number of run stream generations the runner bound to `ctx` has created. - * AICPU streams belong to pipeline slots; AICore streams are reused only while - * their loaded code image is unchanged. Returns 0 on platforms whose runs use - * the persistent bootstrap pair. + * Number of AICore run streams the runner bound to `ctx` has created. AICPU + * streams belong to pipeline slots for the worker's lifetime; each run gets a + * freshly created AICore stream, so this advances once per run. Returns 0 on + * platforms whose runs use the persistent bootstrap pair. */ size_t get_run_stream_set_create_count(DeviceContextHandle ctx); diff --git a/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py b/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py index 049927cdf3..f828ebac98 100644 --- a/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py +++ b/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py @@ -7,21 +7,27 @@ # 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. # ----------------------------------------------------------------------------------------------------------- -"""A2A3 run streams are reused only while their code image is unchanged. +"""A2A3 gives every run a freshly created AICore stream. A2A3 submits each run's AICore and AICPU kernels on the stream set of the selected pipeline slot rather than on the persistent bootstrap pair. The AICPU -stream belongs to the slot, while the AICore stream is also bound to the loaded -code image. Reusing it after different code replaces the image can execute stale -instructions because the platform has no explicit AICore I-cache invalidation. +stream carries no instruction-cache state and belongs to the slot for the +worker's lifetime. The AICore stream belongs to one run: the platform offers no +instruction-cache invalidation for code replaced at a reused GM address, and +creating the stream is the only operation known to leave a core free of the +previous image's instructions — selecting an already-existing stream is not, so +reuse cannot be made safe by tracking which image a stream last ran. -`Worker.run_stream_set_create_count` reports how many sets the bound runner has -built, including every fresh AICore stream created for a code transition. +`Worker.run_stream_set_create_count` counts those creations, so it advances once +per run. """ +import itertools + import pytest import torch from simpler.task_interface import ArgDirection as D +from simpler.worker import Worker from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test from simpler_setup.scene_test import _build_chip_task_args, _compare_outputs @@ -63,7 +69,7 @@ class _SubtractCallable(SceneTestCase): @scene_test(level=2, runtime="host_build_graph") class TestRunStreamReuseHbg(SceneTestCase): - """Repeated runs on one worker must share a single run stream set.""" + """Each run gets its own AICore stream; its pipeline slot owns the rest.""" RTOL = 1e-5 ATOL = 1e-5 @@ -99,7 +105,7 @@ class TestRunStreamReuseHbg(SceneTestCase): CASES = [ { "name": "repeated_runs", - "platforms": ["a2a3"], + "platforms": ["a2a3", "a2a3sim"], "config": {"aicpu_thread_num": 4, "block_dim": 3}, "params": {}, }, @@ -117,18 +123,19 @@ def compute_golden(self, args, params): a, b = args.a, args.b args.f[:] = (a + b + 1) * (a + b + 2) - def test_one_stream_set_serves_repeated_runs(self, st_platform, st_worker): - """N runs on one worker build one stream set, and every result is right.""" + def test_every_run_creates_its_own_aicore_stream(self, st_platform, st_worker): + """N runs create N AICore streams, and every result is right.""" if st_platform != "a2a3": pytest.skip("run stream sets are an a2a3 onboard resource") callable_obj = self.build_callable(st_platform) self._run_and_validate_l2(st_worker, callable_obj, self.CASES[0], rounds=1) after_first = st_worker.run_stream_set_create_count - self._run_and_validate_l2(st_worker, callable_obj, self.CASES[0], rounds=_REPEATED_RUNS - 1) + rounds = _REPEATED_RUNS - 1 + self._run_and_validate_l2(st_worker, callable_obj, self.CASES[0], rounds=rounds) - assert st_worker.run_stream_set_create_count == after_first, ( - f"same-image runs advanced stream generation after the first run: " + assert st_worker.run_stream_set_create_count == after_first + rounds, ( + f"{rounds} runs did not each create an AICore stream: " f"{after_first} -> {st_worker.run_stream_set_create_count}" ) @@ -143,23 +150,138 @@ def _run_registered(self, worker, handle, *, subtract): worker.run(handle, chip_args, config=self._build_config(self.CASES[0]["config"])) _compare_outputs(test_args, golden_args, output_names, self.RTOL, self.ATOL) - def test_aicore_stream_tracks_code_image(self, st_platform, st_worker): + # The st_worker fixture is shared by every test in this class, and so is the + # generation each slot has last admitted. Tests draw from one increasing + # counter so they stay independent of the order pytest runs them in. + _generation_counter = itertools.count(1) + + @classmethod + def _next_generation(cls): + return next(cls._generation_counter) + + def _run_registered_with_lease(self, worker, handle, *, slot_id, generation): + params = self.CASES[0]["params"] + test_args = self.generate_args(params) + golden_args = test_args.clone() + self.compute_golden(golden_args, params) + chip_args, output_names = _build_chip_task_args(test_args, self.CALLABLE["orchestration"]["signature"]) + state = worker._resolve_handle(handle) + worker._chip_worker._run_slot_with_pipeline_lease( + state.slot_id, + chip_args, + slot_id, + generation, + config=self._build_config(self.CASES[0]["config"]), + ) + _compare_outputs(test_args, golden_args, output_names, self.RTOL, self.ATOL) + + def test_alternating_code_images_never_execute_stale_instructions(self, st_platform, st_worker): + """A→B→A→B at one reused GM code address always runs the selected image. + + Each result is compared against the golden for the image that run asked + for, so an AICore that executed the previous image's instructions shows + up as wrong data rather than as a stream-count mismatch. + """ if st_platform != "a2a3": - pytest.skip("AICore stream code generations are an a2a3 onboard resource") + pytest.skip("AICore code images are an a2a3 onboard resource") add_handle = st_worker.register(self.build_callable(st_platform)) sub_handle = st_worker.register(_SubtractCallable.compile_chip_callable(st_platform)) try: - self._run_registered(st_worker, add_handle, subtract=False) - after_add = st_worker.run_stream_set_create_count + for handle, subtract in ( + (add_handle, False), + (sub_handle, True), + (add_handle, False), + (sub_handle, True), + ): + before = st_worker.run_stream_set_create_count + self._run_registered(st_worker, handle, subtract=subtract) + assert st_worker.run_stream_set_create_count == before + 1 + finally: + st_worker.unregister(sub_handle) + st_worker.unregister(add_handle) + + def test_depth_two_slots_own_separate_resources(self, st_platform, st_worker): + """Each slot owns its own host Runtime buffer and its own arena bank. + + Runs on simulation too: sim implements the same depth, so HBG — whose + GM heap is HOST_PER_RUN — must commit a distinct bank per slot there as + well. This is the only case that exercises two *committed* banks; the + TMR suite cannot, because its arenas are DEVICE_SCRATCH and both slots + correctly resolve to bank 0. + """ + if not st_platform.startswith("a2a3"): + pytest.skip("pipeline slot banks are an a2a3 resource") + + chip_worker = st_worker._chip_worker + assert chip_worker.pipeline_depth == 2 + assert chip_worker.runtime_slot_count == 2 + + addrs = chip_worker.runtime_buffer_addrs + assert len(addrs) == 2, addrs + assert addrs[0] != addrs[1], f"both pipeline slots stage into one Runtime buffer: {addrs}" + add_handle = st_worker.register(self.build_callable(st_platform)) + try: self._run_registered(st_worker, add_handle, subtract=False) - assert st_worker.run_stream_set_create_count == after_add + self._run_registered_with_lease(st_worker, add_handle, slot_id=1, generation=self._next_generation()) - for handle, subtract in ((sub_handle, True), (add_handle, False), (sub_handle, True)): - before_transition = st_worker.run_stream_set_create_count - self._run_registered(st_worker, handle, subtract=subtract) - assert st_worker.run_stream_set_create_count == before_transition + 1 + # hbg declares its GM heap HOST_PER_RUN, so a run on slot 1 must + # commit a second device allocation rather than reuse slot 0's. + bank0 = chip_worker.arena_bank_gm_heap_base(0) + bank1 = chip_worker.arena_bank_gm_heap_base(1) + assert bank0 != 0 and bank1 != 0, f"a served bank is uncommitted: {bank0:#x}, {bank1:#x}" + assert bank0 != bank1, f"both arena banks resolve to one GM heap: {bank0:#x}" + + # hbg stages device args directly, never through the retained + # temporary buffer, so neither slot should hold one. + assert chip_worker.retained_temp_addr(0) == 0 + assert chip_worker.retained_temp_addr(1) == 0 + finally: + st_worker.unregister(add_handle) + + def test_unleased_runs_do_not_consume_lease_generations(self, st_platform, st_worker): + """A slot pool's first lease is admitted however many runs precede it. + + `PipelineSlotPool` is the only mint, and it starts every slot at + generation 1. A consumer that minted generations for its own unleased + runs would have advanced past that, and the first real lease it was + ever handed would be rejected as stale. This needs a worker whose + generations no other test has touched, hence a second one. + """ + if st_platform != "a2a3": + pytest.skip("pipeline slot leases are an a2a3 onboard resource") + + other = Worker(level=2, **st_worker._config) + other.init() + try: + handle = other.register(self.build_callable(st_platform)) + try: + for _ in range(3): + self._run_registered(other, handle, subtract=False) + self._run_registered_with_lease(other, handle, slot_id=0, generation=1) + finally: + other.unregister(handle) + finally: + other.close() + + def test_depth_two_slot_is_generation_safe(self, st_platform, st_worker): + if st_platform != "a2a3": + pytest.skip("pipeline slot banks are an a2a3 onboard resource") + + add_handle = st_worker.register(self.build_callable(st_platform)) + superseded = self._next_generation() + generation = self._next_generation() + try: + self._run_registered_with_lease(st_worker, add_handle, slot_id=1, generation=generation) + + # One run may dispatch more than once under its lease. + self._run_registered_with_lease(st_worker, add_handle, slot_id=1, generation=generation) + + # A rejected lease must not reach launch, so it creates no stream. + stream_sets = st_worker.run_stream_set_create_count + with pytest.raises(RuntimeError, match="generation is stale"): + self._run_registered_with_lease(st_worker, add_handle, slot_id=1, generation=superseded) + assert st_worker.run_stream_set_create_count == stream_sets finally: - st_worker.unregister(sub_handle) st_worker.unregister(add_handle) diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/pipeline_slots/test_pipeline_slots.py b/tests/st/a2a3/tensormap_and_ringbuffer/pipeline_slots/test_pipeline_slots.py new file mode 100644 index 0000000000..7d5a112ff3 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/pipeline_slots/test_pipeline_slots.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""tensormap_and_ringbuffer's depth-two slot topology differs from host_build_graph's. + +TMR classifies its GM heap, GM shared memory, and runtime image as +`DEVICE_SCRATCH`, so every slot shares arena bank 0 — only its task args are +`HOST_PER_RUN`. What still separates per slot is the host `Runtime` staging +buffer and the retained temporary buffer TMR stages device args through. +host_build_graph exercises the opposite split, so neither runtime's coverage +substitutes for the other's. +""" + +import itertools + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.scene_test import _build_chip_task_args, _compare_outputs + +_VECTOR_KERNELS = "../../../../../examples/a2a3/tensormap_and_ringbuffer/vector_example/kernels" + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestPipelineSlotsTmr(SceneTestCase): + """A leased run on slot 1 must not disturb slot 0's per-run state.""" + + RTOL = 1e-5 + ATOL = 1e-5 + + CALLABLE = { + "orchestration": { + "source": f"{_VECTOR_KERNELS}/orchestration/example_orchestration.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{_VECTOR_KERNELS}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{_VECTOR_KERNELS}/aiv/kernel_add_scalar.cpp", + "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], + }, + ], + } + + CASES = [ + { + "name": "slot_topology", + "platforms": ["a2a3", "a2a3sim"], + "config": {"aicpu_thread_num": 4, "block_dim": 3}, + "params": {}, + }, + ] + + # Shared st_worker means a shared per-slot high-water mark; draw from one + # increasing counter so tests do not depend on pytest's ordering. + _generation_counter = itertools.count(1) + + def generate_args(self, params): + size = 128 * 128 + return TaskArgsBuilder( + Tensor("a", torch.full((size,), 2.0, dtype=torch.float32)), + Tensor("b", torch.full((size,), 3.0, dtype=torch.float32)), + Tensor("f", torch.zeros(size, dtype=torch.float32)), + ) + + def compute_golden(self, args, params): + a, b = args.a, args.b + args.f[:] = (a + b + 1) * (a + b + 2) + (a + b) + + def _run(self, worker, handle, *, slot_id=None): + """One run, on the ordinary path or through a lease, result checked.""" + test_args = self.generate_args({}) + golden = test_args.clone() + self.compute_golden(golden, {}) + chip_args, output_names = _build_chip_task_args(test_args, self.CALLABLE["orchestration"]["signature"]) + config = self._build_config(self.CASES[0]["config"]) + if slot_id is None: + worker.run(handle, chip_args, config=config) + else: + state = worker._resolve_handle(handle) + worker._chip_worker._run_slot_with_pipeline_lease( + state.slot_id, chip_args, slot_id, next(self._generation_counter), config=config + ) + _compare_outputs(test_args, golden, output_names, self.RTOL, self.ATOL) + + def test_slot_one_keeps_tmr_per_run_state_separate(self, st_platform, st_worker): + chip_worker = st_worker._chip_worker + assert chip_worker.pipeline_depth == 2 + assert chip_worker.runtime_slot_count == 2 + + addrs = chip_worker.runtime_buffer_addrs + assert len(addrs) == 2, addrs + assert addrs[0] != addrs[1], f"both slots stage into one host Runtime: {addrs}" + + handle = st_worker.register(self.build_callable(st_platform)) + try: + # Interleave the slots: if slot 1 shared slot 0's host Runtime or + # its retained temporary buffer, the run after the switch back + # would read arguments the other slot staged. + self._run(st_worker, handle) + self._run(st_worker, handle, slot_id=1) + self._run(st_worker, handle) + self._run(st_worker, handle, slot_id=1) + + # Sequential runs re-stage their arguments every round, so correct + # output alone would survive both slots sharing one staging buffer. + # Read the addresses instead. + temp0 = chip_worker.retained_temp_addr(0) + temp1 = chip_worker.retained_temp_addr(1) + assert temp0 != 0 and temp1 != 0, f"a slot that ran staged nothing: slot0={temp0:#x}, slot1={temp1:#x}" + assert temp0 != temp1, f"both slots stage through one retained buffer: {temp0:#x}" + + # TMR declares its three pooled regions DEVICE_SCRATCH, so both + # slots must still resolve to the one committed bank. + bank0 = chip_worker.arena_bank_gm_heap_base(0) + bank1 = chip_worker.arena_bank_gm_heap_base(1) + assert bank0 != 0, "the shared device-scratch bank was never committed" + assert bank1 == 0, f"slot 1 committed a second device-scratch bank: {bank1:#x}" + finally: + st_worker.unregister(handle) diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index 2e8cf1fcf7..0a5cee98ae 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -368,6 +368,19 @@ add_hierarchical_test(test_remote_wire hierarchical/test_remote_wire.cpp) add_hierarchical_test(test_remote_endpoint hierarchical/test_remote_endpoint.cpp) add_hierarchical_test(test_pipeline_contract hierarchical/test_pipeline_contract.cpp) +# Run-stream slot state machine: fresh AICore stream per run, and the handle a +# failed destroy must keep. Header-only with injected create/destroy, so it +# needs no runtime objects and no device. +add_executable(test_run_stream_slots hierarchical/test_run_stream_slots.cpp) +target_include_directories(test_run_stream_slots PRIVATE + ${GTEST_INCLUDE_DIRS} + ${CMAKE_SOURCE_DIR}/../../../src/common/platform/include + ${CMAKE_SOURCE_DIR}/../../../src/common/worker +) +target_link_libraries(test_run_stream_slots PRIVATE ${GTEST_MAIN_LIB} ${GTEST_LIB} pthread) +add_test(NAME test_run_stream_slots COMMAND test_run_stream_slots) +set_tests_properties(test_run_stream_slots PROPERTIES LABELS "no_hardware") + # --------------------------------------------------------------------------- # Types / task_interface tests (src/common/task_interface/) # --------------------------------------------------------------------------- diff --git a/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp b/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp index bd795a2a27..b4d8fedc32 100644 --- a/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp +++ b/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp @@ -12,6 +12,7 @@ #include #include "pipeline_contract.h" +#include "pipeline_slot_pool.h" namespace { @@ -31,52 +32,62 @@ PipelineContract accepted_contract() { TEST(PipelineContract, AcceptsADeclarationThisBuildCanHonor) { const PipelineContract c = accepted_contract(); - EXPECT_TRUE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_TRUE(is_valid_pipeline_contract(&c)); } // A runtime that exports no contract is handled by the caller, not here. -TEST(PipelineContract, RejectsNull) { EXPECT_FALSE(is_valid_depth1_pipeline_contract(nullptr)); } +TEST(PipelineContract, RejectsNull) { EXPECT_FALSE(is_valid_pipeline_contract(nullptr)); } TEST(PipelineContract, AcceptsAnEmptyResourceList) { PipelineContract c = accepted_contract(); c.resource_count = 0; - EXPECT_TRUE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_TRUE(is_valid_pipeline_contract(&c)); } TEST(PipelineContract, RejectsAnotherAbiVersion) { PipelineContract c = accepted_contract(); c.abi_version = PTO_PIPELINE_CONTRACT_ABI_VERSION + 1; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_FALSE(is_valid_pipeline_contract(&c)); c.abi_version = 0; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_FALSE(is_valid_pipeline_contract(&c)); } TEST(PipelineContract, RejectsMoreResourcesThanFit) { PipelineContract c = accepted_contract(); c.resource_count = PTO_PIPELINE_MAX_RESOURCES + 1; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_FALSE(is_valid_pipeline_contract(&c)); } -// The depth-1 gate is what keeps this build from honoring a contract whose -// resources it would have to replicate. -TEST(PipelineContract, RejectsADepthOtherThanOne) { +TEST(PipelineContract, AcceptsDepthTwoAndDerivesResourceCopies) { PipelineContract c = accepted_contract(); - c.pipeline_depth = 0; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)); c.pipeline_depth = 2; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)); - c.pipeline_depth = PTO_PIPELINE_MAX_DEPTH; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)); + ASSERT_TRUE(is_valid_pipeline_contract(&c)); + EXPECT_EQ(pipeline_resource_copy_count(c, c.resources[0]), 2u); + EXPECT_EQ(pipeline_resource_copy_count(c, c.resources[1]), 1u); + EXPECT_EQ(pipeline_resource_copy_count(c, c.resources[2]), 2u); + + const PipelineSlotLease second_slot{1, 0, 7}; + EXPECT_EQ(pipeline_resource_slot(c, c.resources[0], second_slot), 1u); + EXPECT_EQ(pipeline_resource_slot(c, c.resources[1], second_slot), 0u); + EXPECT_EQ(pipeline_resource_slot(c, c.resources[2], second_slot), 1u); +} + +TEST(PipelineContract, RejectsDepthOutsideSupportedRange) { + PipelineContract c = accepted_contract(); + c.pipeline_depth = 0; + EXPECT_FALSE(is_valid_pipeline_contract(&c)); + c.pipeline_depth = PTO_PIPELINE_MAX_DEPTH + 1; + EXPECT_FALSE(is_valid_pipeline_contract(&c)); } TEST(PipelineContract, RejectsAnOutOfRangeKindOrClass) { PipelineContract c = accepted_contract(); c.resources[0].kind = PTO_PIPELINE_AICORE_STREAM + 1; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_FALSE(is_valid_pipeline_contract(&c)); c = accepted_contract(); c.resources[0].resource_class = PTO_PIPELINE_EXEC_HANDLE + 1; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_FALSE(is_valid_pipeline_contract(&c)); } // bytes_per_copy is reserved: nothing sizes anything from it yet, so a runtime @@ -84,13 +95,13 @@ TEST(PipelineContract, RejectsAnOutOfRangeKindOrClass) { TEST(PipelineContract, RejectsANonZeroReservedSize) { PipelineContract c = accepted_contract(); c.resources[1].bytes_per_copy = 4096; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_FALSE(is_valid_pipeline_contract(&c)); } TEST(PipelineContract, RejectsAnUnspecifiedKind) { PipelineContract c = accepted_contract(); c.resources[1].kind = PTO_PIPELINE_KIND_UNSPECIFIED; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_FALSE(is_valid_pipeline_contract(&c)); } // A resource_count larger than the entries a runtime filled in leaves trailing @@ -106,9 +117,9 @@ TEST(PipelineContract, RejectsAResourceCountPastTheFilledEntries) { c.resources[0] = {filled, PTO_PIPELINE_HOST_PER_RUN, 0}; c.resource_count = 2; // overstates by exactly one - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)) << "filled kind " << filled; + EXPECT_FALSE(is_valid_pipeline_contract(&c)) << "filled kind " << filled; c.resource_count = 3; - EXPECT_FALSE(is_valid_depth1_pipeline_contract(&c)) << "filled kind " << filled; + EXPECT_FALSE(is_valid_pipeline_contract(&c)) << "filled kind " << filled; } } @@ -117,7 +128,7 @@ TEST(PipelineContract, RejectsAResourceCountPastTheFilledEntries) { TEST(PipelineContract, AcceptsARepeatedKind) { PipelineContract c = accepted_contract(); c.resources[3].kind = PTO_PIPELINE_AICPU_STREAM; - EXPECT_TRUE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_TRUE(is_valid_pipeline_contract(&c)); } TEST(PipelineContract, AcceptsEveryKindOnce) { @@ -128,7 +139,97 @@ TEST(PipelineContract, AcceptsEveryKindOnce) { for (uint32_t kind = PTO_PIPELINE_GM_HEAP; kind <= PTO_PIPELINE_AICORE_STREAM; ++kind) { c.resources[kind - 1] = {kind, PTO_PIPELINE_HOST_PER_RUN, 0}; } - EXPECT_TRUE(is_valid_depth1_pipeline_contract(&c)); + EXPECT_TRUE(is_valid_pipeline_contract(&c)); +} + +// GM heap, GM shared memory, and the runtime image are committed by one +// setup_static_arena call into one bank, so a contract that gives them +// different copy counts describes a layout no runner can build. +TEST(PipelineContract, RejectsArenaKindsThatDisagreeOnCopyCount) { + PipelineContract c{PTO_PIPELINE_CONTRACT_ABI_VERSION, 2, 2, {}}; + c.resources[0] = {PTO_PIPELINE_GM_HEAP, PTO_PIPELINE_HOST_PER_RUN, 0}; + c.resources[1] = {PTO_PIPELINE_GM_SM, PTO_PIPELINE_DEVICE_SCRATCH, 0}; + ASSERT_TRUE(is_valid_pipeline_contract(&c)); + EXPECT_FALSE(has_serviceable_arena_topology(c)); + + c.resources[1].resource_class = PTO_PIPELINE_HOST_PER_RUN; + EXPECT_TRUE(has_serviceable_arena_topology(c)); + + // At depth one every class collapses to one copy, so the same declaration + // is serviceable there. + c.resources[1].resource_class = PTO_PIPELINE_DEVICE_SCRATCH; + c.pipeline_depth = 1; + EXPECT_TRUE(has_serviceable_arena_topology(c)); +} + +// Only the first entry of a kind is ever read, so a second one is a silent +// misdeclaration for the kinds that select a bank. +TEST(PipelineContract, RejectsARepeatedArenaKind) { + PipelineContract c{PTO_PIPELINE_CONTRACT_ABI_VERSION, 2, 2, {}}; + c.resources[0] = {PTO_PIPELINE_GM_HEAP, PTO_PIPELINE_HOST_PER_RUN, 0}; + c.resources[1] = {PTO_PIPELINE_GM_HEAP, PTO_PIPELINE_HOST_PER_RUN, 0}; + ASSERT_TRUE(is_valid_pipeline_contract(&c)); + EXPECT_FALSE(has_serviceable_arena_topology(c)); +} + +// A repeated stream kind still selects nothing, so it stays legal. +TEST(PipelineContract, AcceptsARepeatedNonArenaKind) { + PipelineContract c{PTO_PIPELINE_CONTRACT_ABI_VERSION, 2, 2, {}}; + c.resources[0] = {PTO_PIPELINE_AICPU_STREAM, PTO_PIPELINE_EXEC_HANDLE, 0}; + c.resources[1] = {PTO_PIPELINE_AICPU_STREAM, PTO_PIPELINE_EXEC_HANDLE, 0}; + ASSERT_TRUE(is_valid_pipeline_contract(&c)); + EXPECT_TRUE(has_serviceable_arena_topology(c)); +} + +// Both shipped A2/A3 contracts must be serviceable as declared. +TEST(PipelineContract, ShippedArenaTopologiesAreServiceable) { + PipelineContract hbg{PTO_PIPELINE_CONTRACT_ABI_VERSION, 3, 2, {}}; + hbg.resources[0] = {PTO_PIPELINE_GM_HEAP, PTO_PIPELINE_HOST_PER_RUN, 0}; + hbg.resources[1] = {PTO_PIPELINE_GM_SM, PTO_PIPELINE_HOST_PER_RUN, 0}; + hbg.resources[2] = {PTO_PIPELINE_RUNTIME_IMAGE, PTO_PIPELINE_HOST_PER_RUN, 0}; + EXPECT_TRUE(has_serviceable_arena_topology(hbg)); + + PipelineContract tmr{PTO_PIPELINE_CONTRACT_ABI_VERSION, 3, 2, {}}; + tmr.resources[0] = {PTO_PIPELINE_GM_HEAP, PTO_PIPELINE_DEVICE_SCRATCH, 0}; + tmr.resources[1] = {PTO_PIPELINE_GM_SM, PTO_PIPELINE_DEVICE_SCRATCH, 0}; + tmr.resources[2] = {PTO_PIPELINE_RUNTIME_IMAGE, PTO_PIPELINE_DEVICE_SCRATCH, 0}; + EXPECT_TRUE(has_serviceable_arena_topology(tmr)); +} + +TEST(PipelineSlotPool, DepthOneKeepsLegacySingleSlotBehavior) { + PipelineSlotPool pool(1); + auto first = pool.try_acquire(); + ASSERT_TRUE(first.has_value()); + EXPECT_EQ(first->slot_id, 0u); + EXPECT_FALSE(pool.try_acquire().has_value()); + EXPECT_TRUE(pool.release(*first)); +} + +TEST(PipelineSlotPool, DepthTwoProvidesExactlyTwoIndependentLeases) { + PipelineSlotPool pool(2); + auto first = pool.try_acquire(); + auto second = pool.try_acquire(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(second.has_value()); + EXPECT_NE(first->slot_id, second->slot_id); + EXPECT_FALSE(pool.try_acquire().has_value()); + EXPECT_TRUE(pool.owns(*first)); + EXPECT_TRUE(pool.owns(*second)); +} + +TEST(PipelineSlotPool, StaleGenerationCannotAccessOrReleaseAReusedSlot) { + PipelineSlotPool pool(1); + const PipelineSlotLease first = *pool.try_acquire(); + ASSERT_TRUE(pool.release(first)); + EXPECT_TRUE(pool.release(first)); // idempotent before reuse + + const PipelineSlotLease replacement = *pool.try_acquire(); + ASSERT_EQ(replacement.slot_id, first.slot_id); + ASSERT_GT(replacement.generation, first.generation); + EXPECT_FALSE(pool.owns(first)); + EXPECT_FALSE(pool.release(first)); + EXPECT_TRUE(pool.owns(replacement)); + EXPECT_TRUE(pool.release(replacement)); } } // namespace diff --git a/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp b/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp new file mode 100644 index 0000000000..250b645663 --- /dev/null +++ b/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp @@ -0,0 +1,180 @@ +/* + * 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 + +#include "host/run_stream_slots.h" + +namespace { + +// Hands out distinct fake handles and can be told to fail the next N destroys, +// which is what makes the destroy-failure path reachable without a device. +class FakeStreams { +public: + int create(void **out) { + if (create_failures_ > 0) { + --create_failures_; + return -7; + } + *out = reinterpret_cast(++next_handle_); + live_.push_back(*out); + return 0; + } + + int destroy(void *stream) { + if (destroy_failures_ > 0) { + --destroy_failures_; + return -13; + } + for (auto it = live_.begin(); it != live_.end(); ++it) { + if (*it == stream) { + live_.erase(it); + return 0; + } + } + ADD_FAILURE() << "destroyed a handle that was not live"; + return -1; + } + + void fail_next_destroys(int n) { destroy_failures_ = n; } + void fail_next_creates(int n) { create_failures_ = n; } + size_t live_count() const { return live_.size(); } + +private: + std::vector live_; + uintptr_t next_handle_{0}; + int destroy_failures_{0}; + int create_failures_{0}; +}; + +RunStreamSlots make_slots(FakeStreams &fake) { + return RunStreamSlots( + [&fake](void **out) { + return fake.create(out); + }, + [&fake](void *s) { + return fake.destroy(s); + } + ); +} + +// The AICPU stream is the slot's for the runner's lifetime; the AICore stream +// belongs to one run, so the count advances once per acquire. +TEST(RunStreamSlots, EveryAcquireCreatesAnAicoreStreamAndKeepsTheAicpuOne) { + FakeStreams fake; + RunStreamSlots slots = make_slots(fake); + + ASSERT_EQ(slots.acquire(0), 0); + void *aicpu = slots.aicpu(0); + void *first_aicore = slots.aicore(0); + ASSERT_NE(aicpu, nullptr); + ASSERT_NE(first_aicore, nullptr); + EXPECT_EQ(slots.created_count(), 1u); + + ASSERT_EQ(slots.retire_aicore(0), 0); + EXPECT_EQ(slots.aicore(0), nullptr); + + ASSERT_EQ(slots.acquire(0), 0); + EXPECT_EQ(slots.aicpu(0), aicpu) << "the AICPU stream must not be recreated"; + EXPECT_NE(slots.aicore(0), first_aicore) << "the AICore stream must be a new one"; + EXPECT_EQ(slots.created_count(), 2u); +} + +// The three consequences a failed destroy must have. +TEST(RunStreamSlots, AFailedDestroyReportsKeepsTheHandleAndLocksTheSlot) { + FakeStreams fake; + RunStreamSlots slots = make_slots(fake); + + ASSERT_EQ(slots.acquire(0), 0); + void *stranded = slots.aicore(0); + + // 1. the error reaches the caller + fake.fail_next_destroys(1); + EXPECT_EQ(slots.retire_aicore(0), -13); + + // 2. the handle survives, so teardown still has something to reclaim + EXPECT_EQ(slots.aicore(0), stranded); + + // 3. the next run is refused rather than handed a second live stream + EXPECT_NE(slots.acquire(0), 0); + EXPECT_EQ(slots.aicore(0), stranded); + EXPECT_EQ(slots.created_count(), 1u) << "a refused acquire must not create anything"; + + // and teardown retries it + EXPECT_EQ(slots.destroy_all(), 0); + EXPECT_EQ(slots.aicore(0), nullptr); + EXPECT_EQ(fake.live_count(), 0u); +} + +// A slot locked by a failed destroy must not take the whole runner with it. +TEST(RunStreamSlots, AStrandedSlotDoesNotBlockItsPeer) { + FakeStreams fake; + RunStreamSlots slots = make_slots(fake); + + ASSERT_EQ(slots.acquire(0), 0); + fake.fail_next_destroys(1); + ASSERT_NE(slots.retire_aicore(0), 0); + + ASSERT_EQ(slots.acquire(1), 0); + EXPECT_NE(slots.aicore(1), nullptr); + EXPECT_EQ(slots.retire_aicore(1), 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) { + FakeStreams fake; + RunStreamSlots slots = make_slots(fake); + ASSERT_EQ(slots.acquire(0), 0); + ASSERT_EQ(slots.acquire(1), 0); + + fake.fail_next_destroys(1); + EXPECT_EQ(slots.destroy_all(), -13); + EXPECT_EQ(fake.live_count(), 1u) << "only the stream whose destroy failed may survive"; + + EXPECT_EQ(slots.destroy_all(), 0); + EXPECT_EQ(fake.live_count(), 0u); +} + +// A failed create leaves nothing half-owned behind. +TEST(RunStreamSlots, AFailedCreateLeavesTheSlotEmpty) { + FakeStreams fake; + RunStreamSlots slots = make_slots(fake); + + fake.fail_next_creates(1); + EXPECT_NE(slots.acquire(0), 0); + EXPECT_EQ(slots.aicpu(0), nullptr); + EXPECT_EQ(slots.aicore(0), nullptr); + EXPECT_EQ(slots.created_count(), 0u); + + // The AICPU stream lands, the AICore one fails: the slot stays unusable + // rather than half-ready. + fake.fail_next_creates(0); + ASSERT_EQ(slots.acquire(0), 0); + ASSERT_EQ(slots.retire_aicore(0), 0); + fake.fail_next_creates(1); + EXPECT_NE(slots.acquire(0), 0); + EXPECT_FALSE(slots.ready(0)); +} + +TEST(RunStreamSlots, RejectsASlotOutsideTheContractDepth) { + FakeStreams fake; + RunStreamSlots slots = make_slots(fake); + const unsigned past_end = static_cast(RunStreamSlots::capacity()); + EXPECT_NE(slots.acquire(past_end), 0); + EXPECT_NE(slots.retire_aicore(past_end), 0); + EXPECT_EQ(slots.aicore(past_end), nullptr); +} + +} // namespace