/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file data_type_rewriter.cc
 * \brief Rewrite the data type of expressions.
 */

#include "data_type_rewriter.h"

#include <tvm/ffi/cast.h>
#include <tvm/ir/op.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>

#include <algorithm>
#include <utility>

#include "./functor_common.h"
#include "tvm/ir/expr.h"
#include "tvm/tirx/expr.h"
#include "tvm/tirx/stmt.h"
#include "tvm/tirx/var.h"

namespace tvm {
namespace tirx {

Stmt DataTypeLegalizer::VisitStmt_(const ForNode* op) {
  Stmt s = StmtExprMutator::VisitStmt_(op);
  op = s.as<ForNode>();
  TVM_FFI_ICHECK(op != nullptr) << "Expected type to be ForNode, but get " << s->GetTypeKey();
  PrimExpr e = VisitExpr(op->loop_var).as_or_throw<PrimExpr>();
  Var var = e.as_or_throw<Var>();
  auto n = CopyOnWrite(op);
  PrimType var_ty = var->ty.as_or_throw<PrimType>();
  n->min = cast(var_ty, op->min);
  n->extent = cast(var_ty, op->extent);
  if (op->step.has_value()) {
    n->step = cast(var_ty, *op->step);
  }
  return For(n);
}

Stmt DataTypeLegalizer::VisitStmt_(const SBlockRealizeNode* op) {
  SBlockRealize realize = StmtExprMutator::VisitStmt_(op).as_or_throw<SBlockRealize>();
  ffi::Array<PrimExpr> new_iter_values;
  bool changed = false;
  for (int i = 0; i < static_cast<int>(op->iter_values.size()); ++i) {
    PrimType dtype = realize->block->iter_vars[i]->var.ty();
    if (op->iter_values[i].ty() != dtype) {
      new_iter_values.push_back(cast(dtype, realize->iter_values[i]));
      changed = true;
    } else {
      new_iter_values.push_back(realize->iter_values[i]);
    }
  }
  if (changed) {
    realize.CopyOnWrite()->iter_values = std::move(new_iter_values);
  }
  return realize;
}

Stmt DataTypeLegalizer::VisitStmt_(const SBlockNode* op) {
  SBlock new_block = StmtExprMutator::VisitStmt_(op).as_or_throw<SBlock>();
  ffi::Array<IterVar> new_iter_vars =
      MutateArray(new_block->iter_vars, [/*this*/](const IterVar& iter) {
        PrimType dtype = iter->var.ty();
        if (iter->dom->min.ty() != dtype || iter->dom->extent.ty() != dtype) {
          IterVar new_iter = iter;
          new_iter.CopyOnWrite()->dom =
              Range(cast(dtype, iter->dom->min), cast(dtype, iter->dom->extent));
          return new_iter;
        } else {
          return iter;
        }
      });
  if (!op->iter_vars.same_as(new_iter_vars)) {
    new_block.CopyOnWrite()->iter_vars = std::move(new_iter_vars);
  }
  return new_block;
}

Stmt DataTypeLegalizer::VisitStmt_(const AttrStmtNode* op) {
  if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread) {
    Stmt s = StmtExprMutator::VisitStmt_(op);
    op = s.as<AttrStmtNode>();
    TVM_FFI_ICHECK(op != nullptr) << "Expected type to be AttrStmtNode"
                                  << ", but get " << s->GetTypeKey();
    const IterVarNode* iv = op->node.as<IterVarNode>();
    TVM_FFI_ICHECK(iv != nullptr) << "Expected type to be IterVarNode"
                                  << ", but get " << op->node.GetTypeKey();
    PrimExpr e = VisitExpr(iv->var).as_or_throw<PrimExpr>();
    PrimVar var = e.as_or_throw<PrimVar>();
    if (ivmap_.find(iv) == ivmap_.end()) {
      Range dom = iv->dom;
      if (dom.defined()) {
        PrimExpr extend = dom->extent;
        PrimType extend_ty = extend.ty();
        PrimType var_ty = var.ty();
        TVM_FFI_ICHECK(extend_ty.MatchesCode(DLDataTypeCode::kDLInt) &&
                       var_ty.MatchesCode(DLDataTypeCode::kDLInt));
        if (var_ty.bits() != extend_ty.bits()) {
          dom = Range(cast(var_ty, dom->min), cast(var_ty, extend), dom->span);
        }
      }
      ivmap_[iv] = IterVar(dom, var, iv->iter_type, iv->thread_tag);
    }
    return AttrStmt(ivmap_[iv], op->attr_key, cast(var.ty(), op->value), op->body);
  }
  return StmtExprMutator::VisitStmt_(op);
}

Expr DataTypeLegalizer::VisitExpr_(const LetNode* op) {
  PrimExpr value = this->VisitPrimExpr(op->value);
  Var var = op->var;

  if (value.ty() != op->var->ty.as_or_throw<PrimType>()) {
    var = op->var.copy_with_dtype(value.ty());
    var_remap_[op->var.get()] = var;
  }

  PrimExpr new_body = this->VisitPrimExpr(op->body);

  if (value.same_as(op->value) && new_body.same_as(op->body)) {
    return ffi::GetRef<PrimExpr>(op);
  } else {
    return Let(var, value, new_body, op->span);
  }
}

Stmt DataTypeLegalizer::VisitStmt_(const BindNode* op) {
  Expr value = this->VisitExpr(op->value);
  Var var = op->var;

  if (auto prim_value = value.as<PrimExpr>()) {
    if (prim_value.value().ty() != op->var->ty.as_or_throw<PrimType>()) {
      var = op->var.copy_with_dtype(prim_value.value().ty());
      var_remap_[op->var.get()] = var;
    }
  }

  if (value.same_as(op->value) && var.same_as(op->var)) {
    return ffi::GetRef<Stmt>(op);
  } else {
    return Bind(var, value, op->span);
  }
}

Expr DataTypeLegalizer::VisitExpr_(const VarNode* op) {
  if (auto it = var_remap_.find(op); it != var_remap_.end()) {
    return it->second;
  }
  return ffi::GetRef<Var>(op);
}

Expr DataTypeLegalizer::VisitExpr_(const SelectNode* op) {
  PrimExpr condition = this->VisitPrimExpr(op->condition);
  PrimExpr true_value = this->VisitPrimExpr(op->true_value);
  PrimExpr false_value = this->VisitPrimExpr(op->false_value);
  if (condition.same_as(op->condition) && true_value.same_as(op->true_value) &&
      false_value.same_as(op->false_value) && true_value.ty() == false_value.ty()) {
    return ffi::GetRef<PrimExpr>(op);
  } else {
    PrimType true_dtype = true_value.ty();
    PrimType false_dtype = false_value.ty();
    int bits = std::max(true_dtype.bits(), false_dtype.bits());
    PrimType dtype = true_dtype.WithBits(bits);
    if (true_dtype != dtype) true_value = cast(dtype, true_value);
    if (false_dtype != dtype) false_value = cast(dtype, false_value);
    return Select(condition, true_value, false_value);
  }
}

Expr DataTypeLegalizer::VisitExpr_(const RampNode* op) {
  PrimExpr base = VisitPrimExpr(op->base);
  PrimExpr stride = VisitPrimExpr(op->stride);
  if (base.same_as(op->base) && stride.same_as(op->stride) && base.ty() == stride.ty()) {
    return ffi::GetRef<PrimExpr>(op);
  } else {
    PrimType base_dtype = base.ty();
    PrimType stride_dtype = stride.ty();
    TVM_FFI_ICHECK(base_dtype.MatchesCode(DLDataTypeCode::kDLInt) &&
                   stride_dtype.MatchesCode(DLDataTypeCode::kDLInt));
    int bits = std::max(base_dtype.bits(), stride_dtype.bits());
    PrimType dtype = base_dtype.WithBits(bits);
    if (base_dtype->dtype != dtype->dtype) base = cast(dtype, base);
    if (stride_dtype->dtype != dtype->dtype) stride = cast(dtype, stride);
    return Ramp(base, stride, op->lanes);
  }
}

Expr DataTypeLegalizer::VisitExpr_(const CastNode* op) { return StmtExprMutator::VisitExpr_(op); }

#define TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(OP, FUNC)       \
  Expr DataTypeLegalizer::VisitExpr_(const OP* op) {                \
    PrimExpr a = this->VisitPrimExpr(op->a);                        \
    PrimExpr b = this->VisitPrimExpr(op->b);                        \
    if (op->a.same_as(a) && op->b.same_as(b) && a.ty() == b.ty()) { \
      return ffi::GetRef<PrimExpr>(op);                             \
    } else {                                                        \
      return FUNC(a, b);                                            \
    }                                                               \
  }

TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(AddNode, operator+);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(SubNode, operator-);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MulNode, operator*);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(DivNode, div);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(ModNode, truncmod);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(FloorDivNode, floordiv);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(FloorModNode, floormod);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MinNode, min);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MaxNode, max);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(EQNode, operator==);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(NENode, operator!=);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(LENode, operator<=);
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(LTNode, operator<);  // NOLINT(*)
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(GTNode, operator>);  // NOLINT(*)
TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(GENode, operator>=);

