干货｜CUDA内存分配深度优化：从池化、预取到虚拟内存管理实战

在前文探讨了CUDA内存管理与“内存预取”优化方法之后，本文将更深入地分析CUDA的整体内存优化策略。我们会按照不同的层次来拆解与之相关的技术，并探讨如何协调各层次的优化，以期实现整体性能的最大化。

影响性能的限制条件

抛开硬件的物理限制，软件的性能瓶颈往往出现在各层级之间以及接口的数据传递上。对于开发者而言，在运行环境（如CPU或GPU）性能达标之后，数据管理操作就成了关键。其中，内存的管理与应用是最首要的限制条件。不当的（显）存分配与管理会导致无谓的性能损耗，这也正是内存需要被“增强”或优化的前提。

在CUDA编程中，内存相关的优化可以分为三个层次：

1. 主机端内存管理
   任何有过操作系统学习经验的人都明白，内存资源总是紧张的。虚拟内存、交换技术等的出现，本质上都是为了缓解内存压力。特别是当缓存命中率失效导致数据重新加载时，性能会急剧下降。因此，如果在主机端能将关键内存“锁定”，防止其被交换到硬盘，或者提高其缓存命中率，再辅以DMA等技术，就能从整体上提升数据处理速度。异步I/O（IO overlapped）技术的应用也能显著提高资源利用率，尽管它确实增加了编程的复杂性。
   在CUDA的发展历程中，从早期的 cudaHostAlloc()、cudaMallocHost()，到CUDA 11.2之后的 cudaMemcpyAsync() 系列函数，都能实现物理页的锁定。特别是CUDA 11.2引入的流顺序内存分配器，通过 cudaMallocAsync() 和 cudaFreeAsync() 等接口，将内存分配与CUDA流深度绑定，实现了更精细的异步内存管理。

2. GPU的内存优化与重用
   有过大规模软件开发经验的开发者通常熟悉内存池技术。内存的分配与回收本身是耗时且复杂的操作，如果能够分配一块内存并反复使用，就能大幅减少系统在这方面的开销。
   因此，GPU生态中也提供了内存池化技术，例如RAPIDS内存管理器。其常见实现方式是预先分配一大块内存，然后在其内部进行细分和复用。除了具备减少内存碎片等传统优势，它还提升了主机与设备间内存操作的异步性。当然，池化也有其缺点，例如在面对复杂多变的内存分配需求时可能导致内存浪费。
   而CUDA提供的统一内存及其预取功能，进一步提升了内存使用效率。与内存池技术结合，能更好地发挥两者优势。类似地，统一内存中的内存建议机制，为托管内存编程提供了可控性，有助于在长期运行中提升性能。
   此外，池化的内存还可以支持流感知的分配，即将内存池与特定的CUDA流绑定，实现内存管理与流执行的同步。

3. 新技术的应用
   硬件在不断迭代，软件平台也必须随之演进。例如，随着NVIDIA H200等新一代显卡的推出，内存管理技术也需要相应调整，甚至可能发生变革。
   传统的NUMA内存管理方式通常将GPU显存作为一个节点暴露给操作系统管理，这可能导致CPU与GPU间频繁的数据交换，引发性能波动。为此，CUDA推出了CDMM模式，即基于驱动程序的连贯内存管理。它主要面向Grace Hopper/Blackwell架构，不再直接将GPU显存暴露给系统，而是通过驱动程序来管理，从而进行更精细的粒度控制，提升性能。在容器化部署场景下，CDMM模式展现出了更好的适应性。

AI技术快速迭代，作为底层基础软件之一的CUDA也必须紧跟硬件与上层应用的发展。技术体系正是在传统与创新的不断融合、量变到质变的过程中，推动着整个应用系统向前。

内存优化的整体技术及在CUDA中的应用

分析了内存对CUDA应用的影响后，我们来看看有哪些通用的内存优化技术，以及它们在CUDA中的具体体现：

1. 内存池
   内存池化能有效控制分配与回收的性能损耗和瓶颈，促进内存重用，并显著减少内存碎片。这也是为什么从应用层、中间件、C/C++库到底层操作系统都广泛采用该技术的原因。前文提到的流顺序内存分配器便显式地提供了池化操作特性，并与CUDA流感知绑定。

2. 零拷贝与显存直接操作
   在研究网络高并发和数据库底层读写时，Direct I/O和DMA技术被反复提及。它们对于GPU显存操作同样有效。GDS正是这一思路的发展方向，它允许数据在NVMe存储和GPU显存之间直接传输，绕过CPU和主机内存。

3. 内存预取、建议与写合并内存
   这三者很好理解：预取是为了减少数据加载的等待时间；建议是让内存管理器根据当前场景优化访问策略；写合并则是减少对内存的写入次数。这些都是非常传统且有效的内存使用增强手段。

