Implement memory allocation and deallocation logic

test/core/test_allocator.cc

@@ -1,74 +1,116 @@
-#include "core/graph.h"
-#include "core/kernel.h"
-#include "core/runtime.h"
-#include "operators/unary.h"
+size_t Allocator::alloc(size_t size) {
+    IT_ASSERT(this->ptr == nullptr);
+    // pad the size to the multiple of alignment
+    size = this->getAlignedSize(size);
 
-#include "test.h"
-
-namespace infini
-{
-    TEST(Allocator, testAlloc)
-    {
-        Shape shape = Shape{1, 2, 2, 3};
-        Runtime runtime = NativeCpuRuntimeObj::getInstance();
-        Tensor a = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Tensor b = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Tensor c = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Tensor d = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Allocator allocator = Allocator(runtime);
-        // allocate a->b->c
-        size_t offsetA = allocator.alloc(a->getBytes());
-        size_t offsetB = allocator.alloc(b->getBytes());
-        size_t offsetC = allocator.alloc(c->getBytes());
-        // free b, then allocate d
-        allocator.free(offsetB, b->getBytes());
-        size_t offsetD = allocator.alloc(d->getBytes());
-        // expected to be a->d->c
-        EXPECT_EQ(offsetB, offsetD);
-        ASSERT_FALSE(offsetA == 0 && offsetB == 0 && offsetC == 0 && offsetD == 0);
+    // ============== 作业 ========================
+    // TODO: 设计一个算法来分配内存，返回起始地址偏移量
+
+    // 首先尝试在空闲块中查找合适大小的块
+    // 使用最佳适配算法：找到大于等于所需大小的最小块
+    auto it = freeBlocksBySize.lower_bound(size);
+
+    if (it != freeBlocksBySize.end()) {
+        // 找到合适的空闲块
+        size_t blockSize = it->first;
+        auto& addrSet = it->second;
+
+        // 获取第一个可用的起始地址
+        size_t startAddr = *addrSet.begin();
+
+        // 从空闲块集合中移除这个块
+        addrSet.erase(startAddr);
+        if (addrSet.empty()) {
+            freeBlocksBySize.erase(it);
+        }
+
+        // 从起始地址映射中移除
+        freeBlocksByStart.erase(startAddr);
+
+        // 如果块大小大于所需大小，将剩余部分作为新的空闲块
+        if (blockSize > size) {
+            size_t remainingSize = blockSize - size;
+            size_t remainingAddr = startAddr + size;
+
+            // 添加剩余部分到空闲块集合
+            freeBlocksByStart[remainingAddr] = remainingSize;
+            freeBlocksBySize[remainingSize].insert(remainingAddr);
+        }
+
+        return startAddr;
     }
+
+    // 如果没有合适的空闲块，从当前偏移量分配
+    size_t startAddr = currentOffset;
+    currentOffset += size;
+
+    return startAddr;
+
+    // ============== 作业 ========================
+}
 
-    TEST(Allocator, testAllocWithEndFreeBlock)
-    {
-        Shape shape = Shape{1, 2, 2, 3};
-        Runtime runtime = NativeCpuRuntimeObj::getInstance();
-        Tensor a = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Tensor b = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Tensor c = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Tensor d =
-            make_ref<TensorObj>(Shape{2, 2, 2, 3}, DataType::Float32, runtime);
-        Allocator allocator = Allocator(runtime);
-        // allocate a->b->c
-        allocator.alloc(a->getBytes());
-        allocator.alloc(b->getBytes());
-        size_t offsetC = allocator.alloc(c->getBytes());
-        allocator.info();
-        // free c, then allocate d
-        allocator.free(offsetC, c->getBytes());
-        size_t offsetD = allocator.alloc(d->getBytes());
-        allocator.info();
-        // expected to be a->b->d, with no free block between b and c
-        EXPECT_EQ(offsetC, offsetD);
-    }
+void Allocator::free(size_t addr, size_t size) {
+    IT_ASSERT(this->ptr == nullptr);
+    size = getAlignedSize(size);
 
-    TEST(Allocator, testGetPtr)
-    {
-        Shape shape = Shape{1, 2, 2, 3};
-        Runtime runtime = NativeCpuRuntimeObj::getInstance();
-        Tensor a = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Tensor b = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Tensor c = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Tensor d = make_ref<TensorObj>(shape, DataType::Float32, runtime);
-        Allocator allocator = Allocator(runtime);
-        // allocate a->b->c->d
-        allocator.alloc(a->getBytes());
-        allocator.alloc(b->getBytes());
-        allocator.alloc(c->getBytes());
-        allocator.alloc(d->getBytes());
-        // multiple calls to the getPtr() function should return the same pointer
-        void *ptr1 = allocator.getPtr();
-        void *ptr2 = allocator.getPtr();
-        EXPECT_EQ(ptr1, ptr2);
-    }
+    // ============== 作业 ========================
+    // TODO: 设计一个算法来回收内存
+
+    // 将释放的块添加到空闲块集合
+    freeBlocksByStart[addr] = size;
+    freeBlocksBySize[size].insert(addr);
+
+    // 合并相邻的空闲块
+    mergeFreeBlocks();
+
+    // ============== 作业 ========================
+}
 
-} // namespace infini
+void Allocator::mergeFreeBlocks() {
+    if (freeBlocksByStart.empty()) return;
+
+    // 按起始地址排序后，合并相邻的空闲块
+    auto it = freeBlocksByStart.begin();
+    auto nextIt = std::next(it);
+
+    while (nextIt != freeBlocksByStart.end()) {
+        size_t currentAddr = it->first;
+        size_t currentSize = it->second;
+        size_t nextAddr = nextIt->first;
+        size_t nextSize = nextIt->second;
+
+        // 检查当前块和下一个块是否相邻
+        if (currentAddr + currentSize == nextAddr) {
+            // 合并两个块
+            size_t mergedSize = currentSize + nextSize;
+
+            // 从空闲块集合中移除原来的两个块
+            // 从大小映射中移除
+            freeBlocksBySize[currentSize].erase(currentAddr);
+            if (freeBlocksBySize[currentSize].empty()) {
+                freeBlocksBySize.erase(currentSize);
+            }
+
+            freeBlocksBySize[nextSize].erase(nextAddr);
+            if (freeBlocksBySize[nextSize].empty()) {
+                freeBlocksBySize.erase(nextSize);
+            }
+
+            // 从起始地址映射中移除
+            freeBlocksByStart.erase(nextAddr);
+
+            // 更新当前块的大小
+            it->second = mergedSize;
+
+            // 添加到大小映射
+            freeBlocksBySize[mergedSize].insert(currentAddr);
+
+            // 重新开始合并，因为合并后可能还能和下一个块合并
+            nextIt = std::next(it);
+        } else {
+            // 不合并，继续检查下一对
+            ++it;
+            ++nextIt;
+        }
+    }
+}