/*
 * 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.
 */
#include "./concrete_schedule.h"

#include <tvm/ffi/cast.h>
#include <tvm/runtime/logging.h>

#include <random>

namespace tvm {
namespace s_tir {
using namespace tvm::tirx;

Schedule Schedule::Concrete(IRModule mod, LinearCongruentialEngine::TRandState seed, int debug_mask,
                            ScheduleErrorRenderLevel error_render_level, bool enable_check) {
  ffi::ObjectPtr<ConcreteScheduleNode> n = ffi::make_object<ConcreteScheduleNode>();
  n->state_ = ScheduleState(mod, debug_mask, enable_check);
  n->error_render_level_ = error_render_level;
  n->symbol_table_ = {};
  n->analyzer_ = arith::Analyzer();
  n->Seed(seed);
  GlobalVar gv;
  if (FindEntryFunc(mod, &gv) != nullptr) {
    n->func_working_on_ = gv;
  } else {
    n->func_working_on_ = std::nullopt;
  }
  return Schedule(std::move(n));
}

/******** Copy ********/

/*! \brief Helper class to perform a deep copy of the sref tree */
class ScheduleCopier {
  using TSymbolTable = ConcreteScheduleNode::TSymbolTable;
  template <class K, class V>
  using UMap = std::unordered_map<K, V>;
  template <class K, class V>
  using SMap = std::unordered_map<K, V, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>;

 public:
  static void Copy(const ConcreteScheduleNode* self, ScheduleState* new_state,
                   TSymbolTable* new_symbol_table) {
    const ScheduleState& src_state = self->state_;
    ScheduleCopier copier(src_state);
    ffi::ObjectPtr<ScheduleStateNode> n = ffi::make_object<ScheduleStateNode>();
    n->mod = src_state->mod;
    n->block_info = copier.Copy(src_state->block_info);
    n->stmt2ref = copier.Copy(src_state->stmt2ref);
    n->debug_mask = src_state->debug_mask;
    n->enable_check = src_state->enable_check;
    *new_state = ScheduleState(std::move(n));
    *new_symbol_table = copier.Copy(self->symbol_table_);
  }

 private:
  /*! \brief Create the copier and properly set up the `old2new_` table */
  explicit ScheduleCopier(const ScheduleState& state) {
    // Create SRef tree without parents
    for (const auto& kv : state->stmt2ref) {
      const StmtSRefNode* sref = kv.second.operator->();
      old2new_.emplace(sref,                          // the old StmtSRef
                       StmtSRef(/*stmt=*/sref->stmt,  // the new StmtSRef
                                /*parent=*/nullptr,   // parent is not set yet
                                /*seq_index=*/sref->seq_index));
    }
    // Fill in the parent field
    // Find out the root along the way
    for (auto& kv : old2new_) {
      const StmtSRefNode* parent = kv.first->parent;
      StmtSRef& sref = kv.second;
      sref->parent = parent ? old2new_.at(parent).get() : nullptr;
    }
  }

  /*! \brief Copy StmtSRef */
  StmtSRef Copy(const StmtSRef& sref) { return old2new_.at(sref.operator->()); }

  /*! \brief Copy StmtSRefNode */
  StmtSRef Copy(const StmtSRefNode* sref) {
    if (old2new_.count(sref)) {
      return old2new_.at(sref);
    }
    // Handle expired sref
    return old2new_[sref] = StmtSRef(nullptr, nullptr, -1);
  }

  /*! \brief Copy ffi::Array<StmtSRef> */
  ffi::Array<StmtSRef> Copy(const ffi::Array<StmtSRef>& list) {
    ffi::Array<StmtSRef> result;
    result.reserve(list.size());
    for (const StmtSRef& elem : list) {
      result.push_back(Copy(elem));
    }
    return result;
  }

  /*! \brief Copy ffi::Array<Dependency> */
  ffi::Array<Dependency> Copy(const ffi::Array<Dependency>& list) {
    ffi::Array<Dependency> result;
    result.reserve(list.size());
    for (const Dependency& elem : list) {
      result.push_back(Dependency(Copy(elem->src), Copy(elem->dst), elem->kind));
    }
    return result;
  }

  /*! \brief Copy SMap<StmtSRef, ffi::Array<Dependency>> */
  SMap<StmtSRef, ffi::Array<Dependency>> Copy(const SMap<StmtSRef, ffi::Array<Dependency>>& map) {
    SMap<StmtSRef, ffi::Array<Dependency>> result;
    result.reserve(map.size());
    for (const auto& kv : map) {
      result[Copy(kv.first)] = Copy(kv.second);
    }
    return result;
  }

  /*! \brief Copy SMap<Buffer, ffi::Array<StmtSRef>> */
  SMap<Buffer, ffi::Array<StmtSRef>> Copy(const SMap<Buffer, ffi::Array<StmtSRef>>& map) {
    SMap<Buffer, ffi::Array<StmtSRef>> result;
    result.reserve(map.size());
    for (const auto& kv : map) {
      result[kv.first] = Copy(kv.second);
    }
    return result;
  }

  /*! \brief Copy SMap<StmtSRef, Scope> */
  SMap<StmtSRef, SBlockInfo> Copy(const SMap<StmtSRef, SBlockInfo>& scopes) {
    SMap<StmtSRef, SBlockInfo> result;
    for (const auto& kv : scopes) {
      const StmtSRef& old_sref = kv.first;
      const SBlockInfo& old_info = kv.second;
      SBlockInfo new_info = old_info;
      ffi::ObjectPtr<SBlockScopeNode> scope = ffi::make_object<SBlockScopeNode>();
      scope->src2deps = Copy(old_info.scope->src2deps);
      scope->dst2deps = Copy(old_info.scope->dst2deps);
      scope->buffer_writers = Copy(old_info.scope->buffer_writers);
      new_info.scope = SBlockScope(std::move(scope));
      result[Copy(old_sref)] = std::move(new_info);
    }
    return result;
  }

