[mlir] Model MemRef memory space as Attribute

[vinograd47]

Summary

Based on the following discussion:
https://llvm.discourse.group/t/rfc-memref-memory-shape-as-attribute/2229

The goal of the change is to make memory space property to have more
expressive representation, rather then "magic" integer values.

It will allow to have more clean ASM form:

gpu.func @test(%arg0: memref<100xf32, "workgroup">)

// instead of

gpu.func @test(%arg0: memref<100xf32, 3>)

Explanation for Attribute choice instead of plain string:

• Attribute classes allow to use more type safe API based on RTTI.
• Attribute classes provides faster comparison operator based on pointer comparison in contrast to generic string comparison.
• Attribute allows to store more complex things, like structs or dictionaries. It will allows to have more complex memory space hierarchy.

This commit preserve old integer-based API and implements it on top
of the new one.

Depends on D97476

[mehdi_amini]

mlir/include/mlir/IR/BuiltinTypes.h

305

What is the plan to get rid of the integer API? What prevents us from removing it right now?

Also wouldn't the natural name for this API be getMemorySpace? Will you rename getMemorySpaceAttr into getMemorySpace "at some point"?

[rriddle]

mlir/include/mlir/IR/BuiltinTypes.h

12

Please avoid including this here, use forward declarations instead.

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

1806

Add a comment to this funciton.

mlir/lib/IR/BuiltinTypes.cpp

478

nit: Spell out auto here.

479

nit: !attr

632

Use emitOptionalError here.

681

nit: memorySpace &&

681

Drop the trivial braces here.

mlir/lib/Parser/TypeParser.cpp

244

nit: Spell out auto here.

248–255

[bondhugula]

Please remove RFC from the commit title. Also,
"MemRef memory space as Attribute" -> "Model MemRef memory space as Attribute" or "Switch MemRef memory space: unsigned -> Attribute".

mlir/lib/IR/BuiltinTypes.cpp

513–524

All static methods should have doc comments as well.

630

