chore: import upstream snapshot with attribution

python/tvm/s_tir/__init__.py

@@ -0,0 +1,57 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""S-TIR namespace for scheduable TensorIR"""
+
+from tvm.tirx.function import TensorIntrin
+
+# dlight depends on compiler-only C++ functions (e.g. s_tir.schedule.GetSBlockRealize),
+# so skip it in runtime-only builds.
+from tvm.base import _RUNTIME_ONLY
+
+from . import _ffi_api
+from . import backend
+from . import pipeline
+from . import transform
+from . import schedule
+from .schedule import StmtSRef, SBlockScope, ScheduleState, Schedule, ScheduleError, Trace
+from .sblock_dependence_info import SBlockDependenceInfo
+from .data_layout import SLayout, SBijectiveLayout, sbijective_layout, slayout
+
+if not _RUNTIME_ONLY:
+    from . import analysis
+    from . import meta_schedule
+    from . import dlight
+
+
+def renew_defs(func):
+    """Re-generate the definition nodes for a TIR, including VarDef, BufferDef.
+    This pass works as a simple DeepCopy to duplicate a function with different Vars and
+    Buffers but the same behavior
+
+    Parameters
+    ----------
+    func: PrimFunc
+        The input function
+
+    Returns
+    -------
+    result : PrimFunc
+        The new generated func.
+    """
+    return _ffi_api.RenewDefs(func)

python/tvm/s_tir/_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"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("s_tir", __name__)

python/tvm/s_tir/analysis/__init__.py

@@ -0,0 +1,216 @@
+# 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.
+"""Analysis utilities for Schedulable TensorIR (S-TIR)."""
+
+# pylint: disable=invalid-name
+from typing import Optional, Union
+
+import tvm
+from tvm.ir import IRModule
+from tvm.tirx.expr import Var
+from tvm.tirx.stmt import SBlock, BufferRegion
+
+from tvm.tirx import Buffer, Stmt
+from tvm.tirx.function import PrimFunc
+from . import _ffi_api
+
+
+def get_sblock_access_region(
+    block: SBlock, buffer_var_map: dict[Var, Buffer]
+) -> list[list[BufferRegion]]:
+    """Detect which regions of tensors in this block are read or written to.
+       Regions are sorted by order of appearance in the AST.
+
+    Parameters
+    ----------
+    block: tvm.tirx.SBlock
+        The block in which we are detecting read/write regions.
+
+    buffer_var_map : Dict[Var, Buffer]
+        The outside buffers which may access the block. Mapping from buffer var to the buffer
+
+    Returns
+    -------
+    result : List[List[BufferRegion]]
+        Array of access regions. There are three arrays of BufferRegion:
+            - first: read regions
+            - second: write regions
+            - third: opaque regions
+    """
+    return _ffi_api.GetSBlockAccessRegion(block, buffer_var_map)  # type: ignore
+
+
+def get_sblock_read_write_region(
+    block: SBlock, buffer_var_map: dict[Var, Buffer]
+) -> list[list[BufferRegion]]:
+    """Auto detect the block read/write region according to its body stmt.
+       An opaque access will be counted as both a read and a write access
+
+    Parameters
+    ----------
+    block: tvm.tirx.SBlock
+        The block in which we are detecting read/write regions.
+
+    buffer_var_map : Dict[Var, Buffer]
+        The outside buffers which may access the block. Mapping from buffer var to the buffer
+
+    Returns
+    -------
+    result : List[List[BufferRegion]]
+        An array only consisting of the read regions and write regions of the input block
+    """
+    return _ffi_api.GetSBlockReadWriteRegion(block, buffer_var_map)  # type: ignore
+
+
+def detect_buffer_access_lca(func: PrimFunc) -> dict[Buffer, Stmt]:
+    """Detect the lowest common ancestor(LCA) of buffer access, including both high-level
+    access (BufferLoad, BufferStore) and low-level access (BufferLoad, BufferStore and opaque
+    access).
+    The LCA may be a For loop or a Block.
+
+    Parameters
+    ----------
+    func: tvm.tirx.PrimFunc
+        The function to be detected.
+
+    Returns
+    -------
+    result : Dict[Buffer, Stmt]
+        Map from buffer to the LCA of all access to it.
+    """
+    return _ffi_api.detect_buffer_access_lca(func)  # type: ignore # pylint: disable=no-member
+
+
+def find_anchor_sblock(mod: IRModule) -> SBlock | None:
+    """Find the "anchor block" of the given module.
+
+    We define the anchor block to be the block with (1) an init statement and (2) having
+    the biggest flops count. The latter condition is only used when there are multiple blocks
+    with an init statement.
+
+    For example, if the input module is conv2d + fused spatial blocks, conv2d is the anchor block.
+    The input module may not contain more than one such block. For example, a module having
+    two conv2d is not allowed as an input.
+
+    However, a module created from winograd convolution has multiple blocks with an init statement
+    (input transform, batched GEMM, and output transform). We use the second condition, the flops
+    count, to determine that the batched GEMM block is the anchor block.
+
+    Parameters
+    ----------
+    mod: tvm.ir.IRModule
+        The input TIR module.
+    Returns
+    -------
+    anchor_block: Optional[SBlock]
+        The anchor block if found, None otherwise.
+    """
+    return _ffi_api.find_anchor_sblock(mod)  # type: ignore # pylint: disable=no-member
+
+
+def verify_gpu_code(func: PrimFunc, constraints: dict[str, int]) -> bool:
+    """Verify if module contains illegal host side direct memory access.
+
+    Parameters
+    ----------
+    func: tvm.tirx.PrimFunc
+        The module to be verified.
+
+    constraints : Dict[str, int]
+        The attribute constraints.
+
+    Returns
+    -------
+    result : bool
+        The result of verification.
+    """
+    return _ffi_api.verify_gpu_code(func, constraints)  # type: ignore
+
+
+def calculate_allocated_bytes(
+    func_or_mod: PrimFunc | IRModule,
+) -> dict[str, dict[str, int]]:
+    """Calculate allocated memory per memory scope required by TIR PrimFuncs.
+
+    Parameters
+    ----------
+    func_or_mod: Union[PrimFunc, IRModule]
+        The function or module to be detected. If a module is passed, allocated
+        memory is calculated for all PrimFuncs inside the module
+
+    Returns
+    -------
+    result : Dict[str, Dict[str, int]]
+        Allocated memory size per scope in bytes for each function in the IRModule returned as a
+        dict with function names as keys and a dict of allocated sizes as values. If a single
+        PrimFunc is passed, the function name is returned as "main"
+    """
+    if not isinstance(func_or_mod, PrimFunc | IRModule):
+        raise TypeError(
+            f"Expected argument to be PrimFunc or IRModule, but received {type(func_or_mod)}"
+        )
+    return _ffi_api.calculate_allocated_bytes(func_or_mod)  # type: ignore
+
+
+def estimate_tir_flops(stmt_or_mod: Stmt | IRModule) -> float:
+    """Estimate the FLOPs of a TIR fragment.
+
+    Parameters
+    ----------
+    stmt_or_mod: Union[Stmt, IRModule]
+        The TIR fragment or IRModule to be estimated.
+
+    Returns
+    -------
+    flops: float
+        The estimated FLOPs.
+    """
+    return _ffi_api.EstimateTIRFlops(stmt_or_mod)  # type: ignore # pylint: disable=no-member
+
+
+def OOBChecker():
+    """Detect out of bounds memory access in arrays.
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.OOBChecker()  # type: ignore
+
+
+def get_vtcm_compaction_passes() -> list[tvm.transform.Pass]:
+    """Utility function to get the list of lowering passes to be applied to calculate the compacted
+    VTCM allocation size
+
+    Returns
+    -------
+    result : List[tvm.transform.Pass]
+        returns list of passes
+    """
+    return _ffi_api.get_vtcm_compaction_passes()  # type: ignore # pylint: disable=no-member
+
+
+def is_pure_function(func: PrimFunc) -> bool:
+    """Checks if the function is a pure function"""
+    return _ffi_api.is_pure_function(func, False)  # type: ignore # pylint: disable=no-member
+
+
+def assert_pure_function(func: PrimFunc) -> bool:
+    """Asserts that the function is a pure function"""
+    return _ffi_api.is_pure_function(func, True)  # type: ignore # pylint: disable=no-member

python/tvm/s_tir/analysis/_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.analysis"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("s_tir.analysis", __name__)

python/tvm/s_tir/backend/__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.
+"""S-TIR backend compilation pipeline and other passes."""
+
+from . import adreno

python/tvm/s_tir/backend/adreno/__init__.py

@@ -0,0 +1,21 @@
+# 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 S-TIR Adreno pipeline passes"""
+
+from . import pipeline
+from . import transform

python/tvm/s_tir/backend/adreno/pipeline.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.
+
+# pylint: disable=invalid-name
+"""The TIR backend compilation pipeline for Adreno"""
+
+import tvm
+from tvm import s_tir, tirx
+from tvm.tirx import compilation_pipeline as tir_pipeline
+
+
+def default_tir_pipeline():
+    """The default tirx pipeline used in tvm.tirx.build"""
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
+        """The default lowering passes for TIR backend."""
+        pass_ctx = tvm.transform.PassContext.current()
+        config = pass_ctx.config
+        passes = [
+            s_tir.backend.adreno.transform.TextureFlatten(),
+            s_tir.transform.CanonicalizeLoop(),
+            s_tir.transform.LowerCrossThreadReduction(),
+            s_tir.transform.LowerInitBlock(),
+            s_tir.transform.PlanAndUpdateBufferAllocationLocation(),
+            s_tir.transform.ConvertBlocksToOpaque(),
+            s_tir.transform.LiftThreadBinding(),
+            s_tir.transform.ManifestSharedMemoryLocalStage(),
+            s_tir.transform.CompactBufferAllocation(),
+            s_tir.transform.LowerAutoCopy(),
+            s_tir.transform.UnifyThreadBinding(),
+            s_tir.transform.LowerMatchBuffer(),
+            tirx.transform.StmtSimplify(),
+            s_tir.transform.InjectPermutedLayout(),
+            s_tir.transform.AnnotateIrregularLoop(),
+            s_tir.transform.InjectSoftwarePipeline(),
+            s_tir.transform.TransformMmaBufferLayout(),
+            s_tir.transform.LowerOpaqueBlock(),
+            s_tir.backend.adreno.transform.InjectTextureAlloc(),
+            tirx.transform.FlattenBuffer(),
+            tirx.transform.BF16ComputeLegalize(),
+            tirx.transform.NarrowDataType(32),
+            s_tir.transform.LoopPartition(),
+            tirx.transform.VectorizeLoop(not bool(config.get("tirx.disable_vectorize", False))),
+            s_tir.transform.InjectVirtualThread(),
+            s_tir.transform.InjectDoubleBuffer(),
+        ]
+        if not bool(config.get("tirx.disable_storage_rewrite", False)):
+            passes.append(tirx.transform.StorageRewrite())
+        if config.get("tirx.use_async_copy", False):
+            passes.append(s_tir.transform.LowerAsyncDMA())
+        passes.extend(
+            [
+                s_tir.transform.HoistIfThenElse(),
+                tirx.transform.UnrollLoop(),
+                s_tir.transform.RenormalizeSplitPattern(),
+                tirx.transform.StmtSimplify(),
+                tirx.transform.RemoveNoOp(),
+                s_tir.transform.RewriteUnsafeSelect(),
+            ]
+        )
+        # Additional passes based on configuration.
+        if bool(config.get("tirx.instrument_bound_checkers", False)):
+            passes.append(s_tir.transform.InstrumentBoundCheckers())
+        if bool(config.get("tirx.ptx.ldg32", False)):
+            passes.append(s_tir.transform.InjectPTXLDG32(True))
+        if not bool(config.get("tirx.disable_cse_tir", False)):
+            passes.append(tirx.transform.CommonSubexprElim())
+        if bool(config.get("tirx.instrument_lwp", False)):
+            passes.append(s_tir.transform.InstrumentProfileIntrinsics())
+        passes.extend(
+            [
+                # Bind the target first so that target-specific attributes are available.
+                tirx.transform.FP8ComputeLegalize(),
+                # VerifyVTCMLimit must occur before LowerVtcmAlloc.
+                s_tir.transform.VerifyVTCMLimit(),
+                s_tir.transform.LowerVtcmAlloc(),
+                tirx.transform.VerifyMemory(),
+                tirx.transform.AnnotateEntryFunc(),
+            ]
+        )
+        passes.extend(
+            [
+                s_tir.transform.ThreadSync("shared"),
+                s_tir.transform.ThreadSync("shared.dyn"),
+                s_tir.transform.ThreadSync("warp"),
+                s_tir.transform.InferFragment(),
+                s_tir.transform.LowerThreadAllreduce(),
+            ]
+        )
+        if bool(config.get("tirx.use_async_copy", False)):
+            passes.append(s_tir.transform.InjectPTXAsyncCopy())
+        if bool(config.get("tirx.ptx.ldg32", False)):
+            passes.append(s_tir.transform.InjectPTXLDG32())
+        passes.extend(
+            [
+                s_tir.transform.MergeSharedMemoryAllocations(),
+                tirx.transform.SplitHostDevice(),
+                tirx.transform.MakePackedAPI(),
+                tirx.transform.FP8StorageLegalize(),
+                tirx.transform.BF16StorageLegalize(),
+            ]
+        )
+        mod = tvm.ir.transform.Sequential(passes)(mod)
+        return mod
+
+    return _pipeline
+
+
+tir_pipeline.PIPELINE_MAP["adreno"] = default_tir_pipeline

python/tvm/s_tir/backend/adreno/transform/__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.
+"""Adreno TIR transformations."""
+
+from .transform import (
+    InjectTextureAlloc,
+    TextureFlatten,
+)

python/tvm/s_tir/backend/adreno/transform/_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.backend.adreno.transform"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("s_tir.backend.adreno.transform", __name__)

python/tvm/s_tir/backend/adreno/transform/transform.py

