tt-mlir documentation

d2m.composite_view (tt::d2m::CompositeViewOp)

Piecewise view over multiple input tensors

Represents a piecewise-affine view that aggregates multiple input tensors along a certain dimension.

Traits: AlwaysSpeculatableImplTrait

Interfaces: BufferizableOpInterface, ConditionallySpeculatable, D2M_ViewOpInterface, NoMemoryEffect (MemoryEffectOpInterface), ViewOpInterface

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

d2m.create_global_semaphore (tt::d2m::CreateGlobalSemaphoreOp)

Global semaphore allocation operation (D2M).

Syntax:

operation ::= `d2m.create_global_semaphore` `(` $input `)` attr-dict `:` type($input) `->` type($result)

Create an uninitialized global semaphore with the specified core range.

Interfaces: BufferizableOpInterface

Attributes:

Operands:

Results:

d2m.empty (tt::d2m::EmptyOp)

Empty tensor allocation operation (D2M).

Syntax:

operation ::= `d2m.empty` `(` `)` attr-dict `:` type($result)

Create an uninitialized tensor with the specified shape, element type and encoding.

Optionally, the virtualGridInverseMapping attribute stores the inverse affine map (physical to virtual grid coordinates) and virtualGridForwardMapping stores the forward affine map (virtual to physical grid coordinates).

Interfaces: BufferizableOpInterface, MemoryEffectOpInterface

Attributes:

Results:

d2m.generic (tt::d2m::GenericOp)

Generically dispatch work to a grid of cores (D2M).

Syntax:

operation ::= `d2m.generic` attr-dict `\n`
              ` ` ` ` ` ` ` ` `ins` `(` $inputs (`:` type($inputs)^)? `)` `\n`
              ` ` ` ` ` ` ` ` `outs` `(` $outputs  `:` type($outputs) `)` `\n`
              (` ` ` ` ` ` ` ` `additionalArgs` `(` $additionalArgs `:` type($additionalArgs)^ `)` `\n`)? $regions (`:`  type($results)^ )?

Same semantics as D2M generic; carries regions for compute/datamovement to be consumed by the metal path.

Traits: AttrSizedOperandSegments, NoTerminator

Interfaces: BufferizableOpInterface, DestinationStyleOpInterface, MemoryEffectOpInterface, OpAsmOpInterface

Attributes:

Operands:

Results:

d2m.mesh_shard (tt::d2m::MeshShardOp)

Mesh shard operation (D2M).

Syntax:

operation ::= `d2m.mesh_shard` $input attr-dict `:` type($input) `->` type($result)

MeshShard op shards the inputs (FullToShard) or concatenates the outputs (ShardToFull) for ccl ops.

Traits: AlwaysSpeculatableImplTrait

Interfaces: BufferizableOpInterface, ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

d2m.reset_global_semaphore (tt::d2m::ResetGlobalSemaphoreOp)

Global semaphore reset operation (D2M).

Syntax:

operation ::= `d2m.reset_global_semaphore` `(` $semaphore `)` attr-dict `:` type($semaphore)

Reset global semaphore to a specific value. This operation blocks until the value has been set.

Attributes:

Operands:

d2m.spatial (tt::d2m::SpatialOp)

Execute multiple d2m.generic ops simultaneously on disjoint grid ranges of the device.

Syntax:

operation ::= `d2m.spatial` attr-dict `\n`
              ` ` ` ` ` ` ` ` `ins` `(` $inputs (`:` type($inputs)^)? `)` `\n`
              ` ` ` ` ` ` ` ` `outs` `(` $outputs  `:` type($outputs) `)` ` `  $regions (`:`  type($results)^ )?

Execute multiple d2m.generic ops simultaneously on disjoint parts of the physical device. Each region corresponds to one grid range in grid_ranges and contains exactly one d2m.generic op. Grid ranges must be non-overlapping and contained in the device.

Traits: AttrSizedOperandSegments, NoTerminator

Interfaces: BufferizableOpInterface, DestinationStyleOpInterface, MemoryEffectOpInterface

Attributes:

Operands:

Results:

d2m.to_device (tt::d2m::ToDeviceOp)

Transfer data from host to device.

Syntax:

operation ::= `d2m.to_device` $input `,` $output `layout` `=` $layout attr-dict `:` type($input) `into` type($output) (`->` type($results)^)?

ToDevice operation transfers tensor data from host (system) memory to device memory. This is the explicit host-to-device data movement operation that replaces the overloaded system transfer case of ToLayoutOp.

The layout attribute stores the device layout information needed for the transfer.

#layout = #ttcore.metal_layout<8192x128x1, undef, <1x1>, memref<64x128xf32, #l1_>>
%1 = d2m.to_device %arg0, %0 layout = #layout : tensor<64x128xf32> into tensor<64x128xf32, #layout>

Interfaces: BufferizableOpInterface, MemoryEffectOpInterface

Attributes:

