RustからCUDAのカーネルを呼び出してみた

#はじめに
RustからCUDAのホスト側の処理を書く方法を、ぐぐっていたのですが、なかなか出てきませんでした。
それもそのはずで、単にFFIでCUDAのライブラリを呼び出すだけなので、RustでCUDAのホスト側になる特別な手順というものは無いのです。

とはいえ、FFIでCUDAカーネルを呼び出す方法も、なかなか分かりづらいものです。C/C++で書くときにnvccがうまいことしてくれるfunc<<<grid, block>>>();の表記は、他では使えないので、そこを変えてやる必要があります。
また、bindgenを使う練習になったり、build.rsを書く練習になったりと、意外といい練習になりました。

なお、この記事は、CUDA Runtime APIを使って、Rustで、CUDAの ホスト側の処理 を書くものです。
この記事を読んでも、デバイス側の処理(CUDAカーネル)は書けるようになりませんので、ご注意ください。
(RustでCUDAカーネルを書く方法については http://qiita.com/termoshtt/items/b98d5c46ab9c1ab1f7b6 が詳しいです)

CUDAのカーネル自体は、Rustではなく、普通にC/C++で書いてある、なんたら.cuファイルが用意されているものとして、今回はcargoでそれをコンパイルできるようにしてみました。
(コンパイル済みのカーネルでも、extern "C"をつけてコンパイルしてあって、Rustからリンクできるようになっていれば、呼び出せるはずですが、この記事では触れません)

また、build.rsを、自分の環境にベッタリな感じで書いてますので、Arch Linux以外の環境では、結構がんばっていじる感じになると思います。(他の部分は動くはずですが)

#bindgenでcuda_runtime.hをRust用に変換しよう
bindgenは、build.rsの中でも使えますが、今回はコマンドラインツールとして使うことにしました。
https://rust-lang-nursery.github.io/rust-bindgen/command-line-usage.html

1. cargo install bindgen でインストール
2. bindgen /opt/cuda/include/cuda_runtime.h -o src/cuda_runtime.rs のような感じで、Rustに変換
3. コンパイル時に警告が鬱陶しいので、できあがったものの1行目に #![allow(dead_code, non_camel_case_types, non_snake_case, non_upper_case_globals)]を付け加えた方がいい

#適当なラッパーを用意しよう
cuda_runtimeが使えるようになったので、ぶっちゃけ、これでもう書き始められるのですが。
かっこいいエラー処理とか、そういうのを使うために、Rustを使いたいわけで、unsafeでCの関数を呼び出して、返される値がcudaSuccessかどうかでエラーかどうかを判定して、みたいなことをしたいなら、最初からCで書きますよね。
ということで、適当にラッパーを用意しましょう。

このへん、どこまでやるかは趣味の世界なので。ここではあまり触れません。
せめて、結果はResult型で返して欲しいですね。

#カーネルを呼び出そう
カーネルを呼び出す前に、普通はcudaMallocやcudaMemcpyなどの事前準備も必要ですが、この記事を読んでいる方ならきっと、かっこよく作ったラッパーでスマートにRustらしく準備されていることと思いますので、ここでは触れません。
さっそく呼び出してみましょう。

ここでは、カーネルとして、CUDAのサンプルのvectorAdd.cuから、カーネル部分のみを切り出して、また、extern "C" { ... }で囲ったものを使います。

extern "C" {
__global__ void
vectorAdd(const float *A, const float *B, float *C, int numElements)
{
    // 略
}
}

また、Rust内でも、どっかでexternしておいてください。

extern "C" {
    fn vectorAdd(a: *const f32, b: *const f32, c: *mut f32, n: c_int) -> c_void;
}

##1. 非推奨のAPIを使う方法
cudaConfigureCallを使う方法がありますが、CUDA 7.0からdeprecated扱いになったようです。

cudaConfigureCallしてから呼び出し

    let threads_per_block = 256usize;
    let blockdim = cuda_ffi::usize_to_dim3(threads_per_block);
    let griddim = cuda_ffi::usize_to_dim3((n + threads_per_block - 1) / threads_per_block);
    let sharedmem = 0usize;
    let n_int = n as c_int;

    cuda_ffi::configure_call(griddim, blockdim, sharedmem).unwrap();
    unsafe {
            vectorAdd(d_a, d_b, d_c, n_int);
    }
    cuda_ffi::last_error().unwrap();

カーネルの呼び出し前にわざわざcudaConfigureCallを呼ばないといけないなど、たしかに、ちょっと嫌な感じはします。けど、これ、deprecatedじゃない方法よりもだいぶマシなんですよねぇ。。。
ちなみに、ラッパー側で、関係ありそうなところはこんな感じ。

cuda_ffi.rs

pub struct Error {
    raw: cudaError_t,
}
pub type Result<T> = result::Result<T, Error>;


pub fn usize_to_dim3(x: usize) -> dim3 {
    dim3 {
        x: x as raw::c_uint,
        y: 1,
        z: 1,
    }
}

pub fn configure_call(grid_dim: dim3, block_dim: dim3, shared_mem: usize) -> Result<()> {
    let cuda_error =
        unsafe { cuda_runtime::cudaConfigureCall(grid_dim, block_dim, shared_mem, null_mut()) };
    if cuda_error == cuda_runtime::cudaError::cudaSuccess {
        Ok(())
    } else {
        Err(Error { raw: cuda_error })
    }
}

pub fn last_error() -> Result<()> {
    let cuda_error = unsafe { cuda_runtime::cudaGetLastError() };
    if cuda_error == cuda_runtime::cudaError::cudaSuccess {
        Ok(())
    } else {
        Err(Error { raw: cuda_error })
    }
}

##2. cudaLaunchKernelを使う方法
cudaLaunchKernelを使えば、cudaConfigureCallを呼ばずに済みます。
けれど……この引数は……
__host__ ​cudaError_t cudaLaunchKernel ( const void* func, dim3 gridDim, dim3 blockDim, void** args, size_t sharedMem, cudaStream_t stream )
嫌な予感しかしません……　だって、void**で引数リストやて。ほんまかいな。

素直に書くと、やっぱり、こうなりました。

cudaLaunchKernelから呼び出し

    cuda_ffi::launch(vectorAdd as *const c_void,
                     griddim,
                     blockdim,
                     &mut [&mut d_a as *mut *mut f32 as *mut c_void,
                           &mut d_b as *mut *mut f32 as *mut c_void,
                           &mut d_c as *mut *mut f32 as *mut c_void,
                           &n_int as *const c_int as *mut c_int as *mut c_void],
                     sharedmem)
        .unwrap();

ちなみに。streamはnullでいい、みたいなことがドキュメントに書いてあったので、ラッパー側でnullにしてます。

cuda_ffi.rs

pub fn launch(func: *const raw::c_void,
              grid_dim: dim3,
              block_dim: dim3,
              args: &mut [*mut raw::c_void],
              shared_mem: usize)
              -> Result<()> {
    let cuda_error = unsafe {
        cuda_runtime::cudaLaunchKernel(func,
                                       grid_dim,
                                       block_dim,
                                       args.as_mut_ptr(),
                                       shared_mem,
                                       null_mut())
    };
    if cuda_error == cuda_runtime::cudaError::cudaSuccess {
        Ok(())
    } else {
        Err(Error { raw: cuda_error })
    }
}

#build.rsを用意する
build.rsについては、
http://doc.crates.io/build-script.html
に詳しく載っていますが、早い話が、ビルドするときに実行されるコードで、うまく書けば、外部コマンドを呼び出したり、リンクするライブラリを増やしたりできます。

Cargo.tomlの[package]セクションに、以下の一行を追加します。
build = "build.rs"

コンパイルします。

nvccを動かす

use std::env;
use std::path::Path;
use std::process::Command;

fn main() {
    let out_dir = env::var("OUT_DIR").unwrap();

    Command::new("nvcc")
        .args(&["-c", "-arch=sm_20", "src/vectorAdd.cu", "-Xcompiler", "-fPIC", "-o"])
        .arg(&format!("{}/vectorAdd.o", out_dir))
        .status()
        .unwrap();

これだけじゃダメらしくて、staticライブラリに変換します。

arを動かす

    Command::new("ar")
        .args(&["crus", "libvectorAdd.a", "vectorAdd.o"])
        .current_dir(&Path::new(&out_dir))
        .status()
        .unwrap();

そして、おまじないをstdoutに出力します。(ハードコードされてるパスとかは、環境によって違うかもしれないです)

おまじない

    println!("cargo:rustc-link-search=native=/opt/cuda/lib64");
    println!("cargo:rustc-link-lib=cudart");
    println!("cargo:rustc-link-search=native={}", out_dir);
    println!("cargo:rustc-link-lib=static=vectorAdd");
}

これで、cargo buildで、GPUのコードごとコンパイルできるはずです。

#おわりに
今回使ったコード全体は
https://github.com/gyu-don/cuda_from_rust
においています。