[Mlir-commits] [mlir] [NVGPU] Fix crash on ui64 memory space attributes (#204580) (PR #204676)

[llvmlistbot at llvm.org]

https://github.com/Naaz30 updated https://github.com/llvm/llvm-project/pull/204676

>From b37f32cb23e70b7631633df66ee139e4c8c2963f Mon Sep 17 00:00:00 2001
From: Naazni Yahya <naazniyahya at Naaznis-MacBook-Air.local>
Date: Fri, 19 Jun 2026 01:24:42 +0530
Subject: [PATCH] [NVGPU] Fix crash on ui64 memory space attributes (#204580)

---
 mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp         |  3 ++-
 .../test/Dialect/NVGPU/optimize-shared-memory.mlir | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
index 6e1ed05c9d0d0..b53d95e92d7d1 100644
--- a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
+++ b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
@@ -50,7 +50,8 @@ bool NVGPUDialect::isSharedMemoryAddressSpace(Attribute memorySpace) {
   if (!memorySpace)
     return false;
   if (auto intAttr = llvm::dyn_cast<IntegerAttr>(memorySpace))
-    return intAttr.getInt() == NVGPUDialect::kSharedMemoryAddressSpace;
+    return intAttr.getValue().getZExtValue() ==
+         NVGPUDialect::kSharedMemoryAddressSpace;
   if (auto gpuAttr = llvm::dyn_cast<gpu::AddressSpaceAttr>(memorySpace))
     return gpuAttr.getValue() == gpu::AddressSpace::Workgroup;
   return false;
diff --git a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
index 128c4c5d918d9..e6c31ddb3f856 100644
--- a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
+++ b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
@@ -270,3 +270,17 @@ func.func @test_vector_element_type() {
   %alloc = memref.alloc() : memref<16x1xvector<16xf16>, 3>
   return
 }
+
+
+// -----
+
+/// Ensure ui64 memory space attributes do not crash
+// nvgpu-optimize-shared-memory (issue #204580).
+
+// CHECK-LABEL: func @test_int_conversion
+module {
+  func.func @test_int_conversion() {
+    %alloc = memref.alloc() {alignment = 64 : i64} : memref<10xf32, 1 : ui64>
+    return
+  }
+}