code: #include <torch/csrc/jit/tensorexpr/ir_mutator.h>
  
  #include <torch/csrc/jit/tensorexpr/eval.h>
  #include <torch/csrc/jit/tensorexpr/ir.h>
  #include <torch/csrc/jit/tensorexpr/ir_simplifier.h>
  
  namespace torch {
  namespace jit {
  namespace tensorexpr {
  
  template <typename Op>
  static const Expr* mutate_binary_op(
      const BinaryOpNode<Op>* v,
      IRMutator* mutator,
      bool option = false) {
    const Expr* lhs = v->lhs();
    const Expr* rhs = v->rhs();
    const Expr* lhs_new = lhs->accept_mutator(mutator);
    const Expr* rhs_new = rhs->accept_mutator(mutator);
    if (lhs == lhs_new && rhs == rhs_new) {
      return v;
    }
    IRNodeType expr_type = v->expr_type();
    switch (expr_type) {
      case IRNodeType::kAdd:
        return new Add(lhs_new, rhs_new);
      case IRNodeType::kSub:
        return new Sub(lhs_new, rhs_new);
      case IRNodeType::kMul:
        return new Mul(lhs_new, rhs_new);
      case IRNodeType::kDiv:
        return new Div(lhs_new, rhs_new);
      case IRNodeType::kMod:
        return new Mod(lhs_new, rhs_new);
      case IRNodeType::kMax:
        return new Max(lhs_new, rhs_new, option);
      case IRNodeType::kMin:
        return new Min(lhs_new, rhs_new, option);
      case IRNodeType::kAnd:
        return new And(lhs_new, rhs_new);
      case IRNodeType::kOr:
        return new Or(lhs_new, rhs_new);
      case IRNodeType::kXor:
        return new Xor(lhs_new, rhs_new);
      case IRNodeType::kLshift:
        return new Lshift(lhs_new, rhs_new);
      case IRNodeType::kRshift:
        return new Rshift(lhs_new, rhs_new);
      default:
        throw unsupported_dtype();
    }
  }
  
  const Expr* IRMutator::mutate(const Add* v) {
    return mutate_binary_op(v, this);
  }
  
  const Expr* IRMutator::mutate(const Sub* v) {
    return mutate_binary_op(v, this);
  }
  
  const Expr* IRMutator::mutate(const Mul* v) {
    return mutate_binary_op(v, this);
  }
  
  const Expr* IRMutator::mutate(const Div* v) {
    return mutate_binary_op(v, this);
  }
  
  const Expr* IRMutator::mutate(const Mod* v) {
    return mutate_binary_op(v, this);
  }
  
  const Expr* IRMutator::mutate(const And* v) {
    return mutate_binary_op(v, this);
  }
  
  const Expr* IRMutator::mutate(const Or* v) {
    return mutate_binary_op(v, this);
  }
  
  const Expr* IRMutator::mutate(const Xor* v) {
    return mutate_binary_op(v, this);
  }
  
  const Expr* IRMutator::mutate(const Lshift* v) {
    return mutate_binary_op(v, this);
  }
  
  const Expr* IRMutator::mutate(const Rshift* v) {
    return mutate_binary_op(v, this);
  }
  
  const Expr* IRMutator::mutate(const Max* v) {
    return mutate_binary_op(v, this, v->propagate_nans());
  }
  
  const Expr* IRMutator::mutate(const Min* v) {
    return mutate_binary_op(v, this, v->propagate_nans());
  }
  
  const Expr* IRMutator::mutate(const CompareSelect* v) {
    const Expr* lhs = v->lhs();
    const Expr* rhs = v->rhs();
    const Expr* retval1 = v->ret_val1();
    const Expr* retval2 = v->ret_val2();
    const Expr* lhs_new = lhs->accept_mutator(this);
    const Expr* rhs_new = rhs->accept_mutator(this);
    const Expr* retval1_new = retval1->accept_mutator(this);
    const Expr* retval2_new = retval2->accept_mutator(this);
    if (lhs == lhs_new && rhs == rhs_new && retval1 == retval1_new &&
        retval2 == retval2_new) {
      return v;
    }
    return CompareSelect::make(
               ExprHandle(lhs_new),
               ExprHandle(rhs_new),
               ExprHandle(retval1_new),
               ExprHandle(retval2_new),
               v->compare_select_op())
        .node();
  }
  
  // NOLINTNEXTLINE
  #define IMM_MUTATE_DEFINE(_1, Name)                   \
    const Expr* IRMutator::mutate(const Name##Imm* v) { \
      return v;                                         \
    }
  AT_FORALL_SCALAR_TYPES_AND2(Bool, Half, IMM_MUTATE_DEFINE);
  #undef IMM_MUTATE_DEFINE
  
  const Expr* IRMutator::mutate(const Cast* v) {
    const Expr* src_value = v->src_value();
    const Expr* src_value_new = src_value->accept_mutator(this);
    if (src_value_new == v->src_value()) {
      return v;
    }
    return new Cast(v->dtype(), src_value_new);
  }
  
  const Expr* IRMutator::mutate(const Var* v) {
    return v;
  }
  
  const Expr* IRMutator::mutate(const Let* v) {
    const Expr* var = v->var();
    const Expr* value = v->value();
    const Expr* body = v->body();
    const Expr* var_new = var->accept_mutator(this);
    const Expr* value_new = value->accept_mutator(this);
    const Expr* body_new = body->accept_mutator(this);
    if ((var == var_new) && (value == value_new) && (body == body_new)) {
      return v;
    }
    return new Let(var_new, value_new, body_new);
  }
  
  Stmt* IRMutator::mutate(const LetStmt* v) {
    const Var* var = v->var();
    const Expr* value = v->value();
    Stmt* body = v->body();
    const Var* var_new = dynamic_cast<const Var*>(var->accept_mutator(this));
    if (var_new == nullptr) {
      throw std::runtime_error("LetStmt var must be variable");
    }
    const Expr* value_new = value->accept_mutator(this);
    Stmt* body_new = body->accept_mutator(this);
    if ((var == var_new) && (value == value_new) && (body == body_new)) {
      return (Stmt*)v;
    }
    return new LetStmt(var_new, value_new, body_new);
  }
  
  const Expr* IRMutator::mutate(const Ramp* v) {
    const Expr* base = v->base();
    const Expr* stride = v->stride();
    const Expr* base_new = base->accept_mutator(this);
    const Expr* stride_new = stride->accept_mutator(this);
    if (base == base_new && stride == stride_new) {
      return v;
    }
    return new Ramp(base_new, stride_new, v->lanes());
  }
  
  const Expr* IRMutator::mutate(const Load* v) {
    Dtype dtype = v->dtype();
    const Var* base_handle = v->base_handle();
    const Expr* index = v->index();
    const Expr* mask = v->mask();
    const Expr* base_handle_expr = base_handle->accept_mutator(this);
    const Var* base_handle_new = dynamic_cast<const Var*>(base_handle_expr);
    const Expr* index_new = index->accept_mutator(this);
    const Expr* mask_new = mask->accept_mutator(this);
    if (base_handle == base_handle_new && index == index_new &&
        mask == mask_new) {
      return v;
    }
    return new Load(dtype, base_handle_new, index_new, mask_new);
  }
  
  const Expr* IRMutator::mutate(const Broadcast* v) {
    const Expr* value = v->value();
    int lanes = v->lanes();
    const Expr* value_new = value->accept_mutator(this);
    if (value == value_new) {
      return v;
    }
    return new Broadcast(value_new, lanes);
  }
  
  const Expr* IRMutator::mutate(const IfThenElse* v) {
    const Expr* condition = v->condition();
    const Expr* true_value = v->true_value();
    const Expr* false_value = v->false_value();
    const Expr* condition_new = condition->accept_mutator(this);
    const Expr* true_value_new = true_value->accept_mutator(this);
    const Expr* false_value_new = false_value->accept_mutator(this);
  
    if (condition == condition_new && true_value == true_value_new &&
        false_value == false_value_new) {
      return v;
    }
  
    return new IfThenElse(condition_new, true_value_new, false_value_new);
  }
  
  const Expr* IRMutator::mutate(const Intrinsics* v) {
    const BaseCallNode* base = v;
    return this->mutate(base);
  }
  
  const Expr* IRMutator::mutate(const FunctionCall* v) {
    const BaseCallNode* base = v;
    return this->mutate(base);
  }
  
  const Expr* IRMutator::mutate(const Term* v) {
    const Expr* newScalar = v->scalar()->accept_mutator(this);
  
    std::vector<const Expr*> variables;
    for (const auto* t : v->variables()) {
      variables.push_back(t->accept_mutator(this));
    }
    return new Term(v->hasher(), newScalar, variables);
  }
  
  const Expr* IRMutator::mutate(const Polynomial* v) {
    const Expr* newScalar = v->scalar()->accept_mutator(this);
  
    std::vector<const Term*> variables;
    for (const auto* t : v->variables()) {
      variables.push_back(static_cast<const Term*>(t->accept_mutator(this)));
    }
    return new Polynomial(v->hasher(), newScalar, variables);
  }
  
  const Expr* IRMutator::mutate(const BaseCallNode* v) {
    std::vector<const Expr*> params(v->nparams());
    bool any_change = false;
    for (int i = 0; i < v->nparams(); i++) {
      const Expr* value = v->param(i);
      const Expr* value_new = value->accept_mutator(this);
      if (value != value_new) {
        any_change = true;
      }
      params[i] = value_new;
    }
    if (!any_change) {
      return v;
    }
    return v->DefaultMutator(params);
  }
  
  Stmt* IRMutator::mutate(const For* v) {
    const Expr* var = v->var();
    const Expr* start = v->start();
    const Expr* stop = v->stop();
    Stmt* body = v->body();
    LoopOptions loop_options = v->loop_options();
    const Expr* var_new_expr = var->accept_mutator(this);
    const Var* var_new = dynamic_cast<const Var*>(var_new_expr);
    const Expr* start_new = start->accept_mutator(this);
    const Expr* stop_new = stop->accept_mutator(this);
    Stmt* body_new = body->accept_mutator(this);
    if (!body_new) {
      return nullptr;
    }
    if (var == var_new && start == start_new && stop == stop_new &&
        body == body_new) {
      return (Stmt*)v;
    }
    if (body_new == body) {
      body_new = Stmt::clone(body);
    }
    return new For(var_new, start_new, stop_new, body_new, loop_options);
  }
  
  Stmt* IRMutator::mutate(const Block* v) {
    bool any_change = false;
    std::vector<Stmt*> stmts;
    for (Stmt* stmt : v->stmts()) {
      Stmt* stmt_new = stmt->accept_mutator(this);
      if (stmt != stmt_new) {
        any_change = true;
      } else {
        stmt_new = Stmt::clone(stmt);
      }
      if (stmt_new) {
        stmts.push_back(stmt_new);
      }
    }
    if (!any_change) {
      return (Stmt*)v;
    }
    return Block::make(stmts);
  }
  
  Stmt* IRMutator::mutate(const Store* v) {
    const Var* base_handle = v->base_handle();
    const Expr* index = v->index();
    const Expr* value = v->value();
    const Expr* mask = v->mask();
    const Expr* base_handle_expr = base_handle->accept_mutator(this);
    const Var* base_handle_new = dynamic_cast<const Var*>(base_handle_expr);
    const Expr* index_new = index->accept_mutator(this);
    const Expr* value_new = value->accept_mutator(this);
    const Expr* mask_new = mask->accept_mutator(this);
    if (base_handle == base_handle_new && index == index_new &&
        value == value_new && mask == mask_new) {
      return (Stmt*)v;
    }
    return new Store(base_handle_new, index_new, value_new, mask_new);
  }
  
  Stmt* IRMutator::mutate(const Allocate* v) {
    const Var* buffer_var_old = v->buffer_var();
    const Var* buffer_var_new =
        dynamic_cast<const Var*>(buffer_var_old->accept_mutator(this));
    bool any_change = buffer_var_new == buffer_var_old;
  
    std::vector<const Expr*> dims_old = v->dims();
    std::vector<const Expr*> dims_new(dims_old.size());
    for (size_t i = 0; i < dims_old.size(); i++) {
      dims_new[i] = dims_old[i]->accept_mutator(this);
      any_change |= (dims_new[i] == dims_old[i]);
    }
  
    if (!any_change) {
      return (Stmt*)v;
    }
  
    return new Allocate(buffer_var_new, v->dtype(), dims_new);
  }
  
  Stmt* IRMutator::mutate(const Free* v) {
    const Expr* buffer_var_old = v->buffer_var();
    const Var* buffer_var_new =
        dynamic_cast<const Var*>(buffer_var_old->accept_mutator(this));
    if (buffer_var_new == buffer_var_old) {
      return (Stmt*)v;
    }
  
    return new Free(buffer_var_new);
  }
  
  Stmt* IRMutator::mutate(const Cond* v) {
    const Expr* cond_old = v->condition();
    Stmt* true_old = v->true_stmt();
    Stmt* false_old = v->false_stmt();
  
    const Expr* cond_new = cond_old->accept_mutator(this);
    Stmt* true_new = true_old ? true_old->accept_mutator(this) : true_old;
    Stmt* false_new = false_old ? false_old->accept_mutator(this) : false_old;
  
    if (cond_old == cond_new && true_old == true_new && false_old == false_new) {
      return (Stmt*)v;
    }
  
    if (true_old && true_new == true_old) {
      true_new = Stmt::clone(true_old);
    }
    if (false_old && false_new == false_old) {
      false_new = Stmt::clone(false_old);
    }
  
    return new Cond(cond_new, true_new, false_new);
  }
  
  const Expr* IRMutator::DefaultMutator(
      const BaseCallNode* v,
      std::vector<const Expr*>& params) {
    return v->DefaultMutator(params);
  }
  
  class StmtClone : public IRMutator {
   public:
    Stmt* mutate(const LetStmt* v) override;
    Stmt* mutate(const For* v) override;
    Stmt* mutate(const Block* v) override;
    Stmt* mutate(const Store* v) override;
    Stmt* mutate(const Allocate* v) override;
    Stmt* mutate(const Free* v) override;
    Stmt* mutate(const Cond* v) override;
  };
  
  Stmt* StmtClone::mutate(const LetStmt* v) {
    // Expressions are immutable => we don't need to clone them
    const Var* var = v->var();
    const Expr* value = v->value();
  
    // Statements are mutable => we need to clone them
    Stmt* body_new = v->body()->accept_mutator(this);
  
    // Create a new LetStmt with the cloned statement for body
    return new LetStmt(var, value, body_new);
  }
  
  Stmt* StmtClone::mutate(const For* v) {
    // Only body needs to be cloned as only statements are mutable
    Stmt* body_new = v->body()->accept_mutator(this);
  
    return new For(v->var(), v->start(), v->stop(), body_new, v->loop_options());
  }
  
  Stmt* StmtClone::mutate(const Block* v) {
    std::vector<Stmt*> stmts;
    for (Stmt* stmt : v->stmts()) {
      stmts.push_back(stmt->accept_mutator(this));
    }
    return new Block(stmts);
  }
  
  Stmt* StmtClone::mutate(const Store* v) {
    return new Store(v->base_handle(), v->index(), v->value(), v->mask());
  }
  
  Stmt* StmtClone::mutate(const Allocate* v) {
    return new Allocate(v->buffer_var(), v->dtype(), v->dims());
  }
  
  Stmt* StmtClone::mutate(const Free* v) {
    return new Free(v->buffer_var());
  }
  
  Stmt* StmtClone::mutate(const Cond* v) {
    const Expr* cond_old = v->condition();
    Stmt* true_old = v->true_stmt();
    Stmt* false_old = v->false_stmt();
  
    Stmt* true_new = true_old ? true_old->accept_mutator(this) : true_old;
    Stmt* false_new = false_old ? false_old->accept_mutator(this) : false_old;
  
    return new Cond(v->condition(), true_new, false_new);
  }
  
  Stmt* Stmt::clone(Stmt* s) {
    StmtClone clone_mutator;
    Stmt* cloned = s->accept_mutator(&clone_mutator);
    set_parent(cloned, nullptr);
    return cloned;
  }
  
  } // namespace tensorexpr
  } // namespace jit
  } // namespace torch
  