一般来说，PyTorch 算子开发时，都是用 Python 当上层入口，只有在性能关键路径上，才用 C++/CUDA 写自定义算子，然后注册进 PyTorch dispatcher．这样，可以享受到 PyTorch 提供的其他服务，如自动微分、torch.compile 等等．

工程规范

通常，一个 C++/CUDA 自定义算子可以这样进行组织：

1
2
3
4
5
6
7
8
9
myext/
├── setup.py
├── myext
│   └── __init__.py
├── csrc
│   ├── ops.cpp
│   ├── ops_cuda.cu
│   └── ops_cpu.cpp
└── test.py

接下来，我们简单看看如何编写算子，并进行注册．

算子注册

假设我们的 ops.cpp 已经编写好了如下代码，主要是根据 Tensor 的后端，判断使用 CUDA 还是 CPU 进行计算．

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#include <torch/extension.h>
#include <vector>

torch::Tensor myadd_cuda(torch::Tensor a, torch::Tensor b);
torch::Tensor myadd_cpu(torch::Tensor a, torch::Tensor b);

torch::Tensor myadd(torch::Tensor a, torch::Tensor b) {
    TORCH_CHECK(a.sizes() == b.sizes(), "a and b must have same shape");
    TORCH_CHECK(a.scalar_type() == b.scalar_type(), "dtype mismatch");
    TORCH_CHECK(a.device() == b.device(), "device mismatch");

    if (a.is_cuda()) {
        return myadd_cuda(a, b);
    }
    return myadd_cpu(a, b);
}

接下来，我们需要在 ops.cpp 文件里，将其注册为算子．

1
2
3
4
5
6
7
8
9
10
11
12
13
14
// 定义 schema
TORCH_LIBRARY(myops, m) {
    m.def("myadd(Tensor a, Tensor b) -> Tensor");
}

// 注册 CPU 实现
TORCH_LIBRARY_IMPL(myops, CPU, m) {
    m.impl("myadd", &myadd_cpu);
}

// 注册 CUDA 实现
TORCH_LIBRARY_IMPL(myops, CUDA, m) {
    m.impl("myadd", &myadd_cuda);
}

setup.py 编译并构建

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name="myops",
    ext_modules=[
        CUDAExtension(
            name="myops_ext",
            sources=[
                "csrc/ops.cpp",
                "csrc/ops_cpu.cpp",
                "csrc/ops_cuda.cu",
            ],
            extra_compile_args={
                "cxx": ["-O3", "-std=c++17"],
                "nvcc": ["-O3", "--use_fast_math"],
            },
        )
    ],
    cmdclass={"build_ext": BuildExtension},
)

torch::Tensor

在 C++ 侧，最核心的数据结构就是 torch::Tensor．它基本对应 Python 里的 torch.Tensor，背后走的是 ATen / PyTorch dispatcher 这一整套张量与算子系统．官方 C++ API 文档和自定义算子教程都基于这套接口．

License

本文采用 署名-非商业性使用-相同方式共享 4.0 国际 许可协议，转载请注明出处。