[Mlir-commits] [mlir] [mlir][xegpu] Add support for 1D SLM case in vector-to-xepgu (PR #208276)
Charitha Saumya
llvmlistbot at llvm.org
Fri Jul 10 09:54:05 PDT 2026
https://github.com/charithaintc updated https://github.com/llvm/llvm-project/pull/208276
>From ea777cd2577eb4be81ea0358db31474bce64afa6 Mon Sep 17 00:00:00 2001
From: Charitha Saumya <charitha.saumya.gusthinna.waduge at intel.com>
Date: Wed, 8 Jul 2026 17:48:00 +0000
Subject: [PATCH] add 1D SLM support in vector to xegpu
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 12 ++---
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 44 +++++++++++--------
.../transfer-write-to-xegpu.mlir | 20 ++++++---
3 files changed, 46 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 8dcc9771d8db4..ee9327caab002 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -560,10 +560,10 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
bool isSharedMemory = xegpu::XeGPUDialect::isSharedMemory(readMemTy);
// Handle the SLM case.
if (isSharedMemory) {
- // If the memref is SLM only support 2D case for now.
- if (loadedVecTy.getRank() != 2)
+ // load_matrix supports 1D and 2D loads from SLM.
+ if (loadedVecTy.getRank() != 1 && loadedVecTy.getRank() != 2)
return rewriter.notifyMatchFailure(
- readOp, "Only 2D vector loads are supported for SLM");
+ readOp, "Only 1D and 2D vector loads are supported for SLM");
AffineMap readMap = readOp.getPermutationMap();
if (!readMap.isMinorIdentity())
return rewriter.notifyMatchFailure(
@@ -700,10 +700,10 @@ struct TransferWriteLowering
// For shared local memory (address space 3), use create_mem_desc +
// store_matrix
if (isSharedMemory) {
- // Only support 2D case for now.
- if (vecTy.getRank() != 2)
+ // store_matrix supports 1D and 2D stores to SLM.
+ if (vecTy.getRank() != 1 && vecTy.getRank() != 2)
return rewriter.notifyMatchFailure(
- writeOp, "Only 2D vector stores are supported for SLM");
+ writeOp, "Only 1D and 2D vector stores are supported for SLM");
// Create mem_desc for SLM
auto memDescType =
xegpu::MemDescType::get(rewriter.getContext(), writeMemTy.getShape(),
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index f64bcc27940ea..c6793a2e5ca26 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -598,7 +598,7 @@ gpu.func @load_2D_vector_addrspace3(%source: memref<16x32xf32, 3>,
// -----
gpu.module @xevm_module {
-gpu.func @load_1D_vector_addrspace3_unsupported(%source: memref<32xf32, 3>,
+gpu.func @load_1D_vector_addrspace3(%source: memref<32xf32, 3>,
%offset: index) -> vector<8xf32> {
%c0 = arith.constant 0.0 : f32
%0 = vector.transfer_read %source[%offset], %c0
@@ -606,11 +606,19 @@ gpu.func @load_1D_vector_addrspace3_unsupported(%source: memref<32xf32, 3>,
gpu.return %0 : vector<8xf32>
}
-// LOAD-ND-LABEL: @load_1D_vector_addrspace3_unsupported
-// LOAD-ND: vector.transfer_read
+// LOAD-ND-LABEL: @load_1D_vector_addrspace3
+// LOAD-ND-SAME: %[[SOURCE:.+]]: memref<32xf32, 3>
+// LOAD-ND-SAME: %[[OFFSET:.+]]: index
+// LOAD-ND: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<32xf32, 3> -> !xegpu.mem_desc<32xf32>
+// LOAD-ND: %[[DATA:.+]] = xegpu.load_matrix %[[MEM_DESC]][%[[OFFSET]]] : !xegpu.mem_desc<32xf32>, index -> vector<8xf32>
+// LOAD-ND: gpu.return %[[DATA]] : vector<8xf32>
-// LOAD-GATHER-LABEL: @load_1D_vector_addrspace3_unsupported
-// LOAD-GATHER: vector.transfer_read
+// LOAD-GATHER-LABEL: @load_1D_vector_addrspace3
+// LOAD-GATHER-SAME: %[[SOURCE:.+]]: memref<32xf32, 3>
+// LOAD-GATHER-SAME: %[[OFFSET:.+]]: index
+// LOAD-GATHER: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<32xf32, 3> -> !xegpu.mem_desc<32xf32>
+// LOAD-GATHER: %[[DATA:.+]] = xegpu.load_matrix %[[MEM_DESC]][%[[OFFSET]]] : !xegpu.mem_desc<32xf32>, index -> vector<8xf32>
+// LOAD-GATHER: gpu.return %[[DATA]] : vector<8xf32>
}
@@ -642,10 +650,10 @@ gpu.func @load_2D_vector_alloca_promoted_to_slm(%offset: index)
// -----
// memref.alloca is unconditionally promoted to SLM (address space 3). A 1D
-// transfer_read on it cannot be lowered to xegpu.load_matrix (which requires
-// rank 2), so the transfer_read is left as-is on the SLM memref.
+// transfer_read on it is lowered to xegpu.load_matrix, which supports the
+// 1D SLM case.
gpu.module @xevm_module {
-gpu.func @load_1D_vector_alloca_promoted_unsupported(%offset: index)
+gpu.func @load_1D_vector_alloca_promoted_to_slm(%offset: index)
-> vector<8xf32> {
%buf = memref.alloca() : memref<16xf32>
%c0 = arith.constant 0.0 : f32
@@ -654,17 +662,15 @@ gpu.func @load_1D_vector_alloca_promoted_unsupported(%offset: index)
gpu.return %0 : vector<8xf32>
}
-// LOAD-ND-LABEL: @load_1D_vector_alloca_promoted_unsupported
-// LOAD-ND: memref.alloca() : memref<16xf32, 3>
-// LOAD-ND: vector.transfer_read
-// LOAD-ND-NOT: xegpu.create_mem_desc
-// LOAD-ND-NOT: xegpu.load_matrix
+// LOAD-ND-LABEL: @load_1D_vector_alloca_promoted_to_slm
+// LOAD-ND: %[[BUF:.+]] = memref.alloca() : memref<16xf32, 3>
+// LOAD-ND: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[BUF]] : memref<16xf32, 3> -> !xegpu.mem_desc<16xf32>
+// LOAD-ND: xegpu.load_matrix %[[MEM_DESC]]
-// LOAD-GATHER-LABEL: @load_1D_vector_alloca_promoted_unsupported
-// LOAD-GATHER: memref.alloca() : memref<16xf32, 3>
-// LOAD-GATHER: vector.transfer_read
-// LOAD-GATHER-NOT: xegpu.create_mem_desc
-// LOAD-GATHER-NOT: xegpu.load_matrix
+// LOAD-GATHER-LABEL: @load_1D_vector_alloca_promoted_to_slm
+// LOAD-GATHER: %[[BUF:.+]] = memref.alloca() : memref<16xf32, 3>
+// LOAD-GATHER: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[BUF]] : memref<16xf32, 3> -> !xegpu.mem_desc<16xf32>
+// LOAD-GATHER: xegpu.load_matrix %[[MEM_DESC]]
}
@@ -679,4 +685,4 @@ gpu.func @load_0D_memref_unsupported(%source: memref<f16>) -> vector<f16> {
// CHECK-LABEL: @load_0D_memref_unsupported
// CHECK: vector.transfer_read
-}
\ No newline at end of file
+}
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index a7cd3d7652d85..27a83922fab37 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -397,7 +397,7 @@ gpu.func @store_2D_vector_addrspace3(%vec: vector<8x16xf32>,
// -----
gpu.module @xevm_module {
-gpu.func @store_1D_vector_addrspace3_unsupported(%vec: vector<8xf32>,
+gpu.func @store_1D_vector_addrspace3(%vec: vector<8xf32>,
%source: memref<32xf32, 3>, %offset: index) {
vector.transfer_write %vec, %source[%offset]
{in_bounds = [true]}
@@ -405,10 +405,20 @@ gpu.func @store_1D_vector_addrspace3_unsupported(%vec: vector<8xf32>,
gpu.return
}
-// STORE-ND-LABEL: @store_1D_vector_addrspace3_unsupported
-// STORE-ND: vector.transfer_write
+// STORE-ND-LABEL: @store_1D_vector_addrspace3
+// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf32>
+// STORE-ND-SAME: %[[SOURCE:.+]]: memref<32xf32, 3>
+// STORE-ND-SAME: %[[OFFSET:.+]]: index
+// STORE-ND: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<32xf32, 3> -> !xegpu.mem_desc<32xf32>
+// STORE-ND: xegpu.store_matrix %[[VEC]], %[[MEM_DESC]][%[[OFFSET]]] : vector<8xf32>, !xegpu.mem_desc<32xf32>, index
+// STORE-ND: gpu.return
-// STORE-SCATTER-LABEL: @store_1D_vector_addrspace3_unsupported
-// STORE-SCATTER: vector.transfer_write
+// STORE-SCATTER-LABEL: @store_1D_vector_addrspace3
+// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8xf32>
+// STORE-SCATTER-SAME: %[[SOURCE:.+]]: memref<32xf32, 3>
+// STORE-SCATTER-SAME: %[[OFFSET:.+]]: index
+// STORE-SCATTER: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<32xf32, 3> -> !xegpu.mem_desc<32xf32>
+// STORE-SCATTER: xegpu.store_matrix %[[VEC]], %[[MEM_DESC]][%[[OFFSET]]] : vector<8xf32>, !xegpu.mem_desc<32xf32>, index
+// STORE-SCATTER: gpu.return
}
More information about the Mlir-commits
mailing list