allocator.cc

#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 策略：遍历所有空闲块，找到第一个足够大的块
        for (auto it = free_blocks.begin(); it != free_blocks.end(); ++it)
        {
            size_t block_addr = it->first;
            size_t block_size = it->second;
            
            // 情况1: 找到足够大的空闲块，直接使用
            if (block_size >= size)
            {
                // 从空闲列表中移除该块
                free_blocks.erase(it);
                
                // 如果有剩余空间，将剩余部分重新加入空闲列表
                if (block_size > size)
                {
                    free_blocks[block_addr + size] = block_size - size;
                }
                
                // 更新已使用内存
                this->used += size;
                
                return block_addr;
            }
            // 情况2: 该块在内存末尾（block_addr + block_size == peak）
            // 即使块不够大，也可以扩展使用
            else if (block_addr + block_size == this->peak)
            {
                // 从空闲列表中移除该块
                free_blocks.erase(it);
                
                // 计算还需要多少额外空间
                size_t extra_needed = size - block_size;
                
                // 扩展 peak
                this->peak += extra_needed;
                
                // 更新已使用内存
                this->used += size;
                
                return block_addr;
            }
        }
        
        // 没有找到任何空闲块，在内存末尾分配新空间
        size_t block_addr = this->peak;
        this->used += size;
        this->peak += size;
        
        return block_addr;
    }

    void Allocator::free(size_t addr, size_t size)
    {
        IT_ASSERT(this->ptr == nullptr);
        size = getAlignedSize(size);
        
        // =================================== 作业 ===================================
        // TODO: 设计一个算法来回收内存
        // =================================== 作业 ===================================
        
        // 更新已使用内存
        this->used -= size;
        
        // 将释放的块加入空闲列表
        free_blocks[addr] = size;
        
        // 获取当前释放块的迭代器
        auto current = free_blocks.find(addr);
        
        // 尝试与后面的块合并
        auto next = std::next(current);
        if (next != free_blocks.end())
        {
            // 如果当前块的结尾恰好是下一个块的开始
            if (current->first + current->second == next->first)
            {
                // 合并：扩展当前块的大小
                current->second += next->second;
                // 删除下一个块
                free_blocks.erase(next);
            }
        }
        
        // 尝试与前面的块合并
        if (current != free_blocks.begin())
        {
            auto prev = std::prev(current);
            // 如果前一个块的结尾恰好是当前块的开始
            if (prev->first + prev->second == current->first)
            {
                // 合并：扩展前一个块的大小
                prev->second += current->second;
                // 删除当前块
                free_blocks.erase(current);
            }
        }
    }

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