4. 虚拟内存管理
   这里的虚拟内存管理不仅指传统概念，更倾向于一种管理机制。它通过虚拟内存的抽象，实现了CUDA内存更灵活、高效的管理，对内存池等技术更加友好，尤其在大内存分配时优势明显。
   由于虚拟地址空间与物理显存地址空间解耦，两者之间的地址映射（包括多级映射）、分段管理、内存共享等都变得更加方便快捷。

5. 合并应用
   这属于一种典型的工程处理机制。例如，在C++智能指针中，可以通过自定义删除器来适配不同的内存回收策略。CUDA中也提供了类似的机制，但其功能的完善程度和易用性需要在实际应用中检验。

“增强内存分配”对CUDA而言是一个持续演进的方向。它从最初简单粗暴的管理，逐步发展到支持流感知、异步处理和内存池化，不断适配新的软硬件平台和库，有机地融合各种技术以提升GPU显存的利用率和整体性能。

例程：虚拟内存管理实战

下面通过一个具体的例程来演示CUDA中的虚拟内存管理：

#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include <stdio.h>
#include <cuda.h>
#include <stdlib.h>

int main() {
    cuInit(0);

    CUdevice cuDev;
    cuDeviceGet(&cuDev, 0);

    CUcontext cuCtx;
    cuCtxCreate(&cuCtx, 0, cuDev);

    // query allocation granularity
    CUmemAllocationProp prop = {};
    prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;         // fix memory allocation type
    prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;  // the memory will be allocated on the cuDev
    prop.location.id = cuDev;                          // specify the cuDev for allocation

    size_t gran;
    cuMemGetAllocationGranularity(&gran, &prop,
        CU_MEM_ALLOC_GRANULARITY_RECOMMENDED);
    printf("alloc granularity: %zu bytes\n", gran);

    // set allocation size (must be a multiple of the granularity)
    size_t allocSize = 1024 * 1024; // 1 MB
    allocSize = (allocSize + gran - 1) / gran * gran;
    printf("alloc size (aligned): %zu bytes\n", allocSize);

    CUmemGenericAllocationHandle cuMemHandle;
    cuMemCreate(&cuMemHandle, allocSize, &prop, 0);

    CUdeviceptr cuDevPtr;
    cuMemAddressReserve(&cuDevPtr, allocSize, 0, 0, 0);
    printf("Reserved virtual address: 0x%llx\n", (unsigned long long)cuDevPtr);

    // pyhsically map the allocated memory to the reserved virtual address range
    cuMemMap(cuDevPtr, allocSize, 0, cuMemHandle, 0);

    //  set access permissions for the allocated memory: allow read/write access from the cuDev
    CUmemAccessDesc cuAccesDesc = {};
    cuAccesDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
    cuAccesDesc.location.id = cuDev;
    cuAccesDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;  // read/write access

    cuMemSetAccess(cuDevPtr, allocSize, &cuAccesDesc, 1);

    // use the allocated memory - set it to a specific value using cuMemsetD8
    cuMemsetD8(cuDevPtr, 0xCB, allocSize);

    // copy and check
    unsigned char* hostBuf = (unsigned char*)malloc(allocSize);
    cuMemcpyDtoH(hostBuf, cuDevPtr, allocSize);

    // display the first few bytes to verify
    printf("display mem content after memset: ");
    for (int i = 0; i < 16; ++i) {
        printf("%02x ", hostBuf[i]);
    }
    printf("\n");

    // clean up
    free(hostBuf);

    //umap
    cuMemUnmap(cuDevPtr, allocSize);

    // free virtual address
    cuMemAddressFree(cuDevPtr, allocSize);

    // free physical memory allocation
    cuMemRelease(cuMemHandle);

    // destroy cuCtx
    cuCtxDestroy(cuCtx);

    return 0;
}

这段代码简要演示了如何使用CUDA驱动API进行物理内存分配、虚拟地址预留、地址映射以及内存操作。如果对某个API不清楚，可以直接查阅NVIDIA官方文档获取详细描述。

编译说明：在Visual Studio 2022中，需要在项目属性 -> 链接器 -> 输入 -> 附加依赖项中添加 cuda.lib。

总结

通过以上分析，我们可以看到，解决问题的核心思想往往是相通的。这正是抽象理论指导具体实践的例证。GPU显存的处理，与传统内存管理和优化的指导原则并无本质不同。优秀的设计者，往往是在深刻理解传统思想的基础上，在特定平台上做出独到的发展和应用的。

而实践的过程，又会反过来推动传统思想的进步与发展，二者形成有机的统一，共同推进整个技术生态的演进！如果你对高性能计算和底层优化有更多兴趣，欢迎在云栈社区与更多开发者交流探讨。