/*
 * 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 <tvm/ffi/cast.h>

#include "../utils.h"

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

/*! \brief Append a new predicate to the each child of type BlockRealize (not recursively) */
class BlockPredicateAppender : public StmtMutator {
 public:
  /*!
   * \brief Constructor
   * \param to_append The predicate to be appended to BlockRealizeNode
   */
  explicit BlockPredicateAppender(const PrimExpr& to_append) : to_append_(to_append) {}

 private:
  // For each direct child of type BlockRealizeNode, append the predicate
  Stmt VisitStmt_(const SBlockRealizeNode* realize) final {
    // We do not recursively do this
    ffi::ObjectPtr<SBlockRealizeNode> n = CopyOnWrite(realize);
    n->predicate = n->predicate && to_append_;
    return SBlockRealize(n);
  }

  /*! \brief The predicate to be appended */
  const PrimExpr& to_append_;
};

/*! \brief Substitute vars and collect the reuse mapping of opaque blocks */
class SubstituteVarAndCollectOpaqueBlock : public StmtExprMutator {
 public:
  explicit SubstituteVarAndCollectOpaqueBlock(std::function<ffi::Optional<Expr>(const Var&)> vmap,
                                              ffi::Map<SBlock, SBlock>* opaque_blocks)
      : vmap_(vmap), opaque_blocks_(opaque_blocks) {}

 private:
  Expr VisitExpr_(const VarNode* op) final {
    Var var = ffi::GetRef<Var>(op);
    if (ffi::Optional<Expr> ret = vmap_(var)) {
      return tvm::cast(var->ty.as_or_throw<PrimType>(), ret.value().as_or_throw<PrimExpr>());
    } else {
      return var;
    }
  }

  Stmt VisitStmt_(const SBlockRealizeNode* op) final {
    SBlockRealize realize = StmtMutator::VisitStmt_(op).as_or_throw<SBlockRealize>();
    if (realize->block->iter_vars.empty()) {
      opaque_blocks_->Set(op->block, realize->block);
    }
    return realize;
  }

  /*! \brief The substitute function */
  std::function<ffi::Optional<Expr>(const Var&)> vmap_;
  /*! \brief The reuse mapping of opaque blocks */
  ffi::Map<SBlock, SBlock>* opaque_blocks_;
};

/*! \brief Simplify the binding of block realize and update the opaque block reuse mapping */
class IterMapSimplifyBlockBinding : public StmtExprMutator {
 public:
  explicit IterMapSimplifyBlockBinding(ffi::MapObj* opaque_blocks,
                                       ffi::Map<PrimVar, Range> loop_var2extent,
                                       bool preserve_unit_iters)
      : opaque_blocks_(opaque_blocks),
        loop_var2extent_(loop_var2extent),
        preserve_unit_iters_(preserve_unit_iters) {}

  static For SimplifyBindings(Stmt stmt, const ffi::Array<StmtSRef>& loop_srefs,
                              ffi::MapObj* opaque_blocks, bool preserve_unit_iters) {
    ffi::Map<PrimVar, Range> loop_var2extent;
    for (const StmtSRef& sref : loop_srefs) {
      const ForNode* loop = TVM_SREF_TO_FOR(sref);
      loop_var2extent.Set(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent));
    }
    return IterMapSimplifyBlockBinding(opaque_blocks, std::move(loop_var2extent),
                                       preserve_unit_iters)(std::move(stmt))
        .as_or_throw<For>();
  }

 private:
  Stmt VisitStmt_(const ForNode* op) final {
    loop_var2extent_.Set(op->loop_var, Range::FromMinExtent(op->min, op->extent));
    Stmt res = StmtMutator::VisitStmt_(op);
    loop_var2extent_.erase(op->loop_var);
    return res;
  }

  Stmt VisitStmt_(const SBlockRealizeNode* op) final {
    // skip opaque block and update mapping
    if (op->iter_values.empty()) {
      SBlock block = op->block;
      SBlockRealize realize = StmtMutator::VisitStmt_(op).as_or_throw<SBlockRealize>();
      for (const auto& entry : *opaque_blocks_) {
        if (entry.second.same_as(block)) {
          opaque_blocks_->at(entry.first) = realize->block;
          break;
        }
      }
      return realize;
    }
    ffi::Array<PrimExpr> v =
        arith::IterMapSimplify(/*indices=*/op->iter_values,
                               /*input_iters=*/loop_var2extent_,
                               /*input_pred=*/op->predicate,
                               /*check_level=*/arith::IterMapLevel::Surjective,
                               /*analyzer=*/analzyer_,
                               /*simplify_trivial_iterators=*/!preserve_unit_iters_);
    if (v.same_as(op->iter_values)) {
      return ffi::GetRef<Stmt>(op);
    } else {
      ffi::ObjectPtr<SBlockRealizeNode> n = CopyOnWrite(op);
      n->iter_values = std::move(v);
      return Stmt(n);
    }
  }

  /*! \brief The reuse mapping */
  ffi::MapObj* opaque_blocks_;
  /*! \brief The range of loops */
  ffi::Map<PrimVar, Range> loop_var2extent_;
  /*! \brief Internal analyzer */
  arith::Analyzer analzyer_;
  /*! \brief Whether or not to simplify unit iterators */
  bool preserve_unit_iters_;
};

class BlockPropertyError : public ScheduleError {
 public:
  /*!
   * \brief Check that all the blocks under the specific stmt have affine bindings
   *     wrt top loop sref and only have data-parallel or reduction block iters
   * \param self The state of the schedule
   * \param sref The sref to the specific stmt
   */
  static void CheckBlockIterTypeAndAffineBinding(const ScheduleState& self, const StmtSRefNode* top,
                                                 const StmtSRefNode* sref) {
    class BlockIterTypeAndAffineBindingChecker : public StmtVisitor {
     public:
      explicit BlockIterTypeAndAffineBindingChecker(const ScheduleState& state,
                                                    const StmtSRefNode* top)
          : state_(state), top_(top) {}

     private:
      void VisitStmt_(const SBlockNode* op) final {
        for (const IterVar& iter_var : op->iter_vars) {
          if (iter_var->iter_type != kDataPar && iter_var->iter_type != kCommReduce) {
            throw BlockPropertyError(state_->mod, ffi::GetRef<SBlock>(op));
          }
          ffi::Optional<StmtSRef> high_exclusive = top_->parent
                                                       ? ffi::GetRef<StmtSRef>(top_->parent)
                                                       : ffi::Optional<StmtSRef>(std::nullopt);
          CheckPartialAffineBinding(state_, ffi::GetRef<SBlock>(op), high_exclusive);
        }
      }
      const ScheduleState& state_;
      const StmtSRefNode* top_;
    };

    BlockIterTypeAndAffineBindingChecker checker(self, top);
    checker(ffi::GetRef<Stmt>(sref->stmt));
  }

  explicit BlockPropertyError(IRModule mod, SBlock block) : mod_(mod), block_(std::move(block)) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: The block under the loops to be reordered have block iter type other "
           "than data-parallel or reduction";
  }

  ffi::String DetailRenderTemplate() const final {
    return "The block {0} under the loops to be reordered have block iter type other than "
           "data-parallel or reduction";
  }

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

  IRModule mod_;
  SBlock block_;
};

