CUDAでBoostを使ってみたかったんです

よくよく考えると、最近Boostをほとんどつかってないなーと思いました。ちゃんと言うと使いどころが出てこなかったんだと思います。
最近はCUDAをさわりつつ、という感じだったのでBoost.MPIとか使えると嬉しいなと思っています。なので少しずつBoost.MPIのドキュメント読んだり書いてみたりして使い方を覚えていく予定。

Host, Deviceの扱いが面倒ですね

CUDAってホストメモリとデバイスメモリと区別されていて、わりかしその辺気をつけないと簡単にエラーになったりとか、メモリ転送はDeviceToHost?HostToDevice?とか色々めんどくさいです。
見かけ上は、どちらもただのポインタとして振る舞うので非常に分かり辛い。そこはやはりC++、型を使っていい感じに扱いやすくしてほしいものです。

と言うわけで、結構適当な気もしますが書きました。Boostがうれしい。

環境は、nvcc 4.1, gcc 4.6.1 です。
3月4日追記)どうも勘違いだったのかもしれない。今やったらビルドできなくなった。gcc 4.5はサポートしているらしい?

#include <cstdio>
#include <boost/interprocess/smart_ptr/unique_ptr.hpp>
#include <boost/type_traits/arithmetic_traits.hpp>
#include <boost/type_traits/is_same.hpp>
#include <boost/static_assert.hpp>

namespace cuda { namespace detail {
  // detail used boost::interprocess.
  namespace ipc = boost::interprocess;

  // device, host delete.
  template<class T>
  struct host_deleter {
    void operator()(T* p) const {
      std::free(p);
    }
  };
  template<class T>
  struct device_deleter {
    void operator()(T* p) const {
      cudaFree(p);
    }
  };

  // device, host memory type getting.
  template<class T>
  struct memory_type {
    // T is arithmetic type.
    BOOST_STATIC_ASSERT(boost::is_arithmetic<T>::value);
    typedef ipc::unique_ptr<T, host_deleter<T> >   host_memory;
    typedef ipc::unique_ptr<T, device_deleter<T> > device_memory;
    typedef T element_type;
  };

  // unique_ptr index access.
  // allocate Host and Device memory.
  template<class T>
  typename memory_type<T>::host_memory make_hostmemory(std::size_t size) {
    // T is arithmetic type.
    BOOST_STATIC_ASSERT(boost::is_arithmetic<T>::value);
    return typename memory_type<T>::host_memory(
        reinterpret_cast<T*>(std::malloc(sizeof(T) * size))
    );
  }
  template<class T>
  typename memory_type<T>::device_memory make_devicememory(std::size_t size) {
    // T is arithmetic type.
    BOOST_STATIC_ASSERT(boost::is_arithmetic<T>::value);
    T* devp;
    cudaMalloc(reinterpret_cast<void**>(&devp), sizeof(T) * size);
    return typename memory_type<T>::device_memory(devp);
  }

  // memory data forward.
  template<class ElementType, class Dest, class Src>
  struct get_memcpy_kind {};
  template<class ElementType>
  struct get_memcpy_kind<ElementType
    , typename memory_type<ElementType>::host_memory
    , typename memory_type<ElementType>::device_memory
  > {
    static cudaMemcpyKind apply() { return cudaMemcpyDeviceToHost; };
  };
  template<class ElementType>
  struct get_memcpy_kind<ElementType
    , typename memory_type<ElementType>::device_memory
    , typename memory_type<ElementType>::host_memory
  > {
    static cudaMemcpyKind apply() { return cudaMemcpyHostToDevice; };
  };
  template<class ElementType>
  struct get_memcpy_kind<ElementType
    , typename memory_type<ElementType>::device_memory
    , typename memory_type<ElementType>::device_memory
  > {
    static cudaMemcpyKind apply() { return cudaMemcpyDeviceToDevice; };
  };
  template<class MemoryTypeL, class MemoryTypeR>
  void forward(MemoryTypeL const& dest, MemoryTypeR const& src, std::size_t size) {
    BOOST_STATIC_ASSERT((
      boost::is_same<typename MemoryTypeR::element_type,
                     typename MemoryTypeL::element_type>::value
    ));
    typedef typename MemoryTypeR::element_type element_type;
    typedef get_memcpy_kind<element_type, MemoryTypeL, MemoryTypeR> kind;
    cudaMemcpy(dest.get(), src.get(), sizeof(element_type) * size, kind::apply());
  }

