From 7a17464672ee78982c6b4fc70642d642e7d3027b Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Mon, 27 Jul 2026 21:16:46 +0800 Subject: [PATCH 1/2] Fix: bind AICore run streams to code images Compute a stable identity from each callable's AICore child binaries and function mapping, then carry it through callable registration. Reuse a slot's AICPU stream independently, but recreate its AICore stream whenever the loaded image changes. Extend hardware coverage with alternating add/sub images and retain same-image reuse checks. --- python/simpler/task_interface.py | 2 +- python/simpler/worker.py | 10 +-- .../platform/onboard/host/device_runner.cpp | 46 ++++++---- .../platform/onboard/host/device_runner.h | 11 +-- .../host_build_graph/host/runtime_maker.cpp | 3 +- .../host/runtime_maker.cpp | 3 +- .../host_build_graph/host/runtime_maker.cpp | 3 +- .../host/runtime_maker.cpp | 3 +- .../platform/onboard/host/c_api_shared.cpp | 10 +-- .../onboard/host/device_runner_base.cpp | 12 +-- .../onboard/host/device_runner_base.h | 17 ++-- .../task_interface/chip_callable_layout.h | 19 +++- .../task_interface/prepare_callable_common.h | 4 +- src/common/utils/fnv1a_64.h | 14 +-- src/common/worker/chip_worker.h | 8 +- src/common/worker/pto_runtime_c_api.h | 9 +- .../kernels/aiv/kernel_sub.cpp | 64 ++++++++++++++ .../run_stream_reuse/test_run_stream_reuse.py | 86 ++++++++++++++++--- .../test_chip_callable_upload_immutable.cpp | 32 +++++++ 19 files changed, 278 insertions(+), 78 deletions(-) create mode 100644 tests/st/a2a3/host_build_graph/run_stream_reuse/kernels/aiv/kernel_sub.cpp diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index b18142329c..5c20d9bd9e 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -1264,7 +1264,7 @@ def host_dlopen_count(self): @property def run_stream_set_create_count(self): - """Number of run stream sets the bound runner has created.""" + """Number of run stream generations the bound runner has created.""" return self._impl.run_stream_set_create_count def malloc(self, size): diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 2c3fa7bb6c..bcb70e9970 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -6155,12 +6155,12 @@ def host_dlopen_count(self) -> int: @property def run_stream_set_create_count(self) -> int: - """L2 only: number of run stream sets the bound runner has created. + """L2 only: number of run stream generations the 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). + 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). """ 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 0ab734f798..5cd2536015 100644 --- a/src/a2a3/platform/onboard/host/device_runner.cpp +++ b/src/a2a3/platform/onboard/host/device_runner.cpp @@ -248,9 +248,16 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return rc; } - rc = ensure_run_stream_set(kPipelineSlot); + 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); if (rc != 0) { - LOG_ERROR("ensure_run_stream_set(%u) failed: %d", kPipelineSlot, rc); + LOG_ERROR("ensure_run_stream_set(%u, image=0x%lx) failed: %d", kPipelineSlot, aicore_image_hash, rc); return rc; } @@ -480,19 +487,12 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return 0; } -int DeviceRunner::ensure_run_stream_set(unsigned slot) { +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 && 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) { @@ -501,16 +501,31 @@ int DeviceRunner::ensure_run_stream_set(unsigned slot) { return rc; } } - if (streams.aicore == nullptr) { - int rc = rtStreamCreate(&streams.aicore, 0); + 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("rtStreamCreate (run AICore slot %u) failed: %d", slot, rc); - ACL_LOG_ERROR_DETAIL(rc); + LOG_ERROR("rtStreamDestroy (retired AICore slot %u) failed: %d", slot, rc); return rc; } + streams.aicore = nullptr; + streams.has_aicore_image = false; + } + + 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); + streams.aicore = nullptr; + return rc; } + streams.aicore_image_hash = aicore_image_hash; + streams.has_aicore_image = true; ++run_stream_sets_created_; - LOG_INFO_V0("DeviceRunner: run stream set %u created", slot); + LOG_INFO_V0("DeviceRunner: run stream generation %zu created for slot %u", run_stream_sets_created_, slot); return 0; } @@ -539,6 +554,7 @@ int DeviceRunner::destroy_run_stream_sets() { } capture(destroy_rc); streams.aicore = nullptr; + streams.has_aicore_image = false; } } return rc; diff --git a/src/a2a3/platform/onboard/host/device_runner.h b/src/a2a3/platform/onboard/host/device_runner.h index be258712b0..cc5544415f 100644 --- a/src/a2a3/platform/onboard/host/device_runner.h +++ b/src/a2a3/platform/onboard/host/device_runner.h @@ -228,17 +228,18 @@ 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; - // 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. + // 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); + int ensure_run_stream_set(unsigned slot, uint64_t aicore_image_hash); int destroy_run_stream_sets(); // The kernel submission boundary is separate from the stream wait and the 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 f8de55082c..80cc399a0c 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -558,7 +558,8 @@ register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(const LOG_INFO_V0("Registering %d kernel(s) in register_callable_impl", callable->child_count()); if (upload_and_collect_child_addrs( - callable, upload_fn, &out->kernel_addrs, &out->chip_buffer_dev, &out->chip_buffer_hash + callable, upload_fn, &out->kernel_addrs, &out->chip_buffer_dev, &out->chip_buffer_hash, + &out->aicore_image_hash ) != 0) { LOG_ERROR("Failed to upload ChipCallable buffer"); return -1; 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 c7ce87c6c4..9bb1b611cd 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -416,7 +416,8 @@ register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(const LOG_INFO_V0("Registering %d kernel(s) in register_callable_impl", callable->child_count()); if (upload_and_collect_child_addrs( - callable, upload_fn, &out->kernel_addrs, &out->chip_buffer_dev, &out->chip_buffer_hash + callable, upload_fn, &out->kernel_addrs, &out->chip_buffer_dev, &out->chip_buffer_hash, + &out->aicore_image_hash ) != 0) { LOG_ERROR("Failed to upload ChipCallable buffer"); return -1; diff --git a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp index e3b55486aa..ff55ae60ad 100644 --- a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp @@ -322,7 +322,8 @@ int register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(c LOG_INFO_V0("Registering %d kernel(s) in register_callable_impl", callable->child_count()); if (upload_and_collect_child_addrs( - callable, upload_fn, &out->kernel_addrs, &out->chip_buffer_dev, &out->chip_buffer_hash + callable, upload_fn, &out->kernel_addrs, &out->chip_buffer_dev, &out->chip_buffer_hash, + &out->aicore_image_hash ) != 0) { LOG_ERROR("Failed to upload ChipCallable buffer"); return -1; diff --git a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index 187a66d91f..07754e0693 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -393,7 +393,8 @@ register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(const LOG_INFO_V0("Registering %d kernel(s) in register_callable_impl", callable->child_count()); if (upload_and_collect_child_addrs( - callable, upload_fn, &out->kernel_addrs, &out->chip_buffer_dev, &out->chip_buffer_hash + callable, upload_fn, &out->kernel_addrs, &out->chip_buffer_dev, &out->chip_buffer_hash, + &out->aicore_image_hash ) != 0) { LOG_ERROR("Failed to upload ChipCallable buffer"); return -1; diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index bd8c668e95..03e3f0934a 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -434,17 +434,17 @@ int simpler_register_callable(DeviceContextHandle ctx, int32_t callable_id, cons bool needs_aicpu_register = false; if (artifacts.host_dlopen_handle != nullptr) { rc = runner->record_host_orch_callable( - callable_id, artifacts.chip_buffer_hash, artifacts.host_dlopen_handle, artifacts.host_orch_func_ptr, - std::move(kernel_addrs), std::move(artifacts.signature) + callable_id, artifacts.chip_buffer_hash, artifacts.aicore_image_hash, artifacts.host_dlopen_handle, + artifacts.host_orch_func_ptr, std::move(kernel_addrs), std::move(artifacts.signature) ); if (rc != 0) return rc; host_dlopen_guard.dismiss(); chip_buffer_guard.dismiss(); } else { rc = runner->record_device_orch_callable( - callable_id, artifacts.chip_buffer_hash, artifacts.chip_buffer_dev, artifacts.orch_so_data, - artifacts.orch_so_size, artifacts.func_name.c_str(), artifacts.config_name.c_str(), - std::move(kernel_addrs), std::move(artifacts.signature) + callable_id, artifacts.chip_buffer_hash, artifacts.aicore_image_hash, artifacts.chip_buffer_dev, + artifacts.orch_so_data, artifacts.orch_so_size, artifacts.func_name.c_str(), + artifacts.config_name.c_str(), std::move(kernel_addrs), std::move(artifacts.signature) ); if (rc != 0) return rc; chip_buffer_guard.dismiss(); diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index 0af3d711c2..30c2bcd240 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -735,9 +735,9 @@ int DeviceRunnerBase::launch_device_register(int32_t callable_id) { } int DeviceRunnerBase::record_device_orch_callable( - int32_t callable_id, uint64_t chip_buffer_hash, uint64_t chip_dev, const void *orch_so_data, size_t orch_so_size, - const char *func_name, const char *config_name, std::vector> kernel_addrs, - std::vector signature + int32_t callable_id, uint64_t chip_buffer_hash, uint64_t aicore_image_hash, uint64_t chip_dev, + const void *orch_so_data, size_t orch_so_size, const char *func_name, const char *config_name, + std::vector> kernel_addrs, std::vector signature ) { // The AICPU executor reserves `orch_so_table_[MAX_REGISTERED_CALLABLE_IDS]` // (declared in src/common/task_interface/callable_protocol.h) and indexes @@ -767,6 +767,7 @@ int DeviceRunnerBase::record_device_orch_callable( CallableState state; state.hash = hash; state.chip_buffer_hash = chip_buffer_hash; + state.aicore_image_hash = aicore_image_hash; state.dev_orch_so_addr = chip_dev + offsetof(ChipCallable, storage_); state.dev_orch_so_size = orch_so_size; state.func_name = (func_name != nullptr) ? func_name : ""; @@ -782,8 +783,8 @@ int DeviceRunnerBase::record_device_orch_callable( } int DeviceRunnerBase::record_host_orch_callable( - int32_t callable_id, uint64_t chip_buffer_hash, void *host_dlopen_handle, void *host_orch_func_ptr, - std::vector> kernel_addrs, std::vector signature + int32_t callable_id, uint64_t chip_buffer_hash, uint64_t aicore_image_hash, void *host_dlopen_handle, + void *host_orch_func_ptr, std::vector> kernel_addrs, std::vector signature ) { if (callable_id < 0 || callable_id >= MAX_REGISTERED_CALLABLE_IDS) { LOG_ERROR( @@ -806,6 +807,7 @@ int DeviceRunnerBase::record_host_orch_callable( CallableState state; state.chip_buffer_hash = chip_buffer_hash; + state.aicore_image_hash = aicore_image_hash; state.host_dlopen_handle = host_dlopen_handle; state.host_orch_func_ptr = host_orch_func_ptr; state.kernel_addrs = std::move(kernel_addrs); diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index b4101ec957..4953ee8d84 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -303,8 +303,8 @@ class DeviceRunnerBase { * @return 0 on success, negative on failure. */ int record_device_orch_callable( - int32_t callable_id, uint64_t chip_buffer_hash, uint64_t chip_dev, const void *orch_so_data, - size_t orch_so_size, const char *func_name, const char *config_name, + int32_t callable_id, uint64_t chip_buffer_hash, uint64_t aicore_image_hash, uint64_t chip_dev, + const void *orch_so_data, size_t orch_so_size, const char *func_name, const char *config_name, std::vector> kernel_addrs, std::vector signature ); @@ -319,8 +319,9 @@ class DeviceRunnerBase { * dlclose'd by `unregister_callable`. Increments `host_dlopen_total_`. */ int record_host_orch_callable( - int32_t callable_id, uint64_t chip_buffer_hash, void *host_dlopen_handle, void *host_orch_func_ptr, - std::vector> kernel_addrs, std::vector signature + int32_t callable_id, uint64_t chip_buffer_hash, uint64_t aicore_image_hash, void *host_dlopen_handle, + void *host_orch_func_ptr, std::vector> kernel_addrs, + std::vector signature ); /** @@ -420,10 +421,9 @@ 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. + * Number of run stream generations this runner has created. AICPU streams + * belong to pipeline slots, while an AICore stream is reused only for the + * same AICore image. Arches whose runs use the persistent pair report 0. */ virtual size_t run_stream_set_create_count() const { return 0; } @@ -806,6 +806,7 @@ class DeviceRunnerBase { // chip_buffer_hash, which keys the retained buffer. uint64_t hash{0}; uint64_t chip_buffer_hash{0}; + uint64_t aicore_image_hash{0}; uint64_t dev_orch_so_addr{0}; size_t dev_orch_so_size{0}; std::string func_name; diff --git a/src/common/task_interface/chip_callable_layout.h b/src/common/task_interface/chip_callable_layout.h index a0d338a59b..e19f945592 100644 --- a/src/common/task_interface/chip_callable_layout.h +++ b/src/common/task_interface/chip_callable_layout.h @@ -30,9 +30,10 @@ #include "utils/fnv1a_64.h" struct ChipCallableLayout { - size_t header_size; // offsetof(ChipCallable, storage_) - size_t total_size; // header_size + storage_used (matches make_callable()) - uint64_t content_hash; // FNV-1a 64 over [callable, total_size) + size_t header_size; // offsetof(ChipCallable, storage_) + size_t total_size; // header_size + storage_used (matches make_callable()) + uint64_t content_hash; // FNV-1a 64 over [callable, total_size) + uint64_t aicore_image_hash; // FNV-1a 64 over func ids and child binaries }; /** @@ -45,15 +46,25 @@ struct ChipCallableLayout { inline ChipCallableLayout compute_chip_callable_layout(const ChipCallable *callable) { constexpr size_t kHeaderSize = offsetof(ChipCallable, storage_); size_t storage_used = static_cast(callable->binary_size()); + const int32_t child_count = callable->child_count(); + uint64_t aicore_image_hash = simpler::common::utils::fnv1a_64(&child_count, sizeof(child_count)); for (int32_t i = 0; i < callable->child_count(); ++i) { const CoreCallable &c = callable->child(i); + const int32_t func_id = callable->child_func_id(i); + const uint32_t binary_size = c.binary_size(); + aicore_image_hash = simpler::common::utils::fnv1a_64_append(aicore_image_hash, &func_id, sizeof(func_id)); + aicore_image_hash = + simpler::common::utils::fnv1a_64_append(aicore_image_hash, &binary_size, sizeof(binary_size)); + aicore_image_hash = simpler::common::utils::fnv1a_64_append( + aicore_image_hash, c.binary_data(), static_cast(binary_size) + ); size_t child_total = CoreCallable::binary_data_offset() + static_cast(c.binary_size()); size_t end = static_cast(callable->child_offset(i)) + child_total; if (end > storage_used) storage_used = end; } const size_t total_size = kHeaderSize + storage_used; const uint64_t hash = simpler::common::utils::fnv1a_64(reinterpret_cast(callable), total_size); - return ChipCallableLayout{kHeaderSize, total_size, hash}; + return ChipCallableLayout{kHeaderSize, total_size, hash, aicore_image_hash}; } /** diff --git a/src/common/task_interface/prepare_callable_common.h b/src/common/task_interface/prepare_callable_common.h index d1e428255e..5c760ecb2c 100644 --- a/src/common/task_interface/prepare_callable_common.h +++ b/src/common/task_interface/prepare_callable_common.h @@ -65,6 +65,7 @@ struct CallableArtifacts { void *host_dlopen_handle{nullptr}; // hbg only void *host_orch_func_ptr{nullptr}; // hbg only uint64_t chip_buffer_hash{0}; // FNV-1a hash for the whole ChipCallable buffer + uint64_t aicore_image_hash{0}; // FNV-1a hash for func ids and AICore child binaries uint64_t chip_buffer_dev{0}; // device address of the ChipCallable header const void *orch_so_data{nullptr}; // trb only; host view used for validation/hash only size_t orch_so_size{0}; // trb only @@ -90,7 +91,7 @@ struct CallableArtifacts { */ inline int upload_and_collect_child_addrs( const ChipCallable *callable, uint64_t (*upload_fn)(const void *), std::vector *out, - uint64_t *out_chip_dev = nullptr, uint64_t *out_chip_hash = nullptr + uint64_t *out_chip_dev = nullptr, uint64_t *out_chip_hash = nullptr, uint64_t *out_aicore_image_hash = nullptr ) { if (callable == nullptr || upload_fn == nullptr || out == nullptr) return -1; out->clear(); @@ -100,6 +101,7 @@ inline int upload_and_collect_child_addrs( if (chip_dev == 0) return -1; if (out_chip_dev != nullptr) *out_chip_dev = chip_dev; if (out_chip_hash != nullptr) *out_chip_hash = layout.content_hash; + if (out_aicore_image_hash != nullptr) *out_aicore_image_hash = layout.aicore_image_hash; out->reserve(static_cast(callable->child_count())); for (int32_t i = 0; i < callable->child_count(); ++i) { diff --git a/src/common/utils/fnv1a_64.h b/src/common/utils/fnv1a_64.h index b9b6763408..7be0e71724 100644 --- a/src/common/utils/fnv1a_64.h +++ b/src/common/utils/fnv1a_64.h @@ -20,15 +20,19 @@ namespace simpler::common::utils { // FNV-1a 64-bit content hash. Deterministic, allocation-free, ~µs / MB. // Used as a generic content-keyed dedup key (ChipCallable buffer hashing in // DeviceRunner, ELF Build-ID fallback in elf_build_id.h, etc.). -inline uint64_t fnv1a_64(const void *data, std::size_t len) { +inline uint64_t fnv1a_64_append(uint64_t hash, const void *data, std::size_t len) { constexpr uint64_t kPrime = 0x00000100000001b3ULL; - uint64_t h = 0xcbf29ce484222325ULL; const auto *p = static_cast(data); for (std::size_t i = 0; i < len; ++i) { - h ^= p[i]; - h *= kPrime; + hash ^= p[i]; + hash *= kPrime; } - return h; + return hash; +} + +inline uint64_t fnv1a_64(const void *data, std::size_t len) { + constexpr uint64_t kOffsetBasis = 0xcbf29ce484222325ULL; + return fnv1a_64_append(kOffsetBasis, data, len); } } // namespace simpler::common::utils diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 419137ea4c..7ead68ff33 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -90,11 +90,9 @@ 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. + /// 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. size_t run_stream_set_create_count() const; uint64_t malloc(size_t size); diff --git a/src/common/worker/pto_runtime_c_api.h b/src/common/worker/pto_runtime_c_api.h index 8d659ee284..6d20c3b878 100644 --- a/src/common/worker/pto_runtime_c_api.h +++ b/src/common/worker/pto_runtime_c_api.h @@ -307,11 +307,10 @@ 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. + * 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. */ size_t get_run_stream_set_create_count(DeviceContextHandle ctx); diff --git a/tests/st/a2a3/host_build_graph/run_stream_reuse/kernels/aiv/kernel_sub.cpp b/tests/st/a2a3/host_build_graph/run_stream_reuse/kernels/aiv/kernel_sub.cpp new file mode 100644 index 0000000000..2226f5b815 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/run_stream_reuse/kernels/aiv/kernel_sub.cpp @@ -0,0 +1,64 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *src0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *src1_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + + __gm__ float *src0 = reinterpret_cast<__gm__ float *>(src0_tensor->buffer.addr) + src0_tensor->start_offset; + __gm__ float *src1 = reinterpret_cast<__gm__ float *>(src1_tensor->buffer.addr) + src1_tensor->start_offset; + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + constexpr int kRows = 128; + constexpr int kCols = 128; + using Shape5D = Shape<1, 1, 1, kRows, kCols>; + using Stride5D = Stride<1, 1, 1, kCols, 1>; + using GlobalData = GlobalTensor; + using TileData = Tile; + + TileData src0_tile(kRows, kCols); + TileData src1_tile(kRows, kCols); + TileData dst_tile(kRows, kCols); + TASSIGN(src0_tile, 0x0); + TASSIGN(src1_tile, 0x10000); + TASSIGN(dst_tile, 0x20000); + + GlobalData src0_global(src0); + GlobalData src1_global(src1); + GlobalData dst_global(out); + TLOAD(src0_tile, src0_global); + TLOAD(src1_tile, src1_global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TSUB(dst_tile, src0_tile, src1_tile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(dst_global, dst_tile); + pipe_sync(); +} diff --git a/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py b/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py index 8e0444ca4a..049927cdf3 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,16 +7,16 @@ # 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 run streams are reused only while their code image is unchanged. 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. +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. `Worker.run_stream_set_create_count` reports how many sets the bound runner has -built, which is what makes that invariant assertable from here. +built, including every fresh AICore stream created for a code transition. """ import pytest @@ -24,11 +24,43 @@ 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 = "../vector_example/kernels" _REPEATED_RUNS = 4 +@scene_test(level=2, runtime="host_build_graph") +class _SubtractCallable(SceneTestCase): + 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": "kernels/aiv/kernel_sub.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], + }, + ], + } + + @scene_test(level=2, runtime="host_build_graph") class TestRunStreamReuseHbg(SceneTestCase): """Repeated runs on one worker must share a single run stream set.""" @@ -91,9 +123,43 @@ def test_one_stream_set_serves_repeated_runs(self, st_platform, st_worker): 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) + 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) - 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" + assert st_worker.run_stream_set_create_count == after_first, ( + f"same-image runs advanced stream generation after the first run: " + f"{after_first} -> {st_worker.run_stream_set_create_count}" ) + + def _run_registered(self, worker, handle, *, subtract): + params = self.CASES[0]["params"] + test_args = self.generate_args(params) + chip_args, output_names = _build_chip_task_args(test_args, self.CALLABLE["orchestration"]["signature"]) + golden_args = test_args.clone() + a, b = golden_args.a, golden_args.b + base = a - b if subtract else a + b + golden_args.f[:] = (base + 1) * (base + 2) + 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): + if st_platform != "a2a3": + pytest.skip("AICore stream code generations 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 + + self._run_registered(st_worker, add_handle, subtract=False) + assert st_worker.run_stream_set_create_count == after_add + + 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 + finally: + st_worker.unregister(sub_handle) + st_worker.unregister(add_handle) diff --git a/tests/ut/cpp/types/test_chip_callable_upload_immutable.cpp b/tests/ut/cpp/types/test_chip_callable_upload_immutable.cpp index 1ef5301c85..7bd6b63bde 100644 --- a/tests/ut/cpp/types/test_chip_callable_upload_immutable.cpp +++ b/tests/ut/cpp/types/test_chip_callable_upload_immutable.cpp @@ -36,6 +36,7 @@ #include #include "callable.h" +#include "chip_callable_layout.h" #include "utils/fnv1a_64.h" namespace { @@ -135,3 +136,34 @@ TEST(ChipCallableUploadImmutable, ChildOffsetsAreCallableAligned) { << " is not a multiple of CALLABLE_ALIGN=" << CALLABLE_ALIGN; } } + +TEST(ChipCallableImageHash, IgnoresOrchestrationBytes) { + auto first = build_test_chip_callable(); + auto second = first; + auto *second_callable = reinterpret_cast(second.data()); + ASSERT_GT(second_callable->binary_size(), 0u); + second_callable->storage_[0] ^= 0xff; + + const auto first_layout = compute_chip_callable_layout(reinterpret_cast(first.data())); + const auto second_layout = compute_chip_callable_layout(second_callable); + EXPECT_NE(first_layout.content_hash, second_layout.content_hash); + EXPECT_EQ(first_layout.aicore_image_hash, second_layout.aicore_image_hash); +} + +TEST(ChipCallableImageHash, ChangesWithKernelBytesOrFunctionMapping) { + auto original = build_test_chip_callable(); + auto changed_binary = original; + auto *binary_callable = reinterpret_cast(changed_binary.data()); + auto *binary_child = reinterpret_cast(binary_callable) + offsetof(ChipCallable, storage_) + + binary_callable->child_offset(0) + CoreCallable::binary_data_offset(); + *binary_child ^= 0xff; + + auto changed_mapping = original; + auto *mapping_callable = reinterpret_cast(changed_mapping.data()); + mapping_callable->child_func_ids_[0] += 1; + + const auto original_hash = + compute_chip_callable_layout(reinterpret_cast(original.data())).aicore_image_hash; + EXPECT_NE(original_hash, compute_chip_callable_layout(binary_callable).aicore_image_hash); + EXPECT_NE(original_hash, compute_chip_callable_layout(mapping_callable).aicore_image_hash); +} From a89fb11a9e4ca2c1892f90b44adec5f85552b9b1 Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Mon, 27 Jul 2026 20:27:21 +0800 Subject: [PATCH 2/2] Feat: add launch acceptance flight fence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Separate the point at which a run's dispatches have been launched from the point at which they complete, so a later submit() can build its graph while the previous run is still executing on the device. RunState carries pending_accepts: submit_impl raises it by the dispatch's group size, and each endpoint dispatch lowers it once after its launch is accepted. A2A3 onboard publishes acceptance after both the AICore and AICPU launches succeed, and the parent observes it inside the poll it already runs, without releasing mailbox ownership before TASK_DONE. Endpoints with no launch ACK — the remote protocol, A5, sim — keep the conservative fallback: WorkerThread satisfies the fence when the dispatch returns, so it degrades to completion rather than hanging. Acceptance is a sticky mailbox word, not a MailboxState. A state carrying it is lost whenever the child reaches TASK_DONE between two parent polls, which is exactly the short-task case the fence exists to pipeline; the loss is silent and leaves the next submit waiting for completion. The child sets the word before TASK_DONE and the parent clears it only when it publishes the next TASK_READY, and the parent reads it after the state so a TASK_DONE observation implies the word is visible. Acceptance is elected separately from completion on RunHandle. worker.run() is submit().wait(), so a thread parked on the previous handle's completion is the normal case, and queueing the acceptance drain behind that election would make every drain a completion drain. A waiter that races the run's release observes a released run as accepted rather than raising. The fence advances from a WorkerThread, where an escaping exception ends the process, so the accept path is non-throwing throughout: an unknown run or slot is ignored, a count mismatch is recorded as the run's error the way decrement_run_tasks already does, and the post-dispatch fallback cannot propagate out of the thread. It advances even when the dispatch failed, or a waiter would block forever. The decrement runs under completion_mu, the mutex a waiter evaluates its predicate on. Without it the notify can land between that predicate check and the block and be lost, leaving the waiter parked with the fence satisfied. WorkerManager::start takes on_accept without a default: an omitted callback leaves pending_accepts non-zero forever, which is a hang rather than a degraded mode. This is an acceptance bridge, not multi-run admission. There is one mailbox frame, no generation-safe frame reuse and no K=2 slot lease; for a dependent DAG, acceptance can land close to completion. --- docs/orchestrator.md | 11 +- docs/task-flow.md | 9 ++ docs/worker-manager.md | 39 ++++- python/bindings/task_interface.cpp | 7 +- python/bindings/worker_bind.h | 6 + python/simpler/orchestrator.py | 6 + python/simpler/worker.py | 77 ++++++++- .../platform/onboard/host/device_runner.cpp | 4 + src/common/hierarchical/orchestrator.cpp | 67 ++++++++ src/common/hierarchical/orchestrator.h | 9 ++ src/common/hierarchical/types.h | 1 + src/common/hierarchical/worker.cpp | 12 +- src/common/hierarchical/worker_manager.cpp | 96 +++++++++-- src/common/hierarchical/worker_manager.h | 31 +++- .../platform/onboard/host/c_api_shared.cpp | 9 ++ .../onboard/host/device_runner_base.cpp | 12 ++ .../onboard/host/device_runner_base.h | 7 + src/common/worker/chip_worker.cpp | 41 ++++- src/common/worker/chip_worker.h | 8 + src/common/worker/pto_runtime_c_api.h | 7 + .../ut/cpp/hierarchical/test_orchestrator.cpp | 66 ++++++++ tests/ut/cpp/hierarchical/test_scheduler.cpp | 152 ++++++++++++++++-- tests/ut/py/test_worker/test_host_worker.py | 37 +++++ 23 files changed, 655 insertions(+), 59 deletions(-) diff --git a/docs/orchestrator.md b/docs/orchestrator.md index aeb09b8d3d..f1457c7e3f 100644 --- a/docs/orchestrator.md +++ b/docs/orchestrator.md @@ -55,6 +55,8 @@ public: RunId begin_run(); void close_run_submission(RunId run_id); void fail_run_submission(RunId run_id, std::exception_ptr error); + void wait_run_accepted(RunId run_id); + bool run_accepted(RunId run_id) const; void wait_run(RunId run_id); bool wait_run_for(RunId run_id, double timeout_seconds); bool run_done(RunId run_id) const; @@ -96,9 +98,12 @@ worker.submit(orchestration, args, config).wait() ``` The current L2 backend is synchronous, so L2 `submit()` executes the existing -blocking path and returns an already-completed handle. L3 asynchronous return -does not yet imply overlapping runs: a later submit waits for the prior run's -fence and cleanup before building the next DAG. +blocking path and returns an already-completed handle. At L3 and above, graph +callbacks remain serialized, but a later submit waits only for every dispatch +in the prior run to cross its endpoint acceptance boundary. On A2A3 onboard, +that boundary is after both device kernels are enqueued and before stream +synchronization. Endpoints without an earlier signal fall back to completion. +Each run still owns its completion error, keepalives, and cleanup independently. Remote L3 submit adds two hidden pieces of metadata: final eligible worker-id sets and optional `RemoteTaskArgsSidecar` entries aligned by tensor index. diff --git a/docs/task-flow.md b/docs/task-flow.md index 9ed4e25d22..944f0f0028 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -348,6 +348,15 @@ the forked child decodes it. Remote NEXT_LEVEL dispatch through `RemoteL3Endpoint` serializes the same logical payload into a framed TASK request instead. +Every dispatched group member contributes one run-acceptance obligation. For +an A2A3 onboard chip endpoint, the child-side native runner writes +`TASK_ACCEPTED` after its AICore and AICPU kernels are both enqueued; the parent +observes it without releasing the mailbox. Other endpoint paths satisfy the +same obligation conservatively when their completion returns. Once submission +is closed and all obligations are satisfied, the next serialized orchestration +callback may build its DAG even though the prior run has not reached its +completion fence. + Local mailbox path: ```text diff --git a/docs/worker-manager.md b/docs/worker-manager.md index 8c544b58b7..8565f3c65e 100644 --- a/docs/worker-manager.md +++ b/docs/worker-manager.md @@ -46,7 +46,8 @@ public: void add_sub (void *mailbox); // Lifecycle - void start(Ring *ring, OnCompleteFn on_complete); // starts all WorkerThreads + void start(Ring *ring, OnCompleteFn on_complete, + OnAcceptFn on_accept); // starts all WorkerThreads void stop(); // Scheduler API @@ -100,8 +101,9 @@ struct WorkerDispatch { class WorkerThread { public: - void start(Ring *ring, WorkerManager *manager, + void start(Ring *ring, const std::function &on_complete, + const std::function &on_accept, std::unique_ptr endpoint); void stop(); void dispatch(WorkerDispatch d); // slot id + group sub-index @@ -118,13 +120,14 @@ private: std::condition_variable cv_; void loop(); - WorkerCompletion dispatch_process(WorkerDispatch d); + WorkerCompletion dispatch_process(WorkerDispatch d, + const std::function &on_accept); }; ``` The WorkerThread's `std::thread` pumps the internal queue and calls -`endpoint->run(...)` once per dispatch. `LocalMailboxEndpoint::run` drives the -shm handshake — one mailbox round trip per dispatch. The forked child loop +`endpoint->run_with_accept(...)` once per dispatch. `LocalMailboxEndpoint` +drives the shm handshake — one mailbox round trip per dispatch. The forked child loop that consumes the mailbox lives in Python (`_chip_process_loop` / `_sub_worker_loop` in `python/simpler/worker.py`); the parent does not fork children. @@ -152,7 +155,11 @@ and the child polls the mailbox for the lifetime of the worker. ### 3.1 Parent-side dispatch ```cpp -WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, WorkerDispatch d) { +// `run(ring, d)` forwards here with an empty hook; the acceptance-aware +// overload is what WorkerThread calls. +WorkerCompletion LocalMailboxEndpoint::run_with_accept( + Ring *ring, const WorkerDispatch &d, const std::function &on_accept +) { TaskSlotState &s = *ring->slot_state(d.task_slot); char *m = static_cast(mailbox_); @@ -175,7 +182,18 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, WorkerDispatch d) { // kChildLivenessPollPeriod steady-clock period rather than an iteration // count, which no longer maps to a bounded time at spin speed. auto next_liveness_check = steady_clock::now() + kChildLivenessPollPeriod; - while (read_state(mailbox_) != MailboxState::TASK_DONE) { + bool acceptance_observed = false; + while (true) { + MailboxState state = read_state(mailbox_); + // Launch acceptance is a sticky word, not a state: the child sets it + // before TASK_DONE and nothing clears it until the next TASK_READY, so + // a task that finishes between two polls cannot lose the ACK. Read + // after the state so a TASK_DONE observation implies it is visible. + if (!acceptance_observed && read_accepted(mailbox_)) { + acceptance_observed = true; + if (on_accept) on_accept(); + } + if (state == MailboxState::TASK_DONE) break; auto now = steady_clock::now(); if (now >= next_liveness_check) { next_liveness_check = now + kChildLivenessPollPeriod; @@ -206,6 +224,8 @@ Parent-side cost per dispatch: ([codestyle](../.claude/rules/codestyle.md) rule 5). Child liveness is sampled on a steady-clock period, so a dead child ends the wait with `ENDPOINT_FAILURE` instead of spinning the parent forever +- A sticky launch-acceptance word, observed before completion and cleared only + when the parent publishes the next `TASK_READY` - One explicit completion outcome: success, task failure, or endpoint failure Total ~nanoseconds overhead; the wait is dominated by actual kernel execution. @@ -216,7 +236,10 @@ The child loop lives in Python — see `_chip_process_loop` and `_sub_worker_loop` in `python/simpler/worker.py`. Each child polls `MAILBOX_OFF_STATE`, decodes the digest-prefixed args blob on `TASK_READY`, resolves the digest to its private local slot/callable, writes back any error, -and publishes `TASK_DONE`. +and publishes `TASK_DONE`. An A2A3 onboard chip child also exposes its mailbox +state word to the native runner, which publishes `TASK_ACCEPTED` after both +kernel launches succeed. SUB, remote, A5, and failure paths use WorkerThread's +completion-time acceptance fallback. The child inherits the parent's full address space at fork time, so: - ChipCallable objects (pre-fork allocated) are COW-visible at the same VA diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 16c109d92d..1aac74a091 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -1433,7 +1433,7 @@ NB_MODULE(_task_interface, m) { .def( "run_from_blob", [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity, - const CallConfig &config) { + const CallConfig &config, uint64_t accepted_state_addr, int32_t accepted_value) { // 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,9 +1441,12 @@ 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); + self.run( + callable_id, view, config, 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, "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 " diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 63b0f7e43f..f110c9a164 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -367,6 +367,12 @@ inline void bind_worker(nb::module_ &m) { "_wait_run", &Orchestrator::wait_run, nb::arg("run_id"), nb::call_guard(), "Block until one run is terminal and raise only that run's error." ) + .def( + "_wait_run_accepted", &Orchestrator::wait_run_accepted, nb::arg("run_id"), + nb::call_guard(), + "Block until every dispatch in one closed run has crossed its endpoint acceptance boundary." + ) + .def("_run_accepted", &Orchestrator::run_accepted, nb::arg("run_id")) .def( "_wait_run_for", &Orchestrator::wait_run_for, nb::arg("run_id"), nb::arg("timeout_seconds"), nb::call_guard(), "Wait up to timeout_seconds for one run to become terminal." diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index 789918ec26..d1ed08344f 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -566,6 +566,12 @@ def _fail_run_submission(self, run_id: int, message: str = "") -> None: def _wait_run(self, run_id: int) -> None: self._o._wait_run(run_id) + def _wait_run_accepted(self, run_id: int) -> None: + self._o._wait_run_accepted(run_id) + + def _run_accepted(self, run_id: int) -> bool: + return bool(self._o._run_accepted(run_id)) + def _wait_run_for(self, run_id: int, timeout_seconds: float) -> bool: return bool(self._o._wait_run_for(run_id, timeout_seconds)) diff --git a/python/simpler/worker.py b/python/simpler/worker.py index bcb70e9970..ba9a183c68 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -186,7 +186,13 @@ def my_l4_orch(orch, args, config): # MAILBOX_ARGS_CAPACITY mirrors the C++ constexpr in worker_manager.h so the # Python reader can bounds-check incoming args blobs. Source-of-truth for the # constants on the right is the nanobind binding (cannot drift). -_MAILBOX_ARGS_CAPACITY = MAILBOX_SIZE - _OFF_TASK_ARGS_BLOB - MAILBOX_ERROR_MSG_SIZE +# Mirrors MAILBOX_OFF_ACCEPTED / MAILBOX_TASK_ACCEPTED: launch acceptance is a +# sticky word rather than a MailboxState, because a state carrying it is lost +# whenever the child reaches TASK_DONE between two parent polls. The parent +# clears it when it publishes the next TASK_READY. +_OFF_ACCEPTED = MAILBOX_SIZE - MAILBOX_ERROR_MSG_SIZE - 8 +_TASK_ACCEPTED = 1 +_MAILBOX_ARGS_CAPACITY = MAILBOX_SIZE - _OFF_TASK_ARGS_BLOB - MAILBOX_ERROR_MSG_SIZE - 8 _OFF_CONTROL_CALLABLE_HASH = _OFF_ARGS + 32 # MAILBOX_OFF_ERROR_MSG / MAILBOX_ERROR_MSG_SIZE come from the C++ # nanobind module so the two sides cannot drift. @@ -1534,7 +1540,14 @@ def handle_task() -> tuple[int, str]: # it in Python is N×40B of avoidable work and a permanent # opportunity to drop a field. C++ reinterpret_cast # is the source of truth. - cw._impl.run_from_blob(cid, mailbox_addr + _OFF_TASK_ARGS_BLOB, _MAILBOX_ARGS_CAPACITY, cfg) + cw._impl.run_from_blob( + cid, + mailbox_addr + _OFF_TASK_ARGS_BLOB, + _MAILBOX_ARGS_CAPACITY, + cfg, + mailbox_addr + _OFF_ACCEPTED, + _TASK_ACCEPTED, + ) except Exception as e: # noqa: BLE001 code = 1 msg = _format_exc(f"chip_process dev={device_id}", e) @@ -1967,6 +1980,8 @@ def __init__( self._resources = resources if resources is not None else _RunResources() self._cv = threading.Condition() self._wait_in_progress = False + self._accept_wait_in_progress = False + self._launch_accepted = False self._terminal = False self._error: BaseException | None = None @@ -1979,6 +1994,8 @@ def _completed(cls, worker: Worker) -> RunHandle: handle._resources = _RunResources() handle._cv = threading.Condition() handle._wait_in_progress = False + handle._accept_wait_in_progress = False + handle._launch_accepted = True handle._terminal = True handle._error = None return handle @@ -2064,6 +2081,20 @@ def wait(self, timeout: float | None = None) -> None: self._cv.notify_all() raise TimeoutError("RunHandle.wait() timed out") + # An acceptance waiter that already captured this run id must leave the + # native wait before finalize releases that id. It is terminal now, so + # this is only a short ownership hand-off and does not serialize device + # execution with acceptance waiting. + try: + with self._cv: + while self._accept_wait_in_progress: + self._cv.wait(timeout=_RUN_HANDLE_WAIT_RECHECK_S) + except BaseException: + self._wait_in_progress = False + with self._cv: + self._cv.notify_all() + raise + # Cleanup runs exactly once, on this waiter, and its outcome IS the # handle's result: an interruption mid-finalize is cached like any other # error rather than lost, so the publication below is unconditional. @@ -2074,8 +2105,10 @@ def wait(self, timeout: float | None = None) -> None: self._error = error self._run_id = None self._keepalive = None + self._launch_accepted = True self._terminal = True self._wait_in_progress = False + self._accept_wait_in_progress = False with self._cv: self._cv.notify_all() if error is not None: @@ -2093,6 +2126,30 @@ def _wait_for_serialization(self) -> None: # The result remains cached for this handle's public wait/result. pass + def _wait_for_acceptance(self) -> None: + """Wait until this run's dispatches cross their acceptance boundary.""" + with self._cv: + while not self._terminal and self._accept_wait_in_progress: + self._cv.wait(timeout=_RUN_HANDLE_WAIT_RECHECK_S) + if self._terminal or self._launch_accepted: + return + self._accept_wait_in_progress = True + run_id = self._run_id + + assert run_id is not None + try: + self._worker._wait_run_handle_accepted(run_id) + except BaseException: + self._accept_wait_in_progress = False + with self._cv: + self._cv.notify_all() + raise + + self._launch_accepted = True + self._accept_wait_in_progress = False + with self._cv: + self._cv.notify_all() + def _forked_child_main(buf: memoryview, label: str, setup, serve, make_group_leader: bool = False) -> None: """Run a forked child to completion, always terminating via ``os._exit``. @@ -5989,8 +6046,9 @@ def submit(self, callable, args=None, config=None) -> RunHandle: ``args`` : TaskArgs (optional) ``config``: CallConfig (optional, default-constructed if None) - Only one live device run is admitted: a later submission waits for the - previous handle's fence and cleanup before building its DAG. + Graph construction remains serialized. A later submission waits until + prior dispatches are accepted; completion and cleanup stay attached to + each returned handle. """ with self._operation_lease("submit"): return self._submit_locked(callable, args, config) @@ -6017,13 +6075,12 @@ def _submit_locked(self, callable, args, config) -> RunHandle: return RunHandle._completed(self) with self._submit_mu: - # Cleanup is Worker-global while only one live device run is - # admitted. Drain prior handles before a new callback can mutate - # those resources; errors remain attached only to their origin. + # Graph callbacks are serialized, but accepted runs may remain live: + # their Python resources are isolated in each RunHandle. with self._hierarchical_start_cv: prior_handles = tuple(self._accepted_run_handles) for handle in prior_handles: - handle._wait_for_serialization() + handle._wait_for_acceptance() return self._submit_l3_locked(callable, args, cfg) def _submit_l3_locked(self, callable, args, cfg: CallConfig) -> RunHandle: @@ -6080,6 +6137,10 @@ def _wait_run_handle(self, run_id: int, timeout: float | None) -> bool: return True return self._orch._wait_run_for(run_id, timeout) + def _wait_run_handle_accepted(self, run_id: int) -> None: + assert self._orch is not None + self._orch._wait_run_accepted(run_id) + def _finalize_run_handle( self, handle: RunHandle, run_id: int, native_error: BaseException | None ) -> BaseException | None: diff --git a/src/a2a3/platform/onboard/host/device_runner.cpp b/src/a2a3/platform/onboard/host/device_runner.cpp index 5cd2536015..ae35b65156 100644 --- a/src/a2a3/platform/onboard/host/device_runner.cpp +++ b/src/a2a3/platform/onboard/host/device_runner.cpp @@ -618,6 +618,10 @@ int DeviceRunner::launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_ return rc; } + // Both kernels are enqueued. Publish before reap_run synchronizes either + // stream; the parent retains mailbox ownership until TASK_DONE. + publish_task_accepted(); + return 0; } diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index 930bac6e19..649a25bc7b 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -34,6 +34,11 @@ void Orchestrator::init( bool Orchestrator::is_terminal(RunPhase phase) { return phase == RunPhase::COMPLETED || phase == RunPhase::FAILED; } +bool Orchestrator::acceptance_ready(const std::shared_ptr &run) { + return run->submission_closed && (run->pending_accepts.load(std::memory_order_acquire) == 0 || + is_terminal(run->phase.load(std::memory_order_acquire))); +} + std::shared_ptr Orchestrator::find_run(RunId run_id) const { std::lock_guard lk(runs_mu_); auto it = runs_.find(run_id); @@ -103,6 +108,7 @@ void Orchestrator::close_run_submission(RunId run_id) { run->phase.store(RunPhase::EXECUTING, std::memory_order_release); } } + run->completion_cv.notify_all(); finish_run_if_ready(run); } @@ -124,9 +130,27 @@ void Orchestrator::fail_run_submission(RunId run_id, std::exception_ptr error) { run->phase.store(RunPhase::EXECUTING, std::memory_order_release); } } + run->completion_cv.notify_all(); finish_run_if_ready(run); } +void Orchestrator::wait_run_accepted(RunId run_id) { + // A released run is past acceptance by definition, so a waiter that raced + // the run's completion returns instead of throwing. + auto run = find_run(run_id); + if (run == nullptr) return; + std::unique_lock lk(run->completion_mu); + run->completion_cv.wait(lk, [&run] { + return acceptance_ready(run); + }); +} + +bool Orchestrator::run_accepted(RunId run_id) const { + auto run = get_run(run_id); + std::lock_guard lk(run->completion_mu); + return acceptance_ready(run); +} + void Orchestrator::wait_run(RunId run_id) { auto run = get_run(run_id); std::exception_ptr error; @@ -213,6 +237,39 @@ void Orchestrator::decrement_run_tasks(RunId run_id) { if (remaining <= 0) finish_run_if_ready(run); } +void Orchestrator::increment_run_accepts(RunId run_id, int32_t count) { + if (count <= 0) throw std::logic_error("Orchestrator: run acceptance count must be positive"); + get_run(run_id)->pending_accepts.fetch_add(count, std::memory_order_relaxed); +} + +void Orchestrator::decrement_run_accepts(RunId run_id) { + auto run = find_run(run_id); + if (run == nullptr) return; + bool notify = false; + bool underflow = false; + { + // acceptance_ready() is evaluated under this mutex, so the decrement + // that can satisfy it happens under the same one — otherwise it lands + // between the waiter's predicate check and its block, and the notify + // below is lost. + std::lock_guard lk(run->completion_mu); + int32_t remaining = run->pending_accepts.fetch_sub(1, std::memory_order_acq_rel) - 1; + if (remaining < 0) { + run->pending_accepts.store(0, std::memory_order_release); + underflow = true; + remaining = 0; + } + notify = remaining == 0; + } + // record_run_error takes completion_mu, so it runs outside the block above. + if (underflow) { + record_run_error( + run, std::make_exception_ptr(std::logic_error("Orchestrator: run acceptance count underflow")) + ); + } + if (notify) run->completion_cv.notify_all(); +} + void Orchestrator::record_run_error(const std::shared_ptr &run, std::exception_ptr error) { if (run == nullptr || !error) return; std::lock_guard lk(run->completion_mu); @@ -229,6 +286,15 @@ void Orchestrator::report_task_error(TaskSlot slot, const std::string &message) record_run_error(task.run_id, std::make_exception_ptr(std::runtime_error(message))); } +void Orchestrator::mark_task_accepted(TaskSlot slot) { + // Reached from a WorkerThread, where an escaping exception is fatal to the + // process, so an unknown slot degrades to "no fence advance" instead of + // throwing; the dispatch's completion still reports the failure. + TaskSlotState *task = allocator_->slot_state(slot); + if (task == nullptr) return; + decrement_run_accepts(task->run_id); +} + uint64_t Orchestrator::malloc(int worker_id, size_t size) { auto *wt = manager_->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); if (!wt) throw std::runtime_error("Orchestrator::malloc: invalid worker_id"); @@ -481,6 +547,7 @@ SubmitResult Orchestrator::submit_impl( if (scope_ref > 0) scope_->register_task(slot); increment_run_tasks(run->id); + increment_run_accepts(run->id, s.group_size()); if (poisoned_by_failed_producer) { if (poison_message.empty()) poison_message = "producer task failed"; diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index 5273def13c..d131bc0dd8 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -120,6 +120,8 @@ class Orchestrator { RunId begin_run(); void close_run_submission(RunId run_id); void fail_run_submission(RunId run_id, std::exception_ptr error = nullptr); + void wait_run_accepted(RunId run_id); + bool run_accepted(RunId run_id) const; void wait_run(RunId run_id); bool wait_run_for(RunId run_id, double timeout_seconds); bool run_done(RunId run_id) const; @@ -148,6 +150,10 @@ class Orchestrator { // Attach a scheduler/endpoint failure to the task's originating run. void report_task_error(TaskSlot slot, const std::string &message); + // Called once for each dispatched group member after its endpoint has + // accepted the launch, or conservatively at endpoint completion. + void mark_task_accepted(TaskSlot slot); + // Called by Scheduler (via Worker) when a task becomes CONSUMED: // erases TensorMap entries, releases the allocator slot (and implicitly // the slot's heap slab via last_alive). @@ -186,11 +192,14 @@ class Orchestrator { std::shared_ptr current_building_run() const; static void finish_run_if_ready(const std::shared_ptr &run); static bool is_terminal(RunPhase phase); + static bool acceptance_ready(const std::shared_ptr &run); // Callers hold runs_mu_. bool quiescent_locked() const; void compact_if_quiescent(); void increment_run_tasks(RunId run_id); void decrement_run_tasks(RunId run_id); + void increment_run_accepts(RunId run_id, int32_t count); + void decrement_run_accepts(RunId run_id); static void record_run_error(const std::shared_ptr &run, std::exception_ptr error); void record_run_error(RunId run_id, std::exception_ptr error); diff --git a/src/common/hierarchical/types.h b/src/common/hierarchical/types.h index 5819f88ae3..cdce0f1200 100644 --- a/src/common/hierarchical/types.h +++ b/src/common/hierarchical/types.h @@ -141,6 +141,7 @@ struct RunState { RunId id{INVALID_RUN_ID}; std::atomic phase{RunPhase::BUILDING}; std::atomic active_tasks{0}; + std::atomic pending_accepts{0}; mutable std::mutex completion_mu; std::condition_variable completion_cv; std::exception_ptr first_error; diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index 457276583f..90efac32a3 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -98,9 +98,15 @@ void Worker::init() { // Start WorkerManager first — creates WorkerThreads. // The on_complete callback routes through the Scheduler's worker_done(). - manager_.start(&allocator_, [this](WorkerCompletion completion) { - scheduler_.worker_done(std::move(completion)); - }); + manager_.start( + &allocator_, + [this](WorkerCompletion completion) { + scheduler_.worker_done(std::move(completion)); + }, + [this](WorkerDispatch dispatch) { + orchestrator_.mark_task_accepted(dispatch.task_slot); + } + ); ready_next_level_queues_.reset(manager_.next_level_worker_ids()); orchestrator_.init( &tensormap_, &allocator_, &scope_, &ready_sub_queue_, &ready_next_level_queues_, &manager_, [this] { diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 979c42aa41..ad2466985b 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -139,6 +139,13 @@ void WorkerEndpoint::control_l3_l2_region_release(uint64_t) { throw_unsupported_control("control_l3_l2_region_release"); } +WorkerCompletion +WorkerEndpoint::run_with_accept(Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept) { + WorkerCompletion completion = run(ring, dispatch); + if (on_accept) on_accept(); + return completion; +} + // ============================================================================= // LocalMailboxEndpoint — mailbox helpers // ============================================================================= @@ -205,6 +212,19 @@ void LocalMailboxEndpoint::write_mailbox_state(MailboxState s) { #endif } +bool LocalMailboxEndpoint::read_task_accepted() const { + const int32_t *ptr = reinterpret_cast(mbox() + MAILBOX_OFF_ACCEPTED); + int32_t v = 0; + __atomic_load(ptr, &v, __ATOMIC_ACQUIRE); + return v == MAILBOX_TASK_ACCEPTED; +} + +void LocalMailboxEndpoint::clear_task_accepted() { + int32_t *ptr = reinterpret_cast(mbox() + MAILBOX_OFF_ACCEPTED); + int32_t v = 0; + __atomic_store(ptr, &v, __ATOMIC_RELEASE); +} + void LocalMailboxEndpoint::shutdown_child() { write_mailbox_state(MailboxState::SHUTDOWN); } // ============================================================================= @@ -212,11 +232,13 @@ void LocalMailboxEndpoint::shutdown_child() { write_mailbox_state(MailboxState:: // ============================================================================= void WorkerThread::start( - Ring *ring, const std::function &on_complete, std::unique_ptr endpoint + Ring *ring, const std::function &on_complete, + const std::function &on_accept, std::unique_ptr endpoint ) { if (!endpoint) throw std::invalid_argument("WorkerThread::start: null endpoint"); ring_ = ring; on_complete_ = on_complete; + on_accept_ = on_accept; endpoint_ = std::move(endpoint); shutdown_ = false; idle_.store(true, std::memory_order_relaxed); @@ -268,8 +290,14 @@ void WorkerThread::loop() { } WorkerCompletion completion; + bool accepted = false; + auto accept_once = [&]() { + if (accepted) return; + accepted = true; + if (on_accept_) on_accept_(d); + }; try { - completion = dispatch_process(d); + completion = dispatch_process(d, accept_once); } catch (const std::exception &e) { completion.task_slot = d.task_slot; completion.group_index = d.group_index; @@ -281,18 +309,42 @@ void WorkerThread::loop() { completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; completion.error_message = "WorkerThread endpoint failed with unknown exception"; } + // The fence must advance even when the dispatch failed — a waiter on + // wait_run_accepted would otherwise block forever — and it must not + // throw out of this thread, where an escaping exception terminates the + // process. mark_task_accepted is non-throwing by construction; this + // keeps that a local property rather than a remote assumption. + try { + accept_once(); + } catch (const std::exception &e) { + if (completion.outcome == EndpointOutcome::SUCCESS) { + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = std::string("WorkerThread accept fence failed: ") + e.what(); + } + } catch (...) { + if (completion.outcome == EndpointOutcome::SUCCESS) { + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "WorkerThread accept fence failed with unknown exception"; + } + } idle_.store(true, std::memory_order_release); on_complete_(std::move(completion)); } } -WorkerCompletion WorkerThread::dispatch_process(WorkerDispatch d) { +WorkerCompletion WorkerThread::dispatch_process(WorkerDispatch d, const std::function &on_accept) { if (!endpoint_) throw std::runtime_error("WorkerThread::dispatch_process: null endpoint"); - return endpoint_->run(ring_, d); + return endpoint_->run_with_accept(ring_, d, on_accept); } WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dispatch) { + return run_with_accept(ring, dispatch, {}); +} + +WorkerCompletion LocalMailboxEndpoint::run_with_accept( + Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept +) { if (ring == nullptr) throw std::invalid_argument("LocalMailboxEndpoint::run: null ring"); TaskSlotState &s = *ring->slot_state(dispatch.task_slot); int32_t group_index = dispatch.group_index; @@ -361,17 +413,31 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis ); } + // Clear the sticky ACK before the child can set it, so the previous task's + // acceptance cannot satisfy this one's fence. + clear_task_accepted(); + // Signal child process. write_mailbox_state(MailboxState::TASK_READY); - // Spin-poll until child signals TASK_DONE. The task's latency runs through - // this wait, so it never sleeps (codestyle rule 5). A child that dies - // without publishing TASK_DONE would otherwise leave the loop spinning - // forever, so its liveness is sampled on a kChildLivenessPollPeriod wall - // clock rather than an iteration count, which no longer maps to a wall - // time once the loop runs at spin speed. + // Spin-poll until child signals TASK_DONE, observing the sticky launch ACK + // on the way. The task's latency runs through this wait, so it never sleeps + // (codestyle rule 5). A child that dies without publishing TASK_DONE would + // otherwise leave the loop spinning forever, so its liveness is sampled on a + // kChildLivenessPollPeriod wall clock rather than an iteration count, which + // no longer maps to a wall time once the loop runs at spin speed. auto next_liveness_check = std::chrono::steady_clock::now() + kChildLivenessPollPeriod; - while (read_mailbox_state() != MailboxState::TASK_DONE) { + bool acceptance_observed = false; + while (true) { + MailboxState state = read_mailbox_state(); + // Read the ACK after the state. The child sets the sticky word before + // TASK_DONE, so observing TASK_DONE here implies this load sees it — + // a task that finishes between two polls cannot lose its acceptance. + if (!acceptance_observed && read_task_accepted()) { + acceptance_observed = true; + if (on_accept) on_accept(); + } + if (state == MailboxState::TASK_DONE) break; auto now = std::chrono::steady_clock::now(); if (now >= next_liveness_check) { next_liveness_check = now + kChildLivenessPollPeriod; @@ -425,7 +491,7 @@ void WorkerManager::add_next_level_endpoint(std::unique_ptr endp void WorkerManager::add_sub(void *mailbox, int child_pid) { sub_entries_.push_back(LocalSubEntry{mailbox, child_pid}); } -void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete) { +void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete, const OnAcceptFn &on_accept) { if (ring == nullptr) throw std::invalid_argument("WorkerManager::start: null ring"); std::vector next_level_worker_ids; @@ -453,7 +519,7 @@ void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete) { for (const auto &entry : next_level_entries_) { auto wt = std::make_unique(); auto endpoint = std::make_unique(entry.worker_id, entry.mailbox, entry.child_pid); - wt->start(ring, on_complete, std::move(endpoint)); + wt->start(ring, on_complete, on_accept, std::move(endpoint)); next_level_threads_.push_back(std::move(wt)); } }; @@ -464,14 +530,14 @@ void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete) { auto endpoint = std::make_unique( static_cast(i), entries[i].mailbox, entries[i].child_pid ); - wt->start(ring, on_complete, std::move(endpoint)); + wt->start(ring, on_complete, on_accept, std::move(endpoint)); threads.push_back(std::move(wt)); } }; make_next_level_threads(); for (auto &endpoint : next_level_endpoint_entries_) { auto wt = std::make_unique(); - wt->start(ring, on_complete, std::move(endpoint)); + wt->start(ring, on_complete, on_accept, std::move(endpoint)); next_level_threads_.push_back(std::move(wt)); } next_level_endpoint_entries_.clear(); diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index c8a73f6393..6d00a95519 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -111,6 +111,13 @@ static_assert( ); static constexpr ptrdiff_t MAILBOX_OFF_ERROR_MSG = static_cast(MAILBOX_SIZE) - static_cast(MAILBOX_ERROR_MSG_SIZE); +// Launch acceptance is sticky, not a MailboxState: a state word carrying it +// loses the ACK whenever the child reaches TASK_DONE between two parent polls, +// which is exactly the short-task case the fence exists to pipeline. The child +// sets it after both launches; the parent clears it when it publishes the next +// TASK_READY, so nothing else can overwrite it in between. +static constexpr int32_t MAILBOX_TASK_ACCEPTED = 1; +static constexpr ptrdiff_t MAILBOX_OFF_ACCEPTED = MAILBOX_OFF_ERROR_MSG - 8; static constexpr ptrdiff_t MAILBOX_OFF_TASK_CALLABLE_HASH = MAILBOX_OFF_ARGS; static constexpr ptrdiff_t MAILBOX_OFF_TASK_ARGS_BLOB = MAILBOX_OFF_TASK_CALLABLE_HASH + static_cast(CALLABLE_HASH_DIGEST_SIZE); @@ -118,7 +125,7 @@ static constexpr size_t CTRL_SHM_NAME_BYTES = 32; static constexpr ptrdiff_t MAILBOX_OFF_CONTROL_CALLABLE_HASH = MAILBOX_OFF_ARGS + static_cast(CTRL_SHM_NAME_BYTES); static constexpr size_t MAILBOX_ARGS_CAPACITY = - MAILBOX_SIZE - static_cast(MAILBOX_OFF_TASK_ARGS_BLOB) - MAILBOX_ERROR_MSG_SIZE; + MAILBOX_SIZE - static_cast(MAILBOX_OFF_TASK_ARGS_BLOB) - MAILBOX_ERROR_MSG_SIZE - 8; static_assert( MAILBOX_ARGS_CAPACITY >= TASK_ARGS_BLOB_HEADER_SIZE + static_cast(CHIP_MAX_TENSOR_ARGS) * sizeof(Tensor) + static_cast(CHIP_MAX_SCALAR_ARGS) * sizeof(uint64_t), @@ -202,6 +209,10 @@ class WorkerEndpoint { virtual const WorkerEndpointCaps &caps() const = 0; virtual WorkerCompletion run(Ring *ring, const WorkerDispatch &dispatch) = 0; + // Endpoints with an earlier launch boundary override this. The default is a + // conservative completion-time acceptance for remote, SUB and A5 paths. + virtual WorkerCompletion + run_with_accept(Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept); virtual void shutdown_child() {} virtual uint64_t control_malloc(size_t size); @@ -257,6 +268,8 @@ class LocalMailboxEndpoint : public WorkerEndpoint { const WorkerEndpointCaps &caps() const override { return caps_; } WorkerCompletion run(Ring *ring, const WorkerDispatch &dispatch) override; + WorkerCompletion + run_with_accept(Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept) override; void shutdown_child() override; uint64_t control_malloc(size_t size) override; @@ -316,6 +329,10 @@ class LocalMailboxEndpoint : public WorkerEndpoint { char *mbox() const { return static_cast(mailbox_); } MailboxState read_mailbox_state() const; void write_mailbox_state(MailboxState s); + // Sticky launch acceptance, cleared only when this endpoint publishes the + // next TASK_READY. See MAILBOX_OFF_ACCEPTED. + bool read_task_accepted() const; + void clear_task_accepted(); void run_control_command(const char *op_name, double timeout_s = -1.0); // Returns a description of the child's death, or an empty string while it @@ -361,7 +378,8 @@ class WorkerThread { // on_complete(completion) is called (in the WorkerThread) after each // endpoint run(). void start( - Ring *ring, const std::function &on_complete, std::unique_ptr endpoint + Ring *ring, const std::function &on_complete, + const std::function &on_accept, std::unique_ptr endpoint ); // Enqueue a dispatch for the worker. Non-blocking. @@ -451,6 +469,7 @@ class WorkerThread { Ring *ring_{nullptr}; std::unique_ptr endpoint_; std::function on_complete_; + std::function on_accept_; std::thread thread_; std::queue queue_; @@ -460,7 +479,7 @@ class WorkerThread { std::atomic idle_{true}; void loop(); - WorkerCompletion dispatch_process(WorkerDispatch d); + WorkerCompletion dispatch_process(WorkerDispatch d, const std::function &on_accept); }; // ============================================================================= @@ -470,6 +489,7 @@ class WorkerThread { class WorkerManager { public: using OnCompleteFn = std::function; + using OnAcceptFn = std::function; // Register a worker. `mailbox` is a MAILBOX_SIZE-byte MAP_SHARED // region; the real worker (a `ChipWorker` for NEXT_LEVEL, a Python @@ -481,7 +501,10 @@ class WorkerManager { void add_next_level_endpoint(std::unique_ptr endpoint); void add_sub(void *mailbox, int child_pid = -1); - void start(Ring *ring, const OnCompleteFn &on_complete); + // `on_accept` advances the run's launch fence; pass an empty function only + // when nothing waits on that fence, since an omitted callback leaves + // pending_accepts non-zero forever. No default: the choice is the caller's. + void start(Ring *ring, const OnCompleteFn &on_complete, const OnAcceptFn &on_accept); void stop(); WorkerThread *get_worker_by_id(WorkerType type, int32_t worker_id) const; diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index 03e3f0934a..101118d6e8 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -650,6 +650,15 @@ int simpler_run( } } +int set_task_accepted_state_ctx(DeviceContextHandle ctx, volatile int32_t *state, int32_t accepted_value) { + if (ctx == NULL) return -1; + try { + return static_cast(ctx)->set_task_accepted_state(state, accepted_value); + } catch (...) { + return -1; + } +} + int 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 30c2bcd240..e1ef5115c9 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -1430,3 +1430,15 @@ void DeviceRunnerBase::teardown_shared_collectors_after_run() { scope_stats_collector_.write_jsonl(output_prefix_); } } + +int DeviceRunnerBase::set_task_accepted_state(volatile int32_t *state, int32_t accepted_value) { + task_accepted_state_ = state; + task_accepted_value_ = accepted_value; + return 0; +} + +void DeviceRunnerBase::publish_task_accepted() const { + if (task_accepted_state_ != nullptr) { + __atomic_store_n(task_accepted_state_, task_accepted_value_, __ATOMIC_RELEASE); + } +} diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index 4953ee8d84..96b9c87ab9 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -87,6 +87,9 @@ class DeviceRunnerBase { DeviceRunnerBase(DeviceRunnerBase &&) = delete; DeviceRunnerBase &operator=(DeviceRunnerBase &&) = delete; + /** Bind this runner's launch-acceptance publication target. */ + int set_task_accepted_state(volatile int32_t *state, int32_t accepted_value); + /** Allocate / free / copy on the per-Worker `MemoryAllocator` + CANN runtime. */ void *allocate_tensor(std::size_t bytes); void free_tensor(void *dev_ptr); @@ -571,6 +574,8 @@ class DeviceRunnerBase { const std::string &output_prefix() const { return output_prefix_; } protected: + /** Publish launch acceptance for the currently bound target, if any. */ + void publish_task_accepted() const; // Ctor is protected: this class is for inheritance only — direct // instantiation (`new DeviceRunnerBase()`) is a compile error. The // public virtual dtor above lets the shared c_api delete through a @@ -840,6 +845,8 @@ class DeviceRunnerBase { // Same re-register semantics as `aicpu_dlopen_total_`, but for hbg // variants. size_t host_dlopen_total_{0}; + volatile int32_t *task_accepted_state_{nullptr}; + int32_t task_accepted_value_{0}; // ---- State shared by both a2a3 and a5 --------------------------------- // diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index b71e96ecb1..68e48e2ce6 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -117,6 +117,8 @@ 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"); + 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"); unregister_callable_fn_ = load_symbol(handle, "simpler_unregister_callable"); get_aicpu_dlopen_count_fn_ = load_symbol(handle, "get_aicpu_dlopen_count"); @@ -216,6 +218,7 @@ void ChipWorker::init( simpler_init_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + set_task_accepted_state_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; @@ -256,6 +259,7 @@ void ChipWorker::init( simpler_init_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + set_task_accepted_state_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; @@ -326,6 +330,7 @@ void ChipWorker::finalize() { get_runtime_size_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + set_task_accepted_state_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; @@ -365,21 +370,53 @@ void ChipWorker::register_callable(int32_t callable_id, const void *callable) { } void ChipWorker::run(int32_t callable_id, TaskArgsView args, const CallConfig &config) { + run(callable_id, args, config, nullptr, 0); +} + +void ChipWorker::run( + int32_t callable_id, TaskArgsView args, const CallConfig &config, volatile int32_t *accepted_state, + int32_t accepted_value +) { ChipStorageTaskArgs chip_storage = view_to_chip_storage(args); - run(callable_id, &chip_storage, config); + run(callable_id, &chip_storage, config, accepted_state, accepted_value); } void ChipWorker::run(int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config) { + run(callable_id, args, config, nullptr, 0); +} + +void ChipWorker::run( + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, 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(); + if (accepted_state != nullptr && set_task_accepted_state_fn_ != nullptr) { + int bind_rc = set_task_accepted_state_fn_(device_ctx_, accepted_state, accepted_value); + if (bind_rc != 0) { + throw std::runtime_error("set_task_accepted_state_ctx failed with code " + std::to_string(bind_rc)); + } + } + auto clear_accepted_state = [&]() { + if (accepted_state != nullptr && set_task_accepted_state_fn_ != nullptr) { + (void)set_task_accepted_state_fn_(device_ctx_, nullptr, 0); + } + }; // Per-stage timing is emitted by the platform as `[STRACE]` log markers, not // returned (see chip_worker.h::run). CallConfig is threaded through to the C // ABI as a single pointer rather than unpacked into per-field args. - int rc = run_fn_(device_ctx_, rt, callable_id, args, &config); + int rc = -1; + try { + rc = run_fn_(device_ctx_, rt, callable_id, args, &config); + } catch (...) { + clear_accepted_state(); + throw; + } + clear_accepted_state(); if (rc != 0) { throw std::runtime_error("run failed with code " + std::to_string(rc)); } diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 7ead68ff33..7e0dbb84fc 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -69,10 +69,16 @@ class ChipWorker { // platform as `[STRACE]` log markers — see src/common/log/.../strace.h — not // returned, so the L3 dispatcher and L2 child are observed uniformly. void run(int32_t callable_id, TaskArgsView args, const CallConfig &config); + void + run(int32_t callable_id, TaskArgsView args, const CallConfig &config, volatile int32_t *accepted_state, + int32_t accepted_value); // Same launch, but the caller already holds the runtime.so-ABI POD — // skip the view→storage memcpy and hand the pointer straight to the C ABI. // Used by the ChipStorageTaskArgs path in the nanobind binding. void run(int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config); + void + run(int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, + volatile int32_t *accepted_state, int32_t accepted_value); // Per-callable_id preparation. Requires init() first and a callable_id // in [0, MAX_REGISTERED_CALLABLE_IDS) (cap 64). @@ -163,6 +169,7 @@ 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 GetPipelineContractFn = const PipelineContract *(*)(); using SimplerUnregisterCallableFn = int (*)(void *, int32_t); using GetAicpuDlopenCountFn = size_t (*)(void *); @@ -210,6 +217,7 @@ class ChipWorker { SimplerInitFn simpler_init_fn_ = nullptr; SimplerRegisterCallableFn register_callable_fn_ = nullptr; SimplerRunFn run_fn_ = nullptr; + SetTaskAcceptedStateFn set_task_accepted_state_fn_ = nullptr; SimplerUnregisterCallableFn unregister_callable_fn_ = nullptr; GetAicpuDlopenCountFn get_aicpu_dlopen_count_fn_ = nullptr; GetAicpuDlopenCountFn get_host_dlopen_count_fn_ = nullptr; diff --git a/src/common/worker/pto_runtime_c_api.h b/src/common/worker/pto_runtime_c_api.h index 6d20c3b878..0c77507fbe 100644 --- a/src/common/worker/pto_runtime_c_api.h +++ b/src/common/worker/pto_runtime_c_api.h @@ -271,6 +271,13 @@ int simpler_run( DeviceContextHandle ctx, RuntimeHandle runtime, int32_t callable_id, const void *args, const CallConfig *config ); +/** + * Bind an optional host state word that the runner publishes after both device + * kernels have been enqueued. Onboard runtimes may export this symbol; callers + * must fall back to completion when it is absent. + */ +int set_task_accepted_state_ctx(DeviceContextHandle ctx, volatile int32_t *state, int32_t accepted_value); + /** * Drop the prepared state for `callable_id` and release the per-id share of * the device orch SO buffer. The buffer itself is freed only when its diff --git a/tests/ut/cpp/hierarchical/test_orchestrator.cpp b/tests/ut/cpp/hierarchical/test_orchestrator.cpp index 519e170dbd..8dd4791c13 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -473,6 +473,36 @@ TEST_F(OrchestratorFixture, EmptyRunCompletesWhenSubmissionCloses) { orch.release_run(run_id); } +TEST_F(OrchestratorFixture, RunAcceptanceWaitsForEveryDispatchedGroupMember) { + std::vector args{ + single_tensor_args(0x8010, TensorArgType::OUTPUT), + single_tensor_args(0x8020, TensorArgType::OUTPUT), + }; + auto result = orch.submit_next_level_group(C(80), args, cfg, {0, 1}); + + orch.close_run_submission(run_id); + EXPECT_FALSE(orch.run_accepted(run_id)); + + orch.mark_task_accepted(result.task_slot); + EXPECT_FALSE(orch.run_accepted(run_id)); + + orch.mark_task_accepted(result.task_slot); + EXPECT_TRUE(orch.run_accepted(run_id)); + EXPECT_NO_THROW(orch.wait_run_accepted(run_id)); +} + +TEST_F(OrchestratorFixture, TerminalFailureUnblocksRunAcceptance) { + auto result = orch.submit_next_level(C(80), single_tensor_args(0x8030, TensorArgType::OUTPUT), cfg, 0); + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("graph build failed"))); + EXPECT_FALSE(orch.run_accepted(run_id)); + + S(result.task_slot).state.store(TaskState::FAILED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(result.task_slot)); + + EXPECT_TRUE(orch.run_accepted(run_id)); + EXPECT_NO_THROW(orch.wait_run_accepted(run_id)); +} + TEST_F(OrchestratorFixture, TimedWaitCanRetryAfterTimeout) { EXPECT_FALSE(orch.wait_run_for(run_id, 0.0)); EXPECT_FALSE(orch.run_done(run_id)); @@ -496,6 +526,42 @@ TEST_F(OrchestratorFixture, OneTaskRunCompletesAfterConsumption) { orch.release_run(run_id); } +// mark_task_accepted runs on a WorkerThread, where an escaping exception ends +// the process. Neither an unknown slot nor an over-count may throw out of it; a +// count mismatch becomes the run's error, the way decrement_run_tasks does it. +TEST_F(OrchestratorFixture, AcceptanceFenceNeverThrowsOutOfAWorkerThread) { + auto result = orch.submit_next_level(C(90), single_tensor_args(0x9000, TensorArgType::OUTPUT), cfg, 0); + orch.close_run_submission(run_id); + + EXPECT_NO_THROW(orch.mark_task_accepted(TaskSlot{-1})); + EXPECT_NO_THROW(orch.mark_task_accepted(result.task_slot)); + EXPECT_TRUE(orch.run_accepted(run_id)); + EXPECT_FALSE(orch.run_failed(run_id)); + + // One accept past the submitted count: reported, not thrown. + EXPECT_NO_THROW(orch.mark_task_accepted(result.task_slot)); + EXPECT_TRUE(orch.run_failed(run_id)); + + S(result.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(result.task_slot)); + orch.release_run(run_id); +} + +// A waiter that races the run's release must return, not throw: a released run +// is past acceptance by definition. +TEST_F(OrchestratorFixture, AcceptanceWaitOnAReleasedRunReturns) { + auto result = orch.submit_next_level(C(92), single_tensor_args(0x9200, TensorArgType::OUTPUT), cfg, 0); + orch.close_run_submission(run_id); + orch.mark_task_accepted(result.task_slot); + S(result.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(result.task_slot)); + RunId released = run_id; + orch.release_run(released); + + EXPECT_NO_THROW(orch.wait_run_accepted(released)); + run_id = orch.begin_run(); +} + TEST_F(OrchestratorFixture, SequentialRunsHaveDistinctIdsAndErrorsDoNotLeak) { auto failed_task = orch.submit_next_level(C(81), single_tensor_args(0x8100, TensorArgType::OUTPUT), cfg, 0); orch.report_task_error(failed_task.task_slot, "run one failed"); diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 395d0137cc..31ca70b436 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -116,6 +116,13 @@ struct MockMailboxWorker { run_cv.notify_one(); } + // The child publishes acceptance into the sticky word, not the state. + void write_task_accepted() { + auto *ptr = reinterpret_cast(static_cast(mailbox_ptr()) + MAILBOX_OFF_ACCEPTED); + int32_t v = MAILBOX_TASK_ACCEPTED; + __atomic_store(ptr, &v, __ATOMIC_RELEASE); + } + void wait_running(int timeout_ms = 500) { auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); while (!is_running.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) { @@ -288,9 +295,15 @@ struct SchedulerFixture : public ::testing::Test { mock_worker.start(); manager.add_next_level(mock_worker.mailbox_ptr()); - manager.start(&allocator, [this](WorkerCompletion completion) { - sched.worker_done(std::move(completion)); - }); + manager.start( + &allocator, + [this](WorkerCompletion completion) { + sched.worker_done(std::move(completion)); + }, + [this](WorkerDispatch dispatch) { + orch.mark_task_accepted(dispatch.task_slot); + } + ); rq_next_level.reset(manager.next_level_worker_ids()); orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [this] { sched.notify_ready(); @@ -351,7 +364,7 @@ TEST(WorkerManagerTest, StartRejectsDuplicateNextLevelWorkerId) { bool threw = false; try { - manager.start(&allocator, [](WorkerCompletion) {}); + manager.start(&allocator, [](WorkerCompletion) {}, {}); } catch (const std::runtime_error &e) { threw = true; EXPECT_NE(std::string(e.what()).find("duplicate NEXT_LEVEL worker_id 0"), std::string::npos); @@ -362,6 +375,99 @@ TEST(WorkerManagerTest, StartRejectsDuplicateNextLevelWorkerId) { EXPECT_TRUE(threw); } +TEST(WorkerManagerTest, LocalMailboxPublishesAcceptanceBeforeCompletion) { + MockMailboxWorker child; + child.start(); + + Ring allocator; + allocator.init(/*heap_bytes=*/0); + AllocResult ar = allocator.alloc(/*heap_bytes=*/0, /*depth=*/0); + ASSERT_NE(ar.slot, INVALID_SLOT); + TaskSlotState *slot = allocator.slot_state(ar.slot); + ASSERT_NE(slot, nullptr); + slot->reset(); + slot->callable.digest[0] = 0x42; + + LocalMailboxEndpoint endpoint(/*worker_id=*/0, child.mailbox_ptr()); + std::promise result; + auto done = result.get_future(); + std::atomic accepted{false}; + std::thread caller([&] { + result.set_value(endpoint.run_with_accept(&allocator, WorkerDispatch{ar.slot, 0}, [&] { + accepted.store(true, std::memory_order_release); + })); + }); + + child.wait_running(); + EXPECT_TRUE(child.is_running.load(std::memory_order_acquire)); + child.write_task_accepted(); + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3); + while (!accepted.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) {} + EXPECT_TRUE(accepted.load(std::memory_order_acquire)); + EXPECT_EQ(done.wait_for(std::chrono::milliseconds(0)), std::future_status::timeout); + + // Non-fatal from here on: a fatal assertion would return with `caller` + // joinable, and ~std::thread would terminate the whole test binary. + child.complete(); + EXPECT_EQ(done.wait_for(std::chrono::seconds(3)), std::future_status::ready); + if (done.valid() && done.wait_for(std::chrono::seconds(0)) == std::future_status::ready) { + EXPECT_EQ(done.get().outcome, EndpointOutcome::SUCCESS); + } + caller.join(); + allocator.shutdown(); +} + +// The ACK must not be carried by anything TASK_DONE can overwrite. The child +// publishes acceptance only into the sticky word — never into the state — and +// then completes immediately, so an endpoint that looks for acceptance in the +// state word observes none at all. +// +// This does not force the parent to skip a poll between the two writes: the +// parent is already spinning by then and nothing here can stop it. What it +// pins is the property that makes that interleaving harmless — acceptance is +// readable after TASK_DONE, so losing a poll cannot lose the ACK. +TEST(WorkerManagerTest, AcceptanceIsReadableAfterTaskDone) { + MockMailboxWorker child; + child.start(); + + Ring allocator; + allocator.init(/*heap_bytes=*/0); + AllocResult ar = allocator.alloc(/*heap_bytes=*/0, /*depth=*/0); + ASSERT_NE(ar.slot, INVALID_SLOT); + TaskSlotState *slot = allocator.slot_state(ar.slot); + ASSERT_NE(slot, nullptr); + slot->reset(); + slot->callable.digest[0] = 0x42; + + LocalMailboxEndpoint endpoint(/*worker_id=*/0, child.mailbox_ptr()); + std::promise result; + auto done = result.get_future(); + std::atomic accepted{false}; + + std::thread caller([&] { + result.set_value(endpoint.run_with_accept(&allocator, WorkerDispatch{ar.slot, 0}, [&] { + accepted.store(true, std::memory_order_release); + })); + }); + + child.wait_running(); + EXPECT_TRUE(child.is_running.load(std::memory_order_acquire)); + // Back to back, with no parent poll in between. + child.write_task_accepted(); + child.complete(); + + // Non-fatal: a fatal assertion would return with `caller` joinable, and + // ~std::thread would terminate the whole test binary. + EXPECT_EQ(done.wait_for(std::chrono::seconds(3)), std::future_status::ready); + if (done.valid() && done.wait_for(std::chrono::seconds(0)) == std::future_status::ready) { + EXPECT_EQ(done.get().outcome, EndpointOutcome::SUCCESS); + } + EXPECT_TRUE(accepted.load(std::memory_order_acquire)) + << "the endpoint lost the launch ACK to a task that completed first"; + caller.join(); + allocator.shutdown(); +} + // A child that dies without publishing CONTROL_DONE must be reported, not // waited on forever. The mailbox stays at CONTROL_REQUEST exactly as it would // if the real `_chip_process_loop` had crashed mid-command. Run in a worker @@ -416,7 +522,7 @@ TEST(WorkerManagerTest, ControlPrepareUsesStableNextLevelWorkerId) { manager.add_next_level_endpoint(std::make_unique(7, &worker7_prepares)); manager.add_next_level_endpoint(std::make_unique(3, &worker3_prepares)); - manager.start(&allocator, [](WorkerCompletion) {}); + manager.start(&allocator, [](WorkerCompletion) {}, {}); std::array digest{}; manager.control_prepare(3, digest.data()); @@ -550,9 +656,15 @@ struct GroupSchedulerFixture : public ::testing::Test { worker_b.start(); manager.add_next_level(worker_a.mailbox_ptr()); manager.add_next_level(worker_b.mailbox_ptr()); - manager.start(&allocator, [this](WorkerCompletion completion) { - sched.worker_done(std::move(completion)); - }); + manager.start( + &allocator, + [this](WorkerCompletion completion) { + sched.worker_done(std::move(completion)); + }, + [this](WorkerDispatch dispatch) { + orch.mark_task_accepted(dispatch.task_slot); + } + ); rq_next_level.reset(manager.next_level_worker_ids()); orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [this] { sched.notify_ready(); @@ -911,9 +1023,15 @@ TEST(SchedulerWorkerTargetTest, NextLevelTargetUsesWorkerIdNotVectorIndex) { worker_b.start(); manager.add_next_level_at(7, worker_a.mailbox_ptr()); manager.add_next_level_at(9, worker_b.mailbox_ptr()); - manager.start(&allocator, [&sched](WorkerCompletion completion) { - sched.worker_done(std::move(completion)); - }); + manager.start( + &allocator, + [&sched](WorkerCompletion completion) { + sched.worker_done(std::move(completion)); + }, + [&orch](WorkerDispatch dispatch) { + orch.mark_task_accepted(dispatch.task_slot); + } + ); rq_next_level.reset(manager.next_level_worker_ids()); orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [&sched] { sched.notify_ready(); @@ -1039,9 +1157,15 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { sub_worker.start(); manager.add_next_level(next_level_worker.mailbox_ptr()); manager.add_sub(sub_worker.mailbox_ptr()); - manager.start(&allocator, [this](WorkerCompletion completion) { - sched.worker_done(std::move(completion)); - }); + manager.start( + &allocator, + [this](WorkerCompletion completion) { + sched.worker_done(std::move(completion)); + }, + [this](WorkerDispatch dispatch) { + orch.mark_task_accepted(dispatch.task_slot); + } + ); rq_next_level.reset(manager.next_level_worker_ids()); orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [this] { sched.notify_ready(); diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index ad7a4cbc61..6b22e5011f 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -1595,6 +1595,43 @@ def interrupted_wait(_self, _timeout=None): with pytest.raises(KeyboardInterrupt): handle._wait_for_serialization() + def test_acceptance_wait_does_not_block_completion_waiter(self): + acceptance_entered = threading.Event() + acceptance_release = threading.Event() + completion_entered = threading.Event() + completion_release = threading.Event() + + class FakeWorker: + def _wait_run_handle_accepted(self, run_id): + assert run_id == 1 + acceptance_entered.set() + assert acceptance_release.wait(5.0) + + def _wait_run_handle(self, run_id, timeout): + assert run_id == 1 + completion_entered.set() + assert completion_release.wait(5.0) + return True + + def _finalize_run_handle(self, handle, run_id, error): + return error + + handle = RunHandle(cast(Worker, FakeWorker.__new__(FakeWorker)), 1, ()) + completion_thread = threading.Thread(target=handle.wait) + acceptance_thread = threading.Thread(target=handle._wait_for_acceptance) + completion_thread.start() + assert completion_entered.wait(3.0) + acceptance_thread.start() + try: + assert acceptance_entered.wait(3.0) + finally: + acceptance_release.set() + completion_release.set() + acceptance_thread.join(5.0) + completion_thread.join(5.0) + assert not acceptance_thread.is_alive() + assert not completion_thread.is_alive() + def test_done_query_cannot_race_native_run_release(self): done_entered = threading.Event() done_release = threading.Event()