tvm
random_engine.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
26 #ifndef TVM_SUPPORT_RANDOM_ENGINE_H_
27 #define TVM_SUPPORT_RANDOM_ENGINE_H_
28 
29 #include <tvm/runtime/logging.h>
30 
31 #include <cstdint> // for uint64_t
32 #include <random>
33 
34 namespace tvm {
35 namespace support {
36 
48  public:
53  using result_type = uint64_t;
54  using TRandState = int64_t;
55 
57  static constexpr TRandState multiplier = 48271;
58 
60  static constexpr TRandState increment = 0;
61 
63  static constexpr TRandState modulus = 2147483647;
64 
69  static constexpr result_type min() { return 0; }
70 
75  static constexpr result_type max() { return modulus - 1; }
76 
81  static TRandState DeviceRandom() { return (std::random_device()()) % modulus; }
82 
94  (*rand_state_ptr_) = ((*rand_state_ptr_) * multiplier + increment) % modulus;
95  return *rand_state_ptr_;
96  }
97 
102  void Seed(TRandState rand_state) {
103  if (rand_state == -1) {
104  rand_state = DeviceRandom();
105  } else if (rand_state == 0) {
106  rand_state = 1;
107  }
108  ICHECK(rand_state >= 0) << "The random state should be nonnegative";
109  ICHECK(rand_state_ptr_ != nullptr);
110  *rand_state_ptr_ = rand_state % modulus;
111  }
112 
118  // In order for reproducibility, we compute the new seed using RNG's random state and a
119  // different set of parameters. Note that both 32767 and 1999999973 are prime numbers.
120  return ((*this)() * 32767) % 1999999973;
121  }
122 
130  explicit LinearCongruentialEngine(TRandState* rand_state_ptr) {
131  rand_state_ptr_ = rand_state_ptr;
132  }
133 
134  private:
135  TRandState* rand_state_ptr_;
136 };
137 
138 } // namespace support
139 } // namespace tvm
140 
141 #endif // TVM_SUPPORT_RANDOM_ENGINE_H_
static constexpr TRandState modulus
The modulus.
Definition: random_engine.h:63
runtime implementation for LibTorch/TorchScript.
Definition: analyzer.h:36
void Seed(TRandState rand_state)
Change the start random state of RNG with the seed of a new random state value.
Definition: random_engine.h:102
static constexpr TRandState multiplier
The multiplier.
Definition: random_engine.h:57
int64_t TRandState
Definition: random_engine.h:54
LinearCongruentialEngine(TRandState *rand_state_ptr)
Construct a random number generator with a random state pointer.
Definition: random_engine.h:130
This linear congruential engine is a drop-in replacement for std::minstd_rand. It strictly correspond...
Definition: random_engine.h:47
uint64_t result_type
The result type is defined as uint64_t here to avoid overflow.
Definition: random_engine.h:53
TRandState ForkSeed()
Fork a new seed for another RNG from current random state.
Definition: random_engine.h:117
static constexpr TRandState increment
The increment.
Definition: random_engine.h:60
result_type operator()()
Operator to move the random state to the next and return the new random state. According to definitio...
Definition: random_engine.h:93
static constexpr result_type max()
The maximum possible value of random state here.
Definition: random_engine.h:75
static TRandState DeviceRandom()
Get a device random state.
Definition: random_engine.h:81
static constexpr result_type min()
The minimum possible value of random state here.
Definition: random_engine.h:69