@@ -0,0 +1,46 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Wrapping existing transformations."""
+# pylint: disable=invalid-name, unsupported-binary-operation
+
+from . import _ffi_api
+
+
+def InjectTextureAlloc():
+    """Inject Texture Allocation Intrinsic to make sure appropriate lowering
+    via alloc_nd/alloc_free calls
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.InjectTextureAlloc()  # type: ignore
+
+
+def TextureFlatten():
+    """Flatten the multi-dimensional read/write to 2D.
+
+
+    Parameters
+    ----------
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.TextureFlatten()  # type: ignore

python/tvm/s_tir/data_layout.py

@@ -0,0 +1,213 @@
+# 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.
+"""Data layout."""
+
+import tvm_ffi
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("s_tir.SLayout")
+class SLayout(Object):
+    """SLayout is composed of upper cases, lower cases and numbers,
+    where upper case indicates a primal axis and
+    the corresponding lower case with factor size indicates the subordinate axis.
+    For example, NCHW16c can describe a 5-D tensor of
+    [batch_size, channel, height, width, channel_block].
+    Here subordinate axis channel_block=16 is the factor size of the primal axis C (channel).
+
+    See Also
+    --------
+    slayout : Declare a layout
+    """
+
+    def __len__(self):
+        return _ffi_api.SLayoutNdim(self)  # type: ignore
+
+    def __contains__(self, axis):
+        # Note: We do a weaker check for packed axis assuming layout is valid
+        return not any(bkt in axis for bkt in "[]") and axis in self.name
+
+    def __getitem__(self, index):
+        if index >= len(self):
+            raise IndexError("SLayout index out of range")
+        return _ffi_api.SLayoutGetItem(self, index)  # type: ignore
+
+    def index_of(self, axis):
+        """Get the index of an axis
+
+        Parameters
+        ----------
+        axis : str
+            The axis name, needs to be [a-z,A-Z] or a packed axis
+
+        Returns
+        -------
+        index : int
+            The index of the axis, -1 if not found.
+        """
+        return _ffi_api.SLayoutIndexOf(self, axis)  # type: ignore
+
+    def factor_of(self, axis):
+        """Get the factor size of the subordinate axis.
+
+        Parameters
+        ----------
+        axis : str
+            The axis name, need to be [a-z,A-Z]
+
+        Returns
+        -------
+        factor : int
+            the size of the subordinate-axis of axis (if axis is a primal-axis),
+            or the size of axis itself (if axis is a subordinate-axis).
+            Return -1 if axis is not in the layout.
+        """
+        return _ffi_api.SLayoutFactorOf(self, axis)  # type: ignore
+
+
+@tvm_ffi.register_object("s_tir.SBijectiveLayout")
+class SBijectiveLayout(Object):
+    """Bijective mapping for two layouts (src-layout and dst-layout).
+    It provides shape and index conversion between each other.
+
+    Do not construct directly, use :any:`sbijective_layout` instead.
+    See the documentation of :any:`sbijective_layout` for more details.
+
+    Parameters
+    ----------
+    src_layout : str or SLayout
+        source layout.
+
+    dst_layout : str or SLayout
+        destination layout.
+
+    See Also
+    --------
+    sbijective_layout : Declare a layout
+    """
+
+    def forward_index(self, index):
+        """Given the indices of the src-layout, infer the dst index.
+
+        Parameters
+        ----------
+        index: Array of Expr
+            The indices in src-layout.
+
+        Returns
+        -------
+        dst_index: Array of Expr
+            The inferred indices in dst-layout.
+        """
+        return _ffi_api.SBijectiveLayoutForwardIndex(self, index)  # type: ignore
+
+    def backward_index(self, index):
+        """Given the indices of the dst-layout, infer the src index.
+
+        Parameters
+        ----------
+        index: Array of Expr
+            The indices in dst-layout.
+
+        Returns
+        -------
+        src_index: Array of Expr
+            The inferred indices in src-layout.
+        """
+        return _ffi_api.SBijectiveLayoutBackwardIndex(self, index)  # type: ignore
+
+    def forward_shape(self, shape):
+        """Given the shape of the src-layout, infer the dst shape.
+
+        Parameters
+        ----------
+        shape: Array of Expr
+            The shape in src-layout.
+
+        Returns
+        -------
+        dst_shape: Array of Expr
+            The inferred shape in dst-layout.
+        """
+        return _ffi_api.SBijectiveLayoutForwardShape(self, shape)  # type: ignore
+
+    def backward_shape(self, shape):
+        """Given the shape of the dst-layout, infer the src shape.
+
+        Parameters
+        ----------
+        shape: Array of Expr
+            The shape in dst-layout.
+
+        Returns
+        -------
+        src_shape: Array of Expr
+            The inferred shape in src-layout.
+        """
+        return _ffi_api.SBijectiveLayoutBackwardShape(self, shape)  # type: ignore
+
+
+def slayout(layout_str: str, dtype: str = "int32") -> SLayout:
+    """Create a layout node from a string.
+
+    Parameters
+    ----------
+    layout_str : str
+        A layout representation is composed of upper cases, lower cases and numbers,
+        where upper case indicates a primal axis and
+        the corresponding lower case with factor size indicates the subordinate axis.
+        For example, NCHW16c can describe a 5-D tensor of
+        [batch_size, channel, height, width, channel_block].
+        Here subordinate axis channel_block=16 is the factor size of
+        the primal axis C (channel).
+
+    dtype : str
+        The dtype of generated axes vars in the returned layout.
+        It is required to be integer type.
+
+    Returns
+    -------
+    layout : SLayout
+        The created layout
+    """
+    return _ffi_api.SLayout(layout_str, dtype)  # type: ignore
+
+
+def sbijective_layout(src_layout: str | SLayout, dst_layout: str | SLayout) -> SBijectiveLayout:
+    """Create a bijective layout mapping.
+
+    Parameters
+    ----------
+    src_layout : str or SLayout
+        source layout.
+
+    dst_layout : str or SLayout
+        destination layout.
+
+    Returns
+    -------
+    sbijective_layout : SBijectiveLayout
+        The created bijective layout
+    """
+    if isinstance(src_layout, str):
+        src_layout = slayout(src_layout)
+    if isinstance(dst_layout, str):
+        dst_layout = slayout(dst_layout)
+    return _ffi_api.SBijectiveLayout(src_layout, dst_layout)  # type: ignore

python/tvm/s_tir/dlight/__init__.py

@@ -0,0 +1,33 @@
+# 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.
+"""DLight package provides efficient schedules out-of-box for deep learning workloads."""
+
+from . import gpu
+from . import adreno
+from . import cpu
+from .analysis import (
+    SBlockInfo,
+    IterInfo,
+    normalize_prim_func,
+)
+from .base import (
+    ApplyDefaultSchedule,
+    ScheduleRule,
+    try_inline,
+    try_inline_contiguous_spatial,
+)

python/tvm/s_tir/dlight/adreno/__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.
+"""
+Adreno schedule rules.
+"""
+
+from .convolution import Conv2d
+from .layout_transform import LayoutTransform
+from .fallback import Fallback
+from .pool import Pool2D

python/tvm/s_tir/dlight/adreno/base.py

@@ -0,0 +1,41 @@
+# 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.
+"""Base schedule rule for Adreno operators."""
+
+from tvm.target import Target
+
+from ..base import ScheduleRule
+
+
+class AdrenoScheduleRule(ScheduleRule):  # pylint: disable=too-few-public-methods
+    """The Schedule Rule specific to Adreno targets,
+    will return None if the target is not Adreno."""
+
+    def is_target_available(self, target: Target) -> bool:
+        """Check whether the target is available for Adreno rule.
+
+        Parameters
+        ----------
+        target : Target
+            The compilation target to check.
+
+        Returns
+        -------
+        available : bool
+            Whether the target is available for this rule.
+        """
+        return super().is_target_available(target) and "adreno" in target.keys

python/tvm/s_tir/dlight/adreno/convolution.py

@@ -0,0 +1,99 @@
+# 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, invalid-name
+"""A Conv2d schedule rule for Adreno GPU operators."""
+
+from tvm import s_tir, tirx
+from tvm.target import Target
+
+from .. import analysis
+from .base import AdrenoScheduleRule
+from .utils import schedule_default, schedule_inline_blocks
+
+
+class Conv2d(AdrenoScheduleRule):
+    """The schedule rule for convolution computation"""
+
+    @staticmethod
+    def schedule_conv2d(sch: s_tir.Schedule, blk: s_tir.schedule.SBlockRV):
+        n, oc, oh, ow, ob, ic, kh, kw = sch.get_loops(blk)
+
+        bz, vz, tz = sch.split(oc, [None, 8, 1], preserve_unit_iters=True)
+        by, vy, ty = sch.split(oh, [None, 1, 16], preserve_unit_iters=True)
+        bx, vx, tx = sch.split(ow, [None, 1, 16], preserve_unit_iters=True)
+
+        bz = sch.fuse(n, bz, preserve_unit_iters=True)
+        sch.reorder(bz, by, bx, vz, vy, vx, tz, ty, tx, ob)
+        sch.bind(bz, "blockIdx.z")
+        sch.bind(by, "blockIdx.y")
+        sch.bind(bx, "blockIdx.x")
+        sch.bind(vz, "vthread.z")
+        sch.bind(vy, "vthread.y")
+        sch.bind(vx, "vthread.x")
+        sch.bind(tz, "threadIdx.z")
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+        rblk = sch.cache_read(blk, 0, "local")
+        ico, icb = sch.split(ic, [None, 4], preserve_unit_iters=True)
+        sch.reorder(ico, kh, kw, icb, ob)
+
+        sch.compute_at(rblk, kw, preserve_unit_loops=True)
+        sch.vectorize(sch.get_loops(rblk)[-1])
+        wblk = sch.cache_write(blk, 0, "local")
+        sch.reverse_compute_at(wblk, tx, preserve_unit_loops=True)
+        sch.vectorize(sch.get_loops(wblk)[-1])
+        init_blk = sch.decompose_reduction(blk, tx)
+        sch.vectorize(sch.get_loops(init_blk)[-1])
+
+    def apply(  # pylint: disable=too-many-locals,missing-docstring
+        self,
+        func: tirx.PrimFunc | s_tir.Schedule,
+        target: Target,
+        _: bool,
+    ) -> s_tir.Schedule | None:
+        if not (isinstance(func, tirx.PrimFunc | s_tir.Schedule)) or not self.is_target_available(
+            target
+        ):
+            return None
+
+        if isinstance(func, tirx.PrimFunc):
+            sch = s_tir.Schedule(func)
+            sch.work_on("main")
+        elif isinstance(func, s_tir.Schedule):
+            sch = func
+
+        root_block = analysis.get_root_block(sch, sch.func_working_on)
+        blocks = sch.get_child_blocks(root_block)
+        reduction_blocks = list(
+            filter(lambda block: analysis.get_sblock_info(sch, block).is_reduction(), blocks)
+        )
+        remaining_blocks = [blk for blk in blocks if blk not in reduction_blocks]
+
+        def is_convolution(blk):
+            block_info = analysis.get_sblock_info(sch, blk)
+            return "conv2d_NCHWc" in block_info.name
+
+        if len(reduction_blocks) != 1 or not is_convolution(reduction_blocks[0]):
+            return None
+
+        conv_blk = reduction_blocks[0]
+        Conv2d.schedule_conv2d(sch, conv_blk)
+        remaining_blocks = schedule_inline_blocks(sch, remaining_blocks)
+        schedule_default(sch, remaining_blocks)
+
+        return sch

python/tvm/s_tir/dlight/adreno/fallback.py

@@ -0,0 +1,192 @@
+# 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.
+
+# 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.
+"""Dlight Adreno Fallback Schedules"""
+
+from tvm import s_tir, tirx
+from tvm.target import Target
+
+from .. import analysis
+from .base import AdrenoScheduleRule
+
+
+def _assert_gpu_target(target: Target):
+    if "gpu" not in target.keys:
+        raise ValueError(f"Expect a GPU target, but got {target}")
+
+
+def get_max_threads_per_block(target: Target) -> int:
+    _assert_gpu_target(target)
+    max_threads_per_block = None
+    for name in ["max_threads_per_block", "max_num_threads"]:
+        if max_threads_per_block is None:
+            max_threads_per_block = target.attrs.get(name, None)
+    if max_threads_per_block is None:
+        max_threads_per_block = 64
+    return int(max_threads_per_block)
+
+
+# pylint: disable=invalid-name,missing-function-docstring,unused-variable,unused-import
+class Fallback(AdrenoScheduleRule):
+    """Texture Based Fallback Schedule(s) for Adreno"""
+
+    @staticmethod
+    def schedule_inline_blocks(
+        sch: s_tir.Schedule, blocks: list[s_tir.schedule.SBlockRV]
+    ) -> list[s_tir.schedule.SBlockRV]:
+        """
+        Auto Inlines Injective and Element-wise Operations while trying to omit data pad blocks...
+        """
+
+        if blocks is None:
+            root_blk = analysis.get_root_block(sch)
+            blocks = sch.get_child_blocks(root_blk)
+
+        remaining_blocks = []
+        for blk in blocks:
+            block_info = analysis.get_sblock_info(sch, blk)
+            if block_info.is_injective() and not block_info.is_data_pad(sch):
+                if len(sch.get_consumers(blk)) == 1:
+                    try:
+                        sch.compute_inline(blk)
+                    except Exception:  # pylint: disable=broad-exception-caught
+                        remaining_blocks.append(blk)
+                elif len(sch.get_producers(blk)) == 1:
+                    inlined_once = False
+                    try:
+                        # Would cause an issue inlining to producer with multiple consumers
+                        while (
+                            len(sch.get_producers(blk)) == 1
+                            and len(sch.get_consumers(sch.get_producers(blk)[0])) == 1
+                        ):
+                            sch.reverse_compute_inline(blk)
+                            inlined_once = True
+                    except Exception:  # pylint: disable=broad-exception-caught
+                        break
+                    if not inlined_once:
+                        remaining_blocks.append(blk)
+                else:
+                    remaining_blocks.append(blk)
+            else:
+                remaining_blocks.append(blk)
+        return remaining_blocks
+
+    @staticmethod
+    def schedule_default(sch: s_tir.Schedule, blk: s_tir.schedule.SBlockRV):
+        block_info = analysis.get_sblock_info(sch, blk)
+
+        s_loops, r_loops, o_loops = [], [], []
+        v_loop = block_info.write_bufs(sch)[0].assoc_lps[-1]
+
+        for iter_info in block_info.iters:
+            if sch.get(iter_info.loop_rv) == sch.get(v_loop):
+                continue
+            {"S": s_loops, "R": r_loops, "O": o_loops}.get(iter_info.kind).append(iter_info.loop_rv)
+
+        iter_vars = analysis.collect_block_iter_vars_used_in_access_region(
+            sch.get(blk), block_info.write_bufs(sch)[0].buf_region.region
+        )
+        o_outer = [lp for lp in o_loops if sch.get(lp).var in iter_vars]
+        o_inner = [lp for lp in o_loops if sch.get(lp).var not in iter_vars]
+
+        # Can't change loop order for opaque loops
+        if o_loops != o_outer + o_inner:
+            return
+
+        o_outer.append(v_loop)
+        sch.reorder(*s_loops, *o_outer, *r_loops, *o_inner)
+
+        assert s_loops
+        tgt = Target.current(allow_none=True)
+
+        b = sch.fuse(*s_loops)
+        tx_extent = get_max_threads_per_block(tgt) if tgt is not None else 256
+        bx, tx = sch.split(b, [None, tx_extent])
+        sch.bind(bx, "blockIdx.x")
+        sch.bind(tx, "threadIdx.x")
+
+        if len(r_loops) > 1:
+            lp = [*s_loops, *o_outer][-1]
+            init_block = sch.decompose_reduction(blk, lp)
+            wblk = sch.cache_write(blk, 0, "local")
+            sch.compute_at(wblk, lp)
+            if v_loop:
+                sch.vectorize(sch.get_loops(init_block)[-1])
+                sch.vectorize(sch.get_loops(wblk)[-1])
+        elif v_loop is not None:
+            sch.vectorize(v_loop)
+
+    @staticmethod
+    def schedule_fallback(sch):
+        root_block = analysis.get_root_block(sch)
+        blocks = sch.get_child_blocks(root_block)
+
+        schedule_blocks = [
+            blk
+            for blk in blocks
+            if analysis.get_sblock_info(sch, blk).is_reduction()
+            or analysis.get_sblock_info(sch, blk).is_data_pad(sch)
+        ]
+        remaining_blocks = [blk for blk in blocks if blk not in schedule_blocks]
+
+        for blk in schedule_blocks:
+            Fallback.schedule_default(sch, blk)
+        remaining_blocks = Fallback.schedule_inline_blocks(sch, remaining_blocks)
+        # TODO: Analyze unscheduled blocks to schedule instead of relying on remaining
+        for blk in remaining_blocks:
+            Fallback.schedule_default(sch, blk)
+
+    def apply(  # pylint: disable=too-many-locals
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+        # pylint: disable=invalid-name
+
+        if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
+            return None
+
+        sch = s_tir.Schedule(func)
+        root_block = analysis.get_root_block(sch)
+        blocks = sch.get_child_blocks(root_block)
+
+        if any(len(sch.get_child_blocks(block)) != 0 for block in blocks):
+            return None
+
+        block_infos = [analysis.get_sblock_info(sch, block) for block in blocks]
+        if not any("texture" in block.write_bufs(sch)[0].get_scope() for block in block_infos):
+            return None
+
+        Fallback.schedule_fallback(sch)
+        return sch

python/tvm/s_tir/dlight/adreno/layout_transform.py

@@ -0,0 +1,146 @@
+# 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: F841
+
+# Licensed to the apache software foundation (asf) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# pylint: disable=invalid-name, unused-variable
+
+"Schedules for Texture Based Layout Transforms"
+
+from tvm import s_tir, tirx
+from tvm.target import Target
+
+from .. import analysis
+from .base import AdrenoScheduleRule
+
+
+class LayoutTransform(AdrenoScheduleRule):
+    """Texture based Layout Transform Dlight Schedule for Adreno"""
+
+    def __init__(self, use_op_name=True):
+        self.use_op_name = use_op_name
+
+    # TODO: Try using Coalesced Writes...
+    def apply(  # pylint: disable=too-many-locals
+        self,
+        func: tirx.PrimFunc | s_tir.Schedule,
+        target: Target,
+        _: bool,
+    ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+        # pylint: disable=invalid-name
+        if not (isinstance(func, tirx.PrimFunc | s_tir.Schedule)) or not self.is_target_available(
+            target
+        ):
+            return None
+
+        if isinstance(func, tirx.PrimFunc):
+            sch = s_tir.Schedule(func)
+            sch.work_on("main")
+        elif isinstance(func, s_tir.Schedule):
+            sch = func
+
+        root_block = analysis.get_root_block(sch, sch.func_working_on)
+
+        if len(sch.get_child_blocks(root_block)) != 1:
+            return None
+
+        blk = sch.get_child_blocks(root_block)[0]
+        block_info = analysis.get_sblock_info(sch, blk)
+        if not (
+            (self.use_op_name and block_info.name == "te_layout_transform")
+            or (not self.use_op_name and block_info.is_layout_transform(sch))
+        ):
+            return None
+
+        read_buf, write_buf = (block_info.read_bufs(sch)[0], block_info.write_bufs(sch)[0])
+        lps = block_info.get_loops()
+        lpv_read, lpv_write = (
+            read_buf.assoc_lps[-1],
+            write_buf.assoc_lps[-1],
+        )
+
+        if lpv_read is None or lpv_write is None:
+            return None
+
+        vlen_read, vlen_write = read_buf.get_vecsize(), write_buf.get_vecsize()
+        local_cache = sch.get(lpv_read) != sch.get(lpv_write) or vlen_read != vlen_write
+        block_loops = [
+            lp
+            for lp in lps
+            if sch.get(lp) != sch.get(lpv_read) and sch.get(lp) != sch.get(lpv_write)
+        ]
+        vec_loops = (
+            [lpv_read, lpv_write] if sch.get(lpv_read) != sch.get(lpv_write) else (lpv_read,)
+        )
+        sch.reorder(*block_loops, *vec_loops)
+        if local_cache:
+            if sch.get(lpv_read) != sch.get(lpv_write):
+                blp_read, vlp_read = sch.split(
+                    lpv_read, [None, vlen_read], preserve_unit_iters=True
+                )
+                blp_write, vlp_write = sch.split(
+                    lpv_write, [None, vlen_write], preserve_unit_iters=True
+                )
+                sch.reorder(blp_read, blp_write, vlp_read, vlp_write)
+                block_loops += [blp_read, blp_write]
+                rblk = sch.cache_read(blk, 0, "local")
+                sch.compute_at(rblk, block_loops[-1], preserve_unit_loops=True)
+                sch.vectorize(sch.get_loops(rblk)[-1])
+                sch.vectorize(vlp_write)
+            else:
+                if vlen_read > vlen_write:
+                    read_lp, vec_lp = sch.split(blk, [None, vlen_write], preserve_unit_iters=True)
+                    rblk = sch.cache_read(blk, 0, "local")
+                    sch.compute_at(rblk, read_lp, preserve_unit_loops=True)
+                    sch.vectorize(sch.get_loops(rblk)[-1])
+                    sch.vectorize(vec_lp)
+                else:
+                    rblk = sch.cache_read(blk, 0, "local")
+                    sch.compute_at(rblk, block_loops[-1], preserve_unit_loops=True)
+                    _, vread_lp = sch.split(
+                        sch.get_loops(rblk)[-1], vlen_read, preserve_unit_iters=True
+                    )
+                    sch.vectorize(vread_lp)
+                    sch.vectorize(vlp_write)
+        else:
+            blp, vlp = sch.split(lpv_read, [None, vlen_read], preserve_unit_iters=True)
+            block_loops += [blp]
+            sch.vectorize(vlp)
+
+        b = sch.fuse(*block_loops)
+        tx_extent = min(sch.get(b).extent, 256)
+        candidates = [1, 2, 4, 8, 16, 32]
+        bx, tx = sch.split(b, [None, 256], preserve_unit_iters=True)
+        sch.bind(bx, "blockIdx.x")
+        sch.bind(tx, "threadIdx.x")
+        return sch

python/tvm/s_tir/dlight/adreno/pool.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.
+# pylint: disable=missing-docstring
+# ruff: noqa: F841
+"""Pool schedule rule for Adreno operators."""
+
+from tvm import s_tir, tirx
+from tvm.target import Target
+
+from .. import analysis
+from .base import AdrenoScheduleRule
+
+
+# pylint: disable=invalid-name, unused-variable
+class Pool2D(AdrenoScheduleRule):
+    def apply(  # pylint: disable=too-many-locals,missing-docstring
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> s_tir.Schedule:
+        sch = s_tir.Schedule(func)
+        root = sch.get_sblock(name="root", func_name="main")
+
+        blocks = sch.get_child_blocks(root)
+        blocks_names = [sch.get(blk).name_hint for blk in blocks]
+
+        if "adaptive_pool_sum" not in blocks_names and "pool_max" not in blocks_names:
+            return None
+
+        def schedule_pad(blk: s_tir.schedule.SBlockRV):
+            lps, veclp = sch.get_loops(blk)[:-1], sch.get_loops(blk)[-1]
+            sch.vectorize(veclp)
+            b = sch.fuse(*lps)
+            tx_extent = min(int(sch.get(b).extent) & ~int(sch.get(b).extent - 1), 256)
+            bx, tx = sch.split(b, [None, tx_extent])
+            sch.bind(bx, "blockIdx.x")
+            sch.bind(tx, "threadIdx.x")
+
+        def schedule_max_pool(blk: s_tir.schedule.SBlockRV):
+            block_info = analysis.get_sblock_info(sch, blk)
+            iters_kind = "".join([_iter.kind for _iter in block_info.iters])
+            if iters_kind != "SSSSSRR":
+                return None
+
+            lps = sch.get_loops(blk)
+            block_lps, vec_lp, red_lps = lps[:4], lps[4], lps[5:]
+            write_blk = sch.cache_write(blk, 0, "local")
+            sch.reverse_compute_at(write_blk, vec_lp)
+            b = sch.fuse(*block_lps)
+            tx_extent = min(int(sch.get(b).extent) & ~int(sch.get(b).extent - 1), 256)
+            bx, tx = sch.split(b, [None, tx_extent])
+            sch.bind(bx, "blockIdx.x")
+            sch.bind(tx, "threadIdx.x")
+            sch.vectorize(vec_lp)
+
+            return True
+
+        passed_reduction = False
+        for blk in blocks:
+            if sch.get(blk).name_hint == "pad_temp":
+                schedule_pad(blk)
+            elif (
+                sch.get(blk).name_hint == "adaptive_pool_sum"
+                or sch.get(blk).name_hint == "pool_max"
+            ):
+                ok = schedule_max_pool(blk)
+                if not ok:
+                    return None
+                passed_reduction = True
+            else:
+                try:
+                    if passed_reduction:
+                        sch.reverse_compute_inline(blk)
+                    else:
+                        sch.compute_inline(blk)
+                except Exception:  # pylint: disable=broad-except
+                    pass
+        return sch

python/tvm/s_tir/dlight/adreno/utils.py

@@ -0,0 +1,90 @@
+# 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.
+"""Utilis for Adreno operators."""
+
+# pylint: disable=import-outside-toplevel, unused-argument, invalid-name, missing-function-docstring
+
+from tvm import s_tir
+from tvm.target import Target
+
+from ..analysis import SBlockInfo
+
+
+def get_texture_storage(block_info: SBlockInfo):
+    """
+    Returns the texture layout acceptable for the shape
+
+    Parameters
+    ----------
+    shape: array
+        Shape of the tensor to be packed to texture
+    """
+    # certain limitation of the Qualcomm devices. Subject to be determined for certain device
+    # individually, but until we have access to remote device during compilation, we have to
+    # define it uniformly for all target devices
+    # spatial_limit = 16384, depth_limit = 2048
+    # TODO: Check Write Bufs.
+    shape = block_info.write_bufs[0].buf_region.buffer.shape
+
+    spatial_limit = Target.current().attrs["texture_spatial_limit"]
+    depth_limit = Target.current().attrs["texture_depth_limit"]
+
+    if len(shape) > 4:
+        if shape[0] < spatial_limit and shape[1] * shape[2] * shape[3] < spatial_limit:
+            return "global.texture-weight"
+        elif shape[0] < depth_limit and shape[2] * shape[3] < spatial_limit:
+            return "global.texture-nhwc"
+        elif (
+            shape[0] * shape[1] < depth_limit
+            and shape[2] < spatial_limit
+            and shape[3] < spatial_limit
+        ):
+            return "global.texture"
+    elif len(shape) > 3:
+        if shape[0] < spatial_limit and shape[1] * shape[2] < spatial_limit:
+            return "global.texture-weight"
+        elif shape[0] < depth_limit and shape[1] < spatial_limit and shape[2] < spatial_limit:
+            return "global.texture"
+    elif len(shape) == 3:
+        if shape[0] < spatial_limit and shape[1] < spatial_limit:
+            return "global.texture-weight"
+
+    return "global"
+
+
+def schedule_inline_blocks(
+    sch: s_tir.Schedule, blocks: list[s_tir.schedule.SBlockRV] | None = None
+):
+    from .fallback import Fallback
+
+    return Fallback.schedule_inline_blocks(sch, blocks)
+
+
+def schedule_default(sch, blocks: list[s_tir.schedule.SBlockRV] | None = None):
+    from .fallback import Fallback
+
+    ret = []
+    for blk in blocks:
+        ret.append(Fallback.schedule_default(sch, blk))
+
+    return ret
+
+
+def schedule_fallback(sch, blk):
+    from .fallback import Fallback
+
+    return Fallback.schedule_fallback(sch)

python/tvm/s_tir/dlight/analysis/__init__.py

@@ -0,0 +1,35 @@
+# 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.
+"""Base infra"""
+
+from .common_analysis import (
+    SBlockInfo,
+    IterInfo,
+    collect_block_iter_vars_used_in_access_region,
+    collect_vars_used_in_prim_expr,
+    detect_dominant_read,
+    is_broadcast_epilogue,
+    normalize_prim_func,
+    get_root_block,
+    get_sblock_info,
+    get_max_shared_memory_per_block,
+)
+from .gemv import (
+    is_gemv,
+    normalize,
+)

python/tvm/s_tir/dlight/analysis/common_analysis.py

@@ -0,0 +1,460 @@
+# 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-function-docstring, missing-class-docstring
+# pylint: disable=unused-argument, unused-variable
+"""Analysis on TIR blocks, loops and functions."""
+
+import logging
+from collections import namedtuple
+from typing import Literal
+
+from tvm_ffi import get_global_func
+
+from tvm import ir, s_tir, tirx
+from tvm.s_tir import Schedule
+from tvm.s_tir.schedule import SBlockRV
+from tvm.target.target import Target
+
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+
+class IterInfo:
+    """Information about a loop/iter var."""
+
+    kind: Literal["S", "R", "O"]
+    var: tirx.Var
+    _dom: tirx.Expr
+    loop_rv: s_tir.schedule.LoopRV
+
+    def __init__(
+        self,
+        kind: Literal["S", "R", "O"],
+        var: tirx.Var,
+        dom: tirx.Expr,
+        loop_rv: s_tir.schedule.LoopRV,
+    ):
+        """Construct an IterInfo object."""
+        self.kind = kind
+        self.var = var
+        self._dom = dom
+        self.loop_rv = loop_rv
+
+    @property
+    def dom(self) -> int | tirx.Expr:
+        """The iteration domain of the loop."""
+        return int(self._dom) if isinstance(self._dom, tirx.IntImm) else self._dom
+
+    def __str__(self) -> str:
+        return f'Iter("{self.kind}", {self.dom})'
+
+    def __repr__(self) -> str:
+        return str(self)
+
+
+get_sblockrealize = get_global_func("s_tir.schedule.GetSBlockRealize")
+# BufferIndex Types
+Index = namedtuple("Index", ["sub"])  # c
+RemIndex = namedtuple("RemIndex", ["sub", "div"])  # c%len
+DivIndex = namedtuple("DivIndex", ["sub", "div"])  # c//len
+MergeIndex = namedtuple("MulIndex", ["dom", "mul", "sub"])  # co*len + cb
+BufIndex = list[Index | RemIndex | DivIndex | MergeIndex | None]
+
+
+class BufferInfo:
+    "Information about Buffer. Provides useful analysis"
+
+    buf_region: tirx.BufferRegion
+    shape: tuple[int]
+    assoc_lps: list[s_tir.schedule.LoopRV | None]
+    assoc_lps_info: list[tirx.For | None]
+
+    def __init__(
+        self,
+        sch: s_tir.Schedule,
+        block_rv: s_tir.schedule.SBlockRV,
+        buf_region: tirx.BufferRegion,
+        lps: list[s_tir.schedule.LoopRV] | None,
+    ):
+        block = sch.get(block_rv)
+        if lps is None:
+            lps = sch.get_loops(block_rv)
+        loops = [sch.get(lp) for lp in lps]
+        iter_vars = [Var.var for Var in block.iter_vars]
+        iter_values = get_sblockrealize(sch, block_rv).iter_values
+        lpvar_lp = dict([loop.loop_var, lp] for loop, lp in zip(loops, lps))
+        var_lp = dict(zip(iter_vars, [lpvar_lp.get(val, None) for val in iter_values]))
+
+        def extract_index_types(buf: tirx.BufferRegion) -> BufIndex:
+            buf_index = []
+            for expr in buf.region:
+                expr = expr.min
+                dim = None
+                if isinstance(expr, tirx.expr.Add) and isinstance(expr.b, tirx.expr.Var):
+                    var_add = expr.b
+                    if (
+                        isinstance(expr, tirx.expr.Mul)
+                        and isinstance(expr.a, tirx.expr.Var)
+                        and isinstance(expr.b, tirx.expr.IntImm)
+                    ):
+                        mul = expr.b
+                        var_mul = expr.a
+                        dim = MergeIndex(var_mul, mul, var_add)
+                elif (
+                    isinstance(expr, tirx.expr.FloorMod)
+                    and isinstance(expr.a, tirx.expr.Var)
+                    and isinstance(expr.b, tirx.expr.IntImm)
+                ):
+                    dim = RemIndex(expr.a, expr.b)
+                elif (
+                    isinstance(expr, tirx.expr.FloorDiv)
+                    and isinstance(expr.a, tirx.expr.Var)
+                    and isinstance(expr.b, tirx.expr.IntImm)
+                ):
+                    dim = DivIndex(expr.a, expr.b)
+                elif isinstance(expr, tirx.expr.Var):
+                    dim = Index(expr)
+                buf_index.append(dim)
+            return buf_index
+
+        indexes = extract_index_types(buf_region)
+        assoc_lps = [
+            (
+                var_lp.get(getattr(idx, "sub"), None)
+                if not isinstance(idx, DivIndex) and idx is not None
+                else None
+            )
+            for idx in indexes
+        ]
+
+        self.buf_region = buf_region
+        self.assoc_lps = assoc_lps
+        self.assoc_lps_info = [(sch.get(lp) if lp is not None else None) for lp in assoc_lps]
+        self.shape = buf_region.buffer.shape
+
+    def get_scope(self) -> str:
+        return self.buf_region.buffer.scope()
+
+    def get_vecsize(self, buf_index: int = 0, vbits: int = 128):
+        if self.assoc_lps_info[-1] is None:
+            return None
+
+        vlp_extent = int(self.assoc_lps_info[-1].extent) & ~(
+            int(self.assoc_lps_info[-1].extent) - 1
+        )
+        vbuf_extent = int(self.shape[-1]) & ~(int(self.shape[-1]) - 1)
+
+        return min(vlp_extent, vbuf_extent, vbits // self.buf_region.buffer.dtype.dtype.bits)
+
+    def __str__(self) -> str:
+        return f"BufferInfo({self.buf_region})"
+
+    def __repr__(self) -> str:
+        return str(self)
+
+
+class SBlockInfo:
+    """Information about a TIR block."""
+
+    name: str
+    iters: list[IterInfo]
+    block_rv: s_tir.schedule.SBlockRV
+    _reduction_block: bool
+
+    def __init__(
+        self,
+        name: str,
+        iters: list[IterInfo],
+        block_rv: s_tir.schedule.SBlockRV,
+        reduction_block: bool = False,
+    ):
+        """Construct a SBlockInfo object."""
+        self.name = name
+        self.block_rv = block_rv
+        self.iters = iters
+        self._reduction_block = reduction_block
+
+    def dom(self) -> list[int | tirx.Expr]:
+        """The iteration domain of the block."""
+        return [i.dom for i in self.iters]
+
+    def read_bufs(self, sch: s_tir.Schedule) -> list[BufferInfo]:
+        block_stmt = sch.get(self.block_rv)
+        lps = sch.get_loops(self.block_rv)
+        return [BufferInfo(sch, self.block_rv, buf, lps) for buf in block_stmt.reads]
+
+    def write_bufs(self, sch: s_tir.Schedule) -> list[BufferInfo]:
+        block_stmt = sch.get(self.block_rv)
+        lps = sch.get_loops(self.block_rv)
+        return [BufferInfo(sch, self.block_rv, buf, lps) for buf in block_stmt.writes]
+
+    def dom_kind(self) -> str:
+        """The iteration domain kind of the block, for example, SSSS, SSSR."""
+        return "".join(i.kind for i in self.iters)
+
+    def is_injective(self) -> bool:
+        """Whether the SBlock is injective, i.e. all its iteration domains are injective."""
+        return all(k == "S" for k in self.dom_kind())
+
+    def is_elementwise(self, sch: s_tir.Schedule) -> bool:
+        """Whether the SBlock is elementwise, i.e. trivial mapping between read/write region"""
+
+        def _check_unit_var_range(dom: ir.Range, var: tirx.Var) -> bool:
+            return dom.min.same_as(var) and dom.extent == 1
+
+        if not self.is_injective():
+            return False
+        block = sch.get(self.block_rv)
+        if len(block.reads) != 1 or len(block.writes) != 1:
+            return False
+        r_region = block.reads[0].region
+        w_region = block.writes[0].region
+        if len(r_region) != len(w_region):
+            return False
+        for var, r_dom, w_dom in zip(block.iter_vars, r_region, w_region):
+            if not _check_unit_var_range(r_dom, var) or not _check_unit_var_range(w_dom, var):
+                return False
+        return True
+
+    def get_loops(self) -> list[s_tir.schedule.LoopRV]:
+        return [iter_info.loop_rv for iter_info in self.iters]
+
+    def is_reduction(self) -> bool:
+        """Whether the SBlock is a reduction workload."""
+        # TODO(@junrushao): distinguish GEMV and reduction
+        return self._reduction_block
+
+    def is_layout_transform(self, sch: s_tir.Schedule) -> bool:
+        """Whether the SBlock can be considered having a Layout Transform Pattern"""
+        return (
+            all(k == "S" for k in self.dom_kind())
+            and len(self.write_bufs(sch)) == 1
+            and len(self.read_bufs(sch)) == 1
+            and not self.is_elementwise(sch)
+            and not get_global_func("s_tir.schedule.HasIfThenElse")(sch.get(self.block_rv))
+        )
+
+    def is_data_pad(self, sch: s_tir.Schedule) -> bool:
+        """Whether the SBlock can be considered having a data pad pattern"""
+        return (
+            all(k == "S" for k in self.dom_kind())
+            and len(self.write_bufs(sch)) == 1
+            and len(self.read_bufs(sch)) == 1
+            and not self.is_elementwise(sch)
+            and len(self.write_bufs(sch)[0].buf_region.region)
+            == len(self.read_bufs(sch)[0].buf_region.region)
+            and get_global_func("s_tir.schedule.HasIfThenElse")(sch.get(self.block_rv))
+        )
+
+    def is_convolution(self) -> bool:
+        """Whether a SBlock can be considered having Convolution Pattern"""
+        raise NotImplementedError
+
+    def is_pool(self) -> bool:
+        """Whether a SBlock can be considered having Pooling Pattern"""
+        raise NotImplementedError
+
+    def is_gemv(self) -> bool:
+        """Whether the SBlock is a GEMV workload."""
+        raise NotImplementedError
+
+    def is_gemm(self) -> bool:
+        """Whether the SBlock is a GEMM workload."""
+        raise NotImplementedError
+
+    def __str__(self) -> str:
+        return f'SBlockInfo("{self.name}", "{self.dom_kind()}", {self.dom()})'
+
+    def __repr__(self) -> str:
+        return str(self)
+
+
+_normalize_prim_func = get_global_func("s_tir.schedule.NormalizePrimFunc")
+
+
+def normalize_prim_func(sch: s_tir.Schedule) -> list[SBlockInfo] | None:
+    """Normalize the primfunc to normal form"""
+    try:
+        result = _normalize_prim_func(sch)
+        if result is None:
+            return None
+    except Exception:  # pylint: disable=broad-except
+        return None
+
+    def _iter_kind(i: tirx.IterVar) -> str:
+        return {
+            tirx.IterVar.DataPar: "S",
+            tirx.IterVar.CommReduce: "R",
+        }.get(i.iter_type, "O")
+
+    blocks: list[SBlockInfo] = []
+    for block, loops, iters, is_reduction in zip(*result):
+        blocks.append(
+            SBlockInfo(
+                name=sch.get(block).name_hint,
+                iters=[
+                    IterInfo(
+                        kind=_iter_kind(iter),  # type: ignore
+                        var=iter.var,
+                        dom=iter.dom.extent,
+                        loop_rv=loop,
+                    )
+                    for loop, iter in zip(loops, iters)
+                ],
+                block_rv=block,
+                reduction_block=is_reduction,
+            )
+        )
+    return blocks
+
+
+def get_sblock_info(sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV) -> SBlockInfo:
+    def _iter_kind(loop: tirx.IterVar) -> str:
+        return {tirx.IterVar.DataPar: "S", tirx.IterVar.CommReduce: "R"}.get(loop.iter_type, "O")
+
+    def _is_reduction_block(block: s_tir.schedule.SBlockRV):
+        for iter_var in sch.get(block).iter_vars:
+            if _iter_kind(iter_var) == "R":
+                return True
+        return False
+
+    return SBlockInfo(
+        name=sch.get(block).name_hint,
+        iters=[
+            IterInfo(
+                kind=_iter_kind(iter_var),
+                var=iter_var.var,
+                dom=iter_var.dom.extent,
+                loop_rv=loop_rv,
+            )
+            for loop_rv, iter_var in zip(sch.get_loops(block), sch.get(block).iter_vars)
+        ],
+        block_rv=block,
+        reduction_block=_is_reduction_block(block),
+    )
+
+
+def _assert_gpu_target(target: Target):
+    if "gpu" not in target.keys:
+        raise ValueError(f"Expect a GPU target, but got {target}")
+
+
+def get_max_threads_per_block(target: Target) -> int:
+    _assert_gpu_target(target)
+    max_threads_per_block = None
+    for name in ["max_threads_per_block", "max_num_threads"]:
+        if max_threads_per_block is None:
+            max_threads_per_block = target.attrs.get(name, None)
+    if max_threads_per_block is None:
+        max_threads_per_block = 64
+    return int(max_threads_per_block)
+
+
+TARGET_KIND_TO_DEFAULT_MAX_SMEM = {
+    "cuda": 49152,
+    "rocm": 65536,
+    "metal": 32768,
+    "opencl": 16384,
+    "vulkan": 16384,
+}
+
+
+def get_max_shared_memory_per_block(target: Target) -> int:
+    _assert_gpu_target(target)
+    max_shared_memory_per_block = target.attrs.get("max_shared_memory_per_block", None)
+    if max_shared_memory_per_block is not None:
+        return int(max_shared_memory_per_block)
+
+    # Layered fallback strategy for targets that do not carry this attribute
+    # 1) Use explicit target attrs provided (handled above).
+    # 2) Fall back to backend defaults matching target-kind defaults/tag defaults.
+    # 3) Use a conservative GPU default as last resort.
+    default_smem = TARGET_KIND_TO_DEFAULT_MAX_SMEM.get(target.kind.name, 16384)
+    logger.warning(
+        "Target %s missing 'max_shared_memory_per_block'; using %d bytes.",
+        target.kind.name,
+        default_smem,
+    )
+    return int(default_smem)
+
+
+def get_root_block(sch: Schedule, func_name: str = "main") -> SBlockRV:
+    try:
+        block = sch.mod[func_name].body.block
+    except Exception:
+        raise ValueError(
+            f"The function body is expected to be the root block, but got:\n"
+            f"{sch.mod[func_name].body}"
+        )
+    return sch.get_sblock(block.name_hint)
+
+
+def collect_block_iter_vars_used_in_access_region(
+    block: tirx.SBlock, region: list[ir.Range]
+) -> set[tirx.Var]:
+    """Collect the block iter variables used in the access region of a buffer region."""
+    tir_vars = set()
+    for expr in region:
+        assert expr.extent == 1
+        tir_vars |= collect_vars_used_in_prim_expr(expr.min)
+    tir_vars &= set(iter_var.var for iter_var in block.iter_vars)
+    return tir_vars
+
+
+def collect_vars_used_in_prim_expr(expr: tirx.Expr) -> set[tirx.Var]:
+    """Collect the variables used in the Expr."""
+    tir_vars = set()
+
+    def _collect_tir_var(expr):
+        if isinstance(expr, tirx.Var):
+            tir_vars.add(expr)
+
+    tirx.stmt_functor.post_order_visit(expr, _collect_tir_var)
+    return tir_vars
+
+
+def detect_dominant_read(block: tirx.SBlock) -> tirx.Expr:
+    """Detect the dominant read indices in the block."""
+    dominant_read = None
+    num_read_iters = -1
+    for buffer_region in block.reads:
+        tir_vars = collect_block_iter_vars_used_in_access_region(block, buffer_region.region)
+        if num_read_iters < len(tir_vars):
+            num_read_iters = len(tir_vars)
+            dominant_read = buffer_region
+    assert dominant_read is not None
+    (result,) = dominant_read.buffer.offset_of([e.min for e in dominant_read.region])
+    return result
+
+
+def is_broadcast_epilogue(
+    sch: s_tir.Schedule,
+    block: s_tir.schedule.SBlockRV,
+    epilogue: s_tir.schedule.SBlockRV,
+) -> bool:
+    """Check if the epilogue block is a broadcast pattern"""
+    write_buffers = {r.buffer for r in sch.get(block).writes}
+    epilogue_iters = {i.var: i for i in sch.get(epilogue).iter_vars if i.dom != 1}
+    for buffer_region in sch.get(epilogue).reads:
+        if buffer_region.buffer not in write_buffers:
+            continue
+        tir_vars = collect_block_iter_vars_used_in_access_region(
+            sch.get(epilogue), buffer_region.region
+        )
+        if len(tir_vars) < len(epilogue_iters):
+            return True
+    return False

python/tvm/s_tir/dlight/analysis/gemv.py

@@ -0,0 +1,163 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Analysis for GEMV."""
+
+import tvm_ffi
+
+from tvm import arith, s_tir, tirx
+
+from .common_analysis import (
+    SBlockInfo,
+    collect_block_iter_vars_used_in_access_region,
+    collect_vars_used_in_prim_expr,
+    detect_dominant_read,
+)
+
+
+def get_reduction_expr(block: tirx.SBlock) -> tirx.Expr | None:
+    """Extracts the reduction expression from a TIR block.
+
+    This function checks whether the given TIR block follows a reduction pattern
+    of the form `X[...] = X[...] + Y` and returns `Y` as the reduction expression.
+
+    Parameters:
+    ----------
+    block : tirx.SBlock
+        The TIR block to analyze.
+
+    Returns:
+    -------
+    Optional[tirx.Expr]
+        The reduction expression (`Y`) if detected, otherwise None.
+    """
+
+    buffer_store = block.body
+    if not isinstance(buffer_store, tirx.BufferStore):
+        return None
+    if not isinstance(buffer_store.value, tirx.Add):
+        return None
+    if not tvm_ffi.structural_equal(
+        buffer_store.value.a,
+        tirx.BufferLoad(buffer_store.buffer, block.body.indices),
+        map_free_vars=True,
+    ):
+        return None
+    return buffer_store.value.b
+
+
+def is_gemv(sch: s_tir.Schedule, block_info: SBlockInfo) -> list[tirx.Buffer] | None:
+    """Check if the block is a GEMV.
+
+    Parameters
+    ----------
+
+    sch : s_tir.Schedule
+        The schedule
+
+    block_info : SBlockInfo
+        The block info to be checked
+
+
+    Returns
+    -------
+    ret : Optional[List[tirx.Buffer]]
+        The vector buffers used in the GEMV if it is a GEMV, otherwise None.
+    """
+    block = block_info.block_rv
+    block_stmt = sch.get(block)
+    conditions = []
+    conditions.append(block_info.is_reduction())
+    conditions.append(len(block_stmt.reads) >= 2)
+    conditions.append(len(block_stmt.writes) == 1)
+    conditions.append(get_reduction_expr(block_stmt) is not None)
+    conditions.append(
+        len(collect_block_iter_vars_used_in_access_region(block_stmt, block_stmt.writes[0].region))
+        > 0
+    )
+    if not all(conditions):
+        return None
+
+    iter_num = len(block_stmt.iter_vars)
+    ret = [
+        read.buffer
+        for read in block_stmt.reads
+        if len(collect_block_iter_vars_used_in_access_region(block_stmt, read.region)) < iter_num
+        and len(collect_block_iter_vars_used_in_access_region(block_stmt, read.region)) > 0
+    ]
+    return ret if 0 < len(ret) < len(block_stmt.reads) else None
+
+
+def normalize(
+    sch: s_tir.Schedule,
+    block_info: SBlockInfo,
+) -> bool | None:
+    """Normalize the main block."""
+    block_stmt: tirx.SBlock = sch.get(block_info.block_rv)
+    access = arith.normalize_to_iter_sum(
+        detect_dominant_read(block_stmt),
+        input_iters={i.var: i.dom for i in block_stmt.iter_vars},
+    )
+    buffers_use_vars = [
+        collect_block_iter_vars_used_in_access_region(block_stmt, buf.region)
+        for buf in block_stmt.writes
+    ]
+    buffers_use_vars.extend(
+        [
+            collect_block_iter_vars_used_in_access_region(block_stmt, buf.region)
+            for buf in block_stmt.reads
+        ]
+    )
+    if collect_vars_used_in_prim_expr(access.base) & set(
+        iter_var.var for iter_var in block_stmt.iter_vars
+    ):
+        return None
+    iter_to_info = {i.var: i for i in block_info.iters}
+    batch_loops, s_loops, r_loops, c_loops = [], [], [], []
+    inner_axis = access.args[-1].source.source
+    is_inner_reduction = iter_to_info[inner_axis].kind == "R"
+
+    for split_expr in access.args:
+        var = split_expr.source.source
+        info = iter_to_info.get(var)
+        loop = info.loop_rv
+        is_reduction = info.kind == "R"
+        if split_expr.lower_factor > 1:
+            if c_loops:
+                return None
+            loop, c_loop = sch.split(loop, factors=[None, split_expr.lower_factor])
+            # we only support the reduction dim being grouped atm
+            if not is_reduction:
+                return None
+            c_loops.append(c_loop)
+        if is_reduction:
+            r_loops.append(loop)
+        elif all([var in buf_vars for buf_vars in buffers_use_vars]):
+            batch_loops.append(loop)
+        else:
+            s_loops.append(loop)
+
+    assert s_loops
+    assert r_loops
+    if not c_loops:
+        c_loops = [sch.add_unit_loop(block_info.block_rv)]
+    if not batch_loops:
+        batch_loops = [sch.add_unit_loop(block_info.block_rv)]
+    sch.reorder(*batch_loops, *s_loops, *r_loops, *c_loops)
+    sch.fuse(*batch_loops)
+    sch.fuse(*s_loops)
+    sch.fuse(*r_loops)
+    return is_inner_reduction

