[Mlir-commits] [mlir] [MLIR][XeGPU] Don't assert on a multi-dimensional load_matrix result (PR #214432)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 7 02:06:58 PDT 2026
https://github.com/LouisLu060211 updated https://github.com/llvm/llvm-project/pull/214432
>From aae70d92e1222272127238874762a0ed2492d65d Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Thu, 6 Aug 2026 16:29:40 +0800
Subject: [PATCH 1/3] [MLIR][XeGPU] Don't assert on a multi-dimensional
load_matrix result
`LoadStoreMatrixToXeVMPattern` emits a single contiguous access, so it flattens
the result of `xegpu.load_matrix` into a 1D vector. It asserted that the result
had at most one non-unit dimension rather than checking, so a valid op with a
genuinely 2D result, such as a `vector<4x8xf32>`, aborted the compiler.
Report a match failure instead. The op is left in place and the conversion
reports it as unlegalized, which is the normal way an unsupported combination
surfaces.
Fixes #208902
---
.../lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 9 ++++++---
.../XeGPUToXeVM/loadstore_matrix_invalid.mlir | 18 ++++++++++++++++++
2 files changed, 24 insertions(+), 3 deletions(-)
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix_invalid.mlir
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 78d99cf88b768..5cb66dc662c21 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -753,9 +753,12 @@ class LoadStoreMatrixToXeVMPattern : public OpConversionPattern<OpType> {
// Some transforms may leave unit dimension in the 2D vector, adaptors do
// not catch it for results.
if (auto vecType = dyn_cast<VectorType>(resType)) {
- assert(llvm::count_if(vecType.getShape(),
- [](int64_t d) { return d != 1; }) <= 1 &&
- "Expected either 1D vector or nD with unit dimensions");
+ // Only a result flattenable to 1D can be a single contiguous access.
+ if (llvm::count_if(vecType.getShape(),
+ [](int64_t d) { return d != 1; }) > 1)
+ return rewriter.notifyMatchFailure(
+ op, "Expected either a 1D result or an nD result with unit "
+ "dimensions.");
resType = VectorType::get({vecType.getNumElements()},
vecType.getElementType());
}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix_invalid.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix_invalid.mlir
new file mode 100644
index 0000000000000..bb5c426c4d72a
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix_invalid.mlir
@@ -0,0 +1,18 @@
+// RUN: mlir-opt -split-input-file -convert-xegpu-to-xevm -verify-diagnostics %s
+
+// A result with more than one non-unit dimension cannot be flattened into the
+// single contiguous access this lowering emits. Report it instead of tripping
+// an assertion.
+// See https://github.com/llvm/llvm-project/issues/208902.
+
+gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
+ gpu.func @load_matrix_2d_result(%arg0: memref<4x8xf32>) kernel {
+ %alloca = memref.alloca() : memref<4x8xf32, 3>
+ %mdesc = xegpu.create_mem_desc %alloca : memref<4x8xf32, 3> -> !xegpu.mem_desc<4x8xf32>
+ // expected-error at +1 {{failed to legalize operation 'xegpu.load_matrix'}}
+ %res = xegpu.load_matrix %mdesc[0, 0] : !xegpu.mem_desc<4x8xf32> -> vector<4x8xf32>
+ %c0 = arith.constant 0 : index
+ vector.store %res, %arg0[%c0, %c0] : memref<4x8xf32>, vector<4x8xf32>
+ gpu.return
+ }
+}
>From 70d9ef3e5becb646d95d087a1d825cf853f8a775 Mon Sep 17 00:00:00 2001
From: LouisLu060211 <38174270+LouisLu060211 at users.noreply.github.com>
Date: Fri, 7 Aug 2026 02:04:41 -0700
Subject: [PATCH 2/3] Modify loadstore_matrix.mlir for diagnostics and new
kernel
Updated the test to verify diagnostics for conversion from XeGPU to XeVM, and added a new kernel for loading a 2D matrix with error reporting.
---
.../XeGPUToXeVM/loadstore_matrix.mlir | 20 ++++++++++++++++++-
1 file changed, 19 insertions(+), 1 deletion(-)
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
index 07fb09fa2c24b..6c80cbf8ddd3c 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -split-input-file -convert-xegpu-to-xevm %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -convert-xegpu-to-xevm -verify-diagnostics %s | FileCheck %s
gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
@@ -333,3 +333,21 @@ gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
}
}
+
+// -----
+
+// A result with more than one non-unit dimension cannot be flattened into the
+// single contiguous access this lowering emits. Report it instead of tripping
+// an assertion.
+
+gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
+ gpu.func @load_matrix_2d_result(%arg0: memref<4x8xf32>) kernel {
+ %alloca = memref.alloca() : memref<4x8xf32, 3>
+ %mdesc = xegpu.create_mem_desc %alloca : memref<4x8xf32, 3> -> !xegpu.mem_desc<4x8xf32>
+ // expected-error at +1 {{failed to legalize operation 'xegpu.load_matrix'}}
+ %res = xegpu.load_matrix %mdesc[0, 0] : !xegpu.mem_desc<4x8xf32> -> vector<4x8xf32>
+ %c0 = arith.constant 0 : index
+ vector.store %res, %arg0[%c0, %c0] : memref<4x8xf32>, vector<4x8xf32>
+ gpu.return
+ }
+}
>From fe7b315b8eb305adeb658f241c968a6c54acf6a3 Mon Sep 17 00:00:00 2001
From: LouisLu060211 <38174270+LouisLu060211 at users.noreply.github.com>
Date: Fri, 7 Aug 2026 17:06:47 +0800
Subject: [PATCH 3/3] Delete
mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix_invalid.mlir
---
.../XeGPUToXeVM/loadstore_matrix_invalid.mlir | 18 ------------------
1 file changed, 18 deletions(-)
delete mode 100644 mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix_invalid.mlir
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix_invalid.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix_invalid.mlir
deleted file mode 100644
index bb5c426c4d72a..0000000000000
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix_invalid.mlir
+++ /dev/null
@@ -1,18 +0,0 @@
-// RUN: mlir-opt -split-input-file -convert-xegpu-to-xevm -verify-diagnostics %s
-
-// A result with more than one non-unit dimension cannot be flattened into the
-// single contiguous access this lowering emits. Report it instead of tripping
-// an assertion.
-// See https://github.com/llvm/llvm-project/issues/208902.
-
-gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
- gpu.func @load_matrix_2d_result(%arg0: memref<4x8xf32>) kernel {
- %alloca = memref.alloca() : memref<4x8xf32, 3>
- %mdesc = xegpu.create_mem_desc %alloca : memref<4x8xf32, 3> -> !xegpu.mem_desc<4x8xf32>
- // expected-error at +1 {{failed to legalize operation 'xegpu.load_matrix'}}
- %res = xegpu.load_matrix %mdesc[0, 0] : !xegpu.mem_desc<4x8xf32> -> vector<4x8xf32>
- %c0 = arith.constant 0 : index
- vector.store %res, %arg0[%c0, %c0] : memref<4x8xf32>, vector<4x8xf32>
- gpu.return
- }
-}
More information about the Mlir-commits
mailing list