Skip to content

perf(runtime): cut TMARB preamble ~72us -> ~9us via report-first + batched handshake #1323

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
Jul 10, 2026
Merged

perf(runtime): cut TMARB preamble ~72us -> ~9us via report-first + batched handshake #1323

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
13 changes: 13 additions & 0 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 @@ -432,6 +432,19 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) {
// slow-launch / 207001 wedge was measured on a5; this mirror is UNVERIFIED on
// a2a3 silicon (the dev box is a5-only), relying on CI. See
// docs/investigations/2026-06-pa-unroll-207001-optimeout-window.md.
// The AICore publishes aicore_done on launch (gated by nothing), and the
// workers region persists across runs in the pooled arena. Clearing each
// worker's aicore_done before the AICore kernel launches keeps the AICPU's
// handshake sweep from reading a prior run's report — which would open a
// window on that run's physical_core_id. Only aicore_done needs clearing; the
// AICore overwrites physical_core_id/core_type in the same report.
{
Handshake *workers = runtime.get_workers();
for (int i = 0; i < num_aicore; i++) {
workers[i].aicore_done = 0;
}
}

LOG_INFO_V0("=== launch_aicore_kernel ===");
// Launch AICore kernel (pass device copy of KernelArgs)
rc = launch_aicore_kernel(stream_aicore_, kernel_args_.device_k_args_);
Expand Down
24 changes: 12 additions & 12 deletions src/a2a3/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,36 +69,36 @@ __aicore__ __attribute__((always_inline)) static void execute_task(__gm__ PTO2Di
__aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, int block_idx, CoreType core_type) {
__gm__ Handshake *my_hank = (__gm__ Handshake *)(&runtime->dev.workers[block_idx]);

// Phase 1: Wait for AICPU initialization signal
while (my_hank->aicpu_ready == 0) {
dcci(my_hank, SINGLE_CACHE_LINE);
SPIN_WAIT_HINT();
}

// Phase 2: Report physical core ID + core type and signal done in one write.
// The AICPU opens this core's register window (platform_init_aicore_regs:
// FAST_PATH + DATA_MAIN_BASE=IDLE) only after it observes aicore_done, so a
// single report suffices — there is no separate aicpu_regs_ready round-trip.
// Phase 1: report physical core ID + core type and signal done in one write,
// with no wait for the AICPU — both fields are self-known. The AICPU opens
// this core's register window only after it observes aicore_done, so a single
// report suffices. The host clears aicore_done before this kernel launches,
// so the value the AICPU reads is this run's report, never a stale prior one.
my_hank->physical_core_id = get_physical_core_id();
my_hank->core_type = core_type;
OUT_OF_ORDER_STORE_BARRIER();
my_hank->aicore_done = block_idx + 1; // Signal ready (use block_idx + 1 to avoid 0)
dcci(my_hank, SINGLE_CACHE_LINE, CACHELINE_OUT);

// Phase 3: Wait for the AICPU to open our register window. A kernel launch
// Phase 2: Wait for the AICPU to open our register window. A kernel launch
// resets DATA_MAIN_BASE to 0 (verified on a2a3 silicon); the AICPU writes
// DATA_MAIN_BASE = AICPU_IDLE_TASK_ID (non-zero) as it opens FAST_PATH, so a
// non-zero read means the window is open and reads/writes are valid. The
// AICPU runs assign_cores_to_threads (µs) between opening the window and the
// first dispatch, so this IDLE is observed long before any task_id lands —
// the poll cannot miss it and mistake a later task for the reset value.
// Window-open is the sync point for everything the AICPU publishes (task
// pointer, swimlane head): the AICPU writes those before opening the window.
while (read_reg(RegId::DATA_MAIN_BASE) == 0) {
SPIN_WAIT_HINT();
}
// Report initial idle status via register (FAST_PATH is now open).
write_reg(RegId::COND, AICORE_IDLE_VALUE);

// Cache per-core dispatch payload pointer (set by AICPU before aicpu_ready)
// The AICPU writes task after observing our report (so our CACHELINE_OUT flush
// above cannot clobber it) and before opening the window; dcci to read its
// fresh value here.
dcci(my_hank, SINGLE_CACHE_LINE);
__gm__ PTO2DispatchPayload *payload = reinterpret_cast<__gm__ PTO2DispatchPayload *>(my_hank->task);

uint32_t enable_profiling_flag = get_aicore_profiling_flag();
Expand Down
136 changes: 84 additions & 52 deletions src/a2a3/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 @@ -673,19 +673,12 @@ void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32
const int32_t lo = static_cast<int32_t>((static_cast<int64_t>(tidx) * total) / nthreads);
const int32_t hi = static_cast<int32_t>((static_cast<int64_t>(tidx + 1) * total) / nthreads);

// Step 1: publish this slice's payload pointers, then release its cores. One
// barrier (not one per core) suffices: it guarantees every task store is
// globally visible before any aicpu_ready store, which is the only ordering
// AICore relies on (it reads task only after observing aicpu_ready==1). Each
// core is written by exactly one thread, so the slices never overlap.
for (int32_t i = lo; i < hi; i++) {
all_handshakes[i].task = reinterpret_cast<uint64_t>(&payload_per_core_[i][0]);
}
OUT_OF_ORDER_STORE_BARRIER();
for (int32_t i = lo; i < hi; i++) {
all_handshakes[i].aicpu_ready = 1;
}
OUT_OF_ORDER_STORE_BARRIER();
// The AICore publishes {physical_core_id, core_type, aicore_done} on launch,
// gated by nothing. task is not published here: the AICore's aicore_done
// report flushes its whole handshake cache line, so a task stored before the
// report would be clobbered. task is written per core in the sweep below,
// after that core's aicore_done is observed and before its window opens (the
// point the AICore reads task).

// Get platform physical cores count for validation
uint32_t max_physical_cores_count = platform_get_physical_cores_count();
Expand All @@ -709,6 +702,22 @@ void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32
uint64_t *regs = reinterpret_cast<uint64_t *>(regs_);
bool core_serviced[RUNTIME_MAX_WORKER] = {false};

// Every core publishes aicore_done on launch, so the whole slice is already
// reported when the AICPU sweeps it. The reported cores are collected first,
// then serviced in batched phases (publish tasks, open windows, store
// CoreExecStates); each phase issues its stores without interleaving another
// phase's, so posted MMIO STRs and write-through GM stores do not serialize.
struct ReadyCore {
int32_t i;
uint32_t pcid;
uint64_t reg_addr;
CoreType core_type;
};
ReadyCore ready[RUNTIME_MAX_WORKER];
int32_t n_ready = 0;

// Phase 1: collect every reported core in this slice and prefetch its
// CoreExecState line for write, so the Phase 4 struct store hits a warm line.
for (int32_t remaining = hi - lo; remaining > 0;) {
for (int32_t i = lo; i < hi; i++) {
if (core_serviced[i]) continue;
Expand All @@ -717,7 +726,6 @@ void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32
SPIN_WAIT_HINT();
continue;
}

uint32_t physical_core_id = hank->physical_core_id;
if (physical_core_id >= max_physical_cores_count) {
LOG_ERROR(
Expand All @@ -729,25 +737,48 @@ void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32
remaining--;
continue;
}
__builtin_prefetch(&core_exec_states_[i], 1, 3);
ready[n_ready++] = {i, physical_core_id, regs[physical_core_id], hank->core_type};
core_serviced[i] = true;
remaining--;
}
}

// Open the window; the IDLE write releases the core's DATA_MAIN_BASE poll.
uint64_t reg_addr = regs[physical_core_id];
platform_init_aicore_regs(reg_addr);
OUT_OF_ORDER_STORE_BARRIER();
// Phase 2: publish every task pointer, then ONE barrier. The core reads task
// only after its window opens (Phase 3); a single barrier orders all task
// stores before any window STR. Writing task now (after the report) also
// keeps the core's CACHELINE_OUT report flush from clobbering it.
for (int32_t r = 0; r < n_ready; r++) {
all_handshakes[ready[r].i].task = reinterpret_cast<uint64_t>(&payload_per_core_[ready[r].i][0]);
}
OUT_OF_ORDER_STORE_BARRIER();

core_exec_states_[i].reg_addr = reg_addr;
core_exec_states_[i].cond_ptr = get_reg_ptr(reg_addr, RegId::COND);
#if SIMPLER_DFX
physical_core_ids_[i] = physical_core_id;
#endif
// Phase 3: open every window. platform_init_aicore_regs' STRs are posted
// Device-nGnRE writes, issued back-to-back with no interleaved GM stores.
for (int32_t r = 0; r < n_ready; r++) {
platform_init_aicore_regs(ready[r].reg_addr);
}

// Phase 4: publish each CoreExecState with a single (prefetched) struct store.
// core_exec_states_ is AICPU-private (the scheduler reads it, never the core),
// so it may be written after the windows open.
for (int32_t r = 0; r < n_ready; r++) {
int32_t i = ready[r].i;
CoreExecState st{};
st.reg_addr = ready[r].reg_addr;
st.cond_ptr = get_reg_ptr(ready[r].reg_addr, RegId::COND);
st.running_reg_task_id = AICPU_TASK_INVALID;
st.pending_reg_task_id = AICPU_TASK_INVALID;
#if !SIMPLER_DFX
core_exec_states_[i].worker_id = i;
core_exec_states_[i].physical_core_id = physical_core_id;
core_exec_states_[i].core_type = hank->core_type;
st.worker_id = i;
st.physical_core_id = ready[r].pcid;
st.core_type = ready[r].core_type;
#endif
core_exec_states_[i] = st;
core_type_compact_[i] = static_cast<uint8_t>(ready[r].core_type);
#if SIMPLER_DFX
physical_core_ids_[i] = ready[r].pcid;
#endif
core_serviced[i] = true;
remaining--;
}
}
OUT_OF_ORDER_STORE_BARRIER();
}
Expand Down Expand Up @@ -775,10 +806,8 @@ bool SchedulerContext::assign_cores_to_threads() {
active_sched_threads_, aic_count_, aiv_count_
);

