CUDA——编写并调用自定义CUDA Pytorch算子

0. CUDA算子编译方式

编译方式主要分为三种：

1. 基于pytorch的jit（每次运行代码时编译，多次执行即多次编译）；
2. 基于python自身的setup（在安装算子包时手动编译一次，后续可重复使用）；
3. 基于C++ 的cmake（一次编译，可重复使用；但存在较多CmakeList中配置问题需要解决）；

---

在这里，我推荐选择第二种基于python自身的setup的编译方式，该方式具有编译过程简单、快速且可以复用的特点。

因此，本文将主要setup的编译方式进行讲解。

• 参考文章：详解PyTorch编译并调用自定义CUDA算子的三种方式

1. 编写CUDA核函数和主函数

• 首先我们需要编写最基本的CUDA核函数：

__global__ void add2_kernel(float* c,
                            const float* a,
                            const float* b,
                            int n) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
            i < n; i += gridDim.x * blockDim.x) {
        c[i] = a[i] + b[i];
    }
}

• 主函数：就是输入Tensor，然后调用核函数的地方：

void launch_add2(float* c,
                 const float* a,
                 const float* b,
                 int n) {
    dim3 grid((n + 1023) / 1024);
    dim3 block(1024);
    add2_kernel<<<grid, block>>>(c, a, b, n);
}

2. 编写基于Setup编译的配置文件

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name="add2",
    include_dirs=["include"],
    ext_modules=[
        CUDAExtension(
            "add2",		# 最终库的名称
            ["pytorch/add2_ops.cpp", "kernel/add2_kernel.cu"],	# 算子
        )
    ],
    cmdclass={
        "build_ext": BuildExtension
    }
)

3. 在Python中调用编译好的CUDA算子

import time
import torch
import sobel

img_size = (640,640)
a = torch.rand(img_size, device="cuda:0")	# input
b = a.clone()		# output
sobel.torch_launch_sobel(a, b, a.shape[0], a.shape[1])

参考文献：

• PyTorch编译并调用自定义CUDA算子(by setup.py)