torch allocator的阅读笔记

• DataPtr
  • 包含c10::detail::UniqueVoidPtr和device

  • Device是一个type和一个index

    • type就是比如CPU，GPU等

    • index就是表示第几个，比如cuda0, cuda1这种

  • UniqueVoidPtr

    • 一个data表示用户使用的数据

    • 一个unique ptr表示这个数据的owner

    • 用户拿到的指针不一定代表这块内存

    • 而且unique_ptr在null的时候不会调用deleter，这里也可以支持这种情况。data为空，但是存在真正的对象需要释放

    • 一般情况下认为这两个相同就行

• Allocator

  • allocate，给定size，返回DataPtr
• SetAllocator/GetAllocator，针对device设置allocator

• MemoryReportingInfoBase

  • Thread local的memory usage reporter

  • 应该是每次分配做一次上报

Allocator

DefaultCPUAllocator

• posix_memalign做内存分配，就是做地址对齐的malloc

• 带上了profiledCPUMemoryReporter的wrapper，会记录每次申请和释放，以及out of memory

MapAllocator

Map allocator不继承自allocator，而是一个helper
外部使用主要是makeDataPtr，比如StorageUtils中new_shm_fd_storage会用到这个makeDataPtr
一块内存就对应一个Allocator
有一堆flag来控制具体的行为，比如：

• 是否用shm_open来走匿名的shm

• 使用O_EXCL，避免名字冲突

然后来理一下这里的常见的flag的能力，以及对应的分配行为

• 第一步open，需要先从什么地方搞一个fd出来，用来做mmap
  • FROMFD，就是外部传入的fd直接open

  • SHARED，把磁盘文件作为shared mapping，这里会走open

  • SHAREDMEM，走shm_open，对应是在/dev/shm或者tmpfs上创建匿名的shm对象

• 第二步是resize

  • 先通过fstat获取当前对象的大小

  • 通过ftruncate进行resize。这里的语义是修改对应的fd，所以如果是read only场景是不会resize的

• 第三步是mmap，给定fd，通过mmap把指针map出来

  • MAP_SHARED，写入对其他进程可见，这样多个进程会读到相同的值（即写回文件）

  • MAP_PRIVATE，是copy-on-write的，不影响原文件

• 收尾

  • 首先理解：

  mmap 之后，文件/shm 对象在内核里的存活只看两件事：还有没有进程持有 fd 引用，或者还有没有进程持有 mmap 引用（页面表里指向这个 inode/shm 段）。只要任意一种引用 ≥1，对象的物理页就不会被回收，即使你已经 unlink / shm_unlink 把"路径名"移除掉了。

  • KEEPFD：把 fd 保留在 fd_ 里（不 close）。这给后续 share_fd_cpu 拿出来 dup + 通过 Unix domain socket 发给别的进程提供了素材。默认是 close 的——因为对 mmap 来说 fd 已经不需要了，省一个 fd。

  • UNLINK：mmap 之后立刻把名字从命名空间里删掉。结合上面的语义，这意味着"对象本身"没有路径名能再被打开了，但因为 mmap 的引用还在，物理页继续存活；等所有进程 munmap 后内核自动回收。

  • 还有一个madvise，POSIX_FADV_SEQUENTIAL，使用更大的block size，减少读大块内存的时候陷入内核态的次数

• 释放

  • 如果有fd，就close(fd)

  • munmap对应的指针，这段内存对应的物理页就会被释放

• 使用场景
  • at::new_shm_fd_storage，对应_share_fd_cpu/share_memory_
    • SHAREDMEM | EXCLUSIVE | KEEPFD | UNLINK

    • 创建 shm → 立刻 unlink → 仅靠 fd 维持存活，方便发 fd 给别人

  • THPStorage_newSharedFd

    • SHAREDMEM | NOCREATE | KEEPFD | FROMFD

    • 接收端从已经 dup 来的 fd mmap，名字在发送端早已 unlink

FD共享流程

