[CUDA] 手写一个PyTorch的算子
（其实是本人之前上过的分布式机器学习课程的一个作业，这里简单记录一下）
我们都知道，PyTorch里的算子是跑在GPU上的。虽然最外层的接口是python，最内部的实现其实是CUDA。那么，一个python代码是如何一步步的调用内层的CUDA代码的呢？这里用一个简单的例子来讲解一下：
自定义nn.Module
我们想要实现一个自定义的LayerNorm算子，其前向传播公式如下
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 = epsself.elementwise_affine = elementwise_affineif 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):@staticmethoddef forward(ctx, input, normalized_shape, eps):# 将输入存下来供反向传播使用ctx.save_for_backward(input)ctx.normalized_shape = normalized_shapectx.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, CUDAExtensionsetup(// 要创建的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};
}
编译
我们使用setup.py来编译我们刚写的C++和CUDA文件，将其作为一个Python库。
python setup.py install --user
最后，只需要在前面写的myLayerNorm中import mylayer_cuda即可。
