PyTorch自定义C++/CUDA函数运行触发Segmentation fault如何解决

故障核心原因 #

• 过时的CUDA张量校验API调用：你在simple_cuda.cpp里定义的CHECK_CUDA宏使用了x.type().is_cuda()，该API在PyTorch 1.5及之后版本已经被移除，直接调用会触发非法内存访问，是本次段错误的核心诱因。
• CUDA核函数启动配置错误：你当前核函数启动参数为<<<1,1>>>，仅启动1个CUDA线程，只能处理长度为1的张量，虽然不会直接触发段错误，但完全无法满足10万长度的张量计算需求。
• 过时的张量数据指针获取API：你在核函数传参时使用了x.data<scalar_t>()，该旧API在部分PyTorch版本中会出现类型访问异常，存在内存访问风险。
• 缺少输入张量维度校验：你的代码没有校验两个输入张量x和y的长度是否一致，如果y的长度小于x，核函数访问y的时候会出现越界，也可能触发段错误。
• CUDA错误检查逻辑不完整：核函数是异步执行的，你直接调用cudaGetLastError只能拿到启动阶段的错误，无法捕获核函数执行过程中的错误，会导致问题排查遗漏。
• 额外兼容性检查：确认你当前安装的PyTorch版本编译使用的CUDA版本和本地的CUDA 11.1.1一致，版本不匹配也可能导致接口调用异常触发段错误，可以用print(torch.version.cuda)命令查看PyTorch对应的CUDA版本。

修复代码示例 #

修改simple_cuda.cpp #

#include <torch/extension.h>

torch::Tensor cuda_vec_sum(torch::Tensor x, torch::Tensor y);

// 修改过时的CHECK_CUDA宏
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

torch::Tensor vec_sum(torch::Tensor x, torch::Tensor y) {
        CHECK_INPUT(x);
        CHECK_INPUT(y);
        // 新增输入维度校验
        TORCH_CHECK(x.size(0) == y.size(0), "x and y must have the same length");
        return cuda_vec_sum(x, y);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("sum", &vec_sum, "VECSUM (CUDA)");
}

修改simple_cuda_kernel.cu #

#include <cuda.h>
#include <cuda_runtime.h>

template <typename scalar_t>
__global__ void cuda_vec_sum_kernel(
    const scalar_t* __restrict__ x,
    const scalar_t* __restrict__ y,
    scalar_t* __restrict__ z,
    size_t N) {

    const int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index < N) {
        z[index] = x[index] + y[index];
    }
}

torch::Tensor cuda_vec_sum(torch::Tensor x, torch::Tensor y) {
    auto z = torch::zeros_like(x);
    // 动态计算核函数启动配置
    const int threads_per_block = 256;
    const int blocks = (x.size(0) + threads_per_block - 1) / threads_per_block;

    // 修改过时的type()调用为scalar_type()，修改data()为data_ptr()
    AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "cuda_vec_sum", ([&] {
        cuda_vec_sum_kernel<scalar_t><<<blocks, threads_per_block>>>(
            x.data_ptr<scalar_t>(),
            y.data_ptr<scalar_t>(),
            z.data_ptr<scalar_t>(),
            x.size(0));
    }));
    // 同步设备后再检查错误
    cudaDeviceSynchronize();
    printf("CUDA Error: %s\n", cudaGetErrorString(cudaGetLastError()));
    return z;
}

修复完成后重新执行python simple_setup.py install编译扩展，再运行测试脚本即可。

内容的提问来源于stack exchange，提问作者Szocinsk