讨论 2020-11-26
=
###### tags: `tutorials` `Fudan` `2020`
## CUDA编程简单介绍

Compute Unified Device Architecture (CUDA) 是N卡支持的一套并行语言，我们现在用的pytorch，tensorflow等平台都是在其基础之上搭建的，CUDA整体和C++很接近，主要区别在于它天然支持并行。今天主要做一些基本介绍，以及如何在pytorch中添加自定义的CUDA模块。

### 如何在pytorch里使用一个cuda kernel
首先，我们先介绍一下，如何在pytorch中使用cuda kernel，pytorch有提供比较详细的例子，我们这里只做简述
https://pytorch.org/tutorials/advanced/cpp_extension.html

通常来讲，我们需要三个部分，他们通常在三个不同文件中，第一个是 **.cu** 文件，他表示了我们用cuda语言所写的内容，即计算流程。第二个是 **.cpp** 文件，它负责把cuda程序封装成一个python兼容，特别是pytorch兼容的python module。通常来讲，这个文件就是个参数表，没有任何实际计算流程，所以可以按套路编写。第三个文件则是 **.py** 文件，这个也可以写在自己的主程序中，这一步负责把上一步得到的python module封装成一个 pytorch module，其实上一步的结果就是python可运行的，但为了使用pytorch的autograd功能，我们还需要第三步。具体而言就是需要定义一个pytorch operator，这涉及到它的forward和backward，以及存储。

### 基础概念
目前N卡大概有两级并行，一是很多thread组成了block，二是很多block组成了grid。thread是最小的计算单元，每个block拥有的thread数量是有限的，一般为1024，同一个block里面的thread有共享存储空间（需要显示使用shared memory）。

### 读写冲突
CUDA程序一般是以计算为主的，所以与多进程不同，通常不会采用锁和等待的机制，这就要求我们写的程序要考虑到读写冲突。也就是要避免冲突，即每个被写入的元素只被一个线程写入，同时不存在任何先写后读的情况，如果要先写后读，要开多个kernel

## 如何测量效率信息

### 测正确性
用allclose检查forward和backward的正确性，也就是用pytorch实现一遍，和自己写的kernel做对比，注意allclose有时太严格，要手动调整精度误差的许可范围，或者使用double类型。
### 测速
测速时用time即可，但一定要注意排除最开始的几次，并且每次执行都torch.cuda.synchronize()同步结果。另外测速受输入变量的形状影响较大，需要考虑实际使用的场景。

### 测显存
用这个torch.cuda.max_memory_reserved才是pytorch在gpu上占的总显存数，allocated为Tensor所占，不是全部。
### profile
python端推荐cProfile，cuda端推荐nvprof

## 自动求导

### Custom Op
为了兼容pytorch的自动求导框架，我们需要继承 autograd.Function，并且为它编写 forward和backward函数，ctx.save_for_backward可以把前向的变量存下来，以便反向时使用。注意反向的input只有上层传下来的gradient，原始输入是要自己保存的。

## 一个完整的例子，两个向量的外积

#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
namespace {
    // (b, n, 1) * (b, m, 1) = (b, n, m) n>m
template <typename scalar_t>
__global__ void op_kernel(const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, scalar_t* __restrict__ C, const int b, const int n, const int m) {
    int i = blockIdx.x;
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    for (int j=i; j<b; j+= gridDim.x) {
        for (int x=tx; x<n; x+= blockDim.x){
            for (int y=ty; y<m; y+= blockDim.y) {
                C[(j*n+x)*m+y] = A[j*n+x] * B[j*m+y];
            }
        }
    }
}


}//end of namespace
torch::Tensor op_cuda(
    const torch::Tensor &A,
    const torch::Tensor &B) {
    // A: (b, n), B: (b, m)
    const auto b = B.size(0);
    const auto n = A.size(1);
    const auto m = B.size(1);

    auto y = torch::zeros({b, n, m}, A.options());

    const dim3 threads((n<(1024/m-1))?n:(1024/m-1) , m);
    const dim3 blocks(2048);

    AT_DISPATCH_FLOATING_TYPES(A.scalar_type(), "outer", ([&] {
        op_kernel<scalar_t><<<blocks, threads>>>(
            A.data_ptr<scalar_t>(),
            B.data_ptr<scalar_t>(),
            y.data_ptr<scalar_t>(),
            b, n, m);
    }));

    return y;
}

#include <torch/extension.h>
#include <vector>
#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(),#x "must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x "must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

torch::Tensor op_cuda(const torch::Tensor&A,const torch::Tensor&B);

torch::Tensor op(const torch::Tensor&A,const torch::Tensor&B){
    CHECK_INPUT(A);
    CHECK_INPUT(B);
    return op_cuda(A,B);
}




PYBIND11_MODULE(TORCH_EXTENSION_NAME,m){
    m.def("demo_op",&op,"A demo");

}

import torch
from torch.autograd import Function
import demo_op
import time
class DEMOfunc(Function):
    @staticmethod
    def forward(ctx, A, B):
        with torch.no_grad():
            outs = demo_op.demo_op(A, B)
            ctx.save_for_backward(A, B)
        return outs

    @staticmethod
    def backward(ctx, grad):
        with torch.no_grad():
            A, B = ctx.saved_variables
            dA = torch.matmul(grad, B[:,:,None])
            dB = torch.matmul(grad.permute(0, 2,1), A[:,:,None])
        return dA[:,:,0], dB[:,:,0]

Fop = DEMOfunc.apply

if __name__=='__main__':
    b, N, M = 50, 200, 300
    A = torch.randn(b, N).double().cuda()
    A.requires_grad=True
    A1 = A.clone().detach()
    A1.requires_grad=True
    B = torch.randn(b, M).double().cuda()
    B.requires_grad=True
    B1 = B.clone().detach()
    B1.requires_grad=True
    C = torch.matmul(A[:,:,None], B[:,None,:])
    D = Fop(A1, B1)
    print('Forward Check')
    torch.allclose(C,D,1e-5)
    loss = torch.sum(C)
    loss1 = torch.sum(D)
    loss.backward()
    loss1.backward()
    print('Backward Check')
    torch.allclose(A.grad, A1.grad, 1e-5)
    torch.allclose(B.grad, B1.grad, 1e-5)

    baseline = []
    TT = 1000
    for i in range(TT):
        if i>5:
            A = torch.randn(b, N).cuda()
            B = torch.randn(b, M).cuda()
            torch.cuda.synchronize()
            st = time.time()
            C = torch.matmul(A[:,:,None], B[:,None,:])
            torch.cuda.synchronize()
            baseline.append(time.time()-st)
    print('Baseline', sum(baseline), sum(baseline)/TT)

    ours = []
    TT = 1000
    for i in range(TT):
        if i>5:
            A = torch.randn(b, N).cuda()
            B = torch.randn(b, M).cuda()
            torch.cuda.synchronize()
            st = time.time()
            C = Fop(A1, B1)
            torch.cuda.synchronize()
            ours.append(time.time()-st)
    print('Ours', sum(ours), sum(ours)/TT)

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
    name='demo_op',
    version="0.0.0",
    ext_modules=[
        CUDAExtension(
            'demo_op',[
                'demo1.cpp',
                'demo.cu',
            ]
        ),
    ],
    cmdclass={
        'build_ext': BuildExtension.with_options(use_ninja=False)
    }
)