From 79923edce5e7c96c44fc025657f96c65fe22a515 Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Wed, 1 Jul 2026 14:28:54 +0800 Subject: [PATCH] Support: reuse resident prebuilt runtime arenas - Skip cached runtime arena image uploads on bind cache hits - Reset mutable AICPU runtime state in-place before reuse --- .../aicpu/aicpu_executor.cpp | 21 ++- .../host/runtime_maker.cpp | 7 +- .../runtime/pto_async_wait.h | 6 + .../runtime/pto_orchestrator.h | 4 + .../runtime/pto_ring_buffer.h | 19 +++ .../runtime/pto_runtime2.h | 1 + .../runtime/pto_tensormap.h | 23 ++- .../runtime/scheduler/pto_scheduler.h | 11 ++ .../runtime/shared/pto_runtime2_init.cpp | 131 +++++++++++++++++- .../runtime/shared/pto_tensormap.cpp | 33 +++++ .../aicpu/aicpu_executor.cpp | 21 ++- .../host/runtime_maker.cpp | 7 +- .../runtime/pto_async_wait.h | 6 + .../runtime/pto_orchestrator.h | 4 + .../runtime/pto_ring_buffer.h | 19 +++ .../runtime/pto_runtime2.h | 1 + .../runtime/pto_tensormap.h | 23 ++- .../runtime/scheduler/pto_scheduler.h | 11 ++ .../runtime/shared/pto_runtime2_init.cpp | 130 ++++++++++++++++- .../runtime/shared/pto_tensormap.cpp | 33 +++++ 20 files changed, 473 insertions(+), 38 deletions(-) diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp index 983d478308..152f82fc44 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp @@ -480,13 +480,9 @@ int32_t AicpuExecutor::run(Runtime *runtime) { sm_ptr = runtime->get_gm_sm_ptr(); } - // Prebuilt-arena fast path. Host has pre-populated the entire - // runtime arena (PTO2Runtime + orchestrator/scheduler/tensor_map - // sub-regions + sm_handle wrapper + mailbox) and uploaded it via - // rtMemcpy into the pooled runtime_arena buffer. We attach to it, - // wire arena-internal pointers to their device addresses, reset - // the SM, and finalize the few device-only fields the host could - // not know at image-build time. + // Prebuilt-arena fast path. Host uploads the runtime arena image + // on cache miss; cache hits reuse the resident device arena. AICPU + // re-wires arena-internal pointers to device addresses below. { AicpuPhaseScope arena_wire(AicpuPhase::ArenaWire); void *prebuilt_arena = runtime->get_prebuilt_arena_base(); @@ -513,10 +509,7 @@ int32_t AicpuExecutor::run(Runtime *runtime) { } // Reset SM state. setup_pointers + init_header_per_ring restore - // ring flow-control counters, layout metadata, error flags, and - // the per-slot ring->slot_states[] (bind_ring + reset_for_reuse + - // fanin_count/active_mask zero — previously done inside - // RingSchedState::init). + // ring flow-control counters, layout metadata, and error flags. { AicpuPhaseScope sm_reset(AicpuPhase::SmReset); memset(rt->sm_handle, 0, sizeof(*rt->sm_handle)); @@ -529,6 +522,12 @@ int32_t AicpuExecutor::run(Runtime *runtime) { runtime_init_ready_.store(true, std::memory_order_release); return -1; } + if (!runtime_reset_for_reuse(runtime_arena_, rt->prebuilt_layout, rt)) { + LOG_ERROR("Thread %d: runtime_reset_for_reuse failed", thread_idx); + rt = nullptr; + runtime_init_ready_.store(true, std::memory_order_release); + return -1; + } // AICore completion mailbox lives in the pooled arena, so its // head/tail/seq survive across runs and stay monotonic. We do diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index 4b7205db83..fb21356a7b 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -627,11 +627,8 @@ static int bind_cached_runtime_image( } runtime->set_orch_args(device_args); - int rc_upload = api->copy_to_device(runtime_arena_dev, cached_image, cached_image_size); - if (rc_upload != 0) { - LOG_ERROR("Failed to rtMemcpy cached prebuilt runtime arena to device (rc=%d)", rc_upload); - return -1; - } + (void)cached_image; + (void)cached_image_size; runtime->set_gm_sm_ptr(sm_ptr); runtime->set_prebuilt_arena(runtime_arena_dev, runtime_off); return 0; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h index 65608ad2fc..42a9474187 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h @@ -152,6 +152,12 @@ struct AsyncWaitList { // Read by scheduler shutdown / l2 perf summary; not on the hot path. std::atomic mpsc_skipped_count{0}; + void reset_for_reuse() { + busy.store(0, std::memory_order_relaxed); + count = 0; + mpsc_skipped_count.store(0, std::memory_order_relaxed); + } + bool try_lock() { int32_t expected = 0; return busy.compare_exchange_strong(expected, 1, std::memory_order_acquire, std::memory_order_relaxed); diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h index 8ffe39b310..5ceb9af855 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h @@ -161,6 +161,10 @@ struct PTO2OrchestratorState { const PTO2OrchestratorLayout &layout, DeviceArena &arena, void *sm_dev_base, void *gm_heap, const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH] ); + bool reset_for_reuse( + const PTO2OrchestratorLayout &layout, void *sm_dev_base, void *gm_heap, + const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH] + ); // Phase 3b: write the arena-internal pointer fields (scope_tasks, // scope_begins, rings[].fanin_pool.base, tensor_map.{buckets,entry_pool, diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h index a93f39f7cb..97f318d40f 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h @@ -521,6 +521,15 @@ struct PTO2FaninPool { error_code_ptr = in_error_code_ptr; } + void reset_for_reuse(std::atomic *in_error_code_ptr) { + top = 1; + tail = 1; + high_water = 0; + reclaim_task_cursor = 0; + base[0].slot_state = nullptr; + error_code_ptr = in_error_code_ptr; + } + void reclaim(PTO2SharedMemoryRingHeader &ring, int32_t sm_last_task_alive); bool ensure_space(PTO2SharedMemoryRingHeader &ring, int32_t needed); @@ -688,6 +697,16 @@ struct PTO2DepListPool { error_code_ptr = in_error_code_ptr; } + void reset_for_reuse(std::atomic *in_error_code_ptr) { + top = 1; + tail = 1; + high_water = 0; + last_reclaimed = 0; + base[0].slot_state = nullptr; + base[0].next = nullptr; + error_code_ptr = in_error_code_ptr; + } + /** * Reclaim dead entries based on scheduler's slot state dep_pool_mark. * Safe to call multiple times — only advances tail forward. diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h index 058f90f40c..64f4c63196 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h @@ -229,6 +229,7 @@ PTO2Runtime *runtime_init_data_from_layout( * addresses) sides. */ void runtime_wire_arena_pointers(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt); +bool runtime_reset_for_reuse(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt); /** * AICPU-only Phase 4 — fill in the few fields the host could not know at diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h index dea57b7b6b..33673b29c1 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h @@ -74,9 +74,11 @@ struct Segment { */ struct PTO2TensorMapLayout { size_t off_buckets; + size_t off_bucket_epochs; size_t off_entry_pool; size_t off_free_entry_list; size_t off_task_entry_heads[PTO2_MAX_RING_DEPTH]; + size_t off_task_entry_head_epochs[PTO2_MAX_RING_DEPTH]; int32_t num_buckets; int32_t pool_size; int32_t task_window_sizes[PTO2_MAX_RING_DEPTH]; @@ -359,7 +361,8 @@ static_assert( struct PTO2TensorMap { // Hash table buckets (fixed size, power of 2) PTO2TensorMapEntry **buckets; // Array of offsets into entry_pool (-1 = empty) - int32_t num_buckets; // Must be power of 2 for fast modulo + uint32_t *bucket_epochs; + int32_t num_buckets; // Must be power of 2 for fast modulo // Entry pool as ring buffer PTO2TensorMapEntry *entry_pool; // Ring buffer of entries @@ -371,7 +374,9 @@ struct PTO2TensorMap { // Per-ring per-task entry tracking (for efficient bucket cleanup) // Indexed by [ring_id][local_id & (task_window_sizes[ring_id] - 1)] PTO2TensorMapEntry **task_entry_heads[PTO2_MAX_RING_DEPTH]; + uint32_t *task_entry_head_epochs[PTO2_MAX_RING_DEPTH]; int32_t task_window_sizes[PTO2_MAX_RING_DEPTH]; // Per-ring task window size (for slot masking) + uint32_t current_epoch{1}; // Per-ring validity threshold (for lazy invalidation) int32_t last_task_alives[PTO2_MAX_RING_DEPTH]; // Cached from shared memory per ring @@ -421,7 +426,6 @@ struct PTO2TensorMap { } always_assert(next_entry_idx < pool_size); PTO2TensorMapEntry *res = &entry_pool[next_entry_idx++]; - debug_assert(res->bucket_index == -1); return res; } @@ -479,6 +483,7 @@ struct PTO2TensorMap { * a host arena that holds the prebuilt image. */ bool init_data_from_layout(const PTO2TensorMapLayout &layout, DeviceArena &arena); + void reset_for_reuse(const PTO2TensorMapLayout &layout); /** * Phase 3b: write the arena-internal pointer fields. Idempotent; @@ -518,6 +523,9 @@ struct PTO2TensorMap { template void lookup(const Tensor &tensor, Fn &&on_match) { uint32_t bucket_index = hash(tensor.buffer.addr); + if (bucket_epochs[bucket_index] != current_epoch) { + return; + } PTO2TensorMapEntry *cur_entry = buckets[bucket_index]; #if PTO2_TENSORMAP_PROFILING @@ -598,6 +606,9 @@ struct PTO2TensorMap { // Iterate through retired tasks on this ring and remove their entries for (int32_t local_id = old_last_task_alive; local_id < new_last_task_alive; local_id++) { int32_t task_slot = local_id & (task_window_sizes[ring_id] - 1); + if (task_entry_head_epochs[ring_id][task_slot] != current_epoch) { + continue; + } PTO2TensorMapEntry *cur_entry = task_entry_heads[ring_id][task_slot]; while (cur_entry != nullptr) { @@ -649,6 +660,10 @@ struct PTO2TensorMap { entry->producer_task_id = producer_task_id; // Insert at head of hash bucket + if (bucket_epochs[bucket_index] != current_epoch) { + buckets[bucket_index] = nullptr; + bucket_epochs[bucket_index] = current_epoch; + } entry->bucket_index = bucket_index; entry->next_in_bucket = buckets[bucket_index]; if (entry->next_in_bucket != nullptr) { @@ -658,6 +673,10 @@ struct PTO2TensorMap { entry->prev_in_bucket = nullptr; // Link to task's entry list + if (task_entry_head_epochs[ring_id][task_slot] != current_epoch) { + task_entry_heads[ring_id][task_slot] = nullptr; + task_entry_head_epochs[ring_id][task_slot] = current_epoch; + } entry->next_in_task = task_entry_heads[ring_id][task_slot]; entry->prev_in_task = nullptr; if (entry->next_in_task != nullptr) { diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index f84d40f358..a37eb0d43d 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -126,6 +126,8 @@ struct alignas(64) PTO2ReadyQueue { return (e >= d) ? (e - d) : 0; } + void reset_for_reuse() {} + bool push(PTO2TaskSlotState *slot_state) { uint64_t pos; PTO2ReadyQueueSlot *slot; @@ -485,6 +487,13 @@ struct alignas(64) PTO2SpscQueue { buffer_ = static_cast(arena.region_ptr(buffer_off)); } + void reset_for_reuse() { + uint64_t h = head_.load(std::memory_order_relaxed); + tail_.store(h, std::memory_order_relaxed); + tail_cached_ = h; + head_cached_ = h; + } + // Arena owns the buffer; here we only forget our pointer. void destroy() { buffer_ = nullptr; } @@ -604,6 +613,7 @@ struct PTO2SchedulerState { // the device address of the SM ring header — computed via offset // arithmetic, no SM dereference. bool init_data_from_layout(void *sm_dev_base, int32_t ring_id); + void reset_for_reuse(void *sm_dev_base, int32_t ring_id, std::atomic *orch_err); void destroy(); void sync_to_sm() { ring->fc.last_task_alive.store(last_task_alive, std::memory_order_release); } @@ -1425,6 +1435,7 @@ struct PTO2SchedulerState { // scheduler only needs the SM header / ring header base addresses, // both window-size-independent.) bool init_data_from_layout(const PTO2SchedulerLayout &layout, DeviceArena &arena, void *sm_dev_base); + void reset_for_reuse(const PTO2SchedulerLayout &layout, void *sm_dev_base); // Phase 3b: write the arena-internal pointer fields // (ready_queues[].slots, dummy_ready_queue.slots, dep_pool.base for each diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp index c63c94f92a..1561acc563 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp @@ -13,7 +13,7 @@ * * Lives under runtime/shared/ so it is included in both the host_runtime.so * build (host pre-populates the prebuilt arena image) and the aicpu_runtime - * build (AICPU runs wire_arena_pointers + destroy after attach). The + * build (AICPU runs wire_arena_pointers + reset_for_reuse after attach). The * device-only parts of pto_runtime2.cpp / pto_orchestrator.cpp / pto_scheduler.cpp * (ops table, scope/submit/dispatch business logic, profiling) stay in their * original files and the aicpu build only. @@ -104,6 +104,20 @@ bool PTO2SchedulerState::RingSchedState::init_data_from_layout(void *sm_dev_base return true; } +void PTO2SchedulerState::RingSchedState::reset_for_reuse( + void *sm_dev_base, int32_t ring_id, std::atomic *orch_err +) { + ring = pto2_sm_layout::ring_header_addr(sm_dev_base, ring_id); + last_task_alive = 0; + advance_lock.store(0, std::memory_order_relaxed); + dep_deadlock_reported = false; + dep_pool.reset_for_reuse(orch_err); +#if PTO2_PROFILING + dep_pool_snapshot_tail.store(1, std::memory_order_relaxed); + dep_pool_snapshot_top.store(1, std::memory_order_relaxed); +#endif +} + void PTO2SchedulerState::RingSchedState::destroy() { ring = nullptr; } PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena, int32_t dep_pool_capacity) { @@ -190,6 +204,35 @@ bool PTO2SchedulerState::init_data_from_layout( return true; } +void PTO2SchedulerState::reset_for_reuse(const PTO2SchedulerLayout &layout, void *sm_dev_base) { + PTO2SchedulerState *sched = this; + sched->sm_header = reinterpret_cast(sm_dev_base); +#if PTO2_SCHED_PROFILING + sched->tasks_completed.store(0, std::memory_order_relaxed); + sched->tasks_consumed.store(0, std::memory_order_relaxed); +#endif + + auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + sched->ring_sched_states[r].reset_for_reuse(sm_dev_base, r, orch_err); + } + + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + sched->ready_queues[i].reset_for_reuse(); + } + sched->dummy_ready_queue.reset_for_reuse(); + sched->early_dispatch_queue.reset_for_reuse(); + + sched->wiring.queue.reset_for_reuse(); + sched->wiring.batch_count = 0; + sched->wiring.batch_index = 0; + sched->wiring.backoff_counter = 0; + sched->wiring.orch_needs_drain.store(false, std::memory_order_relaxed); + sched->wiring.producer_blocked.store(0, std::memory_order_relaxed); + sched->async_wait_list.reset_for_reuse(); + (void)layout; +} + void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, DeviceArena &arena) { PTO2SchedulerState *sched = this; for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { @@ -348,6 +391,66 @@ bool PTO2OrchestratorState::init_data_from_layout( return true; } +bool PTO2OrchestratorState::reset_for_reuse( + const PTO2OrchestratorLayout &layout, void *sm_dev_base, void *gm_heap, + const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH] +) { + auto *orch = this; + orch->sm_header = reinterpret_cast(sm_dev_base); + orch->gm_heap_base = gm_heap; + uint64_t total_heap_size = 0; + if (!sum_ring_heap_sizes(heap_sizes, &total_heap_size)) { + return false; + } + orch->gm_heap_size = total_heap_size; + orch->fatal = false; + orch->inline_completed_tasks = 0; + + uint32_t next_epoch = orch->fanin_seen_current_epoch + 1; + if (next_epoch == 0) { + next_epoch = 1; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + memset( + orch->fanin_seen_epoch[r], 0, + static_cast(layout.tensor_map.task_window_sizes[r]) * sizeof(uint32_t) + ); + } + } + orch->fanin_seen_current_epoch = next_epoch; + + auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); + uint64_t heap_offset = 0; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + void *ring_heap_base = reinterpret_cast(gm_heap) + heap_offset; + auto *task_descs_dev = pto2_sm_layout::ring_task_descriptors_addr(sm_dev_base, task_window_sizes, r); + auto *slot_states_dev = pto2_sm_layout::ring_slot_states_addr(sm_dev_base, task_window_sizes, r); + auto *cur_idx_dev = pto2_sm_layout::ring_current_task_index_addr(sm_dev_base, r); + auto *last_alive_dev = pto2_sm_layout::ring_last_task_alive_addr(sm_dev_base, r); + + orch->rings[r].task_allocator.init( + task_descs_dev, static_cast(task_window_sizes[r]), cur_idx_dev, last_alive_dev, ring_heap_base, + heap_sizes[r], orch_err, slot_states_dev, 0, static_cast(r) + ); + heap_offset += heap_sizes[r]; + orch->rings[r].fanin_pool.reset_for_reuse(orch_err); + } + + orch->tensor_map.reset_for_reuse(layout.tensor_map); + orch->scope_tasks_size = 0; + orch->scope_tasks_capacity = layout.scope_tasks_cap; + orch->scope_stack_top = -1; + orch->scope_stack_capacity = layout.scope_stack_capacity; + orch->manual_begin_depth = PTO2_MAX_SCOPE_DEPTH; + orch->total_cluster_count = 0; + orch->total_aiv_count = 0; +#if PTO2_PROFILING + orch->tasks_submitted = 0; + orch->buffers_allocated = 0; + orch->bytes_allocated = 0; +#endif + return true; +} + void PTO2OrchestratorState::wire_arena_pointers( const PTO2OrchestratorLayout &layout, DeviceArena &arena, PTO2SchedulerState *scheduler_arg ) { @@ -472,6 +575,32 @@ void runtime_wire_arena_pointers(DeviceArena &arena, const PTO2RuntimeArenaLayou rt->scheduler.wire_arena_pointers(layout.offsets.sched, arena); } +bool runtime_reset_for_reuse(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt) { + (void)arena; + if (rt == nullptr) { + return false; + } + + rt->pending_scope_mode = PTO2ScopeMode::AUTO; + rt->total_cycles = 0; + rt->gm_heap_owned = false; + + uint64_t total_heap_size = 0; + if (!sum_ring_heap_sizes(layout.sizing.heap_sizes, &total_heap_size)) { + return false; + } + rt->gm_heap_size = total_heap_size; + + if (!rt->orchestrator.reset_for_reuse( + layout.offsets.orch, rt->sm_handle->sm_base, rt->gm_heap, layout.sizing.heap_sizes, + layout.sizing.task_window_sizes + )) { + return false; + } + rt->scheduler.reset_for_reuse(layout.offsets.sched, rt->sm_handle->sm_base); + return true; +} + void runtime_destroy(PTO2Runtime *rt, DeviceArena & /*arena*/) { // Arena buffer is pooled across runs by DeviceRunner — never freed here. if (!rt) return; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp index b99c672333..fb22bb4d2b 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp @@ -64,6 +64,8 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout( layout.off_buckets = arena.reserve( static_cast(new_num_buckets) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) ); + layout.off_bucket_epochs = + arena.reserve(static_cast(new_num_buckets) * sizeof(uint32_t), alignof(uint32_t)); layout.off_entry_pool = arena.reserve(static_cast(new_pool_size) * sizeof(PTO2TensorMapEntry), alignof(PTO2TensorMapEntry)); layout.off_free_entry_list = @@ -72,6 +74,8 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout( layout.off_task_entry_heads[r] = arena.reserve( static_cast(new_task_window_sizes[r]) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) ); + layout.off_task_entry_head_epochs[r] = + arena.reserve(static_cast(new_task_window_sizes[r]) * sizeof(uint32_t), alignof(uint32_t)); } return layout; } @@ -88,12 +92,14 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev // Address arena regions for data writes; do not store these in struct // fields (wire_arena_pointers does that). auto *buckets_arena = static_cast(arena.region_ptr(layout.off_buckets)); + auto *bucket_epochs_arena = static_cast(arena.region_ptr(layout.off_bucket_epochs)); auto *entry_pool_arena = static_cast(arena.region_ptr(layout.off_entry_pool)); auto *free_list_arena = static_cast(arena.region_ptr(layout.off_free_entry_list)); // buckets[]: empty == nullptr. for (int32_t i = 0; i < num_buckets; i++) { buckets_arena[i] = nullptr; + bucket_epochs_arena[i] = 0; } // entry_pool: zero-init equivalent to the previous calloc(entry_pool, ...). @@ -118,8 +124,10 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { auto *heads_arena = static_cast(arena.region_ptr(layout.off_task_entry_heads[r])); + auto *head_epochs_arena = static_cast(arena.region_ptr(layout.off_task_entry_head_epochs[r])); for (int32_t i = 0; i < layout.task_window_sizes[r]; i++) { heads_arena[i] = nullptr; + head_epochs_arena[i] = 0; } task_window_sizes[r] = layout.task_window_sizes[r]; last_task_alives[r] = 0; @@ -129,12 +137,35 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev return true; } +void PTO2TensorMap::reset_for_reuse(const PTO2TensorMapLayout &layout) { + num_buckets = layout.num_buckets; + pool_size = layout.pool_size; + next_entry_idx = 0; + free_num = 0; + current_epoch++; + if (current_epoch == 0) { + current_epoch = 1; + memset(bucket_epochs, 0, static_cast(layout.num_buckets) * sizeof(uint32_t)); + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + memset(task_entry_head_epochs[r], 0, static_cast(layout.task_window_sizes[r]) * sizeof(uint32_t)); + } + } + + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + task_window_sizes[r] = layout.task_window_sizes[r]; + last_task_alives[r] = 0; + last_cleanup[r] = 0; + } +} + void PTO2TensorMap::wire_arena_pointers(const PTO2TensorMapLayout &layout, DeviceArena &arena) { buckets = static_cast(arena.region_ptr(layout.off_buckets)); + bucket_epochs = static_cast(arena.region_ptr(layout.off_bucket_epochs)); entry_pool = static_cast(arena.region_ptr(layout.off_entry_pool)); free_entry_list = static_cast(arena.region_ptr(layout.off_free_entry_list)); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { task_entry_heads[r] = static_cast(arena.region_ptr(layout.off_task_entry_heads[r])); + task_entry_head_epochs[r] = static_cast(arena.region_ptr(layout.off_task_entry_head_epochs[r])); } } @@ -143,10 +174,12 @@ void PTO2TensorMap::destroy() { // stray post-destroy access trips a nullptr dereference instead of reading // a recycled allocation. buckets = nullptr; + bucket_epochs = nullptr; entry_pool = nullptr; free_entry_list = nullptr; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { task_entry_heads[r] = nullptr; + task_entry_head_epochs[r] = nullptr; } } diff --git a/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp b/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp index b40d7fda3b..69863d9cc9 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp @@ -477,13 +477,9 @@ int32_t AicpuExecutor::run(Runtime *runtime) { sm_ptr = runtime->get_gm_sm_ptr(); } - // Prebuilt-arena fast path. Host has pre-populated the entire - // runtime arena (PTO2Runtime + orchestrator/scheduler/tensor_map - // sub-regions + sm_handle wrapper + mailbox) and uploaded it via - // rtMemcpy into the pooled runtime_arena buffer. We attach to it, - // wire arena-internal pointers to their device addresses, reset - // the SM, and finalize the few device-only fields the host could - // not know at image-build time. + // Prebuilt-arena fast path. Host uploads the runtime arena image + // on cache miss; cache hits reuse the resident device arena. AICPU + // re-wires arena-internal pointers to device addresses below. { AicpuPhaseScope arena_wire(AicpuPhase::ArenaWire); void *prebuilt_arena = runtime->get_prebuilt_arena_base(); @@ -510,10 +506,7 @@ int32_t AicpuExecutor::run(Runtime *runtime) { } // Reset SM state. setup_pointers + init_header_per_ring restore - // ring flow-control counters, layout metadata, error flags, and - // the per-slot ring->slot_states[] (bind_ring + reset_for_reuse + - // fanin_count/active_mask zero — previously done inside - // RingSchedState::init). + // ring flow-control counters, layout metadata, and error flags. { AicpuPhaseScope sm_reset(AicpuPhase::SmReset); memset(rt->sm_handle, 0, sizeof(*rt->sm_handle)); @@ -526,6 +519,12 @@ int32_t AicpuExecutor::run(Runtime *runtime) { runtime_init_ready_.store(true, std::memory_order_release); return -1; } + if (!runtime_reset_for_reuse(runtime_arena_, rt->prebuilt_layout, rt)) { + LOG_ERROR("Thread %d: runtime_reset_for_reuse failed", thread_idx); + rt = nullptr; + runtime_init_ready_.store(true, std::memory_order_release); + return -1; + } // AICore completion mailbox lives in the pooled arena, so its // head/tail/seq survive across runs and stay monotonic. We do diff --git a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index 4b7205db83..fb21356a7b 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -627,11 +627,8 @@ static int bind_cached_runtime_image( } runtime->set_orch_args(device_args); - int rc_upload = api->copy_to_device(runtime_arena_dev, cached_image, cached_image_size); - if (rc_upload != 0) { - LOG_ERROR("Failed to rtMemcpy cached prebuilt runtime arena to device (rc=%d)", rc_upload); - return -1; - } + (void)cached_image; + (void)cached_image_size; runtime->set_gm_sm_ptr(sm_ptr); runtime->set_prebuilt_arena(runtime_arena_dev, runtime_off); return 0; diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h index 65608ad2fc..42a9474187 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h @@ -152,6 +152,12 @@ struct AsyncWaitList { // Read by scheduler shutdown / l2 perf summary; not on the hot path. std::atomic mpsc_skipped_count{0}; + void reset_for_reuse() { + busy.store(0, std::memory_order_relaxed); + count = 0; + mpsc_skipped_count.store(0, std::memory_order_relaxed); + } + bool try_lock() { int32_t expected = 0; return busy.compare_exchange_strong(expected, 1, std::memory_order_acquire, std::memory_order_relaxed); diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h index a8ed3817fb..8850151f46 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h @@ -160,6 +160,10 @@ struct PTO2OrchestratorState { const PTO2OrchestratorLayout &layout, DeviceArena &arena, void *sm_dev_base, void *gm_heap, const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH] ); + bool reset_for_reuse( + const PTO2OrchestratorLayout &layout, void *sm_dev_base, void *gm_heap, + const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH] + ); // Phase 3b: write the arena-internal pointer fields (scope_tasks, // scope_begins, rings[].fanin_pool.base, tensor_map.{buckets,entry_pool, diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h index a93f39f7cb..97f318d40f 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h @@ -521,6 +521,15 @@ struct PTO2FaninPool { error_code_ptr = in_error_code_ptr; } + void reset_for_reuse(std::atomic *in_error_code_ptr) { + top = 1; + tail = 1; + high_water = 0; + reclaim_task_cursor = 0; + base[0].slot_state = nullptr; + error_code_ptr = in_error_code_ptr; + } + void reclaim(PTO2SharedMemoryRingHeader &ring, int32_t sm_last_task_alive); bool ensure_space(PTO2SharedMemoryRingHeader &ring, int32_t needed); @@ -688,6 +697,16 @@ struct PTO2DepListPool { error_code_ptr = in_error_code_ptr; } + void reset_for_reuse(std::atomic *in_error_code_ptr) { + top = 1; + tail = 1; + high_water = 0; + last_reclaimed = 0; + base[0].slot_state = nullptr; + base[0].next = nullptr; + error_code_ptr = in_error_code_ptr; + } + /** * Reclaim dead entries based on scheduler's slot state dep_pool_mark. * Safe to call multiple times — only advances tail forward. diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h index 187de69a14..fbb65a693b 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h @@ -230,6 +230,7 @@ PTO2Runtime *runtime_init_data_from_layout( * addresses) sides. */ void runtime_wire_arena_pointers(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt); +bool runtime_reset_for_reuse(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt); /** * AICPU-only Phase 4 — fill in the few fields the host could not know at diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h index dea57b7b6b..33673b29c1 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h @@ -74,9 +74,11 @@ struct Segment { */ struct PTO2TensorMapLayout { size_t off_buckets; + size_t off_bucket_epochs; size_t off_entry_pool; size_t off_free_entry_list; size_t off_task_entry_heads[PTO2_MAX_RING_DEPTH]; + size_t off_task_entry_head_epochs[PTO2_MAX_RING_DEPTH]; int32_t num_buckets; int32_t pool_size; int32_t task_window_sizes[PTO2_MAX_RING_DEPTH]; @@ -359,7 +361,8 @@ static_assert( struct PTO2TensorMap { // Hash table buckets (fixed size, power of 2) PTO2TensorMapEntry **buckets; // Array of offsets into entry_pool (-1 = empty) - int32_t num_buckets; // Must be power of 2 for fast modulo + uint32_t *bucket_epochs; + int32_t num_buckets; // Must be power of 2 for fast modulo // Entry pool as ring buffer PTO2TensorMapEntry *entry_pool; // Ring buffer of entries @@ -371,7 +374,9 @@ struct PTO2TensorMap { // Per-ring per-task entry tracking (for efficient bucket cleanup) // Indexed by [ring_id][local_id & (task_window_sizes[ring_id] - 1)] PTO2TensorMapEntry **task_entry_heads[PTO2_MAX_RING_DEPTH]; + uint32_t *task_entry_head_epochs[PTO2_MAX_RING_DEPTH]; int32_t task_window_sizes[PTO2_MAX_RING_DEPTH]; // Per-ring task window size (for slot masking) + uint32_t current_epoch{1}; // Per-ring validity threshold (for lazy invalidation) int32_t last_task_alives[PTO2_MAX_RING_DEPTH]; // Cached from shared memory per ring @@ -421,7 +426,6 @@ struct PTO2TensorMap { } always_assert(next_entry_idx < pool_size); PTO2TensorMapEntry *res = &entry_pool[next_entry_idx++]; - debug_assert(res->bucket_index == -1); return res; } @@ -479,6 +483,7 @@ struct PTO2TensorMap { * a host arena that holds the prebuilt image. */ bool init_data_from_layout(const PTO2TensorMapLayout &layout, DeviceArena &arena); + void reset_for_reuse(const PTO2TensorMapLayout &layout); /** * Phase 3b: write the arena-internal pointer fields. Idempotent; @@ -518,6 +523,9 @@ struct PTO2TensorMap { template void lookup(const Tensor &tensor, Fn &&on_match) { uint32_t bucket_index = hash(tensor.buffer.addr); + if (bucket_epochs[bucket_index] != current_epoch) { + return; + } PTO2TensorMapEntry *cur_entry = buckets[bucket_index]; #if PTO2_TENSORMAP_PROFILING @@ -598,6 +606,9 @@ struct PTO2TensorMap { // Iterate through retired tasks on this ring and remove their entries for (int32_t local_id = old_last_task_alive; local_id < new_last_task_alive; local_id++) { int32_t task_slot = local_id & (task_window_sizes[ring_id] - 1); + if (task_entry_head_epochs[ring_id][task_slot] != current_epoch) { + continue; + } PTO2TensorMapEntry *cur_entry = task_entry_heads[ring_id][task_slot]; while (cur_entry != nullptr) { @@ -649,6 +660,10 @@ struct PTO2TensorMap { entry->producer_task_id = producer_task_id; // Insert at head of hash bucket + if (bucket_epochs[bucket_index] != current_epoch) { + buckets[bucket_index] = nullptr; + bucket_epochs[bucket_index] = current_epoch; + } entry->bucket_index = bucket_index; entry->next_in_bucket = buckets[bucket_index]; if (entry->next_in_bucket != nullptr) { @@ -658,6 +673,10 @@ struct PTO2TensorMap { entry->prev_in_bucket = nullptr; // Link to task's entry list + if (task_entry_head_epochs[ring_id][task_slot] != current_epoch) { + task_entry_heads[ring_id][task_slot] = nullptr; + task_entry_head_epochs[ring_id][task_slot] = current_epoch; + } entry->next_in_task = task_entry_heads[ring_id][task_slot]; entry->prev_in_task = nullptr; if (entry->next_in_task != nullptr) { diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index 5989ded317..ffa47a3009 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -125,6 +125,8 @@ struct alignas(64) PTO2ReadyQueue { return (e >= d) ? (e - d) : 0; } + void reset_for_reuse() {} + bool push(PTO2TaskSlotState *slot_state) { uint64_t pos; PTO2ReadyQueueSlot *slot; @@ -484,6 +486,13 @@ struct alignas(64) PTO2SpscQueue { buffer_ = static_cast(arena.region_ptr(buffer_off)); } + void reset_for_reuse() { + uint64_t h = head_.load(std::memory_order_relaxed); + tail_.store(h, std::memory_order_relaxed); + tail_cached_ = h; + head_cached_ = h; + } + // Arena owns the buffer; here we only forget our pointer. void destroy() { buffer_ = nullptr; } @@ -602,6 +611,7 @@ struct PTO2SchedulerState { // the device address of the SM ring header — computed via offset // arithmetic, no SM dereference. bool init_data_from_layout(void *sm_dev_base, int32_t ring_id); + void reset_for_reuse(void *sm_dev_base, int32_t ring_id, std::atomic *orch_err); void destroy(); void sync_to_sm() { ring->fc.last_task_alive.store(last_task_alive, std::memory_order_release); } @@ -1223,6 +1233,7 @@ struct PTO2SchedulerState { // scheduler only needs the SM header / ring header base addresses, // both window-size-independent.) bool init_data_from_layout(const PTO2SchedulerLayout &layout, DeviceArena &arena, void *sm_dev_base); + void reset_for_reuse(const PTO2SchedulerLayout &layout, void *sm_dev_base); // Phase 3b: write the arena-internal pointer fields // (ready_queues[].slots, dummy_ready_queue.slots, dep_pool.base for each diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp index 6301a8a26f..7072576263 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp @@ -13,7 +13,7 @@ * * Lives under runtime/shared/ so it is included in both the host_runtime.so * build (host pre-populates the prebuilt arena image) and the aicpu_runtime - * build (AICPU runs wire_arena_pointers + destroy after attach). The + * build (AICPU runs wire_arena_pointers + reset_for_reuse after attach). The * device-only parts of pto_runtime2.cpp / pto_orchestrator.cpp / pto_scheduler.cpp * (ops table, scope/submit/dispatch business logic, profiling) stay in their * original files and the aicpu build only. @@ -104,6 +104,20 @@ bool PTO2SchedulerState::RingSchedState::init_data_from_layout(void *sm_dev_base return true; } +void PTO2SchedulerState::RingSchedState::reset_for_reuse( + void *sm_dev_base, int32_t ring_id, std::atomic *orch_err +) { + ring = pto2_sm_layout::ring_header_addr(sm_dev_base, ring_id); + last_task_alive = 0; + advance_lock.store(0, std::memory_order_relaxed); + dep_deadlock_reported = false; + dep_pool.reset_for_reuse(orch_err); +#if PTO2_PROFILING + dep_pool_snapshot_tail.store(1, std::memory_order_relaxed); + dep_pool_snapshot_top.store(1, std::memory_order_relaxed); +#endif +} + void PTO2SchedulerState::RingSchedState::destroy() { ring = nullptr; } PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena, int32_t dep_pool_capacity) { @@ -184,6 +198,34 @@ bool PTO2SchedulerState::init_data_from_layout( return true; } +void PTO2SchedulerState::reset_for_reuse(const PTO2SchedulerLayout &layout, void *sm_dev_base) { + PTO2SchedulerState *sched = this; + sched->sm_header = reinterpret_cast(sm_dev_base); +#if PTO2_SCHED_PROFILING + sched->tasks_completed.store(0, std::memory_order_relaxed); + sched->tasks_consumed.store(0, std::memory_order_relaxed); +#endif + + auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + sched->ring_sched_states[r].reset_for_reuse(sm_dev_base, r, orch_err); + } + + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + sched->ready_queues[i].reset_for_reuse(); + } + sched->dummy_ready_queue.reset_for_reuse(); + + sched->wiring.queue.reset_for_reuse(); + sched->wiring.batch_count = 0; + sched->wiring.batch_index = 0; + sched->wiring.backoff_counter = 0; + sched->wiring.orch_needs_drain.store(false, std::memory_order_relaxed); + sched->wiring.producer_blocked.store(0, std::memory_order_relaxed); + sched->async_wait_list.reset_for_reuse(); + (void)layout; +} + void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, DeviceArena &arena) { PTO2SchedulerState *sched = this; for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { @@ -340,6 +382,66 @@ bool PTO2OrchestratorState::init_data_from_layout( return true; } +bool PTO2OrchestratorState::reset_for_reuse( + const PTO2OrchestratorLayout &layout, void *sm_dev_base, void *gm_heap, + const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH] +) { + auto *orch = this; + orch->sm_header = reinterpret_cast(sm_dev_base); + orch->gm_heap_base = gm_heap; + uint64_t total_heap_size = 0; + if (!sum_ring_heap_sizes(heap_sizes, &total_heap_size)) { + return false; + } + orch->gm_heap_size = total_heap_size; + orch->fatal = false; + orch->inline_completed_tasks = 0; + + uint32_t next_epoch = orch->fanin_seen_current_epoch + 1; + if (next_epoch == 0) { + next_epoch = 1; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + memset( + orch->fanin_seen_epoch[r], 0, + static_cast(layout.tensor_map.task_window_sizes[r]) * sizeof(uint32_t) + ); + } + } + orch->fanin_seen_current_epoch = next_epoch; + + auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); + uint64_t heap_offset = 0; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + void *ring_heap_base = reinterpret_cast(gm_heap) + heap_offset; + auto *task_descs_dev = pto2_sm_layout::ring_task_descriptors_addr(sm_dev_base, task_window_sizes, r); + auto *slot_states_dev = pto2_sm_layout::ring_slot_states_addr(sm_dev_base, task_window_sizes, r); + auto *cur_idx_dev = pto2_sm_layout::ring_current_task_index_addr(sm_dev_base, r); + auto *last_alive_dev = pto2_sm_layout::ring_last_task_alive_addr(sm_dev_base, r); + + orch->rings[r].task_allocator.init( + task_descs_dev, static_cast(task_window_sizes[r]), cur_idx_dev, last_alive_dev, ring_heap_base, + heap_sizes[r], orch_err, slot_states_dev, 0, static_cast(r) + ); + heap_offset += heap_sizes[r]; + orch->rings[r].fanin_pool.reset_for_reuse(orch_err); + } + + orch->tensor_map.reset_for_reuse(layout.tensor_map); + orch->scope_tasks_size = 0; + orch->scope_tasks_capacity = layout.scope_tasks_cap; + orch->scope_stack_top = -1; + orch->scope_stack_capacity = layout.scope_stack_capacity; + orch->manual_begin_depth = PTO2_MAX_SCOPE_DEPTH; + orch->total_cluster_count = 0; + orch->total_aiv_count = 0; +#if PTO2_PROFILING + orch->tasks_submitted = 0; + orch->buffers_allocated = 0; + orch->bytes_allocated = 0; +#endif + return true; +} + void PTO2OrchestratorState::wire_arena_pointers( const PTO2OrchestratorLayout &layout, DeviceArena &arena, PTO2SchedulerState *scheduler_arg ) { @@ -463,6 +565,32 @@ void runtime_wire_arena_pointers(DeviceArena &arena, const PTO2RuntimeArenaLayou rt->scheduler.wire_arena_pointers(layout.offsets.sched, arena); } +bool runtime_reset_for_reuse(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt) { + (void)arena; + if (rt == nullptr) { + return false; + } + + rt->pending_scope_mode = PTO2ScopeMode::AUTO; + rt->total_cycles = 0; + rt->gm_heap_owned = false; + + uint64_t total_heap_size = 0; + if (!sum_ring_heap_sizes(layout.sizing.heap_sizes, &total_heap_size)) { + return false; + } + rt->gm_heap_size = total_heap_size; + + if (!rt->orchestrator.reset_for_reuse( + layout.offsets.orch, rt->sm_handle->sm_base, rt->gm_heap, layout.sizing.heap_sizes, + layout.sizing.task_window_sizes + )) { + return false; + } + rt->scheduler.reset_for_reuse(layout.offsets.sched, rt->sm_handle->sm_base); + return true; +} + void runtime_destroy(PTO2Runtime *rt, DeviceArena & /*arena*/) { // Arena buffer is pooled across runs by DeviceRunner — never freed here. if (!rt) return; diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp index b99c672333..fb22bb4d2b 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp @@ -64,6 +64,8 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout( layout.off_buckets = arena.reserve( static_cast(new_num_buckets) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) ); + layout.off_bucket_epochs = + arena.reserve(static_cast(new_num_buckets) * sizeof(uint32_t), alignof(uint32_t)); layout.off_entry_pool = arena.reserve(static_cast(new_pool_size) * sizeof(PTO2TensorMapEntry), alignof(PTO2TensorMapEntry)); layout.off_free_entry_list = @@ -72,6 +74,8 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout( layout.off_task_entry_heads[r] = arena.reserve( static_cast(new_task_window_sizes[r]) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) ); + layout.off_task_entry_head_epochs[r] = + arena.reserve(static_cast(new_task_window_sizes[r]) * sizeof(uint32_t), alignof(uint32_t)); } return layout; } @@ -88,12 +92,14 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev // Address arena regions for data writes; do not store these in struct // fields (wire_arena_pointers does that). auto *buckets_arena = static_cast(arena.region_ptr(layout.off_buckets)); + auto *bucket_epochs_arena = static_cast(arena.region_ptr(layout.off_bucket_epochs)); auto *entry_pool_arena = static_cast(arena.region_ptr(layout.off_entry_pool)); auto *free_list_arena = static_cast(arena.region_ptr(layout.off_free_entry_list)); // buckets[]: empty == nullptr. for (int32_t i = 0; i < num_buckets; i++) { buckets_arena[i] = nullptr; + bucket_epochs_arena[i] = 0; } // entry_pool: zero-init equivalent to the previous calloc(entry_pool, ...). @@ -118,8 +124,10 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { auto *heads_arena = static_cast(arena.region_ptr(layout.off_task_entry_heads[r])); + auto *head_epochs_arena = static_cast(arena.region_ptr(layout.off_task_entry_head_epochs[r])); for (int32_t i = 0; i < layout.task_window_sizes[r]; i++) { heads_arena[i] = nullptr; + head_epochs_arena[i] = 0; } task_window_sizes[r] = layout.task_window_sizes[r]; last_task_alives[r] = 0; @@ -129,12 +137,35 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev return true; } +void PTO2TensorMap::reset_for_reuse(const PTO2TensorMapLayout &layout) { + num_buckets = layout.num_buckets; + pool_size = layout.pool_size; + next_entry_idx = 0; + free_num = 0; + current_epoch++; + if (current_epoch == 0) { + current_epoch = 1; + memset(bucket_epochs, 0, static_cast(layout.num_buckets) * sizeof(uint32_t)); + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + memset(task_entry_head_epochs[r], 0, static_cast(layout.task_window_sizes[r]) * sizeof(uint32_t)); + } + } + + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + task_window_sizes[r] = layout.task_window_sizes[r]; + last_task_alives[r] = 0; + last_cleanup[r] = 0; + } +} + void PTO2TensorMap::wire_arena_pointers(const PTO2TensorMapLayout &layout, DeviceArena &arena) { buckets = static_cast(arena.region_ptr(layout.off_buckets)); + bucket_epochs = static_cast(arena.region_ptr(layout.off_bucket_epochs)); entry_pool = static_cast(arena.region_ptr(layout.off_entry_pool)); free_entry_list = static_cast(arena.region_ptr(layout.off_free_entry_list)); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { task_entry_heads[r] = static_cast(arena.region_ptr(layout.off_task_entry_heads[r])); + task_entry_head_epochs[r] = static_cast(arena.region_ptr(layout.off_task_entry_head_epochs[r])); } } @@ -143,10 +174,12 @@ void PTO2TensorMap::destroy() { // stray post-destroy access trips a nullptr dereference instead of reading // a recycled allocation. buckets = nullptr; + bucket_epochs = nullptr; entry_pool = nullptr; free_entry_list = nullptr; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { task_entry_heads[r] = nullptr; + task_entry_head_epochs[r] = nullptr; } }