chore: import upstream snapshot with attribution

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