[mlir][NFC] Rename `MemRefType::getMemorySpace` to `getMemorySpaceAsInt`

[vinograd47]

Summary

Just a pure method renaming.

It is a preparation step for replacing "memory space as raw integer"
with more generic "memory space as attribute", which will be done in
separate commit.

The MemRefType::getMemorySpace method will return Attribute and
become the main API, while getMemorySpaceAsInt will be declared as
deprecated and will be replaced in all in-tree dialects (also in separate
commits).

[rriddle]

mlir/include/mlir/IR/BuiltinTypes.h

296

nit: I'd probably go with getMemorySpaceAsInteger, but I don't think it matters that much if it is going away soon.

[vinograd47]

Use getMemorySpaceAsInt instead of getMemorySpaceInteger.

[vinograd47]

mlir/include/mlir/IR/BuiltinTypes.h

296

@rriddle I've replaced it with getMemorySpaceAsInt, which is a bit shorter version of getMemorySpaceAsInteger. Are you OK with this?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

public:
   operand_range getSrcIndices() {
     return {operand_begin() + getSrcMemRefOperandIndex() + 1,
             operand_begin() + getSrcMemRefOperandIndex() + 1 +
                 getSrcMap().getNumInputs()};
   }
 
   /// Returns the memory space of the src memref.
   unsigned getSrcMemorySpace() {
-    return getSrcMemRef().getType().cast<MemRefType>().getMemorySpace();
+    return getSrcMemRef().getType().cast<MemRefType>().getMemorySpaceAsInt();
   }
 
   /// Returns the operand index of the dst memref.
   unsigned getDstMemRefOperandIndex() {
     return getSrcMemRefOperandIndex() + 1 + getSrcMap().getNumInputs();
   }
 
   /// Returns the destination MemRefType for this DMA operations.
   Value getDstMemRef() { return getOperand(getDstMemRefOperandIndex()); }
   MemRefType getDstMemRefType() {
     return getDstMemRef().getType().cast<MemRefType>();
   }
 
   /// Returns the rank (number of indices) of the destination MemRefType.
   unsigned getDstMemRefRank() {
     return getDstMemRef().getType().cast<MemRefType>().getRank();
   }
 
   /// Returns the memory space of the src memref.
   unsigned getDstMemorySpace() {
-    return getDstMemRef().getType().cast<MemRefType>().getMemorySpace();
+    return getDstMemRef().getType().cast<MemRefType>().getMemorySpaceAsInt();
   }
 
   /// Returns the affine map used to access the dst memref.
   AffineMap getDstMap() { return getDstMapAttr().getValue(); }
   AffineMapAttr getDstMapAttr() {
     return (*this)->getAttr(getDstMapAttrName()).cast<AffineMapAttr>();
   }
 

mlir/include/mlir/Dialect/StandardOps/IR/Ops.h

public:
 
   // Returns the destination MemRefType for this DMA operations.
   Value getDstMemRef() { return getOperand(1 + getSrcMemRefRank()); }
   // Returns the rank (number of indices) of the destination MemRefType.
   unsigned getDstMemRefRank() {
     return getDstMemRef().getType().cast<MemRefType>().getRank();
   }
   unsigned getSrcMemorySpace() {
-    return getSrcMemRef().getType().cast<MemRefType>().getMemorySpace();
+    return getSrcMemRef().getType().cast<MemRefType>().getMemorySpaceAsInt();
   }
   unsigned getDstMemorySpace() {
-    return getDstMemRef().getType().cast<MemRefType>().getMemorySpace();
+    return getDstMemRef().getType().cast<MemRefType>().getMemorySpaceAsInt();
   }
 
   // Returns the destination memref indices for this DMA operation.
   operand_range getDstIndices() {
     return {(*this)->operand_begin() + 1 + getSrcMemRefRank() + 1,
             (*this)->operand_begin() + 1 + getSrcMemRefRank() + 1 +
                 getDstMemRefRank()};
   }

mlir/include/mlir/IR/BuiltinTypes.h

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.
-  unsigned getMemorySpace() const;
+  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.getMemorySpace()) {}
+          memorySpace(other.getMemorySpaceAsInt()) {}
 
     // Build from scratch.
     Builder(ArrayRef<int64_t> shape, Type elementType)
         : shape(shape), elementType(elementType), affineMaps(), memorySpace(0) {
     }
 
     Builder &setShape(ArrayRef<int64_t> newShape) {
       shape = newShape;

mlir/lib/CAPI/IR/BuiltinTypes.cpp

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>().getMemorySpace();
+  return unwrap(type).cast<MemRefType>().getMemorySpaceAsInt();
 }
 
 bool mlirTypeIsAUnrankedMemRef(MlirType type) {
   return unwrap(type).isa<UnrankedMemRefType>();
 }
 
 MlirType mlirUnrankedMemRefTypeGet(MlirType elementType, unsigned memorySpace) {
   return wrap(UnrankedMemRefType::get(unwrap(elementType), memorySpace));
 }
 
 MlirType mlirUnrankedMemRefTypeGetChecked(MlirLocation loc,
                                           MlirType elementType,
                                           unsigned memorySpace) {
   return wrap(UnrankedMemRefType::getChecked(unwrap(loc), unwrap(elementType),
                                              memorySpace));
 }
 
 unsigned mlirUnrankedMemrefGetMemorySpace(MlirType type) {
-  return unwrap(type).cast<UnrankedMemRefType>().getMemorySpace();
+  return unwrap(type).cast<UnrankedMemRefType>().getMemorySpaceAsInt();
 }
 
 //===----------------------------------------------------------------------===//
 // Tuple type.
 //===----------------------------------------------------------------------===//
 
 bool mlirTypeIsATuple(MlirType type) { return unwrap(type).isa<TupleType>(); }
 

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

