Beaver.MLIR.Dialect.XeGPU (beaver v0.4.8)

Copy Markdown

Summary

Functions

Return op name xegpu.alloc_nbarrier as a bitstring.

xegpu.alloc_nbarrier - It allocates a set of named barriers.

Return op name xegpu.atomic_rmw as a bitstring.

xegpu.atomic_rmw - Atomic read-modify-write operation on the TensorDesc.

Return op name xegpu.convert_layout as a bitstring.

xegpu.convert_layout - Convert the layout of the input operand

Return op name xegpu.create_mem_desc as a bitstring.

xegpu.create_mem_desc - Create a memory descriptor.

Return op name xegpu.create_nd_tdesc as a bitstring.

xegpu.create_nd_tdesc - Create nd-tensor descriptor operation

Return op name xegpu.dpas as a bitstring.

xegpu.dpas - It performs mma computation

Return op name xegpu.dpas_mx as a bitstring.

xegpu.dpas_mx - It performs scaled mma computation

Return op name xegpu.fence as a bitstring.

xegpu.fence - It synchronizes memory accesses.

Return op name xegpu.init_nbarrier as a bitstring.

xegpu.init_nbarrier - It assigns a named barrier to the current thread.

Return op name xegpu.lane_shuffle as a bitstring.

xegpu.lane_shuffle - Re-distributes a subgroup's fragments across its lanes.

Return op name xegpu.load as a bitstring.

xegpu.load - load a set of scattered data points from memory.

Return op name xegpu.load_matrix as a bitstring.

xegpu.load_matrix

Return op name xegpu.load_nd as a bitstring.

xegpu.load_nd - loads a n-D block from memory (represented by TensorDesc)to registers (represented by vector)

Return op name xegpu.nbarrier_arrive as a bitstring.

xegpu.nbarrier_arrive - It signals the arrival at the named barrier.

Return op name xegpu.nbarrier_wait as a bitstring.

xegpu.nbarrier_wait - It waits for a named barrier.

Return op name xegpu.prefetch as a bitstring.

xegpu.prefetch - prefetches a set of scattered data points to cache

Return op name xegpu.prefetch_nd as a bitstring.

xegpu.prefetch_nd - prefetches a n-D block to cache

Return op name xegpu.store as a bitstring.

xegpu.store - store data to scattered memory locations.

Return op name xegpu.store_matrix as a bitstring.

xegpu.store_matrix

Return op name xegpu.store_nd as a bitstring.

xegpu.store_nd - stores a n-D block register region back to memory

Return op name xegpu.truncf as a bitstring.

xegpu.truncf - It performs floating point truncation from higher precision to lower precision.

Functions

alloc_nbarrier()

Return op name xegpu.alloc_nbarrier as a bitstring.

alloc_nbarrier(ssa)

xegpu.alloc_nbarrier - It allocates a set of named barriers.

Attributes

  • nbarrier_num - Single, I64Attr, 64-bit signless integer attribute

Description

AllocNbarrier is to create a set of named barriers as specified by nbarrier_num. Named barriers are workgroup level resources,

and are shared by all threads in the workgroup. For example, there are
up to 32 barriers (range 0-31) for each XeCore on PVC. A typical use case
is that a workgroup is partitioned into N subgroups of threads (N <= 32),
and each subgroup coordinating their work with a separate barrier with id
range from 0 to N respectively.

atomic_rmw()

Return op name xegpu.atomic_rmw as a bitstring.

atomic_rmw(ssa)

xegpu.atomic_rmw - Atomic read-modify-write operation on the TensorDesc.

Attributes

  • kind - Single, AtomicRMWKindAttr, allowed 64-bit signless integer cases: 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15
  • layout - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance

Operands

  • tensorDesc - Single, XeGPU_TensorDesc, TensorDesc describing regions of interested data.
  • mask - Single, XeGPU_MaskType, fixed-length vector of 1-bit signless integer values
  • value - Single, XeGPU_ValueType, vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4/5/6/7/8

Results

  • result - Single, XeGPU_ValueType, vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4/5/6/7/8

Description

The xegpu.atomic_rmw operation provides a way to perform a read-modify-write operation on the region described by the TensorDesc free from data races. The kind enumeration specifies the modification to be performed, The mask operand has the same shape with TensorDesc, and is used to enable or disable specific data points of the TensorDesc. The value operand represents the new value to be applied during the modification.

This operation serves as an anchor through which users assign a layout attribute to govern computation distribution.

Arguments:

  • kind: An attribute that specifies the atomic operation to be performed (e.g., add, min, max, exchange, etc.).

  • tensorDesc: A TensorDesc describing the memory region on which the atomic read-modify-write is performed.

  • mask: A predicate mask with the same shape as tensorDesc. Only elements with a true (non-zero) mask value participate in the atomic operation; masked-out elements are not modified.

  • value: The input values used by the atomic operation. It must have the same shape and element type as tensorDesc and result.

  • layout: [optional] An attribute that identifies the operation as an anchor, enabling users to assign a layout that governs distribution at the subgroup and/or lane level. Only valid at workgroup and subgroup levels.

convert_layout()

Return op name xegpu.convert_layout as a bitstring.

convert_layout(ssa)

xegpu.convert_layout - Convert the layout of the input operand

Attributes

  • input_layout - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance
  • target_layout - Single, DistributeLayoutAttr, DistributeLayoutAttr instance

Operands

  • source - Single, XeGPU_VectorOrScalarType, vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type or index values of ranks 1/2/3/4/5/6/7/8 or 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type

Results

  • result - Single, XeGPU_VectorOrScalarType, vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type or index values of ranks 1/2/3/4/5/6/7/8 or 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type

Description

convert_layout redistribute data across subgroups and/or lanes from the input_layout to the target_layout. Both input_layout and target_layout must correspond to the same programming scope, such as workgroup level (wg) or subgroup level (sg) code. This operation is not valid once the IR is lowered to WI level because that is the end result of all distributions.

