chore: import upstream snapshot with attribution

tests/python/s_tir/analysis/test_s_tir_analysis_calculate_allocated_memory.py

@@ -0,0 +1,141 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
+# ruff: noqa: F401
+import pytest
+
+import tvm
+from tvm import tirx
+from tvm.script import tirx as T
+
+# fmt: off
+# pylint: disable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
+
+@tvm.script.ir_module
+class Module:
+    @T.prim_func(s_tir=True)
+    def scale_by_two(a: T.Buffer((128,), "int8"), c: T.Buffer((128,), "int8")):
+        for i in T.serial(128):
+            with T.sblock("C"):
+                c[i] = a[i] * T.int8(2)
+
+
+    @T.prim_func(s_tir=True)
+    def scale_by_two_three(a: T.Buffer((128,), "int8"), c: T.Buffer((128,), "int8")):
+        B = T.sblock_alloc_buffer([128], dtype="int8", scope="global.vtcm")
+        for i in T.serial(128):
+            with T.sblock("B"):
+                B[i] = a[i] * T.int8(2)
+        for i in T.serial(128):
+            with T.sblock("C"):
+                c[i] = B[i] * T.int8(3)
+
+# pylint: enable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
+# fmt: on
+
+
+@pytest.mark.parametrize(
+    "primFunc,size", [(Module["scale_by_two"], 128), (Module["scale_by_two_three"], 256)]
+)
+def test_scale_by(primFunc, size):
+    """Test calculate allocated bytes per scope"""
+    mod = tvm.IRModule.from_expr(primFunc.with_attr("global_symbol", "main"))
+    sch = tvm.s_tir.Schedule(mod, debug_mask="all")
+    block_c = sch.get_sblock("C")
+    (flat,) = sch.get_loops(block_c)
+    cache_block = sch.cache_read(block_c, 0, storage_scope="global.vtcm")
+    sch.compute_at(cache_block, flat)
+
+    mod = sch.mod
+    mod = tvm.s_tir.transform.ConvertBlocksToOpaque()(mod)
+    mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+    sizes = tvm.s_tir.analysis.calculate_allocated_bytes(mod["main"])
+    assert "main" in sizes, 'Calls with PrimFunc is expected to return with function key as "main"'
+    sizes = sizes["main"]
+    assert sizes.get("global.vtcm", 0) == size
+
+
+@T.prim_func(s_tir=True)
+def matmul_mix_scope(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128], scope="global")
+    B = T.match_buffer(b, [128, 128], scope="global")
+    C = T.match_buffer(c, [128, 128], scope="global")
+    A_allocated = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="global.texture")
+    B_allocated = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="global.texture")
+    C_allocated = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="global")
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("A.allocated"):
+            A_allocated[i, j] = A[i, j]
+    for i, j in T.grid(128, 128):
+        with T.sblock("B.allocated"):
+            B_allocated[i, j] = B[i, j]
+
+    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_allocated[vi, vj] = 0.0
+            C_allocated[vi, vj] = C[vi, vj] + A_allocated[vi, vk] * B_allocated[vj, vk]
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            C[i, j] = C_allocated[i, j]
+
+
+@pytest.mark.parametrize(
+    "scope,size", [("global", 65536), ("global.texture", 131072), ("global.texture-nhwc", 0)]
+)
+def test_matmul_mix_scope(scope, size):
+    """Test calculate allocated bytes per scope"""
+    mod = tvm.IRModule({"main": matmul_mix_scope})
+    mod = tvm.s_tir.transform.LowerInitBlock()(mod)
+    mod = tvm.s_tir.transform.ConvertBlocksToOpaque()(mod)
+    mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+    sizes = tvm.s_tir.analysis.calculate_allocated_bytes(mod["main"])
+    assert "main" in sizes, 'Calls with PrimFunc is expected to return with function key as "main"'
+    sizes = sizes["main"]
+    assert sizes.get(scope, 0) == size
+
+
+def test_full_mod_calculator():
+    def apply_schedule(sch, func_name):
+        sch.work_on(func_name)
+        block_c = sch.get_sblock("C")
+        sch.cache_read(block_c, 0, storage_scope="global.vtcm")
+
+    sch = tvm.s_tir.Schedule(Module, debug_mask="all")
+    apply_schedule(sch, "scale_by_two")
+    apply_schedule(sch, "scale_by_two_three")
+    mod = tvm.s_tir.transform.ConvertBlocksToOpaque()(sch.mod)
+    mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+    sizes = tvm.s_tir.analysis.calculate_allocated_bytes(mod)
+    assert "scale_by_two" in sizes, "Values for scale_by_two not found"
+    scale_by_two_sizes = sizes["scale_by_two"]
+    assert "global.vtcm" in scale_by_two_sizes, (
+        "Expected global.vtcm allocation to be calculated scale_by_two"
+    )
+    assert scale_by_two_sizes["global.vtcm"] == 128, "Expected the calculated size to be 128"
+    scale_by_two_three_sizes = sizes["scale_by_two_three"]
+    assert "global.vtcm" in scale_by_two_three_sizes, (
+        "Expected global.vtcm allocation to be calculated scale_by_two_three"
+    )
+    assert scale_by_two_three_sizes["global.vtcm"] == 256, "Expected the calculated size to be 256"
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py

