-
Notifications
You must be signed in to change notification settings - Fork 0
/
tutorial.cu
42 lines (33 loc) · 1.11 KB
/
tutorial.cu
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
#define EIGEN_USE_GPU
#include <cuda.h>
#include <stdio.h>
#include "tensorflow/core/util/cuda_device_functions.h"
#include "tutorial.h"
using GPUDevice = Eigen::GpuDevice;
template <typename dtype> __global__ void AddKernel(const dtype* a, const dtype* b, dtype* c, int N){
for(int index : tensorflow::CudaGridRangeX(N))
{
c[index] = a[index] + b[index];
}
}
template <typename dtype>
struct launchAddKernel<GPUDevice, dtype> {
void operator()(const GPUDevice& d, const dtype* a, const dtype* b, dtype* c, int N) {
const int kThreadsPerBlock = 1024;
AddKernel<dtype><<<(N + kThreadsPerBlock - 1) / kThreadsPerBlock,
kThreadsPerBlock, 0, d.stream()>>>(
a, b, c, N);
cudaError_t cudaerr = cudaDeviceSynchronize();
if (cudaerr != cudaSuccess)
printf("kernel launch failed with error \"%s\".\n",
cudaGetErrorString(cudaerr));
}
};
//forward declaration for all the types needed
typedef Eigen::GpuDevice GPUDevice;
#define ADD_KERNEL_TYPE(type) \
template struct launchAddKernel<GPUDevice, type>; \
ADD_KERNEL_TYPE(int);
ADD_KERNEL_TYPE(float);
ADD_KERNEL_TYPE(double);
#undef ADD_KERNEL_TYPE