feature: 补充作业一

LearningInfiniTensor · Jul 31, 2024 · 4d575de · 4d575de
1 parent c80e1d9
commit 4d575de
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 57 deletions.
diff --git a/docs/训练营作业介绍.md b/docs/训练营作业介绍.md
@@ -66,6 +66,27 @@
     }
 ````
 
+### step3 计算图的内存分配部分
+
+需要实现的代码块位置：`src/core/graph.cc`
+
+完善计算图的内存分配部分：
+
+````c++
+    void GraphObj::dataMalloc()
+    {
+        // topological sorting first
+        IT_ASSERT(topo_sort() == true);
+
+        // =================================== 作业 ===================================
+        // TODO：利用 allocator 给计算图分配内存
+        // HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存
+        // =================================== 作业 ===================================
+
+        allocator.info();
+    }
+````
+
 ## 作业二：transpose 算子形状推导
 
 难度：⭐

diff --git a/src/core/graph.cc b/src/core/graph.cc
@@ -144,64 +144,11 @@ namespace infini
     {
         // topological sorting first
         IT_ASSERT(topo_sort() == true);
-        // count the number of times all tensors are used
-        std::unordered_map<TensorObj *, size_t> tensorToRefCount;
-        // record the memory address offsets of all tensors to be allocated
-        std::unordered_map<TensorObj *, size_t> tensorToOffset;
 
-        // record all constant tensors, including weight tensors and input tensors
-        std::unordered_set<TensorObj *> constTensor;
-        for (auto &tensor : tensors)
-        {
-            if (tensor.get()->getSource() == nullptr)
-            {
-                // allocate memory for all constant tensors first, and this memory
-                // will not be reused later
-                constTensor.insert(tensor.get());
-                tensorToOffset[tensor.get()] = allocator.alloc(tensor->getBytes());
-            }
-            else
-            {
-                tensorToRefCount[tensor.get()] = tensor->getTargets().size();
-            }
-        }
-        // traverse in topological order and simulate memory allocation
-        for (auto &op : ops)
-        {
-            // memory should be allocated for the output first
-            auto outputs = op->getOutputs();
-            for (auto &tensor : outputs)
-            {
-                tensorToOffset[tensor.get()] = allocator.alloc(tensor->getBytes());
-            }
-            auto inputs = op->getInputs();
-            for (auto &tensor : inputs)
-            {
-                if (constTensor.find(tensor.get()) == constTensor.end())
-                {
-                    auto tensorIter = tensorToRefCount.find(tensor.get());
-                    IT_ASSERT(tensorIter != tensorToRefCount.end());
-                    tensorToRefCount[tensor.get()] -= 1;
-                    if (tensorToRefCount[tensor.get()] == 0)
-                    {
-                        // indicate that this tensor will no longer be used and
-                        // perform memory free
-                        tensorToRefCount.erase(tensor.get());
-                        allocator.free(tensorToOffset[tensor.get()],
-                                       tensor->getBytes());
-                    }
-                }
-            }
-        }
-
-        // perform actual memory allocation
-        for (auto &tensor : tensors)
-        {
-            IT_ASSERT(tensorToOffset.find(tensor.get()) != tensorToOffset.end());
-            tensor->setDataBlob(make_ref<BlobObj>(
-                tensor->runtime, static_cast<uint8_t *>(allocator.getPtr()) +
-                                     tensorToOffset[tensor.get()]));
-        }
+        // =================================== 作业 ===================================
+        // TODO：利用 allocator 给计算图分配内存
+        // HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存
+        // =================================== 作业 ===================================
 
         allocator.info();
     }