class HasAnnotationOrThreadBindingError : public ScheduleError {
 public:
  explicit HasAnnotationOrThreadBindingError(IRModule mod, For loop)
      : mod_(mod), loop_(std::move(loop)) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: The primitive can't be applied because the loop has annotation or "
           "thread binding";
  }

  ffi::String DetailRenderTemplate() const final {
    return "The primitive can't be applied because the loop {0} has annotation or thread binding";
  }

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

  IRModule mod_;
  For loop_;
};

class OuterNotInnerParent : public ScheduleError {
 public:
  explicit OuterNotInnerParent(IRModule mod, For outer, For inner)
      : mod_(mod), outer_(std::move(outer)), inner_(std::move(inner)) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: The outer loop is not the parent of the inner loop";
  }

  ffi::String DetailRenderTemplate() const final {
    return "The loops can't be fused because the outer loop {0} is not the parent of the inner "
           "loop {1}";
  }

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

  IRModule mod_;
  For outer_;
  For inner_;
};

class NotOnlyChildError : public ScheduleError {
 public:
  explicit NotOnlyChildError(IRModule mod, For outer, For inner)
      : mod_(mod), outer_(std::move(outer)), inner_(std::move(inner)) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: The inner loop is not the only child of outer loop";
  }

  ffi::String DetailRenderTemplate() const final {
    return "The loops can't be fused because the inner loop {1} is not the only child of outer "
           "loop {0}.";
  }

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

  IRModule mod_;
  For outer_;
  For inner_;
};

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 WrongFactorProductError : public ScheduleError {
 public:
  explicit WrongFactorProductError(IRModule mod, For loop) : mod_(mod), loop_(std::move(loop)) {}

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

  ffi::String DetailRenderTemplate() const final {
    return "The product of factors is not 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_;
};

class LoopMultiAppearanceError : public ScheduleError {
 public:
  explicit LoopMultiAppearanceError(IRModule mod, For loop) : mod_(mod), loop_(std::move(loop)) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: Some loop appears in the input array for multiple times.";
  }

  ffi::String DetailRenderTemplate() const final {
    return "Loop {0} appears in the input array for multiple times.";
  }

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

  IRModule mod_;
  For loop_;
};

class LoopsNotAChainError : public ScheduleError {
 public:
  enum class ProblemKind { kNotUnderAScope, kHaveNonSingleBranchStmt };

  explicit LoopsNotAChainError(IRModule mod, ffi::Optional<Stmt> problematic_loop, ProblemKind kind)
      : mod_(mod), problematic_loop_(std::move(problematic_loop)), kind_(kind) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: the loops are not in a chain";
  }

  ffi::String DetailRenderTemplate() const final {
    std::stringstream ss;
    ss << "The loops are not in a chain because";
    if (kind_ == ProblemKind::kNotUnderAScope) {
      ss << " they are not under the same scope.";
    } else {
      ss << " there is a non-single-branch stmt in between. Problematic stmt: {0}";
    }
    return ss.str();
  }

  IRModule mod() const final { return mod_; }
  ffi::Array<ffi::ObjectRef> LocationsOfInterest() const final {
    if (kind_ == ProblemKind::kNotUnderAScope) {
      return {};
    } else {
      TVM_FFI_ICHECK(problematic_loop_.has_value());
      return {problematic_loop_.value()};
    }
  }

  IRModule mod_;
  ffi::Optional<Stmt> problematic_loop_;
  ProblemKind kind_;
};

class DependentLoopError : public ScheduleError {
 public:
  enum class PrimitiveKind { kFuse, kReorder };
  explicit DependentLoopError(IRModule mod, For loop, ffi::String inner_var, PrimitiveKind kind)
      : mod_(mod), loop_(std::move(loop)), inner_var_(std::move(inner_var)), kind_(kind) {}

  ffi::String FastErrorString() const final {
    if (kind_ == PrimitiveKind::kReorder) {
      return "ScheduleError: An outer loop's `min` or `extent` is dependent on an inner loop "
             "in the new order";
    } else {
      return "ScheduleError: A loop's `extent` is dependent on another loop";
    }
  }

  ffi::String DetailRenderTemplate() const final {
    if (kind_ == PrimitiveKind::kReorder) {
      return "Outer Loop {0}'s `min` or `extent` is dependent on an inner loop " + inner_var_ +
             " in the new order";
    } else {
      return "A loop {0}'s `extent` is dependent on another loop " + inner_var_;
    }
  }

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

  IRModule mod_;
  For loop_;
  ffi::String inner_var_;
  PrimitiveKind kind_;
};

ffi::Array<StmtSRef> Split(ScheduleState self, const StmtSRef& loop_sref,
                           const ffi::Array<PrimExpr>& factors, bool preserve_unit_iters,
                           bool disable_predication) {
  // Invariance
  // - The total repeat number has not changed for each direct child block with updating predicate.
  // - The execution order has not changed. (The block executes with the same args and the same
  // order with before.
  // Step 1. Check correctness
  const ForNode* loop = TVM_SREF_TO_FOR(loop_sref);
  if (!loop->annotations.empty() || loop->thread_binding.has_value()) {
    throw HasAnnotationOrThreadBindingError(self->mod, ffi::GetRef<For>(loop));
  }
  // Currently, loops not starting with 0 are not supported
  arith::Analyzer analyzer;
  CheckLoopStartsWithZero(self, loop_sref, analyzer.get());

  // Find the most common dtype
  PrimType dtype = PrimType::Int(32);
  {
    int bits = loop->loop_var.ty().bits();
    for (const PrimExpr& factor : factors) {
      bits = std::max(bits, factor.ty().bits());
    }
    dtype = PrimType::Int(bits);
  }
  int n = factors.size();
  PrimExpr substitute_value = IntImm(dtype, 0);
  std::vector<Var> new_loop_vars;
  new_loop_vars.reserve(n);
  for (int i = 0; i < n; i++) {
    const PrimExpr& factor = factors[i];
    Var var = loop->loop_var.CopyWithSuffix("_" + std::to_string(i)).CopyWithDType(dtype);
    substitute_value = substitute_value * factor + var.as_or_throw<PrimExpr>();
    analyzer->Bind(var, Range::FromMinExtent(IntImm(dtype, 0), tvm::cast(dtype, factor)));
    new_loop_vars.emplace_back(std::move(var));
  }
  ffi::Map<SBlock, SBlock> opaque_block_reuse;
  Stmt new_stmt = loop->body;
  new_stmt = SubstituteVarAndCollectOpaqueBlock(
      [&](const Var& v) -> ffi::Optional<Expr> {
        if (v.same_as(loop->loop_var)) {
          return substitute_value;
        } else {
          return std::nullopt;
        }
      },
      &opaque_block_reuse)(std::move(new_stmt));
  // Step 3. Update predicate to guard the loop
  PrimExpr predicate = substitute_value < loop->extent;
  if (!disable_predication &&
      !analyzer->CanProve(predicate, arith::ProofStrength::kSymbolicBound)) {
    new_stmt = BlockPredicateAppender(/*predicate=*/predicate)(std::move(new_stmt));
  }
  // Step 4. Generate nested loops to replace the original loop and simplify the binding
  for (int i = n - 1; i >= 0; i--) {
    new_stmt =
        For(new_loop_vars[i].as_or_throw<PrimVar>(), 0, factors[i], ForKind::kSerial, new_stmt);
  }
  new_stmt = IterMapSimplifyBlockBinding::SimplifyBindings(std::move(new_stmt), GetLoops(loop_sref),
                                                           opaque_block_reuse.CopyOnWrite(),
                                                           preserve_unit_iters);
  self->Replace(loop_sref, new_stmt, opaque_block_reuse);
  ffi::Array<StmtSRef> result_srefs;
  result_srefs.reserve(n);
  for (int i = 0; i < n; i++) {
    result_srefs.push_back(self->stmt2ref.at(new_stmt.get()));
    const ForNode* outer_loop = TVM_TYPE_AS(new_stmt, ForNode);
    new_stmt = outer_loop->body;
  }
  return result_srefs;
}