LowerToLLVMOptions options = {/*useBarePtrCallConv =*/false,
                                   /*useAlignedAlloc =*/false};
 
     /// MemRef conversion for GPU to NVVM lowering. The GPU dialect uses memory
     /// space 5 for private memory attributions, but NVVM represents private
     /// memory allocations as local `alloca`s in the default address space. This
     /// converter drops the private memory space to support the use case above.
     LLVMTypeConverter converter(m.getContext(), options);
     converter.addConversion([&](MemRefType type) -> Optional<Type> {
-      if (type.getMemorySpace() != gpu::GPUDialect::getPrivateAddressSpace())
+      if (type.getMemorySpaceAsInt() !=
+          gpu::GPUDialect::getPrivateAddressSpace())
         return llvm::None;
       return converter.convertType(MemRefType::Builder(type).setMemorySpace(0));
     });
 
     OwningRewritePatternList patterns, llvmPatterns;
 
     // Apply in-dialect lowering first. In-dialect lowering will replace ops
     // which need to be lowered further, which is not supported by a single

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 LLVMTypeConverter::getMemRefDescriptorFields(MemRefType type,
                                              bool unpackAggregates) {
   assert(isStrided(type) &&
          "Non-strided layout maps must have been normalized away");
 
   Type elementType = unwrap(convertType(type.getElementType()));
   if (!elementType)
     return {};
-  auto ptrTy = LLVM::LLVMPointerType::get(elementType, type.getMemorySpace());
+  auto ptrTy =
+      LLVM::LLVMPointerType::get(elementType, type.getMemorySpaceAsInt());
   auto indexTy = getIndexType();
 
   SmallVector<Type, 5> results = {ptrTy, ptrTy, indexTy};
   auto rank = type.getRank();
   if (rank == 0)
     return results;
 
   if (unpackAggregates)
if (ShapedType::isDynamicStrideOrOffset(stride))
       return {};
 
   if (ShapedType::isDynamicStrideOrOffset(offset))
     return {};
 
   Type elementType = unwrap(convertType(type.getElementType()));
   if (!elementType)
     return {};
-  return LLVM::LLVMPointerType::get(elementType, type.getMemorySpace());
+  return LLVM::LLVMPointerType::get(elementType, type.getMemorySpaceAsInt());
 }
 
 /// Convert an n-D vector type to an LLVM vector type via (n-1)-D array type
 /// when n > 1. For example, `vector<4 x f32>` remains as is while,
 /// `vector<4x8x16xf32>` converts to `!llvm.array<4xarray<8 x vector<16xf32>>>`.
 Type LLVMTypeConverter::convertVectorType(VectorType type) {
   auto elementType = unwrap(convertType(type.getElementType()));
   if (!elementType)
bool ConvertToLLVMPattern::isConvertibleAndHasIdentityMaps(
   return type.getAffineMaps().empty() ||
          llvm::all_of(type.getAffineMaps(),
                       [](AffineMap map) { return map.isIdentity(); });
 }
 
 Type ConvertToLLVMPattern::getElementPtrType(MemRefType type) const {
   auto elementType = type.getElementType();
   auto structElementType = unwrap(typeConverter->convertType(elementType));
-  return LLVM::LLVMPointerType::get(structElementType, type.getMemorySpace());
+  return LLVM::LLVMPointerType::get(structElementType,
+                                    type.getMemorySpaceAsInt());
 }
 
 void ConvertToLLVMPattern::getMemRefDescriptorSizes(
     Location loc, MemRefType memRefType, ValueRange dynamicSizes,
     ConversionPatternRewriter &rewriter, SmallVectorImpl<Value> &sizes,
     SmallVectorImpl<Value> &strides, Value &sizeBytes) const {
   assert(isConvertibleAndHasIdentityMaps(memRefType) &&
          "layout maps must have been normalized away");
auto allocFuncOp = LLVM::lookupOrCreateMallocFn(
         allocOp->getParentOfType<ModuleOp>(), getIndexType());
     auto results = createLLVMCall(rewriter, loc, allocFuncOp, {sizeBytes},
                                   getVoidPtrType());
     Value allocatedPtr =
         rewriter.create<LLVM::BitcastOp>(loc, elementPtrType, results[0]);
 
     Value alignedPtr = allocatedPtr;
     if (alignment) {
-      auto intPtrType = getIntPtrType(memRefType.getMemorySpace());
+      auto intPtrType = getIntPtrType(memRefType.getMemorySpaceAsInt());
       // Compute the aligned type pointer.
       Value allocatedInt =
           rewriter.create<LLVM::PtrToIntOp>(loc, intPtrType, allocatedPtr);
       Value alignmentInt =
           createAligned(rewriter, loc, allocatedInt, alignment);
       alignedPtr =
           rewriter.create<LLVM::IntToPtrOp>(loc, elementPtrType, alignmentInt);
     }
if (!global.isExternal() && !global.isUninitialized()) {
       // For scalar memrefs, the global variable created is of the element type,
       // so unpack the elements attribute to extract the value.
       if (type.getRank() == 0)
         initialValue = elementsAttr.getValue({});
     }
 
     rewriter.replaceOpWithNewOp<LLVM::GlobalOp>(
         global, arrayTy, global.constant(), linkage, global.sym_name(),
-        initialValue, type.getMemorySpace());
+        initialValue, type.getMemorySpaceAsInt());
     return success();
   }
 };
 
 /// GetGlobalMemrefOp is lowered into a Memref descriptor with the pointer to
 /// the first element stashed into the descriptor. This reuses
 /// `AllocLikeOpLowering` to reuse the Memref descriptor construction.
 struct GetGlobalMemrefOpLowering : public AllocLikeOpLowering {
   GetGlobalMemrefOpLowering(LLVMTypeConverter &converter)
       : AllocLikeOpLowering(GetGlobalMemrefOp::getOperationName(), converter) {}
 
   /// Buffer "allocation" for get_global_memref op is getting the address of
   /// the global variable referenced.
   std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                           Location loc, Value sizeBytes,
                                           Operation *op) const override {
     auto getGlobalOp = cast<GetGlobalMemrefOp>(op);
     MemRefType type = getGlobalOp.result().getType().cast<MemRefType>();
-    unsigned memSpace = type.getMemorySpace();
+    unsigned memSpace = type.getMemorySpaceAsInt();
 
     Type arrayTy = convertGlobalMemrefTypeToLLVM(type, *getTypeConverter());
     auto addressOf = rewriter.create<LLVM::AddressOfOp>(
         loc, LLVM::LLVMPointerType::get(arrayTy, memSpace), getGlobalOp.name());
 
     // Get the address of the first element in the array by creating a GEP with
     // the address of the GV as the base, and (rank + 1) number of 0 indices.
     Type elementType =
if (operandType.isa<MemRefType>()) {
     *allocatedPtr = desc.allocatedPtr(rewriter, loc);
     *alignedPtr = desc.alignedPtr(rewriter, loc);
     if (offset != nullptr)
       *offset = desc.offset(rewriter, loc);
     return;
   }
 
   unsigned memorySpace =
-      operandType.cast<UnrankedMemRefType>().getMemorySpace();
+      operandType.cast<UnrankedMemRefType>().getMemorySpaceAsInt();
   Type elementType = operandType.cast<UnrankedMemRefType>().getElementType();
   Type llvmElementType = unwrap(typeConverter.convertType(elementType));
   Type elementPtrPtrType = LLVM::LLVMPointerType::get(
       LLVM::LLVMPointerType::get(llvmElementType, memorySpace));
 
   // Extract pointer to the underlying ranked memref descriptor and cast it to
   // ElemType**.
   UnrankedMemRefDescriptor unrankedDesc(convertedOperand);
convertSourceMemRefToDescriptor(ConversionPatternRewriter &rewriter,
     // The shape is a rank-1 tensor with unknown length.
     Location loc = reshapeOp.getLoc();
     MemRefDescriptor shapeDesc(adaptor.shape());
     Value resultRank = shapeDesc.size(rewriter, loc, 0);
 
     // Extract address space and element type.
     auto targetType =
         reshapeOp.getResult().getType().cast<UnrankedMemRefType>();
-    unsigned addressSpace = targetType.getMemorySpace();
+    unsigned addressSpace = targetType.getMemorySpaceAsInt();
     Type elementType = targetType.getElementType();
 
     // Create the unranked memref descriptor that holds the ranked one. The
     // inner descriptor is allocated on stack.
     auto targetDesc = UnrankedMemRefDescriptor::undef(
         rewriter, loc, unwrap(typeConverter->convertType(targetType)));
     targetDesc.setRank(rewriter, loc, resultRank);
     SmallVector<Value, 4> sizes;
Value extractSizeOfUnrankedMemRef(Type operandType, DimOp dimOp,
                                     ArrayRef<Value> operands,
                                     ConversionPatternRewriter &rewriter) const {
     Location loc = dimOp.getLoc();
     DimOp::Adaptor transformed(operands);
 
     auto unrankedMemRefType = operandType.cast<UnrankedMemRefType>();
     auto scalarMemRefType =
         MemRefType::get({}, unrankedMemRefType.getElementType());
-    unsigned addressSpace = unrankedMemRefType.getMemorySpace();
+    unsigned addressSpace = unrankedMemRefType.getMemorySpaceAsInt();
 
     // Extract pointer to the underlying ranked descriptor and bitcast it to a
     // memref<element_type> descriptor pointer to minimize the number of GEP
     // operations.
     UnrankedMemRefDescriptor unrankedDesc(transformed.memrefOrTensor());
     Value underlyingRankedDesc = unrankedDesc.memRefDescPtr(rewriter, loc);
     Value scalarMemRefDescPtr = rewriter.create<LLVM::BitcastOp>(
         loc,
matchAndRewrite(SubViewOp subViewOp, ArrayRef<Value> operands,
     MemRefDescriptor sourceMemRef(operands.front());
     auto targetMemRef = MemRefDescriptor::undef(rewriter, loc, targetDescTy);
 
     // Copy the buffer pointer from the old descriptor to the new one.
     Value extracted = sourceMemRef.allocatedPtr(rewriter, loc);
     Value bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc,
         LLVM::LLVMPointerType::get(targetElementTy,
-                                   viewMemRefType.getMemorySpace()),
+                                   viewMemRefType.getMemorySpaceAsInt()),
         extracted);
     targetMemRef.setAllocatedPtr(rewriter, loc, bitcastPtr);
 
     // Copy the aligned pointer from the old descriptor to the new one.
     extracted = sourceMemRef.alignedPtr(rewriter, loc);
     bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc,
         LLVM::LLVMPointerType::get(targetElementTy,
-                                   viewMemRefType.getMemorySpace()),
+                                   viewMemRefType.getMemorySpaceAsInt()),
         extracted);
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
     auto shape = viewMemRefType.getShape();
     auto inferredShape = inferredType.getShape();
     size_t inferredShapeRank = inferredShape.size();
     size_t resultShapeRank = shape.size();
     llvm::SmallDenseSet<unsigned> unusedDims =
matchAndRewrite(ViewOp viewOp, ArrayRef<Value> operands,
     auto targetMemRef = MemRefDescriptor::undef(rewriter, loc, targetDescTy);
 
     // Field 1: Copy the allocated pointer, used for malloc/free.
     Value allocatedPtr = sourceMemRef.allocatedPtr(rewriter, loc);
     auto srcMemRefType = viewOp.source().getType().cast<MemRefType>();
     Value bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc,
         LLVM::LLVMPointerType::get(targetElementTy,
-                                   srcMemRefType.getMemorySpace()),
+                                   srcMemRefType.getMemorySpaceAsInt()),
         allocatedPtr);
     targetMemRef.setAllocatedPtr(rewriter, loc, bitcastPtr);
 
     // Field 2: Copy the actual aligned pointer to payload.
     Value alignedPtr = sourceMemRef.alignedPtr(rewriter, loc);
     alignedPtr = rewriter.create<LLVM::GEPOp>(loc, alignedPtr.getType(),
                                               alignedPtr, adaptor.byte_shift());
     bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc,
         LLVM::LLVMPointerType::get(targetElementTy,
-                                   srcMemRefType.getMemorySpace()),
+                                   srcMemRefType.getMemorySpaceAsInt()),
         alignedPtr);
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
     // Field 3: The offset in the resulting type must be 0. This is because of
     // the type change: an offset on srcType* may not be expressible as an
     // offset on dstType*.
     targetMemRef.setOffset(rewriter, loc,
                            createIndexConstant(rewriter, loc, offset));

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

 }
 
 /// Returns true if the allocations of type `t` can be lowered to SPIR-V.
 static bool isAllocationSupported(MemRefType t) {
   // Currently only support workgroup local memory allocations with static
   // shape and int or float or vector of int or float element type.
   if (!(t.hasStaticShape() &&
         SPIRVTypeConverter::getMemorySpaceForStorageClass(
-            spirv::StorageClass::Workgroup) == t.getMemorySpace()))
+            spirv::StorageClass::Workgroup) == t.getMemorySpaceAsInt()))
     return false;
   Type elementType = t.getElementType();
   if (auto vecType = elementType.dyn_cast<VectorType>())
     elementType = vecType.getElementType();
   return elementType.isIntOrFloat();
 }
 
 /// Returns the scope to use for atomic operations use for emulating store
 /// operations of unsupported integer bitwidths, based on the memref
 /// type. Returns None on failure.
 static Optional<spirv::Scope> getAtomicOpScope(MemRefType t) {
   Optional<spirv::StorageClass> storageClass =
-      SPIRVTypeConverter::getStorageClassForMemorySpace(t.getMemorySpace());
+      SPIRVTypeConverter::getStorageClassForMemorySpace(
+          t.getMemorySpaceAsInt());
   if (!storageClass)
     return {};
   switch (*storageClass) {
   case spirv::StorageClass::StorageBuffer:
     return spirv::Scope::Device;
   case spirv::StorageClass::Workgroup:
     return spirv::Scope::Workgroup;
   default: {

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 static LogicalResult getIndexedPtrs(ConversionPatternRewriter &rewriter,
                                     Location loc, Value memref, Value base,
                                     Value index, MemRefType memRefType,
                                     VectorType vType, Value &ptrs) {
   int64_t offset;
   SmallVector<int64_t, 4> strides;
   auto successStrides = getStridesAndOffset(memRefType, strides, offset);
   if (failed(successStrides) || strides.back() != 1 ||
-      memRefType.getMemorySpace() != 0)
+      memRefType.getMemorySpaceAsInt() != 0)
     return failure();
   auto pType = MemRefDescriptor(memref).getElementPtrType();
   auto ptrsType = LLVM::getFixedVectorType(pType, vType.getDimSize(0));
   ptrs = rewriter.create<LLVM::GEPOp>(loc, ptrsType, base, index);
   return success();
 }
 
 // Casts a strided element pointer to a vector pointer.  The vector pointer
 // will be in the same address space as the incoming memref type.
 static Value castDataPtr(ConversionPatternRewriter &rewriter, Location loc,
                          Value ptr, MemRefType memRefType, Type vt) {
-  auto pType = LLVM::LLVMPointerType::get(vt, memRefType.getMemorySpace());
+  auto pType = LLVM::LLVMPointerType::get(vt, memRefType.getMemorySpaceAsInt());
   return rewriter.create<LLVM::BitcastOp>(loc, pType, ptr);
 }
 
 static LogicalResult
 replaceTransferOpWithLoadOrStore(ConversionPatternRewriter &rewriter,
                                  LLVMTypeConverter &typeConverter, Location loc,
                                  TransferReadOp xferOp,
                                  ArrayRef<Value> operands, Value dataPtr) {

mlir/lib/Conversion/VectorToROCDL/VectorToROCDL.cpp

matchAndRewrite(ConcreteOp xferOp, ArrayRef<Value> operands,
 
     // Obtain dataPtr and elementType from the memref.
     auto memRefType = xferOp.getShapedType().template dyn_cast<MemRefType>();
     if (!memRefType)
       return failure();
     // MUBUF instruction operate only on addresspace 0(unified) or 1(global)
     // In case of 3(LDS): fall back to vector->llvm pass
     // In case of 5(VGPR): wrong
-    if ((memRefType.getMemorySpace() != 0) &&
-        (memRefType.getMemorySpace() != 1))
+    if ((memRefType.getMemorySpaceAsInt() != 0) &&
+        (memRefType.getMemorySpaceAsInt() != 1))
       return failure();
 
     // Note that the dataPtr starts at the offset address specified by
     // indices, so no need to calculate offset size in bytes again in
     // the MUBUF instruction.
     Value dataPtr = this->getStridedElementPtr(
         loc, memRefType, adaptor.source(), adaptor.indices(), rewriter);
 

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

NDTransferOpHelper(PatternRewriter &rewriter, ConcreteOp xferOp,
     leadingRank = xferOp.getLeadingShapedRank();
     majorVectorType =
         VectorType::get(vectorType.getShape().take_front(majorRank),
                         vectorType.getElementType());
     minorVectorType =
         VectorType::get(vectorType.getShape().take_back(minorRank),
                         vectorType.getElementType());
     /// Memref of minor vector type is used for individual transfers.
-    memRefMinorVectorType = MemRefType::get(
-        majorVectorType.getShape(), minorVectorType, {},
-        xferOp.getShapedType().template cast<MemRefType>().getMemorySpace());
+    memRefMinorVectorType =
+        MemRefType::get(majorVectorType.getShape(), minorVectorType, {},
+                        xferOp.getShapedType()
+                            .template cast<MemRefType>()
+                            .getMemorySpaceAsInt());
   }
 
   LogicalResult doReplace();
 
 private:
   /// Creates the loop nest on the "major" dimensions and calls the
   /// `loopBodyBuilder` lambda in the context of the loop nest.
   void

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 static LogicalResult verifyAttributions(Operation *op,
                                         ArrayRef<BlockArgument> attributions,
                                         unsigned memorySpace) {
   for (Value v : attributions) {
     auto type = v.getType().dyn_cast<MemRefType>();
     if (!type)
       return op->emitOpError() << "expected memref type in attribution";
 
-    if (type.getMemorySpace() != memorySpace) {
+    if (type.getMemorySpaceAsInt() != memorySpace) {
       return op->emitOpError()
              << "expected memory space " << memorySpace << " in attribution";
     }
   }
   return success();
 }
 
 /// Verifies the body of the function.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

static LogicalResult verifyCast(DialectCastOp op, Type llvmType, Type type,
 
   if (auto memrefType = type.dyn_cast<MemRefType>()) {
     // Bare pointer convention: statically-shaped memref is compatible with an
     // LLVM pointer to the element type.
     if (auto ptrType = llvmType.dyn_cast<LLVMPointerType>()) {
       if (!memrefType.hasStaticShape())
         return op->emitOpError(
             "unexpected bare pointer for dynamically shaped memref");
-      if (memrefType.getMemorySpace() != ptrType.getAddressSpace())
+      if (memrefType.getMemorySpaceAsInt() != ptrType.getAddressSpace())
         return op->emitError("invalid conversion between memref and pointer in "
                              "different memory spaces");
 
       return verifyCast(op, ptrType.getElementType(),
                         memrefType.getElementType(), /*isElement=*/true);
     }
 
     // Otherwise, memrefs are convertible to a descriptor, which is a structure
     // type.
     auto structType = llvmType.dyn_cast<LLVMStructType>();
     if (!structType)
       return op->emitOpError("invalid cast between a memref and a type other "
                              "than pointer or memref descriptor");
 
     unsigned expectedNumElements = memrefType.getRank() == 0 ? 3 : 5;
     if (structType.getBody().size() != expectedNumElements) {
       return op->emitOpError() << "expected memref descriptor with "
                                << expectedNumElements << " elements";
     }
 
     // The first two elements are pointers to the element type.
     auto allocatedPtr = structType.getBody()[0].dyn_cast<LLVMPointerType>();
     if (!allocatedPtr ||
-        allocatedPtr.getAddressSpace() != memrefType.getMemorySpace())
+        allocatedPtr.getAddressSpace() != memrefType.getMemorySpaceAsInt())
       return op->emitOpError("expected first element of a memref descriptor to "
                              "be a pointer in the address space of the memref");
     if (failed(verifyCast(op, allocatedPtr.getElementType(),
                           memrefType.getElementType(), /*isElement=*/true)))
       return failure();
 
     auto alignedPtr = structType.getBody()[1].dyn_cast<LLVMPointerType>();
     if (!alignedPtr ||
-        alignedPtr.getAddressSpace() != memrefType.getMemorySpace())
+        alignedPtr.getAddressSpace() != memrefType.getMemorySpaceAsInt())
       return op->emitOpError(
           "expected second element of a memref descriptor to "
           "be a pointer in the address space of the memref");
     if (failed(verifyCast(op, alignedPtr.getElementType(),
                           memrefType.getElementType(), /*isElement=*/true)))
       return failure();
 
     // The second element (offset) is an equivalent of index.

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

static Optional<Type> convertTensorType(const spirv::TargetEnv &targetEnv,
   }
 
   return spirv::ArrayType::get(*arrayElemType, arrayElemCount, *arrayElemSize);
 }
 
 static Optional<Type> convertMemrefType(const spirv::TargetEnv &targetEnv,
                                         MemRefType type) {
   Optional<spirv::StorageClass> storageClass =
-      SPIRVTypeConverter::getStorageClassForMemorySpace(type.getMemorySpace());
+      SPIRVTypeConverter::getStorageClassForMemorySpace(
+          type.getMemorySpaceAsInt());
   if (!storageClass) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot convert memory space\n");
     return llvm::None;
   }
 
   Optional<Type> arrayElemType;
   Type elementType = type.getElementType();

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

if (aT.getAffineMaps() != bT.getAffineMaps()) {
                 b == MemRefType::getDynamicStrideOrOffset() || a == b);
       };
       if (!checkCompatible(aOffset, bOffset))
         return false;
       for (auto aStride : enumerate(aStrides))
         if (!checkCompatible(aStride.value(), bStrides[aStride.index()]))
           return false;
     }
-    if (aT.getMemorySpace() != bT.getMemorySpace())
+    if (aT.getMemorySpaceAsInt() != bT.getMemorySpaceAsInt())
       return false;
 
     // They must have the same rank, and any specified dimensions must match.
     if (aT.getRank() != bT.getRank())
       return false;
 
     for (unsigned i = 0, e = aT.getRank(); i != e; ++i) {
       int64_t aDim = aT.getDimSize(i), bDim = bT.getDimSize(i);
if (aT && bT) {
     if (uaT && ubT)
       return false;
 
     auto aEltType = (aT) ? aT.getElementType() : uaT.getElementType();
     auto bEltType = (bT) ? bT.getElementType() : ubT.getElementType();
     if (aEltType != bEltType)
       return false;
 
-    auto aMemSpace = (aT) ? aT.getMemorySpace() : uaT.getMemorySpace();
-    auto bMemSpace = (bT) ? bT.getMemorySpace() : ubT.getMemorySpace();
+    auto aMemSpace =
+        (aT) ? aT.getMemorySpaceAsInt() : uaT.getMemorySpaceAsInt();
+    auto bMemSpace =
+        (bT) ? bT.getMemorySpaceAsInt() : ubT.getMemorySpaceAsInt();
     if (aMemSpace != bMemSpace)
       return false;
 
     return true;
   }
 
   return false;
 }

 }
 
 // TODO: ponder whether we want to allow missing trailing sizes/strides that are
 // completed automatically, like we have for subview and subtensor.
 static LogicalResult verify(MemRefReinterpretCastOp op) {
   // The source and result memrefs should be in the same memory space.
   auto srcType = op.source().getType().cast<BaseMemRefType>();
   auto resultType = op.getType().cast<MemRefType>();
-  if (srcType.getMemorySpace() != resultType.getMemorySpace())
+  if (srcType.getMemorySpaceAsInt() != resultType.getMemorySpaceAsInt())
     return op.emitError("different memory spaces specified for source type ")
            << srcType << " and result memref type " << resultType;
   if (srcType.getElementType() != resultType.getElementType())
     return op.emitError("different element types specified for source type ")
            << srcType << " and result memref type " << resultType;
 
   // Match sizes in result memref type and in static_sizes attribute.
   for (auto &en :
for (auto it : llvm::zip(sourceStrides, staticStrides)) {
     targetStrides.push_back(Wrapper(sourceStride) * staticStride);
   }
 
   // The type is now known.
   return MemRefType::get(
       staticSizes, sourceMemRefType.getElementType(),
       makeStridedLinearLayoutMap(targetStrides, targetOffset,
                                  sourceMemRefType.getContext()),
-      sourceMemRefType.getMemorySpace());
+      sourceMemRefType.getMemorySpaceAsInt());
 }
 
 Type SubViewOp::inferResultType(MemRefType sourceMemRefType,
                                 ArrayRef<OpFoldResult> leadingStaticOffsets,
                                 ArrayRef<OpFoldResult> leadingStaticSizes,
                                 ArrayRef<OpFoldResult> leadingStaticStrides) {
   SmallVector<int64_t> staticOffsets, staticSizes, staticStrides;
   SmallVector<Value> dynamicOffsets, dynamicSizes, dynamicStrides;
for (unsigned pos = 0, e = shape.size(); pos < e; ++pos)
         projectedShape.push_back(shape[pos]);
 
     AffineMap map;
     auto maps = inferredType.getAffineMaps();
     if (!maps.empty() && maps.front())
       map = getProjectedMap(maps.front(), dimsToProject);
     inferredType =
         MemRefType::get(projectedShape, inferredType.getElementType(), map,
-                        inferredType.getMemorySpace());
+                        inferredType.getMemorySpaceAsInt());
   }
   return inferredType;
 }
 
 Type SubViewOp::inferRankReducedResultType(
     unsigned resultRank, MemRefType sourceRankedTensorType,
     ArrayRef<OpFoldResult> leadingStaticOffsets,
     ArrayRef<OpFoldResult> leadingStaticSizes,
isRankReducedType(Type originalType, Type candidateReducedType,
 
   // We are done for the tensor case.
   if (originalType.isa<RankedTensorType>())
     return SubViewVerificationResult::Success;
 
   // Strided layout logic is relevant for MemRefType only.
   MemRefType original = originalType.cast<MemRefType>();
   MemRefType candidateReduced = candidateReducedType.cast<MemRefType>();
-  if (original.getMemorySpace() != candidateReduced.getMemorySpace())
+  if (original.getMemorySpaceAsInt() != candidateReduced.getMemorySpaceAsInt())
     return SubViewVerificationResult::MemSpaceMismatch;
 
   llvm::SmallDenseSet<unsigned> unusedDims = optionalUnusedDimsMask.getValue();
   auto inferredType =
       getProjectedMap(getStridedLinearLayoutMap(original), unusedDims);
   AffineMap candidateLayout;
   if (candidateReduced.getAffineMaps().empty())
     candidateLayout = getStridedLinearLayoutMap(candidateReduced);

 }
 
 /// Verifier for SubViewOp.
 static LogicalResult verify(SubViewOp op) {
   MemRefType baseType = op.getSourceType();
   MemRefType subViewType = op.getType();
 
   // The base memref and the view memref should be in the same memory space.
-  if (baseType.getMemorySpace() != subViewType.getMemorySpace())
+  if (baseType.getMemorySpaceAsInt() != subViewType.getMemorySpaceAsInt())
     return op.emitError("different memory spaces specified for base memref "
                         "type ")
            << baseType << " and subview memref type " << subViewType;
 
   // Verify that the base memref type has a strided layout map.
   if (!isStrided(baseType))
     return op.emitError("base type ") << baseType << " is not strided";
 
static LogicalResult verify(ViewOp op) {
 
   // The result memref should have identity layout map (or none).
   if (viewType.getAffineMaps().size() > 1 ||
       (viewType.getAffineMaps().size() == 1 &&
        !viewType.getAffineMaps()[0].isIdentity()))
     return op.emitError("unsupported map for result memref type ") << viewType;
 
   // The base memref and the view memref should be in the same memory space.
-  if (baseType.getMemorySpace() != viewType.getMemorySpace())
+  if (baseType.getMemorySpaceAsInt() != viewType.getMemorySpaceAsInt())
     return op.emitError("different memory spaces specified for base memref "
                         "type ")
            << baseType << " and view memref type " << viewType;
 
   // Verify that we have the correct number of sizes for the result type.
   unsigned numDynamicDims = viewType.getNumDynamicDims();
   if (op.sizes().size() != numDynamicDims)
     return op.emitError("incorrect number of size operands for type ")

mlir/lib/Dialect/Vector/VectorOps.cpp

 void TypeCastOp::build(OpBuilder &builder, OperationState &result,
                        Value source) {
   result.addOperands(source);
   MemRefType memRefType = source.getType().cast<MemRefType>();
   VectorType vectorType =
       VectorType::get(extractShape(memRefType),
                       getElementTypeOrSelf(getElementTypeOrSelf(memRefType)));
   result.addTypes(
-      MemRefType::get({}, vectorType, {}, memRefType.getMemorySpace()));
+      MemRefType::get({}, vectorType, {}, memRefType.getMemorySpaceAsInt()));
 }
 
 static LogicalResult verify(TypeCastOp op) {
   MemRefType canonicalType = canonicalizeStridedLayout(op.getMemRefType());
   if (!canonicalType.getAffineMaps().empty())
     return op.emitOpError("expects operand to be a memref with no layout");
   if (!op.getResultMemRefType().getAffineMaps().empty())
     return op.emitOpError("expects result to be a memref with no layout");
-  if (op.getResultMemRefType().getMemorySpace() !=
-      op.getMemRefType().getMemorySpace())
+  if (op.getResultMemRefType().getMemorySpaceAsInt() !=
+      op.getMemRefType().getMemorySpaceAsInt())
     return op.emitOpError("expects result in same memory space");
 
   auto sourceType = op.getMemRefType();
   auto resultType = op.getResultMemRefType();
   if (getElementTypeOrSelf(getElementTypeOrSelf(sourceType)) !=
       getElementTypeOrSelf(getElementTypeOrSelf(resultType)))
     return op.emitOpError(
                "expects result and operand with same underlying scalar type: ")

mlir/lib/IR/AsmPrinter.cpp

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.getMemorySpace())
-          os << ", " << memrefTy.getMemorySpace();
+        if (memrefTy.getMemorySpaceAsInt())
+          os << ", " << memrefTy.getMemorySpaceAsInt();
         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.getMemorySpace())
-          os << ", " << memrefTy.getMemorySpace();
+        if (memrefTy.getMemorySpaceAsInt())
+          os << ", " << memrefTy.getMemorySpaceAsInt();
         os << '>';
       })
       .Case<ComplexType>([&](ComplexType complexTy) {
         os << "complex<";
         printType(complexTy.getElementType());
         os << '>';
       })
       .Case<TupleType>([&](TupleType tupleTy) {

mlir/lib/IR/BuiltinTypes.cpp

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.getMemorySpace());
+    b.setMemorySpace(other.getMemorySpaceAsInt());
     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.getMemorySpace());
