手把手教你如何高效地在 MMCV 中贡献算子

// 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}

// 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);

// 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);

// 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}

// 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

// 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}

# 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

# Copyright (c) OpenMMLab. All rights reserved.
2...
3from .tensor_add import tensor_add
4
5__all__ = [
6    ..., 
7    'tensor_add'
8]

import 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)