[Mlir-commits] [mlir] [MLIR][XeGPU] Promote memref.alloca to SLM in convert-vector-to-xegpu (PR #197978)
Nishant Patel
llvmlistbot at llvm.org
Thu Jun 4 13:53:03 PDT 2026
https://github.com/nbpatel updated https://github.com/llvm/llvm-project/pull/197978
>From 8780c570183368cb074530b97b3fe696f017d02d Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 28 Apr 2026 18:55:22 +0000
Subject: [PATCH 1/2] Promote local memref allocations to SLM
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 130 ++++++++++++++++++
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 53 +++++++
2 files changed, 183 insertions(+)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 0aae5cb5bb6ad..60bf4f850cf18 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -928,9 +928,139 @@ struct ContractionLowering : public OpRewritePattern<vector::ContractionOp> {
}
};
+// Returns `memrefTy` with its memory space replaced by `newMemSpace`.
+static MemRefType withMemorySpace(MemRefType memrefTy, Attribute newMemSpace) {
+ return MemRefType::get(memrefTy.getShape(), memrefTy.getElementType(),
+ memrefTy.getLayout(), newMemSpace);
+}
+
+// For each `memref.alloca`/`memref.alloc` whose result is not already in shared
+// local memory (SLM) and which is consumed by a `vector.transfer_read` or
+// `vector.transfer_write` that would lower to `xegpu.load_matrix` /
+// `xegpu.store_matrix`, rewrite the allocation to be in SLM (address space 3)
+// and propagate the new memory space through any memref-producing users
+// (memref.cast, memref.subview, memref.expand_shape, memref.collapse_shape,
+// memref.reinterpret_cast, memref.transpose, memref.view). Consumers that take
+// a memref operand but produce a non-memref result (e.g. vector.transfer_read,
+// vector.load) are left untouched: their operand type simply reflects the new
+// memory space.
+//
+// This makes `xegpu.load_matrix`/`xegpu.store_matrix` lowering work end-to-end
+// for IR coming from bufferization, which by default assigns memory space 0/1
+// to allocations. See discussion in the XeGPU lighthouse "softmax" issue.
+static void promoteAllocasToSLM(Operation *root) {
+ MLIRContext *ctx = root->getContext();
+ Attribute slmAttr = IntegerAttr::get(IntegerType::get(ctx, 64), 3);
+
+ auto isMemrefResultOp = [](Operation *op) {
+ return isa<memref::CastOp, memref::SubViewOp, memref::ExpandShapeOp,
+ memref::CollapseShapeOp, memref::ReinterpretCastOp,
+ memref::TransposeOp, memref::ViewOp>(op);
+ };
+
+ // Returns true if `xferOp` (a vector.transfer_read or vector.transfer_write)
+ // would lower to `xegpu.load_matrix`/`xegpu.store_matrix`: must operate on a
+ // 2D vector with a minor-identity permutation map and no out-of-bounds dims.
+ auto isSLMMatrixTransfer = [](Operation *xferOp) {
+ VectorType vecTy;
+ AffineMap permMap;
+ bool hasOOB = false;
+ if (auto rd = dyn_cast<vector::TransferReadOp>(xferOp)) {
+ vecTy = rd.getVectorType();
+ permMap = rd.getPermutationMap();
+ hasOOB = rd.hasOutOfBoundsDim();
+ } else if (auto wr = dyn_cast<vector::TransferWriteOp>(xferOp)) {
+ vecTy = wr.getVectorType();
+ permMap = wr.getPermutationMap();
+ hasOOB = wr.hasOutOfBoundsDim();
+ } else {
+ return false;
+ }
+ return vecTy && vecTy.getRank() == 2 && permMap.isMinorIdentity() &&
+ !hasOOB;
+ };
+
+ // Returns true if `v` (a memref) is transitively consumed by a
+ // transfer_read/transfer_write that would lower to load_matrix/store_matrix.
+ // Walks forward through memref-producing aliasing ops.
+ std::function<bool(Value, llvm::SmallPtrSetImpl<Operation *> &)> feedsMatrix =
+ [&](Value v, llvm::SmallPtrSetImpl<Operation *> &visited) -> bool {
+ for (Operation *user : v.getUsers()) {
+ if (!visited.insert(user).second)
+ continue;
+ if (isSLMMatrixTransfer(user))
+ return true;
+ if (isMemrefResultOp(user)) {
+ for (Value result : user->getResults())
+ if (feedsMatrix(result, visited))
+ return true;
+ }
+ }
+ return false;
+ };
+
+ // Update `v`'s type to have SLM memory space, then walk forward through
+ // memref-producing users and update their result types accordingly.
+ std::function<void(Value)> propagate = [&](Value v) {
+ auto memrefTy = dyn_cast<MemRefType>(v.getType());
+ if (!memrefTy || xegpu::XeGPUDialect::isSharedMemory(memrefTy))
+ return;
+ v.setType(withMemorySpace(memrefTy, slmAttr));
+ for (Operation *user : v.getUsers()) {
+ if (!isMemrefResultOp(user))
+ continue;
+ for (Value result : user->getResults())
+ propagate(result);
+ }
+ };
+
+ SmallVector<Operation *> allocs;
+ root->walk([&](Operation *op) {
+ if (!isa<memref::AllocaOp, memref::AllocOp>(op))
+ return;
+ auto memrefTy = dyn_cast<MemRefType>(op->getResult(0).getType());
+ if (!memrefTy || xegpu::XeGPUDialect::isSharedMemory(memrefTy))
+ return;
+ llvm::SmallPtrSet<Operation *, 8> visited;
+ if (!feedsMatrix(op->getResult(0), visited))
+ return;
+ allocs.push_back(op);
+ });
+
+ for (Operation *op : allocs) {
+ OpBuilder builder(op);
+ auto memrefTy = cast<MemRefType>(op->getResult(0).getType());
+ auto newTy = withMemorySpace(memrefTy, slmAttr);
+ Operation *newOp;
+ if (auto alloca = dyn_cast<memref::AllocaOp>(op)) {
+ newOp = memref::AllocaOp::create(
+ builder, alloca.getLoc(), newTy, alloca.getDynamicSizes(),
+ alloca.getSymbolOperands(), alloca.getAlignmentAttr());
+ } else {
+ auto alloc = cast<memref::AllocOp>(op);
+ newOp = memref::AllocOp::create(
+ builder, alloc.getLoc(), newTy, alloc.getDynamicSizes(),
+ alloc.getSymbolOperands(), alloc.getAlignmentAttr());
+ }
+ op->getResult(0).replaceAllUsesWith(newOp->getResult(0));
+ op->erase();
+ // Propagate the new memory space through memref-producing consumers.
+ for (Operation *user : newOp->getResult(0).getUsers()) {
+ if (!isMemrefResultOp(user))
+ continue;
+ for (Value result : user->getResults())
+ propagate(result);
+ }
+ }
+}
+
struct ConvertVectorToXeGPUPass
: public impl::ConvertVectorToXeGPUBase<ConvertVectorToXeGPUPass> {
void runOnOperation() override {
+ // Promote local allocations to SLM (address space 3) so that
+ // load_matrix/store_matrix lowerings have well-typed memref operands.
+ promoteAllocasToSLM(getOperation());
+
RewritePatternSet patterns(&getContext());
populateVectorToXeGPUConversionPatterns(patterns);
populatePrepareVectorToMMAPatterns(patterns);
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index 52b47cb2111c7..d0b94ccb9f6ac 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -573,3 +573,56 @@ gpu.func @load_1D_vector_addrspace3_unsupported(%source: memref<32xf32, 3>,
// LOAD-GATHER: vector.transfer_read
}
+
+// -----
+// memref.alloca with the default address space is promoted to SLM
+// (address space 3) so that the transfer_read can be lowered to
+// xegpu.load_matrix.
+gpu.module @xevm_module {
+gpu.func @load_2D_vector_alloca_promoted_to_slm(%offset: index)
+ -> vector<8x16xf32> {
+ %buf = memref.alloca() : memref<16x32xf32>
+ %c0 = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %buf[%offset, %offset], %c0
+ {in_bounds = [true, true]} : memref<16x32xf32>, vector<8x16xf32>
+ gpu.return %0 : vector<8x16xf32>
+}
+
+// LOAD-ND-LABEL: @load_2D_vector_alloca_promoted_to_slm
+// LOAD-ND: %[[BUF:.+]] = memref.alloca() : memref<16x32xf32, 3>
+// LOAD-ND: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[BUF]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
+// LOAD-ND: xegpu.load_matrix %[[MEM_DESC]]
+
+// LOAD-GATHER-LABEL: @load_2D_vector_alloca_promoted_to_slm
+// LOAD-GATHER: %[[BUF:.+]] = memref.alloca() : memref<16x32xf32, 3>
+// LOAD-GATHER: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[BUF]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
+// LOAD-GATHER: xegpu.load_matrix %[[MEM_DESC]]
+
+}
+
+// -----
+// memref.alloca consumed only by a 1D transfer_read is NOT promoted to SLM,
+// because that transfer would not lower to xegpu.load_matrix.
+gpu.module @xevm_module {
+gpu.func @load_1D_vector_alloca_not_promoted(%offset: index)
+ -> vector<8xf32> {
+ %buf = memref.alloca() : memref<16xf32>
+ %c0 = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %buf[%offset], %c0
+ {in_bounds = [true]} : memref<16xf32>, vector<8xf32>
+ gpu.return %0 : vector<8xf32>
+}
+
+// LOAD-ND-LABEL: @load_1D_vector_alloca_not_promoted
+// LOAD-ND: memref.alloca() : memref<16xf32>
+// LOAD-ND-NOT: memref<16xf32, 3>
+// LOAD-ND-NOT: xegpu.create_mem_desc
+// LOAD-ND-NOT: xegpu.load_matrix
+
+// LOAD-GATHER-LABEL: @load_1D_vector_alloca_not_promoted
+// LOAD-GATHER: memref.alloca() : memref<16xf32>
+// LOAD-GATHER-NOT: memref<16xf32, 3>
+// LOAD-GATHER-NOT: xegpu.create_mem_desc
+// LOAD-GATHER-NOT: xegpu.load_matrix
+
+}
>From 63664c34b83d279701be67720eb2246ee5f07455 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Thu, 21 May 2026 20:37:06 +0000
Subject: [PATCH 2/2] Address feedback
---
.../Conversion/VectorToXeGPU/VectorToXeGPU.cpp | 17 +++++++++++------
1 file changed, 11 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 3de4d765eae79..9038bf35b6b15 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -20,6 +20,7 @@
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/TypeSwitch.h"
@@ -954,9 +955,8 @@ static MemRefType withMemorySpace(MemRefType memrefTy, Attribute newMemSpace) {
// Rewrite every `memref.alloca` not already in shared local memory (SLM) to
// be in SLM (address space 3), and propagate the new memory space through
-// memref-producing aliasing users (memref.cast, memref.subview,
-// memref.expand_shape, memref.collapse_shape, memref.reinterpret_cast,
-// memref.transpose, memref.view). Consumers that take a memref operand but
+// memref-producing aliasing users (e.g. memref.cast, memref.subview,
+// memref.expand_shape, ...). Consumers that take a memref operand but
// produce a non-memref result (e.g. vector.transfer_read, vector.load) are
// left untouched: their operand type simply reflects the new memory space.
//
@@ -967,10 +967,15 @@ static void promoteAllocasToSLM(Operation *root) {
MLIRContext *ctx = root->getContext();
Attribute slmAttr = IntegerAttr::get(IntegerType::get(ctx, 64), 3);
+ // A user is treated as a memref-producing alias (e.g. memref.cast,
+ // memref.subview, memref.expand_shape, ...) if it is side-effect free and
+ // produces at least one memref result. This excludes ops like memref.copy
+ // that have memory effects.
auto isMemrefResultOp = [](Operation *op) {
- return isa<memref::CastOp, memref::SubViewOp, memref::ExpandShapeOp,
- memref::CollapseShapeOp, memref::ReinterpretCastOp,
- memref::TransposeOp, memref::ViewOp>(op);
+ if (!isMemoryEffectFree(op))
+ return false;
+ return llvm::any_of(op->getResultTypes(),
+ [](Type t) { return isa<MemRefType>(t); });
};
// Update `v`'s type to have SLM memory space, then walk forward through
More information about the Mlir-commits
mailing list