Skip to content

Update: elide A5 terminal releases after orchestration #2070

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

Open
yanghaoran29 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Update: elide A5 terminal releases after orchestration #2070

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
86 changes: 85 additions & 1 deletion docs/tensormap-and-ringbuffer-a2a3-vs-a5.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ This document describes the substantive differences in the current code under
`src/{a2a3,a5}/runtime/tensormap_and_ringbuffer/`.

> **Maintenance baseline:** The source layout and classifications were verified
> on 2026-08-17. Recompute the counts and update the affected sections whenever
> on 2026-09-03. Recompute the counts and update the affected sections whenever
> the files or constants described here change.

## Comparison Boundary and Classification
Expand Down Expand Up @@ -112,6 +112,7 @@ The functional differences group into the following themes:
| URMA completion | A5-specific implementation and product capability gate | Yes, for now | Retain the A5 path; do not claim that URMA is available in the default build |
| Next-block prefetch | A2/A3-only performance optimization | No | Retain on A2/A3; validate on A5 before considering a port |
| Scheduler progress publication | AICPU topology and measured publication cost | No | Retain A5's 16-task batching; keep per-advance publication on A2/A3, where the portable implementation showed no significant benefit |
| Terminal task release | Measured end-of-run scheduler cost | No | A5 traces show per-task release blocking the tail after task submission has ended, so successful A5 runs elide deferred release after the graph seal; retain incremental release on A2/A3 because no tail release blocking was found there |
| Fatal teardown | Software reliability strategy | No | Retain the current implementations; decide whether to converge after measuring the worst-case A5 teardown time |
| Scheduler trace attribution | Software diagnostic strategy | No | Preserve the current traces; converge only after comparing generated timelines |