  /*! \brief Copy the stmt2ref */
  UMap<const StmtNode*, StmtSRef> Copy(const UMap<const StmtNode*, StmtSRef>& stmt2ref) {
    UMap<const StmtNode*, StmtSRef> result;
    result.reserve(stmt2ref.size());
    for (const auto& kv : stmt2ref) {
      const StmtNode* stmt = kv.first;
      const StmtSRef& sref = kv.second;
      result.emplace(stmt, Copy(sref));
    }
    return result;
  }

  /*! \brief Copy the symbol table */
  TSymbolTable Copy(const TSymbolTable& tab) {
    TSymbolTable result;
    for (const auto& kv : tab) {
      ffi::ObjectRef entry = kv.second;
      if (const auto* sref = entry.as<StmtSRefNode>()) {
        entry = Copy(sref);
      }
      result.Set(kv.first, entry);
    }
    return result;
  }

 private:
  std::unordered_map<const StmtSRefNode*, StmtSRef> old2new_;
};

void ConcreteScheduleNode::WorkOn(const ffi::String& func_name) {
  this->func_working_on_ = this->state_->mod->GetGlobalVar(func_name);
}

void ConcreteScheduleNode::Copy(ScheduleState* new_state, TSymbolTable* new_symbol_table) const {
  ScheduleCopier::Copy(this, new_state, new_symbol_table);
  new_state->get()->DebugVerify();
}

Schedule ConcreteScheduleNode::Copy() {
  ffi::ObjectPtr<ConcreteScheduleNode> n = ffi::make_object<ConcreteScheduleNode>();
  n->func_working_on_ = this->func_working_on_;
  n->error_render_level_ = this->error_render_level_;
  ConcreteScheduleNode::Copy(&n->state_, &n->symbol_table_);
  n->analyzer_ = arith::Analyzer();  // new analyzer needed because it is stateful
  n->rand_state_ = ForkSeed();
  return Schedule(std::move(n));
}

/*! \brief Macro that guards the beginning of each invocation of TensorIR schedule primitive */
#define TVM_TIR_SCHEDULE_BEGIN() try {
/*!
 * \brief Macro that pairs with `TVM_TIR_SCHEDULE_BEGIN`, handling potential errors and error
 * message rendering
 * \param level An ScheduleErrorRenderLevel enum, level of error rendering
 * \sa ScheduleErrorRenderLevel
 */
#define TVM_TIR_SCHEDULE_END(primitive, level)                       \
  }                                                                  \
  catch (const ScheduleError& error) {                               \
    if ((level) == ScheduleErrorRenderLevel::kDetail) {              \
      TVM_FFI_THROW(ScheduleError) << error.RenderReport(primitive); \
    } else if ((level) == ScheduleErrorRenderLevel::kFast) {         \
      TVM_FFI_THROW(ScheduleError) << error.FastErrorString();       \
    } else if ((level) == ScheduleErrorRenderLevel::kNone) {         \
      TVM_FFI_THROW(ScheduleError) << "(not rendered)";              \
    }                                                                \
  }

/******** Schedule: Schedule: Sampling ********/

void ConcreteScheduleNode::Seed(LinearCongruentialEngine::TRandState seed) {
  this->rand_state_ = LinearCongruentialEngine::NormalizeSeed(seed);
}

LinearCongruentialEngine::TRandState ConcreteScheduleNode::ForkSeed() {
  return LinearCongruentialEngine(&rand_state_).ForkSeed();
}

ExprRV ConcreteScheduleNode::SampleCategorical(const ffi::Array<int64_t>& candidates,
                                               const ffi::Array<FloatImm>& probs,
                                               ffi::Optional<int64_t> decision) {
  TVM_TIR_SCHEDULE_BEGIN();
  return CreateRV(s_tir::SampleCategorical(&this->rand_state_, candidates, probs, &decision));
  TVM_TIR_SCHEDULE_END("sample-categorical", this->error_render_level_);
  throw;
}

ffi::Array<ExprRV> ConcreteScheduleNode::SamplePerfectTile(
    const LoopRV& loop_rv, int n, int max_innermost_factor,
    ffi::Optional<ffi::Array<int64_t>> decision) {
  TVM_TIR_SCHEDULE_BEGIN();
  // use None RV object to denotes auto-infer tile factors.
  return CreateRV(s_tir::SamplePerfectTile(&this->rand_state_, this->GetSRef(loop_rv), n,
                                           max_innermost_factor, &decision),
                  /*convert_negone_to_none=*/true);
  TVM_TIR_SCHEDULE_END("sample-perfect-tile", this->error_render_level_);
  throw;
}

ffi::Array<ExprRV> ConcreteScheduleNode::SamplePartitionedTile(
    const LoopRV& loop_rv, int n, int partition_pos, int innerpart_factor,
    ffi::Optional<ffi::Array<int64_t>> decision) {
  TVM_TIR_SCHEDULE_BEGIN();
  return CreateRV(s_tir::SamplePartitionedTile(&this->rand_state_, this->GetSRef(loop_rv), n,
                                               partition_pos, innerpart_factor, &decision));
  TVM_TIR_SCHEDULE_END("sample-partitioned-tile", this->error_render_level_);
  throw;
}

LoopRV ConcreteScheduleNode::SampleComputeLocation(const SBlockRV& block_rv,
                                                   ffi::Optional<int64_t> decision) {
  TVM_TIR_SCHEDULE_BEGIN();
  return CreateRV<LoopRV>(
      s_tir::SampleComputeLocation(state_, &this->rand_state_, this->GetSRef(block_rv), &decision));
  TVM_TIR_SCHEDULE_END("sample-compute-location", this->error_render_level_);
  throw;
}

/******** Schedule: Get blocks & loops ********/

SBlockRV ConcreteScheduleNode::GetSBlock(const ffi::String& name,
                                         const ffi::Optional<ffi::String>& func_name) {
  class NotSingleResult : public ScheduleError {
   public:
    explicit NotSingleResult(ffi::String name, IRModule mod, const ffi::Array<StmtSRef>& blocks)
        : name_(name), mod_(mod), blocks_{} {
      blocks_.reserve(blocks.size());
      for (const StmtSRef& block_sref : blocks) {
        const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
        blocks_.push_back(ffi::GetRef<SBlock>(block));
      }
    }

    IRModule mod() const final { return mod_; }
    ffi::Array<ffi::ObjectRef> LocationsOfInterest() const final {
      return {blocks_.begin(), blocks_.end()};
    }

    ffi::String DetailRenderTemplate() const final {
      if (blocks_.empty()) {
        return "Cannot find a block with the name: " + name_;
      } else {
        return "Found  " + std::to_string(blocks_.size()) + " blocks with the name: " + name_;
      }
    }

    ffi::String FastErrorString() const final {
      if (blocks_.empty()) {
        return "ScheduleError: Cannot find a block with the specified name";
      } else {
        return "ScheduleError: Found multiple blocks with the specified name";
      }
    }

    ffi::String name_;
    IRModule mod_;
    ffi::Array<SBlock> blocks_;
  };
  GlobalVar gv;
  if (func_name.has_value()) {
    gv = state_->mod->GetGlobalVar(func_name.value());
  } else if (func_working_on_.has_value()) {
    gv = this->func_working_on_.value();
  } else {
    TVM_FFI_THROW(ValueError)
        << "`get_sblock` does not know which function to be working on. Please "
           "specify the function name explicitly, or call `work_on` to specify the function "
           "before using `get_sblock`.";
  }
  ffi::Array<StmtSRef> blocks = s_tir::GetSBlocks(this->state_, name, gv);
  if (blocks.size() != 1) {
    TVM_TIR_SCHEDULE_BEGIN();
    throw NotSingleResult(name, this->state_->mod, blocks);
    TVM_TIR_SCHEDULE_END("get-block", this->error_render_level_);
  }
  return CreateRV<SBlockRV>(blocks[0]);
}

ffi::Array<LoopRV> ConcreteScheduleNode::GetLoops(const SBlockRV& block_rv) {
  return CreateRV<LoopRV>(s_tir::GetLoops(this->GetSRef(block_rv)));
}

ffi::Array<SBlockRV> ConcreteScheduleNode::GetChildBlocks(const SBlockRV& block_rv) {
  ffi::Array<SBlockRV> result;
  TVM_TIR_SCHEDULE_BEGIN();
  result = CreateRV<SBlockRV>(s_tir::GetChildBlocks(state_, this->GetSRef(block_rv)));
  TVM_TIR_SCHEDULE_END("get-child-blocks", this->error_render_level_);
  this->state_->DebugVerify();
  return result;
}

ffi::Array<SBlockRV> ConcreteScheduleNode::GetChildBlocks(const LoopRV& loop_rv) {
  ffi::Array<SBlockRV> result;
  TVM_TIR_SCHEDULE_BEGIN();
  result = CreateRV<SBlockRV>(s_tir::GetChildBlocks(state_, this->GetSRef(loop_rv)));
  TVM_TIR_SCHEDULE_END("get-child-blocks", this->error_render_level_);
  this->state_->DebugVerify();
  return result;
}

ffi::Array<SBlockRV> ConcreteScheduleNode::GetProducers(const SBlockRV& block_rv) {
  TVM_TIR_SCHEDULE_BEGIN();
  return CreateRV<SBlockRV>(s_tir::GetProducers(state_, this->GetSRef(block_rv)));
  TVM_TIR_SCHEDULE_END("get-producers", this->error_render_level_);
  throw;
}

ffi::Array<SBlockRV> ConcreteScheduleNode::GetConsumers(const SBlockRV& block_rv) {
  TVM_TIR_SCHEDULE_BEGIN();
  return CreateRV<SBlockRV>(s_tir::GetConsumers(state_, this->GetSRef(block_rv)));
  TVM_TIR_SCHEDULE_END("get-consumers", this->error_render_level_);
  throw;
}

ffi::Array<SBlockRV> ConcreteScheduleNode::GetOutputBlocks(const SBlockRV& scope_block_rv) {
  TVM_TIR_SCHEDULE_BEGIN();
  return CreateRV<SBlockRV>(s_tir::GetOutputBlocks(state_, this->GetSRef(scope_block_rv)));
  TVM_TIR_SCHEDULE_END("get-output-blocks", this->error_render_level_);
  throw;
}

/******** Schedule: Transform loops ********/

LoopRV ConcreteScheduleNode::Merge(const ffi::Array<LoopRV>& loop_rvs) {
  TVM_FFI_CHECK(loop_rvs.size() > 1, ValueError) << "'merge' requires at least 2 loop(s)";
  ffi::Array<StmtSRef> loop_srefs = this->GetSRefs(loop_rvs);
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::Merge(state_, loop_srefs);
  TVM_TIR_SCHEDULE_END("merge", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<LoopRV>(result);
}

LoopRV ConcreteScheduleNode::Fuse(const ffi::Array<LoopRV>& loop_rvs, bool preserve_unit_iters) {
  TVM_FFI_CHECK(!loop_rvs.empty(), ValueError) << "'fuse' requires at least 1 loop(s)";
  ffi::Array<StmtSRef> loop_srefs = this->GetSRefs(loop_rvs);
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::Fuse(state_, loop_srefs, preserve_unit_iters);
  TVM_TIR_SCHEDULE_END("fuse", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<LoopRV>(result);
}

class NotSingleInferFactorError : public ScheduleError {
 public:
  explicit NotSingleInferFactorError(IRModule mod) : mod_(mod) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: only one factor can be specified as -1 or none";
  }

  ffi::String DetailRenderTemplate() const final {
    return "Only one factor can be specified as -1 or none";
  }

  IRModule mod() const final { return mod_; }
  ffi::Array<ffi::ObjectRef> LocationsOfInterest() const final { return {}; }

  IRModule mod_;
};

class WrongFactorError : public ScheduleError {
 public:
  explicit WrongFactorError(IRModule mod, For loop, bool product)
      : mod_(mod), loop_(std::move(loop)), product_(product) {}

  ffi::String FastErrorString() const final {
    if (product_)
      return "ScheduleError: The product of factors is not larger than or equal to the extent of "
             "loop";
    else
      return "ScheduleError: The sum of factors is larger than or equal to the extent of loop";
  }

  ffi::String DetailRenderTemplate() const final {
    if (product_)
      return "The product of factors is not larger than or equal to the extent of loop {0}";
    else
      return "The sum of factors is larger than or equal to the extent of loop {0}";
  }

  IRModule mod() const final { return mod_; }
  ffi::Array<ffi::ObjectRef> LocationsOfInterest() const final { return {loop_}; }

  IRModule mod_;
  For loop_;
  bool product_;
};

class NonPositiveFactorError : public ScheduleError {
 public:
  explicit NonPositiveFactorError(IRModule mod, int64_t factor, size_t idx)
      : mod_(std::move(mod)), factor_(factor), idx_(idx) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: All the constant factors are required to be positive. However, some "
           "constant input factor is zero or negative.";
  }
  ffi::String DetailRenderTemplate() const final {
    std::ostringstream os;
    os << "All the constant factors are required to be positive. However, the factor at position "
       << idx_ << " is " << factor_;
    return os.str();
  }
  IRModule mod() const final { return mod_; }
  ffi::Array<ffi::ObjectRef> LocationsOfInterest() const final { return {}; }

 private:
  IRModule mod_;
  int64_t factor_;
  size_t idx_;
};

ffi::Array<LoopRV> ConcreteScheduleNode::Split(const LoopRV& loop_rv,
                                               const ffi::Array<ffi::Optional<ExprRV>>& factor_rvs,
                                               bool preserve_unit_iters, bool disable_predication) {
  // Prepare for the splitting
  StmtSRef loop_sref = this->GetSRef(loop_rv);
  const ForNode* loop = TVM_SREF_TO_FOR(loop_sref);
  ffi::Array<PrimExpr> factors;
  factors.reserve(factor_rvs.size());
  int infer_index = -1;
  PrimExpr tot_length = 1;
  ffi::Array<StmtSRef> results;
  TVM_TIR_SCHEDULE_BEGIN();
  // infer factor if needed and check validity of factors
  for (size_t i = 0; i < factor_rvs.size(); i++) {
    if (!factor_rvs[i].has_value()) {
      factors.push_back(IntImm::Int32(-1));
      if (infer_index != -1) {
        throw NotSingleInferFactorError(state_->mod);
      }
      infer_index = i;
    } else {
      PrimExpr factor = this->Get(factor_rvs[i].value());
      if (is_const_int(factor) && !is_positive_const(factor)) {
        throw NonPositiveFactorError(state_->mod, factor.as<IntImmNode>()->value, i);
      }
      if (factor.ty().bits() > loop->extent.ty().bits()) {
        factor = cast(loop->extent.ty(), factor);
      }
      factors.push_back(factor);
      tot_length *= factor;
    }
  }
  if (infer_index != -1) {
    factors.Set(infer_index,
                this->analyzer_->Simplify(floordiv(loop->extent + tot_length - 1, tot_length)));
  } else if (!this->analyzer_->CanProve(tot_length >= loop->extent)) {
    throw WrongFactorError(state_->mod, ffi::GetRef<For>(loop), true);
  }
  results = s_tir::Split(state_, loop_sref, factors, preserve_unit_iters, disable_predication);
  TVM_TIR_SCHEDULE_END("split", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<LoopRV>(results);
}

ffi::Array<LoopRV> ConcreteScheduleNode::LoopPartition(
    const LoopRV& loop_rv, const ffi::Array<ffi::Optional<ExprRV>>& factor_rvs,
    bool preserve_unit_iters) {
  class SymbolicShapeError : public ScheduleError {
   public:
    explicit SymbolicShapeError(IRModule mod, For loop) : mod_(mod), loop_(std::move(loop)) {}

    ffi::String FastErrorString() const final {
      return "ScheduleError: The min and extent values of the loop are required to be known at "
             "compile time. However, dynamic shape has been detected.";
    }

    ffi::String DetailRenderTemplate() const final {
      return "Detected dynamic shape in either min or extent of a loop {0}";
    }

    IRModule mod() const final { return mod_; }
    ffi::Array<ffi::ObjectRef> LocationsOfInterest() const final { return {loop_}; }

    IRModule mod_;
    For loop_;
  };

  // Prepare for the loop_partitioning
  StmtSRef loop_sref = this->GetSRef(loop_rv);
  const ForNode* loop = TVM_SREF_TO_FOR(loop_sref);
  ffi::Array<PrimExpr> factors;
  factors.reserve(factor_rvs.size());
  int infer_index = -1;
  PrimExpr tot_length = 0;
  ffi::Array<StmtSRef> results;
  TVM_TIR_SCHEDULE_BEGIN();
  if (!is_const_number(loop->min) || !is_const_number(loop->extent)) {
    throw SymbolicShapeError(state_->mod, ffi::GetRef<For>(loop));
  }
  // infer factor if needed and check validity of factors
  for (size_t i = 0; i < factor_rvs.size(); i++) {
    if (!factor_rvs[i].has_value()) {
      factors.push_back(IntImm::Int32(-1));
      if (infer_index != -1) {
        throw NotSingleInferFactorError(state_->mod);
      }
      infer_index = i;
    } else {
      PrimExpr factor = this->Get(factor_rvs[i].value());
      if (is_const_int(factor) && !is_positive_const(factor)) {
        throw NonPositiveFactorError(state_->mod, factor.as<IntImmNode>()->value, i);
      }
      if (factor.ty().bits() > loop->extent.ty().bits()) {
        factor = cast(loop->extent.ty(), factor);
      }
      factors.push_back(factor);
      tot_length += factor;
    }
  }
  if (this->analyzer_->CanProve(tot_length >= loop->extent)) {
    throw WrongFactorError(state_->mod, ffi::GetRef<For>(loop), false);
  }
  if (infer_index != -1) {
    // if there is a 'None' in the factor list, 'None' becomes the difference between the extent and
    // the sum of the factors excluding 'None' specified in the partition directive.
    factors.Set(infer_index, loop->extent - tot_length);
  }
  for (size_t i = 1; i < factor_rvs.size(); i++) {
    factors.Set(i, factors[i] + factors[i - 1]);
  }
  if (infer_index == -1) {
    factors.push_back(loop->extent);
  }
  results = s_tir::LoopPartition(state_, loop_sref, factors, preserve_unit_iters);
  TVM_TIR_SCHEDULE_END("loop_partition", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<LoopRV>(results);
}

void ConcreteScheduleNode::Reorder(const ffi::Array<LoopRV>& ordered_loop_rvs) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Reorder(state_, GetSRefs(ordered_loop_rvs));
  TVM_TIR_SCHEDULE_END("reorder", this->error_render_level_);
  this->state_->DebugVerify();
}

void ConcreteScheduleNode::ReorderBlockIterVar(const SBlockRV& block_rv,
                                               const ffi::Array<int64_t> new_order) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::ReorderBlockIterVar(state_, GetSRef(block_rv), new_order);
  TVM_TIR_SCHEDULE_END("reorder_block_iter_var", this->error_render_level_);
  this->state_->DebugVerify();
}

LoopRV ConcreteScheduleNode::AddUnitLoop(const SBlockRV& block_rv) {
  LoopRV result{ffi::UnsafeInit()};
  TVM_TIR_SCHEDULE_BEGIN();
  result = CreateRV<LoopRV>(s_tir::AddUnitLoop(state_, GetSRef(block_rv)));
  TVM_TIR_SCHEDULE_END("add-unit-loop", this->error_render_level_);
  this->state_->DebugVerify();
  return result;
}

LoopRV ConcreteScheduleNode::AddUnitLoop(const LoopRV& loop_rv) {
  LoopRV result{ffi::UnsafeInit()};
  TVM_TIR_SCHEDULE_BEGIN();
  result = CreateRV<LoopRV>(s_tir::AddUnitLoop(state_, GetSRef(loop_rv)));
  TVM_TIR_SCHEDULE_END("add-unit-loop", this->error_render_level_);
  this->state_->DebugVerify();
  return result;
}

/******** Schedule: Manipulate ForKind ********/

void ConcreteScheduleNode::Parallel(const LoopRV& loop_rv) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Parallel(state_, this->GetSRef(loop_rv));
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("parallel", this->error_render_level_);
}

