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: 设计一个算法来分配内存，返回起始地址偏移量
        // =================================== 作业 ===================================
        //std::cout<<" ******************* map size "<<freeByStart.size()<<std::endl;
        used += size;
        if(freeByStart.empty()){
            peak += size;
            return peak - size;
        }

        bool freeBlkOk = false;  //is there a free block large enough to alloc?
        auto it = freeByStart.begin();
        for(; it != freeByStart.end(); it++){
            if(it->second >= size){
                freeBlkOk = true;
                if(it->second > size)
                    freeByStart.insert({it->first + size, it->second - size});
                break;
            }
        }

        if(!freeBlkOk){
            it--;
            size_t moreMem = size - it->second;
            peak += moreMem;
        }
        freeByStart.erase(it);
        return it->first;

    }

    void Allocator::free(size_t addr, size_t size)
    {
        IT_ASSERT(this->ptr == nullptr);
        size = getAlignedSize(size);

        // =================================== 作业 ===================================
        // TODO: 设计一个算法来回收内存
        // =================================== 作业 ===================================
                auto it = freeByStart.begin();
                int flag = 0;
                for(; it != freeByStart.end(); it++){
                    if(it->first + it->second == addr){
                        flag = -1;
                        break;
                    }
                    else if(it->first == addr + size){
                        flag = 1;
                        break;
                    }
                }
                if(flag == -1){
                    freeByStart.insert({it->first, it->second + size});
                    freeByStart.erase(it);
                }
                else if(flag == 0){
                    freeByStart.insert({addr, size});
                }
                else if(flag == 1){
                    freeByStart.insert({addr, size + it->second});
                    freeByStart.erase(it);
                }
                used -= size;
    }

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