【CUDAでGPUプログラミング】CUDAをかいてPyTorchのカスタムモデルの一部をGPU上で並列計算させる

# はじめに

### この記事から得られる知識
この記事を読むことで，以下を知ることができます．

- PyTorchモデルの一部をC++/CUDAで書き直して並列化する方法
- GPUプログラミングの全体像(スレッド・ブロックの分割方法含む)

「CUDAを書く」ということはなんとなく難しそうと感じていましたが、PyTorchが提供するAPIのおかげで思った以上に簡単と感じましたので，実際に自分で実装をしてみて詰まってしまった点（次回記事）などまとめて少しでもこれからやる人の御役に立てれば幸いです．

[PyTorchのC++/CUDA APIの公式チュートリアル](https://pytorch.org/tutorials/advanced/cpp_extension.html#writing-a-mixed-c-cuda-extension)と[公式実装](https://github.com/pytorch/extension-cpp/tree/master/cuda)を参考にしていますが，そちらを読んだ上で説明が不足していると感じた部分を補えるような記事にしました．全体の見通しを良くするため，この記事では極めて簡単な並列計算（行列の足し算）を例に説明をしていきます．



この記事を読むにあたって、ある程度PyTorchとPyTorch C++ APIを知っているとより理解が捗ると思います．以下のリンクをご参考にしてください．

- [PyTorch C++ Extension](https://pytorch.org/tutorials/advanced/cpp_extension.html#)
- [PyTorch の一部だけをC++で書き換えて高速化する - PyTorchカスタムC++エクステンション -](https://qiita.com/windfall/items/5f653802c469ab6c9f90)

# GPUプログラミングの基礎知識

### 基本用語
[こちら](https://qiita.com/Keech/items/381495ed90e012d69f1a)の記事が大変わかりやすいので，あまりGPUに詳しくない方はまずはそちらをご参照ください．以下，この記事を読むにあたっては，最低限以下の単語を知る必要があります．

- ホスト：処理を呼び出す側．CPU
- デバイス：処理の依頼を受け取る側・GPU
- カーネル：デバイス上で行われる処理
- スレッド：処理を実行する最小単位
- ブロック：スレッドの集まり
- グリッド：ブロックの集まり

### GPUプログラミングのイメージ
まずはGPUをつかった並列計算の仕組みのイメージを持ちましょう．
大雑把に言ってしまえば，GPUはCPUから**カーネル**と呼ばれる計算の処理をする依頼をうけ，GPUにある膨大な数のコアが依頼された計算を並列にし，その結果をCPUに返す（正確にはCPUはGPUの処理が完了するのを待たない）ことで並列計算を行います．
![qiita_cuda_カーネル.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/367429/ae8c81eb-7652-6fce-d263-68d6c09096f9.png)

データの転送や待機などの面倒な処理はCUDAが提供するAPIが行ってくれますので，**実際にプログラマがコードするものはカーネルが主になります．**後述しますが，カーネルの関数をかくときにはコアの位置のような情報を含む変数を使用することができますので，その変数を使用することですべてのコアに対して適用可能な抽象的な処理を記述することができます．

# PyTorch C++/CUDA APIで行列の足し算をする

行列の足し算は，２つの行列を同じ成分の要素ごとに足し算をするだけなので，並列計算が可能です．C++/CUDAで行列の足し算をGPUで行うモジュール`my_add`をビルドし，Pythonからそのモジュールを呼び出し実行することを最終目標とします．

### 環境準備
PyTorchのC++版を[こちら](https://pytorch.org/)よりインストールしてください．今回はGPUが手元にあるという想定なので，computer platformでCUDAを選択します．

次に，gccのバージョンを適切なものにします．ターミナルで

```shell
gcc -v
```
としてバージョンを確認します．筆者のOSはUbuntu20.04LTSですが，Ubuntu20はgccの9系が標準となります．CUDAをコンパイルする際、gcc８系までしかサポートしていないとエラーが吐かれるかと思いますので，Ubuntu20ユーザーは以下を参考に8系のgccを用意するとよいでしょう．

- [Ubuntu20.04でgcc-9からgcc-8に変更する方法](https://0h-n0.github.io/blog/2020/05/07/ubuntu-gcc.html)

### CUDAファイルの作成
まずはCUDAファイル`my_add_kernel.cu`というCUDAファイルを用意し，以下のように大きさが100x100テンソル`A,B`を受け取って足し算の結果であるテンソル`out`を返す関数`my_add_cuda_forward`を記述します．この関数`my_add_cuda_forward`はまだカーネルではなく，カーネルを起動するための関数になります．

```my_add_kernel.cu
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

std::vector<torch::Tensor> my_add_cuda_forward(
  torch::Tensor A,
  torch::Tensor B
){
  // 返り値のplaceholderを先に宣言しておく
  torch::Tensor out = torch::zeros_like(A);

  // blockとthreadの分割設定
  int n_row = 100;  // 行列の行数
  int n_col = 100;  // 行列の列数
  const int threads = 1024;  //これが最大値, number of threads per block
  const dim3 blocks((n_col + threads - 1) / threads, n_row);  // number of blocks per girid
  // カーネルをGPUに依頼
  AT_DISPATCH_FLOATING_TYPES(A.type(), "my_add_forward_cuda", ([&] {
    my_add_cuda_forward_kernel<scalar_t><<<blocks, threads>>>(
      A.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>(),
      B.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>(),
      out.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>()
    );
  }));
  return {out};
}
```
基本的にPyTorch C++ APIと同じ記述ができます．

まず最初に返り値となるテンソル`out`を生成しています．計算結果はこのテンソル`out`に格納されていきます．以降はカーネルを起動していく部分になります．

### 使用するスレッド・ブロックの分割設定
`n_row`を行列の行数（縦方向の大きさ），`n_col`を列数（横方向の大きさ）としました．続く部分では起動するスレッド, ブロックの設定に必要な変数を設定しています．

```cuda
const int threads = 1024;  //これが最大値, nubmer of threads per block 
const dim3 blocks((n_col + threads - 1) / threads, n_row);  // number of blocks per girid
```
ここでは一つのブロックを(x方向に)1024個のスレッドに分割し，さらにグリッドをx方向に`(n_col + threads - 1)/threads`，y方向に`n_row`個のblockに分割するという設定をしています．x方向の数の設定方法がややこしいですが，これは`n_col`が0のときは0を返し，1以上のときは「`n_col`の数のスレッドをカバーできる最小ブロック数」を返させるためです．

このように`threads`と`blocks`を設定するのが[最適というわけではありません](https://stackoverflow.com/questions/33245737/cuda-estimating-threads-per-blocks-and-block-numbers-for-2d-grid-data)が，いわゆるテンソルの「バッチ方向」をブロックのy方向に対応させることで各カーネルが担当すべき行列の要素のインデックスが明快になる（後述）ため，無難な書き方と言えます．

### カーネルの起動
次の以下の部分でCUDAカーネルに行ってほしい処理の依頼・実行が行われます．

```cuda
  AT_DISPATCH_FLOATING_TYPES(A.type(), "my_add_forward_cuda", ([&] {
    my_add_cuda_forward_kernel<scalar_t><<<blocks, threads>>>(
      A.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>(),
      B.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>(),
      out.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>()
    );
  }));
```
少しややこしく見えるので，以下で簡単に解説していきます．`AT_DISPATCH_FLOATING_TYPES`の引数には順番に扱うテンソルの型，関数名（エラーメッセージのため），ラムダ式を渡します．参考までにC++でのラムダ式は`[キャプチャリスト](パラメータリスト){処理;}`という形をとり、例えば２つの整数の和を返すラムダ式は

```cpp
auto sum = [](int a, int b){return a + b;}
```
のように書くことができます．`[&]`というキャプチャ記法が指定されるとラムダ式の外にある変数を参照して（コピーではなく）ラムダ式の中で使用することができます．また、パラメータリストは省略することもできます．

そして，ラムダ式の中で呼ばれている関数`my_add_cuda_forward_kernel`がこれから定義していくカーネルになります．カーネルを呼ぶときには`(関数名)<型エイリアス><<<(blockとthreadを設定する変数)>>>(引数)`という特殊な形をとります．
引数にはテンソルオブジェクトそのものではなく，ポインタのように参照渡しで変数を渡す必要があります．普通にポインタを渡すのも可能なようですが，PyTorch APIが提供する`.packed_accessor`メソッドを使ってアクセッサーオブジェクトを渡してやることで，カーネル定義時によりテンソルが扱いやすくなります．アクセッサの`<>`の中の数字はテンソルの次元です．

### カーネルの定義
以下はカーネルの定義部分です．カーネルはtemplate関数を用いて定義します．`__global__`修飾子を関数の型の前につける必要がありますが，これはCUDA特有のもので，この関数がカーネルでありデバイスとホスト上の両方で実行される関数であることを示しています．また，カーネルでは基本的にポインタ操作をするため返り値がなく，故に`void`型となります．引数としてアクセッサオブジェクトを受け取るので，以下のようにして引数の型宣言をしましょう．

```my_add_kernel.cu
template <typename scalar_t>
__global__ void my_add_cuda_forward_kernel(
  const torch::PackedTensorAccessor<scalar_t, 2, torch::RestrictPtrTraits, size_t> A,
  const torch::PackedTensorAccessor<scalar_t, 2, torch::RestrictPtrTraits, size_t> B,
  torch::PackedTensorAccessor<scalar_t, 2, torch::RestrictPtrTraits, size_t> out
){    
    const int n = blockIdx.y;  // y方向インデックス
    const int c = blockIdx.x * blockDim.x + threadIdx.x;  // x方向インデックス

    if (c < A.size(1)){
      out[n][c] = A[n][c] + B[n][c];
    }
  }
```
前述の通り，カーネル（正確には`__global__,__device__`がつく関数）は以下の変数が自動定義されます．

| 型 変数名 | 意味 |
|:-:|:-:|
|  `dim3 gridDim` | グリッドがそれぞれの方向に何個のブロックに分割されているか  |
|  `dim3 blockDim` | ブロックがそれぞれの方向に何個のスレッドに分割されているか  |
|  `dim3 blockIdx` |  グリッド内のブロックのインデックス |
| `dim3 threadIdx`  |  ブロック内のスレッドのインデックス |

これらの変数を使うことで，各スレッドのGPU内での位置のようなものを以下のように分かりやすい形で取得することができます．

```cuda
    const int n = blockIdx.y;  // y方向インデックス
    const int c = blockIdx.x * blockDim.x + threadIdx.x;  // x方向インデックス
```
スレッド位置の計算方法が納得できない，という方はGPUが以下のような構造をしているというイメージを持つと分かりやすいです．こちらの図は実際のGPUのアーキテクチャからはかけ離れた（こんな単純じゃない）ものですが，今回のようなブロック・スレッド分割方法を指定しておけば，プログラマーからはこのようなイメージでthreadの位置を捉えることができます．

![qiita_cuda_block_image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/367429/8c239b45-b0da-511d-ed89-a1d7bbf89efa.png)

実はスレッドは３次元方向に分割可能ですが，今回は1024とだけ指定したので定義時に`(x,y,z)=(1024, 1, 1)`と分割されました．よって，各ブロックはx方向に多数のスレッドに分割されていると思えばよく，X方向の位置を`blockIdx.x * blockDim.x + threadIdx.x`と表すことができます．この表現方法はよく使われるので，とりあえずｘ位置をこれで計算しておく，くらいの気持ちでいると良いでしょう．また，y方向には単純にブロックのy方向のインデックスを返せばいいことが上の図より分かります．

ここで各スレッドは行列の`[n, c]`成分の足し算のみを行えばよいので，以下のように記述します．

```cuda
if (c < A.size(1)){
  out[n][c] = A[n][c] + B[n][c];
}
```
if文によって，x位置が行列の横方向のサイズより大きいスレッドが処理を行わないようにしています．以上が行列の各成分の足し算を並列に行うカーネルの定義方法になります．

### Backward関数の定義とC++ファイルの作成
以上まででCUDAのPyTorch APIを使った基本的なカーネルの書き方について解説しました．あとは以下をすることでPyTorchから呼び出せるモジュールを作成することができます．

1. CUDAファイル`my_add_kernel.cu`にbackward関数も定義する
1. C++ファイル`my_add.cpp`を作成し，C++からCUDAで定義した関数を呼びだす
1. Pythonファイル`setup.py`を作成し，実行してPythonにC++/CUDAで作成したモジュールをバインドする

基本的にはお決まりの記述をするだけなので，残りはファイル全体と簡単な解説をつけるのみとします．参考リンクも貼りましたので，不明な点はリンク先を参照してもらればと思います．

以下のように３つのファイル全てを同じディレクトリ`my_add`上にファイルを置きます．

```sh
./my_add
├── my_add.cpp
├── my_add_kernel.cu
└── setup.py
```

#### CUDAファイル`my_add_kernel.cu`の全体

行列の足し算のbackward用の関数`my_add_cuda_backward`を用意します．行列の足し算`A + B = out`において勾配計算時，`A`と`B`の勾配は共に`out`の勾配に一致します．backward関数の定義をしたことがない，そもそも自動微分の仕組みが分からないという方は，PyTorchでの自動微分の解説をしている以下の記事を参照してください．

- [【PyTorch】自作関数の勾配計算式（backward関数）の書き方①](https://qiita.com/windfall/items/073cbb4ffdfab356e495)
- [【PyTorch】自作関数の勾配計算式（backward関数）の書き方② 〜多変数出力の場合〜](https://qiita.com/windfall/items/0de42850be355c43ec3e)


```my_add_kernel.cu
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

template <typename scalar_t>
__global__ void my_add_cuda_forward_kernel(
  const torch::PackedTensorAccessor<scalar_t, 2, torch::RestrictPtrTraits, size_t> A,
  const torch::PackedTensorAccessor<scalar_t, 2, torch::RestrictPtrTraits, size_t> B,
  torch::PackedTensorAccessor<scalar_t, 2, torch::RestrictPtrTraits, size_t> out
){
    //batch index
    const int n = blockIdx.y;
    //column index
    const int c = blockIdx.x * blockDim.x + threadIdx.x;

    if (c < A.size(1)){
      out[n][c] = A[n][c] + B[n][c];
    }
  }

/************************************
FORWARD関数、BACKWARD関数を書く
*************************************/

std::vector<torch::Tensor> my_add_cuda_forward(
  torch::Tensor A,
  torch::Tensor B
){
  // 返り値のplaceholderを先に宣言しておく
  torch::Tensor out = torch::zeros_like(A);

  // blockとthreadの分割設定
  int n_row = 100;
  int n_col = 100;
  const int threads = 1024;  //これが最大値, nubmer of threads per block
  const dim3 blocks((n_col + threads - 1) / threads, n_row);  // number of blocks per girid
  // カーネルをGPUに依頼
  AT_DISPATCH_FLOATING_TYPES(A.type(), "my_add_forward_cuda", ([&] {
    my_add_cuda_forward_kernel<scalar_t><<<blocks, threads>>>(
      A.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>(),
      B.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>(),
      out.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>()
    );
  }));
  return {out};
}


std::vector<torch::Tensor> my_add_cuda_backward(
  torch::Tensor grad_out
){
  // grad_A, grad_Bともにgrad_outに一致する
  return {grad_out, grad_out};
}
```

#### C++ファイル`my_add.cpp`の全体
C++ファイルはPythonとCUDAをつなぐような役割をし，CUDA関数を呼び込んだC++関数をビルドして生成したモジュールをPythonから呼ぶといった流れになります．このC++関数をビルドしてPythonのPyTorchモデルでよびこむ方法はPyTochカスタムC++エクステンションと呼ばれます．

ここではC++ファイルの書き方の詳細については触れませんが，C++エクステンションの使用方法が頭に入っているとここでの各関数の意味がより分かると思うので，本記事冒頭で紹介したC++エクステンションに関するリンクも是非参考にしてください．

<!-- - [PyTorch C++エクステンションの公式チュートリアル](https://pytorch.org/tutorials/advanced/cpp_extension.html#)
- [PyTorch の一部だけをC++で書き換えて高速化する - PyTorchカスタムC++エクステンション -](https://qiita.com/windfall/items/5f653802c469ab6c9f90) -->


```my_add.cpp
#include <torch/extension.h>
#include <iostream>
#include <vector>
// テンソルをチェックするためのマクロを定義
#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
/***************************
 CUDAの関数の名前だけ宣言しておく
****************************/
std::vector<torch::Tensor> my_add_cuda_forward(
  torch::Tensor A,
  torch::Tensor B);
std::vector<torch::Tensor> my_add_cuda_backward(
  torch::Tensor grad_out);

/***********************************************
 FORWARD, BACKWARD関数の定義
 それぞれPythonでのmy_add.forward(), my_add.backward()で呼び出される
*************************************************/
std::vector<at::Tensor> my_add_forward(
  torch::Tensor A,
  torch::Tensor B
){
    CHECK_INPUT(A);
    CHECK_INPUT(B);
    return my_add_cuda_forward(A, B);
  }

std::vector<torch::Tensor> my_add_backward(
  torch::Tensor grad_out
){
    grad_out = grad_out.contiguous();
    CHECK_INPUT(grad_out);
    return my_add_cuda_backward(grad_out);
  }
/*********************************************
 以下はPythonから呼び出せるモジュールにするために必要
********************************************/
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &my_add_forward, "myadd forward (CUDA)");
    m.def("backward", &my_add_backward, "myadd backward (CUDA)");
  }

```

はじめにCUDA関数の名前だけ宣言をし，その後，関数`my_add_forward`,`my_add_backward`でそれぞれ対応するCUDA関数を呼び出しています．引数となるテンソルは`CHECK_INPUT()`によりCUDA関数に渡す前に，1. GPUに転送されているか 2. contiguousかを確認し，条件が満たされない場合にエラーを返します．

#### Pythonファイル`setup.py`でC++/CUDA関数をビルド & Pythonにbindする
以下のように記述してください．

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

setup(
    name='my_add',
    ext_modules=[
        CUDAExtension('my_add', [
            'my_add.cpp',
            'my_add_kernel.cu',
        ])
    ],
    cmdclass={
        'build_ext': BuildExtension
    })
```

以上で準備は完了です．`my_add`ディレクトリ上で以下を実行し，長いwarningの後に`Finished processing dependencies for ...`の言葉が出たら成功です．

```sh
python setup.py install
```

# 実際にテストする

以下のようなpythonファイルを用意します．例によって[他の記事](https://qiita.com/windfall/items/5f653802c469ab6c9f90)で解説済みなので不明点はそちらを参照してください．

```main.py
import torch
import torch.nn as nn
import my_add
from pdb import set_trace as db


class MyaddFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A, B):
        output = my_add.forward(A, B)
        ctx.save_for_backward(*output[1:], A, B)
        return output[0]

    @staticmethod
    def backward(ctx, d_out):
        saved_tensors = ctx.saved_tensors
        d_input = my_add.backward(d_out)
        dA, dB = d_input
        return dA, dB

class Myadd(nn.Module):
    def __init__(self):
        super(Myadd, self).__init__()
        pass

    def forward(self, A, B):
        return MyaddFunction.apply(A, B)

if __name__=='__main__':

    use_cuda = torch.cuda.is_available()
    if use_cuda:
        device = torch.device('cuda')
    else:
        raise ValueError("CUDA is not available!!")


    model = Myadd()
    model.to(device)
    _A = torch.rand(100, 100).requires_grad_()
    A = _A.to(device)
    B = torch.rand(100, 100).to(device)

    out = model(A, B)
    print('model out: ', out)
    loss = torch.sum(out)
    loss.backward()
    print("_A.grad: ", _A.grad)  # None --> tensor([1., ...])に変わっている!!
```
正しく実行できました．

# 終わりに
PyTorch CUDA APIの使い方を簡単な例を用いて紹介しました．ただ，実際にCUDAを書こうとしている方はただの足し算のような簡単な演算ではなく，より複雑な処理を並列計算させたい人が多いのではないでしょうか？次回，if文や型変換などより複雑な並列計算をさせるために知っておきたいことなどをまとめて紹介します．