#undef TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH

Expr DataTypeLegalizer::VisitExpr_(const CallNode* op) {
  Call before = ffi::GetRef<Call>(op);
  Expr e = StmtExprMutator::VisitExpr_(op);
  op = e.as<CallNode>();
  TVM_FFI_ICHECK(op != nullptr) << "Expected type to be CallNode"
                                << ", but get " << e->GetTypeKey();
  if (!op->ty.as<PrimTypeNode>()) {
    return e;
  }
  PrimExpr prim_e = e.as_or_throw<PrimExpr>();
  if (op->op.same_as(builtin::shift_right())) {
    return op->args[0].as_or_throw<PrimExpr>() >> op->args[1].as_or_throw<PrimExpr>();
  } else if (op->op.same_as(builtin::shift_left())) {
    return op->args[0].as_or_throw<PrimExpr>() << op->args[1].as_or_throw<PrimExpr>();
  } else if (op->op.same_as(builtin::bitwise_and())) {
    return op->args[0].as_or_throw<PrimExpr>() & op->args[1].as_or_throw<PrimExpr>();
  } else if (op->op.same_as(builtin::bitwise_or())) {
    return op->args[0].as_or_throw<PrimExpr>() | op->args[1].as_or_throw<PrimExpr>();
  } else if (op->op.same_as(builtin::bitwise_xor())) {
    return op->args[0].as_or_throw<PrimExpr>() ^ op->args[1].as_or_throw<PrimExpr>();
  }
  static const Op& pow_op = Op::Get("tirx.pow");
  static const Op& clz_op = Op::Get("tirx.clz");
  if (op->op.same_as(pow_op)) {
    return pow(op->args[0].as_or_throw<PrimExpr>(), op->args[1].as_or_throw<PrimExpr>());
  } else if (op->op.same_as(builtin::if_then_else())) {
    return Call(op->ty.as_or_throw<PrimType>(), op->op,
                {op->args[0].as_or_throw<PrimExpr>(), op->args[1].as_or_throw<PrimExpr>(),
                 op->args[2].as_or_throw<PrimExpr>()},
                op->attrs, {}, op->span)
        .as_or_throw<PrimExpr>();
  } else if (op->op.same_as(clz_op)) {
    PrimType before_dtype = before->args[0].as_or_throw<PrimExpr>().ty();
    PrimType after_dtype = op->args[0].as_or_throw<PrimExpr>().ty();
    TVM_FFI_ICHECK((before_dtype.code() == DLDataTypeCode::kDLInt ||
                    before_dtype.code() == DLDataTypeCode::kDLUInt) &&
                   (before_dtype.bits() == 32 || before_dtype.bits() == 64))
        << "clz only supports 32 or 64 bit integer types, but get type before legalizing: "
        << before_dtype;
    TVM_FFI_ICHECK((after_dtype.code() == DLDataTypeCode::kDLInt ||
                    after_dtype.code() == DLDataTypeCode::kDLUInt) &&
                   (after_dtype.bits() == 32 || after_dtype.bits() == 64))
        << "clz only supports 32 or 64 bit integer types, but get type after legalizing: "
        << after_dtype;
    return prim_e - after_dtype.bits() + before_dtype.bits();
  }
  return prim_e;
}

