C++ 高级显存池设计：在 C++ AI 框架中利用虚拟内存管理技术实现跨算子生命周期的显存复用协议

各位同仁，下午好！

今天，我们将深入探讨一个在现代AI框架中至关重要，却又极具挑战性的主题：C++高级显存池设计，如何利用虚拟内存管理技术实现跨算子生命周期的显存复用协议。 随着AI模型规模的指数级增长，显存已成为制约性能和模型大小的关键瓶颈。高效的显存管理不再是可选项，而是必须攻克的堡垒。我们将从基础概念出发，逐步构建一个利用GPU虚拟内存能力的显存池，旨在最大化显存利用率、减少碎片化，并提升整体框架性能。

---

1. 引言：AI框架显存管理的痛点与挑战

在深度学习领域，无论是训练还是推理，大规模的张量（Tensors）操作是核心。这些张量通常驻留在GPU显存中。传统的显存管理方式，如每次算子执行时简单地调用 cudaMalloc 分配、cudaFree 释放，带来了诸多问题：

1. 高开销的系统调用： cudaMalloc 和 cudaFree 是同步的GPU驱动调用，开销巨大。频繁调用会显著降低性能。
2. 显存碎片化： 频繁、不规则的分配与释放会导致显存空间出现大量不连续的小块空闲区域，即使总空闲显存充足，也可能无法满足大张量的分配请求，导致OOM (Out Of Memory)。
3. 跨算子生命周期管理困难： 一个张量可能由某个算子产生，作为中间结果被后续多个算子消费，其生命周期横跨多个操作。简单地在算子结束时释放，可能导致数据被过早销毁；延迟释放又可能导致显存占用过高。
4. 难以复用： 许多中间张量在计算完成后立即变得无用，其显存本可以被后续算子复用，但传统方法难以高效实现。
5. 动态形状与不确定性： 许多AI模型支持动态输入形状，导致运行时张量大小不确定，进一步增加了显存预分配和复用的难度。
6. 多流并发： 在异步执行和多流场景下，显存的分配和释放需要严格的同步机制，否则可能出现数据竞争或使用已释放的显存。

为了解决这些问题，显存池（Memory Pool）应运而生。但我们今天要讨论的，不仅仅是简单的显存池，而是如何利用GPU提供的虚拟内存管理能力，构建一个更智能、更高效的显存复用协议。

---

2. GPU虚拟内存管理基础

在深入显存池设计之前，我们首先需要理解GPU上的虚拟内存概念。与CPU类似，现代GPU也支持虚拟内存。NVIDIA CUDA 11.2 及更高版本提供了 CUDA Memory API，允许开发者更精细地控制GPU内存，包括虚拟内存地址空间管理。核心API包括：

• cudaMemAddressReserve(): 预留一个GPU虚拟地址范围。这只是预留地址，并没有实际分配物理显存。
• cudaMemCreate(): 创建一个可映射的物理内存区域（cudaMemHandle）。
• cudaMemMap(): 将一个或多个 cudaMemHandle 映射到之前预留的虚拟地址范围中的某个子区域。
• cudaMemUnmap(): 解除虚拟地址与物理内存的映射。
• cudaMemRelease(): 释放通过 cudaMemCreate() 创建的物理内存。
• cudaMemAddressFree(): 释放通过 cudaMemAddressReserve() 预留的虚拟地址范围。

为何要使用这些API？

传统 cudaMalloc 是直接分配并映射物理显存到虚拟地址。而 cudaMemAddressReserve/cudaMemMap 允许我们将虚拟地址与物理显存解耦。这意味着：

1. 减少碎片化： 我们可以预留一个巨大的连续虚拟地址空间。物理显存可以按页（通常是64KB）进行分配，并按需映射到这个虚拟空间。即使物理显存不连续，通过虚拟地址空间，我们可以给应用程序提供一个连续的视图。
2. 更快的映射/解映射： cudaMemMap 和 cudaMemUnmap 通常比 cudaMalloc 和 cudaFree 快得多，因为它们不涉及物理内存的实际分配或归还给驱动，而只是修改页表。
3. 灵活的复用： 不同的物理内存块可以在不同时间映射到同一个虚拟地址。或者，同一个物理内存块可以被映射到不同的虚拟地址（例如用于零拷贝或共享内存）。这为显存复用提供了极大的灵活性。

