[Performance] HBG: up to 99% host-side overhead reduction

src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp

@@ -34,6 +34,7 @@
 #include <sys/time.h>
 #include <unistd.h>
 
+#include <atomic>
 #include <cerrno>
 #include <cinttypes>
 #include <cstddef>
@@ -43,6 +44,7 @@
 #include <cstdlib>
 #include <cstring>
 #include <limits>
+#include <memory>
 #include <string>
 #include <type_traits>
 #include <vector>
@@ -456,8 +458,15 @@ int32_t run_host_orchestration(
     // The dep_gen graph belongs to the orchestration that is about to run.
     dep_gen_host_graph_begin_capture();
 
-    std::vector<uint8_t> host_sm_buf(sm_size, 0);
-    void *host_sm = host_sm_buf.data();
+    // Init-on-write: descriptors, payloads, slot_states and completion_flags are
+    // each written per task at submit and read only for [0, total_tasks). Zero
+    // only the fixed-size header here; the per-slot segments are initialized in
+    // orch::prepare_task and shipped bounded to total_tasks below.
+    const pto2_sm_layout::PTO2RingSegmentOffsets sm_segs =
+        pto2_sm_layout::ring_segment_offsets(eff_task_window_sizes[0]);
+    std::unique_ptr<uint8_t[]> host_sm_buf(new uint8_t[sm_size]);
+    void *host_sm = host_sm_buf.get();
+    std::memset(host_sm, 0, sm_segs.descriptors);
 
     // Re-point the orchestrator half at the host SM (scheduler keeps device SM).
     // init_data_from_layout resets the orchestrator state, so this is safe.
@@ -514,6 +523,13 @@ int32_t run_host_orchestration(
 
     int32_t total_tasks = pto2_sm_layout::ring_current_task_index_addr(host_sm)->load(std::memory_order_acquire);
 
+    // total_tasks sizes the bounded per-segment H2D copies below; a value outside
+    // [0, task_window] would make those copies read/write out of bounds.
+    if (total_tasks < 0 || static_cast<uint64_t>(total_tasks) > eff_task_window_sizes[0]) {
+        LOG_ERROR("host-orch: total_tasks %d out of range [0, %" PRIu64 "]", total_tasks, eff_task_window_sizes[0]);
+        return -1;
+    }
+
     // Relocate the host-DDR cross-task pointers to their final DEVICE addresses
     // on the host, before the SM and arena leave for the device. Pointers into
     // the SM shift by sm_delta; pointers into the arena (fanout adjacency, wiring
@@ -531,7 +547,23 @@ int32_t run_host_orchestration(
         return -1;
     }
 
-    if (api->copy_to_device(device_sm, host_sm, sm_size) != 0) {
+    // Ship only the live prefix of each segment: the device reads no slot past
+    // total_tasks, so upload header + descriptors[0,N), payloads[0,N),
+    // slot_states[0,N) and completion_flags[0,N) — never the ring-sized tails.
+    // header + descriptors[0,N) are contiguous, so that is a single copy.
+    const uint64_t nt = static_cast<uint64_t>(total_tasks);
+    const uint64_t hdr_desc_bytes = sm_segs.descriptors + nt * sizeof(PTO2TaskDescriptor);
+    char *host_base = static_cast<char *>(host_sm);
+    char *dev_base = static_cast<char *>(device_sm);
+    if (api->copy_to_device(dev_base, host_base, hdr_desc_bytes) != 0 ||
+        api->copy_to_device(dev_base + sm_segs.payloads, host_base + sm_segs.payloads, nt * sizeof(PTO2TaskPayload)) !=
+            0 ||
+        api->copy_to_device(
+            dev_base + sm_segs.slot_states, host_base + sm_segs.slot_states, nt * sizeof(PTO2TaskSlotState)
+        ) != 0 ||
+        api->copy_to_device(
+            dev_base + sm_segs.completion_flags, host_base + sm_segs.completion_flags, nt * sizeof(std::atomic<uint8_t>)
+        ) != 0) {
         LOG_ERROR("host-orch: H2D of populated SM failed");
         return -1;
     }
@@ -879,7 +911,52 @@ extern "C" int bind_callable_to_runtime_impl(
     // *before* it can dereference the image.
     rt->prebuilt_layout = layout;
 
-    int rc_upload = api->copy_to_device(runtime_arena_dev, host_arena.base(), layout.arena_size);
+    // Skip uploading the orchestrator block (fanin_seen_epoch / scope / tensormap,
+    // ~8.5 MB): it is host-only dep-computation scratch that the AICPU scheduler
+    // never reads. Ship [0, orch_start) (sm_handle); skip the host-only orch block;
+    // then each big ready-queue's slots bounded to a live prefix; then the rest
+    // (small early queues + runtime header + mailbox) whole. orch_start is the first
+    // orch reservation (runtime_reserve_layout order: sm_handle -> orch -> sched).
+    //
+    // Prefix length = total_tasks + 1, clamped to capacity. A queue takes at most
+    // total_tasks pushes, so its enqueue_pos never exceeds total_tasks; but
+    // PTO2ReadyQueue::pop_batch_tagged reads one slot past dequeue_pos — the slot at
+    // enqueue_pos — to detect the empty boundary, and a batched dequeue that finds a
+    // stale (too-large) sequence there spins retrying forever. Seeding that sentinel
+    // slot to its empty sequence makes the boundary read resolve to "empty", so the
+    // ring-sized tail beyond it is never touched.
+    const auto &sq = layout.sched;
+    const size_t orch_start = layout.orch.off_fanin_seen_epoch;
+    const size_t slot_sz = sizeof(PTO2ReadyQueueSlot);
+    const uint64_t q_span = static_cast<uint64_t>(runtime->host_total_tasks) + 1;
+    const uint64_t q_prefix = q_span < sq.ready_queue_capacity ? q_span : sq.ready_queue_capacity;
+    const size_t big_q_slots[] = {
+        sq.off_ready_queue_slots[0],      sq.off_ready_queue_slots[1],      sq.off_ready_queue_slots[2],
+        sq.off_ready_sync_queue_slots[0], sq.off_ready_sync_queue_slots[1], sq.off_ready_sync_queue_slots[2],
+        sq.off_dummy_ready_queue_slots,
+    };
+    // Small early queues start right after the big queues; from there to arena end is
+    // shipped whole (early queues are tiny; runtime header + mailbox are per-run state).
+    const size_t rest_start = sq.off_early_dispatch_queue_slots[0];
+    char *arena_host = static_cast<char *>(host_arena.base());
+    char *arena_dev = static_cast<char *>(runtime_arena_dev);
+    // Seed the big-queue slots [0, q_prefix) that the build skipped (it did headers
+    // only). Matches ready_queue_init_data_from_layout: sequence=i, empty slot_state.
+    for (size_t qi = 0; qi < sizeof(big_q_slots) / sizeof(big_q_slots[0]); qi++) {
+        auto *slots = reinterpret_cast<PTO2ReadyQueueSlot *>(arena_host + big_q_slots[qi]);
+        for (uint64_t i = 0; i < q_prefix; i++) {
+            slots[i].sequence.store(static_cast<int64_t>(i), std::memory_order_relaxed);
+            slots[i].slot_state = nullptr;
+        }
+    }
+    int rc_upload = api->copy_to_device(arena_dev, arena_host, orch_start);
+    for (size_t i = 0; rc_upload == 0 && i < sizeof(big_q_slots) / sizeof(big_q_slots[0]); i++) {
+        rc_upload = api->copy_to_device(arena_dev + big_q_slots[i], arena_host + big_q_slots[i], q_prefix * slot_sz);
+    }
+    if (rc_upload == 0) {
+        rc_upload =
+            api->copy_to_device(arena_dev + rest_start, arena_host + rest_start, layout.arena_size - rest_start);
+    }
     if (rc_upload != 0) {
         LOG_ERROR("Failed to rtMemcpy prebuilt runtime arena to device (rc=%d)", rc_upload);
         return -1;

src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp

@@ -402,6 +402,14 @@ static bool prepare_task(
     out->task = &orch->sm_header->ring.task_descriptors[out->alloc_result.slot];
     out->payload = &orch->sm_header->ring.task_payloads[out->alloc_result.slot];
 
+    // Init-on-write: reset this slot's dynamic scheduling fields and clear its
+    // completion flag as the orchestrator claims it, rather than pre-zeroing the
+    // whole task window in init_header_per_ring. whole-graph-resident hbg claims
+    // slots [0, total_tasks) once, and the device reads no slot past
+    // total_tasks, so the SM's per-slot init cost tracks the task count.
+    out->slot_state->reset_for_reuse();
+    orch->sm_header->ring.completion_flags[out->alloc_result.slot].store(0, std::memory_order_relaxed);
+
     out->payload->prefetch(args.tensor_count(), args.scalar_count());
 
     // Re-bind payload/task pointers each submit. Value is per-slot constant
@@ -418,7 +426,7 @@ static bool prepare_task(
     // early-dispatch fields) is initialized in PTO2TaskPayload::init, the
     // single payload-init point, which runs before Orch-side wiring publish.
 
-    // Fields already zeroed by reset_for_reuse() at slot init:
+    // Fields already zeroed by the reset_for_reuse() above:
     //   wake_list_head=nullptr, next_in_wake_list=nullptr,
     //   any_subtask_deferred=false, completed_subtasks=0, next_block_idx=0
     // Fields immutable after RingSchedState::init():
@@ -435,8 +443,8 @@ static bool prepare_task(
     out->slot_state->task_attrs = task_attrs;
     // Reclaim gate: seed last_consumer to self, so a producer with no consumers
     // is retirable once completed_watermark >= its own id. Each fanin edge bumps
-    // it in append_fanin_or_fail. completion_flags for this slot are already 0
-    // (zeroed once at init; whole-graph-resident hbg never reuses a slot).
+    // it in append_fanin_or_fail. completion_flags for this slot were cleared
+    // above (whole-graph-resident hbg never reuses a slot).
     out->slot_state->last_consumer_local_id = static_cast<int32_t>(out->task_id.local());
     // payload.fanin_count is set in submit_task_common's STEP 6.
     scope_tasks_push(orch, out->slot_state);

src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h

@@ -394,7 +394,15 @@ size_t ready_queue_reserve_layout(DeviceArena &arena, uint64_t capacity);
 // arena.region_ptr(slots_off) only to address the slot array for writes;
 // does NOT store the pointer in `queue->slots`. Call
 // `ready_queue_wire_arena_pointers` afterwards to set the field itself.
-bool ready_queue_init_data_from_layout(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off, uint64_t capacity);
+// slot_init_count bounds the per-slot Vyukov-sequence init loop (defaults to the
+// full capacity). host_build_graph passes 0 for the big ready queues and seeds
+// only [0, total_tasks) of them after orchestration, since the device reads no
+// slot past its push count; the queue header (capacity/mask/positions) is always
+// initialized.
+bool ready_queue_init_data_from_layout(
+    PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off, uint64_t capacity,
+    uint64_t slot_init_count = ~static_cast<uint64_t>(0)
+);
 // Stores queue->slots = arena.region_ptr(slots_off). Idempotent.
 void ready_queue_wire_arena_pointers(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off);
 void ready_queue_destroy(PTO2ReadyQueue *queue);

src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp

@@ -55,7 +55,9 @@ size_t ready_queue_reserve_layout(DeviceArena &arena, uint64_t capacity) {
     return arena.reserve(capacity * sizeof(PTO2ReadyQueueSlot), PTO2_ALIGN_SIZE);
 }
 
-bool ready_queue_init_data_from_layout(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off, uint64_t capacity) {
+bool ready_queue_init_data_from_layout(
+    PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off, uint64_t capacity, uint64_t slot_init_count
+) {
     // Address the slots region for data writes without storing the pointer in
     // queue->slots — that field is set by ready_queue_wire_arena_pointers.
     auto *slots_arena = static_cast<PTO2ReadyQueueSlot *>(arena.region_ptr(slots_off));
@@ -64,7 +66,8 @@ bool ready_queue_init_data_from_layout(PTO2ReadyQueue *queue, DeviceArena &arena
     queue->enqueue_pos.store(0, std::memory_order_relaxed);
     queue->dequeue_pos.store(0, std::memory_order_relaxed);
 
-    for (uint64_t i = 0; i < capacity; i++) {
+    const uint64_t n = slot_init_count < capacity ? slot_init_count : capacity;
+    for (uint64_t i = 0; i < n; i++) {
         slots_arena[i].sequence.store((int64_t)i, std::memory_order_relaxed);
         slots_arena[i].slot_state = nullptr;
     }
@@ -136,21 +139,24 @@ bool PTO2SchedulerState::init_data_from_layout(
     }
 
     for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) {
+        // slot_init_count = 0: big-queue slots are seeded post-orchestration,
+        // bounded to [0, total_tasks) (see runtime_maker.cpp).
         if (!ready_queue_init_data_from_layout(
-                &sched->ready_queues[i], arena, layout.off_ready_queue_slots[i], layout.ready_queue_capacity
+                &sched->ready_queues[i], arena, layout.off_ready_queue_slots[i], layout.ready_queue_capacity, 0
             )) {
             return false;
         }
     }
     for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) {
         if (!ready_queue_init_data_from_layout(
-                &sched->ready_sync_queues[i], arena, layout.off_ready_sync_queue_slots[i], layout.ready_queue_capacity
+                &sched->ready_sync_queues[i], arena, layout.off_ready_sync_queue_slots[i], layout.ready_queue_capacity,
+                0
             )) {
             return false;
         }
     }
     if (!ready_queue_init_data_from_layout(
-            &sched->dummy_ready_queue, arena, layout.off_dummy_ready_queue_slots, layout.ready_queue_capacity
+            &sched->dummy_ready_queue, arena, layout.off_dummy_ready_queue_slots, layout.ready_queue_capacity, 0
         )) {
         return false;
     }

src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp

@@ -178,22 +178,12 @@ void PTO2SharedMemoryHandle::init_header_per_ring(
     header->sched_error_code.store(PTO2_ERROR_NONE, std::memory_order_relaxed);
     header->sched_error_thread.store(-1, std::memory_order_relaxed);
 
-    // Per-ring slot_states reset. Previously lived in
-    // PTO2SchedulerState::RingSchedState::init(), but it writes into
-    // ring->slot_states[] which is SM-side storage — keeping it here lets
-    // host-side prebuilt-arena init skip all SM dereferences.
-    // reset_for_reuse() prepares dynamic fanout/refcount fields so the first
-    // submit doesn't need an explicit reset.
-    auto &ring = header->ring;
-    for (uint64_t i = 0; i < task_window_sizes[0]; i++) {
-        ring.slot_states[i].reset_for_reuse();
-        ring.slot_states[i].active_mask = ActiveMask{};
-    }
-
-    // Polling completion flags: 0 = pending. Shared memory is not guaranteed
-    // zero on device; stale non-zero bytes would make consumers observe a
-    // producer as already completed. Zero the whole per-ring array once.
-    __builtin_memset((void *)ring.completion_flags, 0, task_window_sizes[0] * sizeof(std::atomic<uint8_t>));
+    // Per-slot init (slot_states.reset_for_reuse() + active_mask, and clearing
+    // the completion flag) is done init-on-write in orch::prepare_task as each
+    // slot [0, total_tasks) is claimed — not swept over the whole task window
+    // here — so the SM init/upload cost tracks the task count, not ring capacity.
+    // The device reads no slot past total_tasks, so the unclaimed tail needs no
+    // reset.
 }
 
 // =============================================================================

src/a2a3/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp

@@ -94,13 +94,11 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev
     // The pool's persistent invariant after init is "bucket_index == -1 means
     // not linked", set explicitly below.
     memset(entry_pool_arena, 0, static_cast<size_t>(pool_size) * sizeof(PTO2TensorMapEntry));
+    // The memset already zeroed every field (all four link pointers -> nullptr,
+    // producer_task_id -> {}); only bucket_index needs its non-zero "not linked"
+    // marker, so the per-entry loop writes just that.
     for (int32_t i = 0; i < pool_size; i++) {
         entry_pool_arena[i].bucket_index = -1;
-        entry_pool_arena[i].next_in_bucket = nullptr;
-        entry_pool_arena[i].prev_in_bucket = nullptr;
-        entry_pool_arena[i].next_in_task = nullptr;
-        entry_pool_arena[i].prev_in_task = nullptr;
-        entry_pool_arena[i].producer_task_id = PTO2TaskId{};
     }
 
     // free_entry_list: zeroed (was calloc'd before); contents become meaningful