class BufferIndicesMapExtractor : public StmtExprVisitor {
 public:
  explicit BufferIndicesMapExtractor(Var loop_var) : loop_var_(loop_var) {}

  static ffi::Map<ffi::String, ffi::Array<ffi::String>> Extract(Var loop_var, SBlock& block) {
    BufferIndicesMapExtractor extractor(loop_var);
    extractor(std::move(block->body));
    return extractor.buffer_indices_map;
  }

 private:
  void VisitStmt_(const BufferStoreNode* store) final {
    ffi::Array<ffi::String> indices;
    bool check_ = false;
    for (size_t i = 0; i < store->indices.size(); i++) {
      auto var = store->indices[i].as<PrimVar>();
      if (!var.has_value()) {
        check_ = true;
        break;
      }
      indices.push_back(var.value()->name);
    }
    if (buffer_indices_map.find(store->buffer.name()) == buffer_indices_map.end() && !check_)
      buffer_indices_map.Set(store->buffer.name(), indices);
    StmtExprVisitor::VisitStmt_(store);
  }

  void VisitExpr_(const BufferLoadNode* load) final {
    ffi::Array<ffi::String> indices;
    bool check_ = false;
    for (size_t i = 0; i < load->indices.size(); i++) {
      auto var = load->indices[i].as<PrimVar>();
      if (!var.has_value()) {
        check_ = true;
        break;
      }
      indices.push_back(var.value()->name);
    }
    if (buffer_indices_map.find(load->buffer.name()) == buffer_indices_map.end() && !check_)
      buffer_indices_map.Set(load->buffer.name(), indices);
    StmtExprVisitor::VisitExpr_(load);
  }

  void VisitStmt_(const SBlockNode* op) final { StmtVisitor::VisitStmt_(op); }

  Var loop_var_;
  ffi::Map<ffi::String, ffi::Array<ffi::String>> buffer_indices_map;
};

ffi::Array<BufferRegion> MutateBufferRegion(
    ffi::Map<ffi::String, ffi::Array<ffi::String>> buffer_indices_map,
    ffi::Map<ffi::String, Range> index_range_map, ffi::Array<BufferRegion> region_arr) {
  // Update the region with new Ranges and return new BufferRegion
  ffi::Array<BufferRegion> new_region_arr =
      MutateArray(region_arr, [&buffer_indices_map, &index_range_map](const BufferRegion& region) {
        BufferRegion new_region = region;
        auto it = buffer_indices_map.find(new_region->buffer.name());
        if (it == buffer_indices_map.end()) return new_region;

        ffi::Array<ffi::String> old_indices = buffer_indices_map[new_region->buffer.name()];
        ffi::Array<Range> new_ranges;
        for (size_t i = 0; i < old_indices.size(); i++) {
          new_ranges.push_back(index_range_map[old_indices[i]]);
        }
        new_region.CopyOnWrite()->region = std::move(new_ranges);
        return new_region;
      });
  return new_region_arr;
}

class BlockMutator : public StmtExprMutator {
 public:
  explicit BlockMutator(Var new_loop_var, PrimExpr min, PrimExpr extent)
      : new_loop_var_(new_loop_var), min_(min), extent_(extent) {}

 private:
  Stmt VisitStmt_(const SBlockNode* _op) final {
    SBlock new_block = StmtMutator::VisitStmt_(_op).as_or_throw<SBlock>();

    // If iter_vars.size() is 0, then the block most probably be an Opaque block
    if (new_block->iter_vars.size() == 0 || inner_iter_var_index == -1) {
      new_block.CopyOnWrite()->name_hint =
          new_block.CopyOnWrite()->name_hint + "_" + new_loop_var_->name;
      return new_block;
    }

    Var iter_var_ = new_block->iter_vars[inner_iter_var_index]->var;
    inner_iter_var_index = -1;
    // As we are working on cloned block, we need to create new instances of iter_var
    ffi::Array<IterVar> new_iter_vars =
        MutateArray(new_block->iter_vars, [this, &iter_var_](const IterVar& iter) {
          auto dtype = iter->var.ty();
          // Create new Var instance for each IterVar
          Var new_var = Var(iter->var->name, iter->var.ty());
          IterVar new_iter = iter;
          new_iter.CopyOnWrite()->var = new_var.as_or_throw<PrimVar>();
          // Change the domain of IterVar corresponding to partitioned loop_var
          if (iter_var_.same_as(iter->var)) {
            new_iter.CopyOnWrite()->dom = Range(tvm::cast(dtype, min_), tvm::cast(dtype, extent_));
          }
          return new_iter;
        });

    // Update the IterVars of new_block
    if (!new_block->iter_vars.same_as(new_iter_vars)) {
      new_block.CopyOnWrite()->iter_vars = std::move(new_iter_vars);
      new_block.CopyOnWrite()->name_hint =
          new_block.CopyOnWrite()->name_hint + "_" + new_loop_var_->name;
    }

    // Get the (iter_var, new Range) map
    ffi::Map<ffi::String, Range> index_range_map;
    for (size_t i = 0; i < new_block->iter_vars.size(); i++) {
      IterVar iter = new_block->iter_vars[i];
      index_range_map.Set(iter->var->name, iter->dom);
    }

    // Get the (BufferVar, indices) map
    ffi::Map<ffi::String, ffi::Array<ffi::String>> buffer_indices_map =
        BufferIndicesMapExtractor::Extract(new_loop_var_, new_block);
    ffi::Array<BufferRegion> new_writes =
        MutateBufferRegion(buffer_indices_map, index_range_map, new_block->writes);
    if (!new_block->writes.same_as(new_writes)) {
      // Update the writes with new_writes
      new_block.CopyOnWrite()->writes = std::move(new_writes);
    }
    ffi::Array<BufferRegion> new_reads =
        MutateBufferRegion(buffer_indices_map, index_range_map, new_block->reads);
    if (!new_block->reads.same_as(new_reads)) {
      // Update the reads with new_reads
      new_block.CopyOnWrite()->reads = std::move(new_reads);
    }

    ffi::Map<Var, Var> var_map;
    for (size_t i = 0; i < new_block->iter_vars.size(); i++) {
      var_map.Set(_op->iter_vars[i]->var, new_block->iter_vars[i]->var);
    }

    // Update all instances of old iter_vars in the block with new iter_vars
    auto block_stmt = tirx::Substitute(new_block, var_map);
    return block_stmt;
  }

