1604 lines
62 KiB
TableGen
1604 lines
62 KiB
TableGen
//===- XeGPUOps.td - XeGPU dialect operations definition ----*- tablegen -*-===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_DIALECT_XEGPU_IR_XEGPUOPS_TD
|
|
#define MLIR_DIALECT_XEGPU_IR_XEGPUOPS_TD
|
|
|
|
include "mlir/Dialect/Arith/IR/ArithBase.td"
|
|
include "mlir/Dialect/XeGPU/IR/XeGPUAttrs.td"
|
|
include "mlir/Dialect/XeGPU/IR/XeGPUDialect.td"
|
|
include "mlir/Dialect/XeGPU/IR/XeGPUTypes.td"
|
|
include "mlir/Interfaces/ShapedOpInterfaces.td"
|
|
include "mlir/Interfaces/SideEffectInterfaces.td"
|
|
include "mlir/Interfaces/ViewLikeInterface.td"
|
|
|
|
|
|
// Base class for dialect operations. This operation inherits from the base
|
|
// `Op` class in OpBase.td, and provides:
|
|
// * The parent dialect of the operation.
|
|
// * The mnemonic for the operation, or the name without the dialect prefix.
|
|
// * A list of traits for the operation.
|
|
class XeGPU_Op<string mnemonic, list<Trait> traits = []>:
|
|
Op<XeGPU_Dialect, mnemonic, traits> {
|
|
|
|
code extraBaseClassDeclaration = [{
|
|
void printProperties(::mlir::MLIRContext *ctx,
|
|
::mlir::OpAsmPrinter &p, const Properties &prop,
|
|
::mlir::ArrayRef<::llvm::StringRef> elidedProps) {
|
|
|
|
DictionaryAttr propAttr = dyn_cast_if_present<mlir::DictionaryAttr>(getPropertiesAsAttr(ctx, prop));
|
|
|
|
// filter out the elidedProps from propAttr, and get the resultAttr
|
|
mlir::SmallVector<mlir::NamedAttribute> filteredAttrs;
|
|
if (propAttr) {
|
|
for (auto namedAttr : propAttr.getValue()) {
|
|
if (llvm::is_contained(elidedProps, namedAttr.getName().strref()))
|
|
continue;
|
|
filteredAttrs.push_back(namedAttr);
|
|
}
|
|
}
|
|
|
|
if (!filteredAttrs.empty()) {
|
|
p << "<" << DictionaryAttr::get(ctx, filteredAttrs) << ">";
|
|
}
|
|
}
|
|
|
|
static ::mlir::ParseResult parseProperties(::mlir::OpAsmParser &parser,
|
|
::mlir::OperationState &result) {
|
|
if (mlir::succeeded(parser.parseOptionalLess())) {
|
|
if (parser.parseAttribute(result.propertiesAttr) || parser.parseGreater())
|
|
return failure();
|
|
}
|
|
return success();
|
|
}
|
|
|
|
}];
|
|
}
|
|
|
|
|
|
def XeGPU_CreateNdDescOp: XeGPU_Op<"create_nd_tdesc", [Pure, ViewLikeOpInterface, AttrSizedOperandSegments]> {
|
|
|
|
let summary = "Create nd-tensor descriptor operation";
|
|
let 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):
|
|
```mlir
|
|
%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):
|
|
```mlir
|
|
%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):
|
|
```mlir
|
|
%0 = ... : ui64
|
|
%c1 = arith.constant 1 : index
|
|
%1 = xegpu.create_nd_tdesc %0, shape:[%h, %w], strides:[%w, %c1]: ui64 -> TensorDesc<8x16xf32>
|
|
```
|
|
}];
|
|
|
|
let arguments = (ins
|
|
XeGPU_BaseAddrType: $source,
|
|
Variadic<Index>: $shape,
|
|
Variadic<Index>: $strides,
|
|
OptionalAttr<DenseI64ArrayAttr>: $const_shape,
|
|
OptionalAttr<DenseI64ArrayAttr>: $const_strides
|
|
);
|
|
|
|
let assemblyFormat = [{
|
|
$source ``
|
|
(`,` `shape` `:` custom<DynamicIndexList>($shape, $const_shape)^
|
|
`,` `strides``:` custom<DynamicIndexList>($strides, $const_strides))?
|
|
attr-dict `:` type($source) `->` qualified(type($TensorDesc))
|
|
}];
|
|
|
|
let results = (outs XeGPU_TensorDesc: $TensorDesc);
|
|
|
|
let hasVerifier = 1;
|
|
|
|
let builders = [
|
|
OpBuilder<(ins "Type": $tdesc, "TypedValue<MemRefType>": $source)>,
|
|
|
|
OpBuilder<(ins "Type": $tdesc, "Value ": $source,
|
|
"llvm::ArrayRef<OpFoldResult>": $shape,
|
|
"llvm::ArrayRef<OpFoldResult>": $strides)>
|
|
];
|
|
|
|
let extraClassDeclaration = extraBaseClassDeclaration # [{
|
|
/// Returns the type of the source memref operand.
|
|
Type getSourceType() {
|
|
return getSource().getType();
|
|
}
|
|
|
|
/// Returns the type of the result TensorDesc.
|
|
xegpu::TensorDescType getType() {
|
|
return getTensorDesc().getType();
|
|
}
|
|
|
|
/// Return the element type of the TensorDesc
|
|
Type getElementType() {
|
|
return getType().getElementType();
|
|
}
|
|
|
|
/// Return the shape of the TensorDesc
|
|
llvm::ArrayRef<int64_t> getTensorDescShape() {
|
|
return getType().getShape();
|
|
}
|
|
|
|
SmallVector<OpFoldResult> getMixedSizes() {
|
|
SmallVector<int64_t> statics;
|
|
|
|
/// Get the static sizes/shape, the value passed to const_shape
|
|
/// will overide the value in memref shape.
|
|
if (auto memrefTy = llvm::dyn_cast<MemRefType>(getSourceType()))
|
|
statics = llvm::to_vector(memrefTy.getShape());
|
|
if (auto attr = getConstShapeAttr())
|
|
statics = llvm::to_vector(attr.asArrayRef());
|
|
|
|
return getMixedValues(statics, getShape(), getContext());
|
|
}
|
|
|
|
SmallVector<OpFoldResult> getMixedStrides() {
|
|
SmallVector<int64_t> statics;
|
|
|
|
/// Get the static strides, the value passed to const_strides
|
|
/// will overide the value in memref.
|
|
if (auto memrefTy = llvm::dyn_cast<MemRefType>(getSourceType()))
|
|
statics = memrefTy.getStridesAndOffset().first;
|
|
if (auto attr = getConstStridesAttr())
|
|
statics = llvm::to_vector(attr.asArrayRef());
|
|
|
|
return getMixedValues(statics, getStrides(), getContext());
|
|
}
|
|
|
|
/// Return the number of leading operands before the `offsets`,
|
|
/// `shape` and `strides` operands.
|
|
static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
|
|
|
|
mlir::Value getViewSource() { return getSource(); }
|
|
|
|
unsigned getSourceMemorySpace() {
|
|
auto srcTy = getSourceType();
|
|
if (auto memrefTy = llvm::dyn_cast<mlir::MemRefType>(srcTy)) {
|
|
auto attr = memrefTy.getMemorySpace();
|
|
if (attr) {
|
|
if (auto intAttr = llvm::dyn_cast<mlir::IntegerAttr>(attr)) {
|
|
return static_cast<unsigned>(intAttr.getInt());
|
|
}
|
|
if (auto memSpaceAttr = llvm::dyn_cast<MemorySpaceAttr>(attr))
|
|
return static_cast<unsigned>(memSpaceAttr.getValue());
|
|
}
|
|
}
|
|
// take global as default memory scope.
|
|
return static_cast<unsigned>(MemorySpace::Global);
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getDescLayoutAttr() {
|
|
return dyn_cast_if_present<xegpu::DistributeLayoutAttr>(getType().getLayout());
|
|
}
|
|
|
|
ArrayRef<int64_t> getDataShape() {
|
|
return getTensorDescShape();
|
|
}
|
|
|
|
}];
|
|
}
|
|
|
|
def XeGPU_PrefetchNdOp : XeGPU_Op<"prefetch_nd", [AnchorLayoutInterface]> {
|
|
let summary = "prefetches a n-D block to cache";
|
|
let 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):
|
|
```mlir
|
|
%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):
|
|
```mlir
|
|
// Prefetch 4 independent 8x16 blocks
|
|
xegpu.prefetch_nd %tdesc[0, 0, 0] <{l1_hint = #xegpu.cache_hint<cached>}>
|
|
: !xegpu.tensor_desc<4x8x16xf16>
|
|
```
|
|
|
|
}];
|
|
|
|
let arguments = (ins XeGPU_TensorDesc: $TensorDesc,
|
|
Variadic<Index>: $offsets,
|
|
DenseI64ArrayAttr: $const_offsets,
|
|
OptionalAttr<XeGPU_CacheHintAttr>: $l1_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>: $l2_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>: $l3_hint,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout);
|
|
|
|
let extraClassDeclaration = extraBaseClassDeclaration # [{
|
|
xegpu::TensorDescType getTensorDescType() {
|
|
return getTensorDesc().getType();
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayout().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutAttr(anchorLayout);
|
|
}
|
|
|
|
SmallVector<OpFoldResult> getMixedOffsets() {
|
|
return getMixedValues(getConstOffsets(), getOffsets(), getContext());
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getDescLayoutAttr() {
|
|
return dyn_cast_if_present<xegpu::DistributeLayoutAttr>(getTensorDescType().getLayout());
|
|
}
|
|
|
|
ArrayRef<int64_t> getDataShape() {
|
|
return getTensorDescType().getShape();
|
|
}
|
|
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$TensorDesc ``
|
|
custom<DynamicIndexList>($offsets, $const_offsets)
|
|
prop-dict attr-dict `:` qualified(type($TensorDesc))
|
|
}];
|
|
|
|
let builders = [
|
|
OpBuilder<(ins "Value": $TensorDesc,
|
|
"ArrayRef<OpFoldResult>": $offsets,
|
|
"xegpu::CachePolicyAttr": $l1_hint,
|
|
"xegpu::CachePolicyAttr": $l2_hint,
|
|
"xegpu::CachePolicyAttr": $l3_hint,
|
|
"xegpu::DistributeLayoutAttr": $layout)>
|
|
];
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
|
|
def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [
|
|
AllElementTypesMatch<["value", "TensorDesc"]>, MemoryEffects<[MemRead]>, AnchorLayoutInterface
|
|
]> {
|
|
let summary = "loads a n-D block from memory (represented by TensorDesc)"
|
|
"to registers (represented by vector)";
|
|
let 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):
|
|
```mlir
|
|
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):
|
|
```mlir
|
|
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):
|
|
```mlir
|
|
// 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>
|
|
```
|
|
|
|
}];
|
|
|
|
let arguments = (ins XeGPU_TensorDesc: $TensorDesc,
|
|
Variadic<Index>: $offsets,
|
|
DenseI64ArrayAttr: $const_offsets,
|
|
OptionalAttr<UnitAttr>: $packed,
|
|
OptionalAttr<DenseI64ArrayAttr>: $transpose,
|
|
OptionalAttr<XeGPU_CacheHintAttr>: $l1_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>: $l2_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>: $l3_hint,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout);
|
|
|
|
let results = (outs XeGPU_ValueType: $value);
|
|
|
|
let extraClassDeclaration = extraBaseClassDeclaration # [{
|
|
VectorType getType() {
|
|
return llvm::dyn_cast<VectorType>(getValue().getType());
|
|
}
|
|
|
|
xegpu::TensorDescType getTensorDescType() {
|
|
return getTensorDesc().getType();
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayout().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutAttr(anchorLayout);
|
|
}
|
|
|
|
SmallVector<OpFoldResult> getMixedOffsets() {
|
|
return getMixedValues(getConstOffsets(), getOffsets(), getContext());
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getDescLayoutAttr() {
|
|
return dyn_cast_if_present<xegpu::DistributeLayoutAttr>(getTensorDescType().getLayout());
|
|
}
|
|
|
|
ArrayRef<int64_t> getDataShape() {
|
|
return getTensorDescType().getShape();
|
|
}
|
|
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$TensorDesc ``
|
|
custom<DynamicIndexList>($offsets, $const_offsets)
|
|
prop-dict attr-dict `:` qualified(type($TensorDesc)) `->` type($value)
|
|
}];
|
|
|
|
let builders = [
|
|
OpBuilder<(ins "Type": $value, "Value": $TensorDesc,
|
|
"ArrayRef<OpFoldResult>": $offsets,
|
|
"UnitAttr": $packed, "DenseI64ArrayAttr": $transpose,
|
|
"xegpu::CachePolicyAttr": $l1_hint,
|
|
"xegpu::CachePolicyAttr": $l2_hint,
|
|
"xegpu::CachePolicyAttr": $l3_hint,
|
|
"xegpu::DistributeLayoutAttr": $layout)>
|
|
];
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [
|
|
AllElementTypesMatch<["value", "TensorDesc"]>, MemoryEffects<[MemWrite]>, AnchorLayoutInterface
|
|
]> {
|
|
let summary = "stores a n-D block register region back to memory";
|
|
|
|
let 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):
|
|
```mlir
|
|
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):
|
|
```mlir
|
|
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):
|
|
```mlir
|
|
// 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>
|
|
```
|
|
|
|
}];
|
|
|
|
let arguments = (ins XeGPU_ValueType: $value,
|
|
XeGPU_TensorDesc: $TensorDesc,
|
|
Variadic<Index>: $offsets,
|
|
DenseI64ArrayAttr: $const_offsets,
|
|
OptionalAttr<XeGPU_CacheHintAttr>: $l1_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>: $l2_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>: $l3_hint,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout);
|
|
|
|
let extraClassDeclaration = extraBaseClassDeclaration # [{
|
|
VectorType getValueType() {
|
|
return llvm::dyn_cast<VectorType>(getValue().getType());
|
|
}
|
|
|
|
xegpu::TensorDescType getTensorDescType() {
|
|
return getTensorDesc().getType();
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayout().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutAttr(anchorLayout);
|
|
}
|
|
|
|
SmallVector<OpFoldResult> getMixedOffsets() {
|
|
return getMixedValues(getConstOffsets(), getOffsets(), getContext());
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getDescLayoutAttr() {
|
|
return dyn_cast_if_present<xegpu::DistributeLayoutAttr>(getTensorDescType().getLayout());
|
|
}
|
|
|
|
ArrayRef<int64_t> getDataShape() {
|
|
return getTensorDescType().getShape();
|
|
}
|
|
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$value `,`
|
|
$TensorDesc ``
|
|
custom<DynamicIndexList>($offsets, $const_offsets)
|
|
prop-dict attr-dict `:` type($value) `,` qualified(type($TensorDesc))
|
|
}];
|
|
|
|
let builders = [
|
|
OpBuilder<(ins "Value": $value, "Value": $TensorDesc,
|
|
"ArrayRef<OpFoldResult>": $offsets,
|
|
"xegpu::CachePolicyAttr": $l1_hint,
|
|
"xegpu::CachePolicyAttr": $l2_hint,
|
|
"xegpu::CachePolicyAttr": $l3_hint,
|
|
"xegpu::DistributeLayoutAttr": $layout)>
|
|
];
|
|
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
def XeGPU_PrefetchOp : XeGPU_Op<"prefetch", [AnchorLayoutInterface]> {
|
|
let summary = "prefetches a set of scattered data points to cache";
|
|
|
|
let 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):
|
|
```mlir
|
|
%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).
|
|
```mlir
|
|
%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.
|
|
```mlir
|
|
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.
|
|
```mlir
|
|
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>
|
|
```
|
|
|
|
}];
|
|
|
|
let arguments = (ins XeGPU_GatherScatterBaseAddrType:$source,
|
|
AnyTypeOf<[XeGPU_OffsetType, Index]>:$offsets,
|
|
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
|
|
OptionalAttr<I64Attr>:$offset_align_byte,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout);
|
|
|
|
let extraClassDeclaration = extraBaseClassDeclaration # [{
|
|
Type getSourceType() {
|
|
return getSource().getType();
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayout().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutAttr(anchorLayout);
|
|
}
|
|
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$source `[` $offsets `]`
|
|
prop-dict
|
|
attr-dict `:` type(operands)
|
|
}];
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayoutInterface]> {
|
|
let summary = "load a set of scattered data points from memory.";
|
|
|
|
let 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).
|
|
```mlir
|
|
%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.
|
|
```mlir
|
|
%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>
|
|
```
|
|
|
|
}];
|
|
|
|
let arguments = (ins XeGPU_GatherScatterBaseAddrType:$source,
|
|
AnyTypeOf<[XeGPU_OffsetType, Index]>:$offsets,
|
|
AnyTypeOf<[XeGPU_MaskType, I1]>:$mask, OptionalAttr<I64Attr>:$chunk_size,
|
|
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout);
|
|
let results = (outs XeGPU_ValueOrScalarType:$value);
|
|
|
|
let extraClassDeclaration = extraBaseClassDeclaration # [{
|
|
|
|
Type getSourceType() {
|
|
return getSource().getType();
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayout().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutAttr(anchorLayout);
|
|
}
|
|
|
|
mlir::Type getElementType() {
|
|
auto type = getValue().getType();
|
|
return getElementTypeOrSelf(type);
|
|
}
|
|
|
|
VectorType getValueType() {
|
|
return llvm::dyn_cast<VectorType>(getValue().getType());
|
|
}
|
|
|
|
Type getMaskType() {
|
|
return getMask().getType();
|
|
}
|
|
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$source `[` $offsets `]` `,`
|
|
$mask prop-dict
|
|
attr-dict `:` type(operands) `->` type($value)
|
|
}];
|
|
|
|
let builders = [
|
|
OpBuilder<(ins "Type": $value, "Value": $source,
|
|
"ArrayRef<OpFoldResult>": $offsets, "Value": $mask,
|
|
"IntegerAttr": $chunk_size,
|
|
"xegpu::CachePolicyAttr": $l1_hint,
|
|
"xegpu::CachePolicyAttr": $l2_hint,
|
|
"xegpu::CachePolicyAttr": $l3_hint)>,
|
|
OpBuilder<(ins "Type": $value, "Value": $source,
|
|
"ArrayRef<OpFoldResult>": $offsets, "Value": $mask,
|
|
"IntegerAttr": $chunk_size,
|
|
"xegpu::CachePolicyAttr": $l1_hint,
|
|
"xegpu::CachePolicyAttr": $l2_hint,
|
|
"xegpu::CachePolicyAttr": $l3_hint,
|
|
"xegpu::DistributeLayoutAttr": $layout)>
|
|
];
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorLayoutInterface]> {
|
|
let summary = "store data to scattered memory locations.";
|
|
let 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).
|
|
```mlir
|
|
%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.
|
|
```mlir
|
|
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>
|
|
```
|
|
|
|
}];
|
|
|
|
let arguments = (ins XeGPU_ValueOrScalarType:$value,
|
|
XeGPU_GatherScatterBaseAddrType:$dest,
|
|
AnyTypeOf<[XeGPU_OffsetType, Index]>:$offsets,
|
|
AnyTypeOf<[XeGPU_MaskType, I1]>:$mask, OptionalAttr<I64Attr>:$chunk_size,
|
|
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
|
|
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout);
|
|
|
|
let extraClassDeclaration = extraBaseClassDeclaration#[{
|
|
Type getDestType() {
|
|
return getDest().getType();
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayout().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutAttr(anchorLayout);
|
|
}
|
|
|
|
mlir::Type getElementType() {
|
|
auto type = getValue().getType();
|
|
return getElementTypeOrSelf(type);
|
|
}
|
|
|
|
VectorType getValueType() {
|
|
return llvm::dyn_cast<VectorType>(getValue().getType());
|
|
}
|
|
|
|
Type getMaskType() {
|
|
return getMask().getType();
|
|
}
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$value `,`
|
|
$dest `[` $offsets `]` `,`
|
|
$mask
|
|
prop-dict
|
|
attr-dict `:` type(operands)
|
|
}];
|
|
|
|
let builders = [
|
|
OpBuilder<(ins "Value": $value, "Value": $dest,
|
|
"ArrayRef<OpFoldResult>": $offsets, "Value": $mask,
|
|
"IntegerAttr": $chunk_size,
|
|
"xegpu::CachePolicyAttr": $l1_hint,
|
|
"xegpu::CachePolicyAttr": $l2_hint,
|
|
"xegpu::CachePolicyAttr": $l3_hint)>,
|
|
OpBuilder<(ins "Value": $value, "Value": $dest,
|
|
"ArrayRef<OpFoldResult>": $offsets, "Value": $mask,
|
|
"IntegerAttr": $chunk_size,
|
|
"xegpu::CachePolicyAttr": $l1_hint,
|
|
"xegpu::CachePolicyAttr": $l2_hint,
|
|
"xegpu::CachePolicyAttr": $l3_hint,
|
|
"xegpu::DistributeLayoutAttr": $layout)>
|
|
];
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
def XeGPU_DpasOp : XeGPU_Op<"dpas", [Pure, AllElementTypesMatch<["lhs", "rhs"]>, AnchorLayoutInterface]> {
|
|
let summary = "It performs mma computation";
|
|
|
|
let 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>
|
|
```
|
|
}];
|
|
|
|
let arguments = (ins
|
|
XeGPU_DpasOprType : $lhs,
|
|
XeGPU_DpasOprType : $rhs,
|
|
Optional<XeGPU_DpasResType>: $acc,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout_a,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout_b,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout_cd
|
|
);
|
|
let results = (outs XeGPU_DpasResType: $result);
|
|
|
|
let extraClassDeclaration = [{
|
|
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayoutCd().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutCdAttr(anchorLayout);
|
|
}
|
|
|
|
VectorType getLhsType() {
|
|
return getLhs().getType();
|
|
}
|
|
|
|
VectorType getRhsType() {
|
|
return getRhs().getType();
|
|
}
|
|
|
|
VectorType getAccType() {
|
|
if (getAcc())
|
|
return getAcc().getType();
|
|
return {};
|
|
}
|
|
|
|
VectorType getResultType() {
|
|
return getResult().getType();
|
|
}
|
|
|
|
bool hasAcc() {
|
|
return getAcc() != nullptr;
|
|
}
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$lhs `,` $rhs (`,` $acc^)? attr-dict `:` type($lhs)`,` type($rhs) (`,` type($acc)^)? `->` type($result)
|
|
}];
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
def XeGPU_AtomicRMWOp: XeGPU_Op<"atomic_rmw", [Pure,
|
|
MemoryEffects<[MemRead, MemWrite]>,
|
|
AllElementTypesMatch<["tensorDesc", "value", "result"]>,
|
|
AllShapesMatch<["tensorDesc", "value", "result"]>,
|
|
AnchorLayoutInterface]> {
|
|
let summary = "Atomic read-modify-write operation on the TensorDesc. ";
|
|
|
|
let 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.
|
|
}];
|
|
|
|
let arguments = (ins
|
|
AtomicRMWKindAttr:$kind,
|
|
XeGPU_TensorDesc:$tensorDesc,
|
|
XeGPU_MaskType:$mask,
|
|
XeGPU_ValueType:$value,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout);
|
|
|
|
let extraClassDeclaration = [{
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayout().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutAttr(anchorLayout);
|
|
}
|
|
|
|
}];
|
|
|
|
let results = (outs XeGPU_ValueType:$result);
|
|
|
|
let assemblyFormat = [{
|
|
$kind $tensorDesc `,` $mask `,` $value attr-dict `:`
|
|
qualified(type($tensorDesc)) `,` type($mask) `,` type($value) `->` type($result)
|
|
}];
|
|
}
|
|
|
|
def XeGPU_AllocNbarrierOp: XeGPU_Op<"alloc_nbarrier", []> {
|
|
let summary = "It allocates a set of named barriers.";
|
|
let 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.}];
|
|
let arguments = (ins I64Attr: $nbarrier_num);
|
|
let assemblyFormat = "$nbarrier_num attr-dict";
|
|
}
|
|
|
|
def XeGPU_InitNbarrierOp: XeGPU_Op<"init_nbarrier", []> {
|
|
let summary = "It assigns a named barrier to the current thread.";
|
|
let 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}];
|
|
|
|
let arguments = (ins I8: $nbarrier_id,
|
|
I8: $participant_thread_num);
|
|
let results = (outs XeGPU_Nbarrier: $result);
|
|
let assemblyFormat = [{
|
|
$nbarrier_id `,` $participant_thread_num attr-dict `:`
|
|
type($nbarrier_id) `,` type($participant_thread_num) `->` qualified(type($result))
|
|
}];
|
|
}
|
|
|
|
def XeGPU_NbarrierArriveOp: XeGPU_Op<"nbarrier_arrive", []> {
|
|
let summary = "It signals the arrival at the named barrier.";
|
|
let 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.}];
|
|
|
|
let arguments = (ins XeGPU_Nbarrier: $nbarrier);
|
|
let assemblyFormat = [{ $nbarrier attr-dict `:` qualified(type($nbarrier))}];
|
|
}
|
|
|
|
def XeGPU_NbarrierWaitOp: XeGPU_Op<"nbarrier_wait", []> {
|
|
let summary = "It waits for a named barrier.";
|
|
let 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.}];
|
|
let arguments = (ins XeGPU_Nbarrier: $nbarrier);
|
|
let assemblyFormat = [{ $nbarrier attr-dict `:` qualified(type($nbarrier)) }];
|
|
}
|
|
|
|
def XeGPU_FenceOp: XeGPU_Op<"fence", []> {
|
|
let summary = "It synchronizes memory accesses.";
|
|
let 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.
|
|
}];
|
|
let arguments = (ins XeGPU_MemorySpaceAttr: $memory_kind,
|
|
XeGPU_FenceScopeAttr: $fence_scope);
|
|
let assemblyFormat = [{`memory_kind` `=` $memory_kind `,` `fence_scope` `=` $fence_scope attr-dict}];
|
|
let extraClassDeclaration = extraBaseClassDeclaration;
|
|
}
|
|
|
|
def XeGPU_ConvertLayoutOp: XeGPU_Op<"convert_layout", [AllTypesMatch<["source", "result"]>,
|
|
MemoryEffects<[MemRead, MemWrite]>,
|
|
AnchorLayoutInterface]> {
|
|
let summary = "Convert the layout of the input operand";
|
|
let 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`: The layout attribute describing the current distribution of `source`
|
|
across subgroups and/or lanes.
|
|
- `target_layout`: The layout attribute describing the desired distribution of the result
|
|
across subgroups and/or lanes.
|
|
|
|
Example (Subgroup level):
|
|
```mlir
|
|
%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>
|
|
```
|
|
}];
|
|
let arguments = (ins XeGPU_VectorOrScalarType: $source,
|
|
DistributeLayoutAttr: $input_layout,
|
|
DistributeLayoutAttr: $target_layout);
|
|
let results = (outs XeGPU_VectorOrScalarType: $result);
|
|
let assemblyFormat = [{
|
|
$source prop-dict attr-dict `:` type($source)
|
|
}];
|
|
let extraClassDeclaration = [{
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getTargetLayout();
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setTargetLayoutAttr(anchorLayout);
|
|
}
|
|
|
|
}];
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
class SizeInBits<string name> :
|
|
StrFunc<"llvm::cast<mlir::ShapedType>($" # name # ".getType()).getNumElements()"
|
|
"*llvm::cast<mlir::ShapedType>($" # name # ".getType()).getElementTypeBitWidth()">;
|
|
class AllMemSizesMatch<list<string> names> :
|
|
AllMatchSameOperatorTrait<names, SizeInBits<"_self">.result,
|
|
"size in bits">;
|
|
|
|
def XeGPU_CreateMemDescOp: XeGPU_Op<"create_mem_desc", [Pure,
|
|
AllMemSizesMatch<["source", "mem_desc"]>]> {
|
|
let summary = "Create a memory descriptor.";
|
|
let 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` : 1D or 2D statically shaped memref, representing the raw SLM buffer. The provided memref must be contiguous.
|
|
|
|
Results:
|
|
- `mem_desc` : the memory descriptor (1D or higher).
|
|
|
|
Example (1D):
|
|
```mlir
|
|
%mdesc = xegpu.create_mem_desc %mref
|
|
: memref<128xi8, 3> -> !xegpu.mem_desc<64xf16>
|
|
```
|
|
|
|
Example (2D with layout):
|
|
```mlir
|
|
%mdesc = xegpu.create_mem_desc %mref
|
|
: memref<4096xi8, 3>
|
|
-> !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>
|
|
```
|
|
|
|
}];
|
|
let arguments = (ins AnyTypeOf<[StaticShared1DMemRefOf<[XeGPU_ScalarType]>, StaticShared2DMemRefOf<[XeGPU_ScalarType]>]>:$source);
|
|
let results = (outs XeGPU_MemDesc:$mem_desc);
|
|
let assemblyFormat = "$source prop-dict attr-dict `` `:` type($source) `->` qualified(type($mem_desc))";
|
|
}
|
|
|
|
def XeGPU_LoadMatrixOp: XeGPU_Op<"load_matrix", [MemoryEffects<[MemRead]>,
|
|
AllElementTypesMatch<["mem_desc", "res"]>, AnchorLayoutInterface]> {
|
|
let arguments = (ins XeGPU_MemDesc:$mem_desc,
|
|
Variadic<Index>: $offsets,
|
|
DenseI64ArrayAttr: $const_offsets,
|
|
OptionalAttr<UnitAttr>:$subgroup_block_io,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout
|
|
);
|
|
let results = (outs XeGPU_ValueOrScalarType:$res);
|
|
let assemblyFormat = [{
|
|
$mem_desc `` custom<DynamicIndexList>($offsets, $const_offsets)
|
|
prop-dict attr-dict `` `:` type(operands) `->` type(results)
|
|
}];
|
|
|
|
let 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):
|
|
```mlir
|
|
%1 = xegpu.load_matrix %0[%c0] : !xegpu.mem_desc<64xf16>, index -> vector<16xf16>
|
|
```
|
|
|
|
Example (Workgroup level, 2D):
|
|
```mlir
|
|
%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>
|
|
```
|
|
}];
|
|
|
|
let builders = [
|
|
OpBuilder<(ins "Type":$res, "TypedValue<MemDescType>": $mem_desc,
|
|
"llvm::ArrayRef<OpFoldResult>": $offsets, "DistributeLayoutAttr": $layout)>,
|
|
];
|
|
let extraClassDeclaration = [{
|
|
SmallVector<OpFoldResult> getMixedOffsets() {
|
|
return getMixedValues(getConstOffsets(), getOffsets(), getContext());
|
|
}
|
|
|
|
ArrayRef<int64_t> getDataShape() {
|
|
auto resTy = getRes().getType();
|
|
if (auto vecTy = llvm::dyn_cast<VectorType>(resTy))
|
|
return vecTy.getShape();
|
|
return {};
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayout().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutAttr(anchorLayout);
|
|
}
|
|
|
|
}];
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
def XeGPU_StoreMatrixOp: XeGPU_Op<"store_matrix", [MemoryEffects<[MemWrite]>,
|
|
AllElementTypesMatch<["mem_desc", "data"]>, AnchorLayoutInterface]> {
|
|
let arguments = (ins
|
|
XeGPU_ValueOrScalarType:$data,
|
|
XeGPU_MemDesc:$mem_desc,
|
|
Variadic<Index>: $offsets,
|
|
DenseI64ArrayAttr: $const_offsets,
|
|
OptionalAttr<UnitAttr>:$subgroup_block_io,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout
|
|
);
|
|
let assemblyFormat = [{ $data `,` $mem_desc `` custom<DynamicIndexList>($offsets, $const_offsets)
|
|
prop-dict attr-dict `` `:` type(operands)}];
|
|
let 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):
|
|
```mlir
|
|
xegpu.store_matrix %1, %0[%c0] : vector<16xf16>, !xegpu.mem_desc<64xf16>, index
|
|
```
|
|
|
|
Example (Workgroup level, 2D):
|
|
```mlir
|
|
%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
|
|
```
|
|
}];
|
|
let builders = [
|
|
OpBuilder<(ins "Value" : $data, "TypedValue<MemDescType>": $mem_desc,
|
|
"llvm::ArrayRef<OpFoldResult>": $offsets, "DistributeLayoutAttr": $layout)>,
|
|
];
|
|
let extraClassDeclaration = [{
|
|
SmallVector<OpFoldResult> getMixedOffsets() {
|
|
return getMixedValues(getConstOffsets(), getOffsets(), getContext());
|
|
}
|
|
|
|
ArrayRef<int64_t> getDataShape() {
|
|
auto DataTy = getData().getType();
|
|
if (auto vecTy = llvm::dyn_cast<VectorType>(DataTy))
|
|
return vecTy.getShape();
|
|
return {};
|
|
}
|
|
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayout().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutAttr(anchorLayout);
|
|
}
|
|
|
|
}];
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
def XeGPU_TruncfOp
|
|
: XeGPU_Op<"truncf", [AllRanksMatch<["source", "result"]>,
|
|
AllShapesMatch<["source", "result"]>, Pure]> {
|
|
let summary = "It performs floating point truncation from higher precision "
|
|
"to lower precision.";
|
|
|
|
let 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>
|
|
```
|
|
|
|
}];
|
|
|
|
let arguments = (ins FixedVectorOfNonZeroRankOf<[XeGPU_FloatType]>:$source);
|
|
let results = (outs FixedVectorOfNonZeroRankOf<[XeGPU_FloatType]>:$result);
|
|
|
|
let extraClassDeclaration = [{
|
|
|
|
Type getSourceType() {
|
|
return getSource().getType();
|
|
}
|
|
|
|
Type getResultType() {
|
|
return getResult().getType();
|
|
}
|
|
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$source attr-dict `:` type(operands) `->` type(results)
|
|
}];
|
|
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments,
|
|
AnchorLayoutInterface]> {
|
|
let summary = "It performs scaled mma computation";
|
|
|
|
let 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.
|
|
|
|
}];
|
|
|
|
let arguments = (ins XeGPU_DpasOprType:$a, XeGPU_DpasOprType:$b,
|
|
Optional<XeGPU_DpasResType>:$acc,
|
|
Optional<AnyTypeOf<[F8E8M0FNU,
|
|
VectorOfRankAndType<[1, 2, 3, 4], [F8E8M0FNU]>]>>:$scale_a,
|
|
Optional<AnyTypeOf<[F8E8M0FNU,
|
|
VectorOfRankAndType<[1, 2, 3, 4], [F8E8M0FNU]>]>>:$scale_b,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout_a,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout_b,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout_cd,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout_a_scale,
|
|
OptionalAttr<DistributeLayoutAttr>:$layout_b_scale);
|
|
let results = (outs XeGPU_DpasResType:$result);
|
|
let extraClassDeclaration = [{
|
|
|
|
xegpu::DistributeLayoutAttr getAnchorLayout() {
|
|
return getLayoutCd().value_or(nullptr);
|
|
}
|
|
|
|
void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) {
|
|
setLayoutCdAttr(anchorLayout);
|
|
}
|
|
|
|
VectorType getAType() {
|
|
return getA().getType();
|
|
}
|
|
|
|
VectorType getBType() {
|
|
return getB().getType();
|
|
}
|
|
|
|
VectorType getAccType() {
|
|
return getAcc().getType();
|
|
}
|
|
|
|
Type getScaleAType() {
|
|
return getScaleA().getType();
|
|
}
|
|
|
|
Type getScaleBType() {
|
|
return getScaleB().getType();
|
|
}
|
|
|
|
VectorType getResultType() {
|
|
return getResult().getType();
|
|
}
|
|
|
|
}];
|
|
let assemblyFormat = [{
|
|
$a `,` $b (`,` $acc^)?
|
|
(`scale_a` `=` $scale_a^)?
|
|
(`scale_b` `=` $scale_b^)?
|
|
attr-dict `:` functional-type(operands, results)
|
|
}];
|
|
let hasVerifier = 1;
|
|
}
|
|
|
|
#endif // MLIR_DIALECT_XEGPU_IR_XEGPUOPS_TD
|