Zenn
👏

Rustで書いたCUDAカーネルで画像処理してみる

2022/12/24に公開
Rust
CUDA
tech

概要

C++の後継と言われているRustですが、現状でまだC++の方が勝っているのではないかという点の一つが、GPUなどのハードウェアアクセラレーションまわりかなと思っています。
ハードウェアアクセラレーションの中でも特にCUDAは近年、画像処理やディープラーニングの発展によって使用される機会が増えてきています。
現状では、画像処理で使われるOpenCVやディープラーニング系のフレームワークもCUDAによる処理を記述するためにC++を用いており、それをPython等にバインディングしていますが、こういったところにもRustが入ってきてより使いやすいものになると、いよいよC++の立場も危うくなってきそうです。

この記事ではRustでもCUDAをうまく扱っていくために、実用的な用途として考えられそうな画像処理を行うコードを書いてみようと思います。

Rust-CUDA

今回RustでCUDAを使うために使用するパッケージはこちらです。
https://github.com/Rust-GPU/Rust-CUDA

RustでCUDAを扱うために必要なクレートを集めたエコシステムのようなパッケージになっていて、こちらのパッケージ内のクレートを使用することで、CUDAのカーネルの部分も含めてRustで記述することができます。
上記パッケージではベクトルの足し算やレイトレーシングのサンプルが置いてあるのですが、画像処理のサンプルは用意されていませんでした。

今回は画像処理の中でも基本的な処理である、グレースケール変換とガウシアンブラーのコードをRust-CUDAを使って書いてみました。
今回、私が作成したコードを以下のリポジトリに置いています。
https://github.com/neka-nat/cuimage

パッケージ構成

https://crates.io/ には登録できていないのですが、ライブラリっぽい使い方ができるような構成にしています。

Rust-CUDAに習って、ホスト用とデバイス用の2つのクレートを持つワークスペースになっており、以下のような形でcpuフォルダにホスト側のプログラムを、gpuフォルダにデバイス側のプログラムを記述するようになっています。

cuimage
├── Cargo.lock
├── Cargo.toml
├── Dockerfile
├── README.md
├── cpu
│   ├── Cargo.toml
│   ├── build.rs
│   ├── examples
│   │   ├── gaussian_blur.rs
│   │   ├── gray.rs
│   │   └── lenna.png
│   └── src
│       ├── buffer.rs
│       ├── context.rs
│       ├── gaussian_blur.rs
│       ├── gray.rs
│       └── lib.rs
├── gpu
│   ├── Cargo.toml
│   └── src
│       ├── gaussian_blur.rs
│       ├── gray.rs
│       └── lib.rs
├── resources
└── rust-toolchain

デバイスのコード

以下にグレースケール変換を行うカーネルのコードを示します。
入力はそれぞれの値がuint8のRGB(もしくはRGBA)、出力はuint8のグレースケール画像を想定しています。
デバイス側のRustコードはコンパイルされてPTXが出力されます。

#![cfg_attr(
    target_os = "cuda",
    no_std,
    feature(register_attr),
    register_attr(nvvm_internal)
)]

use cuda_std::prelude::*;

#[kernel]
#[allow(improper_ctypes_definitions)]
pub unsafe fn gray(src: &[u8], dst: *mut u8, num_of_channels: u32) {
    let idx = thread::index_1d() as usize;
    if idx < src.len() {
        let num_of_channels = num_of_channels as usize;
        let r = src[num_of_channels * idx] as f32;
        let g = src[num_of_channels * idx + 1] as f32;
        let b = src[num_of_channels * idx + 2] as f32;
        let gray = 0.299 * r + 0.587 * g + 0.114 * b;
        let elem = &mut *dst.add(idx);
        *elem = gray as u8;
    }
}

ホストのコード

画像として扱いやすくするために、以下のようなデバイス側のメモリで保存されたImageクラスを作成しました。

pub struct CuImage<P: Pixel>
where
    <P as Pixel>::Subpixel: DeviceCopy,
{
    pub width: u32,
    pub height: u32,
    pub num_of_channels: u32,
    pub data: DeviceBuffer<<P as Pixel>::Subpixel>,
    _phantom: PhantomData<P>,
}