python/tvm/s_tir/dlight/base/__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.
+"""Base infra"""
+
+from .common_schedules import try_inline, try_inline_contiguous_spatial
+from .schedule_rule import ScheduleRule
+from .transform import ApplyDefaultSchedule
+from .utils import (
+    auto_vectorize,
+    get_bytes,
+    get_extent,
+    max_threads_per_block,
+    suggest_threads_per_block,
+)

python/tvm/s_tir/dlight/base/common_schedules.py

@@ -0,0 +1,100 @@
+# 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: E722
+"""Common schedule strategies for TIR."""
+
+from collections.abc import Callable
+
+from tvm import s_tir
+
+from ..analysis import SBlockInfo
+
+
+def try_inline(
+    sch: s_tir.Schedule,
+    blocks: list[SBlockInfo],
+) -> list[SBlockInfo]:
+    """Try to inline as many blocks as possible, and return the remaining blocks.
+
+    Parameters
+    ----------
+    sch : s_tir.Schedule
+        The TIR schedule used to inline blocks.
+    blocks : List[SBlockInfo]
+        The blocks to be inlined.
+
+    Returns
+    -------
+    remaining : List[SBlockInfo]
+        The remaining blocks that cannot be inlined.
+    """
+
+    def _trial(func: Callable):
+        for i, block in enumerate(blocks):
+            try:
+                func(block.block_rv)
+            except:  # pylint: disable=bare-except
+                continue
+            return i
+        return None
+
+    while True:
+        i = _trial(sch.compute_inline)
+        if i is None:
+            i = _trial(sch.reverse_compute_inline)
+        if i is None:
+            break
+        blocks.pop(i)
+    return blocks
+
+
+def try_inline_contiguous_spatial(
+    sch: s_tir.Schedule,
+    block_infos: list[SBlockInfo],
+) -> list[SBlockInfo]:
+    """Try to inline contiguous spatial blocks in a schedule
+
+    Parameters
+    ----------
+    sch : s_tir.Schedule
+        The TIR schedule used to inline blocks.
+    block_infos : List[SBlockInfo]
+        The blocks to be try.
+
+    Returns
+    -------
+    remaining : List[SBlockInfo]
+        The remaining blocks that cannot be inlined.
+    """
+
+    if block_infos is None:
+        return None
+    results = []
+    spatial_blocks = []
+    block: SBlockInfo
+    for block in block_infos:
+        if block.is_injective():
+            spatial_blocks.append(block)
+        elif spatial_blocks:
+            results.extend(try_inline(sch, spatial_blocks))
+            results.append(block)
+            spatial_blocks = []
+        else:
+            results.append(block)
+    if spatial_blocks:
+        results.extend(try_inline(sch, spatial_blocks))
+    return results

python/tvm/s_tir/dlight/base/schedule_rule.py

