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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
|
#if GOOGLE_CUDA
#define EIGEN_USE_GPU
#include "tensorflow/core/kernels/random_op.h"
#include <stdio.h>
#include <assert.h>
#include "tensorflow/core/lib/random/philox_random.h"
#include "tensorflow/core/lib/random/random_distributions.h"
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
namespace tensorflow {
class OpKernelContext;
namespace functor {
typedef Eigen::GpuDevice GPUDevice;
template <class Distribution, bool VariableSamplesPerOutput>
struct FillPhiloxRandomKernel;
// A cuda kernel to fill the data with random numbers from the specified
// distribution. Each output takes a fixed number of samples.
template <class Distribution>
struct FillPhiloxRandomKernel<Distribution, false> {
typedef typename Distribution::ResultElementType T;
PHILOX_DEVICE_FUNC void Run(random::PhiloxRandom gen, T* data, int64 size) {
Distribution dist;
const int kGroupSize = Distribution::kResultElementCount;
const int32 thread_id = blockIdx.x * blockDim.x + threadIdx.x;
const int32 total_thread_count = gridDim.x * blockDim.x;
int32 offset = thread_id * kGroupSize;
gen.Skip(thread_id);
while (offset < size) {
typename Distribution::ResultType samples = dist(&gen);
for (int i = 0; i < kGroupSize; ++i) {
if (offset >= size) {
return;
}
data[offset] = samples[i];
++offset;
}
offset += (total_thread_count - 1) * kGroupSize;
gen.Skip(total_thread_count - 1);
}
}
};
// A cuda kernel to fill the data with random numbers from the specified
// distribution. Each output takes a variable number of samples.
template <class Distribution>
struct FillPhiloxRandomKernel<Distribution, true> {
typedef typename Distribution::ResultElementType T;
PHILOX_DEVICE_FUNC void Run(const random::PhiloxRandom& base_gen, T* data,
int64 size) {
using random::PhiloxRandom;
using random::SingleSampleAdapter;
const int kReservedSamplesPerOutput = 256;
const int kGroupSize = Distribution::kResultElementCount;
const int kGeneratorSkipPerOutputGroup = kGroupSize *
kReservedSamplesPerOutput /
PhiloxRandom::kResultElementCount;
const int32 thread_id = blockIdx.x * blockDim.x + threadIdx.x;
const int32 total_thread_count = gridDim.x * blockDim.x;
int64 group_index = thread_id;
int64 offset = group_index * kGroupSize;
Distribution dist;
while (offset < size) {
// Since each output takes a variable number of samples, we need to
// realign the generator to the beginning for the current output group
PhiloxRandom gen = base_gen;
gen.Skip(group_index * kGeneratorSkipPerOutputGroup);
SingleSampleAdapter<PhiloxRandom> single_samples(&gen);
typename Distribution::ResultType samples = dist(&single_samples);
for (int i = 0; i < kGroupSize; ++i) {
if (offset >= size) {
return;
}
data[offset] = samples[i];
++offset;
}
offset += (total_thread_count - 1) * kGroupSize;
group_index += total_thread_count;
}
}
};
// A simple launch pad to call the correct function templates to fill the data
template <class Distribution>
__global__ void __launch_bounds__(1024)
FillPhiloxRandomKernelLaunch(random::PhiloxRandom base_gen,
typename Distribution::ResultElementType* data,
int64 size) {
FillPhiloxRandomKernel<Distribution,
Distribution::kVariableSamplesPerOutput>()
.Run(base_gen, data, size);
}
// Partial specialization for GPU
template <class Distribution>
struct FillPhiloxRandom<GPUDevice, Distribution> {
typedef typename Distribution::ResultElementType T;
typedef GPUDevice Device;
void operator()(OpKernelContext*, const Device& d, random::PhiloxRandom gen,
T* data, int64 size) {
const int32 block_size = d.maxCudaThreadsPerBlock();
const int32 num_blocks =
(d.getNumCudaMultiProcessors() * d.maxCudaThreadsPerMultiProcessor()) /
block_size;
FillPhiloxRandomKernelLaunch<
Distribution><<<num_blocks, block_size, 0, d.stream()>>>(gen, data,
size);
}
};
// Explicit instantiation of the GPU distributions functors
// clang-format off
// NVCC cannot handle ">>" properly
template struct FillPhiloxRandom<
GPUDevice, random::UniformDistribution<random::PhiloxRandom, float> >;
template struct FillPhiloxRandom<
GPUDevice, random::UniformDistribution<random::PhiloxRandom, double> >;
template struct FillPhiloxRandom<
GPUDevice, random::NormalDistribution<random::PhiloxRandom, float> >;
template struct FillPhiloxRandom<
GPUDevice, random::NormalDistribution<random::PhiloxRandom, double> >;
template struct FillPhiloxRandom<
GPUDevice, random::TruncatedNormalDistribution<
random::SingleSampleAdapter<random::PhiloxRandom>, float> >;
template struct FillPhiloxRandom<
GPUDevice, random::TruncatedNormalDistribution<
random::SingleSampleAdapter<random::PhiloxRandom>, double> >;
// clang-format on
} // namespace functor
} // namespace tensorflow
#endif // GOOGLE_CUDA
|