核心思想： 我们将维护一个巨大的虚拟地址池和一个物理显存页池。当需要显存时，我们从虚拟地址池中分配一个虚拟地址范围，从物理显存页池中获取（或分配新的）物理页，然后将它们映射起来。当显存不再需要时，我们解除映射，将物理页归还到物理显存页池中，并将虚拟地址范围标记为空闲。

---

3. 传统显存池的局限性

在深入我们高级设计之前，快速回顾几种常见的显存池模式及其局限性：

1. 简单Arena/Bump Allocator:
   • 原理： 预先分配一大块显存，然后通过简单地移动一个“指针”来分配子块。释放操作通常是批量进行，或者不提供单个释放。
   • 优点： 分配极快。
   • 缺点： 无法有效处理不同大小的分配请求，容易造成内部碎片。不支持灵活的单个释放，生命周期管理困难。
2. Buddy System Allocator:
   • 原理： 将显存块递归地二分为“伙伴”，直到找到足够小的块。释放时，如果伙伴块也空闲，则合并。
   • 优点： 减少外部碎片，支持不同大小的分配。
   • 缺点： 内部碎片（例如，分配100KB需要128KB块），实现复杂，分配/释放速度不如Arena。
3. Slab Allocator:
   • 原理： 针对固定大小的对象，预先分配多个“slab”（页），每个slab包含多个相同大小的对象。
   • 优点： 对固定大小对象分配极快，无内部碎片。
   • 缺点： 仅适用于固定大小对象，不通用。
4. CUDA cudaMemPool_t (Stream-Ordered Allocator):
   • 原理： CUDA 11.2 引入的显存池，可以创建设备内存池，并从其中分配。它支持流有序的分配和回收，显著降低了 cudaMalloc 的开销，并能自动处理跨流同步。
   • 优点： 简单易用，性能提升明显，自动同步。
   • 缺点： 仍然是基于物理内存的分配，可能仍面临碎片化问题（尽管内部优化减少了），且对显存复用的控制粒度不如我们自定义的虚拟内存方案精细。它主要优化了 cudaMalloc/cudaFree 的性能，而非彻底解决碎片和跨算子生命周期复用。

我们的目标是超越这些传统方案，利用GPU虚拟内存的强大能力，构建一个更具弹性、更高效的显存池。

---

4. 高级显存池设计：基于虚拟内存管理的复用协议

我们的高级显存池将围绕以下核心组件和理念构建：

1. 统一虚拟地址空间 (UVAS): 预留一个巨大的GPU虚拟地址范围，作为所有显存分配的基石。
2. 物理显存页管理器 (PMM): 负责实际物理显存页的分配、释放和管理。
3. 虚拟地址空间管理器 (VAM): 负责虚拟地址范围的分配和释放。
4. 映射管理器 (MM): 负责将物理显存页映射到虚拟地址空间，以及解除映射。
5. 显存块追踪器 (MBT): 追踪所有已分配显存块的元数据，包括所有者、生命周期、依赖关系、流信息等。
6. 跨算子生命周期复用协议: 定义一套机制，允许框架和算子明确地声明显存的生命周期和依赖，从而实现智能复用。

4.1 核心数据结构设计

我们将定义几个关键的数据结构来管理虚拟和物理显存。

// 显存页大小，通常为64KB，根据实际GPU硬件和CUDA版本可能有所不同
constexpr size_t GPU_PAGE_SIZE = 64 * 1024; // 64KB

// 物理显存块的抽象
struct PhysicalMemBlock {
    cudaMemHandle handle; // CUDA物理内存句柄
    size_t size;         // 实际分配的物理显存大小 (可能大于请求，按页对齐)
    CUdeviceptr base_ptr; // 如果是直接cudaMalloc，记录其地址
    bool is_from_memhandle; // 标记是否通过cudaMemCreate创建
    std::atomic<int> ref_count; // 引用计数，用于物理显存块的生命周期管理
    // ... 其他元数据，如是否空闲，分配时间等
};

// 虚拟地址空间中的一个区域
struct VirtualAddressRegion {
    CUdeviceptr virtual_base_addr; // 虚拟基地址
    size_t size;                   // 区域大小
    bool is_free;                  // 是否空闲
    std::vector<PhysicalMemBlock*> mapped_physical_blocks; // 映射到的物理块
    // ... 其他元数据，如所属的UVAS ID
};

