HalideとCUDAの連携

Halideについてはこちら

今回はCUDA前提の小話を書きます。
例えば、Halideでこういったスタティックライブラリを作ったとしましょう。

void GenerateCPU()
{
  Halide::ImageParam input(Halide::type_of<int>(), 1);
  Halide::Func f;
  Halide::Var x;
  f(x) = x + 1;

  f.compile_to_static_library("libsample", {input}, "sample");
}

関数sampleを、すでに確保してあるメモリ内容に適用することを考えます。

CPU側のメモリなら、Halideではこう書けます。

#include <HalideRuntime.h>
#include <HalideBuffer.h>
#include "libsample.h"

void funcCPU()
{
  int h_src[10];
  int h_dst[10];

  for(int i = 0; i < 10; ++i)
    h_src[i] = i;

  Halide::Runtime::Buffer<int> hal_src(h_src, 10);
  Halide::Runtime::Buffer<int> hal_dst(h_dst, 10);

  sample(hal_src, hal_dst);
}

Halide::Runtime::Bufferのコンストラクタにポインタをとるものがあるので、これを使えば問題ありません。
(Halide::Bufferでも同様です)

では、すでに確保してあるGPUメモリを使うにはどうしたらよいのでしょうか。
上述のポインタをとるコンストラクタは、CPUメモリ用のものであり、そのままGPUメモリを指すポインタを渡しても動作しません。
正しくは、以下のように書きます。

void GenerateGPU()
{
  Halide::ImageParam input(Halide::type_of<int>(), 1);
  Halide::Func f;
  Halide::Var x;
  f(x) = x + 1;
  
  Var xo, xi;
  f.gpu_tile(x, xo, xi, 32, Halide::TailStrategy::GuardWithIf);
  f.compile_to_static_library("libsampleGPU", {input}, "sampleGPU", Halide::get_host_target().with_feature(Halide::Target::CUDA));
}

#include <cuda_runtime.h>
#include <HalideRuntime.h>
#include <HalideBuffer.h>
#include <HalideRuntimeCuda.h>
#include "libsampleGPU.h"

void funcGPU()
{
  int h_src[10];
  int h_dst[10];

  for(int i = 0; i < 10; ++i)
    h_src[i] = i;

  int* d_src;
  int* d_dst;
  cudaMalloc(&d_src, sizeof(int)*10);
  cudaMalloc(&d_dst, sizeof(int)*10);
  cudaMemcpy(d_src, h_src, sizeof(int)*10, cudaMemcpyHostToDevice);

  Halide::Runtime::Buffer<int> hal_src(nullptr, 10);
  hal_src.device_wrap_native(halide_cuda_device_interface(), (uint64_t)d_src);
  Halide::Runtime::Buffer<int> hal_dst(nullptr, 10);
  hal_dst.device_wrap_native(halide_cuda_device_interface(), (uint64_t)d_dst);

  sampleGPU(hal_src, hal_dst);

  cudaMemcpy(h_dst, d_dst, sizeof(int)*10, cudaMemcpyDeviceToHost);
}

ちょっと手間になりますが、device_wrap_nativeを用いるとうまくできます。
スタティックライブラリ生成時にgpu_tile指定をしておくことも忘れずに。
ミソはコンストラクタにnullptrを渡してやることでしょうか。
halide_cuda_device_interfaceに当たる部分は、ここに該当するヘッダがいろいろ含まれているので、CUDAだけではなくOpenCLやMetalにも対応可能だと思われます。

メモリコピーは思ったよりも速くありません。高速なプログラムを書くためにも、コピーは可能な限り避けましょう。