从这里再来看一下multiprocessing是怎么做的共享tensor

• _share_fd_cpu_调用到THPStorage_shareFd，然后调用到at::new_shm_fd_storage(nbytes)，对应上面说的流程，创建保留fd的shmem

• at::storage_copy，把数据拷贝进shm，然后原storage的data_ptr替换成新的，旧的就销毁掉了

• 返回fd, size给python

• 然后python层，通过dupfd，把fd发给其他进程。两个进程的fd不一定相同，但是指向同一个对象，此时shm的fd refcount + 1

• 接收端，_new_shared_fd_cpu，拿到新的fd，走的是上面提到的THPStorage_newSharedFd

父进程                                内核                              子进程
─────────                            ─────                             ─────
shm_open("/torch_xxx",O_CREAT|O_EXCL) ──► 创建 inode，nlink=1, i_count=1(fd)
ftruncate / mmap                     ──► i_count=2 (fd+mmap)
shm_unlink("/torch_xxx")             ──► nlink=0  ← 名字消失，inode 仍活着
                                                                       ...
sendmsg(SCM_RIGHTS, fd) ─────────────────────────────────────────────► recvmsg
                                          子进程拿到 fd，i_count=3
                                                                       dup
                                                                       i_count=4 (短暂)
                                                                       close(tmp_fd) → i_count=3
                                                                       mmap
                                                                       i_count=4 (fd+mmap)
~MapAllocator (父端)
  munmap                             ──► i_count=3
  close(fd) (KEEPFD)                 ──► i_count=2
                                                                       ~MapAllocator (子端)
                                                                       munmap → i_count=1
                                                                       close(fd) → i_count=0
                                          nlink=0 && i_count=0
                                          → 内核回收 inode + 物理页

• 这里有一个小细节，就是子进程拿到内核传过来的fd之后，给cpp层的时候会dupfd一下。python层的tmpfd无论如何都会被close掉。

• 主要是分层设计，底层对fd的行为比较多，如果和上层耦合起来上面就需要有一堆判断条件来决定这个fd是否做close。异常处理也比较复杂

CUDACachingHostAllocator

Host allocator相对简单，先看这个。pin memory相关的用的就是这里的Allocator
对应at::getHostAllocator(at::kCUDA)

                         REGISTER_HOST_ALLOCATOR(at::kCUDA, &caching_host_allocator)
                                          │
                                          ▼
┌──────────────────────────────────────────────────────────────────────┐
│  CachingHostAllocatorInterface<T, deleter>   ← 对外接口（at::HostAllocator）│
│  · 实现 allocate() / record_event() / empty_cache() / get_stats() ...   │
│  · 把 (ptr, block*) 包装成 at::DataPtr（device 标成 kCPU）              │
└──────────────────────────────────────────────────────────────────────┘
                                          │
                                          ▼
┌──────────────────────────────────────────────────────────────────────┐
│  CachingHostAllocatorImpl<S, E, B>           ← 通用 caching 逻辑（CRTP 基类）│
│  · 模板参数: S=Stream, E=Event, B=HostBlock<S>                       │
│  · free list / 事件队列 / size_class 分桶 / private pool / 后台线程    │
│  · 不依赖任何具体 backend                                              │
└──────────────────────────────────────────────────────────────────────┘
                                          │ 虚函数 hook
                                          ▼
┌──────────────────────────────────────────────────────────────────────┐
│  CUDACachingHostAllocatorImpl                ← CUDA 特化（仅 ~280 行）│
│  · allocate_host_memory     → cudaHostAlloc / cudaHostRegister       │
│  · free_block               → cudaFreeHost / cudaHostUnregister      │
│  · record_stream/query_event → cudaEventRecord/cudaEventQuery        │
│  · PinnedReserveSegment / multi-thread page register                  │
└──────────────────────────────────────────────────────────────────────┘

HostAllocator继承了at::Allocator，是一个虚基类
在CachingHostAllocator.h中讲了这里的设计。

