zxybazh commented on a change in pull request #8642: URL: https://github.com/apache/tvm/pull/8642#discussion_r682287668
########## File path: include/tvm/support/random_engine.h ########## @@ -0,0 +1,116 @@ +/* + * 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. + */ + +/*! + * \file random_engine.h + * \brief Random number generator, for Sampler and Sampling functions. + */ + +#ifndef TVM_SUPPORT_RANDOM_ENGINE_H_ +#define TVM_SUPPORT_RANDOM_ENGINE_H_ + +#include <tvm/runtime/logging.h> + +#include <cstdint> // for int64_t + +namespace tvm { +namespace support { + +/*! + * \brief This linear congruential engine is a drop-in replacement for and strictly corresponds to + * std::minstd_rand but designed to be serializable and strictly reproducible. Specifically + * implemented for meta schedule but also reusable for other purposes. + * \note Part of std::linear_congruential_engine's member functions are not included, for full + * member functions of std::minstd_rand, please check out the following link: + * https://en.cppreference.com/w/cpp/numeric/random/linear_congruential_engine + */ +class LinearCongruentialEngine { + public: + /*! + * \brief The result type is defined as int64_t here for meta_schedule sampler usage. + * \note The type name is not in Google style because it is used in STL's distribution inferface. + */ + using result_type = int64_t; + + /*! \brief The multiplier */ + static constexpr result_type multiplier = 48271; + + /*! \brief The increment */ + static constexpr result_type increment = 0; + + /*! \brief The modulus */ + static constexpr result_type modulus = 2147483647; + + /*! + * \brief The minimum possible value of random state here. + * \note The function name is uncapilized because it is used in STL's distribution inferface. + */ + result_type min() { return 0; } + + /*! + * \brief The maximum possible value of random state here. + * \note The function name is uncapilized because it is used in STL's distribution inferface. + */ + result_type max() { return modulus - 1; } + + /*! + * \brief Operator to move the random state to the next and return the new random state. According + * to definition of linear congruential engine, the new random state value is computed as + * new_random_state = (current_random_state * multiplier + increment) % modulus. + * \return The next current random state value in the type of result_type. + * \note In case of potential overflow, please use Schrage multiplication algorithm to implement. + * We also assume the given rand state is not nullptr here. + */ + result_type operator()() { + (*rand_state_ptr_) = ((*rand_state_ptr_) * multiplier + increment) % modulus; + return *rand_state_ptr_; + } + + /*! + * \brief Change the start random state of RNG with the seed of a new random state value. + * \param rand_state The random state given in result_type. + */ + void Seed(result_type 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. Review comment: IMHO if the starting `rand_state` is non-zero and not a multiple of the modulus, it won't be zero after any iterations. Thus, we only take care of the situation when seeding by changing the starting state to the default, which is `1`. Also, even tho 0 is not the expected usage here, we can still guarantee reproducibility. Therefore it won't be a problem. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