Nit: if (memorySpace && ...

682

Everything else can be?! (For eg., I would add IntegerSet to this prohibited list as well. )

mlir/lib/Parser/TypeParser.cpp

243–245

Please end all comments with a period. Here, above, and below.

[vinograd47]

Rebased and fix review remarks

[vinograd47]

@rriddle @mehdi_amini @bondhugula I've updated the change according to your comments.

mlir/include/mlir/IR/BuiltinTypes.h

12

I've replaced BuiltInAttributes.h header with Attributes.h. I can't replace it with forward declaration for Attribute class, because it is used as class member for MemRef::Builder.

305

The simple replacement of unsigned with Attribute breaks existing code in various dialects (std, affine, gpu). It might also break separate projects (like TensorFlow). As far as I understand, MLIR tries to avoid such breaking changes in single commit and handle compatibility with old code for some period of time to allow dependent projects to switch to the new API. I've followed this approach.

mlir/lib/IR/BuiltinTypes.cpp

682

I've updated this part. Now only IntegerAttr, StringAttr, DictionaryAttr and custom dialect attributes are allowed.

[mehdi_amini]

mlir/include/mlir/IR/BuiltinTypes.h

305

As far as I understand, MLIR tries to avoid such breaking changes in single commit and handle compatibility with old code for some period of time to allow dependent projects to switch to the new API.

Actually this is explicitly a non-goal: https://mlir.llvm.org/getting_started/DeveloperGuide/#breaking-changes

Anyway, I think my questions about naming and plan still stand even if you deemed easier to update the dialects in-tree in multiple steps.

[bondhugula]

Looks fine for the minor comments part I had.

[vinograd47]

Removed old integer-based API from MemRefType. Updated internal dialects to use new API.

[vinograd47]

Rebased and fixed conflicts with new code.

[vinograd47]

@mehdi_amini I've updated the change and removed old integer-based API from the MemRefType. I've fixed internal dialects by switching to the new API. I still used IntegerAttr for them, because switch to another representation (like StringAttr) will require more deep modification of the dialects and lowering passes.

[mehdi_amini]

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

200 ↗

(On Diff #323947)

How do you see the plan for moving away from IntegerAttr?
This kind of code becomes subject to crashing later.

[vinograd47]

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

200 ↗

(On Diff #323947)

@mehdi_amini In my understanding, the transition from IntegerAttr should be done step-by-step for separate dialects. For example:

1. Switch from IntegerAttr to StringAttr (or something else) inside GPU dialect:
   • Adjust its operation to work with new Attribute type.
   • Add type conversion to fallback to IntegerAttr when lowering to other dialects (SPIRV, LLVM). For example, replace "local" string attribute with 3 integer attribute during the lowering pass.
2. Switch from IntegerAttr to StringAttr inside SPIRV dialect:
   • Adjust its operation to work with new Attribute type.
   • Update lowering passes which converts to SPIRV dialect. For example, remove the type conversion from the above example or replace it with string/enum <-> string/enum conversion.
   • Update lowering passes from SPIRV dialect. Fallback to IntegerAttr as described above.
3. etc.

Each step can be done with separate change and by separate authors who are more familiar with specific dialect.

At least that how I see it. Are you OK with such plan? Actually I'd like to delegate this work since I am not familiar with any of this dialects (GPU, SPIRV, LLVM).

[vinograd47]

@mehdi_amini Actually, while I've seen positive feedback from the forum discussion (integer memory space leads to non-informative IRs, GPU dialect is good example for this), doing such change in atomic commit by one developer is quite hard task. It is breaking change, which affects various dialects, I believe, not only inside MLIR source tree. That's why I've submitted 2 variants of this change:

1. Non-breaking, with preserving original integer-based API in MemRefType. While it introduces dome duplication and non-convenient new API (getMemorySpaceAttr), it allows to avoid the break of the dialects. So, they can switch to the new API more smoothly.
2. Breaking change in MemRefType API, preserving integer semantic in dependent dialects (the current change version).

Both variants have own pros and cons, I'd like to get MLIR architects/maintainers opinion regarding which path is better.

[vinograd47]

BTW, as I can see the CI failed with the following error:

/mnt/disks/ssd0/agent/llvm-project/mlir/lib/Bindings/Python/PybindUtils.h:16:10: error: 'pybind11/pybind11.h' file not found [clang-diagnostic-error]
#include <pybind11/pybind11.h>
         ^

It doesn't look like related to the change itself, but some kind of infrastructure issue. I'm not sure, how to deal with it.

[mehdi_amini]

In D96145#2568366, @vinograd47 wrote:

@mehdi_amini Actually, while I've seen positive feedback from the forum discussion (integer memory space leads to non-informative IRs, GPU dialect is good example for this), doing such change in atomic commit by one developer is quite hard task. It is breaking change, which affects various dialects, I believe, not only inside MLIR source tree. That's why I've submitted 2 variants of this change:

1. Non-breaking, with preserving original integer-based API in MemRefType. While it introduces dome duplication and non-convenient new API (getMemorySpaceAttr), it allows to avoid the break of the dialects. So, they can switch to the new API more smoothly.
2. Breaking change in MemRefType API, preserving integer semantic in dependent dialects (the current change version).

Both variants have own pros and cons, I'd like to get MLIR architects/maintainers opinion regarding which path is better.

I'm fine doing it in multiple steps if it helps, I mostly pointed out that we don't make our life upstream harder for the convenience of downstream projects. Now these steps needs to be carefully crafted to control the amount of technical debt included and the potential to end up in a fragile situation (which is what I raised as a comment inline). Having some amount of extra APIs for the transition is OK, as long as there is a robust plan to get rid of them. For example when you introduce a getMemorySpaceAttr which should ultimately be named getMemorySpace, that seems to make it harder to later end up in the desired state.

What about this incremental plan:

1. rename getMemorySpace into getMemorySpaceInteger as a NFC refactoring.
2. Then change the storage to be a StringAttribute, introduce a getMemorySpace that returns a StringAttribute, and mark getMemorySpaceInteger deprecated (it could continue to return an integer when the string is "0", "1", etc. for now).
3. Migrate all in-tree uses to getMemorySpace.
4. When there is no use of getMemorySpaceInteger left in-tree it can be removed.

(I would think that going through all of these steps upstream shouldn't exceed a couple of weeks, or a month)

[vinograd47]

@mehdi_amini Submitted separate change for the first part (method renaming) : https://reviews.llvm.org/D97476

[vinograd47]

Re-implemented the feature on top of separate "renaming" change.

[vinograd47]

Update related section in LangRef documentation.

[vinograd47]

Rebased to latest main

[vinograd47]

@ftynse @rriddle @mehdi_amini The dependency for this change (old method renaming) was landed. Could you please take a look at this change one more time?

[ftynse]

LGTM in principle. Deferring to @rriddle for approval.

mlir/include/mlir/IR/BuiltinTypes.h

304

Could you actually mark deprecated functions with the [[deprecated]] attribute?

mlir/lib/IR/BuiltinTypes.cpp

494

Nit: is allowed

513

Nit: we use /// for top-level comments

[rriddle]

LGTM for code (after resolving outstanding comments)

First though, can you confirm if you are going with the migration plan proposed by @mehdi_amini ? Just want to make sure there is a clear path for removing the old API.

mlir/include/mlir/IR/BuiltinTypes.h

304

You'll want to use LLVM_ATTRIBUTE_DEPRECATED.

mlir/lib/IR/BuiltinTypes.cpp

503

I'm assuming this is for checking against the builtin dialect? You should be able to do something like isa<BuiltinDialect>(memorySpace.getDialect()) instead.

524–529

[ftynse]

mlir/include/mlir/IR/BuiltinTypes.h

304

The macro itself looks deprecated :)

https://github.com/llvm/llvm-project/blob/10c256ccaf520eed766f594f897b5bafdc8061ae/llvm/include/llvm/Support/Compiler.h#L318

[rriddle]

mlir/include/mlir/IR/BuiltinTypes.h

304

Yes, disregard me. I remember that there a recent discussion, but never actually checked the outcome.

[vinograd47]

Mark old comments as done.

[vinograd47]

Applied review remarks

[vinograd47]

@ftynse @rriddle I've applied your review remarks except deprecated attribute. I didn't add it intentionally, since the main code base was not updated to the new method, this attribute will generate dozen of warnings during the build.

@rriddle Regarding the plan - I'm not very familiar with dialects like GPU, SPIRV, LLVM (which use memorySpace intensively). So for me it will take a lot of time to adopt them, also it will be done in low priority in background to main work :) It would be great if this task is shared with developers more familiar with this dialects specifics.

[vinograd47]

Rebased

[vinograd47]

Rebased

[vinograd47]

Rebased and adjusted to latest changes in upstream.

[vinograd47]

@ftynse @rriddle @mehdi_amini Do you have any more comments for this change?

[mehdi_amini]

In D96145#2597733, @vinograd47 wrote:

@rriddle Regarding the plan - I'm not very familiar with dialects like GPU, SPIRV, LLVM (which use memorySpace intensively). So for me it will take a lot of time to adopt them, also it will be done in low priority in background to main work :) It would be great if this task is shared with developers more familiar with this dialects specifics.

Can you file individual bugs against each of these components to update to richer attributes?

[vinograd47]

In D96145#2614966, @mehdi_amini wrote:

Can you file individual bugs against each of these components to update to richer attributes?

Done:

• GPU Dialect : https://bugs.llvm.org/show_bug.cgi?id=49506
• SPIR-V Dialect : https://bugs.llvm.org/show_bug.cgi?id=49507

mlir/docs/LangRef.md

 ```
 
 The core syntax and representation of a layout specification is a
 [semi-affine map](Dialects/Affine.md#semi-affine-maps). Additionally, syntactic
 sugar is supported to make certain layout specifications more intuitive to read.
 For the moment, a `memref` supports parsing a strided form which is converted to
 a semi-affine map automatically.
 
-The memory space of a memref is specified by a target-specific integer index. If
-no memory space is specified, then the default memory space (0) is used. The
-default space is target specific but always at index 0.
-
-TODO: MLIR will eventually have target-dialects which allow symbolic use of
-memory hierarchy names (e.g. L3, L2, L1, ...) but we have not spec'd the details
-of that mechanism yet. Until then, this document pretends that it is valid to
-refer to these memories by `bare-id`.
+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.
+The empty memory space (attribute is None) is target specific.
 
 The notionally dynamic value of a memref value includes the address of the
 buffer allocated, as well as the symbols referred to by the shape, layout map,
 and index maps.
 
 Examples of memref static type
 
 ```mlir

mlir/include/mlir-c/BuiltinTypes.h

 /// Checks whether the given type is an UnrankedMemRef type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAUnrankedMemRef(MlirType type);
 
 /// Creates a MemRef type with the given rank and shape, a potentially empty
 /// list of affine layout maps, the given memory space and element type, in the
 /// same context as element type. The type is owned by the context.
 MLIR_CAPI_EXPORTED MlirType mlirMemRefTypeGet(
     MlirType elementType, intptr_t rank, const int64_t *shape, intptr_t numMaps,
-    MlirAffineMap const *affineMaps, unsigned memorySpace);
+    MlirAffineMap const *affineMaps, MlirAttribute memorySpace);
 
 /// Same as "mlirMemRefTypeGet" but returns a nullptr-wrapping MlirType o
 /// illegal arguments, emitting appropriate diagnostics.
 MLIR_CAPI_EXPORTED MlirType mlirMemRefTypeGetChecked(
     MlirLocation loc, MlirType elementType, intptr_t rank, const int64_t *shape,
-    intptr_t numMaps, MlirAffineMap const *affineMaps, unsigned memorySpace);
+    intptr_t numMaps, MlirAffineMap const *affineMaps,
+    MlirAttribute memorySpace);
 
 /// Creates a MemRef type with the given rank, shape, memory space and element
 /// type in the same context as the element type. The type has no affine maps,
 /// i.e. represents a default row-major contiguous memref. The type is owned by
 /// the context.
-MLIR_CAPI_EXPORTED MlirType mlirMemRefTypeContiguousGet(MlirType elementType,
-                                                        intptr_t rank,
-                                                        const int64_t *shape,
-                                                        unsigned memorySpace);
+MLIR_CAPI_EXPORTED MlirType
+mlirMemRefTypeContiguousGet(MlirType elementType, intptr_t rank,
+                            const int64_t *shape, MlirAttribute memorySpace);
 
 /// Same as "mlirMemRefTypeContiguousGet" but returns a nullptr wrapping
 /// MlirType on illegal arguments, emitting appropriate diagnostics.
 MLIR_CAPI_EXPORTED MlirType mlirMemRefTypeContiguousGetChecked(
     MlirLocation loc, MlirType elementType, intptr_t rank, const int64_t *shape,
-    unsigned memorySpace);
+    MlirAttribute memorySpace);
 
 /// Creates an Unranked MemRef type with the given element type and in the given
 /// memory space. The type is owned by the context of element type.
-MLIR_CAPI_EXPORTED MlirType mlirUnrankedMemRefTypeGet(MlirType elementType,
-                                                      unsigned memorySpace);
+MLIR_CAPI_EXPORTED MlirType
+mlirUnrankedMemRefTypeGet(MlirType elementType, MlirAttribute memorySpace);
 
 /// Same as "mlirUnrankedMemRefTypeGet" but returns a nullptr wrapping
 /// MlirType on illegal arguments, emitting appropriate diagnostics.
 MLIR_CAPI_EXPORTED MlirType mlirUnrankedMemRefTypeGetChecked(
-    MlirLocation loc, MlirType elementType, unsigned memorySpace);
+    MlirLocation loc, MlirType elementType, MlirAttribute memorySpace);
 
 /// Returns the number of affine layout maps in the given MemRef type.
 MLIR_CAPI_EXPORTED intptr_t mlirMemRefTypeGetNumAffineMaps(MlirType type);
 
 /// Returns the pos-th affine map of the given MemRef type.
 MLIR_CAPI_EXPORTED MlirAffineMap mlirMemRefTypeGetAffineMap(MlirType type,
                                                             intptr_t pos);
 
 /// Returns the memory space of the given MemRef type.
