compute/cker/include/cker/operation/StatelessRandomUniform.h

   1 /*
   2  * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
   3  * Copyright 2017 The TensorFlow Authors. All Rights Reserved.
   4  *
   5  * Licensed under the Apache License, Version 2.0 (the "License");
   6  * you may not use this file except in compliance with the License.
   7  * You may obtain a copy of the License at
   8  *
   9  *      http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 #ifndef __NNFW_CKER_STATELESS_RANDOM_UNIFORM_H__
  19 #define __NNFW_CKER_STATELESS_RANDOM_UNIFORM_H__
  20
  21 #include "cker/Types.h"
  22 #include "cker/Shape.h"
  23 #include "cker/Utils.h"
  24
  25 #include "cker/eigen/EigenSupport.h"
  26
  27 #include "cker/operation/Helper/Tensor.h"
  28 #include "cker/operation/Helper/PhiloxRandom.h"
  29 #include "cker/operation/Helper/RandomOpCpu.h"
  30 #include "cker/operation/Helper/RandomDistributions.h"
  31
  32 namespace nnfw
  33 {
  34 namespace cker
  35 {
  36
  37 void GenerateKey(Tensor seed, random::PhiloxRandom::Key *out_key,
  38                  random::PhiloxRandom::ResultType *out_counter)
  39 {
  40   // Grab the two seeds
  41   uint32_t seed0;
  42   uint32_t seed1;
  43
  44   const auto seed_vals = seed.flat<int32_t>();
  45
  46   seed0 = seed_vals(0);
  47   seed1 = seed_vals(1);
  48   // Scramble the seeds so that the user doesn't need to worry about which
  49   // part of the seed needs to be strong.
  50   (*out_key)[0] = 0x3ec8f720;
  51   (*out_key)[1] = 0x02461e29;
  52   (*out_counter)[0] = static_cast<uint32_t>(seed0);
  53   (*out_counter)[1] = (*out_counter)[3] = 0;
  54   (*out_counter)[2] = static_cast<uint32_t>(seed1);
  55   const auto mix = random::PhiloxRandom(*out_counter, *out_key)();
  56   (*out_key)[0] = mix[0];
  57   (*out_key)[1] = mix[1];
  58   (*out_counter)[0] = (*out_counter)[1] = 0;
  59   (*out_counter)[2] = mix[2];
  60   (*out_counter)[3] = mix[3];
  61 }
  62
  63 template <typename Device, class Distribution>
  64 void Fill(random::PhiloxRandom random, Tensor *output)
  65 {
  66   // Build distribution
  67   typedef typename Distribution::ResultElementType T;
  68
  69   auto flat = output->flat<T>();
  70   // Reuse the compute kernels from the stateful random ops
  71   functor::FillPhiloxRandom<Device, Distribution>()(random, flat.data(), flat.size(),
  72                                                     Distribution());
  73 }
  74
  75 inline void StatelessRandomUniform(const Shape &shape_shape, const int *shape_data,
  76                                    const Shape &seed_shape, const int *seed_data,
  77                                    const Shape &output_shape, float *output_data)
  78 {
  79   Tensor shape_t;
  80   Tensor seed_t;
  81
  82   shape_t.shape.ReplaceWith(shape_shape.DimensionsCount(), shape_shape.DimsData());
  83   shape_t.buffer = (void *)shape_data;
  84
  85   seed_t.shape.ReplaceWith(seed_shape.DimensionsCount(), seed_shape.DimsData());
  86   seed_t.buffer = (void *)seed_data;
  87
  88   Tensor output_t;
  89   output_t.shape.ReplaceWith(output_shape.DimensionsCount(), output_shape.DimsData());
  90   output_t.buffer = output_data;
  91
  92   random::PhiloxRandom::Key key;
  93   random::PhiloxRandom::ResultType counter;
  94
  95   GenerateKey(seed_t, &key, &counter);
  96
  97   Fill<Eigen::ThreadPoolDevice, random::UniformDistribution<random::PhiloxRandom, float>>(
  98     random::PhiloxRandom(counter, key), &output_t);
  99 }
 100 } // namespace cker
 101 } // namespace nnfw
 102
 103 #endif // __NNFW_CKER_STATELESS_RANDOM_UNIFORM_H__