This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 545bd7b3c7 Phase out Relax-specific Id aliases (#19959)
545bd7b3c7 is described below
commit 545bd7b3c727a3bd64e9de26df0c13467d13bc15
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Jul 7 06:28:57 2026 +0800
Phase out Relax-specific Id aliases (#19959)
Remove the Relax-specific Id indirection and use Var/DataflowVar object
identity directly.
Type-changing rewrites now remap definitions, uses, and binding lookups
coherently while preserving reflection, serialization, and the
DataflowVar subtype. Existing tests are adjusted for the API change; no
new test files or test cases are added.
Validation: the compiler and C++ tests build successfully; focused C++
coverage passes 4/4; the affected existing Python matrices pass 549
tests with 2 expected xfails; source censuses, diff checks, and
applicable hooks pass.
---
include/tvm/relax/expr.h | 86 +++-------------------
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 +++---
.../tvm/relax/transform/lazy_transform_params.py | 2 +-
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 | 10 +--
.../distributed/transform/propagate_sharding.cc | 17 +++--
src/relax/ir/binding_rewrite.cc | 2 +-
src/relax/ir/block_builder.cc | 16 ++--
src/relax/ir/dataflow_matcher.cc | 2 +-
src/relax/ir/expr.cc | 54 +++-----------
src/relax/ir/expr_functor.cc | 32 +++++---
src/relax/ir/py_expr_functor.cc | 6 +-
src/relax/ir/transform.cc | 8 +-
src/relax/script/builder/frame.cc | 10 +--
src/relax/script/builder/ir.cc | 12 +--
src/relax/script/builder/utils.h | 6 +-
src/relax/script/printer/utils.h | 2 +-
src/relax/training/utils.cc | 6 +-
src/relax/transform/bind_params.cc | 4 +-
src/relax/transform/bundle_model_params.cc | 2 +-
src/relax/transform/canonicalize_bindings.cc | 36 ++++-----
src/relax/transform/convert_layout.cc | 2 +-
src/relax/transform/eliminate_common_subexpr.cc | 2 +-
src/relax/transform/expand_tuple_arguments.cc | 2 +-
src/relax/transform/fuse_ops.cc | 18 ++---
src/relax/transform/fuse_tir.cc | 6 +-
src/relax/transform/gradient.cc | 55 +++++++-------
src/relax/transform/gradient_simplifier.cc | 2 +-
src/relax/transform/inline_functions.cc | 2 +-
src/relax/transform/lambda_lift.cc | 11 +--
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 | 13 ++--
.../specialize_primfunc_based_on_callsite.cc | 2 +-
src/relax/transform/to_mixed_precision.cc | 14 ++--
src/relax/transform/to_non_dataflow.cc | 4 +-
src/relax/transform/update_param_type.cc | 2 +-
src/relax/transform/utils.h | 14 ++--
src/relax/utils.cc | 2 +-
tests/python/relax/test_expr_functor.py | 4 +-
.../relax/test_transform_canonicalize_bindings.py | 4 +-
51 files changed, 229 insertions(+), 341 deletions(-)
diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h
index 51b4ed43bd..b8ed8cfc23 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<IdNode>().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<VarNode>().def_ro("vid", &VarNode::vid);
- // customize structural equal and hash to include ty
- refl::TypeAttrDef<VarNode>()
- .def("__s_equal__", &VarNode::SEqual)
- .def("__s_hash__", &VarNode::SHash);
+ refl::ObjectDef<VarNode>().def_ro("name_hint", &VarNode::name_hint,
+ refl::AttachFieldFlag::SEqHashIgnore());
}
- bool SEqual(const VarNode* other,
- ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal)
const {
- return equal(vid, other->vid, false, "vid") && equal(ty, other->ty, false,
"ty");
- }
-
- int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView,
int64_t, bool)> 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<Type>
ty_annotation, Span span = Span())
- : Var(Id(name_hint), ty_annotation, span) {}
-
- TVM_DLL explicit Var(Id vid, ffi::Optional<Type> ty_annotation, Span span =
Span());
+ TVM_DLL explicit Var(ffi::String name_hint, ffi::Optional<Type>
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<DataflowVarNode>();
}
- 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<Type>
ty_annotation,
- Span span = Span())
- : DataflowVar(Id(name_hint), ty_annotation, span) {}
-
- TVM_DLL explicit DataflowVar(Id vid, ffi::Optional<Type> 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 05261d8c5e..7f1388803d 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<Id, Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
var_remap_;
+ std::unordered_map<Var, Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
var_remap_;
private:
using TSelf = ExprMutator;
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index 59f8b81695..c28ca1b2c0 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -26,7 +26,6 @@ from tvm.ir import Call
from .expr import (
Expr,
Span,
- Id,
GlobalVar,
Var,
DataflowVar,
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index 18f7d80312..c92a08e868 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -30,7 +30,6 @@ import tvm.ir
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 @@ class Var(ExprWithOp):
"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 @@ class DataflowVar(Var):
"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 9f0e398424..39da280eee 100644
--- a/python/tvm/relax/expr_functor.py
+++ b/python/tvm/relax/expr_functor.py
@@ -40,7 +40,6 @@ from .expr import (
ExternFunc,
Function,
GlobalVar,
- Id,
If,
MatchCast,
SeqExpr,
@@ -1453,26 +1452,28 @@ class PyExprMutator:
"""
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 @@ class PyExprMutator:
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 4eb3b39619..2711e3287d 100644
--- a/python/tvm/relax/transform/lazy_transform_params.py
+++ b/python/tvm/relax/transform/lazy_transform_params.py
@@ -312,7 +312,7 @@ class LazyOutputMutator(PyExprMutator):
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 3af778efe4..a5e30c8139 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<ffi::String, ffi::Array<Binding>> 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<VarBinding>(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<MatchCast>(binding));
}
};
diff --git a/src/relax/backend/contrib/codegen_c/codegen_c.h
b/src/relax/backend/contrib/codegen_c/codegen_c.h
index b183fd85f9..58f76c92a8 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 06c4733e70..7c9f1a1ea7 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<NodeEntries> {
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<JSONGraphNode>(param->name_hint(),
"input" /* op_type_ */);
+ auto node_ptr = std::make_shared<JSONGraphNode>(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 2894b75685..a299289266 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<OutputType>,
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 380671473d..c27a39279f 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<Instruction::Arg(const
Expr&)> {
ffi::Array<ffi::String> 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 439a855e59..ee515489da 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<ffi::Optional<Expr>(const Expr&)> {
ffi::Array<ffi::String> 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 bc2f5d485c..3fb8ea0434 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 2f6b79ecb3..01522183bb 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<VarNode>()) {
- 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<ConstantNode>()) {
for (const auto& spec :
old_ty.as_or_throw<DTensorType>()->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<VarNode>()) {
- var_remap_[var->vid] = new_var;
+ var_remap_[ffi::GetRef<Var>(var)] = new_var;
} else {
tuple_getitem_remap_[old_expr.as_or_throw<TupleGetItem>()] = 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<VarNode>()) {
- var_remap_[var->vid] = scatter_var;
+ var_remap_[ffi::GetRef<Var>(var)] = scatter_var;
} else {
tuple_getitem_remap_[old_expr.as_or_throw<TupleGetItem>()] =
scatter_var;
}
@@ -174,7 +174,7 @@ class DistIRSharder : public ExprMutator {
ffi::Array<Var> new_params;
for (const Var& var : func->params) {
Var new_param = ShardInputParamTensorAndConstant(var).as_or_throw<Var>();
- 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<TupleGetItem>(val))) {
- var_remap_[binding->var->vid] =
tuple_getitem_remap_[ffi::GetRef<TupleGetItem>(val)];
+ var_remap_[binding->var] =
tuple_getitem_remap_[ffi::GetRef<TupleGetItem>(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 28db86b0a3..738d5dc0b2 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<VarNode>()) {
- 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<ConstantNode>()) {
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<VarNode>()) {
- var_remap_[binding->var->vid] = ffi::GetRef<Var>(var);
+ var_remap_[binding->var] = ffi::GetRef<Var>(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<TupleTypeNode>();
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<int>(inferred_tuple_ty->fields.size());
i++) {
if (!ffi::StructuralEqual()(
DTensorType(inferred_tuple_ty->fields[i].as_or_throw<DTensorType>()->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<TupleGetItem>(val))) {
- var_remap_[binding->var->vid] =
tuple_getitem_remap_[ffi::GetRef<TupleGetItem>(val)];
+ var_remap_[binding->var] =
tuple_getitem_remap_[ffi::GetRef<TupleGetItem>(val)];
} else {
ExprMutator::VisitBinding_(binding, val);
}
}
Expr VisitExpr_(const VarNode* var) final {
- auto it = input_tensor_remap_.find(ffi::GetRef<Var>(var));
+ Var var_ref = ffi::GetRef<Var>(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 af21c05673..44348e6cef 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 8048cda6b2..60cefd081d 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<Expr> 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<MatchCastNode>()) {
if (!cur_frame->is_dataflow) {
TVM_FFI_ICHECK(!match_cast->var.as<DataflowVarNode>())
@@ -342,7 +342,7 @@ class BlockBuilderImpl : public BlockBuilderNode {
std::vector<ScopeFrame> scope_stack_;
/*! \brief A binding table that maps var to value. */
- std::unordered_map<Id, Expr, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
binding_table_;
+ std::unordered_map<Var, Expr, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
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
ExprFunctor<Expr(const Expr&
Expr VisitVar_(const typename T::ContainerType* var) {
// Parameters and free-vars must be present with type
// Other vars must have already been normalized through binding
- TVM_FFI_ICHECK(!var->ty.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>(var);
}
diff --git a/src/relax/ir/dataflow_matcher.cc b/src/relax/ir/dataflow_matcher.cc
index 3f7edc6ae8..ef881e317b 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<VarNode>()) {
// "" 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 dc54fef78f..3da61a48fb 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<IdNode> n = ffi::make_object<IdNode>();
- 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<IfNode> n = ffi::make_object<IfNode>();
n->cond = std::move(cond);
@@ -146,9 +139,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-Var::Var(Id vid, ffi::Optional<Type> ty_annotation, Span span) {
+Var::Var(ffi::String name_hint, ffi::Optional<Type> ty_annotation, Span span) {
ffi::ObjectPtr<VarNode> n = ffi::make_object<VarNode>();
- 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<Type> 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<VarNode> node;
- if (auto dataflow_var = as<DataflowVarNode>()) {
- node = ffi::make_object<DataflowVarNode>(*dataflow_var);
- } else {
- node = ffi::make_object<VarNode>(*(operator->()));
- }
- ffi::ObjectPtr<ffi::Object>(std::move(node)).swap(data_);
- }
- return static_cast<VarNode*>(data_.get());
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def("relax.Var", [](ffi::String name_hint, ffi::Optional<Type>
ty_annotation,
- Span span) { return Var(name_hint, ty_annotation,
span); })
- .def("relax.VarFromId", [](Id vid, ffi::Optional<Type> ty_annotation,
Span span) {
- return Var(vid, ty_annotation, span);
- });
+ refl::GlobalDef().def("relax.Var", [](ffi::String name_hint,
ffi::Optional<Type> ty_annotation,
+ Span span) { return Var(name_hint,
ty_annotation, span); });
}
-DataflowVar::DataflowVar(Id vid, ffi::Optional<Type> ty_annotation, Span span)
{
+DataflowVar::DataflowVar(ffi::String name_hint, ffi::Optional<Type>
ty_annotation, Span span) {
ffi::ObjectPtr<DataflowVarNode> n = ffi::make_object<DataflowVarNode>();
- 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<Type>
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<Type> ty_annotation, Span
span) {
- return DataflowVar(name_hint, ty_annotation, span);
- })
- .def("relax.DataflowVarFromId", [](Id vid, ffi::Optional<Type>
ty_annotation, Span span) {
- return DataflowVar(vid, ty_annotation, span);
- });
+ refl::GlobalDef().def("relax.DataflowVar",
+ [](ffi::String name_hint, ffi::Optional<Type>
ty_annotation, Span span) {
+ return DataflowVar(name_hint, ty_annotation, span);
+ });
}
Constant::Constant(runtime::Tensor data, ffi::Optional<Type> ty_annotation,
Span span) {
diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc
index 7f9ff73759..963d302936 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<Var>(op);
+ auto it = var_remap_.find(var);
if (it != var_remap_.end()) {
return it->second;
}
// default case return self.
- return ffi::GetRef<Expr>(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<DataflowVarNode>()) {
- 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>(var);
} else {
- return Var(var->vid, ty, var->span);
+ return Var(var->name_hint, ty, var->span);
}
} else {
return ffi::GetRef<Var>(var);
@@ -830,6 +837,9 @@ Expr ExprMutator::VisitWithInnerScope(const Expr& expr) {
}
ffi::Optional<Expr> 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<DataflowVarNode>() ? DataflowVar(var->vid, ty,
var->span)
- : Var(var->vid, ty, var->span);
+ Var new_var = var.as<DataflowVarNode>() ? 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 1257948663..be64ae6459 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 8b1815812e..75f3873ce7 100644
--- a/src/relax/ir/transform.cc
+++ b/src/relax/ir/transform.cc
@@ -230,7 +230,7 @@ class DataflowBlockMutator : public ExprMutator {
}
}
if (!var.as<DataflowVarNode>()) {
- 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<DataflowVarNode>() &&
global_scope_vars.count(var->name_hint()) > 0) {
- TVM_FFI_ICHECK(var.same_as(global_scope_vars[var->name_hint()]))
+ if (!var.as<DataflowVarNode>() &&
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 1aa7eb9e20..80fec50fce 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<tvm::relax::Id, tvm::relax::Var, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>
+ std::unordered_map<tvm::relax::Var, tvm::relax::Var, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>
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<tvm::relax::Var> new_output_vars;
- std::unordered_map<tvm::relax::Id, tvm::relax::Var, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>
+ std::unordered_map<tvm::relax::Var, tvm::relax::Var, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>
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 35b513b07b..32f9f66d59 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<tvm::relax::VarNode>([](const ffi::ObjectRef& node,
ffi::String name) -> void {
using tvm::relax::VarNode;
- using tvm::relax::IdNode;
- const VarNode* var = node.as<VarNode>();
- IdNode* vid = const_cast<IdNode*>(var->vid.get());
- vid->name_hint = name;
+ VarNode* var = const_cast<VarNode*>(node.as<VarNode>());
+ var->name_hint = name;
});
TVM_STATIC_IR_FUNCTOR(Namer, vtable)
.set_dispatch<tvm::relax::DataflowVarNode>([](const ffi::ObjectRef& node,
ffi::String name) -> void {
using tvm::relax::DataflowVarNode;
- using tvm::relax::IdNode;
- const DataflowVarNode* var = node.as<DataflowVarNode>();
- IdNode* vid = const_cast<IdNode*>(var->vid.get());
- vid->name_hint = name;
+ DataflowVarNode* var =
const_cast<DataflowVarNode*>(node.as<DataflowVarNode>());
+ var->name_hint = name;
});
/////////////////////////////// Function ////////////////////////////////
diff --git a/src/relax/script/builder/utils.h b/src/relax/script/builder/utils.h
index 3a137b0ba4..f353465ee3 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<tvm::relax::DataflowVarNode>())
<< "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<tvm::relax::BindingBlock>
new_blocks(frame->binding_blocks.begin(),
@@ -107,8 +107,8 @@ inline tvm::relax::SeqExpr GetSeqExprForBranch(const
SeqExprFrame& frame, ffi::S
ffi::Array<tvm::relax::Binding>
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<tvm::relax::VarBindingNode>();
diff --git a/src/relax/script/printer/utils.h b/src/relax/script/printer/utils.h
index cced3c9705..b5d575666a 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<ExprDoc> TypeAsAnn(const relax::Var& v, const AccessPath&
v_p,
diff --git a/src/relax/training/utils.cc b/src/relax/training/utils.cc
index 33376e6d58..a279972e95 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 293537888d..04d5f33c66 100644
--- a/src/relax/transform/bind_params.cc
+++ b/src/relax/transform/bind_params.cc
@@ -92,7 +92,7 @@ std::tuple<ffi::Map<Var, Expr>, ffi::Map<tirx::Var,
PrimExpr>> NormalizeBindings
std::unordered_map<std::string, ffi::Array<relax::Var>> string_lookup;
std::unordered_set<const relax::VarNode*> 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<Var, Expr>, ffi::Map<tirx::Var,
PrimExpr>> 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 8ddc5539e1..7aa4c8acbe 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<Var>(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 9ad3a6dda4..cc12031d5f 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<Id, Var> replace_usage;
- ffi::Map<Id, Var> replace_binding;
- std::unordered_set<Id, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
bindings_to_remove;
- ffi::Map<Id, Constant> inline_constant;
+ ffi::Map<Var, Var> replace_usage;
+ ffi::Map<Var, Var> replace_binding;
+ std::unordered_set<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
bindings_to_remove;
+ ffi::Map<Var, Constant> 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<DataflowVarNode>() &&
!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>(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 fe38bbd3ae..28c30b708d 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 b2805bf9a2..0542d02910 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 8294e2c0f0..fa37f9ca07 100644
--- a/src/relax/transform/expand_tuple_arguments.cc
+++ b/src/relax/transform/expand_tuple_arguments.cc
@@ -50,7 +50,7 @@ ffi::Optional<Function> ExpandParams(Function func) {
ffi::Array<Expr> internal_tuple;
for (size_t i = 0; i < ty->fields.size(); i++) {
auto name = static_cast<const std::stringstream&>(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 dae9d0d36e..15d8493092 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<TupleType>();
ffi::Array<Expr> 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<VarBindingNode>();
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<ConstantNode>())) {
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<Expr> 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 7dff781df2..6aa6da3c1a 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<TensorTypeNode>()) {
// 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<PrimTypeNode>()) {
if (const auto* var = arg.as<VarNode>()) {
- tir_vars.push_back(tirx::Var(var->name_hint(),
tvm::PrimType(prim_value->dtype))
- .as_or_throw<PrimExpr>());
+ tir_vars.push_back(
+ tirx::Var(var->name_hint,
tvm::PrimType(prim_value->dtype)).as_or_throw<PrimExpr>());
} else if (auto literal = arg.as<PrimExpr>()) {
tir_vars.push_back(literal.value());
} else {
diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc
index 056f765ae9..f11d6af8c2 100644
--- a/src/relax/transform/gradient.cc
+++ b/src/relax/transform/gradient.cc
@@ -46,7 +46,7 @@ namespace relax {
// We will use NestedMsg<Expr> to handle adjoint updates involving tuple
handling
using AdjointMsg = NestedMsg<Expr>;
-using VarIdSet = std::unordered_set<Id, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>;
+using VarSet = std::unordered_set<Var, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>;
// Used in CallTIRWithGradCollector. call_tir -> call_tir_with_grad
using CallTIRWithGradInfo = std::unordered_map<Call, Call, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>;
@@ -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<Function>();
@@ -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<Id, Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
var_mapping;
+ std::unordered_map<Var, Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
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<VarNode>()) {
+ Var var_ref = ffi::GetRef<Var>(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<DataflowVarNode>() &&
var->IsInstance<DataflowVarNode>()) {
// 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<Var>&
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<Var> CheckAndMapRequireGrads(const ffi::Array<Var>&
require_grads,
const ffi::Map<Var, Var>&
var_map,
const ffi::String& func_name)
{
- VarIdSet var_set;
+ VarSet var_set;
ffi::Array<Var> 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 17d72c3ac2..13fda06faf 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<VarNode>()) {
- var_remap_[binding->var->vid] =
prev_call_node->args[0].as_or_throw<Var>();
+ var_remap_[binding->var] = prev_call_node->args[0].as_or_throw<Var>();
return;
} else {
return reemit_and_return();
diff --git a/src/relax/transform/inline_functions.cc
b/src/relax/transform/inline_functions.cc
index 4e03e83de8..a15d896457 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<Type>());
+ Var param_var(func->params[i]->name_hint, args[i]->ty.as<Type>());
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 2bdfb7178e..6a418ec815 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<Var> typed_captured_vars;
ffi::Map<Var, Expr> 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<CallNode>()) {
+ auto bound_value = LookupBinding(var);
+ if (IsClosure(var) && bound_value.as<CallNode>()) {
// if the original op was pure, we should use invoke_pure_closure
- Call orig_call =
builder_->LookupBinding(var).value().as_or_throw<Call>();
+ Call orig_call = bound_value.value().as_or_throw<Call>();
bool is_pure = [&]() -> bool {
if (auto op = orig_call->op.as<Op>()) {
static const auto& purity_map = Op::GetAttrMap<bool>("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 b5a5918c4a..02fd26c05f 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 a5d295f09a..3907219968 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 2897c0e909..5e0a990b7d 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<DataflowVarNode>()) {
- 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 3b53e3ac3b..585ff6bdeb 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<Type>()));
- var_remap_[input->vid] = new_var;
+ auto new_var = Var(input->name_hint,
VisitExprDepTypeField(input->ty.as_or_throw<Type>()));
+ var_remap_[ffi::GetRef<Var>(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<Var>(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>(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 b867ff967d..8445446e68 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<relax::VarNode>()) {
- name = args[i].as_or_throw<Var>()->name_hint();
+ name = args[i].as_or_throw<Var>()->name_hint;
} else {
name = std::string({static_cast<char>('A' + i)});
}
diff --git a/src/relax/transform/to_mixed_precision.cc
b/src/relax/transform/to_mixed_precision.cc
index ad2e285015..0764fd268f 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<TensorTypeNode>()) {
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 700521eb9c..ae279e4cd5 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<DataflowVarNode>()) {
- 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 7ba696639f..dcebbd4e6f 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 0163b93275..105dbff5e9 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<Id, Var, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>;
+ using VarMap = std::unordered_map<Var, Var, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>;
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<Var>(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>(var)] = copied_var;
relax_var_map_.Set(ffi::GetRef<Var>(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>(var)] = copied_var;
relax_var_map_.Set(ffi::GetRef<Var>(var), copied_var);
return copied_var;
}
diff --git a/src/relax/utils.cc b/src/relax/utils.cc
index a8b148aa89..63f0d08636 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 cf004a879a..6a09ee8083 100644
--- a/tests/python/relax/test_expr_functor.py
+++ b/tests/python/relax/test_expr_functor.py
@@ -303,7 +303,7 @@ class ASTPostPrinterMutator(PyExprMutator):
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 @@ class ASTPostPrinterMutator(PyExprMutator):
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 4650483546..17b4a466b8 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.