  // return reference
  template<class HostMemoryType>
  typename HostMemoryType::element_type& at(HostMemoryType const& host, int index) {
    // T is host_memory.
    BOOST_STATIC_ASSERT((
      boost::is_same<HostMemoryType, 
        typename memory_type<typename HostMemoryType::element_type>::host_memory
      >::value)
    );
    return host.get()[index];
  }
  // return value
  template<class HostMemoryType>
  typename HostMemoryType::element_type at_c(HostMemoryType const& host, int index) {
    // T is host_memory.
    BOOST_STATIC_ASSERT((
      boost::is_same<HostMemoryType, 
        typename memory_type<typename HostMemoryType::element_type>::host_memory
      >::value)
    );
    return host.get()[index];
  }

}; // namespace detail

using detail::memory_type;
using detail::make_hostmemory;
using detail::make_devicememory;
using detail::at;
using detail::at_c;
using detail::forward;

}; // namespace cuda

template<class T>
__global__
void axpy_kernel(T a, T const* x, T const* y, T* z) {
  int tid = threadIdx.x;
  z[tid] = a * x[tid] + y[tid];
}

int main(int argc, char** argv) {
  static const int N = 11;

  // error! memory_type<T>, T is arithmetic type.
  // typedef cuda::memory_type<void>::host_memory void_host_memory;

  typedef cuda::memory_type<float>::host_memory   host_memory;
  typedef cuda::memory_type<float>::device_memory device_memory;

  host_memory x = cuda::make_hostmemory<float>(N);
  host_memory y = cuda::make_hostmemory<float>(N);
  for(int i = 0 ; i < N ; ++i) {
    cuda::at(x,i) = cuda::at(y,i) = i + 1;
  }

  device_memory devx = cuda::make_devicememory<float>(N);
  device_memory devy = cuda::make_devicememory<float>(N);
  cuda::forward(devx, x, N);
  cuda::forward(devy, y, N);

  // error! devx is device memory.
  // cuda::at_c(devx, 0);

  device_memory devz = cuda::make_devicememory<float>(N);
  axpy_kernel<<<1, N>>>(1.f, devx.get(), devy.get(), devz.get());

  host_memory z = cuda::make_hostmemory<float>(N);
  cuda::forward(z, devz, N);

  std::printf("result --- \n");
  for(int i = 0 ; i < N ; ++i) {
    std::printf("%f\n", cuda::at_c(z,i));
  }

  // deleted. z, devz, devy, devx, y, x
}

memory_typeというクラステンプレートでhost_memoryとdevice_memoryの型を取り出せます。typedef templateという奴ですか。
内部が若干「んー?」という感じになっているのはtypedefされたテンプレートの型を受け取って、そのtypedefされた型からテンプレートパラメータが欲しかったんですが、できないっぽかったので泣く泣くstatic assertなどで型チェックをしています。

C++11を使いたい所ですが、nvcc 4.1でもC++11の利用は難しいようなので、boost.interprocessからunique_ptrを引っ張ってホストメモリとデバイスメモリの型にしています。
この辺はうまいこと隠して、host_memoryとdevice_memoryはget()メンバ関数ぐらいしか使えない、コピーできない、とかドキュメントに書いておけば何の心配もないですね。

APIのベタ書きは一番嫌いな私からでした。