RustでGPGPUしてみた(レイトレーシング,vector_add)

## はじめに
この記事は[三重大学 計算研 Advent Calendar 2019](https://qiita.com/advent-calendar/2019/calc-mie) 6日目の記事です．
この記事ではいつも通りRustの布教をしていきたいと思います．今回はRustでGPGPUする話です．

記事の内容ですが，手探りな部分が多いので，間違っているところや，改善点，コメントなどがあったら教えてください．

(2019/12/23 16:34 リポジトリが非公開になっていたのを直しました)

## 概要

RustではLLVMのnvptx backendを使ってPTXコード(nvidiaのGPU向け中間言語)を生成できます．

RustでGPGPUするライブラリとしては[accel](https://github.com/rust-accel/accel)[1][10]などがありますが，今回は1から作ってみました．

初めはコードをすべて解説するつもりだったのですが、まとまりが悪かったので，CUDAのライブラリ周りのメモリ管理やカーネルのlaunchをまとめたcrateを作りました．
この記事ではRay Tracing in One Weekendのレイトレの実装を通して動作・実行時間の確認，vector-addの実装を通してcrateの使い方の解説，作るのに必要だった知識，困っていること・今後の課題について書いていきます．

今回作ったcrateは https://github.com/spica314/cuda-tools にあります．
また，サンプルコードは https://github.com/spica314/cuda-tools-examples にあります．

## cuda-tools-examples/ray-tracing

ソースコード: 
- https://github.com/spica314/cuda-tools-examples/tree/master/ray-tracing
- https://github.com/spica314/cuda-tools-examples/tree/master/ray-tracing-kernel

Ray Tracing in One Weekend[2]の最後の画像をレンダリングしてみます．

できた画像
![t.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280971/ff4dac10-9c33-fca2-f0e9-7e1998646ca7.png)
ちゃんと動いてそうですね．

縦400画素，横600画素，1画素あたり128レイ，485オブジェクトで，出力を含めた実行時間が5秒ぐらいとなりました(CPU:Ryzen 5 3600, GPU: GeForce GTX 1050Ti)．はじめて書いたときのCPU向けコード(同等ではない．おそらく遅め)だと1スレッドで460秒とかなのでたぶんはやいです．GPUすごい．

## examplesのvector-addの解説
この章では[cuda-tools-examples](https://github.com/spica314/cuda-tools-examples)の[vector-add](https://github.com/spica314/cuda-tools-examples/tree/master/vector-add)を通して[cuda-tools](https://github.com/spica314/cuda-tools)の使い方を解説します．

### crateの構成
`cuda-tools`を使う場合は，ホスト側(CPU側)用のcrateとデバイス側(GPU側)用のcrateの2つを用意します．
vector-addの場合は，[vector-add](https://github.com/spica314/cuda-tools-examples/tree/master/vector-add)がホスト側用のcrate、[vector-add-kernel](https://github.com/spica314/cuda-tools-examples/tree/master/vector-add-kernel)がデバイス側用のcrateとなります。

### buildの流れ
`vector-add`側で`cargo build`(もしくは`cargo run`)を実行すると、`build.rs`に基づいて`vector-add-kernel`の`nvptx64-nvidia-cuda`向けのコンパイルが走ります．そして，コンパイルされたPTXファイルを`vector-add`内の`include_kernel!()`マクロによって，`vector-add`内に文字列として展開します．
展開されたPTXコードはカーネルのlaunchに使います(正確にはmoduleのロード)．

### vector-add-kernel/src/lib.rs のコード
url: https://github.com/spica314/cuda-tools-examples/blob/master/vector-add-kernel/src/lib.rs

まずデバイス側のコードから説明していきます．

```rust
#![no_std]
#![feature(abi_ptx)]
#![feature(stdsimd)]
```

nvptx/nvptx64向けのコンパイルはstdが使えないので`#![no_std]`する必要があります．
`#![feature(abi_ptx)]`はCUDAの__global\_\_関数相当の関数を書くために必要なfeatureです．
`#![feature(stdsimd)]`はときどきいるので書いときます．

```rust
#[macro_use]
extern crate cuda_tools;

use cuda_tools::cuda_slice::CUDASlice;
use core::cell::UnsafeCell;
```
extern crateとuseです．
cuda_tools::cuda_sliceについては後述します．
core::cell::UnsafeCellも後述しますがメモリへの書き込み時に使います．

```rust
pub struct Arguments<'a> {
    pub xs: CUDASlice<'a, f32>,
    pub ys: CUDASlice<'a, f32>,
    pub zs: CUDASlice<'a, UnsafeCell<f32>>,
}
```
引数用の構造体です．
`cuda-tools`ではカーネルに渡す引数は1つのみとしているので，基本的には引数用の構造体を作る必要があります．
今回のvector_addでは`zs[i] = xs[i] + ys[i]`を計算します．メンバの`xs`,`ys`,`zs`がそれぞれ対応します．
`CUDASlice<'_, f32>`は`&[f32]`相当です．nvptx向けのコンパイル時には`&CUDASlice<'_, f32>`から`&[f32]`への`Deref`が使えます．

```rust

#[no_mangle]
#[cfg(not(target_arch = "nvptx64"))]
pub extern "ptx-kernel" fn vector_add(args: &Arguments) {}

#[no_mangle]
#[cfg(target_arch = "nvptx64")]
pub extern "ptx-kernel" fn vector_add(args: &Arguments) {
    let i = unsafe {
        use core::arch::nvptx::*;
        _block_dim_x() * _block_idx_x() + _thread_idx_x()
    } as usize;
    unsafe {
        if i < args.zs.len() && i < args.xs.len() && i < args.ys.len() {
            *args.zs[i].get() = args.xs[i] + args.ys[i];
        }
    }
}

```

`extern "ptx-kernel"`をつけると，\_\_global\_\_相当の関数になります．つまり，その関数からGPUでの計算を開始できるようになります．
1つめの`vector_add`はホスト側で型情報を使うためだけに書いてます．将来的にはcustum attributeで自動生成したいです．
2つめの`vector_add`がデバイス用のコードです．
まずそのスレッドが担当するインデックスを計算します．CUDAの`blockIdx`/`blockDim`/`threadIdx`相当の関数が`core::arch::nvptx`にあるので，それを使います．
次に，`zs[i] = xs[i] + ys[i]`の計算をします．`&CUDASlice<'_, f32>`は`Deref`で`&[f32]`になるので，スライスと同じ書き方で書けます．`zs`に関しては`UnsafeCell`内の値を書き換えるので，`UnsafeCell::get()`で`*mut f32`を手に入れて代入の処理をします．生ポインタ越しの書き込みなのでunsafeブロックが必要になります．
`if i < args.zs.len() && i < args.xs.len() && i < args.ys.len() {`のインデックスの範囲のチェックですが，この記述がない場合，sliceのbound checkに失敗したとき用の関数を呼び出すPTXコードが生成されます．しかし，その関数の実装がPTXコードに含まれないので，読み込もうとするとPTXコードのロードに失敗します( CUDA\_ERROR\_INVALID\_PTX )．`#[lang=...]`とかで実装するんだと思いますが，私はまだよくわかってないので毎回境界チェックを書くことにしています．

### vector-add/build.rs のコード
url: https://github.com/spica314/cuda-tools-examples/blob/master/vector-add/build.rs

ここからはホスト側のコードの説明です．

```rust
#[macro_use]
extern crate cuda_tools;

fn main() {    
    build_kernel!("../vector-add-kernel", "vector-add-kernel");
}
```
`build_kernel!()`の第一引数がデバイス用のクレートのパス，第二引数がそのクレートのクレート名です．
`build_kernel!()`ではデバイス用のクレートをPTXコードにコンパイルしたり，cargoのrerun-if-changed [7]を設定しています．
rerun-if-changedの設定は，デバイス用クレートのプログラム変更時にbuildが走るようにするためです．この設定をしないとデバイス用クレートのプログラムを変更してもbuildが走らず面倒です．

第二引数の情報は`cargo`をライブラリとして使えば第一引数から手に入る情報のはずなのでそのうちなくしたいです．

### vector-add/src/bin/vector-add.rs のコード

url: https://github.com/spica314/cuda-tools-examples/blob/master/vector-add/src/bin/vector-add.rs

```rust
fn main() {
    vector_add::run();
}
```
ライブラリ用のコードとバイナリ用のコードを分けています．ここではrun関数を呼び出すだけです．

### vector-add/src/lib.rs

url: https://github.com/spica314/cuda-tools-examples/blob/master/vector-add/src/lib.rs

```rust
#![feature(abi_ptx)]
```
extern "ptx-kernel"な関数を使うので書きます．

```rust
#[macro_use]
extern crate cuda_tools;

use core::cell::UnsafeCell;
use rand::prelude::*;
```
extern crateと各種useです．

```rust
const KERNEL: &str = include_kernel!();
```
`build.rs`が走ったときにできたPTXコードを`cuda-tools`の`include_kernel!()`マクロで文字列として展開します．

```rust
const N: usize = 1<<24;
```
今回の配列のサイズです．16777216要素です．

```rust
pub fn run() {
    let mut runtime = cuda_tools::runtime::Runtime::new(0, KERNEL).unwrap();
```
ホスト側のコードの本体です．
`cuda_tools::runtime::Runtime::new()`の第一引数にGPU番号，第二引数にPTXファイルの文字列を渡します．

```rust     
    runtime.record_function_name(vector_add_kernel::vector_add, "vector_add");
```
関数ポインタと関数名の対応をRuntimeに覚えさせます．将来的には裏でいい感じにしたいです．

```rust
    let mut rng = rand::thread_rng();
    let mut xs = vec![];
    for _ in 0..N {
        xs.push(rng.gen());
    }
    let mut ys = vec![];
    for _ in 0..N {
        ys.push(rng.gen());
    }
```
入力する配列をホスト側にいったん作ります．

```rust
    let mut zs = vec![];
    for i in 0..N {
        zs.push(UnsafeCell::new(0.0));
    }
```
出力用の配列を作るために，いったん`zs`を適当に作ります．


```rust
    let xs_d = runtime.alloc_slice(&xs).unwrap();
    let ys_d = runtime.alloc_slice(&ys).unwrap();
    let zs_d = runtime.alloc_slice(&zs).unwrap(); 
```
`Runtime::alloc_slice()`でホスト側のスライスを渡すと，デバイス側に必要なメモリ確保をした後，各要素の値をコピーします．


```rust
    let args = vector_add_kernel::Arguments {
        xs: xs_d,
        ys: ys_d,
        zs: zs_d,
    };
```
引数用の構造体の値を作ります．

```rust
    runtime.launch(vector_add_kernel::vector_add, &args, N/256, 1, 1, 256, 1, 1).unwrap();
```
カーネルを走らせます．第一引数が関数ポインタ，第二引数が引数，第三引数から3つが`gridDim{x,y,z}`，次の3つが`blockDim{x,y,z}`です．

```rust
    let zs = args.zs.to_host().unwrap();
```

`CUDASlice<'_, f32>::to_host()`でホスト側にすべての要素の値をコピーします．

```rust

    let zs: Vec<f32> = zs.into_iter().map(|x| x.into_inner()).collect();
```
各要素のUnsafeCellを外します．

```rust
    for i in 0..N {
        assert!((zs[i] - (xs[i] + ys[i])).abs() < 1e-5);
    }
    println!("ok");
}
```
正しく計算できているかをチェックします．

## 必要だった知識など

### CUDA Driver APIを用いたホストからのkernelの呼び出し
CUDAを書くときには`kernel<<<1,1>>>(arg)`のようなCUDA用の記法がよく用いられます．しかし，Rustにはこのような記法はないので，他の方法でkernelを呼び出す必要があります．ここで使うのが，CUDA Driver APIの`cuLaunchKernel`関数です．CUDA Driver APIはCUDA Runtime APIより低レイヤよりのライブラリで，見分け方としては，CUDA Runtime APIの関数が`cuda`から始まるのに対して，CUDA Driver APIの関数は`cu`から始まります．

CUDA Driver APIを用いて，PTXコードのファイルからkernelを呼び出す手順は，

1. `cuModuleLoad`関数を使ってPTXコードのファイルから`CUmodule`型の値を取得
1. `cuModuleGetFunction`関数を使って`CUmodule`型の値から，PTXコードのえファイルに含まれる指定した関数名の関数(`CUfunction`型)を取得
1. `cuLaunchKernel`関数を使って指定した関数を呼び出し

という手順です．Unix系のOSで`dlopen`/`dlsym`関数を使って共有ライブラリを動的にロードする方法に似ています．

また，CUDA Driver APIを使うには初期化の手順が必要で，今回であれば，

1. `cuInit`関数を呼び出す
2. (`cuDeviceGetCount`関数でGPUの台数を確認)
3. `cuCtxCreate_v2`関数で計算するGPU用のコンテキストを作成

という手順を取ります．

参考:[3][4][5]

### rustupでのnvptx向けtargetの追加方法
```shell-session
$ rustup target add nvptx64-nvidia-cuda
```
を実行します．

### 重いスレッドのTIMEOUT
GUIの描画に使っているGPUを使う場合，1つのスレッドの実行時間が長いとTIMEOUTします．私の場合は1スレッド/1pixelにして1スレッドで128レイ分とかの処理をし始めたころ(だったはず)にTIMEOUTするようになりました．これは1スレッド/1rayにすることで解消しました．

### nvptx向けのコードのcrateの分割
crateをまたぐと，別crateの関数のPTXコードを含んでくれないようです．解決策としては常に`#[inline(always)]`をつけておけばよさそうです．

参考: [8]

## 困っていること・今後の課題

### ホスト/カーネルのcrateの分割
crateを1つにまとめたかったのですが，build.rsを使おうとすると，ホスト側のbuild時にtargetがロックされて(?)カーネル側のbuildが固まってbuildできない現象が起きてコンパイル出来ませんでした．やり方がまずかったのかそもそもできないのかは不明．

### Shared Memory
Shared Memoryを使いたくなるのですが，いい感じのAPIが思いつかなかったので，まだ`cuda-tools`には入れてないです．
Shared Memoryの使用自体は，PTXファイルのグローバル部分に`.extern .shared .align 8 .b8 xs[];`みたいにして宣言を入れて，load/storeに`ld.shared.*`/`st.shared.*`を使えばできたはず(コードどっかいった)です．ただし，`.extern`の部分をPTXコードに含めるのに`global_asm!`を使うとなぜか2回挿入されたので，PTXコードの文字列を直接処理することになりました．

そのうち対応したいですね．

参考: [8]

### CUDA Math API
CUDA Math APIの関数を使いたくなるのですが，自動でいい感じにする方法がわかりませんでした．
手動でやる場合は，普通にCでCUDAのコードを書いて，PTXコードにコンパイルして，それっぽいところのPTXコードをもってきて`asm!`で実装すればいいです．

### デバイス側のコードでのメモリへの書き込み
メモリへの書き込みする場合(vector-addの場合はzs)に関して，はじめは`&mut`かなぁと思っていたのですが，冷静に考えるとメモリを専有してないことに気づいて，とりあえず`UnsafeCell`にしました．まだ確信を持てていないので，生ポインタを使わない場合にどれを使えばいいのかをきちんと理解していきたいです．

## 感想
- この記事の主題はなんだったんだろうになってる．ぶれずに書くのはむずかしいね．
- 冷静に考えるとバイナリじゃないのでtoolsって名前はどうなんだ．なんかいいかんじの名前にしたいね．
- まとまったらcrates.ioに上げたいね
- ちょっと低レイヤな話はわくわくしますね．なんもわからんが
- GPUが使えるとだいぶ計算の幅が広がってうれしい(妄想)．ありとあらゆる計算を実装していきたい(妄想)．
- Rustはいいぞ

## 参考文献
1. "rust-accel/accel: GPGPU Framework for Rust",  https://github.com/rust-accel/accel, (2019年12月13日 参照)
2. "Ray Tracing in One Weekend",  https://drive.google.com/drive/u/0/folders/14yayBb9XiL16lmuhbYhhvea8mKUUK77W, (2019年12月13日 参照)
3. "CUDA Driver APIでカーネル作成と実行まで - Qiita", https://qiita.com/shisoromi/items/b39742e1bc1b159280f8, (2019年12月13日 参照)
4. "いまXcodeでCUDAをはじめる多分いちばん簡単かもしれない方法 &#8211; 日曜研究室", https://peta.okechan.net/blog/archives/2083, (2019年12月13日 参照)
5. "CUDA Driver API :: CUDA Toolkit Documentation", https://docs.nvidia.com/cuda/cuda-driver-api/index.html, (2019年12月13日 参照)
6. "RustでCUDAカーネルを書く - Qiita", https://qiita.com/termoshtt/items/b98d5c46ab9c1ab1f7b6, (2019年12月13日 参照)
7. "Build Scripts - The Cargo Book", https://doc.rust-lang.org/cargo/reference/build-scripts.html, (2019年12月13日 参照)
8. "NVPTX backend metabug · Issue #38789 · rust-lang/rust",  https://github.com/rust-lang/rust/issues/38789, (2019年12月22日参照)
9. "CUDA Math API :: CUDA Toolkit Documentation",  https://docs.nvidia.com/cuda/cuda-math-api/index.html, (2019年12月22日参照)
10. "Accel: GPGPU framework for Rust - Qiita", https://qiita.com/termoshtt/items/41b4e23c4ce5e822319c, (2019年12月23日参照)