[Mlir-commits] [mlir] [MLIR][XeGPU] Extend SGMapAttr (PR #132425)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Mar 24 09:24:36 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
@llvm/pr-subscribers-mlir
Author: Chao Chen (chencha3)
<details>
<summary>Changes</summary>
This PR enhances the SGMapAttr to support workgroup-level programming, marking the initial step in extending the XeGPU dialect from subgroup to workgroup levels. Since it now encompasses more than just subgroup distribution, it has been renamed to LayoutAttr. Technically, four fields have been added to the previous attribute definition: scope, sg_layout, sg_data, and order. Additionally, a conv_layout operation has been introduced to facilitate conversion between different LayoutAttr configurations, which is essential for both workgroup and subgroup level programming. Per users' requirement, we also renamed wi to lane for clarification.
---
Patch is 116.09 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/132425.diff
7 Files Affected:
- (modified) mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td (+93-19)
- (modified) mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td (+41-20)
- (modified) mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td (+29-24)
- (modified) mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp (+82-108)
- (modified) mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp (+113-45)
- (modified) mlir/test/Dialect/XeGPU/invalid.mlir (+63-47)
- (modified) mlir/test/Dialect/XeGPU/ops.mlir (+125-105)
``````````diff
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index 0136b18ccfa94..7bb59796af36e 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
@@ -154,33 +154,107 @@ def XeGPU_FenceScopeAttr:
let assemblyFormat = "$value";
}
-def XeGPU_SGMapAttr : XeGPUAttr<"SGMap", "sg_map"> {
+def XeGPU_ScopeWG: I32EnumAttrCase<"WG", 0, "wg">; // workgroup level code
+def XeGPU_ScopeSG: I32EnumAttrCase<"SG", 1, "sg">; // subgroup level code
+def XeGPU_ScopeLane: I32EnumAttrCase<"Lane", 2, "lane">; // simt level code
+
+def XeGPU_ScopeEnums : I32EnumAttr<"Scope", "enumeration of scope",
+ [XeGPU_ScopeWG, XeGPU_ScopeSG, XeGPU_ScopeLane]> {
+ let genSpecializedAttr = 0;
+ let cppNamespace = "::mlir::xegpu";
+}
+
+def XeGPU_ScopeAttr
+ : EnumAttr<XeGPU_Dialect, XeGPU_ScopeEnums, "Scope"> {
+ let summary = [{Defines the programming scope of the IR,
+ where WG represents the workgroup level,
+ SG represents the subgroup level, and
+ Lane represents the work-item level}];
+
+ let assemblyFormat = "``$value";
+}
+
+def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout"> {
let summary = [{
- Describes the mapping between work item (WI) and the 2D tensor specified by the tensor descriptor.
+ Describes the data distribution to subgroups and work-items for a tensor
+ specified by the tensor descriptor.
}];
let description = [{
- To distribute the XeGPU operation to work items, the tensor_desc must be specified with the sg_map
- attribute at the tensor description creation time.
- Within the `sg_map`, `wi_layout` specifies the layout of work items,
- describing the mapping of work items to the tensor.
- wi_layout[0] x wi_layout[1] must be equal to the total number of work items within a subgroup.
- `wi_data` specifies the minimum number of data elements assigned to each work item for a single distribution.
-
- E.g., #xegpu.sg_map<wi_layout = [1, 16], wi_data = [1, 1]>
- In this example, the subgroup has 16 work items in wi_layout=[1, 16],
- each accessing 1 element as specified by wi_data=[1, 1].
-
- `wi_data[0] * wi_data[1]` can be greater than 1, meaning that each work item operates on multiple elements,
- which is eventually lowered to "SIMT-flavor" vector, like SPIR-V vector or llvm vector, or packed to a storage data type.
- The multiple elements indicated by `wi_data` can only be from one dimension and must be contiguous in the memory along either dimension.
+ 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, categorized into three groups:
+
+ ### Group 1:
+ * scope: Defines the scope of the code, which can be `wg` (workgroup), `sg` (subgroup),
+ or `lane` (work-item). It is mandatory for subgroup-level programming but optional
+ for workgroup and work-item levels. By default:
+ - If sg_layout is included, the layout is treated as workgroup level.
+ - If only `lane_layout` and `lane_data` are included, it is considered work-item level
+
+ ### Group 2:
+ * sg_layout (optional): 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, and
+ the scope must be empty or set to `wg`.
+ * sg_data (optional): Defines the data size accessed per subgroup. It must be used with sg_layout or
+ left empty, in which case it can be derived from `lane_layout` and `lane_data` using the formula:
+ `sg_data[i] = lane_layout[i] * lane_data[i]`.
+ * order (optional): Specifies the dimension order used to linearize n-dimensional sbugroup IDs to
+ 1-dimensional IDs. The first dimension in the order list is the fastest-changing dimension.
+
+ ### Group 3:
+ * lane_layout (required): Specifies the total number of work-items and their layout within a subgroup
+ * lane_data: (required): Specifies the data size accessed per work-item for a single distribution.
+
+ `lane_data[0] * lane_data[1]` can be greater than 1, indicating that each work item operates on multiple
+ elements. These elements are eventually lowered to a "SIMT-flavor" vector, such as a SPIR-V vector or
+ an LLVM vector, or packed into a storage data type. The multiple elements specified by lane_data must
+ come from a single dimension and be contiguous in memory along either dimension.
+
+ ### Examples:
+ 1. Work-item level layout:
+ ```mlir
+ #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ ```
+ In this example, the subgroup consists of 16 work items arranged as lane_layout=[1, 16], with
+ each work item accessing a single element as defined by lane_data=[1, 1].
+
+ 2. Workgroup level layout:
+ ```mlir
+ #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+ ```
+ In this example, the layout represents a workgroup distribution. A workgroup consists of 8 subgroups
+ arranged in a 2x4 layout. Each subgroup accesses a 16x16 block per instruction, which is further
+ distributed to 16 work items as described above.
}];
+
let parameters = (ins
- ArrayRefParameter<"uint32_t">:$wi_layout,
- ArrayRefParameter<"uint32_t">:$wi_data
+ OptionalParameter<"ScopeAttr">: $scope,
+ OptionalParameter<"DenseI32ArrayAttr">: $sg_layout,
+ OptionalParameter<"DenseI32ArrayAttr">: $sg_data,
+ OptionalParameter<"DenseI32ArrayAttr">: $order,
+ "DenseI32ArrayAttr": $lane_layout,
+ "DenseI32ArrayAttr": $lane_data
);
+ let extraClassDeclaration = [{
+ bool isForWorkgroupLevel() {
+ if (!getScope())
+ return getSgLayout() && getSgData();
+ return getScope() == ScopeAttr::get(getContext(), Scope::WG);
+ }
- let hasCustomAssemblyFormat = 1;
+ bool isForSubgroupLevel() {
+ return getScope() == ScopeAttr::get(getContext(), Scope::SG);
+ }
+
+ bool isForWorkItemLevel() {
+ if (!getScope())
+ return !getSgLayout() && !getSgData() && !getOrder();
+ return getScope() == ScopeAttr::get(getContext(), Scope::Lane);
+ }
+ }];
+
+ let assemblyFormat = "`<` struct(params) `>`";
let genVerifyDecl = 1;
}
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 56b836d707a7d..7188f74815943 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -80,7 +80,7 @@ def XeGPU_CreateNdDescOp: XeGPU_Op<"create_nd_tdesc", [Pure, ViewLikeOpInterface
information e.g., memref<?x?xf16>, the strides information has to be explicitly
passed via the "strides" and "const_strides" argument.
- In SIMT mode, tensor descriptor is augmented with `SGMapAttr` which describes the
+ In SIMT mode, tensor descriptor is augmented with `LayoutAttr` which describes the
mapping of the tensor descriptor to the work items.
Example 1 (suppose the tensor shape inferred by the compiler is 8x16):
@@ -113,7 +113,7 @@ def XeGPU_CreateNdDescOp: XeGPU_Op<"create_nd_tdesc", [Pure, ViewLikeOpInterface
%c0 = arith.constant 0 : index
%c1 = arith.constant 8 : index
%1 = xegpu.create_nd_tdesc %0[%c0, %c0] : memref<1024x1024xf32>
- -> !xegpu.tensor_desc<8x16xf32, #xegpu.sg_map<wi_layout = [1, 16], wi_data = [1, 1]>>
+ -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
```
}];
@@ -306,7 +306,7 @@ def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [
fp32 or fp64. It implies that vnni and transpose cannot exit at the
same time.
- In SIMT mode, LoadNdOp expects the tensor descriptor to be augmented with `SGMapAttr`
+ In SIMT mode, LoadNdOp expects the tensor descriptor to be augmented with `LayoutAttr`
which describes the mapping of the tensor to the work items. In this case, result
vector represents the data to be loaded by each work-item.
@@ -323,7 +323,7 @@ def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [
xegpu.load_nd %1 {l1_hint = #xegpu.cache_hint<cached>,
l2_hint = #xegpu.cache_hint<uncached>}>
: !xegpu.tensor_desc<8x16xf32,
- #xegpu.sg_map<wi_layout = [1, 16], wi_data = [1, 1]>> -> vector<8x1xf32>
+ #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x1xf32>
```
@@ -364,7 +364,7 @@ def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [
of cache, L1, L2 and L3. If hardware does not have a correspoding cache,
Corresponding cache hint attribute will be masked.
- In SIMT mode, StoreNdOp expects the tensor descriptor to be augmented with `SGMapAttr`
+ In SIMT mode, StoreNdOp expects the tensor descriptor to be augmented with `LayoutAttr`
which describes the mapping of the tensor to the work items. In this case, input
vector represents the data to be stored by each work-item.
@@ -381,7 +381,7 @@ def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [
l2_hint = #xegpu.cache_hint<write_back>,
l3_hint = #xegpu.cache_hint<write_through>}
: vector<8x1xf16>, !xegpu.tensor_desc<8x16xf16,
- #xegpu.sg_map<wi_layout = [1, 16], wi_data = [1, 1]>>
+ #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
```
@@ -422,7 +422,7 @@ def XeGPU_UpdateNdOffsetOp : XeGPU_Op<"update_nd_offset",
Example 2 (SIMT mode):
```
%2 = xegpu.update_nd_offset %1, [0, 16]:
- !xegpu.tensor_desc<8x16xf32, #xegpu.sg_map<wi_layout = [1, 16], wi_data = [1, 1]>>
+ !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
```
}];
@@ -482,7 +482,7 @@ def XeGPU_CreateDescOp: XeGPU_Op<"create_tdesc", [Pure, ViewLikeOpInterface]> {
the chunk_size if the chunk size is larger than 1.
In SIMT mode, similar to `create_nd_tdesc` the resulting tensor descriptor is augmented
- with `SGMapAttr` which describes the mapping of the tensor descriptor to the work items.
+ with `LayoutAttr` which describes the mapping of the tensor descriptor to the work items.
In this case, the first dimension of the tensor descriptor represents the work-items, and
the second dimension represents the chunk size.
@@ -517,7 +517,7 @@ def XeGPU_CreateDescOp: XeGPU_Op<"create_tdesc", [Pure, ViewLikeOpInterface]> {
%off = arith.constant dense<[0, 16, 32, 64]> : vector<4xindex>
%1 = xegpu.create_tdesc %0, %off : memref<1024xf32>, vector<4xindex>
-> TensorDesc<4x8xf32, #xegpu.scattered_tdesc_attr<chunk_size = 8>,
- #xegpu.sg_map<wi_layout = [4, 1], wi_data = [1, 1]>>
+ #xegpu.layout<lane_layout = [4, 1], lane_data = [1, 1]>>
```
}];
@@ -623,7 +623,7 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [
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. It applies to slots of SIMD lanes.
- In SIMT mode, LoadGatherOp expects the tensor descriptor to be augmented with `SGMapAttr`
+ In SIMT mode, LoadGatherOp expects the tensor descriptor to be augmented with `LayoutAttr`
which describes the mapping of the tensor to the work items. In this case, result vector
represents the data to be loaded by each work-item. Each work-item recieves a `chunk_size`
number of elements.
@@ -653,7 +653,7 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [
l2_hint = #xegpu.cache_hint<uncached>,
l3_hint = #xegpu.cache_hint<uncached>}
: !xegpu.tensor_desc<16x8xf32, #xegpu.scatter_tdesc_attr<memory_space=global, chunk_size=8>,
- !xegpu.sg_map<wi_layout = [16, 1], wi_data = [1, 1]>>
+ !xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>>
vector<16xi1> -> vector<8x1xf32>
```
@@ -704,7 +704,7 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [
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 SIMT mode, StoreScatterOp expects the tensor descriptor to be augmented with `SGMapAttr`
+ In SIMT mode, StoreScatterOp expects the tensor descriptor to be augmented with `LayoutAttr`
which describes the mapping of the tensor to the work items. In this case, input vector
represents the data to be stored by each work-item. Each work-item recieves a `chunk_size`
number of elements.
@@ -732,7 +732,7 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [
l2_hint = #xegpu.cache_hint<write_back>,
l3_hint = #xegpu.cache_hint<write_through>}
: vector<8x1xf32>, !xegpu.tensor_desc<16x8xf32, #xegpu.scattered_tdesc_attr<chunk_size=8>,
- !xegpu.sg_map<wi_layout = [16, 1], wi_data = [1, 1]>> vector<16xi1>
+ !xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>> vector<16xi1>
```
}];
@@ -790,7 +790,7 @@ def XeGPU_UpdateOffsetOp: XeGPU_Op<"update_offset",
%off = arith.constant dense<[32, 32, 32, 32]> : vector<4xindex>
%2 = xegpu.update_offset %1, %off :
!xegpu.tensor_desc<4x2xf32, #xegpu.scattered_tdesc_attr<chunk_size=2>,
- #xegpu.sg_map<wi_layout = [4, 1], wi_data = [1, 1]>>, vector<4xindex>
+ #xegpu.layout<lane_layout = [4, 1], lane_data = [1, 1]>>, vector<4xindex>
```
}];
@@ -840,9 +840,9 @@ def XeGPU_DpasOp : XeGPU_Op<"dpas", [Pure, AllElementTypesMatch<["lhs", "rhs"]>]
factor, which is computed as `32/bit_width_of_elem_type`. Thus, `B: vector<16x16xf16>`
can be represented as `B: vector<8x16x2xf16>`.
- In SIMT mode, DpasOp expects attributes `sg_map_a`, `sg_map_b`, and `sg_map_c`
- which descibes the data fragment owned by each work-item w.r.t. the tensor
- descriptor these data are loaded from.
+ In SIMT mode, DpasOp expects layout attributes `a`, `b`, and `c` (only if acc is used)
+ which descibe the data fragment owned by each work-item w.r.t. the tensor descriptor
+ these data are loaded from.
Note: on PVC, the hardware can perform load with VNNI transformation when data
element type is 16-bit or lower precision, taking 2 or 4 elements from
@@ -853,9 +853,9 @@ def XeGPU_DpasOp : XeGPU_Op<"dpas", [Pure, AllElementTypesMatch<["lhs", "rhs"]>]
XeGPU_DpasOpType : $lhs,
XeGPU_DpasOpType : $rhs,
Optional<XeGPU_Vector2DType>: $acc,
- OptionalAttr<XeGPU_SGMapAttr>:$sg_map_a,
- OptionalAttr<XeGPU_SGMapAttr>:$sg_map_b,
- OptionalAttr<XeGPU_SGMapAttr>:$sg_map_c);
+ OptionalAttr<XeGPU_LayoutAttr>:$a_layout,
+ OptionalAttr<XeGPU_LayoutAttr>:$b_layout,
+ OptionalAttr<XeGPU_LayoutAttr>:$c_layout);
let results = (outs XeGPU_Vector2DType: $result);
let extraClassDeclaration = [{
@@ -876,6 +876,10 @@ def XeGPU_DpasOp : XeGPU_Op<"dpas", [Pure, AllElementTypesMatch<["lhs", "rhs"]>]
VectorType getResultType() {
return getResult().getType();
}
+
+ bool hasAcc() {
+ return getAcc() != nullptr;
+ }
}];
let assemblyFormat = [{
@@ -979,4 +983,21 @@ def XeGPU_FenceOp: XeGPU_Op<"fence", []> {
let extraClassDeclaration = extraBaseClassDeclaration;
}
+def XeGPU_ConvertLayoutOp: XeGPU_Op<"convert_layout", [Pure, AllTypesMatch<["source", "result"]>]> {
+ let summary = "Convert the sg layout of the input operand";
+ let description = [{
+ convert_layout remaps the distribution of data across workitems by updating the LayoutAttr.
+ }];
+ let arguments = (ins XeGPU_Vector2DType: $source,
+ XeGPU_LayoutAttr: $srcMap,
+ XeGPU_LayoutAttr: $resMap
+ );
+ let results = (outs XeGPU_Vector2DType: $result);
+ let assemblyFormat = [{
+ $source attr-dict `:` type($source)
+ }];
+
+ let hasVerifier = 1;
+}
+
#endif // MLIR_DIALECT_XEGPU_IR_XEGPUOPS_TD
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
index ccd91a928e1dd..8559f4beb2c03 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
@@ -34,27 +34,24 @@ def XeGPU_TensorDesc: XeGPUTypeDef<"TensorDesc", "tensor_desc",
[ShapedTypeInterface], "::mlir::TensorType"> {
let summary = "TensorDesc describing regions of interested data.";
let description = [{
- TensorDesc is a type designed to describe regions of the interested data as well as some
- features that are unique to Intel hardware. Different with the builtin tensor type in MLIR,
- it essentially only contains the meta data, and doesn't hold the data by itself. It is designed
- to mainly support 2D block load/store and DPAS (matrix multiplication instruction) on Intel GPU.
- It encodes the following information:
+ TensorDesc is a type designed to describe regions of interest in data, as well as some features
+ unique to Intel hardware. Unlike the built-in tensor type in MLIR, it essentially contains only
+ metadata and does not hold the data itself. It is primarily designed to support 2D block load/store
+ and DPAS (matrix multiplication instruction) on Intel GPUs. It encodes the following information:
* shape: the sizes/shape of the intereted data block, e.g., 8x16 means 8 rows
and each row contains 16 contiguous data element. The rows could be
- either contiguous or not, depends on whether the encoding attribute
- is set or not.
- * element_type: the data type of the data element, e.g., f16, f32.
+ either contiguous or not, depends on the encoding attribute. If the
+ encoding is a BlockTensorDescAttr, rows are contiguous. If the encoding
+ is a ScatterTensorDescAttr, rows are not necessary to be contiguous. If
+ encoding is not set, it is considered as a default BlockTensorDescAttr.
- Similar to the builtin tensor, it also provides an optinal attribute to encoding
- the following information via the TensorDescAttr object:
- * memory_space (xegpu::MemorySpace): [optional] where the data is located,
- global memory or shared memory. It is default to Global.
- * array_length (int): [optional] The number of contiguous blocks with size as `shape`,
- that will be loaded by block load at a time. It is default to 1.
- * boundary_check (bool): [optional] indicates whether the operation detects the boundary
- and pads with zero for out-of-boundary access. It is default to do boundary check.
+ * element_type: the data type of the data element, e.g., f16, f32.
+ Similar to the built-in tensor, it also provides optional attributes for encoding
+ additional information via either BlockTensorDescAttr or ScatterTensorDescAttr, or
+ supporting Workgroup, Subgroup, and workitem (or SIMT) level programmings via the
+ Layout attribute. Please check their definition for details.
Syntax:
@@ -63,7 +60,9 @@ def XeGPU_TensorDesc: XeGPUTypeDef<"TensorDesc", "tensor_desc",
element-type ::= float-type | integer-type | index-type
dim-list := (static-dim-list `x`)?
static-dim-list ::= decimal-literal `x` decimal-literal
- attr-list = (, memory_space = value)? (, arr_len = value)? (, boundary_check = value)? (, scattered = value)? (, sg_map `<` wi_layout = value, wi_data = value `>`)?
+ attr-list = (, encoding-attr)? (, layout-attr)?
+ enconding-attr = (, memory_space = value)? (, arr_len = value)? (, boundary_check = value)? (, scattered = value)?
+ layout-attr = (, layout `<` (scope = value,)? (sg_layout = value, sg_data = value, order = value)? lane_layout = value, lane_data = value `>`)?
```
Examples:
@@ -78,15 +77,21 @@ def XeGPU_TensorDesc: XeGPUTypeDef<"TensorDesc", "tensor_desc",
// A...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/132425
More information about the Mlir-commits
mailing list