4. RDMA读GPU内存

任务目标

本案例实现一个最小可复现的RDMA读操作：客户端（主动方）通过InfiniBand RDMA Read从服务端（被动方）的GPU显存中读取数据。服务端将数据写入GPU显存并注册为可读内存区域，客户端发起单边RDMA Read操作直接读取该数据，无需服务端CPU参与数据拷贝。

环境与依赖

• 硬件：InfiniBand HCA（支持RDMA Read）、NVIDIA GPU（同一节点或不同节点均可）
• 软件：Linux内核含InfiniBand驱动、libibverbs、librdmacm、CUDA Toolkit（≥11.0）
• 编译依赖：cuda_runtime.h、infiniband/verbs.h、rdma/rdma_cma.h
• 连接方式：服务端监听，客户端主动连接后发起RDMA Read

核心实现

代码包含服务端与客户端两个角色，通过rdma_cm建立连接，交换QP信息与GPU内存的RKey。服务端将数据写入GPU显存并注册为可读内存区域，客户端发起单边RDMA Read操作直接读取该数据，无需服务端CPU参与数据拷贝。服务端使用cudaMalloc分配显存，cudaMemcpy写入数据，再通过ibv_reg_mr注册该显存区域。客户端使用ibv_post_srq接收请求，ibv_post_read发起RDMA Read。为简化流程，代码使用同步轮询完成事件处理，不使用异步回调。

关键完整代码

// rdma_read_gpu.c — 编译: mpicc -o rdma_read_gpu rdma_read_gpu.c -lrdmacm -libverbs -lcuda
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <infiniband/verbs.h>
#include <rdma/rdma_cma.h>
#include <cuda_runtime.h>

#define DATA_SIZE 4096
#define BUFFER_SIZE (DATA_SIZE + 4096) // 对齐

struct conn_info {
    uint64_t addr;
    uint32_t rkey;
    uint32_t qp_num;
};

static int post_read(struct rdma_cm_id *id, uint64_t remote_addr, uint32_t rkey,
                     void *local_buf, uint32_t length, struct ibv_mr *mr) {
    struct ibv_sge sge = {.addr = (uint64_t)local_buf, .length = length, .lkey = mr->lkey};
    struct ibv_send_wr wr = {
        .opcode = IBV_WR_RDMA_READ,
        .send_flags = IBV_SEND_SIGNALED,
        .wr.rdma.remote_addr = remote_addr,
        .wr.rdma.rkey = rkey,
        .sg_list = &sge,
        .num_sge = 1
    };
    struct ibv_send_wr *bad_wr = NULL;
    return ibv_post_send(id->qp, &wr, &bad_wr);
}

static int run_server(struct rdma_cm_id *listen_id) {
    struct rdma_cm_id *conn_id = NULL;
    struct rdma_conn_param param = {0};
    struct rdma_conn_param accepted = {0};
    struct ibv_mr *gpu_mr = NULL;
    char *gpu_buf = NULL;
    struct conn_info local_info, remote_info;
    int ret = 0;

    if (rdma_get_request(listen_id, &conn_id)) { perror("get_request"); return -1; }

    // 分配GPU内存并写入数据
    if (cudaMalloc(&gpu_buf, BUFFER_SIZE) != cudaSuccess) { fprintf(stderr, "cudaMalloc failed\n"); return -1; }
    char *host_data = malloc(DATA_SIZE);
    memset(host_data, 'R', DATA_SIZE);
    if (cudaMemcpy(gpu_buf, host_data, DATA_SIZE, cudaMemcpyHostToDevice) != cudaSuccess) {
        fprintf(stderr, "cudaMemcpy failed\n"); free(host_data); return -1;
    }
    free(host_data);

    // 注册GPU显存
    gpu_mr = ibv_reg_mr(conn_id->pd, gpu_buf, BUFFER_SIZE,
                        IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ);
    if (!gpu_mr) { perror("reg_mr"); return -1; }

    // 交换地址信息
    local_info.addr = (uint64_t)gpu_buf;
    local_info.rkey = gpu_mr->rkey;
    local_info.qp_num = conn_id->qp->qp_num;
    param.private_data = &local_info;
    param.private_data_len = sizeof(local_info);
    param.responder_resources = 1;
    param.initiator_depth = 1;
    param.retry_count = 7;

    if (rdma_accept(conn_id, &param)) { perror("accept"); return -1; }

    struct rdma_cm_event *event;
    if (rdma_get_cm_event(listen_id->channel, &event)) { perror("get_event"); return -1; }
    if (event->event != RDMA_CM_EVENT_ESTABLISHED) { fprintf(stderr, "Unexpected event\n"); return -1; }
    memcpy(&remote_info, event->param.conn.private_data, sizeof(remote_info));
    rdma_ack_cm_event(event);

    // 等待客户端完成读取
    struct ibv_wc wc;
    while (rdma_get_recv_comp(conn_id, &wc) != 0) { /* 轮询 */ }
    if (wc.status != IBV_WC_SUCCESS) { fprintf(stderr, "WC error: %s\n", ibv_wc_status_str(wc.status)); ret = -1; }

    // 验证数据（读回GPU内存检查）
    char *verify = malloc(DATA_SIZE);
    cudaMemcpy(verify, gpu_buf, DATA_SIZE, cudaMemcpyDeviceToHost);
    int ok = 1;
    for (int i = 0; i < DATA_SIZE; i++) if (verify[i] != 'R') { ok = 0; break; }
    printf("Server: data %s\n", ok ? "OK" : "MISMATCH");
    free(verify);

    rdma_disconnect(conn_id);
    ibv_dereg_mr(gpu_mr);
    cudaFree(gpu_buf);
    rdma_destroy_id(conn_id);
    return ret;
}

