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: 设计一个算法来分配内存，返回起始地址偏移量
        // Use First Fit algorithm: find the first free block that is large enough
        for (auto it = freeBlocksByAddr.begin(); it != freeBlocksByAddr.end(); ++it)
        {
            size_t addr = it->first;
            size_t blockSize = it->second;
            
            if (blockSize >= size)
            {
                // Found a suitable block
                // Remove this block from both maps
                freeBlocksByAddr.erase(it);
                freeBlocksByEnd.erase(addr + blockSize);
                
                // If the block is larger than needed, add the remaining part back
                if (blockSize > size)
                {
                    size_t newAddr = addr + size;
                    size_t newSize = blockSize - size;
                    freeBlocksByAddr[newAddr] = newSize;
                    freeBlocksByEnd[newAddr + newSize] = newSize;
                }
                
                // Update memory usage statistics
                used += size;
                if (used > peak)
                {
                    peak = used;
                }
                
                return addr;
            }
        }
        
        // No suitable free block found, allocate at the end
        size_t addr = peak;
        used += size;
        peak += size;
        
        return addr;
        // =================================== 作业 ===================================

        return 0;
    }

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

        // =================================== 作业 ===================================
        // TODO: 设计一个算法来回收内存
        // Update memory usage
        used -= size;
        
        size_t blockStart = addr;
        size_t blockEnd = addr + size;
        size_t blockSize = size;
        
        // Special case: if freeing the block at the end, just reduce peak
        if (blockEnd == peak)
        {
            // Check if we can merge with a previous free block at the end
            auto prevIt = freeBlocksByEnd.find(blockStart);
            if (prevIt != freeBlocksByEnd.end())
            {
                // Merge with the previous block and reduce peak further
                size_t prevSize = prevIt->second;
                size_t prevStart = blockStart - prevSize;
                
                // Remove the previous block from both maps
                freeBlocksByAddr.erase(prevStart);
                freeBlocksByEnd.erase(blockStart);
                
                // Reduce peak to the start of the merged block
                peak = prevStart;
            }
            else
            {
                // Just reduce peak
                peak = blockStart;
            }
            return;
        }
        
        // Try to merge with the previous adjacent free block
        auto prevIt = freeBlocksByEnd.find(blockStart);
        if (prevIt != freeBlocksByEnd.end())
        {
            // Found a previous adjacent block
            size_t prevSize = prevIt->second;
            size_t prevStart = blockStart - prevSize;
            
            // Remove the previous block from both maps
            freeBlocksByAddr.erase(prevStart);
            freeBlocksByEnd.erase(blockStart);
            
            // Merge: extend the current block backwards
            blockStart = prevStart;
            blockSize += prevSize;
        }
        
        // Try to merge with the next adjacent free block
        auto nextIt = freeBlocksByAddr.find(blockEnd);
        if (nextIt != freeBlocksByAddr.end())
        {
            // Found a next adjacent block
            size_t nextSize = nextIt->second;
            
            // Remove the next block from both maps
            freeBlocksByAddr.erase(blockEnd);
            freeBlocksByEnd.erase(blockEnd + nextSize);
            
            // Merge: extend the current block forwards
            blockSize += nextSize;
            blockEnd += nextSize;
        }
        
        // Add the merged (or original) free block to both maps
        freeBlocksByAddr[blockStart] = blockSize;
        freeBlocksByEnd[blockEnd] = blockSize;
        // =================================== 作业 ===================================
    }

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