hw-native-sys · yanghaoran29 · Jul 29, 2026
diff --git a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md
@@ -343,7 +343,7 @@ This guarantees lookup only traverses valid entries — O(valid_entries_in_bucke
 
 Every time the orchestrator submits a task (Step 0 of `PTO2OrchestratorState::submit_task`), it calls `PTO2TensorMap::sync_tensormap`. When `last_task_alive` has advanced by more than `PTO2_TENSORMAP_CLEANUP_INTERVAL` (default 64) tasks since the last cleanup, `PTO2TensorMap::cleanup_retired` runs:
 
-This uses the **per-task entry chain** (`task_entry_head[task_slot]`) — each task's entries are doubly-linked together at insert time via `next_in_task`/`prev_in_task`, allowing O(entries_per_task) cleanup without scanning the entire pool or all buckets. Freed entries are returned to `free_entry_list` for immediate reuse.
+This uses the **per-task entry chain** (`task_entry_head[task_slot]`) — each task's entries are doubly-linked together at insert time via `next_in_task`/`prev_in_task`. A slot's chain can hold more than one task's entries: a task at `local_id + N * window` reuses the slot and prepends to the chain already there, and cleanup can lag that reuse. Cleanup therefore walks the chain and frees only the entries whose `producer_task_id` matches the retiring task, unlinking each and leaving the rest linked — O(entries_in_slot), with no scan of the entire pool or all buckets. Freed entries are returned to `free_entry_list` for immediate reuse.
 
 **Layer 3 — Back-Pressure on Pool Exhaustion** (blocking):
 

diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_tensormap.h b/src/a2a3/runtime/host_build_graph/runtime/pto_tensormap.h
@@ -588,19 +588,27 @@ struct PTO2TensorMap {
         // Iterate through retired tasks 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_size - 1);
+            // A slot's task chain may also hold entries from a newer task that
+            // reused the slot (local_id + N * window) before this cleanup ran.
+            // Free only entries produced by the retiring local_id, unlinking
+            // each from the chain; entries from other tasks stay linked.
+            PTO2TaskId retired_task = PTO2TaskId::make(0, static_cast<uint32_t>(local_id));
             PTO2TensorMapEntry *cur_entry = task_entry_heads[task_slot];
-
             while (cur_entry != nullptr) {
-                PTO2TensorMapEntry *next_entry = cur_entry->next_in_task;  // Save before clearing
-                // Only remove if this entry belongs to the retiring task
-                // (slot may have been reused by a newer task)
-                debug_assert(cur_entry->producer_task_id == PTO2TaskId::make(0, static_cast<uint32_t>(local_id)));
-                free_entry(*cur_entry);
+                PTO2TensorMapEntry *next_entry = cur_entry->next_in_task;  // free_entry clears it
+                if (cur_entry->producer_task_id == retired_task) {
+                    if (cur_entry->prev_in_task != nullptr) {
+                        cur_entry->prev_in_task->next_in_task = next_entry;
+                    } else {
+                        task_entry_heads[task_slot] = next_entry;
+                    }
+                    if (next_entry != nullptr) {
+                        next_entry->prev_in_task = cur_entry->prev_in_task;
+                    }
+                    free_entry(*cur_entry);
+                }
                 cur_entry = next_entry;
             }
-
-            // Clear task's entry head (slot will be reused by local_id + task_window_size)
-            task_entry_heads[task_slot] = nullptr;
         }
     }
 

diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md
@@ -338,7 +338,7 @@ This guarantees lookup only traverses valid entries — O(valid_entries_in_bucke
 
 Every time the orchestrator submits a task (Step 0 of `PTO2OrchestratorState::submit_task`), it calls `PTO2TensorMap::sync_tensormap`. When `last_task_alive` has advanced by more than `PTO2_TENSORMAP_CLEANUP_INTERVAL` (default 64) tasks since the last cleanup, `PTO2TensorMap::cleanup_retired` runs:
 
-This uses the **per-task entry chain** (`task_entry_head[task_slot]`) — each task's entries are doubly-linked together at insert time via `next_in_task`/`prev_in_task`, allowing O(entries_per_task) cleanup without scanning the entire pool or all buckets. Freed entries are returned to `free_entry_list` for immediate reuse.
+This uses the **per-task entry chain** (`task_entry_head[task_slot]`) — each task's entries are doubly-linked together at insert time via `next_in_task`/`prev_in_task`. A slot's chain can hold more than one task's entries: a task at `local_id + N * window` reuses the slot and prepends to the chain already there, and cleanup can lag that reuse. Cleanup therefore walks the chain and frees only the entries whose `producer_task_id` matches the retiring task, unlinking each and leaving the rest linked — O(entries_in_slot), with no scan of the entire pool or all buckets. Freed entries are returned to `free_entry_list` for immediate reuse.
 
 **Layer 3 — Back-Pressure on Pool Exhaustion** (blocking):
 

diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h
@@ -609,22 +609,27 @@ struct PTO2TensorMap {
             if (task_entry_head_epochs[ring_id][task_slot] != current_epoch) {
                 continue;
             }
+            // A slot's task chain may also hold entries from a newer task that
+            // reused the slot (local_id + N * window) before this cleanup ran.
+            // Free only entries produced by the retiring local_id, unlinking
+            // each from the chain; entries from other tasks stay linked.
+            PTO2TaskId retired_task = PTO2TaskId::make(static_cast<uint8_t>(ring_id), static_cast<uint32_t>(local_id));
             PTO2TensorMapEntry *cur_entry = task_entry_heads[ring_id][task_slot];
-
             while (cur_entry != nullptr) {
-                PTO2TensorMapEntry *next_entry = cur_entry->next_in_task;  // Save before clearing
-                // Only remove if this entry belongs to the retiring task
-                // (slot may have been reused by a newer task)
-                debug_assert(
-                    cur_entry->producer_task_id ==
-                    PTO2TaskId::make(static_cast<uint8_t>(ring_id), static_cast<uint32_t>(local_id))
-                );
-                free_entry(*cur_entry);
+                PTO2TensorMapEntry *next_entry = cur_entry->next_in_task;  // free_entry clears it
+                if (cur_entry->producer_task_id == retired_task) {
+                    if (cur_entry->prev_in_task != nullptr) {
+                        cur_entry->prev_in_task->next_in_task = next_entry;
+                    } else {
+                        task_entry_heads[ring_id][task_slot] = next_entry;
+                    }
+                    if (next_entry != nullptr) {
+                        next_entry->prev_in_task = cur_entry->prev_in_task;
+                    }
+                    free_entry(*cur_entry);
+                }
                 cur_entry = next_entry;
             }
-
-            // Clear task's entry head (slot will be reused by local_id + task_window_sizes[ring_id])
-            task_entry_heads[ring_id][task_slot] = nullptr;
         }
     }
 