void ConcreteScheduleNode::Vectorize(const LoopRV& loop_rv) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Vectorize(state_, this->GetSRef(loop_rv));
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("vectorize", this->error_render_level_);
}

void ConcreteScheduleNode::Bind(const LoopRV& loop_rv, const ffi::String& thread_axis) {
  if (thread_axis == "vthread") {
    LOG(WARNING) << "`vthread` is legacy behavior and is going to be deprecated. Please use "
                    "`vthread.x`, `vthread.y` and `vthread.z` instead";
  }
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Bind(state_, this->GetSRef(loop_rv), thread_axis);
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("bind", this->error_render_level_);
}

void ConcreteScheduleNode::Unroll(const LoopRV& loop_rv) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Unroll(state_, this->GetSRef(loop_rv));
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("unroll", this->error_render_level_);
}

/******** Schedule: Insert cache stages ********/

SBlockRV ConcreteScheduleNode::CacheRead(const SBlockRV& block_rv, int read_buffer_index,
                                         const ffi::String& storage_scope,
                                         const ffi::Array<SBlockRV> consumer_blocks) {
  StmtSRef result{nullptr};
  // Create a new array of SRefs from the consumer block list.
  ffi::Array<StmtSRef> consumer_block_refs = {};
  for (SBlockRV block : consumer_blocks) {
    consumer_block_refs.push_back(this->GetSRef(block));
  }
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::CacheRead(state_, this->GetSRef(block_rv), read_buffer_index, storage_scope,
                            consumer_block_refs);
  TVM_TIR_SCHEDULE_END("cache-read", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<SBlockRV>(result);
}

SBlockRV ConcreteScheduleNode::CacheWrite(const SBlockRV& block_rv, int write_buffer_index,
                                          const ffi::String& storage_scope,
                                          const ffi::Array<SBlockRV> consumer_blocks) {
  StmtSRef result{nullptr};
  // Create a new array of SRefs from the consumer block list.
  ffi::Array<StmtSRef> consumer_block_refs = {};
  for (SBlockRV block : consumer_blocks) {
    consumer_block_refs.push_back(this->GetSRef(block));
  }
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::CacheWrite(state_, this->GetSRef(block_rv), write_buffer_index, storage_scope,
                             consumer_block_refs);
  TVM_TIR_SCHEDULE_END("cache-write", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<SBlockRV>(result);
}

SBlockRV ConcreteScheduleNode::ReindexCacheRead(const SBlockRV& block_rv, int read_buffer_index,
                                                const ffi::String& storage_scope,
                                                const IndexMap& index_map) {
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::ReindexCacheRead(state_, this->GetSRef(block_rv), read_buffer_index,
                                   storage_scope, index_map);
  TVM_TIR_SCHEDULE_END("reverse-cache-read", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<SBlockRV>(result);
}

SBlockRV ConcreteScheduleNode::ReindexCacheWrite(const SBlockRV& block_rv, int write_buffer_index,
                                                 const ffi::String& storage_scope,
                                                 const IndexMap& index_map) {
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::ReindexCacheWrite(state_, this->GetSRef(block_rv), write_buffer_index,
                                    storage_scope, index_map);
  TVM_TIR_SCHEDULE_END("reverse-cache-write", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<SBlockRV>(result);
}

ffi::Array<SBlockRV> ConcreteScheduleNode::CacheInplace(const SBlockRV& block_rv,
                                                        int write_buffer_index,
                                                        const ffi::String& storage_scope) {
  ffi::Array<StmtSRef> results;
  TVM_TIR_SCHEDULE_BEGIN();
  results = s_tir::CacheInplace(state_, this->GetSRef(block_rv), write_buffer_index, storage_scope);
  TVM_TIR_SCHEDULE_END("cache-buffer", this->error_render_level_);
  this->state_->DebugVerify();
  ffi::Array<SBlockRV> return_blocks;
  return_blocks.push_back(CreateRV<SBlockRV>(results[0]));
  return_blocks.push_back(CreateRV<SBlockRV>(results[1]));
  return return_blocks;
}

ffi::Array<SBlockRV> ConcreteScheduleNode::CacheIndex(const SBlockRV& block_rv,
                                                      const ffi::String& storage_scope,
                                                      int cse_thresh) {
  ffi::Array<StmtSRef> result;
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::CacheIndex(state_, this->GetSRef(block_rv), storage_scope, cse_thresh);
  TVM_TIR_SCHEDULE_END("cache-index", this->error_render_level_);
  this->state_->DebugVerify();
  ffi::Array<SBlockRV> return_blocks;
  for (const StmtSRef& blockrv : result) {
    return_blocks.push_back(CreateRV<SBlockRV>(blockrv));
  }
  return return_blocks;
}

SBlockRV ConcreteScheduleNode::ReIndex(const SBlockRV& block_rv, int buffer_index,
                                       BufferIndexType buffer_index_type) {
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::ReIndex(state_, this->GetSRef(block_rv), buffer_index, buffer_index_type);
  TVM_TIR_SCHEDULE_END("reindex", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<SBlockRV>(result);
}

/******** Schedule: Data movement ********/

SBlockRV ConcreteScheduleNode::ReadAt(const LoopRV& loop_rv, const SBlockRV& block_rv,
                                      int read_buffer_index, const ffi::String& storage_scope) {
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::ReadAt(state_, this->GetSRef(loop_rv), this->GetSRef(block_rv), read_buffer_index,
                         storage_scope);
  TVM_TIR_SCHEDULE_END("read-at", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<SBlockRV>(result);
}

SBlockRV ConcreteScheduleNode::WriteAt(const LoopRV& loop_rv, const SBlockRV& block_rv,
                                       int write_buffer_index, const ffi::String& storage_scope) {
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::WriteAt(state_, this->GetSRef(loop_rv), this->GetSRef(block_rv),
                          write_buffer_index, storage_scope);
  TVM_TIR_SCHEDULE_END("write-at", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<SBlockRV>(result);
}

/******** Schedule: Compute location ********/

void ConcreteScheduleNode::ComputeAt(const SBlockRV& block_rv, const LoopRV& loop_rv,
                                     bool preserve_unit_loops, int index) {
  static StmtSRef inline_mark = StmtSRef::InlineMark();
  static StmtSRef root_mark = StmtSRef::RootMark();
  StmtSRef loop_sref = this->GetSRef(loop_rv);
  if (loop_sref.same_as(root_mark)) {
    // do nothing
  } else if (loop_sref.same_as(inline_mark)) {
    TVM_TIR_SCHEDULE_BEGIN();
    s_tir::ComputeInline(state_, this->GetSRef(block_rv));
    TVM_TIR_SCHEDULE_END("compute-at", this->error_render_level_);
  } else {
    TVM_TIR_SCHEDULE_BEGIN();
    s_tir::ComputeAt(state_, this->GetSRef(block_rv), loop_sref, preserve_unit_loops, index);
    TVM_TIR_SCHEDULE_END("compute-at", this->error_render_level_);
  }
  this->state_->DebugVerify();
}

void ConcreteScheduleNode::ReverseComputeAt(const SBlockRV& block_rv, const LoopRV& loop_rv,
                                            bool preserve_unit_loops, int index) {
  static StmtSRef inline_mark = StmtSRef::InlineMark();
  static StmtSRef root_mark = StmtSRef::RootMark();
  StmtSRef loop_sref = this->GetSRef(loop_rv);
  if (loop_sref.same_as(root_mark)) {
    // do nothing
  } else if (loop_sref.same_as(inline_mark)) {
    TVM_TIR_SCHEDULE_BEGIN();
    s_tir::ReverseComputeInline(state_, this->GetSRef(block_rv));
    TVM_TIR_SCHEDULE_END("reverse-compute-at", this->error_render_level_);
  } else {
    TVM_TIR_SCHEDULE_BEGIN();
    s_tir::ReverseComputeAt(state_, this->GetSRef(block_rv), loop_sref, preserve_unit_loops, index);
    TVM_TIR_SCHEDULE_END("reverse-compute-at", this->error_render_level_);
  }
  this->state_->DebugVerify();
}

void ConcreteScheduleNode::ComputeInline(const SBlockRV& block_rv) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::ComputeInline(state_, this->GetSRef(block_rv));
  TVM_TIR_SCHEDULE_END("compute-inline", this->error_render_level_);
  this->state_->DebugVerify();
}

void ConcreteScheduleNode::ReverseComputeInline(const SBlockRV& block_rv) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::ReverseComputeInline(state_, this->GetSRef(block_rv));
  TVM_TIR_SCHEDULE_END("reverse-compute-inline", this->error_render_level_);
  this->state_->DebugVerify();
}

void ConcreteScheduleNode::FuseReductionEpilogue(const SBlockRV& reduction_block_rv,
                                                 const SBlockRV& epilogue_block_rv) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::FuseReductionEpilogue(state_, this->GetSRef(reduction_block_rv),
                               this->GetSRef(epilogue_block_rv));
  TVM_TIR_SCHEDULE_END("fuse-reduction-epilogue", this->error_render_level_);
  this->state_->DebugVerify();
}

/******** Schedule: SBlock Annotation ********/

void ConcreteScheduleNode::StorageAlign(const SBlockRV& block_rv, int buffer_index, int axis,
                                        int factor, int offset) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::StorageAlign(state_, this->GetSRef(block_rv), buffer_index, axis, factor, offset);
  TVM_TIR_SCHEDULE_END("storage-align", this->error_render_level_);
  this->state_->DebugVerify();
}

void ConcreteScheduleNode::SetScope(const SBlockRV& block_rv, int buffer_index,
                                    const ffi::String& storage_scope) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::SetScope(state_, this->GetSRef(block_rv), buffer_index, storage_scope);
  TVM_TIR_SCHEDULE_END("set-scope", this->error_render_level_);
  this->state_->DebugVerify();
}

void ConcreteScheduleNode::UnsafeSetDType(const SBlockRV& block_rv, int buffer_index,
                                          const ffi::String& dtype) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::UnsafeSetDType(state_, this->GetSRef(block_rv), buffer_index, dtype);
  TVM_TIR_SCHEDULE_END("set-dtype", this->error_render_level_);
  this->state_->DebugVerify();
}

/******** Schedule: Reduction ********/

SBlockRV ConcreteScheduleNode::DecomposeReduction(const SBlockRV& block_rv, const LoopRV& loop_rv) {
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::DecomposeReduction(state_, this->GetSRef(block_rv), this->GetSRef(loop_rv));
  TVM_TIR_SCHEDULE_END("decompose-reduction", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<SBlockRV>(result);
}

SBlockRV ConcreteScheduleNode::RFactor(const LoopRV& loop_rv, int factor_axis) {
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::RFactor(state_, this->GetSRef(loop_rv), factor_axis);
  TVM_TIR_SCHEDULE_END("rfactor", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<SBlockRV>(result);
}

/******** Schedule: Blockize & Tensorize ********/
SBlockRV ConcreteScheduleNode::Blockize(const LoopRV& loop_rv, bool preserve_unit_iters) {
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::Blockize(state_, this->GetSRef(loop_rv), preserve_unit_iters);
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("blockize", this->error_render_level_);
  return CreateRV<SBlockRV>(result);
}

SBlockRV ConcreteScheduleNode::Blockize(const ffi::Array<SBlockRV>& blocks,
                                        bool preserve_unit_iters) {
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::Blockize(state_, this->GetSRefs(blocks), preserve_unit_iters);
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("blockize", this->error_render_level_);
  return CreateRV<SBlockRV>(result);
}

void ConcreteScheduleNode::Tensorize(const LoopRV& loop_rv, const ffi::String& intrin,
                                     bool preserve_unit_iters) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Tensorize(state_, this->GetSRef(loop_rv), tirx::TensorIntrin::Get(intrin).value(),
                   preserve_unit_iters);
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("tensorize", this->error_render_level_);
}

void ConcreteScheduleNode::Tensorize(const SBlockRV& block_rv, const ffi::String& intrin,
                                     bool preserve_unit_iters) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Tensorize(state_, this->GetSRef(block_rv), tirx::TensorIntrin::Get(intrin).value(),
                   preserve_unit_iters);
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("tensorize", this->error_render_level_);
}

/******** Schedule: Annotation ********/

Any ConcreteScheduleNode::CheckAndGetAnnotationValue(const ffi::Any& ann_val) {
  if (auto opt_str = ann_val.try_cast<ffi::String>()) {
    return *std::move(opt_str);
  }
  if (ann_val.type_index() < ffi::TypeIndex::kTVMFFISmallStr) {
    return ann_val;
  }
  // prefer to return int/float literals for annotations
  if (auto opt_intimm = ann_val.try_cast<IntImm>()) {
    return (*std::move(opt_intimm))->value;
  }
  if (auto opt_floatimm = ann_val.try_cast<FloatImm>()) {
    return (*std::move(opt_floatimm))->value;
  }

  if (auto expr = ann_val.as<PrimExpr>()) {
    TVM_FFI_CHECK(!expr.value().as<StringImmNode>(), TypeError)
        << "ffi::String is expected, but gets StringImm";
    auto res_expr = this->Get(expr.value());
    // prefer to return int/float literals for annotations
    if (auto opt_intimm = res_expr.as<IntImm>()) {
      return (*std::move(opt_intimm))->value;
    }
    if (auto opt_floatimm = res_expr.as<FloatImm>()) {
      return (*std::move(opt_floatimm))->value;
    }
    return res_expr;
  }
  if (const auto* arr = ann_val.as<ffi::ArrayObj>()) {
    ffi::Array<Any> result;
    result.reserve(arr->size());
    for (size_t i = 0; i < arr->size(); i++) {
      result.push_back(CheckAndGetAnnotationValue(arr->at(i)));
    }
    return result;
  }
  if (const auto* dict = ann_val.as<ffi::MapObj>()) {
    ffi::Map<ffi::String, ffi::Any> result;
    for (auto it = dict->begin(); it != dict->end(); ++it) {
      const auto& key = it->first;
      auto value = CheckAndGetAnnotationValue(it->second);
      if (const StringImmNode* imm = key.as<StringImmNode>()) {
        result.Set(imm->value, value);
      } else if (auto opt_str = key.try_cast<ffi::String>()) {
        result.Set(opt_str.value(), value);
      } else {
        TVM_FFI_THROW(TypeError) << "annotation dict key expect to be ffi::String or StringImm";
      }
    }
    return result;
  }
  TVM_FFI_THROW(TypeError)
      << "Only strings, integers, floats, ExprRVs and Arrays are supported for now, but "
      << "gets: " << ann_val.GetTypeKey();
  TVM_FFI_UNREACHABLE();
}

void ConcreteScheduleNode::Annotate(const LoopRV& loop_rv, const ffi::String& ann_key,
                                    const Any& ann_val) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Annotate(state_, this->GetSRef(loop_rv), ann_key,
                  this->CheckAndGetAnnotationValue(ann_val));
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("annotate", this->error_render_level_);
}

