Skip to content

Refactor: submit every a2a3 run on one AICore/AICPU stream pair #1852

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChaoWao merged 1 commit into from
Aug 17, 2026
Merged

Refactor: submit every a2a3 run on one AICore/AICPU stream pair #1852

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 22 additions & 10 deletions docs/task-flow.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -290,16 +290,28 @@ that window needs the admission layer to gate dispatch on `pool.owns(lease)`
before handing work down, which is whole-run admission's job, not this
layer's.

AICore streams are outside that lease. Registered callables keep content-hash
deduplicated GM allocations: simultaneously resident code images occupy
different allocations, while unregister frees an allocation that a later
registration may reuse. A dedup miss is the only path that publishes new
AICore instruction bytes. After that H2D copy succeeds, every retained slot
stream is marked stale; its next acquire destroys and recreates the AICore
stream before launch. Without a new publication, completed slot streams remain
reusable even when two resident callables alternate. Unproven completion still
destroys the stream conservatively. The AICPU stream carries no instruction
cache state and stays with its slot.
Run streams are outside that lease, and there is one pair of them per runner
rather than one per slot. A slot indexes the resources *preparation* mutates,
and preparing a run writes nothing to a stream: only launch submits, and launch
holds the exclusive execution claim, so runs reach the device one at a time and
the stream orders them. The two streams stay distinct because the AICPU Run
kernel spins in the handshake waiting for the AICore workers — one queue would
leave the AICore submission behind a spin that never ends.

The AICore stream carries instruction-cache state. Registered callables keep
content-hash deduplicated GM allocations: simultaneously resident code images
occupy different allocations, while unregister frees an allocation that a later
registration may reuse. A dedup miss is the only repeatable path that publishes
new AICore instruction bytes; after that H2D copy succeeds the pair is marked
stale, and the next launch destroys the AICore stream and creates a replacement.
(The `kernel_entry` ELF that `rtRegisterAllKernel` publishes needs no such mark:
CANN offers no unregister, so it is registered once per runner and released with
the streams at finalize.) Without a new publication the pair stays warm even
when two resident callables alternate. Unproven completion still destroys the
AICore stream conservatively, and only the run that submitted the pair may
retire it — a prepared successor overlaps its predecessor's execution and must
leave the live pair alone. The AICPU stream carries no instruction cache state
and lives for the runner.

#### Whole-run FIFO admission