This operation serves as an anchor through which users assign a layout attribute to govern computation distribution.

Arguments:

  • source: The input vector whose data is to be redistributed. The source and result types must match.
  • input_layout: [optional] The layout attribute describing the current distribution of source across subgroups and/or lanes. When omitted, it defaults to target_layout, i.e. the source is assumed to already be in the target distribution.
  • target_layout: The layout attribute describing the desired distribution of the result across subgroups and/or lanes.

Example (Subgroup level):

    %coop_a = xegpu.convert_layout %a <{
          input_layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>,
          target_layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}>
      : vector<128x128xf16>

create_mem_desc()

Return op name xegpu.create_mem_desc as a bitstring.

create_mem_desc(ssa)

xegpu.create_mem_desc - Create a memory descriptor.

Operands

  • source - Single, anonymous/composite constraint, reside in share memory and statically shaped memref of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values

Results

  • mem_desc - Single, XeGPU_MemDesc, MemDesc describing the data in SLM

Description

Creates a memory descriptor from a shared local memory (SLM) buffer, and xegpu specific memory layout. The resulting memory descriptor has to have the same size as the underlying shared local memory.

Arguments:

  • source : a statically shaped memref of any rank, representing the raw SLM buffer. The provided memref must be contiguous.

Results:

  • mem_desc : the memory descriptor (1D or higher).

Example (1D):

  %mdesc = xegpu.create_mem_desc %mref
    : memref<128xi8, 3> -> !xegpu.mem_desc<64xf16>

Example (2D with layout):

  %mdesc = xegpu.create_mem_desc %mref
    : memref<4096xi8, 3>
      -> !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>

create_nd_tdesc()

Return op name xegpu.create_nd_tdesc as a bitstring.

create_nd_tdesc(ssa)

xegpu.create_nd_tdesc - Create nd-tensor descriptor operation

Attributes

  • const_shape - Optional, DenseI64ArrayAttr, i64 dense array attribute
  • const_strides - Optional, DenseI64ArrayAttr, i64 dense array attribute

Operands

  • source - Single, XeGPU_BaseAddrType, non-0-ranked.memref of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values or 64-bit unsigned integer or 32-bit unsigned integer or 64-bit signless integer or 32-bit signless integer
  • shape - Variadic, Index, variadic of index
  • strides - Variadic, Index, variadic of index

Results

  • TensorDesc - Single, XeGPU_TensorDesc, TensorDesc describing regions of interested data.

Description

The "create_nd_tdesc" operation creates a TensorDescType which represents a sub-view of a 1D/2D memory region inside the one or two innermost dimensions of the source. (It can be extended to support n-D memory region if needed in future). Elements in the subview continuous in each dimension. It encodes the following important information for supporting Intel hardware features:

Arguments:

  • source: an object representing (starting address/pointer of) a memory region. It can be either a memref object, or simply a pointer represented by uint64_t type. For the case of dynamic memrefs or pointer, the shape and layout information of the memory region should be explicitly passed via shape and strides parameters.

  • shape: the shape information of the memory region pointed by the "source". It is typically encoded via the MemRefType of the source, e.g., memref<4096x4096xf16>. But if "source" is simply a pointer represented as uint64_t type, or a memref type without shape information e.g., memref<?x?xf16>, the shape information has to be explicitly passed via the "shape" and "const_shape" arguments.

  • strides: the strides of the memory region pointed by the "source". Similar to shape, it is typically encoded via the MemRefType of the source too. But if "source" is simply a pointer represented as uint64_t type, or a memref type without shape information e.g., memref<?x?xf16>, the strides information has to be explicitly passed via the "strides" and "const_strides" argument.

Results:

  • res: nd tensor descriptor

Example 1 (suppose the tensor shape inferred by the compiler is 8x16):

%0 = memref.alloc() : memref<1024x1024xf32>
%1 = xegpu.create_nd_tdesc %0 : memref<1024x1024xf32> -> TensorDesc<8x16xf32>

Example 2 (suppose the tensor shape inferred by the compiler is 8x16):

%0 = memref.alloc(%h, %w) : memref<?x?xf32>
%c1 = arith.constant 1 : index
%1 = xegpu.create_nd_tdesc %0, shape:[%h, %w], strides:[%w, %c1]: memref<?x?xf32> -> TensorDesc<8x16xf32>

Example 3 (suppose the tensor shape inferred by the compiler is 8x16):

%0 = ... : ui64
%c1 = arith.constant 1 : index
%1 = xegpu.create_nd_tdesc %0, shape:[%h, %w], strides:[%w, %c1]: ui64 -> TensorDesc<8x16xf32>

dpas()

Return op name xegpu.dpas as a bitstring.

dpas(ssa)

xegpu.dpas - It performs mma computation

Attributes

  • layout_a - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance
  • layout_b - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance
  • layout_cd - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance

Operands

  • lhs - Single, XeGPU_DpasOprType, fixed-length vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4
  • rhs - Single, XeGPU_DpasOprType, fixed-length vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4
  • acc - Optional, XeGPU_DpasResType, fixed-length vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4

Results

  • result - Single, XeGPU_DpasResType, fixed-length vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4

Description

DPAS performs matrix multiplication on matrix A of mxk

size, B of `kxn` size, and accumulate on matrix C of `mxn` to the same size
matrix.

The operands can be 2D, 3D, or 4D vectors. When the vectors have more than 2
dimensions, the leading dimensions are treated as **batch dimensions** that
must match across all operands (lhs, rhs, acc, and result). The batch
dimensions represent independent matrix multiplications that are executed in
parallel. For example:
- 2D: `A: vector<8x16xf16>`, `B: vector<16x16xf16>` -> `C: vector<8x16xf32>`
- 3D: `A: vector<4x8x16xf16>`, `B: vector<4x16x16xf16>` -> `C: vector<4x8x16xf32>`
  (4 independent 8x16 matrix multiplications)
