[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 created https://github.com/llvm/llvm-project/pull/214432

`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

Assisted by: Claude Opus 5

>From 565ed8929f14c35f38613b0f9a89948cf0ea2bef Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Thu, 6 Aug 2026 16:29:40 +0800
Subject: [PATCH] [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 | 12 +++++++++---
 .../XeGPUToXeVM/loadstore_matrix_invalid.mlir  | 18 ++++++++++++++++++
 2 files changed, 27 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..cb2359df2f652 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -753,9 +753,15 @@ 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");
+        // This pattern emits a single contiguous access, so the result has to
+        // be flattenable to a 1D vector. A result with more than one non-unit
+        // dimension is valid IR but is not something this lowering can
+        // express, so bail out instead of asserting.
+        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
+  }
+}