chore: import upstream snapshot with attribution

python/tvm/ir/__init__.py

@@ -0,0 +1,40 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-import
+"""Common data structures across all IR variants."""
+
+from . import instrument, transform
+from .attrs import Attrs, DictAttrs, make_node
+from .base import (
+    EnvFunc,
+    Node,
+    SourceName,
+    Span,
+    SequentialSpan,
+    assert_structural_equal,
+    load_json,
+    save_json,
+)
+from .expr import Call, Expr, GlobalVar, Range, is_prim_expr
+from .function import BaseFunc, CallingConv
+from .global_info import GlobalInfo, DummyGlobalInfo, VDevice
+from .module import IRModule
+from .op import Op, register_intrin_lowering, register_op_attr
+from .type import FuncType, PointerType, PrimType, TupleType, Type
+
+from tvm_ffi import Array, Map

python/tvm/ir/_ffi_api.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI APIs for tvm.ir"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("ir", __name__)

python/tvm/ir/_ffi_instrument_api.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI APIs for tvm.instrument"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("instrument", __name__)

python/tvm/ir/_ffi_transform_api.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI APIs for tvm.transform"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("transform", __name__)

python/tvm/ir/_overload_prim_expr.py

@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Primitive-expression overloads for shared IR expressions."""
+
+
+def __add__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __radd__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __sub__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rsub__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __mul__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rmul__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __div__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rdiv__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __truediv__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rtruediv__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __floordiv__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rfloordiv__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __mod__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rmod__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __neg__(_value):
+    return NotImplemented
+
+
+def __lshift__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rlshift__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rshift__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rrshift__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __and__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rand__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __or__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __ror__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __xor__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rxor__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __invert__(_value):
+    return NotImplemented
+
+
+def __lt__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __le__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __eq__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __ne__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __gt__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __ge__(_lhs, _rhs):
+    return NotImplemented
+
+
+def equal(_lhs, _rhs, _span=None):
+    return NotImplemented
+
+
+def astype(_value, _dtype, _span=None):
+    return NotImplemented

python/tvm/ir/_tensor_expr_overload.py

@@ -0,0 +1,113 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tensor-expression overload hooks for shared IR expressions."""
+
+
+def __add__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __radd__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __sub__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rsub__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __mul__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rmul__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __div__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rdiv__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __truediv__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rtruediv__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __floordiv__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rfloordiv__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __mod__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rmod__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __pow__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __rpow__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __neg__(_value):
+    return NotImplemented
+
+
+def __lt__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __le__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __gt__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __ge__(_lhs, _rhs):
+    return NotImplemented
+
+
+def __call__(_value, *_args, **_kwargs):
+    return NotImplemented
+
+
+def __getitem__(_value, _index):
+    return NotImplemented
+
+
+def astype(_value, _dtype, _span=None):
+    return NotImplemented

python/tvm/ir/attrs.py