+    b.setMemorySpace(other.getMemorySpaceAsInt());
     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.getMemorySpace());
+    return UnrankedMemRefType::get(elementType, other.getMemorySpaceAsInt());
   }
 
   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
 //===----------------------------------------------------------------------===//
 
-unsigned BaseMemRefType::getMemorySpace() const {
+unsigned BaseMemRefType::getMemorySpaceAsInt() const {
   return static_cast<ImplType *>(impl)->memorySpace;
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
 
 /// Get or create a new MemRefType based on shape, element type, affine

mlir/lib/Transforms/LoopFusion.cpp

static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
   // Create 'newMemRefType' using 'newShape' from MemRefRegion accessed
   // by 'srcStoreOpInst'.
   uint64_t bufSize =
       getMemRefEltSizeInBytes(oldMemRefType) * numElements.getValue();
   unsigned newMemSpace;
   if (bufSize <= localBufSizeThreshold && fastMemorySpace.hasValue()) {
     newMemSpace = fastMemorySpace.getValue();
   } else {
-    newMemSpace = oldMemRefType.getMemorySpace();
+    newMemSpace = oldMemRefType.getMemorySpaceAsInt();
   }
   auto newMemRefType = MemRefType::get(newShape, oldMemRefType.getElementType(),
                                        {}, newMemSpace);
 
   // Create new private memref for fused loop 'forOp'. 'newShape' is always
   // a constant shape.
   // TODO: Create/move alloc ops for private memrefs closer to their
   // consumer loop nests to reduce their live range. Currently they are added

mlir/lib/Transforms/Utils/LoopUtils.cpp

uint64_t mlir::affineDataCopyGenerate(Block::iterator begin,
   // To check for errors when walking the block.
   bool error = false;
 
   // Walk this range of operations  to gather all memory regions.
   block->walk(begin, end, [&](Operation *opInst) {
     // Gather regions to allocate to buffers in faster memory space.
     if (auto loadOp = dyn_cast<AffineLoadOp>(opInst)) {
       if ((filterMemRef.hasValue() && filterMemRef != loadOp.getMemRef()) ||
-          (loadOp.getMemRefType().getMemorySpace() !=
+          (loadOp.getMemRefType().getMemorySpaceAsInt() !=
            copyOptions.slowMemorySpace))
         return;
     } else if (auto storeOp = dyn_cast<AffineStoreOp>(opInst)) {
       if ((filterMemRef.hasValue() && filterMemRef != storeOp.getMemRef()) ||
-          storeOp.getMemRefType().getMemorySpace() !=
+          storeOp.getMemRefType().getMemorySpaceAsInt() !=
               copyOptions.slowMemorySpace)
         return;
     } else {
       // Neither load nor a store op.
       return;
     }
 
     // Compute the MemRefRegion accessed.