[Mlir-commits] [mlir] [MLIR][XeGPU] Fix XeGPUToXeVM crash on non-integer memref memory spaces (PR #211053)
Hamza Qureshi
llvmlistbot at llvm.org
Wed Jul 29 07:44:05 PDT 2026
https://github.com/hamzaqureshi5 updated https://github.com/llvm/llvm-project/pull/211053
>From b36650073b9be2827d4f1386c78855d6bad639a8 Mon Sep 17 00:00:00 2001
From: Hamza Qureshi <hamza7771.861 at gmail.com>
Date: Tue, 21 Jul 2026 22:02:32 +0500
Subject: [PATCH] [MLIR][XeGPU] Fix XeGPUToXeVM crash on non-integer memref
memory spaces
LoadStoreToXeVMPattern and PrefetchToXeVMPattern computed the LLVM
pointer address space for load/store/prefetch ops via the deprecated
MemRefType::getMemorySpaceAsInt(), which asserts unless the memory
space is an IntegerAttr. Memref memory spaces are not required to be
integers (e.g. SPIR-V storage classes, GPU dialect address spaces),
so a memref such as memref<1024xf32, #spirv.storage_class<StorageBuffer>>
crashed -convert-xegpu-to-xevm instead of being handled or rejected
cleanly.
Add getNumericMemorySpace(), which maps IntegerAttr, xevm::AddrSpaceAttr,
gpu::AddressSpaceAttr, and spirv::StorageClassAttr memory spaces onto
XeVM's address space numbering (which follows the OpenCL/SPIR-V
convention), and fails cleanly for anything else so callers can
notifyMatchFailure instead of asserting. isSharedMemRef is rewritten
in terms of this helper for consistency.
---
.../lib/Conversion/XeGPUToXeVM/CMakeLists.txt | 2 +
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 77 +++++++++++++++----
.../XeGPUToXeVM/failed_conversion.mlir | 16 ++++
.../XeGPUToXeVM/loadstoreprefetch.mlir | 35 +++++++++
4 files changed, 113 insertions(+), 17 deletions(-)
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt b/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
index dd9edc43a1657..317b5b38f1c30 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
+++ b/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
@@ -15,6 +15,8 @@ add_mlir_conversion_library(MLIRXeGPUToXeVM
MLIRGPUDialect
MLIRLLVMCommonConversion
MLIRLLVMDialect
+ MLIRSPIRVAttrToLLVMConversion
+ MLIRSPIRVDialect
MLIRXeVMDialect
MLIRVectorDialect
MLIRArithDialect
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 6144d7c0c1a15..6f0ee24ecbfd0 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -11,6 +11,7 @@
#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Conversion/SPIRVCommon/AttrToLLVMConverter.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
@@ -20,6 +21,7 @@
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
@@ -67,16 +69,45 @@ static int32_t getNumericXeVMAddrSpace(xegpu::MemorySpace xeGpuMemspace) {
llvm_unreachable("Unknown XeGPU memory space");
}
+/// Translates a memref memory space attribute into XeVM's numeric address
+/// space, which follows the OpenCL/SPIR-V convention (0 = private, 1 =
+/// global, 2 = constant, 3 = shared/local, 4 = generic). A null attribute,
+/// meaning the memory space was left unspecified, maps to the default space
+/// 0. Returns failure if `memSpace` is a representation this pass does not
+/// know how to translate (e.g. an arbitrary string or dictionary attribute),
+/// rather than assuming it is an `IntegerAttr` and asserting.
+static FailureOr<unsigned> getNumericMemorySpace(Attribute memSpace) {
+ if (!memSpace)
+ return 0u;
+ if (auto intAttr = llvm::dyn_cast<IntegerAttr>(memSpace))
+ return static_cast<unsigned>(intAttr.getInt());
+ if (auto xevmSpace = llvm::dyn_cast<xevm::AddrSpaceAttr>(memSpace))
+ return static_cast<unsigned>(xevmSpace.getValue());
+ if (auto gpuSpace = llvm::dyn_cast<gpu::AddressSpaceAttr>(memSpace)) {
+ switch (gpuSpace.getValue()) {
+ case gpu::AddressSpace::Global:
+ return static_cast<unsigned>(xevm::AddrSpace::GLOBAL);
+ case gpu::AddressSpace::Workgroup:
+ return static_cast<unsigned>(xevm::AddrSpace::SHARED);
+ case gpu::AddressSpace::Private:
+ return static_cast<unsigned>(xevm::AddrSpace::PRIVATE);
+ case gpu::AddressSpace::Constant:
+ return static_cast<unsigned>(xevm::AddrSpace::CONSTANT);
+ }
+ llvm_unreachable("Unknown GPU address space");
+ }
+ if (auto storageClass = llvm::dyn_cast<spirv::StorageClassAttr>(memSpace))
+ return storageClassToAddressSpace(spirv::ClientAPI::OpenCL,
+ storageClass.getValue());
+ return failure();
+}
+
/// Checks if the given MemRefType refers to shared memory.
static bool isSharedMemRef(const MemRefType &memrefTy) {
- Attribute attr = memrefTy.getMemorySpace();
- if (!attr)
- return false;
- if (auto intAttr = llvm::dyn_cast<IntegerAttr>(attr))
- return intAttr.getInt() == static_cast<int>(xevm::AddrSpace::SHARED);
- if (auto xevmSpace = llvm::dyn_cast<xevm::AddrSpaceAttr>(attr))
- return xevmSpace.getValue() == xevm::AddrSpace::SHARED;
- return gpu::GPUDialect::isWorkgroupMemoryAddressSpace(attr);
+ FailureOr<unsigned> addrSpace =
+ getNumericMemorySpace(memrefTy.getMemorySpace());
+ return succeeded(addrSpace) &&
+ *addrSpace == static_cast<unsigned>(xevm::AddrSpace::SHARED);
}
// Get same bitwidth flat vector type of new element type.
@@ -592,16 +623,24 @@ class LoadStoreToXeVMPattern : public OpConversionPattern<OpType> {
if constexpr (std::is_same_v<OpType, xegpu::LoadGatherOp>) {
basePtrI64 = adaptor.getSource();
if (auto memRefTy = dyn_cast<MemRefType>(op.getSource().getType())) {
- auto addrSpace = memRefTy.getMemorySpaceAsInt();
- if (addrSpace != 0)
- ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, addrSpace);
+ FailureOr<unsigned> addrSpace =
+ getNumericMemorySpace(memRefTy.getMemorySpace());
+ if (failed(addrSpace))
+ return rewriter.notifyMatchFailure(
+ op, "Unsupported memref memory space attribute.");
+ if (*addrSpace != 0)
+ ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, *addrSpace);
}
} else {
basePtrI64 = adaptor.getDest();
if (auto memRefTy = dyn_cast<MemRefType>(op.getDest().getType())) {
- auto addrSpace = memRefTy.getMemorySpaceAsInt();
- if (addrSpace != 0)
- ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, addrSpace);
+ FailureOr<unsigned> addrSpace =
+ getNumericMemorySpace(memRefTy.getMemorySpace());
+ if (failed(addrSpace))
+ return rewriter.notifyMatchFailure(
+ op, "Unsupported memref memory space attribute.");
+ if (*addrSpace != 0)
+ ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, *addrSpace);
}
}
// Base pointer is passed as i32 or i64 by adaptor, cast to i64 if needed.
@@ -859,9 +898,13 @@ class PrefetchToXeVMPattern : public OpConversionPattern<xegpu::PrefetchOp> {
ctxt, getNumericXeVMAddrSpace(xegpu::MemorySpace::Global));
// If source is a memref, we use its memory space.
if (auto memRefTy = dyn_cast<MemRefType>(op.getSource().getType())) {
- auto addrSpace = memRefTy.getMemorySpaceAsInt();
- if (addrSpace != 0)
- ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, addrSpace);
+ FailureOr<unsigned> addrSpace =
+ getNumericMemorySpace(memRefTy.getMemorySpace());
+ if (failed(addrSpace))
+ return rewriter.notifyMatchFailure(
+ op, "Unsupported memref memory space attribute.");
+ if (*addrSpace != 0)
+ ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, *addrSpace);
}
// Convert base pointer (i64) to LLVM pointer type.
Value ptrLLVM =
diff --git a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
index 95211dcff250c..3e41c6b8c9b32 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
@@ -12,3 +12,19 @@ gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
return
}
}
+
+// -----
+
+// Verify that xegpu.store with a memref memory space attribute that has no
+// known numeric address space (a bare string, here) is rejected during
+// XeGPUToXeVM conversion rather than crashing.
+
+gpu.module @test_kernel {
+ gpu.func @store_scatter_unsupported_memspace(%src: memref<1024xf32, "foo">, %offset: vector<1xindex>, %mask: vector<1xi1>) {
+ %0 = arith.constant dense<2.9> : vector<1xf32>
+ // expected-error at +1 {{failed to legalize operation 'xegpu.store' that was explicitly marked illegal}}
+ xegpu.store %0, %src[%offset], %mask <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : vector<1xf32>, memref<1024xf32, "foo">, vector<1xindex>, vector<1xi1>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir
index 39be929978d1e..680d7b5a7acff 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir
@@ -135,3 +135,38 @@ gpu.func @load_gather_from_dyn_memref_subview(%dyn: memref<?xf16>, %offset: vect
gpu.return
}
}
+
+// -----
+
+// A memref memory space is not always an IntegerAttr: SPIR-V storage classes
+// (and other dialect attributes) are also legal memory spaces. Storage
+// classes with no known numeric address space mapping (e.g. StorageBuffer)
+// fall back to the default global address space instead of asserting.
+gpu.module @test {
+// CHECK-LABEL: @store_scatter_spirv_default_memspace
+gpu.func @store_scatter_spirv_default_memspace(%src: memref<1024xf32, #spirv.storage_class<StorageBuffer>>, %offset: vector<1xindex>, %mask: vector<1xi1>) {
+ %0 = arith.constant dense<2.9>: vector<1xf32>
+ // CHECK: %[[PTR:.*]] = llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<1>
+ // CHECK: llvm.store %{{.*}}, %[[PTR]] {{.*}} : f32, !llvm.ptr<1>
+ xegpu.store %0, %src[%offset], %mask <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : vector<1xf32>, memref<1024xf32, #spirv.storage_class<StorageBuffer>>, vector<1xindex>, vector<1xi1>
+ gpu.return
+}
+}
+
+// -----
+
+// SPIR-V storage classes that do have a known mapping (here Workgroup, SPIR-V's
+// local/shared memory class) translate to the matching XeVM address space
+// (3 = shared), not just the default.
+gpu.module @test {
+// CHECK-LABEL: @store_scatter_spirv_workgroup_memspace
+gpu.func @store_scatter_spirv_workgroup_memspace(%src: memref<1024xf32, #spirv.storage_class<Workgroup>>, %offset: vector<1xindex>, %mask: vector<1xi1>) {
+ %0 = arith.constant dense<2.9>: vector<1xf32>
+ // CHECK: %[[PTR:.*]] = llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<3>
+ // CHECK: llvm.store %{{.*}}, %[[PTR]] {{.*}} : f32, !llvm.ptr<3>
+ xegpu.store %0, %src[%offset], %mask <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : vector<1xf32>, memref<1024xf32, #spirv.storage_class<Workgroup>>, vector<1xindex>, vector<1xi1>
+ gpu.return
+}
+}
More information about the Mlir-commits
mailing list