@@ -0,0 +1,169 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""TVM Attribute module, which is mainly used for defining attributes of operators."""
+
+import tvm_ffi
+import tvm_ffi._ffi_api as _tvm_ffi_api
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("ir.Attrs")
+class Attrs(Object):
+    """Attribute node, which is mainly use for defining attributes of operators.
+
+    Used by function registered in python side, such as compute, schedule and alter_layout.
+    Attrs is passed as the first argument to these functions.
+    """
+
+    def get_int_tuple(self, key):
+        """Get a python int tuple of a key
+
+        Parameters
+        ----------
+        key: str
+
+        Returns
+        -------
+        value: Tuple of int
+        """
+        return tuple(x if isinstance(x, int) else x.value for x in getattr(self, key))
+
+    def get_int(self, key):
+        """Get a python int value of a key
+
+        Parameters
+        ----------
+        key: str
+
+        Returns
+        -------
+        value: int
+        """
+        return getattr(self, key)
+
+    def get_str(self, key):
+        """Get a python int value of a key
+
+        Parameters
+        ----------
+        key: str
+
+        Returns
+        -------
+        value: int
+        """
+        return getattr(self, key)
+
+    def __getitem__(self, item):
+        return getattr(self, item)
+
+
+@tvm_ffi.register_object("ir.DictAttrs")
+class DictAttrs(Attrs):
+    """Dictionary attributes."""
+
+    @property
+    def __dict__(self):
+        """Return the underlying key-value map as a Python dict.
+
+        Defined explicitly so that tvm_ffi's _add_class_attrs skips registering
+        the C++ reflection field named '__dict__' (Python forbids adding a class
+        attribute named '__dict__' via setattr on extension-type subclasses).
+        """
+        return dict(self._dict())
+
+    def _dict(self):
+        """Get internal dict"""
+        return _ffi_api.DictAttrsGetDict(self)
+
+    def keys(self):
+        """Get list of names in the attribute.
+
+        Returns
+        -------
+        keys : list of str
+            List of keys
+        """
+        return [k for k, _ in self.items()]
+
+    def __getitem__(self, k):
+        return self._dict().__getitem__(k)
+
+    def get(self, key, default=None):
+        """Get an element with a default value."""
+        return self._dict().get(key, default)
+
+    def __contains__(self, k):
+        return self._dict().__contains__(k)
+
+    def __getattr__(self, name):
+        try:
+            return self._dict().__getitem__(name)
+        except KeyError:
+            raise AttributeError(f"DictAttrs has no attribute {name}")
+
+    def items(self):
+        """Get items from the map."""
+        return self._dict().items()
+
+    def __len__(self):
+        return self._dict().__len__()
+
+
+def make_node(type_key, **kwargs):
+    """Make a new IR node by its type key and fields
+
+    Parameters
+    ----------
+    type_key : str
+        The type key of the node.
+
+    **kwargs : dict
+        The fields of the node.
+
+    Returns
+    -------
+    node : Node
+        The corresponding IR Node
+
+    Note
+    ----
+    If the created node is instance of AttrsNode, then
+    the creator function will also run bound checks and
+    default value setup as supported by Attrs.
+
+    Example
+    -------
+    The following code constructs a IntImm object
+
+    .. code-block:: python
+
+       x = tvm.ir.make_node("ir.IntImm", dtype="int32", value=10, span=None)
+       assert isinstance(x, tvm.tirx.IntImm)
+       assert x.value == 10
+    """
+    if type_key == "ir.DictAttrs":
+        # DictAttrs stores kwargs as a key-value dict, not as named fields.
+        # MakeObjectFromPackedArgs would look for a field named "__dict__".
+        return _tvm_ffi_api.MakeObjectFromPackedArgs("ir.DictAttrs", "__dict__", kwargs)
+    args = [type_key]
+    for k, v in kwargs.items():
+        args += [k, v]
+    return _tvm_ffi_api.MakeObjectFromPackedArgs(*args)

python/tvm/ir/base.py

@@ -0,0 +1,234 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Common base structures."""
+
+import tvm_ffi
+from tvm_ffi import get_global_func, register_object
+from tvm_ffi.serialization import from_json_graph_str, to_json_graph_str
+
+from tvm.runtime import Object
+
+from ..libinfo import __version__
+from . import _ffi_api, json_compact
+
+
+class Node(Object):
+    """Base class of all IR Nodes."""
+
+    def __repr__(self) -> str:
+        from tvm.runtime.script_printer import _script
+
+        try:
+            return _script(self, None)
+        except Exception:
+            return super().__repr__()
+
+
+@register_object("ir.SourceMap")
+class SourceMap(Object):
+    def add(self, name, content):
+        return get_global_func("SourceMapAdd")(self, name, content)
+
+
+@register_object("ir.SourceName")
+class SourceName(Object):
+    """A identifier for a source location.
+
+    Parameters
+    ----------
+    name : str
+        The name of the source.
+    """
+
+    def __init__(self, name):
+        self.__init_handle_by_constructor__(_ffi_api.SourceName, name)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("ir.Span")
+class Span(Object):
+    """Specifies a location in a source program.
+
+    Parameters
+    ----------
+    source : SourceName
+        The source name.
+
+    lineno : int
+        The line number.
+
+    col_offset : int
+        The column offset of the location.
+    """
+
+    def __init__(self, source_name, line, end_line, column, end_column):
+        self.__init_handle_by_constructor__(
+            _ffi_api.Span,
+            source_name,
+            line,
+            end_line,
+            column,
+            end_column,  # type: ignore # pylint: disable=no-member
+        )
+
+
+@register_object("ir.SequentialSpan")
+class SequentialSpan(Object):
+    """A sequence of source spans
+
+    This span is specific for an expression, which is from multiple expressions
+    after an IR transform.
+
+    Parameters
+    ----------
+    spans : Array
+        The array of spans.
+    """
+
+    def __init__(self, spans):
+        self.__init_handle_by_constructor__(_ffi_api.SequentialSpan, spans)
+
+
+@register_object("ir.EnvFunc")
+class EnvFunc(Object):
+    """Environment function.
+
+    This is a global function object that can be serialized by its name.
+    """
+
+    def __call__(self, *args):
+        return _ffi_api.EnvFuncCall(self, *args)  # type: ignore # pylint: disable=no-member
+
+    @property
+    def func(self):
+        return _ffi_api.EnvFuncGetFunction(self)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def get(name):
+        """Get a static env function
+
+        Parameters
+        ----------
+        name : str
+            The name of the function.
+        """
+        return _ffi_api.EnvFuncGet(name)  # type: ignore # pylint: disable=no-member
+
+
+def load_json(json_str) -> Object:
+    """Load tvm object from json_str.
+
+    Parameters
+    ----------
+    json_str : str
+        The json string
+
+    Returns
+    -------
+    node : Object
+        The loaded tvm node.
+    """
+
+    json_str = json_compact.upgrade_json(json_str)
+    return from_json_graph_str(json_str)
+
+
+def save_json(node) -> str:
+    """Save tvm object as json string.
+
+    Parameters
+    ----------
+    node : Object
+        A TVM object to be saved.
+
+    Returns
+    -------
+    json_str : str
+        Saved json string.
+    """
+    return to_json_graph_str(node, {"tvm_version": __version__})
+
+
+def assert_structural_equal(lhs, rhs, map_free_vars=False):
+    """Assert lhs and rhs are structurally equal to each other.
+
+    Parameters
+    ----------
+    lhs : Object
+        The left operand.
+
+    rhs : Object
+        The left operand.
+
+    map_free_vars : bool
+        Whether or not shall we map free vars that does
+        not bound to any definitions as equal to each other.
+
+    Raises
+    ------
+    ValueError : if assertion does not hold.
+
+    See Also
+    --------
+    tvm_ffi.structural_equal
+    """
+    first_mismatch = tvm_ffi.get_first_structural_mismatch(lhs, rhs, map_free_vars)
+    if first_mismatch is not None:
+        from tvm.runtime.script_printer import (  # pylint: disable=import-outside-toplevel
+            PrinterConfig,
+            _script,
+        )
+
+        lhs_path, rhs_path = first_mismatch
+        lhs_script = _script(lhs, PrinterConfig(syntax_sugar=False, path_to_underline=[lhs_path]))
+        rhs_script = _script(rhs, PrinterConfig(syntax_sugar=False, path_to_underline=[rhs_path]))
+        raise ValueError(
+            f"StructuralEqual check failed, caused by lhs at {lhs_path}:\n"
+            f"{lhs_script}\n"
+            f"and rhs at {rhs_path}:\n"
+            f"{rhs_script}"
+        )
+
+
+def deprecated(
+    method_name: str,
+    new_method_name: str,
+):
+    """A decorator to indicate that a method is deprecated
+
+    Parameters
+    ----------
+    method_name : str
+        The name of the method to deprecate
+    new_method_name : str
+        The name of the new method to use instead
+    """
+    import functools  # pylint: disable=import-outside-toplevel
+    import warnings  # pylint: disable=import-outside-toplevel
+
+    def _deprecate(func):
+        @functools.wraps(func)
+        def _wrapper(*args, **kwargs):
+            warnings.warn(
+                f"{method_name} is deprecated, use {new_method_name} instead",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return func(*args, **kwargs)
+
+        return _wrapper
+
+    return _deprecate

python/tvm/ir/expr.py

@@ -0,0 +1,421 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Common expressions data structures in the IR."""
+
+from numbers import Number
+
+import tvm_ffi
+
+import tvm
+
+from ..runtime import Object, Scriptable
+from . import _ffi_api, _overload_prim_expr, _tensor_expr_overload
+from .base import Node, Span
+
+
+@tvm_ffi.register_object("ir.Expr")
+class Expr(Node):
+    """Base class of all the expressions."""
+
+    span: Span | None
+    ty: "tvm.ir.Type"
+
+
+def is_prim_expr(value: object) -> bool:
+    """Return whether an expression has a primitive result type."""
+    return isinstance(value, Expr) and isinstance(value.ty, tvm.ir.PrimType)
+
+
+@tvm_ffi.register_object("ir.GlobalVar")
+class GlobalVar(Expr):
+    """A global variable in the IR.
+
+    GlobalVar is used to refer to the global functions
+    stored in the IRModule.
+
+    Parameters
+    ----------
+    name_hint: str
+        The name of the variable.
+    """
+
+    name_hint: str
+
+    def __init__(self, name_hint: str):
+        self.__init_handle_by_constructor__(_ffi_api.GlobalVar, name_hint)
+
+    def __call__(self, *args: Expr) -> Expr:
+        """Call the global variable.
+
+        Parameters
+        ----------
+        args: List[Expr]
+            The arguments to the call.
+
+        Returns
+        -------
+        call: Expr
+            A call taking the variable as a function.
+        """
+        from .type import PointerType
+
+        def is_tir_arg(x):
+            return (
+                isinstance(x, Number)
+                or is_prim_expr(x)
+                or (isinstance(x, Expr) and isinstance(x.ty, PointerType))
+            )
+
+        if args and all(is_tir_arg(x) for x in args):
+            return tvm.tirx.call_tir(self, *args)
+
+        if all(isinstance(x, Expr) for x in args):
+            return Call(self, args)
+
+        arg_types = [type(x) for x in args]
+        raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for types {arg_types}")
+
+
+@tvm_ffi.register_object("ir.Call")
+class Call(Expr, Scriptable):
+    """Core function call node."""
+
+    __hash__ = Expr.__hash__
+
+    op: Expr
+    args: list[Expr]
+    attrs: "tvm.ir.Attrs | None"
+    ty_args: list["tvm.ir.Type"]
+    span: Span | None
+
+    def __init__(
+        self,
+        op: Expr | str,
+        args: list[Expr] | tuple[Expr, ...],
+        attrs: "tvm.ir.Attrs | dict | None" = None,
+        ty_args: list["tvm.ir.Type"] | tuple["tvm.ir.Type", ...] | None = None,
+        span: Span | None = None,
+        ret_ty: "tvm.ir.Type | str | None" = None,
+    ) -> None:
+        # pylint: disable=import-outside-toplevel
+        from .attrs import DictAttrs
+        from .op import Op
+        from .type import PointerType, PrimType, Type
+
+        if isinstance(op, str):
+            op = Op.get(op)
+        if attrs is not None and isinstance(attrs, dict):
+            attrs = DictAttrs(attrs)
+        if ret_ty is None:
+            ret_ty = Type.missing()
+        if isinstance(ret_ty, str) and ret_ty == "handle":
+            ret_ty = PointerType(PrimType("void"))
+        elif ret_ty is not None and not isinstance(ret_ty, Type):
+            ret_ty = PrimType(ret_ty)
+        if ty_args is None:
+            ty_args = []
+        self.__init_handle_by_constructor__(_ffi_api.Call, ret_ty, op, args, attrs, ty_args, span)
+
+    def expr_ty(self):
+        """Return this expression's primitive result type."""
+        if is_prim_expr(self):
+            return self.ty
+        raise TypeError(f"Expected primitive-valued Call, but result type is {self.ty}")
+
+    def __add__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__add__(self, other)
+        return _tensor_expr_overload.__add__(self, other)
+
+    def __radd__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__radd__(self, other)
+        return _tensor_expr_overload.__radd__(self, other)
+
+    def __sub__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__sub__(self, other)
+        return _tensor_expr_overload.__sub__(self, other)
+
+    def __rsub__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rsub__(self, other)
+        return _tensor_expr_overload.__rsub__(self, other)
+
+    def __mul__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__mul__(self, other)
+        return _tensor_expr_overload.__mul__(self, other)
+
+    def __rmul__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rmul__(self, other)
+        return _tensor_expr_overload.__rmul__(self, other)
+
+    def __div__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__div__(self, other)
+        return _tensor_expr_overload.__div__(self, other)
+
+    def __rdiv__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rdiv__(self, other)
+        return _tensor_expr_overload.__rdiv__(self, other)
+
+    def __truediv__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__truediv__(self, other)
+        return _tensor_expr_overload.__truediv__(self, other)
+
+    def __rtruediv__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rtruediv__(self, other)
+        return _tensor_expr_overload.__rtruediv__(self, other)
+
+    def __floordiv__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__floordiv__(self, other)
+        return _tensor_expr_overload.__floordiv__(self, other)
+
+    def __rfloordiv__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rfloordiv__(self, other)
+        return _tensor_expr_overload.__rfloordiv__(self, other)
+
+    def __mod__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__mod__(self, other)
+        return _tensor_expr_overload.__mod__(self, other)
+
+    def __rmod__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rmod__(self, other)
+        return _tensor_expr_overload.__rmod__(self, other)
+
+    def __pow__(self, other):
+        if is_prim_expr(self):
+            return NotImplemented
+        return _tensor_expr_overload.__pow__(self, other)
+
+    def __rpow__(self, other):
+        if is_prim_expr(self):
+            return NotImplemented
+        return _tensor_expr_overload.__rpow__(self, other)
+
+    def __neg__(self):
+        if is_prim_expr(self):
+            result = _overload_prim_expr.__neg__(self)
+            if result is NotImplemented:
+                raise TypeError("Primitive expression overload __neg__ is not registered")
+            return result
+        result = _tensor_expr_overload.__neg__(self)
+        if result is NotImplemented:
+            raise TypeError("Tensor expression overload negative is not registered")
+        return result
+
+    def __lshift__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__lshift__(self, other)
+        return NotImplemented
+
+    def __rlshift__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rlshift__(self, other)
+        return NotImplemented
+
+    def __rshift__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rshift__(self, other)
+        return NotImplemented
+
+    def __rrshift__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rrshift__(self, other)
+        return NotImplemented
+
+    def __and__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__and__(self, other)
+        return NotImplemented
+
+    def __rand__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rand__(self, other)
+        return NotImplemented
+
+    def __or__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__or__(self, other)
+        return NotImplemented
+
+    def __ror__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__ror__(self, other)
+        return NotImplemented
+
+    def __xor__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__xor__(self, other)
+        return NotImplemented
+
+    def __rxor__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__rxor__(self, other)
+        return NotImplemented
+
+    def __invert__(self):
+        if is_prim_expr(self):
+            result = _overload_prim_expr.__invert__(self)
+            if result is NotImplemented:
+                raise TypeError("Primitive expression overload __invert__ is not registered")
+            return result
+        return NotImplemented
+
+    def __lt__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__lt__(self, other)
+        return _tensor_expr_overload.__lt__(self, other)
+
+    def __le__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__le__(self, other)
+        return _tensor_expr_overload.__le__(self, other)
+
+    def __eq__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__eq__(self, other)
+        return Object.__eq__(self, other)
+
+    def __ne__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__ne__(self, other)
+        return Object.__ne__(self, other)
+
+    def __gt__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__gt__(self, other)
+        return _tensor_expr_overload.__gt__(self, other)
+
+    def __ge__(self, other):
+        if is_prim_expr(self):
+            return _overload_prim_expr.__ge__(self, other)
+        return _tensor_expr_overload.__ge__(self, other)
+
+    def __nonzero__(self):
+        raise ValueError(
+            "Cannot use and / or / not operator to Expr, hint: use tvm.tirx.all / "
+            "tvm.tirx.any, if it is None checking, use node is not None"
+        )
+
+    def __bool__(self):
+        return self.__nonzero__()
+
+    def equal(self, other, span=None):
+        result = _overload_prim_expr.equal(self, other, span)
+        if result is NotImplemented:
+            raise TypeError("Primitive expression overload equal is not registered")
+        return result
+
+    def astype(self, dtype, span=None):
+        if is_prim_expr(self):
+            result = _overload_prim_expr.astype(self, dtype, span)
+            if result is NotImplemented:
+                raise TypeError("Primitive expression overload astype is not registered")
+            return result
+        result = _tensor_expr_overload.astype(self, dtype, span)
+        if result is NotImplemented:
+            raise TypeError("Tensor expression overload astype is not registered")
+        return result
+
+    def __call__(self, *args, attrs=None):
+        if is_prim_expr(self):
+            raise TypeError("A primitive-valued Call cannot be called")
+        result = _tensor_expr_overload.__call__(self, *args, attrs=attrs)
+        if result is NotImplemented:
+            raise TypeError("Tensor expression overload __call__ is not registered")
+        return result
+
+    def __getitem__(self, index):
+        if is_prim_expr(self):
+            raise TypeError("A primitive-valued Call cannot be indexed")
+        result = _tensor_expr_overload.__getitem__(self, index)
+        if result is NotImplemented:
+            raise TypeError("Tensor expression overload __getitem__ is not registered")
+        return result
+
+
+@tvm_ffi.register_object("ir.Range")
+class Range(Node, Scriptable):
+    """Represent a range in TVM.
+
+    You do not need to create a Range explicitly.
+    Python lists and tuples will be converted automatically to a Range in API functions.
+
+    Parameters
+    ----------
+    begin : Expr
+        The begin value of the range when end is None.
+        Otherwise it is the length of the range.
+
+    end : Optional[Expr]
+        The end value of the range.
+
+    span : Optional[Span]
+        The location of this node in the source code.
+
+    Note
+    ----
+    The constructor creates the range `[begin, end)`
+    if the end argument is not None. Otherwise, it creates `[0, begin)`.
+    """
+
+    min: Expr
+    extent: Expr
+    span: Span | None
+
+    def __init__(self, begin: Expr, end: Expr | None = None, span: Span | None = None) -> None:
+        self.__init_handle_by_constructor__(_ffi_api.Range, begin, end, span)
+
+    @staticmethod
+    def from_min_extent(min_value: Expr, extent: Expr, span: Span | None = None) -> "Range":
+        """Construct a Range by min and extent.
+
+        This constructs a range in [min_value, min_value + extent)
+
+        Parameters
+        ----------
+        min_value : Expr
+            The minimum value of the range.
+
+        extent : Expr
+            The extent of the range.
+
+        span : Optional[Span]
+            The location of this node in the source code.
+
+        Returns
+        -------
+        rng : Range
+            The constructed range.
+        """
+        return _ffi_api.Range_from_min_extent(min_value, extent, span)
+
+    def __eq__(self, other: Object) -> bool:
+        return tvm_ffi.structural_equal(self, other)
+
+    def __ne__(self, other: Object) -> bool:
+        return not self.__eq__(other)

python/tvm/ir/function.py

@@ -0,0 +1,108 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""Function definitions."""
+
+from enum import IntEnum
+
+import tvm_ffi
+
+import tvm.runtime
+from tvm.runtime import Object
+
+from . import _ffi_api
+from .attrs import DictAttrs
+from .expr import Expr
+
+
+class CallingConv(IntEnum):
+    """Possible kinds of calling conventions."""
+
+    DEFAULT = 0
+    C_PACKED_FUNC = 1
+    DEVICE_KERNEL_LAUNCH = 2
+
+
+@tvm_ffi.register_object("ir.BaseFunc")
+class BaseFunc(Expr):
+    """Base class of all functions."""
+
+    @property
+    def attrs(self):
+        """Return the attrs member of the function."""
+        return _ffi_api.BaseFunc_Attrs(self)
+
+    def with_attr(self, attr_key_or_dict, attr_value=None) -> "BaseFunc":
+        """Create a new copy of the function and update the attribute.
+
+        Parameters
+        ----------
+        attr_key_or_dict : Union[str, dict]
+            The attribute key to use or a dict containing multiple key value pairs.
+
+        attr_value : Object
+            The new attribute value.
+
+        Returns
+        -------
+        func : BaseFunc
+            A new copy of the function
+        """
+        # make sure we first copy so that we can safely do copy on write
+        # for multiple updates.
+        res = _ffi_api.BaseFuncCopy(self)
+
+        if isinstance(attr_key_or_dict, dict):
+            for key, val in attr_key_or_dict.items():
+                res = _ffi_api.BaseFuncWithAttr(res._move(), key, tvm.runtime.convert(val))
+            return res
+
+        return _ffi_api.BaseFuncWithAttr(
+            res._move(), attr_key_or_dict, tvm.runtime.convert(attr_value)
+        )
+
+    def with_attrs(self, attr_map: DictAttrs | dict[str, Object]) -> "BaseFunc":
+        """Copy the IRModule and add the given attribute map to it.
+        Parameters
+        ----------
+        attr_map: Union[DictAttrs, Dict[str, Object]]
+            The attribute map
+        Returns
+        -------
+        func : BaseFunc
+            A new copy of the function
+        """
+        if isinstance(attr_map, tvm.ir.DictAttrs):
+            attr_map = attr_map._dict()
+
+        return _ffi_api.BaseFuncWithAttrs(self, attr_map)
+
+    def without_attr(self, attr_key: str) -> "BaseFunc":
+        """Create a new copy of the function with an attribute without provided key.
+
+        Parameters
+        ----------
+        attr_key : str
+            The attribute key to delete from the attrubte pairs.
+
+
+        Returns
+        -------
+        func : BaseFunc
+            A new copy of the function
+        """
+        return _ffi_api.BaseFuncWithoutAttr(self, attr_key)

python/tvm/ir/global_info.py

@@ -0,0 +1,67 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Global Info."""
+
+import tvm_ffi
+
+import tvm
+from tvm.runtime import Device, Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("ir.GlobalInfo")
+class GlobalInfo(Object):
+    """Base node for all global info that can appear in the IR"""
+
+    def __eq__(self, other):
+        """Compare two global info objects for structural equivalence."""
+        return tvm_ffi.structural_equal(self, other)
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def same_as(self, other):
+        """Overload with structural equality."""
+        return super().__eq__(other)
+
+
+@tvm_ffi.register_object("ir.DummyGlobalInfo")
+class DummyGlobalInfo(GlobalInfo):
+    """DummyGlobalInfo"""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.DummyGlobalInfo,
+        )
+
+
+@tvm_ffi.register_object("ir.VDevice")
+class VDevice(GlobalInfo):
+    """VDevice"""
+
+    def __init__(
+        self,
+        target=None,
+        vdevice_id: int = 0,
+        memory_scope: str = "global",
+    ) -> None:
+        if isinstance(target, dict | str):
+            target = tvm.target.Target(tvm.runtime.convert(target))
+        if isinstance(target, Device):
+            target = tvm.target.Target.from_device(target)
+        self.__init_handle_by_constructor__(_ffi_api.VDevice, target, vdevice_id, memory_scope)