-MLIR_CAPI_EXPORTED unsigned mlirMemRefTypeGetMemorySpace(MlirType type);
+MLIR_CAPI_EXPORTED MlirAttribute mlirMemRefTypeGetMemorySpace(MlirType type);
 
 /// Returns the memory spcae of the given Unranked MemRef type.
-MLIR_CAPI_EXPORTED unsigned mlirUnrankedMemrefGetMemorySpace(MlirType type);
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirUnrankedMemrefGetMemorySpace(MlirType type);
 
 //===----------------------------------------------------------------------===//
 // Tuple type.
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the given type is a tuple type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsATuple(MlirType type);
 

mlir/include/mlir/IR/BuiltinTypes.h

 //===- BuiltinTypes.h - MLIR Builtin Type Classes ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_BUILTINTYPES_H
 #define MLIR_IR_BUILTINTYPES_H
 
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/Types.h"
 
 namespace llvm {
 struct fltSemantics;
 } // namespace llvm
 
 namespace mlir {
 class AffineExpr;
public:
 
   /// Return true if the specified element type is ok in a memref.
   static bool isValidElementType(Type type);
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type);
 
   /// Returns the memory space in which data referred to by this memref resides.
+  Attribute getMemorySpace() const;
+
+  /// Checks if the memorySpace has supported Attribute type.
+  static bool isSupportedMemorySpace(Attribute memorySpace);
+
+  /// [deprecated] Returns the memory space in old raw integer representation.
+  /// New `Attribute getMemorySpace()` method should be used instead.
   unsigned getMemorySpaceAsInt() const;
 };
 
 //===----------------------------------------------------------------------===//
 // MemRefType
 
 /// MemRef types represent a region of memory that have a shape with a fixed
 /// number of dimensions. Each shape element can be a non-negative integer or
 /// unknown (represented by -1). MemRef types also have an affine map
 /// composition, represented as an array AffineMap pointers.
 class MemRefType : public Type::TypeBase<MemRefType, BaseMemRefType,
                                          detail::MemRefTypeStorage> {
 public:
   /// This is a builder type that keeps local references to arguments. Arguments
   /// that are passed into the builder must out-live the builder.
   class Builder {
   public:
     // Build from another MemRefType.
     explicit Builder(MemRefType other)
         : shape(other.getShape()), elementType(other.getElementType()),
           affineMaps(other.getAffineMaps()),
-          memorySpace(other.getMemorySpaceAsInt()) {}
+          memorySpace(other.getMemorySpace()) {}
 
     // Build from scratch.
     Builder(ArrayRef<int64_t> shape, Type elementType)
-        : shape(shape), elementType(elementType), affineMaps(), memorySpace(0) {
-    }
+        : shape(shape), elementType(elementType), affineMaps() {}
 
     Builder &setShape(ArrayRef<int64_t> newShape) {
       shape = newShape;
       return *this;
     }
 
     Builder &setElementType(Type newElementType) {
       elementType = newElementType;
       return *this;
     }
 
     Builder &setAffineMaps(ArrayRef<AffineMap> newAffineMaps) {
       affineMaps = newAffineMaps;
       return *this;
     }
 
-    Builder &setMemorySpace(unsigned newMemorySpace) {
+    Builder &setMemorySpace(Attribute newMemorySpace) {
       memorySpace = newMemorySpace;
       return *this;
     }
 
+    // [deprecated] `setMemorySpace(Attribute)` should be used instead.
+    Builder &setMemorySpace(unsigned newMemorySpace);
+
     operator MemRefType() {
       return MemRefType::get(shape, elementType, affineMaps, memorySpace);
     }
 
   private:
     ArrayRef<int64_t> shape;
     Type elementType;
     ArrayRef<AffineMap> affineMaps;
-    unsigned memorySpace;
+    Attribute memorySpace;
   };
 
   using Base::Base;
   using Base::getChecked;
 
   /// Get or create a new MemRefType based on shape, element type, affine
   /// map composition, and memory space.  Assumes the arguments define a
   /// well-formed MemRef type.  Use getChecked to gracefully handle MemRefType
   /// construction failures.
   static MemRefType get(ArrayRef<int64_t> shape, Type elementType,
                         ArrayRef<AffineMap> affineMapComposition = {},
-                        unsigned memorySpace = 0);
+                        Attribute memorySpace = {});
+  // [deprecated] `Attribute`-based form should be used instead.
+  static MemRefType get(ArrayRef<int64_t> shape, Type elementType,
+                        ArrayRef<AffineMap> affineMapComposition,
+                        unsigned memorySpace);
 
   /// Get or create a new MemRefType based on shape, element type, affine
   /// map composition, and memory space. If the MemRefType defined by the
   /// arguments would be ill-formed, an error is emitted to `emitError` and a
   /// null type is returned.
   static MemRefType getChecked(function_ref<InFlightDiagnostic()> emitError,
                                ArrayRef<int64_t> shape, Type elementType,
                                ArrayRef<AffineMap> affineMapComposition,
+                               Attribute memorySpace);
+  // [deprecated] `Attribute`-based form should be used instead.
+  static MemRefType getChecked(function_ref<InFlightDiagnostic()> emitError,
+                               ArrayRef<int64_t> shape, Type elementType,
+                               ArrayRef<AffineMap> affineMapComposition,
                                unsigned memorySpace);
 
   ArrayRef<int64_t> getShape() const;
 
   /// Returns an array of affine map pointers representing the memref affine
   /// map composition.
   ArrayRef<AffineMap> getAffineMaps() const;
 
   // TODO: merge these two special values in a single one used everywhere.
   // Unfortunately, uses of `-1` have crept deep into the codebase now and are
   // hard to track.
   static int64_t getDynamicStrideOrOffset() {
     return ShapedType::kDynamicStrideOrOffset;
   }
 
 private:
   /// Get or create a new MemRefType defined by the arguments.  If the resulting
   /// type would be ill-formed, return nullptr.
   static MemRefType getImpl(ArrayRef<int64_t> shape, Type elementType,
                             ArrayRef<AffineMap> affineMapComposition,
-                            unsigned memorySpace,
+                            Attribute memorySpace,
                             function_ref<InFlightDiagnostic()> emitError);
   using Base::getImpl;
 };
 
 //===----------------------------------------------------------------------===//
 // UnrankedMemRefType
 
 /// Unranked MemRef type represent multi-dimensional MemRefs that
 /// have an unknown rank.
 class UnrankedMemRefType
     : public Type::TypeBase<UnrankedMemRefType, BaseMemRefType,
                             detail::UnrankedMemRefTypeStorage> {
 public:
   using Base::Base;
   using Base::getChecked;
 
   /// Get or create a new UnrankedMemRefType of the provided element
   /// type and memory space
+  static UnrankedMemRefType get(Type elementType, Attribute memorySpace);
+  // [deprecated] `Attribute`-based form should be used instead.
   static UnrankedMemRefType get(Type elementType, unsigned memorySpace);
 
   /// Get or create a new UnrankedMemRefType of the provided element
   /// type and memory space. If the UnrankedMemRefType defined by the arguments
   /// would be ill-formed, an error is emitted to `emitError` and a null type is
   /// returned.
   static UnrankedMemRefType
   getChecked(function_ref<InFlightDiagnostic()> emitError, Type elementType,
+             Attribute memorySpace);
+  // [deprecated] `Attribute`-based form should be used instead.
+  static UnrankedMemRefType
+  getChecked(function_ref<InFlightDiagnostic()> emitError, Type elementType,
              unsigned memorySpace);
 
   /// Verify the construction of a unranked memref type.
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
-                              Type elementType, unsigned memorySpace);
+                              Type elementType, Attribute memorySpace);
 
   ArrayRef<int64_t> getShape() const { return llvm::None; }
 };
 } // end namespace mlir
 
 //===----------------------------------------------------------------------===//
 // Tablegen Type Declarations
 //===----------------------------------------------------------------------===//