@@ -0,0 +1,145 @@
+# 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
+
+import pytest
+
+import tvm.testing
+from tvm.ir import IRModule
+from tvm.s_tir.analysis import estimate_tir_flops
+from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
+from tvm.script import tirx as T
+
+
+@pytest.mark.parametrize(
+    "workload, flops",
+    [
+        ("C1D", 6291456),
+        ("C2D", 236027904),
+        ("C3D", 13217562624),
+        ("CAP", 75497472),
+        ("DEP", 7225344),
+        ("DIL", 223552896),
+        ("GMM", 4194304),
+        ("GRP", 28901376),
+        ("T2D", 268435456),
+        ("CBR", 239239168),
+        ("TBG", 25165824),
+        ("NRM", 131072),
+        ("SFM", 262144),
+    ],
+)
+def test_te_workload(workload, flops):
+    te_workload = create_te_workload(workload, 0)
+    mod = IRModule({"main": te_workload})
+    assert float(flops) == estimate_tir_flops(mod)
+
+
+@T.prim_func(s_tir=True)
+def flops_with_let(a: T.Buffer(16, "float32")):
+    for i in range(8):
+        j = i + 8
+        a[j] = a[i]
+
+
+def test_flops_with_let():
+    flops = estimate_tir_flops(IRModule({"main": flops_with_let}))
+    assert flops == 8
+
+
+@T.prim_func(s_tir=True)
+def flops_with_if(a: T.Buffer(16, "float32"), b: T.Buffer(16, "float32")):
+    for i in range(16):
+        if i % 2 == 0:
+            a[i] = b[i]
+        else:
+            if i % 3 == 0:
+                a[i] = b[i - 1] + b[i - 2]
+
+
+def test_flops_with_if():
+    flops = estimate_tir_flops(IRModule({"main": flops_with_if}))
+    assert flops == 16
+
+
+@T.prim_func(s_tir=True)
+def flops_with_forloop_as_expression(A: T.Buffer(1)):
+    for i in T.serial(0, 16):
+        for k in T.serial(0, i):
+            A[0] = A[0] + 1
+
+
+@T.prim_func(s_tir=True)
+def flops_override(A: T.Buffer(16, "float32")):
+    T.func_attr({"estimated_flops": 32})
+    for i in range(16):
+        A[0] = A[0] + 1
+
+
+def test_estimate_flops_forloop_as_expression():
+    flops = estimate_tir_flops(
+        IRModule({"main": flops_with_forloop_as_expression.with_attr("estimated_flops", 32)})
+    )
+    assert flops == 32
+
+    # test whether the user estimated flop would over ride
+    flops = estimate_tir_flops(IRModule({"main": flops_override}))
+    assert flops == 32
+
+
+def test_estimate_flops_with_decl_buffer():
+    def make_func(use_decl_buffer):
+        buffer_func = T.decl_buffer if use_decl_buffer else T.Buffer
+
+        @T.prim_func(s_tir=True)
+        def func(A_data: T.handle("float32")):
+            A = buffer_func(16, "float32", data=A_data)
+            for i in range(16):
+                A[0] = A[0] + 1
+
+        return func
+
+    flops_with_decl_buffer = estimate_tir_flops(IRModule.from_expr(make_func(True)))
+    flops_without_decl_buffer = estimate_tir_flops(IRModule.from_expr(make_func(True)))
+    assert flops_with_decl_buffer == flops_without_decl_buffer
+
+
+@T.prim_func(s_tir=True)
+def flops_with_nonint_extent(a: T.Buffer(16, "float32")):
+    for i in range(4 + 4):
+        a[i] = 2 * a[i]
+
+
+def test_flops_with_nonint_extent():
+    assert estimate_tir_flops(IRModule({"main": flops_with_nonint_extent})) == 8
+
+
+@T.prim_func(s_tir=True)
+def flops_with_variable_extent(a: T.Buffer(16, "float32")):
+    for i in range(4 + 4):
+        for j in range(i + 8):
+            a[j] = 2 * a[i]
+
+
+def test_flops_with_variable_extent():
+    assert estimate_tir_flops(IRModule({"main": flops_with_variable_extent})) == 120
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py

