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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
|
#include "caffe2/operators/tile_op.h"
#include <array>
#include "caffe2/core/context_gpu.h"
#include "caffe2/utils/math.h"
namespace caffe2 {
namespace {
template <typename T>
__global__ void TileCopyCUDAKernel(
const int total_size,
const int inner_size,
const int tiles,
const T* X,
T* Y) {
const int x = blockIdx.x * CAFFE_CUDA_NUM_THREADS + threadIdx.x;
if (x < total_size) {
const int r = x / inner_size / tiles;
const int c = x % inner_size;
#if __CUDA_ARCH__ >= 350 || defined(USE_ROCM)
Y[x] = __ldg(X + r * inner_size + c);
#else
Y[x] = X[r * inner_size + c];
#endif
}
}
} // namespace
template <>
template <typename T>
bool TileOp<CUDAContext>::DoTile(
const int outer_size,
const int inner_size,
const T* X,
T* Y) {
const std::int64_t total_size = static_cast<std::int64_t>(outer_size) *
static_cast<std::int64_t>(tiles_) * static_cast<std::int64_t>(inner_size);
const int M = math::DivUp<std::int64_t>(total_size, CAFFE_CUDA_NUM_THREADS);
TileCopyCUDAKernel<T>
<<<M, CAFFE_CUDA_NUM_THREADS, 0, context_.cuda_stream()>>>(
total_size, inner_size, tiles_, X, Y);
C10_CUDA_KERNEL_LAUNCH_CHECK();
return true;
}
template <>
template <typename T>
bool TileGradientOp<CUDAContext>::DoTileGradient(
const int outer_size,
const int inner_size,
const T* dY,
T* dX) {
const std::array<int, 3> dY_dims = {outer_size, tiles_, inner_size};
const std::array<int, 3> dX_dims = {outer_size, 1, inner_size};
math::ReduceSum<T, CUDAContext>(
3, dY_dims.data(), dX_dims.data(), T(1), dY, dX, &context_);
return true;
}
template <>
template <>
bool TileGradientOp<CUDAContext>::DoTileGradient<float>(
const int outer_size,
const int inner_size,
const float* dY,
float* dX) {
if (inner_size == 1) {
const std::array<int, 2> dY_dims = {outer_size, tiles_};
const std::array<int, 2> dX_dims = {outer_size, 1};
math::ReduceSum<float, CUDAContext>(
2, dY_dims.data(), dX_dims.data(), 1.0f, dY, dX, &context_);
} else {
ReinitializeTensor(&ones_, tiles_, at::dtype<float>().device(CUDA));
math::Set<float, CUDAContext>(
tiles_, 1.0f, ones_.template mutable_data<float>(), &context_);
math::GemmStridedBatched<float, CUDAContext>(
CblasTrans,
CblasNoTrans,
outer_size,
inner_size,
1,
tiles_,
1.0f,
dY,
tiles_ * inner_size,
ones_.template data<float>(),
0,
0.0f,
dX,
inner_size,
&context_);
}
return true;
}
REGISTER_CUDA_OPERATOR(Tile, TileOp<CUDAContext>);
REGISTER_CUDA_OPERATOR(TileGradient, TileGradientOp<CUDAContext>);
} // namespace caffe2
|