mlir/lib/Bindings/Python/IRModules.cpp

public:
   using PyConcreteType::PyConcreteType;
 
   PyMemRefLayoutMapList getLayout();
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
          "get",
          [](std::vector<int64_t> shape, PyType &elementType,
-            std::vector<PyAffineMap> layout, unsigned memorySpace,
+            std::vector<PyAffineMap> layout, PyAttribute *memorySpace,
             DefaultingPyLocation loc) {
            SmallVector<MlirAffineMap> maps;
            maps.reserve(layout.size());
            for (PyAffineMap &map : layout)
              maps.push_back(map);
 
+           MlirAttribute memSpaceAttr = {};
+           if (memorySpace)
+             memSpaceAttr = *memorySpace;
+
            MlirType t = mlirMemRefTypeGetChecked(loc, elementType, shape.size(),
                                                  shape.data(), maps.size(),
-                                                 maps.data(), memorySpace);
+                                                 maps.data(), memSpaceAttr);
            // TODO: Rework error reporting once diagnostic engine is exposed
            // in C API.
            if (mlirTypeIsNull(t)) {
              throw SetPyError(
                  PyExc_ValueError,
                  Twine("invalid '") +
                      py::repr(py::cast(elementType)).cast<std::string>() +
                      "' and expected floating point, integer, vector or "
                      "complex "
                      "type.");
            }
            return PyMemRefType(elementType.getContext(), t);
          },
          py::arg("shape"), py::arg("element_type"),
-         py::arg("layout") = py::list(), py::arg("memory_space") = 0,
+         py::arg("layout") = py::list(), py::arg("memory_space") = py::none(),
          py::arg("loc") = py::none(), "Create a memref type")
         .def_property_readonly("layout", &PyMemRefType::getLayout,
                                "The list of layout maps of the MemRef type.")
         .def_property_readonly(
             "memory_space",
-            [](PyMemRefType &self) -> unsigned {
-              return mlirMemRefTypeGetMemorySpace(self);
+            [](PyMemRefType &self) -> PyAttribute {
+              MlirAttribute a = mlirMemRefTypeGetMemorySpace(self);
+              return PyAttribute(self.getContext(), a);
             },
             "Returns the memory space of the given MemRef type.");
   }
 };
 
 /// A list of affine layout maps in a memref type. Internally, these are stored
 /// as consecutive elements, random access is cheap. Both the type and the maps
 /// are owned by the context, no need to worry about lifetime extension.

 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedMemRef;
   static constexpr const char *pyClassName = "UnrankedMemRefType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
          "get",
-         [](PyType &elementType, unsigned memorySpace,
+         [](PyType &elementType, PyAttribute *memorySpace,
             DefaultingPyLocation loc) {
+           MlirAttribute memSpaceAttr = {};
+           if (memorySpace)
+             memSpaceAttr = *memorySpace;
+
            MlirType t =
-               mlirUnrankedMemRefTypeGetChecked(loc, elementType, memorySpace);
+               mlirUnrankedMemRefTypeGetChecked(loc, elementType, memSpaceAttr);
            // TODO: Rework error reporting once diagnostic engine is exposed
            // in C API.
            if (mlirTypeIsNull(t)) {
              throw SetPyError(
                  PyExc_ValueError,
                  Twine("invalid '") +
                      py::repr(py::cast(elementType)).cast<std::string>() +
                      "' and expected floating point, integer, vector or "
                      "complex "
                      "type.");
            }
            return PyUnrankedMemRefType(elementType.getContext(), t);
          },
          py::arg("element_type"), py::arg("memory_space"),
          py::arg("loc") = py::none(), "Create a unranked memref type")
         .def_property_readonly(
             "memory_space",
-            [](PyUnrankedMemRefType &self) -> unsigned {
-              return mlirUnrankedMemrefGetMemorySpace(self);
+            [](PyUnrankedMemRefType &self) -> PyAttribute {
+              MlirAttribute a = mlirMemRefTypeGetMemorySpace(self);
+              return PyAttribute(self.getContext(), a);
             },
             "Returns the memory space of the given Unranked MemRef type.");
   }
 };
 
 /// Tuple Type subclass - TupleType.
 class PyTupleType : public PyConcreteType<PyTupleType> {
 public:

mlir/lib/CAPI/IR/BuiltinTypes.cpp

 // Ranked / Unranked MemRef type.
 //===----------------------------------------------------------------------===//
 
 bool mlirTypeIsAMemRef(MlirType type) { return unwrap(type).isa<MemRefType>(); }
 
 MlirType mlirMemRefTypeGet(MlirType elementType, intptr_t rank,
                            const int64_t *shape, intptr_t numMaps,
                            MlirAffineMap const *affineMaps,
-                           unsigned memorySpace) {
+                           MlirAttribute memorySpace) {
   SmallVector<AffineMap, 1> maps;
   (void)unwrapList(numMaps, affineMaps, maps);
   return wrap(
       MemRefType::get(llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
-                      unwrap(elementType), maps, memorySpace));
+                      unwrap(elementType), maps, unwrap(memorySpace)));
 }
 
 MlirType mlirMemRefTypeGetChecked(MlirLocation loc, MlirType elementType,
                                   intptr_t rank, const int64_t *shape,
                                   intptr_t numMaps,
                                   MlirAffineMap const *affineMaps,
-                                  unsigned memorySpace) {
+                                  MlirAttribute memorySpace) {
   SmallVector<AffineMap, 1> maps;
   (void)unwrapList(numMaps, affineMaps, maps);
   return wrap(MemRefType::getChecked(
       unwrap(loc), llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
-      unwrap(elementType), maps, memorySpace));
+      unwrap(elementType), maps, unwrap(memorySpace)));
 }
 
 MlirType mlirMemRefTypeContiguousGet(MlirType elementType, intptr_t rank,
                                      const int64_t *shape,
-                                     unsigned memorySpace) {
+                                     MlirAttribute memorySpace) {
   return wrap(
       MemRefType::get(llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
-                      unwrap(elementType), llvm::None, memorySpace));
+                      unwrap(elementType), llvm::None, unwrap(memorySpace)));
 }
 
 MlirType mlirMemRefTypeContiguousGetChecked(MlirLocation loc,
                                             MlirType elementType, intptr_t rank,
                                             const int64_t *shape,
-                                            unsigned memorySpace) {
+                                            MlirAttribute memorySpace) {
   return wrap(MemRefType::getChecked(
       unwrap(loc), llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
-      unwrap(elementType), llvm::None, memorySpace));
+      unwrap(elementType), llvm::None, unwrap(memorySpace)));
 }
 
 intptr_t mlirMemRefTypeGetNumAffineMaps(MlirType type) {
   return static_cast<intptr_t>(
       unwrap(type).cast<MemRefType>().getAffineMaps().size());
 }
 
 MlirAffineMap mlirMemRefTypeGetAffineMap(MlirType type, intptr_t pos) {
   return wrap(unwrap(type).cast<MemRefType>().getAffineMaps()[pos]);
 }
 
