PyTorch中conv2d的C++实现位置查询：对应文件路径

• 顶层C++入口文件：aten/src/ATen/native/Convolution.cpp
  这里是conv2d的核心逻辑分发层，负责输入参数校验，根据设备（CPU/GPU）、卷积类型（普通/分组/深度可分离）、硬件支持的优化选项，选择对应的底层实现路径。

• CPU平台具体实现

  • 原生CPU核函数：aten/src/ATen/native/cpu/Conv2dKernel.cpp，包含基础CPU卷积实现，针对不同输入尺寸、通道数做了分支优化。
  • MKL-DNN优化实现：aten/src/ATen/native/mkldnn/Conv.cpp，若PyTorch编译时启用MKL-DNN，会优先调用这里的优化逻辑，利用Intel计算库加速卷积运算。

• GPU平台具体实现

  • CUDA原生核函数：aten/src/ATen/native/cuda/Conv2d.cu，实现了不依赖CuDNN的CUDA卷积核，适合自定义需求场景。
  • CuDNN加速实现：aten/src/ATen/native/cudnn/Conv.cpp，这是PyTorch GPU卷积的默认路径，会调用CuDNN提供的高度优化卷积算法（如Winograd、FFT、直接卷积等，根据输入自动选最优）。

• 核心计算逻辑流程

  1. 参数校验：确保输入、权重张量维度匹配，检查stride、padding、dilation等参数合法性。
  2. 输入预处理：根据padding参数填充输入张量（或用隐式填充避免内存拷贝），处理dilation带来的权重扩张逻辑。
  3. 算法选择：依据输入尺寸、通道数、硬件性能，挑选最优卷积计算算法。
  4. 核函数调用：将预处理后的输入、权重传入对应CPU/GPU核函数，完成实际卷积运算。
  5. 输出处理：裁剪输出张量到目标尺寸，返回最终结果。

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