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_S_TIR_RANDOM_ENGINE_H_
#define TVM_S_TIR_RANDOM_ENGINE_H_
 
#include <tvm/ffi/error.h>
#include <tvm/ir/prim/expr.h>
 
#include <cstdint>
#include <random>
 
namespace tvm {
namespace s_tir {
class LinearCongruentialEngine {
 public:
  using TRandState = int64_t;
  using result_type = uint64_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; }
 
  static TRandState DeviceRandom() {
    std::random_device rd;
    return rd() % modulus;
  }
 
  result_type operator()() {
    (*rand_state_ptr_) = ((*rand_state_ptr_) * multiplier + increment) % modulus;
    return *rand_state_ptr_;
  }
  static TRandState NormalizeSeed(TRandState rand_state) {
    if (rand_state == -1) {
      rand_state = DeviceRandom();
    } else {
      rand_state %= modulus;
    }
    if (rand_state == 0) {
      rand_state = 1;
    }
    if (rand_state < 0) {
      TVM_FFI_THROW(ValueError) << "Random seed must be non-negative";
    }
    return rand_state;
  }
  void Seed(TRandState rand_state) {
    TVM_FFI_ICHECK(rand_state_ptr_ != nullptr);
    *rand_state_ptr_ = NormalizeSeed(rand_state);
  }
 
  TRandState ForkSeed() {
    // In order for reproducibility, we compute 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 s_tir
}  // namespace tvm
 
#endif  // TVM_S_TIR_RANDOM_ENGINE_H_