// 实际分配给用户（算子）的显存块信息
struct AllocatedMemoryBlock {
    CUdeviceptr virtual_ptr;         // 返回给用户的虚拟地址
    size_t requested_size;           // 用户请求的大小
    std::string owner_op_name;       // 拥有该显存的算子名称
    cudaStream_t creation_stream;    // 创建时的CUDA流
    std::atomic<int> usage_ref_count; // 用于跨算子生命周期复用的引用计数
    std::vector<cudaEvent_t> completion_events; // 记录使用该显存的算子的完成事件
    VirtualAddressRegion* v_region;  // 指向对应的虚拟地址区域
    // ... 其他元数据，如分配时间、调试信息
};

4.2 GPUMemoryPool 类结构

#include <cuda_runtime.h>
#include <vector>
#include <map>
#include <string>
#include <mutex>
#include <atomic>
#include <list>
#include <algorithm>
#include <stdexcept>

// 辅助函数：将大小向上对齐到页
size_t align_to_page(size_t size) {
    return (size + GPU_PAGE_SIZE - 1) / GPU_PAGE_SIZE * GPU_PAGE_SIZE;
}

class GPUMemoryPool {
public:
    GPUMemoryPool(size_t total_virtual_size, int device_id);
    ~GPUMemoryPool();

    // 分配显存给算子
    // 参数：requested_size - 请求大小，stream - 当前操作所在的流，owner_op_name - 算子名称
    // 返回：指向GPU显存的虚拟地址
    CUdeviceptr allocate(size_t requested_size, cudaStream_t stream, const std::string& owner_op_name);

    // 释放显存
    // 参数：ptr - 待释放的虚拟地址，stream - 当前操作所在的流，owner_op_name - 算子名称
    // 返回：无
    void release(CUdeviceptr ptr, cudaStream_t stream, const std::string& owner_op_name);

    // 声明显存依赖关系
    // 用于更精细的生命周期管理。表示 `dependent_op_name` 依赖 `dependency_ptr`
    void declare_dependency(CUdeviceptr dependency_ptr, const std::string& dependent_op_name);

    // 检查并回收不再使用的物理显存
    void garbage_collect();

    // 获取当前池状态 (调试用)
    void print_stats();

private:
    int device_id_;
    CUdeviceptr reserved_virtual_base_addr_; // 预留的虚拟地址空间基地址
    size_t total_virtual_size_;              // 预留的虚拟地址空间总大小

    std::mutex pool_mutex_; // 保护池内部数据结构

    // 虚拟地址空间管理
    std::list<VirtualAddressRegion> virtual_address_free_list_; // 空闲虚拟地址区域列表
    std::map<CUdeviceptr, VirtualAddressRegion*> virtual_address_map_; // 虚拟地址到区域的映射 (用于快速查找)

    // 物理显存管理
    std::list<PhysicalMemBlock> physical_mem_free_list_; // 空闲物理显存块列表
    std::vector<PhysicalMemBlock> all_physical_mem_blocks_; // 追踪所有物理显存块，便于管理句柄

    // 已分配给算子的显存块追踪
    std::map<CUdeviceptr, AllocatedMemoryBlock> active_allocations_; // 虚拟地址到分配块信息的映射

    // 内部帮助函数
    CUdeviceptr reserve_virtual_address_space(size_t size);
    void free_virtual_address_space(CUdeviceptr base_addr, size_t size);

    PhysicalMemBlock* acquire_physical_mem_block(size_t size);
    void release_physical_mem_block(PhysicalMemBlock* block);

    void map_physical_to_virtual(CUdeviceptr virtual_addr, size_t size, PhysicalMemBlock* physical_block);
    void unmap_physical_from_virtual(CUdeviceptr virtual_addr, size_t size);

    // 查找并合并空闲的虚拟地址区域
    void coalesce_virtual_regions();
};

4.3 GPUMemoryPool 构造与析构

构造函数：
在构造函数中，我们首先预留一大块连续的GPU虚拟地址空间，并将其作为一个大的空闲区域添加到虚拟地址自由列表中。