python/tvm/ir/instrument.py

@@ -0,0 +1,341 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name,unused-argument
+"""Common pass instrumentation across IR variants."""
+
+import functools
+import inspect
+import re
+import shutil
+from pathlib import Path
+
+import tvm_ffi
+
+import tvm.runtime
+
+from . import _ffi_instrument_api
+
+
+@tvm_ffi.register_object("instrument.PassInstrument")
+class PassInstrument(tvm.runtime.Object):
+    """A pass instrument implementation.
+
+    To use, a user class can either subclass from PassInstrument
+    directly, or can apply the :py:func:`pass_instrument` wrapper.  In
+    either case, the `enter_pass_ctx`, `exit_pass_ctx`, `should_run`,
+    `run_before_pass`, and `run_after_pass` methods can be defined to
+    adjust the instrument's behavior.  See the no-op implementations
+    in this class definition for more information on each.
+
+    """
+
+    def __init__(self):
+        # initialize handle in case pi_cls creation failed.
+        cls = type(self)
+
+        # If the child class declared the method, then use it.
+        # Otherwise, pass None to avoid a C++ -> Python round trip for
+        # a no-op.
+        def get_child_method(name):
+            if getattr(cls, name) is getattr(PassInstrument, name):
+                return None
+
+            return getattr(self, name)
+
+        # Create runtime pass instrument object.
+        # register instance's enter_pass_ctx,exit_pass_ctx, should_run, run_before_pass and
+        # run_after_pass methods to it if present.
+        self.__init_handle_by_constructor__(
+            _ffi_instrument_api.PassInstrument,
+            cls.__name__,
+            get_child_method("enter_pass_ctx"),
+            get_child_method("exit_pass_ctx"),
+            get_child_method("should_run"),
+            get_child_method("run_before_pass"),
+            get_child_method("run_after_pass"),
+        )
+
+    def enter_pass_ctx(self):
+        """Called when entering the instrumented context.
+
+        Returns
+        -------
+        None
+        """
+
+    def exit_pass_ctx(self):
+        """Called when exiting the instrumented context.
+
+        Returns
+        -------
+        None
+        """
+
+    def should_run(self, mod, info):
+        """Determine whether to run the pass or not.
+
+        Called once for each pass that is run while the instrumented
+        context is active.
+
+        Parameters
+        ----------
+        mod : tvm.ir.module.IRModule
+
+            The module on which an optimization pass is being run.
+
+        info : tvm.transform.PassInfo
+
+            The pass information.
+
+        Returns
+        -------
+        should_run : bool
+
+            True to run the pass, or False to skip the pass.
+        """
+
+    def run_before_pass(self, mod, info):
+        """Instrument before the pass runs.
+
+        Called once for each pass that is run while the instrumented
+        context is active.
+
+        Parameters
+        ----------
+        mod : tvm.ir.module.IRModule
+
+            The module on which an optimization pass is being run.
+
+        info : tvm.transform.PassInfo
+
+            The pass information.
+
+        Returns
+        -------
+        None
+        """
+
+    def run_after_pass(self, mod, info):
+        """Instrument after the pass runs.
+
+        Called once for each pass that is run while the instrumented
+        context is active.
+
+        Parameters
+        ----------
+        mod : tvm.ir.module.IRModule
+
+            The module on which an optimization pass is being run.
+
+        info : tvm.transform.PassInfo
+
+            The pass information.
+
+        Returns
+        -------
+        None
+        """
+
+
+def _wrap_class_pass_instrument(pi_cls):
+    """Wrap a python class as pass instrument"""
+
+    # No additional wrapping needed if the user class already
+    # inherits.
+    if issubclass(pi_cls, PassInstrument):
+        return pi_cls
+
+    class PyPassInstrument(pi_cls, PassInstrument):
+        """Internal wrapper class to create a class instance."""
+
+        def __init__(self, *args, **kwargs):
+            # initialize handle in case pi_cls creation failed.
+            pi_cls.__init__(self, *args, **kwargs)
+            PassInstrument.__init__(self)
+
+    functools.update_wrapper(PyPassInstrument.__init__, pi_cls.__init__)
+    PyPassInstrument.__name__ = pi_cls.__name__
+    PyPassInstrument.__doc__ = pi_cls.__doc__
+    PyPassInstrument.__module__ = pi_cls.__module__
+    return PyPassInstrument
+
+
+def pass_instrument(pi_cls=None):
+    """Decorate a pass instrument.
+
+    Parameters
+    ----------
+    pi_class : class
+        Instrument class. See example below.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        @tvm.instrument.pass_instrument
+        class SkipPass:
+            def __init__(self, skip_pass_name):
+                self.skip_pass_name = skip_pass_name
+
+            # Uncomment to customize
+            # def enter_pass_ctx(self):
+            #    pass
+
+            # Uncomment to customize
+            # def exit_pass_ctx(self):
+            #    pass
+
+            # If pass name contains keyword, skip it by return False. (return True: not skip)
+            def should_run(self, mod, pass_info)
+                if self.skip_pass_name in pass_info.name:
+                    return False
+                return True
+
+            # Uncomment to customize
+            # def run_before_pass(self, mod, pass_info):
+            #    pass
+
+            # Uncomment to customize
+            # def run_after_pass(self, mod, pass_info):
+            #    pass
+
+        skip_annotate = SkipPass("AnnotateSpans")
+        with tvm.transform.PassContext(instruments=[skip_annotate]):
+            tvm.compile(mod, "llvm")
+    """
+
+    def create_pass_instrument(pi_cls):
+        if not inspect.isclass(pi_cls):
+            raise TypeError("pi_cls must be a class")
+
+        return _wrap_class_pass_instrument(pi_cls)
+
+    if pi_cls:
+        return create_pass_instrument(pi_cls)
+    return create_pass_instrument
+
+
+@tvm_ffi.register_object("instrument.PassInstrument")
+class PassTimingInstrument(tvm.runtime.Object):
+    """A wrapper to create a passes time instrument that implemented in C++"""
+
+    def __init__(self):
+        self.__init_handle_by_constructor__(_ffi_instrument_api.MakePassTimingInstrument)
+
+    @staticmethod
+    def render():
+        """Retrieve rendered time profile result
+        Returns
+        -------
+        string : string
+            The rendered string result of time profiles
+
+        Examples
+        --------
+
+        .. code-block:: python
+
+            timing_inst = PassTimingInstrument()
+            with tvm.transform.PassContext(instruments=[timing_inst]):
+                relax_mod = relax.transform.FuseOps()(relax_mod)
+                # before exiting the context, get profile results.
+                profiles = timing_inst.render()
+        """
+        return _ffi_instrument_api.RenderTimePassProfiles()
+
+
+@pass_instrument
+class PassPrintingInstrument:
+    """A pass instrument to print if before or
+    print ir after each element of a named pass."""
+
+    def __init__(self, print_before_pass_names, print_after_pass_names):
+        self.print_before_pass_names = print_before_pass_names
+        self.print_after_pass_names = print_after_pass_names
+
+    def run_before_pass(self, mod, pass_info):
+        if pass_info.name in self.print_before_pass_names:
+            print(f"Print IR before: {pass_info.name}\n{mod}\n\n")
+
+    def run_after_pass(self, mod, pass_info):
+        if pass_info.name in self.print_after_pass_names:
+            print(f"Print IR after: {pass_info.name}\n{mod}\n\n")
+
+
+@pass_instrument
+class PrintAfterAll:
+    """Print the name of the pass, the IR, only after passes execute."""
+
+    def run_after_pass(self, mod, info):
+        print(f"After Running Pass: {info}")
+        print(mod)
+
+
+@pass_instrument
+class PrintBeforeAll:
+    """Print the name of the pass, the IR, only before passes execute."""
+
+    def run_before_pass(self, mod, info):
+        print(f"Before Running Pass: {info}")
+        print(mod)
+
+
+@pass_instrument
+class DumpIR:
+    """Dump the IR after the pass runs."""
+
+    def __init__(self, dump_dir: Path | str, refresh: bool = False):
+        if isinstance(dump_dir, Path):
+            self.dump_dir = dump_dir
+        else:
+            self.dump_dir = Path(dump_dir)
+        self.counter = 0
+        if refresh and self.dump_dir.is_dir():
+            self._safe_remove_dump_dir()
+
+    def _safe_remove_dump_dir(self):
+        """Remove dump directory only if it contains only dumped IR files."""
+        # Pattern for dumped files: {counter:03d}_{pass_name}.py
+        dump_pattern = re.compile(r"^\d{3}_.*\.py$")
+
+        # Check all files in the directory
+        for item in self.dump_dir.iterdir():
+            # If there's a subdirectory or a file that doesn't match the pattern, abort
+            if item.is_dir() or not dump_pattern.match(item.name):
+                print(
+                    f"WARNING: Skipping removal of {self.dump_dir} as it contains "
+                    f"non-dumped files or directories. Please clean it manually."
+                )
+                return
+
+        # Safe to remove - only contains dumped files
+        try:
+            shutil.rmtree(self.dump_dir)
+        except OSError as e:
+            print(f"WARNING: Failed to remove directory {self.dump_dir}: {e}")
+
+    def run_after_pass(self, mod, info):
+        self.dump_dir.mkdir(parents=True, exist_ok=True)
+        try:
+            sanitized_pass_name = re.sub(r'[<>:"/\\|?*]', "_", info.name)
+            with open(self.dump_dir / f"{self.counter:03d}_{sanitized_pass_name}.py", "w") as f:
+                f.write(mod.script())
+        except Exception:  # pylint: disable=broad-exception-caught
+            print(f"WARNING: Failed to dump IR for pass {info.name}")
+        finally:
+            self.counter += 1