Stmt IndexDataTypeRewriter::VisitStmt_(const AttrStmtNode* op) {
  if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread) {
    bool is_enabled = is_enabled_;
    is_enabled_ = true;
    auto stmt = DataTypeLegalizer::VisitStmt_(op);
    is_enabled_ = is_enabled;
    return stmt;
  }
  return DataTypeLegalizer::VisitStmt_(op);
}

Buffer IndexDataTypeRewriter::VisitBufferDef(const Buffer& buffer, bool alloc_data) {
  bool is_enabled = is_enabled_;
  is_enabled_ = true;
  Buffer new_buf = StmtMutator::VisitBufferDef(buffer, alloc_data);
  is_enabled_ = is_enabled;
  return new_buf;
}

Buffer IndexDataTypeRewriter::VisitBufferUse(const Buffer& buffer) {
  return StmtMutator::VisitBufferUse(buffer);
}

Stmt IndexDataTypeRewriter::VisitStmt_(const SBlockRealizeNode* op) {
  bool is_condition = is_condition_;
  is_condition_ = true;
  auto new_predicate = VisitPrimExpr(op->predicate);
  is_condition_ = is_condition;

  bool is_enabled = is_enabled_;
  is_enabled_ = true;
  auto new_iter_values =
      op->iter_values.Map([this](const PrimExpr& e) { return this->VisitPrimExpr(e); });
  is_enabled_ = is_enabled;
  SBlock new_body = this->VisitStmt(op->block).as_or_throw<SBlock>();
  if (!new_predicate.same_as(op->predicate) || !new_iter_values.same_as(op->iter_values) ||
      !new_body.same_as(op->block)) {
    SBlockRealize new_block_realize = ffi::GetRef<SBlockRealize>(op);
    auto* n = new_block_realize.CopyOnWrite();
    n->predicate = std::move(new_predicate);
    n->iter_values = std::move(new_iter_values);
    n->block = std::move(new_body);
    return new_block_realize;

  } else {
    return ffi::GetRef<Stmt>(op);
  }
}