  Stmt VisitStmt_(const SBlockRealizeNode* realize) final {
    ffi::Array<PrimExpr> iter_values = realize->iter_values;
    for (size_t i = 0; i < iter_values.size(); i++) {
      if (new_loop_var_.same_as(iter_values[i])) {
        // Get the iter_var index corresponding to loop_var iter_value index
        inner_iter_var_index = i;
        break;
      }
    }
    SBlockRealize stmt = StmtExprMutator::VisitStmt_(realize).as_or_throw<SBlockRealize>();
    return stmt;
  }

  Stmt VisitStmt_(const ForNode* op) final {
    For res = StmtMutator::VisitStmt_(op).as_or_throw<For>();
    Var new_var = Var(op->loop_var->name, op->loop_var.ty());

    if (!op->loop_var.same_as(new_var)) {
      // If the partioned loop contains nested for loop, then create new iteration variable instance
      res.CopyOnWrite()->body = tirx::Substitute(res->body, {{op->loop_var, new_var}});
      res.CopyOnWrite()->loop_var = new_var.as_or_throw<PrimVar>();
    }
    return res;
  }

  Var new_loop_var_;
  PrimExpr min_, extent_;
  int inner_iter_var_index = -1;
};

const ffi::String get_sblock_name(Stmt loop_body) {
  const SBlockRealizeNode* blk_realize = loop_body.as<SBlockRealizeNode>();
  if (blk_realize == nullptr) {
    return get_sblock_name(loop_body.as<ForNode>()->body);
  }
  return blk_realize->block->name_hint;
}

ffi::Array<StmtSRef> LoopPartition(ScheduleState self, const StmtSRef& loop_sref,
                                   const ffi::Array<PrimExpr>& factors, bool preserve_unit_iters) {
  const ForNode* loop = TVM_SREF_TO_FOR(loop_sref);
  if (!loop->annotations.empty() || loop->thread_binding.has_value()) {
    throw HasAnnotationOrThreadBindingError(self->mod, ffi::GetRef<For>(loop));
  }

  arith::Analyzer analyzer;
  // Find the most common dtype
  PrimType dtype = PrimType::Int(32);
  {
    int bits = loop->loop_var.ty().bits();
    for (const PrimExpr& factor : factors) {
      bits = std::max(bits, factor.ty().bits());
    }
    dtype = PrimType::Int(bits);
  }

  ffi::String block_name = get_sblock_name(loop->body) + "_" + loop->loop_var->name;
  int n = factors.size();
  PrimExpr min_value = loop->min;
  PrimExpr extent_value;

  ffi::Array<Stmt> block_partitions;
  block_partitions.reserve(n);

  // Iterate over each pair of factors and create partition
  for (int i = 0; i < n; i++) {
    extent_value = analyzer->Simplify(factors[i]);
    Var new_loop_var = loop->loop_var.CopyWithSuffix(std::to_string(i)).CopyWithDType(dtype);
    Stmt loop_body = tirx::Substitute(loop->body, {{loop->loop_var, new_loop_var}});

    // Create new block with new reference to each variable/stmt/expr in the existing block
    loop_body = BlockMutator(new_loop_var, min_value, extent_value)(std::move(loop_body));
    // Create new for loop with appropriate range
    auto for_node = For(new_loop_var.as_or_throw<PrimVar>(), min_value, extent_value - min_value,
                        ForKind::kSerial, loop_body);

    const auto& partition_block_name = block_name + std::to_string(i) + "_partition";
    // Create partition_block for the partitioned for loop
    SBlockRealize partition_block({}, extent_value > 0,
                                  SBlock({}, {}, {}, partition_block_name, for_node));
    block_partitions.push_back(partition_block);

    min_value = extent_value;
  }

  // Create common block with all the partitioned blocks as its children blocks
  SBlockRealize common({}, IntImm::Bool(true),
                       SBlock({}, {}, {}, block_name + "_common", tirx::SeqStmt(block_partitions)));

  // Replace existing loop with the newly created common block
  self->Replace(loop_sref, common, {});
  StmtSRef scope_sref = self->stmt2ref.at(common->block.get());
  StmtSRef scope_root = GetScopeRoot(self, scope_sref, /*require_stage_pipeline=*/false);
  bool scope_block_affine_binding = self->IsAffineBlockBinding(scope_root);
  // Update the SRefTree for the newly created common block
  self->UpdateScopeSBlockInfo(GetSBlockRealize(self, scope_root));
  self->block_info[scope_root].affine_binding = scope_block_affine_binding;

  // Collect the SRef for each partitioned loop and return
  ffi::Array<StmtSRef> partition_srefs;
  partition_srefs.reserve(n);
  for (int i = 0; i < n; i++) {
    StmtSRef partition_loop_sref =
        self->stmt2ref.at(block_partitions[i].as<SBlockRealizeNode>()->block->body.get());
    partition_srefs.push_back(partition_loop_sref);
  }
  return partition_srefs;
}

class LoopReconstructor : private StmtMutator {
 public:
  explicit LoopReconstructor(SBlock scope_root, const std::vector<std::vector<For>>& loops)
      : scope_root_(scope_root), loops_(loops) {}

  using StmtMutator::operator();

  /*!
   * \brief Create the new nest loops induced by the given loops
   */
  void MakeNewLoop() {
    ffi::Array<Var> new_loop_vars;
    ffi::Array<PrimExpr> new_loop_extents;
    ffi::Array<Stmt> new_stmts;
    for (size_t i = 0; i < loops_.size(); i++) {
      ffi::Map<Var, PrimExpr> var_map;
      for (size_t j = 0; j < loops_[i].size(); j++) {
        if (i == 0) {
          Var merged_var = loops_[i][j]->loop_var.CopyWithSuffix("_m");
          new_loop_vars.push_back(merged_var);
          new_loop_extents.push_back(loops_[i][j]->extent);
        }
        var_map.Set(loops_[i][j]->loop_var, new_loop_vars[j].as_or_throw<PrimExpr>());
      }
      auto new_stmt = Substitute(loops_[i][0]->body, var_map);
      new_stmts.push_back(new_stmt);
      this->need_remove_loop_.push_back(loops_[i].back());
    }
    auto new_loop = For(new_loop_vars[0].as_or_throw<PrimVar>(), IntImm::Int32(0),
                        new_loop_extents[0], ForKind::kSerial, SeqStmt(std::move(new_stmts)));
    this->new_inner_loop_ = new_loop;
    for (size_t i = 1; i < new_loop_vars.size(); ++i) {
      const Var& loop_var = new_loop_vars[i];
      const PrimExpr& loop_extent = new_loop_extents[i];
      new_loop = For(loop_var.as_or_throw<PrimVar>(), IntImm::Int32(0), loop_extent,
                     ForKind::kSerial, new_loop);
    }
    this->new_outer_loop_ = new_loop;
  }

