[Mlir-commits] [mlir] 0a925a8 - [mlir][NFC] Promote memory space to BaseMemRefType

[Diego Caballero]

Author: Diego Caballero
Date: 2020-09-24T13:54:06-07:00
New Revision: 0a925a813a5003049bf12330111d61cb823b8814

URL: https://github.com/llvm/llvm-project/commit/0a925a813a5003049bf12330111d61cb823b8814
DIFF: https://github.com/llvm/llvm-project/commit/0a925a813a5003049bf12330111d61cb823b8814.diff

LOG: [mlir][NFC] Promote memory space to BaseMemRefType

This patch moves the memory space field from MemRefType and UnrankedMemRefType
to their base class BaseMemRefType so that it can be retrieved from it without
downcasting it to the specific memref.

Reviewed By: silvas

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/StandardTypes.h
    mlir/lib/IR/StandardTypes.cpp
    mlir/lib/IR/TypeDetail.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/StandardTypes.h b/mlir/include/mlir/IR/StandardTypes.h
index e309595415d1..dc050f135390 100644
--- a/mlir/include/mlir/IR/StandardTypes.h
+++ b/mlir/include/mlir/IR/StandardTypes.h
@@ -31,6 +31,7 @@ struct ShapedTypeStorage;
 struct VectorTypeStorage;
 struct RankedTensorTypeStorage;
 struct UnrankedTensorTypeStorage;
+struct BaseMemRefTypeStorage;
 struct MemRefTypeStorage;
 struct UnrankedMemRefTypeStorage;
 struct ComplexTypeStorage;
@@ -451,6 +452,7 @@ class UnrankedTensorType
 /// Base MemRef for Ranked and Unranked variants
 class BaseMemRefType : public ShapedType {
 public:
+  using ImplType = detail::BaseMemRefTypeStorage;
   using ShapedType::ShapedType;
 
   /// Return true if the specified element type is ok in a memref.
@@ -460,6 +462,9 @@ class BaseMemRefType : public ShapedType {
 
   /// 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;
 };
 
 //===----------------------------------------------------------------------===//
@@ -544,9 +549,6 @@ class MemRefType : public Type::TypeBase<MemRefType, BaseMemRefType,
   /// map composition.
   ArrayRef<AffineMap> getAffineMaps() const;
 
-  /// Returns the memory space in which data referred to by this memref resides.
-  unsigned getMemorySpace() 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.
@@ -592,9 +594,6 @@ class UnrankedMemRefType
                                                     unsigned memorySpace);
 
   ArrayRef<int64_t> getShape() const { return llvm::None; }
-
-  /// Returns the memory space in which data referred to by this memref resides.
-  unsigned getMemorySpace() const;
 };
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/IR/StandardTypes.cpp b/mlir/lib/IR/StandardTypes.cpp
index 8eb9025b7735..ea4fb74ad487 100644
--- a/mlir/lib/IR/StandardTypes.cpp
+++ b/mlir/lib/IR/StandardTypes.cpp
@@ -359,6 +359,14 @@ UnrankedTensorType::verifyConstructionInvariants(Location loc,
   return checkTensorElementType(loc, elementType);
 }
 
+//===----------------------------------------------------------------------===//
+// BaseMemRefType
+//===----------------------------------------------------------------------===//
+
+unsigned BaseMemRefType::getMemorySpace() const {
+  return static_cast<ImplType *>(impl)->memorySpace;
+}
+
 //===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
@@ -449,8 +457,6 @@ ArrayRef<AffineMap> MemRefType::getAffineMaps() const {
   return getImpl()->getAffineMaps();
 }
 
-unsigned MemRefType::getMemorySpace() const { return getImpl()->memorySpace; }
-
 //===----------------------------------------------------------------------===//
 // UnrankedMemRefType
 //===----------------------------------------------------------------------===//
@@ -466,10 +472,6 @@ UnrankedMemRefType UnrankedMemRefType::getChecked(Type elementType,
   return Base::getChecked(location, elementType, memorySpace);
 }
 
-unsigned UnrankedMemRefType::getMemorySpace() const {
-  return getImpl()->memorySpace;
-}
-
 LogicalResult
 UnrankedMemRefType::verifyConstructionInvariants(Location loc, Type elementType,
                                                  unsigned memorySpace) {

diff  --git a/mlir/lib/IR/TypeDetail.h b/mlir/lib/IR/TypeDetail.h
index b5019f00d222..954d0745239a 100644
--- a/mlir/lib/IR/TypeDetail.h
+++ b/mlir/lib/IR/TypeDetail.h
@@ -207,13 +207,21 @@ struct UnrankedTensorTypeStorage : public ShapedTypeStorage {
   }
 };
 
-struct MemRefTypeStorage : public ShapedTypeStorage {
+struct BaseMemRefTypeStorage : public ShapedTypeStorage {
+  BaseMemRefTypeStorage(Type elementType, unsigned memorySpace)
+      : ShapedTypeStorage(elementType), memorySpace(memorySpace) {}
+
+  /// Memory space in which data referenced by memref resides.
+  const unsigned memorySpace;
+};
+
+struct MemRefTypeStorage : public BaseMemRefTypeStorage {
   MemRefTypeStorage(unsigned shapeSize, Type elementType,
                     const int64_t *shapeElements, const unsigned numAffineMaps,
                     AffineMap const *affineMapList, const unsigned memorySpace)
-      : ShapedTypeStorage(elementType), shapeElements(shapeElements),
-        shapeSize(shapeSize), numAffineMaps(numAffineMaps),
-        affineMapList(affineMapList), memorySpace(memorySpace) {}
+      : BaseMemRefTypeStorage(elementType, memorySpace),
+        shapeElements(shapeElements), shapeSize(shapeSize),
+        numAffineMaps(numAffineMaps), affineMapList(affineMapList) {}
 
   /// The hash key used for uniquing.
   // MemRefs are uniqued based on their shape, element type, affine map
@@ -257,16 +265,14 @@ struct MemRefTypeStorage : public ShapedTypeStorage {
   const unsigned numAffineMaps;
   /// List of affine maps in the memref's layout/index map composition.
   AffineMap const *affineMapList;
-  /// Memory space in which data referenced by memref resides.
-  const unsigned memorySpace;
 };
 
 /// Unranked MemRef is a MemRef with unknown rank.
 /// Only element type and memory space are known
-struct UnrankedMemRefTypeStorage : public ShapedTypeStorage {
+struct UnrankedMemRefTypeStorage : public BaseMemRefTypeStorage {
 
   UnrankedMemRefTypeStorage(Type elementTy, const unsigned memorySpace)
-      : ShapedTypeStorage(elementTy), memorySpace(memorySpace) {}
+      : BaseMemRefTypeStorage(elementTy, memorySpace) {}
 
   /// The hash key used for uniquing.
   using KeyTy = std::tuple<Type, unsigned>;
@@ -282,8 +288,6 @@ struct UnrankedMemRefTypeStorage : public ShapedTypeStorage {
     return new (allocator.allocate<UnrankedMemRefTypeStorage>())
         UnrankedMemRefTypeStorage(std::get<0>(key), std::get<1>(key));
   }
-  /// Memory space in which data referenced by memref resides.
-  const unsigned memorySpace;
 };
 
 /// Complex Type Storage.