[CUDA] 手写一个PyTorch的算子

（其实是本人之前上过的分布式机器学习课程的一个作业，这里简单记录一下）

我们都知道，PyTorch里的算子是跑在GPU上的。虽然最外层的接口是python，最内部的实现其实是CUDA。那么，一个python代码是如何一步步的调用内层的CUDA代码的呢？这里用一个简单的例子来讲解一下：

自定义nn.Module

我们想要实现一个自定义的LayerNorm算子，其前向传播公式如下

y = \frac{x - E [x]}{\sqrt{V a r [x] + ϵ}}

$y=\frac {x-\mathrm{E}[x]} {\sqrt{\mathrm{Var}[x]+\epsilon}}$

PyTorch的官方LayerNorm算子的接口可以参考https://pytorch.org/docs/1.5.0/nn.html#layernorm。

为了简便，这里只考虑normlized_shape为最后一维 dim size 且elementwize_affine=False的情况。

OK，我们按照官方的接口，自己写一个myLayerNorm类，继承nn.Module。

class myLayerNorm(nn.Module):
    __constants__ = ['normalized_shape', 'eps', 'elementwise_affine']

    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=False):
        super(myLayerNorm, self).__init__()
        if isinstance(normalized_shape, numbers.Integral):
            normalized_shape = (normalized_shape,)
        self.normalized_shape = tuple(normalized_shape)
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        if self.elementwise_affine:
            print("Do not support elementwise_affine=True")
            exit(1)
        else:
            # 注册需要进行更新的参数
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)

    def forward(self, input):
        return myLayerNormFunction.apply(
            input, self.normalized_shape, self.eps)

另外定义myLayerNormFunction函数，执行具体的前向/反向传播

class myLayerNormFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, normalized_shape, eps):
        # 将输入存下来供反向传播使用
        ctx.save_for_backward(input)
        ctx.normalized_shape = normalized_shape
        ctx.eps = eps

        # 调用外部的cuda方法
        output = mylayer_cuda.forward(input, *normalized_shape, eps)

        return output[0]

Python调用C++&CUDA

下一步，我们要实现从python调用c++&cuda的函数。具体做法是通过setuptools。我们编写一个setup.py脚本，内容如下：

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

setup(
    // 要创建的python类
    name='mylayer_cuda',
    ext_modules=[
        CUDAExtension('mylayer_cuda', [
            'mylayer_cuda.cpp',
            'mylayer_cuda_kernel.cu',
        ])
    ],
    cmdclass={
        'build_ext': BuildExtension
    })

其中，mylayer_cuda是我们要创建的一个python类，它的具体实现在mylayer_cuda.cpp和mylayer_cuda_kernel.cu中。其中的CUDAExtension和BuildExtension是pytorch为我们提供好的两个拓展，专门用于编译CUDA与PyTorch相关的库。

另一边，我们需要在mylayer_cuda.cpp中，提供一套用于Python访问的接口，代码大致如下：

// 调用pytorch的C++拓展库
#include <torch/extension.h>

// 函数的返回值是若干个Tensor的tuple
std::vector<torch::Tensor> mylayer_forward(
    torch::Tensor input,
    int normalized_shape,
    float eps) 
{
    // 具体的CUDA实现
    return mylayer_cuda_forward(input, normalized_shape, eps);
}

// 对应前面python代码中的mylayer_cuda.forward(...)
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &mylayer_forward, "myLayerNorm forward (CUDA)");
}

编写CUDA代码

最后，我们只需要用CUDA写出myLayerNorm的具体实现即可，然后在C++中调用它。

编写mylayer_cuda_kernel.cu代码：

#include <torch/extension.h>

#include <cuda.h>
#include <cuda_runtime.h>

// 具体的kernel函数
template <typename scalar_t>
__global__ void layer_forward_kernel(
	scalar_t* A,
    scalar_t eps,
    const int M, 
    const int N) 
{
    const int row = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (row >= M) {
        return ;
    }

    scalar_t l1_sum = 0;
    scalar_t l2_sum = 0;
    for (int i = row * N; i < (row + 1) * N; i ++) {
        l1_sum += A[i];
        l2_sum += A[i] * A[i];
    }

    scalar_t avg = l1_sum / N;
    scalar_t var = (l2_sum / N - avg * avg);
    scalar_t mul = 1.0 / sqrtf(var + eps);
    scalar_t add = - avg * mul;
    for (int i = row * N; i < (row + 1) * N; i ++) {
        A[i] = A[i] * mul + add;
    }
}

std::vector<torch::Tensor> mylayer_cuda_forward(
    torch::Tensor input,
    int normalized_shape,
    float eps)
{
    const int N = input.size(-1);
    const int M = input.numel() / N;

    const int block = 256;
    const int grid = (M - 1) / 256 + 1;

   	// AT_DISPATCH_FLOATING_TYPES是PyTorch提供的工具
    // 它可以根据输入的数据类型，调度相应的CUDA内核
    AT_DISPATCH_FLOATING_TYPES(input.type(), "mylayer_cuda_forward", ([&] {
        // 调用kernel函数
        layer_forward_kernel<scalar_t><<<grid, block>>>(
            input.data<scalar_t>(),
            (scalar_t)eps,
            M,
            N);
        }));
    
    return {input};
}