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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: LouisLu060211

<details>
<summary>Changes</summary>

`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

---
Full diff: https://github.com/llvm/llvm-project/pull/214432.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp (+9-3) 
- (added) mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix_invalid.mlir (+18) 


``````````diff
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
+  }
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/214432