Make Schedule::Copy non-const, fork RND seed in Copy

Author: masahi

include/tvm/support/random_engine.h

@@ -115,7 +115,7 @@ class LinearCongruentialEngine {
    * \return The forked seed.
    */
   TRandState ForkSeed() {
-    // In order for reproducibility, we computer the new seed using RNG's random state and a
+    // 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;
   }

include/tvm/tir/schedule/schedule.h

@@ -123,7 +123,7 @@ class ScheduleNode : public runtime::Object {
    * 3) All the random variables are valid in the copy, pointing to the corresponding sref
    * reconstructed
    */
-  virtual Schedule Copy() const = 0;
+  virtual Schedule Copy() = 0;
   /*!
    * \brief Seed the randomness
    * \param seed The new random seed, -1 if use device random, otherwise non-negative

src/tir/schedule/concrete_schedule.cc

@@ -182,11 +182,12 @@ void ConcreteScheduleNode::Copy(ScheduleState* new_state, TSymbolTable* new_symb
   new_state->get()->DebugVerify();
 }
 
-Schedule ConcreteScheduleNode::Copy() const {
+Schedule ConcreteScheduleNode::Copy() {
   ObjectPtr<ConcreteScheduleNode> n = make_object<ConcreteScheduleNode>();
   n->error_render_level_ = this->error_render_level_;
   ConcreteScheduleNode::Copy(&n->state_, &n->symbol_table_);
   n->analyzer_ = std::make_unique<arith::Analyzer>();  // new analyzer needed because it is stateful
+  n->rand_state_ = ForkSeed();
   return Schedule(std::move(n));
 }
 

src/tir/schedule/concrete_schedule.h

@@ -61,7 +61,7 @@ class ConcreteScheduleNode : public ScheduleNode {
  public:
   ScheduleState state() const final { return state_; }
   Optional<Trace> trace() const override { return NullOpt; }
-  Schedule Copy() const override;
+  Schedule Copy() override;
   void Seed(support::LinearCongruentialEngine::TRandState seed = -1) final;
   support::LinearCongruentialEngine::TRandState ForkSeed() final;
 

src/tir/schedule/traced_schedule.cc

@@ -33,7 +33,7 @@ Schedule Schedule::Traced(IRModule mod, support::LinearCongruentialEngine::TRand
   return Schedule(std::move(n));
 }
 
-Schedule TracedScheduleNode::Copy() const {
+Schedule TracedScheduleNode::Copy() {
   ObjectPtr<TracedScheduleNode> n = make_object<TracedScheduleNode>();
   n->error_render_level_ = this->error_render_level_;
   ConcreteScheduleNode::Copy(&n->state_, &n->symbol_table_);

src/tir/schedule/traced_schedule.h

@@ -43,7 +43,7 @@ class TracedScheduleNode : public ConcreteScheduleNode {
 
  public:
   Optional<Trace> trace() const final { return trace_; }
-  Schedule Copy() const final;
+  Schedule Copy() final;
 
  public:
   /******** Schedule: Sampling ********/