[mlir][Types] Remove the subclass data from Type

[rriddle]

Summary

Subclass data is useful when a certain amount of memory is allocated, but not all of it is used. In the case of Type, that hasn't been the case for a while and the subclass is just taking up a full unsigned. Removing this frees up ~8 bytes for almost every type instance.

Depends On D82597

[rriddle]

Rebase

[mravishankar]

Nice! Thanks for the cleanup!

[rriddle]

Rebase

[mehdi_amini]

Wow that was pretty bad actually, do I miss something or there was zero upside to the subclass data before?

[rriddle]

In D85348#2198569, @mehdi_amini wrote:

Wow that was pretty bad actually, do I miss something or there was zero upside to the subclass data before?

A long long time ago there was an upside, but for at least the past 10 months its just been a waste of space.

[rriddle]

Rebase

[rriddle]

Rebase

mlir/include/mlir/IR/StandardTypes.h

 
 /// Integer types can have arbitrary bitwidth up to a large fixed limit.
 class IntegerType
     : public Type::TypeBase<IntegerType, Type, detail::IntegerTypeStorage> {
 public:
   using Base::Base;
 
   /// Signedness semantics.
-  enum SignednessSemantics {
+  enum SignednessSemantics : uint32_t {
     Signless, /// No signedness semantics
     Signed,   /// Signed integer
     Unsigned, /// Unsigned integer
   };
 
   /// Get or create a new IntegerType of the given width within the context.
   /// The created IntegerType is signless (i.e., no signedness semantics).
   /// Assume the width is within the allowed range and assert on failures. Use

mlir/include/mlir/IR/TypeSupport.h

 
 public:
   /// Return the abstract type descriptor for this type.
   const AbstractType &getAbstractType() {
     assert(abstractType && "Malformed type storage object.");
     return *abstractType;
   }
 
-  /// Get the subclass data.
-  unsigned getSubclassData() const { return subclassData; }
-
-  /// Set the subclass data.
-  void setSubclassData(unsigned val) { subclassData = val; }
-
 protected:
   /// This constructor is used by derived classes as part of the TypeUniquer.
-  TypeStorage(unsigned subclassData = 0)
-      : abstractType(nullptr), subclassData(subclassData) {}
+  TypeStorage() : abstractType(nullptr) {}
 
 private:
   /// Set the abstract type for this storage instance. This is used by the
   /// TypeUniquer when initializing a newly constructed type storage object.
   void initialize(const AbstractType &abstractTy) {
     abstractType = &abstractTy;
   }
 
   /// The abstract description for this type.
   const AbstractType *abstractType;
-
-  /// Space for subclasses to store data.
-  unsigned subclassData;
 };
 
 /// Default storage type for types that require no additional initialization or
 /// storage.
 using DefaultTypeStorage = TypeStorage;
 
 //===----------------------------------------------------------------------===//
 // TypeStorageAllocator

mlir/include/mlir/IR/Types.h

#include "DialectSymbolRegistry.def"
   bool isIntOrIndexOrFloat();
 
   /// Print the current type.
   void print(raw_ostream &os);
   void dump();
 
   friend ::llvm::hash_code hash_value(Type arg);
 
-  unsigned getSubclassData() const;
-  void setSubclassData(unsigned val);
-
   /// Methods for supporting PointerLikeTypeTraits.
   const void *getAsOpaquePointer() const {
     return static_cast<const void *>(impl);
   }
   static Type getFromOpaquePointer(const void *pointer) {
     return Type(reinterpret_cast<ImplType *>(const_cast<void *>(pointer)));
   }
 
class FunctionType
     : public Type::TypeBase<FunctionType, Type, detail::FunctionTypeStorage> {
 public:
   using Base::Base;
 
   static FunctionType get(TypeRange inputs, TypeRange results,
                           MLIRContext *context);
 
   // Input types.
+<<<<<<< dest:   de5c6c6c3887 - riverriddle: [mlir] Add a utility class, Threa...
+  unsigned getNumInputs() const { return getSubclassData(); }
+||||||| base
   unsigned getNumInputs() const { return getSubclassData(); }
+
+=======
+  unsigned getNumInputs() const;
+
+>>>>>>> source: f412125b3eaf - riverriddle: [mlir][Types] Remove the subclass...
   Type getInput(unsigned i) const { return getInputs()[i]; }
   ArrayRef<Type> getInputs() const;
 
   // Result types.
   unsigned getNumResults() const;
   Type getResult(unsigned i) const { return getResults()[i]; }
   ArrayRef<Type> getResults() const;
 

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

struct spirv::detail::ArrayTypeStorage : public TypeStorage {
   using KeyTy = std::tuple<Type, unsigned, unsigned>;
 
   static ArrayTypeStorage *construct(TypeStorageAllocator &allocator,
                                      const KeyTy &key) {
     return new (allocator.allocate<ArrayTypeStorage>()) ArrayTypeStorage(key);
   }
 
   bool operator==(const KeyTy &key) const {
-    return key == KeyTy(elementType, getSubclassData(), stride);
+    return key == KeyTy(elementType, elementCount, stride);
   }
 
   ArrayTypeStorage(const KeyTy &key)
-      : TypeStorage(std::get<1>(key)), elementType(std::get<0>(key)),
+      : elementType(std::get<0>(key)), elementCount(std::get<1>(key)),
         stride(std::get<2>(key)) {}
 
   Type elementType;
+  unsigned elementCount;
   unsigned stride;
 };
 
 ArrayType ArrayType::get(Type elementType, unsigned elementCount) {
   assert(elementCount && "ArrayType needs at least one element");
   return Base::get(elementType.getContext(), TypeKind::Array, elementType,
                    elementCount, /*stride=*/0);
 }
 
 ArrayType ArrayType::get(Type elementType, unsigned elementCount,
                          unsigned stride) {
   assert(elementCount && "ArrayType needs at least one element");
   return Base::get(elementType.getContext(), TypeKind::Array, elementType,
                    elementCount, stride);
 }
 
-unsigned ArrayType::getNumElements() const {
-  return getImpl()->getSubclassData();
-}
+unsigned ArrayType::getNumElements() const { return getImpl()->elementCount; }
 
 Type ArrayType::getElementType() const { return getImpl()->elementType; }
 
 unsigned ArrayType::getArrayStride() const { return getImpl()->stride; }
 
 void ArrayType::getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                               Optional<StorageClass> storage) {
   getElementType().cast<SPIRVType>().getExtensions(extensions, storage);
struct spirv::detail::CooperativeMatrixTypeStorage : public TypeStorage {
 
   static CooperativeMatrixTypeStorage *
   construct(TypeStorageAllocator &allocator, const KeyTy &key) {
     return new (allocator.allocate<CooperativeMatrixTypeStorage>())
         CooperativeMatrixTypeStorage(key);
   }
 
   bool operator==(const KeyTy &key) const {
-    return key == KeyTy(elementType, getScope(), rows, columns);
+    return key == KeyTy(elementType, scope, rows, columns);
   }
 
   CooperativeMatrixTypeStorage(const KeyTy &key)
-      : TypeStorage(static_cast<unsigned>(std::get<1>(key))),
-        elementType(std::get<0>(key)), rows(std::get<2>(key)),
-        columns(std::get<3>(key)) {}
-
-  Scope getScope() const { return static_cast<Scope>(getSubclassData()); }
+      : elementType(std::get<0>(key)), rows(std::get<2>(key)),
+        columns(std::get<3>(key)), scope(std::get<1>(key)) {}
 
   Type elementType;
   unsigned rows;
   unsigned columns;
+  Scope scope;
 };
 
 CooperativeMatrixNVType CooperativeMatrixNVType::get(Type elementType,
                                                      Scope scope, unsigned rows,
                                                      unsigned columns) {
   return Base::get(elementType.getContext(), TypeKind::CooperativeMatrix,
                    elementType, scope, rows, columns);
 }
 
 Type CooperativeMatrixNVType::getElementType() const {
   return getImpl()->elementType;
 }
 
-Scope CooperativeMatrixNVType::getScope() const {
-  return getImpl()->getScope();
-}
+Scope CooperativeMatrixNVType::getScope() const { return getImpl()->scope; }
 
 unsigned CooperativeMatrixNVType::getRows() const { return getImpl()->rows; }
 
 unsigned CooperativeMatrixNVType::getColumns() const {
   return getImpl()->columns;
 }
 
 void CooperativeMatrixNVType::getExtensions(

 }
 template <> constexpr unsigned getNumBits<ImageFormat>() {
   static_assert((1 << 6) > getMaxEnumValForImageFormat(),
                 "Not enough bits to encode ImageFormat value");
   return 6;
 }
 
 struct spirv::detail::ImageTypeStorage : public TypeStorage {
-private:
-  /// Define a bit-field struct to pack the enum values
-  union EnumPack {
-    struct {
-      unsigned dimEncoding : getNumBits<Dim>();
-      unsigned depthInfoEncoding : getNumBits<ImageDepthInfo>();
-      unsigned arrayedInfoEncoding : getNumBits<ImageArrayedInfo>();
-      unsigned samplingInfoEncoding : getNumBits<ImageSamplingInfo>();
-      unsigned samplerUseInfoEncoding : getNumBits<ImageSamplerUseInfo>();
-      unsigned formatEncoding : getNumBits<ImageFormat>();
-    } data;
-    unsigned storage;
-  };
-
 public:
   using KeyTy = std::tuple<Type, Dim, ImageDepthInfo, ImageArrayedInfo,
                            ImageSamplingInfo, ImageSamplerUseInfo, ImageFormat>;
 
   static ImageTypeStorage *construct(TypeStorageAllocator &allocator,
                                      const KeyTy &key) {
     return new (allocator.allocate<ImageTypeStorage>()) ImageTypeStorage(key);
   }
 
   bool operator==(const KeyTy &key) const {
-    return key == KeyTy(elementType, getDim(), getDepthInfo(), getArrayedInfo(),
-                        getSamplingInfo(), getSamplerUseInfo(),
-                        getImageFormat());
-  }
-
-  Dim getDim() const {
-    EnumPack v;
-    v.storage = getSubclassData();
-    return static_cast<Dim>(v.data.dimEncoding);
-  }
-  void setDim(Dim dim) {
-    EnumPack v;
-    v.storage = getSubclassData();
-    v.data.dimEncoding = static_cast<unsigned>(dim);
-    setSubclassData(v.storage);
-  }
-
-  ImageDepthInfo getDepthInfo() const {
-    EnumPack v;
-    v.storage = getSubclassData();
-    return static_cast<ImageDepthInfo>(v.data.depthInfoEncoding);
-  }
-  void setDepthInfo(ImageDepthInfo depthInfo) {
-    EnumPack v;
-    v.storage = getSubclassData();
-    v.data.depthInfoEncoding = static_cast<unsigned>(depthInfo);
-    setSubclassData(v.storage);
-  }
-
-  ImageArrayedInfo getArrayedInfo() const {
-    EnumPack v;
-    v.storage = getSubclassData();
-    return static_cast<ImageArrayedInfo>(v.data.arrayedInfoEncoding);
-  }
-  void setArrayedInfo(ImageArrayedInfo arrayedInfo) {
-    EnumPack v;
-    v.storage = getSubclassData();
-    v.data.arrayedInfoEncoding = static_cast<unsigned>(arrayedInfo);
-    setSubclassData(v.storage);
-  }
-
-  ImageSamplingInfo getSamplingInfo() const {
-    EnumPack v;
-    v.storage = getSubclassData();
-    return static_cast<ImageSamplingInfo>(v.data.samplingInfoEncoding);
-  }
-  void setSamplingInfo(ImageSamplingInfo samplingInfo) {
-    EnumPack v;
-    v.storage = getSubclassData();
-    v.data.samplingInfoEncoding = static_cast<unsigned>(samplingInfo);
-    setSubclassData(v.storage);
-  }
-
-  ImageSamplerUseInfo getSamplerUseInfo() const {
-    EnumPack v;
-    v.storage = getSubclassData();
-    return static_cast<ImageSamplerUseInfo>(v.data.samplerUseInfoEncoding);
-  }
-  void setSamplerUseInfo(ImageSamplerUseInfo samplerUseInfo) {
-    EnumPack v;
-    v.storage = getSubclassData();
-    v.data.samplerUseInfoEncoding = static_cast<unsigned>(samplerUseInfo);
-    setSubclassData(v.storage);
-  }
-
-  ImageFormat getImageFormat() const {
-    EnumPack v;
-    v.storage = getSubclassData();
-    return static_cast<ImageFormat>(v.data.formatEncoding);
-  }
-  void setImageFormat(ImageFormat format) {
-    EnumPack v;
-    v.storage = getSubclassData();
-    v.data.formatEncoding = static_cast<unsigned>(format);
-    setSubclassData(v.storage);
-  }
-
-  ImageTypeStorage(const KeyTy &key) : elementType(std::get<0>(key)) {
-    static_assert(sizeof(EnumPack) <= sizeof(getSubclassData()),
-                  "EnumPack size greater than subClassData type size");
-    setDim(std::get<1>(key));
-    setDepthInfo(std::get<2>(key));
-    setArrayedInfo(std::get<3>(key));
-    setSamplingInfo(std::get<4>(key));
-    setSamplerUseInfo(std::get<5>(key));
-    setImageFormat(std::get<6>(key));
+    return key == KeyTy(elementType, dim, depthInfo, arrayedInfo, samplingInfo,
+                        samplerUseInfo, format);
   }
 
+  ImageTypeStorage(const KeyTy &key)
+      : elementType(std::get<0>(key)), dim(std::get<1>(key)),
+        depthInfo(std::get<2>(key)), arrayedInfo(std::get<3>(key)),
+        samplingInfo(std::get<4>(key)), samplerUseInfo(std::get<5>(key)),
+        format(std::get<6>(key)) {}
+
   Type elementType;
+  Dim dim : getNumBits<Dim>();
+  ImageDepthInfo depthInfo : getNumBits<ImageDepthInfo>();
+  ImageArrayedInfo arrayedInfo : getNumBits<ImageArrayedInfo>();
+  ImageSamplingInfo samplingInfo : getNumBits<ImageSamplingInfo>();
+  ImageSamplerUseInfo samplerUseInfo : getNumBits<ImageSamplerUseInfo>();
+  ImageFormat format : getNumBits<ImageFormat>();
 };
 
 ImageType
 ImageType::get(std::tuple<Type, Dim, ImageDepthInfo, ImageArrayedInfo,
                           ImageSamplingInfo, ImageSamplerUseInfo, ImageFormat>
                    value) {
   return Base::get(std::get<0>(value).getContext(), TypeKind::Image, value);
 }
 
 Type ImageType::getElementType() const { return getImpl()->elementType; }
 
-Dim ImageType::getDim() const { return getImpl()->getDim(); }
+Dim ImageType::getDim() const { return getImpl()->dim; }
 
-ImageDepthInfo ImageType::getDepthInfo() const {
-  return getImpl()->getDepthInfo();
-}
+ImageDepthInfo ImageType::getDepthInfo() const { return getImpl()->depthInfo; }
 
 ImageArrayedInfo ImageType::getArrayedInfo() const {
-  return getImpl()->getArrayedInfo();
+  return getImpl()->arrayedInfo;
 }
 
 ImageSamplingInfo ImageType::getSamplingInfo() const {
-  return getImpl()->getSamplingInfo();
+  return getImpl()->samplingInfo;
 }
 
 ImageSamplerUseInfo ImageType::getSamplerUseInfo() const {
-  return getImpl()->getSamplerUseInfo();
+  return getImpl()->samplerUseInfo;
 }
 
-ImageFormat ImageType::getImageFormat() const {
-  return getImpl()->getImageFormat();
-}
+ImageFormat ImageType::getImageFormat() const { return getImpl()->format; }
 
 void ImageType::getExtensions(SPIRVType::ExtensionArrayRefVector &,
                               Optional<StorageClass>) {
   // Image types do not require extra extensions thus far.
 }
 
 void ImageType::getCapabilities(
     SPIRVType::CapabilityArrayRefVector &capabilities, Optional<StorageClass>) {
struct spirv::detail::PointerTypeStorage : public TypeStorage {
 
   static PointerTypeStorage *construct(TypeStorageAllocator &allocator,
                                        const KeyTy &key) {
     return new (allocator.allocate<PointerTypeStorage>())
         PointerTypeStorage(key);
   }
 
   bool operator==(const KeyTy &key) const {
-    return key == KeyTy(pointeeType, getStorageClass());
+    return key == KeyTy(pointeeType, storageClass);
   }
 
   PointerTypeStorage(const KeyTy &key)
-      : TypeStorage(static_cast<unsigned>(key.second)), pointeeType(key.first) {
-  }
-
-  StorageClass getStorageClass() const {
-    return static_cast<StorageClass>(getSubclassData());
-  }
+      : pointeeType(key.first), storageClass(key.second) {}
 
   Type pointeeType;
+  StorageClass storageClass;
 };
 
 PointerType PointerType::get(Type pointeeType, StorageClass storageClass) {
   return Base::get(pointeeType.getContext(), TypeKind::Pointer, pointeeType,
                    storageClass);
 }
 
 Type PointerType::getPointeeType() const { return getImpl()->pointeeType; }
 
 StorageClass PointerType::getStorageClass() const {
-  return getImpl()->getStorageClass();
+  return getImpl()->storageClass;
 }
 
 void PointerType::getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                                 Optional<StorageClass> storage) {
   // Use this pointer type's storage class because this pointer indicates we are
   // using the pointee type in that specific storage class.
   getPointeeType().cast<SPIRVType>().getExtensions(extensions,
                                                    getStorageClass());
struct spirv::detail::RuntimeArrayTypeStorage : public TypeStorage {
 
   static RuntimeArrayTypeStorage *construct(TypeStorageAllocator &allocator,
                                             const KeyTy &key) {
     return new (allocator.allocate<RuntimeArrayTypeStorage>())
         RuntimeArrayTypeStorage(key);
   }
 
   bool operator==(const KeyTy &key) const {
-    return key == KeyTy(elementType, getSubclassData());
+    return key == KeyTy(elementType, stride);
   }
 
   RuntimeArrayTypeStorage(const KeyTy &key)
-      : TypeStorage(key.second), elementType(key.first) {}
+      : elementType(key.first), stride(key.second) {}
 
   Type elementType;
+  unsigned stride;
 };
 
 RuntimeArrayType RuntimeArrayType::get(Type elementType) {
   return Base::get(elementType.getContext(), TypeKind::RuntimeArray,
                    elementType, /*stride=*/0);
 }
 
 RuntimeArrayType RuntimeArrayType::get(Type elementType, unsigned stride) {
   return Base::get(elementType.getContext(), TypeKind::RuntimeArray,
                    elementType, stride);
 }
 
 Type RuntimeArrayType::getElementType() const { return getImpl()->elementType; }
 
-unsigned RuntimeArrayType::getArrayStride() const {
-  return getImpl()->getSubclassData();
-}
+unsigned RuntimeArrayType::getArrayStride() const { return getImpl()->stride; }
 
 void RuntimeArrayType::getExtensions(
     SPIRVType::ExtensionArrayRefVector &extensions,
     Optional<StorageClass> storage) {
   getElementType().cast<SPIRVType>().getExtensions(extensions, storage);
 }
 
 void RuntimeArrayType::getCapabilities(

 // StructType
 //===----------------------------------------------------------------------===//
 
 struct spirv::detail::StructTypeStorage : public TypeStorage {
   StructTypeStorage(
       unsigned numMembers, Type const *memberTypes,
       StructType::OffsetInfo const *layoutInfo, unsigned numMemberDecorations,
       StructType::MemberDecorationInfo const *memberDecorationsInfo)
-      : TypeStorage(numMembers), memberTypes(memberTypes),
-        offsetInfo(layoutInfo), numMemberDecorations(numMemberDecorations),
+      : memberTypes(memberTypes), offsetInfo(layoutInfo),
+        numMembers(numMembers), numMemberDecorations(numMemberDecorations),
         memberDecorationsInfo(memberDecorationsInfo) {}
 
   using KeyTy = std::tuple<ArrayRef<Type>, ArrayRef<StructType::OffsetInfo>,
                            ArrayRef<StructType::MemberDecorationInfo>>;
   bool operator==(const KeyTy &key) const {
     return key ==
            KeyTy(getMemberTypes(), getOffsetInfo(), getMemberDecorationsInfo());
   }
if (!std::get<2>(key).empty()) {
       memberDecorationList = allocator.copyInto(keyMemberDecorations).data();
     }
     return new (allocator.allocate<StructTypeStorage>())
         StructTypeStorage(keyTypes.size(), typesList, offsetInfoList,
                           numMemberDecorations, memberDecorationList);
   }
 
   ArrayRef<Type> getMemberTypes() const {
-    return ArrayRef<Type>(memberTypes, getSubclassData());
+    return ArrayRef<Type>(memberTypes, numMembers);
   }
 
   ArrayRef<StructType::OffsetInfo> getOffsetInfo() const {
     if (offsetInfo) {
-      return ArrayRef<StructType::OffsetInfo>(offsetInfo, getSubclassData());
+      return ArrayRef<StructType::OffsetInfo>(offsetInfo, numMembers);
     }
     return {};
   }
 
   ArrayRef<StructType::MemberDecorationInfo> getMemberDecorationsInfo() const {
     if (memberDecorationsInfo) {
       return ArrayRef<StructType::MemberDecorationInfo>(memberDecorationsInfo,
                                                         numMemberDecorations);
     }
     return {};
   }
 
   Type const *memberTypes;
   StructType::OffsetInfo const *offsetInfo;
+  unsigned numMembers;
   unsigned numMemberDecorations;
   StructType::MemberDecorationInfo const *memberDecorationsInfo;
 };
 
 StructType
 StructType::get(ArrayRef<Type> memberTypes,
                 ArrayRef<StructType::OffsetInfo> offsetInfo,
                 ArrayRef<StructType::MemberDecorationInfo> memberDecorations) {
   assert(!memberTypes.empty() && "Struct needs at least one member type");
   // Sort the decorations.
   SmallVector<StructType::MemberDecorationInfo, 4> sortedDecorations(
       memberDecorations.begin(), memberDecorations.end());
   llvm::array_pod_sort(sortedDecorations.begin(), sortedDecorations.end());
   return Base::get(memberTypes.vec().front().getContext(), TypeKind::Struct,
                    memberTypes, offsetInfo, sortedDecorations);
 }
 
 StructType StructType::getEmpty(MLIRContext *context) {
   return Base::get(context, TypeKind::Struct, ArrayRef<Type>(),
                    ArrayRef<StructType::OffsetInfo>(),
                    ArrayRef<StructType::MemberDecorationInfo>());
 }
 
-unsigned StructType::getNumElements() const {
-  return getImpl()->getSubclassData();
-}
+unsigned StructType::getNumElements() const { return getImpl()->numMembers; }
 
 Type StructType::getElementType(unsigned index) const {
   assert(getNumElements() > index && "member index out of range");
   return getImpl()->memberTypes[index];
 }
 
 StructType::ElementTypeRange StructType::getElementTypes() const {
   return ElementTypeRange(getImpl()->memberTypes, getNumElements());

mlir/lib/IR/StandardTypes.cpp

IntegerType::verifyConstructionInvariants(Location loc, unsigned width,
                                           SignednessSemantics signedness) {
   if (width > IntegerType::kMaxWidth) {
     return emitError(loc) << "integer bitwidth is limited to "
                           << IntegerType::kMaxWidth << " bits";
   }
   return success();
 }
 
-unsigned IntegerType::getWidth() const { return getImpl()->getWidth(); }
+unsigned IntegerType::getWidth() const { return getImpl()->width; }
 
 IntegerType::SignednessSemantics IntegerType::getSignedness() const {
-  return getImpl()->getSignedness();
+  return getImpl()->signedness;
 }
 
 //===----------------------------------------------------------------------===//
 // Float Type
 //===----------------------------------------------------------------------===//
 
 unsigned FloatType::getWidth() {
   switch (getKind()) {

mlir/lib/IR/TypeDetail.h

struct OpaqueTypeStorage : public TypeStorage {
   // The parser type data for this opaque type.
   StringRef typeData;
 };
 
 /// Integer Type Storage and Uniquing.
 struct IntegerTypeStorage : public TypeStorage {
   IntegerTypeStorage(unsigned width,
                      IntegerType::SignednessSemantics signedness)
-      : TypeStorage(packKeyBits(width, signedness)) {}
+      : width(width), signedness(signedness) {}
 
   /// The hash key used for uniquing.
   using KeyTy = std::pair<unsigned, IntegerType::SignednessSemantics>;
 
   static llvm::hash_code hashKey(const KeyTy &key) {
-    return llvm::hash_value(packKeyBits(key.first, key.second));
+    return llvm::hash_value(key);
   }
 
   bool operator==(const KeyTy &key) const {
-    return getSubclassData() == packKeyBits(key.first, key.second);
+    return KeyTy(width, signedness) == key;
   }
 
   static IntegerTypeStorage *construct(TypeStorageAllocator &allocator,
                                        KeyTy key) {
     return new (allocator.allocate<IntegerTypeStorage>())
         IntegerTypeStorage(key.first, key.second);
   }
 
-  struct KeyBits {
-    unsigned width : 30;
-    unsigned signedness : 2;
-  };
-
-  /// Pack the given `width` and `signedness` as a key.
-  static unsigned packKeyBits(unsigned width,
-                              IntegerType::SignednessSemantics signedness) {
-    KeyBits bits{width, static_cast<unsigned>(signedness)};
-    return llvm::bit_cast<unsigned>(bits);
-  }
-
-  static KeyBits unpackKeyBits(unsigned bits) {
-    return llvm::bit_cast<KeyBits>(bits);
-  }
-
-  unsigned getWidth() { return unpackKeyBits(getSubclassData()).width; }
-
-  IntegerType::SignednessSemantics getSignedness() {
-    return static_cast<IntegerType::SignednessSemantics>(
-        unpackKeyBits(getSubclassData()).signedness);
-  }
+  IntegerType::SignednessSemantics signedness : 2;
 };
 
 /// Function Type Storage and Uniquing.
 struct FunctionTypeStorage : public TypeStorage {
   FunctionTypeStorage(unsigned numInputs, unsigned numResults,
                       Type const *inputsAndResults)
-      : TypeStorage(numInputs), numResults(numResults),
+      : numInputs(numInputs), numResults(numResults),
         inputsAndResults(inputsAndResults) {}
 
   /// The hash key used for uniquing.
   using KeyTy = std::pair<TypeRange, TypeRange>;
   bool operator==(const KeyTy &key) const {
     return key == KeyTy(getInputs(), getResults());
   }
 
static FunctionTypeStorage *construct(TypeStorageAllocator &allocator,
     auto typesList = allocator.copyInto(ArrayRef<Type>(types));
 
     // Initialize the memory using placement new.
     return new (allocator.allocate<FunctionTypeStorage>())
         FunctionTypeStorage(inputs.size(), results.size(), typesList.data());
   }
 
   ArrayRef<Type> getInputs() const {
-    return ArrayRef<Type>(inputsAndResults, getSubclassData());
+    return ArrayRef<Type>(inputsAndResults, numInputs);
   }
   ArrayRef<Type> getResults() const {
-    return ArrayRef<Type>(inputsAndResults + getSubclassData(), numResults);
+    return ArrayRef<Type>(inputsAndResults + numInputs, numResults);
   }
 
+  unsigned numInputs;
   unsigned numResults;
   Type const *inputsAndResults;
 };
 
 /// Shaped Type Storage.
 struct ShapedTypeStorage : public TypeStorage {
-  ShapedTypeStorage(Type elementTy, unsigned subclassData = 0)
-      : TypeStorage(subclassData), elementType(elementTy) {}
+  ShapedTypeStorage(Type elementTy) : elementType(elementTy) {}
 
   /// The hash key used for uniquing.
   using KeyTy = Type;
   bool operator==(const KeyTy &key) const { return key == elementType; }
 
   Type elementType;
 };
 
 /// Vector Type Storage and Uniquing.
 struct VectorTypeStorage : public ShapedTypeStorage {
   VectorTypeStorage(unsigned shapeSize, Type elementTy,
                     const int64_t *shapeElements)
-      : ShapedTypeStorage(elementTy, shapeSize), shapeElements(shapeElements) {}
+      : ShapedTypeStorage(elementTy), shapeElements(shapeElements),
+        shapeSize(shapeSize) {}
 
   /// The hash key used for uniquing.
   using KeyTy = std::pair<ArrayRef<int64_t>, Type>;
   bool operator==(const KeyTy &key) const {
     return key == KeyTy(getShape(), elementType);
   }
 
   /// Construction.
   static VectorTypeStorage *construct(TypeStorageAllocator &allocator,
                                       const KeyTy &key) {
     // Copy the shape into the bump pointer.
     ArrayRef<int64_t> shape = allocator.copyInto(key.first);
 
     // Initialize the memory using placement new.
     return new (allocator.allocate<VectorTypeStorage>())
         VectorTypeStorage(shape.size(), key.second, shape.data());
   }
 
   ArrayRef<int64_t> getShape() const {
-    return ArrayRef<int64_t>(shapeElements, getSubclassData());
+    return ArrayRef<int64_t>(shapeElements, shapeSize);
   }
 
   const int64_t *shapeElements;
+  unsigned shapeSize;
 };
 
 struct RankedTensorTypeStorage : public ShapedTypeStorage {
   RankedTensorTypeStorage(unsigned shapeSize, Type elementTy,
                           const int64_t *shapeElements)
-      : ShapedTypeStorage(elementTy, shapeSize), shapeElements(shapeElements) {}
+      : ShapedTypeStorage(elementTy), shapeElements(shapeElements),
+        shapeSize(shapeSize) {}
 
   /// The hash key used for uniquing.
   using KeyTy = std::pair<ArrayRef<int64_t>, Type>;
   bool operator==(const KeyTy &key) const {
     return key == KeyTy(getShape(), elementType);
   }
 
   /// Construction.
   static RankedTensorTypeStorage *construct(TypeStorageAllocator &allocator,
                                             const KeyTy &key) {
     // Copy the shape into the bump pointer.
     ArrayRef<int64_t> shape = allocator.copyInto(key.first);
 
     // Initialize the memory using placement new.
     return new (allocator.allocate<RankedTensorTypeStorage>())
         RankedTensorTypeStorage(shape.size(), key.second, shape.data());
   }
 
   ArrayRef<int64_t> getShape() const {
-    return ArrayRef<int64_t>(shapeElements, getSubclassData());
+    return ArrayRef<int64_t>(shapeElements, shapeSize);
   }
 
   const int64_t *shapeElements;
+  unsigned shapeSize;
 };
 
 struct UnrankedTensorTypeStorage : public ShapedTypeStorage {
   using ShapedTypeStorage::KeyTy;
   using ShapedTypeStorage::ShapedTypeStorage;
 
   /// Construction.
   static UnrankedTensorTypeStorage *construct(TypeStorageAllocator &allocator,
                                               Type elementTy) {
     return new (allocator.allocate<UnrankedTensorTypeStorage>())
         UnrankedTensorTypeStorage(elementTy);
   }
 };
 
 struct MemRefTypeStorage : public ShapedTypeStorage {
   MemRefTypeStorage(unsigned shapeSize, Type elementType,
                     const int64_t *shapeElements, const unsigned numAffineMaps,
                     AffineMap const *affineMapList, const unsigned memorySpace)
-      : ShapedTypeStorage(elementType, shapeSize), shapeElements(shapeElements),
-        numAffineMaps(numAffineMaps), affineMapList(affineMapList),
-        memorySpace(memorySpace) {}
+      : ShapedTypeStorage(elementType), shapeElements(shapeElements),
+        shapeSize(shapeSize), numAffineMaps(numAffineMaps),
+        affineMapList(affineMapList), memorySpace(memorySpace) {}
 
   /// 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);
static MemRefTypeStorage *construct(TypeStorageAllocator &allocator,
     // Initialize the memory using placement new.
     return new (allocator.allocate<MemRefTypeStorage>())
         MemRefTypeStorage(shape.size(), std::get<1>(key), shape.data(),
                           affineMapComposition.size(),
                           affineMapComposition.data(), std::get<3>(key));
   }
 
   ArrayRef<int64_t> getShape() const {
-    return ArrayRef<int64_t>(shapeElements, getSubclassData());
+    return ArrayRef<int64_t>(shapeElements, shapeSize);
   }
 
   ArrayRef<AffineMap> getAffineMaps() const {
     return ArrayRef<AffineMap>(affineMapList, numAffineMaps);
   }
 
   /// 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;
 };
 

 };
 
 /// A type representing a collection of other types.
 struct TupleTypeStorage final
     : public TypeStorage,
       public llvm::TrailingObjects<TupleTypeStorage, Type> {
   using KeyTy = TypeRange;
 
-  TupleTypeStorage(unsigned numTypes) : TypeStorage(numTypes) {}
+  TupleTypeStorage(unsigned numTypes) : numElements(numTypes) {}
 
   /// Construction.
   static TupleTypeStorage *construct(TypeStorageAllocator &allocator,
                                      TypeRange key) {
     // Allocate a new storage instance.
     auto byteSize = TupleTypeStorage::totalSizeToAlloc<Type>(key.size());
     auto rawMem = allocator.allocate(byteSize, alignof(TupleTypeStorage));
     auto result = ::new (rawMem) TupleTypeStorage(key.size());
 
     // Copy in the element types into the trailing storage.
     std::uninitialized_copy(key.begin(), key.end(),
                             result->getTrailingObjects<Type>());
     return result;
   }
 
   bool operator==(const KeyTy &key) const { return key == getTypes(); }
 
   /// Return the number of held types.
-  unsigned size() const { return getSubclassData(); }
+  unsigned size() const { return numElements; }
 
   /// Return the held types.
   ArrayRef<Type> getTypes() const {
     return {getTrailingObjects<Type>(), size()};
   }
+
+  /// The number of tuple elements.
+  unsigned numElements;
 };
 
 } // namespace detail
 } // namespace mlir
 
 #endif // TYPEDETAIL_H_

mlir/lib/IR/Types.cpp

 unsigned Type::getKind() const { return impl->getKind(); }
 
 Dialect &Type::getDialect() const {
   return impl->getAbstractType().getDialect();
 }
 
 MLIRContext *Type::getContext() const { return getDialect().getContext(); }
 
-unsigned Type::getSubclassData() const { return impl->getSubclassData(); }
-void Type::setSubclassData(unsigned val) { impl->setSubclassData(val); }
-
 //===----------------------------------------------------------------------===//
 // FunctionType
 //===----------------------------------------------------------------------===//
 
 FunctionType FunctionType::get(TypeRange inputs, TypeRange results,
                                MLIRContext *context) {
   return Base::get(context, Type::Kind::Function, inputs, results);
 }
 
+unsigned FunctionType::getNumInputs() const { return getImpl()->numInputs; }
+
 ArrayRef<Type> FunctionType::getInputs() const {
   return getImpl()->getInputs();
 }
 
 unsigned FunctionType::getNumResults() const { return getImpl()->numResults; }
 
 ArrayRef<Type> FunctionType::getResults() const {
   return getImpl()->getResults();