tvm_tirx.backend.trn.transform.TrnPrivateTensorAlloc

class tvm_tirx.backend.trn.transform.TrnPrivateTensorAlloc(*args, **kwargs)

Bases: FunctionPass

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)