Stmt IndexDataTypeRewriter::VisitStmt_(const SBlockNode* op) {
  ffi::Array<Buffer> new_alloc_buffers = op->alloc_buffers.Map(
      [this](const Buffer& buffer) { return this->VisitBufferDef(buffer, /*alloc_data=*/true); });
  ffi::Array<MatchBufferRegion> new_match_buffers =
      op->match_buffers.Map([this](const MatchBufferRegion& match_buffer_region) {
        Buffer new_buffer = this->VisitBufferDef(match_buffer_region->buffer, /*alloc_data=*/true);
        BufferRegion new_buffer_region = this->VisitBufferRegion(match_buffer_region->source);
        if (!new_buffer.same_as(match_buffer_region->buffer) ||
            !new_buffer_region.same_as(match_buffer_region->source)) {
          return MatchBufferRegion(new_buffer, new_buffer_region);
        } else {
          return match_buffer_region;
        }
      });
  ffi::Array<BufferRegion> new_reads = op->reads.Map(
      [this](const BufferRegion& buffer_region) { return this->VisitBufferRegion(buffer_region); });
  ffi::Array<BufferRegion> new_writes = op->writes.Map(
      [this](const BufferRegion& buffer_region) { return this->VisitBufferRegion(buffer_region); });
  ffi::Array<IterVar> new_iter_vars =
      op->iter_vars.Map([this](const IterVar& iter_var) { return this->VisitIterVar(iter_var); });
  ffi::Optional<Stmt> new_init = std::nullopt;
  if (op->init.has_value()) {
    new_init = this->VisitStmt(op->init.value());
  }
  ffi::Map<ffi::String, ffi::Any> new_annotations = VisitBlockAnnotations(op->annotations);
  Stmt new_body = this->VisitStmt(op->body);

  if (!new_init.same_as(op->init) || !new_body.same_as(op->body) ||
      !new_alloc_buffers.same_as(op->alloc_buffers) ||
      !new_match_buffers.same_as(op->match_buffers) || !new_reads.same_as(op->reads) ||
      !new_writes.same_as(op->writes) || new_iter_vars.same_as(op->iter_vars) ||
      !new_annotations.same_as(op->annotations)) {
    SBlock new_block = ffi::GetRef<SBlock>(op);
    SBlockNode* n = new_block.CopyOnWrite();
    n->alloc_buffers = std::move(new_alloc_buffers);
    n->match_buffers = std::move(new_match_buffers);
    n->reads = std::move(new_reads);
    n->writes = std::move(new_writes);
    n->iter_vars = std::move(new_iter_vars);
    n->init = std::move(new_init);
    n->annotations = std::move(new_annotations);
    n->body = std::move(new_body);
    return new_block;
  }
  return ffi::GetRef<Stmt>(op);
}

