chore: import upstream snapshot with attribution

python/tvm/backend/trn/__init__.py

@@ -0,0 +1,68 @@
+# 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.
+"""Trainium-owned TIRx modules."""
+
+from importlib import import_module
+
+_LAZY_SUBMODULES = {"layout", "op", "operator", "pipeline", "script", "target_tags", "transform"}
+
+
+def register_backend():
+    """Register Trainium-owned Python semantics."""
+    from tvm.tirx import compilation_pipeline  # pylint: disable=import-outside-toplevel
+    from tvm.tirx.script.builder import ir as builder_ir  # pylint: disable=import-outside-toplevel
+
+    for name, namespace in script_namespaces().items():
+        builder_ir.register_script_namespace(name, namespace)
+
+    import_module(f"{__name__}.operator.tile_primitive")
+    trn_pipeline = import_module(f"{__name__}.pipeline")
+    import_module(f"{__name__}.target_tags")
+    import_module(f"{__name__}.transform")
+    compilation_pipeline.register_tir_pipeline("trn", trn_pipeline.trn_pipeline)
+
+
+def script_namespace(op_wrapper=None):
+    """Return the Trainium TVMScript namespace object."""
+    from .script import NKINamespace  # pylint: disable=import-outside-toplevel
+
+    return NKINamespace(op_wrapper)
+
+
+def script_namespaces(op_wrapper=None, **_):
+    """Return Trainium-owned TVMScript namespaces."""
+    return {"nki": script_namespace(op_wrapper)}
+
+
+def __getattr__(name: str):
+    if name in _LAZY_SUBMODULES:
+        return import_module(f"{__name__}.{name}")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "layout",
+    "op",
+    "operator",
+    "pipeline",
+    "register_backend",
+    "script",
+    "script_namespace",
+    "script_namespaces",
+    "target_tags",
+    "transform",
+]

python/tvm/backend/trn/layout.py

@@ -0,0 +1,123 @@
+# 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.
+"""Trainium-specific TIRx layout helpers."""
+
+from __future__ import annotations
+
+import functools
+import operator
+import re
+
+import tvm
+from tvm.tirx.expr import Expr
+from tvm.tirx.layout import Axis, Iter, Layout, S, TileLayout
+
+_TRN_MEMORY_AXES = {"F", "P", "Bank"}
+_PSUM_MAX_ELEM_PER_BANK = 512
+
+
+def is_trainium_layout(layout: Layout | None) -> bool:
+    """Return whether a layout uses only Trainium memory axes."""
+    if not isinstance(layout, TileLayout):
+        return False
+    return not any(
+        iter.axis.is_memory() and iter.axis.name not in _TRN_MEMORY_AXES for iter in layout.shard
+    )
+
+
+def trainium_layout(annotation: str, shape: tuple[Expr], is_psum: bool = False) -> TileLayout:
+    """Create a Trainium tile layout from a PF annotation string and logical shape."""
+    analyzer = tvm.arith.Analyzer()
+    assert re.fullmatch(r"[PF]*", annotation), (
+        f"annotation {annotation} must be a string of 'P' and 'F'"
+    )
+    assert len(annotation) == len(shape), (
+        f"annotation {annotation} and shape {shape} must have the same length"
+    )
+    num_p_dim = annotation.count("P")
+    if num_p_dim == 1:
+        p_idx = annotation.index("P")
+        p_dim = shape[p_idx]
+        assert analyzer.can_prove(p_dim <= 128 or p_dim % 128 == 0), (
+            f"There is only 1 P in the annotation. Partition size {p_dim} must be less than "
+            "or equal to 128 or a multiple of 128"
+        )
+        if analyzer.can_prove(p_dim > 128):
+            annotation = "F" + annotation
+            shape = (p_dim // 128, *shape[:p_idx], 128, *shape[p_idx + 1 :])
+    elif num_p_dim > 1:
+        p_dim_prod = functools.reduce(
+            operator.mul, [s for s, c in zip(shape, annotation) if c == "P"]
+        )
+        assert analyzer.can_prove(p_dim_prod <= 128), (
+            f"There are {num_p_dim} Ps in the annotation. Partition size {p_dim_prod} must be "
+            "less than or equal to 128"
+        )
+
+    f_shape = [s for s, c in zip(shape, annotation) if c == "F"]
+    p_shape = [s for s, c in zip(shape, annotation) if c == "P"]
+    f_strides = Layout._get_default_strides(f_shape, 1)  # pylint: disable=protected-access
+    p_strides = Layout._get_default_strides(p_shape, 1)  # pylint: disable=protected-access
+    f_tile_layout = TileLayout(S[tuple(f_shape) : tuple(s @ Axis.F for s in f_strides)])
+    p_tile_layout = TileLayout(S[tuple(p_shape) : tuple(s @ Axis.P for s in p_strides)])
+    result = []
+    f_index = p_index = 0
+
+    for char in annotation:
+        if char == "F":
+            result.append(f_tile_layout.shard[f_index])
+            f_index += 1
+        else:
+            result.append(p_tile_layout.shard[p_index])
+            p_index += 1
+    if num_p_dim == 1 and analyzer.can_prove(p_dim > 128):
+        higher_p = result[0]
+        result = result[1:]
+        result = [*result[:p_idx], higher_p, *result[p_idx:]]
+
+    res = TileLayout.from_iters(result, [], {})
+    if is_psum:
+        res = to_psum_layout(res)
+    return res
+
+
+def to_psum_layout(layout: TileLayout) -> TileLayout:
+    """Convert a Trainium sbuf layout to its psum physical-bank layout."""
+    analyzer = tvm.arith.Analyzer()
+    shard = []
+    for iter in layout.shard:
+        if iter.axis.name == "F":
+            if analyzer.can_prove(iter.stride % _PSUM_MAX_ELEM_PER_BANK == 0):
+                stride = analyzer.simplify(iter.stride // _PSUM_MAX_ELEM_PER_BANK)
+                shard.append(Iter(iter.extent, stride, Axis.get("Bank")))
+            elif analyzer.can_prove(_PSUM_MAX_ELEM_PER_BANK % iter.stride == 0):
+                c = analyzer.simplify(_PSUM_MAX_ELEM_PER_BANK // iter.stride)
+                if analyzer.can_prove(iter.extent < c):
+                    shard.append(iter)
+                elif analyzer.can_prove(iter.extent % c == 0):
+                    shard.append(Iter(analyzer.simplify(iter.extent // c), 1, Axis.get("Bank")))
+                    shard.append(Iter(c, iter.stride, Axis.get("F")))
+                else:
+                    raise ValueError(f"layout {layout} can not be converted to psum layout")
+            else:
+                raise ValueError(f"layout {layout} can not be converted to psum layout")
+        else:
+            shard.append(iter)
+    return TileLayout.from_iters(shard, [], {})
+
+
+__all__ = ["is_trainium_layout", "to_psum_layout", "trainium_layout"]

python/tvm/backend/trn/op.py

@@ -0,0 +1,153 @@
+# 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.
+"""Trainium-owned NKI intrinsic Python wrappers."""
+
+from __future__ import annotations
+
+from tvm.tirx.op import call_intrin
+
+
+def nki_load(res, data):
+    return call_intrin("", "tirx.nki.load", res, data)
+
+
+def nki_store(res, data):
+    return call_intrin("", "tirx.nki.store", res, data)
+
+
+def nki_tensor_copy(res, data):
+    return call_intrin("", "tirx.nki.tensor_copy", res, data)
+
+
+def nki_matmul(res, lhs, rhs, accum=True):
+    return call_intrin("", "tirx.nki.matmul", res, lhs, rhs, accum)
+
+
+def nki_activation(result, data, opcode, bias=0.0, scale=1.0):
+    return call_intrin("", "tirx.nki.activation", result, data, opcode, bias, scale)
+
+
+def nki_reciprocal(result, data):
+    return call_intrin("", "tirx.nki.reciprocal", result, data)
+
+
+def nki_tensorreduce(result, data, opcode, negate, *axes):
+    return call_intrin("", "tirx.nki.tensorreduce", result, data, opcode, negate, *axes)
+
+
+def nki_tensortensor(result, operand0, operand1, opcode):
+    return call_intrin("", "tirx.nki.tensortensor", result, operand0, operand1, opcode)
+
+
+def nki_tensorscalar(result, operand0, operand1, opcode, reverse=False):
+    return call_intrin("", "tirx.nki.tensorscalar", result, operand0, operand1, opcode, reverse)
+
+
+def nki_memset(result, value):
+    return call_intrin("", "tirx.nki.memset", result, value)
+
+
+def nki_activation_reduce(reduce_res, act_res, data, opcode, reduce_opcode, bias=0.0, scale=1.0):
+    return call_intrin(
+        "",
+        "tirx.nki.activation_reduce",
+        reduce_res,
+        act_res,
+        data,
+        opcode,
+        reduce_opcode,
+        bias,
+        scale,
+    )
+
+
+def nki_tensorscalar_reduce(
+    reduce_res, tensorscalar_res, operand0, operand1, opcode, reduce_opcode, reverse=False
+):
+    return call_intrin(
+        "",
+        "tirx.nki.tensorscalar_reduce",
+        reduce_res,
+        tensorscalar_res,
+        operand0,
+        operand1,
+        opcode,
+        reduce_opcode,
+        reverse,
+    )
+
+
+def nki_identity(result, size):
+    return call_intrin("", "tirx.nki.identity", result, size)
+
+
+def nki_scalar_tensor_tensor(
+    result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
+):
+    return call_intrin(
+        "",
+        "tirx.nki.scalar_tensor_tensor",
+        result,
+        data,
+        operand0,
+        operand1,
+        opcode0,
+        opcode1,
+        reverse0,
+        reverse1,
+    )
+
+
+def nki_scalar_tensor_scalar(
+    result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
+):
+    return call_intrin(
+        "",
+        "tirx.nki.scalar_tensor_scalar",
+        result,
+        data,
+        operand0,
+        operand1,
+        opcode0,
+        opcode1,
+        reverse0,
+        reverse1,
+    )
+
+
+def nki_affine_select(result, pred, true_value, false_value):
+    return call_intrin("", "tirx.nki.affine_select", result, pred, true_value, false_value)
+
+
+__all__ = [
+    "nki_activation",
+    "nki_activation_reduce",
+    "nki_affine_select",
+    "nki_identity",
+    "nki_load",
+    "nki_matmul",
+    "nki_memset",
+    "nki_reciprocal",
+    "nki_scalar_tensor_scalar",
+    "nki_scalar_tensor_tensor",
+    "nki_store",
+    "nki_tensor_copy",
+    "nki_tensorreduce",
+    "nki_tensorscalar",
+    "nki_tensorscalar_reduce",
+    "nki_tensortensor",
+]

python/tvm/backend/trn/operator/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+"""Trainium backend operator package.
+
+Loaded by the Trainium backend registration hook.
+"""
+
+__all__ = ["tile_primitive"]

python/tvm/backend/trn/operator/tile_primitive/__init__.py

@@ -0,0 +1,25 @@
+# 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.
+
+from .binary import *
+from .compose_op import *
+from .copy import *
+from .gemm import *
+from .private_alloc import *
+from .reduction import *
+from .select import *
+from .unary import *

python/tvm/backend/trn/operator/tile_primitive/binary/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+
+from .default import *
+from .utils import *

python/tvm/backend/trn/operator/tile_primitive/binary/default.py

@@ -0,0 +1,124 @@
+# 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.
+
+"""Implementation of binary operator dispatches."""
+
+from tvm.script import tirx as T
+from tvm.tirx import FloatImm, PrimFunc
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import MapOpType
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import init_analyzer, nki_dim
+from ..instruction_generator import InstructionGenerator
+from .utils import InstType, binary_map_ops, try_find_inst_nary
+
+
+def binary_trn(
+    op: TilePrimitiveCall, binary_op: MapOpType, sctx: DispatchContext
+) -> PrimFunc | None:
+    """Generate a binary operation schedule for Trainium."""
+    if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
+        fail("requires Trainium target and thread exec_scope")
+
+    assert binary_op in binary_map_ops, f"Unsupported binary operation {binary_op}"
+
+    # Initialize analyzer and buffer regions
+    analyzer = init_analyzer(sctx)
+    _dst, _src1, _src2 = op.args
+
+    # Find instruction parameters
+    inst_gen = InstructionGenerator([_dst, _src1, _src2], analyzer)
+    inst_repr, inst_types, reverse = try_find_inst_nary(_dst, [_src1, _src2], analyzer, inst_gen)
+    # Handle operand swapping if needed
+    if reverse[0]:
+        _src1, _src2 = _src2, _src1
+
+    # Extract buffers and constants
+    CONST = _src2 if isinstance(_src2, FloatImm) else None
+    dst, src1 = _dst.buffer, _src1.buffer
+    src2 = None if CONST is not None else _src2.buffer
+
+    p_var = T.Var("P", "int32")
+    b_var = T.Var("B", "int32")
+    f_var = T.Var("F", "int32")
+    p_size = dst.layout.size("P")
+    inst_size_limit = op.config.get("max_inst_size", 512)
+    inst_repr.bound_inst_size(inst_size_limit, analyzer)
+    inst_gen.bind_inst_iter(_dst, p_var, p_size, 1, False)
+    inst_gen.bind_inst_iter(_dst, f_var, inst_repr.size, inst_repr.stride, True)
+    b_extent = inst_gen.fill_in_block_dim(_dst, b_var)
+    # Setup execution parameters
+    opcode = binary_map_ops[binary_op]
+
+    # Select appropriate NKI function based on instruction type
+    _func = T.nki.tensortensor if inst_types[0] == InstType.TENSOR_TENSOR else T.nki.tensorscalar
+
+    def func(*args):
+        return _func(*args, reverse[0]) if inst_types[0] == InstType.TENSOR_SCALAR else _func(*args)
+
+    # Define the implementation function
+    @T.prim_func
+    def impl():
+        for b_loop in T.serial(0, b_extent):
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                    for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
+                        inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, b_var: b_loop})
+
+                        if inst_gen.make_guard(_dst):
+                            dst_indices = T.meta_var(inst_gen.generate_indices(_dst))
+                            src1_indices = T.meta_var(inst_gen.generate_indices(_src1))
+                            if CONST is None:
+                                src2_indices = T.meta_var(inst_gen.generate_indices(_src2))
+                                T.evaluate(
+                                    func(
+                                        dst[tuple(dst_indices)],
+                                        src1[tuple(src1_indices)],
+                                        src2[tuple(src2_indices)],
+                                        opcode,
+                                    )
+                                )
+                            else:
+                                T.evaluate(
+                                    func(
+                                        dst[tuple(dst_indices)],
+                                        src1[tuple(src1_indices)],
+                                        CONST,
+                                        opcode,
+                                    )
+                                )
+
+    return impl
+
+
+# ---------------------------------------------------------------------------
+# Registration: bind each binary op name to its TRN schedule candidates.
+# ---------------------------------------------------------------------------
+from tvm.tirx.operator.tile_primitive import register_dispatch  # noqa: E402
+
+for _op_name, _op_type in {
+    "add": MapOpType.ADD,
+    "sub": MapOpType.SUB,
+    "mul": MapOpType.MUL,
+    "maximum": MapOpType.MAX,
+    "minimum": MapOpType.MIN,
+}.items():
+
+    @register_dispatch(_op_name, "trn", variant="binary", priority=0)
+    def _binary_dispatch(op, sctx, _ty=_op_type):
+        return binary_trn(op, _ty, sctx)

python/tvm/backend/trn/operator/tile_primitive/binary/utils.py

@@ -0,0 +1,227 @@
+# 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.
+
+"""Shared helpers for binary operator dispatches on TRN targets."""
+
+from enum import Enum
+
+from tvm.arith.analyzer import Analyzer
+from tvm.backend.trn.layout import is_trainium_layout
+from tvm.tirx import BufferRegion, FloatImm
+from tvm.tirx.operator.tile_primitive.common import MapOpType
+
+from ..dim_utils import get_ewise_dim_map
+from ..instruction_generator import InstructionGenerator
+
+binary_map_ops = {
+    MapOpType.ADD: "add",
+    MapOpType.SUB: "sub",
+    MapOpType.MUL: "mul",
+    MapOpType.MAX: "max",
+    MapOpType.MIN: "min",
+}
+
+
+class InstType(Enum):
+    TENSOR_TENSOR = 0
+    TENSOR_SCALAR = 1
+
+
+def try_find_inst_nary(
+    _dst: BufferRegion,
+    _srcs: list[BufferRegion | FloatImm],
+    analyzer: Analyzer,
+    inst_gen: InstructionGenerator,
+    allowed_f_dim_dst: tuple[int] | None = None,
+    allowed_f_dim_srcs: tuple[tuple[int]] | None = None,
+    allow_first_op_tensortensor: bool = True,
+):
+    """Find instruction parameters for n-ary operations."""
+    # Validate inputs and handle source swapping if needed
+    assert not (isinstance(_srcs[0], FloatImm) and isinstance(_srcs[1], FloatImm)), (
+        "Nary operation does not support taking all FloatImm sources"
+    )
+    assert 2 <= len(_srcs) <= 3, "Only 2-3 sources are supported for nary operation"
+
+    if isinstance(_srcs[0], FloatImm):
+        _srcs[0], _srcs[1] = _srcs[1], _srcs[0]
+        reverse = [True] + [False] * (len(_srcs) - 2)
+    else:
+        reverse = [False] * (len(_srcs) - 1)
+
+    # Extract buffers and validate properties
+    dst, srcs = (
+        _dst.buffer,
+        [_src.buffer if isinstance(_src, BufferRegion) else None for _src in _srcs],
+    )
+    dst_region = _dst.region
+
+    valid_buffers = all(
+        [
+            dst.layout and all(src.layout for src in srcs if src is not None),
+            is_trainium_layout(dst.layout),
+            all(is_trainium_layout(src.layout) for src in srcs if src is not None),
+            dst.scope() == "trn.sbuf",
+            all(src.scope() in ["trn.sbuf", "trn.psum"] for src in srcs if src is not None),
+        ]
+    )
+
+    if not valid_buffers:
+        raise ValueError(f"Invalid buffer region: dst: {_dst}, srcs: {_srcs}")
+
+    # Check non-unit extents
+    dst_non_unit_extent = [r.extent for r in dst_region if r.extent != 1]
+
+    # Handle broadcasting between first two sources
+    if not isinstance(_srcs[1], FloatImm):
+        src0_extent = [r.extent for r in _srcs[0].region]
+        src1_extent = [r.extent for r in _srcs[1].region]
+        shared_dim_num = min(len(src0_extent), len(src1_extent))
+
+        # Check for various broadcasting patterns and swap sources if needed
+        dims_equal = all(
+            analyzer.can_prove(e0 == e1)
+            for e0, e1 in zip(src0_extent[-shared_dim_num:], src1_extent[-shared_dim_num:])
+        )
+        if dims_equal:
+            if len(src0_extent) < len(src1_extent) and not all(
+                analyzer.can_prove(e1 == 1) for e1 in src1_extent[:-shared_dim_num]
+            ):
+                _srcs[0], _srcs[1] = _srcs[1], _srcs[0]
+                reverse[0] = True
+        elif all(
+            analyzer.can_prove(e0 == e1) or analyzer.can_prove(e0 == 1)
+            for e0, e1 in zip(src0_extent[-shared_dim_num:], src1_extent[-shared_dim_num:])
+        ):
+            _srcs[0], _srcs[1] = _srcs[1], _srcs[0]
+            reverse[0] = True
+            assert shared_dim_num == len(src0_extent) or all(
+                analyzer.can_prove(e0 == 1) for e0 in src0_extent[:-shared_dim_num]
+            ), f"Shape mismatch: src0: {_srcs[0]}, src1: {_srcs[1]}"
+        elif all(
+            analyzer.can_prove(e0 == e1) or analyzer.can_prove(e1 == 1)
+            for e0, e1 in zip(src0_extent[-shared_dim_num:], src1_extent[-shared_dim_num:])
+        ):
+            assert shared_dim_num == len(src1_extent) or all(
+                analyzer.can_prove(e1 == 1) for e1 in src1_extent[:-shared_dim_num]
+            ), f"Shape mismatch: src0: {_srcs[0]}, src1: {_srcs[1]}"
+        else:
+            raise ValueError(f"Shape mismatch: src0: {_srcs[0]}, src1: {_srcs[1]}")
+
+    # Verify src0 and dst have matching non-unit dimensions
+    src0_non_unit_extent = [r.extent for r in _srcs[0].region if r.extent != 1]
+    valid_shapes = all(
+        [
+            len(src0_non_unit_extent) == len(dst_non_unit_extent),
+            all(
+                analyzer.can_prove_equal(s, d)
+                for s, d in zip(src0_non_unit_extent, dst_non_unit_extent)
+            ),
+        ]
+    )
+
+    assert valid_shapes, "the larger between src0 and src1 must have the same shape as dst"
+
+    # Identify broadcast dimensions for each source after src0
+    src0_extent = [r.extent for r in _srcs[0].region]
+    dst_to_src0_dim_map = get_ewise_dim_map(_dst, _srcs[0], analyzer)
+    inst_gen.link_buffer_regions(_dst, _srcs[0], dst_to_src0_dim_map)
+
+    for src in _srcs[1:]:
+        if isinstance(src, FloatImm):
+            continue
+
+        src_extent = [r.extent for r in src.region]
+
+        # Check extra dimensions
+        assert len(src_extent) <= len(src0_extent) or all(
+            analyzer.can_prove(src_extent[i] == 1)
+            for i in range(len(src_extent) - len(src0_extent))
+        )
+
+        # Find broadcast dimensions
+        broadcast_dims = []
+        for i in range(1, min(len(src_extent), len(src0_extent)) + 1):
+            if analyzer.can_prove(src_extent[-i] != 1) and analyzer.can_prove(
+                src_extent[-i] != src0_extent[-i]
+            ):
+                raise ValueError(f"Shape mismatch: src0: {_srcs[0]}, src: {src}")
+            elif analyzer.can_prove(src_extent[-i] != src0_extent[-i]):
+                broadcast_dims.append(len(src0_extent) - i)
+
+        # Add leading dimensions
+        broadcast_dims += list(range(0, len(src0_extent) - len(src_extent)))
+
+        # Create dimension mapping and verify partition
+        src0_to_src_dim_map = {
+            i: i + len(src_extent) - len(src0_extent)
+            for i in range(len(src0_extent))
+            if i not in broadcast_dims
+        }
+        inst_gen.link_buffer_regions(_srcs[0], src, src0_to_src_dim_map)
+        assert inst_gen.check_partition_dim_match(_srcs[0], src), (
+            f"partition dimension mismatch: src0: {_srcs[0]}, src: {src}"
+        )
+
+    # Find instruction parameters for each source
+    inst_types = []
+    allowed_f_dim_srcs = [None] * len(_srcs) if allowed_f_dim_srcs is None else allowed_f_dim_srcs
+    inst_repr = inst_gen.find_max_inst_size_from_one_region(_dst, allowed_f_dim_dst)
+    for i, src in enumerate(_srcs):
+        if isinstance(src, FloatImm):
+            inst_types.append(InstType.TENSOR_SCALAR)
+            continue
+
+        allow_tt = allow_first_op_tensortensor or i != 0
+        inst_repr_non_bcast = inst_gen.fit_inst_tile_to_region(
+            inst_repr, src, allowed_f_dim_srcs[i]
+        )
+        inst_repr_bcast = inst_gen.fit_inst_tile_to_region(
+            inst_repr, src, allowed_f_dim_srcs[i], broadcast=True
+        )
+        if i == 0:
+            inst_repr = inst_repr_non_bcast
+            continue
+        plan = None
+        if not allow_tt:
+            plan = "tensorscalar"
+        else:
+            if (
+                inst_repr_bcast.stride == 1
+                and inst_repr_non_bcast.stride > 1
+                and inst_repr_bcast.size > 1
+            ):
+                plan = "tensorscalar"
+            elif (
+                inst_repr_bcast.stride > 1
+                and inst_repr_non_bcast.stride == 1
+                and inst_repr_non_bcast.size > 1
+            ):
+                plan = "tensortensor"
+            elif inst_repr_bcast.size > inst_repr_non_bcast.size:
+                plan = "tensorscalar"
+            else:
+                plan = "tensortensor"
+        if plan == "tensorscalar":
+            inst_type = InstType.TENSOR_SCALAR
+            inst_repr = inst_repr_bcast
+        else:
+            inst_type = InstType.TENSOR_TENSOR
+            inst_repr = inst_repr_non_bcast
+        inst_types.append(inst_type)
+
+    return inst_repr, inst_types, reverse

python/tvm/backend/trn/operator/tile_primitive/common.py

@@ -0,0 +1,43 @@
+# 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.
+
+"""Common utilities for TRN operator scheduling."""
+
+from tvm.arith.analyzer import Analyzer
+from tvm.tirx.operator.tile_primitive import DispatchContext
+
+# Used to generate the correct [:, None] for mask/predicate
+nki_dim = "nki_dim"
+
+
+def init_analyzer(sctx: DispatchContext):
+    """Initialize an analyzer with the dispatch context.
+
+    Parameters
+    ----------
+    sctx : DispatchContext
+        The dispatch context
+
+    Returns
+    -------
+    Analyzer :
+        The initialized analyzer
+    """
+    analyzer = Analyzer()
+    for v, r in sctx.var_range_map.items():
+        analyzer.bind(v, r)
+    return analyzer

python/tvm/backend/trn/operator/tile_primitive/compose_op/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+
+from .binary_chain import *
+from .binary_reduce import *
+from .compose_op import *
+from .reduce_negate import *
+from .unary_reduce import *

python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_chain.py

@@ -0,0 +1,125 @@
+# 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.
+
+"""Implementation of BinaryChain dispatch."""
+
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, PrimFunc, TilePrimitiveCall
+from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+from tvm.tirx.operator.tile_primitive.ops import BinaryChain
+
+from ..binary.utils import InstType, try_find_inst_nary
+from ..common import init_analyzer, nki_dim
+from ..instruction_generator import InstructionGenerator
+from .utils import opcode_table
+
+
+def binary_chain_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    """Generate a TRN schedule for binary chain operations."""
+    op = TilePrimitiveCall.downcast(op)
+    assert isinstance(op, BinaryChain), f"invalid operator downcast: {op}"
+
+    # Extract operation components
+    output = op.dsts[0]
+    srcs = op.srcs
+    reverse = [False, op.reverse1]
+    analyzer = init_analyzer(sctx)
+
+    # Find instruction patterns
+    inst_gen = InstructionGenerator([output, *srcs], analyzer)
+    inst_result = try_find_inst_nary(
+        output, srcs, analyzer, inst_gen, allow_first_op_tensortensor=False
+    )
+    inst_repr, inst_types, _reverse = inst_result
+
+    # Generate axes and validate
+    assert inst_types[0] == InstType.TENSOR_SCALAR, (
+        "The first operator must be a tensor scalar operator"
+    )
+
+    # Handle input reversal if needed
+    reverse[0] = _reverse[0]
+    if reverse[0]:
+        srcs[0], srcs[1] = srcs[1], srcs[0]
+
+    p_var = T.Var("P", "int32")
+    b_var = T.Var("B", "int32")
+    f_var = T.Var("F", "int32")
+    p_size = output.buffer.layout.size("P")
+    inst_size_limit = op.config.get("max_inst_size", 512)
+    inst_repr.bound_inst_size(inst_size_limit, analyzer)
+    inst_gen.bind_inst_iter(output, p_var, p_size, 1, False)
+    inst_gen.bind_inst_iter(output, f_var, inst_repr.size, inst_repr.stride, True)
+    b_extent = inst_gen.fill_in_block_dim(output, b_var)
+
+    # Extract buffers and opcodes
+    _src, dst = srcs[0].buffer, output.buffer
+    opcode0, opcode1 = opcode_table[op.op0], opcode_table[op.op1]
+
+    # Determine operation function based on instruction type
+    func = (
+        T.nki.scalar_tensor_scalar
+        if inst_types[1] == InstType.TENSOR_SCALAR
+        else T.nki.scalar_tensor_tensor
+    )
+
+    # Helper function to get source indices
+    def get_srcs(inst_gen):
+        return [
+            (
+                srcs[i].buffer[inst_gen.generate_indices(srcs[i])]
+                if isinstance(srcs[i], BufferRegion)
+                else srcs[i]
+            )
+            for i in range(len(srcs))
+        ]
+
+    # Create implementation
+    # fmt: off
+    @T.prim_func
+    def impl():
+        for b_loop in T.serial(0, b_extent):
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                    for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
+                        inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, b_var: b_loop})
+                        dst_indices = T.meta_var(inst_gen.generate_indices(output))
+                        srcs = T.meta_var(get_srcs(inst_gen))
+                        if inst_gen.make_guard(output):
+                            T.evaluate(func(dst[tuple(dst_indices)], *srcs, opcode0, opcode1, reverse[0], reverse[1]))  # noqa: E501
+    # fmt: on
+
+    return impl
+
+
+@register_dispatch(
+    "binary_chain",
+    "trn",
+    variant="default",
+    priority=10,
+    when=[
+        predicate(
+            "exec_scope",
+            lambda op, sctx: (
+                sctx.scope_kind == "thread",
+                f"unsupported exec_scope {sctx.scope_kind}",
+            ),
+        )
+    ],
+)
+def binary_chain_trn_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return binary_chain_trn(op, sctx)

python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_reduce.py

@@ -0,0 +1,168 @@
+# 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.
+
+"""Implementation of BinaryReduce dispatch."""
+
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, PrimFunc, TilePrimitiveCall
+from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+from tvm.tirx.operator.tile_primitive.ops import BinaryReduce
+
+from ..binary.utils import InstType, try_find_inst_nary
+from ..common import init_analyzer, nki_dim
+from ..dim_utils import get_reduction_dim_map
+from ..instruction_generator import InstructionGenerator
+from ..reduction.utils import generate_intermediate_buffer
+from .utils import opcode_table
+
+
+def binary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    """Generate a TRN schedule for binary reduction operations."""
+    op = TilePrimitiveCall.downcast(op)
+    assert isinstance(op, BinaryReduce), f"invalid operator downcast: {op}"
+
+    # Extract operation components
+    binary_output, reduce_output = op.dsts
+    binary_input1, binary_input2 = op.srcs
+    reduce_axes = op.reduce_axes
+    analyzer = init_analyzer(sctx)
+
+    # Normalize negative axes
+    reduce_axes = [i if i >= 0 else len(binary_output.buffer.shape) + i for i in reduce_axes]
+
+    # Find instruction patterns
+    inst_gen = InstructionGenerator(
+        [binary_output, binary_input1, binary_input2, reduce_output], analyzer
+    )
+    reduce_dim_map = get_reduction_dim_map(binary_output, reduce_output, reduce_axes, analyzer)
+    inst_gen.link_buffer_regions(binary_output, reduce_output, reduce_dim_map)
+    inst_repr, inst_type, reverse = try_find_inst_nary(
+        binary_output,
+        [binary_input1, binary_input2],
+        analyzer,
+        inst_gen,
+        allowed_f_dim_dst=reduce_axes,
+        allow_first_op_tensortensor=False,
+    )
+
+    # Apply instruction size limits
+    inst_size_limit = op.config.get("max_inst_size", None)
+    inst_repr.bound_inst_size(inst_size_limit, analyzer)
+
+    # Generate axes and validate
+    assert inst_type[0] == InstType.TENSOR_SCALAR, (
+        f"TensorTensor is not supported for vector reduce: {op}"
+    )
+
+    # Handle input reversal if needed
+    if reverse[0]:
+        binary_input1, binary_input2 = binary_input2, binary_input1
+
+    # Generate intermediate buffer for reduction if needed
+    p_var = T.Var("P", "int32")
+    f_var = T.Var("F", "int32")
+    reduction_b_var = T.Var("rB", "int32")
+    spatial_b_var = T.Var("sB", "int32")
+    p_size = binary_output.buffer.layout.size("P")
+    inst_gen.bind_inst_iter(binary_output, p_var, p_size, 1, False)
+    inst_gen.bind_inst_iter(binary_output, f_var, inst_repr.size, inst_repr.stride, True)
+    reduction_b_extent = inst_gen.fill_in_block_dim(binary_output, reduction_b_var, reduce_axes)
+    spatial_b_extent = inst_gen.fill_in_block_dim(binary_output, spatial_b_var)
+    if reduction_b_extent != 1:
+        intermediate_buffer = generate_intermediate_buffer(
+            reduce_output, reduction_b_extent, op.workspace, sctx
+        )
+
+    # Handle source 2 (either buffer region or constant)
+    CONST = binary_input2 if not isinstance(binary_input2, BufferRegion) else None
+    # Extract buffers and opcodes
+    src1, src2 = (
+        binary_input1.buffer,
+        (binary_input2.buffer if isinstance(binary_input2, BufferRegion) else None),
+    )
+    dst1, dst2 = binary_output.buffer, reduce_output.buffer
+    binary_opcode, reduce_opcode = opcode_table[op.binary_op], opcode_table[op.reduce_op]
+    # Create appropriate implementation based on intermediate buffer requirement
+    if reduction_b_extent == 1:
+        # Direct implementation without intermediate buffer
+        # fmt: off
+        @T.prim_func
+        def impl():
+            for b_loop in T.serial(0, spatial_b_extent):
+                with T.attr(0, "tensorized_nki_instruction", 1):
+                    for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                        for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
+                            inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, spatial_b_var: b_loop})  # noqa: E501
+                            src_1_indices = T.meta_var(inst_gen.generate_indices(binary_input1))
+                            vec_dst_idx = T.meta_var(inst_gen.generate_indices(binary_output))
+                            reduce_dst_idx = T.meta_var(inst_gen.generate_indices(reduce_output))
+                            if inst_gen.make_guard(binary_output):
+                                if CONST is None:
+                                    src_2_indices = T.meta_var(inst_gen.generate_indices(binary_input2))  # noqa: E501
+                                    T.nki.tensorscalar_reduce(dst2[tuple(reduce_dst_idx)], dst1[tuple(vec_dst_idx)], src1[tuple(src_1_indices)], src2[tuple(src_2_indices)], binary_opcode, reduce_opcode, reverse[0])  # noqa: E501
+                                else:
+                                    T.nki.tensorscalar_reduce(dst2[tuple(reduce_dst_idx)], dst1[tuple(vec_dst_idx)], src1[tuple(src_1_indices)], CONST, binary_opcode, reduce_opcode, reverse[0])  # noqa: E501
+        # fmt: on
+    else:
+        # Implementation with intermediate buffer
+        # fmt: off
+        @T.prim_func
+        def impl():
+            for b_loop in T.serial(0, spatial_b_extent):
+                for reduction_b_loop in T.serial(0, reduction_b_extent):
+                    with T.attr(0, "tensorized_nki_instruction", 1):
+                        for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                            for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
+                                inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, spatial_b_var: b_loop, reduction_b_var: reduction_b_loop})  # noqa: E501
+                                if inst_gen.make_guard(binary_output):
+                                    src_1_indices = T.meta_var(inst_gen.generate_indices(binary_input1))  # noqa: E501
+                                    vec_dst_idx = T.meta_var(inst_gen.generate_indices(binary_output))  # noqa: E501
+                                    if CONST is None:
+                                        src_2_indices = T.meta_var(inst_gen.generate_indices(binary_input2))  # noqa: E501
+                                        T.nki.tensorscalar_reduce(intermediate_buffer[p_loop, reduction_b_loop], dst1[tuple(vec_dst_idx)], src1[tuple(src_1_indices)], src2[tuple(src_2_indices)], binary_opcode, reduce_opcode, reverse[0])  # noqa: E501
+                                    else:
+                                        T.nki.tensorscalar_reduce(intermediate_buffer[p_loop, reduction_b_loop], dst1[tuple(vec_dst_idx)], src1[tuple(src_1_indices)], CONST, binary_opcode, reduce_opcode, reverse[0])  # noqa: E501
+                with T.attr(0, "tensorized_nki_instruction", 1):
+                    for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                        for f_loop in T.serial(0, reduction_b_extent, annotations={nki_dim: "F"}):
+                            inst_gen.set_bind_map_all({p_var: p_loop, spatial_b_var: b_loop})
+                            if inst_gen.make_guard(reduce_output):
+                                dst_2_indices = T.meta_var(inst_gen.generate_indices(reduce_output))
+                                T.nki.tensorreduce(dst2[tuple(dst_2_indices)], intermediate_buffer[p_loop, f_loop], reduce_opcode, False, -1)  # noqa: E501
+        # fmt: on
+
+    return impl
+
+
+# Rich dispatcher variants for TRN compose ops
+@register_dispatch(
+    "binary_reduce",
+    "trn",
+    variant="default",
+    priority=10,
+    when=[
+        predicate(
+            "exec_scope",
+            lambda op, sctx: (
+                sctx.scope_kind == "thread",
+                f"unsupported exec_scope {sctx.scope_kind}",
+            ),
+        )
+    ],
+)
+def binary_reduce_trn_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return binary_reduce_trn(op, sctx)

