HalideからGPUを利用する #C++

HalideはCPU上の並列化だけでなく、CUDA/OpenCL/MetalなどのGPUバックエンドにも対応しています。

Halide上で必要なことは、

GPUスケジューリングの使用
GPUターゲットを含んだJIT/AOTコンパイル
HostPC/GPU間のデータ転送

の3点です。
OpenCVとも連携してみましょう。
前回の投稿と同様に、グレイスケール化をGPUで行います。

なお、筆者の手持ちのGPUがRadeonである都合上、OpenCLを使用します。
CUDAの場合もほぼ一緒です。

#include <Halide.h>
#include <opencv2/opencv.hpp>

int main()
{
    // GPU環境の取得
    Halide::Target env = Halide::get_jit_target_from_environment().with_feature(Halide::Target::OpenCL);

    // 入力・出力先のMat
    cv::Mat cv_src = cv::imread("rgb.png");
    cv::Mat cv_dst(cv_src.size(), CV_8UC1);

    // データポインタを利用してのBufferの初期化
    Halide::Buffer<uint8_t> hal_src = Halide::Buffer<uint8_t>::make_interleaved(cv_src.data, cv_src.cols, cv_src.rows, 3);  
    Halide::Buffer<uint8_t> hal_dst(cv_dst.data, cv_dst.cols, cv_dst.rows);

    hal_src.set_host_dirty();       // ポインタ初期化した場合、dirtyflagが立っていないので立てる
    hal_src.copy_to_device(env);    // HostPC -> GPUの転送

    // GrayScale化のアルゴリズム
    Halide::Func GrayScale;
    Halide::Var x, y, xo, yo, xi, yi;

    Halide::Expr gray =
        0.114f * hal_src(x, y, 0) +     //B
        0.587f * hal_src(x, y, 1) +     //G
        0.299f * hal_src(x, y, 2);      //R
    GrayScale(x, y) = Halide::cast<uint8_t>(Halide::min(gray, 255));

    // GPUスケジューリングの指定
    GrayScale.gpu_tile(x, y, xo, yo, xi, yi, 64, 2);

    // GPU環境でのJITコンパイル
    GrayScale.compile_jit(env);

    GrayScale.realize(hal_dst);     // グレイスケール化の実行
    hal_dst.copy_to_host();         // GPU -> HostPCの転送

    cv::imwrite("gray.png", cv_dst);
}

多少長くなりましたが、アルゴリズム部は変化はしていないことがわかると思います。
生CUDA/OpenCLを書くよりはずっと楽です。

GPUスケジューリングの使用

GrayScale.gpu_tile(x, y, xo, yo, xi, yi, 64, 2);

ここで、GPUによるタイリングを行うことを明示する必要があります。
CUDAでいうところのxo,yoがブロック、xi,yiがスレッドに相当し、64,2がスレッド数になります。
(OpenCLの用語がわからない・・・)
gpu_プレフィックスのないtileもありますが、こちらはCPU用で、使い分ける必要があります。

GPUターゲットを含んだJITコンパイル

Halide::Target env = Halide::get_jit_target_from_environment().with_feature(Halide::Target::OpenCL);
~
GrayScale.compile_jit(env);

ここでは、プログラム動作時点での実行環境+OpenCLを用意し、それにむけてグレイスケール処理をJITコンパイルさせています。
CUDAで動作させたい方はここをTarget::CUDAにすればOK。

HostPC/GPU間のデータ転送

hal_src.set_host_dirty();       // ポインタ初期化した場合、dirtyflagが立っていないので立てる
hal_src.copy_to_device(env);    // HostPC -> GPUの転送
~
hal_dst.copy_to_host();         // GPU -> HostPCの転送

一般に、HostPCとGPUはメモリ空間が切り離されているので、GPUで計算させたいならGPU側に画像データを転送する必要があります。(iGPUやJetsonなんかはあてはまりませんが)

copy_to_〇〇は読んで字の通りです。
set_host_dirtyは「HostPC上のメモリをGPUと同期させる必要がある」フラグになります。
Bufferをポインタ初期化した場合、このフラグが立たないので、手動で操作する必要があります。
set_device_dirtyは必要ないのか？というと、今回は必要ありません。
HalideがBufferを操作する場合、自動的にフラグが設定されているようです。

OpenCLバックエンドの罠

実はOpenCLのJIT/AOTコンパイルは「OpenCLランタイムでGPUカーネルをコンパイルするプログラム」を生成します。
Func.realize()の時点で初めて実行するカーネルを作るので、初回呼び出しには大きなオーバーヘッドがあります。(筆者の環境では初回呼び出しは50msくらい追加でかかります)
大量の画像に対して同じ処理をする用途にはあまり支障はありませんが、残念な仕様です。
その点、CUDAの場合は直接GPUカーネルまで作るので、速度的にはCUDAが優勢です。

速度

Radeon RX580を使用、画像サイズは1920x1080 8bit
CodeXLを使ってGPUカーネルだけの時間を調べたところ0.0776msでした。
おおよそ100GB/sくらいでしょうか。もうちょっと出る気もしますが。