 private:
  Stmt VisitStmt_(const SBlockNode* block) final {
    if (block != scope_root_.get()) {
      return ffi::GetRef<SBlock>(block);
    }
    return StmtMutator::VisitStmt_(block);
  }

  Stmt VisitStmt_(const ForNode* loop) final {
    if (ffi::GetRef<For>(loop) == need_remove_loop_.back()) {
      return new_outer_loop_;
    } else if (std::count(need_remove_loop_.begin(), need_remove_loop_.end(),
                          ffi::GetRef<For>(loop))) {
      return Evaluate(0);
    }
    return StmtMutator::VisitStmt_(loop);
  }

  Stmt VisitStmt_(const SeqStmtNode* seq_stmt) final {
    auto ret = StmtMutator::VisitSeqStmt_(seq_stmt, true).as_or_throw<SeqStmt>();
    ffi::Array<Stmt> filtered;
    for (Stmt stmt : ret->seq) {
      if (!is_no_op(stmt)) {
        filtered.push_back(std::move(stmt));
      }
    }
    ret = SeqStmt(filtered);
    if (ret->size() == 0) {
      return Evaluate(0);
    } else if (ret->size() == 1) {
      return ret->seq[0];
    } else {
      return ret;
    }
  }

 public:
  /*! \brief The root block of the block scope */
  SBlock scope_root_;
  /*! \brief The given loops to be merge */
  const std::vector<std::vector<For>>& loops_;
  /*! \brief The outermost new loop to replace the original loop */
  For new_outer_loop_{nullptr};
  /*! \brief The innermost new loop to replace the original loop */
  For new_inner_loop_{nullptr};
  /*! \brief The loops to be removed */
  std::vector<For> need_remove_loop_;
};

StmtSRef Merge(ScheduleState self, const ffi::Array<StmtSRef>& loop_srefs) {
  // Invariance
  // - The total repeat number has not changed for each direct child block.
  // - The execution order has not changed. (The block executes with the same
  //   args and the same order with before.)
  arith::Analyzer analyzer;
  StmtSRef scope_root_sref;
  StmtSRef lca = GetSRefLowestCommonAncestor(loop_srefs);
  std::vector<std::vector<For>> lca_nest_loops;
  // Step 1. check correctness
  std::vector<For> nest_loop_loops;
  std::vector<PrimExpr> nest_loop_extents;
  for (size_t i = 0; i < loop_srefs.size(); i++) {
    const StmtSRef& sref = loop_srefs[i];
    auto scope_root_sref_ = GetScopeRoot(self, sref, /*require_stage_pipeline=*/false);
    std::vector<PrimExpr> nest_loop_i_extents;
    std::vector<For> nest_loop_i_loops;
    for (auto p = sref.get(); p != lca.get(); p = p->parent) {
      if (auto loop = p->StmtAs<ForNode>()) {
        if (!loop->annotations.empty() || loop->thread_binding.has_value()) {
          throw HasAnnotationOrThreadBindingError(self->mod, ffi::GetRef<For>(loop));
        }
        CheckLoopStartsWithZero(self, ffi::GetRef<StmtSRef>(p), analyzer.get());
        nest_loop_i_loops.push_back(ffi::GetRef<For>(loop));
        nest_loop_i_extents.push_back(loop->extent);
      }
    }
    lca_nest_loops.push_back(nest_loop_i_loops);
    const ForNode* outer_loop = nullptr;
    for (auto iter = nest_loop_i_loops.rbegin(); iter != nest_loop_i_loops.rend(); ++iter) {
      if (outer_loop && !outer_loop->body.same_as(*iter)) {
        throw NotOnlyChildError(self->mod, ffi::GetRef<For>(outer_loop), *iter);
      }
      outer_loop = (*iter).get();
    }
    if (i == 0) {
      scope_root_sref = scope_root_sref_;
      nest_loop_loops = nest_loop_i_loops;
      nest_loop_extents = nest_loop_i_extents;
    } else {
      if (scope_root_sref_.get() != scope_root_sref.get()) {
        TVM_FFI_THROW(ScheduleError) << "Expected the loops to be under the same block scope.";
        throw;
      }
      if (nest_loop_i_extents.size() != nest_loop_extents.size()) {
        TVM_FFI_THROW(ScheduleError) << "Merge loop's nesting depth must be same, but not.";
        throw;
      } else {
        for (size_t j = 0; j < nest_loop_i_extents.size(); j++) {
          if (!analyzer->CanProveEqual(nest_loop_i_extents[j], nest_loop_extents[j])) {
            TVM_FFI_THROW(ScheduleError) << "Merge loop's `extent` must be same, but not."
                                         << " extent=[" << j << "," << nest_loop_extents[j] << ","
                                         << nest_loop_i_extents[j] << "]";
            throw;
          }
        }
      }
    }
  }
  // Step 2. Create merged loops and replace the original loops
  SBlock scope_root = ffi::GetRef<SBlock>(scope_root_sref->StmtAs<SBlockNode>());
  LoopReconstructor reconstructor(scope_root, lca_nest_loops);
  reconstructor.MakeNewLoop();
  SBlock new_scope_root = reconstructor(scope_root).as_or_throw<SBlock>();
  // Step 3. Do the actual replacement
  self->Replace(scope_root_sref, new_scope_root, {{scope_root, new_scope_root}});
  return self->stmt2ref.at(reconstructor.new_inner_loop_.get());
}