python/tvm/backend/trn/operator/tile_primitive/compose_op/compose_op.py

@@ -0,0 +1,47 @@
+# 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.
+
+"""Implementation of ComposeOp dispatch."""
+
+from tvm.tirx import PrimFunc, TilePrimitiveCall
+from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+
+
+def compose_op_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    """Generate a TRN schedule for compose operations."""
+    raise NotImplementedError(
+        "Generic compose_op must be lowered to specific compose ops before operator-level passes"
+    )
+
+
+@register_dispatch(
+    "compose_op",
+    "trn",
+    variant="default",
+    priority=10,
+    when=[
+        predicate(
+            "exec_scope",
+            lambda op, sctx: (
+                sctx.scope_kind == "thread",
+                f"unsupported exec_scope {sctx.scope_kind}",
+            ),
+        )
+    ],
+)
+def compose_op_trn_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return compose_op_trn(op, sctx)

python/tvm/backend/trn/operator/tile_primitive/compose_op/reduce_negate.py

@@ -0,0 +1,51 @@
+# 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.
+
+"""Implementation of ReduceNegate dispatch."""
+
+from tvm.tirx import PrimFunc, TilePrimitiveCall
+from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+from tvm.tirx.operator.tile_primitive.ops import ReduceNegate
+
+from ..reduction.utils import reduction_trn
+from .utils import optype_table
+
+
+def reduce_negate_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    """Generate a TRN schedule for reduce negate operations."""
+    op = TilePrimitiveCall.downcast(op)
+    assert isinstance(op, ReduceNegate), f"invalid operator downcast: {op}"
+    return reduction_trn(op, optype_table[op.reduce_op], sctx, negate=True)
+
+
+@register_dispatch(
+    "reduce_negate",
+    "trn",
+    variant="default",
+    priority=10,
+    when=[
+        predicate(
+            "exec_scope",
+            lambda op, sctx: (
+                sctx.scope_kind == "thread",
+                f"unsupported exec_scope {sctx.scope_kind}",
+            ),
+        )
+    ],
+)
+def reduce_negate_trn_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return reduce_negate_trn(op, sctx)

python/tvm/backend/trn/operator/tile_primitive/compose_op/unary_reduce.py