-unsigned mlirMemRefTypeGetMemorySpace(MlirType type) {
-  return unwrap(type).cast<MemRefType>().getMemorySpaceAsInt();
+MlirAttribute mlirMemRefTypeGetMemorySpace(MlirType type) {
+  return wrap(unwrap(type).cast<MemRefType>().getMemorySpace());
 }
 
 bool mlirTypeIsAUnrankedMemRef(MlirType type) {
   return unwrap(type).isa<UnrankedMemRefType>();
 }
 
-MlirType mlirUnrankedMemRefTypeGet(MlirType elementType, unsigned memorySpace) {
-  return wrap(UnrankedMemRefType::get(unwrap(elementType), memorySpace));
+MlirType mlirUnrankedMemRefTypeGet(MlirType elementType,
+                                   MlirAttribute memorySpace) {
+  return wrap(
+      UnrankedMemRefType::get(unwrap(elementType), unwrap(memorySpace)));
 }
 
 MlirType mlirUnrankedMemRefTypeGetChecked(MlirLocation loc,
                                           MlirType elementType,
-                                          unsigned memorySpace) {
+                                          MlirAttribute memorySpace) {
   return wrap(UnrankedMemRefType::getChecked(unwrap(loc), unwrap(elementType),
-                                             memorySpace));
+                                             unwrap(memorySpace)));
 }
 
