diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h index 51b4ed43bdb9..b8ed8cfc2342 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,28 @@ 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; - - /*! \return The name hint of the variable */ - const ffi::String& name_hint() const { return vid->name_hint; } + /*! + * \brief The hint to the variable name. + * \note Each variable is uniquely identified by its address. + */ + ffi::String 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); + refl::ObjectDef().def_ro("name_hint", &VarNode::name_hint, + refl::AttachFieldFlag::SEqHashIgnore()); } - 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; - } - - 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 +169,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/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 2f6b79ecb30c..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) { @@ -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..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); @@ -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/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 8048cda6b287..60cefd081d53 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: @@ -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 dc54fef78ff2..3da61a48fba2 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..963d302936aa 100644 --- a/src/relax/ir/expr_functor.cc +++ b/src/relax/ir/expr_functor.cc @@ -550,13 +550,14 @@ Expr ExprMutator::VisitExpr(const Expr& 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 +572,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 +661,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 +681,9 @@ 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; + 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 +693,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 +704,9 @@ 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; + 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 +738,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 +751,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 +837,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 +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->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/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 1aa7eb9e2088..80fec50fceb2 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)); + 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..32f9f66d59e4 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/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 33376e6d58ab..a279972e958e 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/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 9ad3a6dda414..cc12031d5f81 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 @@ -179,15 +179,15 @@ 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->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/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 dae9d0d36ee2..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; @@ -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. @@ -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)) { @@ -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. @@ -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 056f765ae9f3..f11d6af8c2b8 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); } } @@ -230,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); } @@ -249,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; } @@ -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 @@ -519,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; @@ -762,21 +763,21 @@ 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) - << "Var " << var->name_hint() << " appears more than once"; - var_set.emplace(var->vid); + << "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_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/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/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 2bdfb7178efa..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); } @@ -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/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 2897c0e909c9..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->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..585ff6bdebd8 100644 --- a/src/relax/transform/rewrite_cuda_graph.cc +++ b/src/relax/transform/rewrite_cuda_graph.cc @@ -137,8 +137,8 @@ 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; + 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); } 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()) { @@ -870,8 +871,8 @@ 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; + 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 ad2e285015d1..0764fd268f39 100644 --- a/src/relax/transform/to_mixed_precision.cc +++ b/src/relax/transform/to_mixed_precision.cc @@ -279,11 +279,11 @@ 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 { - 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,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..ae279e4cd5bf 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..105dbff5e920 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; } } @@ -303,16 +303,16 @@ 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 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; } 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 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; } 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.