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