-unsigned mlirUnrankedMemrefGetMemorySpace(MlirType type) {
-  return unwrap(type).cast<UnrankedMemRefType>().getMemorySpaceAsInt();
+MlirAttribute mlirUnrankedMemrefGetMemorySpace(MlirType type) {
+  return wrap(unwrap(type).cast<UnrankedMemRefType>().getMemorySpace());
 }
 
 //===----------------------------------------------------------------------===//
 // Tuple type.
 //===----------------------------------------------------------------------===//
 
 bool mlirTypeIsATuple(MlirType type) { return unwrap(type).isa<TupleType>(); }
 

mlir/lib/IR/AsmPrinter.cpp

void ModulePrinter::printDenseStringElementsAttr(DenseStringElementsAttr attr) {
   auto printFn = [&](unsigned index) {
     os << "\"";
     printEscapedString(data[index], os);
     os << "\"";
   };
   printDenseElementsAttrImpl(attr.isSplat(), attr.getType(), os, printFn);
 }
 
 void ModulePrinter::printType(Type type) {
   if (!type) {
     os << "<<NULL TYPE>>";
     return;
   }
 
   // Try to print an alias for this type.
   if (state && succeeded(state->getAliasState().getAlias(type, os)))
     return;
TypeSwitch<Type>(type)
           os << 'x';
         }
         printType(memrefTy.getElementType());
         for (auto map : memrefTy.getAffineMaps()) {
           os << ", ";
           printAttribute(AffineMapAttr::get(map));
         }
         // Only print the memory space if it is the non-default one.
-        if (memrefTy.getMemorySpaceAsInt())
-          os << ", " << memrefTy.getMemorySpaceAsInt();
+        if (memrefTy.getMemorySpace()) {
+          os << ", ";
+          printAttribute(memrefTy.getMemorySpace(), AttrTypeElision::May);
+        }
         os << '>';
       })
       .Case<UnrankedMemRefType>([&](UnrankedMemRefType memrefTy) {
         os << "memref<*x";
         printType(memrefTy.getElementType());
         // Only print the memory space if it is the non-default one.
-        if (memrefTy.getMemorySpaceAsInt())
-          os << ", " << memrefTy.getMemorySpaceAsInt();
+        if (memrefTy.getMemorySpace()) {
+          os << ", ";
+          printAttribute(memrefTy.getMemorySpace(), AttrTypeElision::May);
+        }
         os << '>';
       })
       .Case<ComplexType>([&](ComplexType complexTy) {
         os << "complex<";
         printType(complexTy.getElementType());
         os << '>';
       })
       .Case<TupleType>([&](TupleType tupleTy) {

mlir/lib/IR/BuiltinTypes.cpp

 //===- BuiltinTypes.cpp - MLIR Builtin Type Classes -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "TypeDetail.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/Twine.h"
 
 using namespace mlir;
if (auto other = dyn_cast<MemRefType>()) {
     MemRefType::Builder b(other);
     b.setShape(shape);
     b.setElementType(elementType);
     return b;
   }
 
   if (auto other = dyn_cast<UnrankedMemRefType>()) {
     MemRefType::Builder b(shape, elementType);
-    b.setMemorySpace(other.getMemorySpaceAsInt());
+    b.setMemorySpace(other.getMemorySpace());
     return b;
   }
 
   if (isa<TensorType>())
     return RankedTensorType::get(shape, elementType);
 
   if (isa<VectorType>())
     return VectorType::get(shape, elementType);
 
   llvm_unreachable("Unhandled ShapedType clone case");
 }
 
 ShapedType ShapedType::clone(ArrayRef<int64_t> shape) {
   if (auto other = dyn_cast<MemRefType>()) {
     MemRefType::Builder b(other);
     b.setShape(shape);
     return b;
   }
 
   if (auto other = dyn_cast<UnrankedMemRefType>()) {
     MemRefType::Builder b(shape, other.getElementType());
     b.setShape(shape);
-    b.setMemorySpace(other.getMemorySpaceAsInt());
+    b.setMemorySpace(other.getMemorySpace());
     return b;
   }
 
   if (isa<TensorType>())
     return RankedTensorType::get(shape, getElementType());
 
   if (isa<VectorType>())
     return VectorType::get(shape, getElementType());
 
   llvm_unreachable("Unhandled ShapedType clone case");
 }
 
 ShapedType ShapedType::clone(Type elementType) {
   if (auto other = dyn_cast<MemRefType>()) {
     MemRefType::Builder b(other);
     b.setElementType(elementType);
     return b;
   }
 
   if (auto other = dyn_cast<UnrankedMemRefType>()) {
-    return UnrankedMemRefType::get(elementType, other.getMemorySpaceAsInt());
+    return UnrankedMemRefType::get(elementType, other.getMemorySpace());
   }
 
   if (isa<TensorType>()) {
     if (hasRank())
       return RankedTensorType::get(getShape(), elementType);
     return UnrankedTensorType::get(elementType);
   }
 
UnrankedTensorType::verify(function_ref<InFlightDiagnostic()> emitError,
   return checkTensorElementType(emitError, elementType);
 }
 
 //===----------------------------------------------------------------------===//
 // BaseMemRefType
 //===----------------------------------------------------------------------===//
 
 unsigned BaseMemRefType::getMemorySpaceAsInt() const {
+  Attribute attr = getMemorySpace();
+
+  if (!attr)
+    return 0;
+
+  assert(attr.isa<IntegerAttr>() &&
+         "Using `getMemorySpaceInteger` with non-Integer attribute");
+
+  return static_cast<unsigned>(attr.cast<IntegerAttr>().getInt());
+}
+
+Attribute BaseMemRefType::getMemorySpace() const {
   return static_cast<ImplType *>(impl)->memorySpace;
 }
 
+bool BaseMemRefType::isSupportedMemorySpace(Attribute memorySpace) {
+  // Empty attribute is allowed as default memory space.
+  if (!memorySpace)
+    return true;
+
+  // Supported built-in attributes.
+  if (memorySpace.isa<IntegerAttr, StringAttr, DictionaryAttr>())
+    return true;
+
+  // Allow custom dialect attributes.
+  if (!::mlir::isa<BuiltinDialect>(memorySpace.getDialect()))
+    return true;
+
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
 
+/// Wraps deprecated integer memory space to the new Attribute form.
+static Attribute wrapIntegerMemorySpace(unsigned memorySpace,
+                                        MLIRContext *ctx) {
+  if (memorySpace == 0)
+    return nullptr;
+
+  return IntegerAttr::get(IntegerType::get(ctx, 64), memorySpace);
+}
+
+/// Replaces default memorySpace (integer == `0`) with empty Attribute.
+static Attribute skipDefaultMemorySpace(Attribute memorySpace) {
+  IntegerAttr intMemorySpace = memorySpace.dyn_cast_or_null<IntegerAttr>();
+  if (intMemorySpace && intMemorySpace.getValue() == 0)
+    return nullptr;
+
+  return memorySpace;
+}
+
+MemRefType::Builder &
+MemRefType::Builder::setMemorySpace(unsigned newMemorySpace) {
+  memorySpace =
+      wrapIntegerMemorySpace(newMemorySpace, elementType.getContext());
+  return *this;
+}
+
 /// Get or create a new MemRefType based on shape, element type, affine
 /// map composition, and memory space.  Assumes the arguments define a
 /// well-formed MemRef type.  Use getChecked to gracefully handle MemRefType
 /// construction failures.
 MemRefType MemRefType::get(ArrayRef<int64_t> shape, Type elementType,
                            ArrayRef<AffineMap> affineMapComposition,
-                           unsigned memorySpace) {
+                           Attribute memorySpace) {
   auto result =
       getImpl(shape, elementType, affineMapComposition, memorySpace, [=] {
         return emitError(UnknownLoc::get(elementType.getContext()));
       });
   assert(result && "Failed to construct instance of MemRefType.");
   return result;
 }
+MemRefType MemRefType::get(ArrayRef<int64_t> shape, Type elementType,
+                           ArrayRef<AffineMap> affineMapComposition,
+                           unsigned memorySpace) {
+  Attribute memorySpaceAttr =
+      wrapIntegerMemorySpace(memorySpace, elementType.getContext());
+  return get(shape, elementType, affineMapComposition, memorySpaceAttr);
+}
 
 /// Get or create a new MemRefType based on shape, element type, affine
 /// map composition, and memory space declared at the given location.
 /// If the location is unknown, the last argument should be an instance of
 /// UnknownLoc.  If the MemRefType defined by the arguments would be
 /// ill-formed, emits errors (to the handler registered with the context or to
 /// the error stream) and returns nullptr.
 MemRefType MemRefType::getChecked(function_ref<InFlightDiagnostic()> emitError,
                                   ArrayRef<int64_t> shape, Type elementType,
                                   ArrayRef<AffineMap> affineMapComposition,
-                                  unsigned memorySpace) {
+                                  Attribute memorySpace) {
   return getImpl(shape, elementType, affineMapComposition, memorySpace,
                  emitError);
 }
+MemRefType MemRefType::getChecked(function_ref<InFlightDiagnostic()> emitError,
+                                  ArrayRef<int64_t> shape, Type elementType,
+                                  ArrayRef<AffineMap> affineMapComposition,
+                                  unsigned memorySpace) {
+  Attribute memorySpaceAttr =
+      wrapIntegerMemorySpace(memorySpace, elementType.getContext());
+  return getChecked(emitError, shape, elementType, affineMapComposition,
+                    memorySpaceAttr);
+}
 
 /// Get or create a new MemRefType defined by the arguments.  If the resulting
 /// type would be ill-formed, return nullptr.  If the location is provided,
 /// emit detailed error messages.  To emit errors when the location is unknown,
 /// pass in an instance of UnknownLoc.
 MemRefType MemRefType::getImpl(ArrayRef<int64_t> shape, Type elementType,
                                ArrayRef<AffineMap> affineMapComposition,
-                               unsigned memorySpace,
+                               Attribute memorySpace,
                                function_ref<InFlightDiagnostic()> emitError) {
   auto *context = elementType.getContext();
 
   if (!BaseMemRefType::isValidElementType(elementType))
     return (emitError() << "invalid memref element type", MemRefType());
 
   for (int64_t s : shape) {
     // Negative sizes are not allowed except for `-1` that means dynamic size.
MemRefType MemRefType::getImpl(ArrayRef<int64_t> shape, Type elementType,
   // implicit identity.
   SmallVector<AffineMap, 2> cleanedAffineMapComposition;
   for (const auto &map : affineMapComposition) {
     if (map.isIdentity())
       continue;
     cleanedAffineMapComposition.push_back(map);
   }
 
+  if (!isSupportedMemorySpace(memorySpace)) {
+    emitError() << "unsupported memory space Attribute";
+    return nullptr;
+  }
+
   return Base::get(context, shape, elementType, cleanedAffineMapComposition,
-                   memorySpace);
+                   skipDefaultMemorySpace(memorySpace));
 }
 
 ArrayRef<int64_t> MemRefType::getShape() const { return getImpl()->getShape(); }
 
 ArrayRef<AffineMap> MemRefType::getAffineMaps() const {
   return getImpl()->getAffineMaps();
 }
 
 //===----------------------------------------------------------------------===//
 // UnrankedMemRefType
 //===----------------------------------------------------------------------===//
 
 UnrankedMemRefType UnrankedMemRefType::get(Type elementType,
+                                           Attribute memorySpace) {
+  return Base::get(elementType.getContext(), elementType,
+                   skipDefaultMemorySpace(memorySpace));
+}
+UnrankedMemRefType UnrankedMemRefType::get(Type elementType,
                                            unsigned memorySpace) {
-  return Base::get(elementType.getContext(), elementType, memorySpace);
+  Attribute memorySpaceAttr =
+      wrapIntegerMemorySpace(memorySpace, elementType.getContext());
+  return get(elementType, memorySpaceAttr);
 }
 
 UnrankedMemRefType
 UnrankedMemRefType::getChecked(function_ref<InFlightDiagnostic()> emitError,
-                               Type elementType, unsigned memorySpace) {
+                               Type elementType, Attribute memorySpace) {
   return Base::getChecked(emitError, elementType.getContext(), elementType,
-                          memorySpace);
+                          skipDefaultMemorySpace(memorySpace));
+}
+UnrankedMemRefType
+UnrankedMemRefType::getChecked(function_ref<InFlightDiagnostic()> emitError,
+                               Type elementType, unsigned memorySpace) {
+  Attribute memorySpaceAttr =
+      wrapIntegerMemorySpace(memorySpace, elementType.getContext());
+  return getChecked(emitError, elementType, memorySpaceAttr);
 }
 
 LogicalResult
 UnrankedMemRefType::verify(function_ref<InFlightDiagnostic()> emitError,
-                           Type elementType, unsigned memorySpace) {
+                           Type elementType, Attribute memorySpace) {
   if (!BaseMemRefType::isValidElementType(elementType))
     return emitError() << "invalid memref element type";
+
+  if (!isSupportedMemorySpace(memorySpace))
+    return emitError() << "unsupported memory space Attribute";
+
   return success();
 }
 
 // Fallback cases for terminal dim/sym/cst that are not part of a binary op (
 // i.e. single term). Accumulate the AffineExpr into the existing one.
 static void extractStridesFromTerm(AffineExpr e,
                                    AffineExpr multiplicativeFactor,
                                    MutableArrayRef<AffineExpr> strides,

mlir/lib/IR/TypeDetail.h

struct UnrankedTensorTypeStorage : public ShapedTypeStorage {
   static UnrankedTensorTypeStorage *construct(TypeStorageAllocator &allocator,
                                               Type elementTy) {
     return new (allocator.allocate<UnrankedTensorTypeStorage>())
         UnrankedTensorTypeStorage(elementTy);
   }
 };
 
 struct BaseMemRefTypeStorage : public ShapedTypeStorage {
-  BaseMemRefTypeStorage(Type elementType, unsigned memorySpace)
+  BaseMemRefTypeStorage(Type elementType, Attribute memorySpace)
       : ShapedTypeStorage(elementType), memorySpace(memorySpace) {}
 
   /// Memory space in which data referenced by memref resides.
-  const unsigned memorySpace;
+  Attribute memorySpace;
 };
 
 struct MemRefTypeStorage : public BaseMemRefTypeStorage {
   MemRefTypeStorage(unsigned shapeSize, Type elementType,
                     const int64_t *shapeElements, const unsigned numAffineMaps,
-                    AffineMap const *affineMapList, const unsigned memorySpace)
+                    AffineMap const *affineMapList, Attribute memorySpace)
       : BaseMemRefTypeStorage(elementType, memorySpace),
         shapeElements(shapeElements), shapeSize(shapeSize),
         numAffineMaps(numAffineMaps), affineMapList(affineMapList) {}
 
   /// The hash key used for uniquing.
   // MemRefs are uniqued based on their shape, element type, affine map
   // composition, and memory space.
   using KeyTy =
-      std::tuple<ArrayRef<int64_t>, Type, ArrayRef<AffineMap>, unsigned>;
+      std::tuple<ArrayRef<int64_t>, Type, ArrayRef<AffineMap>, Attribute>;
   bool operator==(const KeyTy &key) const {
     return key == KeyTy(getShape(), elementType, getAffineMaps(), memorySpace);
   }
 
   /// Construction.
   static MemRefTypeStorage *construct(TypeStorageAllocator &allocator,
                                       const KeyTy &key) {
     // Copy the shape into the bump pointer.
struct MemRefTypeStorage : public BaseMemRefTypeStorage {
   /// List of affine maps in the memref's layout/index map composition.
   AffineMap const *affineMapList;
 };
 
 /// Unranked MemRef is a MemRef with unknown rank.
 /// Only element type and memory space are known
 struct UnrankedMemRefTypeStorage : public BaseMemRefTypeStorage {
 
-  UnrankedMemRefTypeStorage(Type elementTy, const unsigned memorySpace)
+  UnrankedMemRefTypeStorage(Type elementTy, Attribute memorySpace)
       : BaseMemRefTypeStorage(elementTy, memorySpace) {}
 
   /// The hash key used for uniquing.
-  using KeyTy = std::tuple<Type, unsigned>;
+  using KeyTy = std::tuple<Type, Attribute>;
   bool operator==(const KeyTy &key) const {
     return key == KeyTy(elementType, memorySpace);
   }
 
   /// Construction.
   static UnrankedMemRefTypeStorage *construct(TypeStorageAllocator &allocator,
                                               const KeyTy &key) {
 

mlir/lib/Parser/TypeParser.cpp

if (!elementType)
     return nullptr;
 
   // Check that memref is formed from allowed types.
   if (!BaseMemRefType::isValidElementType(elementType))
     return emitError(typeLoc, "invalid memref element type"), nullptr;
 
   // Parse semi-affine-map-composition.
   SmallVector<AffineMap, 2> affineMapComposition;
-  Optional<unsigned> memorySpace;
+  Attribute memorySpace;
   unsigned numDims = dimensions.size();
 
   auto parseElt = [&]() -> ParseResult {
-    // Check for the memory space.
-    if (getToken().is(Token::integer)) {
-      if (memorySpace)
-        return emitError("multiple memory spaces specified in memref type");
-      memorySpace = getToken().getUnsignedIntegerValue();
-      if (!memorySpace.hasValue())
-        return emitError("invalid memory space in memref type");
-      consumeToken(Token::integer);
-      return success();
-    }
-    if (isUnranked)
-      return emitError("cannot have affine map for unranked memref type");
-    if (memorySpace)
-      return emitError("expected memory space to be last in memref type");
-
     AffineMap map;
     llvm::SMLoc mapLoc = getToken().getLoc();
+
+    // Check for AffineMap as offset/strides.
     if (getToken().is(Token::kw_offset)) {
       int64_t offset;
       SmallVector<int64_t, 4> strides;
       if (failed(parseStridedLayout(offset, strides)))
         return failure();
       // Construct strided affine map.
       map = makeStridedLinearLayoutMap(strides, offset, state.context);
     } else {
-      // Parse an affine map attribute.
-      auto affineMap = parseAttribute();
-      if (!affineMap)
+      // Either it is AffineMapAttr or memory space attribute.
+      Attribute attr = parseAttribute();
+      if (!attr)
         return failure();
-      auto affineMapAttr = affineMap.dyn_cast<AffineMapAttr>();
-      if (!affineMapAttr)
-        return emitError("expected affine map in memref type");
+
+      if (AffineMapAttr affineMapAttr = attr.dyn_cast<AffineMapAttr>()) {
         map = affineMapAttr.getValue();
+      } else if (memorySpace) {
+        return emitError("multiple memory spaces specified in memref type");
+      } else {
+        memorySpace = attr;
+        return success();
+      }
     }
 
+    if (isUnranked)
+      return emitError("cannot have affine map for unranked memref type");
+    if (memorySpace)
+      return emitError("expected memory space to be last in memref type");
+
     if (map.getNumDims() != numDims) {
       size_t i = affineMapComposition.size();
       return emitError(mapLoc, "memref affine map dimension mismatch between ")
              << (i == 0 ? Twine("memref rank") : "affine map " + Twine(i))
              << " and affine map" << i + 1 << ": " << numDims
              << " != " << map.getNumDims();
     }
     numDims = map.getNumResults();
     affineMapComposition.push_back(map);
     return success();
   };
 
   // Parse a list of mappings and address space if present.
   if (!consumeIf(Token::greater)) {
     // Parse comma separated list of affine maps, followed by memory space.
     if (parseToken(Token::comma, "expected ',' or '>' in memref type") ||
         parseCommaSeparatedListUntil(Token::greater, parseElt,
                                      /*allowEmptyList=*/false)) {
       return nullptr;
     }
   }
 
-  if (isUnranked)
-    return UnrankedMemRefType::get(elementType, memorySpace.getValueOr(0));
+  if (isUnranked) {
+    return UnrankedMemRefType::getChecked(
+        [&]() -> InFlightDiagnostic { return emitError(); }, elementType,
+        memorySpace);
+  }
 
-  return MemRefType::get(dimensions, elementType, affineMapComposition,
-                         memorySpace.getValueOr(0));
+  return MemRefType::getChecked(
+      [&]() -> InFlightDiagnostic { return emitError(); }, dimensions,
+      elementType, affineMapComposition, memorySpace);
 }
 
 /// Parse any type except the function type.
 ///
 ///   non-function-type ::= integer-type
 ///                       | index-type
 ///                       | float-type
 ///                       | extended-type

mlir/test/Bindings/Python/ir_types.py

 run(testUnrankedTensorType)
 
 # CHECK-LABEL: TEST: testMemRefType
 def testMemRefType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     shape = [2, 3]
     loc = Location.unknown()
-    memref = MemRefType.get(shape, f32, memory_space=2)
+    memref = MemRefType.get(shape, f32, memory_space=Attribute.parse("2"))
     # CHECK: memref type: memref<2x3xf32, 2>
     print("memref type:", memref)
     # CHECK: number of affine layout maps: 0
     print("number of affine layout maps:", len(memref.layout))
     # CHECK: memory space: 2
     print("memory space:", memref.memory_space)
 
     layout = AffineMap.get_permutation([1, 0])
     memref_layout = MemRefType.get(shape, f32, [layout])
     # CHECK: memref type: memref<2x3xf32, affine_map<(d0, d1) -> (d1, d0)>>
     print("memref type:", memref_layout)
     assert len(memref_layout.layout) == 1
     # CHECK: memref layout: (d0, d1) -> (d1, d0)
     print("memref layout:", memref_layout.layout[0])
-    # CHECK: memory space: 0
+    # CHECK: memory space: <<NULL ATTRIBUTE>>
     print("memory space:", memref_layout.memory_space)
 
     none = NoneType.get()
     try:
       memref_invalid = MemRefType.get(shape, none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
 run(testMemRefType)
 
 # CHECK-LABEL: TEST: testUnrankedMemRefType
 def testUnrankedMemRefType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     loc = Location.unknown()
-    unranked_memref = UnrankedMemRefType.get(f32, 2)
+    unranked_memref = UnrankedMemRefType.get(f32, Attribute.parse("2"))
     # CHECK: unranked memref type: memref<*xf32, 2>
     print("unranked memref type:", unranked_memref)
     try:
       invalid_rank = unranked_memref.rank
     except ValueError as e:
       # CHECK: calling this method requires that the type has a rank.
       print(e)
     else:
with Context(), Location.unknown():
     except ValueError as e:
       # CHECK: calling this method requires that the type has a rank.
       print(e)
     else:
       print("Exception not produced")
 
     none = NoneType.get()
     try:
-      memref_invalid = UnrankedMemRefType.get(none, 2)
+      memref_invalid = UnrankedMemRefType.get(none, Attribute.parse("2"))
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
 run(testUnrankedMemRefType)

mlir/test/CAPI/ir.c

if (!mlirTypeIsATensor(unrankedTensor) ||
       !mlirTypeIsAUnrankedTensor(unrankedTensor) ||
       mlirShapedTypeHasRank(unrankedTensor))
     return 17;
   mlirTypeDump(unrankedTensor);
   fprintf(stderr, "\n");
   // CHECK: tensor<*xf32>
 
   // MemRef type.
+  MlirAttribute memSpace2 = mlirIntegerAttrGet(mlirIntegerTypeGet(ctx, 64), 2);
   MlirType memRef = mlirMemRefTypeContiguousGet(
-      f32, sizeof(shape) / sizeof(int64_t), shape, 2);
+      f32, sizeof(shape) / sizeof(int64_t), shape, memSpace2);
   if (!mlirTypeIsAMemRef(memRef) ||
       mlirMemRefTypeGetNumAffineMaps(memRef) != 0 ||
-      mlirMemRefTypeGetMemorySpace(memRef) != 2)
+      !mlirAttributeEqual(mlirMemRefTypeGetMemorySpace(memRef), memSpace2))
     return 18;
   mlirTypeDump(memRef);
   fprintf(stderr, "\n");
   // CHECK: memref<2x3xf32, 2>
 
   // Unranked MemRef type.
-  MlirType unrankedMemRef = mlirUnrankedMemRefTypeGet(f32, 4);
+  MlirAttribute memSpace4 = mlirIntegerAttrGet(mlirIntegerTypeGet(ctx, 64), 4);
+  MlirType unrankedMemRef = mlirUnrankedMemRefTypeGet(f32, memSpace4);
   if (!mlirTypeIsAUnrankedMemRef(unrankedMemRef) ||
       mlirTypeIsAMemRef(unrankedMemRef) ||
-      mlirUnrankedMemrefGetMemorySpace(unrankedMemRef) != 4)
+      !mlirAttributeEqual(mlirUnrankedMemrefGetMemorySpace(unrankedMemRef),
+                          memSpace4))
     return 19;
   mlirTypeDump(unrankedMemRef);
   fprintf(stderr, "\n");
   // CHECK: memref<*xf32, 4>
 
   // Tuple type.
   MlirType types[] = {unrankedMemRef, f32};
   MlirType tuple = mlirTupleTypeGet(ctx, 2, types);

mlir/test/IR/invalid.mlir

 // Test no map in memref type.
 func @memrefs(memref<2x4xi8, >) // expected-error {{expected list element}}
 
 // -----
 // Test non-existent map in memref type.
 func @memrefs(memref<2x4xi8, #map7>) // expected-error {{undefined symbol alias id 'map7'}}
 
 // -----
-// Test non affine map in memref type.
-func @memrefs(memref<2x4xi8, i8>) // expected-error {{expected affine map in memref type}}
+// Test unsupported memory space.
+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)>
 
 func @memrefs(memref<2x4xi8, #map0, #map8>) // expected-error {{undefined symbol alias id 'map8'}}
 
 // -----

mlir/test/IR/parser.mlir

 func private @memrefs_drop_triv_id_multiple(memref<2xi8, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>>)
 
 // These maps appeared before, so they must be uniqued and hoisted to the beginning.
 // Identity map should be removed.
 // CHECK: func private @memrefs_compose_with_id(memref<2x2xi8, #map{{[0-9]+}}>)
 func private @memrefs_compose_with_id(memref<2x2xi8, affine_map<(d0, d1) -> (d0, d1)>,
                                              affine_map<(d0, d1) -> (d1, d0)>>)
 
+// Test memref with custom memory space
+
+// CHECK: func private @memrefs_nomap_nospace(memref<5x6x7xf32>)
+func private @memrefs_nomap_nospace(memref<5x6x7xf32>)
+
+// CHECK: func private @memrefs_map_nospace(memref<5x6x7xf32, #map{{[0-9]+}}>)
+func private @memrefs_map_nospace(memref<5x6x7xf32, #map3>)
+
+// CHECK: func private @memrefs_nomap_intspace(memref<5x6x7xf32, 3>)
+func private @memrefs_nomap_intspace(memref<5x6x7xf32, 3>)
+
+// CHECK: func private @memrefs_map_intspace(memref<5x6x7xf32, #map{{[0-9]+}}, 5>)
+func private @memrefs_map_intspace(memref<5x6x7xf32, #map3, 5>)
+
+// CHECK: func private @memrefs_nomap_strspace(memref<5x6x7xf32, "local">)
+func private @memrefs_nomap_strspace(memref<5x6x7xf32, "local">)
+
+// CHECK: func private @memrefs_map_strspace(memref<5x6x7xf32, #map{{[0-9]+}}, "private">)
+func private @memrefs_map_strspace(memref<5x6x7xf32, #map3, "private">)
+
+// CHECK: func private @memrefs_nomap_dictspace(memref<5x6x7xf32, {memSpace = "special", subIndex = 1 : i64}>)
+func private @memrefs_nomap_dictspace(memref<5x6x7xf32, {memSpace = "special", subIndex = 1}>)
+
+// CHECK: func private @memrefs_map_dictspace(memref<5x6x7xf32, #map{{[0-9]+}}, {memSpace = "special", subIndex = 3 : i64}>)
+func private @memrefs_map_dictspace(memref<5x6x7xf32, #map3, {memSpace = "special", subIndex = 3}>)
 
 // CHECK: func private @complex_types(complex<i1>) -> complex<f32>
 func private @complex_types(complex<i1>) -> complex<f32>
 
-
 // CHECK: func private @memref_with_index_elems(memref<1x?xindex>)
 func private @memref_with_index_elems(memref<1x?xindex>)
 
 // CHECK: func private @memref_with_complex_elems(memref<1x?xcomplex<f32>>)
 func private @memref_with_complex_elems(memref<1x?xcomplex<f32>>)
 
 // CHECK: func private @memref_with_vector_elems(memref<1x?xvector<10xf32>>)
 func private @memref_with_vector_elems(memref<1x?xvector<10xf32>>)

mlir/unittests/IR/ShapedTypeTest.cpp

 //===- ShapedTypeTest.cpp - ShapedType unit tests -------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectInterface.h"
 #include "llvm/ADT/SmallVector.h"
 #include "gtest/gtest.h"
 #include <cstdint>
 
 using namespace mlir;
 using namespace mlir::detail;
 
 namespace {
 TEST(ShapedTypeTest, CloneMemref) {
   MLIRContext context;
 
   Type i32 = IntegerType::get(&context, 32);
   Type f32 = FloatType::getF32(&context);
-  int memSpace = 7;
+  Attribute memSpace = IntegerAttr::get(IntegerType::get(&context, 64), 7);
   Type memrefOriginalType = i32;
   llvm::SmallVector<int64_t> memrefOriginalShape({10, 20});
   AffineMap map = makeStridedLinearLayoutMap({2, 3}, 5, &context);
 
   ShapedType memrefType =
       MemRefType::Builder(memrefOriginalShape, memrefOriginalType)
           .setMemorySpace(memSpace)
           .setAffineMaps(map);