[mlir] Remove the memory space attribute restriction on memref types (#222995)

Follow-up to #187682. `isSupportedMemorySpace` only rejected builtin
attributes other than IntegerAttr/StringAttr/DictionaryAttr. This is a
leftover from when memory spaces were restricted to integers, and is now
inconsistent.

GitOrigin-RevId: 4388400c3d0a95b75a36f81c7edd2eb95c6d63ea

include/mlir/IR/BuiltinTypes.td

diff --git a/include/mlir/IR/BuiltinTypes.td b/include/mlir/IR/BuiltinTypes.td
index d5d53dc..ffb0e54 100644
--- a/include/mlir/IR/BuiltinTypes.td
+++ b/include/mlir/IR/BuiltinTypes.td
@@ -894,7 +894,8 @@
     form which is converted to a semi-affine map automatically.
 
     The memory space of a memref is specified by a target-specific attribute.
-    It might be an integer value, string, dictionary or custom dialect attribute.
+    It might be an integer value, string, dictionary or custom dialect
+    attribute; no restriction is placed on the kind of attribute used.
     The empty memory space (attribute is None) is target specific.
 
     The notionally dynamic value of a memref value includes the address of the

lib/IR/BuiltinTypes.cpp

diff --git a/lib/IR/BuiltinTypes.cpp b/lib/IR/BuiltinTypes.cpp
index 1608b24..36353ed 100644
--- a/lib/IR/BuiltinTypes.cpp
+++ b/lib/IR/BuiltinTypes.cpp
@@ -608,28 +608,6 @@
   return SliceVerificationResult::Success;
 }
 
-bool mlir::detail::isSupportedMemorySpace(Attribute memorySpace) {
-  // Empty attribute is allowed as default memory space.
-  if (!memorySpace)
-    return true;
-
-  // Supported built-in attributes.
-  if (llvm::isa<IntegerAttr, StringAttr, DictionaryAttr>(memorySpace))
-    return true;
-
-  // Allow opaque attributes if unregistered dialects are allowed.
-  // They hold unregistered custom dialect attributes.
-  if (memorySpace.getContext()->allowsUnregisteredDialects() &&
-      isa<OpaqueAttr>(memorySpace))
-    return true;
-
-  // Allow custom dialect attributes.
-  if (!isa<BuiltinDialect>(memorySpace.getDialect()))
-    return true;
-
-  return false;
-}
-
 Attribute mlir::detail::wrapIntegerMemorySpace(unsigned memorySpace,
                                                MLIRContext *ctx) {
   if (memorySpace == 0)
@@ -785,9 +763,6 @@
   if (failed(layout.verifyLayout(shape, emitError)))
     return failure();
 
-  if (!isSupportedMemorySpace(memorySpace))
-    return emitError() << "unsupported memory space Attribute";
-
   return success();
 }
 
@@ -914,9 +889,6 @@
   if (!BaseMemRefType::isValidElementType(elementType))
     return emitError() << "invalid memref element type";
 
-  if (!isSupportedMemorySpace(memorySpace))
-    return emitError() << "unsupported memory space Attribute";
-
   return success();
 }
 

lib/IR/TypeDetail.h

diff --git a/lib/IR/TypeDetail.h b/lib/IR/TypeDetail.h
index 0e952d5..36620ae 100644
--- a/lib/IR/TypeDetail.h
+++ b/lib/IR/TypeDetail.h
@@ -135,9 +135,6 @@
   unsigned numElements;
 };
 
-/// Checks if the memorySpace has supported Attribute type.
-bool isSupportedMemorySpace(Attribute memorySpace);
-
 /// Wraps deprecated integer memory space to the new Attribute form.
 Attribute wrapIntegerMemorySpace(unsigned memorySpace, MLIRContext *ctx);
 

test/IR/invalid-builtin-types.mlir

diff --git a/test/IR/invalid-builtin-types.mlir b/test/IR/invalid-builtin-types.mlir
index ef34124..8207b18 100644
--- a/test/IR/invalid-builtin-types.mlir
+++ b/test/IR/invalid-builtin-types.mlir
@@ -34,10 +34,6 @@
 func.func @memrefs(memref<2x4xi8, #map7>) // expected-error {{undefined symbol alias id 'map7'}}
 
 // -----
-// Test unsupported memory space.
-func.func @memrefs(memref<2x4xi8, i8>) // expected-error {{unsupported memory space Attribute}}
-
-// -----
 // Test non-existent map in map composition of memref type.
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 

test/IR/parser.mlir

diff --git a/test/IR/parser.mlir b/test/IR/parser.mlir
index 6cd07f6..f3f6378 100644
--- a/test/IR/parser.mlir
+++ b/test/IR/parser.mlir
@@ -133,6 +133,12 @@
 // CHECK: func private @memrefs_map_opaquespace(memref<5x6x7xf32, #map{{[0-9]*}}, #unknown_dialect.unknown_attr>)
 func.func private @memrefs_map_opaquespace(memref<5x6x7xf32, #map3, #unknown_dialect.unknown_attr>)
 
+// CHECK: func private @memrefs_nomap_typespace(memref<5x6x7xf32, i8>)
+func.func private @memrefs_nomap_typespace(memref<5x6x7xf32, i8>)
+
+// CHECK: func private @memrefs_map_typespace(memref<5x6x7xf32, #map{{[0-9]*}}, i8>)
+func.func private @memrefs_map_typespace(memref<5x6x7xf32, #map3, i8>)
+
 // CHECK: func private @complex_types(complex<i1>) -> complex<f32>
 func.func private @complex_types(complex<i1>) -> complex<f32>