StmtSRef Fuse(ScheduleState self, const ffi::Array<StmtSRef>& loop_srefs,
              bool preserve_unit_iters) {
  // Invariance
  // - The total repeat number has not changed for each direct child block.
  // - The execution order has not changed. (The block executes with the same
  //   args and the same order with before.)
  std::vector<const ForNode*> loops;
  loops.reserve(loop_srefs.size());
  StmtSRef outer_loop_sref{nullptr};
  const ForNode* outer_loop = nullptr;
  arith::Analyzer analyzer;
  std::unordered_set<const VarNode*> outer_loop_vars;
  // Step 1. check correctness
  for (const StmtSRef& sref : loop_srefs) {
    const ForNode* loop = TVM_SREF_TO_FOR(sref);
    if (!loop->annotations.empty() || loop->thread_binding.has_value()) {
      throw HasAnnotationOrThreadBindingError(self->mod, ffi::GetRef<For>(loop));
    }
    if (outer_loop_sref.defined()) {
      if (sref->parent != outer_loop_sref.get()) {
        throw OuterNotInnerParent(self->mod, ffi::GetRef<For>(outer_loop), ffi::GetRef<For>(loop));
      }
      if (!outer_loop->body.same_as(ffi::GetRef<For>(loop))) {
        throw NotOnlyChildError(self->mod, ffi::GetRef<For>(outer_loop), ffi::GetRef<For>(loop));
      }
    }
    outer_loop_sref = sref;
    outer_loop = loop;
    CheckLoopStartsWithZero(self, sref, analyzer.get());
    const VarNode* used_var = nullptr;
    auto f_contain = [&outer_loop_vars, &used_var](const VarNode* var) {
      if (outer_loop_vars.count(var)) {
        used_var = var;
        return true;
      }
      return false;
    };
    if (UsesVar(loop->extent, f_contain)) {
      throw DependentLoopError(self->mod, ffi::GetRef<For>(loop), used_var->name,
                               DependentLoopError::PrimitiveKind::kFuse);
    }
    outer_loop_vars.insert(loop->loop_var.get());
    loops.push_back(loop);
  }
  // Step 2. Create fused loop var and replace the original loop vars
  std::string suffix;
  int n = loops.size();
  int bits = loops[0]->loop_var.ty().bits();
  for (int i = 1; i < n; i++) {
    suffix += "_" + loops[i]->loop_var->name;
    bits = std::max(bits, loops[i]->loop_var.ty().bits());
  }
  suffix += "_fused";

  Var fused_var = loops[0]->loop_var.CopyWithSuffix(suffix).CopyWithDType(PrimType::Int(bits));
  ffi::Array<PrimExpr> substitute_value;
  substitute_value.resize(loops.size());
  PrimExpr lower = 1;
  for (int i = static_cast<int>(loops.size()) - 1; i > 0; i--) {
    PrimExpr next_lower = analyzer->canonical_simplify(loops[i]->extent * lower);
    substitute_value.Set(
        i, is_one(loops[i]->extent)
               ? PrimExpr(0)
               : floordiv(floormod(fused_var.as_or_throw<PrimExpr>(), next_lower), lower));
    lower = next_lower;
  }
  substitute_value.Set(0, is_one(loops[0]->extent)
                              ? PrimExpr(0)
                              : floordiv(fused_var.as_or_throw<PrimExpr>(), lower));
  Stmt new_stmt = loops.back()->body;
  ffi::Map<SBlock, SBlock> opaque_block_reuse;
  auto f_substitute = [&](const Var& v) -> ffi::Optional<Expr> {
    for (int i = 0; i < n; i++) {
      if (v.same_as(loops[i]->loop_var)) {
        return substitute_value[i];
      }
    }
    return std::nullopt;
  };
  new_stmt =
      SubstituteVarAndCollectOpaqueBlock(f_substitute, &opaque_block_reuse)(std::move(new_stmt));
  // Step 3. Generate a loop to replace the original loops
  PrimExpr fused_extent = 1;
  for (int i = 0; i < n; i++) {
    fused_extent *= loops[i]->extent;
  }
  fused_extent = analyzer->Simplify(fused_extent);
  new_stmt = For(fused_var.as_or_throw<PrimVar>(), 0, fused_extent, ForKind::kSerial, new_stmt);
  new_stmt = IterMapSimplifyBlockBinding::SimplifyBindings(
      std::move(new_stmt), GetLoops(loop_srefs[0]), opaque_block_reuse.CopyOnWrite(),
      preserve_unit_iters);
  self->Replace(loop_srefs[0], new_stmt, opaque_block_reuse);
  return self->stmt2ref.at(new_stmt.get());
}

/*!
 * \brief Collect an array of loop srefs into a set
 * \param self The schedule state
 * \param ordered_loop_srefs The array of loop srefs
 * \return A set containing all loops in the array
 * \throws ScheduleError If there are duplicate loops in the array
 */
std::unordered_set<const StmtSRefNode*> CollectLoopsIntoSet(
    const ScheduleState& self, const ffi::Array<StmtSRef>& ordered_loop_srefs) {
  std::unordered_set<const StmtSRefNode*> loop_srefs;
  loop_srefs.reserve(ordered_loop_srefs.size());
  for (const StmtSRef& loop_sref : ordered_loop_srefs) {
    auto inserted = loop_srefs.insert(loop_sref.get());
    if (!inserted.second) {
      const ForNode* loop = TVM_SREF_TO_FOR(loop_sref);
      throw LoopMultiAppearanceError(self->mod, ffi::GetRef<For>(loop));
    }
  }
  return loop_srefs;
}

/*!
 * \brief Get the top and bottom boundary of reorder range (which should be a chain)
 * \param self The schedule state
 * \param loop_srefs The set containing the srefs to the loops to be reordered
 * \return A pair containing the top and bottom boundary of the reorder range
 * \throws ScheduleError If the loops to be reordered is not in a chain
 */
std::pair<const StmtSRefNode*, const StmtSRefNode*> GetBoundaryOfReorderRange(
    const ScheduleState& self, const std::unordered_set<const StmtSRefNode*>& loop_srefs) {
  const StmtSRefNode* top = nullptr;
  const StmtSRefNode* bottom = *loop_srefs.begin();
  std::unordered_set<const StmtSRefNode*> visited;
  bool scope_block_visited = false;
  bool first_traversal = true;
  for (const StmtSRefNode* loop_sref : loop_srefs) {
    if (visited.count(loop_sref)) {
      continue;
    }
    for (const StmtSRefNode* v = loop_sref;; v = v->parent) {
      // Case 1. If `v` corresponds to a block, stop traversal.
      if (v->stmt->IsInstance<SBlockNode>()) {
        if (scope_block_visited) {
          throw LoopsNotAChainError(self->mod, std::nullopt,
                                    LoopsNotAChainError::ProblemKind::kNotUnderAScope);
        }
        scope_block_visited = true;
        break;
      }
      // Case 2. If `v` corresponds to a previously-visited loop, stop traversal and update
      // `bottom`.
      if (visited.count(v)) {
        if (v != bottom) {
          throw LoopsNotAChainError(self->mod, ffi::GetRef<Stmt>(v->stmt),
                                    LoopsNotAChainError::ProblemKind::kHaveNonSingleBranchStmt);
        }
        bottom = loop_sref;
        break;
      }
      // Case 3. Add `v` into `visited`
      visited.insert(v);
      // If it's the first traversal and the loop corresponding to `v` is in the input array,
      // update `top`.
      if (first_traversal && loop_srefs.count(v)) {
        top = v;
      }
    }
    first_traversal = false;
  }
  return std::make_pair(top, bottom);
}

/*!
 * \brief Get all the loops in the reorder range
 * \param self The schedule state
 * \param top The top boundary of the reorder range
 * \param bottom The bottom boundary of the reorder range
 * \return An array containing all the loops in the reorder range
 * \throws ScheduleError If some loop in the reorder range is not single-branch
 */
std::vector<const StmtSRefNode*> GetLoopsInReorderRange(const ScheduleState& self,
                                                        const StmtSRefNode* top,
                                                        const StmtSRefNode* bottom) {
  std::vector<const StmtSRefNode*> chain;
  for (const StmtSRefNode* loop_sref = bottom; loop_sref != top;) {
    const StmtSRefNode* parent_loop_sref = loop_sref->parent;
    const ForNode* outer = parent_loop_sref->StmtAs<ForNode>();
    const ForNode* inner = loop_sref->StmtAs<ForNode>();
    TVM_FFI_ICHECK(outer != nullptr && inner != nullptr);
    if (outer->body.get() != inner) {
      throw LoopsNotAChainError(self->mod, ffi::GetRef<For>(outer),
                                LoopsNotAChainError::ProblemKind::kHaveNonSingleBranchStmt);
    }
    chain.push_back(loop_sref);
    loop_sref = parent_loop_sref;
  }
  chain.push_back(top);
  return chain;
}

