MXNet Storage代码分析

本篇主要对mxnet/src/storage目录下的代码进行分析记录。对应的mxnet版本为0.11.0.

storage目录下的文件结构如下：

mxnet
|_ src
|    |_ storage
|        |_ cpu_device_storage.h
|        |_ gpu_device_storage.h
|        |_ navie_storage_manager.h
|        |_ pinned_memory_storage.h
|        |_ pooled_storage_manager.h
|        |_ storage_manager.h
|        |_ storage.cc
|_ include
|_ mxnet
|_ storage.h

这里还涉及到include/mxnet/storage.h文件，故将其列入。

首先看storage.h中storage类基本的定义：

class Storage {
public:
struct Handle {
void* dptr; // Pointer to the data
size_t size; // Size of the storage.
Context ctx; // Context information about device and ID.
};
virtual Handle Alloc(size_t size, Context ctx) = 0;

virtual void Free(Handle handle) = 0;

virtual void DirectFree(Handle handle) = 0;

virtual ~Storage() {}

static Storage* Get();

static std::shared_ptr<Storage> _GetSharedRef();
};

1、storage由地址dptr、大小size和上下文信息ctx组成。

2、主要的函数：Alloc()、Free()、DirectFree()，均为纯虚函数，由子类负责实现。(声明了纯虚函数的类是一个抽象类。用户不能创建类的实例，只能创建它的派生类的实例。)

Free()和DirectFree()的区别在于：DirectFree()会直接将storage释放，而不会将其放入内存池(memory pool)中。

3、注意到这个函数为抽象的基类，且Get()为静态函数，意味着单例模式。

在storage.cc中定义了StorageImpl类，继承了storage类，是storage的implementation. 它继承并实现了storage中的Alloc()、Free()、DirectFree()三个函数。

class StorageImpl : public Storage {
public:
Handle Alloc(size_t size, Context ctx) override;
void Free(Handle handle) override;
void DirectFree(Handle handle) override;
StorageImpl() {}
virtual ~StorageImpl() = default;

private:
static constexpr size_t kMaxNumberOfDevices = Context::kMaxDevType + 1;
static constexpr size_t kMaxNumberOfDeviceIDs = Context::kMaxDevID + 1;
#if MXNET_USE_CUDA
static int num_gpu_device;
#endif  // MXNET_USE_CUDA

static void ActivateDevice(Context ctx) {
......
}
// internal storage managers
std::array<common::LazyAllocArray<storage::StorageManager>,
kMaxNumberOfDevices> storage_managers_;
};

Alloc()、Free()、DirectFree()这三个函数中实际上都是通过调用StorageManager这个存储管理器类来实现内存的分配和释放。以DirecFree为例：

void StorageImpl::DirectFree(Storage::Handle handle) {
const Context &ctx = handle.ctx;
auto&& device = storage_managers_.at(ctx.dev_type);
std::shared_ptr<storage::StorageManager> manager = device.Get(
ctx.dev_id, []() {
LOG(FATAL) <<  "Cannot Free space to a device you have not allocated";
return nullptr;
});
this->ActivateDevice(ctx);
// directly free ths data.
manager->DirectFree(handle.dptr, handle.size);
}

可以看到StorageImpl仅作为提供handle信息(包含数据地址指针、大小和上下文，见上面Storage类中的成员结构体Handle)给StorageManager，由manager去执行实际的释放操作。

在storage_manager.h中定义了StorageManager类，这也是一个”抽象的基类”，只包含三个纯虚函数Alloc()、Free()、DirectFree()，全部由子类实现。

class StorageManager {
public:
virtual void* Alloc(size_t size) = 0;
virtual void Free(void* ptr, size_t size) = 0;
virtual void DirectFree(void* ptr, size_t size) = 0;
/*!
* \brief Destructor.
*/
virtual ~StorageManager() = default;
};

由上面的storage目录的结构图可以看出，StorageManager有两个派生类NaiveStorageManager、GPUPooledStorageManager分别在navie_storage_manager.h和pooled_storage_manager.h中。

主要关注用于管理GPU内存的GPUPooledStorageManager这个类，如名字所示应该是维护了一个memory pool. 这个类的声明如下：

class GPUPooledStorageManager final : public StorageManager {
public:
GPUPooledStorageManager() {
reserve_ = dmlc::GetEnv("MXNET_GPU_MEM_POOL_RESERVE", 5);
}
~GPUPooledStorageManager() {
ReleaseAll();
}

void* Alloc(size_t raw_size) override;
void Free(void* ptr, size_t raw_size) override;
void DirectFree(void* ptr, size_t raw_size) override {
cudaError_t err = cudaFree(ptr);
size_t size = raw_size + NDEV;
// ignore unloading error, as memory has already been recycled
if (err != cudaSuccess && err != cudaErrorCudartUnloading) {
LOG(FATAL) << "CUDA: " << cudaGetErrorString(err);
}
used_memory_ -= size;
}

private:
void ReleaseAll();
std::mutex mutex_; // internal mutex
size_t used_memory_ = 0; // used memory
int reserve_; // percentage of reserved memory
const int NDEV = 32; // number of devices
std::unordered_map<size_t, std::vector<void*>> memory_pool_; // memory pool

DISALLOW_COPY_AND_ASSIGN(GPUPooledStorageManager);
};

其中Alloc()、Free()、DirectFree()是和StorageManager、Storage中的同名函数是一脉相承的。

GPUPooledStorageManager中声明了一个Hash Map用作内存池(内存的size和一个元素为地址指针的vector的映射，这个vector中会包含多个dptr指针)，这个内存池中保存的是空闲(Free)的内存块。

接下来对每个函数进行说明：

- Alloc(): 在memory_pool_中寻找raw_size对应的vector(该vector存放的是空闲内存块的指针)，如果未找到这样的vector或者vector的size==0，就会调用cudaMalloc()申请新的内存块，同时used_memory_会增加；否则，会将vector中最后一个地址指针pop出并返回。

- Free(): 如果调用Free()来”释放”内存，仅仅是将地址指针ptr放入 memory_pool_中与size对应的vector中，used_memory_并不会减少，在vector中维护的就是空闲的内存指针，如果这个size不存在会自动创建(map operator[]的特性)。

- DirectFree(): 如果调用DirectFree()来释放内存，会真正调用cudaFree()来讲指针ptr对应的内存释放掉，used_memory_并会减少。

- ReleaseAll(): 该函数会将memory_pool_中保存的空闲内存地址全部释放掉(对每个指针调用DirectFree()函数)