[Mlir-commits] [mlir] [MLIR][XeGPU] Fix XeGPUToXeVM crash on non-integer memref memory spaces (PR #211053)
Hamza Qureshi
llvmlistbot at llvm.org
Thu Jul 30 07:32:09 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 1/2] [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
+}
+}
>From 823aa66b10f50156ff8322a02830734aff5bcc06 Mon Sep 17 00:00:00 2001
From: Hamza Qureshi <hamza7771.861 at gmail.com>
Date: Thu, 30 Jul 2026 18:02:13 +0500
Subject: [PATCH 2/2] Address review: drop SPIR-V storage class handling
SPIR-V can be a lowering target for XeGPU but is not supported as a
memory space alongside the XeGPU dialect, so spirv::StorageClassAttr
should fail conversion rather than be mapped to an address space.
- Remove the spirv::StorageClassAttr branch from getNumericMemorySpace()
and the two SPIR-V includes / CMake link deps it needed
- Drop the SPIR-V storage class test cases from loadstoreprefetch.mlir
- Use #spirv.storage_class<StorageBuffer> as the unsupported memory
space in failed_conversion.mlir
---
.../lib/Conversion/XeGPUToXeVM/CMakeLists.txt | 2 --
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 9 ++---
.../XeGPUToXeVM/failed_conversion.mlir | 8 ++---
.../XeGPUToXeVM/loadstoreprefetch.mlir | 35 -------------------
4 files changed, 6 insertions(+), 48 deletions(-)
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt b/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
index 317b5b38f1c30..dd9edc43a1657 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
+++ b/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
@@ -15,8 +15,6 @@ 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 6f0ee24ecbfd0..78d99cf88b768 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -11,7 +11,6 @@
#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"
@@ -21,7 +20,6 @@
#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"
@@ -74,8 +72,8 @@ static int32_t getNumericXeVMAddrSpace(xegpu::MemorySpace xeGpuMemspace) {
/// 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.
+/// know how to translate (e.g. a SPIR-V storage class or an arbitrary string
+/// attribute), rather than assuming it is an `IntegerAttr` and asserting.
static FailureOr<unsigned> getNumericMemorySpace(Attribute memSpace) {
if (!memSpace)
return 0u;
@@ -96,9 +94,6 @@ static FailureOr<unsigned> getNumericMemorySpace(Attribute memSpace) {
}
llvm_unreachable("Unknown GPU address space");
}
- if (auto storageClass = llvm::dyn_cast<spirv::StorageClassAttr>(memSpace))
- return storageClassToAddressSpace(spirv::ClientAPI::OpenCL,
- storageClass.getValue());
return failure();
}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
index 3e41c6b8c9b32..cabc65aa0e41d 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
@@ -16,15 +16,15 @@ gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
// -----
// 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.
+// known numeric address space (a SPIR-V storage class, 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>) {
+ gpu.func @store_scatter_unsupported_memspace(%src: memref<1024xf32, #spirv.storage_class<StorageBuffer>>, %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>
+ : vector<1xf32>, memref<1024xf32, #spirv.storage_class<StorageBuffer>>, vector<1xindex>, vector<1xi1>
gpu.return
}
}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir
index 680d7b5a7acff..39be929978d1e 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir
@@ -135,38 +135,3 @@ 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