python/tvm/ir/json_compact.py

@@ -0,0 +1,87 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tool to upgrade json from historical versions."""
+
+import json
+
+
+def get_version(jgraph):
+    """
+    Get the tvm version from the json graph.
+
+    Parameters
+    ----------
+    jgraph : dict
+        The json graph.
+    """
+    return jgraph["metadata"]["tvm_version"]
+
+
+def create_updater(node_map, from_ver, to_ver):
+    """Create an updater to update json loaded data.
+
+    Parameters
+    ----------
+    node_map : Map[str, Function]
+        Map from type_key to updating function
+
+    from_ver : str
+        Prefix of version that we can accept,
+
+    to_ver : str
+        The target version.
+
+    Returns
+    -------
+    fupdater : function
+        The updater function
+    """
+
+    def _updater(data):
+        assert get_version(data).startswith(from_ver)
+        nodes = data["nodes"]
+        for idx, item in enumerate(nodes):
+            f = node_map.get(item["type"], None)
+            if isinstance(f, list):
+                for fpass in f:
+                    item = fpass(item, nodes)
+            elif f:
+                item = f(item, nodes)
+            nodes[idx] = item
+        data["metadata"]["tvm_version"] = to_ver
+        return data
+
+    return _updater
+
+
+def upgrade_json(json_str):
+    """Update json from a historical version.
+
+    Parameters
+    ----------
+    json_str : str
+        A historical json file.
+
+    Returns
+    -------
+    updated_json : str
+        The updated version.
+    """
+    data = json.loads(json_str)
+    if "metadata" not in data and "attrs" in data:
+        raise ValueError("Legacy json graph format detected, we don't support it anymore.")
+    return json.dumps(data, indent=2)

python/tvm/ir/module.py

@@ -0,0 +1,284 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""IRModule that holds the functions and type definitions."""
+
+import tvm_ffi
+
+import tvm
+from tvm.runtime import Object, Scriptable
+
+from . import _ffi_api
+from . import expr as _expr
+from .attrs import DictAttrs
+from .base import Node
+from .function import BaseFunc
+
+
+@tvm_ffi.register_object("ir.IRModule")
+class IRModule(Node, Scriptable):
+    """IRModule that holds functions and type definitions.
+
+    IRModule is the basic unit for all IR transformations across the stack.
+
+    Parameters
+    ----------
+    functions: Optional[dict].
+        Map of global var to BaseFunc
+    """
+
+    def __init__(self, functions=None, attrs=None, global_infos=None):
+        if functions is None:
+            functions = {}
+        elif isinstance(functions, dict):
+            mapped_funcs = {}
+            for k, v in functions.items():
+                if isinstance(k, str):
+                    k = _expr.GlobalVar(k)
+                if not isinstance(k, _expr.GlobalVar):
+                    raise TypeError("Expect functions to be Dict[GlobalVar, Function]")
+                mapped_funcs[k] = v
+            functions = mapped_funcs
+
+        attrs = None if not attrs else attrs
+        if attrs is not None:
+            attrs = tvm.ir.make_node("ir.DictAttrs", **attrs)
+        if global_infos is None:
+            global_infos = {}
+        self.__init_handle_by_constructor__(
+            _ffi_api.IRModule,
+            functions,
+            attrs,
+            global_infos,
+        )
+        self.pyfuncs = {}
+
+    def clone(self) -> "IRModule":
+        return _ffi_api.Module_Clone(self)
+
+    def functions_items(self):
+        """Get items in self.functions.items() in alphabetical order.
+
+        Returns
+        -------
+        items: List[Tuple[GlobalVar, Function]]
+            The functions items.
+        """
+        items = list(self.functions.items())
+        items.sort(key=lambda item: str(item[0].name_hint))
+        return items
+
+    def __setitem__(self, var, val):
+        """Add a mapping to the module.
+
+        Parameters
+        ---------
+        var: GlobalVar
+            The global variable.
+
+        val: Union[Function, Type]
+            The value.
+        """
+        return self._add(var, val, True)
+
+    def _add(self, var, val, update=True):
+        if isinstance(val, BaseFunc):
+            if isinstance(var, str):
+                if _ffi_api.Module_ContainGlobalVar(self, var):
+                    var = _ffi_api.Module_GetGlobalVar(self, var)
+                else:
+                    var = _expr.GlobalVar(var)
+            _ffi_api.Module_Add(self, var, val, update)
+
+    def __getitem__(self, var):
+        """Lookup a global definition by name or by variable.
+
+        Parameters
+        ----------
+        var: Union[String, GlobalVar, GlobalTypeVar]
+            The name or global variable.
+
+        Returns
+        -------
+        val: Union[Function, Type]
+            The definition referenced by :code:`var` (either a function or type).
+        """
+        if isinstance(var, str):
+            return _ffi_api.Module_Lookup_str(self, var)
+        assert isinstance(var, _expr.GlobalVar)
+        return _ffi_api.Module_Lookup(self, var)
+
+    def __delitem__(self, var: str | _expr.GlobalVar):
+        _ffi_api.Module_Remove(self, var)
+
+    def __contains__(self, var: str | _expr.GlobalVar) -> bool:
+        return _ffi_api.Module_Contains(self, var)
+
+    def update(self, other):
+        """Insert functions in another Module to current one.
+
+        Parameters
+        ----------
+        other: IRModule
+            The module to merge into the current Module.
+        """
+        if isinstance(other, dict):
+            other = IRModule(other)
+
+        return _ffi_api.Module_Update(self, other)
+
+    def update_func(self, var, func):
+        """Update the function corresponding to a global variable in the
+        module.
+
+        Parameters
+        ----------
+        var: GlobalVar
+            The global variable.
+
+        func: tvm.ir.BaseFunc
+            The function to be inserted.
+        """
+        return _ffi_api.Module_UpdateFunction(self, var, func)
+
+    def update_global_info(self, name, global_info):
+        """Update global info in the module
+
+        Parameters
+        ----------
+        name: str
+            The name for the global info.
+
+        global_info: List[GlobalInfo]
+            The global info to be updated.
+        """
+        return _ffi_api.Module_UpdateGlobalInfo(self, name, global_info)
+
+    def get_global_var(self, name):
+        """Get a global variable in the function by name.
+
+        Parameters
+        ----------
+        name: str
+            The name of the global variable.
+
+        Returns
+        -------
+        global_var: GlobalVar
+            The global variable mapped to :code:`name`.
+
+        Raises
+        ------
+        RuntimeError if we cannot find corresponding global var.
+        """
+        return _ffi_api.Module_GetGlobalVar(self, name)
+
+    def get_global_vars(self):
+        """Collect all global vars defined in this module.
+
+        Returns
+        -------
+        global_vars: Array[GlobalVar]
+            An array of global vars.
+        """
+        return _ffi_api.Module_GetGlobalVars(self)
+
+    @staticmethod
+    def from_expr(expr, functions=None):
+        """Construct a module from a standalone expression.
+
+        Parameters
+        ----------
+        expr: Expr
+            The starting expression
+
+        global_funcs: Optional[dict]
+            Map of global vars to function definitions
+
+        Returns
+        -------
+        mod: Module
+            A module containing the passed definitions,
+            where expr is set as the entry point
+            (wrapped in a function if necessary)
+        """
+        funcs = functions if functions is not None else {}
+        return _ffi_api.Module_FromExpr(expr, funcs)
+
+    def get_attr(self, attr_key):
+        """Get the IRModule attribute.
+
+        Parameters
+        ----------
+        attr_key : str
+            The attribute key.
+
+        Returns
+        -------
+        attr_value : Any
+            Attribute value
+        """
+
+        return _ffi_api.Module_GetAttr(self, attr_key)
+
+    def with_attr(self, attr_key, attr_value):
+        """Copy the IRModule and add an attribute to it.
+
+        Parameters
+        ----------
+        attr_key : str
+            The attribute key.
+
+        attr_value : Object
+            The new attribute value.
+
+        Returns
+        -------
+        mod : IRModule
+            A new copy of the IRModule with the attribute
+        """
+
+        return _ffi_api.Module_WithAttr(self, attr_key, attr_value)
+
+    def without_attr(self, attr_key: str) -> "IRModule":
+        """Copy the IRModule and remove an attribute key and its associated value.
+        Parameters
+        ----------
+        attr_key : str
+            The attribute key.
+        Returns
+        -------
+        mod : IRModule
+            A new copy of the IRModule without the attribute
+        """
+
+        return _ffi_api.Module_WithoutAttr(self, attr_key)
+
+    def with_attrs(self, attr_map: DictAttrs | dict[str, Object]) -> "IRModule":
+        """Copy the IRModule and add the given attribute map to it.
+        Parameters
+        ----------
+        attr_map: Union[DictAttrs, Dict[str, Object]]
+            The attribute map
+        Returns
+        -------
+        mod : IRModule
+            A new copy of the IRModule with the attribute
+        """
+        if isinstance(attr_map, tvm.ir.DictAttrs):
+            attr_map = attr_map._dict()
+
+        return _ffi_api.Module_WithAttrs(self, attr_map)

