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

[llvmlistbot at llvm.org]

Author: Sang Ik Lee
Date: 2026-06-17T09:02:22-07:00
New Revision: 5fcb285e3c85eef40398b9a28d5a418b1a048e40

URL: https://github.com/llvm/llvm-project/commit/5fcb285e3c85eef40398b9a28d5a418b1a048e40
DIFF: https://github.com/llvm/llvm-project/commit/5fcb285e3c85eef40398b9a28d5a418b1a048e40.diff

LOG: [MLIR][XeGPU] Fix load_matrix lowering for non-LLVM element types (#203629)

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).

Added: 
    

Modified: 
    mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
    mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir

Removed: 
    


################################################################################
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>
+  }
 
 }