CUDAデバイスメモリ用のunique_ptr

先日、他の人の書かれたCUDAコードをいじる機会があり、その際に気になったことのメモ。

CUDAでホスト側からデバイスメモリ(GPU側メモリ)を確保する際には、cudaMallocとcudaFreeをセットで使用しますが、free操作は自動でやっていただきたいのでunique_ptrを使いたくなります。
幸い、unique_ptrはテンプレートパラメータにカスタムデリータを指定できますのでそれを使います。
VC++のdefault_delete実装をパクっ(参考にし)てやります。

template<typename Type>
struct cuda_deleter
{
	constexpr cuda_deleter() noexcept = default;

	template<typename Type2, std::enable_if_t< std::is_convertible<Type*, Type2*>::value == true, nullptr_t > = nullptr>
	cuda_deleter(const cuda_deleter<Type2>&) noexcept
	{
	}

	//指定した型からポインタ型の導出が出来ない場合でもなければ無くてもいい
	using pointer = Type*;

	void operator()(pointer ptr) {
	static_assert(0 < sizeof(Type), "can't delete an incomplete type");
		cudaFree(ptr);
	}
};

カスタムデリータのキモはオーバーロードした関数呼び出し演算子(operator()(pointer ptr))です。そこにポインタを受け取り、開放する操作を入れ込みます。
なお、nullptrのチェックはunique_ptr側のデストラクタにおいて行われたうえで呼び出されるので不要です。完全型かのチェックをsizeofが適用可能かでチェックしておきます(VC++2015の実装はこうなってるけど、不完全型に対するsizeofは失敗するだけな気がする)。
なんかごちゃごちゃしてるコピーコンストラクタっぽいものは、変換可能な型のカスタムデリータを受け入れるためのコピーコンストラクタです。
これをunique_ptrの二つ目のテンプレートパラメータに指定してやればやりたいことは終了です。

std::unique_ptr<unsigned char, cuda_deleter<unsigned char>> ptr{};

しかし、同じ型を二か所に入れこまなければいけないのは面倒極まりない、きっと3回も書くと嫌になってくる。ので、エイリアステンプレートを使いましょう。

template<typename Type>
using cuda_unique_ptr = std::unique_ptr<Type, cuda_deleter<Type>>;

cuda_unique_ptr<unsigned char> ptr{};

すっきりしました。名前はお好きにつけてください。
次なる問題は、どうやってこのunique_ptrにcudaMallocの確保メモリの所有権を移すのかという事です。当然、make_uniqueを使うことはできません。
cudaMallocは第一引数に確保した領域へのアドレスを受け取るポインタへのアドレスを指定します、つまりはポインタのポインタです。そのためにはunique_ptrの保持するポインタそのもののアドレスが必要になるわけです。
しかし、get()で取得できるのはunique_ptrの保持するポインタに格納されているアドレスであり、欲しいものはプライベートメンバのため取得できません。
一応、unique_ptrはポインタとデリータオブジェクトしかメンバに持たず、デリータはサイズ0のためEBOを利用して領域を圧縮しているはずです。つまりはunique_ptrのアドレス=保持するポインタへのアドレスとなるはずです。

cuda_unique_ptr<unsigned char> ptr{};
cudaMalloc((void **)&ptr, sizeof(unsigned char) * 10);

おそらくは正しいのですが、すべてのケースでこれが成立するかは自信が持てません。また、やはりmake_unique欲しいなあと思えて来たので、make_uniqueを作りつつ別の手段を取ることにします。
unique_ptrはコンストラクタの一つに生ポインタの所有権を受け取るものがあるので、まず別のポインタでcudaMallocを受け、それをunique_ptrに入れ込んでやりましょう。

template<typename Type>
cuda_unique_ptr<Type> make_unique(std::size_t length) {
	Type* p_tmp = nullptr;
	cudaError_t err = cudaMalloc((void **)&p_tmp, length);
	if (err != cudaError::cudaSuccess) {
		if (p_tmp != nullptr) cudaFree(p_tmp);
		return nullptr;
	}
	return cuda_unique_ptr<Type>{p_tmp};
}

cudaMallocに残り必要なものは確保する領域のサイズだけなのでそれを引数に貰います。一応、エラーが起きた時に何も確保されていないことをチェックしておき、必要なら解放の上でnullptrを返します。
最後に確保したアドレスをunique_ptrに入れ込み、リターンして終了。cudaMallocが例外を投げないのであればnoexceptになりえますがよくワカラナイ。
おまけに、二次元領域確保版を書いておきます。

template<typename Type>
cuda_unique_ptr<Type> make_unique(size_t* pitch, size_t width, std::size_t height) {
	Type* p_tmp = nullptr;
	cudaError_t err = cudaMallocPitch((void **)&p_tmp, pitch, width, height);
	if (err != cudaError::cudaSuccess)) {
		if (p_tmp != nullptr) cudaFree(p_tmp);
		return nullptr;
	}
	return cuda_unique_ptr<Type>{p_tmp};
}

基本的には同じ感じです、引数が増えたくらい。邪悪なC形式キャストには目をつむってください・・・。あと、名前空間は適宜入れておいた方がいいでしょう。

こんな感じでCUDAデバイスメモリをunique_ptrで扱うことが出来るようになりました。まあ、cudaDeviceReset()が全部なんとかしてくれる説があるらしいのですが・・・

参考文献
default_delete - cpprefjp C++日本語リファレンス
NVIDIA CUDA Library: cudaMalloc