fix(kernel): 没有引用的输出张量也要分配，但计算完后立即释放

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Feb 19, 2024 · e3febd9 · e3febd9
1 parent d076c20
commit e3febd9
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/04kernel/src/allocators/reusable_allocator.cpp b/src/04kernel/src/allocators/reusable_allocator.cpp
@@ -23,7 +23,6 @@ namespace refactor::kernel {
         std::vector<runtime::Edge> edges_(edges.size(), {nullptr, SIZE_MAX});
         for (auto [nodeIdx, inputs, outputs] : topology) {
             for (auto outputIdx : outputs) {
-                if (!edgeRc[outputIdx]) { continue; }
                 if (globalOutputs.contains(outputIdx)) {
                     edges_[outputIdx].stackOffset--;
                 } else {
@@ -34,6 +33,11 @@ namespace refactor::kernel {
                 auto wsSize = node.workspaceOffset;
                 calculator.free(node.workspaceOffset = calculator.alloc(wsSize), wsSize);
             }
+            for (auto outputIdx : outputs) {
+                if (!edgeRc[outputIdx]) {
+                    calculator.free(edges_[outputIdx].stackOffset, edges[outputIdx].size);
+                }
+            }
             for (auto inputIdx : inputs) {
                 ASSERT(edgeRc[inputIdx], "double free");
                 if (!--edgeRc[inputIdx]) {