diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 15a20ca6c..15e95109b 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -1465,6 +1465,14 @@ NB_MODULE(_task_interface, m) { "host_build_graph variants. Mirrors aicpu_dlopen_count for the " "host-orchestration path; 0 on device-orch variants." ) + .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." + ) .def("malloc", &ChipWorker::malloc, nb::arg("size")) .def("free", &ChipWorker::free, nb::arg("ptr")) .def("copy_to", &ChipWorker::copy_to, nb::arg("dst"), nb::arg("src"), nb::arg("size")) diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index 597d2ca27..10e1acde6 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -1196,6 +1196,11 @@ def host_dlopen_count(self): """Number of host-side orch SO dlopens (host_build_graph variants).""" return self._impl.host_dlopen_count + @property + def run_stream_set_create_count(self): + """Number of run stream sets the bound runner has created.""" + return self._impl.run_stream_set_create_count + 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 12260f3fd..940ca1e53 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -5955,6 +5955,19 @@ def host_dlopen_count(self) -> int: return 0 return self._chip_worker.host_dlopen_count + @property + def run_stream_set_create_count(self) -> int: + """L2 only: 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. + 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 + return self._chip_worker.run_stream_set_create_count + # ------------------------------------------------------------------ # close # ------------------------------------------------------------------ diff --git a/src/a2a3/platform/onboard/host/device_runner.cpp b/src/a2a3/platform/onboard/host/device_runner.cpp index 484196b8c..094ca08f5 100644 --- a/src/a2a3/platform/onboard/host/device_runner.cpp +++ b/src/a2a3/platform/onboard/host/device_runner.cpp @@ -192,6 +192,7 @@ int DeviceRunner::destroy_comm_stream(void *stream) { // `src/common/platform/onboard/host/device_runner_base.cpp`. int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { + constexpr unsigned kPipelineSlot = 0; // 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); @@ -224,6 +225,12 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return rc; } + rc = ensure_run_stream_set(kPipelineSlot); + if (rc != 0) { + LOG_ERROR("ensure_run_stream_set(%u) failed: %d", kPipelineSlot, rc); + return rc; + } + ensure_device_wall_buffer(); if (block_dim < 1) { @@ -432,10 +439,10 @@ 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); + rc = launch_run(runtime, num_aicore, launch_aicpu_num, kPipelineSlot); if (rc != 0) return rc; - rc = reap_run(); + rc = reap_run(kPipelineSlot); if (rc != 0) return rc; // Print handshake results (reads from device memory, must be before free) @@ -444,10 +451,80 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return 0; } -int DeviceRunner::launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_num) { +int DeviceRunner::ensure_run_stream_set(unsigned slot) { + 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 && streams.aicore != nullptr) { + return 0; + } + + // Reached only on the first run for this slot, or after a partial create + // rolled one stream back; rtStreamCreate costs ~300 us and the set carries + // no per-run content, so it must not be rebuilt per run. + 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) { + int rc = rtStreamCreate(&streams.aicore, 0); + if (rc != 0) { + LOG_ERROR("rtStreamCreate (run AICore slot %u) failed: %d", slot, rc); + ACL_LOG_ERROR_DETAIL(rc); + return rc; + } + } + ++run_stream_sets_created_; + LOG_INFO_V0("DeviceRunner: run stream set %u created", slot); + return 0; +} + +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; + } + } + return rc; +} + +int DeviceRunner::launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_num, unsigned slot) { // 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) { + LOG_ERROR("launch_run: stream set %u is not ready", slot); + return -1; + } + RunStreamSet &streams = run_stream_sets_[slot]; int rc = 0; // Launch the AICore worker BEFORE the AICPU Run task — mirrors the a5 path @@ -475,7 +552,7 @@ int DeviceRunner::launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_ LOG_INFO_V0("=== launch_aicore_kernel ==="); // Launch AICore kernel (pass device copy of KernelArgs) - rc = launch_aicore_kernel(stream_aicore_, kernel_args_.device_k_args_); + rc = launch_aicore_kernel(streams.aicore, kernel_args_.device_k_args_); if (rc != 0) { LOG_ERROR("launch_aicore_kernel failed: %d", rc); recover_device_or_mark_unusable(rc); @@ -484,7 +561,7 @@ int DeviceRunner::launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_ LOG_INFO_V0("=== launch_aicpu_kernel %s ===", host::KernelNames::RunName); int aicpu_launch_n = (runtime.get_aicpu_launch_count() > 0) ? runtime.get_aicpu_launch_count() : launch_aicpu_num; - rc = launch_aicpu_kernel(stream_aicpu_, &kernel_args_.args, host::KernelNames::RunName, aicpu_launch_n); + rc = launch_aicpu_kernel(streams.aicpu, &kernel_args_.args, host::KernelNames::RunName, aicpu_launch_n); if (rc != 0) { LOG_ERROR("launch_aicpu_kernel (main) failed: %d", rc); // The AICore worker was already launched above and is now spinning in @@ -499,10 +576,15 @@ int DeviceRunner::launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_ return 0; } -int DeviceRunner::reap_run() { - int rc = sync_run_streams(); +int DeviceRunner::reap_run(unsigned slot) { + if (slot >= kRunStreamSetCount) { + 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); if (rc != 0) { - // sync_run_streams surfaces the AICore op-timeout (STARS-reaped op -> + // The pair wait surfaces the AICore op-timeout (STARS-reaped op -> // 507000/507018/507046 at AICPU/AICore stream sync). The op-timeout // leaves the device context poisoned for the SAME DeviceRunner's next // run, so attempt recovery / mark-unusable here too, not only on the @@ -803,10 +885,16 @@ int DeviceRunner::finalize() { // for the no-run-since-init case. finalize_collectors(); + // The per-run stream sets are this subclass's own RTS-owning member, so + // they are released here, while RTS is live and before the device reset + // below — the same window finalize_common() uses for the bootstrap pair. + int stream_rc = destroy_run_stream_sets(); + // Shared cleanup body — streams, kernel_args, callable/orch maps, // chip-callable buffer pool, the three arenas, device_wall, // mem_alloc_.finalize(), and cached arena sizes. rc = finalize_common(); + if (rc == 0) rc = stream_rc; // Reset device AFTER all device memory is freed. Two paths: // diff --git a/src/a2a3/platform/onboard/host/device_runner.h b/src/a2a3/platform/onboard/host/device_runner.h index f5a2231d8..80801721c 100644 --- a/src/a2a3/platform/onboard/host/device_runner.h +++ b/src/a2a3/platform/onboard/host/device_runner.h @@ -37,6 +37,7 @@ #include "callable.h" #include "prepare_callable_common.h" +#include "pto_runtime_c_api.h" // PTO_PIPELINE_MAX_DEPTH #include "common/kernel_args.h" #include "common/memory_barrier.h" #include "common/l2_swimlane_profiling.h" @@ -187,6 +188,8 @@ 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_; } + private: // Most lifecycle state (device_id_, block_dim_, cores_per_blockdim_, // worker_count_, executor + dispatcher bytes, aicore_bin_handle_, @@ -219,10 +222,26 @@ class DeviceRunner : public DeviceRunnerBase { // recovery. See run() and recover_device_or_mark_unusable(). bool device_unusable_{false}; - // Keep the kernel submission boundary separate from stream synchronization - // and teardown. run() still invokes these back-to-back in this change. - int launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_num); - int reap_run(); + struct RunStreamSet { + rtStream_t aicpu{nullptr}; + rtStream_t aicore{nullptr}; + }; + // One slot per in-flight run the pipeline contract can declare. + static constexpr unsigned kRunStreamSetCount = PTO_PIPELINE_MAX_DEPTH; + + // A stream set carries no per-run content, so it is created on first use + // and reused for every later run on this slot; `destroy_run_stream_sets()` + // in finalize() is the sole release point, as for the persistent bootstrap + // pair. Only slot 0 is selected while the contract is K=1. + RunStreamSet run_stream_sets_[kRunStreamSetCount]{}; + size_t run_stream_sets_created_{0}; + int ensure_run_stream_set(unsigned slot); + int destroy_run_stream_sets(); + + // The kernel submission boundary is separate from the stream wait and the + // post-run teardown; run() invokes the two back-to-back. + int launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_num, unsigned slot); + int reap_run(unsigned slot); // On an AICore launch/sync error, best-effort drain the device so a later // run() on the same DeviceRunner can recover in place; if the drain itself diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index 89be1376d..455d54442 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -671,6 +671,15 @@ size_t get_host_dlopen_count(DeviceContextHandle ctx) { } } +size_t get_run_stream_set_create_count(DeviceContextHandle ctx) { + if (ctx == NULL) return 0; + try { + return static_cast(ctx)->run_stream_set_create_count(); + } catch (...) { + return 0; + } +} + int simpler_provision_dma_workspace(DeviceContextHandle ctx, uint32_t required_mask) { 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 1fadff783..1c42b3fee 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -1273,9 +1273,11 @@ void DeviceRunnerBase::resolve_task_binary_addrs(Runtime &runtime) { LOG_DEBUG(""); } -int DeviceRunnerBase::sync_run_streams() { - LOG_INFO_V0("=== aclrtSynchronizeStreamWithTimeout stream_aicpu_ ==="); - int rc = aclrtSynchronizeStreamWithTimeout(stream_aicpu_, timeout_config_.stream_sync_timeout_ms); +int DeviceRunnerBase::sync_run_streams() { return sync_stream_pair(stream_aicpu_, stream_aicore_); } + +int DeviceRunnerBase::sync_stream_pair(rtStream_t aicpu_stream, rtStream_t aicore_stream) { + LOG_INFO_V0("=== aclrtSynchronizeStreamWithTimeout AICPU stream ==="); + int rc = aclrtSynchronizeStreamWithTimeout(aicpu_stream, timeout_config_.stream_sync_timeout_ms); if (rc == ACL_ERROR_RT_STREAM_SYNC_TIMEOUT) { LOG_ERROR( "Stream sync timeout: stream=AICPU timeout_ms=%d device_id=%d block_dim=%d", @@ -1290,8 +1292,8 @@ int DeviceRunnerBase::sync_run_streams() { return rc; } - LOG_INFO_V0("=== aclrtSynchronizeStreamWithTimeout stream_aicore_ ==="); - rc = aclrtSynchronizeStreamWithTimeout(stream_aicore_, timeout_config_.stream_sync_timeout_ms); + LOG_INFO_V0("=== aclrtSynchronizeStreamWithTimeout AICore stream ==="); + rc = aclrtSynchronizeStreamWithTimeout(aicore_stream, timeout_config_.stream_sync_timeout_ms); if (rc == ACL_ERROR_RT_STREAM_SYNC_TIMEOUT) { LOG_ERROR( "Stream sync timeout: stream=AICore timeout_ms=%d device_id=%d block_dim=%d", diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index ec0c6f770..b768c9953 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -419,6 +419,14 @@ class DeviceRunnerBase { */ size_t host_dlopen_count() const { return host_dlopen_total_; } + /** + * Number of run stream sets this runner has created. A set belongs to a + * pipeline slot and is reused for every run on that slot, so a runner that + * has served any number of runs on one slot reports 1. Arches whose runs + * use the persistent pair report 0. + */ + virtual size_t run_stream_set_create_count() const { return 0; } + /** * Device-orchestration callable registration used internally by * simpler_register_callable(): launches `simpler_aicpu_register_callable` with a @@ -681,6 +689,9 @@ class DeviceRunnerBase { */ int sync_run_streams(); + /** Wait for an explicit AICPU/AICore stream pair. */ + int sync_stream_pair(rtStream_t aicpu_stream, rtStream_t aicore_stream); + /** * Pull the device wall (ns) back from `device_wall_dev_ptr_` and * cache it on `device_wall_ns_`. D2H copy failure is a soft warn — @@ -911,7 +922,9 @@ class DeviceRunnerBase { // Persistent AICPU / AICore streams created in // `ensure_device_initialized()` and torn down in the subclass's - // `finalize()`. `nullptr` before init. + // `finalize()`. A2A3 reserves these for bootstrap/control operations and + // submits runs on its own per-slot stream sets; A5 submits runs on these. + // `nullptr` before init. rtStream_t stream_aicpu_{nullptr}; rtStream_t stream_aicore_{nullptr}; KernelArgsHelper kernel_args_; diff --git a/src/common/platform/sim/host/c_api_shared.cpp b/src/common/platform/sim/host/c_api_shared.cpp index ca57cd3a7..a1cceb231 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -618,6 +618,12 @@ size_t get_aicpu_dlopen_count(DeviceContextHandle ctx) { } } +size_t get_run_stream_set_create_count(DeviceContextHandle ctx) { + // Simulation has no ACL streams, so it owns no run stream sets. + (void)ctx; + return 0; +} + int simpler_provision_dma_workspace(DeviceContextHandle ctx, uint32_t required_mask) { // Simulation provides no async-DMA workspaces; a non-empty request fails // fast so an SDMA-enabled Worker cannot come up on sim. diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 251d5bd5b..b71e96ecb 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -121,6 +121,8 @@ void ChipWorker::init( unregister_callable_fn_ = load_symbol(handle, "simpler_unregister_callable"); get_aicpu_dlopen_count_fn_ = load_symbol(handle, "get_aicpu_dlopen_count"); get_host_dlopen_count_fn_ = load_symbol(handle, "get_host_dlopen_count"); + get_run_stream_set_create_count_fn_ = + load_symbol(handle, "get_run_stream_set_create_count"); simpler_provision_dma_workspace_fn_ = load_symbol(handle, "simpler_provision_dma_workspace"); finalize_device_fn_ = load_symbol(handle, "finalize_device"); @@ -217,6 +219,7 @@ void ChipWorker::init( unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; + get_run_stream_set_create_count_fn_ = nullptr; simpler_provision_dma_workspace_fn_ = nullptr; finalize_device_fn_ = nullptr; ensure_acl_ready_fn_ = nullptr; @@ -256,6 +259,7 @@ void ChipWorker::init( unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; + get_run_stream_set_create_count_fn_ = nullptr; simpler_provision_dma_workspace_fn_ = nullptr; finalize_device_fn_ = nullptr; ensure_acl_ready_fn_ = nullptr; @@ -325,6 +329,7 @@ void ChipWorker::finalize() { unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; + get_run_stream_set_create_count_fn_ = nullptr; simpler_provision_dma_workspace_fn_ = nullptr; finalize_device_fn_ = nullptr; ensure_acl_ready_fn_ = nullptr; @@ -404,6 +409,13 @@ size_t ChipWorker::host_dlopen_count() const { return get_host_dlopen_count_fn_(device_ctx_); } +size_t ChipWorker::run_stream_set_create_count() const { + if (!initialized_) { + return 0; + } + return get_run_stream_set_create_count_fn_(device_ctx_); +} + void *ChipWorker::create_comm_stream_checked(const char *op_name) { int rc = ensure_acl_ready_fn_(device_ctx_, device_id_); if (rc != 0) { diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 472f9e1ed..419137ea4 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -90,6 +90,13 @@ 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 sets the bound runner has created. A set belongs + /// to a pipeline slot and is reused for every run on that slot, so a + /// runner that has served any number of runs on one slot reports 1; + /// platforms whose runs use the persistent bootstrap pair report 0. Used + /// by tests to assert that repeated runs do not rebuild the set per run. + size_t run_stream_set_create_count() const; + uint64_t malloc(size_t size); void free(uint64_t ptr); void copy_to(uint64_t dst, uint64_t src, size_t size); @@ -208,6 +215,7 @@ class ChipWorker { SimplerUnregisterCallableFn unregister_callable_fn_ = nullptr; GetAicpuDlopenCountFn get_aicpu_dlopen_count_fn_ = nullptr; GetAicpuDlopenCountFn get_host_dlopen_count_fn_ = nullptr; + GetAicpuDlopenCountFn get_run_stream_set_create_count_fn_ = nullptr; SimplerProvisionDmaWorkspaceFn simpler_provision_dma_workspace_fn_ = nullptr; FinalizeDeviceFn finalize_device_fn_ = nullptr; EnsureAclReadyFn ensure_acl_ready_fn_ = nullptr; diff --git a/src/common/worker/pto_runtime_c_api.h b/src/common/worker/pto_runtime_c_api.h index 29c236412..d13773398 100644 --- a/src/common/worker/pto_runtime_c_api.h +++ b/src/common/worker/pto_runtime_c_api.h @@ -25,6 +25,7 @@ * copy_to_device_ctx, copy_from_device_ctx * - prepared run: simpler_register_callable, simpler_run, unregister_callable, * get_aicpu_dlopen_count, get_host_dlopen_count, + * get_run_stream_set_create_count, * simpler_provision_dma_workspace * - ACL/stream: ensure_acl_ready_ctx, create_comm_stream_ctx, * destroy_comm_stream_ctx @@ -305,6 +306,15 @@ size_t get_aicpu_dlopen_count(DeviceContextHandle ctx); */ size_t get_host_dlopen_count(DeviceContextHandle ctx); +/** + * Number of run stream sets the runner bound to `ctx` has created. A set + * belongs to a pipeline slot and is reused for every run on that slot, so a + * runner that has served any number of runs on one slot reports 1. Returns 0 + * on platforms whose runs use the persistent bootstrap pair. Used by tests to + * assert that repeated `simpler_run` calls do not rebuild the set per run. + */ +size_t get_run_stream_set_create_count(DeviceContextHandle ctx); + /** * Provision the async-DMA workspaces named in `required_mask` (a bitmask of * DmaWorkspaceKind bits) once at Worker init, latching their device addresses 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 new file mode 100644 index 000000000..8e0444ca4 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py @@ -0,0 +1,99 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""A2A3 run stream sets belong to a pipeline slot, not to a run. + +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. A set +carries no per-run content, so it is created on first use and reused; rebuilding +it per run costs two rtStreamCreate plus two rtStreamDestroy (~1.2 ms) on the +synchronous host path around KernelLaunch and buys nothing. + +`Worker.run_stream_set_create_count` reports how many sets the bound runner has +built, which is what makes that invariant assertable from here. +""" + +import pytest +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +_VECTOR_KERNELS = "../vector_example/kernels" +_REPEATED_RUNS = 4 + + +@scene_test(level=2, runtime="host_build_graph") +class TestRunStreamReuseHbg(SceneTestCase): + """Repeated runs on one worker must share a single run stream set.""" + + RTOL = 1e-5 + ATOL = 1e-5 + + CALLABLE = { + "orchestration": { + "source": f"{_VECTOR_KERNELS}/orchestration/example_orch.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": "repeated_runs", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "block_dim": 3}, + "params": {}, + }, + ] + + 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) + + 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.""" + 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=_REPEATED_RUNS) + + assert st_worker.run_stream_set_create_count == 1, ( + f"expected 1 run stream set for {_REPEATED_RUNS} runs, got " + f"{st_worker.run_stream_set_create_count} — the set is being rebuilt per run" + )