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: 设计一个算法来分配内存，返回起始地址偏移量
        // =================================== 作业 ===================================

    //算法1：first fit 缺点：碎片化严重。
        //我们使用的map来保存block,其中key是起始地址，value是block的大小
        //我们遍历整个map，找到第一个value大于size的块
        //将这个块从map中移除出去
        //否则我们增加peak
        //无论如何，我们要增加used

        for(auto it = free_blocks.begin();it!=free_blocks.end();it++){
            size_t block_addr = it->first;
            size_t block_size  = it->second;
            if(block_size>=size){
                this->used += size;
                free_blocks.erase(it);
                if(block_size>size){
                    free_blocks[block_addr+size] = block_size-size;
                }
                return block_addr;
            }else if (block_addr + block_size == this->peak){
                this->used += size;
                size_t needed_extra = size - block_size; // 还需要向系统借多少？
                this->peak += needed_extra;              // 推高 peak (历史水位线)
                free_blocks.erase(it);                   // 消耗掉这个末尾块
                return block_addr;
            }   
        }
        //如果找不到合适的
        size_t block_addr = this->peak;
        this->peak += size;
        this->used += 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 it = free_blocks.find(addr);
    auto next_it = std::next(it);
    if(next_it!=free_blocks.end()){
        if(it->first + it->second == next_it->first){
            it->second += next_it->second;
            free_blocks.erase(next_it);
        }
    }
    if(it!=free_blocks.begin()){
        auto prev_it = std::prev(it);
        if(prev_it->first + prev_it->second == it->first){
            prev_it->second += it->second;
            free_blocks.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;
    }
}