chore: import upstream snapshot with attribution

python/tvm/s_tir/schedule/__init__.py

@@ -0,0 +1,28 @@
+# 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
+"""Namespace for the TensorIR schedule API."""
+
+from ..sblock_scope import SBlockScope, Dependency, DepKind, StmtSRef
+from .instruction import Instruction, InstructionKind
+from .schedule import SBlockRV, ExprRV, LoopRV, Schedule, ScheduleError
+from .state import ScheduleDebugMask, ScheduleState
+from .trace import Trace
+
+from . import analysis
+from . import transform

python/tvm/s_tir/schedule/_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.s_tir.schedule"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("s_tir.schedule", __name__)  # pylint: disable=protected-access

python/tvm/s_tir/schedule/_type_checker.py

@@ -0,0 +1,370 @@
+# 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
+"""Type checking functionality"""
+
+import collections
+import collections.abc
+import functools
+import inspect
+import types
+import typing
+from collections.abc import Callable
+from typing import Any, TypeVar, Union
+
+
+def _is_none_type(type_: Any) -> bool:
+    return type_ is None or type_ is type(None)
+
+
+def _get_subtypes(type_: Any) -> Any:
+    # TODO(@tvm-team): This is hot fix to support subtle difference between python versions
+    #                  Would be nice to find a better way if possible
+    if hasattr(typing, "_SpecialGenericAlias"):
+        if hasattr(typing, "get_args"):
+            subtypes = typing.get_args(type_)  # type: ignore
+        else:
+            subtypes = type_.__args__
+    else:
+        subtypes = type_.__args__
+    return subtypes
+
+
+if hasattr(typing, "_GenericAlias"):
+    # For python versions 3.7 onward, check the __origin__ attribute.
+
+    class _Subtype:
+        @staticmethod
+        def _origin(type_: Any) -> Any:
+            # In Python 3.14+, check if the type has __origin__ attribute directly
+            if hasattr(type_, "__origin__"):
+                return type_.__origin__
+
+            if hasattr(typing, "_SpecialGenericAlias"):
+                if isinstance(type_, typing._SpecialGenericAlias):  # type: ignore # pylint: disable=protected-access
+                    return type_.__origin__
+
+            if isinstance(type_, typing._GenericAlias):  # type: ignore # pylint: disable=protected-access
+                return type_.__origin__
+            return None
+
+        @staticmethod
+        def list_(type_: Any) -> Any:
+            if _Subtype._origin(type_) is list:
+                if hasattr(typing, "get_args"):
+                    args = typing.get_args(type_)  # type: ignore
+                else:
+                    args = type_.__args__
+                if len(args) == 1:
+                    return [args[0]]
+                # Handle list[X | Y] where get_args may return individual types
+                return [Union[args]]  # noqa: UP007 (runtime use)
+            return None
+
+        @staticmethod
+        def dict_(type_: Any) -> Any:
+            if _Subtype._origin(type_) is dict:
+                if hasattr(typing, "get_args"):
+                    (ktype, vtype) = typing.get_args(type_)  # type: ignore
+                else:
+                    (ktype, vtype) = type_.__args__
+                return [ktype, vtype]
+            return None
+
+        @staticmethod
+        def tuple_(type_: Any) -> list[type] | None:
+            if _Subtype._origin(type_) is tuple:
+                subtypes = _get_subtypes(type_)
+                return subtypes
+            return None
+
+        @staticmethod
+        def optional(  # pylint: disable=missing-function-docstring
+            type_: Any,
+        ) -> list[type] | None:
+            if _Subtype._origin(type_) is Union:
+                subtypes = _get_subtypes(type_)
+                if len(subtypes) == 2 and _is_none_type(subtypes[1]):
+                    return [subtypes[0]]
+            # PEP 604: X | None
+            if isinstance(type_, types.UnionType):
+                subtypes = type_.__args__
+                if len(subtypes) == 2 and _is_none_type(subtypes[1]):
+                    return [subtypes[0]]
+            return None
+
+        @staticmethod
+        def union(type_: Any) -> list[type] | None:  # pylint: disable=missing-function-docstring
+            if _Subtype._origin(type_) is Union:
+                subtypes = _get_subtypes(type_)
+                if len(subtypes) != 2 or not _is_none_type(subtypes[1]):
+                    return list(subtypes)
+            # PEP 604: X | Y
+            if isinstance(type_, types.UnionType):
+                subtypes = type_.__args__
+                if len(subtypes) != 2 or not _is_none_type(subtypes[1]):
+                    return list(subtypes)
+            return None
+
+        @staticmethod
+        def callable(type_: Any) -> list[type] | None:
+            if _Subtype._origin(type_) is collections.abc.Callable:
+                subtypes = _get_subtypes(type_)
+                return subtypes
+            return None
+
+elif hasattr(typing, "_Union"):
+    # For python 3.6 and below, check the __name__ attribute, or CallableMeta.
+
+    class _Subtype:  # type: ignore
+        @staticmethod
+        def list_(type_: Any) -> list[type] | None:
+            if isinstance(type_, typing.GenericMeta):  # type: ignore # pylint: disable=no-member
+                if type_.__name__ == "List":
+                    args = type_.__args__  # type: ignore # pylint: disable=no-member
+                    if len(args) == 1:
+                        return [args[0]]
+                    # Handle list[X | Y] where args may return individual types
+                    return [Union[args]]  # noqa: UP007 (runtime use)
+            return None
+
+        @staticmethod
+        def dict_(type_: Any) -> list[type] | None:
+            if isinstance(type_, typing.GenericMeta):  # type: ignore # pylint: disable=no-member
+                if type_.__name__ == "Dict":
+                    (ktype, vtype) = type_.__args__  # type: ignore # pylint: disable=no-member
+                    return [ktype, vtype]
+            return None
+
+        @staticmethod
+        def tuple_(type_: Any) -> list[type] | None:
+            if isinstance(type_, typing.GenericMeta):  # type: ignore # pylint: disable=no-member
+                if type_.__name__ == "Tuple":
+                    subtypes = type_.__args__  # type: ignore # pylint: disable=no-member
+                    return subtypes
+            return None
+
+        @staticmethod
+        def optional(type_: Any) -> list[type] | None:
+            if isinstance(type_, typing._Union):  # type: ignore # pylint: disable=no-member,protected-access
+                subtypes = type_.__args__
+                if len(subtypes) == 2 and _is_none_type(subtypes[1]):
+                    return [subtypes[0]]
+            return None
+
+        @staticmethod
+        def union(type_: Any) -> list[type] | None:
+            if isinstance(type_, typing._Union):  # type: ignore # pylint: disable=no-member,protected-access
+                subtypes = type_.__args__
+                if len(subtypes) != 2 or not _is_none_type(subtypes[1]):
+                    return list(subtypes)
+            return None
+
+        @staticmethod
+        def callable(type_: Any) -> list[type] | None:
+            if isinstance(type_, typing.CallableMeta):  # type: ignore # pylint: disable=no-member,protected-access
+                subtypes = type_.__args__
+                return subtypes
+            return None
+
+
+def _dispatcher(type_: Any) -> tuple[str, list[type]]:
+    if _is_none_type(type_):
+        return "none", []
+
+    subtype = _Subtype.list_(type_)
+    if subtype is not None:
+        return "list", subtype
+
+    subtype = _Subtype.dict_(type_)
+    if subtype is not None:
+        return "dict", subtype
+
+    subtype = _Subtype.tuple_(type_)
+    if subtype is not None:
+        return "tuple", subtype
+
+    subtype = _Subtype.optional(type_)
+    if subtype is not None:
+        return "optional", subtype
+
+    subtype = _Subtype.union(type_)
+    if subtype is not None:
+        return "union", subtype
+
+    subtype = _Subtype.callable(type_)
+    if subtype is not None:
+        return "callable", subtype
+
+    return "atomic", [type_]
+
+
+def callable_str(*subtypes):
+    if subtypes:
+        *arg_types, return_type = subtypes
+        arg_str = ", ".join(_type2str(arg_type) for arg_type in arg_types)
+        return_type_str = _type2str(return_type)
+        return f"Callable[[{arg_str}], {return_type_str}]"
+    else:
+        return "Callable"
+
+
+_TYPE2STR: dict[Any, Callable] = {
+    "none": lambda: "None",
+    "atomic": lambda t: str(t.__name__),
+    "callable": callable_str,
+    "list": lambda t: f"List[{_type2str(t)}]",
+    "dict": lambda k, v: f"Dict[{_type2str(k)}, {_type2str(v)}]",
+    "tuple": lambda *t: f"Tuple[{', '.join([_type2str(x) for x in t])}]",
+    "optional": lambda t: f"Optional[{_type2str(t)}]",
+    "union": lambda *t: f"Union[{', '.join([_type2str(x) for x in t])}]",
+}
+
+
+def _type2str(type_: Any) -> str:
+    key, subtypes = _dispatcher(type_)
+    return _TYPE2STR[key](*subtypes)
+
+
+def _val2type(value: Any):
+    if isinstance(value, list):
+        types = set(_val2type(x) for x in value)
+        if len(types) == 1:
+            return list[types.pop()]  # type: ignore
+
+        return list[tuple(types)]  # type: ignore
+
+    if isinstance(value, tuple):
+        types = tuple(_val2type(x) for x in value)  # type: ignore
+        return tuple[types]
+
+    return type(value)
+
+
+def _type_check_err(x: Any, name: str, expected: Any) -> str:
+    return (
+        f'"{name}" has wrong type. '
+        f'Expected "{_type2str(expected)}", '
+        f'but gets: "{_type2str(_val2type(x))}"'
+    )
+
+
+def _type_check_vtable() -> dict[str, Callable]:
+    def _type_check_none(v: Any, name: str) -> str | None:
+        return None if v is None else _type_check_err(v, name, None)
+
+    def _type_check_atomic(v: Any, name: str, type_: Any) -> str | None:
+        return None if isinstance(v, type_) else _type_check_err(v, name, type_)
+
+    def _type_check_callable(v: Any, name: str, *_subtypes: Any) -> str | None:
+        # Current implementation only validates that the argument is
+        # callable, and doesn't validate the arguments accepted by the
+        # callable, if any.
+        return None if callable(v) else _type_check_err(v, name, Callable)
+
+    def _type_check_list(v: list[Any], name: str, type_: Any) -> str | None:
+        if not isinstance(v, list | tuple):
+            return _type_check_err(v, name, list)
+        for i, x in enumerate(v):
+            error_msg = _type_check(x, f"{name}[{i}]", type_)
+            if error_msg is not None:
+                return error_msg
+        return None
+
+    def _type_check_dict(dict_obj: dict[Any, Any], name: str, *types: Any) -> str | None:
+        ktype_, vtype_ = types
+        if not isinstance(dict_obj, dict):
+            return _type_check_err(dict_obj, name, dict)
+        for k, v in dict_obj.items():
+            error_msg = _type_check(k, f"{name}[{k}]", ktype_)
+            if error_msg is not None:
+                return error_msg
+            error_msg = _type_check(v, f"{name}[{k}]", vtype_)
+            if error_msg is not None:
+                return error_msg
+        return None
+
+    def _type_check_tuple(v: Any, name: str, *types: Any) -> str | None:
+        if not isinstance(v, tuple):
+            return _type_check_err(v, name, tuple[types])
+        if len(types) != len(v):
+            return _type_check_err(v, name, tuple[types])
+        for i, (x, type_) in enumerate(zip(v, types)):
+            error_msg = _type_check(x, f"{name}[{i}]", type_)
+            if error_msg is not None:
+                return error_msg
+        return None
+
+    def _type_check_optional(v: Any, name: str, type_: Any) -> str | None:
+        return None if v is None else _type_check(v, name, type_)
+
+    def _type_check_union(v: Any, name: str, *types: Any) -> str | None:
+        for type_ in types:
+            error_msg = _type_check(v, name, type_)
+            if error_msg is None:
+                return None
+        return _type_check_err(v, name, Union[types])  # noqa: UP007 (runtime use)
+
+    return {
+        "none": _type_check_none,
+        "atomic": _type_check_atomic,
+        "callable": _type_check_callable,
+        "list": _type_check_list,
+        "dict": _type_check_dict,
+        "tuple": _type_check_tuple,
+        "optional": _type_check_optional,
+        "union": _type_check_union,
+    }
+
+
+_TYPE_CHECK: dict[Any, Callable] = _type_check_vtable()
+
+
+def _type_check(v: Any, name: str, type_: Any) -> str | None:
+    key, subtypes = _dispatcher(type_)
+    return _TYPE_CHECK[key](v, name, *subtypes)
+
+
+FType = TypeVar("FType", bound=Callable[..., Any])
+
+
+def type_checked(func: FType) -> FType:
+    """Type check the input arguments of a function."""
+    sig = inspect.signature(func)
+    try:
+        hints = typing.get_type_hints(func)
+    except Exception:
+        hints = {}
+
+    @functools.wraps(func)
+    def wrap(*args, **kwargs):
+        bound_args = sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+        for param in sig.parameters.values():
+            type_hint = hints.get(param.name, inspect.Parameter.empty)
+            if type_hint != inspect.Parameter.empty:
+                error_msg = _type_check(
+                    bound_args.arguments[param.name],
+                    param.name,
+                    type_hint,
+                )
+                if error_msg is not None:
+                    error_msg = f'In "{func.__qualname__}", {error_msg}'
+                    raise TypeError(error_msg)
+        return func(*args, **kwargs)
+
+    return wrap  # type: ignore