@@ -0,0 +1,170 @@
+# 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.
+
+"""Implementation of UnaryReduce dispatch."""
+
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, PrimFunc, TilePrimitiveCall
+from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+from tvm.tirx.operator.tile_primitive.ops import UnaryReduce
+
+from ..binary.utils import try_find_inst_nary
+from ..common import init_analyzer, nki_dim
+from ..dim_utils import get_reduction_dim_map
+from ..instruction_generator import InstructionGenerator
+from ..reduction.utils import generate_intermediate_buffer
+from ..unary.utils import get_const_bias_tensor, try_find_inst_unary
+from .utils import opcode_table
+
+
+def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    """Generate a TRN schedule for unary reduction operations."""
+    op = TilePrimitiveCall.downcast(op)
+    assert isinstance(op, UnaryReduce), f"invalid operator downcast: {op}"
+
+    # Extract operation components
+    unary_output, reduce_output = op.dsts
+    unary_input, bias, scale = op.srcs
+    analyzer = init_analyzer(sctx)
+
+    # Normalize axes and default values
+    reduce_axes = [i if i >= 0 else len(unary_output.buffer.shape) + i for i in op.reduce_axes]
+    scale = 1.0 if scale is None else scale
+    bias = 0.0 if bias is None else bias
+
+    inst_gen = InstructionGenerator([unary_output, unary_input, bias, reduce_output], analyzer)
+    reduce_dim_map = get_reduction_dim_map(unary_output, reduce_output, reduce_axes, analyzer)
+    inst_gen.link_buffer_regions(unary_output, reduce_output, reduce_dim_map)
+    # Find instruction patterns based on bias type
+    if isinstance(bias, BufferRegion):
+        inst_repr, _, _ = try_find_inst_nary(
+            unary_output,
+            [unary_input, bias],
+            analyzer,
+            inst_gen,
+            allow_first_op_tensortensor=False,
+            allowed_f_dim_dst=reduce_axes,
+        )
+    else:
+        inst_repr = try_find_inst_unary(
+            unary_output, unary_input, analyzer, inst_gen, allowed_f_dim_dst=reduce_axes
+        )
+
+    # Apply instruction size limits
+    inst_size_limit = op.config.get("max_inst_size", None)
+    inst_repr.bound_inst_size(inst_size_limit, analyzer)
+
+    p_var = T.Var("P", "int32")
+    f_var = T.Var("F", "int32")
+    reduction_b_var = T.Var("rB", "int32")
+    spatial_b_var = T.Var("sB", "int32")
+    p_size = unary_output.buffer.layout.size("P")
+    inst_gen.bind_inst_iter(unary_output, p_var, p_size, 1, False)
+    inst_gen.bind_inst_iter(unary_output, f_var, inst_repr.size, inst_repr.stride, True)
+    reduction_b_extent = inst_gen.fill_in_block_dim(unary_output, reduction_b_var, reduce_axes)
+    spatial_b_extent = inst_gen.fill_in_block_dim(unary_output, spatial_b_var)
+    if reduction_b_extent != 1:
+        intermediate_buffer = generate_intermediate_buffer(
+            reduce_output, reduction_b_extent, op.workspace, sctx
+        )
+    # Extract buffers and opcodes
+    src, dst1, dst2 = unary_input.buffer, unary_output.buffer, reduce_output.buffer
+    unary_opcode = opcode_table[op.unary_op]
+    reduce_opcode = opcode_table[op.reduce_op]
+
+    # Handle bias buffer
+    bias_buffer = (
+        bias.buffer
+        if isinstance(bias, BufferRegion)
+        else get_const_bias_tensor(bias, (p_size, inst_repr.size), dst1.dtype, op.workspace, sctx)
+    )
+
+    # Create appropriate implementation based on intermediate buffer requirement
+    if reduction_b_extent == 1:
+        # Direct implementation without intermediate buffer
+        # fmt: off
+        @T.prim_func
+        def impl():
+            for b_loop in T.serial(0, spatial_b_extent):
+                with T.attr(0, "tensorized_nki_instruction", 1):
+                    for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                        for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
+                            inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, spatial_b_var: b_loop})  # noqa: E501
+                            src_1_indices = T.meta_var(inst_gen.generate_indices(unary_input))
+                            dst_1_indices = T.meta_var(inst_gen.generate_indices(unary_output))
+                            dst_2_indices = T.meta_var(inst_gen.generate_indices(reduce_output))
+                            if inst_gen.make_guard(unary_output):
+                                if isinstance(bias, BufferRegion):
+                                    src_bias_indices = T.meta_var(inst_gen.generate_indices(bias))
+                                    T.evaluate(T.nki.activation_reduce(dst2[tuple(dst_2_indices)], dst1[tuple(dst_1_indices)], src[tuple(src_1_indices)], unary_opcode, reduce_opcode, bias_buffer[tuple(src_bias_indices)], scale))  # noqa: E501
+                                else:
+                                    T.evaluate(T.nki.activation_reduce(dst2[tuple(dst_2_indices)], dst1[tuple(dst_1_indices)], src[tuple(src_1_indices)], unary_opcode, reduce_opcode, bias_buffer[p_loop, f_loop], scale))  # noqa: E501
+        # fmt: on
+
+        import tvm
+
+        mod = tvm.IRModule({"main": impl})
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        return mod["main"]
+    else:
+        # fmt: off
+        @T.prim_func
+        def impl():
+            for b_loop in T.serial(0, spatial_b_extent):
+                for reduction_b_loop in T.serial(0, reduction_b_extent):
+                    with T.attr(0, "tensorized_nki_instruction", 1):
+                        for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                            for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
+                                inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, spatial_b_var: b_loop, reduction_b_var: reduction_b_loop})  # noqa: E501
+                                src_1_indices = T.meta_var(inst_gen.generate_indices(unary_input))
+                                dst_1_indices = T.meta_var(inst_gen.generate_indices(unary_output))
+                                if inst_gen.make_guard(unary_output):
+                                    if isinstance(bias, BufferRegion):
+                                        src_bias_indices = T.meta_var(inst_gen.generate_indices(bias))  # noqa: E501
+                                        T.evaluate(T.nki.activation_reduce(intermediate_buffer[p_loop, reduction_b_loop], dst1[tuple(dst_1_indices)], src[tuple(src_1_indices)], unary_opcode, reduce_opcode, bias_buffer[tuple(src_bias_indices)], scale))  # noqa: E501
+                                    else:
+                                        T.evaluate(T.nki.activation_reduce(intermediate_buffer[p_loop, reduction_b_loop], dst1[tuple(dst_1_indices)], src[tuple(src_1_indices)], unary_opcode, reduce_opcode, bias_buffer[p_loop, f_loop], scale))  # noqa: E501
+                with T.attr(0, "tensorized_nki_instruction", 1):
+                    for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                        for f_loop in T.serial(0, reduction_b_extent, annotations={nki_dim: "F"}):
+                            inst_gen.set_bind_map_all({p_var: p_loop, spatial_b_var: b_loop})
+                            if inst_gen.make_guard(reduce_output):
+                                dst_2_indices = T.meta_var(inst_gen.generate_indices(reduce_output))
+                                # TODO: we should use nki.activation_reduce as second stage reduction  # noqa: E501
+                                T.evaluate(T.nki.tensorreduce(dst2[tuple(dst_2_indices)], intermediate_buffer[p_loop, f_loop], reduce_opcode, False, -1))  # noqa: E501
+        # fmt: on
+
+        return impl
+
+
+@register_dispatch(
+    "unary_reduce",
+    "trn",
+    variant="default",
+    priority=10,
+    when=[
+        predicate(
+            "exec_scope",
+            lambda op, sctx: (
+                sctx.scope_kind == "thread",
+                f"unsupported exec_scope {sctx.scope_kind}",
+            ),
+        )
+    ],
+)
+def unary_reduce_trn_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return unary_reduce_trn(op, sctx)

python/tvm/backend/trn/operator/tile_primitive/compose_op/utils.py

@@ -0,0 +1,41 @@
+# 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.
+
+"""Shared helpers for compose operator dispatches."""
+
+from tvm.ir import Op
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
+
+# Operation code mappings
+opcode_table = {
+    Op.get("tirx.tile.add"): "add",
+    Op.get("tirx.tile.sub"): "sub",
+    Op.get("tirx.tile.mul"): "mul",
+    Op.get("tirx.tile.maximum"): "max",
+    Op.get("tirx.tile.minimum"): "min",
+    Op.get("tirx.tile.sqrt"): "sqrt",
+    Op.get("tirx.tile.sum"): "add",
+    Op.get("tirx.tile.max"): "max",
+    Op.get("tirx.tile.min"): "min",
+    Op.get("tirx.tile.exp"): "exp",
+}
+
+optype_table = {
+    Op.get("tirx.tile.sum"): ReduceOpType.SUM,
+    Op.get("tirx.tile.max"): ReduceOpType.MAX,
+    Op.get("tirx.tile.min"): ReduceOpType.MIN,
+}

python/tvm/backend/trn/operator/tile_primitive/copy/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+from .default import *

python/tvm/backend/trn/operator/tile_primitive/copy/default.py

