自定义GPU Op访问输入值时触发Segmentation Fault问题求助

我之前也碰到过类似的GPU Op段错误问题，尤其是和CPU版本行为不一致的时候，大概率是GPU内存访问或者数据拷贝的逻辑出了问题。结合你描述的情况，给你几个针对性的排查方向：

CPU Op可以直接通过input_tensor.flat<int32>()直接访问内存，但GPU端绝对不能这么干！GPU核函数只能访问设备内存，如果你的代码里直接用了主机端的指针去访问，必然触发段错误——因为GPU无法直接读取主机内存（除非启用统一内存，但TensorFlow默认不这么配置）。

正确的做法是在GPU Op的Compute函数里，通过TensorFlow的GPUDevice接口获取设备指针：

// 获取输入张量的设备端指针
const int32* input_device_ptr = context->eigen_device<GPUDevice>().template flat<int32>(input_tensor).data();
// 获取输出张量的设备端指针
float* output_device_ptr = context->eigen_device<GPUDevice>().template flat<float>(*output_tensor).data();

然后把这些设备指针传入CUDA核函数，在核函数内进行内存访问。

段错误也可能来自核函数的block/grid大小设置错误，导致线程索引超出张量的有效范围。比如你可以这样计算合理的启动参数：

int tensor_size = input_tensor.flat<int32>().size();
dim3 block_size(256); // 常用的256线程/block
dim3 grid_size((tensor_size + block_size.x - 1) / block_size.x);
// 启动核函数时一定要检查索引边界
InterfaceKernel<<<grid_size, block_size>>>(input_device_ptr, output_device_ptr, tensor_size);

同时在核函数开头加上边界判断：

__global__ void InterfaceKernel(const int32* input, float* output, int size) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= size) return; // 避免越界访问
  // 你的业务逻辑
}

GDB对GPU核函数的调试确实不太友好，不如直接在代码里加入CUDA错误检查，能快速定位问题：

// 核函数启动后立即检查错误
cudaError_t err = cudaGetLastError();
OP_REQUIRES(context, err == cudaSuccess,
            errors::Internal("InterfaceKernel launch failed: ", cudaGetErrorString(err)));

这能帮你区分是核函数启动参数错误，还是内存访问本身的问题。

虽然ShapeFn在主机端执行，但如果输出形状推断错误，会导致GPU端内存分配异常，进而触发段错误。比如你是不是在ShapeFn里没有正确匹配输入输出的维度？可以先把ShapeFn简化成直接沿用输入形状，测试是否还会报错。

把你的GPU Op简化到极致：比如输入一个固定大小的int32张量，输出一个相同大小的float32张量，只做简单的类型转换拷贝。如果这个最小用例能正常运行，再逐步添加你的业务逻辑，这样能快速定位是哪部分代码导致的问题。

给你一个简单的正确GPU Op框架参考，你可以对比自己的代码：

// CUDA核函数
__global__ void InterfaceKernel(const int32* input_ptr, float* output_ptr, int size) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < size) {
    output_ptr[idx] = static_cast<float>(input_ptr[idx]);
  }
}

// GPU Op类
class InterfaceGpuOp : public OpKernel {
 public:
  explicit InterfaceGpuOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // 获取输入张量
    const Tensor& input_tensor = context->input(0);
    int tensor_size = input_tensor.flat<int32>().size();

    // 分配输出张量
    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(), &output_tensor));

    // 获取设备指针
    const int32* input_device_ptr = context->eigen_device<GPUDevice>().template flat<int32>(input_tensor).data();
    float* output_device_ptr = context->eigen_device<GPUDevice>().template flat<float>(*output_tensor).data();

    // 启动核函数
    dim3 block_size(256);
    dim3 grid_size((tensor_size + block_size.x - 1) / block_size.x);
    InterfaceKernel<<<grid_size, block_size>>>(input_device_ptr, output_device_ptr, tensor_size);

    // 检查CUDA错误
    OP_REQUIRES(context, cudaGetLastError() == cudaSuccess,
                errors::Internal("InterfaceKernel launch failed: ", cudaGetErrorString(cudaGetLastError())));
  }
};

// 注册GPU Op
REGISTER_KERNEL_BUILDER(Name("Interface").Device(DEVICE_GPU), InterfaceGpuOp);

核心要记住：GPU编程的内存模型和CPU完全不同，核函数只能访问设备内存，所有数据交互都要遵循主机-设备的拷贝规则，千万不能混用指针！

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