首页 > 代码库 > mxnet 源码阅读 6

mxnet 源码阅读 6

mxnet/src/storage/storage.cc

mxnet/include/mxnet/storage.h

mxnet/include/mxnet/base.h

 

以上三个文件共同描述了Storage虚类以及其实例化的结果,其中storage.h中定义了Storage抽象接口(虚函数) Alloc Free DirectFree 以及静态方法Get ,其中定义了以后用以操作和管理内存的Handle句柄,其内部有两个部分组成一个是描述内存本身属性的包括地址 void *dptr 和大小 size_t size 以及描述内存所在设备的Context ctx; 其中Context结构在base.h中定义,其成员有两个一个是 dev_type (1-CPU 2-GPU 3-CPU_PINNED ) 另一是 dev_id ,同时在base中还有一些构建Context的相关函数,如GPU(...) CPU(...) CPUPinned(,,,) FromString(...) 等可以直接构建相关设备环境。

StorageImpl是 Stroage的具体继承类接口有 返回Handle 的 Alloc 函数 Free 函数 和 DirectFree函数,同时多了静态函数 ActivateDevice用以切换GPU设备,CPU设备其实不用切换,所以实现也只是直接break而已。

从Alloc的实现看,其申请的内存包括 CPU内存,CPUpinned 内存 和 GPU内存池,没有单纯的申请一块GPU内存的接口。

同时总结一下mxnet/src/storage 目录中文件和类的情况:

cpu_device_storage.h gpu_device_storage.h pinned_memory_storage.h 分别构建了负责CPU GPU 和 CPUpinned 内存申请的类。

storage_manager.h 是纯粹的抽象的接口类。naive_storage_manager.h 和  pooled_storage_manager.h 是抽象类storage_manager是具体实现。实质使用是通过C++多态,分别管理了 CPU(包括CPUpinned)和GPU(内存池)。

 

整个storage目录中的内容在mxnet框架中属于底层模块python等接口涉及不到该接口。

总结起来讲,其提供了以Handle为操作的句柄的内存申请和释放接口,即 Handle Alloc(size_t size, Context ctx),void Free(Handle handle),void DirectFree(Handle handle) 函数,统一了CPU GPU 接口。CPU 和CPUpinned 是随时申请随时释放模式,GPU内存采用内存池的模式。

猜测gpu_device_storage.h是为了以后临时资源的申请之用。

mxnet 源码阅读 6