tt-mlir documentation

ArchAttr

TT Arch

Syntax:

#tt.arch<
  ::mlir::tt::Arch   # value
>

Enum cases:

• grayskull (Grayskull)
• wormhole_b0 (WormholeB0)
• blackhole (Blackhole)

Parameters:

ArgumentAllocationAttr

Argument allocation attribute in TT dialect

Syntax:

#tt.arg_alloc<
  uint64_t,   # address
  uint64_t,   # size
  MemorySpace   # memorySpace
>

Holds the metadata for the allocation of an function argument i.e. for graph inputs.

Parameters:

ArgumentTypeAttr

Argument Type

Syntax:

#tt.argument_type<
  ::mlir::tt::ArgumentType   # value
>

Enum cases:

• input (Input)
• parameter (Parameter)
• constant (Constant)

Parameters:

CPUDescAttr

TT cpu_desc attribute

Syntax:

#tt.cpu_desc<
  CPURole,   # role
  StringAttr   # target_triple
>

TT cpu_desc attribute

Parameters:

CPURoleAttr

TT CPU Role

Syntax:

#tt.cpu_role<
  ::mlir::tt::CPURole   # value
>

Enum cases:

• host (Host)
• device (Device)

Parameters:

ChipChannelAttr

TT chip_channel attribute

Syntax:

#tt.chip_channel<
  unsigned,   # deviceId0
  ::llvm::ArrayRef<int64_t>,   # ethernetCoreCoord0
  unsigned,   # deviceId1
  ::llvm::ArrayRef<int64_t>   # ethernetCoreCoord1
>

TT chip_channel attribute

Parameters:

ChipCoordAttr

TT chip_coord attribute

Syntax:

#tt.chip_coord<
  unsigned,   # rack
  unsigned,   # shelf
  unsigned,   # y
  unsigned   # x
>

TT chip_coord attribute

Parameters:

ChipDescAttr

TT chip_desc attribute

Syntax:

#tt.chip_desc<
  ArchAttr,   # arch
  ::llvm::ArrayRef<int64_t>,   # grid
  unsigned,   # l1Size
  unsigned,   # numDramChannels
  unsigned,   # dramChannelSize
  unsigned,   # nocL1AddressAlignBytes
  unsigned,   # pcieAddressAlignBytes
  unsigned,   # nocDRAMAddressAlignBytes
  unsigned,   # l1UnreservedBase
  unsigned,   # eriscL1UnreservedBase
  unsigned,   # dramUnreservedBase
  unsigned,   # dramUnreservedEnd
  ChipPhysicalCoresAttr,   # chipPhysicalCores
  ::llvm::ArrayRef<DataTypeAttr>,   # supportedDataTypes
  ::llvm::ArrayRef<TileSizeAttr>,   # supportedTileSizes
  unsigned   # numCBs
>

TT chip_desc attribute

Parameters:

ChipPhysicalCoresAttr

TT chip_physical_cores attribute

Syntax:

#tt.chip_physical_cores<
  ::llvm::ArrayRef<CoreCoordAttr>,   # worker
  ::llvm::ArrayRef<CoreCoordAttr>,   # dram
  ::llvm::ArrayRef<CoreCoordAttr>,   # eth
  ::llvm::ArrayRef<CoreCoordAttr>   # eth_inactive
>

TT chip_physical_cores attribute containing arrays of physical cores by core type in order of logical cores.

Parameters:

CoreCoordAttr

TT core_coord attribute

Syntax:

#tt.core_coord<
  int64_t,   # y
  int64_t   # x
>

TT core_coord attribute containing a single physical core coordinate.

Parameters:

DataTypeAttr

TT DataTypes

Syntax:

#tt.supportedDataTypes<
  ::mlir::tt::DataType   # value
>

Enum cases:

