chore: import upstream snapshot with attribution

python/tvm/s_tir/meta_schedule/__init__.py

@@ -0,0 +1,56 @@
+# 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.
+"""Package `tvm.s_tir.meta_schedule`. The meta schedule infrastructure."""
+
+from . import (
+    arg_info,
+    builder,
+    cost_model,
+    database,
+    feature_extractor,
+    measure_callback,
+    mutator,
+    postproc,
+    relax_integration,
+    runner,
+    schedule,
+    schedule_rule,
+    search_strategy,
+    space_generator,
+    tir_integration,
+    trace_apply,
+    post_optimization,
+)
+from .builder import Builder
+from .cost_model import CostModel
+from .database import Database
+from .extracted_task import ExtractedTask
+from .feature_extractor import FeatureExtractor
+from .measure_callback import MeasureCallback
+from .mutator import Mutator
+from .postproc import Postproc
+from .profiler import Profiler
+from .runner import Runner
+from .schedule_rule import ScheduleRule
+from .search_strategy import MeasureCandidate, SearchStrategy
+from .space_generator import SpaceGenerator
+from .task_scheduler import TaskScheduler
+from .tir_integration import tune_tir
+from .tune import tune_tasks
+from .tune_context import TuneContext
+from .post_optimization import post_opt

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

python/tvm/s_tir/meta_schedule/arg_info.py

@@ -0,0 +1,124 @@
+# 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.
+"""The argument information"""
+
+from typing import Any
+
+from tvm_ffi import Shape, register_object
+
+from tvm.ir import IRModule
+from tvm.runtime import DataType, Object
+from tvm.tirx import PrimFunc
+
+from . import _ffi_api
+from .utils import _json_de_tvm
+
+
+@register_object("s_tir.meta_schedule.ArgInfo")
+class ArgInfo(Object):
+    """Argument information"""
+
+    def as_json(self) -> Any:
+        """Converts the ArgInfo to its corresponding JSON representation."""
+        return _json_de_tvm(_ffi_api.ArgInfoAsJSON(self))  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def from_json(json_obj: Any) -> "ArgInfo":
+        """Parse the argument information from a JSON object.
+
+        Parameters
+        ----------
+        json_obj : Any
+            The json object to parse.
+
+        Returns
+        -------
+        parsed : ArgInfo
+            The argument information parsed.
+        """
+        return _ffi_api.ArgInfoFromJSON(json_obj)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def from_prim_func(func: PrimFunc) -> list["ArgInfo"]:
+        """Extract a list of the argument information from PrimFunc.
+
+        Parameters
+        ----------
+        func : PrimFunc
+            The PrimFunc to get argument information from.
+
+        Returns
+        -------
+        extracted : List[ArgInfo]
+            An array of the argument information derived.
+        """
+        return _ffi_api.ArgInfoFromPrimFunc(func)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def from_entry_func(mod: IRModule, remove_preproc: bool = True) -> list["ArgInfo"]:
+        """Extract a list of the argument information from the entry func of an IRModule.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to get argument information from.
+        remove_preproc : bool
+            Whether to remove the preprocessing blocks.
+
+        Returns
+        -------
+        extracted : List[ArgInfo]
+            An array of the argument information derived.
+        """
+        return _ffi_api.ArgInfoFromEntryFunc(mod, remove_preproc)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("s_tir.meta_schedule.TensorInfo")
+class TensorInfo(ArgInfo):
+    """Tensor argument information
+
+    Parameters
+    ----------
+    dtype : DataType
+        The data type of the tensor.
+    shape : Shape
+        The shape of the tensor.
+    """
+
+    dtype: DataType
+    shape: Shape
+
+    def __init__(
+        self,
+        dtype: DataType,
+        shape: Shape | list[int],
+    ) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        dtype : DataType
+            The data type of the tensor.
+        shape : Shape
+            The shape of the tensor.
+        """
+        shape_tuple = shape if isinstance(shape, Shape) else Shape(shape)
+        self.__init_handle_by_constructor__(
+            _ffi_api.TensorInfo,  # type: ignore # pylint: disable=no-member
+            dtype,
+            shape_tuple,
+        )

python/tvm/s_tir/meta_schedule/builder/__init__.py

@@ -0,0 +1,25 @@
+# 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.
+"""
+The tvm.s_tir.meta_schedule.builder package.
+Meta Schedule builders that translate IRModule to runtime.Module,
+and then export
+"""
+
+from .builder import Builder, BuilderInput, BuilderResult, PyBuilder, create
+from .local_builder import LocalBuilder

python/tvm/s_tir/meta_schedule/builder/builder.py

@@ -0,0 +1,203 @@
+# 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.
+# ruff: noqa: RUF012
+"""Meta Schedule builders that translate IRModule to runtime.Module, and then export"""
+
+from collections.abc import Callable
+from typing import Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+from tvm_ffi import register_object
+
+from tvm.ir import IRModule
+from tvm.runtime import Object, Tensor
+from tvm.target import Target
+
+from .. import _ffi_api
+
+
+@register_object("s_tir.meta_schedule.BuilderInput")
+class BuilderInput(Object):
+    """The builder's input.
+
+    Parameters
+    ----------
+    mod : IRModule
+        The IRModule to be built.
+    target : Target
+        The target to be built for.
+    params: Optional[Dict[str, Tensor]]
+        The parameters for Relax build module
+    """
+
+    mod: IRModule
+    target: Target
+    params: dict[str, Tensor] | None
+
+    def __init__(
+        self,
+        mod: IRModule,
+        target: Target,
+        params: dict[str, Tensor] | None = None,
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be built.
+        target : Target
+            The target to be built for.
+        params: Optional[Dict[str, Tensor]]
+            The parameters for Relax build module
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.BuilderInput,  # type: ignore # pylint: disable=no-member
+            mod,
+            target,
+            params,
+        )
+
+
+@register_object("s_tir.meta_schedule.BuilderResult")
+class BuilderResult(Object):
+    """The builder's result.
+
+    Parameters
+    ----------
+    artifact_path : Optional[str]
+        The path to the artifact.
+    error_msg : Optional[str]
+        The error message.
+    """
+
+    artifact_path: str | None
+    error_msg: str | None
+
+    def __init__(
+        self,
+        artifact_path: str | None,
+        error_msg: str | None,
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        artifact_path : Optional[str]
+            The path to the artifact.
+        error_msg : Optional[str]
+            The error message.
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.BuilderResult,  # type: ignore # pylint: disable=no-member
+            artifact_path,
+            error_msg,
+        )
+
+
+@register_object("s_tir.meta_schedule.Builder")
+class Builder(Object):
+    """The abstract builder interface."""
+
+    BuilderType = Union["Builder", Literal["local"]]
+
+    def build(self, build_inputs: list[BuilderInput]) -> list[BuilderResult]:
+        """Build the given inputs.
+
+        Parameters
+        ----------
+        build_inputs : List[BuilderInput]
+            The inputs to be built.
+        Returns
+        -------
+        build_results : List[BuilderResult]
+            The results of building the given inputs.
+        """
+        return _ffi_api.BuilderBuild(self, build_inputs)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def create(  # pylint: disable=keyword-arg-before-vararg
+        kind: Literal["local"] = "local",
+        *args,
+        **kwargs,
+    ) -> "Builder":
+        """Create a Builder.
+
+        Parameters
+        ----------
+        kind : Literal["local"]
+            The kind of the builder. For now, only "local" is supported.
+
+        Returns
+        -------
+        builder : Builder
+            The builder created.
+        """
+        from . import LocalBuilder  # pylint: disable=import-outside-toplevel
+
+        if kind == "local":
+            return LocalBuilder(*args, **kwargs)  # type: ignore
+        raise ValueError(f"Unknown Builder: {kind}")
+
+
+create = Builder.create  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.PyBuilder")
+class _PyBuilder(Builder):
+    """
+    A TVM object builder to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyBuilder
+    """
+
+    def __init__(self, f_build: Callable | None = None):
+        """Constructor."""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.BuilderPyBuilder,  # type: ignore # pylint: disable=no-member
+            f_build,
+        )
+
+
+class PyBuilder:
+    """
+    An abstract builder with customized build method on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {"cls": _PyBuilder, "methods": ["build"]}
+
+    def build(self, build_inputs: list[BuilderInput]) -> list[BuilderResult]:
+        """Build the given inputs.
+
+        Parameters
+        ----------
+        build_inputs : List[BuilderInput]
+            The inputs to be built.
+        Returns
+        -------
+        build_results : List[BuilderResult]
+            The results of building the given inputs.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/builder/local_builder.py

@@ -0,0 +1,300 @@
+# 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.
+# ruff: noqa: F401
+"""Local builder that compile on the local host"""
+
+import os
+import tempfile
+from collections.abc import Callable
+from typing import Optional, Union
+
+from tvm_ffi import register_global_func
+
+from tvm.ir import IRModule
+from tvm.ir.utils import derived_object
+from tvm.runtime import Module, Tensor, load_param_dict, save_param_dict
+from tvm.support.popen_pool import MapResult, PopenPoolExecutor, StatusKind
+from tvm.target import Target
+
+from ..logging import get_logger
+from ..utils import cpu_count, get_global_func_with_default_on_worker
+from .builder import BuilderInput, BuilderResult, PyBuilder
+
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
+
+T_BUILD = Callable[  # pylint: disable=invalid-name
+    [IRModule, Target, dict[str, Tensor] | None], Module
+]
+T_EXPORT = Callable[[Module], str]  # pylint: disable=invalid-name
+
+
+def _serialize_params(params: dict[str, Tensor] | None) -> bytearray | None:
+    if params is None:
+        return None
+    return save_param_dict(params)
+
+
+def _deserialize_params(params: bytearray | None) -> dict[str, Tensor] | None:
+    if params is None:
+        return None
+    return load_param_dict(params)
+
+
+@derived_object
+class LocalBuilder(PyBuilder):
+    """A builder that builds the given input on local host.
+
+    Parameters
+    ----------
+    pool : PopenPoolExecutor
+        The process pool to run the build.
+    max_workers: int
+        The max number of Popen workers.
+    timeout_sec : float
+        The timeout in seconds for the build.
+    initializer: Optional[Callable[[], None]]
+        The initializer function for each popen worker.
+    f_build : Union[None, str, T_BUILD]
+        Name of the build function to be used.
+        Defaults to `meta_schedule.builder.default_build`.
+    f_export : Union[None, str, T_EXPORT]
+        Name of the export function to be used.
+        Defaults to `meta_schedule.builder.default_export`.
+
+    Attributes
+    ----------
+    T_BUILD : typing._GenericAlias
+        The signature of the function `f_build`, which is
+
+        .. code-block:: python
+
+        def default_build(
+            mod: IRModule,
+            target: Target,
+            params: Optional[Dict[str, Tensor]]
+        ) -> Module:
+            ...
+
+    T_EXPORT : typing._GenericAlias
+        The signature of the function `f_export`, which is
+
+        .. code-block:: python
+
+        def default_export(mod: Module) -> str:
+            ...
+
+    Note
+    ----
+    The build function and export function should be registered in the worker process.
+    The worker process is only aware of functions registered in TVM package,
+    if there are extra functions to be registered,
+    please send the registration logic via initializer.
+    """
+
+    max_workers: int
+    timeout_sec: float
+    initializer: Callable[[], None] | None
+    f_build: None | str | T_BUILD
+    f_export: None | str | T_EXPORT
+
+    def __init__(
+        self,
+        *,
+        max_workers: int | None = None,
+        timeout_sec: float = 30.0,
+        f_build: None | str | T_BUILD = None,
+        f_export: None | str | T_EXPORT = None,
+        initializer: Callable[[], None] | None = None,
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        max_workers : Optional[int]
+            The maximum number of worker processes to be used.
+            Defaults to number of CPUs.
+        timeout_sec : float
+            The timeout in seconds for the build.
+        f_build : T_BUILD
+            Name of the build function to be used.
+            Defaults to `meta_schedule.builder.default_build`.
+        f_export : T_EXPORT
+            Name of the export function to be used.
+            Defaults to `meta_schedule.builder.default_export`.
+        initializer : Optional[Callable[[], None]]
+            The initializer to be used for the worker processes.
+        """
+        super().__init__()
+
+        if max_workers is None:
+            max_workers = cpu_count(logical=True)
+        logger.info("LocalBuilder: max_workers = %d", max_workers)
+
+        self.max_workers = max_workers
+        self.timeout_sec = timeout_sec
+        self.initializer = initializer
+        self.f_build = f_build
+        self.f_export = f_export
+        self._sanity_check()
+
+    def build(self, build_inputs: list[BuilderInput]) -> list[BuilderResult]:
+        results: list[BuilderResult] = []
+        map_result: MapResult
+
+        # Here we restart the PopenPool everytime because of a known memory leak issue with the
+        # PopenPool workers after a couple times of usage. We don't apply the same to runners to
+        # avoid potential problem caused by async behaviour.
+        pool = PopenPoolExecutor(
+            max_workers=self.max_workers,
+            timeout=self.timeout_sec,
+            initializer=self.initializer,
+        )
+
+        # Dispatch the build inputs to the worker processes.
+        for map_result in pool.map_with_error_catching(
+            lambda x: _worker_func(*x),
+            [
+                (
+                    self.f_build,
+                    self.f_export,
+                    build_input.mod,
+                    build_input.target,
+                    _serialize_params(build_input.params),
+                )
+                for build_input in build_inputs
+            ],
+        ):
+            if map_result.status == StatusKind.COMPLETE:
+                results.append(BuilderResult(map_result.value, None))
+            elif map_result.status == StatusKind.TIMEOUT:
+                results.append(
+                    BuilderResult(
+                        None,
+                        f"LocalBuilder: Timeout, killed after {self.timeout_sec} seconds",
+                    )
+                )
+            elif map_result.status == StatusKind.EXCEPTION:
+                results.append(
+                    BuilderResult(
+                        None,
+                        "LocalBuilder: An exception occurred\n" + str(map_result.value),
+                    )
+                )
+            else:
+                raise ValueError("Unreachable: unexpected result: {map_result}")
+        pool.shutdown()
+        return results
+
+    def _sanity_check(self) -> None:
+        def _check(f_build, f_export) -> None:
+            get_global_func_with_default_on_worker(name=f_build, default=None)
+            get_global_func_with_default_on_worker(name=f_export, default=None)
+
+        # Same reason for the single use PopenPool as mentioned above
+        pool = PopenPoolExecutor(
+            max_workers=self.max_workers,
+            timeout=self.timeout_sec,
+            initializer=self.initializer,
+        )
+        value = pool.submit(_check, self.f_build, self.f_export)
+        value.result()
+        pool.shutdown()
+
+
+def _worker_func(
+    _f_build: None | str | T_BUILD,
+    _f_export: None | str | T_EXPORT,
+    mod: IRModule,
+    target: Target,
+    params: bytearray | None,
+) -> str:
+    # Step 0. Get the registered functions
+    f_build: T_BUILD = get_global_func_with_default_on_worker(
+        _f_build,
+        default_build,
+    )
+    f_export: T_EXPORT = get_global_func_with_default_on_worker(
+        _f_export,
+        default_export,
+    )
+    # Step 1. Build the IRModule
+    rt_mod: Module = f_build(mod, target, _deserialize_params(params))
+    # Step 2. Export the Module
+    artifact_path: str = f_export(rt_mod)
+    return artifact_path
+
+
+@register_global_func("s_tir.meta_schedule.builder.default_build")
+def default_build(mod: IRModule, target: Target, _params: dict[str, Tensor] | None) -> Module:
+    """Default build function.
+
+    Parameters
+    ----------
+    mod : IRModule
+        The IRModule to be built.
+    target : Target
+        The target to be built.
+    _params : Optional[Dict[str, Tensor]]
+        The parameters to be used for the build. Must be None.
+
+    Returns
+    -------
+    rt_mod : Module
+        The built Module.
+    """
+    # pylint: disable=import-outside-toplevel
+    import tvm.s_tir.tensor_intrin  # pylint: disable=unused-import
+    from tvm.driver import build as tvm_build
+    from tvm.s_tir.transform import RemoveWeightLayoutRewriteBlock
+
+    # pylint: enable=import-outside-toplevel
+    mod = RemoveWeightLayoutRewriteBlock(skip_tensor_rewrite=True)(mod)
+    return tvm_build(mod, target=target)
+
+
+@register_global_func("s_tir.meta_schedule.builder.default_export")
+def default_export(mod: Module) -> str:
+    """Default export function.
+
+    Parameters
+    ----------
+    mod : Module
+        The Module to be exported.
+
+    Returns
+    -------
+    artifact_path : str
+        The path to the exported Module.
+    """
+    from tvm.support.tar import tar  # pylint: disable=import-outside-toplevel
+
+    artifact_path = os.path.join(tempfile.mkdtemp(), "tvm_tmp_mod." + tar.output_format)
+    mod.export_library(artifact_path, fcompile=tar)
+    return artifact_path
+
+
+@register_global_func("s_tir.meta_schedule.builder.get_local_builder")
+def get_local_builder() -> LocalBuilder:
+    """Get the local builder.
+
+    Returns
+    -------
+    builder : LocalBuilder
+        The local builder.
+    """
+    return LocalBuilder()

python/tvm/s_tir/meta_schedule/cost_model/__init__.py

@@ -0,0 +1,24 @@
+# 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.
+"""
+The tvm.s_tir.meta_schedule.cost_model package.
+"""
+
+from .cost_model import CostModel, PyCostModel
+from .random_model import RandomModel
+from .xgb_model import XGBModel

python/tvm/s_tir/meta_schedule/cost_model/cost_model.py

@@ -0,0 +1,260 @@
+# 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.
+# ruff: noqa: RUF012
+"""Meta Schedule CostModel."""
+
+import ctypes
+from collections.abc import Callable
+from typing import Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+import numpy as np  # type: ignore
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+
+from .. import _ffi_api
+from ..runner import RunnerResult
+from ..search_strategy import MeasureCandidate
+from ..tune_context import TuneContext
+
+
+@register_object("s_tir.meta_schedule.CostModel")
+class CostModel(Object):
+    """Cost model."""
+
+    CostModelType = Union["CostModel", Literal["xgb", "mlp", "random"]]
+
+    def load(self, path: str) -> None:
+        """Load the cost model from given file location.
+
+        Parameters
+        ----------
+        path : str
+            The file path.
+        """
+        _ffi_api.CostModelLoad(self, path)  # type: ignore # pylint: disable=no-member
+
+    def save(self, path: str) -> None:
+        """Save the cost model to given file location.
+
+        Parameters
+        ----------
+        path : str
+            The file path.
+        """
+        _ffi_api.CostModelSave(self, path)  # type: ignore # pylint: disable=no-member
+
+    def update(
+        self,
+        context: TuneContext,
+        candidates: list[MeasureCandidate],
+        results: list[RunnerResult],
+    ) -> None:
+        """Update the cost model given running results.
+
+        Parameters
+        ----------
+        context : TuneContext,
+            The tuning context.
+        candidates : List[MeasureCandidate]
+            The measure candidates.
+        results : List[RunnerResult]
+            The running results of the measure candidates.
+        """
+        _ffi_api.CostModelUpdate(self, context, candidates, results)  # type: ignore # pylint: disable=no-member
+
+    def predict(self, context: TuneContext, candidates: list[MeasureCandidate]) -> np.ndarray:
+        """Predict normalized score with the cost model.
+
+        Parameters
+        ----------
+        context : TuneContext,
+            The tuning context.
+        candidates : List[MeasureCandidate]
+            The measure candidates.
+
+        Return
+        ------
+        result : np.ndarray
+            The predicted normalized score.
+        """
+        n = len(candidates)
+        results = np.zeros(shape=(n,), dtype="float64")
+        _ffi_api.CostModelPredict(  # type: ignore # pylint: disable=no-member
+            self,
+            context,
+            candidates,
+            results.ctypes.data_as(ctypes.c_void_p),
+        )
+        return results
+
+    @staticmethod
+    def create(
+        kind: Literal["xgb", "mlp", "random", "none"],
+        *args,
+        **kwargs,
+    ) -> "CostModel":
+        """Create a CostModel.
+
+        Parameters
+        ----------
+        kind : Literal["xgb", "mlp", "random", "none"]
+            The kind of the cost model. Can be "xgb", "mlp", "random" or "none".
+
+        Returns
+        -------
+        cost_model : CostModel
+            The created cost model.
+        """
+        from . import RandomModel, XGBModel  # pylint: disable=import-outside-toplevel
+
+        if kind == "xgb":
+            return XGBModel(*args, **kwargs)  # type: ignore
+
+        # params only relevant to XGBModel
+        _xgb_params = ["num_tuning_cores", "tree_method"]
+
+        for param in _xgb_params:
+            if param in kwargs:
+                kwargs.pop(param)
+
+        if kind == "random":
+            return RandomModel(*args, **kwargs)  # type: ignore
+        if kind == "mlp":
+            from .mlp_model import (  # type: ignore  # pylint: disable=import-outside-toplevel
+                MLPModel,
+            )
+
+            return MLPModel(*args, **kwargs)  # type: ignore
+        if kind == "none":
+            return None  # no cost model required
+        raise ValueError(f"Unknown CostModel: {kind}")
+
+
+create = CostModel.create  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.PyCostModel")
+class _PyCostModel(CostModel):
+    """
+    A TVM object cost model to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyCostModel
+    """
+
+    def __init__(
+        self,
+        f_load: Callable | None = None,
+        f_save: Callable | None = None,
+        f_update: Callable | None = None,
+        predict_func: Callable | None = None,
+    ):
+        """Constructor."""
+
+        def f_predict(context: TuneContext, candidates: list[MeasureCandidate], return_ptr) -> None:
+            n = len(candidates)
+            return_ptr = ctypes.cast(return_ptr, ctypes.POINTER(ctypes.c_double))
+            array_wrapper = np.ctypeslib.as_array(return_ptr, shape=(n,))
+            res = predict_func(context, candidates)
+            array_wrapper[:] = res
+            assert array_wrapper.dtype == "float64", (
+                "ValueError: Invalid data type returned from CostModel Predict!"
+            )
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.CostModelPyCostModel,  # type: ignore # pylint: disable=no-member
+            f_load,
+            f_save,
+            f_update,
+            f_predict,
+        )
+
+
+class PyCostModel:
+    """
+    An abstract cost model with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PyCostModel,
+        "methods": ["load", "save", "update", "predict"],
+    }
+
+    def load(self, path: str) -> None:
+        """Load the cost model from given file location.
+
+        Parameters
+        ----------
+        path : str
+            The file path.
+        """
+        raise NotImplementedError
+
+    def save(self, path: str) -> None:
+        """Save the cost model to given file location.
+
+        Parameters
+        ----------
+        path : str
+            The file path.
+        """
+        raise NotImplementedError
+
+    def update(
+        self,
+        context: TuneContext,
+        candidates: list[MeasureCandidate],
+        results: list[RunnerResult],
+    ) -> None:
+        """Update the cost model given running results.
+
+        Parameters
+        ----------
+        context : TuneContext,
+            The tuning context.
+        candidates : List[MeasureCandidate]
+            The measure candidates.
+        results : List[RunnerResult]
+            The running results of the measure candidates.
+        """
+        raise NotImplementedError
+
+    def predict(self, context: TuneContext, candidates: list[MeasureCandidate]) -> np.ndarray:
+        """Predict given the measure candidates.
+
+        Parameters
+        ----------
+        context : TuneContext,
+            The tuning context.
+        candidates : List[MeasureCandidate]
+            The measure candidates.
+
+        Return
+        ------
+        result : np.ndarray
+            The predicted normalized score.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/cost_model/metric.py

@@ -0,0 +1,40 @@
+# 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.
+"""Cost model metrics for meta schedule"""
+
+import numpy as np  # type: ignore
+
+
+def max_curve(trial_scores: np.ndarray) -> np.ndarray:
+    """f(n) = max([s[i] fo i < n])
+
+    Parameters
+    ----------
+    trial_scores : List[float]
+        the score of i-th trial
+
+    Returns
+    -------
+    curve : np.ndarray
+        A vector, the max-curve function values
+    """
+    ret = np.empty(len(trial_scores))
+    keep = -1e9
+    for i, score in enumerate(trial_scores):
+        keep = max(keep, score)
+        ret[i] = keep
+    return ret

python/tvm/s_tir/meta_schedule/cost_model/random_model.py

@@ -0,0 +1,133 @@
+# 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.
+"""
+Random cost model
+"""
+
+from tvm.ir.utils import derived_object
+
+from ..cost_model import PyCostModel
+from ..runner import RunnerResult
+from ..search_strategy import MeasureCandidate
+from ..tune_context import TuneContext
+
+
+@derived_object
+class RandomModel(PyCostModel):
+    """Random cost model
+
+    Parameters
+    ----------
+    random_state : Union[Tuple[str, np.ndarray, int, int, float], dict]
+        The random state of the random number generator.
+    path : Optional[str]
+        The path of the random cost model.
+    max_range : Optional[int]
+        The maximum range of random results, [0, max_range].
+
+    Reference
+    ---------
+    https://numpy.org/doc/stable/reference/random/generated/numpy.random.get_state.html
+    """
+
+    import numpy as np  # type: ignore # pylint: disable=import-outside-toplevel
+
+    random_state: tuple[str, np.ndarray, int, int, float] | dict
+    path: str | None
+
+    def __init__(
+        self,
+        *,
+        seed: int | None = None,
+        path: str | None = None,
+        max_range: int | None = 100,
+    ):
+        import numpy as np  # type: ignore # pylint: disable=import-outside-toplevel
+
+        super().__init__()
+        if path is not None:
+            self.load(path)
+        else:
+            np.random.seed(seed)
+            self.random_state = np.random.get_state()
+        self.max_range = max_range
+
+    def load(self, path: str) -> None:
+        """Load the cost model from given file location.
+
+        Parameters
+        ----------
+        path : str
+            The file path.
+        """
+        import numpy as np  # type: ignore # pylint: disable=import-outside-toplevel
+
+        self.random_state = tuple(np.load(path, allow_pickle=True))  # type: ignore
+
+    def save(self, path: str) -> None:
+        """Save the cost model to given file location.
+
+        Parameters
+        ----------
+        path : str
+            The file path.
+        """
+        import numpy as np  # type: ignore # pylint: disable=import-outside-toplevel
+
+        np.save(path, np.array(self.random_state, dtype=object), allow_pickle=True)
+
+    def update(
+        self,
+        context: TuneContext,
+        candidates: list[MeasureCandidate],
+        results: list[RunnerResult],
+    ) -> None:
+        """Update the cost model given running results.
+
+        Parameters
+        ----------
+        context : TuneContext,
+            The tuning context.
+        candidates : List[MeasureCandidate]
+            The measure candidates.
+        results : List[RunnerResult]
+            The running results of the measure candidates.
+        """
+
+    def predict(self, context: TuneContext, candidates: list[MeasureCandidate]) -> np.ndarray:  # type: ignore # pylint: disable=used-before-assignment
+        """Update the cost model given running results.
+
+        Parameters
+        ----------
+        context : TuneContext,
+            The tuning context.
+        candidates : List[MeasureCandidate]
+            The measure candidates.
+
+        Return
+        ------
+        result : np.ndarray
+            The predicted running results.
+        """
+        import numpy as np  # type: ignore # pylint: disable=import-outside-toplevel
+
+        np.random.set_state(self.random_state)
+        # TODO(@zxybazh): Use numpy's RandState object:
+        # https://numpy.org/doc/1.16/reference/generated/numpy.random.RandomState.html#numpy.random.RandomState
+        result = np.random.rand(len(candidates)) * self.max_range  # type: ignore
+        self.random_state = np.random.get_state()
+        return result

python/tvm/s_tir/meta_schedule/cost_model/xgb_model.py

@@ -0,0 +1,867 @@
+# 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.
+"""XGBoost-based cost model"""
+
+import os
+import tempfile
+from collections import OrderedDict
+from collections.abc import Callable
+from itertools import chain as itertools_chain
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Optional
+
+import numpy as np  # type: ignore
+
+from tvm.ir.utils import derived_object
+from tvm.support.tar import tar, untar
+
+from ....runtime import Tensor
+from ..cost_model import PyCostModel
+from ..feature_extractor import FeatureExtractor
+from ..logging import get_logger
+from ..runner import RunnerResult
+from ..search_strategy import MeasureCandidate
+from ..utils import cpu_count, shash2hex
+from .metric import max_curve
+
+if TYPE_CHECKING:
+    import xgboost as xgb  # type: ignore
+    from xgboost.callback import TrainingCallback  # type: ignore
+
+    from ..tune_context import TuneContext
+
+
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def make_metric_sorter(focused_metric):
+    """Make sure the focused metric is the first one."""
+
+    def metric_name_for_sort(name):
+        if focused_metric == name:
+            return "!" + name
+        return name
+
+    def sort_key(key):
+        key, _ = key
+        return metric_name_for_sort(key)
+
+    return sort_key
+
+
+class PackSum:
+    """The pack-sum format
+
+    Parameters
+    ----------
+    dmatrix : xgb.DMatrix
+        A float64 array of shape [n, m],
+        where `n` is the packed number of blocks,
+        and `m` is the length of feature vector on each block
+    ids : np.ndarray
+        An int64 array of shape [n] containing nonnegative integers,
+        indicating which the index of a sample that a block belongs to
+    """
+
+    dmatrix: "xgb.DMatrix"  # type: ignore # pylint: disable=invalid-name
+    ids: np.ndarray
+
+    def __init__(
+        self,
+        xs: list[np.ndarray],  # pylint: disable=invalid-name
+        ys: np.ndarray | None,  # pylint: disable=invalid-name
+    ):
+        """Create PackSum format given a batch of samples
+
+        Parameters
+        ----------
+        xs : List[np.ndarray]
+            A batch of input samples
+        ys : Optional[List[float]]
+            A batch of labels. None means no labels available.
+        """
+        import xgboost as xgb  # type: ignore # pylint: disable=import-outside-toplevel
+
+        repeats = [x.shape[0] for x in xs]
+        xs = np.concatenate(xs, axis=0)
+        self.ids = np.concatenate([[i] * repeat for i, repeat in enumerate(repeats)], axis=0)
+        if ys is None:
+            self.dmatrix = xgb.DMatrix(data=xs, label=None)
+        else:
+            ys = np.concatenate([[y] * repeat for y, repeat in zip(ys, repeats)], axis=0)
+            self.dmatrix = xgb.DMatrix(data=xs, label=ys)
+            self.dmatrix.set_weight(ys)
+
+    def predict_with_score(self, pred: np.ndarray) -> np.ndarray:
+        """Predict the labels given the block level prediction scores.
+
+        Parameters
+        ----------
+        pred : np.ndarray
+            The block level predictions
+
+        Returns
+        -------
+        result : np.ndarray
+            The predictions for each candidate.
+        """
+        return np.bincount(self.ids, weights=pred)
+
+    def obj_square_error(self, ys_pred: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        """Implement square error loss on pack-sum format as
+        a custom objective function for xgboost.
+
+        Parameters
+        ----------
+        ys_pred: np.ndarray
+            The predictions
+
+        Returns
+        -------
+        gradient: np.ndarray
+            The gradient according to the xgboost format
+        hessian: np.ndarray
+            The hessian according to the xgboost format
+        """
+        # Making prediction
+        ys_pred = self.predict_with_score(ys_pred)
+        # Propagate prediction to each block
+        ys_pred = ys_pred[self.ids]  # pylint: disable=invalid-sequence-index
+        # The gradient and hessian
+        ys = self.dmatrix.get_label()  # type: ignore # pylint: disable=invalid-name
+        gradient = ys_pred - ys
+        hessian = np.ones_like(gradient)
+        return gradient * ys, hessian * ys
+
+    def rmse(self, ys_pred: np.ndarray) -> tuple[str, float]:
+        """Evaluate RMSE (rooted mean square error) in the pack-sum format
+
+        Parameters
+        ----------
+        ys_pred: np.ndarray
+            The raw predictions
+
+        Returns
+        -------
+        name: str
+            The name of the metric
+        score: float
+            The score of the metric
+        """
+        # Making prediction
+        ys_pred = self.predict_with_score(ys_pred)
+        # Propagate prediction to each block
+        ys_pred = ys_pred[self.ids]  # pylint: disable=invalid-sequence-index
+        # The RMSE
+        ys = self.dmatrix.get_label()  # type: ignore # pylint: disable=invalid-name
+        square_error = np.square(ys_pred - ys)
+        rmse = np.sqrt(square_error.mean())
+        return "p-rmse", rmse
+
+    def average_peak_score(
+        self,
+        ys_pred: np.ndarray,
+        n: int,
+    ) -> tuple[str, float]:
+        """Evaluate average-peak-score@N in the pack-sum format
+
+        Parameters
+        ----------
+        ys_pred: np.ndarray
+            The raw prediction
+        n : int
+            The N in average-peak-score@N
+
+        Returns
+        -------
+        name: str
+            The name of the metric
+        score: float
+            The score of the metric
+        """
+        ys = self.dmatrix.get_label()  # type: ignore # pylint: disable=invalid-name
+        ys = self.predict_with_score(ys)  # type: ignore # pylint: disable=invalid-name
+        ys = ys / np.unique(self.ids, return_counts=True)[1]  # type: ignore # pylint: disable=invalid-name
+        ys_pred = self.predict_with_score(ys_pred)
+        trials = np.argsort(ys_pred)[::-1][:n]
+        trial_scores = ys[trials]
+        curve = max_curve(trial_scores) / np.max(ys)
+        score = np.mean(curve)
+        return f"a-peak@{n}", score
+
+
+class XGBConfig(NamedTuple):
+    """XGBoost model configuration
+
+    Reference: https://xgboost.readthedocs.io/en/stable/parameter.html
+
+    Parameters
+    ----------
+    max_depth : int
+        The maximum depth.
+    gamma : float
+        The gamma.
+    min_child_weight : float
+        The minimum child weight.
+    eta : float
+        The eta, learning rate.
+    seed : int
+        The random seed.
+    nthread : Optional[int],
+        The number of threads to use.
+        Default is None, which means to use physical number of cores.
+    tree_method : Literal["auto", "exact", "approx", "hist", "gpu_hist"]
+        The tree construction algorithm used in XGBoost.
+    """
+
+    max_depth: int = 10
+    gamma: float = 0.001
+    min_child_weight: float = 0
+    eta: float = 0.2
+    seed: int = 43
+    nthread: int | None = None
+    tree_method: Literal["auto", "exact", "approx", "hist", "gpu_hist"] = "auto"
+
+    def to_dict(self):
+        """Convert to dict"""
+
+        return {
+            "max_depth": self.max_depth,
+            "gamma": self.gamma,
+            "min_child_weight": self.min_child_weight,
+            "eta": self.eta,
+            "seed": self.seed,
+            "nthread": self.nthread,
+            "tree_method": self.tree_method,
+        }
+
+
+class FeatureGroup:
+    """Feature group
+
+    Parameters
+    ----------
+    group_hash : str
+        The hash of the group
+    features : List[np.ndarray]
+        The features
+    costs : List[float]
+        The costs
+    min_cost : float
+        The minimum cost
+    """
+
+    group_hash: str
+    features: list[np.ndarray]
+    costs: np.ndarray
+    min_cost: float
+
+    def __init__(
+        self,
+        group_hash: str,
+        features: list[np.ndarray],
+        costs: np.ndarray,
+    ) -> None:
+        self.group_hash = group_hash
+        self.features = features
+        self.costs = costs
+        self.min_cost = np.min(costs)
+
+    def append(
+        self,
+        features: list[np.ndarray],
+        costs: np.ndarray,
+    ) -> None:
+        self.features.extend(features)
+        self.costs = np.append(self.costs, costs)
+        self.min_cost = np.min(self.costs)
+
+
+@derived_object
+class XGBModel(PyCostModel):
+    """XGBoost model
+
+    Parameters
+    ----------
+    extractor : FeatureExtractor
+        The feature extractor for the model.
+    config : XGBConfig
+        The XGBoost model config.
+    num_warmup_samples : int
+        The number of samples that are used for warmup, i.e., the first few samples are predicted
+        with random results.
+    early_stopping_rounds : int
+        The number of rounds for early stopping.
+    verbose_eval : int
+        The verbose level when doing evaluation.
+    average_peak_n : int
+        The number to calculate average peak score.
+    adaptive_training : bool
+        Whether use adaptive training to reduce tuning time.
+    """
+
+    # feature extractor
+    extractor: FeatureExtractor
+    # xgboost model config
+    config: XGBConfig
+    # behavior of randomness
+    num_warmup_samples: int
+    # evaluation
+    early_stopping_rounds: int
+    verbose_eval: int
+    average_peak_n: int
+    # states
+    data: dict[str, FeatureGroup]
+    data_size: int
+    booster: Optional["xgb.Booster"]
+    # adaptive training
+    adaptive_training: bool
+    last_train_size: int
+
+    def __init__(
+        self,
+        *,
+        # feature extractor
+        extractor: FeatureExtractor.FeatureExtractorType = "per-store-feature",
+        # xgboost model config
+        config: XGBConfig = XGBConfig(),
+        # random result before enough samples
+        num_warmup_samples: int = 100,
+        # evaluation
+        early_stopping_rounds: int = 50,
+        verbose_eval: int = 25,
+        average_peak_n: int = 32,
+        adaptive_training: bool = True,
+        num_tuning_cores: int | None = None,
+        tree_method: Literal["auto", "exact", "approx", "hist", "gpu_hist"] | None = None,
+    ):
+        super().__init__()
+        if not isinstance(extractor, FeatureExtractor):
+            extractor = FeatureExtractor.create(extractor)
+        # feature extractor
+        self.extractor = extractor
+        # model-related
+        if config.nthread is None:
+            # use physical core number
+            if num_tuning_cores is None:
+                config = config._replace(nthread=cpu_count(logical=False))
+            else:
+                config = config._replace(nthread=num_tuning_cores)
+
+        if tree_method is not None:
+            config._replace(tree_method=tree_method)
+
+        self.config = config
+        # behavior of randomness
+        self.num_warmup_samples = num_warmup_samples
+        # evaluation
+        self.early_stopping_rounds = early_stopping_rounds
+        self.verbose_eval = verbose_eval
+        self.average_peak_n = average_peak_n
+        # states
+        self.data = OrderedDict()
+        self.data_size = 0
+        self.booster = None
+        # adaptive training
+        self.adaptive_training = adaptive_training
+        self.last_train_size = 0
+
+    def load(self, path: str) -> None:
+        """Load the cost model from given file location.
+
+        Parameters
+        ----------
+        path : str
+            The file path.
+
+        Note
+        ----
+        Since XGBoost model trains from scratch, each time this method loads the model together with
+        previously cached feature vectors and results, so that the subsequent training process could
+        use all the existing data being stored on disk.
+        """
+        import xgboost as xgb  # pylint: disable=import-outside-toplevel
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model_path = os.path.join(tmp_dir, "model.bin")
+            data_path = os.path.join(tmp_dir, "data.npy")
+            # Step 1. Untar
+            untar(path, tmp_dir)
+            # Step 2. Load data
+            data = OrderedDict()
+            data_size = 0
+            for group_hash, features, costs in np.load(data_path, allow_pickle=True):
+                data[group_hash] = FeatureGroup(
+                    group_hash=group_hash,
+                    features=list(features),
+                    costs=costs,
+                )
+                data_size += len(costs)
+            # Step 3. Load the model
+            if os.path.exists(model_path):
+                booster = xgb.Booster()
+                booster.load_model(model_path)
+            else:
+                self.booster = None
+        self.data = data
+        self.data_size = data_size
+        self.booster = booster
+
+    def save(self, path: str) -> None:
+        """Save the cost model to given file location.
+
+        Parameters
+        ----------
+        path : str
+            The file path.
+
+        Note
+        ----
+        Since XGBoost model trains from scratch, each time this method saves the model together with
+        previously cached feature vectors and results, so that the subsequent training process could
+        use all the existing data being stored on disk.
+        """
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model_path = os.path.join(tmp_dir, "model.bin")
+            data_path = os.path.join(tmp_dir, "data.npy")
+            # Step 1. Save the model
+            booster = self.booster
+            if booster is not None:
+                booster.save_model(model_path)
+            else:
+                model_path = None
+            # Step 2. Save data
+            data = [
+                (
+                    g.group_hash,
+                    g.features,
+                    g.costs,
+                )
+                for g in self.data.values()
+            ]
+            np.save(
+                file=data_path,
+                arr=np.array(data, dtype=object),
+            )
+            # Step 3. Tar it
+            tar(path, [x for x in [model_path, data_path] if x is not None])
+            logger.info("Saved XGBModel to %s", path)
+
+    def update(
+        self,
+        context: "TuneContext",
+        candidates: list[MeasureCandidate],
+        results: list[RunnerResult],
+    ) -> None:
+        """Update the cost model given running results.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context.
+        candidates : List[MeasureCandidate]
+            The measure candidates.
+        results : List[RunnerResult]
+            The running results of the measure candidates.
+        """
+        assert len(candidates) == len(results)
+        if len(candidates) == 0:
+            return
+
+        # Step 1. Get the feature group
+        new_group_hash = shash2hex(context.mod)
+        group = self.data.get(new_group_hash, None)
+
+        # Step 2. Extract features
+        def _feature(x: Tensor) -> np.ndarray:
+            return x.numpy().astype("float32")
+
+        def _mean_cost(x: RunnerResult) -> float:
+            if not x.run_secs:
+                return 1e10
+            return float(np.median([float(s) for s in x.run_secs]))
+
+        new_features = [_feature(x) for x in self.extractor.extract_from(context, candidates)]
+        new_mean_costs = [_mean_cost(x) for x in results]
+
+        # Filter instances with no features
+        new_mean_costs = [c for i, c in enumerate(new_mean_costs) if len(new_features[i]) != 0]
+        new_mean_costs_np = np.array(new_mean_costs).astype("float32")
+        new_features = [f for f in new_features if len(f) != 0]
+        if not new_features:
+            return
+
+        # Steps 3. Run validation
+        if group is not None and self.booster is not None:
+            logger.debug(
+                "XGB validation: %s",
+                "\t".join(
+                    f"{key}: {score:.6f}"
+                    for key, score in self._validate(
+                        xs=new_features,
+                        ys=group.min_cost / new_mean_costs_np,
+                    )
+                ),
+            )
+
+        # Step 4. Add the features into the data points
+        if group is None:
+            group = FeatureGroup(
+                group_hash=new_group_hash,
+                features=new_features,
+                costs=new_mean_costs_np,
+            )
+        else:
+            group.append(new_features, new_mean_costs_np)
+        self.data[new_group_hash] = group
+        self.data_size += len(new_features)
+
+        if (
+            self.adaptive_training
+            and self.data_size - self.last_train_size < self.last_train_size / 5
+        ):
+            # Set a training threshold related to `last_train_size` to reduce the training
+            # overhead when there're too many results
+            return
+        self.last_train_size = self.data_size
+
+        # Step 5. Re-train the model
+        with np.errstate(divide="ignore", invalid="ignore"):
+            feature_list = list(
+                itertools_chain.from_iterable([g.features for g in self.data.values()])
+            )
+            cost_ratio_list = [
+                np.divide(g.min_cost, g.costs, out=np.zeros_like(g.costs), where=g.costs != 0)
+                for g in self.data.values()
+            ]
+            cost_ratios = np.concatenate(cost_ratio_list, axis=0)
+
+        self._train(xs=feature_list, ys=cost_ratios)
+
+    def predict(
+        self,
+        context: "TuneContext",
+        candidates: list[MeasureCandidate],
+    ) -> np.ndarray:
+        """Predict the normalized score using the cost model.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context.
+        candidates : List[MeasureCandidate]
+            The measure candidates.
+
+        Return
+        ------
+        result : np.ndarray
+            The predicted normalized score.
+        """
+        if self.data_size >= self.num_warmup_samples and self.booster is not None:
+            ret = self._predict(
+                xs=[
+                    x.numpy().astype("float32")
+                    for x in self.extractor.extract_from(
+                        context,
+                        candidates,
+                    )
+                ]
+            )
+        else:
+            ret = np.random.uniform(
+                low=0,
+                high=1,
+                size=(len(candidates),),
+            )
+        return ret.astype("float64")
+
+    def _train(  # type: ignore # pylint: disable=invalid-name
+        self,
+        xs: list[np.ndarray],
+        ys: np.ndarray,
+    ) -> None:
+        import xgboost as xgb  # type: ignore # pylint: disable=import-outside-toplevel
+
+        self.d_train = PackSum(xs=xs, ys=ys)
+
+        def obj(ys_pred: np.ndarray, d_train: "xgb.DMatrix"):  # type: ignore # pylint: disable = unused-argument
+            return self.d_train.obj_square_error(ys_pred)
+
+        def rmse(ys_pred: np.ndarray, d_train: "xgb.DMatrix"):  # type: ignore # pylint: disable = unused-argument
+            return self.d_train.rmse(ys_pred)
+
+        def avg_peak_score(ys_pred: np.ndarray, d_train: "xgb.DMatrix"):  # type: ignore # pylint: disable = unused-argument
+            return self.d_train.average_peak_score(ys_pred, self.average_peak_n)
+
+        self.booster = xgb.train(
+            self.config.to_dict(),
+            self.d_train.dmatrix,
+            num_boost_round=10000,
+            obj=obj,
+            callbacks=[
+                _get_custom_call_back(
+                    early_stopping_rounds=self.early_stopping_rounds,
+                    verbose_eval=self.verbose_eval,
+                    fevals=[rmse, avg_peak_score],
+                    evals=[(self.d_train.dmatrix, "tr")],
+                    cvfolds=None,
+                )
+            ],
+        )
+
+        del self.d_train
+
+    def _predict(  # type: ignore # pylint: disable=invalid-name
+        self,
+        xs: list[np.ndarray],
+    ) -> np.ndarray:
+        d_test = PackSum(xs=xs, ys=None)
+        pred = self.booster.predict(d_test.dmatrix)
+        ret = d_test.predict_with_score(pred)
+        return ret
+
+    def _validate(  # type: ignore # pylint: disable=invalid-name
+        self,
+        xs: list[np.ndarray],
+        ys: np.ndarray,
+    ) -> list[tuple[str, float]]:
+        """Evaluate the score of inputs.
+
+        Parameters
+        ----------
+        xs : List[np.ndarray]
+            A batch of input samples
+        ys : List[float]
+            A batch of labels
+
+        Returns
+        -------
+        scores: np.ndarray
+            The predicted result for all inputs.
+        """
+        assert self.booster is not None
+
+        d_valid = PackSum(xs=xs, ys=ys)
+
+        def average_peak_score(ys_pred: np.ndarray):
+            return d_valid.average_peak_score(ys_pred, n=self.average_peak_n)
+
+        ys_pred = self.booster.predict(d_valid.dmatrix)
+        eval_result: list[tuple[str, float]] = [
+            feval(ys_pred)
+            for feval in (
+                average_peak_score,
+                d_valid.rmse,
+            )
+        ]
+        eval_result.sort(key=make_metric_sorter("p-rmse"))
+        return eval_result
+
+
+def _get_custom_call_back(
+    early_stopping_rounds: int,
+    verbose_eval: int,
+    fevals: list[Callable],
+    evals: list[tuple["xgb.DMatrix", str]],
+    focused_metric: str = "tr-p-rmse",
+    cvfolds: list["xgb.training.CVPack"] | None = None,
+) -> "TrainingCallback":
+    """Get a customized callback function for XGBoost. Work around xgboost import."""
+
+    def optional_xgboost_callback(cls):
+        """Decorator for importing TrainingCallback from xgboost"""
+        # pylint:disable = import-outside-toplevel
+        try:
+            from xgboost.callback import TrainingCallback  # type: ignore
+        # pylint:enable = import-outside-toplevel
+        except ImportError:
+
+            class TrainingCallback:  # type: ignore
+                pass
+
+        class OptXGBoostCustomCallback(cls, TrainingCallback):  # type: ignore
+            pass
+
+        return OptXGBoostCustomCallback
+
+    @optional_xgboost_callback
+    class XGBoostCustomCallback:
+        """Custom callback class for xgboost to support multiple custom evaluation functions"""
+
+        def __init__(
+            self,
+            early_stopping_rounds: int,
+            verbose_eval: int,
+            fevals: list[Callable],
+            evals: list[tuple["xgb.DMatrix", str]],
+            focused_metric: str = "tr-p-rmse",
+            cvfolds: list["xgb.training.CVPack"] | None = None,
+        ):
+            self.early_stopping_rounds = early_stopping_rounds
+            self.verbose_eval = verbose_eval
+            self.fevals = fevals
+            self.evals = evals
+            self.state: dict[str, Any] = {}
+            self.focused_metric = focused_metric
+            self.sort_key = make_metric_sorter(focused_metric=focused_metric)
+            self.cvfolds = cvfolds
+            if cvfolds is not None:
+                self.aggregated_cv = None
+
+        def __call__(self, env: "xgb.core.CallbackEnv"):
+            # Compatibility with xgboost < 1.3
+            return self.after_iteration(env.model, env.iteration, env.evaluation_result_list)
+
+        def init(self, model: "xgb.Booster"):
+            """Internal function for initialization"""
+            booster: xgb.Booster = model
+            self.state["best_iteration"] = 0
+            self.state["best_score"] = float("inf")
+            if booster is None:
+                assert self.cvfolds is not None
+                return
+            if booster.attr("best_score") is not None:
+                self.state["best_score"] = float(booster.attr("best_score"))
+                self.state["best_iteration"] = int(booster.attr("best_iteration"))
+                self.state["best_msg"] = booster.attr("best_msg")
+            else:
+                booster.set_attr(best_iteration=str(self.state["best_iteration"]))
+                booster.set_attr(best_score=str(self.state["best_score"]))
+
+        def after_iteration(self, model: "xgb.Booster", epoch: int, evals_log: dict):  # pylint: disable = unused-argument
+            """Internal function for after_iteration"""
+            # pylint:disable = import-outside-toplevel
+            try:
+                from xgboost.callback import _fmt_metric  # type: ignore
+            except ImportError:
+                # Compatibility with xgboost >= 1.6
+
+                def _fmt_metric(value, show_stdv=True):
+                    if len(value) == 2:
+                        return f"{value[0]}:{value[1]:.5f}"
+                    if len(value) == 3:
+                        if show_stdv:
+                            return f"{value[0]}:{value[1]:.5f}+{value[2]:.5f}"
+                        return f"{value[0]}:{value[1]:.5f}"
+                    raise ValueError("wrong metric value", value)
+
+            import xgboost as xgb
+
+            # make it compatible with xgboost<1.7
+            try:
+                from xgboost import rabit as collective  # type: ignore
+            except ImportError:
+                from xgboost import collective  # type: ignore
+
+            try:
+                from xgboost.training import aggcv  # type: ignore
+            except ImportError:
+                from xgboost.callback import _aggcv as aggcv  # type: ignore
+
+            # pylint:enable = import-outside-toplevel
+            if not self.state:
+                self.init(model)
+            booster: xgb.Booster = model
+            iteration: int = epoch
+            cvfolds: list[xgb.training.CVPack] = self.cvfolds
+            ##### Evaluation #####
+            # `eval_result` is a list of (key, score)
+            eval_result: list[tuple[str, float]] = []
+            if cvfolds is None:
+                eval_result = list(
+                    itertools_chain.from_iterable(
+                        [
+                            (key, float(value))
+                            for key, value in map(
+                                lambda x: x.split(":"),
+                                booster.eval_set(
+                                    evals=self.evals,
+                                    iteration=iteration,
+                                    feval=feval,
+                                ).split()[1:],
+                            )
+                        ]
+                        for feval in self.fevals
+                    )
+                )
+            else:
+                eval_result = list(
+                    itertools_chain.from_iterable(
+                        [
+                            (key, score)
+                            for key, score, _std in aggcv(
+                                fold.eval(
+                                    iteration=iteration,
+                                    feval=feval,
+                                )
+                                for fold in cvfolds
+                            )
+                        ]
+                        for feval in self.fevals
+                    )
+                )
+            eval_result = list(eval_result)
+            eval_result.sort(key=self.sort_key)
+
+            ##### Print eval result #####
+            if self.verbose_eval and iteration % self.verbose_eval == 0:
+                info = []
+                for key, score in eval_result:
+                    if "null" not in key:
+                        info.append(f"{key}: {score:.6f}")
+                logger.debug("XGB iter %3d: %s", iteration, "\t".join(info))
+
+            ##### Choose score and do early stopping #####
+            score = None
+            for key, _score in eval_result:
+                if key == self.focused_metric:
+                    score = _score
+                    break
+            assert score is not None
+
+            best_score = self.state["best_score"]
+            best_iteration = self.state["best_iteration"]
+            if score < best_score:
+                tab = "\t"  # to work with f-string
+                msg = f"[{epoch}] {tab.join([_fmt_metric(x) for x in eval_result])}"
+                self.state["best_msg"] = msg
+                self.state["best_score"] = score
+                self.state["best_iteration"] = epoch
+                # save the property to attributes, so they will occur in checkpoint.
+                if model is not None:
+                    model.set_attr(
+                        best_score=str(self.state["best_score"]),
+                        best_iteration=str(self.state["best_iteration"]),
+                        best_msg=self.state["best_msg"],
+                    )
+            elif epoch - best_iteration >= self.early_stopping_rounds:
+                best_msg = self.state["best_msg"]
+
+                if self.verbose_eval and collective.get_rank() == 0:
+                    logger.debug("XGB stopped. Best iteration: %s ", best_msg)
+                # instead of raising EarlyStopException, returning True to end the training
+                return True
+            # False to indicate training should not stop.
+            return False
+
+    return XGBoostCustomCallback(
+        early_stopping_rounds=early_stopping_rounds,
+        verbose_eval=verbose_eval,
+        fevals=fevals,
+        evals=evals,
+        focused_metric=focused_metric,
+        cvfolds=cvfolds,
+    )

python/tvm/s_tir/meta_schedule/database/__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.
+"""
+The tvm.s_tir.meta_schedule.database package.
+The database that stores serialized tuning records and workloads
+"""
+
+from .database import Database, PyDatabase, TuningRecord, Workload, create
+from .json_database import JSONDatabase
+from .memory_database import MemoryDatabase
+from .ordered_union_database import OrderedUnionDatabase
+from .schedule_fn_database import ScheduleFnDatabase
+from .union_database import UnionDatabase

python/tvm/s_tir/meta_schedule/database/database.py

@@ -0,0 +1,644 @@
+# 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.
+# ruff: noqa: RUF012
+"""TuningRecord database"""
+
+from collections.abc import Callable
+from typing import Any, Optional, Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+from tvm_ffi import register_object
+
+from tvm.ir.module import IRModule
+from tvm.runtime import Object
+from tvm.s_tir.schedule import Schedule, Trace
+from tvm.target import Target
+
+from .. import _ffi_api
+from ..arg_info import ArgInfo
+from ..utils import _json_de_tvm
+
+
+@register_object("s_tir.meta_schedule.Workload")
+class Workload(Object):
+    """A workload, i.e. an IRModule and its structural hash.
+
+    Parameters
+    ----------
+    mod : IRModule
+        The workload's IRModule
+    """
+
+    mod: IRModule
+
+    def __init__(self, mod: IRModule) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.Workload,  # type: ignore # pylint: disable=no-member
+            mod,
+        )
+
+    def as_json(self) -> Any:
+        """Export the workload to JSON as a python object.
+
+        Returns
+        -------
+        json : Any
+            The JSON serialized as a python object (e.g. a Dict or List).
+            Use json.dumps() to get the associated json string.
+        """
+        return _json_de_tvm(_ffi_api.WorkloadAsJSON(self))  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def from_json(json_obj: Any) -> "Workload":
+        """Create a workload from a json object.
+
+        Parameters
+        ----------
+        json_obj : Any
+            The json object to parse.
+
+        Returns
+        -------
+        tuning_record : TuningRecord
+            The parsed tuning record.
+        """
+        return _ffi_api.WorkloadFromJSON(json_obj)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("s_tir.meta_schedule.TuningRecord")
+class TuningRecord(Object):
+    """The class of tuning records.
+
+    Parameters
+    ----------
+    trace : tvm.ir.Trace
+        The trace of the tuning record.
+    workload : Workload
+        The workload of the tuning record.
+    run_secs : Optional[List[float]]
+        The run time of the tuning record.
+    target : Optional[Target]
+        The target of the tuning record.
+    args_info : Optional[List[ArgInfo]]
+        The argument information of the tuning record.
+    """
+
+    trace: Trace
+    workload: Workload
+    run_secs: list[float] | None
+    target: Target | None
+    args_info: list[ArgInfo] | None
+
+    def __init__(  # type: ignore # pylint: disable=too-many-arguments
+        self,
+        trace: Trace,
+        workload: Workload,
+        run_secs: list[float] | None = None,
+        target: Target | None = None,
+        args_info: list[ArgInfo] | None = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.TuningRecord,  # type: ignore # pylint: disable=no-member
+            trace,
+            workload,
+            run_secs,
+            target,
+            args_info,
+        )
+
+    def as_measure_candidate(self) -> Any:
+        """Generate a measure candidate given an initial IR module and a trace
+        stored in the tuning record.
+
+        Returns
+        -------
+        candidate : MeasureCandidate
+            A generated candidate.
+        """
+        return _ffi_api.TuningRecordAsMeasureCandidate(self)  # type: ignore # pylint: disable=no-member
+
+    def as_json(self) -> Any:
+        """Export the tuning record to a JSON string.
+
+        Returns
+        -------
+        json_str : str
+            The JSON string exported.
+        """
+        return _json_de_tvm(_ffi_api.TuningRecordAsJSON(self))  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def from_json(json_obj: Any, workload: Workload) -> "TuningRecord":
+        """Create a tuning record from a json object.
+
+        Parameters
+        ----------
+        json_obj : Any
+            The json object to parse.
+        workload : Workload
+            The workload.
+
+        Returns
+        -------
+        tuning_record : TuningRecord
+            The parsed tuning record.
+        """
+        return _ffi_api.TuningRecordFromJSON(json_obj, workload)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("s_tir.meta_schedule.Database")
+class Database(Object):
+    """The abstract database interface."""
+
+    DatabaseType = Union["Database", Literal["json", "memory"]]
+
+    def has_workload(self, mod: IRModule) -> bool:
+        """Check if the database has the given workload.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for.
+
+        Returns
+        -------
+        result : bool
+            Whether the database has the given workload.
+        """
+        return _ffi_api.DatabaseHasWorkload(self, mod)  # type: ignore # pylint: disable=no-member
+
+    def commit_workload(self, mod: IRModule) -> Workload:
+        """Commit a workload to the database if missing.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for or added.
+
+        Returns
+        -------
+        workload : Workload
+            The workload corresponding to the given IRModule.
+        """
+        return _ffi_api.DatabaseCommitWorkload(self, mod)  # type: ignore # pylint: disable=no-member
+
+    def commit_tuning_record(self, record: TuningRecord) -> None:
+        """Commit a tuning record to the database.
+
+        Parameters
+        ----------
+        record : TuningRecord
+            The tuning record to add.
+        """
+        _ffi_api.DatabaseCommitTuningRecord(self, record)  # type: ignore # pylint: disable=no-member
+
+    def get_top_k(self, workload: Workload, top_k: int) -> list[TuningRecord]:
+        """Get the top K valid tuning records of given workload from the database.
+
+        Parameters
+        ----------
+        workload : Workload
+            The workload to be searched for.
+        top_k : int
+            The number of top records to get.
+
+        Returns
+        -------
+        top_k_records : List[TuningRecord]
+            The top K records.
+        """
+        return _ffi_api.DatabaseGetTopK(self, workload, top_k)  # type: ignore # pylint: disable=no-member
+
+    def get_all_tuning_records(self) -> list[TuningRecord]:
+        """Get all the tuning records from the database.
+
+        Returns
+        -------
+        tuning_records : List[TuningRecord]
+            All tuning records from the database.
+        """
+        return _ffi_api.DatabaseGetAllTuningRecords(self)  # type: ignore # pylint: disable=no-member
+
+    def __len__(self) -> int:
+        """Get the number of records in the database.
+
+        Returns
+        -------
+        num_records : int
+            The number of records in the database
+        """
+        return _ffi_api.DatabaseSize(self)  # type: ignore # pylint: disable=no-member
+
+    def query_tuning_record(
+        self,
+        mod: IRModule,
+        target: Target,
+        workload_name: str,
+    ) -> TuningRecord | None:
+        """Query the best record of the given workload from the database.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for.
+        target : Target
+            The target to be searched for.
+        workload_name : str
+            The name of the workload to be searched for.
+
+        Returns
+        -------
+        tuning_record : Optional[TuningRecord]
+            The best record of the given workload; None if not found.
+        """
+        return _ffi_api.DatabaseQueryTuningRecord(self, mod, target, workload_name)  # type: ignore # pylint: disable=no-member
+
+    def query_schedule(
+        self,
+        mod: IRModule,
+        target: Target,
+        workload_name: str,
+    ) -> Schedule | None:
+        """Query the best schedule of the given workload from the database.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for.
+        target : Target
+            The target to be searched for.
+        workload_name : str
+            The name of the workload to be searched for.
+
+        Returns
+        -------
+        schedule : Optional[tvm.s_tir.Schedule]
+            The best schedule of the given workload; None if not found.
+        """
+        return _ffi_api.DatabaseQuerySchedule(self, mod, target, workload_name)  # type: ignore # pylint: disable=no-member
+
+    def query_ir_module(
+        self,
+        mod: IRModule,
+        target: Target,
+        workload_name: str,
+    ) -> IRModule | None:
+        """Query the best IRModule of the given workload from the database.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for.
+        target : Target
+            The target to be searched for.
+        workload_name : str
+            The name of the workload to be searched for.
+
+        Returns
+        -------
+        ir_module : Optional[IRModule]
+            The best IRModule of the given workload; None if not found.
+        """
+        return _ffi_api.DatabaseQueryIRModule(self, mod, target, workload_name)  # type: ignore # pylint: disable=no-member
+
+    def dump_pruned(self, destination: "Database") -> None:
+        """Dump the pruned database to files of JSONDatabase format.
+
+        Parameters
+        ----------
+        destination : Database
+            The destination database to be dumped to.
+        """
+        return _ffi_api.DatabaseDumpPruned(  # type: ignore # pylint: disable=no-member
+            self, destination
+        )
+
+    def query(
+        self,
+        mod: IRModule,
+        target: Target,
+        *,
+        workload_name: str = "main",
+        kind: Literal["schedule"] | Literal["record"] | Literal["ir_module"] = "schedule",
+    ) -> Schedule | IRModule | TuningRecord:
+        """Query the database to retrieve the best optimization outcome of the given workload.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for.
+        target : Target
+            The target to be searched for.
+        kind : str = "schedule" | "record" | "ir_module"
+            The kind of the optimization outcome to be returned.
+
+        Returns
+        -------
+        result : Union[tvm.s_tir.Schedule, IRModule, TuningRecord]
+            The best optimization outcome of the given workload.
+        """
+        if kind == "schedule":
+            return self.query_schedule(mod, target, workload_name)
+        if kind == "record":
+            return self.query_tuning_record(mod, target, workload_name)
+        if kind == "ir_module":
+            return self.query_ir_module(mod, target, workload_name)
+        raise ValueError(f'Unknown kind: {kind}. Candidates are: "schedule", "record", "ir_module"')
+
+    def __enter__(self) -> "Database":
+        """Entering the scope of the context manager"""
+        _ffi_api.DatabaseEnterWithScope(self)  # type: ignore # pylint: disable=no-member
+        return self
+
+    def __exit__(self, ptype, value, trace) -> None:
+        """Exiting the scope of the context manager"""
+        _ffi_api.DatabaseExitWithScope(self)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def current() -> Optional["Database"]:
+        """Get the current database under scope."""
+        return _ffi_api.DatabaseCurrent()  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def create(  # pylint: disable=keyword-arg-before-vararg
+        kind: (
+            Literal["json", "memory", "union", "ordered_union"] | Callable[[Schedule], bool]
+        ) = "json",
+        *args,
+        **kwargs,
+    ) -> "Database":
+        """Create a Database.
+
+        Parameters
+        ----------
+        kind : str = "json" | "memory" | "union" | "ordered_union" | Callable[[tvm.s_tir.Schedule],
+        bool]
+            The kind of the database to be created. The following kinds are supported:
+            "json", "memory", "union", "ordered_union", and a custom schedule function.
+
+        Returns
+        -------
+        database : Database
+            The created database.
+        """
+        from . import (  # pylint: disable=import-outside-toplevel
+            JSONDatabase,
+            MemoryDatabase,
+            OrderedUnionDatabase,
+            ScheduleFnDatabase,
+            UnionDatabase,
+        )
+
+        if callable(kind):
+            return ScheduleFnDatabase(kind, *args, **kwargs)  # type: ignore
+        if kind == "json":
+            return JSONDatabase(*args, **kwargs)
+        if kind == "memory":
+            return MemoryDatabase(*args, **kwargs)  # type: ignore
+        if kind == "union":
+            return UnionDatabase(*args, **kwargs)  # type: ignore
+        if kind == "ordered_union":
+            return OrderedUnionDatabase(*args, **kwargs)  # type: ignore
+        raise ValueError(f"Unknown Database: {kind}")
+
+
+create = Database.create  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.PyDatabase")
+class _PyDatabase(Database):
+    """
+    A TVM object database to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyDatabase
+    """
+
+    def __init__(
+        self,
+        f_has_workload: Callable | None = None,
+        f_commit_workload: Callable | None = None,
+        f_commit_tuning_record: Callable | None = None,
+        f_get_top_k: Callable | None = None,
+        f_get_all_tuning_records: Callable | None = None,
+        f_query_tuning_record: Callable | None = None,
+        f_query_schedule: Callable | None = None,
+        f_query_ir_module: Callable | None = None,
+        f_size: Callable | None = None,
+        module_equality: str = "structural",
+    ):
+        """Constructor."""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.DatabasePyDatabase,  # type: ignore  # pylint: disable=no-member
+            f_has_workload,
+            f_commit_workload,
+            f_commit_tuning_record,
+            f_get_top_k,
+            f_get_all_tuning_records,
+            f_query_tuning_record,
+            f_query_schedule,
+            f_query_ir_module,
+            f_size,
+            module_equality,
+        )
+
+
+class PyDatabase:
+    """
+    An abstract database with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PyDatabase,
+        "methods": [
+            "has_workload",
+            "commit_workload",
+            "commit_tuning_record",
+            "get_top_k",
+            "get_all_tuning_records",
+            "query_tuning_record",
+            "query_schedule",
+            "query_ir_module",
+            "__len__",
+        ],
+    }
+
+    def has_workload(self, mod: IRModule) -> bool:
+        """Check if the database has the given workload.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for.
+
+        Returns
+        -------
+        result : bool
+            Whether the database has the given workload.
+        """
+        raise NotImplementedError
+
+    def commit_workload(self, mod: IRModule) -> Workload:
+        """Commit a workload to the database if missing.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for or added.
+
+        Returns
+        -------
+        workload : Workload
+            The workload corresponding to the given IRModule.
+        """
+        raise NotImplementedError
+
+    def commit_tuning_record(self, record: TuningRecord) -> None:
+        """Commit a tuning record to the database.
+
+        Parameters
+        ----------
+        record : TuningRecord
+            The tuning record to add.
+        """
+        raise NotImplementedError
+
+    def get_top_k(self, workload: Workload, top_k: int) -> list[TuningRecord]:
+        """Get the top K tuning records of given workload from the database.
+
+        Parameters
+        ----------
+        workload : Workload
+            The workload to be searched for.
+        top_k : int
+            The number of top records to get.
+
+        Returns
+        -------
+        top_k_records : List[TuningRecord]
+            The top K records.
+        """
+        raise NotImplementedError
+
+    def get_all_tuning_records(self) -> list[TuningRecord]:
+        """Get all the tuning records from the database.
+
+        Returns
+        -------
+        tuning_records : List[TuningRecord]
+            All tuning records from the database.
+        """
+        raise NotImplementedError
+
+    def query_tuning_record(
+        self, mod: IRModule, target: Target, workload_name: str | None = None
+    ) -> TuningRecord | None:
+        """Query a tuning record from the database.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for.
+        target : Target
+            The target to be searched for.
+        workload_name : Optional[str]
+            The workload name to be searched for.
+
+        Returns
+        -------
+        record : Optional[TuningRecord]
+            The tuning record corresponding to the given workload.
+        """
+        # Using self._outer to replace the self pointer
+        return _ffi_api.DatabaseQueryTuningRecord(  # type: ignore # pylint: disable=no-member
+            self._outer(),
+            mod,
+            target,
+            workload_name,  # type: ignore # pylint: disable=no-member
+        )
+
+    def query_schedule(
+        self, mod: IRModule, target: Target, workload_name: str | None = None
+    ) -> Schedule | None:
+        """Query a schedule from the database.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for.
+        target : Target
+            The target to be searched for.
+        workload_name : Optional[str]
+            The workload name to be searched for.
+
+        Returns
+        -------
+        schedule : Optional[Schedule]
+            The schedule corresponding to the given workload.
+        """
+        # Using self._outer to replace the self pointer
+        return _ffi_api.DatabaseQuerySchedule(  # type: ignore # pylint: disable=no-member
+            self._outer(),
+            mod,
+            target,
+            workload_name,  # type: ignore # pylint: disable=no-member
+        )
+
+    def query_ir_module(
+        self, mod: IRModule, target: Target, workload_name: str | None = None
+    ) -> IRModule | None:
+        """Query an IRModule from the database.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The IRModule to be searched for.
+        target : Target
+            The target to be searched for.
+        workload_name : Optional[str]
+            The workload name to be searched for.
+
+        Returns
+        -------
+        mod : Optional[IRModule]
+            The IRModule corresponding to the given workload.
+        """
+        # Using self._outer to replace the self pointer
+        return _ffi_api.DatabaseQueryIRModule(  # type: ignore # pylint: disable=no-member
+            self._outer(),
+            mod,
+            target,
+            workload_name,  # type: ignore # pylint: disable=no-member
+        )
+
+    def __len__(self) -> int:
+        """Get the number of records in the database.
+
+        Returns
+        -------
+        num_records : int
+            The number of records in the database
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/database/json_database.py

@@ -0,0 +1,93 @@
+# 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.
+"""The default database that uses a JSON File to store tuning records"""
+
+import os.path as osp
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .database import Database
+
+
+@register_object("s_tir.meta_schedule.JSONDatabase")
+class JSONDatabase(Database):
+    """Database class backed by JSON.
+
+    Parameters
+    ----------
+    path_workload : str
+        The path to the workload table.
+    path_tuning_record : str
+        The path to the tuning record table.
+    module_equality : Optional[str]
+        A string to specify the module equality testing and hashing method.
+        It must be one of the followings:
+
+          - "structural": Use StructuralEqual/Hash
+          - "ignore-tensor": Same as "structural", but ignore tensor raw data during
+                              equality testing and hashing.
+          - "anchor-block": Apply equality testing and hashing on the anchor block extracted from a
+                            given module. The "ignore-tensor" varint is used for the extracted
+                            blocks or in case no anchor block is found.
+                            For the definition of the anchor block, see tirx/analysis/analysis.py.
+    """
+
+    path_workload: str
+    path_tuning_record: str
+
+    def __init__(
+        self,
+        path_workload: str | None = None,
+        path_tuning_record: str | None = None,
+        *,
+        work_dir: str | None = None,
+        allow_missing: bool = True,
+        module_equality: str = "structural",
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        path_workload : Optional[str] = None
+            The path to the workload table. If not specified,
+            will be generated from `work_dir` as `$work_dir/database_workload.json`.
+        path_tuning_record : Optional[str] = None
+            The path to the tuning record table. If not specified,
+            will be generated from `work_dir` as `$work_dir/database_tuning_record.json`.
+        work_dir : Optional[str] = None
+            The work directory, if specified, will be used to generate `path_tuning_record`
+            and `path_workload`.
+        allow_missing : bool
+            Whether to create new file when the given path is not found.
+        """
+        if work_dir is not None:
+            if path_workload is None:
+                path_workload = osp.join(work_dir, "database_workload.json")
+            if path_tuning_record is None:
+                path_tuning_record = osp.join(work_dir, "database_tuning_record.json")
+        if path_workload is None:
+            raise ValueError("`path_workload` is not specified.")
+        if path_tuning_record is None:
+            raise ValueError("`path_tuning_record` is not specified.")
+        self.__init_handle_by_constructor__(
+            _ffi_api.DatabaseJSONDatabase,  # type: ignore # pylint: disable=no-member
+            path_workload,
+            path_tuning_record,
+            allow_missing,
+            module_equality,
+        )

python/tvm/s_tir/meta_schedule/database/memory_database.py

@@ -0,0 +1,51 @@
+# 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.
+"""A database that stores TuningRecords in memory"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .database import Database
+
+
+@register_object("s_tir.meta_schedule.MemoryDatabase")
+class MemoryDatabase(Database):
+    """An in-memory database
+
+    Parameters
+    ----------
+    module_equality : Optional[str]
+        A string to specify the module equality testing and hashing method.
+        It must be one of the followings:
+
+          - "structural": Use StructuralEqual/Hash
+          - "ignore-tensor": Same as "structural", but ignore tensor raw data during
+                              equality testing and hashing.
+          - "anchor-block": Apply equality testing and hashing on the anchor block extracted from a
+                            given module. The "ignore-tensor" varint is used for the extracted
+                            blocks or in case no anchor block is found.
+                            For the definition of the anchor block, see tirx/analysis/analysis.py.
+    """
+
+    def __init__(
+        self,
+        module_equality: str = "structural",
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.DatabaseMemoryDatabase,  # type: ignore # pylint: disable=no-member,
+            module_equality,
+        )

python/tvm/s_tir/meta_schedule/database/ordered_union_database.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.
+"""A database consists of multiple databases."""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .database import Database
+
+
+@register_object("s_tir.meta_schedule.OrderedUnionDatabase")
+class OrderedUnionDatabase(Database):
+    """A database composed of multiple databases, allowing users to guide IR rewriting using
+    combined knowledge of those databases. To each query, it returns the record from the first
+    database that responds to the query.
+
+    Examples
+    --------
+    Examples below demonstrate the usecases of and difference between UnionDatabase and
+    OrderDatabase.
+
+    Assumption:
+    * db1, db2 do not have tuning records for the target workload.
+    * Each of db3, db4, db5 has tuning records r3, r4, r5 for target workload respectively.
+
+    .. code-block:: python
+
+        #### Case 1. `UnionDatabase`:
+        merged_db = ms.database.UnionDatabase(
+            db1, # no record
+            db2, # no record
+            db3, # has r3
+            db4  # has r4
+        )
+        # returns the better one between r3 and r4
+        merged_db.query_tuning_record(..., target_workload)
+
+        ### Case 2. `OrderedUnionDatabase`
+        merged_db = ms.database.OrderedUnionDatabase(
+            db1, # no record
+            db2, # no record
+            db3, # has r3
+            db4  # has r4
+        )
+        # returns r3
+        merged_db.query_tuning_record(..., target_workload)
+
+        ### Case 3. Mix-use scenario
+        merged_db = ms.database.UnionDatabase(
+            db1, # no record
+            db2, # no record
+            db3, # has r3
+            ms.database.OrderedUnionDatabase( # returns r4
+                db4,  # has r4
+                db5,  # has r5
+            )
+        )
+        # returns the better one between r3 and r4
+        merged_db.query_tuning_record(..., target_workload)
+
+        ### Case 4. Another mix-use scenario
+        merged_db = ms.database.UnionDatabase(
+            db1, # no record
+            db2, # no record
+            db3, # has r3
+            ms.database.UnionDatabase( # returns best one between r4 and r5
+                db4,  # has r4
+                db5,  # has r5
+            )
+        )
+        # returns the best one among r3, r4 and r5
+        merged_db.query_tuning_record(..., target_workload)
+
+        ### Case 5. Yet another mix-use scenario
+        merged_db = ms.database.OrderedUnionDatabase(
+            db1, # no record
+            db2, # no record
+            ms.database.UnionDatabase( # returns best one between r3 and r4
+                db3, # has r3
+                db4,  # has r4
+            )
+            db5,  # has r5
+        )
+        # returns the better one between r3 and r4
+        merged_db.query_tuning_record(..., target_workload)
+    """
+
+    def __init__(self, *databases: Database) -> None:
+        """Construct a merged database from multiple databases.
+
+        Parameters
+        ----------
+        *databases : Database
+            The list of databases to combine.
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.DatabaseOrderedUnionDatabase,  # type: ignore # pylint: disable=no-member
+            databases,
+        )

python/tvm/s_tir/meta_schedule/database/schedule_fn_database.py

@@ -0,0 +1,60 @@
+# 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.
+"""A database for injecting handcrafted schedule functions."""
+
+from collections.abc import Callable
+
+from tvm_ffi import register_object
+
+from tvm.s_tir import Schedule
+
+from .. import _ffi_api
+from .database import Database
+
+
+@register_object("s_tir.meta_schedule.ScheduleFnDatabase")
+class ScheduleFnDatabase(Database):
+    """A database for injecting handcrafted schedule functions.
+
+    Parameters
+    ----------
+    schedule_fn : Callable[[Schedule], bool],
+        The function to do scheduling, which takes a TIR schedule, and returns
+        a boolean indicating if the schedule is committed to the database.
+    module_equality : Optional[str]
+        A string to specify the module equality testing and hashing method.
+        It must be one of the followings:
+
+          - "structural": Use StructuralEqual/Hash
+          - "ignore-tensor": Same as "structural", but ignore tensor raw data during
+                              equality testing and hashing.
+          - "anchor-block": Apply equality testing and hashing on the anchor block extracted from a
+                            given module. The "ignore-tensor" varint is used for the extracted
+                            blocks or in case no anchor block is found.
+                            For the definition of the anchor block, see tirx/analysis/analysis.py.
+    """
+
+    def __init__(
+        self,
+        schedule_fn: Callable[[Schedule], bool],
+        module_equality: str = "structural",
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.DatabaseScheduleFnDatabase,  # type: ignore # pylint: disable=no-member
+            schedule_fn,
+            module_equality,
+        )

python/tvm/s_tir/meta_schedule/database/union_database.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.
+"""A database consists of multiple databases."""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .database import Database
+
+
+@register_object("s_tir.meta_schedule.UnionDatabase")
+class UnionDatabase(Database):
+    """A database composed of multiple databases, allowing users to guide IR rewriting using
+    combined knowledge of those databases. To each query, it returns the best record among all the
+    databases given.
+
+    Examples
+    --------
+    Examples below demonstrate the usecases of and difference between UnionDatabase and
+    OrderDatabase.
+
+    Assumption:
+    * db1, db2 do not have tuning records for the target workload.
+    * Each of db3, db4, db5 has tuning records r3, r4, r5 for target workload respectively.
+
+    .. code-block:: python
+
+        #### Case 1. `UnionDatabase`:
+        merged_db = ms.database.UnionDatabase(
+            db1, # no record
+            db2, # no record
+            db3, # has r3
+            db4  # has r4
+        )
+        # returns the better one between r3 and r4
+        merged_db.query_tuning_record(..., target_workload)
+
+        ### Case 2. `OrderedUnionDatabase`
+        merged_db = ms.database.OrderedUnionDatabase(
+            db1, # no record
+            db2, # no record
+            db3, # has r3
+            db4  # has r4
+        )
+        # returns r3
+        merged_db.query_tuning_record(..., target_workload)
+
+        ### Case 3. Mix-use scenario
+        merged_db = ms.database.UnionDatabase(
+            db1, # no record
+            db2, # no record
+            db3, # has r3
+            ms.database.OrderedUnionDatabase( # returns r4
+                db4,  # has r4
+                db5,  # has r5
+            )
+        )
+        # returns the better one between r3 and r4
+        merged_db.query_tuning_record(..., target_workload)
+
+        ### Case 4. Another mix-use scenario
+        merged_db = ms.database.UnionDatabase(
+            db1, # no record
+            db2, # no record
+            db3, # has r3
+            ms.database.UnionDatabase( # returns best one between r4 and r5
+                db4,  # has r4
+                db5,  # has r5
+            )
+        )
+        # returns the best one among r3, r4 and r5
+        merged_db.query_tuning_record(..., target_workload)
+
+        ### Case 5. Yet another mix-use scenario
+        merged_db = ms.database.OrderedUnionDatabase(
+            db1, # no record
+            db2, # no record
+            ms.database.UnionDatabase( # returns best one between r3 and r4
+                db3, # has r3
+                db4,  # has r4
+            )
+            db5,  # has r5
+        )
+        # returns the better one between r3 and r4
+        merged_db.query_tuning_record(..., target_workload)
+    """
+
+    def __init__(self, *databases: Database) -> None:
+        """Construct a merged database from multiple databases.
+
+        Parameters
+        ----------
+        *databases : Database
+            The list of databases to combine.
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.DatabaseUnionDatabase,  # type: ignore # pylint: disable=no-member
+            databases,
+        )

python/tvm/s_tir/meta_schedule/extracted_task.py

@@ -0,0 +1,66 @@
+# 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.
+"""Extracted tasks from high-level IR."""
+
+from tvm_ffi import register_object
+
+from tvm.ir import IRModule
+from tvm.runtime import Object
+from tvm.target import Target
+
+from . import _ffi_api
+
+
+@register_object("s_tir.meta_schedule.ExtractedTask")
+class ExtractedTask(Object):
+    """A tuning task extracted from the high-level IR
+
+    Parameters
+    ----------
+    task_name : str
+        The name of the task extracted
+    mod : IRModule
+        The high-level IR
+    target: Target
+        Target information
+    dispatched : List[IRModule]
+        A list of low-level IRs that the high-level IR could potentially dispatch to
+    weight : int
+        The weight of the task
+    """
+
+    task_name: str
+    mod: IRModule
+    dispatched: list[IRModule]
+    weight: int
+
+    def __init__(
+        self,
+        task_name: str,
+        mod: IRModule,
+        target: Target,
+        dispatched: list[IRModule],
+        weight: int,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ExtractedTask,  # type: ignore # pylint: disable=no-member
+            task_name,
+            mod,
+            target,
+            dispatched,
+            weight,
+        )

python/tvm/s_tir/meta_schedule/feature_extractor/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+"""
+The tvm.s_tir.meta_schedule.feature_extractor package.
+Meta Schedule feature extractors that extracts features from
+measure candidates for use in cost model.
+"""
+
+from .feature_extractor import FeatureExtractor, PyFeatureExtractor
+from .per_store_feature import PerStoreFeature
+from .random_feature_extractor import RandomFeatureExtractor

python/tvm/s_tir/meta_schedule/feature_extractor/feature_extractor.py

@@ -0,0 +1,128 @@
+# 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.
+# ruff: noqa: RUF012
+"""Meta Schedule FeatureExtractor."""
+
+from collections.abc import Callable
+from typing import Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+from tvm.runtime._tensor import Tensor
+
+from .. import _ffi_api
+from ..search_strategy import MeasureCandidate
+from ..tune_context import TuneContext
+
+
+@register_object("s_tir.meta_schedule.FeatureExtractor")
+class FeatureExtractor(Object):
+    """Extractor for features from measure candidates for use in cost model."""
+
+    FeatureExtractorType = Union[Literal["per-store-feature"], "FeatureExtractor"]
+
+    def extract_from(
+        self, context: TuneContext, candidates: list[MeasureCandidate]
+    ) -> list[Tensor]:
+        """Extract features from the given measure candidate.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for feature extraction.
+        candidates : List[MeasureCandidate]
+            The measure candidates to extract features from.
+
+        Returns
+        -------
+        features : List[Tensor]
+            The feature tvm ndarray extracted.
+        """
+        result = _ffi_api.FeatureExtractorExtractFrom(  # type: ignore # pylint: disable=no-member
+            self, context, candidates
+        )
+        return result
+
+    @staticmethod
+    def create(
+        kind: Literal["per-store-feature"],
+        *args,
+        **kwargs,
+    ) -> "FeatureExtractor":
+        """Create a CostModel."""
+        from . import PerStoreFeature  # pylint: disable=import-outside-toplevel
+
+        if kind == "per-store-feature":
+            return PerStoreFeature(*args, **kwargs)  # type: ignore
+        raise ValueError(f"Unknown CostModel: {kind}")
+
+
+@register_object("s_tir.meta_schedule.PyFeatureExtractor")
+class _PyFeatureExtractor(FeatureExtractor):
+    """
+    A TVM object feature extractor to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyFeatureExtractor
+    """
+
+    def __init__(self, f_extract_from: Callable):
+        """Constructor."""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.FeatureExtractorPyFeatureExtractor,  # type: ignore # pylint: disable=no-member
+            f_extract_from,
+        )
+
+
+class PyFeatureExtractor:
+    """
+    An abstract feature extractor with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PyFeatureExtractor,
+        "methods": ["extract_from"],
+    }
+
+    def extract_from(
+        self, context: TuneContext, candidates: list[MeasureCandidate]
+    ) -> list[Tensor]:
+        """Extract features from the given measure candidate.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for feature extraction.
+        candidates : List[MeasureCandidate]
+            The measure candidates to extract features from.
+
+        Returns
+        -------
+        features : List[Tensor]
+            The feature tvm ndarray extracted.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/feature_extractor/per_store_feature.py

@@ -0,0 +1,68 @@
+# 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
+"""We extract one feature vector per BufferStoreNode statement in a TIR Stmt,
+so we call this feature as "per-store" feature.
+"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .feature_extractor import FeatureExtractor
+
+
+@register_object("s_tir.meta_schedule.PerStoreFeature")
+class PerStoreFeature(FeatureExtractor):
+    """PerStoreFeature extracts one feature vector per BufferStoreNode
+
+    Parameters
+    ----------
+    buffers_per_store : int
+        The number of buffers in each BufferStore; Pad or truncate if necessary.
+    arith_intensity_curve_num_samples : int
+        The number of samples used in the arithmetic intensity curve.
+    cache_line_bytes : int
+        The number of bytes in a cache line.
+    extract_workload : bool
+        Whether to extract features in the workload in tuning context or not.
+    """
+
+    buffers_per_store: int
+    """The number of buffers in each BufferStore; Pad or truncate if necessary."""
+    arith_intensity_curve_num_samples: int  # pylint: disable=invalid-name
+    """The number of samples used in the arithmetic intensity curve."""
+    cache_line_bytes: int
+    """The number of bytes in a cache line."""
+    extract_workload: bool
+    """Whether to extract features in the workload in tuning context or not."""
+    feature_vector_length: int
+    """Length of the feature vector."""
+
+    def __init__(
+        self,
+        buffers_per_store: int = 5,
+        arith_intensity_curve_num_samples: int = 10,
+        cache_line_bytes: int = 64,
+        extract_workload: bool = False,
+    ):
+        self.__init_handle_by_constructor__(
+            _ffi_api.FeatureExtractorPerStoreFeature,  # type: ignore # pylint: disable=no-member
+            buffers_per_store,
+            arith_intensity_curve_num_samples,
+            cache_line_bytes,
+            extract_workload,
+        )

python/tvm/s_tir/meta_schedule/feature_extractor/random_feature_extractor.py

@@ -0,0 +1,64 @@
+# 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.
+"""Random Feature Extractor."""
+
+import numpy as np  # type: ignore
+
+import tvm.runtime
+from tvm.ir.utils import derived_object
+
+from ..feature_extractor import PyFeatureExtractor
+from ..search_strategy import MeasureCandidate
+from ..tune_context import TuneContext
+
+
+@derived_object
+class RandomFeatureExtractor(PyFeatureExtractor):
+    """Random Feature Extractor
+
+    Parameters
+    ----------
+    feature_size : int
+        The size of each block's feature vector.
+    max_block_num : int
+        The maximum number of blocks in each schedule.
+    random_state : Union[Tuple[str, np.ndarray, int, int, float], dict]
+        The current random state of the f
+    """
+
+    feature_size: int
+    max_block_num: int
+    random_state: tuple[str, np.ndarray, int, int, float] | dict
+
+    def __init__(self, *, feature_size: int = 30, max_block_num: int = 5, seed=0):
+        super().__init__()
+        assert max_block_num >= 1, "Max block number must be greater or equal to one!"
+        self.max_block_num = max_block_num
+        self.feature_size = feature_size
+        np.random.seed(seed)
+        self.random_state = np.random.get_state()
+
+    def extract_from(
+        self, context: TuneContext, candidates: list[MeasureCandidate]
+    ) -> list[tvm.runtime.Tensor]:
+        np.random.set_state(self.random_state)
+        result = [
+            np.random.rand(np.random.randint(1, self.max_block_num + 1), self.feature_size)
+            for candidate in candidates
+        ]
+        self.random_state = np.random.get_state()
+        return [tvm.runtime.tensor(x) for x in result]

python/tvm/s_tir/meta_schedule/logging.py

@@ -0,0 +1,266 @@
+# 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.
+"""Logging interface in MetaSchedule"""
+
+import logging
+import logging.config
+import os
+import os.path as osp
+from collections.abc import Callable
+from logging import Logger
+from typing import Any
+
+
+def get_logger(name: str) -> Logger:
+    """Create or get a logger by its name. This is essentially a wrapper of python's native logger.
+
+    Parameters
+    ----------
+    name : str
+        The name of the logger.
+
+    Returns
+    -------
+    logger : Logger
+        The logger instance.
+    """
+    return logging.getLogger(name)
+
+
+def get_logging_func(logger: Logger) -> Callable[[int, str, int, str], None] | None:
+    """Get the logging function.
+
+    Parameters
+    ----------
+    logger : Logger
+        The logger instance.
+    Returns
+    -------
+    result : Optional[Callable]
+        The function to do the specified level of logging.
+    """
+    if logger is None:
+        return None
+
+    level2log = {
+        logging.DEBUG: logger.debug,
+        logging.INFO: logger.info,
+        logging.WARNING: logger.warning,
+        logging.ERROR: logger.error,
+        # logging.FATAL not included
+    }
+
+    def logging_func(level: int, filename: str, lineo: int, msg: str):
+        if level < 0:  # clear the output in notebook / console
+            from IPython.display import (  # type: ignore # pylint: disable=import-outside-toplevel
+                clear_output,
+            )
+
+            clear_output(wait=True)
+        else:
+            level2log[level](f"[{os.path.basename(filename)}:{lineo}] " + msg)
+
+    return logging_func
+
+
+def create_loggers(
+    log_dir: str,
+    params: list[dict[str, Any]],
+    logger_config: dict[str, Any] | None = None,
+    disable_existing_loggers: bool = False,
+):
+    """Create loggers from configuration"""
+    if logger_config is None:
+        config = {}
+    else:
+        config = logger_config
+
+    config.setdefault("loggers", {})
+    config.setdefault("handlers", {})
+    config.setdefault("formatters", {})
+
+    global_logger_name = "tvm.s_tir.meta_schedule"
+    global_logger = logging.getLogger(global_logger_name)
+    if global_logger.level is logging.NOTSET:
+        global_logger.setLevel(logging.DEBUG)
+    console_logging_level = logging._levelToName[  # pylint: disable=protected-access
+        global_logger.level
+    ]
+
+    config["loggers"].setdefault(
+        global_logger_name,
+        {
+            "level": logging.DEBUG,
+            "handlers": [handler.get_name() for handler in global_logger.handlers]
+            + [global_logger_name + ".console", global_logger_name + ".file"],
+            "propagate": False,
+        },
+    )
+    config["loggers"].setdefault(
+        "{logger_name}",
+        {
+            "level": "DEBUG",
+            "handlers": [
+                "{logger_name}.file",
+            ],
+            "propagate": False,
+        },
+    )
+    config["handlers"].setdefault(
+        global_logger_name + ".console",
+        {
+            "class": "logging.StreamHandler",
+            "stream": "ext://sys.stdout",
+            "formatter": "tvm.s_tir.meta_schedule.standard_formatter",
+            "level": console_logging_level,
+        },
+    )
+    config["handlers"].setdefault(
+        global_logger_name + ".file",
+        {
+            "class": "logging.FileHandler",
+            "filename": "{log_dir}/" + __name__ + ".task_scheduler.log",
+            "mode": "a",
+            "level": "DEBUG",
+            "formatter": "tvm.s_tir.meta_schedule.standard_formatter",
+        },
+    )
+    config["handlers"].setdefault(
+        "{logger_name}.file",
+        {
+            "class": "logging.FileHandler",
+            "filename": "{log_dir}/{logger_name}.log",
+            "mode": "a",
+            "level": "DEBUG",
+            "formatter": "tvm.s_tir.meta_schedule.standard_formatter",
+        },
+    )
+    config["formatters"].setdefault(
+        "tvm.s_tir.meta_schedule.standard_formatter",
+        {
+            "format": "%(asctime)s [%(levelname)s] %(message)s",
+            "datefmt": "%Y-%m-%d %H:%M:%S",
+        },
+    )
+
+    # set up dictConfig loggers
+    p_config = {"version": 1, "disable_existing_loggers": disable_existing_loggers}
+    for k, v in config.items():
+        if k in ["formatters", "handlers", "loggers"]:
+            p_config[k] = _batch_parameterize_config(v, params)  # type: ignore
+        else:
+            p_config[k] = v
+    logging.config.dictConfig(p_config)
+
+    # check global logger
+    if global_logger.level not in [logging.DEBUG, logging.INFO]:
+        global_logger.warning(
+            "Logging level set to %s, please set to logging.INFO"
+            " or logging.DEBUG to view full log.",
+            logging._levelToName[global_logger.level],  # pylint: disable=protected-access
+        )
+    global_logger.info("Logging directory: %s", log_dir)
+
+
+def _batch_parameterize_config(
+    config: dict[str, Any],
+    params: list[dict[str, str]],
+) -> dict[str, Any]:
+    """Parameterize the given configuration with multiple parameters sets.
+
+    Parameters
+    ----------
+    config : Dict[str, Any]
+        The given config dict.
+    Params : List[Dict[str, str]]
+        List of the given multiple parameters sets.
+
+    Returns
+    -------
+    result : Dict[str, Any]
+        The parameterized configuration.
+    """
+    results = {}
+    for name, cfg in config.items():
+        for p in params:
+            p_name = name.format(**p)
+            if p_name not in results:
+                p_cfg = _parameterize_config(cfg, p)
+                results[p_name] = p_cfg
+    return results
+
+
+def _parameterize_config(
+    config: dict[str, Any],
+    params: dict[str, str],
+) -> dict[str, Any]:
+    """Parameterize the given configuration.
+
+    Parameters
+    ----------
+    config : Dict[str, Any]
+        The given config dict.
+    Params : Dict[str, str]
+        The given parameters.
+
+    Returns
+    -------
+    result : Dict[str, Any]
+        The parameterized configuration.
+    """
+    result = {}
+    for k, v in config.items():
+        if isinstance(k, str):
+            k = k.format(**params)
+        if isinstance(v, str):
+            v = v.format(**params)
+        elif isinstance(v, dict):
+            v = _parameterize_config(v, params)
+        elif isinstance(v, list):
+            v = [t.format(**params) for t in v]
+        result[k] = v
+    return result
+
+
+def get_loggers_from_work_dir(
+    work_dir: str,
+    task_names: list[str],
+) -> list[Logger]:
+    """Create loggers from work directory
+
+    Parameters
+    ----------
+    work_dir : str
+        The work directory.
+    task_names : List[str]
+        The list of task names.
+
+    Returns
+    -------
+    loggers : List[Logger]
+        The list of loggers.
+    """
+    log_dir = osp.join(work_dir, "logs")
+    os.makedirs(log_dir, exist_ok=True)
+    pattern = __name__ + ".task_{i:0" + f"{len(str(len(task_names) - 1))}" + "d}_{name}"
+    # Very long names may need be clipped to prevent os errors, we use the first 100 characters.
+    loggers = [pattern.format(i=i, name=name[:100]) for i, name in enumerate(task_names)]
+    create_loggers(
+        log_dir=log_dir,
+        params=[{"log_dir": log_dir, "logger_name": logger} for logger in loggers],
+    )
+    return [get_logger(logger) for logger in loggers]

python/tvm/s_tir/meta_schedule/measure_callback/__init__.py

@@ -0,0 +1,23 @@
+# 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.
+"""The tvm.s_tir.meta_schedule.measure_callback package."""
+
+from .add_to_database import AddToDatabase
+from .measure_callback import MeasureCallback, PyMeasureCallback
+from .remove_build_artifact import RemoveBuildArtifact
+from .update_cost_model import UpdateCostModel

python/tvm/s_tir/meta_schedule/measure_callback/add_to_database.py

@@ -0,0 +1,31 @@
+# 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.
+"""A callback that adds the measurement results into the database"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .measure_callback import MeasureCallback
+
+
+@register_object("s_tir.meta_schedule.AddToDatabase")
+class AddToDatabase(MeasureCallback):
+    def __init__(self) -> None:
+        """A callback that adds the measurement results into the database"""
+        self.__init_handle_by_constructor__(
+            _ffi_api.MeasureCallbackAddToDatabase,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/measure_callback/measure_callback.py

@@ -0,0 +1,141 @@
+# 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.
+# ruff: noqa: RUF012
+"""Meta Schedule MeasureCallback."""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+
+from .. import _ffi_api
+from ..builder import BuilderResult
+from ..runner import RunnerResult
+from ..search_strategy import MeasureCandidate
+
+if TYPE_CHECKING:
+    from ..task_scheduler import TaskScheduler
+
+
+@register_object("s_tir.meta_schedule.MeasureCallback")
+class MeasureCallback(Object):
+    """Rules to apply after measure results is available."""
+
+    CallbackListType = Union[list["MeasureCallback"], "MeasureCallback", Literal["default"]]
+
+    def apply(
+        self,
+        task_scheduler: "TaskScheduler",
+        task_id: int,
+        measure_candidates: list[MeasureCandidate],
+        builder_results: list[BuilderResult],
+        runner_results: list[RunnerResult],
+    ) -> None:
+        """Apply a measure callback to the given schedule.
+
+        Parameters
+        ----------
+        task_scheduler: TaskScheduler
+            The task scheduler.
+        task_id: int
+            The task id.
+        measure_candidates: List[MeasureCandidate]
+            The measure candidates.
+        builder_results: List[BuilderResult]
+            The builder results by building the measure candidates.
+        runner_results: List[RunnerResult]
+            The runner results by running the built measure candidates.
+        """
+        return _ffi_api.MeasureCallbackApply(  # type: ignore # pylint: disable=no-member
+            self,
+            task_scheduler,
+            task_id,
+            measure_candidates,
+            builder_results,
+            runner_results,
+        )
+
+    @staticmethod
+    def create(kind: Literal["default"]) -> list["MeasureCallback"]:
+        """Create a list of measure callbacks."""
+        if kind == "default":
+            return _ffi_api.MeasureCallbackDefault()  # type: ignore # pylint: disable=no-member
+        raise ValueError(f"Unknown kind of MeasureCallback list: {kind}")
+
+
+@register_object("s_tir.meta_schedule.PyMeasureCallback")
+class _PyMeasureCallback(MeasureCallback):
+    """
+    A TVM object measure callback to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyMeasureCallback
+    """
+
+    def __init__(self, f_apply: Callable):
+        """Constructor."""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.MeasureCallbackPyMeasureCallback,  # type: ignore # pylint: disable=no-member
+            f_apply,
+        )
+
+
+class PyMeasureCallback:
+    """
+    An abstract measure callback with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PyMeasureCallback,
+        "methods": ["apply"],
+    }
+
+    def apply(
+        self,
+        task_scheduler: "TaskScheduler",
+        task_id: int,
+        measure_candidates: list[MeasureCandidate],
+        builder_results: list[BuilderResult],
+        runner_results: list[RunnerResult],
+    ) -> None:
+        """Apply a measure callback to the given schedule.
+
+        Parameters
+        ----------
+        task_scheduler: TaskScheduler
+            The task scheduler.
+        task_id: int
+            The task id.
+        measure_candidates: List[MeasureCandidate]
+            The measure candidates.
+        builder_results: List[BuilderResult]
+            The builder results by building the measure candidates.
+        runner_results: List[RunnerResult]
+            The runner results by running the built measure candidates.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/measure_callback/remove_build_artifact.py

@@ -0,0 +1,31 @@
+# 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.
+"""A callback that removes the build artifacts from the disk"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .measure_callback import MeasureCallback
+
+
+@register_object("s_tir.meta_schedule.RemoveBuildArtifact")
+class RemoveBuildArtifact(MeasureCallback):
+    def __init__(self) -> None:
+        """A callback that removes the build artifacts from the disk"""
+        self.__init_handle_by_constructor__(
+            _ffi_api.MeasureCallbackRemoveBuildArtifact,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/measure_callback/update_cost_model.py

@@ -0,0 +1,31 @@
+# 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.
+"""A measure callback that updates the cost model"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .measure_callback import MeasureCallback
+
+
+@register_object("s_tir.meta_schedule.UpdateCostModel")
+class UpdateCostModel(MeasureCallback):
+    def __init__(self) -> None:
+        """A measure callback that updates the cost model"""
+        self.__init_handle_by_constructor__(
+            _ffi_api.MeasureCallbackUpdateCostModel,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/mutator/__init__.py

@@ -0,0 +1,29 @@
+# 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.
+"""
+The tvm.s_tir.meta_schedule.mutator package.
+Meta Schedule mutator that mutates the trace to explore the
+design space.
+"""
+
+from .mutator import Mutator, PyMutator
+from .mutate_compute_location import MutateComputeLocation
+from .mutate_tile_size import MutateTileSize
+from .mutate_thread_binding import MutateThreadBinding
+from .mutate_parallel import MutateParallel
+from .mutate_unroll import MutateUnroll

python/tvm/s_tir/meta_schedule/mutator/mutate_compute_location.py

@@ -0,0 +1,32 @@
+# 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.
+"""A mutator that mutates the compute-at location decision of SampleComputeLocation"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .mutator import Mutator
+
+
+@register_object("s_tir.meta_schedule.MutateComputeLocation")
+class MutateComputeLocation(Mutator):
+    """A mutator that mutates the compute-at location decision of SampleComputeLocation"""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.MutatorMutateComputeLocation,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/mutator/mutate_parallel.py

@@ -0,0 +1,34 @@
+# 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.
+"""Mutator that mutates the parallel extent"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .mutator import Mutator
+
+
+@register_object("s_tir.meta_schedule.MutateParallel")
+class MutateParallel(Mutator):
+    """Mutator that mutates the parallel extent"""
+
+    def __init__(self, max_jobs_per_core: int) -> None:
+        """Mutator that mutates the parallel extent"""
+        self.__init_handle_by_constructor__(
+            _ffi_api.MutatorMutateParallel,  # type: ignore # pylint: disable=no-member
+            max_jobs_per_core,
+        )

python/tvm/s_tir/meta_schedule/mutator/mutate_thread_binding.py

@@ -0,0 +1,33 @@
+# 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.
+"""Mutator that mutates the thread binding extent"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .mutator import Mutator
+
+
+@register_object("s_tir.meta_schedule.MutateThreadBinding")
+class MutateThreadBinding(Mutator):
+    """Mutator that mutates the binding extent"""
+
+    def __init__(self) -> None:
+        """Mutator that mutates the binding extent"""
+        self.__init_handle_by_constructor__(
+            _ffi_api.MutateThreadBinding,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/mutator/mutate_tile_size.py

@@ -0,0 +1,32 @@
+# 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.
+"""Mutator that mutates the decision of instruction Sample-Perfect-Tile"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .mutator import Mutator
+
+
+@register_object("s_tir.meta_schedule.MutateTileSize")
+class MutateTileSize(Mutator):
+    """Mutator that mutates the decision of instruction Sample-Perfect-Tile"""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.MutatorMutateTileSize,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/mutator/mutate_unroll.py

@@ -0,0 +1,32 @@
+# 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.
+"""Mutator that mutates auto unroll step"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .mutator import Mutator
+
+
+@register_object("s_tir.meta_schedule.MutateUnroll")
+class MutateUnroll(Mutator):
+    """Mutator that mutates auto unroll step"""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.MutatorMutateUnroll,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/mutator/mutator.py

@@ -0,0 +1,188 @@
+# 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.
+# ruff: noqa: RUF012
+"""Meta Schedule Mutator."""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+from tvm.s_tir.schedule import Trace
+
+from .. import _ffi_api
+
+if TYPE_CHECKING:
+    from ..tune_context import TuneContext
+
+
+class Mutator(Object):
+    """Mutator is designed to mutate the trace to explore the design space."""
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        """Initialize the mutator with a tune context.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for initializing the mutator.
+        """
+        _ffi_api.MutatorInitializeWithTuneContext(  # type: ignore # pylint: disable=no-member
+            self, context
+        )
+
+    def apply(self, trace: Trace) -> Trace | None:
+        """Apply the mutator function to the given trace.
+
+        Parameters
+        ----------
+        trace : Trace
+            The given trace for mutation.
+
+        Returns
+        -------
+        trace : Optional[Trace]
+            None if mutator failed, otherwise return the mutated trace.
+        """
+        return _ffi_api.MutatorApply(self, trace, -1)  # type: ignore # pylint: disable=no-member
+
+    def clone(self) -> "Mutator":
+        """Clone the mutator.
+
+        Returns
+        -------
+        mutator : Mutator
+            The cloned mutator.
+        """
+        return _ffi_api.MutatorClone(self)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def create(
+        kind: Literal[
+            "llvm",
+            "cuda",
+            "cuda-tensorcore",
+            "hexagon",
+        ],
+    ) -> dict["Mutator", float]:
+        """Create a list of default mutators.
+
+        Parameters
+        ----------
+        kind : Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]
+            The kind of mutators.
+
+        Returns
+        -------
+        mutators : List[Mutator]
+            The list of mutators.
+        """
+        funcs = {
+            # pylint: disable=no-member
+            "llvm": _ffi_api.MutatorDefaultLLVM,  # type: ignore
+            "cuda": _ffi_api.MutatorDefaultCUDA,  # type: ignore
+            "cuda-tensorcore": _ffi_api.MutatorDefaultCUDATensorCore,  # type: ignore
+            "hexagon": _ffi_api.MutatorDefaultHexagon,  # type: ignore
+            # pylint: enable=no-member
+        }
+        for k, v in funcs.items():
+            if k == kind:
+                return v()
+        raise ValueError(f"Unsupported kind {kind} for mutator creation.")
+
+
+create = Mutator.create  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.PyMutator")
+class _PyMutator(Mutator):
+    """
+    A TVM object mutator to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyMutator
+    """
+
+    def __init__(
+        self,
+        f_initialize_with_tune_context: Callable | None = None,
+        f_apply: Callable | None = None,
+        f_clone: Callable | None = None,
+    ):
+        """Constructor."""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.MutatorPyMutator,  # type: ignore # pylint: disable=no-member
+            f_initialize_with_tune_context,
+            f_apply,
+            f_clone,
+        )
+
+
+class PyMutator:
+    """
+    An abstract mutator with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PyMutator,
+        "methods": ["_initialize_with_tune_context", "apply", "clone"],
+    }
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        """Initialize the mutator with a tune context.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for initializing the mutator.
+        """
+        raise NotImplementedError
+
+    def apply(self, trace: Trace, _) -> Trace | None:
+        """Apply the mutator function to the given trace.
+
+        Parameters
+        ----------
+        trace : Trace
+            The given trace for mutation.
+
+        Returns
+        -------
+        trace : Optional[Trace]
+            None if mutator failed, otherwise return the mutated trace.
+        """
+        raise NotImplementedError
+
+    def clone(self) -> Mutator:
+        """Clone the mutator.
+
+        Returns
+        -------
+        mutator : Mutator
+            The cloned mutator.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/post_optimization/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+"""
+The tvm.s_tir.meta_schedule.database package.
+The database that stores serialized tuning records and workloads
+"""
+
+from .post_opt import PostOpt
+from .droplet import Droplet
+from .space import Space
+from .utils import write_file, get_time

python/tvm/s_tir/meta_schedule/post_optimization/droplet.py

@@ -0,0 +1,135 @@
+# 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.
+
+"""Droplet algorithm"""
+
+import os
+
+import numpy as np  # type: ignore
+
+from .space import Space
+from .utils import get_time, write_file
+
+
+class Droplet:
+    """Tuner with droplet algorithm in Meta Schedule.
+
+    Parameters
+    ----------
+    json_file: str
+        json format file
+    target:
+        hardware target
+    log: str
+        path to save json file
+    trials: int
+        number of samples, the default is 100
+    pvalue: float
+        statistical value to confidence level, the default is 0.05
+    """
+
+    def __init__(self, json_file, workload_file, target, log, pvalue=0.05) -> None:
+        self.space = Space(json_file, workload_file, target)
+        self.final_log = write_file([json_file], log)
+        self.pvalue = pvalue
+        self.next = [(0, [0] * len(self.space.dims))]
+        best_avg, _ = get_time(log)
+        self.best_choice = [0, [0] * len(self.space.dims), best_avg]
+        self.count, self.execution, self.found_best_pos = 1, 1, True
+        self.total_execution = 1
+        if len(self.space.dims) > 0:
+            self.total_execution = max(self.space.dims)
+        self.dims, self.step = self.space.dims, 1
+        self.visited, self.batch = set([0]), max(os.cpu_count(), len(self.dims))
+
+    def next_batch(self, batch_size):
+        i, json_file_list = 0, []
+        while i < len(self.next):
+            if batch_size > 0 and self.count >= self.trials:
+                break
+            json_file_list.append(self.space.template(values=self.next[i][1], create=False))
+            i, self.count = i + 1, self.count + 1
+        return self.space.run(json_file_list, self.final_log)
+
+    def has_next(self):
+        return len(self.next) > 0 and self.found_best_pos
+
+    def tune(self, n_trial=100):
+        self.trials = n_trial
+        self.speculation()
+        while self.has_next():
+            res = self.next_batch(self.batch)
+            self.update(res)
+
+    def num_to_bin(self, value, factor=1):
+        bin_format = str(0) * (len(self.dims) - len(bin(value)[2:])) + bin(value)[2:]
+        return [int(i) * factor for i in bin_format]
+
+    def search_space(self, factor=1):
+        "create a search space"
+        search_space: list = []
+        for i in range(0, len(self.space.dims)):
+            if len(search_space) > self.batch - len(self.next):
+                break
+            space = self.num_to_bin(2**i, factor)
+            idx = self.space.knob2point(space)
+            if idx not in self.visited:
+                search_space.append(space)
+        return search_space
+
+    def next_pos(self, new_positions):
+        "returns the neighbors of the best solution"
+        next_set = []
+        for p in new_positions:
+            new_p = [
+                (x + y) % self.dims[i] if (x + y > 0) else 0
+                for i, (x, y) in enumerate(zip(p, self.best_choice[1]))
+            ]
+            idx_p = self.space.knob2point(new_p)
+            if idx_p not in self.visited:
+                self.visited.add(idx_p)
+                next_set.append((idx_p, new_p))
+        return next_set
+
+    def speculation(self):
+        # Gradient descending direction prediction and search space filling
+        while len(self.next) < self.batch and self.execution < self.total_execution:
+            self.next += self.next_pos(self.search_space(self.execution))
+            self.execution += self.step
+
+    def update(self, results):
+        """Update the values"""
+        self.found_best_pos, count_valids = False, 0
+        for i, res in enumerate(results):
+            if np.mean(self.best_choice[2]) > np.mean(res):
+                self.best_choice = [self.next[i][0], self.next[i][1], res]
+                self.found_best_pos = True
+            if np.mean(res) != 10000:
+                count_valids += 1
+
+        self.next = []
+
+        # stop, because all neighborhoods are invalid.
+        if count_valids == 0:
+            self.speculation()
+            self.found_best_pos = True
+            return
+
+        if self.found_best_pos:
+            self.next += self.next_pos(self.search_space())
+            self.execution = 1
+            self.speculation()

python/tvm/s_tir/meta_schedule/post_optimization/post_opt.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.
+"""Post optimization method"""
+
+import numpy as np  # type: ignore
+
+from tvm.target import Target
+
+from .droplet import Droplet
+from .utils import clean_file, get_time, read_cfg_file, write_file
+
+
+class PostOpt:
+    """PostOpt class
+
+    Parameters
+    ----------
+    work_dir : str
+        The working directory.
+    target: Target data
+        Target device information
+    trials: integer value
+        Max number of trials to execute the optimization
+    """
+
+    def __init__(self, work_dir: str, target: Target, trials: int = 100) -> None:
+        self.work_dir = work_dir
+        self.target = target
+        self.trials = trials
+
+    def run(self) -> None:
+        """Execute the post optimization"""
+
+        tuning_file = self.work_dir + "/database_tuning_record.json"
+        workload_file = self.work_dir + "/database_workload.json"
+
+        cfg = read_cfg_file(tuning_file, workload_file)
+
+        print("id | time MS (s) | time DPMS (s) | speedup")
+        for idx, layer in enumerate(cfg):
+            time, data, workload = cfg[layer]
+            ms_time = np.mean(time)
+
+            temp_log = f"{self.work_dir}/opt_{idx}.log"
+
+            # Run the exploitation by Droplet
+            droplet = Droplet(data, workload, self.target, temp_log)
+            droplet.tune(self.trials)
+
+            dpms_time, dpm_sol = get_time(temp_log)
+            dpms_time = np.mean(dpms_time)
+
+            speedup = ms_time / dpms_time
+
+            # save the best solution
+            write_file([dpm_sol], tuning_file, mode="a")
+
+            # show the perfomance
+            print(f"{idx:2d} | {ms_time:.10f} | {dpms_time:.10f} | {speedup:.2f}")
+
+            # clean the temporary files
+            clean_file(temp_log)

python/tvm/s_tir/meta_schedule/post_optimization/space.py

@@ -0,0 +1,260 @@
+# 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.
+"""The class of Space used to optimize the Meta parameters"""
+
+import json
+import random
+from copy import deepcopy
+from typing import Any
+
+import numpy as np  # type: ignore
+
+from tvm.s_tir import Schedule
+from tvm.s_tir import meta_schedule as ms
+from tvm.s_tir.meta_schedule.database import TuningRecord, Workload
+from tvm.s_tir.meta_schedule.utils import remove_build_dir
+from tvm.target import Target
+
+from .utils import write_file
+
+
+class Space:
+    """Space class
+
+    Parameters
+    ----------
+    data: json data
+        A json file template
+    workload: json data
+        A json file workload
+    target: Target data
+        Target device information
+    """
+
+    def __init__(self, data: Any, workload: Any, target: Target):
+        self.cfg = deepcopy(data)
+        self._id = data[0]
+        self.workload = Workload.from_json(workload)
+        self.target = target
+        self.dev = self.get_device_type(target)
+        self.total_dims = 0
+        self.dims: list[int] = []
+        self.start: list[int] = []
+        self.config_space: dict[str, list[int]] = dict()
+        self.create_space()
+
+    def __repr__(self) -> str:
+        """Print the config space"""
+        out = ""
+        for key in self.config_space:
+            out += f"{key}: dims={self.config_space[key]}\n"
+        out += f"Total dimensions: {self.total_dims}\n"
+        return out
+
+    def __str__(self) -> str:
+        """Print the config space"""
+        out = ""
+        for key in self.config_space:
+            out += f"{key}: dims={self.config_space[key]}\n"
+        out += f"Total dimensions: {self.total_dims}\n"
+        return out
+
+    def get_value(self, key, pos):
+        """Return the space"""
+        return self.config_space[key][pos]
+
+    def add_space(self, space_list: list, element_list: list, limit=10000) -> list[int]:
+        """Return a list without repeat and with limited value"""
+        new_list = element_list
+        for elem in space_list:
+            if elem not in new_list and elem <= limit:
+                new_list.append(elem)
+        return new_list
+
+    def knob2point(self, knob):
+        """Convert a array to point"""
+        point = 0
+        for j, k in enumerate(knob):
+            point += int(np.prod(self.dims[:j])) * k
+        return point
+
+    def point2knob(self, point):
+        """Convert point form (single integer) to knob (vector)"""
+        knob = []
+        for dim in self.dims:
+            knob.append(point % dim)
+            point //= dim
+        return knob
+
+    def power_of_two(self, min_value: int, max_value: int) -> list:
+        """Return power of two array in interval"""
+        return [1 << i for i in range(min_value, max_value + 1)]
+
+    def get_index(self, array: list, value: int):
+        """returns an index if it finds the value"""
+        for i in range(len(array)):
+            if array[i][0] == value:
+                return i
+        return -1
+
+    def template(self, values=None, create=True):
+        """Generate the template from the values"""
+        idx = -1
+        config = deepcopy(self.cfg[1])
+        for counter, cfg in enumerate(config[0][0]):
+            opt = cfg[0]
+            if opt == "Annotate":
+                ann_key = cfg[2]
+                if ann_key == ["meta_schedule.parallel"]:
+                    interval = self.power_of_two(5, 9)
+                elif ann_key == ["meta_schedule.vectorize"]:
+                    interval = self.power_of_two(4, 8)
+                elif ann_key == ["pragma_auto_unroll_max_step"]:
+                    interval = self.power_of_two(7, 11)
+                elif ann_key == ["meta_schedule.thread_extent_low_inclusive"]:
+                    interval = self.power_of_two(5, 6)
+                elif ann_key == ["meta_schedule.thread_extent_high_inclusive"]:
+                    interval = self.power_of_two(8, 12)
+                else:
+                    continue
+                idx += 1
+                key = f"ann_{idx}"
+                ann_value = cfg[1][1]
+                if create:
+                    self.config_space[key] = self.add_space(interval, [ann_value])
+                else:
+                    cfg[1][1] = self.get_value(key, values[idx])
+            elif opt == "SamplePerfectTile":
+                tile = config[0][1]
+                tile_idx = self.get_index(tile, counter)
+                tile_val = tile[tile_idx][1]
+                interval = self.power_of_two(1, 6)
+                for i in range(len(tile_val)):
+                    idx += 1
+                    key = f"sp_{counter}_{idx}"
+                    split = tile_val[i]
+                    if create:
+                        self.config_space[key] = self.add_space(interval, [split])
+                    else:
+                        config[0][1][tile_idx][1][i] = self.get_value(key, values[idx])
+            elif opt == "TransformLayout":
+                del config[0][0][counter]
+        if create:
+            return None
+        return config
+
+    def create_space(self):
+        """Create the space using Meta's space"""
+        self.template(create=True)
+        # print(self.config_space)
+        self.dims = []
+        for key in self.config_space:
+            self.dims.append(len(self.config_space[key]))
+        self.total_dims = 1
+        if len(self.dims) > 0:
+            for dim in self.dims:
+                self.total_dims *= dim
+
+    def get_device_type(self, target: Target) -> str:
+        """Get the device type string from a target.
+
+        Parameters
+        ----------
+        target : Target
+            The target to get the device type from.
+
+        Returns
+        -------
+        device_type : str
+            The device type string.
+        """
+        if target.kind.name == "llvm":
+            return "cpu"
+        elif target.kind.name == "cuda":
+            return "cuda"
+        else:
+            raise RuntimeError(f"Unsupported target kind for device type: {target.kind.name}")
+
+    def save_log(
+        self,
+        path: str,
+        record: ms.database.TuningRecord,
+        results: ms.runner.RunnerResult,
+    ) -> None:
+        """Save the log file"""
+        new_json = [self._id, record.as_json()]
+        new_json[1][1] = results
+        write_file([new_json], path, "a")
+
+    def run(
+        self,
+        json_file_list,
+        final_log,
+        timeout=10,
+        number=2,
+        repeat=3,
+        min_repeat_ms=0,
+        cpu_cache=False,
+    ):
+        """Execute a log file and save"""
+
+        builder = ms.builder.LocalBuilder(timeout_sec=timeout)
+        runner = ms.runner.LocalRunner(
+            evaluator_config=ms.runner.EvaluatorConfig(
+                number=number,
+                repeat=repeat,
+                min_repeat_ms=min_repeat_ms,
+                enable_cpu_cache_flush=cpu_cache,
+            ),
+        )
+
+        results = np.full(len(json_file_list), [10000], dtype=list)
+        records, mods = [], []
+        for i, cfg in enumerate(json_file_list):
+            try:
+                record = TuningRecord.from_json(json.loads(json.dumps(cfg)), self.workload)
+                sch = Schedule(self.workload.mod)
+                # In some layers this is a heavy impact in time cost, so
+                # I applied this only 25% of the samples.
+                remove_postproc = random.random() > 0.75
+                record.trace.apply_to_schedule(sch, remove_postproc=remove_postproc)
+                mods.append(sch.mod)
+                records.append(record)
+            except Exception:  # pylint: disable=broad-except, invalid-name
+                continue
+
+        builder_res = builder.build([ms.builder.BuilderInput(mod, self.target) for mod in mods])
+
+        for i, record in enumerate(records):
+            try:
+                inp = ms.runner.RunnerInput(
+                    builder_res[i].artifact_path,
+                    device_type=self.dev,
+                    args_info=ms.arg_info.TensorInfo.from_prim_func(mods[i]["main"]),
+                )
+                runner_res = runner.run([inp])[0].result()
+                results[i] = [v.value for v in runner_res.run_secs]  # type: ignore
+            except Exception:  # pylint: disable=broad-except, invalid-name
+                results[i] = [1e10]
+                continue
+
+            # save the solution in json file
+            self.save_log(final_log, record, results[i])
+
+            # clean up
+            remove_build_dir(builder_res[i].artifact_path)
+        return results

python/tvm/s_tir/meta_schedule/post_optimization/utils.py

@@ -0,0 +1,112 @@
+# 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.
+"""Utils file for exploitation schedule"""
+
+import json
+import os
+
+import numpy as np  # type: ignore
+
+
+def write_file(json_list: list, log: str = "/tmp/file.json", mode: str = "w") -> str:
+    """Write the log file
+
+    Parameters
+    ----------
+    json_list: list
+        The list input json
+    log: Optional[str]
+        Path destiny to save the log file
+    mode: Optional[str]
+        Mode save, "a" means append and "w" means write
+
+    Returns
+    -------
+    ret: str
+        log path file
+    """
+    with open(log, mode, encoding="utf-8") as outfile:
+        for j in json_list:
+            outfile.write(json.dumps(j) + "\n")
+    return log
+
+
+def clean_file(filename: str) -> None:
+    """Clean temporary files
+
+    Parameters
+    ----------
+    filename: str
+        The filepath with remove from the system
+    """
+    if os.path.isfile(filename):
+        os.remove(filename)
+
+
+def get_time(log: str) -> list:
+    """Get the time from the log file
+
+    Parameters
+    ----------
+    log: str
+        log file
+
+    Returns
+    -------
+    ret: list
+        A list with the best time and the json data
+    """
+    best_time = [1e10, None]
+    with open(log, encoding="utf-8") as log_file:
+        for line in log_file.readlines():
+            data = json.loads(line)
+            params = data[1]
+            time = params[1]
+            if np.mean(best_time[0]) > np.mean(time):
+                best_time = [time, data]
+    return best_time
+
+
+def read_cfg_file(path_tuning_file: str, path_workload_file: str) -> dict[int, list]:
+    """Colect the info from meta logfile
+
+    Parameters
+    ----------
+    log: str
+        The input log path with the meta parameter
+
+    Returns
+    -------
+    ret: dict[layer, Union[time, dict]]
+        Returns the best time, total time, and data
+    """
+    workload_list = []
+    with open(path_workload_file, encoding="utf-8") as log_file:
+        for line in log_file.readlines():
+            workload_list.append(json.loads(line))
+
+    cfg: dict[int, list] = dict()
+    with open(path_tuning_file, encoding="utf-8") as log_file:
+        for line in log_file.readlines():
+            data = json.loads(line)
+            layer = data[0]
+            params = data[1]
+            time = params[1]
+
+            if layer not in cfg.keys() or np.mean(cfg[layer][0]) > np.mean(time):
+                cfg[layer] = [time, data, workload_list[layer]]
+    return cfg

python/tvm/s_tir/meta_schedule/postproc/__init__.py

@@ -0,0 +1,30 @@
+# 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.
+"""The tvm.s_tir.meta_schedule.postproc package."""
+
+from .disallow_dynamic_loop import DisallowDynamicLoop
+from .disallow_async_strided_mem_copy import DisallowAsyncStridedMemCopy
+from .postproc import Postproc, PyPostproc
+from .rewrite_cooperative_fetch import RewriteCooperativeFetch
+from .rewrite_layout import RewriteLayout
+from .rewrite_parallel_vectorize_unroll import RewriteParallelVectorizeUnroll
+from .rewrite_reduction_block import RewriteReductionBlock
+from .rewrite_tensorize import RewriteTensorize
+from .rewrite_unbound_block import RewriteUnboundBlock
+from .verify_gpu_code import VerifyGPUCode
+from .verify_vtcm_limit import VerifyVTCMLimit

python/tvm/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.py

@@ -0,0 +1,32 @@
+# 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.
+"""A postprocessor that checks if the IRModule has any strided memory copies"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .postproc import Postproc
+
+
+@register_object("s_tir.meta_schedule.DisallowAsyncStridedMemCopy")
+class DisallowAsyncStridedMemCopy(Postproc):
+    """A postprocessor that disallows schedules that use async strided mem copies."""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocDisallowAsyncStridedMemCopy,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/postproc/disallow_dynamic_loop.py

@@ -0,0 +1,32 @@
+# 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.
+"""A postprocessor that checks if the IRModule has any loop with non-constant extent"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .postproc import Postproc
+
+
+@register_object("s_tir.meta_schedule.DisallowDynamicLoop")
+class DisallowDynamicLoop(Postproc):
+    """A postprocessor that checks if the IRModule has any loop with non-constant extent"""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocDisallowDynamicLoop,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/postproc/postproc.py

@@ -0,0 +1,182 @@
+# 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.
+# ruff: noqa: RUF012
+"""Meta Schedule Postproc."""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+from tvm.s_tir.schedule import Schedule
+
+from .. import _ffi_api
+
+if TYPE_CHECKING:
+    from ..tune_context import TuneContext
+
+
+@register_object("s_tir.meta_schedule.Postproc")
+class Postproc(Object):
+    """Rules to apply a postprocessor to a schedule."""
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        """Initialize the postprocessor with a tune context.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for initializing the postprocessor.
+        """
+        _ffi_api.PostprocInitializeWithTuneContext(  # type: ignore # pylint: disable=no-member
+            self, context
+        )
+
+    def apply(self, sch: Schedule) -> bool:
+        """Apply a postprocessor to the given schedule.
+
+        Parameters
+        ----------
+        sch : tvm.s_tir.Schedule
+            The schedule to be post processed.
+
+        Returns
+        -------
+        result : bool
+            Whether the postprocessor was successfully applied.
+        """
+        return _ffi_api.PostprocApply(self, sch)  # type: ignore # pylint: disable=no-member
+
+    def clone(self) -> "Postproc":
+        """Clone the postprocessor.
+
+        Returns
+        -------
+        cloned_postproc : Postproc
+            The cloned postprocessor.
+        """
+        return _ffi_api.PostprocClone(self)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def create(kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]) -> list["Postproc"]:
+        """Create a list of default postprocessors.
+
+        Parameters
+        ----------
+        kind : Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]
+            The kind of the postprocessors.
+
+        Returns
+        -------
+        postprocs : List[Mutator]
+            The list of postprocessors.
+        """
+        funcs = {
+            # pylint: disable=no-member
+            "llvm": _ffi_api.PostprocDefaultLLVM,  # type: ignore
+            "cuda": _ffi_api.PostprocDefaultCUDA,  # type: ignore
+            "cuda-tensorcore": _ffi_api.PostprocDefaultCUDATensorCore,  # type: ignore
+            "hexagon": _ffi_api.PostprocDefaultHexagon,  # type: ignore
+            # pylint: enable=no-member
+        }
+        for k, v in funcs.items():
+            if k == kind:
+                return v()
+        raise ValueError(f"Unsupported kind {kind} for postproc creation.")
+
+
+create = Postproc.create  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.PyPostproc")
+class _PyPostproc(Postproc):
+    """
+    A TVM object post processor to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyPostproc
+    """
+
+    def __init__(
+        self,
+        f_initialize_with_tune_context: Callable | None = None,
+        f_apply: Callable | None = None,
+        f_clone: Callable | None = None,
+    ):
+        """Constructor."""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocPyPostproc,  # type: ignore # pylint: disable=no-member
+            f_initialize_with_tune_context,
+            f_apply,
+            f_clone,
+        )
+
+
+class PyPostproc:
+    """
+    An abstract post processor with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PyPostproc,
+        "methods": ["_initialize_with_tune_context", "apply", "clone"],
+    }
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        """Initialize the postprocessor with a tune context.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for initializing the postprocessor.
+        """
+        raise NotImplementedError
+
+    def apply(self, sch: Schedule) -> bool:
+        """Apply a postprocessor to the given schedule.
+
+        Parameters
+        ----------
+        sch : Schedule
+            The schedule to be post processed.
+
+        Returns
+        -------
+        result : bool
+            Whether the postprocessor was successfully applied.
+        """
+        raise NotImplementedError
+
+    def clone(self) -> Postproc:
+        """Clone the postprocessor.
+
+        Returns
+        -------
+        cloned_postproc : Postproc
+            The cloned postprocessor.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.py

@@ -0,0 +1,35 @@
+# 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.
+"""A postprocessor that rewrites the cooperative fetch annotation to actual
+vectorized cooperative fetching in loop bindings."""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .postproc import Postproc
+
+
+@register_object("s_tir.meta_schedule.RewriteCooperativeFetch")
+class RewriteCooperativeFetch(Postproc):
+    """A postprocessor that rewrites the cooperative fetch annotation to actual vectorized
+    cooperative fetching in loop bindings.
+    """
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocRewriteCooperativeFetch,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/postproc/rewrite_layout.py

@@ -0,0 +1,32 @@
+# 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.
+"""A postprocessor that rewrites the layout of input tensor"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .postproc import Postproc
+
+
+@register_object("s_tir.meta_schedule.RewriteLayout")
+class RewriteLayout(Postproc):
+    """A postprocessor that rewrites the layout of input tensor"""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocRewriteLayout,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.py

@@ -0,0 +1,34 @@
+# 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.
+"""A postprocessor that applies parallelization, vectorization and auto unrolling
+according to the annotation of each block"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .postproc import Postproc
+
+
+@register_object("s_tir.meta_schedule.RewriteParallelVectorizeUnroll")
+class RewriteParallelVectorizeUnroll(Postproc):
+    """A postprocessor that applies parallelization, vectorization and auto unrolling
+    according to the annotation of each block"""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocRewriteParallelVectorizeUnroll,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/postproc/rewrite_reduction_block.py

@@ -0,0 +1,32 @@
+# 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.
+"""A postprocessor that rewrites reduction block by moving the init block out."""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .postproc import Postproc
+
+
+@register_object("s_tir.meta_schedule.RewriteReductionBlock")
+class RewriteReductionBlock(Postproc):
+    """A postprocessor that rewrites reduction block by moving the init block out."""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocRewriteReductionBlock,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/postproc/rewrite_tensorize.py

@@ -0,0 +1,39 @@
+# 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.
+"""A postprocessor that tensorize related components."""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .postproc import Postproc
+
+
+@register_object("s_tir.meta_schedule.RewriteTensorize")
+class RewriteTensorize(Postproc):
+    """A postprocessor that applies tensorization to annotated blocks.
+
+    Parameters
+    ----------
+    vectorize_init_loop : bool
+       Whether or not vectorize the initialization loop produced by DecomposeReduction
+    """
+
+    def __init__(self, vectorize_init_loop=False) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocRewriteTensorize,  # type: ignore # pylint: disable=no-member
+            vectorize_init_loop,
+        )

python/tvm/s_tir/meta_schedule/postproc/rewrite_unbound_block.py

@@ -0,0 +1,33 @@
+# 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.
+"""A postprocessor that adds thread binding to unbound blocks"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .postproc import Postproc
+
+
+@register_object("s_tir.meta_schedule.RewriteUnboundBlock")
+class RewriteUnboundBlock(Postproc):
+    """A postprocessor that adds thread binding to unbound blocks"""
+
+    def __init__(self, max_threadblocks: int = 256) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocRewriteUnboundBlock,  # type: ignore # pylint: disable=no-member
+            max_threadblocks,
+        )

python/tvm/s_tir/meta_schedule/postproc/verify_gpu_code.py

@@ -0,0 +1,32 @@
+# 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.
+"""A postprocessor that verifies if the GPU code is correct"""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .postproc import Postproc
+
+
+@register_object("s_tir.meta_schedule.VerifyGPUCode")
+class VerifyGPUCode(Postproc):
+    """A postprocessor that verifies if the GPU code is correct"""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocVerifyGPUCode,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/postproc/verify_vtcm_limit.py

@@ -0,0 +1,32 @@
+# 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.
+"""A postprocessor that verifies the VTCM usage of a given schedule."""
+
+from tvm_ffi.registry import register_object
+
+from .. import _ffi_api
+from .postproc import Postproc
+
+
+@register_object("s_tir.meta_schedule.VerifyVTCMLimit")
+class VerifyVTCMLimit(Postproc):
+    """Verifies that the VTCM usage of a given schedule is within the provided limit."""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.PostprocVerifyVTCMLimit,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/profiler.py

@@ -0,0 +1,74 @@
+# 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=used-before-assignment
+"""A context manager that profiles tuning time cost for different parts."""
+
+from contextlib import contextmanager
+from typing import Optional
+
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+@register_object("s_tir.meta_schedule.Profiler")
+class Profiler(Object):
+    """Tuning time profiler."""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.Profiler,  # type: ignore # pylint: disable=no-member
+        )
+
+    def get(self) -> dict[str, float]:
+        """Get the profiling results in seconds"""
+        return _ffi_api.ProfilerGet(self)  # type: ignore # pylint: disable=no-member
+
+    def table(self) -> str:
+        """Get the profiling results in a table format"""
+        return _ffi_api.ProfilerTable(self)  # type: ignore # pylint: disable=no-member
+
+    def __enter__(self) -> "Profiler":
+        """Entering the scope of the context manager"""
+        _ffi_api.ProfilerEnterWithScope(self)  # type: ignore # pylint: disable=no-member
+        return self
+
+    def __exit__(self, ptype, value, trace) -> None:
+        """Exiting the scope of the context manager"""
+        _ffi_api.ProfilerExitWithScope(self)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def current() -> Optional["Profiler"]:
+        """Get the current profiler."""
+        return _ffi_api.ProfilerCurrent()  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def timeit(name: str):
+        """Timeit a block of code"""
+
+        @contextmanager
+        def _timeit():
+            try:
+                f = _ffi_api.ProfilerTimedScope(name)  # type: ignore # pylint: disable=no-member
+                yield
+            finally:
+                if f:
+                    f()
+
+        return _timeit()

python/tvm/s_tir/meta_schedule/relax_integration.py

@@ -0,0 +1,494 @@
+# 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.
+"""Meta schedule integration with high-level IR"""
+
+import warnings
+from typing import TYPE_CHECKING, Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+from tvm_ffi import get_global_func, register_global_func
+
+from tvm.ir import IRModule
+from tvm.ir.transform import PassContext
+from tvm.runtime import Tensor
+from tvm.target import Target
+from tvm.tirx.expr import IntImm
+
+from .builder import Builder
+from .cost_model import CostModel
+from .database import Database
+from .extracted_task import ExtractedTask
+from .logging import get_loggers_from_work_dir
+from .measure_callback import MeasureCallback
+from .runner import Runner
+from .search_strategy import SearchStrategy
+from .space_generator import SpaceGenerator
+from .task_scheduler import TaskScheduler
+from .tune import tune_tasks
+from .tune_context import TuneContext
+from .utils import fork_seed
+
+if TYPE_CHECKING:
+    from tvm import relax
+
+_extract_task_func = get_global_func(  # pylint: disable=invalid-name
+    "relax.backend.MetaScheduleExtractTask",
+    allow_missing=True,
+)
+
+
+def extract_tasks(
+    mod: Union[IRModule, "relax.Function"],
+    target: Target,
+    params: dict[str, Tensor] | None = None,
+    module_equality: str = "structural",
+) -> list[ExtractedTask]:
+    """Extract tuning tasks from a relax program.
+
+    Parameters
+    ----------
+    mod : Union[IRModule, relax.Function]
+        The module or function to tune
+    target : tvm.target.Target
+        The compilation target
+    params : Optional[Dict[str, tvm.runtime.Tensor]]
+        The associated parameters of the program
+    module_equality : Optional[str]
+        A string to specify the module equality testing and hashing method.
+        It must be one of the followings:
+
+          - "structural": Use StructuralEqual/Hash
+          - "ignore-tensor": Same as "structural", but ignore tensor raw data during
+                              equality testing and hashing.
+          - "anchor-block": Apply equality testing and hashing on the anchor block extracted from a
+                            given module. The "ignore-tensor" varint is used for the extracted
+                            blocks or in case no anchor block is found.
+                            For the definition of the anchor block, see tirx/analysis/analysis.py.
+
+    Returns
+    -------
+    tasks: List[ExtractedTask]
+        The tasks extracted from this module
+    """
+    # pylint: disable=import-outside-toplevel
+    from tvm.relax.expr import Function as RelaxFunc
+    from tvm.relax.transform import BindParams
+
+    # pylint: enable=import-outside-toplevel
+    if isinstance(mod, RelaxFunc):
+        mod = IRModule({"main": mod})
+    if not isinstance(target, Target):
+        target = Target(target)
+    if params:
+        mod = BindParams("main", params)(mod)
+    return list(_extract_task_func(mod, target, module_equality))
+
+
+def extracted_tasks_to_tune_contexts(
+    extracted_tasks: list[ExtractedTask],
+    work_dir: str,
+    space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+    strategy: SearchStrategy.SearchStrategyType = "evolutionary",
+    num_threads: Literal["physical", "logical"] | int = "physical",
+    seed: int | None = None,
+) -> tuple[list[TuneContext], list[float]]:
+    """Convert ExtractedTask to TuneContext.
+
+    Parameters
+    ----------
+    tasks : List[ExtractedTask]
+        The tasks to be converted
+    work_dir : str
+        The working directory to store logs and databases
+    space : SpaceGenerator.SpaceGeneratorType
+        The space generator to use.
+    strategy : SearchStrategy.SearchStrategyType
+        The search strategy to use.
+    num_threads : Union[Literal["physical", "logical"], int]
+        The number of threads to use in multi-threaded search algorithm.
+    seed : Optional[int]
+        The random seed to use.
+
+    Returns
+    -------
+    tasks : List[TuneContext]
+        The converted tasks
+    task_weights : List[float]
+        The weights of the tasks
+    """
+    tasks: list[TuneContext] = []
+    task_weights: list[float] = []
+    for task, logger, rand_state in zip(
+        extracted_tasks,
+        get_loggers_from_work_dir(work_dir, [t.task_name for t in extracted_tasks]),
+        fork_seed(seed, n=len(extracted_tasks)),
+    ):
+        if task.mod.attrs.get("tirx.is_scheduled", False):
+            warnings.warn("The task {task.task_name} is already scheduled, skipping it.")
+            continue
+        tasks.append(
+            TuneContext(
+                mod=task.dispatched[0],
+                target=task.target,
+                space_generator=space,
+                search_strategy=strategy,
+                task_name=task.task_name,
+                logger=logger,
+                rand_state=rand_state,
+                num_threads=num_threads,
+            ).clone()
+        )
+        task_weights.append(task.weight)
+    return tasks, task_weights
+
+
+def tune_relax(
+    mod: Union[IRModule, "relax.Function"],
+    params: dict[str, Tensor],
+    target: str | Target,
+    work_dir: str,
+    max_trials_global: int,
+    max_trials_per_task: int | None = None,
+    op_names: list[str] | None = None,
+    *,
+    num_trials_per_iter: int = 64,
+    builder: Builder.BuilderType = "local",
+    runner: Runner.RunnerType = "local",
+    database: Database.DatabaseType = "json",
+    cost_model: CostModel.CostModelType = "xgb",
+    measure_callbacks: MeasureCallback.CallbackListType = "default",
+    task_scheduler: TaskScheduler.TaskSchedulerType = "gradient",
+    space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+    strategy: SearchStrategy.SearchStrategyType = "evolutionary",
+    seed: int | None = None,
+    module_equality: str = "structural",
+) -> Database:
+    """Tune a Relax program.
+
+    Parameters
+    ----------
+    mod : Union[IRModule, relax.Function]
+        The module or function to tune
+    params : Optional[Dict[str, tvm.runtime.Tensor]]
+        The associated parameters of the program
+    target : Union[Target, str]
+        The compilation target
+    work_dir : str
+        The working directory to store the tuning records
+    max_trials_global : int
+        The maximum number of trials to run
+    max_trials_per_task : Optional[int]
+        The maximum number of trials to run for each task
+    op_names: Optional[List[str]]
+        A list of operator names to specify which op to tune. When it is None, all operators
+        are tuned.
+    num_trials_per_iter : int
+        The number of trials to run per iteration
+    builder : BuilderType
+        The builder to use
+    runner : RunnerType
+        The runner to use
+    database : DatabaseType
+        The database to use
+    cost_model : CostModelType
+        The cost model to use
+    measure_callbacks : CallbackListType
+        The measure callbacks to use
+    task_scheduler : TaskSchedulerType
+        The task scheduler to use
+    space : SpaceGeneratorType
+        The space generator to use
+    strategy : SearchStrategyType
+        The search strategy to use
+    seed : Optional[int]
+        The random seed
+    module_equality : Optional[str]
+        A string to specify the module equality testing and hashing method.
+        It must be one of the followings:
+
+          - "structural": Use StructuralEqual/Hash
+          - "ignore-tensor": Same as "structural", but ignore tensor raw data during
+                              equality testing and hashing.
+          - "anchor-block": Apply equality testing and hashing on the anchor block extracted from a
+                            given module. The "ignore-tensor" variant is used for the extracted
+                            blocks or in case no anchor block is found.
+                            For the definition of the anchor block, see tirx/analysis/analysis.py.
+
+    Returns
+    -------
+    database : Database
+        The database that contains the tuning records
+    """
+    all_tasks = extract_tasks(mod, target, params, module_equality=module_equality)
+
+    if not op_names:
+        selected_tasks = all_tasks
+    else:
+        selected_tasks = []
+
+        for task in all_tasks:
+            for op_name in op_names:
+                if op_name in task.task_name:
+                    selected_tasks.append(task)
+
+    tasks, task_weights = extracted_tasks_to_tune_contexts(
+        extracted_tasks=selected_tasks,
+        work_dir=work_dir,
+        space=space,
+        strategy=strategy,
+        seed=seed,
+    )
+    return tune_tasks(
+        tasks=tasks,
+        task_weights=task_weights,
+        work_dir=work_dir,
+        max_trials_global=max_trials_global,
+        max_trials_per_task=max_trials_per_task,
+        num_trials_per_iter=num_trials_per_iter,
+        builder=builder,
+        runner=runner,
+        database=database,
+        cost_model=cost_model,
+        measure_callbacks=measure_callbacks,
+        task_scheduler=task_scheduler,
+        module_equality=module_equality,
+    )
+
+
+@register_global_func("tvm.s_tir.meta_schedule.tune_relax")
+def _tune_relax(
+    mod: Union[IRModule, "relax.Function"],
+    params: dict[str, Tensor],
+    target: str | Target,
+    work_dir: str,
+    max_trials_global: int,
+    max_trials_per_task: int | None = None,
+    op_names: list[str] | None = None,
+    *,
+    num_trials_per_iter: int = 64,
+    builder: Builder.BuilderType = "local",
+    runner: Runner.RunnerType = "local",
+    database: Database.DatabaseType = "json",
+    cost_model: CostModel.CostModelType = "xgb",
+    measure_callbacks: MeasureCallback.CallbackListType = "default",
+    task_scheduler: TaskScheduler.TaskSchedulerType = "gradient",
+    space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+    strategy: SearchStrategy.SearchStrategyType = "evolutionary",
+    seed: int | None = None,
+    module_equality: str = "structural",
+) -> Database:
+    """Interface with tuning api to tune a Relax program.
+
+    Parameters
+    ----------
+    mod : Union[IRModule, relax.Function]
+        The module or function to tune
+    params : Optional[Dict[str, tvm.runtime.Tensor]]
+        The associated parameters of the program
+    target : Union[Target, str]
+        The compilation target
+    work_dir : str
+        The working directory to store the tuning records
+    max_trials_global : int
+        The maximum number of trials to run
+    max_trials_per_task : Optional[int]
+        The maximum number of trials to run for each task
+    op_names: Optional[List[str]]
+        A list of operator names to specify which op to tune. When it is None, all operators
+        are tuned.
+    num_trials_per_iter : int
+        The number of trials to run per iteration
+    builder : BuilderType
+        The builder to use
+    runner : RunnerType
+        The runner to use
+    database : DatabaseType
+        The database to use
+    cost_model : CostModelType
+        The cost model to use
+    measure_callbacks : CallbackListType
+        The measure callbacks to use
+    task_scheduler : TaskSchedulerType
+        The task scheduler to use
+    space : SpaceGeneratorType
+        The space generator to use
+    strategy : SearchStrategyType
+        The search strategy to use
+    seed : Optional[int]
+        The random seed
+    module_equality : Optional[str]
+        A string to specify the module equality testing and hashing method.
+        It must be one of the followings:
+
+          - "structural": Use StructuralEqual/Hash
+          - "ignore-tensor": Same as "structural", but ignore tensor raw data during
+                              equality testing and hashing.
+          - "anchor-block": Apply equality testing and hashing on the anchor block extracted from a
+                            given module. The "ignore-tensor" varint is used for the extracted
+                            blocks or in case no anchor block is found.
+                            For the definition of the anchor block, see tirx/analysis/analysis.py.
+
+    Returns
+    -------
+    ret_mod : IRModule
+        IRModule
+    """
+    if isinstance(max_trials_global, IntImm):
+        max_trials_global = int(max_trials_global)
+    if isinstance(max_trials_per_task, IntImm):
+        max_trials_per_task = int(max_trials_per_task)
+
+    tune_relax(
+        mod,
+        params,
+        target,
+        work_dir,
+        max_trials_global,
+        max_trials_per_task=max_trials_per_task,
+        num_trials_per_iter=num_trials_per_iter,
+        op_names=op_names,
+        builder=builder,
+        runner=runner,
+        database=database,
+        cost_model=cost_model,
+        measure_callbacks=measure_callbacks,
+        task_scheduler=task_scheduler,
+        space=space,
+        strategy=strategy,
+        seed=seed,
+        module_equality=module_equality,
+    )
+    # Return original IRModule
+    # This pass only makes optimization decision
+    return mod
+
+
+def compile_relax(
+    database: Database,
+    mod: IRModule,
+    target: Target | str,
+    params: dict[str, Tensor] | None,
+    enable_warning: bool = False,
+) -> "relax.VMExecutable":
+    """Compile a relax program with a MetaSchedule database.
+
+    Parameters
+    ----------
+    database : Database
+        The database to use
+    mod : IRModule
+        The Relax program to be compiled
+    target : tvm.target.Target
+        The compilation target
+    params : Optional[Dict[str, tvm.runtime.Tensor]]
+        The associated parameters of the program
+    enable_warning : bool
+        A boolean value indicating if to print warnings for TIR functions not
+        showing up in the database. By default we don't print warning.
+
+    Returns
+    -------
+    lib : relax.VMExecutable
+        The built runtime module or vm VMExecutable for the given relax workload.
+    """
+    # pylint: disable=import-outside-toplevel
+    import tvm
+    from tvm import relax
+    from tvm.relax import build as relax_build
+    from tvm.relax import pipeline as relax_pipeline_mod
+    from tvm.relax.transform import BindParams, MetaScheduleApplyDatabase
+    from tvm.s_tir import dlight as dl
+
+    # pylint: enable=import-outside-toplevel
+    if not isinstance(target, Target):
+        target = Target(target)
+    if params:
+        mod = BindParams("main", params)(mod)
+
+    # Build a pipeline with the correct ordering:
+    #   1. library_dispatch + LegalizeOps + FuseOps + FuseTIR
+    #      (same preparation as extract_tasks, so database keys match)
+    #   2. MetaScheduleApplyDatabase — replaces tuned fused-TIR functions
+    #   3. DLight fallback — schedules remaining untuned functions
+    #   4. dataflow_lower + finalize passes
+    #
+    # Applying MetaScheduleApplyDatabase BEFORE FuseOps (the original bug)
+    # caused DLight.Matmul to fail on cache-write stages embedded in fused TIR.
+    #
+    # All pass lists are obtained from relax.pipeline.*_passes(target) so that
+    # target-specific helpers (dispatch, finalize) are shared with the default
+    # pipeline rather than duplicated here.
+    try:
+        dispatch_passes = relax_pipeline_mod.library_dispatch_passes(target)
+    except (ValueError, AttributeError):
+        dispatch_passes = []
+
+    try:
+        lower_passes = relax_pipeline_mod.dataflow_lower_passes(target)
+        finalize_passes = relax_pipeline_mod.finalize_passes(target)
+    except (ValueError, AttributeError):
+        # Fallback for targets not yet registered in the pipeline dispatcher
+        lower_passes = [
+            relax.transform.RewriteDataflowReshape(),
+            relax.transform.ToNonDataflow(),
+            relax.transform.RemovePurityChecking(),
+            relax.transform.CallTIRRewrite(),
+        ]
+        finalize_passes = [
+            relax.transform.StaticPlanBlockMemory(),
+            relax.transform.LowerAllocTensor(),
+            relax.transform.KillAfterLastUse(),
+            relax.transform.LowerRuntimeBuiltin(),
+            relax.transform.ComputePrimValue(),
+            relax.transform.VMShapeLower(),
+            relax.transform.AttachGlobalSymbol(),
+        ]
+
+    is_gpu_target = relax_pipeline_mod.BackendDispatcher.is_gpu_target(target)
+
+    @tvm.transform.module_pass(opt_level=3)
+    def _ms_pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
+        fuse_seq = [
+            *dispatch_passes,
+            relax.transform.LegalizeOps(enable_warning=enable_warning),
+            relax.transform.AnnotateTIROpPattern(),
+            relax.transform.FoldConstant(),
+            relax.transform.FuseOps(),
+            relax.transform.FuseTIR(),
+        ]
+        mod = tvm.transform.Sequential(fuse_seq)(mod)
+        mod = MetaScheduleApplyDatabase(enable_warning=enable_warning)(mod)
+        # DLight handles functions not covered by the database.
+        # GPU rules apply only for GPU targets.
+        if is_gpu_target:
+            mod = dl.ApplyDefaultSchedule(
+                dl.gpu.Matmul(),
+                dl.gpu.GEMV(),
+                dl.gpu.Reduction(),
+                dl.gpu.GeneralReduction(),
+                dl.gpu.Fallback(),
+            )(mod)
+        mod = tvm.transform.Sequential(lower_passes + finalize_passes)(mod)
+        return mod
+
+    with target, database, PassContext(opt_level=3):
+        relax_ex = relax_build(mod, target=target, relax_pipeline=_ms_pipeline)
+    return relax_ex

python/tvm/s_tir/meta_schedule/runner/__init__.py

@@ -0,0 +1,34 @@
+# 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.
+"""
+The tvm.s_tir.meta_schedule.runner package.
+Meta Schedule runners that runs an artifact either locally or through the RPC interface
+"""
+
+from .config import EvaluatorConfig, RPCConfig
+from .local_runner import LocalRunner, LocalRunnerFuture
+from .rpc_runner import RPCRunner
+from .runner import (
+    PyRunner,
+    PyRunnerFuture,
+    Runner,
+    RunnerFuture,
+    RunnerInput,
+    RunnerResult,
+    create,
+)

python/tvm/s_tir/meta_schedule/runner/config.py

@@ -0,0 +1,201 @@
+# 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.
+"""Configurations for measurements in the runner"""
+
+import os
+from threading import Thread
+from typing import NamedTuple, Optional
+
+from tvm import rpc
+
+
+class EvaluatorConfig(NamedTuple):
+    """Config Details of Evaluator
+
+    Parameters
+    ----------
+    number: int
+        The number of times to run this function for taking average.
+        We call these runs as one `repeat` of measurement.
+    repeat: int
+        The number of times to repeat the measurement.
+        In total, the function will be invoked (1 + number x repeat) times,
+        where the first one is warm up and will be discarded.
+        The returned result contains `repeat` costs,
+        each of which is an average of `number` costs.
+    min_repeat_ms: int
+        Minimum repeat time in ms. if the execution latency is too short,
+        increase the number of runs to the given time (in ms) to reduce the measurement error.
+    enable_cpu_cache_flush: bool
+        Whether to flush the cache on CPU.
+
+    Note
+    ----
+    The total number of actual executions is 1+number*repeat because we would warm up 1 time before
+    actual run. The number of runs would be increased if run time is below min_repeat_ms.
+    """
+
+    number: int = 3
+    repeat: int = 1
+    min_repeat_ms: int = 100
+    enable_cpu_cache_flush: bool = False
+
+    @staticmethod
+    def _normalized(config: Optional["EvaluatorConfig"]) -> "EvaluatorConfig":
+        if config is None:
+            return EvaluatorConfig()
+        config = EvaluatorConfig(
+            number=config.number,
+            repeat=config.repeat,
+            min_repeat_ms=config.min_repeat_ms,
+            enable_cpu_cache_flush=config.enable_cpu_cache_flush,
+        )
+        return config
+
+
+class RPCConfig(NamedTuple):
+    """RPC configuration
+
+    Parameters
+    ----------
+    tracker_host: str
+        Host of the RPC Tracker
+    tracker_port: int
+        Port of the RPC Tracker
+    tracker_key: str
+        Key of the Tracker
+    session_timeout_sec: float
+        Timeout of the RPC session
+    session_priority: int
+        Priority of the RPC session
+    """
+
+    tracker_host: str | None = None
+    tracker_port: None | int | str = None
+    tracker_key: str | None = None
+    session_priority: int = 1
+    session_timeout_sec: int = 10
+
+    def _sanity_check(self) -> None:
+        err_str = (
+            "RPCConfig.{0} is not provided. Please provide it explicitly,"
+            "or set environment variable {1}"
+        )
+        if self.tracker_host is None:
+            raise ValueError(err_str.format("tracker_host", "TVM_TRACKER_HOST"))
+        if self.tracker_port is None:
+            raise ValueError(err_str.format("tracker_port", "TVM_TRACKER_PORT"))
+        if self.tracker_key is None:
+            raise ValueError(err_str.format("tracker_key", "TVM_TRACKER_KEY"))
+
+    @staticmethod
+    def _normalized(config: Optional["RPCConfig"]) -> "RPCConfig":
+        if config is None:
+            config = RPCConfig()
+        tracker_host = config.tracker_host or os.environ.get("TVM_TRACKER_HOST", None)
+        tracker_port = config.tracker_port or os.environ.get("TVM_TRACKER_PORT", None)
+        tracker_key = config.tracker_key or os.environ.get("TVM_TRACKER_KEY", None)
+        if isinstance(tracker_port, str):
+            tracker_port = int(tracker_port)
+        config = RPCConfig(
+            tracker_host=tracker_host,
+            tracker_port=tracker_port,
+            tracker_key=tracker_key,
+            session_priority=config.session_priority,
+            session_timeout_sec=config.session_timeout_sec,
+        )
+        config._sanity_check()  # pylint: disable=protected-access
+        return config
+
+    def connect_tracker(self) -> rpc.TrackerSession:
+        """Connect to the tracker
+
+        Returns
+        -------
+        tracker : TrackerSession
+            The connected tracker session
+        """
+        tracker: rpc.TrackerSession | None = None
+
+        def _connect():
+            nonlocal tracker
+            tracker = rpc.connect_tracker(self.tracker_host, self.tracker_port)
+
+        t = Thread(target=_connect)
+        t.start()
+        t.join(self.session_timeout_sec)
+        if t.is_alive() or tracker is None:
+            raise ValueError(
+                "Unable to connect to the tracker using the following configuration:\n"
+                f"    tracker host: {self.tracker_host}\n"
+                f"    tracker port: {self.tracker_port}\n"
+                f"    timeout (sec): {self.session_timeout_sec}\n"
+                "Please check the tracker status via the following command:\n"
+                "     python3 -m tvm.exec.query_rpc_tracker "
+                f"--host {self.tracker_host} --port {self.tracker_port}"
+            )
+        return tracker
+
+    def connect_server(self) -> rpc.RPCSession:
+        """Connect to the server
+
+        Returns
+        -------
+        session : RPCSession
+            The connected rpc session
+        """
+        tracker = self.connect_tracker()
+        session: rpc.RPCSession = tracker.request(
+            key=self.tracker_key,
+            priority=self.session_priority,
+            session_timeout=self.session_timeout_sec,
+        )
+        return session
+
+    def count_num_servers(self, allow_missing=True) -> int:
+        """Count the number of servers available in the tracker
+
+        Parameters
+        ----------
+        allow_missing : bool
+            Whether to allow no server to be found.
+
+        Returns
+        -------
+        num_servers : int
+            The number of servers
+        """
+        tracker = self.connect_tracker()
+        tracker_summary = tracker.summary()
+        result: int = 0
+        for item in tracker_summary["server_info"]:
+            _, item_key = item["key"].split(":")
+            if item_key == self.tracker_key:
+                result += 1
+        if result == 0 and not allow_missing:
+            raise ValueError(
+                "Unable to find servers with the specific key using the following configuration:\n"
+                f"    tracker host: {self.tracker_host}\n"
+                f"    tracker port: {self.tracker_port}\n"
+                f"    tracker key: {self.tracker_key}\n"
+                f"    timeout (sec): {self.session_timeout_sec}\n"
+                "Please check the tracker status via the following command:\n"
+                "     python3 -m tvm.exec.query_rpc_tracker "
+                f"--host {self.tracker_host} --port {self.tracker_port}\n"
+                f'and look for key: "{self.tracker_key}"'
+            )
+        return result

python/tvm/s_tir/meta_schedule/runner/local_runner.py

@@ -0,0 +1,404 @@
+# 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.
+"""Local Runner"""
+
+import logging
+import subprocess
+from collections.abc import Callable
+from contextlib import contextmanager
+
+import tvm
+from tvm.ir.utils import derived_object
+from tvm.support.popen_pool import PopenPoolExecutor
+
+from ....runtime import Device, Module
+from ..logging import get_logger
+from ..profiler import Profiler
+from ..utils import get_global_func_with_default_on_worker
+from .config import EvaluatorConfig
+from .runner import PyRunner, PyRunnerFuture, RunnerFuture, RunnerInput, RunnerResult
+from .utils import (
+    T_ARG_INFO_JSON_OBJ_LIST,
+    T_ARGUMENT_LIST,
+    alloc_argument_common,
+    run_evaluator_common,
+)
+
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
+
+T_ALLOC_ARGUMENT = Callable[  # pylint: disable=invalid-name
+    [
+        Device,  # The device on the remote
+        T_ARG_INFO_JSON_OBJ_LIST,  # The metadata information of the arguments to be allocated
+        int,  # The number of repeated allocations to be done
+    ],
+    list[T_ARGUMENT_LIST],  # A list of argument lists
+]
+T_RUN_EVALUATOR = Callable[  # pylint: disable=invalid-name
+    [
+        Module,  # The Module opened on the remote
+        Device,  # The device on the remote
+        EvaluatorConfig,  # The evaluator configuration
+        list[T_ARGUMENT_LIST],  # A list of argument lists
+    ],
+    list[float],  # A list of running time
+]
+T_CLEANUP = Callable[  # pylint: disable=invalid-name
+    [],
+    None,
+]
+
+
+@derived_object
+class LocalRunnerFuture(PyRunnerFuture):
+    """Local based runner future
+
+    Parameters
+    ----------
+    res: Optional[List[float]]
+        The optional result as a list of float.
+    error_message: Optional[str]
+        The optional error message.
+
+    Note
+    ----
+    Only one of the parameters should be None upon the creation
+    of LocalRunnerFuture object
+    """
+
+    res: list[float] | None
+    error_message: str | None
+
+    def __init__(self, res: list[float] | None = None, error_message: str | None = None) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        res: Optional[List[float]]
+            The result of this LocalRunnerFuture
+        error_message: Optional[str]
+            The stringfied error message of any exception during execution
+
+        """
+        super().__init__()
+        self.res = res
+        self.error_message = error_message
+
+        # sanity check upon the creation of LocalRunnerFuture object
+        if (res is None and error_message is None) or (
+            res is not None and error_message is not None
+        ):
+            raise AttributeError(
+                "Only one of the two parameters should be None upon the creation"
+                "of LocalRunnerFuture object."
+            )
+
+    def done(self) -> bool:
+        return True
+
+    def result(self) -> RunnerResult:
+        return RunnerResult(self.res, self.error_message)
+
+
+def _worker_func(
+    _f_alloc_argument: str | None,
+    _f_run_evaluator: str | None,
+    _f_cleanup: str | None,
+    evaluator_config: EvaluatorConfig,
+    alloc_repeat: int,
+    artifact_path: str,
+    device_type: str,
+    args_info: T_ARG_INFO_JSON_OBJ_LIST,
+) -> list[float]:
+    f_alloc_argument: T_ALLOC_ARGUMENT = get_global_func_with_default_on_worker(
+        _f_alloc_argument, default_alloc_argument
+    )
+    f_run_evaluator: T_RUN_EVALUATOR = get_global_func_with_default_on_worker(
+        _f_run_evaluator, default_run_evaluator
+    )
+    f_cleanup: T_CLEANUP = get_global_func_with_default_on_worker(_f_cleanup, default_cleanup)
+
+    @contextmanager
+    def resource_handler():
+        try:
+            yield
+        finally:
+            # Final step. Always clean up
+            with Profiler.timeit("LocalRunner/cleanup"):
+                f_cleanup()
+
+    with resource_handler():
+        # Step 1: create the local runtime module
+        with Profiler.timeit("LocalRunner/load_module"):
+            rt_mod = tvm.runtime.load_module(artifact_path)
+        # Step 2: Allocate input arguments
+        with Profiler.timeit("LocalRunner/alloc_argument"):
+            device = tvm.runtime.device(device_type, 0)
+            repeated_args: list[T_ARGUMENT_LIST] = f_alloc_argument(
+                device,
+                args_info,
+                alloc_repeat,
+            )
+        # Step 3: Run time_evaluator
+        with Profiler.timeit("LocalRunner/run_evaluator"):
+            costs: list[float] = f_run_evaluator(
+                rt_mod,
+                device,
+                evaluator_config,
+                repeated_args,
+            )
+    return costs
+
+
+@derived_object
+class LocalRunner(PyRunner):
+    """Local runner
+
+    Parameters
+    ----------
+    evaluator_config: EvaluatorConfig
+        The evaluator configuration.
+    cooldown_sec: float
+        The cooldown in seconds.
+    alloc_repeat: int
+        The number of times to repeat the allocation.
+    f_alloc_argument: Optional[str, Callable]
+        The function name to allocate the arguments or the function itself.
+    f_run_evaluator: Optional[str, Callable]
+        The function name to run the evaluator or the function itself.
+    f_cleanup: Optional[str, Callable]
+        The function name to cleanup the session or the function itself.
+    pool: PopenPoolExecutor
+        The popen pool executor.
+
+    Attributes
+    ----------
+    T_ALLOC_ARGUMENT : typing._GenericAlias
+        The signature of the function `f_alloc_argument`, which is:
+
+        .. code-block:: python
+
+        def default_alloc_argument(
+            device: Device,
+            args_info: T_ARG_INFO_JSON_OBJ_LIST,
+            alloc_repeat: int,
+        ) -> List[T_ARGUMENT_LIST]:
+            ...
+
+    T_RUN_EVALUATOR : typing._GenericAlias
+        The signature of the function `f_run_evaluator`, which is:
+
+        .. code-block:: python
+
+        def default_run_evaluator(
+            rt_mod: Module,
+            device: Device,
+            evaluator_config: EvaluatorConfig,
+            repeated_args: List[T_ARGUMENT_LIST],
+        ) -> List[float]:
+            ...
+
+    T_CLEANUP : typing._GenericAlias
+        The signature of the function `f_cleanup`, which is:
+
+        .. code-block:: python
+
+        def default_cleanup() -> None:
+            ...
+    """
+
+    timeout_sec: float
+    evaluator_config: EvaluatorConfig
+    cooldown_sec: float
+    alloc_repeat: int
+
+    f_alloc_argument: T_ALLOC_ARGUMENT | str | None
+    f_run_evaluator: T_RUN_EVALUATOR | str | None
+    f_cleanup: T_CLEANUP | str | None
+
+    pool: PopenPoolExecutor
+
+    def __init__(
+        self,
+        timeout_sec: float = 30,
+        evaluator_config: EvaluatorConfig | None = None,
+        cooldown_sec: float = 0.0,
+        alloc_repeat: int = 1,
+        f_alloc_argument: T_ALLOC_ARGUMENT | str | None = None,
+        f_run_evaluator: T_RUN_EVALUATOR | str | None = None,
+        f_cleanup: T_CLEANUP | str | None = None,
+        initializer: Callable[[], None] | None = None,
+    ) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        timeout_sec: float
+            The timeout setting.
+        evaluator_config: EvaluatorConfig
+            The evaluator configuration.
+        cooldown_sec: float
+            The cooldown in seconds.
+        alloc_repeat: int
+            The number of times to random fill the allocation.
+        f_alloc_argument: Union[T_ALLOC_ARGUMENT, str, None]
+            The function name to allocate the arguments or the function itself.
+        f_run_evaluator: Union[T_RUN_EVALUATOR, str, None]
+            The function name to run the evaluator or the function itself.
+        f_cleanup: Union[T_CLEANUP, str, None]
+            The function name to cleanup the session or the function itself.
+        initializer: Optional[Callable[[], None]]
+            The initializer function.
+        """
+        super().__init__()
+        self.timeout_sec = timeout_sec
+        self.evaluator_config = EvaluatorConfig._normalized(evaluator_config)
+        self.cooldown_sec = cooldown_sec
+        self.alloc_repeat = alloc_repeat
+        self.f_alloc_argument = f_alloc_argument
+        self.f_run_evaluator = f_run_evaluator
+        self.f_cleanup = f_cleanup
+
+        err_path = subprocess.DEVNULL
+        if logger.root.level <= logging.DEBUG:
+            err_path = subprocess.STDOUT
+
+        logger.info("LocalRunner: max_workers = 1")
+        self.pool = PopenPoolExecutor(
+            max_workers=1,  # one local worker
+            timeout=timeout_sec,
+            initializer=initializer,
+            stderr=err_path,  # suppress the stderr output
+        )
+        self._sanity_check()
+
+    def run(self, runner_inputs: list[RunnerInput]) -> list[RunnerFuture]:
+        results: list[RunnerFuture] = []
+        for runner_input in runner_inputs:
+            future = self.pool.submit(
+                _worker_func,
+                self.f_alloc_argument,
+                self.f_run_evaluator,
+                self.f_cleanup,
+                self.evaluator_config,
+                self.alloc_repeat,
+                str(runner_input.artifact_path),
+                str(runner_input.device_type),
+                tuple(arg_info.as_json() for arg_info in runner_input.args_info),
+            )
+            try:
+                result: list[float] = future.result()
+                error_message: str = None
+            except TimeoutError:
+                result = None
+                error_message = f"LocalRunner: Timeout, killed after {self.timeout_sec} seconds\n"
+            except Exception as exception:  # pylint: disable=broad-except
+                result = None
+                error_message = "LocalRunner: An exception occurred\n" + str(exception)
+            local_future = LocalRunnerFuture(res=result, error_message=error_message)
+            results.append(local_future)  # type: ignore
+        return results
+
+    def _sanity_check(self) -> None:
+        def _check(
+            f_alloc_argument,
+            f_run_evaluator,
+            f_cleanup,
+        ) -> None:
+            get_global_func_with_default_on_worker(name=f_alloc_argument, default=None)
+            get_global_func_with_default_on_worker(name=f_run_evaluator, default=None)
+            get_global_func_with_default_on_worker(name=f_cleanup, default=None)
+
+        value = self.pool.submit(
+            _check,
+            self.f_alloc_argument,
+            self.f_run_evaluator,
+            self.f_cleanup,
+        )
+        value.result()
+
+
+def default_alloc_argument(
+    device: Device,
+    args_info: T_ARG_INFO_JSON_OBJ_LIST,
+    alloc_repeat: int,
+) -> list[T_ARGUMENT_LIST]:
+    """Default function to allocate the arguments
+
+    Parameters
+    ----------
+    device: Device
+        The device to allocate the arguments
+    args_info: T_ARG_INFO_JSON_OBJ_LIST
+        The arguments info
+    alloc_repeat: int
+        The number of times to repeat the allocation
+
+    Returns
+    -------
+    repeated_args: List[T_ARGUMENT_LIST]
+        The allocation args
+    """
+    f_random_fill = get_global_func_with_default_on_worker(
+        name="tvm.contrib.random.random_fill_for_measure", default=None
+    )
+    return alloc_argument_common(f_random_fill, device, args_info, alloc_repeat)
+
+
+def default_run_evaluator(
+    rt_mod: Module,
+    device: Device,
+    evaluator_config: EvaluatorConfig,
+    repeated_args: list[T_ARGUMENT_LIST],
+) -> list[float]:
+    """Default function to run the evaluator
+
+    Parameters
+    ----------
+    rt_mod: Module
+        The runtime module
+    device: Device
+        The device to run the evaluator
+    evaluator_config: EvaluatorConfig
+        The evaluator config
+    repeated_args: List[T_ARGUMENT_LIST]
+        The repeated arguments
+
+    Returns
+    -------
+    costs: List[float]
+        The evaluator results
+    """
+    return run_evaluator_common(rt_mod, device, evaluator_config, repeated_args)
+
+
+def default_cleanup() -> None:
+    """Default function to clean up the session"""
+    pass  # pylint: disable=unnecessary-pass
+
+
+@tvm.register_global_func("s_tir.meta_schedule.runner.get_local_runner")
+def get_local_builder() -> LocalRunner:
+    """Get the local Runner.
+
+    Returns
+    -------
+    runner: LocalRunner
+        The local runner
+    """
+    return LocalRunner()

python/tvm/s_tir/meta_schedule/runner/rpc_runner.py

@@ -0,0 +1,535 @@
+# 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.
+"""RPC Runner"""
+
+import concurrent.futures
+import os.path as osp
+from collections.abc import Callable
+from contextlib import contextmanager
+
+from tvm.ir.utils import derived_object
+from tvm.rpc import RPCSession
+from tvm.runtime import Device, Module
+from tvm.support.popen_pool import PopenPoolExecutor
+
+from ..logging import get_logger
+from ..profiler import Profiler
+from ..utils import (
+    get_global_func_on_rpc_session,
+    get_global_func_with_default_on_worker,
+)
+from .config import EvaluatorConfig, RPCConfig
+from .runner import PyRunner, PyRunnerFuture, RunnerFuture, RunnerInput, RunnerResult
+from .utils import (
+    T_ARG_INFO_JSON_OBJ_LIST,
+    T_ARGUMENT_LIST,
+    alloc_argument_common,
+    run_evaluator_common,
+)
+
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
+
+T_CREATE_SESSION = Callable[  # pylint: disable=invalid-name
+    [RPCConfig],  # The RPC configuration
+    RPCSession,  # The RPC Session
+]
+T_UPLOAD_MODULE = Callable[  # pylint: disable=invalid-name
+    [
+        RPCSession,  # The RPC Session
+        str,  # local path to the artifact
+        str,  # remote path to the artifact
+    ],
+    Module,  # the Module opened on the remote
+]
+T_ALLOC_ARGUMENT = Callable[  # pylint: disable=invalid-name
+    [
+        RPCSession,  # The RPC Session
+        Device,  # The device on the remote
+        T_ARG_INFO_JSON_OBJ_LIST,  # The metadata information of the arguments to be allocated
+        int,  # The number of repeated allocations to be done
+    ],
+    list[T_ARGUMENT_LIST],  # A list of argument lists
+]
+T_RUN_EVALUATOR = Callable[  # pylint: disable=invalid-name
+    [
+        RPCSession,  # The RPC Session
+        Module,  # The Module opened on the remote
+        Device,  # The device on the remote
+        EvaluatorConfig,  # The evaluator configuration
+        list[T_ARGUMENT_LIST],  # A list of argument lists
+    ],
+    list[float],  # A list of running time
+]
+T_CLEANUP = Callable[  # pylint: disable=invalid-name
+    [
+        RPCSession | None,  # The RPC Session to be cleaned up
+        str | None,  # remote path to the artifact
+    ],
+    None,
+]
+
+
+@derived_object
+class RPCRunnerFuture(PyRunnerFuture):
+    """RPC based runner future
+
+    Parameters
+    ----------
+    future: concurrent.futures.Future
+        The concurrent function to check when the function is done and to return the result.
+    timeout_sec: float
+        The timeout in seconds.
+    """
+
+    future: concurrent.futures.Future
+    timeout_sec: float
+
+    def __init__(self, future: concurrent.futures.Future, timeout_sec: float) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        future: concurrent.futures.Future
+            The concurrent function to check when the function is done and to return the result.
+        timeout_sec: float
+            The timeout in seconds.
+        """
+        super().__init__()
+        self.future = future
+        self.timeout_sec = timeout_sec
+
+    def done(self) -> bool:
+        return self.future.done()
+
+    def result(self) -> RunnerResult:
+        try:
+            run_secs: list[float] = self.future.result()
+        except TimeoutError:
+            return RunnerResult(
+                None,
+                error_msg=f"RPCRunner: Timeout, killed after {self.timeout_sec} seconds",
+            )
+        except Exception as exception:  # pylint: disable=broad-except
+            return RunnerResult(
+                None,
+                error_msg="RPCRunner: An exception occurred\n" + str(exception),
+            )
+        return RunnerResult(run_secs, None)
+
+
+@derived_object
+class RPCRunner(PyRunner):
+    """RPC based runner
+
+    Parameters
+    ----------
+    rpc_config: RPCConfig
+        The rpc configuration.
+    evaluator_config: EvaluatorConfig
+        The evaluator configuration.
+    cooldown_sec: float
+        The cooldown in seconds. TODO(@junrushao1994,@zxybazh): This is not used yet.
+    alloc_repeat: int
+        The number of times to repeat the allocation.
+    f_create_session: Optional[str, Callable]
+        The function name to create the session or the function itself.
+    f_upload_module: Optional[str, Callable]
+        The function name to upload the module or the function itself.
+    f_alloc_argument: Optional[str, Callable]
+        The function name to allocate the arguments or the function itself.
+    f_run_evaluator: Optional[str, Callable]
+        The function name to run the evaluator or the function itself.
+    f_cleanup: Optional[str, Callable]
+        The function name to cleanup the session or the function itself.
+    pool: PopenPoolExecutor
+        The popen pool executor.
+
+    Attributes
+    ----------
+    T_CREATE_SESSION : typing._GenericAlias
+        The signature of the function `f_create_session`, which is:
+
+        .. code-block:: python
+
+        def default_create_session(rpc_config: RPCConfig) -> RPCSession:
+            ...
+
+    T_UPLOAD_MODULE : typing._GenericAlias
+        The signature of the function `f_upload_module`, which is:
+
+        .. code-block:: python
+
+        def default_upload_module(
+            session: RPCSession,
+            local_path: str,
+            remote_path: str,
+        ) -> Module:
+            ...
+
+    T_ALLOC_ARGUMENT : typing._GenericAlias
+        The signature of the function `f_alloc_argument`, which is:
+
+        .. code-block:: python
+
+        def default_alloc_argument(
+            session: RPCSession,
+            device: Device,
+            args_info: T_ARG_INFO_JSON_OBJ_LIST,
+            alloc_repeat: int,
+        ) -> List[T_ARGUMENT_LIST]:
+            ...
+
+    T_RUN_EVALUATOR : typing._GenericAlias
+        The signature of the function `f_run_evaluator`, which is:
+
+        .. code-block:: python
+
+        def default_run_evaluator(
+            session: RPCSession,
+            rt_mod: Module,
+            device: Device,
+            evaluator_config: EvaluatorConfig,
+            repeated_args: List[T_ARGUMENT_LIST],
+        ) -> List[float]:
+            ...
+
+    T_CLEANUP : typing._GenericAlias
+        The signature of the function `f_cleanup`, which is:
+
+        .. code-block:: python
+
+        def default_cleanup(
+            session: Optional[RPCSession],
+            remote_path: Optional[str],
+        ) -> None:
+            ...
+    """
+
+    rpc_config: RPCConfig
+    evaluator_config: EvaluatorConfig
+    cooldown_sec: float
+    alloc_repeat: int
+
+    f_create_session: T_CREATE_SESSION | str | None
+    f_upload_module: T_UPLOAD_MODULE | str | None
+    f_alloc_argument: T_ALLOC_ARGUMENT | str | None
+    f_run_evaluator: T_RUN_EVALUATOR | str | None
+    f_cleanup: T_CLEANUP | str | None
+
+    pool: PopenPoolExecutor
+
+    def __init__(
+        self,
+        rpc_config: RPCConfig | None = None,
+        evaluator_config: EvaluatorConfig | None = None,
+        cooldown_sec: float = 0.0,
+        alloc_repeat: int = 1,
+        f_create_session: T_CREATE_SESSION | str | None = None,
+        f_upload_module: T_UPLOAD_MODULE | str | None = None,
+        f_alloc_argument: T_ALLOC_ARGUMENT | str | None = None,
+        f_run_evaluator: T_RUN_EVALUATOR | str | None = None,
+        f_cleanup: T_CLEANUP | str | None = None,
+        max_workers: int | None = None,
+        initializer: Callable[[], None] | None = None,
+    ) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        rpc_config: RPCConfig
+            The rpc configuration.
+        evaluator_config: EvaluatorConfig
+            The evaluator configuration.
+        cooldown_sec: float
+            The cooldown in seconds.
+        alloc_repeat: int
+            The number of times to random fill the allocation.
+        f_create_session: Union[T_CREATE_SESSION, str, None]
+            The function name to create the session or the function itself.
+        f_upload_module: Union[T_UPLOAD_MODULE, str, None]
+            The function name to upload the module or the function itself.
+        f_alloc_argument: Union[T_ALLOC_ARGUMENT, str, None]
+            The function name to allocate the arguments or the function itself.
+        f_run_evaluator: Union[T_RUN_EVALUATOR, str, None]
+            The function name to run the evaluator or the function itself.
+        f_cleanup: Union[T_CLEANUP, str, None]
+            The function name to cleanup the session or the function itself.
+        max_workers: Optional[int] = None
+            The maximum number of connections. Defaults to 1.
+        initializer: Optional[Callable[[], None]]
+            The initializer function.
+        """
+        super().__init__()
+        self.rpc_config = RPCConfig._normalized(rpc_config)
+        self.evaluator_config = EvaluatorConfig._normalized(evaluator_config)
+        self.cooldown_sec = cooldown_sec
+        self.alloc_repeat = alloc_repeat
+        self.f_create_session = f_create_session
+        self.f_upload_module = f_upload_module
+        self.f_alloc_argument = f_alloc_argument
+        self.f_run_evaluator = f_run_evaluator
+        self.f_cleanup = f_cleanup
+        if max_workers is None:
+            max_workers = 1
+        logger.info("RPCRunner: max_workers = %d", max_workers)
+        self.pool = PopenPoolExecutor(
+            max_workers=max_workers,
+            initializer=initializer,
+        )
+        self._sanity_check()
+
+    def run(self, runner_inputs: list[RunnerInput]) -> list[RunnerFuture]:
+        results: list[RunnerFuture] = []
+        for runner_input in runner_inputs:
+            future = RPCRunnerFuture(
+                future=self.pool.submit(
+                    _worker_func,
+                    self.f_create_session,
+                    self.f_upload_module,
+                    self.f_alloc_argument,
+                    self.f_run_evaluator,
+                    self.f_cleanup,
+                    self.rpc_config,
+                    self.evaluator_config,
+                    self.alloc_repeat,
+                    str(runner_input.artifact_path),
+                    str(runner_input.device_type),
+                    tuple(arg_info.as_json() for arg_info in runner_input.args_info),
+                ),
+                timeout_sec=self.rpc_config.session_timeout_sec,
+            )
+            results.append(future)  # type: ignore
+        return results
+
+    def _sanity_check(self) -> None:
+        def _check(
+            f_create_session,
+            f_upload_module,
+            f_alloc_argument,
+            f_run_evaluator,
+            f_cleanup,
+        ) -> None:
+            get_global_func_with_default_on_worker(name=f_create_session, default=None)
+            get_global_func_with_default_on_worker(name=f_upload_module, default=None)
+            get_global_func_with_default_on_worker(name=f_alloc_argument, default=None)
+            get_global_func_with_default_on_worker(name=f_run_evaluator, default=None)
+            get_global_func_with_default_on_worker(name=f_cleanup, default=None)
+
+        value = self.pool.submit(
+            _check,
+            self.f_create_session,
+            self.f_upload_module,
+            self.f_alloc_argument,
+            self.f_run_evaluator,
+            self.f_cleanup,
+        )
+        value.result()
+
+
+def _worker_func(
+    _f_create_session: T_CREATE_SESSION | str | None,
+    _f_upload_module: T_UPLOAD_MODULE | str | None,
+    _f_alloc_argument: T_ALLOC_ARGUMENT | str | None,
+    _f_run_evaluator: T_RUN_EVALUATOR | str | None,
+    _f_cleanup: T_CLEANUP | str | None,
+    rpc_config: RPCConfig,
+    evaluator_config: EvaluatorConfig,
+    alloc_repeat: int,
+    artifact_path: str,
+    device_type: str,
+    args_info: T_ARG_INFO_JSON_OBJ_LIST,
+) -> list[float]:
+    # Step 0. Get the registered functions
+    f_create_session: T_CREATE_SESSION = get_global_func_with_default_on_worker(
+        _f_create_session, default_create_session
+    )
+    f_upload_module: T_UPLOAD_MODULE = get_global_func_with_default_on_worker(
+        _f_upload_module, default_upload_module
+    )
+    f_alloc_argument: T_ALLOC_ARGUMENT = get_global_func_with_default_on_worker(
+        _f_alloc_argument, default_alloc_argument
+    )
+    f_run_evaluator: T_RUN_EVALUATOR = get_global_func_with_default_on_worker(
+        _f_run_evaluator, default_run_evaluator
+    )
+    f_cleanup: T_CLEANUP = get_global_func_with_default_on_worker(_f_cleanup, default_cleanup)
+    # Managed resources
+    session: RPCSession | None = None
+    remote_path: str | None = None
+
+    @contextmanager
+    def resource_handler():
+        try:
+            yield
+        finally:
+            # Final step. Always clean up
+            with Profiler.timeit("RPCRunner/cleanup"):
+                f_cleanup(session, remote_path)
+
+    with resource_handler():
+        # Step 1. Create session
+        with Profiler.timeit("RPCRunner/create_session"):
+            session = f_create_session(rpc_config)
+            device = session.device(device_type, 0)
+        # Step 2. Upload the module
+        with Profiler.timeit("RPCRunner/upload_module"):
+            _, remote_path = osp.split(artifact_path)
+            local_path: str = artifact_path
+            rt_mod: Module = f_upload_module(session, local_path, remote_path)
+        # Step 3: Allocate input arguments
+        with Profiler.timeit("RPCRunner/alloc_argument"):
+            repeated_args: list[T_ARGUMENT_LIST] = f_alloc_argument(
+                session,
+                device,
+                args_info,
+                alloc_repeat,
+            )
+        # Step 4: Run time_evaluator
+        with Profiler.timeit("LocalRunner/run_evaluator"):
+            costs: list[float] = f_run_evaluator(
+                session,
+                rt_mod,
+                device,
+                evaluator_config,
+                repeated_args,
+            )
+    return costs
+
+
+def default_create_session(rpc_config: RPCConfig) -> RPCSession:
+    """Default function to create the session
+
+    Parameters
+    ----------
+    rpc_config : RPCConfig
+        The configuration of the RPC session
+
+    Returns
+    -------
+    session : RPCSession
+        The created rpc session
+    """
+    return rpc_config.connect_server()
+
+
+def default_upload_module(
+    session: RPCSession,
+    local_path: str,
+    remote_path: str,
+) -> Module:
+    """Default function to upload the module
+
+    Parameters
+    ----------
+    session: RPCSession
+        The session to upload the module
+    local_path: str
+        The local path of the module
+    remote_path: str
+        The remote path to place the module
+
+    Returns
+    -------
+    rt_mod : Module
+        The runtime module
+    """
+    session.upload(local_path, remote_path)
+    rt_mod: Module = session.load_module(remote_path)
+    return rt_mod
+
+
+def default_alloc_argument(
+    session: RPCSession,
+    device: Device,
+    args_info: T_ARG_INFO_JSON_OBJ_LIST,
+    alloc_repeat: int,
+) -> list[T_ARGUMENT_LIST]:
+    """Default function to allocate the arguments
+
+    Parameters
+    ----------
+    session: RPCSession
+        The session to allocate the arguments
+    device: Device
+        The device to allocate the arguments
+    args_info: T_ARG_INFO_JSON_OBJ_LIST
+        The arguments info
+    alloc_repeat: int
+        The number of times to repeat the allocation
+
+    Returns
+    -------
+    repeated_args: List[Args]
+        The allocation args
+    """
+    f_random_fill = get_global_func_on_rpc_session(
+        session,
+        "tvm.contrib.random.random_fill_for_measure",
+        "Please make sure 'USE_RANDOM' is turned ON in the config.cmake on the RPC server.",
+    )
+
+    return alloc_argument_common(f_random_fill, device, args_info, alloc_repeat)
+
+
+def default_run_evaluator(
+    session: RPCSession,  # pylint: disable=unused-argument
+    rt_mod: Module,
+    device: Device,
+    evaluator_config: EvaluatorConfig,
+    repeated_args: list[T_ARGUMENT_LIST],
+) -> list[float]:
+    """Default function to run the evaluator
+
+    Parameters
+    ----------
+    session: RPCSession
+        The session to run the evaluator
+    rt_mod: Module
+        The runtime module
+    device: Device
+        The device to run the evaluator
+    evaluator_config: EvaluatorConfig
+        The evaluator config
+    repeated_args: List[T_ARGUMENT_LIST]
+        The repeated arguments
+
+    Returns
+    -------
+    costs: List[float]
+        The evaluator results
+    """
+    return run_evaluator_common(rt_mod, device, evaluator_config, repeated_args)
+
+
+def default_cleanup(
+    session: RPCSession | None,
+    remote_path: str | None,
+) -> None:
+    """Default function to clean up the session
+
+    Parameters
+    ----------
+    session: RPCSession
+        The session to clean up
+    remote_path: str
+        The remote path to clean up
+    """
+    if session is not None and remote_path is not None:
+        session.remove(remote_path)
+        session.remove(remote_path + ".so")
+        session.remove("")

python/tvm/s_tir/meta_schedule/runner/runner.py

@@ -0,0 +1,257 @@
+# 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.
+# ruff: noqa: RUF012
+"""Runners"""
+
+from collections.abc import Callable
+from typing import Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+
+from .. import _ffi_api
+from ..arg_info import ArgInfo
+
+
+@register_object("s_tir.meta_schedule.RunnerInput")
+class RunnerInput(Object):
+    """The runner's input
+
+    Parameters
+    ----------
+    artifact_path : str
+        The path to the built artifact.
+    device_type : str
+        The device type.
+    args_info : List[ArgInfo]
+        The argument information.
+    """
+
+    artifact_path: str
+    device_type: str
+    args_info: list[ArgInfo]
+
+    def __init__(
+        self,
+        artifact_path: str,
+        device_type: str,
+        args_info: list[ArgInfo],
+    ) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        artifact_path : str
+            The path to the built artifact.
+        device_type : str
+            The device type.
+        args_info : List[ArgInfo]
+            The argument information.
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.RunnerInput,  # type: ignore # pylint: disable=no-member
+            artifact_path,
+            device_type,
+            args_info,
+        )
+
+
+@register_object("s_tir.meta_schedule.RunnerResult")
+class RunnerResult(Object):
+    """The runner's result
+
+    Parameters
+    ----------
+    run_secs : Optional[List[float]]
+        The run time in seconds.
+    error_msg : Optional[str]
+        The error message, if any.
+    """
+
+    run_secs: list[float] | None
+    error_msg: str | None
+
+    def __init__(
+        self,
+        run_secs: list[float] | None,
+        error_msg: str | None,
+    ) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        run_secs : Optional[List[float]]
+            The run time in seconds.
+        error_msg : Optional[str]
+            The error message, if any.
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.RunnerResult,  # type: ignore # pylint: disable=no-member
+            run_secs,
+            error_msg,
+        )
+
+
+@register_object("s_tir.meta_schedule.RunnerFuture")
+class RunnerFuture(Object):
+    """
+    A class to fetch asynchronous runner's output.
+    This is NOT the user facing class for function overloading inheritance.
+    Can be used for general return type of runner.
+
+    See also: PyRunnerFuture
+    """
+
+    def __init__(self, f_done: Callable, f_result: Callable | None = None) -> None:
+        """Constructor"""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.RunnerFuture,  # type: ignore # pylint: disable=no-member
+            f_done,
+            f_result,
+        )
+
+    def done(self) -> bool:
+        """Check whether the runner has finished."""
+        return _ffi_api.RunnerFutureDone(self)  # type: ignore # pylint: disable=no-member
+
+    def result(self) -> RunnerResult:
+        """Fetch the runner's output if it is ready."""
+        return _ffi_api.RunnerFutureResult(self)  # type: ignore # pylint: disable=no-member
+
+
+class PyRunnerFuture:
+    """
+    A class to fetch asynchronous runner's output with customizable function on the python side.
+    This is the user facing class for function overloading inheritance.
+    Can NOT be used for general return type of runner.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    Example::
+
+        @derived_object
+        def LocalRunnerFuture(PyRunnerFuture):
+            ...
+    """
+
+    _tvm_metadata = {
+        "cls": RunnerFuture,
+        "methods": ["done", "result"],
+    }
+
+    def done(self) -> bool:
+        """Check whether the runner has finished."""
+        raise NotImplementedError
+
+    def result(self) -> RunnerResult:
+        """Fetch the runner's output if it is ready."""
+        raise NotImplementedError
+
+
+@register_object("s_tir.meta_schedule.Runner")
+class Runner(Object):
+    """The abstract runner interface"""
+
+    RunnerType = Union["Runner", Literal["local", "rpc"]]
+
+    def run(self, runner_inputs: list[RunnerInput]) -> list[RunnerFuture]:
+        """Run the built artifact and get runner futures.
+
+        Parameters
+        ----------
+        runner_inputs : List[RunnerInput]
+            The inputs to the runner.
+
+        Returns
+        -------
+        runner_futures: List[RunnerFuture]
+            The runner futures.
+        """
+        return _ffi_api.RunnerRun(self, runner_inputs)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def create(  # pylint: disable=keyword-arg-before-vararg
+        kind: Literal["local", "rpc"] = "local",
+        *args,
+        **kwargs,
+    ) -> "Runner":
+        """Create a Runner."""
+        from . import LocalRunner, RPCRunner  # pylint: disable=import-outside-toplevel
+
+        if kind == "local":
+            if "max_workers" in kwargs:
+                kwargs.pop("max_workers")
+            return LocalRunner(*args, **kwargs)  # type: ignore
+        elif kind == "rpc":
+            return RPCRunner(*args, **kwargs)  # type: ignore
+        raise ValueError(f"Unknown Runner: {kind}")
+
+
+create = Runner.create  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.PyRunner")
+class _PyRunner(Runner):
+    """
+    A TVM object runner to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyRunner
+    """
+
+    def __init__(self, f_run: Callable | None = None) -> None:
+        """Constructor"""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.RunnerPyRunner,  # type: ignore # pylint: disable=no-member
+            f_run,
+        )
+
+
+class PyRunner:
+    """
+    An abstract runner with customized run method on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PyRunner,
+        "methods": ["run"],
+    }
+
+    def run(self, runner_inputs: list[RunnerInput]) -> list[RunnerFuture]:
+        """Run the built artifact and get runner futures.
+
+        Parameters
+        ----------
+        runner_inputs : List[RunnerInput]
+            The inputs to the runner.
+
+        Returns
+        -------
+        runner_futures: List[RunnerFuture]
+            The runner futures.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/runner/utils.py

@@ -0,0 +1,124 @@
+# 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.
+"""Runner utility functions"""
+
+import itertools
+from collections.abc import Callable
+from typing import Any
+
+import tvm.runtime
+
+from ....runtime import Device, Module
+from .config import EvaluatorConfig
+
+T_ARG_INFO_JSON_OBJ = list[Any]  # pylint: disable=invalid-name
+T_ARG_INFO_JSON_OBJ_LIST = list[T_ARG_INFO_JSON_OBJ]  # pylint: disable=invalid-name
+T_ARGUMENT = Any  # pylint: disable=invalid-name
+T_ARGUMENT_LIST = list[T_ARGUMENT]  # pylint: disable=invalid-name
+
+
+def alloc_argument_common(
+    f_random_fill: Callable,
+    device: Device,
+    args_info: T_ARG_INFO_JSON_OBJ_LIST,
+    alloc_repeat: int,
+) -> list[T_ARGUMENT_LIST]:
+    """Common function to allocate the arguments
+
+    Parameters
+    ----------
+    f_random_fill: Callable
+        The callable function for random fill
+    device: Device
+        The device to allocate the arguments
+    args_info: T_ARG_INFO_JSON_OBJ_LIST
+        The arguments info
+    alloc_repeat: int
+        The number of times to repeat the allocation
+
+    Returns
+    -------
+    repeated_args: List[T_ARGUMENT_LIST]
+        The allocation args
+    """
+
+    def alloc_tensor(_, dtype, shape) -> tvm.runtime.Tensor:
+        arg = tvm.runtime.empty(shape=shape, dtype=dtype, device=device)
+        f_random_fill(arg)
+        return arg
+
+    def alloc_fail(*arg_info) -> None:
+        raise NotImplementedError(arg_info)
+
+    dispatcher: dict[Any, Callable] = {
+        "TENSOR": alloc_tensor,
+        None: alloc_fail,
+    }
+
+    repeated_args: list[T_ARGUMENT_LIST] = []
+    for _ in range(alloc_repeat):
+        args: T_ARGUMENT_LIST = []
+        arg_info: T_ARG_INFO_JSON_OBJ
+        for arg_info in args_info:
+            arg_type = arg_info[0]
+            arg: Any = dispatcher.get(arg_type, None)(*arg_info)
+            args.append(arg)
+        repeated_args.append(args)
+    return repeated_args
+
+
+def run_evaluator_common(
+    rt_mod: Module,
+    device: Device,
+    evaluator_config: EvaluatorConfig,
+    repeated_args: list[T_ARGUMENT_LIST],
+) -> list[float]:
+    """Common function to run the evaluator
+
+    Parameters
+    ----------
+    rt_mod: Module
+        The runtime module
+    device: Device
+        The device to run the evaluator
+    evaluator_config: EvaluatorConfig
+        The evaluator config
+    repeated_args: List[T_ARGUMENT_LIST]
+        The repeated arguments
+
+    Returns
+    -------
+    costs: List[float]
+        The evaluator results
+    """
+    evaluator = rt_mod.time_evaluator(
+        func_name=rt_mod.entry_name,
+        dev=device,
+        number=evaluator_config.number,
+        repeat=evaluator_config.repeat,
+        min_repeat_ms=evaluator_config.min_repeat_ms,
+        f_preproc="cache_flush_cpu_non_first_arg"
+        if evaluator_config.enable_cpu_cache_flush
+        else "",
+    )
+    repeated_costs: list[list[float]] = []
+    for args in repeated_args:
+        device.sync()
+        profile_result = evaluator(*args)
+        repeated_costs.append(profile_result.results)
+    costs = [float(cost) for cost in itertools.chain.from_iterable(repeated_costs)]
+    return costs

python/tvm/s_tir/meta_schedule/schedule/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+"""Per-block schedule rules in MetaSchedule"""
+
+from . import cpu, cuda, generic, x86

python/tvm/s_tir/meta_schedule/schedule/cpu/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+"""Per-block schedule rules in MetaSchedule for target key 'cpu'"""

python/tvm/s_tir/meta_schedule/schedule/cuda/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+"""Per-block schedule rules in MetaSchedule for target key 'cuda'"""
+
+from . import layout_transform

python/tvm/s_tir/meta_schedule/schedule/cuda/layout_transform.py

@@ -0,0 +1,581 @@
+# 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.
+# ruff: noqa: RUF005
+"""layout_transform scheduling rule for CUDA."""
+
+import math
+from collections import deque
+
+import tvm
+from tvm.s_tir import meta_schedule
+from tvm.s_tir.schedule import ExprRV, LoopRV, SBlockRV, Schedule
+
+## Tiling layout transforms:
+# Assume we have an input shape of [A, B, C, D] and want to layout transform
+# ABCD --> DBAC so the output shape would be [D, B, A, C].
+#
+# Consider reading from the input buffer in a cache-friendly fashion on CPU. We would
+# expect a loop structure like:
+# lAr, lBr, lCr, lDr = T.grid(A, B, C, D)
+#
+# Meanwhile consider writing to the output buffer in a cache-friendly fashion on CPU:
+# lDw, lBw, lAw, lCw = T.grid(D, B, A, C)
+#
+# Clearly in many scenarios it is impossible to guarantee contiguous writes and reads
+# within a single loop due to non-adjacent dimensions. Instead we work on transposing some
+# small sub-tensor of our input writing and then reading from shared memory. We must now
+# construct our submatrix so that reading and writing can both be done with some contiguous
+# access in global memory.
+#
+# Consider the case of a 2D transpose. For example [1024, 2048] -> [2048, 1024].
+# We note that if we deal with a submatrix of shape [32, 32] which corresponds
+# to the dimension of our input tensor, then rows of the submatrix are contiguous
+# in the input tensor. Meanwhile, columns of our submatrix are contiguous in our
+# output vector. Therefore, with this tile shape we have opportunity to read
+# contiguously in our input tensor and write to shared memory, and write contiguously
+# to our output tensor.
+#
+# The multiple dimensional case has a similar analogue. We want to allocate shared
+# memory per block of [`tile_size`, `tile_size`]. We want the inner most dimension
+# of our shared memory to correspond to contiguous reads from the input tensor and
+# the outer dimension to correspond to contiguous writes into the output tensor.
+#
+# In terms of the loop structure reading from the input tensor, the inner most loops
+# of our tile must correspond to the inner most dimensions of the input shape,
+# while the outer dimensions correspond to the inner most dimensions of the output shape.
+# To obtain an inner tile with this loop structure we factor out a contiguous `tile_size`
+# chunk of our loop in the shape of interest.
+#
+# An example is probably best to show this idea:
+# Let's say we want a layout transform of ABCD --> DCAB. With shape
+# [1024_a, 2_b, 32_c, 8_d] --> [8_d, 32_c, 1024_a, 2_b]
+#
+# And tile size 32.
+#
+# Then we initially have a coalesced-read loop pattern of:
+# T.grid(1024_a, 2_b, 32_c, 8_d)
+#
+# To obtain an inner tile of 32, we factor 4 from 32_c and 8 from 8_d:
+# T.grid(1024_a, 2_b, 8_c1, 1_d1, 4_c2t, 8_d2t)
+# T.grid(1024_a, 2_b, 8_cr, 1_dr, 32_dim1)
+#
+# To obtain an outer tile of 32, we factor from B then A to follow contiguous write
+# pattern:
+#
+# T.grid(64_a1, 1_b1, 8_cr, 1_dr, 16_a2t, 2_b2t, 32_dim1)
+# T.grid(64_ar, 1_br, 8_cr, 1_dr, 32_dim0, 32_dim1)
+#
+# Which allows us to read a tile with our wanted properties.
+# For writing we use the existing analysis infrastructure to generate the structure for writing.
+
+
+def tile_layout_transform(
+    sch: Schedule,
+    block_read: SBlockRV,
+    block_write: SBlockRV,
+    src_layout: str,
+    dst_layout: str,
+    input_shape: list[int],
+    tile_size: ExprRV,
+) -> tuple[SBlockRV, SBlockRV]:
+    """
+    High level tiling for layout transform block. Mutates sch in place.
+
+    Parameters
+    ----------
+    sch:
+        The initial schedule. We expect `block_read` and `block_write` to correspond to
+        the blocks which reads and writes from global memory respectively. We also expect
+        block_read's initial loops to follow
+
+    block_read:
+        The block which reads from global memory and writes to shared memory buffer.
+
+    block_write:
+        The block which writes to global memory and reads from shared memory buffer.
+
+    src_layout :
+        The src_layout, each character should appear once and also appear in dst_layout.
+        There should be not numeric characters and refer to potentially implicit reshapes.
+        E.g. the transform NCHW --> NCHW4c really implies NCcHW --> NCHWc. In this case
+        src_layout should be NCcHW.
+
+    dst_layout:
+        The dst_layout. There should not be numeric characters, e.g. NCHW4c becomes NCHWc.
+
+    input_shape:
+        The input shape after applying potentially implicit reshapes. Should match the loop
+        extants corresponding to src_layout.
+
+    tile_size:
+        The tile size of read and writes. There will be tile_size threads per block, each of which
+        reads up to tile_size elements.
+
+    Returns
+    -------
+    ret:
+        A tuple of the block that writes to global memory, and the block that reads from
+        global memory.
+    """
+
+    def pad_dimension_to_at_least_number(loop: LoopRV, requested_size: int):
+        """E.g. if loop has extant of 8 but we want 10, returns size 10 loop with padding."""
+        left, right = sch.split(loop, [None, requested_size])
+        return sch.fuse(left, right)
+
+    def pad_dimension_to_factor_of_tile_size(
+        loop: LoopRV, initial_size: int, tile_size: int = tile_size
+    ) -> tuple[LoopRV, int]:
+        """
+        Pads loop of given size until it is divisible into tile_size.
+        If the given size of the loop is greater than tile size. Do not pad.
+
+        examples:
+        - loop_size = 5 , tile_size = 32. loop_size --> 8
+        - loop_size = 5 , tile_size = 36. loop_size --> 6
+        - loop_size = 8 , tile_size = 32. loop_size --> 8  : since 8 already divides 32.
+        - loop_size = 33, tile_size = 32. loop_size --> 33 : since 33 > 32.
+
+        Returns padded loopRV and the new size.
+        """
+        if tile_size % initial_size == 0:
+            return loop, int(initial_size)
+
+        if initial_size > tile_size or initial_size == tile_size:
+            return loop, int(initial_size)
+
+        # if initial_size > tile_size return without change, factor = 1
+        size = initial_size
+        while (tile_size % size) % tile_size > 0:
+            size += 1
+
+        return pad_dimension_to_at_least_number(loop, size), int(size)
+
+    def spin_out_factor(
+        loops: list[LoopRV], loop_extants: list[int], index: int, factor_needed: int
+    ) -> tuple[list[LoopRV], list[int], int]:
+        """
+        Factor out the requested loop's dimensions to reach the requested factor and
+        places the requested factor as the innermost loop.
+
+        Updates the schedule in-place.
+
+        E.g. say we want to factors which eventually multiply to 32 (factor_needed).
+
+        Say we have the index we chose is a loop with an extant of 8.
+        E.g. loops / loop_extants = [3, 32, 6, 8], factor_needed = 32, index=3 (dim=8)
+            - 8 divides into 32 so we just split up the loop into two loops with extants 1 and 8.
+            - we then keep the 1-loop in place and move the new 8-loop to back of the list of loops
+            - ending loops / loop_extants = [3, 32, 6, 1, 8], remaining_factor_needed = 32 / 8 = 4
+
+        E.g. loops / loop_extants = [3, 32, 6, 8], factor_needed=32, index=0 (dim=3)
+            - 3 does not divide 32, so we pad until the extant divides 32, e.g. 4
+            - we then split up the loop into extants 1 and 4, moving the 4 to the back
+            - ending loops / loop_extants = [1, 32, 6, 8, 4], remaining_factor_needed = 32 / 4 = 8
+
+        E.g. loops / loop_extants = [3, 32, 6, 8], factor_needed=5, index=3 (dim=8)
+            - 8 is larger than 5 so we immediately do the splitting routine.
+            - the 8 extant loop becomes loops with extants 2 and 5
+            - ending loops / loop_extants = [1, 32, 6, 2, 5], remaining_factor_needed = 5 / 5 = 1
+
+        After updating loop ordering in place, returns the new list of loops, extants, and the
+        remaining factor needed.
+        """
+        cur_loop = loops[index]
+        cur_extant = loop_extants[index]
+
+        # Pad loops to divide evenly for factors needed, and split
+        new_loop, new_size = pad_dimension_to_factor_of_tile_size(
+            cur_loop, cur_extant, tile_size=factor_needed
+        )
+
+        split_factor = min(new_size, factor_needed)
+        new_loop_split, factored_loop = sch.split(new_loop, [None, split_factor])
+        factor_needed = factor_needed // split_factor
+
+        # update caching
+        loops[index] = new_loop_split
+        loops.append(factored_loop)
+
+        loop_extants[index] = math.ceil(int(new_size) / int(split_factor))
+        loop_extants.append(split_factor)
+
+        sch.reorder(*loops)
+        return loops, loop_extants, factor_needed
+
+    def factor_dim_in_order(
+        indices: list[int],
+        loops: list[LoopRV],
+        cur_loop_extants: list[int],
+        work_needed_inner_loop: int = tile_size,
+    ) -> tuple[list[LoopRV], list[int]]:
+        """Factors out the loops in the order of indices until we reach needed work.
+
+        Adds new loop factors to the back in reverse order of access. Returns new list
+        of loops and their extants.
+        """
+        for i in indices:
+            loops, cur_loop_extants, work_needed_inner_loop = spin_out_factor(
+                loops, cur_loop_extants, i, work_needed_inner_loop
+            )
+            if work_needed_inner_loop == 1:
+                break
+        return loops, cur_loop_extants
+
+    def get_high_level_loop_structure(
+        block_read: SBlockRV, input_shape: list[int], src_layout: str, dst_layout: str
+    ):
+        """Runs the factorization described above."""
+        # index 0 ... rank - 1 will always correspond to original loops
+        # perhaps after they have been factored.
+        rank = len(input_shape)
+        loops = sch.get_loops(block_read)
+        cur_loop_extants = list(input_shape)
+
+        # Factor dim0 tile size and fuse things together
+        loops, cur_loop_extants = factor_dim_in_order(
+            list(range(rank - 1, -1, -1)),
+            loops,
+            cur_loop_extants,
+            work_needed_inner_loop=tile_size,
+        )
+        # The factors which multiply to tile_size are now in back of our
+        # list of loops. However because we added them by traversing the inner
+        # dimensions, they are actually reversed order to guarantee the best access
+        # so reorder before fusing.
+        loops = loops[:rank] + loops[rank:][::-1]
+        cur_loop_extants = cur_loop_extants[:rank] + cur_loop_extants[rank::-1]
+        sch.reorder(*loops)
+        dim0_loop_tiled = sch.fuse(*loops[rank:])
+        loops = loops[:rank]
+        loops.append(dim0_loop_tiled)
+        cur_loop_extants = cur_loop_extants[:rank]
+        cur_loop_extants.append(tile_size)
+
+        # Same thing with dim1
+        # [:rank + 1], since we placed dim0_loop_tiled in the end which we want to keep
+        loops, cur_loop_extants = factor_dim_in_order(
+            list(
+                src_layout.index(dst_layout[loop_index_dst])
+                for loop_index_dst in range(rank - 1, -1, -1)
+            ),
+            loops,
+            cur_loop_extants,
+            work_needed_inner_loop=tile_size,
+        )
+        loops = loops[: rank + 1] + loops[rank + 1 :][::-1]
+        cur_loop_extants = cur_loop_extants[: rank + 1] + cur_loop_extants[rank + 1 :: -1]
+        sch.reorder(*loops)
+        dim1_loop_tiled = sch.fuse(*loops[rank + 1 :])
+        loops = loops[: rank + 1]
+        loops.append(dim1_loop_tiled)
+        cur_loop_extants = cur_loop_extants[: rank + 1]
+        cur_loop_extants.append(tile_size)
+
+    # After this we have loops: [loop1, loop2, loop3 ... dim0_tiled, dim1_tiled]
+    get_high_level_loop_structure(block_read, input_shape, src_layout, dst_layout)
+
+    # If there are insufficient elements, than dim1_tiled or dim0_tiled might be too small
+    # In all likelihood you should use a smaller tile, but I don't want things to crash.
+    loops = sch.get_loops(block_read)
+    loops[-1] = pad_dimension_to_at_least_number(loops[-1], tile_size)
+    loops[-2] = pad_dimension_to_at_least_number(loops[-2], tile_size)
+
+    # We want the dim0 and dim1 parent loops to be the inner most. Right now dim1 is inner-msot
+    # and we just need to move dim0 in (last dimension of dst).
+    # Recall right now structure is at least [l1 l2 ... ln, dim0_tiled, dim1_tiled]
+    # where n >= 2.
+    dim0_loop_index = src_layout.index(dst_layout[-1])
+    dim0_loop = loops.pop(dim0_loop_index)
+    loops = loops[:-3] + [dim0_loop, loops[-3]] + loops[-2:]
+    sch.reorder(*loops)
+
+    # After this loops are: [outer_loop (block binding), dim0_tiled, dim1_tiled]
+    outer_loop = sch.fuse(*loops[:-2])
+
+    # Now that we have the high level loop structure, we can use reverse_compute_at magic
+    # To get the proper loop structure for writing! This is also as coalesced as possible
+    # already.
+    sch.reverse_compute_at(block_write, outer_loop)
+
+    # Fuse all inner loops for the write into 2 loops, grab inner loops for both read
+    # and write block which have locality (we will bind these to threadIdx)
+    fused_write_loop = sch.fuse(*sch.get_loops(block_write)[1:])
+    _, inner_write_loop = sch.split(fused_write_loop, [None, tile_size])
+    inner_read_loop = sch.get_loops(block_read)[-2]
+
+    sch.bind(loop=outer_loop, thread_axis="blockIdx.x")
+    sch.bind(loop=inner_write_loop, thread_axis="threadIdx.x")
+    sch.bind(loop=inner_read_loop, thread_axis="threadIdx.x")
+
+    return block_write, block_read
+
+
+def create_cached_read(
+    sch: Schedule,
+    block_write: SBlockRV,
+    orig_input_shape: list[int],
+    orig_src_layout: str,
+    orig_dst_layout: str,
+) -> tuple[SBlockRV, list[int], str, str]:
+    """
+    Creates the cached read block with expected structure.
+
+    Loop extants should follow the input shape closely. E.g. if the input is [2, 6, 8], we
+    expect our loop structure to be T.grid(2, 6, 8). Possibly reshape to handle implicit reshapes,
+    in which case we will match the implicit reshape shape.
+
+    Layout transform allows semantics like NCHW --> NCHW4c. Which involves splitting the original C
+    axis into contiguous 4-element chunks. This axis is then moved to the end (NCHWc). This is
+    guaranteed by the operator to be done without additional padding. To handle this we just split
+    the associating axis (prev. type checking ensures C is divisible by 4)in src_layout found in
+    block_read. E.g. NCHW -> NCHW4c now becomes NC4cHW -> NCHW4c.
+
+    Note: NCHW4c --> NCHW is not allowed, so the only numeric digits will be in dst.
+
+    The returned layout strings will be santized and made compatible. E.g. NCHW --> NCHW4c becomes
+    NCcHW --> NCHWc.
+
+    TODO(AndrewZhaoLuo): Investigate using proper memory alignment to avoid bank conflict.
+
+    Parameters
+    ----------
+    sch:
+        The initial schedule. We expect `block_read`. We also expect
+        block_read's initial loops to follow the original input shape.
+
+    block_read:
+        The block which reads from global memory and writes to shared memory buffer.
+
+    orig_input_shape:
+        The input shape of the input buffer to the primfunc.
+
+    orig_src_layout:
+        The original src_layout string.
+
+    orig_dst_layout:
+        The original dst_layout string.
+
+    Returns
+    -------
+    ret:
+        A tuple of the cached read block, new input shape of shared memory buffer,
+        the new src_layout, and new dst_layout string.
+    """
+    # Figure out split dimensions, entries are (loop index in src_layout, split amount)
+    split_dimensions: list[tuple[int, int]] = []
+
+    # This is without numeric digits, e.g. NCHW4c -> NCHWc
+    new_dst_layout = []
+
+    # Use state machine to parse NCHW4c string
+    split_size = 0
+    for char in orig_dst_layout:
+        if char.isnumeric():
+            split_size = split_size * 10 + int(char)
+        else:
+            if char.islower():
+                # hit axis like 'c', need to find parent axis 'C' in src_layout
+                src_layout_index = orig_src_layout.index(char.upper())
+                split_dimensions.append((src_layout_index, split_size))
+            split_size = 0
+            new_dst_layout.append(char)
+
+    # If no splits were detected we are done
+    if len(split_dimensions) == 0:
+        block_read = sch.cache_read(block_write, 0, "shared")
+        return block_read, orig_input_shape, orig_src_layout, orig_dst_layout
+
+    # Calculate final input shapes, each of these are a single element for unsplit dims
+    # and tuples for split dims associated with the two new axis
+    input_shape: list[int | tuple] = list(orig_input_shape)
+    new_src_layout: list[str | tuple] = list(orig_src_layout)
+    for src_layout_split_index, split_factor in split_dimensions:
+        dimension_name = orig_src_layout[src_layout_split_index]
+        new_src_layout[src_layout_split_index] = (dimension_name, dimension_name.lower())
+        input_shape[src_layout_split_index] = (
+            orig_input_shape[src_layout_split_index] // split_factor,
+            split_factor,
+        )
+
+    # Unpack any tuples introduced via appending
+    def unpack_list(target_list) -> list:
+        output: list = []
+        for ele in target_list:
+            if isinstance(ele, tuple):
+                output.extend(ele)
+            else:
+                output.append(ele)
+        return output
+
+    new_src_layout_str = "".join(unpack_list(new_src_layout))
+    new_dst_layout_str = "".join(unpack_list(new_dst_layout))
+
+    # Write block loop extants match
+    dst_to_src_map = [new_dst_layout_str.index(dim) for dim in new_src_layout_str]
+    block_read = sch.reindex_cache_read(
+        block_write,
+        read_buffer_index=0,
+        index_map=tvm.tirx.IndexMap.from_func(
+            lambda *loops: [loops[dst_to_src_map[i]] for i, _ in enumerate(loops)],
+            ndim=len(new_src_layout_str),
+        ),
+        storage_scope="shared",
+    )
+
+    loops_read = sch.get_loops(block_read)
+    sch.reorder(
+        *[loops_read[new_dst_layout_str.index(dst_dim_name)] for dst_dim_name in new_src_layout_str]
+    )
+    return block_read, unpack_list(input_shape), new_src_layout_str, new_dst_layout_str
+
+
+def auto_inline_into(sch: Schedule, start_block: SBlockRV) -> SBlockRV:
+    """
+    Inlines given start_block's consumers and future dependencies into start_block.
+
+    Parameters
+    ----------
+    sch:
+        The initial schedule.
+
+    start_block:
+        The block to inline into, should be a block which reads and writes to global memory, doing
+        layout transform.
+
+    Returns
+    -------
+    ret:
+        The new block inlined into it's consumers.
+    """
+    # Rules defined by DefaultCUDA schedule_rule set.
+    autoinline_rule = meta_schedule.schedule_rule.AutoInline(
+        into_producer=True,
+        into_consumer=False,
+        inline_const_tensor=True,
+        disallow_if_then_else=False,
+        require_injective=False,
+        require_ordered=False,
+    )
+
+    fringe = deque(sch.get_consumers(start_block))
+    visited = set()
+    while len(fringe) > 0:
+        cur_block = fringe.popleft()
+        if cur_block in visited:
+            continue
+
+        visited.add(cur_block)
+        consumer_blocks = sch.get_consumers(cur_block)
+        fringe.extend(consumer_blocks)
+
+        sch = autoinline_rule.apply(sch, cur_block)[0]
+
+
+def get_max_tile_size() -> int:
+    """Returns the max tile size.
+
+    This is assuming only threads in a warp can have coalesced accesses. 32 is the default if
+    no target information can be gotten.
+    """
+    max_tile_size = 32
+    cur_target = tvm.target.Target.current()
+    if cur_target is not None and "thread_warp_size" in cur_target.attrs:
+        max_tile_size = int(cur_target.attrs["thread_warp_size"])
+    return max_tile_size
+
+
+@tvm.register_global_func("s_tir.meta_schedule.cuda.layout_transform")
+def cuda_layout_transform_schedule_rule(
+    sch: Schedule, block: SBlockRV, testing_tile_sizes: list[int] | None = None
+) -> list[Schedule]:
+    """
+    Applies tiling scheme to layout transform task (potentially fused with other injective funcs).
+
+    Returned schedules will be the default schedule, as well as tiled versions with tile_size in
+    the range of 2,3...threads_per_warp.
+
+    This is assuming only threads in a warp can have coalesced accesses. 32 is the default if
+    no target information can be gotten.
+
+    Parameters
+    ----------
+    sch:
+        The initial schedule.
+
+    block:
+        The block corresponding to the layout transform.
+        Should be a block which reads and writes to global memory, doing layout transform.
+
+    testing_tile_sizes:
+        A list of tile sizes to try, overriding normal settings. For testing. None means
+        ignore. Else overrides normal settings of tile sizes to try.
+
+    Returns
+    -------
+    ret:
+        A list of new schedules to try.
+    """
+    # Info needed for tiling
+    src_layout = sch.get_sref(block).stmt.annotations["src_layout"]
+    dst_layout = sch.get_sref(block).stmt.annotations["dst_layout"]
+    input_shape = [int(c) for c in sch.get_sref(block).stmt.annotations["input_shape"]]
+
+    schedules = []
+
+    # Always include the default schedules which will be handled via AutoBind schedule rule
+    # Except during testing
+    if not testing_tile_sizes:
+        schedules.append(sch)
+
+    sch = sch.copy()
+
+    # Inline consumers of the layout transform into the layout transform block.
+    # Normally default for injective schedules but must manually be called in new schedule rule
+    # for consumers of the layout transform. TODO(AndrewZhaoLuo): Figure out why this is the case.
+    auto_inline_into(sch, block)
+
+    # Setup up basic structure of schedule of creating read into shared mem, before applying tiling
+    # Outer loop structure of read block matches that of src_layout
+    # E.g. if input_shape is [4, 6, 8]. Loops for read block will be
+    # for i, j, k in T.grid(4, 6, 8):
+    #     ...
+    # Read block will read from global memory coalesced at the start
+    # Assume write to output global memory is coalesced in block_write
+    #
+    # This also handles the case where there is an implicit reshape going on.
+    # e.g. NCHW -> NCHW4c which is equivalent to reshaping NCHW
+    # to NCcHW and then applying the new layout where the extant of c is 4.
+    # Grab final input shape and src and dst layouts with possible implicit reshape.
+    block_read, input_shape, src_layout, dst_layout = create_cached_read(
+        sch, block, input_shape, src_layout, dst_layout
+    )
+
+    # Try tile size 2,3...threads_per_warp as tile size of 1 has no coaslescing.
+    if testing_tile_sizes is None:
+        tile_sizes = list(range(2, get_max_tile_size() + 1))
+    else:
+        tile_sizes = testing_tile_sizes
+
+    for tile_size in tile_sizes:
+        new_sch = sch.copy()
+        tile_layout_transform(
+            new_sch, block_read, block, src_layout, dst_layout, input_shape, tile_size
+        )
+        schedules.append(new_sch)
+
+    return schedules

python/tvm/s_tir/meta_schedule/schedule/generic/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+"""Per-block schedule rules in MetaSchedule for generic cases"""

python/tvm/s_tir/meta_schedule/schedule/x86/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+"""Per-block schedule rules in MetaSchedule for target key 'x86'"""

python/tvm/s_tir/meta_schedule/schedule_rule/__init__.py

@@ -0,0 +1,38 @@
+# 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.
+"""
+The tvm.s_tir.meta_schedule.schedule_rule package.
+Meta Schedule schedule rules are used for modification of
+blocks in a schedule. See also PostOrderApply.
+"""
+
+from .add_rfactor import AddRFactor
+from .apply_custom_rule import ApplyCustomRule
+from .auto_bind import AutoBind
+from .auto_inline import AutoInline, InlineConstantScalars
+from .cross_thread_reduction import CrossThreadReduction
+from .multi_level_tiling import (
+    MultiLevelTiling,
+    MultiLevelTilingTensorCore,
+    MultiLevelTilingWideVector,
+    MultiLevelTilingWithIntrin,
+    ReuseType,
+)
+from .parallel_vectorize_unroll import ParallelizeVectorizeUnroll
+from .random_compute_location import RandomComputeLocation
+from .schedule_rule import PyScheduleRule, ScheduleRule

python/tvm/s_tir/meta_schedule/schedule_rule/add_rfactor.py

@@ -0,0 +1,48 @@
+# 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.
+"""Add-rfactor Rule that add-rfactor to some blocks if needed"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .schedule_rule import ScheduleRule
+
+
+@register_object("s_tir.meta_schedule.AddRFactor")
+class AddRFactor(ScheduleRule):
+    """Rules for add-rfactor to some blocks if needed.
+
+    Parameters
+    ----------
+    max_jobs_per_core: int
+        The maximum number of jobs to be launched per CPU core. It sets the uplimit of CPU
+        parallelism, i.e. `num_cores * max_jobs_per_core`.
+        Use -1 to disable parallelism.
+    max_innermost_factor: Optional[int] = None
+        The maximum size of the innermost factor. None means no limit.
+    """
+
+    def __init__(
+        self,
+        max_jobs_per_core: int = 16,
+        max_innermost_factor: int | None = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleAddRFactor,  # type: ignore # pylint: disable=no-member
+            max_jobs_per_core,
+            max_innermost_factor,
+        )

python/tvm/s_tir/meta_schedule/schedule_rule/apply_custom_rule.py

@@ -0,0 +1,34 @@
+# 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.
+"""Create a rule that applies customized rules registered using block attribute `schedule_rule`.
+The rule will be dispatched according to target keys."""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .schedule_rule import ScheduleRule
+
+
+@register_object("s_tir.meta_schedule.ApplyCustomRule")
+class ApplyCustomRule(ScheduleRule):
+    """A rule that applies customized rules registered using block attribute `schedule_rule`.
+    The rule will be dispatched according to target keys."""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleApplyCustomRule,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/schedule_rule/auto_bind.py

@@ -0,0 +1,52 @@
+# 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.
+"""Auto-bind Rule that binds blocks to threads if needed"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .schedule_rule import ScheduleRule
+
+
+@register_object("s_tir.meta_schedule.AutoBind")
+class AutoBind(ScheduleRule):
+    """Auto bind loops around the block to BlockIdx and ThreadIdx
+
+    Parameters
+    ----------
+    max_threadblocks: int
+        The maximum number of threadblock on GPU.
+    thread_extents: Optional[List[int]]
+        Candidates of thread axis extent.
+    max_threads_per_block: int
+        The maximum number of threads per block, if it is known when this schedule rule is created.
+    """
+
+    def __init__(
+        self,
+        max_threadblocks: int = 256,
+        thread_extents: list[int] | None = None,
+        max_threads_per_block: int = -1,
+    ) -> None:
+        if thread_extents is None:
+            thread_extents = [32, 64, 128, 256, 512, 1024]
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleAutoBind,  # type: ignore # pylint: disable=no-member
+            max_threadblocks,
+            thread_extents,
+            max_threads_per_block,
+        )

python/tvm/s_tir/meta_schedule/schedule_rule/auto_inline.py

@@ -0,0 +1,83 @@
+# 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.
+"""Auto-Inline. Rule that inlines spatial blocks if it satisfies some conditions"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .schedule_rule import ScheduleRule
+
+
+@register_object("s_tir.meta_schedule.AutoInline")
+class AutoInline(ScheduleRule):
+    """Rule that inlines spatial blocks if it satisfies some conditions
+
+    Parameters
+    ----------
+    into_producer : bool
+        If allows to inline a block into its producer
+    into_consumer : bool
+        If allows to inline a block into its consumer
+    inline_const_tensor : bool
+        Always inline constant tensors
+    disallow_if_then_else : bool
+        Always disallow if-then-else-like constructs
+    require_injective : bool
+        Always require the read-to-write mapping to be ordered
+    require_ordered : bool
+        Always require the read-to-write mapping to be injective
+    disallow_op : Optional[List[str]]
+        The operators that are disallowed in auto inline
+    """
+
+    def __init__(
+        self,
+        into_producer: bool,
+        into_consumer: bool,
+        inline_const_tensor: bool,
+        disallow_if_then_else: bool,
+        require_injective: bool,
+        require_ordered: bool,
+        disallow_op: list[str] | None = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleAutoInline,  # type: ignore # pylint: disable=no-member
+            into_producer,
+            into_consumer,
+            inline_const_tensor,
+            disallow_if_then_else,
+            require_injective,
+            require_ordered,
+            disallow_op,
+        )
+
+
+@register_object("s_tir.meta_schedule.InlineConstantScalars")
+class InlineConstantScalars(ScheduleRule):
+    """Inline blocks that produce a constant scalar.
+
+    Such blocks get in the way of ReverseComputeInline during AutoInline, since they are also
+    counted as a producer block unless they are inlined first. So it is recommended to run
+    InlineConstantScalars before AutoInline.
+    """
+
+    def __init__(
+        self,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleInlineConstantScalars,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.py

@@ -0,0 +1,40 @@
+# 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.
+"""Rules which apply cross-thread reduction to some reduction blocks correspondingly when needed"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .schedule_rule import ScheduleRule
+
+
+@register_object("s_tir.meta_schedule.CrossThreadReduction")
+class CrossThreadReduction(ScheduleRule):
+    """A schedule rule which applies cross-thread reduction to some reduction blocks
+    correspondingly when needed
+
+    Parameters
+    ----------
+    thread_extents: List[int]
+        Candidates of thread axis extent (values are required to be positive).
+    """
+
+    def __init__(self, thread_extents: list[int]) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleCrossThreadReduction,  # type: ignore # pylint: disable=no-member
+            thread_extents,
+        )

python/tvm/s_tir/meta_schedule/schedule_rule/multi_level_tiling.py

@@ -0,0 +1,237 @@
+# 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.
+"""Multi-level tiling with reuse."""
+
+from collections.abc import Callable, Mapping
+from typing import Any, NamedTuple
+
+from tvm_ffi import register_object
+
+from tvm.s_tir.schedule import SBlockRV, Schedule
+
+from .. import _ffi_api
+from .schedule_rule import ScheduleRule
+
+
+class ReuseType(NamedTuple):
+    """Reuse type."""
+
+    req: str
+    levels: list[int]
+    scope: str
+
+    def as_dict(self) -> dict[str, Any]:
+        """Return the dict representation of the reuse type."""
+        return {
+            "req": self.req,
+            "levels": self.levels,
+            "scope": self.scope,
+        }
+
+
+@register_object("s_tir.meta_schedule.MultiLevelTiling")
+class MultiLevelTiling(ScheduleRule):
+    """Multi-level tiling with reuse.
+
+    Parameters
+    ----------
+    structure : str
+        The tiling structure. Recommended:
+        - 'SSRSRS' on CPU
+        - 'SSSRRSRS' on GPU
+    tile_bind : Optional[List[str]]
+        For each level of tiles, which thread axis it is bound to. Recommended:
+        - None on CPU
+        - [blockIdx.x, vthread.x, threadIdx.x] on GPU
+    max_innermost_factor : Optional[int]
+        The maximum size of the innermost factor. None means no limit
+    vector_load_lens : Optional[List[int]]
+        The length of vector lane in vectorized cooperative fetching.
+        None means disable vectorization
+    reuse_read : Optional[ReuseType]
+        Data reuse configuration for reading. None means no reuse.
+    reuse_write : Optional[ReuseType]
+        Data reuse configuration for writing. None means no reuse.
+    filter_fn: Optional[Callable[[Schedule, SBlockRV], bool]]
+        A function that can be passed to overwrite the default condition for applying
+        MultiLevelTiling to a block. This is useful if there is a need to apply MultiLevelTiling
+        to an operation / block which is ignored by default. This function should return True
+        for a block that should be tiled (based on the block name, for example).
+    """
+
+    def __init__(
+        self,
+        structure: str,
+        tile_binds: list[str] | None = None,
+        max_innermost_factor: int | None = None,
+        vector_load_lens: list[int] | None = None,
+        reuse_read: ReuseType | None = None,
+        reuse_write: ReuseType | None = None,
+        filter_fn: Callable[[Schedule, SBlockRV], bool] | None = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleMultiLevelTiling,  # type: ignore # pylint: disable=no-member
+            structure,
+            tile_binds,
+            max_innermost_factor,
+            vector_load_lens,
+            reuse_read.as_dict() if reuse_read is not None else None,
+            reuse_write.as_dict() if reuse_write is not None else None,
+            filter_fn,
+        )
+
+
+@register_object("s_tir.meta_schedule.MultiLevelTilingWithIntrin")
+class MultiLevelTilingWithIntrin(ScheduleRule):
+    """Extension of MultiLevelTiling for auto-tensorizing with a single intrinsic.
+
+    Parameters
+    ----------
+    intrin_name : str
+        The name of a tensor intrinsic, must be registerd via TensorIntrin.register(...) beforehand
+    structure : str
+        The tiling structure. Recommended:
+        - 'SSRSRS' on CPU
+        - 'SSSRRSRS' on GPU
+    tile_bind : Optional[List[str]]
+        For each level of tiles, which thread axis it is bound to. Recommended:
+        - None on CPU
+        - [blockIdx.x, vthread.x, threadIdx.x] on GPU
+    max_innermost_factor : Optional[int]
+        The maximum size of the innermost factor. None means no limit
+    vector_load_lens : Optional[List[int]]
+        The length of vector lane in vectorized cooperative fetching.
+        None means disable vectorization
+    reuse_read : Optional[ReuseType]
+        Data reuse configuration for reading. None means no reuse.
+    reuse_write : Optional[ReuseType]
+        Data reuse configuration for writing. None means no reuse.
+    """
+
+    def __init__(
+        self,
+        intrin_name: str,
+        structure: str,
+        tile_binds: list[str] | None = None,
+        max_innermost_factor: int | None = None,
+        vector_load_lens: list[int] | None = None,
+        reuse_read: ReuseType | None = None,
+        reuse_write: ReuseType | None = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleMultiLevelTilingWithIntrin,  # type: ignore # pylint: disable=no-member
+            intrin_name,
+            structure,
+            tile_binds,
+            max_innermost_factor,
+            vector_load_lens,
+            reuse_read.as_dict() if reuse_read is not None else None,
+            reuse_write.as_dict() if reuse_write is not None else None,
+        )
+
+
+@register_object("s_tir.meta_schedule.MultiLevelTilingTensorCore")
+class MultiLevelTilingTensorCore(ScheduleRule):
+    """Extension of MultiLevelTiling for auto-tensorizing with multiple groups of candidate tensor
+    core intrinsics.
+
+    Parameters
+    ----------
+    intrin_groups : List[Mapping[str, str]]
+        A list of groups of tensor core intrinsics. The map should contains key "init", "load_a",
+        "load_b", "compute", "store", which represent the tensor intrin for initialization,
+        loading operand A, loading operand B, tensor core computation, storing the result.
+        The value of the map should be names of tensor intrinsics, must be registerd via
+        TensorIntrin.register(...) beforehand
+    structure : str
+        The tiling structure. Recommended:
+        - 'SSSRRSRS' on GPU
+    tile_bind : Optional[List[str]]
+        For each level of tiles, which thread axis it is bound to. Recommended:
+        - [blockIdx.y, vthread.x, threadIdx.y] on GPU
+    max_innermost_factor : Optional[int]
+        The maximum size of the innermost factor. None means no limit
+    vector_load_lens : Optional[List[int]]
+        The length of vector lane in vectorized cooperative fetching.
+        None means disable vectorization
+    reuse_read : Optional[ReuseType]
+        Data reuse configuration for reading. None means no reuse.
+    reuse_write : Optional[ReuseType]
+        Data reuse configuration for writing. None means no reuse.
+    use_software_pipeline : bool
+        Whether to use the software pipeline.
+    """
+
+    def __init__(
+        self,
+        intrin_groups: list[Mapping[str, str]],
+        structure: str,
+        tile_binds: list[str] | None = None,
+        max_innermost_factor: int | None = None,
+        vector_load_lens: list[int] | None = None,
+        reuse_read: ReuseType | None = None,
+        reuse_write: ReuseType | None = None,
+        use_software_pipeline: bool = False,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleMultiLevelTilingTensorCore,  # type: ignore # pylint: disable=no-member
+            intrin_groups,
+            structure,
+            tile_binds,
+            max_innermost_factor,
+            vector_load_lens,
+            reuse_read.as_dict() if reuse_read is not None else None,
+            reuse_write.as_dict() if reuse_write is not None else None,
+            use_software_pipeline,
+        )
+
+
+@register_object("s_tir.meta_schedule.MultiLevelTilingWideVector")
+class MultiLevelTilingWideVector(ScheduleRule):
+    """Extension of MultiLevelTiling for backends with wide vectors. The loop over the innermost
+    spatial axis of the output buffer is always vectorized with the maximum vector length.
+
+    Parameters
+    ----------
+    structure : str
+        The tiling structure. 'SSRSRS' is recommended.
+    vector_length_in_bits: int
+        The length of a vector register in bits.
+    max_innermost_factor : Optional[int]
+        The maximum size of the innermost factor. None means no limit
+    reuse_read : Optional[ReuseType]
+        Data reuse configuration for reading. None means no reuse.
+    reuse_write : Optional[ReuseType]
+        Data reuse configuration for writing. None means no reuse.
+    """
+
+    def __init__(
+        self,
+        structure: str,
+        vector_length_in_bits: int,
+        max_innermost_factor: int | None = None,
+        reuse_read: ReuseType | None = None,
+        reuse_write: ReuseType | None = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleMultiLevelTilingWideVector,  # type: ignore # pylint: disable=no-member
+            structure,
+            vector_length_in_bits,
+            max_innermost_factor,
+            reuse_read.as_dict() if reuse_read is not None else None,
+            reuse_write.as_dict() if reuse_write is not None else None,
+        )

python/tvm/s_tir/meta_schedule/schedule_rule/parallel_vectorize_unroll.py

@@ -0,0 +1,63 @@
+# 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.
+"""Rule that mark parallelize, vectorize and unroll to the root block. The mark will be applied to
+each block in a follow-up post processor"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .schedule_rule import ScheduleRule
+
+
+@register_object("s_tir.meta_schedule.ParallelizeVectorizeUnroll")
+class ParallelizeVectorizeUnroll(ScheduleRule):
+    """Rule that mark parallelize, vectorize and unroll to the root block. The mark will be applied
+    to each block in a follow-up post processor
+
+    Parameters
+    ----------
+    max_jobs_per_core: int
+        The maximum number of jobs to be launched per CPU core. It sets the upper limit of CPU
+        parallelism, i.e. `num_cores * max_jobs_per_core`.
+        Use -1 to disable parallelism.
+    max_vectorize_extent: int
+        The maximum extent to be vectorized. It sets the upper limit of the hardware target
+        vectorization.
+        Use -1 to disable vectorization.
+    unroll_max_steps: Optional[List[int]]
+        The options of the maximum number of unroll steps to be done.
+        Use None to disable unroll
+    unroll_explicit: bool
+        Whether to explicitly unroll the loop, or just add an "unroll" pragma
+    """
+
+    def __init__(
+        self,
+        max_jobs_per_core: int = 16,
+        max_vectorize_extent: int = 16,
+        unroll_max_steps: list[int] | None = None,
+        unroll_explicit: bool = True,
+    ) -> None:
+        if unroll_max_steps is None:
+            unroll_max_steps = []
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleParallelizeVectorizeUnroll,  # type: ignore # pylint: disable=no-member
+            max_jobs_per_core,
+            max_vectorize_extent,
+            unroll_max_steps,
+            unroll_explicit,
+        )

python/tvm/s_tir/meta_schedule/schedule_rule/random_compute_location.py

@@ -0,0 +1,32 @@
+# 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.
+"""Rule that randomly select a compute-at location for a free block"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .schedule_rule import ScheduleRule
+
+
+@register_object("s_tir.meta_schedule.RandomComputeLocation")
+class RandomComputeLocation(ScheduleRule):
+    """A rule that randomly select a compute-at location for a free block"""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRuleRandomComputeLocation,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py

@@ -0,0 +1,191 @@
+# 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.
+# ruff: noqa: RUF012
+"""
+Meta Schedule schedule rules are used for modification of
+blocks in a schedule. See also PostOrderApply.
+"""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+from tvm.s_tir.schedule import SBlockRV, Schedule
+
+from .. import _ffi_api
+
+if TYPE_CHECKING:
+    from ..tune_context import TuneContext
+
+
+@register_object("s_tir.meta_schedule.ScheduleRule")
+class ScheduleRule(Object):
+    """Rules to modify a block in a schedule."""
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        """Initialize the schedule rule with a tune context.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for initializing the schedule rule.
+        """
+        _ffi_api.ScheduleRuleInitializeWithTuneContext(  # type: ignore # pylint: disable=no-member
+            self, context
+        )
+
+    def apply(self, sch: Schedule, block: SBlockRV) -> list[Schedule]:
+        """Apply a schedule rule to the specific block in the given schedule.
+
+        Parameters
+        ----------
+        sch : tvm.s_tir.Schedule
+            The schedule to be modified.
+        block : SBlockRV
+            The specific block to apply the schedule rule.
+
+        Returns
+        -------
+        design_spaces : List[tvm.s_tir.Schedule]
+            The list of schedules generated by applying the schedule rule.
+        """
+        return _ffi_api.ScheduleRuleApply(  #  type: ignore # pylint: disable=no-member
+            self, sch, block
+        )
+
+    def clone(self) -> "ScheduleRule":
+        """Deep clone the schedule rule.
+
+        Returns
+        -------
+        cloned_rule : ScheduleRule
+            The cloned schedule rule.
+        """
+        return _ffi_api.ScheduleRuleClone(self)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def create(kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]) -> list["ScheduleRule"]:
+        """Create a list of schedule rules for the given kind.
+
+        Parameters
+        ----------
+        kind : Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]
+            The kind of the schedule rules.
+
+        Returns
+        -------
+        rules : List[ScheduleRule]
+            The list of schedule rules.
+        """
+        funcs = {
+            # pylint: disable=no-member
+            "llvm": _ffi_api.ScheduleRuleDefaultLLVM,  # type: ignore
+            "cuda": _ffi_api.ScheduleRuleDefaultCUDA,  # type: ignore
+            "cuda-tensorcore": _ffi_api.ScheduleRuleDefaultCUDATensorCore,  # type: ignore
+            "hexagon": _ffi_api.ScheduleRuleDefaultHexagon,  # type: ignore
+            # pylint: enable=no-member
+        }
+        for k, v in funcs.items():
+            if k == kind:
+                return v()
+        raise ValueError(f"Unsupported kind {kind} for schedule rule creation.")
+
+
+create = ScheduleRule.create  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.PyScheduleRule")
+class _PyScheduleRule(ScheduleRule):
+    """
+    A TVM object schedule rule to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyScheduleRule
+    """
+
+    def __init__(
+        self,
+        f_initialize_with_tune_context: Callable | None = None,
+        f_apply: Callable | None = None,
+        f_clone: Callable | None = None,
+    ):
+        """Constructor."""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.ScheduleRulePyScheduleRule,  # type: ignore # pylint: disable=no-member
+            f_initialize_with_tune_context,
+            f_apply,
+            f_clone,
+        )
+
+
+class PyScheduleRule:
+    """
+    An abstract schedule rule with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PyScheduleRule,
+        "methods": ["_initialize_with_tune_context", "apply", "clone"],
+    }
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        """Initialize the schedule rule with a tune context.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for initializing the schedule rule.
+        """
+        raise NotImplementedError
+
+    def apply(self, sch: Schedule, block: SBlockRV) -> list[Schedule]:
+        """Apply a schedule rule to the specific block in the given schedule.
+
+        Parameters
+        ----------
+        sch : Schedule
+            The schedule to be modified.
+        block : SBlockRV
+            The specific block to apply the schedule rule.
+
+        Returns
+        -------
+        design_spaces : List[Schedule]
+            The list of schedules generated by applying the schedule rule.
+        """
+        raise NotImplementedError
+
+    def clone(self) -> ScheduleRule:
+        """Deep clone the schedule rule.
+
+        Returns
+        -------
+        cloned_rule : ScheduleRule
+            The cloned schedule rule.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/search_strategy/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+"""
+The tvm.s_tir.meta_schedule.search_strategy package.
+Meta Schedule search strategy utilizes the design spaces given
+to generate measure candidates.
+"""
+
+from .evolutionary_search import EvolutionarySearch
+from .replay_func import ReplayFunc
+from .replay_trace import ReplayTrace
+from .search_strategy import MeasureCandidate, PySearchStrategy, SearchStrategy, create

python/tvm/s_tir/meta_schedule/search_strategy/evolutionary_search.py

@@ -0,0 +1,82 @@
+# 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.
+"""Evolutionary Search Strategy"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .search_strategy import SearchStrategy
+
+
+@register_object("s_tir.meta_schedule.EvolutionarySearch")
+class EvolutionarySearch(SearchStrategy):
+    """
+    Replay Trace Search Strategy is a search strategy that always replays the trace by removing its
+    decisions so that the decisions would be randomly re-generated.
+
+    Parameters
+    ----------
+    population_size : int
+        The initial population of traces from measured samples and randomly generated samples.
+    init_measured_ratio : int
+        The ratio of measured samples in the initial population.
+    init_min_unmeasured : int
+        The minimal size of unmeasured population in the initial sampling.
+    max_fail_count : int
+        The maximum number of failure during initial sampling.
+    genetic_num_iters : int
+        The number of iterations for genetic algorithm.
+    genetic_mutate_prob : float
+        The probability of mutation.
+    genetic_max_fail_count : int
+        The maximum number to retry mutation.
+    eps_greedy : float
+        The ratio of greedy selected samples in the final picks.
+    """
+
+    population_size: int
+    init_measured_ratio: int
+    init_min_unmeasured: int
+    genetic_num_iters: int
+    genetic_mutate_prob: float
+    genetic_max_fail_count: int
+    eps_greedy: float
+
+    def __init__(
+        self,
+        *,
+        population_size: int = 512,
+        init_measured_ratio: float = 0.2,
+        init_min_unmeasured: int = 50,
+        max_fail_count: int = 5,
+        genetic_num_iters: int = 4,
+        genetic_mutate_prob: float = 0.85,
+        genetic_max_fail_count: int = 10,
+        eps_greedy: float = 0.05,
+    ) -> None:
+        """Constructor"""
+        self.__init_handle_by_constructor__(
+            _ffi_api.SearchStrategyEvolutionarySearch,  # type: ignore # pylint: disable=no-member
+            population_size,
+            init_measured_ratio,
+            init_min_unmeasured,
+            max_fail_count,
+            genetic_num_iters,
+            genetic_mutate_prob,
+            genetic_max_fail_count,
+            eps_greedy,
+        )

python/tvm/s_tir/meta_schedule/search_strategy/replay_func.py

@@ -0,0 +1,43 @@
+# 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.
+"""Replay Trace Search Strategy"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .search_strategy import SearchStrategy
+
+
+@register_object("s_tir.meta_schedule.ReplayFunc")
+class ReplayFunc(SearchStrategy):
+    """
+    Replay Func Search Strategy is a search strategy that generates measure candidates by
+    calling a design space generator and transform the design space.
+
+    Parameters
+    ----------
+    num_trials_per_iter : int
+        Number of trials per iteration.
+    max_trials_per_task : int
+        Total number of trials for one task
+    """
+
+    def __init__(self):
+        """Constructor"""
+        self.__init_handle_by_constructor__(
+            _ffi_api.SearchStrategyReplayFunc,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/s_tir/meta_schedule/search_strategy/replay_trace.py

@@ -0,0 +1,44 @@
+# 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.
+"""Replay Trace Search Strategy"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .search_strategy import SearchStrategy
+
+
+@register_object("s_tir.meta_schedule.ReplayTrace")
+class ReplayTrace(SearchStrategy):
+    """
+    Replay Trace Search Strategy is a search strategy that always replays the trace by removing its
+    decisions so that the decisions would be randomly re-generated.
+
+    Parameters
+    ----------
+    max_fail_count : int
+        Max number of failures during trace replaying.
+    """
+
+    max_fail_count: int
+
+    def __init__(self, max_fail_count: int = 100):
+        """Constructor"""
+        self.__init_handle_by_constructor__(
+            _ffi_api.SearchStrategyReplayTrace,  # type: ignore # pylint: disable=no-member
+            max_fail_count,
+        )

python/tvm/s_tir/meta_schedule/search_strategy/search_strategy.py

@@ -0,0 +1,342 @@
+# 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.
+# ruff: noqa: RUF012
+"""
+Meta Schedule search strategy that generates the measure
+candidates for measurement.
+"""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Optional, Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+from tvm.s_tir.schedule import Schedule
+
+from .. import _ffi_api
+from ..arg_info import ArgInfo
+from ..runner import RunnerResult
+
+if TYPE_CHECKING:
+    from ..cost_model import CostModel
+    from ..database import Database
+    from ..tune_context import TuneContext
+
+
+@register_object("s_tir.meta_schedule.MeasureCandidate")
+class MeasureCandidate(Object):
+    """Measure candidate class.
+
+    Parameters
+    ----------
+    sch : tvm.s_tir.Schedule
+        The schedule to be measured.
+    args_info : List[ArgInfo]
+        The argument information.
+    """
+
+    sch: Schedule
+    args_info: list[ArgInfo]
+
+    def __init__(
+        self,
+        sch: Schedule,
+        args_info: list[ArgInfo],
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        sch : tvm.s_tir.Schedule
+            The schedule to be measured.
+        args_info : List[ArgInfo]
+            The argument information.
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.MeasureCandidate,  # type: ignore # pylint: disable=no-member
+            sch,
+            args_info,
+        )
+
+
+@register_object("s_tir.meta_schedule.SearchStrategy")
+class SearchStrategy(Object):
+    """Search strategy is the class that generates the measure candidates."""
+
+    SearchStrategyType = Union[
+        "SearchStrategy",
+        Literal[
+            "replay-func",
+            "replay-trace",
+            "evolutionary",
+        ],
+    ]
+
+    def __new__(cls, *args, **kwargs):  # pylint: disable=unused-argument
+        """Prevent direct instantiation of abstract SearchStrategy class.
+
+        SearchStrategy is an abstract class and cannot be directly instantiated.
+        Use SearchStrategy.create() or a concrete subclass instead.
+        """
+        if cls is SearchStrategy:
+            raise TypeError(
+                "Cannot instantiate abstract class SearchStrategy. "
+                "Use SearchStrategy.create() with a valid strategy type "
+                "(e.g., 'evolutionary', 'replay-trace', 'replay-func') "
+                "or use a concrete subclass instead."
+            )
+        return super().__new__(cls)  # pylint: disable=no-value-for-parameter
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        """Initialize the search strategy with tuning context.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for initialization.
+        """
+        _ffi_api.SearchStrategyInitializeWithTuneContext(  # type: ignore # pylint: disable=no-member
+            self, context
+        )
+
+    def pre_tuning(
+        self,
+        max_trials: int,
+        num_trials_per_iter: int,
+        design_spaces: list[Schedule],
+        database: Optional["Database"] = None,
+        cost_model: Optional["CostModel"] = None,
+    ) -> None:
+        """Pre-tuning for the search strategy.
+
+        Parameters
+        ----------
+        max_trials : int
+            The maximum number of trials.
+        num_trials_per_iter : int
+            The number of trials per iteration.
+        design_spaces : List[tvm.s_tir.Schedule]
+            The design spaces used during tuning process.
+        database : Optional[Database] = None
+            The database used during tuning process.
+        cost_model : Optional[CostModel] = None
+            The cost model used during tuning process.
+        """
+        _ffi_api.SearchStrategyPreTuning(  # type: ignore # pylint: disable=no-member
+            self,
+            max_trials,
+            num_trials_per_iter,
+            design_spaces,
+            database,
+            cost_model,
+        )
+
+    def post_tuning(self) -> None:
+        """Post-tuning for the search strategy."""
+        _ffi_api.SearchStrategyPostTuning(self)  # type: ignore # pylint: disable=no-member
+
+    def generate_measure_candidates(self) -> list[MeasureCandidate] | None:
+        """Generate measure candidates from design spaces for measurement.
+
+        Returns
+        -------
+        measure_candidates : Optional[List[IRModule]]
+            The measure candidates generated, None if finished.
+        """
+        return _ffi_api.SearchStrategyGenerateMeasureCandidates(self)  # type: ignore # pylint: disable=no-member
+
+    def notify_runner_results(
+        self,
+        measure_candidates: list[MeasureCandidate],
+        results: list[RunnerResult],
+    ) -> None:
+        """Update the search strategy with profiling results.
+
+        Parameters
+        ----------
+        measure_candidates : List[MeasureCandidate]
+            The measure candidates for update.
+        results : List[RunnerResult]
+            The profiling results from the runner.
+        """
+        _ffi_api.SearchStrategyNotifyRunnerResults(  # type: ignore # pylint: disable=no-member
+            self,
+            measure_candidates,
+            results,
+        )
+
+    def clone(self) -> "SearchStrategy":
+        """Clone the search strategy.
+
+        Returns
+        -------
+        cloned : SearchStrategy
+            The cloned search strategy.
+        """
+        return _ffi_api.SearchStrategyClone(self)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def create(  # pylint: disable=keyword-arg-before-vararg
+        kind: Literal[
+            "evolutionary",
+            "replay-trace",
+            "replay-func",
+        ] = "evolutionary",
+        *args,
+        **kwargs,
+    ) -> "SearchStrategy":
+        """Create a search strategy."""
+        from . import (  # pylint: disable=import-outside-toplevel
+            EvolutionarySearch,
+            ReplayFunc,
+            ReplayTrace,
+        )
+
+        if kind == "evolutionary":
+            return EvolutionarySearch(*args, **kwargs)
+        if kind == "replay-trace":
+            return ReplayTrace(*args, **kwargs)
+        if kind == "replay-func":
+            return ReplayFunc(*args, **kwargs)  # type: ignore
+        raise ValueError(f"Unknown SearchStrategy: {kind}")
+
+
+create = SearchStrategy.create  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.PySearchStrategy")
+class _PySearchStrategy(SearchStrategy):
+    """
+    A TVM object search strategy to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PySearchStrategy
+    """
+
+    def __init__(
+        self,
+        f_initialize_with_tune_context: Callable | None = None,
+        f_pre_tuning: Callable | None = None,
+        f_post_tuning: Callable | None = None,
+        f_generate_measure_candidates: Callable | None = None,
+        f_notify_runner_results: Callable | None = None,
+        f_clone: Callable | None = None,
+    ):
+        """Constructor."""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.SearchStrategyPySearchStrategy,  # type: ignore # pylint: disable=no-member
+            f_initialize_with_tune_context,
+            f_pre_tuning,
+            f_post_tuning,
+            f_generate_measure_candidates,
+            f_notify_runner_results,
+            f_clone,
+        )
+
+
+class PySearchStrategy:
+    """
+    An abstract search strategy with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PySearchStrategy,
+        "methods": [
+            "_initialize_with_tune_context",
+            "pre_tuning",
+            "post_tuning",
+            "generate_measure_candidates",
+            "notify_runner_results",
+            "clone",
+        ],
+    }
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        """Initialize the search strategy with tuning context.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for initialization.
+        """
+        raise NotImplementedError
+
+    def pre_tuning(
+        self,
+        max_trials: int,
+        num_trials_per_iter: int,
+        design_spaces: list[Schedule],
+        database: Optional["Database"] = None,
+        cost_model: Optional["CostModel"] = None,
+    ) -> None:
+        """Pre-tuning for the search strategy.
+
+        Parameters
+        ----------
+        design_spaces : List[Schedule]
+            The design spaces for pre-tuning.
+        """
+        raise NotImplementedError
+
+    def post_tuning(self) -> None:
+        """Post-tuning for the search strategy."""
+        raise NotImplementedError
+
+    def generate_measure_candidates(self) -> list[MeasureCandidate] | None:
+        """Generate measure candidates from design spaces for measurement.
+
+        Returns
+        -------
+        measure_candidates : Optional[List[IRModule]]
+            The measure candidates generated, None if finished.
+        """
+        raise NotImplementedError
+
+    def notify_runner_results(
+        self,
+        measure_candidates: list[MeasureCandidate],
+        results: list[RunnerResult],
+    ) -> None:
+        """Update the search strategy with profiling results.
+
+        Parameters
+        ----------
+        measure_candidates : List[MeasureCandidate]
+            The measure candidates for update.
+        results : List[RunnerResult]
+            The profiling results from the runner.
+        """
+        raise NotImplementedError
+
+    def clone(self) -> SearchStrategy:
+        """Clone the search strategy.
+
+        Returns
+        -------
+        strategy : SearchStrategy
+            The cloned search strategy.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/space_generator/__init__.py

@@ -0,0 +1,27 @@
+# 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.
+"""
+The tvm.s_tir.meta_schedule.space_generator package.
+Meta Schedule design space generators that generates design
+space for generation of measure candidates.
+"""
+
+from .post_order_apply import PostOrderApply
+from .schedule_fn import ScheduleFn
+from .space_generator import PySpaceGenerator, ScheduleFnType, SpaceGenerator, create
+from .space_generator_union import SpaceGeneratorUnion

python/tvm/s_tir/meta_schedule/space_generator/post_order_apply.py

@@ -0,0 +1,61 @@
+# 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.
+"""Post Order Apply Space Generator."""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .space_generator import (
+    MutatorProbType,
+    PostprocType,
+    ScheduleRuleType,
+    SpaceGenerator,
+    _normalize_rules,
+)
+
+
+@register_object("s_tir.meta_schedule.PostOrderApply")
+class PostOrderApply(SpaceGenerator):
+    """
+    PostOrderApply is the design space generator that generates design spaces by applying schedule
+    rules to blocks in post-DFS order.
+
+    Parameters
+    ----------
+    f_block_filter : Optional[function]
+        An optional callback function that is used to filter which blocks have schedules generated
+        for them. The function should take in a block and return True if a schedule should
+        be generated or False if that block should be skipped. If no function is provided
+        all blocks will have schedules generated.
+    """
+
+    def __init__(
+        self,
+        f_block_filter=None,
+        sch_rules: ScheduleRuleType = "from-target",
+        postprocs: PostprocType = "from-target",
+        mutator_probs: MutatorProbType = "from-target",
+    ):
+        """Constructor"""
+        sch_rules, postprocs, mutator_probs = _normalize_rules(sch_rules, postprocs, mutator_probs)
+        self.__init_handle_by_constructor__(
+            _ffi_api.SpaceGeneratorPostOrderApply,  # type: ignore # pylint: disable=no-member
+            f_block_filter,
+            sch_rules,
+            postprocs,
+            mutator_probs,
+        )

python/tvm/s_tir/meta_schedule/space_generator/schedule_fn.py

@@ -0,0 +1,64 @@
+# 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.
+"""Union of meta Schedule design space generators."""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .space_generator import (
+    MutatorProbType,
+    PostprocType,
+    ScheduleRuleType,
+    SpaceGenerator,
+    _normalize_rules,
+)
+
+
+@register_object("s_tir.meta_schedule.ScheduleFn")
+class ScheduleFn(SpaceGenerator):
+    """Create a design space generator with customized schedule function.
+    The schedule function can have the following signatures:
+    - 1) [Schedule] -> None
+    - 2) [Schedule] -> Schedule
+    - 3) [Schedule] -> List[Schedule]
+    """
+
+    def __init__(
+        self,
+        sch_fn: SpaceGenerator.ScheduleFnType,
+        sch_rules: ScheduleRuleType = "from-target",
+        postprocs: PostprocType = "from-target",
+        mutator_probs: MutatorProbType = "from-target",
+    ):
+        """Constructor.
+
+        Parameters
+        ----------
+        sch_fn : SpaceGenerator.ScheduleFnType
+            The schedule function, which can have the following signatures:
+            - 1) [Schedule] -> None
+            - 2) [Schedule] -> Schedule
+            - 3) [Schedule] -> List[Schedule]
+        """
+        sch_rules, postprocs, mutator_probs = _normalize_rules(sch_rules, postprocs, mutator_probs)
+        self.__init_handle_by_constructor__(
+            _ffi_api.SpaceGeneratorScheduleFn,  # type: ignore # pylint: disable=no-member
+            sch_fn,
+            sch_rules,
+            postprocs,
+            mutator_probs,
+        )

python/tvm/s_tir/meta_schedule/space_generator/space_generator.py

@@ -0,0 +1,264 @@
+# 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.
+# ruff: noqa: RUF012
+"""
+Meta Schedule design space generators that generates design
+space for generation of measure candidates.
+"""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+from tvm_ffi import register_object
+
+from tvm.ir import IRModule
+from tvm.runtime import Object
+from tvm.s_tir.schedule import Schedule
+
+from .. import _ffi_api
+
+if TYPE_CHECKING:
+    from ..mutator import Mutator
+    from ..postproc import Postproc
+    from ..schedule_rule import ScheduleRule
+    from ..tune_context import TuneContext
+
+
+@register_object("s_tir.meta_schedule.SpaceGenerator")
+class SpaceGenerator(Object):
+    """The abstract design space generator interface."""
+
+    ScheduleFnType = (
+        Callable[[Schedule], None]  # No output
+        | Callable[[Schedule], Schedule]  # Single output
+        | Callable[[Schedule], list[Schedule]]  # Multiple outputs
+    )
+
+    SpaceGeneratorType = Union[
+        "SpaceGenerator",
+        ScheduleFnType,
+        Literal["post-order-apply", "union"],
+    ]
+
+    sch_rules: list["ScheduleRule"] | None
+    postprocs: list["Postproc"] | None
+    mutator_probs: dict["Mutator", float] | None
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        """Initialize the design space generator with tuning context.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for initializing the design space generator.
+        """
+        _ffi_api.SpaceGeneratorInitializeWithTuneContext(  # type: ignore # pylint: disable=no-member
+            self, context
+        )
+
+    def generate_design_space(self, mod: IRModule) -> list[Schedule]:
+        """Generate design spaces given a module.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The module used for design space generation.
+
+        Returns
+        -------
+        design_spaces : List[tvm.s_tir.Schedule]
+            The generated design spaces, i.e., schedules.
+        """
+        return _ffi_api.SpaceGeneratorGenerateDesignSpace(self, mod)  # type: ignore # pylint: disable=no-member
+
+    def clone(self) -> "SpaceGenerator":
+        """Clone the design space generator.
+
+        Returns
+        -------
+        cloned_sg : SpaceGenerator
+            The cloned design space generator.
+        """
+        return _ffi_api.SpaceGeneratorClone(self)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def create(  # pylint: disable=keyword-arg-before-vararg
+        kind: Literal["post-order-apply", "union"] | ScheduleFnType = "post-order-apply",
+        *args,
+        **kwargs,
+    ) -> "SpaceGenerator":
+        """Create a design space generator."""
+        from . import (  # pylint: disable=import-outside-toplevel
+            PostOrderApply,
+            ScheduleFn,
+            SpaceGeneratorUnion,
+        )
+
+        if callable(kind):
+
+            def create_schedule_fn(
+                func,
+                sch_rules=[],
+                postprocs=[],
+                mutator_probs={},
+            ):  # pylint: disable=dangerous-default-value
+                return ScheduleFn(func, sch_rules, postprocs, mutator_probs)
+
+            return create_schedule_fn(kind, *args, **kwargs)  # type: ignore
+        if kind == "post-order-apply":
+            return PostOrderApply(*args, **kwargs)
+        if kind == "union":
+            return SpaceGeneratorUnion(*args, **kwargs)
+        if isinstance(kind, str):
+            return PostOrderApply(sch_rules=kind, postprocs=kind, mutator_probs=kind)
+        raise ValueError(f"Unknown SpaceGenerator: {kind}")
+
+
+ScheduleFnType = SpaceGenerator.ScheduleFnType
+ScheduleRuleType = (
+    list["ScheduleRule"] | Literal["llvm", "cuda", "cuda-tensorcore", "hexagon", "from-target"]
+)
+PostprocType = (
+    list["Postproc"] | Literal["llvm", "cuda", "cuda-tensorcore", "hexagon", "from-target"]
+)
+MutatorProbType = (
+    dict["Mutator", float] | Literal["llvm", "cuda", "cuda-tensorcore", "hexagon", "from-target"]
+)
+create = SpaceGenerator.create  # pylint: disable=invalid-name
+
+
+def _normalize_rules(
+    sch_rules: ScheduleRuleType,
+    postprocs: PostprocType,
+    mutator_probs: MutatorProbType,
+) -> tuple[
+    list["ScheduleRule"] | None,
+    list["Postproc"] | None,
+    dict["Mutator", float] | None,
+]:
+    # pylint: disable=import-outside-toplevel
+    from ..mutator import Mutator
+    from ..postproc import Postproc
+    from ..schedule_rule import ScheduleRule
+
+    # pylint: enable=import-outside-toplevel
+    assert sch_rules is not None
+    assert postprocs is not None
+    assert mutator_probs is not None
+
+    if isinstance(sch_rules, str):
+        if sch_rules == "from-target":
+            sch_rules = None
+        else:
+            sch_rules = ScheduleRule.create(sch_rules)
+    if isinstance(postprocs, str):
+        if postprocs == "from-target":
+            postprocs = None
+        else:
+            postprocs = Postproc.create(postprocs)
+    if isinstance(mutator_probs, str):
+        if mutator_probs == "from-target":
+            mutator_probs = None
+        else:
+            mutator_probs = Mutator.create(mutator_probs)
+    return sch_rules, postprocs, mutator_probs  # type: ignore
+
+
+@register_object("s_tir.meta_schedule.PySpaceGenerator")
+class _PySpaceGenerator(SpaceGenerator):
+    """
+    A TVM object space generator to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PySpaceGenerator
+    """
+
+    def __init__(
+        self,
+        sch_rules: ScheduleRuleType = "from-target",
+        postprocs: PostprocType = "from-target",
+        mutator_probs: MutatorProbType = "from-target",
+        f_initialize_with_tune_context: Callable | None = None,
+        f_generate_design_space: Callable | None = None,
+        f_clone: Callable | None = None,
+    ):
+        """Constructor."""
+        sch_rules, postprocs, mutator_probs = _normalize_rules(sch_rules, postprocs, mutator_probs)
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.SpaceGeneratorPySpaceGenerator,  # type: ignore # pylint: disable=no-member
+            sch_rules,
+            postprocs,
+            mutator_probs,
+            f_initialize_with_tune_context,
+            f_generate_design_space,
+            f_clone,
+        )
+
+
+class PySpaceGenerator:
+    """
+    An abstract space generator with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PySpaceGenerator,
+        "fields": ["sch_rules", "postprocs", "mutator_probs"],
+        "methods": ["_initialize_with_tune_context", "generate_design_space", "clone"],
+    }
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        """Initialize the design space generator with tuning context.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for initializing the design space generator.
+        """
+        raise NotImplementedError
+
+    def generate_design_space(self, mod: IRModule) -> list[Schedule]:
+        """Generate design spaces given a module.
+
+        Parameters
+        ----------
+        mod : IRModule
+            The module used for design space generation.
+
+        Returns
+        -------
+        design_spaces : List[tvm.s_tir.Schedule]
+            The generated design spaces, i.e., schedules.
+        """
+        raise NotImplementedError
+
+    def clone(self) -> SpaceGenerator:
+        """Clone the design space generator.
+
+        Returns
+        -------
+        cloned_sg : SpaceGenerator
+            The cloned design space generator.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/space_generator/space_generator_union.py

@@ -0,0 +1,56 @@
+# 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.
+"""Union of meta Schedule design space generators."""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .space_generator import (
+    MutatorProbType,
+    PostprocType,
+    ScheduleRuleType,
+    SpaceGenerator,
+    _normalize_rules,
+)
+
+
+@register_object("s_tir.meta_schedule.SpaceGeneratorUnion")
+class SpaceGeneratorUnion(SpaceGenerator):
+    """Union of design space generators."""
+
+    def __init__(
+        self,
+        space_generators: list[SpaceGenerator],
+        sch_rules: ScheduleRuleType = "from-target",
+        postprocs: PostprocType = "from-target",
+        mutator_probs: MutatorProbType = "from-target",
+    ):
+        """Constructor.
+
+        Parameters
+        ----------
+        space_generators : List[SpaceGenerator]
+            The list of design space generators to be unioned.
+        """
+        sch_rules, postprocs, mutator_probs = _normalize_rules(sch_rules, postprocs, mutator_probs)
+        self.__init_handle_by_constructor__(
+            _ffi_api.SpaceGeneratorSpaceGeneratorUnion,  # type: ignore # pylint: disable=no-member
+            space_generators,
+            sch_rules,
+            postprocs,
+            mutator_probs,
+        )

python/tvm/s_tir/meta_schedule/task_scheduler/__init__.py

@@ -0,0 +1,27 @@
+# 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.
+"""
+The tvm.s_tir.meta_schedule.task_scheduler package.
+Meta Schedule task scheduler that manage the task scheduling
+for measure candidates generation and measurement, then save
+records to the database.
+"""
+
+from .gradient_based import GradientBased
+from .round_robin import RoundRobin
+from .task_scheduler import PyTaskScheduler, TaskScheduler, create

python/tvm/s_tir/meta_schedule/task_scheduler/gradient_based.py

@@ -0,0 +1,56 @@
+# 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.
+"""Gradient Based Task Scheduler"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from ..logging import get_logger, get_logging_func
+from .task_scheduler import TaskScheduler
+
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.GradientBased")
+class GradientBased(TaskScheduler):
+    """Gradient Based Task Scheduler"""
+
+    def __init__(
+        self,
+        *,
+        alpha: float = 0.2,
+        window_size: int = 3,
+        seed: int = -1,
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        alpha : float = 0.2
+            The parameter alpha in gradient computation.
+        window_size : int = 3
+            The parameter to control backward window size in gradient computation.
+        seed : int = -1
+            The random seed.
+        """
+        self.__init_handle_by_constructor__(
+            _ffi_api.TaskSchedulerGradientBased,  # type: ignore # pylint: disable=no-member
+            get_logging_func(logger),
+            alpha,
+            window_size,
+            seed,
+        )

python/tvm/s_tir/meta_schedule/task_scheduler/round_robin.py

@@ -0,0 +1,37 @@
+# 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.
+"""Round Robin Task Scheduler"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from ..logging import get_logger, get_logging_func
+from .task_scheduler import TaskScheduler
+
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.RoundRobin")
+class RoundRobin(TaskScheduler):
+    """Round Robin Task Scheduler"""
+
+    def __init__(self) -> None:
+        """Constructor."""
+        self.__init_handle_by_constructor__(
+            _ffi_api.TaskSchedulerRoundRobin,  # type: ignore # pylint: disable=no-member
+            get_logging_func(logger),
+        )

python/tvm/s_tir/meta_schedule/task_scheduler/task_scheduler.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.
+# ruff: noqa: RUF012
+"""Auto-tuning Task Scheduler"""
+
+from collections.abc import Callable
+from typing import Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+
+from .. import _ffi_api
+from ..builder import Builder, BuilderResult
+from ..cost_model import CostModel
+from ..database import Database
+from ..logging import get_logger, get_logging_func
+from ..measure_callback import MeasureCallback
+from ..runner import Runner, RunnerResult
+from ..search_strategy import MeasureCandidate
+from ..tune_context import TuneContext
+
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.TaskRecord")
+class TaskRecord(Object):
+    """The running record of a task."""
+
+    ctx: TuneContext
+    task_weight: float
+    flop: float
+    is_terminated: bool
+    build_error_count: int
+    run_error_count: int
+    measure_candidates: list[MeasureCandidate]
+    builder_results: list[BuilderResult]
+    runner_results: list[RunnerResult]
+
+
+@register_object("s_tir.meta_schedule.TaskScheduler")
+class TaskScheduler(Object):
+    """The abstract task scheduler interface."""
+
+    tasks_: list[TaskRecord]
+    measure_callbacks_: list[MeasureCallback]
+    database_: Database | None
+    cost_model_: CostModel | None
+    remaining_tasks_: int
+
+    TaskSchedulerType = Union["TaskScheduler", Literal["gradient", "round-robin"]]
+
+    def next_task_id(self) -> int:
+        """Fetch the next task id.
+
+        Returns
+        -------
+        next_task_id : int
+            The next task id.
+        """
+        return _ffi_api.TaskSchedulerNextTaskId(self)  # type: ignore # pylint: disable=no-member
+
+    def join_running_task(self, task_id: int) -> list[RunnerResult]:
+        """Wait until the task is finished.
+
+        Parameters
+        ----------
+        task_id : int
+            The task id to be joined.
+
+        Returns
+        -------
+        results : List[RunnerResult]
+            The list of results.
+        """
+        return _ffi_api.TaskSchedulerJoinRunningTask(self, task_id)  # type: ignore # pylint: disable=no-member
+
+    def tune(
+        self,
+        tasks: list[TuneContext],
+        task_weights: list[float],
+        max_trials_global: int,
+        max_trials_per_task: int,
+        num_trials_per_iter: int,
+        builder: Builder,
+        runner: Runner,
+        measure_callbacks: list[MeasureCallback],
+        database: Database | None,
+        cost_model: CostModel | None,
+    ) -> None:
+        """Auto-tuning.
+
+        Parameters
+        ----------
+        tasks : List[TuneContext]
+            The list of tuning contexts as tasks.
+        task_weights : List[float]
+            The list of task weights.
+        max_trials_global : int
+            The maximum number of trials globally.
+        max_trials_per_task : int
+            The maximum number of trials per task.
+        num_trials_per_iter : int
+            The number of trials per iteration.
+        builder : Builder
+            The builder.
+        runner : Runner
+            The runner.
+        measure_callbacks : List[MeasureCallback]
+            The list of measure callbacks.
+        database : Optional[Database]
+            The database.
+        cost_model : Optional[CostModel]
+            The cost model.
+        """
+        task_weights = [float(w) for w in task_weights]
+        _ffi_api.TaskSchedulerTune(  # type: ignore # pylint: disable=no-member
+            self,
+            tasks,
+            task_weights,
+            max_trials_global,
+            max_trials_per_task,
+            num_trials_per_iter,
+            builder,
+            runner,
+            measure_callbacks,
+            database,
+            cost_model,
+        )
+
+    def terminate_task(self, task_id: int) -> None:
+        """Terminate the task
+
+        Parameters
+        ----------
+        task_id : int
+            The task id to be terminated.
+        """
+        _ffi_api.TaskSchedulerTerminateTask(self, task_id)  # type: ignore # pylint: disable=no-member
+
+    def touch_task(self, task_id: int) -> None:
+        """Touch the task and update its status
+
+        Parameters
+        ----------
+        task_id : int
+            The task id to be checked.
+        """
+        _ffi_api.TaskSchedulerTouchTask(self, task_id)  # type: ignore # pylint: disable=no-member
+
+    def print_tuning_statistics(self) -> None:
+        """Print out a human-readable format of the tuning statistics."""
+        return _ffi_api.TaskSchedulerPrintTuningStatistics(self)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def create(  # pylint: disable=keyword-arg-before-vararg
+        kind: Literal["round-robin", "gradient"] = "gradient",
+        *args,
+        **kwargs,
+    ) -> "TaskScheduler":
+        """Create a task scheduler."""
+        from . import (  # pylint: disable=import-outside-toplevel
+            GradientBased,
+            RoundRobin,
+        )
+
+        if kind == "round-robin":
+            return RoundRobin(*args, **kwargs)  # type: ignore
+        if kind == "gradient":
+            return GradientBased(*args, **kwargs)
+        raise ValueError(f"Unknown TaskScheduler name: {kind}")
+
+
+create = TaskScheduler.create  # pylint: disable=invalid-name
+
+
+@register_object("s_tir.meta_schedule.PyTaskScheduler")
+class _PyTaskScheduler(TaskScheduler):
+    """
+    A TVM object task scheduler to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyTaskScheduler
+    """
+
+    def __init__(
+        self,
+        f_next_task_id: Callable,
+        f_join_running_task: Callable,
+        f_tune: Callable,
+    ):
+        """Constructor."""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.TaskSchedulerPyTaskScheduler,  # type: ignore # pylint: disable=no-member
+            get_logging_func(logger),
+            f_next_task_id,
+            f_join_running_task,
+            f_tune,
+        )
+
+
+class PyTaskScheduler:
+    """
+    An abstract task scheduler with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PyTaskScheduler,
+        "fields": [],
+        "methods": ["next_task_id", "join_running_task", "tune"],
+    }
+
+    def __init__(self): ...
+
+    def tune(
+        self,
+        tasks: list[TuneContext],
+        task_weights: list[float],
+        max_trials_global: int,
+        max_trials_per_task: int,
+        builder: Builder,
+        runner: Runner,
+        measure_callbacks: list[MeasureCallback],
+        database: Database | None,
+        cost_model: CostModel | None,
+    ) -> None:
+        """Auto-tuning."""
+        # Using self._outer to replace the self pointer
+        _ffi_api.TaskSchedulerTune(  # type: ignore # pylint: disable=no-member
+            self._outer(),  # type: ignore # pylint: disable=no-member
+            tasks,
+            task_weights,
+            max_trials_global,
+            max_trials_per_task,
+            builder,
+            runner,
+            measure_callbacks,
+            database,
+            cost_model,
+        )
+
+    def next_task_id(self) -> int:
+        """Fetch the next task id.
+
+        Returns
+        -------
+        next_task_id : int
+            The next task id.
+        """
+        raise NotImplementedError
+
+    def join_running_task(self, task_id: int) -> list[RunnerResult]:
+        """Wait until the task is finished.
+
+        Parameters
+        ----------
+        task_id : int
+            The task id to be joined.
+        """
+        # Using self._outer to replace the self pointer
+        return _ffi_api.TaskSchedulerJoinRunningTask(self._outer(), task_id)  # type: ignore # pylint: disable=no-member

python/tvm/s_tir/meta_schedule/testing/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+"""Testing utilities in meta schedule"""
+
+# NOTE: Do not import any module here by default

python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py

@@ -0,0 +1,54 @@
+# 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.
+"""Customized builder and runner methods"""
+# pylint: disable=import-outside-toplevel
+
+from collections.abc import Callable
+
+import numpy as np  # type: ignore
+
+from tvm.runtime import Executable, Module
+from tvm.s_tir.meta_schedule.runner import RPCConfig
+
+
+def run_module_via_rpc(
+    rpc_config: RPCConfig,
+    lib: Module | Executable,
+    dev_type: str,
+    args: dict[int, np.ndarray] | dict[str, np.ndarray],
+    continuation: Callable,
+):
+    """Execute a tvm.runtime.Module on RPC remote"""
+    # pylint: disable=import-outside-toplevel
+    import os
+    import tempfile
+
+    from tvm.runtime import ndarray
+    from tvm.support.tar import tar
+
+    # pylint: enable=import-outside-toplevel
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        filename = os.path.join(tmp_dir, "tvm_tmp_mod." + tar.output_format)
+        lib.export_library(filename, fcompile=tar)
+        session = rpc_config.connect_server()
+        session.upload(filename)
+        _, filename = os.path.split(filename)
+        rt_mod = session.load_module(filename)
+        dev = session.device(dev_type, 0)
+        nd_args = {k: ndarray.array(v, dev) for k, v in args.items()}
+        return continuation(rt_mod, dev, nd_args)

python/tvm/s_tir/meta_schedule/testing/distributed_measure_candidates.py

@@ -0,0 +1,199 @@
+# 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=missing-docstring
+
+import argparse
+import glob
+import os
+
+from tqdm import tqdm  # type: ignore
+
+from tvm.s_tir import meta_schedule as ms
+from tvm.target import Target
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--candidate_cache_dir", type=str, help="Please provide the full path to the candidates."
+    )
+    parser.add_argument(
+        "--result_cache_dir", type=str, help="Please provide the full path to the result database."
+    )
+    parser.add_argument(
+        "--target",
+        type=str,
+        default="nvidia/nvidia-v100",
+        help="Please specify the target hardware for tuning context.",
+    )
+    parser.add_argument(
+        "--rpc_host", type=str, help="Please provide the private IPv4 address for the tracker."
+    )
+    parser.add_argument(
+        "--rpc_port", type=int, default=4445, help="Please provide the port for the tracker."
+    )
+    parser.add_argument(
+        "--rpc_key",
+        type=str,
+        default="p3.2xlarge",
+        help="Please provide the key for the rpc servers.",
+    )
+    parser.add_argument(
+        "--builder_timeout_sec",
+        type=int,
+        default=10,
+        help="The time for the builder session to time out.",
+    )
+    parser.add_argument(
+        "--min_repeat_ms", type=int, default=100, help="The time for preheating the gpu."
+    )
+    parser.add_argument(
+        "--runner_timeout_sec",
+        type=int,
+        default=100,
+        help="The time for the runner session to time out.",
+    )
+    parser.add_argument(
+        "--cpu_flush", type=bool, default=False, help="Whether to enable cpu cache flush or not."
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=128,
+        help="The batch size of candidates sent to builder and runner each time.",
+    )
+    return parser.parse_args()
+
+
+# pylint: disable=too-many-locals
+def measure_candidates(database, builder, runner):
+    """Send the candidates to builder and runner for distributed measurement,
+    and save the results in a new json database.
+
+    Parameters
+    ----------
+    database : JSONDatabase
+        The database for candidates to be measured.
+    builder : Builder
+        The builder for building the candidates.
+    runner : Runner
+        The runner for measuring the candidates.
+
+    Returns
+    -------
+    None
+    """
+    candidates, runner_results, build_fail_indices, run_fail_indices = [], [], [], []
+    context = ms.TuneContext(target=Target(args.target))
+    tuning_records = database.get_all_tuning_records()
+    for record in tuning_records:
+        candidates.append(record.as_measure_candidate())
+    with ms.Profiler() as profiler:
+        for idx in range(0, len(candidates), args.batch_size):
+            batch_candidates = candidates[idx : idx + args.batch_size]
+            context._set_measure_candidates(batch_candidates)  # pylint: disable=protected-access
+            with ms.Profiler.timeit("build"):
+                context._send_to_builder(builder)  # pylint: disable=protected-access
+            with ms.Profiler.timeit("run"):
+                context._send_to_runner(runner)  # pylint: disable=protected-access
+                batch_runner_results = context._join()  # pylint: disable=protected-access
+            runner_results.extend(batch_runner_results)
+            for i, result in enumerate(context.builder_results):
+                if result.error_msg is None:
+                    ms.utils.remove_build_dir(result.artifact_path)
+                else:
+                    build_fail_indices.append(i + idx)
+            context._clear_measure_state()  # pylint: disable=protected-access
+
+    model_name, workload_name = database.path_workload.split("/")[-2:]
+    record_name = database.path_tuning_record.split("/")[-1]
+    new_database = ms.database.JSONDatabase(
+        path_workload=os.path.join(args.result_cache_dir, model_name, workload_name),
+        path_tuning_record=os.path.join(args.result_cache_dir, model_name, record_name),
+    )
+    workload = tuning_records[0].workload
+    new_database.commit_workload(workload.mod)
+    for i, (record, result) in enumerate(zip(tuning_records, runner_results)):
+        if result.error_msg is None:
+            new_database.commit_tuning_record(
+                ms.database.TuningRecord(
+                    trace=record.trace,
+                    workload=workload,
+                    run_secs=[v.value for v in result.run_secs],
+                    target=Target(args.target),
+                )
+            )
+        else:
+            run_fail_indices.append(i)
+    fail_indices_name = workload_name.replace("_workload.json", "_failed_indices.txt")
+    with open(
+        os.path.join(args.result_cache_dir, model_name, fail_indices_name), "w", encoding="utf8"
+    ) as file:
+        file.write(" ".join([str(n) for n in run_fail_indices]))
+    print(
+        f"Builder time: {profiler.get()['build']}, Runner time: {profiler.get()['run']}\n\
+            Failed number of builds: {len(build_fail_indices)},\
+            Failed number of runs: {len(run_fail_indices)}"
+    )
+
+
+args = _parse_args()  # pylint: disable=invalid-name
+
+
+def main():
+    builder = ms.builder.LocalBuilder(timeout_sec=args.builder_timeout_sec)
+    runner = ms.runner.RPCRunner(
+        rpc_config=ms.runner.RPCConfig(
+            tracker_host=args.rpc_host,
+            tracker_port=args.rpc_port,
+            tracker_key=args.rpc_key,
+            session_timeout_sec=args.runner_timeout_sec,
+        ),
+        evaluator_config=ms.runner.EvaluatorConfig(
+            number=3,
+            repeat=1,
+            min_repeat_ms=args.min_repeat_ms,
+            enable_cpu_cache_flush=args.cpu_flush,
+        ),
+        max_workers=os.cpu_count(),
+    )
+    if not os.path.isdir(args.candidate_cache_dir):
+        raise Exception("Please provide a correct candidate cache dir.")
+    try:
+        os.makedirs(args.result_cache_dir, exist_ok=True)
+    except OSError:
+        print(f"Directory {args.result_cache_dir} cannot be created successfully.")
+    model_dirs = glob.glob(os.path.join(args.candidate_cache_dir, "*"))
+    for model_dir in model_dirs:
+        model_name = model_dir.split("/")[-1]
+        os.makedirs(os.path.join(args.result_cache_dir, model_name), exist_ok=True)
+        all_tasks = glob.glob(os.path.join(model_dir, "*.json"))
+        workload_paths = []
+        for path in all_tasks:
+            if path.endswith("_workload.json"):
+                workload_paths.append(path)
+        for workload_path in tqdm(workload_paths):
+            candidate_path = workload_path.replace("_workload.json", "_candidates.json")
+            database = ms.database.JSONDatabase(
+                path_workload=workload_path,
+                path_tuning_record=candidate_path,
+            )
+            measure_candidates(database, builder, runner)
+
+
+if __name__ == "__main__":
+    main()

python/tvm/s_tir/meta_schedule/testing/dummy_object.py

@@ -0,0 +1,63 @@
+# 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.
+"""Dummy objects for testing."""
+
+import random
+
+from tvm.ir.utils import derived_object
+from tvm.s_tir.schedule import Trace
+
+from ..builder import BuilderInput, BuilderResult, PyBuilder
+from ..mutator import PyMutator
+from ..runner import PyRunner, PyRunnerFuture, RunnerFuture, RunnerInput, RunnerResult
+from ..tune_context import TuneContext  # pylint: disable=unused-import
+
+
+@derived_object
+class DummyRunnerFuture(PyRunnerFuture):
+    def done(self) -> bool:
+        return True
+
+    def result(self) -> RunnerResult:
+        run_secs = [random.uniform(5, 30) for _ in range(random.randint(1, 10))]
+        return RunnerResult(run_secs, None)
+
+
+@derived_object
+class DummyBuilder(PyBuilder):
+    def build(self, build_inputs: list[BuilderInput]) -> list[BuilderResult]:
+        return [BuilderResult("test_path", None) for _ in build_inputs]
+
+
+@derived_object
+class DummyRunner(PyRunner):
+    def run(self, runner_inputs: list[RunnerInput]) -> list[RunnerFuture]:
+        return [DummyRunnerFuture() for _ in runner_inputs]  # type: ignore
+
+
+@derived_object
+class DummyMutator(PyMutator):
+    """Dummy Mutator for testing"""
+
+    def _initialize_with_tune_context(self, context: "TuneContext") -> None:
+        pass
+
+    def apply(self, trace: Trace, _) -> Trace | None:
+        return Trace(trace.insts, {})
+
+    def clone(self):
+        return DummyMutator()

python/tvm/s_tir/meta_schedule/testing/local_rpc.py

@@ -0,0 +1,72 @@
+# 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.
+"""RPC tracker and server running locally"""
+
+from tvm.rpc.server import Server
+from tvm.rpc.tracker import Tracker
+
+
+class LocalRPC:
+    """A pair of RPC tracker/server running locally
+
+    Parameters
+    ----------
+    tracker_host : str
+        The host URL of the tracker
+    tracker_port : int
+        The port of the tracker
+    tracker_key: str
+        The key used in the tracker to refer to a worker
+    """
+
+    tracker_host: str
+    tracker_port: int
+    tracker_key: str
+
+    def __init__(
+        self,
+        tracker_key: str = "key",
+        silent: bool = False,
+        no_fork: bool = False,
+    ) -> None:
+        self.tracker = Tracker(
+            silent=silent,
+            port=9190,
+            port_end=12345,
+        )
+        self.server = Server(
+            host="0.0.0.0",
+            is_proxy=False,
+            tracker_addr=(self.tracker.host, self.tracker.port),
+            key=tracker_key,
+            silent=silent,
+            no_fork=no_fork,
+            port=9190,
+            port_end=12345,
+        )
+        self.tracker_host = self.tracker.host
+        self.tracker_port = self.tracker.port
+        self.tracker_key = tracker_key
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, _type, _value, _traceback):
+        if hasattr(self, "server"):
+            del self.server
+        if hasattr(self, "tracker"):
+            del self.tracker

python/tvm/s_tir/meta_schedule/testing/space_generation.py

@@ -0,0 +1,150 @@
+# 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=missing-module-docstring,missing-function-docstring,missing-class-docstring
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+import tvm_ffi
+
+from tvm.ir import IRModule
+from tvm.s_tir import Schedule
+from tvm.s_tir import meta_schedule as ms
+from tvm.s_tir.schedule import Trace
+from tvm.s_tir.schedule.testing import verify_trace_roundtrip
+from tvm.target import Target
+
+
+def get_rules(
+    kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"],
+    types: type | tuple[type, ...],
+) -> list[ms.ScheduleRule]:
+    """Get default schedule rules"""
+    rules = ms.ScheduleRule.create(kind)
+    return [rule for rule in rules if isinstance(rule, types)]
+
+
+def structural_equal_no_gs(mod1: IRModule, mod2: IRModule) -> bool:
+    """
+    Checks structural equality but ignores global symbols
+    """
+
+    # for every function in the modules, remove global symbols from the attrs and then compare
+    def remove_global_symbols(mod: IRModule) -> IRModule:
+        stripped_mod = IRModule()
+        for global_var in mod.get_global_vars():
+            func = mod[global_var]
+            stripped_mod[global_var] = func.without_attr("global_symbol")
+        return stripped_mod
+
+    return tvm_ffi.structural_equal(remove_global_symbols(mod1), remove_global_symbols(mod2))
+
+
+def generate_design_space(
+    kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"],
+    mod: IRModule,
+    target: Target,
+    types: type | tuple[type, ...],
+    sch_rules: list[ms.ScheduleRule] | None = None,
+) -> list[Schedule]:
+    if sch_rules is None:
+        sch_rules = get_rules(kind, types)
+    else:
+        assert types is None
+    return ms.TuneContext(
+        mod=mod,
+        target=target,
+        space_generator=ms.space_generator.PostOrderApply(
+            sch_rules=sch_rules,
+            postprocs=[],
+            mutator_probs={},
+        ),
+        task_name="test",
+    ).generate_design_space()
+
+
+def _find_match_sketch_id(
+    mod: IRModule,
+    sketches: list[Schedule],
+    expected_mod: IRModule,
+    expected_decision: list[tuple[str, list[int]]],
+    *,
+    debug_mask="all",
+) -> int | None:
+    for sketch_id, sketch in enumerate(sketches):
+        i = 0
+        new_decisions = {}
+        for inst in sketch.trace.insts:
+            if not inst.kind.name.startswith("Sample"):
+                continue
+            assert i < len(expected_decision)
+            if inst.kind.name == expected_decision[i][0]:
+                new_decisions[inst] = expected_decision[i][1]
+                i += 1
+        if len(new_decisions) != len(expected_decision):
+            continue
+        sch = Schedule(mod, debug_mask=debug_mask)
+        Trace(
+            insts=sketch.trace.insts,
+            decisions=new_decisions,
+        ).apply_to_schedule(sch, remove_postproc=True)
+        if structural_equal_no_gs(sch.mod, expected_mod):
+            verify_trace_roundtrip(sch=sch, mod=mod, debug_mask=debug_mask, text_format="json")
+            return sketch_id
+    return None
+
+
+def check_sketches(
+    mod: IRModule,
+    sketches: list[Schedule],
+    expected_mods: list[IRModule],
+    expected_decisions: list[list[tuple[str, list[int]]]],
+    *,
+    debug_mask="all",
+):
+    assert len(expected_mods) == len(expected_decisions)
+    assert len(sketches) == len(expected_mods)
+    expected_mods = [
+        IRModule({"main": m}) if not isinstance(m, IRModule) else m for m in expected_mods
+    ]
+    sketches = list(sketches)
+    for expected_id, (expected_mod, expected_decision) in enumerate(
+        zip(expected_mods, expected_decisions)
+    ):
+        sketch_id = _find_match_sketch_id(
+            mod,
+            sketches,
+            expected_mod,
+            expected_decision,
+            debug_mask=debug_mask,
+        )
+        if sketch_id is None:
+            raise AssertionError(
+                f"Expected sketch #{expected_id} doesn't exist in the generated sketches."
+            )
+        sketches.pop(sketch_id)
+
+
+def print_sketches(sketches: list[Schedule]):
+    for i, sch in enumerate(sketches):
+        print(f"###### {i}")
+        sch.mod.show(black_format=False)
+        for inst in sch.trace.insts:
+            if inst in sch.trace.decisions:
+                print(f'("{inst.kind.name}", {sch.trace.decisions[inst]}),')

python/tvm/s_tir/meta_schedule/testing/te_workload.py

@@ -0,0 +1,837 @@
+# 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.
+# ruff: noqa: E741
+"""Workloads in TE"""
+
+# pylint: disable=missing-docstring
+
+from tvm import te, tirx, topi
+from tvm.target import Target
+
+
+def batch_matmul_nkkm(  # pylint: disable=invalid-name,missing-docstring
+    B: int,
+    N: int,
+    M: int,
+    K: int,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    x = te.placeholder((B, N, K), name="X", dtype=in_dtype)
+    y = te.placeholder((B, K, M), name="Y", dtype=in_dtype)
+    k = te.reduce_axis((0, K), name="k")
+    z = te.compute(  # pylint: disable=invalid-name
+        (B, N, M),
+        lambda b, i, j: te.sum(
+            x[b][i][k].astype(out_dtype) * y[b][k][j].astype(out_dtype),
+            axis=[k],
+        ),
+        name="Z",
+    )
+    return (x, y, z)
+
+
+def conv1d_nlc(  # pylint: disable=invalid-name,missing-docstring
+    N: int,
+    L: int,
+    CI: int,
+    CO: int,
+    kernel_size: int,
+    stride: int = 1,
+    padding: int = 0,
+    dilation: int = 1,
+    groups: int = 1,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    inputs = te.placeholder((N, L, CI), name="inputs", dtype=in_dtype)
+    weight = te.placeholder((kernel_size, CI // groups, CO), name="weight", dtype=in_dtype)
+
+    batch_size, in_len, _ = inputs.shape
+    k_len, channel_per_group, out_channel = weight.shape
+    out_channel_per_group = out_channel // groups
+    out_len = (in_len + 2 * padding - dilation * (k_len - 1) - 1) // stride + 1
+    rc = te.reduce_axis((0, channel_per_group), name="rc")
+    rl = te.reduce_axis((0, k_len), name="rl")
+
+    padded = topi.nn.pad(inputs, [0, padding, 0])
+    output = te.compute(
+        (batch_size, out_len, out_channel),
+        lambda n, l, co: te.sum(
+            (
+                padded[
+                    n,
+                    l * stride + rl * dilation,
+                    co // out_channel_per_group * channel_per_group + rc,
+                ].astype(out_dtype)
+                * weight[rl, rc, co].astype(out_dtype)
+            ),
+            axis=[rl, rc],
+        ),
+        name="conv1d_nlc",
+    )
+    return (inputs, weight, output)
+
+
+def conv2d_nhwc(  # pylint: disable=invalid-name,missing-docstring
+    N: int,
+    H: int,
+    W: int,
+    CI: int,
+    CO: int,
+    kernel_size: int,
+    stride: int = 1,
+    padding: int = 0,
+    dilation: int = 1,
+    groups: int = 1,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    inputs = te.placeholder((N, H, W, CI), name="inputs", dtype=in_dtype)
+    weight = te.placeholder(
+        (kernel_size, kernel_size, CI // groups, CO), name="weight", dtype=in_dtype
+    )
+    batch_size, in_h, in_w, _ = inputs.shape
+    k_h, k_w, channel_per_group, out_channel = weight.shape
+    out_channel_per_group = out_channel // groups
+
+    out_h = (in_h + 2 * padding - dilation * (k_h - 1) - 1) // stride + 1
+    out_w = (in_w + 2 * padding - dilation * (k_w - 1) - 1) // stride + 1
+    rh = te.reduce_axis((0, k_h), name="rh")
+    rw = te.reduce_axis((0, k_w), name="rw")
+    rc = te.reduce_axis((0, channel_per_group), name="rc")
+
+    padded = topi.nn.pad(inputs, [0, padding, padding, 0])
+    output = te.compute(
+        (batch_size, out_h, out_w, out_channel),
+        lambda n, h, w, co: te.sum(
+            (
+                padded[
+                    n,
+                    h * stride + rh * dilation,
+                    w * stride + rw * dilation,
+                    co // out_channel_per_group * channel_per_group + rc,
+                ].astype(out_dtype)
+                * weight[rh, rw, rc, co].astype(out_dtype)
+            ),
+            axis=[rh, rw, rc],
+        ),
+        name="conv2d_nhwc",
+    )
+    return (inputs, weight, output)
+
+
+def conv3d_ndhwc(  # pylint: disable=invalid-name,missing-docstring
+    N: int,
+    D: int,
+    H: int,
+    W: int,
+    CI: int,
+    CO: int,
+    kernel_size: int,
+    stride: int = 1,
+    padding: int = 0,
+    dilation: int = 1,
+    groups: int = 1,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    inputs = te.placeholder((N, D, H, W, CI), name="inputs", dtype=in_dtype)
+    weight = te.placeholder(
+        (kernel_size, kernel_size, kernel_size, CI // groups, CO), name="weight", dtype=in_dtype
+    )
+    batch_size, in_d, in_h, in_w, _ = inputs.shape
+    k_d, k_h, k_w, channel_per_group, out_channel = weight.shape
+    out_channel_per_group = out_channel // groups
+
+    out_d = (in_d + 2 * padding - dilation * (k_d - 1) - 1) // stride + 1
+    out_h = (in_h + 2 * padding - dilation * (k_h - 1) - 1) // stride + 1
+    out_w = (in_w + 2 * padding - dilation * (k_w - 1) - 1) // stride + 1
+    rd = te.reduce_axis((0, k_d), name="rd")
+    rh = te.reduce_axis((0, k_h), name="rh")
+    rw = te.reduce_axis((0, k_w), name="rw")
+    rc = te.reduce_axis((0, channel_per_group), name="rc")
+
+    padded = topi.nn.pad(inputs, [0, padding, padding, padding, 0])
+    output = te.compute(
+        (batch_size, out_d, out_h, out_w, out_channel),
+        lambda n, d, h, w, co: te.sum(
+            (
+                padded[
+                    n,
+                    d * stride + rd * dilation,
+                    h * stride + rh * dilation,
+                    w * stride + rw * dilation,
+                    co // out_channel_per_group * channel_per_group + rc,
+                ].astype(out_dtype)
+                * weight[rd, rh, rw, rc, co].astype(out_dtype)
+            ),
+            axis=[rd, rh, rw, rc],
+        ),
+        name="conv3d_ndhwc",
+    )
+    return (inputs, weight, output)
+
+
+def depthwise_conv2d_nhwc(  # pylint: disable=invalid-name,missing-docstring
+    N: int,
+    H: int,
+    W: int,
+    C: int,
+    kernel_size: int,
+    stride: int = 1,
+    padding: int = 0,
+    dilation: int = 1,
+    factor: int = 1,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    inputs = te.placeholder((N, H, W, C), dtype=in_dtype)
+    weight = te.placeholder((factor, kernel_size, kernel_size, C), dtype=in_dtype)
+    batch_size, in_h, in_w, in_channel = inputs.shape
+    factor, k_h, k_w, in_channel = weight.shape
+    out_channel = in_channel * factor
+    assert int(factor) == 1, "Not optimized for factor != 1"
+    out_h = (in_h + 2 * padding - dilation * (k_h - 1) - 1) // stride + 1
+    out_w = (in_w + 2 * padding - dilation * (k_w - 1) - 1) // stride + 1
+    rh = te.reduce_axis((0, k_h), name="rh")
+    rw = te.reduce_axis((0, k_w), name="rw")
+    padded = topi.nn.pad(inputs, [0, padding, padding, 0])
+    output = te.compute(
+        (batch_size, out_h, out_w, out_channel),
+        lambda n, h, w, c: te.sum(
+            (
+                padded[
+                    n,
+                    h * stride + rh * dilation,
+                    w * stride + rw * dilation,
+                    c // factor,
+                ].astype(out_dtype)
+                * weight[c % factor, rh, rw, c // factor].astype(out_dtype)
+            ),
+            axis=[rh, rw],
+        ),
+        name="depth_conv2d_nhwc",
+    )
+    return (inputs, weight, output)
+
+
+def conv2d_transpose_nhwc(  # pylint: disable=invalid-name,missing-docstring
+    N: int,
+    H: int,
+    W: int,
+    CI: int,
+    CO: int,
+    kernel_size: int,
+    stride: int = 1,
+    padding: int = 0,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    inputs = te.placeholder((N, H, W, CI), name="inputs", dtype=in_dtype)
+    weight = te.placeholder((kernel_size, kernel_size, CI, CO), name="weight", dtype=in_dtype)
+
+    batch, in_h, in_w, in_c = inputs.shape
+    filter_h, filter_w, in_c, out_c = weight.shape
+    stride_h, stride_w = (stride, stride)
+
+    # compute padding
+    fpad_top, fpad_left, fpad_bottom, fpad_right = topi.nn.get_pad_tuple(
+        padding, (filter_h, filter_w)
+    )
+    bpad_top = filter_h - 1 - fpad_top
+    bpad_bottom = filter_h - 1 - fpad_bottom
+    bpad_left = filter_w - 1 - fpad_left
+    bpad_right = filter_w - 1 - fpad_right
+
+    # padding stage
+    padded = topi.nn.pad(
+        inputs,
+        [
+            0,
+            (bpad_top + stride_h - 1) // stride_h,
+            (bpad_left + stride_w - 1) // stride_w,
+            0,
+        ],
+        [
+            0,
+            (bpad_bottom + stride_h - 1) // stride_h,
+            (bpad_right + stride_w - 1) // stride_w,
+            0,
+        ],
+    )
+
+    # remove extra padding introduced by dilatation
+    idx_div = te.indexdiv
+    idx_mod = te.indexmod
+    border_h = idx_mod(stride_h - idx_mod(bpad_top, stride_h), stride_h)
+    border_w = idx_mod(stride_w - idx_mod(bpad_left, stride_w), stride_w)
+
+    # dilation stage
+    strides = [1, stride_h, stride_w, 1]
+    n = len(padded.shape)
+
+    # We should embed this dilation directly into te.compute rather than creating a new te.compute.
+    # Only in this way can we use unroll to eliminate the multiplication of zeros.
+    def _dilate(*indices):
+        not_zero = []
+        index_tuple = []
+        for i in range(n):
+            if not strides[i] == 1:
+                index_tuple.append(idx_div(indices[i], strides[i]))
+                not_zero.append(idx_mod(indices[i], strides[i]).equal(0))
+            else:
+                index_tuple.append(indices[i])
+        if not_zero:
+            not_zero = te.all(*not_zero)
+            return te.if_then_else(not_zero, padded(*index_tuple), tirx.const(0.0, padded.dtype))
+        return padded(*index_tuple)
+
+    # convolution stage
+    out_h = (in_h - 1) * stride_h - fpad_top - fpad_bottom + filter_h
+    out_w = (in_w - 1) * stride_w - fpad_left - fpad_right + filter_w
+    rc = te.reduce_axis((0, in_c), name="rc")
+    rh = te.reduce_axis((0, filter_h), name="rh")
+    rw = te.reduce_axis((0, filter_w), name="rw")
+
+    output = te.compute(
+        (batch, out_h, out_w, out_c),
+        lambda n, h, w, co: te.sum(
+            _dilate(n, h + rh + border_h, w + rw + border_w, rc).astype(out_dtype)
+            * weight[filter_h - 1 - rh, filter_w - 1 - rw, rc, co].astype(out_dtype),
+            axis=[rh, rw, rc],
+        ),
+        name="conv2d_transpose_nhwc",
+    )
+    return (inputs, weight, output)
+
+
+def conv2d_capsule_nhwijc(  # pylint: disable=invalid-name,missing-docstring
+    N: int,
+    H: int,
+    W: int,
+    CI: int,
+    CO: int,
+    kernel_size: int,
+    stride: int = 1,
+    padding: int = 0,
+    capsule_size: int = 4,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    inputs = te.placeholder(
+        (N, H, W, capsule_size, capsule_size, CI), name="inputs", dtype=in_dtype
+    )
+    weight = te.placeholder(
+        (kernel_size, kernel_size, capsule_size, capsule_size, CI, CO),
+        name="weight",
+        dtype=in_dtype,
+    )
+    batch_size, in_h, in_w, _, _, in_channel = inputs.shape
+    k_h, k_w, _, _, _, out_channel = weight.shape
+
+    out_h = (in_h + 2 * padding - kernel_size) // stride + 1
+    out_w = (in_w + 2 * padding - kernel_size) // stride + 1
+
+    rh = te.reduce_axis((0, k_h), name="rh")
+    rw = te.reduce_axis((0, k_w), name="rw")
+    cap_k = te.reduce_axis((0, capsule_size), name="cap_k")
+    rc = te.reduce_axis((0, in_channel), name="rc")
+
+    padded = topi.nn.pad(inputs, [0, padding, padding, 0, 0, 0])
+    output = te.compute(
+        (batch_size, out_h, out_w, capsule_size, capsule_size, out_channel),
+        lambda n, h, w, cap_i, cap_j, co: te.sum(
+            (
+                padded[n, h * stride + rh, w * stride + rw, cap_i, cap_k, rc].astype(out_dtype)
+                * weight[rh, rw, cap_k, cap_j, rc, co].astype(out_dtype)
+            ),
+            axis=[rh, rw, cap_k, rc],
+        ),
+        name="conv2d_capsule_nhwijc",
+    )
+    return (inputs, weight, output)
+
+
+def norm_bmn(  # pylint: disable=invalid-name,missing-docstring
+    B: int,
+    M: int,
+    N: int,
+) -> tuple[te.Tensor, te.Tensor]:
+    a = te.placeholder((B, M, N), name="A")
+    i = te.reduce_axis((0, M), name="i")
+    j = te.reduce_axis((0, N), name="j")
+    c = te.compute(
+        (B,),
+        lambda b: te.sum(a[b][i][j] * a[b][i][j], axis=[i, j]),
+        name="C",
+    )
+    d = te.compute((B,), lambda b: te.sqrt(c[b]), name="D")
+    return (a, d)
+
+
+def conv2d_nhwc_without_layout_rewrite(  # pylint: disable=invalid-name
+    Input: te.Tensor,
+    Filter: te.Tensor,
+    stride: int,
+    padding: int,
+    dilation: int,
+    out_dtype="float32",
+):
+    """A copy of `topi.nn.conv2d_nhwc` but without the 'layout_free` attribute.
+    We use this in single op and subgraph evaluation
+    because we don't want to introduce graph level optimization.
+    """
+    assert isinstance(stride, int) or len(stride) == 2
+    assert isinstance(dilation, int) or len(dilation) == 2
+
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    if isinstance(dilation, int):
+        dilation_h = dilation_w = dilation
+    else:
+        dilation_h, dilation_w = dilation
+
+    batch, in_height, in_width, in_channel = Input.shape  # type: ignore
+    kernel_h, kernel_w, _channel, num_filter = Filter.shape  # type: ignore
+
+    # compute the output shape
+    dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
+    dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
+    pad_top, pad_left, pad_down, pad_right = topi.nn.get_pad_tuple(
+        padding, (dilated_kernel_h, dilated_kernel_w)
+    )
+    out_channel = num_filter
+    out_height = topi.utils.simplify(
+        (in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1
+    )
+    out_width = topi.utils.simplify(
+        (in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1
+    )
+    pad_before = [0, pad_top, pad_left, 0]
+    pad_after = [0, pad_down, pad_right, 0]
+    PaddedInput = topi.nn.pad(Input, pad_before, pad_after, name="PaddedInput")
+    rc = te.reduce_axis((0, in_channel), name="rc")
+    ry = te.reduce_axis((0, kernel_h), name="ry")
+    rx = te.reduce_axis((0, kernel_w), name="rx")
+    Output = te.compute(
+        (batch, out_height, out_width, out_channel),
+        lambda nn, yy, xx, ff: te.sum(
+            PaddedInput[
+                nn, yy * stride_h + ry * dilation_h, xx * stride_w + rx * dilation_w, rc
+            ].astype(out_dtype)
+            * Filter[ry, rx, rc, ff].astype(out_dtype),  # type: ignore
+            axis=[ry, rx, rc],
+        ),
+        name="Conv2dOutput",
+        tag="conv2d_nhwc",
+    )
+    return Output
+
+
+def conv2d_nhwc_bn_relu(  # pylint: disable=invalid-name,missing-docstring
+    N: int,
+    H: int,
+    W: int,
+    CI: int,
+    CO: int,
+    kernel_size: int,
+    strides: int,
+    padding: int,
+    dilation: int = 1,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor, te.Tensor, te.Tensor, te.Tensor]:
+    data = te.placeholder((N, H, W, CI), name="data", dtype=in_dtype)
+    kernel = te.placeholder((kernel_size, kernel_size, CI, CO), name="kernel", dtype=in_dtype)
+    bias = te.placeholder((CO,), name="bias")
+    bn_scale = te.placeholder((CO,), name="bn_scale")
+    bn_offset = te.placeholder((CO,), name="bn_offset")
+    OH = (H + 2 * padding - (kernel_size - 1) * dilation - 1) // strides + 1
+    OW = (W + 2 * padding - (kernel_size - 1) * dilation - 1) // strides + 1
+    conv = conv2d_nhwc_without_layout_rewrite(data, kernel, strides, padding, dilation, out_dtype)
+    conv = te.compute(
+        (N, OH, OW, CO), lambda i, j, k, l: conv[i, j, k, l] + bias[l], name="bias_add"
+    )
+    conv = te.compute(
+        (N, OH, OW, CO), lambda i, j, k, l: conv[i, j, k, l] * bn_scale[l], name="bn_mul"
+    )
+    conv = te.compute(
+        (N, OH, OW, CO), lambda i, j, k, l: conv[i, j, k, l] + bn_offset[l], name="bn_add"
+    )
+    out = topi.nn.relu(conv)
+    return (data, kernel, bias, bn_offset, bn_scale, out)
+
+
+def transpose_batch_matmul(  # pylint: disable=invalid-name,missing-docstring
+    batch: int,
+    seq_len: int,
+    n_head: int,
+    n_dim: int,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    query = te.placeholder((batch, seq_len, n_head, n_dim), name="query", dtype=in_dtype)
+    value = te.placeholder((batch, seq_len, n_head, n_dim), name="value", dtype=in_dtype)
+    query_T = te.compute(
+        (batch, n_head, seq_len, n_dim),
+        lambda b, h, l, d: query[b, l, h, d],
+        name="query_T",
+    )
+    value_T = te.compute(
+        (batch, n_head, n_dim, seq_len),
+        lambda b, h, d, l: value[b, l, h, d],
+        name="value_T",
+    )
+    k = te.reduce_axis((0, n_dim), name="k")
+    out = te.compute(
+        (batch, n_head, seq_len, seq_len),
+        lambda b, h, i, j: te.sum(
+            query_T[b, h, i, k].astype(out_dtype) * value_T[b, h, k, j].astype(out_dtype), axis=[k]
+        ),
+        name="C",
+    )
+    return (query, value, out)
+
+
+def conv2d_winograd_nhwc(  # pylint: disable=invalid-name,missing-docstring
+    N: int,
+    H: int,
+    W: int,
+    CI: int,
+    CO: int,
+    kernel_size: int,
+    stride: int = 1,
+    padding: int = 0,
+    dilation: int = 1,
+    tile_size: int = 4,
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    from tvm.topi.nn.conv2d import (  # pylint: disable=import-outside-toplevel
+        _conv2d_winograd_nhwc_impl,
+    )
+
+    target = Target.current(allow_none=True)
+    if target is not None and target.kind.name == "cuda":
+        write_cache_level = 3
+    else:
+        write_cache_level = 2
+    data = te.placeholder((N, H, W, CI), "float32", name="data")
+    weight = te.placeholder((kernel_size, kernel_size, CO, CI), "float32", name="weight")
+    out = _conv2d_winograd_nhwc_impl(
+        data,
+        weight,
+        stride,
+        padding,
+        dilation,
+        "float32",
+        pre_computed=True,
+        auto_scheduler_rewritten_layout="",
+        meta_schedule_original_shape=None,
+        tile_size=tile_size,
+        write_cache_level=write_cache_level,
+    )
+    return (data, weight, out)
+
+
+def matmul(
+    n: int, m: int, k: int, in_dtype: str = "float32", out_dtype: str = "float32"
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    a = te.placeholder((n, k), name="A", dtype=in_dtype)
+    b = te.placeholder((k, m), name="B", dtype=in_dtype)
+    k = te.reduce_axis((0, k), name="k")
+    c = te.compute(
+        (n, m),
+        lambda i, j: te.sum(a[i, k].astype(out_dtype) * b[k, j].astype(out_dtype), axis=[k]),
+        name="C",
+    )
+    return (a, b, c)
+
+
+def matmul_relu(
+    n: int, m: int, k: int, in_dtype: str = "float32", out_dtype: str = "float32"
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    a = te.placeholder((n, k), name="A", dtype=in_dtype)
+    b = te.placeholder((k, m), name="B", dtype=in_dtype)
+    k = te.reduce_axis((0, k), name="k")
+    c = te.compute(
+        (n, m),
+        lambda i, j: te.sum(a[i, k].astype(out_dtype) * b[k, j].astype(out_dtype), axis=[k]),
+        name="C",
+    )
+    d = topi.nn.relu(c)  # pylint: disable=invalid-name
+    return (a, b, d)
+
+
+def conv2d_nchw(  # pylint: disable=invalid-name
+    n: int,
+    h: int,
+    w: int,
+    ci: int,
+    co: int,
+    kh: int,
+    kw: int,
+    stride: int,
+    padding: int,
+    dilation: int = 1,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
+    x = te.placeholder((n, ci, h, w), name="X", dtype=in_dtype)
+    w = te.placeholder((co, ci, kh, kw), name="W", dtype=in_dtype)
+    y = topi.nn.conv2d_nchw(
+        Input=x, Filter=w, stride=stride, padding=padding, dilation=dilation, out_dtype=out_dtype
+    )
+    return (x, w, y)
+
+
+def conv2d_nchw_bias_bn_relu(  # pylint: disable=invalid-name
+    n: int,
+    h: int,
+    w: int,
+    ci: int,
+    co: int,
+    kh: int,
+    kw: int,
+    stride: int,
+    padding: int,
+    dilation: int = 1,
+    in_dtype: str = "float32",
+    out_dtype: str = "float32",
+) -> tuple[te.Tensor, te.Tensor, te.Tensor, te.Tensor, te.Tensor, te.Tensor]:
+    oh = (h + 2 * padding - (kh - 1) * dilation - 1) // stride + 1  # pylint: disable=invalid-name
+    ow = (w + 2 * padding - (kw - 1) * dilation - 1) // stride + 1  # pylint: disable=invalid-name
+    x = te.placeholder((n, ci, h, w), name="X", dtype=in_dtype)
+    w = te.placeholder((co, ci, kh, kw), name="W", dtype=in_dtype)
+    b = te.placeholder((co, 1, 1), name="B", dtype=out_dtype)
+    bn_scale = te.placeholder((co, 1, 1), name="bn_scale", dtype=out_dtype)
+    bn_offset = te.placeholder((co, 1, 1), name="bn_offset", dtype=out_dtype)
+    y = topi.nn.conv2d_nchw(
+        Input=x, Filter=w, stride=stride, padding=padding, dilation=dilation, out_dtype=out_dtype
+    )
+    y = te.compute((n, co, oh, ow), lambda i, j, k, l: y[i, j, k, l] + b[j, 0, 0], name="bias_add")
+    y = te.compute(
+        (n, co, oh, ow), lambda i, j, k, l: y[i, j, k, l] * bn_scale[j, 0, 0], name="bn_mul"
+    )
+    y = te.compute(
+        (n, co, oh, ow), lambda i, j, k, l: y[i, j, k, l] + bn_offset[j, 0, 0], name="bn_add"
+    )
+    y = topi.nn.relu(y)
+    return (x, w, b, bn_scale, bn_offset, y)
+
+
+def max_pool2d_nchw(  # pylint: disable=invalid-name
+    n: int,
+    h: int,
+    w: int,
+    ci: int,
+    padding: int,
+) -> tuple[te.Tensor, te.Tensor]:  # pylint: disable=invalid-name
+    x = te.placeholder((n, ci, h, w), name="X")
+    y = topi.nn.pool2d(x, [2, 2], [1, 1], [1, 1], [padding, padding, padding, padding], "max")
+    return (x, y)
+
+
+def softmax_mn(m, n) -> tuple[te.Tensor, te.Tensor]:  # pylint: disable=invalid-name
+    a = te.placeholder((m, n), name="A")
+    b = topi.nn.softmax(a, axis=1)
+
+    return (a, b)
+
+
+def create_te_workload(name: str, idx: int) -> tirx.PrimFunc:
+    workload_func, params = CONFIGS[name]
+    return te.create_prim_func(workload_func(*params[idx]))  # type: ignore
+
+
+CONFIGS = {
+    "C1D": (
+        conv1d_nlc,
+        [
+            # derived from conv2d_shapes
+            (1, 256, 64, 128, 3, 2, 1),
+            #    (1, 256, 64, 128, 1, 2, 0),
+            #    (1, 256, 64, 64, 1, 1, 0),
+            #    (1, 128, 128, 256, 3, 2, 1),
+            (1, 128, 128, 256, 1, 2, 0),
+            #    (1, 128, 128, 128, 3, 1, 1),
+            #    (1, 64, 256, 512, 3, 2, 1),
+            #    (1, 64, 256, 512, 1, 2, 0),
+            (1, 64, 256, 256, 5, 1, 2),
+            (1, 32, 512, 512, 3, 1, 1),
+        ],
+    ),
+    "C2D": (
+        conv2d_nhwc,
+        [
+            # all conv2d layers in resnet-18
+            (1, 224, 224, 3, 64, 7, 2, 3),
+            #    (1, 56, 56, 64, 128, 3, 2, 1),
+            #    (1, 56, 56, 64, 128, 1, 2, 0),
+            #    (1, 56, 56, 64, 64, 3, 1, 1),
+            (1, 56, 56, 64, 64, 1, 1, 0),
+            #    (1, 28, 28, 128, 256, 3, 2, 1),
+            #    (1, 28, 28, 128, 256, 1, 2, 0),
+            #    (1, 28, 28, 128, 128, 3, 1, 1),
+            #    (1, 14, 14, 256, 512, 3, 2, 1),
+            #    (1, 14, 14, 256, 512, 1, 2, 0),
+            (1, 14, 14, 256, 256, 3, 1, 1),
+            (1, 7, 7, 512, 512, 3, 1, 1),
+        ],
+    ),
+    "C3D": (
+        conv3d_ndhwc,
+        [
+            # Derived from conv2d_shapes. Use depth=16 for all configurations
+            (1, 16, 224, 224, 3, 64, 7, 2, 3),
+            #    (1, 16, 56, 56, 64, 128, 3, 2, 1),
+            #    (1, 16, 56, 56, 64, 128, 1, 2, 0),
+            #    (1, 16, 56, 56, 64, 64, 3, 1, 1),
+            (1, 16, 56, 56, 64, 64, 1, 1, 0),
+            #    (1, 16, 28, 28, 128, 256, 3, 2, 1),
+            #    (1, 16, 28, 28, 128, 256, 1, 2, 0),
+            #    (1, 16, 28, 28, 128, 128, 3, 1, 1),
+            #    (1, 16, 14, 14, 256, 512, 3, 2, 1),
+            #    (1, 16, 14, 14, 256, 512, 1, 2, 0),
+            (1, 16, 14, 14, 256, 256, 3, 1, 1),
+            (1, 16, 7, 7, 512, 512, 3, 1, 1),
+        ],
+    ),
+    "GMM": (
+        batch_matmul_nkkm,
+        [
+            (1, 128, 128, 128),
+            (1, 512, 32, 512),
+            (1, 512, 512, 512),
+            (1, 1024, 1024, 1024),
+        ],
+    ),
+    "GRP": (
+        conv2d_nhwc,
+        [
+            # Derived from conv2d_shapes. Use group=4 for all configurations
+            (1, 56, 56, 64, 128, 3, 2, 1, 1, 4),
+            #    (1, 56, 56, 64, 128, 1, 2, 0 , 1, 4),
+            #    (1, 56, 56, 64, 64, 3, 1, 1  , 1, 4),
+            (1, 56, 56, 64, 64, 1, 1, 0, 1, 4),
+            #    (1, 28, 28, 128, 256, 3, 2, 1, 1, 4),
+            #    (1, 28, 28, 128, 256, 1, 2, 0, 1, 4),
+            #    (1, 28, 28, 128, 128, 3, 1, 1, 1, 4),
+            #    (1, 14, 14, 256, 512, 3, 2, 1, 1, 4),
+            #    (1, 14, 14, 256, 512, 1, 2, 0, 1, 4),
+            (1, 14, 14, 256, 256, 3, 1, 1, 1, 4),
+            (1, 7, 7, 512, 512, 3, 1, 1, 1, 4),
+        ],
+    ),
+    "DIL": (
+        conv2d_nhwc,
+        [
+            # Derived from conv2d_shapes. Use dilation=2 for all configurations
+            (1, 224, 224, 3, 64, 7, 2, 3, 2),
+            #    (1, 56, 56, 64, 128, 3, 2, 1 , 2),
+            #    (1, 56, 56, 64, 128, 1, 2, 0 , 2),
+            #    (1, 56, 56, 64, 64, 3, 1, 1  , 2),
+            (1, 56, 56, 64, 64, 1, 1, 0, 2),
+            #    (1, 28, 28, 128, 256, 3, 2, 1, 2),
+            #    (1, 28, 28, 128, 256, 1, 2, 0, 2),
+            #    (1, 28, 28, 128, 128, 3, 1, 1, 2),
+            #    (1, 14, 14, 256, 512, 3, 2, 1, 2),
+            #    (1, 14, 14, 256, 512, 1, 2, 0, 2),
+            (1, 14, 14, 256, 256, 3, 1, 1, 2),
+            (1, 7, 7, 512, 512, 3, 1, 1, 2),
+        ],
+    ),
+    "DEP": (
+        depthwise_conv2d_nhwc,
+        [
+            # all depthwise conv2d layers in mobilenet
+            (1, 112, 112, 32, 3, 1, 1),
+            (1, 112, 112, 64, 3, 2, 1),
+            #    (1,  56,  56, 128, 3, 1, 1),
+            #    (1,  56,  56, 128, 3, 2, 1),
+            #    (1,  28,  28, 256, 3, 1, 1),
+            #    (1,  28,  28, 256, 3, 2, 1),
+            #    (1,  14,  14, 512, 3, 1, 1),
+            (1, 14, 14, 512, 3, 2, 1),
+            (1, 7, 7, 1024, 3, 1, 1),
+        ],
+    ),
+    "T2D": (
+        conv2d_transpose_nhwc,
+        [
+            # all conv2d transpose layers in DCGAN
+            (1, 4, 4, 512, 256, 4, 2, 1),
+            (1, 8, 8, 256, 128, 4, 2, 1),
+            (1, 16, 16, 128, 64, 4, 2, 1),
+            (1, 32, 32, 64, 3, 4, 2, 1),
+        ],
+    ),
+    "CAP": (
+        conv2d_capsule_nhwijc,
+        [
+            # all conv2d capsule layers in matrix capsules withemrouting (ICLR 2018)
+            (1, 16, 16, 32, 32, 3, 2, 1),
+            (1, 8, 8, 32, 32, 3, 1, 1),
+            (1, 16, 16, 8, 16, 3, 2, 1),
+            (1, 8, 8, 16, 16, 3, 1, 1),
+        ],
+    ),
+    "NRM": (
+        norm_bmn,
+        [
+            (1, 256, 256),
+            (1, 512, 512),
+            (1, 1024, 1024),
+            (1, 4096, 1024),
+        ],
+    ),
+    "SFM": (
+        softmax_mn,
+        [
+            (256, 256),
+            (512, 512),
+            (1024, 1024),
+            (2048, 2048),
+        ],
+    ),
+    "CBR": (
+        conv2d_nhwc_bn_relu,
+        [
+            (1, 224, 224, 3, 64, 7, 2, 3),
+            (1, 56, 56, 64, 128, 3, 2, 1),
+            (1, 28, 28, 128, 256, 1, 2, 0),
+            (1, 7, 7, 512, 512, 3, 1, 1),
+        ],
+    ),
+    "TBG": (
+        transpose_batch_matmul,
+        [
+            (1, 128, 12, 64),
+            (1, 128, 16, 64),
+            (1, 64, 12, 128),
+            (1, 128, 12, 128),
+        ],
+    ),
+    "C2D_WIN_NHWC": (
+        conv2d_winograd_nhwc,
+        [
+            (1, 14, 14, 128, 128, 6),
+        ],
+    ),
+}

python/tvm/s_tir/meta_schedule/testing/tune_te.py

@@ -0,0 +1,150 @@
+# 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=missing-docstring
+# ruff: noqa: F821
+import argparse
+import logging
+
+import tvm
+from tvm.s_tir import meta_schedule as ms
+from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
+from tvm.support import describe
+from tvm.testing.utils import strtobool
+
+
+def _parse_args():
+    args = argparse.ArgumentParser()
+    args.add_argument(
+        "--workload",
+        type=str,
+        required=True,
+    )
+    args.add_argument(
+        "--target",
+        type=str,
+        required=True,
+    )
+    args.add_argument(
+        "--num-trials",
+        type=int,
+        required=True,
+    )
+    args.add_argument(
+        "--rpc-host",
+        type=str,
+        required=True,
+    )
+    args.add_argument(
+        "--rpc-port",
+        type=int,
+        required=True,
+    )
+    args.add_argument(
+        "--rpc-key",
+        type=str,
+        required=True,
+    )
+    args.add_argument(
+        "--work-dir",
+        type=str,
+        required=True,
+    )
+    args.add_argument(
+        "--number",
+        type=int,
+        default=3,
+    )
+    args.add_argument(
+        "--repeat",
+        type=int,
+        default=1,
+    )
+    args.add_argument(
+        "--min-repeat-ms",
+        type=int,
+        default=100,
+    )
+    args.add_argument(
+        "--adaptive-training",
+        type=lambda x: bool(strtobool(x)),
+        required=False,
+        help="example: True / False",
+        default=True,
+    )
+    args.add_argument(
+        "--cpu-flush",
+        type=lambda x: bool(strtobool(x)),
+        help="example: True / False",
+        required=True,
+    )
+    parsed = args.parse_args()
+    parsed.target = tvm.target.Target(parsed.target)
+    parsed.rpc_config = ms.runner.RPCConfig(
+        tracker_host=parsed.rpc_host,
+        tracker_port=parsed.rpc_port,
+        tracker_key=parsed.rpc_key,
+        session_timeout_sec=60,
+    )
+    return parsed
+
+
+logging.basicConfig(
+    format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+)
+logging.getLogger("tvm.s_tir.meta_schedule").setLevel(logging.DEBUG)
+ARGS = _parse_args()
+
+
+def main():
+    describe()
+    print(f"Workload: {ARGS.workload}")
+    with ms.Profiler() as profiler:
+        sch: s_tir.Schedule | None = ms.tir_integration.tune_tir(
+            mod=create_te_workload(ARGS.workload, 0),
+            target=ARGS.target,
+            work_dir=ARGS.work_dir,
+            max_trials_global=ARGS.num_trials,
+            num_trials_per_iter=64,
+            runner=ms.runner.RPCRunner(  # type: ignore
+                rpc_config=ARGS.rpc_config,
+                evaluator_config=ms.runner.EvaluatorConfig(
+                    number=ARGS.number,
+                    repeat=ARGS.repeat,
+                    min_repeat_ms=ARGS.min_repeat_ms,
+                    enable_cpu_cache_flush=ARGS.cpu_flush,
+                ),
+                alloc_repeat=1,
+            ),
+            cost_model=ms.cost_model.XGBModel(  # type: ignore
+                extractor=ms.feature_extractor.PerStoreFeature(),
+                adaptive_training=ARGS.adaptive_training,
+            ),
+            strategy=ms.search_strategy.EvolutionarySearch(),
+        )
+
+    print("Tuning Time:")
+    print(profiler.table())
+
+    if sch is None:
+        print("No valid schedule found!")
+    else:
+        print(sch.mod.script())
+        print(sch.trace)
+
+
+if __name__ == "__main__":
+    main()