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
|
#include <cmath>
#include "caffe2/core/context_gpu.h"
#include "caffe2/operators/elementwise_op.h"
namespace caffe2 {
template <typename T>
__global__ void SinKernel(const int N, const T* X, T* Y) {
CUDA_1D_KERNEL_LOOP(i, N) {
Y[i] = sin(X[i]);
}
}
template <typename T>
__global__ void SinGradientKernel(const int N, const T* X, const T* dY, T* dX) {
CUDA_1D_KERNEL_LOOP(i, N) {
dX[i] = dY[i] * cos(X[i]);
}
}
struct SinCUDAFunctor {
template <typename T>
inline void
operator()(const int n, const T* x, T* y, CUDAContext* device_context) {
SinKernel<T>
<<<CAFFE_GET_BLOCKS(n),
CAFFE_CUDA_NUM_THREADS,
0,
device_context->cuda_stream()>>>(n, x, y);
return;
}
};
struct SinGradientCUDAFunctor {
template <typename T>
inline void Run(
const int n,
const T* x,
const T* dy,
T* dx,
CUDAContext* device_context) {
SinGradientKernel<T>
<<<CAFFE_GET_BLOCKS(n),
CAFFE_CUDA_NUM_THREADS,
0,
device_context->cuda_stream()>>>(n, x, dy, dx);
return;
}
};
REGISTER_CUDA_OPERATOR(
Sin,
UnaryElementwiseOp<TensorTypes<float>, CUDAContext, SinCUDAFunctor>);
REGISTER_CUDA_OPERATOR(
SinGradient,
BinaryElementwiseOp<
TensorTypes<float>,
CUDAContext,
WithoutBroadcast<SinGradientCUDAFunctor>>);
} // namespace caffe2
|