python/tvm/ir/op.py

@@ -0,0 +1,226 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""Primitive operators in the TVM IR."""
+
+import tvm_ffi
+
+from . import _ffi_api
+from .expr import Expr
+
+
+@tvm_ffi.register_object("ir.Op")
+class Op(Expr):
+    """Primitive operator in the IR."""
+
+    def __init__(self):
+        raise RuntimeError("Cannot create op, use get instead")
+
+    @staticmethod
+    def get(op_name):
+        """Get the Op for a given name
+
+        Parameters
+        ----------
+        op_name : str
+            The operator name
+
+        Returns
+        -------
+        op : Op
+            The op of the corresponding name
+        """
+        return _ffi_api.GetOp(op_name)
+
+    def get_attr(self, attr_name):
+        """Get additional attribute about the operator.
+
+        Parameters
+        ----------
+        attr_name : str
+            The attribute name.
+
+        Returns
+        -------
+        value : object
+            The attribute value
+        """
+        return _ffi_api.OpGetAttr(self, attr_name)
+
+    def has_attr(self, attr_name):
+        """Check whether the operator has additional attribute.
+
+        Parameters
+        ----------
+        attr_name : str
+            The attribute name.
+
+        Returns
+        -------
+        value : bool
+            Whether the operator has additional attribute
+        """
+        return _ffi_api.OpHasAttr(self, attr_name)
+
+    def set_attr(self, attr_name, value, plevel=10):
+        """Set attribute about the operator.
+
+        Parameters
+        ----------
+        attr_name : str
+            The attribute name
+
+        value : object
+            The attribute value
+
+        plevel : int
+            The priority level
+        """
+        _ffi_api.OpSetAttr(self, attr_name, value, plevel)
+
+    def reset_attr(self, attr_name):
+        """Reset attribute about the operator.
+
+        Parameters
+        ----------
+        attr_name : str
+            The attribute name
+        """
+        _ffi_api.OpResetAttr(self, attr_name)
+
+    def add_argument(self, name, type, description):  # pylint: disable=redefined-builtin
+        """Add arguments information to the function.
+
+        Parameters
+        ----------
+        name : str
+            The argument name.
+        type : str
+            The argument type.
+        description : str
+            The argument description.
+        """
+        _ffi_api.OpAddArgument(self, name, type, description)
+
+    def set_support_level(self, level):
+        """Set the support level of op.
+
+        Parameters
+        ----------
+        level : int
+            The support level.
+        """
+        _ffi_api.OpSetSupportLevel(self, level)
+
+    def set_num_inputs(self, n):
+        """Set the support level of op.
+
+        Parameters
+        ----------
+        n : int
+            The input number.
+        """
+        _ffi_api.OpSetNumInputs(self, n)
+
+    def set_attrs_type_key(self, key):
+        """Set the attribute type key of op.
+
+        Parameters
+        ----------
+        key : str
+            The type key.
+        """
+        _ffi_api.OpSetAttrsTypeKey(self, key)
+
+    @staticmethod
+    def list_op_names():
+        """List all the op names in the op registry.
+
+        Returns
+        -------
+        value : List[str]
+            The registered op names
+        """
+        return _ffi_api.ListOpNames()
+
+
+def register_op_attr(op_name, attr_key, value=None, level=10):
+    """Register an operator property of an operator by name.
+
+    Parameters
+    ----------
+    op_name : str
+        The name of operator
+
+    attr_key : str
+        The attribute name.
+
+    value : object, optional
+        The value to set
+
+    level : int, optional
+        The priority level
+
+    Returns
+    -------
+    fregister : function
+        Register function if value is not specified.
+    """
+
+    def _register(v):
+        """internal register function"""
+        _ffi_api.RegisterOpAttr(op_name, attr_key, v, level)
+        return v
+
+    return _register(value) if value is not None else _register
+
+
+def register_intrin_lowering(
+    op_name,
+    target,
+    *,
+    f=None,
+    level=10,
+):
+    """Register Op lowering function
+
+    Parameters
+    ----------
+    op_name : str
+        The op name
+
+    target : str
+        The target string for given intrinsic lowering function
+
+    f : function, optional
+        The function to be registered.
+
+    level : int
+        The priority level
+
+    Returns
+    -------
+    fregister : function
+        Register op lowering function if f is not specified.
+    """
+
+    def _register(f):
+        """internal register function"""
+        _ffi_api.RegisterOpLowerIntrinsic(op_name, f, target, level)
+        return f
+
+    return _register(f) if f is not None else _register

python/tvm/ir/supply.py

@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Suppliers that are used to guarantee uniqueness of names."""
+
+import tvm_ffi
+
+from tvm import Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("ir.UniqueNameSupply")
+class UniqueNameSupply(Object):
+    """UniqueNameSupply that can be used to generate unique names.
+
+    Parameters
+    ----------
+    prefix: The prefix to be added to the generated names.
+    """
+
+    def __init__(self, prefix=""):
+        self.__init_handle_by_constructor__(_ffi_api.UniqueNameSupply, prefix)
+
+    def fresh_name(self, name, add_prefix=True, add_underscore=True):
+        """Generates a unique name from this UniqueNameSupply.
+
+        Parameters
+        ----------
+        name: String
+            The name from which the generated name is derived.
+
+        add_prefix: bool
+            If set to true, then the prefix of this UniqueNameSupply will be prepended to the name.
+
+        add_underscore: bool
+            If set to True, adds '_' between prefix and digit.
+        """
+        return _ffi_api.UniqueNameSupply_FreshName(self, name, add_prefix, add_underscore)
+
+    def reserve_name(self, name, add_prefix=True):
+        """Reserves an existing name with this UniqueNameSupply.
+
+        Parameters
+        ----------
+        name: String
+            The name to be reserved.
+
+        add_prefix: bool
+            If set to true, then the prefix of this UniqueNameSupply will be prepended to the name
+            before reserving it.
+        """
+        return _ffi_api.UniqueNameSupply_ReserveName(self, name, add_prefix)
+
+    def contains_name(self, name, add_prefix=True):
+        """Checks if this UniqueNameSupply already generated a name.
+
+        Parameters
+        ----------
+        name: String
+            The name to check.
+
+        add_prefix: bool
+            If set to true, then the prefix of this UniqueNameSupply will be prepended to the name
+            before checking for it.
+        """
+        return _ffi_api.UniqueNameSupply_ContainsName(self, name, add_prefix)