@@ -0,0 +1,121 @@
+# 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 lightweight wrapper on an arbitrary function that can be used to schedule a TIR PrimFunc."""
+
+from collections.abc import Callable
+
+from tvm import s_tir, tirx
+from tvm.target import Target
+
+
+class ScheduleRule:  # pylint: disable=too-few-public-methods
+    """A thin wrapper on an arbitrary function that can be used to schedule a TIR PrimFunc.
+
+    Given a PrimFunc, a target, and a tunable flag, the apply method of a ScheduleRule
+    returns either a Schedule, a list of Schedules, or None, where None means that the rule
+    is not applicable to the given PrimFunc. If the tunable flag is True, the ScheduleRule is
+    allowed to return either a Schedule or a list of Schedules, and the Schedules are allowed to
+    contain tunable instructions. If the tunable flag is False, the ScheduleRule is only allowed to
+    return a Schedule, and the Schedule is not allowed to contain tunable instructions.
+    """
+
+    def apply(
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        tunable: bool,
+    ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+        """Apply the ScheduleRule to the given PrimFunc.
+
+        Parameters
+        ----------
+        func : tirx.PrimFunc
+            The PrimFunc to apply the ScheduleRule to.
+        target : Target
+            The compilation target the schedule is supposed to be built for.
+        tunable : bool
+            Whether the schedule is allowed to contain tunable instructions.
+
+        Returns
+        -------
+        results : Union[None, s_tir.Schedule, List[s_tir.Schedule]]
+            Either a Schedule, a list of Schedules, or None, where None means that the rule
+            is not applicable to the given PrimFunc.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def from_callable(
+        name,
+    ) -> Callable[
+        [
+            Callable[
+                [tirx.PrimFunc, Target, bool],
+                None | s_tir.Schedule | list[s_tir.Schedule],
+            ],
+        ],
+        "ScheduleRule",
+    ]:
+        """Create a ScheduleRule from a callable.
+
+        Parameters
+        ----------
+        name : str
+
+        Returns
+        -------
+        decorator : Callable
+            A decorator that takes a callable and returns a ScheduleRule.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            @ScheduleRule.from_callable("MyRule")
+            def my_rule(func: tirx.PrimFunc, target: Target, tunable: bool) -> Union[None, Schedule]
+                # Do something with func and target
+        """
+
+        def decorator(f) -> "ScheduleRule":  # pylint: disable=invalid-name
+            class _Rule(ScheduleRule):
+                def apply(
+                    self,
+                    func: tirx.PrimFunc,
+                    target: Target,
+                    tunable: bool,
+                ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+                    return f(func, target, tunable)
+
+            _Rule.__name__ = name
+            return _Rule()
+
+        return decorator
+
+    def is_target_available(self, target: Target) -> bool:  # pylint: disable=unused-argument
+        """Check whether the rule is available for the given target.
+
+        Parameters
+        ----------
+        target : Target
+            The compilation target the schedule is supposed to be built for.
+
+        Returns
+        -------
+        available : bool
+            Whether the rule is available for the given target.
+        """
+        return True

python/tvm/s_tir/dlight/base/transform.py

@@ -0,0 +1,94 @@
+# 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.
+"""
+Apply ScheduleRules onto an IRModule to generate default schedules without tuning,
+or a space for MetaSchedule tuning
+"""
+
+from tvm import s_tir, tirx
+from tvm.ir import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.target import Target
+
+from .schedule_rule import ScheduleRule
+
+
+def _is_scheduled(func: tirx.PrimFunc) -> bool:
+    if not isinstance(func, tirx.PrimFunc):
+        return False
+    if "tirx.is_scheduled" not in func.attrs:
+        return False
+    return func.attrs["tirx.is_scheduled"] == 1
+
+
+def _get_target(func: tirx.PrimFunc) -> Target:
+    target = func.attrs.get("target")
+    if target is None:
+        return Target.current(allow_none=False)
+    else:
+        return target
+
+
+@module_pass(opt_level=0, name="ApplyDefaultSchedule")
+class ApplyDefaultSchedule:  # pylint: disable=too-few-public-methods
+    """A IRModule pass that applies a list of ScheduleRules to all PrimFuncs in the module."""
+
+    def __init__(self, *rules: ScheduleRule):
+        """Construct a new ApplyDefaultSchedule pass.
+
+        Parameters
+        ----------
+        *rules : ScheduleRule
+            The ScheduleRules to apply to all PrimFuncs in the module.
+        """
+        self.rules = list(rules)
+
+    def transform_module(  # pylint: disable=missing-function-docstring
+        self,
+        mod: IRModule,
+        _: PassContext,
+    ) -> IRModule:
+        updated_functions = {}
+        for g_var, func in mod.functions_items():
+            if isinstance(func, tirx.PrimFunc) and not _is_scheduled(func):
+                target = _get_target(func)
+
+                sch = _apply_rules(func, target, self.rules, tunable=False)
+                if sch is not None:
+                    assert len(sch) == 1
+                    updated_functions[g_var] = (
+                        sch[0].mod["main"].with_attr("tirx.is_scheduled", True)
+                    )
+        for g_var, func in updated_functions.items():
+            mod[g_var] = func
+        return mod
+
+
+def _apply_rules(
+    func: tirx.PrimFunc,
+    target: Target,
+    rules: list[ScheduleRule],
+    tunable: bool,
+) -> list[s_tir.Schedule] | None:
+    for rule in rules:
+        space = rule.apply(func, target, tunable)
+        if space is None:
+            continue
+        if isinstance(space, s_tir.Schedule):
+            space = [space]
+        return space
+    return None

python/tvm/s_tir/dlight/base/utils.py

@@ -0,0 +1,115 @@
+# 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
+"""Utility methods for generic GPU."""
+
+from tvm import DataType, s_tir, tirx
+from tvm.ir import PrimType
+from tvm.target import Target
+
+
+def get_bytes(dtype: DataType | PrimType | str) -> int:
+    if isinstance(dtype, PrimType):
+        dtype = dtype.dtype
+    if isinstance(dtype, str):
+        dtype = DataType(dtype)
+    return dtype.itemsize
+
+
+def get_extent(sch: s_tir.Schedule, loop_rv: s_tir.schedule.LoopRV):
+    loop: tirx.For = sch.get(loop_rv)
+    return loop.extent.value if isinstance(loop.extent, tirx.IntImm) else loop.extent
+
+
+def auto_vectorize(sch: s_tir.Schedule, loop: s_tir.schedule.LoopRV, max_vec: int):
+    """Auto vectorize the loop."""
+    extent = get_extent(sch, loop)
+    if not isinstance(extent, int):
+        return
+    v = loop if extent <= max_vec else sch.split(loop, factors=[None, max_vec])[-1]
+    sch.vectorize(v)
+
+
+def max_threads_per_block(target: Target) -> int:
+    """Get the maximum number of threads per block for a given target.
+
+    Parameters
+    ----------
+    target : Target
+        The target to get the maximum number of threads per block for.
+
+    Returns
+    -------
+    max_threads_per_block : int
+        The maximum number of threads per block for the given target.
+    """
+    for name in ["max_threads_per_block", "max_num_threads"]:
+        result = target.attrs.get(name, None)
+        if result is not None:
+            return result
+    if target.kind.name == "cuda":
+        return 1024
+    return 256
+
+
+def suggest_threads_per_block(
+    target: Target,
+    loops: list[tirx.For],
+    max_threads_for_dynamic_loop: int = 32,
+) -> list[int]:
+    if target.kind.name == "cuda":
+        threads = 1024
+    elif target.kind.name == "rocm":
+        threads = 256
+    elif target.kind.name == "metal":
+        threads = 256
+    elif target.kind.name == "opencl":
+        threads = 256
+    else:
+        threads = 64
+    results: list[int | None] = []
+    dynamic: list[int] = []
+    for i, loop in enumerate(loops):
+        loop_extent = loop.extent
+        if isinstance(loop_extent, tirx.IntImm):
+            loop_extent = loop_extent.value
+            extent = 1
+            while extent <= loop_extent and extent <= threads:
+                extent *= 2
+            extent //= 2
+            assert extent >= 1
+            assert threads % extent == 0
+            threads //= extent
+            results.append(extent)
+        else:
+            results.append(None)
+            dynamic.append(i)
+
+    for i in dynamic:
+        extent = 1
+        while extent <= max_threads_for_dynamic_loop and extent <= threads:
+            extent *= 2
+        extent //= 2
+        assert extent >= 1
+        assert threads % extent == 0
+        threads //= extent
+        results[i] = extent
+
+    if dynamic:
+        results[dynamic[0]] *= threads
+
+    return results

python/tvm/s_tir/dlight/benchmark/__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.
+"""Benchmarking dynamic shape workloads"""
+
+from .bench import benchmark, benchmark_prim_func, benchmark_relax_func
+from .extract import (
+    extract_prim_func,
+    extract_from_relax,
+    extract_func_info_from_prim_func,
+    extract_all_func_info_from_relax,
+)

python/tvm/s_tir/dlight/benchmark/bench.py

@@ -0,0 +1,313 @@
+# 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.
+"""Extract self-contained benchmarking scripts for dynamic shape workloads"""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Optional
+
+import tvm
+from tvm import relax
+from tvm.ir import IRModule
+from tvm.s_tir.meta_schedule.runner import EvaluatorConfig
+from tvm.s_tir.meta_schedule.testing.tune_utils import generate_input_data
+from tvm.tirx import PrimFunc
+
+from .extract import extract_all_func_info_from_relax, extract_func_info_from_prim_func
+from .utils import (
+    default_dym_var_sample_func,
+    dym_var_sample_str,
+    get_func_name_from_gv,
+    populuate_input_shape,
+    print_results,
+)
+
+if TYPE_CHECKING:
+    from tvm.s_tir.meta_schedule.runner import RPCConfig
+
+
+def benchmark(
+    mod_or_func: PrimFunc | IRModule,
+    *,
+    dym_var_sample: dict[str, int],
+    args: list[relax.TensorType | tuple[tuple[int | str, ...], str]] | None,
+    target: str | tvm.target.Target | None = None,
+    func_name: str | None = None,
+    evaluator_config: Optional["EvaluatorConfig"] = None,
+    rpc_config: Optional["RPCConfig"] = None,
+) -> tuple[list[tuple[tuple[int, ...], str]], float, float]:
+    """Benchmark a PrimFunc or IRModule with dynamic input shapes.
+
+    Parameters
+    ----------
+    mod_or_func : Union[PrimFunc, IRModule]
+        The PrimFunc or IRModule to be benchmarked.
+    dym_var_sample : Optional[Dict[str, int]]
+        The dynamic shape variable sample, e.g., {"n": 64, "m": 128}.
+    args : Optional[List[Union[relax.TensorType, Tuple[Tuple[Union[int, str], ...], str]]]]
+        The input tensor information, including shape and dtype. If none, will use
+        the input information from the PrimFunc or IRModule.
+    target : Optional[Union[str, tvm.target.Target]]
+        The target to be benchmarked on, if none, will get the target from context.
+    func_name : Optional[str]
+        The name of the function to be benchmarked, will use "main" by default.
+    evaluator_config : Optional["EvaluatorConfig"]
+        The evaluator configuration to use.
+        If none, will use default evaluator configuration.
+    rpc_config : Optional["RPCConfig"]
+        The RPC configuration to connect to the remote device.
+        If none, will use local mode.
+
+    Returns
+    -------
+    input_infos : List[Tuple[Tuple[int, ...], str]]
+        The input tensor information, including shape and dtype.
+    median : float
+        The median of the benchmarking results.
+    std : float
+        The standard deviation of the benchmarking results.
+    """
+    # produce IRModule and function name
+    if isinstance(mod_or_func, PrimFunc):
+        func_name = "main" if func_name is None else func_name
+        mod = IRModule.from_expr(mod_or_func.with_attr("global_symbol", func_name))
+    else:
+        mod = mod_or_func
+        # assume only one global function
+        (func_name,) = mod.get_global_vars()
+        func_name = func_name.name_hint
+    # produce input shapes
+    if args is None:
+        args, _ = extract_func_info_from_prim_func(mod[func_name])
+    # produce target & device
+    target = tvm.target.Target.current() if target is None else tvm.target.Target(target)
+    if target is None:
+        raise ValueError("Target is not specified")
+    if target.kind.name == "llvm":
+        dev = tvm.cpu()
+    elif target.kind.name == "cuda":
+        dev = tvm.cuda()
+    else:
+        raise ValueError(f"Unsupported device type from {target.kind.name}")
+    # populate input shapes
+    input_infos = populuate_input_shape(args, dym_var_sample)
+    # generate input tensors, including scalars
+    # scalars are appended to the end of the list due to parsing order
+    input_tensors: list[tvm.runtime.Tensor | int] = []
+    scalar_input_tensors: list[int] = []
+    for input_shape, input_dtype in input_infos:
+        if input_dtype == "scalar":
+            # special case like [n], generate int value
+            assert len(input_shape) == 1
+            scalar_input_tensors.append(input_shape[0])
+        else:
+            # normal case like [1, n, 128], generate random tensor
+            input_tensors.append(
+                tvm.runtime.tensor(generate_input_data(list(input_shape), input_dtype), device=dev)
+            )
+    # append scalar input tensors for rotary embedding
+    input_tensors.extend(scalar_input_tensors)
+    # build locally
+    rt_mod = tvm.tirx.build(mod, target=target)
+    # set up evaluator config
+    evaluator_config = EvaluatorConfig._normalized(  # pylint: disable=protected-access
+        evaluator_config
+    )
+    # run benchmark
+    if rpc_config is None:
+        profile_result = rt_mod.time_evaluator(
+            func_name,
+            dev=dev,
+            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 ""
+            ),
+        )(*input_tensors)
+    else:
+        from tvm.testing import rpc_run  # pylint: disable=import-outside-toplevel
+
+        _, profile_result = rpc_run(
+            rt_mod,
+            device_type=dev._DEVICE_TYPE_TO_NAME[dev.dlpack_device_type()],
+            args=[w.numpy() if isinstance(w, tvm.runtime.Tensor) else w for w in input_tensors],
+            rpc_config=rpc_config,
+            evaluator_config=evaluator_config,
+        )
+    # return input infos, median, std
+    return input_infos, profile_result.median, profile_result.std
+
+
+def benchmark_prim_func(
+    mod_or_func: PrimFunc | IRModule,
+    *,
+    dym_var_sample_func: Callable[[dict[str, str]], dict[str, int]] = default_dym_var_sample_func,
+    args: list[relax.TensorType | tuple[tuple[int | str, ...], str]] | None = None,
+    dym_var_dict: dict[str, str] | None = None,
+    sample_number: int = 5,
+    target: str | tvm.target.Target | None = None,
+    weight: int | None = 1,
+    relax_func_name: str | None = None,
+    prim_func_name: str | None = None,
+    evaluator_config: Optional["EvaluatorConfig"] = None,
+    rpc_config: Optional["RPCConfig"] = None,
+    sort_by: str | None = None,
+    desc: bool | None = True,
+):
+    """Benchmark a PrimFunc or IRModule with dynamic input shapes and show results.
+
+    Parameters
+    ----------
+    mod_or_func : Union[PrimFunc, IRModule]
+        The PrimFunc or IRModule to be benchmarked.
+    dym_var_sample_func : Callable[[Dict[str, str]], Dict[str, int]]
+        The function to sample dynamic shape variables.
+    dym_var_dict : Optional[Dict[str, str]]
+        Dynamic shape variable dictionary, e.g., {"n": "int32", "m": "int32"}. If none, will use
+        the input information from the PrimFunc or IRModule.
+    args : Optional[List[Union[relax.TensorType, Tuple[Tuple[Union[int, str], ...], str]]]]
+        The input tensor information, including shape and dtype. If none, will use
+        the input information from the PrimFunc or IRModule.
+    sample_number : int
+        The number of times to sample dynamic shape variables.
+    target: Optional[Union[str, tvm.target.Target]]
+        The target to be benchmarked on, if none, will get the target from context.
+    weight : Optional[int]
+        The weight of this PrimFunc.
+    relax_func_name : Optional[str]
+        The name of the relax function.
+    prim_func_name : Optional[str]
+        The name of the PrimFunc.
+    evaluator_config : Optional["EvaluatorConfig"]
+        The evaluator configuration to use.
+        If none, will use default evaluator configuration.
+    rpc_config : Optional["RPCConfig"]
+        The RPC configuration to connect to the remote device.
+        If none, will use local mode.
+    sort_by : Optional[str]
+        Sort results by this key, if None, no sorting.
+    desc : Optional[bool]
+        Whether to sort results in descending order.
+    """
+    results = []
+    if dym_var_dict is None or args is None:
+        args, dym_var_dict = extract_func_info_from_prim_func(mod_or_func)
+    for _ in range(sample_number):
+        dym_var_sample = dym_var_sample_func(dym_var_dict)
+        _, median, std = benchmark(
+            mod_or_func,
+            args=args,
+            dym_var_sample=dym_var_sample,
+            target=target,
+            evaluator_config=evaluator_config,
+            rpc_config=rpc_config,
+        )
+        row = {
+            "InputInfo": ", ".join([f"{k} = {v}" for k, v in dym_var_sample.items()]),
+            "Time(us)": median * 1e6,
+            "Std(us)": std * 1e6,
+        }
+        if relax_func_name is not None:
+            row["RelaxFunc"] = relax_func_name
+        if prim_func_name is not None:
+            row["PrimFunc"] = prim_func_name
+        weight = 1 if weight is None else weight
+        row["Weight"] = weight
+        row["WxTime(ms)"] = weight * median * 1e3
+        results.append(row)
+    print_results(results, sort_by=sort_by, desc=desc)
+
+
+def benchmark_relax_func(
+    mod: tvm.ir.IRModule,
+    relax_func: tvm.ir.GlobalVar | str,
+    sample_number: int = 2,
+    dym_var_sample_func: Callable[
+        [dict[str, str]],
+        dict[str, int],
+    ] = default_dym_var_sample_func,
+    target: str | dict | tvm.target.Target = None,
+    evaluator_config: Optional["EvaluatorConfig"] = None,
+    rpc_config: Optional["RPCConfig"] = None,
+) -> None:
+    """Benchmark a relax function with dynamic input shapes.
+
+    Parameters
+    ----------
+    mod : tvm.ir.IRModule
+        The IRModule to be benchmarked.
+    relax_func : Union[tvm.ir.GlobalVar, str]
+        The relax function to be benchmarked.
+    sample_number : int
+        The number of times to sample dynamic shape variables.
+    dym_var_sample_func : Callable[[Dict[str, str]], Dict[str, int]]
+        The function to sample dynamic shape variables.
+    target : Union[str, tvm.target.Target]
+        The target to be benchmarked on.
+    dev : tvm.runtime.Device
+        The device to be benchmarked on.
+    evaluator_config : Optional["EvaluatorConfig"]
+        The evaluator configuration to use.
+        If none, will use default evaluator configuration.
+    rpc_config : Optional["RPCConfig"]
+        The RPC configuration to connect to the remote device.
+    """
+    if target is None:
+        target = {"kind": "llvm", "num-cores": 4}
+    # extract function information
+    relax_funcs, dynamic_var_dict = extract_all_func_info_from_relax(mod)
+    # find the relax function global var
+    if isinstance(relax_func, str):
+        for gv in relax_funcs:  # pylint: disable=invalid-name
+            if get_func_name_from_gv(gv) == relax_func:
+                relax_func = gv
+                break
+        if not isinstance(relax_func, tvm.ir.GlobalVar):
+            raise ValueError(
+                f"Cannot find relax function with name {relax_func}, "
+                + f"candidates are: {[get_func_name_from_gv(gv) for gv in relax_funcs]}"
+            )
+    # benchmark
+    for _ in range(sample_number):
+        dym_var_sample = dym_var_sample_func(dynamic_var_dict[relax_func])
+        bench_results = []
+        # enumerate all functors
+        for functor in relax_funcs[relax_func]:
+            for args, weight in relax_funcs[relax_func][functor]:
+                _, median, _ = benchmark(
+                    mod[functor],
+                    args=args,
+                    dym_var_sample=dym_var_sample,
+                    target=target,
+                    evaluator_config=evaluator_config,
+                    rpc_config=rpc_config,
+                )
+                bench_results.append(
+                    {
+                        f"PrimFuncs in {get_func_name_from_gv(relax_func)}": get_func_name_from_gv(
+                            functor
+                        ),
+                        f"InputInfo({dym_var_sample_str(dym_var_sample)})": ", ".join(
+                            [str(w) for w in args]
+                        ),
+                        "Time(us)": median * 1e6,
+                        # "Std(us)": std * 1e6,
+                        "Weight": weight,
+                        "WxTime(ms)": median * weight * 1e3,
+                    }
+                )
+        print_results(bench_results)

python/tvm/s_tir/dlight/benchmark/extract.py

@@ -0,0 +1,355 @@
+# 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.
+"""Performance debug tool for dynamic shape workloads"""
+
+from pathlib import Path
+
+import cloudpickle
+import tvm_ffi
+
+import tvm
+from tvm import relax
+
+from .utils import default_dym_var_sample_func, get_func_name_from_gv
+
+SKETCH = """import pickle
+
+import tvm
+from tvm import relax
+from tvm.script import tirx as T
+
+from tvm.s_tir.dlight.benchmark import benchmark_prim_func
+
+MODEL_NAME = "{model_name}"
+RELAX_FUNC_NAME = "{relax_func_name}"
+PRIM_FUNC_NAME = "{prim_func_name}"
+FUNC_HASH = {func_hash}
+WEIGHT = {weight}
+SAMPLE_NUMBER = {sample_number}
+
+DYM_VAR_SAMPLE_FUNC = {dym_var_sample_func}
+
+# None means extract from PrimFunc
+INPUT_ARGS = {input_args}
+DYM_VAR_DICT = {dym_var_dict}
+
+{func_script}
+
+if __name__ == "__main__":
+    target = tvm.target.Target({target})
+    benchmark_prim_func(
+        main,
+        args = INPUT_ARGS,
+        dym_var_dict = DYM_VAR_DICT,
+        dym_var_sample_func = DYM_VAR_SAMPLE_FUNC,
+        sample_number = SAMPLE_NUMBER,
+        target = target,
+        weight = WEIGHT,
+        relax_func_name = RELAX_FUNC_NAME,
+        prim_func_name = PRIM_FUNC_NAME,
+    )
+"""
+
+
+def extract_shape(
+    arg: tuple | list | relax.Tuple | relax.ShapeType,
+) -> list[relax.ShapeType]:
+    """Extract shape information from a relax argument.
+
+    Parameters
+    ----------
+    arg : Union[Tuple, List, relax.Tuple, relax.ShapeType]
+        The relax argument to be extracted.
+
+    Returns
+    -------
+    result : List[relax.ShapeType]
+        The extracted shape information.
+    """
+    if isinstance(arg, tuple | list | tvm.relax.Tuple):
+        results = []
+        for sub_arg in arg:
+            results.extend(extract_shape(sub_arg))
+        return results
+    return [arg.ty]
+
+
+def extract_dynamic_var(
+    func_dict: dict[
+        tvm.ir.GlobalVar,
+        dict[
+            tvm.ir.GlobalVar,
+            list[tuple[list, int]],
+        ],
+    ],
+) -> dict[tvm.ir.GlobalVar, dict[str, str]]:
+    """Extract dynamic shape variables from a relax function dictionary.
+
+    Parameters
+    ----------
+    func_dict : Dict[
+        tvm.ir.GlobalVar,
+        Dict[
+            tvm.ir.GlobalVar,
+            List[Tuple[List, int]],
+        ],
+        The relax function dictionary, containing the input arguments' shape information of each
+        PrimFunc in a Relax function.
+
+    Returns
+    -------
+    result : Dict[tvm.ir.GlobalVar, Dict[str, str]]
+        The dictionary of dynamic shape variables. Given in format {"n": "int32", "m": "int32"}.
+    """
+    dym_var_dict: dict[tvm.ir.GlobalVar, dict[str, str]] = {}
+    for gv in func_dict:  # pylint: disable=invalid-name,too-many-nested-blocks
+        dym_var_dict[gv] = {}
+        for functor in func_dict[gv]:
+            for arg_list, _ in func_dict[gv][functor]:
+                flattened_arg_list = []
+                for arg in arg_list:
+                    if isinstance(arg, relax.TupleType):
+                        flattened_arg_list.extend(arg.fields)
+                    else:
+                        flattened_arg_list.append(arg)
+                for arg in flattened_arg_list:
+                    if isinstance(arg, relax.TensorType):
+                        for val in arg.shape.values:
+                            if isinstance(val, tvm.tirx.Var):
+                                dym_var_dict[gv][str(val)] = str(val.ty)
+                    elif isinstance(arg, relax.ShapeType):
+                        for val in arg.values:
+                            if isinstance(val, tvm.tirx.Var):
+                                dym_var_dict[gv][str(val)] = str(val.ty)
+                    else:
+                        raise NotImplementedError
+    return dym_var_dict
+
+
+def update_records(
+    records: dict[list[relax.ShapeType], int], new_args: list[relax.ShapeType]
+) -> None:
+    """Update the count of a function input argument config.
+
+    Parameters
+    ----------
+    records : Dict[List[relax.ShapeType], int]
+        The dictionary to count how many times a function input argument config appears.
+    new_args : List[relax.ShapeType]
+        The new input argument config.
+    """
+    for i, (args, count) in enumerate(records):
+        if new_args == args:
+            records[i] = (args, count + 1)
+            return
+    records.append((new_args, 1))
+
+
+def extract_func_info_from_prim_func(
+    func: tvm.tirx.PrimFunc,
+) -> tuple[list[tuple[tuple[tvm.tirx.Var | int, ...], str]], dict[str, str]]:
+    """Extract function input information from a PrimFunc.
+
+    Parameters
+    ----------
+    func : tvm.tirx.PrimFunc
+        The PrimFunc to be analyzed.
+
+    Returns
+    -------
+    result : Tuple[
+        List[Tuple[Tuple[Union[tvm.tirx.Var, int], ...], str]],
+        Dict[str, str],
+    ]
+        The function input information and dynamic shape variable dictionary.
+    """
+    func_args = []
+    dym_var = {}
+    for param in func.params:
+        buffer = func.buffer_map[param]
+        shape = []
+        for dim in buffer.shape:
+            if isinstance(dim, tvm.tirx.IntImm):
+                shape.append(dim.value)
+            elif isinstance(dim, tvm.tirx.Var):
+                dym_var[str(dim)] = str(dim.ty)
+                shape.append(dim)
+            else:
+                raise ValueError(f"Unknown shape: {buffer.shape}")
+        func_args.append((tuple(shape), str(buffer.dtype)))
+    return func_args, dym_var
+
+
+def extract_all_func_info_from_relax(
+    mod: tvm.ir.IRModule,
+) -> tuple[
+    dict[tvm.ir.GlobalVar, dict[tvm.ir.GlobalVar, list[tuple[list, int]]]],
+    dict[tvm.ir.GlobalVar, dict[str, str]],
+]:
+    """Extract function input information from a relax module.
+
+    Parameters
+    ----------
+    mod : tvm.ir.IRModule
+        The Relax module to be analyzed.
+
+    Returns
+    -------
+    result : Tuple[
+        Dict[tvm.ir.GlobalVar, Dict[tvm.ir.GlobalVar, List[Tuple[List, int]]]],
+        Dict[tvm.ir.GlobalVar, Dict[str, str]],
+    ]
+        The function input information and dynamic shape variable dictionary.
+    """
+    relax_func_dict: dict[tvm.ir.GlobalVar, dict[tvm.ir.GlobalVar, list[tuple[list, int]]]] = {}
+    for gv, func in mod.functions_items():  # pylint: disable=invalid-name,too-many-nested-blocks
+        if isinstance(func, tvm.relax.Function):
+            for block in func.body.blocks:
+                for binding in block.bindings:
+                    if isinstance(binding.value, tvm.ir.Call):
+                        raw_args = binding.value.args
+                        functor = raw_args[0]
+                        if isinstance(functor, tvm.ir.GlobalVar) and isinstance(
+                            mod.functions[functor], tvm.tirx.PrimFunc
+                        ):
+                            args = extract_shape(raw_args[1:]) + extract_shape(binding.value)
+                            if isinstance(functor, tvm.ir.GlobalVar):
+                                if gv not in relax_func_dict:
+                                    relax_func_dict[gv] = {}
+                                if functor not in relax_func_dict[gv]:
+                                    relax_func_dict[gv][functor] = []
+                                update_records(relax_func_dict[gv][functor], args)
+
+    return relax_func_dict, extract_dynamic_var(relax_func_dict)
+
+
+def extract_prim_func(  # pylint: disable=too-many-arguments
+    model_name: str,
+    relax_func_name: str,
+    prim_func_name: str,
+    func: tvm.tirx.PrimFunc,
+    *,
+    func_args: list[tuple[tuple[tvm.ir.Call | int, ...], str]] | None = None,
+    dym_var_dict: dict[str, str] | None = None,
+    weight: int = 1,
+    sample_number: int = 5,
+    target: str | dict | tvm.target.Target | None = None,
+) -> str:
+    """Extract a self-contained PrimFunc test file from a Relax module.
+
+    Parameters
+    ----------
+    model_name: str
+        The name of the model.
+    relax_func_name: str
+        The name of the Relax function.
+    prim_func_name: str
+        The name of the prim function.
+    func: tvm.tirx.PrimFunc
+        The PrimFunc to be extracted.
+    func_args: Optional[List[Tuple[Tuple[Union[tvm.ir.Call, int], ...], str]]]
+        The arguments of the prim function, including both static and dynamic shape arguments.
+        Given in format [ ..., ((1, n, 128), "float32"), ... ].
+        If not given, the arguments will be extracted from the PrimFunc.
+    dym_var_dict: Optional[Dict[str, str]]
+        The dictionary of dynamic shape variables. Given in format {"n": "int32", "m": "int32"}.
+        If not given, the dictionary will be extracted from the PrimFunc.
+    weight: int
+        The weight of the prim function, by default 1.
+    sample_number: int
+        The number of times to sample dynamic shape variables, by default 5.
+    target: Optional[Union[str, dict, tvm.target.Target]]
+        The target device to run the PrimFunc. If None, will use target from the context.
+
+    Returns
+    -------
+    result : str
+        The extracted PrimFunc test file content.
+    """
+    if target is None:
+        target = tvm.target.Target.current()
+        if target is None:
+            raise ValueError("Target is not specified.")
+    elif isinstance(target, str | dict):
+        target = tvm.target.Target(target)
+    elif not isinstance(target, tvm.target.Target):
+        raise TypeError("Unsupported target type: " + str(type(target)))
+    target_json = str(target)
+
+    return SKETCH.format(
+        **{
+            "model_name": model_name,
+            "relax_func_name": relax_func_name,
+            "prim_func_name": prim_func_name,
+            "func_hash": tvm_ffi.structural_hash(func),
+            "weight": weight,
+            "sample_number": sample_number,
+            "dym_var_dict": f"pickle.loads({cloudpickle.dumps(dym_var_dict)})"
+            if dym_var_dict is not None
+            else "None",
+            "input_args": f"pickle.loads({cloudpickle.dumps(func_args)})" if func_args else "None",
+            "dym_var_sample_func": "pickle.loads("
+            + f"{cloudpickle.dumps(default_dym_var_sample_func)}"
+            + ")",
+            "func_script": func.script(),
+            "target": target_json,
+        }
+    )
+
+
+def extract_from_relax(
+    mod: tvm.ir.IRModule,
+    model_name: str,
+    file_path: str,
+    target: str | dict | tvm.target.Target | None = None,
+) -> None:
+    """Extract self-contained PrimFunc test files from a Relax module.
+
+    Parameters
+    ----------
+    mod: tvm.ir.IRModule
+        The Relax module to be extracted.
+    model_name: str
+        The name of the model.
+    file_path: str
+        The path to store the extracted files.
+    target: Optional[Union[str, tvm.target.Target]]
+        The target device to run the PrimFunc. If None, will use target from the context.
+    """
+    relax_funcs, dym_var_dict = extract_all_func_info_from_relax(mod)
+    Path(file_path).mkdir(parents=True, exist_ok=True)
+    for relax_func_gv in relax_funcs:  # pylint: disable=consider-using-dict-items
+        relax_func_name = get_func_name_from_gv(relax_func_gv)
+        for prim_func_gv in relax_funcs[relax_func_gv]:
+            prim_func_name = get_func_name_from_gv(prim_func_gv)
+            for func_args, weight in relax_funcs[relax_func_gv][prim_func_gv]:
+                with open(
+                    f"{file_path}/{relax_func_name}_{prim_func_name}.py", "w", encoding="utf-8"
+                ) as file:
+                    print(
+                        extract_prim_func(
+                            model_name=model_name,
+                            relax_func_name=relax_func_name,
+                            prim_func_name=prim_func_name,
+                            func=mod[prim_func_gv],
+                            dym_var_dict=dym_var_dict[relax_func_gv],
+                            func_args=func_args,
+                            weight=weight,
+                            target=target,
+                        ),
+                        file=file,
+                    )

python/tvm/s_tir/dlight/benchmark/utils.py

@@ -0,0 +1,178 @@
+# 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.
+"""Util functions for benchmarking dynamic shape workloads"""
+
+from typing import Any
+
+import tvm
+from tvm import relax
+
+INPUT_SHAPE_TYPE = list[tuple[tuple[int, ...], str]]  # pylint: disable=invalid-name
+
+
+def _dtype_str(dtype) -> str:
+    if isinstance(dtype, tvm.ir.PrimType):
+        dtype = dtype.dtype
+    return str(dtype)
+
+
+def get_func_name_from_gv(gv: tvm.ir.GlobalVar) -> str:  # pylint: disable=invalid-name
+    """Get function name from a global variable.
+
+    Parameters
+    ----------
+    gv : tvm.ir.GlobalVar
+        The given global variable.
+
+    Returns
+    -------
+    result : str
+        The global variable name without the prefix "...@".
+    """
+    return gv.name_hint
+
+
+def dym_var_sample_str(sample: dict[str | tvm.ir.Call, int]) -> str:
+    """Convert a variable value sample to a string.
+
+    Parameters
+    ----------
+    sample : Dict[Union[str, tvm.ir.Call], int]
+        Variable value sample, e.g., {n: 64, m: 128} or {"n": 64, "m": 128}
+
+    Returns
+    -------
+    result : str
+        Variable value sample string, e.g., "n=64, m=128"
+    """
+    return ", ".join([f"{k}={v}" for k, v in sample.items()])
+
+
+def populuate_input_shape(
+    input_infos: list[relax.TensorType | tuple[tuple[int | str, ...], str]],
+    dym_var_sample: dict[str, int],
+) -> INPUT_SHAPE_TYPE:
+    """
+    Populate input shapes with dynamic shape variable samples.
+
+    Parameters
+    ----------
+    input_infos : List[Union[relax.TensorType, Tuple[Tuple[Union[int, str], ...], str]]]
+        Input tensor information, including shape and dtype,
+        e.g., [..., Shape(1, n, 128) with dtype="int32", ...]
+    dym_var_sample : Dict[str, int]
+        Dynamic shape variable sample, e.g., {"n": 64}
+
+    Returns
+    -------
+    results : INPUT_SHAPE_TYPE
+        Input shapes with dynamic shape variable samples, e.g.,
+        [..., ((1, 64, 128), "int32"), ...] if n=64 or
+        [..., (128, "scalar"), ...] if n=128 for scalar input
+    """
+    results: INPUT_SHAPE_TYPE = []
+    for input_info in input_infos:
+        shape = []
+        if isinstance(input_info, relax.ShapeType):
+            # scalar input
+            results.append(((dym_var_sample[str(input_info.values[0])],), "scalar"))
+        else:
+            if isinstance(input_info, relax.TensorType):
+                tensor_shape = input_info.shape
+                tensor_dtype = input_info.dtype
+            else:
+                tensor_shape, tensor_dtype = input_info  # type: ignore
+            for dim in tensor_shape:
+                if isinstance(dim, int):
+                    shape.append(dim)
+                elif isinstance(dim, tvm.tirx.IntImm):
+                    shape.append(dim.value)
+                else:
+                    shape.append(dym_var_sample[str(dim)])
+            results.append(((*shape,), _dtype_str(tensor_dtype)))
+    return results
+
+
+def default_dym_var_sample_func(dym_var_dict: dict[str, str]) -> dict[str, int]:
+    """
+    Default dynamic shape variable sample function.
+    Sample a random value for each dynamic shape variable.
+
+    Parameters
+    ----------
+    dym_var_dict : Dict[str, str]
+        Dynamic shape variable dictionary, e.g., {"n": "int32", "m": "int32"}
+
+    Returns
+    -------
+    result : Dict[str, int]
+        Dynamic shape variable sample, e.g., {"n": 64, "m": 128}
+    """
+    results = {}
+    for var in dym_var_dict:
+        if dym_var_dict[var] in ["int32", "int64"]:
+            import random  # pylint: disable=import-outside-toplevel
+
+            results[var] = random.randint(2, 128)
+        else:
+            raise TypeError("Unsupported dynamic shape variable type: " + dym_var_dict[var])
+    return results
+
+
+def print_results(
+    bench_results: list[dict[str, Any]], sort_by: str = "WxTime(ms)", desc: bool = True
+):
+    """Print benchmark results.
+
+    Parameters
+    ----------
+    bench_results : List[Dict[str, Any]]
+        Benchmark results as dictionary list.
+    sort_by : str
+        Sort results by this key, if None, no sorting.
+    desc : bool
+        Whether to sort results in descending order.
+    """
+    # pylint: disable=invalid-name, import-outside-toplevel
+    try:
+        import pandas as pd
+
+        df = pd.DataFrame()
+        for record in bench_results:
+            df = pd.concat(
+                [df, pd.DataFrame(record, index=[0])],
+                ignore_index=True,
+            )
+        if sort_by is not None:
+            if sort_by not in df.columns:
+                raise ValueError(f"sort_by key {sort_by} not in benchmark results")
+            df = df.sort_values(sort_by, ascending=not desc).reset_index().drop("index", axis=1)
+        print(df)
+    except ModuleNotFoundError:
+        print("Pandas not found, printing results in raw format.")
+        keys = []
+        if len(bench_results) > 0:
+            for key in bench_results[0]:
+                keys.append(str(key))
+        print("\t".join(keys))
+        for record in bench_results:
+            values = []
+            for key in keys:
+                values.append(str(record[key]))
+            print("\t".join(values))
+    print("\n")
+    # pylint: enable=invalid-name, import-outside-toplevel

python/tvm/s_tir/dlight/cpu/__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.
+"""
+CPU-generic schedule rules.
+"""
+
+from .gemv import GEMV
+from .reduction import Reduction

python/tvm/s_tir/dlight/cpu/base.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.
+"""Base schedule rule for CPU operators."""
+
+from tvm.target import Target
+
+from ..base import ScheduleRule
+
+
+class CPUScheduleRule(ScheduleRule):  # pylint: disable=too-few-public-methods
+    """The Schedule Rule specific to CPU targets, will return None if the target is not CPU."""
+
+    def is_target_available(self, target: Target) -> bool:
+        """Check whether the target is available for gpu rule.
+
+        Parameters
+        ----------
+        target : Target
+            The compilation target to check.
+
+        Returns
+        -------
+        available : bool
+            Whether the target is available for this rule.
+        """
+        return super().is_target_available(target) and "llvm" == target.kind.name

python/tvm/s_tir/dlight/cpu/gemv.py

@@ -0,0 +1,131 @@
+# 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 rule for GEMV and DecodeGEMV."""
+
+from tvm import s_tir, tirx
+from tvm.target import Target
+
+from ..analysis import SBlockInfo, normalize_prim_func
+from ..analysis.gemv import is_gemv, normalize
+from ..base import get_extent, try_inline_contiguous_spatial
+from .base import CPUScheduleRule
+
+
+class GEMV(CPUScheduleRule):
+    """A rule for GEMV and DecodeGEMV."""
+
+    def apply(  # pylint: disable=too-many-locals,too-many-branches,too-many-return-statements, no-else-return
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+        if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
+            return None
+        sch = s_tir.Schedule(func)
+        block_infos = normalize_prim_func(sch)
+        block_infos = try_inline_contiguous_spatial(sch, block_infos)
+        if block_infos is None:
+            return None
+        if len(block_infos) == 1:
+            epilogue = None
+        elif len(block_infos) == 2:
+            epilogue = block_infos[1]
+            if not epilogue.is_injective():
+                return None
+        else:
+            return None
+
+        block_info = block_infos[0]
+        if len(block_info.iters) not in [2, 3]:
+            # either [B, S, R] = [B, S, R] * [B, R]
+            # or [S, R] = [S, R] * [R]
+            return None
+        block = block_info.block_rv
+        vector_input_buffers = is_gemv(sch, block_info)
+        if vector_input_buffers is None:
+            return None
+
+        # Step 1. Normalize the block, merge spatial and reduction iters
+        is_inner_reduction = normalize(sch, block_info)
+
+        # Step 2. Do the scheduling
+        if is_inner_reduction is None:
+            return None
+        elif is_inner_reduction:
+            return self.sch_inner_reduction(sch, target, block, vector_input_buffers, epilogue)
+        else:
+            # sch_outer reduction
+            return None
+
+    def sch_inner_reduction(  # pylint: disable=too-many-arguments, too-many-positional-arguments, invalid-name, unused-argument
+        self,
+        sch: s_tir.Schedule,
+        target: Target,
+        block: s_tir.schedule.SBlockRV,
+        vector_input_buffers: list[tirx.Buffer],
+        epilogue_info: SBlockInfo | None,
+    ):
+        """Schedule the inner reduction block."""
+
+        def apply(  # pylint: disable=unused-variable, too-many-locals
+            sch: s_tir.Schedule,
+            gemv,
+            vector_width: int = 8,
+            parallel_threads: int = 8,
+            unroll_factor: int = 256,
+        ):
+            batch, s, r, c = sch.get_loops(block)
+            len_batch, len_s, len_r, len_c = (
+                get_extent(sch, batch),
+                get_extent(sch, s),
+                get_extent(sch, r),
+                get_extent(sch, c),
+            )
+            len_S = len_batch * len_s
+            len_R = len_r * len_c
+
+            if isinstance(len_S, int) and isinstance(len_R, int):
+                if len_S > len_R:
+                    tile_s, tile_r = 128, 64  # Larger tiling for s-axis when len_S is larger
+                else:
+                    tile_s, tile_r = 64, 128  # Larger tiling for r-axis when len_R is larger
+            else:
+                tile_s, tile_r = 64, 64  # Default tile sizes for unknown extents
+
+            tile_c = min(vector_width, len_c)  # Ensure c-axis tiling aligns with SIMD vector width
+
+            # Apply loop tiling (improves cache locality)
+            s_outer, s_inner = sch.split(s, factors=[None, tile_s])
+            r_outer, r_inner = sch.split(r, factors=[None, tile_r])
+            c_outer, c_inner = sch.split(c, factors=[None, tile_c])
+
+            # Apply vectorization (SIMD optimization)
+            sch.vectorize(s_inner)  # Vectorize computation along c-axis for AVX/NEON
+
+            # Enable parallel execution
+            sch.parallel(s_outer)  # Parallelize along the s-axis (major computation loop)
+
+            # Apply loop unrolling for better CPU performance
+            sch.annotate(r_outer, "pragma_auto_unroll_max_step", unroll_factor)
+            sch.annotate(r_outer, "pragma_unroll_explicit", 1)
+            return sch
+
+        return apply(
+            sch,
+            gemv=block,
+        )

python/tvm/s_tir/dlight/cpu/reduction.py

@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""CPU reduction rule for operators including softmax, layer norm, RMS norm, etc."""
+
+from tvm import s_tir, tirx
+from tvm.target import Target
+from tvm.target.codegen import llvm_get_vector_width
+
+from ..analysis import normalize_prim_func
+from ..base import get_extent
+from .base import CPUScheduleRule
+
+
+def _get_num_leading_s(dom_kind: str) -> int:
+    """Count leading spatial ('S') axes in a dom_kind string."""
+    return len(dom_kind) - len(dom_kind.lstrip("S"))
+
+
+class Reduction(CPUScheduleRule):
+    """CPU reduction rule for softmax, layer norm, RMS norm, and similar operators.
+
+    Targets patterns with a mix of reduction (SR) and injective (SS) blocks,
+    where all blocks share the same leading spatial axes.
+    Example: softmax = maxelem(SR) -> exp(SS) -> expsum(SR) -> norm(SS).
+
+    Schedule strategy:
+      1. Parallelize leading spatial axes (batch dimension).
+      2. Move all blocks under the spatial loop via compute_at.
+      3. Vectorize injective blocks (exp, delta, norm) on their inner axis.
+      4. Split reduction inner axis to VLEN-sized chunks and annotate for
+         LLVM unrolling, preventing harmful full-unroll by the backend.
+
+    Note: vectorized reduction via rfactor is not used here because TVM's
+    rfactor primitive requires the reduction block to be the first child of
+    its enclosing loop, which is incompatible with compute_at when multiple
+    blocks share the same spatial loop. A follow-up using RVV reduction
+    intrinsics (vfredmax/vfredusum) via tensorize can address this.
+    """
+
+    def apply(  # pylint: disable=too-many-locals,too-many-return-statements,too-many-branches
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+        if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
+            return None
+
+        sch = s_tir.Schedule(func)
+        block_infos = normalize_prim_func(sch)
+        if block_infos is None or len(block_infos) < 2:
+            return None
+
+        # Must have at least one reduction block and last block must be injective.
+        if not any(not bi.is_injective() for bi in block_infos):
+            return None
+        if not block_infos[-1].is_injective():
+            return None
+
+        # Every block must start with at least one spatial axis, and all blocks
+        # must agree on the minimum number of leading spatial axes.
+        num_leading_s = None
+        for bi in block_infos:
+            dk = bi.dom_kind()
+            if not dk or dk[0] != "S":
+                return None
+            n = _get_num_leading_s(dk)
+            num_leading_s = n if num_leading_s is None else min(num_leading_s, n)
+        if not num_leading_s:
+            return None
+
+        # Infer dtype from the last block's write buffer.
+        last_block_stmt = sch.get(block_infos[-1].block_rv)
+        dtype_bits = (
+            last_block_stmt.writes[0].buffer.dtype.dtype.bits if last_block_stmt.writes else 32
+        )
+
+        # Determine vector lanes from target VLEN.
+        vlen_bits = llvm_get_vector_width(target)
+        if vlen_bits <= 0:
+            vlen_bits = 128
+        vec_lanes = max(vlen_bits // dtype_bits, 2)
+
+        # --- Phase 1: Parallelize spatial on the last block ---
+        last_block = block_infos[-1]
+        loops = sch.get_loops(last_block.block_rv)
+        if num_leading_s > 1:
+            spatial = sch.fuse(*loops[:num_leading_s])
+        else:
+            spatial = loops[0]
+        sch.parallel(spatial)
+
+        # --- Phase 2: Vectorize the last (injective) block ---
+        self._vectorize_inner(sch, last_block.block_rv, vec_lanes)
+
+        # --- Phase 3: compute_at all preceding blocks under spatial ---
+        for block_info in reversed(block_infos[:-1]):
+            sch.compute_at(block_info.block_rv, spatial, preserve_unit_loops=True)
+
+        # --- Phase 4: Vectorize injective, split+unroll reduction blocks ---
+        for block_info in block_infos[:-1]:
+            if block_info.is_injective():
+                self._vectorize_inner(sch, block_info.block_rv, vec_lanes)
+            else:
+                self._unroll_reduction_inner(sch, block_info.block_rv, vec_lanes)
+
+        return sch
+
+    @staticmethod
+    def _vectorize_inner(sch, block_rv, vec_lanes):
+        """Split the innermost loop to vec_lanes and vectorize."""
+        block_loops = sch.get_loops(block_rv)
+        if len(block_loops) <= 1:
+            return
+        inner = block_loops[-1]
+        extent = get_extent(sch, inner)
+        if isinstance(extent, int):
+            if extent > vec_lanes:
+                _, vec_loop = sch.split(inner, factors=[None, vec_lanes])
+                sch.vectorize(vec_loop)
+            elif extent >= 2:
+                sch.vectorize(inner)
+        else:
+            _, vec_loop = sch.split(inner, factors=[None, vec_lanes])
+            sch.vectorize(vec_loop)
+
+    @staticmethod
+    def _unroll_reduction_inner(sch, block_rv, vec_lanes):
+        """Split the reduction inner loop and annotate for unrolling."""
+        block_loops = sch.get_loops(block_rv)
+        if len(block_loops) <= 1:
+            return
+        inner = block_loops[-1]
+        extent = get_extent(sch, inner)
+        if isinstance(extent, int) and extent <= vec_lanes:
+            return
+        _, inner_loop = sch.split(inner, factors=[None, vec_lanes])
+        sch.annotate(inner_loop, ann_key="pragma_auto_unroll_max_step", ann_val=vec_lanes)
+        sch.annotate(inner_loop, ann_key="pragma_unroll_explicit", ann_val=1)

python/tvm/s_tir/dlight/gpu/__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.
+"""
+GPU-generic schedule rules.
+For CUDA/ROCm/Vulkan/Metal-specific rules, use `tvm.s_tir.dlight.cuda/rocm/vulkan/metal` instead
+"""
+
+from .gemv import GEMV
+from .low_batch_gemv import LowBatchGEMV
+from .fallback import Fallback
+from .matmul import Matmul
+from .reduction import Reduction
+from .transpose import Transpose
+from .general_reduction import GeneralReduction
+from .rmsnorm import RMSNorm

python/tvm/s_tir/dlight/gpu/base.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.
+"""Base schedule rule for GPU operators."""
+
+from tvm.target import Target
+
+from ..base import ScheduleRule
+
+
+class GPUScheduleRule(ScheduleRule):  # pylint: disable=too-few-public-methods
+    """The Schedule Rule specific to GPU targets, will return None if the target is not GPU."""
+
+    def is_target_available(self, target: Target) -> bool:
+        """Check whether the target is available for gpu rule.
+
+        Parameters
+        ----------
+        target : Target
+            The compilation target to check.
+
+        Returns
+        -------
+        available : bool
+            Whether the target is available for this rule.
+        """
+        return super().is_target_available(target) and "gpu" in target.keys

python/tvm/s_tir/dlight/gpu/fallback.py

@@ -0,0 +1,107 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring
+"""A fallback schedule rule for GPU operators."""
+
+from tvm import s_tir, tirx
+from tvm.target import Target
+
+from .. import base
+from ..analysis import normalize_prim_func
+from ..base import try_inline
+from .base import GPUScheduleRule
+
+
+def _has_internal_thread_env(stmt: tirx.Stmt) -> bool:
+    """Check whether a statement already launches GPU threads internally,
+    e.g. via `T.launch_thread` (AttrStmt "thread_extent") or nested
+    thread-bound loops. Such blocks manage their own thread environment
+    and must not be wrapped in an additional thread binding."""
+    found = False
+
+    def _visit(node):
+        nonlocal found
+        if isinstance(node, tirx.AttrStmt) and node.attr_key in ("thread_extent", "virtual_thread"):
+            found = True
+        elif isinstance(node, tirx.For) and node.kind == tirx.ForKind.THREAD_BINDING:
+            found = True
+
+    tirx.stmt_functor.post_order_visit(stmt, _visit)
+    return found
+
+
+class Fallback(GPUScheduleRule):
+    """
+    A fallback schedule rule for all GPU operators. It will try to inline all the blocks first,
+    and then apply a simple block/grid mapping to the spatial loops on top of the remaining blocks.
+    """
+
+    def apply(  # pylint: disable=too-many-locals,missing-docstring
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> s_tir.Schedule:
+        if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
+            return None
+        max_threads_per_block = base.max_threads_per_block(target)
+
+        sch = s_tir.Schedule(func)
+        block_infos = normalize_prim_func(sch)
+
+        if block_infos is None:
+            return None
+
+        block_infos = try_inline(sch, block_infos)
+        reduction_blocks: list[tuple[s_tir.schedule.SBlockRV, s_tir.schedule.LoopRV]] = []
+        for block in block_infos:
+            s_loops: list[s_tir.schedule.LoopRV] = []
+            r_loops: list[s_tir.schedule.LoopRV] = []
+            o_loops: list[s_tir.schedule.LoopRV] = []
+            dom_kind = block.dom_kind()
+            block = block.block_rv
+
+            if any(
+                [sch.get(loop_rv).thread_binding is not None for loop_rv in sch.get_loops(block)]
+            ):
+                continue
+
+            if len(sch.get_loops(block)) == 0 and _has_internal_thread_env(sch.get(block).body):
+                # The block (e.g. an opaque sort kernel) launches its own
+                # threads; binding an outer loop would conflict with them.
+                continue
+
+            for loop, iter_type in zip(sch.get_loops(block), dom_kind):
+                {"S": s_loops, "R": r_loops, "O": o_loops}[iter_type].append(loop)
+
+            if not s_loops:
+                s_loops.append(sch.add_unit_loop(block))
+            sch.reorder(*s_loops, *r_loops, *o_loops)
+            bx, tx = sch.split(  # pylint: disable=invalid-name
+                sch.fuse(*s_loops),
+                factors=[None, max_threads_per_block],
+            )
+            sch.bind(bx, "blockIdx.x")
+            sch.bind(tx, "threadIdx.x")
+
+            if len(r_loops) > 0:
+                reduction_blocks.append((block, r_loops[0]))
+
+        for block, r_loop in reduction_blocks:
+            sch.decompose_reduction(block, r_loop)
+
+        return sch

python/tvm/s_tir/dlight/gpu/gemv.py

@@ -0,0 +1,689 @@
+# 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, F821
+"""A rule for GEMV and DecodeGEMV."""
+
+from functools import reduce
+
+from tvm import s_tir, tirx
+from tvm.target import Target
+
+from ..analysis import (
+    SBlockInfo,
+    get_max_shared_memory_per_block,
+    is_broadcast_epilogue,
+    is_gemv,
+    normalize,
+    normalize_prim_func,
+)
+from ..base import auto_vectorize, get_bytes, get_extent, try_inline_contiguous_spatial
+from .base import GPUScheduleRule
+
+
+class GEMV(GPUScheduleRule):
+    """A rule for GEMV and DecodeGEMV."""
+
+    def apply(  # pylint: disable=too-many-locals,too-many-branches,too-many-return-statements
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+        if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
+            return None
+        sch = s_tir.Schedule(func)
+        block_infos = normalize_prim_func(sch)
+        block_infos = try_inline_contiguous_spatial(sch, block_infos)
+        if block_infos is None:
+            return None
+        if len(block_infos) == 1:
+            epilogue = None
+        elif len(block_infos) == 2:
+            epilogue = block_infos[1]
+            if not epilogue.is_injective():
+                return None
+        else:
+            return None
+
+        block_info = block_infos[0]
+        if len(block_info.iters) not in [2, 3]:
+            # either [B, S, R] = [B, S, R] * [B, R]
+            # or [S, R] = [S, R] * [R]
+            return None
+        block = block_info.block_rv
+        vector_input_buffers = is_gemv(sch, block_info)
+        if vector_input_buffers is None:
+            return None
+
+        # Step 1. Normalize the block, merge spatial and reduction iters
+        is_inner_reduction = normalize(sch, block_info)
+
+        # Step 2. Do the scheduling
+        if is_inner_reduction is None:
+            return None
+        elif is_inner_reduction:
+            return self.sch_inner_reduction(sch, target, block, vector_input_buffers, epilogue)
+        else:
+            ret = self.sch_outer_reduction(sch, target, block, vector_input_buffers, epilogue)
+            if ret is None:
+                return self.sch_outer_reduction_fallback(
+                    sch, target, block, vector_input_buffers, epilogue
+                )
+            return sch
+
+    def sch_inner_reduction(  # pylint: disable=too-many-arguments, invalid-name, unused-argument
+        self,
+        sch: s_tir.Schedule,
+        target: Target,
+        block: s_tir.schedule.SBlockRV,
+        vector_input_buffers: list[tirx.Buffer],
+        epilogue_info: SBlockInfo | None,
+    ):
+        """Schedule the inner reduction block."""
+
+        def get_max_factor(n, factors):
+            factors = sorted(factors, reverse=True)
+            for factor in factors:
+                if n % factor == 0:
+                    return factor
+            return 1
+
+        def apply(
+            sch: s_tir.Schedule,
+            gemv,
+            TAG_S,
+            TAG_R,
+            TS,
+            TR,
+            TILE_S,
+            TILE_R,
+            VEC_LOAD,
+            VEC_C,
+            LOAD_V_SHARED,
+            LOAD_V_VEC,
+            UNROLL,
+            SUPPORT_WARP_SHUFFLE,
+        ):
+            # rfactor: reduce to tx * vec_c
+            _, s, r, c = sch.get_loops(block=gemv)
+            s = sch.fuse(_, s)
+            r = sch.fuse(r, c)
+            bx, ts, tile_s = sch.split(s, factors=[None, TS, TILE_S], preserve_unit_iters=True)
+            r, tr, tile_r_vec_n, vec_c = sch.split(
+                r, factors=[None, TR, TILE_R // VEC_C, VEC_C], preserve_unit_iters=True
+            )
+            sch.reorder(r, tile_r_vec_n, tr, vec_c)
+            tr_vec_c = sch.fuse(tr, vec_c)
+            rf = sch.rfactor(tr_vec_c, 0)
+
+            # rfactor: reduce to tx
+            bx, ts, tile_s, tr_vec_c = sch.get_loops(block=gemv)
+            tr, vec_c = sch.split(tr_vec_c, factors=[TR, None], preserve_unit_iters=True)
+            rf2 = sch.rfactor(tr, 0)
+
+            # bind, vectorize compute
+            bx, ts, tile_s, r, tile_r_vec_n, tr_vec_c = sch.get_loops(block=rf)
+            tr, vec_c = sch.split(tr_vec_c, factors=[TR, None], preserve_unit_iters=True)
+            sch.reorder(bx, ts, tr, r, tile_s, tile_r_vec_n, vec_c)
+            sch.bind(bx, "blockIdx.x")
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+            sch.vectorize(vec_c)
+
+            shared_mem_usage = 0
+            for buf in vector_input_buffers:
+                dtype_bytes = get_bytes(buf.dtype)
+                buf_size = (
+                    reduce(lambda x, y: x * y, buf.shape, tirx.IntImm(buf.shape[0].ty, 1))
+                    * dtype_bytes
+                )
+                shared_mem_usage += buf_size
+                if not SUPPORT_WARP_SHUFFLE:
+                    # When warp shuffle is not able, cross-thread allreduce
+                    # is implemented with shared memory.
+                    shared_mem_usage += TS * TR * dtype_bytes
+
+            max_smem = get_max_shared_memory_per_block(target)
+            LOAD_V_SHARED = (
+                LOAD_V_SHARED
+                and isinstance(shared_mem_usage, tirx.IntImm)
+                and shared_mem_usage.value <= max_smem
+            )
+
+            # vectorize load A
+            # (TODO) this is now actually problematic since the number of loops is dependent on the
+            # number of dimensions of A_q
+            Aq_local = sch.cache_read(rf, read_buffer_index=1, storage_scope="local")
+            sch.compute_at(Aq_local, r, preserve_unit_loops=True)
+            s_local, r_local = sch.get_loops(block=Aq_local)[-2:]
+            fused_load = sch.fuse(s_local, r_local)
+            aq_vec_len = max(1, VEC_LOAD // get_bytes(sch.get(Aq_local).reads[0].buffer.dtype))
+            fused_load, vec_load = sch.split(
+                fused_load, factors=[None, aq_vec_len], preserve_unit_iters=True
+            )
+            sch.vectorize(vec_load)
+
+            # load vector into shared memory, shape should be the whole vector
+            if LOAD_V_SHARED:
+                if len(vector_input_buffers) != 1:
+                    return None
+                V_shared = sch.cache_read(rf, read_buffer_index=0, storage_scope="shared")
+                sch.compute_at(V_shared, tr, preserve_unit_loops=True)
+                l = sch.get_loops(block=V_shared)[-1]
+                loop: tirx.For = sch.get(l)
+                if isinstance(loop.extent, tirx.IntImm):
+                    # avoid introducing predicates when vector length is too large
+                    vec_length = max(
+                        min(
+                            get_max_factor(
+                                (int)(loop.extent),
+                                [TS * TR * 1, TS * TR * 2, TS * TR * 4, TS * TR * 8],
+                            )
+                            // TS
+                            // TR,
+                            LOAD_V_VEC,
+                        ),
+                        1,
+                    )
+                else:
+                    vec_length = LOAD_V_VEC
+                if TAG_R == "threadIdx.x":
+                    _, ty, tx, vec = sch.split(
+                        l, factors=[None, TS, TR, vec_length], preserve_unit_iters=True
+                    )
+                else:
+                    _, ty, tx, vec = sch.split(
+                        l, factors=[None, TR, TS, vec_length], preserve_unit_iters=True
+                    )
+                sch.bind(ty, "threadIdx.y")
+                sch.bind(tx, "threadIdx.x")
+                sch.vectorize(vec)
+
+            # reduce tile_s * tr * vec to tile_s * tr
+            sch.reverse_compute_at(rf2, loop=bx, preserve_unit_loops=True)
+            tr, vec_c, *ts_tile_s = sch.get_loops(block=rf2)[1:]
+            ts_tile_s = sch.fuse(*ts_tile_s)
+            ts_o, ts_i, tile_s = sch.split(
+                ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True
+            )
+            tile_s, vec_s = sch.split(
+                tile_s,
+                factors=[None, get_max_factor(TILE_S, [1, 2, 4, 8])],
+                preserve_unit_iters=True,
+            )
+            assert sch.get(ts_o).extent.value == 1
+            ts = sch.fuse(ts_o, ts_i)
+            sch.reorder(ts, tr, tile_s, vec_s, vec_c)
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+            sch.vectorize(vec_s)
+
+            # reduce tile_s * tr to tile_s
+            sch.reverse_compute_at(gemv, loop=bx, preserve_unit_loops=True)
+            tr, *ts_tile_s = sch.get_loops(block=gemv)[1:]
+            ts_tile_s = sch.fuse(*ts_tile_s)
+            ts_o, ts_i, tile_s = sch.split(
+                ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True
+            )
+            assert sch.get(ts_o).extent.value == 1
+            ts = sch.fuse(ts_o, ts_i)
+            sch.reorder(tile_s, ts, tr)
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+
+            sch.decompose_reduction(rf, loop=sch.get_loops(block=rf)[3])
+            sch.decompose_reduction(rf2, loop=sch.get_loops(block=rf2)[-1])
+
+            sch.set_scope(rf, buffer_index=0, storage_scope="local")
+            sch.set_scope(rf2, buffer_index=0, storage_scope="local")
+
+            unroll_factor = UNROLL
+
+            sch.annotate(
+                block_or_loop=sch.get_loops(rf)[3],
+                ann_key="pragma_auto_unroll_max_step",
+                ann_val=unroll_factor,
+            )
+            sch.annotate(
+                block_or_loop=sch.get_loops(rf)[3], ann_key="pragma_unroll_explicit", ann_val=1
+            )
+
+            sch.annotate(
+                block_or_loop=sch.get_loops(rf2)[3],
+                ann_key="pragma_auto_unroll_max_step",
+                ann_val=unroll_factor,
+            )
+            sch.annotate(
+                block_or_loop=sch.get_loops(rf2)[3], ann_key="pragma_unroll_explicit", ann_val=1
+            )
+
+            if LOAD_V_SHARED:
+                sch.annotate(
+                    block_or_loop=sch.get_loops(V_shared)[-4],
+                    ann_key="pragma_unroll_explicit",
+                    ann_val=unroll_factor,
+                )
+                sch.annotate(
+                    block_or_loop=sch.get_loops(V_shared)[-4], ann_key="pragma_vectorize", ann_val=1
+                )
+
+            # Schedule epilogue
+            if epilogue_info is not None:
+                epilogue = epilogue_info.block_rv
+                if is_broadcast_epilogue(sch, block, epilogue):
+                    sch.reverse_compute_at(epilogue, bx)
+                    sch.set_scope(block, 0, "shared")
+                    _, _, *s = sch.get_loops(epilogue)  # pylint: disable=invalid-name
+                    _, tx = sch.split(sch.fuse(*s), factors=[None, TX])
+                    sch.bind(tx, "threadIdx.x")
+                else:
+                    sch.reverse_compute_at(epilogue, bx, preserve_unit_loops=True)
+                    ts_tile_s = sch.fuse(*sch.get_loops(epilogue)[1:])
+                    ts_tile_s = sch.get_loops(epilogue)[-1]
+                    ts_o, ts_i, tile_s = sch.split(
+                        ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True
+                    )
+                    assert sch.get(ts_o).extent.value == 1
+                    ts = sch.fuse(ts_o, ts_i)
+                    sch.bind(ts, TAG_S)
+                    sch.set_scope(block, 0, "local")
+            # pylint: enable=invalid-name
+            return sch
+
+        # Specify the `len_tx` and `len_ty` according to the loop extent
+        batch, s, r, c = sch.get_loops(block=block)
+        len_batch, len_s, len_r, len_c = (
+            get_extent(sch, batch),
+            get_extent(sch, s),
+            get_extent(sch, r),
+            get_extent(sch, c),
+        )
+        len_S = len_batch * len_s
+        len_R = len_r * len_c
+
+        TAG_S, TAG_R = "threadIdx.y", "threadIdx.x"
+        SUPPORT_WARP_SHUFFLE = False
+        VEC_LOAD = 1
+        if target.kind.name == "cuda":
+            VEC_C = 4
+            LOAD_V_SHARED = True
+            LOAD_V_VEC = 8
+            VEC_LOAD = 4
+            UNROLL = 256
+            SUPPORT_WARP_SHUFFLE = True
+            if isinstance(len_S, int):
+                TS, TR = 16, 32
+            else:
+                TS, TR = 1, 64
+        elif target.kind.name == "metal":
+            # Note that the following tile size is tuned on M2 Ultra for 7B
+            TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
+            VEC_C = 1
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = -1
+            UNROLL = 256
+            SUPPORT_WARP_SHUFFLE = True
+            if isinstance(len_S, int):
+                if len_S > len_R:
+                    TS, TR = 4, 16
+                else:
+                    TS, TR = 2, 64
+            else:
+                TS, TR = 1, 64
+        elif target.kind.name == "rocm":
+            VEC_C = 4
+            # TODO: set LOAD_V_SHARED = False for now
+            # rocm might have some issues when load/store of shared do not belong to same data type
+            # and only works for certain vector lens, our commonly useful vector lens are in 4
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = 8
+            UNROLL = 256
+            if isinstance(len_S, int):
+                if len_S > len_R:
+                    TS, TR = 1, 128
+                else:
+                    TS, TR = 8, 64
+            else:
+                TS, TR = 1, 64
+        elif target.kind.name == "opencl" and (
+            ("android" in str(target.host)) or ("adreno" in str(target.attrs))
+        ):
+            TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
+            VEC_C = 8
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = -1
+            UNROLL = 8
+            TS, TR = 2, 32
+        elif target.kind.name == "vulkan":
+            VEC_C = 4
+            LOAD_V_SHARED = True
+            LOAD_V_VEC = 4
+            UNROLL = 256
+            if isinstance(len_S, int):
+                if len_S > len_R:
+                    TS, TR = 4, 32
+                else:
+                    TS, TR = 16, 32
+            else:
+                TS, TR = 1, 64
+        elif target.kind.name == "opencl" and "mali" in str(target.attrs):
+            VEC_C = 8
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = -1
+            UNROLL = 64
+            TS, TR = 1, 64
+        else:
+            VEC_C = 1
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = -1
+            UNROLL = 64
+            TS, TR = 1, 64
+
+        while TS * TR > int(target.attrs["max_num_threads"]):
+            if TS > 1:
+                TS //= 2
+            else:
+                TR //= 2
+
+        TILE_S, TILE_R = (
+            1,
+            (
+                len_c
+                if len_c > 1
+                else max(get_max_factor(len_r, [TR * 1, TR * 2, TR * 4, TR * 8]) // TR, 1)
+            ),
+        )
+        VEC_C = min(get_max_factor(TILE_R, [1, 2, 4, 8]), VEC_C)
+
+        return apply(
+            sch,
+            gemv=block,
+            TAG_S=TAG_S,
+            TAG_R=TAG_R,
+            TS=TS,
+            TR=TR,
+            TILE_S=TILE_S,
+            TILE_R=TILE_R,
+            VEC_LOAD=VEC_LOAD,
+            VEC_C=VEC_C,
+            LOAD_V_SHARED=LOAD_V_SHARED,
+            LOAD_V_VEC=LOAD_V_VEC,
+            UNROLL=UNROLL,
+            SUPPORT_WARP_SHUFFLE=SUPPORT_WARP_SHUFFLE,
+        )
+
+    def sch_outer_reduction(  # pylint: disable=too-many-arguments, invalid-name, unused-argument
+        self,
+        sch: s_tir.Schedule,
+        target: Target,
+        block: s_tir.schedule.SBlockRV,
+        vector_input_buffers: list[tirx.Buffer],
+        epilogue_info: SBlockInfo | None,
+    ):
+        """Schedule the outer reduction block."""
+
+        def get_max_factor(n, factors):
+            factors = sorted(factors, reverse=True)
+            for factor in factors:
+                if n % factor == 0:
+                    return factor
+            return 1
+
+        def apply(
+            sch: s_tir.Schedule,
+            gemv,
+            TAG_S,
+            TAG_R,
+            TS,
+            TR,
+            SCALE_PACK,
+            DEC_PACK,
+            VEC_LOAD,
+            VEC_C,
+            LOAD_V_SHARED,
+            LOAD_V_VEC,
+            UNROLL,
+            LOAD_V_TILE,
+        ):
+            # rfactor: reduce to tx * vec_c
+            batch, s, r, c = sch.get_loops(block=gemv)
+            s = sch.fuse(batch, s)
+            r = sch.fuse(r, c)
+            bx, ts = sch.split(s, factors=[None, TS], preserve_unit_iters=True)
+            r, v_tile, tr, tile_r, vec_c = sch.split(
+                r, factors=[None, LOAD_V_TILE, TR, SCALE_PACK, DEC_PACK], preserve_unit_iters=True
+            )
+            sch.reorder(bx, ts, r, v_tile, tile_r, tr, vec_c)
+            tr_vec_c = sch.fuse(tr, vec_c)
+            rf = sch.rfactor(tr_vec_c, 0)
+
+            # rfactor: reduce to tx
+            bx, ts, tr_vec_c = sch.get_loops(block=gemv)
+            tr, vec_c = sch.split(tr_vec_c, factors=[TR, None], preserve_unit_iters=True)
+            rf2 = sch.rfactor(tr, 0)
+
+            # bind, vectorize compute
+            bx, ts, r, v_tile, tile_r, tr_vec_c = sch.get_loops(block=rf)
+            tr, vec_c = sch.split(tr_vec_c, factors=[TR, DEC_PACK])
+            sch.reorder(bx, ts, tr, r, v_tile, tile_r, vec_c)
+            # sch.bind(batch, "blockIdx.z")
+            sch.bind(bx, "blockIdx.x")
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+            auto_vectorize(sch, vec_c, VEC_C)
+
+            # decompose independent scale read to outer loop
+            block_rf_stmt = sch.get(rf)
+            if len(block_rf_stmt.reads) >= 3:
+                As_local = sch.cache_read(rf, read_buffer_index=2, storage_scope="local")
+                sch.compute_at(As_local, v_tile, preserve_unit_loops=True)
+                # *tile_thr, vec_s = sch.get_loops(block=As_local)
+                # sch.vectorize(vec_s)
+
+            Aq_local = sch.cache_read(rf, read_buffer_index=1, storage_scope="local")
+            sch.compute_at(Aq_local, tile_r, preserve_unit_loops=True)
+            # *tile_thr, vec_s = sch.get_loops(block=Aq_local)
+            # sch.vectorize(vec_s)
+
+            if LOAD_V_SHARED:
+                V_shared = sch.cache_read(rf, read_buffer_index=0, storage_scope="shared")
+                sch.compute_at(V_shared, r, preserve_unit_loops=True)
+                l = sch.get_loops(block=V_shared)[-1]
+                _, v_tile, ts, tr, vec = sch.split(
+                    l, factors=[None, LOAD_V_TILE, TS, TR, LOAD_V_VEC], preserve_unit_iters=True
+                )
+                sch.bind(tr, TAG_R)
+                sch.bind(ts, TAG_S)
+                auto_vectorize(sch, vec, LOAD_V_VEC)
+
+            # reduce tile_s * tr * vec to tile_s * tr
+            sch.reverse_compute_at(rf2, loop=bx, preserve_unit_loops=True)
+            tr, vec_c, ts = sch.get_loops(block=rf2)[1:]
+            sch.reorder(ts, tr, vec_c)
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+
+            # reduce tile_s * tr to tile_s
+            sch.reverse_compute_at(gemv, loop=bx, preserve_unit_loops=True)
+            tr, ts = sch.get_loops(block=gemv)[1:]
+            sch.reorder(ts, tr)
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+
+            sch.decompose_reduction(rf, loop=sch.get_loops(block=rf)[2])
+            sch.decompose_reduction(rf2, loop=sch.get_loops(block=rf2)[-1])
+
+            sch.set_scope(rf, buffer_index=0, storage_scope="local")
+            sch.set_scope(rf2, buffer_index=0, storage_scope="local")
+
+            sch.annotate(
+                block_or_loop=sch.get_loops(rf2)[3],
+                ann_key="pragma_auto_unroll_max_step",
+                ann_val=UNROLL,
+            )
+            sch.annotate(
+                block_or_loop=sch.get_loops(rf2)[3], ann_key="pragma_unroll_explicit", ann_val=1
+            )
+
+            # Schedule epilogue
+            if epilogue_info is not None:
+                epilogue = epilogue_info.block_rv
+                if is_broadcast_epilogue(sch, block, epilogue):
+                    sch.reverse_compute_at(epilogue, bx)
+                    sch.set_scope(block, 0, "shared")
+                    _, _, *s = sch.get_loops(epilogue)  # pylint: disable=invalid-name
+                    _, ts = sch.split(sch.fuse(*s), factors=[None, TS])
+                    sch.bind(ts, TAG_S)
+                else:
+                    sch.reverse_compute_at(epilogue, bx, preserve_unit_loops=True)
+                    ts_tile_s = sch.fuse(*sch.get_loops(epilogue)[1:])
+                    ts_tile_s = sch.get_loops(epilogue)[-1]
+                    ts, _ = sch.split(ts_tile_s, factors=[TS, None], preserve_unit_iters=True)
+                    sch.bind(ts, TAG_S)
+                    sch.set_scope(block, 0, "local")
+            return sch
+
+        # Specify the `len_tx` and `len_ty` according to the loop extent
+        batch, s, r, c = sch.get_loops(block=block)
+        _, len_s, len_r, len_c = (
+            get_extent(sch, batch),
+            get_extent(sch, s),
+            get_extent(sch, r),
+            get_extent(sch, c),
+        )
+
+        DEC_PACK = 8
+        SCALE_PACK = 4
+
+        if target.kind.name == "opencl" and (
+            ("android" in str(target.host)) or ("adreno" in str(target.attrs))
+        ):
+            TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
+            VEC_C = 8
+            UNROLL = 8
+            TS, TR = 64, 4
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = 4
+            LOAD_V_TILE = 8
+        elif target.kind.name == "metal":
+            TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
+            VEC_C = 4
+            UNROLL = 8
+            TS, TR = 128, 4
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = 4
+            LOAD_V_TILE = 4
+        else:
+            return None
+
+        if LOAD_V_SHARED is False:
+            LOAD_V_TILE = 1
+
+        if not isinstance(len_r, int) or len_r < LOAD_V_TILE * TR * SCALE_PACK * DEC_PACK:
+            return None
+
+        if not isinstance(len_s, int):
+            TS, TR = 256, 1
+            LOAD_V_SHARED = True
+
+        if isinstance(len_s, int) and len_s > 96000:
+            return None
+
+        _, TILE_R = (
+            1,
+            (
+                len_c
+                if len_c > 1
+                else max(get_max_factor(len_r, [TR * 1, TR * 2, TR * 4, TR * 8]) // TR, 1)
+            ),
+        )
+        LOAD_V_VEC = min(get_max_factor(TILE_R, [1, 2, 4, 8]), LOAD_V_VEC)
+        VEC_LOAD = 1
+
+        return apply(
+            sch,
+            gemv=block,
+            TAG_S=TAG_S,
+            TAG_R=TAG_R,
+            TS=TS,
+            TR=TR,
+            SCALE_PACK=SCALE_PACK,
+            DEC_PACK=DEC_PACK,
+            VEC_LOAD=VEC_LOAD,
+            VEC_C=VEC_C,
+            LOAD_V_SHARED=LOAD_V_SHARED,
+            LOAD_V_VEC=LOAD_V_VEC,
+            UNROLL=UNROLL,
+            LOAD_V_TILE=LOAD_V_TILE,
+        )
+
+    def sch_outer_reduction_fallback(  # pylint: disable=too-many-arguments, invalid-name, unused-argument
+        self,
+        sch: s_tir.Schedule,
+        target: Target,
+        block: s_tir.schedule.SBlockRV,
+        vector_input_buffers: list[tirx.Buffer],
+        epilogue_info: SBlockInfo | None,
+    ):
+        """Schedule the outer reduction block."""
+        # NOTE: Only Android is supported so far
+        if not (
+            target.kind.name == "opencl"
+            and (("android" in str(target.host)) or ("adreno" in str(target.attrs)))
+        ):
+            return None
+        batch, s, r, c = sch.get_loops(block)
+        len_s = get_extent(sch, s)
+
+        # The config is designed for Adreno
+        LOAD_V_SHARED = 1
+        tx_len = 128
+        vec_len = (4 if len_s > 4096 else 2) if isinstance(len_s, int) else 1
+        inner_r = 4
+
+        bx, tx, vec = sch.split(s, factors=[None, tx_len, vec_len])
+        r0, r1 = sch.split(r, factors=[None, inner_r])
+        sch.bind(batch, "blockIdx.y")
+        sch.bind(bx, "blockIdx.x")
+        sch.bind(tx, "threadIdx.x")
+        sch.reorder(bx, tx, r0, r1, c, vec)
+
+        sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=8)
+        sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1)
+
+        if LOAD_V_SHARED:
+            V_shared = sch.cache_read(block, vector_input_buffers[0], storage_scope="shared")
+            sch.compute_at(V_shared, bx, preserve_unit_loops=True)
+            l = sch.get_loops(block=V_shared)[-1]
+            _, tx, vec_r = sch.split(l, factors=[None, tx_len, 8], preserve_unit_iters=True)
+            sch.bind(tx, "threadIdx.x")
+            sch.vectorize(vec_r)
+
+        sch.vectorize(vec)
+
+        # Schedule epilogue
+        if epilogue_info is not None:
+            sch.reverse_compute_at(epilogue_info.block_rv, bx, preserve_unit_loops=True)
+            ts_tile_s = sch.get_loops(epilogue_info.block_rv)[-1]
+            ts, vec = sch.split(ts_tile_s, factors=[tx_len, vec_len], preserve_unit_iters=True)
+            sch.bind(ts, "threadIdx.x")
+            sch.vectorize(vec)
+            sch.set_scope(block, 0, "local")
+
+        sch.decompose_reduction(block, r0)
+
+        return sch

python/tvm/s_tir/dlight/gpu/general_reduction.py

@@ -0,0 +1,183 @@
+# 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
+"""Reduction rule for operators including softmax, layer norm, RMS norm, etc"""
+
+from tvm import arith, s_tir, tirx
+from tvm.target import Target
+
+from ..analysis import normalize_prim_func
+from ..base import try_inline_contiguous_spatial
+from .base import GPUScheduleRule
+
+
+class GeneralReduction(GPUScheduleRule):
+    """General Reduction rule for operators including softmax, layer norm, RMS norm, etc"""
+
+    def apply(  # pylint: disable=too-many-locals
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+        if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
+            return None
+
+        if target.kind.name == "cuda":
+            len_tx = 256
+            unroll_depth = 256
+        elif target.kind.name == "opencl":
+            len_tx = 256
+            unroll_depth = 64
+        else:
+            len_tx = 64
+            unroll_depth = 64
+
+        sch = s_tir.Schedule(func)
+        block_infos = normalize_prim_func(sch)
+        block_infos = try_inline_contiguous_spatial(sch, block_infos)
+        if block_infos is None or len(block_infos) == 0:
+            return None
+
+        dom_kind = block_infos[0].dom_kind()
+        num_leading_s = len(dom_kind) - len(dom_kind.lstrip("S"))
+        num_trailing_r = len(dom_kind) - len(dom_kind.rstrip("R"))
+
+        # Align the number of block iters of the last block.
+        num_last_block_iter = len(block_infos[-1].dom_kind())
+        if num_last_block_iter < len(dom_kind):
+            # If the last block is a scalar value, there is nothing left to
+            # tile/parallelise, and  `iters` is an empty tuple.
+            # Add a unit thread loop so the final write happens inside a valid
+            # GPU thread environment.
+            if num_last_block_iter == 0:
+                # Put every block (both the running reductions and the final
+                # scalar write) inside a trivial GPU thread. The very first block
+                # gets a `blockIdx.x` wrapper so that kernels still have a unique
+                # block scope.
+                for i, info in enumerate(block_infos):
+                    loop_rv = sch.add_unit_loop(info.block_rv)
+                    if i == 0:
+                        sch.bind(loop_rv, "blockIdx.x")
+                    else:
+                        sch.bind(loop_rv, "threadIdx.x")
+
+                return sch
+
+            def f_layout_mapping(*iters):
+                analyzer = arith.Analyzer()
+                # Try to match the iters of last block to the iters of the first block.
+                # For matched positions, use the iter from the input `iters`.
+                # For unmatched positions, use a new iter which is constant 0.
+                num_matched = 0
+                target_layout_iters = []
+                for block_iter in block_infos[0].iters:
+                    if num_matched < len(iters) and analyzer.can_prove_equal(
+                        block_iter.dom, block_infos[-1].iters[num_matched].dom
+                    ):
+                        target_layout_iters.append(iters[num_matched])
+                        num_matched += 1
+                    else:
+                        target_layout_iters.append(tirx.const(0, iters[0].ty))
+
+                # If all the iters of the last block can match, return the new layout.
+                if num_matched == len(iters):
+                    return target_layout_iters
+                # Otherwise, fallback to appending zeros in the beginning.
+                return [tirx.const(0, iters[0].ty)] * (len(dom_kind) - num_last_block_iter) + list(
+                    iters
+                )
+
+            index_map = tirx.IndexMap.from_func(f_layout_mapping, ndim=num_last_block_iter)
+            sch.transform_block_layout(block_infos[-1].block_rv, index_map)
+
+        try:
+            # TODO: fix num_leading_s = 0 case
+            assert num_trailing_r > 0
+            for block in block_infos[1:-1]:
+                assert block.dom_kind() == dom_kind
+            assert block_infos[-1].is_injective()
+            assert len(block_infos[-1].dom_kind()) <= len(dom_kind)
+        except AssertionError:
+            return None
+
+        if "R" not in block_infos[-1].dom_kind():
+            # The final block is a spatial block.
+            # It is possible that the loop order of the last block is not the same as
+            # previous blocks.
+            # Thus we reorder spatial loops to align with reduction loops for followup schedule.
+            # We first collect all the buffers written by reduction blocks,
+            # then in the final block, any index of those buffers are spatial.
+            reduced_buffers = []
+            for block_info in block_infos[:-1]:
+                for buffer_write in sch.get(block_info.block_rv).writes:
+                    reduced_buffers.append(buffer_write.buffer)
+
+            spatial_block = sch.get(block_infos[-1].block_rv)
+            spatial_loops = set()
+            block_var_to_loop_var = {}
+            loops = sch.get_loops(block_infos[-1].block_rv)
+            for block_iter, loop_rv in zip(spatial_block.iter_vars, loops):
+                block_var_to_loop_var[block_iter.var] = sch.get(loop_rv).loop_var
+
+            def _visit_expr(e: tirx.Expr):
+                if isinstance(e, tirx.Var) and e in block_var_to_loop_var:
+                    spatial_loops.add(block_var_to_loop_var[e])
+
+            for buffer_read in spatial_block.reads:
+                buffer = buffer_read.buffer
+                if buffer in reduced_buffers:
+                    for read_range in buffer_read.region:
+                        tirx.stmt_functor.post_order_visit(read_range.min, _visit_expr)
+                        tirx.stmt_functor.post_order_visit(read_range.extent, _visit_expr)
+
+            s_loops = []
+            other_loops = []
+            for loop_rv in loops:
+                loop = sch.get(loop_rv)
+                if loop.loop_var in spatial_loops or loop.extent == 1:
+                    s_loops.append(loop_rv)
+                else:
+                    other_loops.append(loop_rv)
+            sch.reorder(*s_loops, *other_loops)
+
+        loops = sch.get_loops(block_infos[-1].block_rv)
+        bx = sch.fuse(*loops[:num_leading_s])
+        r_loop, tx = sch.split(loops[-1], [None, len_tx])
+        sch.reorder(tx, r_loop)
+        sch.bind(bx, "blockIdx.x")
+        sch.bind(tx, "threadIdx.x")
+        sch.annotate(r_loop, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+        sch.annotate(r_loop, ann_key="pragma_unroll_explicit", ann_val=1)
+
+        for block in reversed(block_infos[:-1]):
+            block = block.block_rv
+            for i, _ in enumerate(sch.get(block).writes):
+                sch.set_scope(block, buffer_index=i, storage_scope="shared")
+            sch.compute_at(block, bx, preserve_unit_loops=True)
+            r_loop = sch.fuse(*sch.get_loops(block)[-num_trailing_r:])
+            r_loop, tx = sch.split(r_loop, [None, len_tx])
+            sch.reorder(tx, r_loop)
+            sch.bind(tx, "threadIdx.x")
+            sch.annotate(r_loop, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+            sch.annotate(r_loop, ann_key="pragma_unroll_explicit", ann_val=1)
+
+        # TODO: It's just a workaround to avoid unroll spatial loops, because of the bug of
+        # the pass lower-thread-allreduce. We should fix it in the future.
+        # sch.annotate(bx, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+        # sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
+        return sch

python/tvm/s_tir/dlight/gpu/low_batch_gemv.py

@@ -0,0 +1,744 @@
+# 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, F821
+"""A rule for low-batch GEMM / decode-GEMM using GEMV schedule."""
+
+from functools import reduce
+from typing import Literal
+
+import tvm_ffi
+
+from tvm import arith, s_tir, tirx
+from tvm.target import Target
+
+from ..analysis import (
+    SBlockInfo,
+    collect_block_iter_vars_used_in_access_region,
+    collect_vars_used_in_prim_expr,
+    get_max_shared_memory_per_block,
+    is_broadcast_epilogue,
+    normalize_prim_func,
+)
+from ..base import auto_vectorize, get_bytes, get_extent, try_inline_contiguous_spatial
+from .base import GPUScheduleRule
+
+
+def _get_reduction_expr(block: tirx.SBlock) -> tirx.Expr | None:
+    # Detect and return `Y` in `X[...] = X[...] + Y`
+    buffer_store = block.body
+    if not isinstance(buffer_store, tirx.BufferStore):
+        return None
+    if not isinstance(buffer_store.value, tirx.Add):
+        return None
+    if not tvm_ffi.structural_equal(
+        buffer_store.value.a,
+        tirx.BufferLoad(buffer_store.buffer, block.body.indices),
+        map_free_vars=True,
+    ):
+        return None
+    return buffer_store.value.b
+
+
+def is_gemv(sch: s_tir.Schedule, block_info: SBlockInfo) -> list[tirx.Buffer] | None:
+    """Check if the block is a low batch GEMM.
+
+    Parameters
+    ----------
+
+    sch : s_tir.Schedule
+        The schedule
+
+    block_info : SBlockInfo
+        The block info to be checked
+
+
+    Returns
+    -------
+    ret : Optional[List[tirx.Buffer]]
+        The vector-like buffers used in the low batch GEMM if it is a low batch GEMM,
+        otherwise None.
+    """
+    block = block_info.block_rv
+    block_stmt = sch.get(block)
+    conditions = []
+    conditions.append(block_info.is_reduction())
+    conditions.append(len(block_stmt.reads) >= 2)
+    conditions.append(len(block_stmt.writes) == 1)
+    conditions.append(_get_reduction_expr(block_stmt) is not None)
+    conditions.append(
+        len(collect_block_iter_vars_used_in_access_region(block_stmt, block_stmt.writes[0].region))
+        > 0
+    )
+    if not all(conditions):
+        return None
+    const_iter_vars = set(
+        iter_var.var
+        for iter_var in block_stmt.iter_vars
+        if isinstance(iter_var.dom.extent, tirx.IntImm)
+    )
+    if len(block_stmt.iter_vars) - len(const_iter_vars) != 1:
+        return None
+    symbolic_iter_var = next(
+        iter_var
+        for iter_var in block_stmt.iter_vars
+        if not isinstance(iter_var.dom.extent, tirx.IntImm)
+    )
+    if symbolic_iter_var.iter_type != tirx.stmt.IterVar.DataPar:
+        return None
+    ret = [
+        read.buffer
+        for read in block_stmt.reads
+        if len(
+            collect_block_iter_vars_used_in_access_region(block_stmt, read.region) & const_iter_vars
+        )
+        < len(const_iter_vars)
+        and len(
+            collect_block_iter_vars_used_in_access_region(block_stmt, read.region) & const_iter_vars
+        )
+        > 0
+    ]
+    return ret if 0 < len(ret) < len(block_stmt.reads) else None
+
+
+def detect_dominant_read(block: tirx.SBlock, const_iter_vars: set[tirx.Var]) -> tirx.Expr:
+    """Detect the dominant read indices in the block."""
+    dominant_read = None
+    num_read_iters = -1
+    for buffer_region in block.reads:
+        tir_vars = (
+            collect_block_iter_vars_used_in_access_region(block, buffer_region.region)
+            & const_iter_vars
+        )
+        if num_read_iters < len(tir_vars):
+            num_read_iters = len(tir_vars)
+            dominant_read = buffer_region
+    assert dominant_read is not None
+    (result,) = dominant_read.buffer.offset_of([e.min for e in dominant_read.region])
+    return result
+
+
+def normalize(
+    sch: s_tir.Schedule,
+    block_info: SBlockInfo,
+) -> bool | None:
+    """Normalize the main block."""
+    block_stmt: tirx.SBlock = sch.get(block_info.block_rv)
+    const_iter_vars = set(
+        iter_var.var
+        for iter_var in block_stmt.iter_vars
+        if isinstance(iter_var.dom.extent, tirx.IntImm)
+    )
+    dynamic_iter_vars = set(
+        iter_var.var for iter_var in block_stmt.iter_vars if iter_var.var not in const_iter_vars
+    )
+    access = arith.normalize_to_iter_sum(
+        detect_dominant_read(block_stmt, const_iter_vars),
+        input_iters={i.var: i.dom for i in block_stmt.iter_vars},
+    )
+    buffers_use_vars = [
+        collect_block_iter_vars_used_in_access_region(block_stmt, buf.region)
+        for buf in block_stmt.writes
+    ]
+    buffers_use_vars.extend(
+        [
+            collect_block_iter_vars_used_in_access_region(block_stmt, buf.region)
+            for buf in block_stmt.reads
+        ]
+    )
+    if collect_vars_used_in_prim_expr(access.base) & set(
+        iter_var.var for iter_var in block_stmt.iter_vars
+    ):
+        return None
+    iter_to_info = {i.var: i for i in block_info.iters}
+    batch_loops, s_loops, r_loops = [], [], []
+    inner_axis = access.args[-1].source.source
+    is_inner_reduction = iter_to_info[inner_axis].kind == "R"
+
+    for split_expr in access.args:
+        var = split_expr.source.source
+        info = iter_to_info.get(var)
+        loop = info.loop_rv
+        is_reduction = info.kind == "R"
+        # No C loops as we do not compute_inline weights into main block
+        if is_reduction:
+            r_loops.append(loop)
+        elif all([var in buf_vars for buf_vars in buffers_use_vars]):
+            batch_loops.append(loop)
+        else:
+            s_loops.append(loop)
+
+    assert s_loops
+    assert r_loops
+    dynamic_loops = [iter_to_info[var].loop_rv for var in dynamic_iter_vars]
+    assert len(dynamic_loops) == 1
+    sch.reorder(*dynamic_loops, *s_loops, *r_loops)
+    sch.fuse(*s_loops)
+    sch.fuse(*r_loops)
+    return is_inner_reduction
+
+
+class LowBatchGEMV(GPUScheduleRule):
+    """A rule for low batch GEMM / decode-GEMM."""
+
+    def __init__(self, bucket=4):
+        self.bucket = bucket
+
+    def apply(  # pylint: disable=too-many-locals,too-many-branches,too-many-return-statements
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+        if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
+            return None
+        sch = s_tir.Schedule(func)
+        block_infos = normalize_prim_func(sch)
+        if block_infos is None:
+            return None
+        reduction_block_infos = [
+            block_info for block_info in block_infos if block_info.is_reduction()
+        ]
+        if len(reduction_block_infos) != 1:
+            return None
+        reduction_block_info = reduction_block_infos[0]
+        vector_input_buffers = is_gemv(sch, reduction_block_info)
+        if vector_input_buffers is None:
+            return None
+        batch_pad = self.bucket
+        pad_value = [
+            iter.dom if isinstance(iter.dom, int) else batch_pad
+            for iter in reduction_block_info.iters
+        ]
+        sch.pad_einsum(reduction_block_info.block_rv, pad_value)
+        block_infos = normalize_prim_func(sch)
+        dequantize_block = None
+        pad_input_block = None
+        for block_info in block_infos:
+            if "dequantize" in block_info.name:
+                dequantize_block = block_info.block_rv
+            elif "pad" in block_info.name and len(sch.get_producers(block_info.block_rv)) == 0:
+                pad_input_block = block_info.block_rv
+        block_infos = [
+            block_info
+            for block_info in block_infos
+            if "pad" not in block_info.name and "dequantize" not in block_info.name
+        ]
+        block_infos = try_inline_contiguous_spatial(sch, block_infos)
+        if len(block_infos) == 1:
+            epilogue = None
+        elif len(block_infos) == 2:
+            epilogue = block_infos[1]
+            if not epilogue.is_injective():
+                return None
+        else:
+            return None
+
+        block_info = block_infos[0]
+        if len(block_info.iters) not in [2, 3]:
+            # either [B, S, R] = [B, S, R] * [B, R]
+            # or [S, R] = [S, R] * [R]
+            return None
+        block = block_info.block_rv
+        vector_input_buffers = is_gemv(sch, block_info)
+        if vector_input_buffers is None:
+            return None
+
+        # Step 1. Normalize the block, merge spatial and reduction iters
+        is_inner_reduction = normalize(sch, block_info)
+        # Step 2. Do the scheduling
+        if is_inner_reduction is None:
+            return None
+        elif is_inner_reduction:
+            self.sch_inner_reduction(
+                sch,
+                target,
+                block,
+                dequantize_block,
+                pad_input_block,
+                vector_input_buffers,
+                epilogue,
+                batch_pad,
+            )
+            return sch
+        elif self.bucket <= 4:
+            self.sch_outer_reduction(
+                sch,
+                target,
+                block,
+                dequantize_block,
+                pad_input_block,
+                vector_input_buffers,
+                epilogue,
+                batch_pad,
+            )
+            return sch
+        else:
+            return None
+
+    def sch_inner_reduction(  # pylint: disable=too-many-arguments, invalid-name, unused-argument
+        self,
+        sch: s_tir.Schedule,
+        target: Target,
+        block: s_tir.schedule.SBlockRV,
+        dequantize_block: s_tir.schedule.SBlockRV | None,
+        pad_input_block: s_tir.schedule.SBlockRV | None,
+        vector_input_buffers: list[tirx.Buffer],
+        epilogue_info: SBlockInfo | None,
+        batch_pad: int,
+    ):
+        """Schedule the inner reduction block."""
+
+        def get_max_factor(n, factors):
+            factors = sorted(factors, reverse=True)
+            for factor in factors:
+                if n % factor == 0:
+                    return factor
+            return 1
+
+        def apply(
+            sch: s_tir.Schedule,
+            gemv,
+            TAG_S,
+            TAG_R,
+            TS,
+            TR,
+            TILE_S,
+            TILE_R,
+            VEC_LOAD,
+            VEC_C,
+            LOAD_V_SHARED,
+            LOAD_V_VEC,
+            UNROLL,
+        ):
+            # rfactor: reduce to tx * vec_c
+
+            _, s, r = sch.get_loops(block=gemv)
+            bx, ts, tile_s = sch.split(s, factors=[None, TS, TILE_S], preserve_unit_iters=True)
+            r, tr, tile_r_vec_n, vec_c = sch.split(
+                r, factors=[None, TR, TILE_R // VEC_C, VEC_C], preserve_unit_iters=True
+            )
+            sch.reorder(r, tile_r_vec_n, tr, vec_c)
+            tr_vec_c = sch.fuse(tr, vec_c)
+            rf = sch.rfactor(tr_vec_c, 0)
+
+            # rfactor: reduce to tx
+            _, bx, ts, tile_s, tr_vec_c = sch.get_loops(block=gemv)
+            tr, vec_c = sch.split(tr_vec_c, factors=[TR, None], preserve_unit_iters=True)
+            rf2 = sch.rfactor(tr, 0)
+            # bind, vectorize compute
+            batch_loop, bx, ts, tile_s, r, tile_r_vec_n, tr_vec_c = sch.get_loops(block=rf)
+            tr, vec_c = sch.split(tr_vec_c, factors=[TR, None], preserve_unit_iters=True)
+            sch.reorder(bx, ts, tr, r, tile_s, tile_r_vec_n, vec_c)
+            sch.bind(bx, "blockIdx.x")
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+            sch.vectorize(vec_c)
+            by, batch = sch.split(batch_loop, factors=[None, batch_pad])
+            sch.bind(by, "blockIdx.y")
+            sch.reorder(bx, ts, tr, r, batch)
+
+            shared_mem_usage = 0
+            for buf in vector_input_buffers:
+                buf_size = reduce(
+                    lambda x, y: x * y, buf.shape, tirx.IntImm(buf.shape[0].ty, 1)
+                ) * get_bytes(buf.dtype)
+                shared_mem_usage += buf_size
+            max_smem = get_max_shared_memory_per_block(target)
+            LOAD_V_SHARED = (
+                LOAD_V_SHARED
+                and isinstance(shared_mem_usage, tirx.IntImm)
+                and shared_mem_usage.value <= max_smem
+            )
+
+            # vectorize load A
+            # (TODO) this is now actually problematic since the number of loops is dependent on the
+            # number of dimensions of A_q
+            if dequantize_block is not None:
+                sch.compute_at(dequantize_block, r, preserve_unit_loops=True)
+                sch.set_scope(dequantize_block, 0, "local")
+
+                s_local, r_local = sch.get_loops(block=dequantize_block)[-2:]
+                s_local, vec_load = sch.split(
+                    s_local, factors=[None, VEC_LOAD], preserve_unit_iters=True
+                )
+                sch.reorder(s_local, r_local, vec_load)  # either s_local or r_local should be 1
+                sch.vectorize(vec_load)
+
+            # load vector into shared memory, shape should be the whole vector
+            if LOAD_V_SHARED:
+                assert len(vector_input_buffers) == 1
+                V_shared = sch.cache_read(rf, read_buffer_index=0, storage_scope="shared")
+                sch.compute_at(V_shared, tr, preserve_unit_loops=True)
+                l = sch.get_loops(block=V_shared)[-1]
+                loop: tirx.For = sch.get(l)
+                if isinstance(loop.extent, tirx.IntImm):
+                    # avoid introducing predicates when vector length is too large
+                    vec_length = max(
+                        min(
+                            get_max_factor(
+                                (int)(loop.extent),
+                                [TS * TR * 1, TS * TR * 2, TS * TR * 4, TS * TR * 8],
+                            )
+                            // TS
+                            // TR,
+                            LOAD_V_VEC,
+                        ),
+                        1,
+                    )
+                else:
+                    vec_length = LOAD_V_VEC
+                if TAG_R == "threadIdx.x":
+                    _, ty, tx, vec = sch.split(
+                        l, factors=[None, TS, TR, vec_length], preserve_unit_iters=True
+                    )
+                else:
+                    _, ty, tx, vec = sch.split(
+                        l, factors=[None, TR, TS, vec_length], preserve_unit_iters=True
+                    )
+                sch.bind(ty, "threadIdx.y")
+                sch.bind(tx, "threadIdx.x")
+                sch.vectorize(vec)
+            if pad_input_block is not None:
+                sch.compute_inline(pad_input_block)
+
+            # reduce tile_s * tr * vec to tile_s * tr
+            sch.reverse_compute_at(rf2, loop=bx, preserve_unit_loops=True)
+            tr, vec_c, batch_loop, *ts_tile_s = sch.get_loops(block=rf2)[2:]
+            ts_tile_s = sch.fuse(*ts_tile_s)
+            ts_o, ts_i, tile_s = sch.split(
+                ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True
+            )
+            tile_s, vec_s = sch.split(
+                tile_s,
+                factors=[None, get_max_factor(TILE_S, [1, 2, 4, 8])],
+                preserve_unit_iters=True,
+            )
+            assert sch.get(ts_o).extent.value == 1
+            ts = sch.fuse(ts_o, ts_i)
+            sch.reorder(ts, tr, tile_s, batch_loop, vec_s, vec_c)
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+            sch.vectorize(vec_s)
+
+            # reduce tile_s * tr to tile_s
+            sch.reverse_compute_at(gemv, loop=bx, preserve_unit_loops=True)
+
+            tr, batch_loop, *ts_tile_s = sch.get_loops(block=gemv)[2:]
+            ts_tile_s = sch.fuse(*ts_tile_s)
+            ts_o, ts_i, tile_s = sch.split(
+                ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True
+            )
+            assert sch.get(ts_o).extent.value == 1
+            ts = sch.fuse(ts_o, ts_i)
+            sch.reorder(tile_s, batch_loop, ts, tr)
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+
+            sch.decompose_reduction(rf, loop=sch.get_loops(block=rf)[4])
+            sch.decompose_reduction(rf2, loop=sch.get_loops(block=rf2)[-1])
+
+            sch.set_scope(rf, buffer_index=0, storage_scope="local")
+            sch.set_scope(rf2, buffer_index=0, storage_scope="local")
+
+            unroll_factor = UNROLL
+
+            sch.annotate(
+                block_or_loop=sch.get_loops(rf)[4],
+                ann_key="pragma_auto_unroll_max_step",
+                ann_val=unroll_factor,
+            )
+            sch.annotate(
+                block_or_loop=sch.get_loops(rf)[4], ann_key="pragma_unroll_explicit", ann_val=1
+            )
+
+            sch.annotate(
+                block_or_loop=sch.get_loops(rf2)[4],
+                ann_key="pragma_auto_unroll_max_step",
+                ann_val=unroll_factor,
+            )
+            sch.annotate(
+                block_or_loop=sch.get_loops(rf2)[4], ann_key="pragma_unroll_explicit", ann_val=1
+            )
+
+            if LOAD_V_SHARED:
+                sch.annotate(
+                    block_or_loop=sch.get_loops(V_shared)[-4],
+                    ann_key="pragma_unroll_explicit",
+                    ann_val=unroll_factor,
+                )
+                sch.annotate(
+                    block_or_loop=sch.get_loops(V_shared)[-4], ann_key="pragma_vectorize", ann_val=1
+                )
+
+            epilogue = sch.get_consumers(gemv)
+            # Schedule epilogue
+            if epilogue:
+                epilogue = epilogue[0]
+                if is_broadcast_epilogue(sch, block, epilogue):
+                    sch.reverse_compute_at(epilogue, bx)
+                    sch.set_scope(block, 0, "shared")
+                    _, _, _, *s = sch.get_loops(epilogue)  # pylint: disable=invalid-name
+                    _, tx = sch.split(sch.fuse(*s), factors=[None, TX])
+                    sch.bind(tx, TAG_S)
+                else:
+                    sch.reverse_compute_at(epilogue, bx, preserve_unit_loops=True)
+                    ts_tile_s = sch.fuse(*sch.get_loops(epilogue)[3:])
+                    ts_tile_s = sch.get_loops(epilogue)[-1]
+                    ts_o, ts_i, tile_s = sch.split(
+                        ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True
+                    )
+                    assert sch.get(ts_o).extent.value == 1
+                    ts = sch.fuse(ts_o, ts_i)
+                    sch.bind(ts, TAG_S)
+                    sch.set_scope(block, 0, "local")
+
+            return sch
+
+        # Specify the `len_tx` and `len_ty` according to the loop extent
+        _, s, r = sch.get_loops(block=block)
+        len_s, len_r = get_extent(sch, s), get_extent(sch, r)
+
+        TAG_S, TAG_R = "threadIdx.y", "threadIdx.x"
+        if target.kind.name == "cuda":
+            VEC_C = 4
+            LOAD_V_SHARED = True
+            LOAD_V_VEC = 8
+            UNROLL = 256
+            if isinstance(len_s, int):
+                if len_s > len_r:
+                    TS, TR = 4, 64
+                else:
+                    TS, TR = 16, 32
+        elif target.kind.name == "metal":
+            VEC_C = 4
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = -1
+            UNROLL = 8
+            if isinstance(len_s, int):
+                if len_s > len_r:
+                    TS, TR = 8, 32
+                else:
+                    TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
+                    TS, TR = 8, 32
+        elif target.kind.name == "rocm":
+            VEC_C = 4
+            LOAD_V_SHARED = True
+            LOAD_V_VEC = 8
+            UNROLL = 256
+            if isinstance(len_s, int):
+                if len_s > len_r:
+                    TS, TR = 1, 128
+                else:
+                    TS, TR = 8, 64
+        elif target.kind.name == "opencl" and "android" in str(target.host):
+            TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
+            VEC_C = 8
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = -1
+            UNROLL = 8
+            TS, TR = 2, 32
+        elif target.kind.name == "vulkan":
+            VEC_C = 4
+            LOAD_V_SHARED = True
+            LOAD_V_VEC = 4
+            UNROLL = 256
+            if isinstance(len_s, int):
+                if len_s > len_r:
+                    TS, TR = 4, 32
+                else:
+                    TS, TR = 16, 32
+        elif target.kind.name == "opencl" and "mali" in str(target.attrs):
+            VEC_C = 8
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = -1
+            UNROLL = 64
+            TS, TR = 1, 64
+        else:
+            VEC_C = 1
+            LOAD_V_SHARED = False
+            LOAD_V_VEC = -1
+            UNROLL = 64
+            TS, TR = 1, 64
+
+        if not isinstance(len_s, int):
+            TS, TR = 1, 64
+
+        while TS * TR > int(target.attrs["max_num_threads"]):
+            if TS > 1:
+                TS //= 2
+            else:
+                TR //= 2
+
+        TILE_S, TILE_R = 2, max(get_max_factor(len_r, [TR * 1, TR * 2, TR * 4, TR * 8]) // TR, 1)
+        VEC_C = min(get_max_factor(TILE_R, [1, 2, 4, 8]), VEC_C)
+        VEC_LOAD = 1
+        return apply(
+            sch,
+            gemv=block,
+            TAG_S=TAG_S,
+            TAG_R=TAG_R,
+            TS=TS,
+            TR=TR,
+            TILE_S=TILE_S,
+            TILE_R=TILE_R,
+            VEC_LOAD=VEC_LOAD,
+            VEC_C=VEC_C,
+            LOAD_V_SHARED=LOAD_V_SHARED,
+            LOAD_V_VEC=LOAD_V_VEC,
+            UNROLL=UNROLL,
+        )
+
+    def sch_outer_reduction(  # pylint: disable=too-many-arguments, invalid-name, unused-argument
+        self,
+        sch: s_tir.Schedule,
+        target: Target,
+        block: s_tir.schedule.SBlockRV,
+        dequantize_block: s_tir.schedule.SBlockRV | None,
+        pad_input_block: s_tir.schedule.SBlockRV | None,
+        vector_input_buffers: list[tirx.Buffer],
+        epilogue_info: SBlockInfo | None,
+        batch_pad: int,
+    ):
+        """Schedule the outer reduction block."""
+
+        # Need to detect from the block
+        DEC_PACK = 8
+        SCALE_PACK = 4
+
+        def apply(
+            sch: s_tir.Schedule,
+            main_block: s_tir.schedule.SBlockRV,
+            TAG_S: Literal["threadIdx.x", "threadIdx.y"],
+            TAG_R: Literal["threadIdx.x", "threadIdx.y"],
+            TS: int,
+            TR: int,
+            VEC: int,
+            UNROLL: int,
+        ):
+            # rfactor: reduce to tx * vec_c
+            b, s, r = sch.get_loops(main_block)
+            by, batch = sch.split(b, [None, batch_pad], preserve_unit_iters=True)
+            bx, ts = sch.split(s, [None, TS], preserve_unit_iters=True)
+            r, tr, scale_c, vec_c = sch.split(
+                r, [None, TR, SCALE_PACK, DEC_PACK], preserve_unit_iters=True
+            )
+            sch.reorder(by, bx, ts, r, batch, scale_c, tr, vec_c)
+            tr_vec_c = sch.fuse(tr, vec_c)
+            rf = sch.rfactor(tr_vec_c, 0)
+
+            # rfactor: reduce to tx
+            by, bx, ts, batch, tr_vec_c = sch.get_loops(block=main_block)
+            tr, vec_c = sch.split(tr_vec_c, [TR, DEC_PACK], preserve_unit_iters=True)
+            rf2 = sch.rfactor(tr, 0)
+
+            # bind, vectorize compute
+            by, bx, ts, r, batch, scale_c, tr_vec_c = sch.get_loops(block=rf)
+            tr, vec_c = sch.split(tr_vec_c, [TR, DEC_PACK], preserve_unit_iters=True)
+            sch.reorder(by, bx, ts, tr, r, scale_c, batch, vec_c)
+            sch.bind(by, "blockIdx.y")
+            sch.bind(bx, "blockIdx.x")
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+            auto_vectorize(sch, vec_c, VEC)
+
+            if dequantize_block is not None:
+                sch.compute_at(dequantize_block, scale_c, preserve_unit_loops=True)
+                sch.set_scope(dequantize_block, 0, "local")
+                auto_vectorize(sch, sch.fuse(*sch.get_loops(dequantize_block)[6:]), VEC)
+
+                B0_local = sch.cache_read(dequantize_block, 0, "local")
+                sch.compute_at(B0_local, r, preserve_unit_loops=True)
+                auto_vectorize(sch, sch.fuse(*sch.get_loops(B0_local)[5:]), VEC)
+
+                B1_local = sch.cache_read(dequantize_block, 1, "local")
+                sch.compute_at(B1_local, r, preserve_unit_loops=True)
+                auto_vectorize(sch, sch.fuse(*sch.get_loops(B1_local)[5:]), VEC)
+            else:
+                # Only support quantized workloads for now
+                sch = None
+                return
+
+            if LOAD_V_SHARED:
+                sch.set_scope(pad_input_block, 0, "shared")
+                sch.compute_at(pad_input_block, r, preserve_unit_loops=True)
+                sch.storage_align(pad_input_block, 0, axis=-2, factor=8, offset=1)
+                tr, ts, v = sch.split(sch.fuse(*sch.get_loops(pad_input_block)[5:]), [TR, TS, None])
+                sch.bind(tr, TAG_R)
+                sch.bind(ts, TAG_S)
+                auto_vectorize(sch, v, VEC)
+            else:
+                sch.compute_inline(pad_input_block)
+
+            # reduce tile_s * tr * vec to tile_s * tr
+            sch.reverse_compute_at(rf2, bx, preserve_unit_loops=True)
+            tr, vec_c, batch, ts = sch.get_loops(rf2)[2:]
+            sch.reorder(ts, tr, batch, vec_c)
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+
+            # reduce tile_s * tr to tile_s
+            sch.reverse_compute_at(main_block, bx, preserve_unit_loops=True)
+            tr, batch, ts = sch.get_loops(main_block)[2:]
+            sch.reorder(batch, ts, tr)
+            sch.bind(ts, TAG_S)
+            sch.bind(tr, TAG_R)
+            # unroll(batch, 1)
+
+            sch.decompose_reduction(rf, loop=sch.get_loops(block=rf)[4])
+            sch.decompose_reduction(rf2, loop=sch.get_loops(block=rf2)[4])
+
+            sch.set_scope(rf, buffer_index=0, storage_scope="local")
+            sch.set_scope(rf2, buffer_index=0, storage_scope="local")
+
+            epilogue = sch.get_consumers(main_block)
+            # Schedule epilogue
+            if epilogue:
+                epilogue = epilogue[0]
+                if is_broadcast_epilogue(  # pylint: disable=no-else-raise
+                    sch, main_block, epilogue
+                ):
+                    raise NotImplementedError
+                else:
+                    sch.reverse_compute_at(epilogue, bx, preserve_unit_loops=True)
+                    batch, ts = sch.get_loops(epilogue)[2:]
+                    sch.bind(ts, TAG_S)
+                    sch.set_scope(main_block, 0, "local")
+
+        if target.kind.name == "metal":
+            TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
+            TS, TR = 64, 4
+            LOAD_V_SHARED = True
+            VEC = 4
+            UNROLL = 8
+        else:
+            # fallback configuration
+            TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
+            TS, TR = 32, 4
+            LOAD_V_SHARED = False
+            VEC = 1
+            UNROLL = 64
+
+        return apply(
+            sch,
+            block,
+            TAG_S,
+            TAG_R,
+            TS,
+            TR,
+            VEC,
+            UNROLL,
+        )

python/tvm/s_tir/dlight/gpu/reduction.py

@@ -0,0 +1,305 @@
+# 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 rule for reduction."""
+
+# TODO: combine reduction rule and general reduction rule into one file.
+from collections.abc import Mapping
+
+import tvm_ffi
+
+from tvm import arith, s_tir, tirx
+from tvm.target import Target
+
+from ..analysis import (
+    SBlockInfo,
+    detect_dominant_read,
+    is_broadcast_epilogue,
+    normalize_prim_func,
+)
+from ..base import suggest_threads_per_block, try_inline_contiguous_spatial
+from .base import GPUScheduleRule
+
+
+def _get_reduction_expr(block: tirx.SBlock) -> tirx.Expr | None:
+    # Detect and return `Y` in `X[...] = X[...] + Y`
+    buffer_store = block.body
+    if not isinstance(buffer_store, tirx.BufferStore):
+        return None
+    if not isinstance(buffer_store.value, tirx.Add):
+        return None
+    if not tvm_ffi.structural_equal(
+        buffer_store.value.a,
+        tirx.BufferLoad(buffer_store.buffer, block.body.indices),
+        map_free_vars=True,
+    ):
+        return None
+    return buffer_store.value.b
+
+
+def _has_reduction_loop(block_info):
+    return any([info.kind == "R" for info in block_info.iters])
+
+
+class Reduction(GPUScheduleRule):
+    """A rule for Reduction."""
+
+    def apply(  # pylint: disable=too-many-locals,too-many-branches,too-many-return-statements
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+        if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
+            return None
+        sch = s_tir.Schedule(func)
+        block_infos = normalize_prim_func(sch)
+        if block_infos is None:
+            return None
+        block_infos = try_inline_contiguous_spatial(sch, block_infos)
+        if len(block_infos) == 1:
+            epilogue = None
+        elif len(block_infos) == 2:
+            epilogue = block_infos[1]
+            if not epilogue.is_injective():
+                return None
+        else:
+            return None
+
+        block_info = block_infos[0]
+        block = block_info.block_rv
+        block_stmt = sch.get(block)
+
+        # Step 1. Check reduction block
+        if (
+            (not block_info.is_reduction())
+            or (not _has_reduction_loop(block_info))
+            or len(block_stmt.writes) != 1
+            or _get_reduction_expr(block_stmt) is None
+        ):
+            return None
+        # Step 2. Normalize the block, merge spatial and reduction iters
+        is_inner_reduction, c_factor, loop_order, s_split_index = self._normalize(
+            sch,
+            block_info,
+            arith.normalize_to_iter_sum(
+                detect_dominant_read(block_stmt),
+                input_iters={i.var: i.dom for i in block_stmt.iter_vars},
+            ),
+        )
+        if is_inner_reduction is None and c_factor is None:
+            return None
+        # Step 3. Do the scheduling
+        if is_inner_reduction:
+            self._sch_inner_reduction(
+                sch, target, block, c_factor, epilogue, loop_order, s_split_index
+            )
+        else:
+            self._sch_inner_spatial(
+                sch, target, block, block_info, c_factor, epilogue, loop_order, s_split_index
+            )
+        return sch
+
+    def _normalize(  # pylint: disable=too-many-branches
+        self,
+        sch: s_tir.Schedule,
+        block_info: SBlockInfo,
+        access: arith.IterSumExpr,
+    ) -> tuple[bool | None, int | None, Mapping[int, int] | None, int | None]:
+        if access.base != 0:
+            return None, None, None, None
+        iter_to_info = {i.var: i for i in block_info.iters}
+        s_loops, r_loops, c_loops, c_factor = [], [], [], None
+        s_split_loop, s_split_index = None, None
+        for split_expr in access.args:
+            var = split_expr.source.source
+            info = iter_to_info.pop(var)
+            loop = info.loop_rv
+            is_inner_reduction = info.kind == "R"
+            if split_expr.lower_factor > 1:
+                if c_loops:
+                    return None, None, None, None
+                s_split_loop = loop
+                s_split_index = len(s_loops)
+                loop, c_loop = sch.split(loop, factors=[None, split_expr.lower_factor])
+                c_loops.append(c_loop)
+                if not is_inner_reduction:
+                    c_factor = split_expr.lower_factor
+            if is_inner_reduction:
+                r_loops.append(loop)
+            else:
+                s_loops.append(loop)
+
+        if iter_to_info:
+            for var, info in iter_to_info.items():
+                if info.kind == "S" and info.dom == 1:
+                    s_loops.append(info.loop_rv)
+                else:
+                    return None, None, None, None
+
+        loop_order = {}
+        s_block_var_loops = []
+        for i in block_info.iters:
+            if i.loop_rv in s_loops or i.loop_rv == s_split_loop:
+                s_block_var_loops.append(i.loop_rv)
+
+        for i in range(len(s_block_var_loops)):
+            for j in range(len(s_loops)):
+                if s_block_var_loops[i] == s_loops[j]:
+                    loop_order[i] = j
+                    break
+                if s_block_var_loops[i] == s_split_loop:
+                    loop_order[i] = s_split_index
+                    break
+
+        assert s_loops
+        assert r_loops
+        if len(s_loops) != len([i for i in block_info.iters if i.kind == "S"]):
+            return None, None, None, None
+        if not c_loops:
+            c_loops = [sch.add_unit_loop(block_info.block_rv)]
+        sch.reorder(*s_loops, *r_loops, *c_loops)
+        sch.fuse(*s_loops)
+        sch.fuse(*r_loops)
+        return is_inner_reduction, c_factor, loop_order, s_split_index
+
+    def _sch_inner_reduction(  # pylint: disable=too-many-arguments
+        self,
+        sch: s_tir.Schedule,
+        target: Target,
+        block: s_tir.schedule.SBlockRV,
+        unroll_spatial_factor: int | None,
+        epilogue_info: SBlockInfo | None,
+        loop_order,
+        s_split_index,
+    ):
+        # pylint: disable=invalid-name
+        _, r, _ = sch.get_loops(block)
+        (len_tx,) = suggest_threads_per_block(  # pylint: disable=unbalanced-tuple-unpacking
+            target, [sch.get(r)]
+        )
+
+        _, tx = sch.split(r, factors=[None, len_tx])
+        # Schedule the RF block
+        rf = sch.rfactor(tx, 0)
+        bx, r, tx, _ = sch.get_loops(rf)
+        sch.reorder(bx, tx, r)
+        sch.bind(bx, "blockIdx.x")
+        sch.bind(tx, "threadIdx.x")
+        sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=256)
+        sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1)
+        sch.set_scope(rf, 0, "local")
+        sch.decompose_reduction(rf, r)
+        # Schedule the write back block
+        sch.reverse_compute_at(block, bx, preserve_unit_loops=True)
+        _, tx, *s = sch.get_loops(block)
+
+        if unroll_spatial_factor:
+            assert len(s) == len(loop_order)
+            new_order_s = [s[loop_order[i]] for i in range(len(s))]
+            sch.reorder(*new_order_s)
+            new_order_s[s_split_index], c = sch.split(
+                new_order_s[s_split_index], factors=[None, unroll_spatial_factor]
+            )
+            sch.reorder(*new_order_s, c)
+            s = sch.fuse(*new_order_s)
+            sch.reorder(s, tx, c)
+        else:
+            s = sch.fuse(*s)
+            sch.reorder(s, tx)
+        sch.bind(tx, "threadIdx.x")
+        # Schedule epilogue
+        if epilogue_info is not None:
+            epilogue = epilogue_info.block_rv
+            sch.reverse_compute_at(epilogue, bx, preserve_unit_loops=True)
+            if is_broadcast_epilogue(sch, block, epilogue):
+                sch.set_scope(block, 0, "shared")
+                _, *s = sch.get_loops(epilogue)  # pylint: disable=invalid-name
+                _, tx = sch.split(sch.fuse(*s), factors=[None, len_tx])
+                sch.bind(tx, "threadIdx.x")
+            else:
+                sch.set_scope(block, 0, "local")
+        # pylint: enable=invalid-name
+
+    def _sch_inner_spatial(
+        self,
+        sch: s_tir.Schedule,
+        _: Target,
+        block: s_tir.schedule.SBlockRV,
+        block_info: SBlockInfo,
+        unroll_spatial_factor: int | None,
+        epilogue_info: SBlockInfo | None,
+        loop_order,
+        s_split_index,
+    ):
+        # pylint: disable=invalid-name
+        s, r, _ = sch.get_loops(block)
+        len_tx, len_ty = 16, 16
+        s_factor = [i.dom for i in block_info.iters if i.kind == "S"][-1]
+        # get perfect spatial factor, spatial factor should be divide the innermost spatial loop so
+        # that the block after r_factor and be reversed compute at the original scope
+        while len_tx > 1:
+            if s_factor % len_tx == 0:
+                break
+            len_tx -= 1
+        _, _ = sch.split(s, factors=[None, len_tx])
+        _, ty = sch.split(r, factors=[None, len_ty])
+        # Schedule the RF block
+        rf = sch.rfactor(ty, 0)
+        bx, tx, r, ty, _ = sch.get_loops(rf)
+        sch.reorder(bx, tx, ty, r)
+        sch.bind(tx, "threadIdx.x")
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(bx, "blockIdx.x")
+        sch.set_scope(rf, 0, "local")
+        sch.decompose_reduction(rf, r)
+        # Schedule the write back block
+        sch.reverse_compute_at(block, bx, preserve_unit_loops=True)
+        _, r, *s = sch.get_loops(block)
+        if unroll_spatial_factor:
+            assert len(s) == len(loop_order)
+            new_order_s = [s[loop_order[i]] for i in range(len(s))]
+            sch.reorder(*new_order_s)
+            new_order_s[s_split_index], c = sch.split(
+                new_order_s[s_split_index], factors=[None, unroll_spatial_factor]
+            )
+            sch.reorder(*new_order_s, c)
+            s = sch.fuse(*new_order_s)
+            sch.reorder(s, c, r)
+        else:
+            s = sch.fuse(*s)
+            sch.reorder(s, r)
+        sch.bind(s, "threadIdx.x")
+        sch.bind(r, "threadIdx.y")
+
+        # Schedule epilogue
+        if epilogue_info is not None:
+            epilogue = epilogue_info.block_rv
+            sch.reverse_compute_at(epilogue, bx, preserve_unit_loops=True)
+            if is_broadcast_epilogue(sch, block, epilogue):
+                sch.set_scope(block, 0, "shared")
+                _, *s = sch.get_loops(epilogue)  # pylint: disable=invalid-name
+                _, tx, ty = sch.split(sch.fuse(*s), factors=[None, len_tx, len_ty])
+                sch.bind(tx, "threadIdx.x")
+                sch.bind(ty, "threadIdx.y")
+            else:
+                # The epilogue is element-wise without broadcasting.
+                # Thus the remaining spatial part should be bind to tx.
+                sch.set_scope(block, 0, "local")
+                _, *s = sch.get_loops(epilogue)  # pylint: disable=invalid-name
+                tx, _ = sch.split(sch.fuse(*s), factors=[len_tx, None])
+                sch.bind(tx, "threadIdx.x")
+        # pylint: enable=invalid-name

python/tvm/s_tir/dlight/gpu/rmsnorm.py

@@ -0,0 +1,143 @@
+# 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
+"""A RMS norm schedule rule for GPU operators."""
+
+import tvm
+from tvm import tirx
+from tvm.ir import Call
+from tvm.target import Target
+from tvm.tirx import BufferStore, SBlock
+from tvm.tirx.expr import BufferLoad, Cast
+
+from ..base import ScheduleRule
+
+
+def identify_cast_or_load_block(block: SBlock) -> bool:
+    if len(block.reads) != 1 or len(block.writes) != 1:
+        return False
+
+    if not isinstance(block.body, BufferStore):
+        return False
+    store = block.body
+
+    # check types
+    if isinstance(store.value, BufferLoad):
+        load = store.value
+    elif isinstance(store.value, Cast):
+        load = store.value.value
+        if not isinstance(load, BufferLoad):
+            return False
+    else:
+        return False
+
+    # check indices
+    if len(load.indices) != len(store.indices):
+        return False
+
+    for lhs, rhs in zip(load.indices, store.indices):
+        if not lhs.same_as(rhs):
+            return False
+
+    return True
+
+
+def identify_rsqrt_block(block: SBlock) -> bool:
+    if len(block.reads) != 1 or len(block.writes) != 1:
+        return False
+
+    if not isinstance(block.body, BufferStore):
+        return False
+    store = block.body
+
+    if not isinstance(store.value, Call):
+        return False
+    call = store.value
+    op = call.op
+
+    return op == tvm.ir.op.Op.get("tirx.rsqrt")
+
+
+class RMSNorm(ScheduleRule):
+    """A rule for RMS norm."""
+
+    def apply(  # pylint: disable=too-many-locals,missing-docstring
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> "tvm.s_tir.Schedule":
+        if target.kind.name == "cuda":
+            num_tx = 512
+        elif target.kind.name == "opencl":
+            num_tx = 256
+        else:
+            num_tx = 64
+
+        sch = tvm.s_tir.Schedule(func)
+        root = sch.get_sblock(name="root", func_name="main")
+
+        blocks = sch.get_child_blocks(root)
+
+        if not any([identify_rsqrt_block(sch.get(block)) for block in blocks]):
+            return None
+
+        read = sch.cache_read(block=blocks[0], read_buffer_index=0, storage_scope="local")
+        write = sch.cache_write(block=blocks[-1], write_buffer_index=0, storage_scope="local")
+
+        for block in blocks:
+            if identify_cast_or_load_block(sch.get(block)):
+                sch.compute_inline(block)
+
+        blocks = sch.get_child_blocks(root)
+
+        read, sqr, redsum, rsqrt, norm, write = blocks
+
+        if not identify_rsqrt_block(sch.get(rsqrt)):
+            return None
+
+        for name in [read, sqr, redsum, rsqrt, norm, write]:
+            loops = sch.get_loops(name)
+            sch.fuse(*loops[:-1])
+
+        block_loop, loops = sch.get_loops(block=read)
+        thread_loop, _, _ = sch.split(
+            loop=loops, factors=[num_tx, None, 8], preserve_unit_iters=True
+        )
+        sch.bind(block_loop, thread_axis="blockIdx.x")
+        sch.bind(thread_loop, thread_axis="threadIdx.x")
+        sch.vectorize(sch.get_loops(block=read)[-1])
+        sch.reverse_compute_at(block=sqr, loop=thread_loop)
+        sch.reverse_compute_at(block=redsum, loop=thread_loop)
+
+        sch.reverse_compute_at(block=rsqrt, loop=block_loop, index=-1)
+        sch.reverse_compute_at(block=norm, loop=block_loop, index=-1)
+        block_loop, loops = sch.get_loops(block=norm)
+        thread_loop, _, _ = sch.split(
+            loop=loops, factors=[num_tx, None, 8], preserve_unit_iters=True
+        )
+        sch.bind(thread_loop, thread_axis="threadIdx.x")
+
+        sch.reverse_compute_at(block=write, loop=thread_loop, index=-1)
+        sch.vectorize(sch.get_loops(block=write)[-1])
+
+        sch.set_scope(block=sqr, buffer_index=0, storage_scope="local")
+        sch.set_scope(block=redsum, buffer_index=0, storage_scope="local")
+        sch.set_scope(block=rsqrt, buffer_index=0, storage_scope="shared")
+        sch.set_scope(block=norm, buffer_index=0, storage_scope="local")
+
+        return sch

python/tvm/s_tir/dlight/gpu/transpose.py

@@ -0,0 +1,129 @@
+# 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.
+"""Reduction rule for operators including softmax, layer norm, RMS norm, etc"""
+
+from tvm import arith, s_tir, tirx
+from tvm.s_tir import Schedule
+from tvm.s_tir.schedule import SBlockRV
+from tvm.target import Target
+
+from ..analysis import detect_dominant_read, normalize_prim_func
+from ..base import try_inline_contiguous_spatial
+from .base import GPUScheduleRule
+
+
+class Transpose(GPUScheduleRule):
+    """Schedule rule for transpose"""
+
+    def is_transpose(self, sch: Schedule, block_rv: SBlockRV):
+        block = sch.get(block_rv)
+        if isinstance(block.body, tirx.BufferStore):
+            rhs = block.body.value
+            if isinstance(rhs, tirx.BufferLoad):
+                lhs_indices = block.body.indices
+                rhs_indices = rhs.indices
+                if list(lhs_indices) != list(rhs_indices) and set(lhs_indices) == set(rhs_indices):
+                    return True
+        return False
+
+    def apply(  # pylint: disable=too-many-locals
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> None | s_tir.Schedule | list[s_tir.Schedule]:
+        # pylint: disable=invalid-name
+        if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
+            return None
+        if target.kind.name == "cuda":
+            len_tx = 16
+            len_ty = 8
+            unroll_depth = 256
+        elif target.kind.name == "opencl":
+            len_tx = 16
+            len_ty = 8
+            unroll_depth = 64
+        else:
+            len_tx = 8
+            len_ty = 4
+            unroll_depth = 64
+        len_vec = 4
+
+        sch = s_tir.Schedule(func)
+        blocks = normalize_prim_func(sch)
+        transpose_block_idx = -1
+        for idx, block in reversed(list(enumerate(blocks))):
+            if self.is_transpose(sch, block.block_rv):
+                transpose_block_idx = idx
+                break
+            if not block.is_injective():
+                return None
+        if transpose_block_idx == -1:
+            return None
+        transpose_block = blocks[transpose_block_idx].block_rv
+
+        prologue = None  # the optional decoding block
+        if transpose_block_idx > 0:
+            spatials = try_inline_contiguous_spatial(sch, blocks[: transpose_block_idx - 1])
+            assert len(spatials) == 0
+            prologue = blocks[transpose_block_idx - 1].block_rv
+
+        loops = sch.get_loops(transpose_block)
+        if len(loops) != 2:
+            # transpose with more than 2 axes is not supported
+            return None
+
+        c_factor = 1
+        if prologue is not None:
+            block_stmt = sch.get(prologue)
+            result = arith.normalize_to_iter_sum(
+                detect_dominant_read(block_stmt),
+                input_iters={i.var: i.dom for i in block_stmt.iter_vars},
+            )
+            if len(result.args) > 0:
+                c_factor = int(result.args[0].lower_factor)
+
+        i, j = loops
+        i, vi = sch.split(i, factors=[None, c_factor], preserve_unit_iters=True)
+        bi, ti = sch.split(i, factors=[None, len_ty], preserve_unit_iters=True)
+        bj, tj = sch.split(j, factors=[None, len_tx], preserve_unit_iters=True)
+        sch.reorder(bi, bj, ti, tj, vi)
+        sch.bind(bi, "blockIdx.y")
+        sch.bind(bj, "blockIdx.x")
+        sch.bind(ti, "threadIdx.y")
+        sch.bind(tj, "threadIdx.x")
+        len_vec = min(len_vec, c_factor)
+        _, vi = sch.split(vi, factors=[None, len_vec])
+        if len_vec > 1:
+            sch.vectorize(vi)
+
+        cache_read = sch.cache_read(transpose_block, read_buffer_index=0, storage_scope="shared")
+        sch.compute_at(cache_read, bj)
+        loops = sch.get_loops(cache_read)[2:]
+        fused = sch.fuse(*loops)
+        _, ty, tx, v = sch.split(fused, factors=[None, len_ty, len_tx, c_factor])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+        sch.unroll(v)
+        sch.storage_align(block=cache_read, buffer_index=0, axis=0, factor=32, offset=1)
+
+        sch.annotate(bi, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+        sch.annotate(bi, ann_key="pragma_unroll_explicit", ann_val=1)
+
+        if prologue is not None:
+            sch.compute_inline(prologue)
+        return sch

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