python/tvm/s_tir/schedule/analysis.py

@@ -0,0 +1,160 @@
+# 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.
+"""Analysis used in TensorIR scheduling"""
+
+import tvm_ffi
+
+from tvm.runtime import Object
+from tvm.tirx.buffer import Buffer
+from tvm.tirx.expr import Expr
+from tvm.tirx.function import IndexMap, PrimFunc
+from tvm.tirx.stmt import For
+
+from . import _ffi_api
+from .schedule import SBlockRV, Schedule
+
+
+def suggest_index_map(
+    buffer: Buffer,
+    indices: list[Expr],
+    loops: list[For],
+    predicate: Expr,
+) -> IndexMap | None:
+    """Provided the access pattern to a buffer, suggest one of the possible layout
+    transformation to maximize the locality of the access pattern.
+
+    Parameters
+    ----------
+    buffer : Buffer
+        The buffer to be transformed.
+    indices : List[Expr]
+        The access pattern to the buffer.
+    loops : List[For]
+        The loops above the buffer.
+    predicate : Expr
+        The predicate of the access.
+
+    Returns
+    -------
+    index_map : Optional[IndexMap]
+        The suggested index map. None if no transformation is suggested.
+    """
+    return _ffi_api.SuggestIndexMap(  # type: ignore # pylint: disable=no-member
+        buffer,
+        indices,
+        loops,
+        predicate,
+    )
+
+
+@tvm_ffi.register_object("s_tir.schedule.TensorizeInfo")
+class TensorizeInfo(Object):
+    """Necessary information used for tensorization."""
+
+
+def get_tensorize_loop_mapping(
+    sch: Schedule, block: SBlockRV, desc_func: PrimFunc, allow_padding: bool = False
+) -> TensorizeInfo | None:
+    """Establish a mapping between loops in a target block and an intrinsic description
+
+    Parameters
+    ----------
+    sch : Schedule
+        The schedule to be tensorized
+    block : SBlockRV
+        The target block to match against
+    desc_func : PrimFunc
+        The prim func describing the computation to be tensorized
+    allow_padding : bool
+        Whether to allow padding the block iters to match the intrinsic description
+    Returns
+    -------
+    tensorize_info : Optional[TensorizeInfo]
+        TensorizeInfo structure if a valid mapping is found, None otherwise
+    """
+    return _ffi_api.GetTensorizeLoopMapping(sch, block, desc_func, allow_padding)  # type: ignore
+
+
+@tvm_ffi.register_object("s_tir.schedule.AutoTensorizeMappingInfo")
+class AutoTensorizeMappingInfo(Object):
+    """Necessary information used to perform transformations for tensorization."""
+
+
+def get_auto_tensorize_mapping_info(
+    sch: Schedule, block: SBlockRV, desc_func: PrimFunc
+) -> AutoTensorizeMappingInfo | None:
+    """Get mapping info between a target block and an intrinsic description including layout
+    transformations to apply.
+
+    Parameters
+    ----------
+    sch : Schedule
+        The schedule to be tensorized
+    block : SBlockRV
+        The compute block for auto tensorization
+    desc_func : PrimFunc
+        The prim func describing the computation to be tensorized
+
+    Returns
+    -------
+    auto_tensorize_mapping_info : Optional[AutoTensorizeMappingInfo]
+        AutoTensorizeMappingInfo structure if potential mappings found, None otherwise.
+
+    Note
+    ----
+    Returning a valid AutoTensorizeMappingInfo doesn't guarantee the block can be tensorized.
+    We will need to apply the suggested layout transformations and then match against the tensor
+    intrinsics.
+    """
+    return _ffi_api.GetAutoTensorizeMappingInfo(sch, block, desc_func)  # type: ignore
+
+
+def has_block(sch: Schedule, block_name: str) -> bool:
+    """Query if the given block name exists in the module associated with the provided schedule.
+
+    Parameters
+    ----------
+    sch : Schedule
+        The schedule
+    block_name : str
+        The name of the block to query
+
+    Returns
+    -------
+    yes/no: bool
+        True if the given block exists in the schedule.
+    """
+    return _ffi_api.HasBlock(sch, block_name)  # type: ignore
+
+
+def is_output_block(sch: Schedule, block: SBlockRV) -> bool:
+    """Check whether the given block is an output block
+
+    Parameters
+    ----------
+    sch : Schedule
+        The schedule object of the block
+    block : SBlockRV
+        The blockRV to be checked
+
+    Returns
+    -------
+    yes/no : bool
+        True if the given block is an output block
+
+    """
+    return _ffi_api.IsOutputBlock(sch, block)  # type: ignore

