MQForceGPU/cuda_util.h at master · FloydZ/MQForceGPU · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
#ifndef CUDA_UTIL_H
#define CUDA_UTIL_H

#include <cuda.h>

#ifndef CUDA_ASSERT
#include <stdio.h>
#define CUDA_ASSERT(ans)                                                       \
  {                                                                            \
    gpuAssert((ans), __FILE__, __LINE__);                                      \
  }
inline void
gpuAssert(cudaError_t code, const char* file, int line, bool abort = true)
{
  if (code != cudaSuccess) {
    fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file,
            line);
    if (abort)
      exit(code);
  }
}
#endif

template <class Type> class cudaData
{
public:
  cudaData(size_t len_host, size_t len_dev = 0);
  ~cudaData();

  size_t size_host();
  size_t size_dev();
  void clear();
  void write(size_t off_src = 0, size_t size = 0, size_t off_des = 0);
  void read(size_t off_src = 0, size_t size = 0, size_t off_des = 0);

  Type* host;
  Type* dev;

private:
  size_t sz_host;
  size_t sz_dev;
};

template <class Type> cudaData<Type>::cudaData(size_t len_host, size_t len_dev)
{
  if (len_dev == 0)
    len_dev = len_host;

  sz_host = len_host * sizeof(Type);
  sz_dev = len_dev * sizeof(Type);

  host = (Type*)malloc(sz_host);

  CUDA_ASSERT(cudaMalloc((void**)&dev, sz_dev));
}

template <class Type> cudaData<Type>::~cudaData()
{
  free(host);
  cudaFree(dev);
}

template <class Type> size_t cudaData<Type>::size_host() { return sz_host; }

template <class Type> size_t cudaData<Type>::size_dev() { return sz_dev; }

template <class Type> void cudaData<Type>::clear() { memset(host, 0, sz_host); }

template <class Type>
void cudaData<Type>::write(size_t off_src, size_t size, size_t off_des)
{
  if (size == 0)
    size = (sz_host <= sz_dev) ? sz_host : sz_dev;

  CUDA_ASSERT(
      cudaMemcpy(&dev[off_src], &host[off_des], size, cudaMemcpyHostToDevice));
}

template <class Type>
void cudaData<Type>::read(size_t off_src, size_t size, size_t off_des)
{
  if (size == 0)
    size = (sz_host <= sz_dev) ? sz_host : sz_dev;

  CUDA_ASSERT(
      cudaMemcpy(&host[off_src], &dev[off_des], size, cudaMemcpyDeviceToHost));
}

#endif