python/tvm/ir/transform.py

@@ -0,0 +1,367 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name,unused-argument
+"""Common pass infrastructure across IR variants."""
+
+import functools
+import inspect
+
+import tvm_ffi
+
+import tvm.runtime
+
+from . import _ffi_transform_api
+
+
+@tvm_ffi.register_object("transform.PassInfo")
+class PassInfo(tvm.runtime.Object):
+    """The class contains the meta data required by a pass. It is the
+    container of information needed by running an optimization or analysis.
+    This class can be extended by adding new members when more meta data is
+    needed.
+
+    Parameters
+    ----------
+    opt_level : int
+        The optimization level of this pass.
+
+    name : str
+        The pass name.
+
+    required : List[str]
+        The list of passes that are required by a certain pass.
+    """
+
+    def __init__(self, opt_level, name, required=None, traceable=False):
+        self.__init_handle_by_constructor__(
+            _ffi_transform_api.PassInfo, opt_level, name, required, traceable
+        )
+
+
+@tvm_ffi.register_object("transform.PassContext")
+class PassContext(tvm.runtime.Object):
+    """The basis where a TVM optimization/analysis runs on.
+    Each pass context contains a number of auxiliary information that is used
+    to help an optimization pass. Such information includes the error reporter
+    to record the errors of during the optimization, etc.
+
+    opt_level : Optional[int]
+        The optimization level of this pass.
+
+    required_pass : Optional[Union[List[str], Set[str], Tuple[str]]]
+        The list of passes that are required by a certain pass.
+
+    disabled_pass : Optional[Union[List[str], Set[str], Tuple[str]]]
+        The list of passes that are disabled.
+
+    instruments : Optional[Sequence[PassInstrument]]
+        The list of pass instrument implementations.
+
+    config : Optional[Dict[str, Object]]
+        Additional configurations for specific passes.
+    """
+
+    def __init__(
+        self,
+        opt_level=2,
+        required_pass=None,
+        disabled_pass=None,
+        instruments=None,
+        config=None,
+    ):
+        required = list(required_pass) if required_pass else []
+        if not isinstance(required, list | tuple):
+            raise TypeError("required_pass is expected to be the type of " + "list/tuple/set.")
+
+        disabled = list(disabled_pass) if disabled_pass else []
+        if not isinstance(disabled, list | tuple):
+            raise TypeError("disabled_pass is expected to be the type of " + "list/tuple/set.")
+
+        instruments = list(instruments) if instruments else []
+        if not isinstance(instruments, list | tuple):
+            raise TypeError("instruments is expected to be the type of " + "list/tuple/set.")
+
+        config = config if config else None
+        self.__init_handle_by_constructor__(
+            _ffi_transform_api.PassContext,
+            opt_level,
+            required,
+            disabled,
+            instruments,
+            config,
+        )
+
+    def __enter__(self):
+        _ffi_transform_api.EnterPassContext(self)
+        return self
+
+    def __exit__(self, ptype, value, trace):
+        _ffi_transform_api.ExitPassContext(self)
+
+    def override_instruments(self, instruments):
+        """Override instruments within this PassContext.
+
+        If there are existing instruments, their ``exit_pass_ctx`` callbacks are called.
+        Then switching to new instruments and calling new ``enter_pass_ctx`` callbacks.
+
+        instruments : Sequence[PassInstrument]
+            The list of pass instrument implementations.
+        """
+        _ffi_transform_api.OverrideInstruments(self, instruments)
+
+    @staticmethod
+    def current():
+        """Return the current pass context."""
+        return _ffi_transform_api.GetCurrentPassContext()
+
+    @staticmethod
+    def list_configs():
+        """List all registered `PassContext` configuration names and metadata.
+
+        Returns
+        -------
+        configs : Dict[str, Dict[str, str]]
+
+        """
+        return _ffi_transform_api.ListConfigs()
+
+
+@tvm_ffi.register_object("transform.Pass")
+class Pass(tvm.runtime.Object):
+    """The base class of all passes. All methods here are just simple wrappers
+    that are implemented in the backend. They are defined for users to
+    conveniently interact with the base class.
+    """
+
+    __slots__ = ("__dict__",)
+
+    @property
+    def info(self):
+        """Get the pass meta."""
+        return _ffi_transform_api.Info(self)
+
+    def __call__(self, mod):
+        """Execute the pass. Note that for sequential pass, the dependency among
+        different passes will be resolved in the backend.
+
+        Parameters
+        ----------
+        mod : tvm.IRModule
+            The module that a certain optimization is performed on.
+
+        Returns
+        -------
+        mod : tvm.IRModule
+            The updated module after applying this pass.
+        """
+        return _ffi_transform_api.RunPass(self, mod)
+
+
+@tvm_ffi.register_object("transform.ModulePass")
+class ModulePass(Pass):
+    """A pass that works on tvm.IRModule. Users don't need to interact with
+    this class directly. Instead, a module pass should be created through
+    `module_pass`, because the design of the `module_pass` API is flexible
+    enough to handle the creation of a module pass in different manners. In
+    addition, all members of a module pass can be accessed from the base class.
+    The same rule applies to FunctionPass as well.
+    """
+
+
+@tvm_ffi.register_object("transform.Sequential")
+class Sequential(Pass):
+    """A pass that works on a sequence of pass objects. Multiple passes can be
+    executed sequentially using this class.
+
+    Note that users can also provide a series of passes that they don't want to
+    apply when running a sequential pass. Pass dependency will be resolved in
+    the backend as well.
+
+    Parameters
+    ----------
+    passes : Optional[List[Pass]]
+        A sequence of passes candidate for optimization.
+
+    opt_level : Optional[int]
+        The optimization level of this sequential pass.
+        The opt_level of a default sequential pass is set to 0.
+        Note that some of the passes within the Sequantial may still not be executed
+        if their opt_level is higher than the provided opt_level.
+
+    name : Optional[str]
+        The name of the sequential pass.
+
+    required : Optional[List[str]]
+        The list of passes that the sequential pass is dependent on.
+    """
+
+    def __init__(self, passes=None, opt_level=0, name="sequential", required=None, traceable=False):
+        passes = passes if passes else []
+        if not isinstance(passes, list | tuple):
+            raise TypeError("passes must be a list of Pass objects.")
+
+        required = required if required else []
+        if not isinstance(required, list | tuple):
+            raise TypeError("Required is expected to be the type of list/tuple.")
+
+        self.__init_handle_by_constructor__(
+            _ffi_transform_api.Sequential, passes, opt_level, name, required, traceable
+        )
+
+
+def _wrap_class_module_pass(pass_cls, pass_info):
+    """Wrap a python class as function pass"""
+
+    class PyModulePass(ModulePass):
+        """Internal wrapper class to create a class instance."""
+
+        def __init__(self, *args, **kwargs):
+            inst = pass_cls(*args, **kwargs)
+
+            # it is important not to capture self to
+            # avoid a cyclic dependency
+            def _pass_func(mod, ctx):
+                return inst.transform_module(mod, ctx)
+
+            self.__init_handle_by_constructor__(
+                _ffi_transform_api.MakeModulePass, _pass_func, pass_info
+            )
+            self._inst = inst
+
+        def __getattr__(self, name):
+            # fall back to instance attribute if there is not any
+            return self._inst.__getattribute__(name)
+
+    functools.update_wrapper(PyModulePass.__init__, pass_cls.__init__)
+    PyModulePass.__name__ = pass_cls.__name__
+    PyModulePass.__doc__ = pass_cls.__doc__
+    PyModulePass.__module__ = pass_cls.__module__
+    return PyModulePass
+
+
+def module_pass(pass_func=None, opt_level=None, name=None, required=None, traceable=False):
+    """Decorate a module pass.
+
+    This function returns a callback when pass_func is provided.
+    Otherwise, it serves a decorator function.
+
+    pass_func can also be a class type with a method transform_module.
+    This function will create a decorated ModulePass using transform_module
+    as the pass function.
+
+    Parameters
+    ----------
+    pass_func : Optional[Callable[(Module, PassContext) ->Module]]
+        The transformation function or class.
+
+    opt_level : int
+        The optimization level of this module pass.
+
+    name : Optional[str]
+        The name of the module pass. The name could be empty. In this case, the
+        name of the optimization function will be used as the pass name.
+
+    required : Optional[List[str]]
+        The list of passes that the module pass is dependent on.
+
+    traceable: Boolean
+        Boolean variable whether the module pass is traceable
+
+    Returns
+    -------
+    create_module_pass : Union[Callable, ModulePass]
+        A decorator will be returned if pass_func is not provided,
+        otherwise return the decorated result.
+        The returned decorator has two behaviors depending on the input:
+        A new ModulePass will be returned when we decorate a pass function.
+        A new ModulePass class will be returned when we decorate a class type.
+
+    Examples
+    --------
+    The following code block decorates a module pass class.
+
+    .. code-block:: python
+
+        @tvm.ir.transform.module_pass
+        class CustomPipeline:
+            def __init__(self, enable_fold):
+                self.enable_fold = enable_fold
+                self.const_fold = relax.transform.FoldConstant()
+
+            def transform_module(self, mod, ctx):
+                if self.enable_fold:
+                    mod = self.const_fold(mod, ctx)
+                return mod
+
+        # create an instance of customized pipeline
+        pipeline = CustomPipeline(enable_fold=False)
+        assert isinstance(pipeline, transform.ModulePass)
+        # run the pipeline.
+        output_module = pipeline(input_module)
+
+    The following code creates a module pass by decorating
+    a user defined transform function.
+
+    .. code-block:: python
+
+        @tvm.ir.transform.module_pass(opt_level=2)
+        def transform(mod, ctx):
+            return relax.transform.FoldConstant(mod)
+
+        module_pass = transform
+        assert isinstance(module_pass, transform.ModulePass)
+        assert module_pass.info.opt_level == 2
+
+        # Given a module m, the optimization could be invoked as the follwoing:
+        updated_mod = module_pass(m)
+        # Now a function abs should be added to the module m.
+    """
+    if opt_level is None:
+        raise ValueError("Please provide opt_level for the module pass.")
+
+    required = required if required else []
+    if not isinstance(required, list | tuple):
+        raise TypeError("Required is expected to be the type of " + "list/tuple.")
+
+    def create_module_pass(pass_arg):
+        """Internal function that creates a module pass"""
+        fname = name if name else pass_arg.__name__
+        info = PassInfo(opt_level, fname, required, traceable)
+        if inspect.isclass(pass_arg):
+            return _wrap_class_module_pass(pass_arg, info)
+        if not callable(pass_arg):
+            raise TypeError("pass_func must be a callable for Module pass")
+        return _ffi_transform_api.MakeModulePass(pass_arg, info)
+
+    if pass_func:
+        return create_module_pass(pass_func)
+    return create_module_pass
+
+
+def PrintIR(header=""):
+    """A special trace pass that prints the header and IR.
+
+    Parameters
+    ----------
+    header : str
+        The header to be displayed along with the dump.
+
+    Returns
+    --------
+    The pass
+    """
+    return _ffi_transform_api.PrintIR(header)