void ConcreteScheduleNode::Unannotate(const LoopRV& loop_rv, const ffi::String& ann_key) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Unannotate(state_, this->GetSRef(loop_rv), ann_key);
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("unannotate", this->error_render_level_);
}

void ConcreteScheduleNode::Annotate(const SBlockRV& block_rv, const ffi::String& ann_key,
                                    const Any& ann_val) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Annotate(state_, this->GetSRef(block_rv), ann_key,
                  this->CheckAndGetAnnotationValue(ann_val));
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("annotate", this->error_render_level_);
}

void ConcreteScheduleNode::Unannotate(const SBlockRV& block_rv, const ffi::String& ann_key) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::Unannotate(state_, this->GetSRef(block_rv), ann_key);
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("unannotate", this->error_render_level_);
}

/******** Schedule: Layout transformation ********/
void ConcreteScheduleNode::TransformLayout(const SBlockRV& block_rv, int buffer_index,
                                           BufferIndexType buffer_index_type,
                                           const IndexMap& index_map,
                                           const ffi::Optional<IndexMap>& pad_value,
                                           bool assume_injective_transform) {
  TVM_TIR_SCHEDULE_BEGIN();
  auto f_subst = [&](const Var& var) -> ffi::Optional<PrimExpr> {
    if (auto opt_expr = symbol_table_.Get(var)) {
      return opt_expr.value().as_or_throw<PrimExpr>();
    } else {
      return std::nullopt;
    }
  };
  auto new_index_map = Substitute(index_map, f_subst);
  s_tir::TransformLayout(state_, this->GetSRef(block_rv), buffer_index, buffer_index_type,
                         new_index_map, pad_value, assume_injective_transform);
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("transform_layout", this->error_render_level_);
}