@@ -0,0 +1,335 @@
+# 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 re
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, StringImm
+
+identify_memcpy = tvm.s_tir.analysis._ffi_api._identify_memcpy
+
+
+def _check_memcpy_results(func, expected):
+    """Check that identify_memcpy returns the expected results."""
+    results = identify_memcpy(func.body)
+
+    if isinstance(expected, str) or (
+        isinstance(expected, tuple) and isinstance(expected[0], BufferRegion)
+    ):
+        expected = [expected]
+
+    assert len(expected) == len(results)
+    for exp, result in zip(expected, results):
+        if isinstance(exp, str):
+            assert isinstance(result, StringImm)
+            assert re.search(exp, result.value)
+        else:
+            tvm.ir.assert_structural_equal(result, exp)
+
+
+def test_1d():
+    """Simplest test case"""
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(1024, "float32"), B: T.Buffer(1024, "float32")):
+        for i in T.serial(1024):
+            B[i] = A[i]
+
+    A, B = func.buffer_map.values()
+    expected = (A[0:1024], B[0:1024])
+    _check_memcpy_results(func, expected)
+
+
+def test_1d_compute():
+    """Like test_1d, but a computation prevents this being a memcpy"""
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(1024, "float32"), B: T.Buffer(1024, "float32")):
+        for i in T.serial(1024):
+            B[i] = A[i] + 1.0
+
+    expected = "Expected BufferStore's value to be BufferLoad"
+    _check_memcpy_results(func, expected)
+
+
+def test_1d_conditional():
+    """Like test_1d, but a conditionals prevents this being a memcpy"""
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(1024, "float32"), B: T.Buffer(1024, "float32")):
+        for i in T.serial(1024):
+            if i < 1024:
+                B[i] = A[i]
+
+    expected = "Expected innermost loop to have BufferStore body"
+    _check_memcpy_results(func, expected)
+
+
+def test_1d_strided_input():
+    """Like test_1d, but strided input prevents this being a memcpy"""
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(2048, "float32"), B: T.Buffer(1024, "float32")):
+        for i in T.serial(1024):
+            B[i] = A[i * 2]
+
+    expected = "Mismatch between loop iterations (.*) and number of src indices"
+    _check_memcpy_results(func, expected)
+
+
+def test_1d_strided_output():
+    """Like test_1d, but strided output prevents this being a memcpy"""
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(1024, "float32"), B: T.Buffer(2048, "float32")):
+        for i in T.serial(1024):
+            B[i * 2] = A[i]
+
+    expected = "Mismatch between loop iterations (.*) and number of dst indices"
+    _check_memcpy_results(func, expected)
+
+
+def test_1d_input_2d_output_fused_loop():
+    """Like test_1d, but the output is written as a 2-d buffer"""
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(1024, "float32"), B: T.Buffer((32, 32), "float32")):
+        for i in T.serial(1024):
+            B[i // 32, i % 32] = A[i]
+
+    A, B = func.buffer_map.values()
+    expected = (A[0:1024], B[0:32, 0:32])
+    _check_memcpy_results(func, expected)
+
+
+def test_2d_input_1d_output_fused_loop():
+    """Like test_1d, but the input is written as a 2-d buffer"""
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer(1024, "float32")):
+        for i in T.serial(1024):
+            B[i] = A[i // 32, i % 32]
+
+    A, B = func.buffer_map.values()
+    expected = (A[0:32, 0:32], B[0:1024])
+    _check_memcpy_results(func, expected)
+
+
+def test_1d_input_1d_output_nested_loop():
+    """Like test_1d, but the iterator is written as a nested loop
+
+    In test cases with more than one loop, each loop is checked to see
+    if could be written as a memcpy.  The C++ utility function
+    operates on individual loops, but for unit testing in Python, it
+    is more convenient to return the results for all loops.
+    """
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(1024, "float32"), B: T.Buffer(1024, "float32")):
+        for i, j in T.grid(32, 32):
+            B[i * 32 + j] = A[i * 32 + j]
+
+    A, B = func.buffer_map.values()
+    i = func.body.loop_var
+    expected = [
+        (A[0:1024], B[0:1024]),
+        (A[i * 32 : i * 32 + 32], B[i * 32 : i * 32 + 32]),
+    ]
+    _check_memcpy_results(func, expected)
+
+
+def test_1d_input_1d_output_nested_loop_equivalent_expressions():
+    """Like test_1d_input_1d_output_nested_loop, but with equivalent indices
+
+    If the expressions are not identical, the loops may still be
+    recognizable as a memcpy, so long as the expressions are
+    equivalent.
+    """
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(1024, "float32"), B: T.Buffer(1024, "float32")):
+        for i, j in T.grid(32, 32):
+            B[i * 32 + j] = A[j + i * 32]
+
+    A, B = func.buffer_map.values()
+    i = func.body.loop_var
+    expected = [
+        (A[0:1024], B[0:1024]),
+        (A[i * 32 : i * 32 + 32], B[i * 32 : i * 32 + 32]),
+    ]
+    _check_memcpy_results(func, expected)
+
+
+def test_1d_input_2d_output_nested_loop():
+    """Like test_1d_input_1d_output_nested_loop, but with a 2-d output buffer"""
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(1024, "float32"), B: T.Buffer((32, 32), "float32")):
+        for i, j in T.grid(32, 32):
+            B[i, j] = A[i * 32 + j]
+
+    A, B = func.buffer_map.values()
+    i = func.body.loop_var
+    expected = [
+        (A[0:1024], B[0:32, 0:32]),
+        (A[i * 32 : i * 32 + 32], B[i, 0:32]),
+    ]
+    _check_memcpy_results(func, expected)
+
+
+def test_2d_input_1d_output_nested_loop():
+    """Like test_1d_input_1d_output_nested_loop, but with a 2-d input buffer"""
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer(1024, "float32")):
+        for i, j in T.grid(32, 32):
+            B[i * 32 + j] = A[i, j]
+
+    A, B = func.buffer_map.values()
+    i = func.body.loop_var
+    expected = [
+        (A[0:32, 0:32], B[0:1024]),
+        (A[i, 0:32], B[i * 32 : i * 32 + 32]),
+    ]
+    _check_memcpy_results(func, expected)
+
+
+def test_2d_input_2d_output_nested_loop():
+    """Like test_1d_input_1d_output_nested_loop, but with 2-d input/output buffers"""
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer((32, 32), "float32")):
+        for i, j in T.grid(32, 32):
+            B[i, j] = A[i, j]
+
+    A, B = func.buffer_map.values()
+    i = func.body.loop_var
+    expected = [
+        (A[0:32, 0:32], B[0:32, 0:32]),
+        (A[i, 0:32], B[i, 0:32]),
+    ]
+    _check_memcpy_results(func, expected)
+
+
+def test_2d_input_2d_output_transpose_output():
+    """test_2d_input_2d_output_nested_loop, but with a transposed output
+
+    This is not recognized as a memcpy, because it results in a transpose.
+    """
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer((32, 32), "float32")):
+        for i, j in T.grid(32, 32):
+            B[j, i] = A[i, j]
+
+    expected = [
+        "different source",
+        "Mismatch .* number of dst indices touched",
+    ]
+    _check_memcpy_results(func, expected)
+
+
+def test_2d_input_2d_output_transpose_input():
+    """test_2d_input_2d_output_nested_loop, but with a transposed input
+
+    This is not recognized as a memcpy, because it results in a transpose.
+    """
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer((32, 32), "float32")):
+        for i, j in T.grid(32, 32):
+            B[i, j] = A[j, i]
+
+    expected = [
+        "different source",
+        "Mismatch .* number of src indices touched",
+    ]
+    _check_memcpy_results(func, expected)
+
+
+def test_2d_input_2d_output_transpose_both():
+    """test_2d_input_2d_output_nested_loop, but with a transposed input
+
+    The inner loop is not recognized as a memcpy, because it has
+    strided access of both the input and output buffers.  However, the
+    outer loop is still recognized as a memcpy, because the full
+    region has been copied over, even though it occurs out of order.
+    """
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer((32, 32), "float32")):
+        for i, j in T.grid(32, 32):
+            B[j, i] = A[j, i]
+
+    A, B = func.buffer_map.values()
+    expected = [
+        (A[0:32, 0:32], B[0:32, 0:32]),
+        "Mismatch .* number of src indices touched",
+    ]
+    _check_memcpy_results(func, expected)
+
+
+def test_cache_read():
+    """Like test_2d_input_2d_output_nested_loop, but with a 1-d
+
+    The inner loop is a memcpy of a single row at a time.  This
+    pattern would appear when B is a read cache of A.
+    """
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer(32, "float32")):
+        for i, j in T.grid(32, 32):
+            B[j] = A[i, j]
+
+    A, B = func.buffer_map.values()
+    i = func.body.loop_var
+    expected = [
+        "does not form a bijective transform",
+        (A[i, 0:32], B[0:32]),
+    ]
+    _check_memcpy_results(func, expected)
+
+
+def test_cache_write():
+    """Like test_2d_input_2d_output_nested_loop, but with a 1-d
+
+    The inner loop is a memcpy of a single row at a time.  This
+    pattern would appear when A is a write cache of B.
+    """
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(32, "float32"), B: T.Buffer((32, 32), "float32")):
+        for i, j in T.grid(32, 32):
+            B[i, j] = A[j]
+
+    A, B = func.buffer_map.values()
+    i = func.body.loop_var
+    expected = [
+        "does not form a bijective transform",
+        (A[0:32], B[i, 0:32]),
+    ]
+    _check_memcpy_results(func, expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py

@@ -0,0 +1,103 @@
+# 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 pytest
+
+import tvm.testing
+from tvm.s_tir.analysis import assert_pure_function, is_pure_function
+from tvm.script import tirx as T
+
+
+class CheckPureFunction:
+    def test_check_purity(self):
+        assert is_pure_function(self.func)
+
+    def test_assert_purity(self):
+        assert_pure_function(self.func)
+
+
+class CheckImpureFunction:
+    def test_check_purity(self):
+        assert not is_pure_function(self.func)
+
+    def test_assert_purity(self):
+        with pytest.raises(AssertionError):
+            assert_pure_function(self.func)
+
+
+class TestNoOp(CheckPureFunction):
+    @T.prim_func(s_tir=True)
+    def func():
+        pass
+
+
+class TestReturnValue(CheckPureFunction):
+    @T.prim_func(s_tir=True)
+    def func() -> T.int32:
+        T.ret(42)
+
+
+class TestComputeValueAndReturn(CheckPureFunction):
+    @T.prim_func(s_tir=True)
+    def func(N: T.int32, M: T.int32) -> T.int32:
+        T.ret(N * M)
+
+
+class TestReadBufferArgument(CheckPureFunction):
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(16, "float32")) -> T.float32:
+        T.ret(A[0])
+
+
+class TestWriteToBufferArgument(CheckImpureFunction):
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer(16, "float32"), B: T.Buffer(16, "float32")):
+        for i in range(16):
+            B[i] = A[i]
+
+
+class TestWriteToInternalAllocation(CheckPureFunction):
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer([16, 16], "float32")) -> T.float32:
+        Sum = T.decl_buffer([], "float32")
+        Sum[()] = 0.0
+        for i, j in T.grid(16, 16):
+            Sum[()] = Sum[()] + A[i, j]
+
+        T.ret(Sum[()])
+
+
+class TestCallPureBuiltin(CheckPureFunction):
+    @T.prim_func(s_tir=True)
+    def func(x: T.float32) -> T.float32:
+        T.ret(T.cos(x))
+
+
+class TestCallPureExtern(CheckPureFunction):
+    @T.prim_func(s_tir=True)
+    def func():
+        T.call_pure_extern("some_pure_extern_func_name", dtype="void")
+
+
+class TestCallImpureExtern(CheckImpureFunction):
+    @T.prim_func(s_tir=True)
+    def func():
+        T.call_extern("some_impure_extern_func_name", dtype="void")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/analysis/test_s_tir_analysis_oob.py

