tvm_tirx.backend.trn.transform

Contents

tvm_tirx.backend.trn.transform#

Trainium-specific TIRX transformations.

Classes#

AnnotateKernelArguments(*args, **kwargs)

Order kernel inputs before outputs and annotate their per-kernel input count.

TrnNaiveAllocator(*args, **kwargs)

Assign consecutive byte offsets to unaddressed tensor allocations.

TrnPrivateTensorAlloc(*args, **kwargs)

Allocate and initialize instruction-private Trainium workspaces.

Functions#

LowerTrainiumLayout()

Lower logical tensor accesses to Trainium physical shapes and indices.

LowerTIRx()

Run TileDispatch followed by LowerTrainiumLayout.
Contents