ホスト側で作成されたImageクラスを使って初期化することができるようになっています。

// ホストで画像読み込み
let img = image::open("lenna.png").unwrap().into_rgba8();
// デバイスに転送
let src_gpu = CuImage::<Rgba<u8>>::from_host_image(&img)?;

次にホストからのカーネルの呼び出し部分です。
画像処理のアルゴリズム毎に構造体を用意し、カーネルを実行する関数に先程のデバイスのImageクラスを入出力として与えられるようにしています。

pub struct Gray<'a, P: Pixel>
where
    <P as Pixel>::Subpixel: DeviceCopy,
{
    context: &'a CuContext,
    _phantom: PhantomData<P>,
}

impl<'a, P: Pixel> Gray<'a, P>
where
    <P as Pixel>::Subpixel: DeviceCopy,
{
    pub fn new(context: &'a CuContext) -> Self {
        Gray {
            context,
            _phantom: PhantomData,
        }
    }
    pub fn run(
        &self,
        src_img: &CuImage<P>,
        dst_img: &mut CuImage<Luma<u8>>,
    ) -> Result<(), Box<dyn Error>> {
        let kernel = self.context.module.get_function("gray").unwrap();
        let (_, block_size) = kernel.suggested_launch_configuration(0, 0.into())?;
        let grid_size = (src_img.width * src_img.height + block_size - 1) / block_size;
        let stream = &self.context.stream;
        unsafe {
            launch!(
                kernel<<<grid_size, block_size, 0, stream>>>(
                    src_img.data.as_device_ptr(),
                    src_img.width * src_img.height,
                    dst_img.data.as_device_ptr(),
                    src_img.num_of_channels,
                )
            )?;
        }
        stream.synchronize()?;
        Ok(())
    }
}

mainでの呼び出し方

ここまでの形で処理をまとめて、実際にmain関数で使用する場合のサンプルは以下になります。

use cuimage::*;
use image::{Luma, Rgba};
use std::error::Error;

fn main() -> Result<(), Box<dyn Error>> {
    let img = image::open("cpu/examples/lenna.png").unwrap().into_rgba8();

    let context = CuContext::default();
    let src_gpu = CuImage::<Rgba<u8>>::from_host_image(&img)?;
    let mut dst_gpu = CuImage::<Luma<u8>>::new(img.width(), img.height());

    let gray = Gray::new(&context);
    gray.run(&src_gpu, &mut dst_gpu)?;

    let out = dst_gpu.to_host_image()?;
    out.save("cpu/examples/lenna_gray.png").unwrap();

    Ok(())
}

比較的シンプルにCUDAを画像処理で使用できる形にできたのではないかと思っています。

ビルド

ビルドの環境ですが、ホストの環境にうまく作れなかったので、Rust-CUDAで用意されているDockerfileを使用してビルドしています。

docker build -t rust-cuda .
docker run -it --gpus all -v $PWD:/root/rust-cuda --entrypoint /bin/bash rust-cuda
cd /root/rust-cuda
cargo run --release --example gray

結果

最後に今回作成したグレースケール変換とガウシアンブラーの結果です。
本来であればパフォーマンス計測も行いたかったのですが、なんか動いて満足してしまったので、そのあたりはまた後日行いたいと思います。

元画像

グレースケール

ガウシアンブラー

まとめ

Rustを使ってCUDAで画像処理を行うプログラムを書いてみました。
ホスト側では画像のピクセルをジェネリクスで扱うように書いていますが、実際はu8のピクセルでしかうまく動かないのが現状ちょっとイマイチです。
C++には生のCUDAを扱う以外にもcubやThrustといったより高次のインターフェースが用意されたライブラリがあったりして、そういった点ではまだCUDAを使うのはC++のほうが扱いやすいかなという印象でした。
あとカーネル部分はno_stdなのでデバッグしづらいです。(この記事とか参考にすれば途中のprintとかもできるのかも)

今後もまだまだRust＋CUDAまわりも整備されてくるかと思いますので、引き続きこのあたり注目していきたいと思います。

Discussion