"PyTorch" 不是一个库,而是至少四层的同心圆。 不同层有不同的职责、不同的依赖、不同的开发者群体。 读源码先认这四层,后面再钻细节就不会迷路。

四层职责对比

native/ 的子目录就是 backend

ATen 真正放算子代码的地方是 aten/src/ATen/native/。 它的子目录名就是该算子针对哪种 backend 写的—— 这种"按目录组织 backend"的模式让你很容易定位 kernel 文件。

PyTorch 有 数千个算子 × 几十种 backend × 多种 dtype, 手写绑定不现实。 所有这些样板都由代码生成器从一个 yaml 文件自动产出—— 理解了这条管线,你就知道源码里 aten/src/ATen/Functions.h 为什么会有上万行, 也知道编译 PyTorch 时为什么会跑很久。

一条真实的 yaml 入口

# aten/src/ATen/native/native_functions.yaml 节选(简化)
- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  device_check: NoCheck
  structured_delegate: add.out
  variants: function, method
  dispatch:
    SparseCPU, SparseCUDA: add_sparse
    MkldnnCPU: mkldnn_add
    ZeroTensor: add_zerotensor
  tags: pointwise

# 同一个 op 的 .out 变体 —— "结构化 kernel" 的真正实现入口
- func: add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)
  device_check: NoCheck
  structured: True             # 声明这是 structured kernel
  structured_inherits: TensorIteratorBase
  ufunc_inner_loop:
    Generic: add (AllAndComplex, BFloat16, Half, ComplexHalf)
    ScalarOnly: add (Bool)
  dispatch:
    SparseCPU: add_out_sparse_cpu
    SparseCUDA: add_out_sparse_cuda
    MkldnnCPU: mkldnn_add_out

每个字段的含义

代码生成管线

Structured Kernels — 现代 op 的标准模式

新加的 op 几乎都写成 structured kernel:把"形状推断 + 输出分配 + 实际计算"分离, 实际计算用 TensorIterator 复用所有 broadcast / type promotion / contiguity 处理。 开发者只需写 ~10 行的核心循环。

// aten/src/ATen/native/BinaryOps.cpp · structured kernel 的"meta"
// 由 yaml 中 structured_delegate / structured: True 触发生成

TORCH_META_FUNC2(add, Tensor)(
    const Tensor& self, const Tensor& other, const Scalar& alpha) {
  build_borrowing_binary_op(maybe_get_output(), self, other);
}

// 实际的 CPU kernel —— 由 ufunc 装饰器生成向量化循环
// aten/src/ATen/native/cpu/BinaryOpsKernel.cpp
void add_kernel(TensorIteratorBase& iter, const Scalar& alpha_scalar) {
  AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBFloat16, iter.dtype(), "add_cpu", [&]() {
    auto alpha = alpha_scalar.to<scalar_t>();
    cpu_kernel_vec(iter,
      [=](scalar_t a, scalar_t b) -> scalar_t { return a + alpha * b; },
      [=](Vectorized<scalar_t> a, Vectorized<scalar_t> b) {
        return vec::fmadd(b, Vectorized<scalar_t>(alpha), a);
      });
  });
}

上面讲的是 PyTorch 内部 op 的"官方通道"(yaml + codegen)。 但你写自己的 op 时不会改 yaml —— 改了要重编译整个 PyTorch,工程量太大。 TORCH_LIBRARY / TORCH_LIBRARY_IMPL 提供了 第二条通道:在外部 C++ 文件里声明 schema、注册 kernel, 运行时挂进 Dispatcher。所有自定义算子都走这条路。

三种典型注册模式

关键的 schema 语法

// 类型签名遵循 IDL —— 不是 C++,也不是 Python
"my_add(Tensor a, Tensor b) -> Tensor"          // 最普通
"my_add(Tensor a, Tensor b, *, Scalar alpha=1) -> Tensor"   // 默认参数 + kwargs-only
"my_add_(Tensor(a!) self, Tensor b) -> Tensor(a!)"   // inplace,用 ! 标 mutable alias
"my_add.out(Tensor a, Tensor b, *, Tensor(a!) out) -> Tensor(a!)"   // .out 变体
"my_split(Tensor a, int chunks) -> Tensor[]"   // 返回 list

// (a!) 是别名注解:
//   (a)  — 表示这个张量是输入 a 的视图(共享 Storage,read-only alias)
//   (a!) — 同上,且会被修改(write alias),用于 inplace / out 变体
//   不写 — 完全独立的新张量