@@ -0,0 +1,77 @@
+# 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 pytest
+
+import tvm
+from tvm.script import tirx as T
+
+
+@T.prim_func(s_tir=True)
+def bad_load(A: T.Buffer((2, 3), "float32"), B: T.Buffer((3, 2), "float32")):
+    B[0, 0] = A[2, 2]
+
+
+@T.prim_func(s_tir=True)
+def bad_load_loop(A: T.Buffer((2, 3), "float32"), B: T.Buffer((3, 2), "float32")):
+    for i in range(3):
+        B[i, 0] = A[i, 2]
+
+
+@T.prim_func(s_tir=True)
+def bad_store(A: T.Buffer((2, 3), "float32"), B: T.Buffer((3, 2), "float32")):
+    B[0, 3] = A[1, 2]
+
+
+@T.prim_func(s_tir=True)
+def bad_store_loop(A: T.Buffer((2, 3), "float32"), B: T.Buffer((3, 2), "float32")):
+    for i in range(3):
+        B[0, i] = A[1, i]
+
+
+@T.prim_func(s_tir=True)
+def unknown_bounds(A: T.Buffer((2, 3), "float32"), B: T.Buffer((3, 2), "float32"), N: T.int32):
+    for i in range(3):
+        B[0, N] = A[1, i]
+
+
+def test_oob_load():
+    with pytest.raises(tvm.s_tir.ScheduleError) as err:
+        tvm.s_tir.analysis.OOBChecker()(tvm.IRModule.from_expr(bad_load))
+    assert "buffer A" in err.value.args[0]
+
+    with pytest.raises(tvm.s_tir.ScheduleError) as err:
+        tvm.s_tir.analysis.OOBChecker()(tvm.IRModule.from_expr(bad_load_loop))
+    assert "buffer A" in err.value.args[0]
+
+
+def test_oob_store():
+    with pytest.raises(tvm.s_tir.ScheduleError) as err:
+        tvm.s_tir.analysis.OOBChecker()(tvm.IRModule.from_expr(bad_store))
+    assert "buffer B" in err.value.args[0]
+
+    with pytest.raises(tvm.s_tir.ScheduleError) as err:
+        tvm.s_tir.analysis.OOBChecker()(tvm.IRModule.from_expr(bad_store_loop))
+    assert "buffer B" in err.value.args[0]
+
+
+def test_unknown_bounds():
+    # This should not return an error as we can't probe that N goes out of bounds
+    tvm.s_tir.analysis.OOBChecker()(tvm.IRModule.from_expr(unknown_bounds))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/analysis/test_sblock_access_region.py

