[Mlir-commits] [mlir] 286d5c7 - [MLIR][XeGPU] Fix XeGPUToXeVM crash on non-integer memref memory spaces (#211053)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 5 11:51:15 PDT 2026
Author: Hamza Qureshi
Date: 2026-08-05T11:51:10-07:00
New Revision: 286d5c760a7710644ad06f354baa69cec5cc7b72
URL: https://github.com/llvm/llvm-project/commit/286d5c760a7710644ad06f354baa69cec5cc7b72
DIFF: https://github.com/llvm/llvm-project/commit/286d5c760a7710644ad06f354baa69cec5cc7b72.diff
LOG: [MLIR][XeGPU] Fix XeGPUToXeVM crash on non-integer memref memory spaces (#211053)
## Summary
`convert-xegpu-to-xevm` asserts/crashes when a memref's memory space
isn't an
`IntegerAttr` — e.g. `memref<1024xf32,
#spirv.storage_class<StorageBuffer>>`.
`LoadStoreToXeVMPattern` and `PrefetchToXeVMPattern` compute the LLVM
pointer
address space via the deprecated `MemRefType::getMemorySpaceAsInt()`,
which
asserts on anything but an integer. Memref memory spaces aren't
integer-only
in general (GPU address spaces and XeVM's own `xevm::AddrSpaceAttr` are
both
legal here) — the file's own `isSharedMemRef` already handles this
correctly
elsewhere, these three call sites just didn't.
## Fix
Add `getNumericMemorySpace()`: maps `IntegerAttr` (unchanged, verbatim),
`xevm::AddrSpaceAttr`, and `gpu::AddressSpaceAttr` onto XeVM's numeric
address spaces, and returns `failure()` for anything else so callers
reject
the op via `notifyMatchFailure` instead of asserting. `isSharedMemRef`
now
just calls this helper instead of duplicating the attribute dispatch.
SPIR-V storage classes are deliberately not translated: SPIR-V is a
lowering
target for XeGPU, not a memory space meant to be used alongside the
XeGPU
dialect, so `#spirv.storage_class<...>` falls through to `failure()` and
produces a clean legalization error.
Not using `TypeConverter::addTypeAttributeConversion` since this pass
overrides `MemRefType` conversion to a bare integer r than an LLVM
struct descriptor, so that hook is never consulted.
## Test plan
- The reported repro no longer crashes — it now fails legalization
cleanly
- `failed_conversion.mlir`: added a `#spirv.storage_class<StorageBuff
case, verifying a legalization failure instead of an assert
- Full XeGPUToXeVM test suite passes
Fixes #210988
Added:
Modified:
mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 6144d7c0c1a15..78d99cf88b768 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -67,16 +67,42 @@ 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. 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;
+ 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");
+ }
+ 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 +618,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 +893,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..cabc65aa0e41d 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 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, #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, #spirv.storage_class<StorageBuffer>>, vector<1xindex>, vector<1xi1>
+ gpu.return
+ }
+}
More information about the Mlir-commits
mailing list