python/tvm/s_tir/schedule/instruction.py

@@ -0,0 +1,168 @@
+# 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.
+"""Schedule instructions each corresponds to a schedule primitive"""
+
+from typing import TYPE_CHECKING, Any
+
+from tvm_ffi import register_object as _register_object
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+if TYPE_CHECKING:
+    from .schedule import RAND_VAR_TYPE
+
+    INPUT_RV_TYPE = RAND_VAR_TYPE | float | int | str | None  # pylint: disable=invalid-name
+    OUTPUT_RV_TYPE = RAND_VAR_TYPE  # pylint: disable=invalid-name
+    ATTR_TYPE = Any
+else:
+    INPUT_RV_TYPE = OUTPUT_RV_TYPE = ATTR_TYPE = Any
+
+
+@_register_object("s_tir.InstructionKind")
+class InstructionKind(Object):
+    """Kind of an instruction, e.g. Split, Reorder, etc.
+    Besides the name, every kind of instruction has its own properties, including:
+    1) A boolean indicating if the instruction is pure, i.e. change nothing in the schedule state
+    2) A functor that applies the instruction to a TensorIR schedule
+    3) A functor that converts the instruction to a statement in python syntax
+    4) A functor that serialize its attributes to JSON
+    5) A functor that deserialize its attributes from JSON
+
+    Unlike `tvm.ir.op`, `InstructionKind` doesn't support unstructured properties,
+    mainly because there is no such usecase yet to add any other property.
+
+    Attributes
+    ----------
+    name : str
+        The name of a kind of instructions
+
+    Note
+    ----
+    The functor properties are not exposed on python side at the moment
+    """
+
+    name: str
+
+    @property
+    def is_pure(self) -> bool:
+        """Indicates if the instruction is pure, i.e. removing it alone doesn't mutate the schedule
+        state. For example, the instruction `GetSBlock` is pure because it changes
+        nothing, while `ComputeInline` is not because removing it leads to a different resulting
+        schedule.
+
+        Returns
+        -------
+        pure : bool
+            The boolean flag indicating if the instruction is pure
+        """
+        return bool(self._is_pure)
+
+    @staticmethod
+    def get(name: str) -> "InstructionKind":
+        """Retrieve an InstructionKind using its name
+
+        Parameters
+        ----------
+        name : str
+            The registered name of the InstructionKind
+
+        Returns
+        -------
+        kind : InstructionKind
+            The InstructionKind retrieved
+        """
+        return _ffi_api.InstructionKindGet(name)  # type: ignore # pylint: disable=no-member
+
+
+@_register_object("s_tir.Instruction")
+class Instruction(Object):
+    """Schedule instructions each corresponds to a schedule primitive
+
+    Attributes
+    ----------
+    kind : InstructionKind
+        The kind of the instruction
+    inputs : List[INPUT_RV_TYPE]
+        The input random variables of the instruction,
+        and the type of each element can be one of the following:
+        - SBlockRV
+        - LoopRV
+        - ExprRV
+        - float
+        - int
+        - str
+        - None
+    attrs : List[ATTR_TYPE]
+        The attributes of the instruction. Similar to attributes of an operator,
+        attributes of an instruction are arbitrary constant metadata required by the instructions.
+        For example, the name of the block to be retrieved in `GetSBlock`.
+    outputs : List[OUTPUT_RV_TYPE]
+        The output random variables of the instruction,
+        and the type of each element can be one of the following:
+        - SBlockRV
+        - LoopRV
+        - ExprRV, atomic variables only, won't be constants or composite Expr
+    """
+
+    kind: InstructionKind
+    inputs: list[INPUT_RV_TYPE]
+    attrs: list[ATTR_TYPE]
+    outputs: list[OUTPUT_RV_TYPE]
+
+    def __init__(
+        self,
+        kind: InstructionKind,
+        inputs: list[INPUT_RV_TYPE],
+        attrs: list[ATTR_TYPE],
+        outputs: list[OUTPUT_RV_TYPE],
+    ) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        kind : InstructionKind
+            The kind of the instruction
+        inputs : List[INPUT_RV_TYPE]
+            The input random variables of the instruction,
+            and the type of each element can be one of the following:
+            - SBlockRV
+            - LoopRV
+            - ExprRV
+            - float
+            - int
+            - str
+            - None
+        attrs : List[ATTR_TYPE]
+            The attributes of the instruction. Similar to attributes of an operator,
+            attributes of an instruction are arbitrary constant metadata required by the
+            instructions. For example, the name of the block to be retrieved in `GetSBlock`.
+        outputs : List[OUTPUT_RV_TYPE]
+            The output random variables of the instruction,
+            and the type of each element can be one of the following:
+            - SBlockRV
+            - LoopRV
+            - ExprRV, atomic variables only, won't be constants or composite Expr
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.Instruction,  # type: ignore # pylint: disable=no-member
+            kind,
+            inputs,
+            attrs,
+            outputs,
+        )

