1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 // Implement the Philox algorithm to generate random numbers in parallel. 17 // Salmon et al. SC 2011. Parallel random numbers: as easy as 1, 2, 3. 18 // http://www.thesalmons.org/john/random123/papers/random123sc11.pdf 19 20 #ifndef TENSORFLOW_CORE_LIB_RANDOM_PHILOX_RANDOM_H_ 21 #define TENSORFLOW_CORE_LIB_RANDOM_PHILOX_RANDOM_H_ 22 23 #include <stdint.h> 24 25 // Function qualifiers that need to work on both CPU and GPU. 26 #if defined(__CUDACC__) 27 // For nvcc. 28 #define PHILOX_DEVICE_FUNC __host__ __device__ 29 #define PHILOX_INLINE __inline__ 30 #else 31 // For non-nvcc. 32 #define PHILOX_DEVICE_FUNC 33 #define PHILOX_INLINE inline 34 #endif 35 #define PHILOX_DEVICE_INLINE PHILOX_DEVICE_FUNC PHILOX_INLINE 36 37 #include <math.h> 38 39 namespace tensorflow { 40 namespace random { 41 42 typedef int16_t int16; 43 typedef uint16_t uint16; 44 typedef int32_t int32; 45 typedef uint32_t uint32; 46 typedef int64_t int64; 47 typedef uint64_t uint64; 48 49 // A class that represents an inline array. It can be used on both CPU and GPU, 50 // and also trivially copyable between CPU and GPU. 51 // Arguments: 52 // T: the array element type; 53 // ElementCount: the fixed size of the array; 54 template <typename T, int ElementCount> 55 class Array { 56 public: Array()57 PHILOX_DEVICE_INLINE Array() { 58 for (int i = 0; i < ElementCount; ++i) { 59 data_[i] = T(0); 60 } 61 } 62 63 PHILOX_DEVICE_INLINE const T& operator[](int index) const { return data_[index]; } 64 65 PHILOX_DEVICE_INLINE T& operator[](int index) { return data_[index]; } 66 size()67 size_t size() const { return ElementCount; } 68 69 private: 70 T data_[ElementCount]; 71 }; 72 73 // A class that encapsulates all the states for a random number generator using 74 // the philox_4x32_10 algorithm. Each invocation returns a 128-bit random bits 75 // in the form of four uint32. 76 // There are multiple variants of this algorithm, we picked the 4x32_10 version 77 // that is most suited for our applications. 78 // Since this class is meant to be copied between CPU to GPU, it maintains a 79 // value semantics. 80 // 81 // For example: To use this class and populate an array of 1024 randoms on CPU 82 // with two threads, 83 // 84 // void Fill(PhiloxRandom rnd, uint32* output, int start, int limit) { 85 // assert(start % 4 == 0); 86 // assert(limit % 4 == 0); 87 // rnd.Skip(start / 4); 88 // for (int i = start; i < limit; i += 4) { 89 // auto sample = rnd(); 90 // ... copy sample[0..3] to output[i..i+3] 91 // } 92 // } 93 // 94 // PhiloxRandom rng(seed); 95 // PhiloxRandom rng_copy = rng; 96 // rng.Skip(1000/4); 97 // 98 // ... schedule Fill(rng_copy, output, 0, 512) in thread 1; 99 // ... schedule Fill(rng_copy, output, 512, 1024) in thread 2; 100 // ... wait for thread 1 & 2 to finish executing Fill(). 101 // 102 // NOTE: 103 // 1. PhiloxRandom is trivially copyable. 104 // 2. PhiloxRandom is compilable by gcc and nvcc. 105 class PhiloxRandom { 106 public: 107 using ResultType = Array<uint32, 4>; 108 using ResultElementType = uint32; 109 // The number of elements that will be returned. 110 static const int kResultElementCount = 4; 111 // Cost of generation of a single element (in cycles). 112 static const int kElementCost = 10; 113 // The type for the 64-bit key stored in the form of two 32-bit uint 114 // that are used in the diffusion process. 115 using Key = Array<uint32, 2>; 116 117 PHILOX_DEVICE_INLINE PhiloxRandom()118 PhiloxRandom() {} 119 120 PHILOX_DEVICE_INLINE PhiloxRandom(uint64 seed)121 explicit PhiloxRandom(uint64 seed) { 122 key_[0] = static_cast<uint32>(seed); 123 key_[1] = static_cast<uint32>(seed >> 32); 124 } 125 126 PHILOX_DEVICE_INLINE PhiloxRandom(uint64 seed_lo,uint64 seed_hi)127 explicit PhiloxRandom(uint64 seed_lo, uint64 seed_hi) { 128 key_[0] = static_cast<uint32>(seed_lo); 129 key_[1] = static_cast<uint32>(seed_lo >> 32); 130 counter_[2] = static_cast<uint32>(seed_hi); 131 counter_[3] = static_cast<uint32>(seed_hi >> 32); 132 } 133 134 PHILOX_DEVICE_INLINE PhiloxRandom(ResultType counter,Key key)135 PhiloxRandom(ResultType counter, Key key) : counter_(counter), key_(key) {} 136 137 // Skip the specified number of samples of 128-bits in the current stream. 138 PHILOX_DEVICE_INLINE Skip(uint64 count)139 void Skip(uint64 count) { 140 const uint32 count_lo = static_cast<uint32>(count); 141 uint32 count_hi = static_cast<uint32>(count >> 32); 142 143 counter_[0] += count_lo; 144 if (counter_[0] < count_lo) { 145 ++count_hi; 146 } 147 148 counter_[1] += count_hi; 149 if (counter_[1] < count_hi) { 150 if (++counter_[2] == 0) { 151 ++counter_[3]; 152 } 153 } 154 } 155 156 // Returns a group of four random numbers using the underlying Philox 157 // algorithm. operator()158 PHILOX_DEVICE_INLINE ResultType operator()() { 159 ResultType counter = counter_; 160 Key key = key_; 161 162 // Run the single rounds for ten times. Manually unrolling the loop 163 // for better performance. 164 counter = ComputeSingleRound(counter, key); 165 RaiseKey(&key); 166 counter = ComputeSingleRound(counter, key); 167 RaiseKey(&key); 168 counter = ComputeSingleRound(counter, key); 169 RaiseKey(&key); 170 counter = ComputeSingleRound(counter, key); 171 RaiseKey(&key); 172 counter = ComputeSingleRound(counter, key); 173 RaiseKey(&key); 174 counter = ComputeSingleRound(counter, key); 175 RaiseKey(&key); 176 counter = ComputeSingleRound(counter, key); 177 RaiseKey(&key); 178 counter = ComputeSingleRound(counter, key); 179 RaiseKey(&key); 180 counter = ComputeSingleRound(counter, key); 181 RaiseKey(&key); 182 counter = ComputeSingleRound(counter, key); 183 184 SkipOne(); 185 186 return counter; 187 } 188 189 private: 190 // We use the same constants as recommended by the original paper. 191 static const uint32 kPhiloxW32A = 0x9E3779B9; 192 static const uint32 kPhiloxW32B = 0xBB67AE85; 193 static const uint32 kPhiloxM4x32A = 0xD2511F53; 194 static const uint32 kPhiloxM4x32B = 0xCD9E8D57; 195 196 // Helper function to skip the next sample of 128-bits in the current stream. SkipOne()197 PHILOX_DEVICE_INLINE void SkipOne() { 198 if (++counter_[0] == 0) { 199 if (++counter_[1] == 0) { 200 if (++counter_[2] == 0) { 201 ++counter_[3]; 202 } 203 } 204 } 205 } 206 207 // Helper function to return the lower and higher 32-bits from two 32-bit 208 // integer multiplications. 209 PHILOX_DEVICE_INLINE MultiplyHighLow(uint32 a,uint32 b,uint32 * result_low,uint32 * result_high)210 static void MultiplyHighLow(uint32 a, uint32 b, uint32* result_low, uint32* result_high) { 211 #ifndef __CUDA_ARCH__ 212 const uint64 product = static_cast<uint64>(a) * b; 213 *result_low = static_cast<uint32>(product); 214 *result_high = static_cast<uint32>(product >> 32); 215 #else 216 *result_low = a * b; 217 *result_high = __umulhi(a, b); 218 #endif 219 } 220 221 // Helper function for a single round of the underlying Philox algorithm. ComputeSingleRound(const ResultType & counter,const Key & key)222 PHILOX_DEVICE_INLINE static ResultType ComputeSingleRound(const ResultType& counter, 223 const Key& key) { 224 uint32 lo0; 225 uint32 hi0; 226 MultiplyHighLow(kPhiloxM4x32A, counter[0], &lo0, &hi0); 227 228 uint32 lo1; 229 uint32 hi1; 230 MultiplyHighLow(kPhiloxM4x32B, counter[2], &lo1, &hi1); 231 232 ResultType result; 233 result[0] = hi1 ^ counter[1] ^ key[0]; 234 result[1] = lo1; 235 result[2] = hi0 ^ counter[3] ^ key[1]; 236 result[3] = lo0; 237 return result; 238 } 239 RaiseKey(Key * key)240 PHILOX_DEVICE_INLINE void RaiseKey(Key* key) { 241 (*key)[0] += kPhiloxW32A; 242 (*key)[1] += kPhiloxW32B; 243 } 244 245 private: 246 ResultType counter_; 247 Key key_; 248 }; 249 250 } // namespace random 251 } // namespace tensorflow 252 253 #endif // TENSORFLOW_CORE_LIB_RANDOM_PHILOX_RANDOM_H_ 254