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 #ifndef TENSORFLOW_KERNELS_GUARDED_PHILOX_RANDOM_H_ 17 #define TENSORFLOW_KERNELS_GUARDED_PHILOX_RANDOM_H_ 18 19 #include "tensorflow/core/framework/op_kernel.h" 20 #include "tensorflow/core/lib/random/philox_random.h" 21 #include "tensorflow/core/platform/macros.h" 22 #include "tensorflow/core/platform/mutex.h" 23 #include "tensorflow/core/platform/types.h" 24 25 namespace tensorflow { 26 27 // A thread safe wrapper around a Philox generator. Example usage: 28 // 29 // GuardedRandomPhilox generator; 30 // generator.Init(context); 31 // 32 // // In thread safe code 33 // const int samples = ...; 34 // auto local_generator = generator.ReserveSamples128(samples); 35 // for (int i = 0; i < samples; i++) 36 // Array<uint32, 4> sample = local_generator(); 37 // // Use sample 38 // } 39 // 40 class GuardedPhiloxRandom { 41 public: 42 // Must call Init to finish initialization 43 GuardedPhiloxRandom() : initialized_(false) {} 44 45 // Initialize the generator from attributes "seed" and "seed2". 46 // If both seeds are unspecified, use random seeds. 47 // Must be called exactly once. 48 Status Init(OpKernelConstruction* context); 49 50 // Initialize with given seeds. 51 void Init(int64 seed, int64 seed2); 52 53 // Reserve a certain number of 128-bit samples. 54 // This function is thread safe. The returned generator is valid for the 55 // given number of samples, and can be used without a lock. 56 random::PhiloxRandom ReserveSamples128(int64 samples); 57 58 // Reserve a certain number of 32-bit samples. 59 random::PhiloxRandom ReserveSamples32(int64 samples) { 60 return ReserveSamples128((samples + 3) / 4); 61 } 62 63 // Reserve enough random samples in the generator for the given output count. 64 random::PhiloxRandom ReserveRandomOutputs(int64 output_count, 65 int multiplier) { 66 int64 conservative_sample_count = output_count * multiplier; 67 return ReserveSamples128(conservative_sample_count); 68 } 69 70 private: 71 mutex mu_; 72 random::PhiloxRandom generator_ GUARDED_BY(mu_); 73 bool initialized_; 74 75 TF_DISALLOW_COPY_AND_ASSIGN(GuardedPhiloxRandom); 76 }; 77 78 } // namespace tensorflow 79 80 #endif // TENSORFLOW_KERNELS_GUARDED_PHILOX_RANDOM_H_ 81
