[mlir][NFC] Promote memory space to BaseMemRefType

[dcaballe]

Summary

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.

[dcaballe]

Ping :)

[bondhugula]

Looks good, overall. One question below.

mlir/lib/IR/TypeDetail.h

214–216

Do we need this? What happens if this is removed?

[dcaballe]

Thanks!

mlir/lib/IR/TypeDetail.h

214–216

I'll try. According to the doc, it shouldn't be needed for basic types.

mlir/include/mlir/IR/StandardTypes.h

 
 namespace detail {
 
 struct IntegerTypeStorage;
 struct ShapedTypeStorage;
 struct VectorTypeStorage;
 struct RankedTensorTypeStorage;
 struct UnrankedTensorTypeStorage;
+struct BaseMemRefTypeStorage;
 struct MemRefTypeStorage;
 struct UnrankedMemRefTypeStorage;
 struct ComplexTypeStorage;
 struct TupleTypeStorage;
 
 } // namespace detail
 
 //===----------------------------------------------------------------------===//

 
 //===----------------------------------------------------------------------===//
 // BaseMemRefType
 //===----------------------------------------------------------------------===//
 
 /// 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.
   static bool isValidElementType(Type type) {
     return type.isIntOrIndexOrFloat() || type.isa<VectorType, ComplexType>();
   }
 
   /// 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;
 };
 
 //===----------------------------------------------------------------------===//
 // 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
static MemRefType getChecked(ArrayRef<int64_t> shape, Type elementType,
                                unsigned memorySpace, Location location);
 
   ArrayRef<int64_t> getShape() const;
 
   /// Returns an array of affine map pointers representing the memref affine
   /// 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.
   static int64_t getDynamicStrideOrOffset() {
     return ShapedType::kDynamicStrideOrOffset;
   }
 
 private:
static UnrankedMemRefType getChecked(Type elementType, unsigned memorySpace,
                                        Location location);
 
   /// Verify the construction of a unranked memref type.
   static LogicalResult verifyConstructionInvariants(Location loc,
                                                     Type elementType,
                                                     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;
 };
 
 //===----------------------------------------------------------------------===//
 // TupleType
 //===----------------------------------------------------------------------===//
 
 /// Tuple types represent a collection of other types. Note: This type merely
 /// provides a common mechanism for representing tuples in MLIR. It is up to

mlir/lib/IR/StandardTypes.cpp

 
 LogicalResult
 UnrankedTensorType::verifyConstructionInvariants(Location loc,
                                                  Type elementType) {
   return checkTensorElementType(loc, elementType);
 }
 
 //===----------------------------------------------------------------------===//
+// BaseMemRefType
+//===----------------------------------------------------------------------===//
+
+unsigned BaseMemRefType::getMemorySpace() const {
+  return static_cast<ImplType *>(impl)->memorySpace;
+}
+
+//===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
 
 /// Get or create a new MemRefType based on shape, element type, affine
 /// map composition, and memory space.  Assumes the arguments define a
 /// well-formed MemRef type.  Use getChecked to gracefully handle MemRefType
 /// construction failures.
 MemRefType MemRefType::get(ArrayRef<int64_t> shape, Type elementType,

 }
 
 ArrayRef<int64_t> MemRefType::getShape() const { return getImpl()->getShape(); }
 
 ArrayRef<AffineMap> MemRefType::getAffineMaps() const {
   return getImpl()->getAffineMaps();
 }
 
-unsigned MemRefType::getMemorySpace() const { return getImpl()->memorySpace; }
-
 //===----------------------------------------------------------------------===//
 // UnrankedMemRefType
 //===----------------------------------------------------------------------===//
 
 UnrankedMemRefType UnrankedMemRefType::get(Type elementType,
                                            unsigned memorySpace) {
   return Base::get(elementType.getContext(), elementType, memorySpace);
 }
 
 UnrankedMemRefType UnrankedMemRefType::getChecked(Type elementType,
                                                   unsigned memorySpace,
                                                   Location location) {
   return Base::getChecked(location, elementType, memorySpace);
 }
 
-unsigned UnrankedMemRefType::getMemorySpace() const {
-  return getImpl()->memorySpace;
-}
-
 LogicalResult
 UnrankedMemRefType::verifyConstructionInvariants(Location loc, Type elementType,
                                                  unsigned memorySpace) {
   if (!BaseMemRefType::isValidElementType(elementType))
     return emitError(loc, "invalid memref element type");
   return success();
 }
 

mlir/lib/IR/TypeDetail.h

struct UnrankedTensorTypeStorage : public ShapedTypeStorage {
   /// Construction.
   static UnrankedTensorTypeStorage *construct(TypeStorageAllocator &allocator,
                                               Type elementTy) {
     return new (allocator.allocate<UnrankedTensorTypeStorage>())
         UnrankedTensorTypeStorage(elementTy);
   }
 };
 
-struct MemRefTypeStorage : public ShapedTypeStorage {
+struct BaseMemRefTypeStorage : public ShapedTypeStorage {
+  BaseMemRefTypeStorage(Type elementType, unsigned memorySpace)
+      : ShapedTypeStorage(elementType), memorySpace(memorySpace) {}
+
+  /// The hash key used for uniquing.
+  using KeyTy = unsigned;
+  bool operator==(const KeyTy &key) const { return key == 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
   // composition, and memory space.
   using KeyTy =
       std::tuple<ArrayRef<int64_t>, Type, ArrayRef<AffineMap>, unsigned>;
   bool operator==(const KeyTy &key) const {
     return key == KeyTy(getShape(), elementType, getAffineMaps(), memorySpace);
struct MemRefTypeStorage : public BaseMemRefTypeStorage {
   /// An array of integers which stores the shape dimension sizes.
   const int64_t *shapeElements;
   /// The number of shape elements.
   unsigned shapeSize;
   /// The number of affine maps in the 'affineMapList' array.
   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>;
   bool operator==(const KeyTy &key) const {
     return key == KeyTy(elementType, memorySpace);
   }
 
   /// Construction.
   static UnrankedMemRefTypeStorage *construct(TypeStorageAllocator &allocator,
                                               const KeyTy &key) {
 
     // Initialize the memory using placement new.
     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.
 struct ComplexTypeStorage : public TypeStorage {
   ComplexTypeStorage(Type elementType) : elementType(elementType) {}
 
   /// The hash key used for uniquing.
   using KeyTy = Type;