• HostAllocator的作用是允许新的设备自己实现host allocator

• CachingHostAllocatorImpl的作用是实现一个比较通用的allocator，允许用户方便的使用这套能力，而不用自己重新写代码

关键结构：

• HostBlock
  • 一个memory block
• FreeBlockList
  • 某一个特定size的block list
• HostBlockPool
  • 维护所有的block ptr，free list，events

  • Free list共64个，对应分配内存从2^1到2^64

• PinnedReserveSegment

  • 通过PYTORCH_CUDA_ALLOC_CONF=pinned_reserve_segment_size_mb:N控制

• allocate(size)
  • 这里针对CUDA Graph有一个单独的pool，暂时先跳过

  • process_events

    • 把已经ready的block放回free list

    • 如果打开了pinned_use_background_threads，会用后台线程做清理

  • get_free_block

    • size向上power of 2取整，从对应的free list中获取block
  • 有free block 就直接返回，否则的话调用allocate_host_memory，按照roundSize申请一个block
    • 这个allocate_host_memory就是继承CRTP的类需要实现的
  • pool中记录block，返回block和对应的ptr

  • 启动后台线程的逻辑也在这里，即首次分配启动线程，100ms一次尝试调用process_events

• free(block)

  • 释放block，先看一下这个block中的stream是否为空（应该是表示是否有stream在使用）。如果非空的话，这里会分配若干个event，然后record_stream()。并把这些event记录到pool.event_队列中
    • 因为这个block释放的时候，有stream可能在异步的使用，所以不能复用这块内存。free的时候记录一下stream的位置，这样使用的event会在free event之前，那么到达free event的时候就可以安全的复用内存了。
  • 如果没有stream在用的话，会把这个block插回free_list中。

• record_event(ptr, block, stream)

  • 记录一下传入的stream使用了当前block。这里没有做record_stream，而是在free的时候才调用的record_stream。

  • 因为可能一个pinned memory block会被使用多次，每次使用都record_stream会有开销。所以这里延迟到free的时候。

• process_events_for_specific_size

  • 对应上面的process_event

  • 从pool的event queue中取一个event，调用query_event，看这个event是否完成。如果完成了就event_size - 1，当event数量为0的时候，归还给free list

  • 如果传入的size是-1，这里会一直query_event，直到遇到没结束的event

  • 如果传入的size是特殊的size，这里会遍历指定size的block的event，直到找到一个可以复用的再返回

• empty_cache

  • process_events，收割event

  • free_from_pool，遍历free list，调用free_block

通用机制看完了，再看一下这里CUDA相关的实现

• allocate_host_memory
  • PinnedReserveSegment，通过配置预先分配好的。这里会按照4K做对齐进行分配。如果超过了配置的数量的话就会返回空指针分配失败。好处是减少cudaHostAlloc的次数

  • 这里还有一个问题是pinned host memory是全局都可以使用的，不像是GPU的内存。但是分配pinned memory需要指定cuda context，是和device绑定的。所以这里会随便找到一个context做分配。

    • 不过一般来说都是一个进程一个GPU，所以这里都是找自己对应的那个cuda context做分配，context只有一个。
  • 如果启动了pinned_use_cuda_host_register，会调用allocWithCudaHostRegister
    • cudaHostAlloc 在 driver 内部有全局锁，分配大块时极慢。所以这里的实现是先posix_memalign分配一段内存，然后用线程池并行的触发page fault。然后调用cudaHostRegister pin住这些内存页

    • free也是对应的实现，或者是使用cudaHostUnregister + free，或者直接调用cudaFreeHost

  • 否则会调用cudaHostAlloc

• record_stream/query_event

  • 有一个全局的CUDAEventPool，从里面创建，event有自己的api，比如event->record(stream)，或者是event->query()

  • event的实现在c10/cuda/CUDAEvent.h中

有一个TODO是这块在CUDAGraph下有单独的行为，有空再来研究一下

CUDACachingAllocator

设计点：

• allocation是和stream绑定的，不能跨stream共享

