在Ascend Transformer Boost加速库(以下简称ATB)的调用流程中,往往需要频繁调用Malloc和Free向设备申请和释放内存空间,这样的操作会打断设备的流水,造成额外的性能开销。
内存池是一种常见的技术手段用于管理和分配设备的内存。其基本功能原理如下:
一个Block至少需要三个元素,用于定位的索引、块大小描述和对应的实际物理地址。
struct MemoryBlock { // blockId,内存块索引,全局唯一 int64_t blockId; // blockSize,单位为字节 size_t blockSize; // 内存块的物理地址 void *address = nullptr; };
最基本的内存池,分别管理每个设备上的内存,一般其创建个数与设备数量对应。通常使用unordered_map的数据结构管理所有的block。它包含三个类方法,AllocateBlock、FreeBlock和GetBlockPtr,用于分配block、释放block和获取block中的内存地址。
内存管理类,负责初始化所有设备的MemoryPool,并将其放入一个数组中统一维护。当不同线程或进程访问该类并调用方法时,通过GetDeviceId方法获得线程对应的设备ID,并返回对应设备内存池的相应方法。
分配内存块方法,主要逻辑分为两个步骤:
释放内存块时,不直接释放内存,只是将其标记为未使用。
实现内存池的首要目的是减少aclrtMalloc和acletFree的调用次数,在本例中主要针对workspace的内存进行管理。
在Node类中会保存有当前已分配的workspace大小和对应的blockId。当需求的workspace大小发生变化时,如果需求大小小于当前已有空间则不进行操作,大于则释放当前内存块并重新申请更大的内存块。