python/tvm/ir/type.py

@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unified type system in the project."""
+
+import tvm_ffi
+
+from tvm.runtime import Scriptable
+
+from . import _ffi_api
+from .base import Node
+
+
+@tvm_ffi.register_object("ir.Type")
+class Type(Node, Scriptable):
+    """The base class of all types."""
+
+    @staticmethod
+    def missing():
+        """Return the sentinel for missing type information."""
+        return _ffi_api.TypeMissing()
+
+    @staticmethod
+    def Missing():
+        """Return the sentinel for missing type information."""
+        return _ffi_api.TypeMissing()
+
+    def is_missing(self):
+        """Return whether this is the missing-type sentinel."""
+        return _ffi_api.TypeIsMissing(self)
+
+    def __eq__(self, other):
+        """Compare two types for structural equivalence."""
+        return bool(tvm_ffi.structural_equal(self, other))
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def same_as(self, other):
+        """Compares two TVM types by referential equality."""
+        return self.is_(other)
+
+
+@tvm_ffi.register_object("ir.PrimType")
+class PrimType(Type):
+    """Primitive data type in the low level IR
+
+    Parameters
+    ----------
+    dtype : str
+        The runtime data type relates to the primtype.
+    """
+
+    def __init__(self, dtype):
+        self.__init_handle_by_constructor__(_ffi_api.PrimType, dtype)
+
+    def __eq__(self, other):
+        if isinstance(other, str):
+            return self.dtype == other
+        return super().__eq__(other)
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __hash__(self):
+        dtype = self.dtype
+        return hash((dtype.type_code, dtype.bits, dtype.lanes))
+
+    def __str__(self):
+        return str(self.dtype)
+
+    def matches_code(self, *codes) -> bool:
+        """Return whether this type has any of the given DLPack dtype codes."""
+        type_code = self.dtype.type_code
+        return any(type_code == int(code) for code in codes)
+
+    def matches_element_type(self, code, bits: int) -> bool:
+        """Return whether this type has the given scalar element code and bits."""
+        dtype = self.dtype
+        return dtype.type_code == int(code) and dtype.bits == bits
+
+    def is_scalar(self) -> bool:
+        """Return whether this type has exactly one fixed lane."""
+        return self.dtype.lanes == 1
+
+
+@tvm_ffi.register_object("ir.PointerType")
+class PointerType(Type):
+    """PointerType used in the low-level TIR.
+
+    Parameters
+    ----------
+    element_type : tvm.ir.Type
+        The type of pointer's element.
+
+    storage_scope : str
+        The storage scope into which the pointer addresses.
+    """
+
+    def __init__(self, element_type, storage_scope=""):
+        self.__init_handle_by_constructor__(_ffi_api.PointerType, element_type, storage_scope)
+
+
+@tvm_ffi.register_object("ir.TupleType")
+class TupleType(Type):
+    """The type of tuple values.
+
+    Parameters
+    ----------
+    fields : List[Type]
+        The fields in the tuple
+    """
+
+    def __init__(self, fields, span=None):
+        self.__init_handle_by_constructor__(_ffi_api.TupleType, fields, span)
+
+
+@tvm_ffi.register_object("ir.FuncType")
+class FuncType(Type):
+    """Function type.
+
+    A function type consists of a list of type parameters to enable
+    the definition of generic functions,
+    a set of type constraints which we omit for the time being,
+    a sequence of argument types, and a return type.
+
+    Parameters
+    ----------
+    arg_types : List[tvm.ir.Type]
+        The argument types
+
+    ret_type : tvm.ir.Type
+        The return type.
+    """
+
+    def __init__(self, arg_types, ret_type):
+        self.__init_handle_by_constructor__(
+            _ffi_api.FuncType,
+            arg_types,
+            ret_type,
+        )
+
+
+@tvm_ffi.register_object("ir.TensorMapType")
+class TensorMapType(Type):
+    """TensorMapType used in the low-level TIR.
+
+    Parameters
+    ----------
+    span : tvm.ir.Span
+        The span information.
+    """
+
+    def __init__(self, span=None):
+        self.__init_handle_by_constructor__(
+            _ffi_api.TensorMapType,
+            span,  # pylint: disable=no-member
+        )

python/tvm/ir/type_relation.py

@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Type relation and function for type checking."""
+
+import tvm_ffi
+
+from . import _ffi_api
+from .type import Type, TypeConstraint
+
+
+@tvm_ffi.register_object("TypeCall")
+class TypeCall(Type):
+    """Type function application.
+
+    Parameters
+    ----------
+    func: tvm.ir.Type
+        The function.
+
+    args: List[tvm.ir.Type]
+        The arguments.
+
+    Returns
+    -------
+    type_call: TypeCall
+        The type function application.
+    """
+
+    def __init__(self, func, args):
+        self.__init_handle_by_constructor__(_ffi_api.TypeCall, func, args)
+
+
+@tvm_ffi.register_object("TypeRelation")
+class TypeRelation(TypeConstraint):
+    """User defined type relation, it is an input-output relation on types.
+
+    TypeRelation is more generalized than TypeCall as it allows inference
+     of both inputs and outputs.
+
+    Parameters
+    ----------
+    func : EnvFunc
+        User defined relation function.
+
+    args : [tvm.ir.Type]
+        List of types to the func.
+
+    num_inputs : int
+        Number of input arguments in args,
+        this act as a hint for type inference.
+
+    attrs : Attrs
+        The attribute attached to the relation information
+
+    Returns
+    -------
+    type_relation : tvm.ir.TypeRelation
+        The type relation.
+    """
+
+    def __init__(self, func, args, num_inputs, attrs):
+        self.__init_handle_by_constructor__(_ffi_api.TypeRelation, func, args, num_inputs, attrs)

