fix: fix re-dataMalloc for weight tensor and use of naive allocator

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:

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:

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;
     }