From 3066a1f050ee682154bcc7e804e344f949ee23f5 Mon Sep 17 00:00:00 2001 From: tqchen Date: Mon, 6 Jul 2026 04:58:21 +0000 Subject: [PATCH 1/3] [Relax] Phase out variable Id indirection Use Relax variable object identity directly and make type-changing rewrites remap definitions, uses, and binding lookups coherently. --- include/tvm/relax/expr.h | 85 +++---------------- include/tvm/relax/expr_functor.h | 2 +- python/tvm/relax/__init__.py | 1 - python/tvm/relax/expr.py | 44 ++-------- python/tvm/relax/expr_functor.py | 21 ++--- .../relax/transform/lazy_transform_params.py | 2 +- .../distributed/transform/lower_distir.cc | 8 +- .../transform/propagate_sharding.cc | 11 +-- src/relax/ir/block_builder.cc | 14 +-- src/relax/ir/expr.cc | 54 +++--------- src/relax/ir/expr_functor.cc | 50 ++++++++--- src/relax/ir/py_expr_functor.cc | 6 +- src/relax/script/builder/frame.cc | 8 +- src/relax/script/builder/ir.cc | 12 +-- src/relax/training/utils.cc | 6 +- src/relax/transform/canonicalize_bindings.cc | 34 ++++---- src/relax/transform/convert_layout.cc | 2 +- .../transform/eliminate_common_subexpr.cc | 2 +- src/relax/transform/fuse_ops.cc | 8 +- src/relax/transform/gradient.cc | 41 ++++----- src/relax/transform/gradient_simplifier.cc | 2 +- src/relax/transform/lambda_lift.cc | 7 +- src/relax/transform/realize_vdevice.cc | 4 +- src/relax/transform/rewrite_cuda_graph.cc | 9 +- src/relax/transform/to_mixed_precision.cc | 12 +-- src/relax/transform/to_non_dataflow.cc | 4 +- src/relax/transform/update_param_type.cc | 2 +- src/relax/transform/utils.h | 10 +-- src/relax/utils.cc | 2 +- tests/python/relax/test_expr_functor.py | 4 +- .../test_transform_canonicalize_bindings.py | 4 +- 31 files changed, 190 insertions(+), 281 deletions(-) diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h index 51b4ed43bdb9..0e67dc277c18 100644 --- a/include/tvm/relax/expr.h +++ b/include/tvm/relax/expr.h @@ -36,44 +36,6 @@ namespace tvm { namespace relax { -/*! - * \brief The unique identifier of variables. - * - * Id is like name to the variables, - * except that id is unique for each Var. - * - * \note Do not create Id directly, they are created in Var. - */ -class IdNode : public ffi::Object { - public: - /*! - * \brief The name of the variable, - * this only acts as a hint to the user, - * and is not used for equality. - */ - ffi::String name_hint; - - static void RegisterReflection() { - namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("name_hint", &IdNode::name_hint, - refl::AttachFieldFlag::SEqHashIgnore()); - } - - static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.Id", IdNode, ffi::Object); -}; - -class Id : public ffi::ObjectRef { - public: - /*! - * \brief The constructor - * \param name_hint The name of the variable. - */ - TVM_DLL explicit Id(ffi::String name_hint); - - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Id, ffi::ObjectRef, IdNode); -}; - /*! \brief Tuple container */ class TupleNode : public ExprNode { public: @@ -173,48 +135,31 @@ class ShapeExpr : public Expr { /*! \brief The variable class for all Relax bindings. */ class VarNode : public ExprNode { public: - /*! \brief The identifier of the variable, which is used for comparing stable equality across - * transformations. */ - Id vid; + /*! + * \brief The hint to the variable name. + * \note Each variable is uniquely identified by its address. + */ + ffi::String name_hint_; /*! \return The name hint of the variable */ - const ffi::String& name_hint() const { return vid->name_hint; } + const ffi::String& name_hint() const { return name_hint_; } static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("vid", &VarNode::vid); - // customize structural equal and hash to include ty - refl::TypeAttrDef() - .def("__s_equal__", &VarNode::SEqual) - .def("__s_hash__", &VarNode::SHash); - } - - bool SEqual(const VarNode* other, - ffi::TypedFunction equal) const { - return equal(vid, other->vid, false, "vid") && equal(ty, other->ty, false, "ty"); - } - - int64_t SHash(int64_t init_hash, ffi::TypedFunction hash) const { - int64_t hash_value = init_hash; - hash_value = hash(vid, hash_value, false); - hash_value = hash(ty, hash_value, false); - return hash_value; + refl::ObjectDef().def_ro("name_hint", &VarNode::name_hint_, + refl::AttachFieldFlag::SEqHashIgnore()); } - static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar; static constexpr const uint32_t _type_child_slots = 1; TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.Var", VarNode, ExprNode); }; class Var : public Expr { public: - TVM_DLL explicit Var(ffi::String name_hint, ffi::Optional ty_annotation, Span span = Span()) - : Var(Id(name_hint), ty_annotation, span) {} - - TVM_DLL explicit Var(Id vid, ffi::Optional ty_annotation, Span span = Span()); + TVM_DLL explicit Var(ffi::String name_hint, ffi::Optional ty_annotation, + Span span = Span()); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Var, Expr, VarNode); - - VarNode* CopyOnWrite(); }; /*! \brief A sub-type of the variable node used to mark dataflow variables from @@ -227,20 +172,16 @@ class DataflowVarNode : public VarNode { refl::ObjectDef(); } - static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.DataflowVar", DataflowVarNode, VarNode); }; class DataflowVar : public Var { public: TVM_DLL explicit DataflowVar(ffi::String name_hint, ffi::Optional ty_annotation, - Span span = Span()) - : DataflowVar(Id(name_hint), ty_annotation, span) {} - - TVM_DLL explicit DataflowVar(Id vid, ffi::Optional ty_annotation, Span span = Span()); + Span span = Span()); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataflowVar, Var, DataflowVarNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(DataflowVarNode); }; /*! diff --git a/include/tvm/relax/expr_functor.h b/include/tvm/relax/expr_functor.h index 05261d8c5ec6..7f1388803df6 100644 --- a/include/tvm/relax/expr_functor.h +++ b/include/tvm/relax/expr_functor.h @@ -559,7 +559,7 @@ class ExprMutator : public ExprMutatorBase { BlockBuilder builder_; /*! \brief Remap a var to a new var in use-site. */ - std::unordered_map var_remap_; + std::unordered_map var_remap_; private: using TSelf = ExprMutator; diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py index 59f8b81695e9..c28ca1b2c0bb 100644 --- a/python/tvm/relax/__init__.py +++ b/python/tvm/relax/__init__.py @@ -26,7 +26,6 @@ from .expr import ( Expr, Span, - Id, GlobalVar, Var, DataflowVar, diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py index 18f7d80312cd..c92a08e868cd 100644 --- a/python/tvm/relax/expr.py +++ b/python/tvm/relax/expr.py @@ -30,7 +30,6 @@ import tvm.relax import tvm.runtime from tvm import DataType -from tvm.runtime import Object from ..ir import BaseFunc, Node, Span from ..runtime import Scriptable @@ -75,18 +74,6 @@ def prim_value(value: Expr | int | float, dtype: str | None = None) -> Expr: raise TypeError(f"Cannot convert {value} with type {type(value)} to `Expr`") -@tvm_ffi.register_object("relax.Id") -class Id(Object): - """Unique identifier(name) used in Var. - Guaranteed to be stable across all passes. - """ - - name_hint: str - - def __init__(self): - raise RuntimeError("Cannot directly construct Id") - - def _relax_type_is_base_of(self: Type, derived: Type) -> bool: """Check if this Relax type is a base of another Relax type.""" @@ -429,7 +416,7 @@ class Var(ExprWithOp): Parameters ---------- - name_hint: str | Id + name_hint: str The name hint of the variable. ty: Optional[Type] @@ -439,12 +426,12 @@ class Var(ExprWithOp): Span that points to original source code """ - vid: Id + name_hint: str span: Span | None def __init__( self, - name_hint: str | Id, + name_hint: str, ty: Type | None = None, span: Span | None = None, ) -> None: @@ -457,18 +444,12 @@ def __init__( "use relax.TensorType(shape, dtype)." ) self.__init_handle_by_constructor__( - _ffi_api.Var if isinstance(name_hint, str) else _ffi_api.VarFromId, # type: ignore + _ffi_api.Var, # type: ignore name_hint, ty, span, ) - @property - def name_hint(self) -> str: - """Get name hint of the current var.""" - name = str(self.vid.name_hint) - return name - @tvm_ffi.register_object("relax.expr.DataflowVar") class DataflowVar(Var): @@ -478,7 +459,7 @@ class DataflowVar(Var): Parameters ---------- - name_hint: str | Id + name_hint: str The name hint of the variable. ty: Optional[Type] @@ -488,12 +469,12 @@ class DataflowVar(Var): Span that points to original source code """ - vid: Id + name_hint: str span: Span | None def __init__( self, - name_hint: str | Id, + name_hint: str, ty: Type | None = None, span: Span | None = None, ) -> None: @@ -507,16 +488,7 @@ def __init__( "use relax.TensorType(shape, dtype)." ) - self.__init_handle_by_constructor__( - ( - _ffi_api.DataflowVar # type: ignore - if isinstance(name_hint, str) - else _ffi_api.DataflowVarFromId - ), # type: ignore - name_hint, - ty, - span, - ) + self.__init_handle_by_constructor__(_ffi_api.DataflowVar, name_hint, ty, span) # type: ignore @tvm_ffi.register_object("relax.expr.StringImm") diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py index 9f0e398424df..39da280eeebb 100644 --- a/python/tvm/relax/expr_functor.py +++ b/python/tvm/relax/expr_functor.py @@ -40,7 +40,6 @@ ExternFunc, Function, GlobalVar, - Id, If, MatchCast, SeqExpr, @@ -1453,26 +1452,28 @@ def visit_expr_post_order(self, expr: Expr) -> Expr: """ return _ffi_api.PyExprMutatorVisitExprPostOrder(self._outer(), expr) # type: ignore - def set_var_remap(self, vid: Id, var: Var) -> None: + def set_var_remap(self, old_var: Var, new_var: Var) -> None: """Remap a var to a new var in use-site. Parameters ---------- - vid : Id - The vid of the old var. - var : Var + old_var : Var + The old var. + new_var : Var The new var. """ # Using self._outer() to ref _PyExprMutator - return _ffi_api.PyExprMutatorSetVarRemap(self._outer(), vid, var) # type: ignore + return _ffi_api.PyExprMutatorSetVarRemap( # type: ignore + self._outer(), old_var, new_var + ) - def get_var_remap(self, vid: Id) -> Var: + def get_var_remap(self, var: Var) -> Var: """Remap a var to a new var in use-site. Parameters ---------- - vid : Id - The vid of the old var + var : Var + The old var. Returns ------- @@ -1480,7 +1481,7 @@ def get_var_remap(self, vid: Id) -> Var: The remapped var. """ # Using self._outer() to ref _PyExprMutator - return _ffi_api.PyExprMutatorGetVarRemap(self._outer(), vid) # type: ignore + return _ffi_api.PyExprMutatorGetVarRemap(self._outer(), var) # type: ignore def visit_with_new_scope(self, expr: Expr) -> Expr: """Rewrite the expr with a new scope, used in a Function's body and the branches of If. diff --git a/python/tvm/relax/transform/lazy_transform_params.py b/python/tvm/relax/transform/lazy_transform_params.py index 4eb3b39619f6..2711e3287d49 100644 --- a/python/tvm/relax/transform/lazy_transform_params.py +++ b/python/tvm/relax/transform/lazy_transform_params.py @@ -312,7 +312,7 @@ def visit_var_binding_(self, binding: relax.VarBinding) -> None: if binding.var == self.func_creator.out_tuple_var: # The function after rewriting returns a empty tuple. func_output = self.builder_.emit(relax.Tuple([])) - self.set_var_remap(binding.var.vid, func_output) + self.set_var_remap(binding.var, func_output) return super().visit_var_binding_(binding) diff --git a/src/relax/distributed/transform/lower_distir.cc b/src/relax/distributed/transform/lower_distir.cc index 2f6b79ecb30c..66aebdcfc371 100644 --- a/src/relax/distributed/transform/lower_distir.cc +++ b/src/relax/distributed/transform/lower_distir.cc @@ -137,7 +137,7 @@ class DistIRSharder : public ExprMutator { if (sharding_spec->kind == PlacementSpecKind::kReplica) { Var new_var = builder_->Emit(broadcast_from_worker0(new_expr)); if (const auto* var = old_expr.as()) { - var_remap_[var->vid] = new_var; + var_remap_[ffi::GetRef(var)] = new_var; } else { tuple_getitem_remap_[old_expr.as_or_throw()] = new_var; } @@ -145,7 +145,7 @@ class DistIRSharder : public ExprMutator { Var scatter_var = builder_->Emit( scatter_from_worker0(new_expr, dtensor_ty->device_mesh->shape[0], sharding_spec->axis)); if (const auto* var = old_expr.as()) { - var_remap_[var->vid] = scatter_var; + var_remap_[ffi::GetRef(var)] = scatter_var; } else { tuple_getitem_remap_[old_expr.as_or_throw()] = scatter_var; } @@ -174,7 +174,7 @@ class DistIRSharder : public ExprMutator { ffi::Array new_params; for (const Var& var : func->params) { Var new_param = ShardInputParamTensorAndConstant(var).as_or_throw(); - var_remap_[var->vid] = new_param; + var_remap_[var] = new_param; new_params.push_back(new_param); } func_ = func; @@ -186,7 +186,7 @@ class DistIRSharder : public ExprMutator { void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val) { if (tuple_getitem_remap_.count(ffi::GetRef(val))) { - var_remap_[binding->var->vid] = tuple_getitem_remap_[ffi::GetRef(val)]; + var_remap_[binding->var] = tuple_getitem_remap_[ffi::GetRef(val)]; } else { ExprMutator::VisitBinding_(binding, val); } diff --git a/src/relax/distributed/transform/propagate_sharding.cc b/src/relax/distributed/transform/propagate_sharding.cc index 28db86b0a315..7e001ee16bc1 100644 --- a/src/relax/distributed/transform/propagate_sharding.cc +++ b/src/relax/distributed/transform/propagate_sharding.cc @@ -566,7 +566,7 @@ class DistributedIRBuilder : public ExprMutator { new_value = InsertRedistribute(new_value, device_mesh, placements[0]); } if (const auto* var = new_value.as()) { - var_remap_[binding->var->vid] = ffi::GetRef(var); + var_remap_[binding->var] = ffi::GetRef(var); } else { ReEmitBinding(binding, builder_->Normalize(new_value)); } @@ -574,7 +574,7 @@ class DistributedIRBuilder : public ExprMutator { const auto* inferred_tuple_ty = new_call->ty.as(); TVM_FFI_ICHECK(inferred_tuple_ty) << new_call; Var new_var = builder_->Emit(new_call); - var_remap_[binding->var->vid] = new_var; + var_remap_[binding->var] = new_var; for (int i = 0; i < static_cast(inferred_tuple_ty->fields.size()); i++) { if (!ffi::StructuralEqual()( DTensorType(inferred_tuple_ty->fields[i].as_or_throw()->tensor_ty, @@ -590,16 +590,17 @@ class DistributedIRBuilder : public ExprMutator { void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val) { if (tuple_getitem_remap_.count(ffi::GetRef(val))) { - var_remap_[binding->var->vid] = tuple_getitem_remap_[ffi::GetRef(val)]; + var_remap_[binding->var] = tuple_getitem_remap_[ffi::GetRef(val)]; } else { ExprMutator::VisitBinding_(binding, val); } } Expr VisitExpr_(const VarNode* var) final { - auto it = input_tensor_remap_.find(ffi::GetRef(var)); + Var var_ref = ffi::GetRef(var); + auto it = input_tensor_remap_.find(var_ref); if (it != input_tensor_remap_.end()) { - var_remap_[var->vid] = (*it).second; + var_remap_[var_ref] = (*it).second; } return ExprMutator::VisitExpr_(var); } diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc index 8048cda6b287..916c71629d27 100644 --- a/src/relax/ir/block_builder.cc +++ b/src/relax/ir/block_builder.cc @@ -152,7 +152,7 @@ class BlockBuilderImpl : public BlockBuilderNode { // Scope management //------------------------------- ffi::Optional LookupBinding(const Var& var) final { - auto it = binding_table_.find(var->vid); + auto it = binding_table_.find(var); if (it == binding_table_.end()) return std::nullopt; return it->second; } @@ -276,7 +276,7 @@ class BlockBuilderImpl : public BlockBuilderNode { TVM_FFI_ICHECK(!var_binding->var->ty.IsMissing()); TVM_FFI_ICHECK(!var_binding->value->ty.IsMissing()); cur_frame->bindings.push_back(binding); - binding_table_[var_binding->var->vid] = var_binding->value; + binding_table_[var_binding->var] = var_binding->value; } else if (const auto* match_cast = binding.as()) { if (!cur_frame->is_dataflow) { TVM_FFI_ICHECK(!match_cast->var.as()) @@ -342,7 +342,7 @@ class BlockBuilderImpl : public BlockBuilderNode { std::vector scope_stack_; /*! \brief A binding table that maps var to value. */ - std::unordered_map binding_table_; + std::unordered_map binding_table_; /*! \brief A unique name supply to get unique names for IR construction. */ UniqueNameSupply name_supply_; @@ -393,7 +393,7 @@ class BlockBuilderImpl : public BlockBuilderNode { CurrentBindingBlockFrame()->bindings.push_back(VarBinding(var, expr)); // update the binding table - binding_table_[var->vid] = expr; + binding_table_[var] = expr; return var; } @@ -408,9 +408,9 @@ class BlockBuilderImpl : public BlockBuilderNode { if (name_hint.empty()) { name_hint = is_dataflow ? "lv" : "gv"; } - Id vid = Id(GetUniqueName(name_hint)); - return is_dataflow ? DataflowVar(vid, /*ty_annotation=*/std::nullopt) - : Var(vid, /*ty_annotation=*/std::nullopt); + name_hint = GetUniqueName(name_hint); + return is_dataflow ? DataflowVar(name_hint, /*ty_annotation=*/std::nullopt) + : Var(name_hint, /*ty_annotation=*/std::nullopt); } private: diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc index dc54fef78ff2..0841f4a17ed2 100644 --- a/src/relax/ir/expr.cc +++ b/src/relax/ir/expr.cc @@ -28,7 +28,6 @@ namespace tvm { namespace relax { TVM_FFI_STATIC_INIT_BLOCK() { - IdNode::RegisterReflection(); TupleNode::RegisterReflection(); TupleGetItemNode::RegisterReflection(); ShapeExprNode::RegisterReflection(); @@ -48,12 +47,6 @@ TVM_FFI_STATIC_INIT_BLOCK() { ExternFuncNode::RegisterReflection(); } -Id::Id(ffi::String name_hint) { - ffi::ObjectPtr n = ffi::make_object(); - n->name_hint = std::move(name_hint); - data_ = std::move(n); -} - If::If(Expr cond, Expr true_branch, Expr false_branch, Span span) { ffi::ObjectPtr n = ffi::make_object(); n->cond = std::move(cond); @@ -146,9 +139,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { }); } -Var::Var(Id vid, ffi::Optional ty_annotation, Span span) { +Var::Var(ffi::String name_hint, ffi::Optional ty_annotation, Span span) { ffi::ObjectPtr n = ffi::make_object(); - n->vid = std::move(vid); + n->name_hint_ = std::move(name_hint); if (ty_annotation.has_value()) { n->ty = ty_annotation.value(); } @@ -156,38 +149,15 @@ Var::Var(Id vid, ffi::Optional ty_annotation, Span span) { data_ = std::move(n); } -VarNode* Var::CopyOnWrite() { - // The `TVM_DEFINE_OBJECT_REF_COW_METHOD` cannot be used for - // Var, because it is the base class for `DataflowBlock`. - // If the `TVM_DEFINE_OBJECT_REF_COW_METHOD` were used, the - // automatic implementation would erroneously convert from a - // `DataflowBlock` to a `Var`. - TVM_FFI_ICHECK(data_ != nullptr); - if (!data_.unique()) { - ffi::ObjectPtr node; - if (auto dataflow_var = as()) { - node = ffi::make_object(*dataflow_var); - } else { - node = ffi::make_object(*(operator->())); - } - ffi::ObjectPtr(std::move(node)).swap(data_); - } - return static_cast(data_.get()); -} - TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef() - .def("relax.Var", [](ffi::String name_hint, ffi::Optional ty_annotation, - Span span) { return Var(name_hint, ty_annotation, span); }) - .def("relax.VarFromId", [](Id vid, ffi::Optional ty_annotation, Span span) { - return Var(vid, ty_annotation, span); - }); + refl::GlobalDef().def("relax.Var", [](ffi::String name_hint, ffi::Optional ty_annotation, + Span span) { return Var(name_hint, ty_annotation, span); }); } -DataflowVar::DataflowVar(Id vid, ffi::Optional ty_annotation, Span span) { +DataflowVar::DataflowVar(ffi::String name_hint, ffi::Optional ty_annotation, Span span) { ffi::ObjectPtr n = ffi::make_object(); - n->vid = std::move(vid); + n->name_hint_ = std::move(name_hint); if (ty_annotation.has_value()) { n->ty = ty_annotation.value(); } @@ -197,14 +167,10 @@ DataflowVar::DataflowVar(Id vid, ffi::Optional ty_annotation, Span span) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef() - .def("relax.DataflowVar", - [](ffi::String name_hint, ffi::Optional ty_annotation, Span span) { - return DataflowVar(name_hint, ty_annotation, span); - }) - .def("relax.DataflowVarFromId", [](Id vid, ffi::Optional ty_annotation, Span span) { - return DataflowVar(vid, ty_annotation, span); - }); + refl::GlobalDef().def("relax.DataflowVar", + [](ffi::String name_hint, ffi::Optional ty_annotation, Span span) { + return DataflowVar(name_hint, ty_annotation, span); + }); } Constant::Constant(runtime::Tensor data, ffi::Optional ty_annotation, Span span) { diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc index 7f9ff73759a4..1869cf2881bd 100644 --- a/src/relax/ir/expr_functor.cc +++ b/src/relax/ir/expr_functor.cc @@ -544,19 +544,36 @@ PrimExpr ExprMutatorBase::VisitTypePrimExprField(const PrimExpr& expr) { // ================== // ExprMutator +namespace { + +template +void RedirectVarRemapTargets(TMap* var_remap, const Var& old_target, const Var& new_target) { + if (old_target.same_as(new_target)) { + return; + } + for (auto& entry : *var_remap) { + if (entry.second.same_as(old_target)) { + entry.second = new_target; + } + } +} + +} // namespace + Expr ExprMutator::VisitExpr(const Expr& expr) { return builder_->Normalize(ExprFunctor::VisitExpr(expr)); } // Visit the use-site of a defined Var Expr ExprMutator::VisitExpr_(const VarNode* op) { - auto it = var_remap_.find(op->vid); + Var var = ffi::GetRef(op); + auto it = var_remap_.find(var); if (it != var_remap_.end()) { return it->second; } // default case return self. - return ffi::GetRef(op); + return var; } // Visit the use-site of a defined DataflowVar @@ -571,7 +588,7 @@ Expr ExprMutator::VisitExpr_(const FunctionNode* op) { Var new_param = this->VisitVarDef(param); params.push_back(new_param); if (!param.same_as(new_param)) { - var_remap_[param->vid] = new_param; + var_remap_[param] = new_param; all_params_unchanged = false; } } @@ -660,6 +677,7 @@ RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(DataTypeImmNode); void ExprMutator::ReEmitBinding(const VarBindingNode* binding, Expr new_value) { Var new_var = this->VisitVarDef(binding->var); + Var visited_var = new_var; // fast path: re-emit binding if nothing changes if (new_var.same_as(binding->var) && new_value.same_as(binding->value)) { @@ -679,8 +697,10 @@ void ExprMutator::ReEmitBinding(const VarBindingNode* binding, Expr new_value) { new_var = temp; } - this->var_remap_[binding->var->vid] = new_var; - this->var_remap_[new_var->vid] = new_var; + RedirectVarRemapTargets(&var_remap_, visited_var, new_var); + this->var_remap_[binding->var] = new_var; + this->var_remap_[visited_var] = new_var; + this->var_remap_[new_var] = new_var; builder_->EmitNormalized(VarBinding(new_var, new_value)); } @@ -690,6 +710,7 @@ void ExprMutator::VisitBinding_(const MatchCastNode* binding) { Type new_ty = this->VisitExprDepTypeField(binding->ty); Var new_var = this->VisitVarDef(binding->var); + Var visited_var = new_var; MatchCast new_binding = [&]() -> MatchCast { if (new_var.same_as(binding->var) && new_value.same_as(binding->value) && @@ -700,8 +721,10 @@ void ExprMutator::VisitBinding_(const MatchCastNode* binding) { new_value = builder_->NormalizeArgument(new_value); new_var = WithType(new_var, new_ty); - var_remap_[binding->var->vid] = new_var; - var_remap_[new_var->vid] = new_var; + RedirectVarRemapTargets(&var_remap_, visited_var, new_var); + var_remap_[binding->var] = new_var; + var_remap_[visited_var] = new_var; + var_remap_[new_var] = new_var; return MatchCast(new_var, new_value, new_ty, binding->span); } @@ -733,7 +756,9 @@ Var ExprMutator::VisitVarDef_(const DataflowVarNode* var) { // provide default behavior in subclasses, we may produce a Var // where we should produce a DataflowVar. if (!output->IsInstance()) { - output = DataflowVar(output->vid, GetType(output), output->span); + Var delegated_output = output; + output = DataflowVar(output->name_hint(), GetType(output), output->span); + var_remap_[delegated_output] = output; } return output; } @@ -744,7 +769,7 @@ Var ExprMutator::VisitVarDef_(const VarNode* var) { if (ty.same_as(var->ty)) { return ffi::GetRef(var); } else { - return Var(var->vid, ty, var->span); + return Var(var->name_hint(), ty, var->span); } } else { return ffi::GetRef(var); @@ -830,6 +855,9 @@ Expr ExprMutator::VisitWithInnerScope(const Expr& expr) { } ffi::Optional ExprMutator::LookupBinding(const Var& var) { + if (auto it = var_remap_.find(var); it != var_remap_.end()) { + return builder_->LookupBinding(it->second); + } return builder_->LookupBinding(var); } @@ -842,8 +870,8 @@ Var ExprMutator::WithType(Var var, Type ty) { if (var->ty.same_as(ty) || ffi::StructuralEqual()(var->ty, ty)) { return var; } else { - Var new_var = var.as() ? DataflowVar(var->vid, ty, var->span) - : Var(var->vid, ty, var->span); + Var new_var = var.as() ? DataflowVar(var->name_hint(), ty, var->span) + : Var(var->name_hint(), ty, var->span); return new_var; } } else { diff --git a/src/relax/ir/py_expr_functor.cc b/src/relax/ir/py_expr_functor.cc index 125794866335..be64ae645914 100644 --- a/src/relax/ir/py_expr_functor.cc +++ b/src/relax/ir/py_expr_functor.cc @@ -680,9 +680,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("relax.PyExprMutatorWithType", [](PyExprMutator mutator, Var var, Type ty) { return mutator->WithType(var, ty); }) .def("relax.PyExprMutatorSetVarRemap", - [](PyExprMutator mutator, Id id, Var var) { return mutator->var_remap_[id] = var; }) + [](PyExprMutator mutator, Var old_var, Var new_var) { + return mutator->var_remap_[old_var] = new_var; + }) .def("relax.PyExprMutatorGetVarRemap", - [](PyExprMutator mutator, Id id) { return mutator->var_remap_[id]; }); + [](PyExprMutator mutator, Var var) { return mutator->var_remap_[var]; }); } TVM_FFI_STATIC_INIT_BLOCK() { diff --git a/src/relax/script/builder/frame.cc b/src/relax/script/builder/frame.cc index 1aa7eb9e2088..0daa192ae66b 100644 --- a/src/relax/script/builder/frame.cc +++ b/src/relax/script/builder/frame.cc @@ -136,7 +136,7 @@ void BindingBlockFrameNode::EnterWithScope() { class VarReplacer : public tvm::relax::ExprMutator { public: explicit VarReplacer( - std::unordered_map + std::unordered_map var_remap) { var_remap_ = std::move(var_remap); } @@ -144,7 +144,7 @@ class VarReplacer : public tvm::relax::ExprMutator { tvm::relax::Var VisitVarDef(const tvm::relax::Var& var) override { // ExprMutator only applies var_remap_ at usage sites. This // applies var_remap_ at each definition site as well. - if (auto it = var_remap_.find(var->vid); it != var_remap_.end()) { + if (auto it = var_remap_.find(var); it != var_remap_.end()) { return it->second; } else { return var; @@ -168,12 +168,12 @@ void BindingBlockFrameNode::ExitWithScope() { if (is_dataflow) { // Step 3.0. Define a map to replace variables ffi::Array new_output_vars; - std::unordered_map + std::unordered_map var_remap; for (const auto& output_var : output_vars) { tvm::relax::Var new_output_var(output_var->name_hint(), GetType(output_var)); new_output_vars.push_back(new_output_var); - var_remap[output_var->vid] = new_output_var; + var_remap[output_var] = new_output_var; } VarReplacer mutator(std::move(var_remap)); diff --git a/src/relax/script/builder/ir.cc b/src/relax/script/builder/ir.cc index 35b513b07b50..3c8681fccdb6 100644 --- a/src/relax/script/builder/ir.cc +++ b/src/relax/script/builder/ir.cc @@ -36,20 +36,16 @@ using tvm::script::ir_builder::details::Namer; TVM_STATIC_IR_FUNCTOR(Namer, vtable) .set_dispatch([](const ffi::ObjectRef& node, ffi::String name) -> void { using tvm::relax::VarNode; - using tvm::relax::IdNode; - const VarNode* var = node.as(); - IdNode* vid = const_cast(var->vid.get()); - vid->name_hint = name; + VarNode* var = const_cast(node.as()); + var->name_hint_ = name; }); TVM_STATIC_IR_FUNCTOR(Namer, vtable) .set_dispatch([](const ffi::ObjectRef& node, ffi::String name) -> void { using tvm::relax::DataflowVarNode; - using tvm::relax::IdNode; - const DataflowVarNode* var = node.as(); - IdNode* vid = const_cast(var->vid.get()); - vid->name_hint = name; + DataflowVarNode* var = const_cast(node.as()); + var->name_hint_ = name; }); /////////////////////////////// Function //////////////////////////////// diff --git a/src/relax/training/utils.cc b/src/relax/training/utils.cc index 33376e6d58ab..78072820f85f 100644 --- a/src/relax/training/utils.cc +++ b/src/relax/training/utils.cc @@ -167,7 +167,7 @@ class AppendLossMutator : private ExprMutator { << " while the corresponding type of parameter of loss function is " << loss_param_ty << ", which is different."; - this->var_remap_[loss_param->vid] = backbone_ret; + this->var_remap_[loss_param] = backbone_ret; } } @@ -190,8 +190,8 @@ class AppendLossMutator : private ExprMutator { for (int i = 0; i < num_backbone_outputs_; ++i) { auto var = backbone_return_arr_[i]; if (other_outputs_var.count(var) == 0) { - auto new_var = DataflowVar(var->vid, GetType(var), var->span); - this->var_remap_[var->vid] = new_var; + auto new_var = DataflowVar(var->name_hint(), GetType(var), var->span); + this->var_remap_[var] = new_var; backbone_return_arr_.Set(i, new_var); } } diff --git a/src/relax/transform/canonicalize_bindings.cc b/src/relax/transform/canonicalize_bindings.cc index 9ad3a6dda414..ca988351b6c0 100644 --- a/src/relax/transform/canonicalize_bindings.cc +++ b/src/relax/transform/canonicalize_bindings.cc @@ -152,10 +152,10 @@ class SymbolicVarCanonicalizer : public ExprMutator { }; struct CanonicalizationPlan { - ffi::Map replace_usage; - ffi::Map replace_binding; - std::unordered_set bindings_to_remove; - ffi::Map inline_constant; + ffi::Map replace_usage; + ffi::Map replace_binding; + std::unordered_set bindings_to_remove; + ffi::Map inline_constant; }; /*! \brief Utility class to identify usage location @@ -181,13 +181,13 @@ class CanonicalizePlanner : public ExprVisitor { if (!var.as() && !visitor.used_outside_home_dataflow_.count(var)) { DataflowVar new_var(var->name_hint(), GetType(var)); - plan.replace_binding.Set(var->vid, new_var); - plan.replace_usage.Set(var->vid, new_var); + plan.replace_binding.Set(var, new_var); + plan.replace_usage.Set(var, new_var); } } for (const auto& [var, constant] : visitor.known_bound_to_constant_) { - plan.inline_constant.Set(var->vid, constant); + plan.inline_constant.Set(var, constant); } for (const auto& binding_iter : visitor.trivial_bindings_) { @@ -200,7 +200,7 @@ class CanonicalizePlanner : public ExprVisitor { // non-trivial binding. bound_to = opt.value(); } - while (auto opt = plan.replace_binding.Get(bound_to->vid)) { + while (auto opt = plan.replace_binding.Get(bound_to)) { // The variable we are binding to may have already been // replaced, if it fell into Case 4 (Var = DataflowVar). In // that case, we check against its replacement instead. @@ -217,8 +217,8 @@ class CanonicalizePlanner : public ExprVisitor { // // For these four cases, the trivial binding can be unwrapped, // using the bound variable directly at the point of use. - plan.replace_usage.Set(bound_var->vid, bound_to); - plan.bindings_to_remove.insert(bound_var->vid); + plan.replace_usage.Set(bound_var, bound_to); + plan.bindings_to_remove.insert(bound_var); } else { // Case 4b: Var = DataflowVar, where the Var is used somewhere // outside the DataflowBlock containing the binding @@ -227,9 +227,9 @@ class CanonicalizePlanner : public ExprVisitor { // use of a DataflowVar outside of a DataflowBlock. Instead, // we replace in the opposite direction, replacing the binding // of the DataflowVar with a binding of the Var. - plan.replace_binding.Set(bound_to->vid, bound_var); - plan.replace_usage.Set(bound_to->vid, bound_var); - plan.bindings_to_remove.insert(bound_var->vid); + plan.replace_binding.Set(bound_to, bound_var); + plan.replace_usage.Set(bound_to, bound_var); + plan.bindings_to_remove.insert(bound_var); } } @@ -433,14 +433,14 @@ class BindingCanonicalizer : public ExprMutator { explicit BindingCanonicalizer(CanonicalizationPlan plan) : plan_(plan) {} void VisitBinding(const Binding& binding) override { - if (!plan_.bindings_to_remove.count(binding->var->vid)) { + if (!plan_.bindings_to_remove.count(binding->var)) { ExprMutator::VisitBinding(binding); } } Var VisitVarDef(const Var& var) override { Var new_var = var; - while (auto opt = plan_.replace_binding.Get(new_var->vid)) { + while (auto opt = plan_.replace_binding.Get(new_var)) { new_var = opt.value(); } @@ -449,10 +449,10 @@ class BindingCanonicalizer : public ExprMutator { Expr VisitExpr_(const VarNode* var) override { Var new_var = ffi::GetRef(var); - while (auto opt = plan_.replace_usage.Get(new_var->vid)) { + while (auto opt = plan_.replace_usage.Get(new_var)) { new_var = opt.value(); } - if (auto opt = plan_.inline_constant.Get(new_var->vid)) { + if (auto opt = plan_.inline_constant.Get(new_var)) { return VisitExpr(opt.value()); } diff --git a/src/relax/transform/convert_layout.cc b/src/relax/transform/convert_layout.cc index fe38bbd3aec6..28c30b708d57 100644 --- a/src/relax/transform/convert_layout.cc +++ b/src/relax/transform/convert_layout.cc @@ -339,7 +339,7 @@ class LayoutConvertMutator : public ExprMutator { } else { Var new_var = builder_->EmitMatchCast(RewriteExpr(binding->value, input_layout), new_ty); var_layout_map_[binding->var] = input_layout; - this->var_remap_[binding->var->vid] = new_var; + this->var_remap_[binding->var] = new_var; } } diff --git a/src/relax/transform/eliminate_common_subexpr.cc b/src/relax/transform/eliminate_common_subexpr.cc index b2805bf9a2c7..0542d029109e 100644 --- a/src/relax/transform/eliminate_common_subexpr.cc +++ b/src/relax/transform/eliminate_common_subexpr.cc @@ -141,7 +141,7 @@ class CommonSubexprEliminator : public ExprMutator { << "and occurrences of " << binding->var << " will be replaced with " << it->second[0]; output_binding = VarBinding(binding->var, it->second[0]); - var_remap_.insert({binding->var->vid, it->second[0]}); + var_remap_.insert({binding->var, it->second[0]}); it->second.push_back(binding->var); } else { diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc index dae9d0d36ee2..0f4a2ad560a4 100644 --- a/src/relax/transform/fuse_ops.cc +++ b/src/relax/transform/fuse_ops.cc @@ -530,7 +530,7 @@ class FunctionCreator : public ExprMutator { auto it = tuple_get_item_remap.find(tuple_get_item->tuple.get()); if (it != tuple_get_item_remap.end()) { TVM_FFI_ICHECK(it->second.find(tuple_get_item->index) != it->second.end()); - var_remap_[var_binding->var->vid] = it->second[tuple_get_item->index]; + var_remap_[var_binding->var] = it->second[tuple_get_item->index]; if (auto output_idx = GetOutputIndex(binding->var)) { outputs.Set(*output_idx, it->second[tuple_get_item->index]); } @@ -545,7 +545,7 @@ class FunctionCreator : public ExprMutator { const auto* var_binding = binding.as(); TVM_FFI_ICHECK_NOTNULL(var_binding); Var output_var = builder_->EmitOutput(VisitExpr(var_binding->value)); - var_remap_[var_binding->var->vid] = output_var; + var_remap_[var_binding->var] = output_var; outputs.Set(*output_idx, output_var); } else { // Case 2. It is an internal binding, add it to the binding list. @@ -881,10 +881,10 @@ class OperatorFusor : public ExprMutator { // available in pending_tuple_get and tuple_get_indices_ respectively. for (const auto& var : pending_tuple_get[group]) { auto tuple_get = TupleGetItem(new_var, tuple_get_indices_[var.get()]); - var_remap_[var->vid] = builder_->Emit(tuple_get); + var_remap_[var] = builder_->Emit(tuple_get); } } else { - var_remap_[var_binding->var->vid] = new_var; + var_remap_[var_binding->var] = new_var; } } // Step 5. Finish the binding block generation. diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc index 056f765ae9f3..499e88f8c520 100644 --- a/src/relax/transform/gradient.cc +++ b/src/relax/transform/gradient.cc @@ -46,7 +46,7 @@ namespace relax { // We will use NestedMsg to handle adjoint updates involving tuple handling using AdjointMsg = NestedMsg; -using VarIdSet = std::unordered_set; +using VarSet = std::unordered_set; // Used in CallTIRWithGradCollector. call_tir -> call_tir_with_grad using CallTIRWithGradInfo = std::unordered_map; @@ -63,7 +63,7 @@ class CallTIRWithGradEliminator : private ExprMutator { * * \param func The original function * \return The function with all start_checkpoint and end_checkpoint bindings removed, and a - * VarIdSet containing all checkpointed vars. + * VarSet containing all checkpointed vars. */ static Function Transform(const Function& func) { return CallTIRWithGradEliminator().VisitExpr(func).as_or_throw(); @@ -109,14 +109,14 @@ class CheckpointCollector : private ExprMutator { } // checkpointed vars - VarIdSet checkpoints; + VarSet checkpoints; // mapping from vars that are wrapped in start_checkpoint or end_checkpoint to the original vars - std::unordered_map var_mapping; + std::unordered_map var_mapping; private: Expr VisitExpr_(const FunctionNode* func) final { for (auto var : func->params) { - checkpoints.insert(var->vid); + checkpoints.insert(var); } return ExprMutator::VisitExpr_(func); @@ -137,13 +137,14 @@ class CheckpointCollector : private ExprMutator { bool all_inner_var_checkpointed = true; PostOrderVisit(var_binding->value, [this, &all_inner_var_checkpointed](const Expr& expr) { if (auto var = expr.as()) { + Var var_ref = ffi::GetRef(var); all_inner_var_checkpointed &= - (checkpoints.count(var->vid) != 0 || e_vars_.count(var->vid) != 0); + (checkpoints.count(var_ref) != 0 || e_vars_.count(var_ref) != 0); } }); if (all_inner_var_checkpointed) { - checkpoints.insert(var_binding->var->vid); + checkpoints.insert(var_binding->var); } } @@ -165,17 +166,17 @@ class CheckpointCollector : private ExprMutator { // Add remapping from binding->var to new_var if (!binding->var.as() && var->IsInstance()) { // For output binding, emit a dummy binding - this->var_remap_[binding->var->vid] = builder_->EmitOutput(orig_var, orig_var->name_hint()); + this->var_remap_[binding->var] = builder_->EmitOutput(orig_var, orig_var->name_hint()); } else { - this->var_remap_[binding->var->vid] = orig_var; + this->var_remap_[binding->var] = orig_var; } - var_mapping[binding->var->vid] = orig_var; + var_mapping[binding->var] = orig_var; if (value->op.same_as(s_cp)) { // mark the original var to be checkpointed - checkpoints.insert(orig_var->vid); + checkpoints.insert(orig_var); } else if (value->op.same_as(e_cp)) { - e_vars_.insert(binding->var->vid); + e_vars_.insert(binding->var); } } else { ExprMutator::VisitBinding_(binding, value); @@ -183,7 +184,7 @@ class CheckpointCollector : private ExprMutator { } // vars that are the output of end_checkpoint - VarIdSet e_vars_; + VarSet e_vars_; }; /*! @@ -205,7 +206,7 @@ class CheckpointGenerator : private ExprMutator { * checkpointed */ CheckpointGenerator(const BlockBuilder& builder, const ffi::Array& orig_params, - const DataflowBlock& forward_block, const VarIdSet& checkpoints) + const DataflowBlock& forward_block, const VarSet& checkpoints) : builder_(builder) { // func params will always be checkpointed for (auto var : orig_params) { @@ -217,7 +218,7 @@ class CheckpointGenerator : private ExprMutator { TVM_FFI_ICHECK(var_binding) << "Now only support VarBindingNode"; auto var = var_binding->var; binding_map_.Set(var, var_binding->value); - if (checkpoints.count(var->vid)) { + if (checkpoints.count(var)) { checkpoint_map_.Set(var, var); } } @@ -499,8 +500,8 @@ class BackwardBindingGenerator : private ExprVisitor { for (Var var : require_grads) { // var might be wrapped in start_checkpoint or end_checkpoint, so we should find the original // var first - if (cp_collector_.var_mapping.count(var->vid)) { - var = cp_collector_.var_mapping[var->vid]; + if (cp_collector_.var_mapping.count(var)) { + var = cp_collector_.var_mapping[var]; } // If the var don't have adjoint var, it do not contribute to the target. So its adjoint is // zeros @@ -762,15 +763,15 @@ class GradientMutator : private ExprMutator { static ffi::Array CheckAndMapRequireGrads(const ffi::Array& require_grads, const ffi::Map& var_map, const ffi::String& func_name) { - VarIdSet var_set; + VarSet var_set; ffi::Array mapped_vars; for (const auto& var : require_grads) { auto it = var_map.find(var); TVM_FFI_ICHECK(it != var_map.end()) << "There is no Var named " << var->name_hint() << " in the function " << func_name; - TVM_FFI_ICHECK_EQ(var_set.count(var->vid), 0) + TVM_FFI_ICHECK_EQ(var_set.count(var), 0) << "Var " << var->name_hint() << " appears more than once"; - var_set.emplace(var->vid); + var_set.emplace(var); mapped_vars.push_back((*it).second); TVM_FFI_ICHECK(IsNestedTensorConditioned(GetType(var), IsFloatTensorType)) diff --git a/src/relax/transform/gradient_simplifier.cc b/src/relax/transform/gradient_simplifier.cc index 17d72c3ac268..13fda06faf44 100644 --- a/src/relax/transform/gradient_simplifier.cc +++ b/src/relax/transform/gradient_simplifier.cc @@ -166,7 +166,7 @@ class GradientSimplifier : private ExprMutator { if (IsTransposeOp(prev_call_node)) { // rewrite rule #1: permute_dims(permute_dims(a)) -> a if (prev_call_node->args[0]->IsInstance()) { - var_remap_[binding->var->vid] = prev_call_node->args[0].as_or_throw(); + var_remap_[binding->var] = prev_call_node->args[0].as_or_throw(); return; } else { return reemit_and_return(); diff --git a/src/relax/transform/lambda_lift.cc b/src/relax/transform/lambda_lift.cc index 2bdfb7178efa..c7b2b1ea2fc1 100644 --- a/src/relax/transform/lambda_lift.cc +++ b/src/relax/transform/lambda_lift.cc @@ -367,9 +367,10 @@ class LambdaLifter : public ExprMutator { // Call "relax.invoke_closure" to invoke closure - if (IsClosure(var) && builder_->LookupBinding(var).as()) { + auto bound_value = LookupBinding(var); + if (IsClosure(var) && bound_value.as()) { // if the original op was pure, we should use invoke_pure_closure - Call orig_call = builder_->LookupBinding(var).value().as_or_throw(); + Call orig_call = bound_value.value().as_or_throw(); bool is_pure = [&]() -> bool { if (auto op = orig_call->op.as()) { static const auto& purity_map = Op::GetAttrMap("FPurity"); @@ -422,7 +423,7 @@ class LambdaLifter : public ExprMutator { if (closures_.count(opt_var.value())) { return true; } - if (auto bound_value = builder_->LookupBinding(opt_var.value())) { + if (auto bound_value = LookupBinding(opt_var.value())) { val = bound_value.value(); } } diff --git a/src/relax/transform/realize_vdevice.cc b/src/relax/transform/realize_vdevice.cc index 2897c0e909c9..7e93209390e0 100644 --- a/src/relax/transform/realize_vdevice.cc +++ b/src/relax/transform/realize_vdevice.cc @@ -363,9 +363,9 @@ class VDeviceTypeUpdater : ExprMutator { }(); if (var->IsInstance()) { - var = DataflowVar(var->vid, new_ty, var->span); + var = DataflowVar(var->name_hint(), new_ty, var->span); } else { - var = Var(var->vid, new_ty, var->span); + var = Var(var->name_hint(), new_ty, var->span); } } } diff --git a/src/relax/transform/rewrite_cuda_graph.cc b/src/relax/transform/rewrite_cuda_graph.cc index 3b53e3ac3b7f..78acb764b7ae 100644 --- a/src/relax/transform/rewrite_cuda_graph.cc +++ b/src/relax/transform/rewrite_cuda_graph.cc @@ -138,7 +138,7 @@ class FuncBuilder : public ExprMutator { // Set up the parameters for (const auto* input : inputs_) { auto new_var = Var(input->name_hint(), VisitExprDepTypeField(input->ty.as_or_throw())); - var_remap_[input->vid] = new_var; + var_remap_[ffi::GetRef(input)] = new_var; params.push_back(new_var); } if (shape_expr) { @@ -846,7 +846,7 @@ class CUDAGraphRewriter : public ExprMutator { LaunchSubgraph(op, subgraph_launches_[op->var.get()]); } if (auto it = var_redef_.find(op->var.get()); - it != var_redef_.end() && !var_remap_.count(op->var->vid)) { + it != var_redef_.end() && !var_remap_.count(op->var)) { EmitRedef(op->var.get(), it->second); return; } @@ -858,7 +858,8 @@ class CUDAGraphRewriter : public ExprMutator { } Expr VisitExpr_(const VarNode* op) final { - if (auto it = var_remap_.find(op->vid); it != var_remap_.end()) { + Var var = ffi::GetRef(op); + if (auto it = var_remap_.find(var); it != var_remap_.end()) { return it->second; } if (auto it = var_redef_.find(op); it != var_redef_.end()) { @@ -871,7 +872,7 @@ class CUDAGraphRewriter : public ExprMutator { Var EmitRedef(const VarNode* var, const Expr& redef) { auto new_var = builder_->Emit(redef, var->name_hint()); - var_remap_[var->vid] = new_var; + var_remap_[ffi::GetRef(var)] = new_var; return new_var; } diff --git a/src/relax/transform/to_mixed_precision.cc b/src/relax/transform/to_mixed_precision.cc index ad2e285015d1..b13772cf0bf8 100644 --- a/src/relax/transform/to_mixed_precision.cc +++ b/src/relax/transform/to_mixed_precision.cc @@ -279,7 +279,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { private: Var GetRemapped(const Var& var) { - auto it = var_remap_.find(var->vid); + auto it = var_remap_.find(var); if (it != var_remap_.end()) { return it->second; } else { @@ -291,8 +291,8 @@ class ToMixedPrecisionRewriter : public ExprMutator { vdev = tensor_ty->vdevice.value(); } TensorType fp16_ty(tensor_ty->shape.value(), PrimType::Float(16), vdev, tensor_ty->span); - Var fp16_var(var->vid, fp16_ty, var->span); - var_remap_[var->vid] = fp16_var; + Var fp16_var(var->name_hint(), fp16_ty, var->span); + var_remap_[var] = fp16_var; return fp16_var; } } @@ -424,13 +424,13 @@ class ToMixedPrecisionRewriter : public ExprMutator { // If cur_var is not rewritten, we don't need to emit a new var if (!rewrite.same_as(cur_var)) { // Emit a new var, and update the var remap - var_remap_[var->vid] = builder_->Emit(rewrite); + var_remap_[var] = builder_->Emit(rewrite); } } Expr VisitVar_(const Var& var) { // We rewrite the remapped var to the original dtype - auto it = var_remap_.find(var->vid); + auto it = var_remap_.find(var); if (it != var_remap_.end()) { return RewriteExpr(it->second, NTypeFrom(var)); } @@ -577,7 +577,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { } for (auto param : params_) { // remove the local version of params - auto it = var_remap_.find(param->vid); + auto it = var_remap_.find(param); if (it != var_remap_.end()) { var_remap_.erase(it); } diff --git a/src/relax/transform/to_non_dataflow.cc b/src/relax/transform/to_non_dataflow.cc index 700521eb9cb8..e8701d5c9c3e 100644 --- a/src/relax/transform/to_non_dataflow.cc +++ b/src/relax/transform/to_non_dataflow.cc @@ -34,8 +34,8 @@ class ToNonDFMutator : public ExprMutator { public: Var VisitVarDef(const Var& var) final { if (var.as()) { - Var new_var = Var(var->vid, GetType(var), var->span); - this->var_remap_[var->vid] = new_var; + Var new_var = Var(var->name_hint(), GetType(var), var->span); + this->var_remap_[var] = new_var; return new_var; } return var; diff --git a/src/relax/transform/update_param_type.cc b/src/relax/transform/update_param_type.cc index 7ba696639fd2..dcebbd4e6f4d 100644 --- a/src/relax/transform/update_param_type.cc +++ b/src/relax/transform/update_param_type.cc @@ -52,7 +52,7 @@ class ParamTypeMutator : public ExprMutator { auto params = op->params.Map([this](Var param) { if (auto new_ty = ty_func_(param)) { auto new_param = WithType(param, new_ty.value()); - var_remap_[param->vid] = new_param; + var_remap_[param] = new_param; return new_param; } else { return param; diff --git a/src/relax/transform/utils.h b/src/relax/transform/utils.h index 0163b9327519..e62351bc99f5 100644 --- a/src/relax/transform/utils.h +++ b/src/relax/transform/utils.h @@ -199,7 +199,7 @@ bool IsNestedTensor(const Expr& expr); // TODO(@bohan): implements some postorder function accepts a visitor closure class VarReplacer : public ExprMutator { public: - using VarMap = std::unordered_map; + using VarMap = std::unordered_map; explicit VarReplacer(const VarMap& var_remap) : var_remap_(var_remap) {} @@ -211,7 +211,7 @@ class VarReplacer : public ExprMutator { private: Expr VisitExpr_(const VarNode* op) final { Var var = ffi::GetRef(op); - auto it = var_remap_.find(var->vid); + auto it = var_remap_.find(var); return it == var_remap_.end() ? var : it->second; } @@ -269,7 +269,7 @@ class SymbolicVarRenewMutator : public ExprMutator, tirx::ExprMutator { Var new_param = this->VisitVarDef(param); params.push_back(new_param); if (!param.same_as(new_param)) { - var_remap_[param->vid] = new_param; + var_remap_[param] = new_param; all_params_unchanged = false; } } @@ -304,7 +304,7 @@ class FunctionCopier : public SymbolicVarRenewMutator { Var VisitVarDef_(const DataflowVarNode* var) override { Var new_var = SymbolicVarRenewMutator::VisitVarDef_(var); Var copied_var = DataflowVar(new_var->name_hint(), GetType(new_var), new_var->span); - var_remap_[var->vid] = copied_var; + var_remap_[ffi::GetRef(var)] = copied_var; relax_var_map_.Set(ffi::GetRef(var), copied_var); return copied_var; } @@ -312,7 +312,7 @@ class FunctionCopier : public SymbolicVarRenewMutator { Var VisitVarDef_(const VarNode* var) override { Var new_var = SymbolicVarRenewMutator::VisitVarDef_(var); Var copied_var = Var(new_var->name_hint(), GetType(new_var), new_var->span); - var_remap_[var->vid] = copied_var; + var_remap_[ffi::GetRef(var)] = copied_var; relax_var_map_.Set(ffi::GetRef(var), copied_var); return copied_var; } diff --git a/src/relax/utils.cc b/src/relax/utils.cc index a8b148aa8943..63f0d08636a3 100644 --- a/src/relax/utils.cc +++ b/src/relax/utils.cc @@ -49,7 +49,7 @@ class ExprBinder : public ExprMutator { Var new_param = this->VisitVarDef(param); params.push_back(new_param); if (!param.same_as(new_param)) { - this->var_remap_[param->vid] = new_param; + this->var_remap_[param] = new_param; all_params_unchanged = false; } } diff --git a/tests/python/relax/test_expr_functor.py b/tests/python/relax/test_expr_functor.py index cf004a879a74..6a09ee80837d 100644 --- a/tests/python/relax/test_expr_functor.py +++ b/tests/python/relax/test_expr_functor.py @@ -303,7 +303,7 @@ def visit_var_binding_(self, binding: VarBinding) -> None: temp = self.with_type(new_var, new_value.ty) if not temp.same_as(new_var): new_var = temp - self.set_var_remap(binding.var.vid, new_var) + self.set_var_remap(binding.var, new_var) self.builder_.emit_normalized(VarBinding(new_var, new_value)) @@ -315,7 +315,7 @@ def visit_match_cast_(self, binding: MatchCast) -> None: temp = self.with_type(new_var, binding.ty) if not temp.same_as(new_var): new_var = temp - self.set_var_remap(binding.var.vid, new_var) + self.set_var_remap(binding.var, new_var) self.log.add("MatchCast") self.builder_.emit_normalized(MatchCast(new_var, new_value, binding.ty)) diff --git a/tests/python/relax/test_transform_canonicalize_bindings.py b/tests/python/relax/test_transform_canonicalize_bindings.py index 4650483546af..17b4a466b85b 100644 --- a/tests/python/relax/test_transform_canonicalize_bindings.py +++ b/tests/python/relax/test_transform_canonicalize_bindings.py @@ -1173,12 +1173,12 @@ def test_canonicalization_causes_ty_update(): better type as a result. If this happens, the In previous implementations, ExprMutator::ReEmitBinding defined a - remap for `binding->var->vid`, even if the derived class defined a + remap for `binding->var`, even if the derived class defined a replacement by overriding `VisitVarDef`. If the derived class defines a new variable binding by overriding `VisitVarDef`, and also causes a variable replacement by overriding `VisitExpr` and returning a type with different type, then `ExprMutator` - must check for both `binding->var->vid` *AND* `new_var->vid`. The + must check for both `binding->var` *AND* `new_var`. The former may be present in the unmodified graph, and the latter may be produced by the derived class before delegating to the base class. From c4e73613441385eae58682a570b2196156cd278b Mon Sep 17 00:00:00 2001 From: tqchen Date: Mon, 6 Jul 2026 16:04:13 +0000 Subject: [PATCH 2/3] [Relax] Keep variable remapping one-hop --- src/relax/ir/expr_functor.cc | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc index 1869cf2881bd..e886c845b8be 100644 --- a/src/relax/ir/expr_functor.cc +++ b/src/relax/ir/expr_functor.cc @@ -544,22 +544,6 @@ PrimExpr ExprMutatorBase::VisitTypePrimExprField(const PrimExpr& expr) { // ================== // ExprMutator -namespace { - -template -void RedirectVarRemapTargets(TMap* var_remap, const Var& old_target, const Var& new_target) { - if (old_target.same_as(new_target)) { - return; - } - for (auto& entry : *var_remap) { - if (entry.second.same_as(old_target)) { - entry.second = new_target; - } - } -} - -} // namespace - Expr ExprMutator::VisitExpr(const Expr& expr) { return builder_->Normalize(ExprFunctor::VisitExpr(expr)); } @@ -697,7 +681,6 @@ void ExprMutator::ReEmitBinding(const VarBindingNode* binding, Expr new_value) { new_var = temp; } - RedirectVarRemapTargets(&var_remap_, visited_var, new_var); this->var_remap_[binding->var] = new_var; this->var_remap_[visited_var] = new_var; this->var_remap_[new_var] = new_var; @@ -721,7 +704,6 @@ void ExprMutator::VisitBinding_(const MatchCastNode* binding) { new_value = builder_->NormalizeArgument(new_value); new_var = WithType(new_var, new_ty); - RedirectVarRemapTargets(&var_remap_, visited_var, new_var); var_remap_[binding->var] = new_var; var_remap_[visited_var] = new_var; var_remap_[new_var] = new_var; From 900bf3caf0aab0935ef7998dacc12245a8a512ea Mon Sep 17 00:00:00 2001 From: tqchen Date: Mon, 6 Jul 2026 17:37:56 +0000 Subject: [PATCH 3/3] [Relax] Use direct variable name hint field --- include/tvm/relax/expr.h | 7 ++----- src/relax/analysis/var2value.cc | 4 ++-- src/relax/backend/contrib/codegen_c/codegen_c.h | 2 +- .../backend/contrib/codegen_json/codegen_json.h | 2 +- src/relax/backend/contrib/cutlass/codegen.cc | 2 +- src/relax/backend/vm/codegen_vm.cc | 2 +- src/relax/backend/vm/codegen_vm_tir.cc | 2 +- src/relax/backend/vm/vm_shape_lower.cc | 2 +- src/relax/distributed/transform/lower_distir.cc | 2 +- .../distributed/transform/propagate_sharding.cc | 6 +++--- src/relax/ir/binding_rewrite.cc | 2 +- src/relax/ir/block_builder.cc | 2 +- src/relax/ir/dataflow_matcher.cc | 2 +- src/relax/ir/expr.cc | 4 ++-- src/relax/ir/expr_functor.cc | 8 ++++---- src/relax/ir/transform.cc | 8 ++++---- src/relax/script/builder/frame.cc | 2 +- src/relax/script/builder/ir.cc | 4 ++-- src/relax/script/builder/utils.h | 6 +++--- src/relax/script/printer/utils.h | 2 +- src/relax/training/utils.cc | 2 +- src/relax/transform/bind_params.cc | 4 ++-- src/relax/transform/bundle_model_params.cc | 2 +- src/relax/transform/canonicalize_bindings.cc | 2 +- src/relax/transform/expand_tuple_arguments.cc | 2 +- src/relax/transform/fuse_ops.cc | 10 +++++----- src/relax/transform/fuse_tir.cc | 6 +++--- src/relax/transform/gradient.cc | 16 ++++++++-------- src/relax/transform/inline_functions.cc | 2 +- src/relax/transform/lambda_lift.cc | 4 ++-- src/relax/transform/lazy_transform_params.cc | 6 +++--- src/relax/transform/lift_transform_params.cc | 4 ++-- src/relax/transform/realize_vdevice.cc | 4 ++-- src/relax/transform/rewrite_cuda_graph.cc | 4 ++-- .../specialize_primfunc_based_on_callsite.cc | 2 +- src/relax/transform/to_mixed_precision.cc | 4 ++-- src/relax/transform/to_non_dataflow.cc | 2 +- src/relax/transform/utils.h | 4 ++-- 38 files changed, 74 insertions(+), 77 deletions(-) diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h index 0e67dc277c18..b8ed8cfc2342 100644 --- a/include/tvm/relax/expr.h +++ b/include/tvm/relax/expr.h @@ -139,14 +139,11 @@ class VarNode : public ExprNode { * \brief The hint to the variable name. * \note Each variable is uniquely identified by its address. */ - ffi::String name_hint_; - - /*! \return The name hint of the variable */ - const ffi::String& name_hint() const { return name_hint_; } + ffi::String name_hint; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("name_hint", &VarNode::name_hint_, + refl::ObjectDef().def_ro("name_hint", &VarNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore()); } diff --git a/src/relax/analysis/var2value.cc b/src/relax/analysis/var2value.cc index 3af778efe4ff..a5e30c813950 100644 --- a/src/relax/analysis/var2value.cc +++ b/src/relax/analysis/var2value.cc @@ -74,12 +74,12 @@ class Name2BindingAnalysis : public relax::ExprVisitor { // so we use standard container for internal usage. std::map> name2bindings_; void VisitBinding_(const VarBindingNode* binding) override { - const auto& vname = binding->var->name_hint(); + const auto& vname = binding->var->name_hint; name2bindings_[vname].push_back(ffi::GetRef(binding)); } void VisitBinding_(const MatchCastNode* binding) override { - const auto& vname = binding->var->name_hint(); + const auto& vname = binding->var->name_hint; name2bindings_[vname].push_back(ffi::GetRef(binding)); } }; diff --git a/src/relax/backend/contrib/codegen_c/codegen_c.h b/src/relax/backend/contrib/codegen_c/codegen_c.h index b183fd85f950..58f76c92a826 100644 --- a/src/relax/backend/contrib/codegen_c/codegen_c.h +++ b/src/relax/backend/contrib/codegen_c/codegen_c.h @@ -282,7 +282,7 @@ class CodegenCBase { for (const auto& arg : args) { const auto& dtype_str = GetDtypeString(arg); - code_stream_ << dtype_str << "* " << arg->name_hint() << ", "; + code_stream_ << dtype_str << "* " << arg->name_hint << ", "; } for (size_t i = 0; i < outs.size() - 1; ++i) { code_stream_ << outs[i].dtype << "* out" << i << ", "; diff --git a/src/relax/backend/contrib/codegen_json/codegen_json.h b/src/relax/backend/contrib/codegen_json/codegen_json.h index 06c4733e7010..7c9f1a1ea750 100644 --- a/src/relax/backend/contrib/codegen_json/codegen_json.h +++ b/src/relax/backend/contrib/codegen_json/codegen_json.h @@ -238,7 +238,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator { void serialize(Function func) { // First we convert all the parameters into input nodes. for (const auto& param : func->params) { - auto node_ptr = std::make_shared(param->name_hint(), "input" /* op_type_ */); + auto node_ptr = std::make_shared(param->name_hint, "input" /* op_type_ */); memo_[param] = AddNode(node_ptr, param); } heads_ = VisitExpr(func->body); diff --git a/src/relax/backend/contrib/cutlass/codegen.cc b/src/relax/backend/contrib/cutlass/codegen.cc index 2894b75685b4..a2992892660e 100644 --- a/src/relax/backend/contrib/cutlass/codegen.cc +++ b/src/relax/backend/contrib/cutlass/codegen.cc @@ -157,7 +157,7 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, void AddParm(Var param) { ext_func_args_.push_back(param); - auto v_name = name_sup_->FreshName(param->name_hint()); + auto v_name = name_sup_->FreshName(param->name_hint); var_name_map_[param.get()] = v_name; } diff --git a/src/relax/backend/vm/codegen_vm.cc b/src/relax/backend/vm/codegen_vm.cc index 380671473d60..c27a39279fc1 100644 --- a/src/relax/backend/vm/codegen_vm.cc +++ b/src/relax/backend/vm/codegen_vm.cc @@ -91,7 +91,7 @@ class CodeGenVM : public ExprFunctor { ffi::Array param_names; for (Var param : func->params) { - param_names.push_back(param->name_hint()); + param_names.push_back(param->name_hint); } builder_->EmitFunction(gsymbol.value(), func->params.size(), param_names); diff --git a/src/relax/backend/vm/codegen_vm_tir.cc b/src/relax/backend/vm/codegen_vm_tir.cc index 439a855e590e..ee515489daf5 100644 --- a/src/relax/backend/vm/codegen_vm_tir.cc +++ b/src/relax/backend/vm/codegen_vm_tir.cc @@ -171,7 +171,7 @@ class CodeGenVMTIR : public ExprFunctor(const Expr&)> { ffi::Array param_names; for (Var param : func->params) { - param_names.push_back(param->name_hint()); + param_names.push_back(param->name_hint); } // declare this function. builder_->DeclareFunction(gsymbol.value(), vm::VMFuncInfo::FuncKind::kVMTIRFunc); diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc index bc2f5d485c1e..3fb8ea04349d 100644 --- a/src/relax/backend/vm/vm_shape_lower.cc +++ b/src/relax/backend/vm/vm_shape_lower.cc @@ -305,7 +305,7 @@ class VMShapeLowerMutator Type ty = GetType(func->params[i]); std::ostringstream err_ctx; err_ctx << "ErrorContext(fn=" << gvar->name_hint << ", loc=param[" << i - << "], param=" << func->params[i]->name_hint() << ", annotation=" << ty << ") "; + << "], param=" << func->params[i]->name_hint << ", annotation=" << ty << ") "; this->CheckMatchCast(ty, func->params[i], true, i >= num_input, err_ctx.str(), &match_todos); } diff --git a/src/relax/distributed/transform/lower_distir.cc b/src/relax/distributed/transform/lower_distir.cc index 66aebdcfc371..01522183bb8c 100644 --- a/src/relax/distributed/transform/lower_distir.cc +++ b/src/relax/distributed/transform/lower_distir.cc @@ -116,7 +116,7 @@ class DistIRSharder : public ExprMutator { Type old_ty = GetType(input); Type new_ty = ConvertType(old_ty, false); if (const auto* var = input.as()) { - Var new_param(var->name_hint(), new_ty); + Var new_param(var->name_hint, new_ty); return new_param; } else if (const auto* constant = input.as()) { for (const auto& spec : old_ty.as_or_throw()->placement->dim_specs) { diff --git a/src/relax/distributed/transform/propagate_sharding.cc b/src/relax/distributed/transform/propagate_sharding.cc index 7e001ee16bc1..738d5dc0b218 100644 --- a/src/relax/distributed/transform/propagate_sharding.cc +++ b/src/relax/distributed/transform/propagate_sharding.cc @@ -285,7 +285,7 @@ class ShardingConflictHandler : public ExprVisitor { if (device_mesh.has_value()) { TVM_FFI_ICHECK(ffi::StructuralEqual()(device_mesh.value(), sharding_spec.first)) - << "Sharding conflict detected for tensor " << var->name_hint() + << "Sharding conflict detected for tensor " << var->name_hint << ": Device Mesh mismatch" << ". Conflict Handling logic will be added in the future."; } else { @@ -294,7 +294,7 @@ class ShardingConflictHandler : public ExprVisitor { if (i >= 0) { int sharding_dim = sharding_spec.second; TVM_FFI_ICHECK(sharded_mesh_dim.count(sharding_dim) == 0) - << "Sharding conflict detected for tensor " << var->name_hint() + << "Sharding conflict detected for tensor " << var->name_hint << ": Replicate sharding device mesh axis " << sharding_dim << ". Conflict Handling logic will be added in the future."; sharded_mesh_dim.insert(sharding_dim); @@ -401,7 +401,7 @@ class DistributedIRBuilder : public ExprMutator { } if (const auto* var = tensor.as()) { - Var new_param(var->name_hint(), new_ty); + Var new_param(var->name_hint, new_ty); return new_param; } else if (const auto* constant = tensor.as()) { Constant new_constant(constant->data, new_ty); diff --git a/src/relax/ir/binding_rewrite.cc b/src/relax/ir/binding_rewrite.cc index af21c056736c..44348e6cefa6 100644 --- a/src/relax/ir/binding_rewrite.cc +++ b/src/relax/ir/binding_rewrite.cc @@ -55,7 +55,7 @@ DataflowBlockRewrite::DataflowBlockRewrite(DataflowBlock dfb, Function root_fn) n->to_users_ = std::move(p.first); n->fn_outputs_ = std::move(p.second); n->name_supply_ = UniqueNameSupply(n->to_users_.begin(), n->to_users_.end(), - [](const auto& p) { return p.first->name_hint(); }); + [](const auto& p) { return p.first->name_hint; }); data_ = std::move(n); } diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc index 916c71629d27..60cefd081d53 100644 --- a/src/relax/ir/block_builder.cc +++ b/src/relax/ir/block_builder.cc @@ -584,7 +584,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctorty.IsMissing()) << "Var " << var->name_hint() << " does not have type."; + TVM_FFI_ICHECK(!var->ty.IsMissing()) << "Var " << var->name_hint << " does not have type."; return ffi::GetRef(var); } diff --git a/src/relax/ir/dataflow_matcher.cc b/src/relax/ir/dataflow_matcher.cc index 3f7edc6ae8ed..ef881e317b98 100644 --- a/src/relax/ir/dataflow_matcher.cc +++ b/src/relax/ir/dataflow_matcher.cc @@ -582,7 +582,7 @@ bool DFPatternMatcher::VisitDFPattern_(const VarPatternNode* op, const Expr& exp // We don't jump for var pattern, as there's no need to access its value to judge it. if (const auto* var_node = expr.as()) { // "" means any name. - return "" == op->name_hint() || op->name_hint() == var_node->name_hint(); + return "" == op->name_hint() || op->name_hint() == var_node->name_hint; } return false; } diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc index 0841f4a17ed2..3da61a48fba2 100644 --- a/src/relax/ir/expr.cc +++ b/src/relax/ir/expr.cc @@ -141,7 +141,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Var::Var(ffi::String name_hint, ffi::Optional ty_annotation, Span span) { ffi::ObjectPtr n = ffi::make_object(); - n->name_hint_ = std::move(name_hint); + n->name_hint = std::move(name_hint); if (ty_annotation.has_value()) { n->ty = ty_annotation.value(); } @@ -157,7 +157,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { DataflowVar::DataflowVar(ffi::String name_hint, ffi::Optional ty_annotation, Span span) { ffi::ObjectPtr n = ffi::make_object(); - n->name_hint_ = std::move(name_hint); + n->name_hint = std::move(name_hint); if (ty_annotation.has_value()) { n->ty = ty_annotation.value(); } diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc index e886c845b8be..963d302936aa 100644 --- a/src/relax/ir/expr_functor.cc +++ b/src/relax/ir/expr_functor.cc @@ -739,7 +739,7 @@ Var ExprMutator::VisitVarDef_(const DataflowVarNode* var) { // where we should produce a DataflowVar. if (!output->IsInstance()) { Var delegated_output = output; - output = DataflowVar(output->name_hint(), GetType(output), output->span); + output = DataflowVar(output->name_hint, GetType(output), output->span); var_remap_[delegated_output] = output; } return output; @@ -751,7 +751,7 @@ Var ExprMutator::VisitVarDef_(const VarNode* var) { if (ty.same_as(var->ty)) { return ffi::GetRef(var); } else { - return Var(var->name_hint(), ty, var->span); + return Var(var->name_hint, ty, var->span); } } else { return ffi::GetRef(var); @@ -852,8 +852,8 @@ Var ExprMutator::WithType(Var var, Type ty) { if (var->ty.same_as(ty) || ffi::StructuralEqual()(var->ty, ty)) { return var; } else { - Var new_var = var.as() ? DataflowVar(var->name_hint(), ty, var->span) - : Var(var->name_hint(), ty, var->span); + Var new_var = var.as() ? DataflowVar(var->name_hint, ty, var->span) + : Var(var->name_hint, ty, var->span); return new_var; } } else { diff --git a/src/relax/ir/transform.cc b/src/relax/ir/transform.cc index 8b1815812eac..75f3873ce781 100644 --- a/src/relax/ir/transform.cc +++ b/src/relax/ir/transform.cc @@ -230,7 +230,7 @@ class DataflowBlockMutator : public ExprMutator { } } if (!var.as()) { - global_scope_vars.Set(var->name_hint(), var); + global_scope_vars.Set(var->name_hint, var); } } @@ -252,10 +252,10 @@ class DataflowBlockMutator : public ExprMutator { } } } - if (!var.as() && global_scope_vars.count(var->name_hint()) > 0) { - TVM_FFI_ICHECK(var.same_as(global_scope_vars[var->name_hint()])) + if (!var.as() && global_scope_vars.count(var->name_hint) > 0) { + TVM_FFI_ICHECK(var.same_as(global_scope_vars[var->name_hint])) << "Error: DataflowBlock Pass should not rewrite any GlobalScope Var."; - global_scope_vars.erase(var->name_hint()); + global_scope_vars.erase(var->name_hint); } } TVM_FFI_ICHECK(global_scope_vars.empty() && symbolic_vars.empty()) diff --git a/src/relax/script/builder/frame.cc b/src/relax/script/builder/frame.cc index 0daa192ae66b..80fec50fceb2 100644 --- a/src/relax/script/builder/frame.cc +++ b/src/relax/script/builder/frame.cc @@ -171,7 +171,7 @@ void BindingBlockFrameNode::ExitWithScope() { std::unordered_map var_remap; for (const auto& output_var : output_vars) { - tvm::relax::Var new_output_var(output_var->name_hint(), GetType(output_var)); + tvm::relax::Var new_output_var(output_var->name_hint, GetType(output_var)); new_output_vars.push_back(new_output_var); var_remap[output_var] = new_output_var; } diff --git a/src/relax/script/builder/ir.cc b/src/relax/script/builder/ir.cc index 3c8681fccdb6..32f9f66d59e4 100644 --- a/src/relax/script/builder/ir.cc +++ b/src/relax/script/builder/ir.cc @@ -37,7 +37,7 @@ TVM_STATIC_IR_FUNCTOR(Namer, vtable) .set_dispatch([](const ffi::ObjectRef& node, ffi::String name) -> void { using tvm::relax::VarNode; VarNode* var = const_cast(node.as()); - var->name_hint_ = name; + var->name_hint = name; }); TVM_STATIC_IR_FUNCTOR(Namer, vtable) @@ -45,7 +45,7 @@ TVM_STATIC_IR_FUNCTOR(Namer, vtable) ffi::String name) -> void { using tvm::relax::DataflowVarNode; DataflowVarNode* var = const_cast(node.as()); - var->name_hint_ = name; + var->name_hint = name; }); /////////////////////////////// Function //////////////////////////////// diff --git a/src/relax/script/builder/utils.h b/src/relax/script/builder/utils.h index 3a137b0ba452..f353465ee376 100644 --- a/src/relax/script/builder/utils.h +++ b/src/relax/script/builder/utils.h @@ -99,7 +99,7 @@ inline tvm::relax::SeqExpr GetSeqExprForBranch(const SeqExprFrame& frame, ffi::S TVM_FFI_ICHECK(!last_binding->var->IsInstance()) << "A non-dataflow var is expected in the last binding of '" << method << "'."; - *var_name = last_binding->var->name_hint(); + *var_name = last_binding->var->name_hint; // Step 3. Re-collect binding blocks to replace the last binding. ffi::Array new_blocks(frame->binding_blocks.begin(), @@ -107,8 +107,8 @@ inline tvm::relax::SeqExpr GetSeqExprForBranch(const SeqExprFrame& frame, ffi::S ffi::Array last_block_bindings(last_block->bindings.begin(), last_block->bindings.end() - 1); - tvm::relax::Var new_var = tvm::relax::Var(last_binding->var->name_hint() + output_var_suffix, - GetType(last_binding->var)); + tvm::relax::Var new_var = + tvm::relax::Var(last_binding->var->name_hint + output_var_suffix, GetType(last_binding->var)); tvm::relax::Expr body; const auto* var_binding = last_binding.as(); diff --git a/src/relax/script/printer/utils.h b/src/relax/script/printer/utils.h index cced3c97057b..b5d575666a1d 100644 --- a/src/relax/script/printer/utils.h +++ b/src/relax/script/printer/utils.h @@ -75,7 +75,7 @@ inline std::string ReprPrintRelax(const ffi::ObjectRef& obj, const PrinterConfig } inline IdDoc DefineVar(const relax::Var& var, const Frame& frame, const IRDocsifier& d) { - return d->Define(var, frame, var->name_hint().empty() ? "v" : var->name_hint()); + return d->Define(var, frame, var->name_hint.empty() ? "v" : var->name_hint); } inline ffi::Optional TypeAsAnn(const relax::Var& v, const AccessPath& v_p, diff --git a/src/relax/training/utils.cc b/src/relax/training/utils.cc index 78072820f85f..a279972e958e 100644 --- a/src/relax/training/utils.cc +++ b/src/relax/training/utils.cc @@ -190,7 +190,7 @@ class AppendLossMutator : private ExprMutator { for (int i = 0; i < num_backbone_outputs_; ++i) { auto var = backbone_return_arr_[i]; if (other_outputs_var.count(var) == 0) { - auto new_var = DataflowVar(var->name_hint(), GetType(var), var->span); + auto new_var = DataflowVar(var->name_hint, GetType(var), var->span); this->var_remap_[var] = new_var; backbone_return_arr_.Set(i, new_var); } diff --git a/src/relax/transform/bind_params.cc b/src/relax/transform/bind_params.cc index 293537888d49..04d5f33c6657 100644 --- a/src/relax/transform/bind_params.cc +++ b/src/relax/transform/bind_params.cc @@ -92,7 +92,7 @@ std::tuple, ffi::Map> NormalizeBindings std::unordered_map> string_lookup; std::unordered_set var_set; for (const auto& param : func->params) { - string_lookup[param->name_hint()].push_back(param); + string_lookup[param->name_hint].push_back(param); var_set.insert(param.get()); } @@ -105,7 +105,7 @@ std::tuple, ffi::Map> NormalizeBindings TVM_FFI_ICHECK(it != string_lookup.end()) << "Function does not have parameter with name \"" << str << "\". " << "Function parameters are named " - << func->params.Map([](const auto& param) { return param->name_hint(); }); + << func->params.Map([](const auto& param) { return param->name_hint; }); TVM_FFI_ICHECK_EQ(it->second.size(), 1) << "Function contains multiple parameters with name \"" << str << "\". " << "The Relax variables " << it->second << " are all named \"" << str << "\""; diff --git a/src/relax/transform/bundle_model_params.cc b/src/relax/transform/bundle_model_params.cc index 8ddc5539e199..7aa4c8acbe18 100644 --- a/src/relax/transform/bundle_model_params.cc +++ b/src/relax/transform/bundle_model_params.cc @@ -77,7 +77,7 @@ class ModelParamBundler : public ExprMutator { Expr VisitExpr_(const VarNode* op) override { auto var = ffi::GetRef(op); if (auto it = var_to_expr_.find(var); it != var_to_expr_.end()) { - return builder_->Emit((*it).second, op->name_hint()); + return builder_->Emit((*it).second, op->name_hint); } else { return ExprMutator::VisitExpr_(op); } diff --git a/src/relax/transform/canonicalize_bindings.cc b/src/relax/transform/canonicalize_bindings.cc index ca988351b6c0..cc12031d5f81 100644 --- a/src/relax/transform/canonicalize_bindings.cc +++ b/src/relax/transform/canonicalize_bindings.cc @@ -179,7 +179,7 @@ class CanonicalizePlanner : public ExprVisitor { // of trivial bindings, then we can replace it with a DataflowVar. for (auto var : visitor.defined_inside_dataflow_) { if (!var.as() && !visitor.used_outside_home_dataflow_.count(var)) { - DataflowVar new_var(var->name_hint(), GetType(var)); + DataflowVar new_var(var->name_hint, GetType(var)); plan.replace_binding.Set(var, new_var); plan.replace_usage.Set(var, new_var); diff --git a/src/relax/transform/expand_tuple_arguments.cc b/src/relax/transform/expand_tuple_arguments.cc index 8294e2c0f028..fa37f9ca07c8 100644 --- a/src/relax/transform/expand_tuple_arguments.cc +++ b/src/relax/transform/expand_tuple_arguments.cc @@ -50,7 +50,7 @@ ffi::Optional ExpandParams(Function func) { ffi::Array internal_tuple; for (size_t i = 0; i < ty->fields.size(); i++) { auto name = static_cast(std::stringstream() - << param->name_hint() << "_" << i) + << param->name_hint << "_" << i) .str(); Var new_param(name, ty->fields[i]); internal_tuple.push_back(new_param); diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc index 0f4a2ad560a4..15d84930920e 100644 --- a/src/relax/transform/fuse_ops.cc +++ b/src/relax/transform/fuse_ops.cc @@ -502,7 +502,7 @@ class FunctionCreator : public ExprMutator { TVM_FFI_ICHECK(!item_indices.empty()); int param_idx = tuple_param_idx_[tuple_arg]; Var param = params_[param_idx]; - ffi::String param_name = params_[param_idx]->name_hint(); + ffi::String param_name = params_[param_idx]->name_hint; TupleType param_ty = tuple_arg->ty.as_or_throw(); ffi::Array item_args; @@ -624,7 +624,7 @@ class FunctionCreator : public ExprMutator { if ((var == nullptr || defined_vars_.count(var) == 0) && (lift_constant_ || !expr->IsInstance())) { ffi::String name = var != nullptr - ? var->name_hint() + ? var->name_hint : ffi::String("param_" + std::to_string(n_param_for_const_++)); Type param_ty = GetType(expr); if (!IsInlinableConstants(expr)) { @@ -931,8 +931,8 @@ class OperatorFusor : public ExprMutator { if (producer_group != cur_group) { for (Group* depgroup : group_deps_[producer_group]) { TVM_FFI_ICHECK(depgroup != cur_group) - << "A cyclic dependency detected between the groups " << binding->var->name_hint() - << " and " << used_var->name_hint() << " are in."; + << "A cyclic dependency detected between the groups " << binding->var->name_hint + << " and " << used_var->name_hint << " are in."; } group_deps_[cur_group].push_back(producer_group); } @@ -1332,7 +1332,7 @@ class CompositeFunctionAnnotator : public ExprMutator { ffi::Array params; for (auto v : func_node->params) { - Var new_v(v->name_hint(), GetType(v)); + Var new_v(v->name_hint, GetType(v)); param_vars.push_back(new_v); params.push_back(new_v); } diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index 7dff781df21f..6aa6da3c1abb 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -968,7 +968,7 @@ class FusedTIRConstructor : public ExprVisitor { << "All tuple parameters should be expanded before this point in FuseTIR. " << "However, parameter " << relax_param << " has type " << ty; - auto name_hint = relax_param->name_hint(); + auto name_hint = relax_param->name_hint; if (const auto* tensor = ty.as()) { // Case 1. The relax param is a Tensor, we directly create a tirx var and buffer @@ -1263,8 +1263,8 @@ class TIRFuseMutator : public ExprMutator { } } else if (const auto* prim_value = ty.as()) { if (const auto* var = arg.as()) { - tir_vars.push_back(tirx::Var(var->name_hint(), tvm::PrimType(prim_value->dtype)) - .as_or_throw()); + tir_vars.push_back( + tirx::Var(var->name_hint, tvm::PrimType(prim_value->dtype)).as_or_throw()); } else if (auto literal = arg.as()) { tir_vars.push_back(literal.value()); } else { diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc index 499e88f8c520..f11d6af8c2b8 100644 --- a/src/relax/transform/gradient.cc +++ b/src/relax/transform/gradient.cc @@ -166,7 +166,7 @@ class CheckpointCollector : private ExprMutator { // Add remapping from binding->var to new_var if (!binding->var.as() && var->IsInstance()) { // For output binding, emit a dummy binding - this->var_remap_[binding->var] = builder_->EmitOutput(orig_var, orig_var->name_hint()); + this->var_remap_[binding->var] = builder_->EmitOutput(orig_var, orig_var->name_hint); } else { this->var_remap_[binding->var] = orig_var; } @@ -231,7 +231,7 @@ class CheckpointGenerator : private ExprMutator { if (it != checkpoint_map_.end()) { return std::make_pair((*it).second, new_value); } - auto new_var = builder_->Emit(new_value, var->name_hint() + "_cp"); + auto new_var = builder_->Emit(new_value, var->name_hint + "_cp"); checkpoint_map_.Set(var, new_var); return std::make_pair(new_var, new_value); } @@ -250,7 +250,7 @@ class CheckpointGenerator : private ExprMutator { if (it != checkpoint_map_.end()) { return (*it).second; } - Var new_var = builder_->Emit(VisitExpr(binding_map_[var]), var->name_hint() + "_cp"); + Var new_var = builder_->Emit(VisitExpr(binding_map_[var]), var->name_hint + "_cp"); checkpoint_map_.Set(var, new_var); return new_var; } @@ -520,9 +520,9 @@ class BackwardBindingGenerator : private ExprVisitor { Var EmitAdjoint(const Var& source_var, const Expr& adjoint, bool is_output) { Var adjoint_var; if (is_output) { - adjoint_var = builder_->EmitOutput(adjoint, source_var->name_hint() + "_adjoint_out"); + adjoint_var = builder_->EmitOutput(adjoint, source_var->name_hint + "_adjoint_out"); } else { - adjoint_var = builder_->Emit(adjoint, source_var->name_hint() + "_adjoint"); + adjoint_var = builder_->Emit(adjoint, source_var->name_hint + "_adjoint"); adjoint_var_map_.Set(source_var, adjoint_var); } return adjoint_var; @@ -768,16 +768,16 @@ class GradientMutator : private ExprMutator { for (const auto& var : require_grads) { auto it = var_map.find(var); TVM_FFI_ICHECK(it != var_map.end()) - << "There is no Var named " << var->name_hint() << " in the function " << func_name; + << "There is no Var named " << var->name_hint << " in the function " << func_name; TVM_FFI_ICHECK_EQ(var_set.count(var), 0) - << "Var " << var->name_hint() << " appears more than once"; + << "Var " << var->name_hint << " appears more than once"; var_set.emplace(var); mapped_vars.push_back((*it).second); TVM_FFI_ICHECK(IsNestedTensorConditioned(GetType(var), IsFloatTensorType)) << "Only Tensors of floating point dtype or Tuples of float " "Tensors can require gradients, but the Type of Var " - << var->name_hint() << " is " << GetType(var); + << var->name_hint << " is " << GetType(var); } return mapped_vars; } diff --git a/src/relax/transform/inline_functions.cc b/src/relax/transform/inline_functions.cc index 4e03e83de8a6..a15d89645718 100644 --- a/src/relax/transform/inline_functions.cc +++ b/src/relax/transform/inline_functions.cc @@ -129,7 +129,7 @@ class FunctionInliner : public ExprMutator { // // This implementation uses Option 4. - Var param_var(func->params[i]->name_hint(), args[i]->ty.as()); + Var param_var(func->params[i]->name_hint, args[i]->ty.as()); param_bindings.push_back(VarBinding(param_var, args[i])); param_map.Set(func->params[i], param_var); } diff --git a/src/relax/transform/lambda_lift.cc b/src/relax/transform/lambda_lift.cc index c7b2b1ea2fc1..6a418ec8152e 100644 --- a/src/relax/transform/lambda_lift.cc +++ b/src/relax/transform/lambda_lift.cc @@ -96,7 +96,7 @@ class LambdaNameCollector : ExprVisitor { lifted_with_global_symbol_.insert({func, public_name}); } - name_stack_.push_back(binding->var->name_hint()); + name_stack_.push_back(binding->var->name_hint); lambda_location_.insert({func, name_stack_}); ExprVisitor::VisitBinding_(binding, func); name_stack_.pop_back(); @@ -285,7 +285,7 @@ class LambdaLifter : public ExprMutator { ffi::Array typed_captured_vars; ffi::Map rebinding_map; for (auto free_var : captured_vars) { - Var var = Var(free_var->name_hint(), GetType(free_var), free_var->span); + Var var = Var(free_var->name_hint, GetType(free_var), free_var->span); typed_captured_vars.push_back(var); rebinding_map.Set(free_var, var); } diff --git a/src/relax/transform/lazy_transform_params.cc b/src/relax/transform/lazy_transform_params.cc index b5a5918c4a18..02fd26c05f25 100644 --- a/src/relax/transform/lazy_transform_params.cc +++ b/src/relax/transform/lazy_transform_params.cc @@ -100,10 +100,10 @@ class LazyInputMutator : public ExprMutator { auto untyped = builder_->Emit(Call(Type::Missing(), plan_->fget_param, { PrimExpr(IntImm::Int64(it->second)), - StringImm(var->name_hint()), + StringImm(var->name_hint), }), - var->name_hint() + "_untyped"); - return builder_->EmitMatchCast(untyped, GetType(var), var->name_hint()); + var->name_hint + "_untyped"); + return builder_->EmitMatchCast(untyped, GetType(var), var->name_hint); } } diff --git a/src/relax/transform/lift_transform_params.cc b/src/relax/transform/lift_transform_params.cc index a5d295f09af0..390721996868 100644 --- a/src/relax/transform/lift_transform_params.cc +++ b/src/relax/transform/lift_transform_params.cc @@ -106,7 +106,7 @@ struct BaseCollectInfo { } for (const auto& var : outputs) { - Var out_var(var->name_hint() + "_output", GetType(var)); + Var out_var(var->name_hint + "_output", GetType(var)); output_var_binding.push_back(VarBinding(out_var, var)); output_exprs.push_back(out_var); } @@ -280,7 +280,7 @@ struct LocalCollectInfo : public BaseCollectInfo { return global_outputs; }(); for (const auto& var : compile_time_outputs) { - Var param_var(var->name_hint(), GetType(var)); + Var param_var(var->name_hint, GetType(var)); bindings.push_back(VarBinding(var, param_var)); params.push_back(param_var); } diff --git a/src/relax/transform/realize_vdevice.cc b/src/relax/transform/realize_vdevice.cc index 7e93209390e0..5e0a990b7d4d 100644 --- a/src/relax/transform/realize_vdevice.cc +++ b/src/relax/transform/realize_vdevice.cc @@ -363,9 +363,9 @@ class VDeviceTypeUpdater : ExprMutator { }(); if (var->IsInstance()) { - var = DataflowVar(var->name_hint(), new_ty, var->span); + var = DataflowVar(var->name_hint, new_ty, var->span); } else { - var = Var(var->name_hint(), new_ty, var->span); + var = Var(var->name_hint, new_ty, var->span); } } } diff --git a/src/relax/transform/rewrite_cuda_graph.cc b/src/relax/transform/rewrite_cuda_graph.cc index 78acb764b7ae..585ff6bdebd8 100644 --- a/src/relax/transform/rewrite_cuda_graph.cc +++ b/src/relax/transform/rewrite_cuda_graph.cc @@ -137,7 +137,7 @@ class FuncBuilder : public ExprMutator { } // Set up the parameters for (const auto* input : inputs_) { - auto new_var = Var(input->name_hint(), VisitExprDepTypeField(input->ty.as_or_throw())); + auto new_var = Var(input->name_hint, VisitExprDepTypeField(input->ty.as_or_throw())); var_remap_[ffi::GetRef(input)] = new_var; params.push_back(new_var); } @@ -871,7 +871,7 @@ class CUDAGraphRewriter : public ExprMutator { } Var EmitRedef(const VarNode* var, const Expr& redef) { - auto new_var = builder_->Emit(redef, var->name_hint()); + auto new_var = builder_->Emit(redef, var->name_hint); var_remap_[ffi::GetRef(var)] = new_var; return new_var; } diff --git a/src/relax/transform/specialize_primfunc_based_on_callsite.cc b/src/relax/transform/specialize_primfunc_based_on_callsite.cc index b867ff967dab..8445446e6850 100644 --- a/src/relax/transform/specialize_primfunc_based_on_callsite.cc +++ b/src/relax/transform/specialize_primfunc_based_on_callsite.cc @@ -94,7 +94,7 @@ class SpecializeTIRCallArgs : ExprMutator { } ffi::String name; if (args[i]->IsInstance()) { - name = args[i].as_or_throw()->name_hint(); + name = args[i].as_or_throw()->name_hint; } else { name = std::string({static_cast('A' + i)}); } diff --git a/src/relax/transform/to_mixed_precision.cc b/src/relax/transform/to_mixed_precision.cc index b13772cf0bf8..0764fd268f39 100644 --- a/src/relax/transform/to_mixed_precision.cc +++ b/src/relax/transform/to_mixed_precision.cc @@ -283,7 +283,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { if (it != var_remap_.end()) { return it->second; } else { - if (fp16_input_names_.count(var->name_hint())) { + if (fp16_input_names_.count(var->name_hint)) { auto ty = GetType(var); if (auto tensor_ty = ty.as()) { VDevice vdev = VDevice(); @@ -291,7 +291,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { vdev = tensor_ty->vdevice.value(); } TensorType fp16_ty(tensor_ty->shape.value(), PrimType::Float(16), vdev, tensor_ty->span); - Var fp16_var(var->name_hint(), fp16_ty, var->span); + Var fp16_var(var->name_hint, fp16_ty, var->span); var_remap_[var] = fp16_var; return fp16_var; } diff --git a/src/relax/transform/to_non_dataflow.cc b/src/relax/transform/to_non_dataflow.cc index e8701d5c9c3e..ae279e4cd5bf 100644 --- a/src/relax/transform/to_non_dataflow.cc +++ b/src/relax/transform/to_non_dataflow.cc @@ -34,7 +34,7 @@ class ToNonDFMutator : public ExprMutator { public: Var VisitVarDef(const Var& var) final { if (var.as()) { - Var new_var = Var(var->name_hint(), GetType(var), var->span); + Var new_var = Var(var->name_hint, GetType(var), var->span); this->var_remap_[var] = new_var; return new_var; } diff --git a/src/relax/transform/utils.h b/src/relax/transform/utils.h index e62351bc99f5..105dbff5e920 100644 --- a/src/relax/transform/utils.h +++ b/src/relax/transform/utils.h @@ -303,7 +303,7 @@ class FunctionCopier : public SymbolicVarRenewMutator { Var VisitVarDef_(const DataflowVarNode* var) override { Var new_var = SymbolicVarRenewMutator::VisitVarDef_(var); - Var copied_var = DataflowVar(new_var->name_hint(), GetType(new_var), new_var->span); + Var copied_var = DataflowVar(new_var->name_hint, GetType(new_var), new_var->span); var_remap_[ffi::GetRef(var)] = copied_var; relax_var_map_.Set(ffi::GetRef(var), copied_var); return copied_var; @@ -311,7 +311,7 @@ class FunctionCopier : public SymbolicVarRenewMutator { Var VisitVarDef_(const VarNode* var) override { Var new_var = SymbolicVarRenewMutator::VisitVarDef_(var); - Var copied_var = Var(new_var->name_hint(), GetType(new_var), new_var->span); + Var copied_var = Var(new_var->name_hint, GetType(new_var), new_var->span); var_remap_[ffi::GetRef(var)] = copied_var; relax_var_map_.Set(ffi::GetRef(var), copied_var); return copied_var;