[mlir][NFC] Replace `getMemorySpaceAsInt` with `getMemorySpace` where possible

[vinograd47]

Summary

Use new MemRefType::getMemorySpace method with generic Attribute
in cases, where there is no specific logic around the memory space.

[mehdi_amini]

Thanks for the cleanup!

[rriddle]

Nice!

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>()
-                            .getMemorySpaceAsInt());
+    memRefMinorVectorType = MemRefType::get(
+        majorVectorType.getShape(), minorVectorType, {},
+        xferOp.getShapedType().template cast<MemRefType>().getMemorySpace());
   }
 
   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/MemRef/IR/MemRefOps.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.getMemorySpaceAsInt() != bT.getMemorySpaceAsInt())
+    if (aT.getMemorySpace() != bT.getMemorySpace())
       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.getMemorySpaceAsInt() : uaT.getMemorySpaceAsInt();
-    auto bMemSpace =
-        (bT) ? bT.getMemorySpaceAsInt() : ubT.getMemorySpaceAsInt();
+    auto aMemSpace = (aT) ? aT.getMemorySpace() : uaT.getMemorySpace();
+    auto bMemSpace = (bT) ? bT.getMemorySpace() : ubT.getMemorySpace();
     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(ReinterpretCastOp 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.getMemorySpaceAsInt() != resultType.getMemorySpaceAsInt())
+  if (srcType.getMemorySpace() != resultType.getMemorySpace())
     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.getMemorySpaceAsInt());
+      sourceMemRefType.getMemorySpace());
 }
 
 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.getMemorySpaceAsInt());
+                        inferredType.getMemorySpace());
   }
   return inferredType;
 }
 
 Type SubViewOp::inferRankReducedResultType(
     unsigned resultRank, MemRefType sourceRankedTensorType,
     ArrayRef<OpFoldResult> leadingStaticOffsets,
     ArrayRef<OpFoldResult> leadingStaticSizes,
isRankReducedType(Type originalType, Type candidateReducedType,
 
   if (originalShapedType.getElementType() !=
       candidateReducedShapedType.getElementType())
     return SubViewVerificationResult::ElemTypeMismatch;
 
   // Strided layout logic is relevant for MemRefType only.
   MemRefType original = originalType.cast<MemRefType>();
   MemRefType candidateReduced = candidateReducedType.cast<MemRefType>();
-  if (original.getMemorySpaceAsInt() != candidateReduced.getMemorySpaceAsInt())
+  if (original.getMemorySpace() != candidateReduced.getMemorySpace())
     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.getMemorySpaceAsInt() != subViewType.getMemorySpaceAsInt())
+  if (baseType.getMemorySpace() != subViewType.getMemorySpace())
     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.getMemorySpaceAsInt() != viewType.getMemorySpaceAsInt())
+  if (baseType.getMemorySpace() != viewType.getMemorySpace())
     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.getMemorySpaceAsInt()));
+      MemRefType::get({}, vectorType, {}, memRefType.getMemorySpace()));
 }
 
 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().getMemorySpaceAsInt() !=
-      op.getMemRefType().getMemorySpaceAsInt())
+  if (op.getResultMemRefType().getMemorySpace() !=
+      op.getMemRefType().getMemorySpace())
     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: ")