CUDA用メモリクラス(素体)

CUDAのプログラムは、思い切りシンプルにすると、こんな感じになります。

const int N = 256;

float* h_A = (float*)malloc(sizeof(float)*N);
float* d_A;
cudaMalloc((void**)&d_A, sizeof(float)*N);

for (int i = 0; i < N; i++)
    h_A[i] = i;

cudaMemcpy(d_A, h_A, sizeof(float)*N, cudaMemcpyHostToDevice);

func<<<1, N>>>(d_A);

cudaMemcpy(h_A, d_A, sizeof(float)*N, cudaMemcpyDeviceToHost);

free(h_A);
cudaFree(d_A);

1. メモリを確保する
2. GPUにデータを送る
3. 計算する
4. GPUからデータを受け取る
5. メモリを開放する

ここで忘れてはいけないのが、最後にきちんとメモリを開放することでしょう。

開放し忘れても、プログラムを終了させたら、ちゃんと開放される………のかな？

そんなわけで、メモリ開放を忘れないためのクラスを書いてみました。

#ifndef        __CUDA_MEM_H
#define        __CUDA_MEM_H

#include        <cuda_runtime_api.h>
#include        <stdexcept>

template<typename T>
class CudaMem
{
  private:
    size_t size_;
    T* host_;
    T* device_;

  public:
    CudaMem(size_t size) : size_(size), host_(NULL), device_(NULL)
    {
        host_ =new T[size_];
        const cudaError err = cudaMalloc(reinterpret_cast<void**>&device_, 
                                         sizeof(T)*size_);
        if (err != cudaSuccess) {
            delete[] host_;
            throw std::bad_alloc(cudaGetErrorString(err));
        }
    }

    ~CudaMem()
    {
        if (device_)
            cudaFree(device_);

        delete[] host_;
    }

    size_t size() const {return size_;}

    T& operator[](size_t index) {return host_[index];}
    const T& operator[](size_t index) const {return host_[index];}

    T* get() {return host_;}
    const T* get() const {return host_;}

    T* get_gpu() {return device_;}
    const T* get_gpu() const {return device_;}

    void to_gpu()
    {
        const cudaError err = cudaMemcpy(device_, host_, sizeof(T)*size_,
                                         cudaMemcpyHostToDevice);
        if (err != cudaSuccess)
            throw std::runtime_error(cudaGetErrorString(err));
    }

    void from_gpu()
    {
        const cudaError err = cudaMemcpy(host_, device_, sizeof(T)*size_,
                                         cudaMemcpyDeviceToHost);
        if (err != cudaSuccess)
            throw std::runtime_error(cudaGetErrorString(err));
    }
    
  private:
    CudaMem(const CudaMem<T>&);
    CudaMem<T>& operator=(const CudaMem<T>&);
};

#endif

最初のサンプルコードを、このクラスを使って書き換えると、こんな感じ。

const int N = 256;
CudaMem a(N);

for (int i = 0; i < N; i++)
    a[i] = i;

a.to_gpu();

func<<<1, N>>>(a.get_gpu());

a.from_gpu();

このクラスのインスタンスはコピーできないようになってますが、必要に応じて、所有権で管理したり（auto_ptr風）、参照カウンタで管理したり（shared_ptr風）すればいいんじゃないかな？

タグ：Cuda