为何PyTorch C++扩展比等效Numba版本慢很多？

我一直在尝试各种方案加速PyTorch中的循环密集型逻辑，常用方案包括Numba或自定义PyTorch C++扩展。

以数字信号处理中的可变长度延迟线为例，该逻辑可用简单但低效的Python原生for循环实现：

def delay_line(samples, delays):
    """
    :param samples: Float tensor of shape (N,)
    :param delays: Int tensor of shape (N,)
    
    核心逻辑：将每个`samples[i]`与`delays[i]`指定的延迟样本混合
    """
    for i in range(len(samples)):
        delay = int(delays[i].item())
        index_delayed = i - delay
        if index_delayed < 0:
            index_delayed = 0

        samples[i] = 0.5 * (samples[i] + samples[index_delayed])

我原本期望通过C实现获得显著性能提升，于是按照PyTorch官方教程将Python代码直接转译为C：

void delay_line(torch::Tensor samples, torch::Tensor delays) {

  int64_t input_size = samples.size(-1);

  for (int64_t i = 0; i < input_size; ++i) {
    int64_t delay = delays[i].item<int64_t>();
    int64_t index_delayed = i - delay;
    if (index_delayed < 0) {
      index_delayed = 0;
    }

    samples[i] = 0.5 * (samples[i] + samples[index_delayed]);
  }
}

同时我还为Python函数添加不同JIT装饰器，得到Numba和TorchScript版本。针对CPU和GPU设备的各版本执行基准测试后，结果令人意外：

说明：样本缓冲区大小固定为1024；结果为100次执行的中位数，排除初始JIT开销；输入数据创建及设备迁移时间未计入测量

最值得注意的是：C版本性能意外低下，接近低效的Python for循环，而Numba版本快两个数量级，说明该问题存在更优解。我已确保编译时启用-O2/-O3优化，空实现的C函数耗时仅2-3μs，说明时间消耗在函数主体中。

请问：

1. 是什么原因导致该PyTorch C++扩展性能差？
2. 在C++端有什么方法可以达到Numba版本的性能？
3. 附加问题：GPU版本远慢于CPU版本是否符合预期？

一、C++版本性能低下的核心原因 #

你的C++代码直接复用了Python的逐元素Tensor访问方式，而**torch::Tensor的[]运算符和.item()方法会带来极高的开销**：

• 每次调用.item()或[]都会触发Tensor的边界检查、类型校验，以及可能的设备同步（尤其是GPU张量）。
• 逐元素访问Tensor本质是通过框架层的间接内存访问，无法利用CPU的缓存优化，也无法触发SIMD向量指令，开销层级和Python循环接近。

而Numba会将函数编译为直接操作原始内存的机器码：它会直接获取Tensor的底层数据指针，跳过所有Tensor框架的额外校验，同时自动启用SIMD优化，所以性能远超你的C++实现。

二、C++端的性能优化方案 #

要达到Numba的性能，需要绕过Tensor的逐元素API，直接操作原始内存：

1. 直接操作Tensor底层数据指针 #

使用torch::Tensor::data_ptr()获取内存地址，避免逐元素API的开销：

void delay_line(torch::Tensor samples, torch::Tensor delays) {
  // 确保张量连续，避免非连续内存的访问开销
  samples = samples.contiguous();
  delays = delays.contiguous();

  float* samples_ptr = samples.data_ptr<float>();
  const int64_t* delays_ptr = delays.data_ptr<int64_t>();
  int64_t input_size = samples.size(-1);

  for (int64_t i = 0; i < input_size; ++i) {
    int64_t delay = delays_ptr[i];
    int64_t index_delayed = i - delay;
    if (index_delayed < 0) {
      index_delayed = 0;
    }

    samples_ptr[i] = 0.5f * (samples_ptr[i] + samples_ptr[index_delayed]);
  }
}

2. 启用编译器优化与SIMD #

编译扩展时添加以下选项，最大化CPU性能：

• -O3：最高级别代码优化
• -march=native：针对当前CPU架构生成最优指令
• -ffast-math：允许数学运算近似优化（适合信号处理场景）

3. GPU版本的优化思路 #

GPU版本不能用CPU风格的串行循环，必须编写CUDA核函数：

• 利用CUDA线程并行性，每个线程处理一个元素
• 避免全局内存随机访问，可使用共享内存缓存常用数据
• 用torch::cuda::kernel封装CUDA核，或直接调用PyTorch CUDA API

示例CUDA核框架：

__global__ void delay_line_cuda_kernel(float* samples, const int64_t* delays, int64_t input_size) {
  int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= input_size) return;

  int64_t delay = delays[i];
  int64_t index_delayed = i - delay;
  if (index_delayed < 0) index_delayed = 0;

  samples[i] = 0.5f * (samples[i] + samples[index_delayed]);
}

void delay_line(torch::Tensor samples, torch::Tensor delays) {
  samples = samples.contiguous().cuda();
  delays = delays.contiguous().cuda();

  int64_t input_size = samples.size(-1);
  dim3 block_size(256);
  dim3 grid_size((input_size + block_size.x - 1) / block_size.x);

  delay_line_cuda_kernel<<<grid_size, block_size>>>(
    samples.data_ptr<float>(),
    delays.data_ptr<int64_t>(),
    input_size
  );
  cudaDeviceSynchronize();
}

三、附加问题：GPU版本慢于CPU是否符合预期？ #

你的测试中GPU版本更慢是符合预期的，原因如下：

• 任务是逐元素依赖的串行逻辑：每个元素的计算依赖前序元素的修改结果，GPU的并行优势无法发挥。
• GPU存在启动开销：对于仅1024个元素的小任务，GPU的线程调度、数据传输开销远大于计算收益，反而不如CPU高效。
• 只有当任务规模足够大（如百万级元素）且逻辑可并行化时，GPU才能体现性能优势。

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