/*!
 * \brief Construct a loop chain in the new order
 * \param self The schedule state
 * \param chain The loops in the reorder range
 * \param ordered_loop_srefs The loop srefs to be reordered
 * \param loop_srefs The set containing loop srefs to be reordered
 * \return The new loop chain
 * \throws ScheduleError If the domain of an outer loop depends on any of the inner loops after
 * reordering
 */
For ConstructNewLoopChain(const ScheduleState& self, std::vector<const StmtSRefNode*> chain,
                          const ffi::Array<StmtSRef>& ordered_loop_srefs,
                          const std::unordered_set<const StmtSRefNode*>& loop_srefs) {
  std::unordered_set<const VarNode*> inner_vars;
  inner_vars.reserve(chain.size());
  For new_loop{nullptr};
  int index = static_cast<int>(ordered_loop_srefs.size()) - 1;
  for (const StmtSRefNode* loop_sref : chain) {
    const ForNode* copy = nullptr;
    if (loop_srefs.count(loop_sref)) {
      copy = ordered_loop_srefs[index]->StmtAs<ForNode>();
      --index;
    } else {
      copy = loop_sref->StmtAs<ForNode>();
    }
    TVM_FFI_ICHECK(copy != nullptr);
    ffi::ObjectPtr<ForNode> n = ffi::make_object<ForNode>(*copy);
    if (new_loop.defined()) {
      n->body = new_loop;
    } else {
      n->body = loop_sref->StmtAs<ForNode>()->body;
    }
    const VarNode* used_var = nullptr;
    auto f_contain = [&inner_vars, &used_var](const VarNode* var) {
      if (inner_vars.count(var)) {
        used_var = var;
        return true;
      }
      return false;
    };
    if (UsesVar(copy->min, f_contain) || UsesVar(copy->extent, f_contain)) {
      throw DependentLoopError(self->mod, ffi::GetRef<For>(copy), used_var->name,
                               DependentLoopError::PrimitiveKind::kReorder);
    }
    inner_vars.insert(copy->loop_var.get());
    new_loop = For(std::move(n));
  }
  return new_loop;
}

void Reorder(ScheduleState self, const ffi::Array<StmtSRef>& ordered_loop_srefs) {
  if (ordered_loop_srefs.size() <= 1) {
    return;
  }
  // Step 1. Check uniqueness and collect the input loop srefs into a set
  std::unordered_set<const StmtSRefNode*> loop_srefs =
      CollectLoopsIntoSet(self, ordered_loop_srefs);
  // Step 2. Gather loops to be reordered
  // For each loop sref in the input sref array, traverse upwards along its parent pointer in the
  // sref tree, and stop on either a block, or a previously-visited loop
  // - the top of the reorder range is the last loop visited in the first traversal which exists in
  //   the input array
  // - the bottom of the reorder range is the last loop in the input array which is not visited in
  // the previous traversals
  auto [top, bottom] = GetBoundaryOfReorderRange(self, loop_srefs);
  // Step 3. Collect all loops in the chain and check the loops are single-branch
  std::vector<const StmtSRefNode*> chain = GetLoopsInReorderRange(self, top, bottom);
  // Step 4. Check the block below has all its block_var to be data-parallel or reduction,
  // and the block has an affine binding wrt top of the loop range.
  BlockPropertyError::CheckBlockIterTypeAndAffineBinding(self, top, bottom);
  // Step 5. Replace the original loops with the reordered loops and check that outer loop is
  // not dependent on inner loop
  For new_loop = ConstructNewLoopChain(self, std::move(chain), ordered_loop_srefs, loop_srefs);
  self->Replace(ffi::GetRef<StmtSRef>(top), new_loop, {});
}

StmtSRef AddUnitLoop(ScheduleState self, StmtSRef sref) {
  if (sref->stmt->IsInstance<ForNode>()) {
    For new_loop =
        For(PrimVar("u", PrimType::Int(32)), 0, 1, ForKind::kSerial, ffi::GetRef<Stmt>(sref->stmt));
    self->Replace(sref, new_loop, {});
    return self->stmt2ref.at(new_loop.get());
  }
  class NewLoopCreator : public StmtMutator {
   public:
    explicit NewLoopCreator(const StmtNode* src_block) : src_block_(src_block) {}

    Stmt VisitStmt_(const SBlockRealizeNode* realize) final {
      if (realize->block.get() == src_block_) {
        new_loop_ = For(PrimVar("u", PrimType::Int(32)), 0, 1, ForKind::kSerial,
                        ffi::GetRef<SBlockRealize>(realize));
        return new_loop_;
      }
      return StmtMutator::VisitStmt_(realize);
    }

    const StmtNode* src_block_;
    For new_loop_{nullptr};
  };

  TVM_FFI_CHECK(sref->parent != nullptr, ValueError) << "Cannot add loops on top of the root block";
  StmtSRef parent_sref = ffi::GetRef<StmtSRef>(sref->parent);
  NewLoopCreator creator(sref->stmt);
  Stmt new_stmt = creator(ffi::GetRef<Stmt>(parent_sref->stmt));
  if (new_stmt->IsInstance<ForNode>()) {
    self->Replace(parent_sref, std::move(new_stmt), {});
  } else {
    SBlock old_parent_block = ffi::GetRef<SBlock>(parent_sref->StmtAs<SBlockNode>());
    SBlock new_parent_block = new_stmt.as_or_throw<SBlock>();
    self->Replace(parent_sref, new_stmt, {{old_parent_block, new_parent_block}});
  }
  return self->stmt2ref.at(creator.new_loop_.get());
}

/******** InstructionKind Registration ********/

struct SplitTraits : public UnpackedInstTraits<SplitTraits> {
  static constexpr const char* kName = "Split";
  static constexpr bool kIsPure = false;

 private:
  static constexpr size_t kNumInputs = 2;
  static constexpr size_t kNumAttrs = 2;
  static constexpr size_t kNumDecisions = 0;

  template <size_t delta>
  TVM_FFI_INLINE static void _SetInputs(AnyView* packed_args, const ffi::Array<Any>& inputs) {
    thread_local Any loop_rv{nullptr};
    thread_local ffi::Array<Any> factors{nullptr};
    loop_rv = inputs[0];
    factors = ffi::Array<Any>{inputs.begin() + 1, inputs.end()};
    packed_args[delta] = loop_rv;
    packed_args[delta + 1] = factors;
  }

  static ffi::Array<LoopRV> UnpackedApplyToSchedule(Schedule sch, LoopRV loop_rv,
                                                    ffi::Array<ffi::Optional<ExprRV>> factors,
                                                    IntImm preserve_unit_iters,
                                                    IntImm disable_predication) {
    return sch->Split(loop_rv, factors, preserve_unit_iters->value != 0,
                      disable_predication->value != 0);
  }

  static ffi::String UnpackedAsPython(ffi::Array<ffi::String> outputs, ffi::String loop_rv,
                                      ffi::Array<Any> factors, IntImm preserve_unit_iters,
                                      IntImm disable_predication) {
    PythonAPICall py("split");
    py.Input("loop", loop_rv);
    py.Input("factors", factors);
    py.Input("preserve_unit_iters", preserve_unit_iters->value != 0);
    py.Input("disable_predication", disable_predication->value != 0);
    py.OutputList(outputs);
    return py.Str();
  }