• allocator会尝试使用smallest cached block，如果block比较大的话，会尝试切分block，如果没有block的话，会使用cudaMalloc

• 如果cudaMalloc失败，allocator会尝试释放一个cached block，然后重试。如果再失败，会尝试释放所有cached block，再重试

• >1MB的allocation，和小的allocation是分到不同的pool中的。小的allocation会pack到2MB的buffer中

• 1MB到10MB的request会分配并切分一个20MB的block

• 大于max_split_size的block是不能切分的，用于避免fragmentation

• DeviceAllocator，继承自c10::Allocator，额外的接口有
  • emptyCache

  • recordStream等

• CUDAAllocator继承自DeviceAllocator

数据结构

DeviceCachingAllocator (每个 GPU 一个)
├── small_blocks : BlockPool                      ← 默认 small 池
├── large_blocks : BlockPool                      ← 默认 large 池
├── active_blocks : flat_hash_set<Block*>         ← 当前已分配出去的 block
├── expandable_segments_ : vector<ExpandableSegment*>
├── cuda_events : map<Stream, deque<(Event, Block*)>>   ← 跨 stream 同步队列

Cuda graph相关的我去掉了
外层有一个NativeCachingAllocator，把N个DeviceCachingAllocator通过device index串起来。所以外面直接调用的是NativeCachingAllocator的接口

• Block
  • 基础属性：device, stream(分配他的stream), size/request_size/ptr

  • stream_uses，通过record_stream登记的其他stream

  • pool，small/large，对应上面的 1MB分割

  • prev/next指针，表示的是相同segment内的相邻block，用来做split/merge

  • registration_counter，每个block创建时+1，用来做block之间的比较

• BlockPool

  • blocks，空闲的block，是一个set
    • 根据stream/size/registeration_counter进行排序，方便在相同stream中根据size找可以用来复用的block
  • is_small，1MB的区分

  • unmapped，没有cuMemMap的段，启动expandable segment的时候会用到。也是一个set

    • 根据stream, ptr排序。方便在相同stream中找连续的空洞
• AllocParams
  • 单次分配的参数，核心是search_key，根据device/stream/size来进行查找

主流程 malloc/free

• malloc
  • process_events，回收block

  • round_size()

    • 最小分配是512B

    • 一般情况，对齐到512B的倍数

    • roundup_power2_divisions逻辑：

      • 比如size变化比较多，但是如果按照512B对齐无法做复用。但是直接按照2的幂次对齐又会有浪费。此时就可以用这个功能。意思是分配的时候，把2的幂次区间做切分，比如[2MB, 4MB]，切分成若干个小块，然后按照这个小块做对齐

      • 比如配置roundup_power2_divisions:[64:8, 256:4, 1024:4, >:1]，表示小于64MB的，切分成8块，64-256MB切分成4块，到1024MB也是4块，更大的区间则是1块

      • 默认都是不切，所以就都是对齐到512

  • get_pool，根据请求的size去判断用small pool还是large pool

  • get_allocation_size，对应上面的设计

    • 如果是small pool的话，小于1MB，分配2MB

    • 如果是小于10MB的话，分配20MB

    • 否则round到2MB分配

  • get_free_block

    • 到pool中进行lower bound，上面说了是根据stream，size做的排序。所以这里会找到第一个大于等于request size的free block。没找到就直接返回false了

    • 如果找到了block，判断一下

      • 如果请求的size小于max_split_size，同时block的size大于max split size。说明尝试在用大的block满足小的请求，此时会返回false，分配失败。（不过默认这个max_split_size是int max)

      • 如果请求的size也大于max_split_size的话，block size不能大于请求size + max_non_split_rounding_size。这两个判断条件主要就是为了启动split的情况下，满足split设计的要求

  • 如果失败了，会trigger_free_memory_callbacks，再尝试get_free_block

    • FreeCudaMemoryCallbacksRegistry中遍历并执行
  • 先看get free block成功的路径
    • should_split
      • 如果是small池，并且剩余的block size大于512B，就会做切分。（512B是最小分配的粒度，所以这里的逻辑就是只要剩下的block还能用就会切分）因为small池分配小，不容易产生碎片

      • 否则的话，就是large池的切分，size < max_split_size，并且remain > 1MB，表示remain也在large池，就会做切分。

    • alloc_found_block

      • 如果本次要进行切分的话，这里会分配一个新的block，把remaining这段放过来。然后pool->insert_into_blocks写回去。

      • 然后维护一下prev/next指针

      • 剩下的都是一堆profile相关的

        • record_trace

        • for_each_selected_stat_type

        • reportMemoryUsageToProfiler

  • 然后是没有复用block的链路

    • alloc_block
      • throw_on_cudamalloc_oom，启动后会先判断一下，当前process使用的显存是不是超了配置的阈值，超了的话就会抛异常。应该是对应多个进程使用同一块GPU的case

      • Expandable segment是单独的路径

      • release_lock_on_cudamalloc启动的话，会释放一个device level的lock，再去调用cudaMallocMaybeCapturing。应该是一个小性能优化，默认也是不开的，因为都是一个进程一个GPU

      • 分配成功了就还是一些统计信息相关的，record_trace等

    • 重试的链路在外层

      • 如果不是oom rejected，说明还是有可能分配出来这块内存的。

      • try_mempool_fallback

      • release_available_cached_blocks，然后alloc_block

      • release_cached_blocks，再尝试alloc_block

    • 如果还是没分配出来，就会走到oom的链路，也是记录一堆统计信息

