[Mlir-commits] [mlir] [MLIR][XeGPU] Don't assert on a multi-dimensional load_matrix result (PR #214432)

[llvmlistbot at llvm.org]

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/2] [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/2] 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
+  }
+}