MLIR

Multi-Level IR Compiler Framework

'amx' Dialect

The Intel Advanced Matrix Extensions (AMX) provide a tile matrix multiply unit (TMUL), a tile control register (TILECFG), and eight tile registers TMM0 through TMM7 (TILEDATA).

This AMX dialect provides a bridge between MLIR concepts such as vectors and memrefs and the lower level LLVM IR support of AMX. The dialect is split into user-facing AMX ops (AMX_Op) and backend-facing intrinsic ops (AMX_IntrOp).

Note that since configuration changes (implicit at dialect level) are costly, it is highly recommended to use the AMX dialect on same-shaped vectors, at least within a single method.

For details, see the Intel documentation: https://software.intel.com/content/www/us/en/develop/articles/intel-sdm.html

Operation definition 

amx.tdpbf16ps (::mlir::amx::x86_amx_tdpbf16ps) 

Operands: 

OperandDescription
«unnamed»integer
«unnamed»integer
«unnamed»integer
«unnamed»LLVM dialect-compatible type
«unnamed»LLVM dialect-compatible type
«unnamed»LLVM dialect-compatible type

Results: 

ResultDescription
resLLVM dialect-compatible type

amx.tdpbssd (::mlir::amx::x86_amx_tdpbssd) 

Operands: 

OperandDescription
«unnamed»integer
«unnamed»integer
«unnamed»integer
«unnamed»LLVM dialect-compatible type
«unnamed»LLVM dialect-compatible type
«unnamed»LLVM dialect-compatible type

Results: 

ResultDescription
resLLVM dialect-compatible type

amx.tdpbsud (::mlir::amx::x86_amx_tdpbsud) 

Operands: 

OperandDescription
«unnamed»integer
«unnamed»integer
«unnamed»integer
«unnamed»LLVM dialect-compatible type
«unnamed»LLVM dialect-compatible type
«unnamed»LLVM dialect-compatible type

Results: 

ResultDescription
resLLVM dialect-compatible type

amx.tdpbusd (::mlir::amx::x86_amx_tdpbusd) 

Operands: 

OperandDescription
«unnamed»integer
«unnamed»integer
«unnamed»integer
«unnamed»LLVM dialect-compatible type
«unnamed»LLVM dialect-compatible type
«unnamed»LLVM dialect-compatible type

Results: 

ResultDescription
resLLVM dialect-compatible type

amx.tdpbuud (::mlir::amx::x86_amx_tdpbuud) 

Operands: 

OperandDescription
«unnamed»integer
«unnamed»integer
«unnamed»integer
«unnamed»LLVM dialect-compatible type
«unnamed»LLVM dialect-compatible type
«unnamed»LLVM dialect-compatible type

Results: 

ResultDescription
resLLVM dialect-compatible type

amx.tileloadd64 (::mlir::amx::x86_amx_tileloadd64) 

Operands: 

OperandDescription
«unnamed»integer
«unnamed»integer
«unnamed»LLVM pointer type
«unnamed»integer

Results: 

ResultDescription
resLLVM dialect-compatible type

amx.tilestored64 (::mlir::amx::x86_amx_tilestored64) 

Operands: 

OperandDescription
«unnamed»integer
«unnamed»integer
«unnamed»LLVM pointer type
«unnamed»integer
«unnamed»LLVM dialect-compatible type

amx.tilezero (::mlir::amx::x86_amx_tilezero) 

Operands: 

OperandDescription
«unnamed»integer
«unnamed»integer

Results: 

ResultDescription
resLLVM dialect-compatible type

amx.tile_load (::mlir::amx::TileLoadOp) 

tile load operation

Syntax:

operation ::= `amx.tile_load` $base `[` $indices `]` attr-dict `:` type($base) `into` type($res)

Loads a tile from memory defined by a base and indices, with the shape defined by the 2-dim vector type of the result. This is eventually lowered into the “tileloadd” instruction with the corresponding tile configuration.

Example:

  %0 = amx.tile_load %arg0[%c0, %c0] : memref<?x?xi8> into vector<16x64xi8>

