MemPool

每个MemPool都靠一个(u64, u64)的二元组来唯一定位。分别是uid和uuid。

• 在aten/src/Aten/cuda/MemPool.cpp中有详细的解释

• uid每次用户申请一个新的mem pool会自增

• CUDAGraph每次创建mem pool会增加uuid

• {0, 0}表示默认的pool

在CUDACachingAllocator中，通过PrivatePool来管理所有的内存池

• 有一个use_count/cudaMalloc_count，相当于引用计数。当变成0的时候才可以释放掉这个pool

• 支持用户自定义一个allocator，类型是CUDAAllocator，用户需要实现它的raw_alloc/raw_delete

  • 在alloc/release block的时候可以看到相关的判断和调用
• 每一个pool也有自己的BlockPool，同样是有large/small的区分。

pool的作用/使用的场景：

• CUDAGraph，graph capture期间使用的内存的虚拟地址会被记录到图中，后续重放的时候会被使用，所以对应的内存不能被复用。
  • 比如同一个batch size的cuda graph用同一个内存池
• 自定义Allocator

• use_on_oom, no_split这些选项允许一定程度的调整caching allocator的行为

然后再简单看一下Pool是怎么使用的：

• torch.cuda.MemPool创建pool
  • 对应到aten/cuda的MemPool对象

  • 会调用到CachingAllocator的create_or_incref_pool，作用是创建PrivatePool，保存在DeviceCachingAllocator中

• torch.cuda.use_mem_pool，是一个context manager，用来使用这个pool

  • 对应到CachingAllocator的beginAllocateToPool

  • 会把pool id，以及一个filter（作用是让pool只针对当前线程生效），记录到captures_underway中

  • 主流程的get_pool就会从captures_underway中把对应的pool找出来使用

CUDAGraph

除去上面的pool的设计，再来看一下CachingAllocator在支持cuda graph上都做了那些适配。
在c10/cuda/CUDACachingAllocator.cpp开头的注释中有相关的设计note

CUDAGraph capture期间，不允许调用和CPU有关联的API，event相关的就不能用。

• 所以在有跨stream复用的情况下
  • 之前的逻辑是通过query event来等待所有stream结束，再去释放对应块的内存

  • 因为query event不能用，所以这里引入了一个deferred block，就是只能等整个pool释放的时候才能释放这些内存。

  • 但是这种方法会导致跨越stream的内存就无法复用了，所以有一个free_safe_blocks_in_capture的优化，可以做到细粒度的追踪。细节就等看cuda graph的时候再追了

• 默认的cuda graph capture模式下，是不允许执行malloc的。

  • 实际上malloc的行为是正常的，在capture模式下会分配新的VA，然后在replay的时候就是dummy的操作，直接使用当时分配的VA即可

  • 所以这里会先切换到cudaStreamCaptureModeRelaxed模式，表示允许做malloc，然后再去调用cuda malloc

• 代码中还有支持cuda graph tree的一些机制，也等待看cuda graph的时候再研究