random.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); you
// may not use this file except in compliance with the License.  You may
// obtain a copy of the License at:
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the license.

#ifndef LBANN_UTILS_RANDOM_HPP
#define LBANN_UTILS_RANDOM_HPP

#include "lbann/base.hpp"
#include "lbann/comm.hpp"
#include "lbann/io/persist.hpp"
#include "lbann/utils/exception.hpp"
#include "lbann/utils/random_number_generators.hpp"

namespace lbann {

enum class probability_distribution
{
  invalid,
  gaussian,
  bernoulli,
  uniform
};

template <typename Generator, typename T>
inline T fast_rand_int(Generator& g, T max)
{
#ifdef LBANN_DEBUG
  if (max == 0) {
    LBANN_ERROR("fast_rand_int called with max=0");
  }
#endif
  typename Generator::result_type x;
  do {
    x = g();
  } while (x >= (Generator::max() - Generator::max() % max));
  return x % max;
}

template <typename Generator, typename T>
inline T fast_rand_int_pow2(Generator& g, T max)
{
  typename Generator::result_type x;
  max -= 1;
  const typename Generator::result_type upper =
    Generator::max() -
    (Generator::max() & (typename Generator::result_type)max);
  do {
    x = g();
  } while (x >= upper);
  return x & ((typename Generator::result_type)max);
}

// Methods for generating uniformly random values in [0, 1).

namespace details {

template <typename Generator, typename T>
struct random_uniform_impl
{
  static T generate(Generator&);
};
template <typename Generator>
struct random_uniform_impl<Generator, float>
{
  static float generate(Generator& g)
  {
    // float has a 24-bit significand, including an implicit bit. See
    // section on converting uint64_ts to doubles in
    // http://xoshiro.di.unimi.it/
    constexpr uint64_t mask32 = 0xFFFFFFFFull;
    const uint64_t r = uint64_t(g()) & mask32;
    return (r >> 8) * (1.0f / 16777216.0f);
  }
};
template <typename Generator>
struct random_uniform_impl<Generator, double>
{
  static double generate(Generator& g)
  {
    // double has a 53-bit significand, including an implicit bit. See
    // section on converting uint64_ts to doubles in
    // http://xoshiro.di.unimi.it/
    constexpr uint64_t mask32 = 0xFFFFFFFFull;
    const uint64_t r = (uint64_t(g()) << 32) | (uint64_t(g()) & mask32);
    return (r >> 11) * (1.0 / 9007199254740992.0);
  }
};

} // namespace details

template <typename T, typename Generator>
inline T random_uniform(Generator& g)
{
  return details::random_uniform_impl<Generator, T>::generate(g);
}

template <typename TensorDataType>
void gaussian_fill(El::AbstractDistMatrix<TensorDataType>& mat,
                   El::Int m,
                   El::Int n,
                   TensorDataType mean = 0.0,
                   TensorDataType stddev = 1.0);
template <typename TensorDataType>
void bernoulli_fill(El::AbstractDistMatrix<TensorDataType>& mat,
                    El::Int m,
                    El::Int n,
                    double p = 0.5);
template <typename TensorDataType>
void uniform_fill(El::AbstractDistMatrix<TensorDataType>& mat,
                  El::Int m,
                  El::Int n,
                  TensorDataType center = 0.0,
                  TensorDataType radius = 1.0);

template <typename TensorDataType>
void gaussian_fill_procdet(El::AbstractDistMatrix<TensorDataType>& mat,
                           El::Int m,
                           El::Int n,
                           TensorDataType mean = 0.0,
                           TensorDataType stddev = 1.0);
template <typename TensorDataType>
void bernoulli_fill_procdet(El::AbstractDistMatrix<TensorDataType>& mat,
                            El::Int m,
                            El::Int n,
                            double p = 0.5);
template <typename TensorDataType>
void uniform_fill_procdet(El::AbstractDistMatrix<TensorDataType>& mat,
                          El::Int m,
                          El::Int n,
                          TensorDataType center = 0.0,
                          TensorDataType radius = 1.0);

template <typename TensorDataType>
void gaussian_fill_parallel(El::AbstractDistMatrix<TensorDataType>& mat,
                            El::Int m,
                            El::Int n,
                            TensorDataType mean = 0.0,
                            TensorDataType stddev = 1.0);

bool save_rng_to_checkpoint_shared(persist& p, lbann_comm* comm);
bool save_rng_to_checkpoint_distributed(persist& p, lbann_comm* comm);
bool load_rng_from_checkpoint(persist& p, const lbann_comm* comm);

template <typename DistType, typename DType = DataType>
class rng
{

private:
  DistType m_dist; // Distribution type

public:
  typename DistType::result_type gen() { return m_dist(get_generator()); }
  rng() {}
  // bernoulli_distribution with prob p
  rng(DType p) : m_dist(p) {}
  // (uniform) real distribution between min/mean and max/stdev
  rng(DType a, DType b) : m_dist(a, b) {}
};

template <typename DType = DataType>
void rng_bernoulli(const float p, DistMat* m)
{

  float scale = 1. / (1. - p);

  //@todo: use seed from parameter
  rng<std::bernoulli_distribution, DType> myrn(p); // magic seed?

  for (int row = 0; row < m->LocalHeight(); ++row) {
    for (int col = 0; col < m->LocalWidth(); ++col) {
      m->Set(row, col, myrn.gen() * scale); // SetLocal?
    }
  }
}

#ifndef LBANN_RANDOM_INSTANTIATE
#define PROTO(T)                                                               \
  extern template void gaussian_fill<T>(El::AbstractDistMatrix<T> & mat,       \
                                        El::Int m,                             \
                                        El::Int n,                             \
                                        T mean,                                \
                                        T stddev);                             \
  extern template void bernoulli_fill<T>(El::AbstractDistMatrix<T> & mat,      \
                                         El::Int m,                            \
                                         El::Int n,                            \
                                         double p);                            \
  extern template void uniform_fill<T>(El::AbstractDistMatrix<T> & mat,        \
                                       El::Int m,                              \
                                       El::Int n,                              \
                                       T center,                               \
                                       T radius);                              \
  extern template void gaussian_fill_procdet<T>(El::AbstractDistMatrix<T> &    \
                                                  mat,                         \
                                                El::Int m,                     \
                                                El::Int n,                     \
                                                T mean,                        \
                                                T stddev);                     \
  extern template void bernoulli_fill_procdet<T>(El::AbstractDistMatrix<T> &   \
                                                   mat,                        \
                                                 El::Int m,                    \
                                                 El::Int n,                    \
                                                 double p);                    \
  extern template void uniform_fill_procdet<T>(El::AbstractDistMatrix<T> &     \
                                                 mat,                          \
                                               El::Int m,                      \
                                               El::Int n,                      \
                                               T center,                       \
                                               T radius);                      \
  extern template void gaussian_fill_parallel<T>(El::AbstractDistMatrix<T> &   \
                                                   mat,                        \
                                                 El::Int m,                    \
                                                 El::Int n,                    \
                                                 T mean,                       \
                                                 T stddev)

#define LBANN_INSTANTIATE_CPU_HALF
#define LBANN_INSTANTIATE_GPU_HALF
#include "lbann/macros/instantiate.hpp"
#undef PROTO
#undef LBANN_INSTANTIATE_CPU_HALF
#undef LBANN_INSTANTIATE_GPU_HALF
#endif // LBANN_RANDOM_INSTANTIATE

} // namespace lbann
#endif // LBANN_UTILS_RANDOM_HPP