Accel: GPGPU framework for Rust #Rust

これはRust Advent Calendar 2017 3日目の記事です
今回は現在開発中であるRustによるGPGPUプログラミングのためのフレームワークAccelを紹介します。

GPUを使った汎用計算の技術(GPGPU)は伝統的なHigh Performance Computing (HPC)業界だけでなく、機械学習等への応用も広がり現代では欠かせない技術です。GPUの利用には大きく分けて３つの段階があります：

高速化されたライブラリを使用する(cuBLAS, cuDNN等)
CPU用のコードに僅かな変更を加えてGPUで動くようにする(OpenACC)
GPU用のコードを設計・開発する(CUDA)

下に行くほど開発難度が増大します。
最近はOpenACCに力を入ているようで、ごく僅かな変更で大幅な高速化が期待できると宣伝しているのをよく見ます ¹。
高速化されたライブラリの使用は基本的に従来のC APIとして利用できるので、多くの言語でラッパーが提供されている一方、最大の自由度を発揮できるCUDAはC/C++/Fortranからしか利用できないのが現状です。
Accelの開発目的はこのリストにRustを入れる事になります²。

Accel: GPGPU framework for Rust

Accelの基本アイディアは以下の３つです：

RustをLLVMを経由してCUDAのアセンブラ相当であるPTXに出力する
proc-macro-attributeを使用して関数を#[kernel]で修飾するだけでCUDAカーネルに変換する
CUDA6で追加された(CUDA8で拡張された)Unified Memoryをラップしたライブラリを提供することでメモリ管理を簡単にする

これらによってC++でCUDAを書く場合よりも快適なGPGPUプログラミングを提供するためのプロジェクトです。前置きが長くなってきたのでコードを載せましょう：

#![feature(proc_macro)]

extern crate accel;
extern crate accel_derive;

use accel_derive::kernel;  // #[macro_use]は使わない
use accel::*;

#[kernel]
#[depends("accel-core" = "0.1")]  // これでCargo.tomlとextern crateに追加
pub unsafe fn add(a: *const f64, b: *const f64, c: *mut f64, n: usize) {
    let i = accel_core::index(); // threadId.x等をラップしたもの
    if (i as usize) < n {
        *c.offset(i) = *a.offset(i) + *b.offset(i);
        // この辺はまだ未完成(- -;)
    }
}

fn main() {
    let n = 8; // debug用に少なく
    // Unified Memory版Vecを用意(0-fill)
    let mut a = UVec::new(n).unwrap();
    let mut b = UVec::new(n).unwrap();
    let mut c = UVec::new(n).unwrap();

    // CPU側で初期化
    for i in 0..n {
        a[i] = i as f64;
        b[i] = 2.0 * i as f64;
    }
    println!("a = {:?}", a.as_slice());
    println!("b = {:?}", b.as_slice());

    let grid = Grid::x(64);
    let block = Block::x(64);
    // CPU -> GPUに転送
    add(grid, block, a.as_ptr(), b.as_ptr(), c.as_mut_ptr(), n);

    device::sync().unwrap(); // 実行を待つ
    // GPU -> CPUに転送
    println!("c = {:?}", c.as_slice());
}

#[kernel]で修飾されている関数addがCUDA kernelとしてコンパイルされます。このproc-macroによってコンパイル時にptx_builderというcrateが作成されて、外部プロセスとして別のコンパイルが走りPTXが生成されます。これはNVPTXが別のアーキテクチャへのクロスコンパイル相当になるため、少しややこしい設計になっています（求：改善案）。生成されたPTXファイルは読みだされてソースコードに文字列として埋め込まれて、元のRustコードがコンパイルされます。
main()内にあるadd関数に引数が増えていることに慧眼な読者諸君は気付かれていると思いますが、これはproc-macroによってコンパイル時に生成された関数に置き換わっているためです。C/CUDAでは

add<<<grid, block>>>(a, b, c, n);

のように実行するGrid/Blockを指定する必要がありましたが、この部分が関数の第１・２引数として実装されています。

メモリはVecのUnified Memory版UVecとして管理してあります。これはCPU/GPUで共通のメモリ空間を持ち、さらに転送されてない状態で読みだすと自動的に転送されます。これによりメモリ管理をいったん考えずにプログラミングできるため、非常に簡単になります。後で転送のヒントを追加していくことでメモリ転送のタイミングを工夫し、逐次的に高速化を実現することも出来ます。
UVecの実態はAccelで定義されたスマートポインタです。
as_slice()によってコピーせずに通常のRustのsliceに変換することで既存のRustのライブラリとの互換性もコストなく保てます。

構成

accel
- Accelの本体、UVec等のCUDAライブラリとPTXへのコンパイルを実行するptx_builderを含む
accel-derive
- proc-macro #[kernel]が定義してある
accel-core
- NVPTX backendのintrinsic (stdsimdに入るそうなので廃止予定)とCUDAカーネルの実装のためのユーティリティ
- 現状nvptxはno_stdでしかコンパイルできないので、これはno_stdで実装してある
cuda-sys
- CUDA Driver/Runtime APIsのラッパー。CUDA本体を配布することはできないのでシステムのCUDAを使う³

最後に

これまでにもRustで数値計算を行うための記事を書いてきましたが、今回ようやくGPGPU計算が可能になり、これでHPC業界でもRustを積極的に使っていく準備ができたと思います(｀・ω・´)

Rustによる数値計算: 線形代数編 (Rust AdC 2016 12/4)
データ並列ライブラリRayonを使ってみた (Rust AdC 2016 12/18)
RustでNumPyを拡張する
RustでCUDAカーネルを書く

皆さんもRustで良い数値計算ライフを(/・ω・)/

実際に早くなるかは知らない(´・ω・｀) ↩
当初はOpenACC相当のものを作るつもりでaccというプロジェクトだったのに、気が付いたらCUDA相当のものを作っていたのは秘密です ↩
CUDAがシステムに入ってないとリンクエラーで落ちる(´・ω・｀) ↩