- 4D: `A: vector<2x4x8x16xf16>`, `B: vector<2x4x16x16xf16>` -> `C: vector<2x4x8x16xf32>`
  (2x4=8 independent 8x16 matrix multiplications)

The last 2 dimensions are always the core matrix multiplication dimensions
(M, K for lhs; K, N for rhs; M, N for result).

In lane level code, each lane from a subgroup holds a data fragment for A, B, C and the result,
which are represented as 1D vectors. Please refer to [OpenCL Intel extentions]
(https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_matrix_multiply_accumulate.html)
for more details about the fragment distribution.

This operation serves as an anchor through which users assign a layout attribute
to govern computation distribution.

Arguments:

- `lhs`: A vector value representing the left-hand-side matrix tile (A) participating in the
  matrix multiply. Can be 1D, 2D, 3D, or 4D where leading dimensions are batch dimensions.

- `rhs`: A vector value representing the right-hand-side matrix tile (B). Can be 1D, 2D, 3D, or 4D
  where leading dimensions are batch dimensions.

- `acc`: [optional] A vector value representing the accumulator matrix tile (C). When present, the
  result is computed as `lhs * rhs + acc`; otherwise, the accumulator is implicitly assumed to be zero.
  Must have the same batch dimensions as lhs and rhs.

- `layout_a`, `layout_b`, `layout_cd`: [optional] Attributes that identify this
  operation as anchor for operands A, B, and the accumulator/result, enabling users to assign layouts
  that govern distribution at the subgroup and/or lane level. Only valid at workgroup and subgroup
  level.

Example 1 (Workgroup level, 2D):

```mlir
  %d = xegpu.dpas %a, %b, %c <{
      layout_a = #xegpu.layout<sg_layout = [4, 8], sg_data = [16, 128]>,
      layout_b = #xegpu.layout<sg_layout = [4, 8], sg_data = [128, 16]>,
      layout_cd = #xegpu.layout<sg_layout = [4, 8], sg_data = [16, 16]>}
      : vector<64x128xf16>, vector<128x128xf16>, vector<64x128xf32> -> vector<64x128xf32>
```

Example 2 (Lane level, 1D):

```mlir
  %d = xegpu.dpas %a, %b, %c
        :  vector<8xf16>, vector<16xf16>, vector<8xf32> -> vector<8xf32>
```

Example 3 (Workgroup level, 3D with batch):

```mlir
  // 4 independent 8x16 matrix multiplications
  %d = xegpu.dpas %a, %b, %c
      : vector<4x8x16xf16>, vector<4x16x16xf16>, vector<4x8x16xf32> -> vector<4x8x16xf32>
```

dpas_mx()

Return op name xegpu.dpas_mx as a bitstring.

dpas_mx(ssa)

xegpu.dpas_mx - It performs scaled mma computation

Attributes

  • layout_a - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance
  • layout_b - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance
  • layout_cd - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance
  • layout_a_scale - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance
  • layout_b_scale - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance

Operands

  • a - Single, XeGPU_DpasOprType, fixed-length vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4
  • b - Single, XeGPU_DpasOprType, fixed-length vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4
  • acc - Optional, XeGPU_DpasResType, fixed-length vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4
  • scale_a - Optional, anonymous/composite constraint, f8E8M0FNU type or vector of f8E8M0FNU type values of ranks 1/2/3/4
  • scale_b - Optional, anonymous/composite constraint, f8E8M0FNU type or vector of f8E8M0FNU type values of ranks 1/2/3/4

Results

  • result - Single, XeGPU_DpasResType, fixed-length vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4

Description

DPAS MX performs matrix multiplication on matrix A and Matrix B

of low precision data type. A is of size`mxk`
size, B is of size `kxn`, and accumulate on matrix C of size `mxn` to the same size
matrix.

The operands can be 2D, 3D, or 4D vectors. When the vectors have more than 2
dimensions, the leading dimensions are treated as **batch dimensions** that
must match across all operands (a, b, acc, result, scale_a, scale_b). The batch
dimensions represent independent matrix multiplications that are executed in
parallel. For example:
- 2D: `A: vector<8x32xf8E5M2>`, `B: vector<32x16xf8E5M2>` -> `C: vector<8x16xbf16>`
- 3D: `A: vector<4x8x32xf8E5M2>`, `B: vector<4x32x16xf8E5M2>` -> `C: vector<4x8x16xbf16>`
  (4 independent 8x16 scaled matrix multiplications)
- 4D: `A: vector<2x4x8x32xf8E5M2>`, `B: vector<2x4x32x16xf8E5M2>` -> `C: vector<2x4x8x16xbf16>`
  (2x4=8 independent 8x16 scaled matrix multiplications)

The last 2 dimensions are always the core matrix multiplication dimensions
(M, K for a; K, N for b; M, N for result). The scale vectors also follow the
same batch dimension structure.

In lane level code, each lane from a subgroup holds a data fragment for A, B, Acc and the result,
which are represented as 1D vectors.

This operation serves as an anchor through which users assign a layout attribute
to govern computation distribution.

Arguments:

- `a`: A vector value representing the left-hand-side matrix tile (A) participating in the
  matrix multiply. Can be 1D, 2D, 3D, or 4D where leading dimensions are batch dimensions.

- `b`: A vector value representing the right-hand-side matrix tile (B). Can be 1D, 2D, 3D, or 4D
  where leading dimensions are batch dimensions.

- `acc`: A vector value representing the accumulator matrix tile (C). The
  result is computed as `a * b + acc`. Must have the same batch dimensions as a and b.

- `scale_a`: A floating point vector/scalar value used to scale `a` for
  matrix multiplication. When a vector, must have matching batch dimensions.

- `scale_b`: A floating point vector/scalar value used to scale `b` for
  matrix multiplication. When a vector, must have matching batch dimensions.