Expand Down Expand Up @@ -293,6 +294,89 @@ measurements instead showed lower Effective time in all eight workloads, with
an unweighted mean reduction of `2.81%`. Full A2/A3 measurements are recorded
in the [PR benchmark follow-up](https://github.com/hw-native-sys/simpler/pull/1575#issuecomment-5310909143).

### Terminal Task Release: A5 Seal and Elision

Task completion and task release are separate scheduler operations. Completion
records the finished AICore work and unlocks dependent tasks. Release later
drops the completed task's retained references, advances the ring's reclaim
head across consumed slots, resets reusable slot state, and publishes reclaim
progress. Both platforms defer this release work in a per-scheduler array with
a capacity of 256 entries.

A2/A3 preserves the incremental protocol for the whole run. It drains the
array when it becomes full, during idle cleanup, and when dispatch exits. Every
completed task therefore reaches `on_task_release()` before the scheduler
returns.

A5 follows the same protocol while the Orchestrator can still submit tasks.
After `orchestrator_done_` seals the graph, however, no new task can require a
reclaimed ring slot. At the next existing full-array, idle-drain, or exit-drain
boundary, A5 discards the deferred-release backlog instead of calling
`on_task_release()` once per entry. Shared helper
`drain_or_elide_deferred_releases` owns that decision. The seal is deliberately
not loaded on every scheduler-loop iteration or every completion: completion
and dependency unlocking remain unchanged, and the added acquire loads stay on
boundaries that already perform release bookkeeping. At a sealed capacity
boundary, the overflowing completed slot is also not deferred.

The seal is stored only after a clean orchestration exit (`orch_error_code ==
NONE`). Failed orchestration never sets `orchestrator_done_`, so deferred
release stays exact until emergency teardown. A scheduler or async failure
after a successful seal leaves `orchestrator_done_` true: later idle/exit
drains still elide, and the next-run SM reset closes the lifecycle.

Skipped incremental release does not need a terminal barrier or bulk slot
closure. The next run clears the entire SM with `memset` and rebuilds flow
control via `init_per_ring` → `fc.init()`, and the orchestrator already
self-cleans each reused slot on submit. There is no functional downstream
reader of successful-exit `CONSUMED` watermarks between runs. An earlier
barrier / `TerminalClose` layer was removed as redundant.

This optimization is independent of the A5 K=16 progress-publication policy in
the preceding section. K=16 controls how often an already-advanced reclaim
head is copied to shared memory during the run. Terminal release elision avoids
the per-task reference-count and ring-advance work itself after graph sealing;
its deferred-release array still has capacity 256 and does not impose a
16-task release limit.

| Stage | A2/A3 | A5 |
| ----- | ----- | -- |
| Before graph sealing | Complete tasks, defer release, then incrementally call `on_task_release()` | Same |
| After graph sealing | Continue incremental release | Drop deferred release work at existing release boundaries |
| Successful Scheduler exit | Drain every remaining deferred entry | Drop remaining deferred entries; next run resets SM |
| Orchestration failure (never sealed) | Emergency teardown after the existing error checks | Emergency teardown; `orchestrator_done_` stays false, so deferred release stays exact until teardown |
| Scheduler / async failure after a successful seal | Emergency teardown | `orchestrator_done_` remains true, so exit/idle drains still elide; SM reset on the next run closes the lifecycle |
| Profiling | Release phases | Release phases only when a real drain runs (no `terminal_close`) |

| File | A5-only terminal-release role |
| ---- | ----------------------------- |
| `runtime/async_wait.h`, `runtime/scheduler/scheduler_completion.cpp` | Sync completion reads the graph seal at deferred-release capacity boundaries; async capacity drains keep exact release |
| `runtime/scheduler/scheduler_dispatch.cpp` | Elide sealed idle/exit backlog drains via shared `drain_or_elide_deferred_releases`; skip DFX `release` when elided |
| `runtime/scheduler/scheduler.h` | Define the shared drain-or-elide helper used by deferred-release sites |
| `runtime/scheduler/scheduler_cold_path.cpp` | Seal `orchestrator_done_` only after a clean orchestration exit |

The post-removal A/B was run only on the local A5 system against current
`main`. The seven non-Qwen workloads improved by `9.753%` in Effective
geometric mean; Qwen3 changed by `-0.010%`, and no workload regressed by 5% or
more in Effective time. The largest Scheduler gains remain in the
paged-attention-unroll family (`10.550%` to `14.021%` Effective). Full tables
are recorded under the [PR #2070](https://github.com/hw-native-sys/simpler/pull/2070)
benchmark discussion / local `outputs/pr2070_full_bench/` artifacts.

The platform scope follows the observed bottleneck. On A5, the motivating
timelines contain a visible tail after the Orchestrator has finished submitting
tasks: Schedulers continue executing per-task release work even though no new
task can consume the reclaimed capacity. That release interval extends the
execution critical path, which gives seal-and-elide a direct optimization
target. No tail release blocking was found on A2/A3. Its release work did not
appear as the corresponding post-orchestration critical-path interval, so there
is currently no performance evidence that A2/A3 would benefit from the extra
graph-seal observation. A2/A3 therefore keeps the simpler incremental release
protocol, and this experiment does not run an A2/A3 benchmark or port the
implementation there. This is an evidence-based software decision rather than
an A5 hardware requirement; revisit it if a future A2/A3 timeline exposes the
same tail release blocking.

### Fatal Teardown

The A2/A3 scheduler uses a dedicated fatal latch to elect an owner, broadcasts
Expand Down
46 changes: 34 additions & 12 deletions src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1314,6 +1314,29 @@ struct SchedulerState {
// Scheduler cold-path API is declared as SchedulerState member functions.
// See init()/destroy()/print_stats()/print_queues() below the struct definition.

// Drop deferred releases when release_elided; otherwise drain via on_task_release.
// Callers set release_elided from orchestrator_done_ only at existing release
// boundaries (full buffer / idle / exit); the async path always passes false.
inline void drain_or_elide_deferred_releases(
SchedulerState *sched, ChipTaskSlotState **slots, int32_t &count, bool release_elided
#if SIMPLER_SCHED_PROFILING
,
int32_t thread_idx
#endif
) {
if (release_elided) {
count = 0;
return;
}
while (count > 0) {
#if SIMPLER_SCHED_PROFILING
(void)sched->on_task_release(*slots[--count], thread_idx);
#else
sched->on_task_release(*slots[--count]);
#endif
}
}

// Short-circuit NotDeferred completions seen during drain so they don't grow
// entries[]. Mirrors the a2a3 impl; see that mirror for the rationale.
inline bool
Expand All @@ -1323,16 +1346,15 @@ AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &si
#else
sink.sched->on_task_complete(slot_state);
#endif
// Async path keeps exact deferred release (no graph-seal observation here).
if (*sink.deferred_release_count >= sink.deferred_release_capacity) {
while (*sink.deferred_release_count > 0) {
drain_or_elide_deferred_releases(
sink.sched, sink.deferred_release_slot_states, *sink.deferred_release_count, /*release_elided=*/false
#if SIMPLER_SCHED_PROFILING
(void)sink.sched->on_task_release(
*sink.deferred_release_slot_states[--(*sink.deferred_release_count)], sink.thread_idx
);
#else
sink.sched->on_task_release(*sink.deferred_release_slot_states[--(*sink.deferred_release_count)]);
,
sink.thread_idx
#endif
}
);
}
sink.deferred_release_slot_states[(*sink.deferred_release_count)++] = &slot_state;
sink.inline_completed++;
Expand Down Expand Up @@ -1395,13 +1417,13 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete(
sched->on_task_complete(*entry.slot_state);
#endif
if (deferred_release_count >= deferred_release_capacity) {
while (deferred_release_count > 0) {
drain_or_elide_deferred_releases(
sched, deferred_release_slot_states, deferred_release_count, /*release_elided=*/false
#if SIMPLER_SCHED_PROFILING
(void)sched->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx);
#else
sched->on_task_release(*deferred_release_slot_states[--deferred_release_count]);
,
thread_idx
#endif
}
);
}
deferred_release_slot_states[deferred_release_count++] = entry.slot_state;
result.completed++;
Expand Down
9 changes: 7 additions & 2 deletions src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,10 @@
*/
#include "scheduler_context.h"

#include <algorithm>
#include <cinttypes>
#include <cstdio>
#include <limits>

#include "common/unified_log.h"
#include "aicpu/aicpu_device_config.h"
Expand Down Expand Up @@ -1379,9 +1381,10 @@ void SchedulerContext::on_orchestration_done(
rt->scheduler.tasks_completed.fetch_add(inline_completed, std::memory_order_relaxed);
#endif
}
orchestrator_done_.store(true, std::memory_order_release);

// Check for fatal error from orchestration; if so, shut down immediately.
// Seal only on a clean orchestration exit so deferred-release elision never
// discards lifecycle work after an orch error. Failed runs keep the
// unsealed exact-release path until emergency teardown.
int32_t orch_err = 0;
if (sched_->sm_header) {
orch_err = sched_->sm_header->orch_error_code.load(std::memory_order_relaxed);
Expand All @@ -1390,6 +1393,8 @@ void SchedulerContext::on_orchestration_done(
if (!completed_.exchange(true, std::memory_order_acq_rel)) {
emergency_shutdown(runtime);
}
} else {
orchestrator_done_.store(true, std::memory_order_release);
}

#if SIMPLER_DFX
Expand Down
20 changes: 12 additions & 8 deletions src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,16 +199,20 @@ void SchedulerContext::complete_slot_task(
}
chip_swimlane.phase_complete_count++;
#endif
if (deferred_release_count < DEFERRED_RELEASE_CAP) {
deferred_release_slot_states[deferred_release_count++] = &slot_state;
} else {
while (deferred_release_count > 0) {
// At capacity, elide deferred releases (including the overflowing slot)
// when orchestration is done; otherwise drain then push.
bool release_elided = false;
if (deferred_release_count >= DEFERRED_RELEASE_CAP) {
release_elided = orchestrator_done_.load(std::memory_order_acquire);
drain_or_elide_deferred_releases(
sched_, deferred_release_slot_states, deferred_release_count, release_elided
#if SIMPLER_SCHED_PROFILING
(void)sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx);
#else
sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count]);
,
thread_idx
#endif
}
);
}
if (!release_elided) {
deferred_release_slot_states[deferred_release_count++] = &slot_state;
}
completed_this_turn++;
Expand Down
51 changes: 26 additions & 25 deletions src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1174,19 +1174,18 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_
}
#endif
// Dummy tasks have no subtasks to retire and no fanout pre-conditions
// beyond their own producers; release self-reference so the slot can
// reach CONSUMED once all consumers drain.
// beyond their own producers. While lifecycle reclamation is active,
// release their self-reference so the slot can reach CONSUMED.
deferred_release_slot_states[deferred_release_count++] = &dummy_slot;
if (deferred_release_count >= DEFERRED_RELEASE_CAP) {
while (deferred_release_count > 0) {
bool release_elided = orchestrator_done_.load(std::memory_order_acquire);
drain_or_elide_deferred_releases(
sched_, deferred_release_slot_states, deferred_release_count, release_elided
#if SIMPLER_SCHED_PROFILING
(void)sched_->on_task_release(
*deferred_release_slot_states[--deferred_release_count], thread_idx
);
#else
sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count]);
,
thread_idx
#endif
}
);
}
int32_t prev = completed_tasks_.fetch_add(1, std::memory_order_relaxed);
last_progress_count = prev + 1;
Expand Down Expand Up @@ -1298,15 +1297,18 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_
0;
uint32_t released_count = static_cast<uint32_t>(deferred_release_count);
#endif
while (deferred_release_count > 0) {
// Elide remaining deferred releases when orchestration is done;
// otherwise drain via on_task_release.
bool release_elided = deferred_release_count > 0 && orchestrator_done_.load(std::memory_order_acquire);
drain_or_elide_deferred_releases(
sched_, deferred_release_slot_states, deferred_release_count, release_elided
#if SIMPLER_SCHED_PROFILING
(void)sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx);
#else
sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count]);
,
thread_idx
#endif
}
);
#if SIMPLER_DFX
if (release_t0 != 0) {
if (release_t0 != 0 && !release_elided) {
chip_swimlane_aicpu_record_sched_phase(
thread_idx, ChipSwimlaneSchedPhaseKind::Release, release_t0, get_sys_cnt_aicpu(),
chip_swimlane.sched_loop_count, released_count
Expand Down Expand Up @@ -1384,18 +1386,17 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_
}
}

// Drain any entries left in the deferred-release batch. The in-loop flush
// only fires on idle iterations and on buffer-full; a loop exit while the
// last iteration made progress can leave entries un-released. Drop them
// here so every consumed producer slot completes its on_task_release
// regardless of which loop-exit path fired.
while (deferred_release_count > 0) {
// Exact release on every exit while orchestration is still running.
// After orchestrator_done_, drop remaining deferred releases — the next
// run memset+reinit clears SM, and reused slots self-clean on submit.
bool release_elided = deferred_release_count > 0 && orchestrator_done_.load(std::memory_order_acquire);
drain_or_elide_deferred_releases(
sched_, deferred_release_slot_states, deferred_release_count, release_elided
#if SIMPLER_SCHED_PROFILING
(void)sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx);
#else
sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count]);
,
thread_idx
#endif
}
);

#if SIMPLER_DFX
// Final-drain: emit any pop_hit / pop_miss accrued since the last
Expand Down
Loading
Loading