memorypool.cpp

#include <atb/types.h>
#include <acl/acl.h>
#include "memorypool.h"
#include "utils/log.h"
#include "utils/utils.h"

constexpr size_t POOL_SIZE = 104857600;  // Pre-allocate 100 MiB memory. The size can be determined based on the actual requirements.

MemoryPool::MemoryPool(size_t poolSize = POOL_SIZE)
{
    // Call the acl API to pre-allocate memory.
    CHECK_RET(aclrtMalloc(&baseMemPtr_, poolSize, ACL_MEM_MALLOC_HUGE_FIRST),
        "malloc huge size memory " + std::to_string(poolSize) + " bytes fail");
    curMemPtr_ = baseMemPtr_;
    remainSize_ = poolSize;
}

MemoryPool::~MemoryPool()
{
    // Release the memory through the acl API during destruction.
    if (baseMemPtr_ != nullptr) {
        CHECK_RET(aclrtFree(baseMemPtr_), "free huge memory fail");
    }
    LOG_INFO("release MemoryPool success");
}

uint64_t MemoryPool::GenerateBlocksId()
{
    return static_cast<uint64_t>(id_.fetch_add(1, std::memory_order_relaxed));
}

void MemoryPool::AllocateBlock(uint32_t size, int &blockId)
{
    std::unique_lock<std::mutex> lock(blockMutex_);

    size_t alignSize = ((size + 31) & ~31) + 32; // The space allocated by must be 32-byte aligned plus 32 bytes.

    // Check whether there are sufficient free memory blocks.
    for (auto it = freeBlocks_.begin(); it != freeBlocks_.end(); it++) {
        if (it->second.blockSize >= alignSize) {
            blockId = it->second.blockId;
            usedBlocks_.insert(*it);
            freeBlocks_.erase(it);
            LOG_INFO("find free block id " + std::to_string(blockId) + " to allocate");
            return;
        }
    }

    // No memory block is found. A new memory block needs to be created from the remaining memory space.
    if (remainSize_ > alignSize) {
        blockId = GenerateBlocksId();
        uint64_t curMemPtrAlign = (reinterpret_cast<uint64_t>(curMemPtr_) + 63) & ~63;  // The memory address must be 64-byte aligned.
        remainSize_ -= (curMemPtrAlign - reinterpret_cast<uint64_t>(curMemPtr_));
        curMemPtr_ = reinterpret_cast<void *>(curMemPtrAlign);

        MemoryBlock block = {blockId, alignSize, curMemPtr_};
        usedBlocks_.insert({blockId, block});
        remainSize_ -= alignSize;
        curMemPtr_ = reinterpret_cast<uint8_t *>(curMemPtr_) + alignSize;
        LOG_INFO("allocate block id " + std::to_string(blockId) + " for size " + std::to_string(alignSize));
        return;
    }

    // Failed to allocate a memory block due to insufficient remaining memory space.
    LOG_ERROR("allocate block fail");
}

void MemoryPool::FreeBlock(int blockId)
{
    std::unique_lock<std::mutex> lock(blockMutex_);

    // Check the validity of the memory block index.
    if (blockId < 0) {
        LOG_INFO("skip over the invalid block id " + std::to_string(blockId));
        return;
    }

    // Mark the memory block to be released as idle.
    auto it = usedBlocks_.find(blockId);
    if (it != usedBlocks_.end()) {
        freeBlocks_.insert(*it);
        usedBlocks_.erase(it);
    } else {
        LOG_ERROR("Double free block id " + std::to_string(blockId));
    }
}

void MemoryPool::GetBlockPtr(int blockId, void *&addr)
{
    std::unique_lock<std::mutex> lock(blockMutex_);

    if (blockId < 0) {
        LOG_INFO("Invalid block id " + std::to_string(blockId) + "to get ptr");
        return;
    }

    // Search for a memory block and return the physical memory address.
    auto it = usedBlocks_.find(blockId);
    if (it != usedBlocks_.end()) {
        addr = it->second.address;
    } else {
        LOG_ERROR("Get block address error, block id " + std::to_string(blockId));
    }
}