Skip to content

Sync: host_build_graph scheduler up a generation to tensormap_and_ringbuffer #1327

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 11, 2026
Merged

Sync: host_build_graph scheduler up a generation to tensormap_and_ringbuffer #1327

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
64 changes: 35 additions & 29 deletions src/a2a3/runtime/host_build_graph/aicore/aicore_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,15 +47,17 @@ __aicore__ __attribute__((always_inline)) static void execute_task(__gm__ PTO2Di
* AICore main execution loop
*
* Implements the AICPU-AICore register-based dispatch protocol:
* 1. Wait for AICPU ready signal via handshake buffer
* 2. Report physical core ID and core type, signal AICore ready
* 1. Report physical core ID and core type, signal aicore_done (no AICPU wait)
* 2. Wait for the AICPU to open our register window (DATA_MAIN_BASE != 0)
* 3. Cache per-core PTO2DispatchPayload pointer from hank->task
* 4. Poll DATA_MAIN_BASE register for task dispatch until exit signal
*
* AICPU writes &s_payload_per_core[i] to hank->task before setting
* aicpu_ready=1. AICore caches this pointer and reads function_bin_addr +
* args pointer from it on each dispatch. reg_val is a monotonically
* increasing task ID used only for dispatch signaling and ACK/FIN protocol.
* AICore reports on launch; the AICPU writes &s_payload_per_core[i] to
* hank->task and then opens the register window (DATA_MAIN_BASE = IDLE), which
* is itself the acknowledgement. AICore caches this pointer and reads
* function_bin_addr + args pointer from it on each dispatch. reg_val is a
* monotonically increasing task ID used only for dispatch signaling and
* ACK/FIN protocol.
*
* Profiling state (enable flag, L2 swimlane rotation channel) is published into the platform
* via set_aicore_profiling_flag / set_aicore_l2_swimlane_ring at kernel entry —
Expand All @@ -69,32 +71,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->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, signal ready
// 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();
OUT_OF_ORDER_STORE_BARRIER();
my_hank->aicore_regs_ready = 1;
dcci(&my_hank->aicore_regs_ready, SINGLE_CACHE_LINE, CACHELINE_OUT);
while (my_hank->aicpu_regs_ready == 0) {
dcci(&my_hank->aicpu_regs_ready, SINGLE_CACHE_LINE);
SPIN_WAIT_HINT();
}
// Report initial idle status via register
write_reg(RegId::COND, AICORE_IDLE_VALUE);

// Phase 3: Report core type, signal ready
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);

// Cache per-core dispatch payload pointer (set by AICPU before aicpu_ready)
// 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);

// 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 All @@ -103,9 +109,9 @@ __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, in
bool pmu_enabled = SIMPLER_GET_DFX_FLAG(enable_profiling_flag, SIMPLER_DFX_FLAG_PMU);

// Per-core L2SwimlaneActiveHead channel. AICPU completes
// `l2_swimlane_aicpu_init` before writing `aicpu_ready = 1` in
// `handshake_all_cores`, and Phase 1 above has already observed
// `aicpu_ready == 1`, so the rotation-table slot is populated and the
// `l2_swimlane_aicpu_init` (in pre_handshake_init) before any thread opens a
// register window in `handshake_partition`, and Phase 2 above has already
// observed our window open, so the rotation-table slot is populated and the
// first deref is safe here — off the dispatch→start critical path.
__gm__ L2SwimlaneActiveHead *l2_swimlane_head = l2_swimlane_enabled ? get_l2_swimlane_aicore_head() : nullptr;
// cached_buf_seq must start != AICPU's initial head.current_buf_seq (0)
Expand Down
142 changes: 100 additions & 42 deletions src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,11 +84,20 @@ struct AicpuExecutor {

// ===== Thread management state =====
std::atomic<int32_t> thread_idx_{0};
std::atomic<bool> initialized_{false};
std::atomic<bool> init_done_{false};
std::atomic<bool> init_failed_{false};
std::atomic<bool> finished_{false};

// Parallel-handshake coordination (see AicpuExecutor::init). hs_setup_done_
// is published by the leader once the shared pre-handshake setup is visible;
// hs_arrived_ is the barrier counting threads that finished their core slice.
// hs_thread_seq_ hands out a distinct [0, nthreads) index when the platform
// exposes no affinity idx (sim, where platform_aicpu_affinity_thread_idx()
// is -1 during init) so the threads don't all collapse to leader 0.
std::atomic<bool> hs_setup_done_{false};
std::atomic<int32_t> hs_arrived_{0};
std::atomic<int32_t> hs_thread_seq_{0};

int32_t aicpu_thread_num_{0};

// ===== Task queue state (managed by scheduler ready queues) =====
Expand All @@ -115,42 +124,88 @@ static AicpuExecutor g_aicpu_executor;
// ===== AicpuExecutor Method Implementations =====

int32_t AicpuExecutor::init(Runtime *runtime) {
bool expected = false;
if (!initialized_.compare_exchange_strong(expected, true, std::memory_order_acq_rel, std::memory_order_acquire)) {
return 0;
}

LOG_INFO_V0("AicpuExecutor: Initializing");

if (runtime == nullptr) {
LOG_ERROR("runtime is nullptr");
init_failed_.store(true, std::memory_order_release);
return -1;
}

// Read execution parameters from runtime. The 0 → 1 fixup runs before the
// sched_thread_num_ derivation so a zero input doesn't leave the scheduler
// count at -1.
aicpu_thread_num_ = runtime->aicpu_thread_num;
if (aicpu_thread_num_ == 0) aicpu_thread_num_ = 1;
sched_thread_num_ = aicpu_thread_num_ - 1;
orch_to_sched_ = runtime->orch_to_sched;

if (aicpu_thread_num_ < 1 || aicpu_thread_num_ > MAX_AICPU_THREADS) {
LOG_ERROR("Invalid aicpu_thread_num: %d", aicpu_thread_num_);
// All AICPU threads enter init. The per-core AICore handshake is the
// dominant preamble cost (serial MMIO, ~217 µs of ~283 µs for 72 cores), so
// it is parallelized: the leader (tidx 0) does the shared setup, every
// thread handshakes a disjoint slice of cores, then the leader finishes init
// after a barrier. Non-leaders spin on init_done_.
int32_t nthreads = runtime->aicpu_thread_num;
if (nthreads == 0) nthreads = 1;
if (nthreads < 1 || nthreads > MAX_AICPU_THREADS) {
LOG_ERROR("Invalid aicpu_thread_num: %d", nthreads);
init_failed_.store(true, std::memory_order_release);
return -1;
}

if (sched_ctx_.init(runtime, aicpu_thread_num_, sched_thread_num_, orch_to_sched_, get_platform_regs()) != 0) {
init_failed_.store(true, std::memory_order_release);
// Each thread needs a distinct index in [0, nthreads) to pick the leader and
// partition the cores. Onboard the gate filter assigns it (exec_idx); sim's
// gate does not, so platform_aicpu_affinity_thread_idx() is -1 here for every
// thread — hand those a distinct index from a counter (mirrors run()'s
// thread_idx_++ fallback) instead of collapsing them all to leader 0, which
// would run pre_/post_handshake_init on every thread and race the shared
// scheduler state. Exactly nthreads threads reach init (the gate drops the
// rest), so the counter yields a gap-free [0, nthreads).
int32_t tidx = platform_aicpu_affinity_thread_idx();
if (tidx < 0) tidx = hs_thread_seq_.fetch_add(1, std::memory_order_acq_rel);
// A thread whose index still falls outside [0, nthreads) owns no core slice:
// handshake_partition would compute lo/hi past cores_total_num_ and index
// all_handshakes[]/core_exec_states_ out of bounds. Reject it here (mirrors
// the bounds guard already in run()). Fail only this thread and do NOT set
// init_failed_ — that would make the valid peers abort before their
// hs_arrived_ increment and hang the leader at the barrier below.
if (tidx >= nthreads) {
LOG_ERROR("AICPU affinity thread idx %d out of range [0,%d) in init", tidx, nthreads);
return -1;
}
Comment thread
ChaoWao marked this conversation as resolved.
const bool is_leader = (tidx == 0);

if (is_leader) {
LOG_INFO_V0("AicpuExecutor: Initializing");
// The 0 → 1 fixup already applied above; derive scheduler count from it.
aicpu_thread_num_ = nthreads;
sched_thread_num_ = nthreads - 1;
orch_to_sched_ = runtime->orch_to_sched;

hs_arrived_.store(0, std::memory_order_relaxed);
if (sched_ctx_.pre_handshake_init(runtime, aicpu_thread_num_, sched_thread_num_, get_platform_regs()) != 0) {
init_failed_.store(true, std::memory_order_release);
hs_setup_done_.store(true, std::memory_order_release);
return -1;
}
hs_setup_done_.store(true, std::memory_order_release);
} else {
while (!hs_setup_done_.load(std::memory_order_acquire)) {
if (init_failed_.load(std::memory_order_acquire)) return -1;
}
if (init_failed_.load(std::memory_order_acquire)) return -1;
}

finished_count_.store(0, std::memory_order_release);

init_done_.store(true, std::memory_order_release);
LOG_INFO_V0("AicpuExecutor: Init complete");
// All threads: handshake this thread's slice of cores in parallel.
sched_ctx_.handshake_partition(runtime, tidx, nthreads);

// Barrier: leader waits for every slice to finish, then completes init.
hs_arrived_.fetch_add(1, std::memory_order_acq_rel);
if (is_leader) {
while (hs_arrived_.load(std::memory_order_acquire) < nthreads) {}
Comment thread
ChaoWao marked this conversation as resolved.
finished_count_.store(0, std::memory_order_release);
if (sched_ctx_.post_handshake_init(runtime) != 0) {
init_failed_.store(true, std::memory_order_release);
init_done_.store(true, std::memory_order_release);
return -1;
}
init_done_.store(true, std::memory_order_release);
LOG_INFO_V0("AicpuExecutor: Init complete");
} else {
while (!init_done_.load(std::memory_order_acquire)) {
if (init_failed_.load(std::memory_order_acquire)) return -1;
}
if (init_failed_.load(std::memory_order_acquire)) return -1;
}
return 0;
}

Expand Down Expand Up @@ -206,11 +261,12 @@ int32_t AicpuExecutor::run(Runtime *runtime) {

sched_ctx_.bind_runtime(rt);

// Publish orchestration completion (sets total_tasks_ + orchestrator_done_)
// BEFORE releasing the scheduler threads. Otherwise they unblock with
// total_tasks_=0 and can race to an early exit before the host task
// count is published (host-orch has no concurrent orchestrator to
// keep them alive).
// Latch the host-built task count (on_orchestration_done sets total_tasks_)
// BEFORE the runtime_init_ready_ release below — that store is the barrier
// that unblocks the scheduler threads. Otherwise they would acquire
// runtime_init_ready_ with total_tasks_=0 and race to an early exit before
// the host task count is visible (host-orch has no concurrent orchestrator
// to keep them alive).
// NOTE: do NOT call rt_orchestration_done(rt) here. The HOST already
// called it in run_host_orchestration; the orchestrator's own
// task-allocator pointers are intentionally NOT relocated (only the
Expand Down Expand Up @@ -297,9 +353,11 @@ void AicpuExecutor::deinit(Runtime *runtime) {

LOG_INFO_V0("DeInit: Runtime execution state reset");

initialized_.store(false, std::memory_order_release);
init_done_.store(false, std::memory_order_release);
init_failed_.store(false, std::memory_order_release);
hs_setup_done_.store(false, std::memory_order_release);
hs_arrived_.store(0, std::memory_order_release);
hs_thread_seq_.store(0, std::memory_order_release);
thread_idx_.store(0, std::memory_order_release);
finished_.store(false, std::memory_order_release);

Expand All @@ -324,10 +382,11 @@ extern "C" int32_t aicpu_prewarm_callable(Runtime *runtime) {
*
* This is called by DynTileFwkBackendKernelServer in kernel.cpp.
* Orchestrates the complete task runtime execution:
* 1. Initialize executor (thread-safe, first thread only)
* 2. Wait for initialization to complete
* 3. Execute tasks on managed cores
* 4. Cleanup when last thread finishes
* 1. Initialize executor: all threads enter init(), which handshakes the cores
* in parallel and barriers internally until init is complete (or a thread
* failed); its return value is authoritative on every thread.
* 2. Execute tasks on managed cores
* 3. Cleanup when last thread finishes
*
* @param runtime Pointer to Runtime structure
* @return 0 on success, non-zero on error
Expand All @@ -340,13 +399,12 @@ extern "C" int32_t aicpu_execute(Runtime *runtime) {

LOG_INFO_V0("%s", "aicpu_execute: Starting AICPU kernel execution");

g_aicpu_executor.init(runtime);

while (!g_aicpu_executor.init_done_.load(std::memory_order_acquire)) {
if (g_aicpu_executor.init_failed_.load(std::memory_order_acquire)) {
LOG_ERROR("%s", "aicpu_execute: Initialization failed, aborting execution");
return -1;
}
// init() barriers every thread internally until init is complete on the
// leader (or a thread failed), then returns the status — so a non-zero
// return is authoritative on all threads and no extra spin is needed.
if (g_aicpu_executor.init(runtime) != 0) {
LOG_ERROR("%s", "aicpu_execute: Initialization failed, aborting execution");
return -1;
}

int32_t rc = g_aicpu_executor.run(runtime);
Expand Down
Loading
Loading