[Mlir-commits] [mlir] [MLIR][XeGPU] Relax LoadStoreMatrixToXeVMPattern to accept non-distributed load_matrix/store_matrix (PR #210595)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-gpu

Author: Naaz30

<details>
<summary>Changes</summary>

`LoadStoreMatrixToXeVMPattern` currently rejects `xegpu.load_matrix` and `xegpu.store_matrix` operations in two situations that are not required for the plain (non-subgroup_block_io) lowering path:

1. **Vector shape restriction** : The pass asserts that the vector type has at most one non-unit dimension, assuming that xegpu-sg-to-lane-distribute has       already reduced the operation to a per-lane fragment. This causes the pass to fail (or behave incorrectly in NDEBUG builds) for valid multi-dimensional tiles such as vector<4x8xf32>.

2. **Chip-specific target check** : The pass requires the operation to reside inside a gpu.module with an xevm.target attribute identifying a supported architecture (PVC, BMG, or CRI) via xegpu::getChipStr(). However, this requirement is unnecessary for the plain lowering path, which only generates chip-independent LLVM::LoadOp and LLVM::StoreOp operations.

**Changes**
Remove the unnecessary vector shape restriction for the plain lowering path.
Remove the chip target requirement for the plain lowering path.

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


2 Files Affected:

- (modified) mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp (+7-7) 
- (modified) mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir (+32-9) 


``````````diff
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 6144d7c0c1a15..03808d2d0f8ff 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -719,9 +719,9 @@ 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");
+        // Flatten to 1D
+        // Accepts multi-dim tile as one flat, contiguous run. 
+        // This is only valid when the underlying mem_desc region is contiguous in memory
         resType = VectorType::get({vecType.getNumElements()},
                                   vecType.getElementType());
       }
@@ -787,9 +787,9 @@ class LoadStoreMatrixToXeVMPattern : public OpConversionPattern<OpType> {
 
     if (valOrResVecTy.getNumElements() >= 1) {
       auto chipOpt = xegpu::getChipStr(op);
-      if (!chipOpt ||
-          (*chipOpt != "pvc" && *chipOpt != "bmg" && *chipOpt != "cri")) {
-        // the lowering for chunk load only works for pvc, bmg or cri
+      // reject an explicitly unsupported chip
+      if (chipOpt && *chipOpt != "pvc" && *chipOpt != "bmg" &&
+          *chipOpt != "cri") {
         return rewriter.notifyMatchFailure(
             op, "The lowering is specific to pvc, bmg or cri.");
       }
@@ -1626,4 +1626,4 @@ void mlir::populateXeGPUToXeVMConversionPatterns(
   patterns.add<DpasMxToXeVMPattern>(typeConverter, patterns.getContext());
   patterns.add<ExtfToXeVMPattern, TruncfToXeVMPattern>(typeConverter,
                                                        patterns.getContext());
-}
+}
\ No newline at end of file
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
index 07fb09fa2c24b..50f7263b65e05 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
@@ -1,10 +1,11 @@
 // RUN: mlir-opt  -split-input-file -convert-xegpu-to-xevm %s | FileCheck %s
 
-gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
+gpu.module @test_kernel[#xevm.target<chip = "pvc">] {
 
   // e.g. for mem_desc<32x32xf16, @strides=[1, 16]>
-  // its memory layout tuple is (blocked shape = [1,1,32,32],strides=[1024,1024,32,1])
-  //CHECK-LABEL: load_store_matrix_plain
+  // its memory layout tuple is (blocked shape =
+  // [1,1,32,32],strides=[1024,1024,32,1])
+  // CHECK-LABEL: load_store_matrix_plain
   gpu.func @load_store_matrix_plain(%arg0: memref<4096xi8, 3>) -> f32 {
 
     //CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %arg0 : memref<4096xi8, 3> -> index
@@ -324,12 +325,34 @@ gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
 
   //CHECK-LABEL: load_matrix_f8e8m0_vector
   gpu.func @load_matrix_f8e8m0_vector(%arg0: memref<1024xi8, 3>) -> vector<8xf8E8M0FNU> {
-    %c0 = arith.constant 0 : index
-    %0 = xegpu.create_mem_desc %arg0 : memref<1024xi8, 3> -> !xegpu.mem_desc<32x32xf8E8M0FNU>
-    //CHECK: %[[LOADEDV:.*]] = llvm.load %{{.*}} : !llvm.ptr<3> -> vector<8xi8>
-    //CHECK: vector.bitcast %[[LOADEDV]] : vector<8xi8> to vector<8xf8E8M0FNU>
-    %1 = xegpu.load_matrix %0[%c0, %c0]: !xegpu.mem_desc<32x32xf8E8M0FNU>, index, index -> vector<8xf8E8M0FNU>
-    gpu.return %1 : vector<8xf8E8M0FNU>
+    % c0 = arith.constant 0 : index % 0 =
+               xegpu.create_mem_desc %
+               arg0 : memref<1024xi8, 3>->!xegpu.mem_desc<32x32xf8E8M0FNU>
+               // CHECK: %[[LOADEDV:.*]] = llvm.load %{{.*}} : !llvm.ptr<3> ->
+               // vector<8xi8> CHECK: vector.bitcast %[[LOADEDV]] : vector<8xi8>
+               // to vector<8xf8E8M0FNU>
+               % 1 = xegpu.load_matrix % 0 [% c0, % c0]
+        : !xegpu.mem_desc<32x32xf8E8M0FNU>,
+      index, index->vector<8xf8E8M0FNU> gpu.return % 1 : vector<8xf8E8M0FNU>
   }
+}
 
+// -----
+
+// CHECK-LABEL: func.func @m(
+// CHECK: memref.extract_aligned_pointer_as_index
+// CHECK: llvm.inttoptr
+// CHECK: llvm.load
+// CHECK-SAME: vector<32xf32>
+
+module {
+  func.func @m(% arg0 : index) {
+    % alloca =
+        memref.alloca()
+        : memref<4x8xf32, 3> %
+        0 = xegpu.create_mem_desc %
+            alloca : memref<4x8xf32, 3>->!xegpu.mem_desc<4x8xf32> % 1 =
+                xegpu.load_matrix %
+                0 [0, 0] : !xegpu.mem_desc<4x8xf32>->vector<4x8xf32> return
+  }
 }

``````````

</details>


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