@@ -0,0 +1,415 @@
+# 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 pytest
+
+import tvm
+import tvm.testing
+from tvm import s_tir
+from tvm.ir import Range
+from tvm.script import tirx as T
+
+
+@T.prim_func(s_tir=True)
+def func() -> None:
+    A = T.sblock_alloc_buffer((128, 128), "float32")
+    B = T.sblock_alloc_buffer((128, 128), "float32")
+    C = T.sblock_alloc_buffer((128, 128), "float32")
+    D = T.sblock_alloc_buffer((128, 128), "float32")
+    with T.sblock():
+        # Need add read/write region manually to avoid triggering block access region detector
+        T.reads([B[0, 0], C[0:16, 0:16], A[4:12, 4:12]])
+        T.writes([A[0:12, 0:12]])
+        for i, j in T.grid(8, 8):
+            A[i, j] = B[0, 0] + C[0, 0]
+        for i, j in T.grid(2, 2):
+            with T.sblock():
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads([A[vi * 4 + 4 : vi * 4 + 8, vj * 4 + 4 : vj * 4 + 8], C[12:16, 12:16]])
+                T.writes([A[vi * 4 + 4 : vi * 4 + 8, vj * 4 + 4 : vj * 4 + 8]])
+                for i, j in T.grid(4, 4):
+                    A[vi * 4 + 4 + i, vj * 4 + 4 + j] += C[i + 12, j + 12]
+        T.evaluate(D.data)
+
+
+@T.prim_func(s_tir=True)
+def match_buffer_func() -> None:
+    with T.sblock("root"):
+        A = T.sblock_alloc_buffer((128, 128), "float32")
+        B = T.sblock_alloc_buffer((128, 128), "float32")
+        T.reads([])
+        T.writes([])
+        # Need add read/write region manually to avoid triggering block access region detector
+        for i, j in T.grid(8, 8):
+            with T.sblock("block"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                T.reads(B[vi * 16 + 2 : vi * 16 + 12, vj * 16 + 2 : vj * 16 + 16])
+                T.writes(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+                AA = T.match_buffer(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16], (16, 16))
+                B0 = T.match_buffer(B[vi * 16 + 2 : vi * 16 + 6, vj * 16 + 2 : vj * 16 + 6], (4, 4))
+                B1 = T.match_buffer(
+                    B[vi * 16 + 8 : vi * 16 + 12, vj * 16 + 8 : vj * 16 + 16], (4, 8)
+                )
+                for ii, jj in T.grid(16, 16):
+                    with T.sblock("AAA"):
+                        vii, vjj = T.axis.remap("SS", [ii, jj])
+                        T.reads([])
+                        T.writes(AA[vii, vjj])
+                        AAA = T.match_buffer(AA[vii, vjj], ())
+                        AAA[()] = 1.0
+                T.evaluate(B0.data)
+                T.evaluate(B1.data)
+
+
+@T.prim_func(s_tir=True)
+def opaque_block_func() -> None:
+    with T.sblock("root"):
+        A = T.sblock_alloc_buffer((16, 16), "float32")
+        B = T.sblock_alloc_buffer((16, 16), "float32")
+        T.reads([])
+        T.writes([])
+        # Need add read/write region manually to avoid triggering block access region detector
+        for i in range(0, 16):
+            with T.sblock():
+                T.reads(A[i, 0:16])
+                T.writes([B[i, 0:16]])
+                for j in range(0, 16):
+                    with T.sblock():
+                        T.reads(A[i, j])
+                        T.writes(B[i, j])
+                        B[i, j] = A[i, j] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def opaque_access_func() -> None:
+    A = T.sblock_alloc_buffer([1024])
+    B = T.sblock_alloc_buffer([1024])
+    for i in T.serial(0, 8):
+        with T.sblock():
+            v = T.axis.S(8, i)
+            T.reads([A[v * 128 : v * 128 + 128]])
+            T.writes([B[v * 128 : v * 128 + 128]])
+            T.evaluate(
+                T.call_extern("test", B.data, v * 128, 128, A.data, v * 128, 128, dtype="float32")
+            )
+
+
+@T.prim_func(s_tir=True)
+def opaque_access_with_tvm_access_ptr_func() -> None:
+    A = T.sblock_alloc_buffer([1024])
+    B = T.sblock_alloc_buffer([1024])
+    C = T.sblock_alloc_buffer([1024])
+    with T.sblock("opaque"):
+        T.reads(A[0:1024], C[0:1024])
+        T.writes(B[0:1024], C[0:1024])
+        T.evaluate(A.access_ptr("r"))
+        T.evaluate(B.access_ptr("w"))
+        T.evaluate(C.access_ptr("rw"))
+
+
+@T.prim_func(s_tir=True)
+def access_in_if_then_else_func() -> None:
+    A = T.sblock_alloc_buffer([8])
+    B = T.sblock_alloc_buffer([8])
+    with T.sblock():
+        T.reads([A[0:5]])
+        T.writes([B[0:8]])
+        for i in T.serial(0, 8):
+            B[i] = T.if_then_else(i < 5, A[i], 0.0, dtype="float32")
+
+
+@T.prim_func(s_tir=True)
+def access_in_branch_func() -> None:
+    A = T.sblock_alloc_buffer([8])
+    B = T.sblock_alloc_buffer([8])
+    with T.sblock():
+        T.reads([A[0:7]])
+        T.writes([B[0:8]])
+        for i in T.serial(0, 8):
+            if i < 5:
+                B[i] = A[i] + 1.0
+            else:
+                B[i] = A[i - 1]
+
+
+@T.prim_func(s_tir=True)
+def gemm() -> None:
+    A = T.sblock_alloc_buffer([16, 16], "float32")
+    B = T.sblock_alloc_buffer([16, 16], "float32")
+    C = T.sblock_alloc_buffer([16, 16], "float32")
+    for i, j, k, ii, jj in T.grid(4, 4, 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)
+            T.reads(A[vi, vk], B[vj, vk])
+            T.writes(C[vi, vj])
+            with T.init():
+                C[vi, vj] = 0
+            C[vi, vj] += A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def decomposed_gemm() -> None:
+    A = T.sblock_alloc_buffer([16, 16], "float32")
+    B = T.sblock_alloc_buffer([16, 16], "float32")
+    C = 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)
+                T.reads([])
+                T.writes(C[vi, vj])
+                C[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)
+                T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
+                T.writes(C[vi, vj])
+                C[vi, vj] += A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def access_of_padding_pattern() -> None:
+    X = T.sblock_alloc_buffer([28, 28])
+    X_pad = T.sblock_alloc_buffer([32, 32])
+    Y = T.sblock_alloc_buffer([28, 28])
+    for i, j in T.grid(32, 32):
+        with T.sblock("padding"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([X[vi - 2, vj - 2]])
+            T.writes([X_pad[vi, vj]])
+            X_pad[vi, vj] = T.if_then_else(
+                2 <= vi and vi < 30 and 2 <= vj and vj < 30, X[vi - 2, vj - 2], 0.0, dtype="float32"
+            )
+        with T.sblock("padding_reverse"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads([X_pad[vi, vj]])
+            T.writes([Y[vi - 2, vj - 2]])
+            if 2 <= vi and vi < 30 and 2 <= vj and vj < 30:
+                Y[vi - 2, vj - 2] = X_pad[vi, vj]
+
+
+def test_block_access_region_detector():
+    block = func.body.block.body.block
+    alloc_buffers = func.body.block.alloc_buffers
+    buffer_var_map = {buf.data: buf for buf in alloc_buffers}
+    ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+
+    tvm.ir.assert_structural_equal(block.reads, ret[0])
+    tvm.ir.assert_structural_equal(block.writes, ret[1])
+    D = alloc_buffers[-1]
+    tvm.ir.assert_structural_equal(
+        [tvm.tirx.BufferRegion(D, [Range(0, 128), Range(0, 128)])], ret[2]
+    )
+
+
+def test_opaque_block():
+    alloc_buffers = opaque_block_func.body.block.alloc_buffers
+    buffer_var_map = {buf.data: buf for buf in alloc_buffers}
+
+    block0 = opaque_block_func.body.block.body.body.block
+    ret = s_tir.analysis.get_sblock_access_region(block0, buffer_var_map)
+    tvm.ir.assert_structural_equal(block0.reads, ret[0])
+    tvm.ir.assert_structural_equal(block0.writes, ret[1])
+
+    block1 = block0.body.body.block
+    ret = s_tir.analysis.get_sblock_access_region(block1, buffer_var_map)
+    tvm.ir.assert_structural_equal(block1.reads, ret[0])
+    tvm.ir.assert_structural_equal(block1.writes, ret[1])
+
+
+def test_opaque_access():
+    block = opaque_access_func.body.block.body.body.block
+    alloc_buffers = opaque_access_func.body.block.alloc_buffers
+    buffer_var_map = {buf.data: buf for buf in alloc_buffers}
+
+    ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map)
+    ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+    with pytest.raises(ValueError):
+        tvm.ir.assert_structural_equal(ret0[0], ret1[0])
+    with pytest.raises(ValueError):
+        tvm.ir.assert_structural_equal(ret0[1], ret1[1])
+
+
+def test_opaque_access_with_tvm_access_ptr():
+    block = opaque_access_with_tvm_access_ptr_func.body.block.body.block
+    alloc_buffers = opaque_access_with_tvm_access_ptr_func.body.block.alloc_buffers
+    buffer_var_map = {buf.data: buf for buf in alloc_buffers}
+
+    ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map)
+    ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+    tvm.ir.assert_structural_equal(block.reads, ret0[0])
+    tvm.ir.assert_structural_equal(block.writes, ret0[1])
+    with pytest.raises(ValueError):
+        tvm.ir.assert_structural_equal(ret0[0], ret1[0])
+    with pytest.raises(ValueError):
+        tvm.ir.assert_structural_equal(ret0[1], ret1[1])
+
+
+def test_match_buffer():
+    root_block = match_buffer_func.body.block
+    block = root_block.body.body.body.block
+    block_inner = block.body[0].body.body.block
+    alloc_buffers = match_buffer_func.body.block.alloc_buffers
+    buffer_var_map = {buf.data: buf for buf in alloc_buffers}
+
+    # Check block
+    ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+    tvm.ir.assert_structural_equal(block.writes, ret[1])
+    # B is opaque access
+    tvm.ir.assert_structural_equal(block.reads, ret[2])
+
+    # Check inner block AAA without updating buffer_var_map
+    ret = s_tir.analysis.get_sblock_access_region(block_inner, buffer_var_map)
+    # Since AA is not in the buffer_var_map, region of AA will not be collected.
+    tvm.ir.assert_structural_equal([], ret[1])
+
+    # Check inner block AAA
+    for match_buffer in block.match_buffers:
+        target_buffer = match_buffer.buffer
+        buffer_var_map[target_buffer.data] = target_buffer
+
+    ret = s_tir.analysis.get_sblock_access_region(block_inner, buffer_var_map)
+    tvm.ir.assert_structural_equal(block_inner.reads, ret[0])
+    tvm.ir.assert_structural_equal(block_inner.writes, ret[1])
+
+
+def test_access_in_if_then_else_func():
+    block = access_in_if_then_else_func.body.block.body.block
+    alloc_buffers = access_in_if_then_else_func.body.block.alloc_buffers
+    buffer_var_map = {buf.data: buf for buf in alloc_buffers}
+    ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map)
+    ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+    tvm.ir.assert_structural_equal(ret0[0], ret1[0])
+    tvm.ir.assert_structural_equal(ret0[1], ret1[1])
+
+
+def test_access_in_branch_func():
+    block = access_in_branch_func.body.block.body.block
+    alloc_buffers = access_in_branch_func.body.block.alloc_buffers
+    buffer_var_map = {buf.data: buf for buf in alloc_buffers}
+    ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map)
+    ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+    tvm.ir.assert_structural_equal(ret0[0], ret1[0])
+    tvm.ir.assert_structural_equal(ret0[1], ret1[1])
+
+
+def test_access_of_padding_pattern():
+    s = tvm.s_tir.schedule.Schedule(access_of_padding_pattern)
+    alloc_buffers = s.get_sref(s.get_sblock("root")).stmt.alloc_buffers
+    buffer_var_map = {buf.data: buf for buf in alloc_buffers}
+
+    def do_compare_buffer_region(region, expect):
+        assert region.buffer == expect.buffer
+        analyzer = tvm.arith.Analyzer()
+        for observed_range, expected_range in zip(region.region, expect.region):
+            analyzer.can_prove_equal(observed_range.min, expected_range.min)
+            analyzer.can_prove_equal(observed_range.extent, expected_range.extent)
+
+    def do_check_block(block_name):
+        block = s.get_sref(s.get_sblock(block_name)).stmt
+        expect_reads = block.reads
+        expect_writes = block.writes
+        ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+        for i, read in enumerate(ret[0]):
+            do_compare_buffer_region(read, expect_reads[i])
+        for i, write in enumerate(ret[1]):
+            do_compare_buffer_region(write, expect_writes[i])
+
+    do_check_block("padding")
+    do_check_block("padding_reverse")
+
+
+def test_access_of_reduction():
+    block = gemm.body.block.body.body.body.body.body.body.block
+    alloc_buffers = gemm.body.block.alloc_buffers
+    buffer_var_map = {buf.data: buf for buf in alloc_buffers}
+    ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+    tvm.ir.assert_structural_equal(block.reads, ret[0])
+    tvm.ir.assert_structural_equal(block.writes, ret[1])
+
+
+def test_access_of_decompose_reduction():
+    init = decomposed_gemm.body.block.body.body.body[0].body.body.block
+    update = decomposed_gemm.body.block.body.body.body[1].body.body.body.block
+    alloc_buffers = decomposed_gemm.body.block.alloc_buffers
+    buffer_var_map = {buf.data: buf for buf in alloc_buffers}
+    for block in [init, update]:
+        ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+        tvm.ir.assert_structural_equal(block.reads, ret[0])
+        tvm.ir.assert_structural_equal(block.writes, ret[1])
+
+
+def test_buffer_access_with_let_binding():
+    @T.prim_func(s_tir=True)
+    def func(
+        storage: T.Buffer((16, 16, 16), "float32"),
+        seq_slot_ids: T.Buffer((16,), "int32"),
+        history_slot_ids: T.Buffer((16,), "int32"),
+        output: T.Buffer((16, 16), "float32"),
+    ):
+        for i, s in T.grid(16, 16):
+            with T.sblock("copy"):
+                vi, vs = T.axis.remap("SS", [i, s])
+                T.reads(
+                    seq_slot_ids[vi],
+                    history_slot_ids[vi],
+                    storage[seq_slot_ids[vi], history_slot_ids[vi], vs],
+                )
+                T.writes(output[vi, vs])
+                seq_id: T.let[T.int32] = seq_slot_ids[vi]
+                history_id: T.let[T.int32] = history_slot_ids[vi]
+                output[vi, vs] = storage[seq_id, history_id, vs]
+
+    block = func.body.block.body.body.body.block
+    buffer_var_map = {buf.data: buf for buf in func.buffer_map.values()}
+    ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+    tvm.ir.assert_structural_equal(block.reads, ret[0])
+    tvm.ir.assert_structural_equal(block.writes, ret[1])
+
+
+def test_buffer_access_with_nested_let_binding():
+    @T.prim_func(s_tir=True)
+    def func(
+        A: T.Buffer((16, 16), "float32"),
+        B: T.Buffer((16, 16), "float32"),
+        C: T.Buffer((16, 16), "float32"),
+    ):
+        for i, s in T.grid(16, 16):
+            with T.sblock("copy"):
+                vi, vs = T.axis.remap("SS", [i, s])
+                T.reads(A[vi, vs], B[vi, vs])
+                T.writes(C[vi, vs])
+                vi1: T.let[T.int32] = vi
+                vi2: T.let[T.int32] = vi1
+                vs1: T.let[T.int32] = vs
+                vs2: T.let[T.int32] = vs1
+                vs3: T.let[T.int32] = vs2
+                C[vi, vs1] = A[vi1, vs2] + B[vi2, vs3]
+
+    block = func.body.block.body.body.body.block
+    buffer_var_map = {buf.data: buf for buf in func.buffer_map.values()}
+    ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map)
+    tvm.ir.assert_structural_equal(block.reads, ret[0])
+    tvm.ir.assert_structural_equal(block.writes, ret[1])
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/analysis/test_sblock_buffer_access_lca.py