@@ -0,0 +1,304 @@
+# 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.
+
+"""Implementation of copy operator dispatchs."""
+
+from tvm.backend.trn.layout import is_trainium_layout
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc
+from tvm.tirx.operator.tile_primitive import (
+    DispatchContext,
+    fail,
+    predicate,
+    register_dispatch,
+)
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import init_analyzer, nki_dim
+from ..dim_utils import get_ewise_dim_map
+from ..instruction_generator import InstructionGenerator
+from ..workspace_utils import check_workspace_buffer, largest_psum_per_bank, max_psum_banks
+
+
+def transpose_schedule(
+    op: TilePrimitiveCall, inst_gen: InstructionGenerator, sctx: DispatchContext
+) -> PrimFunc | None:
+    dst_region, src_region = op.args
+    assert src_region.buffer.scope() != "trn.psum", "Transpose on psum buffer is not supported"
+
+    inst_repr_dst, inst_repr_src = inst_gen.find_max_inst_size_transpose(dst_region, src_region)
+
+    lhs_f = T.Var("lhs_F", "int32")
+    lhs_p = T.Var("lhs_P", "int32")
+    dst_f = T.Var("dst_F", "int32")
+    b_var = T.Var("B", "int32")
+    extend_b = T.Var("extend_B", "int32")
+    p_size = src_region.buffer.layout.size("P")
+    lhs_f_size = dst_region.buffer.layout.size("P")
+    rhs_f_size = p_size
+    inst_gen.bind_inst_iter(
+        src_region, lhs_f, inst_repr_src.size, inst_repr_src.stride, is_free_dim=True
+    )
+    inst_gen.bind_inst_iter(
+        dst_region,
+        dst_f,
+        inst_repr_dst.size,
+        inst_repr_dst.stride,
+        is_free_dim=True,
+        no_propagate=True,
+    )
+    inst_gen.bind_inst_iter(src_region, lhs_p, p_size, 1, is_free_dim=False, no_propagate=True)
+    if dst_region.buffer.scope() == "trn.sbuf":
+        max_extend_num = (
+            inst_gen.find_max_inst_size_from_one_region(
+                dst_region, min_stride=inst_repr_dst.stride
+            ).size
+            // rhs_f_size
+        )
+        max_elem_in_a_bank = largest_psum_per_bank // rhs_f_size
+        if max_extend_num < max_elem_in_a_bank:
+            extend_len = max_extend_num
+        elif max_extend_num % max_elem_in_a_bank == 0:
+            extend_len = max_elem_in_a_bank
+        else:
+            extend_len = 1
+        inst_gen.bind_inst_iter(
+            dst_region,
+            extend_b,
+            extend_len,
+            inst_repr_dst.stride * inst_repr_dst.size,
+            is_free_dim=True,
+        )
+    b_extent = inst_gen.fill_in_block_dim(dst_region, b_var)
+
+    if "identity" not in op.workspace:
+        assert sctx.alloc_only, (
+            "Identity tensor must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first."  # noqa: E501
+        )
+        identity_tensor = T.buffer(
+            (p_size, rhs_f_size), src_region.buffer.dtype, scope="trn.sbuf", buffer_name="identity"
+        )
+        sctx.add_alloc_buffer(identity_tensor)
+
+        @T.prim_func
+        def identity_init():
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                    for rhs_f_loop in T.serial(0, rhs_f_size, annotations={nki_dim: "F"}):
+                        T.evaluate(T.nki.identity(identity_tensor[p_loop, rhs_f_loop], p_size))
+            T.tvm_kernel_replace_point()
+
+        sctx.add_init_stmt(identity_init.body)
+    else:
+        identity_tensor = op.workspace["identity"]
+        check_workspace_buffer(identity_tensor, (p_size, rhs_f_size), "trn.sbuf")
+
+    dst_buffer = dst_region.buffer
+    src_buffer = src_region.buffer
+    if dst_buffer.scope() == "trn.psum":
+
+        @T.prim_func
+        def transpose_psum_output():
+            for b_loop in T.serial(0, b_extent):
+                with T.attr(0, "tensorized_nki_instruction", 1):
+                    for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                        for lhs_f_loop in T.serial(0, lhs_f_size, annotations={nki_dim: "lhs_F"}):
+                            for rhs_f_loop in T.serial(
+                                0, rhs_f_size, annotations={nki_dim: "rhs_F"}
+                            ):
+                                inst_gen.set_bind_map(
+                                    dst_region,
+                                    {b_var: b_loop, lhs_f: lhs_f_loop, dst_f: rhs_f_loop},
+                                )
+                                inst_gen.set_bind_map(
+                                    src_region, {b_var: b_loop, lhs_f: lhs_f_loop, lhs_p: p_loop}
+                                )
+                                src_indices = T.meta_var(inst_gen.generate_indices(src_region))
+                                dst_indices = T.meta_var(inst_gen.generate_indices(dst_region))
+                                src_guard = T.meta_var(inst_gen.make_guard(src_region))
+                                dst_guard = T.meta_var(inst_gen.make_guard(dst_region))
+                                if src_guard and dst_guard:
+                                    T.evaluate(
+                                        T.nki.matmul(
+                                            dst_buffer[tuple(dst_indices)],
+                                            src_buffer[tuple(src_indices)],
+                                            identity_tensor[p_loop, rhs_f_loop],
+                                        )
+                                    )
+
+        return transpose_psum_output
+
+    if "acc_psum" not in op.workspace:
+        assert sctx.alloc_only, (
+            "Accumulation psum buffer must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first."  # noqa: E501
+        )
+        acc_psum = T.buffer(
+            (max_psum_banks, p_size, largest_psum_per_bank),
+            "float32",
+            scope="trn.psum",
+            allocated_addr=(0, 0),
+            buffer_name="acc_psum",
+        )
+        sctx.add_alloc_buffer(acc_psum)
+        max_psum_slots = max_psum_banks
+    else:
+        acc_psum = op.workspace["acc_psum"]
+        check_workspace_buffer(acc_psum, (p_size, largest_psum_per_bank), "trn.psum")
+        max_psum_slots = acc_psum.shape[0]
+
+    # fmt: off
+    @T.prim_func
+    def transpose_sbuf_output():
+        for b_loop in T.serial(0, b_extent):
+            for extend_b_loop in T.serial(0, extend_len):
+                with T.attr(0, "tensorized_nki_instruction", 1):
+                    for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                        for lhs_f_loop in T.serial(0, lhs_f_size, annotations={nki_dim: "lhs_F"}):
+                            for rhs_f_loop in T.serial(0, rhs_f_size, annotations={nki_dim: "rhs_F"}):  # noqa: E501
+                                inst_gen.set_bind_map(src_region, {b_var: b_loop, lhs_f: lhs_f_loop, lhs_p: p_loop, extend_b: extend_b_loop})  # noqa: E501
+                                src_indices = T.meta_var(inst_gen.generate_indices(src_region))
+                                src_guard = T.meta_var(inst_gen.make_guard(src_region))
+                                if src_guard:
+                                    T.evaluate(T.nki.matmul(acc_psum[b_loop % max_psum_slots, lhs_f_loop,extend_b_loop * rhs_f_size + rhs_f_loop], src_buffer[tuple(src_indices)], identity_tensor[p_loop, rhs_f_loop]))  # noqa: E501
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                    for f_loop in T.serial(0, rhs_f_size * extend_len, annotations={nki_dim: "F"}):
+                        inst_gen.set_bind_map(dst_region, {b_var: b_loop, lhs_f: p_loop, dst_f: f_loop % rhs_f_size, extend_b: f_loop // rhs_f_size})  # noqa: E501
+                        dst_guard = T.meta_var(inst_gen.make_guard(dst_region))
+                        dst_indices = T.meta_var(inst_gen.generate_indices(dst_region))
+                        if dst_guard:
+                            T.evaluate(T.nki.tensor_copy(dst_buffer[tuple(dst_indices)], acc_psum[b_loop % max_psum_slots, p_loop, f_loop]))  # noqa: E501
+    # fmt: on
+    return transpose_sbuf_output
+
+
+def copy_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    """Schedule copy operation between global and shared memory on CUDA."""
+    # Basic validation checks
+    if sctx.scope_kind != "thread":
+        fail("requires thread exec_scope for TRN copy")
+
+    dst_region, src_region = op.args
+    src, dst = src_region.buffer, dst_region.buffer
+
+    # Check for valid buffer configurations
+    valid_config = all(
+        [
+            src.layout and dst.layout,
+            src.scope() in ["global", "trn.sbuf", "trn.psum"],
+            dst.scope() in ["global", "trn.sbuf", "trn.psum"],
+            src.scope() != "global" or dst.scope() != "global",
+            (src.scope() == "global" and isinstance(src.layout, T.TileLayout))
+            or (src.scope() in ["trn.sbuf", "trn.psum"] and is_trainium_layout(src.layout)),
+            (dst.scope() == "global" and isinstance(dst.layout, T.TileLayout))
+            or (dst.scope() in ["trn.sbuf", "trn.psum"] and is_trainium_layout(dst.layout)),
+        ]
+    )
+
+    if not valid_config:
+        raise ValueError("Invalid buffer layout/scope for copy operation.")
+
+    analyzer = init_analyzer(sctx)
+    src_extent = [r.extent for r in src_region.region]
+    dst_extent = [r.extent for r in dst_region.region]
+
+    # Validate non-unit dimensions match
+    src_non_unit = [e for e in src_extent if e != 1]
+    dst_non_unit = [e for e in dst_extent if e != 1]
+    dims_match = len(src_non_unit) == len(dst_non_unit) and all(
+        analyzer.can_prove_equal(s, d) for s, d in zip(src_non_unit, dst_non_unit)
+    )
+
+    if not dims_match:
+        fail("shape mismatch between src and dst for TRN copy")
+
+    dim_map = get_ewise_dim_map(src_region, dst_region, analyzer)
+    inst_gen = InstructionGenerator([src_region, dst_region], analyzer)
+    inst_gen.link_buffer_regions(src_region, dst_region, dim_map)
+
+    if not inst_gen.check_partition_dim_match(src_region, dst_region):
+        return transpose_schedule(op, inst_gen, sctx)
+
+    if is_trainium_layout(src.layout):
+        inst = inst_gen.find_max_inst_size_from_one_region(src_region)
+        inst = inst_gen.fit_inst_tile_to_region(inst, dst_region)
+        src_to_dst = True
+    else:
+        inst = inst_gen.find_max_inst_size_from_one_region(dst_region)
+        inst = inst_gen.fit_inst_tile_to_region(inst, src_region)
+        src_to_dst = False
+
+    if src.scope() == "global":
+        func = T.nki.load
+    elif dst.scope() == "global":
+        func = T.nki.store
+    else:
+        func = T.nki.tensor_copy
+
+    if func == T.nki.tensor_copy:
+        inst_size_limit = op.config.get("max_inst_size", 512)
+        inst.bound_inst_size(inst_size_limit, analyzer)
+    else:
+        assert "max_inst_size" not in op.config, "max_inst_size is not supported for load/store"
+
+    p_var = T.Var("P", "int32")
+    f_var = T.Var("F", "int32")
+    b_var = T.Var("B", "int32")
+    if src_to_dst:
+        from_region, _to_region = src_region, dst_region
+    else:
+        from_region, _to_region = dst_region, src_region
+    p_size = from_region.buffer.layout.size("P")
+    inst_gen.bind_inst_iter(from_region, p_var, p_size, 1, is_free_dim=False)
+    inst_gen.bind_inst_iter(from_region, f_var, inst.size, inst.stride, is_free_dim=True)
+    b_extent = inst_gen.fill_in_block_dim(from_region, b_var)
+
+    # fmt: off
+    @T.prim_func
+    def impl():
+        # the additional b loop is to satisfy hardware instuction size limit
+        for b_loop in T.serial(0, b_extent):
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                    for f_loop in T.serial(0, inst.size, annotations={nki_dim: "F"}):
+                        inst_gen.set_bind_map_all({b_var: b_loop, p_var: p_loop, f_var: f_loop})
+                        if inst_gen.make_guard(dst_region):
+                            src_indices = T.meta_var(inst_gen.generate_indices(src_region))
+                            dst_indices = T.meta_var(inst_gen.generate_indices(dst_region))
+                            func(dst[tuple(dst_indices)], src[tuple(src_indices)])
+    # fmt: on
+    return impl
+
+
+# Rich dispatcher variant for TRN copy
+@register_dispatch(
+    "copy",
+    "trn",
+    variant="default",
+    priority=10,
+    when=[
+        predicate(
+            "exec_scope",
+            lambda op, sctx: (
+                sctx.scope_kind == "thread",
+                f"unsupported exec_scope {sctx.scope_kind}",
+            ),
+        )
+    ],
+)
+def copy_trn_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return copy_trn(op, sctx)

python/tvm/backend/trn/operator/tile_primitive/dim_utils.py

@@ -0,0 +1,262 @@
+# 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.
+
+"""Dimension mapping utilities for TRN operator scheduling."""
+
+from collections import namedtuple
+
+from tvm.arith.analyzer import Analyzer
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion
+
+# Represents the part of data iter covered by the buffer region
+RangeInfo = namedtuple(
+    "RangeInfo", ["start", "extent", "dim_in_data_iter", "dim_in_shape", "dim_type"]
+)
+
+
+def normalize_and_group(layout, shape):
+    """Normalize a layout with a given shape.
+
+    Parameters
+    ----------
+    layout : Union[T.TrainiumLayout, T.TileLayout]
+        The layout to normalize
+    shape : List[int]
+        The shape to normalize with
+
+    Returns
+    -------
+    Tuple[Union[T.TrainiumLayout, T.TileLayout], List[int]] :
+        Normalized layout and separators
+
+    Raises
+    ------
+    ValueError :
+        If layout is not a valid layout type
+    """
+    if isinstance(layout, T.TileLayout):
+        return layout.canonicalize().group(shape)
+    else:
+        raise ValueError("Invalid layout")
+
+
+def get_ewise_dim_map(
+    buffer_region: BufferRegion, second_buffer_region: BufferRegion, analyzer: Analyzer
+):
+    """Get the dimension map between two elementwise buffer regions.
+
+    Parameters
+    ----------
+    buffer_region : BufferRegion
+        The first buffer region
+    second_buffer_region : BufferRegion
+        The second buffer region
+    analyzer : Analyzer
+        The analyzer to use
+
+    Returns
+    -------
+    Dict[int, int] :
+        A dimension map from first to second buffer region
+
+    Raises
+    ------
+    AssertionError :
+        If dimensions do not match
+    """
+    extent_1 = [r.extent for r in buffer_region.region]
+    extent_2 = [r.extent for r in second_buffer_region.region]
+    extent_1_non_unit = [e for e in extent_1 if e != 1]
+    extent_2_non_unit = [e for e in extent_2 if e != 1]
+    assert all(
+        [
+            len(extent_1_non_unit) == len(extent_2_non_unit),
+            all(
+                analyzer.can_prove_equal(s, d) for s, d in zip(extent_1_non_unit, extent_2_non_unit)
+            ),
+        ]
+    )
+    dim_map = {}
+    i = 0
+    j = 0
+    while i < len(extent_1) and j < len(extent_2):
+        if analyzer.can_prove_equal(extent_1[i], 1):
+            i += 1
+            continue
+        if analyzer.can_prove_equal(extent_2[j], 1):
+            j += 1
+            continue
+        dim_map[i] = j
+        i += 1
+        j += 1
+    return dim_map
+
+
+def get_reduction_dim_map(
+    src_buffer_region: BufferRegion,
+    dst_buffer_region: BufferRegion,
+    axes: tuple[int],
+    analyzer: Analyzer,
+):
+    """Get the dimension map between source and destination buffer regions for reduction.
+
+    Parameters
+    ----------
+    src_buffer_region : BufferRegion
+        The source buffer region
+    dst_buffer_region : BufferRegion
+        The destination buffer region
+    axes : Tuple[int]
+        The reduction axes
+    analyzer : Analyzer
+        The analyzer to use
+
+    Returns
+    -------
+    Dict[int, int] :
+        A dimension map from source to destination buffer region
+
+    Raises
+    ------
+    AssertionError :
+        If dimensions do not match
+    """
+    dst_region = dst_buffer_region.region
+    dst_extent = [r.extent for r in dst_region]
+    dst_non_unit_extent_ = [(i, e) for i, e in enumerate(dst_extent) if e != 1]
+    src_region = src_buffer_region.region
+    src_extent = [r.extent for r in src_region]
+    src_non_unit_extent_ = [(i, e) for i, e in enumerate(src_extent) if e != 1]
+    src_non_reduction_extents = [(i, e) for i, e in src_non_unit_extent_ if i not in axes]
+    assert len(src_non_reduction_extents) == len(dst_non_unit_extent_), (
+        f"Source and destination must have the same number of non-reduction extents: {len(src_non_reduction_extents)} != {len(dst_non_unit_extent_)}"  # noqa: E501
+    )
+    for i in range(len(src_non_reduction_extents)):
+        assert analyzer.can_prove_equal(
+            src_non_reduction_extents[i][1], dst_non_unit_extent_[i][1]
+        ), (
+            f"Source and destination must have the same extent for non-reduction axes: {src_non_reduction_extents[i][1]} != {dst_non_unit_extent_[i][1]}"  # noqa: E501
+        )
+    dim_map = {s[0]: d[0] for s, d in zip(src_non_reduction_extents, dst_non_unit_extent_)}
+    return dim_map
+
+
+class DimensionMapper:
+    """
+    A class to manage dimension mappings between tensors.
+
+    A dimension mapping (dim_map) has type Dict[int, int]. dim_map[i] = j means
+    dimension i in the first tensor should be mapped to dimension j in the second tensor.
+    """
+
+    def __init__(self):
+        self.mappings = {}  # Dictionary to store mappings between tensors
+
+    def register_dim_map(self, first_tensor, second_tensor, dim_map):
+        """
+        Register a dimension mapping between two tensors.
+
+        Args:
+            first_tensor: The first tensor
+            second_tensor: The second tensor
+            dim_map: A dictionary mapping dimensions from first_tensor to second_tensor
+        """
+        # Initialize dictionaries if they don't exist
+        if first_tensor not in self.mappings:
+            self.mappings[first_tensor] = {}
+
+        # Register the mapping
+        self.mappings[first_tensor][second_tensor] = dim_map
+
+        # Register the reverse mapping
+        reverse_dim_map = {dim_map[i]: i for i in dim_map}
+
+        if second_tensor not in self.mappings:
+            self.mappings[second_tensor] = {}
+
+        self.mappings[second_tensor][first_tensor] = reverse_dim_map
+
+    def compose_mappings(self, map1, map2):
+        """
+        Compose two mappings: map1 followed by map2.
+
+        Args:
+            map1: The first mapping
+            map2: The second mapping
+
+        Returns:
+            A composition of the two mappings, or None if the composition is empty
+        """
+        result = {}
+        for i, j in map1.items():
+            if j in map2:
+                result[i] = map2[j]
+
+        # If the result is empty, return None
+        return result if result else None
+
+    def get_dim_map(self, first_tensor, second_tensor):
+        """
+        Get the dimension mapping between two tensors.
+
+        Args:
+            first_tensor: The first tensor
+            second_tensor: The second tensor
+
+        Returns:
+            A dictionary mapping dimensions from first_tensor to second_tensor,
+            or {} if no mapping exists
+        """
+        # Check if there is a direct mapping
+        if first_tensor in self.mappings and second_tensor in self.mappings[first_tensor]:
+            return self.mappings[first_tensor][second_tensor]
+
+        # No direct mapping, try to find a path using BFS
+        visited = {first_tensor}
+        queue = []
+
+        # Add all direct neighbors of the first tensor to the queue
+        if first_tensor in self.mappings:
+            for neighbor, direct_mapping in self.mappings[first_tensor].items():
+                visited.add(neighbor)
+                queue.append((neighbor, direct_mapping))
+
+        while queue:
+            current_tensor, mapping_from_first = queue.pop(0)
+
+            if current_tensor == second_tensor:
+                # Found a path to the second tensor
+                self.register_dim_map(first_tensor, second_tensor, mapping_from_first)
+                return mapping_from_first
+
+            if current_tensor not in self.mappings:
+                continue
+
+            for neighbor, direct_mapping in self.mappings[current_tensor].items():
+                if neighbor not in visited:
+                    visited.add(neighbor)
+
+                    # Compose the mappings: first_tensor -> current_tensor -> neighbor
+                    composed_mapping = self.compose_mappings(mapping_from_first, direct_mapping)
+
+                    # Only add to the queue if the composed mapping is not None
+                    if composed_mapping is not None:
+                        queue.append((neighbor, composed_mapping))
+
+        # No mapping found
+        return {}

python/tvm/backend/trn/operator/tile_primitive/gemm/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+from .default import *

python/tvm/backend/trn/operator/tile_primitive/gemm/default.py

@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Implementation of copy operator dispatchs."""
+
+import functools
+import operator
+
+from tvm.arith.analyzer import Analyzer
+from tvm.backend.trn.layout import is_trainium_layout
+from tvm.ir import assert_structural_equal
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, PrimFunc
+from tvm.tirx.operator.tile_primitive import (
+    DispatchContext,
+    fail,
+    predicate,
+    register_dispatch,
+)
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import init_analyzer
+from ..dim_utils import normalize_and_group
+from ..instruction_generator import InstructionGenerator
+from ..workspace_utils import check_workspace_buffer, largest_psum_per_bank, max_psum_banks
+
+
+class OperatorKind:
+    A = 0
+    B = 1
+    C = 2
+
+
+def get_pf_dim_from_buffer_region(
+    buffer_region: BufferRegion,
+    analyzer: Analyzer,
+    operator_kind: OperatorKind,
+    transposed: bool = False,
+):
+    """Extract partition and free dimensions from buffer region."""
+    # Find non-unit dimensions
+    non_unit_dims = [
+        i
+        for i in range(len(buffer_region.buffer.shape))
+        if not analyzer.can_prove_equal(buffer_region.region[i].extent, 1)
+    ]
+    assert len(non_unit_dims) == 2, "Only 2D matrix is supported for gemm"
+
+    layout, seps = normalize_and_group(buffer_region.buffer.layout, buffer_region.buffer.shape)
+    # Determine partition and free dimensions based on operator kind
+    if operator_kind == OperatorKind.A:
+        p_dim, f_dim = non_unit_dims[1], non_unit_dims[0]
+    elif operator_kind == OperatorKind.B:
+        p_dim, f_dim = non_unit_dims[0], non_unit_dims[1]
+    else:
+        assert not transposed, (
+            "Transposed C is implemented by swapping lhs and rhs. No need to specify by user."
+        )
+        # For C, determine dimensions based on layout
+        has_partition = any(
+            layout.shard[i].axis.name == "P"
+            for i in range(seps[non_unit_dims[0]], seps[non_unit_dims[0] + 1])
+        )
+        p_dim, f_dim = (
+            (non_unit_dims[0], non_unit_dims[1])
+            if has_partition
+            else (non_unit_dims[1], non_unit_dims[0])
+        )
+
+    # Swap dimensions if transposed
+    if transposed:
+        p_dim, f_dim = f_dim, p_dim
+
+    # Validate partition dimension
+    p_exts = [
+        layout.shard[i].extent
+        for i in range(seps[p_dim], seps[p_dim + 1])
+        if layout.shard[i].axis.name == "P"
+    ]
+
+    assert functools.reduce(operator.mul, p_exts, 1) == layout.size("P"), (
+        f"Accumulation dimension and output non-streaming dimension must contain whole P dimension. "  # noqa: E501
+        f"However, the {p_dim} dimension of {buffer_region} does not."
+    )
+
+    # Validate free dimension
+    assert all(
+        layout.shard[i].axis.name in ["F", "Bank"] or layout.shard[i].extent == 1
+        for i in range(seps[f_dim], seps[f_dim + 1])
+    ), (
+        f"Spatial dimension must not contain P. However, the {f_dim} dimension of {buffer_region} does."  # noqa: E501
+    )
+
+    return p_dim, f_dim
+
+
+def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    """Schedule GEMM operation on Trainium."""
+    # Basic validation checks
+    if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
+        fail("requires Trainium target and thread exec_scope")
+
+    # Extract arguments
+    (
+        D_buffer_region,
+        A_buffer_region,
+        B_buffer_region,
+        C_buffer_region,
+        transpose_A,
+        transpose_B,
+        alpha,
+        beta,
+    ) = op.args
+    analyzer = init_analyzer(sctx)
+    A, B, C, _D = (
+        A_buffer_region.buffer,
+        B_buffer_region.buffer,
+        C_buffer_region.buffer,
+        D_buffer_region.buffer,
+    )
+
+    # Validate alpha, beta
+    assert analyzer.can_prove_equal(alpha, 1) and analyzer.can_prove_equal(beta, 0), (
+        "Only alpha=1 and beta=0 are supported"
+    )
+
+    # D and C must be the same buffer region
+    assert_structural_equal(D_buffer_region, C_buffer_region)
+
+    # Validate buffer properties
+    assert all(
+        [
+            A.layout and B.layout and C.layout,
+            A.dtype == B.dtype,
+            A.scope() == "trn.sbuf" and B.scope() == "trn.sbuf",
+            C.scope() == "trn.psum" or C.scope() == "trn.sbuf",
+            is_trainium_layout(A.layout),
+            is_trainium_layout(B.layout),
+            is_trainium_layout(C.layout),
+            A.layout.size("P") == B.layout.size("P"),
+        ]
+    ), "Invalid buffer layout and scope"
+
+    p_size = A.layout.size("P")
+    assert p_size == B.layout.size("P"), "Partition size mismatch"
+
+    # Get partition and free dimensions
+    lhs_p_dim, lhs_f_dim = get_pf_dim_from_buffer_region(
+        A_buffer_region, analyzer, OperatorKind.A, transpose_A
+    )
+    rhs_p_dim, rhs_f_dim = get_pf_dim_from_buffer_region(
+        B_buffer_region, analyzer, OperatorKind.B, transpose_B
+    )
+    acc_p_dim, acc_f_dim = get_pf_dim_from_buffer_region(C_buffer_region, analyzer, OperatorKind.C)
+    # Swap LHS and RHS if needed based on accumulator dimensions
+    swap_lhs_rhs = acc_p_dim > acc_f_dim
+    if swap_lhs_rhs:
+        lhs_p_dim, rhs_p_dim = rhs_p_dim, lhs_p_dim
+        lhs_f_dim, rhs_f_dim = rhs_f_dim, lhs_f_dim
+        A, B = B, A
+        A_buffer_region, B_buffer_region = B_buffer_region, A_buffer_region
+
+    # Validate dimension compatibility
+    assert analyzer.can_prove(
+        A_buffer_region.region[lhs_p_dim].extent == B_buffer_region.region[rhs_p_dim].extent
+    ), (
+        f"Reduction dimension must match, but the {lhs_p_dim} dimension of {A_buffer_region} != the {rhs_p_dim} dimension of {B_buffer_region}"  # noqa: E501
+    )
+
+    assert analyzer.can_prove(
+        A_buffer_region.region[lhs_f_dim].extent == C_buffer_region.region[acc_p_dim].extent
+    ), (
+        f"Spatial dimension must match, but the {lhs_f_dim} dimension of {A_buffer_region} != the {acc_p_dim} dimension of {C_buffer_region}"  # noqa: E501
+    )
+
+    assert analyzer.can_prove(
+        B_buffer_region.region[rhs_f_dim].extent == C_buffer_region.region[acc_f_dim].extent
+    ), (
+        f"Spatial dimension must match, but the {rhs_f_dim} dimension of {B_buffer_region} != the {acc_f_dim} dimension of {C_buffer_region}"  # noqa: E501
+    )
+
+    inst_gen = InstructionGenerator([A_buffer_region, B_buffer_region, C_buffer_region], analyzer)
+    inst_gen.link_buffer_regions(A_buffer_region, B_buffer_region, {lhs_p_dim: rhs_p_dim})
+    inst_gen.link_buffer_regions(B_buffer_region, C_buffer_region, {rhs_f_dim: acc_f_dim})
+    inst_gen.link_buffer_regions(A_buffer_region, C_buffer_region, {lhs_f_dim: acc_p_dim})
+    inst_repr = inst_gen.find_max_inst_size_from_one_region(B_buffer_region, [rhs_f_dim])
+    inst_repr = inst_gen.fit_inst_tile_to_region(inst_repr, C_buffer_region, [acc_f_dim])
+    inst_repr.bound_inst_size(512, analyzer)
+    rhs_f = T.Var("rhs_f", "int32")
+    lhs_f = T.Var("lhs_f", "int32")
+    p = T.Var("p", "int32")
+    reduction_b = T.Var("reduction_b", "int32")
+    lhs_b = T.Var("lhs_b", "int32")
+    rhs_b = T.Var("rhs_b", "int32")
+    lhs_f_size = C.layout.size("P")
+    inst_gen.bind_inst_iter(
+        B_buffer_region, rhs_f, inst_repr.size, inst_repr.stride, is_free_dim=True
+    )
+    inst_gen.bind_inst_iter(C_buffer_region, lhs_f, lhs_f_size, 1, is_free_dim=False)
+    inst_gen.bind_inst_iter(A_buffer_region, p, A.layout.size("P"), 1, is_free_dim=False)
+    reduction_b_extent = inst_gen.fill_in_block_dim(A_buffer_region, reduction_b, [lhs_p_dim])
+    lhs_b_extent = inst_gen.fill_in_block_dim(A_buffer_region, lhs_b, [lhs_f_dim])
+    rhs_b_extent = inst_gen.fill_in_block_dim(B_buffer_region, rhs_b, [rhs_f_dim])
+
+    # FIXME: we need to lower the guard to things like matmul(lhs[...][lhs_guard], rhs[...][rhs_guard], mask=p_guard)  # noqa: E501
+    # so we need to separate the guard for lhs_f, rhs_f and p
+    # fmt: off
+    @T.inline
+    def matmul_inst_macro(lhs_b_loop, rhs_b_loop, reduction_b_loop, acc, C_as_output, max_psum_slots):  # noqa: E501
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(0, p_size, annotations={"nki_dim": "P"}):
+                for lhs_f_loop in T.serial(0, lhs_f_size, annotations={"nki_dim": "lhs_F"}):
+                    for rhs_f_loop in T.serial(0, inst_repr.size, annotations={"nki_dim": "rhs_F"}):
+                        b_idx = T.meta_var(lhs_b_loop * rhs_b_extent + rhs_b_loop)
+                        inst_gen.set_bind_map(A_buffer_region, {lhs_b: lhs_b_loop, lhs_f: lhs_f_loop, p: p_loop, reduction_b: reduction_b_loop})  # noqa: E501
+                        inst_gen.set_bind_map(B_buffer_region, {rhs_b: rhs_b_loop, rhs_f: rhs_f_loop, p: p_loop, reduction_b: reduction_b_loop})  # noqa: E501
+                        inst_gen.set_bind_map(C_buffer_region, {lhs_f: lhs_f_loop, rhs_f: rhs_f_loop, lhs_b: lhs_b_loop, rhs_b: rhs_b_loop})  # noqa: E501
+                        lhs_indices = T.meta_var(inst_gen.generate_indices(A_buffer_region))
+                        rhs_indices = T.meta_var(inst_gen.generate_indices(B_buffer_region))
+                        C_indices = T.meta_var(inst_gen.generate_indices(C_buffer_region))
+                        if inst_gen.make_guard(A_buffer_region) and inst_gen.make_guard(B_buffer_region):  # noqa: E501
+                            if C_as_output:
+                                T.evaluate(T.nki.matmul(acc[C_indices], A[lhs_indices], B[rhs_indices]))  # noqa: E501
+                            else:
+                                T.evaluate(T.nki.matmul(acc[b_idx % max_psum_slots, lhs_f_loop, rhs_f_loop], A[lhs_indices], B[rhs_indices]))  # noqa: E501
+
+    if C.scope() == "trn.psum":
+        @T.prim_func
+        def impl_C_psum():
+            for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(lhs_b_extent, rhs_b_extent, reduction_b_extent):  # noqa: E501
+                matmul_inst_macro(lhs_b_loop, rhs_b_loop, reduction_b_loop, C, True, None)
+        return impl_C_psum
+
+    # todo: generalize the process of generating composite matmul + another_op pattern
+    # by generating TIR op and reusing existing dispatch rule
+
+    # we will support matmul + epilogue as a user-specified pattern
+    # and a matmul fusion pass can help infer the pattern
+
+    acc_psum_shape = (max_psum_banks, p_size, largest_psum_per_bank)
+    if "acc_psum" not in op.workspace:
+        assert sctx.alloc_only, "Accumulation psum buffer must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first."  # noqa: E501
+        acc_psum = T.buffer(
+                acc_psum_shape,
+                "float32",
+                scope="trn.psum",
+                allocated_addr=(0, 0),
+                buffer_name="acc_psum"
+            )
+        sctx.add_alloc_buffer(acc_psum)
+        max_psum_slots = max_psum_banks
+    else:
+        acc_psum = op.workspace["acc_psum"]
+        check_workspace_buffer(acc_psum, (p_size, largest_psum_per_bank), "trn.psum")
+        max_psum_slots = acc_psum.shape[0]
+
+    @T.prim_func
+    def impl_C_sbuf():
+        for lhs_b_loop, rhs_b_loop in T.grid(lhs_b_extent, rhs_b_extent):
+            for reduction_b_loop in T.serial(0, reduction_b_extent):
+                matmul_inst_macro(lhs_b_loop, rhs_b_loop, reduction_b_loop, acc_psum, False, max_psum_slots)  # noqa: E501
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for lhs_f_loop in T.serial(0, lhs_f_size, annotations={"nki_dim": "P"}):
+                    for rhs_f_loop in T.serial(0, inst_repr.size, annotations={"nki_dim": "F"}):
+                        b_idx = T.meta_var(lhs_b_loop * rhs_b_extent + rhs_b_loop)
+                        inst_gen.set_bind_map(C_buffer_region, {lhs_f: lhs_f_loop, rhs_f: rhs_f_loop, lhs_b: lhs_b_loop, rhs_b: rhs_b_loop})  # noqa: E501
+                        if inst_gen.make_guard(C_buffer_region):
+                            acc_indices = T.meta_var(inst_gen.generate_indices(C_buffer_region))
+                            T.evaluate(T.nki.tensor_copy(C[acc_indices], acc_psum[b_idx % max_psum_slots, lhs_f_loop, rhs_f_loop]))  # noqa: E501
+    # fmt: on
+    return impl_C_sbuf
+
+
+# Rich dispatcher variant for TRN gemm
+@register_dispatch(
+    "gemm",
+    "trn",
+    variant="default",
+    priority=10,
+    when=[
+        predicate(
+            "exec_scope",
+            lambda op, sctx: (
+                sctx.scope_kind == "thread",
+                f"unsupported exec_scope {sctx.scope_kind}",
+            ),
+        )
+    ],
+)
+def gemm_trn_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return matmul_trn(op, sctx)

python/tvm/backend/trn/operator/tile_primitive/instruction_generator.py

@@ -0,0 +1,732 @@
+# 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.
+
+"""Instruction generation utilities for TRN operator scheduling."""
+
+import itertools
+from dataclasses import dataclass
+from functools import reduce
+from math import gcd
+from operator import mul
+
+import tvm
+from tvm.arith.analyzer import Analyzer
+from tvm.backend.trn.layout import is_trainium_layout
+from tvm.ir import Range
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, Expr, Var
+from tvm.tirx.expr_functor import ExprMutator
+from tvm.tirx.layout import Iter
+
+from .dim_utils import DimensionMapper, RangeInfo, normalize_and_group
+
+
+@dataclass
+class LogicalIterDim:
+    logical_stride: int
+    extent: int
+    bind_expr: Expr
+
+    @staticmethod
+    def default():
+        return LogicalIterDim(1, 1, T.int32(0))
+
+
+LogicalIterList = tuple[tuple[tuple[LogicalIterDim]]]
+
+
+def to_int_list(intimm_list: list[T.IntImm]):
+    return [int(i) for i in intimm_list]
+
+
+class VarReplacer(ExprMutator):
+    def __init__(self, var_map: dict[Var, Expr]):
+        super().__init__()
+        self.var_map = var_map
+
+    def visit_var_(self, op):
+        if op in self.var_map:
+            return self.var_map[op]
+        return op
+
+    @staticmethod
+    def replace_vars(expr: Expr, var_map: dict[Var, Expr]) -> Expr:
+        return VarReplacer(var_map).visit_expr(expr)
+
+
+@dataclass
+class InstructionRepr:
+    buffer_region: BufferRegion
+    size: int
+    stride: int
+    selected_data_iter_ids: list[int]
+
+    def __init__(
+        self,
+        buffer_region: BufferRegion,
+        inst_size: int,
+        inst_stride: int,
+        selected_data_iter_ids: list[int],
+    ):
+        self.buffer_region = buffer_region
+        self.size = inst_size if inst_size is not None else 1
+        self.stride = inst_stride if inst_stride is not None else 1
+        self.selected_data_iter_ids = selected_data_iter_ids
+
+    def bound_inst_size(self, max_inst_size: int | None, analyzer: Analyzer):
+        if max_inst_size is None:
+            return
+        if analyzer.can_prove(self.size <= max_inst_size):
+            return
+        assert analyzer.can_prove(self.size % max_inst_size == 0), (
+            f"The instruction size {self.size} is not a multiple of the max instruction size {max_inst_size}"  # noqa: E501
+        )
+        self.size = max_inst_size
+        self.selected_data_iter_ids = None
+
+
+class InstructionGenerator:
+    def __init__(self, buffer_regions: tuple[BufferRegion], analyzer: Analyzer):
+        self.buffer_regions = []
+        self.analyzer = analyzer
+        self.split_shape_views = {}
+        self.split_layout_views = {}
+        self.seps = {}
+        self.bound_regions = {}
+        self.bind_iters: dict[BufferRegion, LogicalIterList] = None
+        self.bind_maps: dict[BufferRegion, dict[Var, Expr]] = {}
+        for buffer_region in buffer_regions:
+            if not isinstance(buffer_region, BufferRegion):
+                continue
+            self.buffer_regions.append(buffer_region)
+            bound_buffer_region = self._bound_buffer_region(buffer_region)
+            layout, seps = self._get_sub_layout(bound_buffer_region)
+            self.split_shape_views[buffer_region] = self._get_flattened_shape_view_from_layout_seps(
+                layout, seps
+            )
+            self.split_layout_views[buffer_region] = layout
+            self.seps[buffer_region] = seps
+        self.dim_mapper = DimensionMapper()
+
+    def _bound_buffer_region(self, buffer_region: BufferRegion):
+        region = []
+        changed = False
+        for r in buffer_region.region:
+            bound = self.analyzer.const_int_bound(r.extent)
+            if not self.analyzer.can_prove_equal(bound.max_value, r.extent):
+                changed = True
+            region.append(Range.from_min_extent(r.min, bound.max_value))
+        if changed:
+            bound_region = BufferRegion(buffer_region.buffer, region)
+            self.bound_regions[buffer_region] = bound_region
+            return bound_region
+        return buffer_region
+
+    def _get_sub_layout(self, buffer_region: BufferRegion):
+        layout = buffer_region.buffer.layout
+        layout, seps = normalize_and_group(layout, buffer_region.buffer.shape)
+        tiled_range_infos_per_dim = []
+        new_shard = []
+        new_seps = [0]
+        for i in range(len(seps) - 1):
+            r = buffer_region.region[i]
+            st = r.min
+            ext = r.extent
+            reversed_shard = []
+            for j in reversed(range(seps[i], seps[i + 1])):
+                if self.analyzer.can_prove_equal(ext, 1):
+                    break
+                if layout.shard[j].axis.name == "P" and (
+                    not self.analyzer.can_prove(st % layout.shard[j].extent == 0)
+                    or not self.analyzer.can_prove(ext % layout.shard[j].extent == 0)
+                ):
+                    assert False, "Invalid layout"
+                if self.analyzer.can_prove(
+                    ext % layout.shard[j].extent == 0
+                ) and self.analyzer.can_prove(st % layout.shard[j].extent == 0):
+                    st = st // layout.shard[j].extent
+                    ext = ext // layout.shard[j].extent
+                    tiled_range_infos_per_dim.append(
+                        RangeInfo(0, layout.shard[j].extent, j, i, layout.shard[j].axis)
+                    )
+                    reversed_shard.append(layout.shard[j])
+                    continue
+                if self.analyzer.can_prove(st + ext <= layout.shard[j].extent):
+                    tiled_range_infos_per_dim.append(RangeInfo(st, ext, j, i, layout.shard[j].axis))
+                    reversed_shard.append(Iter(ext, layout.shard[j].stride, layout.shard[j].axis))
+                    break
+                assert False, f"Cannot analyze physical tensor region for: {buffer_region}"
+            new_shard += reversed(reversed_shard)
+            new_seps.append(len(reversed_shard) + new_seps[-1])
+        new_tile_layout = tvm.tirx.layout.TileLayout.from_iters(  # pylint: disable=no-member
+            new_shard, [], dict()
+        )
+        return new_tile_layout, new_seps
+
+    def _init_bind_iters(self):
+        self.bind_iters = {}
+        for buffer_region in self.buffer_regions:
+            seps = self.seps[buffer_region]
+            self.bind_iters[buffer_region] = [
+                [[] for _ in range(seps[i], seps[i + 1])] for i in range(len(buffer_region.region))
+            ]
+
+    def _normalize_bind_iters(self):
+        for buffer_region in self.buffer_regions:
+            seps = self.seps[buffer_region]
+            self.bind_iters[buffer_region] = [
+                [
+                    sorted(
+                        self.bind_iters[buffer_region][i][j - seps[i]],
+                        key=lambda x: (x.logical_stride, x.extent),
+                    )
+                    for j in range(seps[i], seps[i + 1])
+                ]
+                for i in range(len(buffer_region.region))
+            ]
+
+    def _get_flattened_shape_view_from_layout_seps(self, layout, seps):
+        return [
+            [layout.shard[j].extent for j in range(seps[i], seps[i + 1])]
+            for i in range(len(seps) - 1)
+        ]
+
+    def common_factor(self, shape_a, shape_b):
+        """
+        Return the finest common factor shape of two compatible shapes.
+
+        A "common factor" shape `C` satisfies:
+        1. ∏shape_a == ∏shape_b == ∏C   (same total #elements)
+        2. `C` can be obtained from `shape_a` **only** by splitting (never merging)
+            dimensions, and likewise for `shape_b`.
+
+        Parameters
+        ----------
+        shape_a, shape_b : tuple[int] | list[int]
+            Two equally-sized shapes.
+
+        Returns
+        -------
+        tuple[int]
+            The common-factor shape.
+
+        Raises
+        ------
+        AssertionError
+            - if the shapes have different element counts
+            - or if a common-factor decomposition does not exist
+            (which only happens if the two shapes do not share a
+            compatible prime-factor ordering).
+        """
+        if len(shape_a) == 0 and len(shape_b) == 0:
+            return shape_a
+        shape_a = to_int_list(shape_a)
+        shape_b = to_int_list(shape_b)
+        # 1. identical element count
+        size_a = reduce(mul, shape_a, 1)
+        size_b = reduce(mul, shape_b, 1)
+        assert size_a == size_b, "Shapes hold different numbers of elements"
+
+        i, j = 0, 0
+        rem_a, rem_b = shape_a[0], shape_b[0]
+        out = []
+
+        while i < len(shape_a) and j < len(shape_b):
+            g = gcd(rem_a, rem_b)
+            assert g > 1 or (rem_a == rem_b == 1), "Incompatible factor ordering"
+            out.append(g)
+
+            # consume g from the current "head" factors
+            rem_a //= g
+            rem_b //= g
+
+            # advance whenever a remainder has been completely consumed
+            if rem_a == 1:
+                i += 1
+                rem_a = shape_a[i] if i < len(shape_a) else 1
+            if rem_b == 1:
+                j += 1
+                rem_b = shape_b[j] if j < len(shape_b) else 1
+
+        # sanity check
+        assert i == len(shape_a) and j == len(shape_b), "Did not exhaust both shapes"
+
+        return tuple(out)
+
+    def _link_buffer_regions(
+        self, buffer_region: BufferRegion, to_link: BufferRegion, dim_map: dict[int, int]
+    ):
+        split_shape_view_1 = self.split_shape_views[buffer_region]
+        split_layout_view_1 = self.split_layout_views[buffer_region]
+        split_shape_view_2 = self.split_shape_views[to_link]
+
+        # adapt to the shape view of the to_link buffer region
+        new_split_shape_view_1 = [
+            (
+                self.common_factor(split_shape_view_2[dim_map[i]], split_shape_view_1[i])
+                if i in dim_map
+                else split_shape_view_1[i]
+            )
+            for i in range(len(buffer_region.region))
+        ]
+        flattened_shape_view_1 = list(itertools.chain(*new_split_shape_view_1))
+        layout, tiled_seps = normalize_and_group(split_layout_view_1, flattened_shape_view_1)
+        actual_seps = [0]
+        ptr = 0
+        for i in range(len(buffer_region.region)):
+            ptr += len(new_split_shape_view_1[i])
+            actual_seps.append(tiled_seps[ptr])
+        self.split_shape_views[buffer_region] = self._get_flattened_shape_view_from_layout_seps(
+            layout, actual_seps
+        )
+        self.split_layout_views[buffer_region] = layout
+        self.seps[buffer_region] = actual_seps
+
+    def _get_reverse_dim_map(self, dim_map: dict[int, int]) -> dict[int, int]:
+        return {dim_map[i]: i for i in dim_map}
+
+    def link_buffer_regions(
+        self, buffer_region: BufferRegion, to_link: BufferRegion, dim_map: dict[int, int]
+    ):
+        self.dim_mapper.register_dim_map(buffer_region, to_link, dim_map)
+        for r in self.buffer_regions:
+            if r == to_link:
+                continue
+            dim_map = self.dim_mapper.get_dim_map(r, to_link)
+            reverse_dim_map = self._get_reverse_dim_map(dim_map)
+            self._link_buffer_regions(r, to_link, dim_map)
+            self._link_buffer_regions(to_link, r, reverse_dim_map)
+            seps_1 = self.seps[r]
+            seps_2 = self.seps[to_link]
+            for i, j in dim_map.items():
+                assert seps_1[i + 1] - seps_1[i] == seps_2[j + 1] - seps_2[j], (
+                    f"The number of data iters at dim {i} of {buffer_region.buffer.name} is not equal to the number of data iters at dim {j} of {to_link.buffer.name}"  # noqa: E501
+                )
+
+    def bind_inst_iter(
+        self,
+        buffer_region: BufferRegion,
+        bind: Var,
+        inst_size: int,
+        inst_stride: int,
+        is_free_dim: bool,
+        no_propagate: bool = False,
+    ):
+        logical_iter_list = self._get_inst_logical_iter_list(
+            buffer_region, bind, inst_stride, inst_size, is_free_dim
+        )
+        self._add_bind_iter_list(buffer_region, logical_iter_list)
+        if no_propagate:
+            return
+        self._propagate_bind_iter(buffer_region, logical_iter_list)
+
+    def _propagate_bind_iter(self, buffer_region: BufferRegion, logical_iter_list: LogicalIterList):
+        for to_propagate in self.buffer_regions:
+            if to_propagate == buffer_region:
+                continue
+            dim_map = self.dim_mapper.get_dim_map(buffer_region, to_propagate)
+            reverse_dim_map = self._get_reverse_dim_map(dim_map)
+            seps = self.seps[to_propagate]
+            propagated_logical_iter = [
+                (
+                    logical_iter_list[reverse_dim_map[i]]
+                    if i in reverse_dim_map
+                    else [[] for _ in range(seps[i], seps[i + 1])]
+                )
+                for i in range(len(to_propagate.region))
+            ]
+            self._add_bind_iter_list(to_propagate, propagated_logical_iter)
+
+    def _add_bind_iter_list(self, buffer_region: BufferRegion, bind_iter_list: LogicalIterList):
+        if self.bind_iters is None:
+            self._init_bind_iters()
+        seps = self.seps[buffer_region]
+        for i in range(len(buffer_region.region)):
+            for j in range(seps[i], seps[i + 1]):
+                self.bind_iters[buffer_region][i][j - seps[i]].extend(
+                    bind_iter_list[i][j - seps[i]]
+                )
+
+    def fill_in_block_dim(
+        self, buffer_region: BufferRegion, bind: Var, dims: list[int] | None = None
+    ):
+        # fixme: be cautious of the min of buffer region. This implementation is not correct.
+        #        we need to first take a view of sub-layout (keep strides, but reduce the extent
+        #        then we analyze the relationship between data iter of sub-layout
+        dims = dims or list(range(len(buffer_region.buffer.shape)))
+        layout = self.split_layout_views[buffer_region]
+        shards = layout.shard
+        self._normalize_bind_iters()
+        bind_iters = self.bind_iters[buffer_region]
+        seps = self.seps[buffer_region]
+        logical_iter_list_block = [
+            [[] for _ in range(seps[i], seps[i + 1])] for i in range(len(buffer_region.region))
+        ]
+        acc_block_ext = 1
+        for i in reversed(dims):
+            for j in reversed(range(seps[i], seps[i + 1])):
+                it = shards[j]
+                is_partition = it.axis.name == "P" if is_trainium_layout(layout) else False
+                logical_iter_dims = bind_iters[i][j - seps[i]]
+                for d in range(-1, len(logical_iter_dims)):
+                    next_logical_stride = (
+                        logical_iter_dims[d + 1].logical_stride
+                        if d + 1 < len(logical_iter_dims)
+                        else it.extent
+                    )
+                    cur = (
+                        logical_iter_dims[d].logical_stride * logical_iter_dims[d].extent
+                        if d >= 0
+                        else 1
+                    )
+                    assert next_logical_stride % cur == 0, (
+                        f"Fail to infer block dim for {buffer_region.buffer.name} at dim {i}"
+                    )
+                    gap = next_logical_stride // cur
+                    if is_partition:
+                        assert gap == 1, (
+                            f"Fail to propagate partition dim. The propagated dim does not cover the whole partition on {buffer_region.buffer.name} at dim {i}"  # noqa: E501
+                        )
+                    elif gap > 1:
+                        new_acc_block_ext = acc_block_ext * gap
+                        logical_iter_list_block[i][j - seps[i]].append(
+                            LogicalIterDim(cur, gap, bind % new_acc_block_ext // acc_block_ext)
+                        )
+                        acc_block_ext = new_acc_block_ext
+        self._add_bind_iter_list(buffer_region, logical_iter_list_block)
+        self._propagate_bind_iter(buffer_region, logical_iter_list_block)
+        return acc_block_ext
+
+    def _check_bind_iter_coverage(self, buffer_region: BufferRegion):
+        self._normalize_bind_iters()
+        seps = self.seps[buffer_region]
+        iters = self.split_layout_views[buffer_region].shard
+        bind_iters = self.bind_iters[buffer_region]
+        for i in range(len(buffer_region.region)):
+            for j in range(seps[i], seps[i + 1]):
+                it = iters[j]
+                logical_iter_dims = bind_iters[i][j - seps[i]]
+                for d in range(len(logical_iter_dims)):
+                    next_logical_stride = (
+                        logical_iter_dims[d + 1].logical_stride
+                        if d + 1 < len(logical_iter_dims)
+                        else it.extent
+                    )
+                    assert (
+                        next_logical_stride
+                        % (logical_iter_dims[d].logical_stride * logical_iter_dims[d].extent)
+                        == 0
+                    ), f"Fail to infer block dim for {buffer_region.buffer.name} at dim {i}"
+                    gap = next_logical_stride // (
+                        logical_iter_dims[d].logical_stride * logical_iter_dims[d].extent
+                    )
+                    assert gap == 1, "Call fill_in_block_dim() before calling generate_indices()"
+
+    def set_bind_map(self, buffer_region: BufferRegion, bind_map: dict[Var, Expr]):
+        self.bind_maps[buffer_region] = bind_map
+
+    def set_bind_map_all(self, bind_map: dict[Var, Expr]):
+        for buffer_region in self.buffer_regions:
+            self.set_bind_map(buffer_region, bind_map)
+
+    def generate_axes(self, buffer_region: BufferRegion) -> list[Expr]:
+        self._check_bind_iter_coverage(buffer_region)
+        layout = self.split_layout_views[buffer_region]
+        iters = layout.shard
+        bind_iters = self.bind_iters[buffer_region]
+        seps = self.seps[buffer_region]
+        axes = []
+        for i in range(len(bind_iters)):
+            index = 0
+            acc_logical_stride = 1
+            for j in reversed(range(seps[i], seps[i + 1])):
+                logical_iter_dims = bind_iters[i][j - seps[i]]
+                for d in reversed(logical_iter_dims):
+                    if d.extent == 1:
+                        continue
+                    index += (
+                        d.logical_stride
+                        * VarReplacer.replace_vars(d.bind_expr, self.bind_maps[buffer_region])
+                        * acc_logical_stride
+                    )
+                acc_logical_stride *= iters[j].extent
+            axes.append(index)
+        return axes
+
+    def generate_indices(self, buffer_region: BufferRegion) -> list[Expr]:
+        axes = self.generate_axes(buffer_region)
+        return [axes[i] + r.min for i, r in enumerate(buffer_region.region)]
+
+    def _get_inst_logical_iter_list(
+        self,
+        buffer_region: BufferRegion,
+        bind: Var,
+        stride: int,
+        size: int,
+        is_free_dim: bool = True,
+    ) -> LogicalIterList:
+        layout = self.split_layout_views[buffer_region]
+        assert is_trainium_layout(layout), (
+            " Cannot propagate instruction information from HBM tensor"
+        )
+        iters = layout.shard
+        seps = self.seps[buffer_region]
+        ret = [[[] for _ in range(seps[i], seps[i + 1])] for i in range(len(buffer_region.region))]
+        for i in range(len(buffer_region.region)):
+            for j in range(seps[i], seps[i + 1]):
+                if (iters[j].axis.name in ["F", "Bank"]) ^ is_free_dim:
+                    continue
+                it = iters[j]
+                if it.stride * it.extent <= stride or it.stride >= size * stride:
+                    continue
+                if it.stride * it.extent < size * stride and stride <= it.stride:
+                    assert (size * stride) % (
+                        it.stride * it.extent
+                    ) == 0 and it.stride % stride == 0
+                    ret[i][j - seps[i]].append(
+                        LogicalIterDim(
+                            1,
+                            it.extent,
+                            bind % (it.stride * it.extent // stride) // (it.stride // stride),
+                        )
+                    )
+                elif it.stride * it.extent < size * stride and stride > it.stride:
+                    assert (size * stride) % (
+                        it.stride * it.extent
+                    ) == 0 and stride % it.stride == 0
+                    ret[i][j - seps[i]].append(
+                        LogicalIterDim(
+                            stride // it.stride,
+                            it.stride * it.extent // stride,
+                            bind % (it.stride * it.extent // stride),
+                        )
+                    )
+                elif it.stride * it.extent >= size * stride and stride <= it.stride:
+                    assert (it.stride * it.extent) % (
+                        size * stride
+                    ) == 0 and it.stride % stride == 0
+                    ret[i][j - seps[i]].append(
+                        LogicalIterDim(1, size * stride // it.stride, bind // (it.stride // stride))
+                    )
+        return ret
+
+    def make_guard(self, buffer_region: BufferRegion):
+        if buffer_region not in self.bound_regions:
+            return True
+        bound_region = self.bound_regions[buffer_region]
+        relaxed_dims = [
+            i
+            for i, (r1, r2) in enumerate(zip(bound_region.region, buffer_region.region))
+            if not self.analyzer.can_prove(r1.extent == r2.extent)
+        ]
+        axes = self.generate_axes(buffer_region)
+        guard = reduce(
+            T.And,
+            [axes[i] < r.extent for i, r in enumerate(buffer_region.region) if i in relaxed_dims],
+            True,
+        )
+        return guard
+
+    def _find_max_linear_inst(self, indexed_data_iters, min_stride: int | None = None):
+        min_stride = min_stride or 1
+        indexed_data_iters = sorted(indexed_data_iters, key=lambda x: x[1].stride)
+        inst_size = 1
+        inst_stride = None
+        idx_list = []
+        for idx, data_iter in indexed_data_iters:
+            if data_iter.extent == 1 or data_iter.stride * data_iter.extent < min_stride:
+                continue
+            assert data_iter.stride % min_stride == 0 or min_stride % data_iter.stride == 0, (
+                f"Invalid instruction stride {min_stride}"
+            )
+            if inst_stride is not None and inst_stride * inst_size != data_iter.stride:
+                # the stride of the found data iter is not compatible with previous data iters
+                break
+            elif inst_stride is None:
+                inst_stride = max(min_stride, data_iter.stride)
+            if min_stride % data_iter.stride == 0:
+                inst_size = data_iter.extent * data_iter.stride // inst_stride
+            else:
+                inst_size *= data_iter.extent
+            idx_list.append(idx)
+        return inst_size, inst_stride, idx_list
+
+    def find_max_inst_size_from_one_region(
+        self,
+        buffer_region: BufferRegion,
+        allowed_f_dim: tuple[int] | None = None,
+        min_stride: int | None = None,
+    ):
+        allowed_f_dim = allowed_f_dim or tuple(range(len(buffer_region.region)))
+        layout = self.split_layout_views[buffer_region]
+        seps = self.seps[buffer_region]
+        allowed_data_iter_idx = itertools.chain.from_iterable(
+            range(seps[dim], seps[dim + 1]) for dim in allowed_f_dim
+        )
+        filtered_data_iters = [
+            (i, layout.shard[i])
+            for i in allowed_data_iter_idx
+            if layout.shard[i].axis.name in ["F", "Bank"]
+        ]
+        inst_size, inst_stride, idx_list = self._find_max_linear_inst(
+            filtered_data_iters, min_stride
+        )
+        return InstructionRepr(buffer_region, inst_size, inst_stride, idx_list)
+
+    def fit_inst_tile_to_region(
+        self,
+        inst_repr: InstructionRepr,
+        to_region: BufferRegion,
+        allowed_to_f_dim: tuple[int] | None = None,
+        broadcast: bool = False,
+    ):
+        allowed_to_f_dim = allowed_to_f_dim or tuple(range(len(to_region.region)))
+        from_region = inst_repr.buffer_region
+        from_layout = self.split_layout_views[from_region]
+        to_layout = self.split_layout_views[to_region]
+        from_seps = self.seps[from_region]
+        to_seps = self.seps[to_region]
+        dim_map = self.dim_mapper.get_dim_map(from_region, to_region)
+        dim_map = {i: j for i, j in dim_map.items() if j in allowed_to_f_dim}
+        data_iter_map = {
+            from_seps[i] + idx: to_seps[j] + idx
+            for i, j in dim_map.items()
+            for idx in range(from_seps[i + 1] - from_seps[i])
+        }
+        if broadcast:
+            data_iter_idx_to_dim = {
+                from_seps[i] + j: i
+                for i in range(len(from_region.region))
+                for j in range(from_seps[i + 1] - from_seps[i])
+            }
+            indexed_selected_shard = [
+                (i, from_layout.shard[i])
+                for i in inst_repr.selected_data_iter_ids
+                if data_iter_idx_to_dim[i] not in dim_map
+            ]
+            inst_size, inst_stride, idx_list = self._find_max_linear_inst(indexed_selected_shard)
+            return InstructionRepr(from_region, inst_size, inst_stride, idx_list)
+        indexed_selected_shard = [
+            (i, from_layout.shard[i]) for i in inst_repr.selected_data_iter_ids
+        ]
+        indexed_selected_shard = sorted(indexed_selected_shard, key=lambda x: x[1].stride)
+        inst_size = 1
+        inst_stride_from = None
+        inst_stride_to = None
+        idx_list = []
+        for i, data_iter in indexed_selected_shard:
+            if i not in data_iter_map:
+                if inst_stride_from is None:
+                    continue
+                break
+            mapped_data_iter = to_layout.shard[data_iter_map[i]]
+            if inst_stride_from is None:
+                inst_stride_from = data_iter.stride
+                if not is_trainium_layout(to_layout) and mapped_data_iter.stride != 1:
+                    # dma copy must be contiguous on hbm
+                    break
+                inst_stride_to = mapped_data_iter.stride
+            elif inst_stride_to * inst_size != mapped_data_iter.stride:
+                break
+            inst_size *= data_iter.extent
+            idx_list.append(i)
+        return InstructionRepr(from_region, inst_size, inst_stride_from, idx_list)
+
+    def check_partition_dim_match(
+        self, buffer_region_1: BufferRegion, buffer_region_2: BufferRegion
+    ):
+        dim_map = self.dim_mapper.get_dim_map(buffer_region_1, buffer_region_2)
+        layout_1 = self.split_layout_views[buffer_region_1]
+        layout_2 = self.split_layout_views[buffer_region_2]
+        if not is_trainium_layout(layout_1) or not is_trainium_layout(layout_2):
+            return True
+        seps_1 = self.seps[buffer_region_1]
+        seps_2 = self.seps[buffer_region_2]
+        for i, j in dim_map.items():
+            for k in range(seps_1[i + 1] - seps_1[i]):
+                if (
+                    layout_1.shard[seps_1[i] + k].axis.name
+                    != layout_2.shard[seps_2[j] + k].axis.name
+                ):
+                    return False
+                if layout_1.shard[seps_1[i] + k].axis.name in ["F", "Bank"]:
+                    continue
+                if layout_1.shard[seps_1[i] + k].stride != layout_2.shard[seps_2[j] + k].stride:
+                    return False
+                if layout_1.shard[seps_1[i] + k].extent != layout_2.shard[seps_2[j] + k].extent:
+                    return False
+        return True
+
+    def find_max_inst_size_transpose(
+        self, buffer_region_1: BufferRegion, buffer_region_2: BufferRegion
+    ):
+        dim_map = self.dim_mapper.get_dim_map(buffer_region_1, buffer_region_2)
+        layout_1 = self.split_layout_views[buffer_region_1]
+        layout_2 = self.split_layout_views[buffer_region_2]
+        iters_1 = layout_1.shard
+        iters_2 = layout_2.shard
+        seps_1 = self.seps[buffer_region_1]
+        seps_2 = self.seps[buffer_region_2]
+        indexed_iters_1 = []
+        indexed_iters_2 = []
+        print(iters_1, seps_1)
+        print(iters_2, seps_2)
+        print(dim_map)
+        for i, j in dim_map.items():
+            for k in range(seps_1[i + 1] - seps_1[i]):
+                if iters_1[seps_1[i] + k].axis.name == iters_2[seps_2[j] + k].axis.name:
+                    if iters_1[seps_1[i] + k].axis.name in ["F", "Bank"]:
+                        continue
+                    raise ValueError("Transpose only part of P dimension is not supported")
+                if iters_1[seps_1[i] + k].axis.name == "P":
+                    indexed_iters_2.append((seps_2[j] + k, iters_2[seps_2[j] + k]))
+                else:
+                    indexed_iters_1.append((seps_1[i] + k, iters_1[seps_1[i] + k]))
+        inst_repr_1 = InstructionRepr(buffer_region_1, *self._find_max_linear_inst(indexed_iters_1))
+        inst_repr_2 = InstructionRepr(buffer_region_2, *self._find_max_linear_inst(indexed_iters_2))
+        assert inst_repr_1.size == layout_2.size("P"), (
+            f"The instruction size of {buffer_region_1.buffer.name} does not match the partition size of {buffer_region_2.buffer.name}"  # noqa: E501
+        )
+        assert inst_repr_2.size == layout_1.size("P"), (
+            f"The instruction size of {buffer_region_2.buffer.name} does not match the partition size of {buffer_region_1.buffer.name}"  # noqa: E501
+        )
+        return inst_repr_1, inst_repr_2
+
+    def restrict_inst_to_one_dim(self, inst_repr: InstructionRepr):
+        region = inst_repr.buffer_region
+        layout = self.split_layout_views[region]
+        iters = layout.shard
+        seps = self.seps[region]
+        indexed_selected_iters = [(i, iters[i]) for i in inst_repr.selected_data_iter_ids]
+        indexed_selected_iters = sorted(indexed_selected_iters, key=lambda x: x[1].stride)
+        iter_idx_to_dim = {
+            seps[j]: i for i in range(len(region.buffer.shape)) for j in range(seps[i], seps[i + 1])
+        }
+        last_dim = None
+        inst_size = 1
+        selected_data_iter_ids = []
+        for i, it in indexed_selected_iters:
+            if last_dim is None:
+                inst_size *= it.extent
+                last_dim = iter_idx_to_dim[i]
+                selected_data_iter_ids.append(i)
+                continue
+            if iter_idx_to_dim[i] != last_dim:
+                break
+            inst_size *= it.extent
+            selected_data_iter_ids.append(i)
+        return InstructionRepr(region, inst_size, inst_repr.stride, selected_data_iter_ids)

python/tvm/backend/trn/operator/tile_primitive/private_alloc.py

@@ -0,0 +1,208 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Any
+
+from tvm.backend.trn.operator.tile_primitive.common import init_analyzer, nki_dim
+from tvm.backend.trn.operator.tile_primitive.dim_utils import get_ewise_dim_map
+from tvm.backend.trn.operator.tile_primitive.instruction_generator import InstructionGenerator
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, FloatImm, Stmt
+from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext
+from tvm.tirx.operator.tile_primitive.ops import (
+    BinaryReduce,
+    Copy,
+    Gemm,
+    ReduceOp,
+    UnaryOpWithBiasScale,
+    UnaryReduce,
+)
+from tvm.tirx.operator.tile_primitive.registry import f_op_dispatcher
+from tvm.tirx.stmt import TilePrimitiveCall
+
+
+def _scalar_dtype(scalar) -> str:
+    dtype = getattr(scalar, "dtype", None)
+    if dtype is not None:
+        return str(dtype)
+    ty = getattr(scalar, "ty", None)
+    dtype = getattr(ty, "dtype", None)
+    if dtype is None:
+        raise AttributeError(f"{type(scalar).__name__} has no dtype-bearing PrimType")
+    return str(dtype)
+
+
+def alloc_const_bias_trn(
+    op: TilePrimitiveCall, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: DispatchContext
+) -> dict[str, Any]:
+    bias = op.bias if op.bias is not None else FloatImm(op.dsts[0].buffer.dtype, 0.0)
+    if "const_bias" in op.workspace:
+        return {}
+    if not isinstance(bias, (FloatImm)):
+        return {}
+    par_size = op.dsts[0].buffer.layout.size("P")
+    max_inst_size = op.config.get("max_inst_size", 512)
+    if ("const_bias", bias.value) in buffer_dict:
+        bias_buffer, bias_init_stmt = buffer_dict[("const_bias", bias.value)]
+        old_shape = bias_buffer.shape
+        new_shape = [max(par_size, old_shape[0]), max(max_inst_size, old_shape[1])]
+        if new_shape[0] == old_shape[0] and new_shape[1] == old_shape[1]:
+            return {"const_bias": ("const_bias", bias.value)}
+    else:
+        new_shape = (par_size, max_inst_size)
+    new_buffer = T.buffer(
+        new_shape, dtype=_scalar_dtype(bias), scope="trn.sbuf", buffer_name="const_bias"
+    )
+
+    @T.prim_func
+    def const_bias_init():
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(0, par_size, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, max_inst_size, annotations={nki_dim: "F"}):
+                    T.evaluate(T.nki.memset(new_buffer[p_loop, f_loop], bias))
+        T.tvm_kernel_replace_point()
+
+    buffer_dict[("const_bias", bias.value)] = (new_buffer, const_bias_init.body)
+    return {"const_bias": ("const_bias", bias.value)}
+
+
+def alloc_partial_reduce_trn(
+    op: TilePrimitiveCall, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: DispatchContext
+) -> dict[str, Any]:
+    if "partial_reduce" in op.workspace:
+        return {}
+    f_op_dispatcher(op, sctx)
+    partial_reduce_buffer = None
+    if DispatchContext.kPrivateAlloc not in sctx.callbacks:
+        return {}
+    for buffer in sctx.callbacks[DispatchContext.kPrivateAlloc]:
+        if buffer.name == "partial_reduce":
+            partial_reduce_buffer = buffer
+            break
+    if partial_reduce_buffer is None:
+        return {}
+    # no reuse opportunity
+    buffer_dict[partial_reduce_buffer] = (partial_reduce_buffer, None)
+    return {"partial_reduce": partial_reduce_buffer}
+
+
+def alloc_identity_trn(
+    op: TilePrimitiveCall, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: DispatchContext
+) -> dict[str, Any]:
+    if "identity" in op.workspace:
+        return {}
+    par_size = op.srcs[0].buffer.layout.size("P")
+    if "identity" in buffer_dict:
+        identity_buffer, identity_init_stmt = buffer_dict["identity"]
+        old_shape = identity_buffer.shape
+        new_shape = [max(par_size, old_shape[0]), max(par_size, old_shape[1])]
+        if new_shape[0] == old_shape[0] and new_shape[1] == old_shape[1]:
+            return {"identity": "identity"}
+    else:
+        new_shape = (par_size, par_size)
+    new_buffer = T.buffer(
+        new_shape, dtype=op.srcs[0].buffer.dtype, scope="trn.sbuf", buffer_name="identity"
+    )
+
+    @T.prim_func
+    def identity_init():
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(0, par_size, annotations={nki_dim: "P"}):
+                for rhs_f_loop in T.serial(0, par_size, annotations={nki_dim: "F"}):
+                    T.evaluate(T.nki.identity(new_buffer[p_loop, rhs_f_loop], par_size))
+        T.tvm_kernel_replace_point()
+
+    buffer_dict["identity"] = (new_buffer, identity_init.body)
+    return {"identity": "identity"}
+
+
+def alloc_acc_psum_trn(
+    op: TilePrimitiveCall, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: DispatchContext
+) -> dict[str, Any]:
+    if "acc_psum" in op.workspace or op.dsts[0].buffer.scope() == "trn.psum":
+        return {}
+    par_size = op.dsts[0].buffer.layout.size("P")
+    acc_psum = T.buffer(
+        (8, par_size, 512),
+        "float32",
+        scope="trn.psum",
+        allocated_addr=(0, 0),
+        buffer_name="acc_psum",
+    )
+    # no reuse opportunity
+    buffer_dict[acc_psum] = (acc_psum, None)
+    return {"acc_psum": acc_psum}
+
+
+def alloc_copy_trn(
+    op: TilePrimitiveCall, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: DispatchContext
+) -> dict[str, Buffer]:
+    src_region = op.srcs[0]
+    dst_region = op.dsts[0]
+    analyzer = init_analyzer(sctx)
+    dim_map = get_ewise_dim_map(src_region, dst_region, analyzer)
+    inst_gen = InstructionGenerator([src_region, dst_region], analyzer)
+    inst_gen.link_buffer_regions(src_region, dst_region, dim_map)
+    if inst_gen.check_partition_dim_match(src_region, dst_region):
+        return {}
+
+    identity_dict = alloc_identity_trn(op, buffer_dict, sctx)
+    acc_psum_dict = alloc_acc_psum_trn(op, buffer_dict, sctx)
+    return identity_dict | acc_psum_dict
+
+
+def alloc_unary_reduce_trn(
+    op: TilePrimitiveCall, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: DispatchContext
+) -> dict[str, Buffer]:
+    if "max_inst_size" in op.config:
+        partial_reduce_dict = alloc_partial_reduce_trn(op, buffer_dict, sctx)
+        const_bias_dict = alloc_const_bias_trn(op, buffer_dict, sctx)
+        return partial_reduce_dict | const_bias_dict
+    else:
+        if "const_bias" in op.workspace and "partial_reduce" in op.workspace:
+            return {}
+        f_op_dispatcher(op, sctx)
+        partial_reduce_buffer = None
+        const_bias_buffer = None
+        if DispatchContext.kPrivateAlloc not in sctx.callbacks:
+            return {}
+        for buffer in sctx.callbacks[DispatchContext.kPrivateAlloc]:
+            if buffer.name == "partial_reduce":
+                partial_reduce_buffer = buffer
+            elif buffer.name == "const_bias":
+                const_bias_buffer = buffer
+        # no reuse opportunity
+        workspace_dict = {}
+        if partial_reduce_buffer is not None and "partial_reduce" not in op.workspace:
+            buffer_dict[partial_reduce_buffer] = (partial_reduce_buffer, None)
+            workspace_dict["partial_reduce"] = partial_reduce_buffer
+        if const_bias_buffer is not None and "const_bias" not in op.workspace:
+            assert len(sctx.callbacks[DispatchContext.kDeviceInitStmt]) == 1, (
+                "const_bias should have init"
+            )
+            init_stmt = sctx.callbacks[DispatchContext.kDeviceInitStmt][0]
+            buffer_dict[const_bias_buffer] = (const_bias_buffer, init_stmt)
+            workspace_dict["const_bias"] = const_bias_buffer
+        return workspace_dict
+
+
+UnaryOpWithBiasScale.get_private_buffers_trn = alloc_const_bias_trn
+ReduceOp.get_private_buffers_trn = alloc_partial_reduce_trn
+Copy.get_private_buffers_trn = alloc_copy_trn
+Gemm.get_private_buffers_trn = alloc_acc_psum_trn
+BinaryReduce.get_private_buffers_trn = alloc_partial_reduce_trn
+UnaryReduce.get_private_buffers_trn = alloc_unary_reduce_trn

python/tvm/backend/trn/operator/tile_primitive/reduction/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+from .default import *

python/tvm/backend/trn/operator/tile_primitive/reduction/default.py

@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Reduction dispatch variant registrations."""
+
+from tvm.tirx.operator.tile_primitive import register_dispatch
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
+
+from .utils import reduction_trn
+
+for _op_name, _op_type in {
+    "sum": ReduceOpType.SUM,
+    "max": ReduceOpType.MAX,
+    "min": ReduceOpType.MIN,
+}.items():
+
+    @register_dispatch(_op_name, "trn", variant="reduction", priority=0)
+    def _reduction_dispatch(op, sctx, _ty=_op_type):
+        return reduction_trn(op, _ty, sctx)

python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py

@@ -0,0 +1,167 @@
+# 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.
+
+"""Shared helpers for reduction schedules."""
+
+from tvm.backend.trn.layout import is_trainium_layout
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import init_analyzer, nki_dim
+from ..dim_utils import get_reduction_dim_map
+from ..instruction_generator import InstructionGenerator
+from ..workspace_utils import check_workspace_buffer
+
+reduce_ops = {ReduceOpType.SUM: "add", ReduceOpType.MAX: "max", ReduceOpType.MIN: "min"}
+
+
+def generate_intermediate_buffer(
+    dst_buffer_region: int, rfactor_size: int, workspace, sctx: DispatchContext
+):
+    """Generate an intermediate buffer for two-stage reduction if needed.
+
+    Returns:
+        Tuple[Optional[buffer], int]: The intermediate buffer and reduction factor size.
+    """
+    intermediate_shape = [dst_buffer_region.buffer.layout.size("P"), rfactor_size]
+
+    if "partial_reduce" in workspace:
+        intermediate_buffer = workspace["partial_reduce"]
+        check_workspace_buffer(intermediate_buffer, intermediate_shape, "trn.sbuf")
+    else:
+        assert sctx.alloc_only, (
+            "Partial reduce buffer must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first."  # noqa: E501
+        )
+        intermediate_buffer = T.buffer(
+            intermediate_shape,
+            dtype=dst_buffer_region.buffer.dtype,
+            scope="trn.sbuf",
+            buffer_name="partial_reduce",
+        )
+        sctx.add_alloc_buffer(intermediate_buffer)
+
+    return intermediate_buffer
+
+
+def reduction_trn(
+    op: TilePrimitiveCall, reduce_op: ReduceOpType, sctx: DispatchContext, negate: bool = False
+) -> PrimFunc | None:
+    """Schedule reduction operation on Trainium.
+
+    Args:
+        op: The operation call.
+        reduce_op: The reduction operation type.
+        sctx: The dispatch context.
+        negate: Whether to negate the result.
+
+    Returns:
+        Optional[PrimFunc]: The scheduled function, or None if not applicable.
+    """
+    if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
+        fail("requires Trainium target and thread exec_scope")
+
+    dst_buffer_region, src_buffer_region, axes, accum = op.args[:4]
+    assert not accum, "Accumulation is not supported for reduction on Trainium"
+    analyzer = init_analyzer(sctx)
+    assert reduce_op in reduce_ops, f"Unsupported reduce operation {reduce_op}"
+
+    # Extract buffers
+    dst = dst_buffer_region.buffer
+    src = src_buffer_region.buffer
+    axes = [i if i >= 0 else len(src.shape) + i for i in axes]
+    dim_map = get_reduction_dim_map(src_buffer_region, dst_buffer_region, axes, analyzer)
+
+    # Layout validation
+    assert all(
+        [
+            src.layout and dst.layout,
+            src.scope() == "trn.sbuf" or src.scope() == "trn.psum",
+            dst.scope() == "trn.sbuf",
+            is_trainium_layout(src.layout),
+            is_trainium_layout(dst.layout),
+            src.layout.size("P") == dst.layout.size("P"),
+        ]
+    ), "Invalid layout"
+
+    # Find maximum instruction size
+    inst_gen = InstructionGenerator([src_buffer_region, dst_buffer_region], analyzer)
+    inst_gen.link_buffer_regions(src_buffer_region, dst_buffer_region, dim_map)
+    inst_repr = inst_gen.find_max_inst_size_from_one_region(src_buffer_region, axes)
+    inst_size_limit = op.config.get("max_inst_size", None)
+    inst_repr.bound_inst_size(inst_size_limit, analyzer)
+    assert analyzer.can_prove(inst_repr.size > 1), "Instruction size must be greater than 1"
+
+    # Get partition size and extents
+    p_size = src.layout.size("P")
+    f_var = T.Var("F", "int32")
+    p_var = T.Var("P", "int32")
+    spatial_b_var = T.Var("sB", "int32")
+    reduction_b_var = T.Var("rB", "int32")
+    inst_gen.bind_inst_iter(src_buffer_region, f_var, inst_repr.size, inst_repr.stride, True)
+    inst_gen.bind_inst_iter(src_buffer_region, p_var, p_size, 1, False)
+    reduction_b_extent = inst_gen.fill_in_block_dim(src_buffer_region, reduction_b_var, axes)
+    spatial_b_extent = inst_gen.fill_in_block_dim(src_buffer_region, spatial_b_var)
+    # Get reduction operation code
+    opcode = reduce_ops[reduce_op]
+
+    # Generate intermediate buffer if needed
+    if reduction_b_extent != 1:
+        intermediate_buffer = generate_intermediate_buffer(
+            dst_buffer_region, reduction_b_extent, op.workspace, sctx
+        )
+
+    # fmt: off
+    # Single-stage reduction implementation
+    if reduction_b_extent == 1:
+        @T.prim_func
+        def impl():
+            for b_loop in T.serial(0, spatial_b_extent):
+                with T.attr(0, "tensorized_nki_instruction", 1):
+                    for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                        for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
+                            inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, spatial_b_var: b_loop})  # noqa: E501
+                            if inst_gen.make_guard(src_buffer_region):
+                                src_indices = T.meta_var(inst_gen.generate_indices(src_buffer_region))  # noqa: E501
+                                dst_indices = T.meta_var(inst_gen.generate_indices(dst_buffer_region))  # noqa: E501
+                                T.evaluate(T.nki.tensorreduce(dst[tuple(dst_indices)], src[tuple(src_indices)], opcode, negate, -1))  # noqa: E501
+        return impl
+    # Two-stage reduction implementation
+    else:
+        @T.prim_func
+        def two_stage_reduction():
+            for b_loop in T.serial(0, spatial_b_extent):
+                for reduction_b_loop in T.serial(0, reduction_b_extent):
+                    with T.attr(0, "tensorized_nki_instruction", 1):
+                        for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                            for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
+                                inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, spatial_b_var: b_loop, reduction_b_var: reduction_b_loop})  # noqa: E501
+                                if inst_gen.make_guard(src_buffer_region):
+                                    src_indices = T.meta_var(inst_gen.generate_indices(src_buffer_region))  # noqa: E501
+                                    T.evaluate(T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop], src[src_indices], opcode, False, -1))  # noqa: E501
+                with T.attr(0, "tensorized_nki_instruction", 1):
+                    for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                        for f_loop in T.serial(0, reduction_b_extent, annotations={nki_dim: "F"}):
+                            inst_gen.set_bind_map(src_buffer_region, {p_var: p_loop, f_var: 0, spatial_b_var: b_loop, reduction_b_var: f_loop})  # noqa: E501
+                            inst_gen.set_bind_map(dst_buffer_region, {p_var: p_loop, spatial_b_var: b_loop})  # noqa: E501
+                            if inst_gen.make_guard(src_buffer_region):
+                                dst_indices = T.meta_var(inst_gen.generate_indices(dst_buffer_region))  # noqa: E501
+                                T.evaluate(T.nki.tensorreduce(dst[dst_indices], intermediate_buffer[p_loop, f_loop], opcode, negate, -1))  # noqa: E501
+        return two_stage_reduction
+    # fmt: on

