1 /* Copyright 2016 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 // See docs in ../ops/random_ops.cc. 17 18 #define EIGEN_USE_THREADS 19 20 #include "tensorflow/core/kernels/random_poisson_op.h" 21 22 #include <algorithm> 23 #include <cmath> 24 #include <limits> 25 #include <memory> 26 27 #include "tensorflow/core/framework/op_kernel.h" 28 #include "tensorflow/core/framework/register_types.h" 29 #include "tensorflow/core/framework/tensor.h" 30 #include "tensorflow/core/framework/tensor_shape.h" 31 #include "tensorflow/core/lib/random/random_distributions.h" 32 #include "tensorflow/core/lib/random/simple_philox.h" 33 #include "tensorflow/core/util/guarded_philox_random.h" 34 #include "tensorflow/core/util/work_sharder.h" 35 36 #if EIGEN_COMP_GNUC && __cplusplus > 199711L 37 #define DISABLE_FLOAT_EQUALITY_WARNING \ 38 _Pragma("GCC diagnostic push") \ 39 _Pragma("GCC diagnostic ignored \"-Wfloat-equal\"") 40 #define ENABLE_FLOAT_EQUALITY_WARNING _Pragma("GCC diagnostic pop") 41 #else 42 #define DISABLE_FLOAT_EQUALITY_WARNING 43 #define ENABLE_FLOAT_EQUALITY_WARNING 44 #endif 45 46 #define UNIFORM(X) \ 47 if (uniform_remaining == 0) { \ 48 uniform_remaining = Uniform::kResultElementCount; \ 49 uniform_result = uniform(&gen); \ 50 } \ 51 uniform_remaining--; \ 52 CT X = uniform_result[uniform_remaining] 53 54 namespace tensorflow { 55 namespace { 56 57 static constexpr int kReservedSamplesPerOutput = 256; 58 59 typedef Eigen::ThreadPoolDevice CPUDevice; 60 61 template <typename T> 62 struct PoissonComputeType { 63 typedef double ComputeType; 64 }; 65 66 } // namespace 67 68 namespace functor { 69 70 template <typename Device, typename T, typename U> 71 struct PoissonFunctor { 72 void operator()(OpKernelContext* ctx, const Device& d, const T* rate_flat, 73 int num_rate, int num_samples, 74 const random::PhiloxRandom& rng, U* samples_flat); 75 }; 76 77 template <typename T, typename U> 78 struct PoissonFunctor<CPUDevice, T, U> { 79 void operator()(OpKernelContext* ctx, const CPUDevice& d, const T* rate_flat, 80 int num_rate, int num_samples, 81 const random::PhiloxRandom& rng, U* samples_flat) { 82 // Two different algorithms are employed, depending on the size of 83 // rate. 84 // If rate < 10, we use an algorithm attributed to Knuth: 85 // Seminumerical Algorithms. Art of Computer Programming, Volume 2. 86 // 87 // This algorithm runs in O(rate) time, and will require O(rate) 88 // uniform variates. 89 // 90 // If rate >= 10 we use a transformation-rejection algorithm from 91 // pairs of uniform random variables due to Hormann. 92 // http://www.sciencedirect.com/science/article/pii/0167668793909974 93 // 94 // The algorithm has an acceptance rate of ~89% for the smallest rate 95 // (~10), 96 // and higher accept rates for higher rate, so runtime is 97 // O(NumRate * NumSamples * k) with k ~ 1 / 0.89. 98 // 99 // We partition work first across rates then across 100 // samples-per-rate to 101 // avoid a couple flops which can be done on a per-rate basis. 102 103 typedef random::UniformDistribution<random::PhiloxRandom, CT> Uniform; 104 105 auto DoWork = [num_samples, num_rate, &rng, samples_flat, rate_flat]( 106 int start_output, int limit_output) { 107 // Capturing "rng" by value would only make a copy for the _shared_ 108 // lambda. Since we want to let each worker have its own copy, we pass 109 // "rng" by reference and explicitly do a copy assignment. 110 111 Uniform uniform; 112 typename Uniform::ResultType uniform_result; 113 for (int64 output_idx = start_output; output_idx < limit_output; 114 /* output_idx incremented within inner loop below */) { 115 const int64 rate_idx = output_idx / num_samples; 116 117 // Several calculations can be done on a per-rate basis. 118 const CT rate = CT(rate_flat[rate_idx]); 119 120 auto samples_rate_output = samples_flat + rate_idx; 121 122 if (rate < CT(10)) { 123 // Knuth's algorithm for generating Poisson random variates. 124 // Given a Poisson process, the time between events is exponentially 125 // distributed. If we have a Poisson process with rate lambda, then, 126 // the time between events is distributed Exp(lambda). If X ~ 127 // Uniform(0, 1), then Y ~ Exp(lambda), where Y = -log(X) / lambda. 128 // Thus to simulate a Poisson draw, we can draw X_i ~ Exp(lambda), 129 // and N ~ Poisson(lambda), where N is the least number such that 130 // \sum_i^N X_i > 1. 131 const CT exp_neg_rate = Eigen::numext::exp(-rate); 132 133 // Compute the rest of the samples for the current rate value. 134 for (int64 sample_idx = output_idx % num_samples; 135 sample_idx < num_samples && output_idx < limit_output; 136 sample_idx++, output_idx++) { 137 random::PhiloxRandom gen = rng; 138 gen.Skip(kReservedSamplesPerOutput * output_idx); 139 int16 uniform_remaining = 0; 140 141 CT prod = 1; 142 CT x = 0; 143 144 // Keep trying until we surpass e^(-rate). This will take 145 // expected time proportional to rate. 146 while (true) { 147 UNIFORM(u); 148 prod = prod * u; 149 if (prod <= exp_neg_rate && 150 x <= CT(Eigen::NumTraits<U>::highest())) { 151 samples_rate_output[sample_idx * num_rate] = U(x); 152 break; 153 } 154 x += 1; 155 } 156 } 157 continue; 158 } 159 // Transformed rejection due to Hormann. 160 // 161 // Given a CDF F(x), and G(x), a dominating distribution chosen such 162 // that it is close to the inverse CDF F^-1(x), compute the following 163 // steps: 164 // 165 // 1) Generate U and V, two independent random variates. Set U = U - 0.5 166 // (this step isn't strictly necessary, but is done to make some 167 // calculations symmetric and convenient. Henceforth, G is defined on 168 // [-0.5, 0.5]). 169 // 170 // 2) If V <= alpha * F'(G(U)) * G'(U), return floor(G(U)), else return 171 // to step 1. alpha is the acceptance probability of the rejection 172 // algorithm. 173 // 174 // For more details on transformed rejection, see: 175 // http://citeseer.ist.psu.edu/viewdoc/citations;jsessionid=1BEB35946CC807879F55D42512E5490C?doi=10.1.1.48.3054. 176 // 177 // The dominating distribution in this case: 178 // 179 // G(u) = (2 * a / (2 - |u|) + b) * u + c 180 181 using Eigen::numext::log; 182 const CT log_rate = log(rate); 183 184 // Constants used to define the dominating distribution. Names taken 185 // from Hormann's paper. Constants were chosen to define the tightest 186 // G(u) for the inverse Poisson CDF. 187 const CT b = CT(0.931) + CT(2.53) * Eigen::numext::sqrt(rate); 188 const CT a = CT(-0.059) + CT(0.02483) * b; 189 190 // This is the inverse acceptance rate. At a minimum (when rate = 10), 191 // this corresponds to ~75% acceptance. As the rate becomes larger, this 192 // approaches ~89%. 193 const CT inv_alpha = CT(1.1239) + CT(1.1328) / (b - CT(3.4)); 194 195 // Compute the rest of the samples for the current rate value. 196 for (int64 sample_idx = output_idx % num_samples; 197 sample_idx < num_samples && output_idx < limit_output; 198 sample_idx++, output_idx++) { 199 random::PhiloxRandom gen = rng; 200 gen.Skip(kReservedSamplesPerOutput * output_idx); 201 int16 uniform_remaining = 0; 202 203 while (true) { 204 UNIFORM(u); 205 u -= CT(0.5); 206 UNIFORM(v); 207 208 CT u_shifted = CT(0.5) - Eigen::numext::abs(u); 209 CT k = Eigen::numext::floor((CT(2) * a / u_shifted + b) * u + rate + 210 CT(0.43)); 211 212 if (k > CT(Eigen::NumTraits<U>::highest())) { 213 // retry in case of overflow. 214 continue; 215 } 216 217 // When alpha * f(G(U)) * G'(U) is close to 1, it is possible to 218 // find a rectangle (-u_r, u_r) x (0, v_r) under the curve, such 219 // that if v <= v_r and |u| <= u_r, then we can accept. 220 // Here v_r = 0.9227 - 3.6224 / (b - 2) and u_r = 0.43. 221 if (u_shifted >= CT(0.07) && 222 v <= CT(0.9277) - CT(3.6224) / (b - CT(2))) { 223 samples_rate_output[sample_idx * num_rate] = U(k); 224 break; 225 } 226 227 if (k < 0 || (u_shifted < CT(0.013) && v > u_shifted)) { 228 continue; 229 } 230 231 // The expression below is equivalent to the computation of step 2) 232 // in transformed rejection (v <= alpha * F'(G(u)) * G'(u)). 233 CT s = log(v * inv_alpha / (a / (u_shifted * u_shifted) + b)); 234 CT t = -rate + k * log_rate - Eigen::numext::lgamma(k + 1); 235 if (s <= t) { 236 samples_rate_output[sample_idx * num_rate] = U(k); 237 break; 238 } 239 } 240 } 241 } 242 }; 243 244 // This will depend on rate. 245 // For rate < 10, on average, O(rate) calls to uniform are 246 // needed, with that 247 // many multiplies. ~10 uniform calls on average with ~25 cost op calls. 248 // 249 // Very roughly, for rate >= 10, the single call to log + call to 250 // lgamma 251 // occur for ~60 percent of samples. 252 // 2 x 100 (64-bit cycles per log) * 0.62 = ~124 253 // Additionally, there are ~10 other ops (+, *, /, ...) at 3-6 cycles each: 254 // 40 * .62 = ~25. 255 // 256 // Finally, there are several other ops that are done every loop along with 257 // 2 uniform generations along with 5 other ops at 3-6 cycles each. 258 // ~15 / .89 = ~16 259 // 260 // In total this should be ~165 + 2 * Uniform::kElementCost. 261 // We assume that half the tensor has rate < 10, so on average 6 262 // uniform's 263 // will be needed. We will upper bound the other op cost by the one for 264 // rate > 10. 265 static const int kElementCost = 165 + 6 * Uniform::kElementCost + 266 6 * random::PhiloxRandom::kElementCost; 267 auto worker_threads = *(ctx->device()->tensorflow_cpu_worker_threads()); 268 Shard(worker_threads.num_threads, worker_threads.workers, 269 num_rate * num_samples, kElementCost, DoWork); 270 } 271 272 private: 273 typedef typename PoissonComputeType<T>::ComputeType CT; 274 }; 275 276 } // namespace functor 277 278 namespace { 279 280 // Samples from one or more Poisson distributions. 281 template <typename T, typename U> 282 class RandomPoissonOp : public OpKernel { 283 public: 284 explicit RandomPoissonOp(OpKernelConstruction* context) : OpKernel(context) { 285 OP_REQUIRES_OK(context, generator_.Init(context)); 286 } 287 288 void Compute(OpKernelContext* ctx) override { 289 const Tensor& shape_t = ctx->input(0); 290 const Tensor& rate_t = ctx->input(1); 291 292 TensorShape samples_shape; 293 OP_REQUIRES_OK(ctx, MakeShape(shape_t, &samples_shape)); 294 const int64 num_samples = samples_shape.num_elements(); 295 296 samples_shape.AppendShape(rate_t.shape()); 297 // Allocate output samples. 298 Tensor* samples_t = nullptr; 299 OP_REQUIRES_OK(ctx, ctx->allocate_output(0, samples_shape, &samples_t)); 300 if (num_samples == 0) return; 301 302 const auto rate_flat = rate_t.flat<T>().data(); 303 const int64 num_rate = rate_t.NumElements(); 304 auto samples_flat = samples_t->flat<U>().data(); 305 random::PhiloxRandom rng = generator_.ReserveRandomOutputs( 306 num_samples * num_rate, kReservedSamplesPerOutput); 307 308 functor::PoissonFunctor<CPUDevice, T, U>()( 309 ctx, ctx->eigen_device<CPUDevice>(), rate_flat, num_rate, num_samples, 310 rng, samples_flat); 311 } 312 313 private: 314 GuardedPhiloxRandom generator_; 315 316 TF_DISALLOW_COPY_AND_ASSIGN(RandomPoissonOp); 317 }; 318 } // namespace 319 320 #undef UNIFORM 321 322 #define REGISTER(TYPE) \ 323 REGISTER_KERNEL_BUILDER( \ 324 Name("RandomPoisson").Device(DEVICE_CPU).TypeConstraint<TYPE>("dtype"), \ 325 RandomPoissonOp<TYPE, TYPE>); 326 327 TF_CALL_half(REGISTER); 328 TF_CALL_float(REGISTER); 329 TF_CALL_double(REGISTER); 330 331 #define REGISTER_V2(RTYPE, OTYPE) \ 332 REGISTER_KERNEL_BUILDER(Name("RandomPoissonV2") \ 333 .Device(DEVICE_CPU) \ 334 .TypeConstraint<RTYPE>("R") \ 335 .TypeConstraint<OTYPE>("dtype"), \ 336 RandomPoissonOp<RTYPE, OTYPE>); 337 338 #define REGISTER_ALL(RTYPE) \ 339 REGISTER_V2(RTYPE, Eigen::half); \ 340 REGISTER_V2(RTYPE, float); \ 341 REGISTER_V2(RTYPE, double); \ 342 REGISTER_V2(RTYPE, int32); \ 343 REGISTER_V2(RTYPE, int64); 344 345 REGISTER_ALL(Eigen::half); 346 REGISTER_ALL(float); 347 REGISTER_ALL(double); 348 REGISTER_ALL(int32); 349 REGISTER_ALL(int64); 350 351 #undef REGISTER_ALL 352 #undef REGISTER_V2 353 #undef REGISTER 354 355 } // end namespace tensorflow 356