Gaussian Splatting项目中多GPU环境下的CUDA内存访问问题解析

问题背景

在使用Gaussian Splatting项目中的Simple-KNN模块时，开发者发现当输入张量位于第二个GPU(cuda:1)上时，distCUDA2函数会抛出CUDA非法内存访问错误，而同样的代码在第一个GPU(cuda:0)上却能正常运行。这种现象在多GPU环境中并不罕见，但需要深入理解CUDA编程模型才能有效解决。

问题分析

CUDA设备上下文管理

在CUDA编程中，每个设备(cuda:0, cuda:1等)都有自己的内存空间和执行上下文。当我们在PyTorch中创建张量并指定设备时，PyTorch会自动处理设备间的数据传输和同步。然而，当直接调用CUDA原生API或使用thrust库时，我们需要显式管理设备上下文。

根本原因

distCUDA2函数内部使用了SimpleKNN::knn函数，后者可能直接调用了CUDA内核或thrust算法。问题发生时，虽然输入张量位于cuda:1设备上，但CUDA运行时可能仍然默认使用cuda:0设备执行计算，导致跨设备内存访问违规。

解决方案

设备同步机制

正确的做法是在调用任何CUDA操作前，确保CUDA运行时设备与张量所在设备一致。修改后的代码如下：

#include "spatial.h"
#include "simple_knn.h"
#include <cuda_runtime.h>

torch::Tensor distCUDA2(const torch::Tensor& points) {
    const int P = points.size(0);
    
    // 获取输入张量所在的设备索引
    auto device = points.device();
    int device_index = device.index();
    
    // 显式设置CUDA设备
    cudaSetDevice(device_index);
    
    auto float_opts = points.options().dtype(torch::kFloat32);
    torch::Tensor means = torch::full({P}, 0.0, float_opts);
    
    SimpleKNN::knn(P, (float3*)points.contiguous().data<float>(), 
                  means.contiguous().data<float>());
    
    return means;
}

关键改进点

1. 设备查询：通过points.device()获取输入张量所在的设备信息
2. 设备同步：使用cudaSetDevice()显式设置当前线程的CUDA设备
3. 内存连续性：保持使用contiguous()确保内存布局符合CUDA内核要求

深入理解

CUDA设备上下文

CUDA设备上下文是线程本地的，每个主机线程维护一个当前设备栈。当调用cudaSetDevice()时，会改变当前线程的设备上下文，影响后续所有CUDA操作。

PyTorch与CUDA交互

PyTorch封装了大部分设备管理逻辑，但在混合使用PyTorch和原生CUDA代码时，开发者需要特别注意：

1. PyTorch张量的设备属性
2. CUDA运行时的当前设备设置
3. 内存指针的有效性范围

最佳实践

在多GPU环境中开发CUDA扩展时，建议：

1. 始终检查输入张量的设备属性
2. 在执行CUDA操作前显式设置设备
3. 避免直接在不同设备间传递原始指针
4. 使用PyTorch提供的跨设备数据传输方法
5. 在文档中明确标注函数的设备要求

总结

Gaussian Splatting项目中遇到的这个问题很好地展示了混合使用PyTorch和原生CUDA代码时的潜在陷阱。通过显式管理CUDA设备上下文，我们不仅解决了当前的问题，也为项目在多GPU环境中的稳定运行奠定了基础。理解这些底层机制对于开发高性能的计算机视觉和图形学应用至关重要。