- `layout_a`, `layout_b`, `layout_cd`: [optional] Attributes that identify this
  operation as anchor for operands A, B, and the accumulator/result, enabling users to assign layouts
  that govern distribution at the subgroup and/or lane level. Only valid at workgroup and subgroup
  level.

fence()

Return op name xegpu.fence as a bitstring.

fence(ssa)

xegpu.fence - It synchronizes memory accesses.

Attributes

  • memory_kind - Single, XeGPU_MemorySpaceAttr, Describe the location of data described by a TensorDesc:
               Global device memory (`Global`) or Shared local memory (`SLM`).
  • fence_scope - Single, XeGPU_FenceScopeAttr, Describes the scope of fence.
                  "workgroup" means that the scope is within each work group.
                  "gpu" means the scope is across work groups within the gpu.

Description

It synchronizes the memory access between

write and following read or write.
1. `Memory_kind` describes the memory kind. "global" means the global memory,
    "slm" means the share local memory.
2. `Fence_scope` describes the scope of fence. "Workgroup" means that the scope would be
    within each workgroup. "GPU" means the scope would be across workgroups within the GPU.

init_nbarrier()

Return op name xegpu.init_nbarrier as a bitstring.

init_nbarrier(ssa)

xegpu.init_nbarrier - It assigns a named barrier to the current thread.

Operands

  • nbarrier_id - Single, I8, 8-bit signless integer
  • participant_thread_num - Single, I8, 8-bit signless integer

Results

  • result - Single, XeGPU_Nbarrier, !xegpu.nbarrier a custom XeGPU type representing a barrier.

Description

InitNbarrierOp assigns the named barrier with the specified

  barrier ID (0~31) to the current thread. Multiple threads may bind to the
  same named barrier, and the `participant_thread_num` specifies the total
  number of threads associated with the nbarrier. It returns an object of
  NbarrierType representing the barrier

lane_shuffle()

Return op name xegpu.lane_shuffle as a bitstring.

lane_shuffle(ssa)

xegpu.lane_shuffle - Re-distributes a subgroup's fragments across its lanes.

Attributes

  • mode - Single, XeGPU_LaneShuffleModeAttr, Describes the direction of a lane shuffle.
                  "pack" gathers a lane's elements into consecutive logical
                  positions. "unpack" scatters them back out, strided by the
                  subgroup size.

Operands

  • source - Single, XeGPU_LaneShuffleType, fixed-length vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1

Results

  • result - Single, XeGPU_LaneShuffleType, fixed-length vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1

Description

xegpu.lane_shuffle is a lane-level operation. The source operand is the fragment (a 1D vector) held by one lane, and the result is the fragment that lane holds afterwards. It re-distributes the elements across the lanes of the subgroup without changing their type, so source and result have the same type.

Let S be the subgroup size and N the number of elements per lane, so the subgroup collectively holds S * N elements. Number these elements by their logical position 0 .. S*N-1. The mode selects which of the two orientations of the S x N (lane, element) index grid the operation maps between:

  • pack: element j of lane i holds logical position j * S + i before the operation and i * N + j after it. A lane's elements end up at consecutive logical positions.
  • unpack: the reverse mapping. A lane's elements end up strided by S.

The two modes are exact inverses, so shuffling one way and then the other yields the original fragments.

For example, consider vector<2xi16> with a subgroup size of 8, denoting a lane's fragment as <a,b>. In pack mode:

lane:       0      1      2      3      4      5      6      7
source:   <0,8>  <1,9>  <2,10> <3,11> <4,12> <5,13> <6,14> <7,15>
result:   <0,1>  <2,3>  <4,5>  <6,7>  <8,9> <10,11> <12,13> <14,15>

In general each lane's result elements come from N different lanes' source fragments -- that is the cross-lane shuffle.

This operation implements the xegpu.convert_layout semantics at the lane level, for the case where the lane_layout is the same but the lane_data is different.

The operation assumes a full subgroup (the subgroup size equals the maximum subgroup size) and that all lanes execute the same dynamic instance of the operation; otherwise the behavior is undefined.

Arguments:

  • source: The 1D vector fragment held by one lane.
  • mode: The direction of the shuffle, either pack or unpack.

load()

Return op name xegpu.load as a bitstring.

load(ssa)

xegpu.load - load a set of scattered data points from memory.

Attributes

  • chunk_size - Optional, I64Attr, 64-bit signless integer attribute
  • l1_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l2_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l3_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • layout - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance
  • contiguity - Optional, I64Attr, 64-bit signless integer attribute

Operands

  • source - Single, XeGPU_GatherScatterBaseAddrType, 1D memref of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values or 64-bit unsigned integer or 32-bit unsigned integer or 64-bit signless integer or 32-bit signless integer
  • offsets - Single, anonymous/composite constraint, fixed-length vector of index values or index
  • mask - Single, anonymous/composite constraint, fixed-length vector of 1-bit signless integer values or 1-bit signless integer

Results

  • value - Single, XeGPU_ValueOrScalarType, vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4/5/6/7/8 or 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type

Description

It (aka. load) load data per each lane. The output describes the data being loaded at the subgroup level, so its size is consistent with the number of lanes in a subgroup. When the chunk size is larger than 2, the output vector is a 2D vector, with dim-0 correspoding to lanes, and dim-1 corresponding to the chunk size loaded by each lane. The mask operand masks out memory access so that it is safe to pass out-of-boundary addresses/offsets as long as they are masked. Each mask element applies to one lane.

In lane level, the result is a 1D vector that represents the data to be loaded by each lane. If size is not 1, size should be equal to the chunk size.

This operation serves as an anchor through which users assign a layout attribute to govern computation distribution.

Arguments:

  • source: represents the memory region to be loaded from, which can be either a 1D memref or pointer (ui64, ui32, i64 or i32).

  • offsets: represents offsets from source. offsets is a vector of index type and vector length is either the subgroup size or 1 at lane level. scalar offset is also valid for lane level.

  • mask: is a vector of i1 type, which is used to mask out the memory access. mask is a vector of size equal to the subgroup size, or 1 at lane level. scalar mask is also valid for lane level.

  • chunk_size: [optional] represents contiguous number of elements to load from per work item.

  • l1_hint, l2_hint, l3_hint: [optional] cache hints for each level of cache.

  • layout: [optional] Describes the expected layout of the tensor_desc operand or the result of load. Only valid at workgroup and subgroup levels.

