random_engine.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); 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 TVM_SUPPORT_RANDOM_ENGINE_H_
#define TVM_SUPPORT_RANDOM_ENGINE_H_

#include <tvm/runtime/logging.h>

#include <cstdint>  // for uint64_t

namespace tvm {
namespace support {

class LinearCongruentialEngine {
 public:
  using result_type = uint64_t;
  using TRandState = int64_t;

  static constexpr TRandState multiplier = 48271;

  static constexpr TRandState increment = 0;

  static constexpr TRandState modulus = 2147483647;

  static constexpr result_type min() { return 0; }

  static constexpr result_type max() { return modulus - 1; }

  result_type operator()() {
    (*rand_state_ptr_) = ((*rand_state_ptr_) * multiplier + increment) % modulus;
    return *rand_state_ptr_;
  }

  void Seed(TRandState rand_state = 1) {
    rand_state %= modulus;  // Make sure the seed is within the range of modulus.
    if (rand_state == 0)
      rand_state = 1;  // Avoid getting all 0 given the current parameter set.
    else if (rand_state < 0)
      rand_state += modulus;             // Make sure the rand state is non-negative.
    ICHECK(rand_state_ptr_ != nullptr);  // Make sure the pointer is not null.
    *rand_state_ptr_ = rand_state;       // Change pointed random state to given random state value.
  }

  TRandState ForkSeed() {
    // In order for reproducibility, we computer the new seed using RNG's random state and a
    // different set of parameters. Note that both 32767 and 1999999973 are prime numbers.
    return ((*this)() * 32767) % 1999999973;
  }

  explicit LinearCongruentialEngine(TRandState* rand_state_ptr) {
    rand_state_ptr_ = rand_state_ptr;
  }

 private:
  TRandState* rand_state_ptr_;
};

}  // namespace support
}  // namespace tvm

#endif  // TVM_SUPPORT_RANDOM_ENGINE_H_