chore: import upstream snapshot with attribution

tests/python/tirx/operator/tile_primitive/trn/test_binary_trn.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.
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+Tx_func_map = {"add": Tx.add, "sub": Tx.sub, "mul": Tx.mul, "min": Tx.minimum, "max": Tx.maximum}
+
+
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "min", "max"])
+@pytest.mark.parametrize(
+    "operands_type",
+    [
+        "region_region",
+        "const_region",
+        "region_const",
+        "region_broadcast_lhs",
+        "region_broadcast_rhs",
+    ],
+)
+def test_simple_binary(op_type, operands_type):
+    const = T.float32(3.0)
+    src1_shape = [128, 512] if operands_type != "region_broadcast_lhs" else [128, 1]
+    src1_layout = TileLayout(S[src1_shape : (1 @ P, 1 @ F)])
+    src2_shape = [128, 512] if operands_type != "region_broadcast_rhs" else [128, 1]
+    src2_layout = TileLayout(S[src2_shape : (1 @ P, 1 @ F)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    Tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def binary() ->None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        if operands_type == "region_region" or operands_type.startswith("region_broadcast"):
+            Tx_func(C_sbuf, A_sbuf, B_sbuf)
+        elif operands_type == "const_region":
+            Tx_func(C_sbuf, const, A_sbuf)
+        elif operands_type == "region_const":
+            Tx_func(C_sbuf, A_sbuf, const)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer(src1_shape, scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer(src2_shape, scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer(dst_shape, scope="trn.sbuf")
+        for b_loop in T.serial(0, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if operands_type == "region_region":
+                        T.nki.tensortensor(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], B_sbuf[p_loop, f_loop], op_type)  # noqa: E501
+                    elif operands_type == "region_const":
+                        T.nki.tensorscalar(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], T.float32(3.0), op_type, T.bool(False))  # noqa: E501
+                    elif operands_type == "const_region":
+                        T.nki.tensorscalar(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], T.float32(3.0), op_type, T.bool(True))  # noqa: E501
+                    elif operands_type == "region_broadcast_rhs":
+                        T.nki.tensorscalar(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], B_sbuf[p_loop, 0], op_type, T.bool(False))  # noqa: E501
+                    elif operands_type == "region_broadcast_lhs":
+                        T.nki.tensorscalar(C_sbuf[p_loop, f_loop], B_sbuf[p_loop, f_loop], A_sbuf[p_loop, 0], op_type, T.bool(True))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "min", "max"])
+@pytest.mark.parametrize(
+    "operands_type",
+    [
+        "region_region",
+        "const_region",
+        "region_const",
+        "region_broadcast_lhs",
+        "region_broadcast_rhs",
+    ],
+)
+def test_binary_complex(op_type, operands_type):
+    src1_shape = [1024, 512] if operands_type != "region_broadcast_lhs" else [1024, 4]
+    src1_layout_data_iter = (128, 4096) if operands_type != "region_broadcast_lhs" else (128, 32)
+    src1_layout = TileLayout(S[src1_layout_data_iter : (1 @ P, 1 @ F)])
+    src2_shape = [512, 512] if operands_type != "region_broadcast_rhs" else [128, 512]
+    src2_layout_data_iter = (128, 2048) if operands_type != "region_broadcast_rhs" else (128, 512)
+    src2_layout = TileLayout(S[src2_layout_data_iter : (1 @ P, 1 @ F)])
+
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
+    const = T.float32(3.0)
+    Tx_func = Tx_func_map[op_type]
+
+    src1_view_shape = [128, 8, 512]
+    src2_view_shape = [128, 4, 512] if operands_type != "region_broadcast_rhs" else [128, 1, 512]
+    dst_view_shape = [128, 4, 512]
+    if operands_type == "region_broadcast_lhs":
+        src1_view_shape = [128, 8, 4, 1]
+        src2_view_shape = [128, 4, 4, 128]
+        dst_view_shape = [128, 4, 4, 128]
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        A_sbuf_view = A_sbuf.view(*src1_view_shape)
+        B_sbuf_view = B_sbuf.view(*src2_view_shape)
+        C_sbuf_view = C_sbuf.view(*dst_view_shape)
+        for i in range(4):
+            if operands_type == "region_region":
+                Tx_func(C_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :], B_sbuf_view[:, i, :])
+            elif operands_type == "region_const":
+                Tx_func(C_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :], const)
+            elif operands_type == "const_region":
+                Tx_func(C_sbuf_view[:, i, :], const, A_sbuf_view[:, i * 2, :])
+            elif operands_type == "region_broadcast_rhs":
+                Tx_func(C_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :], B_sbuf_view[:, 0, :])
+            elif operands_type == "region_broadcast_lhs":
+                Tx_func(C_sbuf_view[:, i, :, :], A_sbuf_view[:, i*2,:, :], B_sbuf_view[:, i, :, :])
+
+    f_extent = 128 if operands_type == "region_broadcast_lhs" else 512
+    b_extent = 4 if operands_type == "region_broadcast_lhs" else 1
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer(src1_layout_data_iter, scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer(src2_layout_data_iter, scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        A_sbuf_view = T.decl_buffer(src1_layout_data_iter, data=A_sbuf.data, scope="trn.sbuf", layout=None)  # noqa: E501
+        B_sbuf_view = T.decl_buffer(src2_layout_data_iter, data=B_sbuf.data, scope="trn.sbuf", layout=None)  # noqa: E501
+        C_sbuf_view = T.decl_buffer((128, 2048), data=C_sbuf.data, scope="trn.sbuf", layout=None)
+        for i, b_loop in T.grid(4, b_extent):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, f_extent, annotations={"nki_dim":"F"}):
+                    if operands_type == "region_region":
+                        T.nki.tensortensor(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], B_sbuf_view[p_loop, i * 512 + f_loop], op_type)  # noqa: E501
+                    elif operands_type == "const_region":
+                        T.nki.tensorscalar(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], T.float32(3.0), op_type, T.bool(True))  # noqa: E501
+                    elif operands_type == "region_const":
+                        T.nki.tensorscalar(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], T.float32(3.0), op_type, T.bool(False))  # noqa: E501
+                    elif operands_type == "region_broadcast_lhs":
+                        T.nki.tensorscalar(C_sbuf_view[p_loop, i * 512 + b_loop * 128 + f_loop], B_sbuf_view[p_loop, i * 512 + b_loop * 128 + f_loop], A_sbuf_view[p_loop, i * 8 + b_loop], op_type, T.bool(True))  # noqa: E501
+                    elif operands_type == "region_broadcast_rhs":
+                        T.nki.tensortensor(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], B_sbuf_view[p_loop, f_loop], op_type)  # noqa: E501
+
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_broadcast1():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (1 @ F, 32 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(512, 128) : (1 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.add(C_sbuf, A_sbuf, B_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for b_loop in T.serial(0, 512):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.tensorscalar(C_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], B_sbuf[p_loop, b_loop], "add", T.bool(False))  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_broadcast2():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.add(C_sbuf, A_sbuf, B_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for b_loop in T.serial(0, 128):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                    T.nki.tensortensor(C_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 128 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 128 + f_loop], B_sbuf[p_loop, b_loop % 4 * 128 + f_loop], "add")  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_broadcast3():
+    src1_shape = [128, 512]
+    src1_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
+    src2_shape = [32, 128, 512]
+    src2_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.add(C_sbuf, A_sbuf, B_sbuf[0])
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in T.serial(0, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                    T.nki.tensortensor(C_sbuf[p_loop, b_loop * 128 + f_loop], A_sbuf[p_loop, b_loop * 128 + f_loop], B_sbuf[p_loop, b_loop * 4096 + f_loop], "add")  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_with_guard():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for j in range(4):
+            Tx.add(C_sbuf[:, :, 0:j*128], A_sbuf[:, :, 0:j*128], B_sbuf[:, 0:j*128])
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for j, b_loop in T.grid(4, 96):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                    if b_loop % 3 - j < 0:
+                        T.nki.tensortensor(C_sbuf[p_loop, b_loop % 3 * 4096 + b_loop // 3 * 128 + f_loop], A_sbuf[p_loop, b_loop % 3 * 4096 + b_loop // 3 * 128 + f_loop], B_sbuf[p_loop, b_loop % 3 * 128 + f_loop], "add")  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_compose_op_trn.py

@@ -0,0 +1,772 @@
+# 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.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+def test_simple_activation_reduce():
+    A_shape = (128, 512)
+    A_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    B_shape = (128, 512)
+    B_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    C_shape = (128, 1)
+    C_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        Tx.unary_reduce(B, C, A, "sqrt", "sum", reduce_axes=1)
+
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for b_loop in range(1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.activation_reduce(C[p_loop, 0], B[p_loop, f_loop], A[p_loop, f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_in_loop():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
+    C_shape = (16, 128)
+    C_layout = TileLayout(S[(2, 4, 2, 128) : (2 @ F, 4 @ F, 1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=1)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 16):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 + b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop % 8 // 2 * 2048 + b_loop // 8 * 1024 + b_loop % 2 * 512 + f_loop], A[p_loop, i * 8192 + b_loop * 512 + f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop])  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_in_loop2():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(16 * 512, 128) : (1 @ F, 1 @ P)])
+    C_shape = (16, 128)
+    C_layout = TileLayout(S[(2, 4, 2, 128) : (2 @ F, 4 @ F, 1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=1)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 16):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 + b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop * 512 + f_loop], A[p_loop, i * 8192 + b_loop * 512 + f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop])  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_two_stage():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
+    C_shape = (1, 128)
+    C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        partial_reduce = T.alloc_buffer((128, 8), scope="trn.sbuf")
+        const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 1):
+            for reduction_b_loop in range(8):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                        T.nki.activation_reduce(partial_reduce[p_loop, reduction_b_loop], B[p_loop, reduction_b_loop % 4 * 2048 + reduction_b_loop // 4 * 1024 + f_loop], A[p_loop, i * 8192 + reduction_b_loop * 1024 + f_loop], "sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(8, annotations={"nki_dim": "F"}):
+                    T.nki.tensorreduce(C[p_loop, 0], partial_reduce[p_loop, f_loop], "add", T.bool(False), -1)  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_with_bias_scale():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(16 * 512, 128) : (1 @ F, 1 @ P)])
+    C_shape = (16, 128)
+    C_layout = TileLayout(S[(2, 4, 2, 128) : (2 @ F, 4 @ F, 1 @ F, 1 @ P)])
+    bias_shape = 128
+    bias_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        bias = T.alloc_buffer(bias_shape, dtype="float32", scope="trn.sbuf", layout=bias_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=1, bias=bias, scale=2.0)  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        bias = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 16):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 + b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop * 512 + f_loop], A[p_loop, i * 8192 + b_loop * 512 + f_loop], "sqrt", "add", bias[p_loop, 0], T.float32(2.0))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_simple_tensor_scalar_reduce():
+    A_shape = (128, 512)
+    A_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    B_shape = (128, 512)
+    B_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    C_shape = (128, 1)
+    C_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        Tx.binary_reduce(B, C, A, 1.0, "add", "sum", reduce_axes=1)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "tensor_scalar_reduce"})
+        A = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for b_loop in range(1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.tensorscalar_reduce(C[p_loop, 0], B[p_loop, f_loop], A[p_loop, f_loop], T.float32(1.0), "add", "add", T.bool(False))  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": tensor_scalar_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_tensor_tensor_reduce_fail():
+    A_shape = (128, 512)
+    A_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    B_shape = (128, 512)
+    B_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    D_shape = (128, 512)
+    D_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    C_shape = (128, 1)
+    C_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        D = T.alloc_buffer(D_shape, dtype="float32", scope="trn.sbuf", layout=D_layout)
+        Tx.binary_reduce(B, C, A, D, "add", "sum", reduce_axes=1)
+
+        # fmt: off
+    with pytest.raises(Exception):
+        with target:
+            mod = tvm.IRModule({"main": tensor_scalar_reduce})
+            mod = tvm.tirx.transform.LowerTIRx()(mod)
+
+
+def test_tensor_scalar_reduce_complex():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+    reduce_dst_shape = [128, 512]
+    reduce_dst_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        D_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        Tx.binary_reduce(C_sbuf, D_sbuf, B_sbuf, A_sbuf, "add", "sum", reduce_axes=0)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "tensor_scalar_reduce"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        D_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in range(512):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.tensorscalar_reduce(D_sbuf[p_loop, b_loop % 4 * 128 + b_loop // 4], C_sbuf[p_loop, b_loop % 4 * 4096 + f_loop * 128 + b_loop // 4], A_sbuf[p_loop, b_loop % 4 * 4096 + f_loop * 128 + b_loop // 4], B_sbuf[p_loop, b_loop % 4 * 128 + b_loop // 4], "add", "add", T.bool(True))  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": tensor_scalar_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_tensor_scalar_reduce_two_stage():
+    src1_shape = [512, 1024, 4]
+    src1_layout = TileLayout(S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)])
+    dst1_shape = src1_shape
+    dst1_layout = src1_layout
+    reduce_dst_shape = [512]
+    reduce_dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(dst1_shape, "float32", scope="trn.sbuf", layout=dst1_layout)
+        C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        Tx.binary_reduce(B_sbuf, C_sbuf, A_sbuf, 1.0, "add", "sum", reduce_axes=(1, 2))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "tensor_scalar_reduce"})
+        partial_reduce = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for b_loop in range(4):
+            for reduction_b_loop in range(4):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                        T.nki.tensorscalar_reduce(partial_reduce[p_loop, reduction_b_loop], B_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], A_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], T.float32(1.0), "add", "add", T.bool(False))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(4, annotations={"nki_dim": "F"}):
+                    T.nki.tensorreduce(C_sbuf[p_loop, b_loop], partial_reduce[p_loop, f_loop], "add", T.bool(False), -1)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": tensor_scalar_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_vector_chain():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (1 @ F, 32 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(512, 128) : (1 @ F, 1 @ P)])
+    src3_shape = [512]
+    src3_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        _C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        D_sbuf = T.alloc_buffer(src3_shape, "float32", scope="trn.sbuf", layout=src3_layout)
+        E_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.binary_chain(E_sbuf, A_sbuf, B_sbuf, D_sbuf, "add", "add", reverse1=True)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        _C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        D_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        E_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for b_loop in T.serial(0, 512):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.scalar_tensor_scalar(E_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], B_sbuf[p_loop, b_loop], D_sbuf[p_loop, b_loop % 4], "add", "add", T.bool(False), T.bool(True))  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_vector_chain_2():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (1 @ F, 32 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(512, 128) : (1 @ F, 1 @ P)])
+    src3_shape = src1_shape
+    src3_layout = src1_layout
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        _C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        D_sbuf = T.alloc_buffer(src3_shape, "float32", scope="trn.sbuf", layout=src3_layout)
+        E_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.binary_chain(E_sbuf, A_sbuf, B_sbuf, D_sbuf, "add", "add", reverse1=True)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        _C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        D_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        E_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for b_loop in T.serial(0, 512):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.scalar_tensor_tensor(E_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], B_sbuf[p_loop, b_loop], D_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], "add", "add", T.bool(False), T.bool(True))  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduce_negate():
+    src_shape = [128, 512, 4]
+    src_layout = TileLayout(S[(128, 512, 4) : (1 @ P, 4 @ F, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.reduce_negate(B_sbuf[:, i], A_sbuf[:, :, i], reduce_op="sum", reduce_axes=-2)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, i], A_sbuf[p_loop, f_loop * 4 + i], "add", True, -1)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_reduce_guard():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    reduce_dst_shape = [512]
+    reduce_dst_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def binary_reduce() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        for j in range(4):
+            for i in range(4):
+                Tx.binary_reduce(B_sbuf[0:128*(j+1), 0:128*(i+1)], C_sbuf[0:128*(j+1)], A_sbuf[0:128*(j+1), 0:128*(i+1)], 0.0, "add", "sum", [-1])  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary_reduce"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for j, i, b_loop in T.grid(4, 4, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if b_loop - j < 1 and f_loop < i * 128 + 128:
+                        T.nki.tensorscalar_reduce(C_sbuf[p_loop, b_loop], B_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0), "add", "add", T.bool(False))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": binary_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_unary_reduce_guard():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    reduce_dst_shape = [512]
+    reduce_dst_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def unary_reduce() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        for j in range(4):
+            for i in range(4):
+                Tx.unary_reduce(B_sbuf[0:128*(j+1), 0:128*(i+1)], C_sbuf[0:128*(j+1)], A_sbuf[0:128*(j+1), 0:128*(i+1)], "sqrt", "sum", reduce_axes=[-1])  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary_reduce"})
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for j, i, b_loop in T.grid(4, 4, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    if b_loop - j < 1 and f_loop < i * 128 + 128:
+                        T.nki.activation_reduce(C_sbuf[p_loop, b_loop], B_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], "sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0))  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_chain_guard():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    src2_shape = [512, 1]
+    src2_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def binary_chain() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for j in range(4):
+            for i in range(4):
+                Tx.binary_chain(C_sbuf[0:128*(j+1), 0:128*(i+1)], A_sbuf[0:128*(j+1), 0:128*(i+1)], B_sbuf[0:128*(j+1), 0], 1.0, "add", "sub", reverse1=True)  # noqa: E501
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary_chain"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for j, i, b_loop in T.grid(4, 4, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if b_loop - j < 1 and f_loop < i * 128 + 128:
+                        T.nki.scalar_tensor_scalar(C_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], B_sbuf[p_loop, b_loop], T.float32(1.0), "add", "sub", T.bool(False), T.bool(True))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": binary_chain})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_two_stage_workspace():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
+    C_shape = (1, 128)
+    C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        intermediate_buffer = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1), workspace={"partial_reduce": intermediate_buffer})  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        intermediate_buffer = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 1):
+            for reduction_b_loop in range(8):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                        T.nki.activation_reduce(intermediate_buffer[p_loop, reduction_b_loop], B[p_loop, reduction_b_loop % 4 * 2048 + reduction_b_loop // 4 * 1024 + f_loop], A[p_loop, i * 8192 + reduction_b_loop * 1024 + f_loop], "sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(8, annotations={"nki_dim": "F"}):
+                    T.nki.tensorreduce(C[p_loop, 0], intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1)  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_tensor_scalar_reduce_two_stage_workspace():
+    src1_shape = [512, 1024, 4]
+    src1_layout = TileLayout(S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)])
+    dst1_shape = src1_shape
+    dst1_layout = src1_layout
+    reduce_dst_shape = [512]
+    reduce_dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce() -> None:
+        T.device_entry()
+        intermediate_buffer = T.alloc_buffer((128, 8), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(dst1_shape, "float32", scope="trn.sbuf", layout=dst1_layout)
+        C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        Tx.binary_reduce(B_sbuf, C_sbuf, A_sbuf, 1.0, "add", "sum", reduce_axes=(1, 2), workspace={"partial_reduce": intermediate_buffer})  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "tensor_scalar_reduce"})
+        intermediate_buffer = T.alloc_buffer((128, 8), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for b_loop in range(4):
+            for reduction_b_loop in range(4):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                        T.nki.tensorscalar_reduce(intermediate_buffer[p_loop, reduction_b_loop], B_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], A_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], T.float32(1.0), "add", "add", T.bool(False))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(4, annotations={"nki_dim": "F"}):
+                    T.nki.tensorreduce(C_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": tensor_scalar_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_unary_reduce_complex():
+    # fmt: off
+    @T.prim_func
+    def unary_reduce():
+        T.device_entry()
+        p = T.alloc_buffer((128, 8192), "float16", scope="trn.sbuf", layout="PF")
+        rowsum_p = T.alloc_buffer((2, 128, 1), scope="trn.sbuf", layout="FPF")
+        qk = T.alloc_buffer((2, 128, 8192), scope="trn.sbuf", layout="FPF")
+        running_max = T.alloc_buffer((16384, 1), dtype="float32", scope="trn.sbuf", layout="PF")
+        for i in range(4):
+            Tx.unary_reduce(p[0:128, 0:8192], rowsum_p[i % 2, 0:128, 0], qk[i % 2, 0:128, 0:8192], "exp", "sum", bias=running_max[i * 128:i * 128 + 128, 0])  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary_reduce"})
+        p = T.alloc_buffer((128, 8192), "float16", scope="trn.sbuf")
+        rowsum_p = T.alloc_buffer((128, 2), scope="trn.sbuf")
+        qk = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        running_max = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(8192, annotations={"nki_dim": "F"}):
+                    T.nki.activation_reduce(rowsum_p[p_loop, i % 2], p[p_loop, f_loop], qk[p_loop, i % 2 * 8192 + f_loop], "exp", "add", running_max[p_loop, i], T.float32(1.0))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_copy_trn.py

@@ -0,0 +1,841 @@
+# 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.testing
+from tvm.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+def test_simple_copy():
+    src_shape = [128, 512]
+    src_layout = T.TileLayout(T.S[(128, 512) : (512, 1)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(A_sbuf, A)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (128, 512), layout=None)
+        A_1 = T.decl_buffer((65536,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in T.serial(0, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.load(A_sbuf[p_loop, f_loop], A_1[p_loop * 512 + f_loop])
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_simple_copy_2():
+    src_shape = [128, 512]
+    src_layout = TileLayout(S[(128, 4, 128) : (512, 128, 1)])
+
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 4, 128) : (4 @ F, 1 @ F, 1 @ P)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(A_sbuf, A)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (128, 512), layout=None)
+        A_1 = T.decl_buffer((65536,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in T.serial(0, 512):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 1, annotations={"nki_dim": "F"}):
+                    T.nki.load(A_sbuf[p_loop, b_loop], A_1[b_loop * 128 + p_loop])
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_in_a_loop():
+    src_shape = [512, 512]
+    src_layout = T.TileLayout(T.S[(4, 128, 512) : (512 * 128, 512, 1)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.copy(A_sbuf[i * 128 : i * 128 + 128, :], A[i * 128 : i * 128 + 128, :])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (512, 512), layout=None)
+        A_1 = T.decl_buffer((262144,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.load(
+                        A_sbuf[p_loop, i * 512 + f_loop], A_1[i * 65536 + p_loop * 512 + f_loop]
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_in_a_loop_2():
+    src_shape = [512, 512]
+    src_layout = T.TileLayout(T.S[(128, 2048) : (2048, 1)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        A_sbuf_view = A_sbuf.view(128, 4, 512)
+        A_view = A.view(128, 4, 512)
+        for i in range(4):
+            Tx.copy(A_sbuf_view[:, i, :], A_view[:, i, :])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (512, 512), layout=None)
+        _A_flat = T.decl_buffer((262144,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        A_sbuf_view = T.decl_buffer((128, 2048), data=A_sbuf.data, scope="trn.sbuf", layout=None)
+        A_view = T.decl_buffer((262144,), data=A.data, layout=None)
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.load(
+                        A_sbuf_view[p_loop, i * 512 + f_loop],
+                        A_view[p_loop * 2048 + i * 512 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod.show()
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_transpose():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(B_sbuf, A_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for b_loop in range(16):
+            for extend_b_loop in range(1):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
+                        for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
+                            T.nki.matmul(acc_psum[b_loop % 8, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + b_loop], acc_psum[b_loop % 8, p_loop, f_loop])  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_transpose_2():
+    src_shape = [65536]
+    src_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    dst_shape = [4, 65536]
+    dst_layout = TileLayout(S[(4, 128, 128, 4) : (4 @ F, 16 @ F, 1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.copy(B_sbuf[i, :], A_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for i in range(4):
+            for b_loop in range(4):
+                for extend_b_loop in range(1):
+                    T.attr(0, "tensorized_nki_instruction", 1)
+                    for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                        for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
+                            for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
+                                T.nki.matmul(acc_psum[b_loop, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, lhs_f_loop * 4 + b_loop], identity[p_loop, rhs_f_loop], T.bool(True))  # noqa: E501
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                        T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + i * 4 + b_loop], acc_psum[b_loop, p_loop, f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_different_f():
+    src_shape = [512, 64]
+    src_layout = TileLayout(S[(4, 128, 4, 4, 4) : (64 @ F, 1 @ P, 16 @ F, 4 @ F, 1 @ F)])
+    dst_shape = [512, 64]
+    dst_layout = TileLayout(S[(4, 128, 4, 4, 4) : (64 @ F, 1 @ P, 4 @ F, 16 @ F, 1 @ F)])
+
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(B_sbuf, A_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        A_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
+        for b_loop in T.serial(0, 64):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 4, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(
+                        B_sbuf[
+                            p_loop,
+                            b_loop // 16 * 64 + b_loop % 4 * 16 + b_loop % 16 // 4 * 4 + f_loop,
+                        ],
+                        A_sbuf[p_loop, b_loop * 4 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_different_shape():
+    src_shape = [512, 64]
+    src_layout = TileLayout(S[(4, 128, 4, 4, 4) : (64 @ F, 1 @ P, 16 @ F, 4 @ F, 1 @ F)])
+    dst_shape = [4, 128, 4]
+    dst_layout = TileLayout(S[(4, 128, 4) : (4 @ F, 1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        B_sbuf_view = B_sbuf.view(512, 4)
+        Tx.copy(B_sbuf_view, A_sbuf[:, 0:4])
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        A_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        _B_sbuf_view = T.decl_buffer((128, 16), data=B_sbuf.data, scope="trn.sbuf", layout=None)
+        for b_loop in T.serial(0, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 4, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(
+                        B_sbuf[p_loop, b_loop * 4 + f_loop],
+                        A_sbuf[p_loop, b_loop * 64 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_irregular_shape():
+    src_shape = [128, 10000]
+    src_layout = TileLayout(S[(128, 10000) : (10000, 1)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.copy(A[:, i * 512 : i * 512 + 512], A_sbuf)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (128, 10000), layout=None)
+        A_1 = T.decl_buffer((1280000,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.store(A_1[p_loop * 10000 + i * 512 + f_loop], A_sbuf[p_loop, f_loop])
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_different_shape_dim():
+    src_shape = [32, 128, 512]
+    src_layout = TileLayout(S[(32, 128, 512) : (128 * 512, 128, 1)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(32):
+            Tx.copy(A_sbuf, A[i, :, :])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (32, 128, 512), layout=None)
+        A_1 = T.decl_buffer((2097152,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for i, b_loop in T.grid(32, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.load(A_sbuf[p_loop, f_loop], A_1[i * 65536 + p_loop * 128 + f_loop])
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_offset():
+    src_shape = [256, 512]
+    src_layout = TileLayout(S[(256, 512) : (512, 1)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(2):
+            Tx.copy(A_sbuf[i * 256 : i * 256 + 256, :], A)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (256, 512), layout=None)
+        A_1 = T.decl_buffer((131072,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 2):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.load(
+                        A_sbuf[p_loop, i * 1024 + b_loop * 512 + f_loop],
+                        A_1[b_loop * 65536 + p_loop * 512 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_large_dma_copy():
+    src_shape = [512, 4096]
+    src_layout = T.TileLayout(T.S[(4, 128, 4096) : (4096 * 128, 4096, 1)])
+    dst_shape = [512, 4096]
+    dst_layout = TileLayout(S[(4, 128, 4096) : (4096 @ F, 1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.copy(A_sbuf[i * 128 : i * 128 + 128, :], A[i * 128 : i * 128 + 128, :])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (512, 4096), layout=None)
+        A_1 = T.decl_buffer((2097152,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 4096, annotations={"nki_dim": "F"}):
+                    T.nki.load(
+                        A_sbuf[p_loop, i * 4096 + f_loop],
+                        A_1[i * 524288 + p_loop * 4096 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_inst_size_limit():
+    src_shape = [512, 4096]
+    src_layout = dst_layout = TileLayout(S[(4, 128, 4096) : (4096 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        B_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.copy(A_sbuf[i * 128 : i * 128 + 128, :], B_sbuf[i * 128 : i * 128 + 128, :])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+        B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(
+                        A_sbuf[p_loop, i * 4096 + b_loop * 512 + f_loop],
+                        B_sbuf[p_loop, i * 4096 + b_loop * 512 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_complex_index():
+    A_shape = [4096, 4096]
+    A_layout = T.TileLayout(T.S[(4096, 4096) : (1, 4096)])
+    A_sbuf_shape = (2, 2048, 1024)
+    A_sbuf_layout = TileLayout(S[(2, 2048, 8, 128) : (16384 @ F, 1 @ F, 2048 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle, ) -> None:
+        A = T.match_buffer(A_ptr, A_shape, "float32", layout=A_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(A_sbuf_shape, "float32", scope="trn.sbuf", layout=A_sbuf_layout)
+        Tx.copy(A_sbuf[1, 0:2048, 0:1024], A[2048: 4096, 3072:4096])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (4096, 4096), layout=None)
+        A_1 = T.decl_buffer((16777216,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 32768), scope="trn.sbuf")
+        for b_loop in T.serial(0, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
+                    T.nki.load(A_sbuf[p_loop, b_loop * 2048 + f_loop + 16384], A_1[b_loop * 524288 + p_loop * 4096 + f_loop + 12584960])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_complex_index_2():
+    A_sbuf_shape = [4096, 4096]
+    A_sbuf_layout = T.TileLayout(T.S[(4096, 32, 128) : (1 @ F, 4096 @ F, 1 @ P)])
+    A_shape = (2, 2048, 1024)
+    A_layout = T.TileLayout(T.S[(2, 2048, 1024) : (2048 * 1024, 1, 2048)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle, ) -> None:
+        A = T.match_buffer(A_ptr, A_shape, "float32", layout=A_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(A_sbuf_shape, "float32", scope="trn.sbuf", layout=A_sbuf_layout)
+        Tx.copy(A_sbuf[2048: 4096, 3072:4096], A[1, 0:2048, 0:1024])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (2, 2048, 1024), layout=None)
+        A_1 = T.decl_buffer((4194304,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 131072), scope="trn.sbuf")
+        for b_loop in T.serial(0, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
+                    T.nki.load(A_sbuf[p_loop, b_loop * 4096 + f_loop + 100352], A_1[b_loop * 262144 + p_loop * 2048 + f_loop + 2097152])  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_transpose_with_workspace():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        identity = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((1, 128, 512), "float32", scope="trn.psum", allocated_addr=(0, 0))
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        Tx.copy(B_sbuf, A_sbuf, workspace={"identity": identity, "acc_psum": acc_psum})
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((1, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        for b_loop in range(16):
+            for extend_b_loop in range(1):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
+                        for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
+                            T.nki.matmul(acc_psum[0, lhs_f_loop, extend_b_loop * 128 + rhs_f_loop], A_sbuf[p_loop, b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + b_loop], acc_psum[0, p_loop, f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_guard():
+    src_shape = [512, 512]
+    src_layout = T.TileLayout(T.S[(4, 128, 512) : (512 * 128, 512, 1)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for j in range(4):
+            for i in range(4):
+                Tx.copy(A_sbuf[i * 128 : i * 128 + 128, 0:128*j], A[i * 128 : i * 128 + 128, 0:128*j])  # noqa: E501
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (512, 512), layout=None)
+        A_1 = T.decl_buffer((262144,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for j, i, b_loop in T.grid(4, 4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 384, annotations={"nki_dim":"F"}):
+                    if f_loop < j * 128:
+                        T.nki.load(A_sbuf[p_loop, i * 512 + f_loop], A_1[i * 65536 + p_loop * 512 + f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_guard_2():
+    src_shape = [512, 512]
+    src_layout = T.TileLayout(T.S[(4, 128, 512) : (512 * 128, 512, 1)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for j in range(4):
+            for i in range(4):
+                Tx.copy(A_sbuf[0:128*j, 0:128*i], A[0:128*j, 0:128*i])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (512, 512), layout=None)
+        A_1 = T.decl_buffer((262144,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for j, i, b_loop in T.grid(4, 4, 3):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 384, annotations={"nki_dim":"F"}):
+                    if b_loop - j < 0 and f_loop < i * 128:
+                        T.nki.load(A_sbuf[p_loop, b_loop * 512 + f_loop], A_1[b_loop * 65536 + p_loop * 512 + f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_transpose_with_guard():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            for j in range(4):
+                Tx.copy(B_sbuf[i * 128 : i * 128 + 128, 0:128*j], A_sbuf[i * 128 : i * 128 + 128, 0:128*j])  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for i, j, b_loop in T.grid(4, 4, 3):
+            for extend_b_loop in range(1):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
+                        for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
+                            if b_loop - j < 0:
+                                T.nki.matmul(acc_psum[b_loop, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, i * 512 + b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    if b_loop - j < 0:
+                        T.nki.tensor_copy(B_sbuf[p_loop, i * 512 + f_loop * 4 + b_loop], acc_psum[b_loop, p_loop, f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_specified_max_inst_size():
+    src_shape = [128, 512]
+    src_layout = "PF"
+    dst_shape = src_shape
+    dst_layout = src_layout
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(A_sbuf, B_sbuf, max_inst_size=128)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf", layout=None)
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf", layout=None)
+        for b_loop in T.serial(0, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(A_sbuf[p_loop, b_loop * 128 + f_loop], B_sbuf[p_loop, b_loop * 128 + f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_transpose_with_extended_f():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((128, 2048), "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer((128, 2048), "float32", scope="trn.sbuf", layout="FP")
+        Tx.copy(B_sbuf, A_sbuf)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for b_loop in range(4):
+            for extend_b_loop in range(4):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
+                        for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
+                            T.nki.matmul(acc_psum[b_loop, lhs_f_loop, extend_b_loop * 128 + rhs_f_loop], A_sbuf[p_loop, b_loop * 512 + extend_b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(B_sbuf[p_loop, b_loop * 512 + f_loop], acc_psum[b_loop, p_loop, f_loop])  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_gemm_trn.py

@@ -0,0 +1,583 @@
+# 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.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+def test_simple_gemm():
+    A_layout = TileLayout(S[(128, 128) : (1 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)])
+
+    C_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((128, 128), "float32", scope="trn.psum", layout=C_layout)
+        Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((1, 128, 128), scope="trn.psum")
+        for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(1, 1, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[0, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, lhs_f_loop], B_sbuf[p_loop, rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_larger_gemm():
+    A_layout = TileLayout(S[(2, 128, 4, 128) : (512 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(2, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((256, 512), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((256, 256), "float32", scope="trn.psum", layout=C_layout)
+        Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((1, 128, 512), scope="trn.psum")
+        for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 1, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[0, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, lhs_b_loop * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, reduction_b_loop * 256 + rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_in_a_loop():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_psum[256 * i : 256 * i + 256, :],
+                )
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 1, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_with_stride():
+    A_layout = TileLayout(S[(4, 128, 128, 8) : (1024 @ F, 1 @ F, 1 @ P, 128 @ F)])
+    B_layout = TileLayout(S[(128, 8, 2, 128) : (1 @ P, 512 @ F, 256 @ F, 2 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 512, 2), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((512, 2, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, :, k],
+                    B_sbuf[:, k, :],
+                    C_psum[256 * i : 256 * i + 256, :],
+                )
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4095), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 1, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + reduction_b_loop * 256 + k * 128 + lhs_f_loop], B_sbuf[p_loop, reduction_b_loop * 1024 + k * 512 + rhs_f_loop * 2], True)  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_swap_lhs_rhs():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_psum[256 * i : 256 * i + 256, :],
+                )
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 2, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[i, lhs_f_loop, rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_with_sbuf_output():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                )
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        buffer = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        for i, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2):
+            for reduction_b_loop in range(4):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                  for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                    for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                        T.nki.matmul(buffer[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop], True)  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                T.nki.tensor_copy(C_sbuf[lhs_f_loop, i * 512 + rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], buffer[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_different_shape():
+    A_layout = TileLayout(S[(2, 4, 128, 8, 128) : (4096 @ F, 1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((2, 512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[1, 256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_psum[256 * i : 256 * i + 256, :],
+                )
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 2, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[i, lhs_f_loop, rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop + 4096], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_too_large_f_size():
+    A_layout = TileLayout(S[(256, 128) : (1 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(128, 1024) : (1 @ P, 1 @ F)])
+
+    C_layout = TileLayout(S[(2, 128, 1024) : (1024 @ F, 1 @ P, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((256, 128), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((128, 1024), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((256, 1024), "float32", scope="trn.psum", layout=C_layout)
+        Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((4, 128, 512), scope="trn.psum")
+        for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 512, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, lhs_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, rhs_b_loop * 512 + rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_sbuf_output_with_workspace():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
+        C_psum = T.alloc_buffer((1, 128, 512), "float32", scope="trn.psum", allocated_addr=(0, 0))
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                    workspace={"acc_psum": C_psum}
+                )
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((1, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        for i, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2):
+            for reduction_b_loop in range(4):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                  for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                    for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                        T.nki.matmul(C_psum[0, lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop], True)  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                T.nki.tensor_copy(C_sbuf[lhs_f_loop, i * 512 + rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], C_psum[0, lhs_f_loop, rhs_f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_pf_mismatch_fail():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(2, 128, 8, 128) : (128 @ F, 1 @ F, 256 @ F, 1 @ P)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((256, 1024), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[:, 512 * k : 512 * k + 512],
+                    C_psum[256 * i : 256 * i + 256, :],
+                )
+        # fmt: on
+    with pytest.raises(Exception):
+        with target:
+            mod = tvm.IRModule({"main": gemm})
+            mod = tvm.tirx.transform.LowerTIRx()(mod)
+
+
+def test_gemm_transpose_AB():
+    A_layout = TileLayout(S[(8, 128, 4, 128) : (128 @ F, 1 @ P, 1024 @ F, 1 @ F)])
+    B_layout = TileLayout(S[(2, 128, 8, 128) : (128 @ F, 1 @ F, 256 @ F, 1 @ P)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((1024, 512), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((256, 1024), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[512 * k : 512 * k + 512, 256 * i : 256 * i + 256],
+                    B_sbuf[:, 512 * k : 512 * k + 512],
+                    C_psum[256 * i : 256 * i + 256, :],
+                    transpose_A=True,
+                    transpose_B=True,
+                )
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 1, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + rhs_f_loop], True)  # noqa: E501
+
+            #fmt: off
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_guard():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            for j in range(2):
+                for k in range(2):
+                    Tx.gemm(
+                        C_sbuf[0: 256 * i, 0: 128 * (j + 1)],
+                        A_sbuf[0: 256 * i, 0: 512 * (k + 1)],
+                        B_sbuf[0: 512 * (k + 1), 0: 128 * (j + 1)],
+                        C_sbuf[0: 256 * i, 0: 128 * (j + 1)],
+                    )
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        for i, j, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2, 2):
+            for reduction_b_loop in range(8):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                  for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                    for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                        if reduction_b_loop - k * 4 < 4 and lhs_b_loop - j < 1 and 0 < i and reduction_b_loop - k * 4 < 4:  # noqa: E501
+                            T.nki.matmul(acc_psum[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, rhs_b_loop * 1024 + reduction_b_loop * 128 + rhs_f_loop], True)  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                if 0 < i and lhs_b_loop - j < 1:
+                    T.nki.tensor_copy(C_sbuf[lhs_f_loop, rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], acc_psum[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_guard2():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for j in range(4):
+            for i in range(2):
+                for k in range(2):
+                    Tx.gemm(
+                        C_psum[256 * i : 256 * i + 256, :],
+                        A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + (j+1) * 128],
+                        B_sbuf[512 * k : 512 * k + (j+1) * 128, :],
+                        C_psum[256 * i : 256 * i + 256, :],
+                    )
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for j, i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(4, 2, 2, 2, 1, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
+                    if reduction_b_loop - j < 1 and reduction_b_loop - j < 1:
+                        T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_private_alloc_trn.py

@@ -0,0 +1,402 @@
+# 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.testing
+from tvm.ir import assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.trn.transform import TrnPrivateBufferAlloc
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def test_copy_transpose():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(B_sbuf, A_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        A_sbuf = T.alloc_buffer((512, 512), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 2048) : (1 @ P, 1@F)]))
+        B_sbuf = T.alloc_buffer((512, 512), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(2048, 128) : (1@F, 1@P)]))
+        Tx.copy(B_sbuf[0:512, 0:512], A_sbuf[0:512, 0:512], workspace={"acc_psum": acc_psum, "identity": identity})  # noqa: E501
+
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_normal_copy():
+    src_shape = [128, 512]
+    src_layout = TileLayout(S[(128, 512) : (512, 1)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(A_sbuf, A)
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], copy)
+
+
+def test_unary_with_bias_scale():
+    src_shape = [512, 1024]
+    src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    bias = T.float32(1.0)
+    scale = T.float32(2.0)
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.exp(C_sbuf, A_sbuf, bias=bias, scale=scale)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        T.device_entry()
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(1.0))
+        A_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096) : (1@P, 1@F)]))
+        C_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096) : (1@P, 1@F)]))
+        Tx.exp(C_sbuf[0:512, 0:1024], A_sbuf[0:512, 0:1024], T.float32(1.0), T.float32(2.0), workspace={"const_bias": const_bias})  # noqa: E501
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_two_stage():
+    src_shape = [128, 32, 4, 32]
+    src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.sum(B_sbuf, A_sbuf, axes=(1, 3))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        T.device_entry()
+        partial_reduce = T.alloc_buffer((128, 32), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 32, 4, 32), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 32 * 32 * 4) : (1@P, 1@F)]))
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4) : (1@P, 1@F)]))
+        Tx.sum(B_sbuf[0:128, 0:4], A_sbuf[0:128, 0:32, 0:4, 0:32], [1, 3], False, workspace={"partial_reduce": partial_reduce})  # noqa: E501
+
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 1 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                )
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        T.device_entry()
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        A_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(4, 128, 8, 128) : (1024@F, 1@F, 1@F, 1@P)]))  # noqa: E501
+        B_sbuf = T.alloc_buffer((1024, 256), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(8, 128, 2, 128) : (256@F, 1@P, 128@F, 1@F)]))  # noqa: E501
+        C_sbuf = T.alloc_buffer((512, 256), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(4, 128, 2, 128) : (256@F, 1@F, 128@F, 1@P)]))  # noqa: E501
+        for i, k in T.grid(2, 2):
+            Tx.gemm(C_sbuf[256 * i:256 * i + 256, 0:256], A_sbuf[256 * i:256 * i + 256, 512 * k:512 * k + 512], B_sbuf[512 * k:512 * k + 512, 0:256], C_sbuf[256 * i:256 * i + 256, 0:256], False, False, T.float32(1.0), T.float32(0.0), workspace={"acc_psum": acc_psum})  # noqa: E501
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_reduce_two_stage():
+    src1_shape = [512, 1024, 4]
+    src1_layout = TileLayout(S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)])
+    dst1_shape = src1_shape
+    dst1_layout = src1_layout
+    reduce_dst_shape = [512]
+    reduce_dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(dst1_shape, "float32", scope="trn.sbuf", layout=dst1_layout)
+        C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        Tx.binary_reduce(B_sbuf, C_sbuf, A_sbuf, 1.0, "add", "sum", reduce_axes=(1, 2))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "tensor_scalar_reduce"})
+        T.device_entry()
+        partial_reduce = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((512, 1024, 4), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)]))
+        B_sbuf = T.alloc_buffer((512, 1024, 4), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)]))
+        C_sbuf = T.alloc_buffer((512,), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4) : (1 @ P, 1 @ F)]))
+        Tx.binary_reduce(B_sbuf[0:512, 0:1024, 0:4], C_sbuf[0:512], A_sbuf[0:512, 0:1024, 0:4], T.float32(1.0), "add", "sum", [1, 2], workspace={"partial_reduce": partial_reduce})  # noqa: E501
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": tensor_scalar_reduce})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_two_stage():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
+    C_shape = (1, 128)
+    C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        T.device_entry()
+        partial_reduce = T.alloc_buffer((128, 8), scope="trn.sbuf")
+        const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A = T.alloc_buffer((32, 512, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(16 * 1024, 128) : (1@F, 1@P)]))
+        B = T.alloc_buffer((16, 512, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(2, 4, 1024, 128) : (1024@F, 2048@F, 1@F, 1@P)]))
+        C = T.alloc_buffer((1, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(1, 128) : (1@F, 1@P)]))
+        for i in range(2):
+            Tx.unary_reduce(B[0:16, 0:512, 0:128], C[0, 0:128], A[i * 16:i * 16 + 16, 0:512, 0:128], "sqrt", "sum", None, None, [0, 1], workspace={"const_bias": const_bias, "partial_reduce": partial_reduce})  # noqa: E501
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_partial_workspace_specify():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
+    C_shape = (1, 128)
+    C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        partial_reduce = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1), workspace={"partial_reduce": partial_reduce})  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        T.device_entry()
+        const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        partial_reduce = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        A = T.alloc_buffer((32, 512, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(16 * 1024, 128) : (1@F, 1@P)]))
+        B = T.alloc_buffer((16, 512, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(2, 4, 1024, 128) : (1024@F, 2048@F, 1@F, 1@P)]))
+        C = T.alloc_buffer((1, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(1, 128) : (1@F, 1@P)]))
+        for i in range(2):
+            Tx.unary_reduce(B[0:16, 0:512, 0:128], C[0, 0:128], A[i * 16:i * 16 + 16, 0:512, 0:128], "sqrt", "sum", None, None, [0, 1], workspace={"const_bias": const_bias, "partial_reduce": partial_reduce})  # noqa: E501
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_workspace_reuse():
+    src_shape = [512, 1024]
+    src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    scale = T.float32(2.0)
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.exp(C_sbuf, A_sbuf, bias=0.0, scale=scale, max_inst_size=1024)
+        Tx.exp(C_sbuf, C_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        T.device_entry()
+        const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096) : (1 @ P, 1 @ F)]))
+        C_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096) : (1 @ P, 1 @ F)]))
+        Tx.exp(C_sbuf[0:512, 0:1024], A_sbuf[0:512, 0:1024], T.float32(0.0), T.float32(2.0), workspace={"const_bias": const_bias}, max_inst_size=1024)  # noqa: E501
+        Tx.exp(C_sbuf[0:512, 0:1024], C_sbuf[0:512, 0:1024], None, None, workspace={"const_bias": const_bias})  # noqa: E501
+
+        # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_no_rewrite_with_existing_workspace():
+    src_shape = [128, 32, 4, 32]
+    src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        intermediate_buffer = T.alloc_buffer((128, 64), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.sum(B_sbuf, A_sbuf, axes=(1, 3), workspace={"partial_reduce": intermediate_buffer})
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], reduction)
+
+
+def test_no_rewrite_with_psum_output():
+    A_layout = TileLayout(S[(128, 128) : (1 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)])
+
+    C_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((128, 128), "float32", scope="trn.psum", layout=C_layout)
+        Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], gemm)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_reduction_trn.py

@@ -0,0 +1,283 @@
+# 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.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+opcode_map = {"sum": "add", "max": "max", "min": "min"}
+
+Tx_func_map = {"sum": Tx.sum, "max": Tx.max, "min": Tx.min}
+
+
+@pytest.mark.parametrize("op_type", ["sum", "max", "min"])
+def test_simple_reduction(op_type):
+    src_shape = [128, 512]
+    src_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 1]
+    dst_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
+
+    opcode = opcode_map[op_type]
+    tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def reduction() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        tx_func(B_sbuf, A_sbuf, axes=-1)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for b_loop in range(1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, 0], A_sbuf[p_loop, f_loop], opcode, False, -1)
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_with_multiple_axes():
+    src_shape = [128, 512, 4]
+    src_layout = TileLayout(S[(128, 512, 4) : (1 @ P, 1 @ F, 512 @ F)])
+    dst_shape = [128]
+    dst_layout = TileLayout(S[128 : 1 @ P])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.sum(B_sbuf, A_sbuf, axes=(1, 2), max_inst_size=2048)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for b_loop in range(1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, 0], A_sbuf[p_loop, f_loop], "add", False, -1)
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_in_loop():
+    src_shape = [128, 512, 4]
+    src_layout = TileLayout(S[(128, 512, 4) : (1 @ P, 4 @ F, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.sum(B_sbuf[:, i], A_sbuf[:, :, i], axes=-2)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, i], A_sbuf[p_loop, f_loop * 4 + i], "add", False, -1)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_two_stage():
+    src_shape = [128, 32, 4, 32]
+    src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.sum(B_sbuf, A_sbuf, axes=(1, 3))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        intermediate_buffer = T.alloc_buffer((128, 32), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for b_loop in range(4):
+            for reduction_b_loop in range(32):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                    for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                        T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop], A_sbuf[p_loop, reduction_b_loop * 128 + b_loop * 32 + f_loop], "add", False, -1)  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", False, -1)  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_with_guard():
+    src_shape = [512, 2048]
+    src_layout = TileLayout(S[(4, 128, 2048) : (2048 @ F, 1 @ P, 1 @ F)])
+    dst_shape = [512, 1]
+    dst_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            for j in range(4):
+                Tx.sum(B_sbuf[0: (i+1) * 128, 0], A_sbuf[0: (i+1) * 128, 0: (j+1) * 256], max_inst_size=512)  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        intermediate_buffer = T.alloc_buffer((128, 2), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for i, j in T.grid(4, 4):
+            for b_loop in range(4):
+                for reduction_b_loop in range(2):
+                    T.attr(0, "tensorized_nki_instruction", 1)
+                    for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                        for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                            if (
+                                b_loop - i < 1
+                                and reduction_b_loop * 512 + f_loop < j * 256 + 256
+                            ):
+                                T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop], A_sbuf[p_loop, b_loop * 2048 + reduction_b_loop * 512 + f_loop], "add", T.bool(False), -1)  # noqa: E501
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(2, annotations={"nki_dim": "F"}):
+                        if b_loop - i < 1 and f_loop * 2 - j < 1:
+                            T.nki.tensorreduce(B_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_two_stage_workspace():
+    src_shape = [128, 32, 4, 32]
+    src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        intermediate_buffer = T.alloc_buffer((128, 64), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.sum(B_sbuf, A_sbuf, axes=(1, 3), workspace={"partial_reduce": intermediate_buffer})
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        intermediate_buffer = T.alloc_buffer((128, 64), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for b_loop in range(4):
+            for reduction_b_loop in range(32):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                    for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                        T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop], A_sbuf[p_loop, reduction_b_loop * 128 + b_loop * 32 + f_loop], "add", False, -1)  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", False, -1)  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_select_trn.py

@@ -0,0 +1,186 @@
+# 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.testing
+from tvm.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+def test_select():
+    src_shape = [128, 512]
+    src_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def select() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.select(B_sbuf, A_sbuf, 0.0, lambda i, j: i < j)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "select"})
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in T.serial(0, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.affine_select(B_sbuf[p_loop, f_loop], p_loop < f_loop, A_sbuf[p_loop, f_loop], T.float32(0.0))  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": select})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_select_in_loop():
+    src_shape = [32, 128, 512]
+    src_layout = TileLayout(S[(32, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def select() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(2):
+            Tx.select(B_sbuf, A_sbuf[i*16, :, :], 0.0, lambda a, b: (i+1)* a < b)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "select"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.affine_select(B_sbuf[p_loop, f_loop], (i + 1) * p_loop < f_loop, A_sbuf[p_loop, i * 8192 + f_loop], T.float32(0.0))  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": select})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_select_expr_affine():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+
+    # fmt: off
+    @T.prim_func
+    def select() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.select(B_sbuf, A_sbuf, 0.0, lambda i, j: i < j)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "select"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for b_loop in T.serial(0, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.affine_select(B_sbuf[p_loop, b_loop * 512 + f_loop], b_loop * 128 + p_loop < f_loop, A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": select})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_select_with_guard():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+
+    # fmt: off
+    @T.prim_func
+    def select() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            for j in range(4):
+                Tx.select(B_sbuf[0: (i+1) * 128, 0: (j+1) * 128], A_sbuf[0: (i+1) * 128, 0: (j+1) * 128], 0.0, lambda a, b: a < b)  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "select"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for i, j, b_loop in T.grid(4, 4, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if b_loop - i < 1 and f_loop < j * 128 + 128:
+                        T.nki.affine_select(B_sbuf[p_loop, b_loop * 512 + f_loop], b_loop * 128 + p_loop < f_loop, A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": select})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_unary_trn.py

@@ -0,0 +1,288 @@
+# 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.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+Tx_func_map = {"reciprocal": Tx.reciprocal, "sqrt": Tx.sqrt, "memset": Tx.memset, "exp": Tx.exp}
+
+
+@pytest.mark.parametrize("op_type", ["reciprocal", "memset"])
+def test_simple_unary(op_type):
+    src_shape = [128, 512]
+    src_layout = T.TileLayout(T.S[(128, 512) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 512]
+    dst_layout = T.TileLayout(T.S[(128, 512) : (1 @ P, 1 @ F)])
+    tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        if op_type == "memset":
+            tx_func(B_sbuf, T.float32(0.0))
+        else:
+            tx_func(B_sbuf, A_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in T.serial(0, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if op_type == "reciprocal":
+                        T.nki.reciprocal(
+                            B_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop]
+                        )
+                    elif op_type == "memset":
+                        T.nki.memset(B_sbuf[p_loop, f_loop], 0.0)
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.parametrize("op_type", ["reciprocal", "memset"])
+def test_unary_in_a_loop(op_type):
+    src_shape = [1024, 512]
+    src_layout = T.TileLayout(T.S[(128, 4096) : (1 @ P, 1 @ F)])
+    dst_shape = [512, 512]
+    dst_layout = T.TileLayout(T.S[(128, 2048) : (1 @ P, 1 @ F)])
+
+    Tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        A_sbuf_view = A_sbuf.view(128, 8, 512)
+        B_sbuf_view = B_sbuf.view(128, 4, 512)
+        for i in range(4):
+            if op_type == "memset":
+                Tx_func(B_sbuf_view[:, i, :], T.float32(0.0))
+            else:
+                Tx_func(B_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :])
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        A_sbuf_view = T.decl_buffer((128, 4096), data=A_sbuf.data, scope="trn.sbuf", layout=None)
+        B_sbuf_view = T.decl_buffer((128, 2048), data=B_sbuf.data, scope="trn.sbuf", layout=None)
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if op_type == "reciprocal":
+                        T.nki.reciprocal(B_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop])  # noqa: E501
+                    elif op_type == "memset":
+                        T.nki.memset(B_sbuf[p_loop, i * 512 + f_loop], 0.0)
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_unary_complex1():
+    dst_layout = TileLayout(S[(32, 128, 256) : (256 @ F, 1 @ P, 1 @ F)])
+    dst_shape = [4096, 256]
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.memset(A_sbuf, T.float32(0.0))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        A_sbuf = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        for b_loop in T.serial(0, 16):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.memset(A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0))
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.parametrize("op_type", ["sqrt", "exp"])
+def test_unary_with_bias_scale(op_type):
+    src_shape = [512, 1024]
+    src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    bias_shape = [512, 1]
+    bias_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+    scale = T.float32(2.0)
+    tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(bias_shape, "float32", scope="trn.sbuf", layout=bias_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        tx_func(C_sbuf, A_sbuf, bias=B_sbuf, scale=scale)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        for b_loop in T.serial(0, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.activation(C_sbuf[p_loop, b_loop * 512  + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], op_type, B_sbuf[p_loop, b_loop//2], T.float32(2.0))  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.parametrize("op_type", ["sqrt", "exp"])
+def test_unary_with_bias_scale_2(op_type):
+    src_shape = [512, 1024]
+    src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    bias = T.float32(1.0)
+    scale = T.float32(2.0)
+    tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        tx_func(C_sbuf, A_sbuf, bias=bias, scale=scale)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(1.0))
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        for b_loop in T.serial(0, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.activation(C_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], op_type, const_bias[p_loop, f_loop], T.float32(2.0))  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_unary_with_guard():
+    src_shape = [512, 1024]
+    src_layout = TileLayout(S[(4, 128, 1024) : (1024 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    bias_shape = [512, 1]
+    bias_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+    scale = T.float32(2.0)
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(bias_shape, "float32", scope="trn.sbuf", layout=bias_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            for j in range(4):
+                Tx.sqrt(C_sbuf[0: (i+1) * 128, 0: (j+1)*256], A_sbuf[0: (i+1) * 128, 0: (j+1)*256], bias=B_sbuf[0: (i+1) * 128, 0], scale=scale)  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        for i, j, b_loop in T.grid(4, 4, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if b_loop // 2 - i < 1 and b_loop % 2 * 512 + f_loop < j * 256 + 256:
+                        T.nki.activation(C_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], "sqrt", B_sbuf[p_loop, b_loop // 2], T.float32(2.0))  # noqa: E501
+             # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()