Some high level goals:
  - **Logical shapes**: Store 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.
  - **Collapsed dims**: We may collapse dimensions during transformation, but it
    is important we capture this information such that it is not lost during tensor
    transformation.  The collapsed_intervals field stores the collapses performed
    during conversion from logical_shape to physical tensor shape.
  - **Padding**: store the desired alignments s.t. padding can be simply encoded;
  dim_alignments field represents alignment along each logical dim during collapse.
  - **Memref translation**: ensure we have all necessary info s.t. we can trivially
    lower a tensor into a memref without any intermediate passes.

For a logical tensor of shape [H, W] distributed across a grid [GY, GX], the tensor shape would be:
- Without tiling: [GY, GX, H/GY, W/GX]
- With tiling: [GY, GX, H/GY/TH, W/GX/TW, TH, TW] where TH,TW are tile dimensions

This makes the representation 1:1 with memrefs and eliminates the need for shape conversion passes.

Examples:
```mlir
// Logical 8x300 tensor distributed across 1x2 grid:
// tensor<1x2x8x150xf32, #tt.metal_layout<logical_shape=8x300, ...>>

// Logical 1024x1024 tensor distributed across 2x2 grid with 32x32 tiles:
// tensor<2x2x16x16x!ttcore.tile<32x32xf32>, #tt.metal_layout<logical_shape=1024x1024, ...>>
```

Operands:

Results:

d2m.to_host (tt::d2m::ToHostOp)

Transfer data from device to host.

Syntax:

operation ::= `d2m.to_host` $input `,` $output `layout` `=` $layout attr-dict `:` type($input) `into` type($output) (`->` type($results)^)?

ToHost operation transfers tensor data from device memory to host (system) memory. This is the explicit device-to-host data movement operation that replaces the overloaded system transfer case of ToLayoutOp.

The layout attribute stores the device layout information needed for the transfer.

#layout = #ttcore.metal_layout<8192x128x1, undef, <1x1>, memref<64x128xf32, #l1_>>
%1 = d2m.to_host %arg0, %0 layout = #layout : tensor<64x128xf32, #layout> into tensor<64x128xf32>

Interfaces: BufferizableOpInterface, MemoryEffectOpInterface

Attributes:

Some high level goals:
  - **Logical shapes**: Store 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.
  - **Collapsed dims**: We may collapse dimensions during transformation, but it
    is important we capture this information such that it is not lost during tensor
    transformation.  The collapsed_intervals field stores the collapses performed
    during conversion from logical_shape to physical tensor shape.
  - **Padding**: store the desired alignments s.t. padding can be simply encoded;
  dim_alignments field represents alignment along each logical dim during collapse.
  - **Memref translation**: ensure we have all necessary info s.t. we can trivially
    lower a tensor into a memref without any intermediate passes.

For a logical tensor of shape [H, W] distributed across a grid [GY, GX], the tensor shape would be:
- Without tiling: [GY, GX, H/GY, W/GX]
- With tiling: [GY, GX, H/GY/TH, W/GX/TW, TH, TW] where TH,TW are tile dimensions

This makes the representation 1:1 with memrefs and eliminates the need for shape conversion passes.

Examples:
```mlir
// Logical 8x300 tensor distributed across 1x2 grid:
// tensor<1x2x8x150xf32, #tt.metal_layout<logical_shape=8x300, ...>>

// Logical 1024x1024 tensor distributed across 2x2 grid with 32x32 tiles:
// tensor<2x2x16x16x!ttcore.tile<32x32xf32>, #tt.metal_layout<logical_shape=1024x1024, ...>>
```

Operands:

Results:

d2m.to_layout (tt::d2m::ToLayoutOp)

Layout op.

Syntax:

operation ::= `d2m.to_layout` $input `,` $output `:` type($input) `into` type($output) attr-dict (`->` type($results)^)?

ToLayout operation, transition tensors from one layout to another. Some examples include:

• Transitioning between different memory spaces, e.g. DRAM to L1.
• Transitioning between different data types, e.g. f32 to f16.
• Transitioning between different tile sizes, e.g. 1x16 to 32x32
• Transitioning between different tensor sharding
• Some combination of the above

#layout = #ttcore.metal_layout<8192x128x1, undef, <1x1>, memref<64x128xf32, #system>>
#layout1 = #ttcore.metal_layout<8192x128x1, undef, <1x1>, memref<64x128xf32, #l1_>>
%1 = "d2m.to_layout"(%arg0, %0) : (tensor<64x128xf32, #layout>, tensor<64x128xf32, #layout1>) -> tensor<64x128xf32, #layout1>

Interfaces: BufferizableOpInterface, MemoryEffectOpInterface

Operands:

Results:

d2m.view_layout (tt::d2m::ViewLayoutOp)

View Layout op (D2M subset)

Syntax:

operation ::= `d2m.view_layout` $input `remapping` `=` $remapping attr-dict `:` type($input) `->` type($result)

Create a representational view of a tensor/memref with a different layout. This is a no-op for codegen; consumers are expected to compose layouts.

The remapping attr describes the exact remapping this op represents.

Traits: AlwaysSpeculatableImplTrait

Interfaces: BufferizableOpInterface, ConditionallySpeculatable, D2M_ViewOpInterface, NoMemoryEffect (MemoryEffectOpInterface), OpAsmOpInterface

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results: