Skip to content

Refactor: merge the hbg task table into one ChipTaskStorage array #2073

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
poursoul merged 1 commit into from
Aug 31, 2026
Merged

Refactor: merge the hbg task table into one ChipTaskStorage array #2073

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
1 change: 0 additions & 1 deletion 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 @@ -307,7 +307,6 @@ int32_t AicpuExecutor::run(Runtime *runtime) {

if (boot_ok) {
runtime_bind_ops(rt);
runtime->set_slot_states_ptr(nullptr);

sched_ctx_.bind_runtime(rt);

Expand Down
20 changes: 10 additions & 10 deletions src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -507,8 +507,7 @@ bool bind_graph_definitions(

for (size_t index = 0; index < count; ++index) {
std::optional<GraphHostUpload> upload = graph_host_upload(graph_state, index);
if (!upload.has_value() || upload->outer_slot == nullptr || upload->outer_slot->task_kind != TaskKind::GRAPH ||
upload->outer_slot->task == nullptr || upload->outer_slot->payload == nullptr) {
if (!upload.has_value() || upload->outer_slot == nullptr || upload->outer_slot->task_kind != TaskKind::GRAPH) {
LOG_ERROR("host-orch: invalid pending Graph task");
return false;
}
Expand All @@ -534,21 +533,22 @@ bool bind_graph_definitions(
definition->scalar_arg_count, &storage_layout
) ||
storage_layout.total_bytes != definition->execution_storage_bytes ||
upload->outer_slot->payload->tensor_count != static_cast<int32_t>(definition->boundary_count) ||
upload->outer_slot->payload->scalar_count != static_cast<int32_t>(definition->boundary_scalar_count)) {
upload->outer_slot->to_payload().tensor_count != static_cast<int32_t>(definition->boundary_count) ||
upload->outer_slot->to_payload().scalar_count != static_cast<int32_t>(definition->boundary_scalar_count)) {
LOG_ERROR("host-orch: invalid Graph Definition for task");
return false;
}
const uintptr_t outer_base = reinterpret_cast<uintptr_t>(upload->outer_slot->task->packed_buffer_base);
const uintptr_t outer_end = reinterpret_cast<uintptr_t>(upload->outer_slot->task->packed_buffer_end);
const uintptr_t outer_base =
reinterpret_cast<uintptr_t>(upload->outer_slot->to_descriptor().packed_buffer_base);
const uintptr_t outer_end = reinterpret_cast<uintptr_t>(upload->outer_slot->to_descriptor().packed_buffer_end);
if (outer_end < outer_base || definition->required_heap > UINTPTR_MAX - outer_base ||
storage_layout.total_bytes > outer_end - outer_base ||
definition->required_heap > outer_end - outer_base - storage_layout.total_bytes) {
LOG_ERROR("host-orch: Graph runtime storage does not fit its outer task heap");
return false;
}
const uintptr_t storage_addr = outer_base + definition->required_heap;
if (storage_addr % alignof(InGraphTaskStorage) != 0) {
if (storage_addr % alignof(ChipTaskStorage) != 0) {
LOG_ERROR("host-orch: Graph runtime storage address is misaligned");
return false;
}
Expand Down Expand Up @@ -616,7 +616,7 @@ int32_t run_host_orchestration(
LOG_ERROR("host-orch: host SM mirror of %" PRIu64 " bytes unavailable", sm_size);
return PTO_RUNTIME_ERR_INTERNAL;
}
std::memset(host_sm, 0, sm_segs.descriptors);
std::memset(host_sm, 0, sm_segs.storage);

// Re-point the orchestrator half at the host SM (scheduler keeps device SM).
// Host-owned and destroyed with this frame, so rt->orchestrator is dropped on
Expand Down Expand Up @@ -909,15 +909,15 @@ int32_t run_host_orchestration(
reinterpret_cast<uint64_t>(gm_heap) < HEAP_VIRTUAL_BASE && "device memory reaches into the virtual heap window"
);
// The alignment bind_graph_definitions checked on the virtual base — a Graph
// task's runtime storage must land on alignof(InGraphTaskStorage) — carries to
// task's runtime storage must land on alignof(ChipTaskStorage) — carries to
// the real base only while the two are congruent: both are aligned to
// kDefaultBaseAlign, and that covers the storage's own requirement.
static_assert(
HEAP_VIRTUAL_BASE % DeviceArena::kDefaultBaseAlign == 0,
"the virtual heap base must share the committed region's alignment"
);
static_assert(
alignof(InGraphTaskStorage) <= DeviceArena::kDefaultBaseAlign,
alignof(ChipTaskStorage) <= DeviceArena::kDefaultBaseAlign,
"an in-graph task's storage alignment must be covered by the heap region's base alignment"
);
always_assert(reinterpret_cast<uint64_t>(gm_heap) % DeviceArena::kDefaultBaseAlign == 0);
Expand Down
34 changes: 18 additions & 16 deletions src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -559,7 +559,7 @@ struct SchedulerState {
} else {
ResourceShape shape = slot_state->active_mask.to_shape();
if (shape == ResourceShape::DUMMY ||
(slot_state->task_attrs.has_predicate() && !slot_state->payload->predicate.pass())) {
(slot_state->task_attrs.has_predicate() && !slot_state->to_payload().predicate.pass())) {
pushed = dummy_ready_queue.push(slot_state);
} else if (slot_state->task_attrs.requires_sync_start()) {
pushed = ready_sync_queues[static_cast<int32_t>(shape)].push(slot_state);
Expand Down Expand Up @@ -594,7 +594,7 @@ struct SchedulerState {
// lists. The decision is terminal: tasks are never re-polled; a producer's
// completion re-scans its waiters via on_mixed_task_complete's wake drain.
int classify_fanin_state(const ChipTaskSlotState *s) const {
const TaskPayload &p = *s->payload;
const TaskPayload &p = s->to_payload();
const SharedMemoryTaskHeader &tasks = *task_view.tasks;
const int32_t *fanin = p.fanin_data();
for (int32_t i = p.fanin_count - 1; i >= 0; i--) {
Expand Down Expand Up @@ -624,7 +624,7 @@ struct SchedulerState {
push_ready_routed(consumer);
return;
}
producer = &tasks.get_slot_state_by_task_id(consumer->payload->fanin_data()[state]);
producer = &tasks.get_slot_state_by_task_id(consumer->to_payload().fanin_data()[state]);
}
}

Expand All @@ -635,7 +635,7 @@ struct SchedulerState {
// watermark >= producer.last_consumer_local_id). Whole-graph-resident hbg
// has no device slot reclaim, so nothing advances a reclaim cursor here.
void on_mixed_task_complete(ChipTaskSlotState &slot_state) {
const int32_t task_id = static_cast<int32_t>(simpler::hbg::task_local_id(slot_state.task->task_id));
const int32_t task_id = static_cast<int32_t>(simpler::hbg::task_local_id(slot_state.to_descriptor().task_id));
SharedMemoryTaskHeader &tasks = *task_view.tasks;

slot_state.mark_completed(); // host-visible mirror (task_state = COMPLETED)
Expand All @@ -644,7 +644,7 @@ struct SchedulerState {
ChipTaskSlotState *waiter = slot_state.wake_list_head.exchange(WAKE_LIST_SENTINEL, std::memory_order_acq_rel);
while (waiter != nullptr && waiter != WAKE_LIST_SENTINEL) {
ChipTaskSlotState *next = waiter->next_in_wake_list;
if (waiter->payload->fanin_count == 1) {
if (waiter->to_payload().fanin_count == 1) {
push_ready_routed(waiter); // single-fanin waiter was waiting only on us
waiter = next;
continue;
Expand All @@ -653,7 +653,7 @@ struct SchedulerState {
if (state < 0) {
push_ready_routed(waiter);
} else {
register_wake(&tasks.get_slot_state_by_task_id(waiter->payload->fanin_data()[state]), waiter);
register_wake(&tasks.get_slot_state_by_task_id(waiter->to_payload().fanin_data()[state]), waiter);
}
waiter = next;
}
Expand Down Expand Up @@ -775,15 +775,15 @@ struct SchedulerState {

inline void record_published_blocks(ChipTaskSlotState &slot_state, int32_t count) {
if (count <= 0 || !slot_state.task_attrs.allow_early_resolve()) return;
slot_state.payload->published_block_count.fetch_add(static_cast<int16_t>(count), std::memory_order_seq_cst);
slot_state.to_payload().published_block_count.fetch_add(static_cast<int16_t>(count), std::memory_order_seq_cst);
}

// Ring one sync_start cohort from its stable staged_core_mask. The caller owns
// the NONE->RINGING launch latch and invokes this exactly once after local or
// global staging completes, while the corresponding per-core table entries are live.
inline void ring_all_staged_doorbells(ChipTaskSlotState &slot_state) {
for (int w = 0; w < EARLY_DISPATCH_CORE_MASK_WORDS; w++) {
uint64_t bits = slot_state.payload->staged_core_mask[w].load(std::memory_order_seq_cst);
uint64_t bits = slot_state.to_payload().staged_core_mask[w].load(std::memory_order_seq_cst);
while (bits != 0) {
int core_id = w * 64 + __builtin_ctzll(bits);
bits &= bits - 1;
Expand Down Expand Up @@ -816,14 +816,16 @@ struct SchedulerState {

inline void cancel_early_sync_drain(ChipTaskSlotState &slot_state) {
uint8_t previous =
slot_state.payload->early_sync_drain_state.exchange(EARLY_SYNC_DRAIN_NONE, std::memory_order_seq_cst);
slot_state.to_payload().early_sync_drain_state.exchange(EARLY_SYNC_DRAIN_NONE, std::memory_order_seq_cst);
if ((previous & EARLY_SYNC_DRAIN_OWNER) == 0) return;
if ((previous & EARLY_SYNC_DRAIN_READY) != 0) {
push_ready_routed(&slot_state);
return;
}
if (slot_state.payload->early_dispatch_state.load(std::memory_order_seq_cst) == EARLY_DISPATCH_STAGING) {
early_sync_start_queue.push_tagged(&slot_state, static_cast<uint64_t>(slot_state.task->task_id.raw));
if (slot_state.to_payload().early_dispatch_state.load(std::memory_order_seq_cst) == EARLY_DISPATCH_STAGING) {
early_sync_start_queue.push_tagged(
&slot_state, static_cast<uint64_t>(slot_state.to_descriptor().task_id.raw)
);
}
}

Expand Down Expand Up @@ -852,19 +854,19 @@ struct SchedulerState {
// happens-before its final store. Read the seed first, then the mask, so
// observing the final count cannot be paired with a partially published
// mask when producer release races staging.
int32_t running_cores = slot_state.payload->running_slot_count.load(std::memory_order_seq_cst);
int32_t running_cores = slot_state.to_payload().running_slot_count.load(std::memory_order_seq_cst);
int32_t staged_cores = 0;
for (int w = 0; w < EARLY_DISPATCH_CORE_MASK_WORDS; w++)
staged_cores +=
__builtin_popcountll(slot_state.payload->staged_core_mask[w].load(std::memory_order_seq_cst));
__builtin_popcountll(slot_state.to_payload().staged_core_mask[w].load(std::memory_order_seq_cst));
if (staged_cores == 0) return false;
if (running_cores != staged_cores) return false;
if (slot_state.payload->early_dispatch_state.load(std::memory_order_seq_cst) != EARLY_DISPATCH_DISPATCHED)
if (slot_state.to_payload().early_dispatch_state.load(std::memory_order_seq_cst) != EARLY_DISPATCH_DISPATCHED)
return false;
if (!try_claim_early_dispatch_launch(*slot_state.payload)) return false;
if (!try_claim_early_dispatch_launch(slot_state.to_payload())) return false;
ring_all_staged_doorbells(slot_state);
wmb();
slot_state.payload->early_dispatch_launch_state.store(
slot_state.to_payload().early_dispatch_launch_state.store(
EARLY_DISPATCH_LAUNCH_COMPLETE, std::memory_order_release
);
return true;
Expand Down
46 changes: 25 additions & 21 deletions src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,8 +149,8 @@ void format_core_status(
int64_t task_id_raw = -1;
if (core_state && core_state->running_slot_state) {
int32_t subslot = static_cast<int32_t>(core_state->running_subslot);
kernel = core_state->running_slot_state->task->kernel_id[subslot];
task_id_raw = static_cast<int64_t>(core_state->running_slot_state->task->task_id.raw);
kernel = core_state->running_slot_state->to_descriptor().kernel_id[subslot];
task_id_raw = static_cast<int64_t>(core_state->running_slot_state->to_descriptor().task_id.raw);
}
uint64_t cond_reg = read_reg(reg_addr_for_cond, RegId::COND);
int32_t hw_state = EXTRACT_TASK_STATE(cond_reg);
Expand Down Expand Up @@ -221,18 +221,18 @@ void SchedulerContext::log_stall_diagnostics(
// Polling: no fanin_refcount. Recompute met/total from the inline
// fanin ids vs the completion_flags (rc = satisfied producers,
// fi = raw producer count) so the stall dump still shows readiness.
int32_t fi = slot_state.payload != nullptr ? slot_state.payload->fanin_count : 0;
int32_t fi = slot_state.to_payload().fanin_count;
int32_t rc = 0;
if (slot_state.payload != nullptr) {
const int32_t *fanin = slot_state.payload->fanin_data();
{
const int32_t *fanin = slot_state.to_payload().fanin_data();
for (int32_t k = 0; k < fi; k++) {
if (tasks.is_completion_flag_set(fanin[k], std::memory_order_relaxed)) rc++;
}
}
int32_t kid_aic = slot_state.task->kernel_id[0];
int32_t kid_aiv0 = slot_state.task->kernel_id[1];
int32_t kid_aiv1 = slot_state.task->kernel_id[2];
int64_t task_id = static_cast<int64_t>(slot_state.task->task_id.raw);
int32_t kid_aic = slot_state.to_descriptor().kernel_id[0];
int32_t kid_aiv0 = slot_state.to_descriptor().kernel_id[1];
int32_t kid_aiv1 = slot_state.to_descriptor().kernel_id[2];
int64_t task_id = static_cast<int64_t>(slot_state.to_descriptor().task_id.raw);
if (st >= CHIP_TASK_COMPLETED) continue;
// task_state has no intermediate ready/running value — it
// stays PENDING until the worker stores COMPLETED. Classify
Expand Down Expand Up @@ -364,19 +364,23 @@ int32_t SchedulerContext::handle_timeout_exit(
// Capture the in-flight kernels' partial output before signalling the
// cores to exit, so the dump reflects the live stuck state.
if (is_dump_args_enabled()) {
dump_running_task_outputs<SUBTASK_SLOT_COUNT>(
thread_idx, cores_total_num_,
dump_running_task_outputs(
cores_total_num_,
[this](int32_t cid) {
return core_exec_states_[cid].running_slot_state;
},
[](ActiveMask active_mask, int raw_subtask_id) {
return active_mask.subtask_active(static_cast<SubtaskSlot>(raw_subtask_id));
},
[this](int32_t func_id) {
return get_function_bin_addr(func_id);
},
[](const ChipTaskSlotState &slot_state) {
return &slot_state.payload->dump_metadata;
[this, thread_idx](const ChipTaskSlotState &slot_state) {
dump_args_for_task<SUBTASK_SLOT_COUNT>(
thread_idx, slot_state.to_descriptor(), slot_state.to_payload(), slot_state.active_mask,
ArgsDumpStage::AFTER_COMPLETION,
[](ActiveMask active_mask, int raw_subtask_id) {
return active_mask.subtask_active(static_cast<SubtaskSlot>(raw_subtask_id));
},
[this](int32_t func_id) {
return get_function_bin_addr(func_id);
},
&slot_state.to_payload().dump_metadata
);
}
);
}
Expand Down Expand Up @@ -1086,13 +1090,13 @@ void SchedulerContext::classify_partition(int32_t thread_idx, int32_t nthreads)
ChipTaskSlotState &slot = tasks.get_slot_state_by_task_id(id);
if (slot.task_kind == TaskKind::GRAPH) {
if (graph_execution_localize(slot) == nullptr) slot.graph_context = nullptr;
if (!sched_->push_graph_prepare(&slot, slot.task->task_id.raw, thread_idx)) return;
if (!sched_->push_graph_prepare(&slot, slot.to_descriptor().task_id.raw, thread_idx)) return;
}
int32_t state = sched_->classify_fanin_state(&slot);
if (state < 0) {
sched_->push_ready_routed(&slot);
} else {
int32_t prod_local = slot.payload->fanin_data()[state];
int32_t prod_local = slot.to_payload().fanin_data()[state];
sched_->register_wake(&tasks.get_slot_state_by_task_id(prod_local), &slot);
}
}
Expand Down
Loading
Loading