memory_utils.cpp

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
#include <acl/acl.h>
#include "memory_utils.h"
#include "utils/log.h"
#include "utils/utils.h"

// 全局MemoryManager实例
static MemoryManager g_memoryManager;

MemoryManager::MemoryManager()
{}

void MemoryManager::CreateMemoryPool(size_t poolSize)
{
    uint32_t deviceCount = 0;

    // 获取全部Device的数量
    CHECK_RET(aclrtGetDeviceCount(&deviceCount), "get devicecount fail");
    for (size_t i = 0; i < deviceCount; i++) {

        // 指定操作的Device
        aclrtSetDevice(i);

        // 创建内存池,poolSize参数指定预分配空间大小
        std::shared_ptr<MemoryPool> memoryPool = std::make_shared<MemoryPool>(poolSize);
        memoryPools_.push_back(memoryPool);
        LOG_INFO("create mempool for device " + std::to_string(i) + " success");
    }
}

int32_t MemoryManager::GetDeviceId()
{
    int32_t deviceId = -1;
    CHECK_RET(aclrtGetDevice(&deviceId), "get device ID fail");
    return deviceId;
}

std::shared_ptr<MemoryPool> &MemoryManager::GetMemoryPool()
{
    // 获取当前操作的Device,返回对应的内存池
    size_t deviceId = static_cast<size_t>(GetDeviceId());
    CHECK_RET(deviceId >= memoryPools_.size(), "Invalid device id " + deviceId);
    return memoryPools_[deviceId];
}

void MemoryManager::AllocateBlock(uint32_t size, int &blockId)
{
    GetMemoryPool()->AllocateBlock(size, blockId);
}

void MemoryManager::FreeBlock(int blockId)
{
    GetMemoryPool()->FreeBlock(blockId);
}

void MemoryManager::GetBlockPtr(int blockId, void *&addr)
{
    GetMemoryPool()->GetBlockPtr(blockId, addr);
}

MemoryManager &GetMemoryManager()
{
    return g_memoryManager;
}