理论看完,亲手写一个完整的自定义算子:CUDA kernel + C++ 调度 + Python 绑定 + Autograd 支持。 不到 100 行代码,用 torch.utils.cpp_extension.load_inline 可以不写 setup.py、直接在 notebook 里热编译, 30 秒拿到一个能跑的算子。

步骤 1 — 写 CUDA kernel

// my_add_kernel.cu (在下一段代码里以字符串形式传入)
#include <torch/extension.h>
#include <cuda_runtime.h>

template <typename scalar_t>
__global__ void my_add_kernel(
    const scalar_t* __restrict__ a,
    const scalar_t* __restrict__ b,
    scalar_t* __restrict__ out,
    int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) out[i] = a[i] + b[i];
}

// C++ 入口:校验形状、launch kernel、返回结果
torch::Tensor my_add_cuda(const torch::Tensor& a, const torch::Tensor& b) {
  TORCH_CHECK(a.is_cuda() && b.is_cuda(), "inputs must be CUDA");
  TORCH_CHECK(a.sizes() == b.sizes(),    "shape mismatch");
  auto a_c = a.contiguous();
  auto b_c = b.contiguous();
  auto out = torch::empty_like(a_c);

  int n = a_c.numel();
  int threads = 256;
  int blocks  = (n + threads - 1) / threads;

  AT_DISPATCH_FLOATING_TYPES(a_c.scalar_type(), "my_add_cuda", [&] {
    my_add_kernel<scalar_t><<<blocks, threads>>>(
        a_c.data_ptr<scalar_t>(),
        b_c.data_ptr<scalar_t>(),
        out.data_ptr<scalar_t>(),
        n);
  });
  return out;
}

// 注册到 Dispatcher —— 两个宏
TORCH_LIBRARY(myops, m) {
  m.def("my_add(Tensor a, Tensor b) -> Tensor");
}
TORCH_LIBRARY_IMPL(myops, CUDA, m) {
  m.impl("my_add", my_add_cuda);
}

步骤 2 — 用 load_inline 热编译

import torch
from torch.utils.cpp_extension import load_inline

cuda_src = open("my_add_kernel.cu").read()
cpp_src  = ""   # 没有额外 .cpp,全部代码在 .cu 文件里

myops = load_inline(
    name="myops",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    is_python_module=False,        # 走 TORCH_LIBRARY,不走 pybind
    verbose=True,                  # 看 nvcc 输出
)

# 现在算子已经挂在 torch.ops.myops 下
a = torch.randn(1024, device="cuda")
b = torch.randn(1024, device="cuda")
c = torch.ops.myops.my_add(a, b)

print(torch.allclose(c, a + b))   # True ✓

步骤 3 — 加上 Autograd 支持

上面注册的是无梯度版本。如果想让 my_add 能被 backward() 反向传播, 最简单的方法是在 Python 端用 torch.autograd.Function 包一层—— 不必写 C++ Autograd 类。

class MyAddFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)        # 反向需要的输入
        return torch.ops.myops.my_add(a, b)

    @staticmethod
    def backward(ctx, grad_output):
        # 加法的局部梯度 = 1,直接把上游梯度传下去
        return grad_output, grad_output

# 包装成 Python 函数
def my_add(a, b):
    return MyAddFunction.apply(a, b)

# --- 测试 backward 是否正常 ---
a = torch.randn(3, device="cuda", requires_grad=True)
b = torch.randn(3, device="cuda", requires_grad=True)
y = my_add(a, b).sum()
y.backward()

print(a.grad, b.grad)
# tensor([1., 1., 1.], device='cuda:0')
# tensor([1., 1., 1.], device='cuda:0')
# ✓ 与 PyTorch 内置 add 完全一致

三种构建方式对比

1. 画一张 PyTorch 分层图(c10 / ATen / torch/csrc / torch),标出每层目录路径和不能跨层引用的依赖方向。
2. 用 grep 在 PyTorch 源码里找到 at::relu 的 CUDA 实现文件路径——演练"从 yaml 到 kernel"的检索流程。
3. 解释为什么 structured kernel 比传统写法节省代码,并举例说明 meta 函数与 kernel 函数的分工。
4. 用 cpp_extension.load_inline 实现一个 my_relu(forward+backward),与 PyTorch 内置 ReLU 对比梯度。
5. 设想给 PyTorch 加一个名为 NPU 的新硬件,你需要分别在 c10 / ATen / torch/csrc 哪些位置做改动?哪些可以完全不动?