python/tvm/ir/utils.py

@@ -0,0 +1,161 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Utilities shared across TVM IR packages."""
+
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def derived_object(cls: type[T]) -> type[T]:
+    """A decorator to register derived subclasses for TVM objects.
+
+    Parameters
+    ----------
+    cls : type
+        The derived class to be registered.
+
+    Returns
+    -------
+    cls : type
+        The decorated TVM object.
+
+    Example
+    -------
+    .. code-block:: python
+
+        @register_object("s_tir.meta_schedule.PyRunner")
+        class _PyRunner(meta_schedule.Runner):
+            def __init__(self, f_run: Callable = None):
+                self.__init_handle_by_constructor__(_ffi_api.RunnerPyRunner, f_run)
+
+        class PyRunner:
+            _tvm_metadata = {
+                "cls": _PyRunner,
+                "methods": ["run"]
+            }
+            def run(self, runner_inputs):
+                raise NotImplementedError
+
+        @derived_object
+        class LocalRunner(PyRunner):
+            def run(self, runner_inputs):
+                ...
+    """
+
+    import functools  # pylint: disable=import-outside-toplevel
+    import weakref  # pylint: disable=import-outside-toplevel
+
+    def _extract(inst: type, name: str):
+        """Extract function from intrinsic class."""
+
+        def method(*args, **kwargs):
+            return getattr(inst, name)(*args, **kwargs)
+
+        for inherit_cls, base_cls in zip(cls.__mro__, cls.__mro__[1:]):
+            # extract functions that differ from the base class
+            if not hasattr(base_cls, name):
+                continue
+            if getattr(base_cls, name) is getattr(inherit_cls, name):
+                continue
+            return method
+
+        # for task scheduler return None means calling default function
+        # otherwise it will trigger a RuntimeError of method not implemented
+        # on the c++ side when you call the method
+        return None
+
+    assert isinstance(cls.__base__, type)
+    if hasattr(cls, "_type") and cls._type == "TVMDerivedObject":  # type: ignore
+        raise TypeError(
+            f"Inheritance from a decorated object `{cls.__name__}` is not allowed. "
+            f"Please inherit from `{cls.__name__}._cls`."
+        )
+    assert hasattr(cls, "_tvm_metadata"), (
+        "Please use the user-facing method overriding class, i.e., PyRunner."
+    )
+
+    base = cls.__base__
+    metadata = getattr(base, "_tvm_metadata")
+    fields = metadata.get("fields", [])
+    methods = metadata.get("methods", [])
+
+    base_cls = metadata["cls"]
+    slots = []
+    if getattr(base_cls, "__dictoffset__", 0) == 0:
+        slots.append("__dict__")
+    if getattr(base_cls, "__weakrefoffset__", 0) == 0:
+        slots.append("__weakref__")
+
+    class TVMDerivedObject(base_cls):  # type: ignore
+        """The derived object to avoid cyclic dependency."""
+
+        __slots__ = tuple(slots)
+
+        _cls = cls
+        _type = "TVMDerivedObject"
+
+        def __init__(self, *args, **kwargs):
+            """Constructor."""
+            self._inst = cls(*args, **kwargs)
+
+            super().__init__(
+                # the constructor's parameters, builder, runner, etc.
+                *[getattr(self._inst, name) for name in fields],
+                # the function methods, init_with_tune_context, build, run, etc.
+                *[_extract(self._inst, name) for name in methods],
+            )
+
+            # for task scheduler hybrid funcs in c++ & python side
+            # using weakref to avoid cyclic dependency
+            self._inst._outer = weakref.ref(self)
+
+        def __getattr__(self, name):
+            import inspect  # pylint: disable=import-outside-toplevel
+
+            try:
+                # fall back to instance attribute if there is not any
+                # return self._inst.__getattribute__(name)
+                result = self._inst.__getattribute__(name)
+            except AttributeError:
+                result = super().__getattr__(name)
+
+            if inspect.ismethod(result):
+
+                def method(*args, **kwargs):
+                    return result(*args, **kwargs)
+
+                # set __own__ to aviod implicit deconstruction
+                setattr(method, "__own__", self)
+                return method
+
+            return result
+
+        def __setattr__(self, name, value):
+            if name not in ["_inst", "key", "handle"]:
+                self._inst.__setattr__(name, value)
+            else:
+                super().__setattr__(name, value)
+
+    functools.update_wrapper(TVMDerivedObject.__init__, cls.__init__)  # type: ignore
+    TVMDerivedObject.__name__ = cls.__name__
+    TVMDerivedObject.__doc__ = cls.__doc__
+    TVMDerivedObject.__module__ = cls.__module__
+    for key, value in cls.__dict__.items():
+        if isinstance(value, classmethod | staticmethod):
+            setattr(TVMDerivedObject, key, value)
+    return TVMDerivedObject