Skip to content

a5 host_build_graph: converge scheduler_dispatch + runtime_maker with a2a3 #1789

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

a5 host_build_graph: converge scheduler_dispatch + runtime_maker with a2a3 #1789

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
89 changes: 41 additions & 48 deletions src/a5/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 @@ -55,8 +55,8 @@
#include "../runtime/common.h"
#include "../runtime/dep_gen_host_graph.h"
#include "../runtime/graph_execution.h"
#include "../runtime/graph_host_state.h"
#include "../runtime/host_tensor_access.h"
#include "../runtime/graph_host_state.h"
#include "../runtime/pto_orchestrator.h"
#include "../runtime/pto_runtime2.h"
#include "../runtime/pto_shared_memory.h"
Expand Down Expand Up @@ -378,7 +378,7 @@ struct HostOrchEntryPoints {
// (slot_state.task/.payload, fanin_inline_slot_states[], dep-entry.slot_state,
// ready-queue slot.slot_state) and the arena block (slot_state.fanout_head,
// dep-entry.next point into the SM but live in the arena).
// Rather than track which delta each field needs, reloc() classifies every
// Rather than track which delta each field needs, relocate() classifies every
// pointer by the region it points INTO and applies that region's delta; foreign
// and null pointers pass through untouched. The fanout adjacency is wired inline
Comment on lines +381 to 383

@coderabbitai coderabbitai Bot Aug 11, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Correct the relocation contract comment.

Lines 381-383 state that foreign pointers pass through unchanged. Lines 420-427 reject a non-null pointer outside both windows and make relocation fail. State that null pointers pass through unchanged and external pointers cause failure.

Proposed fix 
-// pointer by the region it points INTO and applies that region's delta; foreign
-// and null pointers pass through untouched. The fanout adjacency is wired inline
+// pointer by the region it points INTO and applies that region's delta. Null
+// pointers pass through unchanged. External pointers cause relocation failure.
+// The fanout adjacency is wired inline
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// Rather than track which delta each field needs, relocate() classifies every
// pointer by the region it points INTO and applies that region's delta; foreign
// and null pointers pass through untouched. The fanout adjacency is wired inline
// Rather than track which delta each field needs, relocate() classifies every
// pointer by the region it points INTO and applies that region's delta. Null
// pointers pass through unchanged. External pointers cause relocation failure.
// The fanout adjacency is wired inline
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/a5/runtime/host_build_graph/host/runtime_maker.cpp` around lines 381 -
383, Update the relocation contract comment near relocate() to state that null
pointers pass through unchanged, while non-null pointers outside both recognized
regions cause relocation to fail; remove the claim that foreign pointers pass
through untouched.

// during host submit, so dep_pool/ready are already populated here.
Expand All @@ -390,16 +390,12 @@ struct HostOrchEntryPoints {
// (returns false) rather than shipping un-relocated host pointers to the device.
// Returns false on any unrelocatable pointer so the caller can fail the prepare.
static bool relocate_host_orch_image(
PTO2SharedMemoryHandle &host_sm_handle, [[maybe_unused]] PTO2Runtime *rt, uint64_t host_sm, uint64_t sm_size,
int64_t sm_delta, uint64_t host_arena, uint64_t arena_size, int64_t arena_delta
PTO2SharedMemoryHandle &host_sm_handle, uint64_t host_sm, uint64_t sm_size, int64_t sm_delta, uint64_t host_arena,
uint64_t arena_size, int64_t arena_delta
) {
// host_build_graph is single-ring; the loops below iterate the lone ring and
// index header->ring (singular). If the ring depth ever grows, those loops
// would relocate the same ring N times (applying the delta repeatedly =
// corruption), so pin the assumption here.
static_assert(PTO2_MAX_RING_DEPTH == 1, "relocate_host_orch_image assumes a single ring");

// SM and arena windows must not overlap — reloc classifies a pointer by
// SM and arena windows must not overlap — relocate classifies a pointer by
// which window it falls in, so an overlap would misclassify and apply the
// wrong delta. Both are independent malloc-backed host buffers in practice;
// assert it so a future shared-buffer layout can't silently corrupt.
Expand All @@ -412,41 +408,39 @@ static bool relocate_host_orch_image(
}

bool ok = true;
auto reloc = [&](auto *&p) {
using Ptr = std::remove_reference_t<decltype(p)>;
uint64_t v = reinterpret_cast<uint64_t>(p);
if (v == 0) {
return;
}
if (v >= host_sm && v < host_sm + sm_size) {
p = reinterpret_cast<Ptr>(static_cast<uintptr_t>(v + sm_delta));
} else if (v >= host_arena && v < host_arena + arena_size) {
p = reinterpret_cast<Ptr>(static_cast<uintptr_t>(v + arena_delta));
auto relocate = [&](auto *&pointer) {
using Pointer = std::remove_reference_t<decltype(pointer)>;
const uint64_t address = reinterpret_cast<uint64_t>(pointer);
if (address == 0) return;
if (address >= host_sm && address < host_sm + sm_size) {
pointer = reinterpret_cast<Pointer>(static_cast<uintptr_t>(address + sm_delta));
} else if (address >= host_arena && address < host_arena + arena_size) {
pointer = reinterpret_cast<Pointer>(static_cast<uintptr_t>(address + arena_delta));
} else {
// A non-null pointer in neither window is an external/host address
// the device would dereference verbatim after H2D. No field should
// legitimately carry one; latch fatal rather than ship a host VA to
// the device (silent AICPU corruption otherwise).
LOG_ERROR("host-orch: pointer %#lx is outside both SM and arena windows; cannot relocate for device", v);
LOG_ERROR(
"host-orch: pointer %#lx is outside both SM and arena windows; cannot relocate for device", address
);
ok = false;
}
};

PTO2SharedMemoryHeader *header = host_sm_handle.header;
if (header != nullptr) {
for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) {
PTO2SharedMemoryRingHeader &ring = header->ring;
int32_t count = ring.fc.current_task_index.load(std::memory_order_acquire);
for (int32_t slot = 0; slot < count; slot++) {
PTO2TaskSlotState *ss = &ring.slot_states[slot];
// Polling: fanin is a flat array of position-independent local-id
// integers on the payload, so only the two per-slot arena/SM
// pointers need relocating. There is no fanout_head/dep_pool graph
// and no host-seeded ready queue (the device boot scan classifies),
// so those relocation passes are gone.
reloc(ss->task);
reloc(ss->payload);
}
PTO2SharedMemoryRingHeader &ring = header->ring;
const int32_t count = ring.fc.current_task_index.load(std::memory_order_acquire);
for (int32_t slot = 0; slot < count; ++slot) {
PTO2TaskSlotState *state = &ring.slot_states[slot];
// Polling: fanin is a flat array of position-independent local-id
// integers on the payload, so only the two per-slot arena/SM
// pointers need relocating. There is no fanout_head/dep_pool graph
// and no host-seeded ready queue (the device boot scan classifies),
// so those relocation passes are gone.
relocate(state->task);
relocate(state->payload);
}
}
return ok;
Expand Down Expand Up @@ -559,7 +553,7 @@ int32_t run_host_orchestration(

GraphHostStatePtr graph_state = make_graph_host_state();
if (!graph_state) {
LOG_ERROR("host-orch: failed to allocate Graph recording state");
LOG_ERROR("host-orch: failed to allocate Graph host state");
return -1;
}
GraphHostStateBinding graph_binding(rt->orchestrator, graph_state.get());
Expand All @@ -583,27 +577,26 @@ int32_t run_host_orchestration(
// context_lens/block_table) whether or not the platform maps device memory
// into the host address space.

const HostOrchEntryPoints *eps = reinterpret_cast<const HostOrchEntryPoints *>(host_orch_func_ptr);
rt->active_callable_hash = reinterpret_cast<uint64_t>(eps->entry);
if (eps->bind != nullptr) {
rt->tensor_access = &tensor_access;
// Binds the orchestration .so's own framework_current_runtime, which its
// inline rt_submit_* read. The host library links a same-named copy from
// orchestration/common.cpp, but nothing outside the .so includes
// pto_orchestration_api.h, so nothing reads that one — rt_scope_* and
// rt_orchestration_done take the runtime as an argument.
eps->bind(rt);
} else {
const auto *entry_points = reinterpret_cast<const HostOrchEntryPoints *>(host_orch_func_ptr);
if (entry_points->bind == nullptr) {
LOG_ERROR("host-orch: orch .so framework_bind_runtime was not resolved");
return -1;
}
rt->active_callable_hash = reinterpret_cast<uint64_t>(entry_points->entry);
rt->tensor_access = &tensor_access;
// Binds the orchestration .so's own framework_current_runtime, which its
// inline rt_submit_* read. The host library links a same-named copy from
// orchestration/common.cpp, but nothing outside the .so includes
// pto_orchestration_api.h, so nothing reads that one — rt_scope_* and
// rt_orchestration_done take the runtime as an argument.
entry_points->bind(rt);

rt_scope_begin(rt);
eps->entry(orch_l2);
entry_points->entry(orch_l2);
rt_scope_end(rt);
rt_orchestration_done(rt);

int32_t total_tasks = pto2_sm_layout::ring_current_task_index_addr(host_sm)->load(std::memory_order_acquire);
const int32_t total_tasks = pto2_sm_layout::ring_current_task_index_addr(host_sm)->load(std::memory_order_acquire);
if (!upload_graph_submissions(runtime, api, *graph_state)) return -1;

// total_tasks sizes the bounded per-segment H2D copies below; a value outside
Expand All @@ -623,7 +616,7 @@ int32_t run_host_orchestration(
const int64_t arena_delta = static_cast<int64_t>(reinterpret_cast<uint64_t>(device_arena)) -
static_cast<int64_t>(reinterpret_cast<uint64_t>(host_arena.base()));
if (!relocate_host_orch_image(
host_sm_handle, rt, reinterpret_cast<uint64_t>(host_sm), sm_size, sm_delta,
host_sm_handle, reinterpret_cast<uint64_t>(host_sm), sm_size, sm_delta,
reinterpret_cast<uint64_t>(host_arena.base()), layout.arena_size, arena_delta
)) {
LOG_ERROR("host-orch: relocation failed; refusing to H2D an image with unrelocated host pointers");
Expand Down
27 changes: 7 additions & 20 deletions src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -896,15 +896,9 @@ int32_t SchedulerContext::run_resolution_thread(Runtime *runtime, int32_t thread
while (true) {
if (completed_.load(std::memory_order_acquire)) break;

// Propagate a fatal error latched by the orchestrator (host) or a
// scheduler thread; mirror resolve_and_dispatch's exit behavior.
if (header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE ||
header->sched_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE) {
if (!completed_.exchange(true, std::memory_order_acq_rel)) {
emergency_shutdown(runtime);
}
int32_t published_task_count = 0;
if (handle_orchestrator_exit(thread_idx, header, runtime, published_task_count) == LoopAction::BREAK_LOOP)
break;
}

int32_t resolved_this_pass = 0;
bool resolved_any = false;
Expand Down Expand Up @@ -976,22 +970,15 @@ int32_t SchedulerContext::run_resolution_thread(Runtime *runtime, int32_t thread
if (completed_.load(std::memory_order_acquire)) break;

if (resolved_any) {
int32_t new_total = completed_tasks_.load(std::memory_order_relaxed);
if (resolved_this_pass > 0) {
new_total =
completed_tasks_.fetch_add(resolved_this_pass, std::memory_order_relaxed) + resolved_this_pass;
completed_tasks_.fetch_add(resolved_this_pass, std::memory_order_relaxed);
#if SIMPLER_SCHED_PROFILING
// P owns the completion accounting, so it owns the profiling mirror too
// (the S threads' completed_this_turn no longer feeds it in P mode).
sched_->tasks_completed.fetch_add(resolved_this_pass, std::memory_order_relaxed);
#endif
}
last_progress_ts = get_sys_cnt_aicpu();
if (total_tasks_ > 0 && new_total >= total_tasks_) {
completed_.store(true, std::memory_order_release);
LOG_INFO("Thread %d: P resolved all tasks %d/%d", thread_idx, new_total, total_tasks_);
break;
}
continue; // fast re-drain while work keeps arriving
}

Expand All @@ -1004,11 +991,12 @@ int32_t SchedulerContext::run_resolution_thread(Runtime *runtime, int32_t thread
// task — a genuine forward-progress stall / pre-dispatch deadlock.
uint64_t now = get_sys_cnt_aicpu();
if (now - last_progress_ts > scheduler_timeout_cycles) {
bool outstanding = total_tasks_ > 0 && completed_tasks_.load(std::memory_order_relaxed) < total_tasks_;
const int32_t total = total_tasks_;
bool outstanding = total > 0 && completed_tasks_.load(std::memory_order_relaxed) < total;
if (outstanding && no_thread_owns_running_task()) {
LOG_ERROR(
"Thread %d: P resolution stall (%d/%d resolved)", thread_idx,
completed_tasks_.load(std::memory_order_relaxed), total_tasks_
completed_tasks_.load(std::memory_order_relaxed), total
);
int32_t expected = PTO2_ERROR_NONE;
header->sched_error_code.compare_exchange_strong(
Expand Down Expand Up @@ -1352,8 +1340,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_
}

// Phase 3 (dependency-only dummy / predicate-failed retirement) runs on
// the resolution thread P, not here — see run_resolution_thread. The
// scheduler loop goes straight from completion detection to dispatch.
// the resolution thread P, not here — see run_resolution_thread.

// Phase 4: MIX-strict-priority dispatch with phase-split and
// cross-thread idle gating. See dispatch_ready_tasks for the policy.
Expand Down
Loading