Trainium Authoring and Support APIs

Trainium Authoring and Support APIs#

NKI authoring namespace#

The Trainium backend installs Tx.nki. Its current operations are load, store, tensor_copy, matmul, activation, activation_reduce, reciprocal, tensorreduce, tensortensor, tensorscalar, tensorscalar_reduce, scalar_tensor_tensor, scalar_tensor_scalar, memset, identity, and affine_select.

class tvm_tirx.backend.trn.script.NKINamespace

The NKI instructions submodule.

static tensorized_instruction()

Tensorize the NKI instructions in the region body.

Layout helpers#

Trainium-specific TIRx layout helpers.

tvm_tirx.backend.trn.layout.is_trainium_layout(layout: Layout | None) → bool

Return whether a layout uses only Trainium memory axes.

Compilation pipeline#

Trainium TIRX pipeline entrypoints.

tvm_tirx.backend.trn.pipeline.finalize_device_passes_trn()

Construct the statement-simplification stage for Trainium device IR.

tvm_tirx.backend.trn.pipeline.trn_pipeline(*, target=None)

Construct the Trainium lowering and host/device finalization stages.

Parameters:

target (Target or None) – Target to bind before lowering, or None to retain existing targets.

Returns:

stages – The lowering ModulePass, host-finalization factory, and device-finalization factory consumed by tvm_tirx.tirx.compile.

Return type:

tuple

Transforms#

Trainium-specific TIRX transformations.

tvm_tirx.backend.trn.transform.LowerTIRx()

Run TileDispatch followed by LowerTrainiumLayout.

Run TrnPrivateTensorAlloc and TrnNaiveAllocator first when tiles need workspaces or allocation addresses. Tile dispatch requires a Trainium target on the function or in the current target scope.

tvm_tirx.backend.trn.transform.LowerTrainiumLayout()

Lower logical tensor accesses to Trainium physical shapes and indices.

SBUF uses partition/free dimensions; PSUM adds a bank dimension. Global memory is flattened. Parameter ABI shapes stay logical. Run after tile dispatch.

Examples

import tvm_tirx as tvm
from tvm_tirx.script import tirx as T
from tvm_tirx.backend.trn.axis import P, F
from tvm_tirx.backend.trn.transform import LowerTrainiumLayout

@T.ir_module
class Before:
    @T.function
    def main():
        T.device_entry()
        A = T.alloc_tensor(ty_args=[T.Tensor(
            (128, 2, 2), "float32", scope="trn.sbuf",
            layout=T.TileLayout(T.S[(128, 4):(1 @ P, 1 @ F)]))])
        A[0, 1, 0] = T.float32(1)

@T.ir_module
class Expected:
    @T.function
    def main():
        T.device_entry()
        A = T.alloc_tensor(ty_args=[T.Tensor((128, 4), "float32", scope="trn.sbuf")])
        A[0, 2] = T.float32(1)

after = LowerTrainiumLayout()(Before)
tvm.ir.assert_structural_equal(after, Expected)
class tvm_tirx.backend.trn.transform.TrnNaiveAllocator(*args, **kwargs)

Assign consecutive byte offsets to unaddressed tensor allocations.

Handles constant-sized trn.sbuf, shared, and shared.dyn tensors; SBUF sizes count bytes per partition. New offsets follow existing allocations, without lifetime-based reuse.

Examples

import tvm_tirx as tvm
from tvm_tirx.script import tirx as T
from tvm_tirx.backend.trn.transform import TrnNaiveAllocator

@T.ir_module
class Before:
    @T.function
    def main():
        A = T.alloc_tensor(ty_args=[T.Tensor((128, 4), "float32", scope="trn.sbuf")])

@T.ir_module
class Expected:
    @T.function
    def main():
        A = T.alloc_tensor((0,), ty_args=[T.Tensor((128, 4), "float32", scope="trn.sbuf")])

after = TrnNaiveAllocator()(Before)
tvm.ir.assert_structural_equal(after, Expected)
class tvm_tirx.backend.trn.transform.TrnPrivateTensorAlloc(*args, **kwargs)

Allocate and initialize instruction-private Trainium workspaces.

Insert workspaces at the device-entry root and attach them to tile calls. Run before tile dispatch with a Trainium target. Calls needing no workspace are unchanged.

Examples

import tvm_tirx as tvm
from tvm_tirx.script import tirx as T
from tvm_tirx.backend.trn.axis import P, F
from tvm_tirx.backend.trn.transform import TrnPrivateTensorAlloc

tile = T.Tensor((128, 512), "float32", scope="trn.sbuf",
                layout=T.TileLayout(T.S[(128, 512):(1 @ P, 1 @ F)]))

@T.ir_module
class Before:
    @T.function
    def main():
        T.device_entry()
        A = T.alloc_tensor(ty_args=[tile])
        C = T.alloc_tensor(ty_args=[tile])
        T.trn.tile.activation(C, A, bias=T.float32(1), opcode="exp")

@T.ir_module
class Expected:
    @T.function
    def main():
        T.device_entry()
        const_bias = T.alloc_tensor(
            ty_args=[T.Tensor((128, 512), "float32", scope="trn.sbuf")])
        with T.nki.tensorized_instruction():
            for p in T.serial(128, annotations={"nki_dim": "P"}):
                for f in T.serial(512, annotations={"nki_dim": "F"}):
                    T.nki.memset(const_bias[p, f], T.float32(1))
        A = T.alloc_tensor(ty_args=[tile])
        C = T.alloc_tensor(ty_args=[tile])
        T.trn.tile.activation(C, A, bias=T.float32(1), const_bias=const_bias, opcode="exp")

with tvm.target.Target("aws/trn1/trn1.2xlarge"):
    after = TrnPrivateTensorAlloc()(Before)
tvm.ir.assert_structural_equal(after, Expected)
Contents