[Mlir-commits] [mlir] 18a2f47 - [mlir][NFC] Replace `getMemorySpaceAsInt` with `getMemorySpace` where possible

[Vladislav Vinogradov]

Author: Vladislav Vinogradov
Date: 2021-03-24T13:23:59+03:00
New Revision: 18a2f479bf475c7cb94b742b39290bef98e3d305

URL: https://github.com/llvm/llvm-project/commit/18a2f479bf475c7cb94b742b39290bef98e3d305
DIFF: https://github.com/llvm/llvm-project/commit/18a2f479bf475c7cb94b742b39290bef98e3d305.diff

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

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

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D99154

Added: 
    

Modified: 
    mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
    mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
    mlir/lib/Dialect/Vector/VectorOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
index 1c8e05b2d623..f95193f76508 100644
--- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
+++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
@@ -116,11 +116,9 @@ class NDTransferOpHelper {
         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();

diff  --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index e0e273d85669..546c43a97407 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -420,7 +420,7 @@ bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
         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.
@@ -447,10 +447,8 @@ bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
     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;
 
@@ -1204,7 +1202,7 @@ 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("
diff erent memory spaces specified for source type ")
            << srcType << " and result memref type " << resultType;
   if (srcType.getElementType() != resultType.getElementType())
@@ -1389,7 +1387,7 @@ Type SubViewOp::inferResultType(MemRefType sourceMemRefType,
       staticSizes, sourceMemRefType.getElementType(),
       makeStridedLinearLayoutMap(targetStrides, targetOffset,
                                  sourceMemRefType.getContext()),
-      sourceMemRefType.getMemorySpaceAsInt());
+      sourceMemRefType.getMemorySpace());
 }
 
 Type SubViewOp::inferResultType(MemRefType sourceMemRefType,
@@ -1435,7 +1433,7 @@ Type SubViewOp::inferRankReducedResultType(
       map = getProjectedMap(maps.front(), dimsToProject);
     inferredType =
         MemRefType::get(projectedShape, inferredType.getElementType(), map,
-                        inferredType.getMemorySpaceAsInt());
+                        inferredType.getMemorySpace());
   }
   return inferredType;
 }
@@ -1613,7 +1611,7 @@ isRankReducedType(Type originalType, Type candidateReducedType,
   // 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();
@@ -1687,7 +1685,7 @@ static LogicalResult verify(SubViewOp op) {
   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("
diff erent memory spaces specified for base memref "
                         "type ")
            << baseType << " and subview memref type " << subViewType;
@@ -1979,7 +1977,7 @@ static LogicalResult verify(ViewOp op) {
     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("
diff erent memory spaces specified for base memref "
                         "type ")
            << baseType << " and view memref type " << viewType;

diff  --git a/mlir/lib/Dialect/Vector/VectorOps.cpp b/mlir/lib/Dialect/Vector/VectorOps.cpp
index d1703caccc46..9079f99ade8b 100644
--- a/mlir/lib/Dialect/Vector/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/VectorOps.cpp
@@ -3235,7 +3235,7 @@ void TypeCastOp::build(OpBuilder &builder, OperationState &result,
       VectorType::get(extractShape(memRefType),
                       getElementTypeOrSelf(getElementTypeOrSelf(memRefType)));
   result.addTypes(
-      MemRefType::get({}, vectorType, {}, memRefType.getMemorySpaceAsInt()));
+      MemRefType::get({}, vectorType, {}, memRefType.getMemorySpace()));
 }
 
 static LogicalResult verify(TypeCastOp op) {
@@ -3244,8 +3244,8 @@ static LogicalResult verify(TypeCastOp op) {
     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();