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

はじめに

この記事は三重大学計算研 Advent Calendar 2019 6日目の記事です．
この記事ではいつも通りRustの布教をしていきたいと思います．今回はRustでGPGPUする話です．

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

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

概要

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

RustでGPGPUするライブラリとしては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]の最後の画像をレンダリングしてみます．

できた画像

ちゃんと動いてそうですね．

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

examplesのvector-addの解説

この章ではcuda-tools-examplesのvector-addを通してcuda-toolsの使い方を解説します．

crateの構成

cuda-toolsを使う場合は，ホスト側(CPU側)用のcrateとデバイス側(GPU側)用のcrateの2つを用意します．
vector-addの場合は，vector-addがホスト側用のcrate、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

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

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

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

#[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も後述しますがメモリへの書き込み時に使います．

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が使えます．


#[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

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

#[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

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

#![feature(abi_ptx)]

extern "ptx-kernel"な関数を使うので書きます．

#[macro_use]
extern crate cuda_tools;

use core::cell::UnsafeCell;
use rand::prelude::*;

extern crateと各種useです．

const KERNEL: &str = include_kernel!();

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

const N: usize = 1<<24;

今回の配列のサイズです．16777216要素です．

pub fn run() {
    let mut runtime = cuda_tools::runtime::Runtime::new(0, KERNEL).unwrap();

ホスト側のコードの本体です．
cuda_tools::runtime::Runtime::new()の第一引数にGPU番号，第二引数にPTXファイルの文字列を渡します．

    runtime.record_function_name(vector_add_kernel::vector_add, "vector_add");

関数ポインタと関数名の対応をRuntimeに覚えさせます．将来的には裏でいい感じにしたいです．

    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());
    }

入力する配列をホスト側にいったん作ります．

    let mut zs = vec![];
    for i in 0..N {
        zs.push(UnsafeCell::new(0.0));
    }

出力用の配列を作るために，いったんzsを適当に作ります．

    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()でホスト側のスライスを渡すと，デバイス側に必要なメモリ確保をした後，各要素の値をコピーします．

    let args = vector_add_kernel::Arguments {
        xs: xs_d,
        ys: ys_d,
        zs: zs_d,
    };

引数用の構造体の値を作ります．

    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}です．

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

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


    let zs: Vec<f32> = zs.into_iter().map(|x| x.into_inner()).collect();

各要素のUnsafeCellを外します．

    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を呼び出す手順は，

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

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

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

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

という手順を取ります．

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

rustupでのnvptx向けtargetの追加方法

$ 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はいいぞ

参考文献

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