[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

184

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

1811

Add a comment to this funciton.

mlir/lib/IR/BuiltinTypes.cpp

441

nit: Spell out auto here.

442

nit: !attr

540

Use emitOptionalError here.

560

nit: memorySpace &&

560

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

457–468

All static methods should have doc comments as well.

538

Nit: if (memorySpace && ...

561

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.

184

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

561

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

[mehdi_amini]

mlir/include/mlir/IR/BuiltinTypes.h

184

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

183

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

mlir/lib/IR/BuiltinTypes.cpp

454

Nit: is allowed

457

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

183

You'll want to use LLVM_ATTRIBUTE_DEPRECATED.

mlir/lib/IR/BuiltinTypes.cpp

463

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

468–473

[ftynse]

mlir/include/mlir/IR/BuiltinTypes.h

183

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

183

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;
+
+  /// [deprecated] Returns the memory space in old raw integer representation.
+  /// New `Attribute getMemorySpace()` method should be used instead.
   unsigned getMemorySpaceAsInt() const;
 };
 
 } // end namespace mlir
 
 //===----------------------------------------------------------------------===//
 // Tablegen Type Declarations
 //===----------------------------------------------------------------------===//
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/IR/BuiltinTypes.h.inc"
 
 namespace mlir {
 //===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
 
 /// This is a builder type that keeps local references to arguments. Arguments
 /// that are passed into the builder must out-live the builder.
 class MemRefType::Builder {
 public:
   // Build from another MemRefType.
   explicit Builder(MemRefType other)
       : shape(other.getShape()), elementType(other.getElementType()),
-        affineMaps(other.getAffineMaps()),
-        memorySpace(other.getMemorySpaceAsInt()) {}
+        affineMaps(other.getAffineMaps()), 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;
 };
 
 //===----------------------------------------------------------------------===//
 // Deferred Method Definitions
 //===----------------------------------------------------------------------===//
 
 inline bool BaseMemRefType::classof(Type type) {
   return type.isa<MemRefType, UnrankedMemRefType>();

mlir/include/mlir/IR/BuiltinTypes.td

let description = [{
     ```
     memref-type ::= `memref` `<` dimension-list-ranked type
                     (`,` layout-specification)? (`,` memory-space)? `>`
 
     stride-list ::= `[` (dimension (`,` dimension)*)? `]`
     strided-layout ::= `offset:` dimension `,` `strides: ` stride-list
     semi-affine-map-composition ::= (semi-affine-map `,` )* semi-affine-map
     layout-specification ::= semi-affine-map-composition | strided-layout
-    memory-space ::= integer-literal /* | TODO: address-space-id */
+    memory-space ::= attribute-value
     ```
 
     A `memref` type is a reference to a region of memory (similar to a buffer
     pointer, but more powerful). The buffer pointed to by a memref can be
     allocated, aliased and deallocated. A memref can be used to read and write
     data from/to the memory region which it references. Memref types use the
     same shape specifier as tensor types. Note that `memref<f32>`,
     `memref<0 x f32>`, `memref<1 x 0 x f32>`, and `memref<0 x 1 x f32>` are all
let description = [{
     ```
 
     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
let description = [{
     are constructed with special ops, yet to be introduced. Strided memrefs are
     a special subclass of memrefs with generic semi-affine map and correspond to
     a normalized memref descriptor when lowering to LLVM.
   }];
   let parameters = (ins
     ArrayRefParameter<"int64_t">:$shape,
     "Type":$elementType,
     ArrayRefParameter<"AffineMap">:$affineMaps,
-    "unsigned":$memorySpaceAsInt
+    "Attribute":$memorySpace
   );
-
   let builders = [
     TypeBuilderWithInferredContext<(ins
       "ArrayRef<int64_t>":$shape, "Type":$elementType,
       CArg<"ArrayRef<AffineMap>", "{}">:$affineMaps,
-      CArg<"unsigned", "0">:$memorySpace
+      CArg<"Attribute", "{}">:$memorySpace
     ), [{
       // Drop identity maps from the composition. This may lead to the
       // composition becoming empty, which is interpreted as an implicit
       // identity.
       auto nonIdentityMaps = llvm::make_filter_range(affineMaps,
         [](AffineMap map) { return !map.isIdentity(); });
+      // Drop default memory space value and replace it with empty attribute.
+      Attribute nonDefaultMemorySpace = skipDefaultMemorySpace(memorySpace);
       return $_get(elementType.getContext(), shape, elementType,
-                   llvm::to_vector<4>(nonIdentityMaps), memorySpace);
+                   llvm::to_vector<4>(nonIdentityMaps), nonDefaultMemorySpace);
+    }]>,
+    /// [deprecated] `Attribute`-based form should be used instead.
+    TypeBuilderWithInferredContext<(ins
+      "ArrayRef<int64_t>":$shape, "Type":$elementType,
+      "ArrayRef<AffineMap>":$affineMaps,
+      "unsigned":$memorySpace
+    ), [{
+      // Convert deprecated integer-like memory space to Attribute.
+      Attribute memorySpaceAttr =
+          wrapIntegerMemorySpace(memorySpace, elementType.getContext());
+      return MemRefType::get(shape, elementType, affineMaps, memorySpaceAttr);
     }]>
   ];
   let extraClassDeclaration = [{
     /// 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;
 
+    /// [deprecated] Returns the memory space in old raw integer representation.
+    /// New `Attribute getMemorySpace()` method should be used instead.
+    unsigned getMemorySpaceAsInt() 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;
     }
   }];
   let skipDefaultBuilders = 1;

 
 def Builtin_UnrankedMemRef : Builtin_Type<"UnrankedMemRef", "BaseMemRefType"> {
   let summary = "Shaped reference, with unknown rank, to a region of memory";
   let description = [{
     Syntax:
 
     ```
     unranked-memref-type ::= `memref` `<*x` type (`,` memory-space)? `>`
-    memory-space ::= integer-literal /* | TODO: address-space-id */
+    memory-space ::= attribute-value
     ```
 
     A `memref` type with an unknown rank (e.g. `memref<*xf32>`). The purpose of
     unranked memrefs is to allow external library functions to receive memref
     arguments of any rank without versioning the functions based on the rank.
     Other uses of this type are disallowed or will have undefined behavior.
 
     See [MemRefType](#builtin_memref-memreftype) for more information on
     memref types.
 
     Examples:
 
     ```mlir
     memref<*f32>
 
     // An unranked memref with a memory space of 10.
     memref<*f32, 10>
     ```
   }];
-  let parameters = (ins "Type":$elementType, "unsigned":$memorySpaceAsInt);
+  let parameters = (ins "Type":$elementType, "Attribute":$memorySpace);
 
   let builders = [
     TypeBuilderWithInferredContext<(ins "Type":$elementType,
+                                        "Attribute":$memorySpace), [{
+      // Drop default memory space value and replace it with empty attribute.
+      Attribute nonDefaultMemorySpace = skipDefaultMemorySpace(memorySpace);
+      return $_get(elementType.getContext(), elementType, nonDefaultMemorySpace);
+    }]>,
+    /// [deprecated] `Attribute`-based form should be used instead.
+    TypeBuilderWithInferredContext<(ins "Type":$elementType,
                                         "unsigned":$memorySpace), [{
-      return $_get(elementType.getContext(), elementType, memorySpace);
+      // Convert deprecated integer-like memory space to Attribute.
+      Attribute memorySpaceAttr =
+          wrapIntegerMemorySpace(memorySpace, elementType.getContext());
+      return UnrankedMemRefType::get(elementType, memorySpaceAttr);
     }]>
   ];
   let extraClassDeclaration = [{
     ArrayRef<int64_t> getShape() const { return llvm::None; }
+
+    /// [deprecated] Returns the memory space in old raw integer representation.
+    /// New `Attribute getMemorySpace()` method should be used instead.
+    unsigned getMemorySpaceAsInt() const;
   }];
   let skipDefaultBuilders = 1;
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // UnrankedTensorType
 //===----------------------------------------------------------------------===//

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"
 #include "llvm/ADT/TypeSwitch.h"
 
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,
                            Type elementType) {
   return checkTensorElementType(emitError, elementType);
 }
 
 //===----------------------------------------------------------------------===//
 // BaseMemRefType
 //===----------------------------------------------------------------------===//
 
+Attribute BaseMemRefType::getMemorySpace() const {
+  if (auto rankedMemRefTy = dyn_cast<MemRefType>())
+    return rankedMemRefTy.getMemorySpace();
+  return cast<UnrankedMemRefType>().getMemorySpace();
+}
+
 unsigned BaseMemRefType::getMemorySpaceAsInt() const {
   if (auto rankedMemRefTy = dyn_cast<MemRefType>())
     return rankedMemRefTy.getMemorySpaceAsInt();
   return cast<UnrankedMemRefType>().getMemorySpaceAsInt();
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
 
+bool mlir::detail::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;
+}
+
+Attribute mlir::detail::wrapIntegerMemorySpace(unsigned memorySpace,
+                                               MLIRContext *ctx) {
+  if (memorySpace == 0)
+    return nullptr;
+
+  return IntegerAttr::get(IntegerType::get(ctx, 64), memorySpace);
+}
+
+Attribute mlir::detail::skipDefaultMemorySpace(Attribute memorySpace) {
+  IntegerAttr intMemorySpace = memorySpace.dyn_cast_or_null<IntegerAttr>();
+  if (intMemorySpace && intMemorySpace.getValue() == 0)
+    return nullptr;
+
+  return memorySpace;
+}
+
+unsigned mlir::detail::getMemorySpaceAsInt(Attribute memorySpace) {
+  if (!memorySpace)
+    return 0;
+
+  assert(memorySpace.isa<IntegerAttr>() &&
+         "Using `getMemorySpaceInteger` with non-Integer attribute");
+
+  return static_cast<unsigned>(memorySpace.cast<IntegerAttr>().getInt());
+}
+
+MemRefType::Builder &
+MemRefType::Builder::setMemorySpace(unsigned newMemorySpace) {
+  memorySpace =
+      wrapIntegerMemorySpace(newMemorySpace, elementType.getContext());
+  return *this;
+}
+
+unsigned MemRefType::getMemorySpaceAsInt() const {
+  return detail::getMemorySpaceAsInt(getMemorySpace());
+}
+
 LogicalResult MemRefType::verify(function_ref<InFlightDiagnostic()> emitError,
                                  ArrayRef<int64_t> shape, Type elementType,
                                  ArrayRef<AffineMap> affineMapComposition,
-                                 unsigned memorySpace) {
+                                 Attribute memorySpace) {
   if (!BaseMemRefType::isValidElementType(elementType))
     return emitError() << "invalid memref element type";
 
     // Negative sizes are not allowed except for `-1` that means dynamic size.
   for (int64_t s : shape)
     if (s < -1)
       return emitError() << "invalid memref size";
 
   // Check that the structure of the composition is valid, i.e. that each
   // subsequent affine map has as many inputs as the previous map has results.
   // Take the dimensionality of the MemRef for the first map.
   size_t dim = shape.size();
   for (auto it : llvm::enumerate(affineMapComposition)) {
     AffineMap map = it.value();
     if (map.getNumDims() == dim) {
       dim = map.getNumResults();
       continue;
     }
     return emitError() << "memref affine map dimension mismatch between "
                        << (it.index() == 0 ? Twine("memref rank")
                                            : "affine map " + Twine(it.index()))
                        << " and affine map" << it.index() + 1 << ": " << dim
                        << " != " << map.getNumDims();
   }
+
+  if (!isSupportedMemorySpace(memorySpace)) {
+    return emitError() << "unsupported memory space Attribute";
+  }
+
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // UnrankedMemRefType
 //===----------------------------------------------------------------------===//
 
+unsigned UnrankedMemRefType::getMemorySpaceAsInt() const {
+  return detail::getMemorySpaceAsInt(getMemorySpace());
+}
+
 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 TupleTypeStorage final
   ArrayRef<Type> getTypes() const {
     return {getTrailingObjects<Type>(), size()};
   }
 
   /// The number of tuple elements.
   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);
+
+/// Replaces default memorySpace (integer == `0`) with empty Attribute.
+Attribute skipDefaultMemorySpace(Attribute memorySpace);
+
+/// [deprecated] Returns the memory space in old raw integer representation.
+/// New `Attribute getMemorySpace()` method should be used instead.
+unsigned getMemorySpaceAsInt(Attribute memorySpace);
+
 } // namespace detail
 } // namespace mlir
 
 #endif // TYPEDETAIL_H_

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);