ffi::Map<ffi::String, ffi::Any> IndexDataTypeRewriter::VisitBlockAnnotations(
    const ffi::Map<ffi::String, ffi::Any>& annotations) {
  auto new_annotations = annotations;

  std::function<Any(const Any&)> f_mutate_obj = [this, &f_mutate_obj](const Any& obj) -> Any {
    if (obj == nullptr) {
      return obj;
    }
    if (obj.as<BufferNode>()) {
      Buffer buffer = obj.as_or_throw<Buffer>();
      if (Buffer new_buffer = VisitBufferUse(buffer); !new_buffer.same_as(buffer)) {
        return new_buffer;
      }
    } else if (obj.as<ffi::ArrayObj>()) {
      return obj.as_or_throw<ffi::Array<Any>>().Map(f_mutate_obj);
    }
    return obj;
  };
  for (const auto& [key, value] : annotations) {
    if (auto opt_object_ref = value.as<ffi::ObjectRef>()) {
      auto new_value = f_mutate_obj(*opt_object_ref);
      if (!new_value.same_as(*opt_object_ref)) {
        new_annotations.Set(key, new_value);
      }
    }
  }
  return new_annotations;
}

IterVar IndexDataTypeRewriter::VisitIterVar(const IterVar& iter_var) {
  bool is_enabled = is_enabled_;
  is_enabled_ = true;
  PrimVar new_var = VisitPrimExpr(iter_var->var).as_or_throw<PrimVar>();
  PrimExpr min = VisitPrimExpr(iter_var->dom->min);
  PrimExpr extent = VisitPrimExpr(iter_var->dom->extent);
  is_enabled_ = is_enabled;
  if (!new_var.same_as(iter_var->var) || !min.same_as(iter_var->dom->min) ||
      !extent.same_as(iter_var->dom->extent)) {
    IterVar new_iter_var = iter_var;
    IterVarNode* n = new_iter_var.CopyOnWrite();
    n->var = std::move(new_var);
    n->dom = Range(min, extent);
    return new_iter_var;
  }
  return iter_var;
}