  template <typename>
  friend struct ::tvm::s_tir::UnpackedInstTraits;
};

struct LoopPartitionTraits : public UnpackedInstTraits<LoopPartitionTraits> {
  static constexpr const char* kName = "LoopPartition";
  static constexpr bool kIsPure = false;

 private:
  static constexpr size_t kNumInputs = 2;
  static constexpr size_t kNumAttrs = 1;
  static constexpr size_t kNumDecisions = 0;

  template <size_t delta>
  TVM_FFI_INLINE static void _SetInputs(AnyView* packed_args, const ffi::Array<Any>& inputs) {
    thread_local Any loop_rv{nullptr};
    thread_local ffi::Array<Any> factors{nullptr};
    loop_rv = inputs[0];
    factors = ffi::Array<Any>{inputs.begin() + 1, inputs.end()};
    packed_args[delta] = loop_rv;
    packed_args[delta + 1] = factors;
  }

  static ffi::Array<LoopRV> UnpackedApplyToSchedule(Schedule sch, LoopRV loop_rv,
                                                    ffi::Array<ffi::Optional<ExprRV>> factors,
                                                    IntImm preserve_unit_iters) {
    return sch->LoopPartition(loop_rv, factors, preserve_unit_iters->value != 0);
  }

  static ffi::String UnpackedAsPython(ffi::Array<ffi::String> outputs, ffi::String loop_rv,
                                      ffi::Array<Any> factors, IntImm preserve_unit_iters) {
    PythonAPICall py("loop_partition");
    py.Input("loop", loop_rv);
    py.Input("factors", factors);
    py.Input("preserve_unit_iters", preserve_unit_iters->value != 0);
    py.OutputList(outputs);
    return py.Str();
  }

  template <typename>
  friend struct ::tvm::s_tir::UnpackedInstTraits;
};

struct MergeTraits : public UnpackedInstTraits<MergeTraits> {
  static constexpr const char* kName = "Merge";
  static constexpr bool kIsPure = false;

 private:
  static constexpr size_t kNumInputs = 1;
  static constexpr size_t kNumAttrs = 0;
  static constexpr size_t kNumDecisions = 0;

  template <size_t delta>
  TVM_FFI_INLINE static void _SetInputs(AnyView* packed_args, const ffi::Array<Any>& inputs) {
    packed_args[delta] = inputs;
  }

  static LoopRV UnpackedApplyToSchedule(Schedule sch, ffi::Array<LoopRV> loop_rvs) {
    return sch->Merge(loop_rvs);
  }

  static ffi::String UnpackedAsPython(ffi::Array<ffi::String> outputs,
                                      ffi::Array<ffi::String> loop_rvs) {
    PythonAPICall py("merge");
    for (const ffi::String& loop_rv : loop_rvs) {
      py.Input("", loop_rv);
    }
    py.SingleOutput(outputs);
    return py.Str();
  }

  template <typename>
  friend struct ::tvm::s_tir::UnpackedInstTraits;
};

struct FuseTraits : public UnpackedInstTraits<FuseTraits> {
  static constexpr const char* kName = "Fuse";
  static constexpr bool kIsPure = false;

 private:
  static constexpr size_t kNumInputs = 1;
  static constexpr size_t kNumAttrs = 1;
  static constexpr size_t kNumDecisions = 0;

  template <size_t delta>
  TVM_FFI_INLINE static void _SetInputs(AnyView* packed_args, const ffi::Array<Any>& inputs) {
    packed_args[delta] = inputs;
  }

  static LoopRV UnpackedApplyToSchedule(Schedule sch, ffi::Array<LoopRV> loop_rvs,
                                        IntImm preserve_unit_iters) {
    return sch->Fuse(loop_rvs, preserve_unit_iters->value != 0);
  }

  static ffi::String UnpackedAsPython(ffi::Array<ffi::String> outputs,
                                      ffi::Array<ffi::String> loop_rvs,
                                      IntImm preserve_unit_iters) {
    PythonAPICall py("fuse");
    for (const ffi::String& loop_rv : loop_rvs) {
      py.Input("", loop_rv);
    }
    py.Input("preserve_unit_iters", preserve_unit_iters->value != 0);
    py.SingleOutput(outputs);
    return py.Str();
  }

  template <typename>
  friend struct ::tvm::s_tir::UnpackedInstTraits;
};

struct ReorderTraits : public UnpackedInstTraits<ReorderTraits> {
  static constexpr const char* kName = "Reorder";
  static constexpr bool kIsPure = false;

 private:
  static constexpr size_t kNumInputs = 1;
  static constexpr size_t kNumAttrs = 0;
  static constexpr size_t kNumDecisions = 0;

  template <size_t delta>
  TVM_FFI_INLINE static void _SetInputs(AnyView* packed_args, const ffi::Array<Any>& inputs) {
    packed_args[delta] = inputs;
  }

  static void UnpackedApplyToSchedule(Schedule sch, ffi::Array<LoopRV> loop_rvs) {
    return sch->Reorder(loop_rvs);
  }

  static ffi::String UnpackedAsPython(ffi::Array<ffi::String> outputs,
                                      ffi::Array<ffi::String> loop_rvs) {
    PythonAPICall py("reorder");
    for (const ffi::String& loop_rv : loop_rvs) {
      py.Input("", loop_rv);
    }
    return py.Str();
  }

  template <typename>
  friend struct ::tvm::s_tir::UnpackedInstTraits;
};

struct AddUnitLoopTraits : public UnpackedInstTraits<AddUnitLoopTraits> {
  static constexpr const char* kName = "AddUnitLoop";
  static constexpr bool kIsPure = false;

 private:
  static constexpr size_t kNumInputs = 1;
  static constexpr size_t kNumAttrs = 0;
  static constexpr size_t kNumDecisions = 0;

  static LoopRV UnpackedApplyToSchedule(Schedule sch, ffi::ObjectRef rv) {
    if (auto block = rv.as<SBlockRV>()) {
      return sch->AddUnitLoop(block.value());
    } else if (auto loop = rv.as<LoopRV>()) {
      return sch->AddUnitLoop(loop.value());
    } else {
      TVM_FFI_THROW(TypeError) << "AddUnitLoop expects a loop or block";
      throw;
    }
  }

  static ffi::String UnpackedAsPython(ffi::Array<ffi::String> outputs, ffi::String rv) {
    PythonAPICall py("add_unit_loop");
    py.Input("block_or_loop", rv);
    py.SingleOutput(outputs);
    return py.Str();
  }

  template <typename>
  friend struct ::tvm::s_tir::UnpackedInstTraits;
};

TVM_REGISTER_INST_KIND_TRAITS(SplitTraits);
TVM_REGISTER_INST_KIND_TRAITS(LoopPartitionTraits);
TVM_REGISTER_INST_KIND_TRAITS(MergeTraits);
TVM_REGISTER_INST_KIND_TRAITS(FuseTraits);
TVM_REGISTER_INST_KIND_TRAITS(ReorderTraits);
TVM_REGISTER_INST_KIND_TRAITS(AddUnitLoopTraits);

}  // namespace s_tir
}  // namespace tvm
