0.前言
不知道大家在使用 MMCV 的过程中有没有遇到这种情况：MMCV 没有提供自己需要的 CPU/CUDA 算子，于是希望提一个 PR（Pull Request），将这个算子加入 MMCV，但是又不知从何处下手。本文以最简单的 TensorAdd 算子为例，向大家展示为 MMCV 贡献算子的全过程，希望能够帮助大家更好地理解 MMCV 算子的
目录结构，以便更高效地贡献算子。
注意：如果您不太了解提 PR 的流程，可以阅读 https://mmcv.readthedocs.io/zh_CN/latest/community/pr.html。
1.更新文档
在 docs/en/understand_mmcv/ops.md 以及 docs/zh_cn/understand_mmcv/ops.md 里添加 TensorAdd，该文档用于维护 MMCV 已支持的算子。
2.算子实现
在这篇文章里我们提供了 CPU 算子和 CUDA 算子的实现方法，您可以选择只提供 CPU 算子或者只提供 CUDA 算子或者两种算子都提供。
2.1 提供算子 C 接口
在 mmcv/ops/csrc/pytorch/ 目录添加 tensor_add.cpp 文件。
代码语言：javascript复制 1// Copyright(c) OpenMMLab.All rights reserved.
2#include "pytorch_cpp_helper.hpp"
3#include "pytorch_device_registry.hpp"
4
5void tensor_add_impl(const Tensor input1, const Tensor input2, Tensor output) {
6 DISPATCH_DEVICE_IMPL(tensor_add_impl, input1, input2, output);
7}
8
9void tensor_add(const Tensor input1, const Tensor input2, Tensor output) {
10 tensor_add_impl(input1, input2, output);
11}tensor_add 是算子在 C 层的接口，而 tensor_add_impl 中的 DISPATCH_DEVICE_IMPL 宏会根据 Tensor 参数的设备类型自动选择 CPU 或 CUDA 的算子实现。DISPATCH_DEVICE_IMPL 宏依赖于REGISTER_DEVICE_IMPL 宏，我们会在下面的 CPU 算子实现和 CUDA 算子实现提到。
DISPATCH_DEVICE_IMPL 和 REGISTER_DEVICE_IMPL 都用于算子分发，更多细节见 兼容性文档，PR-1463 以及 Dispatcher 机制解析。
兼容性文档：
https://mmcv.readthedocs.io/zh_CN/latest/compatibility.html#v1-3-18
PR-1463：
https://github.com/open-mmlab/mmcv/pull/1463
Dispatcher 机制解析：
https://zhuanlan.zhihu.com/p/451671838
2.2 CPU 算子实现
在 mmcv/ops/csrc/pytorch/cpu/ 目录添加 tensor_add_cpu.cpp 文件。
代码语言：javascript复制 1// Copyright(c) OpenMMLab.All rights reserved.
2#include "pytorch_cpp_helper.hpp"
3#include "pytorch_device_registry.hpp"
4
5template <typename T>
6void tensor_add_cpu_kernel(int N, const T* input1, const T* input2, T* output) {
7 for (int i = 0; i < N; i ) {
8 output[i] = input1[i] input2[i];
9 }
10}
11
12void TensorAddCPUKernelLaucher(const Tensor input1, const Tensor input2,
13 Tensor output) {
14 int N = input1.size(0);
15 AT_DISPATCH_FLOATING_TYPES_AND_HALF(
16 input1.scalar_type(), "tensor_add_cpu_kernel", [&] {
17 tensor_add_cpu_kernel<scalar_t>(N, input1.data_ptr<scalar_t>(),
18 input2.data_ptr<scalar_t>(),
19 output.data_ptr<scalar_t>());
20 });
21}
22
23void tensor_add_cpu(const Tensor input1, const Tensor input2, Tensor output) {
24 TensorAddCPUKernelLaucher(input1, input2, output);
25}
26void tensor_add_impl(const Tensor input1, const Tensor input2, Tensor output);
27
28REGISTER_DEVICE_IMPL(tensor_add_impl, CPU, tensor_add_cpu);我们看到最后一行 REGISTER_DEVICE_IMPL 宏将 tensor_add_impl 和 tensor_add_cpu 绑定在一起。
tensor_add_cpu 调用 TensorAddCPUKernelLaucher，TensorAddCPUKernelLaucher 启动 tensor_add_cpu_kernel 完成算子的计算。
2.3 CUDA 算子实现
CUDA 算子的调用过程和 CPU 算子类似，但是在代码安排上略有不同，比如很多 CUDA Kernel 被放在 mmcv/ops/csrc/common/cuda/ 目录，因为该目录负责管理后端无关可共享的代码（cuda kernel，mlu kernel等）。
2.3.1 算子绑定
在 mmcv/ops/csrc/pytorch/cuda/cudabind.cpp 里添加 TensorAdd 的 CUDA 算子绑定。
代码语言：javascript复制 1// Copyright (c) OpenMMLab. All rights reserved.
2...
3void TensorAddCUDAKernelLauncher(const Tensor input1, const Tensor input2,
4 const Tensor output);
5
6void tensor_add_cuda(const Tensor input1, const Tensor input2, Tensor output) {
7 TensorAddCUDAKernelLauncher(input1, input2, output);
8}
9
10void tensor_add_impl(const Tensor input1, const Tensor input2, Tensor output);
11
12REGISTER_DEVICE_IMPL(tensor_add_impl, CUDA, tensor_add_cuda);我们看到最后一行 REGISTER_DEVICE_IMPL 宏将 tensor_add_impl 和 tensor_add_cuda 绑定在一起。
tensor_add_cuda 调用 TensorAddCUDAKernelLaucher。
2.3.2 KernelLaucher
在 mmcv/ops/csrc/pytorch/cuda/ 目录添加 tensor_add_cuda.cu 文件。
代码语言：javascript复制 1// Copyright (c) OpenMMLab. All rights reserved.
2#include <torch/types.h>
3
4#include "pytorch_cuda_helper.hpp"
5#include "tensor_add_cuda_kernel.cuh"
6
7void TensorAddCUDAKernelLauncher(const Tensor input1, const Tensor input2,
8 const Tensor output) {
9 int N = input1.size(0);
10 at::cuda::CUDAGuard device_guard(input1.device());
11 cudaStream_t stream = at::cuda::getCurrentCUDAStream();
12
13 AT_DISPATCH_FLOATING_TYPES_AND_HALF(
14 input1.scalar_type(), "tensor_add_cuda_kernel", ([&] {
15 dim3 blocks(GET_BLOCKS(N));
16 dim3 threads(THREADS_PER_BLOCK);
17 tensor_add_cuda_kernel<<<blocks, threads, 0, stream>>>(
18 N, input1.data_ptr<scalar_t>(), input2.data_ptr<scalar_t>(),
19 output.data_ptr<scalar_t>());
20 }));
21
22 AT_CUDA_CHECK(cudaGetLastError());
23}TensorAddCUDAKernelLauncher 会启动 tensor_add_cuda_kernel 完成算子的具体操作。其中使用 AT_DISPATCH_FLOATING_TYPES_AND_HALF 宏启动 CUDA Kernel ，该宏内部包装了一个 switch 语句来完成针对张量类型的分派，更多这类宏可见：https://github.com/pytorch/pytorch/blob/HEAD/aten/src/ATen/Dispatch.h#L262。
2.3.3 Kernel
在 mmcv/ops/csrc/common/cuda/ 下添加 tensor_add_cuda_kernel.cuh。
代码语言：javascript复制 1// Copyright (c) OpenMMLab. All rights reserved.
2#ifndef TENSOR_ADD_CUDA_KERNEL_CUH
3#define TENSOR_ADD_CUDA_KERNEL_CUH
4
5#ifdef MMCV_USE_PARROTS
6#include "parrots_cuda_helper.hpp"
7#else
8#include "pytorch_cuda_helper.hpp"
9#endif
10
11template <typename T>
12__global__ void tensor_add_cuda_kernel(const int N, const T* input1,
13 const T* input2, T* output) {
14 CUDA_1D_KERNEL_LOOP(i, N) { output[i] = input1[i] input2[i]; }
15}
16#endif // TENSOR_ADD_CUDA_KERNEL_CUH在这里实现了算子的具体操作。其中 CUDA_1D_KERNEL_LOOP 是 MMCV 提供的一个简写 Kernel Loop 的宏，更多这类宏可见：https://github.com/open-mmlab/mmcv/blob/master/mmcv/ops/csrc/common/cuda/common_cuda_helper.hpp#L6。
3. 提供 Python 接口
在完成 C /CUDA 的算子后，我们需要在 mmcv/ops/csrc/pytorch/pybind.cpp 里实现 C 接口和 Python 接口的绑定，从而提供一个 Python 可以调用的接口。
代码语言：javascript复制1// Copyright (c) OpenMMLab. All rights reserved.
2...
3void tensor_add(const Tensor input1, const Tensor input2, Tensor output);
4
5PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
6 ...
7 m.def("tensor_add", &tensor_add, "tensor_add", py::arg("input1"),
8 py::arg("input2"), py::arg("output"));
9}4. Python 调用算子
在 mmcv/ops/ 下添加 tensor_add.py。
代码语言：javascript复制 1# Copyright (c) OpenMMLab. All rights reserved.
2import torch
3from torch.autograd import Function
4from torch.autograd.function import once_differentiable
5
6from ..utils import ext_loader
7
8ext_module = ext_loader.load_ext('_ext', ['tensor_add'])
9
10
11class TensorAdd(Function):
12
13 @staticmethod
14 def forward(ctx, input1: torch.Tensor,
15 input2: torch.Tensor) -> torch.Tensor:
16 """Add two tensor.
17
18 Args:
19 input1 (torch.Tensor): shape (N).
20 input2 (torch.Tensor): shape (N).
21
22 Returns:
23 torch.Tensor: shape (N), tensor of input1 input2.
24 """
25 assert input1.is_contiguous() and input2.is_contiguous()
26 assert input1.dim() == 1 and input2.dim() == 1
27 assert input1.size(0) == input2.size(0)
28 output = torch.zeros(
29 input1.size(0), dtype=input1.dtype, device=input1.device)
30 ext_module.tensor_add(input1, input2, output)
31 return output
32
33 @staticmethod
34 @once_differentiable
35 def backward(ctx, grad_output):
36 return grad_output, grad_output
37
38
39tensor_add = TensorAdd.apply在 mmcv/ops/__init__.py 文件添加对外的接口。
1# Copyright (c) OpenMMLab. All rights reserved.
2...
3from .tensor_add import tensor_add
4
5__all__ = [
6 ...,
7 'tensor_add'
8]5. 编译 MMCV
参考 从源码编译 MMCV 的步骤重新编译 MMCV，更多关于 C /CUDA 算子实现和调用全流程的内容可见：PyTorch 源码解读之 cpp_extension。
从源码编译 MMCV：
https://mmcv.readthedocs.io/zh_CN/latest/get_started/build.html
PyTorch 源码解读之 cpp_extension：
https://zhuanlan.zhihu.com/p/348555597
6. 添加单元测试
在 tests/test_ops/ 下添加 test_tensor_add.py。
代码语言：javascript复制 1import numpy as np
2import pytest
3import torch
4
5from mmcv.ops import tensor_add
6
7
8@pytest.mark.parametrize('device', [
9 'cpu',
10 pytest.param(
11 'cuda',
12 marks=pytest.mark.skipif(
13 not torch.cuda.is_available(), reason='requires CUDA support'))
14])
15@pytest.mark.parametrize('dtype', [torch.float, torch.half])
16def test_tensor_add(device, dtype):
17 n = 1024 * 1024
18 input1 = torch.rand(n).type(dtype).to(device).requires_grad_()
19 input2 = torch.rand(n).type(dtype).to(device).requires_grad_()
20 expected_output = (input1 input2).cpu()
21 output = tensor_add(input1, input2)
22 output.backward(torch.ones_like(output))
23 assert np.allclose(
24 output.detach().cpu(), expected_output.detach(), atol=1e-4)
25 assert np.allclose(input1.grad.cpu().detach(), 1, atol=1e-4)在终端通过 pytest tests/test_ops/test_tensor_add.py 测试 tensor_add 是否正确。
7. 需要注意的点
1.文件需要添加 copyright。
2.DISPATCH_DEVICE_IMPL 接口参数至少要有一个是 Tensor，否则无法根据设备类型做算子分发。
3.目录安排、代码风格和命名规范等需符合 MMCV 算子的规范，例如 tensor_add_cpu_kernel 不能写成 tensor_add_kernel、tensor_add_cpukernel 等其他形式。
4.如果算子需要反向传播，请编写反向传播算子并根据 自定义 C 和 CUDA 扩展 组织 Python 代码。
8. 函数调用流程图
横实线区别 Python 层和 C 层，两者之间通过 pybind11 绑定。
竖实线区别 C 算子和 CUDA 算子，算子的分发由 DISPATCH_DEVICE_IMPL 决定。
9. 总结
希望本篇文章让您更为深入地了解了如何在 MMCV 中添加自定义算子，如果对文档有任何疑问或者修改建议，欢迎提交 Issue 或 PR