• f32 (Float32)
• f16 (Float16)
• bf16 (BFloat16)
• bfp_f8 (BFP_Float8)
• bfp_bf8 (BFP_BFloat8)
• bfp_f4 (BFP_Float4)
• bfp_bf4 (BFP_BFloat4)
• bfp_f2 (BFP_Float2)
• bfp_bf2 (BFP_BFloat2)
• u32 (UInt32)
• u16 (UInt16)
• u8 (UInt8)

Parameters:

DeviceAttr

Device attribute in TT dialect.

Syntax:

#tt.device<
  ::mlir::tt::GridAttr,   # workerGrid
  AffineMap,   # l1Map
  AffineMap,   # dramMap
  ::llvm::ArrayRef<int64_t>,   # meshShape
  ::llvm::ArrayRef<unsigned>   # chipIds
>

Describes the physical layout of a device in the system and is made up of a few components:

• A grid attribute that describes the device's compute grid shape. It not only describes the shape of the compute grid, but also carries an affine map that describes how the logical grid maps to the physical grid.
• Two affine maps that describe how a tensor layout's linear attribute maps to the L1 and DRAM memory spaces.
• A mesh shape that describes the virtual layout of the chips with respect to each other. Note that in a multi-chip system, this grid encapsulates the entire system's grid shape, e.g. 8x16 grid could be made up of a 1x2 mesh of chips side-by-side. The mesh attribute configures how the above grid/map attributes are created such that they implement this mesh topology.
• An array of chip ids that this device is made up of. This array's length must match the volume of the mesh shape and should be interpreted in row-major order.

Parameters:

GridAttr

TT grid attribute

Syntax:

#tt.grid<
  ::llvm::ArrayRef<int64_t>,   # shape
  AffineMap   # mapping
>

TT grid attribute

Parameters:

IteratorTypeAttr

TT IteratorType

Syntax:

#tt.iterator_type<
  ::mlir::tt::IteratorType   # value
>

Enum cases:

• parallel (Parallel)
• systolic (Systolic)
• broadcast (Broadcast)
• reduction (Reduction)

Parameters:

MemorySpaceAttr

TT MemorySpace

Syntax:

#tt.memory_space<
  ::mlir::tt::MemorySpace   # value
>

Enum cases:

• system (System)
• mmio (SystemMMIO)
• dram (DeviceDRAM)
• l1 (DeviceL1)

Parameters:

MeshShardDirectionAttr

TT MeshShardDirection

Syntax:

#tt.shard_direction<
  ::mlir::tt::MeshShardDirection   # value
>

Enum cases:

• full_to_shard (FullToShard)
• shard_to_full (ShardToFull)

Parameters:

MeshShardTypeAttr

TT MeshShardType

Syntax:

#tt.shard_type<
  ::mlir::tt::MeshShardType   # value
>

Enum cases:

• manual (Manual)
• replicate (Replicate)
• maximal (Maximal)
• devices (Devices)

Parameters:

MetalLayoutAttr

Tensor layout attribute

Syntax:

#tt.metal_layout<
  AffineMap,   # linear
  OOBVal,   # oob_val
  GridAttr,   # grid
  MemRefType,   # memref
  TensorMemoryLayout   # mem_layout
>

The tensor layout attribute captures how tensor data is sharded across a grid of devices, cores, and is laid out in memory.

Some high level goals

• Logical shapes: Keep the original tensor shape and rank intact and agnostic to underlying storage layout. Keeping the logical shapes not only makes some graph transformations vastly simpler, in particular convs, but it makes the lowered IR much easier to read and reason about. The original tensor shapes leave breadcrumbs that make it much easier to map back to the input representation.
• Flexible sharding: Enable flexibility in choosing grid shape, to get better parallelization and avoid resharding. This is particularly important in cases where tensor shapes are not clean powers of two and would otherwise force our hand in choosing non-optimal grid shapes.
• Logical-Physical Isomorphism: Encode this information with just a few attributes to enable derived conversions from logical to physical layout and back.
• Explicit: A single source of truth.
• Enable a direct way to query padded regions.

Please refer to the Tensor Layout Spec for more in depth documentation.

Examples:

