986 lines
45 KiB
TableGen
986 lines
45 KiB
TableGen
//===- XeGPUAttrs.td - XeGPU dialect attributes 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_XEGPUATTRS_TD
|
|
#define MLIR_DIALECT_XEGPU_IR_XEGPUATTRS_TD
|
|
|
|
include "mlir/Dialect/XeGPU/IR/XeGPUDialect.td"
|
|
include "mlir/IR/AttrTypeBase.td"
|
|
include "mlir/IR/EnumAttr.td"
|
|
|
|
class XeGPUAttr<string name, string attrMnemonic, list<Trait> traits = [],
|
|
string baseCppClass = "::mlir::Attribute">
|
|
: AttrDef<XeGPU_Dialect, name, traits, baseCppClass> {
|
|
let mnemonic = attrMnemonic;
|
|
}
|
|
|
|
class XeGPU_TensorDescAttr<string name, string attrMnemonic, list<Trait> traits = [],
|
|
string baseCppClass = "::mlir::Attribute">
|
|
: XeGPUAttr<name, attrMnemonic, traits, baseCppClass> {
|
|
let assemblyFormat = "`<` struct(params) `>`";
|
|
}
|
|
|
|
def XeGPU_BlockTensorDescAttr: XeGPU_TensorDescAttr<"BlockTensorDesc", "block_tdesc_attr"> {
|
|
let summary = [{a composite attribute for `TensorDescType`}];
|
|
let description = [{`BlockTensorDesc` (or `block_tdesc_attr`) is a composite
|
|
attribute defined for `TensorDescType` for describing following
|
|
properties of a `TensorDesc`.
|
|
1. `memory_space`: It describes where the data block described by the
|
|
TensorDesc is located, `Global` device memory or `Shared` local memory.
|
|
It is default to `Global`.
|
|
2. `array_length`: It describes how many horizontally consecutive blocks
|
|
will be loaded by a hardware load instruction. If the TensorDesc shape
|
|
is 8x16, with array_length = 2. The loaded block shape will be actually
|
|
8x32. Its default value is 1.
|
|
3. `boundary_check`: It is used to indicates the hardware whether to do
|
|
out-of-boundary check. The default value is true.
|
|
}];
|
|
|
|
let parameters = (ins
|
|
DefaultValuedParameter<
|
|
"MemorySpaceAttr",
|
|
"MemorySpaceAttr::get($_ctxt, xegpu::MemorySpace::Global)",
|
|
"Data memory location">: $memory_space,
|
|
DefaultValuedParameter<
|
|
"IntegerAttr",
|
|
"IntegerAttr::get(IntegerType::get($_ctxt, 64), 1)",
|
|
"Number of continuous blocks to load">: $array_length,
|
|
DefaultValuedParameter<
|
|
"BoolAttr",
|
|
"BoolAttr::get($_ctxt, true)",
|
|
"Checking the out of boundary access">: $boundary_check
|
|
);
|
|
|
|
let builders = [
|
|
AttrBuilder<(ins
|
|
CArg<"xegpu::MemorySpace", "xegpu::MemorySpace::Global">:$memory_space,
|
|
CArg<"int", "1">:$array_length,
|
|
CArg<"bool", "true">: $boundary_check
|
|
)>
|
|
];
|
|
|
|
let extraClassDeclaration = [{
|
|
// return true if all fields of the BlockTensorDescAttr are set with
|
|
// default values.
|
|
bool hasDefaultsOnly();
|
|
}];
|
|
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// XeGPU Memory Scope Enums.
|
|
//===----------------------------------------------------------------------===//
|
|
def XeGPU_MemorySpaceGlobal: I32EnumAttrCase<"Global", 0, "global">;
|
|
def XeGPU_MemorySpaceShared: I32EnumAttrCase<"SLM", 3, "slm">;
|
|
def XeGPU_MemorySpace: I32EnumAttr<"MemorySpace",
|
|
"The address space of the memory the tensor descritor is created for",
|
|
[XeGPU_MemorySpaceGlobal, XeGPU_MemorySpaceShared]> {
|
|
let genSpecializedAttr = 0;
|
|
let cppNamespace = "::mlir::xegpu";
|
|
}
|
|
|
|
def XeGPU_MemorySpaceAttr:
|
|
EnumAttr<XeGPU_Dialect, XeGPU_MemorySpace, "memory_space"> {
|
|
let summary = [{Describe the location of data described by a `TensorDesc`:
|
|
Global device memory (`Global`) or Shared local memory (`SLM`).}];
|
|
let assemblyFormat = "$value";
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// XeGPU Cache Enums.
|
|
//===----------------------------------------------------------------------===//
|
|
def XeGPU_CachePolicyCached: I32EnumAttrCase<"CACHED", 0, "cached">; // valid for read and write
|
|
def XeGPU_CachePolicyUncached: I32EnumAttrCase<"UNCACHED", 1, "uncached">; // valid for read and write
|
|
def XeGPU_CachePolicyStreaming: I32EnumAttrCase<"STREAMING", 2, "streaming">; // valid for read only
|
|
def XeGPU_CachePolicyInvalid: I32EnumAttrCase<"READ_INVALIDATE", 3, "read_invalidate">; // valid for read only
|
|
def XeGPU_CachePolicyWriteBack: I32EnumAttrCase<"WRITE_BACK", 4, "write_back">; // valid for write only
|
|
def XeGPU_CachePolicyWriteThrough: I32EnumAttrCase<"WRITE_THROUGH", 5, "write_through">; // valid for write only
|
|
|
|
def XeGPU_CachePolicyEnums : I32EnumAttr<"CachePolicy", "Cache policy",
|
|
[XeGPU_CachePolicyCached, XeGPU_CachePolicyUncached,
|
|
XeGPU_CachePolicyStreaming, XeGPU_CachePolicyInvalid,
|
|
XeGPU_CachePolicyWriteBack, XeGPU_CachePolicyWriteThrough]> {
|
|
let genSpecializedAttr = 0;
|
|
let cppNamespace = "::mlir::xegpu";
|
|
}
|
|
|
|
def XeGPU_CacheHintAttr
|
|
: EnumAttr<XeGPU_Dialect, XeGPU_CachePolicyEnums, "cache_hint"> {
|
|
let summary = [{Describe the cache settings for prefetch/load/store operators}];
|
|
let assemblyFormat = "`<` $value `>`";
|
|
}
|
|
|
|
def XeGPU_FenceScopeWorkgroup: I32EnumAttrCase<"Workgroup", 0, "workgroup">;
|
|
def XeGPU_FenceScopeGPU: I32EnumAttrCase<"GPU", 1, "gpu">;
|
|
def XeGPU_FenceScope: I32EnumAttr<"FenceScope",
|
|
"The enumeration for the scope of fence operation.",
|
|
[XeGPU_FenceScopeWorkgroup, XeGPU_FenceScopeGPU]> {
|
|
let genSpecializedAttr = 0;
|
|
let cppNamespace = "::mlir::xegpu";
|
|
}
|
|
|
|
def XeGPU_FenceScopeAttr:
|
|
EnumAttr<XeGPU_Dialect, XeGPU_FenceScope, "fence_scope"> {
|
|
let summary = [{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.}];
|
|
let assemblyFormat = "$value";
|
|
}
|
|
|
|
def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
|
|
let cppNamespace = "::mlir::xegpu";
|
|
let description = [{
|
|
Common trait for all XeGPU layouts.
|
|
}];
|
|
|
|
let methods = [
|
|
InterfaceMethod<"Check the availability of workgroup level layouts",
|
|
"bool",
|
|
"isForWorkgroup">,
|
|
InterfaceMethod<"Check the availability of subgroup level layouts",
|
|
"bool",
|
|
"isForSubgroup">,
|
|
InterfaceMethod<"Check the availability of lane level layouts",
|
|
"bool",
|
|
"isForLane">,
|
|
InterfaceMethod<"Get the rank of attribute",
|
|
"int64_t",
|
|
"getRank">,
|
|
InterfaceMethod<"Get the num of effective subgroups",
|
|
"int64_t",
|
|
"getNumSubgroups">,
|
|
InterfaceMethod<"Get the order of the layout attribute",
|
|
"DenseI32ArrayAttr",
|
|
"getOrder">,
|
|
InterfaceMethod<"Get the effective SgLayout of the layout attribute as integer array",
|
|
"SmallVector<int64_t>",
|
|
"getEffectiveSgLayoutAsInt">,
|
|
InterfaceMethod<"Get the effective SgData of the layout attribute as integer array",
|
|
"SmallVector<int64_t>",
|
|
"getEffectiveSgDataAsInt">,
|
|
InterfaceMethod<"Get the effective InstData of the layout attribute as integer array",
|
|
"SmallVector<int64_t>",
|
|
"getEffectiveInstDataAsInt">,
|
|
InterfaceMethod<"Get the effective LaneLayout of the layout attribute as integer array",
|
|
"SmallVector<int64_t>",
|
|
"getEffectiveLaneLayoutAsInt">,
|
|
InterfaceMethod<"Get the effective LaneData of the layout attribute as integer array",
|
|
"SmallVector<int64_t>",
|
|
"getEffectiveLaneDataAsInt">,
|
|
InterfaceMethod<"Get the effective order of the layout attribute as integer array",
|
|
"SmallVector<int64_t>",
|
|
"getEffectiveOrderAsInt">,
|
|
InterfaceMethod<"Derive a new layout by dropping sgLayout and sgData",
|
|
"xegpu::DistributeLayoutAttr",
|
|
"dropSgLayoutAndData">,
|
|
InterfaceMethod<"Derive a new layout by dropping InstData",
|
|
"xegpu::DistributeLayoutAttr",
|
|
"dropInstData">,
|
|
InterfaceMethod<"Derive a new layout with sg_data, inst_data and lane_data set to 1 for the specified unit dims",
|
|
"xegpu::DistributeLayoutAttr",
|
|
"setUnitDimData",
|
|
/*args=*/(ins "const SmallVector<int64_t>": $unitDims)>,
|
|
InterfaceMethod<"Derive a new layout with sg_lane and lane_layout set to 1 for the specified unit dims",
|
|
"xegpu::DistributeLayoutAttr",
|
|
"setUnitDimLayout",
|
|
/*args=*/(ins "const SmallVector<int64_t>": $unitDims)>,
|
|
InterfaceMethod<[{Delinearizes a linear ID into its multidimensional
|
|
indices based on the effective layout level.}],
|
|
"FailureOr<SmallVector<Value>>",
|
|
"delinearizeId",
|
|
(ins "OpBuilder &": $builder, "Location":$loc, "Value":$linearId)>,
|
|
InterfaceMethod<[{Derive a new layout with sg_data, inst_data and lane_data set to the
|
|
specified values for the given dimension. Passing -1 for any parameter
|
|
preserves its original value.}],
|
|
"xegpu::DistributeLayoutAttr",
|
|
"setDimData",
|
|
(ins "int64_t": $dim,
|
|
"int64_t": $sgData,
|
|
"int64_t": $instData,
|
|
"int64_t": $laneData)>,
|
|
InterfaceMethod<[{Derive a new layout by removing dimensions.
|
|
`dimGroup` specifies a group of dimensions to be removed in the derived layout.}],
|
|
"xegpu::DistributeLayoutAttr",
|
|
"dropDims",
|
|
(ins "SmallVector<int64_t>": $dimGroup)>,
|
|
InterfaceMethod<[{Derive a new layout by collapsing dimensions.
|
|
`dimGroup` specifies a group of adjacent dimensions that are collapsed into
|
|
a single dimension in the derived layout.}],
|
|
"xegpu::DistributeLayoutAttr",
|
|
"collapseDims",
|
|
(ins "SmallVector<int64_t>": $dimGroup)>,
|
|
InterfaceMethod<[{Derive a new layout by expanding a single dimension into
|
|
multiple adjacent dimensions whose extents are given by `targetShape`
|
|
(`product(targetShape) == originalSize` on `dim`). The original layout's
|
|
per-dim values on `dim` are distributed across the new dims:
|
|
`sg_layout` and `lane_layout` spread outer-to-inner (capped per dim by
|
|
`targetShape[i]`); `sg_data`, `lane_data`, and `inst_data` fill innermost-
|
|
first; `inst_data` is seeded from `lane_layout * lane_data` per new dim.
|
|
`order`, when present, is rewritten so the expanded dims appear
|
|
innermost-fastest in the original dim's slot; if the input has no
|
|
`order`, the result has no `order` either (the default FCD-first
|
|
order `[rank-1, ..., 0]` is assumed).
|
|
Example: layout<sg_layout=[2, 16], sg_data=[1, 1]> with
|
|
expandDim(1, [2, 8]) -> layout<sg_layout=[2, 2, 8], sg_data=[1, 1, 1]>
|
|
(no order in, no order out).}],
|
|
"xegpu::DistributeLayoutAttr",
|
|
"expandDim",
|
|
(ins "int64_t": $dim, "ArrayRef<int64_t>": $targetShape)>,
|
|
InterfaceMethod<[{Derive a new layout by trasnposing it using `permutation`.}],
|
|
"xegpu::DistributeLayoutAttr",
|
|
"transposeDims",
|
|
(ins "ArrayRef<int64_t>": $permutation)>,
|
|
InterfaceMethod<[{Generates instructions to compute multidimensional coordinates for dist units
|
|
assigned to a level identified by linearId. The shape parameter
|
|
represents the higher-level problem size. Each level may access
|
|
multiple blocks according to round-robin distribution rules.}],
|
|
"FailureOr<SmallVector<SmallVector<Value>>>",
|
|
"computeDistributedCoords",
|
|
(ins "OpBuilder &": $builder, "Location":$loc, "Value":$linearId, "ArrayRef<int64_t>":$shape)>,
|
|
InterfaceMethod<[{Statically computes multidimensional coordinates for all dist units
|
|
assigned to a compute unit identified by `linearId`. This is the
|
|
compile-time counterpart of `computeDistributedCoords`: it performs
|
|
the same delinearization and round-robin enumeration but operates
|
|
entirely on static integer values. Returns a list of coordinate
|
|
vectors, one per dist unit.}],
|
|
/*retTy=*/"SmallVector<SmallVector<int64_t>>",
|
|
/*methodName=*/"computeStaticDistributedCoords",
|
|
/*args=*/(ins "int64_t":$linearId, "ArrayRef<int64_t>":$shape)>,
|
|
InterfaceMethod<[{Computes the distributed shape for each compute unit by dividing each
|
|
dimension of `shape` by the corresponding layout factor (sg_layout or
|
|
lane_layout).
|
|
The distributed shape represents the per-compute-unit tile. Each
|
|
distribution unit is defined as the combination of layout factors and
|
|
per-unit data (`subshape`, e.g., sg_data or lane_data). When `shape`
|
|
spans multiple distribution units, the distributed shape may contain
|
|
multiple such units.
|
|
For wrap-around dimensions where the division is uneven, the tensor tile
|
|
is broadcast to all subgroups/lanes.}],
|
|
/*retTy=*/"FailureOr<SmallVector<int64_t>>",
|
|
/*methodName=*/"computeDistributedShape",
|
|
/*args=*/(ins "SmallVector<int64_t>":$shape),
|
|
/*methodBody=*/[{
|
|
SmallVector<int64_t> layout;
|
|
SmallVector<int64_t> subShape;
|
|
if ($_self.isForWorkgroup()) {
|
|
layout = $_self.getEffectiveSgLayoutAsInt();
|
|
subShape = $_self.getEffectiveSgDataAsInt();
|
|
} else if ($_self.isForSubgroup()) {
|
|
layout = $_self.getEffectiveLaneLayoutAsInt();
|
|
subShape = $_self.getEffectiveLaneDataAsInt();
|
|
} else {
|
|
return failure();
|
|
}
|
|
// sgdata or lanedata cannot be empty for distributed shape computation
|
|
if (subShape.empty())
|
|
return failure();
|
|
SmallVector<int64_t> distributedShape(shape.size());
|
|
for (auto [i, dim] : llvm::enumerate(shape)) {
|
|
int64_t distriUnit = layout[i]*subShape[i];
|
|
if ((dim % distriUnit) == 0) {
|
|
// Evenly divisible case, divide the dimension by the layout factor.
|
|
distributedShape[i] = dim / layout[i];
|
|
if (distributedShape[i] % subShape[i] != 0)
|
|
return failure();
|
|
} else {
|
|
// wrap around case, the dimension size must be equal to subShape value
|
|
if(dim != subShape[i])
|
|
return failure();
|
|
distributedShape[i] = dim;
|
|
}
|
|
}
|
|
return distributedShape;
|
|
}]>,
|
|
InterfaceMethod<[{Checks if the given shape can be distributed by the layout}],
|
|
/*retTy=*/"bool",
|
|
/*methodName=*/"isDistributable",
|
|
/*args=*/(ins "SmallVector<int64_t>":$shape),
|
|
/*methodBody=*/[{
|
|
DistributeLayoutAttr curLayoutAttr = $_self;
|
|
SmallVector<int64_t> curShape = shape;
|
|
// Phase 1: Distribute across subgroups (sg_layout + sg_data).
|
|
if (curLayoutAttr.isForWorkgroup()) {
|
|
auto maybeSgShape = curLayoutAttr.computeDistributedShape(curShape);
|
|
if (failed(maybeSgShape))
|
|
return false;
|
|
curShape = maybeSgShape.value();
|
|
curLayoutAttr = curLayoutAttr.dropSgLayoutAndData();
|
|
if (!curLayoutAttr)
|
|
return true;
|
|
}
|
|
// Phase 2: Distribute across instruction data (inst_data).
|
|
if (curLayoutAttr.isForSubgroup() && !curLayoutAttr.isForLane()) {
|
|
SmallVector<int64_t> instData = curLayoutAttr.getEffectiveInstDataAsInt();
|
|
for (size_t i = 0; i < curShape.size(); ++i) {
|
|
if (curShape[i] % instData[i] != 0)
|
|
return false;
|
|
}
|
|
// inst_data becomes the new shape for next phase
|
|
curShape = instData;
|
|
curLayoutAttr = curLayoutAttr.dropInstData();
|
|
if (!curLayoutAttr)
|
|
return true;
|
|
}
|
|
// Phase 3: Distribute across lanes (lane_layout + lane_data).
|
|
auto maybeLaneShape = curLayoutAttr.computeDistributedShape(curShape);
|
|
return succeeded(maybeLaneShape);
|
|
}]>,
|
|
InterfaceMethod</*desc=*/[{Check if this layout is a slice of another layout.}],
|
|
/*retTy=*/"bool",
|
|
/*methodName=*/"isSliceOf",
|
|
/*args=*/(ins "const xegpu::DistributeLayoutAttr&": $other)>,
|
|
InterfaceMethod</*desc=*/[{Check if this layout is a transpose of
|
|
the other layout according to given permutation of (0...n-1).}],
|
|
/*retTy=*/"bool",
|
|
/*methodName=*/"isTransposeOf",
|
|
/*args=*/(ins "const xegpu::DistributeLayoutAttr&": $other,
|
|
"ArrayRef<int64_t>": $perm,
|
|
"xegpu::LayoutKind": $kind)>,
|
|
InterfaceMethod</*desc=*/[{Check if this layout is compatible with another layout
|
|
at a specific level of the layout hierarchy regarding a given shape. }],
|
|
/*retTy=*/"bool",
|
|
/*methodName=*/"isCompatibleWith",
|
|
/*args=*/(ins "const xegpu::DistributeLayoutAttr&": $other,
|
|
"SmallVector<int64_t>": $shape,
|
|
"xegpu::LayoutKind": $level)>,
|
|
InterfaceMethod</*desc=*/[{Check if this layout is equal to another layout.
|
|
For LayoutAttr, this compares all fields.
|
|
For SliceAttr, this requires the same parent and same sliced dims.}],
|
|
/*retTy=*/"bool",
|
|
/*methodName=*/"isEqualTo",
|
|
/*args=*/(ins "const xegpu::DistributeLayoutAttr&": $other)>
|
|
];
|
|
}
|
|
|
|
def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout", [DistributeLayoutAttr]> {
|
|
let summary = [{
|
|
Describes the data distribution to subgroups and work-items for a tensor
|
|
specified by the tensor descriptor.
|
|
}];
|
|
let description = [{
|
|
XeGPU operations use `LayoutAttr` to define how data is distributed across subgroups and work-items.
|
|
This attribute is specified in tensor descriptors during tensor description creation. `LayoutAttr`
|
|
includes the following parameters:
|
|
|
|
* `sg_layout`: Specifies the total number of subgroups and their layout within a workgroup.
|
|
It is mandatory for workgroup-level programming. Its presence implies workgroup-level code.
|
|
* `sg_data`: Defines the data size accessed per subgroup. It is optionally used with `sg_layout`
|
|
for workgroup-level programming. When it is left empty, the size accessed per subgroup can be
|
|
derived from the tensor shape and `sg_layout` using the formula:
|
|
`sg_data[i] = tensor_shape[i] / sg_layout[i]`.
|
|
* `inst_data`: Specifies the data size that is processed by an instruction. It is optionally
|
|
used with lane_layout. When it is left empty, the data size per instruction is equivalent to
|
|
the sg_data for workgroup-level programming or equivalent to tensor shape for subgroup-level
|
|
programming.
|
|
* `lane_layout` : Specifies the total number of work-items and their arrangement within a subgroup.
|
|
It is mandatory for subgroup-level programming and optional for workgroup-level programming.
|
|
* `lane_data` : Specifies the shape of the tensor fragment that each lane accesses. It defines a single,
|
|
minimal distribution unit. Processing the entire tensor may require one or more distribution units per
|
|
hardware instruction.
|
|
* `order`: Specifies the dimension order used to linearize n-dimensional sg_layout and lane_layout to
|
|
1-dimensional layout. The first dimension in the order list is the fastest-changing dimension. If it
|
|
is not present, the default value is [1, 0].
|
|
|
|
Examples:
|
|
|
|
1. Subgroup level layout:
|
|
```mlir
|
|
#xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>
|
|
```
|
|
In this example, there are 16 work-items per subgroup, and is organized as
|
|
[[0, 1, 2, .., 7],[8, 9, .., 15]]. The distribution unit is 1x1.
|
|
|
|
2. Subgroup level layout with order:
|
|
```mlir
|
|
#xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1], order = [0, 1]>
|
|
```
|
|
In this example, there are 16 work-items per subgroup, and is organized as
|
|
[[0, 2, 4, ..., 14], [1, 3, 5, ..., 15]]. The distribution unit is 1x1.
|
|
|
|
3. Subgroup level layout with inst_data
|
|
```mlir
|
|
#xegpu.layout<inst_data = [8, 16], lane_layout = [2, 8], lane_data = [2, 2]>
|
|
```
|
|
In this example, the original problem size is partitioned into smaller subproblems of dimensions [8, 16],
|
|
which are then distributed among 16 work-items arranged as [[0, 1, 2, ..., 7], [8, 9, ..., 15]]. Each
|
|
work-item is assigned four 2x2 blocks in a round-robin manner.
|
|
|
|
4. Workgroup level layout:
|
|
```mlir
|
|
#xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16], lane_layout = [2, 8], lane_data = [1, 1]>
|
|
```
|
|
In this example, the layout represents a workgroup distribution. A workgroup consists of 8 subgroups
|
|
arranged as [[0, 1, 2, 3], [4, 5, 6, 7]]. Each subgroup accesses a 16x16 block per instruction, which
|
|
is further distributed to 16 work items which is organized as [[0, 1, 2, .., 7],[8, 9, .., 15]].
|
|
|
|
5. Workgroup level layout with order:
|
|
```mlir
|
|
#xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16], lane_layout = [2, 8], lane_data = [1, 1], order = [0, 1]>
|
|
```
|
|
In this example, the layout represents a workgroup distribution. A workgroup consists of 8 subgroups
|
|
arranged as [[0, 2, 4, 6], [1, 3, 5, 7]]. Each subgroup accesses a 16x16 block per instruction, which
|
|
is further distributed to 16 work items which is organized as [[0, 2, 4, ..., 14], [1, 3, 5, ..., 15]].
|
|
|
|
6. Workgroup level layout with inst_data:
|
|
```mlir
|
|
#xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16], inst_data = [8, 16], lane_layout = [2, 8], lane_data = [1, 1]>
|
|
```
|
|
This example is similar to the previous ones, but the `inst_data` parameter divides `sg_data` into two instructions,
|
|
each processing an 8x16 block. These blocks are further distributed across 16 work-items with a distribution unit of 1x1.
|
|
Unlike the 2x2 distribution unit in example 3, which results in accessing contiguous 2x2 blocks, the 1x1 distribution
|
|
unit may result in non-contiguous access.
|
|
}];
|
|
|
|
let parameters = (ins
|
|
OptionalParameter<"DenseI32ArrayAttr">: $sg_layout,
|
|
OptionalParameter<"DenseI32ArrayAttr">: $sg_data,
|
|
OptionalParameter<"DenseI32ArrayAttr">: $inst_data,
|
|
OptionalParameter<"DenseI32ArrayAttr">: $lane_layout,
|
|
OptionalParameter<"DenseI32ArrayAttr">: $lane_data,
|
|
OptionalParameter<"DenseI32ArrayAttr">: $order
|
|
);
|
|
|
|
let builders = [
|
|
AttrBuilder<(ins "llvm::ArrayRef<int32_t>": $inst_data),
|
|
[{
|
|
auto sg_layout = DenseI32ArrayAttr();
|
|
auto sg_data = DenseI32ArrayAttr();
|
|
auto order = DenseI32ArrayAttr();
|
|
auto lane_layout = DenseI32ArrayAttr();
|
|
auto lane_data = DenseI32ArrayAttr();
|
|
return $_get($_ctxt, sg_layout, sg_data,
|
|
DenseI32ArrayAttr::get($_ctxt, inst_data),
|
|
lane_layout, lane_data, order);
|
|
}]>,
|
|
AttrBuilder<(ins "llvm::ArrayRef<int32_t>": $inst_data,
|
|
"llvm::ArrayRef<int32_t>": $lane_layout,
|
|
"llvm::ArrayRef<int32_t>": $lane_data),
|
|
[{
|
|
auto sg_layout = DenseI32ArrayAttr();
|
|
auto sg_data = DenseI32ArrayAttr();
|
|
auto order = DenseI32ArrayAttr();
|
|
return $_get($_ctxt, sg_layout, sg_data,
|
|
DenseI32ArrayAttr::get($_ctxt, inst_data),
|
|
DenseI32ArrayAttr::get($_ctxt, lane_layout),
|
|
DenseI32ArrayAttr::get($_ctxt, lane_data), order);
|
|
}]>,
|
|
AttrBuilder<(ins "llvm::ArrayRef<int32_t>": $lane_layout,
|
|
"llvm::ArrayRef<int32_t>": $lane_data),
|
|
[{
|
|
auto sg_layout = DenseI32ArrayAttr();
|
|
auto sg_data = DenseI32ArrayAttr();
|
|
auto inst_data = DenseI32ArrayAttr();
|
|
auto order = DenseI32ArrayAttr();
|
|
return $_get($_ctxt, sg_layout, sg_data, inst_data,
|
|
DenseI32ArrayAttr::get($_ctxt, lane_layout),
|
|
DenseI32ArrayAttr::get($_ctxt, lane_data), order);
|
|
}]>,
|
|
AttrBuilder<(ins "DenseI32ArrayAttr": $lane_layout,
|
|
"DenseI32ArrayAttr": $lane_data,
|
|
"DenseI32ArrayAttr": $order),
|
|
[{
|
|
return $_get($_ctxt, /*sg_layout =*/ nullptr, /*sg_data =*/ nullptr,
|
|
/*inst_data =*/ nullptr, lane_layout, lane_data, order);
|
|
}]>
|
|
];
|
|
|
|
let extraClassDeclaration = [{
|
|
bool isForWorkgroup() {
|
|
return getSgLayout() != nullptr;
|
|
}
|
|
|
|
bool isForSubgroup() {
|
|
return !isForWorkgroup();
|
|
}
|
|
|
|
bool isForLane() {
|
|
return !isForWorkgroup() && (getInstData() == nullptr);
|
|
}
|
|
|
|
int64_t getRank() const {
|
|
if (auto attr = getSgLayout())
|
|
return attr.size();
|
|
if (auto attr = getInstData())
|
|
return attr.size();
|
|
if (auto attr = getLaneLayout())
|
|
return attr.size();
|
|
return 0;
|
|
}
|
|
|
|
int64_t getNumSubgroups() const {
|
|
auto sgLayout = getEffectiveSgLayoutAsInt();
|
|
if (!sgLayout.empty())
|
|
return computeProduct(sgLayout);
|
|
return 0;
|
|
}
|
|
|
|
LayoutAttr dropSgLayoutAndData() const{
|
|
// avoid every field of the attribute is nullptr, which may lead to segment fault
|
|
if (!getInstData() && !getLaneLayout())
|
|
return nullptr;
|
|
// Only preserve order if lane_layout remains, since order requires
|
|
// sg_layout or lane_layout to be present.
|
|
auto order = getLaneLayout() ? getOrder() : nullptr;
|
|
return LayoutAttr::get(getContext(), nullptr, nullptr, getInstData(),
|
|
getLaneLayout(), getLaneData(), order);
|
|
}
|
|
|
|
LayoutAttr dropInstData() const{
|
|
// avoid every field of the attribute is nullptr, which may lead to segment fault
|
|
if (!getSgLayout() && !getLaneLayout())
|
|
return nullptr;
|
|
return LayoutAttr::get(getContext(), getSgLayout(), getSgData(), nullptr,
|
|
getLaneLayout(), getLaneData(), getOrder());
|
|
}
|
|
|
|
SmallVector<int64_t> getEffectiveSgLayoutAsInt() const {
|
|
if (DenseI32ArrayAttr layout = getSgLayout())
|
|
return llvm::to_vector_of<int64_t>(layout.asArrayRef());
|
|
return {};
|
|
}
|
|
|
|
SmallVector<int64_t> getEffectiveSgDataAsInt() const {
|
|
if (DenseI32ArrayAttr data = getSgData())
|
|
return llvm::to_vector_of<int64_t>(data.asArrayRef());
|
|
return {};
|
|
}
|
|
|
|
SmallVector<int64_t> getEffectiveInstDataAsInt() const {
|
|
if (DenseI32ArrayAttr inst = getInstData())
|
|
return llvm::to_vector_of<int64_t>(inst.asArrayRef());
|
|
return {};
|
|
}
|
|
|
|
SmallVector<int64_t> getEffectiveLaneLayoutAsInt() const {
|
|
if (DenseI32ArrayAttr layout = getLaneLayout())
|
|
return llvm::to_vector_of<int64_t>(layout.asArrayRef());
|
|
return {};
|
|
}
|
|
|
|
SmallVector<int64_t> getEffectiveLaneDataAsInt() const {
|
|
if (DenseI32ArrayAttr data = getLaneData())
|
|
return llvm::to_vector_of<int64_t>(data.asArrayRef());
|
|
return {};
|
|
}
|
|
|
|
// Returns the order as integers. If order is not set, returns the default
|
|
// value [rank-1, ..., 0] (row-major, last dim is fastest).
|
|
SmallVector<int64_t> getEffectiveOrderAsInt() const {
|
|
if (DenseI32ArrayAttr order = getOrder())
|
|
return llvm::to_vector_of<int64_t>(order.asArrayRef());
|
|
int64_t rank = getRank();
|
|
SmallVector<int64_t> defaultOrder(rank);
|
|
std::iota(defaultOrder.rbegin(), defaultOrder.rend(), 0);
|
|
return defaultOrder;
|
|
}
|
|
|
|
//set the layout for the sepcified unit dims: sg_data, inst_data and lane_data to 1
|
|
DistributeLayoutAttr setUnitDimData(SmallVector<int64_t> unitDims) const;
|
|
|
|
//set the layout for the sepcified unit dims: sg_lane and lane_layout to 1
|
|
DistributeLayoutAttr setUnitDimLayout(SmallVector<int64_t> unitDims) const;
|
|
|
|
// Derive a new layout with sg_data, inst_data and lane_data set to the
|
|
// specified values for the given dimension. Passing -1 for any parameter
|
|
// preserves its original value.
|
|
DistributeLayoutAttr setDimData(int64_t dim, int64_t sgData, int64_t instData, int64_t laneData);
|
|
|
|
// Derive a new layout by removing dimensions.
|
|
// `dimGroup` specifies a group of dimensions to be removed in the derived layout.
|
|
DistributeLayoutAttr dropDims(SmallVector<int64_t> dimGroup);
|
|
|
|
// Derive a new layout by collapsing dimensions.
|
|
// `dimGroup` specifies a group of adjacent dimensions
|
|
// that are collapsed into a single dimension in the derived layout.
|
|
DistributeLayoutAttr collapseDims(SmallVector<int64_t> dimGroup);
|
|
|
|
// Derive a new layout by expanding a single dimension `dim` into
|
|
// multiple adjacent dimensions whose shape is given by `targetShape`.
|
|
// See the interface method documentation for the per-field distribution
|
|
// policy.
|
|
DistributeLayoutAttr expandDim(int64_t dim, ArrayRef<int64_t> targetShape);
|
|
|
|
// Derive a new layout by transposing the layout using `permutation`.
|
|
DistributeLayoutAttr transposeDims(ArrayRef<int64_t> permutation);
|
|
|
|
/// Delinearizes a linear ID into its multidimensional indices
|
|
/// based on the effective level of the layout.
|
|
FailureOr<SmallVector<Value>>
|
|
delinearizeId(OpBuilder &builder, Location loc, Value linearId);
|
|
|
|
/// Generates instructions to compute multidimensional coordinates for dist units
|
|
/// assigned to a level identified by linearId. The shape parameter
|
|
/// represents the higher-level problem size. Each `level` may access
|
|
/// multiple blocks according to round-robin distribution rules.
|
|
FailureOr<SmallVector<SmallVector<Value>>>
|
|
computeDistributedCoords(OpBuilder &builder, Location loc, Value linearId, ArrayRef<int64_t> shape);
|
|
|
|
///Statically computes multidimensional coordinates for all dist units
|
|
///assigned to a compute unit identified by `linearId`. This is the
|
|
///compile-time counterpart of `computeDistributedCoords`.
|
|
SmallVector<SmallVector<int64_t>>
|
|
computeStaticDistributedCoords(int64_t linearId, ArrayRef<int64_t> shape);
|
|
|
|
/// Check if this is slice of some other layout.
|
|
bool isSliceOf(const xegpu::DistributeLayoutAttr &other) { return false; }
|
|
|
|
/// Check if this layout is equal to another layout.
|
|
bool isEqualTo(const xegpu::DistributeLayoutAttr &other);
|
|
|
|
/// Check if this layout is compatible with another layout
|
|
/// at a specific level of the layout hierarchy regarding a given shape.
|
|
bool isCompatibleWith(const xegpu::DistributeLayoutAttr &other,
|
|
SmallVector<int64_t> shape,
|
|
xegpu::LayoutKind level);
|
|
|
|
/// Check if this layout is a transpose of another layout.
|
|
bool isTransposeOf(const xegpu::DistributeLayoutAttr &other, ArrayRef<int64_t> perm, const xegpu::LayoutKind kind);
|
|
}];
|
|
|
|
let assemblyFormat = "`<` struct(params) `>`";
|
|
let genVerifyDecl = 1;
|
|
}
|
|
|
|
|
|
def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
|
|
let summary = [{Describes the data distribution and sharing among subgroups or work-items.}];
|
|
|
|
let description = [{
|
|
Like LayoutAttr, SliceAttr describes data distribution among subgroups or work-items.
|
|
However, whereas LayoutAttr requires the data to have the same rank as the attribute,
|
|
SliceAttr permits the data to have a lower rank. In this case, compute units in the
|
|
specified dimensions (given by `$dims`) share the data, provided that the remaining
|
|
ranks match the data rank. SliceAttr is commonly used by operations such as
|
|
vector.multi_reduction and vector.broadcast.
|
|
|
|
Example:
|
|
```
|
|
#l = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>
|
|
#r = #xegpu.slice<#l, dim = [0]>
|
|
|
|
%exp = math.exp %input {layout_result_0 = #l}: vector<256x128xf32>
|
|
%red = vector.multi_reduction<add>, %exp, %acc [0] {layout_result_0 = #r}: vector<256x128xf32> to vector<128xf32>
|
|
%bcast = vector.broadcast %red {layout_result_0 = #l} : vector<128xf32> to vector<256x128xf32>
|
|
```
|
|
In this example, %red is conceptually divided into 4 vectors of type vector<32xf32>, each assigned to
|
|
a group of subgroups. Each group consists of 8 subgroups from the same column of sg_layout, sharing a
|
|
single reduction result of type vector<32xf32>.
|
|
|
|
}];
|
|
|
|
let parameters = (ins
|
|
"xegpu::DistributeLayoutAttr": $parent,
|
|
"DenseI64ArrayAttr": $dims
|
|
);
|
|
|
|
let extraClassDeclaration = [{
|
|
|
|
int64_t getRank() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
return parent.getRank() - attr.getDims().size();
|
|
}
|
|
|
|
int64_t getNumSubgroups() const {
|
|
return getParent().getNumSubgroups();
|
|
}
|
|
|
|
DenseI32ArrayAttr getOrder() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
return parent.getOrder();
|
|
}
|
|
|
|
bool isForWorkgroup() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
return parent.isForWorkgroup();
|
|
}
|
|
|
|
bool isForSubgroup() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
return parent.isForSubgroup();
|
|
}
|
|
|
|
bool isForLane() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
return parent.isForLane();
|
|
}
|
|
|
|
/// Returns the SgLayout of the attribute, computed by applying
|
|
/// the slice dimensions to the underlying LayoutAttr.
|
|
SmallVector<int64_t> getEffectiveSgLayoutAsInt() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
auto layout = parent.getEffectiveSgLayoutAsInt();
|
|
if (layout.size()) {
|
|
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
|
|
return XeGPUDialect::slice(ArrayRef<int64_t>(layout), dims);
|
|
}
|
|
return {};
|
|
}
|
|
|
|
/// Returns the SgData of the attribute, computed by applying
|
|
/// the slice dimensions to the underlying LayoutAttr.
|
|
SmallVector<int64_t> getEffectiveSgDataAsInt() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
auto data = parent.getEffectiveSgDataAsInt();
|
|
if (data.size()) {
|
|
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
|
|
return XeGPUDialect::slice(ArrayRef<int64_t>(data), dims);
|
|
}
|
|
return {};
|
|
}
|
|
|
|
/// Returns the InstData of the attribute, computed by applying
|
|
/// the slice dimensions to the underlying LayoutAttr.
|
|
SmallVector<int64_t> getEffectiveInstDataAsInt() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
auto inst = parent.getEffectiveInstDataAsInt();
|
|
if (inst.size()) {
|
|
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
|
|
return XeGPUDialect::slice(llvm::ArrayRef<int64_t>(inst), dims);
|
|
}
|
|
return {};
|
|
}
|
|
|
|
/// Returns the LaneLayout of the attribute, computed by applying
|
|
/// the slice dimensions to the underlying LayoutAttr.
|
|
SmallVector<int64_t> getEffectiveLaneLayoutAsInt() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
auto layout = parent.getEffectiveLaneLayoutAsInt();
|
|
if (layout.size()) {
|
|
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
|
|
return XeGPUDialect::slice(llvm::ArrayRef<int64_t>(layout), dims);
|
|
}
|
|
return {};
|
|
}
|
|
|
|
/// Returns the LaneData of the attribute, computed by applying
|
|
/// the slice dimensions to the underlying LayoutAttr.
|
|
SmallVector<int64_t> getEffectiveLaneDataAsInt() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
auto data = parent.getEffectiveLaneDataAsInt();
|
|
if (data.size()) {
|
|
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
|
|
return XeGPUDialect::slice(llvm::ArrayRef<int64_t>(data), dims);
|
|
}
|
|
return {};
|
|
}
|
|
|
|
/// Returns the effective order of the attribute. The parent's effective
|
|
/// order (or default [rank-1,...,0]) is filtered by removing entries that
|
|
/// correspond to sliced dimensions, then renumbered to form a valid
|
|
/// permutation over the remaining dimensions.
|
|
/// Example: parent order = [2, 0, 1, 3], slice dims = [0, 2] -> result = [1, 0]
|
|
SmallVector<int64_t> getEffectiveOrderAsInt() const {
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
auto order = parent.getEffectiveOrderAsInt();
|
|
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
|
|
SmallVector<int64_t> result;
|
|
for (int64_t d : order) {
|
|
if (llvm::is_contained(dims, d))
|
|
continue;
|
|
int64_t offset = llvm::count_if(dims, [&](int64_t s) { return s < d; });
|
|
result.push_back(d - offset);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
SliceAttr dropSgLayoutAndData() const{
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
parent = parent.dropSgLayoutAndData();
|
|
if (!parent)
|
|
return nullptr;
|
|
return SliceAttr::get(getContext(), parent, attr.getDims());
|
|
}
|
|
|
|
SliceAttr dropInstData() const{
|
|
SliceAttr attr = flatten();
|
|
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
|
|
parent = parent.dropInstData();
|
|
if (!parent)
|
|
return nullptr;
|
|
return SliceAttr::get(getContext(), parent, attr.getDims());
|
|
}
|
|
|
|
//set the layout for the sepcified unit dims: sg_data, inst_data and lane_data to 1
|
|
DistributeLayoutAttr setUnitDimData(SmallVector<int64_t> unitDims) const;
|
|
|
|
//set the layout for the sepcified unit dims: sg_lane and lane_layout to 1
|
|
DistributeLayoutAttr setUnitDimLayout(SmallVector<int64_t> unitDims) const;
|
|
|
|
// Derive a new layout with sg_data, inst_data and lane_data set to the
|
|
// specified values for the given dimension. Passing -1 for any parameter
|
|
// preserves its original value.
|
|
DistributeLayoutAttr setDimData(int64_t dim, int64_t sgData, int64_t instData, int64_t laneData);
|
|
|
|
// Derive a new layout by removing dimensions.
|
|
// `dimGroup` specifies a group of dimensions to be removed in the derived layout.
|
|
DistributeLayoutAttr dropDims(SmallVector<int64_t> dimGroup);
|
|
|
|
// Derive a new layout by collapsing dimensions.
|
|
// `dimGroup` specifies a group of adjacent dimensions
|
|
// that are collapsed into a single dimension in the derived layout.
|
|
DistributeLayoutAttr collapseDims(SmallVector<int64_t> dimGroup);
|
|
|
|
// Derive a new layout by expanding a single dimension `dim` into
|
|
// multiple adjacent dimensions whose shape is given by `targetShape`.
|
|
// See the interface method documentation for the per-field distribution
|
|
// policy.
|
|
DistributeLayoutAttr expandDim(int64_t dim, ArrayRef<int64_t> targetShape);
|
|
|
|
// Derive a new layout by transposing the layout using `permutation`.
|
|
DistributeLayoutAttr transposeDims(ArrayRef<int64_t> permutation);
|
|
|
|
/// flatten a nested SliceAttr, e.g., for 2-level nested SliceAttr
|
|
/// #xegpu.slice<#xegpu.slice<#xegpu.layout<sg_layout = [4, 8, 12]>, dims = [0]>, dims = [0]>
|
|
/// it will coalese two slice operations and return a simplified SliceAttr
|
|
/// #xegpu.slice<#xegpu.layout<sg_layout = [4, 8, 12]>, dims = [0, 1]>
|
|
SliceAttr flatten() const;
|
|
|
|
/// Delinearizes a linear subgroup ID into its multidimensional indices
|
|
/// based on the effective subgroup layout.
|
|
FailureOr<SmallVector<Value>>
|
|
delinearizeId(OpBuilder &builder, Location loc, Value linearId);
|
|
|
|
/// Generates instructions to compute multidimensional coordinates for blocks
|
|
/// assigned to a subgroup identified by linearId. The shape parameter
|
|
/// represents the workgroup-level problem size. Each subgroup may access
|
|
/// multiple blocks according to round-robin distribution rules.
|
|
FailureOr<SmallVector<SmallVector<Value>>>
|
|
computeDistributedCoords(OpBuilder &builder, Location loc, Value linearId, ArrayRef<int64_t> shape);
|
|
|
|
///Statically computes multidimensional coordinates for all dist units
|
|
///assigned to a compute unit identified by `linearId`. This is the
|
|
///compile-time counterpart of `computeDistributedCoords`.
|
|
SmallVector<SmallVector<int64_t>>
|
|
computeStaticDistributedCoords(int64_t linearId, ArrayRef<int64_t> shape);
|
|
|
|
/// Check if this is slice of some other layout.
|
|
bool isSliceOf(const xegpu::DistributeLayoutAttr &other);
|
|
|
|
/// Check if this layout is equal to another layout.
|
|
bool isEqualTo(const xegpu::DistributeLayoutAttr &other);
|
|
|
|
/// Check if this layout is compatible with another layout
|
|
/// at a specific level of the layout hierarchy regarding a given shape.
|
|
bool isCompatibleWith(const xegpu::DistributeLayoutAttr &other,
|
|
SmallVector<int64_t> shape,
|
|
xegpu::LayoutKind level);
|
|
|
|
/// Check if this layout is a transpose of another layout.
|
|
bool isTransposeOf(const xegpu::DistributeLayoutAttr &other, ArrayRef<int64_t> perm, const xegpu::LayoutKind kind);
|
|
|
|
/// Drop the slice dims to get the original layout.
|
|
SliceAttr dropSliceDims(ArrayRef<int64_t> sliceDimsToDrop);
|
|
}];
|
|
|
|
let assemblyFormat = "`<` qualified($parent) `,` `dims` `=` $dims `>`";
|
|
let genVerifyDecl = 1;
|
|
}
|
|
|
|
def XeGPU_RangeAttr : XeGPUAttr<"Range", "range"> {
|
|
let summary = [{Specifies a half-open range}];
|
|
let description = [{
|
|
`RangeAttr` is an attribute that defines a half-open range [start, end).
|
|
The range is inclusive of the start value and exclusive of the end value.
|
|
One usage of this attribute can be to specify the subgroup id range.
|
|
The subgroup id range can be specified using this attribute,
|
|
and it can be attached to a scf.if op like
|
|
```mlir
|
|
scf.if %cond {
|
|
// some operations
|
|
} {sg_id_range = #xegpu.range<[2, 4]>}
|
|
```
|
|
In this case, the scf.if op will only be executed for subgroup IDs 2 and 3.
|
|
}];
|
|
|
|
let parameters = (ins
|
|
"IntegerAttr": $start,
|
|
"IntegerAttr": $end
|
|
);
|
|
|
|
let builders = [
|
|
AttrBuilder<(ins "int":$start, "int":$end)>
|
|
];
|
|
|
|
let assemblyFormat = "`<` `[`$start `,` $end `]` `>`";
|
|
let genVerifyDecl = 1;
|
|
}
|
|
|
|
def XeGPU_MemLayoutAttr : XeGPUAttr<"MemLayout", "mem_layout"> {
|
|
let summary = [{Specifies memory layouts with named attributes.}];
|
|
|
|
let description = [{
|
|
This attribute stores a collection of named attributes that describe
|
|
memory layout properties such as stride, block, etc.
|
|
}];
|
|
|
|
let parameters = (ins "DictionaryAttr": $attrs);
|
|
let hasCustomAssemblyFormat = 1;
|
|
|
|
let extraClassDeclaration = [{
|
|
/// Get a specific attribute by name
|
|
Attribute getAttr(StringRef name) const {
|
|
return getAttrs().get(name);
|
|
}
|
|
|
|
/// Check if a specific attribute exists
|
|
bool hasAttr(StringRef name) const {
|
|
return getAttrs().contains(name);
|
|
}
|
|
|
|
ArrayAttr getStrideAttr() {
|
|
return getAttrs().getAs<ArrayAttr>("stride");
|
|
}
|
|
|
|
ArrayAttr getBlockAttr() {
|
|
return getAttrs().getAs<ArrayAttr>("block");
|
|
}
|
|
|
|
}];
|
|
|
|
}
|
|
|
|
def AnchorLayoutInterface : OpInterface<"AnchorLayoutInterface"> {
|
|
let cppNamespace = "::mlir::xegpu";
|
|
|
|
let description = [{
|
|
An attribute interface for accessing anchor layout information.
|
|
This interface provides a method to set and retrieve the anchor layout
|
|
from operations that implement it.
|
|
}];
|
|
|
|
let methods = [
|
|
InterfaceMethod<
|
|
/*desc=*/"Get the anchor layout attribute.",
|
|
/*retTy=*/"xegpu::DistributeLayoutAttr",
|
|
/*methodName=*/"getAnchorLayout",
|
|
/*args=*/(ins)
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/"Set the anchor layout attribute.",
|
|
/*retTy=*/"void",
|
|
/*methodName=*/"setAnchorLayout",
|
|
/*args=*/(ins "xegpu::DistributeLayoutAttr":$layout)
|
|
>,
|
|
];
|
|
|
|
}
|
|
|
|
#endif // MLIR_DIALECT_XEGPU_IR_XEGPUATTRS_TD
|