📚 课程目录第3章 GPU内存分配与注册

3. GPU内存分配与注册

任务目标

本案例演示如何在GPU显存中分配内存，并将其注册为RDMA可访问的内存区域。完成本案例后，您将掌握使用cudaMalloc分配GPU内存、通过ibv_reg_mr注册内存区域的基本流程，为后续GPU与RDMA的直接数据传输（GPUDirect RDMA）打下基础。

环境依赖

NVIDIA GPU（支持UVA，Compute Capability 2.0及以上）
CUDA Toolkit（≥10.0）
RDMA网卡（支持GPUDirect RDMA，如Mellanox ConnectX系列）
libibverbs-dev、librdmacm-dev
编译命令需包含CUDA与libibverbs链接参数

核心实现

以下代码完成：分配GPU显存 → 获取设备指针 → 注册为RDMA内存区域 → 验证注册结果 → 清理资源。关键点在于ibv_reg_mr需要传入GPU内存的设备虚拟地址（即cudaMalloc返回的指针），且内存区域必须保持有效直到注销。

#include <cuda_runtime.h>
#include <infiniband/verbs.h>
#include <cstdio>
#include <cstdlib>

#define CHECK_CUDA(call) do { \
    cudaError_t err = (call); \
    if (err != cudaSuccess) { \
        fprintf(stderr, "CUDA error %s at %s:%d\n", \
                cudaGetErrorString(err), __FILE__, __LINE__); \
        exit(EXIT_FAILURE); \
    } \
} while (0)

#define CHECK_IB(call) do { \
    int ret = (call); \
    if (ret) { \
        fprintf(stderr, "IB error %d at %s:%d\n", ret, __FILE__, __LINE__); \
        exit(EXIT_FAILURE); \
    } \
} while (0)

int main() {
    const size_t buf_size = 4096;  // 4KB GPU内存
    void *gpu_buf = nullptr;

    // 1. 分配GPU显存
    CHECK_CUDA(cudaMalloc(&gpu_buf, buf_size));
    printf("GPU buffer allocated at %p, size %zu bytes\n", gpu_buf, buf_size);

    // 2. 打开RDMA设备（此处使用第一个可用设备）
    struct ibv_device **dev_list = ibv_get_device_list(nullptr);
    if (!dev_list || !dev_list[0]) {
        fprintf(stderr, "No RDMA device found\n");
        return EXIT_FAILURE;
    }
    struct ibv_context *ctx = ibv_open_device(dev_list[0]);
    if (!ctx) {
        fprintf(stderr, "Failed to open RDMA device\n");
        return EXIT_FAILURE;
    }
    printf("Opened RDMA device: %s\n", ibv_get_device_name(dev_list[0]));
    ibv_free_device_list(dev_list);

    // 3. 创建保护域（PD）
    struct ibv_pd *pd = ibv_alloc_pd(ctx);
    if (!pd) {
        fprintf(stderr, "Failed to allocate protection domain\n");
        return EXIT_FAILURE;
    }

    // 4. 注册GPU内存为RDMA内存区域
    //    注意：ibv_reg_mr的addr参数直接使用cudaMalloc返回的GPU虚拟地址
    struct ibv_mr *mr = ibv_reg_mr(pd, gpu_buf, buf_size,
                                   IBV_ACCESS_LOCAL_WRITE |
                                   IBV_ACCESS_REMOTE_READ |
                                   IBV_ACCESS_REMOTE_WRITE);
    if (!mr) {
        fprintf(stderr, "Failed to register GPU memory for RDMA\n");
        return EXIT_FAILURE;
    }
    printf("Registered MR: lkey=0x%x, rkey=0x%x, addr=%p, length=%zu\n",
           mr->lkey, mr->rkey, mr->addr, mr->length);

    // 5. 验证：检查注册的地址与长度是否与分配一致
    if (mr->addr != gpu_buf || mr->length != buf_size) {
        fprintf(stderr, "MR registration mismatch!\n");
        return EXIT_FAILURE;
    }
    printf("Verification passed: MR address and length match GPU allocation\n");

    // 6. 清理资源
    CHECK_IB(ibv_dereg_mr(mr));
    ibv_dealloc_pd(pd);
    ibv_close_device(ctx);
    CHECK_CUDA(cudaFree(gpu_buf));

    printf("All resources cleaned up successfully\n");
    return EXIT_SUCCESS;
}

编译与运行

nvcc -o gpu_mr_register gpu_mr_register.c -libverbs -lcuda
sudo ./gpu_mr_register   # 需要root权限或CAP_SYS_NICE以锁定内存

验证方法

程序执行后，观察输出是否包含以下关键信息：

GPU buffer allocated：显示GPU内存分配成功及地址
Opened RDMA device：RDMA设备正常打开
Registered MR：显示lkey/rkey及注册地址，其中addr应与GPU分配地址一致
Verification passed：程序自动校验MR注册信息与GPU分配一致
All resources cleaned up：资源释放无错误

若程序退出码为0且包含上述输出，则判定通过。若出现CUDA error或IB error，请检查GPU驱动、RDMA驱动及权限设置。

常见问题

注册失败（MR为NULL）：确认GPU内存已分配且未释放；检查RDMA网卡是否支持GPUDirect RDMA（使用ibv_devinfo -v查看）；确认进程有足够内存锁定额度（ulimit -l）
CUDA错误"invalid argument"：cudaMalloc返回的指针必须直接用于ibv_reg_mr，不能经过任何地址转换
编译错误找不到libibverbs：安装libibverbs-dev包，并确认链接路径