• Free

  • record_trace, FREE_REQUESTED

  • 如果没有跨stream使用，直接free_block

    • free_block中对应的事件是FREE_COMPLETED，所以这里可以看出来因为跨stream的free延迟

    • try_merge_blocks

      • 判断一下邻居是否可用，需要stream_uses.empty(), event_count 0等。

      • 把邻居的block释放掉，dst的block的size增加

    • pool.insert_into_blocks

  • 跨stream使用的话，insert_events

    • 和host allocator一样，把每一个当前正在使用的stream都拿出来，然后创建一个event，并记录到stream级别的event队列上
• process_events
  • 遍历所有的stream的event队列，做cudaEventQuery

  • 如果处理完，block的event count 0，则会free掉这个block

  • 如果没ready就会break，因为这里event的插入是按照时间顺序的，所以按顺序做query即可。

    • 怎么感觉host allocator这样实现也没啥开销，反而还更简单，比现在这样一个队列前后pop更好理解

retry链路整理

整理一下malloc链路中的几种获取显存的方法

• get_free_block，池内做best fit + split，不会有device/host sync，也不会把显存还给driver，速度最快

• trigger_free_memory_callbacks，是让注册了free memory callback的系统释放资源。比如ProcessGroupNCCL中的资源，可以注册在这里并释放。也是软件层的操作，开销比较小

• garbage_collect_cached_blocks，释放非split块。发生在alloc_block之前，用来提高alloc block申请显存的成功率

  • 调用release_block，内部调用cudaFree()，并把block释放掉
• try_mempool_fallback
  • 从私有的池子中尝试借用。

  • 对应torch.cuda.MemPool(use_on_oom=True)

• release_available_cached_blocks

  • 针对max_split_size设计的，尝试释放掉大于max_split_size的块
• release_cached_blocks
  • synchronize_and_free_events，同步等待事件，然后释放block。相当于把所有record stream的内存块全都归还回来。因为这些是能确定的可以回收的内存

  • release_blocks，释放large blocks/ small blocks中所有的block

• 在release_cached_blocks之后的alloc_block的retry标记才为true。也就是说平常看到的日志中的CUDA malloc retry是走到了这个release_cached_blocks，一旦打印说明已经有大量synchronize了

Max split size设计

• CUDACachingAllocator中的split设计有一个问题，就是如果size变化比较多的话，会导致很多大块被切成小块，碎片率比较高。