for (int32_t i = 0; i < RUNTIME_MAX_WORKER; i++) {
core_exec_states_[i].running_reg_task_id = AICPU_TASK_INVALID;
core_exec_states_[i].pending_reg_task_id = AICPU_TASK_INVALID;
}
// running_reg_task_id / pending_reg_task_id for every serviced core are reset
// in handshake_partition's sweep.

// Count clusters per thread first (round-robin may distribute unevenly)
int32_t clusters_per_thread[MAX_AICPU_THREADS] = {};
Expand Down Expand Up @@ -908,26 +937,26 @@ int32_t SchedulerContext::post_handshake_init(Runtime *runtime) {
return -1;
}

// Build cluster-ordered AIC/AIV worker-id lists from the discovered cores.
// Serial and MMIO-free — the expensive per-core handshake already ran in
// parallel across threads. Core-index order matches the original
// single-thread handshake so assign_cores_to_threads forms identical
// clusters. core_type / physical_core_id are read from the Handshake struct
// (stable after the handshake) so this stays correct under SIMPLER_DFX,
// where they are not mirrored into CoreExecState.
Handshake *all_handshakes = reinterpret_cast<Handshake *>(runtime->dev.workers);
// Build the AIC/AIV worker-id lists in core-index order, which
// assign_cores_to_threads pairs into clusters. core_type is read from the
// contiguously packed core_type_compact_ the sweep filled, not the 64B-aligned
// per-core Handshake struct. aic_worker_ids_/aiv_worker_ids_ store through to
// HBM, so the lists are built in local (cached) buffers and published with two
// wide memcpys rather than element by element.
int32_t local_aic[RUNTIME_MAX_WORKER];
int32_t local_aiv[RUNTIME_MAX_WORKER];
int32_t la = 0, lv = 0;
for (int32_t i = 0; i < cores_total_num_; i++) {
CoreType type = all_handshakes[i].core_type;
uint32_t physical_core_id = all_handshakes[i].physical_core_id;
uint64_t reg_addr = core_exec_states_[i].reg_addr;
if (type == CoreType::AIC) {
aic_worker_ids_[aic_count_++] = i;
LOG_INFO_V0("Core %d: AIC, physical_id=%u, reg_addr=0x%lx", i, physical_core_id, reg_addr);
if (static_cast<CoreType>(core_type_compact_[i]) == CoreType::AIC) {
local_aic[la++] = i;
} else {
aiv_worker_ids_[aiv_count_++] = i;
LOG_INFO_V0("Core %d: AIV, physical_id=%u, reg_addr=0x%lx", i, physical_core_id, reg_addr);
local_aiv[lv++] = i;
}
}
memcpy(aic_worker_ids_, local_aic, static_cast<size_t>(la) * sizeof(int32_t));
memcpy(aiv_worker_ids_, local_aiv, static_cast<size_t>(lv) * sizeof(int32_t));
aic_count_ = la;
aiv_count_ = lv;
LOG_INFO_V0("Core discovery complete: %d AIC, %d AIV", aic_count_, aiv_count_);

if (!assign_cores_to_threads()) {
Expand Down Expand Up @@ -982,9 +1011,12 @@ int32_t SchedulerContext::post_handshake_init(Runtime *runtime) {
// Device orchestration: the orchestrator thread flips this when the graph is built.
orchestrator_done_.store(false, std::memory_order_release);

// Clear per-core dispatch payloads
memset(payload_per_core_, 0, sizeof(payload_per_core_));
memset(deferred_slab_per_core_, 0, sizeof(deferred_slab_per_core_));
// prepare_subtask_to_core fully writes a per-core payload / deferred-slab slot
// before the AICore is told to read it: build_payload sets
// function_bin_addr/args/local_context/not_ready, and deferred_slab->count/
// error_code are reset inline on every dispatch. An AICore reads a slot only
// after a dispatch targets it (DATA_MAIN_BASE), so a prior round's bytes in an
// untouched slot are never observed.

// Initialize per-core GlobalContext (sub_block_id) based on cluster position.
// This is done once at startup and never modified afterwards.
Expand Down
6 changes: 6 additions & 0 deletions src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,12 @@ class SchedulerContext {
int32_t aic_count_{0};
int32_t aiv_count_{0};

// Compact per-core CoreType, packed contiguously (~2 cache lines total) so
// post_handshake_init's ordered discovery scan reads it instead of taking a
// per-core volatile GM load from the 64B-aligned Handshake struct. Filled by
// each handshake thread for its own [lo,hi) slice during the parallel sweep.
uint8_t core_type_compact_[RUNTIME_MAX_WORKER]{};

// Set by any thread whose slice hits an invalid physical_core_id in
// handshake_partition; checked by the leader in post_handshake_init.
std::atomic<bool> handshake_failed_{false};
Expand Down
13 changes: 13 additions & 0 deletions src/a5/platform/onboard/host/device_runner.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -363,6 +363,19 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) {
// removes the slow launch itself, so the wedge cannot return if the timeout
// is ever tightened or a slower device pushes the launch past it. See
// docs/investigations/2026-06-pa-unroll-207001-optimeout-window.md.
// The AICore publishes aicore_done on launch (gated by nothing), and the
// workers region persists across runs in the pooled arena. Clearing each
// worker's aicore_done before the AICore kernel launches keeps the AICPU's
// handshake sweep from reading a prior run's report — which would open a
// window on that run's physical_core_id. Only aicore_done needs clearing; the
// AICore overwrites physical_core_id/core_type in the same report.
{
Handshake *workers = runtime.get_workers();
for (int i = 0; i < num_aicore; i++) {
workers[i].aicore_done = 0;
}
}

LOG_INFO_V0("=== launch_aicore_kernel ===");
rc = kernel_args_.init_device_kernel_args(mem_alloc_);
if (rc != 0) {
Expand Down
24 changes: 12 additions & 12 deletions src/a5/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,37 +65,37 @@ __aicore__ __attribute__((always_inline)) static void execute_task(__gm__ PTO2Di
__aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, int s_block_idx, CoreType core_type) {
__gm__ Handshake *my_hank = (__gm__ Handshake *)(&runtime->dev.workers[s_block_idx]);

// Phase 1: Wait for AICPU initialization signal
while (my_hank->aicpu_ready == 0) {
dcci(my_hank, SINGLE_CACHE_LINE);
SPIN_WAIT_HINT();
}

// Phase 2: Report physical core ID + core type and signal done in one write.
// The AICPU opens this core's register window (platform_init_aicore_regs:
// FAST_PATH + DATA_MAIN_BASE=IDLE) only after it observes aicore_done, so a
// single report suffices — there is no separate aicpu_regs_ready round-trip.
// Phase 1: report physical core ID + core type and signal done in one write,
// with no wait for the AICPU — both fields are self-known. The AICPU opens
// this core's register window only after it observes aicore_done, so a single
// report suffices. The host clears aicore_done before this kernel launches,
// so the value the AICPU reads is this run's report, never a stale prior one.
my_hank->physical_core_id = get_physical_core_id();
my_hank->core_type = core_type;
OUT_OF_ORDER_STORE_BARRIER();
my_hank->aicore_done = s_block_idx + 1; // Signal ready (use s_block_idx + 1 to avoid 0)
dcci(my_hank, SINGLE_CACHE_LINE, CACHELINE_OUT);

// Phase 3: Wait for the AICPU to open our register window. A kernel launch
// Phase 2: Wait for the AICPU to open our register window. A kernel launch
// resets DATA_MAIN_BASE to 0 (verified on a2a3 silicon; a5 shares this
// register protocol and relies on CI); the AICPU writes DATA_MAIN_BASE =
// AICPU_IDLE_TASK_ID (non-zero) as it opens FAST_PATH, so a non-zero read
// means the window is open and reads/writes are valid. The AICPU runs
// assign_cores_to_threads (µs) between opening the window and the first
// dispatch, so this IDLE is observed long before any task_id lands — the
// poll cannot miss it and mistake a later task for the reset value.
// Window-open is the sync point for everything the AICPU publishes (task
// pointer, swimlane head): the AICPU writes those before opening the window.
while (read_reg(RegId::DATA_MAIN_BASE) == 0) {
SPIN_WAIT_HINT();
}
// Report initial idle status via register (FAST_PATH is now open).
write_reg(RegId::COND, AICORE_IDLE_VALUE);

// Cache per-core dispatch payload pointer (set by AICPU before aicpu_ready)
// The AICPU writes task after observing our report (so our CACHELINE_OUT flush
// above cannot clobber it) and before opening the window; dcci to read its
// fresh value here.
dcci(my_hank, SINGLE_CACHE_LINE);
__gm__ PTO2DispatchPayload *payload = reinterpret_cast<__gm__ PTO2DispatchPayload *>(my_hank->task);

// Cache profiling state once after Phase 3. The L2 / PMU rings and the
Expand Down
Loading
Loading