Results:

  • res: represents loaded data

Example 1 (Subgroup level): A variant accepts memref as base pointer or the source operand could be a raw pointer (ui64, ui32, i64, i32).

   %a = memref.alloc() : memref<1024xf32>
   %offsets = vector.step : vector<16xindex>
   %mask = vector.constant_mask [16]: vector<16xi1>
   %val = xegpu.load %a[%offsets], %mask {l1_hint = #xegpu.cache_hint<cached>,
                          l2_hint = #xegpu.cache_hint<cached>,
                          l3_hint = #xegpu.cache_hint<cached>,
                          layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
     : memref<1024xf32>, vector<16xi1>, vector<16xindex> -> vector<16xf32>

Example 2 (lane level): lane level only accepts the offsets variant. chunk_size can be inferred from result type. In this example, chunk_size is 8.

   %2 = xegpu.load %1[%2], %0 <{l1_hint = #xegpu.cache_hint<cached>,
                            l2_hint = #xegpu.cache_hint<uncached>,
                            l3_hint = #xegpu.cache_hint<uncached>}>
         : memref<128xf32>, vector<1xindex>, vector<1xi1> -> vector<8xf32>

load_matrix()

Return op name xegpu.load_matrix as a bitstring.

load_matrix(ssa)

xegpu.load_matrix

Attributes

  • const_offsets - Single, DenseI64ArrayAttr, i64 dense array attribute
  • subgroup_block_io - Optional, UnitAttr, unit attribute
  • layout - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance

Operands

  • mem_desc - Single, XeGPU_MemDesc, MemDesc describing the data in SLM
  • offsets - Variadic, Index, variadic of index

Results

  • res - Single, XeGPU_ValueOrScalarType, vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4/5/6/7/8 or 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type

Description

This operation loads a block of data from shared local memory (SLM) as specified by the provided mem_desc. Memory descriptors of any rank (1D or higher) are supported.

This operation serves as an anchor through which users assign a layout attribute to govern computation distribution.

Arguments:

  • mem_desc: the memory descriptor identifying the SLM region.
  • offsets: the coordinates within the memory descriptor to read from.
  • subgroup_block_io: [optional] An attribute indicating that the operation can be lowered to a subgroup block load. When this attribute is present, the offsets are subgroup-uniform across all lanes. Only used on subgroup and lane level.
  • layout: [optional] Describes the expected layout of the mem_desc operand as well as the result of load (they are identical). Only valid at workgroup and subgroup levels.

Results:

  • res: the elements loaded from SLM.

Example (1D):

    %1 = xegpu.load_matrix %0[%c0] : !xegpu.mem_desc<64xf16>, index -> vector<16xf16>

Example (Workgroup level, 2D):

    %c0 = arith.constant 0 : index
    %1 = xegpu.load_matrix %0[%c0, %c0] <{
            layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [32, 16]> }>
      : !xegpu.mem_desc<128x128xf16, #xegpu.mem_layout<stride = [1, 128], block = [16, 16]>>
      , index, index -> vector<128x128xf16>

load_nd()

Return op name xegpu.load_nd as a bitstring.

load_nd(ssa)

xegpu.load_nd - loads a n-D block from memory (represented by TensorDesc)to registers (represented by vector)

Attributes

  • const_offsets - Single, DenseI64ArrayAttr, i64 dense array attribute
  • packed - Optional, UnitAttr, unit attribute
  • transpose - Optional, DenseI64ArrayAttr, i64 dense array attribute
  • l1_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l2_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l3_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • layout - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance

Operands

  • TensorDesc - Single, XeGPU_TensorDesc, TensorDesc describing regions of interested data.
  • offsets - Variadic, Index, variadic of index

Results

  • value - Single, XeGPU_ValueType, vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4/5/6/7/8

Description

LoadNdOp essentially mimics the hardware block read instruction to read a block of data from memory to register. It takes a set of optional cache hints for each level of cache, L1, L2 and L3. If hardware does not have a correspoding cache, Corresponding cache hint attribute will be masked.

The tensor descriptor can be 1D, 2D, 3D, or higher dimensional. When the tensor descriptor has more than 2 dimensions, the leading dimensions are treated as batch dimensions that are unrolled to unit dimensions during lowering. At the subgroup level, the load_nd operation executes at 2D granularity to match the 2D block IO hardware support. The result vector has the same shape as the tensor descriptor. The number of offset indices must match the rank of the tensor descriptor.

On Intel GPUs, hardware-supported packing rearranges data elements during the load of the B operand when the element bit-width is less than 32 bits (for example, fp16). The transpose feature reorders data during the load when the element type is fp32 or fp64. These two features are mutually exclusive and shall not be enabled simultaneously. Both features support only 2D blocked tensor_desc.

At lane level, result vector represents the data to be loaded by each lane.

This operation serves as an anchor through which users assign a layout attribute to govern computation distribution.

Arguments:

  • TensorDesc: A tensor descriptor specifying the base nd-region of memory and the tensor tile to be loaded. Can be 1D, 2D, 3D, or higher dimensional where leading dimensions are batch dimensions.

  • offsets: Index values representing per-dimension offsets from the base position encoded in TensorDesc. They are encoded via offsets and const_offsets. The number of offsets must match the tensor descriptor rank.

  • packed: [optional] A unit attribute indicating that packing is applied during the load when supported by the hardware. Only valid at lane level.

  • transpose: [optional] An attribute describing a hardware-supported transpose to be applied during the load. Only valid at Lane level.

  • l1_hint, l2_hint, l3_hint: [optional] Cache-hint attributes indicating the desired behavior at the L1, L2, and L3 cache levels.

  • layout: [optional] Describes the expected layout of the tensor_desc operand as well as the result of the load (they are identical). Only valid at workgroup and subgroup levels.

Example 1 (Workgroup level, 2D):

  xegpu.load_nd %1 {transpose = [1, 0],
                    l1_hint = #xegpu.cache_hint<cached>,
                    l2_hint = #xegpu.cache_hint<uncached>,
                    l3_hint = #xegpu.cache_hint<streaming>,
                    layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [8, 32]>}
          : !xegpu.tensor_desc<32x256xf32> -> vector<32x256xf32>

Example 2 (lane level, 2D):

  xegpu.load_nd %1 {l1_hint = #xegpu.cache_hint<cached>,
                    l2_hint = #xegpu.cache_hint<uncached>}>
    : !xegpu.tensor_desc<8x16xf32> -> vector<8xf32>

Example 3 (3D with batch dimension):

  // Load 4 independent 8x16 blocks
  %result = xegpu.load_nd %tdesc[0, 0, 0] <{l1_hint = #xegpu.cache_hint<cached>}>
    : !xegpu.tensor_desc<4x8x16xf16> -> vector<4x8x16xf16>

nbarrier_arrive()

Return op name xegpu.nbarrier_arrive as a bitstring.

nbarrier_arrive(ssa)

xegpu.nbarrier_arrive - It signals the arrival at the named barrier.

Operands

  • nbarrier - Single, XeGPU_Nbarrier, !xegpu.nbarrier a custom XeGPU type representing a barrier.

Description

NbarrierArriveOp signals the hardware (or other threads)

that the current thread has produced its data for the consumer threads. When
the hardware signalled by `participant_thread_num` threads for the named barrier,
it will notify the threads waiting for the named barrier to continue their work.

nbarrier_wait()

Return op name xegpu.nbarrier_wait as a bitstring.

nbarrier_wait(ssa)

xegpu.nbarrier_wait - It waits for a named barrier.

Operands

  • nbarrier - Single, XeGPU_Nbarrier, !xegpu.nbarrier a custom XeGPU type representing a barrier.

Description

NbarrierWaitOp signals the hardware which named barrier

the current thread is waiting for, such that it can get notified when the
named barrier is completed.

prefetch()

Return op name xegpu.prefetch as a bitstring.

prefetch(ssa)

xegpu.prefetch - prefetches a set of scattered data points to cache

Attributes

  • l1_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l2_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l3_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • offset_align_byte - Optional, I64Attr, 64-bit signless integer attribute
  • layout - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance

Operands

  • source - Single, XeGPU_GatherScatterBaseAddrType, 1D memref of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values or 64-bit unsigned integer or 32-bit unsigned integer or 64-bit signless integer or 32-bit signless integer
  • offsets - Single, anonymous/composite constraint, fixed-length vector of index values or index

Description

It issues instructions to prefetch a set of scattered data points from memory to each level of the cache based on their cache policy. As compared to prefetch_nd, which works on non-scattered TensorDesc, it works on scattered TensorDesc instead.

This operation serves as an anchor through which users assign a layout attribute to govern computation distribution.

Arguments:

  • source: represents the memory region to be loaded from, which can be either a 1D memref or pointer (ui64, ui32, i64 or i32).

  • offsets: represents offsets from source. offsets is a vector of index type and vector length is either the subgroup size or 1 at lane level. scalar offset is also valid for lane level.

  • l1_hint, l2_hint, l3_hint: [optional] cache hints for each level of cache.

  • offset_align_byte: [optional] required if source is a pointer. If source is not a pointer, it is not allowed. Represents the alignment in bytes of each offset in offsets.

  • layout: [optional] Describes the expected layout of the tensor_desc or offsets operand. Only valid at workgroup and subgroup levels.

Example 1 (Workgroup level):

  %a = memref.alloc() : memref<256xf16>
  %offsets = arith.constant dense<[0, 1, ..., 255]> : vector<256xindex>
  xegpu.prefetch %a[%offsets] {l1_hint = #xegpu.cache_hint<cached>,
                         l2_hint = #xegpu.cache_hint<cached>,
                         l3_hint = #xegpu.cache_hint<cached>,
                         layout = #xegpu.layout<sg_layout = [8], sg_data = [32]>
                         }
    : memref<256xf16>, vector<256xindex>

Example 2 (lane level): A variant accepts memref or integer (raw pointer) as base and offsets directly. The source operand could be a raw pointer (ui64, ui32, i64, i32).

  %a = memref.alloc() : memref<1024xf32>
  %0 = arith.constant dense<[0, 16, 32, 64]> : vector<4xindex>
  xegpu.prefetch %a[%0] {l1_hint = #xegpu.cache_hint<cached>,
                         l2_hint = #xegpu.cache_hint<cached>,
                         l3_hint = #xegpu.cache_hint<cached>}
    : memref<1024xf32>, vector<4xindex>

Example 3 (lane level): lane level only accepts the offsets variant.

  xegpu.prefetch %0[%1] {l1_hint = #xegpu.cache_hint<cached>,
                         l2_hint = #xegpu.cache_hint<cached>,
                         l3_hint = #xegpu.cache_hint<cached>}
    : memref<256xf32>, vector<1xindex>

Example 4 (lane level): lane level only accepts the offsets variant.

  xegpu.prefetch %0[%1] {l1_hint = #xegpu.cache_hint<cached>,
                         l2_hint = #xegpu.cache_hint<cached>,
                         l3_hint = #xegpu.cache_hint<cached>,
                         offset_align_byte = 2}
    : i64, vector<1xindex>

prefetch_nd()

Return op name xegpu.prefetch_nd as a bitstring.

prefetch_nd(ssa)

xegpu.prefetch_nd - prefetches a n-D block to cache

Attributes

  • const_offsets - Single, DenseI64ArrayAttr, i64 dense array attribute
  • l1_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l2_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l3_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • layout - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance

Operands

  • TensorDesc - Single, XeGPU_TensorDesc, TensorDesc describing regions of interested data.
  • offsets - Variadic, Index, variadic of index

Description

It issues an instruction to prefetch a block of data from continuous memory regions to each level of the cache based on their cache policy.

The tensor descriptor can be n-dimensional. When the tensor descriptor has more than 2 dimensions, the leading dimensions are treated as batch dimensions that are unrolled to unit dimensions during lowering. At the subgroup level, the prefetch_nd operation executes at 2D granularity to match the 2D block IO hardware support. The number of offset indices must match the rank of the tensor descriptor.

This operation serves as an anchor through which users assign a layout attribute to govern computation distribution.

Arguments:

  • TensorDesc: A tensor descriptor specifying the base nd-region of memory and tensor tile to be prefetched. Can be 1D, 2D, 3D, or higher dimensional where leading dimensions are batch dimensions.

  • offsets: index values representing per-dimension offsets from the base position encoded in TensorDesc. It is encoded via "offsets" and "const_offsets". The number of offsets must match the tensor descriptor rank.

  • l1_hint, l2_hint, l3_hint: [optional] An cache-hint attribute indicating the desired behavior at the L1, L2, and L3 cache levels.

  • layout: [optional] Describes the expected layout of the tensor_desc operand. Only valid at the workgroup and subgroup levels.

Example 1 (Workgroup level, 2D):

  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  xegpu.prefetch_nd %tdesc[%c0, %c1] {l1_hint = #xegpu.cache_hint<cached>,
                            l2_hint = #xegpu.cache_hint<cached>,
                            l3_hint = #xegpu.cache_hint<cached>,
                            layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [8, 32]> }
    : !xegpu.tensor_desc<32x256xf16>

Example 2 (3D with batch dimension):

  // Prefetch 4 independent 8x16 blocks
  xegpu.prefetch_nd %tdesc[0, 0, 0] <{l1_hint = #xegpu.cache_hint<cached>}>
    : !xegpu.tensor_desc<4x8x16xf16>

store()

Return op name xegpu.store as a bitstring.

store(ssa)

xegpu.store - store data to scattered memory locations.

Attributes

  • chunk_size - Optional, I64Attr, 64-bit signless integer attribute
  • l1_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l2_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l3_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • layout - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance
  • contiguity - Optional, I64Attr, 64-bit signless integer attribute

Operands

  • value - Single, XeGPU_ValueOrScalarType, vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4/5/6/7/8 or 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type
  • dest - Single, XeGPU_GatherScatterBaseAddrType, 1D memref of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values or 64-bit unsigned integer or 32-bit unsigned integer or 64-bit signless integer or 32-bit signless integer
  • offsets - Single, anonymous/composite constraint, fixed-length vector of index values or index
  • mask - Single, anonymous/composite constraint, fixed-length vector of 1-bit signless integer values or 1-bit signless integer

Description

It (aka. store) stores data to scattered memory locations. The value is typically a 1D vector. But when the chunk size of the TensorDesc is larger than 1, it will be a 2D vector instead. For the later case, dim-1 of the value correspods to the simd lanes and the dim-0 of the value corresponds to the chunk size stored per lane. So store_scatter has transpose effect, which is similar to load_gather. Therefore, a transpose attribute is introduced on purpose, making sure users are aware of this implicit transformation.

In lane level, the result is a 1D vector that represents the data to be stored by each lane. If size is not 1, size should be equal to the chunk size.

This operation serves as an anchor through which users assign a layout attribute to govern computation distribution.

Arguments:

  • value: represents the data to be stored.

  • dest: represents the memory region to be stored to, which can be either a 1D memref or pointer (ui64, ui32, i64 or i32).

  • offsets: represents offsets from dest. offsets is a vector of index type and vector length is either the subgroup size or 1 at lane level. scalar offset is also valid for lane level.

  • mask: is a vector of i1 type, which is used to mask out the memory access. mask is a vector of size equal to the subgroup size, or 1 at lane level. scalar mask is also valid for lane level.

  • chunk_size: [optional] represents contiguous number of elements to store to per work item.

  • l1_hint, l2_hint, l3_hint: [optional] cache hints for each level of cache.

  • layout: [optional] Describes the expected layout of the tensor_desc operand or the value to be stored. Only valid at workgroup and subgroup levels.

Example 1 (Subgroup level): A variant accepts memref as base pointer and an offset. The dest operand could be a raw pointer (uint64_t).

   %a = memref.alloc() : memref<1024xf32>
   %val = arith.constant dense<0.0> : vector<16xf32>
   %offsets = vector.step : vector<16xindex>
   %mask = vector.constant_mask [16]: vector<16xi1>
   xegpu.store %val, %a[%offsets], %mask {l1_hint = #xegpu.cache_hint<cached>,
                          l2_hint = #xegpu.cache_hint<cached>,
                          l3_hint = #xegpu.cache_hint<cached>,
                          layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
     : vector<16xf32>, memref<1024xf32>, vector<16xi1>, vector<16xindex>

Example 2 (Lane level): Lane level IR only accepts the offsets variant. chunk_size can be inferred from value type. In this example, chunk_size is 8.

   xegpu.store %0, %1[%2], %3 <{l1_hint = #xegpu.cache_hint<uncached>,
                            l2_hint = #xegpu.cache_hint<write_back>,
                            l3_hint = #xegpu.cache_hint<write_through>}>
         : vector<8xf32>, memref<256xf32>, vector<1xindex>, vector<1xi1>

store_matrix()

Return op name xegpu.store_matrix as a bitstring.

store_matrix(ssa)

xegpu.store_matrix

Attributes

  • const_offsets - Single, DenseI64ArrayAttr, i64 dense array attribute
  • subgroup_block_io - Optional, UnitAttr, unit attribute
  • layout - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance

Operands

  • data - Single, XeGPU_ValueOrScalarType, vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4/5/6/7/8 or 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type
  • mem_desc - Single, XeGPU_MemDesc, MemDesc describing the data in SLM
  • offsets - Variadic, Index, variadic of index

Description

This operation stores a data fragment into the shared local memory region specified by a mem_desc. Memory descriptors of any rank (1D or higher) are supported.

This operation serves as an anchor through which users assign a layout attribute to govern computation distribution.

Arguments:

  • mem_desc: the memory descriptor specifying the SLM region.
  • offsets: the coordinates within the memory descriptor where the data will be written.
  • data: the values to be stored.
  • subgroup_block_io: [optional] An attribute indicating that the operation can be lowered to a subgroup block store. When this attribute is present, the offsets are subgroup-uniform across all lanes. Only used on subgroup and lane level.
  • layout: [optional] Describes the expected layout of the mem_desc operand as well as the value to be stored (they are identical). Only valid at workgroup and subgroup levels.

Example (1D):

    xegpu.store_matrix %1, %0[%c0] : vector<16xf16>, !xegpu.mem_desc<64xf16>, index

Example (Workgroup level, 2D):

    %c0 = arith.constant 0 : index
    xegpu.store_matrix %1, %0[%c0, %c0] <{
            layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [32, 16]> }>
      : vector<128x128xf16>, !xegpu.mem_desc<128x128xf16>>, index, index

store_nd()

Return op name xegpu.store_nd as a bitstring.

store_nd(ssa)

xegpu.store_nd - stores a n-D block register region back to memory

Attributes

  • const_offsets - Single, DenseI64ArrayAttr, i64 dense array attribute
  • l1_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l2_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • l3_hint - Optional, XeGPU_CacheHintAttr, Describe the cache settings for prefetch/load/store operators
  • layout - Optional, DistributeLayoutAttr, DistributeLayoutAttr instance

Operands

  • value - Single, XeGPU_ValueType, vector of 1-bit signless integer or 4-bit signless integer or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 1-bit signed integer or 8-bit signed integer or 16-bit signed integer or 32-bit signed integer or 64-bit signed integer or 1-bit unsigned integer or 8-bit unsigned integer or 16-bit unsigned integer or 32-bit unsigned integer or 64-bit unsigned integer or f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values of ranks 1/2/3/4/5/6/7/8
  • TensorDesc - Single, XeGPU_TensorDesc, TensorDesc describing regions of interested data.
  • offsets - Variadic, Index, variadic of index

Description

StoreNdOp essentially mimics the hardware block write instruction io write a block of data from register into the memory region as described by the TensorDesc. It takes a set of optional cache hints for each level of cache, L1, L2 and L3. If hardware does not have a correspoding cache, Corresponding cache hint attribute will be masked.

The tensor descriptor can be 1D, 2D, 3D, or higher dimensional. When the tensor descriptor has more than 2 dimensions, the leading dimensions are treated as batch dimensions that are unrolled to unit dimensions during lowering. At the subgroup level, the store_nd operation executes at 2D granularity to match the 2D block IO hardware support. The value vector must have the same shape as the tensor descriptor. The number of offset indices must match the rank of the tensor descriptor.

At lane level, the input vector represents the data to be stored by each lane.

This operation serves as an anchor through which users assign a layout attribute to govern computation distribution.

Arguments:

  • value: A vector value representing the tensor tile to be stored. Can be 1D, 2D, 3D, or higher dimensional where leading dimensions are batch dimensions.

  • TensorDesc: A tensor descriptor specifying the base nd-region of memory and the tensor tile to be stored. Can be 1D, 2D, 3D, or higher dimensional where leading dimensions are batch dimensions.

  • offsets: Index values representing per-dimension offsets from the base position encoded in TensorDesc. They are encoded via offsets and const_offsets. The number of offsets must match the tensor descriptor rank.

  • l1_hint, l2_hint, l3_hint: [optional] Cache-hint attributes indicating the desired behavior at the L1, L2, and L3 cache levels.

  • layout: [optional] Describes the expected layout of the tensor_desc operand as well as the value to be stored (they are identical). Only valid at workgroup and subgroup levels.

Example 1 (Workgroup level, 2D):

  xegpu.store_nd %3, %2 {l1_hint = #xegpu.cache_hint<uncached>,
                         l2_hint = #xegpu.cache_hint<write_back>,
                         l3_hint = #xegpu.cache_hint<write_through>,
                         layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [8, 32]>}
                         : vector<32x256xf16>, !xegpu.tensor_desc<32x256xf16>

Example 2 (lane level, 2D):

  xegpu.store_nd %3, %2 {l1_hint = #xegpu.cache_hint<uncached>,
                         l2_hint = #xegpu.cache_hint<write_back>,
                         l3_hint = #xegpu.cache_hint<write_through>}
                         : vector<8xf16>, !xegpu.tensor_desc<8x16xf16>

Example 3 (3D with batch dimension):

  // Store 4 independent 8x16 blocks
  xegpu.store_nd %value, %tdesc[0, 0, 0] <{l1_hint = #xegpu.cache_hint<write_back>}>
    : vector<4x8x16xf16>, !xegpu.tensor_desc<4x8x16xf16>

truncf()

Return op name xegpu.truncf as a bitstring.

truncf(ssa)

xegpu.truncf - It performs floating point truncation from higher precision to lower precision.

Operands

  • source - Single, anonymous/composite constraint, fixed-length vector of f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values

Results

  • result - Single, anonymous/composite constraint, fixed-length vector of f4E2M1FN type or f8E4M3FN type or f8E5M2 type or f8E8M0FNU type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or tf32 type values

Description

The xegpu.truncf operation truncates floating point values from a higher

precision type to a lower precision type.
Converts `f16` and `bf16` to microscaling float types.
Rounding mode defaults to round to nearest even.

Example:
```mlir
  %res = xegpu.truncf %src : vector<16xf16> -> vector<16xf8E5M2>
```