TinyInfiniTensor-ljc

src/core/allocator.cc

@@ -1,69 +1,121 @@
-#include "core/allocator.h"
-#include <utility>
-
-namespace infini
-{
-    Allocator::Allocator(Runtime runtime) : runtime(runtime)
-    {
-        used = 0;
-        peak = 0;
-        ptr = nullptr;
-
-        // 'alignment' defaults to sizeof(uint64_t), because it is the length of
-        // the longest data type currently supported by the DataType field of
-        // the tensor
-        alignment = sizeof(uint64_t);
-    }
-
-    Allocator::~Allocator()
-    {
-        if (this->ptr != nullptr)
-        {
-            runtime->dealloc(this->ptr);
-        }
-    }
-
-    size_t Allocator::alloc(size_t size)
-    {
-        IT_ASSERT(this->ptr == nullptr);
-        // pad the size to the multiple of alignment
-        size = this->getAlignedSize(size);
-
-        // =================================== 作业 ===================================
-        // TODO: 设计一个算法来分配内存，返回起始地址偏移量
-        // =================================== 作业 ===================================
-
-        return 0;
-    }
-
-    void Allocator::free(size_t addr, size_t size)
-    {
-        IT_ASSERT(this->ptr == nullptr);
-        size = getAlignedSize(size);
-
-        // =================================== 作业 ===================================
-        // TODO: 设计一个算法来回收内存
-        // =================================== 作业 ===================================
-    }
-
-    void *Allocator::getPtr()
-    {
-        if (this->ptr == nullptr)
-        {
-            this->ptr = runtime->alloc(this->peak);
-            printf("Allocator really alloc: %p %lu bytes\n", this->ptr, peak);
-        }
-        return this->ptr;
-    }
-
-    size_t Allocator::getAlignedSize(size_t size)
-    {
-        return ((size - 1) / this->alignment + 1) * this->alignment;
-    }
-
-    void Allocator::info()
-    {
-        std::cout << "Used memory: " << this->used
-                  << ", peak memory: " << this->peak << std::endl;
-    }
-}
+#include "core/allocator.h"
+#include <utility>
+
+namespace infini
+{
+    Allocator::Allocator(Runtime runtime) : runtime(runtime)
+    {
+        used = 0;
+        peak = 0;
+        ptr = nullptr;
+
+        // 'alignment' defaults to sizeof(uint64_t), because it is the length of
+        // the longest data type currently supported by the DataType field of
+        // the tensor
+        alignment = sizeof(uint64_t);
+    }
+
+    Allocator::~Allocator()
+    {
+        if (this->ptr != nullptr)
+        {
+            runtime->dealloc(this->ptr);
+        }
+    }
+
+    size_t Allocator::alloc(size_t size)
+    {
+        IT_ASSERT(this->ptr == nullptr);
+        // pad the size to the multiple of alignment
+        size = this->getAlignedSize(size);
+
+        // =================================== 作业 ===================================
+        // TODO: 设计一个算法来分配内存，返回起始地址偏移量
+        // =================================== 作业 ===================================
+        // First-fit: search free blocks for a suitable one
+        for (auto it = freeBlocks.begin(); it != freeBlocks.end(); ++it) {
+            if (it->second >= size) {
+                size_t addr = it->first;
+                size_t remaining = it->second - size;
+                freeBlocks.erase(it);
+                if (remaining > 0) {
+                    freeBlocks[addr + size] = remaining;
+                }
+                return addr;
+            }
+        }
+        // No suitable free block found, allocate at the end
+        size_t addr = used;
+        used += size;
+        peak = std::max(peak, used);
+        return addr;
+    }
+
+    void Allocator::free(size_t addr, size_t size)
+    {
+        IT_ASSERT(this->ptr == nullptr);
+        size = getAlignedSize(size);
+
+        // =================================== 作业 ===================================
+        // TODO: 设计一个算法来回收内存
+        // =================================== 作业 ===================================
+        // If the block is at the end of used memory, shrink used
+        if (addr + size == used) {
+            used = addr;
+            // Remove any trailing free blocks that are now past the new used
+            while (!freeBlocks.empty()) {
+                auto last = std::prev(freeBlocks.end());
+                if (last->first + last->second == addr) {
+                    addr = last->first;
+                    size += last->second;
+                    freeBlocks.erase(last);
+                } else {
+                    break;
+                }
+            }
+            used = addr;
+            return;
+        }
+
+        // Insert into free blocks
+        auto [it, inserted] = freeBlocks.insert({addr, size});
+        IT_ASSERT(inserted);
+
+        // Merge with next block if adjacent
+        auto next = std::next(it);
+        if (next != freeBlocks.end() && addr + size == next->first) {
+            it->second += next->second;
+            freeBlocks.erase(next);
+        }
+
+        // Merge with previous block if adjacent
+        if (it != freeBlocks.begin()) {
+            auto prev = std::prev(it);
+            if (prev->first + prev->second == addr) {
+                prev->second += it->second;
+                freeBlocks.erase(it);
+            }
+        }
+    }
+
+    void *Allocator::getPtr()
+    {
+        if (this->ptr == nullptr)
+        {
+            this->ptr = runtime->alloc(this->peak);
+            printf("Allocator really alloc: %p %lu bytes\n", this->ptr, peak);
+        }
+        return this->ptr;
+    }
+
+    size_t Allocator::getAlignedSize(size_t size)
+    {
+        return ((size - 1) / this->alignment + 1) * this->alignment;
+    }
+
+    void Allocator::info()
+    {
+        std::cout << "Used memory: " << this->used
+                  << ", peak memory: " << this->peak << std::endl;
+    }
+}

src/core/data_type.cc

@@ -1,22 +1,22 @@
-#include "core/data_type.h"
-
-namespace infini {
-// Move implementation here to avoid compile time error on some platform
-// to be consistent with onnx
-// https://github.com/onnx/onnx/blob/aeb21329122b96df1d3ef33b500a35ca140b1431/onnx/onnx.proto#L484
-const DataType DataType::Undefine(0);
-const DataType DataType::Float32(1);
-const DataType DataType::UInt8(2);
-const DataType DataType::Int8(3);
-const DataType DataType::UInt16(4);
-const DataType DataType::Int16(5);
-const DataType DataType::Int32(6);
-const DataType DataType::Int64(7);
-const DataType DataType::String(8);
-const DataType DataType::Bool(9);
-const DataType DataType::Float16(10);
-const DataType DataType::Double(11);
-const DataType DataType::UInt32(12);
-const DataType DataType::UInt64(13);
-const DataType DataType::BFloat16(16);
-} // namespace infini
+#include "core/data_type.h"
+
+namespace infini {
+// Move implementation here to avoid compile time error on some platform
+// to be consistent with onnx
+// https://github.com/onnx/onnx/blob/aeb21329122b96df1d3ef33b500a35ca140b1431/onnx/onnx.proto#L484
+const DataType DataType::Undefine(0);
+const DataType DataType::Float32(1);
+const DataType DataType::UInt8(2);
+const DataType DataType::Int8(3);
+const DataType DataType::UInt16(4);
+const DataType DataType::Int16(5);
+const DataType DataType::Int32(6);
+const DataType DataType::Int64(7);
+const DataType DataType::String(8);
+const DataType DataType::Bool(9);
+const DataType DataType::Float16(10);
+const DataType DataType::Double(11);
+const DataType DataType::UInt32(12);
+const DataType DataType::UInt64(13);
+const DataType DataType::BFloat16(16);
+} // namespace infini