Phase out Relax-specific Id aliases

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);
 };
 
 /*!

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;

python/tvm/relax/__init__.py

@@ -26,7 +26,6 @@
 from .expr import (
     Expr,
     Span,
-    Id,
     GlobalVar,
     Var,
     DataflowVar,

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")

python/tvm/relax/expr_functor.py

@@ -40,7 +40,6 @@
     ExternFunc,
     Function,
     GlobalVar,
-    Id,
     If,
     MatchCast,
     SeqExpr,
@@ -1453,34 +1452,36 @@ 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
         -------
         var : 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.

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)

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));
   }
 };

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 << ", ";

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);

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;
   }
 

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);

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);

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);
       }