tensor<8x300xf32,
  #tt.metal_layout<(d0, d1) -> (d0, d1),
    undef,
    <1x2>,
    memref<8x150xf32, #tt.memory_space<l1>>
  >
>

tensor<8x96x32xf32,
  #tt.metal_layout<(d0, d1, d2) -> (d0 * 96 + d1, d2),
    undef,
    <2x1>,
    memref<384x32xf32, #tt.memory_space<l1>>
  >
>

tensor<8x96x32xf32,
  #tt.metal_layout<(d0, d1, d2) -> (d0 * 96 + d1, d1, d2),
    undef,
    <2x1x2>,
    memref<384x96x16xf32, #tt.memory_space<l1>>
  >
>

tensor<5x3x2x2x7x32x32xf32,
  #tt.metal_layout<
    (d0, d1, d2, d3, d4, d5, d6)
      -> (d0 * 2688 + d1 * 896 + d2 * 448 + d3 * 224 + d4 * 32 + d5, d4, d5, d6),
    undef,
    <3x2x2x2>,
    memref<4480x4x16x16xf32, #tt.memory_space<l1>>
  >
>

Parameters:

OOBValAttr

TT OOBVal

Syntax:

#tt.oob_val<
  ::mlir::tt::OOBVal   # value
>

Enum cases:

• undef (Undef)
• zero (Zero)
• one (One)
• inf (Inf)
• neginf (NegInf)

Parameters:

ReduceTypeAttr

TT Reduce Type

Syntax:

#tt.reduce_type<
  ::mlir::tt::ReduceType   # value
>

Enum cases:

• sum (Sum)
• mean (Mean)
• max (Max)
• min (Min)
• std (Std)
• var (Var)

Parameters:

StreamLayoutAttr

Stream layout attribute in TT dialect

Syntax:

#tt.stream<
  AffineMap,   # affineMap
  StreamMode,   # streamMode
  uint32_t   # numBuffers
>

Describes stream mode- and buffering-aware layout of a memref buffer.

• AffineMap: Provides affine map indexing into the associated data stream.
• StreamMode: Encodes 'Alias' or 'Stream' mode.
• NumBuffers: Number of back buffers used for double buffering, I/O latency hiding, etc (always 1 in 'Alias' stream mode).

Parameters:

StreamModeAttr

TT Stream Mode

Syntax:

#tt.stream_mode<
  ::mlir::tt::StreamMode   # value
>

Describes mode of a data stream:

• Alias: Data stream is a buffer that aliases data directly in L1. No data movement is necessary to co-locate data with kernel calculations.
• Stream: Data stream references remote memory (e.g., in DRAM or another core's L1). Data movement is necessary to get remote data into local buffer(s), possibly via a sequence of steps that use local buffer(s) as temporary destination(s).

Parameters:

SystemDescAttr

TT system_desc attribute

Syntax:

#tt.system_desc<
  ::llvm::ArrayRef<CPUDescAttr>,   # cpuDescs
  ::llvm::ArrayRef<ChipDescAttr>,   # chipDescs
  ::llvm::ArrayRef<unsigned>,   # chipDescIndices
  ::llvm::ArrayRef<ChipCapabilityAttr>,   # chipCapabilities
  ::llvm::ArrayRef<ChipCoordAttr>,   # chipCoords
  ::llvm::ArrayRef<ChipChannelAttr>   # chipChannels
>

TT system_desc attribute

Parameters:

TensorMemoryLayoutAttr

TT TensorMemoryLayout

Syntax:

#tt.tensor_memory_layout<
  ::mlir::tt::TensorMemoryLayout   # value
>

Enum cases:

• none (None)
• interleaved (Interleaved)
• single_bank (SingleBank)
• height_sharded (HeightSharded)
• width_sharded (WidthSharded)
• block_sharded (BlockSharded)

Parameters:

TileSizeAttr

TT tile_size attribute

Syntax:

#tt.tile_size<
  int64_t,   # y
  int64_t   # x
>

TT tile_size attribute containing a supported Tensix tile shape.

Parameters: