From dc6befb549f03e3487dfd140d92e888b9a6bed47 Mon Sep 17 00:00:00 2001
From: kilinchange <kilinchange@163.com>
Date: Fri, 29 Dec 2023 17:27:36 +0800
Subject: [PATCH] fix: fix re-dataMalloc for weight tensor and use of naive
 allocator

---
 include/core/tensor.h                     | 12 ++++++++++--
 pyinfinitensor/src/pyinfinitensor/onnx.py |  9 +++++----
 src/core/graph.cc                         |  5 ++++-
 3 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/include/core/tensor.h b/include/core/tensor.h
index cb09261a..68d073c1 100644
--- a/include/core/tensor.h
+++ b/include/core/tensor.h
@@ -42,8 +42,16 @@ class TensorObj : public TensorBaseObj {
     bool isOutput() const { return tensorType == TensorType::output; }
     bool isOthers() const { return tensorType == TensorType::others; }
     void setWeight() { tensorType = TensorType::weight; }
-    void setInput() { tensorType = TensorType::input; }
-    void setOutput() { tensorType = TensorType::output; }
+    void setInput() {
+        if (!this->isWeight()) {
+            tensorType = TensorType::input;
+        }
+    }
+    void setOutput() {
+        if (!this->isWeight()) {
+            tensorType = TensorType::output;
+        }
+    }
     string tensorTypeToString() const {
         switch (tensorType) {
         case TensorType::weight:
diff --git a/pyinfinitensor/src/pyinfinitensor/onnx.py b/pyinfinitensor/src/pyinfinitensor/onnx.py
index 4a062ec7..c15feed5 100644
--- a/pyinfinitensor/src/pyinfinitensor/onnx.py
+++ b/pyinfinitensor/src/pyinfinitensor/onnx.py
@@ -37,7 +37,7 @@ class OnnxStub:
     It can be generated from an Onnx model object.
     """
 
-    def __init__(self, model: ModelProto, runtime):
+    def __init__(self, model: ModelProto, runtime, use_naive_allocator: bool = False):
         # We use some user-defined operators for distributed inference
         try:
             # onnx simplifier performs inplace simplify
@@ -54,6 +54,7 @@ class OnnxStub:
         self.tensors: Dict[str, backend.Tensor] = {}
         self.tensor_node_map: Dict[str, str] = {}
         self.initializer: Dict[int, TensorProto] = {}
+        self.use_naive_allocator: bool = use_naive_allocator
         try:
             model = infer_shapes(model)
         except:
@@ -885,7 +886,7 @@ class OnnxStub:
         ################################
         # Allocate memory space for data
         ################################
-        self.handler.data_malloc()
+        self.handler.data_malloc(self.use_naive_allocator)
 
         #################################
         # Copy in data to tensor objects
@@ -1223,7 +1224,7 @@ class OnnxStub:
         return ctx.build(name)
 
     def init(self) -> None:
-        self.handler.data_malloc()
+        self.handler.data_malloc(self.use_naive_allocator)
 
     def optimize(self) -> None:
         self.handler.optimize()
@@ -1239,7 +1240,7 @@ class OnnxStub:
             oldTensor = self.inputs[oldInput]
             self.handler.change_shape(newInput, oldTensor.fuid())
         self.handler.shape_infer()
-        self.handler.data_malloc()
+        self.handler.data_malloc(self.use_naive_allocator)
 
     def getShape(self, name: str) -> List[int]:
         if name in self.inputs:
diff --git a/src/core/graph.cc b/src/core/graph.cc
index 19831d65..ac90344a 100644
--- a/src/core/graph.cc
+++ b/src/core/graph.cc
@@ -167,7 +167,10 @@ void GraphObj::dataMalloc(bool useNaiveAllocator, size_t memPoolSize) {
         // note: behavior may not match running in non-naive mode, and it may
         // not reproduce the bug
         for (auto &tensor : tensors) {
-            tensor->dataMalloc();
+            if (!tensor->isWeight() ||
+                (tensor->isWeight() && !weightAllocated)) {
+                tensor->dataMalloc();
+            }
         }
         return;
     }