compute/cker/include/cker/operation/Helper/RandomOpCpu.h

   1 /*
   2  * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
   3  * Copyright 2019 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_HELPER_RANDOM_OP_CPU_H__
  19 #define __NNFW_CKER_HELPER_RANDOM_OP_CPU_H__
  20
  21 #define EIGEN_USE_THREADS
  22
  23 #include <algorithm>
  24 #include <cmath>
  25 #include <memory>
  26
  27 #include "cker/Types.h"
  28 #include "cker/Shape.h"
  29 #include "cker/Utils.h"
  30
  31 #include "cker/eigen/EigenSupport.h"
  32
  33 #include "cker/operation/Helper/PhiloxRandom.h"
  34 #include "cker/operation/Helper/RandomOp.h"
  35 #include "cker/operation/Helper/RandomDistributions.h"
  36
  37 #if EIGEN_COMP_GNUC && __cplusplus > 199711L
  38 #define DISABLE_FLOAT_EQUALITY_WARNING \
  39   _Pragma("GCC diagnostic push") _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 namespace nnfw
  47 {
  48 namespace cker
  49 {
  50
  51 typedef Eigen::ThreadPoolDevice CPUDevice;
  52
  53 namespace functor
  54 {
  55 using random::PhiloxRandom;
  56 using random::SingleSampleAdapter;
  57
  58 // The default implementation of the functor, which should never be invoked
  59 // But we still need to provide implementation for now for the linker to work,
  60 // since we do not support all the distributions yet.
  61 template <typename Device, class Distribution> struct FillPhiloxRandom
  62 {
  63   typedef typename Distribution::ResultElementType T;
  64   void operator()() {}
  65 };
  66
  67 // A class to fill a specified range of random groups
  68 template <class Distribution, bool VariableSamplesPerOutput> struct FillPhiloxRandomTask;
  69
  70 // Specialization for distribution that takes a fixed number of samples for
  71 // each output.
  72 template <class Distribution> struct FillPhiloxRandomTask<Distribution, false>
  73 {
  74   typedef typename Distribution::ResultElementType T;
  75   static void Run(random::PhiloxRandom gen, T *data, int64_t size, Distribution dist)
  76   {
  77     const int kGroupSize = Distribution::kResultElementCount;
  78     gen.Skip(0);
  79     int64_t offset = 0;
  80
  81     // First fill all the full-size groups
  82     int64_t limit_group_full = size / kGroupSize;
  83     for (int64_t index = 0; index < limit_group_full; ++index)
  84     {
  85       auto samples = dist(&gen);
  86       std::copy(&samples[0], &samples[0] + kGroupSize, data + offset);
  87       offset += kGroupSize;
  88     }
  89
  90     int64_t remaining_size = size - limit_group_full * kGroupSize;
  91
  92     // If there are any remaining elements that need to be filled, process them
  93     if (remaining_size > 0)
  94     {
  95       auto samples = dist(&gen);
  96       std::copy(&samples[0], &samples[0] + remaining_size, data + offset);
  97     }
  98   }
  99 };
 100
 101 // Specialization for distribution that takes a variable number of samples for
 102 // each output. This will be slower due to the generality.
 103 template <class Distribution> struct FillPhiloxRandomTask<Distribution, true>
 104 {
 105   typedef typename Distribution::ResultElementType T;
 106   static constexpr int64_t kReservedSamplesPerOutput = 256;
 107
 108   static void Run(random::PhiloxRandom base_gen, T *data, int64_t size, Distribution dist)
 109   {
 110     const int kGroupSize = Distribution::kResultElementCount;
 111     static const int kGeneratorSkipPerOutputGroup =
 112         kGroupSize * kReservedSamplesPerOutput / PhiloxRandom::kResultElementCount;
 113
 114     int64_t offset = 0;
 115
 116     // First fill all the full-size groups
 117     int64_t limit_group_full = size / kGroupSize;
 118     int64_t group_index;
 119     for (group_index = 0; group_index < limit_group_full; ++group_index)
 120     {
 121       // Reset the generator to the beginning of the output group region
 122       // This is necessary if we want the results to be independent of order
 123       // of work
 124       PhiloxRandom gen = base_gen;
 125       gen.Skip(group_index * kGeneratorSkipPerOutputGroup);
 126       SingleSampleAdapter<PhiloxRandom> single_samples(&gen);
 127
 128       auto samples = dist(&single_samples);
 129       std::copy(&samples[0], &samples[0] + kGroupSize, data + offset);
 130       offset += kGroupSize;
 131     }
 132
 133     int64_t remaining_size = size - limit_group_full * kGroupSize;
 134     // If there are any remaining elements that need to be filled, process them
 135     if (remaining_size > 0)
 136     {
 137       PhiloxRandom gen = base_gen;
 138       gen.Skip(group_index * kGeneratorSkipPerOutputGroup);
 139       SingleSampleAdapter<PhiloxRandom> single_samples(&gen);
 140
 141       auto samples = dist(&single_samples);
 142       std::copy(&samples[0], &samples[0] + remaining_size, data + offset);
 143     }
 144   }
 145 };
 146
 147 // Partial specialization for CPU to fill the entire region with randoms
 148 // It splits the work into several tasks and run them in parallel
 149 template <class Distribution>
 150 void FillPhiloxRandom<CPUDevice, Distribution>::
 151 operator()(random::PhiloxRandom gen, typename Distribution::ResultElementType *data, int64_t size,
 152            Distribution dist)
 153 {
 154   FillPhiloxRandomTask<Distribution, Distribution::kVariableSamplesPerOutput>::Run(gen, data, size,
 155                                                                                    dist);
 156 }
 157
 158 } // namespace functor
 159
 160 } // end namespace tensorflow
 161 }
 162
 163 #endif // __NNFW_CKER_HELPER_RANDOM_OP_CPU_H__