diff --git a/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md
@@ -335,7 +335,7 @@ This guarantees lookup only traverses valid entries — O(valid_entries_in_bucke
 
 Every time the orchestrator submits a task (Step 0 of `PTO2OrchestratorState::submit_task`), it calls `PTO2TensorMap::sync_tensormap`. When `last_task_alive` has advanced by more than `PTO2_TENSORMAP_CLEANUP_INTERVAL` (default 64) tasks since the last cleanup, `PTO2TensorMap::cleanup_retired` runs:
 
-This uses the **per-task entry chain** (`task_entry_head[task_slot]`) — each task's entries are doubly-linked together at insert time via `next_in_task`/`prev_in_task`, allowing O(entries_per_task) cleanup without scanning the entire pool or all buckets. Freed entries are returned to `free_entry_list` for immediate reuse.
+This uses the **per-task entry chain** (`task_entry_head[task_slot]`) — each task's entries are doubly-linked together at insert time via `next_in_task`/`prev_in_task`. A slot's chain can hold more than one task's entries: a task at `local_id + N * window` reuses the slot and prepends to the chain already there, and cleanup can lag that reuse. Cleanup therefore walks the chain and frees only the entries whose `producer_task_id` matches the retiring task, unlinking each and leaving the rest linked — O(entries_in_slot), with no scan of the entire pool or all buckets. Freed entries are returned to `free_entry_list` for immediate reuse.
 
 **Layer 3 — Back-Pressure on Pool Exhaustion** (blocking):
 

diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h
@@ -609,22 +609,27 @@ struct PTO2TensorMap {
             if (task_entry_head_epochs[ring_id][task_slot] != current_epoch) {
                 continue;
             }
+            // A slot's task chain may also hold entries from a newer task that
+            // reused the slot (local_id + N * window) before this cleanup ran.
+            // Free only entries produced by the retiring local_id, unlinking
+            // each from the chain; entries from other tasks stay linked.
+            PTO2TaskId retired_task = PTO2TaskId::make(static_cast<uint8_t>(ring_id), static_cast<uint32_t>(local_id));
             PTO2TensorMapEntry *cur_entry = task_entry_heads[ring_id][task_slot];
-
             while (cur_entry != nullptr) {
-                PTO2TensorMapEntry *next_entry = cur_entry->next_in_task;  // Save before clearing
-                // Only remove if this entry belongs to the retiring task
-                // (slot may have been reused by a newer task)
-                debug_assert(
-                    cur_entry->producer_task_id ==
-                    PTO2TaskId::make(static_cast<uint8_t>(ring_id), static_cast<uint32_t>(local_id))
-                );
-                free_entry(*cur_entry);
+                PTO2TensorMapEntry *next_entry = cur_entry->next_in_task;  // free_entry clears it
+                if (cur_entry->producer_task_id == retired_task) {
+                    if (cur_entry->prev_in_task != nullptr) {
+                        cur_entry->prev_in_task->next_in_task = next_entry;
+                    } else {
+                        task_entry_heads[ring_id][task_slot] = next_entry;
+                    }
+                    if (next_entry != nullptr) {
+                        next_entry->prev_in_task = cur_entry->prev_in_task;
+                    }
+                    free_entry(*cur_entry);
+                }
                 cur_entry = next_entry;
             }
-
-            // Clear task's entry head (slot will be reused by local_id + task_window_sizes[ring_id])
-            task_entry_heads[ring_id][task_slot] = nullptr;
         }
     }
 

diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt
@@ -653,6 +653,12 @@ add_a2a3_test(test_task_timing_slots a2a3/test_task_timing_slots.cpp)
 # PTO2 runtime-linked tests
 add_a2a3_runtime_test(test_task_allocator   a2a3/test_task_allocator.cpp)
 add_a2a3_hbg_runtime_test(test_hbg_task_allocator a2a3/test_task_allocator.cpp)
+add_a2a3_hbg_runtime_test(test_hbg_tensormap a2a3/test_hbg_tensormap.cpp)
+# PTO2TensorMap's out-of-line members (reserve_layout / init / valid_count) live
+# in this .cpp; no other add_a2a3_hbg_runtime_test target needs them.
+target_sources(test_hbg_tensormap PRIVATE
+    ${CMAKE_SOURCE_DIR}/../../../src/a2a3/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp
+)
 add_a2a3_runtime_test(test_dep_list_pool    a2a3/test_dep_list_pool.cpp)
 add_a2a3_runtime_test(test_scheduler_state  a2a3/test_scheduler_state.cpp)
 add_a2a3_runtime_test(test_task_state       a2a3/test_task_state.cpp)

