chore: import upstream snapshot with attribution

tests/python/s_tir/schedule/test_tir_schedule_analysis.py

@@ -0,0 +1,459 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring
+# ruff: noqa: F401, F841
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.s_tir import Schedule
+from tvm.s_tir.meta_schedule.testing import te_workload
+from tvm.s_tir.schedule.analysis import (
+    TensorizeInfo,
+    get_auto_tensorize_mapping_info,
+    get_tensorize_loop_mapping,
+    is_output_block,
+    suggest_index_map,
+)
+from tvm.s_tir.tensor_intrin.cuda import (
+    WMMA_SYNC_16x16x16_f16f16f16_INTRIN,
+    WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
+)
+from tvm.s_tir.tensor_intrin.x86 import dot_product_16x4_u8i8i32_desc
+from tvm.script import tirx as T
+from tvm.te import create_prim_func
+from tvm.tirx import (
+    Evaluate,
+    For,
+    ForKind,
+    IndexMap,
+    Var,
+    decl_buffer,
+    floordiv,
+    floormod,
+)
+from tvm.tirx.analysis import expr_deep_equal
+from tvm.tirx.function import TensorIntrin
+from tvm.tirx.stmt_functor import pre_order_visit
+
+
+def _make_vars(*args: str) -> list[Var]:
+    return [Var(arg, dtype="int32") for arg in args]
+
+
+def _make_loops(loop_vars: list[Var], extents: list[int]) -> list[For]:
+    assert len(loop_vars) == len(extents)
+    return [
+        For(
+            loop_var=loop_var,
+            min=0,
+            extent=extent,
+            kind=ForKind.SERIAL,
+            body=Evaluate(0),
+        )
+        for loop_var, extent in zip(loop_vars, extents)
+    ]
+
+
+def test_suggest_index_map_simple():
+    i, j = _make_vars("i", "j")
+    index_map = suggest_index_map(
+        buffer=decl_buffer(shape=[8, 256]),
+        indices=[
+            floordiv(i, 16) * 4 + floordiv(j, 16),
+            floormod(i, 16) * 16 + floormod(j, 16),
+        ],
+        loops=_make_loops(
+            loop_vars=[i, j],
+            extents=[32, 64],
+        ),
+        predicate=True,
+    )
+    expected_index_map = IndexMap.from_func(
+        lambda x, y: [
+            floordiv(x, 4),
+            floordiv(y, 16),
+            floormod(x, 4),
+            floormod(y, 16),
+        ],
+    )
+    assert index_map.is_equivalent_to(expected_index_map)
+
+
+def test_suggest_index_map_bijective():
+    i, j = _make_vars("i", "j")
+    index_map = suggest_index_map(
+        buffer=decl_buffer(shape=[8]),
+        indices=[floormod(j, 4) * 2 + i],
+        loops=_make_loops(
+            loop_vars=[i, j],
+            extents=[2, 32],
+        ),
+        predicate=True,
+    )
+    expected_index_map = IndexMap.from_func(
+        lambda x: [
+            floormod(x, 2),
+            floordiv(x, 2),
+        ],
+    )
+    assert index_map.is_equivalent_to(expected_index_map)
+
+
+def test_suggest_index_map_winograd():
+    """use case in winograd conv where the indices are complicated"""
+    fused_outer, i3_3_fused, i4_0, i4_1 = _make_vars("fused_outer", "i3_3_fused", "i4_0", "i4_1")
+    eps = floordiv(fused_outer, 336) * 2 + floordiv(floormod(fused_outer, 16), 8)
+    nu = floordiv(floormod(fused_outer, 336), 112) * 2 + floordiv(floormod(fused_outer, 8), 4)
+    co = floormod(fused_outer, 4) * 32 + i3_3_fused
+    ci = (i4_0 * 32) + i4_1
+    buffer = decl_buffer(shape=[6, 6, 128, 128])
+    index_map = suggest_index_map(
+        buffer=buffer,
+        indices=[eps, nu, co, ci],
+        loops=_make_loops(
+            loop_vars=[fused_outer, i3_3_fused, i4_0, i4_1],
+            extents=[1008, 32, 4, 32],
+        ),
+        predicate=True,
+    )
+    expected_index_map = IndexMap.from_func(
+        lambda i0, i1, i2, i3: (
+            floordiv(i0, 2),
+            floordiv(i1, 2),
+            floormod(i0, 2),
+            floormod(i1, 2) * 4 + floordiv(i2, 32),
+            floormod(i2, 32),
+            floordiv(i3, 32),
+            floormod(i3, 32),
+        )
+    )
+    assert index_map.is_equivalent_to(expected_index_map)
+    inverse_index_map = index_map.inverse(buffer.shape)
+    expected_inverse_index_map = IndexMap.from_func(
+        lambda i0, i1, i2, i3, i4, i5, i6: (
+            ((i0 * 2) + i2),
+            i1 * 2 + floordiv(i3, 4),
+            floormod(i3, 4) * 32 + i4,
+            ((i5 * 32) + i6),
+        )
+    )
+    assert inverse_index_map.is_equivalent_to(expected_inverse_index_map)
+
+
+@tvm.script.ir_module
+class DenseTIRModule:
+    @T.prim_func(s_tir=True)
+    def main(
+        placeholder: T.Buffer((1024, 1024), "uint8"),
+        placeholder_1: T.Buffer((64, 256, 16, 4), "int8"),
+        compute: T.Buffer((1024, 1024), "int32"),
+    ) -> None:
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        with T.sblock("root"):
+            T.reads()
+            T.writes()
+            for i0, i1, i2 in T.grid(1024, 1024, 1024):
+                with T.sblock("compute"):
+                    i, j, k = T.axis.remap("SSR", [i0, i1, i2])
+                    T.reads(placeholder[i, k], placeholder_1[j // 16, k // 4, j % 16, k % 4])
+                    T.writes(compute[i, j])
+                    with T.init():
+                        compute[i, j] = 0
+                    compute[i, j] = compute[i, j] + T.cast(placeholder[i, k], "int32") * T.cast(
+                        placeholder_1[j // 16, k // 4, j % 16, k % 4], "int32"
+                    )
+
+
+@tvm.script.ir_module
+class Conv2dNCHWcTIRModule:
+    @T.prim_func(s_tir=True)
+    def main(
+        placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"),
+        placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"),
+        conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32"),
+    ) -> None:
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i0, i1, i2, i3, i4, i5, i6, i7, i8, i9 in T.grid(1, 16, 56, 56, 16, 1, 1, 4, 4, 4):
+            with T.sblock("conv2d_NCHWc_int8"):
+                (
+                    n,
+                    oc_chunk,
+                    oh,
+                    ow,
+                    oc_block,
+                    kh,
+                    kw,
+                    ic_outer,
+                    ic_f_inner,
+                    ic_s_inner,
+                ) = T.axis.remap("SSSSSRRRRR", [i0, i1, i2, i3, i4, i5, i6, i7, i8, i9])
+                T.reads(
+                    placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
+                    placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
+                )
+                T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block])
+                with T.init():
+                    conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = 0
+                conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = conv2d_NCHWc_int8[
+                    n, oc_chunk, oh, ow, oc_block
+                ] + T.cast(
+                    placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
+                    "int32",
+                ) * T.cast(
+                    placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
+                    "int32",
+                )
+
+
+def collect_loops(prim_func):
+    loops = []
+
+    def callback(node):
+        if isinstance(node, tvm.tirx.For):
+            loops.append(node)
+        return True
+
+    pre_order_visit(prim_func.body, callback)
+
+    return loops
+
+
+def test_get_tensorize_loop_mapping_dense_16x4():
+    s = Schedule(DenseTIRModule)
+    block = s.get_sblock("compute")
+
+    info = get_tensorize_loop_mapping(s, block, dot_product_16x4_u8i8i32_desc)
+
+    assert isinstance(info, TensorizeInfo)
+
+    desc_loop_to_sref = dict((v, k) for k, v in info.loop_map.items())
+
+    desc_loops = collect_loops(dot_product_16x4_u8i8i32_desc)
+    _, loop_j, loop_k = s.get_loops(block)
+
+    assert desc_loops[0] in desc_loop_to_sref and desc_loops[1] in desc_loop_to_sref
+    assert s.get(desc_loop_to_sref[desc_loops[0]]) == s.get(loop_j)
+    assert s.get(desc_loop_to_sref[desc_loops[1]]) == s.get(loop_k)
+
+
+def test_get_tensorize_loop_mapping_conv2d_nchwc_16x4():
+    s = Schedule(Conv2dNCHWcTIRModule)
+    block = s.get_sblock("conv2d_NCHWc_int8")
+
+    info = get_tensorize_loop_mapping(s, block, dot_product_16x4_u8i8i32_desc)
+
+    desc_loop_to_sref = dict((v, k) for k, v in info.loop_map.items())
+
+    desc_loops = collect_loops(dot_product_16x4_u8i8i32_desc)
+
+    # i4 corresonds to the inner output channel axis of the NCHWc output tensor
+    # for i0, i1, i2, i3, i4, i5, i6, i7, i8, i9 in T.grid(1, 16, 56, 56, 16, 1, 1, 4, 4, 4):
+    _, _, _, _, i4, _, _, _, _, i9 = s.get_loops(block)
+
+    assert desc_loops[0] in desc_loop_to_sref and desc_loops[1] in desc_loop_to_sref
+    assert s.get(desc_loop_to_sref[desc_loops[0]]) == s.get(i4)
+    assert s.get(desc_loop_to_sref[desc_loops[1]]) == s.get(i9)
+
+
+def test_get_tensorize_loop_mapping_matmul_mma():
+    @T.prim_func(s_tir=True)
+    def matmul_16x16x16xf16f16f16_desc(
+        A: T.Buffer((16, 16), "float16", align=64, offset_factor=1),
+        B: T.Buffer((16, 16), "float16", align=64, offset_factor=1),
+        C: T.Buffer((16, 16), "float16", align=64, offset_factor=1),
+    ) -> None:
+        with T.sblock("root"):
+            T.reads(C[0:16, 0:16], A[0:16, 0:16], B[0:16, 0:16])
+            T.writes(C[0:16, 0:16])
+            for i, j, k in T.grid(16, 16, 16):
+                with T.sblock("update"):
+                    vii, vjj, vkk = T.axis.remap("SSR", [i, j, k])
+                    C[vii, vjj] = C[vii, vjj] + A[vii, vkk] * B[vjj, vkk]
+
+    matmul = create_prim_func(
+        te_workload.matmul_relu(
+            n=512,
+            m=512,
+            k=512,
+        )
+    )
+
+    s = Schedule(matmul)
+    block = s.get_sblock("C")
+    i0, i1, i2 = s.get_loops(block)
+    desc_loops = collect_loops(matmul_16x16x16xf16f16f16_desc)
+
+    for do_reorder in [False, True]:
+        # Mapping should be invariant to the loop permutation
+        if do_reorder:
+            s.reorder(i2, i0, i1)
+
+        info = get_tensorize_loop_mapping(s, block, matmul_16x16x16xf16f16f16_desc)
+        assert info is not None
+        desc_loop_to_sref = dict((v, k) for k, v in info.loop_map.items())
+
+        for i in range(3):
+            assert desc_loops[i] in desc_loop_to_sref
+
+        assert s.get(desc_loop_to_sref[desc_loops[0]]) == s.get(i0)
+        assert s.get(desc_loop_to_sref[desc_loops[1]]) == s.get(i1)
+        assert s.get(desc_loop_to_sref[desc_loops[2]]) == s.get(i2)
+
+
+def test_get_tensorize_loop_mapping_padding_matmul():
+    matmul = create_prim_func(
+        te_workload.matmul_relu(
+            n=127,
+            m=256,
+            k=65,
+            in_dtype="float16",
+            out_dtype="float16",
+        )
+    )
+    s = Schedule(matmul)
+    block = s.get_sblock("C")
+
+    desc = TensorIntrin.get(WMMA_SYNC_16x16x16_f16f16f16_INTRIN).desc
+    info = get_tensorize_loop_mapping(s, block, desc, allow_padding=True)
+    assert info is not None
+    expected_padding = [16, 1, 16]
+    actual_padding = info.block_iter_paddings
+    assert actual_padding is not None
+    assert len(actual_padding) == len(expected_padding)
+    for actual, expected in zip(actual_padding, expected_padding):
+        assert actual == expected
+
+
+def check_index_map(workload, block_name, intrin_name, expected_index_map):
+    s = Schedule(workload)
+    block = s.get_sblock(block_name)
+    desc_func = TensorIntrin.get(intrin_name).desc
+    info = get_auto_tensorize_mapping_info(s, block, desc_func)
+    if expected_index_map is None:
+        assert info is None
+        return
+    assert len(info.mappings) == 1
+    assert IndexMap.from_func(expected_index_map).is_equivalent_to(info.mappings[0])
+
+
+def test_get_auto_tensorize_mapping_info_conv2d():
+    conv2d = create_prim_func(
+        te_workload.conv2d_nhwc(4, 16, 16, 64, 64, 3, 1, 1, in_dtype="float16", out_dtype="float32")
+    )
+    check_index_map(
+        conv2d,
+        "conv2d_nhwc",
+        WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
+        lambda n, h, w, c, rh, rw, rc: (n * 256 + h * 16 + w, c, rh * 192 + rw * 64 + rc),
+    )
+
+
+def test_get_auto_tensorize_mapping_info_conv2d_unit_batch():
+    conv2d = create_prim_func(
+        te_workload.conv2d_nhwc(1, 16, 16, 64, 64, 3, 1, 1, in_dtype="float16", out_dtype="float32")
+    )
+    check_index_map(
+        conv2d,
+        "conv2d_nhwc",
+        WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
+        lambda n, h, w, c, rh, rw, rc: (n * 256 + h * 16 + w, c, rh * 192 + rw * 64 + rc),
+    )
+
+
+@pytest.mark.parametrize("b,m,n,k", [(1, 512, 512, 512), (16, 32, 32, 32)])
+def test_get_auto_tensorize_mapping_info_batch_matmul(b, m, n, k):
+    matmul = create_prim_func(
+        te_workload.batch_matmul_nkkm(b, m, n, k, in_dtype="float16", out_dtype="float32")
+    )
+    check_index_map(
+        matmul, "Z", WMMA_SYNC_16x16x16_f16f16f32_INTRIN, lambda b, m, n, k: (b, m, n, k)
+    )
+
+
+@pytest.mark.parametrize(
+    "n,m,k,expected",
+    [
+        (
+            512,
+            512,
+            512,
+            lambda n, m, k: (
+                n,
+                m,
+                k,
+            ),
+        ),
+        (1, 32, 32, lambda n, m, k: (n, m, k)),
+    ],
+)
+def test_get_auto_tensorize_mapping_info_matmul(n, m, k, expected):
+    matmul = create_prim_func(te_workload.matmul(n, m, k, in_dtype="float16", out_dtype="float32"))
+    check_index_map(matmul, "C", WMMA_SYNC_16x16x16_f16f16f32_INTRIN, expected)
+
+
+def test_is_output_block():
+    @T.prim_func(s_tir=True)
+    def two_elementwise(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (128, 128), "float32")
+        B = T.sblock_alloc_buffer((128, 128), "float32")
+        C = T.match_buffer(c, (128, 128), "float32")
+        for i, j in T.grid(128, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    sch = tvm.s_tir.Schedule(two_elementwise)
+    block_rv = sch.get_sblock("C")
+    assert is_output_block(sch, block_rv)
+
+
+def test_empty_grid():
+    @T.prim_func(s_tir=True)
+    def foo(out: T.Buffer((T.int64(1), T.int64(8), T.int64(8)), "int32")):
+        act = T.sblock_alloc_buffer((1, 8, 8), "int32")
+        for z2, y2, x2 in T.grid(1, 8, 8):
+            with T.sblock("b0"):
+                az, ay, ax = T.axis.remap("SSS", [z2, y2, x2])
+                T.writes(act[az, ay, ax])
+                act[az, ay, az] = T.int32(0)
+        # Empty grid:
+        for z1, y1, x1 in T.grid(0, 8, 8):
+            with T.sblock("b1"):
+                az, ay, ax = T.axis.remap("SSS", [z1, y1, x1])
+                T.reads(act[az + 1, ay, ax])
+                T.writes(out[az, ay, ax])
+                out[az, ay, ax] = act[az + 1, ay, ax]
+        # The block below is not needed to show the bug, but the 'out'
+        # buffer would be undefined without it.
+        for z2, y2, x2 in T.grid(1, 8, 8):
+            with T.sblock("b2"):
+                az, ay, ax = T.axis.remap("SSS", [z2, y2, x2])
+                T.writes(out[az, ay, ax])
+                out[az, ay, az] = T.int32(0)
+
+    # This caused a crash before.
+    sch = tvm.s_tir.Schedule(foo)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_annotate_buffer_access.py

@@ -0,0 +1,333 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E501, F401
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+
+def test_annotate_read_buffer_access():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
+        B = T.sblock_alloc_buffer((128, 128), "float32")
+        for i, j in T.grid(128, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
+        B = T.sblock_alloc_buffer((128, 128), "float32")
+        for i, j in T.grid(128, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(A[vi - 1 : vi - 1 + 2, vj - 1 : vj - 1 + 2])
+                T.writes(B[vi, vj])
+                T.sblock_attr({"explicit_read_region": [0]})
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    block = sch.get_sblock("B")
+    sch.annotate_buffer_access(
+        block, 0, "read", lambda vi, vj: ((vi - 1, vi + 1), (vj - 1, vj + 1))
+    )
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], expected)
+    verify_trace_roundtrip(sch=sch, mod=before)
+
+
+def test_annotate_write_buffer_access():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
+        B = T.sblock_alloc_buffer((128, 128), "float32")
+        for i, j in T.grid(128, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
+        B = T.sblock_alloc_buffer((128, 128), "float32")
+        for i, j in T.grid(128, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(A[vi, vj])
+                T.writes(B[vi : vi + 2, vj : vj + 2])
+                T.sblock_attr({"explicit_write_region": [0]})
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    block = sch.get_sblock("B")
+    sch.annotate_buffer_access(block, 0, "write", lambda vi, vj: ((vi, vi + 2), (vj, vj + 2)))
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], expected)
+    verify_trace_roundtrip(sch=sch, mod=before)
+
+
+def test_annotate_buffer_access_for_resize():
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def resize_before(x: T.Buffer((1, 1, 32, 32), "float16"), resize: T.Buffer((1, 1, 16, 16), "float16")):
+        for i0, i1, i2, i3 in T.grid(1, 1, 16, 16):
+            with T.sblock("resize"):
+                v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                T.reads(x[v_i0, v_i1, 0:32, 0:32])
+                T.writes(resize[v_i0, v_i1, v_i2, v_i3])
+                resize[v_i0, v_i1, v_i2, v_i3] = T.Cast("float16", T.Cast("float32", x[v_i0, v_i1, T.max(T.min(T.Cast("int32", T.floor((T.Cast("float32", v_i2) + T.float32(0.5)) * T.float32(2) - T.float32(0.5) + T.float32(1.0000000000000001e-05))), 31), 0), T.max(T.min(T.Cast("int32", T.floor((T.Cast("float32", v_i3) + T.float32(0.5)) * T.float32(2) - T.float32(0.5) + T.float32(1.0000000000000001e-05))), 31), 0)]))
+
+    @T.prim_func(s_tir=True)
+    def resize_expected(x: T.Buffer((1, 1, 32, 32), "float16"), resize: T.Buffer((1, 1, 16, 16), "float16")):
+        for i0, i1, i2, i3 in T.grid(1, 1, 16, 16):
+            with T.sblock("resize"):
+                v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                T.reads(x[v_i0, v_i1, v_i2 * 2 - 3:v_i2 * 2 + 3, v_i3 * 2 - 3:v_i3 * 2 + 3])
+                T.writes(resize[v_i0, v_i1, v_i2, v_i3])
+                T.sblock_attr({"explicit_read_region": [0]})
+                resize[v_i0, v_i1, v_i2, v_i3] = T.Cast("float16", T.Cast("float32", x[v_i0, v_i1, T.max(T.min(T.Cast("int32", T.floor((T.Cast("float32", v_i2) + T.float32(0.5)) * T.float32(2) - T.float32(0.5) + T.float32(1.0000000000000001e-05))), 31), 0), T.max(T.min(T.Cast("int32", T.floor((T.Cast("float32", v_i3) + T.float32(0.5)) * T.float32(2) - T.float32(0.5) + T.float32(1.0000000000000001e-05))), 31), 0)]))
+    # fmt: on
+    sch = tvm.s_tir.Schedule(resize_before, debug_mask="all")
+    block = sch.get_sblock("resize")
+    sch.annotate_buffer_access(
+        block,
+        0,
+        "read",
+        gen_new_ranges=lambda v_i0, v_i1, v_i2, v_i3: [
+            v_i0,
+            v_i1,
+            (v_i2 * 2 - 3, v_i2 * 2 + 3),
+            (v_i3 * 2 - 3, v_i3 * 2 + 3),
+        ],
+    )
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], resize_expected)
+    verify_trace_roundtrip(sch=sch, mod=resize_before)
+
+
+def test_annotate_buffer_access_read_and_write():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
+        B = T.sblock_alloc_buffer((128, 128), "float32")
+        for i, j in T.grid(128, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(A[vi, vj])
+                T.writes(B[vi, vj])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(B[vi, vj])
+                T.writes(C[vi, vj])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
+        B = T.sblock_alloc_buffer((128, 128), "float32")
+        for i, j in T.grid(128, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(A[vi - 1 : vi + 2, vj - 1 : vj + 2])
+                T.writes(B[vi : vi + 2, vj : vj + 2])
+                T.sblock_attr({"explicit_read_region": [0], "explicit_write_region": [0]})
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(B[vi, vj])
+                T.writes(C[vi, vj])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    block = sch.get_sblock("B")
+
+    sch.annotate_buffer_access(
+        block, 0, "read", lambda vi, vj: ((vi - 1, vi + 2), (vj - 1, vj + 2))
+    )
+
+    sch.annotate_buffer_access(block, 0, "write", lambda vi, vj: ((vi, vi + 2), (vj, vj + 2)))
+
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], expected)
+    verify_trace_roundtrip(sch=sch, mod=before)
+
+
+def test_double_annotate_buffer_access_read():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
+        B = T.sblock_alloc_buffer((128, 128), "float32")
+        for i, j in T.grid(128, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(A[vi, vj])
+                T.writes(B[vi, vj])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(B[vi, vj])
+                T.writes(C[vi, vj])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
+        B = T.sblock_alloc_buffer((128, 128), "float32")
+        for i, j in T.grid(128, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(A[vi - 2 : vi + 3, vj - 2 : vj + 3])
+                T.writes(B[vi, vj])
+                T.sblock_attr({"explicit_read_region": [0]})
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(B[vi, vj])
+                T.writes(C[vi, vj])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    block = sch.get_sblock("B")
+
+    sch.annotate_buffer_access(
+        block, 0, "read", lambda vi, vj: ((vi - 1, vi + 2), (vj - 1, vj + 2))
+    )
+
+    sch.annotate_buffer_access(
+        block, 0, "read", lambda vi, vj: ((vi - 2, vi + 3), (vj - 2, vj + 3))
+    )
+
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], expected)
+    verify_trace_roundtrip(sch=sch, mod=before)
+
+
+def test_annotate_buffer_access_with_compute_at_for_resize():
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def before(x: T.Buffer((1, 3, 200, 200), "float32"), y: T.Buffer((1, 3, 100, 100), "float32")):
+        x_global = T.sblock_alloc_buffer([1, 3, 200, 200], dtype="float32")
+        for ax0, ax1, ax2, ax3 in T.grid(1, 3, 200, 200):
+            with T.sblock("cache"):
+                v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
+                x_global[v0, v1, v2, v3] = x[v0, v1, v2, v3]
+        for i0, i1, i2, i3 in T.grid(1, 3, 100, 100):
+            with T.sblock("resize"):
+                v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                y[v_i0, v_i1, v_i2, v_i3] = x_global[v_i0, v_i1, T.Cast("int32", T.floor(v_i2 * 2 + 0.5)), T.Cast("int32", T.floor(v_i3 * 2 + 0.5))]
+
+    @T.prim_func(s_tir=True)
+    def after(x: T.Buffer((1, 3, 200, 200), "float32"), y: T.Buffer((1, 3, 100, 100), "float32")):
+        x_global = T.sblock_alloc_buffer((1, 3, 200, 200))
+        for i0, i1, i2_0, i3_0 in T.grid(1, 3, 10, 10):
+            for ax0, ax1 in T.grid(24, 24):
+                with T.sblock("cache"):
+                    v0 = T.axis.spatial(1, 0)
+                    v1 = T.axis.spatial(3, i1)
+                    v2 = T.axis.spatial(200, i2_0 * 20 - 3 + ax0)
+                    v3 = T.axis.spatial(200, i3_0 * 20 - 3 + ax1)
+                    T.where(3 <= i2_0 * 20 + ax0 and i2_0 * 20 + ax0 < 203 and 3 <= i3_0 * 20 + ax1 and i3_0 * 20 + ax1 < 203)
+                    T.reads(x[v0, v1, v2, v3])
+                    T.writes(x_global[v0, v1, v2, v3])
+                    x_global[v0, v1, v2, v3] = x[v0, v1, v2, v3]
+            for i2_1, i3_1 in T.grid(10, 10):
+                with T.sblock("resize"):
+                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+                    v_i2 = T.axis.spatial(100, i2_0 * 10 + i2_1)
+                    v_i3 = T.axis.spatial(100, i3_0 * 10 + i3_1)
+                    T.reads(x_global[v_i0, v_i1, v_i2 * 2 - 3:v_i2 * 2 - 3 + 6, v_i3 * 2 - 3:v_i3 * 2 - 3 + 6])
+                    T.writes(y[v_i0, v_i1, v_i2, v_i3])
+                    T.sblock_attr({"explicit_read_region": [0]})
+                    y[v_i0, v_i1, v_i2, v_i3] = x_global[v_i0, v_i1, T.Cast("int32", T.floor(T.Cast("float32", v_i2 * 2) + T.float32(0.5))), T.Cast("int32", T.floor(T.Cast("float32", v_i3 * 2) + T.float32(0.5)))]
+
+    @T.prim_func(s_tir=True)
+    def after_without_annotate_buffer_access(x: T.Buffer((1, 3, 200, 200), "float32"), y: T.Buffer((1, 3, 100, 100), "float32")):
+        x_global = T.sblock_alloc_buffer((1, 3, 200, 200))
+        for i0, i1, i2_0, i3_0 in T.grid(1, 3, 10, 10):
+            for ax0, ax1 in T.grid(200, 200):
+                with T.sblock("cache"):
+                    v0 = T.axis.spatial(1, 0)
+                    v1, v2, v3 = T.axis.remap("SSS", [i1, ax0, ax1])
+                    T.reads(x[v0, v1, v2, v3])
+                    T.writes(x_global[v0, v1, v2, v3])
+                    x_global[v0, v1, v2, v3] = x[v0, v1, v2, v3]
+            for i2_1, i3_1 in T.grid(10, 10):
+                with T.sblock("resize"):
+                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+                    v_i2 = T.axis.spatial(100, i2_0 * 10 + i2_1)
+                    v_i3 = T.axis.spatial(100, i3_0 * 10 + i3_1)
+                    T.reads(x_global[v_i0, v_i1, 0:200, 0:200])
+                    T.writes(y[v_i0, v_i1, v_i2, v_i3])
+                    y[v_i0, v_i1, v_i2, v_i3] = x_global[v_i0, v_i1, T.Cast("int32", T.floor(T.Cast("float32", v_i2 * 2) + T.float32(0.5))), T.Cast("int32", T.floor(T.Cast("float32", v_i3 * 2) + T.float32(0.5)))]
+    # fmt: on
+
+    # Schedule with annotate_buffer_access
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    block = sch.get_sblock("resize")
+    cache_block = sch.get_sblock("cache")
+
+    # Annotate buffer access
+    sch.annotate_buffer_access(
+        block,
+        0,
+        "read",
+        lambda vn, vc, vh, vw: (vn, vc, (vh * 2 - 3, vh * 2 + 3), (vw * 2 - 3, vw * 2 + 3)),
+    )
+
+    h, w = sch.get_loops(block)[-2:]
+    ho, hi = sch.split(h, factors=[10, 10])
+    wo, wi = sch.split(w, factors=[10, 10])
+    sch.reorder(ho, wo, hi, wi)
+    sch.compute_at(cache_block, wo)
+
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], after)
+    verify_trace_roundtrip(sch=sch, mod=before)
+
+    # Schedule without annotate_buffer_access
+    sch_without_annotate = tvm.s_tir.Schedule(before, debug_mask="all")
+    block_without_annotate = sch_without_annotate.get_sblock("resize")
+    cache_block_without_annotate = sch_without_annotate.get_sblock("cache")
+
+    h, w = sch_without_annotate.get_loops(block_without_annotate)[-2:]
+    ho, hi = sch_without_annotate.split(h, factors=[10, 10])
+    wo, wi = sch_without_annotate.split(w, factors=[10, 10])
+    sch_without_annotate.reorder(ho, wo, hi, wi)
+    sch_without_annotate.compute_at(cache_block_without_annotate, wo)
+
+    assert_structural_equal_ignore_global_symbol(
+        sch_without_annotate.mod["main"], after_without_annotate_buffer_access
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_block_scope.py

@@ -0,0 +1,157 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import sys
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import s_tir, tirx
+from tvm.s_tir.schedule import DepKind
+from tvm.script import tirx as T
+from tvm.tirx.stmt_functor import post_order_visit
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def elementwise(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128), "float32")
+    C = T.match_buffer(c, (128, 128), "float32")
+    B = T.sblock_alloc_buffer((128, 128), "float32")
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j in T.grid(128, 128):
+        with T.sblock("init"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = T.float32(0)
+        for k in range(0, 128):
+            with T.sblock("update"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def war_dependency(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+# pylint: disable=invalid-name
+
+
+def _get_sblock(s: s_tir.ScheduleState, name_hint: str) -> s_tir.StmtSRef:
+    result = None
+
+    def f_visit(node):
+        nonlocal result
+        if isinstance(node, tvm.tirx.SBlock) and node.name_hint == name_hint:
+            result = node
+
+    func = s.mod["main"]
+    post_order_visit(func.body, f_visit)
+    assert result is not None and isinstance(result, tvm.tirx.SBlock)
+    return s.get_sref(result)
+
+
+def test_elementwise_dependency():
+    s = tvm.s_tir.ScheduleState(elementwise, debug_mask="all")
+    root = _get_sblock(s, "root")
+    block_b = _get_sblock(s, "B")
+    block_c = _get_sblock(s, "C")
+    # Check get_deps_by_src
+    (dep,) = s.get_sblock_scope(root).get_deps_by_src(block_b)
+    assert dep.src.same_as(block_b)
+    assert dep.dst.same_as(block_c)
+    assert dep.kind == DepKind.RAW
+    # Check get_deps_by_dst
+    (dep,) = s.get_sblock_scope(root).get_deps_by_dst(block_c)
+    assert dep.src.same_as(block_b)
+    assert dep.dst.same_as(block_c)
+    assert dep.kind == DepKind.RAW
+
+
+def test_matmul_dependency():
+    s = tvm.s_tir.ScheduleState(matmul, debug_mask="all")
+    root = _get_sblock(s, "root")
+    init = _get_sblock(s, "init")
+    update = _get_sblock(s, "update")
+    # Check get_deps_by_src
+    p0, p1 = s.get_sblock_scope(root).get_deps_by_src(init)
+    assert p0.src.same_as(init)
+    assert p0.dst.same_as(update)
+    assert p1.src.same_as(init)
+    assert p1.dst.same_as(update)
+    assert (p0.kind == DepKind.RAW and p1.kind == DepKind.WAW) or (
+        p0.kind == DepKind.WAW and p1.kind == DepKind.RAW
+    )
+    # Check get_deps_by_dst
+    p0, p1 = s.get_sblock_scope(root).get_deps_by_dst(update)
+    assert p0.src.same_as(init)
+    assert p0.dst.same_as(update)
+    assert p1.src.same_as(init)
+    assert p1.dst.same_as(update)
+    assert (p0.kind == DepKind.RAW and p1.kind == DepKind.WAW) or (
+        p0.kind == DepKind.WAW and p1.kind == DepKind.RAW
+    )
+
+
+def test_war_dependency():
+    s = tvm.s_tir.ScheduleState(war_dependency, debug_mask="all")
+    root = _get_sblock(s, "root")
+    block_c = _get_sblock(s, "C")
+    block_b = _get_sblock(s, "B")
+    # Check get_deps_by_src
+    (dep,) = s.get_sblock_scope(root).get_deps_by_src(block_c)
+    assert dep.src.same_as(block_c)
+    assert dep.dst.same_as(block_b)
+    assert dep.kind == DepKind.WAR
+    # Check get_deps_by_dst
+    (dep,) = s.get_sblock_scope(root).get_deps_by_dst(block_b)
+    assert dep.src.same_as(block_c)
+    assert dep.dst.same_as(block_b)
+    assert dep.kind == DepKind.WAR
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_blockize.py

@@ -0,0 +1,375 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401, F841
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import verify_trace_roundtrip
+from tvm.script import tirx as T
+
+# fmt: off
+# pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks
+
+@T.prim_func(s_tir=True)
+def single_elementwise(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")):
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+
+# fmt: on
+# pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks
+
+
+def test_blockize_outer():
+    @T.prim_func(s_tir=True)
+    def after_blockize_outer(
+        A: T.Buffer((128, 128), "float32"),
+        B: T.Buffer((128, 128), "float32"),
+    ) -> None:
+        with T.sblock("blockized_B"):
+            vio = T.axis.spatial(1, 0)
+            vjo = T.axis.spatial(1, 0)
+            for i, j in T.grid(128, 128):
+                with T.sblock("B"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    B[vi, vj] = A[vi, vj] * 2.0
+
+    func = single_elementwise
+    s = tvm.s_tir.Schedule(func, debug_mask="all")
+    x, _ = s.get_loops(s.get_sblock("B"))
+    s.blockize(x)
+    tvm.ir.assert_structural_equal(
+        s.mod["main"], after_blockize_outer.with_attr("global_symbol", "single_elementwise")
+    )
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_blockize_inner():
+    @T.prim_func(s_tir=True)
+    def after_blockize_inner(
+        A: T.Buffer((128, 128), "float32"),
+        B: T.Buffer((128, 128), "float32"),
+    ) -> None:
+        for i in T.serial(128):
+            with T.sblock("blockized_B"):
+                vi = T.axis.spatial(128, i)
+                vjo = T.axis.spatial(1, 0)
+                for j in T.serial(128):
+                    with T.sblock("B"):
+                        vj = T.axis.remap("S", [j])
+                        B[vi, vj] = A[vi, vj] * 2.0
+
+    func = single_elementwise
+    s = tvm.s_tir.Schedule(func, debug_mask="all")
+    _, y = s.get_loops(s.get_sblock("B"))
+    s.blockize(y)
+    tvm.ir.assert_structural_equal(
+        s.mod["main"], after_blockize_inner.with_attr("global_symbol", "single_elementwise")
+    )
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_two_elementwise_blockize_reverse_compute_at():
+    @T.prim_func(s_tir=True)
+    def before_blockize_rca(
+        A: T.Buffer((128, 128), "float32"),
+        C: T.Buffer((128, 128), "float32"),
+    ) -> None:
+        B = T.sblock_alloc_buffer([128, 128], dtype="float32")
+        for i, j in T.grid(8, 8):
+            with T.sblock("B_o"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+                T.writes(B[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+                for i_1, j_1 in T.grid(16, 16):
+                    with T.sblock("B"):
+                        vi_i, vj_i = T.axis.remap("SS", [i_1, j_1])
+                        T.reads(A[vi * 16 + vi_i, vj * 16 + vj_i])
+                        T.writes(B[vi * 16 + vi_i, vj * 16 + vj_i])
+                        B[vi * 16 + vi_i, vj * 16 + vj_i] = A[vi * 16 + vi_i, vj * 16 + vj_i] * 2.0
+            for ax0, ax1 in T.grid(16, 16):
+                with T.sblock("C"):
+                    vi = T.axis.spatial(128, i * 16 + ax0)
+                    vj = T.axis.spatial(128, j * 16 + ax1)
+                    T.reads(B[vi, vj])
+                    T.writes(C[vi, vj])
+                    C[vi, vj] = B[vi, vj] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def after_blockize_rca(
+        A: T.Buffer((128, 128), "float32"),
+        C: T.Buffer((128, 128), "float32"),
+    ) -> None:
+        B = T.sblock_alloc_buffer([128, 128], dtype="float32")
+        for i, j in T.grid(8, 8):
+            with T.sblock("B_o"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+                T.writes(B[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+                for i_1, j_1 in T.grid(16, 16):
+                    with T.sblock("B"):
+                        vi_i, vj_i = T.axis.remap("SS", [i_1, j_1])
+                        T.reads(A[vi * 16 + vi_i, vj * 16 + vj_i])
+                        T.writes(B[vi * 16 + vi_i, vj * 16 + vj_i])
+                        B[vi * 16 + vi_i, vj * 16 + vj_i] = A[vi * 16 + vi_i, vj * 16 + vj_i] * 2.0
+            with T.sblock("C_o"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(B[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+                T.writes(C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+                for ax0, ax1 in T.grid(16, 16):
+                    with T.sblock("C"):
+                        vi_i, vj_i = T.axis.remap("SS", [ax0, ax1])
+                        T.reads(B[vi * 16 + vi_i, vj * 16 + vj_i])
+                        T.writes(C[vi * 16 + vi_i, vj * 16 + vj_i])
+                        C[vi * 16 + vi_i, vj * 16 + vj_i] = B[vi * 16 + vi_i, vj * 16 + vj_i] + 1.0
+
+    func = before_blockize_rca
+    s = tvm.s_tir.Schedule(func, debug_mask="all")
+    _, _, x, _ = s.get_loops(s.get_sblock("C"))
+    s.blockize(x)
+    tvm.ir.assert_structural_equal(
+        s.mod["main"], after_blockize_rca.with_attr("global_symbol", "before_blockize_rca")
+    )
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_two_elementwise_blockize_compute_at():
+    @T.prim_func(s_tir=True)
+    def before_blockize_compute_at(
+        A: T.Buffer((128, 128), "float32"),
+        C: T.Buffer((128, 128), "float32"),
+    ) -> None:
+        # body
+        # with T.sblock("root")
+        B = T.sblock_alloc_buffer([128, 128], dtype="float32")
+        for i_0, j_0 in T.grid(8, 8):
+            for ax0, ax1 in T.grid(16, 16):
+                with T.sblock("B"):
+                    vi = T.axis.spatial(128, i_0 * 16 + ax0)
+                    vj = T.axis.spatial(128, j_0 * 16 + ax1)
+                    T.reads(A[vi, vj])
+                    T.writes(B[vi, vj])
+                    B[vi, vj] = A[vi, vj] * 2.0
+            with T.sblock("C_o"):
+                vi_o, vj_o = T.axis.remap("SS", [i_0, j_0])
+                T.reads(B[vi_o * 16 : vi_o * 16 + 16, vj_o * 16 : vj_o * 16 + 16])
+                T.writes(C[vi_o * 16 : vi_o * 16 + 16, vj_o * 16 : vj_o * 16 + 16])
+                for i_1, j_1 in T.grid(16, 16):
+                    with T.sblock("C"):
+                        vi_i, vj_i = T.axis.remap("SS", [i_1, j_1])
+                        T.reads(B[vi_o * 16 + vi_i, vj_o * 16 + vj_i])
+                        T.writes(C[vi_o * 16 + vi_i, vj_o * 16 + vj_i])
+                        C[vi_o * 16 + vi_i, vj_o * 16 + vj_i] = (
+                            B[vi_o * 16 + vi_i, vj_o * 16 + vj_i] + 1.0
+                        )
+
+    @T.prim_func(s_tir=True)
+    def after_blockize_compute_at(
+        A: T.Buffer((128, 128), "float32"),
+        C: T.Buffer((128, 128), "float32"),
+    ) -> None:
+        B = T.sblock_alloc_buffer([128, 128], dtype="float32")
+        for i_0, j_0 in T.grid(8, 8):
+            with T.sblock("B_o"):
+                vi_o, vj_o = T.axis.remap("SS", [i_0, j_0])
+                T.reads(A[vi_o * 16 : vi_o * 16 + 16, vj_o * 16 : vj_o * 16 + 16])
+                T.writes(B[vi_o * 16 : vi_o * 16 + 16, vj_o * 16 : vj_o * 16 + 16])
+                for ax0, ax1 in T.grid(16, 16):
+                    with T.sblock("B"):
+                        vi_i, vj_i = T.axis.remap("SS", [ax0, ax1])
+                        T.reads(A[vi_o * 16 + vi_i, vj_o * 16 + vj_i])
+                        T.writes(B[vi_o * 16 + vi_i, vj_o * 16 + vj_i])
+                        B[vi_o * 16 + vi_i, vj_o * 16 + vj_i] = (
+                            A[vi_o * 16 + vi_i, vj_o * 16 + vj_i] * 2.0
+                        )
+            with T.sblock("C_o"):
+                vi_o, vj_o = T.axis.remap("SS", [i_0, j_0])
+                T.reads(B[vi_o * 16 : vi_o * 16 + 16, vj_o * 16 : vj_o * 16 + 16])
+                T.writes(C[vi_o * 16 : vi_o * 16 + 16, vj_o * 16 : vj_o * 16 + 16])
+                for i_1, j_1 in T.grid(16, 16):
+                    with T.sblock("C"):
+                        vi_i, vj_i = T.axis.remap("SS", [i_1, j_1])
+                        T.reads(B[vi_o * 16 + vi_i, vj_o * 16 + vj_i])
+                        T.writes(C[vi_o * 16 + vi_i, vj_o * 16 + vj_i])
+                        C[vi_o * 16 + vi_i, vj_o * 16 + vj_i] = (
+                            B[vi_o * 16 + vi_i, vj_o * 16 + vj_i] + 1.0
+                        )
+
+    func = before_blockize_compute_at
+    s = tvm.s_tir.Schedule(func, debug_mask="all")
+    _, _, x, _ = s.get_loops(s.get_sblock("B"))
+    s.blockize(x)
+    tvm.ir.assert_structural_equal(
+        s.mod["main"],
+        after_blockize_compute_at.with_attr("global_symbol", "before_blockize_compute_at"),
+    )
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_blockize_init_loops():
+    @T.prim_func(s_tir=True)
+    def rowsum(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128,), "float32")) -> None:
+        for k, i in T.grid(128, 128):
+            with T.sblock("B"):
+                vk, vi = T.axis.remap("RS", [k, i])
+                with T.init():
+                    B[vi] = 0.0
+                B[vi] = B[vi] + A[vi, vk]
+
+    @T.prim_func(s_tir=True)
+    def after_rowsum_blockize(
+        A: T.Buffer((128, 128), "float32"),
+        B: T.Buffer((128,), "float32"),
+    ) -> None:
+        with T.sblock("blockized_B"):
+            vko = T.axis.R(1, 0)
+            vio = T.axis.S(1, 0)
+            with T.init():
+                for i1 in T.serial(0, 128):
+                    with T.sblock("B_init"):
+                        vi_init = T.axis.S(128, i1)
+                        B[vi_init] = T.float32(0)
+            for i0, i1_1 in T.grid(128, 128):
+                with T.sblock("B"):
+                    vk, vi = T.axis.remap("RS", [i0, i1_1])
+                    B[vi] = B[vi] + A[vi, vk]
+
+    s = tvm.s_tir.Schedule(rowsum, debug_mask="all")
+    k, _ = s.get_loops(s.get_sblock("B"))
+    s.blockize(k)
+    tvm.ir.assert_structural_equal(
+        s.mod["main"], after_rowsum_blockize.with_attr("global_symbol", "rowsum")
+    )
+    verify_trace_roundtrip(sch=s, mod=rowsum)
+
+
+@pytest.mark.parametrize("preserve_unit_iters", [True, False])
+def test_blockize_outer_int64_shape(preserve_unit_iters):
+    @T.prim_func(s_tir=True)
+    def single_elementwise_int64(
+        A: T.Buffer((T.int64(16), T.int64(128)), "float32"),
+        B: T.Buffer((T.int64(16), T.int64(128)), "float32"),
+    ) -> None:
+        for i0, j0, i1, j1 in T.grid(T.int64(1), T.int64(8), T.int64(16), T.int64(16)):
+            with T.sblock("B"):
+                vi = T.axis.S(T.int64(16), i0 * T.int64(16) + i1)
+                vj = T.axis.S(T.int64(128), j0 * T.int64(16) + j1)
+                B[vi, vj] = A[vi, vj] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def after_single_elementwise_int64_blockize(
+        A: T.Buffer((T.int64(16), T.int64(128)), "float32"),
+        B: T.Buffer((T.int64(16), T.int64(128)), "float32"),
+    ) -> None:
+        for i0, j0 in T.grid(T.int64(1), T.int64(8)):
+            with T.sblock("B_o"):
+                vi_o = T.axis.spatial(T.int64(1), T.int64(0))
+                vj_o = T.axis.spatial(T.int64(8), j0)
+                for i1, j1 in T.grid(T.int64(16), T.int64(16)):
+                    with T.sblock("B"):
+                        vi_i, vj_i = T.axis.remap("SS", [i1, j1])
+                        B[vi_i, vj_o * T.int64(16) + vj_i] = A[
+                            vi_i, vj_o * T.int64(16) + vj_i
+                        ] + T.float32(1)
+
+    @T.prim_func(s_tir=True)
+    def after_single_elementwise_int64_blockize_preserve_unit_iters(
+        A: T.Buffer((T.int64(16), T.int64(128)), "float32"),
+        B: T.Buffer((T.int64(16), T.int64(128)), "float32"),
+    ) -> None:
+        for i0, j0 in T.grid(T.int64(1), T.int64(8)):
+            with T.sblock("B_o"):
+                vi_o = T.axis.spatial(T.int64(1), i0)
+                vj_o = T.axis.spatial(T.int64(8), j0)
+                for i1, j1 in T.grid(T.int64(16), T.int64(16)):
+                    with T.sblock("B"):
+                        vi_i, vj_i = T.axis.remap("SS", [i1, j1])
+                        B[vi_i, vj_o * T.int64(16) + vj_i] = A[
+                            vi_i, vj_o * T.int64(16) + vj_i
+                        ] + T.float32(1)
+
+    s = tvm.s_tir.Schedule(single_elementwise_int64, debug_mask="all")
+    _, _, i1, _ = s.get_loops(s.get_sblock("B"))
+    s.blockize(i1, preserve_unit_iters=preserve_unit_iters)
+    expected = (
+        after_single_elementwise_int64_blockize_preserve_unit_iters
+        if preserve_unit_iters
+        else after_single_elementwise_int64_blockize
+    )
+    tvm.ir.assert_structural_equal(
+        s.mod["main"], expected.with_attr("global_symbol", "single_elementwise_int64")
+    )
+    verify_trace_roundtrip(sch=s, mod=single_elementwise_int64)
+
+
+def test_blockize_blocks():
+    @T.prim_func(s_tir=True)
+    def blocks_func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")) -> None:
+        for m in T.serial(6):
+            for i, j in T.grid(3, 1):
+                with T.sblock("B"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    T.reads(A[vi, vj])
+                    T.writes(B[vi, vj])
+                    B[vi, vj] = A[vi, vj] * 2.0
+
+            for i, j in T.grid(128, 64):
+                with T.sblock("C"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    T.reads(A[vi, vj + 64])
+                    T.writes(B[vi, vj + 64])
+                    B[vi, vj + 64] = A[vi, vj + 64] * 3.0
+
+    @T.prim_func(s_tir=True)
+    def after_blocks_blockize(
+        A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")
+    ) -> None:
+        for m in range(6):
+            with T.sblock("outer_B_C_"):
+                vi_o = T.axis.spatial(1, 0)
+                vj_o = T.axis.spatial(1, 0)
+                T.reads(A[0:128, 0:128])
+                T.writes(B[0:128, 0:128])
+                for i, j in T.grid(3, 1):
+                    with T.sblock("B"):
+                        vi_i = T.axis.spatial(3, i)
+                        T.reads(A[vi_i, 0])
+                        T.writes(B[vi_i, 0])
+                        B[vi_i, 0] = A[vi_i, 0] * T.float32(2)
+                for i, j in T.grid(128, 64):
+                    with T.sblock("C"):
+                        vi_i, vj_i = T.axis.remap("SS", [i, j])
+                        T.reads(A[vi_i, vj_i + 64])
+                        T.writes(B[vi_i, vj_i + 64])
+                        B[vi_i, vj_i + 64] = A[vi_i, vj_i + 64] * T.float32(3)
+
+    s = tvm.s_tir.Schedule(blocks_func, debug_mask="all")
+    blocks = [s.get_sblock("B"), s.get_sblock("C")]
+    s.blockize(blocks, preserve_unit_iters=False)
+    expected = after_blocks_blockize
+    tvm.ir.assert_structural_equal(
+        s.mod["main"], expected.with_attr("global_symbol", "blocks_func")
+    )
+    verify_trace_roundtrip(sch=s, mod=blocks_func)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_cache_index.py

@@ -0,0 +1,478 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import sys
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import verify_trace_roundtrip
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+########## Function before schedule ##########
+
+
+@T.prim_func(s_tir=True)
+def resize(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (1, 3, 40, 40))
+    B = T.match_buffer(b, (1, 3, 80, 80))
+    for i0, i1, i2, i3 in T.grid(1, 3, 80, 80):
+        with T.sblock("A"):
+            n, c, vi, vj = T.axis.remap("SSSS", [i0, i1, i2, i3])
+            B[n, c, vi, vj] = A[n, c, vi // 4 + vj // 4, vj // 2]
+
+
+@T.prim_func(s_tir=True)
+def resize_cache_index(
+    A: T.Buffer((1, 3, 40, 40), "float32"), B: T.Buffer((1, 3, 80, 80), "float32")
+) -> None:
+    index_var_0 = T.sblock_alloc_buffer([80, 80], dtype="int32", strides=[1])
+    index_var_1 = T.sblock_alloc_buffer([80], dtype="int32", strides=[1])
+    for ax0, ax1 in T.grid(80, 80):
+        with T.sblock("index_0"):
+            v0, v1 = T.axis.remap("SS", [ax0, ax1])
+            T.reads()
+            T.writes(index_var_0[v0, v1])
+            index_var_0[v0, v1] = v0 // 4 + v1 // 4
+    for ax0 in T.serial(80):
+        with T.sblock("index_1"):
+            v0 = T.axis.spatial(80, ax0)
+            T.reads()
+            T.writes(index_var_1[v0])
+            index_var_1[v0] = v0 // 2
+    for i0, i1, i2, i3 in T.grid(1, 3, 80, 80):
+        with T.sblock("A"):
+            n, c, vi, vj = T.axis.remap("SSSS", [i0, i1, i2, i3])
+            T.reads(A[n, c, vi // 4 + vj // 4, vj // 2])
+            T.writes(B[n, c, vi, vj])
+            B[n, c, vi, vj] = A[n, c, index_var_0[vi, vj], index_var_1[vj]]
+
+
+@T.prim_func(s_tir=True)
+def bilinear_resize(
+    x: T.Buffer((1, 3, 40, 40), "float16"), resize: T.Buffer((1, 3, 80, 80), "float16")
+):
+    for i0, i1, i2, i3 in T.grid(1, 3, 80, 80):
+        with T.sblock("resize"):
+            i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+            T.reads(x[i0_1, i1_1, 0:40, 0:40])
+            T.writes(resize[i0_1, i1_1, i2_1, i3_1])
+            resize[i0_1, i1_1, i2_1, i3_1] = T.Cast(
+                "float16",
+                (
+                    T.Cast(
+                        "float32",
+                        x[
+                            i0_1,
+                            i1_1,
+                            T.max(
+                                T.min(
+                                    T.Cast(
+                                        "int32",
+                                        T.floor(
+                                            (T.Cast("float32", i2_1) + T.float32(0.5))
+                                            * T.float32(0.5)
+                                            - T.float32(0.5),
+                                            dtype="float32",
+                                        ),
+                                    ),
+                                    39,
+                                ),
+                                0,
+                            ),
+                            T.max(
+                                T.min(
+                                    T.Cast(
+                                        "int32",
+                                        T.floor(
+                                            (T.Cast("float32", i3_1) + T.float32(0.5))
+                                            * T.float32(0.5)
+                                            - T.float32(0.5),
+                                            dtype="float32",
+                                        ),
+                                    ),
+                                    39,
+                                ),
+                                0,
+                            ),
+                        ],
+                    )
+                    * (
+                        T.float32(1)
+                        - (
+                            (T.Cast("float32", i3_1) + T.float32(0.5)) * T.float32(0.5)
+                            - T.float32(0.5)
+                            - T.Cast(
+                                "float32",
+                                T.Cast(
+                                    "int32",
+                                    T.floor(
+                                        (T.Cast("float32", i3_1) + T.float32(0.5)) * T.float32(0.5)
+                                        - T.float32(0.5),
+                                        dtype="float32",
+                                    ),
+                                ),
+                            )
+                        )
+                    )
+                    + T.Cast(
+                        "float32",
+                        x[
+                            i0_1,
+                            i1_1,
+                            T.max(
+                                T.min(
+                                    T.Cast(
+                                        "int32",
+                                        T.floor(
+                                            (T.Cast("float32", i2_1) + T.float32(0.5))
+                                            * T.float32(0.5)
+                                            - T.float32(0.5),
+                                            dtype="float32",
+                                        ),
+                                    ),
+                                    39,
+                                ),
+                                0,
+                            ),
+                            T.max(
+                                T.min(
+                                    T.Cast(
+                                        "int32",
+                                        T.floor(
+                                            (T.Cast("float32", i3_1) + T.float32(0.5))
+                                            * T.float32(0.5)
+                                            - T.float32(0.5),
+                                            dtype="float32",
+                                        ),
+                                    )
+                                    + 1,
+                                    39,
+                                ),
+                                0,
+                            ),
+                        ],
+                    )
+                    * (
+                        (T.Cast("float32", i3_1) + T.float32(0.5)) * T.float32(0.5)
+                        - T.float32(0.5)
+                        - T.Cast(
+                            "float32",
+                            T.Cast(
+                                "int32",
+                                T.floor(
+                                    (T.Cast("float32", i3_1) + T.float32(0.5)) * T.float32(0.5)
+                                    - T.float32(0.5),
+                                    dtype="float32",
+                                ),
+                            ),
+                        )
+                    )
+                )
+                * (
+                    T.float32(1)
+                    - (
+                        (T.Cast("float32", i2_1) + T.float32(0.5)) * T.float32(0.5)
+                        - T.float32(0.5)
+                        - T.Cast(
+                            "float32",
+                            T.Cast(
+                                "int32",
+                                T.floor(
+                                    (T.Cast("float32", i2_1) + T.float32(0.5)) * T.float32(0.5)
+                                    - T.float32(0.5),
+                                    dtype="float32",
+                                ),
+                            ),
+                        )
+                    )
+                )
+                + (
+                    T.Cast(
+                        "float32",
+                        x[
+                            i0_1,
+                            i1_1,
+                            T.max(
+                                T.min(
+                                    T.Cast(
+                                        "int32",
+                                        T.floor(
+                                            (T.Cast("float32", i2_1) + T.float32(0.5))
+                                            * T.float32(0.5)
+                                            - T.float32(0.5),
+                                            dtype="float32",
+                                        ),
+                                    )
+                                    + 1,
+                                    39,
+                                ),
+                                0,
+                            ),
+                            T.max(
+                                T.min(
+                                    T.Cast(
+                                        "int32",
+                                        T.floor(
+                                            (T.Cast("float32", i3_1) + T.float32(0.5))
+                                            * T.float32(0.5)
+                                            - T.float32(0.5),
+                                            dtype="float32",
+                                        ),
+                                    ),
+                                    39,
+                                ),
+                                0,
+                            ),
+                        ],
+                    )
+                    * (
+                        T.float32(1)
+                        - (
+                            (T.Cast("float32", i3_1) + T.float32(0.5)) * T.float32(0.5)
+                            - T.float32(0.5)
+                            - T.Cast(
+                                "float32",
+                                T.Cast(
+                                    "int32",
+                                    T.floor(
+                                        (T.Cast("float32", i3_1) + T.float32(0.5)) * T.float32(0.5)
+                                        - T.float32(0.5),
+                                        dtype="float32",
+                                    ),
+                                ),
+                            )
+                        )
+                    )
+                    + T.Cast(
+                        "float32",
+                        x[
+                            i0_1,
+                            i1_1,
+                            T.max(
+                                T.min(
+                                    T.Cast(
+                                        "int32",
+                                        T.floor(
+                                            (T.Cast("float32", i2_1) + T.float32(0.5))
+                                            * T.float32(0.5)
+                                            - T.float32(0.5),
+                                            dtype="float32",
+                                        ),
+                                    )
+                                    + 1,
+                                    39,
+                                ),
+                                0,
+                            ),
+                            T.max(
+                                T.min(
+                                    T.Cast(
+                                        "int32",
+                                        T.floor(
+                                            (T.Cast("float32", i3_1) + T.float32(0.5))
+                                            * T.float32(0.5)
+                                            - T.float32(0.5),
+                                            dtype="float32",
+                                        ),
+                                    )
+                                    + 1,
+                                    39,
+                                ),
+                                0,
+                            ),
+                        ],
+                    )
+                    * (
+                        (T.Cast("float32", i3_1) + T.float32(0.5)) * T.float32(0.5)
+                        - T.float32(0.5)
+                        - T.Cast(
+                            "float32",
+                            T.Cast(
+                                "int32",
+                                T.floor(
+                                    (T.Cast("float32", i3_1) + T.float32(0.5)) * T.float32(0.5)
+                                    - T.float32(0.5),
+                                    dtype="float32",
+                                ),
+                            ),
+                        )
+                    )
+                )
+                * (
+                    (T.Cast("float32", i2_1) + T.float32(0.5)) * T.float32(0.5)
+                    - T.float32(0.5)
+                    - T.Cast(
+                        "float32",
+                        T.Cast(
+                            "int32",
+                            T.floor(
+                                (T.Cast("float32", i2_1) + T.float32(0.5)) * T.float32(0.5)
+                                - T.float32(0.5),
+                                dtype="float32",
+                            ),
+                        ),
+                    )
+                ),
+            )
+
+
+@T.prim_func(s_tir=True)
+def cached_bilinear_resize(
+    x: T.Buffer((1, 3, 40, 40), "float16"), resize: T.Buffer((1, 3, 80, 80), "float16")
+):
+    index_var_0 = T.sblock_alloc_buffer([80], dtype="float32", strides=[1])
+    index_var_1 = T.sblock_alloc_buffer([80], dtype="int32", strides=[1])
+    index_var_2 = T.sblock_alloc_buffer([80], dtype="int32", strides=[1])
+    for ax0 in T.serial(80):
+        with T.sblock("index_0"):
+            v0 = T.axis.spatial(80, ax0)
+            T.reads()
+            T.writes(index_var_0[v0])
+            index_var_0[v0] = (
+                (T.Cast("float32", v0) + T.float32(0.5)) * T.float32(0.5)
+                - T.float32(0.5)
+                - T.Cast(
+                    "float32",
+                    T.Cast(
+                        "int32",
+                        T.floor(
+                            (T.Cast("float32", v0) + T.float32(0.5)) * T.float32(0.5)
+                            - T.float32(0.5),
+                            dtype="float32",
+                        ),
+                    ),
+                )
+            )
+    for ax0 in T.serial(80):
+        with T.sblock("index_1"):
+            v0 = T.axis.spatial(80, ax0)
+            T.reads()
+            T.writes(index_var_1[v0])
+            index_var_1[v0] = T.Cast(
+                "int32",
+                T.floor(
+                    (T.Cast("float32", v0) + T.float32(0.5)) * T.float32(0.5) - T.float32(0.5),
+                    dtype="float32",
+                ),
+            )
+    for ax0 in T.serial(80):
+        with T.sblock("index_2"):
+            v0 = T.axis.spatial(80, ax0)
+            T.reads()
+            T.writes(index_var_2[v0])
+            index_var_2[v0] = T.Cast(
+                "int32",
+                T.floor(
+                    (T.Cast("float32", v0) + T.float32(0.5)) * T.float32(0.5) - T.float32(0.5),
+                    dtype="float32",
+                ),
+            )
+    for i0, i1, i2, i3 in T.grid(1, 3, 80, 80):
+        with T.sblock("resize"):
+            i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+            T.reads(x[i0_1, i1_1, 0:40, 0:40])
+            T.writes(resize[i0_1, i1_1, i2_1, i3_1])
+            resize[i0_1, i1_1, i2_1, i3_1] = T.Cast(
+                "float16",
+                (
+                    T.Cast(
+                        "float32",
+                        x[
+                            i0_1,
+                            i1_1,
+                            T.max(T.min(index_var_1[i2_1], 39), 0),
+                            T.max(T.min(index_var_2[i3_1], 39), 0),
+                        ],
+                    )
+                    * (T.float32(1) - index_var_0[i3_1])
+                    + T.Cast(
+                        "float32",
+                        x[
+                            i0_1,
+                            i1_1,
+                            T.max(T.min(index_var_1[i2_1], 39), 0),
+                            T.max(T.min(index_var_2[i3_1] + 1, 39), 0),
+                        ],
+                    )
+                    * index_var_0[i3_1]
+                )
+                * (
+                    T.float32(1)
+                    - (
+                        (T.Cast("float32", i2_1) + T.float32(0.5)) * T.float32(0.5)
+                        - T.float32(0.5)
+                        - T.Cast("float32", index_var_1[i2_1])
+                    )
+                )
+                + (
+                    T.Cast(
+                        "float32",
+                        x[
+                            i0_1,
+                            i1_1,
+                            T.max(T.min(index_var_1[i2_1] + 1, 39), 0),
+                            T.max(T.min(index_var_2[i3_1], 39), 0),
+                        ],
+                    )
+                    * (T.float32(1) - index_var_0[i3_1])
+                    + T.Cast(
+                        "float32",
+                        x[
+                            i0_1,
+                            i1_1,
+                            T.max(T.min(index_var_1[i2_1] + 1, 39), 0),
+                            T.max(T.min(index_var_2[i3_1] + 1, 39), 0),
+                        ],
+                    )
+                    * index_var_0[i3_1]
+                )
+                * (
+                    (T.Cast("float32", i2_1) + T.float32(0.5)) * T.float32(0.5)
+                    - T.float32(0.5)
+                    - T.Cast("float32", index_var_1[i2_1])
+                ),
+            )
+
+
+def test_basic_cache_index():
+    sch = tvm.s_tir.Schedule(resize, debug_mask="all")
+    block = sch.get_sblock("A")
+    sch.cache_index(block, "global")
+    tvm.ir.assert_structural_equal(
+        resize_cache_index, sch.mod["main"].with_attr("global_symbol", "resize_cache_index")
+    )
+    verify_trace_roundtrip(sch=sch, mod=resize)
+
+
+def test_resize_bilinear_cache_index():
+    sch = tvm.s_tir.Schedule(bilinear_resize, debug_mask="all")
+    block = sch.get_sblock("resize")
+    sch.cache_index(block, "global", 4)
+    tvm.ir.assert_structural_equal(
+        sch.mod["main"], cached_bilinear_resize.with_attr("global_symbol", "bilinear_resize")
+    )
+    verify_trace_roundtrip(sch=sch, mod=bilinear_resize)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py

@@ -0,0 +1,424 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+
+def check_decompose_padding(origin, scheduled, expected, check_run=False):
+    assert_structural_equal_ignore_global_symbol(scheduled, expected)
+    if check_run:
+        in_buffer = origin.buffer_map[origin.params[0]]
+        out_buffer = origin.buffer_map[origin.params[1]]
+        in_shape = [int(_) for _ in in_buffer.shape]
+        out_shape = [int(_) for _ in out_buffer.shape]
+        x = tvm.runtime.tensor(np.random.uniform(0, 64, in_shape).astype(in_buffer.dtype))
+        y0 = tvm.runtime.tensor(np.zeros(out_shape).astype(out_buffer.dtype))
+        y1 = tvm.runtime.tensor(np.zeros(out_shape).astype(out_buffer.dtype))
+        f_origin = tvm.compile(origin)
+        f_scheduled = tvm.compile(scheduled)
+        f_origin(x, y0)
+        f_scheduled(x, y1)
+        tvm.testing.assert_allclose(y0.numpy(), y1.numpy())
+
+
+def test_int64_indices_batch_decompose_padding():
+    @T.prim_func(s_tir=True)
+    def before_decompose(
+        x: T.Buffer((T.int64(1), T.int64(128), T.int64(128)), "int32"),
+        y: T.Buffer((T.int64(1), T.int64(140), T.int64(128)), "int32"),
+    ):
+        for b, i, j in T.grid(T.int64(1), T.int64(140), T.int64(128)):
+            with T.sblock("block"):
+                vb, vi, vj = T.axis.remap("SSS", [b, i, j])
+                y[vb, vi, vj] = T.if_then_else(vi < T.int64(128), x[vb, vi, vj], 0)
+
+    @T.prim_func(s_tir=True)
+    def after_decompose(
+        x: T.Buffer((T.int64(1), T.int64(128), T.int64(128)), "int32"),
+        y: T.Buffer((T.int64(1), T.int64(140), T.int64(128)), "int32"),
+    ):
+        # with T.sblock("root"):
+        for b, i in T.grid(T.int64(1), T.int64(140)):
+            # Use T.serial(T.int64(0), T.int64(128)) so iter_var dom.min is int64
+            # (matches schedule output; `range(T.int64(...))` would emit an int32 min).
+            for j in T.serial(T.int64(0), T.int64(128)):
+                with T.sblock("block_pad_const"):
+                    vb = T.axis.spatial(T.int64(1), T.int64(0))
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    T.reads()
+                    T.writes(y[vb, vi, vj])
+                    y[vb, vi, vj] = 0
+            for j in T.serial(T.int64(0), T.int64(128)):
+                with T.sblock("block"):
+                    vb = T.axis.spatial(T.int64(1), T.int64(0))
+                    vi = T.axis.spatial(T.int64(128), i)
+                    vj = T.axis.spatial(T.int64(128), j)
+                    T.where(i < T.int64(128))
+                    T.reads(x[vb, vi, vj])
+                    T.writes(y[vb, vi, vj])
+                    y[vb, vi, vj] = x[vb, vi, vj]
+
+    sch = tvm.s_tir.Schedule(before_decompose, debug_mask="all")
+    block = sch.get_sblock("block")
+    sch.decompose_padding(block, sch.get_loops(block)[2])
+    check_decompose_padding(before_decompose, sch.mod["main"], after_decompose, check_run=False)
+
+
+def test_1d_decompose_padding():
+    @T.prim_func(s_tir=True)
+    def before_decompose(x: T.Buffer(128, "int32"), y: T.Buffer(140, "int32")):
+        for i in range(140):
+            with T.sblock("block"):
+                vi = T.axis.remap("S", [i])
+                y[vi] = T.if_then_else(vi >= 6 and vi < 134, x[vi - 6], 0, dtype="int32")
+
+    @T.prim_func(s_tir=True)
+    def after_decompose(x: T.Buffer(128, "int32"), y: T.Buffer(140, "int32")):
+        for i in T.serial(140):
+            with T.sblock("block_pad_const"):
+                vi = T.axis.spatial(140, i)
+                T.reads()
+                T.writes(y[vi])
+                y[vi] = 0
+        for i in T.serial(128):
+            with T.sblock("block"):
+                vi = T.axis.spatial(128, i)
+                T.reads(x[vi])
+                T.writes(y[vi + 6])
+                y[vi + 6] = x[vi]
+
+    sch = tvm.s_tir.Schedule(before_decompose, debug_mask="all")
+    block = sch.get_sblock("block")
+    sch.decompose_padding(block, sch.get_loops(block)[0])
+    check_decompose_padding(before_decompose, sch.mod["main"], after_decompose, check_run=False)
+
+
+@T.prim_func(s_tir=True)
+def sum_pool_2d(
+    x: T.Buffer((1, 16, 225, 225), "int8"), tensor: T.Buffer((1, 16, 225, 225), "int8")
+):
+    pad_temp = T.sblock_alloc_buffer([1, 16, 231, 231], dtype="int8")
+    for i0, i1, i2, i3 in T.grid(1, 16, 231, 231):
+        with T.sblock("pad_temp"):
+            ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+            pad_temp[ax0, ax1, ax2, ax3] = T.if_then_else(
+                3 <= ax2 and ax2 < 228 and 3 <= ax3 and ax3 < 228,
+                x[ax0, ax1, ax2 - 3, ax3 - 3],
+                T.int8(0),
+                dtype="int8",
+            )
+    for i0, i1, i2, i3, i4, i5 in T.grid(1, 16, 225, 225, 7, 7):
+        with T.sblock("tensor"):
+            ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+            with T.init():
+                tensor[ax0, ax1, ax2, ax3] = T.int8(0)
+            tensor[ax0, ax1, ax2, ax3] = (
+                tensor[ax0, ax1, ax2, ax3] + pad_temp[ax0, ax1, ax2 + rv0, ax3 + rv1]
+            )
+
+
+def test_decompose_hw_padding_direct():
+    """Case 0. direct decompose"""
+
+    @T.prim_func(s_tir=True)
+    def pooling_decompose_0(
+        x: T.Buffer((1, 16, 225, 225), "int8"), tensor: T.Buffer((1, 16, 225, 225), "int8")
+    ):
+        pad_temp = T.sblock_alloc_buffer([1, 16, 231, 231], dtype="int8")
+        for i0, i1, i2, i3 in T.grid(1, 16, 231, 231):
+            with T.sblock("pad_temp_pad_const"):
+                ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                pad_temp[ax0, ax1, ax2, ax3] = T.int8(0)
+        for i0, i1, i2, i3 in T.grid(1, 16, 225, 225):
+            with T.sblock("pad_temp"):
+                ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                pad_temp[ax0, ax1, ax2 + 3, ax3 + 3] = x[ax0, ax1, ax2, ax3]
+        for i0, i1, i2, i3, i4, i5 in T.grid(1, 16, 225, 225, 7, 7):
+            with T.sblock("tensor"):
+                ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+                with T.init():
+                    tensor[ax0, ax1, ax2, ax3] = T.int8(0)
+                tensor[ax0, ax1, ax2, ax3] = (
+                    tensor[ax0, ax1, ax2, ax3] + pad_temp[ax0, ax1, ax2 + rv0, ax3 + rv1]
+                )
+
+    sch = tvm.s_tir.Schedule(sum_pool_2d, debug_mask="all")
+    pad = sch.get_sblock("pad_temp")
+    sch.decompose_padding(pad, sch.get_loops(pad)[0])
+    check_decompose_padding(sum_pool_2d, sch.mod["main"], pooling_decompose_0, check_run=True)
+
+
+def test_decompose_hw_padding_tiled():
+    """Case 1. tiling and then decompose"""
+
+    @T.prim_func(s_tir=True)
+    def pooling_decompose_1(
+        x: T.Buffer((1, 16, 225, 225), "int8"), tensor: T.Buffer((1, 16, 225, 225), "int8")
+    ) -> None:
+        pad_temp = T.sblock_alloc_buffer([1, 16, 231, 231], dtype="int8")
+        for i0, i2_0, i3_0 in T.grid(1, 3, 3):
+            for ax0, ax1, ax2 in T.grid(16, 81, 81):
+                with T.sblock("pad_temp_pad_const"):
+                    ax0_1 = T.axis.spatial(1, 0)
+                    ax1_1 = T.axis.spatial(16, ax0)
+                    ax2_1 = T.axis.spatial(231, i2_0 * 75 + ax1)
+                    ax3 = T.axis.spatial(231, i3_0 * 75 + ax2)
+                    T.reads()
+                    T.writes(pad_temp[ax0_1, ax1_1, ax2_1, ax3])
+                    pad_temp[ax0_1, ax1_1, ax2_1, ax3] = T.int8(0)
+            for ax0, ax1, ax2 in T.grid(16, 81, 81):
+                with T.sblock("pad_temp"):
+                    ax0_2 = T.axis.spatial(1, 0)
+                    ax1_2 = T.axis.spatial(16, ax0)
+                    ax2_2 = T.axis.spatial(225, i2_0 * 75 + ax1 - 3)
+                    ax3 = T.axis.spatial(225, i3_0 * 75 + ax2 - 3)
+                    T.where(
+                        3 <= i2_0 * 75 + ax1
+                        and i2_0 * 75 + ax1 < 228
+                        and 3 <= i3_0 * 75 + ax2
+                        and i3_0 * 75 + ax2 < 228
+                    )
+                    T.reads(x[ax0_2, ax1_2, ax2_2, ax3])
+                    T.writes(pad_temp[ax0_2, ax1_2, ax2_2 + 3, ax3 + 3])
+                    pad_temp[ax0_2, ax1_2, ax2_2 + 3, ax3 + 3] = x[ax0_2, ax1_2, ax2_2, ax3]
+            for i1, i2_1, i3_1, i4, i5 in T.grid(16, 75, 75, 7, 7):
+                with T.sblock("tensor"):
+                    ax0_3, ax1_3 = T.axis.remap("SS", [i0, i1])
+                    ax2_3 = T.axis.spatial(225, i2_0 * 75 + i2_1)
+                    ax3 = T.axis.spatial(225, i3_0 * 75 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.reads(pad_temp[ax0_3, ax1_3, ax2_3 + rv0, ax3 + rv1])
+                    T.writes(tensor[ax0_3, ax1_3, ax2_3, ax3])
+                    with T.init():
+                        tensor[ax0_3, ax1_3, ax2_3, ax3] = T.int8(0)
+                    tensor[ax0_3, ax1_3, ax2_3, ax3] = (
+                        tensor[ax0_3, ax1_3, ax2_3, ax3]
+                        + pad_temp[ax0_3, ax1_3, ax2_3 + rv0, ax3 + rv1]
+                    )
+
+    sch = tvm.s_tir.Schedule(sum_pool_2d, debug_mask="all")
+    block = sch.get_sblock("tensor")
+    pad = sch.get_sblock("pad_temp")
+    n, c, h, w, kh, kw = sch.get_loops(block)
+    ho, hi = sch.split(h, [3, 75])
+    wo, wi = sch.split(w, [3, 75])
+    sch.reorder(n, ho, wo, c, hi, wi, kh, kw)
+    sch.compute_at(sch.get_sblock("pad_temp"), wo)
+    sch.decompose_padding(pad, sch.get_loops(pad)[3])
+    check_decompose_padding(sum_pool_2d, sch.mod["main"], pooling_decompose_1, check_run=True)
+
+
+def test_decompose_hw_padding_tiled_and_lift_pad():
+    """Case 2. tiling and then decompose, lift const pad values to outer loop"""
+
+    @T.prim_func(s_tir=True)
+    def pooling_decompose_2(
+        x: T.Buffer((1, 16, 225, 225), "int8"), tensor: T.Buffer((1, 16, 225, 225), "int8")
+    ) -> None:
+        pad_temp = T.sblock_alloc_buffer([1, 16, 231, 231], dtype="int8")
+        for i0, i2_0, i3_0, ax0, ax1, ax2 in T.grid(1, 3, 3, 16, 81, 81):
+            with T.sblock("pad_temp_pad_const"):
+                ax0_1 = T.axis.spatial(1, 0)
+                ax1_1 = T.axis.spatial(16, ax0)
+                ax2_1 = T.axis.spatial(231, i2_0 * 75 + ax1)
+                ax3 = T.axis.spatial(231, i3_0 * 75 + ax2)
+                T.reads()
+                T.writes(pad_temp[ax0_1, ax1_1, ax2_1, ax3])
+                pad_temp[ax0_1, ax1_1, ax2_1, ax3] = T.int8(0)
+        for i0, i2_0, i3_0 in T.grid(1, 3, 3):
+            for ax0, ax1, ax2 in T.grid(16, 81, 81):
+                with T.sblock("pad_temp"):
+                    ax0_2 = T.axis.spatial(1, 0)
+                    ax1_2 = T.axis.spatial(16, ax0)
+                    ax2_2 = T.axis.spatial(225, i2_0 * 75 + ax1 - 3)
+                    ax3 = T.axis.spatial(225, i3_0 * 75 + ax2 - 3)
+                    T.where(
+                        3 <= i2_0 * 75 + ax1
+                        and i2_0 * 75 + ax1 < 228
+                        and 3 <= i3_0 * 75 + ax2
+                        and i3_0 * 75 + ax2 < 228
+                    )
+                    T.reads(x[ax0_2, ax1_2, ax2_2, ax3])
+                    T.writes(pad_temp[ax0_2, ax1_2, ax2_2 + 3, ax3 + 3])
+                    pad_temp[ax0_2, ax1_2, ax2_2 + 3, ax3 + 3] = x[ax0_2, ax1_2, ax2_2, ax3]
+            for i1, i2_1, i3_1, i4, i5 in T.grid(16, 75, 75, 7, 7):
+                with T.sblock("tensor"):
+                    ax0_3, ax1_3 = T.axis.remap("SS", [i0, i1])
+                    ax2_3 = T.axis.spatial(225, i2_0 * 75 + i2_1)
+                    ax3 = T.axis.spatial(225, i3_0 * 75 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.reads(pad_temp[ax0_3, ax1_3, ax2_3 + rv0, ax3 + rv1])
+                    T.writes(tensor[ax0_3, ax1_3, ax2_3, ax3])
+                    with T.init():
+                        tensor[ax0_3, ax1_3, ax2_3, ax3] = T.int8(0)
+                    tensor[ax0_3, ax1_3, ax2_3, ax3] = (
+                        tensor[ax0_3, ax1_3, ax2_3, ax3]
+                        + pad_temp[ax0_3, ax1_3, ax2_3 + rv0, ax3 + rv1]
+                    )
+
+    sch = tvm.s_tir.Schedule(sum_pool_2d, debug_mask="all")
+    block = sch.get_sblock("tensor")
+    pad = sch.get_sblock("pad_temp")
+    n, c, h, w, kh, kw = sch.get_loops(block)
+    ho, hi = sch.split(h, [3, 75])
+    wo, wi = sch.split(w, [3, 75])
+    sch.reorder(n, ho, wo, c, hi, wi, kh, kw)
+    sch.compute_at(sch.get_sblock("pad_temp"), wo)
+    sch.decompose_padding(pad, sch.get_loops(pad)[0])
+    check_decompose_padding(sum_pool_2d, sch.mod["main"], pooling_decompose_2, check_run=True)
+
+
+def test_decompose_hw_padding_non_perfect_tiled():
+    """Case 3. non-perfect tiling and then decompose"""
+
+    @T.prim_func(s_tir=True)
+    def pooling_decompose_3(
+        x: T.Buffer((1, 16, 225, 225), "int8"), tensor: T.Buffer((1, 16, 225, 225), "int8")
+    ) -> None:
+        pad_temp = T.sblock_alloc_buffer([1, 16, 231, 231], dtype="int8")
+        for i0, i2_0, i3_0 in T.grid(1, 3, 3):
+            for ax0, ax1, ax2 in T.grid(16, 86, 86):
+                with T.sblock("pad_temp_pad_const"):
+                    ax0_1 = T.axis.spatial(1, 0)
+                    ax1_1 = T.axis.spatial(16, ax0)
+                    ax2_1 = T.axis.spatial(231, i2_0 * 80 + ax1)
+                    ax3 = T.axis.spatial(231, i3_0 * 80 + ax2)
+                    T.where(i2_0 * 80 + ax1 < 231 and i3_0 * 80 + ax2 < 231)
+                    T.reads()
+                    T.writes(pad_temp[ax0_1, ax1_1, ax2_1, ax3])
+                    pad_temp[ax0_1, ax1_1, ax2_1, ax3] = T.int8(0)
+            for ax0, ax1, ax2 in T.grid(16, 86, 86):
+                with T.sblock("pad_temp"):
+                    ax0_2 = T.axis.spatial(1, 0)
+                    ax1_2 = T.axis.spatial(16, ax0)
+                    ax2_2 = T.axis.spatial(225, i2_0 * 80 + ax1 - 3)
+                    ax3 = T.axis.spatial(225, i3_0 * 80 + ax2 - 3)
+                    T.where(
+                        3 <= i2_0 * 80 + ax1
+                        and i2_0 * 80 + ax1 < 228
+                        and 3 <= i3_0 * 80 + ax2
+                        and i3_0 * 80 + ax2 < 228
+                        and i2_0 * 80 + ax1 < 231
+                        and i3_0 * 80 + ax2 < 231
+                    )
+                    T.reads(x[ax0_2, ax1_2, ax2_2, ax3])
+                    T.writes(pad_temp[ax0_2, ax1_2, ax2_2 + 3, ax3 + 3])
+                    pad_temp[ax0_2, ax1_2, ax2_2 + 3, ax3 + 3] = x[ax0_2, ax1_2, ax2_2, ax3]
+            for i1, i2_1, i3_1, i4, i5 in T.grid(16, 80, 80, 7, 7):
+                with T.sblock("tensor"):
+                    ax0_3, ax1_3 = T.axis.remap("SS", [i0, i1])
+                    ax2_3 = T.axis.spatial(225, i2_0 * 80 + i2_1)
+                    ax3 = T.axis.spatial(225, i3_0 * 80 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.where(i2_0 * 80 + i2_1 < 225 and i3_0 * 80 + i3_1 < 225)
+                    T.reads(pad_temp[ax0_3, ax1_3, ax2_3 + rv0, ax3 + rv1])
+                    T.writes(tensor[ax0_3, ax1_3, ax2_3, ax3])
+                    with T.init():
+                        tensor[ax0_3, ax1_3, ax2_3, ax3] = T.int8(0)
+                    tensor[ax0_3, ax1_3, ax2_3, ax3] = (
+                        tensor[ax0_3, ax1_3, ax2_3, ax3]
+                        + pad_temp[ax0_3, ax1_3, ax2_3 + rv0, ax3 + rv1]
+                    )
+
+    sch = tvm.s_tir.Schedule(sum_pool_2d, debug_mask="all")
+    block = sch.get_sblock("tensor")
+    pad = sch.get_sblock("pad_temp")
+    n, c, h, w, kh, kw = sch.get_loops(block)
+    ho, hi = sch.split(h, [None, 80])
+    wo, wi = sch.split(w, [None, 80])
+    sch.reorder(n, ho, wo, c, hi, wi, kh, kw)
+    sch.compute_at(sch.get_sblock("pad_temp"), wo)
+    sch.decompose_padding(pad, sch.get_loops(pad)[3])
+    check_decompose_padding(sum_pool_2d, sch.mod["main"], pooling_decompose_3, check_run=True)
+
+
+def test_decompose_wrt_single_child_subtree():
+    """Test the case when the decompose position is under the single child subtree"""
+
+    @T.prim_func(s_tir=True)
+    def pad_op(
+        x: T.Buffer((1, 16, 225, 225), "int8"),
+        y: T.Buffer((1, 16, 231, 231), dtype="int8"),
+    ):
+        for i0, i1, i2, i3 in T.grid(1, 16, 231, 231):
+            with T.sblock("pad_temp"):
+                ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                y[ax0, ax1, ax2, ax3] = T.if_then_else(
+                    3 <= ax2 and ax2 < 228 and 3 <= ax3 and ax3 < 228,
+                    x[ax0, ax1, ax2 - 3, ax3 - 3],
+                    T.int8(0),
+                    dtype="int8",
+                )
+
+    @T.prim_func(s_tir=True)
+    def pad_op_after(
+        x: T.Buffer((1, 16, 225, 225), "int8"), y: T.Buffer((1, 16, 231, 231), "int8")
+    ):
+        for i0, i1 in T.grid(1, 16):
+            for i2, i3 in T.grid(231, 231):
+                with T.sblock("pad_temp_pad_const"):
+                    ax0 = T.axis.spatial(1, 0)
+                    ax1, ax2, ax3 = T.axis.remap("SSS", [i1, i2, i3])
+                    y[ax0, ax1, ax2, ax3] = T.int8(0)
+            for i2, i3 in T.grid(225, 225):
+                with T.sblock("pad_temp"):
+                    ax0 = T.axis.spatial(1, 0)
+                    ax1, ax2, ax3 = T.axis.remap("SSS", [i1, i2, i3])
+                    y[ax0, ax1, ax2 + 3, ax3 + 3] = x[ax0, ax1, ax2, ax3]
+
+    sch = tvm.s_tir.Schedule(pad_op, debug_mask="all")
+    pad = sch.get_sblock("pad_temp")
+    _, _, h, _ = sch.get_loops(pad)
+    sch.decompose_padding(pad, h)
+    check_decompose_padding(pad_op, sch.mod["main"], pad_op_after, check_run=True)
+
+
+def test_not_to_decompose_trivial_predicate():
+    """Test the case when the padding condition is trivial"""
+
+    @T.prim_func(s_tir=True)
+    def trivial_pad(
+        x: T.Buffer((1, 16, 225, 225), "int8"), y: T.Buffer([1, 16, 225, 225], dtype="int8")
+    ):
+        for i0, i1, i2, i3 in T.grid(1, 16, 225, 225):
+            with T.sblock("pad_temp"):
+                ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                y[ax0, ax1, ax2, ax3] = T.if_then_else(
+                    0 <= ax2 and ax2 < 225 and 0 <= ax3 and ax3 < 225,
+                    x[ax0, ax1, ax2, ax3],
+                    T.int8(0),
+                    dtype="int8",
+                )
+
+    sch = tvm.s_tir.Schedule(trivial_pad, debug_mask="all")
+    pad = sch.get_sblock("pad_temp")
+    _, _, h, _ = sch.get_loops(pad)
+    assert not sch.can_decompose_padding(pad, h)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_error.py

@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import s_tir, tirx
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j in T.grid(128, 128):
+        with T.sblock("init"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = T.float32(0)
+        for k in range(128):
+            with T.sblock("update"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def two_kernels(var_A: T.handle, var_B: T.handle, seq_len: T.int32):
+    T.func_attr({"tirx.noalias": True})
+    A = T.match_buffer(var_A, (1, seq_len * 8), "int32")
+    B = T.match_buffer(var_B, (1, seq_len * 8), "int32", align=8)
+    with T.sblock("exclusive_scan"):
+        T.reads()
+        T.writes()
+        s8: T.int32 = seq_len * 8
+        if s8 == 0:
+            blockIdx_x = T.launch_thread("blockIdx.x", 1)
+        else:
+            with T.launch_thread("threadIdx.x", 1024) as threadIdx_x:
+                blockIdx_x = T.launch_thread("blockIdx.x", T.ceildiv(s8, 1024))
+                i: T.int32 = blockIdx_x * 1024 + threadIdx_x
+                if i < s8:
+                    B[i // s8, i % s8] = A[i // s8, i % s8]
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+
+def test_tir_schedule_error_detail():
+    sch = tvm.s_tir.Schedule(matmul, debug_mask="all", error_render_level="detail")
+    with pytest.raises(tvm.s_tir.ScheduleError) as excinfo:
+        sch.get_sblock("wrong_name")
+    (msg,) = excinfo.value.args
+    assert "Cannot find a block with the name: wrong_name" in msg
+
+
+def test_tir_schedule_error_fast():
+    sch = tvm.s_tir.Schedule(matmul, debug_mask="all", error_render_level="fast")
+    with pytest.raises(tvm.s_tir.ScheduleError) as excinfo:
+        sch.get_sblock("wrong_name")
+    (msg,) = excinfo.value.args
+    assert "Cannot find a block with the specified name" in msg
+
+
+def test_tir_schedule_error_none():
+    sch = tvm.s_tir.Schedule(matmul, debug_mask="all", error_render_level="none")
+    with pytest.raises(tvm.s_tir.ScheduleError) as excinfo:
+        sch.get_sblock("wrong_name")
+    (msg,) = excinfo.value.args
+    assert "(not rendered)" in msg
+
+
+def test_tir_schedule_attribute_error():
+    sch = tvm.s_tir.Schedule(matmul)
+    with pytest.raises(AttributeError):
+        sch.non_existent_field()
+
+
+def test_tir_schedule_two_kernels():
+    s_tir.Schedule(two_kernels)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_for_kind.py

@@ -0,0 +1,706 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E741, F401
+import sys
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def element_wise(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128, 128))
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_parallelized(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128, 128))
+    for i0 in T.parallel(0, 128):
+        for i1 in T.serial(0, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i0, i1])
+                B[vi, vj] = A[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_i_bound(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128, 128))
+    for i0 in T.thread_binding(0, 128, thread="threadIdx.x"):
+        for i1 in T.serial(0, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i0, i1])
+                B[vi, vj] = A[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_compute_at_split(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+    B = T.sblock_alloc_buffer((128, 128))
+    for i in T.serial(0, 128):
+        for j0 in T.serial(0, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j0])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for j1o, j1i in T.grid(32, 4):
+            with T.sblock("C"):
+                vi = T.axis.S(128, i)
+                vj = T.axis.S(128, j1o * 4 + j1i)
+                C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_compute_at_split_vectorized(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+    B = T.sblock_alloc_buffer((128, 128))
+    for i in T.serial(0, 128):
+        for j0 in T.serial(0, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j0])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for j1o in T.serial(0, 32):
+            for j1i in T.vectorized(0, 4):
+                with T.sblock("C"):
+                    vi = T.axis.S(128, i)
+                    vj = T.axis.S(128, j1o * 4 + j1i)
+                    C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_split_predicate(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    for i, j_0, j_1 in T.grid(128, 13, 10):
+        with T.sblock("B"):
+            T.where(j_0 * 10 + j_1 < 128)
+            vi = T.axis.S(128, i)
+            vj = T.axis.S(128, j_0 * 10 + j_1)
+            B[vi, vj] = A[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_split_predicate_parallelized(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    for i in T.serial(0, 128):
+        for j_0 in T.parallel(0, 13):
+            for j_1 in T.serial(0, 10):
+                with T.sblock("B"):
+                    T.where(j_0 * 10 + j_1 < 128)
+                    vi = T.axis.S(128, i)
+                    vj = T.axis.S(128, j_0 * 10 + j_1)
+                    B[vi, vj] = A[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_split_predicate_vectorized(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    for i in T.vectorized(0, 128):
+        for j_0, j_1 in T.grid(13, 10):
+            with T.sblock("B"):
+                T.where(j_0 * 10 + j_1 < 128)
+                vi = T.axis.S(128, i)
+                vj = T.axis.S(128, j_0 * 10 + j_1)
+                B[vi, vj] = A[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_compute_at_split_j0_j1o_bound(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+    B = T.sblock_alloc_buffer((128, 128))
+    for i in T.serial(0, 128):
+        for j0 in T.thread_binding(0, 128, thread="threadIdx.x"):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j0])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for j1o in T.thread_binding(0, 32, thread="threadIdx.x"):
+            for j1i in T.serial(0, 4):
+                with T.sblock("C"):
+                    vi = T.axis.S(128, i)
+                    vj = T.axis.S(128, j1o * 4 + j1i)
+                    C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("C"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def rowsum(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128,))
+
+    for i, k in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vk = T.axis.remap("SR", [i, k])
+            with T.init():
+                B[vi] = 0.0
+            B[vi] = B[vi] + A[vi, vk]
+
+
+@T.prim_func(s_tir=True)
+def rowsum_unrolled(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128,))
+    for i0 in T.unroll(0, 128):
+        for i1 in T.serial(0, 128):
+            with T.sblock("B"):
+                vi, vk = T.axis.remap("SR", [i0, i1])
+                with T.init():
+                    B[vi] = 0.0
+                B[vi] = B[vi] + A[vi, vk]
+
+
+@T.prim_func(s_tir=True)
+def rowsum_not_quasi_affine(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128,))
+
+    for i, k in T.grid(128, 16):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i)
+            vk = T.axis.R(128, T.floordiv(k * k, 2))
+            with T.init():
+                B[vi] = 0.0
+            B[vi] = B[vi] + A[vi, vk]
+
+
+@T.prim_func(s_tir=True)
+def rowsum_not_compact_data_flow(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128,))
+
+    for i, k in T.grid(128, 16):
+        with T.sblock("B"):
+            vi, vk = T.axis.remap("SR", [i, k])
+            with T.init():
+                B[vk] = 0.0
+            B[vk] = B[vk] + A[vi, vk]
+
+
+@T.prim_func(s_tir=True)
+def rowsum_cross_thread_reduction(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128,))
+    for i0 in T.serial(0, 128):
+        for i1 in T.thread_binding(0, 128, thread="threadIdx.x"):
+            with T.sblock("B"):
+                vi, vk = T.axis.remap("SR", [i0, i1])
+                with T.init():
+                    B[vi] = 0.0
+                B[vi] = B[vi] + A[vi, vk]
+
+
+@T.prim_func(s_tir=True)
+def opaque_block(a: T.handle) -> None:
+    A = T.match_buffer(a, (16,))
+    for i in T.serial(0, 15):
+        with T.sblock("opaque"):
+            A[i + 1] = A[i + 1] + A[i]
+
+
+@T.prim_func(s_tir=True)
+def block_inside_init(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128, 128], dtype="float32")
+    B = T.match_buffer(b, [128, 128], dtype="float32")
+    for i in T.serial(0, 128):
+        with T.sblock("outer"):
+            vi = T.axis.S(128, i)
+            with T.init():
+                for j in T.serial(0, 128):
+                    with T.sblock("init"):
+                        vj = T.axis.S(128, j)
+                        B[vi, vj] = 0.0
+            for k in T.serial(0, 128):
+                for j in T.serial(0, 128):
+                    with T.sblock("inner"):
+                        vj, vk = T.axis.remap("SR", [j, k])
+                        B[vi, vj] = B[vi, vj] + A[vi, vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def thread_bound_block_inside_init(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128, 128], dtype="float32")
+    B = T.match_buffer(b, [128, 128], dtype="float32")
+    for i in T.thread_binding(0, 128, thread="threadIdx.x"):
+        with T.sblock("outer"):
+            vi = T.axis.S(128, i)
+            with T.init():
+                for j in T.serial(0, 128):
+                    with T.sblock("init"):
+                        vj = T.axis.S(128, j)
+                        B[vi, vj] = 0.0
+            for k in T.serial(0, 128):
+                for j in T.serial(0, 128):
+                    with T.sblock("inner"):
+                        vj, vk = T.axis.remap("SR", [j, k])
+                        B[vi, vj] = B[vi, vj] + A[vi, vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def decomposed_gemm(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    C: T.Buffer((16, 16), "float32"),
+):
+    local = T.sblock_alloc_buffer((16, 16), "float32")
+    for i, j in T.grid(4, 4):
+        for ii, jj in T.grid(4, 4):
+            with T.sblock("init"):
+                vi = T.axis.S(16, i * 4 + ii)
+                vj = T.axis.S(16, j * 4 + jj)
+                local[vi, vj] = 0
+        for k, ii, jj in T.grid(16, 4, 4):
+            with T.sblock("update"):
+                vi = T.axis.S(16, i * 4 + ii)
+                vj = T.axis.S(16, j * 4 + jj)
+                vk = T.axis.R(16, k)
+                local[vi, vj] += A[vi, vk] * B[vj, vk]
+        for ii, jj in T.grid(4, 4):
+            with T.sblock("C"):
+                vi = T.axis.S(16, i * 4 + ii)
+                vj = T.axis.S(16, j * 4 + jj)
+                C[vi, vj] = local[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def decomposed_gemm_after_vectorize(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    C: T.Buffer((16, 16), "float32"),
+):
+    local = T.sblock_alloc_buffer((16, 16), "float32")
+    for i, j in T.grid(4, 4):
+        for ii, jj in T.grid(4, 4):
+            with T.sblock("init"):
+                vi = T.axis.S(16, i * 4 + ii)
+                vj = T.axis.S(16, j * 4 + jj)
+                local[vi, vj] = 0
+        for k, ii, jj in T.grid(16, 4, 4):
+            with T.sblock("update"):
+                vi = T.axis.S(16, i * 4 + ii)
+                vj = T.axis.S(16, j * 4 + jj)
+                vk = T.axis.R(16, k)
+                local[vi, vj] += A[vi, vk] * B[vj, vk]
+        for ii in range(4):
+            for jj in T.vectorized(4):
+                with T.sblock("C"):
+                    vi = T.axis.S(16, i * 4 + ii)
+                    vj = T.axis.S(16, j * 4 + jj)
+                    C[vi, vj] = local[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def nested_block_bind(
+    A: T.Buffer((16, 16, 16, 16), "float32"), B: T.Buffer((16, 16, 16), "float32")
+):
+    for i, j in T.grid(16, 16):
+        with T.sblock("outer"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            for k, l in T.grid(16, 16):
+                with T.sblock("inner"):
+                    vk, vl = T.axis.remap("SR", [k, l])
+                    with T.init():
+                        B[vi, vj, vk] = 0.0
+                    B[vi, vj, vk] = B[vi, vj, vk] + A[vi, vj, vk, vl]
+
+
+@T.prim_func(s_tir=True)
+def thread_bound_nested_block(
+    A: T.Buffer((16, 16, 16, 16), "float32"), B: T.Buffer((16, 16, 16), "float32")
+) -> None:
+    for i in T.serial(16):
+        for j in T.thread_binding(16, thread="blockIdx.x"):
+            with T.sblock("outer"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                for k in T.serial(16):
+                    for l in T.thread_binding(16, thread="threadIdx.x"):
+                        with T.sblock("inner"):
+                            vk, vl = T.axis.remap("SR", [k, l])
+                            with T.init():
+                                B[vi, vj, vk] = T.float32(0)
+                            B[vi, vj, vk] = B[vi, vj, vk] + A[vi, vj, vk, vl]
+
+
+@T.prim_func(s_tir=True)
+def nested_block_bind_after_cache_read(
+    A: T.Buffer((16, 16), "float32"), B: T.Buffer((16,), "float32")
+) -> None:
+    for i in T.serial(16):
+        with T.sblock("outer"):
+            vi = T.axis.spatial(16, i)
+            A_shared = T.sblock_alloc_buffer([1, 16], dtype="float32", scope="shared")
+            for ax0, ax1 in T.grid(1, 16):
+                with T.sblock("A_shared"):
+                    v0 = T.axis.spatial(16, vi + ax0)
+                    v1 = T.axis.spatial(16, ax1)
+                    A_shared[v0, v1] = A[v0, v1]
+            for j in T.serial(16):
+                with T.sblock("inner"):
+                    vj = T.axis.reduce(16, j)
+                    with T.init():
+                        B[vi] = T.float32(0)
+                    B[vi] = B[vi] + A_shared[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def thread_bound_nested_block_after_cache_read(
+    A: T.Buffer((16, 16), "float32"), B: T.Buffer((16,), "float32")
+) -> None:
+    for i in T.thread_binding(16, thread="blockIdx.x"):
+        with T.sblock("outer"):
+            vi = T.axis.spatial(16, i)
+            A_shared = T.sblock_alloc_buffer([1, 16], dtype="float32", scope="shared")
+            for ax0, ax1 in T.grid(1, 16):
+                with T.sblock("A_shared"):
+                    v0 = T.axis.spatial(16, vi + ax0)
+                    v1 = T.axis.spatial(16, ax1)
+                    A_shared[v0, v1] = A[v0, v1]
+            for j in T.thread_binding(16, thread="threadIdx.x"):
+                with T.sblock("inner"):
+                    vj = T.axis.reduce(16, j)
+                    with T.init():
+                        B[vi] = T.float32(0)
+                    B[vi] = B[vi] + A_shared[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def decomposed_gemm_parallelize_init(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    C: T.Buffer((16, 16), "float32"),
+) -> None:
+    local = T.sblock_alloc_buffer([16, 16], dtype="float32")
+    for i, j in T.grid(4, 4):
+        for ii in T.serial(4):
+            for jj in T.vectorized(4):
+                with T.sblock("init"):
+                    vi = T.axis.spatial(16, i * 4 + ii)
+                    vj = T.axis.spatial(16, j * 4 + jj)
+                    T.reads()
+                    T.writes(local[vi, vj])
+                    local[vi, vj] = 0
+        for k, ii, jj in T.grid(16, 4, 4):
+            with T.sblock("update"):
+                vi = T.axis.spatial(16, i * 4 + ii)
+                vj = T.axis.spatial(16, j * 4 + jj)
+                vk = T.axis.reduce(16, k)
+                T.reads(local[vi, vj], A[vi, vk], B[vj, vk])
+                T.writes(local[vi, vj])
+                local[vi, vj] = local[vi, vj] + A[vi, vk] * B[vj, vk]
+        for ii, jj in T.grid(4, 4):
+            with T.sblock("C"):
+                vi = T.axis.spatial(16, i * 4 + ii)
+                vj = T.axis.spatial(16, j * 4 + jj)
+                T.reads(local[vi, vj])
+                T.writes(C[vi, vj])
+                C[vi, vj] = local[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def scatter_compute(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
+    for i in T.grid(8):
+        with T.sblock("first_half"):
+            vi = T.axis.spatial(16, 8 + i)
+            B[vi] = A[vi - 8]
+
+    for i in T.grid(8):
+        with T.sblock("last_half"):
+            vi = T.axis.spatial(16, i)
+            B[vi] = A[vi + 8]
+
+
+@T.prim_func(s_tir=True)
+def scatter_compute_parallelize(
+    A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")
+) -> None:
+    # body
+    # with T.sblock("root")
+    for i in T.parallel(8):
+        with T.sblock("first_half"):
+            vi = T.axis.spatial(16, 8 + i)
+            T.reads(A[vi - 8])
+            T.writes(B[vi])
+            B[vi] = A[vi - 8]
+    for i in T.parallel(8):
+        with T.sblock("last_half"):
+            vi = T.axis.spatial(16, i)
+            T.reads(A[vi + 8])
+            T.writes(B[vi])
+            B[vi] = A[vi + 8]
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+
+def test_parallel():
+    s = tvm.s_tir.Schedule(element_wise, debug_mask="all")
+    i, _ = s.get_loops(s.get_sblock("B"))
+    s.parallel(i)
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], element_wise_parallelized)
+    verify_trace_roundtrip(s, mod=element_wise)
+
+
+def test_parallel_predicate():
+    s = tvm.s_tir.Schedule(element_wise_split_predicate, debug_mask="all")
+    _, j, _ = s.get_loops(s.get_sblock("B"))
+    s.parallel(j)
+    assert_structural_equal_ignore_global_symbol(
+        s.mod["main"], element_wise_split_predicate_parallelized
+    )
+    verify_trace_roundtrip(s, mod=element_wise_split_predicate)
+
+
+def test_parallel_reduction_block_iter():
+    s = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    _, _, k = s.get_loops(s.get_sblock("C"))
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.parallel(k)
+
+
+def test_parallel_not_quasi_affine():
+    s = tvm.s_tir.Schedule(rowsum_not_quasi_affine, debug_mask="all")
+    i, _ = s.get_loops(s.get_sblock("B"))
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.parallel(i)
+
+
+def test_parallel_not_compact_data_flow():
+    s = tvm.s_tir.Schedule(rowsum_not_compact_data_flow, debug_mask="all")
+    i, _ = s.get_loops(s.get_sblock("B"))
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.parallel(i)
+
+
+def test_vectorize():
+    s = tvm.s_tir.Schedule(element_wise_compute_at_split, debug_mask="all")
+    _, _, j1i = s.get_loops(s.get_sblock("C"))
+    s.vectorize(j1i)
+    assert_structural_equal_ignore_global_symbol(
+        s.mod["main"], element_wise_compute_at_split_vectorized
+    )
+    verify_trace_roundtrip(s, mod=element_wise_compute_at_split)
+
+
+def test_vectorize_predicate():
+    s = tvm.s_tir.Schedule(element_wise_split_predicate, debug_mask="all")
+    i, _, _ = s.get_loops(s.get_sblock("B"))
+    s.vectorize(i)
+    assert_structural_equal_ignore_global_symbol(
+        s.mod["main"], element_wise_split_predicate_vectorized
+    )
+    verify_trace_roundtrip(s, mod=element_wise_split_predicate)
+
+
+def test_vectorize_opaque_block():
+    s = tvm.s_tir.Schedule(opaque_block, debug_mask="all")
+    (i,) = s.get_loops(s.get_sblock("opaque"))
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.vectorize(i)
+
+
+def test_unroll():
+    s = tvm.s_tir.Schedule(rowsum, debug_mask="all")
+    i, _ = s.get_loops(s.get_sblock("B"))
+    s.unroll(i)
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], rowsum_unrolled)
+    verify_trace_roundtrip(s, mod=rowsum)
+
+
+def test_unroll_after_bind():
+    s = tvm.s_tir.Schedule(rowsum, debug_mask="all")
+    i, _ = s.get_loops(s.get_sblock("B"))
+    s.bind(i, "blockIdx.x")
+    s.unroll(i)
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], rowsum_unrolled)
+    verify_trace_roundtrip(s, mod=rowsum)
+
+
+def test_bind1():
+    s = tvm.s_tir.Schedule(element_wise, debug_mask="all")
+    i, _ = s.get_loops(s.get_sblock("B"))
+    s.bind(i, "threadIdx.x")
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], element_wise_i_bound)
+    verify_trace_roundtrip(s, mod=element_wise)
+
+
+def test_bind2():
+    s = tvm.s_tir.Schedule(element_wise_compute_at_split, debug_mask="all")
+    _, j0 = s.get_loops(s.get_sblock("B"))
+    _, j1o, _ = s.get_loops(s.get_sblock("C"))
+    s.bind(j0, "threadIdx.x")
+    s.bind(j1o, "threadIdx.x")
+    assert_structural_equal_ignore_global_symbol(
+        s.mod["main"], element_wise_compute_at_split_j0_j1o_bound
+    )
+    verify_trace_roundtrip(s, mod=element_wise_compute_at_split)
+
+
+def test_bind_cross_thread_reduction():
+    s = tvm.s_tir.Schedule(rowsum, debug_mask="all")
+    _, k = s.get_loops(s.get_sblock("B"))
+    s.bind(k, "threadIdx.x")
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], rowsum_cross_thread_reduction)
+    verify_trace_roundtrip(s, mod=rowsum)
+
+
+def test_bind_not_cross_thread_reduction():
+    s = tvm.s_tir.Schedule(rowsum, debug_mask="all")
+    _, k = s.get_loops(s.get_sblock("B"))
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.bind(k, "blockIdx.x")
+
+
+def test_bind_after_bind():
+    s = tvm.s_tir.Schedule(element_wise, debug_mask="all")
+    i, _ = s.get_loops(s.get_sblock("B"))
+    s.bind(i, "blockIdx.x")
+    s.bind(i, "threadIdx.x")
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], element_wise_i_bound)
+    verify_trace_roundtrip(s, mod=element_wise)
+
+
+def test_block_inside_init():
+    s = tvm.s_tir.Schedule(block_inside_init, debug_mask="all")
+    (i,) = s.get_loops(s.get_sblock("outer"))
+    s.bind(i, "threadIdx.x")
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], thread_bound_block_inside_init)
+    verify_trace_roundtrip(s, mod=block_inside_init)
+
+
+def test_vectorize_after_decompose():
+    s = tvm.s_tir.Schedule(decomposed_gemm, debug_mask="all")
+    jj = s.get_loops(s.get_sblock("C"))[-1]
+    s.vectorize(jj)
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], decomposed_gemm_after_vectorize)
+    verify_trace_roundtrip(s, mod=decomposed_gemm)
+
+
+def test_nested_block_bind():
+    s = tvm.s_tir.Schedule(nested_block_bind)
+    block_outer = s.get_sblock("outer")
+    block_inner = s.get_sblock("inner")
+    _, j = s.get_loops(block_outer)
+    _, l = s.get_loops(block_inner)
+    s.bind(l, "threadIdx.x")
+    s.bind(j, "blockIdx.x")
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], thread_bound_nested_block)
+    verify_trace_roundtrip(s, mod=nested_block_bind)
+
+
+def test_nexted_block_bind_after_cache_read():
+    s = tvm.s_tir.Schedule(nested_block_bind_after_cache_read)
+    block_outer = s.get_sblock("outer")
+    block_inner = s.get_sblock("inner")
+    (i,) = s.get_loops(block_outer)
+    (j,) = s.get_loops(block_inner)
+    s.bind(i, "blockIdx.x")
+    s.bind(j, "threadIdx.x")
+    assert_structural_equal_ignore_global_symbol(
+        s.mod["main"], thread_bound_nested_block_after_cache_read
+    )
+    verify_trace_roundtrip(s, mod=nested_block_bind_after_cache_read)
+
+
+def test_vectorize_init():
+    s = tvm.s_tir.Schedule(decomposed_gemm, debug_mask="all")
+    init_blk = s.get_sblock("init")
+    upd_blk = s.get_sblock("update")
+    _, _, ii_0, jj_0 = s.get_loops(init_blk)
+    _, _, k_1, ii_1, jj_1 = s.get_loops(upd_blk)
+    s.vectorize(jj_0)
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], decomposed_gemm_parallelize_init)
+    verify_trace_roundtrip(s, mod=decomposed_gemm)
+
+
+def test_scatter_parallelize():
+    s = tvm.s_tir.Schedule(scatter_compute, debug_mask="all")
+    first = s.get_sblock("first_half")
+    last = s.get_sblock("last_half")
+    (i_0,) = s.get_loops(first)
+    (i_1,) = s.get_loops(last)
+    s.parallel(i_0)
+    s.parallel(i_1)
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], scatter_compute_parallelize)
+    verify_trace_roundtrip(s, mod=scatter_compute)
+
+
+def test_bind_thread_iter_var_dtype():
+    @T.prim_func(private=True, s_tir=True)
+    def before(
+        A: T.Buffer((T.int64(128), T.int64(128))),
+        B: T.Buffer((T.int64(128), T.int64(128))),
+    ) -> None:
+        for i, j in T.grid(T.int64(128), T.int64(128)):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                B[vi, vj] = A[vi, vj] * 2.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(
+        A: T.Buffer((T.int64(128), T.int64(128))),
+        B: T.Buffer((T.int64(128), T.int64(128))),
+    ) -> None:
+        for i0 in T.thread_binding(T.int64(128), thread="threadIdx.x"):
+            # Use T.serial with explicit int64 min so the inner sblock iter_var dom
+            # is all-int64 (matches what `s.bind` emits; `range(T.int64(128))` parses
+            # min as int32 even when extent is int64).
+            for i1 in T.serial(T.int64(0), T.int64(128)):
+                with T.sblock("B"):
+                    vi, vj = T.axis.remap("SS", [i0, i1])
+                    B[vi, vj] = A[vi, vj] * 2.0
+
+    s = tvm.s_tir.Schedule(before, debug_mask="all")
+    i, _ = s.get_loops(s.get_sblock("B"))
+    s.bind(i, "threadIdx.x")
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], expected)
+    verify_trace_roundtrip(s, mod=before)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue.py

@@ -0,0 +1,279 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401, F841
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_before(
+    A: T.Buffer((16, 16), "int8"),
+    B: T.Buffer((16, 16), "int8"),
+    C: T.Buffer((16, 16), "int32"),
+    D: T.Buffer((16, 16), "int32"),
+) -> None:
+    temp = T.sblock_alloc_buffer((16, 16), dtype="int32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("multiply"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                temp[vi, vj] = T.int32(0)
+            temp[vi, vj] = temp[vi, vj] + T.cast(A[vi, vk], "int32") * T.cast(B[vj, vk], "int32")
+    for i, j in T.grid(16, 16):
+        with T.sblock("add"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = temp[vi, vj] + C[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_expected(
+    A: T.Buffer((16, 16), "int8"),
+    B: T.Buffer((16, 16), "int8"),
+    C: T.Buffer((16, 16), "int32"),
+    D: T.Buffer((16, 16), "int32"),
+) -> None:
+    temp = T.sblock_alloc_buffer((16, 16), dtype="int32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("multiply"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
+            T.writes(D[vi, vj])
+            with T.init():
+                D[vi, vj] = C[vi, vj]
+            D[vi, vj] = D[vi, vj] + T.cast(A[vi, vk], "int32") * T.cast(B[vj, vk], "int32")
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_fp32_before(
+    A: T.Buffer((32, 32), "float32"),
+    B: T.Buffer((32, 32), "float32"),
+    C: T.Buffer((32, 32), "float32"),
+    D: T.Buffer((32, 32), "float32"),
+) -> None:
+    temp = T.sblock_alloc_buffer((32, 32), dtype="float32")
+    for i, j, k in T.grid(32, 32, 32):
+        with T.sblock("multiply"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                temp[vi, vj] = T.float32(0)
+            temp[vi, vj] = temp[vi, vj] + A[vi, vk] * B[vj, vk]
+    for i, j in T.grid(32, 32):
+        with T.sblock("add"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = temp[vi, vj] + C[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_fp32_expected(
+    A: T.Buffer((32, 32), "float32"),
+    B: T.Buffer((32, 32), "float32"),
+    C: T.Buffer((32, 32), "float32"),
+    D: T.Buffer((32, 32), "float32"),
+) -> None:
+    temp = T.sblock_alloc_buffer((32, 32), dtype="float32")
+    for i, j, k in T.grid(32, 32, 32):
+        with T.sblock("multiply"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
+            T.writes(D[vi, vj])
+            with T.init():
+                D[vi, vj] = C[vi, vj]
+            D[vi, vj] = D[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_multiple_epilogue_before(
+    A: T.Buffer((16, 16), "int8"),
+    B: T.Buffer((16, 16), "int8"),
+    C: T.Buffer((16, 16), "int32"),
+    D: T.Buffer((16, 16), "int32"),
+    E: T.Buffer((16, 16), "int32"),
+) -> None:
+    temp = T.sblock_alloc_buffer((16, 16), dtype="int32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("multiply"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                temp[vi, vj] = T.int32(0)
+            temp[vi, vj] = temp[vi, vj] + T.cast(A[vi, vk], "int32") * T.cast(B[vj, vk], "int32")
+    for i, j in T.grid(16, 16):
+        with T.sblock("add"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = temp[vi, vj] + C[vi, vj]
+    for i, j in T.grid(16, 16):
+        with T.sblock("add2"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            E[vi, vj] = temp[vi, vj] + C[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_multiple_epilogue_expected(
+    A: T.Buffer((16, 16), "int8"),
+    B: T.Buffer((16, 16), "int8"),
+    C: T.Buffer((16, 16), "int32"),
+    D: T.Buffer((16, 16), "int32"),
+    E: T.Buffer((16, 16), "int32"),
+) -> None:
+    temp = T.sblock_alloc_buffer((16, 16), dtype="int32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("multiply"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
+            T.writes(D[vi, vj])
+            with T.init():
+                D[vi, vj] = C[vi, vj]
+            D[vi, vj] = D[vi, vj] + T.cast(A[vi, vk], "int32") * T.cast(B[vj, vk], "int32")
+    for i, j in T.grid(16, 16):
+        with T.sblock("add2"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(temp[vi, vj], C[vi, vj])
+            T.writes(E[vi, vj])
+            E[vi, vj] = temp[vi, vj] + C[vi, vj]
+
+
+def test_fuse_reduction_epilogue_basic():
+    sch = tvm.s_tir.Schedule(matmul_bias_before, debug_mask="all")
+    sch.fuse_reduction_epilogue("multiply", "add")
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], matmul_bias_expected)
+    verify_trace_roundtrip(sch=sch, mod=matmul_bias_before)
+
+
+def test_fuse_reduction_epilogue_fp32():
+    sch = tvm.s_tir.Schedule(matmul_bias_fp32_before, debug_mask="all")
+    sch.fuse_reduction_epilogue("multiply", "add")
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], matmul_bias_fp32_expected)
+    verify_trace_roundtrip(sch=sch, mod=matmul_bias_fp32_before)
+
+
+def test_fuse_reduction_epilogue_numerical_correctness():
+    sch_original = tvm.s_tir.Schedule(matmul_bias_before, debug_mask="all")
+    mod_original = tvm.compile(sch_original.mod["main"], target="llvm")
+
+    sch_fused = tvm.s_tir.Schedule(matmul_bias_before, debug_mask="all")
+    sch_fused.fuse_reduction_epilogue("multiply", "add")
+    mod_fused = tvm.compile(sch_fused.mod["main"], target="llvm")
+
+    A_np = np.random.randint(-128, 127, size=(16, 16), dtype="int8")
+    B_np = np.random.randint(-128, 127, size=(16, 16), dtype="int8")
+    C_np = np.random.randint(-1000, 1000, size=(16, 16), dtype="int32")
+
+    expected = (A_np.astype("int32") @ B_np.T.astype("int32")) + C_np
+
+    D_original_tvm = tvm.runtime.tensor(np.zeros((16, 16), dtype="int32"))
+    D_fused_tvm = tvm.runtime.tensor(np.zeros((16, 16), dtype="int32"))
+
+    mod_original(
+        tvm.runtime.tensor(A_np), tvm.runtime.tensor(B_np), tvm.runtime.tensor(C_np), D_original_tvm
+    )
+
+    mod_fused(
+        tvm.runtime.tensor(A_np), tvm.runtime.tensor(B_np), tvm.runtime.tensor(C_np), D_fused_tvm
+    )
+
+    D_original = D_original_tvm.numpy()
+    D_fused = D_fused_tvm.numpy()
+
+    tvm.testing.assert_allclose(D_original, expected, rtol=1e-5)
+    tvm.testing.assert_allclose(D_fused, expected, rtol=1e-5)
+    tvm.testing.assert_allclose(D_fused, D_original, rtol=1e-5)
+
+
+def test_fuse_reduction_epilogue_multiple_epilogue():
+    sch = tvm.s_tir.Schedule(matmul_bias_multiple_epilogue_before, debug_mask="all")
+    sch.fuse_reduction_epilogue("multiply", "add")
+    assert_structural_equal_ignore_global_symbol(
+        sch.mod["main"], matmul_bias_multiple_epilogue_expected
+    )
+    verify_trace_roundtrip(sch=sch, mod=matmul_bias_multiple_epilogue_before)
+
+    mod = tvm.compile(sch.mod["main"], target="llvm")
+    assert mod is not None
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_invalid_multiple_use_before(
+    A: T.Buffer((16, 16), "int8"),
+    B: T.Buffer((16, 16), "int8"),
+    C1: T.Buffer((16, 16), "int32"),
+    C2: T.Buffer((16, 16), "int32"),
+    D: T.Buffer((16, 16), "int32"),
+) -> None:
+    """Epilogue uses the reduction result twice; fusion must be rejected."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="int32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("multiply"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                temp[vi, vj] = T.int32(0)
+            temp[vi, vj] = temp[vi, vj] + T.cast(A[vi, vk], "int32") * T.cast(B[vj, vk], "int32")
+    for i, j in T.grid(16, 16):
+        with T.sblock("bad_epilogue"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            # temp[vi, vj] is used twice in the epilogue expression
+            D[vi, vj] = (temp[vi, vj] + C1[vi, vj]) * (temp[vi, vj] + C2[vi, vj])
+
+
+def test_fuse_reduction_epilogue_reject_multiple_use():
+    """fusion should be rejected when the reduction result appears more than once."""
+    sch = tvm.s_tir.Schedule(matmul_bias_invalid_multiple_use_before, debug_mask="all")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.fuse_reduction_epilogue("multiply", "bad_epilogue")
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_invalid_scaling_before(
+    A: T.Buffer((16, 16), "int8"),
+    B: T.Buffer((16, 16), "int8"),
+    C: T.Buffer((16, 16), "int32"),
+    D: T.Buffer((16, 16), "int32"),
+) -> None:
+    """Epilogue scales the reduction result; fusion must be rejected."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="int32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("multiply"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                temp[vi, vj] = T.int32(0)
+            temp[vi, vj] = temp[vi, vj] + T.cast(A[vi, vk], "int32") * T.cast(B[vj, vk], "int32")
+    for i, j in T.grid(16, 16):
+        with T.sblock("scaled_epilogue"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            # temp[vi, vj] is scaled by 2 before adding bias; this must not be fused.
+            D[vi, vj] = temp[vi, vj] * T.int32(2) + C[vi, vj]
+
+
+def test_fuse_reduction_epilogue_reject_scaling():
+    """fusion should be rejected when the reduction result is scaled by Mul/Div/Mod."""
+    sch = tvm.s_tir.Schedule(matmul_bias_invalid_scaling_before, debug_mask="all")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.fuse_reduction_epilogue("multiply", "scaled_epilogue")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py

@@ -0,0 +1,273 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E501, F401, F841
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def matmul_clipping_before(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    D: T.Buffer((16, 16), "float32"),
+    lower: T.float32,
+    upper: T.float32,
+) -> None:
+    """Original function with separate reduction and clipping epilogue blocks."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="float32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                temp[vi, vj] = T.float32(0)
+            temp[vi, vj] = temp[vi, vj] + A[vi, vk] * B[vj, vk]
+
+    for i, j in T.grid(16, 16):
+        with T.sblock("clipping"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = T.min(T.max(temp[vi, vj], lower), upper)
+
+
+@T.prim_func(s_tir=True)
+def matmul_clipping_expected(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    D: T.Buffer((16, 16), "float32"),
+    lower: T.float32,
+    upper: T.float32,
+) -> None:
+    """Expected function after fusion (Clipping)."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="float32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            T.reads(A[vi, vk], B[vj, vk])
+            T.writes(D[vi, vj])
+            with T.init():
+                D[vi, vj] = T.min(T.max(T.float32(0), lower), upper)
+            D[vi, vj] = T.min(T.max(D[vi, vj] + A[vi, vk] * B[vj, vk], lower), upper)
+
+
+def test_matmul_clipping():
+    """Test fusion of matmul with clipping epilogue."""
+    sch = tvm.s_tir.Schedule(matmul_clipping_before, debug_mask="all")
+    sch.fuse_reduction_epilogue("matmul", "clipping")
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], matmul_clipping_expected)
+    verify_trace_roundtrip(sch=sch, mod=matmul_clipping_before)
+
+
+@T.prim_func(s_tir=True)
+def matmul_clipping_before_per_iteration(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    D: T.Buffer((16, 16), "float32"),
+) -> None:
+    """Original function with per-iteration clipping (same semantics as fused)."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="float32")
+    lower = T.float32(-5.0)
+    upper = T.float32(5.0)
+    for i, j in T.grid(16, 16):
+        with T.sblock("init"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            temp[vi, vj] = T.min(T.max(T.float32(0), lower), upper)  # Clip init
+
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            # Per-iteration clipping
+            temp[vi, vj] = T.min(T.max(temp[vi, vj] + A[vi, vk] * B[vj, vk], lower), upper)
+
+    for i, j in T.grid(16, 16):
+        with T.sblock("copy"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = temp[vi, vj]
+
+
+def test_matmul_clipping_correctness_unified():
+    """Test that original and fused produce identical results with per-iteration clipping."""
+    A_np = np.random.randn(16, 16).astype("float32")
+    B_np = np.random.randn(16, 16).astype("float32")
+    lower = -5.0
+    upper = 5.0
+
+    # NumPy reference for per-iteration clipping
+    D_ref = np.clip(0.0, lower, upper)  # init with clipping
+    for k in range(16):
+        D_ref = np.clip(D_ref + np.outer(A_np[:, k], B_np[:, k]), lower, upper)
+
+    # TVM execution (original with per-iteration clipping)
+    mod_original = tvm.compile(matmul_clipping_before_per_iteration, target="llvm")
+    D_original_tvm = tvm.runtime.tensor(np.zeros((16, 16), dtype="float32"))
+    mod_original(
+        tvm.runtime.tensor(A_np),
+        tvm.runtime.tensor(B_np),
+        D_original_tvm,
+    )
+
+    # TVM execution (fused)
+    sch = tvm.s_tir.Schedule(matmul_clipping_before)
+    sch.fuse_reduction_epilogue("matmul", "clipping")
+    mod_fused = tvm.compile(sch.mod["main"], target="llvm")
+    D_fused_tvm = tvm.runtime.tensor(np.zeros((16, 16), dtype="float32"))
+    # Pass scalar values directly as Python floats
+    mod_fused(
+        tvm.runtime.tensor(A_np),
+        tvm.runtime.tensor(B_np),
+        D_fused_tvm,
+        lower,
+        upper,
+    )
+
+    D_original = D_original_tvm.numpy()
+    D_fused = D_fused_tvm.numpy()
+
+    # Now both should match exactly
+    np.testing.assert_allclose(D_original, D_ref, rtol=1e-5, atol=1e-6)
+    np.testing.assert_allclose(D_fused, D_ref, rtol=1e-5, atol=1e-6)
+    np.testing.assert_allclose(D_original, D_fused, rtol=1e-5, atol=1e-6)
+
+
+@T.prim_func(s_tir=True)
+def matmul_clipping_multiple_epilogue_before(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    D: T.Buffer((16, 16), "float32"),
+    E: T.Buffer((16, 16), "float32"),
+    lower: T.float32,
+    upper: T.float32,
+) -> None:
+    """Original function with separate reduction and multiple epilogue blocks (one with clipping, one without)."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="float32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                temp[vi, vj] = T.float32(0)
+            temp[vi, vj] = temp[vi, vj] + A[vi, vk] * B[vj, vk]
+
+    for i, j in T.grid(16, 16):
+        with T.sblock("clipping"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = T.min(T.max(temp[vi, vj], lower), upper)
+
+    for i, j in T.grid(16, 16):
+        with T.sblock("copy"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            E[vi, vj] = temp[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def matmul_clipping_multiple_epilogue_expected(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    D: T.Buffer((16, 16), "float32"),
+    E: T.Buffer((16, 16), "float32"),
+    lower: T.float32,
+    upper: T.float32,
+) -> None:
+    """Expected function after fusion (Clipping) with multiple epilogue blocks."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="float32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            T.reads(A[vi, vk], B[vj, vk])
+            T.writes(D[vi, vj])
+            with T.init():
+                D[vi, vj] = T.min(T.max(T.float32(0), lower), upper)
+            D[vi, vj] = T.min(T.max(D[vi, vj] + A[vi, vk] * B[vj, vk], lower), upper)
+    for i, j in T.grid(16, 16):
+        with T.sblock("copy"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(temp[vi, vj])
+            T.writes(E[vi, vj])
+            E[vi, vj] = temp[vi, vj]
+
+
+def test_matmul_clipping_multiple_epilogue():
+    """Test fusion with multiple epilogue blocks - one with clipping, one without.
+
+    Following the same pattern as test_fuse_reduction_epilogue_multiple_epilogue,
+    this test verifies that fusion works correctly when there are multiple
+    epilogue blocks. The temp buffer is kept because the second epilogue block
+    still needs it.
+    """
+    sch = tvm.s_tir.Schedule(matmul_clipping_multiple_epilogue_before, debug_mask="all")
+    sch.fuse_reduction_epilogue("matmul", "clipping")
+    assert_structural_equal_ignore_global_symbol(
+        sch.mod["main"], matmul_clipping_multiple_epilogue_expected
+    )
+    verify_trace_roundtrip(sch=sch, mod=matmul_clipping_multiple_epilogue_before)
+
+    mod = tvm.compile(sch.mod["main"], target="llvm")
+    assert mod is not None
+
+
+# Test commutative variants of clipping patterns
+@pytest.mark.parametrize(
+    "pattern_func",
+    [
+        lambda temp, lower, upper: T.min(T.max(temp, lower), upper),  # min(max(temp, lower), upper)
+        lambda temp, lower, upper: T.min(upper, T.max(temp, lower)),  # min(upper, max(temp, lower))
+        lambda temp, lower, upper: T.min(T.max(lower, temp), upper),  # min(max(lower, temp), upper)
+        lambda temp, lower, upper: T.max(T.min(temp, upper), lower),  # max(min(temp, upper), lower)
+        lambda temp, lower, upper: T.max(lower, T.min(temp, upper)),  # max(lower, min(temp, upper))
+    ],
+)
+def test_matmul_clipping_commutative_variants(pattern_func):
+    """Test that all commutative variants of clipping patterns are recognized."""
+    lower = -5.0
+    upper = 5.0
+
+    @T.prim_func(s_tir=True)
+    def test_func(
+        A: T.Buffer((8, 8), "float32"),
+        B: T.Buffer((8, 8), "float32"),
+        D: T.Buffer((8, 8), "float32"),
+    ) -> None:
+        temp = T.sblock_alloc_buffer((8, 8), dtype="float32")
+        for i, j, k in T.grid(8, 8, 8):
+            with T.sblock("matmul"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    temp[vi, vj] = T.float32(0)
+                temp[vi, vj] = temp[vi, vj] + A[vi, vk] * B[vj, vk]
+
+        for i, j in T.grid(8, 8):
+            with T.sblock("clipping"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                D[vi, vj] = pattern_func(temp[vi, vj], T.float32(lower), T.float32(upper))
+
+    sch = tvm.s_tir.Schedule(test_func, debug_mask="all")
+    # Should not raise an error - all variants should be recognized
+    sch.fuse_reduction_epilogue("matmul", "clipping")
+    verify_trace_roundtrip(sch=sch, mod=test_func)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py

@@ -0,0 +1,231 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E501, F401, F841
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_relu_before(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    C: T.Buffer((16, 16), "float32"),
+    D: T.Buffer((16, 16), "float32"),
+) -> None:
+    """Original function with separate reduction and epilogue blocks (Bias + ReLU)."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="float32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                temp[vi, vj] = T.float32(0)
+            temp[vi, vj] = temp[vi, vj] + A[vi, vk] * B[vj, vk]
+
+    for i, j in T.grid(16, 16):
+        with T.sblock("bias_relu"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = T.max(temp[vi, vj] + C[vi, vj], T.float32(0))
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_relu_before_per_iteration(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    C: T.Buffer((16, 16), "float32"),
+    D: T.Buffer((16, 16), "float32"),
+) -> None:
+    """Original function with per-iteration ReLU (same semantics as fused)."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="float32")
+    for i, j in T.grid(16, 16):
+        with T.sblock("init"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            temp[vi, vj] = T.max(C[vi, vj], T.float32(0))  # ReLU on bias
+
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            # Per-iteration ReLU
+            temp[vi, vj] = T.max(temp[vi, vj] + A[vi, vk] * B[vj, vk], T.float32(0))
+
+    for i, j in T.grid(16, 16):
+        with T.sblock("copy"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = temp[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_relu_expected(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    C: T.Buffer((16, 16), "float32"),
+    D: T.Buffer((16, 16), "float32"),
+) -> None:
+    """Expected function after fusion (Bias + ReLU)."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="float32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
+            T.writes(D[vi, vj])
+            with T.init():
+                D[vi, vj] = T.max(C[vi, vj], T.float32(0))
+            D[vi, vj] = T.max(D[vi, vj] + A[vi, vk] * B[vj, vk], T.float32(0))
+
+
+def test_matmul_bias_relu():
+    """Test fusion of matmul with bias + ReLU epilogue."""
+    sch = tvm.s_tir.Schedule(matmul_bias_relu_before, debug_mask="all")
+    sch.fuse_reduction_epilogue("matmul", "bias_relu")
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], matmul_bias_relu_expected)
+    verify_trace_roundtrip(sch=sch, mod=matmul_bias_relu_before)
+
+
+def test_matmul_bias_relu_correctness_unified():
+    """Test that original and fused produce identical results with per-iteration ReLU."""
+    A_np = np.random.randn(16, 16).astype("float32")
+    B_np = np.random.randn(16, 16).astype("float32")
+    C_np = np.random.randn(16, 16).astype("float32")
+
+    # NumPy reference for per-iteration ReLU
+    # Simulate per-iteration ReLU behavior
+    # Original code computes A[vi, vk] * B[vj, vk] which is A[i, k] * B[j, k]
+    # For each k: add outer product of A[:, k] and B[:, k]
+    D_ref = np.maximum(C_np, 0)  # init with ReLU on bias
+    for k in range(16):
+        # A[:, k] is shape (16,), B[:, k] is shape (16,)
+        # Outer product: A[:, k] * B[:, k] for all i, j = A[i, k] * B[j, k]
+        # Using broadcasting: A[:, k:k+1] * B[:, k:k+1].T gives (16, 1) * (1, 16) = (16, 16)
+        D_ref = np.maximum(D_ref + np.outer(A_np[:, k], B_np[:, k]), 0)
+
+    # TVM execution (original with per-iteration ReLU)
+    mod_original = tvm.compile(matmul_bias_relu_before_per_iteration, target="llvm")
+    D_original_tvm = tvm.runtime.tensor(np.zeros((16, 16), dtype="float32"))
+    mod_original(
+        tvm.runtime.tensor(A_np),
+        tvm.runtime.tensor(B_np),
+        tvm.runtime.tensor(C_np),
+        D_original_tvm,
+    )
+
+    # TVM execution (fused)
+    sch = tvm.s_tir.Schedule(matmul_bias_relu_before)
+    sch.fuse_reduction_epilogue("matmul", "bias_relu")
+    mod_fused = tvm.compile(sch.mod["main"], target="llvm")
+    D_fused_tvm = tvm.runtime.tensor(np.zeros((16, 16), dtype="float32"))
+    mod_fused(
+        tvm.runtime.tensor(A_np),
+        tvm.runtime.tensor(B_np),
+        tvm.runtime.tensor(C_np),
+        D_fused_tvm,
+    )
+
+    D_original = D_original_tvm.numpy()
+    D_fused = D_fused_tvm.numpy()
+
+    # Now both should match exactly
+    np.testing.assert_allclose(D_original, D_ref, rtol=1e-5, atol=1e-6)
+    np.testing.assert_allclose(D_fused, D_ref, rtol=1e-5, atol=1e-6)
+    np.testing.assert_allclose(D_original, D_fused, rtol=1e-5, atol=1e-6)
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_relu_multiple_epilogue_before(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    C: T.Buffer((16, 16), "float32"),
+    D: T.Buffer((16, 16), "float32"),
+    E: T.Buffer((16, 16), "float32"),
+) -> None:
+    """Original function with separate reduction and multiple epilogue blocks (one with ReLU, one without)."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="float32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                temp[vi, vj] = T.float32(0)
+            temp[vi, vj] = temp[vi, vj] + A[vi, vk] * B[vj, vk]
+
+    for i, j in T.grid(16, 16):
+        with T.sblock("bias_relu"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = T.max(temp[vi, vj] + C[vi, vj], T.float32(0))
+
+    for i, j in T.grid(16, 16):
+        with T.sblock("bias"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            E[vi, vj] = temp[vi, vj] + C[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def matmul_bias_relu_multiple_epilogue_expected(
+    A: T.Buffer((16, 16), "float32"),
+    B: T.Buffer((16, 16), "float32"),
+    C: T.Buffer((16, 16), "float32"),
+    D: T.Buffer((16, 16), "float32"),
+    E: T.Buffer((16, 16), "float32"),
+) -> None:
+    """Expected function after fusion (Bias + ReLU) with multiple epilogue blocks."""
+    temp = T.sblock_alloc_buffer((16, 16), dtype="float32")
+    for i, j, k in T.grid(16, 16, 16):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
+            T.writes(D[vi, vj])
+            with T.init():
+                D[vi, vj] = T.max(C[vi, vj], T.float32(0))
+            D[vi, vj] = T.max(D[vi, vj] + A[vi, vk] * B[vj, vk], T.float32(0))
+    for i, j in T.grid(16, 16):
+        with T.sblock("bias"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(temp[vi, vj], C[vi, vj])
+            T.writes(E[vi, vj])
+            E[vi, vj] = temp[vi, vj] + C[vi, vj]
+
+
+def test_matmul_bias_relu_multiple_epilogue():
+    """Test fusion with multiple epilogue blocks - one with ReLU, one without.
+
+    Following the same pattern as test_fuse_reduction_epilogue_multiple_epilogue,
+    this test verifies that fusion works correctly when there are multiple
+    epilogue blocks. The temp buffer is kept because the second epilogue block
+    still needs it.
+    """
+    sch = tvm.s_tir.Schedule(matmul_bias_relu_multiple_epilogue_before, debug_mask="all")
+    sch.fuse_reduction_epilogue("matmul", "bias_relu")
+    assert_structural_equal_ignore_global_symbol(
+        sch.mod["main"], matmul_bias_relu_multiple_epilogue_expected
+    )
+    verify_trace_roundtrip(sch=sch, mod=matmul_bias_relu_multiple_epilogue_before)
+
+    mod = tvm.compile(sch.mod["main"], target="llvm")
+    assert mod is not None
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_instruction.py

@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# mypy: ignore-errors
+# ruff: noqa: F401
+import sys
+
+import pytest
+
+import tvm.testing
+from tvm.s_tir.schedule import Instruction, InstructionKind, LoopRV, SBlockRV
+
+
+def test_inst_kind_get():
+    kind = InstructionKind.get("EnterPostproc")
+    assert not kind.is_pure
+    assert kind.name == "EnterPostproc"
+
+
+def test_inst_construct_1():
+    block = SBlockRV()
+    loop0 = LoopRV()
+    loop1 = LoopRV()
+    inst = Instruction(
+        kind=InstructionKind.get("GetLoops"),
+        inputs=[block],
+        attrs=[],
+        outputs=[loop0, loop1],
+    )
+    assert str(inst) == "_, _ = sch.get_loops(block=_)"
+    assert len(inst.inputs) == 1
+    assert len(inst.attrs) == 0
+    assert len(inst.outputs) == 2
+    assert inst.kind.same_as(InstructionKind.get("GetLoops"))
+    assert inst.inputs[0].same_as(block)
+    assert inst.outputs[0].same_as(loop0)
+    assert inst.outputs[1].same_as(loop1)
+
+
+def test_inst_construct_2():
+    block = SBlockRV()
+    inst = Instruction(
+        kind=InstructionKind.get("ComputeInline"),
+        inputs=[block],
+        attrs=[],
+        outputs=[],
+    )
+    assert str(inst) == "sch.compute_inline(block=_)"
+    assert len(inst.inputs) == 1
+    assert len(inst.attrs) == 0
+    assert len(inst.outputs) == 0
+    assert inst.kind.same_as(InstructionKind.get("ComputeInline"))
+    assert inst.inputs[0].same_as(block)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_merge.py

@@ -0,0 +1,278 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def elementwise(a: T.handle, c: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+    D = T.match_buffer(d, (64, 64))
+    B = T.sblock_alloc_buffer((128, 128))
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(A[vi, vj])
+            T.writes(B[vi, vj])
+            B[vi, vj] = A[vi, vj] * T.float32(2)
+    for i_0, j_0, i_1, j_1 in T.grid(8, 8, 16, 16):
+        with T.sblock("C"):
+            vi = T.axis.spatial(128, i_0 * 16 + i_1)
+            vj = T.axis.spatial(128, j_0 * 16 + j_1)
+            T.reads(B[vi, vj])
+            T.writes(C[vi, vj])
+            C[vi, vj] = B[vi, vj] + T.float32(1)
+    for i_0, j_0, i_1, j_1 in T.grid(8, 8, 8, 8):
+        with T.sblock("D"):
+            vi = T.axis.spatial(64, i_0 * 8 + i_1)
+            vj = T.axis.spatial(64, j_0 * 8 + j_1)
+            T.reads(B[vi, vj])
+            T.writes(D[vi, vj])
+            D[vi, vj] = B[vi, vj] + T.float32(2)
+
+
+@T.prim_func(s_tir=True)
+def elementwise_merged(a: T.handle, c: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+    D = T.match_buffer(d, (64, 64))
+    B = T.sblock_alloc_buffer((128, 128))
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(A[vi, vj])
+            T.writes(B[vi, vj])
+            B[vi, vj] = A[vi, vj] * T.float32(2)
+    for i_0_m in range(8):
+        for j_0, i_1, j_1 in T.grid(8, 16, 16):
+            with T.sblock("C"):
+                vi = T.axis.spatial(128, i_0_m * 16 + i_1)
+                vj = T.axis.spatial(128, j_0 * 16 + j_1)
+                T.reads(B[vi, vj])
+                T.writes(C[vi, vj])
+                C[vi, vj] = B[vi, vj] + T.float32(1)
+        for j_0, i_1, j_1 in T.grid(8, 8, 8):
+            with T.sblock("D"):
+                vi = T.axis.spatial(64, i_0_m * 8 + i_1)
+                vj = T.axis.spatial(64, j_0 * 8 + j_1)
+                T.reads(B[vi, vj])
+                T.writes(D[vi, vj])
+                D[vi, vj] = B[vi, vj] + T.float32(2)
+
+
+@T.prim_func(s_tir=True)
+def elementwise_merged2(a: T.handle, c: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+    D = T.match_buffer(d, (64, 64))
+    B = T.sblock_alloc_buffer((128, 128))
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(A[vi, vj])
+            T.writes(B[vi, vj])
+            B[vi, vj] = A[vi, vj] * T.float32(2)
+    for i_0_m, j_0_m in T.grid(8, 8):
+        for i_1, j_1 in T.grid(16, 16):
+            with T.sblock("C"):
+                vi = T.axis.spatial(128, i_0_m * 16 + i_1)
+                vj = T.axis.spatial(128, j_0_m * 16 + j_1)
+                T.reads(B[vi, vj])
+                T.writes(C[vi, vj])
+                C[vi, vj] = B[vi, vj] + T.float32(1)
+        for i_1, j_1 in T.grid(8, 8):
+            with T.sblock("D"):
+                vi = T.axis.spatial(64, i_0_m * 8 + i_1)
+                vj = T.axis.spatial(64, j_0_m * 8 + j_1)
+                T.reads(B[vi, vj])
+                T.writes(D[vi, vj])
+                D[vi, vj] = B[vi, vj] + T.float32(2)
+
+
+def test_merge():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_c = sch.get_sblock("C")
+    block_d = sch.get_sblock("D")
+    i = sch.get_loops(block_c)[0]
+    j = sch.get_loops(block_d)[0]
+    sch.merge(i, j)
+    assert_structural_equal_ignore_global_symbol(elementwise_merged, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+
+def test_merge2():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_c = sch.get_sblock("C")
+    block_d = sch.get_sblock("D")
+    i = sch.get_loops(block_c)[1]
+    j = sch.get_loops(block_d)[1]
+    sch.merge(i, j)
+    assert_structural_equal_ignore_global_symbol(elementwise_merged2, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+
+def test_merge_fail_not_only_child():
+    @T.prim_func(s_tir=True)
+    def elementwise_with_seq(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (128, 128, 128))
+        C = T.match_buffer(c, (128, 128, 128))
+        B = T.sblock_alloc_buffer((128, 128, 128))
+        D = T.sblock_alloc_buffer((128, 128, 128))
+        for i, j in T.grid(128, 128):
+            for k in T.serial(0, 128):
+                with T.sblock("D"):
+                    vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                    D[vi, vj, vk] = A[vi, vj, vk] * 2.0
+            for k in T.serial(0, 128):
+                with T.sblock("B"):
+                    vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                    B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+        for i, j in T.grid(128, 128):
+            for k in T.serial(0, 128):
+                with T.sblock("C"):
+                    vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                    C[vi, vj, vk] = B[vi, vj, vk] * 2.0
+
+    sch = tvm.s_tir.Schedule(elementwise_with_seq, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, _, b = sch.get_loops(block_b)
+    block_c = sch.get_sblock("C")
+    _, _, c = sch.get_loops(block_c)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.merge(b, c)
+
+
+def test_merge_fail_not_start_with_zero():
+    @T.prim_func(s_tir=True)
+    def elementwise_loops_not_start_with_zero(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (128, 128, 128))
+        C = T.match_buffer(c, (128, 128, 128))
+        B = T.sblock_alloc_buffer((128, 128, 128))
+        for i, j in T.grid(128, 128):
+            for k in T.serial(1, 128):
+                with T.sblock("B"):
+                    vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                    B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+        for i, j in T.grid(128, 128):
+            for k in T.serial(0, 128):
+                with T.sblock("C"):
+                    vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                    C[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+    sch = tvm.s_tir.Schedule(elementwise_loops_not_start_with_zero, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, _, b = sch.get_loops(block_b)
+    block_c = sch.get_sblock("C")
+    _, _, c = sch.get_loops(block_c)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.merge(b, c)
+
+
+def test_merge_fail_not_same_extent():
+    @T.prim_func(s_tir=True)
+    def elementwise_loops_not_same_extent(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (128, 128, 128))
+        C = T.match_buffer(c, (128, 128, 128))
+        B = T.sblock_alloc_buffer((64, 128, 128))
+        for i, j in T.grid(64, 128):
+            for k in T.serial(0, 128):
+                with T.sblock("B"):
+                    vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                    B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+        for i, j in T.grid(128, 128):
+            for k in T.serial(0, 128):
+                with T.sblock("C"):
+                    vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                    C[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+    sch = tvm.s_tir.Schedule(elementwise_loops_not_same_extent, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, _, b = sch.get_loops(block_b)
+    block_c = sch.get_sblock("C")
+    _, _, c = sch.get_loops(block_c)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.merge(b, c)
+
+
+def test_merge_fail_not_same_level():
+    @T.prim_func(s_tir=True)
+    def elementwise_not_same_level(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (128, 128, 128))
+        C = T.match_buffer(c, (128, 128, 128))
+        B = T.sblock_alloc_buffer((128, 128, 128))
+        for i, j in T.grid(128, 128):
+            for k in T.serial(0, 128):
+                with T.sblock("B"):
+                    vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                    B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+        for i, j in T.grid(128, 128):
+            for k in T.serial(0, 128):
+                with T.sblock("C"):
+                    vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                    C[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+    sch = tvm.s_tir.Schedule(elementwise_not_same_level, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, b, _ = sch.get_loops(block_b)
+    block_c = sch.get_sblock("C")
+    _, _, c = sch.get_loops(block_c)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.merge(b, c)
+
+
+def test_merge_fail_with_different_scope():
+    @T.prim_func(s_tir=True)
+    def elementwise_with_different_scope(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (128, 128, 128))
+        C = T.match_buffer(c, (128, 128, 128))
+        B = T.sblock_alloc_buffer((128, 128, 128))
+        with T.sblock("A"):
+            for i, j in T.grid(128, 128):
+                for k in T.serial(0, 128):
+                    with T.sblock("B"):
+                        vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                        B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+        for i, j in T.grid(128, 128):
+            for k in T.serial(0, 128):
+                with T.sblock("C"):
+                    vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                    C[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+    sch = tvm.s_tir.Schedule(elementwise_with_different_scope, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, _, b = sch.get_loops(block_b)
+    block_c = sch.get_sblock("C")
+    _, _, c = sch.get_loops(block_c)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.merge(b, c)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_pad_einsum.py

@@ -0,0 +1,307 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+
+@T.prim_func(s_tir=True)
+def matmul_before(
+    A: T.Buffer((128, 127), "float32"),
+    B: T.Buffer((127, 127), "float32"),
+    C: T.Buffer((128, 127), "float32"),
+) -> None:
+    A_shared = T.sblock_alloc_buffer((128, 127), "float32", scope="shared")
+    B_shared = T.sblock_alloc_buffer((127, 127), "float32", scope="shared")
+    C_shared = T.sblock_alloc_buffer((128, 127), "float32", scope="shared")
+    for i0, i1 in T.grid(128, 127):
+        with T.sblock("A"):
+            i, j = T.axis.remap("SS", [i0, i1])
+            A_shared[i, j] = A[i, j]
+    for i0, i1 in T.grid(127, 127):
+        with T.sblock("B"):
+            i, j = T.axis.remap("SS", [i0, i1])
+            B_shared[i, j] = B[i, j]
+    for i0, i1, i2 in T.grid(128, 127, 127):
+        with T.sblock("C_shared"):
+            i, j, k = T.axis.remap("SSR", [i0, i1, i2])
+            with T.init():
+                C_shared[i, j] = T.float32(0)
+            C_shared[i, j] = C_shared[i, j] + A_shared[i, k] * B_shared[k, j]
+    for i0, i1 in T.grid(128, 127):
+        with T.sblock("C"):
+            i, j = T.axis.remap("SS", [i0, i1])
+            C[i, j] = C_shared[i, j]
+
+
+@T.prim_func(s_tir=True)
+def matmul_expected(
+    A: T.Buffer((128, 127), "float32"),
+    B: T.Buffer((127, 127), "float32"),
+    C: T.Buffer((128, 127), "float32"),
+) -> None:
+    A_shared_padded = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="shared")
+    B_shared_padded = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="shared")
+    C_shared_padded = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="shared")
+    for i0, i1 in T.grid(128, 128):
+        with T.sblock("A"):
+            i, j = T.axis.remap("SS", [i0, i1])
+            T.reads(A[i, j])
+            T.writes(A_shared_padded[i, j])
+            A_shared_padded[i, j] = T.if_then_else(j < 127, A[i, j], T.float32(0), dtype="float32")
+    for i0, i1 in T.grid(128, 128):
+        with T.sblock("B"):
+            i, j = T.axis.remap("SS", [i0, i1])
+            T.reads(B[i, j])
+            T.writes(B_shared_padded[i, j])
+            B_shared_padded[i, j] = T.if_then_else(
+                i < 127 and j < 127, B[i, j], T.float32(0), dtype="float32"
+            )
+    for i0, i1, i2 in T.grid(128, 128, 128):
+        with T.sblock("C_shared"):
+            i, j, k = T.axis.remap("SSR", [i0, i1, i2])
+            T.reads(A_shared_padded[i, k], B_shared_padded[k, j])
+            T.writes(C_shared_padded[i, j])
+            with T.init():
+                C_shared_padded[i, j] = T.float32(0)
+            C_shared_padded[i, j] = (
+                C_shared_padded[i, j] + A_shared_padded[i, k] * B_shared_padded[k, j]
+            )
+    for i0, i1 in T.grid(128, 127):
+        with T.sblock("C"):
+            i, j = T.axis.remap("SS", [i0, i1])
+            T.reads(C_shared_padded[i, j])
+            T.writes(C[i, j])
+            C[i, j] = C_shared_padded[i, j]
+
+
+# pylint: enable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+
+def test_pad_matmul():
+    # pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+    @T.prim_func(s_tir=True)
+    def matmul_before(
+        a: T.handle,
+        b: T.handle,
+        c: T.handle,
+    ) -> None:
+        n = T.int32()
+        A = T.match_buffer(a, (128, 128), "float32")
+        B = T.match_buffer(b, (n, 128), "float32")
+        C = T.match_buffer(c, (128, n), "float32")
+        for i0, i1, i2 in T.grid(128, n, 128):
+            with T.sblock("C"):
+                i, j, k = T.axis.remap("SSR", [i0, i1, i2])
+                with T.init():
+                    C[i, j] = T.float32(0)
+                C[i, j] = C[i, j] + A[i, k] * B[j, k]
+
+    @T.prim_func(s_tir=True)
+    def matmul_after(
+        a: T.handle,
+        b: T.handle,
+        c: T.handle,
+    ):
+        n = T.int32()
+        A = T.match_buffer(a, (128, 128), "float32")
+        B = T.match_buffer(b, (n, 128), "float32")
+        C = T.match_buffer(c, (128, n), "float32")
+        B_pad = T.sblock_alloc_buffer(((n + 31) // 32 * 32, 128))
+        C_pad = T.sblock_alloc_buffer((128, (n + 31) // 32 * 32))
+        for i0, i1 in T.grid((n + 31) // 32 * 32, 128):
+            with T.sblock("B_pad"):
+                v0, v1 = T.axis.remap("SS", [i0, i1])
+                B_pad[v0, v1] = T.if_then_else(v0 < n, B[v0, v1], T.float32(0))
+        for i0, i1, i2 in T.grid(128, (n + 31) // 32 * 32, 128):
+            with T.sblock("C"):
+                i, j, k = T.axis.remap("SSR", [i0, i1, i2])
+                T.reads(A[i, k], B_pad[j, k])
+                T.writes(C_pad[i, j])
+                with T.init():
+                    C_pad[i, j] = T.float32(0)
+                C_pad[i, j] = C_pad[i, j] + A[i, k] * B_pad[j, k]
+        for i0, i1 in T.grid(128, n):
+            with T.sblock("C_pad"):
+                v0, v1 = T.axis.remap("SS", [i0, i1])
+                C[v0, v1] = C_pad[v0, v1]
+
+    sch = tvm.s_tir.Schedule(matmul_before, debug_mask="all")
+    C = sch.get_sblock("C")
+    sch.pad_einsum(C, [32, 32, 32])
+    assert_structural_equal_ignore_global_symbol(matmul_after, sch.mod["main"])
+    verify_trace_roundtrip(sch, mod=matmul_before)
+
+
+def test_pad_matmul_2():
+    @T.prim_func(s_tir=True)
+    def before(
+        a: T.handle,
+        b: T.handle,
+        m: T.handle,
+        d: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        n = T.int32()
+        A = T.match_buffer(a, (1, n, 4096))
+        B = T.match_buffer(b, (11008, 4096))
+        M = T.match_buffer(m, (1, n, 11008))
+        D = T.match_buffer(d, (1, n, 11008))
+        C = T.sblock_alloc_buffer((1, n, 11008))
+        for i0, i1, i2, k in T.grid(1, n, 11008, 4096):
+            with T.sblock("C"):
+                v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+                T.reads(A[v_i0, v_i1, v_k], B[v_i2, v_k])
+                T.writes(C[v_i0, v_i1, v_i2])
+                with T.init():
+                    C[v_i0, v_i1, v_i2] = T.float32(0)
+                C[v_i0, v_i1, v_i2] = C[v_i0, v_i1, v_i2] + A[v_i0, v_i1, v_k] * B[v_i2, v_k]
+        for ax0, ax1, ax2 in T.grid(1, n, 11008):
+            with T.sblock("D"):
+                v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+                D[v_ax0, v_ax1, v_ax2] = M[v_ax0, v_ax1, v_ax2] * C[v_ax0, v_ax1, v_ax2]
+
+    @T.prim_func(s_tir=True)
+    def after(a: T.handle, b: T.handle, m: T.handle, d: T.handle):
+        T.func_attr({"tirx.noalias": True})
+        n = T.int32()
+        A = T.match_buffer(a, (1, n, 4096))
+        B = T.match_buffer(b, (11008, 4096))
+        M = T.match_buffer(m, (1, n, 11008))
+        D = T.match_buffer(d, (1, n, 11008))
+        # with T.sblock("root"):
+        C = T.sblock_alloc_buffer((1, n, 11008))
+        A_pad = T.sblock_alloc_buffer((1, (n + 31) // 32 * 32, 4096))
+        C_pad = T.sblock_alloc_buffer((1, (n + 31) // 32 * 32, 11008))
+        for i0, i1, i2 in T.grid(1, (n + 31) // 32 * 32, 4096):
+            with T.sblock("A_pad"):
+                v0, v1, v2 = T.axis.remap("SSS", [i0, i1, i2])
+                A_pad[v0, v1, v2] = T.if_then_else(v1 < n, A[v0, v1, v2], T.float32(0))
+        for i0, i1, i2, k in T.grid(1, (n + 31) // 32 * 32, 11008, 4096):
+            with T.sblock("C"):
+                v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+                T.reads(A_pad[v_i0, v_i1, v_k], B[v_i2, v_k])
+                T.writes(C_pad[v_i0, v_i1, v_i2])
+                with T.init():
+                    C_pad[v_i0, v_i1, v_i2] = T.float32(0)
+                C_pad[v_i0, v_i1, v_i2] = (
+                    C_pad[v_i0, v_i1, v_i2] + A_pad[v_i0, v_i1, v_k] * B[v_i2, v_k]
+                )
+        for i0, i1, i2 in T.grid(1, n, 11008):
+            with T.sblock("C_pad"):
+                v0, v1, v2 = T.axis.remap("SSS", [i0, i1, i2])
+                C[v0, v1, v2] = C_pad[v0, v1, v2]
+        for ax0, ax1, ax2 in T.grid(1, n, 11008):
+            with T.sblock("D"):
+                v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+                D[v_ax0, v_ax1, v_ax2] = M[v_ax0, v_ax1, v_ax2] * C[v_ax0, v_ax1, v_ax2]
+
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    C = sch.get_sblock("C")
+    sch.pad_einsum(C, [1, 32, 32, 32])
+    assert_structural_equal_ignore_global_symbol(after, sch.mod["main"])
+    verify_trace_roundtrip(sch, mod=before)
+
+
+def test_pad_rms():
+    @T.prim_func(s_tir=True)
+    def before(
+        a: T.handle,
+        w: T.handle,
+        r: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        n = T.int32()
+        A = T.match_buffer(a, (1, n, 4096))
+        W = T.match_buffer(w, (4096,), "float32")
+        R = T.match_buffer(r, (1, n, 4096), "float32")
+        S = T.sblock_alloc_buffer((1, n), "float32")
+        for bsz, i, k in T.grid(1, n, 4096):
+            with T.sblock("S"):
+                v_bsz, v_i, v_k = T.axis.remap("SSR", [bsz, i, k])
+                T.reads(A[v_bsz, v_i, v_k])
+                T.writes(S[v_bsz, v_i])
+                with T.init():
+                    S[v_bsz, v_i] = T.float32(0)
+                S[v_bsz, v_i] = S[v_bsz, v_i] + A[v_bsz, v_i, v_k] * A[v_bsz, v_i, v_k]
+        for bsz, i, k in T.grid(1, n, 4096):
+            with T.sblock("R"):
+                v_bsz, v_i, v_k = T.axis.remap("SSS", [bsz, i, k])
+                R[v_bsz, v_i, v_k] = W[v_k] * (
+                    A[v_bsz, v_i, v_k]
+                    / T.sqrt(S[v_bsz, v_i] * T.float32(0.000244140625) + T.float32(1e-6))
+                )
+
+    @T.prim_func(s_tir=True)
+    def after(a: T.handle, w: T.handle, r: T.handle):
+        T.func_attr({"tirx.noalias": True})
+        n = T.int32()
+        A = T.match_buffer(a, (1, n, 4096))
+        W = T.match_buffer(w, (4096,), "float32")
+        R = T.match_buffer(r, (1, n, 4096))
+        S = T.sblock_alloc_buffer((1, n))
+        A_pad = T.sblock_alloc_buffer((1, (n + 31) // 32 * 32, 4096))
+        S_pad = T.sblock_alloc_buffer((1, (n + 31) // 32 * 32))
+        for i0, i1, i2 in T.grid(1, (n + 31) // 32 * 32, 4096):
+            with T.sblock("A_pad"):
+                v0, v1, v2 = T.axis.remap("SSS", [i0, i1, i2])
+                A_pad[v0, v1, v2] = T.if_then_else(v1 < n, A[v0, v1, v2], T.float32(0))
+        for bsz, i, k in T.grid(1, (n + 31) // 32 * 32, 4096):
+            with T.sblock("S"):
+                v_bsz, v_i, v_k = T.axis.remap("SSR", [bsz, i, k])
+                T.reads(A_pad[v_bsz, v_i, v_k])
+                T.writes(S_pad[v_bsz, v_i])
+                with T.init():
+                    S_pad[v_bsz, v_i] = T.float32(0)
+                S_pad[v_bsz, v_i] = (
+                    S_pad[v_bsz, v_i] + A_pad[v_bsz, v_i, v_k] * A_pad[v_bsz, v_i, v_k]
+                )
+        for i0, i1 in T.grid(1, n):
+            with T.sblock("S_pad"):
+                v0, v1 = T.axis.remap("SS", [i0, i1])
+                S[v0, v1] = S_pad[v0, v1]
+        for bsz, i, k in T.grid(1, n, 4096):
+            with T.sblock("R"):
+                v_bsz, v_i, v_k = T.axis.remap("SSS", [bsz, i, k])
+                R[v_bsz, v_i, v_k] = W[v_k] * (
+                    A[v_bsz, v_i, v_k]
+                    / T.sqrt(S[v_bsz, v_i] * T.float32(0.000244140625) + T.float32(1e-6))
+                )
+
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    C = sch.get_sblock("S")
+    sch.pad_einsum(C, [1, 32, 1])
+    assert_structural_equal_ignore_global_symbol(after, sch.mod["main"])
+    verify_trace_roundtrip(sch, mod=before)
+
+
+if __name__ == "__main__":
+    test_pad_matmul()
+    test_pad_matmul_2()
+    test_pad_rms()

tests/python/s_tir/schedule/test_tir_schedule_partition.py

@@ -0,0 +1,462 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401, F841
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import te, tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+from tvm.tirx.expr import IntImm
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def elementwise(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("B"):
+            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_symbolic(a: T.handle, b: T.handle, n: T.int32) -> None:
+    A = T.match_buffer(a, (128, 128, n))
+    B = T.match_buffer(b, (128, 128, n))
+    for i, j, k in T.grid(128, 128, n):
+        with T.sblock("B"):
+            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_with_anno(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j in T.grid(128, 128):
+        for k in T.serial(0, 128, annotations={"useless_annotation": True}):
+            with T.sblock("B"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                T.reads([A[vi, vj, vk]])
+                T.writes([B[vi, vj, vk]])
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_with_thread_binding(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j in T.grid(128, 128):
+        for k in T.thread_binding(0, 128, thread="threadIdx.x"):
+            with T.sblock("B"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                T.reads([A[vi, vj, vk]])
+                T.writes([B[vi, vj, vk]])
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_with_opaque_block(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("opaque"):
+            T.reads([A[i, j, k]])
+            T.writes([B[i, j, k]])
+            with T.sblock("B"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                T.reads([A[vi, vj, vk]])
+                T.writes([B[vi, vj, vk]])
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_partition_with_opaque_block(a: T.handle, b: T.handle) -> None:
+    B = T.match_buffer(b, [128, 128, 128])
+    A = T.match_buffer(a, [128, 128, 128])
+    with T.sblock("root"):
+        T.reads()
+        T.writes()
+        with T.sblock("opaque_i_common"):
+            T.reads()
+            T.writes()
+            with T.sblock("opaque_i0_partition"):
+                T.reads()
+                T.writes()
+                for i0, j, k in T.grid(112, 128, 128):
+                    with T.sblock("opaque_i0"):
+                        T.reads(A[i0, j, k])
+                        T.writes(B[i0, j, k])
+                        with T.sblock("B_i0"):
+                            vi, vj, vk = T.axis.remap("SSS", [i0, j, k])
+                            T.reads(A[0:112, 0:128, 0:128])
+                            T.writes(B[0:112, 0:128, 0:128])
+                            B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+            with T.sblock("opaque_i1_partition"):
+                T.reads()
+                T.writes()
+                for i1 in range(112, 128):
+                    for j, k in T.grid(128, 128):
+                        with T.sblock("opaque_i1"):
+                            T.reads(A[i1, j, k])
+                            T.writes(B[i1, j, k])
+                            with T.sblock("B_i1"):
+                                vi, vj, vk = T.axis.remap("SSS", [i1, j, k])
+                                T.reads(A[112:128, 0:128, 0:128])
+                                T.writes(B[112:128, 0:128, 0:128])
+                                B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+
+
+@T.prim_func(s_tir=True)
+def elementwise_loop_partition_case0(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128, 128])
+    B = T.match_buffer(b, [128, 128, 128])
+    with T.sblock("root"):
+        T.reads()
+        T.writes()
+        with T.sblock("B_i_common"):
+            T.reads()
+            T.writes()
+            with T.sblock("B_i0_partition"):
+                T.reads()
+                T.writes()
+                for i0 in range(2):
+                    with T.sblock("B_i0_j_common"):
+                        T.reads()
+                        T.writes()
+                        with T.sblock("B_i0_j0_partition"):
+                            T.reads()
+                            T.writes()
+                            for j0, k in T.grid(4, 128):
+                                with T.sblock("B_i0_j0"):
+                                    vi, vj, vk = T.axis.remap("SSS", [i0, j0, k])
+                                    T.reads(A[0:2, 0:4, 0:128])
+                                    T.writes(B[0:2, 0:4, 0:128])
+                                    B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+                        with T.sblock("B_i0_j1_partition"):
+                            T.reads()
+                            T.writes()
+                            for j1 in range(4, 36):
+                                for k in range(128):
+                                    with T.sblock("B_i0_j1"):
+                                        vi, vj, vk = T.axis.remap("SSS", [i0, j1, k])
+                                        T.reads(A[0:2, 4:36, 0:128])
+                                        T.writes(B[0:2, 4:36, 0:128])
+                                        B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+                        with T.sblock("B_i0_j2_partition"):
+                            T.reads()
+                            T.writes()
+                            for j2 in range(36, 128):
+                                for k in range(128):
+                                    with T.sblock("B_i0_j2"):
+                                        vi, vj, vk = T.axis.remap("SSS", [i0, j2, k])
+                                        T.reads(A[0:2, 36:128, 0:128])
+                                        T.writes(B[0:2, 36:128, 0:128])
+                                        B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+            with T.sblock("B_i1_partition"):
+                T.reads()
+                T.writes()
+                for i1 in range(2, 3):
+                    for j, k in T.grid(128, 128):
+                        with T.sblock("B_i1"):
+                            vi, vj, vk = T.axis.remap("SSS", [i1, j, k])
+                            T.reads(A[2, 0:128, 0:128])
+                            T.writes(B[2, 0:128, 0:128])
+                            B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+            with T.sblock("B_i2_partition"):
+                T.reads()
+                T.writes()
+                for i2 in range(3, 67):
+                    for j, k in T.grid(128, 128):
+                        with T.sblock("B_i2"):
+                            vi, vj, vk = T.axis.remap("SSS", [i2, j, k])
+                            T.reads(A[3:67, 0:128, 0:128])
+                            T.writes(B[3:67, 0:128, 0:128])
+                            B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+            with T.sblock("B_i3_partition"):
+                T.reads()
+                T.writes()
+                for i3 in range(67, 128):
+                    for j, k in T.grid(128, 128):
+                        with T.sblock("B_i3"):
+                            vi, vj, vk = T.axis.remap("SSS", [i3, j, k])
+                            T.reads(A[67:128, 0:128, 0:128])
+                            T.writes(B[67:128, 0:128, 0:128])
+                            B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+
+
+@T.prim_func(s_tir=True)
+def elementwise_loop_partition_case1(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128, 128])
+    B = T.match_buffer(b, [128, 128, 128])
+    with T.sblock("root"):
+        T.reads()
+        T.writes()
+        with T.sblock("B_i_common"):
+            T.reads()
+            T.writes()
+            with T.sblock("B_i0_partition"):
+                T.reads()
+                T.writes()
+                for i0, j, k in T.grid(63, 128, 128):
+                    with T.sblock("B_i0"):
+                        vi, vj, vk = T.axis.remap("SSS", [i0, j, k])
+                        T.reads(A[0:63, 0:128, 0:128])
+                        T.writes(B[0:63, 0:128, 0:128])
+                        B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+            with T.sblock("B_i1_partition"):
+                T.reads()
+                T.writes()
+                for i1 in range(63, 64):
+                    for j in range(128):
+                        with T.sblock("B_i1_k_common"):
+                            T.reads()
+                            T.writes()
+                            with T.sblock("B_i1_k0_partition"):
+                                T.reads()
+                                T.writes()
+                                for k0 in range(1):
+                                    with T.sblock("B_i1_k0"):
+                                        vi, vj, vk = T.axis.remap("SSS", [i1, j, k0])
+                                        T.reads(A[63, 0:128, 0])
+                                        T.writes(B[63, 0:128, 0])
+                                        B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+                            with T.sblock("B_i1_k1_partition"):
+                                T.reads()
+                                T.writes()
+                                for k1 in range(1, 65):
+                                    with T.sblock("B_i1_k1"):
+                                        vi, vj, vk = T.axis.remap("SSS", [i1, j, k1])
+                                        T.reads(A[63, 0:128, 1:65])
+                                        T.writes(B[63, 0:128, 1:65])
+                                        B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+                            with T.sblock("B_i1_k2_partition"):
+                                T.reads()
+                                T.writes()
+                                for k2 in range(65, 128):
+                                    with T.sblock("B_i1_k2"):
+                                        vi, vj, vk = T.axis.remap("SSS", [i1, j, k2])
+                                        T.reads(A[63, 0:128, 65:128])
+                                        T.writes(B[63, 0:128, 65:128])
+                                        B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+            with T.sblock("B_i2_partition"):
+                T.reads()
+                T.writes()
+                for i2 in range(64, 128):
+                    for j, k in T.grid(128, 128):
+                        with T.sblock("B_i2"):
+                            vi, vj, vk = T.axis.remap("SSS", [i2, j, k])
+                            T.reads(A[64:128, 0:128, 0:128])
+                            T.writes(B[64:128, 0:128, 0:128])
+                            B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2)
+
+
+@T.prim_func(s_tir=True)
+def opaque_access(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [16, 16], "float32")
+    B = T.match_buffer(b, [16, 16], "float32")
+    for i, j in T.grid(16, 16):
+        with T.sblock("A"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([])
+            T.writes([A[0:16, 0:16]])
+            A[vi, vj] = 1
+    for i, j in T.grid(16, 16):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([])
+            T.writes([B[0:16, 0:16]])
+            T.evaluate(T.tvm_fill_fragment(B.data, 16, 16, 16, 0, vi * 16 + vj, dtype="handle"))
+
+
+@T.prim_func(s_tir=True)
+def opaque_access_loop_partition(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16))
+    B = T.match_buffer(b, (16, 16))
+    for i in range(16):
+        with T.sblock("A_j_common"):
+            T.reads()
+            T.writes()
+            with T.sblock("A_j0_partition"):
+                T.reads()
+                T.writes()
+                for j0 in range(12):
+                    with T.sblock("A_j0"):
+                        vi, vj = T.axis.remap("SS", [i, j0])
+                        T.reads()
+                        T.writes(A[0:16, 0:12])
+                        A[vi, vj] = T.float32(1)
+            with T.sblock("A_j1_partition"):
+                T.reads()
+                T.writes()
+                for j1 in range(12, 16):
+                    with T.sblock("A_j1"):
+                        vi, vj = T.axis.remap("SS", [i, j1])
+                        T.reads()
+                        T.writes(A[0:16, 12:16])
+                        A[vi, vj] = T.float32(1)
+    for i in range(16):
+        with T.sblock("B_j_common"):
+            T.reads()
+            T.writes()
+            with T.sblock("B_j0_partition"):
+                T.reads()
+                T.writes()
+                for j0 in range(12):
+                    with T.sblock("B_j0"):
+                        vi, vj = T.axis.remap("SS", [i, j0])
+                        T.reads()
+                        T.writes(B[0:16, 0:16])
+                        T.tvm_fill_fragment(B.data, 16, 16, 16, 0, vi * 16 + vj)
+            with T.sblock("B_j1_partition"):
+                T.reads()
+                T.writes()
+                for j1 in range(12, 16):
+                    with T.sblock("B_j1"):
+                        vi, vj = T.axis.remap("SS", [i, j1])
+                        T.reads()
+                        T.writes(B[0:16, 0:16])
+                        T.tvm_fill_fragment(B.data, 16, 16, 16, 0, vi * 16 + vj)
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+
+def test_loop_partition():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    sch.loop_partition(i, factors=[2, 1, 64])
+
+    block_b_partition = sch.get_sblock("B_i0")
+    i, j, k = sch.get_loops(block_b_partition)
+    loops = sch.loop_partition(j, factors=[4, 32])
+
+    assert_structural_equal_ignore_global_symbol(elementwise_loop_partition_case0, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+
+def test_partition_with_inferred_factor():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    sch.loop_partition(i, factors=[None, 1, 64])
+
+    block_b_partition = sch.get_sblock("B_i1")
+    i, j, k = sch.get_loops(block_b_partition)
+    sch.loop_partition(k, factors=[1, 64, None])
+
+    assert_structural_equal_ignore_global_symbol(elementwise_loop_partition_case1, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+
+def test_partition_with_opaque_block():
+    sch = tvm.s_tir.Schedule(elementwise_with_opaque_block, debug_mask="all")
+    block_opaque = sch.get_sblock("opaque")
+    i, _, _ = sch.get_loops(block_opaque)
+    sch.loop_partition(i, factors=[None, 16])
+    assert_structural_equal_ignore_global_symbol(
+        elementwise_partition_with_opaque_block, sch.mod["main"]
+    )
+    verify_trace_roundtrip(sch=sch, mod=elementwise_with_opaque_block)
+
+
+def test_partition_with_opaque_access():
+    sch = tvm.s_tir.Schedule(opaque_access, debug_mask="all")
+    block_a = sch.get_sblock("A")
+    _, j = sch.get_loops(block_a)
+    sch.loop_partition(j, factors=[None, 4])
+    block_b = sch.get_sblock("B")
+    _, j = sch.get_loops(block_b)
+    sch.loop_partition(j, factors=[None, 4])
+    assert_structural_equal_ignore_global_symbol(opaque_access_loop_partition, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=opaque_access)
+
+
+def test_partition_int64_extent_with_mixed_factors():
+    def _create_prim_func():
+        m = te.const(384, "int64")
+        A = te.placeholder((m,), name="A", dtype="float32")
+        B = te.compute((m,), lambda i: A[i] + 1, name="B")
+        return te.create_prim_func([A, B])
+
+    mod = _create_prim_func()
+    sch = tvm.s_tir.Schedule(mod, debug_mask="all")
+    (i,) = sch.get_loops(sch.get_sblock("B"))
+    sch.loop_partition(
+        i,
+        factors=[
+            te.const(1, "int64"),
+            te.const(51, "int32"),
+        ],
+    )
+
+
+def test_partition_fail_symbolic():
+    sch = tvm.s_tir.Schedule(elementwise_symbolic, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, _, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.loop_partition(k, factors=[10, None])
+
+
+def test_partition_fail_out_of_bound():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.loop_partition(i, factors=[1000, 2, 3])
+
+
+def test_partition_with_non_positive_factors():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.loop_partition(i, factors=[-2, -64])
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.loop_partition(j, factors=[0, None])
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.loop_partition(k, factors=[None, -16])
+
+
+def test_partition_fail_with_annotation():
+    sch = tvm.s_tir.Schedule(elementwise_with_anno, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.loop_partition(k, factors=[None, 10])
+
+
+def test_partition_fail_with_thread_binding():
+    sch = tvm.s_tir.Schedule(elementwise_with_thread_binding, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.loop_partition(k, factors=[None, 10])
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py

@@ -0,0 +1,241 @@
+# 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: E501, F401
+
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+import sys
+
+import pytest
+
+import tvm
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# fmt: off
+# pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks,not-callable
+
+@T.prim_func(s_tir=True)
+def cuda_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:  # pylint: disable=undefined-loop-variable
+    A = T.match_buffer(a, [2048, 2048], "float32")
+    B = T.match_buffer(b, [2048, 2048], "float32")
+    C = T.match_buffer(c, [2048, 2048], "float32")
+    for by in T.thread_binding(0, 32, thread = "blockIdx.y"):
+        for bx in T.thread_binding(0, 32, thread = "blockIdx.x"):
+            for vy in T.thread_binding(0, 2, thread = "vthread.y"):
+                for vx in T.thread_binding(0, 2, thread = "vthread.x"):
+                    for ty in T.thread_binding(0, 8, thread = "threadIdx.y"):
+                        for tx in T.thread_binding(0, 8, thread = "threadIdx.x"):
+                            for k0 in T.serial(0, 256):
+                                for k1 in T.unroll(0, 8):
+                                    for _, i, j in T.grid(1, 4, 4):
+                                        with T.sblock("C"):
+                                            vi = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
+                                            vj = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
+                                            vk = T.axis.R(2048, k0 * 8 + k1)
+                                            T.reads([C[vi, vj], A[vi, vk], B[vk, vj]])
+                                            T.writes([C[vi, vj]])
+                                            with T.init():
+                                                C[vi, vj] = 0.0
+                                            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
+
+
+@T.prim_func(s_tir=True)
+def cuda_matmul_read_at_a(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [2048, 2048], dtype="float32")
+    B = T.match_buffer(b, [2048, 2048], dtype="float32")
+    C = T.match_buffer(c, [2048, 2048], dtype="float32")
+    A_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
+    for by in T.thread_binding(0, 32, thread="blockIdx.y"):
+        for bx in T.thread_binding(0, 32, thread="blockIdx.x"):
+            for vy in T.thread_binding(0, 2, thread="vthread.y"):
+                for vx in T.thread_binding(0, 2, thread="vthread.x"):
+                    for ty in T.thread_binding(0, 8, thread="threadIdx.y"):
+                        for tx in T.thread_binding(0, 8, thread="threadIdx.x"):
+                            for k0 in T.serial(0, 256):
+                                with T.sblock("A_shared"):
+                                    v0 = T.axis.S(32, by)
+                                    v1 = T.axis.S(256, k0)
+                                    T.reads([A[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
+                                    T.writes([A_shared[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
+                                    T.sblock_attr({"auto_copy": True})
+                                    for ax0, ax1 in T.grid(64, 8):
+                                        A_shared[v0 * 64 + ax0, v1 * 8 + ax1] = A[v0 * 64 + ax0, v1 * 8 + ax1]
+                                for k1 in T.unroll(0, 8):
+                                    for v_, i, j in T.grid(1, 4, 4):
+                                        with T.sblock("C"):
+                                            vi = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
+                                            vj = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
+                                            vk = T.axis.R(2048, k0 * 8 + k1)
+                                            T.reads([C[vi, vj], A_shared[vi, vk], B[vk, vj]])
+                                            T.writes([C[vi, vj]])
+                                            with T.init():
+                                                C[vi, vj] = T.float32(0)
+                                            C[vi, vj] = C[vi, vj] + A_shared[vi, vk] * B[vk, vj]
+
+
+@T.prim_func(s_tir=True)
+def cuda_matmul_read_at_ab(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [2048, 2048], dtype="float32")
+    B = T.match_buffer(b, [2048, 2048], dtype="float32")
+    C = T.match_buffer(c, [2048, 2048], dtype="float32")
+    A_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
+    B_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
+    for by in T.thread_binding(0, 32, thread="blockIdx.y"):
+        for bx in T.thread_binding(0, 32, thread="blockIdx.x"):
+            for vy in T.thread_binding(0, 2, thread="vthread.y"):
+                for vx in T.thread_binding(0, 2, thread="vthread.x"):
+                    for ty in T.thread_binding(0, 8, thread="threadIdx.y"):
+                        for tx in T.thread_binding(0, 8, thread="threadIdx.x"):
+                            for k0 in T.serial(0, 256):
+                                with T.sblock("A_shared"):
+                                    v0 = T.axis.S(32, by)
+                                    v1 = T.axis.S(256, k0)
+                                    T.reads([A[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
+                                    T.writes([A_shared[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
+                                    T.sblock_attr({"auto_copy": True})
+                                    for ax0, ax1 in T.grid(64, 8):
+                                        A_shared[v0 * 64 + ax0, v1 * 8 + ax1] = A[v0 * 64 + ax0, v1 * 8 + ax1]
+                                with T.sblock("B_shared"):
+                                    v0 = T.axis.S(256, k0)
+                                    v1 = T.axis.S(32, bx)
+                                    T.reads([B[v0 * 8 : v0 * 8 + 8, v1 * 64 : v1 * 64 + 64]])
+                                    T.writes([B_shared[v0 * 8 : v0 * 8 + 8, v1 * 64 : v1 * 64 + 64]])
+                                    T.sblock_attr({"auto_copy": True})
+                                    for ax0, ax1 in T.grid(8, 64):
+                                        B_shared[v0 * 8 + ax0, v1 * 64 + ax1] = B[v0 * 8 + ax0, v1 * 64 + ax1]
+                                for k1 in T.unroll(0, 8):
+                                    for v_, i, j in T.grid(1, 4, 4):
+                                        with T.sblock("C"):
+                                            vi = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
+                                            vj = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
+                                            vk = T.axis.R(2048, k0 * 8 + k1)
+                                            T.reads([C[vi, vj], A_shared[vi, vk], B_shared[vk, vj]])
+                                            T.writes([C[vi, vj]])
+                                            with T.init():
+                                                C[vi, vj] = T.float32(0)
+                                            C[vi, vj] = C[vi, vj] + A_shared[vi, vk] * B_shared[vk, vj]
+
+@T.prim_func(s_tir=True)
+def cuda_matmul_write_at_c(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [2048, 2048], dtype="float32")
+    B = T.match_buffer(b, [2048, 2048], dtype="float32")
+    C = T.match_buffer(c, [2048, 2048], dtype="float32")
+    A_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
+    B_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
+    C_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
+    for by in T.thread_binding(0, 32, thread="blockIdx.y"):
+        for bx in T.thread_binding(0, 32, thread="blockIdx.x"):
+            for vy in T.thread_binding(0, 2, thread="vthread.y"):
+                for vx in T.thread_binding(0, 2, thread="vthread.x"):
+                    for ty in T.thread_binding(0, 8, thread="threadIdx.y"):
+                        for tx in T.thread_binding(0, 8, thread="threadIdx.x"):
+                            for k0 in T.serial(0, 256):
+                                with T.sblock("A_shared"):
+                                    v0 = T.axis.S(32, by)
+                                    v1 = T.axis.S(256, k0)
+                                    T.reads([A[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
+                                    T.writes([A_shared[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
+                                    T.sblock_attr({"auto_copy": True})
+                                    for ax0, ax1 in T.grid(64, 8):
+                                        A_shared[v0 * 64 + ax0, v1 * 8 + ax1] = A[v0 * 64 + ax0, v1 * 8 + ax1]
+                                with T.sblock("B_shared"):
+                                    v0 = T.axis.S(256, k0)
+                                    v1 = T.axis.S(32, bx)
+                                    T.reads([B[v0 * 8 : v0 * 8 + 8, v1 * 64 : v1 * 64 + 64]])
+                                    T.writes([B_shared[v0 * 8 : v0 * 8 + 8, v1 * 64 : v1 * 64 + 64]])
+                                    T.sblock_attr({"auto_copy": True})
+                                    for ax0, ax1 in T.grid(8, 64):
+                                        B_shared[v0 * 8 + ax0, v1 * 64 + ax1] = B[v0 * 8 + ax0, v1 * 64 + ax1]
+                                for k1 in T.unroll(0, 8):
+                                    for v_, i, j in T.grid(1, 4, 4):
+                                        with T.sblock("C"):
+                                            vi = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
+                                            vj = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
+                                            vk = T.axis.R(2048, k0 * 8 + k1)
+                                            T.reads([C_shared[vi, vj], A_shared[vi, vk], B_shared[vk, vj]])
+                                            T.writes([C_shared[vi, vj]])
+                                            with T.init():
+                                                C_shared[vi, vj] = T.float32(0)
+                                            C_shared[vi, vj] = C_shared[vi, vj] + A_shared[vi, vk] * B_shared[vk, vj]
+                            with T.sblock("C_shared"):
+                                v0 = T.axis.S(32, by)
+                                v1 = T.axis.S(32, bx)
+                                T.reads([C_shared[v0 * 64 : v0 * 64 + 64, v1 * 64 : v1 * 64 + 64]])
+                                T.writes([C[v0 * 64 : v0 * 64 + 64, v1 * 64 : v1 * 64 + 64]])
+                                T.sblock_attr({"auto_copy": True})
+                                for ax0, ax1 in T.grid(64, 64):
+                                    C[v0 * 64 + ax0, v1 * 64 + ax1] = C_shared[v0 * 64 + ax0, v1 * 64 + ax1]
+
+
+# pylint: enable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks,not-callable
+# fmt: on
+
+
+def test_read_at_global_to_shared_a():
+    sch = tvm.s_tir.Schedule(cuda_matmul, debug_mask="all")
+    block = sch.get_sblock("C")
+    # pylint: disable=invalid-name
+    _by, _bx, _vy, _vx, _ty, _tx, k0, _k1, _, _i, _j = sch.get_loops(block)
+    # pylint: enable=invalid-name
+    sch.read_at(k0, block, 1, "shared")
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], cuda_matmul_read_at_a)
+    verify_trace_roundtrip(sch, cuda_matmul)
+
+
+def test_read_at_global_to_shared_ab():
+    sch = tvm.s_tir.Schedule(cuda_matmul_read_at_a, debug_mask="all")
+    block = sch.get_sblock("C")
+    # pylint: disable=invalid-name
+    _by, _bx, _vy, _vx, _ty, _tx, k0, _k1, _, _i, _j = sch.get_loops(block)
+    # pylint: enable=invalid-name
+    sch.read_at(k0, block, 2, "shared")
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], cuda_matmul_read_at_ab)
+    verify_trace_roundtrip(sch, cuda_matmul_read_at_a)
+
+
+def test_read_at_local_to_shared_c():
+    sch = tvm.s_tir.Schedule(cuda_matmul_read_at_ab, debug_mask="all")
+    block = sch.get_sblock("C")
+    # pylint: disable=invalid-name
+    _by, _bx, _vy, _vx, _ty, tx, _k0, _k1, _, _i, _j = sch.get_loops(block)
+    # pylint: enable=invalid-name
+    sch.write_at(tx, block, 0, "shared")
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], cuda_matmul_write_at_c)
+    verify_trace_roundtrip(sch, cuda_matmul_read_at_ab)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_reduction.py

@@ -0,0 +1,395 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import sys
+
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+
+@T.prim_func(s_tir=True)
+def rowsum_blockized(a: T.handle, b: T.handle) -> None:
+    B = T.match_buffer(b, [32, 4])
+    A = T.match_buffer(a, [32, 4, 128])
+    for i0, i2_0 in T.grid(32, 16):
+        with T.sblock("blockized_B"):
+            io, ko = T.axis.remap("SR", [i0, i2_0])
+            with T.init():
+                for i1 in T.serial(0, 4):
+                    with T.sblock("B_init"):
+                        ii_init = T.axis.S(4, i1)
+                        B[io, ii_init] = 0.0
+            for i1_1, i2_1 in T.grid(4, 8):
+                with T.sblock("B"):
+                    ii = T.axis.S(4, i1_1)
+                    k = T.axis.R(128, ko * 8 + i2_1)
+                    B[io, ii] = B[io, ii] + A[io, ii, k]
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def matmul_decompose0(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("init"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = 0.0
+
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def matmul_decompose1(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [32, 4, 128], elem_offset=0, align=64, offset_factor=1)
+    B = T.match_buffer(b, [32, 4], elem_offset=0, align=64, offset_factor=1)
+
+    for i0 in T.serial(0, 32):
+        with T.sblock("blockized_B_init"):
+            io = T.axis.S(32, i0)
+            for i1 in T.serial(0, 4):
+                with T.sblock("B_init"):
+                    ii = T.axis.S(4, i1)
+                    B[io, ii] = T.float32(0)
+    for i0, i2_o in T.grid(32, 16):
+        with T.sblock("blockized_B_update"):
+            io, ko = T.axis.remap("SR", [i0, i2_o])
+            for i1, i2_i in T.grid(4, 8):
+                with T.sblock("B"):
+                    ii = T.axis.S(4, i1)
+                    k = T.axis.R(128, ko * 8 + i2_i)
+                    B[io, ii] = B[io, ii] + A[io, ii, k]
+
+
+@T.prim_func(s_tir=True)
+def matmul_decompose2(a: T.handle, b: T.handle, c: T.handle) -> None:
+    C = T.match_buffer(c, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    B = T.match_buffer(b, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    A = T.match_buffer(a, [128, 128], elem_offset=0, align=64, offset_factor=1)
+
+    for i0, i1 in T.grid(128, 128):
+        with T.sblock("update_init"):
+            vi_init, vj_init = T.axis.remap("SS", [i0, i1])
+            C[vi_init, vj_init] = T.float32(0)
+        for i2 in T.serial(0, 128):
+            with T.sblock("update_update"):
+                vi, vj, vk = T.axis.remap("SSR", [i0, i1, i2])
+                C[vi, vj] = C[vi, vj] + (A[vi, vk] * B[vj, vk])
+
+
+@T.prim_func(s_tir=True)
+def matmul_decompose_fail3(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+
+    for i, k, j in T.grid(128, 128, 128):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def matmul_decompose4(a: T.handle, b: T.handle, c: T.handle) -> None:
+    C = T.match_buffer(c, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    B = T.match_buffer(b, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    A = T.match_buffer(a, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    # body
+    with T.sblock("root"):
+        T.reads([])
+        T.writes([])
+        for i0_0 in T.serial(0, 16):
+            for i0_1_init, i1_init in T.grid(8, 128):
+                with T.sblock("update_init"):
+                    vi_init = T.axis.S(128, i0_0 * 8 + i0_1_init)
+                    vj_init = T.axis.S(128, i1_init)
+                    C[vi_init, vj_init] = T.float32(0)
+            for i0_1, i1, i2_0, i2_1 in T.grid(8, 128, 19, 7):
+                with T.sblock("update_update"):
+                    T.where(((i2_0 * 7) + i2_1) < 128)
+                    vi = T.axis.S(128, i0_0 * 8 + i0_1)
+                    vj = T.axis.S(128, i1)
+                    vk = T.axis.R(128, i2_0 * 7 + i2_1)
+                    C[vi, vj] = C[vi, vj] + (A[vi, vk] * B[vj, vk])
+
+
+@T.prim_func(s_tir=True)
+def matmul_with_annotation(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("update"):
+            T.sblock_attr({"test_annotation": 1})
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def matmul_decompose_with_annotation(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("init"):
+            T.sblock_attr({"test_annotation": 1})
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = 0.0
+
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("update"):
+            T.sblock_attr({"test_annotation": 1})
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def colsum_with_vectorization(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 32], dtype="float32")
+    B = T.match_buffer(b, [32], dtype="float32")
+    for k in T.serial(0, 128):
+        for i in T.vectorized(0, 32):
+            with T.sblock("B"):
+                vk, vi = T.axis.remap("RS", [k, i])
+                with T.init():
+                    B[vi] = T.float32(0)
+                B[vi] = B[vi] + A[vk, vi]
+
+
+@T.prim_func(s_tir=True)
+def colsum_decompose_with_vectorization(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 32], dtype="float32")
+    B = T.match_buffer(b, [32], dtype="float32")
+    for i in T.vectorized(0, 32):
+        with T.sblock("B_init"):
+            vi = T.axis.S(32, i)
+            B[vi] = T.float32(0)
+    for k in T.serial(0, 128):
+        for i in T.vectorized(0, 32):
+            with T.sblock("B"):
+                vk, vi = T.axis.remap("RS", [k, i])
+                B[vi] = B[vi] + A[vk, vi]
+
+
+# pylint: enable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+use_block_name = tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True})
+
+
+def test_reduction_decompose0(use_block_name):
+    s = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    C = "update" if use_block_name else s.get_sblock("update")
+    i, j, k = s.get_loops(C)
+    s.decompose_reduction(C, i)
+    assert_structural_equal_ignore_global_symbol(matmul_decompose0, s.mod["main"])
+    verify_trace_roundtrip(s, mod=matmul)
+
+
+def test_reduction_decompose1(use_block_name):
+    s = tvm.s_tir.Schedule(rowsum_blockized, debug_mask="all")
+    blockized_B = "blockized_B" if use_block_name else s.get_sblock("blockized_B")
+    io, ko = s.get_loops(blockized_B)
+    s.decompose_reduction(blockized_B, io)
+    assert_structural_equal_ignore_global_symbol(matmul_decompose1, s.mod["main"])
+    verify_trace_roundtrip(s, mod=rowsum_blockized)
+
+
+def test_reduction_decompose2():
+    s = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    C = s.get_sblock("update")
+    i, j, k = s.get_loops(C)
+    s.decompose_reduction(C, k)
+    assert_structural_equal_ignore_global_symbol(matmul_decompose2, s.mod["main"])
+    verify_trace_roundtrip(s, mod=matmul)
+
+
+def test_reduction_decompose3():
+    s = tvm.s_tir.Schedule(matmul_decompose_fail3, debug_mask="all")
+    C = s.get_sblock("update")
+    i, j, k = s.get_loops(C)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.decompose_reduction(C, k)
+
+
+def test_reduction_decompose4():
+    s = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    C = s.get_sblock("update")
+    i, j, k = s.get_loops(C)
+    io, ii = s.split(i, factors=[16, 8])
+    ko, ki = s.split(k, factors=[19, 7])
+    s.decompose_reduction(C, ii)
+    assert_structural_equal_ignore_global_symbol(matmul_decompose4, s.mod["main"])
+    verify_trace_roundtrip(s, mod=matmul)
+
+
+def test_reduction_decompose_with_annotation():
+    s = tvm.s_tir.Schedule(matmul_with_annotation, debug_mask="all")
+    C = s.get_sblock("update")
+    i, j, k = s.get_loops(C)
+    s.decompose_reduction(C, i)
+    assert_structural_equal_ignore_global_symbol(matmul_decompose_with_annotation, s.mod["main"])
+    verify_trace_roundtrip(s, mod=matmul_with_annotation)
+
+
+def test_reduction_decompose_with_different_for_kind():
+    s = tvm.s_tir.Schedule(colsum_with_vectorization, debug_mask="all")
+    B = s.get_sblock("B")
+    k, _ = s.get_loops(B)
+    B_init = s.decompose_reduction(B, k)
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], colsum_decompose_with_vectorization)
+    assert s.get(B).same_as(s.get(s.get_sblock("B_update")))
+    assert s.get(B_init).same_as(s.get(s.get_sblock("B_init")))
+    verify_trace_roundtrip(s, mod=colsum_with_vectorization)
+
+
+def test_decompose_reduction_ref_hash_check():
+    mod = tvm.IRModule.from_expr(matmul.with_attr("global_symbol", "main"))
+    mod_bak = mod
+    hash_before = tvm_ffi.structural_hash(mod_bak)
+    s = tvm.s_tir.Schedule(mod["main"], debug_mask="all")
+    C = s.get_sblock("update")
+    i, j, k = s.get_loops(C)
+    s.decompose_reduction(C, k)
+    hash_after = tvm_ffi.structural_hash(mod_bak)
+    assert hash_before == hash_after
+
+
+def test_decompose_reduction_nested_block():
+    @T.prim_func(s_tir=True)
+    def nested_block(A: T.Buffer((1, 64), "float32"), B: T.Buffer((1,), "float32")):
+        for i, ko in T.grid(1, 2):
+            with T.sblock("outer"):
+                vi, vko = T.axis.remap("SR", [i, ko])
+                C = T.sblock_alloc_buffer((32,), dtype="float32")
+                with T.init():
+                    B[vi] = T.float32(0)
+                for ki in T.serial(32):
+                    with T.sblock("inner_1"):
+                        vki = T.axis.remap("S", [ki])
+                        C[vki] = A[vi, vko * 32 + vki]
+                for ki in T.serial(32):
+                    with T.sblock("inner_2"):
+                        vki = T.axis.remap("R", [ki])
+                        B[vi] += C[vki]
+
+    @T.prim_func(s_tir=True)
+    def decomposed_nested_block(A: T.Buffer((1, 64), "float32"), B: T.Buffer((1,), "float32")):
+        for i in range(1):
+            with T.sblock("outer_init"):
+                vi = T.axis.spatial(1, i)
+                T.reads()
+                T.writes(B[vi])
+                B[vi] = T.float32(0)
+            for ko in range(2):
+                with T.sblock("outer_update"):
+                    vi, vko = T.axis.remap("SR", [i, ko])
+                    T.reads(B[vi], A[vi, vko * 32 : vko * 32 + 32])
+                    T.writes(B[vi])
+                    C = T.sblock_alloc_buffer((32,))
+                    for ki in range(32):
+                        with T.sblock("inner_1"):
+                            vki = T.axis.spatial(32, ki)
+                            T.reads(A[vi, vko * 32 + vki])
+                            T.writes(C[vki])
+                            C[vki] = A[vi, vko * 32 + vki]
+                    for ki in range(32):
+                        with T.sblock("inner_2"):
+                            vki = T.axis.reduce(32, ki)
+                            T.reads(B[vi], C[vki])
+                            T.writes(B[vi])
+                            B[vi] = B[vi] + C[vki]
+
+    sch = tvm.s_tir.Schedule(nested_block, debug_mask="all")
+    outer = sch.get_sblock("outer")
+    i, ko = sch.get_loops(outer)
+    sch.decompose_reduction(outer, ko)
+
+    assert_structural_equal_ignore_global_symbol(decomposed_nested_block, sch.mod["main"])
+    verify_trace_roundtrip(sch, mod=nested_block)
+
+
+def test_decompose_reduction_with_thread_binding():
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((32, 16), "float32"), B: T.Buffer((32,), "float32")):
+            for t in T.thread_binding(0, 32, thread="threadIdx.x"):
+                for r in T.serial(16):
+                    with T.sblock("B"):
+                        vi, vr = T.axis.remap("SR", [t, r])
+                        with T.init():
+                            B[vi] = T.float32(0)
+                        B[vi] += A[vi, vr]
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((32, 16), "float32"), B: T.Buffer((32,), "float32")):
+            for t_init in T.thread_binding(0, 32, thread="threadIdx.x"):
+                with T.sblock("B_init"):
+                    vi = T.axis.remap("S", [t_init])
+                    B[vi] = T.float32(0)
+            for t in T.thread_binding(0, 32, thread="threadIdx.x"):
+                for r in T.serial(16):
+                    with T.sblock("B"):
+                        vi, vr = T.axis.remap("SR", [t, r])
+                        B[vi] += A[vi, vr]
+
+    sch = tvm.s_tir.Schedule(Before)
+    t, _ = sch.get_loops("B")
+    sch.decompose_reduction("B", t)
+    After = sch.mod
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_reindex.py

@@ -0,0 +1,355 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.schedule import ScheduleError
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+
+@T.prim_func(s_tir=True)
+def transpose_elementwise(
+    A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")
+) -> None:
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vj, vi] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def transpose_elementwise_reindex_read(
+    A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")
+) -> None:
+    A_reindex = T.sblock_alloc_buffer((128, 128), "float32")
+    for i, j in T.grid(128, 128):
+        with T.sblock("A_reindex"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            A_reindex[vi, vj] = A[vj, vi]
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A_reindex[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def conv2d_nhwc(
+    Input: T.Buffer((1, 224, 224, 3), "float32"),
+    Weight: T.Buffer((7, 7, 3, 64), "float32"),
+    Conv2d_nhwc: T.Buffer((1, 112, 112, 64), "float32"),
+) -> None:
+    PadInput = T.sblock_alloc_buffer([1, 230, 230, 3], dtype="float32")
+    for i0, i1, i2, i3 in T.grid(1, 230, 230, 3):
+        with T.sblock("PadInput"):
+            i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+            PadInput[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(
+                ((((i1_1 >= 3) and (i1_1 < 227)) and (i2_1 >= 3)) and (i2_1 < 227)),
+                Input[i0_1, (i1_1 - 3), (i2_1 - 3), i3_1],
+                T.float32(0),
+                dtype="float32",
+            )
+    for i0, i1, i2, i3, i4, i5, i6 in T.grid(1, 112, 112, 64, 7, 7, 3):
+        with T.sblock("conv2d_nhwc"):
+            n, h, w, co, rh, rw, rc = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
+            with T.init():
+                Conv2d_nhwc[n, h, w, co] = T.float32(0)
+            Conv2d_nhwc[n, h, w, co] = Conv2d_nhwc[n, h, w, co] + (
+                PadInput[n, ((h * 2) + rh), ((w * 2) + rw), ((T.floordiv(co, 64) * 3) + rc)]
+                * Weight[rh, rw, rc, co]
+            )
+
+
+@T.prim_func(s_tir=True)
+def conv2d_nhwc_reindex_data(
+    Input: T.Buffer((1, 224, 224, 3), "float32"),
+    Weight: T.Buffer((7, 7, 3, 64), "float32"),
+    Conv2d_nhwc: T.Buffer((1, 112, 112, 64), "float32"),
+) -> None:
+    PadInput = T.sblock_alloc_buffer([1, 230, 230, 3], dtype="float32")
+    ReindexInput = T.sblock_alloc_buffer([1, 112, 112, 7, 7, 3], dtype="float32")
+    for i0, i1, i2, i3 in T.grid(1, 230, 230, 3):
+        with T.sblock("PadInput"):
+            i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+            PadInput[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(
+                ((((i1_1 >= 3) and (i1_1 < 227)) and (i2_1 >= 3)) and (i2_1 < 227)),
+                Input[i0_1, (i1_1 - 3), (i2_1 - 3), i3_1],
+                T.float32(0),
+                dtype="float32",
+            )
+    for i0, i1, i2, i3, i4, i5 in T.grid(1, 112, 112, 7, 7, 3):
+        with T.sblock("ReindexInput"):
+            n, h, w, rh, rw, rc = T.axis.remap("SSSSSS", [i0, i1, i2, i3, i4, i5])
+            ReindexInput[n, h, w, rh, rw, rc] = PadInput[n, ((h * 2) + rh), ((w * 2) + rw), rc]
+    for i0, i1, i2, i3, i4, i5, i6 in T.grid(1, 112, 112, 64, 7, 7, 3):
+        with T.sblock("conv2d_nhwc"):
+            n, h, w, co, rh, rw, rc = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
+            with T.init():
+                Conv2d_nhwc[n, h, w, co] = T.float32(0)
+            Conv2d_nhwc[n, h, w, co] = Conv2d_nhwc[n, h, w, co] + (
+                ReindexInput[n, h, w, rh, rw, rc] * Weight[rh, rw, rc, co]
+            )
+
+
+@T.prim_func(s_tir=True)
+def conv2d_nhwc_reindex_weight(
+    var_inputs: T.handle, var_weight: T.handle, var_conv2d_nhwc: T.handle
+) -> None:
+    inputs = T.match_buffer(var_inputs, [1, 224, 224, 3], dtype="float32")
+    weight = T.match_buffer(var_weight, [7, 7, 3, 64], dtype="float32")
+    conv2d_nhwc = T.match_buffer(var_conv2d_nhwc, [1, 112, 112, 64], dtype="float32")
+    PadInput = T.sblock_alloc_buffer([1, 230, 230, 3], dtype="float32")
+    weight_reindex = T.sblock_alloc_buffer([64, 7, 7, 3], dtype="float32")
+    for i0, i1, i2, i3 in T.grid(1, 230, 230, 3):
+        with T.sblock("PadInput"):
+            i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+            T.reads(inputs[i0_1, i1_1 - 3, i2_1 - 3, i3_1])
+            T.writes(PadInput[i0_1, i1_1, i2_1, i3_1])
+            PadInput[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(
+                i1_1 >= 3 and i1_1 < 227 and i2_1 >= 3 and i2_1 < 227,
+                inputs[i0_1, i1_1 - 3, i2_1 - 3, i3_1],
+                T.float32(0),
+                dtype="float32",
+            )
+    for ax3, ax4, ax5, ax6 in T.grid(64, 7, 7, 3):
+        with T.sblock("weight_reindex"):
+            v3, v4, v5, v6 = T.axis.remap("SSSS", [ax3, ax4, ax5, ax6])
+            T.reads(weight[v4, v5, v6, v3])
+            T.writes(weight_reindex[v3, v4, v5, v6])
+            weight_reindex[v3, v4, v5, v6] = weight[v4, v5, v6, v3]
+    for i0, i1, i2, i3, i4, i5, i6 in T.grid(1, 112, 112, 64, 7, 7, 3):
+        with T.sblock("conv2d_nhwc"):
+            n, h, w, co, rh, rw, rc = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
+            T.reads(
+                PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc],
+                weight_reindex[co, rh, rw, rc],
+            )
+            T.writes(conv2d_nhwc[n, h, w, co])
+            with T.init():
+                conv2d_nhwc[n, h, w, co] = T.float32(0)
+            conv2d_nhwc[n, h, w, co] = (
+                conv2d_nhwc[n, h, w, co]
+                + PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc]
+                * weight_reindex[co, rh, rw, rc]
+            )
+
+
+@T.prim_func(s_tir=True)
+def matmul(
+    A: T.Buffer((512, 512), "float32"),
+    B: T.Buffer((512, 512), "float32"),
+    C: T.Buffer((512, 512), "float32"),
+) -> None:
+    for i0, i1, i2 in T.grid(512, 512, 512):
+        with T.sblock("matmul"):
+            i, j, k = T.axis.remap("SSR", [i0, i1, i2])
+            T.reads(C[i, j], A[i, k], B[k, j])
+            T.writes(C[i, j])
+            with T.init():
+                C[i, j] = T.float32(0)
+            C[i, j] = C[i, j] + A[i, k] * B[k, j]
+
+
+@T.prim_func(s_tir=True)
+def matmul_reindex_write(
+    A: T.Buffer((512, 512), "float32"),
+    B: T.Buffer((512, 512), "float32"),
+    C: T.Buffer((512, 512), "float32"),
+) -> None:
+    C_reindex = T.sblock_alloc_buffer([512, 512], dtype="float32")
+    for i0, i1, i2 in T.grid(512, 512, 512):
+        with T.sblock("matmul"):
+            i, j, k = T.axis.remap("SSR", [i0, i1, i2])
+            T.reads(C_reindex[i, j], A[i, k], B[k, j])
+            T.writes(C_reindex[i, j])
+            with T.init():
+                C_reindex[i, j] = T.float32(0)
+            C_reindex[i, j] = C_reindex[i, j] + A[i, k] * B[k, j]
+    for i0, i1 in T.grid(512, 512):
+        with T.sblock("C_reindex"):
+            v0, v1 = T.axis.remap("SS", [i0, i1])
+            T.reads(C_reindex[v0, v1])
+            T.writes(C[v0, v1])
+            C[v0, v1] = C_reindex[v0, v1]
+
+
+@T.prim_func(s_tir=True)
+def multiple_read(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")) -> None:
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vj, vi] + A[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def mixed_dtype(
+    p0: T.Buffer((T.int64(2), 1280), "float16"),
+    p1: T.Buffer((1280, 1280), "float16"),
+    T_matmul_NT: T.Buffer((T.int64(2), 1280), "float16"),
+) -> None:
+    for i0, i1, i2 in T.grid(T.int64(2), 1280, 1280):
+        with T.sblock("T_matmul_NT"):
+            i = T.axis.spatial(T.int64(2), i0)
+            j, k = T.axis.remap("SR", [i1, i2])
+            T.reads(p0[i, k], p1[j, k])
+            T.writes(T_matmul_NT[i, j])
+            with T.init():
+                T_matmul_NT[i, j] = T.float16(0)
+            T_matmul_NT[i, j] = T_matmul_NT[i, j] + p0[i, k] * p1[j, k]
+
+
+@T.prim_func(s_tir=True)
+def mixed_dtype_reindex_write(
+    p0: T.Buffer((T.int64(2), 1280), "float16"),
+    p1: T.Buffer((1280, 1280), "float16"),
+    T_matmul_NT: T.Buffer((T.int64(2), 1280), "float16"),
+) -> None:
+    T_matmul_NT_reindex = T.sblock_alloc_buffer([T.int64(2), 1280], dtype="float16")
+    for i0, i1, i2 in T.grid(T.int64(2), 1280, 1280):
+        with T.sblock("T_matmul_NT"):
+            i = T.axis.spatial(T.int64(2), i0)
+            j, k = T.axis.remap("SR", [i1, i2])
+            T.reads(p0[i, k], p1[j, k])
+            T.writes(T_matmul_NT_reindex[i, j])
+            with T.init():
+                T_matmul_NT_reindex[i, j] = T.float16(0)
+            T_matmul_NT_reindex[i, j] = T_matmul_NT_reindex[i, j] + p0[i, k] * p1[j, k]
+    for ax0, ax1 in T.grid(T.int64(2), 1280):
+        with T.sblock("T_matmul_NT_reindex"):
+            v0 = T.axis.spatial(T.int64(2), ax0)
+            v1 = T.axis.remap("S", [ax1])
+            T.reads(T_matmul_NT_reindex[v0, v1])
+            T.writes(T_matmul_NT[v0, v1])
+            T_matmul_NT[v0, v1] = T_matmul_NT_reindex[v0, v1]
+
+
+@T.prim_func(s_tir=True)
+def matmul_unit_dim(
+    A: T.Buffer((1, 512), "float32"),
+    B: T.Buffer((512, 1), "float32"),
+    C: T.Buffer((1, 1), "float32"),
+) -> None:
+    for i0, i1, i2 in T.grid(1, 1, 512):
+        with T.sblock("matmul"):
+            i, j, k = T.axis.remap("SSR", [i0, i1, i2])
+            T.reads(C[i, j], A[i, k], B[k, j])
+            T.writes(C[i, j])
+            with T.init():
+                C[i, j] = T.float32(0)
+            C[i, j] = C[i, j] + A[i, k] * B[k, j]
+
+
+@T.prim_func(s_tir=True)
+def matmul_unit_dim_reindex_write(
+    A: T.Buffer((1, 512), "float32"),
+    B: T.Buffer((512, 1), "float32"),
+    C: T.Buffer((1, 1), "float32"),
+) -> None:
+    C_reindex = T.sblock_alloc_buffer([1, 1], dtype="float32")
+    for i0, i1, i2 in T.grid(1, 1, 512):
+        with T.sblock("matmul"):
+            i, j, k = T.axis.remap("SSR", [i0, i1, i2])
+            T.reads(C_reindex[i, j], A[i, k], B[k, j])
+            T.writes(C_reindex[i, j])
+            with T.init():
+                C_reindex[i, j] = T.float32(0)
+            C_reindex[i, j] = C_reindex[i, j] + A[i, k] * B[k, j]
+    for i0, i1 in T.grid(1, 1):
+        with T.sblock("C_reindex"):
+            v0, v1 = T.axis.remap("SS", [i0, i1])
+            T.reads(C_reindex[v0, v1])
+            T.writes(C[v0, v1])
+            C[v0, v1] = C_reindex[v0, v1]
+
+
+use_block_name = tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True})
+use_buffer_name = tvm.testing.parameter(by_dict={"buffer_index": False, "buffer_name": True})
+
+
+def test_reindex_read_basic(use_block_name, use_buffer_name):
+    sch = tvm.s_tir.Schedule(transpose_elementwise)
+    block = "B" if use_block_name else sch.get_sblock("B")
+    buf = "A" if use_buffer_name else ("read", 0)
+    sch.reindex(block, buf)
+    assert_structural_equal_ignore_global_symbol(
+        transpose_elementwise_reindex_read, sch.mod["main"]
+    )
+    verify_trace_roundtrip(sch=sch, mod=transpose_elementwise)
+
+
+def test_conv2d_reindex_weight(use_block_name, use_buffer_name):
+    sch = tvm.s_tir.Schedule(conv2d_nhwc)
+    block = "conv2d_nhwc" if use_block_name else sch.get_sblock("conv2d_nhwc")
+    buf = "Weight" if use_buffer_name else ("read", 1)
+    sch.reindex(block, buf)
+    assert_structural_equal_ignore_global_symbol(conv2d_nhwc_reindex_weight, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=conv2d_nhwc)
+
+
+def test_conv2d_reindex_data(use_block_name, use_buffer_name):
+    sch = tvm.s_tir.Schedule(conv2d_nhwc)
+    block = "conv2d_nhwc" if use_block_name else sch.get_sblock("conv2d_nhwc")
+    buf = "PadInput" if use_buffer_name else ("read", 0)
+    sch.reindex(block, buf)
+    assert_structural_equal_ignore_global_symbol(conv2d_nhwc_reindex_data, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=conv2d_nhwc)
+
+
+def test_matmul_reindex_write(use_block_name, use_buffer_name):
+    sch = tvm.s_tir.Schedule(matmul)
+    block = "matmul" if use_block_name else sch.get_sblock("matmul")
+    buf = "C" if use_buffer_name else ("write", 0)
+    sch.reindex(block, buf)
+    assert_structural_equal_ignore_global_symbol(matmul_reindex_write, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=matmul)
+
+
+def test_reindex_fail_multiple_read(use_block_name, use_buffer_name):
+    sch = tvm.s_tir.Schedule(multiple_read)
+    block = "B" if use_block_name else sch.get_sblock("B")
+    buf = "A" if use_buffer_name else ("read", 0)
+    with pytest.raises(ScheduleError):
+        sch.reindex(block, buf)
+
+
+def test_reindex_mixed_dtype(use_block_name, use_buffer_name):
+    sch = tvm.s_tir.Schedule(mixed_dtype)
+    block = "T_matmul_NT" if use_block_name else sch.get_sblock("T_matmul_NT")
+    buf = "T_matmul_NT" if use_buffer_name else ("write", 0)
+    sch.reindex(block, buf)
+    assert_structural_equal_ignore_global_symbol(mixed_dtype_reindex_write, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=mixed_dtype)
+
+
+def test_matmul_unit_dim_reindex_write(use_block_name, use_buffer_name):
+    sch = tvm.s_tir.Schedule(matmul_unit_dim)
+    block = "matmul" if use_block_name else sch.get_sblock("matmul")
+    buf = "C" if use_buffer_name else ("write", 0)
+    sch.reindex(block, buf)
+    assert_structural_equal_ignore_global_symbol(matmul_unit_dim_reindex_write, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=matmul_unit_dim)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_reorder.py

@@ -0,0 +1,403 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E741, F401
+import sys
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def elementwise(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128, 128))
+    for i, j, k, l in T.grid(128, 128, 128, 128):
+        with T.sblock("B"):
+            vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+            B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_not_affine(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128, 128))
+    for i, j, k, l in T.grid(128, 128, 128, 8):
+        with T.sblock("B"):
+            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+            vl = T.axis.S(128, l * 16)
+            B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_dependent_loop(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128, 128))
+    for i in T.serial(0, 128):
+        for j, k, l in T.grid(128, i, 128):
+            with T.sblock("B"):
+                vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_predicate(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128, 128))
+    for i, j, k, l in T.grid(128, 128, 128, 128):
+        with T.sblock("B"):
+            T.where(i * 2097152 + j * 16384 + k * 128 + l < 100)
+            vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+            B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_non_single_branch(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    C = T.sblock_alloc_buffer((128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j in T.grid(128, 128):
+        for k in T.serial(0, 128):
+            with T.sblock("C"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                C[vi, vj, vk] = A[vi, vj, vk] * 2.0
+        for k in T.serial(0, 128):
+            with T.sblock("B"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                B[vi, vj, vk] = C[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_with_loops_not_same_scope(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j in T.grid(128, 128):
+        with T.sblock("A"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            for k in T.serial(0, 128):
+                with T.sblock("B"):
+                    vk = T.axis.S(128, k)
+                    T.reads([A[vi, vj, vk]])
+                    T.writes([B[vi, vj, vk]])
+                    B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_with_wrong_block_var_type(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            vk = T.axis.scan(128, k)
+            T.reads([A[vi, vj, vk]])
+            T.writes([B[vi, vj, vk]])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_reordered(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128, 128))
+    for l, j, k, i in T.grid(128, 128, 128, 128):
+        with T.sblock("B"):
+            vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+            B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_reordered2(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128, 128))
+    for k, j, i, l in T.grid(128, 128, 128, 128):
+        with T.sblock("B"):
+            vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+            B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_reordered_with_predicate(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128, 128))
+    for l, j, k, i in T.grid(128, 128, 128, 128):
+        with T.sblock("B"):
+            T.where(i * 2097152 + j * 16384 + k * 128 + l < 100)
+            vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+            B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def opaque_access(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [16, 16], "float32")
+    B = T.match_buffer(b, [16, 16], "float32")
+    for i, j in T.grid(16, 16):
+        with T.sblock("A"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([])
+            T.writes([A[0:16, 0:16]])
+            A[vi, vj] = 1
+    for i, j in T.grid(16, 16):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([])
+            T.writes([B[0:16, 0:16]])
+            T.evaluate(T.tvm_fill_fragment(B.data, 16, 16, 16, 0, vi * 16 + vj, dtype="handle"))
+
+
+@T.prim_func(s_tir=True)
+def opaque_access_reorder(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [16, 16], "float32")
+    B = T.match_buffer(b, [16, 16], "float32")
+    for j, i in T.grid(16, 16):
+        with T.sblock("A"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([])
+            T.writes([A[0:16, 0:16]])
+            A[vi, vj] = 1
+    for j, i in T.grid(16, 16):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([])
+            T.writes([B[0:16, 0:16]])
+            T.evaluate(T.tvm_fill_fragment(B.data, 16, 16, 16, 0, vi * 16 + vj, dtype="handle"))
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+
+def test_reorder():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k, l = sch.get_loops(block_b)
+    sch.reorder(l, i)
+    assert_structural_equal_ignore_global_symbol(elementwise_reordered, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+
+def test_reorder2():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k, l = sch.get_loops(block_b)
+    sch.reorder(k, i, l)
+    assert_structural_equal_ignore_global_symbol(elementwise_reordered2, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+
+def test_reorder_with_opaque_access():
+    sch = tvm.s_tir.Schedule(opaque_access, debug_mask="all")
+    block_a = sch.get_sblock("A")
+    i, j = sch.get_loops(block_a)
+    sch.reorder(j, i)
+    block_b = sch.get_sblock("B")
+    i, j = sch.get_loops(block_b)
+    sch.reorder(j, i)
+    assert_structural_equal_ignore_global_symbol(opaque_access_reorder, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=opaque_access)
+
+
+def test_reorder_overlapped_access():
+    @T.prim_func(s_tir=True)
+    def overlapped_access(A: T.Buffer((14, 4), "float32"), B: T.Buffer((14, 4), "float32")):
+        # example to write first axis multiple times
+        for v0, v1, v2 in T.grid(6, 4, 4):
+            with T.sblock("block"):
+                i = T.axis.spatial(14, v0 * 2 + v1)
+                j = T.axis.spatial(4, v2)
+                B[i, j] = A[i, j] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def overlapped_access_reorder(A: T.Buffer((14, 4), "float32"), B: T.Buffer((14, 4), "float32")):
+        # example to write first axis multiple times
+        for v0, v2, v1 in T.grid(6, 4, 4):
+            with T.sblock("block"):
+                i = T.axis.spatial(14, v0 * 2 + v1)
+                j = T.axis.spatial(4, v2)
+                B[i, j] = A[i, j] + 1.0
+
+    sch = tvm.s_tir.Schedule(overlapped_access, debug_mask="all")
+    v0, v1, v2 = sch.get_loops(sch.get_sblock("block"))
+    sch.reorder(v0, v2, v1)
+    assert_structural_equal_ignore_global_symbol(overlapped_access_reorder, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=overlapped_access)
+
+
+def test_reorder_with_partial_affineness():
+    @T.prim_func(s_tir=True)
+    def non_affine_func(A: T.Buffer((14, 4), "float32"), B: T.Buffer((14, 4), "float32")):
+        for v0, v1, v2 in T.grid(6, 4, 4):
+            with T.sblock("block"):
+                i = T.axis.spatial(14, v0 * v0 + v1)
+                j = T.axis.spatial(4, v2)
+                B[i, j] = A[i, j] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def non_affine_func_reorder(A: T.Buffer((14, 4), "float32"), B: T.Buffer((14, 4), "float32")):
+        for v0, v2, v1 in T.grid(6, 4, 4):
+            with T.sblock("block"):
+                i = T.axis.spatial(14, v0 * v0 + v1)
+                j = T.axis.spatial(4, v2)
+                B[i, j] = A[i, j] + 1.0
+
+    sch = tvm.s_tir.Schedule(non_affine_func, debug_mask="all")
+    v0, v1, v2 = sch.get_loops(sch.get_sblock("block"))
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder(v0, v2, v1)
+
+    sch.reorder(v2, v1)
+    assert_structural_equal_ignore_global_symbol(non_affine_func_reorder, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=non_affine_func)
+
+
+def test_reorder_with_cascade_tiled_ops():
+    @T.prim_func(s_tir=True)
+    def cascade_pool_ops(
+        x: T.Buffer((1, 16, 112, 112), "float32"), y2: T.Buffer((1, 16, 108, 108), "float32")
+    ) -> None:
+        y1 = T.sblock_alloc_buffer([1, 16, 110, 110], dtype="float32")
+        for n, c, h, w, kh, kw in T.grid(1, 16, 110, 110, 3, 3):
+            with T.sblock("pool_0"):
+                ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [n, c, h, w, kh, kw])
+                with T.init():
+                    y1[ax0, ax1, ax2, ax3] = 0.0
+                y1[ax0, ax1, ax2, ax3] = y1[ax0, ax1, ax2, ax3] + x[ax0, ax1, ax2 + rv0, ax3 + rv1]
+        for n, c, h, w, kh, kw in T.grid(1, 16, 108, 108, 3, 3):
+            with T.sblock("pool_1"):
+                ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [n, c, h, w, kh, kw])
+                with T.init():
+                    y2[ax0, ax1, ax2, ax3] = 0.0
+                y2[ax0, ax1, ax2, ax3] = y2[ax0, ax1, ax2, ax3] + y1[ax0, ax1, ax2 + rv0, ax3 + rv1]
+
+    @T.prim_func(s_tir=True)
+    def cascade_pool_ops_tile_reordered(
+        x: T.Buffer((1, 16, 112, 112), "float32"), y2: T.Buffer((1, 16, 108, 108), "float32")
+    ) -> None:
+        y1 = T.sblock_alloc_buffer([1, 16, 110, 110], dtype="float32")
+        for n, c, h_o in T.grid(1, 16, 27):
+            for w, h_i, kh, kw in T.grid(110, 6, 3, 3):
+                with T.sblock("pool_0"):
+                    ax0 = T.axis.spatial(1, 0)
+                    ax1 = T.axis.spatial(16, c)
+                    ax2 = T.axis.spatial(110, h_o * 4 + h_i)
+                    ax3, rv0, rv1 = T.axis.remap("SRR", [w, kh, kw])
+                    with T.init():
+                        y1[ax0, ax1, ax2, ax3] = 0.0
+                    y1[ax0, ax1, ax2, ax3] = (
+                        y1[ax0, ax1, ax2, ax3] + x[ax0, ax1, ax2 + rv0, ax3 + rv1]
+                    )
+            for h_i, w, kh, kw in T.grid(4, 108, 3, 3):
+                with T.sblock("pool_1"):
+                    ax0 = T.axis.spatial(1, n)
+                    ax1 = T.axis.spatial(16, c)
+                    ax2 = T.axis.spatial(108, h_o * 4 + h_i)
+                    ax3, rv0, rv1 = T.axis.remap("SRR", [w, kh, kw])
+                    with T.init():
+                        y2[ax0, ax1, ax2, ax3] = 0.0
+                    y2[ax0, ax1, ax2, ax3] = (
+                        y2[ax0, ax1, ax2, ax3] + y1[ax0, ax1, ax2 + rv0, ax3 + rv1]
+                    )
+
+    sch = tvm.s_tir.schedule.Schedule(cascade_pool_ops)
+    pool_0 = sch.get_sblock("pool_0")
+    pool_1 = sch.get_sblock("pool_1")
+    _, _, h, w, _, _ = sch.get_loops(pool_1)
+    ho, _ = sch.split(h, factors=[None, 4])
+    sch.compute_at(pool_0, ho)
+    _, _, _, h_i, w, _, _ = sch.get_loops(pool_0)
+    sch.reorder(w, h_i)
+    assert_structural_equal_ignore_global_symbol(
+        cascade_pool_ops_tile_reordered, sch.mod["main"], True
+    )
+    verify_trace_roundtrip(sch=sch, mod=cascade_pool_ops)
+
+
+def test_reorder_with_predicate():
+    sch = tvm.s_tir.Schedule(elementwise_predicate, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k, l = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder(l, i)
+
+
+def test_reorder_fail_with_multi_appearance_loops():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k, l = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder(k, i, i)
+
+
+def test_reorder_fail_with_non_single_branch_loop():
+    sch = tvm.s_tir.Schedule(elementwise_non_single_branch, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder(k, i)
+    sch = tvm.s_tir.Schedule(elementwise_non_single_branch, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    block_c = sch.get_sblock("C")
+    i, j, k1 = sch.get_loops(block_b)
+    _, _, k2 = sch.get_loops(block_c)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder(k1, i, k2)
+
+
+def test_reorder_fail_with_loops_not_under_same_scope():
+    sch = tvm.s_tir.Schedule(elementwise_with_loops_not_same_scope, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    block_a = sch.get_sblock("A")
+    i, j = sch.get_loops(block_a)
+    k = sch.get_loops(block_b)[0]
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder(k, i)
+
+
+def test_reorder_fail_with_wrong_block_var_type():
+    sch = tvm.s_tir.Schedule(elementwise_with_wrong_block_var_type, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder(k, i)
+
+
+def test_reorder_fail_with_dependent_loops():
+    sch = tvm.s_tir.Schedule(elementwise_dependent_loop, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k, l = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder(l, i)
+
+
+def test_reorder_fail_not_affine_bindings():
+    sch = tvm.s_tir.Schedule(elementwise_not_affine, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k, l = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder(l, i)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_reorder_block_iter_var.py

@@ -0,0 +1,90 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import verify_trace_roundtrip
+from tvm.script import tirx as T
+
+
+@T.prim_func(s_tir=True)
+def matmul(
+    A: T.Buffer((128, 128), "float32"),
+    B: T.Buffer((128, 128), "float32"),
+    C: T.Buffer((128, 128), "float32"),
+) -> None:
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("C"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def matmul_after_reorder_block_iter_var(
+    A: T.Buffer((128, 128), "float32"),
+    B: T.Buffer((128, 128), "float32"),
+    C: T.Buffer((128, 128), "float32"),
+):
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("C"):
+            vk, vj, vi = T.axis.remap("RSS", [k, j, i])
+            T.reads(A[vi, vk], B[vj, vk])
+            T.writes(C[vi, vj])
+            with T.init():
+                C[vi, vj] = T.float32(0)
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+def test_reorder_block_iter_var():
+    sch = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    C = sch.get_sblock("C")
+    sch.reorder_block_iter_var(C, [2, 1, 0])
+    tvm.ir.assert_structural_equal(
+        matmul_after_reorder_block_iter_var.with_attr("global_symbol", "matmul"), sch.mod["main"]
+    )
+    verify_trace_roundtrip(sch=sch, mod=matmul)
+
+
+def test_reorder_block_iter_var_fail_not_full():
+    sch = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    C = sch.get_sblock("C")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder_block_iter_var(C, [2, 1])
+
+
+def test_reorder_block_iter_var_fail_not_within_bound():
+    sch = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    C = sch.get_sblock("C")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder_block_iter_var(C, [-1, 3, 2])
+
+
+def test_reorder_block_iter_var_fail_not_unique():
+    sch = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    C = sch.get_sblock("C")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.reorder_block_iter_var(C, [0, 0, 2])
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py

@@ -0,0 +1,577 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import s_tir, tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+
+def check_rolling_buffer(
+    sch: s_tir.Schedule, origin: tirx.PrimFunc, expected: tirx.PrimFunc, check_run=False
+):
+    scheduled = sch.mod["main"]
+    assert_structural_equal_ignore_global_symbol(scheduled, expected)
+    verify_trace_roundtrip(sch, origin)
+    if check_run:
+        in_buffer = origin.buffer_map[origin.params[0]]
+        out_buffer = origin.buffer_map[origin.params[1]]
+        in_shape = [int(_) for _ in in_buffer.shape]
+        out_shape = [int(_) for _ in out_buffer.shape]
+        x = tvm.runtime.tensor(np.random.uniform(0, 64, in_shape).astype(in_buffer.dtype))
+        y0 = tvm.runtime.tensor(np.zeros(out_shape).astype(out_buffer.dtype))
+        y1 = tvm.runtime.tensor(np.zeros(out_shape).astype(out_buffer.dtype))
+        f_origin = tvm.compile(origin)
+        f_scheduled = tvm.compile(scheduled)
+        f_origin(x, y0)
+        f_scheduled(x, y1)
+        tvm.testing.assert_allclose(y0.numpy(), y1.numpy())
+
+
+def _tile_nd(s, tile, block_name):
+    outer_indices = []
+    inner_indices = []
+    block = s.get_sblock(block_name)
+    loops = s.get_loops(block)
+    for i, size in enumerate(tile):
+        outer, inner = s.split(loops[i], [None, size])
+        outer_indices.append(outer)
+        inner_indices.append(inner)
+
+    s.reorder(*outer_indices, *inner_indices)
+    return outer_indices, inner_indices
+
+
+def test_1d_rolling_buffer():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((4, 12), "int32"), C: T.Buffer((4, 8), "int32")):
+        B = T.sblock_alloc_buffer((4, 10), "int32")
+        for c in T.serial(4):
+            for i in T.serial(0, 10):
+                for k in T.serial(3):
+                    with T.sblock("B"):
+                        cc, vi, vk = T.axis.remap("SSR", [c, i, k])
+                        with T.init():
+                            B[cc, vi] = 0
+                        B[cc, vi] = B[cc, vi] + A[cc, vi + vk]
+            for i in T.serial(0, 8):
+                for k in T.serial(3):
+                    with T.sblock("C"):
+                        cc, vi, vk = T.axis.remap("SSR", [c, i, k])
+                        with T.init():
+                            C[cc, vi] = 0
+                        C[cc, vi] = C[cc, vi] + B[cc, vi + vk]
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((4, 12), "int32"), C: T.Buffer((4, 8), "int32")):
+        B = T.sblock_alloc_buffer([4, 6], dtype="int32")
+        for c, i_0 in T.grid(4, 2):
+            for ax0, ax1 in T.grid(6, 3):
+                with T.sblock("B"):
+                    T.where(i_0 < 1 or 2 <= ax0)
+                    cc = T.axis.spatial(4, c)
+                    vi = T.axis.opaque(10, i_0 * 4 + ax0)
+                    vk = T.axis.reduce(3, ax1)
+                    T.reads(A[cc, vi + vk])
+                    T.writes(B[cc, vi % 6])
+                    with T.init():
+                        B[cc, vi % 6] = 0
+                    B[cc, vi % 6] = B[cc, vi % 6] + A[cc, vi + vk]
+            for i_1, k in T.grid(4, 3):
+                with T.sblock("C"):
+                    cc = T.axis.spatial(4, c)
+                    vi = T.axis.opaque(8, i_0 * 4 + i_1)
+                    vk = T.axis.reduce(3, k)
+                    T.reads(B[cc, (vi + vk) % 6])
+                    T.writes(C[cc, vi])
+                    with T.init():
+                        C[cc, vi] = 0
+                    C[cc, vi] = C[cc, vi] + B[cc, (vi + vk) % 6]
+
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    _, i, _ = sch.get_loops(sch.get_sblock("C"))
+    io, _ = sch.split(i, [2, 4])
+    sch.compute_at(sch.get_sblock("B"), io)
+    sch.rolling_buffer(sch.get_sblock("B"), 0)
+    check_rolling_buffer(sch, before, expected, check_run=True)
+
+
+@T.prim_func(s_tir=True)
+def cascade_2_max_pool2d(A: T.Buffer((1, 12, 12, 16), "int8"), C: T.Buffer((1, 8, 8, 16), "int8")):
+    B = T.sblock_alloc_buffer([1, 10, 10, 16], dtype="int8")
+    for i0, i1, i2, i3, i4, i5 in T.grid(1, 10, 10, 16, 3, 3):
+        with T.sblock("B"):
+            ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+            with T.init():
+                B[ax0, ax1, ax2, ax3] = T.int8(-128)
+            B[ax0, ax1, ax2, ax3] = T.max(B[ax0, ax1, ax2, ax3], A[ax0, ax1 + rv0, ax2 + rv1, ax3])
+    for i0, i1, i2, i3, i4, i5 in T.grid(1, 8, 8, 16, 3, 3):
+        with T.sblock("C"):
+            ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+            with T.init():
+                C[ax0, ax1, ax2, ax3] = T.int8(-128)
+            C[ax0, ax1, ax2, ax3] = T.max(C[ax0, ax1, ax2, ax3], B[ax0, ax1 + rv0, ax2 + rv1, ax3])
+
+
+@T.prim_func(s_tir=True)
+def cascade_3_max_pool2d_with_stride(
+    A: T.Buffer((1, 24, 24, 16), "int8"), C: T.Buffer((1, 8, 8, 16), "int8")
+):
+    B_0 = T.sblock_alloc_buffer([1, 22, 22, 16], dtype="int8")
+    B_1 = T.sblock_alloc_buffer([1, 10, 10, 16], dtype="int8")
+    for i0, i1, i2, i3, i4, i5 in T.grid(1, 22, 22, 16, 3, 3):
+        with T.sblock("B_0"):
+            ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+            with T.init():
+                B_0[ax0, ax1, ax2, ax3] = T.int8(-128)
+            B_0[ax0, ax1, ax2, ax3] = T.max(
+                B_0[ax0, ax1, ax2, ax3], A[ax0, ax1 + rv0, ax2 + rv1, ax3]
+            )
+    for i0, i1, i2, i3, i4, i5 in T.grid(1, 10, 10, 16, 3, 3):
+        with T.sblock("B_1"):
+            ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+            with T.init():
+                B_1[ax0, ax1, ax2, ax3] = T.int8(-128)
+            B_1[ax0, ax1, ax2, ax3] = T.max(
+                B_1[ax0, ax1, ax2, ax3], B_0[ax0, ax1 * 2 + rv0, ax2 * 2 + rv1, ax3]
+            )
+    for i0, i1, i2, i3, i4, i5 in T.grid(1, 8, 8, 16, 3, 3):
+        with T.sblock("C"):
+            ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+            with T.init():
+                C[ax0, ax1, ax2, ax3] = T.int8(-128)
+            C[ax0, ax1, ax2, ax3] = T.max(
+                C[ax0, ax1, ax2, ax3], B_1[ax0, ax1 + rv0, ax2 + rv1, ax3]
+            )
+
+
+def test_cascade_max_pool2d_w_tiled():
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((1, 12, 12, 16), "int8"), C: T.Buffer((1, 8, 8, 16), "int8")):
+        B = T.sblock_alloc_buffer([1, 10, 6, 16], dtype="int8")
+        for i0_0, i1_0, i2_0, i3_0 in T.grid(1, 1, 2, 1):
+            for ax0, ax1, ax2, ax3, ax4 in T.grid(10, 6, 16, 3, 3):
+                with T.sblock("B"):
+                    T.where(i2_0 < 1 or 2 <= ax1)
+                    ax0_1 = T.axis.spatial(1, 0)
+                    ax1_1 = T.axis.spatial(10, ax0)
+                    ax2_1 = T.axis.opaque(10, i2_0 * 4 + ax1)
+                    ax3_1, rv0, rv1 = T.axis.remap("SRR", [ax2, ax3, ax4])
+                    T.reads(A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1])
+                    T.writes(B[ax0_1, ax1_1, ax2_1 % 6, ax3_1])
+                    with T.init():
+                        B[ax0_1, ax1_1, ax2_1 % 6, ax3_1] = T.int8(-128)
+                    B[ax0_1, ax1_1, ax2_1 % 6, ax3_1] = T.max(
+                        B[ax0_1, ax1_1, ax2_1 % 6, ax3_1], A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1]
+                    )
+            for i0_1, i1_1, i2_1, i3_1, i4, i5 in T.grid(1, 8, 4, 16, 3, 3):
+                with T.sblock("C"):
+                    ax0 = T.axis.spatial(1, i0_0 + i0_1)
+                    ax1 = T.axis.spatial(8, i1_0 * 8 + i1_1)
+                    ax2 = T.axis.opaque(8, i2_0 * 4 + i2_1)
+                    ax3 = T.axis.spatial(16, i3_0 * 16 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.reads(B[ax0, ax1 + rv0, (ax2 + rv1) % 6, ax3])
+                    T.writes(C[ax0, ax1, ax2, ax3])
+                    with T.init():
+                        C[ax0, ax1, ax2, ax3] = T.int8(-128)
+                    C[ax0, ax1, ax2, ax3] = T.max(
+                        C[ax0, ax1, ax2, ax3], B[ax0, ax1 + rv0, (ax2 + rv1) % 6, ax3]
+                    )
+
+    sch = tvm.s_tir.Schedule(cascade_2_max_pool2d, debug_mask="all")
+    oi, _ = _tile_nd(sch, [1, 8, 4, 16], "C")
+    sch.compute_at(sch.get_sblock("B"), oi[-1])
+    sch.rolling_buffer(sch.get_sblock("B"), 0)
+    check_rolling_buffer(sch, cascade_2_max_pool2d, expected, check_run=True)
+
+
+def test_cascade_max_pool2d_h_tiled():
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((1, 12, 12, 16), "int8"), C: T.Buffer((1, 8, 8, 16), "int8")):
+        B = T.sblock_alloc_buffer([1, 6, 10, 16], dtype="int8")
+        for i0_0, i1_0, i2_0, i3_0 in T.grid(1, 2, 1, 1):
+            for ax0, ax1, ax2, ax3, ax4 in T.grid(6, 10, 16, 3, 3):
+                with T.sblock("B"):
+                    T.where(i1_0 < 1 or 2 <= ax0)
+                    ax0_1 = T.axis.spatial(1, 0)
+                    ax1_1 = T.axis.opaque(10, i1_0 * 4 + ax0)
+                    ax2_1 = T.axis.spatial(10, ax1)
+                    ax3_1, rv0, rv1 = T.axis.remap("SRR", [ax2, ax3, ax4])
+                    T.reads(A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1])
+                    T.writes(B[ax0_1, ax1_1 % 6, ax2_1, ax3_1])
+                    with T.init():
+                        B[ax0_1, ax1_1 % 6, ax2_1, ax3_1] = T.int8(-128)
+                    B[ax0_1, ax1_1 % 6, ax2_1, ax3_1] = T.max(
+                        B[ax0_1, ax1_1 % 6, ax2_1, ax3_1], A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1]
+                    )
+            for i0_1, i1_1, i2_1, i3_1, i4, i5 in T.grid(1, 4, 8, 16, 3, 3):
+                with T.sblock("C"):
+                    ax0 = T.axis.spatial(1, i0_0 + i0_1)
+                    ax1 = T.axis.opaque(8, i1_0 * 4 + i1_1)
+                    ax2 = T.axis.spatial(8, i2_0 * 8 + i2_1)
+                    ax3 = T.axis.spatial(16, i3_0 * 16 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.reads(B[ax0, (ax1 + rv0) % 6, ax2 + rv1, ax3])
+                    T.writes(C[ax0, ax1, ax2, ax3])
+                    with T.init():
+                        C[ax0, ax1, ax2, ax3] = T.int8(-128)
+                    C[ax0, ax1, ax2, ax3] = T.max(
+                        C[ax0, ax1, ax2, ax3], B[ax0, (ax1 + rv0) % 6, ax2 + rv1, ax3]
+                    )
+
+    sch = tvm.s_tir.Schedule(cascade_2_max_pool2d, debug_mask="all")
+    io, _ = _tile_nd(sch, [1, 4, 8, 16], "C")
+    sch.compute_at(sch.get_sblock("B"), io[-1])
+    sch.rolling_buffer(sch.get_sblock("B"), 0)
+    check_rolling_buffer(sch, cascade_2_max_pool2d, expected, check_run=True)
+
+
+def test_cascade_max_pool2d_h_w_c_tiled():
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((1, 12, 12, 16), "int8"), C: T.Buffer((1, 8, 8, 16), "int8")):
+        B = T.sblock_alloc_buffer([1, 6, 10, 16], dtype="int8")
+        for i0_0, i1_0, i2_0, i3_0 in T.grid(1, 2, 2, 2):
+            for ax0, ax1, ax2, ax3, ax4 in T.grid(6, 6, 8, 3, 3):
+                with T.sblock("B"):
+                    T.where((i1_0 < 1 or 2 <= ax0) and (i2_0 < 1 or 2 <= ax1))
+                    ax0_1 = T.axis.spatial(1, 0)
+                    ax1_1 = T.axis.opaque(10, i1_0 * 4 + ax0)
+                    ax2_1 = T.axis.spatial(10, i2_0 * 4 + ax1)
+                    ax3_1 = T.axis.spatial(16, i3_0 * 8 + ax2)
+                    rv0, rv1 = T.axis.remap("RR", [ax3, ax4])
+                    T.reads(A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1])
+                    T.writes(B[ax0_1, ax1_1 % 6, ax2_1, ax3_1])
+                    with T.init():
+                        B[ax0_1, ax1_1 % 6, ax2_1, ax3_1] = T.int8(-128)
+                    B[ax0_1, ax1_1 % 6, ax2_1, ax3_1] = T.max(
+                        B[ax0_1, ax1_1 % 6, ax2_1, ax3_1], A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1]
+                    )
+            for i0_1, i1_1, i2_1, i3_1, i4, i5 in T.grid(1, 4, 4, 8, 3, 3):
+                with T.sblock("C"):
+                    ax0 = T.axis.spatial(1, i0_0 + i0_1)
+                    ax1 = T.axis.opaque(8, i1_0 * 4 + i1_1)
+                    ax2 = T.axis.spatial(8, i2_0 * 4 + i2_1)
+                    ax3 = T.axis.spatial(16, i3_0 * 8 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.reads(B[ax0, (ax1 + rv0) % 6, ax2 + rv1, ax3])
+                    T.writes(C[ax0, ax1, ax2, ax3])
+                    with T.init():
+                        C[ax0, ax1, ax2, ax3] = T.int8(-128)
+                    C[ax0, ax1, ax2, ax3] = T.max(
+                        C[ax0, ax1, ax2, ax3], B[ax0, (ax1 + rv0) % 6, ax2 + rv1, ax3]
+                    )
+
+    sch = tvm.s_tir.Schedule(cascade_2_max_pool2d, debug_mask="all")
+    io, _ = _tile_nd(sch, [1, 4, 4, 8], "C")
+    sch.compute_at(sch.get_sblock("B"), io[-1])
+    sch.rolling_buffer(sch.get_sblock("B"), 0)
+    check_rolling_buffer(sch, cascade_2_max_pool2d, expected, check_run=True)
+
+
+def test_cascade_max_pool2d_non_perfect_tiled():
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((1, 12, 12, 16), "int8"), C: T.Buffer((1, 8, 8, 16), "int8")) -> None:
+        B = T.sblock_alloc_buffer([1, 8, 10, 16], dtype="int8")
+        for i0_0, i1_0, i2_0, i3_0 in T.grid(1, 2, 2, 1):
+            for ax0, ax1, ax2, ax3, ax4 in T.grid(8, 8, 16, 3, 3):
+                with T.sblock("B"):
+                    T.where(
+                        i1_0 * 6 + ax0 < 10
+                        and i2_0 * 6 + ax1 < 10
+                        and (i1_0 < 1 or 2 <= ax0)
+                        and (i2_0 < 1 or 2 <= ax1)
+                    )
+                    ax0_1 = T.axis.spatial(1, 0)
+                    ax1_1 = T.axis.opaque(10, i1_0 * 6 + ax0)
+                    ax2_1 = T.axis.spatial(10, i2_0 * 6 + ax1)
+                    ax3_1, rv0, rv1 = T.axis.remap("SRR", [ax2, ax3, ax4])
+                    T.reads(A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1])
+                    T.writes(B[ax0_1, ax1_1 % 8, ax2_1, ax3_1])
+                    with T.init():
+                        B[ax0_1, ax1_1 % 8, ax2_1, ax3_1] = T.int8(-128)
+                    B[ax0_1, ax1_1 % 8, ax2_1, ax3_1] = T.max(
+                        B[ax0_1, ax1_1 % 8, ax2_1, ax3_1], A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1]
+                    )
+            for i0_1, i1_1, i2_1, i3_1, i4, i5 in T.grid(1, 6, 6, 16, 3, 3):
+                with T.sblock("C"):
+                    T.where(i1_0 * 6 + i1_1 < 8 and i2_0 * 6 + i2_1 < 8)
+                    ax0 = T.axis.spatial(1, i0_0 + i0_1)
+                    ax1 = T.axis.opaque(8, i1_0 * 6 + i1_1)
+                    ax2 = T.axis.spatial(8, i2_0 * 6 + i2_1)
+                    ax3 = T.axis.spatial(16, i3_0 * 16 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.reads(B[ax0, (ax1 + rv0) % 8, ax2 + rv1, ax3])
+                    T.writes(C[ax0, ax1, ax2, ax3])
+                    with T.init():
+                        C[ax0, ax1, ax2, ax3] = T.int8(-128)
+                    C[ax0, ax1, ax2, ax3] = T.max(
+                        C[ax0, ax1, ax2, ax3], B[ax0, (ax1 + rv0) % 8, ax2 + rv1, ax3]
+                    )
+
+    sch = tvm.s_tir.Schedule(cascade_2_max_pool2d, debug_mask="all")
+    io, _ = _tile_nd(sch, [1, 6, 6, 16], "C")
+    sch.compute_at(sch.get_sblock("B"), io[-1])
+    sch.rolling_buffer(sch.get_sblock("B"), 0)
+    check_rolling_buffer(sch, cascade_2_max_pool2d, expected, check_run=True)
+
+
+def test_cascade_3_max_pool2d_with_stride():
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((1, 24, 24, 16), "int8"), C: T.Buffer((1, 8, 8, 16), "int8")) -> None:
+        B_0 = T.sblock_alloc_buffer([1, 13, 22, 16], dtype="int8")
+        B_1 = T.sblock_alloc_buffer([1, 6, 10, 16], dtype="int8")
+        for i0_0, i1_0, i2_0, i3_0 in T.grid(1, 2, 2, 1):
+            for ax0, ax1, ax2, ax3, ax4 in T.grid(13, 13, 16, 3, 3):
+                with T.sblock("B_0"):
+                    T.where((i1_0 < 1 or 5 <= ax0) and (i2_0 < 1 or 5 <= ax1))
+                    ax0_1 = T.axis.spatial(1, 0)
+                    ax1_1 = T.axis.opaque(22, i1_0 * 8 + ax0)
+                    ax2_1 = T.axis.spatial(22, i2_0 * 8 + ax1)
+                    ax3_1, rv0, rv1 = T.axis.remap("SRR", [ax2, ax3, ax4])
+                    T.reads(A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1])
+                    T.writes(B_0[ax0_1, ax1_1 % 13, ax2_1, ax3_1])
+                    with T.init():
+                        B_0[ax0_1, ax1_1 % 13, ax2_1, ax3_1] = T.int8(-128)
+                    B_0[ax0_1, ax1_1 % 13, ax2_1, ax3_1] = T.max(
+                        B_0[ax0_1, ax1_1 % 13, ax2_1, ax3_1],
+                        A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1],
+                    )
+            for ax0, ax1, ax2, ax3, ax4 in T.grid(6, 6, 16, 3, 3):
+                with T.sblock("B_1"):
+                    T.where((i1_0 < 1 or 2 <= ax0) and (i2_0 < 1 or 2 <= ax1))
+                    ax0_2 = T.axis.spatial(1, 0)
+                    ax1_2 = T.axis.opaque(10, i1_0 * 4 + ax0)
+                    ax2_2 = T.axis.spatial(10, i2_0 * 4 + ax1)
+                    ax3_2, rv0, rv1 = T.axis.remap("SRR", [ax2, ax3, ax4])
+                    T.reads(B_0[ax0_2, (ax1_2 * 2 + rv0) % 13, ax2_2 * 2 + rv1, ax3_2])
+                    T.writes(B_1[ax0_2, ax1_2 % 6, ax2_2, ax3_2])
+                    with T.init():
+                        B_1[ax0_2, ax1_2 % 6, ax2_2, ax3_2] = T.int8(-128)
+                    B_1[ax0_2, ax1_2 % 6, ax2_2, ax3_2] = T.max(
+                        B_1[ax0_2, ax1_2 % 6, ax2_2, ax3_2],
+                        B_0[ax0_2, (ax1_2 * 2 + rv0) % 13, ax2_2 * 2 + rv1, ax3_2],
+                    )
+            for i0_1, i1_1, i2_1, i3_1, i4, i5 in T.grid(1, 4, 4, 16, 3, 3):
+                with T.sblock("C"):
+                    ax0_3 = T.axis.spatial(1, i0_0 + i0_1)
+                    ax1_3 = T.axis.opaque(8, i1_0 * 4 + i1_1)
+                    ax2_3 = T.axis.spatial(8, i2_0 * 4 + i2_1)
+                    ax3_3 = T.axis.spatial(16, i3_0 * 16 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.reads(B_1[ax0_3, (ax1_3 + rv0) % 6, ax2_3 + rv1, ax3_3])
+                    T.writes(C[ax0_3, ax1_3, ax2_3, ax3_3])
+                    with T.init():
+                        C[ax0_3, ax1_3, ax2_3, ax3_3] = T.int8(-128)
+                    C[ax0_3, ax1_3, ax2_3, ax3_3] = T.max(
+                        C[ax0_3, ax1_3, ax2_3, ax3_3],
+                        B_1[ax0_3, (ax1_3 + rv0) % 6, ax2_3 + rv1, ax3_3],
+                    )
+
+    sch = tvm.s_tir.Schedule(cascade_3_max_pool2d_with_stride, debug_mask="all")
+    io, _ = _tile_nd(sch, [1, 4, 4, 16], "C")
+    sch.compute_at(sch.get_sblock("B_1"), io[-1])
+    sch.compute_at(sch.get_sblock("B_0"), io[-1])
+    sch.rolling_buffer(sch.get_sblock("B_0"), 0)
+    sch.rolling_buffer(sch.get_sblock("B_1"), 0)
+    check_rolling_buffer(sch, cascade_3_max_pool2d_with_stride, expected, check_run=True)
+
+
+def test_upscale():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((1, 16, 16, 16), "int8"), C: T.Buffer((1, 24, 24, 16), "int8")) -> None:
+        B = T.sblock_alloc_buffer([1, 14, 14, 16], dtype="int8")
+        for i0_0, i1_0, i2_0, i3_0 in T.grid(1, 5, 5, 1):
+            for ax0, ax1, ax2, ax3, ax4 in T.grid(5, 5, 16, 3, 3):
+                with T.sblock("B"):
+                    T.where(i1_0 * 5 // 2 + ax0 < 14 and i2_0 * 5 // 2 + ax1 < 14)
+                    ax0_1 = T.axis.spatial(1, 0)
+                    ax1_1 = T.axis.spatial(14, i1_0 * 5 // 2 + ax0)
+                    ax2_1 = T.axis.spatial(14, i2_0 * 5 // 2 + ax1)
+                    ax3_1 = T.axis.spatial(16, ax2)
+                    rv0, rv1 = T.axis.remap("RR", [ax3, ax4])
+                    T.reads(A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1])
+                    T.writes(B[ax0_1, ax1_1, ax2_1, ax3_1])
+                    with T.init():
+                        B[ax0_1, ax1_1, ax2_1, ax3_1] = T.int8(-128)
+                    B[ax0_1, ax1_1, ax2_1, ax3_1] = T.max(
+                        B[ax0_1, ax1_1, ax2_1, ax3_1], A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1]
+                    )
+            for i0_1, i1_1, i2_1, i3_1, i4, i5 in T.grid(1, 5, 5, 16, 3, 3):
+                with T.sblock("C"):
+                    T.where(i1_0 * 5 + i1_1 < 24 and i2_0 * 5 + i2_1 < 24)
+                    ax0 = T.axis.spatial(1, i0_0 + i0_1)
+                    ax1 = T.axis.spatial(24, i1_0 * 5 + i1_1)
+                    ax2 = T.axis.spatial(24, i2_0 * 5 + i2_1)
+                    ax3 = T.axis.spatial(16, i3_0 * 16 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.reads(B[ax0, ax1 // 2 + rv0, ax2 // 2 + rv1, ax3])
+                    T.writes(C[ax0, ax1, ax2, ax3])
+                    with T.init():
+                        C[ax0, ax1, ax2, ax3] = T.int8(-128)
+                    C[ax0, ax1, ax2, ax3] = T.max(
+                        C[ax0, ax1, ax2, ax3], B[ax0, ax1 // 2 + rv0, ax2 // 2 + rv1, ax3]
+                    )
+
+    @T.prim_func(s_tir=True)
+    def expected(
+        A: T.Buffer((1, 16, 16, 16), "int8"), C: T.Buffer((1, 24, 24, 16), "int8")
+    ) -> None:
+        B = T.sblock_alloc_buffer([1, 5, 14, 16], dtype="int8")
+        for i0_0, i1_0, i2_0, i3_0 in T.grid(1, 5, 5, 1):
+            for ax0, ax1, ax2, ax3, ax4 in T.grid(5, 5, 16, 3, 3):
+                with T.sblock("B"):
+                    T.where(
+                        i1_0 * 5 // 2 + ax0 < 14
+                        and i2_0 * 5 // 2 + ax1 < 14
+                        and (i1_0 < 1 or 2 <= ax0)
+                        and (i2_0 < 1 or 2 <= ax1)
+                    )
+                    ax0_1 = T.axis.spatial(1, 0)
+                    ax1_1 = T.axis.opaque(14, i1_0 * 5 // 2 + ax0)
+                    ax2_1 = T.axis.spatial(14, i2_0 * 5 // 2 + ax1)
+                    ax3_1 = T.axis.spatial(16, ax2)
+                    rv0, rv1 = T.axis.remap("RR", [ax3, ax4])
+                    T.reads(A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1])
+                    T.writes(B[ax0_1, ax1_1 % 5, ax2_1, ax3_1])
+                    with T.init():
+                        B[ax0_1, ax1_1 % 5, ax2_1, ax3_1] = T.int8(-128)
+                    B[ax0_1, ax1_1 % 5, ax2_1, ax3_1] = T.max(
+                        B[ax0_1, ax1_1 % 5, ax2_1, ax3_1], A[ax0_1, ax1_1 + rv0, ax2_1 + rv1, ax3_1]
+                    )
+            for i0_1, i1_1, i2_1, i3_1, i4, i5 in T.grid(1, 5, 5, 16, 3, 3):
+                with T.sblock("C"):
+                    T.where(i1_0 * 5 + i1_1 < 24 and i2_0 * 5 + i2_1 < 24)
+                    ax0 = T.axis.spatial(1, i0_0 + i0_1)
+                    ax1 = T.axis.opaque(24, i1_0 * 5 + i1_1)
+                    ax2 = T.axis.spatial(24, i2_0 * 5 + i2_1)
+                    ax3 = T.axis.spatial(16, i3_0 * 16 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.reads(B[ax0, (ax1 // 2 + rv0) % 5, ax2 // 2 + rv1, ax3])
+                    T.writes(C[ax0, ax1, ax2, ax3])
+                    with T.init():
+                        C[ax0, ax1, ax2, ax3] = T.int8(-128)
+                    C[ax0, ax1, ax2, ax3] = T.max(
+                        C[ax0, ax1, ax2, ax3], B[ax0, (ax1 // 2 + rv0) % 5, ax2 // 2 + rv1, ax3]
+                    )
+
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    sch.rolling_buffer(sch.get_sblock("B"), 0)
+    check_rolling_buffer(sch, before, expected, check_run=True)
+
+
+def test_fail_rolling_buffer_multi_writers():
+    @T.prim_func(s_tir=True)
+    def func_multi_writers(
+        A: T.Buffer((1, 12, 12, 16), "int8"), C: T.Buffer((1, 12, 12, 16), "int8")
+    ):
+        B = T.sblock_alloc_buffer([1, 12, 12, 16], dtype="int8")
+        for i0, i1, i2, i3 in T.grid(1, 3, 3, 1):
+            for ax0, ax1, ax2 in T.grid(6, 6, 16):
+                with T.sblock("B_writer_0"):
+                    ax0_1 = T.axis.spatial(1, i0)
+                    ax1_1 = T.axis.spatial(12, i1 * 4 + ax0)
+                    ax2_1 = T.axis.spatial(12, i2 * 4 + ax1)
+                    ax3_1 = T.axis.spatial(16, ax2)
+                    with T.init():
+                        B[ax0_1, ax1_1, ax2_1, ax3_1] = T.int8(-128)
+                    B[ax0_1, ax1_1, ax2_1, ax3_1] = A[ax0_1, ax1_1, ax2_1, ax3_1] + T.int8(1)
+            for ax0, ax1, ax2 in T.grid(6, 6, 16):
+                with T.sblock("B_writer_1"):
+                    ax0_2 = T.axis.spatial(1, i0)
+                    ax1_2 = T.axis.spatial(12, i1 * 4 + ax0)
+                    ax2_2 = T.axis.spatial(12, i2 * 4 + ax1)
+                    ax3_2 = T.axis.spatial(16, ax2)
+                    with T.init():
+                        B[ax0_2, ax1_2, ax2_2, ax3_2] = T.int8(-128)
+                    B[ax0_2, ax1_2, ax2_2, ax3_2] = B[ax0_2, ax1_2, ax2_2, ax3_2] + A[
+                        ax0_2, ax1_2, ax2_2, ax3_2
+                    ] * T.int8(2)
+            for ax0, ax1, ax2, ax3, ax4, ax5 in T.grid(1, 4, 4, 16, 3, 3):
+                with T.sblock("C"):
+                    ax0_3 = T.axis.spatial(1, i0 + ax0)
+                    ax1_3 = T.axis.spatial(12, i1 * 4 + ax1)
+                    ax2_3 = T.axis.spatial(12, i2 * 4 + ax2)
+                    ax3_3 = T.axis.spatial(16, i3 * 16 + ax3)
+                    rv0, rv1 = T.axis.remap("RR", [ax4, ax5])
+                    with T.init():
+                        C[ax0_3, ax1_3, ax2_3, ax3_3] = T.int8(-128)
+                    C[ax0_3, ax1_3, ax2_3, ax3_3] = T.max(
+                        C[ax0_3, ax1_3, ax2_3, ax3_3], B[ax0_3, ax1_3 + rv0, ax2_3 + rv1, ax3_3]
+                    )
+
+    sch = tvm.s_tir.Schedule(func_multi_writers, debug_mask="all")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.rolling_buffer(sch.get_sblock("B_writer_0"), 0)
+
+
+def test_fail_rolling_buffer_not_match():
+    @T.prim_func(s_tir=True)
+    def func_non_overlap(
+        A: T.Buffer((1, 12, 12, 16), "int8"), C: T.Buffer((1, 12, 12, 16), "int8")
+    ):
+        B = T.sblock_alloc_buffer([1, 12, 12, 16], dtype="int8")
+        for i0_0, i1_0, i2_0, i3_0 in T.grid(1, 3, 3, 1):
+            for ax0, ax1, ax2 in T.grid(4, 4, 16):
+                with T.sblock("B"):
+                    ax0_1 = T.axis.spatial(1, 0)
+                    ax1_1 = T.axis.spatial(12, i1_0 * 4 + ax0)
+                    ax2_1 = T.axis.spatial(12, i2_0 * 4 + ax1)
+                    ax3 = T.axis.spatial(16, ax2)
+                    T.reads(A[ax0_1, ax1_1, ax2_1, ax3])
+                    T.writes(B[ax0_1, ax1_1, ax2_1, ax3])
+                    with T.init():
+                        B[ax0_1, ax1_1, ax2_1, ax3] = T.int8(-128)
+                    B[ax0_1, ax1_1, ax2_1, ax3] = A[ax0_1, ax1_1, ax2_1, ax3]
+            for i0_1, i1_1, i2_1, i3_1, i4, i5 in T.grid(1, 4, 4, 16, 1, 1):
+                with T.sblock("C"):
+                    ax0 = T.axis.spatial(1, i0_0 + i0_1)
+                    ax1 = T.axis.spatial(12, i1_0 * 4 + i1_1)
+                    ax2 = T.axis.spatial(12, i2_0 * 4 + i2_1)
+                    ax3 = T.axis.spatial(16, i3_0 * 16 + i3_1)
+                    rv0, rv1 = T.axis.remap("RR", [i4, i5])
+                    T.reads(B[ax0, ax1 + rv0, ax2 + rv1, ax3])
+                    T.writes(C[ax0, ax1, ax2, ax3])
+                    with T.init():
+                        C[ax0, ax1, ax2, ax3] = T.int8(-128)
+                    C[ax0, ax1, ax2, ax3] = T.max(
+                        C[ax0, ax1, ax2, ax3], B[ax0, ax1 + rv0, ax2 + rv1, ax3]
+                    )
+
+    sch = tvm.s_tir.Schedule(func_non_overlap, debug_mask="all")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.rolling_buffer(sch.get_sblock("B"), 0)
+
+
+def test_fail_rolling_buffer_injection_invalid():
+    sch = tvm.s_tir.Schedule(cascade_2_max_pool2d, debug_mask="all")
+    # Block B is not compute_at to Block C, so rolling_buffer injection is invalid.
+    _, _ = _tile_nd(sch, [1, 4, 8, 16], "C")
+    _, _ = _tile_nd(sch, [1, 4, 8, 16], "B")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.rolling_buffer(sch.get_sblock("B"), 0)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_sampling.py

@@ -0,0 +1,244 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+import sys
+from collections import defaultdict
+
+import numpy
+import pytest
+
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import verify_trace_roundtrip
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def elementwise(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 257, 1470))
+    B = T.match_buffer(b, (128, 257, 1470))
+    for i, j, k in T.grid(128, 257, 1470):
+        with T.sblock("B"):
+            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def tiled_conv2d_with_padding(
+    inputs: T.Buffer((1, 224, 224, 3), "float32"),
+    weight: T.Buffer((7, 7, 3, 64), "float32"),
+    conv2d_nhwc: T.Buffer((1, 112, 112, 64), "float32"),
+) -> None:
+    PadInput = T.sblock_alloc_buffer([1, 230, 230, 3], dtype="float32")
+    for i0, i1, i2, i3 in T.grid(1, 230, 230, 3):
+        with T.sblock("PadInput"):
+            i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+            T.reads(inputs[i0_1, i1_1 - 3, i2_1 - 3, i3_1])
+            T.writes(PadInput[i0_1, i1_1, i2_1, i3_1])
+            PadInput[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(
+                3 <= i1_1 and i1_1 < 227 and 3 <= i2_1 and i2_1 < 227,
+                inputs[i0_1, i1_1 - 3, i2_1 - 3, i3_1],
+                T.float32(0),
+                dtype="float32",
+            )
+    for (
+        i0_0,
+        i1_0,
+        i2_0,
+        i3_0,
+        i0_1_1,
+        i1_1_1,
+        i2_1_1,
+        i3_1_1,
+        i4_0,
+        i5_0,
+        i6_0,
+        i0_2,
+        i1_2,
+        i2_2,
+        i3_2,
+        i4_1,
+        i5_1,
+        i6_1,
+        i0_3,
+        i1_3,
+        i2_3,
+        i3_3,
+    ) in T.grid(1, 1, 4, 1, 1, 2, 4, 1, 7, 7, 1, 1, 1, 1, 1, 1, 1, 3, 1, 56, 7, 64):
+        with T.sblock("conv2d_nhwc"):
+            n = T.axis.spatial(1, 0)
+            h = T.axis.spatial(112, i1_1_1 * 56 + i1_3)
+            w = T.axis.spatial(112, i2_0 * 28 + i2_1_1 * 7 + i2_3)
+            co, rh, rw, rc = T.axis.remap("SRRR", [i3_3, i4_0, i5_0, i6_1])
+            T.reads(
+                conv2d_nhwc[n, h, w, co],
+                PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc],
+                weight[rh, rw, rc, co],
+            )
+            T.writes(conv2d_nhwc[n, h, w, co])
+            with T.init():
+                conv2d_nhwc[n, h, w, co] = T.float32(0)
+            conv2d_nhwc[n, h, w, co] = (
+                conv2d_nhwc[n, h, w, co]
+                + PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc] * weight[rh, rw, rc, co]
+            )
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+
+def test_sample_categorical():
+    """Test sample categorical sampling function"""
+    n = 1000
+    sch = tvm.s_tir.Schedule(elementwise, seed=42, debug_mask="all")
+    counter = defaultdict(int)
+    candidates = [5, 2, 7, 1]
+    probs = [0.15, 0.55, 0.05, 0.25]
+    for _ in range(n):
+        v = sch.get(sch.sample_categorical(candidates, probs))
+        counter[v] += 1
+    for i, prob in enumerate(probs):
+        assert (prob - 0.07) * n <= counter[candidates[i]] <= (prob + 0.07) * n
+    verify_trace_roundtrip(sch, mod=elementwise)
+
+
+def test_sample_categorical_copy():
+    """Check the random variable sampling results after schedule copy"""
+    n = 100
+    sch = tvm.s_tir.Schedule(elementwise, seed=42, debug_mask="all")
+    candidates = [1, 2, 3, 4]
+    probs = [0.1, 0.2, 0.3, 0.4]
+    rv_decisions = []
+    for _ in range(n):
+        rv = sch.sample_categorical(candidates, probs)  # pylint: disable=invalid-name
+        rv_decisions.append((rv, sch.get(rv)))
+    sch_copy = sch.copy()
+    for rv, decision in rv_decisions:  # pylint: disable=invalid-name
+        decision_copy = sch_copy.get(rv)
+        assert int(decision) == int(decision_copy)
+
+
+def test_sample_categorical_serialize():
+    """Check the random variable sampling results after schedule serialization"""
+    n = 100
+    sch = tvm.s_tir.Schedule(elementwise, seed=42, debug_mask="all")
+    candidates = [5, 6, 7, 8]
+    probs = [0.23, 0.19, 0.37, 0.21]
+    decisions = []
+    for _ in range(n):
+        rv = sch.get(sch.sample_categorical(candidates, probs))  # pylint: disable=invalid-name
+        decisions.append(rv)
+    new_sch = verify_trace_roundtrip(sch, mod=elementwise)
+    for i, new_inst in enumerate(new_sch.trace.insts):
+        assert decisions[i] == candidates[new_sch.trace.decisions[new_inst]]
+
+
+def test_sample_perfect_tile_power_of_two():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    i, _, _ = sch.get_loops(sch.get_sblock("B"))
+    factors = sch.sample_perfect_tile(i, n=4)
+    factors = [sch.get(i) for i in factors]
+    prod = factors[0] * factors[1] * factors[2] * factors[3]
+    assert prod == 128
+    verify_trace_roundtrip(sch, mod=elementwise)
+
+
+def test_sample_perfect_tile_prime():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    _, i, _ = sch.get_loops(sch.get_sblock("B"))
+    factors = sch.sample_perfect_tile(i, n=4)
+    factors = [sch.get(i) for i in factors]
+    prod = factors[0] * factors[1] * factors[2] * factors[3]
+    assert prod == 257
+    verify_trace_roundtrip(sch, mod=elementwise)
+
+
+def test_sample_perfect_tile_composite():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    _, _, i = sch.get_loops(sch.get_sblock("B"))
+    factors = sch.sample_perfect_tile(i, n=4)
+    factors = [sch.get(i) for i in factors]
+    prod = factors[0] * factors[1] * factors[2] * factors[3]
+    assert prod == 1470
+    verify_trace_roundtrip(sch, mod=elementwise)
+
+
+use_sugared_block = tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True})
+
+
+def test_sample_compute_location(use_sugared_block):
+    n = 100
+    sch = tvm.s_tir.Schedule(tiled_conv2d_with_padding, seed=42, debug_mask="all")
+    if use_sugared_block:
+        pad_input = "PadInput"
+    else:
+        pad_input = sch.get_sblock("PadInput")
+    decision_dict = dict()
+    for _ in range(n):
+        _ = sch.sample_compute_location(pad_input)  # pylint: disable=invalid-name
+        decision = sch.trace.decisions[sch.trace.insts[-1]]
+        decision_dict[decision] = decision_dict[decision] + 1 if decision in decision_dict else 1
+
+    n_candidates = 8
+    expected_rate = 1.0 / n_candidates
+    for _, cnt in decision_dict.items():
+        numpy.testing.assert_allclose(expected_rate, cnt / n, atol=0.04)
+
+
+def test_sample_perfect_tile_after_copy():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    sch_copy = sch.copy()
+    _, _, i = sch.get_loops(sch.get_sblock("B"))
+    sch.sample_perfect_tile(i, n=4)
+
+    _, _, i = sch_copy.get_loops(sch_copy.get_sblock("B"))
+    # Hangs if ForkSeed is not invoked when copying a schedule
+    sch_copy.sample_perfect_tile(i, n=4)
+
+
+def test_sample_perfect_tile_on_dynamic_loops():
+    """Currently dynamic loop is trivially tiled"""
+
+    @T.prim_func(s_tir=True)
+    def workload(a: T.handle) -> None:
+        n = T.int32()
+        A = T.match_buffer(a, (n, 1024))
+        for i, j in T.grid(n, 1024):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                A[vi, vj] = 1.0
+
+    sch = tvm.s_tir.Schedule(workload, debug_mask="all")
+    di, si = sch.get_loops(sch.get_sblock("B"))
+
+    factors = sch.sample_perfect_tile(si, n=4)
+    factors = [sch.get(i) for i in factors]
+    prod = factors[0] * factors[1] * factors[2] * factors[3]
+    assert prod == 1024
+
+    factors = sch.sample_perfect_tile(di, n=4)
+    assert factors[0] is None
+    factors = [sch.get(i) for i in factors[1:]]
+    prod = factors[0] * factors[1] * factors[2]
+    assert prod == 1
+    verify_trace_roundtrip(sch, mod=workload)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_set_axis_separator.py

@@ -0,0 +1,208 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E501, F401
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.tirx import IndexMap
+
+# fmt: off
+# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+@T.prim_func(s_tir=True)
+def element_wise(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B = T.sblock_alloc_buffer((128, 128), dtype="float32")
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_set_axis_separator(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B = T.sblock_alloc_buffer([128, 128], dtype="float32", axis_separators=[1])
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * T.float32(2)
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + T.float32(1)
+
+
+@T.prim_func(s_tir=True)
+def element_wise_set_axis_separator_input_buffer(A: T.Buffer(shape=(128, 128), dtype="float32", axis_separators=(1,)), C: T.Buffer((128, 128), "float32")) -> None:
+    B = T.sblock_alloc_buffer([128, 128], dtype="float32")
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * T.float32(2)
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + T.float32(1)
+
+
+@T.prim_func(s_tir=True)
+def element_wise_subregion_match(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B = T.sblock_alloc_buffer((128, 128), dtype="float32")
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_subregion0 = T.match_buffer(B[vi, vj], [], offset_factor=1)
+            B_subregion0[()] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_subregion1 = T.match_buffer(B[vi, vj], [], offset_factor=1)
+            C[vi, vj] = B_subregion1[()] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_subregion_match_set_axis_separator(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B = T.sblock_alloc_buffer([128, 128], dtype="float32", axis_separators=[1])
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_subregion0 = T.match_buffer(B[vi, vj], [], dtype="float32", offset_factor=1, axis_separators=[0])
+            B_subregion0[()] = A[vi, vj] * T.float32(2)
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_subregion1 = T.match_buffer(B[vi, vj], [], dtype="float32", offset_factor=1, axis_separators=[0])
+            C[vi, vj] = B_subregion1[()] + T.float32(1)
+
+
+# pylint: enable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+argument_style = tvm.testing.parameter('set_axis_separators',
+                                      'transform_layout_named',
+                                      'transform_layout_buffer_object',
+                                      )
+
+
+def test_set_axis_separator(argument_style):
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+
+    if argument_style=='set_axis_separators':
+        s.set_axis_separator(s.get_sblock("B"), ("write",0), [1])
+    elif argument_style=='transform_layout_named':
+        s.transform_layout(block='B', buffer='B', index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j])
+    elif argument_style =='transform_layout_buffer_object':
+        B = s.get(s.get_sblock('B')).writes[0].buffer
+        s.transform_layout(block='B', buffer=B, index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j])
+    else:
+        raise ValueError(f'Unexpected argument_style: {argument_style}')
+
+    assert_structural_equal_ignore_global_symbol(element_wise_set_axis_separator, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_set_scope_fail_on_index_out_of_bound():
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    with pytest.raises(AssertionError):
+        s.set_axis_separator(s.get_sblock("B"), ("write",1),[1])
+    with pytest.raises(AssertionError):
+        s.set_axis_separator(s.get_sblock("B"), ("read",-1),[1])
+
+
+def test_set_axis_separator_input_buffer(argument_style):
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+
+    if argument_style=='set_axis_separators':
+        s.set_axis_separator(s.get_sblock("B"), ("read",0), [1])
+    elif argument_style=='transform_layout_named':
+        s.transform_layout(block='B', buffer='A', index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j])
+    elif argument_style =='transform_layout_buffer_object':
+        A = s.get(s.get_sblock('B')).reads[0].buffer
+        s.transform_layout(block='B', buffer=A, index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j])
+    else:
+        raise ValueError(f'Unexpected argument_style: {argument_style}')
+
+
+    assert_structural_equal_ignore_global_symbol(element_wise_set_axis_separator_input_buffer, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_set_axis_separator_subregion(argument_style):
+    func = element_wise_subregion_match
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+
+    if argument_style=='set_axis_separators':
+        s.set_axis_separator(s.get_sblock("B"), ("write",0), [1])
+    elif argument_style=='transform_layout_named':
+        s.transform_layout(block='B', buffer='B', index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j])
+    elif argument_style =='transform_layout_buffer_object':
+        B = s.get(s.get_sblock('B')).writes[0].buffer
+        s.transform_layout(block='B', buffer=B, index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j])
+    else:
+        raise ValueError(f'Unexpected argument_style: {argument_style}')
+
+    assert_structural_equal_ignore_global_symbol(element_wise_subregion_match_set_axis_separator, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+def test_indexed_lookup():
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.sblock_alloc_buffer([4,4], dtype="int32")
+            B = T.sblock_alloc_buffer([1,1], dtype="int32")
+            for j in T.serial(4):
+                with T.sblock('block'):
+                    A[B[0,0],j] = 0
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.sblock_alloc_buffer([4,4], dtype="int32")
+            B = T.sblock_alloc_buffer([1,1], dtype="int32", axis_separators=[1])
+            for j in T.serial(4):
+                with T.sblock('block'):
+                    A[B[0,0],j] = 0
+
+    sch = tvm.s_tir.Schedule(Before)
+    sch.set_axis_separator('block', 'B', [1])
+    After = sch.mod
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_set_dtype.py

@@ -0,0 +1,130 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E501, F401
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# fmt: off
+# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+@T.prim_func(s_tir=True)
+def element_wise(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B = T.sblock_alloc_buffer((128, 128), dtype="float32")
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+@T.prim_func(s_tir=True)
+def element_wise_set_dtype(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
+    B = T.sblock_alloc_buffer((128, 128), "float16")
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(A[vi, vj])
+            T.writes(B[vi, vj])
+            B[vi, vj] = T.cast(A[vi, vj] * 2.0, "float16")
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(B[vi, vj])
+            T.writes(C[vi, vj])
+            C[vi, vj] = T.cast(B[vi, vj], "float32") + 1.0
+
+@T.prim_func(s_tir=True)
+def element_wise_subregion_match(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B = T.sblock_alloc_buffer((128, 128), dtype="float32")
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_subregion0 = T.match_buffer(B[vi, vj], [], offset_factor=1)
+            B_subregion0[()] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_subregion1 = T.match_buffer(B[vi, vj], [], offset_factor=1)
+            C[vi, vj] = B_subregion1[()] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_subregion_match_set_dtype(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B = T.sblock_alloc_buffer((128, 128), "float16")
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(A[vi, vj])
+            T.writes(B[vi, vj])
+            B_subregion0 = T.match_buffer(B[vi, vj], (), "float16", offset_factor=1)
+            B_subregion0[()] = T.cast(A[vi, vj] * 2.0, "float16")
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(B[vi, vj])
+            T.writes(C[vi, vj])
+            B_subregion1 = T.match_buffer(B[vi, vj], (), "float16", offset_factor=1)
+            C[vi, vj] = T.cast(B_subregion1[()], "float32") + 1.0
+
+
+use_block_name = tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True})
+
+def test_set_dtype(use_block_name):
+    func = element_wise
+    sch = tvm.s_tir.Schedule(func, debug_mask="all")
+    sch.unsafe_set_dtype("B" if use_block_name else sch.get_sblock("B"), 0, "float16")
+    assert_structural_equal_ignore_global_symbol(element_wise_set_dtype, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=func)
+
+def test_set_dtype_fail_on_output_buffer(use_block_name):
+    func = element_wise
+    sch = tvm.s_tir.Schedule(func, debug_mask='all')
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.unsafe_set_dtype('C' if use_block_name else sch.get_sblock("C"), 0, "float16")
+
+def test_set_dtype_fail_on_index_out_of_bound():
+    func = element_wise
+    sch = tvm.s_tir.Schedule(func, debug_mask='all')
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.unsafe_set_dtype(sch.get_sblock("B"), 1, "float64")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.unsafe_set_dtype(sch.get_sblock("B"), -1, "float64")
+
+def test_set_dtype_subregion():
+    func = element_wise_subregion_match
+    sch = tvm.s_tir.Schedule(func, debug_mask='all')
+    sch.unsafe_set_dtype(sch.get_sblock("B"), 0, "float16")
+    assert_structural_equal_ignore_global_symbol(element_wise_subregion_match_set_dtype, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=func)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_set_scope.py

@@ -0,0 +1,138 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E501, F401
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# fmt: off
+# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+@T.prim_func(s_tir=True)
+def element_wise(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B = T.sblock_alloc_buffer((128, 128), dtype="float32")
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_set_scope(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B_shared = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="shared")
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_shared[vi, vj] = A[vi, vj] * T.float32(2)
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B_shared[vi, vj] + T.float32(1)
+
+
+@T.prim_func(s_tir=True)
+def element_wise_subregion_match(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B = T.sblock_alloc_buffer((128, 128), dtype="float32")
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_subregion0 = T.match_buffer(B[vi, vj], [], offset_factor=1)
+            B_subregion0[()] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_subregion1 = T.match_buffer(B[vi, vj], [], offset_factor=1)
+            C[vi, vj] = B_subregion1[()] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_subregion_match_set_scope(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None:
+    B_shared = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="shared")
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_subregion0_shared = T.match_buffer(B_shared[vi, vj], [], dtype="float32", scope="shared", offset_factor=1)
+            B_subregion0_shared[()] = A[vi, vj] * T.float32(2)
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B_subregion1_shared = T.match_buffer(B_shared[vi, vj], [], dtype="float32", scope="shared", offset_factor=1)
+            C[vi, vj] = B_subregion1_shared[()] + T.float32(1)
+
+
+# pylint: enable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+
+use_block_name = tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True})
+use_buffer_name = tvm.testing.parameter(by_dict={"buffer_index": False, "buffer_name": True})
+
+def test_set_scope(use_block_name, use_buffer_name):
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    s.set_scope('B' if use_block_name else s.get_sblock("B"), 'B' if use_buffer_name else 0, "shared")
+    assert_structural_equal_ignore_global_symbol(element_wise_set_scope, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_set_scope_fail_on_output_buffer(use_block_name, use_buffer_name):
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.set_scope('C' if use_block_name else s.get_sblock("C"), 'C' if use_buffer_name else 0, "shared")
+
+
+def test_set_scope_fail_on_index_out_of_bound():
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.set_scope(s.get_sblock("B"), 1, "shared")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.set_scope(s.get_sblock("B"), -1, "shared")
+
+
+def test_set_scope_fail_on_invalid_scope():
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.set_scope(s.get_sblock("B"), 0, "test_scope")
+
+
+def test_set_scope_subregion():
+    func = element_wise_subregion_match
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    s.set_scope(s.get_sblock("B"), 0, "shared")
+    assert_structural_equal_ignore_global_symbol(element_wise_subregion_match_set_scope, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_split_fuse.py

@@ -0,0 +1,748 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401, F841
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import te, tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+from tvm.tirx.expr import IntImm
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def elementwise(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("B"):
+            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_dependent_loops(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i in T.serial(0, 128):
+        for j, k in T.grid(i, 128):
+            with T.sblock("B"):
+                vi = T.axis.S(128, i)
+                vj = T.axis.S(i, j)
+                vk = T.axis.S(128, k)
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_symbolic(a: T.handle, b: T.handle, n: T.int32) -> None:
+    A = T.match_buffer(a, (128, 128, n))
+    B = T.match_buffer(b, (128, 128, n))
+    for i, j, k in T.grid(128, 128, n):
+        with T.sblock("B"):
+            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_symbolic_fused(a: T.handle, b: T.handle, n: T.int32) -> None:
+    A = T.match_buffer(a, (128, 128, n))
+    B = T.match_buffer(b, (128, 128, n))
+    for i_j_k_fused in T.serial(0, (n * 16384)):
+        with T.sblock("B"):
+            vi = T.axis.S(128, T.floordiv(i_j_k_fused, n * 128))
+            vj = T.axis.S(128, T.floordiv(T.floormod(i_j_k_fused, n * 128), n))
+            vk = T.axis.S(n, T.floormod(i_j_k_fused, n))
+            T.reads([A[vi, vj, vk]])
+            T.writes([B[vi, vj, vk]])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_symbolic_split(a: T.handle, b: T.handle, n: T.int32) -> None:
+    A = T.match_buffer(a, (128, 128, n))
+    B = T.match_buffer(b, (128, 128, n))
+    for i, j, k0, k1 in T.grid(128, 128, 10, T.floordiv((n + 9), 10)):
+        with T.sblock("B"):
+            T.where(((k0 * T.floordiv((n + 9), 10)) + k1) < n)
+            vi, vj = T.axis.remap("SS", [i, j])
+            vk = T.axis.S(n, k0 * T.floordiv(n + 9, 10) + k1)
+            T.reads([A[vi, vj, vk]])
+            T.writes([B[vi, vj, vk]])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_with_seq(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    C = T.sblock_alloc_buffer((128, 128, 128))
+    for i, j in T.grid(128, 128):
+        for k in T.serial(0, 128):
+            with T.sblock("C"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                C[vi, vj, vk] = A[vi, vj, vk] * 2.0
+        for k in T.serial(0, 128):
+            with T.sblock("B"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                B[vi, vj, vk] = C[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_with_anno(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j in T.grid(128, 128):
+        for k in T.serial(0, 128, annotations={"useless_annotation": True}):
+            with T.sblock("B"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                T.reads([A[vi, vj, vk]])
+                T.writes([B[vi, vj, vk]])
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_with_thread_binding(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j in T.grid(128, 128):
+        for k in T.thread_binding(0, 128, thread="threadIdx.x"):
+            with T.sblock("B"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                T.reads([A[vi, vj, vk]])
+                T.writes([B[vi, vj, vk]])
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_with_starting_point(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j in T.grid(128, 128):
+        for k in T.serial(10, 128):
+            with T.sblock("B"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                T.reads([A[vi, vj, vk]])
+                T.writes([B[vi, vj, vk]])
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_with_opaque_block(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("opaque"):
+            T.reads([A[i, j, k]])
+            T.writes([B[i, j, k]])
+            with T.sblock("B"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                T.reads([A[vi, vj, vk]])
+                T.writes([B[vi, vj, vk]])
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_fused(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128, 128))
+    B = T.match_buffer(b, (128, 128, 128))
+    for fused in T.serial(0, 2097152):
+        with T.sblock("B"):
+            vi = T.axis.S(128, T.floordiv(fused, 16384))
+            vj = T.axis.S(128, T.floordiv(T.floormod(fused, 16384), 128))
+            vk = T.axis.S(128, T.floormod(fused, 128))
+            T.reads([A[vi, vj, vk]])
+            T.writes([B[vi, vj, vk]])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_split_case0(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128, 128])
+    B = T.match_buffer(b, [128, 128, 128])
+    for i1, i2, i3, j1, j2, k1, k2 in T.grid(2, 1, 64, 4, 32, 16, 8):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i1 * 64 + i2 * 64 + i3)
+            vj = T.axis.S(128, j1 * 32 + j2)
+            vk = T.axis.S(128, k1 * 8 + k2)
+            T.reads([A[vi, vj, vk]])
+            T.writes([B[vi, vj, vk]])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_split_case1(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128, 128])
+    B = T.match_buffer(b, [128, 128, 128])
+    for i1, i2, i3, j1, j2, j3, k1, k2, k3 in T.grid(2, 1, 64, 2, 1, 64, 2, 1, 64):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i1 * 64 + i2 * 64 + i3)
+            vj = T.axis.S(128, j1 * 64 + j2 * 64 + j3)
+            vk = T.axis.S(128, k1 * 64 + k2 * 64 + k3)
+            T.reads([A[vi, vj, vk]])
+            T.writes([B[vi, vj, vk]])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_split_with_predicate(a: T.handle, b: T.handle) -> None:
+    B = T.match_buffer(b, [128, 128, 128])
+    A = T.match_buffer(a, [128, 128, 128])
+    for i0, i1, i2, j0, j1, k0, k1 in T.grid(1000, 2, 3, 1, 129, 3, 43):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i0 * 6 + i1 * 3 + i2)
+            vj = T.axis.S(128, j0 * 129 + j1)
+            vk = T.axis.S(128, k0 * 43 + k1)
+            T.where((i0 * 2 + i1) * 3 + i2 < 128 and j0 * 129 + j1 < 128 and k0 * 43 + k1 < 128)
+            T.reads([A[vi, vj, vk]])
+            T.writes([B[vi, vj, vk]])
+            B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_fuse_with_opaque_block(a: T.handle, b: T.handle) -> None:
+    B = T.match_buffer(b, [128, 128, 128])
+    A = T.match_buffer(a, [128, 128, 128])
+    for i_j_k_fused in T.serial(0, 2097152):
+        with T.sblock("opaque"):
+            T.reads(
+                [
+                    A[
+                        T.floordiv(i_j_k_fused, 16384),
+                        T.floordiv(T.floormod(i_j_k_fused, 16384), 128),
+                        T.floormod(i_j_k_fused, 128),
+                    ]
+                ]
+            )
+            T.writes(
+                [
+                    B[
+                        T.floordiv(i_j_k_fused, 16384),
+                        T.floordiv(T.floormod(i_j_k_fused, 16384), 128),
+                        T.floormod(i_j_k_fused, 128),
+                    ]
+                ]
+            )
+            with T.sblock("B"):
+                vi = T.axis.S(128, T.floordiv(i_j_k_fused, 16384))
+                vj = T.axis.S(128, T.floordiv(T.floormod(i_j_k_fused, 16384), 128))
+                vk = T.axis.S(128, T.floormod(i_j_k_fused, 128))
+                T.reads([A[vi, vj, vk]])
+                T.writes([B[vi, vj, vk]])
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_split_with_opaque_block(a: T.handle, b: T.handle) -> None:
+    B = T.match_buffer(b, [128, 128, 128])
+    A = T.match_buffer(a, [128, 128, 128])
+
+    for i0, i1, j, k in T.grid(8, 16, 128, 128):
+        with T.sblock("opaque"):
+            T.reads([A[i0 * 16 + i1, j, k]])
+            T.writes([B[i0 * 16 + i1, j, k]])
+            with T.sblock("B"):
+                vi = T.axis.S(128, i0 * 16 + i1)
+                vj, vk = T.axis.remap("SS", [j, k])
+                T.reads([A[vi, vj, vk]])
+                T.writes([B[vi, vj, vk]])
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def opaque_access(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [16, 16], "float32")
+    B = T.match_buffer(b, [16, 16], "float32")
+    for i, j in T.grid(16, 16):
+        with T.sblock("A"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([])
+            T.writes([A[0:16, 0:16]])
+            A[vi, vj] = 1
+    for i, j in T.grid(16, 16):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([])
+            T.writes([B[0:16, 0:16]])
+            T.evaluate(T.tvm_fill_fragment(B.data, 16, 16, 16, 0, vi * 16 + vj, dtype="handle"))
+
+
+@T.prim_func(s_tir=True)
+def opaque_access_fused(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [16, 16])
+    B = T.match_buffer(b, [16, 16])
+    for i_j_fused in T.serial(0, 256):
+        with T.sblock("A"):
+            vi = T.axis.S(16, T.floordiv(i_j_fused, 16))
+            vj = T.axis.S(16, T.floormod(i_j_fused, 16))
+            T.reads([])
+            T.writes([A[0:16, 0:16]])
+            A[vi, vj] = 1
+    for i_j_fused in T.serial(0, 256):
+        with T.sblock("B"):
+            vi = T.axis.S(16, T.floordiv(i_j_fused, 16))
+            vj = T.axis.S(16, T.floormod(i_j_fused, 16))
+            T.reads([])
+            T.writes([B[0:16, 0:16]])
+            T.evaluate(T.tvm_fill_fragment(B.data, 16, 16, 16, 0, ((vi * 16) + vj), dtype="handle"))
+
+
+@T.prim_func(s_tir=True)
+def opaque_access_split(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16))
+    B = T.match_buffer(b, (16, 16))
+    for i, j0, j1 in T.grid(16, 4, 4):
+        with T.sblock("A"):
+            vi = T.axis.S(16, i)
+            vj = T.axis.S(16, j0 * 4 + j1)
+            T.reads([])
+            T.writes([A[0:16, 0:16]])
+            A[vi, vj] = 1
+    for i, j0, j1 in T.grid(16, 4, 4):
+        with T.sblock("B"):
+            vi = T.axis.S(16, i)
+            vj = T.axis.S(16, j0 * 4 + j1)
+            T.reads([])
+            T.writes([B[0:16, 0:16]])
+            T.evaluate(T.tvm_fill_fragment(B.data, 16, 16, 16, 0, ((vi * 16) + vj), dtype="handle"))
+
+
+@T.prim_func(s_tir=True)
+def elementwise_not_affine(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (127, 128))
+    B = T.match_buffer(b, (127, 128))
+    for i in T.serial(0, 4):
+        for j, k in T.grid(T.min(31, 126 - i * 32) + 1, 128):
+            with T.sblock("B"):
+                vi = T.axis.S(127, i * 32 + j)
+                vj = T.axis.S(128, k)
+                B[vi, vj] = A[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def elementwise_not_affine_fused(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [127, 128])
+    B = T.match_buffer(b, [127, 128])
+    for i in T.grid(4):
+        for j_k_fused in T.serial(0, T.min(31, 126 - i * 32) * 128 + 128):
+            with T.sblock("B"):
+                vi = T.axis.S(
+                    127,
+                    i * 32 + T.floordiv(j_k_fused, 128),
+                )
+                vj = T.axis.S(128, T.floormod(j_k_fused, 128))
+                T.reads([A[vi, vj]])
+                T.writes([B[vi, vj]])
+                B[vi, vj] = A[vi, vj]
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+
+def test_fuse():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    sch.fuse(i, j, k)
+    assert_structural_equal_ignore_global_symbol(elementwise_fused, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+
+@pytest.mark.parametrize("disable_predication", [True, False])
+def test_split(disable_predication):
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    sch.split(i, factors=[2, 1, 64], disable_predication=disable_predication)
+    sch.split(j, factors=[4, 32], disable_predication=disable_predication)
+    sch.split(k, factors=[16, 8], disable_predication=disable_predication)
+    assert_structural_equal_ignore_global_symbol(elementwise_split_case0, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+
+def test_split_with_inferred_factor():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    sch.split(i, factors=[None, 1, 64])
+    sch.split(j, factors=[2, None, 64])
+    sch.split(k, factors=[2, 1, None])
+    assert_structural_equal_ignore_global_symbol(elementwise_split_case1, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+
+def test_split_with_dynamic_inferred_factor():
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle) -> None:
+        N = T.int32()
+        M = T.int32()
+        A = T.match_buffer(a, (N, 128, M))
+        B = T.match_buffer(b, (N, 128, M))
+        for i, j, k in T.grid(N, 128, M):
+            with T.sblock("B"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                B[vi, vj, vk] = A[vi, vj, vk] * 2.0
+
+    @T.prim_func(s_tir=True)
+    def expected(a: T.handle, b: T.handle) -> None:
+        N, M = T.int32(), T.int32()
+        A = T.match_buffer(a, (N, 128, M))
+        B = T.match_buffer(b, (N, 128, M))
+        for i_0, i_1, j_0, j_1, k_0, k_1 in T.grid((N + 15) // 16, 16, 4, 32, 16, (M + 15) // 16):
+            with T.sblock("B"):
+                vi = T.axis.spatial(N, i_0 * 16 + i_1)
+                vj = T.axis.spatial(128, j_0 * 32 + j_1)
+                vk = T.axis.spatial(M, k_0 * ((M + 15) // 16) + k_1)
+                T.where(i_0 * 16 + i_1 < N and k_0 * ((M + 15) // 16) + k_1 < M)
+                B[vi, vj, vk] = A[vi, vj, vk] * T.float32(2.0)
+
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    sch.split(i, factors=[None, 16])
+    sch.split(j, factors=[4, 32])
+    sch.split(k, factors=[16, None])
+    assert_structural_equal_ignore_global_symbol(expected, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=before)
+
+
+def test_split_with_predicate():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    sch.split(i, factors=[1000, 2, 3])
+    sch.split(j, factors=[None, 129])
+    sch.split(k, factors=[3, None])
+    assert_structural_equal_ignore_global_symbol(elementwise_split_with_predicate, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+
+def test_fuse_fail_not_only_child():
+    sch = tvm.s_tir.Schedule(elementwise_with_seq, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.fuse(j, k)
+
+
+def test_fuse_split_fail_with_annotation():
+    sch = tvm.s_tir.Schedule(elementwise_with_anno, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.fuse(j, k)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.split(k, factors=[None, 10])
+
+
+def test_fuse_split_fail_not_start_with_zero():
+    sch = tvm.s_tir.Schedule(elementwise_with_anno, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.fuse(j, k)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.split(k, factors=[None, 10])
+
+
+def test_fuse_with_opaque_block():
+    sch = tvm.s_tir.Schedule(elementwise_with_opaque_block, debug_mask="all")
+    block_opaque = sch.get_sblock("opaque")
+    i, j, k = sch.get_loops(block_opaque)
+    sch.fuse(i, j, k)
+    assert_structural_equal_ignore_global_symbol(
+        elementwise_fuse_with_opaque_block, sch.mod["main"]
+    )
+    verify_trace_roundtrip(sch=sch, mod=elementwise_with_opaque_block)
+
+
+def test_fuse_with_opaque_access():
+    sch = tvm.s_tir.Schedule(opaque_access, debug_mask="all")
+    block_a = sch.get_sblock("A")
+    i, j = sch.get_loops(block_a)
+    sch.fuse(i, j)
+    block_b = sch.get_sblock("B")
+    i, j = sch.get_loops(block_b)
+    sch.fuse(i, j)
+    assert_structural_equal_ignore_global_symbol(opaque_access_fused, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=opaque_access)
+
+
+def test_split_with_opaque_block():
+    sch = tvm.s_tir.Schedule(elementwise_with_opaque_block, debug_mask="all")
+    block_opaque = sch.get_sblock("opaque")
+    i, _, _ = sch.get_loops(block_opaque)
+    sch.split(i, factors=[None, 16])
+    assert_structural_equal_ignore_global_symbol(
+        elementwise_split_with_opaque_block, sch.mod["main"]
+    )
+    verify_trace_roundtrip(sch=sch, mod=elementwise_with_opaque_block)
+
+
+def test_split_with_opaque_access():
+    sch = tvm.s_tir.Schedule(opaque_access, debug_mask="all")
+    block_a = sch.get_sblock("A")
+    _, j = sch.get_loops(block_a)
+    sch.split(j, factors=[None, 4])
+    block_b = sch.get_sblock("B")
+    _, j = sch.get_loops(block_b)
+    sch.split(j, factors=[None, 4])
+    assert_structural_equal_ignore_global_symbol(opaque_access_split, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=opaque_access)
+
+
+def test_split_with_non_positive_factors():
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.split(i, factors=[-2, -64])
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.split(j, factors=[0, None])
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.split(k, factors=[None, -16])
+
+
+def test_fuse_split_fail_with_thread_binding():
+    sch = tvm.s_tir.Schedule(elementwise_with_thread_binding, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, j, k = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.fuse(j, k)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.split(k, factors=[None, 10])
+
+
+def test_fuse_symbolic():
+    sch = tvm.s_tir.Schedule(elementwise_symbolic, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, k = sch.get_loops(block_b)
+    sch.fuse(i, j, k)
+    assert_structural_equal_ignore_global_symbol(elementwise_symbolic_fused, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise_symbolic)
+
+
+def test_split_symbolic():
+    sch = tvm.s_tir.Schedule(elementwise_symbolic, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, _, k = sch.get_loops(block_b)
+    sch.split(k, factors=[10, None])
+    assert_structural_equal_ignore_global_symbol(elementwise_symbolic_split, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise_symbolic)
+
+
+def test_fuse_fail_with_dependent_loops():
+    sch = tvm.s_tir.Schedule(elementwise_dependent_loops, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j, _ = sch.get_loops(block_b)
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        sch.fuse(i, j)
+
+
+def test_fuse_not_affine():
+    sch = tvm.s_tir.Schedule(elementwise_not_affine, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, j, k = sch.get_loops(block_b)
+    sch.fuse(j, k)
+    assert_structural_equal_ignore_global_symbol(elementwise_not_affine_fused, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=elementwise_not_affine)
+
+
+def test_add_unit_loop_above_block():
+    @T.prim_func(s_tir=True)
+    def zero_dim(
+        A: T.Buffer((), "int32"),
+        B: T.Buffer((), "int32"),
+        C: T.Buffer((), "int32"),
+    ) -> None:
+        with T.sblock("C"):
+            vi = T.axis.spatial(1, 0)
+            C[()] = A[()] + B[()]
+
+    @T.prim_func(s_tir=True)
+    def zero_dim_added(
+        A: T.Buffer((), "int32"),
+        B: T.Buffer((), "int32"),
+        C: T.Buffer((), "int32"),
+    ) -> None:
+        for u in range(1):
+            with T.sblock("C"):
+                vi = T.axis.spatial(1, 0)
+                C[()] = A[()] + B[()]
+
+    sch = tvm.s_tir.Schedule(zero_dim, debug_mask="all")
+    block = sch.get_sblock("C")
+    sch.add_unit_loop(block)
+    assert_structural_equal_ignore_global_symbol(zero_dim_added, sch.mod["main"])
+
+
+def test_add_unit_loop_above_loop():
+    @T.prim_func(s_tir=True)
+    def zero_dim(
+        A: T.Buffer((), "int32"),
+        B: T.Buffer((), "int32"),
+        C: T.Buffer((), "int32"),
+    ) -> None:
+        for u in range(1):
+            with T.sblock("C"):
+                vi = T.axis.spatial(1, 0)
+                C[()] = A[()] + B[()]
+
+    @T.prim_func(s_tir=True)
+    def zero_dim_added(
+        A: T.Buffer((), "int32"),
+        B: T.Buffer((), "int32"),
+        C: T.Buffer((), "int32"),
+    ) -> None:
+        for u1, u2 in T.grid(1, 1):
+            with T.sblock("C"):
+                vi = T.axis.spatial(1, 0)
+                C[()] = A[()] + B[()]
+
+    sch = tvm.s_tir.Schedule(zero_dim, debug_mask="all")
+    block = sch.get_sblock("C")
+    (loop,) = sch.get_loops(block)
+    sch.add_unit_loop(loop)
+    assert_structural_equal_ignore_global_symbol(zero_dim_added, sch.mod["main"])
+
+
+@pytest.mark.skip("Pending fix in affine analysis")
+def test_fuse_int64():
+    def _create_prim_func():
+        n = te.const(16, "int32")
+        m = te.const(32, "int64")
+        A = te.placeholder((n, m), name="A", dtype="int32")
+        B = te.compute((n, m), lambda i, j: A[i, j] + 1, name="B")
+        return te.create_prim_func([A, B])
+
+    mod = _create_prim_func()
+    sch = tvm.s_tir.Schedule(mod, debug_mask="all")
+    i, j = sch.get_loops(sch.get_sblock("B"))
+    sch.fuse(i, j)
+    verify_trace_roundtrip(sch=sch, mod=mod)
+
+
+def test_split_int64_extent_with_mixed_factors():
+    def _create_prim_func():
+        m = te.const(384, "int64")
+        A = te.placeholder((m,), name="A", dtype="float32")
+        B = te.compute((m,), lambda i: A[i] + 1, name="B")
+        return te.create_prim_func([A, B])
+
+    mod = _create_prim_func()
+    sch = tvm.s_tir.Schedule(mod, debug_mask="all")
+    (i,) = sch.get_loops(sch.get_sblock("B"))
+    sch.split(
+        i,
+        factors=[
+            te.const(1, "int64"),
+            te.const(512, "int32"),
+        ],
+    )
+
+
+def test_split_int64_extent_with_int32_factors():
+    def _create_prim_func():
+        m = te.const(12, "int64")
+        A = te.placeholder((m,), name="A", dtype="float32")
+        B = te.compute((m,), lambda i: A[i] + 1, name="B")
+        return te.create_prim_func([A, B])
+
+    mod = _create_prim_func()
+    sch = tvm.s_tir.Schedule(mod, debug_mask="all")
+    (i,) = sch.get_loops(sch.get_sblock("B"))
+    sch.split(
+        i,
+        factors=[
+            te.const(1, "int32"),
+            te.const(1, "int32"),
+            te.const(3, "int32"),
+            te.const(1, "int32"),
+            te.const(4, "int32"),
+        ],
+    )
+
+
+def test_split_int64_factors():
+    sch = tvm.s_tir.Schedule(elementwise_symbolic, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    _, _, k = sch.get_loops(block_b)
+    sch.split(k, factors=[IntImm(dtype="int64", value=10), None])
+    assert_structural_equal_ignore_global_symbol(elementwise_symbolic_split, sch.mod["main"])
+
+
+def test_unsupported_target_scalable_split():
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle):
+        A = T.match_buffer(a, (128,), "float32")
+        T.func_attr({"global_symbol": "my_module", "tirx.noalias": True})
+        for i in T.serial(128):
+            with T.sblock("A"):
+                v_i = T.axis.remap("S", [i])
+                A[v_i] = 1.0
+
+    sch = tvm.s_tir.Schedule(before)
+    (a,) = sch.get_loops("A")
+
+    err_msg = "The product of factors is not larger than or equal to the extent of loop tirx.For#0"
+    with pytest.raises(tvm.s_tir.schedule.ScheduleError, match=err_msg):
+        sch.split(a, factors=[T.ceildiv(128, 4 * T.vscale()), 4 * T.vscale()])
+
+
+def test_fused_symbolic_2D_tiling():
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle, M: T.int32, N: T.int32) -> None:
+        A = T.match_buffer(a, (M, N))
+        B = T.match_buffer(b, (M, N))
+        for i, j in T.grid(M, N):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                B[vi, vj] = A[vi, vj] * 2.0
+
+    @T.prim_func(s_tir=True)
+    def expected(a: T.handle, b: T.handle, M: T.int32, N: T.int32) -> None:
+        A = T.match_buffer(a, (M, N))
+        B = T.match_buffer(b, (M, N))
+        for i_0_j_0_fused, i_1, j_1 in T.grid(((M + 63) // 64) * ((N + 15) // 16), 64, 16):
+            with T.sblock("B"):
+                vi = T.axis.spatial(M, i_0_j_0_fused // ((N + 15) // 16) * 64 + i_1)
+                vj = T.axis.spatial(N, i_0_j_0_fused % ((N + 15) // 16) * 16 + j_1)
+                T.where(
+                    i_0_j_0_fused // ((N + 15) // 16) * 64 + i_1 < M
+                    and i_0_j_0_fused % ((N + 15) // 16) * 16 + j_1 < N
+                )
+                B[vi, vj] = A[vi, vj] * T.float32(2.0)
+
+    sch = tvm.s_tir.Schedule(before, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    i, j = sch.get_loops(block_b)
+    i0, i1 = sch.split(i, factors=[None, 64])
+    j0, j1 = sch.split(j, factors=[None, 16])
+    sch.reorder(i0, j0, i1, j1)
+    sch.fuse(i0, j0)
+    assert_structural_equal_ignore_global_symbol(expected, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=before)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_state.py

@@ -0,0 +1,358 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import gc
+import sys
+
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.ir import IRModule
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def elementwise(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128), "float32")
+    C = T.match_buffer(c, (128, 128), "float32")
+    B = T.sblock_alloc_buffer((128, 128), "float32")
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j in T.grid(128, 128):
+        with T.sblock("init"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = T.float32(0)
+        for k in range(0, 128):
+            with T.sblock("update"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def block_in_opaque_block(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128), "float32")
+    B = T.match_buffer(b, (128, 128), "float32")
+    for i in range(128):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i)
+            T.reads([A[0:128, 0:128]])
+            T.writes([B[0:128, 0:128]])
+            B[vi, 0] = A[vi, 0]
+            if A[vi, 0] == 0.0:
+                with T.sblock("C"):
+                    T.reads([A[0:128, 0:128]])
+                    T.writes([B[0:128, 0:128]])
+                    for j in range(128):
+                        with T.sblock("D"):
+                            vj = T.axis.S(128, j)
+                            B[vi, vj] = A[vi, vj] * 3.0
+            else:
+                with T.sblock("E"):
+                    T.reads([A[0:128, 0:128]])
+                    T.writes([B[0:128, 0:128]])
+                    for j in range(128):
+                        with T.sblock("F"):
+                            vj = T.axis.S(128, j)
+                            B[vi, vj] = A[vi, vj] * 2.0
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+
+def replace_ir_builder(deep_copy=False, realize=False):
+    new_func = tvm.script.from_source(elementwise.script())
+    s = tvm.s_tir.ScheduleState(new_func, debug_mask="all")
+    target = tvm.tirx.SBlock(
+        iter_vars=[],
+        reads=[],
+        writes=[],
+        name_hint="target",
+        body=s.mod["main"].body.block.body[1],
+        init=None,
+        alloc_buffers=None,
+        match_buffers=None,
+        annotations=None,
+    )
+    if realize:
+        target = tvm.tirx.SBlockRealize(
+            iter_values=[],
+            predicate=True,
+            block=target,
+        )
+    if deep_copy:
+        target.__setstate__(target.__getstate__())
+    gc.collect()
+    return s, target
+
+
+def replace_ir_builder_module(deep_copy=False, realize=False):
+    new_func = tvm.script.from_source(elementwise.script())
+    other_func = tvm.script.from_source(elementwise.script())
+    mod = IRModule(functions={"main": new_func, "other": other_func})
+    s = tvm.s_tir.ScheduleState(mod, debug_mask="all")
+    target = tvm.tirx.SBlock(
+        iter_vars=[],
+        reads=[],
+        writes=[],
+        name_hint="target",
+        body=s.mod["main"].body.block.body[1],
+        init=None,
+        alloc_buffers=None,
+        match_buffers=None,
+        annotations=None,
+    )
+    if realize:
+        target = tvm.tirx.SBlockRealize(
+            iter_values=[],
+            predicate=True,
+            block=target,
+        )
+    if deep_copy:
+        target.__setstate__(target.__getstate__())
+    gc.collect()
+    return s, target
+
+
+def replace_ir_builder_with_opaque():
+    func = tvm.script.from_source(block_in_opaque_block.script())
+    s = tvm.s_tir.ScheduleState(func, debug_mask="all")
+    gc.collect()
+    return s
+
+
+def test_replace_direct_write0():
+    s, target = replace_ir_builder(realize=True)
+    old_hash = s.mod["main"].__hash__()
+    sref = s.get_sref(s.mod["main"].body.block.body[1])
+    s.replace(sref, target)
+    # Check the replaced part is equal to the target
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block.body[1], target)
+    # There is no other reference so the AST node can be written directly
+    assert old_hash == s.mod["main"].__hash__()
+    # The target reuse the stmt of the sref, so the sref won't be None
+    assert sref.stmt is not None
+
+
+def test_replace_direct_write1():
+    s, target = replace_ir_builder(realize=True)
+    old_hash = s.mod["main"].body.block.body.__hash__()
+    hold_ref = s.mod["main"].body.block.body[1]
+    sref = s.get_sref(s.mod["main"].body.block.body[1])
+    s.replace(sref, target)
+    # There is no other reference so the AST node can be written directly
+    assert old_hash == s.mod["main"].body.block.body.__hash__()
+    assert not tvm_ffi.structural_equal(hold_ref.body, target)
+    # Check the replaced part is equal to the target
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block.body[1], target)
+    # The target reuse `sref.stmt`, so the sref won't be None
+    assert sref.stmt is not None
+
+
+def test_replace_copy():
+    s, target = replace_ir_builder(deep_copy=True, realize=True)
+    old_hash = s.mod["main"].__hash__()
+    # We hold another reference of func
+    old_func = s.mod["main"]
+    sref = s.get_sref(s.mod["main"].body.block.body[0])
+    s.replace(sref, target)
+    # We need to copy the whole func to remain the old_func unchanged
+    assert old_hash != s.mod["main"].__hash__()
+    assert not tvm_ffi.structural_equal(old_func.body, s.mod["main"].body)
+    assert old_hash == old_func.__hash__()
+    # Check the replaced part is equal to the target
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block.body[0], target)
+    # The replaced AST node will be deleted, so the ref will be None
+    assert sref.stmt is None
+
+
+def test_replace_partial_copy0():
+    s, target = replace_ir_builder(deep_copy=True, realize=True)
+    func_old_hash = s.mod["main"].__hash__()
+    hold_ref = s.mod["main"].body.block.body[0]
+    ref_old_hash = hold_ref.__hash__()
+    sref = s.get_sref(s.mod["main"].body.block.body[0].body)
+    other_part_hash = s.mod["main"].body.block.body[1].__hash__()
+    s.replace(sref, target)
+    # The stmt is held by `hold_sref`, so it will be coped in copy-on-write
+    # because the ref count is not unique
+    assert ref_old_hash != s.mod["main"].body.block.body[0].__hash__()
+    assert not tvm_ffi.structural_equal(hold_ref.body, target)
+    # The function and the other part stmt can be directly written
+    assert func_old_hash == s.mod["main"].__hash__()
+    assert other_part_hash == s.mod["main"].body.block.body[1].__hash__()
+    # Check the replaced part is equal to the target
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block.body[0].body, target)
+    # The replaced AST node will be deleted, so the ref will be None
+    assert sref.stmt is None
+
+
+def test_replace_partial_copy1():
+    s, target = replace_ir_builder(deep_copy=True)
+    func_old_hash = s.mod["main"].__hash__()
+    hold_ref = s.mod["main"].body.block.body[0].body
+    stmt_old_hash = s.mod["main"].body.block.body[0].__hash__()
+    sref = s.get_sref(s.mod["main"].body.block.body[0].body.body.block)
+    other_part_hash = s.mod["main"].body.block.body[1].__hash__()
+    s.replace(sref, target)
+    # The parent stmt will change since there is only one reference
+    assert stmt_old_hash == s.mod["main"].body.block.body[0].__hash__()
+    assert not tvm_ffi.structural_equal(hold_ref.body, target)
+    # The function and the other part stmt can be directly written
+    assert func_old_hash == s.mod["main"].__hash__()
+    assert other_part_hash == s.mod["main"].body.block.body[1].__hash__()
+    # Check the replaced part is equal to the target
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block.body[0].body.body.block, target)
+    # The replaced AST node will be deleted, so the ref will be None
+    assert sref.stmt is None
+
+
+def test_replace_root_write():
+    s, target = replace_ir_builder()
+    old_hash = s.mod["main"].__hash__()
+    sref = s.get_sref(s.mod["main"].body.block)
+    s.replace(sref, target)
+    # Check no copy and the new body equals to target
+    assert old_hash == s.mod["main"].__hash__()
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block, target)
+
+
+def test_replace_root_copy0():
+    s, target = replace_ir_builder(deep_copy=True)
+    old_hash = s.mod["main"].__hash__()
+    func_ref = s.mod["main"]
+    sref = s.get_sref(s.mod["main"].body.block)
+    s.replace(sref, target)
+    # Check the new body equals to target
+    assert old_hash != s.mod["main"].__hash__()
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block, target)
+    # Check the original func remains unchanged
+    assert old_hash == func_ref.__hash__()
+    assert not tvm_ffi.structural_equal(func_ref.body, target)
+
+
+def test_replace_root_copy1():
+    s, target = replace_ir_builder(deep_copy=True, realize=True)
+    old_hash = s.mod["main"].body.block.__hash__()
+    func_ref = s.mod["main"].body.block
+    sref = s.get_sref(s.mod["main"].body.block.body[0])
+    s.replace(sref, target)
+    # Check the new body equals to target
+    assert old_hash != s.mod["main"].body.block.__hash__()
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block.body[0], target)
+    # Check the original func remains unchanged
+    assert old_hash == func_ref.__hash__()
+    assert not tvm_ffi.structural_equal(func_ref.body, target)
+
+
+def test_replace_root_copy2():
+    s, target = replace_ir_builder(deep_copy=True)
+    old_hash = s.mod.functions.__hash__()
+    func_ref = s.mod.functions
+    sref = s.get_sref(s.mod["main"].body.block)
+    s.replace(sref, target)
+    # Check the new body equals to target
+    assert old_hash != s.mod.functions.__hash__()
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block, target)
+    # Check the original func remains unchanged
+    assert old_hash == func_ref.__hash__()
+    for _, v in func_ref.items():
+        assert not tvm_ffi.structural_equal(v.body.block, target)
+
+
+def test_replace_root_copy3():
+    s, target = replace_ir_builder(deep_copy=True)
+    old_hash = s.mod.__hash__()
+    func_ref = s.mod
+    sref = s.get_sref(s.mod["main"].body.block)
+    s.replace(sref, target)
+    # Check the new body equals to target
+    assert old_hash != s.mod.__hash__()
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block, target)
+    # Check the original func remains unchanged
+    assert old_hash == func_ref.__hash__()
+    assert not tvm_ffi.structural_equal(func_ref["main"].body.block, target)
+
+
+def test_replace_block_remap():
+    func = elementwise
+    s = tvm.s_tir.ScheduleState(func, debug_mask="all")
+    # The target stmt
+    target = matmul.body.block.body.body.body[0].block
+    sref = s.get_sref(s.mod["main"].body.block.body[0].body.body.block)
+    s.replace(sref, target, {sref.stmt: target})
+    sref_new = s.get_sref(s.mod["main"].body.block.body[0].body.body.block)
+    # Check the original sref has been remapped
+    assert sref.__hash__() == sref_new.__hash__()
+    tvm.ir.assert_structural_equal(sref.stmt, target)
+
+
+def test_replace_block_in_opaque_block():
+    s = replace_ir_builder_with_opaque()
+    root_hash = s.mod["main"].__hash__()
+    for_loop = s.mod["main"].body.block.body.body.block.body[1].then_case.block.body
+    sref = s.get_sref(for_loop)
+    new_for_loop = tirx.For(
+        loop_var=for_loop.loop_var,
+        min=0,
+        extent=128,
+        kind=tirx.ForKind.SERIAL,
+        body=tirx.Evaluate(0),
+        thread_binding=None,
+        annotations=None,
+    )
+    s.replace(sref, new_for_loop)
+    assert root_hash == s.mod["main"].__hash__()
+    tvm.ir.assert_structural_equal(sref.stmt, new_for_loop)
+
+
+def test_replace_ir_module():
+    s, target = replace_ir_builder_module(deep_copy=True)
+    old_hash = s.mod["main"].__hash__()
+    other_func_hash = s.mod["other"].__hash__()
+    func_ref = s.mod["main"]
+    sref = s.get_sref(s.mod["main"].body.block)
+    s.replace(sref, target)
+    # Check the new body equals to target
+    assert old_hash != s.mod["main"].__hash__()
+    tvm.ir.assert_structural_equal(s.mod["main"].body.block, target)
+    # Check the original func remains unchanged
+    assert old_hash == func_ref.__hash__()
+    assert not tvm_ffi.structural_equal(func_ref.body, target)
+    assert other_func_hash == s.mod["other"].__hash__()
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py

@@ -0,0 +1,901 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E501, E741, F401
+import sys
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import s_tir, tirx
+from tvm.s_tir.schedule.state import CachedFlags
+from tvm.script import tirx as T
+from tvm.tirx.stmt_functor import post_order_visit
+
+# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+# fmt: off
+
+@T.prim_func(s_tir=True)
+def elementwise(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128), "float32")
+    C = T.match_buffer(c, (128, 128), "float32")
+    B = T.sblock_alloc_buffer((128, 128), "float32")
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j in T.grid(128, 128):
+        with T.sblock("init"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = 0.0
+        for k in range(0, 128):
+            with T.sblock("update"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def block_in_opaque_block(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128), "float32")
+    B = T.match_buffer(b, (128, 128), "float32")
+    for i in range(128):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i)
+            T.reads([A[0:128, 0:128]])
+            T.writes([B[0:128, 0:128]])
+            B[vi, 0] = A[vi, 0]
+            if A[vi, 0] == 0.0:
+                with T.sblock("C"):
+                    T.reads([A[0:128, 0:128]])
+                    T.writes([B[0:128, 0:128]])
+                    for j in range(128):
+                        with T.sblock("D"):
+                            vj = T.axis.S(128, j)
+                            B[vi, vj] = A[vi, vj] * 3.0
+            else:
+                with T.sblock("E"):
+                    T.reads([A[0:128, 0:128]])
+                    T.writes([B[0:128, 0:128]])
+                    for j in range(128):
+                        with T.sblock("F"):
+                            vj = T.axis.S(128, j)
+                            B[vi, vj] = A[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def write_after_read(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def loop_carried_dependency(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128,))
+    B = T.match_buffer(b, (128,))
+    C = T.match_buffer(c, (128,))
+    for i in range(0, 128):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i)
+            B[vi] = A[vi] * 2.0
+        with T.sblock("C"):
+            vi = T.axis.S(128, i)
+            C[vi] = T.if_then_else(vi >= 1, B[vi - 1] + 1.0, 0.0, dtype="float32")
+
+
+@T.prim_func(s_tir=True)
+def concatenate_multi_producer(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128,))
+    B = T.match_buffer(b, (128,))
+    for i in range(0, 64):
+        with T.sblock("A_0"):
+            vi = T.axis.S(64, i)
+            A[vi] = vi + 1
+    for i in range(0, 64):
+        with T.sblock("A_1"):
+            vi = T.axis.S(64, i + 64)
+            A[vi] = vi + 2
+    for i in range(0, 128):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i)
+            B[vi] = A[vi] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def concatenate_multi_producer_uncovered(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128,))
+    B = T.match_buffer(b, (128,))
+    for i in range(0, 63):
+        with T.sblock("A_0"):
+            vi = T.axis.S(63, i)
+            A[vi] = vi + 1
+    for i in range(0, 64):
+        with T.sblock("A_1"):
+            vi = T.axis.S(64, i + 64)
+            A[vi] = vi + 2
+    for i in range(0, 128):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i)
+            B[vi] = A[vi] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def lca_at_loop(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128,))
+    B = T.match_buffer(b, (128,))
+    C = T.match_buffer(c, (128,))
+    for i in range(0, 128):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i)
+            B[vi] = A[vi] * 2.0
+        with T.sblock("C"):
+            vi = T.axis.S(128, i)
+            C[vi] = B[vi] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def multi_producer_consumer(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128,))
+    B = T.match_buffer(b, (128,))
+    for i in range(0, 64):
+        with T.sblock("A_0"):
+            vi = T.axis.S(64, i)
+            A[vi] = vi + 1
+    for i in range(0, 64):
+        with T.sblock("A_1"):
+            vi = T.axis.S(64, i + 64)
+            A[vi] = vi + 2
+    for i in range(0, 64):
+        with T.sblock("B_0"):
+            vi = T.axis.S(64, i)
+            B[vi] = A[vi] + 2.0
+    for i in range(0, 64):
+        with T.sblock("B_1"):
+            vi = T.axis.S(64, i + 64)
+            B[vi] = A[vi] + 3.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_affine_producer(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128), "float32")
+    C = T.match_buffer(c, (128, 128), "float32")
+    B = T.sblock_alloc_buffer((128, 128), "float32")
+    for i, j, k, l in T.grid(16, 2, 32, 16):
+        with T.sblock("B"):
+            vi = T.axis.S(128, i * 8 + j * 4 + k // 8)
+            vj = T.axis.S(128, k % 8 * 16 + l)
+            B[vi, vj] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_subblock(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128), "float32")
+    C = T.match_buffer(c, (128, 128), "float32")
+    B = T.sblock_alloc_buffer((128, 128), "float32")
+    for i, j in T.grid(32, 32):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([A[vi * 4 : vi * 4 + 4, vj * 4 : vj * 4 + 4]])
+            T.writes([B[vi * 4 : vi * 4 + 4, vj * 4 : vj * 4 + 4]])
+            for ii, jj in T.grid(4, 4):
+                with T.sblock("B_sub"):
+                    vi_i, vj_i = T.axis.remap("SS", [ii, jj])
+                    B[vi * 4 + vi_i, vj * 4 + vj_i] = A[vi * 4 + vi_i, vj * 4 + vj_i] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_subblock_uncovered(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128), "float32")
+    C = T.match_buffer(c, (128, 128), "float32")
+    B = T.sblock_alloc_buffer((128, 128), "float32")
+    for i, j in T.grid(32, 32):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([A[vi * 4 : vi * 4 + 2, vj * 4 : vj * 4 + 2]])
+            T.writes([B[vi * 4 : vi * 4 + 2, vj * 4 : vj * 4 + 2]])
+            for ii, jj in T.grid(2, 2):
+                with T.sblock("B_sub"):
+                    vi_i, vj_i = T.axis.remap("SS", [ii, jj])
+                    B[vi * 4 + vi_i, vj * 4 + vj_i] = A[vi * 4 + vi_i, vj * 4 + vj_i] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def bound_to_thread(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    B = T.sblock_alloc_buffer([128, 128], scope="shared")
+    for i in T.thread_binding(0, 128, thread="threadIdx.x"):
+        for j in T.serial(0, 128):
+            with T.sblock("B"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for j in T.serial(0, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vj, vi] = B[vj, vi] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def equal_ranked_threads(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    B = T.sblock_alloc_buffer([128, 128], scope="shared")
+    for i_o in T.thread_binding(0, 16, thread="threadIdx.x"):
+        for i_i in T.thread_binding(0, 8, thread="threadIdx.y"):
+            for j in T.serial(0, 128):
+                with T.sblock("B"):
+                    vi = T.axis.S(128, i_o * 8 + i_i)
+                    vj = T.axis.S(128, j)
+                    B[vi, vj] = A[vi, vj] * 2.0
+            for j in T.serial(0, 128):
+                with T.sblock("C"):
+                    vi = T.axis.S(128, i_o * 8 + i_i)
+                    vj = T.axis.S(128, j)
+                    C[vj, vi] = B[vj, vi] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def warp_memory(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    B = T.sblock_alloc_buffer([128, 4, 32], scope="warp")
+    for i_o in T.thread_binding(0, 4, thread="threadIdx.y"):
+        for i_i in T.thread_binding(0, 32, thread="threadIdx.x"):
+            for j in T.serial(0, 128):
+                with T.sblock("B"):
+                    warp_id, lane_id, vj = T.axis.remap("SSS", [i_o, i_i, j])
+                    B[vj, warp_id, lane_id] = A[warp_id * 32 + lane_id, vj] * 2.0
+            for j in T.serial(0, 128):
+                with T.sblock("C"):
+                    warp_id, lane_id, vj = T.axis.remap("SSS", [i_o, i_i, j])
+                    C[warp_id * 32 + lane_id, vj] = B[vj, warp_id, lane_id] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def warp_memory_negative(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    B = T.sblock_alloc_buffer([128, 4, 32], scope="warp")
+    for i_o in T.thread_binding(0, 4, thread="threadIdx.y"):
+        for i_i in T.thread_binding(0, 32, thread="threadIdx.x"):
+            for j in T.serial(0, 128):
+                with T.sblock("B"):
+                    warp_id, lane_id, vj = T.axis.remap("SSS", [i_o, i_i, j])
+                    B[vj, warp_id, lane_id] = A[warp_id * 32 + lane_id, vj] * 2.0
+            for i_o_prime in T.thread_binding(0, 4, thread="threadIdx.y"):
+                for j in T.serial(0, 128):
+                    with T.sblock("C"):
+                        _warp_id, warp_id, lane_id, vj = T.axis.remap(
+                            "SSSS", [i_o, i_i, i_o_prime, j]
+                        )
+                        C[warp_id * 32 + lane_id, vj] = B[vj, warp_id, lane_id] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def non_perfect_tiling_cache(a: T.handle, b: T.handle) -> None:
+    X = T.match_buffer(a, [224, 224], dtype="float32")
+    Y = T.match_buffer(b, [224, 224], dtype="float32")
+    cache = T.sblock_alloc_buffer([224, 224], dtype="float32")
+    for hh_0, ww_0 in T.grid(28, 28):
+        for ax0 in T.serial(0, 10):
+            for ax1 in T.serial(0, 10):
+                with T.sblock("cache"):
+                    h = T.axis.spatial(224, hh_0 * 8 - 1 + ax0)
+                    w = T.axis.spatial(224, ww_0 * 8 - 1 + ax1)
+                    T.where(
+                        1 <= hh_0 * 8 + ax0
+                        and hh_0 * 8 + ax0 < 225
+                        and 1 <= ww_0 * 8 + ax1
+                        and ww_0 * 8 + ax1 < 225
+                    )
+                    cache[h, w] = X[h, w]
+        for hh_1, ww_1, khh, kww in T.grid(8, 8, 3, 3):
+            with T.sblock("compute"):
+                h = T.axis.spatial(224, hh_0 * 8 + hh_1)
+                w = T.axis.spatial(224, ww_0 * 8 + ww_1)
+                kh, kw = T.axis.remap("RR", [khh, kww])
+                with T.init():
+                    Y[h, w] = 0.0
+                Y[h, w] = T.max(
+                    Y[h, w],
+                    T.if_then_else(
+                        T.likely(1 <= h + kh, dtype="bool")
+                        and T.likely(h + kh < 225, dtype="bool")
+                        and T.likely(1 <= w + kw, dtype="bool")
+                        and T.likely(w + kw < 225, dtype="bool"),
+                        cache[h + kh - 1, w + kw - 1],
+                        0.0,
+                        dtype="float32",
+                    ),
+                )
+
+
+@T.prim_func(s_tir=True)
+def uncovered_producer_region(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+    for i in range(120):
+        with T.sblock("producer"):
+            vi = T.axis.S((0, 120), i)
+            A[vi] = 1.0
+    for i in range(120):
+        with T.sblock("consumer"):
+            vi = T.axis.S((8, 128), i + 8)
+            B[vi] = A[vi]
+
+
+@T.prim_func(s_tir=True)
+def matmul_relu_padding(A: T.Buffer((127, 127), "float16"), B: T.Buffer((127, 127), "float16"), compute: T.Buffer((127, 127), "float32")) -> None:
+    # function attr dict
+    T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+    # body
+    # with T.sblock("root")
+    C = T.sblock_alloc_buffer([127, 127], dtype="float32")
+    A_reindex = T.sblock_alloc_buffer([128, 128], dtype="float16")
+    B_reindex = T.sblock_alloc_buffer([128, 128], dtype="float16")
+    C_reindex_shared = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="shared")
+    C_reindex_shared_wmma_accumulator = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="wmma.accumulator")
+    for ax0, ax1, ax2 in T.grid(128, 1, 128):
+        with T.sblock("A_reindex"):
+            v0, v1, v2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+            T.reads(A[v0, v2])
+            T.writes(A_reindex[v0, v2])
+            A_reindex[v0, v2] = T.if_then_else(v0 < 127 and v2 < 127, A[v0, v2], T.float16(0), dtype="float16")
+    for ax0, ax1, ax2 in T.grid(1, 128, 128):
+        with T.sblock("B_reindex"):
+            v0, v1, v2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+            T.reads(B[v2, v1])
+            T.writes(B_reindex[v2, v1])
+            B_reindex[v2, v1] = T.if_then_else(v2 < 127 and v1 < 127, B[v2, v1], T.float16(0), dtype="float16")
+    for ax0_0_0_ax1_0_0_fused in T.thread_binding(2, thread="blockIdx.y"):
+        for ax0_0_1_ax1_0_1_fused in T.thread_binding(1, thread="blockIdx.x"):
+            for ax0_0_2_ax1_0_2_fused in T.thread_binding(16, thread="threadIdx.y"):
+                for ax2_0_0, ax2_0_1, ax0_0_3, ax1_0_3, ax2_0_2, ax0_0_4, ax1_0_4 in T.grid(2, 2, 1, 2, 2, 1, 1):
+                    with T.sblock("C_o"):
+                        v0_o = T.axis.spatial(8, ax0_0_2_ax1_0_2_fused // 2 + ax0_0_3 + ax0_0_4)
+                        v1_o = T.axis.spatial(8, ax1_0_4 + ax0_0_0_ax1_0_0_fused * 4 + ax0_0_2_ax1_0_2_fused % 2 * 2 + ax1_0_3)
+                        v2_o = T.axis.reduce(8, ax2_0_0 * 4 + ax2_0_1 * 2 + ax2_0_2)
+                        T.reads(A_reindex[v0_o * 16 : v0_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16], B_reindex[v2_o * 16 : v2_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16])
+                        T.writes(C_reindex_shared_wmma_accumulator[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16])
+                        T.sblock_attr({"meta_schedule.auto_tensorize":"wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init":"wmma_fill_16x16x16_f32", "warp_execution":1})
+                        with T.init():
+                            for ax0_1, ax1_1 in T.grid(16, 16):
+                                with T.sblock("C_init"):
+                                    v0_i_init, v1_i_init = T.axis.remap("SS", [ax0_1, ax1_1])
+                                    T.reads()
+                                    T.writes(C_reindex_shared_wmma_accumulator[v0_o * 16 + v0_i_init, v1_o * 16 + v1_i_init])
+                                    C_reindex_shared_wmma_accumulator[v0_o * 16 + v0_i_init, v1_o * 16 + v1_i_init] = T.float32(0)
+                        for ax0_1, ax1_1, ax2_1 in T.grid(16, 16, 16):
+                            with T.sblock("C"):
+                                v0_i, v1_i, v2_i = T.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
+                                T.reads(C_reindex_shared_wmma_accumulator[v0_o * 16 + v0_i, v1_o * 16 + v1_i], A_reindex[v0_o * 16 + v0_i, v2_o * 16 + v2_i], B_reindex[v2_o * 16 + v2_i, v1_o * 16 + v1_i])
+                                T.writes(C_reindex_shared_wmma_accumulator[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
+                                T.sblock_attr({"meta_schedule.tiling_structure":"SSSRRSRS"})
+                                C_reindex_shared_wmma_accumulator[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = C_reindex_shared_wmma_accumulator[v0_o * 16 + v0_i, v1_o * 16 + v1_i] + T.cast(A_reindex[v0_o * 16 + v0_i, v2_o * 16 + v2_i], "float32") * T.cast(B_reindex[v2_o * 16 + v2_i, v1_o * 16 + v1_i], "float32")
+                for ax0, ax1 in T.grid(16, 32):
+                    with T.sblock("C_reindex_shared_wmma.accumulator"):
+                        v0 = T.axis.spatial(128, ax0_0_2_ax1_0_2_fused // 2 * 16 + ax0)
+                        v1 = T.axis.spatial(128, ax0_0_0_ax1_0_0_fused * 64 + ax0_0_2_ax1_0_2_fused % 2 * 32 + ax1)
+                        T.reads(C_reindex_shared_wmma_accumulator[v0, v1])
+                        T.writes(C_reindex_shared[v0, v1])
+                        C_reindex_shared[v0, v1] = C_reindex_shared_wmma_accumulator[v0, v1]
+            for ax0, ax1 in T.grid(128, 64):
+                with T.sblock("C_reindex_shared"):
+                    v0 = T.axis.spatial(128, ax0)
+                    v1 = T.axis.spatial(128, ax0_0_0_ax1_0_0_fused * 64 + ax1)
+                    T.where(ax0 < 127 and ax0_0_0_ax1_0_0_fused * 64 + ax1 < 127)
+                    T.reads(C_reindex_shared[v0, v1])
+                    T.writes(C[v0, v1])
+                    T.sblock_attr({"meta_schedule.cooperative_fetch":3})
+                    C[v0, v1] = C_reindex_shared[v0, v1]
+    for i0, i1 in T.grid(127, 127):
+        with T.sblock("compute"):
+            i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+            T.reads(C[i0_1, i1_1])
+            T.writes(compute[i0_1, i1_1])
+            compute[i0_1, i1_1] = T.max(C[i0_1, i1_1], T.float32(0))
+
+
+@T.prim_func(s_tir=True)
+def splitted_square_sum_with_predicate(
+    A: T.Buffer((1, 7, 7, 512), "float32"), B: T.Buffer((1, 1, 1, 512), "float32")
+) -> None:
+    for i0_i1_i2_i3_0_fused, ax0, ax1, ax2, ax3 in T.grid(2, 1, 1, 1, 256):
+        for ax4_ax5_fused_0, ax4_ax5_fused_1 in T.grid(1, 256):
+            with T.sblock("B"):
+                T.where(ax4_ax5_fused_0 * 256 + ax4_ax5_fused_1 < 49)
+                ax0_1, ax1_1, ax2_1 = T.axis.remap("SSS", [ax0, ax1, ax2])
+                ax3_1 = T.axis.spatial(512, i0_i1_i2_i3_0_fused * 256 + ax3)
+                rv0 = T.axis.reduce(7, (ax4_ax5_fused_0 * 256 + ax4_ax5_fused_1) // 7)
+                rv1 = T.axis.reduce(7, (ax4_ax5_fused_0 * 256 + ax4_ax5_fused_1) % 7)
+                T.reads(A[ax0_1, ax1_1 * 7 + rv0, ax2_1 * 7 + rv1, ax3_1])
+                T.writes(B[ax0_1, ax1_1, ax2_1, ax3_1])
+                with T.init():
+                    B[ax0_1, ax1_1, ax2_1, ax3_1] = T.float32(0)
+                B[ax0_1, ax1_1, ax2_1, ax3_1] += A[ax0_1, ax1_1 * 7 + rv0, ax2_1 * 7 + rv1, ax3_1]
+
+
+# pylint: enable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
+# fmt: on
+
+
+def _get_sblock(s: s_tir.ScheduleState, name_hint: str) -> s_tir.StmtSRef:
+    result = None
+
+    def f_visit(node):
+        nonlocal result
+        if isinstance(node, tvm.tirx.SBlock) and node.name_hint == name_hint:
+            result = node
+
+    func = s.mod["main"]
+    post_order_visit(func.body, f_visit)
+    assert result is not None and isinstance(result, tvm.tirx.SBlock)
+    return s.get_sref(result)
+
+
+def test_elementwise():
+    s = tvm.s_tir.ScheduleState(elementwise, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_matmul():
+    s = tvm.s_tir.ScheduleState(matmul, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "init")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "update")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_block_in_opaque_block():
+    s = tvm.s_tir.ScheduleState(block_in_opaque_block, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "E")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "F")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_write_after_read():
+    s = tvm.s_tir.ScheduleState(write_after_read, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=False,
+    )
+    # pylint: enable=protected-access
+
+
+def test_loop_carried_dependency():
+    s = tvm.s_tir.ScheduleState(loop_carried_dependency, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=False,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=False,
+    )
+    # pylint: enable=protected-access
+
+
+def test_concatenate_multi_producer_covered():  # pylint: disable=invalid-name
+    s = tvm.s_tir.ScheduleState(concatenate_multi_producer, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "A_0")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "A_1")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_concatenate_multi_producer_uncovered():  # pylint: disable=invalid-name
+    s = tvm.s_tir.ScheduleState(concatenate_multi_producer_uncovered, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "A_0")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "A_1")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=False,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=False,
+    )
+    # pylint: enable=protected-access
+
+
+def test_lca_at_loop():
+    s = tvm.s_tir.ScheduleState(lca_at_loop, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_multi_producer_consumer():
+    s = tvm.s_tir.ScheduleState(multi_producer_consumer, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "A_0")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "A_1")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B_0")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B_1")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_elementwise_affine_producer():
+    s = tvm.s_tir.ScheduleState(elementwise_affine_producer, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_subblock():
+    s = tvm.s_tir.ScheduleState(elementwise_subblock, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B_sub")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_subblock_uncovered():
+    s = tvm.s_tir.ScheduleState(elementwise_subblock_uncovered, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=False,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B_sub")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=False,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_thread_binding():
+    s = tvm.s_tir.ScheduleState(bound_to_thread, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_equal_ranked_threads():
+    s = tvm.s_tir.ScheduleState(equal_ranked_threads, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_warp_memory():
+    s = tvm.s_tir.ScheduleState(warp_memory, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_warp_memory_negative():
+    s = tvm.s_tir.ScheduleState(warp_memory_negative, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "root")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=False,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "C")) == CachedFlags(
+        affine_binding=True,
+        region_cover=False,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_non_perfect_tiling_cache():
+    s = tvm.s_tir.ScheduleState(non_perfect_tiling_cache, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "cache")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    assert s._get_cached_flags(_get_sblock(s, "compute")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_uncovered_producer_region():
+    s = tvm.s_tir.ScheduleState(uncovered_producer_region, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "consumer")) == CachedFlags(
+        affine_binding=True,
+        region_cover=False,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_matmul_relu_padding():
+    s = tvm.s_tir.ScheduleState(matmul_relu_padding, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "C_reindex_shared")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+def test_splitted_square_sum_with_predicate():
+    s = tvm.s_tir.ScheduleState(splitted_square_sum_with_predicate, debug_mask="all")
+    # pylint: disable=protected-access
+    assert s._get_cached_flags(_get_sblock(s, "B")) == CachedFlags(
+        affine_binding=True,
+        region_cover=True,
+        stage_pipeline=True,
+    )
+    # pylint: enable=protected-access
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_storage_align.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=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import pytest
+
+import tvm
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# fmt: off
+# pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name
+
+@T.prim_func(s_tir=True)
+def element_wise(a: T.handle, c: T.handle) -> None:
+    C = T.match_buffer(c, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    A = T.match_buffer(a, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    # body
+    with T.sblock("root"):
+        T.reads([])
+        T.writes([])
+        B = T.sblock_alloc_buffer([128, 128], elem_offset=0, align=64, offset_factor=1)
+        for i0 in T.serial(0, 128):
+            for ax1 in T.serial(0, 128):
+                with T.sblock("B"):
+                    vi, vj = T.axis.remap("SS", [i0, ax1])
+                    T.reads([A[vi, vj]])
+                    T.writes([B[vi, vj]])
+                    B[vi, vj] = (A[vi, vj]*T.float32(2))
+            for i1 in T.serial(0, 128):
+                with T.sblock("C"):
+                    vi_1, vj_1 = T.axis.remap("SS", [i0, i1])
+                    T.reads([B[vi_1, vj_1]])
+                    T.writes([C[vi_1, vj_1]])
+                    C[vi_1, vj_1] = (B[vi_1, vj_1] + T.float32(1))
+
+
+@T.prim_func(s_tir=True)
+def element_wise_storage_align(a: T.handle, c: T.handle) -> None:
+    C = T.match_buffer(c, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    A = T.match_buffer(a, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    # body
+    with T.sblock("root"):
+        T.reads([])
+        T.writes([])
+        B = T.sblock_alloc_buffer([128, 128], elem_offset=0, align=64, offset_factor=1)
+        for i0 in T.serial(0, 128):
+            for ax1 in T.serial(0, 128):
+                with T.sblock("B"):
+                    vi, vj = T.axis.remap("SS", [i0, ax1])
+                    T.reads([A[vi, vj]])
+                    T.writes([B[vi, vj]])
+                    T.sblock_attr({"buffer_dim_align":[[0, 0, 128, 127]]})
+                    B[vi, vj] = (A[vi, vj]*T.float32(2))
+            for i1 in T.serial(0, 128):
+                with T.sblock("C"):
+                    vi_1, vj_1 = T.axis.remap("SS", [i0, i1])
+                    T.reads([B[vi_1, vj_1]])
+                    T.writes([C[vi_1, vj_1]])
+                    C[vi_1, vj_1] = (B[vi_1, vj_1] + T.float32(1))
+
+
+@T.prim_func(s_tir=True)
+def element_wise_invalid_annotation(a: T.handle, c: T.handle) -> None:
+    C = T.match_buffer(c, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    A = T.match_buffer(a, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    # body
+    with T.sblock("root"):
+        T.reads([])
+        T.writes([])
+        B = T.sblock_alloc_buffer([128, 128], elem_offset=0, align=64, offset_factor=1)
+        for i0 in T.serial(0, 128):
+            for ax1 in T.serial(0, 128):
+                with T.sblock("B"):
+                    T.sblock_attr({"buffer_dim_align": [0]})
+                    vi, vj = T.axis.remap("SS", [i0, ax1])
+                    T.reads([A[vi, vj]])
+                    T.writes([B[vi, vj]])
+                    B[vi, vj] = (A[vi, vj]*T.float32(2))
+            for i1 in T.serial(0, 128):
+                with T.sblock("C"):
+                    vi_1, vj_1 = T.axis.remap("SS", [i0, i1])
+                    T.reads([B[vi_1, vj_1]])
+                    T.writes([C[vi_1, vj_1]])
+                    C[vi_1, vj_1] = (B[vi_1, vj_1] + T.float32(1))
+
+
+use_block_name = tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True})
+
+def test_storage_align(use_block_name):
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    B = 'B' if use_block_name else s.get_sblock("B")
+    s.storage_align(B, 0, axis=0, factor=128, offset=127)
+    assert_structural_equal_ignore_global_symbol(element_wise_storage_align, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_storage_align_update():
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    B = s.get_sblock("B")
+    s.storage_align(B, 0, axis=0, factor=128, offset=0)
+    s.storage_align(B, 0, axis=0, factor=128, offset=127)
+    assert_structural_equal_ignore_global_symbol(element_wise_storage_align, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_storage_align_invalid_factor1():
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    B = s.get_sblock("B")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.storage_align(B, 0, axis=0, factor=0, offset=127)
+
+
+def test_storage_align_invalid_factor2():
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    B = s.get_sblock("B")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.storage_align(B, 0, axis=0, factor=-1, offset=127)
+
+
+def test_storage_align_invalid_buffer():
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    C = s.get_sblock("C")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.storage_align(C, 0, axis=0, factor=128, offset=127)
+
+
+def test_storage_align_invalid_buffer_index():
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    B = s.get_sblock("B")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.storage_align(B, 2, axis=0, factor=128, offset=127)
+
+
+def test_storage_align_invalid_axis():
+    func = element_wise
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    B = s.get_sblock("B")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.storage_align(B, 0, axis=2, factor=128, offset=127)
+
+
+def test_storage_align_invalid_annotation():
+    func = element_wise_invalid_annotation
+    s = tvm.s_tir.Schedule(func, debug_mask='all')
+    B = s.get_sblock("B")
+    with pytest.raises(tvm.s_tir.ScheduleError):
+        s.storage_align(B, 0, axis=2, factor=128, offset=127)
+
+
+if __name__ == "__main__":
+    test_storage_align()
+    test_storage_align_update()
+    test_storage_align_invalid_factor1()
+    test_storage_align_invalid_factor2()
+    test_storage_align_invalid_buffer()
+    test_storage_align_invalid_buffer_index()
+    test_storage_align_invalid_axis()
+    test_storage_align_invalid_annotation()

tests/python/s_tir/schedule/test_tir_schedule_tensorize.py

@@ -0,0 +1,961 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E501
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import te, tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.s_tir.tensor_intrin.arm_cpu import (
+    DP4A_S8S8S32_INTRIN,
+    DP4A_S8U8S32_INTRIN,
+    DP4A_U8S8S32_INTRIN,
+    DP4A_U8U8U32_INTRIN,
+    ARM_DOT_4x4_i8_NEON_INTRIN,
+    ARM_DOT_4x4_i8_SDOT_INTRIN,
+)
+from tvm.s_tir.tensor_intrin.hexagon import VDMPY_i16i16i32_INTRIN, VRMPY_u8u8i32_INTRIN
+from tvm.s_tir.tensor_intrin.rocm import AMDGPU_SDOT4_INTRIN
+from tvm.s_tir.tensor_intrin.x86 import AVX512_DOT_16x4_INTRIN, VNNI_DOT_16x4_INTRIN
+from tvm.script import tirx as T
+
+# fmt: off
+# pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks
+
+@T.prim_func(s_tir=True)
+def mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), align=64, offset_factor=1)
+    B = T.match_buffer(b, (16, 16), align=64, offset_factor=1)
+    C = T.match_buffer(c, (16, 16), align=64, offset_factor=1)
+
+    with T.sblock("root"):
+        T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16])
+        T.writes(C[0 : 16, 0 : 16])
+        for i, j, k in T.grid(16, 16, 16):
+            with T.sblock("update"):
+                vii, vjj, vkk = T.axis.remap("SSR", [i, j, k])
+                C[vii, vjj] = C[vii, vjj] + A[vii, vkk] * B[vjj, vkk]
+
+
+@T.prim_func(s_tir=True)
+def mma_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), align=64, offset_factor=1)
+    B = T.match_buffer(b, (16, 16), align=64, offset_factor=1)
+    C = T.match_buffer(c, (16, 16), align=64, offset_factor=1)
+
+    with T.sblock("root"):
+        T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16])
+        T.writes(C[0 : 16, 0 : 16])
+        T.evaluate(
+            T.tvm_mma_sync(
+                C.data,
+                C.elem_offset // 256,
+                A.data,
+                A.elem_offset // 256,
+                B.data,
+                B.elem_offset // 256,
+                C.data,
+                C.elem_offset // 256,
+                dtype="handle",
+            )
+        )
+
+
+@T.prim_func(s_tir=True)
+def dot_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (4,))
+    B = T.match_buffer(b, (4,))
+    C = T.match_buffer(c, ())
+
+    with T.sblock("root"):
+        T.reads(C[()], A[0 : 4], B[0 : 4])
+        T.writes(C[()])
+        for i in range(0, 4):
+            with T.sblock("update"):
+                vi = T.axis.remap("R", [i])
+                C[()] = C[()] + A[vi] * B[vi]
+
+
+@T.prim_func(s_tir=True)
+def dot_product_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (4,), offset_factor=1)
+    B = T.match_buffer(b, (4,), offset_factor=1)
+    C = T.match_buffer(c, (), offset_factor=1)
+
+    with T.sblock("root"):
+        T.reads(C[()], A[0 : 4], B[0 : 4])
+        T.writes(C[()])
+        T.evaluate(
+            T.call_extern(
+                "vec4add",
+                C.data,
+                C.elem_offset,
+                A.data,
+                A.elem_offset,
+                B.data,
+                B.elem_offset,
+                dtype="int32",
+            )
+        )
+
+
+@T.prim_func(s_tir=True)
+def dot_product_intrin_annotated(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (4,), offset_factor=1)
+    B = T.match_buffer(b, (4,), offset_factor=1)
+    C = T.match_buffer(c, (), offset_factor=1)
+
+    with T.sblock("root"):
+        T.reads(C[()], A[0 : 4], B[0 : 4])
+        T.writes(C[()])
+        T.sblock_attr({"test_annotation": True})
+        T.evaluate(
+            T.call_extern(
+                "vec4add",
+                C.data,
+                C.elem_offset,
+                A.data,
+                A.elem_offset,
+                B.data,
+                B.elem_offset,
+                dtype="int32",
+            )
+        )
+
+
+@T.prim_func(s_tir=True)
+def outer_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 1), offset_factor=1)
+    B = T.match_buffer(b, (16, 1), offset_factor=1)
+    C = T.match_buffer(c, (16, 16), offset_factor=1)
+
+    with T.sblock("root"):
+        T.reads(
+            C[0 : 16, 0 : 16],
+            A[0 : 16, 0 : 1],
+            B[0 : 16, 0 : 1],
+        )
+        T.writes(C[0 : 16, 0 : 16])
+        for i, j in T.grid(16, 16):
+            with T.sblock("update"):
+                vii, vjj = T.axis.remap("SS", [i, j])
+                C[vii, vjj] = C[vii, vjj] + A[vii, 0] * B[vjj, 0]
+
+
+@T.prim_func(s_tir=True)
+def outer_product_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 1), offset_factor=1)
+    B = T.match_buffer(b, (16, 1), offset_factor=1)
+    C = T.match_buffer(c, (16, 16), offset_factor=1)
+
+    with T.sblock("root"):
+        T.reads(
+            C[0 : 16, 0 : 16],
+            A[0 : 16, 0 : 1],
+            B[0 : 16, 0 : 1],
+        )
+        T.writes(C[0 : 16, 0 : 16])
+        T.evaluate(
+            T.call_extern(
+                "outer_product",
+                C.data,
+                C.elem_offset,
+                A.data,
+                A.elem_offset,
+                B.data,
+                B.elem_offset,
+                dtype="int32",
+            )
+        )
+
+
+@T.prim_func(s_tir=True)
+def matmul(
+    A: T.Buffer((128, 128), "float32"),
+    B: T.Buffer((128, 128), "float32"),
+    C: T.Buffer((128, 128), "float32"),
+) -> None:
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = T.float32(0)
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def tensorized_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    C = T.match_buffer(c, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    B = T.match_buffer(b, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    A = T.match_buffer(a, [128, 128], elem_offset=0, align=64, offset_factor=1)
+
+    for i_outer, j_outer in T.grid(8, 8):
+        for i_inner_init, j_inner_init in T.grid(16, 16):
+            with T.sblock("init"):
+                vi_init = T.axis.S(128, ((i_outer * 16) + i_inner_init))
+                vj_init = T.axis.S(128, ((j_outer * 16) + j_inner_init))
+                C[vi_init, vj_init] = T.float32(0)
+        for k_outer in T.grid(8):
+            with T.sblock("update"):
+                vi, vj, vk = T.axis.remap("SSR", [i_outer, j_outer, k_outer])
+                T.reads(
+                    [
+                        C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
+                        A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
+                        B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
+                    ]
+                )
+                T.writes(C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+                A_elem_offset = T.int32()
+                B_elem_offset = T.int32()
+                C_elem_offset = T.int32()
+                A_sub = T.match_buffer(
+                    A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
+                    [16, 16],
+                    elem_offset=A_elem_offset,
+                )
+                B_sub = T.match_buffer(
+                    B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
+                    [16, 16],
+                    elem_offset=B_elem_offset,
+                )
+                C_sub = T.match_buffer(
+                    C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
+                    [16, 16],
+                    elem_offset=C_elem_offset,
+                )
+                T.evaluate(
+                    T.tvm_mma_sync(
+                        C_sub.data,
+                        T.floordiv(C_sub.elem_offset, 256),
+                        A_sub.data,
+                        T.floordiv(A_sub.elem_offset, 256),
+                        B_sub.data,
+                        T.floordiv(B_sub.elem_offset, 256),
+                        C_sub.data,
+                        T.floordiv(C_sub.elem_offset, 256),
+                        dtype="handle",
+                    )
+                )
+
+
+@T.prim_func(s_tir=True)
+def batch_matmul(
+    A: T.Buffer((16, 128, 128), "float32"),
+    B: T.Buffer((16, 128, 128), "float32"),
+    C: T.Buffer((16, 128, 128), "float32"),
+) -> None:
+    for n, i, j in T.grid(16, 128, 128):
+        with T.sblock("init"):
+            vn, vi, vj = T.axis.remap("SSS", [n, i, j])
+            C[vn, vi, vj] = T.float32(0)
+
+    for n, i, j, k in T.grid(16, 128, 128, 128):
+        with T.sblock("update"):
+            vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])
+            C[vn, vi, vj] = C[vn, vi, vj] + A[vn, vi, vk] * B[vn, vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def tensorized_batch_matmul_mma(
+    A: T.Buffer((16, 128, 128), "float32"),
+    B: T.Buffer((16, 128, 128), "float32"),
+    C: T.Buffer((16, 128, 128), "float32"),
+) -> None:
+    for n, i, j in T.grid(16, 128, 128):
+        with T.sblock("init"):
+            vn, vi, vj = T.axis.remap("SSS", [n, i, j])
+            T.reads()
+            T.writes(C[vn, vi, vj])
+            C[vn, vi, vj] = T.float32(0)
+    for n in range(0, 16):
+        for i, j, k in T.grid(8, 8, 8):
+            with T.sblock("update"):
+                vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])
+                T.reads(
+                    C[vn : vn + 1, vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
+                    A[vn : vn + 1, vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
+                    B[vn : vn + 1, vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
+                )
+                T.writes(C[vn : vn + 1, vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+                A_elem_offset = T.int32()
+                B_elem_offset = T.int32()
+                C_elem_offset = T.int32()
+                A_sub = T.match_buffer(
+                    A[vn : vn + 1, vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
+                    (16, 16),
+                    elem_offset=A_elem_offset,
+                )
+                B_sub = T.match_buffer(
+                    B[vn : vn + 1, vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
+                    (16, 16),
+                    elem_offset=B_elem_offset,
+                )
+                C_sub = T.match_buffer(
+                    C[vn : vn + 1, vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
+                    (16, 16),
+                    elem_offset=C_elem_offset,
+                )
+                T.evaluate(
+                    T.tvm_mma_sync(
+                        C_sub.data,
+                        T.floordiv(C_sub.elem_offset, 256),
+                        A_sub.data,
+                        T.floordiv(A_sub.elem_offset, 256),
+                        B_sub.data,
+                        T.floordiv(B_sub.elem_offset, 256),
+                        C_sub.data,
+                        T.floordiv(C_sub.elem_offset, 256),
+                        dtype="handle",
+                    )
+                )
+
+
+@T.prim_func(s_tir=True)
+def tensorized_batch_matmul_dot_product(
+    A: T.Buffer((16, 128, 128), "float32"),
+    B: T.Buffer((16, 128, 128), "float32"),
+    C: T.Buffer((16, 128, 128), "float32"),
+) -> None:
+    for n, i, j in T.grid(16, 128, 128):
+        with T.sblock("init"):
+            vn, vi, vj = T.axis.remap("SSS", [n, i, j])
+            T.reads()
+            T.writes(C[vn, vi, vj])
+            C[vn, vi, vj] = T.float32(0)
+    for n, i, j, k_0 in T.grid(16, 128, 128, 32):
+        with T.sblock("blockized_update"):
+            vn, vi, vj, vko = T.axis.remap("SSSR", [n, i, j, k_0])
+            T.reads(
+                C[vn, vi, vj], A[vn, vi, vko * 4 : vko * 4 + 4], B[vn, vj, vko * 4 : vko * 4 + 4]
+            )
+            T.writes(C[vn, vi, vj])
+            A_1 = T.match_buffer(
+                A[vn, vi, vko * 4 : vko * 4 + 4], [4], dtype="float32", offset_factor=1
+            )
+            B_1 = T.match_buffer(
+                B[vn, vj, vko * 4 : vko * 4 + 4], [4], dtype="float32", offset_factor=1
+            )
+            C_1 = T.match_buffer(C[vn, vi, vj], [], dtype="float32", offset_factor=1)
+            T.evaluate(
+                T.call_extern(
+                    "vec4add",
+                    C_1.data,
+                    C_1.elem_offset,
+                    A_1.data,
+                    A_1.elem_offset,
+                    B_1.data,
+                    B_1.elem_offset,
+                    dtype="int32",
+                )
+            )
+
+
+@T.prim_func(s_tir=True)
+def tensorized_batch_matmul_outer_product(
+    A: T.Buffer((16, 128, 128), "float32"),
+    B: T.Buffer((16, 128, 128), "float32"),
+    C: T.Buffer((16, 128, 128), "float32"),
+) -> None:
+    for n, i, j in T.grid(16, 128, 128):
+        with T.sblock("init"):
+            vn, vi, vj = T.axis.remap("SSS", [n, i, j])
+            T.reads()
+            T.writes(C[vn, vi, vj])
+            C[vn, vi, vj] = T.float32(0)
+    for n, i_0, j_0, k in T.grid(16, 8, 8, 128):
+        with T.sblock("blockized_update"):
+            vn, vio, vjo, vk = T.axis.remap("SSSR", [n, i_0, j_0, k])
+            T.reads(
+                C[vn, vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16],
+                A[vn, vio * 16 : vio * 16 + 16, vk],
+                B[vn, vjo * 16 : vjo * 16 + 16, vk],
+            )
+            T.writes(C[vn, vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16])
+            A_1 = T.match_buffer(A[vn, vio * 16 : vio * 16 + 16, vk], [16, 1], dtype="float32", offset_factor=1)
+            B_1 = T.match_buffer(B[vn, vjo * 16 : vjo * 16 + 16, vk], [16, 1], dtype="float32", offset_factor=1
+            )
+            C_1 = T.match_buffer(
+                C[vn, vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16], [16, 16], dtype="float32", offset_factor=1
+            )
+            T.evaluate(
+                T.call_extern("outer_product", C_1.data, C_1.elem_offset, A_1.data, A_1.elem_offset,
+                              B_1.data, B_1.elem_offset, dtype="int32"
+                )
+            )
+
+
+@T.prim_func(s_tir=True)
+def annotated_mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), align=64, offset_factor=1)
+    B = T.match_buffer(b, (16, 16), align=64, offset_factor=1)
+    C = T.match_buffer(c, (16, 16), align=64, offset_factor=1)
+
+    with T.sblock("root"):
+        T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16])
+        T.writes(C[0 : 16, 0 : 16])
+        for i, j, k in T.grid(16, 16, 16):
+            with T.sblock("update"):
+                T.sblock_attr({"test_annotation": True})
+                vii, vjj, vkk = T.axis.remap("SSR", [i, j, k])
+                C[vii, vjj] = C[vii, vjj] + A[vii, vkk] * B[vjj, vkk]
+
+
+@T.prim_func(s_tir=True)
+def annotated_matmul(
+    A: T.Buffer((128, 128), "float32"),
+    B: T.Buffer((128, 128), "float32"),
+    C: T.Buffer((128, 128), "float32"),
+) -> None:
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            T.sblock_attr({"test_annotation": True})
+            with T.init():
+                C[vi, vj] = T.float32(0)
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def annotated_tensorized_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    C = T.match_buffer(c, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    B = T.match_buffer(b, [128, 128], elem_offset=0, align=64, offset_factor=1)
+    A = T.match_buffer(a, [128, 128], elem_offset=0, align=64, offset_factor=1)
+
+    for i_outer, j_outer in T.grid(8, 8):
+        for i_inner_init, j_inner_init in T.grid(16, 16):
+            with T.sblock("init"):
+                vi_init = T.axis.S(128, ((i_outer * 16) + i_inner_init))
+                vj_init = T.axis.S(128, ((j_outer * 16) + j_inner_init))
+                T.sblock_attr({"test_annotation": True})
+                C[vi_init, vj_init] = T.float32(0)
+        for k_outer in T.grid(8):
+            with T.sblock("update"):
+                vi, vj, vk = T.axis.remap("SSR", [i_outer, j_outer, k_outer])
+                T.reads(
+                    [
+                        C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
+                        A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
+                        B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
+                    ]
+                )
+                T.writes(C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+                A_elem_offset = T.int32()
+                B_elem_offset = T.int32()
+                C_elem_offset = T.int32()
+                A_sub = T.match_buffer(
+                    A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
+                    [16, 16],
+                    elem_offset=A_elem_offset,
+                )
+                B_sub = T.match_buffer(
+                    B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
+                    [16, 16],
+                    elem_offset=B_elem_offset,
+                )
+                C_sub = T.match_buffer(
+                    C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
+                    [16, 16],
+                    elem_offset=C_elem_offset,
+                )
+                T.evaluate(
+                    T.tvm_mma_sync(
+                        C_sub.data,
+                        T.floordiv(C_sub.elem_offset, 256),
+                        A_sub.data,
+                        T.floordiv(A_sub.elem_offset, 256),
+                        B_sub.data,
+                        T.floordiv(B_sub.elem_offset, 256),
+                        C_sub.data,
+                        T.floordiv(C_sub.elem_offset, 256),
+                        dtype="handle",
+                    )
+                )
+
+
+# fmt: off
+# pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks
+
+tirx.TensorIntrin.register("test_mma_intrin", mma_desc, mma_intrin)
+tirx.TensorIntrin.register("test_annotated_mma_intrin", annotated_mma_desc, mma_intrin)
+tirx.TensorIntrin.register("test_dot_product_intrin", dot_product_desc, dot_product_intrin)
+tirx.TensorIntrin.register("test_outer_product_intrin", outer_product_desc, outer_product_intrin)
+tirx.TensorIntrin.register("test_dot_product_intrin_annotated", dot_product_desc, dot_product_intrin_annotated)
+
+
+def test_tensorize_matmul():
+    func = matmul
+    # schedule
+    s = tvm.s_tir.Schedule(func, debug_mask="all")
+    update = s.get_sblock("update")
+    i, j, k = s.get_loops(update)
+    io, ii = s.split(i, factors=[None, 16])
+    jo, ji = s.split(j, factors=[None, 16])
+    ko, ki = s.split(k, factors=[None, 16])
+    s.reorder(io, jo, ko, ii, ji, ki)
+    s.decompose_reduction(update, ko)
+    s.tensorize(ii, "test_mma_intrin")
+    assert_structural_equal_ignore_global_symbol(tensorized_matmul, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_tensorize_batch_matmul():
+    func = batch_matmul
+    s = tvm.s_tir.Schedule(func, debug_mask="all")
+    update = s.get_sblock("update")
+    _, i, j, k = s.get_loops(update)
+    io, ii = s.split(i, factors=[None, 16])
+    jo, ji = s.split(j, factors=[None, 16])
+    ko, ki = s.split(k, factors=[None, 16])
+    s.reorder(io, jo, ko, ii, ji, ki)
+    s.tensorize(ii, "test_mma_intrin")
+    assert_structural_equal_ignore_global_symbol(tensorized_batch_matmul_mma, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=batch_matmul)
+
+
+def test_tensorize_dot_product():
+    func = batch_matmul
+    s = tvm.s_tir.Schedule(func, debug_mask="all")
+    C = s.get_sblock("update")
+    _, _, _, k = s.get_loops(C)
+    _, ki = s.split(k, factors=[None, 4])
+    s.tensorize(ki, "test_dot_product_intrin")
+    assert_structural_equal_ignore_global_symbol(tensorized_batch_matmul_dot_product, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_tensorize_outer_product():
+    func = batch_matmul
+    s = tvm.s_tir.Schedule(func, debug_mask="all")
+    C = s.get_sblock("update")
+    _, i, j, k = s.get_loops(C)
+    io, ii = s.split(i, factors=[None, 16])
+    jo, ji = s.split(j, factors=[None, 16])
+    s.reorder(io, jo, k, ii, ji)
+    s.tensorize(ii, "test_outer_product_intrin")
+    assert_structural_equal_ignore_global_symbol(tensorized_batch_matmul_outer_product, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_tensorize_with_annotation():
+    func = annotated_matmul
+    s = tvm.s_tir.Schedule(func, debug_mask="all")
+    update = s.get_sblock("update")
+    i, j, k = s.get_loops(update)
+    io, ii = s.split(i, factors=[None, 16])
+    jo, ji = s.split(j, factors=[None, 16])
+    ko, ki = s.split(k, factors=[None, 16])
+    s.reorder(io, jo, ko, ii, ji, ki)
+    s.decompose_reduction(update, ko)
+    s.tensorize(ii, "test_annotated_mma_intrin")
+    assert_structural_equal_ignore_global_symbol(annotated_tensorized_matmul, s.mod["main"])
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def test_tensorize_intrinsic_with_annotation():
+    func = matmul
+    s = tvm.s_tir.Schedule(func, debug_mask="all")
+    update = s.get_sblock("update")
+    _, _, k = s.get_loops(update)
+    ko, ki = s.split(k, factors=[None, 4])
+    s.decompose_reduction(update, ko)
+    s.tensorize(ki, "test_dot_product_intrin_annotated")
+
+    b = s.get(s.get_sblock("update_update_o"))
+    assert b.annotations["test_annotation"] == T.bool(True)
+    verify_trace_roundtrip(sch=s, mod=func)
+
+
+def get_matmul_packed(m, n, k, lhs_type, rhs_dtype="int8"):
+    X = te.placeholder((m, k), name="X", dtype=lhs_type)
+    W = te.placeholder((n, k), name="W", dtype=rhs_dtype)
+
+    ak = te.reduce_axis((0, k), name="k")
+    matmul = te.compute(
+        (m, n),
+        lambda i, j: te.sum(
+            X[i, ak].astype("int32") * W[j, ak].astype("int32"),
+            axis=ak,
+        ),
+        name="compute",
+    )
+
+    return te.create_prim_func([X, W, matmul])
+
+
+def tensorize_16x4_test(intrin=VNNI_DOT_16x4_INTRIN):
+    m, n, k = 128, 128, 128
+
+    func = get_matmul_packed(m, n, k, "uint8")
+
+    sch = tvm.s_tir.Schedule(func, debug_mask="all")
+    block = sch.get_sblock("compute")
+    sch.transform_layout(block, "W", lambda i, j: [i//16, j//4, i%16, j%4])
+    _, j, k = sch.get_loops(block)
+
+    _, ji = sch.split(j, factors=[None, 16])
+    ko, ki = sch.split(k, factors=[None, 4])
+    sch.reorder(ko, ji, ki)
+
+    sch.decompose_reduction(block, ko)
+    sch.tensorize(ji, intrin)
+
+    verify_trace_roundtrip(sch=sch, mod=func)
+
+
+def test_tensorize_vnni():
+    tensorize_16x4_test()
+
+
+def test_tensorize_avx512():
+    tensorize_16x4_test(AVX512_DOT_16x4_INTRIN)
+
+
+def test_tensorize_arm_dot():
+    m, n, k = 128, 128, 128
+
+    func = get_matmul_packed(m, n, k, "int8")
+
+    for intrin in [ARM_DOT_4x4_i8_SDOT_INTRIN, ARM_DOT_4x4_i8_NEON_INTRIN]:
+        sch = tvm.s_tir.Schedule(func, debug_mask="all")
+        block = sch.get_sblock("compute")
+        sch.transform_layout(block, "W", lambda i, j: [i//4, j//4, i%4, j%4])
+        _, j, k = sch.get_loops(block)
+
+        _, ji = sch.split(j, factors=[None, 4])
+        ko, ki = sch.split(k, factors=[None, 4])
+        sch.reorder(ko, ji, ki)
+
+        sch.decompose_reduction(block, ko)
+        sch.tensorize(ji, intrin)
+
+        verify_trace_roundtrip(sch=sch, mod=func)
+
+
+def test_tensorize_vrmpy():
+    m, n, k = 128, 128, 128
+
+    func = get_matmul_packed(m, n, k, "uint8", "uint8")
+
+    sch = tvm.s_tir.Schedule(func, debug_mask="all")
+    block = sch.get_sblock("compute")
+    sch.transform_layout(block, "W", lambda i, j: [i//32, j//4, i%32, j%4])
+    _, j, k = sch.get_loops(block)
+
+    _, ji = sch.split(j, factors=[None, 32])
+    ko, ki = sch.split(k, factors=[None, 4])
+    sch.reorder(ko, ji, ki)
+
+    sch.decompose_reduction(block, ko)
+    sch.tensorize(ji, VRMPY_u8u8i32_INTRIN)
+
+    verify_trace_roundtrip(sch=sch, mod=func)
+
+
+def test_tensorize_vdmpy():
+    m, n, k = 128, 128, 128
+
+    func = get_matmul_packed(m, n, k, "int16", "int16")
+
+    sch = tvm.s_tir.Schedule(func, debug_mask="all")
+    block = sch.get_sblock("compute")
+    sch.transform_layout(block, "W", lambda i, j: [i//32, j//2, i%32, j%2])
+    _, j, k = sch.get_loops(block)
+
+    _, ji = sch.split(j, factors=[None, 32])
+    ko, ki = sch.split(k, factors=[None, 2])
+    sch.reorder(ko, ji, ki)
+
+    sch.decompose_reduction(block, ko)
+    sch.tensorize(ji, VDMPY_i16i16i32_INTRIN)
+
+    verify_trace_roundtrip(sch=sch, mod=func)
+
+
+def test_tensorize_dp4a():
+    # pylint: disable=too-many-locals
+    def _test_intrin(dtype_a, dtype_b, dtype_c, intrin):
+        m, n, k = 128, 128, 128
+        X = te.placeholder((m, k), name="X", dtype=dtype_a)
+        W = te.placeholder((n, k), name="W", dtype=dtype_b)
+        ak = te.reduce_axis((0, k), name="k")
+
+        matmul = te.compute(
+            (m, n),
+            lambda i, j: te.sum(
+                X[i, ak].astype(dtype_c) * W[j, ak].astype(dtype_c),
+                axis=ak,
+            ),
+            name="compute",
+        )
+
+        func = te.create_prim_func([X, W, matmul])
+
+        sch = tvm.s_tir.Schedule(func, debug_mask="all")
+        block = sch.get_sblock("compute")
+        i, j, k = sch.get_loops(block)
+
+        by, ty, yi = sch.split(i, factors=sch.sample_perfect_tile(i, n=3))
+        bx, tx, xi = sch.split(j, factors=sch.sample_perfect_tile(j, n=3))
+        ko, ki = sch.split(k, [None, 4])
+        ko, kt = sch.split(ko, factors=sch.sample_perfect_tile(ko, n=2))
+
+        sch.reorder(by, bx, ty, tx, yi, xi)
+        CC = sch.cache_write(block, 0, "local")
+        sch.reverse_compute_at(CC, tx)
+
+        def fetch_to_shared(block, idx):
+            block_read = sch.cache_read(block, idx, "shared")
+            sch.compute_at(block_read, ko, True)
+            return block_read
+
+        fetch_to_shared(block, 0)
+        fetch_to_shared(block, 1)
+
+        sch.decompose_reduction(block, ko)
+        sch.tensorize(ki, intrin)
+
+        verify_trace_roundtrip(sch=sch, mod=func)
+
+    for args in [
+        ("int8", "int8", "int32", AMDGPU_SDOT4_INTRIN),
+        ("int8", "int8", "int32", DP4A_S8S8S32_INTRIN),
+        ("int8", "uint8", "int32", DP4A_S8U8S32_INTRIN),
+        ("uint8", "int8", "int32", DP4A_U8S8S32_INTRIN),
+        ("uint8", "uint8", "uint32", DP4A_U8U8U32_INTRIN),
+    ]:
+        _test_intrin(*args)
+
+
+def test_tensor_intrin_look_up():
+    intrin_name = 'non_existent_intrin'
+    assert tirx.TensorIntrin.get(intrin_name, allow_missing=True) is None
+    with pytest.raises(ValueError):
+        tirx.TensorIntrin.get(intrin_name)
+
+
+def test_tensorize_matmul_mixed_dtype():
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def matmul_int64_shape(
+        A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+        B: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+        C: T.Buffer((T.int64(128), T.int64(128)), "float32")
+    ) -> None:
+        for i_0, j_0 in T.grid(T.int64(8), T.int64(8)):
+            for i_1_init, j_1_init in T.grid(T.int64(16), T.int64(16)):
+                with T.sblock("init"):
+                    vi = T.axis.spatial(T.int64(128), i_0 * T.int64(16) + i_1_init)
+                    vj = T.axis.spatial(T.int64(128), j_0 * T.int64(16) + j_1_init)
+                    C[vi, vj] = T.float32(0)
+            for k_0, i_1, j_1, k_1 in T.grid(T.int64(8), T.int64(16), T.int64(16), T.int64(16)):
+                with T.sblock("update"):
+                    vi = T.axis.spatial(T.int64(128), i_0 * T.int64(16) + i_1)
+                    vj = T.axis.spatial(T.int64(128), j_0 * T.int64(16) + j_1)
+                    vk = T.axis.reduce(T.int64(128), k_0 * T.int64(16) + k_1)
+                    C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+    @T.prim_func(s_tir=True)
+    def tensorized_matmul_int64_shape(
+        A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+        B: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+        C: T.Buffer((T.int64(128), T.int64(128)), "float32")
+    ) -> None:
+        for i_outer, j_outer in T.grid(T.int64(8), T.int64(8)):
+            for i_inner_init, j_inner_init in T.grid(T.int64(16), T.int64(16)):
+                with T.sblock("init"):
+                    vi = T.axis.spatial(T.int64(128), i_outer * T.int64(16) + i_inner_init)
+                    vj = T.axis.spatial(T.int64(128), j_outer * T.int64(16) + j_inner_init)
+                    C[vi, vj] = T.float32(0)
+            for k_outer in T.grid(T.int64(8)):
+                with T.sblock("update"):
+                    vi, vj, vk = T.axis.remap("SSR", [i_outer, j_outer, k_outer])
+                    T.reads(
+                        [
+                            C[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vj * T.int64(16) : vj * T.int64(16) + T.int64(16)],
+                            A[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vk * T.int64(16) : vk * T.int64(16) + T.int64(16)],
+                            B[vj * T.int64(16) : vj * T.int64(16) + T.int64(16), vk * T.int64(16) : vk * T.int64(16) + T.int64(16)],
+                        ]
+                    )
+                    T.writes(C[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vj * T.int64(16) : vj * T.int64(16) + T.int64(16)])
+                    A_elem_offset = T.int64()
+                    B_elem_offset = T.int64()
+                    C_elem_offset = T.int64()
+                    A_sub = T.match_buffer(
+                        A[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vk * T.int64(16) : vk * T.int64(16) + T.int64(16)],
+                        [T.int64(16), T.int64(16)],
+                        elem_offset=A_elem_offset,
+                    )
+                    B_sub = T.match_buffer(
+                        B[vj * T.int64(16) : vj * T.int64(16) + T.int64(16), vk * T.int64(16) : vk * T.int64(16) + T.int64(16)],
+                        [T.int64(16), T.int64(16)],
+                        elem_offset=B_elem_offset,
+                    )
+                    C_sub = T.match_buffer(
+                        C[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vj * T.int64(16) : vj * T.int64(16) + T.int64(16)],
+                        [T.int64(16), T.int64(16)],
+                        elem_offset=C_elem_offset,
+                    )
+                    T.evaluate(
+                        T.tvm_mma_sync(
+                            C_sub.data,
+                            T.floordiv(C_sub.elem_offset, T.int64(256)),
+                            A_sub.data,
+                            T.floordiv(A_sub.elem_offset, T.int64(256)),
+                            B_sub.data,
+                            T.floordiv(B_sub.elem_offset, T.int64(256)),
+                            C_sub.data,
+                            T.floordiv(C_sub.elem_offset, T.int64(256)),
+                            dtype="handle",
+                        )
+                    )
+    # fmt: on
+
+    s = tvm.s_tir.Schedule(matmul_int64_shape, debug_mask="all")
+    update = s.get_sblock("update")
+    ii = s.get_loops(update)[-3]
+    s.tensorize(ii, "test_mma_intrin")
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], tensorized_matmul_int64_shape)
+    verify_trace_roundtrip(sch=s, mod=matmul_int64_shape)
+
+def _tir_packed_int_to_int_to_float(storage_nbit: int):
+    storage_dtype = "int" + str(storage_nbit)
+
+    def f_convert(nbit: int, val: tirx.Expr, pos: tirx.Expr, dtype: str):
+        assert val.ty.dtype == storage_dtype
+        mask = tirx.const((1 << nbit) - 1, "int32")
+        unextended = (val >> (pos.astype("int32") * tirx.const(nbit, "int32"))) & mask
+        return tirx.Cast(dtype, (unextended << tirx.const(32 - nbit, "int32")) >> tirx.const(32 - nbit, "int32"))
+
+    return f_convert
+
+@T.prim_func(s_tir=True)
+def decode_i4s_to_f16_desc(compressed: T.handle, decompressed: T.handle) -> None:
+    Compressed = T.match_buffer(
+        compressed,
+        [
+            1,
+        ],
+        dtype="int32",
+        scope="local",
+    )
+    Decompressed = T.match_buffer(
+        decompressed,
+        [
+            8,
+        ],
+        dtype="float16",
+        scope="local",
+    )
+
+    with T.sblock("root"):
+        T.reads(Compressed[0:1])
+        T.writes(Decompressed[0:8])
+        for i in T.grid(8):
+            with T.sblock("decode"):
+                vi = T.axis.remap("S", [i])
+                Decompressed[vi] = _tir_packed_int_to_int_to_float(32)(
+                    4,
+                    Compressed[vi // 8],
+                    vi % 8,
+                    dtype="float16",
+                )
+
+@T.prim_func(s_tir=True)
+def decode_i4s_to_f16_impl(compressed: T.handle, decompressed: T.handle) -> None:
+    Compressed = T.match_buffer(
+        compressed,
+        [
+            1,
+        ],
+        dtype="int32",
+        scope="local",
+    )
+    Decompressed = T.match_buffer(
+        decompressed,
+        [
+            8,
+        ],
+        dtype="float16",
+        scope="local",
+    )
+
+    with T.sblock("root"):
+        T.reads(Compressed[0:1])
+        T.writes(Decompressed[0:8])
+        T.call_extern(
+            "handle",
+            "test_decode_i4s_to_f16",
+            Compressed.data,
+            Decompressed.data,
+            8,
+        )
+
+tirx.TensorIntrin.register("test_decode_i4s_to_f16_intrin", decode_i4s_to_f16_desc, decode_i4s_to_f16_impl)
+
+def test_tensorize_arith_simplification():
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def decode_i4s_to_int32_to_f16():
+        B_decode_local = T.sblock_alloc_buffer((16384, 16384), "float16", scope="local")
+        B_local = T.sblock_alloc_buffer((16384, 2048), "int32", scope="local")
+        for ax0_0 in T.thread_binding(8192, thread="blockIdx.x"):
+            for ax0_1 in T.thread_binding(2, thread="threadIdx.y"):
+                for ax1_0 in range(32):
+                    for ax1_1 in T.thread_binding(64, thread="threadIdx.x"):
+                        for ax0, ax1 in T.grid(1, 8):
+                            with T.sblock("B_decode_local"):
+                                v0 = T.axis.spatial(16384, ax0_0 * 2 + ax0_1 + ax0)
+                                v1 = T.axis.spatial(16384, ax1_0 * 512 + ax1_1 * 8 + ax1)
+                                T.reads(B_local[v0, v1 // 8])
+                                T.writes(B_decode_local[v0, v1])
+                                B_decode_local[v0, v1] = T.Cast("float16", T.shift_right(T.shift_left(T.bitwise_and(T.shift_right(B_local[v0, v1 // 8], v1 % 8 * 4), 15), 28), 28))
+
+    @T.prim_func(s_tir=True)
+    def tensorized_decode_i4s_to_int32_to_f16():
+        B_decode_local = T.sblock_alloc_buffer((16384, 16384), "float16", scope="local")
+        B_local = T.sblock_alloc_buffer((16384, 2048), "int32", scope="local")
+        for ax0_0 in T.thread_binding(8192, thread="blockIdx.x"):
+            for ax0_1 in T.thread_binding(2, thread="threadIdx.y"):
+                for ax1_0 in range(32):
+                    for ax1_1 in T.thread_binding(64, thread="threadIdx.x"):
+                        for ax0 in range(1):
+                            with T.sblock("B_decode_local_o"):
+                                v0_o = T.axis.spatial(16384, ax0_0 * 2 + ax0_1 + ax0)
+                                v1_o = T.axis.spatial(2048, ax1_0 * 64 + ax1_1)
+                                T.reads(B_local[v0_o, v1_o])
+                                T.writes(B_decode_local[v0_o, v1_o * 8:v1_o * 8 + 8])
+                                Compressed = T.match_buffer(B_local[v0_o, v1_o], (1,), "int32", scope="local")
+                                Decompressed = T.match_buffer(B_decode_local[v0_o, v1_o * 8:v1_o * 8 + 8], (8,), "float16", scope="local")
+                                T.call_extern("handle", "test_decode_i4s_to_f16", Compressed.data, Decompressed.data, 8)
+
+    s = tvm.s_tir.Schedule(decode_i4s_to_int32_to_f16, debug_mask="all")
+    update = s.get_sblock("B_decode_local")
+    ii = s.get_loops(update)[-1]
+    s.tensorize(ii, "test_decode_i4s_to_f16_intrin")
+    assert_structural_equal_ignore_global_symbol(s.mod["main"], tensorized_decode_i4s_to_int32_to_f16)
+    verify_trace_roundtrip(sch=s, mod=decode_i4s_to_int32_to_f16)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py

@@ -0,0 +1,464 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import te
+from tvm.s_tir.tensor_intrin.cuda import (
+    LDMATRIX_e4m3_A_INTRIN,
+    LDMATRIX_e4m3_B_TRANS_INTRIN,
+    LDMATRIX_e5m2_A_INTRIN,
+    LDMATRIX_e5m2_B_TRANS_INTRIN,
+    LDMATRIX_f16_A_INTRIN,
+    LDMATRIX_f16_B_INTRIN,
+    LDMATRIX_f16_B_TRANS_INTRIN,
+    LDMATRIX_i8_A_INTRIN,
+    LDMATRIX_i8_B_INTRIN,
+    LDMATRIX_i8_B_TRANS_INTRIN,
+    MMA_e4m3e4m3f32_TRANS_B_INTRIN,
+    MMA_e5m2e5m2f32_TRANS_B_INTRIN,
+    MMA_f16f16f16_INTRIN,
+    MMA_f16f16f16_TRANS_B_INTRIN,
+    MMA_f16f16f32_INTRIN,
+    MMA_f16f16f32_TRANS_B_INTRIN,
+    MMA_fill_16x16_f16_INTRIN,
+    MMA_fill_16x16_f32_INTRIN,
+    MMA_fill_16x16_i32_INTRIN,
+    MMA_i8i8i32_INTRIN,
+    MMA_i8i8i32_TRANS_B_INTRIN,
+    MMA_store_16x16_f16_global_INTRIN,
+    MMA_store_16x16_f32_global_INTRIN,
+    MMA_store_16x16_i32_global_INTRIN,
+    shared_16x16_to_ldmatrix_32x8_layout,
+    shared_16x32_to_ldmatrix_32x16_layout,
+    shared_32x16_to_ldmatrix_32x16_layout,
+)
+from tvm.testing import env
+from tvm.testing.tir import mma_schedule
+
+M = 4096
+N = 4096
+K = 4096
+measure_perf = False
+gflops = (N * M * K) * 2 / 1e9
+
+
+def matmul(m, n, k, in_dtype, out_dtype, b_transposed):
+    b_shape = (n, k) if b_transposed else (k, n)
+    a = te.placeholder((m, k), name="A", dtype=in_dtype)
+    b = te.placeholder(b_shape, name="B", dtype=in_dtype)
+    k = te.reduce_axis((0, k), name="k")
+
+    def maybe_cast(v):
+        if in_dtype != out_dtype:
+            return tvm.tirx.Cast(out_dtype, v)
+        return v
+
+    def maybe_swap(i, j):
+        if b_transposed:
+            return j, i
+        return i, j
+
+    c = te.compute(
+        (m, n),
+        lambda i, j: te.sum(maybe_cast(a[i, k]) * maybe_cast(b[maybe_swap(k, j)]), axis=[k]),
+        name="C",
+    )
+    return (a, b, c)
+
+
+def run_test(
+    k_inner,
+    in_dtype,
+    out_dtype,
+    b_transposed,
+    i_factors,
+    j_factors,
+    k_factors,
+    index_map_A,
+    index_map_B,
+    index_map_C,
+    ldmatrix_a_intrin,
+    ldmatrix_b_intrin,
+    mma_intrin,
+    mma_fill_intrin,
+    mma_store_intrin,
+):
+    sch = mma_schedule(
+        te.create_prim_func(matmul(M, N, K, in_dtype, out_dtype, b_transposed)),
+        k_inner,
+        in_dtype,
+        b_transposed,
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map_A,
+        index_map_B,
+        index_map_C,
+        ldmatrix_a_intrin,
+        ldmatrix_b_intrin,
+        mma_intrin,
+        mma_fill_intrin,
+        mma_store_intrin,
+    )
+
+    f = tvm.compile(sch.mod["main"], target="cuda")
+
+    if in_dtype == "float16":
+        a_np = np.random.normal(size=(M, K)).astype("float16")
+
+        if b_transposed:
+            b_np = np.random.normal(size=(N, K)).astype("float16")
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32").transpose()).astype(
+                out_dtype
+            )
+        else:
+            b_np = np.random.normal(size=(K, N)).astype("float16")
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32")).astype(out_dtype)
+    elif in_dtype in ["float8_e4m3fn", "float8_e5m2"]:
+        typemap = {
+            "float8_e4m3fn": "float8_e4m3fn",
+            "float8_e5m2": "float8_e5m2",
+        }
+        a_np = (
+            np.random.uniform(low=-5, high=5, size=(M * K))
+            .reshape((M, K))
+            .astype(typemap[in_dtype])
+        )
+        if b_transposed:
+            b_np = (
+                np.random.uniform(low=-5, high=5, size=(N * K))
+                .reshape((N, K))
+                .astype(typemap[in_dtype])
+            )
+            c_np = np.dot(a_np.astype("float32"), b_np.T.astype("float32")).astype(out_dtype)
+        else:
+            b_np = (
+                np.random.uniform(low=-5, high=5, size=(N * K))
+                .reshape((K, N))
+                .astype(typemap[in_dtype])
+            )
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32")).astype(out_dtype)
+    else:
+        a_np = np.random.randint(-128, 128, (M, K)).astype("int8")
+
+        if b_transposed:
+            b_np = np.random.randint(-128, 128, (N, K)).astype("int8")
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32").transpose()).astype(
+                "int32"
+            )
+        else:
+            b_np = np.random.randint(-128, 128, (K, N)).astype("int8")
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32")).astype("int32")
+
+    def run_and_check(measure=False):
+        dev = tvm.device("cuda", 0)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        c = tvm.runtime.tensor(np.zeros((M, N), dtype=out_dtype), dev)
+        if measure:
+            return f.time_evaluator(f.entry_name, dev, number=500)(a, b, c)
+        f(a, b, c)
+        dev.sync()
+        if out_dtype != "float16" and in_dtype not in ["float8_e4m3fn", "float8_e5m2"]:
+            tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-2, atol=1e-2)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+    return lambda: tvm.testing.run_with_gpu_lock(run_and_check, True)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
+def test_f16f16f32_m16n16k16():
+    def index_map(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
+        )
+
+    k_inner = 16
+    in_dtype = "float16"
+    out_dtype = "float32"
+    i_factors, j_factors, k_factors = [4, 8, 2, 4, 1], [1, 64, 2, 1, 2], [128, 2, 1]
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        False,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map,
+        index_map,
+        index_map,
+        LDMATRIX_f16_A_INTRIN,
+        LDMATRIX_f16_B_INTRIN,
+        MMA_f16f16f32_INTRIN,
+        MMA_fill_16x16_f32_INTRIN,
+        MMA_store_16x16_f32_global_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("f16f16f32_m16n16k16: %f GFLOPS" % (gflops / (timer().mean)))
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        True,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map,
+        index_map,
+        index_map,
+        LDMATRIX_f16_A_INTRIN,
+        LDMATRIX_f16_B_TRANS_INTRIN,
+        MMA_f16f16f32_TRANS_B_INTRIN,
+        MMA_fill_16x16_f32_INTRIN,
+        MMA_store_16x16_f32_global_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("f16f16f32_m16n16k16_trans: %f GFLOPS" % (gflops / (timer().mean)))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
+def test_f16f16f16_m16n16k16():
+    def index_map(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
+        )
+
+    k_inner = 16
+    in_dtype = "float16"
+    out_dtype = "float16"
+    i_factors, j_factors, k_factors = [16, 2, 1, 4, 2], [16, 2, 2, 1, 4], [128, 2, 1]
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        False,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map,
+        index_map,
+        index_map,
+        LDMATRIX_f16_A_INTRIN,
+        LDMATRIX_f16_B_INTRIN,
+        MMA_f16f16f16_INTRIN,
+        MMA_fill_16x16_f16_INTRIN,
+        MMA_store_16x16_f16_global_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("f16f16f16_m16n16k16: %f GFLOPS" % (gflops / (timer().mean)))
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        True,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map,
+        index_map,
+        index_map,
+        LDMATRIX_f16_A_INTRIN,
+        LDMATRIX_f16_B_TRANS_INTRIN,
+        MMA_f16f16f16_TRANS_B_INTRIN,
+        MMA_fill_16x16_f16_INTRIN,
+        MMA_store_16x16_f16_global_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("f16f16f16_m16n16k16_trans: %f GFLOPS" % (gflops / (timer().mean)))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
+def test_i8i8i32_m16n16k32():
+    def index_map_A(i, j):
+        return (
+            i // 16,
+            j // 32,
+            *shared_16x32_to_ldmatrix_32x16_layout(i % 16, j % 32),
+        )
+
+    def index_map_B(i, j):
+        return (
+            i // 32,
+            j // 16,
+            *shared_32x16_to_ldmatrix_32x16_layout(i % 32, j % 16),
+        )
+
+    def index_map_C(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
+        )
+
+    k_inner = 32
+    in_dtype = "int8"
+    out_dtype = "int32"
+    i_factors, j_factors, k_factors = [1, 32, 1, 4, 2], [8, 4, 4, 2, 1], [32, 2, 2]
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        False,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map_A,
+        index_map_B,
+        index_map_C,
+        LDMATRIX_i8_A_INTRIN,
+        LDMATRIX_i8_B_INTRIN,
+        MMA_i8i8i32_INTRIN,
+        MMA_fill_16x16_i32_INTRIN,
+        MMA_store_16x16_i32_global_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("i8i8i32_m16n16k32: %f GOPS" % (gflops / (timer().mean)))
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        True,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map_A,
+        index_map_A,
+        index_map_C,
+        LDMATRIX_i8_A_INTRIN,
+        LDMATRIX_i8_B_TRANS_INTRIN,
+        MMA_i8i8i32_TRANS_B_INTRIN,
+        MMA_fill_16x16_i32_INTRIN,
+        MMA_store_16x16_i32_global_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("i8i8i32_m16n16k32_trans: %f GOPS" % (gflops / (timer().mean)))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda compute >= 8.9")
+def test_e4m3e4m3f32_m16n16k32():
+    def index_map_A(i, j):
+        return (
+            i // 16,
+            j // 32,
+            *shared_16x32_to_ldmatrix_32x16_layout(i % 16, j % 32),
+        )
+
+    def index_map_C(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
+        )
+
+    k_inner = 32
+    in_dtype = "float8_e4m3fn"
+    out_dtype = "float32"
+    i_factors, j_factors, k_factors = [1, 32, 1, 4, 2], [8, 4, 4, 2, 1], [32, 2, 2]
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        True,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map_A,
+        index_map_A,
+        index_map_C,
+        LDMATRIX_e4m3_A_INTRIN,
+        LDMATRIX_e4m3_B_TRANS_INTRIN,
+        MMA_e4m3e4m3f32_TRANS_B_INTRIN,
+        MMA_fill_16x16_f32_INTRIN,
+        MMA_store_16x16_f32_global_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("e4m3e4m3f32_m16n16k32_trans: %f GOPS" % (gflops / (timer().mean)))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda compute >= 8.9")
+def test_e5m2e5m2f32_m16n16k32():
+    def index_map_A(i, j):
+        return (
+            i // 16,
+            j // 32,
+            *shared_16x32_to_ldmatrix_32x16_layout(i % 16, j % 32),
+        )
+
+    def index_map_C(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
+        )
+
+    k_inner = 32
+    in_dtype = "float8_e5m2"
+    out_dtype = "float32"
+    i_factors, j_factors, k_factors = [1, 32, 1, 4, 2], [8, 4, 4, 2, 1], [32, 2, 2]
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        True,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map_A,
+        index_map_A,
+        index_map_C,
+        LDMATRIX_e5m2_A_INTRIN,
+        LDMATRIX_e5m2_B_TRANS_INTRIN,
+        MMA_e5m2e5m2f32_TRANS_B_INTRIN,
+        MMA_fill_16x16_f32_INTRIN,
+        MMA_store_16x16_f32_global_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("e5m2e5m2f32_m16n16k32_trans: %f GOPS" % (gflops / (timer().mean)))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py

@@ -0,0 +1,319 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import te
+from tvm.s_tir.tensor_intrin.rocm import (
+    ROCM_MFMA_f16f16f32_INTRIN,
+    ROCM_MFMA_f32f32f32_INTRIN,
+    ROCM_MFMA_fill_16x16_f32_INTRIN,
+    ROCM_MFMA_fill_16x16_i32_INTRIN,
+    ROCM_MFMA_LOAD_16x4_A_SHARED_f32_INTRIN,
+    ROCM_MFMA_LOAD_16x4_B_SHARED_f32_INTRIN,
+    ROCM_MFMA_LOAD_16x16_A_SHARED_f16_INTRIN,
+    ROCM_MFMA_LOAD_16x16_A_SHARED_s8_INTRIN,
+    ROCM_MFMA_LOAD_16x16_B_SHARED_f16_INTRIN,
+    ROCM_MFMA_LOAD_16x16_B_SHARED_s8_INTRIN,
+    ROCM_MFMA_s8s8s32_INTRIN,
+    ROCM_MFMA_STORE_16x16_f32_INTRIN,
+    ROCM_MFMA_STORE_16x16_s32_INTRIN,
+    shared_4x16_to_local_64x1_layout_B,
+    shared_16x4_to_local_64x1_layout_A,
+    shared_16x16_to_local_64x4_layout_A,
+    shared_16x16_to_local_64x4_layout_B,
+    shared_16x16_to_local_64x4_layout_C,
+)
+from tvm.testing import env
+from tvm.testing.tir import mfma_schedule
+
+M = 1024
+N = 1024
+K = 1024
+measure_perf = False
+gflops = (N * M * K) * 2 / 1e9
+
+
+def matmul(m, n, k, in_dtype, out_dtype, b_transposed):
+    b_shape = (n, k) if b_transposed else (k, n)
+    a = te.placeholder((m, k), name="A", dtype=in_dtype)
+    b = te.placeholder(b_shape, name="B", dtype=in_dtype)
+    k = te.reduce_axis((0, k), name="k")
+
+    def maybe_cast(v):
+        if in_dtype != out_dtype:
+            return tvm.tirx.Cast(out_dtype, v)
+        return v
+
+    def maybe_swap(i, j):
+        if b_transposed:
+            return j, i
+        return i, j
+
+    c = te.compute(
+        (m, n),
+        lambda i, j: te.sum(maybe_cast(a[i, k]) * maybe_cast(b[maybe_swap(k, j)]), axis=[k]),
+        name="C",
+    )
+    return (a, b, c)
+
+
+def run_test(
+    k_inner,
+    in_dtype,
+    out_dtype,
+    b_transposed,
+    i_factors,
+    j_factors,
+    k_factors,
+    index_map_A,
+    index_map_B,
+    index_map_C,
+    ldmatrix_a_intrin,
+    ldmatrix_b_intrin,
+    mma_intrin,
+    mma_fill_intrin,
+    mma_store_intrin,
+):
+    sch = mfma_schedule(
+        te.create_prim_func(matmul(M, N, K, in_dtype, out_dtype, b_transposed)),
+        k_inner,
+        in_dtype,
+        b_transposed,
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map_A,
+        index_map_B,
+        index_map_C,
+        ldmatrix_a_intrin,
+        ldmatrix_b_intrin,
+        mma_intrin,
+        mma_fill_intrin,
+        mma_store_intrin,
+    )
+
+    f = tvm.compile(sch.mod["main"], target="rocm")
+
+    if in_dtype == "float32":
+        a_np = np.random.uniform(size=(M, K)).astype("float32")
+
+        if b_transposed:
+            b_np = np.random.uniform(size=(N, K)).astype("float32")
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32").transpose()).astype(
+                out_dtype
+            )
+        else:
+            b_np = np.random.uniform(size=(K, N)).astype("float32")
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32")).astype(out_dtype)
+    elif in_dtype == "float16":
+        a_np = np.random.uniform(size=(M, K)).astype("float16")
+
+        if b_transposed:
+            b_np = np.random.uniform(size=(N, K)).astype("float16")
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32").transpose()).astype(
+                out_dtype
+            )
+        else:
+            b_np = np.random.uniform(size=(K, N)).astype("float16")
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32")).astype(out_dtype)
+    else:
+        a_np = np.random.randint(-128, 128, (M, K)).astype("int8")
+
+        if b_transposed:
+            b_np = np.random.randint(-128, 128, (N, K)).astype("int8")
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32").transpose()).astype(
+                "int32"
+            )
+        else:
+            b_np = np.random.randint(-128, 128, (K, N)).astype("int8")
+            c_np = np.dot(a_np.astype("float32"), b_np.astype("float32")).astype("int32")
+
+    def run_and_check(measure=False):
+        dev = tvm.device("rocm", 0)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        c = tvm.runtime.tensor(np.zeros((M, N), dtype=out_dtype), dev)
+        if measure:
+            return f.time_evaluator(f.entry_name, dev, number=500)(a, b, c)
+        f(a, b, c)
+        dev.sync()
+        if in_dtype != "float16":
+            tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-2, atol=1e-2)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+    return lambda: tvm.testing.run_with_gpu_lock(run_and_check, True)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_matrixcore(), reason="need matrixcore")
+def test_i8i8i32_m16n16k16():
+    def index_map_A(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_local_64x4_layout_A(i % 16, j % 16),
+        )
+
+    def index_map_B(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_local_64x4_layout_B(i % 16, j % 16),
+        )
+
+    def index_map_C(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_local_64x4_layout_C(i % 16, j % 16),
+        )
+
+    k_inner = 16
+    in_dtype = "int8"
+    out_dtype = "int32"
+    i_factors, j_factors, k_factors = [1, 8, 2, 4, 1], [1, 16, 2, 1, 2], [32, 2, 1]
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        False,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map_A,
+        index_map_B,
+        index_map_C,
+        ROCM_MFMA_LOAD_16x16_A_SHARED_s8_INTRIN,
+        ROCM_MFMA_LOAD_16x16_B_SHARED_s8_INTRIN,
+        ROCM_MFMA_s8s8s32_INTRIN,
+        ROCM_MFMA_fill_16x16_i32_INTRIN,
+        ROCM_MFMA_STORE_16x16_s32_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("test_i8i8i32_m16n16k16: %f GFLOPS" % (gflops / (timer().mean)))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_matrixcore(), reason="need matrixcore")
+def test_f16f16f32_m16n16k16():
+    def index_map_A(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_local_64x4_layout_A(i % 16, j % 16),
+        )
+
+    def index_map_B(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_local_64x4_layout_B(i % 16, j % 16),
+        )
+
+    def index_map_C(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_local_64x4_layout_C(i % 16, j % 16),
+        )
+
+    k_inner = 16
+    in_dtype = "float16"
+    out_dtype = "float32"
+    i_factors, j_factors, k_factors = [1, 8, 2, 4, 1], [1, 16, 2, 1, 2], [32, 2, 1]
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        False,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map_A,
+        index_map_B,
+        index_map_C,
+        ROCM_MFMA_LOAD_16x16_A_SHARED_f16_INTRIN,
+        ROCM_MFMA_LOAD_16x16_B_SHARED_f16_INTRIN,
+        ROCM_MFMA_f16f16f32_INTRIN,
+        ROCM_MFMA_fill_16x16_f32_INTRIN,
+        ROCM_MFMA_STORE_16x16_f32_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("f16f16f32_m16n16k16: %f GFLOPS" % (gflops / (timer().mean)))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_matrixcore(), reason="need matrixcore")
+def test_f32f32f32_m16n16k4():
+    def index_map_A(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x4_to_local_64x1_layout_A(i % 16, j % 16),
+        )
+
+    def index_map_B(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_4x16_to_local_64x1_layout_B(i % 16, j % 16),
+        )
+
+    def index_map_C(i, j):
+        return (
+            i // 16,
+            j // 16,
+            *shared_16x16_to_local_64x4_layout_C(i % 16, j % 16),
+        )
+
+    k_inner = 4
+    in_dtype = "float32"
+    out_dtype = "float32"
+    i_factors, j_factors, k_factors = [4, 2, 1, 4, 2], [4, 2, 2, 1, 4], [128, 2, 1]
+
+    timer = run_test(
+        k_inner,
+        in_dtype,
+        out_dtype,
+        False,  # b_transposed
+        i_factors,
+        j_factors,
+        k_factors,
+        index_map_A,
+        index_map_B,
+        index_map_C,
+        ROCM_MFMA_LOAD_16x4_A_SHARED_f32_INTRIN,
+        ROCM_MFMA_LOAD_16x4_B_SHARED_f32_INTRIN,
+        ROCM_MFMA_f32f32f32_INTRIN,
+        ROCM_MFMA_fill_16x16_f32_INTRIN,
+        ROCM_MFMA_STORE_16x16_f32_INTRIN,
+    )
+
+    if measure_perf and timer:
+        print("test_f32f32f32_m16n16k4: %f GFLOPS" % (gflops / (timer().mean)))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_trace.py

@@ -0,0 +1,399 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# mypy: ignore-errors
+# ruff: noqa: E501, F401
+import sys
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import s_tir, tirx
+from tvm.s_tir.schedule import Instruction, InstructionKind, LoopRV, SBlockRV, Trace
+from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def elementwise(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.sblock_alloc_buffer((128, 128))
+    C = T.match_buffer(c, (128, 128))
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def elementwise_inlined(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = A[vi, vj] * 2.0 + 1.0
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+
+def _make_get_sblock(name, output):
+    return Instruction(
+        kind=InstructionKind.get("GetSBlock"),
+        inputs=[],
+        attrs=[name, "main"],
+        outputs=[output],
+    )
+
+
+def _make_get_loops(input, outputs):  # pylint: disable=redefined-builtin
+    return Instruction(
+        kind=InstructionKind.get("GetLoops"),
+        inputs=[input],
+        attrs=[],
+        outputs=outputs,
+    )
+
+
+def _make_compute_inline(input):  # pylint: disable=redefined-builtin
+    return Instruction(
+        kind=InstructionKind.get("ComputeInline"),
+        inputs=[input],
+        attrs=[],
+        outputs=[],
+    )
+
+
+def _make_split(inputs, outputs):  # pylint: disable=redefined-builtin
+    return Instruction(
+        kind=InstructionKind.get("Split"),
+        inputs=inputs,
+        attrs=[T.bool(True), T.bool(False)],
+        outputs=outputs,
+    )
+
+
+def _make_enter_postproc():
+    return Instruction(
+        kind=InstructionKind.get("EnterPostproc"),
+        inputs=[],
+        attrs=[],
+        outputs=[],
+    )
+
+
+def _make_annotate(block: SBlockRV, annotation: str):
+    return Instruction(
+        kind=InstructionKind.get("Annotate"),
+        inputs=[block, annotation],
+        attrs=["meta_schedule.auto_tensorize"],
+        outputs=[],
+    )
+
+
+def _make_trace_1(b0, l1, l2):  # pylint: disable=invalid-name
+    return Trace(
+        insts=[
+            _make_get_sblock(name="block", output=b0),
+            _make_get_loops(input=b0, outputs=[l1, l2]),
+        ],
+        decisions={},
+    )
+
+
+def _make_trace_2(b0):  # pylint: disable=invalid-name
+    return Trace(
+        insts=[
+            _make_get_sblock(name="B", output=b0),
+            _make_compute_inline(input=b0),
+        ],
+        decisions={},
+    )
+
+
+def _make_trace_3(b0, b1, add_postproc):  # pylint: disable=invalid-name
+    if add_postproc:
+        insts = [
+            _make_get_sblock(name="B", output=b0),
+            _make_compute_inline(input=b0),
+            _make_get_sblock(name="C", output=b1),
+            _make_enter_postproc(),
+            _make_compute_inline(input=b1),
+        ]
+    else:
+        insts = [
+            _make_get_sblock(name="B", output=b0),
+            _make_compute_inline(input=b0),
+            _make_get_sblock(name="C", output=b1),
+        ]
+    return Trace(insts=insts, decisions={})
+
+
+def _make_trace_4(b0, l1, l2, l3):  # pylint: disable=invalid-name
+    return Trace(
+        insts=[
+            _make_get_sblock(name="B", output=b0),
+            _make_get_loops(input=b0, outputs=[l1]),
+            _make_split([l1, None, T.int32(32)], [l2, l3]),
+        ],
+        decisions={},
+    )
+
+
+def test_trace_construct_1():
+    trace = _make_trace_1(SBlockRV(), LoopRV(), LoopRV())
+    assert str(trace) == "\n".join(
+        (
+            "# from tvm import s_tir",
+            "def apply_trace(sch: s_tir.Schedule) -> None:",
+            '  b0 = sch.get_sblock(name="block", func_name="main")',
+            "  l1, l2 = sch.get_loops(block=b0)",
+        )
+    )
+    assert len(trace.insts) == 2
+    assert len(trace.decisions) == 0
+
+
+def test_trace_construct_get_decision_1():
+    trace = _make_trace_1(SBlockRV(), LoopRV(), LoopRV())
+    assert trace.get_decision(trace.insts[0]) is None
+    assert trace.get_decision(trace.insts[1]) is None
+
+
+def test_trace_construct_append_1():
+    trace = _make_trace_1(SBlockRV(), LoopRV(), LoopRV())
+    trace.append(inst=_make_get_sblock("block2", SBlockRV()))
+    assert str(trace) == "\n".join(
+        (
+            "# from tvm import s_tir",
+            "def apply_trace(sch: s_tir.Schedule) -> None:",
+            '  b0 = sch.get_sblock(name="block", func_name="main")',
+            "  l1, l2 = sch.get_loops(block=b0)",
+            '  b3 = sch.get_sblock(name="block2", func_name="main")',
+        )
+    )
+
+
+def test_trace_construct_pop_1():
+    trace = _make_trace_1(SBlockRV(), LoopRV(), LoopRV())
+    last_inst = trace.insts[-1]
+    assert trace.pop().same_as(last_inst)
+    assert str(trace) == "\n".join(
+        (
+            "# from tvm import s_tir",
+            "def apply_trace(sch: s_tir.Schedule) -> None:",
+            '  b0 = sch.get_sblock(name="block", func_name="main")',
+        )
+    )
+
+
+def test_trace_construct_pop_2():
+    trace = Trace([], {})
+    assert str(trace) == "\n".join(
+        (
+            "# from tvm import s_tir",
+            "def apply_trace(sch: s_tir.Schedule) -> None:",
+            "  pass",
+        )
+    )
+    assert trace.pop() is None
+    assert str(trace) == "\n".join(
+        (
+            "# from tvm import s_tir",
+            "def apply_trace(sch: s_tir.Schedule) -> None:",
+            "  pass",
+        )
+    )
+
+
+def test_trace_apply_to_schedule():
+    trace = _make_trace_2(SBlockRV())
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    trace.apply_to_schedule(sch, remove_postproc=False, decision_provider=None)
+    assert_structural_equal_ignore_global_symbol(elementwise_inlined, sch.mod["main"])
+
+
+def test_trace_as_json_1():
+    trace = _make_trace_1(SBlockRV(), LoopRV(), LoopRV())
+    obj = trace.as_json()
+    assert obj == [
+        [
+            ["GetSBlock", [], ["block", "main"], ["b0"]],
+            ["GetLoops", ["b0"], [], ["l1", "l2"]],
+        ],
+        [],
+    ]
+
+
+def test_trace_simplified_1():
+    trace = _make_trace_3(SBlockRV(), SBlockRV(), add_postproc=True)
+    assert str(trace) == "\n".join(
+        (
+            "# from tvm import s_tir",
+            "def apply_trace(sch: s_tir.Schedule) -> None:",
+            '  b0 = sch.get_sblock(name="B", func_name="main")',
+            "  sch.compute_inline(block=b0)",
+            '  b1 = sch.get_sblock(name="C", func_name="main")',
+            "  sch.enter_postproc()",
+            "  sch.compute_inline(block=b1)",
+        )
+    )
+    trace = trace.simplified(remove_postproc=True)
+    assert str(trace) == "\n".join(
+        (
+            "# from tvm import s_tir",
+            "def apply_trace(sch: s_tir.Schedule) -> None:",
+            '  b0 = sch.get_sblock(name="B", func_name="main")',
+            "  sch.compute_inline(block=b0)",
+        )
+    )
+
+
+def test_trace_simplified_2():
+    trace = _make_trace_3(SBlockRV(), SBlockRV(), add_postproc=True)
+    assert str(trace) == "\n".join(
+        (
+            "# from tvm import s_tir",
+            "def apply_trace(sch: s_tir.Schedule) -> None:",
+            '  b0 = sch.get_sblock(name="B", func_name="main")',
+            "  sch.compute_inline(block=b0)",
+            '  b1 = sch.get_sblock(name="C", func_name="main")',
+            "  sch.enter_postproc()",
+            "  sch.compute_inline(block=b1)",
+        )
+    )
+    trace = trace.simplified(remove_postproc=False)
+    assert str(trace) == "\n".join(
+        (
+            "# from tvm import s_tir",
+            "def apply_trace(sch: s_tir.Schedule) -> None:",
+            '  b0 = sch.get_sblock(name="B", func_name="main")',
+            "  sch.compute_inline(block=b0)",
+            '  b1 = sch.get_sblock(name="C", func_name="main")',
+            "  sch.enter_postproc()",
+            "  sch.compute_inline(block=b1)",
+        )
+    )
+
+
+def test_trace_simplified_3():
+    trace = _make_trace_4(SBlockRV(), LoopRV(), LoopRV(), LoopRV()).simplified(
+        remove_postproc=False
+    )
+    assert str(trace) == "\n".join(
+        (
+            "# from tvm import s_tir",
+            "def apply_trace(sch: s_tir.Schedule) -> None:",
+            '  b0 = sch.get_sblock(name="B", func_name="main")',
+            "  l1, = sch.get_loops(block=b0)",
+            "  l2, l3 = sch.split(loop=l1, factors=[None, 32], preserve_unit_iters=True, disable_predication=False)",
+        )
+    )
+
+
+def test_apply_json_to_schedule_1():
+    trace = _make_trace_2(SBlockRV())
+    json_obj = trace.as_json()
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    Trace.apply_json_to_schedule(json_obj, sch)
+    assert_structural_equal_ignore_global_symbol(elementwise_inlined, sch.mod["main"])
+
+
+def test_apply_json_to_schedule_sample_categorical():
+    var = tirx.Var("v", "int32")
+    trace1 = Trace(
+        insts=[
+            Instruction(
+                kind=InstructionKind.get("SampleCategorical"),
+                inputs=[],
+                attrs=[[tvm.tirx.IntImm("int32", 3)], [tvm.tirx.FloatImm("float32", 1.0)]],
+                outputs=[var],
+            )
+        ],
+        decisions={},
+    )
+    json = trace1.as_json()
+    assert str(json) == "[[['SampleCategorical', [], [[3], [T.float32(1.0)]], ['v0']]], []]"
+
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    # As long as the application does not fail, it is fine.
+    Trace.apply_json_to_schedule(json, sch)
+    python_str = sch.trace.as_python()
+    assert len(python_str) == 1
+    assert python_str[0] == "v0 = sch.sample_categorical(candidates=[3], probs=[1], decision=0)"
+
+
+def _test_apply_annotation_trace_from_json(annotation: str):
+    """Test applying an annotation works without crashing.
+
+    Designed to handle some previously failing edge cases like the
+    empty string.
+    """
+    b0 = SBlockRV()
+    trace = Trace(
+        insts=[
+            _make_get_sblock(name="B", output=b0),
+            _make_annotate(block=b0, annotation=annotation),
+        ],
+        decisions={},
+    )
+    json_obj = trace.as_json()
+    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
+    Trace.apply_json_to_schedule(json_obj, sch)
+
+    @T.prim_func(s_tir=True)
+    def elementwise_expected(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (128, 128))
+        B = T.sblock_alloc_buffer((128, 128))
+        C = T.match_buffer(c, (128, 128))
+        for i, j in T.grid(128, 128):
+            with T.sblock("B"):
+                T.sblock_attr({"meta_schedule.auto_tensorize": annotation})
+                vi, vj = T.axis.remap("SS", [i, j])
+                B[vi, vj] = A[vi, vj] * 2.0
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vi, vj] = B[vi, vj] + 1.0
+
+    assert_structural_equal_ignore_global_symbol(elementwise_expected, sch.mod["main"])
+
+
+def test_apply_annotation_from_json():
+    # Something reasonable
+    _test_apply_annotation_trace_from_json("SSRSSR")
+
+    # The empty string
+    _test_apply_annotation_trace_from_json("")
+
+    # A string of two quotation marks
+    _test_apply_annotation_trace_from_json('""')
+
+    # A string of one quotation mark
+    _test_apply_annotation_trace_from_json('"')
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/schedule/test_tir_schedule_transform.py

@@ -0,0 +1,179 @@
+# 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.
+import tvm
+from tvm.s_tir import Schedule
+from tvm.s_tir.schedule.transform import tile_with_tensor_intrin
+from tvm.s_tir.tensor_intrin.x86 import AVX512_DOT_16x4_INTRIN, VNNI_DOT_16x4_INTRIN
+from tvm.script import tirx as T
+
+
+@tvm.script.ir_module
+class DenseTIRModule:
+    @T.prim_func(s_tir=True)
+    def main(
+        placeholder: T.Buffer((1024, 1024), "uint8"),
+        placeholder_1: T.Buffer((64, 256, 16, 4), "int8"),
+        compute: T.Buffer((1024, 1024), "int32"),
+    ) -> None:
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        with T.sblock("root"):
+            T.reads()
+            T.writes()
+            for i0, i1, i2 in T.grid(1024, 1024, 1024):
+                with T.sblock("compute"):
+                    i, j, k = T.axis.remap("SSR", [i0, i1, i2])
+                    T.reads(placeholder[i, k], placeholder_1[j // 16, k // 4, j % 16, k % 4])
+                    T.writes(compute[i, j])
+                    with T.init():
+                        compute[i, j] = 0
+                    compute[i, j] = compute[i, j] + T.cast(placeholder[i, k], "int32") * T.cast(
+                        placeholder_1[j // 16, k // 4, j % 16, k % 4], "int32"
+                    )
+
+
+@tvm.script.ir_module
+class DenseTIRModuleTiled:
+    @T.prim_func(s_tir=True)
+    def main(
+        placeholder: T.Buffer((1024, 1024), "uint8"),
+        placeholder_1: T.Buffer((64, 256, 16, 4), "int8"),
+        compute: T.Buffer((1024, 1024), "int32"),
+    ) -> None:
+        # function attr dict
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        # body
+        # with T.sblock("root")
+        for i0, i1_0, i2_0, i1_1, i2_1 in T.grid(1024, 64, 256, 16, 4):
+            with T.sblock("compute"):
+                i = T.axis.spatial(1024, i0)
+                j = T.axis.spatial(1024, i1_0 * 16 + i1_1)
+                k = T.axis.reduce(1024, i2_0 * 4 + i2_1)
+                T.reads(placeholder[i, k], placeholder_1[j // 16, k // 4, j % 16, k % 4])
+                T.writes(compute[i, j])
+                with T.init():
+                    compute[i, j] = 0
+                compute[i, j] = compute[i, j] + T.cast(placeholder[i, k], "int32") * T.cast(
+                    placeholder_1[j // 16, k // 4, j % 16, k % 4], "int32"
+                )
+
+
+@tvm.script.ir_module
+class Conv2dNCHWcTIRModule:
+    @T.prim_func(s_tir=True)
+    def main(
+        placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"),
+        placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"),
+        conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32"),
+    ) -> None:
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i0, i1, i2, i3, i4, i5, i6, i7, i8, i9 in T.grid(1, 16, 56, 56, 16, 1, 1, 4, 4, 4):
+            with T.sblock("conv2d_NCHWc_int8"):
+                (
+                    n,
+                    oc_chunk,
+                    oh,
+                    ow,
+                    oc_block,
+                    kh,
+                    kw,
+                    ic_outer,
+                    ic_f_inner,
+                    ic_s_inner,
+                ) = T.axis.remap("SSSSSRRRRR", [i0, i1, i2, i3, i4, i5, i6, i7, i8, i9])
+                T.reads(
+                    placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
+                    placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
+                )
+                T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block])
+                with T.init():
+                    conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = 0
+                conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = conv2d_NCHWc_int8[
+                    n, oc_chunk, oh, ow, oc_block
+                ] + T.cast(
+                    placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
+                    "int32",
+                ) * T.cast(
+                    placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
+                    "int32",
+                )
+
+
+@tvm.script.ir_module
+class Conv2dNCHWcTIRModuleTiled:
+    @T.prim_func(s_tir=True)
+    def main(
+        placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"),
+        placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"),
+        conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32"),
+    ) -> None:
+        # function attr dict
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        # body
+        # with T.sblock("root")
+        for i0, i1, i2, i3, i4_0, i5, i6, i7, i8, i9_0, i4_1, i9_1 in T.grid(
+            1, 16, 56, 56, 1, 1, 1, 4, 4, 1, 16, 4
+        ):
+            with T.sblock("conv2d_NCHWc_int8"):
+                n, oc_chunk, oh, ow = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                oc_block = T.axis.spatial(16, i4_0 * 16 + i4_1)
+                kh, kw, ic_outer, ic_f_inner = T.axis.remap("RRRR", [i5, i6, i7, i8])
+                ic_s_inner = T.axis.reduce(4, i9_0 * 4 + i9_1)
+                T.reads(
+                    placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
+                    placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
+                )
+                T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block])
+                with T.init():
+                    conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = 0
+                conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = conv2d_NCHWc_int8[
+                    n, oc_chunk, oh, ow, oc_block
+                ] + T.cast(
+                    placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
+                    "int32",
+                ) * T.cast(
+                    placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
+                    "int32",
+                )
+
+
+def test_tile_with_tensor_intrin_dense(intrin=VNNI_DOT_16x4_INTRIN):
+    s = Schedule(DenseTIRModule)
+    block = s.get_sblock("compute")
+
+    tiled_loop = tile_with_tensor_intrin(s, block, intrin)
+
+    _, _, _, i1_1, _ = s.get_loops(block)
+
+    assert s.get(tiled_loop) == s.get(i1_1)
+    tvm.ir.assert_structural_equal(s.mod, DenseTIRModuleTiled)
+
+
+def test_tile_with_tensor_intrin_conv2d_nchwc(intrin=VNNI_DOT_16x4_INTRIN):
+    s = Schedule(Conv2dNCHWcTIRModule)
+    block = s.get_sblock("conv2d_NCHWc_int8")
+    tiled_loop = tile_with_tensor_intrin(s, block, intrin)
+    tiled_loops = s.get_loops(block)
+    assert len(tiled_loops) == 12
+    assert s.get(tiled_loop) == s.get(tiled_loops[-2])
+    tvm.ir.assert_structural_equal(s.mod, Conv2dNCHWcTIRModuleTiled)
+
+
+if __name__ == "__main__":
+    test_tile_with_tensor_intrin_dense()
+    test_tile_with_tensor_intrin_dense(AVX512_DOT_16x4_INTRIN)
+    test_tile_with_tensor_intrin_conv2d_nchwc()
+    test_tile_with_tensor_intrin_conv2d_nchwc(AVX512_DOT_16x4_INTRIN)

tests/python/s_tir/schedule/test_tir_schedule_utilities.py

@@ -0,0 +1,427 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401, F841
+import sys
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.ir import IRModule
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j in T.grid(128, 128):
+        with T.sblock("init"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = 0.0
+        for k in range(0, 128):
+            with T.sblock("update"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def matmul_relu(a: T.handle, b: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (1024, 1024))
+    B = T.match_buffer(b, (1024, 1024))
+    C = T.sblock_alloc_buffer((1024, 1024))
+    D = T.match_buffer(d, (1024, 1024))
+    for i, j, k in T.grid(1024, 1024, 1024):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
+    for i, j in T.grid(1024, 1024):
+        with T.sblock("relu"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = T.max(C[vi, vj], 0.0)
+
+
+@T.prim_func(s_tir=True)
+def matmul_relu_ann1(a: T.handle, b: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (1024, 1024))
+    B = T.match_buffer(b, (1024, 1024))
+    C = T.sblock_alloc_buffer((1024, 1024))
+    D = T.match_buffer(d, (1024, 1024))
+    for i in T.serial(0, 1024, annotations={"test1": "aaa", "test4": {"arr": [0, 0], "key": 3}}):
+        for j in T.serial(0, 1024, annotations={"test2": 612, "test3": ["aa", 1]}):
+            for k in T.serial(0, 1024):
+                with T.sblock("matmul"):
+                    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                    with T.init():
+                        C[vi, vj] = 0.0
+                    C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
+    for i, j in T.grid(1024, 1024):
+        with T.sblock("relu"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            D[vi, vj] = T.max(C[vi, vj], 0.0)
+
+
+@T.prim_func(s_tir=True)
+def matmul_relu_ann2(a: T.handle, b: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (1024, 1024))
+    B = T.match_buffer(b, (1024, 1024))
+    C = T.sblock_alloc_buffer((1024, 1024))
+    D = T.match_buffer(d, (1024, 1024))
+    for i, j, k in T.grid(1024, 1024, 1024):
+        with T.sblock("matmul"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            T.sblock_attr({"test1": "aaa", "test4": {"arr": [0, 0], "key": 3}})
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
+    for i, j in T.grid(1024, 1024):
+        with T.sblock("relu"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.sblock_attr({"test2": 0.22, "test3": ["aa", 1]})
+            D[vi, vj] = T.max(C[vi, vj], 0.0)
+
+
+@tvm.script.ir_module
+class ModuleWithMultipleFuncs:
+    @T.prim_func(s_tir=True)
+    def vector_add(
+        A: T.Buffer(128, "float32"),
+        B: T.Buffer(128, "float32"),
+    ) -> None:
+        for i in range(128):
+            with T.sblock("init"):
+                vi = T.axis.remap("S", [i])
+                B[vi] = A[vi]
+
+    @T.prim_func(s_tir=True)
+    def vector_add_2(
+        A: T.Buffer(128, "float32"),
+        B: T.Buffer(128, "float32"),
+    ) -> None:
+        for i in range(128):
+            with T.sblock("init"):
+                vi = T.axis.remap("S", [i])
+                B[vi] = A[vi]
+
+
+@T.prim_func(s_tir=True)
+def tuple_reduction(data: T.Buffer((4, 32), "float32"), T_add: T.Buffer((4,), "float32")) -> None:
+    # function attr dict
+    T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+    # body
+    with T.sblock("root"):
+        T.reads()
+        T.writes()
+        data_red_temp_v0 = T.sblock_alloc_buffer([4], dtype="float32")
+        data_red_temp_v1 = T.sblock_alloc_buffer([4], dtype="float32")
+        for i0, i1 in T.grid(4, 32):
+            with T.sblock("data_red_temp"):
+                ax0, k1 = T.axis.remap("SR", [i0, i1])
+                T.reads(data[ax0, k1])
+                T.writes(data_red_temp_v0[ax0], data_red_temp_v1[ax0])
+                with T.init():
+                    data_red_temp_v0[ax0] = T.float32(0)
+                    data_red_temp_v1[ax0] = T.float32(0)
+                v_data_red_temp_v0: T.let[T.float32] = data_red_temp_v0[ax0] + data[ax0, k1]
+                v_data_red_temp_v1: T.let[T.float32] = (
+                    data_red_temp_v1[ax0] + data[ax0, k1] * data[ax0, k1]
+                )
+                data_red_temp_v0[ax0] = v_data_red_temp_v0
+                data_red_temp_v1[ax0] = v_data_red_temp_v1
+        for i0 in range(4):
+            with T.sblock("T_add"):
+                ax0 = T.axis.remap("S", [i0])
+                T.reads(data_red_temp_v0[ax0], data_red_temp_v1[ax0])
+                T.writes(T_add[ax0])
+                T_add[ax0] = data_red_temp_v0[ax0] + data_red_temp_v1[ax0]
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+use_block_name = tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True})
+
+
+def test_tir_schedule_creation():
+    # Tests:
+    # - Schedule.__init__ for PrimFunc and IRModule
+    # - Schedule.mod
+    # - Schedule.state
+    sch_1 = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    sch_2 = tvm.s_tir.Schedule(IRModule({"main": matmul}), debug_mask="all")
+    assert sch_1.mod["main"].same_as(sch_2.mod["main"])
+    assert sch_1.state.mod["main"].same_as(sch_2.state.mod["main"])
+
+
+def test_tir_schedule_get_sblock():
+    # Tests:
+    # - Schedule.get_sblock
+    # - Schedule.get_sref
+    # - Schedule.get
+    sch = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    block_rv = sch.get_sblock(name="update")
+    block_sref = sch.get_sref(block_rv)
+    block = sch.get(block_rv)
+    assert block.name_hint == "update"
+    assert block_sref.stmt.same_as(block)
+    assert sch.state.get_sref(block).same_as(block_sref)
+    assert block.same_as(matmul.body.block.body.body.body[1].body.block)
+
+
+def test_tir_schedule_work_on():
+    sch = tvm.s_tir.Schedule(ModuleWithMultipleFuncs, debug_mask="all")
+    with pytest.raises(ValueError, match="does not know which function to be working on"):
+        sch.get_sblock(name="init")
+    sch.work_on(func_name="vector_add")
+    sch.get_sblock(name="init")
+    assert sch.func_working_on == sch.mod.get_global_var("vector_add")
+
+
+def test_tir_schedule_get_loops(use_block_name):
+    # Tests:
+    # - Schedule.get_loops
+    # - Schedule.get
+    sch = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    block = "update" if use_block_name else sch.get_sblock(name="update")
+    i, j, k = sch.get_loops(block)
+    assert sch.get(i).loop_var.name == "i"
+    assert sch.get(j).loop_var.name == "j"
+    assert sch.get(k).loop_var.name == "k"
+
+
+def test_tir_schedule_copy_1(use_block_name):
+    # Tests:
+    # - Schedule.copy
+    sch_1 = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    block_rv = sch_1.get_sblock(name="update")
+    i, j, k = sch_1.get_loops(block="update" if use_block_name else block_rv)
+    assert sch_1.get(i).loop_var.name == "i"
+    assert sch_1.get(j).loop_var.name == "j"
+    assert sch_1.get(k).loop_var.name == "k"
+
+    sch_2 = sch_1.copy()
+    assert sch_2.get(block_rv).name_hint == "update"
+    assert sch_2.get(i).loop_var.name == "i"
+    assert sch_2.get(j).loop_var.name == "j"
+    assert sch_2.get(k).loop_var.name == "k"
+
+
+def test_tir_schedule_copy_2():
+    sch = tvm.s_tir.Schedule(mod=matmul, debug_mask="all")
+    i, j, k = sch.get_loops(sch.get_sblock("update"))
+    sch_copy = sch.copy()
+    assert not sch.get_sref(i).same_as(sch_copy.get_sref(i))
+    assert not sch.get_sref(j).same_as(sch_copy.get_sref(j))
+    assert not sch.get_sref(k).same_as(sch_copy.get_sref(k))
+    assert sch.get_sref(i).stmt.same_as(sch_copy.get_sref(i).stmt)
+    assert sch.get_sref(j).stmt.same_as(sch_copy.get_sref(j).stmt)
+    assert sch.get_sref(k).stmt.same_as(sch_copy.get_sref(k).stmt)
+    i_0, i_1 = sch.split(i, factors=[None, 64])
+    j_0, j_1 = sch_copy.split(j, factors=[None, 32])
+
+    assert sch.get_sref(i_0).stmt.extent == 2
+    assert sch.get_sref(i_1).stmt.extent == 64
+    with pytest.raises(IndexError):
+        sch_copy.get_sref(i_0)
+    with pytest.raises(IndexError):
+        sch_copy.get_sref(i_1)
+
+    with pytest.raises(IndexError):
+        sch.get_sref(j_0)
+    with pytest.raises(IndexError):
+        sch.get_sref(j_1)
+    assert sch_copy.get_sref(j_0).stmt.extent == 4
+    assert sch_copy.get_sref(j_1).stmt.extent == 32
+    verify_trace_roundtrip(sch, mod=matmul)
+    verify_trace_roundtrip(sch_copy, mod=matmul)
+
+
+def test_tir_schedule_remove_rv():
+    # Tests:
+    # - Schedule.remove_rv
+    sch = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    block_rv = sch.get_sblock(name="update")
+    assert sch.get(block_rv).name_hint == "update"
+    sch.remove_rv(block_rv)
+    with pytest.raises(IndexError):
+        sch.get(block_rv)
+
+
+def test_get_child_blocks():
+    s = tvm.s_tir.Schedule(matmul, debug_mask="all")
+    init = s.get_sblock("init")
+    update = s.get_sblock("update")
+    # loop
+    blocks = s.get_child_blocks(s.get_loops(init)[0])
+    assert len(blocks) == 2
+    assert s.get(init) == s.get(blocks[0])
+    assert s.get(update) == s.get(blocks[1])
+    # block
+    root = s.get_sblock("root")
+    blocks = s.get_child_blocks(root)
+    assert len(blocks) == 2
+    assert s.get(init) == s.get(blocks[0])
+    assert s.get(update) == s.get(blocks[1])
+
+
+def test_get_producers(use_block_name):
+    sch = tvm.s_tir.Schedule(mod=matmul_relu, debug_mask="all")
+    block = "relu" if use_block_name else sch.get_sblock("relu")
+    (producer,) = sch.get_producers(block)
+    tvm.ir.assert_structural_equal(
+        sch.get_sref(producer).stmt,
+        sch.get_sref(sch.get_sblock("matmul")).stmt,
+    )
+    verify_trace_roundtrip(sch, mod=matmul_relu)
+
+
+def test_get_producers_multiple_buffer_depdencies(use_block_name):
+    sch = tvm.s_tir.Schedule(mod=tuple_reduction, debug_mask="all")
+    block = "T_add" if use_block_name else sch.get_sblock("T_add")
+    (producer,) = sch.get_producers(block)
+    tvm.ir.assert_structural_equal(
+        sch.get_sref(producer).stmt,
+        sch.get_sref(sch.get_sblock("data_red_temp")).stmt,
+    )
+
+
+def test_get_consumers(use_block_name):
+    sch = tvm.s_tir.Schedule(mod=matmul_relu, debug_mask="all")
+    block = "matmul" if use_block_name else sch.get_sblock("matmul")
+    (consumer,) = sch.get_consumers(block)
+    tvm.ir.assert_structural_equal(
+        sch.get_sref(consumer).stmt,
+        sch.get_sref(sch.get_sblock("relu")).stmt,
+    )
+    verify_trace_roundtrip(sch, mod=matmul_relu)
+
+
+def test_get_consumers_multiple_buffer_depdencies(use_block_name):
+    sch = tvm.s_tir.Schedule(mod=tuple_reduction, debug_mask="all")
+    block = "data_red_temp" if use_block_name else sch.get_sblock("data_red_temp")
+    (consumer,) = sch.get_consumers(block)
+    tvm.ir.assert_structural_equal(
+        sch.get_sref(consumer).stmt,
+        sch.get_sref(sch.get_sblock("T_add")).stmt,
+    )
+
+
+def test_annotate_unannotate_loop():
+    sch = tvm.s_tir.Schedule(mod=matmul_relu, debug_mask="all")
+    matmul = sch.get_sblock("matmul")
+    relu = sch.get_sblock("relu")
+    sch.annotate(sch.get_loops(matmul)[0], "test1", "aaa")
+    sch.annotate(sch.get_loops(matmul)[1], "test2", 612)
+    sch.annotate(sch.get_loops(matmul)[1], "test3", ["aa", 1])
+    sch.annotate(sch.get_loops(matmul)[0], "test4", {"arr": [0, 0], "key": 3})
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], matmul_relu_ann1)
+    verify_trace_roundtrip(sch=sch, mod=matmul_relu)
+    sch.unannotate(sch.get_loops(matmul)[0], "test1")
+    sch.unannotate(sch.get_loops(matmul)[1], "test2")
+    sch.unannotate(sch.get_loops(matmul)[1], "test3")
+    sch.unannotate(sch.get_loops(matmul)[0], "test4")
+    verify_trace_roundtrip(sch=sch, mod=matmul_relu)
+
+
+def test_annotate_unannotate_block():
+    sch = tvm.s_tir.Schedule(mod=matmul_relu, debug_mask="all")
+    matmul = sch.get_sblock("matmul")
+    relu = sch.get_sblock("relu")
+    sch.annotate(matmul, "test1", "aaa")
+    sch.annotate(relu, "test2", 0.22)
+    sch.annotate(relu, "test3", ["aa", 1])
+    sch.annotate(matmul, "test4", {"arr": [0, 0], "key": 3})
+    assert_structural_equal_ignore_global_symbol(sch.mod["main"], matmul_relu_ann2)
+    verify_trace_roundtrip(sch=sch, mod=matmul_relu)
+    sch.unannotate(matmul, "test1")
+    sch.unannotate(relu, "test2")
+    sch.unannotate(relu, "test3")
+    sch.unannotate(matmul, "test4")
+    verify_trace_roundtrip(sch=sch, mod=matmul_relu)
+
+
+def test_get_output_blocks_single_output():
+    sch = tvm.s_tir.Schedule(mod=matmul_relu, debug_mask="all")
+    output_blocks = sch.get_output_blocks("root")
+    assert len(output_blocks) == 1, "Unexpected number of blocks when 1 was expected"
+    block = sch.get(output_blocks[0])
+    assert block.name_hint == "relu"
+    relu_block = sch.get_sblock("relu")
+    assert sch.get(relu_block).same_as(block)
+
+
+def test_get_output_blocks_multiple_outputs():
+    sch = tvm.s_tir.Schedule(mod=matmul, debug_mask="all")
+    output_blocks = sch.get_output_blocks("root")
+    assert len(output_blocks) == 2, "Unexpected number of blocks when 2 were expected"
+    block_1 = sch.get(output_blocks[0])
+    assert block_1.name_hint == "init"
+    block_2 = sch.get(output_blocks[1])
+    assert block_2.name_hint == "update"
+    init_block = sch.get_sblock("init")
+    assert sch.get(init_block).same_as(block_1)
+    update_block = sch.get_sblock("update")
+    assert sch.get(update_block).same_as(block_2)
+
+
+def test_get_output_blocks_nested():
+    @T.prim_func(s_tir=True)
+    def blockized(
+        A: T.Buffer((128, 128), "float32"),
+        B: T.Buffer((128, 128), "float32"),
+    ) -> None:
+        with T.sblock("blockized_B"):
+            vio = T.axis.spatial(1, 0)
+            vjo = T.axis.spatial(1, 0)
+            for i, j in T.grid(128, 128):
+                with T.sblock("B"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    B[vi, vj] = A[vi, vj] * 2.0
+
+    sch = tvm.s_tir.Schedule(mod=blockized, debug_mask="all")
+    output_blocks = sch.get_output_blocks("root")
+    assert len(output_blocks) == 2, "Unexpected number of blocks when 2 were expected"
+    block_1 = sch.get(output_blocks[0])
+    assert block_1.name_hint == "blockized_B"
+    block_2 = sch.get(output_blocks[1])
+    assert block_2.name_hint == "B"
+    blockized_block = sch.get_sblock("blockized_B")
+    assert sch.get(blockized_block).same_as(block_1)
+    b_block = sch.get_sblock("B")
+    assert sch.get(b_block).same_as(block_2)
+
+    sch = tvm.s_tir.Schedule(mod=blockized, debug_mask="all")
+    output_blocks = sch.get_output_blocks("blockized_B")
+    assert len(output_blocks) == 1, "Unexpected number of blocks when 1 were expected"
+    block = sch.get(output_blocks[0])
+    assert block.name_hint == "B"
+    b_block = sch.get_sblock("B")
+    assert sch.get(b_block).same_as(block)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()