[Mlir-commits] [mlir] [MLIR][XeGPU] Fix load_matrix lowering for non-LLVM element types (PR #203629)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Sang Ik Lee (silee2)

<details>
<summary>Changes</summary>

LoadStoreMatrixToXeVMPattern built the llvm.load result from the raw op result element type. For element types without a direct LLVM representation (e.g. f8E8M0FNU) this produced an illegal op: 'llvm.load' op result #<!-- -->0 must be LLVM type with size, but got 'f8E8M0FNU'.

Derive the load result type from the type converter instead. This maps such element types to an integer storage type of the same bit width, collapses single-element vectors to a scalar, and flattens multi-element vectors. The store path already used the converted operand and is unchanged; the XeVM type converter's materialization casts bridge the loaded value back to the original vector type for downstream consumers.

Add load_matrix regression tests for f8E8M0FNU (scalar and vector).

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


2 Files Affected:

- (modified) mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp (+8-10) 
- (modified) mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir (+25-1) 


``````````diff
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 7d006c9e69f89..c7d593c74f264 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -796,16 +796,14 @@ class LoadStoreMatrixToXeVMPattern : public OpConversionPattern<OpType> {
     }
 
     if constexpr (std::is_same_v<OpType, xegpu::LoadMatrixOp>) {
-      // if the size of valOrResVecTy is 1, it lowers to a scalar load/store
-      // operation. LLVM load/store does not support vector of size 1, so we
-      // need to handle this case separately.
-      auto scalarTy = valOrResVecTy.getElementType();
-      LLVM::LoadOp loadOp;
-      if (valOrResVecTy.getNumElements() == 1)
-        loadOp = LLVM::LoadOp::create(rewriter, loc, scalarTy, basePtrLLVM);
-      else
-        loadOp =
-            LLVM::LoadOp::create(rewriter, loc, valOrResVecTy, basePtrLLVM);
+      // The load result type is taken from the type converter. This maps
+      // element types that are not directly representable in LLVM (e.g.
+      // f8E8M0FNU) to an integer storage type of the same bit width, and
+      // collapses single-element vectors to a scalar, since LLVM load/store
+      // does not support vectors of size 1.
+      Type loadTy =
+          this->getTypeConverter()->convertType(op.getResult().getType());
+      auto loadOp = LLVM::LoadOp::create(rewriter, loc, loadTy, basePtrLLVM);
       rewriter.replaceOp(op, loadOp);
     } else {
       LLVM::StoreOp::create(rewriter, loc, adaptor.getData(), basePtrLLVM);
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
index fa683175693be..07fb09fa2c24b 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
@@ -256,7 +256,7 @@ gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
     //CHECK: %[[ptr:.*]] = llvm.inttoptr %[[final_ptr]] : i32 to !llvm.ptr<3>
     //CHECK: %[[blockload:.*]] = xevm.blockload %[[ptr]] : (!llvm.ptr<3>) -> vector<8xi16>
     //CHECK: %[[loaded:.*]] = vector.bitcast %[[blockload]] : vector<8xi16> to vector<8xf16>
-    
+
     %0 = xegpu.create_mem_desc %arg0 : memref<4096xi8, 3> -> !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>
 
     %c16 = arith.constant 16 : index
@@ -307,5 +307,29 @@ gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
 
   }
 
+  // f8E8M0FNU has no native LLVM representation. The load must use the integer
+  // storage type from the XeVM type converter (i8) and then materialize back
+  // to f8E8M0FNU, rather than emitting an illegal llvm.load of f8E8M0FNU.
+  //CHECK-LABEL: load_matrix_f8e8m0_scalar
+  gpu.func @load_matrix_f8e8m0_scalar(%arg0: memref<1024xi8, 3>) -> vector<1x1xf8E8M0FNU> {
+    %c0 = arith.constant 0 : index
+    %0 = xegpu.create_mem_desc %arg0 : memref<1024xi8, 3> -> !xegpu.mem_desc<32x32xf8E8M0FNU>
+    %tid_x = gpu.thread_id x
+    //CHECK: %[[LOADED:.*]] = llvm.load %{{.*}} : !llvm.ptr<3> -> i8
+    //CHECK: %[[BCAST:.*]] = arith.bitcast %[[LOADED]] : i8 to f8E8M0FNU
+    //CHECK: vector.broadcast %[[BCAST]] : f8E8M0FNU to vector<1x1xf8E8M0FNU>
+    %1 = xegpu.load_matrix %0[%c0, %tid_x]: !xegpu.mem_desc<32x32xf8E8M0FNU>, index, index -> vector<1x1xf8E8M0FNU>
+    gpu.return %1 : vector<1x1xf8E8M0FNU>
+  }
+
+  //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>
+  }
 
 }

``````````

</details>


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