python/tvm/s_tir/schedule/state.py

@@ -0,0 +1,238 @@
+# 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
+"""This file defines ScheduleState, the core data structure of TensorIR scheduling."""
+
+from collections import namedtuple
+from enum import IntEnum
+
+from tvm_ffi import register_object
+
+from tvm.ir import IRModule
+from tvm.runtime import Object
+from tvm.tirx import For, PrimFunc, SBlock, SBlockRealize
+
+from ..sblock_scope import SBlockScope, StmtSRef
+from . import _ffi_api
+
+CachedFlags = namedtuple("CachedFlags", ["affine_binding", "region_cover", "stage_pipeline"])
+
+
+class ScheduleDebugMask(IntEnum):
+    """The bitmask of the `debug_mask` flag in the ScheduleState class.
+
+    If the `debug_mask` flag has a certain bit on, then the correpsonding
+    verification pass will be conducted. For example, if `(debug_mask & VERIFY_SREF_TREE) != 0`,
+    then the correctness of the sref tree will be verified after each schedule instruction.
+
+    Attributes
+    ----------
+    VERIFY_SREF_TREE : int = 1
+        Verify the correctness of the sref tree
+    VERIFY_CACHED_FLAGS : int = 2
+        Verify the correctness of affine_binding, region_cover and stage_pipeline
+    """
+
+    VERIFY_SREF_TREE = 1
+    VERIFY_CACHED_FLAGS = 2
+
+
+def _parse_mod(mod: PrimFunc | IRModule) -> IRModule:
+    if isinstance(mod, PrimFunc):
+        mod = IRModule({"main": mod})
+    if not isinstance(mod, IRModule):
+        raise TypeError(f"Expected `mod` to be PrimFunc or IRModule, but gets: {mod}")
+    return mod
+
+
+def _parse_debug_mask(debug_mask: str | int) -> int:
+    if isinstance(debug_mask, str):
+        if debug_mask == "all":
+            debug_mask = ScheduleDebugMask.VERIFY_SREF_TREE | ScheduleDebugMask.VERIFY_CACHED_FLAGS
+        elif debug_mask == "none":
+            debug_mask = 0
+        else:
+            raise ValueError(f"Unrecognizable `debug_mask`: {debug_mask}")
+    if not isinstance(debug_mask, bool) and not isinstance(debug_mask, int):
+        raise TypeError(f"`debug_mask` should be integer or boolean, but gets: {debug_mask}")
+    return debug_mask
+
+
+def _parse_enable_checks(enable_checks: bool) -> bool:
+    if not isinstance(enable_checks, bool):
+        raise TypeError(f"enable_checks only accepts bool value, got {type(enable_checks)} instead")
+    return enable_checks
+
+
+@register_object("s_tir.ScheduleState")
+class ScheduleState(Object):
+    """The state of scheduling, which exposes a `Replace` method as
+    the primary resort for all the scheduling primitives to manipulate the TensorIR.
+
+    The data structure contains the following information
+    1) The AST being scheduled (mod)
+    2) The sref tree of schedulable statements (indicated by the srefs)
+    3) The dependency information of each block scope (block_info)
+    4) A reverse mapping from the AST nodes to that in the sref tree (get_sref)
+    5) A debug flag, if set, extra checking is enabled (debug_mask)
+    6) A enable check flag, if False, some prerequisite checks are disabled.
+
+    Parameters
+    ----------
+    mod : IRModule
+        The AST of the module being scheduled
+    debug_mask : int
+        Do extra correctness checking after the object construction
+        and each time after calling the Replace method.
+    enable_check : bool
+        Indicates whether we enable prerequisite checks for some schedule primitives or not,
+        defaults to `True`.
+    """
+
+    mod: IRModule
+    debug_mask: int
+
+    def __init__(
+        self,
+        mod: PrimFunc | IRModule,
+        *,
+        debug_mask: str | int = "none",
+        enable_check: bool = True,
+    ) -> None:
+        """Construct a schedule state from an IRModule or a PrimFunc
+
+        Parameters
+        ----------
+        mod : Union[PrimFunc, IRModule]
+            The IRModule or PrimFunc to be scheduled
+        debug_mask : Union[str, int]
+            Do extra correctness checking after the class creation and each time
+            after calling the Replace method.
+            Possible choices of `debug_mask`:
+            1) "all" - Turn on all the checks
+            2) "none" - Turn off all the checks
+            3) An integer - Turn on checks according to the bitmasks provided in ScheduleDebugMask
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleState,  # type: ignore # pylint: disable=no-member
+            _parse_mod(mod),
+            _parse_debug_mask(debug_mask),
+            _parse_enable_checks(enable_check),
+        )
+
+    def get_sref(self, stmt: SBlock | For) -> StmtSRef | None:
+        """Return the corresponding sref that points to the stmt
+
+        Parameters
+        ----------
+        stmt : Union[Block, For]
+            The schedulable statement in the TensorIR to be retrieved for its sref
+
+        Returns
+        -------
+        sref : StmtSRef
+            The corresponding sref
+        """
+        return _ffi_api.ScheduleStateGetSRef(self, stmt)  # type: ignore # pylint: disable=no-member
+
+    def get_sblock_scope(self, block_sref: StmtSRef) -> SBlockScope:
+        """Get the SBlockScope correpsonding to the block sref
+
+        Parameters
+        ----------
+        block_sref : StmtSRef
+            The block sref to be retrieved
+
+        Returns
+        -------
+        sref : StmtSRef
+            The corresponding sref
+        """
+        return _ffi_api.ScheduleStateGetSBlockScope(  # type: ignore # pylint: disable=no-member
+            self, block_sref
+        )
+
+    def _get_cached_flags(self, block_sref: StmtSRef) -> CachedFlags:
+        """Get the cached flags of the corresponding block
+
+        Parameters
+        ----------
+        block_sref : StmtSRef
+            The block sref to be retrieved
+
+        Returns
+        -------
+        flags : CachedFlags
+            Three flags: affine_binding, region_cover, stage_pipeline
+
+        Note
+        ----
+        It is an API intended for internal testing use.
+        """
+        (
+            affine_binding,
+            region_cover,
+            stage_pipeline,
+        ) = _ffi_api.ScheduleStateGetCachedFlags(  # type: ignore # pylint: disable=no-member
+            self, block_sref
+        )
+        return CachedFlags(
+            affine_binding=bool(affine_binding.value),
+            region_cover=bool(region_cover.value),
+            stage_pipeline=bool(stage_pipeline.value),
+        )
+
+    def replace(
+        self,
+        src_sref: StmtSRef,
+        tgt_stmt: SBlock | For | SBlockRealize,
+        block_sref_reuse: dict[SBlock, SBlock] | None = None,
+    ) -> None:
+        """
+        Replace the part of the AST, as being pointed to by `src_sref`,
+        with a specific statement `tgt_stmt`, and maintain the sref tree accordingly.
+        Replace will try to perform copy on write as much as possible when the ScheduleState holds
+        the only copy to the IRModule and IR nodes.
+
+        Only 3 types of replacements are allowed: from `src_sref->stmt` to `tgt_stmt`.
+        1) SBlock -> SBlock
+        2) Loop -> Loop
+        3) Loop -> BlockRealize
+
+        Parameters
+        ----------
+        src_sref : StmtSRef
+            The sref to the statement to be replaced in the TensorIR AST
+
+        tgt_stmt : Union[Block, For, BlockRealize]
+            The statement to be replaced to
+
+        block_sref_reuse : Optional[Dict[Block, Block]] = None
+            Maps an old block (to be replaced in the subtree under `src_sref->stmt`)
+            to a new block (replaced to, in the subtree under `tgt_stmt`), and enforces
+            reuse of srefs between them (rather than create new srefs) i.e. after being replaced,
+            the sref that points to the old block will point to the new one
+
+        Note
+        ----
+        The reuse of loop srefs are detected automatically according to the reuse of loop vars.
+        """
+        if block_sref_reuse is None:
+            block_sref_reuse = {}
+        _ffi_api.ScheduleStateReplace(  # type: ignore # pylint: disable=no-member
+            self, src_sref, tgt_stmt, block_sref_reuse
+        )

python/tvm/s_tir/schedule/testing.py

@@ -0,0 +1,107 @@
+# 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=dangerous-default-value
+"""Testing utilities for the TensorIR schedule API"""
+
+from collections.abc import Sequence
+from typing import Any
+
+import tvm
+from tvm.ir import IRModule, assert_structural_equal
+from tvm.s_tir.schedule import Schedule, Trace
+from tvm.tirx import PrimFunc
+
+
+def assert_structural_equal_ignore_global_symbol(
+    func1: PrimFunc,
+    func2: PrimFunc,
+    *args: Any,
+    **kwargs: Any,
+) -> None:
+    """
+    Asserts that PrimFuncs func1 and func2 are structurally equal, setting both
+    their global symbol attributes to main so that the global symbol
+    will not be a point of comparison.
+    """
+    assert_structural_equal(
+        func1.with_attr("global_symbol", "main"),
+        func2.with_attr("global_symbol", "main"),
+        *args,
+        **kwargs,
+    )
+
+
+def verify_trace_roundtrip(
+    sch: Schedule,
+    mod: PrimFunc | IRModule,
+    *,
+    debug_mask: str | int = "all",
+    text_format: str | Sequence[str] = ["python", "json"],
+) -> Schedule:
+    """Serialize a traced schedule to JSON, then replay the JSON trace by applying to
+    a fresh new schedule, verifying the reproducibility of scheduling.
+
+    Parameters
+    ----------
+    sch : s_tir.Schedule
+        The traced TensorIR schedule to be verified
+    mod : Union[PrimFunc, IRModule]
+        The IRModule or PrimFunc to construct the fresh new schedule
+    debug_mask : Union[str, int]
+        Do extra correctness checking after the class creation and each time
+        after calling the Replace method.
+        Possible choices of `debug_mask`:
+        1) "all" - Turn on all the checks
+        2) "none" - Turn off all the checks
+        3) An integer - Turn on checks according to the bitmasks provided in ScheduleDebugMask
+    text_format: Union[str, Sequence[str]]
+        The text format or formats whose round-trip behavior should be
+        validated.  If a single string, validate round-trips through
+    """
+    from tvm.script import tirx as T  # pylint: disable=import-outside-toplevel
+
+    if not isinstance(text_format, str):
+        for opt in text_format:
+            new_sch = verify_trace_roundtrip(sch, mod, debug_mask=debug_mask, text_format=opt)
+        return new_sch
+
+    trace = sch.trace
+    assert trace is not None
+
+    # Step 1. Perform a round-trip through the text-format
+    new_sch = Schedule(mod=mod, debug_mask=debug_mask)
+    if text_format == "json":
+        json_obj = trace.as_json()
+        Trace.apply_json_to_schedule(json_obj=json_obj, sch=new_sch)
+    elif text_format == "python":
+        py_trace = "\n".join(trace.as_python())
+        vars_dict = {"T": T}
+        vars_dict.update(tvm.tirx.__dict__)
+        exec(py_trace, vars_dict, {"sch": new_sch})  # pylint: disable=exec-used
+    else:
+        assert text_format in ("json", "python"), f"Unknown text format: {text_format}"
+
+    # Step 2. Verify that the round-trip produced the same scheduling
+    assert_structural_equal(new_sch.mod, sch.mod)
+
+    # Step 3. Check the consistency of the text format between the old and new traces
+    py_repr = "\n".join(trace.as_python())
+    new_py_repr = "\n".join(new_sch.trace.as_python())
+    assert py_repr == new_py_repr
+
+    # Step 4. Return the new schedule in case it could be useful
+    return new_sch

python/tvm/s_tir/schedule/trace.py

@@ -0,0 +1,297 @@
+# 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.
+"""An execution trace of a scheduling program"""
+
+import os
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from tvm_ffi import Array, Map
+from tvm_ffi import register_object as _register_object
+
+from tvm.runtime import Object
+from tvm.tirx.expr import FloatImm, IntImm
+from tvm.tirx.function import IndexMap
+
+from ...ir import save_json
+from . import _ffi_api
+from .instruction import ATTR_TYPE, INPUT_RV_TYPE, Instruction
+
+if TYPE_CHECKING:
+    from .schedule import Schedule
+
+
+DECISION_TYPE = Any
+JSON_TYPE = Any
+
+
+def _json_from_tvm(obj):
+    if obj is None:
+        return None
+    elif isinstance(obj, bool | int | float | str):
+        return obj
+    elif isinstance(obj, Array):
+        return [_json_from_tvm(i) for i in obj]
+    elif isinstance(obj, Map):
+        return {_json_from_tvm(k): _json_from_tvm(v) for k, v in obj.items()}
+    elif isinstance(obj, str):
+        return str(obj)
+    elif isinstance(obj, IntImm | FloatImm):
+        return obj
+    elif isinstance(obj, IndexMap):
+        return save_json(obj)
+    else:
+        raise TypeError("Not supported type: " + str(type(obj)))
+
+
+@_register_object("s_tir.Trace")
+class Trace(Object):
+    """An execution trace of a scheduling program.
+
+    A trace has two parts:
+    1) The instructions invoked so far
+    2) The random decisions made upon those instructions, if any
+
+    A trace can be serialized to:
+    1) Roundtrippable JSON format: can be saved to file and loaded back
+    2) Python syntax: allows users to copy-paste the trace to reproduce the scheduling process
+
+    A trace can be applied to a TensorIR schedule by re-applying all its instructions possibly with
+    their decisions accordingly. Re-sampling is invoked if a sampling instruction doesn't have its
+    corresponding decision; Otherwise the existing decision will be reused accordingly.
+
+    Attributes
+    ----------
+    insts : List[Instruction]
+        The instructions invoked so far in the program execution
+    decisions : Dict[Instruction, DECISION_TYPE]
+        The random decisions made upon those instructions
+    """
+
+    insts: list[Instruction]
+    decisions: dict[Instruction, DECISION_TYPE]
+
+    def __init__(
+        self,
+        insts: list[Instruction],
+        decisions: dict[Instruction, DECISION_TYPE],
+    ) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        insts : List[Instruction]
+            The instructions invoked so far in the program execution
+        decisions : Dict[Instruction, DECISION_TYPE]
+            The random decisions made upon those instructions
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.Trace,  # type: ignore # pylint: disable=no-member
+            insts,
+            decisions,
+        )
+
+    def get_decision(self, inst: Instruction) -> DECISION_TYPE | None:
+        """Retrieve the decision made on a specific instruction
+
+        Parameters
+        ----------
+        insts : Instruction
+            The instruction whose decision is to be retrieved
+
+        Returns
+        -------
+        decision : Optional[DECISION_TYPE]
+            The corresponding decision; None if there is no decision made on the instruction
+        """
+        return _ffi_api.TraceGetDecision(self, inst)  # type: ignore # pylint: disable=no-member
+
+    def append(
+        self,
+        inst: Instruction,
+        decision: DECISION_TYPE | None = None,
+    ) -> None:
+        """Append a new instruction to the trace
+
+        Parameters
+        ----------
+        insts : Instruction
+            The new instruction to be appended
+        decision : Optional[DECISION_TYPE] = None
+            The random decision made on this instruction
+        """
+        _ffi_api.TraceAppend(self, inst, decision)  # type: ignore # pylint: disable=no-member
+
+    def pop(self) -> Instruction | None:
+        """Remove the last instruction, along with the decision made on that instruction, if any
+
+        Returns
+        -------
+        popped_inst : Instruction
+            Returns the instruction removed; std::nullopt if the trace is empty
+        """
+        return _ffi_api.TracePop(self)  # type: ignore # pylint: disable=no-member
+
+    def apply_to_schedule(
+        self,
+        sch: "Schedule",
+        remove_postproc: bool,
+        decision_provider: (
+            Callable[
+                [Instruction, list[INPUT_RV_TYPE], list[ATTR_TYPE], DECISION_TYPE], DECISION_TYPE
+            ]
+            | None
+        ) = None,
+    ) -> None:
+        """Apply the trace to a TensorIR schedule
+
+        Parameters
+        ----------
+        sch : Schedule
+            The schedule to be applied onto
+        remove_postproc : bool
+            If postprocessing instructions are removed
+        decision_provider: Optional[Callable] = None
+            A callback that allows users to mutate decisions on the fly when applying instructions.
+            The signature of the callback is:
+            - The 1st argument: The instruction
+            - The 2nd argument: The input random variables
+            - The 3rd argument: The attributes
+            - The 4th argument: The decision
+            - Return: A new decision
+        """
+        _ffi_api.TraceApplyToSchedule(  # type: ignore # pylint: disable=no-member
+            self,
+            sch,
+            remove_postproc,
+            decision_provider,
+        )
+
+    def as_json(self, remove_postproc: bool = False) -> JSON_TYPE:
+        """Serialize the trace as a JSON-style object
+
+        Parameters
+        ----------
+        remove_postproc : bool = False
+            If postprocessing instructions are removed
+
+        Returns
+        -------
+        json: JSON_TYPE
+            The JSON-style object
+        """
+        obj = _ffi_api.TraceAsJSON(self, remove_postproc)  # type: ignore # pylint: disable=no-member
+        return _json_from_tvm(obj)
+
+    def as_python(self, remove_postproc: bool = False) -> list[str]:
+        """Serialize the trace as a sequence of python statements
+
+        Parameters
+        ----------
+        remove_postproc : bool = False
+            If postprocessing instructions are removed
+
+        Returns
+        -------
+        py_stmts: List[str]
+            A sequence of python statements
+        """
+        return _ffi_api.TraceAsPython(self, remove_postproc)  # type: ignore # pylint: disable=no-member
+
+    def with_decision(
+        self,
+        inst: Instruction,
+        decision: DECISION_TYPE,
+        remove_postproc: bool,
+    ) -> "Trace":
+        """Create a new trace with an instruction whose decision is changed,
+        assuming this instruction exists in the resulting trace
+
+        Parameters
+        ----------
+        inst : Instruction
+            The instruction whose decision is to be changed
+        decision : DECISION_TYPE
+            The decision to be changed to
+        remove_postproc : bool
+            If postprocessing instructions are removed
+
+        Returns
+        -------
+        trace: Trace
+            The new trace with the decision changed
+        """
+        return _ffi_api.TraceWithDecision(  # type: ignore # pylint: disable=no-member
+            self,
+            inst,
+            decision,
+            remove_postproc,
+        )
+
+    def simplified(self, remove_postproc: bool) -> "Trace":
+        """Simplify the trace with dead-code elimination
+
+        Parameters
+        ----------
+        remove_postproc : bool
+            If postprocessing instructions are removed
+
+        Returns
+        -------
+        trace: Trace
+            A simplified trace
+        """
+        return _ffi_api.TraceSimplified(self, remove_postproc)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def apply_json_to_schedule(json_obj: JSON_TYPE, sch: "Schedule") -> None:
+        """Apply a JSON-serialized trace to a TensorIR schedule
+
+        Parameters
+        ----------
+        json_obj : JSON_TYPE
+            The JSON-serialized trace
+        sch : Schedule
+            The TensorIR schedule
+        """
+        _ffi_api.TraceApplyJSONToSchedule(json_obj, sch)  # type: ignore # pylint: disable=no-member
+
+    def show(self, style: str | None = None, black_format: bool = False) -> None:
+        """A sugar for print highlighted TVM script.
+
+        Parameters
+        ----------
+        style : str, optional
+
+            Pygmentize printing style, auto-detected if None.  See
+            `tvm.script.highlight.cprint` for more details.
+
+        black_format: bool
+
+            If true, use the formatter Black to format the TVMScript.
+            If None, determine based on the "TVM_BLACK_FORMAT" environment
+            variable.
+        """
+        from tvm.script.highlight import (  # pylint: disable=import-outside-toplevel
+            cprint,
+        )
+
+        if black_format is None:
+            env = os.environ.get("TVM_BLACK_FORMAT")
+            black_format = bool(env and int(env))
+
+        cprint(str(self), style=style, black_format=black_format)

python/tvm/s_tir/schedule/transform.py

@@ -0,0 +1,46 @@
+# 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.
+"""Transformation on TIR schedule."""
+
+from tvm.s_tir.schedule import LoopRV, SBlockRV, Schedule
+
+from . import _ffi_api
+
+
+def tile_with_tensor_intrin(
+    sch: Schedule, block: SBlockRV, intrin_name: str, allow_padding: bool = False
+) -> LoopRV | None:
+    """Tile a subset of loops in the block according to the given tensor intrinsic.
+
+    Parameters
+    ----------
+    sch : Schedule
+        The schedule to which tiling is applied
+    block : SBlockRV
+        The block whose subset of loops will be tiled
+    intrin_name : str
+        The name of a tensor intrinsic, must be registerd via TensorIntrin.register(...) beforehand
+    allow_padding : bool
+        Whether to allow padding when tiling
+
+    Returns
+    -------
+    tiled_loop_rv : Optional[LoopRV]
+        LoopRV corresponding to the outermost loop of a block tiled according to the given intrin
+        std::nullopt if no valid loop mapping is found
+    """
+    return _ffi_api.TileWithTensorIntrin(sch, block, intrin_name, allow_padding)  # type: ignore