void ConcreteScheduleNode::TransformBlockLayout(const SBlockRV& block_rv,
                                                const IndexMap& index_map) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::TransformBlockLayout(state_, this->GetSRef(block_rv), index_map);
  this->state_->DebugVerify();
  TVM_TIR_SCHEDULE_END("transform_block_layout", this->error_render_level_);
}

void ConcreteScheduleNode::SetAxisSeparator(const SBlockRV& block_rv, int buffer_index,
                                            BufferIndexType buffer_index_type,
                                            const ffi::Array<IntImm>& axis_separators) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::SetAxisSeparator(state_, this->GetSRef(block_rv), buffer_index, buffer_index_type,
                          axis_separators);
  TVM_TIR_SCHEDULE_END("set-axis-separator", this->error_render_level_);
  this->state_->DebugVerify();
}

/******** Schedule: Padding ********/

SBlockRV ConcreteScheduleNode::DecomposePadding(const SBlockRV& block_rv, const LoopRV& loop_rv) {
  StmtSRef result{nullptr};
  TVM_TIR_SCHEDULE_BEGIN();
  result = s_tir::DecomposePadding(state_, this->GetSRef(block_rv), this->GetSRef(loop_rv));
  TVM_TIR_SCHEDULE_END("decompose-padding", this->error_render_level_);
  this->state_->DebugVerify();
  return CreateRV<SBlockRV>(result);
}