GPUMemoryPool::GPUMemoryPool(size_t total_virtual_size, int device_id)
    : device_id_(device_id), total_virtual_size_(align_to_page(total_virtual_size)) {

    cudaSetDevice(device_id_);

    // 1. 预留整个虚拟地址空间
    cudaMemAllocationProp prop = {};
    prop.type = cudaMemAllocationType_Pinned; // 必须是Pinned，因为我们要映射到设备
    prop.location.type = cudaMemLocationType_Device;
    prop.location.id = device_id_;

    cudaError_t err = cudaMemAddressReserve(&reserved_virtual_base_addr_, total_virtual_size_, 0, 0, 0);
    if (err != cudaSuccess) {
        throw std::runtime_error("Failed to reserve GPU virtual address space: " + std::string(cudaGetErrorString(err)));
    }

    // 2. 将整个预留空间作为初始空闲区域
    VirtualAddressRegion initial_region = {reserved_virtual_base_addr_, total_virtual_size_, true};
    virtual_address_free_list_.push_back(initial_region);
    virtual_address_map_[initial_region.virtual_base_addr] = &virtual_address_free_list_.back();

    std::cout << "GPU Memory Pool initialized on device " << device_id_ 
              << ". Reserved virtual address space: 0x" << std::hex << reserved_virtual_base_addr_ 
              << " size: " << std::dec << total_virtual_size_ / (1024 * 1024) << " MB." << std::endl;
}

析构函数：
在析构函数中，我们需要释放所有预留的虚拟地址空间和所有实际分配的物理显存。

GPUMemoryPool::~GPUMemoryPool() {
    cudaSetDevice(device_id_);

    // 1. 确保所有映射都已解除
    for (auto const& [ptr, alloc_block] : active_allocations_) {
        // 理论上，在析构前所有分配都应该被释放，这里是防御性编程
        unmap_physical_from_virtual(alloc_block.virtual_ptr, alloc_block.requested_size);
        // 如果有未释放的事件，需要等待
        for (cudaEvent_t event : alloc_block.completion_events) {
            cudaEventSynchronize(event);
            cudaEventDestroy(event);
        }
    }

    // 2. 释放所有通过 cudaMemCreate 创建的物理内存句柄
    for (PhysicalMemBlock& block : all_physical_mem_blocks_) {
        if (block.is_from_memhandle && block.handle) {
            cudaMemRelease(block.handle);
        } else if (block.base_ptr) { // 如果是直接cudaMalloc的，释放
            cudaFree(reinterpret_cast<void*>(block.base_ptr));
        }
    }
    all_physical_mem_blocks_.clear();

    // 3. 释放预留的虚拟地址空间
    cudaError_t err = cudaMemAddressFree(reserved_virtual_base_addr_, total_virtual_size_);
    if (err != cudaSuccess) {
        std::cerr << "Warning: Failed to free reserved GPU virtual address space: " << cudaGetErrorString(err) << std::endl;
    }

    std::cout << "GPU Memory Pool on device " << device_id_ << " destroyed." << std::endl;
}

4.4 allocate 函数实现

allocate 函数是核心。它需要：

1. 找到一个足够大的空闲虚拟地址区域。
2. 获取一个或多个物理显存块（从池中复用或新分配）。
3. 将物理块映射到虚拟地址区域。
4. 记录分配信息。

CUdeviceptr GPUMemoryPool::allocate(size_t requested_size, cudaStream_t stream, const std::string& owner_op_name) {
    std::lock_guard<std::mutex> lock(pool_mutex_);
    size_t aligned_size = align_to_page(requested_size);

    // 1. 查找空闲虚拟地址区域
    VirtualAddressRegion* v_region = nullptr;
    auto it = virtual_address_free_list_.begin();
    while (it != virtual_address_free_list_.end()) {
        if (it->is_free && it->size >= aligned_size) {
            v_region = &(*it);
            break;
        }
        ++it;
    }

    if (!v_region) {
        // 考虑动态扩容虚拟地址空间，但此处为简化暂不实现
        throw std::runtime_error("No sufficient contiguous virtual address space available for " + std::to_string(aligned_size) + " bytes.");
    }

    // 从找到的空闲区域中切分出所需大小的区域
    CUdeviceptr allocated_virtual_addr = v_region->virtual_base_addr;
    if (v_region->size > aligned_size) {
        // 分割空闲区域
        VirtualAddressRegion new_free_region = {
            v_region->virtual_base_addr + aligned_size,
            v_region->size - aligned_size,
            true
        };
        virtual_address_free_list_.push_back(new_free_region);
        virtual_address_map_[new_free_region.virtual_base_addr] = &virtual_address_free_list_.back();

        v_region->size = aligned_size;
    }
    v_region->is_free = false;

    // 2. 获取物理显存块
    PhysicalMemBlock* p_block = acquire_physical_mem_block(aligned_size);
    if (!p_block) {
        // 如果物理显存分配失败，需要回滚虚拟地址区域的状态
        v_region->is_free = true; // 简单回滚
        throw std::runtime_error("Failed to acquire physical memory block for " + std::to_string(aligned_size) + " bytes.");
    }

    // 3. 将物理块映射到虚拟地址
    map_physical_to_virtual(allocated_virtual_addr, aligned_size, p_block);
    v_region->mapped_physical_blocks.push_back(p_block); // 关联虚拟区域与物理块

    // 4. 记录分配信息
    active_allocations_[allocated_virtual_addr] = {
        allocated_virtual_addr,
        requested_size, // 记录请求大小，而非对齐后的大小
        owner_op_name,
        stream,
        1, // 初始引用计数为1
        {}, // 初始无完成事件
        v_region
    };

    return allocated_virtual_addr;
}

