4. NVLink与PCIe混合通信路径选择

任务说明

在同时具备NVLink和PCIe连接的GPU对之间，小数据量传输时PCIe延迟更低，大数据量时NVLink带宽优势明显。本案例实现一个运行时路径选择器，根据待传输数据大小自动选择延迟最优或带宽最优的通信路径，并验证选择逻辑的正确性。

环境依赖

• CUDA 11.0+（需支持NVLink检测API）
• 双GPU环境（通过NVLink互联，同时挂载在PCIe总线上）
• 编译参数：-lnvidia-ml（NVML库）

核心实现

#include <cuda_runtime.h>
#include <nvml.h>
#include <cstdio>
#include <cstdlib>
#include <vector>
#include <chrono>

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

#define CHECK_NVML(call) do { \
    nvmlReturn_t ret = (call); \
    if (ret != NVML_SUCCESS) { \
        printf("NVML error %s:%d: %s\n", __FILE__, __LINE__, \
               nvmlErrorString(ret)); \
        exit(1); \
    } \
} while(0)

// 路径类型
enum class PathType { PCIE, NVLink };

// 路径选择器：根据数据大小动态选择通信路径
class PathSelector {
public:
    PathSelector(int dev0, int dev1, size_t threshold_bytes)
        : dev0_(dev0), dev1_(dev1), threshold_(threshold_bytes) {}

