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

tile_load(ssa)

amx.tile_load - tile load operation

Operands

  • base - Single, AnyMemRef, memref of any type values
  • indices - 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>

tile_mulf(ssa)

amx.tile_mulf - tile multiplication operation (floating-point)

Operands

  • lhs - Single, AMXTileF16OrBF16, tile of 16-bit float or bfloat16 type values
  • rhs - Single, AMXTileF16OrBF16, tile of 16-bit float or bfloat16 type values
  • acc - 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>

tile_muli(ssa)

amx.tile_muli - tile multiplication operation (integer)

Attributes

  • isZextLhs - Optional, UnitAttr, unit attribute
  • isZextRhs - Optional, UnitAttr, unit attribute

Operands

  • lhs - Single, AMXTileI8, tile of 8-bit signless integer values
  • rhs - Single, AMXTileI8, tile of 8-bit signless integer values
  • acc - 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>

tile_store(ssa)

amx.tile_store - tile store operation

Operands

  • base - Single, AnyMemRef, memref of any type values
  • indices - Variadic, Index, variadic of index
  • val - 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>

tile_zero(ssa)

amx.tile_zero