BufferRegion IndexDataTypeRewriter::VisitBufferRegion(const BufferRegion& buffer_region) {
  Buffer remapped_buffer = VisitBufferUse(buffer_region->buffer);

  bool is_enabled = is_enabled_;
  is_enabled_ = true;
  auto new_region = buffer_region->region.Map([&](const Range& range) {
    return Range::FromMinExtent(this->VisitPrimExpr(range->min),
                                this->VisitPrimExpr(range->extent));
  });
  is_enabled_ = is_enabled;

  if (!remapped_buffer.same_as(buffer_region->buffer) ||
      !new_region.same_as(buffer_region->region)) {
    return BufferRegion(remapped_buffer, new_region);
  } else {
    return buffer_region;
  }
}

Stmt IndexDataTypeRewriter::VisitStmt_(const BufferStoreNode* op) {
  BufferStore store = ffi::GetRef<BufferStore>(op);

  Buffer new_buffer = VisitBufferUse(op->buffer);
  auto value = this->VisitPrimExpr(op->value);
  PrimType value_dtype = value.ty();
  if (new_buffer->dtype != value_dtype && value_dtype.IsScalar()) {
    value = cast(new_buffer->dtype, value);
  }
  auto indices = VisitIndices(op->indices);

  if (!new_buffer.same_as(op->buffer) || !value.same_as(op->value) ||
      !indices.same_as(op->indices)) {
    auto writer = store.CopyOnWrite();
    writer->buffer = new_buffer;
    writer->value = value;
    writer->indices = indices;
  }

  return store;
}

Expr IndexDataTypeRewriter::VisitExpr_(const BufferLoadNode* op) {
  BufferLoad load = ffi::GetRef<BufferLoad>(op);

  Buffer new_buffer = VisitBufferUse(op->buffer);
  auto indices = VisitIndices(op->indices);

  if (!new_buffer.same_as(op->buffer) || !indices.same_as(op->indices)) {
    auto writer = load.CopyOnWrite();
    writer->indices = indices;
    writer->buffer = new_buffer;
  }

  return load;
}

ffi::Array<PrimExpr> IndexDataTypeRewriter::VisitIndices(ffi::Array<PrimExpr> indices) {
  bool is_enabled = is_enabled_;
  is_enabled_ = true;

  auto fmutate = [this](const PrimExpr& index) { return this->VisitPrimExpr(index); };
  indices.MutateByApply(fmutate);

  is_enabled_ = is_enabled;

  return indices;
}

Stmt IndexDataTypeRewriter::VisitStmt_(const IfThenElseNode* op) {
  bool is_condition = is_condition_;
  is_condition_ = true;
  PrimExpr cond = VisitPrimExpr(op->condition);
  is_condition_ = is_condition;

  Stmt then_case = VisitStmt(op->then_case);
  ffi::Optional<Stmt> else_case = op->else_case.has_value()
                                      ? ffi::Optional<Stmt>{VisitStmt(op->else_case.value())}
                                      : std::nullopt;
  if (!cond.same_as(op->condition) || !then_case.same_as(op->then_case) ||
      !else_case.same_as(op->else_case)) {
    IfThenElse new_stmt = ffi::GetRef<IfThenElse>(op);
    auto* n = new_stmt.CopyOnWrite();
    n->condition = std::move(cond);
    n->then_case = std::move(then_case);
    n->else_case = std::move(else_case);
    return new_stmt;
  }
  return ffi::GetRef<Stmt>(op);
}