diff --git a/tests/ut/cpp/a2a3/test_hbg_tensormap.cpp b/tests/ut/cpp/a2a3/test_hbg_tensormap.cpp
@@ -0,0 +1,125 @@
+/*
+ * Copyright (c) PyPTO Contributors.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ * -----------------------------------------------------------------------------------------------------------
+ */
+/**
+ * cleanup_retired tests for the host_build_graph copy of PTO2TensorMap.
+ *
+ * This is a distinct type from the tensormap_and_ringbuffer PTO2TensorMap
+ * covered by a2a3/test_tensormap.cpp: single-ring, no entry epochs. Only the
+ * per-task entry reclamation path is covered here — the hash / overlap /
+ * lazy-invalidation surface is shared logic already exercised by that suite.
+ */
+
+#include <gtest/gtest.h>
+
+#include <vector>
+
+#include "utils/device_arena.h"
+#include "pto_tensormap.h"
+
+namespace {
+
+struct TestLookupResult {
+    struct Entry {
+        PTO2TensorMapEntry *entry;
+        OverlapStatus overlap_status;
+    };
+    std::vector<Entry> entries;
+    int count = 0;
+};
+
+void run_lookup(PTO2TensorMap &tmap, const Tensor &tensor, TestLookupResult &out) {
+    tmap.lookup(tensor, [&](PTO2TensorMapEntry &e, OverlapStatus s) -> bool {
+        out.entries.push_back({&e, s});
+        out.count++;
+        return true;
+    });
+}
+
+Tensor make_test_tensor(uint64_t addr, uint32_t shape0) {
+    uint32_t shapes[MAX_TENSOR_DIMS] = {shape0};
+    return make_tensor_external(reinterpret_cast<void *>(addr), shapes, 1, DataType::FLOAT32, false, 0);
+}
+
+class HbgTensorMapTest : public ::testing::Test {
+protected:
+    static constexpr int32_t NUM_BUCKETS = 16;
+    static constexpr int32_t POOL_SIZE = 64;
+    static constexpr int32_t WINDOW_SIZE = 32;
+
+    PTO2TensorMap tmap{};
+    DeviceArena arena;
+
+    void SetUp() override {
+        auto layout = PTO2TensorMap::reserve_layout(arena, NUM_BUCKETS, POOL_SIZE, WINDOW_SIZE);
+        ASSERT_NE(arena.commit(), nullptr);
+        ASSERT_TRUE(tmap.init_data_from_layout(layout, arena));
+        tmap.wire_arena_pointers(layout, arena);
+    }
+
+    void TearDown() override {
+        tmap.destroy();
+        arena.release();
+    }
+};
+
+TEST_F(HbgTensorMapTest, CleanupRetiredRemovesEntriesForRetiredTasks) {
+    Tensor t = make_test_tensor(0x1000, 256);
+    tmap.insert(t, PTO2TaskId::make(0, 0));
+    tmap.insert(t, PTO2TaskId::make(0, 1));
+    tmap.insert(t, PTO2TaskId::make(0, 2));
+    EXPECT_EQ(tmap.valid_count(), 3);
+
+    tmap.cleanup_retired(0, 2);
+
+    EXPECT_EQ(tmap.valid_count(), 1);
+    TestLookupResult result;
+    run_lookup(tmap, t, result);
+    ASSERT_EQ(result.count, 1);
+    EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, 2));
+}
+
+TEST_F(HbgTensorMapTest, CleanupRetiredFreesEveryOutputOfOneTask) {
+    Tensor t1 = make_test_tensor(0x1000, 256);
+    Tensor t2 = make_test_tensor(0x2000, 128);
+    PTO2TaskId tid = PTO2TaskId::make(0, 5);
+
+    tmap.insert(t1, tid);
+    tmap.insert(t2, tid);
+    EXPECT_EQ(tmap.valid_count(), 2);
+
+    tmap.cleanup_retired(5, 6);
+    EXPECT_EQ(tmap.valid_count(), 0);
+    EXPECT_EQ(tmap.free_num, 2);
+}
+
+// A later task that reuses a slot (local_id + WINDOW_SIZE) before cleanup has
+// run on the earlier task chains its entries under the same task_entry_head.
+// cleanup_retired retiring only the earlier task must free that earlier task's
+// entries alone and leave the still-live (later) task's entries intact.
+TEST_F(HbgTensorMapTest, CleanupRetiredSparesLaterTaskReusingSlot) {
+    Tensor t = make_test_tensor(0x1000, 256);
+    // Task 0 and task 0 + WINDOW_SIZE share slot 0 (local_id & (WINDOW_SIZE-1)).
+    tmap.insert(t, PTO2TaskId::make(0, 0));
+    tmap.insert(t, PTO2TaskId::make(0, WINDOW_SIZE));
+    ASSERT_EQ(tmap.valid_count(), 2);
+
+    // Retire only task 0.
+    tmap.cleanup_retired(0, 1);
+
+    // Only task 0's entry is freed; task WINDOW_SIZE's entry survives.
+    EXPECT_EQ(tmap.valid_count(), 1);
+    TestLookupResult result;
+    run_lookup(tmap, t, result);
+    ASSERT_EQ(result.count, 1);
+    EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, WINDOW_SIZE));
+}
+
+}  // namespace
diff --git a/tests/ut/cpp/a2a3/test_tensormap.cpp b/tests/ut/cpp/a2a3/test_tensormap.cpp
@@ -404,6 +404,28 @@ TEST_F(TensorMapTest, CleanupRetiredFreesEntriesToPool) {
     EXPECT_EQ(tmap.next_entry_idx, 1) << "Should reuse freed entry, not allocate new";
 }
 
