chore: import upstream snapshot with attribution

tests/python/s_tir/transform/test_s_tir_transform_annotate_irregular_loop.py

@@ -0,0 +1,271 @@
+# 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
+
+"""Tests for AnnotateIrregularLoop"""
+
+import tvm
+import tvm.testing
+from tvm import s_tir, tirx
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_handle_irrgular_unit_loop():
+    """Dedicated testcase to check the unitloop with loop jump not simplified"""
+
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((10,), "int32")):
+        for i in T.serial(1):
+            if A[i] > 5:
+                break
+            A[i] = A[i] + 1
+        for j in T.serial(1):
+            if A[j] > 5:
+                continue
+            A[j] = A[j] + 1
+        for k in T.serial(1):
+            A[k] = A[k] + 1
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((10,), "int32")):
+        for i in T.serial(1, annotations={"irregular_loop_mark": 1}):
+            if A[i] > 5:
+                break
+            A[i] = A[i] + 1
+        for j in T.serial(1, annotations={"irregular_loop_mark": 1}):
+            if A[j] > 5:
+                continue
+            A[j] = A[j] + 1
+        A[0] = A[0] + 1
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = tvm.s_tir.transform.AnnotateIrregularLoop()(mod)
+    mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+    tvm.ir.assert_structural_equal(mod["before"].with_attr("global_symbol", "expected"), expected)
+
+
+def test_annotate_loop_with_break():
+    """Test that loops containing break statements are annotated as irregular."""
+    transform = s_tir.transform.AnnotateIrregularLoop()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10,), "int32")):
+            for i in T.serial(10):
+                if A[i] > 5:
+                    break
+                A[i] = A[i] + 1
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10,), "int32")):
+            for i in T.serial(10, annotations={"irregular_loop_mark": 1}):
+                if A[i] > 5:
+                    break
+                A[i] = A[i] + 1
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_annotate_loop_with_continue():
+    """Test that loops containing continue statements are annotated as irregular."""
+    transform = s_tir.transform.AnnotateIrregularLoop()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10,), "int32")):
+            for i in T.serial(10):
+                if A[i] < 0:
+                    continue
+                A[i] = A[i] * 2
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10,), "int32")):
+            for i in T.serial(10, annotations={"irregular_loop_mark": 1}):
+                if A[i] < 0:
+                    continue
+                A[i] = A[i] * 2
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_nested_irregular_both_loops():
+    """Test nested loops where both loops have break/continue."""
+    transform = s_tir.transform.AnnotateIrregularLoop()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10, 10), "int32")):
+            for i in T.serial(10):
+                if i > 7:
+                    break
+                for j in T.serial(10):
+                    if A[i, j] < 0:
+                        continue
+                    A[i, j] = A[i, j] + 1
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10, 10), "int32")):
+            for i in T.serial(10, annotations={"irregular_loop_mark": 1}):
+                if i > 7:
+                    break
+                for j in T.serial(10, annotations={"irregular_loop_mark": 1}):
+                    if A[i, j] < 0:
+                        continue
+                    A[i, j] = A[i, j] + 1
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_while_loop_with_break():
+    """Test that while loops with break/continue are not annotated (while loops don't have annotations)."""
+    transform = s_tir.transform.AnnotateIrregularLoop()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10,), "int32")):
+            i = T.int32(0)
+            while i < 10:
+                if A[i] > 5:
+                    break
+                A[i] = A[i] + 1
+                i = i + 1
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10,), "int32")):
+            i = T.int32(0)
+            while i < 10:
+                if A[i] > 5:
+                    break
+                A[i] = A[i] + 1
+                i = i + 1
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_break_in_nested_conditional():
+    """Test break statement deeply nested in conditional blocks."""
+    transform = s_tir.transform.AnnotateIrregularLoop()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10,), "int32"), flag1: T.int32, flag2: T.int32):
+            for i in T.serial(10):
+                if flag1 > 0:
+                    if flag2 > 0:
+                        if A[i] > 5:
+                            break
+                A[i] = A[i] + 1
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10,), "int32"), flag1: T.int32, flag2: T.int32):
+            for i in T.serial(10, annotations={"irregular_loop_mark": 1}):
+                if flag1 > 0:
+                    if flag2 > 0:
+                        if A[i] > 5:
+                            break
+                A[i] = A[i] + 1
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_while_loop_with_break_standalone():
+    """Test that while loops with break/continue are not annotated (while loops don't have annotations)."""
+    transform = s_tir.transform.AnnotateIrregularLoop()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10,), "int32")):
+            i = T.int32(0)
+            while i < 10:
+                if A[i] > 5:
+                    break
+                A[i] = A[i] + 1
+                i = i + 1
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((10,), "int32")):
+            i = T.int32(0)
+            while i < 10:
+                if A[i] > 5:
+                    break
+                A[i] = A[i] + 1
+                i = i + 1
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_nested_irregular_loop_standalone():
+    """Test deeply nested loops with irregular control flow only in innermost loop."""
+    transform = s_tir.transform.AnnotateIrregularLoop()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((5, 5, 5), "int32")):
+            for i in T.serial(5):
+                for j in T.serial(5):
+                    for k in T.serial(5):
+                        if A[i, j, k] > 10:
+                            break
+                        if A[i, j, k] < 0:
+                            continue
+                        A[i, j, k] = A[i, j, k] + 1
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((5, 5, 5), "int32")):
+            for i in T.serial(5):
+                for j in T.serial(5):
+                    for k in T.serial(5, annotations={"irregular_loop_mark": 1}):
+                        if A[i, j, k] > 10:
+                            break
+                        if A[i, j, k] < 0:
+                            continue
+                        A[i, j, k] = A[i, j, k] + 1
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_canonicalize_loop.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
+from tvm import s_tir, tirx
+from tvm.script import tirx as T
+
+
+def test_canonicalize_loop():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+        T.func_attr({"global_symbol": "main"})
+        for i in range(1, 128, 5):
+            B[i] = A[i] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+        T.func_attr({"global_symbol": "main"})
+        for i in T.serial(0, 26):
+            B[i * 5 + 1] = A[i * 5 + 1] + 1.0
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = s_tir.transform.CanonicalizeLoop()(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_canonicalize_nested_loop():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")):
+        T.func_attr({"global_symbol": "main"})
+        for i in range(1, 128, 5):
+            for j in range(2, 128, 3):
+                B[i, j] = A[i, j] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")):
+        T.func_attr({"global_symbol": "main"})
+        for i in T.serial(0, 26):
+            for j in T.serial(0, 42):
+                B[i * 5 + 1, j * 3 + 2] = A[i * 5 + 1, j * 3 + 2] + 1.0
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = s_tir.transform.CanonicalizeLoop()(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_canonicalize_negative_step():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+        T.func_attr({"global_symbol": "main"})
+        for i in T.serial(0, 127, step=-3):
+            B[i] = A[i] + 1.0
+
+    mod = tvm.IRModule.from_expr(before)
+    with pytest.raises(tvm.error.InternalError):
+        mod = s_tir.transform.CanonicalizeLoop()(mod)
+
+
+def test_canonicalize_dynamic_step():
+    """Currently we report error for dynamic step since we could not prove it is positive"""
+
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32"), step: T.int32):
+        T.func_attr({"global_symbol": "main"})
+        for i in T.serial(0, 128, step=step):
+            B[i] = A[i] + 1.0
+
+    mod = tvm.IRModule.from_expr(before)
+    with pytest.raises(tvm.error.InternalError):
+        mod = s_tir.transform.CanonicalizeLoop()(mod)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py

@@ -0,0 +1,97 @@
+# 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 s_tir, te, tirx
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def _check(original, transformed):
+    func = original
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.ConvertBlocksToOpaque()(mod)
+    mod = tvm.tirx.transform.StmtSimplify()(mod)
+    tvm.ir.assert_structural_equal(mod["main"], transformed.with_attr("global_symbol", "main"))
+
+
+@T.prim_func(s_tir=True)
+def elementwise_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), "float32")
+    C = T.match_buffer(c, (16, 16), "float32")
+    for i in range(0, 16):
+        with T.sblock():
+            T.reads(A[i, 0:16])
+            T.writes(C[i, 0:16])
+            B = T.sblock_alloc_buffer((16, 16), "float32")
+            for j in range(0, 16):
+                with T.sblock():
+                    vi = T.axis.S(16, i)
+                    vj = T.axis.S(16, j)
+                    B[vi, vj] = A[vi, vj] + 1.0
+            for j in range(0, 16):
+                with T.sblock():
+                    vi = T.axis.S(16, i)
+                    vj = T.axis.S(16, j)
+                    C[vi, vj] = B[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def substituted_elementwise_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), "float32")
+    C = T.match_buffer(c, (16, 16), "float32")
+    for i in range(0, 16):
+        with T.sblock():
+            T.reads(A[i, 0:16])
+            T.writes(C[i, 0:16])
+            B = T.sblock_alloc_buffer([16, 16], "float32")
+            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
+            for j in range(0, 16):
+                with T.sblock():
+                    T.reads([B[i, j]])
+                    T.writes([C[i, j]])
+                    C[i, j] = B[i, j] * 2.0
+
+
+def test_elementwise():
+    _check(elementwise_func, substituted_elementwise_func)
+
+
+def test_error_if_predicate_uses_block_variables():
+    @I.ir_module(check_well_formed=False, s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(8, "int32")):
+            for i in T.serial(8):
+                with T.sblock():
+                    vi = T.axis.remap("S", [i])
+                    T.where(vi < 6)
+                    T.evaluate(0)
+
+    with pytest.raises(RuntimeError):
+        tvm.s_tir.transform.ConvertBlocksToOpaque()(Before)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_decorate_device_scope.py

@@ -0,0 +1,29 @@
+# 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
+
+
+def test_decorate_device():
+    x = tvm.tirx.Var("x", "int32")
+    mod = tvm.IRModule.from_expr(tvm.tirx.PrimFunc([x], tvm.tirx.Evaluate(x)))
+
+    stmt = tvm.s_tir.transform.DecorateDeviceScope()(mod)["main"].body
+    assert stmt.attr_key == "device_scope"
+
+
+if __name__ == "__main__":
+    test_decorate_device()

tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py

@@ -0,0 +1,605 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name,,missing-function-docstring
+# ruff: noqa: E501, F841
+import tvm
+import tvm.testing
+from tvm.s_tir.transform import DefaultGPUSchedule
+from tvm.script import tirx as T
+
+
+def test_broadcast_to_symbolic():
+    # pylint: disable=no-self-argument,missing-class-docstring,line-too-long
+    # fmt: off
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def broadcast_to(
+            rxplaceholder: T.Buffer((T.int64(3), T.int64(1)), "float32"),
+            var_T_broadcast_to: T.handle,
+        ):
+            T.func_attr({"tirx.noalias": True})
+            x_0 = T.int64()
+            x_1 = T.int64()
+            T_broadcast_to = T.match_buffer(var_T_broadcast_to, (x_0, x_1))
+            # with T.sblock("root"):
+            for ax0, ax1 in T.grid(x_0, x_1):
+                with T.sblock("T_broadcast_to"):
+                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+                    T.reads(rxplaceholder[v_ax0, T.int64(0)])
+                    T.writes(T_broadcast_to[v_ax0, v_ax1])
+                    T_broadcast_to[v_ax0, v_ax1] = rxplaceholder[v_ax0, T.int64(0)]
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def broadcast_to(rxplaceholder: T.Buffer((T.int64(3), T.int64(1)), "float32"), var_T_broadcast_to: T.handle):
+            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
+            x_0, x_1 = T.int64(), T.int64()
+            T_broadcast_to = T.match_buffer(var_T_broadcast_to, (x_0, x_1))
+            for ax0_ax1_fused_1 in T.thread_binding(T.int64(256), thread="blockIdx.x"):
+                for ax0_ax1_fused_2 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
+                    for ax0_ax1_fused_0 in range((x_0 * x_1 + T.int64(262143)) // T.int64(262144)):
+                        with T.sblock("T_broadcast_to"):
+                            v_ax0 = T.axis.spatial(x_0, (ax0_ax1_fused_0 * T.int64(262144) + ax0_ax1_fused_1 * T.int64(1024) + ax0_ax1_fused_2) // x_1)
+                            v_ax1 = T.axis.spatial(x_1, (ax0_ax1_fused_0 * T.int64(262144) + ax0_ax1_fused_1 * T.int64(1024) + ax0_ax1_fused_2) % x_1)
+                            T.where((ax0_ax1_fused_0 * T.int64(256) + ax0_ax1_fused_1) * T.int64(1024) + ax0_ax1_fused_2 < x_0 * x_1)
+                            T_broadcast_to[v_ax0, v_ax1] = rxplaceholder[v_ax0, T.int64(0)]
+    # fmt: on
+    # pylint: enable=no-self-argument,missing-class-docstring,line-too-long
+    target = tvm.target.Target("nvidia/geforce-rtx-3070")
+    with target, tvm.transform.PassContext(opt_level=3):
+        After = DefaultGPUSchedule()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_matmul():
+    # pylint: disable=no-self-argument,missing-class-docstring,line-too-long
+    # fmt: off
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def matmul(
+            A: T.Buffer((32, 32), "float16"),
+            B: T.Buffer((32, 32), "float16"),
+            C: T.Buffer((32, 32), "float16"),
+        ):
+            T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i, j, k in T.grid(32, 32, 32):
+                with T.sblock("C"):
+                    v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
+                    T.reads(A[v_i, v_k], B[v_k, v_j])
+                    T.writes(C[v_i, v_j])
+                    with T.init():
+                        C[v_i, v_j] = T.float16(0)
+                    C[v_i, v_j] = C[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
+
+        @T.prim_func(s_tir=True)
+        def matmul_gpu(
+            A: T.Buffer((32, 32), "float16"),
+            B: T.Buffer((32, 32), "float16"),
+            C: T.Buffer((32, 32), "float16"),
+        ):
+            T.func_attr({"global_symbol": "main",
+                         "target": T.target({"arch": "sm_86",
+                                             "keys": ["cuda", "gpu"],
+                                             "kind": "cuda",
+                                             "max_num_threads": 1024,
+                                             "tag": "",
+                                             "thread_warp_size": 32}),
+                         "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i, j, k in T.grid(32, 32, 32):
+                with T.sblock("C"):
+                    v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
+                    T.reads(A[v_i, v_k], B[v_k, v_j])
+                    T.writes(C[v_i, v_j])
+                    with T.init():
+                        C[v_i, v_j] = T.float16(0)
+                    C[v_i, v_j] = C[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
+
+        @T.prim_func(s_tir=True)
+        def matmul_cpu(
+            A: T.Buffer((32, 32), "float16"),
+            B: T.Buffer((32, 32), "float16"),
+            C: T.Buffer((32, 32), "float16"),
+        ):
+            T.func_attr({"global_symbol": "main",
+                         "target": T.target({"keys": ["cpu"], "kind": "llvm", "tag": ""}),
+                        "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i, j, k in T.grid(32, 32, 32):
+                with T.sblock("C"):
+                    v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
+                    T.reads(A[v_i, v_k], B[v_k, v_j])
+                    T.writes(C[v_i, v_j])
+                    with T.init():
+                        C[v_i, v_j] = T.float16(0)
+                    C[v_i, v_j] = C[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def matmul(
+            A: T.Buffer((32, 32), "float16"),
+            B: T.Buffer((32, 32), "float16"),
+            C: T.Buffer((32, 32), "float16"),
+        ):
+            T.func_attr({"tirx.is_scheduled": True, "global_symbol": "main", "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i_j_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
+                for i_j_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
+                    for k in range(32):
+                        with T.sblock("C"):
+                            v_i = T.axis.spatial(
+                                32, (i_j_fused_0 * 1024 + i_j_fused_1) // 32
+                            )
+                            v_j = T.axis.spatial(
+                                32, (i_j_fused_0 * 1024 + i_j_fused_1) % 32
+                            )
+                            v_k = T.axis.reduce(32, k)
+                            T.reads(A[v_i, v_k], B[v_k, v_j])
+                            T.writes(C[v_i, v_j])
+                            with T.init():
+                                C[v_i, v_j] = T.float16(0)
+                            C[v_i, v_j] = C[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
+
+        @T.prim_func(s_tir=True)
+        def matmul_cpu(A: T.Buffer((32, 32), "float16"), B: T.Buffer((32, 32), "float16"), C: T.Buffer((32, 32), "float16")):
+            T.func_attr({"global_symbol": "main", "target": T.target({"keys": ["cpu"], "kind": "llvm", "tag": ""}), "tirx.is_scheduled": True, "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i, j, k in T.grid(32, 32, 32):
+                with T.sblock("C"):
+                    v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
+                    T.reads(A[v_i, v_k], B[v_k, v_j])
+                    T.writes(C[v_i, v_j])
+                    with T.init():
+                        C[v_i, v_j] = T.float16(0)
+                    C[v_i, v_j] = C[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
+
+        @T.prim_func(s_tir=True)
+        def matmul_gpu(A: T.Buffer((32, 32), "float16"), B: T.Buffer((32, 32), "float16"), C: T.Buffer((32, 32), "float16")):
+            T.func_attr({"global_symbol": "main", "target": T.target({"arch": "sm_86", "keys": ["cuda", "gpu"], "kind": "cuda", "max_num_threads": 1024, "tag": "", "thread_warp_size": 32}), "tirx.is_scheduled": True, "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i_j_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
+                for i_j_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
+                    for k in range(32):
+                        with T.sblock("C"):
+                            v_i = T.axis.spatial(32, (i_j_fused_0 * 1024 + i_j_fused_1) // 32)
+                            v_j = T.axis.spatial(32, (i_j_fused_0 * 1024 + i_j_fused_1) % 32)
+                            v_k = T.axis.reduce(32, k)
+                            T.reads(A[v_i, v_k], B[v_k, v_j])
+                            T.writes(C[v_i, v_j])
+                            with T.init():
+                                C[v_i, v_j] = T.float16(0)
+                            C[v_i, v_j] = C[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
+    # fmt: on
+    # pylint: enable=no-self-argument,missing-class-docstring,line-too-long
+    target = tvm.target.Target("nvidia/geforce-rtx-3070")
+    with target, tvm.transform.PassContext(opt_level=3):
+        After = DefaultGPUSchedule()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_add():
+    # pylint: disable=no-self-argument,missing-class-docstring,line-too-long
+    # fmt: off
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def add(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_add: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+                with T.sblock("T_add"):
+                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                    T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+                    T.writes(T_add[ax0, ax1, ax2, ax3])
+                    T_add[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] + rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def add(
+            rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"),
+            rxplaceholder_1: T.Buffer(
+                (T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"
+            ),
+            T_add: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32"),
+        ):
+            T.func_attr({"tirx.is_scheduled": True,  "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i0_i1_i2_i3_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
+                for i0_i1_i2_i3_fused_1 in T.thread_binding(
+                    T.int64(72), thread="threadIdx.x"
+                ):
+                    with T.sblock("T_add"):
+                        ax0 = T.axis.spatial(
+                            T.int64(4),
+                            (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1)
+                            // T.int64(18),
+                        )
+                        ax1 = T.axis.spatial(
+                            T.int64(3),
+                            (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1)
+                            % T.int64(18)
+                            // T.int64(6),
+                        )
+                        ax2 = T.axis.spatial(
+                            T.int64(2),
+                            (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1)
+                            % T.int64(6)
+                            // T.int64(3),
+                        )
+                        ax3 = T.axis.spatial(
+                            T.int64(3),
+                            (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1)
+                            % T.int64(3),
+                        )
+                        T.reads(
+                            rxplaceholder[T.int64(0), ax2, ax3],
+                            rxplaceholder_1[ax0, ax1, ax2, T.int64(0)],
+                        )
+                        T.writes(T_add[ax0, ax1, ax2, ax3])
+                        T_add[ax0, ax1, ax2, ax3] = (
+                            rxplaceholder[T.int64(0), ax2, ax3]
+                            + rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+                        )
+
+    # fmt: on
+    # pylint: enable=no-self-argument,missing-class-docstring,line-too-long
+    target = tvm.target.Target("nvidia/geforce-rtx-3070")
+    with target, tvm.transform.PassContext(opt_level=3):
+        After = DefaultGPUSchedule()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_full():
+    # pylint: disable=no-self-argument,missing-class-docstring,line-too-long
+    # fmt: off
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def full(rxplaceholder: T.Buffer((), "int32"), T_full: T.Buffer((T.int64(2), T.int64(3)), "int32")):
+            T.func_attr({"tirx.noalias": True})
+            for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+                with T.sblock("T_full"):
+                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
+                    T.reads(rxplaceholder[()])
+                    T.writes(T_full[ax0, ax1])
+                    T_full[ax0, ax1] = rxplaceholder[()]
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def full(
+            rxplaceholder: T.Buffer((), "int32"),
+            T_full: T.Buffer((T.int64(2), T.int64(3)), "int32"),
+        ):
+            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i0_i1_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
+                for i0_i1_fused_1 in T.thread_binding(T.int64(6), thread="threadIdx.x"):
+                    with T.sblock("T_full"):
+                        ax0 = T.axis.spatial(
+                            T.int64(2),
+                            (i0_i1_fused_0 * T.int64(6) + i0_i1_fused_1) // T.int64(3),
+                        )
+                        ax1 = T.axis.spatial(
+                            T.int64(3),
+                            (i0_i1_fused_0 * T.int64(6) + i0_i1_fused_1) % T.int64(3),
+                        )
+                        T.reads(rxplaceholder[()])
+                        T.writes(T_full[ax0, ax1])
+                        T_full[ax0, ax1] = rxplaceholder[()]
+
+    # fmt: on
+    # pylint: enable=no-self-argument,missing-class-docstring,line-too-long
+    target = tvm.target.Target("nvidia/geforce-rtx-3070")
+    with target, tvm.transform.PassContext(opt_level=3):
+        After = DefaultGPUSchedule()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_scheduled():
+    # pylint: disable=no-self-argument,missing-class-docstring,line-too-long
+    # fmt: off
+
+    @tvm.script.ir_module
+    class Scheduled:
+        @T.prim_func(s_tir=True)
+        def full(
+            rxplaceholder: T.Buffer((), "int32"),
+            T_full: T.Buffer((T.int64(2), T.int64(3)), "int32"),
+        ):
+            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i0_i1_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
+                for i0_i1_fused_1 in T.thread_binding(T.int64(6), thread="threadIdx.x"):
+                    with T.sblock("T_full"):
+                        ax0 = T.axis.spatial(
+                            T.int64(2),
+                            (i0_i1_fused_0 * T.int64(6) + i0_i1_fused_1) // T.int64(3),
+                        )
+                        ax1 = T.axis.spatial(
+                            T.int64(3),
+                            (i0_i1_fused_0 * T.int64(6) + i0_i1_fused_1) % T.int64(3),
+                        )
+                        T.reads(rxplaceholder[()])
+                        T.writes(T_full[ax0, ax1])
+                        T_full[ax0, ax1] = rxplaceholder[()]
+
+    # fmt: on
+    # pylint: enable=no-self-argument,missing-class-docstring,line-too-long
+    target = tvm.target.Target("nvidia/geforce-rtx-3070")
+    with target, tvm.transform.PassContext(opt_level=3):
+        # should do nothing
+        After = DefaultGPUSchedule()(Scheduled)
+    tvm.ir.assert_structural_equal(After, Scheduled)
+
+
+def test_multiple():
+    # pylint: disable=no-self-argument,missing-class-docstring,line-too-long
+    # fmt: off
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def add(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_add: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+                with T.sblock("T_add"):
+                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                    T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+                    T.writes(T_add[ax0, ax1, ax2, ax3])
+                    T_add[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] + rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+
+        @T.prim_func(s_tir=True)
+        def full(rxplaceholder: T.Buffer((), "int32"), T_full: T.Buffer((T.int64(2), T.int64(3)), "int32")):
+            T.func_attr({"tirx.noalias": True})
+            for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+                with T.sblock("T_full"):
+                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
+                    T.reads(rxplaceholder[()])
+                    T.writes(T_full[ax0, ax1])
+                    T_full[ax0, ax1] = rxplaceholder[()]
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def add(
+            rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"),
+            rxplaceholder_1: T.Buffer(
+                (T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"
+            ),
+            T_add: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32"),
+        ):
+            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i0_i1_i2_i3_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
+                for i0_i1_i2_i3_fused_1 in T.thread_binding(
+                    T.int64(72), thread="threadIdx.x"
+                ):
+                    with T.sblock("T_add"):
+                        ax0 = T.axis.spatial(
+                            T.int64(4),
+                            (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1)
+                            // T.int64(18),
+                        )
+                        ax1 = T.axis.spatial(
+                            T.int64(3),
+                            (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1)
+                            % T.int64(18)
+                            // T.int64(6),
+                        )
+                        ax2 = T.axis.spatial(
+                            T.int64(2),
+                            (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1)
+                            % T.int64(6)
+                            // T.int64(3),
+                        )
+                        ax3 = T.axis.spatial(
+                            T.int64(3),
+                            (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1)
+                            % T.int64(3),
+                        )
+                        T.reads(
+                            rxplaceholder[T.int64(0), ax2, ax3],
+                            rxplaceholder_1[ax0, ax1, ax2, T.int64(0)],
+                        )
+                        T.writes(T_add[ax0, ax1, ax2, ax3])
+                        T_add[ax0, ax1, ax2, ax3] = (
+                            rxplaceholder[T.int64(0), ax2, ax3]
+                            + rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+                        )
+
+        @T.prim_func(s_tir=True)
+        def full(
+            rxplaceholder: T.Buffer((), "int32"),
+            T_full: T.Buffer((T.int64(2), T.int64(3)), "int32"),
+        ):
+            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
+            # with T.sblock("root"):
+            for i0_i1_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
+                for i0_i1_fused_1 in T.thread_binding(T.int64(6), thread="threadIdx.x"):
+                    with T.sblock("T_full"):
+                        ax0 = T.axis.spatial(
+                            T.int64(2),
+                            (i0_i1_fused_0 * T.int64(6) + i0_i1_fused_1) // T.int64(3),
+                        )
+                        ax1 = T.axis.spatial(
+                            T.int64(3),
+                            (i0_i1_fused_0 * T.int64(6) + i0_i1_fused_1) % T.int64(3),
+                        )
+                        T.reads(rxplaceholder[()])
+                        T.writes(T_full[ax0, ax1])
+                        T_full[ax0, ax1] = rxplaceholder[()]
+    # fmt: on
+    # pylint: enable=no-self-argument,missing-class-docstring,line-too-long
+    target = tvm.target.Target("nvidia/geforce-rtx-3070")
+    with target, tvm.transform.PassContext(opt_level=3):
+        After = DefaultGPUSchedule()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_add_on_metal():
+    # pylint: disable=no-self-argument,missing-class-docstring,line-too-long
+    # fmt: off
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def add(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_add: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+                with T.sblock("T_add"):
+                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+                    T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+                    T.writes(T_add[ax0, ax1, ax2, ax3])
+                    T_add[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] + rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def add(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_add: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")):
+            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
+            for i0_i1_i2_i3_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
+                for i0_i1_i2_i3_fused_1 in T.thread_binding(T.int64(72), thread="threadIdx.x"):
+                    with T.sblock("T_add"):
+                        ax0 = T.axis.spatial(T.int64(4), (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1) // T.int64(18))
+                        ax1 = T.axis.spatial(T.int64(3), (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1) % T.int64(18) // T.int64(6))
+                        ax2 = T.axis.spatial(T.int64(2), (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1) % T.int64(6) // T.int64(3))
+                        ax3 = T.axis.spatial(T.int64(3), (i0_i1_i2_i3_fused_0 * T.int64(72) + i0_i1_i2_i3_fused_1) % T.int64(3))
+                        T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+                        T.writes(T_add[ax0, ax1, ax2, ax3])
+                        T_add[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] + rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+    # fmt: on
+    # pylint: enable=no-self-argument,missing-class-docstring,line-too-long
+    target = tvm.target.Target("apple/m1-gpu")
+    with target, tvm.transform.PassContext(opt_level=0):
+        mod = DefaultGPUSchedule()(Before)
+    tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_scalar_add():
+    # pylint: disable=no-self-argument,missing-class-docstring,line-too-long
+    # fmt: off
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def add(rxplaceholder: T.Buffer((), "int64"), T_add: T.Buffer((), "int64")):
+            T.func_attr({"tirx.noalias": True})
+            with T.sblock("T_add"):
+                vi = T.axis.spatial(1, T.int64(0))
+                T.reads(rxplaceholder[()])
+                T.writes(T_add[()])
+                T_add[()] = rxplaceholder[()] + T.int64(1)
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def add(rxplaceholder: T.Buffer((), "int64"), T_add: T.Buffer((), "int64")):
+            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
+            # with T.sblock("root"):
+            for u_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
+                for u_fused_1 in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("T_add"):
+                        vi = T.axis.spatial(1, T.int64(0))
+                        T.reads(rxplaceholder[()])
+                        T.writes(T_add[()])
+                        T_add[()] = rxplaceholder[()] + T.int64(1)
+    # fmt: on
+    # pylint: enable=no-self-argument,missing-class-docstring,line-too-long
+    target = tvm.target.Target("nvidia/geforce-rtx-3070")
+    with target, tvm.transform.PassContext(opt_level=0):
+        mod = DefaultGPUSchedule()(Before)
+    tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_sum():
+    # sum has two reduction axes and no spatial axis
+    # pylint: disable=no-self-argument,missing-class-docstring,line-too-long
+    # fmt: off
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def sum(A: T.Buffer((T.int64(2), T.int64(2)), "float64"), A_red: T.Buffer((), "float64")):
+            for k0, k1 in T.grid(T.int64(2), T.int64(2)):
+                with T.sblock("A_red"):
+                    v_k0, v_k1 = T.axis.remap("RR", [k0, k1])
+                    with T.init():
+                        A_red[()] = T.float64(0)
+                    A_red[()] = A_red[()] + A[v_k0, v_k1]
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def sum(A: T.Buffer((T.int64(2), T.int64(2)), "float64"), A_red: T.Buffer((), "float64")):
+            T.func_attr({"tirx.is_scheduled": True})
+            # with T.sblock("root"):
+            for u_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
+                for u_fused_1 in T.thread_binding(1, thread="threadIdx.x"):
+                    for k0, k1 in T.grid(T.int64(2), T.int64(2)):
+                        with T.sblock("A_red"):
+                            v_k0, v_k1 = T.axis.remap("RR", [k0, k1])
+                            T.reads(A[v_k0, v_k1])
+                            T.writes(A_red[()])
+                            with T.init():
+                                A_red[()] = T.float64(0)
+                            A_red[()] = A_red[()] + A[v_k0, v_k1]
+    # fmt: on
+    # pylint: enable=no-self-argument,missing-class-docstring,line-too-long
+    target = tvm.target.Target("nvidia/geforce-rtx-3070")
+    with target, tvm.transform.PassContext(opt_level=0):
+        mod = DefaultGPUSchedule()(Before)
+    tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_scalar_block_no_loops():
+    # A PrimFunc whose body is a bare SBlockRealize (e.g. a fully-scalar op)
+    # used to crash DefaultGPUSchedule with "Cannot add loops on top of the
+    # root block" because the realized block was the function's root sref.
+    # pylint: disable=no-self-argument,missing-class-docstring,line-too-long
+    # fmt: off
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def scalar_add(a: T.Buffer((), "float32"), b: T.Buffer((), "float32"), c: T.Buffer((), "float32")):
+            with T.sblock("scalar_add"):
+                c[()] = a[()] + b[()]
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def scalar_add(a: T.Buffer((), "float32"), b: T.Buffer((), "float32"), c: T.Buffer((), "float32")):
+            T.func_attr({"tirx.is_scheduled": True})
+            # with T.sblock("root"):
+            for u_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
+                for u_fused_1 in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("scalar_add"):
+                        vu = T.axis.spatial(1, 0)
+                        T.reads()
+                        T.writes()
+                        c[()] = a[()] + b[()]
+    # fmt: on
+    # pylint: enable=no-self-argument,missing-class-docstring,line-too-long
+    target = tvm.target.Target("nvidia/geforce-rtx-3070")
+    with target, tvm.transform.PassContext(opt_level=0):
+        mod = DefaultGPUSchedule()(Before)
+    tvm.ir.assert_structural_equal(mod, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_hoist_expression.py

@@ -0,0 +1,548 @@
+# 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
+import tvm.testing
+from tvm import s_tir
+from tvm.s_tir.transform import HoistedConditionals, HoistedLetBindings
+from tvm.script import tirx as T
+
+
+def _run_transform(before, hoisted_conditionals, hoisted_let_bindings):
+    """Run HoistExpression transform with the given configuration."""
+    before_mod = tvm.IRModule.from_expr(before)
+
+    config = {
+        "s_tir.HoistExpression": {
+            "hoisted_conditionals": hoisted_conditionals.value,
+            "hoisted_let_bindings": hoisted_let_bindings.value,
+        }
+    }
+
+    with tvm.transform.PassContext(config=config):
+        after_mod = tvm.s_tir.transform.HoistExpression()(before_mod)
+
+    return after_mod["main"]
+
+
+def test_hoist_to_top_if_else_stmt():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((16,), "float32"), n: T.int32):
+        for i in T.serial(16):
+            if n != 0:
+                A[i] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((16,), "float32"), n: T.int32):
+        if n != 0:
+            for i in T.serial(16):
+                A[i] = 0.0
+
+    after = _run_transform(before, HoistedConditionals.IfElseStmt, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_to_top_all():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((16,), "float32"), n: T.int32):
+        for i in T.serial(16):
+            if n != 0:
+                A[i] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((16,), "float32"), n: T.int32):
+        if n != 0:
+            for i in T.serial(16):
+                A[i] = 0.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_suppress_hoist_if_else_never():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((16,), "float32"), n: T.int32):
+        for i in T.serial(16):
+            if n != 0:
+                A[i] = 0.0
+
+    expected = before
+
+    after = _run_transform(before, HoistedConditionals.Never, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_suppress_hoist_if_else_expr_only():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((16,), "float32"), n: T.int32):
+        for i in T.serial(16):
+            if n != 0:
+                A[i] = 0.0
+
+    expected = before
+
+    after = _run_transform(before, HoistedConditionals.IfElseExpr, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_block_var():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((128, 16), "float32"), n: T.int32):
+        i = T.env_thread("threadIdx.x")
+        T.launch_thread(i, 128)
+
+        for j in T.serial(16):
+            if i < 32:
+                A[i, j] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((128, 16), "float32"), n: T.int32):
+        i = T.env_thread("threadIdx.x")
+        T.launch_thread(i, 128)
+
+        if i < 32:
+            for j in T.serial(16):
+                A[i, j] = 0.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_suppress_hoist_block_var():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((128, 16), "float32"), n: T.int32):
+        thread_x = T.env_thread("threadIdx.x")
+        T.launch_thread(thread_x, 128)
+
+        for i in T.thread_binding(0, 128, thread="threadIdx.x"):
+            if i < 32:
+                for j in T.serial(16):
+                    A[i, j] = 0.0
+
+    expected = before
+
+    after = _run_transform(
+        before,
+        HoistedConditionals.All & ~HoistedConditionals.UsingBlockVar,
+        HoistedLetBindings.All,
+    )
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_across_block_var():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((128, 16), "float32"), n: T.int32):
+        thread_x = T.env_thread("threadIdx.x")
+        T.launch_thread(thread_x, 128)
+
+        for i in T.thread_binding(0, 128, thread="threadIdx.x"):
+            if n == 0:
+                for j in T.serial(16):
+                    A[i, j] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((128, 16), "float32"), n: T.int32):
+        thread_x = T.env_thread("threadIdx.x")
+
+        if n == 0:
+            T.launch_thread(thread_x, 128)
+            for i in T.thread_binding(0, 128, thread="threadIdx.x"):
+                for j in T.serial(16):
+                    A[i, j] = 0.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_suppress_hoist_across_block_var():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((128, 16), "float32"), n: T.int32):
+        thread_x = T.env_thread("threadIdx.x")
+        T.launch_thread(thread_x, 128)
+
+        for i in T.thread_binding(0, 128, thread="threadIdx.x"):
+            for j in T.serial(16):
+                if n == 0:
+                    A[i, j] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((128, 16), "float32"), n: T.int32):
+        thread_x = T.env_thread("threadIdx.x")
+
+        T.launch_thread(thread_x, 128)
+        if n == 0:
+            for i in T.thread_binding(0, 128, thread="threadIdx.x"):
+                for j in T.serial(16):
+                    A[i, j] = 0.0
+
+    after = _run_transform(
+        before,
+        HoistedConditionals.All & ~HoistedConditionals.UsingBlockVar,
+        HoistedLetBindings.All,
+    )
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_to_middle():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            for j in T.serial(4):
+                if i < 3:
+                    A[i, j] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            if i < 3:
+                for j in T.serial(4):
+                    A[i, j] = 0.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_with_let():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            for j in T.serial(4):
+                condition: T.let[T.bool] = i < 3
+                if condition:
+                    A[i, j] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            condition: T.let[T.bool] = i < 3  # noqa: F841
+            if i < 3:
+                for j in T.serial(4):
+                    A[i, j] = T.float32(0.0)
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_disable_let():
+    """As test_hoist_with_let, but forbid hoisting of Bind
+
+    Because the condition depends on the let binding, it should no
+    longer be hoisted.  With Bind lifecycle management, the condition
+    var is erased at sequence boundaries, so the if-condition uses
+    the raw expression.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            for j in T.serial(4):
+                condition: T.let[T.bool] = i < 3
+                if condition:
+                    A[i, j] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i, j in T.grid(4, 4):
+            condition: T.let[T.bool] = i < 3  # noqa: F841
+            if i < 3:
+                A[i, j] = T.float32(0.0)
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.Never)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_if_else():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            for j in T.serial(4):
+                if i < 3:
+                    A[i, j] = 0.0
+                else:
+                    A[i, j] = 1.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            if i < 3:
+                for j in T.serial(4):
+                    A[i, j] = 0.0
+            else:
+                for j in T.serial(4):
+                    A[i, j] = 1.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_sequential_assign():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32"), B: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            for j in T.serial(4):
+                if i < 3:
+                    A[i, j] = 0.0
+                    B[i, j] = 0.0
+                else:
+                    A[i, j] = 1.0
+                    B[i, j] = 1.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32"), B: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            if i < 3:
+                for j in T.serial(4):
+                    A[i, j] = 0.0
+                    B[i, j] = 0.0
+            else:
+                for j in T.serial(4):
+                    A[i, j] = 1.0
+                    B[i, j] = 1.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_multi_if():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            for j in T.serial(4):
+                for k in T.serial(4):
+                    if j < 3:
+                        if i < 2:
+                            A[i, j] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            if i < 2:
+                for j in T.serial(4):
+                    if j < 3:
+                        for k in T.serial(4):
+                            A[i, j] = 0.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_complex_conditional():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i, j, k in T.grid(4, 4, 4):
+            if j < 3 and i < 2:
+                A[i, j] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            if i < 2:
+                for j in T.serial(4):
+                    if j < 3:
+                        for k in T.serial(4):
+                            A[i, j] = 0.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_suppress_splitting_conditional():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i, j, k in T.grid(4, 4, 4):
+            if j < 3 and i < 2:
+                A[i, j] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i, j in T.grid(4, 4):
+            if j < 3 and i < 2:
+                for k in T.serial(4):
+                    A[i, j] = 0.0
+
+    after = _run_transform(
+        before,
+        HoistedConditionals.All & ~HoistedConditionals.BooleanExpression,
+        HoistedLetBindings.All,
+    )
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_multi_if_else():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            for j in T.serial(4):
+                for k in T.serial(4):
+                    if j < 3:
+                        if i < 2:
+                            A[i, j] = 0.0
+                        else:
+                            A[i, j] = 1.0
+                    else:
+                        if i < 2:
+                            A[i, j] = 2.0
+                        else:
+                            A[i, j] = 3.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            if i < 2:
+                for j in T.serial(4):
+                    if j < 3:
+                        for k in T.serial(4):
+                            A[i, j] = 0.0
+                    else:
+                        for k in T.serial(4):
+                            A[i, j] = 2.0
+            else:
+                for j in T.serial(4):
+                    if j < 3:
+                        for k in T.serial(4):
+                            A[i, j] = 1.0
+                    else:
+                        for k in T.serial(4):
+                            A[i, j] = 3.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_multi_if_else_different_branches():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            for j in T.serial(4):
+                for k in T.serial(4):
+                    if j < 3:
+                        if i < 2:
+                            A[i, j] = 0.0
+                        else:
+                            A[i, j] = 1.0
+                    else:
+                        if i < 1:
+                            A[i, j] = 2.0
+                        else:
+                            A[i, j] = 3.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            if i < 2:
+                if i < 1:
+                    for j in T.serial(4):
+                        if j < 3:
+                            for k in T.serial(4):
+                                A[i, j] = 0.0
+                        else:
+                            for k in T.serial(4):
+                                A[i, j] = 2.0
+                else:
+                    for j in T.serial(4):
+                        if j < 3:
+                            for k in T.serial(4):
+                                A[i, j] = 0.0
+                        else:
+                            for k in T.serial(4):
+                                A[i, j] = 3.0
+            else:
+                for j in T.serial(4):
+                    if j < 3:
+                        for k in T.serial(4):
+                            A[i, j] = 1.0
+                    else:
+                        for k in T.serial(4):
+                            A[i, j] = 3.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_if_else_expr():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i, j in T.grid(4, 4):
+            A[i, j] = T.if_then_else(i < 2, 1.0, 2.0, dtype="float32")
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            if i < 2:
+                for j in T.serial(4):
+                    A[i, j] = 1.0
+            else:
+                for j in T.serial(4):
+                    A[i, j] = 2.0
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_suppress_hoist_if_else_expr():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i, j in T.grid(4, 4):
+            A[i, j] = T.if_then_else(i < 2, 1.0, 2.0, dtype="float32")
+
+    expected = before
+
+    after = _run_transform(
+        before,
+        HoistedConditionals.All & ~HoistedConditionals.IfElseExpr,
+        HoistedLetBindings.All,
+    )
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_hoist_let_expr():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i, j in T.grid(4, 4):
+            x = T.float32()
+            A[i, j] = T.Let(5.0 * x + T.cast(j, "float32"), where={x: T.cast(i + 1, "float32")})
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32")):
+        for i in T.serial(4):
+            x: T.let[T.float32] = T.cast(i + 1, "float32")  # noqa: F841
+            for j in T.serial(4):
+                A[i, j] = T.float32(5.0) * T.cast(i + 1, "float32") + T.cast(j, "float32")
+
+    after = _run_transform(before, HoistedConditionals.All, HoistedLetBindings.All)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_suppress_hoist_let_expr():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((4, 4), "float32")):
+        for i, j in T.grid(4, 4):
+            x = T.float32()
+            A[i, j] = T.Let(5.0 * x + T.cast(j, "float32"), where={x: T.cast(i + 1, "float32")})
+
+    expected = before
+
+    after = _run_transform(
+        before,
+        HoistedConditionals.All,
+        HoistedLetBindings.All & ~HoistedLetBindings.LetExpr,
+    )
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py

@@ -0,0 +1,569 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E741, F401
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+from tvm import s_tir
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import enabled_targets
+
+var_list = []
+
+
+def verify_structure(stmt, expected_struct):
+    node_dict = {}
+    struct = {}
+
+    def _extract_vars(op):
+        global var_list
+        if isinstance(op, tvm.tirx.Var):
+            var_list.append(op.name)
+
+    def _visit(op):
+        key = op
+        if isinstance(op, tvm.tirx.IfThenElse):
+            global var_list
+            tvm.tirx.stmt_functor.post_order_visit(op.condition, _extract_vars)
+            val = [(op.then_case, op.else_case), ("tirx.IfThenElse", tuple(var_list))]
+            var_list.clear()
+        elif isinstance(op, tvm.tirx.For):
+            val = [(op.body,), ("tirx.For", op.loop_var.name)]
+        elif isinstance(op, tvm.tirx.AttrStmt):
+            val = [(op.body,), ("tirx.AttrStmt", op.attr_key, int(op.value))]
+        else:
+            return
+        node_dict[key] = val
+
+    tvm.tirx.stmt_functor.post_order_visit(stmt, _visit)
+    for key, val in node_dict.items():
+        struct[val[1]] = tuple(
+            node_dict[child][1] if child in node_dict else None for child in val[0]
+        )
+
+    assert struct == expected_struct, (
+        f"Structure mismatch: expect {expected_struct} but got {struct}"
+    )
+    var_list.clear()
+
+
+def _opaque_eval(var):
+    return tvm.tirx.Evaluate(tvm.tirx.call_extern("int32", "dummy", var))
+
+
+def test_hoist_top_for():
+    @T.prim_func(private=True, s_tir=True)
+    def func(l: T.int32, m: T.int32, n: T.int32):
+        for i in T.serial(l):
+            for j in T.serial(m):
+                for k in T.serial(n):
+                    if T.likely(i < 2):
+                        T.evaluate(T.call_extern("int32", "dummy", m))
+                    else:
+                        T.evaluate(T.call_extern("int32", "dummy", n))
+
+    mod = tvm.IRModule.from_expr(func)
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    expected_struct = {
+        ("tirx.For", "k"): (None,),
+        ("tirx.For", "j"): (("tirx.For", "k"),),
+        ("tirx.IfThenElse", ("i",)): (("tirx.For", "j"), ("tirx.For", "j")),
+        ("tirx.For", "i"): (("tirx.IfThenElse", ("i",)),),
+    }
+    verify_structure(new_stmt, expected_struct)
+
+
+def test_hoist_multi_var_if():
+    @T.prim_func(private=True, s_tir=True)
+    def func(l: T.int32, m: T.int32, n: T.int32):
+        for i in T.serial(l):
+            for j in T.serial(m):
+                for k in T.serial(n):
+                    if T.likely(i + j < 2):
+                        T.evaluate(T.call_extern("int32", "dummy", m))
+                    else:
+                        T.evaluate(T.call_extern("int32", "dummy", n))
+
+    mod = tvm.IRModule.from_expr(func)
+    new_mod = tvm.s_tir.transform.HoistIfThenElse()(mod)
+    new_stmt = new_mod["main"].body
+    expected_struct = {
+        ("tirx.For", "k"): (None,),
+        ("tirx.IfThenElse", ("i", "j")): (("tirx.For", "k"), ("tirx.For", "k")),
+        ("tirx.For", "j"): (("tirx.IfThenElse", ("i", "j")),),
+        ("tirx.For", "i"): (("tirx.For", "j"),),
+    }
+    verify_structure(new_stmt, expected_struct)
+
+
+def test_hoist_no_match_for():
+    @T.prim_func(private=True, s_tir=True)
+    def func(data: T.handle("float32"), l: T.int32, m: T.int32, n: T.int32):
+        data_ptr = T.decl_buffer(1, "float32", data=data)
+        for i in T.serial(l):
+            for j in T.serial(m):
+                data_ptr[i * 3 + j] = data_ptr[i * 3 + j] + T.float32(0.5)
+                for k in T.serial(n):
+                    if T.likely(i < 2):
+                        T.evaluate(T.call_extern("int32", "dummy", m))
+                    else:
+                        T.evaluate(T.call_extern("int32", "dummy", n))
+
+    mod = tvm.IRModule.from_expr(func)
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    expected_struct = {
+        ("tirx.For", "k"): (None,),
+        ("tirx.IfThenElse", ("i",)): (("tirx.For", "k"), ("tirx.For", "k")),
+        ("tirx.For", "j"): (None,),
+        ("tirx.For", "i"): (("tirx.For", "j"),),
+    }
+    verify_structure(new_stmt, expected_struct)
+
+
+def test_no_else():
+    @T.prim_func(private=True, s_tir=True)
+    def func(l: T.int32, m: T.int32, n: T.int32):
+        for i in T.serial(l):
+            for j in T.serial(m):
+                for k in T.serial(n):
+                    if T.likely(i < 2):
+                        T.evaluate(T.call_extern("int32", "dummy", m))
+
+    mod = tvm.IRModule.from_expr(func)
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    expected_struct = {
+        ("tirx.For", "k"): (None,),
+        ("tirx.For", "j"): (("tirx.For", "k"),),
+        ("tirx.IfThenElse", ("i",)): (("tirx.For", "j"), None),
+        ("tirx.For", "i"): (("tirx.IfThenElse", ("i",)),),
+    }
+    verify_structure(new_stmt, expected_struct)
+
+
+def test_attr_stmt():
+    dshape = (32, 64)
+
+    @T.prim_func(private=True, s_tir=True)
+    def func(data: T.handle("float32"), l: T.int32, m: T.int32, n: T.int32):
+        data_ptr = T.decl_buffer(1, "float32", data=data)
+        tx = T.launch_thread("threadIdx.x", dshape[0])
+        bx = T.launch_thread("blockIdx.x", dshape[1])
+        for i in T.serial(l):
+            for j in T.serial(m):
+                for k in T.serial(n):
+                    if i < 4 or j >= 8:
+                        data_ptr[bx * j + tx * j * k] = data_ptr[bx * j + tx * j * k] + T.float32(
+                            0.5
+                        )
+                    else:
+                        data_ptr[bx * j + tx * j * k] = data_ptr[bx * j + tx * j * k] + T.float32(
+                            1.0
+                        )
+
+    mod = tvm.IRModule.from_expr(func)
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    expected_struct = {
+        ("tirx.For", "k"): (None,),
+        ("tirx.IfThenElse", ("i", "j")): (("tirx.For", "k"), ("tirx.For", "k")),
+        ("tirx.For", "j"): (("tirx.IfThenElse", ("i", "j")),),
+        ("tirx.For", "i"): (("tirx.For", "j"),),
+        ("tirx.AttrStmt", "thread_extent", 64): (("tirx.For", "i"),),
+        ("tirx.AttrStmt", "thread_extent", 32): (("tirx.AttrStmt", "thread_extent", 64),),
+    }
+    verify_structure(new_stmt, expected_struct)
+
+
+def test_nested_for():
+    @T.prim_func(private=True, s_tir=True)
+    def func(data: T.handle("float32")):
+        data_ptr = T.decl_buffer(1, "float32", data=data)
+        for i in range(5):
+            for j in range(10):
+                if i >= 3:
+                    data_ptr[i * 3 + j] = data_ptr[i * 3 + j] + T.float32(0.5)
+                    for k in range(15):
+                        for l in range(20):
+                            if i < 4 or j >= 8:
+                                data_ptr[i * 3 + j + k + l] = data_ptr[
+                                    i * 3 + j + k + l
+                                ] * T.float32(2)
+                            else:
+                                data_ptr[i * 3 + j + k + l] = data_ptr[
+                                    i * 3 + j + k + l
+                                ] * T.float32(1.5)
+
+    mod = tvm.IRModule.from_expr(func)
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    expected_struct = {
+        ("tirx.For", "l"): (None,),
+        ("tirx.For", "k"): (("tirx.For", "l"),),
+        ("tirx.IfThenElse", ("i", "j")): (("tirx.For", "k"), ("tirx.For", "k")),
+        ("tirx.For", "j"): (None,),
+        ("tirx.IfThenElse", ("i",)): (("tirx.For", "j"), None),
+        ("tirx.For", "i"): (("tirx.IfThenElse", ("i",)),),
+    }
+    verify_structure(new_stmt, expected_struct)
+
+
+def test_if_block():
+    # Use different variable names for second loop nest to avoid dict key collision
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def main(data: T.Buffer((1,), "float32"), n: T.int32):
+            # First loop nest: i, j, k, l
+            for i in T.serial(5):
+                for j in T.serial(10):
+                    if i >= 3:
+                        data[i * 3 + j] = data[i * 3 + j] + T.float32(0.5)
+                        for k in T.serial(15):
+                            for l in T.serial(20):
+                                if i < 4 or j >= 8:
+                                    data[i * 3 + j + k + l] = data[i * 3 + j + k + l] * T.float32(2)
+                                else:
+                                    data[i * 3 + j + k + l] = data[i * 3 + j + k + l] * T.float32(
+                                        1.5
+                                    )
+                                if j < 5:
+                                    data[i * 3 + j + k + l] = data[i * 3 + j + k + l] - T.float32(1)
+
+            # Second loop nest: i2, j2, k2 (different names)
+            for i2 in T.serial(5):
+                for j2 in T.serial(10):
+                    for k2 in T.serial(15):
+                        if n >= 3:
+                            data[i2 * 3 + j2 + k2] = data[i2 * 3 + j2 + k2] + T.float32(0.6)
+
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    # Updated expected_struct with renamed second nest variables
+    expected_struct = {
+        ("tirx.IfThenElse", ("i", "j")): (None, None),
+        ("tirx.IfThenElse", ("j",)): (None, None),
+        ("tirx.For", "l"): (None,),
+        ("tirx.For", "k"): (("tirx.For", "l"),),
+        ("tirx.For", "j"): (None,),
+        ("tirx.IfThenElse", ("i",)): (("tirx.For", "j"), None),
+        ("tirx.For", "i"): (("tirx.IfThenElse", ("i",)),),
+        ("tirx.For", "k2"): (None,),
+        ("tirx.For", "j2"): (("tirx.For", "k2"),),
+        ("tirx.For", "i2"): (("tirx.For", "j2"),),
+        ("tirx.IfThenElse", ("n",)): (("tirx.For", "i2"), None),
+    }
+    verify_structure(new_stmt, expected_struct)
+
+
+def test_multi_if():
+    @T.prim_func(private=True, s_tir=True)
+    def func(data: T.handle("float32")):
+        data_ptr = T.decl_buffer(1, "float32", data=data)
+        for i in range(10):
+            for j in range(10):
+                for k in range(10):
+                    if 3 <= i:
+                        if 3 <= j:
+                            data_ptr[i * 100 + j * 10 + k] = data_ptr[
+                                i * 100 + j * 10 + k
+                            ] + T.float32(0.5)
+
+    mod = tvm.IRModule.from_expr(func)
+    new_mod = tvm.s_tir.transform.HoistIfThenElse()(mod)
+    new_stmt = new_mod["main"].body
+    expected_struct = {
+        ("tirx.For", "k"): (None,),
+        ("tirx.IfThenElse", ("j",)): (("tirx.For", "k"), None),
+        ("tirx.For", "j"): (("tirx.IfThenElse", ("j",)),),
+        ("tirx.IfThenElse", ("i",)): (("tirx.For", "j"), None),
+        ("tirx.For", "i"): (("tirx.IfThenElse", ("i",)),),
+    }
+    verify_structure(new_stmt, expected_struct)
+
+
+def test_no_hoisting_1():
+    @T.prim_func(private=True, s_tir=True)
+    def func(data: T.handle("float32")):
+        data_ptr = T.decl_buffer(1, "float32", data=data)
+        for i in range(10):
+            for j in range(10):
+                for k in range(10):
+                    if k <= 3:
+                        data_ptr[i * 100 + j * 10 + k] = data_ptr[i * 100 + j * 10 + k] + T.float32(
+                            0.5
+                        )
+
+    mod = tvm.IRModule.from_expr(func)
+    stmt = mod["main"].body
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+    with tvm.transform.PassContext(
+        config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting": True}}
+    ):
+        new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+
+def test_no_hoisting_2():
+    @T.prim_func(private=True, s_tir=True)
+    def func(data: T.handle("float32")):
+        data_ptr = T.decl_buffer(1, "float32", data=data)
+        for i in range(10):
+            for j in range(10):
+                for k in range(10):
+                    if i <= 3:
+                        data_ptr[i * 100 + j * 10 + k] = data_ptr[i * 100 + j * 10 + k] + T.float32(
+                            0.3
+                        )
+                    data_ptr[i * 100 + j * 10 + k] = data_ptr[i * 100 + j * 10 + k] + T.float32(0.5)
+
+    mod = tvm.IRModule.from_expr(func)
+    stmt = mod["main"].body
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+    with tvm.transform.PassContext(
+        config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting": True}}
+    ):
+        new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+
+def test_no_hoisting_4():
+    dshape = (32, 64)
+    dshape_inner = (33, 63)
+
+    # Create iter_var for tx (used inside loop with T.attr)
+    tx_var = tvm.tirx.Var("threadIdx.x", "int32")
+    tx_iter = tvm.tirx.IterVar(
+        tvm.ir.Range(0, dshape_inner[0]), tx_var, tvm.tirx.IterVar.ThreadIndex, "threadIdx.x"
+    )
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def main(data: T.Buffer((1,), "float32"), l: T.int32, m: T.int32, n: T.int32):
+            bx = T.launch_thread("blockIdx.x", dshape[1])
+            for i in T.serial(l):
+                for j in T.serial(m):
+                    for k in T.serial(n):
+                        T.attr(tx_iter, "thread_extent", dshape_inner[0])
+                        if tx_var < 3:
+                            data[bx * j + tx_var * j * k] = data[
+                                bx * j + tx_var * j * k
+                            ] + T.float32(0.3)
+                        else:
+                            data[bx * j + tx_var * j * k] = data[
+                                bx * j + tx_var * j * k
+                            ] + T.float32(1.3)
+
+    stmt = Module["main"].body
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+    with tvm.transform.PassContext(
+        config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting": True}}
+    ):
+        new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+
+def test_no_hoisting_6():
+    dshape = (32, 64)
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def main(data: T.Buffer((1,), "float32"), l: T.int32, m: T.int32, n: T.int32):
+            tx = T.launch_thread("threadIdx.x", dshape[0])
+            bx = T.launch_thread("blockIdx.x", dshape[1])
+            for i in T.serial(l):
+                for j in T.serial(m):
+                    for k in T.serial(n):
+                        if tx + k < 3:
+                            data[bx * j + tx * j * k] = data[bx * j + tx * j * k] + T.float32(0.3)
+                        else:
+                            data[bx * j + tx * j * k] = data[bx * j + tx * j * k] + T.float32(1.3)
+
+    stmt = Module["main"].body
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+    with tvm.transform.PassContext(
+        config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting": True}}
+    ):
+        new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+
+def test_no_hoisting_7():
+    dshape = (32, 64)
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def main(data: T.Buffer((1,), "float32"), l: T.int32, m: T.int32, n: T.int32):
+            tx = T.launch_thread("threadIdx.x", dshape[0])
+            bx = T.launch_thread("blockIdx.x", dshape[1])
+            for i in T.serial(l):
+                for j in T.serial(m):
+                    if tx + j < 9:
+                        for k in T.serial(n):
+                            if tx + k < 3:
+                                data[bx * j + tx * j * k] = data[bx * j + tx * j * k] + T.float32(
+                                    0.3
+                                )
+
+    stmt = Module["main"].body
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+    with tvm.transform.PassContext(
+        config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting": True}}
+    ):
+        new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+
+def test_hoisting_block_scope_2():
+    dshape = (32, 64)
+
+    # Create iter_var for bx (used inside loop with T.attr)
+    bx_var = tvm.tirx.Var("blockIdx.x", "int32")
+    bx_iter = tvm.tirx.IterVar(
+        tvm.ir.Range(0, dshape[1]), bx_var, tvm.tirx.IterVar.ThreadIndex, "blockIdx.x"
+    )
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def main(data: T.Buffer((1,), "float32"), l: T.int32, m: T.int32, n: T.int32):
+            tx = T.launch_thread("threadIdx.x", dshape[0])
+            for i in T.serial(l):
+                for j in T.serial(m):
+                    for k in T.serial(n):
+                        T.attr(bx_iter, "thread_extent", dshape[1])
+                        if tx < 3:
+                            data[bx_var * j + tx * j * k] = data[
+                                bx_var * j + tx * j * k
+                            ] + T.float32(0.3)
+                        else:
+                            data[bx_var * j + tx * j * k] = data[
+                                bx_var * j + tx * j * k
+                            ] + T.float32(1.3)
+
+    mod = Module
+    mod = tvm.tirx.transform.StmtSimplify()(mod)
+    mod = tvm.tirx.transform.RemoveNoOp()(mod)
+    stmt = mod["main"].body
+
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+    with tvm.transform.PassContext(
+        config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting": True}}
+    ):
+        new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    assert not tvm_ffi.structural_equal(new_stmt, stmt)
+
+
+def test_hoisting_block_scope_5():
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def main(data: T.Buffer((1,), "float32"), l: T.int32, m: T.int32, n: T.int32, g: T.int32):
+            for i in T.serial(l):
+                for j in T.serial(m):
+                    for k in T.serial(n):
+                        if data[g] < T.float32(3):
+                            data[9 * j + 3 * j * k] = data[9 * j + 3 * j * k] + T.float32(0.3)
+                        else:
+                            data[9 * j + 3 * j * k] = data[9 * j + 3 * j * k] + T.float32(1.3)
+
+    stmt = Module["main"].body
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    assert not tvm_ffi.structural_equal(new_stmt, stmt)
+
+    mod = tvm.IRModule.from_expr(tvm.tirx.PrimFunc([], new_stmt))
+    stmt = new_stmt
+
+    with tvm.transform.PassContext(
+        config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting": True}}
+    ):
+        new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+
+def test_hoisting_block_scope_6():
+    dshape = (32, 64)
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def main(data: T.Buffer((1,), "float32"), l: T.int32, m: T.int32, n: T.int32):
+            tx = T.launch_thread("threadIdx.x", dshape[0])
+            bx = T.launch_thread("blockIdx.x", dshape[1])
+            for i in T.serial(l):
+                for j in T.serial(m):
+                    for k in T.serial(n):
+                        if tx + n < 3:
+                            data[bx * j + tx * j * k] = data[bx * j + tx * j * k] + T.float32(0.3)
+                        else:
+                            data[bx * j + tx * j * k] = data[bx * j + tx * j * k] + T.float32(1.3)
+
+    stmt = Module["main"].body
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+    with tvm.transform.PassContext(
+        config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting": True}}
+    ):
+        new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    assert not tvm_ffi.structural_equal(new_stmt, stmt)
+
+
+def test_hoisting_block_scope_7():
+    dshape = (32, 64)
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def main(data: T.Buffer((1,), "float32"), l: T.int32, m: T.int32, n: T.int32):
+            tx = T.launch_thread("threadIdx.x", dshape[0])
+            bx = T.launch_thread("blockIdx.x", dshape[1])
+            for i in T.serial(l):
+                for j in T.serial(m):
+                    for k in T.serial(n):
+                        if tx + i < 3:
+                            data[bx * j + tx * j * k] = data[bx * j + tx * j * k] + T.float32(0.3)
+                        else:
+                            data[bx * j + tx * j * k] = data[bx * j + tx * j * k] + T.float32(1.3)
+
+    stmt = Module["main"].body
+    new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    tvm.ir.assert_structural_equal(new_stmt, stmt)
+
+    with tvm.transform.PassContext(
+        config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting": True}}
+    ):
+        new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
+    assert not tvm_ffi.structural_equal(new_stmt, stmt)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py

@@ -0,0 +1,165 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F841
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_double_buffer():
+    n = 100
+    m = 4
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def db(A: T.handle("float32"), C: T.handle("float32")):
+            A_buf = T.decl_buffer((n * m,), "float32", data=A)
+            C_buf = T.decl_buffer((m,), "float32", data=C)
+            tx = T.launch_thread("threadIdx.x", 1)
+            for i in range(n):
+                B = T.alloc_buffer((m,), "float32", scope="shared")
+                with T.attr(B.data, "double_buffer_scope", 1):
+                    for j in range(m):
+                        B[j] = A_buf[i * 4 + j]
+                for j in range(m):
+                    C_buf[j] = B[j] + T.float32(1.0)
+
+    mod = Module
+
+    opt = tvm.transform.Sequential(
+        [tvm.s_tir.transform.InjectDoubleBuffer(), tvm.tirx.transform.StmtSimplify()]
+    )
+
+    with tvm.transform.PassContext(config={"s_tir.InjectDoubleBuffer": {"split_loop": 2}}):
+        mod = opt(mod)
+    stmt = mod["db"].body
+
+    # After transformation, the buffer allocation should be doubled
+    allocate_node = None
+
+    def visitor(op):
+        nonlocal allocate_node
+        if isinstance(op, tvm.tirx.AllocBuffer) and "B" in str(op.buffer.data):
+            allocate_node = op
+
+    tvm.tirx.stmt_functor.post_order_visit(stmt, visitor)
+    assert allocate_node is not None
+    assert list(allocate_node.buffer.shape) == [m * 2]
+
+    f = tvm.s_tir.transform.ThreadSync("shared")(mod)["db"]
+    count = [0]
+
+    def count_sync(op):
+        if isinstance(op, tvm.ir.Call) and op.op.same_as(tvm.ir.Op.get("tirx.tvm_storage_sync")):
+            count[0] += 1
+
+    tvm.tirx.stmt_functor.post_order_visit(f.body, count_sync)
+    assert count[0] == 4
+
+
+def test_double_buffer_transform():
+    transform = tvm.ir.transform.Sequential(
+        [
+            tvm.s_tir.transform.InjectDoubleBuffer(),
+            tvm.tirx.transform.StmtSimplify(),
+        ]
+    )
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer([16, 32], "float32"), B: T.Buffer(16, "float32")):
+            for i in range(16):
+                cache = T.alloc_buffer((32,), "float32")
+
+                T.attr(cache.data, "double_buffer_scope", 1)
+
+                for j in range(32):
+                    cache[j] = A[i, j]
+
+                B[i] = 0.0
+                for j in range(32):
+                    B[i] = B[i] + cache[j]
+
+    After = transform(Before)
+    # Verify the double buffer transformation: the allocation should be doubled (32 -> 64)
+    allocate_node = None
+
+    def visitor(op):
+        nonlocal allocate_node
+        if isinstance(op, tvm.tirx.AllocBuffer):
+            allocate_node = op
+
+    tvm.tirx.stmt_functor.post_order_visit(After["main"].body, visitor)
+    assert allocate_node is not None
+    assert list(allocate_node.buffer.shape) == [64]
+
+
+def test_double_buffer_with_decl_buffer():
+    """Like test_double_buffer_transform, but with a declared buffer object"""
+
+    transform = tvm.ir.transform.Sequential(
+        [
+            tvm.s_tir.transform.InjectDoubleBuffer(),
+            tvm.tirx.transform.StmtSimplify(),
+        ]
+    )
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 32), "float32"), B: T.Buffer(16, "float32")):
+            for i in range(16):
+                cache = T.decl_buffer(32, "float32")
+                T.attr(cache.data, "double_buffer_scope", 1)
+
+                for j in range(32):
+                    cache[j] = A[i, j]
+
+                B[i] = 0.0
+                for j in range(32):
+                    B[i] = B[i] + cache[j]
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 32), "float32"), B: T.Buffer(16, "float32")):
+            cache = T.decl_buffer(64, "float32")
+            for j in range(32):
+                cache[j] = A[0, j]
+
+            B[0] = T.float32(0.0)
+            for j in range(32):
+                B[0] = B[0] + cache[j]
+
+            for i_outer in range(15):
+                T.attr(cache.data, "double_buffer_write", 1)
+                for j in range(32):
+                    cache[(i_outer + 1) % 2 * 32 + j] = A[i_outer + 1, j]
+                B[i_outer + 1] = T.float32(0.0)
+                for j in range(32):
+                    B[i_outer + 1] = B[i_outer + 1] + cache[(i_outer + 1) % 2 * 32 + j]
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py

@@ -0,0 +1,353 @@
+# 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, E741
+import tvm
+import tvm.s_tir
+import tvm.testing
+from tvm import IRModule
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc
+
+
+def _check_primfunc_transform(before: PrimFunc, expected: PrimFunc):
+    before_module = IRModule.from_expr(before)
+    after_module = tvm.s_tir.transform.InjectPermutedLayout()(before_module)
+
+    after = after_module["before"].without_attr("global_symbol")
+    expected = expected.without_attr("global_symbol")
+
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+# This pass is adapted from another previous pass, so we need to ensure backward compatibility here
+def test_backward_compatibility_shared_a():
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def before(X: T.Buffer((4096, 4096), "float16")):
+        # with T.sblock("root"):
+        for blockIdx_y in T.thread_binding(256, thread="blockIdx.y"):
+            for threadIdx_y in T.thread_binding(4, thread="threadIdx.y"):
+                for threadIdx_x in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock(""):
+                        T.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, threadIdx_x % 4 * 8:threadIdx_x % 4 * 8 + 4072])
+                        T.writes()
+                        for ax2_0_0 in range(128):
+                            with T.sblock(""):
+                                T.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, ax2_0_0 * 32 + threadIdx_x % 4 * 8:ax2_0_0 * 32 + threadIdx_x % 4 * 8 + 8])
+                                T.writes()
+                                X_reindex_shared_dyn = T.sblock_alloc_buffer((128, 32), "float16", strides=(32, 1), scope="shared.dyn")
+                                with T.sblock("X_reindex_shared.dyn"):
+                                    T.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, ax2_0_0 * 32 + threadIdx_x % 4 * 8:ax2_0_0 * 32 + threadIdx_x % 4 * 8 + 8])
+                                    T.writes(X_reindex_shared_dyn[threadIdx_y * 8 + threadIdx_x // 4:threadIdx_y * 8 + threadIdx_x // 4 + 97, threadIdx_x % 4 * 8:threadIdx_x % 4 * 8 + 8])
+                                    T.sblock_attr({"permuted_layout": "g2s_A"})
+                                    for ax0_ax1_fused_0 in range(4):
+                                        for ax0_ax1_fused_3 in T.vectorized(8):
+                                            X_reindex_shared_dyn[ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, threadIdx_x % 4 * 8 + ax0_ax1_fused_3] = X[blockIdx_y // 8 * 128 + ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, ax2_0_0 * 32 + threadIdx_x % 4 * 8 + ax0_ax1_fused_3]
+                                for ax2_0_1 in range(4):
+                                    with T.sblock(""):
+                                        T.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64:threadIdx_y // 2 * 64 + 64, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
+                                        T.writes()
+                                        X_reindex_shared_dyn_m16n8k8_matrixA = T.sblock_alloc_buffer((64, 8), "float16", scope="m16n8k8.matrixA")
+                                        for ax0_0, ax1_0 in T.grid(2, 1):
+                                            with T.sblock("X_reindex_shared.dyn_m16n8k8.matrixA_o"):
+                                                T.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y // 2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
+                                                T.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
+                                                T.sblock_attr({"permuted_layout": "s2l_A"})
+                                                T.ptx.ldmatrix_legacy("float16", T.bool(False), 4, ".b16", X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8, T.tvm_access_ptr(T.type_annotation("float16"), X_reindex_shared_dyn.data, threadIdx_y // 2 * 2048 + ax0_0 * 1024 + ax2_0_1 * 8, 1024, 1), threadIdx_x * 32)
+
+    @T.prim_func(s_tir=True)
+    def expected(X: T.Buffer((4096, 4096), "float16")):
+        for blockIdx_y in T.thread_binding(256, thread="blockIdx.y"):
+            for threadIdx_y in T.thread_binding(4, thread="threadIdx.y"):
+                for threadIdx_x in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock(""):
+                        for ax2_0_0 in T.serial(128):
+                            with T.sblock(""):
+                                X_reindex_shared_dyn = T.sblock_alloc_buffer((128, 32), "float16", strides=(32, 1), scope="shared.dyn")
+                                with T.sblock("X_reindex_shared.dyn"):
+                                    # annotate the reads and writes because they cannot be inferred from tirx.bitwise_xor
+                                    T.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, ax2_0_0 * 32 + threadIdx_x % 4 * 8:ax2_0_0 * 32 + threadIdx_x % 4 * 8 + 8])
+                                    T.writes(X_reindex_shared_dyn[threadIdx_y * 8 + threadIdx_x // 4:threadIdx_y * 8 + threadIdx_x // 4 + 97, threadIdx_x % 4 * 8:threadIdx_x % 4 * 8 + 8])
+                                    for ax0_ax1_fused_0 in range(4):
+                                        for ax0_ax1_fused_3 in T.vectorized(8):
+                                            X_reindex_shared_dyn[ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, T.bitwise_xor(threadIdx_x % 4, threadIdx_x // 8) * 8 + ax0_ax1_fused_3] = X[blockIdx_y // 8 * 128 + ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, ax2_0_0 * 32 + threadIdx_x % 4 * 8 + ax0_ax1_fused_3]
+                                for ax2_0_1 in T.serial(4):
+                                    with T.sblock(""):
+                                        X_reindex_shared_dyn_m16n8k8_matrixA = T.sblock_alloc_buffer((64, 8), "float16", scope="m16n8k8.matrixA")
+                                        for ax0_0, ax1_0 in T.grid(2, 1):
+                                            with T.sblock("X_reindex_shared.dyn_m16n8k8.matrixA_o"):
+                                                T.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y // 2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
+                                                T.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
+                                                T.ptx.ldmatrix_legacy("float16", T.bool(False), 4, ".b16", X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8, T.tvm_access_ptr(T.type_annotation("float16"), X_reindex_shared_dyn.data, threadIdx_y // 2 * 2048 + ax0_0 * 1024 + threadIdx_x * 32 + T.bitwise_xor(ax2_0_1, threadIdx_x % 8 // 2) * 8, 1024, 1), 0)
+    # fmt: on
+    _check_primfunc_transform(before, expected)
+
+
+def test_backward_compatibility_shared_a_and_b():
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def before(X: T.Buffer((4096, 4096), "float16"), Y: T.Buffer((4096, 4096), "float16")):
+        for blockIdx_x in T.thread_binding(4, thread="blockIdx.x"):
+            for blockIdx_y in T.thread_binding(256, thread="blockIdx.y"):
+                for threadIdx_y in T.thread_binding(4, thread="threadIdx.y"):
+                    for threadIdx_x in T.thread_binding(32, thread="threadIdx.x"):
+                        with T.sblock(""):
+                            for ax2_0_0 in T.serial(128):
+                                with T.sblock(""):
+                                    X_reindex_shared_dyn = T.sblock_alloc_buffer((128, 32), "float16", strides=(32, 1), scope="shared.dyn")
+                                    Y_reindex_shared_dyn = T.sblock_alloc_buffer((32, 128), "float16", strides=(128, 1), scope="shared.dyn")
+                                    with T.sblock("X_reindex_shared.dyn"):
+                                        T.sblock_attr({"permuted_layout": "g2s_A"})
+                                        for ax0_ax1_fused_0 in range(4):
+                                            for ax0_ax1_fused_3 in T.vectorized(8):
+                                                X_reindex_shared_dyn[ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, threadIdx_x % 4 * 8 + ax0_ax1_fused_3] = X[blockIdx_y // 8 * 128 + ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, ax2_0_0 * 32 + threadIdx_x % 4 * 8 + ax0_ax1_fused_3]
+                                    with T.sblock("Y_reindex_shared.dyn"):
+                                        T.sblock_attr({"permuted_layout": "g2s_B"})
+                                        for ax0_ax1_fused_0 in range(4):
+                                            for ax0_ax1_fused_3 in T.vectorized(8):
+                                                Y_reindex_shared_dyn[ax0_ax1_fused_0 * 8 + threadIdx_y * 2 + threadIdx_x // 16, threadIdx_x % 16 * 8 + ax0_ax1_fused_3] = Y[ax2_0_0 * 32 + ax0_ax1_fused_0 * 8 + threadIdx_y * 2 + threadIdx_x // 16, blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8 + ax0_ax1_fused_3]
+                                    for ax2_0_1 in T.serial(4):
+                                        with T.sblock(""):
+                                            X_reindex_shared_dyn_m16n8k8_matrixA = T.sblock_alloc_buffer((64, 8), "float16", scope="m16n8k8.matrixA")
+                                            Y_reindex_shared_dyn_m16n8k8_matrixB = T.sblock_alloc_buffer((8, 64), "float16", scope="m16n8k8.matrixB")
+                                            for ax0_0, ax1_0 in T.grid(2, 1):
+                                                with T.sblock("X_reindex_shared.dyn_m16n8k8.matrixA_o"):
+                                                    T.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y // 2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
+                                                    T.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
+                                                    T.sblock_attr({"permuted_layout": "s2l_A"})
+                                                    T.ptx.ldmatrix_legacy("float16", T.bool(False), 4, ".b16", X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8, T.tvm_access_ptr(T.type_annotation("float16"), X_reindex_shared_dyn.data, threadIdx_y // 2 * 2048 + ax0_0 * 1024 + ax2_0_1 * 8, 1024, 1), threadIdx_x * 32)
+                                            for ax0_0, ax1_0 in T.grid(1, 2):
+                                                with T.sblock("Y_reindex_shared.dyn_m16n8k8.matrixB_o"):
+                                                    T.reads(Y_reindex_shared_dyn[ax2_0_1 * 8:ax2_0_1 * 8 + 8, threadIdx_y % 2 * 64 + ax1_0 * 32:threadIdx_y % 2 * 64 + ax1_0 * 32 + 32])
+                                                    T.writes(Y_reindex_shared_dyn_m16n8k8_matrixB[0:8, ax1_0 * 32:ax1_0 * 32 + 32])
+                                                    T.sblock_attr({"permuted_layout": "s2l_B"})
+                                                    T.ptx.ldmatrix_legacy("float16", T.bool(True), 4, ".b16", Y_reindex_shared_dyn_m16n8k8_matrixB.data, ax1_0 * 8, T.tvm_access_ptr(T.type_annotation("float16"), Y_reindex_shared_dyn.data, ax2_0_1 * 1024 + threadIdx_y % 2 * 64 + ax1_0 * 32, 1024, 1), threadIdx_x % 8 * 128 + threadIdx_x // 8 * 8)
+
+    @T.prim_func(s_tir=True)
+    def expected(X: T.Buffer((4096, 4096), "float16"), Y: T.Buffer((4096, 4096), "float16")):
+        for blockIdx_x in T.thread_binding(4, thread="blockIdx.x"):
+            for blockIdx_y in T.thread_binding(256, thread="blockIdx.y"):
+                for threadIdx_y in T.thread_binding(4, thread="threadIdx.y"):
+                    for threadIdx_x in T.thread_binding(32, thread="threadIdx.x"):
+                        with T.sblock(""):
+                            T.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, threadIdx_x % 4 * 8:threadIdx_x % 4 * 8 + 4072], Y[threadIdx_y * 2 + threadIdx_x // 16:threadIdx_y * 2 + threadIdx_x // 16 + 4089, blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8:blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8 + 8])
+                            T.writes()
+                            for ax2_0_0 in T.serial(128):
+                                with T.sblock(""):
+                                    T.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, ax2_0_0 * 32 + threadIdx_x % 4 * 8:ax2_0_0 * 32 + threadIdx_x % 4 * 8 + 8], Y[ax2_0_0 * 32 + threadIdx_y * 2 + threadIdx_x // 16:ax2_0_0 * 32 + threadIdx_y * 2 + threadIdx_x // 16 + 25, blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8:blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8 + 8])
+                                    T.writes()
+                                    X_reindex_shared_dyn = T.sblock_alloc_buffer((128, 32), "float16", strides=(32, 1), scope="shared.dyn")
+                                    Y_reindex_shared_dyn = T.sblock_alloc_buffer((32, 128), "float16", strides=(128, 1), scope="shared.dyn")
+                                    with T.sblock("X_reindex_shared.dyn"):
+                                        T.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, ax2_0_0 * 32 + threadIdx_x % 4 * 8:ax2_0_0 * 32 + threadIdx_x % 4 * 8 + 8])
+                                        T.writes(X_reindex_shared_dyn[threadIdx_y * 8 + threadIdx_x // 4:threadIdx_y * 8 + threadIdx_x // 4 + 97, threadIdx_x % 4 * 8:threadIdx_x % 4 * 8 + 8])
+                                        for ax0_ax1_fused_0 in range(4):
+                                            for ax0_ax1_fused_3 in T.vectorized(8):
+                                                X_reindex_shared_dyn[ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, T.bitwise_xor(threadIdx_x % 4, threadIdx_x // 8) * 8 + ax0_ax1_fused_3] = X[blockIdx_y // 8 * 128 + ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, ax2_0_0 * 32 + threadIdx_x % 4 * 8 + ax0_ax1_fused_3]
+                                    with T.sblock("Y_reindex_shared.dyn"):
+                                        T.reads(Y[ax2_0_0 * 32 + threadIdx_y * 2 + threadIdx_x // 16:ax2_0_0 * 32 + threadIdx_y * 2 + threadIdx_x // 16 + 25, blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8:blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8 + 8])
+                                        T.writes(Y_reindex_shared_dyn[threadIdx_y * 2 + threadIdx_x // 16:threadIdx_y * 2 + threadIdx_x // 16 + 25, threadIdx_x % 16 * 8:threadIdx_x % 16 * 8 + 8])
+                                        for ax0_ax1_fused_0 in range(4):
+                                            for ax0_ax1_fused_3 in T.vectorized(8):
+                                                Y_reindex_shared_dyn[ax0_ax1_fused_0 * 8 + threadIdx_y * 2 + threadIdx_x // 16, T.bitwise_xor(threadIdx_x % 16, threadIdx_y * 2 + threadIdx_x // 16) * 8 + ax0_ax1_fused_3]   = Y[ax2_0_0 * 32 + ax0_ax1_fused_0 * 8 + threadIdx_y * 2 + threadIdx_x // 16, blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8 + ax0_ax1_fused_3]
+                                    for ax2_0_1 in T.serial(4):
+                                        with T.sblock(""):
+                                            X_reindex_shared_dyn_m16n8k8_matrixA = T.sblock_alloc_buffer((64, 8), "float16", scope="m16n8k8.matrixA")
+                                            Y_reindex_shared_dyn_m16n8k8_matrixB = T.sblock_alloc_buffer((8, 64), "float16", scope="m16n8k8.matrixB")
+                                            for ax0_0, ax1_0 in T.grid(2, 1):
+                                                with T.sblock("X_reindex_shared.dyn_m16n8k8.matrixA_o"):
+                                                    T.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y // 2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
+                                                    T.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
+                                                    T.ptx.ldmatrix_legacy("float16", T.bool(False), 4, ".b16", X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8, T.tvm_access_ptr(T.type_annotation("float16"), X_reindex_shared_dyn.data, threadIdx_y // 2 * 2048 + ax0_0 * 1024 + threadIdx_x * 32 + T.bitwise_xor(ax2_0_1, threadIdx_x % 8 // 2) * 8, 1024, 1), 0)
+                                            for ax0_0, ax1_0 in T.grid(1, 2):
+                                                with T.sblock("Y_reindex_shared.dyn_m16n8k8.matrixB_o"):
+                                                    T.reads(Y_reindex_shared_dyn[ax2_0_1 * 8:ax2_0_1 * 8 + 8, threadIdx_y % 2 * 64 + ax1_0 * 32:threadIdx_y % 2 * 64 + ax1_0 * 32 + 32])
+                                                    T.writes(Y_reindex_shared_dyn_m16n8k8_matrixB[0:8, ax1_0 * 32:ax1_0 * 32 + 32])
+                                                    T.ptx.ldmatrix_legacy("float16", T.bool(True), 4, ".b16", Y_reindex_shared_dyn_m16n8k8_matrixB.data, ax1_0 * 8, T.tvm_access_ptr(T.type_annotation("float16"), Y_reindex_shared_dyn.data, ax2_0_1 * 1024 + threadIdx_x % 8 * 128 + T.bitwise_xor(threadIdx_y % 2 * 8 + ax1_0 * 4 + threadIdx_x // 8, threadIdx_x % 8) * 8, 1024, 1), 0)
+    # fmt: on
+    _check_primfunc_transform(before, expected)
+
+
+def test_buffer_a():
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def before(p_A: T.handle):
+        A = T.match_buffer(p_A, (T.int64(128), T.int64(32)), "float16")
+        A_shared_dyn = T.sblock_alloc_buffer((T.int64(128), T.int64(32)), "float16", scope="shared.dyn")
+        A_warp = T.sblock_alloc_buffer((T.int64(4), T.int64(1), T.int64(32), T.int64(8)), "float16", scope="warp")
+        for threadIdx_z in T.thread_binding(T.int64(2), thread="threadIdx.z"):
+            for threadIdx_y in T.thread_binding(T.int64(2), thread="threadIdx.y"):
+                for threadIdx_x in T.thread_binding(T.int64(32), thread="threadIdx.x"):
+                    for v0 in range(T.int64(4)):
+                        for v1 in T.vectorized(T.int64(8)):
+                            with T.sblock("A_reindex_shared.dyn"):
+                                T.sblock_attr({"permuted_layout": 1})
+                                A_shared_dyn[
+                                    v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4),
+                                    threadIdx_x % T.int64(4) * T.int64(8) + v1
+                                ] = A[
+                                    (v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4)) % T.int64(32),
+                                    threadIdx_x % T.int64(4) * T.int64(8) + v1
+                                ]
+                    for v0, v1 in T.grid(T.int64(2), T.int64(4)):
+                        with T.sblock("A_reindex_shared.dyn_warp_o"):
+                            T.sblock_attr({"permuted_layout": 1})
+                            with T.sblock("A_reindex_shared.dyn_warp_o"):
+                                T.reads(A_shared_dyn[threadIdx_z * T.int64(64) + v1 * T.int64(16):threadIdx_z * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
+                                T.writes(A_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
+                                T.ptx.ldmatrix_legacy("float16", T.bool(False), 4, ".b16",
+                                    A_warp.data,
+                                    v1 * T.int64(256) + threadIdx_x * T.int64(8),
+                                    T.tvm_access_ptr(T.type_annotation("float16"),
+                                        A_shared_dyn.data,
+                                        threadIdx_z * T.int64(2048) + v1 * T.int64(512) + v0 * T.int64(16), T.int64(512),
+                                        1
+                                    ),
+                                    threadIdx_x % T.int64(16) * T.int64(32) + threadIdx_x // T.int64(16) * T.int64(8)
+                                )
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((T.int64(128), T.int64(32)), "float16")):
+        A_shared_dyn = T.sblock_alloc_buffer((T.int64(128), T.int64(32)), "float16", scope="shared.dyn")
+        A_warp = T.sblock_alloc_buffer((T.int64(4), T.int64(1), T.int64(32), T.int64(8)), "float16", scope="warp")
+        for threadIdx_z in T.thread_binding(T.int64(2), thread="threadIdx.z"):
+            for threadIdx_y in T.thread_binding(T.int64(2), thread="threadIdx.y"):
+                for threadIdx_x in T.thread_binding(T.int64(32), thread="threadIdx.x"):
+                    for v0 in range(T.int64(4)):
+                        for v1 in T.vectorized(T.int64(8)):
+                            with T.sblock("A_reindex_shared.dyn"):
+                                T.reads(A[(v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4)) % T.int64(32), threadIdx_x % T.int64(4) * T.int64(8) + v1])
+                                T.writes(A_shared_dyn[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1])
+                                A_shared_dyn[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), T.bitwise_xor(threadIdx_x % T.int64(4), threadIdx_x // T.int64(8)) * T.int64(8) + v1] = A[(v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4)) % T.int64(32), threadIdx_x % T.int64(4) * T.int64(8) + v1]
+                    for v0, v1 in T.grid(T.int64(2), T.int64(4)):
+                        with T.sblock("A_reindex_shared.dyn_warp_o"):
+                            T.reads(A_shared_dyn[threadIdx_z * T.int64(64) + v1 * T.int64(16):threadIdx_z * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
+                            T.writes(A_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
+                            with T.sblock("A_reindex_shared.dyn_warp_o"):
+                                T.reads(A_shared_dyn[threadIdx_z * T.int64(64) + v1 * T.int64(16):threadIdx_z * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
+                                T.writes(A_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
+                                T.ptx.ldmatrix_legacy("float16", T.bool(False), 4, ".b16", A_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8), T.tvm_access_ptr(T.type_annotation("float16"), A_shared_dyn.data, threadIdx_z * T.int64(2048) + v1 * T.int64(512) + threadIdx_x % T.int64(16) * T.int64(32) + T.bitwise_xor(v0 * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(8) // T.int64(2)) * T.int64(8), T.int64(512), 1), T.int64(0))
+
+    # fmt: on
+    _check_primfunc_transform(before, expected)
+
+
+def test_buffer_b():
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def before(B: T.Buffer((T.int64(128), T.int64(32)), "float16")):
+        B_shared_dyn = T.sblock_alloc_buffer((T.int64(128), T.int64(32)), "float16", scope="shared.dyn")
+        for threadIdx_z in T.thread_binding(T.int64(2), thread="threadIdx.z"):
+            for threadIdx_y in T.thread_binding(T.int64(2), thread="threadIdx.y"):
+                for threadIdx_x in T.thread_binding(T.int64(32), thread="threadIdx.x"):
+                    for v0 in range(T.int64(4)):
+                        for v1 in T.vectorized(T.int64(8)):
+                            with T.sblock("B_reindex_shared.dyn"):
+                                T.sblock_attr({"permuted_layout": 1})
+                                B_shared_dyn[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1] = B[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1]
+                    for v0 in range(T.int64(2)):
+                        with T.sblock(""):
+                            B_warp = T.sblock_alloc_buffer((T.int64(4), T.int64(1), T.int64(32), T.int64(8)), "float16", scope="warp")
+                            for v1 in range(T.int64(4)):
+                                with T.sblock("B_reindex_shared.dyn_warp_o"):
+                                    T.sblock_attr({"permuted_layout": 1})
+                                    with T.sblock("B_reindex_shared.dyn_warp_o"):
+                                        T.reads(B_shared_dyn[threadIdx_y * T.int64(64) + v1 * T.int64(16):threadIdx_y * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
+                                        T.writes(B_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
+                                        T.ptx.ldmatrix_legacy("float16", T.bool(False), 4, ".b16", B_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8), T.tvm_access_ptr(T.type_annotation("float16"), B_shared_dyn.data, threadIdx_y * T.int64(2048) + v1 * T.int64(512) + v0 * T.int64(16), T.int64(512), 1), threadIdx_x // T.int64(16) * T.int64(256) + threadIdx_x % T.int64(8) * T.int64(32) + threadIdx_x % T.int64(16) // T.int64(8) * T.int64(8))
+
+    @T.prim_func(s_tir=True)
+    def expected(B: T.Buffer((T.int64(128), T.int64(32)), "float16")):
+        B_shared_dyn = T.sblock_alloc_buffer((T.int64(128), T.int64(32)), "float16", scope="shared.dyn")
+        for threadIdx_z in T.thread_binding(T.int64(2), thread="threadIdx.z"):
+            for threadIdx_y in T.thread_binding(T.int64(2), thread="threadIdx.y"):
+                for threadIdx_x in T.thread_binding(T.int64(32), thread="threadIdx.x"):
+                    for v0 in range(T.int64(4)):
+                        for v1 in T.vectorized(T.int64(8)):
+                            with T.sblock("B_reindex_shared.dyn"):
+                                T.reads(B[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1])
+                                T.writes(B_shared_dyn[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1])
+                                B_shared_dyn[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), T.bitwise_xor(threadIdx_x % T.int64(4), threadIdx_x // T.int64(8)) * T.int64(8) + v1] = B[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1]
+                    for v0 in range(T.int64(2)):
+                        with T.sblock(""):
+                            B_warp = T.sblock_alloc_buffer((T.int64(4), T.int64(1), T.int64(32), T.int64(8)), "float16", scope="warp")
+                            for v1 in range(T.int64(4)):
+                                with T.sblock("B_reindex_shared.dyn_warp_o"):
+                                    T.reads(B_shared_dyn[threadIdx_y * T.int64(64) + v1 * T.int64(16):threadIdx_y * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
+                                    T.writes(B_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
+                                    with T.sblock("B_reindex_shared.dyn_warp_o"):
+                                        T.reads(B_shared_dyn[threadIdx_y * T.int64(64) + v1 * T.int64(16):threadIdx_y * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
+                                        T.writes(B_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
+                                        T.ptx.ldmatrix_legacy("float16", T.bool(False), 4, ".b16", B_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8), T.tvm_access_ptr(T.type_annotation("float16"), B_shared_dyn.data, threadIdx_y * T.int64(2048) + v1 * T.int64(512) + threadIdx_x // T.int64(16) * T.int64(256) + threadIdx_x % T.int64(8) * T.int64(32) + T.bitwise_xor(v0 * T.int64(2) + threadIdx_x % T.int64(16) // T.int64(8), threadIdx_x % T.int64(8) // T.int64(2)) * T.int64(8), T.int64(512), 1), T.int64(0))
+
+    # fmt: on
+    _check_primfunc_transform(before, expected)
+
+
+def test_buffer_c_fp32():
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def before(p_O: T.handle):
+        O = T.match_buffer(p_O, (T.int64(128), T.int64(128)), "float16")
+        O_shared_dyn = T.sblock_alloc_buffer((T.int64(128), T.int64(128)), scope="shared.dyn")
+        O_warp = T.sblock_alloc_buffer((T.int64(4), T.int64(4), T.int64(32), T.int64(8)), scope="warp")
+        for threadIdx_z in T.thread_binding(T.int64(2), thread="threadIdx.z"):
+            for threadIdx_y in T.thread_binding(T.int64(2), thread="threadIdx.y"):
+                for threadIdx_x in T.thread_binding(T.int64(32), thread="threadIdx.x"):
+                    for v0, v1 in T.grid(T.int64(4), T.int64(4)):
+                        with T.sblock("O.dyn_warp_o"):
+                            T.sblock_attr({"permuted_layout": 1})
+                            with T.sblock("O.dyn_warp_o"):
+                                for local_id in range(T.int64(8)):
+                                    O_shared_dyn[threadIdx_z * T.int64(64) + v0 * T.int64(16) + local_id % T.int64(4) // T.int64(2) * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_y * T.int64(64) + v1 * T.int64(16) + local_id // T.int64(4) * T.int64(8) + threadIdx_x % T.int64(4) * T.int64(2) + local_id % T.int64(2)] = O_warp[v0, v1, threadIdx_x, local_id]
+                    for v0 in range(T.int64(16)):
+                        for v1 in T.vectorized(T.int64(8)):
+                            with T.sblock("O.dyn"):
+                                T.sblock_attr({"permuted_layout": 1})
+                                O[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(16) * T.int64(8) + v1] = T.Cast("float16", O_shared_dyn[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(16) * T.int64(8) + v1])
+
+
+    @T.prim_func(s_tir=True)
+    def expected(O: T.Buffer((T.int64(128), T.int64(128)), "float16")):
+        # with T.sblock("root"):
+        O_shared_dyn = T.sblock_alloc_buffer((T.int64(128), T.int64(128)), scope="shared.dyn")
+        O_warp = T.sblock_alloc_buffer((T.int64(4), T.int64(4), T.int64(32), T.int64(8)), scope="warp")
+        for threadIdx_z in T.thread_binding(T.int64(2), thread="threadIdx.z"):
+            for threadIdx_y in T.thread_binding(T.int64(2), thread="threadIdx.y"):
+                for threadIdx_x in T.thread_binding(T.int64(32), thread="threadIdx.x"):
+                    for v0, v1 in T.grid(T.int64(4), T.int64(4)):
+                        with T.sblock("O.dyn_warp_o"):
+                            T.reads(O_warp[v0, v1, threadIdx_x, T.int64(0):T.int64(8)])
+                            T.writes(O_shared_dyn[threadIdx_z * T.int64(64) + v0 * T.int64(16) + threadIdx_x // T.int64(4):threadIdx_z * T.int64(64) + v0 * T.int64(16) + threadIdx_x // T.int64(4) + T.int64(9), threadIdx_y * T.int64(64) + v1 * T.int64(16) + threadIdx_x % T.int64(4) * T.int64(2):threadIdx_y * T.int64(64) + v1 * T.int64(16) + threadIdx_x % T.int64(4) * T.int64(2) + T.int64(10)])
+                            with T.sblock("O.dyn_warp_o"):
+                                T.reads(O_warp[v0, v1, threadIdx_x, T.int64(0):T.int64(8)])
+                                T.writes(O_shared_dyn[threadIdx_z * T.int64(64) + v0 * T.int64(16) + threadIdx_x // T.int64(4):threadIdx_z * T.int64(64) + v0 * T.int64(16) + threadIdx_x // T.int64(4) + T.int64(9), threadIdx_y * T.int64(64) + v1 * T.int64(16) + threadIdx_x % T.int64(4) * T.int64(2):threadIdx_y * T.int64(64) + v1 * T.int64(16) + threadIdx_x % T.int64(4) * T.int64(2) + T.int64(10)])
+                                for local_id in range(T.int64(8)):
+                                    O_shared_dyn[threadIdx_z * T.int64(64) + v0 * T.int64(16) + local_id % T.int64(4) // T.int64(2) * T.int64(8) + threadIdx_x // T.int64(4), T.bitwise_xor(threadIdx_y * T.int64(8) + v1 * T.int64(2) + local_id // T.int64(4), threadIdx_x // T.int64(4)) * T.int64(8) + threadIdx_x % T.int64(4) * T.int64(2) + local_id % T.int64(2)] = O_warp[v0, v1, threadIdx_x, local_id]
+                    for v0 in range(T.int64(16)):
+                        for v1 in T.vectorized(T.int64(8)):
+                            with T.sblock("O.dyn"):
+                                T.reads(O_shared_dyn[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(16) * T.int64(8) + v1])
+                                T.writes(O[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(16) * T.int64(8) + v1])
+                                O[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(16) * T.int64(8) + v1] = T.Cast("float16", O_shared_dyn[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), T.bitwise_xor(threadIdx_x % T.int64(16), threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16)) * T.int64(8) + v1])
+
+    # fmt: on
+    _check_primfunc_transform(before, expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py

@@ -0,0 +1,943 @@
+# 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, F841
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import s_tir
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def test_cp_async_raw_dtype_round_trips():
+    # The raw cp.async form emitted by InjectPTXAsyncCopy carries the element
+    # dtype in Call.dtype and must survive a TVMScript print -> parse round-trip
+    # (it prints dtype-first via tirx.ptx.cp_async_raw). Guards the regression
+    # where the element dtype was dropped after the flat op was phased out.
+    @T.prim_func
+    def f(A: T.Buffer((128,), "float16"), B: T.Buffer((128,), "float16")):
+        T.func_attr({"global_symbol": "f"})
+        for i in T.serial(8):
+            T.ptx.cp_async("float16", B.data, i * 16, A.data, i * 16, 16)
+
+    reparsed = tvm.script.from_source(f.script())
+    tvm.ir.assert_structural_equal(f, reparsed)
+
+
+def count_cp_async(stmt):
+    num_alloc = [0]
+
+    def verify(n):
+        if isinstance(n, tvm.ir.Call) and n.op.name == "tirx.ptx.cp_async_raw":
+            num_alloc[0] += 1
+
+    tvm.tirx.stmt_functor.post_order_visit(stmt, verify)
+    return num_alloc[0]
+
+
+def generate_global_to_shared_vectorized_copy(dtype, vector_size):
+    num_iters = 128 // vector_size
+    vector_size_expr = tvm.runtime.convert(vector_size)
+
+    @T.prim_func(s_tir=True)
+    def ptx_global_to_shared_copy(
+        A: T.Buffer((32, 128), dtype), B: T.Buffer((32, 128), dtype)
+    ) -> None:
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        bx = T.env_thread("blockIdx.x")
+        tx = T.env_thread("threadIdx.x")
+        T.launch_thread(bx, 1)
+        T.launch_thread(tx, 32)
+        with T.sblock():
+            A_shared = T.sblock_alloc_buffer([32, 128], dtype, scope="shared")
+            T.reads(A[0:32, 0:128])
+            T.writes(B[0:32, 0:128])
+
+            T.attr("default", "async_scope", 1)
+            for i in T.serial(num_iters):
+                for j in T.vectorized(vector_size):
+                    A_shared[tx, i * vector_size_expr + j] = A[tx, i * vector_size_expr + j]
+
+            T.evaluate(T.ptx.cp_async.commit_group(dtype=""))
+            T.evaluate(T.ptx.cp_async.wait_group(0, dtype=""))
+
+            for i in range(128):
+                B[tx, i] = A_shared[tx, i]
+
+    return ptx_global_to_shared_copy
+
+
+@T.prim_func(s_tir=True)
+def ptx_global_to_shared_copy_fp32x1(
+    A: T.Buffer((32, 128), "float32"), B: T.Buffer((32, 128), "float32")
+) -> None:
+    T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+    bx = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(bx, 1)
+    T.launch_thread(tx, 32)
+    with T.sblock():
+        A_shared = T.sblock_alloc_buffer([32, 128], "float32", scope="shared")
+        T.reads(A[0:32, 0:128])
+        T.writes(B[0:32, 0:128])
+
+        T.attr("default", "async_scope", 1)
+        for i in T.serial(128):
+            A_shared[tx, i] = A[tx, i]
+
+        T.evaluate(T.ptx.cp_async.commit_group(dtype=""))
+        T.evaluate(T.ptx.cp_async.wait_group(0, dtype=""))
+
+        for i in range(128):
+            B[tx, i] = A_shared[tx, i]
+
+
+@T.prim_func(s_tir=True)
+def ptx_global_to_shared_dyn_copy_fp16x8(
+    A: T.Buffer((32, 128), "float16"),
+    B: T.Buffer((32, 128), "float16"),
+    C: T.Buffer((32, 128), "float16"),
+) -> None:
+    T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+    bx = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(bx, 1)
+    T.launch_thread(tx, 32)
+    with T.sblock():
+        A_shared = T.sblock_alloc_buffer([32, 128], "float16", scope="shared.dyn")
+        B_shared = T.sblock_alloc_buffer([32, 128], "float16", scope="shared.dyn")
+        T.reads(A[0:32, 0:128], B[0:32, 0:128])
+        T.writes(C[0:32, 0:128])
+
+        T.attr("default", "async_scope", 1)
+        for i in T.serial(16):
+            for j in T.vectorized(8):
+                A_shared[tx, i * 8 + j] = A[tx, i * 8 + j]
+                B_shared[tx, i * 8 + j] = B[tx, i * 8 + j]
+
+        T.evaluate(T.ptx.cp_async.commit_group(dtype=""))
+        T.evaluate(T.ptx.cp_async.wait_group(0, dtype=""))
+
+        for i in range(128):
+            C[tx, i] = A_shared[tx, i] + B_shared[tx, i]
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_inject_async_copy():
+    for dtype, vec_size in [("float16", 8), ("float16", 4), ("float32", 4), ("float32", 1)]:
+        if vec_size == 1:
+            f = ptx_global_to_shared_copy_fp32x1
+        else:
+            f = generate_global_to_shared_vectorized_copy(dtype, vec_size)
+
+        mod = tvm.IRModule.from_expr(f)
+        mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+        mod = tvm.tirx.transform.FlattenBuffer()(mod)
+        if vec_size > 1:
+            mod = tvm.tirx.transform.VectorizeLoop()(mod)
+        mod = tvm.s_tir.transform.InjectPTXAsyncCopy()(mod)
+
+        assert count_cp_async(mod["main"].body) == 1
+
+        if not tvm.testing.is_ampere_or_newer():
+            continue
+
+        with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
+            mod = tvm.compile(tvm.IRModule.from_expr(f), target="cuda")
+
+        A_np = np.random.rand(32, 128).astype(dtype)
+        B_np = np.zeros((32, 128)).astype(dtype)
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_nd = tvm.runtime.tensor(A_np, device=dev)
+            B_nd = tvm.runtime.tensor(B_np, device=dev)
+            mod(A_nd, B_nd)
+            tvm.testing.assert_allclose(B_nd.numpy(), A_np)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_inject_async_copy_shared_dyn():
+    f = ptx_global_to_shared_dyn_copy_fp16x8
+
+    mod = tvm.IRModule.from_expr(f)
+    mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+    mod = tvm.tirx.transform.FlattenBuffer()(mod)
+    mod = tvm.tirx.transform.VectorizeLoop()(mod)
+    mod = tvm.s_tir.transform.MergeSharedMemoryAllocations()(mod)
+    mod = tvm.s_tir.transform.InjectPTXAsyncCopy()(mod)
+
+    assert count_cp_async(mod["main"].body) == 2
+
+    if not tvm.testing.is_ampere_or_newer():
+        return
+
+    with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
+        mod = tvm.compile(tvm.IRModule.from_expr(f), target="cuda")
+
+    A_np = np.random.rand(32, 128).astype("float16")
+    B_np = np.random.rand(32, 128).astype("float16")
+    C_np = np.zeros((32, 128)).astype("float16")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_nd = tvm.runtime.tensor(A_np, device=dev)
+        B_nd = tvm.runtime.tensor(B_np, device=dev)
+        C_nd = tvm.runtime.tensor(C_np, device=dev)
+        mod(A_nd, B_nd, C_nd)
+        tvm.testing.assert_allclose(C_nd.numpy(), A_np + B_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# Note: the test_inject_async_copy_barrier case (and its prim_func helper)
+# was removed — it relied on the indexed barrier API
+# (`create_barriers`, `init_barrier_thread_count`, `arrive_barrier`,
+# `wait_barrier`) which fork does not provide; fork uses the
+# `ptx_mbarrier_*` family instead.
+
+
+expected_cuda_script = r"""#include <cuda.h>
+#endif
+
+#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || \
+     (__CUDACC_VER_MAJOR__ > 11))
+#define TVM_ENABLE_L2_PREFETCH 1
+#else
+#define TVM_ENABLE_L2_PREFETCH 0
+#endif
+
+#ifdef _WIN32
+  using uint = unsigned int;
+  using uchar = unsigned char;
+  using ushort = unsigned short;
+  using int64_t = long long;
+  using uint64_t = unsigned long long;
+#else
+  #define uint unsigned int
+  #define uchar unsigned char
+  #define ushort unsigned short
+#endif
+
+__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_0() {
+    asm volatile("cp.async.wait_group 0;");
+}
+
+__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_1() {
+    asm volatile("cp.async.wait_group 1;");
+}
+
+__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_2() {
+    asm volatile("cp.async.wait_group 2;");
+}
+
+__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_5() {
+    asm volatile("cp.async.wait_group 5;");
+}
+
+__forceinline__ __device__ void ptx_cp_async_legacy_pred_ca_4_4_4(void* dst, int dst_off, void* src, int src_off, int predicate) {
+  uint8_t* dst_p = (uint8_t*)dst + dst_off * 4;
+  uint8_t* src_p = (uint8_t*)src + src_off * 4;
+  unsigned int dst_addr = __cvta_generic_to_shared(dst_p);
+  __asm__ __volatile__(
+    "{\n"
+    " .reg .pred p;\n"
+    " setp.eq.u32 p, %3, 1;\n"
+    " @p cp.async.ca.shared.global [%0], [%1], %2;\n"
+    " @!p st.shared.u32 [%0], {%4};\n"
+    "}\n"
+    :: "r"(dst_addr), "l"(src_p), "n"(4), "r"(predicate), "r"(0)
+  );
+}
+
+__forceinline__ __device__ void ptx_cp_async_legacy_ca_4_4_4(void* dst, int dst_off, void* src, int src_off) {
+  uint8_t* dst_p = (uint8_t*)dst + dst_off * 4;
+  uint8_t* src_p = (uint8_t*)src + src_off * 4;
+  unsigned int dst_addr = __cvta_generic_to_shared(dst_p);
+  asm volatile("cp.async.ca.shared.global [%0], [%1], %2;"
+    :: "r"(dst_addr), "l"(src_p), "n"(4));
+}
+
+__forceinline__ __device__ void tvm_builtin_ptx_cp_async_commit_group() {
+    asm volatile("cp.async.commit_group;");
+}
+extern "C" __global__ void __launch_bounds__(16) main_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr, float* __restrict__ C_ptr);
+extern "C" __global__ void __launch_bounds__(16) main_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr, float* __restrict__ C_ptr) {
+  __shared__ alignas(64) float A_shared_ptr[64];
+  __shared__ alignas(64) float B_shared_ptr[64];
+  A_shared_ptr[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/;
+  B_shared_ptr[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/;
+  tvm_builtin_ptx_cp_async_commit_group();
+  int cse_v1 = (((int)threadIdx.x) * 14);
+  int cse_v2 = (((int)threadIdx.x) + 16);
+  ptx_cp_async_legacy_ca_4_4_4(A_shared_ptr, (((int)threadIdx.x) + 16), A_ptr, (((int)threadIdx.x) * 14));
+  ptx_cp_async_legacy_ca_4_4_4(B_shared_ptr, (((int)threadIdx.x) + 16), B_ptr, (((int)threadIdx.x) * 14));
+  tvm_builtin_ptx_cp_async_commit_group();
+  int cse_v3 = (((int)threadIdx.x) + 32);
+  int cse_v6 = ((((int)threadIdx.x) * 14) + 1);
+  ptx_cp_async_legacy_ca_4_4_4(A_shared_ptr, (((int)threadIdx.x) + 32), A_ptr, ((((int)threadIdx.x) * 14) + 1));
+  ptx_cp_async_legacy_ca_4_4_4(B_shared_ptr, (((int)threadIdx.x) + 32), B_ptr, ((((int)threadIdx.x) * 14) + 1));
+  tvm_builtin_ptx_cp_async_commit_group();
+  int cse_v4 = (((int)threadIdx.x) * 16);
+  for (int i = 0; i < 13; ++i) {
+    int cse_v7 = (((((int)threadIdx.x) * 14) + i) + 2);
+    int cse_v9 = ((((i + 3) & 3) * 16) + ((int)threadIdx.x));
+    ptx_cp_async_legacy_pred_ca_4_4_4(A_shared_ptr, ((((i + 3) & 3) * 16) + ((int)threadIdx.x)), A_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12));
+    tvm_builtin_ptx_cp_async_commit_group();
+    tvm_builtin_ptx_cp_async_wait_group_5();
+    __syncthreads();
+    int cse_v8 = (((i & 3) * 16) + ((int)threadIdx.x));
+    C_ptr[((((int)threadIdx.x) * 16) + i)] = (A_shared_ptr[(((i & 3) * 16) + ((int)threadIdx.x))] + B_shared_ptr[(((i & 3) * 16) + ((int)threadIdx.x))]);
+    __syncthreads();
+    ptx_cp_async_legacy_pred_ca_4_4_4(B_shared_ptr, ((((i + 3) & 3) * 16) + ((int)threadIdx.x)), B_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12));
+    tvm_builtin_ptx_cp_async_commit_group();
+  }
+  tvm_builtin_ptx_cp_async_wait_group_2();
+  __syncthreads();
+  C_ptr[((((int)threadIdx.x) * 16) + 13)] = (A_shared_ptr[(((int)threadIdx.x) + 16)] + B_shared_ptr[(((int)threadIdx.x) + 16)]);
+  tvm_builtin_ptx_cp_async_wait_group_1();
+  __syncthreads();
+  C_ptr[((((int)threadIdx.x) * 16) + 14)] = (A_shared_ptr[(((int)threadIdx.x) + 32)] + B_shared_ptr[(((int)threadIdx.x) + 32)]);
+  tvm_builtin_ptx_cp_async_wait_group_0();
+  __syncthreads();
+  int cse_v5 = (((int)threadIdx.x) + 48);
+  C_ptr[((((int)threadIdx.x) * 16) + 15)] = (A_shared_ptr[(((int)threadIdx.x) + 48)] + B_shared_ptr[(((int)threadIdx.x) + 48)]);
+}
+
+"""
+
+
+@pytest.fixture
+def postproc_if_missing_async_support():
+    arch = tvm.support.nvcc.get_target_compute_version()
+    major, _ = tvm.support.nvcc.parse_compute_version(arch)
+    support_async = major >= 8
+
+    func_name = "tvm_callback_cuda_postproc"
+    prev_postproc = tvm.get_global_func(func_name, allow_missing=True)
+
+    # Store the generated code prior to the post-processing.  This
+    # way, even though the generated code doesn't compile on platforms
+    # that do not support async, the comparison against an expected
+    # output can still be performed.  We cannot use
+    # `mod.inspect_source()`, as that contains the source after all
+    # post-processing.
+    original_code = None
+
+    def get_original_code():
+        nonlocal original_code
+        return original_code
+
+    @tvm.register_global_func(func_name, override=True)
+    def tvm_callback_cuda_postproc(code, _):
+        nonlocal original_code
+        original_code = code
+        if support_async:
+            return code
+        else:
+            ret = []
+            for line in code.split("\n"):
+                ret.append(line)
+                ret.append("\n")
+                if line.startswith('extern "C" __global__') and line.endswith("{"):
+                    break
+            ret.append("}")
+            return "".join(ret)
+
+    yield get_original_code
+
+    # Restore previous postproc func to avoid impacting other tests
+    if prev_postproc is None:
+        tvm_ffi.registry.remove_global_func(func_name)
+    else:
+        tvm.register_global_func(func_name, prev_postproc, override=True)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_cp_async_in_if_then_else(postproc_if_missing_async_support):
+    @T.prim_func(s_tir=True)
+    def simple_compute(
+        A: T.Buffer((16, 14), "float32"),
+        B: T.Buffer((16, 14), "float32"),
+        C: T.Buffer((16, 16), "float32"),
+    ):
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+            for i in T.serial(
+                16,
+                annotations={
+                    "software_pipeline_stage": [0, 0, 3],
+                    "software_pipeline_order": [0, 2, 1],
+                    "software_pipeline_async_stages": [0],
+                },
+            ):
+                with T.sblock("compute"):
+                    T.reads(A[tx, i])
+                    T.writes(C[tx, i])
+                    A_shared = T.sblock_alloc_buffer((16, 1), dtype="float32", scope="shared")
+                    B_shared = T.sblock_alloc_buffer((16, 1), dtype="float32", scope="shared")
+                    with T.sblock():
+                        T.reads(A[tx, i])
+                        T.writes(A_shared[tx, 0])
+                        A_shared[tx, 0] = T.if_then_else(
+                            1 <= i and i < 15, A[tx, i - 1], T.float32(0), dtype="float32"
+                        )
+                    with T.sblock():
+                        T.reads(B[tx, i])
+                        T.writes(B_shared[tx, 0])
+                        B_shared[tx, 0] = T.if_then_else(
+                            1 <= i and i < 15, B[tx, i - 1], T.float32(0), dtype="float32"
+                        )
+                    with T.sblock():
+                        T.reads(A_shared[tx, 0], B_shared[tx, 0])
+                        T.writes(C[tx, i])
+                        C[tx, i] = A_shared[tx, 0] + B_shared[tx, 0]
+
+    mod = tvm.IRModule.from_expr(simple_compute)
+    with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
+        tvm.compile(mod, target="cuda")
+    generated_code = postproc_if_missing_async_support()
+    print(generated_code)
+    # Fork emits an NVRTC-aware preamble (`#ifdef __CUDACC_RTC__ ... #else ...`
+    # block) before the apache-style `#include <cuda.h>`; the body after that
+    # block matches the expected snippet, so compare from the kernel-body
+    # onwards instead of byte-for-byte from the start.
+    marker = "#include <cuda.h>"
+    expected_body = expected_cuda_script[expected_cuda_script.index(marker) :]
+    actual_body = generated_code[generated_code.index(marker) :]
+    assert actual_body == expected_body
+
+
+@pytest.mark.skip(
+    reason="This test fails due to an ordering issue with MergeSharedMemoryAllocations "
+    "in device_driver_api.cc. However, fixing this causes failures in MLC. "
+    "This bug should be addressed. See discussion in https://github.com/apache/tvm/pull/16769 "
+    "and https://github.com/apache/tvm/pull/16569#issuecomment-1992720448"
+)
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_vectorize_cp_async_in_if_then_else(postproc_if_missing_async_support):
+    @T.prim_func(s_tir=True)
+    def complex_compute(
+        A: T.Buffer((2, 16, 16, 1280), "float16"),
+        W: T.Buffer((1280, 3, 3, 1280), "float16"),
+        Conv: T.Buffer((512, 1280), "float16"),
+    ):
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        # with T.sblock("root"):
+        data_im2col_reindex_shared_dyn = T.sblock_alloc_buffer(
+            (512, 11520), "float16", scope="shared.dyn"
+        )
+        data_im2col_reindex_shared_dyn_wmma_matrix_a = T.sblock_alloc_buffer(
+            (512, 11520), "float16", scope="wmma.matrix_a"
+        )
+        weight_flatten_reindex_shared_dyn = T.sblock_alloc_buffer(
+            (1280, 11520), "float16", scope="shared.dyn"
+        )
+        weight_flatten_reindex_shared_dyn_wmma_matrix_b = T.sblock_alloc_buffer(
+            (1280, 11520), "float16", scope="wmma.matrix_b"
+        )
+        Conv_reindex_wmma_accumulator = T.sblock_alloc_buffer(
+            (512, 1280), "float16", scope="wmma.accumulator"
+        )
+        for x_0_0 in T.thread_binding(8, thread="blockIdx.y"):
+            for y_0_0 in T.thread_binding(20, thread="blockIdx.x"):
+                for x_0_1 in T.thread_binding(2, thread="threadIdx.y"):
+                    for y_0_1 in T.thread_binding(2, thread="threadIdx.z"):
+                        for x_0_2_init, y_0_2_init in T.grid(2, 2):
+                            with T.sblock("Conv_init_o"):
+                                v_x_o = T.axis.spatial(32, x_0_0 * 4 + x_0_1 * 2 + x_0_2_init)
+                                v_y_o = T.axis.spatial(80, y_0_0 * 4 + y_0_1 * 2 + y_0_2_init)
+                                T.reads()
+                                T.writes(
+                                    Conv_reindex_wmma_accumulator[
+                                        v_x_o * 16 : v_x_o * 16 + 16, v_y_o * 16 : v_y_o * 16 + 16
+                                    ]
+                                )
+                                C_s0 = T.int32()
+                                C_s1 = T.int32()
+                                C = T.match_buffer(
+                                    Conv_reindex_wmma_accumulator[
+                                        v_x_o * 16 : v_x_o * 16 + 16, v_y_o * 16 : v_y_o * 16 + 16
+                                    ],
+                                    (16, 16),
+                                    "float16",
+                                    strides=(C_s0, C_s1),
+                                    scope="wmma.accumulator",
+                                    offset_factor=16,
+                                )
+                                T.tvm_fill_fragment(
+                                    C.data,
+                                    16,
+                                    16,
+                                    16,
+                                    C.elem_offset // C_s0 // 16 * (C_s0 // 16)
+                                    + C.elem_offset % C_s0 // 16,
+                                    T.float32(0),
+                                )
+                        for k_0_0 in T.serial(
+                            180,
+                            annotations={
+                                "software_pipeline_stage": [0, 0, 1],
+                                "software_pipeline_order": [0, 1, 2],
+                                "software_pipeline_async_stages": [0],
+                            },
+                        ):
+                            for ax0_ax1_0_fused_0 in range(4):
+                                for ax0_ax1_0_fused_1 in T.thread_binding(2, thread="threadIdx.z"):
+                                    for ax0_ax1_0_fused_2 in T.thread_binding(
+                                        2, thread="threadIdx.y"
+                                    ):
+                                        for ax0_ax1_0_fused_3 in T.thread_binding(
+                                            32, thread="threadIdx.x"
+                                        ):
+                                            with T.sblock("data_im2col_reindex_shared.dyn_o"):
+                                                v0 = T.axis.spatial(
+                                                    512,
+                                                    x_0_0 * 64
+                                                    + (
+                                                        ax0_ax1_0_fused_0 * 128
+                                                        + ax0_ax1_0_fused_1 * 64
+                                                        + ax0_ax1_0_fused_2 * 32
+                                                        + ax0_ax1_0_fused_3
+                                                    )
+                                                    // 8,
+                                                )
+                                                v1_o = T.axis.spatial(
+                                                    1440,
+                                                    k_0_0 * 8
+                                                    + (
+                                                        ax0_ax1_0_fused_0 * 128
+                                                        + ax0_ax1_0_fused_1 * 64
+                                                        + ax0_ax1_0_fused_2 * 32
+                                                        + ax0_ax1_0_fused_3
+                                                    )
+                                                    % 8,
+                                                )
+                                                T.reads(
+                                                    A[
+                                                        v0 // 256,
+                                                        v1_o // 480 + v0 % 256 // 16 - 1,
+                                                        v1_o % 480 // 160 + v0 % 16 - 1,
+                                                        v1_o % 160 * 8 : v1_o % 160 * 8 + 8,
+                                                    ]
+                                                )
+                                                T.writes(
+                                                    data_im2col_reindex_shared_dyn[
+                                                        v0, v1_o * 8 : v1_o * 8 + 8
+                                                    ]
+                                                )
+                                                for ax1_1 in T.vectorized(8):
+                                                    with T.sblock("data_im2col_reindex_shared.dyn"):
+                                                        v1_i = T.axis.spatial(8, ax1_1)
+                                                        T.reads(
+                                                            A[
+                                                                v0 // 256,
+                                                                v1_o // 480 + v0 % 256 // 16 - 1,
+                                                                v1_o % 480 // 160 + v0 % 16 - 1,
+                                                                v1_o % 160 * 8 + v1_i,
+                                                            ]
+                                                        )
+                                                        T.writes(
+                                                            data_im2col_reindex_shared_dyn[
+                                                                v0, v1_o * 8 + v1_i
+                                                            ]
+                                                        )
+                                                        T.sblock_attr(
+                                                            {"buffer_dim_align": [[0, 0, 32, 8]]}
+                                                        )
+                                                        data_im2col_reindex_shared_dyn[
+                                                            v0, v1_o * 8 + v1_i
+                                                        ] = T.if_then_else(
+                                                            1 <= v1_o // 480 + v0 % 256 // 16
+                                                            and v1_o // 480 + v0 % 256 // 16 < 17
+                                                            and 1 <= v1_o % 480 // 160 + v0 % 16
+                                                            and v1_o % 480 // 160 + v0 % 16 < 17,
+                                                            A[
+                                                                v0 // 256,
+                                                                v1_o // 480 + v0 % 256 // 16 - 1,
+                                                                v1_o % 480 // 160 + v0 % 16 - 1,
+                                                                v1_o % 160 * 8 + v1_i,
+                                                            ],
+                                                            T.float16(0),
+                                                        )
+                            for ax0_ax1_0_fused_0 in range(4):
+                                for ax0_ax1_0_fused_1 in T.thread_binding(2, thread="threadIdx.z"):
+                                    for ax0_ax1_0_fused_2 in T.thread_binding(
+                                        2, thread="threadIdx.y"
+                                    ):
+                                        for ax0_ax1_0_fused_3 in T.thread_binding(
+                                            32, thread="threadIdx.x"
+                                        ):
+                                            for ax1_1 in T.vectorized(8):
+                                                with T.sblock("weight_flatten_reindex_shared.dyn"):
+                                                    v0 = T.axis.spatial(
+                                                        1280,
+                                                        y_0_0 * 64
+                                                        + (
+                                                            ax0_ax1_0_fused_0 * 128
+                                                            + ax0_ax1_0_fused_1 * 64
+                                                            + ax0_ax1_0_fused_2 * 32
+                                                            + ax0_ax1_0_fused_3
+                                                        )
+                                                        // 8,
+                                                    )
+                                                    v1 = T.axis.spatial(
+                                                        11520,
+                                                        k_0_0 * 64
+                                                        + (
+                                                            ax0_ax1_0_fused_0 * 128
+                                                            + ax0_ax1_0_fused_1 * 64
+                                                            + ax0_ax1_0_fused_2 * 32
+                                                            + ax0_ax1_0_fused_3
+                                                        )
+                                                        % 8
+                                                        * 8
+                                                        + ax1_1,
+                                                    )
+                                                    T.reads(
+                                                        W[
+                                                            v0,
+                                                            v1 // 3840,
+                                                            v1 % 3840 // 1280,
+                                                            v1 % 1280,
+                                                        ]
+                                                    )
+                                                    T.writes(
+                                                        weight_flatten_reindex_shared_dyn[v0, v1]
+                                                    )
+                                                    T.sblock_attr(
+                                                        {"buffer_dim_align": [[0, 0, 32, 8]]}
+                                                    )
+                                                    weight_flatten_reindex_shared_dyn[v0, v1] = W[
+                                                        v0,
+                                                        v1 // 1280 // 3,
+                                                        v1 // 1280 % 3,
+                                                        v1 % 1280,
+                                                    ]
+                            for k_0_1 in range(4):
+                                for ax0_0, ax1_0 in T.grid(2, 1):
+                                    with T.sblock("data_im2col_reindex_shared.dyn_wmma.matrix_a_o"):
+                                        v0_o = T.axis.spatial(32, x_0_0 * 4 + x_0_1 * 2 + ax0_0)
+                                        v1_o = T.axis.spatial(720, k_0_0 * 4 + k_0_1 + ax1_0)
+                                        T.reads(
+                                            data_im2col_reindex_shared_dyn[
+                                                v0_o * 16 : v0_o * 16 + 16,
+                                                v1_o * 16 : v1_o * 16 + 16,
+                                            ]
+                                        )
+                                        T.writes(
+                                            data_im2col_reindex_shared_dyn_wmma_matrix_a[
+                                                v0_o * 16 : v0_o * 16 + 16,
+                                                v1_o * 16 : v1_o * 16 + 16,
+                                            ]
+                                        )
+                                        A_s0 = T.int32()
+                                        A_s1 = T.int32()
+                                        A_1 = T.match_buffer(
+                                            data_im2col_reindex_shared_dyn[
+                                                v0_o * 16 : v0_o * 16 + 16,
+                                                v1_o * 16 : v1_o * 16 + 16,
+                                            ],
+                                            (16, 16),
+                                            "float16",
+                                            strides=(A_s0, A_s1),
+                                            scope="shared.dyn",
+                                            offset_factor=16,
+                                        )
+                                        C_s0 = T.int32()
+                                        C_s1 = T.int32()
+                                        C = T.match_buffer(
+                                            data_im2col_reindex_shared_dyn_wmma_matrix_a[
+                                                v0_o * 16 : v0_o * 16 + 16,
+                                                v1_o * 16 : v1_o * 16 + 16,
+                                            ],
+                                            (16, 16),
+                                            "float16",
+                                            strides=(C_s0, C_s1),
+                                            scope="wmma.matrix_a",
+                                            offset_factor=16,
+                                        )
+                                        T.tvm_load_matrix_sync(
+                                            C.data,
+                                            16,
+                                            16,
+                                            16,
+                                            C.elem_offset // C_s0 // 16 * (C_s0 // 16)
+                                            + C.elem_offset % C_s0 // 16,
+                                            T.tvm_access_ptr(
+                                                T.type_annotation("float16"),
+                                                A_1.data,
+                                                A_1.elem_offset,
+                                                A_s0 * 16,
+                                                1,
+                                            ),
+                                            A_s0,
+                                            "row_major",
+                                        )
+                                for ax0_0, ax1_0 in T.grid(2, 1):
+                                    with T.sblock(
+                                        "weight_flatten_reindex_shared.dyn_wmma.matrix_b_o"
+                                    ):
+                                        v0_o = T.axis.spatial(80, y_0_0 * 4 + y_0_1 * 2 + ax0_0)
+                                        v1_o = T.axis.spatial(720, k_0_0 * 4 + k_0_1 + ax1_0)
+                                        T.reads(
+                                            weight_flatten_reindex_shared_dyn[
+                                                v0_o * 16 : v0_o * 16 + 16,
+                                                v1_o * 16 : v1_o * 16 + 16,
+                                            ]
+                                        )
+                                        T.writes(
+                                            weight_flatten_reindex_shared_dyn_wmma_matrix_b[
+                                                v0_o * 16 : v0_o * 16 + 16,
+                                                v1_o * 16 : v1_o * 16 + 16,
+                                            ]
+                                        )
+                                        A_s0 = T.int32()
+                                        A_s1 = T.int32()
+                                        A_1 = T.match_buffer(
+                                            weight_flatten_reindex_shared_dyn[
+                                                v0_o * 16 : v0_o * 16 + 16,
+                                                v1_o * 16 : v1_o * 16 + 16,
+                                            ],
+                                            (16, 16),
+                                            "float16",
+                                            strides=(A_s0, A_s1),
+                                            scope="shared.dyn",
+                                            offset_factor=16,
+                                        )
+                                        C_s0 = T.int32()
+                                        C_s1 = T.int32()
+                                        C = T.match_buffer(
+                                            weight_flatten_reindex_shared_dyn_wmma_matrix_b[
+                                                v0_o * 16 : v0_o * 16 + 16,
+                                                v1_o * 16 : v1_o * 16 + 16,
+                                            ],
+                                            (16, 16),
+                                            "float16",
+                                            strides=(C_s0, C_s1),
+                                            scope="wmma.matrix_b",
+                                            offset_factor=16,
+                                        )
+                                        T.tvm_load_matrix_sync(
+                                            C.data,
+                                            16,
+                                            16,
+                                            16,
+                                            C.elem_offset // C_s0 // 16 * (C_s0 // 16)
+                                            + C.elem_offset % C_s0 // 16,
+                                            T.tvm_access_ptr(
+                                                T.type_annotation("float16"),
+                                                A_1.data,
+                                                A_1.elem_offset,
+                                                A_s0 * 16,
+                                                1,
+                                            ),
+                                            A_s0,
+                                            "col_major",
+                                        )
+                                for x_0_2, y_0_2 in T.grid(2, 2):
+                                    with T.sblock("Conv_update_o"):
+                                        v_x_o = T.axis.spatial(32, x_0_0 * 4 + x_0_1 * 2 + x_0_2)
+                                        v_y_o = T.axis.spatial(80, y_0_0 * 4 + y_0_1 * 2 + y_0_2)
+                                        v_k_o = T.axis.reduce(720, k_0_0 * 4 + k_0_1)
+                                        T.reads(
+                                            Conv_reindex_wmma_accumulator[
+                                                v_x_o * 16 : v_x_o * 16 + 16,
+                                                v_y_o * 16 : v_y_o * 16 + 16,
+                                            ],
+                                            data_im2col_reindex_shared_dyn_wmma_matrix_a[
+                                                v_x_o * 16 : v_x_o * 16 + 16,
+                                                v_k_o * 16 : v_k_o * 16 + 16,
+                                            ],
+                                            weight_flatten_reindex_shared_dyn_wmma_matrix_b[
+                                                v_y_o * 16 : v_y_o * 16 + 16,
+                                                v_k_o * 16 : v_k_o * 16 + 16,
+                                            ],
+                                        )
+                                        T.writes(
+                                            Conv_reindex_wmma_accumulator[
+                                                v_x_o * 16 : v_x_o * 16 + 16,
+                                                v_y_o * 16 : v_y_o * 16 + 16,
+                                            ]
+                                        )
+                                        A_s0 = T.int32()
+                                        A_s1 = T.int32()
+                                        A_1 = T.match_buffer(
+                                            data_im2col_reindex_shared_dyn_wmma_matrix_a[
+                                                v_x_o * 16 : v_x_o * 16 + 16,
+                                                v_k_o * 16 : v_k_o * 16 + 16,
+                                            ],
+                                            (16, 16),
+                                            "float16",
+                                            strides=(A_s0, A_s1),
+                                            scope="wmma.matrix_a",
+                                            offset_factor=16,
+                                        )
+                                        B_s0 = T.int32()
+                                        B_s1 = T.int32()
+                                        B = T.match_buffer(
+                                            weight_flatten_reindex_shared_dyn_wmma_matrix_b[
+                                                v_y_o * 16 : v_y_o * 16 + 16,
+                                                v_k_o * 16 : v_k_o * 16 + 16,
+                                            ],
+                                            (16, 16),
+                                            "float16",
+                                            strides=(B_s0, B_s1),
+                                            scope="wmma.matrix_b",
+                                            offset_factor=16,
+                                        )
+                                        C_s0 = T.int32()
+                                        C_s1 = T.int32()
+                                        C = T.match_buffer(
+                                            Conv_reindex_wmma_accumulator[
+                                                v_x_o * 16 : v_x_o * 16 + 16,
+                                                v_y_o * 16 : v_y_o * 16 + 16,
+                                            ],
+                                            (16, 16),
+                                            "float16",
+                                            strides=(C_s0, C_s1),
+                                            scope="wmma.accumulator",
+                                            offset_factor=16,
+                                        )
+                                        T.tvm_mma_sync(
+                                            C.data,
+                                            C.elem_offset // C_s0 // 16 * (C_s0 // 16)
+                                            + C.elem_offset % C_s0 // 16,
+                                            A_1.data,
+                                            A_1.elem_offset // A_s0 // 16 * (A_s0 // 16)
+                                            + A_1.elem_offset % A_s0 // 16,
+                                            B.data,
+                                            B.elem_offset // B_s0 // 16 * (B_s0 // 16)
+                                            + B.elem_offset % B_s0 // 16,
+                                            C.data,
+                                            C.elem_offset // C_s0 // 16 * (C_s0 // 16)
+                                            + C.elem_offset % C_s0 // 16,
+                                        )
+                        for ax0_0, ax1_0 in T.grid(2, 2):
+                            with T.sblock("Conv_reindex_wmma.accumulator_o"):
+                                v0_o = T.axis.spatial(32, x_0_0 * 4 + x_0_1 * 2 + ax0_0)
+                                v1_o = T.axis.spatial(80, y_0_0 * 4 + y_0_1 * 2 + ax1_0)
+                                T.reads(
+                                    Conv_reindex_wmma_accumulator[
+                                        v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
+                                    ]
+                                )
+                                T.writes(
+                                    Conv[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16]
+                                )
+                                A_s0 = T.int32()
+                                A_s1 = T.int32()
+                                A_1 = T.match_buffer(
+                                    Conv_reindex_wmma_accumulator[
+                                        v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
+                                    ],
+                                    (16, 16),
+                                    "float16",
+                                    strides=(A_s0, A_s1),
+                                    scope="wmma.accumulator",
+                                    offset_factor=16,
+                                )
+                                C_s0 = T.int32()
+                                C_s1 = T.int32()
+                                C = T.match_buffer(
+                                    Conv[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16],
+                                    (16, 16),
+                                    "float16",
+                                    strides=(C_s0, C_s1),
+                                    offset_factor=16,
+                                )
+                                T.tvm_store_matrix_sync(
+                                    A_1.data,
+                                    16,
+                                    16,
+                                    16,
+                                    A_1.elem_offset // A_s0 // 16 * (A_s0 // 16)
+                                    + A_1.elem_offset % A_s0 // 16,
+                                    T.tvm_access_ptr(
+                                        T.type_annotation("float16"),
+                                        C.data,
+                                        C.elem_offset,
+                                        C_s0 * 16,
+                                        2,
+                                    ),
+                                    C_s0,
+                                    "row_major",
+                                )
+
+    mod = tvm.IRModule.from_expr(complex_compute)
+    with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
+        tvm.compile(mod, target="cuda")
+    generated_code = postproc_if_missing_async_support()
+    # generated_code must contain "  setp.ne.b32 p, %0, 0;"
+    assert "setp.ne.b32" in generated_code
+
+
+def test_multiplication_nodes_are_inlined():
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((32, 128), "float16")):
+            tx = T.launch_thread("threadIdx.x", T.int64(32))
+            A_flattened = T.decl_buffer((4096,), "float16", data=A.data)
+            A_shared = T.decl_buffer([4096], "float16", scope="shared")
+
+            T.attr("default", "async_scope", 1)
+            for i in range(16):
+                cse_v1: T.int64 = T.Cast("int64", i)
+                A_shared[T.Ramp(tx * T.int64(128) + cse_v1 * T.int64(8), T.int64(1), 8)] = (
+                    A_flattened[T.Ramp(tx * T.int64(128) + cse_v1 * T.int64(8), T.int64(1), 8)]
+                )
+            T.ptx.cp_async.commit_group()
+            T.ptx.cp_async.wait_group(0)
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((32, 128), "float16")):
+            tx = T.launch_thread("threadIdx.x", T.int64(32))
+            A_flattened = T.decl_buffer((4096,), "float16", data=A.data)
+            A_shared = T.decl_buffer((4096,), "float16", scope="shared")
+            for i in range(16):
+                cse_v1: T.int64 = T.Cast("int64", i)
+                T.ptx.cp_async(
+                    "float16",
+                    A_shared.data,
+                    tx * T.int64(128) + cse_v1 * T.int64(8),
+                    A.data,
+                    tx * T.int64(128) + cse_v1 * T.int64(8),
+                    16,
+                )
+            T.ptx.cp_async.commit_group()
+            T.ptx.cp_async.wait_group(0)
+
+    After = tvm.s_tir.transform.InjectPTXAsyncCopy()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_ldg32.py

@@ -0,0 +1,82 @@
+# 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
+import tvm.testing
+from tvm import s_tir
+from tvm.script import tirx as T
+
+
+def _count_alloc(stmt):
+    num_alloc = [0]
+
+    def visit(n):
+        if isinstance(n, tvm.tirx.AllocBuffer):
+            num_alloc[0] += 1
+
+    tvm.tirx.stmt_functor.post_order_visit(stmt, visit)
+    return num_alloc[0]
+
+
+def _count_ptx_ldg32(stmt):
+    num_call = [0]
+
+    def visit(n):
+        if isinstance(n, tvm.ir.Call) and n.op.name == "tirx.ptx.ldg32":
+            num_call[0] += 1
+
+    tvm.tirx.stmt_functor.post_order_visit(stmt, visit)
+    return num_call[0]
+
+
+@T.prim_func(s_tir=True)
+def where_no_alloc(A: T.Buffer((4,), "float32"), C: T.Buffer((4,), "float32")) -> None:
+    T.func_attr({"global_symbol": "main", "tirx.noalias": True, "target": T.target("cuda")})
+    for i in range(4):
+        C[i] = T.if_then_else(A[i] > T.float32(0), A[i], T.float32(0))
+
+
+@T.prim_func(s_tir=True)
+def where_no_alloc_cpu(A: T.Buffer((4,), "float32"), C: T.Buffer((4,), "float32")) -> None:
+    T.func_attr({"global_symbol": "main", "tirx.noalias": True, "target": T.target("llvm")})
+    for i in range(4):
+        C[i] = T.if_then_else(A[i] > T.float32(0), A[i], T.float32(0))
+
+
+def test_inject_ptx_ldg32_inserts_alloc_for_no_alloc_func():
+    mod = tvm.IRModule.from_expr(where_no_alloc)
+    assert _count_alloc(mod["main"].body) == 0
+
+    mod = tvm.s_tir.transform.InjectPTXLDG32()(mod)
+    assert _count_alloc(mod["main"].body) > 0
+    assert _count_ptx_ldg32(mod["main"].body) == 1
+
+
+def test_inject_ptx_ldg32_skip_non_cuda_target():
+    mod = tvm.IRModule.from_expr(where_no_alloc_cpu)
+    cpu_target = tvm.target.Target("llvm")
+    mod = tvm.IRModule({"main": mod["main"].with_attr("target", cpu_target)})
+    assert _count_alloc(mod["main"].body) == 0
+
+    mod = tvm.s_tir.transform.InjectPTXLDG32()(mod)
+    assert _count_alloc(mod["main"].body) == 0
+    assert _count_ptx_ldg32(mod["main"].body) == 0
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py

@@ -0,0 +1,218 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F841
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_vthread():
+    """Test virtual thread injection with vthread"""
+    n = 100
+    m = 4
+    nthread = 2
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.handle("float32"), C: T.handle("float32")):
+            A_buf = T.decl_buffer((n * nthread,), "float32", data=A)
+            C_buf = T.decl_buffer((n * nthread,), "float32", data=C)
+            for i in range(n):
+                vt_x = T.launch_thread("vthread", nthread)
+                vt_y = T.launch_thread("vthread", nthread)
+                B = T.alloc_buffer((m,), scope="shared")
+                B[i] = A_buf[i * nthread + vt_x]
+                T.evaluate(
+                    T.call_extern(
+                        "int32",
+                        "Run",
+                        B.access_ptr("r"),
+                        T.call_intrin("int32", "tirx.tvm_context_id"),
+                    )
+                )
+                C_buf[i * nthread + vt_x] = B[i] + T.float32(1)
+
+    # For vthread, expected allocation is m * nthread
+    B_expected_alloc = m * nthread
+
+    stmt = tvm.s_tir.transform.InjectVirtualThread()(Module)["main"]
+
+    # Find allocate nodes
+    allocates = []
+
+    def find_allocates(node):
+        if isinstance(node, tvm.tirx.AllocBuffer):
+            allocates.append(node)
+
+    tvm.tirx.stmt_functor.post_order_visit(stmt.body, find_allocates)
+    assert len(allocates) == 1
+    assert list(allocates[0].buffer.shape) == [B_expected_alloc]
+
+
+def test_vthread_extern():
+    """Test virtual thread injection with extern call"""
+    n = 100
+    m = 4
+    nthread = 2
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"global_symbol": "main"})
+            for i in range(n):
+                vt_x = T.launch_thread("vthread", nthread)
+                vt_y = T.launch_thread("vthread", nthread)
+                A = T.alloc_buffer((m,), scope="shared")
+                B = T.alloc_buffer((m,), scope="shared")
+                C = T.alloc_buffer((m,), scope="shared")
+                A[vt_x] = T.Cast("float32", vt_x) + T.float32(1)
+                B[vt_y] = T.Cast("float32", vt_y) + T.float32(1)
+                T.evaluate(
+                    T.call_extern(
+                        "int32",
+                        "Run",
+                        A.access_ptr("r"),
+                        B.access_ptr("r"),
+                        C.access_ptr("rw"),
+                    )
+                )
+
+    # For vthread:
+    # A, B expected allocation is m * nthread (used with single vthread each)
+    A_expected_alloc = m * nthread
+    # C expected allocation is m * nthread * nthread (used in extern with both vthreads)
+    C_expected_alloc = m * nthread * nthread
+
+    stmt = tvm.s_tir.transform.InjectVirtualThread()(Module)["main"]
+
+    # Find allocate nodes
+    allocates = []
+
+    def find_allocates(node):
+        if isinstance(node, tvm.tirx.AllocBuffer):
+            allocates.append(node)
+
+    tvm.tirx.stmt_functor.post_order_visit(stmt.body, find_allocates)
+    assert len(allocates) == 3
+    # Check that we have the expected extents (order may vary)
+    extents = sorted([int(a.buffer.shape[0]) for a in allocates])
+    assert extents == sorted([A_expected_alloc, A_expected_alloc, C_expected_alloc])
+
+
+def test_vthread_if_then_else():
+    """Test virtual thread injection with if-then-else"""
+    nthread = 2
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.handle("float32")):
+            T.func_attr({"global_symbol": "main"})
+            A_buf = T.decl_buffer((100 * nthread,), "float32", data=A)
+            for i in range(100):
+                vt = T.launch_thread("vthread", nthread)
+                B = T.alloc_buffer((128,), scope="shared")
+                if i == 0:
+                    B[i] = A_buf[i * nthread + vt]
+                else:
+                    B[i] = A_buf[i * nthread + vt] + T.float32(1)
+                if i == 0:
+                    B[i] = A_buf[i * nthread + vt] + T.float32(2)
+
+    stmt = tvm.s_tir.transform.InjectVirtualThread()(Module)["main"]
+
+    # Find IfThenElse nodes
+    if_nodes = []
+
+    def find_ifs(node):
+        if isinstance(node, tvm.tirx.IfThenElse):
+            if_nodes.append(node)
+
+    tvm.tirx.stmt_functor.post_order_visit(stmt.body, find_ifs)
+
+    assert len(if_nodes) == 2
+    # First if has else_case, second does not
+    assert if_nodes[0].else_case is not None
+    assert if_nodes[1].else_case is None
+
+
+def test_vthread_simplified():
+    """Indices resulting from vthread injection should simplified
+
+    This ensures that downstream passes that check for Ramp nodes do
+    not need to each simplify the indices.
+    """
+
+    @T.prim_func(s_tir=True)
+    def before_func():
+        vthread = T.env_thread("vthread")
+        T.launch_thread(vthread, 4)
+        B = T.alloc_buffer((4,), "int32", scope="shared")
+        B[0:4] = T.broadcast(vthread, 4)
+
+    @T.prim_func(check_well_formed=False, s_tir=True)
+    def expected_func():
+        B = T.alloc_buffer((16,), "int32", scope="shared")
+        B_1 = T.Buffer([16], "int32", data=B.data, scope="shared")
+        # The indices for B should each be a single Ramp node, and
+        # should not be the sum of a Ramp and Broadcast node.
+        B_1[T.Mul(0, 4) : T.Mul(0, 4) + 4] = T.broadcast(0, 4)
+        B_1[T.Mul(1, 4) : T.Mul(1, 4) + 4] = T.broadcast(1, 4)
+        B_1[T.Mul(2, 4) : T.Mul(2, 4) + 4] = T.broadcast(2, 4)
+        B_1[T.Mul(3, 4) : T.Mul(3, 4) + 4] = T.broadcast(3, 4)
+
+    before_mod = tvm.IRModule.from_expr(before_func.with_attr("global_symbol", "main"))
+    after_mod = tvm.s_tir.transform.InjectVirtualThread()(before_mod)
+    after_func = after_mod["main"]
+
+    tvm.ir.assert_structural_equal(after_func, expected_func.with_attr("global_symbol", "main"))
+
+
+def test_vthread_vectorized():
+    """Use of vthread is compatible with vector allocations"""
+
+    @T.prim_func(s_tir=True)
+    def before_func():
+        vthread = T.env_thread("vthread")
+        T.launch_thread(vthread, 4)
+        B = T.alloc_buffer((4,), "int32", scope="shared")
+        B[0:4] = T.broadcast(vthread, 4)
+
+    before_mod = tvm.IRModule.from_expr(before_func.with_attr("global_symbol", "main"))
+    intermediate_mod = tvm.s_tir.transform.InjectVirtualThread()(before_mod)
+    after_mod = tvm.tirx.transform.StorageRewrite()(intermediate_mod)
+    after_func = after_mod["main"]
+
+    # Verify the vectorized allocation has the expected shape and dtype
+    allocate_node = None
+
+    def visitor(op):
+        nonlocal allocate_node
+        if isinstance(op, tvm.tirx.AllocBuffer) and "shared" in str(op.buffer.data.ty):
+            allocate_node = op
+
+    tvm.tirx.stmt_functor.post_order_visit(after_func.body, visitor)
+    assert allocate_node is not None
+    assert list(allocate_node.buffer.shape) == [4]
+    assert allocate_node.buffer.dtype == "int32x4"
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_lift_thread_binding.py

@@ -0,0 +1,140 @@
+# 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
+import tvm
+from tvm import s_tir, tirx
+from tvm.script import tirx as T
+
+
+def test_lift_tx_beyond_local():
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle, c: T.handle):
+        n = T.int32()
+        A = T.match_buffer(a, (32, 1, 128))
+        B = T.match_buffer(b, (32, n, 128))
+        C = T.match_buffer(c, (32, 1, n))
+        for ax0_ax1_fused in T.thread_binding(n * 32, thread="blockIdx.x"):
+            with T.sblock(""):
+                T.reads(A[ax0_ax1_fused // n, 0, 0:256], B[ax0_ax1_fused // n, ax0_ax1_fused % n, 0:256])
+                T.writes(C[ax0_ax1_fused // n, 0, ax0_ax1_fused % n])
+                D_local = T.sblock_alloc_buffer((32, 1, n), scope="local")
+                D_rf_local = T.sblock_alloc_buffer((256, 32, 1, n), scope="local")
+                for ax2_fused_1 in T.thread_binding(256, thread="threadIdx.x"):
+                    with T.sblock("NT_matmul_rf_init"):
+                        T.reads()
+                        T.writes(D_rf_local[ax2_fused_1, ax0_ax1_fused // n, 0, ax0_ax1_fused % n])
+                        D_rf_local[ax2_fused_1, ax0_ax1_fused // n, 0, ax0_ax1_fused % n] = T.float32(0)
+                    for ax2_fused_0 in range(1):
+                        with T.sblock("NT_matmul_rf_update"):
+                            T.where(ax2_fused_0 * 256 + ax2_fused_1 < 128)
+                            T.reads(D_rf_local[ax2_fused_1, ax0_ax1_fused // n, 0, ax0_ax1_fused % n], A[ax0_ax1_fused // n, 0, ax2_fused_0 * 256 + ax2_fused_1], B[ax0_ax1_fused // n, ax0_ax1_fused % n, ax2_fused_0 * 256 + ax2_fused_1])
+                            T.writes(D_rf_local[ax2_fused_1, ax0_ax1_fused // n, 0, ax0_ax1_fused % n])
+                            D_rf_local[ax2_fused_1, ax0_ax1_fused // n, 0, ax0_ax1_fused % n] = D_rf_local[ax2_fused_1, ax0_ax1_fused // n, 0, ax0_ax1_fused % n] + A[ax0_ax1_fused // n, 0, ax2_fused_0 * 256 + ax2_fused_1] * B[ax0_ax1_fused // n, ax0_ax1_fused % n, ax2_fused_0 * 256 + ax2_fused_1]
+                for ax1_ax2_fused in range(1):
+                    for ax0_fused in T.thread_binding(256, thread="threadIdx.x"):
+                        with T.sblock(""):
+                            T.reads(D_rf_local[ax0_fused, ax0_ax1_fused // n, 0, ax0_ax1_fused % n])
+                            T.writes(D_local[ax0_ax1_fused // n, 0, ax0_ax1_fused % n])
+                            cross_thread_D_local = T.sblock_alloc_buffer((1,), strides=(1,), scope="local")
+                            in_thread_D_local = T.sblock_alloc_buffer((1,), strides=(1,), scope="local")
+                            with T.sblock("NT_matmul_in_thread_init"):
+                                T.reads()
+                                T.writes(in_thread_D_local[0])
+                                in_thread_D_local[0] = T.float32(0)
+                            with T.sblock("NT_matmul_in_thread"):
+                                T.where(0 <= ax0_ax1_fused // n and ax0_ax1_fused // n < 32 and 0 <= ax0_ax1_fused % n and ax0_ax1_fused % n < n)
+                                T.reads(D_rf_local[ax0_fused, ax0_ax1_fused // n, 0, ax0_ax1_fused % n])
+                                T.writes(in_thread_D_local[0])
+                                in_thread_D_local[0] = in_thread_D_local[0] + D_rf_local[ax0_fused, ax0_ax1_fused // n, 0, ax0_ax1_fused % n]
+                            with T.sblock("NT_matmul_cross_thread"):
+                                T.reads(in_thread_D_local[0])
+                                T.writes(cross_thread_D_local[0])
+                                T.attr(T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]), "reduce_scope", T.int32(0))
+                                T.tvm_thread_allreduce(T.uint32(1), in_thread_D_local[0], T.bool(True), cross_thread_D_local[0], ax0_fused)
+                            with T.sblock("NT_matmul_write_back"):
+                                T.where(ax0_fused == 0)
+                                T.reads(cross_thread_D_local[0])
+                                T.writes(D_local[ax0_ax1_fused // n, 0, ax0_ax1_fused % n])
+                                D_local[ax0_ax1_fused // n, 0, ax0_ax1_fused % n] = cross_thread_D_local[0]
+                with T.sblock("T_divide"):
+                    T.where(0 <= ax0_ax1_fused // n and ax0_ax1_fused // n < 32 and 0 <= ax0_ax1_fused % n and ax0_ax1_fused % n < n)
+                    T.reads(D_local[ax0_ax1_fused // n, 0, ax0_ax1_fused % n])
+                    T.writes(C[ax0_ax1_fused // n, 0, ax0_ax1_fused % n])
+                    C[ax0_ax1_fused // n, 0, ax0_ax1_fused % n] = D_local[ax0_ax1_fused // n, 0, ax0_ax1_fused % n] * T.float32(0.088397790055248615)
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((32, 1, 128), "float32"), b: T.handle, c: T.handle):
+        n = T.int32()
+        B = T.match_buffer(b, (32, n, 128))
+        C = T.match_buffer(c, (32, 1, n))
+        # with T.sblock("root"):
+        for blockIdx_x in T.thread_binding(n * 32, thread="blockIdx.x"):
+            for threadIdx_x in T.thread_binding(256, thread="threadIdx.x"):
+                with T.sblock(""):
+                    T.reads(A[blockIdx_x // n, 0, 0:256], B[blockIdx_x // n, blockIdx_x % n, 0:256])
+                    T.writes(C[blockIdx_x // n, 0, blockIdx_x % n])
+                    D_local = T.sblock_alloc_buffer((32, 1, n), scope="local")
+                    D_rf_local = T.sblock_alloc_buffer((256, 32, 1, n), scope="local")
+                    with T.sblock("NT_matmul_rf_init"):
+                        T.reads()
+                        T.writes(D_rf_local[threadIdx_x, blockIdx_x // n, 0, blockIdx_x % n])
+                        D_rf_local[threadIdx_x, blockIdx_x // n, 0, blockIdx_x % n] = T.float32(0)
+                    for ax2_fused_0 in range(1):
+                        with T.sblock("NT_matmul_rf_update"):
+                            T.where(ax2_fused_0 * 256 + threadIdx_x < 128)
+                            T.reads(D_rf_local[threadIdx_x, blockIdx_x // n, 0, blockIdx_x % n], A[blockIdx_x // n, 0, ax2_fused_0 * 256 + threadIdx_x], B[blockIdx_x // n, blockIdx_x % n, ax2_fused_0 * 256 + threadIdx_x])
+                            T.writes(D_rf_local[threadIdx_x, blockIdx_x // n, 0, blockIdx_x % n])
+                            D_rf_local[threadIdx_x, blockIdx_x // n, 0, blockIdx_x % n] = D_rf_local[threadIdx_x, blockIdx_x // n, 0, blockIdx_x % n] + A[blockIdx_x // n, 0, ax2_fused_0 * 256 + threadIdx_x] * B[blockIdx_x // n, blockIdx_x % n, ax2_fused_0 * 256 + threadIdx_x]
+                    for ax1_ax2_fused in range(1):
+                        with T.sblock(""):
+                            T.reads(D_rf_local[threadIdx_x, blockIdx_x // n, 0, blockIdx_x % n])
+                            T.writes(D_local[blockIdx_x // n, 0, blockIdx_x % n])
+                            cross_thread_D_local = T.sblock_alloc_buffer((1,), strides=(1,), scope="local")
+                            in_thread_D_local = T.sblock_alloc_buffer((1,), strides=(1,), scope="local")
+                            with T.sblock("NT_matmul_in_thread_init"):
+                                T.reads()
+                                T.writes(in_thread_D_local[0])
+                                in_thread_D_local[0] = T.float32(0)
+                            with T.sblock("NT_matmul_in_thread"):
+                                T.where(0 <= blockIdx_x // n and blockIdx_x // n < 32 and 0 <= blockIdx_x % n and blockIdx_x % n < n)
+                                T.reads(D_rf_local[threadIdx_x, blockIdx_x // n, 0, blockIdx_x % n])
+                                T.writes(in_thread_D_local[0])
+                                in_thread_D_local[0] = in_thread_D_local[0] + D_rf_local[threadIdx_x, blockIdx_x // n, 0, blockIdx_x % n]
+                            with T.sblock("NT_matmul_cross_thread"):
+                                T.reads(in_thread_D_local[0])
+                                T.writes(cross_thread_D_local[0])
+                                T.attr(T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]), "reduce_scope", T.int32(0))
+                                T.tvm_thread_allreduce(T.uint32(1), in_thread_D_local[0], T.bool(True), cross_thread_D_local[0], threadIdx_x)
+                            with T.sblock("NT_matmul_write_back"):
+                                T.where(threadIdx_x == 0)
+                                T.reads(cross_thread_D_local[0])
+                                T.writes(D_local[blockIdx_x // n, 0, blockIdx_x % n])
+                                D_local[blockIdx_x // n, 0, blockIdx_x % n] = cross_thread_D_local[0]
+                    with T.sblock("T_divide"):
+                        T.where(0 <= blockIdx_x // n and blockIdx_x // n < 32 and 0 <= blockIdx_x % n and blockIdx_x % n < n)
+                        T.reads(D_local[blockIdx_x // n, 0, blockIdx_x % n])
+                        T.writes(C[blockIdx_x // n, 0, blockIdx_x % n])
+                        C[blockIdx_x // n, 0, blockIdx_x % n] = D_local[blockIdx_x // n, 0, blockIdx_x % n] * T.float32(0.088397790055248615)
+    # fmt: on
+    mod = tvm.IRModule({"main": before.with_attr("global_symbol", "main")})
+    after = s_tir.transform.LiftThreadBinding()(mod)
+    tvm.ir.assert_structural_equal(expected.with_attr("global_symbol", "main"), after["main"])
+
+
+if __name__ == "__main__":
+    test_lift_tx_beyond_local()

tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py

@@ -0,0 +1,802 @@
+# 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 numpy
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.ir.module import IRModule
+from tvm.script import tirx as T
+
+
+def collect_visit(stmt, f):
+    ret = []
+    tvm.tirx.stmt_functor.post_order_visit(stmt, lambda x: ret.append(f(x)))
+    return ret
+
+
+def test_multi_loop():
+    @T.prim_func(s_tir=True)
+    def func(n: T.int64, m: T.int64):
+        for i in range(4):
+            for j in T.serial(n):
+                for k in T.serial(m):
+                    if T.likely(i * m + j + k < n):
+                        T.evaluate(m)
+                    else:
+                        T.evaluate(n)
+
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.LoopPartition()(mod)
+    stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
+
+    assert not any(collect_visit(stmt.body[0], lambda x: isinstance(x, tvm.tirx.IfThenElse)))
+
+
+def test_multi_if():
+    @T.prim_func(s_tir=True)
+    def func(n: T.int64, m: T.int64):
+        for i in range(4):
+            for j in T.serial(n):
+                for k in T.serial(m):
+                    if T.likely(i * m + j + k < n):
+                        T.evaluate(m)
+                    else:
+                        T.evaluate(n)
+                    if T.likely(i * m + j - k < n):
+                        T.evaluate(m)
+                    else:
+                        T.evaluate(n)
+
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.LoopPartition()(mod)
+    stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
+
+    assert not any(collect_visit(stmt.body[0], lambda x: isinstance(x, tvm.tirx.IfThenElse)))
+
+
+def test_condition():
+    @T.prim_func(s_tir=True)
+    def func(m: T.int64, n: T.int64):
+        for i in T.serial(T.truncdiv(n + 3, 4)):
+            for j in range(4):
+                T.evaluate(T.Select(T.likely(i * 4 + j < n), m, n))
+
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.LoopPartition()(mod)
+    stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
+
+    assert not any(collect_visit(stmt[0], lambda x: isinstance(x, tvm.tirx.Select)))
+
+
+def test_condition_EQ():
+    @T.prim_func(s_tir=True)
+    def func(m: T.int64, n: T.int64):
+        for i in range(10):
+            T.evaluate(T.Select(T.likely(i == 5), m, n))
+
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    with tvm.transform.PassContext(config={"s_tir.LoopPartition": {"partition_const_loop": True}}):
+        mod = tvm.s_tir.transform.LoopPartition()(mod)
+        stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
+
+    assert not any(collect_visit(stmt[0], lambda x: isinstance(x, tvm.tirx.Select)))
+
+
+def test_everything_during_deduction():
+    @T.prim_func(s_tir=True)
+    def func(m: T.int64, n: T.int64):
+        for i in T.serial(n):
+            for j in range(32):
+                if T.likely(T.truncdiv(i, j) < m):
+                    # this guard will produce everything during deduction
+                    T.evaluate(m)
+
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.LoopPartition()(mod)
+    stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
+
+    assert isinstance(stmt.body.body, tvm.tirx.IfThenElse)
+
+
+def test_oneD_pool():
+    @T.prim_func(s_tir=True)
+    def func(m: T.int64, data: T.handle("float32"), out: T.handle("float32")):
+        data_ptr = T.decl_buffer((16,), "float32", data=data)
+        out_ptr = T.decl_buffer((16,), "float32", data=out)
+        for ow in range(16):
+            for kw in range(3):
+                if T.likely(ow > 0):
+                    if T.likely(ow < 15):
+                        out_ptr[ow] = T.max(out_ptr[ow], data_ptr[ow + kw - 1])
+        for ow in range(16):
+            for kw in range(3):
+                if T.likely(ow < 1):
+                    if T.likely(kw > 0):
+                        out_ptr[ow] = T.max(out_ptr[ow], data_ptr[ow + kw - 1])
+        for ow in range(16):
+            for kw in range(3):
+                if T.likely(ow > 14):
+                    if T.likely(kw < 2):
+                        out_ptr[ow] = T.max(out_ptr[ow], data_ptr[ow + kw - 1])
+
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+
+    with tvm.transform.PassContext(config={"s_tir.LoopPartition": {"partition_const_loop": True}}):
+        mod = tvm.s_tir.transform.LoopPartition()(mod)
+        stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
+
+    assert not any(collect_visit(stmt, lambda x: isinstance(x, tvm.tirx.IfThenElse)))
+
+
+def test_cce_loop_1():
+    n = 514
+    m = 514
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((n * m,), "float16"), B: T.Buffer((n * m,), "float16")):
+        for i in range(11):
+            for j in range(160):
+                if T.likely(i * 160 + j < 1600):
+                    A[(i + 1) * m + j + 1] = (
+                        B[i * m + j + 1] + B[(i + 1) * m + j + 1] + B[(i + 2) * m + j + 1]
+                    )
+
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    with tvm.transform.PassContext(config={"s_tir.LoopPartition": {"partition_const_loop": True}}):
+        mod = tvm.s_tir.transform.LoopPartition()(mod)
+        stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
+
+    assert not any(collect_visit(stmt, lambda x: isinstance(x, tvm.tirx.IfThenElse)))
+
+
+def test_cce_loop_2():
+    length = 112
+    tile = 32
+    loop = (length + tile - 1) // tile
+
+    @T.prim_func(private=True, s_tir=True)
+    def func():
+        for i in range(loop):
+            if T.likely(i * tile + tile > length):
+                T.evaluate(T.call_extern("float32", "cce_intrisic", i * tile, length))
+            else:
+                T.evaluate(T.call_extern("float32", "cce_intrisic", i * tile, i * tile + tile))
+
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    with tvm.transform.PassContext(config={"s_tir.LoopPartition": {"partition_const_loop": True}}):
+        mod = tvm.s_tir.transform.LoopPartition()(mod)
+        stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
+
+    assert not any(collect_visit(stmt, lambda x: isinstance(x, tvm.tirx.IfThenElse)))
+
+
+def test_cce_loop_3():
+    loop1 = 4
+    loop2 = 9998
+    tile = 39991
+
+    @T.prim_func(private=True, s_tir=True)
+    def func():
+        for i in range(loop2):
+            for j in range(loop1):
+                if T.likely(i * loop1 + j < tile):
+                    T.evaluate(T.call_extern("float16", "cce_intrisic", i))
+
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+
+    with tvm.transform.PassContext(config={"s_tir.LoopPartition": {"partition_const_loop": True}}):
+        mod = tvm.s_tir.transform.LoopPartition()(mod)
+        stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
+
+    assert not any(collect_visit(stmt, lambda x: isinstance(x, tvm.tirx.IfThenElse)))
+
+
+@T.prim_func(s_tir=True)
+def partitioned_concat(
+    A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32"), C: T.Buffer((32,), "float32")
+) -> None:
+    T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+    for i in T.serial(0, 16):
+        C[i] = A[i]
+    for i in T.serial(0, 16):
+        C[i + 16] = B[i + 16]
+
+
+def partition_from_scheduled_tir(prim_func, pass_cfg, do_flatten=True):
+    with tvm.transform.PassContext(config=pass_cfg):
+        mod = IRModule.from_expr(prim_func.with_attr("global_symbol", "main"))
+        mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+        if do_flatten:
+            mod = tvm.tirx.transform.FlattenBuffer()(mod)
+        mod = tvm.s_tir.transform.LoopPartition()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        mod = tvm.tirx.transform.RemoveNoOp()(mod)
+        return mod
+
+
+@T.prim_func(s_tir=True)
+def partitioned_concat_3(
+    placeholder: T.Buffer((1, 64, 28, 28), "int8"),
+    placeholder_1: T.Buffer((1, 32, 28, 28), "int8"),
+    placeholder_2: T.Buffer((1, 32, 28, 28), "int8"),
+    T_concat: T.Buffer((1, 128, 28, 28), "int8"),
+) -> None:
+    placeholder_flat = T.decl_buffer([50176], "int8", data=placeholder.data)
+    placeholder_1_flat = T.decl_buffer([25088], "int8", data=placeholder_1.data)
+    placeholder_2_flat = T.decl_buffer([25088], "int8", data=placeholder_2.data)
+    T_concat_flat = T.decl_buffer([100352], "int8", data=T_concat.data)
+    for i1, i2, i3 in T.grid(64, 28, 28):
+        T_concat_flat[i1 * 784 + i2 * 28 + i3] = placeholder_flat[i1 * 784 + i2 * 28 + i3]
+    for i1, i2, i3 in T.grid(32, 28, 28):
+        T_concat_flat[i1 * 784 + i2 * 28 + i3 + 50176] = placeholder_1_flat[i1 * 784 + i2 * 28 + i3]
+    for i1, i2, i3 in T.grid(32, 28, 28):
+        T_concat_flat[i1 * 784 + i2 * 28 + i3 + 75264] = placeholder_2_flat[i1 * 784 + i2 * 28 + i3]
+
+
+@T.prim_func(s_tir=True)
+def concat_func_3(
+    placeholder: T.Buffer((1, 64, 28, 28), "int8"),
+    placeholder_1: T.Buffer((1, 32, 28, 28), "int8"),
+    placeholder_2: T.Buffer((1, 32, 28, 28), "int8"),
+    T_concat: T.Buffer((1, 128, 28, 28), "int8"),
+) -> None:
+    placeholder_flat = T.decl_buffer([50176], "int8", data=placeholder.data)
+    placeholder_1_flat = T.decl_buffer([25088], "int8", data=placeholder_1.data)
+    placeholder_2_flat = T.decl_buffer([25088], "int8", data=placeholder_2.data)
+    T_concat_flat = T.decl_buffer([100352], "int8", data=T_concat.data)
+    for i1 in T.serial(128, annotations={"pragma_loop_partition_hint": 1}):
+        for i2, i3 in T.grid(28, 28):
+            if 96 <= i1:
+                T_concat_flat[i1 * 784 + i2 * 28 + i3] = placeholder_2_flat[
+                    i1 * 784 + i2 * 28 + i3 - 75264
+                ]
+            if 64 <= i1 and i1 < 96:
+                T_concat_flat[i1 * 784 + i2 * 28 + i3] = placeholder_1_flat[
+                    i1 * 784 + i2 * 28 + i3 - 50176
+                ]
+            if i1 < 64:
+                T_concat_flat[i1 * 784 + i2 * 28 + i3] = placeholder_flat[i1 * 784 + i2 * 28 + i3]
+
+
+def test_condition_mutually_exclusive():
+    mod = partition_from_scheduled_tir(
+        concat_func_3, {"s_tir.LoopPartition": {"partition_const_loop": True}}
+    )
+    tvm.ir.assert_structural_equal(
+        mod["main"], partitioned_concat_3.with_attr("global_symbol", "main")
+    )
+
+
+def test_loop_partition_unroll_hint():
+    @T.prim_func(s_tir=True)
+    def main(
+        A_arg: T.Buffer((1, 3, 224, 224), "int8"), B_arg: T.Buffer((1, 224, 7, 16), "int8")
+    ) -> None:
+        A = T.decl_buffer(150528, "int8", data=A_arg.data)
+        B = T.decl_buffer(25088, "int8", data=B_arg.data)
+        for ax0 in T.serial(
+            112,
+            annotations={"pragma_loop_partition_hint": True},
+        ):
+            for ax1, ax2, ax3 in T.grid(224, 7, 16):
+                if 3 <= ax0 * 2 + ax2 and ax0 * 2 + ax2 < 227 and ax3 < 3:
+                    B[ax1 * 112 + ax2 * 16 + ax3] = A[ax3 * 50176 + ax1 * 224 + ax0 * 2 + ax2 - 3]
+
+    @T.prim_func(s_tir=True)
+    def partitioned_main(
+        A_arg: T.Buffer((1, 3, 224, 224), "int8"), B_arg: T.Buffer((1, 224, 7, 16), "int8")
+    ) -> None:
+        A = T.decl_buffer(150528, dtype="int8", data=A_arg.data)
+        B = T.decl_buffer(25088, dtype="int8", data=B_arg.data)
+        # body
+        for ax1, ax2, ax3 in T.grid(224, 7, 16):
+            if 3 <= ax2 and ax3 < 3:
+                B[ax1 * 112 + ax2 * 16 + ax3] = A[ax3 * 50176 + ax1 * 224 + ax2 - 3]
+        for ax1, ax2, ax3 in T.grid(224, 7, 16):
+            if 1 <= ax2 and ax3 < 3:
+                B[ax1 * 112 + ax2 * 16 + ax3] = A[ax3 * 50176 + ax1 * 224 + ax2 - 1]
+        for ax0, ax1, ax2, ax3 in T.grid(109, 224, 7, 16):
+            if ax3 < 3:
+                B[ax1 * 112 + ax2 * 16 + ax3] = A[ax3 * 50176 + ax1 * 224 + ax0 * 2 + ax2 + 1]
+        for ax1, ax2, ax3 in T.grid(224, 7, 16):
+            if ax2 < 5 and ax3 < 3:
+                B[ax1 * 112 + ax2 * 16 + ax3] = A[ax3 * 50176 + ax1 * 224 + ax2 + 219]
+
+    mod = partition_from_scheduled_tir(
+        main,
+        {
+            "s_tir.LoopPartition": {
+                "partition_const_loop": True,
+                "unroll_loop_with_partition_hint_no_interval": True,
+            }
+        },
+    )
+    mod = tvm.tirx.transform.UnrollLoop()(mod)
+    mod = tvm.tirx.transform.RemoveNoOp()(mod)
+    mod = tvm.tirx.transform.StmtSimplify()(mod)
+    tvm.ir.assert_structural_equal(mod["main"], partitioned_main.with_attr("global_symbol", "main"))
+
+
+def test_loop_partition_recursive_unroll_hint():
+    @T.prim_func(s_tir=True)
+    def main():
+        placeholder_0_dm = T.decl_buffer([1, 32, 32, 16], dtype="int8")
+        for i3_0 in T.serial(5, annotations={"pragma_loop_partition_hint": 1}):
+            for i2_0 in T.serial(2, annotations={"pragma_loop_partition_hint": 1}):
+                pad_temp = T.decl_buffer([1, 16, 16, 16], dtype="int8")
+                for ax0, ax1, ax2 in T.grid(16, 16, 16):
+                    if (
+                        6 <= i2_0 * 4 + ax0
+                        and i2_0 * 4 + ax0 < 26
+                        and 6 <= i3_0 * 4 + ax1
+                        and i3_0 * 4 + ax1 < 26
+                    ):
+                        pad_temp[
+                            0,
+                            i2_0 * 4 + ax0 - 6 + 6 - i2_0 * 4,
+                            i3_0 * 4 + ax1 - 6 + 6 - i3_0 * 4,
+                            ax2,
+                        ] = placeholder_0_dm[
+                            0,
+                            i2_0 * 4 + ax0 - 6 - -6,
+                            i3_0 * 4 + ax1 - 6 - -6,
+                            ax2,
+                        ]
+
+    @T.prim_func(s_tir=True)
+    def partitioned_main():
+        placeholder_0_dm = T.decl_buffer((16384,), "int8")
+        for i3_0 in T.unroll(2):
+            for i2_0 in T.unroll(2):
+                pad_temp = T.decl_buffer((4096,), "int8")
+                for ax0, ax1, ax2 in T.grid(16, 16, 16):
+                    if 6 <= i2_0 * 4 + ax0 and 6 <= i3_0 * 4 + ax1:
+                        pad_temp[ax0 * 256 + ax1 * 16 + ax2] = placeholder_0_dm[
+                            i2_0 * 2048 + ax0 * 512 + i3_0 * 64 + ax1 * 16 + ax2
+                        ]
+        for i2_0 in T.unroll(2):
+            pad_temp = T.decl_buffer((4096,), "int8")
+            for ax0, ax1, ax2 in T.grid(16, 16, 16):
+                if 6 <= i2_0 * 4 + ax0:
+                    pad_temp[ax0 * 256 + ax1 * 16 + ax2] = placeholder_0_dm[
+                        i2_0 * 2048 + ax0 * 512 + ax1 * 16 + ax2 + 128
+                    ]
+        for i3_0 in T.unroll(2):
+            for i2_0 in T.unroll(2):
+                pad_temp = T.decl_buffer((4096,), "int8")
+                for ax0, ax1, ax2 in T.grid(16, 16, 16):
+                    if 6 <= i2_0 * 4 + ax0 and i3_0 * 4 + ax1 < 14:
+                        pad_temp[ax0 * 256 + ax1 * 16 + ax2] = placeholder_0_dm[
+                            i2_0 * 2048 + ax0 * 512 + i3_0 * 64 + ax1 * 16 + ax2 + 192
+                        ]
+
+    mod = partition_from_scheduled_tir(
+        main,
+        {
+            "s_tir.LoopPartition": {
+                "partition_const_loop": True,
+                "unroll_loop_with_partition_hint_no_interval": True,
+            }
+        },
+    )
+    tvm.ir.assert_structural_equal(mod["main"], partitioned_main.with_attr("global_symbol", "main"))
+
+
+def test_loop_partition_keep_loop_annotations():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer(160, "int32"), B: T.Buffer(160, "int32")) -> None:
+        for i in T.serial(
+            160,
+            annotations={"pragma_loop_partition_hint": True, "key": "value"},
+        ):
+            if i < 10:
+                B[i] = A[i] + 1
+            elif 10 <= i and i < 150:
+                B[i] = A[i] + 2
+            else:
+                B[i] = A[i] + 3
+
+    @T.prim_func(s_tir=True)
+    def after(A: T.Buffer(160, "int32"), B: T.Buffer(160, "int32")) -> None:
+        A_1 = T.decl_buffer((160,), "int32", data=A.data)
+        B_1 = T.decl_buffer((160,), "int32", data=B.data)
+        for i in T.serial(10, annotations={"key": "value"}):
+            B_1[i] = A_1[i] + 1
+        for i in T.serial(140, annotations={"key": "value"}):
+            B_1[i + 10] = A_1[i + 10] + 2
+        for i in T.serial(10, annotations={"key": "value"}):
+            B_1[i + 150] = A_1[i + 150] + 3
+
+    mod = partition_from_scheduled_tir(
+        before,
+        {
+            "s_tir.LoopPartition": {
+                "partition_const_loop": True,
+            }
+        },
+    )
+    tvm.ir.assert_structural_equal(mod["main"], after.with_attr("global_symbol", "main"))
+
+
+def test_loop_partition_with_unit_loop_in_condition():
+    @T.prim_func(s_tir=True)
+    def before(
+        placeholder: T.Buffer((50176,), "int8"),
+        placeholder_1: T.Buffer((25088,), "int8"),
+        placeholder_2: T.Buffer((25088,), "int8"),
+        T_concat: T.Buffer((100352,), "int8"),
+    ) -> None:
+        for k in range(1, annotations={"preserve_unit_loop": True}):
+            for i1 in T.serial(128, annotations={"pragma_loop_partition_hint": 1}):
+                for i2, i3 in T.grid(28, 28):
+                    if 96 <= k * 128 + i1:
+                        T_concat[k * i1 * 784 + i2 * 28 + i3] = placeholder_2[
+                            i1 * 784 + i2 * 28 + i3 - 75264
+                        ]
+                    if 64 <= k * 128 + i1 and k * 128 + i1 < 96:
+                        T_concat[i1 * 784 + i2 * 28 + i3] = placeholder_1[
+                            i1 * 784 + i2 * 28 + i3 - 50176
+                        ]
+                    if k * 128 + i1 < 64:
+                        T_concat[i1 * 784 + i2 * 28 + i3] = placeholder[i1 * 784 + i2 * 28 + i3]
+
+    @T.prim_func(s_tir=True)
+    def after(
+        placeholder: T.Buffer(50176, "int8"),
+        placeholder_1: T.Buffer(25088, "int8"),
+        placeholder_2: T.Buffer(25088, "int8"),
+        T_concat: T.Buffer(100352, "int8"),
+    ) -> None:
+        placeholder_3 = T.decl_buffer((50176,), "int8", data=placeholder.data)
+        placeholder_1_1 = T.decl_buffer((25088,), "int8", data=placeholder_1.data)
+        placeholder_2_1 = T.decl_buffer((25088,), "int8", data=placeholder_2.data)
+        T_concat_1 = T.decl_buffer((100352,), "int8", data=T_concat.data)
+        for k in T.serial(1, annotations={"preserve_unit_loop": True}):
+            for i1, i2, i3 in T.grid(64, 28, 28):
+                T_concat_1[i1 * 784 + i2 * 28 + i3] = placeholder_3[i1 * 784 + i2 * 28 + i3]
+            for i1, i2, i3 in T.grid(32, 28, 28):
+                T_concat_1[i1 * 784 + i2 * 28 + i3 + 50176] = placeholder_1_1[
+                    i1 * 784 + i2 * 28 + i3
+                ]
+            for i1, i2, i3 in T.grid(32, 28, 28):
+                T_concat_1[i2 * 28 + i3] = placeholder_2_1[i1 * 784 + i2 * 28 + i3]
+
+    mod = partition_from_scheduled_tir(
+        before,
+        {
+            "s_tir.LoopPartition": {
+                "partition_const_loop": True,
+            }
+        },
+    )
+    tvm.ir.assert_structural_equal(mod["main"], after.with_attr("global_symbol", "main"))
+
+
+@T.prim_func(s_tir=True)
+def concat_func_single_point(
+    placeholder: T.Buffer((28, 64), "int8"),
+    placeholder_1: T.Buffer((28, 1), "int8"),
+    placeholder_2: T.Buffer((28, 63), "int8"),
+    T_concat: T.Buffer((28, 128), "int8"),
+) -> None:
+    for i0 in range(28):
+        for i1 in T.serial(128, annotations={"pragma_loop_partition_hint": 1}):
+            if i1 > 63:
+                T_concat[i0, i1] = placeholder[i0, i1 - 64]
+            elif i1 == 63:
+                T_concat[i0, i1] = placeholder_1[i0, i1 - 63]
+            else:
+                T_concat[i0, i1] = placeholder_2[i0, i1]
+
+
+@T.prim_func(s_tir=True)
+def expected_partitioned_concat_single_point(
+    placeholder: T.Buffer((28, 64), "int8"),
+    placeholder_1: T.Buffer((28, 1), "int8"),
+    placeholder_2: T.Buffer((28, 63), "int8"),
+    T_concat: T.Buffer((28, 128), "int8"),
+):
+    placeholder_3 = T.decl_buffer((1792,), "int8", data=placeholder.data)
+    placeholder_1_1 = T.decl_buffer((28,), "int8", data=placeholder_1.data)
+    placeholder_2_1 = T.decl_buffer((1764,), "int8", data=placeholder_2.data)
+    T_concat_1 = T.decl_buffer((3584,), "int8", data=T_concat.data)
+    for i0 in range(28):
+        for i1 in range(63):
+            T_concat_1[i0 * 128 + i1] = placeholder_2_1[i0 * 63 + i1]
+        T_concat_1[i0 * 128 + 63] = placeholder_1_1[i0]
+        for i1 in range(64):
+            T_concat_1[i0 * 128 + i1 + 64] = placeholder_3[i0 * 64 + i1]
+
+
+@T.prim_func(s_tir=True)
+def concat_func_start_point_equality(
+    placeholder: T.Buffer((28, 64), "int8"),
+    placeholder_1: T.Buffer((28, 1), "int8"),
+    placeholder_2: T.Buffer((28, 63), "int8"),
+    T_concat: T.Buffer((28, 128), "int8"),
+) -> None:
+    for i0 in range(28):
+        for i1 in range(128, annotations={"pragma_loop_partition_hint": 1}):
+            if i1 == 0:
+                # Special case for i1 == 0
+                T_concat[i0, i1] = placeholder_1[i0, 0]
+            elif i1 < 64:
+                # Normal case for i1 in [1, 63]
+                T_concat[i0, i1] = placeholder_2[i0, i1]
+            else:
+                # Case for i1 in [64, 127]
+                T_concat[i0, i1] = placeholder[i0, i1 - 64]
+
+
+@T.prim_func(s_tir=True)
+def concat_func_start_point_equality_expected(
+    placeholder: T.Buffer((28, 64), "int8"),
+    placeholder_1: T.Buffer((28, 1), "int8"),
+    placeholder_2: T.Buffer((28, 63), "int8"),
+    T_concat: T.Buffer((28, 128), "int8"),
+):
+    placeholder_3 = T.decl_buffer((1792,), "int8", data=placeholder.data)
+    placeholder_1_1 = T.decl_buffer((28,), "int8", data=placeholder_1.data)
+    placeholder_2_1 = T.decl_buffer((1764,), "int8", data=placeholder_2.data)
+    T_concat_1 = T.decl_buffer((3584,), "int8", data=T_concat.data)
+    for i0 in range(28):
+        T_concat_1[i0 * 128] = placeholder_1_1[i0]
+        for i1 in range(63):
+            T_concat_1[i0 * 128 + i1 + 1] = placeholder_2_1[i0 * 63 + i1 + 1]
+        for i1 in range(64):
+            T_concat_1[i0 * 128 + i1 + 64] = placeholder_3[i0 * 64 + i1]
+
+
+@T.prim_func(s_tir=True)
+def concat_func_end_point_equality(
+    placeholder: T.Buffer((28, 64), "int8"),
+    placeholder_1: T.Buffer((28, 1), "int8"),
+    placeholder_2: T.Buffer((28, 63), "int8"),
+    T_concat: T.Buffer((28, 128), "int8"),
+) -> None:
+    for i0 in range(28):
+        for i1 in range(128, annotations={"pragma_loop_partition_hint": 1}):
+            if i1 == 127:
+                # Explicit equality check for the end point i1 == 127
+                T_concat[i0, i1] = placeholder_1[i0, 0]
+            elif i1 >= 64:
+                # Case for i1 in [64, 126]
+                T_concat[i0, i1] = placeholder[i0, i1 - 64]
+            else:
+                # Case for i1 in [0, 63]
+                T_concat[i0, i1] = placeholder_2[i0, i1]
+
+
+@T.prim_func(s_tir=True)
+def concat_func_end_point_equality_expected(
+    placeholder: T.Buffer((28, 64), "int8"),
+    placeholder_1: T.Buffer((28, 1), "int8"),
+    placeholder_2: T.Buffer((28, 63), "int8"),
+    T_concat: T.Buffer((28, 128), "int8"),
+):
+    placeholder_3 = T.decl_buffer((1792,), "int8", data=placeholder.data)
+    placeholder_1_1 = T.decl_buffer((28,), "int8", data=placeholder_1.data)
+    placeholder_2_1 = T.decl_buffer((1764,), "int8", data=placeholder_2.data)
+    T_concat_1 = T.decl_buffer((3584,), "int8", data=T_concat.data)
+    for i0 in range(28):
+        for i1 in range(64):
+            T_concat_1[i0 * 128 + i1] = placeholder_2_1[i0 * 63 + i1]
+        for i1 in range(63):
+            T_concat_1[i0 * 128 + i1 + 64] = placeholder_3[i0 * 64 + i1]
+        T_concat_1[i0 * 128 + 127] = placeholder_1_1[i0]
+
+
+@T.prim_func(s_tir=True)
+def concat_func_edge_equalities(
+    placeholder: T.Buffer((28, 64), "int8"),
+    placeholder_1: T.Buffer((28, 1), "int8"),
+    placeholder_2: T.Buffer((28, 1), "int8"),
+    T_concat: T.Buffer((28, 66), "int8"),
+) -> None:
+    for i0 in range(28):
+        for i1 in range(
+            66, annotations={"pragma_loop_partition_hint": 1}
+        ):  # Loop from 0 to 65 inclusive
+            if i1 == 0:
+                # Handle equality at the start of the range: i1 == 0
+                T_concat[i0, i1] = placeholder_2[i0, 0]
+            elif i1 == 65:
+                # Handle equality at the end of the range: i1 == 65
+                T_concat[i0, i1] = placeholder_1[i0, 0]
+            else:
+                # Copying from placeholder (from 0 to 63)
+                T_concat[i0, i1] = placeholder[i0, i1 - 1]
+
+
+@T.prim_func(s_tir=True)
+def concat_func_edge_equalities_expected(
+    placeholder: T.Buffer((28, 64), "int8"),
+    placeholder_1: T.Buffer((28, 1), "int8"),
+    placeholder_2: T.Buffer((28, 1), "int8"),
+    T_concat: T.Buffer((28, 66), "int8"),
+):
+    placeholder_3 = T.decl_buffer((1792,), "int8", data=placeholder.data)
+    placeholder_1_1 = T.decl_buffer((28,), "int8", data=placeholder_1.data)
+    placeholder_2_1 = T.decl_buffer((28,), "int8", data=placeholder_2.data)
+    T_concat_1 = T.decl_buffer((1848,), "int8", data=T_concat.data)
+    for i0 in range(28):
+        T_concat_1[i0 * 66] = placeholder_2_1[i0]
+        for i1 in range(64):
+            T_concat_1[i0 * 66 + i1 + 1] = placeholder_3[i0 * 64 + i1]
+        T_concat_1[i0 * 66 + 65] = placeholder_1_1[i0]
+
+
+@T.prim_func(s_tir=True)
+def concat_five_buffers_with_equalities(
+    buffer_a: T.Buffer((28, 1), "int8"),  # Used for i1 == 0
+    buffer_b: T.Buffer((28, 63), "int8"),  # Fills i1 from 1 to 63
+    buffer_c: T.Buffer((28, 1), "int8"),  # Used for i1 == 64
+    buffer_d: T.Buffer((28, 63), "int8"),  # Fills i1 from 65 to 128
+    buffer_e: T.Buffer((28, 1), "int8"),  # Used for i1 == 129
+    T_concat: T.Buffer((28, 129), "int8"),
+) -> None:
+    for i0 in range(28):
+        for i1 in range(130, annotations={"pragma_loop_partition_hint": 1}):
+            if i1 == 0:
+                T_concat[i0, i1] = buffer_a[i0, 0]
+            elif i1 == 64:
+                T_concat[i0, i1] = buffer_c[i0, 0]
+            elif i1 == 129:
+                T_concat[i0, i1] = buffer_e[i0, 0]
+            elif i1 < 64:
+                T_concat[i0, i1] = buffer_b[i0, i1 - 1]
+            else:  # i1 > 64 and i1 < 128
+                T_concat[i0, i1] = buffer_d[i0, i1 - 65]
+
+
+@T.prim_func(s_tir=True)
+def concat_five_buffers_with_equalities_expected(
+    buffer_a: T.Buffer((28, 1), "int8"),  # Used for i1 == 0
+    buffer_b: T.Buffer((28, 63), "int8"),  # Fills i1 from 1 to 63
+    buffer_c: T.Buffer((28, 1), "int8"),  # Used for i1 == 64
+    buffer_d: T.Buffer((28, 63), "int8"),  # Fills i1 from 65 to 128
+    buffer_e: T.Buffer((28, 1), "int8"),  # Used for i1 == 129
+    T_concat: T.Buffer((28, 129), "int8"),
+):
+    buffer_a_1 = T.decl_buffer((28,), "int8", data=buffer_a.data)
+    buffer_b_1 = T.decl_buffer((1764,), "int8", data=buffer_b.data)
+    buffer_c_1 = T.decl_buffer((28,), "int8", data=buffer_c.data)
+    buffer_d_1 = T.decl_buffer((1764,), "int8", data=buffer_d.data)
+    buffer_e_1 = T.decl_buffer((28,), "int8", data=buffer_e.data)
+    T_concat_1 = T.decl_buffer((3612,), "int8", data=T_concat.data)
+    for i0 in range(28):
+        T_concat_1[i0 * 129] = buffer_a_1[i0]
+        for i1 in range(63):
+            T_concat_1[i0 * 129 + i1 + 1] = buffer_b_1[i0 * 63 + i1]
+        T_concat_1[i0 * 129 + 64] = buffer_c_1[i0]
+        for i1 in range(64):
+            T_concat_1[i0 * 129 + i1 + 65] = buffer_d_1[i0 * 63 + i1]
+        T_concat_1[i0 * 129 + 129] = buffer_e_1[i0]
+
+
+@T.prim_func(s_tir=True)
+def nested_partition_with_single_points(A: T.Buffer((25,), "int32")):
+    for i in T.serial(5, annotations={"pragma_loop_partition_hint": 1}):
+        if i == 1:
+            for j in T.serial(5, annotations={"pragma_loop_partition_hint": 1}):
+                if j > 2:
+                    A[i * 5 + j] = i * 5 + j
+        else:
+            for j in T.serial(5, annotations={"pragma_loop_partition_hint": 1}):
+                if j > 2:
+                    A[i * 5 + j] = i * 15 + j
+
+
+@T.prim_func(s_tir=True)
+def nested_partition_with_single_points_expected(A: T.Buffer((25,), "int32")):
+    A_1 = T.decl_buffer((25,), "int32", data=A.data)
+    for j in range(2):
+        A_1[j + 3] = j + 3
+    for j in range(2):
+        A_1[j + 8] = j + 8
+    for i, j in T.grid(3, 2):
+        A_1[i * 5 + j + 13] = i * 15 + j + 33
+
+
+@pytest.mark.parametrize(
+    "origin,expected",
+    [
+        (concat_func_single_point, expected_partitioned_concat_single_point),
+        (concat_func_start_point_equality, concat_func_start_point_equality_expected),
+        (concat_func_end_point_equality, concat_func_end_point_equality_expected),
+        (concat_func_edge_equalities, concat_func_edge_equalities_expected),
+        (concat_five_buffers_with_equalities, concat_five_buffers_with_equalities_expected),
+        (nested_partition_with_single_points, nested_partition_with_single_points_expected),
+    ],
+)
+def test_single_point_partition(origin, expected):
+    origin = origin.with_attr({"global_symbol": "main"})
+    expected = expected.with_attr({"global_symbol": "main"})
+    mod = partition_from_scheduled_tir(
+        origin,
+        {
+            "s_tir.LoopPartition": {
+                "partition_const_loop": True,
+                "unroll_loop_with_partition_hint_no_interval": True,
+            }
+        },
+    )
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_equation_on_floordiv():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((2, 2, 20), "int32")):
+        for i in T.serial(5, annotations={"pragma_loop_partition_hint": 1}):
+            if i == 1:
+                for vv in T.vectorized(640, annotations={"pragma_loop_partition_hint": 1}):
+                    if i * 2 + vv // 320 == 3:
+                        A[i - 1, i * 2 + vv // 320 - 3, vv % 320 // 16] = 1
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((2, 2, 20), "int32")):
+        for vv in T.vectorized(320):
+            A[0, 0, vv // 16] = 1
+
+    expected = expected.with_attr({"global_symbol": "main"})
+    after = partition_from_scheduled_tir(
+        before.with_attr("global_symbol", "main"), {}, do_flatten=False
+    )
+    tvm.ir.assert_structural_equal(after["main"], expected)
+
+
+def test_ignore_loop_partition_hint():
+    """Skip unroll body and prologue for pipeline case"""
+
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((10), "float32"), D: T.Buffer((10), "float32")):
+        B = T.decl_buffer([2], "float32")
+        C = T.decl_buffer([2], "float32")
+        for i in T.serial(12, annotations={"pragma_loop_partition_hint": 1}):
+            if T.ignore_loop_partition(i < 10):
+                B[i % 2] = A[i] + 1.0
+            if T.ignore_loop_partition(1 <= i and i < 11):
+                C[(i - 1) % 2] = B[(i - 1) % 2] + 2.0
+            if 2 <= i:
+                D[i - 2] = C[i % 2] + 3.0
+
+    @T.prim_func(s_tir=True)
+    def expected(A: T.Buffer((10), "float32"), D: T.Buffer((10), "float32")):
+        B = T.decl_buffer([2], "float32")
+        C = T.decl_buffer([2], "float32")
+        for i in range(2):
+            B[i] = A[i] + 1.0
+            if i == 1:
+                C[i - 1] = B[i - 1] + 2.0
+        for i in T.serial(10):
+            if i < 8:
+                B[i % 2] = A[i + 2] + 1.0
+            if i < 9:
+                C[(i + 1) % 2] = B[(i + 1) % 2] + 2.0
+            D[i] = C[i % 2] + 3.0
+
+    expected = expected.with_attr({"global_symbol": "main"})
+    after = partition_from_scheduled_tir(
+        before.with_attr({"global_symbol": "main"}), {}, do_flatten=False
+    )
+    tvm.ir.assert_structural_equal(after["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_lower_init_block.py

@@ -0,0 +1,112 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401, F821
+import tvm
+from tvm import s_tir
+from tvm.script import tirx as T
+
+# pylint: disable=no-self-argument
+
+
+@tvm.script.ir_module
+class WithInit:
+    @T.prim_func(s_tir=True)
+    def main(a: T.handle, b: T.handle) -> None:
+        A = T.match_buffer(a, [64, 64, 64])
+        B = T.match_buffer(b, [64])
+
+        for i0, j0 in T.grid(64, 64):
+            for k0 in T.serial(32, 64):
+                with T.sblock():
+                    i, j, k = T.axis.remap("SRR", [i0, j0, k0])
+                    with T.init():
+                        B[i] = T.float32(0)
+                    B[i] += A[i, j, k]
+
+
+@tvm.script.ir_module
+class WithBranch:
+    @T.prim_func(s_tir=True)
+    def main(a: T.handle, b: T.handle) -> None:
+        A = T.match_buffer(a, [64, 64, 64])
+        B = T.match_buffer(b, [64])
+
+        for i0, j0 in T.grid(64, 64):
+            for k0 in T.serial(32, 64):
+                with T.sblock():
+                    i, j, k = T.axis.remap("SRR", [i0, j0, k0])
+                    T.reads(A[i, j, k])
+                    T.writes(B[i])
+                    if (j == 0) and (k == 32):
+                        B[i] = T.float32(0)
+                    B[i] += A[i, j, k]
+
+
+@tvm.script.ir_module
+class InitWithMatchBuffer:
+    @T.prim_func(s_tir=True)
+    def main(a: T.handle, b: T.handle) -> None:
+        A = T.match_buffer(a, [64, 64, 64])
+        B = T.match_buffer(b, [64])
+
+        for i0, j0 in T.grid(64, 64):
+            for k0 in T.serial(32, 64):
+                with T.sblock():
+                    i, j, k = T.axis.remap("SRR", [i0, j0, k0])
+                    BB = T.match_buffer(B[i], ())
+                    AA = T.match_buffer(A[i, 0:64, 0:64], (64, 64))
+                    with T.init():
+                        BB[()] = T.float32(0)
+                    BB[()] += AA[j, k]
+
+
+@tvm.script.ir_module
+class BranchWithMatchBuffer:
+    @T.prim_func(s_tir=True)
+    def main(a: T.handle, b: T.handle) -> None:
+        A = T.match_buffer(a, [64, 64, 64])
+        B = T.match_buffer(b, [64])
+
+        for i0, j0 in T.grid(64, 64):
+            for k0 in T.serial(32, 64):
+                with T.sblock():
+                    i, j, k = T.axis.remap("SRR", [i0, j0, k0])
+                    T.reads(A[i, j, k])
+                    T.writes(B[i])
+                    BB = T.match_buffer(B[i], ())
+                    AA = T.match_buffer(A[i, 0:64, 0:64], (64, 64))
+                    if (j == 0) and (k == 32):
+                        BB[()] = T.float32(0)
+                    BB[()] += AA[j, k]
+
+
+def test_lower_reduction():
+    origin_mod = WithInit
+    mod = tvm.s_tir.transform.LowerInitBlock()(origin_mod)
+    tvm.ir.assert_structural_equal(mod, WithBranch, True)
+
+
+def test_lower_match_buffer():
+    origin_mod = InitWithMatchBuffer
+    mod = tvm.s_tir.transform.LowerInitBlock()(origin_mod)
+    tvm.ir.assert_structural_equal(mod, BranchWithMatchBuffer, True)
+
+
+if __name__ == "__main__":
+    test_lower_reduction()
+    test_lower_match_buffer()
+    test_lower_te()

tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py

@@ -0,0 +1,571 @@
+# 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.s_tir
+import tvm.testing
+from tvm.script import tirx as T
+
+
+def _check(original, transformed):
+    mod = tvm.IRModule.from_expr(original.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.LowerMatchBuffer()(mod)
+    mod = tvm.tirx.transform.StmtSimplify()(mod)
+    tvm.ir.assert_structural_equal(mod["main"], transformed.with_attr("global_symbol", "main"))
+
+
+def _check_fail(original):
+    mod = tvm.IRModule.from_expr(original)
+    with pytest.raises(RuntimeError):
+        mod = tvm.s_tir.transform.LowerMatchBuffer()(mod)
+
+
+@T.prim_func(s_tir=True)
+def buffer_load_store(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16, 16))
+    C = T.match_buffer(c, (16, 16))
+    for i, j, k in T.grid(4, 16, 8):
+        with T.sblock():
+            T.reads(C[i * 4 : i * 4 + 4, k * 2 : k * 2 + 2])
+            T.writes(A[i * 4 : i * 4 + 4, j, k * 2 : k * 2 + 2])
+            sub_A = T.match_buffer(
+                A[i * 4 : i * 4 + 4, j, k * 2 : k * 2 + 2], (4, 1, 2), offset_factor=1
+            )
+            sub_C = T.match_buffer(C[i * 4 : i * 4 + 4, k * 2 : k * 2 + 2], (4, 2), offset_factor=1)
+            for ii, kk in T.grid(4, 2):
+                sub_A[ii, 0, kk] += sub_C[ii, kk]
+
+
+@T.prim_func(s_tir=True)
+def transformed_buffer_load_store(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16, 16))
+    C = T.match_buffer(c, (16, 16))
+    for i, j, k in T.grid(4, 16, 8):
+        with T.sblock():
+            T.reads(C[i * 4 : i * 4 + 4, k * 2 : k * 2 + 2])
+            T.writes(A[i * 4 : i * 4 + 4, j, k * 2 : k * 2 + 2])
+            for ii, kk in T.grid(4, 2):
+                A[i * 4 + ii, j, k * 2 + kk] += C[i * 4 + ii, k * 2 + kk]
+
+
+# Dummy intrinsic whose arguments exercise match_buffer fields.  TVMScript
+# evaluates the call eagerly (to 0), so it must NOT be registered as an op:
+# registering "tirx.intrin_test" only leaves a category-less op in the tirx
+# registry, breaking the exactly-one-category invariant for later tests.
+def intrin_test(data, elem_offset, stride_0, stride_1, shape_0, shape_1):
+    return 0
+
+
+@T.prim_func(s_tir=True)
+def opaque_access(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (32, 64, 128))
+    B = T.match_buffer(b, (64, 64, 64))
+    for i, j, k in T.grid(2, 64, 8):
+        with T.sblock():
+            T.reads([])
+            T.writes(A[i * 16 : i * 16 + 16, j, k * 16 : k * 16 + 16])
+            sub_A = T.match_buffer(
+                A[i * 16 : i * 16 + 16, j, k * 16 : k * 16 + 16],
+                (16, 1, 16),
+                strides=[8192, 128, 1],
+                offset_factor=1,
+            )
+            T.evaluate(
+                intrin_test(
+                    sub_A.data,
+                    sub_A.elem_offset,
+                    sub_A.strides[0],
+                    sub_A.strides[1],
+                    sub_A.shape[0],
+                    sub_A.shape[1],
+                )
+            )
+    for i, j, k in T.grid(64, 2, 8):
+        with T.sblock():
+            Bs_0 = T.int32()
+            Bs_1 = T.int32()
+            T.reads([])
+            T.writes(B[i, j * 32 : j * 32 + 32, k * 8 : k * 8 + 8])
+            sub_B = T.match_buffer(
+                B[i, j * 32 : j * 32 + 32, k * 8 : k * 8 + 8],
+                (32, 8),
+                strides=[Bs_0, Bs_1],
+                offset_factor=1,
+            )
+            T.evaluate(
+                intrin_test(
+                    sub_B.data,
+                    sub_B.elem_offset,
+                    sub_B.strides[0],
+                    sub_B.strides[1],
+                    sub_B.shape[0],
+                    sub_B.shape[1],
+                )
+            )
+
+
+@T.prim_func(s_tir=True)
+def transformed_opaque_access(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (32, 64, 128))
+    B = T.match_buffer(b, (64, 64, 64))
+    for i, j, k in T.grid(2, 64, 8):
+        with T.sblock():
+            T.reads([])
+            T.writes(A[i * 16 : i * 16 + 16, j, k * 16 : k * 16 + 16])
+            T.evaluate(
+                intrin_test(
+                    A.data,
+                    i * 131072 + j * 128 + k * 16,
+                    8192,
+                    128,
+                    16,
+                    1,
+                )
+            )
+    for i, j, k in T.grid(64, 2, 8):
+        with T.sblock():
+            T.reads([])
+            T.writes(B[i, j * 32 : j * 32 + 32, k * 8 : k * 8 + 8])
+            T.evaluate(
+                intrin_test(
+                    B.data,
+                    i * 4096 + j * 2048 + k * 8,
+                    64,
+                    1,
+                    32,
+                    8,
+                )
+            )
+
+
+@T.prim_func(s_tir=True)
+def high_dim_opaque_access(a: T.handle) -> None:
+    A = T.match_buffer(a, (16, 32, 64))
+    for i, j, k in T.grid(16, 2, 4):
+        with T.sblock():
+            As_0 = T.int32()
+            As_1 = T.int32()
+            T.reads([])
+            T.writes(A[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16])
+            sub_A = T.match_buffer(
+                A[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16],
+                (16, 16),
+                strides=[As_0, As_1],
+                offset_factor=1,
+            )
+            T.evaluate(
+                intrin_test(
+                    sub_A.data,
+                    sub_A.elem_offset,
+                    sub_A.strides[0],
+                    sub_A.strides[1],
+                    sub_A.shape[0],
+                    sub_A.shape[1],
+                )
+            )
+
+
+@T.prim_func(s_tir=True)
+def transformed_high_dim_opaque_access(a: T.handle) -> None:
+    A = T.match_buffer(a, (16, 32, 64))
+    for i, j, k in T.grid(16, 2, 4):
+        with T.sblock():
+            T.reads([])
+            T.writes(A[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16])
+            T.evaluate(
+                intrin_test(
+                    A.data,
+                    i * 2048 + j * 1024 + k * 16,
+                    64,
+                    1,
+                    16,
+                    16,
+                )
+            )
+
+
+@T.prim_func(s_tir=True)
+def high_dim_opaque_access_with_source_strides(a: T.handle) -> None:
+    A = T.match_buffer(a, (16, 32, 64), strides=[2576, 80, 1])
+    for i, j, k in T.grid(16, 2, 4):
+        with T.sblock():
+            As_0 = T.int32()
+            As_1 = T.int32()
+            T.reads([])
+            T.writes(A[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16])
+            sub_A = T.match_buffer(
+                A[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16],
+                (16, 16),
+                strides=[As_0, As_1],
+                offset_factor=1,
+            )
+            T.evaluate(
+                intrin_test(
+                    sub_A.data,
+                    sub_A.elem_offset,
+                    sub_A.strides[0],
+                    sub_A.strides[1],
+                    sub_A.shape[0],
+                    sub_A.shape[1],
+                )
+            )
+
+
+@T.prim_func(s_tir=True)
+def transformed_high_dim_opaque_access_with_source_strides(a: T.handle) -> None:
+    A = T.match_buffer(a, (16, 32, 64), strides=[2576, 80, 1])
+    for i, j, k in T.grid(16, 2, 4):
+        with T.sblock():
+            T.reads([])
+            T.writes(A[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16])
+            T.evaluate(
+                intrin_test(
+                    A.data,
+                    i * 2576 + j * 1280 + k * 16,
+                    80,
+                    1,
+                    16,
+                    16,
+                )
+            )
+
+
+@T.prim_func(s_tir=True)
+def recursive_match(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (64, 64, 64))
+    B = T.match_buffer(b, (64, 64, 64))
+    for i, j, k in T.grid(64, 4, 4):
+        with T.sblock():
+            T.reads([])
+            T.writes(
+                [
+                    A[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16],
+                    B[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16],
+                ]
+            )
+            As_0 = T.int32()
+            As_1 = T.int32()
+            sub_A = T.match_buffer(
+                A[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16],
+                (16, 16),
+                strides=[As_0, As_1],
+                offset_factor=1,
+            )
+            sub_B = T.match_buffer(
+                B[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16],
+                (16, 16),
+                offset_factor=1,
+            )
+            for jj, kk in T.grid(4, 4):
+                with T.sblock():
+                    T.reads([])
+                    T.writes(
+                        [
+                            sub_A[jj * 4 : jj * 4 + 4, kk * 4 : kk * 4 + 4],
+                            sub_B[jj * 4 : jj * 4 + 4, kk * 4 : kk * 4 + 4],
+                        ]
+                    )
+                    Ass_0 = T.int32()
+                    Ass_1 = T.int32()
+                    sub_sub_A = T.match_buffer(
+                        sub_A[jj * 4 : jj * 4 + 4, kk * 4 : kk * 4 + 4],
+                        (4, 4),
+                        strides=[Ass_0, Ass_1],
+                        offset_factor=1,
+                    )
+                    sub_sub_B = T.match_buffer(
+                        sub_B[jj * 4 : jj * 4 + 4, kk * 4 : kk * 4 + 4],
+                        (4, 4),
+                        offset_factor=1,
+                    )
+                    T.evaluate(
+                        intrin_test(
+                            sub_sub_A.data,
+                            sub_sub_A.elem_offset,
+                            sub_sub_A.strides[0],
+                            sub_sub_A.strides[1],
+                            sub_sub_A.shape[0],
+                            sub_sub_A.shape[1],
+                        )
+                    )
+                    for jjj, kkk in T.grid(4, 4):
+                        sub_sub_B[jjj, kkk] = 1
+
+
+@T.prim_func(s_tir=True)
+def transformed_recursive_match(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (64, 64, 64))
+    B = T.match_buffer(b, (64, 64, 64))
+    for i, j, k in T.grid(64, 4, 4):
+        with T.sblock():
+            T.reads([])
+            T.writes(
+                [
+                    A[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16],
+                    B[i, j * 16 : j * 16 + 16, k * 16 : k * 16 + 16],
+                ]
+            )
+            for jj, kk in T.grid(4, 4):
+                with T.sblock():
+                    T.reads([])
+                    T.writes(
+                        [
+                            A[
+                                i,
+                                j * 16 + jj * 4 : j * 16 + jj * 4 + 4,
+                                k * 16 + kk * 4 : k * 16 + kk * 4 + 4,
+                            ],
+                            B[
+                                i,
+                                j * 16 + jj * 4 : j * 16 + jj * 4 + 4,
+                                k * 16 + kk * 4 : k * 16 + kk * 4 + 4,
+                            ],
+                        ]
+                    )
+                    T.evaluate(
+                        intrin_test(
+                            A.data,
+                            i * 4096 + j * 1024 + jj * 256 + k * 16 + kk * 4,
+                            64,
+                            1,
+                            4,
+                            4,
+                        )
+                    )
+                    for jjj, kkk in T.grid(4, 4):
+                        B[i, j * 16 + jj * 4 + jjj, k * 16 + kk * 4 + kkk] = 1
+
+
+@T.prim_func(s_tir=True)
+def symbolic_match(a: T.handle, b: T.handle, n: T.int32, m: T.int32) -> None:
+    A = T.match_buffer(a, (n * m, m))
+    B = T.match_buffer(b, (n * 2, m * 4))
+    for i in range(0, n):
+        with T.sblock():
+            T.reads([])
+            T.writes([A[i * m : i * m + n, 0:m], B[i * n : i * n + 2, 0 : m * 4]])
+            Bs_0 = T.int32()
+            Bs_1 = T.int32()
+            sub_A = T.match_buffer(A[i * m : i * m + m, 0:m], (m, m), offset_factor=1)
+            sub_B = T.match_buffer(
+                B[i * n : i * n + 2, 0 : m * 4], (2, m * 4), strides=[Bs_0, Bs_1], offset_factor=1
+            )
+            for ii, jj in T.grid(m, m):
+                sub_A[ii, jj] = 1
+            for j in range(0, 4):
+                T.evaluate(
+                    intrin_test(
+                        sub_B.data,
+                        sub_B.elem_offset,
+                        sub_B.strides[0],
+                        sub_B.strides[1],
+                        sub_B.shape[0],
+                        sub_B.shape[1],
+                    )
+                )
+
+
+@T.prim_func(s_tir=True)
+def transformed_symbolic_match(a: T.handle, b: T.handle, n: T.int32, m: T.int32) -> None:
+    A = T.match_buffer(a, (n * m, m))
+    B = T.match_buffer(b, (n * 2, m * 4))
+    for i in range(0, n):
+        with T.sblock():
+            T.reads([])
+            T.writes([A[i * m : i * m + n, 0:m], B[i * n : i * n + 2, 0 : m * 4]])
+            for ii, jj in T.grid(m, m):
+                A[i * m + ii, jj] = 1
+            for j in range(0, 4):
+                T.evaluate(
+                    intrin_test(
+                        B.data,
+                        i * n * (m * 4),
+                        m * 4,
+                        1,
+                        2,
+                        m * 4,
+                    )
+                )
+
+
+@T.prim_func(s_tir=True)
+def rank0_buffer(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8))
+    B = T.match_buffer(b, (8, 8))
+    for i, j in T.grid(8, 8):
+        with T.sblock():
+            T.reads([])
+            T.writes([A[i, j], B[i, j]])
+            sub_A = T.match_buffer(A[i, j], (), offset_factor=1)
+            sub_B = T.match_buffer(B[i, j], (), offset_factor=1)
+            sub_A[()] = 1
+            T.evaluate(
+                intrin_test(
+                    sub_B.data,
+                    sub_B.elem_offset,
+                    0,
+                    0,
+                    0,
+                    0,
+                )
+            )
+
+
+@T.prim_func(s_tir=True)
+def transformed_rank0_buffer(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8))
+    B = T.match_buffer(b, (8, 8))
+    for i, j in T.grid(8, 8):
+        with T.sblock():
+            T.reads([])
+            T.writes([A[i, j], B[i, j]])
+            A[i, j] = 1
+            T.evaluate(
+                intrin_test(
+                    B.data,
+                    i * 8 + j,
+                    0,
+                    0,
+                    0,
+                    0,
+                )
+            )
+
+
+@T.prim_func(s_tir=True)
+def fail_match_load(a: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8))
+    for i, j in T.grid(8, 8):
+        with T.sblock():
+            T.reads(A[i, j])
+            T.writes([])
+            sub_A = T.match_buffer(A[i, j], (), elem_offset=0)
+            T.evaluate(sub_A[()])
+
+
+@T.prim_func(s_tir=True)
+def fail_match_store(a: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8))
+    for i, j in T.grid(8, 8):
+        with T.sblock():
+            T.reads([])
+            T.writes(A[i, j])
+            sub_A = T.match_buffer(A[i, j], (), elem_offset=0)
+            sub_A[()] = 1
+
+
+# well-formed checker complains about redefinition of a stride variable
+@T.prim_func(check_well_formed=False, s_tir=True)
+def fail_buffer_bind(a: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8))
+    for i, j in T.grid(8, 2):
+        with T.sblock():
+            stride = T.int32()
+            sub_A = T.match_buffer(
+                A[i, j * 4 : j * 4 + 4], (1, 4), strides=[stride, stride], offset_factor=1
+            )
+            for jj in range(0, 4):
+                sub_A[i, j * 4 + jj] = 1
+
+
+# well-formed checker complains about redefinition of a stride variable
+@T.prim_func(check_well_formed=False, s_tir=True)
+def fail_match_func_param(a: T.handle, m: T.int32, n: T.int32) -> None:
+    A = T.match_buffer(a, (8, 8))
+    for i, j in T.grid(8, 2):
+        with T.sblock():
+            sub_A = T.match_buffer(A[i, j * 4 : j * 4 + 4], (1, 4), strides=[m, n], offset_factor=1)
+            for jj in range(0, 4):
+                sub_A[i, j * 4 + jj] = 1
+
+
+def test_buffer_load_store():
+    _check(buffer_load_store, transformed_buffer_load_store)
+
+
+def test_opaque_access():
+    _check(opaque_access, transformed_opaque_access)
+
+
+def test_high_dim_opaque_access():
+    _check(high_dim_opaque_access, transformed_high_dim_opaque_access)
+    _check(
+        high_dim_opaque_access_with_source_strides,
+        transformed_high_dim_opaque_access_with_source_strides,
+    )
+
+
+def test_recursive_match():
+    _check(recursive_match, transformed_recursive_match)
+
+
+def test_symbolic_match():
+    _check(symbolic_match, transformed_symbolic_match)
+
+
+def test_rank0_buffer():
+    _check(rank0_buffer, transformed_rank0_buffer)
+
+
+def test_fail_load_store():
+    _check_fail(fail_match_load)
+    _check_fail(fail_match_store)
+
+
+def test_fail_buffer_bind():
+    _check_fail(fail_buffer_bind)
+
+
+def test_fail_match_func_param():
+    _check_fail(fail_match_func_param)
+
+
+@T.prim_func(s_tir=True)
+def scalar_match_buffer_type_coercion(a: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8))
+    for i, j in T.grid(8, 8):
+        with T.sblock(""):
+            vi = T.axis.spatial(8, i)
+            vj = T.axis.spatial(8, j)
+            T.reads()
+            T.writes(A[vi, vj])
+            # Create scalar match buffer from single element - this triggers type coercion
+            scalar_buf = T.match_buffer(A[vi, vj], (), offset_factor=1)
+            scalar_buf[()] = T.float32(1.0)
+
+
+@T.prim_func(s_tir=True)
+def transformed_scalar_match_buffer_type_coercion(a: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8))
+    for i, j in T.grid(8, 8):
+        with T.sblock(""):
+            vi = T.axis.spatial(8, i)
+            vj = T.axis.spatial(8, j)
+            T.reads()
+            T.writes(A[vi, vj])
+            # Scalar match_buffer eliminated, direct assignment
+            A[vi, vj] = T.float32(1.0)
+
+
+def test_scalar_match_buffer_type_coercion():
+    _check(scalar_match_buffer_type_coercion, transformed_scalar_match_buffer_type_coercion)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py

@@ -0,0 +1,429 @@
+# 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
+import tvm.s_tir
+import tvm.testing
+from tvm.script import tirx as T
+
+
+def _check(original, transformed):
+    func = original
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+    mod = tvm.tirx.transform.StmtSimplify()(mod)
+    tvm.ir.assert_structural_equal(
+        mod["main"], transformed.with_attr("global_symbol", "main"), True
+    )
+
+
+@T.prim_func(s_tir=True)
+def compacted_elementwise_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), "float32")
+    C = T.match_buffer(c, (16, 16), "float32")
+    for i in range(0, 16):
+        with T.sblock():
+            T.reads(A[i, 0:16])
+            T.writes(C[i, 0:16])
+            B = T.sblock_alloc_buffer([1, 16], "float32", scope="global")
+            for j in range(0, 16):
+                with T.sblock():
+                    T.reads(A[i, j])
+                    T.writes(B[0, j])
+                    B[0, j] = A[i, j] + 1.0
+            for j in range(0, 16):
+                with T.sblock():
+                    T.reads(B[0, j])
+                    T.writes(C[i, j])
+                    C[i, j] = B[0, j] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def transformed_elementwise_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), "float32")
+    C = T.match_buffer(c, (16, 16), "float32")
+    for i in T.serial(0, 16):
+        B_new = T.alloc_buffer(
+            [1, 16],
+            "float32",
+            annotations={"buffer_allocated_addr": [], "buffer_data_alignment": 64},
+        )
+        for j in T.serial(0, 16):
+            B_new[0, j] = A[i, j] + 1.0
+        for j in T.serial(0, 16):
+            C[i, j] = B_new[0, j] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def compacted_gpu_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), "float32")
+    C = T.match_buffer(c, (16, 16), "float32")
+    for i0 in T.thread_binding(0, 4, thread="blockIdx.x"):
+        for i1 in T.thread_binding(0, 2, thread="threadIdx.x"):
+            for i2 in T.thread_binding(0, 2, thread="vthread"):
+                with T.sblock():
+                    T.reads(A[i0 * 4 + i1 * 2 + i2, 0:16])
+                    T.writes(C[i0 * 4 + i1 * 2 + i2, 0:16])
+                    B = T.sblock_alloc_buffer([1, 16], "float32", scope="local")
+                    for j in range(0, 16):
+                        with T.sblock():
+                            T.reads(A[i0 * 4 + i1 * 2 + i2, j])
+                            T.writes(B[0, j])
+                            B[0, j] = A[i0 * 4 + i1 * 2 + i2, j] + 1.0
+                    for j in range(0, 16):
+                        with T.sblock():
+                            T.reads(B[0, j])
+                            T.writes(C[i0 * 4 + i1 * 2 + i2, j])
+                            C[i0 * 4 + i1 * 2 + i2, j] = B[0, j] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def transformed_gpu_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), "float32")
+    C = T.match_buffer(c, (16, 16), "float32")
+
+    i0 = T.env_thread("blockIdx.x")
+    i1 = T.env_thread("threadIdx.x")
+    i2 = T.env_thread("vthread")
+
+    T.launch_thread(i0, 4)
+    T.launch_thread(i1, 2)
+    T.launch_thread(i2, 2)
+    B = T.alloc_buffer(
+        [1, 16],
+        "float32",
+        scope="local",
+        annotations={"buffer_allocated_addr": [], "buffer_data_alignment": 64},
+    )
+    for j in range(0, 16):
+        B[0, j] = A[i0 * 4 + i1 * 2 + i2, j] + 1.0
+    for j in range(0, 16):
+        C[i0 * 4 + i1 * 2 + i2, j] = B[0, j] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def compacted_symbolic_func(a: T.handle, c: T.handle, n: T.int32, m: T.int32) -> None:
+    A = T.match_buffer(a, (n, m), "float32")
+    C = T.match_buffer(c, (n, m), "float32")
+
+    for i in range(0, n):
+        with T.sblock():
+            T.reads(A[i, m])
+            T.writes(C[i, m])
+            B = T.sblock_alloc_buffer((m,), "float32", scope="global")
+            for j in range(0, m):
+                with T.sblock():
+                    T.reads(A[i, j])
+                    T.writes(B[j])
+                    B[j] = A[i, j] + 1.0
+            for j in range(0, m):
+                with T.sblock():
+                    T.reads(B[j])
+                    T.writes(C[i, j])
+                    C[i, j] = B[j] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def transformed_symbolic_func(a: T.handle, c: T.handle, n: T.int32, m: T.int32) -> None:
+    A = T.match_buffer(a, (n, m), "float32")
+    C = T.match_buffer(c, (n, m), "float32")
+
+    for i in range(0, n):
+        B = T.alloc_buffer(
+            [m],
+            "float32",
+            annotations={"buffer_allocated_addr": [], "buffer_data_alignment": 64},
+        )
+        for j in range(0, m):
+            B[j] = A[i, j] + 1.0
+        for j in range(0, m):
+            C[i, j] = B[j] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def compacted_predicate_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (32), "float32")
+    C = T.match_buffer(c, (32), "float32")
+
+    for i, j in T.grid(5, 7):
+        with T.sblock():
+            T.reads(A[i * 7 + j])
+            T.writes(C[i * 7 + j])
+            T.where(i * 7 + j < 32)
+            C[i * 7 + j] = A[i * 7 + j] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def transformed_predicate_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (32), "float32")
+    C = T.match_buffer(c, (32), "float32")
+
+    for i, j in T.grid(5, 7):
+        if i * 7 + j < 32:
+            C[i * 7 + j] = A[i * 7 + j] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def compacted_unit_loop_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (32), "float32")
+    C = T.match_buffer(c, (32), "float32")
+
+    for x, y, z in T.grid(4, 1, 8):
+        with T.sblock():
+            T.reads(A[x * 8 + y * 8 + z])
+            T.writes(C[x * 8 + y * 8 + z])
+            C[x * 8 + y * 8 + z] = A[x * 8 + y * 8 + z] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def transformed_unit_loop_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (32), "float32")
+    C = T.match_buffer(c, (32), "float32")
+
+    for x, z in T.grid(4, 8):
+        C[x * 8 + z] = A[x * 8 + z] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def compacted_multi_alloc_func(a: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (32), "float32")
+    D = T.match_buffer(d, (32), "float32")
+
+    for i in range(0, 32):
+        with T.sblock():
+            T.reads(A[i])
+            T.writes(D[i])
+            B = T.sblock_alloc_buffer((32,), scope="global")
+            C = T.sblock_alloc_buffer((32,), scope="global")
+            B[i] = A[i] + 1.0
+            C[i] = A[i] + B[i]
+            D[i] = C[i] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def transformed_multi_alloc_func(a: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (32), "float32")
+    D = T.match_buffer(d, (32), "float32")
+
+    for i in range(0, 32):
+        B = T.alloc_buffer(
+            (32,),
+            "float32",
+            annotations={"buffer_allocated_addr": [], "buffer_data_alignment": 64},
+        )
+        C = T.alloc_buffer(
+            (32,),
+            "float32",
+            annotations={"buffer_allocated_addr": [], "buffer_data_alignment": 64},
+        )
+        B[i] = A[i] + 1.0
+        C[i] = A[i] + B[i]
+        D[i] = C[i] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def compacted_strided_buffer_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), "float32")
+    C = T.match_buffer(c, (16, 16), "float32")
+    for i0 in range(0, 4):
+        with T.sblock():
+            T.reads(A[i0 * 4 : i0 * 4 + 4, 0:16])
+            T.writes(C[i0 * 4 : i0 * 4 + 4, 0:16])
+            B = T.sblock_alloc_buffer([4, 16], "float32", strides=[17, 1], scope="global")
+            for i1 in range(0, 4):
+                for j in range(0, 16):
+                    with T.sblock():
+                        T.reads(A[i0 * 4 + i1, j])
+                        T.writes(B[i1, j])
+                        B[i1, j] = A[i0 * 4 + i1, j] + 1.0
+            for i1 in range(0, 4):
+                for j in range(0, 16):
+                    with T.sblock():
+                        T.reads(B[i1, j])
+                        T.writes(C[i0 * 4 + i1, j])
+                        C[i0 * 4 + i1, j] = B[i1, j] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def transformed_strided_buffer_func(
+    A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")
+) -> None:
+    # body
+    for i0 in T.serial(4):
+        B = T.alloc_buffer(
+            [4, 16],
+            "float32",
+            strides=[17, 1],
+            annotations={"buffer_allocated_addr": [], "buffer_data_alignment": 64},
+        )
+        for i1, j in T.grid(4, 16):
+            B[i1, j] = A[i0 * 4 + i1, j] + T.float32(1)
+        for i1, j in T.grid(4, 16):
+            C[i0 * 4 + i1, j] = B[i1, j] * T.float32(2)
+
+
+@T.prim_func(s_tir=True)
+def compacted_symbolic_strided_buffer_func(a: T.handle) -> None:
+    n = T.int32()
+    A = T.match_buffer(a, (1, n, 10240))
+    padded_size = T.meta_var(T.min((n + 63) // 64 * 64, 96))
+    # with T.sblock("root"):
+    for i, j, k in T.grid(((n + 63) // 64 * 4 + 7) // 8, 2, 160):
+        with T.sblock(""):
+            A_pad_shared_dyn = T.sblock_alloc_buffer(
+                (1, padded_size, 64), strides=(72 * padded_size, 72, 1), scope="shared.dyn"
+            )
+            for ax0, ax1 in T.grid(96, 64):
+                with T.sblock("A_pad_shared.dyn"):
+                    T.where(i * 128 + j * 32 + ax0 < (n + 63) // 64 * 64)
+                    A_pad_shared_dyn[0, ax0, ax1] = T.if_then_else(
+                        i * 128 + j * 32 + ax0 < n,
+                        A[0, i * 128 + j * 32 + ax0, k * 64 + ax1],
+                        T.float32(0),
+                    )
+
+
+@T.prim_func(s_tir=True)
+def transformed_symbolic_strided_buffer_func(a: T.handle):
+    n = T.int32()
+    A = T.match_buffer(a, (1, n, 10240))
+    for i, j, k in T.grid(((n + 63) // 64 * 4 + 7) // 8, 2, 160):
+        A_pad_shared_dyn = T.alloc_buffer(
+            (1, T.min((n + 63) // 64 * 64, 96), 64),
+            strides=(72 * T.min((n + 63) // 64 * 64, 96), 72, 1),
+            scope="shared.dyn",
+            annotations={"buffer_allocated_addr": [], "buffer_data_alignment": 64},
+        )
+        for ax0, ax1 in T.grid(96, 64):
+            if i * 128 + j * 32 + ax0 < (n + 63) // 64 * 64:
+                A_pad_shared_dyn[0, ax0, ax1] = T.if_then_else(
+                    i * 128 + j * 32 + ax0 < n,
+                    A[0, i * 128 + j * 32 + ax0, k * 64 + ax1],
+                    T.float32(0),
+                )
+
+
+@T.prim_func(s_tir=True)
+def annotated_loops(a: T.handle) -> None:
+    A = T.match_buffer(a, (16,), "float32")
+    for i in range(0, 16, annotations={"pragma_1": "str_value", "pragma_2": 1, "pragma_3": 0.0}):
+        A[i] = 0.0
+
+
+@T.prim_func(s_tir=True)
+def boolean_handling_before(a: T.Buffer(10, "bool"), b: T.Buffer(10, "bool")) -> None:
+    for i0 in T.serial(10):
+        with T.sblock("b"):
+            T.reads(a[i0])
+            T.writes(b[i0])
+            b[i0] = a[i0]
+
+
+@T.prim_func(s_tir=True)
+def boolean_handling_after(a: T.Buffer(10, "bool"), b: T.Buffer(10, "bool")) -> None:
+    # body
+    for i0 in T.serial(10):
+        b[i0] = a[i0]
+
+
+def test_elementwise():
+    _check(compacted_elementwise_func, transformed_elementwise_func)
+
+
+def test_gpu_workload():
+    _check(compacted_gpu_func, transformed_gpu_func)
+
+
+def test_symbolic_shape():
+    _check(compacted_symbolic_func, transformed_symbolic_func)
+
+
+def test_predicate():
+    _check(compacted_predicate_func, transformed_predicate_func)
+
+
+def test_unit_loops():
+    _check(compacted_unit_loop_func, transformed_unit_loop_func)
+
+
+def test_multi_alloc():
+    _check(compacted_multi_alloc_func, transformed_multi_alloc_func)
+
+
+def test_strided_buffer():
+    _check(compacted_strided_buffer_func, transformed_strided_buffer_func)
+
+
+def test_symbolic_strided_buffer():
+    _check(compacted_symbolic_strided_buffer_func, transformed_symbolic_strided_buffer_func)
+
+
+def test_annotated_loops():
+    mod = tvm.IRModule.from_expr(annotated_loops.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+    attr1 = mod["main"].body
+    attr2 = attr1.body
+    attr3 = attr2.body
+    assert attr1.attr_key == "pragma_1" and attr1.value == "str_value"
+    assert attr2.attr_key == "pragma_2"
+    tvm.ir.assert_structural_equal(attr2.value, tvm.tirx.IntImm("int32", 1))
+    assert attr3.attr_key == "pragma_3"
+    tvm.ir.assert_structural_equal(attr3.value, tvm.tirx.FloatImm("float32", 0.0))
+
+
+def test_annotated_block():
+    @T.prim_func(s_tir=True)
+    def annotated_block() -> None:
+        with T.sblock():
+            T.sblock_attr({"pragma_1": "str_value", "pragma_2": 1, "pragma_3": 0.0})
+            T.evaluate(0)
+
+    mod = tvm.IRModule.from_expr(annotated_block.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+    attr1 = mod["main"].body
+    attr2 = attr1.body
+    attr3 = attr2.body
+    assert attr1.attr_key == "pragma_1" and attr1.value == "str_value"
+    assert attr2.attr_key == "pragma_2"
+    tvm.ir.assert_structural_equal(attr2.value, tvm.tirx.IntImm("int32", 1))
+    assert attr3.attr_key == "pragma_3"
+    tvm.ir.assert_structural_equal(attr3.value, tvm.tirx.FloatImm("float32", 0.0))
+
+
+def test_preserved_annotations():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer(8, "float32"), B: T.Buffer(8, "float32")):
+        for i in T.serial(8, annotations={"k_0": 1, "k_1": [2, 3], "k_2": 3.14}):
+            with T.sblock("block"):
+                T.sblock_attr({"k_3": "oops"})
+                B[i] = A[i] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def after(A: T.Buffer(8, "float32"), B: T.Buffer(8, "float32")):
+        for i in T.serial(8, annotations={"k_0": 1, "k_1": [2, 3], "k_2": 3.14}):
+            B[i] = A[i] + 1.0
+
+    mod = tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+    tvm.ir.assert_structural_equal(mod["main"], after.with_attr("global_symbol", "main"))
+
+
+def test_boolean_handling():
+    _check(boolean_handling_before, boolean_handling_after)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py

@@ -0,0 +1,523 @@
+# 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, F841
+
+import tvm
+import tvm.testing
+from tvm import s_tir
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def _has_volatile_alloc_buffer(mod):
+    has_volatile_alloc = False
+
+    def visit(node):
+        nonlocal has_volatile_alloc
+        if isinstance(node, tvm.tirx.AllocBuffer) and "tirx.volatile" in node.annotations:
+            has_volatile_alloc = has_volatile_alloc or node.annotations["tirx.volatile"] is True
+
+    tvm.tirx.stmt_functor.post_order_visit(mod["main"].body, visit)
+    return has_volatile_alloc
+
+
+def test_basic():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((128, 32), "float32"), B: T.Buffer(128, "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            A_flat = T.decl_buffer(4096, data=A.data)
+
+            for i in range(128):
+                threadIdx_x = T.launch_thread("threadIdx.x", 32)
+
+                reduce = T.alloc_buffer((1,), scope="local")
+                reduce_1 = T.decl_buffer(1, data=reduce.data, scope="local")
+
+                with T.attr(
+                    T.comm_reducer(lambda x, y: x + y, [T.float32(0)]),
+                    "reduce_scope",
+                    T.int32(0),
+                ):
+                    T.tvm_thread_allreduce(
+                        T.uint32(1),
+                        A_flat[0],
+                        T.bool(True),
+                        reduce_1[0],
+                        threadIdx_x,
+                    )
+                if threadIdx_x == 0:
+                    B[i] = reduce_1[0]
+
+    After = transform(Before)
+    # Verify the transform produces valid output (not checking exact structure
+    # since the output depends on the allreduce lowering implementation)
+    assert After is not None
+    # Run script roundtrip to verify it can be printed and reparsed
+    After_script = After.script()
+    assert "tvm_warp_shuffle" in After_script
+
+
+def test_basic_with_decl_buffer():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((128, 32), "float32"), B: T.Buffer(128, "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            A_flat = T.decl_buffer(4096, data=A.data)
+
+            for i in range(128):
+                threadIdx_x = T.launch_thread("threadIdx.x", 32)
+
+                reduce = T.decl_buffer(1, dtype="float32", scope="local")
+
+                with T.attr(
+                    T.comm_reducer(lambda x, y: x + y, [T.float32(0)]),
+                    "reduce_scope",
+                    T.int32(0),
+                ):
+                    T.tvm_thread_allreduce(
+                        T.uint32(1),
+                        A_flat[0],
+                        T.bool(True),
+                        reduce[0],
+                        threadIdx_x,
+                    )
+                if threadIdx_x == 0:
+                    B[i] = reduce[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    assert "tvm_warp_shuffle" in After_script
+
+
+def test_reduce_summation():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer(128, "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            A_flat = T.decl_buffer(16384, data=A.data)
+
+            for i in range(128):
+                threadIdx_x = T.launch_thread("threadIdx.x", 32)
+
+                normal_reduce = T.alloc_buffer((1,), scope="local")
+                normal_reduce_1 = T.decl_buffer(1, data=normal_reduce.data, scope="local")
+
+                reduce = T.alloc_buffer((1,), scope="local")
+                reduce_1 = T.decl_buffer(1, data=reduce.data, scope="local")
+
+                normal_reduce_1[0] = T.float32(0)
+
+                for ko in range(4):
+                    normal_reduce_1[0] = (
+                        normal_reduce_1[0] + A_flat[i * 128 + ko * 32 + threadIdx_x]
+                    )
+
+                with T.attr(
+                    T.comm_reducer(lambda x, y: x + y, [T.float32(0)]),
+                    "reduce_scope",
+                    T.int32(0),
+                ):
+                    T.tvm_thread_allreduce(
+                        T.uint32(1),
+                        normal_reduce_1[0],
+                        T.bool(True),
+                        reduce_1[0],
+                        threadIdx_x,
+                    )
+                if threadIdx_x == 0:
+                    B[i] = reduce_1[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    assert "tvm_warp_shuffle" in After_script
+
+
+def test_multi_group_reduction():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((32, 32), "float32"), B: T.Buffer((32,), "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            threadIdx_y = T.launch_thread("threadIdx.y", 32)
+            cross_thread_B = T.alloc_buffer((1,), scope="local")
+            threadIdx_x = T.launch_thread("threadIdx.x", 32)
+            cross_thread_B_1 = T.decl_buffer((1,), data=cross_thread_B.data, scope="local")
+            with T.attr(
+                T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                "reduce_scope",
+                T.int32(0),
+            ):
+                A_1 = T.decl_buffer((1024,), data=A.data)
+                T.tvm_thread_allreduce(
+                    T.uint32(1),
+                    A_1[threadIdx_y * 32 + threadIdx_x],
+                    T.bool(True),
+                    cross_thread_B_1[0],
+                    threadIdx_x,
+                )
+            if threadIdx_x == 0:
+                B_1 = T.decl_buffer((32,), data=B.data)
+                B_1[threadIdx_y] = cross_thread_B_1[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    assert "tvm_warp_shuffle" in After_script
+
+
+def test_multi_group_mask1():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((32, 8), "float32"), B: T.Buffer((32,), "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            threadIdx_y = T.launch_thread("threadIdx.y", 32)
+            cross_thread_B = T.alloc_buffer((1,), scope="local")
+            threadIdx_x = T.launch_thread("threadIdx.x", 8)
+            cross_thread_B_1 = T.decl_buffer((1,), data=cross_thread_B.data, scope="local")
+            with T.attr(
+                T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                "reduce_scope",
+                T.int32(0),
+            ):
+                A_1 = T.decl_buffer((256,), data=A.data)
+                T.tvm_thread_allreduce(
+                    T.uint32(1),
+                    A_1[threadIdx_y * 8 + threadIdx_x],
+                    T.bool(True),
+                    cross_thread_B_1[0],
+                    threadIdx_x,
+                )
+            if threadIdx_x == 0:
+                B_1 = T.decl_buffer((32,), data=B.data)
+                B_1[threadIdx_y] = cross_thread_B_1[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    assert "tvm_warp_shuffle" in After_script
+
+
+def test_multi_warp_reduce1():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128,), "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            for i in range(128):
+                threadIdx_x = T.launch_thread("threadIdx.x", 128)
+                cross_thread_B = T.alloc_buffer((1,), scope="local")
+                cross_thread_B_1 = T.decl_buffer((1,), data=cross_thread_B.data, scope="local")
+                with T.attr(
+                    T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                    "reduce_scope",
+                    T.int32(0),
+                ):
+                    A_1 = T.decl_buffer((16384,), data=A.data)
+                    T.tvm_thread_allreduce(
+                        T.uint32(1),
+                        A_1[i * 128 + threadIdx_x],
+                        T.bool(True),
+                        cross_thread_B_1[0],
+                        threadIdx_x,
+                    )
+                if threadIdx_x == 0:
+                    B_1 = T.decl_buffer((128,), data=B.data)
+                    B_1[i] = cross_thread_B_1[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    assert "tvm_warp_shuffle" in After_script
+    assert "tvm_storage_sync" in After_script  # multi-warp needs shared sync
+
+
+def test_multi_warp_reduce2():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((1, 1024), "float32"), B: T.Buffer((1,), "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            threadIdx_x = T.launch_thread("threadIdx.x", 1024)
+            cross_thread_B = T.alloc_buffer((1,), scope="local")
+            cross_thread_B_1 = T.decl_buffer((1,), data=cross_thread_B.data, scope="local")
+            with T.attr(
+                T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                "reduce_scope",
+                T.int32(0),
+            ):
+                A_1 = T.decl_buffer((1024,), data=A.data)
+                T.tvm_thread_allreduce(
+                    T.uint32(1), A_1[threadIdx_x], T.bool(True), cross_thread_B_1[0], threadIdx_x
+                )
+            if threadIdx_x == 0:
+                B_1 = T.decl_buffer((1,), data=B.data)
+                B_1[0] = cross_thread_B_1[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    assert "tvm_warp_shuffle" in After_script
+    assert "tvm_storage_sync" in After_script
+
+
+def test_multi_group_multi_warp_reduction():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((4, 128), "float32"), B: T.Buffer((4,), "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            threadIdx_y = T.launch_thread("threadIdx.y", 4)
+            cross_thread_B = T.alloc_buffer((1,), scope="local")
+            threadIdx_x = T.launch_thread("threadIdx.x", 128)
+            cross_thread_B_1 = T.decl_buffer((1,), data=cross_thread_B.data, scope="local")
+            with T.attr(
+                T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                "reduce_scope",
+                T.int32(0),
+            ):
+                A_1 = T.decl_buffer((512,), data=A.data)
+                T.tvm_thread_allreduce(
+                    T.uint32(1),
+                    A_1[threadIdx_y * 128 + threadIdx_x],
+                    T.bool(True),
+                    cross_thread_B_1[0],
+                    threadIdx_x,
+                )
+            if threadIdx_x == 0:
+                B_1 = T.decl_buffer((4,), data=B.data)
+                B_1[threadIdx_y] = cross_thread_B_1[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    assert "tvm_warp_shuffle" in After_script
+    assert "tvm_storage_sync" in After_script
+
+
+def test_multi_group_multi_warp_predicated_reduction():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((2, 70), "float32"), B: T.Buffer((2,), "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            threadIdx_y = T.launch_thread("threadIdx.y", 2)
+            in_thread_B = T.alloc_buffer((1,), scope="local")
+            cross_thread_B = T.alloc_buffer((1,), scope="local")
+            threadIdx_x = T.launch_thread("threadIdx.x", 512)
+            in_thread_B_1 = T.decl_buffer((1,), data=in_thread_B.data, scope="local")
+            in_thread_B_1[0] = T.float32(0)
+            if threadIdx_x < 70:
+                A_1 = T.decl_buffer((140,), data=A.data)
+                in_thread_B_1[0] = in_thread_B_1[0] + A_1[threadIdx_y * 70 + threadIdx_x]
+            cross_thread_B_1 = T.decl_buffer((1,), data=cross_thread_B.data, scope="local")
+            with T.attr(
+                T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                "reduce_scope",
+                T.int32(0),
+            ):
+                T.tvm_thread_allreduce(
+                    T.uint32(1), in_thread_B_1[0], T.bool(True), cross_thread_B_1[0], threadIdx_x
+                )
+            if threadIdx_x == 0:
+                B_1 = T.decl_buffer((2,), data=B.data)
+                B_1[threadIdx_y] = cross_thread_B_1[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    assert "tvm_warp_shuffle" in After_script
+    assert "tvm_storage_sync" in After_script
+
+
+def test_metal_no_mask():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((1, 1, 2, 128), "float32"), B: T.Buffer((1, 1, 2), "float32")):
+            T.func_attr(
+                {
+                    "target": T.target(
+                        {
+                            "kind": "metal",
+                            "max_threads_per_block": 1024,
+                            "thread_warp_size": 32,
+                            "host": "llvm",
+                        }
+                    ),
+                }
+            )
+            blockIdx_x = T.launch_thread("blockIdx.x", 1)
+            cross_thread_B = T.alloc_buffer((1,), scope="local")
+            threadIdx_z = T.launch_thread("threadIdx.z", 1)
+            threadIdx_y = T.launch_thread("threadIdx.y", 2)
+            threadIdx_x = T.launch_thread("threadIdx.x", 128)
+            cross_thread_B_1 = T.decl_buffer((1,), data=cross_thread_B.data, scope="local")
+            with T.attr(
+                T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                "reduce_scope",
+                T.int32(0),
+            ):
+                A_1 = T.decl_buffer((256,), data=A.data)
+                T.tvm_thread_allreduce(
+                    T.uint32(1),
+                    A_1[threadIdx_y * 128 + threadIdx_x],
+                    T.bool(True),
+                    cross_thread_B_1[0],
+                    threadIdx_x,
+                )
+            if threadIdx_x == 0:
+                B_1 = T.decl_buffer((2,), data=B.data)
+                B_1[threadIdx_y] = cross_thread_B_1[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    # Metal does not use warp masks
+    assert "tvm_warp_shuffle_down" in After_script
+    assert "tvm_storage_sync" in After_script
+
+
+def test_webgpu_warp_reduce():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((128, 32), "float32"), B: T.Buffer(128, "float32")):
+            T.func_attr(
+                {
+                    "target": T.target(
+                        {
+                            "kind": "webgpu",
+                            "supports_subgroups": True,
+                            "host": "llvm",
+                        }
+                    ),
+                }
+            )
+            A_flat = T.decl_buffer(4096, data=A.data)
+
+            for i in range(128):
+                threadIdx_x = T.launch_thread("threadIdx.x", 32)
+
+                reduce_data = T.alloc_buffer((1,), "float32", scope="local")
+                reduce = T.decl_buffer(1, data=reduce_data.data, scope="local")
+
+                with T.attr(
+                    T.comm_reducer(lambda x, y: x + y, [T.float32(0)]),
+                    "reduce_scope",
+                    T.int32(0),
+                ):
+                    T.tvm_thread_allreduce(
+                        T.uint32(1),
+                        A_flat[0],
+                        T.bool(True),
+                        reduce[0],
+                        threadIdx_x,
+                    )
+                if threadIdx_x == 0:
+                    B[i] = reduce[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    assert "tvm_warp_shuffle_down" in After_script
+    assert "tvm_warp_shuffle(" in After_script
+    assert "tvm_storage_sync" not in After_script
+    assert "T.uint32(" not in After_script
+
+
+def test_webgpu_multi_warp_reduce():
+    transform = tvm.s_tir.transform.LowerThreadAllreduce()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((1, 1, 2, 128), "float32"), B: T.Buffer((1, 1, 2), "float32")):
+            T.func_attr(
+                {
+                    "target": T.target(
+                        {
+                            "kind": "webgpu",
+                            "max_num_threads": 1024,
+                            "supports_subgroups": True,
+                            "host": "llvm",
+                        }
+                    ),
+                }
+            )
+            blockIdx_x = T.launch_thread("blockIdx.x", 1)
+            cross_thread_B = T.alloc_buffer((1,), "float32", scope="local")
+            threadIdx_z = T.launch_thread("threadIdx.z", 1)
+            threadIdx_y = T.launch_thread("threadIdx.y", 2)
+            threadIdx_x = T.launch_thread("threadIdx.x", 128)
+            cross_thread_B_1 = T.decl_buffer((1,), data=cross_thread_B.data, scope="local")
+            with T.attr(
+                T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                "reduce_scope",
+                T.int32(0),
+            ):
+                A_1 = T.decl_buffer((256,), data=A.data)
+                T.tvm_thread_allreduce(
+                    T.uint32(1),
+                    A_1[threadIdx_y * 128 + threadIdx_x],
+                    T.bool(True),
+                    cross_thread_B_1[0],
+                    threadIdx_x,
+                )
+            if threadIdx_x == 0:
+                B_1 = T.decl_buffer((2,), data=B.data)
+                B_1[threadIdx_y] = cross_thread_B_1[0]
+
+    After = transform(Before)
+    assert After is not None
+    After_script = After.script()
+    assert "tvm_warp_shuffle_down" in After_script
+    assert "tvm_storage_sync" in After_script
+    assert _has_volatile_alloc_buffer(After)
+    assert "T.uint32(" not in After_script
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_manifest_shared_memory_local_stage.py

@@ -0,0 +1,135 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501, F401
+import tvm
+import tvm.testing
+from tvm import s_tir
+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
+
+
+@tvm.script.ir_module
+class MatmulBefore:
+    @T.prim_func(s_tir=True)
+    def main(A: T.Buffer((1024, 1024), "float32"), B: T.Buffer((1024, 1024), "float32"), C: T.Buffer((1024, 1024), "float32")) -> None:
+        # function attr dict
+        T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+        # body
+        # with T.sblock("root")
+        for blockIdx_y in T.thread_binding(32, thread="blockIdx.y"):
+            for blockIdx_x in T.thread_binding(32, thread="blockIdx.x"):
+                for threadIdx_y in T.thread_binding(2, thread="threadIdx.y"):
+                    for threadIdx_x in T.thread_binding(2, thread="threadIdx.x"):
+                        for k_0 in T.serial(32):
+                            with T.sblock():
+                                T.reads(A[blockIdx_y * 32 : blockIdx_y * 32 + 32, k_0 * 32 : k_0 * 32 + 32], B[k_0 * 32 : k_0 * 32 + 32, blockIdx_x * 32 : blockIdx_x * 32 + 32])
+                                T.writes(C[blockIdx_y * 32 : blockIdx_y * 32 + 32, blockIdx_x * 32 : blockIdx_x * 32 + 32])
+                                A_shared = T.sblock_alloc_buffer([1024, 1024], dtype="float32", scope="shared")
+                                B_shared = T.sblock_alloc_buffer([1024, 1024], dtype="float32", scope="shared")
+                                for ax0_ax1_fused_0 in T.serial(64):
+                                    for ax0_ax1_fused_3 in T.vectorized(4):
+                                        with T.sblock("A_shared"):
+                                            T.reads(A[blockIdx_y * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32])
+                                            T.writes(A_shared[blockIdx_y * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32])
+                                            T.sblock_attr({"tirx.manifest_shared_memory_local_stage":1})
+                                            A_shared[blockIdx_y * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32] = A[blockIdx_y * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32]
+                                for ax0_ax1_fused_0 in T.serial(64):
+                                    for ax0_ax1_fused_3 in T.vectorized(4):
+                                        with T.sblock("B_shared"):
+                                            T.reads(B[k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, blockIdx_x * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32])
+                                            T.writes(B_shared[k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, blockIdx_x * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32])
+                                            T.sblock_attr({"tirx.manifest_shared_memory_local_stage":1})
+                                            B_shared[k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, blockIdx_x * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32] = B[k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, blockIdx_x * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32]
+                                for k_1, i_2, j_2, k_2 in T.grid(2, 16, 16, 16):
+                                    with T.sblock("C"):
+                                        T.reads(A_shared[blockIdx_y * 32 + threadIdx_y * 16 + i_2, k_0 * 32 + k_1 * 16 + k_2], B_shared[k_0 * 32 + k_1 * 16 + k_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2])
+                                        T.writes(C[blockIdx_y * 32 + threadIdx_y * 16 + i_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2])
+                                        if k_0 * 32 + k_1 * 16 + k_2 == 0:
+                                            C[blockIdx_y * 32 + threadIdx_y * 16 + i_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2] = T.float32(0)
+                                        C[blockIdx_y * 32 + threadIdx_y * 16 + i_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2] = C[blockIdx_y * 32 + threadIdx_y * 16 + i_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2] + A_shared[blockIdx_y * 32 + threadIdx_y * 16 + i_2, k_0 * 32 + k_1 * 16 + k_2] * B_shared[k_0 * 32 + k_1 * 16 + k_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2]
+
+
+@tvm.script.ir_module
+class MatmulAfter:
+    @T.prim_func(s_tir=True)
+    def main(A: T.Buffer((1024, 1024), "float32"), B: T.Buffer((1024, 1024), "float32"), C: T.Buffer((1024, 1024), "float32")) -> None:
+        # function attr dict
+        T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+        # body
+        # with T.sblock("root")
+        for blockIdx_y in T.thread_binding(32, thread="blockIdx.y"):
+            for blockIdx_x in T.thread_binding(32, thread="blockIdx.x"):
+                for threadIdx_y in T.thread_binding(2, thread="threadIdx.y"):
+                    for threadIdx_x in T.thread_binding(2, thread="threadIdx.x"):
+                        for k_0 in T.serial(32):
+                            with T.sblock():
+                                T.reads(A[blockIdx_y * 32 : blockIdx_y * 32 + 32, k_0 * 32 : k_0 * 32 + 32], B[k_0 * 32 : k_0 * 32 + 32, blockIdx_x * 32 : blockIdx_x * 32 + 32])
+                                T.writes(C[blockIdx_y * 32 : blockIdx_y * 32 + 32, blockIdx_x * 32 : blockIdx_x * 32 + 32])
+                                A_shared = T.sblock_alloc_buffer([1024, 1024], dtype="float32", scope="shared")
+                                B_shared = T.sblock_alloc_buffer([1024, 1024], dtype="float32", scope="shared")
+                                A_shared_local = T.sblock_alloc_buffer([64, 4], dtype="float32", scope="local")
+                                B_shared_local = T.sblock_alloc_buffer([64, 4], dtype="float32", scope="local")
+                                for ax0_ax1_fused_0 in T.serial(64):
+                                    for ax0_ax1_fused_3 in T.vectorized(4):
+                                        with T.sblock():
+                                            T.reads(A[blockIdx_y * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32])
+                                            T.writes(A_shared_local[ax0_ax1_fused_0, ax0_ax1_fused_3])
+                                            A_shared_local[ax0_ax1_fused_0, ax0_ax1_fused_3] = A[blockIdx_y * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32]
+                                for ax0_ax1_fused_0 in T.serial(64):
+                                    for ax0_ax1_fused_3 in T.vectorized(4):
+                                        with T.sblock("A_shared"):
+                                            T.reads(A[blockIdx_y * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32])
+                                            T.writes(A_shared[blockIdx_y * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32])
+                                            A_shared[blockIdx_y * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32] = A_shared_local[ax0_ax1_fused_0, ax0_ax1_fused_3]
+                                for ax0_ax1_fused_0 in T.serial(64):
+                                    for ax0_ax1_fused_3 in T.vectorized(4):
+                                        with T.sblock():
+                                            T.reads(B[k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, blockIdx_x * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32])
+                                            T.writes(B_shared_local[ax0_ax1_fused_0, ax0_ax1_fused_3])
+                                            B_shared_local[ax0_ax1_fused_0, ax0_ax1_fused_3] = B[k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, blockIdx_x * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32]
+                                for ax0_ax1_fused_0 in T.serial(64):
+                                    for ax0_ax1_fused_3 in T.vectorized(4):
+                                        with T.sblock("B_shared"):
+                                            T.reads(B[k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, blockIdx_x * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32])
+                                            T.writes(B_shared[k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, blockIdx_x * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32])
+                                            B_shared[k_0 * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) // 32, blockIdx_x * 32 + (ax0_ax1_fused_0 * 16 + threadIdx_y * 8 + threadIdx_x * 4 + ax0_ax1_fused_3) % 32] = B_shared_local[ax0_ax1_fused_0, ax0_ax1_fused_3]
+                                for k_1, i_2, j_2, k_2 in T.grid(2, 16, 16, 16):
+                                    with T.sblock("C"):
+                                        T.reads(A_shared[blockIdx_y * 32 + threadIdx_y * 16 + i_2, k_0 * 32 + k_1 * 16 + k_2], B_shared[k_0 * 32 + k_1 * 16 + k_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2])
+                                        T.writes(C[blockIdx_y * 32 + threadIdx_y * 16 + i_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2])
+                                        if k_0 * 32 + k_1 * 16 + k_2 == 0:
+                                            C[blockIdx_y * 32 + threadIdx_y * 16 + i_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2] = T.float32(0)
+                                        C[blockIdx_y * 32 + threadIdx_y * 16 + i_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2] = C[blockIdx_y * 32 + threadIdx_y * 16 + i_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2] + A_shared[blockIdx_y * 32 + threadIdx_y * 16 + i_2, k_0 * 32 + k_1 * 16 + k_2] * B_shared[k_0 * 32 + k_1 * 16 + k_2, blockIdx_x * 32 + threadIdx_x * 16 + j_2]
+
+
+# fmt: on
+# pylint: enable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks
+
+
+def _check(before, expected):
+    after = tvm.s_tir.transform.ManifestSharedMemoryLocalStage()(before)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_transform_matmul():
+    _check(MatmulBefore, MatmulAfter)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py

@@ -0,0 +1,349 @@
+# 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, F841
+import numpy as np
+
+import tvm
+import tvm.testing
+from tvm import s_tir
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.topi.math import cast
+
+
+def test_matmul_t_buffer():
+    """Shared allocations should be merged, preserving DeclBuffer if present
+
+    This test uses a matmul PrimFunc adapted from
+    test_matmul_dyn_shared, using `T.Buffer` (Allocate without
+    DeclBuffer) for the replaced allocations.
+    """
+    transform = tvm.s_tir.transform.MergeSharedMemoryAllocations()
+    buffer_func = T.Buffer
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((1024, 1024), "float16"),
+            B: T.Buffer((1024, 1024), "float16"),
+            matmul: T.Buffer((1024, 1024), "float32"),
+        ):
+            A_flat = T.decl_buffer(1048576, "float16", data=A.data)
+            B_flat = T.decl_buffer(1048576, "float16", data=B.data)
+            matmul_flat = T.decl_buffer(1048576, data=matmul.data)
+
+            threadIdx_x = T.launch_thread("threadIdx.x", 16)
+            C_local = T.alloc_buffer((1,), "float32", scope="local")
+            A_sh = T.alloc_buffer((256,), "float16", scope="shared.dyn")
+            B_sh = T.alloc_buffer((256,), "float16", scope="shared.dyn")
+            C_sh = T.alloc_buffer((256,), "float32", scope="shared.dyn")
+            threadIdx_y = T.launch_thread("threadIdx.y", 16)
+            blockIdx_x = T.launch_thread("blockIdx.x", 64)
+            blockIdx_y = T.launch_thread("blockIdx.y", 64)
+
+            C_local[0] = T.float32(0)
+            for i in range(64):
+                A_sh[threadIdx_y * 16 + threadIdx_x] = A_flat[
+                    blockIdx_y * 16384 + threadIdx_y * 1024 + i * 16 + threadIdx_x
+                ]
+
+                B_sh[threadIdx_y * 16 + threadIdx_x] = B_flat[
+                    i * 16384 + threadIdx_y * 1024 + blockIdx_x * 16 + threadIdx_x
+                ]
+                T.tvm_storage_sync("shared")
+                for k in range(16):
+                    C_local[0] = C_local[0] + T.Cast(
+                        "float32",
+                        A_sh[threadIdx_y * 16 + k] * B_sh[k * 16 + threadIdx_x],
+                    )
+                T.tvm_storage_sync("shared")
+
+            C_sh[threadIdx_y * 16 + threadIdx_x] = C_local[0]
+            T.tvm_storage_sync("shared.dyn")
+
+            matmul_flat[blockIdx_y * 16384 + threadIdx_y * 1024 + blockIdx_x * 16 + threadIdx_x] = (
+                C_sh[threadIdx_y * 16 + threadIdx_x]
+            )
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((1024, 1024), "float16"),
+            B: T.Buffer((1024, 1024), "float16"),
+            matmul: T.Buffer((1024, 1024), "float32"),
+        ):
+            A_flat = T.decl_buffer(1048576, "float16", data=A.data)
+            B_flat = T.decl_buffer(1048576, "float16", data=B.data)
+            matmul_flat = T.decl_buffer(1048576, data=matmul.data)
+
+            threadIdx_x = T.launch_thread("threadIdx.x", 16)
+
+            buf_dyn_shmem = T.alloc_buffer((1024,), "uint8", scope="shared.dyn")
+
+            C_local = T.alloc_buffer((1,), "float32", scope="local")
+            A_sh = T.decl_buffer(256, "float16", data=buf_dyn_shmem.data, scope="shared.dyn")
+            B_sh = T.decl_buffer(256, "float16", data=buf_dyn_shmem.data, scope="shared.dyn")
+            C_sh = T.decl_buffer(256, "float32", data=buf_dyn_shmem.data, scope="shared.dyn")
+
+            threadIdx_y = T.launch_thread("threadIdx.y", 16)
+            blockIdx_x = T.launch_thread("blockIdx.x", 64)
+            blockIdx_y = T.launch_thread("blockIdx.y", 64)
+
+            C_local[0] = T.float32(0)
+            for i in range(64):
+                A_sh[threadIdx_y * 16 + threadIdx_x + 256] = A_flat[
+                    blockIdx_y * 16384 + threadIdx_y * 1024 + i * 16 + threadIdx_x
+                ]
+                B_sh[threadIdx_y * 16 + threadIdx_x] = B_flat[
+                    i * 16384 + threadIdx_y * 1024 + blockIdx_x * 16 + threadIdx_x
+                ]
+                T.tvm_storage_sync("shared")
+                for k in range(16):
+                    C_local[0] = C_local[0] + T.Cast(
+                        "float32",
+                        A_sh[threadIdx_y * 16 + k + 256] * B_sh[k * 16 + threadIdx_x],
+                    )
+                T.tvm_storage_sync("shared")
+
+            C_sh[threadIdx_y * 16 + threadIdx_x] = C_local[0]
+            T.tvm_storage_sync("shared.dyn")
+
+    After = transform(Before)
+    script = After["main"].script()
+    # Verify merged allocation: one shared.dyn buffer of 1024 bytes (256*2 float16 + 256 float32)
+    assert "alloc_buffer((1024,)" in script
+    assert '"uint8"' in script
+    assert '"shared.dyn"' in script
+    # Verify storage sync calls preserved
+    assert "tvm_storage_sync" in script
+    # Verify offset indexing (shared memory merged)
+    assert "+ 256" in script
+
+
+def test_matmul_decl_buffer():
+    """Shared allocations should be merged, preserving DeclBuffer if present
+
+    This test uses a matmul PrimFunc adapted from
+    test_matmul_dyn_shared, using `T.decl_buffer` (Allocate followed by DeclBuffer)
+    for the replaced allocations.
+    """
+    transform = tvm.s_tir.transform.MergeSharedMemoryAllocations()
+    buffer_func = T.decl_buffer
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((1024, 1024), "float16"),
+            B: T.Buffer((1024, 1024), "float16"),
+            matmul: T.Buffer((1024, 1024), "float32"),
+        ):
+            A_flat = T.decl_buffer(1048576, "float16", data=A.data)
+            B_flat = T.decl_buffer(1048576, "float16", data=B.data)
+            matmul_flat = T.decl_buffer(1048576, data=matmul.data)
+
+            threadIdx_x = T.launch_thread("threadIdx.x", 16)
+            C_local = T.alloc_buffer((1,), "float32", scope="local")
+            A_sh = T.alloc_buffer((256,), "float16", scope="shared.dyn")
+            B_sh = T.alloc_buffer((256,), "float16", scope="shared.dyn")
+            C_sh = T.alloc_buffer((256,), "float32", scope="shared.dyn")
+            threadIdx_y = T.launch_thread("threadIdx.y", 16)
+            blockIdx_x = T.launch_thread("blockIdx.x", 64)
+            blockIdx_y = T.launch_thread("blockIdx.y", 64)
+
+            C_local[0] = T.float32(0)
+            for i in range(64):
+                A_sh[threadIdx_y * 16 + threadIdx_x] = A_flat[
+                    blockIdx_y * 16384 + threadIdx_y * 1024 + i * 16 + threadIdx_x
+                ]
+
+                B_sh[threadIdx_y * 16 + threadIdx_x] = B_flat[
+                    i * 16384 + threadIdx_y * 1024 + blockIdx_x * 16 + threadIdx_x
+                ]
+                T.tvm_storage_sync("shared")
+                for k in range(16):
+                    C_local[0] = C_local[0] + T.Cast(
+                        "float32",
+                        A_sh[threadIdx_y * 16 + k] * B_sh[k * 16 + threadIdx_x],
+                    )
+                T.tvm_storage_sync("shared")
+
+            C_sh[threadIdx_y * 16 + threadIdx_x] = C_local[0]
+            T.tvm_storage_sync("shared.dyn")
+
+            matmul_flat[blockIdx_y * 16384 + threadIdx_y * 1024 + blockIdx_x * 16 + threadIdx_x] = (
+                C_sh[threadIdx_y * 16 + threadIdx_x]
+            )
+
+    After = transform(Before)
+    script = After["main"].script()
+    # Verify merged allocation: one shared.dyn buffer of 1024 bytes
+    assert "alloc_buffer((1024,)" in script
+    assert '"uint8"' in script
+    assert '"shared.dyn"' in script
+    assert "tvm_storage_sync" in script
+    assert "+ 256" in script
+
+
+def test_simple_alloc_no_reuse():
+    """Test alloc and free within the same scope."""
+    transform = tvm.s_tir.transform.MergeSharedMemoryAllocations()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            threadIdx_x = T.launch_thread("threadIdx.x", 128)
+            A_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn")
+            B_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn")
+            B_sh[threadIdx_x] = A_sh[threadIdx_x]
+
+    After = transform(Before)
+    script = After["main"].script()
+    # Verify merged allocation: 1024 bytes (128*4 + 128*4)
+    assert "alloc_buffer((1024,)" in script
+    assert '"uint8"' in script
+    assert '"shared.dyn"' in script
+    # Verify offset indexing
+    assert "+ 128" in script
+
+
+def test_simple_alloc_reuse():
+    """Test alloc and free within the same scope with a reuse chance."""
+    transform = tvm.s_tir.transform.MergeSharedMemoryAllocations()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            threadIdx_x = T.launch_thread("threadIdx.x", 128)
+            A_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn")
+            B_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn")
+            A_sh[threadIdx_x] = 0
+            B_sh[threadIdx_x] = 0
+
+    After = transform(Before)
+    script = After["main"].script()
+    # Verify merged allocation: 512 bytes (128*4, reusable)
+    assert "alloc_buffer((512,)" in script
+    assert '"uint8"' in script
+    assert '"shared.dyn"' in script
+
+
+def test_async_copy():
+    """Test async copy in shared memory."""
+    transform = tvm.s_tir.transform.MergeSharedMemoryAllocations()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+            threadIdx_x = T.launch_thread("threadIdx.x", 128)
+            A_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn")
+            B_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn")
+            T.ptx.cp_async("float32", A_sh.data, threadIdx_x, A.data, threadIdx_x, 512)
+            T.ptx.cp_async("float32", B_sh.data, threadIdx_x, B.data, threadIdx_x, 512)
+
+    After = transform(Before)
+    # The pass merges shared.dyn allocations. A_sh and B_sh are accessed
+    # sequentially inside the thread_extent with non-overlapping lifetimes,
+    # so the liveness analysis allows reuse — both fit in 512 bytes
+    # (= 128 elements * 4 bytes).
+    script = After["main"].script()
+    # Verify merged allocation (512 bytes - A_sh and B_sh can be reused)
+    assert '"uint8"' in script and '"shared.dyn"' in script and "(512,)" in script
+    # Verify cp_async uses the merged buffer
+    assert "buf_dyn_shmem" in script
+    assert "threadIdx_x * 4" in script
+
+
+def test_multi_thread_extent_blocks():
+    """Each thread_extent block must get its own merged buffer.
+
+    Reproduces the scoping bug from PR #19605: a single PrimFunc
+    with two sibling thread_extent regions, each containing its
+    own shared.dyn allocations. The merged buffer must be allocated
+    inside each kernel body — not just the first.
+    """
+    transform = tvm.s_tir.transform.MergeSharedMemoryAllocations()
+
+    @I.ir_module(check_well_formed=False)
+    class Before:
+        @T.prim_func(s_tir=True, check_well_formed=False)
+        def main(
+            X: T.Buffer((128,), "float32"),
+            Y: T.Buffer((128,), "float32"),
+        ):
+            X_flat = T.decl_buffer(128, data=X.data)
+            Y_flat = T.decl_buffer(128, data=Y.data)
+
+            # First kernel launch
+            tx0 = T.env_thread("threadIdx.x")
+            with T.attr(tx0, "thread_extent", 128):
+                A_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn")
+                B_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn")
+                A_sh[tx0] = X_flat[tx0]
+                B_sh[tx0] = A_sh[tx0]
+                X_flat[tx0] = B_sh[tx0]
+
+            # Second kernel launch — must NOT see kernel #0's merged buffer.
+            tx1 = T.env_thread("threadIdx.x")
+            with T.attr(tx1, "thread_extent", 128):
+                C_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn")
+                D_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn")
+                C_sh[tx1] = Y_flat[tx1]
+                D_sh[tx1] = C_sh[tx1]
+                Y_flat[tx1] = D_sh[tx1]
+
+    After = transform(Before)
+    script = After["main"].script()
+
+    # Two merged allocations — one per thread_extent body.
+    # Each of the four original 128-float32 buffers (A_sh, B_sh, C_sh, D_sh)
+    # gets merged within its own kernel scope.
+    assert script.count("shared.dyn") >= 2, (
+        "Expected at least two shared.dyn allocations (one per kernel)"
+    )
+    assert script.count("alloc_buffer") >= 2, (
+        "Expected at least two alloc_buffer nodes (one merged buf per kernel)"
+    )
+
+    # Both thread_extent blocks must contain their own merged buffer —
+    # they must NOT share the same buf_dyn_shmem variable.
+    # Structurally verify that the first kernel's body accesses are
+    # not rewritten to the second kernel's buf_dyn_shmem (and vice versa).
+    first_block = script.split("with T.attr(tx1")[0]
+    second_block = script.split("with T.attr(tx1")[1] if "tx1" in script else ""
+    assert "buf_dyn_shmem" in first_block, "Kernel 1 must have a merged buffer"
+    if second_block:
+        assert "buf_dyn_shmem" in second_block, "Kernel 2 must have a merged buffer"
+
+    # End-to-end: post-merge IR must remain well-formed through
+    # the host/device split — this is the exact ordering from
+    # PR #19605 that triggers the scoping bug.
+    target = tvm.target.Target("llvm")
+    mod_with_target = tvm.IRModule({"main": After["main"].with_attr({"target": target})})
+    split = tvm.tirx.transform.SplitHostDevice()
+    # If kernel #1 referenced an undefined buf_dyn_shmem, this
+    # would raise during well-formedness checking inside SplitHostDevice.
+    split(mod_with_target)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_plan_update_buffer_allocation_location.py

@@ -0,0 +1,460 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+import numpy as np
+
+import tvm
+import tvm.testing
+from tvm import s_tir, tirx
+from tvm.s_tir.tensor_intrin.hexagon import VRMPY_u8u8i32_INTRIN
+from tvm.script import tirx as T
+
+
+def _check(original, transformed):
+    func = original
+    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.PlanAndUpdateBufferAllocationLocation()(mod)
+    tvm.ir.assert_structural_equal(mod["main"], transformed.with_attr("global_symbol", "main"))
+
+
+@T.prim_func(s_tir=True)
+def element_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16))
+    C = T.match_buffer(c, (16, 16))
+    B = T.sblock_alloc_buffer((16, 16))
+    for i0 in range(0, 16):
+        for j0 in range(0, 16):
+            with T.sblock():
+                i, j = T.axis.remap("SS", [i0, j0])
+                B[i, j] = A[i, j] + 1.0
+        for j0 in range(0, 16):
+            with T.sblock():
+                i, j = T.axis.remap("SS", [i0, j0])
+                C[i, j] = B[i, j] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def transformed_element_func(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [16, 16])
+    C = T.match_buffer(c, [16, 16])
+
+    for i_0 in range(0, 16):
+        with T.sblock():
+            T.reads([A[i_0, 0:16]])
+            T.writes([C[i_0, 0:16]])
+            B = T.sblock_alloc_buffer([16, 16])
+            for j_0 in T.serial(0, 16):
+                with T.sblock():
+                    i, j = T.axis.remap("SS", [i_0, j_0])
+                    B[i, j] = A[i, j] + 1.0
+            for j_0 in T.serial(0, 16):
+                with T.sblock():
+                    i, j = T.axis.remap("SS", [i_0, j_0])
+                    C[i, j] = B[i, j] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def original_func() -> None:
+    A = T.sblock_alloc_buffer((128, 128), "float32")
+    for i0, j0 in T.grid(128, 128):
+        with T.sblock():
+            i, j = T.axis.remap("SS", [i0, j0])
+            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])
+            B = T.sblock_alloc_buffer((128, 128), "float32")
+            C = T.sblock_alloc_buffer((128, 128), "float32")
+            D = T.sblock_alloc_buffer((128, 128), "float32")
+            if k == 0:
+                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 transformed_func() -> None:
+    A = T.sblock_alloc_buffer([128, 128])
+    for i0, j0 in T.grid(128, 128):
+        with T.sblock():
+            i, j = T.axis.remap("SS", [i0, j0])
+            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])
+            B = T.sblock_alloc_buffer([128, 128])
+            if k == 0:
+                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):
+                with T.sblock(""):
+                    T.reads([B[((i * 4) + ii), ((j * 4) + jj)]])
+                    T.writes([B[((i * 4) + ii), ((j * 4) + jj)]])
+                    C = T.sblock_alloc_buffer([128, 128])
+                    for kk in T.serial(0, 4):
+                        B[((i * 4) + ii), ((j * 4) + jj)] = (
+                            B[((i * 4) + ii), ((j * 4) + jj)] + C[((i * 4) + ii), ((k * 4) + kk)]
+                        )
+                    for kk in T.serial(0, 4):
+                        with T.sblock(""):
+                            T.reads(
+                                [
+                                    B[((i * 4) + ii), ((j * 4) + jj)],
+                                    C[((i * 4) + ii), ((k * 4) + kk)],
+                                ]
+                            )
+                            T.writes([B[((i * 4) + ii), ((j * 4) + jj)]])
+                            D = T.sblock_alloc_buffer([128, 128])
+                            B[((i * 4) + ii), ((j * 4) + jj)] = 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 match_buffer_func() -> None:
+    C = T.sblock_alloc_buffer((128, 128))
+    for i in range(128):
+        with T.sblock():
+            vi = T.axis.S(128, i)
+            C0 = T.match_buffer(C[vi, 0:128], (128))
+            for j in range(128):
+                with T.sblock():
+                    jj = T.axis.S(128, j)
+                    C1 = T.match_buffer(C0[jj], ())
+                    C1[()] = 0
+
+
+@T.prim_func(s_tir=True)
+def transformed_match_buffer_func() -> None:
+    for i in range(0, 128):
+        with T.sblock():
+            vi = T.axis.S(128, i)
+            C = T.sblock_alloc_buffer((128, 128))
+            C0 = T.match_buffer(C[vi, 0:128], (128))
+            for j in range(128):
+                with T.sblock():
+                    jj = T.axis.S(128, j)
+                    C1 = T.match_buffer(C0[jj], ())
+                    C1[()] = 0
+
+
+@T.prim_func(s_tir=True)
+def opaque_access(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [1024])
+    B = T.match_buffer(b, [1024])
+    A_cache = T.sblock_alloc_buffer([1024])
+    for i in T.serial(0, 8):
+        with T.sblock():
+            vi = T.axis.S(8, i)
+            with T.sblock():
+                v = T.axis.S(8, vi)
+                T.reads([A[(v * 128) : ((v * 128) + 128)]])
+                T.writes([A_cache[(v * 128) : ((v * 128) + 128)]])
+                T.evaluate(
+                    T.call_extern(
+                        "test",
+                        A_cache.data,
+                        (v * 128),
+                        128,
+                        A.data,
+                        (v * 128),
+                        128,
+                        dtype="float32",
+                    )
+                )
+            for j in T.serial(0, 128):
+                with T.sblock():
+                    v = T.axis.S(1024, vi * 128 + j)
+                    T.reads([A_cache[v]])
+                    T.writes([B[v]])
+                    B[v] = A_cache[v]
+
+
+@T.prim_func(s_tir=True)
+def transformed_opaque_access(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [1024])
+    B = T.match_buffer(b, [1024])
+    for i in T.serial(0, 8):
+        with T.sblock():
+            vi = T.axis.S(8, i)
+            T.reads(A[vi * 128 : vi * 128 + 128])
+            T.writes(B[vi * 128 : vi * 128 + 128])
+            A_cache = T.sblock_alloc_buffer([1024])
+            with T.sblock():
+                v = T.axis.S(8, vi)
+                T.reads([A[v * 128 : v * 128 + 128]])
+                T.writes([A_cache[v * 128 : v * 128 + 128]])
+                T.evaluate(
+                    T.call_extern(
+                        "test", A_cache.data, v * 128, 128, A.data, v * 128, 128, dtype="float32"
+                    )
+                )
+            for j in T.serial(0, 128):
+                with T.sblock():
+                    v = T.axis.S(1024, vi * 128 + j)
+                    T.reads([A_cache[v]])
+                    T.writes([B[v]])
+                    B[v] = A_cache[v]
+
+
+def test_elementwise():
+    _check(element_func, transformed_element_func)
+
+
+def test_locate_buffer_allocation():
+    _check(original_func, transformed_func)
+
+
+def test_match_buffer_allocation():
+    _check(match_buffer_func, transformed_match_buffer_func)
+
+
+def test_opaque_access():
+    _check(opaque_access, transformed_opaque_access)
+
+
+def test_loop_carried_dependency():
+    """The buffer allocation should be above opaque iter var's loop scopes
+    such that buffer accesses with loop carried dependencies are covered,
+    and the allocate buffer should keep the order."""
+
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((8, 8, 8), "int32"), B: T.Buffer((8, 8, 8), "int32")):
+        C = T.sblock_alloc_buffer([8, 8, 8], dtype="int32")
+        D = T.sblock_alloc_buffer([8, 8, 8], dtype="int32")
+        for i in T.serial(8):
+            for j in T.serial(8):
+                for k in T.serial(8):
+                    with T.sblock("b0"):
+                        vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                        C[vi, vj, vk] = A[vi, vj, vk] + 1
+                for k in T.serial(8):
+                    with T.sblock("b1"):
+                        vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                        D[vi, vj, vk] = A[vi, vj, vk] + 2
+                for k in T.serial(8):
+                    with T.sblock("b2"):
+                        vi, vk = T.axis.remap("SS", [i, k])
+                        vj = T.axis.opaque(8, j)
+                        B[vi, vj, vk] = (
+                            C[vi, vj, vk]
+                            + T.if_then_else(0 < vj, C[vi, vj - 1, vk], 0, dtype="int32")
+                            + D[vi, vj, vk]
+                        )
+
+    @T.prim_func(s_tir=True)
+    def after(A: T.Buffer((8, 8, 8), "int32"), B: T.Buffer((8, 8, 8), "int32")) -> None:
+        for i in T.serial(8):
+            with T.sblock():
+                T.reads(A[i, 0:8, 0:8])
+                T.writes(B[i, 0:8, 0:8])
+                C = T.sblock_alloc_buffer([8, 8, 8], dtype="int32")
+                D = T.sblock_alloc_buffer([8, 8, 8], dtype="int32")
+                for j in T.serial(8):
+                    for k in T.serial(8):
+                        with T.sblock("b0"):
+                            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                            C[vi, vj, vk] = A[vi, vj, vk] + 1
+                    for k in T.serial(8):
+                        with T.sblock("b1"):
+                            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                            D[vi, vj, vk] = A[vi, vj, vk] + 2
+                    for k in T.serial(8):
+                        with T.sblock("b2"):
+                            vi, vk = T.axis.remap("SS", [i, k])
+                            vj = T.axis.opaque(8, j)
+                            B[vi, vj, vk] = (
+                                C[vi, vj, vk]
+                                + T.if_then_else(0 < vj, C[vi, vj - 1, vk], 0, dtype="int32")
+                                + D[vi, vj, vk]
+                            )
+
+    _check(before, after)
+
+
+def test_1D_cascade_op_rolling_buffer():
+    """The intermediate buffer must be allocated above rolling buffer's rolling loop,
+    which is marked as opaque in consumer block's iter mappings."""
+
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((4, 16), "int32"), C: T.Buffer((4, 8), "int32")):
+        B = T.sblock_alloc_buffer((4, 6), "int32")
+        for c in T.serial(4):
+            for i in T.serial(0, 2):
+                for j in T.serial(0, 6):
+                    for k in T.serial(3):
+                        with T.sblock("P1"):
+                            T.where(i < 1 or j >= 2)
+                            cc, vi, vj, vk = T.axis.remap("SSSR", [c, i, j, k])
+                            if vk == 0:
+                                B[cc, T.floormod(vi * 4 + vj, 6)] = 0
+                            B[cc, T.floormod(vi * 4 + vj, 6)] = (
+                                B[cc, T.floormod(vi * 4 + vj, 6)] + A[cc, vi * 4 + vj + vk]
+                            )
+                for j in T.serial(0, 4):
+                    for k in T.serial(3):
+                        with T.sblock("P2"):
+                            vi = T.axis.opaque(2, i)
+                            cc, vj, vk = T.axis.remap("SSR", [c, j, k])
+                            if vk == 0:
+                                C[cc, vi * 4 + vj] = 0
+                            C[cc, vi * 4 + vj] = (
+                                C[cc, vi * 4 + vj] + B[cc, T.floormod(vi * 4 + vj + vk, 6)]
+                            )
+
+    @T.prim_func(s_tir=True)
+    def after(A: T.Buffer((4, 16), "int32"), C: T.Buffer((4, 8), "int32")):
+        for c in T.serial(4):
+            with T.sblock():
+                T.reads(A[c, 0:12], C[c, 0:8])
+                T.writes(C[c, 0:8])
+                B = T.sblock_alloc_buffer([4, 6], dtype="int32")
+                for i in T.serial(2):
+                    for j, k in T.grid(6, 3):
+                        with T.sblock("P1"):
+                            T.where(i < 1 or j >= 2)
+                            cc, vi, vj, vk = T.axis.remap("SSSR", [c, i, j, k])
+                            if vk == 0:
+                                B[cc, (vi * 4 + vj) % 6] = 0
+                            B[cc, (vi * 4 + vj) % 6] = (
+                                B[cc, (vi * 4 + vj) % 6] + A[cc, vi * 4 + vj + vk]
+                            )
+                    for j, k in T.grid(4, 3):
+                        with T.sblock("P2"):
+                            vi = T.axis.opaque(2, i)
+                            cc, vj, vk = T.axis.remap("SSR", [c, j, k])
+                            if vk == 0:
+                                C[cc, vi * 4 + vj] = 0
+                            C[cc, vi * 4 + vj] = C[cc, vi * 4 + vj] + B[cc, (vi * 4 + vj + vk) % 6]
+
+    _check(before, after)
+
+
+def test_buffer_conditional_lowering():
+    """Buffers passed as pointer arguments are unmodified
+
+    Confirm that the `tirx.PlanAndUpdateBufferAllocationLocation` pass
+    leaves (Buffer nodes corresponding to pointer-typed PrimFunc arguments)
+    unchanged, rather than lowering them to `reads`, `writes`, and `alloc_buffer` nodes.
+    """
+
+    @T.prim_func(s_tir=True)
+    def before(A: T.handle("float32")):
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i in range(1):
+            A_1 = T.decl_buffer((1,), data=A)
+            A_1[i] = 0
+
+    after = before
+    _check(before, after)
+
+
+def test_dltensor_buffer_is_unlowered():
+    """Buffers allocated with a Bind are unmodified
+
+    Confirm that the `tirx.PlanAndUpdateBufferAllocationLocation` pass
+    leaves (Buffer nodes corresponding to PrimFunc DLTensor arguments)
+    unchanged, rather than lowering them to `reads`, `writes`, and
+    `alloc_buffer` nodes.
+    """
+
+    @T.prim_func(s_tir=True)
+    def before(dlpack_handle: T.handle, axis: T.int64) -> T.int64:
+        ndim: T.int32 = T.tvm_struct_get(dlpack_handle, 0, 5, "int32")
+        stride_ptr: T.let[T.handle("int64")] = T.tvm_struct_get(
+            dlpack_handle, 0, 4, dtype=T.handle("int64").ty
+        )
+        if T.isnullptr(stride_ptr):
+            shape_ptr: T.let[T.handle("int64")] = T.tvm_struct_get(
+                dlpack_handle, 0, 3, dtype=T.handle("int64").ty
+            )
+            shape = T.decl_buffer(ndim, "int64", data=shape_ptr)
+            product = T.decl_buffer([], "int64")
+            product[()] = 1
+            for dim in range(axis + 1, ndim):
+                product[()] = product[()] * shape[dim]
+            return product[()]
+        else:
+            strides = T.decl_buffer(ndim, "int64", data=stride_ptr)
+            stride: T.int64 = strides[axis]
+            return stride
+
+    after = before
+    _check(before, after)
+
+
+def test_reduce_buffer_dominate_reduce_loops():
+    """Reduction write buffer allocation should dominate all reduce loops"""
+
+    @T.prim_func(s_tir=True)
+    def before(x: T.Buffer((256, 256, 256), "float32"), x_red: T.Buffer((256, 256), "float32")):
+        x_red_ = T.sblock_alloc_buffer((256, 256))
+        for ax0_0, k1_0, ax1_0 in T.grid(4, 4, 4):
+            for ax0_1, k1_1, ax1_1 in T.grid(64, 64, 64):
+                with T.sblock("x_red"):
+                    v_ax0 = T.axis.spatial(256, ax0_0 * 64 + ax0_1)
+                    v_ax1 = T.axis.spatial(256, ax1_0 * 64 + ax1_1)
+                    v_k1 = T.axis.reduce(256, k1_0 * 64 + k1_1)
+                    if v_k1 == 0:
+                        x_red_[v_ax0, v_ax1] = T.float32(0.0)
+                    x_red_[v_ax0, v_ax1] = x_red_[v_ax0, v_ax1] + x[v_ax0, v_k1, v_ax1]
+            for ax0, ax1 in T.grid(64, 64):
+                with T.sblock("x_red_"):
+                    v0 = T.axis.spatial(256, ax0_0 * 64 + ax0)
+                    v1 = T.axis.spatial(256, ax1_0 * 64 + ax1)
+                    x_red[v0, v1] = x_red_[v0, v1]
+
+    @T.prim_func(s_tir=True)
+    def after(x: T.Buffer((256, 256, 256), "float32"), x_red: T.Buffer((256, 256), "float32")):
+        for ax0_0 in range(4):
+            with T.sblock(""):
+                T.reads(x[ax0_0 * 64 : ax0_0 * 64 + 64, 0:256, 0:256])
+                T.writes(x_red[ax0_0 * 64 : ax0_0 * 64 + 64, 0:256])
+                x_red_ = T.sblock_alloc_buffer((256, 256))
+                for k1_0, ax1_0 in T.grid(4, 4):
+                    for ax0_1, k1_1, ax1_1 in T.grid(64, 64, 64):
+                        with T.sblock("x_red"):
+                            v_ax0 = T.axis.spatial(256, ax0_0 * 64 + ax0_1)
+                            v_ax1 = T.axis.spatial(256, ax1_0 * 64 + ax1_1)
+                            v_k1 = T.axis.reduce(256, k1_0 * 64 + k1_1)
+                            T.reads(x_red_[v_ax0, v_ax1], x[v_ax0, v_k1, v_ax1])
+                            T.writes(x_red_[v_ax0, v_ax1])
+                            if v_k1 == 0:
+                                x_red_[v_ax0, v_ax1] = T.float32(0.0)
+                            x_red_[v_ax0, v_ax1] = x_red_[v_ax0, v_ax1] + x[v_ax0, v_k1, v_ax1]
+                    for ax0, ax1 in T.grid(64, 64):
+                        with T.sblock("x_red_"):
+                            v0 = T.axis.spatial(256, ax0_0 * 64 + ax0)
+                            v1 = T.axis.spatial(256, ax1_0 * 64 + ax1)
+                            T.reads(x_red_[v0, v1])
+                            T.writes(x_red[v0, v1])
+                            x_red[v0, v1] = x_red_[v0, v1]
+
+    _check(before, after)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_profiling_instr.py

@@ -0,0 +1,354 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E741, F401
+
+import numpy
+
+import tvm
+import tvm.testing
+from tvm import s_tir
+from tvm.ir.module import IRModule
+from tvm.script import tirx as T
+
+default_lwp_test_config = {
+    "tirx.instrument_lwp": True,
+    "s_tir.lwp_disable_func_prof": True,
+    "s_tir.reset_start_id": True,
+}
+
+
+@T.prim_func(s_tir=True)
+def input1(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8, 128), dtype="int32")
+    B = T.match_buffer(b, (8, 8, 128), dtype="int32")
+    C = T.match_buffer(c, (8, 8, 128), dtype="int32")
+    for i, j in T.grid(8, 8):
+        for k, l in T.grid(8, 16):
+            with T.sblock("B"):
+                vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
+        for k, l in T.grid(8, 16):
+            with T.sblock("C"):
+                vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * 2
+
+
+@T.prim_func(s_tir=True)
+def input2(a: T.handle, b: T.handle, c: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8, 128), dtype="int32")
+    B = T.match_buffer(b, (8, 8, 128), dtype="int32")
+    C = T.match_buffer(c, (8, 8, 128), dtype="int32")
+    D = T.match_buffer(d, (8, 8, 128), dtype="int32")
+    for i in T.serial(0, 8):
+        for j in T.serial(0, 8):
+            for k, l in T.grid(8, 16):
+                with T.sblock("B"):
+                    vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                    B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
+            for k, l in T.grid(8, 16):
+                with T.sblock("B"):
+                    vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                    B[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * D[vi, vj, vk * 16 + vl]
+        for j in T.serial(0, 8):
+            for k, l in T.grid(8, 16):
+                with T.sblock("C"):
+                    vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                    C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] + 2
+            for k, l in T.grid(8, 16):
+                with T.sblock("B"):
+                    vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                    C[vi, vj, vk * 16 + vl] = C[vi, vj, vk * 16 + vl] * D[vi, vj, vk * 16 + vl]
+
+
+@T.prim_func(s_tir=True)
+def input3(a: T.handle, b: T.handle, c: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8, 128), dtype="int32")
+    B = T.match_buffer(b, (8, 8, 128), dtype="int32")
+    C = T.match_buffer(c, (8, 8, 128), dtype="int32")
+    D = T.match_buffer(d, (8, 8, 128), dtype="int32")
+    for i in T.serial(0, 8):
+        for j in T.parallel(0, 8):
+            for k in T.serial(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("B"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
+            for k in T.serial(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("B"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        B[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * D[vi, vj, vk * 16 + vl]
+        for j in T.serial(0, 8):
+            for k in T.parallel(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("C"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] + 2
+            for k in T.parallel(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("B"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        C[vi, vj, vk * 16 + vl] = C[vi, vj, vk * 16 + vl] * D[vi, vj, vk * 16 + vl]
+
+
+@T.prim_func(s_tir=True)
+def test1_expected_output(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8, 128), dtype="int32")
+    B = T.match_buffer(b, (8, 8, 128), dtype="int32")
+    C = T.match_buffer(c, (8, 8, 128), dtype="int32")
+    for i, j in T.grid(8, 8):
+        T.evaluate(T.start_profile_intrinsic(3, dtype="handle"))
+        for k, l in T.grid(8, 16):
+            with T.sblock("B"):
+                vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
+        T.evaluate(T.end_profile_intrinsic(3, dtype="handle"))
+        T.evaluate(T.start_profile_intrinsic(5, dtype="handle"))
+        for k, l in T.grid(8, 16):
+            with T.sblock("C"):
+                vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * 2
+        T.evaluate(T.end_profile_intrinsic(5, dtype="handle"))
+
+
+@T.prim_func(s_tir=True)
+def test2_expected_output(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8, 128), dtype="int32")
+    B = T.match_buffer(b, (8, 8, 128), dtype="int32")
+    C = T.match_buffer(c, (8, 8, 128), dtype="int32")
+    T.evaluate(T.start_profile_intrinsic(1, dtype="handle"))
+    for i in T.serial(0, 8):
+        T.evaluate(T.start_profile_intrinsic(2, dtype="handle"))
+        for j in T.serial(0, 8):
+            for k in T.serial(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("B"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
+            for k in T.serial(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("C"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * 2
+        T.evaluate(T.end_profile_intrinsic(2, dtype="handle"))
+    T.evaluate(T.end_profile_intrinsic(1, dtype="handle"))
+
+
+@T.prim_func(s_tir=True)
+def test3_expected_output(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8, 128), dtype="int32")
+    B = T.match_buffer(b, (8, 8, 128), dtype="int32")
+    C = T.match_buffer(c, (8, 8, 128), dtype="int32")
+    T.evaluate(T.start_profile_intrinsic(1, dtype="handle"))
+    for i in T.serial(0, 8):
+        T.evaluate(T.start_profile_intrinsic(2, dtype="handle"))
+        for j in T.serial(0, 8):
+            T.evaluate(T.start_profile_intrinsic(3, dtype="handle"))
+            for k in T.serial(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("B"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
+            T.evaluate(T.end_profile_intrinsic(3, dtype="handle"))
+            T.evaluate(T.start_profile_intrinsic(5, dtype="handle"))
+            for k in T.serial(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("C"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * 2
+            T.evaluate(T.end_profile_intrinsic(5, dtype="handle"))
+        T.evaluate(T.end_profile_intrinsic(2, dtype="handle"))
+    T.evaluate(T.end_profile_intrinsic(1, dtype="handle"))
+
+
+@T.prim_func(s_tir=True)
+def test4_expected_output(a: T.handle, b: T.handle, c: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8, 128), dtype="int32")
+    B = T.match_buffer(b, (8, 8, 128), dtype="int32")
+    C = T.match_buffer(c, (8, 8, 128), dtype="int32")
+    D = T.match_buffer(d, (8, 8, 128), dtype="int32")
+    for i in T.serial(0, 8):
+        T.evaluate(T.start_profile_intrinsic(2, dtype="handle"))
+        for j in T.serial(0, 8):
+            T.evaluate(T.start_profile_intrinsic(3, dtype="handle"))
+            for k, l in T.grid(8, 16):
+                with T.sblock("B"):
+                    vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                    B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
+            T.evaluate(T.end_profile_intrinsic(3, dtype="handle"))
+            T.evaluate(T.start_profile_intrinsic(5, dtype="handle"))
+            for k, l in T.grid(8, 16):
+                with T.sblock("B"):
+                    vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                    B[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * D[vi, vj, vk * 16 + vl]
+            T.evaluate(T.end_profile_intrinsic(5, dtype="handle"))
+        T.evaluate(T.end_profile_intrinsic(2, dtype="handle"))
+        T.evaluate(T.start_profile_intrinsic(7, dtype="handle"))
+        for j in T.serial(0, 8):
+            T.evaluate(T.start_profile_intrinsic(8, dtype="handle"))
+            for k, l in T.grid(8, 16):
+                with T.sblock("C"):
+                    vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                    C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] + 2
+            T.evaluate(T.end_profile_intrinsic(8, dtype="handle"))
+            T.evaluate(T.start_profile_intrinsic(10, dtype="handle"))
+            for k, l in T.grid(8, 16):
+                with T.sblock("B"):
+                    vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                    C[vi, vj, vk * 16 + vl] = C[vi, vj, vk * 16 + vl] * D[vi, vj, vk * 16 + vl]
+            T.evaluate(T.end_profile_intrinsic(10, dtype="handle"))
+        T.evaluate(T.end_profile_intrinsic(7, dtype="handle"))
+
+
+@T.prim_func(s_tir=True)
+def test5_expected_output(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8, 128), dtype="int32")
+    B = T.match_buffer(b, (8, 8, 128), dtype="int32")
+    C = T.match_buffer(c, (8, 8, 128), dtype="int32")
+    T.evaluate(T.start_profile_intrinsic(1, dtype="handle"))
+    for i in T.serial(0, 8):
+        T.evaluate(T.start_profile_intrinsic(2, dtype="handle"))
+        for j in T.serial(0, 8):
+            for k in T.serial(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("B"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
+            for k in T.serial(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("C"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * 2
+        T.evaluate(T.end_profile_intrinsic(2, dtype="handle"))
+    T.evaluate(T.end_profile_intrinsic(1, dtype="handle"))
+
+
+@T.prim_func(s_tir=True)
+def test6_expected_output(a: T.handle, b: T.handle, c: T.handle, d: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8, 128), dtype="int32")
+    B = T.match_buffer(b, (8, 8, 128), dtype="int32")
+    C = T.match_buffer(c, (8, 8, 128), dtype="int32")
+    D = T.match_buffer(d, (8, 8, 128), dtype="int32")
+    for i in T.serial(0, 8):
+        T.evaluate(T.start_profile_intrinsic(2, dtype="handle"))
+        for j in T.parallel(0, 8):
+            for k in T.serial(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("B"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
+            for k in T.serial(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("B"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        B[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * D[vi, vj, vk * 16 + vl]
+        T.evaluate(T.end_profile_intrinsic(2, dtype="handle"))
+        T.evaluate(T.start_profile_intrinsic(7, dtype="handle"))
+        for j in T.serial(0, 8):
+            T.evaluate(T.start_profile_intrinsic(8, dtype="handle"))
+            for k in T.parallel(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("C"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] + 2
+            T.evaluate(T.end_profile_intrinsic(8, dtype="handle"))
+            T.evaluate(T.start_profile_intrinsic(10, dtype="handle"))
+            for k in T.parallel(0, 8):
+                for l in T.serial(0, 16):
+                    with T.sblock("B"):
+                        vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                        C[vi, vj, vk * 16 + vl] = C[vi, vj, vk * 16 + vl] * D[vi, vj, vk * 16 + vl]
+            T.evaluate(T.end_profile_intrinsic(10, dtype="handle"))
+        T.evaluate(T.end_profile_intrinsic(7, dtype="handle"))
+
+
+# By default, only loops with siblings are instrumented.
+def test1():
+    with tvm.transform.PassContext(config=default_lwp_test_config):
+        mod = tvm.IRModule.from_expr(input1.with_attr("global_symbol", "main"))
+        mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
+    tvm.ir.assert_structural_equal(
+        mod["main"], test1_expected_output.with_attr("global_symbol", "main")
+    )
+
+
+# By default, only loops with siblings are instrumented. Here, 'lwp_max_depth'
+# doesn't have any effect unless 'instr_siblings' is set to False (ex: test3).
+def test2():
+    test2_config = default_lwp_test_config.copy()
+    test2_config.update({"s_tir.lwp_max_depth": 3})
+    with tvm.transform.PassContext(config=test2_config):
+        mod = tvm.IRModule.from_expr(input1.with_attr("global_symbol", "main"))
+        mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
+    tvm.ir.assert_structural_equal(
+        mod["main"], test1_expected_output.with_attr("global_symbol", "main")
+    )
+
+
+# test3: Use 'lwp_max_depth' to instrument loops upto a certain depth. This flag
+# is effective only when 'instr_siblings' is disabled. Also, note that inner-most
+# loops are always excluded from instrumentation unless overwritten using
+# 'lwp_min_height' (ex: test5)
+def test3():
+    test3_config = default_lwp_test_config.copy()
+    test3_config.update({"s_tir.lwp_max_depth": 3, "s_tir.instr_siblings": False})
+    with tvm.transform.PassContext(config=test3_config):
+        mod = tvm.IRModule.from_expr(input1.with_attr("global_symbol", "main"))
+        mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
+    tvm.ir.assert_structural_equal(
+        mod["main"], test3_expected_output.with_attr("global_symbol", "main")
+    )
+
+
+# test4: Use 'lwp_min_height' to exclude inner loops upto a certain height from
+# instrumentation.
+def test4():
+    with tvm.transform.PassContext(config=default_lwp_test_config):
+        mod = tvm.IRModule.from_expr(input2.with_attr("global_symbol", "main"))
+        mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
+    tvm.ir.assert_structural_equal(
+        mod["main"], test4_expected_output.with_attr("global_symbol", "main")
+    )
+
+
+# test5: Use both 'lwp_min_height' and 'lwp_max_depth'.
+# instrumentation.
+def test5():
+    test5_config = default_lwp_test_config.copy()
+    test5_config.update(
+        {"s_tir.lwp_max_depth": 3, "s_tir.instr_siblings": False, "s_tir.lwp_min_height": 2}
+    )
+    with tvm.transform.PassContext(config=test5_config):
+        mod = tvm.IRModule.from_expr(input1.with_attr("global_symbol", "main"))
+        mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
+    tvm.ir.assert_structural_equal(
+        mod["main"], test5_expected_output.with_attr("global_symbol", "main")
+    )
+
+
+# test6: Tests instrumentation for the parallel loops
+def test6():
+    with tvm.transform.PassContext(config=default_lwp_test_config):
+        mod = tvm.IRModule.from_expr(input3.with_attr("global_symbol", "main"))
+        mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
+    tvm.ir.assert_structural_equal(
+        mod["main"], test6_expected_output.with_attr("global_symbol", "main")
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py

@@ -0,0 +1,131 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_remove_store_undef():
+    """Remove a store whose value is T.undef()"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32")):
+            A[0] = T.undef(dtype="int32")
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32")):
+            T.evaluate(0)
+
+    After = tvm.s_tir.transform.RemoveStoreUndef()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_remove_store_undef_expression():
+    """Expressions containing T.undef() are removed"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32")):
+            A[0] = 1 + T.undef(dtype="int32")
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32")):
+            T.evaluate(0)
+
+    After = tvm.s_tir.transform.RemoveStoreUndef()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_keep_other_call_nodes():
+    """Expressions containing other CallNodes are not removed"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32"), n: T.int32):
+            A[0] = T.shift_left(n, 1, dtype="int32")
+
+    Expected = Before
+
+    After = tvm.s_tir.transform.RemoveStoreUndef()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_remove_let_undef():
+    """Remove a store whose value is bound to T.undef()"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32")):
+            val: T.let[T.int32] = T.undef(dtype="int32")
+            A[0] = val
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32")):
+            T.evaluate(0)
+
+    After = tvm.s_tir.transform.RemoveStoreUndef()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_raise_error_for_undef_as_store_indices():
+    """Use of T.undef() as buffer indices is an error"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32")):
+            val: T.let[T.int32] = T.undef(dtype="int32")
+            A[val] = 5
+
+    with pytest.raises(RuntimeError):
+        tvm.s_tir.transform.RemoveStoreUndef()(Before)
+
+
+def test_raise_error_for_undef_as_load_indices():
+    """Use of T.undef() as buffer indices is an error
+
+    Even though this occurs as part of the BufferStore's value, the
+    T.undef() may not appear in a buffer's indices.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32"), B: T.Buffer(1, "int32")):
+            B[0] = A[T.undef(dtype="int32")]
+
+    with pytest.raises(RuntimeError):
+        tvm.s_tir.transform.RemoveStoreUndef()(Before)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_remove_weight_layout_rewrite_block.py

@@ -0,0 +1,92 @@
+# 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 tvm
+from tvm.ir.module import IRModule
+from tvm.script import tirx as T
+from tvm.tirx.function import PrimFunc
+
+
+def _check(before, expect):
+    if isinstance(before, PrimFunc):
+        before = IRModule({"main": before.with_attr("global_symbol", "main")})
+    if isinstance(expect, PrimFunc):
+        expect = IRModule({"main": expect.with_attr("global_symbol", "main")})
+
+    mod = tvm.s_tir.transform.RemoveWeightLayoutRewriteBlock()(before)
+    tvm.ir.assert_structural_equal(mod, expect)
+
+
+def test_matmul():
+    @T.prim_func(s_tir=True)
+    def before(
+        A: T.Buffer((16, 16), "float32"),
+        B: T.Buffer((16, 16), "float32"),
+        C: T.Buffer((16, 16), "float32"),
+    ) -> None:
+        T.func_attr({"layout_free_buffers": [1]})
+        B_ = T.sblock_alloc_buffer([16, 4, 4], dtype="float32")
+        for i0_o, i1_o in T.grid(16, 16):
+            with T.sblock("layout_rewrite"):
+                i0, i1 = T.axis.remap("SS", [i0_o, i1_o])
+                T.reads(B[i0, i1])
+                T.writes(B_[i1, i0 // 4, i0 % 4])
+                T.sblock_attr({"meta_schedule.layout_rewrite_preproc": True})
+                B_[i1, i0 // 4, i0 % 4] = B[i0, i1]
+        for i0, j, k0, i1, k1 in T.grid(4, 16, 4, 4, 4):
+            with T.sblock("matmul"):
+                vi = T.axis.spatial(16, i0 * 4 + i1)
+                vj = T.axis.spatial(16, j)
+                vk = T.axis.reduce(16, k0 * 4 + k1)
+                T.reads(A[vi, vk], B_[vj, vk // 4, vk % 4])
+                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 // 4, vk % 4]
+
+    @T.prim_func(s_tir=True)
+    def after(
+        A: T.Buffer((16, 16), "float32"),
+        B: T.Buffer((16, 4, 4), "float32"),
+        C: T.Buffer((16, 16), "float32"),
+    ) -> None:
+        T.func_attr({"layout_free_buffers": [1]})
+        for i0_o, i1_o in T.grid(16, 16):
+            with T.sblock("layout_rewrite"):
+                i0, i1 = T.axis.remap("SS", [i0_o, i1_o])
+                T.reads()
+                T.writes()
+                T.sblock_attr({"meta_schedule.layout_rewrite_preproc": True})
+                T.evaluate(0)
+        for i0, j, k0, i1, k1 in T.grid(4, 16, 4, 4, 4):
+            with T.sblock("matmul"):
+                vi = T.axis.spatial(16, i0 * 4 + i1)
+                vj = T.axis.spatial(16, j)
+                vk = T.axis.reduce(16, k0 * 4 + k1)
+                T.reads(A[vi, vk], B[vj, vk // 4, vk % 4])
+                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 // 4, vk % 4]
+
+    _check(before, after)
+
+
+if __name__ == "__main__":
+    test_matmul()

tests/python/s_tir/transform/test_s_tir_transform_renormalize_split_pattern.py

@@ -0,0 +1,181 @@
+# 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
+
+import tvm
+import tvm.testing
+from tvm import s_tir
+from tvm.script import tirx as T
+
+# fmt: off
+# pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,redundant-keyword-arg
+
+@tvm.script.ir_module
+class Before:
+    @T.prim_func(s_tir=True)
+    def main(inputs: T.Buffer((1, 4, 4, 512), "float32"), weight: T.Buffer((4, 4, 512, 256), "float32"), conv2d_transpose_nhwc: T.Buffer((1, 8, 8, 256), "float32")) -> None:
+        # function attr dict
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        inputs_flat = T.decl_buffer([8192], dtype="float32", data=inputs.data)
+        weight_flat = T.decl_buffer([2097152], dtype="float32", data=weight.data)
+        conv2d_transpose_nhwc_flat = T.decl_buffer([16384], dtype="float32", data=conv2d_transpose_nhwc.data)
+        # var definition
+        threadIdx_x = T.env_thread("threadIdx.x")
+        blockIdx_x = T.env_thread("blockIdx.x")
+        # body
+        T.launch_thread(blockIdx_x, 64)
+        conv2d_transpose_nhwc_local = T.decl_buffer([8], "float32", scope="local")
+        PadInput_shared = T.decl_buffer([768], "float32", scope="shared")
+        weight_shared = T.decl_buffer([4096], "float32", scope="shared")
+        T.launch_thread(threadIdx_x, 32)
+        for i2_3_init, i1_4_init, i2_4_init in T.grid(2, 2, 2):
+            conv2d_transpose_nhwc_local[i1_4_init * 4 + i2_3_init * 2 + i2_4_init] = T.float32(0)
+        for i6_0 in T.serial(16):
+            for ax0_ax1_ax2_ax3_fused_0 in T.serial(24):
+                PadInput_shared[ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x] = T.if_then_else(128 <= ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x and ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x < 640 and 1 <= blockIdx_x // 32 * 2 + (ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x) % 128 // 32 and blockIdx_x // 32 * 2 + (ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x) % 128 // 32 < 5, inputs_flat[blockIdx_x // 32 * 1024 + ax0_ax1_ax2_ax3_fused_0 * 512 + i6_0 * 32 + threadIdx_x - 2560], T.float32(0), dtype="float32")
+            for ax0_ax1_ax2_ax3_fused_0 in T.serial(32):
+                weight_shared[T.ramp(ax0_ax1_ax2_ax3_fused_0 * 128 + threadIdx_x * 4, 1, 4)] = weight_flat[T.ramp((ax0_ax1_ax2_ax3_fused_0 * 128 + threadIdx_x * 4) // 256 * 131072 + i6_0 * 8192 + (ax0_ax1_ax2_ax3_fused_0 * 128 + threadIdx_x * 4) % 256 // 8 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 2 * 4, 1, 4)]
+            for i6_1, i2_3, i4_2, i5_2, i6_2, i1_4, i2_4 in T.grid(4, 2, 4, 4, 8, 2, 2):
+                conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] = conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] + T.if_then_else((i1_4 + i4_2) % 2 == 0 and (i2_4 + i5_2) % 2 == 0, PadInput_shared[threadIdx_x // 8 * 128 + (i1_4 + i4_2) // 2 * 128 + (i2_4 + i5_2) // 2 * 32 + i2_3 * 32 + i6_1 * 8 + i6_2], T.float32(0), dtype="float32") * weight_shared[i6_1 * 64 + i6_2 * 8 + threadIdx_x % 8 + 3840 - i5_2 * 256 - i4_2 * 1024]
+        for ax1, ax2 in T.grid(2, 4):
+            conv2d_transpose_nhwc_flat[threadIdx_x // 8 * 4096 + ax1 * 2048 + blockIdx_x // 32 * 1024 + ax2 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 8] = conv2d_transpose_nhwc_local[ax1 * 4 + ax2]
+
+
+@tvm.script.ir_module
+class After:
+    @T.prim_func(s_tir=True)
+    def main(inputs: T.Buffer((1, 4, 4, 512), "float32"), weight: T.Buffer((4, 4, 512, 256), "float32"), conv2d_transpose_nhwc: T.Buffer((1, 8, 8, 256), "float32")) -> None:
+        # function attr dict
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        inputs_flat = T.decl_buffer([8192], dtype="float32", data=inputs.data)
+        weight_flat = T.decl_buffer([2097152], dtype="float32", data=weight.data)
+        conv2d_transpose_nhwc_flat = T.decl_buffer([16384], dtype="float32", data=conv2d_transpose_nhwc.data)
+        # var definition
+        threadIdx_x = T.env_thread("threadIdx.x")
+        blockIdx_x = T.env_thread("blockIdx.x")
+        # body
+        T.launch_thread(blockIdx_x, 64)
+        conv2d_transpose_nhwc_local = T.decl_buffer([8], "float32", scope="local")
+        PadInput_shared = T.decl_buffer([768], "float32", scope="shared")
+        weight_shared = T.decl_buffer([4096], "float32", scope="shared")
+        T.launch_thread(threadIdx_x, 32)
+        for i2_3_init, i1_4_init, i2_4_init in T.grid(2, 2, 2):
+            conv2d_transpose_nhwc_local[i1_4_init * 4 + i2_3_init * 2 + i2_4_init] = T.float32(0)
+        for i6_0 in T.serial(16):
+            for ax0_ax1_ax2_ax3_fused_0 in T.serial(24):
+                PadInput_shared[ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x] = T.if_then_else(1 <= (ax0_ax1_ax2_ax3_fused_0 + threadIdx_x // 32) // 4 and (ax0_ax1_ax2_ax3_fused_0 + threadIdx_x // 32) // 20 < 1 and 1 <= blockIdx_x // 32 * 2 + (ax0_ax1_ax2_ax3_fused_0 + threadIdx_x // 32) % 4 and (blockIdx_x // 32 * 2 + (ax0_ax1_ax2_ax3_fused_0 + threadIdx_x // 32) % 4) // 5 < 1, inputs_flat[blockIdx_x // 32 * 1024 + ax0_ax1_ax2_ax3_fused_0 * 512 + i6_0 * 32 + threadIdx_x - 2560], T.float32(0), dtype="float32")
+            for ax0_ax1_ax2_ax3_fused_0 in T.serial(32):
+                weight_shared[T.ramp(ax0_ax1_ax2_ax3_fused_0 * 128 + threadIdx_x * 4, 1, 4)] = weight_flat[T.ramp((ax0_ax1_ax2_ax3_fused_0 + threadIdx_x * 4 // 128) // 2 * 131072 + i6_0 * 8192 + (ax0_ax1_ax2_ax3_fused_0 * 16 + threadIdx_x * 4 // 8) % 32 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 2 * 4, 1, 4)]
+            for i6_1, i2_3, i4_2, i5_2, i6_2, i1_4, i2_4 in T.grid(4, 2, 4, 4, 8, 2, 2):
+                conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] = conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] + T.if_then_else((i1_4 + i4_2) % 2 == 0 and (i2_4 + i5_2) % 2 == 0, PadInput_shared[threadIdx_x // 8 * 128 + (i1_4 + i4_2) // 2 * 128 + (i2_4 + i5_2) // 2 * 32 + i2_3 * 32 + i6_1 * 8 + i6_2], T.float32(0), dtype="float32") * weight_shared[i6_1 * 64 + i6_2 * 8 + threadIdx_x % 8 + 3840 - i5_2 * 256 - i4_2 * 1024]
+        for ax1, ax2 in T.grid(2, 4):
+            conv2d_transpose_nhwc_flat[threadIdx_x // 8 * 4096 + ax1 * 2048 + blockIdx_x // 32 * 1024 + ax2 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 8] = conv2d_transpose_nhwc_local[ax1 * 4 + ax2]
+
+
+@tvm.script.ir_module
+class After_simplified:
+    @T.prim_func(s_tir=True)
+    def main(inputs: T.Buffer((1, 4, 4, 512), "float32"), weight: T.Buffer((4, 4, 512, 256), "float32"), conv2d_transpose_nhwc: T.Buffer((1, 8, 8, 256), "float32")) -> None:
+        # function attr dict
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        # var definition
+        threadIdx_x = T.env_thread("threadIdx.x")
+        blockIdx_x = T.env_thread("blockIdx.x")
+        inputs_flat = T.decl_buffer([8192], dtype="float32", data=inputs.data)
+        weight_flat = T.decl_buffer([2097152], dtype="float32", data=weight.data)
+        conv2d_transpose_nhwc_flat = T.decl_buffer([16384], dtype="float32", data=conv2d_transpose_nhwc.data)
+        # body
+        T.launch_thread(blockIdx_x, 64)
+        conv2d_transpose_nhwc_local = T.decl_buffer([8], "float32", scope="local")
+        PadInput_shared = T.decl_buffer([768], "float32", scope="shared")
+        weight_shared = T.decl_buffer([4096], "float32", scope="shared")
+        T.launch_thread(threadIdx_x, 32)
+        for i2_3_init, i1_4_init, i2_4_init in T.grid(2, 2, 2):
+            conv2d_transpose_nhwc_local[i1_4_init * 4 + i2_3_init * 2 + i2_4_init] = T.float32(0)
+        for i6_0 in T.serial(16):
+            for ax0_ax1_ax2_ax3_fused_0 in T.serial(24):
+                PadInput_shared[ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x] = T.if_then_else(4 <= ax0_ax1_ax2_ax3_fused_0 and ax0_ax1_ax2_ax3_fused_0 < 20 and 1 <= blockIdx_x // 32 * 2 + ax0_ax1_ax2_ax3_fused_0 % 4 and blockIdx_x // 32 * 2 + ax0_ax1_ax2_ax3_fused_0 % 4 < 5, inputs_flat[blockIdx_x // 32 * 1024 + ax0_ax1_ax2_ax3_fused_0 * 512 + i6_0 * 32 + threadIdx_x - 2560], T.float32(0), dtype="float32")
+            for ax0_ax1_ax2_ax3_fused_0 in T.serial(32):
+                weight_shared[T.ramp(ax0_ax1_ax2_ax3_fused_0 * 128 + threadIdx_x * 4, 1, 4)] = weight_flat[T.ramp(ax0_ax1_ax2_ax3_fused_0 // 2 * 131072 + i6_0 * 8192 + ax0_ax1_ax2_ax3_fused_0 % 2 * 4096 + threadIdx_x // 2 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 2 * 4, 1, 4)]
+            for i6_1, i2_3, i4_2, i5_2, i6_2, i1_4, i2_4 in T.grid(4, 2, 4, 4, 8, 2, 2):
+                conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] = conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] + T.if_then_else((i1_4 + i4_2) % 2 == 0 and (i2_4 + i5_2) % 2 == 0, PadInput_shared[threadIdx_x // 8 * 128 + (i1_4 + i4_2) // 2 * 128 + (i2_4 + i5_2) // 2 * 32 + i2_3 * 32 + i6_1 * 8 + i6_2], T.float32(0), dtype="float32") * weight_shared[i6_1 * 64 + i6_2 * 8 + threadIdx_x % 8 + 3840 - i5_2 * 256 - i4_2 * 1024]
+        for ax1, ax2 in T.grid(2, 4):
+            conv2d_transpose_nhwc_flat[threadIdx_x // 8 * 4096 + ax1 * 2048 + blockIdx_x // 32 * 1024 + ax2 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 8] = conv2d_transpose_nhwc_local[ax1 * 4 + ax2]
+
+# pylint: enable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,redundant-keyword-arg
+# fmt: on
+
+
+def test_renormalize_split_pattern():
+    after = tvm.s_tir.transform.RenormalizeSplitPattern()(Before)
+    tvm.ir.assert_structural_equal(after, After)
+    after = tvm.tirx.transform.StmtSimplify()(after)
+    tvm.ir.assert_structural_equal(after, After_simplified)
+
+
+@T.prim_func(s_tir=True)
+def impossible_equality(n: T.int32):
+    # Prior to bugfix, this conditional defined the expression "2" as
+    # equal to zero within the then_case. [min_value=2, max_value=0]
+    if 2 == 0:
+        # Then this expression evaluates n/2, using the min/max values
+        # of "2", which is caught as a divide by zero error.
+        if n // 2 >= 16:
+            T.evaluate(0)
+
+
+@T.prim_func(s_tir=True)
+def impossible_inequality(n: T.int32):
+    # Prior to bugfix, this conditional set up a range of possible
+    # values for the expression "-2" as [0, kPosInf].
+    if -1 < -2:
+        if n // (-2) >= 16:
+            T.evaluate(0)
+
+
+integer_condition = tvm.testing.parameter(
+    impossible_equality,
+    impossible_inequality,
+)
+
+
+def test_analyze_inside_integer_conditional(integer_condition):
+    """Avoid crash occurring in ConstIntBoundAnalyzer.
+
+    Crash occurred when simplifying some expressions with provably
+    false integer expressions.  If the expressions were renormalized
+    before calling Simplify, conditional statements could assign a
+    range of possible values to integers, as if they were variables.
+    This would result in divide by zero throwing an exception,
+    followed by a second exception during stack unwinding causing the
+    program to crash.
+    """
+
+    # Similar issue would occur in most transformations that subclass
+    # IRMutatorWithAnalyzer.  tirx.transform.StmtSimplify() is an
+    # exception, as it rewrites the integer conditionals first.  These
+    # tests are written using RenormalizeSplitPattern as it is the
+    # first case identified.
+    transform = tvm.s_tir.transform.RenormalizeSplitPattern()
+
+    # Issue would result in an error through while applying the transformation.
+    mod = tvm.IRModule.from_expr(integer_condition)
+    transform(mod)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py

@@ -0,0 +1,92 @@
+# 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, F841
+import tvm
+from tvm import s_tir
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_rewrite_Select():
+    @I.ir_module
+    class ModuleY:
+        @T.prim_func(s_tir=True)
+        def main(i: T.int32):
+            A = T.alloc_buffer((100,))
+            T.evaluate(T.Select(i > 1, A[i - 1], T.float32(1.0)))
+
+    yy = tvm.s_tir.transform.RewriteUnsafeSelect()(ModuleY)["main"].body.seq[-1].value
+
+    @I.ir_module
+    class ModuleZ:
+        @T.prim_func(s_tir=True)
+        def main(i: T.int32):
+            A = T.alloc_buffer((100,))
+            T.evaluate(
+                T.Select(
+                    T.Select(i > 1, A[i - 1], T.float32(1.0)) > T.float32(0.0), A[i], T.float32(0.1)
+                )
+            )
+
+    zz = tvm.s_tir.transform.RewriteUnsafeSelect()(ModuleZ)["main"].body.seq[-1].value
+
+    @I.ir_module
+    class ModuleA:
+        @T.prim_func(s_tir=True)
+        def main(i: T.int32):
+            A = T.alloc_buffer((100,))
+            # Inline y and z to avoid Let bindings - outer Select condition is safe (no buffer access)
+            T.evaluate(
+                T.Select(
+                    T.floordiv(i, 4) > 10,
+                    T.Select(i > 1, A[i - 1], T.float32(1.0)),
+                    T.Select(
+                        T.Select(i > 1, A[i - 1], T.float32(1.0)) > T.float32(0.0),
+                        A[i],
+                        T.float32(0.1),
+                    ),
+                )
+            )
+
+    aa = tvm.s_tir.transform.RewriteUnsafeSelect()(ModuleA)["main"].body.seq[-1].value
+    builtin_if_then_else = tvm.ir.Op.get("tirx.if_then_else")
+
+    assert yy.op.same_as(builtin_if_then_else)
+    assert yy.op.same_as(builtin_if_then_else)
+    assert isinstance(aa, tvm.tirx.Select)
+
+
+def test_scalar_address_survives_generic_transforms():
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(value: T.uint32, condition: T.bool):
+            T.evaluate(T.Select(condition, T.isnullptr(T.address_of(value)), T.bool(False)))
+
+    transforms = [
+        tvm.s_tir.transform.RewriteUnsafeSelect(),
+        tvm.s_tir.transform.ThreadSync("shared"),
+        tvm.s_tir.transform.MergeSharedMemoryAllocations(),
+        tvm.tirx.transform.StorageRewrite(),
+    ]
+    for transform in transforms:
+        transformed = transform(Before)
+        assert transformed["main"] is not None
+
+
+if __name__ == "__main__":
+    test_rewrite_Select()

tests/python/s_tir/transform/test_s_tir_transform_thread_sync.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, F821, F841
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import s_tir
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def run_passes(func: tvm.tirx.PrimFunc):
+    mod = tvm.IRModule.from_expr(func)
+
+    cuda_target = tvm.target.Target("cuda", host="llvm")
+
+    mod = tvm.tirx.transform.Apply(
+        lambda f: f.with_attr({"global_symbol": "test", "target": cuda_target})
+    )(mod)
+
+    mod = tvm.tirx.transform.SplitHostDevice()(mod)
+    return tvm.s_tir.transform.ThreadSync("shared")(mod)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_sync_read_thread_id_independent_location():
+    @T.prim_func(check_well_formed=False, s_tir=True)
+    def func(p0_arg: T.Buffer((1, 2, 1, 1), "float32"), p1: T.Buffer(2, "float32")) -> None:
+        threadIdx_x = T.env_thread("threadIdx.x")
+        blockIdx_x = T.env_thread("blockIdx.x")
+        p0 = T.Buffer([2], dtype="float32", data=p0_arg.data)
+        result_local = T.sblock_alloc_buffer([1], dtype="float32", scope="local")
+        temp_shared = T.sblock_alloc_buffer([1], dtype="float32", scope="shared")
+        T.launch_thread(blockIdx_x, 8)
+        T.launch_thread(threadIdx_x, 4)
+        result_local[0] = T.float32(0)
+        if threadIdx_x < 1:
+            temp_shared[0] = p0[0]
+        result_local[0] = result_local[0] + temp_shared[0] * p1[0]
+        if threadIdx_x < 1:
+            temp_shared[0] = p0[1]
+        result_local[0] = result_local[0] + temp_shared[0] * p1[1]
+
+    mod = run_passes(func)
+    assert "T.tvm_storage_sync" in str(mod)
+
+
+def test_sync_shared_dyn():
+    @T.prim_func(private=True, s_tir=True)
+    def func(A: T.Buffer((4, 4), "float32"), E: T.Buffer((4, 4), "float32")):
+        blockIdx_x = T.launch_thread("blockIdx.x", 1)
+        B = T.alloc_buffer((24,), "float32", scope="shared.dyn")
+        C = T.alloc_buffer((1,), "float32", scope="local")
+        D = T.alloc_buffer((16,), "float32", scope="shared.dyn")
+        threadIdx_x = T.launch_thread("threadIdx.x", 16)
+        B_1 = T.decl_buffer((24,), data=B.data, scope="shared.dyn")
+        A_1 = T.decl_buffer((16,), data=A.data)
+        B_1[threadIdx_x // 4 * 6 + threadIdx_x % 4] = A_1[threadIdx_x]
+        C_1 = T.decl_buffer((1,), data=C.data, scope="local")
+        C_1[0] = B_1[threadIdx_x // 4 * 6 + threadIdx_x % 4]
+        D_1 = T.decl_buffer((16,), data=D.data, scope="shared.dyn")
+        D_1[threadIdx_x] = C_1[0]
+        E_1 = T.decl_buffer((16,), data=E.data)
+        E_1[threadIdx_x] = D_1[threadIdx_x]
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((4, 4), "float32"), E: T.Buffer((4, 4), "float32")):
+        blockIdx_x = T.launch_thread("blockIdx.x", 1)
+        B_1 = T.alloc_buffer((24,), "float32", scope="shared.dyn")
+        C_1 = T.alloc_buffer((1,), "float32", scope="local")
+        D_1 = T.alloc_buffer((16,), "float32", scope="shared.dyn")
+        threadIdx_x = T.launch_thread("threadIdx.x", 16)
+        B_1_1 = T.decl_buffer((24,), data=B_1.data, scope="shared.dyn")
+        A_1 = T.decl_buffer((16,), data=A.data)
+        B_1_1[threadIdx_x // 4 * 6 + threadIdx_x % 4] = A_1[threadIdx_x]
+        C_1_1 = T.decl_buffer((1,), data=C_1.data, scope="local")
+        C_1_1[0] = B_1_1[threadIdx_x // 4 * 6 + threadIdx_x % 4]
+        D_1_1 = T.decl_buffer((16,), data=D_1.data, scope="shared.dyn")
+        T.evaluate(T.call_intrin("int32", "tirx.tvm_storage_sync", "shared.dyn"))
+        D_1_1[threadIdx_x] = C_1_1[0]
+        E_1 = T.decl_buffer((16,), data=E.data)
+        E_1[threadIdx_x] = D_1_1[threadIdx_x]
+
+    mod = tvm.IRModule({"main": func})
+    mod = tvm.s_tir.transform.ThreadSync("shared.dyn")(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_sync_bind():
+    @T.prim_func(private=True, s_tir=True)
+    def func(A: T.Buffer((16 * 512), "float32")):
+        blockIdx_x = T.launch_thread("blockIdx.x", 16)
+        A_shared = T.alloc_buffer((512,), "float32", scope="shared")
+        in_thread_A_temp = T.alloc_buffer((1,), "float32", scope="local")
+        cross_thread_A_temp = T.alloc_buffer((1,), "float32", scope="local")
+        threadIdx_x = T.launch_thread("threadIdx.x", 128)
+        A_shared_1 = T.decl_buffer((512,), data=A_shared.data, scope="shared")
+        for ax0 in range(512):
+            A_shared_1[ax0] = A[blockIdx_x * 512 + ax0]
+        in_thread_A_temp_1 = T.decl_buffer((1,), data=in_thread_A_temp.data, scope="local")
+        in_thread_A_temp_1[0] = T.float32(0)
+        A_temp_1 = T.bind(in_thread_A_temp_1[0] + A_shared_1[threadIdx_x])
+        in_thread_A_temp_1[0] = A_temp_1
+        A_temp_2 = T.bind(in_thread_A_temp_1[0] + A_shared_1[threadIdx_x + 128])
+        in_thread_A_temp_1[0] = A_temp_2
+        A_temp_3 = T.bind(in_thread_A_temp_1[0] + A_shared_1[threadIdx_x + 256])
+        in_thread_A_temp_1[0] = A_temp_3
+        A_temp_4 = T.bind(in_thread_A_temp_1[0] + A_shared_1[threadIdx_x + 384])
+        in_thread_A_temp_1[0] = A_temp_4
+        cross_thread_A_temp_1 = T.decl_buffer((1,), data=cross_thread_A_temp.data, scope="local")
+        with T.attr(
+            T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+            "reduce_scope",
+            T.int32(0),
+        ):
+            T.tvm_thread_allreduce(
+                T.uint32(1),
+                in_thread_A_temp_1[0],
+                T.bool(True),
+                cross_thread_A_temp_1[0],
+                threadIdx_x,
+            )
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((8192,), "float32")):
+        blockIdx_x = T.launch_thread("blockIdx.x", 16)
+        A_shared_1 = T.alloc_buffer((512,), "float32", scope="shared")
+        in_thread_A_temp_1 = T.alloc_buffer((1,), "float32", scope="local")
+        cross_thread_A_temp_1 = T.alloc_buffer((1,), "float32", scope="local")
+        threadIdx_x = T.launch_thread("threadIdx.x", 128)
+        A_shared_1_1 = T.decl_buffer((512,), data=A_shared_1.data, scope="shared")
+        for ax0 in range(512):
+            A_shared_1_1[ax0] = A[blockIdx_x * 512 + ax0]
+        in_thread_A_temp_1_1 = T.decl_buffer((1,), data=in_thread_A_temp_1.data, scope="local")
+        in_thread_A_temp_1_1[0] = T.float32(0)
+        T.evaluate(T.call_intrin("int32", "tirx.tvm_storage_sync", "shared"))
+        A_temp_1 = T.bind(in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x])
+        in_thread_A_temp_1_1[0] = A_temp_1
+        A_temp_2 = T.bind(in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x + 128])
+        in_thread_A_temp_1_1[0] = A_temp_2
+        A_temp_3 = T.bind(in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x + 256])
+        in_thread_A_temp_1_1[0] = A_temp_3
+        A_temp_4 = T.bind(in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x + 384])
+        in_thread_A_temp_1_1[0] = A_temp_4
+        cross_thread_A_temp_1_1 = T.decl_buffer(
+            (1,), data=cross_thread_A_temp_1.data, scope="local"
+        )
+        T.attr(
+            T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+            "reduce_scope",
+            T.int32(0),
+        )
+        T.tvm_thread_allreduce(
+            T.uint32(1),
+            in_thread_A_temp_1_1[0],
+            T.bool(True),
+            cross_thread_A_temp_1_1[0],
+            threadIdx_x,
+        )
+
+    mod = tvm.IRModule({"main": func})
+    mod = tvm.s_tir.transform.ThreadSync("shared")(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    test_thread_storage_sync()
+    test_sync_else_branch()
+    test_sync_read_thread_id_independent_location()
+    test_sync_shared_dyn()
+    test_sync_bind()

tests/python/s_tir/transform/test_s_tir_transform_unify_thread_binding.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.
+# ruff: noqa: F401
+import sys
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import s_tir
+from tvm.script import tirx as T
+
+
+def _check(original, transformed):
+    mod = tvm.IRModule.from_expr(original.with_attr("global_symbol", "main"))
+    mod = tvm.s_tir.transform.UnifyThreadBinding()(mod)
+    mod = tvm.tirx.transform.StmtSimplify()(mod)
+    tvm.ir.assert_structural_equal(
+        mod["main"], transformed.with_attr("global_symbol", "main"), True
+    )
+
+
+def _check_fail(original):
+    mod = tvm.IRModule.from_expr(original)
+    with pytest.raises(ValueError):
+        tvm.s_tir.transform.UnifyThreadBinding()(mod)
+
+
+@T.prim_func(s_tir=True)
+def element_wise_thread_x(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 in T.thread_binding(0, 128, "blockIdx.x"):
+        for j0_0 in T.thread_binding(0, 4, "threadIdx.x"):
+            for j0_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    B[i, j0_0 * 32 + j0_1] = A[i, j0_0 * 32 + j0_1] * 2.0
+        for j1_0 in T.thread_binding(0, 4, "threadIdx.x"):
+            for j1_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    C[i, j1_0 * 32 + j1_1] = B[i, j1_0 * 32 + j1_1] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def unified_element_wise_thread_x(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 blockIdx_x in T.thread_binding(0, 128, "blockIdx.x"):
+        for threadIdx_x in T.thread_binding(0, 4, "threadIdx.x"):
+            for j0_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    B[blockIdx_x, threadIdx_x * 32 + j0_1] = (
+                        A[blockIdx_x, threadIdx_x * 32 + j0_1] * 2.0
+                    )
+            for j1_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    C[blockIdx_x, threadIdx_x * 32 + j1_1] = (
+                        B[blockIdx_x, threadIdx_x * 32 + j1_1] + 1.0
+                    )
+
+
+@T.prim_func(s_tir=True)
+def element_wise_thread_x_different_dtype(
+    A: T.Buffer((128, 128), "float32"),
+    B: T.Buffer((128, 128), "float32"),
+    C: T.Buffer((128, 128), "float32"),
+) -> None:
+    for i in T.thread_binding(128, "blockIdx.x"):
+        for j0_0 in T.thread_binding(4, "threadIdx.x"):
+            for j0_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    B[i, j0_0 * 32 + j0_1] = A[i, j0_0 * 32 + j0_1] * 2.0
+        for j1_0 in T.thread_binding(T.int64(4), "threadIdx.x"):
+            for j1_1 in T.serial(T.int64(32)):
+                with T.sblock(""):
+                    C[i, j1_0 * T.int64(32) + j1_1] = B[i, j1_0 * T.int64(32) + j1_1] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def unified_element_wise_thread_x_different_dtype(
+    A: T.Buffer((128, 128), "float32"),
+    B: T.Buffer((128, 128), "float32"),
+    C: T.Buffer((128, 128), "float32"),
+) -> None:
+    for blockIdx_x in T.thread_binding(128, "blockIdx.x"):
+        for threadIdx_x in T.thread_binding(4, "threadIdx.x"):
+            for j0_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    B[blockIdx_x, threadIdx_x * 32 + j0_1] = (
+                        A[blockIdx_x, threadIdx_x * 32 + j0_1] * 2.0
+                    )
+            for j1_1 in T.serial(T.int64(32)):
+                with T.sblock(""):
+                    C[blockIdx_x, T.cast(threadIdx_x, "int64") * T.int64(32) + j1_1] = (
+                        B[blockIdx_x, T.cast(threadIdx_x, "int64") * T.int64(32) + j1_1] + 1.0
+                    )
+
+
+@T.prim_func(s_tir=True)
+def element_wise_env_thread_x(a: T.handle, b: T.handle, c: T.handle) -> None:
+    j1_0 = T.env_thread("threadIdx.x")
+    j0_0 = T.env_thread("threadIdx.x")
+    i = T.env_thread("blockIdx.x")
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    T.launch_thread(i, 128)
+    T.launch_thread(j0_0, 4)
+    T.launch_thread(j1_0, 4)
+
+    for j0_1 in T.serial(0, 32):
+        with T.sblock(""):
+            B[i, j0_0 * 32 + j0_1] = A[i, j0_0 * 32 + j0_1] * 2.0
+    for j1_1 in T.serial(0, 32):
+        with T.sblock(""):
+            C[i, j1_0 * 32 + j1_1] = B[i, j1_0 * 32 + j1_1] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def unified_element_wise_env_thread_x(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 blockIdx_x in T.thread_binding(0, 128, "blockIdx.x"):
+        for threadIdx_x in T.thread_binding(0, 4, "threadIdx.x"):
+            for j0_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    B[blockIdx_x, threadIdx_x * 32 + j0_1] = (
+                        A[blockIdx_x, threadIdx_x * 32 + j0_1] * 2.0
+                    )
+            for j1_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    C[blockIdx_x, threadIdx_x * 32 + j1_1] = (
+                        B[blockIdx_x, threadIdx_x * 32 + j1_1] + 1.0
+                    )
+
+
+@T.prim_func(s_tir=True)
+def element_wise_vthread_x(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    for i_0 in T.thread_binding(0, 2, "vthread.x"):
+        for i_1 in T.thread_binding(0, 64, "threadIdx.x"):
+            for j_0 in T.thread_binding(0, 2, "vthread.x"):
+                for j_1 in T.serial(0, 64):
+                    with T.sblock(""):
+                        B[i_0 * 64 + i_1, j_0 * 64 + j_1] = A[i_0 * 64 + i_1, j_0 * 64 + j_1] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def unified_element_wise_vthread_x(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    for vthread_x in T.thread_binding(0, 2, "vthread.x"):
+        for threadIdx_x in T.thread_binding(0, 64, "threadIdx.x"):
+            for j_1 in T.serial(0, 64):
+                with T.sblock(""):
+                    B[vthread_x * 64 + threadIdx_x, vthread_x * 64 + j_1] = (
+                        A[vthread_x * 64 + threadIdx_x, vthread_x * 64 + j_1] * 2.0
+                    )
+
+
+@T.prim_func(s_tir=True)
+def element_wise_two_thread_x_in_same_kernel_not_equal(
+    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, 64])
+    for i in T.thread_binding(0, 128, "blockIdx.x"):
+        for j0 in T.thread_binding(0, 128, "threadIdx.x"):
+            B[i, j0] = A[i, j0] * 2.0
+        for j1 in T.thread_binding(0, 64, "threadIdx.x"):
+            C[i, j1] = A[i, j1] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_kernels_with_different_size(
+    a: T.handle, b: T.handle, c: T.handle, d: T.handle
+) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [256, 256])
+    D = T.match_buffer(d, [256, 256])
+    for i0 in T.thread_binding(0, 128, "blockIdx.x"):
+        for j0 in T.thread_binding(0, 128, "threadIdx.x"):
+            B[i0, j0] = A[i0, j0] * 2.0
+    for i1 in T.thread_binding(0, 256, "blockIdx.x"):
+        for j1 in T.thread_binding(0, 256, "threadIdx.x"):
+            D[i1, j1] = C[i1, j1] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def unified_element_wise_kernels_with_different_size(
+    a: T.handle, b: T.handle, c: T.handle, d: T.handle
+) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [256, 256])
+    D = T.match_buffer(d, [256, 256])
+    for blockIdx_x in T.thread_binding(0, 128, "blockIdx.x"):
+        for threadIdx_x in T.thread_binding(0, 128, "threadIdx.x"):
+            B[blockIdx_x, threadIdx_x] = A[blockIdx_x, threadIdx_x] * 2.0
+    for blockIdx_x in T.thread_binding(0, 256, "blockIdx.x"):
+        for threadIdx_x in T.thread_binding(0, 256, "threadIdx.x"):
+            D[blockIdx_x, threadIdx_x] = C[blockIdx_x, threadIdx_x] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_implicit_block(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 in T.thread_binding(0, 128, "threadIdx.y"):
+        for j0_0 in T.thread_binding(0, 4, "threadIdx.x"):
+            for j0_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    B[i, j0_0 * 32 + j0_1] = A[i, j0_0 * 32 + j0_1] * 2.0
+        for j1_0 in T.thread_binding(0, 4, "threadIdx.x"):
+            for j1_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    C[i, j1_0 * 32 + j1_1] = B[i, j1_0 * 32 + j1_1] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def unified_element_wise_implicit_block(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 blockIdx_x in T.thread_binding(0, 128, "threadIdx.y"):
+        for threadIdx_x in T.thread_binding(0, 4, "threadIdx.x"):
+            for j0_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    B[blockIdx_x, threadIdx_x * 32 + j0_1] = (
+                        A[blockIdx_x, threadIdx_x * 32 + j0_1] * 2.0
+                    )
+            for j1_1 in T.serial(0, 32):
+                with T.sblock(""):
+                    C[blockIdx_x, threadIdx_x * 32 + j1_1] = (
+                        B[blockIdx_x, threadIdx_x * 32 + j1_1] + 1.0
+                    )
+
+
+def test_thread_x():
+    _check(element_wise_thread_x, unified_element_wise_thread_x)
+
+
+def test_thread_x_different_dtype():
+    _check(element_wise_thread_x_different_dtype, unified_element_wise_thread_x_different_dtype)
+
+
+def test_env_thread_x():
+    _check(element_wise_env_thread_x, unified_element_wise_env_thread_x)
+
+
+def test_vthread_x():
+    _check(element_wise_vthread_x, unified_element_wise_vthread_x)
+
+
+def test_two_thread_x_in_same_kernel_not_equal():
+    _check_fail(element_wise_two_thread_x_in_same_kernel_not_equal)
+
+
+def test_kernels_with_different_size():
+    _check(
+        element_wise_kernels_with_different_size, unified_element_wise_kernels_with_different_size
+    )
+
+
+def test_implicit_block():
+    _check(element_wise_implicit_block, unified_element_wise_implicit_block)
+
+
+def test_inner_binding_with_annotation():
+    @T.prim_func(s_tir=True)
+    def inner_binding_with_annotation(A: T.Buffer((64,), "float32"), B: T.Buffer((64,), "float32")):
+        for bx in T.thread_binding(32, "blockIdx.x"):
+            for tx in T.thread_binding(2, "threadIdx.x", annotations={"my_annotation": 1}):
+                with T.sblock("block"):
+                    v = T.axis.spatial(64, bx * 2 + tx)
+                    B[v] = A[v]
+
+    @T.prim_func(s_tir=True)
+    def unified_inner_binding_with_annotation(
+        A: T.Buffer((64,), "float32"), B: T.Buffer((64,), "float32")
+    ):
+        for blockIdx_x in T.thread_binding(32, thread="blockIdx.x"):
+            for threadIdx_x in T.thread_binding(2, thread="threadIdx.x"):
+                for var in T.serial(1, annotations={"my_annotation": 1}):
+                    with T.sblock("block"):
+                        v = T.axis.spatial(64, blockIdx_x * 2 + threadIdx_x)
+                        T.reads(A[v])
+                        T.writes(B[v])
+                        B[v] = A[v]
+
+    _check(inner_binding_with_annotation, unified_inner_binding_with_annotation)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()