python/tvm/backend/trn/operator/tile_primitive/select/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+from .default import *

python/tvm/backend/trn/operator/tile_primitive/select/default.py

@@ -0,0 +1,145 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Implementation of select schedules."""
+
+from tvm.backend.trn.layout import is_trainium_layout
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, FloatImm, PrimFunc, TilePrimitiveCall
+from tvm.tirx.operator.tile_primitive import (
+    DispatchContext,
+    fail,
+    predicate,
+    register_dispatch,
+)
+from tvm.tirx.operator.tile_primitive.ops import Select
+
+from ..common import init_analyzer, nki_dim
+from ..dim_utils import get_ewise_dim_map
+from ..instruction_generator import InstructionGenerator
+
+
+def select_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    """Generate schedule for select operation on Trainium."""
+    if sctx.scope_kind != "thread":
+        fail("requires thread exec_scope for TRN select")
+
+    op = TilePrimitiveCall.downcast(op)
+    assert isinstance(op, Select), f"{op} is not a Select"
+
+    # Unpack operands
+    dst, true_value, false_value = *op.dsts, *op.srcs
+    pred = op.predicate
+
+    # Check that one of the sources is a float immediate
+    assert isinstance(true_value, FloatImm) or isinstance(false_value, FloatImm), (
+        f"{op} expects one of the source to be a float"
+    )
+
+    # Ensure true_value is the buffer and false_value is the float immediate
+    if isinstance(true_value, FloatImm):
+        pred = not pred
+        true_value, false_value = false_value, true_value
+
+    assert isinstance(true_value, BufferRegion), f"{op} expects one of the source to be a buffer"
+
+    # Initialize analyzer and validate buffers
+    analyzer = init_analyzer(sctx)
+
+    # Validate buffer layout and scope
+    buffer_conditions = [
+        dst.buffer.layout and true_value.buffer.layout,
+        dst.buffer.scope() == "trn.sbuf" and true_value.buffer.scope() == "trn.sbuf",
+        is_trainium_layout(true_value.buffer.layout),
+        is_trainium_layout(dst.buffer.layout),
+    ]
+
+    if not all(buffer_conditions):
+        assert False, f"scope or layout mismatch, {dst} vs {true_value}"
+
+    # Extract regions and validate dimensions
+    dst_extent = [r.extent for r in dst.region]
+    dst_extent_non_unit = [e for e in dst_extent if e != 1]
+    true_value_extent = [r.extent for r in true_value.region]
+    true_value_extent_non_unit = [e for e in true_value_extent if e != 1]
+
+    # Validate non-unit dimensions match
+    dims_match = len(true_value_extent_non_unit) == len(dst_extent_non_unit) and all(
+        analyzer.can_prove_equal(s, d)
+        for s, d in zip(true_value_extent_non_unit, dst_extent_non_unit)
+    )
+
+    if not dims_match:
+        assert False, f"shape or dimension mismatch, {dst} vs {true_value}"
+
+    # Bound buffer regions and find instruction size
+    inst_gen = InstructionGenerator([dst, true_value], analyzer)
+    dim_map = get_ewise_dim_map(dst, true_value, analyzer)
+    inst_gen.link_buffer_regions(dst, true_value, dim_map)
+    inst_repr = inst_gen.find_max_inst_size_from_one_region(dst)
+    inst_repr = inst_gen.fit_inst_tile_to_region(inst_repr, true_value)
+    inst_repr = inst_gen.restrict_inst_to_one_dim(inst_repr)
+    inst_repr.bound_inst_size(op.config.get("max_inst_size", 512), analyzer)
+
+    p_var = T.Var("p", "int32")
+    b_var = T.Var("b", "int32")
+    f_var = T.Var("f", "int32")
+    p_size = dst.buffer.layout.size("P")
+    inst_gen.bind_inst_iter(dst, f_var, inst_repr.size, inst_repr.stride, True)
+    inst_gen.bind_inst_iter(dst, p_var, p_size, 1, False)
+    b_extent = inst_gen.fill_in_block_dim(dst, b_var)
+
+    # Get buffer references and guard function
+    dst_buffer = dst.buffer
+    true_value_buffer = true_value.buffer
+
+    # fmt: off
+    @T.prim_func
+    def impl():
+        for b_loop in T.serial(0, b_extent):
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                    for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
+                        inst_gen.set_bind_map_all({f_var: f_loop, p_var: p_loop, b_var: b_loop})
+                        if inst_gen.make_guard(dst):
+                            dst_indices = T.meta_var(inst_gen.generate_indices(dst))
+                            true_value_indices = T.meta_var(inst_gen.generate_indices(true_value))
+                            pred = T.meta_var(analyzer.simplify(op.predicate.apply(inst_gen.generate_axes(dst))))  # noqa: E501
+                            T.evaluate(T.nki.affine_select(dst_buffer[tuple(dst_indices)], pred, true_value_buffer[tuple(true_value_indices)], false_value))  # noqa: E501
+    # fmt: on
+
+    return impl
+
+
+# Rich dispatcher variant for TRN select
+@register_dispatch(
+    "select",
+    "trn",
+    variant="default",
+    priority=10,
+    when=[
+        predicate(
+            "exec_scope",
+            lambda op, sctx: (
+                sctx.scope_kind == "thread",
+                f"unsupported exec_scope {sctx.scope_kind}",
+            ),
+        )
+    ],
+)
+def select_trn_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return select_trn(op, sctx)

python/tvm/backend/trn/operator/tile_primitive/unary/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+
+from .default import *
+from .utils import *
+from .with_bias_scale import *

python/tvm/backend/trn/operator/tile_primitive/unary/default.py

@@ -0,0 +1,89 @@
+# 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.
+
+"""Implementation of default unary operator dispatches."""
+
+from tvm.tirx import FloatImm, PrimFunc
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import MapOpType
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import init_analyzer
+from ..instruction_generator import InstructionGenerator
+from .utils import (
+    const_input_ops,
+    generate_unary_func,
+    non_activation_unary_map_ops,
+    try_find_inst_unary,
+)
+
+
+def unary_trn(op: TilePrimitiveCall, unary_op: MapOpType, sctx: DispatchContext) -> PrimFunc | None:
+    """Schedule unary operation on Trainium."""
+    # Check execution environment
+    if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
+        fail("requires Trainium target and thread exec_scope")
+
+    # Extract operation arguments
+    dst_buffer_region, _src = op.args
+
+    # Handle constant or buffer source
+    if isinstance(_src, FloatImm):
+        if unary_op not in const_input_ops:
+            assert False, f"Unsupported unary operation {unary_op} taking const as input"
+        CONST = _src
+        src_buffer_region = None
+    else:
+        CONST = None
+        src_buffer_region = _src
+
+    # Initialize analyzer and validate operation type
+    analyzer = init_analyzer(sctx)
+    assert unary_op in non_activation_unary_map_ops, f"Unsupported unary operation {unary_op}"
+
+    inst_gen = InstructionGenerator([dst_buffer_region, _src], analyzer)
+    # Find instruction parameters
+    if CONST is None:
+        inst_repr = try_find_inst_unary(dst_buffer_region, src_buffer_region, analyzer, inst_gen)
+    else:
+        inst_repr = try_find_inst_unary(dst_buffer_region, dst_buffer_region, analyzer, inst_gen)
+    # Generate and return the implementation function
+    return generate_unary_func(
+        dst_buffer_region,
+        _src,
+        inst_gen,
+        inst_repr,
+        unary_op,
+        None,  # No bias
+        None,  # No scale
+        analyzer,
+        op.workspace,
+        op.config,
+        sctx,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Registration: bind each default unary op name to its TRN schedule candidates.
+# ---------------------------------------------------------------------------
+from tvm.tirx.operator.tile_primitive import register_dispatch  # noqa: E402
+
+for _op_name, _op_type in {"reciprocal": MapOpType.RECIPROCAL, "memset": MapOpType.FILL}.items():
+
+    @register_dispatch(_op_name, "trn", variant="unary", priority=0)
+    def _unary_dispatch(op, sctx, _ty=_op_type):
+        return unary_trn(op, _ty, sctx)

python/tvm/backend/trn/operator/tile_primitive/unary/utils.py

@@ -0,0 +1,190 @@
+# 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.
+
+"""Shared helpers, op tables, and validation functions for unary operator dispatches."""
+
+from tvm.arith.analyzer import Analyzer
+from tvm.backend.trn.layout import is_trainium_layout
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, FloatImm
+from tvm.tirx.operator.tile_primitive.common import MapOpType
+
+from ..common import nki_dim
+from ..dim_utils import get_ewise_dim_map
+from ..instruction_generator import InstructionGenerator
+from ..workspace_utils import check_workspace_buffer
+
+# Operation type classifications
+non_activation_unary_map_ops = [MapOpType.RECIPROCAL, MapOpType.FILL]
+activation_map_ops = [MapOpType.SQRT, MapOpType.EXP]
+
+# Operation code table for instructions
+opcode_table = {MapOpType.SQRT: "sqrt", MapOpType.EXP: "exp"}
+
+# Operations that take constants as input
+const_input_ops = [MapOpType.FILL]
+
+
+def try_find_inst_unary(
+    dst_buffer_region: BufferRegion,
+    src_buffer_region: BufferRegion,
+    analyzer: Analyzer,
+    inst_gen: InstructionGenerator,
+    allowed_f_dim_dst: tuple[int] | None = None,
+    allowed_f_dim_src: tuple[int] | None = None,
+):
+    """Find instruction parameters for a unary operation."""
+    dst = dst_buffer_region.buffer
+    src = src_buffer_region.buffer
+
+    # Validate buffer layouts and scopes
+    valid_layout_scope = all(
+        [
+            src.layout and dst.layout,
+            src.scope() in ("trn.sbuf", "trn.psum"),
+            dst.scope() == "trn.sbuf",
+            is_trainium_layout(src.layout),
+            is_trainium_layout(dst.layout),
+        ]
+    )
+
+    if not valid_layout_scope:
+        assert False, (
+            f"scope or layout mismatch, src: {src_buffer_region}, dst: {dst_buffer_region}"
+        )
+
+    # Extract and validate dimensions
+    dst_region = dst_buffer_region.region
+    src_region = src_buffer_region.region
+
+    dst_extent = [r.extent for r in dst_region]
+    src_extent = [r.extent for r in src_region]
+
+    dst_extent_nonunit = [e for e in dst_extent if e != 1]
+    src_extent_nonunit = [e for e in src_extent if e != 1]
+
+    # Verify dimensions match
+    dims_match = len(src_extent_nonunit) == len(dst_extent_nonunit) and all(
+        analyzer.can_prove_equal(s, d) for s, d in zip(src_extent_nonunit, dst_extent_nonunit)
+    )
+
+    if not dims_match:
+        assert False, (
+            f"shape or dimension mismatch, src: {src_buffer_region}, dst: {dst_buffer_region}"
+        )
+    dim_map = get_ewise_dim_map(src_buffer_region, dst_buffer_region, analyzer)
+    inst_gen.link_buffer_regions(src_buffer_region, dst_buffer_region, dim_map)
+    # Find optimal instruction parameters
+    inst_repr = inst_gen.find_max_inst_size_from_one_region(dst_buffer_region, allowed_f_dim_dst)
+    inst_repr = inst_gen.fit_inst_tile_to_region(inst_repr, src_buffer_region, allowed_f_dim_src)
+    return inst_repr
+
+
+def get_const_bias_tensor(bias, shape, dtype, workspace, sctx):
+    """Create or retrieve a constant bias tensor."""
+    if "const_bias" not in workspace:
+        assert sctx.alloc_only, (
+            "Constant bias tensor must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first."  # noqa: E501
+        )
+        # Create new bias buffer
+        bias_buffer = T.buffer(shape, dtype, scope="trn.sbuf", buffer_name="const_bias")
+        sctx.add_alloc_buffer(bias_buffer)
+
+        @T.prim_func
+        def const_bias_init():
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for p_loop in T.serial(0, shape[0], annotations={nki_dim: "P"}):
+                    for f_loop in T.serial(0, shape[1], annotations={nki_dim: "F"}):
+                        T.evaluate(T.nki.memset(bias_buffer[p_loop, f_loop], bias))
+            T.tvm_kernel_replace_point()
+
+        sctx.add_init_stmt(const_bias_init.body)
+    else:
+        # Use existing bias buffer
+        bias_buffer = workspace["const_bias"]
+        check_workspace_buffer(bias_buffer, shape, "trn.sbuf")
+
+    return bias_buffer
+
+
+def generate_unary_func(
+    dst_buffer_region,
+    _src,
+    inst_gen: InstructionGenerator,
+    inst_repr,
+    unary_op,
+    bias,
+    scale,
+    analyzer,
+    workspace,
+    config,
+    sctx,
+):
+    """Generate a function that implements a unary operation."""
+    # Prepare parameters
+    p_size = dst_buffer_region.buffer.layout.size("P")
+
+    # Apply instruction size limits if specified
+    inst_size_limit = config.get("max_inst_size", 512)
+    inst_repr.bound_inst_size(inst_size_limit, analyzer)
+
+    f_var = T.Var("F", "int32")
+    p_var = T.Var("P", "int32")
+    b_var = T.Var("B", "int32")
+    inst_gen.bind_inst_iter(dst_buffer_region, f_var, inst_repr.size, inst_repr.stride, True)
+    inst_gen.bind_inst_iter(dst_buffer_region, p_var, p_size, 1, False)
+    b_extent = inst_gen.fill_in_block_dim(dst_buffer_region, b_var)
+
+    # Get operation code if available
+    opcode = opcode_table.get(unary_op, None)
+
+    # Extract buffers
+    dst = dst_buffer_region.buffer
+    src = _src.buffer if isinstance(_src, BufferRegion) else None
+
+    # Handle bias tensor
+    if isinstance(bias, FloatImm | float):
+        bias_buffer = get_const_bias_tensor(
+            bias, (p_size, inst_repr.size), dst.dtype, workspace, sctx
+        )
+    elif isinstance(bias, BufferRegion):
+        bias_buffer = bias.buffer
+
+    # fmt: off
+    @T.prim_func
+    def impl():
+        for b_loop in T.serial(0, b_extent):
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
+                    for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
+                        inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, b_var: b_loop})
+                        dst_indices = T.meta_var(inst_gen.generate_indices(dst_buffer_region))
+                        if inst_gen.make_guard(dst_buffer_region):
+                            if unary_op == MapOpType.FILL:
+                                T.evaluate(T.nki.memset(dst[tuple(dst_indices)], _src))
+                            else:
+                                src_indices = T.meta_var(inst_gen.generate_indices(_src))
+                                if unary_op == MapOpType.RECIPROCAL:
+                                    T.evaluate(T.nki.reciprocal(dst[tuple(dst_indices)], src[tuple(src_indices)]))  # noqa: E501
+                                elif isinstance(bias, BufferRegion):
+                                    bias_indices = T.meta_var(inst_gen.generate_indices(bias))
+                                    T.evaluate(T.nki.activation(dst[tuple(dst_indices)], src[tuple(src_indices)], opcode, scale=scale, bias=bias_buffer[tuple(bias_indices)]))  # noqa: E501
+                                else:
+                                    T.evaluate(T.nki.activation(dst[tuple(dst_indices)], src[tuple(src_indices)], opcode, scale=scale, bias=bias_buffer[p_loop, f_loop]))  # noqa: E501
+    # fmt: on
+
+    return impl

python/tvm/backend/trn/operator/tile_primitive/unary/with_bias_scale.py

@@ -0,0 +1,87 @@
+# 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.
+
+"""Implementation of unary with bias and scale operator dispatches."""
+
+from tvm.tirx import BufferRegion, PrimFunc
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import MapOpType
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..binary import try_find_inst_nary
+from ..common import init_analyzer
+from ..instruction_generator import InstructionGenerator
+from .utils import activation_map_ops, generate_unary_func, try_find_inst_unary
+
+
+def unary_with_bias_scale_trn(
+    op: TilePrimitiveCall, unary_op: MapOpType = MapOpType.SQRT, sctx: DispatchContext = None
+) -> PrimFunc | None:
+    """Schedule unary operation with bias and scale on Trainium."""
+    # Check execution environment
+    if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
+        fail("requires Trainium target and thread exec_scope")
+
+    # Extract operation arguments with defaults
+    dst_buffer_region, src_buffer_region, _bias, scale = op.args
+    scale = 1.0 if scale is None else scale
+    _bias = 0.0 if _bias is None else _bias
+
+    # Initialize analyzer and validate operation type
+    analyzer = init_analyzer(sctx)
+    assert unary_op in activation_map_ops, f"Unsupported activation operation {unary_op}"
+
+    # Find instruction parameters
+    inst_gen = InstructionGenerator([dst_buffer_region, src_buffer_region, _bias], analyzer)
+    if isinstance(_bias, BufferRegion):
+        inst_repr, _, _ = try_find_inst_nary(
+            dst_buffer_region,
+            [src_buffer_region, _bias],
+            analyzer,
+            inst_gen,
+            allow_first_op_tensortensor=False,
+        )
+    else:
+        # Handle scalar bias
+        inst_repr = try_find_inst_unary(dst_buffer_region, src_buffer_region, analyzer, inst_gen)
+
+    # Generate and return the implementation function
+    return generate_unary_func(
+        dst_buffer_region,
+        src_buffer_region,
+        inst_gen,
+        inst_repr,
+        unary_op,
+        _bias,
+        scale,
+        analyzer,
+        op.workspace,
+        op.config,
+        sctx,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Registration: bind each unary_with_bias_scale op name to its TRN schedule candidates.
+# ---------------------------------------------------------------------------
+from tvm.tirx.operator.tile_primitive import register_dispatch  # noqa: E402
+
+for _op_name, _op_type in {"sqrt": MapOpType.SQRT, "exp": MapOpType.EXP}.items():
+
+    @register_dispatch(_op_name, "trn", variant="unary_with_bias_scale", priority=0)
+    def _unary_bs_dispatch(op, sctx, _ty=_op_type):
+        return unary_with_bias_scale_trn(op, _ty, sctx)

python/tvm/backend/trn/operator/tile_primitive/workspace_utils.py

@@ -0,0 +1,54 @@
+# 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.
+
+"""Workspace buffer utilities for TRN operator scheduling."""
+
+from tvm.tirx import Buffer
+
+largest_psum_per_bank = 512
+max_psum_banks = 8
+
+
+def check_workspace_buffer(buffer: Buffer, shape: tuple[int], scope: str):
+    """Check if a workspace buffer is valid.
+
+    Parameters
+    ----------
+    buffer : Buffer
+        The workspace buffer to check
+    shape : Tuple[int]
+        The required shape
+    scope : str
+        The required scope
+
+    Raises
+    ------
+    AssertionError :
+        If the buffer is invalid
+    """
+    assert buffer.scope() == scope, f"workspace buffer must be a {scope} buffer"
+    assert buffer.layout is None, "workspace buffer must not have a layout"
+    if scope == "trn.psum":
+        # the number of psum banks used is inferred from the shape
+        # only check p and f dims
+        assert all(x >= y for x, y in zip(buffer.shape[1:], shape)), (
+            f"workspace buffer must have enough size, {buffer.shape[1:]} cannot cover {shape}"
+        )
+    else:
+        assert all(x >= y for x, y in zip(buffer.shape, shape)), (
+            f"workspace buffer must have enough size, {buffer.shape} cannot cover {shape}"
+        )

python/tvm/backend/trn/pipeline.py

@@ -0,0 +1,58 @@
+# 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.
+"""Trainium TIRX pipeline entrypoints."""
+
+import tvm
+from tvm import tirx
+from tvm.tirx.compilation_pipeline import finalize_host_passes
+
+from . import transform as trn_transform
+
+
+def trn_pipeline():
+    """The Trainium pipeline used in tvm.tirx.build."""
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
+        """Lower TIRx for the Trainium backend."""
+        passes = [
+            trn_transform.TrnPrivateBufferAlloc(),
+            trn_transform.TrnNaiveAllocator(),
+            tirx.transform.TilePrimitiveDispatch(),
+            trn_transform.LowerTrainiumLayout(),
+            tvm.s_tir.transform.DecorateDeviceScope(),
+            tirx.transform.StmtSimplify(),
+            tirx.transform.LowerTIRxOpaque(),
+            tvm.s_tir.transform.LoopPartition(),
+            tvm.s_tir.transform.HoistIfThenElse(),
+            tirx.transform.StmtSimplify(),
+            tirx.transform.RemoveNoOp(),
+            tirx.transform.AnnotateEntryFunc(),
+            tirx.transform.SplitHostDevice(),
+            tirx.transform.MakePackedAPI(),
+        ]
+        return tvm.ir.transform.Sequential(passes)(mod)
+
+    return _pipeline, finalize_host_passes, finalize_device_passes_trn
+
+
+def finalize_device_passes_trn():  # pylint: disable=unused-argument
+    """The finalization passes for the Trainium backend."""
+    return tvm.ir.transform.Sequential([tirx.transform.StmtSimplify()])
+
+
+__all__ = ["finalize_device_passes_trn", "trn_pipeline"]

python/tvm/backend/trn/script.py

@@ -0,0 +1,58 @@
+# 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.
+"""Trainium TVMScript namespaces."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from . import op as _trn_op
+
+OpWrapper = Callable[[Callable[..., Any]], Callable[..., Any]]
+
+
+def _default_op_wrapper() -> OpWrapper:
+    from tvm.tirx.script.builder.ir import _op_wrapper  # pylint: disable=import-outside-toplevel
+
+    return _op_wrapper
+
+
+class NKINamespace:
+    """The NKI instructions submodule."""
+
+    def __init__(self, op_wrapper: OpWrapper | None = None):
+        wrap = op_wrapper or _default_op_wrapper()
+        self.load = wrap(_trn_op.nki_load)
+        self.store = wrap(_trn_op.nki_store)
+        self.tensor_copy = wrap(_trn_op.nki_tensor_copy)
+        self.matmul = wrap(_trn_op.nki_matmul)
+        self.activation = wrap(_trn_op.nki_activation)
+        self.activation_reduce = wrap(_trn_op.nki_activation_reduce)
+        self.reciprocal = wrap(_trn_op.nki_reciprocal)
+        self.tensorreduce = wrap(_trn_op.nki_tensorreduce)
+        self.tensortensor = wrap(_trn_op.nki_tensortensor)
+        self.tensorscalar = wrap(_trn_op.nki_tensorscalar)
+        self.tensorscalar_reduce = wrap(_trn_op.nki_tensorscalar_reduce)
+        self.scalar_tensor_tensor = wrap(_trn_op.nki_scalar_tensor_tensor)
+        self.scalar_tensor_scalar = wrap(_trn_op.nki_scalar_tensor_scalar)
+        self.memset = wrap(_trn_op.nki_memset)
+        self.identity = wrap(_trn_op.nki_identity)
+        self.affine_select = wrap(_trn_op.nki_affine_select)
+
+
+__all__ = ["NKINamespace", "OpWrapper"]