Expand Down
9 changes: 5 additions & 4 deletions python/bindings/task_interface.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2304,10 +2304,11 @@ NB_MODULE(_task_interface, m) {
)
.def_prop_ro(
"run_stream_set_create_count", &ChipWorker::run_stream_set_create_count,
"Number of AICore run streams the bound runner has created. The AICPU "
"stream belongs to a pipeline slot for the worker's lifetime. A completed "
"AICore stream is reused until a new code upload makes it stale. Platforms "
"using the persistent bootstrap pair report 0."
"Number of AICore run streams the bound runner has created. One AICPU "
"+ AICore pair serves every run for the runner's lifetime. The AICPU "
"stream persists; the AICore stream is recreated when a new code upload "
"makes it stale, and destroyed when an unproven completion retires it. "
"Platforms using the persistent bootstrap pair report 0."
Comment thread
coderabbitai[bot] marked this conversation as resolved.
)
.def_prop_ro(
"committed_device_memory", &ChipWorker::committed_device_memory,
Expand Down
11 changes: 7 additions & 4 deletions python/simpler/worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11131,10 +11131,13 @@ def host_dlopen_count(self) -> int:
def run_stream_set_create_count(self) -> int:
"""L2 only: number of AICore run streams the runner has created.

AICPU streams belong to pipeline slots for the worker's lifetime. A
completed AICore stream is reused until a new code upload makes it
stale. Returns 0 on non-L2 workers and on platforms whose runs use the
persistent bootstrap stream pair (simulation, a5).
One AICPU + AICore pair serves every run for the runner's lifetime. The
AICPU stream persists; the AICore stream is recreated when a new code
upload makes it stale, and destroyed when an unproven completion retires
it, so this advances per publication or unproven retirement rather than
once per run or per pipeline slot. 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
Expand Down
74 changes: 31 additions & 43 deletions src/a2a3/platform/onboard/host/device_runner.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,29 +217,16 @@ void DeviceRunner::set_dep_gen_enabled(bool enable) {
dep_gen_host_graph_set_enabled(enable);
}

int DeviceRunner::provision_native_run_resources(uint32_t pipeline_slot) {
return ensure_run_stream_set(pipeline_slot);
}

int DeviceRunner::abandon_native_run_resources(uint32_t pipeline_slot) {
return retire_run_aicore_stream(pipeline_slot, RunStreamSlots::CompletionStatus::Unproven);
}

int DeviceRunner::prepare_execution(
Runtime &runtime, const CallConfig &config, uint32_t pipeline_slot, const NativeRunIdentity &identity,
std::unique_ptr<PreparedExecution> *prepared
) {
if (prepared == nullptr || *prepared != nullptr) return -1;
const unsigned selected_pipeline_slot = pipeline_slot;
auto execution = std::make_unique<PreparedExecution>(identity, runtime, config, pipeline_slot);
execution->resources_owned = true;
auto prepare_rollback = RAIIScopeGuard([this, &execution]() {
cleanup_execution(*execution, /*retire_aicore=*/false);
});
if (!run_stream_slots_.ready(selected_pipeline_slot)) {
LOG_ERROR("run stream set %u was not provisioned during native prepare", selected_pipeline_slot);
return -1;
}
const int block_dim = runtime.get_worker_count() / cores_per_blockdim_;
int launch_aicpu_num = config.aicpu_thread_num;
// A prior AICore launch/sync error poisoned the device context and the
Expand Down Expand Up @@ -444,21 +431,19 @@ int DeviceRunner::prepare_execution(

int DeviceRunner::poll_execution(const ActiveExecution &active) {
if (active.prepared == nullptr) return SIMPLER_NATIVE_RUN_POLL_ERROR;
const uint32_t pipeline_slot = active.prepared->pipeline_slot;
return run_stream_slots_.poll(pipeline_slot, [](void *aicpu, void *aicore) {
return run_streams_.poll([](void *aicpu, void *aicore) {
return query_stream_pair_nonblocking(static_cast<rtStream_t>(aicpu), static_cast<rtStream_t>(aicore));
});
}

int DeviceRunner::drain_execution(ActiveExecution &active) {
if (active.prepared == nullptr || !active.prepared->resources_owned) return -1;
PreparedExecution &prepared = *active.prepared;
const uint32_t pipeline_slot = prepared.pipeline_slot;
auto drain_cleanup = RAIIScopeGuard([this, &prepared]() {
cleanup_execution(prepared, /*retire_aicore=*/true);
});

int rc = reap_run(pipeline_slot);
int rc = reap_run();
if (rc != 0) {
// The device/sync error remains authoritative over teardown errors.
return rc;
Expand All @@ -468,7 +453,7 @@ int DeviceRunner::drain_execution(ActiveExecution &active) {
// stale. Publish retirement so cleanup does not replace it with an
// unproven state after the device result has already been established.
prepared.aicore_retirement_attempted = true;
rc = retire_run_aicore_stream(pipeline_slot, RunStreamSlots::CompletionStatus::Complete);
rc = retire_run_aicore_stream(&prepared, RunStreamPair::CompletionStatus::Complete);
if (rc != 0) return rc;

// Reads device memory, so it must precede KernelArgs/runtime cleanup.
Expand Down Expand Up @@ -507,7 +492,7 @@ void DeviceRunner::cleanup_execution(PreparedExecution &prepared, bool retire_ai
}
if (retire_aicore && !abandon && !prepared.aicore_retirement_attempted) {
prepared.aicore_retirement_attempted = true;
(void)retire_run_aicore_stream(prepared.pipeline_slot, RunStreamSlots::CompletionStatus::Unproven);
(void)retire_run_aicore_stream(&prepared, RunStreamPair::CompletionStatus::Unproven);
}
prepared.resources_owned = false;
}
Expand Down Expand Up @@ -547,30 +532,30 @@ int DeviceRunner::create_run_stream(void **out) {
return 0;
}

int DeviceRunner::ensure_run_stream_set(unsigned slot) {
int rc = run_stream_slots_.acquire(slot);
int DeviceRunner::ensure_run_streams() {
int rc = run_streams_.ensure();
if (rc != 0) {
LOG_ERROR("ensure_run_stream_set(%u) failed: %d", slot, rc);
LOG_ERROR("ensure_run_streams failed: %d", rc);
ACL_LOG_ERROR_DETAIL(rc);
}
return rc;
}

int DeviceRunner::retire_run_aicore_stream(unsigned slot, RunStreamSlots::CompletionStatus completion_status) {
int rc = run_stream_slots_.retire_aicore(slot, completion_status);
int DeviceRunner::retire_run_aicore_stream(const void *owner, RunStreamPair::CompletionStatus completion_status) {
int rc = run_streams_.retire(completion_status, owner);
if (rc != 0) {
LOG_ERROR("rtStreamDestroy (run AICore slot %u) failed: %d, slot is now unusable", slot, rc);
LOG_ERROR("rtStreamDestroy (run AICore stream) failed: %d, the handle is retained for teardown", rc);
}
return rc;
}

int DeviceRunner::destroy_run_stream_sets() {
int DeviceRunner::destroy_run_streams() {
// No pre-destroy sync, for the reason finalize_common() documents for the
// bootstrap pair: rtStreamDestroy is the supported teardown for a stream
// left in the error state by an op-timeout.
int rc = run_stream_slots_.destroy_all();
int rc = run_streams_.destroy();
if (rc != 0) {
LOG_ERROR("destroy_run_stream_sets: a stream survived teardown: %d", rc);
LOG_ERROR("destroy_run_streams: a stream survived teardown: %d", rc);
}
return rc;
}
Expand All @@ -597,15 +582,18 @@ LaunchTransactionResult DeviceRunner::launch_run(PreparedExecution &prepared, La
Runtime &runtime = *prepared.runtime;
const int num_aicore = prepared.num_aicore;
const int launch_aicpu_num = prepared.launch_aicpu_num;
const unsigned slot = prepared.pipeline_slot;
// KernelLaunch is the pipeline boundary: this method clears the handshake
// consumed by the launch and submits exactly the AICore and AICPU kernels.
// It intentionally performs no stream synchronization or per-run cleanup.
if (!run_stream_slots_.ready(slot)) {
LOG_ERROR("launch_run: stream set %u is not ready", slot);
//
// The pair is readied here rather than at prepare because this is the first
// point the caller holds the execution claim: a prepared successor overlaps
// its predecessor's execution, so replacing a stale AICore stream during
// preparation would destroy a stream the predecessor is still running on.
if (ensure_run_streams() != 0) {
return LaunchTransactionResult{};
}
RunStreamSet streams{run_stream_slots_.aicpu(slot), run_stream_slots_.aicore(slot)};
RunStreamSet streams{static_cast<rtStream_t>(run_streams_.aicpu()), static_cast<rtStream_t>(run_streams_.aicore())};
LaunchTransactionResult result = exact_launch_transaction(
prepared.identity, std::move(permit),
[&]() {
Expand Down Expand Up @@ -655,9 +643,9 @@ LaunchTransactionResult DeviceRunner::launch_run(PreparedExecution &prepared, La

// Publishing query/drain ownership is a state change, so it sits past
// the arming block: everything the block covers is rollback-free.
int rc = run_stream_slots_.mark_submitted(slot);
int rc = run_streams_.mark_submitted(&prepared);
if (rc != 0) {
LOG_ERROR("launch_run: failed to publish stream set %u: %d", slot, rc);
LOG_ERROR("launch_run: failed to publish the run stream pair: %d", rc);
return rc;
}

Expand Down Expand Up @@ -685,12 +673,12 @@ LaunchTransactionResult DeviceRunner::launch_run(PreparedExecution &prepared, La
return result;
}

int DeviceRunner::reap_run(unsigned slot) {
if (!run_stream_slots_.ready(slot)) {
LOG_ERROR("reap_run: invalid stream set %u", slot);
int DeviceRunner::reap_run() {
if (!run_streams_.ready()) {
LOG_ERROR("reap_run: the run stream pair is not ready");
return -1;
}
int rc = sync_stream_pair(run_stream_slots_.aicpu(slot), run_stream_slots_.aicore(slot));
int rc = sync_stream_pair(run_streams_.aicpu(), run_streams_.aicore());
if (rc != 0) {
// The pair wait surfaces the AICore op-timeout (STARS-reaped op ->
// 507000/507018/507046 at AICPU/AICore stream sync). The op-timeout
Expand Down Expand Up @@ -1044,7 +1032,7 @@ int DeviceRunner::finalize() {
);
}

run_stream_slots_.abandon_all();
run_streams_.abandon();
int abandon_rc = abandon_common_after_device_failure();

// Only finalize the ACL owner after force reset established a clean
Expand Down Expand Up @@ -1082,10 +1070,10 @@ int DeviceRunner::finalize() {
// for the no-run-since-init case.
finalize_collectors();

// The per-run stream sets are this subclass's own RTS-owning member, so
// they are released here, while RTS is live and before the device reset
// below — the same window finalize_common() uses for the bootstrap pair.
int stream_rc = destroy_run_stream_sets();
// The run stream pair is this subclass's own RTS-owning member, so it is
// released here, while RTS is live and before the device reset below — the
// same window finalize_common() uses for the bootstrap pair.
int stream_rc = destroy_run_streams();

// Shared cleanup body — streams, kernel_args, callable/orch maps,
// chip-callable buffer pool, the three arenas, device_wall,
Expand Down
42 changes: 21 additions & 21 deletions src/a2a3/platform/onboard/host/device_runner.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
#include "callable.h"
#include "prepare_callable_common.h"
#include "pto_runtime_c_api.h" // PTO_PIPELINE_MAX_DEPTH
#include "host/run_stream_slots.h"
#include "host/run_stream_pair.h"
#include "common/kernel_args.h"
#include "common/memory_barrier.h"
#include "common/chip_swimlane_profiling.h"
Expand Down Expand Up @@ -103,9 +103,10 @@ class DeviceRunner : public DeviceRunnerBase {
int poll_execution(const ActiveExecution &active) override;
int drain_execution(ActiveExecution &active) override;
bool can_accept_run() const override { return !device_unusable_.load(std::memory_order_acquire); }
int provision_native_run_resources(uint32_t pipeline_slot) override;
int abandon_native_run_resources(uint32_t pipeline_slot) override;
void mark_run_streams_stale() override { run_stream_slots_.mark_all_stale(); }
// provision/abandon_native_run_resources keep the base no-op: preparation
// owns no stream, so there is nothing for a prepared run to provision or
// hand back. launch_run() readies the pair under the execution claim.
void mark_run_streams_stale() override { run_streams_.mark_stale(); }

// Map/unmap a device buffer into host address space via
// halHostRegister(DEV_SVM_MAP_HOST) / halHostUnregister. The returned host
Expand Down Expand Up @@ -182,7 +183,7 @@ class DeviceRunner : public DeviceRunnerBase {
// `aicpu_dlopen_count`, and `host_dlopen_count` are inherited from
// `DeviceRunnerBase`.

size_t run_stream_set_create_count() const override { return run_stream_slots_.created_count(); }
size_t run_stream_set_create_count() const override { return run_streams_.created_count(); }

private:
// Most lifecycle state (device_id_, block_dim_, cores_per_blockdim_,
Expand Down Expand Up @@ -223,14 +224,12 @@ class DeviceRunner : public DeviceRunnerBase {
rtStream_t aicore{nullptr};
};

// One slot per in-flight run the pipeline contract can declare.
static constexpr unsigned kRunStreamSetCount = PTO_PIPELINE_MAX_DEPTH;

// AICPU streams belong to slots for the worker's lifetime. A completed
// AICore stream remains reusable until a new code upload marks every slot
// stale. Stream lifetimes live in RunStreamSlots so publication and
// failed-destroy states are testable without a device.
RunStreamSlots run_stream_slots_{
// One pair carries every run: the execution claim is exclusive, so runs
// reach the device one at a time and the stream orders them. Pipeline slots
// index resources that preparation mutates, and preparing a run writes
// nothing to a stream. Stream lifetimes live in RunStreamPair so
// publication and failed-destroy states are testable without a device.
RunStreamPair run_streams_{
[this](void **out) {
return create_run_stream(out);
},
Expand All @@ -239,13 +238,14 @@ class DeviceRunner : public DeviceRunnerBase {
}
};
int create_run_stream(void **out);
int ensure_run_stream_set(unsigned slot);
// Destroys this run's AICore stream. Returns the driver's error and KEEPS
// the handle when the destroy fails: the stream may still hold the previous
// image's instructions, so the slot must refuse the next run until finalize
// reclaims it.
int retire_run_aicore_stream(unsigned slot, RunStreamSlots::CompletionStatus completion_status);
int destroy_run_stream_sets();
int ensure_run_streams();
// Retires the run pair on behalf of `owner`, the PreparedExecution that
// submitted it. Returns the driver's error and KEEPS the handle when an
// unproven destroy fails, so finalize can reclaim it. A caller that never
// submitted retires nothing: a prepared successor overlaps its
// predecessor's execution and must not destroy the live pair.
int retire_run_aicore_stream(const void *owner, RunStreamPair::CompletionStatus completion_status);
int destroy_run_streams();

// Release execution-owned resources in collector, runtime-argument,
// register-buffer, then stream order. The collectors this releases were
Expand All @@ -257,7 +257,7 @@ class DeviceRunner : public DeviceRunnerBase {
// The kernel submission boundary is separate from the stream wait and
// post-run teardown: launch_run() submits and drain_execution() reaps.
LaunchTransactionResult launch_run(PreparedExecution &prepared, LaunchPermit permit);
int reap_run(unsigned slot);
int reap_run();

// On an AICore launch/sync error, best-effort drain the device so a later
// enqueue on the same DeviceRunner can recover in place; if the drain itself
Expand Down
Loading
Loading