    // 检测GPU间是否存在NVLink连接
    bool hasNVLink() {
        nvmlDevice_t nvml0, nvml1;
        CHECK_NVML(nvmlDeviceGetHandleByIndex(dev0_, &nvml0));
        CHECK_NVML(nvmlDeviceGetHandleByIndex(dev1_, &nvml1));

        unsigned int linkCount = 0;
        nvmlReturn_t ret = nvmlDeviceGetNvLinkCount(nvml0, &linkCount);
        if (ret != NVML_SUCCESS || linkCount == 0) return false;

        for (unsigned int i = 0; i < linkCount; i++) {
            unsigned int state = 0;
            if (nvmlDeviceGetNvLinkState(nvml0, i, &state) == NVML_SUCCESS
                && state == NVML_GPU_NVLINK_STATE_ACTIVE) {
                // 确认对端是dev1
                nvmlPciInfo_t pci0, pci1;
                nvmlDeviceGetPciInfo(nvml0, &pci0);
                nvmlDeviceGetPciInfo(nvml1, &pci1);
                // 简化：通过NVML查询对端设备
                nvmlDevice_t remote;
                if (nvmlDeviceGetNvLinkRemoteDevice(nvml0, i, &remote) == NVML_SUCCESS) {
                    nvmlPciInfo_t remotePci;
                    nvmlDeviceGetPciInfo(remote, &remotePci);
                    if (remotePci.busIdLegacy[0] != '\0' &&
                        strcmp(remotePci.busIdLegacy, pci1.busIdLegacy) == 0) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

    // 根据数据大小选择路径
    PathType select(size_t bytes) {
        return (bytes < threshold_) ? PathType::PCIE : PathType::NVLink;
    }

    // 执行一次传输并测量带宽（简化版：仅做cudaMemcpy）
    double benchmark(size_t bytes, PathType path) {
        char *h_src, *h_dst;
        char *d_src, *d_dst;
        CHECK_CUDA(cudaMallocHost(&h_src, bytes));
        CHECK_CUDA(cudaMallocHost(&h_dst, bytes));
        CHECK_CUDA(cudaMemset(h_src, 0xAB, bytes));

        CHECK_CUDA(cudaSetDevice(dev0_));
        CHECK_CUDA(cudaMalloc(&d_src, bytes));
        CHECK_CUDA(cudaSetDevice(dev1_));
        CHECK_CUDA(cudaMalloc(&d_dst, bytes));

        // 选择路径：NVLink路径使用cudaMemcpyPeer，PCIe路径使用cudaMemcpy
        auto start = std::chrono::high_resolution_clock::now();
        if (path == PathType::NVLink) {
            CHECK_CUDA(cudaMemcpyPeer(d_dst, dev1_, d_src, dev0_, bytes));
        } else {
            CHECK_CUDA(cudaMemcpy(d_dst, d_src, bytes, cudaMemcpyDefault));
        }
        CHECK_CUDA(cudaDeviceSynchronize());
        auto end = std::chrono::high_resolution_clock::now();

        double secs = std::chrono::duration<double>(end - start).count();
        double bw = (double)bytes / secs / 1e9;  // GB/s

        CHECK_CUDA(cudaFree(d_src));
        CHECK_CUDA(cudaFree(d_dst));
        CHECK_CUDA(cudaFreeHost(h_src));
        CHECK_CUDA(cudaFreeHost(h_dst));
        return bw;
    }

private:
    int dev0_, dev1_;
    size_t threshold_;
};

int main() {
    CHECK_NVML(nvmlInit());

    int devCount = 0;
    CHECK_CUDA(cudaGetDeviceCount(&devCount));
    if (devCount < 2) {
        printf("需要至少2块GPU\n");
        nvmlShutdown();
        return 1;
    }

    // 使用前两块GPU
    int dev0 = 0, dev1 = 1;
    cudaDeviceProp prop0, prop1;
    CHECK_CUDA(cudaGetDeviceProperties(&prop0, dev0));
    CHECK_CUDA(cudaGetDeviceProperties(&prop1, dev1));

    // 阈值：NVLink带宽通常为PCIe的5-10倍，取1MB作为切换点
    const size_t THRESHOLD = 1 * 1024 * 1024;  // 1MB

    PathSelector selector(dev0, dev1, THRESHOLD);
    bool has_nvlink = selector.hasNVLink();
    printf("GPU0: %s, GPU1: %s\n", prop0.name, prop1.name);
    printf("NVLink连接: %s\n", has_nvlink ? "是" : "否");
    printf("切换阈值: %zu bytes (%.2f MB)\n", THRESHOLD, THRESHOLD/1024.0/1024.0);

    if (!has_nvlink) {
        printf("未检测到NVLink，将始终使用PCIe路径\n");
    }

    // 测试不同大小数据的路径选择
    std::vector<size_t> test_sizes = {
        64 * 1024,          // 64KB - 应走PCIe
        512 * 1024,         // 512KB - 应走PCIe
        2 * 1024 * 1024,    // 2MB - 应走NVLink
        16 * 1024 * 1024    // 16MB - 应走NVLink
    };

    printf("\n%-12s %-10s %-12s %-12s\n", "数据大小", "选择路径", "带宽(GB/s)", "验证");
    for (size_t sz : test_sizes) {
        PathType chosen = selector.select(sz);
        const char* path_name = (chosen == PathType::NVLink) ? "NVLink" : "PCIe";

        // 实际执行所选路径的传输
        double bw = selector.benchmark(sz, chosen);

        // 验证：如果NVLink可用且数据大于阈值，应选择NVLink
        bool correct = true;
        if (has_nvlink) {
            correct = (sz < THRESHOLD) ? (chosen == PathType::PCIE)
                                       : (chosen == PathType::NVLink);
        } else {
            correct = (chosen == PathType::PCIE);
        }

        printf("%-12zu %-10s %-12.2f %-12s\n",
               sz, path_name, bw, correct ? "PASS" : "FAIL");
    }

    CHECK_NVML(nvmlShutdown());
    return 0;
}

编译与运行

nvcc -o path_selector path_selector.cu -lnvidia-ml
./path_selector

验证方法

• 程序启动时通过NVML检测GPU0与GPU1之间是否存在活跃的NVLink连接，并打印检测结果。
• 对4组不同大小的测试数据（64KB、512KB、2MB、16MB），程序调用select()选择路径，并实际执行该路径的cudaMemcpyPeer（NVLink）或cudaMemcpy（PCIe）传输，测量带宽。
• 判定标准：当NVLink存在时，小于1MB的数据应选择PCIe路径，大于等于1MB的数据应选择NVLink路径；每行输出PASS表示选择逻辑正确。若NVLink不存在，所有数据应选择PCIe且全部输出PASS。
• 预期输出中带宽数值会随硬件环境变化，但路径选择结果（NVLink/PCIe列）和PASS/FAIL状态应与上述规则严格一致。

关键说明

• 阈值1MB为经验值，实际最优切换点取决于具体GPU的NVLink与PCIe带宽比，可在真实环境用benchmark()对不同大小数据分别测速后确定。
• NVLink路径使用cudaMemcpyPeer，PCIe路径使用cudaMemcpy（默认使用当前设备到目标设备的PCIe路径）。
• 本案例聚焦路径选择逻辑，实际工程中可替换为cudaMemcpyPeerAsync与流操作以重叠通信与计算。