acquire_physical_mem_block 辅助函数：

PhysicalMemBlock* GPUMemoryPool::acquire_physical_mem_block(size_t size) {
    // 1. 尝试从空闲列表中查找可复用的物理块
    for (auto it = physical_mem_free_list_.begin(); it != physical_mem_free_list_.end(); ++it) {
        if (it->size >= size) {
            // 找到一个足够大的块，可以复用
            // 考虑分割块，但为简化，直接使用整个块
            it->ref_count = 1; // 重新计数
            PhysicalMemBlock* block = &(*it);
            physical_mem_free_list_.erase(it); // 从空闲列表移除
            return block;
        }
    }

    // 2. 如果没有可复用的，则创建新的物理内存块
    cudaMemAllocationProp prop = {};
    prop.type = cudaMemAllocationType_Pinned;
    prop.location.type = cudaMemLocationType_Device;
    prop.location.id = device_id_;

    cudaMemHandle handle;
    cudaError_t err = cudaMemCreate(&handle, size, &prop, 0);
    if (err != cudaSuccess) {
        std::cerr << "Failed to create physical memory handle: " << cudaGetErrorString(err) << std::endl;
        // 尝试回退到 cudaMalloc 模式，如果 cudaMemCreate 失败
        void* ptr;
        err = cudaMalloc(&ptr, size);
        if (err == cudaSuccess) {
            PhysicalMemBlock new_block = {nullptr, size, reinterpret_cast<CUdeviceptr>(ptr), false, 1};
            all_physical_mem_blocks_.push_back(new_block);
            return &all_physical_mem_blocks_.back();
        }
        return nullptr; // 物理内存分配失败
    }

    PhysicalMemBlock new_block = {handle, size, 0, true, 1};
    all_physical_mem_blocks_.push_back(new_block);
    return &all_physical_mem_blocks_.back();
}

map_physical_to_virtual 辅助函数：

void GPUMemoryPool::map_physical_to_virtual(CUdeviceptr virtual_addr, size_t size, PhysicalMemBlock* physical_block) {
    cudaError_t err;
    if (physical_block->is_from_memhandle) {
        err = cudaMemMap(virtual_addr, size, 0, physical_block->handle, 0);
    } else {
        // 如果物理块是直接cudaMalloc的，需要将其内容拷贝到虚拟地址对应的物理页
        // 实际上，cudaMemMap 也可以将 cudaMalloc 的地址作为 handle 使用，但更推荐 MemHandle
        // 对于直接 cudaMalloc 的块，需要通过 cudaMemMapExternalArray 这样的高级API，或者直接使用 cudaMemPool
        // 为了简化，我们假设所有物理块都通过 cudaMemCreate 获取 MemHandle
        // 真实场景中，如果混合使用，需要更复杂的逻辑
        throw std::runtime_error("Mapping cudaMalloc'd memory without handle is not directly supported by this simplified example.");
    }

    if (err != cudaSuccess) {
        throw std::runtime_error("Failed to map physical memory to virtual address: " + std::string(cudaGetErrorString(err)));
    }
}

4.5 release 函数实现

release 函数的逻辑要复杂得多，因为它需要考虑跨算子生命周期复用和流同步。

1. 查找分配记录： 根据虚拟地址找到对应的 AllocatedMemoryBlock。
2. 更新引用计数： 减少 usage_ref_count。
3. 流同步： 记录当前流的完成事件。如果该显存块被多个算子使用，需要确保所有使用它的算子都完成后才能真正释放。
4. 真正释放条件： 当 usage_ref_count 降到 0，并且所有相关的CUDA事件都已完成时，才能进行物理内存的解除映射和回收。

