Support: reuse resident prebuilt runtime arenas

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

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;

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<uint64_t> 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);

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,

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<int32_t> *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<int32_t> *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.

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

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 <typename Fn>
     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) {

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<PTO2TaskSlotState **>(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<int32_t> *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