[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 = 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};
}

編譯

我們使用setup.py來編譯我們剛寫的C++和CUDA文件，將其作為一個Python庫。

python setup.py install --user

最后，只需要在前面寫的myLayerNorm中import mylayer_cuda即可。