1 /* Copyright 2021 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_CORE_LIB_RANDOM_RANDOM_DISTRIBUTIONS_UTILS_H_
17 #define TENSORFLOW_CORE_LIB_RANDOM_RANDOM_DISTRIBUTIONS_UTILS_H_
18
19 #include <string.h>
20
21 #include <cstdint>
22
23 #include "tensorflow/core/lib/random/philox_random.h"
24
25 #ifndef M_PI
26 #define M_PI (3.14159265358979323846)
27 #endif
28
29 namespace tensorflow {
30 namespace random {
31
32 // Helper function to convert an 32-bit integer to a float between [0..1).
Uint32ToFloat(uint32_t x)33 PHILOX_DEVICE_INLINE float Uint32ToFloat(uint32_t x) {
34 // IEEE754 floats are formatted as follows (MSB first):
35 // sign(1) exponent(8) mantissa(23)
36 // Conceptually construct the following:
37 // sign == 0
38 // exponent == 127 -- an excess 127 representation of a zero exponent
39 // mantissa == 23 random bits
40 const uint32_t man = x & 0x7fffffu; // 23 bit mantissa
41 const uint32_t exp = static_cast<uint32_t>(127);
42 const uint32_t val = (exp << 23) | man;
43
44 // Assumes that endian-ness is same for float and uint32_t.
45 float result;
46 memcpy(&result, &val, sizeof(val));
47 return result - 1.0f;
48 }
49
50 // Helper function to convert two 32-bit integers to a double between [0..1).
Uint64ToDouble(uint32_t x0,uint32_t x1)51 PHILOX_DEVICE_INLINE double Uint64ToDouble(uint32_t x0, uint32_t x1) {
52 // IEEE754 doubles are formatted as follows (MSB first):
53 // sign(1) exponent(11) mantissa(52)
54 // Conceptually construct the following:
55 // sign == 0
56 // exponent == 1023 -- an excess 1023 representation of a zero exponent
57 // mantissa == 52 random bits
58 const uint32_t mhi = x0 & 0xfffffu; // upper 20 bits of mantissa
59 const uint32_t mlo = x1; // lower 32 bits of mantissa
60 const uint64_t man = (static_cast<uint64_t>(mhi) << 32) | mlo; // mantissa
61 const uint64_t exp = static_cast<uint64_t>(1023);
62 const uint64_t val = (exp << 52) | man;
63 // Assumes that endian-ness is same for double and uint64_t.
64 double result;
65 memcpy(&result, &val, sizeof(val));
66 return result - 1.0;
67 }
68
69 // Helper function to convert two 32-bit uniform integers to two floats
70 // under the unit normal distribution.
71 PHILOX_DEVICE_INLINE
BoxMullerFloat(uint32_t x0,uint32_t x1,float * f0,float * f1)72 void BoxMullerFloat(uint32_t x0, uint32_t x1, float* f0, float* f1) {
73 // This function implements the Box-Muller transform:
74 // http://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform#Basic_form
75 // Do not send a really small number to log().
76 // We cannot mark "epsilon" as "static const" because NVCC would complain
77 const float epsilon = 1.0e-7f;
78 float u1 = Uint32ToFloat(x0);
79 if (u1 < epsilon) {
80 u1 = epsilon;
81 }
82 const float v1 = 2.0f * M_PI * Uint32ToFloat(x1);
83 const float u2 = sqrt(-2.0f * log(u1));
84 #if !defined(__linux__)
85 *f0 = sin(v1);
86 *f1 = cos(v1);
87 #else
88 sincosf(v1, f0, f1);
89 #endif
90 *f0 *= u2;
91 *f1 *= u2;
92 }
93
94 } // namespace random
95 } // namespace tensorflow
96
97 #endif // TENSORFLOW_CORE_LIB_RANDOM_RANDOM_DISTRIBUTIONS_UTILS_H_
98