【CUDAプログラミング】PyTorch C++/CUDA APIを使って躓いたところ&解決策集

# はじめに
[前回記事](https://qiita.com/windfall/items/5031d70c649b06a1534f)の続きという位置づけの記事です．

前回記事でPyTorchのCUDAエクステンションの基本的な使い方を紹介しました．本記事ではより発展的な処理をしたい人たちに向けて，単なる足し算以上の複雑な処理をするために知っておきたいことまとめます．扱う内容は以下です．

- `__device__`修飾子
- CUDA関数内での変数の型変換(`float` &rarr; `scalar_t`)
- `torch::Tensor`型変数の形の取得方法
- CUDAテンソルの生成方法
- テンソルのcontigous化

### CUDAカーネルから呼びだせる関数を定義したい
CUDAではあらかじめ定義した関数をカーネルから呼び出すことができます．ただし，ここで定義する関数は**スカラー**テンソル同士の演算という想定です．例えば，シグモイド関数は以下のように定義します．

```cuda
template <typename scalar_t>
__device__ __forceinline__ scalar_t sigmoid(scalar_t z) {
  auto out = 1.0 / (1.0 + exp(-z));
  return out;
}
```
ポイントは以下です．

- テンプレート関数で定義する
- 修飾子として`__device__`と`__forceinline__`をつける
- **テンソルの型は`scalar_t`とする**（実行時に自動的にテンソルとして処理してくれる）
- `exp()`や`tanh()`などの標準的な関数が使える
- `scalar_t`と数字(ここでの`1.0`など)は演算可能（**ただし，明確に`float x = 1.0;`のように定義された変数との演算はエラーを起こす**）

修飾子について，CUDAではホストとデバイス両方で実行されるカーネルは`__global__`という修飾子をつけましたが，デバイスのみで実行される関数には`__device__`，ホストのみで実行される関数には`__host__`をつけます．

また，scalar_t型変数と数字との演算が可能なのはかなり便利であり，不等号の演算結果もそのままbool値として扱うことができます（[C++で明確に変数を`torch::Tensor`型と宣言してしまうとこれができなかった](https://qiita.com/windfall/items/5f653802c469ab6c9f90#%E5%A4%89%E6%95%B0%E3%81%AE%E5%9E%8B%E5%A4%89%E6%8F%9B)）．以下はスカラテンソルの絶対値を返すabs関数の例です．

```cuda
  template <typename scalar_t>
  __device__ __forceinline__ scalar_t abs(scalar_t x){
    if (x > 0.0){
      return x;
    }else{
      return -x;
    }
  }
```

定義した関数は以下のようにカーネルの中から呼び出せます．

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

template <typename scalar_t>
__device__ __forceinline__ scalar_t sigmoid(scalar_t z) {
  return 1.0 / (1.0 + exp(-z));
}

/******************
  カーネル
******************/
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] = sigmoid(A[n][c] + B[n][c]);  // 定義したsigmoid関数を呼び出す
    }
  }
```


### テンソル以外の型の変数を扱いたい
複雑な関数を書こうとすると，テンソル以外の引数も受け取りたくなることがあります．当然，カーネルの引数にアクセッサーオブジェクト以外を渡すことは可能ですので，`int`なのか`float`なのか等，きちんと型を明確にして渡します．

問題は，`float`などの変数と`scalar_t`型変数の演算を行うときです．前述のとおり，そのままではエラーを起こしますので，以下のように変数の前に`(scalar_t)`とつけてやることで型変換をします．

```cuda
  template <typename scalar_t>
  __device__ __forceinline__ scalar_t add(scalar_t x, float y){
    return x + (scalar_t)y;
  }
```

### 関数内部でテンソルの形を取得したい
事前にテンソルの形がわからない時，関数の内部でテンソルの形を受け取りたくなるときがあります．`torch::Tensor`型の変数の形は以下の方法で取得できます．

```cuda
// 方法１
torch::Tensor X = torch::zeros({3, 10, 10});
int n = X.size(0);  // 3
int h = X.size(1);  // 10
int w = X.size(2);  // 10
```

```cuda
// 方法２：配列として受け取りたい時
c10::ArrayRef<long int> tensor_shape = X.sizes();  // size{s}なので注意
int n = tensor_shape[0];
```



# トラブルシューティング
### CUDAテンソルを入力したはずが返り値がCPUテンソルになっていた
ビルドは通ったのに，いざ使おうとしたら勾配計算時にGPUに転送したはずのテンソル(cuda テンソル)がいつの間にかCPUテンソルになっていてエラーを起こす，ということがあります．これはC++/CUDA関数内部で生成したテンソルがCPUテンソルな可能性があります．

例えば，前記事で行列`A`と`B`を足すという操作をするとき，先に返り値となるテンソル`out`を`torch::zeros_like()`を用いて用意しました．この`torch::zeros_like()`はテンソルの形を形を変えることはできませんが，テンソルのデバイス情報も引き継いでくれるためテンソル`out`のデバイスは入力テンソル`A`と同じものになり，「デバイスが異なる」というエラーを未然に防いでくれます．

```cuda
std::vector<torch::Tensor> my_add_cuda_forward(
  torch::Tensor A,
  torch::Tensor B
){
  // 返り値のplaceholderを先に宣言しておく
  torch::Tensor out = torch::zeros_like(A);
  ...  // 
```
一方で自由な形のテンソルを生成したいというときもあります．例えばテンソル`out`を`torch::zeros({100, 100})`を使って

```cuda
torch::Tensor out = torch::zeros({100, 100});
```
のように初期化してしまうとこれは**必ずCPUテンソル**になってしまい，エラーの原因になります．デバイスをCUDAにするためには初期化の際に明示的にデバイスを指定してやる必要があります．

```cuda
torch::Device device = torch::kCUDA;
torch::Tensor out = torch::zeros({100, 100}, device);
```

### 変数がcontiguousじゃないというエラー`RuntimeError: tensor must be contiguous `が出た

チュートリアルにしたがってC++/CUDA APIを使っていると以下のマクロを設定すると思いますが，

```cpp
#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)
```
この`CHECK_CONTIGUOUS(x)`に引っかかり``RuntimeError: (チェックした変数) must be contiguous`というエラーになることがあります．こういうときは引っかかった変数`x`を`x = x.contiguous();`とすることでcontigousにしてやることができます．

# 終わりに
自分がやや複雑な処理をCUDAで並列計算させようとしたとき，PyTorch CUDA APIを使って躓いたところをまとめました．これからPyTorch CUDA APIを使おうとする方のお役に立てると幸いです．