Beaver.MLIR.Dialect.AMX (beaver v0.4.7)
Summary
Functions
amx.tile_load - tile load operation
amx.tile_mulf - tile multiplication operation (floating-point)
amx.tile_muli - tile multiplication operation (integer)
amx.tile_store - tile store operation
amx.tile_zero
Functions
amx.tile_load - tile load operation
Operands
base- Single,AnyMemRef, memref of any type valuesindices- Variadic,Index, variadic of index
Results
res- Single,AnyAMXTile, tile of 32-bit float or 16-bit float or bfloat16 type or 32-bit signless integer or 8-bit signless integer values
Description
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. With memory-effects, each "tileload" operation serves as a compilation hint to use a separate tile register.
Example:
%0 = amx.tile_load %arg0[%c0, %c0] : memref<?x?xi8> into !amx.tile<16x64xi8>
amx.tile_mulf - tile multiplication operation (floating-point)
Operands
lhs- Single,AMXTileF16OrBF16, tile of 16-bit float or bfloat16 type valuesrhs- Single,AMXTileF16OrBF16, tile of 16-bit float or bfloat16 type valuesacc- Single,AMXTileF32, tile of 32-bit float values
Results
res- Single,AMXTileF32, tile of 32-bit float values
Description
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
: !amx.tile<16x32xbf16>, !amx.tile<16x32xbf16>, !amx.tile<16x16xf32>
amx.tile_muli - tile multiplication operation (integer)
Attributes
isZextLhs- Optional,UnitAttr, unit attributeisZextRhs- Optional,UnitAttr, unit attribute
Operands
lhs- Single,AMXTileI8, tile of 8-bit signless integer valuesrhs- Single,AMXTileI8, tile of 8-bit signless integer valuesacc- Single,AMXTileI32, tile of 32-bit signless integer values
Results
res- Single,AMXTileI32, tile of 32-bit signless integer values
Description
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
: !amx.tile<16x64xi8>, !amx.tile<16x64xi8>, !amx.tile<16x16xi32>
amx.tile_store - tile store operation
Operands
base- Single,AnyMemRef, memref of any type valuesindices- Variadic,Index, variadic of indexval- Single,AnyAMXTile, tile of 32-bit float or 16-bit float or bfloat16 type or 32-bit signless integer or 8-bit signless integer values
Description
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>, !amx.tile<16x64xi8>
amx.tile_zero