void GPUMemoryPool::release(CUdeviceptr ptr, cudaStream_t stream, const std::string& owner_op_name) {
    std::lock_guard<std::mutex> lock(pool_mutex_);

    auto it_alloc = active_allocations_.find(ptr);
    if (it_alloc == active_allocations_.end()) {
        std::cerr << "Warning: Attempted to release unmanaged or already released memory at 0x" << std::hex << ptr << std::endl;
        return;
    }

    AllocatedMemoryBlock& alloc_block = it_alloc->second;

    // 1. 记录当前流的完成事件 (表示当前操作者已完成对该内存的使用)
    cudaEvent_t event;
    cudaEventCreate(&event);
    cudaEventRecord(event, stream);
    alloc_block.completion_events.push_back(event);

    // 2. 减少引用计数
    int current_ref_count = --alloc_block.usage_ref_count;

    // 3. 如果引用计数降到0，可以考虑异步回收
    if (current_ref_count == 0) {
        // 等待所有使用该显存的算子完成
        for (cudaEvent_t comp_event : alloc_block.completion_events) {
            cudaEventSynchronize(comp_event); // 等待事件完成
            cudaEventDestroy(comp_event);     // 销毁事件
        }
        alloc_block.completion_events.clear();

        // 解除物理内存映射
        unmap_physical_from_virtual(alloc_block.virtual_ptr, align_to_page(alloc_block.requested_size));

        // 将对应的虚拟地址区域标记为空闲
        VirtualAddressRegion* v_region = alloc_block.v_region;
        if (v_region) {
            v_region->is_free = true;
            v_region->mapped_physical_blocks.clear(); // 清除映射关系
            coalesce_virtual_regions(); // 尝试合并空闲虚拟区域
        }

        // 释放物理显存块 (归还到空闲列表或真正释放)
        // 这里需要找到物理块，然后减少其引用计数。如果物理块引用计数也为0，则可以放回物理池或释放。
        // 为了简化，我们假设每个虚拟区域只映射到一个物理块，且物理块生命周期与虚拟区域一致。
        // 真实情况可能一个物理块被多个虚拟区域共享，需要独立的物理块引用计数。
        // 这里我们复用 PhysicalMemBlock 里的 ref_count 作为物理块的引用计数
        if (!v_region->mapped_physical_blocks.empty()) {
            PhysicalMemBlock* p_block = v_region->mapped_physical_blocks[0]; // 假设只有一个
            if (--p_block->ref_count == 0) {
                release_physical_mem_block(p_block);
            }
        }

        active_allocations_.erase(it_alloc);
    }
    // 如果引用计数不为0，则只是标记当前算子不再使用，内存块仍然活跃
}

unmap_physical_from_virtual 辅助函数：

void GPUMemoryPool::unmap_physical_from_virtual(CUdeviceptr virtual_addr, size_t size) {
    cudaError_t err = cudaMemUnmap(virtual_addr, size);
    if (err != cudaSuccess) {
        std::cerr << "Warning: Failed to unmap physical memory from virtual address 0x" << std::hex << virtual_addr 
                  << ": " << cudaGetErrorString(err) << std::endl;
    }
}

release_physical_mem_block 辅助函数：

void GPUMemoryPool::release_physical_mem_block(PhysicalMemBlock* block) {
    // 将物理块归还到空闲列表
    physical_mem_free_list_.push_back(*block);
    // 实际的 cudaMemRelease 可以在垃圾回收阶段进行，而不是立即执行
    // 这样可以避免频繁的 cudaMemCreate/Release
}

coalesce_virtual_regions 辅助函数：

void GPUMemoryPool::coalesce_virtual_regions() {
    // 遍历空闲列表，查找相邻的空闲区域并合并
    virtual_address_free_list_.sort([](const VirtualAddressRegion& a, const VirtualAddressRegion& b) {
        return a.virtual_base_addr < b.virtual_base_addr;
    });

    auto it = virtual_address_free_list_.begin();
    while (it != virtual_address_free_list_.end()) {
        if (!it->is_free) {
            ++it;
            continue;
        }

        auto next_it = std::next(it);
        if (next_it != virtual_address_free_list_.end() && next_it->is_free &&
            it->virtual_base_addr + it->size == next_it->virtual_base_addr) {

            it->size += next_it->size;
            virtual_address_map_.erase(next_it->virtual_base_addr);
            it = virtual_address_free_list_.erase(next_it); // 移除并更新迭代器
        } else {
            ++it;
        }
    }
}