static int run_client(struct rdma_cm_id *id) {
    struct rdma_conn_param param = {0};
    struct conn_info local_info, remote_info;
    char *local_buf = malloc(DATA_SIZE);
    struct ibv_mr *local_mr = NULL;
    struct ibv_wc wc;
    int ret = 0;

    memset(local_buf, 0, DATA_SIZE);
    local_mr = ibv_reg_mr(id->pd, local_buf, DATA_SIZE, IBV_ACCESS_LOCAL_WRITE);
    if (!local_mr) { perror("reg_mr"); return -1; }

    local_info.addr = 0; // 客户端不需要暴露内存
    local_info.rkey = 0;
    local_info.qp_num = id->qp->qp_num;
    param.private_data = &local_info;
    param.private_data_len = sizeof(local_info);
    param.responder_resources = 1;
    param.initiator_depth = 1;
    param.retry_count = 7;

    if (rdma_connect(id, &param)) { perror("connect"); return -1; }

    struct rdma_cm_event *event;
    if (rdma_get_cm_event(id->channel, &event)) { perror("get_event"); return -1; }
    if (event->event != RDMA_CM_EVENT_ESTABLISHED) { fprintf(stderr, "Unexpected event\n"); return -1; }
    memcpy(&remote_info, event->param.conn.private_data, sizeof(remote_info));
    rdma_ack_cm_event(event);

    // 发起RDMA Read
    if (post_read(id, remote_info.addr, remote_info.rkey, local_buf, DATA_SIZE, local_mr)) {
        perror("post_read"); ret = -1; goto out;
    }

    // 等待完成
    if (rdma_get_send_comp(id, &wc) != 0) { perror("get_send_comp"); ret = -1; goto out; }
    if (wc.status != IBV_WC_SUCCESS) {
        fprintf(stderr, "Read WC error: %s\n", ibv_wc_status_str(wc.status)); ret = -1; goto out;
    }

    // 验证读取结果
    int ok = 1;
    for (int i = 0; i < DATA_SIZE; i++) if (local_buf[i] != 'R') { ok = 0; break; }
    printf("Client: read %d bytes, data %s\n", DATA_SIZE, ok ? "OK" : "MISMATCH");

out:
    rdma_disconnect(id);
    ibv_dereg_mr(local_mr);
    free(local_buf);
    return ret;
}

int main(int argc, char **argv) {
    if (argc != 3) {
        fprintf(stderr, "Usage: %s <server|client> <addr>\n", argv[0]);
        return 1;
    }
    struct rdma_cm_id *listen_id = NULL, *conn_id = NULL;
    struct rdma_event_channel *ec = rdma_create_event_channel();
    if (!ec) { perror("create_event_channel"); return 1; }

    if (strcmp(argv[1], "server") == 0) {
        if (rdma_create_id(ec, &listen_id, NULL, RDMA_PS_TCP)) { perror("create_id"); return 1; }
        struct sockaddr_in addr = {.sin_family = AF_INET, .sin_port = htons(7471)};
        addr.sin_addr.s_addr = htonl(INADDR_ANY);
        if (rdma_bind_addr(listen_id, (struct sockaddr *)&addr)) { perror("bind"); return 1; }
        if (rdma_listen(listen_id, 1)) { perror("listen"); return 1; }
        printf("Server listening on port 7471\n");
        run_server(listen_id);
        rdma_destroy_id(listen_id);
    } else {
        if (rdma_create_id(ec, &conn_id, NULL, RDMA_PS_TCP)) { perror("create_id"); return 1; }
        struct sockaddr_in addr = {.sin_family = AF_INET, .sin_port = htons(7471)};
        inet_pton(AF_INET, argv[2], &addr.sin_addr);
        if (rdma_resolve_addr(conn_id, NULL, (struct sockaddr *)&addr, 2000)) { perror("resolve_addr"); return 1; }
        if (rdma_resolve_route(conn_id, 2000)) { perror("resolve_route"); return 1; }
        run_client(conn_id);
        rdma_destroy_id(conn_id);
    }
    rdma_destroy_event_channel(ec);
    return 0;
}

运行与验证

1. 编译：mpicc -o rdma_read_gpu rdma_read_gpu.c -lrdmacm -libverbs -lcuda（若无mpicc可用gcc替代）
2. 启动服务端：./rdma_read_gpu server 0.0.0.0（监听端口7471）
3. 启动客户端：./rdma_read_gpu client <服务端IP>
4. 通过条件：服务端输出Server: data OK，客户端输出Client: read 4096 bytes, data OK，且两端均无错误信息

关键说明

• 服务端使用cudaMalloc分配显存，ibv_reg_mr注册时需指定IBV_ACCESS_REMOTE_READ权限，否则客户端无法读取
• 客户端通过IBV_WR_RDMA_READ操作码发起单边读，服务端CPU不参与数据拷贝
• 代码使用同步轮询等待完成事件（rdma_get_send_comp/rdma_get_recv_comp），实际工程中应使用异步事件通道
• 若GPU与HCA不在同一NUMA节点，性能会下降，但不影响功能正确性