chore: import upstream snapshot with attribution

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