Refactor: isolate A2A3 run stream sets per pipeline slot

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"))

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)))

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
     # ------------------------------------------------------------------

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:
     //

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

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<DeviceRunnerBase *>(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 {

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",

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_;

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.