• 引入max_split_size就是保证超过这个大小的块不再被切，让大块内存的复用率会更高

实现点是4块：

• should_split中，保证oversize请求不能被split

• get_free_block中的时候，保证小的请求不能使用oversize块，同时大的请求在复用oversize块的时候，也有一个碎片的阈值，避免大块碎片太多

• OOM时，优先释放oversize块，因为这些块不容易复用

从而保证：

• 一个 1GB 的 oversize 块永远是 1GB 的整块状态——要么完整空闲、要么完整在用，不会变成"800MB used + 200MB free 卡住"。

• 碎片只会发生在 < max_split_size 的小块世界里，最坏情况下的碎片浪费也被限制在阈值附近。

通过这些配置开启：

PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512,max_non_split_rounding_mb:200

ExpandableSegment

torch2.x新引入的技术，在代码中搜Expandable Segments就可以看到他的设计思路

场景

• 模型第一次 forward 用 batch=N，allocator 已经 cudaMalloc 了一堆精确按 N*A、N*B 大小切好的 segment。

• 后面 batch=N+1，请求略大于原 segment。N*A 段装不下 (N+1)*A，allocator 不得不再 cudaMalloc 一个新段。

• 但旧的 N*A 段不会立刻还给驱动（要 release_cached_blocks 才会），它只是变成"半空"的碎片留在 cache 里。

• 50 层 layer × 多次这样的"扩张"，慢慢就把显存吃光了。

传统 caching allocator 之所以解决不了，是因为它的一个 segment 物理地址和虚拟地址是绑定的——cudaMalloc 给你的是一段连续的 VA + 连续的物理页，不能把后面"长出来"的物理页 append 到前一段后面，所以也就没法把多个相邻 segment 真正拼起来。

Expandable segment 用的是 CUDA 10.2 引入的low level的虚拟内存管理 API，把虚拟地址分配和物理内存分配彻底拆开：

• ExpandableSegment
  • 创建的时候会预留显存乘上 (1 + 1/8)的VA，通过cuMemAddressReserve来做

  • Small pool的物理segment size是2MB，large pool是20MB。就会和上面reserve的VA对应上。注释中也提到了，map的开销和对应物理页的数量是正比的，所以不能都选小的segment size，否则map/unmap的开销会比较大。

    • 小的segment size的好处就是更灵活
  • 维护了一个vector<optional>的结构，为空表示没有做map，handle内保存的是CUmemGenericAllocationHandle，就是create物理页出来的

• map

  • 传入一个ptr, size，表示把这段VA对应的物理页分配了

  • 这里会根据VA的边界按照segment size对齐，然后找出所有空的handle，调用cuMemCreate

  • 然后在cuMemMap，把handle和对应的VA映射上

• unmap

  • cudaStreamSynchronize一下，然后cuMemUnmap断开VA的映射，再把对应range的handle通过cuMemRelease释放掉

然后再来看一下这些API是怎么和DeviceCachingAllocator配合的

• alloc_block
  • find_expandable_block
    • 从unmapped block中找，按照地址的顺序来找，找到第一个block串起来可以大于期望size的block
      • block可能是切分的，所以这里会调用block->next把切分的block都计算上。尝试凑一段连续的地址空间

      • allocated表示是否被外部用到，所以需要找allocated=false的

      • mapped表示是否映射了物理页

    • 如果找不到连续的地址空间的话，就new ExpandableSegment，申请一段新的地址空间。然后申请一个block，block的大小是整个segment的大小，加入到unmapped列表中，再返回

  • map_block

    • ExpandableSegment::map，分配内存并做映射

    • pool.unmapped中erase掉这个block

    • 检查如果mapped的大小小于整个block的size的话，会把剩下的部分拆分成一个新的unmapped block，写回到pool.unmapped中，并维护双向链表

    • 尝试把新map出来的块看一下prev/next，尝试合并成一个新的大块

  • try_allocate_expandable_block

    • 这里会尝试先find_expandable_block

    • 然后通过map_block开始分配物理页，核心逻辑是不断往右找block，尝试map出来，然后把mapped block合并成1个，最后返回