@@ -0,0 +1,189 @@
+# 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 tvm
+from tvm import s_tir
+from tvm.script import tirx as T
+
+
+@T.prim_func(s_tir=True)
+def buffer_load_store_func(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128), "float32")
+    B = T.match_buffer(b, (128, 128), "float32")
+    C = T.sblock_alloc_buffer((128, 128), "float32")
+    D = T.sblock_alloc_buffer((128, 128), "float32")
+    for ii, jj in T.grid(128, 128):
+        with T.sblock():
+            i, j = T.axis.remap("SS", [ii, jj])
+            A[i, j] = T.float32(0)
+    for i0, j0, k0 in T.grid(32, 32, 32):
+        with T.sblock():
+            i, j, k = T.axis.remap("SSR", [i0, j0, k0])
+            with T.init():
+                for ii, jj in T.grid(4, 4):
+                    B[i * 4 + ii, j * 4 + jj] = A[i * 4 + ii, j * 4 + jj]
+            for ii, jj in T.grid(4, 4):
+                for kk in range(0, 4):
+                    B[i * 4 + ii, j * 4 + jj] += C[i * 4 + ii, k * 4 + kk]
+                for kk in range(0, 4):
+                    B[i * 4 + ii, j * 4 + jj] += (
+                        D[j * 4 + jj, k * 4 + kk] * C[i * 4 + ii, k * 4 + kk]
+                    )
+
+
+@T.prim_func(s_tir=True)
+def buffer_opaque_access(b: T.handle, c: T.handle) -> None:
+    B = T.match_buffer(b, [16, 16], "float32")
+    C = T.match_buffer(c, [16, 16], "float32")
+
+    with T.sblock():
+        T.reads([])
+        T.writes(B[0:16, 0:16])
+        A = T.decl_buffer([256], "float32")
+        for i, j in T.grid(16, 16):
+            A[i * 16 + j] = 1
+        for i in range(0, 16):
+            for j in range(0, 16):
+                T.evaluate(A[i * 16 + j])
+            for j in range(0, 16):
+                T.evaluate(T.tvm_fill_fragment(B.data, 16, 16, 16, 0, T.float32(0), dtype="handle"))
+
+    for i, j in T.grid(16, 16):
+        with T.sblock():
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def lca_is_func_root(a: T.handle) -> None:
+    A = T.match_buffer(a, [0, 0], "float32")
+    A[0, 0] = 1.0
+
+
+@T.prim_func(s_tir=True)
+def match_buffer_func(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, j in T.grid(8, 8):
+        with T.sblock("block"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            T.reads(B[vi * 16 + 2 : vi * 16 + 12, vj * 16 + 2 : vj * 16 + 16])
+            T.writes(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
+            B0 = T.match_buffer(B[vi * 16 + 2 : vi * 16 + 6, vj * 16 + 2 : vj * 16 + 6], (4, 4))
+            B1 = T.match_buffer(B[vi * 16 + 8 : vi * 16 + 12, vj * 16 + 8 : vj * 16 + 16], (4, 8))
+            for ii, jj in T.grid(16, 16):
+                with T.sblock("AAA"):
+                    vii, vjj = T.axis.remap("SS", [ii, jj])
+                    AA = T.match_buffer(A[vii, vjj], ())
+                    AA[()] = 1.0
+            T.evaluate(B0.data)
+            T.evaluate(B1.data)
+
+
+@T.prim_func(s_tir=True)
+def global_buffer_with_blockidx(
+    a: T.Buffer((1, 32), "int32"), b: T.Buffer((1, 32), "int32")
+) -> None:
+    for i0 in T.thread_binding(0, 1, thread="blockIdx.x"):
+        for i1 in T.thread_binding(0, 32, thread="threadIdx.x"):
+            with T.sblock("copy"):
+                i, j = T.axis.remap("SS", [i0, i1])
+                T.reads(a[i, j])
+                T.writes(b[i, j])
+                b[i, j] = a[i, j]
+
+
+def test_buffer_load_store():
+    func = buffer_load_store_func
+    A, B = [func.buffer_map[x] for x in func.params]
+    C, D = func.body.block.alloc_buffers
+    lca = s_tir.analysis.detect_buffer_access_lca(func)
+
+    # LCA of Buffer A is root
+    root_block = func.body.block
+    assert lca[A] == func.body.block
+
+    # LCA of Buffer B is the loop dominate all reduction loop
+    reduce_dom_loop = root_block.body[1].body
+    reduce_block = reduce_dom_loop.body.body.block
+    assert lca[B] == reduce_dom_loop
+
+    # LCA of Buffer C is the second loop kk
+    loop_jj = reduce_block.body.body
+    assert lca[C] == loop_jj
+
+    # LCA of Buffer D is loop jj
+    loop_kk = loop_jj.body[1]
+    assert lca[D] == loop_kk
+
+
+def test_opaque_access():
+    func = buffer_opaque_access
+    B, C = [func.buffer_map[x] for x in func.params]
+    lca = s_tir.analysis.detect_buffer_access_lca(func)
+
+    # Cannot detect buffer A since it is define by low-level Allocate
+
+    # LCA of Buffer B is root
+    root_block = func.body.block
+    assert lca[B] == func.body.block
+
+    # LCA of Buffer C is the correspond block
+    assert lca[C] == root_block.body[1].body.body.block
+
+
+def test_lca_func_root():
+    func = lca_is_func_root
+    (A,) = [func.buffer_map[x] for x in func.params]
+    lca = s_tir.analysis.detect_buffer_access_lca(func)
+    assert lca[A] is None
+
+
+def test_match_buffer():
+    func = match_buffer_func
+    A, B = [func.buffer_map[x] for x in func.params]
+    lca = s_tir.analysis.detect_buffer_access_lca(func)
+
+    root_block = func.body.block
+    block = root_block.body.body.body.block
+    block_inner = block.body[0].body.body.block
+
+    # LCA of Buffer C is the inner block
+    assert lca[A] == block_inner
+
+    # LCA of Buffer C is the main block
+    assert lca[B] == block
+
+
+def test_global_buffer_with_blockidx():
+    func = global_buffer_with_blockidx
+    A, B = [func.buffer_map[x] for x in func.params]
+    lca = s_tir.analysis.detect_buffer_access_lca(func)
+
+    root_block = func.body.block
+    blockidx_loop = root_block.body
+    # LCA of both A and B should be the loop bound to `blockIdx`
+    assert lca[A] == blockidx_loop
+    assert lca[B] == blockidx_loop
+
+
+if __name__ == "__main__":
+    test_buffer_load_store()
+    test_opaque_access()
+    test_lca_func_root()
+    test_match_buffer()
+    test_global_buffer_with_blockidx()