+// A later task that reuses a slot (local_id + WINDOW_SIZE) before cleanup has
+// run on the earlier task chains its entries under the same task_entry_head.
+// cleanup_retired retiring only the earlier task must free that earlier task's
+// entries alone and leave the still-live (later) task's entries intact.
+TEST_F(TensorMapTest, CleanupRetiredSparesLaterTaskReusingSlot) {
+    Tensor t = make_test_tensor(0x1000, 256);
+    // Task 0 and task 0 + WINDOW_SIZE share slot 0 (local_id & (WINDOW_SIZE-1)).
+    tmap.insert(t, PTO2TaskId::make(0, 0));
+    tmap.insert(t, PTO2TaskId::make(0, WINDOW_SIZE));
+    ASSERT_EQ(tmap.valid_count(), 2);
+
+    // Retire only task 0.
+    tmap.cleanup_retired(0, 0, 1);
+
+    // Only task 0's entry is freed; task WINDOW_SIZE's entry survives.
+    EXPECT_EQ(tmap.valid_count(), 1);
+    TestLookupResult result;
+    run_lookup(tmap, t, result);
+    ASSERT_EQ(result.count, 1);
+    EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, WINDOW_SIZE));
+}
+
 // =============================================================================
 // Multi-ring isolation
 // =============================================================================

diff --git a/tests/ut/cpp/a5/test_tensormap.cpp b/tests/ut/cpp/a5/test_tensormap.cpp
@@ -405,6 +405,28 @@ TEST_F(TensorMapTest, CleanupRetiredFreesEntriesToPool) {
     EXPECT_EQ(tmap.next_entry_idx, 1) << "Should reuse freed entry, not allocate new";
 }
 
+// A later task that reuses a slot (local_id + WINDOW_SIZE) before cleanup has
+// run on the earlier task chains its entries under the same task_entry_head.
+// cleanup_retired retiring only the earlier task must free that earlier task's
+// entries alone and leave the still-live (later) task's entries intact.
+TEST_F(TensorMapTest, CleanupRetiredSparesLaterTaskReusingSlot) {
+    Tensor t = make_test_tensor(0x1000, 256);
+    // Task 0 and task 0 + WINDOW_SIZE share slot 0 (local_id & (WINDOW_SIZE-1)).
+    tmap.insert(t, PTO2TaskId::make(0, 0));
+    tmap.insert(t, PTO2TaskId::make(0, WINDOW_SIZE));
+    ASSERT_EQ(tmap.valid_count(), 2);
+
+    // Retire only task 0.
+    tmap.cleanup_retired(0, 0, 1);
+
+    // Only task 0's entry is freed; task WINDOW_SIZE's entry survives.
+    EXPECT_EQ(tmap.valid_count(), 1);
+    TestLookupResult result;
+    run_lookup(tmap, t, result);
+    ASSERT_EQ(result.count, 1);
+    EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, WINDOW_SIZE));
+}
+
 // =============================================================================
 // Multi-ring isolation
 // =============================================================================