• get_free_block

  • 这里在做best fit选择的时候，有一个区别是不是选大于size的最小的block，而是找到大于size的block之后，尝试往后看一些，看expandable size。根据上面find_expandable_block的流程看，一个block如果可以扩张的话，对应的条件是block->next存在，并且是unmapped。这里会尝试选择一个更小的expandable size的block，从而让在未来希望申请更大内存的时候，可以用到这段VA
• release_blocks
  • 对于expandable segment来说，释放block代表unmap block

  • Expandable segment可以单独进行unmap，不需要像普通路径已经按照segment来释放。

  • unmap给定的ptr range，可能有几种情况。

    • 最简单的就是还有其他人还在使用这个物理页，不需要执行任何的释放操作

    • 或者是释放了部分物理页

    unmap 之前的 block 覆盖区域：
    ┌──────────────────────────────────────────────────────────┐
    │ block (mapped, free, in pool.blocks)                     │
    └──────────────────────────────────────────────────────────┘
       ↑                       ↑                          ↑
     block->ptr             unmapped.ptr               block 末尾
    
    按 page 对齐后能 unmap 的范围：
                              ┌──────────────────────┐
                              │  unmapped (整页对齐)  │
                              └──────────────────────┘
    
    切完之后变成三段：
    ┌──────────────────────┐   ┌──────────────────────┐   ┌──────────────┐
    │ before_free (mapped) │   │ block (unmapped)     │   │ after_free   │
    │  ∈ pool.blocks       │   │  ∈ pool.unmapped     │   │  ∈ pool.blocks│
    └──────────────────────┘   └──────────────────────┘   └──────────────┘
    

    • 此时会把这个block拆成3份，完整释放的unmapped页，写入到unmapped中，后续find block的时候使用

    • 前后两个没有完整覆盖的，把部分块作为free block写回去。

  • 尝试合并block，比如前面的block也是unmapped，则合并成一个大块的unmapped block

Invariant

1. VA 区间是一次 reserve 的连续段：所有同段 block 的 ptr 都落在 [seg->ptr(), seg->ptr() + seg->size()) 内。

2. mapped 部分等于 handles_ 中已贴页的 page 范围之并：mapped block 覆盖的 page 在 handles_[i] 必有值；unmapped block 覆盖的 page handles_[i] == nullopt。

3. mapped 与 unmapped block 各自 page 对齐：因为 map/unmap 都以 segment_size_ 为粒度操作。但 mapped block 内部可被进一步 split（用户态切出小 block），切分点不需要 page 对齐——这是 user 层细粒度切分和 driver 层粗粒度 map 的分层。

cudaFree的同步

在看ExpandableSegment的时候发现有这样一段注释，说cudaFree会做同步，而cuMemUnmap不会做同步。所以在unmap的时候会主动调用一次cudaStreamSynchronize
可能会有一个疑问，free block不是会被process event做了event query才会允许释放么，为什么这里还需要同步？

• Process event只处理record stream的场景，即跨stream的复用。

• 当一个block没有跨stream复用的时候，free就会立刻把这个block放到free block列表中

• 此时如果其他请求需要使用内存，可以复用这个block，后续的kernel launch会使用这段复用的内存，不会有问题

• 如果希望返还这段内存，比如给其他的stream，或者还给driver给其他进程用，此时就会调用cudaFree。这里其实有一个隐式的同步，cudaFree会等所有引用这块内存的请求跑完才会返回。所以比较类似一个cudaStreamSynchronize

• 而cuMemUnmap没有这个保证，立刻unmap可能会影响到inflight的kernel，所以这里在unmap的时候主动调用了一下同步，保证真没有kernel使用这段内存了，才会释放。

TODO

• mempool，用户自定义池？

• 统计信息, trace相关

• Graph capture

• IPC相关