はじめに

前回記事の続きという位置づけの記事です．

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

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

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型と宣言してしまうとこれができなかった）．以下はスカラテンソルの絶対値を返すabs関数の例です．

  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)とつけてやることで型変換をします．

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

関数内部でテンソルの形を取得したい

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

// 方法１
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

// 方法２：配列として受け取りたい時
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と同じものになり，「デバイスが異なる」というエラーを未然に防いでくれます．

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})を使って

torch::Tensor out = torch::zeros({100, 100});

のように初期化してしまうとこれは必ずCPUテンソルになってしまい，エラーの原因になります．デバイスをCUDAにするためには初期化の際に明示的にデバイスを指定してやる必要があります．

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

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

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

#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を使おうとする方のお役に立てると幸いです．