This is an automated email from the ASF dual-hosted git repository.
tqchen 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 c584f233ce [IR][TVMScript] Share DataTypeImm values across script
dialects (#20462)
c584f233ce is described below
commit c584f233ce60536c8ebe7a6099a41051b80b58d4
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 26 10:03:33 2026 -0400
[IR][TVMScript] Share DataTypeImm values across script dialects (#20462)
## Summary
DataType values are ordinary expression arguments, so they now use a
shared `DataTypeImm` constant with a typed `DLDataType` value. The node
owns its root `AnyType` expression type and supports direct C++ value
extraction, reflection, structural operations, and serialization.
The shared `dtype` builder constructs `DataTypeImm` across the `I`, `T`,
and `R` script namespaces. Printing uses `T.dtype(...)` by default,
including Relax functions, and `I.dtype(...)` in shared IR module
contexts. `R.dtype` remains accepted as input; scalar constructors keep
their existing meaning. Non-dtype `GenericConst` payloads retain their
behavior.
Relax allocation, view, and VM paths consume typed dtype values. TIRx
traversal admits the shared literal as an operation argument. `AnyType`
remains in the root IR type layer, with Relax aliases preserved.
---
docs/reference/api/python/relax/relax.rst | 2 +-
docs/reference/api/python/script/ir_builder.rst | 2 +-
include/tvm/ir/expr.h | 22 ++++++++-
include/tvm/ir/expr_functor.h | 6 +++
include/tvm/ir/type.h | 23 ++++++++++
include/tvm/relax/expr_functor.h | 6 +++
include/tvm/relax/type.h | 32 +++----------
python/tvm/ir/__init__.py | 3 +-
python/tvm/ir/expr.py | 18 ++++++++
python/tvm/ir/type.py | 8 ++++
python/tvm/relax/expr_functor.py | 44 +++++++++++++++++-
python/tvm/relax/op/builtin/builtin.py | 7 ++-
python/tvm/relax/op/memory/memory.py | 11 ++---
python/tvm/relax/op/memory/view.py | 13 ++----
python/tvm/relax/op/vm/vm.py | 11 ++---
python/tvm/relax/script/ir_builder/ir.py | 17 +------
python/tvm/relax/testing/ast_printer.py | 3 ++
python/tvm/relax/type.py | 11 +----
python/tvm/script/ir_builder/__init__.py | 5 +-
python/tvm/script/ir_builder/ir.py | 12 ++++-
python/tvm/tirx/script/ir_builder/__init__.py | 2 +
src/ir/expr.cc | 53 ++++++++++++++++++++++
src/ir/expr_functor.cc | 9 ++++
src/ir/type.cc | 17 +++++++
src/relax/backend/vm/codegen_vm.cc | 3 ++
src/relax/backend/vm/codegen_vm_tir.cc | 3 ++
src/relax/backend/vm/lower_runtime_builtin.cc | 4 +-
src/relax/backend/vm/vm_shape_lower.cc | 8 ++--
src/relax/ir/block_builder.cc | 1 +
src/relax/ir/dependent_type.cc | 35 --------------
src/relax/ir/expr_functor.cc | 5 ++
src/relax/ir/py_expr_functor.cc | 33 ++++++++++----
src/relax/op/memory/view.cc | 12 ++---
src/relax/op/op.cc | 50 ++++++++++----------
src/relax/op/op_common.h | 6 +--
src/relax/script/ir_builder/frame.cc | 2 +-
src/relax/script/printer/dependent_type.cc | 8 ++--
src/relax/script/printer/expr.cc | 7 +--
src/relax/script/printer/function.cc | 2 +-
src/relax/transform/allocate_workspace.cc | 2 +-
src/relax/transform/call_tir_rewrite.cc | 2 +-
src/relax/transform/fuse_ops.cc | 3 +-
src/relax/transform/lower_alloc_tensor.cc | 10 ++--
src/relax/transform/rewrite_cuda_graph.cc | 4 +-
src/relax/transform/static_plan_block_memory.cc | 18 ++++----
src/script/printer/ir/misc.cc | 33 ++++++++++++++
src/tirx/ir/function.cc | 4 +-
src/tirx/ir/tir_visitor_with_path.cc | 2 +
src/tirx/ir/tir_visitor_with_path.h | 6 +++
.../relax/script/test_relax_script_printer.py | 4 +-
tests/python/relax/test_ast_printer.py | 4 +-
tests/python/relax/test_expr.py | 4 +-
tests/python/script/test_basic_usage.py | 13 +++++-
53 files changed, 419 insertions(+), 206 deletions(-)
diff --git a/docs/reference/api/python/relax/relax.rst
b/docs/reference/api/python/relax/relax.rst
index 7ae3359987..a20975cf7a 100644
--- a/docs/reference/api/python/relax/relax.rst
+++ b/docs/reference/api/python/relax/relax.rst
@@ -21,4 +21,4 @@ tvm.relax
.. automodule:: tvm.relax
:members:
:imported-members:
- :exclude-members: BlockBuilder, Call, Tuple, TupleGetItem, Var, Span,
GlobalVar, SourceName, TupleType, Type, FuncType
+ :exclude-members: AnyType, ObjectType, BlockBuilder, Call, Tuple,
TupleGetItem, Var, Span, GlobalVar, SourceName, TupleType, Type, FuncType
diff --git a/docs/reference/api/python/script/ir_builder.rst
b/docs/reference/api/python/script/ir_builder.rst
index 53f6041e09..b17ba6de31 100644
--- a/docs/reference/api/python/script/ir_builder.rst
+++ b/docs/reference/api/python/script/ir_builder.rst
@@ -23,7 +23,7 @@ tvm.script.ir_builder
.. automodule:: tvm.script.ir_builder
:members:
:imported-members:
- :exclude-members: GenericConst, Range, StringImm, StringType
+ :exclude-members: DataTypeImm, GenericConst, Range, StringImm, StringType
tvm.relax.script.ir_builder.distributed
***************************************
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index dea7902280..7fa6dbf532 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -517,7 +517,7 @@ class Call : public Expr {
/*! \brief Base node for literal constants. */
class ConstantNode : public ExprNode {
public:
- static constexpr uint32_t _type_child_slots = 4;
+ static constexpr uint32_t _type_child_slots = 5;
static void RegisterReflection() {
ffi::reflection::ObjectDef<ConstantNode>(); }
TVM_FFI_DECLARE_OBJECT_INFO("ir.Constant", ConstantNode, ExprNode);
};
@@ -568,6 +568,26 @@ class StringImm : public Constant {
TVM_DEFINE_OBJECT_REF_COW_METHOD(StringImmNode);
};
+/*! \brief A data type literal with AnyType as its expression type. */
+class DataTypeImmNode : public ConstantNode {
+ public:
+ DLDataType value;
+
+ static void RegisterReflection() {
+ ffi::reflection::ObjectDef<DataTypeImmNode>().def_ro("value",
&DataTypeImmNode::value);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DataTypeImm", DataTypeImmNode,
ConstantNode);
+};
+
+/*! \brief Managed reference to a data type literal. */
+class DataTypeImm : public Constant {
+ public:
+ TVM_DLL explicit DataTypeImm(DLDataType value, Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataTypeImm, Constant,
DataTypeImmNode);
+ TVM_DEFINE_OBJECT_REF_COW_METHOD(DataTypeImmNode);
+};
+
/*!
* \brief Constant integer literals in the program.
* \sa IntImm
diff --git a/include/tvm/ir/expr_functor.h b/include/tvm/ir/expr_functor.h
index a3b7671643..cc735b4db2 100644
--- a/include/tvm/ir/expr_functor.h
+++ b/include/tvm/ir/expr_functor.h
@@ -100,6 +100,9 @@ class ExprFunctor<R(const Expr&, Args...)> {
virtual R Dispatch_(const GenericConstNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
+ virtual R Dispatch_(const DataTypeImmNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
virtual R Dispatch_(const IntImmNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
@@ -236,6 +239,7 @@ class ExprFunctor<R(const Expr&, Args...)> {
SetDispatch<TSelf, GlobalVarNode>(vtable);
SetDispatch<TSelf, CallNode>(vtable);
SetDispatch<TSelf, GenericConstNode>(vtable);
+ SetDispatch<TSelf, DataTypeImmNode>(vtable);
SetDispatch<TSelf, IntImmNode>(vtable);
SetDispatch<TSelf, FloatImmNode>(vtable);
SetDispatch<TSelf, OpNode>(vtable);
@@ -343,6 +347,7 @@ class TVM_DLL ExprVisitor : public ObjectVisitor {
virtual ffi::Optional<VisitInterrupt> Visit_(const GlobalVarNode* node);
virtual ffi::Optional<VisitInterrupt> Visit_(const CallNode* node);
virtual ffi::Optional<VisitInterrupt> Visit_(const GenericConstNode* node);
+ virtual ffi::Optional<VisitInterrupt> Visit_(const DataTypeImmNode* node);
virtual ffi::Optional<VisitInterrupt> Visit_(const IntImmNode* node);
virtual ffi::Optional<VisitInterrupt> Visit_(const FloatImmNode* node);
virtual ffi::Optional<VisitInterrupt> Visit_(const OpNode* node);
@@ -468,6 +473,7 @@ class TVM_DLL ExprMutator : public ObjectMutator {
virtual UnchangedOr<Expr> Mutate_(const GlobalVarNode* node, InplaceMode
inplace_mode);
virtual UnchangedOr<Expr> Mutate_(const CallNode* node, InplaceMode
inplace_mode);
virtual UnchangedOr<Expr> Mutate_(const GenericConstNode* node, InplaceMode
inplace_mode);
+ virtual UnchangedOr<Expr> Mutate_(const DataTypeImmNode* node, InplaceMode
inplace_mode);
virtual UnchangedOr<PrimExpr> Mutate_(const IntImmNode* node, InplaceMode
inplace_mode);
virtual UnchangedOr<PrimExpr> Mutate_(const FloatImmNode* node, InplaceMode
inplace_mode);
virtual UnchangedOr<Expr> Mutate_(const OpNode* node, InplaceMode
inplace_mode);
diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h
index 0c148f1095..21132604c4 100644
--- a/include/tvm/ir/type.h
+++ b/include/tvm/ir/type.h
@@ -34,6 +34,29 @@
namespace tvm {
+/*!
+ * \brief The top type, which admits any value.
+ *
+ * AnyType types values that have no more specific static type, such as
+ * DataType-valued generic constants and opaque runtime objects.
+ */
+class AnyTypeNode : public TypeNode {
+ public:
+ static void RegisterReflection() {
ffi::reflection::ObjectDef<AnyTypeNode>(); }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.AnyType", AnyTypeNode, TypeNode);
+};
+
+/*!
+ * \brief Managed reference to AnyTypeNode.
+ * \sa AnyTypeNode
+ */
+class AnyType : public Type {
+ public:
+ TVM_DLL AnyType(Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AnyType, Type, AnyTypeNode);
+};
+
/*! \brief Semantic string type; its physical representation is context
dependent. */
class StringTypeNode : public TypeNode {
public:
diff --git a/include/tvm/relax/expr_functor.h b/include/tvm/relax/expr_functor.h
index 6ae7d50610..cc722a6d8d 100644
--- a/include/tvm/relax/expr_functor.h
+++ b/include/tvm/relax/expr_functor.h
@@ -163,6 +163,7 @@ class ExprFunctor<R(const Expr& n, Args...)> {
virtual R VisitExpr_(const OpNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const TupleGetItemNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const StringImmNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const DataTypeImmNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const prim::LetNode* op, Args...) EXPR_FUNCTOR_DISABLED;
virtual R VisitExpr_(const prim::AddNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const prim::LShiftNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
@@ -251,6 +252,7 @@ class ExprFunctor<R(const Expr& n, Args...)> {
RELAX_EXPR_FUNCTOR_DISPATCH(OpNode);
RELAX_EXPR_FUNCTOR_DISPATCH(TupleGetItemNode);
RELAX_EXPR_FUNCTOR_DISPATCH(StringImmNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(DataTypeImmNode);
vtable.Finalize();
return vtable;
}
@@ -285,6 +287,7 @@ class ExprVisitor : public ExprFunctor<void(const Expr&)> {
void VisitExpr_(const OpNode* op) override;
void VisitExpr_(const TupleGetItemNode* op) override;
void VisitExpr_(const StringImmNode* op) override;
+ void VisitExpr_(const DataTypeImmNode* op) override;
void VisitExpr_(const prim::AddNode* op) override;
void VisitExpr_(const prim::LShiftNode* op) override;
void VisitExpr_(const prim::RShiftNode* op) override;
@@ -341,6 +344,7 @@ class ExprVisitor : public ExprFunctor<void(const Expr&)> {
virtual void VisitBinding_(const VarBindingNode* binding, const
TupleGetItemNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const ExprNode*
val);
virtual void VisitBinding_(const VarBindingNode* binding, const
StringImmNode* val);
+ virtual void VisitBinding_(const VarBindingNode* binding, const
DataTypeImmNode* val);
/*!
* \brief Generic dispatcher for binding blocks.
* \param block The binding block to be visited.
@@ -443,6 +447,7 @@ class ExprMutatorBase : public ExprFunctor<Expr(const
Expr&)> {
Expr VisitExpr_(const OpNode* op) override;
Expr VisitExpr_(const TupleGetItemNode* op) override;
Expr VisitExpr_(const StringImmNode* op) override;
+ Expr VisitExpr_(const DataTypeImmNode* op) override;
Expr VisitExpr_(const prim::AddNode* op) override;
Expr VisitExpr_(const prim::LShiftNode* op) override;
Expr VisitExpr_(const prim::RShiftNode* op) override;
@@ -594,6 +599,7 @@ class ExprMutator : public ExprMutatorBase {
virtual void VisitBinding_(const VarBindingNode* binding, const
TupleGetItemNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const ExprNode*
val);
virtual void VisitBinding_(const VarBindingNode* binding, const
StringImmNode* val);
+ virtual void VisitBinding_(const VarBindingNode* binding, const
DataTypeImmNode* val);
/*!
* \brief Generic dispatcher for binding blocks.
* \param block The binding block to be visited.
diff --git a/include/tvm/relax/type.h b/include/tvm/relax/type.h
index cb3c0e8196..5b3b652024 100644
--- a/include/tvm/relax/type.h
+++ b/include/tvm/relax/type.h
@@ -46,9 +46,15 @@ class BlockBuilder;
/*! \brief Indicates the number of dimensions of a tensor is unknown at
compile time. */
static constexpr int kUnknownNDim = -1;
+using tvm::AnyType;
+using tvm::AnyTypeNode;
using tvm::TupleType;
using tvm::TupleTypeNode;
+// Compatibility aliases for existing C++ callers. New code should use
AnyType.
+using ObjectTypeNode = AnyTypeNode;
+using ObjectType = AnyType;
+
class PackedFuncTypeNode : public TypeNode {
public:
static void RegisterReflection() {
@@ -89,32 +95,6 @@ class PackedFuncType : public Type {
* normalized through NormalizeArg. This invariant is checked in constructors
* and simplifies assumptions during type deduction.
*/
-/*!
- * \brief Any Relax value.
- */
-class AnyTypeNode : public TypeNode {
- public:
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AnyTypeNode>();
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.AnyType", AnyTypeNode, TypeNode);
-};
-
-/*!
- * \brief Managed reference to AnyTypeNode.
- * \sa AnyTypeNode
- */
-class AnyType : public Type {
- public:
- TVM_DLL AnyType(Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AnyType, Type, AnyTypeNode);
-};
-
-// Compatibility aliases for existing C++ callers. New code should use
AnyType.
-using ObjectTypeNode = AnyTypeNode;
-using ObjectType = AnyType;
/*!
* \brief Type of shape value.
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index 17e2c7f147..de84cecf66 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -33,11 +33,12 @@ from .base import (
# Register Type before Expr. Expr's reflected ``ty`` field otherwise creates
# an auto-generated Type wrapper before the concrete Python class is available.
-from .type import FuncType, OpaqueType, PointerType, PrimType, StringType,
TupleType, Type
+from .type import AnyType, FuncType, OpaqueType, PointerType, PrimType,
StringType, TupleType, Type
from .expr import (
Call,
Constant,
GenericConst,
+ DataTypeImm,
StringImm,
Expr,
ExprOperand,
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 7cd7680ea8..42b8f258da 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -463,6 +463,24 @@ class GenericConst(_ExprCallable, Constant):
return True
+@tvm_ffi.register_object("ir.DataTypeImm")
+class DataTypeImm(Constant):
+ """A data type literal whose expression type is AnyType.
+
+ Parameters
+ ----------
+ value : str or tvm.DataType
+ The represented data type.
+ span : Span, optional
+ The source span of the literal.
+ """
+
+ value: tvm.DataType
+
+ def __init__(self, value: str | tvm.DataType, span: Span | None = None) ->
None:
+ self.__init_handle_by_constructor__(_ffi_api.DataTypeImm, value, span)
+
+
@tvm_ffi.register_object("ir.StringImm")
class StringImm(Constant):
"""A string literal with StringType."""
diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py
index 43060f95a6..e0fbf6cd09 100644
--- a/python/tvm/ir/type.py
+++ b/python/tvm/ir/type.py
@@ -54,6 +54,14 @@ class Type(Node, Scriptable):
return self.is_(other)
+@tvm_ffi.register_object("ir.AnyType")
+class AnyType(Type):
+ """The top type, which admits any value."""
+
+ def __init__(self, span=None) -> None:
+ self.__init_handle_by_constructor__(_ffi_api.AnyType, span)
+
+
@tvm_ffi.register_object("ir.OpaqueType")
class OpaqueType(Type):
"""Type marker for opaque values that must be removed from finished IR."""
diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py
index ec7a773958..a2ba021939 100644
--- a/python/tvm/relax/expr_functor.py
+++ b/python/tvm/relax/expr_functor.py
@@ -23,7 +23,7 @@ from collections.abc import Callable
import tvm_ffi
from tvm import tirx as _tirx
-from tvm.ir import Call, GenericConst, Op, StringImm, TensorLoad, is_prim_expr
+from tvm.ir import Call, DataTypeImm, GenericConst, Op, StringImm, TensorLoad,
is_prim_expr
from tvm.ir.utils import derived_object
from tvm.runtime import Object
@@ -221,6 +221,8 @@ class ExprFunctor:
ret = self.visit_op_(expr)
elif isinstance(expr, TupleGetItem):
ret = self.visit_tuple_getitem_(expr)
+ elif isinstance(expr, DataTypeImm):
+ ret = self.visit_data_type_imm_(expr)
elif isinstance(expr, StringImm):
ret = self.visit_string_imm_(expr)
elif isinstance(expr, _tirx.Let | _tirx.Reduce):
@@ -375,6 +377,9 @@ class ExprFunctor:
def visit_string_imm_(self, op: StringImm):
raise NotImplementedError()
+ def visit_data_type_imm_(self, op: DataTypeImm):
+ raise NotImplementedError()
+
def visit_var_binding_(self, binding: VarBinding):
raise NotImplementedError()
@@ -447,6 +452,7 @@ class _PyExprVisitor(tvm_ffi.core.Object):
f_visit_tuple_getitem_: Callable | None = None,
f_visit_expr_fallback_: Callable | None = None,
f_visit_string_imm_: Callable | None = None,
+ f_visit_data_type_imm_: Callable | None = None,
f_visit_binding: Callable | None = None,
f_visit_var_binding_: Callable | None = None,
f_visit_match_cast_: Callable | None = None,
@@ -478,6 +484,7 @@ class _PyExprVisitor(tvm_ffi.core.Object):
f_visit_tuple_getitem_,
f_visit_expr_fallback_,
f_visit_string_imm_,
+ f_visit_data_type_imm_,
f_visit_binding,
f_visit_var_binding_,
f_visit_match_cast_,
@@ -572,6 +579,7 @@ class PyExprVisitor:
"visit_tuple_getitem_",
"visit_expr_fallback_",
"visit_string_imm_",
+ "visit_data_type_imm_",
"visit_binding",
"visit_var_binding_",
"visit_match_cast_",
@@ -829,6 +837,19 @@ class PyExprVisitor:
# Using self._outer() to ref _PyExprVisitor
return _ffi_api.ExprVisitorVisitExpr(self._outer(), op) # type: ignore
+ def visit_data_type_imm_(self, op: DataTypeImm) -> None:
+ """Visit DataTypeImm.
+ Users can customized this function to overwrite VisitExpr_(const
DataTypeImmNode* op)
+ on the C++ side.
+
+ Parameters
+ ----------
+ op : DataTypeImm
+ The DataTypeImm to be visited.
+ """
+ # Using self._outer() to ref _PyExprVisitor
+ return _ffi_api.ExprVisitorVisitExpr(self._outer(), op) # type: ignore
+
def visit_var_binding_(self, binding: VarBinding) -> None:
"""Visit VarBinding.
Users can customized this function to overwrite VisitBinding_(const
VarBindingNode* binding)
@@ -950,6 +971,7 @@ class _PyExprMutator(Object):
f_visit_tuple_getitem_: Callable | None = None,
f_visit_expr_fallback_: Callable | None = None,
f_visit_string_imm_: Callable | None = None,
+ f_visit_data_type_imm_: Callable | None = None,
f_visit_binding: Callable | None = None,
f_visit_var_binding_: Callable | None = None,
f_visit_match_cast_: Callable | None = None,
@@ -982,6 +1004,7 @@ class _PyExprMutator(Object):
f_visit_tuple_getitem_,
f_visit_expr_fallback_,
f_visit_string_imm_,
+ f_visit_data_type_imm_,
f_visit_binding,
f_visit_var_binding_,
f_visit_match_cast_,
@@ -1092,6 +1115,7 @@ class PyExprMutator:
"visit_tuple_getitem_",
"visit_expr_fallback_",
"visit_string_imm_",
+ "visit_data_type_imm_",
"visit_binding",
"visit_var_binding_",
"visit_match_cast_",
@@ -1443,6 +1467,24 @@ class PyExprMutator:
# Using self._outer() to ref _PyExprMutator
return _ffi_api.ExprMutatorVisitExpr(self._outer(), op) # type: ignore
+ def visit_data_type_imm_(self, op: DataTypeImm) -> Expr:
+ """Visit DataTypeImm.
+ Users can customized this function to overwrite VisitExpr_(const
DataTypeImmNode* op)
+ on the C++ side.
+
+ Parameters
+ ----------
+ op : DataTypeImm
+ The DataTypeImm to be visited.
+
+ Returns
+ -------
+ result : Expr
+ The Expr after transformation
+ """
+ # Using self._outer() to ref _PyExprMutator
+ return _ffi_api.ExprMutatorVisitExpr(self._outer(), op) # type: ignore
+
def visit_var_binding_(self, binding: VarBinding) -> None:
"""Visit VarBinding.
Users can customized this function to overwrite VisitBinding_(const
VarBindingNode* binding)
diff --git a/python/tvm/relax/op/builtin/builtin.py
b/python/tvm/relax/op/builtin/builtin.py
index bc5d0db2d7..5592e0d402 100644
--- a/python/tvm/relax/op/builtin/builtin.py
+++ b/python/tvm/relax/op/builtin/builtin.py
@@ -15,8 +15,7 @@
# specific language governing permissions and limitations
"""The builtin Relax operators."""
-import tvm
-from tvm.ir import Call, GenericConst, StringImm
+from tvm.ir import Call, DataTypeImm, GenericConst, StringImm
from ...expr import Expr, prim_value
from ...utils import convert_to_expr
@@ -53,8 +52,8 @@ def alloc_tensor(
"""
if not isinstance(shape, Expr):
shape = convert_to_expr(shape)
- if isinstance(dtype, str):
- dtype = GenericConst(tvm.DataType(dtype), tvm.relax.AnyType())
+ if isinstance(dtype, str | GenericConst):
+ dtype = DataTypeImm(dtype.value if isinstance(dtype, GenericConst)
else dtype)
if isinstance(runtime_device_index, int):
runtime_device_index = prim_value(runtime_device_index)
if isinstance(storage_scope, str):
diff --git a/python/tvm/relax/op/memory/memory.py
b/python/tvm/relax/op/memory/memory.py
index 7b0e22f78f..2cdcbdfe7b 100644
--- a/python/tvm/relax/op/memory/memory.py
+++ b/python/tvm/relax/op/memory/memory.py
@@ -15,8 +15,7 @@
# specific language governing permissions and limitations
"""Relax memory primitives."""
-import tvm
-from tvm.ir import Call, GenericConst, StringImm
+from tvm.ir import Call, DataTypeImm, GenericConst, StringImm
from ...expr import Expr, prim_value
from ...utils import convert_to_expr
@@ -53,8 +52,8 @@ def alloc_storage(
A relax Call, which gets the allocated storage.
"""
size = convert_to_expr(size)
- if isinstance(dtype, str):
- dtype = GenericConst(tvm.DataType(dtype), tvm.relax.AnyType())
+ if isinstance(dtype, str | GenericConst):
+ dtype = DataTypeImm(dtype.value if isinstance(dtype, GenericConst)
else dtype)
if isinstance(storage_scope, str):
storage_scope = StringImm(storage_scope)
if isinstance(virtual_device_index, int):
@@ -97,8 +96,8 @@ def alloc_tensor(
if isinstance(offset, int):
offset = prim_value(offset)
shape = convert_to_expr(shape)
- if isinstance(dtype, str):
- dtype = GenericConst(tvm.DataType(dtype), tvm.relax.AnyType())
+ if isinstance(dtype, str | GenericConst):
+ dtype = DataTypeImm(dtype.value if isinstance(dtype, GenericConst)
else dtype)
if isinstance(runtime_device_ind, int):
runtime_device_ind = prim_value(runtime_device_ind)
return _ffi_api.alloc_tensor(storage, offset, shape, dtype,
runtime_device_ind) # type: ignore
diff --git a/python/tvm/relax/op/memory/view.py
b/python/tvm/relax/op/memory/view.py
index 5874c1e5e9..b3ecc57842 100644
--- a/python/tvm/relax/op/memory/view.py
+++ b/python/tvm/relax/op/memory/view.py
@@ -28,9 +28,8 @@ while keeping the same underlying data.
from collections.abc import Sequence
-from tvm import DataType
-from tvm.ir import GenericConst
-from tvm.relax import AnyType, Expr, ShapeExpr
+from tvm.ir import DataTypeImm, GenericConst
+from tvm.relax import Expr, ShapeExpr
from tvm.relax.expr import prim_value
from ..base import null_value
@@ -91,11 +90,9 @@ def view(
return relax_cls(expr)
shape = _normalize(shape, ShapeExpr)
- dtype = (
- null_value()
- if dtype is None
- else _normalize(dtype, lambda value: GenericConst(DataType(value),
AnyType()))
- )
+ if isinstance(dtype, GenericConst):
+ dtype = DataTypeImm(dtype.value)
+ dtype = null_value() if dtype is None else _normalize(dtype, DataTypeImm)
relative_byte_offset = (
relative_byte_offset
if relative_byte_offset is None or isinstance(relative_byte_offset,
Expr)
diff --git a/python/tvm/relax/op/vm/vm.py b/python/tvm/relax/op/vm/vm.py
index d5c3bf88a6..51c9b4bdda 100644
--- a/python/tvm/relax/op/vm/vm.py
+++ b/python/tvm/relax/op/vm/vm.py
@@ -15,8 +15,7 @@
# specific language governing permissions and limitations
"""Relax vm primitives."""
-import tvm
-from tvm.ir import Call, GenericConst, StringImm
+from tvm.ir import Call, DataTypeImm, GenericConst, StringImm
from ...expr import Expr, Tuple, prim_value
from ...utils import convert_to_expr
@@ -53,8 +52,8 @@ def alloc_storage(
A relax Call, which gets the allocated storage.
"""
shape = convert_to_expr(shape)
- if isinstance(dtype, str):
- dtype = GenericConst(tvm.DataType(dtype), tvm.relax.AnyType())
+ if isinstance(dtype, str | GenericConst):
+ dtype = DataTypeImm(dtype.value if isinstance(dtype, GenericConst)
else dtype)
if isinstance(storage_scope, str):
storage_scope = StringImm(storage_scope)
if isinstance(runtime_device_index, int):
@@ -97,8 +96,8 @@ def alloc_tensor(
if isinstance(offset, int):
offset = prim_value(offset)
shape = convert_to_expr(shape)
- if isinstance(dtype, str):
- dtype = GenericConst(tvm.DataType(dtype), tvm.relax.AnyType())
+ if isinstance(dtype, str | GenericConst):
+ dtype = DataTypeImm(dtype.value if isinstance(dtype, GenericConst)
else dtype)
if isinstance(runtime_device_ind, int):
runtime_device_ind = prim_value(runtime_device_ind)
return _ffi_api.alloc_tensor(storage, offset, shape, dtype,
runtime_device_ind) # type: ignore
diff --git a/python/tvm/relax/script/ir_builder/ir.py
b/python/tvm/relax/script/ir_builder/ir.py
index 481fc74afc..b560176f34 100644
--- a/python/tvm/relax/script/ir_builder/ir.py
+++ b/python/tvm/relax/script/ir_builder/ir.py
@@ -25,8 +25,8 @@ import numbers as _numbers
import re as _re
import tvm
-from tvm import DataType, relax
from tvm import ir as _ir
+from tvm import relax
from tvm import relax as _relax
from tvm.ir import IRModule
from tvm.relax import Expr, ExternFunc, ShapeExpr, TupleGetItem, const
@@ -52,6 +52,7 @@ from tvm.runtime._tensor import (
from tvm.script.ir_builder.base import SpanEntry as _SpanEntry
from tvm.script.ir_builder.base import resolve_global_info_args as
_resolve_global_info_args
from tvm.script.ir_builder.ir import _global_infos
+from tvm.script.ir_builder.ir import dtype as dtype
from tvm.script.ir_builder.parser_protocol import resolve_global_info_ as
_resolve_global_info
from . import _ffi_api
@@ -291,20 +292,6 @@ def str(value: py_str) -> Expr:
return tvm.ir.StringImm(value) # type: ignore[attr-defined] # pylint:
disable=no-member
-def dtype(value: py_str | DataType) -> Expr:
- """Create a dtype imm expression.
- Parameters
- ----------
- value : dtype
- The value of the dtype.
- Returns
- -------
- res : Expr
- The result dtype.
- """
- return tvm.ir.GenericConst(tvm.DataType(value), tvm.relax.AnyType()) #
type: ignore[attr-defined] # pylint: disable=no-member
-
-
@_resolve_global_info_args("vdevice", resolver=resolve_global_info_)
def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
"""Construct a Relax tensor type.
diff --git a/python/tvm/relax/testing/ast_printer.py
b/python/tvm/relax/testing/ast_printer.py
index 5d7fdcb55d..9a467eb577 100644
--- a/python/tvm/relax/testing/ast_printer.py
+++ b/python/tvm/relax/testing/ast_printer.py
@@ -210,6 +210,9 @@ class ASTPrinter(ExprFunctor):
def visit_string_imm_(self, op: tvm.ir.StringImm) -> str:
return self.build_expr(op, "StringImm", value=wrap_quotes(op.value))
+ def visit_data_type_imm_(self, op: tvm.ir.DataTypeImm) -> str:
+ return self.build_expr(op, "DataTypeImm", value=str(op.value))
+
def visit_op_(self, op: tvm.ir.Op) -> str:
# TODO: List other attributes?
# op is not actually a Relax expr and does not have
diff --git a/python/tvm/relax/type.py b/python/tvm/relax/type.py
index b5018ec7ea..4ed8e7a2ff 100644
--- a/python/tvm/relax/type.py
+++ b/python/tvm/relax/type.py
@@ -21,22 +21,13 @@
import tvm_ffi
from tvm_ffi import Array
-from tvm.ir import EnvFunc, PrimType, Span, TupleType
+from tvm.ir import AnyType, EnvFunc, PrimType, Span, TupleType
from tvm.relax.global_info import VDevice
from . import _ffi_api
from .expr import Expr, ShapeExpr, Type
from .ty import PackedFuncType
-
-@tvm_ffi.register_object("relax.AnyType")
-class AnyType(Type):
- """Type of any Relax value."""
-
- def __init__(self, span: Span = None) -> None:
- self.__init_handle_by_constructor__(_ffi_api.AnyType, span) # type:
ignore
-
-
ObjectType = AnyType
diff --git a/python/tvm/script/ir_builder/__init__.py
b/python/tvm/script/ir_builder/__init__.py
index ecaecca8e6..f1faadf35e 100644
--- a/python/tvm/script/ir_builder/__init__.py
+++ b/python/tvm/script/ir_builder/__init__.py
@@ -16,7 +16,7 @@
# under the License.
"""Shared TVMScript construction APIs and lazy language variant builders."""
-from tvm.ir import GenericConst, Range, StringImm, StringType
+from tvm.ir import DataTypeImm, GenericConst, Range, StringImm, StringType
from .base import (
MISSING,
@@ -30,6 +30,7 @@ from .frame import IRModuleFrame
from .ir import _get_dialect_builder as __getattr__
from .ir import (
constexpr,
+ dtype,
dynamic,
meta_var,
module_attrs,
@@ -49,6 +50,7 @@ from .parser_protocol import (
__all__ = [
"MISSING",
"AlreadyEmitted",
+ "DataTypeImm",
"GenericConst",
"IRBuilder",
"IRModuleFrame",
@@ -59,6 +61,7 @@ __all__ = [
"check_well_formed_",
"constexpr",
"decl_function",
+ "dtype",
"dynamic",
"ir_module",
"meta_var",
diff --git a/python/tvm/script/ir_builder/ir.py
b/python/tvm/script/ir_builder/ir.py
index 8afe2567be..6b31d7b174 100644
--- a/python/tvm/script/ir_builder/ir.py
+++ b/python/tvm/script/ir_builder/ir.py
@@ -18,7 +18,8 @@
from typing import NoReturn, TypeVar
-from tvm.ir import GlobalInfo, Span, Var
+from tvm import DataType
+from tvm.ir import DataTypeImm, GlobalInfo, Span, Var
from tvm.runtime import Object as tvm_Object
from . import _ffi_api
@@ -28,6 +29,15 @@ from .frame import IRModuleFrame
T = TypeVar("T")
+def dtype(value: str | DataType) -> DataTypeImm:
+ """Construct a DataType-valued constant for an ordinary operation argument.
+
+ Unlike scalar constructors such as ``T.float32``, this constructs a dtype
+ value, rather than a scalar value or variable of that dtype.
+ """
+ return DataTypeImm(value)
+
+
def dynamic(name: str, dtype: str = "int64", *, span: SpanEntry | Span | None
= None) -> Var:
"""Create a fresh primitive symbolic variable, independently of builder
scope.
diff --git a/python/tvm/tirx/script/ir_builder/__init__.py
b/python/tvm/tirx/script/ir_builder/__init__.py
index 34b2815f20..c4eb0558ea 100644
--- a/python/tvm/tirx/script/ir_builder/__init__.py
+++ b/python/tvm/tirx/script/ir_builder/__init__.py
@@ -17,6 +17,7 @@
"""Concrete TIRx construction operations over the shared native IRBuilder
stack."""
from tvm.script.ir_builder.ir import constexpr as constexpr
+from tvm.script.ir_builder.ir import dtype as dtype
from tvm.script.ir_builder.ir import dynamic as dynamic
from tvm.tirx.lang.alloc_pool import SMEMPool as SMEMPool
from tvm.tirx.lang.alloc_pool import TMEMPool as TMEMPool
@@ -66,6 +67,7 @@ __all__ = [
*_op_exports,
*_protocol_exports,
"dynamic",
+ "dtype",
"constexpr",
"SMEMPool",
"TMEMPool",
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index b5218cc51c..fca78bafc0 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -49,6 +49,43 @@ TVMFFIAny ConstantVisit(ffi::StructuralVisitorObj* visitor,
ffi::AnyView value)
return ffi::AnyView(nullptr).CopyToTVMFFIAny();
}
+TVMFFIAny DataTypeImmVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ const DataTypeImmNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
DataTypeImmNode>(value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny DataTypeImmMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ const DataTypeImmNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
DataTypeImmNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+ mutator->MutateExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<DLDataType>, mapped_value,
+ mutator->MutateExpected(self->value));
+ if (mapped_ty.UnchangedOrSameAs(self->ty) &&
mapped_value.UnchangedOrSameAs(self->value)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<DataTypeImmNode> copy =
ffi::make_object<DataTypeImmNode>(*self);
+ copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty));
+ copy->value = std::move(mapped_value).ValueOrUnchanged(copy->value);
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny DataTypeImmMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ DataTypeImmNode* self = const_cast<DataTypeImmNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
DataTypeImmNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+ mutator->MutateExpected(self->ty,
ffi::InplaceMode::kAllow));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<DLDataType>, mapped_value,
+ mutator->MutateExpected(self->value,
ffi::InplaceMode::kAllow));
+ if (!mapped_ty.IsUnchanged()) self->ty =
std::move(mapped_ty).ValueUnchecked();
+ if (!mapped_value.IsUnchanged()) self->value =
std::move(mapped_value).ValueUnchecked();
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
template <typename TNode>
TVMFFIAny ConstantMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
const TNode* self =
@@ -764,6 +801,14 @@ StringImm::StringImm(ffi::String value, Span span) {
data_ = std::move(node);
}
+DataTypeImm::DataTypeImm(DLDataType value, Span span) {
+ auto node = ffi::make_object<DataTypeImmNode>();
+ node->value = std::move(value);
+ node->ty = AnyType();
+ node->span = std::move(span);
+ data_ = std::move(node);
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
ConstantNode::RegisterReflection();
@@ -783,11 +828,19 @@ TVM_FFI_STATIC_INIT_BLOCK() {
reinterpret_cast<void*>(&ConstantMutate<StringImmNode>))
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
reinterpret_cast<void*>(&ConstantMaybeInplaceMutate<StringImmNode>));
+ DataTypeImmNode::RegisterReflection();
+ refl::TypeAttrDef<DataTypeImmNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&DataTypeImmVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&DataTypeImmMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&DataTypeImmMaybeInplaceMutate));
refl::GlobalDef()
.def("ir.GenericConst",
[](ffi::Any value, Type ty, Span span) {
return GenericConst(std::move(value), std::move(ty),
std::move(span));
})
+ .def("ir.DataTypeImm",
+ [](DLDataType value, Span span) { return DataTypeImm(value,
std::move(span)); })
.def("ir.StringImm", [](ffi::String value, Span span) {
return StringImm(std::move(value), std::move(span));
});
diff --git a/src/ir/expr_functor.cc b/src/ir/expr_functor.cc
index e076b5c566..046508e987 100644
--- a/src/ir/expr_functor.cc
+++ b/src/ir/expr_functor.cc
@@ -35,6 +35,7 @@ void ExprVisitor::InitVTable(VTable* vtable) {
SetDispatch<ExprVisitor, FloatImmNode>(vtable);
SetDispatch<ExprVisitor, OpNode>(vtable);
SetDispatch<ExprVisitor, StringImmNode>(vtable);
+ SetDispatch<ExprVisitor, DataTypeImmNode>(vtable);
SetDispatch<ExprVisitor, prim::CastNode>(vtable);
SetDispatch<ExprVisitor, prim::AddNode>(vtable);
SetDispatch<ExprVisitor, prim::LShiftNode>(vtable);
@@ -147,6 +148,9 @@ ffi::Optional<VisitInterrupt> ExprVisitor::Visit_(const
OpNode* node) { return s
ffi::Optional<VisitInterrupt> ExprVisitor::Visit_(const StringImmNode* node) {
return std::nullopt;
}
+ffi::Optional<VisitInterrupt> ExprVisitor::Visit_(const DataTypeImmNode* node)
{
+ return std::nullopt;
+}
ffi::Optional<VisitInterrupt> ExprVisitor::Visit_(const prim::CastNode* node) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(node->value));
@@ -344,6 +348,7 @@ void ExprMutator::InitVTable(VTable* vtable) {
SetDispatch<ExprMutator, FloatImmNode>(vtable);
SetDispatch<ExprMutator, OpNode>(vtable);
SetDispatch<ExprMutator, StringImmNode>(vtable);
+ SetDispatch<ExprMutator, DataTypeImmNode>(vtable);
SetDispatch<ExprMutator, prim::CastNode>(vtable);
SetDispatch<ExprMutator, prim::AddNode>(vtable);
SetDispatch<ExprMutator, prim::LShiftNode>(vtable);
@@ -535,6 +540,10 @@ UnchangedOr<Expr> ExprMutator::Mutate_(const
StringImmNode* node, InplaceMode in
// Registry atoms and constant leaves do not descend into metadata or types.
return ffi::Unchanged();
}
+UnchangedOr<Expr> ExprMutator::Mutate_(const DataTypeImmNode* node,
InplaceMode inplace_mode) {
+ // Registry atoms and constant leaves do not descend into metadata or types.
+ return ffi::Unchanged();
+}
UnchangedOr<PrimExpr> ExprMutator::Mutate_(const prim::CastNode* node,
InplaceMode inplace_mode) {
auto value_u = Mutate(node->value, inplace_mode);
diff --git a/src/ir/type.cc b/src/ir/type.cc
index 525ba4fe8d..ba5996858d 100644
--- a/src/ir/type.cc
+++ b/src/ir/type.cc
@@ -255,6 +255,23 @@ TVM_FFI_STATIC_INIT_BLOCK() {
bool Type::IsMissing() const { return this->same_as(Type::Missing()); }
+AnyType::AnyType(Span span) : Type(ffi::UnsafeInit{}) {
+ ffi::ObjectPtr<AnyTypeNode> n = ffi::make_object<AnyTypeNode>();
+ n->span = std::move(span);
+ data_ = std::move(n);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ AnyTypeNode::RegisterReflection();
+ refl::TypeAttrDef<AnyTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&TypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&TypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&TypeMaybeInplaceMutate));
+ refl::GlobalDef().def("ir.AnyType", [](Span span) { return AnyType(span); });
+}
+
OpaqueType::OpaqueType() : Type(ffi::UnsafeInit{}) { data_ =
ffi::make_object<OpaqueTypeNode>(); }
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/relax/backend/vm/codegen_vm.cc
b/src/relax/backend/vm/codegen_vm.cc
index 87c4d69b74..53d01f3f31 100644
--- a/src/relax/backend/vm/codegen_vm.cc
+++ b/src/relax/backend/vm/codegen_vm.cc
@@ -260,6 +260,9 @@ class CodeGenVM : public ExprFunctor<Instruction::Arg(const
Expr&)> {
Instruction::Arg VisitExpr_(const StringImmNode* op) final {
return builder_->ConvertConstant(op->value);
}
+ Instruction::Arg VisitExpr_(const DataTypeImmNode* op) final {
+ return builder_->ConvertConstant(op->value);
+ }
Instruction::Arg VisitExpr_(const TupleNode* op) final {
Tuple tuple = ffi::GetRef<Tuple>(op);
diff --git a/src/relax/backend/vm/codegen_vm_tir.cc
b/src/relax/backend/vm/codegen_vm_tir.cc
index d2571780eb..2537f9b440 100644
--- a/src/relax/backend/vm/codegen_vm_tir.cc
+++ b/src/relax/backend/vm/codegen_vm_tir.cc
@@ -358,6 +358,9 @@ class CodeGenVMTIR : public
ExprFunctor<ffi::Optional<Expr>(const Expr&)> {
ffi::Optional<Expr> VisitExpr_(const StringImmNode* op) final {
return ConstListGet(builder_->ConvertConstant(op->value).value());
}
+ ffi::Optional<Expr> VisitExpr_(const DataTypeImmNode* op) final {
+ return ConstListGet(builder_->ConvertConstant(op->value).value());
+ }
ffi::Optional<Expr> VisitExpr_(const TupleNode* op) final {
Tuple tuple = ffi::GetRef<Tuple>(op);
diff --git a/src/relax/backend/vm/lower_runtime_builtin.cc
b/src/relax/backend/vm/lower_runtime_builtin.cc
index 1e0d7bb696..f96cd1a8a2 100644
--- a/src/relax/backend/vm/lower_runtime_builtin.cc
+++ b/src/relax/backend/vm/lower_runtime_builtin.cc
@@ -86,14 +86,14 @@ class LowerRuntimeBuiltinMutator : public ExprMutator {
Expr MakeMemAllocStorage(const Call& call) {
PrimExpr runtime_device_index = call->args[1].as_or_throw<PrimExpr>();
StringImm storage_scope = call->args[2].as_or_throw<StringImm>();
- GenericConst output_dtype = GenericConst((DLDataType{kDLUInt, 8, 1}),
AnyType());
+ DataTypeImm output_dtype = DataTypeImm((DLDataType{kDLUInt, 8, 1}));
return Call(Type::Missing(), vm_alloc_storage_op_,
{call->args[0], runtime_device_index, output_dtype,
storage_scope}, Attrs());
}
Expr MakeMemAllocTensor(const Call& call) {
PrimExpr offset = call->args[1].as_or_throw<PrimExpr>();
- GenericConst dtype = call->args[3].as_or_throw<GenericConst>();
+ DataTypeImm dtype = call->args[3].as_or_throw<DataTypeImm>();
ffi::Array<Expr> call_args = {call->args[0], offset, call->args[2], dtype};
if (5 == call->args.size()) {
diff --git a/src/relax/backend/vm/vm_shape_lower.cc
b/src/relax/backend/vm/vm_shape_lower.cc
index 0469dff885..a8807e2728 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -804,8 +804,7 @@ class VMShapeLowerMutator
if (always_check || !IsBaseOf(PrimType(op->dtype), GetType(value))) {
// check_shape_info(value, ndim, err_ctx)
Call call(Type::Missing(), builtin_check_prim_value_info_,
- {value, GenericConst(op->dtype, AnyType()),
GetErrContext(err_ctx)}, Attrs(),
- {void_ty_});
+ {value, DataTypeImm(op->dtype), GetErrContext(err_ctx)},
Attrs(), {void_ty_});
builder_->Emit(call, "_");
}
}
@@ -840,9 +839,8 @@ class VMShapeLowerMutator
}
if (always_check || !IsBaseOf(TensorType(op->dtype, op->ndim),
GetType(value))) {
// check_tensor_info(value, ndim, dtype, err_ctx)
- Expr dtype_arg = op->IsUnknownDtype()
- ? Expr(Call(Type::Missing(), null_value_op_, {}))
- : Expr(GenericConst(op->dtype.value()->dtype,
AnyType()));
+ Expr dtype_arg = op->IsUnknownDtype() ? Expr(Call(Type::Missing(),
null_value_op_, {}))
+ :
Expr(DataTypeImm(op->dtype.value()->dtype));
Call call(Type::Missing(), builtin_check_tensor_info_,
{value, IntImm::Int64(op->ndim), dtype_arg,
GetErrContext(err_ctx)}, Attrs(),
{void_ty_});
diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc
index 413f41e5bd..14f92b7f8d 100644
--- a/src/relax/ir/block_builder.cc
+++ b/src/relax/ir/block_builder.cc
@@ -578,6 +578,7 @@ class Normalizer : public BlockBuilderImpl, private
ExprFunctor<Expr(const Expr&
RELAX_EXPR_NORMALIZER_LEAF(GenericConstNode);
RELAX_EXPR_NORMALIZER_LEAF(ShapeExprNode);
RELAX_EXPR_NORMALIZER_LEAF(StringImmNode);
+ RELAX_EXPR_NORMALIZER_LEAF(DataTypeImmNode);
Expr VisitExprDefault_(const ffi::Object* op) final {
Expr expr = ffi::GetRef<Expr>(static_cast<const ExprNode*>(op));
diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc
index 3b7ed49017..406ca9c654 100644
--- a/src/relax/ir/dependent_type.cc
+++ b/src/relax/ir/dependent_type.cc
@@ -35,18 +35,6 @@ using namespace tvm::prim;
namespace {
-TVMFFIAny AnyTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
-}
-
-TVMFFIAny AnyTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
-}
-
-TVMFFIAny AnyTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView)
noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
-}
-
TVMFFIAny ShapeTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
// skips: ndim (scalar)
const ShapeTypeNode* self =
@@ -188,29 +176,6 @@ TVMFFIAny
FuncTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
} // namespace
-AnyType::AnyType(Span span) : Type(ffi::UnsafeInit{}) {
- ffi::ObjectPtr<AnyTypeNode> n = ffi::make_object<AnyTypeNode>();
- n->span = span;
- data_ = std::move(n);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- AnyTypeNode::RegisterReflection();
- refl::TypeAttrDef<AnyTypeNode>()
- .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&AnyTypeVisit))
- .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&AnyTypeMutate))
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&AnyTypeMaybeInplaceMutate));
-
- refl::GlobalDef().def("relax.AnyType", [](Span span) { return AnyType(span);
});
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("relax.ObjectType", [](Span span) { return
AnyType(span); });
-}
-
// Shape
ShapeType::ShapeType(ffi::Array<PrimExpr> values, Span span) :
Type(ffi::UnsafeInit{}) {
ffi::ObjectPtr<ShapeTypeNode> n = ffi::make_object<ShapeTypeNode>();
diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc
index f7ec9d1816..63dab4874f 100644
--- a/src/relax/ir/expr_functor.cc
+++ b/src/relax/ir/expr_functor.cc
@@ -55,6 +55,7 @@
RELAX_VISIT_BINDING_DISPATCH(OpNode);
\
RELAX_VISIT_BINDING_DISPATCH(TupleGetItemNode);
\
RELAX_VISIT_BINDING_DISPATCH(StringImmNode);
\
+ RELAX_VISIT_BINDING_DISPATCH(DataTypeImmNode);
\
return vtable;
\
}
\
void Type::VisitBinding_(const VarBindingNode* binding) {
\
@@ -320,6 +321,7 @@ void ExprVisitor::VisitExprFallback_(const ExprNode* op) {
}
void ExprVisitor::VisitExpr_(const StringImmNode* op) {
this->VisitSpan(op->span); }
+void ExprVisitor::VisitExpr_(const DataTypeImmNode* op) {
this->VisitSpan(op->span); }
void ExprVisitor::VisitSpan(const Span& span) {}
@@ -342,6 +344,7 @@ RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(OpNode);
RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(TupleGetItemNode);
RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(ExprNode);
RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(StringImmNode);
+RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(DataTypeImmNode);
void ExprVisitor::VisitBinding_(const MatchCastNode* binding) {
this->VisitExpr(binding->value);
@@ -675,6 +678,7 @@ Expr ExprMutatorBase::VisitExprFallback_(const ExprNode*
op) {
}
Expr ExprMutatorBase::VisitExpr_(const StringImmNode* op) { return
ffi::GetRef<Expr>(op); }
+Expr ExprMutatorBase::VisitExpr_(const DataTypeImmNode* op) { return
ffi::GetRef<Expr>(op); }
Expr ExprMutatorBase::VisitExpr_(const ShapeExprNode* op) {
auto values = op->values.Map(
@@ -859,6 +863,7 @@ RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(OpNode);
RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(TupleGetItemNode);
RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(ExprNode);
RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(StringImmNode);
+RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(DataTypeImmNode);
void ExprMutator::ReEmitBinding(const VarBindingNode* binding, Expr new_value)
{
Var new_var = this->VisitVarDef(binding->var);
diff --git a/src/relax/ir/py_expr_functor.cc b/src/relax/ir/py_expr_functor.cc
index 45d014ee11..a31a31820f 100644
--- a/src/relax/ir/py_expr_functor.cc
+++ b/src/relax/ir/py_expr_functor.cc
@@ -69,6 +69,8 @@ class PyExprVisitorNode : public ffi::Object, public
ExprVisitor {
ffi::Function f_visit_expr_fallback_{nullptr};
/*! \brief The packed function to the `VisitExpr_(const StringImmNode* op)`
function. */
ffi::Function f_visit_string_imm_{nullptr};
+ /*! \brief The packed function to the `VisitExpr_(const DataTypeImmNode*
op)` function. */
+ ffi::Function f_visit_data_type_imm_{nullptr};
/*! \brief The packed function to the `VisitBinding(const Binding& binding)`
function. */
ffi::Function f_visit_binding{nullptr};
/*! \brief The packed function to the `VisitBinding_(const VarBindingNode*
binding)`
@@ -171,6 +173,7 @@ class PyExprVisitorNode : public ffi::Object, public
ExprVisitor {
PY_EXPR_VISITOR_DISPATCH(OpNode, f_visit_op_);
PY_EXPR_VISITOR_DISPATCH(TupleGetItemNode, f_visit_tuple_getitem_);
PY_EXPR_VISITOR_DISPATCH(StringImmNode, f_visit_string_imm_);
+ PY_EXPR_VISITOR_DISPATCH(DataTypeImmNode, f_visit_data_type_imm_);
vtable.Finalize();
return vtable;
}
@@ -203,6 +206,7 @@ class PyExprVisitor : public ffi::ObjectRef {
* \param f_visit_tuple_getitem_ The packed function of `VisitExpr_(const
TupleGetItemNode* op)`.
* \param f_visit_expr_fallback_ The packed function of the generic
expression fallback.
* \param f_visit_string_imm_ The packed function of `VisitExpr_(const
StringImmNode* op)`.
+ * \param f_visit_data_type_imm_ The packed function of `VisitExpr_(const
DataTypeImmNode* op)`.
* \param f_visit_binding The packed function of `VisitBinding(const
Binding& binding)`.
* \param f_visit_var_binding_ The packed function of `VisitBinding_(const
VarBindingNode*
* binding)`.
@@ -229,11 +233,12 @@ class PyExprVisitor : public ffi::ObjectRef {
ffi::Function f_visit_call_, ffi::Function f_visit_seq_expr_,
ffi::Function f_visit_if_,
ffi::Function f_visit_op_, ffi::Function f_visit_tuple_getitem_,
ffi::Function f_visit_expr_fallback_, ffi::Function f_visit_string_imm_,
- ffi::Function f_visit_binding, ffi::Function f_visit_var_binding_,
- ffi::Function f_visit_match_cast_, ffi::Function f_visit_binding_block,
- ffi::Function f_visit_binding_block_, ffi::Function
f_visit_dataflow_block_,
- ffi::Function f_visit_var_def, ffi::Function f_visit_var_def_,
- ffi::Function f_visit_dataflow_var_def_, ffi::Function f_visit_span) {
+ ffi::Function f_visit_data_type_imm_, ffi::Function f_visit_binding,
+ ffi::Function f_visit_var_binding_, ffi::Function f_visit_match_cast_,
+ ffi::Function f_visit_binding_block, ffi::Function
f_visit_binding_block_,
+ ffi::Function f_visit_dataflow_block_, ffi::Function f_visit_var_def,
+ ffi::Function f_visit_var_def_, ffi::Function f_visit_dataflow_var_def_,
+ ffi::Function f_visit_span) {
ffi::ObjectPtr<PyExprVisitorNode> n =
ffi::make_object<PyExprVisitorNode>();
n->f_visit_expr = f_visit_expr;
n->f_visit_binding = f_visit_binding;
@@ -255,6 +260,7 @@ class PyExprVisitor : public ffi::ObjectRef {
n->f_visit_tuple_getitem_ = f_visit_tuple_getitem_;
n->f_visit_expr_fallback_ = f_visit_expr_fallback_;
n->f_visit_string_imm_ = f_visit_string_imm_;
+ n->f_visit_data_type_imm_ = f_visit_data_type_imm_;
n->f_visit_var_binding_ = f_visit_var_binding_;
n->f_visit_match_cast_ = f_visit_match_cast_;
n->f_visit_binding_block_ = f_visit_binding_block_;
@@ -308,6 +314,8 @@ class PyExprMutatorNode : public ffi::Object, public
ExprMutator {
ffi::Function f_visit_expr_fallback_{nullptr};
/*! \brief The packed function to the `VisitExpr_(const StringImmNode* op)`
function. */
ffi::Function f_visit_string_imm_{nullptr};
+ /*! \brief The packed function to the `VisitExpr_(const DataTypeImmNode*
op)` function. */
+ ffi::Function f_visit_data_type_imm_{nullptr};
/*! \brief The packed function to the `VisitBinding(const Binding& binding)`
function. */
ffi::Function f_visit_binding{nullptr};
/*! \brief The packed function to the `VisitBinding_(const VarBindingNode*
binding)`
@@ -437,6 +445,7 @@ class PyExprMutatorNode : public ffi::Object, public
ExprMutator {
PY_EXPR_MUTATOR_DISPATCH(OpNode, f_visit_op_);
PY_EXPR_MUTATOR_DISPATCH(TupleGetItemNode, f_visit_tuple_getitem_);
PY_EXPR_MUTATOR_DISPATCH(StringImmNode, f_visit_string_imm_);
+ PY_EXPR_MUTATOR_DISPATCH(DataTypeImmNode, f_visit_data_type_imm_);
vtable.Finalize();
return vtable;
}
@@ -459,6 +468,7 @@ class PyExprMutatorNode : public ffi::Object, public
ExprMutator {
PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(OpNode);
PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(TupleGetItemNode);
PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(StringImmNode);
+ PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(DataTypeImmNode);
post_order_vtable.Finalize();
return post_order_vtable;
}
@@ -491,6 +501,7 @@ class PyExprMutator : public ffi::ObjectRef {
* \param f_visit_tuple_getitem_ The packed function of `VisitExpr_(const
TupleGetItemNode* op)`.
* \param f_visit_expr_fallback_ The packed function of the generic
expression fallback.
* \param f_visit_string_imm_ The packed function of `VisitExpr_(const
StringImmNode* op)`.
+ * \param f_visit_data_type_imm_ The packed function of `VisitExpr_(const
DataTypeImmNode* op)`.
* \param f_visit_binding The packed function of `VisitBinding(const
Binding& binding)`.
* \param f_visit_var_binding_ The packed function of `VisitBinding_(const
VarBindingNode*
* binding)`.
@@ -517,11 +528,12 @@ class PyExprMutator : public ffi::ObjectRef {
ffi::Function f_visit_call_, ffi::Function f_visit_seq_expr_,
ffi::Function f_visit_if_,
ffi::Function f_visit_op_, ffi::Function f_visit_tuple_getitem_,
ffi::Function f_visit_expr_fallback_, ffi::Function f_visit_string_imm_,
- ffi::Function f_visit_binding, ffi::Function f_visit_var_binding_,
- ffi::Function f_visit_match_cast_, ffi::Function f_visit_binding_block,
- ffi::Function f_visit_binding_block_, ffi::Function
f_visit_dataflow_block_,
- ffi::Function f_visit_var_def, ffi::Function f_visit_var_def_,
- ffi::Function f_visit_dataflow_var_def_, ffi::Function f_visit_span) {
+ ffi::Function f_visit_data_type_imm_, ffi::Function f_visit_binding,
+ ffi::Function f_visit_var_binding_, ffi::Function f_visit_match_cast_,
+ ffi::Function f_visit_binding_block, ffi::Function
f_visit_binding_block_,
+ ffi::Function f_visit_dataflow_block_, ffi::Function f_visit_var_def,
+ ffi::Function f_visit_var_def_, ffi::Function f_visit_dataflow_var_def_,
+ ffi::Function f_visit_span) {
ffi::ObjectPtr<PyExprMutatorNode> n =
ffi::make_object<PyExprMutatorNode>();
n->builder_ = builder_;
n->f_visit_expr = f_visit_expr;
@@ -540,6 +552,7 @@ class PyExprMutator : public ffi::ObjectRef {
n->f_visit_tuple_getitem_ = f_visit_tuple_getitem_;
n->f_visit_expr_fallback_ = f_visit_expr_fallback_;
n->f_visit_string_imm_ = f_visit_string_imm_;
+ n->f_visit_data_type_imm_ = f_visit_data_type_imm_;
n->f_visit_binding = f_visit_binding;
n->f_visit_var_binding_ = f_visit_var_binding_;
n->f_visit_match_cast_ = f_visit_match_cast_;
diff --git a/src/relax/op/memory/view.cc b/src/relax/op/memory/view.cc
index 3b2a83ff04..39353d2a35 100644
--- a/src/relax/op/memory/view.cc
+++ b/src/relax/op/memory/view.cc
@@ -121,23 +121,23 @@ Type InferTypeView(const Call& call, const BlockBuilder&
ctx) {
// In general, Type inference should only depend on the
// Type of the arguments, and not on the arguments
- // themselves. However, dtype GenericConst uses
+ // themselves. However, dtype DataTypeImm uses
// `AnyType`, so we need to inspect the argument itself
// in this case.
- if (auto dtype_imm = arg_value.as<GenericConstNode>()) {
+ if (auto dtype_imm = arg_value.as<DataTypeImmNode>()) {
// We know the datatype for the view.
- TVM_FFI_CHECK(!PrimType(dtype_imm->value.cast<DLDataType>()).IsVoid(),
TypeError)
+ TVM_FFI_CHECK(!PrimType(dtype_imm->value).IsVoid(), TypeError)
<< "Operator " << call->op
<< " expects the dtype argument to be a concrete tensor datatype, "
<< "but received void. Use relax.null_value() if no dtype change is
requested.";
- return PrimType(dtype_imm->value.cast<DLDataType>());
+ return PrimType(dtype_imm->value);
} else if (ty.as<AnyTypeNode>()) {
// The view changes the datatype, but we don't know what it is
// being changed into.
return std::nullopt;
} else {
TVM_FFI_THROW(TypeError) << "Operator " << call->op
- << " expects the dtype argument to be a dtype
GenericConst, "
+ << " expects the dtype argument to be a dtype
DataTypeImm, "
<< "but received " << arg_dtype << " with type
" << ty;
}
}();
@@ -376,7 +376,7 @@ Expr LowerBuiltinView(const BlockBuilder& bb, const Call&
call) {
<< "or the input dtype is known. "
<< "However, in expression " << call << ", no output dtype is
specified, "
<< "and the input " << data << " of type " << data->ty << " has
unknown dtype.";
- dtype = GenericConst(data_tensor_ty->dtype.value()->dtype, AnyType());
+ dtype = DataTypeImm(data_tensor_ty->dtype.value()->dtype);
}
if (HasVoidType(relative_byte_offset)) {
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index 3826a83e5a..4c344c18ca 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -1120,12 +1120,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Type InferTypeAllocateTensor(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_ICHECK(call->args[0].as<ShapeExprNode>())
<< "must be ShapeExpr, but got " << call->args[0]->GetTypeKey();
- TVM_FFI_ICHECK(call->args[1].as<GenericConstNode>())
- << "must be GenericConst, but got " << call->args[1]->GetTypeKey();
+ TVM_FFI_ICHECK(call->args[1].as<DataTypeImmNode>())
+ << "must be DataTypeImm, but got " << call->args[1]->GetTypeKey();
PrimType out_dtype = PrimType::Void();
- if (const auto* dtype_node = call->args[1].as<GenericConstNode>()) {
- const GenericConst dtype_imm = ffi::GetRef<GenericConst>(dtype_node);
- out_dtype = PrimType(dtype_imm->value.cast<DLDataType>());
+ if (const auto* dtype_node = call->args[1].as<DataTypeImmNode>()) {
+ const DataTypeImm dtype_imm = ffi::GetRef<DataTypeImm>(dtype_node);
+ out_dtype = PrimType(dtype_imm->value);
}
int64_t vdevice_index = -1;
if (const auto* int_imm = call->args[2].as<IntImmNode>()) {
@@ -1142,7 +1142,7 @@ Type InferTypeAllocateTensor(const Call& call, const
BlockBuilder& ctx) {
TVM_REGISTER_OP("relax.builtin.alloc_tensor")
.set_num_inputs(4)
.add_argument("shape", "Expr", "The shape of the tensor to allocate.")
- .add_argument("dtype", "GenericConst", "The dtype of the tensor to
allocate.")
+ .add_argument("dtype", "DataTypeImm", "The dtype of the tensor to
allocate.")
.add_argument("runtime_device_index", "PrimExpr",
"The device index indicating on which device the tensor is
to be "
"allocated at runtime. Index -1 is reserved for the host
device.")
@@ -1153,7 +1153,7 @@ TVM_REGISTER_OP("relax.builtin.alloc_tensor")
.set_attr<bool>("FPurity", true)
.set_attr<bool>("TAllocator", true);
-Expr MakeAllocTensor(Expr shape, GenericConst dtype, PrimExpr
runtime_device_index,
+Expr MakeAllocTensor(Expr shape, DataTypeImm dtype, PrimExpr
runtime_device_index,
StringImm storage_scope) {
static const Op& op = Op::Get("relax.builtin.alloc_tensor");
return Call(Type::Missing(), op, {shape, dtype, runtime_device_index,
storage_scope}, Attrs(),
@@ -1176,14 +1176,14 @@ TVM_REGISTER_OP("relax.memory.alloc_storage")
"Index -1 is reserved for the host device.")
.add_argument("storage_scope", "StringImm",
"The storage scope of the storage to allocate. Default is
global.")
- .add_argument("dtype", "GenericConst", "The dtype of the tensor to
allocate.")
+ .add_argument("dtype", "DataTypeImm", "The dtype of the tensor to
allocate.")
.set_attr<FInferType>("FInferType", ReturnAnyType)
// memory allocation isn't considered a "visible effect" as far as purity
is concerned
.set_attr<bool>("FPurity", true)
.set_attr<bool>("TAllocator", true);
Expr MakeAllocStorage(Expr size, PrimExpr virtual_device_index, StringImm
storage_scope,
- GenericConst dtype) {
+ DataTypeImm dtype) {
static const Op& op = Op::Get("relax.memory.alloc_storage");
return Call(Type::Missing(), op, {size, virtual_device_index, storage_scope,
dtype}, Attrs(), {});
}
@@ -1199,9 +1199,9 @@ Type InferTypeMemAllocTensor(const Call& call, const
BlockBuilder& ctx) {
TVM_FFI_ICHECK(GetTypeAs<ShapeTypeNode>(call->args[2]))
<< "must be a Expr of ShapeType, but got " <<
call->args[1]->GetTypeKey();
PrimType out_dtype = PrimType::Void();
- if (const auto* dtype_node = call->args[3].as<GenericConstNode>()) {
- const GenericConst dtype_imm = ffi::GetRef<GenericConst>(dtype_node);
- out_dtype = PrimType(dtype_imm->value.cast<DLDataType>());
+ if (const auto* dtype_node = call->args[3].as<DataTypeImmNode>()) {
+ const DataTypeImm dtype_imm = ffi::GetRef<DataTypeImm>(dtype_node);
+ out_dtype = PrimType(dtype_imm->value);
}
if (call->args.size() == 5) {
@@ -1223,7 +1223,7 @@ TVM_REGISTER_OP("relax.memory.alloc_tensor")
.add_argument("storage", "Expr", "The storage to allocate the tensor to.")
.add_argument("offset", "PrimExpr", "Storage offset to allocate the
tensor.")
.add_argument("shape", "Expr", "The shape of the tensor to allocate.")
- .add_argument("dtype", "GenericConst", "The dtype of the tensor to
allocate.")
+ .add_argument("dtype", "DataTypeImm", "The dtype of the tensor to
allocate.")
.add_argument("runtime_device_index", "PrimExpr",
"The device index indicating on which device the tensor is
to be "
"allocated at runtime. Index -1 is reserved for the host
device.")
@@ -1232,7 +1232,7 @@ TVM_REGISTER_OP("relax.memory.alloc_tensor")
.set_attr<bool>("FPurity", true)
.set_attr<bool>("TAllocator", true);
-Expr MakeMemAllocTensor(Expr storage, PrimExpr offset, Expr shape,
GenericConst dtype,
+Expr MakeMemAllocTensor(Expr storage, PrimExpr offset, Expr shape, DataTypeImm
dtype,
PrimExpr virtual_device_index) {
static const Op& op = Op::Get("relax.memory.alloc_tensor");
return Call(Type::Missing(), op, {storage, offset, shape, dtype,
virtual_device_index}, Attrs(),
@@ -1245,11 +1245,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"relax.op.memory.alloc_tensor", [](ffi::PackedArgs args, ffi::Any* ret) {
if (args.size() == 5) {
*ret = MakeMemAllocTensor(args[0].cast<Expr>(),
args[1].cast<PrimExpr>(),
- args[2].cast<Expr>(),
args[3].cast<GenericConst>(),
+ args[2].cast<Expr>(),
args[3].cast<DataTypeImm>(),
args[4].cast<PrimExpr>());
} else {
*ret = MakeMemAllocTensor(args[0].cast<Expr>(),
args[1].cast<PrimExpr>(),
- args[2].cast<Expr>(),
args[3].cast<GenericConst>(),
+ args[2].cast<Expr>(),
args[3].cast<DataTypeImm>(),
IntImm::Int64(0));
}
});
@@ -1298,7 +1298,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TVM_REGISTER_OP("relax.vm.alloc_storage")
.set_num_inputs(4)
.add_argument("size", "Expr", "The size of the storage to allocate.")
- .add_argument("dtype", "GenericConst", "The dtype of the tensor to
allocate.")
+ .add_argument("dtype", "DataTypeImm", "The dtype of the tensor to
allocate.")
.add_argument("runtime_device_index", "PrimExpr",
"The device index indicating on which device the tensor is "
"to be allocated at runtime.")
@@ -1309,7 +1309,7 @@ TVM_REGISTER_OP("relax.vm.alloc_storage")
.set_attr<bool>("FPurity", true)
.set_attr<bool>("TAllocator", true);
-Expr MakeVMAllocStorage(Expr size, PrimExpr runtime_device_index, GenericConst
dtype,
+Expr MakeVMAllocStorage(Expr size, PrimExpr runtime_device_index, DataTypeImm
dtype,
StringImm storage_scope) {
static const Op& op = Op::Get("relax.vm.alloc_storage");
return Call(Type::Missing(), op, {size, runtime_device_index, dtype,
storage_scope}, Attrs(), {});
@@ -1324,9 +1324,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Type InferTypeVMAllocTensor(const Call& call, const BlockBuilder& ctx) {
PrimType out_dtype = PrimType::Void();
- if (const auto* dtype_node = call->args[3].as<GenericConstNode>()) {
- const GenericConst dtype_imm = ffi::GetRef<GenericConst>(dtype_node);
- out_dtype = PrimType(dtype_imm->value.cast<DLDataType>());
+ if (const auto* dtype_node = call->args[3].as<DataTypeImmNode>()) {
+ const DataTypeImm dtype_imm = ffi::GetRef<DataTypeImm>(dtype_node);
+ out_dtype = PrimType(dtype_imm->value);
}
int64_t vdevice_index = -1;
if (const auto* int_imm = call->args[4].as<IntImmNode>()) {
@@ -1351,7 +1351,7 @@ TVM_REGISTER_OP("relax.vm.alloc_tensor")
.add_argument("storage", "Expr", "The storage to allocate the tensor to.")
.add_argument("offset", "PrimExpr", "Storage offset to allocate the
tensor.")
.add_argument("shape", "Expr", "The shape of the tensor to allocate.")
- .add_argument("dtype", "GenericConst", "The dtype of the tensor to
allocate.")
+ .add_argument("dtype", "DataTypeImm", "The dtype of the tensor to
allocate.")
.add_argument("runtime_device_index", "PrimExpr",
"The device index indicating on which device the tensor is "
"to be allocated at runtime.")
@@ -1360,7 +1360,7 @@ TVM_REGISTER_OP("relax.vm.alloc_tensor")
.set_attr<bool>("FPurity", true)
.set_attr<bool>("TAllocator", true);
-Expr MakeVMAllocTensor(Expr storage, PrimExpr offset, Expr shape, GenericConst
dtype,
+Expr MakeVMAllocTensor(Expr storage, PrimExpr offset, Expr shape, DataTypeImm
dtype,
PrimExpr runtime_device_index) {
static const Op& op = Op::Get("relax.vm.alloc_tensor");
return Call(Type::Missing(), op, {storage, offset, shape, dtype,
runtime_device_index}, Attrs(),
@@ -1372,10 +1372,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def_packed("relax.op.vm.alloc_tensor", [](ffi::PackedArgs
args, ffi::Any* ret) {
if (args.size() == 5) {
*ret = MakeVMAllocTensor(args[0].cast<Expr>(), args[1].cast<PrimExpr>(),
args[2].cast<Expr>(),
- args[3].cast<GenericConst>(),
args[4].cast<PrimExpr>());
+ args[3].cast<DataTypeImm>(),
args[4].cast<PrimExpr>());
} else {
*ret = MakeVMAllocTensor(args[0].cast<Expr>(), args[1].cast<PrimExpr>(),
args[2].cast<Expr>(),
- args[3].cast<GenericConst>(), IntImm::Int64(0));
+ args[3].cast<DataTypeImm>(), IntImm::Int64(0));
}
});
}
diff --git a/src/relax/op/op_common.h b/src/relax/op/op_common.h
index da4013e899..953892c1f1 100644
--- a/src/relax/op/op_common.h
+++ b/src/relax/op/op_common.h
@@ -606,12 +606,12 @@ inline ffi::Optional<ShapeExpr>
CheckNdimPerLayoutAndGetShape(const Call& call,
return std::nullopt;
}
-Expr MakeVMAllocStorage(Expr size, PrimExpr runtime_device_index, GenericConst
dtype,
+Expr MakeVMAllocStorage(Expr size, PrimExpr runtime_device_index, DataTypeImm
dtype,
StringImm storage_scope = StringImm("global"));
-Expr MakeVMAllocTensor(Expr storage, PrimExpr offset, Expr shape, GenericConst
dtype,
+Expr MakeVMAllocTensor(Expr storage, PrimExpr offset, Expr shape, DataTypeImm
dtype,
PrimExpr runtime_device_index);
-Expr MakeAllocTensor(Expr shape, GenericConst dtype, PrimExpr
runtime_device_index,
+Expr MakeAllocTensor(Expr shape, DataTypeImm dtype, PrimExpr
runtime_device_index,
StringImm storage_scope = StringImm("global"));
/**
diff --git a/src/relax/script/ir_builder/frame.cc
b/src/relax/script/ir_builder/frame.cc
index 0c5ec90534..fb47653d2f 100644
--- a/src/relax/script/ir_builder/frame.cc
+++ b/src/relax/script/ir_builder/frame.cc
@@ -86,7 +86,7 @@ void FunctionFrameNode::ExitWithScope() {
RelaxFrameNode::ExitWithScope();
block_builder->EndScope();
function =
- tvm::relax::Function::CreateEmpty(params,
ret_ty.value_or(tvm::relax::AnyType()),
+ tvm::relax::Function::CreateEmpty(params,
ret_ty.value_or(tvm::AnyType()),
is_pure.value_or(true),
DictAttrs(attrs), source_span);
if (local) {
auto ty = tvm::relax::GetType(function.value());
diff --git a/src/relax/script/printer/dependent_type.cc
b/src/relax/script/printer/dependent_type.cc
index 026819dcd9..433932502a 100644
--- a/src/relax/script/printer/dependent_type.cc
+++ b/src/relax/script/printer/dependent_type.cc
@@ -26,8 +26,8 @@ namespace script {
namespace printer {
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::AnyType>( //
- "", [](relax::AnyType n, AccessPath n_p, IRDocsifier d) -> Doc { return
Relax(d, "Any"); });
+ IRDocsifier::vtable().set_dispatch<AnyType>( //
+ "", [](AnyType n, AccessPath n_p, IRDocsifier d) -> Doc { return
Relax(d, "Any"); });
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -125,7 +125,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::Array<ffi::String> keys;
ffi::Array<ExprDoc, void> values;
- if (!n->ret->IsInstance<relax::AnyTypeNode>()) {
+ if (!n->ret->IsInstance<AnyTypeNode>()) {
keys.push_back("ret");
values.push_back(ret_doc);
}
@@ -151,7 +151,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(relax::AnyTypeNode, ReprPrintRelax);
+TVM_REGISTER_SCRIPT_AS_REPR(AnyTypeNode, ReprPrintRelax);
TVM_REGISTER_SCRIPT_AS_REPR(relax::ShapeTypeNode, ReprPrintRelax);
TVM_REGISTER_SCRIPT_AS_REPR(relax::TensorTypeNode, ReprPrintRelax);
TVM_REGISTER_SCRIPT_AS_REPR(relax::FuncTypeNode, ReprPrintRelax);
diff --git a/src/relax/script/printer/expr.cc b/src/relax/script/printer/expr.cc
index c1901459e5..176d01ae83 100644
--- a/src/relax/script/printer/expr.cc
+++ b/src/relax/script/printer/expr.cc
@@ -130,9 +130,9 @@ ffi::Optional<ExprDoc> SpecialScalar(const runtime::Tensor&
n, const AccessPath&
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<::tvm::GenericConst>( //
- "", [](::tvm::GenericConst n, AccessPath n_p, IRDocsifier d) -> Doc {
- if (auto dtype = n->value.as<DLDataType>()) {
- return Relax(d, "dtype")->Call({LiteralDoc::DataType(*dtype,
n_p->Attr("value"))});
+ "relax", [](::tvm::GenericConst n, AccessPath n_p, IRDocsifier d) -> Doc
{
+ if (n->value.as<DLDataType>()) {
+ return IRDocsifier::vtable()("", n, n_p, d);
}
auto data = n->value.cast<runtime::Tensor>();
if (ffi::Optional<ExprDoc> s = SpecialScalar(data,
n_p->Attr("value"))) {
@@ -179,6 +179,7 @@ TVM_REGISTER_SCRIPT_AS_REPR(relax::ShapeExprNode,
ReprPrintRelax);
TVM_REGISTER_SCRIPT_AS_REPR(VarNode, ReprPrintVar);
TVM_REGISTER_SCRIPT_AS_REPR(relax::DataflowVarNode, ReprPrintRelax);
TVM_REGISTER_SCRIPT_AS_REPR(::tvm::GenericConstNode, ReprPrintRelax);
+TVM_REGISTER_SCRIPT_AS_REPR(::tvm::DataTypeImmNode, ReprPrintRelax);
} // namespace printer
} // namespace script
diff --git a/src/relax/script/printer/function.cc
b/src/relax/script/printer/function.cc
index bd27d16ba5..7e5cc28b8f 100644
--- a/src/relax/script/printer/function.cc
+++ b/src/relax/script/printer/function.cc
@@ -27,7 +27,7 @@ namespace printer {
static bool HasDefaultExternFuncType(const relax::ExternFunc& n) {
const auto* ty = n->ty.as<relax::FuncTypeNode>();
if (ty == nullptr || ty->params.has_value() || ty->purity ||
- !ty->ret->IsInstance<relax::AnyTypeNode>()) {
+ !ty->ret->IsInstance<AnyTypeNode>()) {
return false;
}
return true;
diff --git a/src/relax/transform/allocate_workspace.cc
b/src/relax/transform/allocate_workspace.cc
index 0937dc83e1..ab80be5119 100644
--- a/src/relax/transform/allocate_workspace.cc
+++ b/src/relax/transform/allocate_workspace.cc
@@ -150,7 +150,7 @@ class WorkspaceProvider : ExprMutator {
builder_->BeginDataflowBlock();
if (!workspace_var_main_.defined()) {
auto shape = ShapeExpr({IntImm::Int32(max_workspace_size_)});
- auto ty = GenericConst((DLDataType{kDLUInt, 8, 1}), AnyType());
+ auto ty = DataTypeImm((DLDataType{kDLUInt, 8, 1}));
auto workspace = MakeAllocTensor(shape, ty, IntImm::Int64(0));
workspace_var_main_ = builder_->Emit(workspace, "workspace_main");
}
diff --git a/src/relax/transform/call_tir_rewrite.cc
b/src/relax/transform/call_tir_rewrite.cc
index 67cf180cb7..20c7027c97 100644
--- a/src/relax/transform/call_tir_rewrite.cc
+++ b/src/relax/transform/call_tir_rewrite.cc
@@ -176,7 +176,7 @@ class CallTIRMutator : public ExprMutator {
return builder_->Emit(Call(Type::Missing(), alloc_tensor_op,
{tensor_ty->shape.value().as_or_throw<ShapeExpr>(),
- GenericConst(tensor_ty->dtype.value()->dtype,
AnyType()),
+ DataTypeImm(tensor_ty->dtype.value()->dtype),
IntImm::Int64(dev_index), StringImm(scope)},
Attrs(), {tensor_ty}),
"alloc");
diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc
index 354b08fe8f..937f4cd7c6 100644
--- a/src/relax/transform/fuse_ops.cc
+++ b/src/relax/transform/fuse_ops.cc
@@ -279,7 +279,7 @@ class GraphCreator : public ExprVisitor {
if (!leaf_expr.as<ShapeExprNode>() && !leaf_expr.as<VarNode>() &&
!leaf_expr.as<GenericConstNode>() && !leaf_expr.as<PrimExpr>() &&
- !leaf_expr.as<StringImmNode>()) {
+ !leaf_expr.as<StringImmNode>() && !leaf_expr.as<DataTypeImmNode>()) {
// Skip GlobalVar, ExternFunc, OpNode.
return;
}
@@ -1213,6 +1213,7 @@ class PatternBasedPartitioner : ExprVisitor {
}
void VisitExpr_(const GenericConstNode* op) final { group_map_[op] =
arena_->make<Group>(); }
+ void VisitExpr_(const DataTypeImmNode* op) final { group_map_[op] =
arena_->make<Group>(); }
void VisitBinding_(const VarBindingNode* binding, const CallNode* call)
final {
VisitVarDef(binding->var);
diff --git a/src/relax/transform/lower_alloc_tensor.cc
b/src/relax/transform/lower_alloc_tensor.cc
index dd46331bf8..ff382f4de6 100644
--- a/src/relax/transform/lower_alloc_tensor.cc
+++ b/src/relax/transform/lower_alloc_tensor.cc
@@ -49,7 +49,7 @@ class Mutator : public ExprMutator {
<< "However, received " << ffi::GetRef<Call>(op);
auto shape_arg = op->args[0];
- auto dtype = op->args[1].as_or_throw<GenericConst>();
+ auto dtype = op->args[1].as_or_throw<DataTypeImm>();
PrimExpr runtime_device_index = op->args[2].as_or_throw<PrimExpr>();
StringImm storage_scope = op->args[3].as_or_throw<StringImm>();
@@ -73,7 +73,7 @@ class Mutator : public ExprMutator {
}();
PrimExpr nbytes = [&]() -> PrimExpr {
- PrimType dtype_ty(dtype->value.cast<DLDataType>());
+ PrimType dtype_ty(dtype->value);
TVM_FFI_ICHECK(!dtype_ty.IsScalableVector())
<< "Cannot statically compute allocation size for scalable vector
dtype " << dtype_ty;
PrimExpr nbytes =
IntImm::Int64(static_cast<int64_t>(dtype_ty.StorageBytes()));
@@ -115,9 +115,9 @@ class Mutator : public ExprMutator {
auto offset = IntImm::Int64(0);
- Expr storage = Call(Type::Missing(), mem_alloc_storage_op,
- {size, runtime_device_index, storage_scope,
- GenericConst((DLDataType{kDLUInt, 8, 1}),
AnyType())});
+ Expr storage = Call(
+ Type::Missing(), mem_alloc_storage_op,
+ {size, runtime_device_index, storage_scope,
DataTypeImm((DLDataType{kDLUInt, 8, 1}))});
storage = builder_->Emit(storage, "storage");
Expr tensor = Call(Type::Missing(), mem_alloc_tensor_op,
{storage, offset, shape_arg, dtype, op->args[2]});
diff --git a/src/relax/transform/rewrite_cuda_graph.cc
b/src/relax/transform/rewrite_cuda_graph.cc
index ba6a6e9e63..a2995ea0f7 100644
--- a/src/relax/transform/rewrite_cuda_graph.cc
+++ b/src/relax/transform/rewrite_cuda_graph.cc
@@ -505,8 +505,8 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
bool IsStatic(const Expr& expr, std::vector<const VarNode*>* vars_collector
= nullptr,
std::vector<PrimVar>* tir_vars_collector = nullptr) {
- if (expr->IsInstance<GenericConstNode>() ||
expr->IsInstance<StringImmNode>() ||
- expr->IsInstance<GlobalVarNode>()) {
+ if (expr->IsInstance<GenericConstNode>() ||
expr->IsInstance<DataTypeImmNode>() ||
+ expr->IsInstance<StringImmNode>() ||
expr->IsInstance<GlobalVarNode>()) {
return true;
}
if (const auto* var = expr.as<VarNode>();
diff --git a/src/relax/transform/static_plan_block_memory.cc
b/src/relax/transform/static_plan_block_memory.cc
index 1871a905a6..53059da3e6 100644
--- a/src/relax/transform/static_plan_block_memory.cc
+++ b/src/relax/transform/static_plan_block_memory.cc
@@ -653,8 +653,7 @@ class StorageAllocatorInit : public
StorageAllocatorBaseVisitor {
const auto* shape = ty->shape.as<ShapeExprNode>();
TVM_FFI_ICHECK_NOTNULL(shape);
TVM_FFI_ICHECK(!ty->IsUnknownDtype());
- TVM_FFI_ICHECK(ty->dtype.value()->dtype ==
-
call->args[1].as_or_throw<GenericConst>()->value.cast<DLDataType>());
+ TVM_FFI_ICHECK(ty->dtype.value()->dtype ==
call->args[1].as_or_throw<DataTypeImm>()->value);
TVM_FFI_ICHECK(!token_map_.count(call));
// Use the upper bounds of TIR vars as their values. The upper bound shape
can still be dynamic
@@ -960,7 +959,7 @@ class StorageAllocationRewriter : public ExprMutator {
DLDataType dtype = token->dtype;
Call alloc_storage(Type::Missing(), mem_alloc_storage,
{std::move(size), virtual_device_index,
StringImm(token->storage_scope),
- GenericConst(dtype, AnyType())},
+ DataTypeImm(dtype)},
Attrs());
storage_var = builder_->Emit(alloc_storage, "storage");
token2storage_var_[token.get()] = storage_var;
@@ -971,10 +970,9 @@ class StorageAllocationRewriter : public ExprMutator {
// And always create a `memory.alloc_tensor` for the old
`builtin.alloc_tensor`.
PrimExpr offset = IntImm::Int64(0);
DLDataType dtype = ty->dtype.value()->dtype;
- return Call(
- Type::Missing(), mem_alloc_tensor,
- {storage_var, offset, ty->shape.value(), GenericConst(dtype,
AnyType()), call->args[2]},
- Attrs());
+ return Call(Type::Missing(), mem_alloc_tensor,
+ {storage_var, offset, ty->shape.value(), DataTypeImm(dtype),
call->args[2]},
+ Attrs());
} else if (plan_dynamic_output_ && call->op.same_as(alloc_tensor_op)) {
// Case 2. For a `alloc_tensor` that is not planned for memory reuse,
// we would still like to allocate **static** memory for the tensor.
@@ -991,7 +989,7 @@ class StorageAllocationRewriter : public ExprMutator {
if (!IsStaticShape(shape->values)) {
TVM_FFI_ICHECK(!ty->IsUnknownDtype());
TVM_FFI_ICHECK_EQ(ty->dtype.value()->dtype,
-
call->args[1].as_or_throw<GenericConst>()->value.cast<DLDataType>());
+ call->args[1].as_or_throw<DataTypeImm>()->value);
PrimExpr bytes = upper_bounded_shape[0];
for (int i = 1; i < static_cast<int>(upper_bounded_shape.size()); ++i)
{
bytes *= upper_bounded_shape[i];
@@ -1005,13 +1003,13 @@ class StorageAllocationRewriter : public ExprMutator {
{/*size=*/ShapeExpr({bytes}),
/*virtual_device_index=*/call->args[2].as_or_throw<PrimExpr>(),
/*storage_scope=*/call->args[3].as_or_throw<StringImm>(), //
- /*dtype=*/GenericConst(dtype, AnyType())});
+ /*dtype=*/DataTypeImm(dtype)});
Var storage = builder_->Emit(alloc_storage, "storage");
return Call(Type::Missing(), mem_alloc_tensor,
{storage, //
/*offset=*/IntImm::Int64(0),
/*shape=*/ffi::GetRef<ShapeExpr>(shape), //
- /*dtype=*/GenericConst(dtype, AnyType()),
+ /*dtype=*/DataTypeImm(dtype),
/*vdevice_index=*/call->args[2]});
}
}
diff --git a/src/script/printer/ir/misc.cc b/src/script/printer/ir/misc.cc
index 53e57f8afb..400cb4fe5d 100644
--- a/src/script/printer/ir/misc.cc
+++ b/src/script/printer/ir/misc.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/container/shape.h>
+#include <tvm/runtime/tensor.h>
#include "./utils.h"
@@ -24,6 +25,38 @@ namespace tvm {
namespace script {
namespace printer {
+TVM_FFI_STATIC_INIT_BLOCK() {
+ IRDocsifier::vtable().set_dispatch<DataTypeImm>(
+ "", [](DataTypeImm n, AccessPath p, IRDocsifier d) -> Doc {
+ return TIR(d, "dtype")->Call({LiteralDoc::DataType(n->value,
p->Attr("value"))});
+ });
+ IRDocsifier::vtable().set_dispatch<DataTypeImm>(
+ "ir", [](DataTypeImm n, AccessPath p, IRDocsifier d) -> Doc {
+ return IR(d, "dtype")->Call({LiteralDoc::DataType(n->value,
p->Attr("value"))});
+ });
+ IRDocsifier::vtable().set_dispatch<GenericConst>(
+ "", [](GenericConst n, AccessPath p, IRDocsifier d) -> Doc {
+ if (auto dtype = n->value.as<DLDataType>()) {
+ return TIR(d, "dtype")->Call({LiteralDoc::DataType(*dtype,
p->Attr("value"))});
+ }
+ if (n->value.as<runtime::Tensor>()) {
+ // Tensor-valued constants are Relax constants and keep the Relax
spelling.
+ return IRDocsifier::vtable()("relax", n, p, d);
+ }
+ return d->AddMetadata(n);
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ IRDocsifier::vtable().set_dispatch<GenericConst>(
+ "ir", [](GenericConst n, AccessPath p, IRDocsifier d) -> Doc {
+ if (auto dtype = n->value.as<DLDataType>()) {
+ return IR(d, "dtype")->Call({LiteralDoc::DataType(*dtype,
p->Attr("value"))});
+ }
+ return IRDocsifier::vtable()("", n, p, d);
+ });
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<ffi::Array<Any>>( //
"", [](ffi::Array<Any> array, AccessPath p, IRDocsifier d) -> Doc {
diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc
index 3b87c1a62c..ccc2b25228 100644
--- a/src/tirx/ir/function.cc
+++ b/src/tirx/ir/function.cc
@@ -53,7 +53,7 @@ tvm::Type InferType(const PrimFunc& prim_func) {
// same Relax-facing wildcard semantics that opaque handle parameters had
// before pointers became exact IR types.
if (param->ty.as<PointerTypeNode>()) {
- return relax::AnyType();
+ return AnyType();
}
return param->ty;
@@ -67,7 +67,7 @@ tvm::Type InferType(const PrimFunc& prim_func) {
} else if (IsVoidType(prim_func->ret_type)) {
return relax::TupleType(ffi::Array<tvm::Type>{});
} else {
- return relax::AnyType();
+ return AnyType();
}
}();
diff --git a/src/tirx/ir/tir_visitor_with_path.cc
b/src/tirx/ir/tir_visitor_with_path.cc
index 6757e1b01a..95ad243f88 100644
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ b/src/tirx/ir/tir_visitor_with_path.cc
@@ -364,6 +364,8 @@ DEFINE_BINOP_VISIT_(prim::BitwiseXorNode);
void TIRVisitorWithPath::Dispatch_(const IntImmNode* op, AccessPath path) {}
void TIRVisitorWithPath::Dispatch_(const FloatImmNode* op, AccessPath path) {}
void TIRVisitorWithPath::Dispatch_(const StringImmNode* op, AccessPath path) {}
+void TIRVisitorWithPath::Dispatch_(const DataTypeImmNode* op, AccessPath path)
{}
+void TIRVisitorWithPath::Dispatch_(const GenericConstNode* op, AccessPath
path) {}
void TIRVisitorWithPath::Dispatch_(const prim::CastNode* op, AccessPath path) {
Visit(op->value, path->Attr("value"));
diff --git a/src/tirx/ir/tir_visitor_with_path.h
b/src/tirx/ir/tir_visitor_with_path.h
index 56305d8698..101951a77b 100644
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ b/src/tirx/ir/tir_visitor_with_path.h
@@ -64,8 +64,12 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
virtual inline void Visit(const Expr& obj, ffi::reflection::AccessPath path)
{
if (auto prim = obj.as<PrimExpr>()) {
Visit(prim.value(), path);
+ } else if (auto* constant = obj.as<GenericConstNode>()) {
+ Dispatch_(constant, path);
} else if (auto* str = obj.as<StringImmNode>()) {
Dispatch_(str, path);
+ } else if (auto* str = obj.as<DataTypeImmNode>()) {
+ Dispatch_(str, path);
} else if (auto* var = obj.as<VarNode>()) {
Dispatch_(var, path);
} else if (auto* call = obj.as<CallNode>()) {
@@ -167,6 +171,8 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
void Dispatch_(const IntImmNode* op, ffi::reflection::AccessPath path)
override;
void Dispatch_(const FloatImmNode* op, ffi::reflection::AccessPath path)
override;
void Dispatch_(const StringImmNode* op, ffi::reflection::AccessPath path)
override;
+ void Dispatch_(const DataTypeImmNode* op, ffi::reflection::AccessPath path)
override;
+ void Dispatch_(const GenericConstNode* op, ffi::reflection::AccessPath path)
override;
void Dispatch_(const prim::LetNode* op, ffi::reflection::AccessPath path)
override;
void Dispatch_(const prim::AddNode* op, ffi::reflection::AccessPath path)
override;
void Dispatch_(const prim::SubNode* op, ffi::reflection::AccessPath path)
override;
diff --git a/tests/python/relax/script/test_relax_script_printer.py
b/tests/python/relax/script/test_relax_script_printer.py
index 184969514d..89fbc969c1 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -554,8 +554,8 @@ def
test_primitive_bindings_roundtrip_without_prim_value_marker():
def test_data_type_imm():
- obj = tvm.ir.GenericConst(tvm.DataType("float32"), tvm.relax.AnyType())
- _assert_print_lines(obj, 'R.dtype("float32")')
+ obj = tvm.ir.DataTypeImm(tvm.DataType("float32"))
+ _assert_print_lines(obj, 'T.dtype("float32")')
def test_var():
diff --git a/tests/python/relax/test_ast_printer.py
b/tests/python/relax/test_ast_printer.py
index 97a6960b9b..efea39ea82 100644
--- a/tests/python/relax/test_ast_printer.py
+++ b/tests/python/relax/test_ast_printer.py
@@ -681,11 +681,11 @@ def test_string_imm():
def test_datatype_imm():
- data_type_imm = tvm.ir.GenericConst(tvm.DataType("int32"),
tvm.relax.AnyType())
+ data_type_imm = tvm.ir.DataTypeImm(tvm.DataType("int32"))
data_type_str = strip_whitespace(dump_ast(data_type_imm))
assert data_type_str == strip_whitespace(
"""
- GenericConst(
+ DataTypeImm(
value=int32,
ty=AnyType()
)
diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py
index 93d554424e..4eaa2162ff 100644
--- a/tests/python/relax/test_expr.py
+++ b/tests/python/relax/test_expr.py
@@ -343,8 +343,8 @@ def test_string_imm():
def test_datatype_imm():
- d0 = tvm.ir.GenericConst(tvm.DataType("int32"), tvm.relax.AnyType())
- d1 = tvm.ir.GenericConst(tvm.DataType("int32"), tvm.relax.AnyType())
+ d0 = tvm.ir.DataTypeImm(tvm.DataType("int32"))
+ d1 = tvm.ir.DataTypeImm(tvm.DataType("int32"))
assert d0.value == "int32"
_check_equal(d0, d1)
_check_json_roundtrip(d0)
diff --git a/tests/python/script/test_basic_usage.py
b/tests/python/script/test_basic_usage.py
index a1787e14fa..75bc8208a5 100644
--- a/tests/python/script/test_basic_usage.py
+++ b/tests/python/script/test_basic_usage.py
@@ -26,7 +26,7 @@ import pytest
# ruff: noqa: F841
from minilang import Value
-from tvm import ir
+from tvm import DataType, ir
from tvm.ir import prim
from tvm.ir._overload_prim_expr import EqualOp
from tvm.script import ir as I
@@ -460,6 +460,17 @@ def
test_ir_branch_incoming_read_and_rebinding_obeys_python_scope(language):
return y
+def test_datatype_value_roundtrip(language):
+ # A DataType value needs only shared IR: the minimal language parses its
shared spelling
+ # into a root-typed constant, and the shared printer reproduces that
spelling.
+ spelling = 'I.dtype("float32")'
+ source = f"@M.function\ndef main():\n M.record({spelling})\n"
+ value = entry.parse(source, extra_vars={"M": language.M, "I":
I}).body[0][1]
+ ir.assert_structural_equal(value, ir.DataTypeImm(DataType("float32")))
+ printed = ir.IRModule(attrs={"dtype": value}).script()
+ assert f'I.module_attrs({{"dtype": {spelling}}})' in printed
+
+
def test_mutating_stores_and_loop_control_keep_effect_order(language):
# Stores evaluate RHS before target/index; augmented stores load the
target before their RHS.
M = language.M