void ConcreteScheduleNode::PadEinsum(const SBlockRV& block_rv, const ffi::Array<int64_t>& padding) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::PadEinsum(state_, this->GetSRef(block_rv), padding);
  TVM_TIR_SCHEDULE_END("pad-einsum", this->error_render_level_);
  this->state_->DebugVerify();
}

/******** Schedule: Buffer Transformation ********/

void ConcreteScheduleNode::RollingBuffer(const SBlockRV& block_rv, int write_buffer_index) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::RollingBuffer(state_, this->GetSRef(block_rv), write_buffer_index);
  TVM_TIR_SCHEDULE_END("rolling-buffer", this->error_render_level_);
  this->state_->DebugVerify();
}

/******** Schedule: Misc ********/

void ConcreteScheduleNode::UnsafeHideBufferAccess(const SBlockRV& block_rv,
                                                  const ffi::String& buf_type,
                                                  const ffi::Array<IntImm>& buf_index_array) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::UnsafeHideBufferAccess(state_, this->GetSRef(block_rv), buf_type, buf_index_array);
  TVM_TIR_SCHEDULE_END("hide-buffer-access", this->error_render_level_);
  this->state_->DebugVerify();
}

void ConcreteScheduleNode::AnnotateBufferAccess(const SBlockRV& block_rv, int buffer_index,
                                                BufferIndexType buffer_index_type,
                                                const IndexMap& index_map) {
  TVM_TIR_SCHEDULE_BEGIN();
  s_tir::AnnotateBufferAccess(state_, this->GetSRef(block_rv), buffer_index, buffer_index_type,
                              index_map);
  TVM_TIR_SCHEDULE_END("annotate-buffer-access", this->error_render_level_);
  this->state_->DebugVerify();
}

}  // namespace s_tir
}  // namespace tvm
