[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 GitOrigin-RevId: 286d5c760a7710644ad06f354baa69cec5cc7b72
diff --git a/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp index 6144d7c..78d99cf 100644 --- a/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp +++ b/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -67,16 +67,42 @@ 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 @@ 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 @@ 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/test/Conversion/XeGPUToXeVM/failed_conversion.mlir b/test/Conversion/XeGPUToXeVM/failed_conversion.mlir index 95211dc..cabc65a 100644 --- a/test/Conversion/XeGPUToXeVM/failed_conversion.mlir +++ b/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
@@ -12,3 +12,19 @@ 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@+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 + } +}