Interfaces: NoSideEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands: 

OperandDescription
basememref of any type values
indicesindex

Results: 

ResultDescription
resvector of 32-bit float or bfloat16 type or 32-bit signless integer or 8-bit signless integer values of ranks 2

amx.tile_mulf (::mlir::amx::TileMulFOp) 

tile multiplication operation (floating-point)

Syntax:

operation ::= `amx.tile_mulf` $lhs `,` $rhs `,` $acc attr-dict `:` type($lhs) `,` type($rhs) `,` type($acc)

Multiplies a “m x k” tile with a “k x n” tile and accumulates the results into a “m x n” destination tile. Supports “f32 <- bf16 x bf16” (with pairs of “bf16”). The operation is eventually lowered into the “tdpbf16ps” instruction with the corresponding tile configuration.

Example:

  %0 = amx.tile_mulf %a, %b, %c
    : vector<16x32xbf16>, vector<16x32xbf16>, vector<16x16xf32>

Interfaces: NoSideEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands: 

OperandDescription
lhsvector of 32-bit float or bfloat16 type values of ranks 2
rhsvector of 32-bit float or bfloat16 type values of ranks 2
accvector of 32-bit float or bfloat16 type values of ranks 2

Results: 

ResultDescription
resvector of 32-bit float or bfloat16 type values of ranks 2

amx.tile_muli (::mlir::amx::TileMulIOp) 

tile multiplication operation (integer)

Syntax:

operation ::= `amx.tile_muli` $lhs (`zext` $isZextLhs^)? `,` $rhs (`zext` $isZextRhs^)? `,` $acc attr-dict `:` type($lhs) `,` type($rhs) `,` type($acc)

Multiplies a “m x k” tile with a “k x n” tile and accumulates the results into a “m x n” destination tile. Supports all “si32 <- s/ui8 x s/ui8” combinations (4 bytes packed into dwords in the columns of both the source operand tiles; the zero or sign extension is specified with the attributes and default to sign extended). The operation is eventually lowered into one of the “tdpbssd”, “tdpbsud”, “tdpbusd”, or “tdpbuud” instructions with the corresponding tile configuration.

Example:

  %0 = amx.tile_muli %a zext, %b zext, %c
    : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>

Interfaces: NoSideEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes: 

AttributeMLIR TypeDescription
isZextLhs::mlir::UnitAttrunit attribute
isZextRhs::mlir::UnitAttrunit attribute

Operands: 

OperandDescription
lhsvector of 32-bit signless integer or 8-bit signless integer values of ranks 2
rhsvector of 32-bit signless integer or 8-bit signless integer values of ranks 2
accvector of 32-bit signless integer or 8-bit signless integer values of ranks 2

Results: 

ResultDescription
resvector of 32-bit signless integer or 8-bit signless integer values of ranks 2

amx.tile_store (::mlir::amx::TileStoreOp) 

tile store operation

Syntax:

operation ::= `amx.tile_store` $base `[` $indices `]` `,` $val attr-dict `:` type($base) `,` type($val)

Stores a tile to memory defined by a base and indices, with the shape defined by the 2-dim vector type of the value. This is eventually lowered into the “tilestored” instruction with the corresponding tile configuration.

Example:

  amx.tile_store %arg1[%c0, %c0], %0 : memref<?x?xi8>, vector<16x64xi8>

Operands: 

OperandDescription
basememref of any type values
indicesindex
valvector of 32-bit float or bfloat16 type or 32-bit signless integer or 8-bit signless integer values of ranks 2

amx.tile_zero (::mlir::amx::TileZeroOp) 

tile zero operation

Syntax:

operation ::= `amx.tile_zero` attr-dict `:` type($res)

Zeroes the destination tile, with the shape defined by the 2-dim vector type of the result. This is eventually lowered into the “tilezero” instruction with the corresponding tile configuration.

Example:

  %0 = amx.tile_zero : vector<16x16xbf16>

Interfaces: NoSideEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Results: 

ResultDescription
resvector of 32-bit float or bfloat16 type or 32-bit signless integer or 8-bit signless integer values of ranks 2