Stmt IndexDataTypeRewriter::VisitStmt_(const ForNode* op) {
  bool is_enabled = is_enabled_;
  is_enabled_ = true;
  PrimVar new_loop_var = VisitPrimExpr(op->loop_var).as_or_throw<PrimVar>();
  PrimExpr min = VisitPrimExpr(op->min);
  PrimExpr extent = VisitPrimExpr(op->extent);
  is_enabled_ = is_enabled;

  Stmt new_body = VisitStmt(op->body);

  if (!new_loop_var.same_as(op->loop_var) || !min.same_as(op->min) || !extent.same_as(op->extent) ||
      !new_body.same_as(op->body)) {
    For new_for = ffi::GetRef<For>(op);
    auto* n = new_for.CopyOnWrite();
    n->loop_var = new_loop_var;
    n->min = cast(new_loop_var.ty(), min);
    n->extent = cast(new_loop_var.ty(), extent);
    if (op->thread_binding.has_value()) {
      auto old_thread_binding = op->thread_binding.value();
      auto* ptr = old_thread_binding.CopyOnWrite();
      ptr->var = old_thread_binding->var.copy_with_dtype(new_loop_var.ty());
      n->thread_binding = ffi::Optional<IterVar>(std::move(old_thread_binding));
    }
    n->body = new_body;
    return new_for;

  } else {
    return ffi::GetRef<Stmt>(op);
  }
}

Stmt IndexDataTypeRewriter::VisitStmt_(const BindNode* op) {
  Bind bind_stmt = DataTypeLegalizer::VisitStmt_(op).as_or_throw<Bind>();
  if (var_remap_.find(bind_stmt->var.get()) == var_remap_.end()) {
    return bind_stmt;
  }
  bool is_enabled = is_enabled_;
  is_enabled_ = true;
  PrimExpr value = VisitExpr(op->value).as_or_throw<PrimExpr>();
  Var var = var_remap_[bind_stmt->var.get()];
  is_enabled_ = is_enabled;
  TVM_FFI_ICHECK(value.ty() == var->ty.as_or_throw<PrimType>());
  return Bind(var, value, bind_stmt->span);
}

#define TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(OP, FUNC)                       \
  Expr IndexDataTypeRewriter::VisitExpr_(const OP* op) {                             \
    bool is_enabled = is_enabled_;                                                   \
    is_enabled_ = is_condition_ && op->a.ty().MatchesCode(DLDataTypeCode::kDLInt) && \
                  op->b.ty().MatchesCode(DLDataTypeCode::kDLInt);                    \
    auto result = Parent::VisitExpr_(op);                                            \
    is_enabled_ = is_enabled;                                                        \
    return result;                                                                   \
  }

TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(EQNode, operator==);
TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(NENode, operator!=);
TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(LENode, operator<=);
TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(LTNode, operator<);  // NOLINT(*)
TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(GTNode, operator>);  // NOLINT(*)
TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(GENode, operator>=);

Expr IndexDataTypeRewriter::VisitExpr_(const CallNode* op) {
  // handle if_then_else condition
  if (op->op.same_as(builtin::if_then_else())) {
    bool is_condition = is_condition_;
    is_condition_ = true;
    PrimExpr cond = VisitPrimExpr(op->args[0].as_or_throw<PrimExpr>());
    is_condition_ = is_condition;
    return Call(op->ty.as_or_throw<PrimType>(), op->op,
                {cond, VisitPrimExpr(op->args[1].as_or_throw<PrimExpr>()),
                 VisitPrimExpr(op->args[2].as_or_throw<PrimExpr>())},
                op->attrs, {}, op->span)
        .as_or_throw<PrimExpr>();
  }
  return Parent::VisitExpr_(op);
}

Expr IndexDataTypeRewriter::VisitExpr_(const SelectNode* op) {
  bool is_condition = true;
  std::swap(is_condition_, is_condition);
  PrimExpr condition = this->VisitPrimExpr(op->condition);
  std::swap(is_condition_, is_condition);
  PrimExpr true_value = this->VisitPrimExpr(op->true_value);
  PrimExpr false_value = this->VisitPrimExpr(op->false_value);

  if (condition.same_as(op->condition) && true_value.same_as(op->true_value) &&
      false_value.same_as(op->false_value) && true_value.ty() == false_value.ty()) {
    return ffi::GetRef<PrimExpr>(op);
  } else {
    PrimType true_dtype = true_value.ty();
    PrimType false_dtype = false_value.ty();
    int bits = std::max(true_dtype.bits(), false_dtype.bits());
    PrimType dtype = true_dtype.WithBits(bits);
    if (true_dtype->dtype != dtype->dtype) true_value = cast(dtype, true_value);
    if (false_dtype->dtype != dtype->dtype) false_value = cast(dtype, false_value);
    return Select(condition, true_value, false_value);
  }
}

#undef TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH

IndexDataTypeNormalizer::IndexDataTypeNormalizer(PrimType target_data_type)
    : target_data_type_(std::move(target_data_type)) {}

PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) {
  // collect var remap
  VisitStmt(std::move(func->body));
  buffer_remap_.clear();
  ivmap_.clear();
  // start rewrite
  ffi::Map<Var, Buffer> new_buffer_map = func->buffer_map;
  for (const auto& [var, buffer] : func->buffer_map) {
    new_buffer_map.Set(var, VisitBufferDef(buffer, /*alloc_data=*/true));
  }
  // remap params
  bool is_enabled = true;
  std::swap(is_enabled_, is_enabled);
  ffi::Array<Var> params = func->params.Map([this](Var param) {
    if (auto param_ty = param->ty.as<PrimType>();
        param_ty && param_ty.value().MatchesCode(DLDataTypeCode::kDLInt)) {
      return this->VisitPrimExpr(param.as_or_throw<PrimExpr>()).as_or_throw<Var>();
    } else {
      return param;
    }
  });
  std::swap(is_enabled_, is_enabled);

  PrimFuncNode* new_func = func.CopyOnWrite();
  new_func->params = std::move(params);
  new_func->buffer_map = std::move(new_buffer_map);
  new_func->body = VisitStmt(std::move(new_func->body));
  return func;
}

bool IndexDataTypeNormalizer::CanRewriteDType(PrimType dtype) const {
  return dtype.code() == DLDataTypeCode::kDLInt && dtype.bits() >= 32;
}

Expr IndexDataTypeNormalizer::VisitExpr_(const IntImmNode* op) {
  if (is_enabled_ && CanRewriteDType(op->ty.as_or_throw<PrimType>())) {
    TVM_FFI_ICHECK_LE(op->value, max_value(target_data_type_).as_or_throw<IntImm>()->value);
    return cast(target_data_type_, ffi::GetRef<IntImm>(op));
  }
  return ffi::GetRef<IntImm>(op);
}

Expr IndexDataTypeNormalizer::VisitExpr_(const VarNode* op) {
  auto dtype_opt = op->ty.as<PrimType>();
  if (!dtype_opt) {
    return DataTypeLegalizer::VisitExpr_(op);
  }
  PrimType dtype = dtype_opt.value();
  if (is_enabled_ && CanRewriteDType(dtype) && dtype->dtype != target_data_type_->dtype &&
      !var_remap_.count(op)) {
    var_remap_[op] = ffi::GetRef<Var>(op).copy_with_dtype(target_data_type_);
  }
  return DataTypeLegalizer::VisitExpr_(op);
}

Expr IndexDataTypeNormalizer::VisitExpr_(const CastNode* op) {
  // Unwrap the cast only when the dtype of this cast is integer dtype.
  // When the dtype of this cast is not integer dtype, it means that this cast
  // has some other purpose, and we should not unwrap the cast.
  PrimType dtype = op->ty.as_or_throw<PrimType>();
  if (is_enabled_ && CanRewriteDType(dtype)) {
    PrimExpr value = IndexDataTypeNormalizer::VisitPrimExpr(op->value);
    return value.ty() == target_data_type_ ? value : Cast(target_data_type_, value);
  }
  return IndexDataTypeRewriter::VisitExpr_(op);
}

}  // namespace tirx
}  // namespace tvm
