[Mlir-commits] [mlir] 456906f - [MLIR][XeGPU]Extend load_matrix/store_matrix to support 1D SLM access (#198652)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri May 22 21:45:39 PDT 2026
Author: Jianhui Li
Date: 2026-05-22T21:45:34-07:00
New Revision: 456906f06d4f3f4107ece2b9c522cc6f855c7dd1
URL: https://github.com/llvm/llvm-project/commit/456906f06d4f3f4107ece2b9c522cc6f855c7dd1
DIFF: https://github.com/llvm/llvm-project/commit/456906f06d4f3f4107ece2b9c522cc6f855c7dd1.diff
LOG: [MLIR][XeGPU]Extend load_matrix/store_matrix to support 1D SLM access (#198652)
This PR extended xegpu.load_matrix and xegpu.store_matrix to support 1D
mem_desc for contiguous SLM access
- Added unit tests for 1D load/store (valid ops and invalid cases)
- Added integration test verifying both 1D (<4096xbf16>) and 2D
(<64x128xbf16>), correctly lower through the full WG→SG→WI→XeVM pipeline
---------
Co-authored-by: Claude Opus 4.6 (1M context) <noreply at anthropic.com>
Added:
mlir/test/Integration/Dialect/XeGPU/WG/load_store_matrix.mlir
Modified:
mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
mlir/test/Dialect/XeGPU/invalid.mlir
mlir/test/Dialect/XeGPU/ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index f0c380752306d..fce14999b4011 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -1204,12 +1204,18 @@ def XeGPU_CreateMemDescOp: XeGPU_Op<"create_mem_desc", [Pure,
as the underlying shared local memory.
Arguments:
- - `source` : 1D or 2D statically shape memref, representing the raw SLM buffer. The provided memref must be contiguous.
+ - `source` : 1D or 2D statically shaped memref, representing the raw SLM buffer. The provided memref must be contiguous.
Results:
- - `mem_desc` : the memory descriptor.
+ - `mem_desc` : the memory descriptor (1D or higher).
- Example:
+ 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>
@@ -1237,15 +1243,15 @@ def XeGPU_LoadMatrixOp: XeGPU_Op<"load_matrix", [MemoryEffects<[MemRead]>,
}];
let description = [{
- This operation loads an nD block of data from shared local memory (SLM) as specified
- by the provided nD `mem_desc`. Memory descriptors of any rank are supported.
+ 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 matrix to read from.
+ - `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.
@@ -1254,9 +1260,14 @@ def XeGPU_LoadMatrixOp: XeGPU_Op<"load_matrix", [MemoryEffects<[MemRead]>,
Only valid at workgroup and subgroup levels.
Results:
- - `res`: the matrix elements loaded from SLM.
+ - `res`: the elements loaded from SLM.
- Example (Workgroup level):
+ 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] <{
@@ -1308,23 +1319,28 @@ def XeGPU_StoreMatrixOp: XeGPU_Op<"store_matrix", [MemoryEffects<[MemWrite]>,
let assemblyFormat = [{ $data `,` $mem_desc `` custom<DynamicIndexList>($offsets, $const_offsets)
prop-dict attr-dict `` `:` type(operands)}];
let description = [{
- This operation stores an nD `data` fragment into the shared local memory region
- specified by an nD `mem_desc`. Memory descriptors of any rank are supported.
+ 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 matrix where the data will be written.
- - `data`: the values to be stored in the matrix.
+ - `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 load. When this attribute is present, the offsets are subgroup-uniform
+ 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 `tensor_desc` operand as well as
+ - `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 (Workgroup level):
+ 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] <{
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 69db066982de9..56d340141ee8f 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -126,9 +126,6 @@ IsValidMatrixOpParams(VectorType dataTy, MemDescType mdescTy,
return success();
}
- if (mdescTy.getRank() < 2)
- return emitError() << "mem_desc must be 2D or greater.";
-
ArrayRef<int64_t> dataShape = dataTy.getShape();
ArrayRef<int64_t> mdescShape = mdescTy.getShape();
@@ -163,23 +160,16 @@ IsValidMatrixOpParams(VectorType dataTy, MemDescType mdescTy,
SmallVector<int64_t>(dataShape.begin(), dataShape.end())))
return emitError() << "Value shape is not distributable with the layout";
- if (dataShape.size() == 2) {
+ if (dataShape.size() == mdescShape.size()) {
if (llvm::any_of(llvm::zip_equal(dataShape, mdescShape),
[](auto p) { return std::get<0>(p) > std::get<1>(p); }))
return emitError() << "data shape must not exceed mem_desc shape.";
- } else {
- // if the subgroup_block_io attribute is set, mdescTy must have block
- // attribute
- if (subgroup_block_io && !blockShape.size())
- return emitError() << "mem_desc must have block attribute when "
- "subgroup_block_io is set.";
- // if the subgroup_block_io attribute is set, the memdesc should be row
- // major
- if (subgroup_block_io && mdescTy.isColMajor())
- return emitError() << "mem_desc should be row major when "
- "subgroup_block_io is set.";
}
-
+ // if the subgroup_block_io attribute is set, mdescTy must have block
+ // attribute
+ if (subgroup_block_io && !blockShape.size())
+ return emitError() << "mem_desc must have block attribute when "
+ "subgroup_block_io is set.";
return success();
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 80dec0f33ee0e..b8a6206d07d47 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1170,11 +1170,11 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
/// load matrix lowers to load gather and 1d block load. All of them share the
/// same layout setup logic.
/// For Subgroup layout, uses the consumer layout directly.
-/// non-chunked loads:
+/// non-chunked loads (1D or 2D):
/// InstData = {1, ..., min(consumer, maxLaneLoadSize * subgroupSize)}
/// LaneLayout = {1, ..., subgroupSize}
/// lane_data = {1, ..., min(consumer, maxLaneLoadSize)}
-/// chunked loads:
+/// chunked loads (2D only):
/// InstData = {subgroupSize, min(consumer, maxLaneLoadSize)}
/// LaneLayout = {subgroupSize, 1}
/// lane_data={1,min(consumer, maxLaneLoadSize)}
@@ -1268,12 +1268,12 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
/// Sets up the anchor layout for store scatter and store matrix operation.
/// store matrix lowers to store scatter and 1d block store. All of them share
-/// the same layout setup logic. For Subgroup layout, not support yet.
-/// non-chunked stores:
+/// the same layout setup logic. For Subgroup layout, not supported yet.
+/// non-chunked stores (1D or 2D):
/// InstData = {1, ..., subgroupSize}
/// LaneLayout = {1, ..., subgroupSize}
/// lane_data = {1, ..., 1}
-/// chunked stores:
+/// chunked stores (2D only):
/// InstData = {subgroupSize, min(srcVec, maxLaneStoreSize)}
/// LaneLayout = {subgroupSize, 1}
/// lane_data={1,min(srcVec, maxLaneStoreSize)}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 4237fcf93ae0e..323caa1bd5738 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1323,7 +1323,6 @@ void LayoutInfoPropagation::visitLoadMatrixOp(
}
}
-// Store matrix is a flavor of scattered store for 2D shapes.
void LayoutInfoPropagation::visitStoreMatrixOp(
xegpu::StoreMatrixOp storeMatrix, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index c469a40310607..341427e7e1231 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -635,9 +635,9 @@ func.func @load_mem_desc_invalid_result_size(%arg0: !xegpu.mem_desc<16x64xf16>)
}
// -----
-func.func @load_mem_desc_invalid_rank(%arg0: !xegpu.mem_desc<64xf16>) {
- // expected-error at +1 {{mem_desc must be 2D or greater}}
- %data = xegpu.load_matrix %arg0[16]: !xegpu.mem_desc<64xf16> -> vector<16xf16>
+func.func @load_mem_desc_1d_exceeds_shape(%arg0: !xegpu.mem_desc<16xf16>) {
+ // expected-error at +1 {{data shape must not exceed mem_desc shape}}
+ %data = xegpu.load_matrix %arg0[0]: !xegpu.mem_desc<16xf16> -> vector<32xf16>
return
}
@@ -656,9 +656,9 @@ func.func @store_mem_desc_invalid_data_size(%arg0: !xegpu.mem_desc<16x64xf16>, %
}
// -----
-func.func @store_mem_desc_invalid_rank(%arg0: !xegpu.mem_desc<64xf16>, %arg1: vector<32xf16>) {
- // expected-error at +1 {{mem_desc must be 2D or greater}}
- xegpu.store_matrix %arg1, %arg0[32] : vector<32xf16>, !xegpu.mem_desc<64xf16>
+func.func @store_mem_desc_1d_exceeds_shape(%arg0: !xegpu.mem_desc<16xf16>, %arg1: vector<32xf16>) {
+ // expected-error at +1 {{data shape must not exceed mem_desc shape}}
+ xegpu.store_matrix %arg1, %arg0[0] : vector<32xf16>, !xegpu.mem_desc<16xf16>
return
}
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index f28fb39fd278b..198dc64d2814b 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -600,6 +600,41 @@ gpu.func @simt_store_matrix_vector(%arg0: !xegpu.mem_desc<16x64xf16, #xegpu.mem_
gpu.return
}
+// CHECK-LABEL: gpu.func @load_matrix_1d
+gpu.func @load_matrix_1d(%arg0: !xegpu.mem_desc<64xf16>) {
+ // CHECK: xegpu.load_matrix %{{.+}}[16] : !xegpu.mem_desc<64xf16> -> vector<16xf16>
+ %data = xegpu.load_matrix %arg0[16]: !xegpu.mem_desc<64xf16> -> vector<16xf16>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @store_matrix_1d
+gpu.func @store_matrix_1d(%arg0: !xegpu.mem_desc<64xf16>, %arg1: vector<16xf16>) {
+ // CHECK: xegpu.store_matrix %{{.+}}, %{{.+}}[0] : vector<16xf16>, !xegpu.mem_desc<64xf16>
+ xegpu.store_matrix %arg1, %arg0[0]: vector<16xf16>, !xegpu.mem_desc<64xf16>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @load_matrix_1d_block_io
+gpu.func @load_matrix_1d_block_io(%arg0: !xegpu.mem_desc<64xf16, #xegpu.mem_layout<block = [16]>>) {
+ // CHECK: xegpu.load_matrix %{{.+}}[0] <{subgroup_block_io}>: !xegpu.mem_desc<64xf16, #xegpu.mem_layout<block = [16]>> -> vector<16xf16>
+ %data = xegpu.load_matrix %arg0[0] <{subgroup_block_io}>: !xegpu.mem_desc<64xf16, #xegpu.mem_layout<block = [16]>> -> vector<16xf16>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @store_matrix_1d_block_io
+gpu.func @store_matrix_1d_block_io(%arg0: !xegpu.mem_desc<64xf16, #xegpu.mem_layout<block = [16]>>, %arg1: vector<16xf16>) {
+ // CHECK: xegpu.store_matrix %{{.+}}, %{{.+}}[0] <{subgroup_block_io}>: vector<16xf16>, !xegpu.mem_desc<64xf16, #xegpu.mem_layout<block = [16]>>
+ xegpu.store_matrix %arg1, %arg0[0] <{subgroup_block_io}>: vector<16xf16>, !xegpu.mem_desc<64xf16, #xegpu.mem_layout<block = [16]>>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @simt_load_matrix_1d
+gpu.func @simt_load_matrix_1d(%arg0: !xegpu.mem_desc<64xf16>) {
+ // CHECK: xegpu.load_matrix %{{.+}}[0] : !xegpu.mem_desc<64xf16> -> vector<1xf16>
+ %data = xegpu.load_matrix %arg0[0]: !xegpu.mem_desc<64xf16> -> vector<1xf16>
+ gpu.return
+}
+
// CHECK-LABEL: gpu.func @truncf
gpu.func @truncf(%a: vector<8x16xf16>) {
// CHECK: %{{.+}} = xegpu.truncf %{{.+}} : vector<8x16xf16> -> vector<8x16xf8E5M2>
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/load_store_matrix.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/load_store_matrix.mlir
new file mode 100644
index 0000000000000..598991c7401c6
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/load_store_matrix.mlir
@@ -0,0 +1,40 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(gpu.module(xegpu-wg-to-sg-distribute, xegpu-blocking, xegpu-sg-to-wi-distribute-experimental), xevm-attach-target{chip=pvc}, gpu.module(convert-xegpu-to-xevm))" \
+// RUN: | FileCheck %s
+
+#layout_1d = #xegpu.layout<sg_layout = [16], sg_data = [256], inst_data = [16], lane_layout = [16], lane_data = [1]>
+#layout_2d = #xegpu.layout<sg_layout = [4, 4], sg_data = [16, 32], inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+
+gpu.module @test {
+
+ // CHECK-LABEL: gpu.func @test_load_store_matrix_1d
+ // CHECK: memref.extract_aligned_pointer_as_index %arg0 : memref<8192xi8, 3> -> index
+ // CHECK: gpu.subgroup_id
+ // CHECK: gpu.lane_id
+ // CHECK-COUNT-16: llvm.load %{{.*}} : !llvm.ptr<3> -> bf16
+ // CHECK-COUNT-16: llvm.store %{{.*}}, %{{.*}} : bf16, !llvm.ptr<3>
+ // CHECK-NOT: xegpu.load_matrix
+ // CHECK-NOT: xegpu.store_matrix
+ gpu.func @test_load_store_matrix_1d(%src: memref<8192xi8, 3>) {
+ %c0 = arith.constant 0 : index
+ %mdesc = xegpu.create_mem_desc %src : memref<8192xi8, 3> -> !xegpu.mem_desc<4096xbf16>
+ %data = xegpu.load_matrix %mdesc[%c0] {layout = #layout_1d} : !xegpu.mem_desc<4096xbf16>, index -> vector<4096xbf16>
+ xegpu.store_matrix %data, %mdesc[%c0] {layout = #layout_1d} : vector<4096xbf16>, !xegpu.mem_desc<4096xbf16>, index
+ gpu.return
+ }
+
+ // CHECK-LABEL: gpu.func @test_load_store_matrix_2d
+ // CHECK: memref.extract_aligned_pointer_as_index %arg0 : memref<16384xi8, 3> -> index
+ // CHECK: gpu.subgroup_id
+ // CHECK: gpu.lane_id
+ // CHECK-COUNT-32: llvm.load %{{.*}} : !llvm.ptr<3> -> bf16
+ // CHECK-COUNT-32: llvm.store %{{.*}}, %{{.*}} : bf16, !llvm.ptr<3>
+ // CHECK-NOT: xegpu.load_matrix
+ // CHECK-NOT: xegpu.store_matrix
+ gpu.func @test_load_store_matrix_2d(%src: memref<16384xi8, 3>) {
+ %c0 = arith.constant 0 : index
+ %mdesc = xegpu.create_mem_desc %src : memref<16384xi8, 3> -> !xegpu.mem_desc<64x128xbf16>
+ %data = xegpu.load_matrix %mdesc[%c0, %c0] {layout = #layout_2d} : !xegpu.mem_desc<64x128xbf16>, index, index -> vector<64x128xbf16>
+ xegpu.store_matrix %data, %mdesc[%c0, %c0] {layout = #layout_2d} : vector<64x128xbf16>, !xegpu.mem_desc<64x128xbf16>, index, index
+ gpu.return
+ }
+}
More information about the Mlir-commits
mailing list