strorage.cpp 1.78 KB
Newer Older
thatPepe's avatar
thatPepe committed
1
#include "../allocator.hpp"
PanZezhong's avatar
init  
PanZezhong committed
2
3
4
#include "../tensor.hpp"

std::shared_ptr<Storage> Storage::create(size_t size) {
wooway777's avatar
wooway777 committed
5
    auto storage = std::shared_ptr<Storage>(new Storage());
PanZezhong's avatar
PanZezhong committed
6
7
8
    RUN_INFINI(infinirtMalloc(&storage->_memory, size));
    storage->_size = size;
    RUN_INFINI(infinirtGetDevice(&storage->_device_type, &storage->_device_id));
PanZezhong's avatar
init  
PanZezhong committed
9
10
11
12
    return storage;
}

std::shared_ptr<Storage> Storage::createAsync(size_t size, infinirtStream_t stream) {
wooway777's avatar
wooway777 committed
13
    auto storage = std::shared_ptr<Storage>(new Storage());
PanZezhong's avatar
PanZezhong committed
14
15
16
    RUN_INFINI(infinirtMallocAsync(&storage->_memory, size, stream));
    storage->_size = size;
    RUN_INFINI(infinirtGetDevice(&storage->_device_type, &storage->_device_id));
PanZezhong's avatar
init  
PanZezhong committed
17
18
19
    return storage;
}

thatPepe's avatar
thatPepe committed
20
std::shared_ptr<Storage> Storage::createFromPool(size_t size, std::shared_ptr<MemoryPool> pool) {
wooway777's avatar
wooway777 committed
21
    auto storage = std::shared_ptr<Storage>(new Storage());
PanZezhong's avatar
PanZezhong committed
22
    storage->_memory_pool = pool;
thatPepe's avatar
thatPepe committed
23
    if (pool) {
PanZezhong's avatar
PanZezhong committed
24
        storage->_memory = pool->alloc(size);
thatPepe's avatar
thatPepe committed
25
    } else {
PanZezhong's avatar
PanZezhong committed
26
        RUN_INFINI(infinirtMalloc(&storage->_memory, size));
thatPepe's avatar
thatPepe committed
27
    }
PanZezhong's avatar
PanZezhong committed
28
29
    storage->_size = size;
    RUN_INFINI(infinirtGetDevice(&storage->_device_type, &storage->_device_id));
thatPepe's avatar
thatPepe committed
30
31
32
    return storage;
}

PanZezhong's avatar
init  
PanZezhong committed
33
std::shared_ptr<Storage> Storage::createHost(size_t size) {
wooway777's avatar
wooway777 committed
34
    auto storage = std::shared_ptr<Storage>(new Storage());
PanZezhong's avatar
PanZezhong committed
35
36
37
38
39
    RUN_INFINI(infinirtMallocHost(&storage->_memory, size));
    storage->_size = size;
    storage->_device_type = INFINI_DEVICE_CPU;
    storage->_device_id = 0;
    storage->_memory_pool = nullptr; // No pool for host memory
PanZezhong's avatar
init  
PanZezhong committed
40
41
42
43
    return storage;
}

Storage::~Storage() {
PanZezhong's avatar
PanZezhong committed
44
    if (_memory_pool) {
wooway777's avatar
wooway777 committed
45
        _memory_pool->release(_memory);
PanZezhong's avatar
init  
PanZezhong committed
46
    } else {
PanZezhong's avatar
PanZezhong committed
47
48
        if (_device_type == INFINI_DEVICE_CPU) {
            RUN_INFINI(infinirtFreeHost(_memory));
thatPepe's avatar
thatPepe committed
49
        } else {
PanZezhong's avatar
PanZezhong committed
50
            RUN_INFINI(infinirtFree(_memory));
thatPepe's avatar
thatPepe committed
51
        }
PanZezhong's avatar
init  
PanZezhong committed
52
53
    }
}