python/tvm/backend/trn/target_tags.py

@@ -0,0 +1,36 @@
+# 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.
+"""AWS Trainium target tags."""
+
+from tvm.target import register_tag
+
+
+def _register_aws_trn1_tag(name, cores):
+    register_tag(
+        name,
+        {
+            "kind": "trn",
+            "num-cores": cores,
+            "partition_size": 128,
+            "max_sbuf_size_per_partition": 196608,
+            "max_psum_size_per_partition": 16384,
+        },
+    )
+
+
+_register_aws_trn1_tag("aws/trn1/trn1.2xlarge", 2)
+_register_aws_trn1_tag("aws/trn1/trn1.32xlarge", 32)

python/tvm/backend/trn/transform/__init__.py

@@ -0,0 +1,53 @@
+# 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.
+"""Trainium-specific TIRX transformations."""
+# pylint: disable=invalid-name
+
+from tvm_ffi import get_global_func
+
+import tvm
+from tvm import tirx
+
+_LAZY_TRANSFORMS = {
+    "TrnNaiveAllocator": ".naive_allocator",
+    "TrnPrivateBufferAlloc": ".private_buffer_alloc",
+}
+
+
+def LowerTrainiumLayout():
+    """Lower Trainium layouts to backend physical buffer shapes and indices."""
+    return get_global_func("tirx.backend.trn.transform.LowerTrainiumLayout")()
+
+
+def LowerTIRx():
+    """Lower TIRx tile primitive calls for the Trainium backend."""
+    return tvm.ir.transform.Sequential(
+        [tirx.transform.TilePrimitiveDispatch(), LowerTrainiumLayout()],
+        name="tirx.backend.trn.LowerTIRx",
+    )
+
+
+def __getattr__(name):
+    target = _LAZY_TRANSFORMS.get(name)
+    if target is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    from importlib import import_module  # pylint: disable=import-outside-toplevel
+
+    return getattr(import_module(target, __name__), name)
+
+
+__all__ = ["LowerTIRx", "LowerTrainiumLayout", "TrnNaiveAllocator", "TrnPrivateBufferAlloc"]