4.6 跨算子生命周期复用协议

这是显存池智能化的关键。

1. 引用计数 (usage_ref_count): 最简单的协议。当一个算子需要某个张量时，它“引用”该张量对应的显存块，引用计数加1。当算子完成对该张量的使用时，引用计数减1。只有当引用计数为0时，显存块才被视为可回收。
2. 显式依赖声明 (declare_dependency): 更高级的协议。AI框架的计算图可以提供算子之间的依赖关系。当算子A的输出是算子B的输入时，算子B就依赖于算子A的输出显存。
   • declare_dependency(producer_ptr, consumer_op_name): 显式声明 consumer_op_name 将使用 producer_ptr 指向的显存。这可以增加 producer_ptr 对应显存块的引用计数。
   • 在计算图编译阶段，框架可以分析出每个张量的生命周期，从而在运行时告知显存池何时增加/减少引用计数。

declare_dependency 实现：

void GPUMemoryPool::declare_dependency(CUdeviceptr dependency_ptr, const std::string& dependent_op_name) {
    std::lock_guard<std::mutex> lock(pool_mutex_);
    auto it_alloc = active_allocations_.find(dependency_ptr);
    if (it_alloc == active_allocations_.end()) {
        std::cerr << "Warning: Attempted to declare dependency on unmanaged memory at 0x" << std::hex << dependency_ptr << std::endl;
        return;
    }
    // 增加引用计数，表示有新的使用者
    it_alloc->second.usage_ref_count++;
    // 可以在这里记录依赖关系，例如一个列表 `std::vector<std::string> dependent_ops;`
    // 但核心是引用计数来驱动生命周期
}

工作流示例：

1. Op1 allocate(size_A, stream1, "Op1_Output") -> ptr_A (ref_count = 1)
2. Op2 (需要 ptr_A 作为输入) declare_dependency(ptr_A, "Op2_Input") -> ptr_A (ref_count = 2)
3. Op1 完成计算，执行 release(ptr_A, stream1, "Op1_Output") -> ptr_A (ref_count = 1, Op1_event_recorded)
4. Op3 (需要 ptr_A 作为输入) declare_dependency(ptr_A, "Op3_Input") -> ptr_A (ref_count = 2)
5. Op2 完成计算，执行 release(ptr_A, stream2, "Op2_Input") -> ptr_A (ref_count = 1, Op2_event_recorded)
6. Op3 完成计算，执行 release(ptr_A, stream3, "Op3_Input") -> ptr_A (ref_count = 0, Op3_event_recorded)
   • 此时，ptr_A 的 usage_ref_count 变为0。
   • 显存池会等待 Op1_event, Op2_event, Op3_event 全部完成。
   • 所有事件完成后，解除 ptr_A 的物理映射，并将虚拟地址区域和物理显存块标记为空闲，以便复用。

4.7 垃圾回收 (garbage_collect)

garbage_collect 函数可以周期性运行，负责清理那些引用计数已归零但尚未被回收的物理显存块。这些块可能因为等待CUDA事件完成而被延迟回收。

void GPUMemoryPool::garbage_collect() {
    std::lock_guard<std::mutex> lock(pool_mutex_);

    // 遍历 physical_mem_free_list_，寻找那些 ref_count 真的为 0，并且可以被 cudaMemRelease 的块
    // 或者，定期清理那些长时间未被使用的物理内存，将其真正的 cudaMemRelease 掉
    // 这是一个策略问题：是尽可能持有物理内存以备复用，还是及时释放给OS？
    // 通常在AI框架中，会倾向于持有，因为物理内存分配昂贵。
    // 这里的 garbage_collect 可以更侧重于合并空闲的物理块，或者根据某种LRU策略释放物理块。

    // 示例：合并相邻的空闲物理块
    // (需要 physical_mem_free_list_ 按照地址排序)
    physical_mem_free_list_.sort([](const PhysicalMemBlock& a, const PhysicalMemBlock& b) {
        return a.base_ptr < b.base_ptr; // 如果是 cudaMalloc 的块
        // 对于 MemHandle 的块，合并更复杂，因为它们是逻辑块而非连续地址
    });

    // 实际的 GC 策略会更加复杂，可能包括：
    // - 释放超过 TTL (Time To Live) 的物理块
    // - 根据显存压力动态调整物理块的释放策略
}

4.8 性能与碎片化分析

特性	简单 cudaMalloc/Free	cudaMemPool_t	本文高级虚拟池方案
分配/释放性能	低	高 (池内)	极高 (映射/解映射)
外部碎片化	高	中	极低 (虚拟地址连续)
内部碎片化	低	中 (对齐问题)	低 (按页对齐)
跨算子复用	困难	需额外逻辑	内置引用计数/依赖图
显存利用率	较低	中	极高
实现复杂度	低	中	高
同步开销	高 (隐式)	自动 (流有序)	显式事件/引用计数
GPU虚拟地址利用	无	部分 (内部实现)	核心机制

显存复用流程图示 (简化):

[预留巨大虚拟地址空间]
        |
        V
[PMM: 物理显存页池] <-----> [VAM: 虚拟地址区域池]
        |                       |
        | 分配物理页            | 分配虚拟区域
        V                       V
[Op1 请求显存 (Size A)] --map--> [虚拟地址 VA1]
        |                           ^
        |                           |
        |                           |
[Op2 依赖 VA1, 增加 RefCount]      |
        |                           |
[Op1 释放 VA1, 减少 RefCount, 记录事件]
        |                           |
[Op3 依赖 VA1, 增加 RefCount]      |
        |                           |
[Op2 释放 VA1, 减少 RefCount, 记录事件]
        |                           |
[Op3 释放 VA1, RefCount = 0, 记录事件]
        |
        V
[等待所有事件完成]
        |
        V
[解除 VA1 映射]
        |
        V
[将物理页归还 PMM, 虚拟区域归还 VAM]

---

5. 高级考量与未来方向

5.1 多GPU支持
在多GPU系统中，显存池需要管理每个设备的虚拟地址空间和物理显存。cudaMemcpyPeer 可以用于GPU之间的数据传输，但内存池本身需要支持跨设备分配，或者为每个设备维护一个独立的池。更高级的方案是使用 cudaMallocManaged 配合 cudaMemAdvise 和 cudaMemPrefetchAsync 实现统一内存，但 cudaMemAddressReserve/cudaMemMap 仍提供了更细粒度的控制。

5.2 显存局部性与NUMA
对于拥有多个显存控制器（HBM stacks）的GPU，显存局部性可能影响性能。我们的池可以尝试根据计算核的亲和性来分配物理显存，或者通过分析访问模式进行优化。

5.3 框架集成与图优化
最好的显存复用发生在计算图编译阶段。AI框架（如TensorFlow XLA, PyTorch Inductor, ONNX Runtime）可以通过静态分析计算图，确定每个中间张量的生命周期，生成更优化的显存分配和复用计划。我们的显存池可以作为这些图优化器的后端，提供底层的虚拟内存管理能力。

5.4 性能监控与调试
显存泄漏和碎片化是常见问题。需要集成工具来监控显存池的利用率、碎片程度、活跃分配、引用计数等，以便及时发现并解决问题。CUDA提供了一些API用于查询显存信息。

5.5 错误处理与鲁棒性
显存分配失败（OOM）是AI训练中常见的错误。显存池需要健壮的错误处理机制，包括回滚分配、提供详细错误信息，甚至在OOM时尝试进行更激进的垃圾回收策略。

5.6 动态调整与自适应
显存池可以根据运行时负载和显存压力动态调整其策略，例如在显存充足时倾向于保留物理显存以备复用，在显存紧张时更积极地释放物理显存。

---

6. 总结与展望

我们今天深入探讨了如何利用GPU的虚拟内存管理能力，构建一个高级C++显存池，以实现跨算子生命周期的显存复用。通过将虚拟地址空间与物理显存解耦，我们能够有效对抗碎片化，显著提升显存分配和回收的性能，并为AI框架提供更灵活、更智能的显存复用机制。这种设计虽然实现复杂，但在大规模深度学习模型训练和部署中，其带来的性能和显存利用率提升是巨大的。

未来，随着GPU架构和CUDA技术的不断演进，显存管理将继续向更智能、更自适应的方向发展。将显存池与AI框架的计算图优化深度融合，将是实现极致性能的关键。