python/tvm/backend/trn/transform/naive_allocator.py

@@ -0,0 +1,99 @@
+# 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 functools
+
+from tvm.tirx import AllocBuffer, IntImm
+from tvm.tirx.buffer import Buffer
+from tvm.tirx.stmt_functor import StmtVisitor
+from tvm.tirx.transform.common import BufferReplacer
+from tvm.tirx.transform.function_pass import prim_func_pass
+
+
+def is_const_shape(buffer: Buffer) -> bool:
+    for i in buffer.shape:
+        if not isinstance(i, IntImm):
+            return False
+    return True
+
+
+def get_buffer_size(buffer: Buffer) -> int:
+    if buffer.scope() == "trn.sbuf":
+        if buffer.layout is None:
+            # the first dimension is partition size
+            num_elem = functools.reduce(lambda x, y: x * y, buffer.shape[1:])
+        else:
+            par_size = buffer.layout.size("P")
+            num_elem = functools.reduce(lambda x, y: x * y, buffer.shape) // par_size
+    elif buffer.scope().startswith("shared"):
+        num_elem = functools.reduce(lambda x, y: x * y, buffer.shape)
+    else:
+        return None
+    if not is_const_shape(buffer):
+        raise ValueError(
+            f"Buffer {buffer.name} has non-constant shape. Do not know how to allocate it."
+        )
+    return int(num_elem * buffer.dtype.dtype.itemsize)
+
+
+class AllocInfoCollector(StmtVisitor):
+    def __init__(self):
+        super().__init__()
+        self.alloc_pool_start = 0
+
+    def visit_alloc_buffer_(self, op: AllocBuffer):
+        super().visit_alloc_buffer_(op)
+        buffer = op.buffer
+        if len(buffer.allocated_addr) == 0:
+            return op
+        buffer_size = get_buffer_size(buffer)
+        if buffer_size is None:
+            return op
+        self.alloc_pool_start = max(self.alloc_pool_start, buffer.allocated_addr[-1] + buffer_size)
+
+
+class AllocMutator(BufferReplacer):
+    def __init__(self, alloc_pool_start: int):
+        super().__init__()
+        self.alloc_offset = alloc_pool_start
+
+    def visit_alloc_buffer_(self, op: AllocBuffer):
+        changed = False
+        buffer = op.buffer
+        buffer_size = get_buffer_size(buffer)
+        if len(buffer.allocated_addr) > 0 or buffer_size is None:
+            pass
+        else:
+            new_buffer = buffer.with_allocated_addr([self.alloc_offset])
+            self.buffer_map[buffer] = new_buffer
+            changed = True
+            self.alloc_offset += buffer_size
+
+        op = super().visit_alloc_buffer_(op)
+        if changed:
+            return AllocBuffer(new_buffer, op.annotations, op.span)
+        return op
+
+
+@prim_func_pass(opt_level=0, name="TrnNaiveAllocator")
+class TrnNaiveAllocator:
+    def transform_function(self, func, mod, ctx):
+        collector = AllocInfoCollector()
+        collector(func.body)
+        mutator = AllocMutator(collector.alloc_pool_start)
+        new_body = mutator(func.body)
+        return func.with_body(new_body)

python/tvm/backend/trn/transform/private_buffer_alloc.py

@@ -0,0 +1,138 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+from tvm.ir import Range
+from tvm.target import Target
+from tvm.tirx.buffer import Buffer
+from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext
+from tvm.tirx.stmt import (
+    AllocBuffer,
+    AttrStmt,
+    For,
+    SeqStmt,
+    Stmt,
+    TilePrimitiveCall,
+)
+from tvm.tirx.stmt_functor import StmtMutator, StmtVisitor
+from tvm.tirx.transform.common import seek_kernel_replace_point
+from tvm.tirx.transform.function_pass import prim_func_pass
+
+
+class PrivateAllocCollector(StmtVisitor):
+    def __init__(self, target: Target):
+        super().__init__()
+        self.target = target
+        self.launch_params = {}
+        self.var_range_map = {}
+        self.buffer_dict = {}
+        self.private_buf_refs = {}
+
+    def visit_attr_(self, op: AttrStmt):
+        if op.attr_key == "thread_extent":
+            self.launch_params[op.node.thread_tag] = op.value
+        super().visit_attr_(op)
+
+    def visit_for_(self, op: For):
+        self.var_range_map[op.loop_var] = Range.from_min_extent(op.min, op.extent)
+        super().visit_for_(op)
+
+    def visit_op_call_(self, op: TilePrimitiveCall):
+        # Scope is a per-call field on the node; read it directly.
+        exec_scope = op.scope
+        scope_kind = op.scope.name
+        sctx = DispatchContext(
+            target=self.target,
+            exec_scope=exec_scope,
+            launch_params=self.launch_params,
+            var_range_map=self.var_range_map,
+            alloc_only=True,
+            scope_kind=scope_kind,
+        )
+        op = TilePrimitiveCall.downcast(op)
+        private_buf_refs = op.get_private_buffers(self.buffer_dict, sctx)
+        self.private_buf_refs[op] = private_buf_refs
+
+
+class PrivateAllocMutator(StmtMutator):
+    def __init__(
+        self,
+        alloc_buffers: list[Buffer],
+        init_stmts: list[Stmt],
+        added_workspace: dict[TilePrimitiveCall, dict[str, Buffer]],
+    ):
+        super().__init__()
+        self.alloc_buffers = alloc_buffers
+        self.init_stmts = init_stmts
+        self.added_workspace = added_workspace
+        self.is_outer_block = True
+
+    def visit_attr_(self, op: AttrStmt):
+        # AttrStmt(kDeviceEntry) marks the device-region root: inject the
+        # collected init stmts + alloc_buffers into its body.
+        if op.attr_key == "tirx.device_entry":
+            is_outer_block = self.is_outer_block
+            self.is_outer_block = False
+            op = super().visit_attr_(op)
+            if is_outer_block:
+                body = op.body
+                for stmt in self.init_stmts:
+                    body = seek_kernel_replace_point(stmt, body)
+                for buffer in reversed(self.alloc_buffers):
+                    body = SeqStmt([AllocBuffer(buffer), body])
+                return AttrStmt(op.node, op.attr_key, op.value, body)
+            return op
+        return super().visit_attr_(op)
+
+    def visit_op_call_(self, op):
+        if op not in self.added_workspace:
+            return op
+        new_workspace = dict(op.workspace)
+        new_workspace.update(self.added_workspace[op])
+        op = TilePrimitiveCall.downcast(op).with_workspace(new_workspace)
+        return op
+
+
+def private_alloc(stmt: Stmt, target: Target) -> Stmt:
+    collector = PrivateAllocCollector(target)
+    collector(stmt)
+
+    alloc_buffers = [buffer for buffer, _ in collector.buffer_dict.values()]
+    init_stmts = [stmt for _, stmt in collector.buffer_dict.values() if stmt is not None]
+    added_workspace = {
+        op: {
+            name: collector.buffer_dict[ref][0]
+            for name, ref in collector.private_buf_refs[op].items()
+        }
+        for op in collector.private_buf_refs
+    }
+
+    mutator = PrivateAllocMutator(alloc_buffers, init_stmts, added_workspace)
+    return mutator(stmt)
+
+
+@prim_func_pass(opt_level=0, name="TrnPrivateBufferAlloc")
+class TrnPrivateBufferAlloc:
+    """Generate private buffer allocations for each TilePrimitiveCall"""
+
+    def transform_function(self, func, mod, ctx):
+        target = func.attrs.get("target", None)
+        if target is None:
+            target = Target.current(allow_none=False)
+        new_body = private_alloc(func.body, target)
+        new_func = func.with_body(new_body)
+        return new_func