[mlir][StorageUniquer] Properly call the destructor on non-trivially destructible storage instances

[rriddle]

Summary

This allows for storage instances to store data that isn't uniqued in the context, or contain otherwise non-trivial logic, in the rare situations that they occur. Storage instances with trivial destructors will still have their destructor skipped. A consequence of this is that the storage instance definition must be visible from the place that registers the type.

[ftynse]

A consequence of this is that the storage instance definition must be visible from the place that registers the type.

This is worth putting in the documentation.

I wonder if there is an easy way to test a destructor actually being called. An idea that comes to bind is to report_fatal_error in it and have a death unit test, but that sounds hacky and expensive.

[rriddle]

Add test

[ftynse]

Thanks!

mlir/examples/toy/Ch7/mlir/Dialect.cpp

struct ToyInlinerInterface : public DialectInlinerInterface {
   Operation *materializeCallConversion(OpBuilder &builder, Value input,
                                        Type resultType,
                                        Location conversionLoc) const final {
     return builder.create<CastOp>(conversionLoc, resultType, input);
   }
 };
 
 //===----------------------------------------------------------------------===//
-// ToyDialect
-//===----------------------------------------------------------------------===//
-
-/// Dialect creation, the instance will be owned by the context. This is the
-/// point of registration of custom types and operations for the dialect.
-ToyDialect::ToyDialect(mlir::MLIRContext *ctx)
-    : mlir::Dialect(getDialectNamespace(), ctx, TypeID::get<ToyDialect>()) {
-  addOperations<
-#define GET_OP_LIST
-#include "toy/Ops.cpp.inc"
-      >();
-  addInterfaces<ToyInlinerInterface>();
-  addTypes<StructType>();
-}
-
-mlir::Operation *ToyDialect::materializeConstant(mlir::OpBuilder &builder,
-                                                 mlir::Attribute value,
-                                                 mlir::Type type,
-                                                 mlir::Location loc) {
-  if (type.isa<StructType>())
-    return builder.create<StructConstantOp>(loc, type,
-                                            value.cast<mlir::ArrayAttr>());
-  return builder.create<ConstantOp>(loc, type,
-                                    value.cast<mlir::DenseElementsAttr>());
-}
-
-//===----------------------------------------------------------------------===//
 // Toy Operations
 //===----------------------------------------------------------------------===//
 
 /// A generalized parser for binary operations. This parses the different forms
 /// of 'printBinaryOp' below.
 static mlir::ParseResult parseBinaryOp(mlir::OpAsmParser &parser,
                                        mlir::OperationState &result) {
   SmallVector<mlir::OpAsmParser::OperandType, 2> operands;

 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "toy/Ops.cpp.inc"
+
+//===----------------------------------------------------------------------===//
+// ToyDialect
+//===----------------------------------------------------------------------===//
+
+/// Dialect creation, the instance will be owned by the context. This is the
+/// point of registration of custom types and operations for the dialect.
+ToyDialect::ToyDialect(mlir::MLIRContext *ctx)
+    : mlir::Dialect(getDialectNamespace(), ctx, TypeID::get<ToyDialect>()) {
+  addOperations<
+#define GET_OP_LIST
+#include "toy/Ops.cpp.inc"
+      >();
+  addInterfaces<ToyInlinerInterface>();
+  addTypes<StructType>();
+}
+
+mlir::Operation *ToyDialect::materializeConstant(mlir::OpBuilder &builder,
+                                                 mlir::Attribute value,
+                                                 mlir::Type type,
+                                                 mlir::Location loc) {
+  if (type.isa<StructType>())
+    return builder.create<StructConstantOp>(loc, type,
+                                            value.cast<mlir::ArrayAttr>());
+  return builder.create<ConstantOp>(loc, type,
+                                    value.cast<mlir::DenseElementsAttr>());
+}

mlir/include/mlir/Dialect/PDL/IR/PDLDialect.td

string description = [{
     how the input operation is specified structurally, similarly to how it would
     look in memory. This is a simple example and pdl provides support for many
     other features such as applying external constraints or external generator
     methods. These features and more are detailed below.
   }];
 
   let name = "pdl";
   let cppNamespace = "::mlir::pdl";
+  let extraClassDeclaration = [{
+    void registerTypes();
+  }];
 }
 
 #endif // MLIR_DIALECT_PDL_IR_PDLDIALECT

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td

def SPIRV_Dialect : Dialect {
 
   let cppNamespace = "::mlir::spirv";
   let hasConstantMaterializer = 1;
   let hasOperationAttrVerify = 1;
   let hasRegionArgAttrVerify = 1;
   let hasRegionResultAttrVerify = 1;
 
   let extraClassDeclaration = [{
+    void registerAttributes();
+    void registerTypes();
+
     //===------------------------------------------------------------------===//
     // Attribute
     //===------------------------------------------------------------------===//
 
     /// Returns the attribute name to use when specifying decorations on results
     /// of operations.
     static std::string getAttributeName(Decoration decoration);
   }];

mlir/include/mlir/IR/BuiltinDialect.td

 include "mlir/IR/OpBase.td"
 
 def Builtin_Dialect : Dialect {
   let summary =
     "A dialect containing the builtin Attributes, Operations, and Types";
 
   let name = "";
   let cppNamespace = "::mlir";
+  let extraClassDeclaration = [{
+  private:
+    // Register the builtin Attributes.
+    void registerAttributes();
+    // Register the builtin Location Attributes.
+    void registerLocationAttributes();
+    // Register the builtin Types.
+    void registerTypes();
+
+  public:
+  }];
 }
 
 #endif // BUILTIN_BASE

mlir/include/mlir/Support/StorageUniquer.h

public:
 
   /// Set the flag specifying if multi-threading is disabled within the uniquer.
   void disableMultithreading(bool disable = true);
 
   /// Register a new parametric storage class, this is necessary to create
   /// instances of this class type. `id` is the type identifier that will be
   /// used to identify this type when creating instances of it via 'get'.
   template <typename Storage> void registerParametricStorageType(TypeID id) {
-    registerParametricStorageTypeImpl(id);
+    // If the storage is trivially destructible, we don't need a destructor
+    // function.
+    if (std::is_trivially_destructible<Storage>::value)
+      return registerParametricStorageTypeImpl(id, nullptr);
+    registerParametricStorageTypeImpl(id, [](BaseStorage *storage) {
+      reinterpret_cast<Storage *>(storage)->~Storage();
+    });
   }
   /// Utility override when the storage type represents the type id.
   template <typename Storage> void registerParametricStorageType() {
     registerParametricStorageType<Storage>(TypeID::get<Storage>());
   }
   /// Register a new singleton storage class, this is necessary to get the
   /// singletone instance. `id` is the type identifier that will be used to
   /// access the singleton instance via 'get'. An optional initialization
private:
   /// Implementation for getting/creating an instance of a derived type with
   /// parametric storage.
   BaseStorage *getParametricStorageTypeImpl(
       TypeID id, unsigned hashValue,
       function_ref<bool(const BaseStorage *)> isEqual,
       function_ref<BaseStorage *(StorageAllocator &)> ctorFn);
 
   /// Implementation for registering an instance of a derived type with
-  /// parametric storage.
-  void registerParametricStorageTypeImpl(TypeID id);
+  /// parametric storage. This method takes an optional destructor function that
+  /// destructs storage instances when necessary.
+  void registerParametricStorageTypeImpl(
+      TypeID id, function_ref<void(BaseStorage *)> destructorFn);
 
   /// Implementation for getting an instance of a derived type with default
   /// storage.
   BaseStorage *getSingletonImpl(TypeID id);
 
   /// Implementation for registering an instance of a derived type with default
   /// storage.
   void

mlir/lib/Dialect/ArmSVE/IR/ArmSVEDialect.cpp

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 
+#define GET_OP_CLASSES
+#include "mlir/Dialect/ArmSVE/ArmSVE.cpp.inc"
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/ArmSVE/ArmSVETypes.cpp.inc"
+
 void arm_sve::ArmSVEDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/ArmSVE/ArmSVE.cpp.inc"
       >();
   addTypes<
 #define GET_TYPEDEF_LIST
 #include "mlir/Dialect/ArmSVE/ArmSVETypes.cpp.inc"
       >();
 }
 
-#define GET_OP_CLASSES
-#include "mlir/Dialect/ArmSVE/ArmSVE.cpp.inc"
-
-#define GET_TYPEDEF_CLASSES
-#include "mlir/Dialect/ArmSVE/ArmSVETypes.cpp.inc"
-
 //===----------------------------------------------------------------------===//
 // ScalableVectorType
 //===----------------------------------------------------------------------===//
 
 Type arm_sve::ArmSVEDialect::parseType(DialectAsmParser &parser) const {
   llvm::SMLoc typeLoc = parser.getCurrentLocation();
   {
     Type genType;

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

 //===- LLVMDialect.cpp - LLVM IR Ops and Dialect registration -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the types and operation details for the LLVM IR dialect in
 // MLIR, and the LLVM IR dialect.  It also registers the dialect.
 //
 //===----------------------------------------------------------------------===//
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "TypeDetail.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/MLIRContext.h"
 

mlir/lib/Dialect/PDL/IR/PDL.cpp

 // PDLDialect
 //===----------------------------------------------------------------------===//
 
 void PDLDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/PDL/IR/PDLOps.cpp.inc"
       >();
-  addTypes<
-#define GET_TYPEDEF_LIST
-#include "mlir/Dialect/PDL/IR/PDLOpsTypes.cpp.inc"
-      >();
+  registerTypes();
 }
 
 /// Returns true if the given operation is used by a "binding" pdl operation
 /// within the main matcher body of a `pdl.pattern`.
 static LogicalResult
 verifyHasBindingUseInMatcher(Operation *op,
                              StringRef bindableContextStr = "`pdl.operation`") {
   // If the pattern is not a pattern, there is nothing to do.

mlir/lib/Dialect/PDL/IR/PDLTypes.cpp

 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/PDL/IR/PDLOpsTypes.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // PDLDialect
 //===----------------------------------------------------------------------===//
 
+void PDLDialect::registerTypes() {
+  addTypes<
+#define GET_TYPEDEF_LIST
+#include "mlir/Dialect/PDL/IR/PDLOpsTypes.cpp.inc"
+      >();
+}
+
 static Type parsePDLType(DialectAsmParser &parser) {
   StringRef typeTag;
   if (parser.parseKeyword(&typeTag))
     return Type();
   {
     Type genType;
     auto parseResult = generatedTypeParser(parser.getBuilder().getContext(),
                                            parser, typeTag, genType);

mlir/lib/Dialect/SPIRV/IR/SPIRVAttributes.cpp

 //===- SPIRVAttributes.cpp - SPIR-V attribute definitions -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/IR/Builders.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // TableGen'erated attribute utility functions
 //===----------------------------------------------------------------------===//
spirv::TargetEnvAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                              spirv::Vendor /*vendorID*/,
                              spirv::DeviceType /*deviceType*/,
                              uint32_t /*deviceID*/, DictionaryAttr limits) {
   if (!limits.isa<spirv::ResourceLimitsAttr>())
     return emitError() << "expected spirv::ResourceLimitsAttr for limits";
 
   return success();
 }
+
+//===----------------------------------------------------------------------===//
+// SPIR-V Dialect
+//===----------------------------------------------------------------------===//
+
+void spirv::SPIRVDialect::registerAttributes() {
+  addAttributes<InterfaceVarABIAttr, TargetEnvAttr, VerCapExtAttr>();
+}

mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp

 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // SPIR-V Dialect
 //===----------------------------------------------------------------------===//
 
 void SPIRVDialect::initialize() {
-  addTypes<ArrayType, CooperativeMatrixNVType, ImageType, MatrixType,
-           PointerType, RuntimeArrayType, SampledImageType, StructType>();
-
-  addAttributes<InterfaceVarABIAttr, TargetEnvAttr, VerCapExtAttr>();
+  registerAttributes();
+  registerTypes();
 
   // Add SPIR-V ops.
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.cpp.inc"
       >();
 
   addInterfaces<SPIRVInlinerInterface>();

mlir/lib/Dialect/SPIRV/IR/SPIRVTypes.cpp

void MatrixType::getCapabilities(
   {
     static const Capability caps[] = {Capability::Matrix};
     ArrayRef<Capability> ref(caps, llvm::array_lengthof(caps));
     capabilities.push_back(ref);
   }
   // Add any capabilities associated with the underlying vectors (i.e., columns)
   getColumnType().cast<SPIRVType>().getCapabilities(capabilities, storage);
 }
+
+//===----------------------------------------------------------------------===//
+// SPIR-V Dialect
+//===----------------------------------------------------------------------===//
+
+void SPIRVDialect::registerTypes() {
+  addTypes<ArrayType, CooperativeMatrixNVType, ImageType, MatrixType,
+           PointerType, RuntimeArrayType, SampledImageType, StructType>();
+}

mlir/lib/Dialect/Vector/VectorOps.cpp

static bool isSupportedCombiningKind(CombiningKind combiningKind,
   case CombiningKind::OR:
   case CombiningKind::XOR:
     return elementType.isIntOrIndex();
   }
   return false;
 }
 
 //===----------------------------------------------------------------------===//
-// VectorDialect
-//===----------------------------------------------------------------------===//
-
-void VectorDialect::initialize() {
-  addAttributes<CombiningKindAttr>();
-
-  addOperations<
-#define GET_OP_LIST
-#include "mlir/Dialect/Vector/VectorOps.cpp.inc"
-      >();
-}
-
-/// Materialize a single constant operation from a given attribute value with
-/// the desired resultant type.
-Operation *VectorDialect::materializeConstant(OpBuilder &builder,
-                                              Attribute value, Type type,
-                                              Location loc) {
-  return builder.create<ConstantOp>(loc, type, value);
-}
-
-IntegerType vector::getVectorSubscriptType(Builder &builder) {
-  return builder.getIntegerType(64);
-}
-
-ArrayAttr vector::getVectorSubscriptAttr(Builder &builder,
-                                         ArrayRef<int64_t> values) {
-  return builder.getI64ArrayAttr(values);
-}
-
-//===----------------------------------------------------------------------===//
 // CombiningKindAttr
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace vector {
 namespace detail {
 struct BitmaskEnumStorage : public AttributeStorage {
   using KeyTy = uint64_t;
void VectorDialect::printAttribute(Attribute attr,
                                    DialectAsmPrinter &os) const {
   if (auto ck = attr.dyn_cast<CombiningKindAttr>())
     ck.print(os);
   else
     llvm_unreachable("Unknown attribute type");
 }
 
 //===----------------------------------------------------------------------===//
+// VectorDialect
+//===----------------------------------------------------------------------===//
+
+void VectorDialect::initialize() {
+  addAttributes<CombiningKindAttr>();
+
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/Vector/VectorOps.cpp.inc"
+      >();
+}
+
+/// Materialize a single constant operation from a given attribute value with
+/// the desired resultant type.
+Operation *VectorDialect::materializeConstant(OpBuilder &builder,
+                                              Attribute value, Type type,
+                                              Location loc) {
+  return builder.create<ConstantOp>(loc, type, value);
+}
+
+IntegerType vector::getVectorSubscriptType(Builder &builder) {
+  return builder.getIntegerType(64);
+}
+
+ArrayAttr vector::getVectorSubscriptAttr(Builder &builder,
+                                         ArrayRef<int64_t> values) {
+  return builder.getI64ArrayAttr(values);
+}
+
+//===----------------------------------------------------------------------===//
 // ReductionOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ReductionOp op) {
   // Verify for 1-D vector.
   int64_t rank = op.getVectorType().getRank();
   if (rank != 1)
     return op.emitOpError("unsupported reduction rank: ") << rank;

mlir/lib/IR/BuiltinAttributes.cpp

 //===- BuiltinAttributes.cpp - MLIR Builtin Attribute Classes -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/BuiltinAttributes.h"
 #include "AttributeDetail.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DecodeAttributesInterfaces.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Endian.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 /// Tablegen Attribute Definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/IR/BuiltinAttributes.cpp.inc"
 
 //===----------------------------------------------------------------------===//
+// BuiltinDialect
+//===----------------------------------------------------------------------===//
+
+void BuiltinDialect::registerAttributes() {
+  addAttributes<AffineMapAttr, ArrayAttr, DenseIntOrFPElementsAttr,
+                DenseStringElementsAttr, DictionaryAttr, FloatAttr,
+                SymbolRefAttr, IntegerAttr, IntegerSetAttr, OpaqueAttr,
+                OpaqueElementsAttr, SparseElementsAttr, StringAttr, TypeAttr,
+                UnitAttr>();
+}
+
+//===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
 /// Helper function that does either an in place sort or sorts from source array
 /// into destination. If inPlace then storage is both the source and the
 /// destination, else value is the source and storage destination. Returns
 /// whether source was sorted.
 template <bool inPlace>

mlir/lib/IR/BuiltinDialect.cpp

if (auto tupleType = type.dyn_cast<TupleType>()) {
       }
     }
     return failure();
   }
 };
 } // end anonymous namespace.
 
 void BuiltinDialect::initialize() {
-  addTypes<ComplexType, BFloat16Type, Float16Type, Float32Type, Float64Type,
-           Float80Type, Float128Type, FunctionType, IndexType, IntegerType,
-           MemRefType, UnrankedMemRefType, NoneType, OpaqueType,
-           RankedTensorType, TupleType, UnrankedTensorType, VectorType>();
-  addAttributes<AffineMapAttr, ArrayAttr, DenseIntOrFPElementsAttr,
-                DenseStringElementsAttr, DictionaryAttr, FloatAttr,
-                SymbolRefAttr, IntegerAttr, IntegerSetAttr, OpaqueAttr,
-                OpaqueElementsAttr, SparseElementsAttr, StringAttr, TypeAttr,
-                UnitAttr>();
-  addAttributes<CallSiteLoc, FileLineColLoc, FusedLoc, NameLoc, OpaqueLoc,
-                UnknownLoc>();
+  registerTypes();
+  registerAttributes();
+  registerLocationAttributes();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/IR/BuiltinOps.cpp.inc"
       >();
   addInterfaces<BuiltinOpAsmDialectInterface>();
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/IR/BuiltinTypes.cpp

 //===- BuiltinTypes.cpp - MLIR Builtin Type Classes -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "TypeDetail.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 /// Tablegen Type Definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/IR/BuiltinTypes.cpp.inc"
 
 //===----------------------------------------------------------------------===//
+// BuiltinDialect
+//===----------------------------------------------------------------------===//
+
+void BuiltinDialect::registerTypes() {
+  addTypes<
+#define GET_TYPEDEF_LIST
+#include "mlir/IR/BuiltinTypes.cpp.inc"
+      >();
+}
+
+//===----------------------------------------------------------------------===//
 /// ComplexType
 //===----------------------------------------------------------------------===//
 
 /// Verify the construction of an integer type.
 LogicalResult ComplexType::verify(function_ref<InFlightDiagnostic()> emitError,
                                   Type elementType) {
   if (!elementType.isIntOrFloat())
     return emitError() << "invalid element type for complex";

 
 LogicalResult MemRefType::verify(function_ref<InFlightDiagnostic()> emitError,
                                  ArrayRef<int64_t> shape, Type elementType,
                                  ArrayRef<AffineMap> affineMapComposition,
                                  unsigned memorySpace) {
   if (!BaseMemRefType::isValidElementType(elementType))
     return emitError() << "invalid memref element type";
 
-    // Negative sizes are not allowed except for `-1` that means dynamic size.
   for (int64_t s : shape)
     if (s < -1)
       return emitError() << "invalid memref size";
 
   // Check that the structure of the composition is valid, i.e. that each
   // subsequent affine map has as many inputs as the previous map has results.
   // Take the dimensionality of the MemRef for the first map.
   size_t dim = shape.size();

mlir/lib/IR/Location.cpp

 //===- Location.cpp - MLIR Location Classes -------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Location.h"
+#include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Identifier.h"
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 /// Tablegen Attribute Definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/IR/BuiltinLocationAttributes.cpp.inc"
 
 //===----------------------------------------------------------------------===//
+// BuiltinDialect
+//===----------------------------------------------------------------------===//
+
+void BuiltinDialect::registerLocationAttributes() {
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "mlir/IR/BuiltinLocationAttributes.cpp.inc"
+      >();
+}
+
+//===----------------------------------------------------------------------===//
 // LocationAttr
 //===----------------------------------------------------------------------===//
 
 /// Methods for support type inquiry through isa, cast, and dyn_cast.
 bool LocationAttr::classof(Attribute attr) {
   return attr.isa<CallSiteLoc, FileLineColLoc, FusedLoc, NameLoc, OpaqueLoc,
                   UnknownLoc>();
 }

mlir/lib/Support/StorageUniquer.cpp

getOrCreateUnsafe(Shard &shard, LookupKey &key,
                     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
     auto existing = shard.instances.insert_as({key.hashValue}, key);
     BaseStorage *&storage = existing.first->storage;
     if (existing.second)
       storage = ctorFn(shard.allocator);
     return storage;
   }
 
+  /// Destroy all of the storage instances within the given shard.
+  void destroyShardInstances(Shard &shard) {
+    if (!destructorFn)
+      return;
+    for (HashedStorage &instance : shard.instances)
+      destructorFn(instance.storage);
+  }
+
 public:
 #if LLVM_ENABLE_THREADS != 0
   /// Initialize the storage uniquer with a given number of storage shards to
-  /// use. The provided shard number is required to be a valid power of 2.
-  ParametricStorageUniquer(size_t numShards = 8)
-      : shards(new std::atomic<Shard *>[numShards]), numShards(numShards) {
+  /// use. The provided shard number is required to be a valid power of 2. The
+  /// destructor function is used to destroy any allocated storage instances.
+  ParametricStorageUniquer(function_ref<void(BaseStorage *)> destructorFn,
+                           size_t numShards = 8)
+      : shards(new std::atomic<Shard *>[numShards]), numShards(numShards),
+        destructorFn(destructorFn) {
     assert(llvm::isPowerOf2_64(numShards) &&
            "the number of shards is required to be a power of 2");
     for (size_t i = 0; i < numShards; i++)
       shards[i].store(nullptr, std::memory_order_relaxed);
   }
   ~ParametricStorageUniquer() {
     // Free all of the allocated shards.
-    for (size_t i = 0; i != numShards; ++i)
-      if (Shard *shard = shards[i].load())
+    for (size_t i = 0; i != numShards; ++i) {
+      if (Shard *shard = shards[i].load()) {
+        destroyShardInstances(*shard);
         delete shard;
       }
+    }
+  }
   /// Get or create an instance of a parametric type.
   BaseStorage *
   getOrCreate(bool threadingIsEnabled, unsigned hashValue,
               function_ref<bool(const BaseStorage *)> isEqual,
               function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
     Shard &shard = getShard(hashValue);
     ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
     if (!threadingIsEnabled)
private:
   /// A set of uniquer shards to allow for further bucketing accesses for
   /// instances of this storage type. Each shard is lazily initialized to reduce
   /// the overhead when only a small amount of shards are in use.
   std::unique_ptr<std::atomic<Shard *>[]> shards;
 
   /// The number of available shards.
   size_t numShards;
 
+  /// Function to used to destruct any allocated storage instances.
+  function_ref<void(BaseStorage *)> destructorFn;
+
 #else
   /// If multi-threading is disabled, ignore the shard parameter as we will
-  /// always use one shard.
-  ParametricStorageUniquer(size_t numShards = 0) {}
+  /// always use one shard. The destructor function is used to destroy any
+  /// allocated storage instances.
+  ParametricStorageUniquer(function_ref<void(BaseStorage *)> destructorFn,
+                           size_t numShards = 0)
+      : destructorFn(destructorFn) {}
+  ~ParametricStorageUniquer() { destroyShardInstances(shard); }
 
   /// Get or create an instance of a parametric type.
   BaseStorage *
   getOrCreate(bool threadingIsEnabled, unsigned hashValue,
               function_ref<bool(const BaseStorage *)> isEqual,
               function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
     ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
     return getOrCreateUnsafe(shard, lookupKey, ctorFn);
   }
   /// Run a mutation function on the provided storage object in a thread-safe
   /// way.
   LogicalResult
   mutate(bool threadingIsEnabled, BaseStorage *storage,
          function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
     return mutationFn(shard.allocator);
   }
 
 private:
   /// The main uniquer shard that is used for allocating storage instances.
   Shard shard;
+
+  /// Function to used to destruct any allocated storage instances.
+  function_ref<void(BaseStorage *)> destructorFn;
 #endif
 };
 } // end anonymous namespace
 
 namespace mlir {
 namespace detail {
 /// This is the implementation of the StorageUniquer class.
 struct StorageUniquerImpl {
auto StorageUniquer::getParametricStorageTypeImpl(
     TypeID id, unsigned hashValue,
     function_ref<bool(const BaseStorage *)> isEqual,
     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) -> BaseStorage * {
   return impl->getOrCreate(id, hashValue, isEqual, ctorFn);
 }
 
 /// Implementation for registering an instance of a derived type with
 /// parametric storage.
-void StorageUniquer::registerParametricStorageTypeImpl(TypeID id) {
+void StorageUniquer::registerParametricStorageTypeImpl(
+    TypeID id, function_ref<void(BaseStorage *)> destructorFn) {
   impl->parametricUniquers.try_emplace(
-      id, std::make_unique<ParametricStorageUniquer>());
+      id, std::make_unique<ParametricStorageUniquer>(destructorFn));
 }
 
 /// Implementation for getting an instance of a derived type with default
 /// storage.
 auto StorageUniquer::getSingletonImpl(TypeID id) -> BaseStorage * {
   return impl->getSingleton(id);
 }
 

mlir/test/lib/Dialect/Test/TestAttributes.cpp

 
 #define GET_ATTRDEF_CLASSES
 #include "TestAttrDefs.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
+void TestDialect::registerAttributes() {
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "TestAttrDefs.cpp.inc"
+      >();
+}
+
 Attribute TestDialect::parseAttribute(DialectAsmParser &parser,
                                       Type type) const {
   StringRef attrTag;
   if (failed(parser.parseKeyword(&attrTag)))
     return Attribute();
   {
     Attribute attr;
     auto parseResult =

mlir/test/lib/Dialect/Test/TestDialect.cpp

 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
 void TestDialect::initialize() {
+  registerAttributes();
+  registerTypes();
   addOperations<
 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
-  addAttributes<
-#define GET_ATTRDEF_LIST
-#include "TestAttrDefs.cpp.inc"
-      >();
   addInterfaces<TestOpAsmInterface, TestDialectFoldInterface,
                 TestInlinerInterface>();
-  addTypes<TestType, TestRecursiveType,
-#define GET_TYPEDEF_LIST
-#include "TestTypeDefs.cpp.inc"
-           >();
   allowUnknownOperations();
 }
 
 Operation *TestDialect::materializeConstant(OpBuilder &builder, Attribute value,
                                             Type type, Location loc) {
   return builder.create<TestOpConstant>(loc, type, value);
 }
 

mlir/test/lib/Dialect/Test/TestOps.td

def Test_Dialect : Dialect {
   let name = "test";
   let cppNamespace = "::mlir::test";
   let hasConstantMaterializer = 1;
   let hasOperationAttrVerify = 1;
   let hasRegionArgAttrVerify = 1;
   let hasRegionResultAttrVerify = 1;
 
   let extraClassDeclaration = [{
+    void registerAttributes();
+    void registerTypes();
+
     Attribute parseAttribute(DialectAsmParser &parser,
                              Type type) const override;
     void printAttribute(Attribute attr,
                         DialectAsmPrinter &printer) const override;
   }];
 }
 
 class TEST_Op<string mnemonic, list<OpTrait> traits = []> :

mlir/test/lib/Dialect/Test/TestTypes.cpp

 
 #define GET_TYPEDEF_CLASSES
 #include "TestTypeDefs.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
+void TestDialect::registerTypes() {
+  addTypes<TestType, TestRecursiveType,
+#define GET_TYPEDEF_LIST
+#include "TestTypeDefs.cpp.inc"
+           >();
+}
+
 static Type parseTestType(MLIRContext *ctxt, DialectAsmParser &parser,
                           llvm::SetVector<Type> &stack) {
   StringRef typeTag;
   if (failed(parser.parseKeyword(&typeTag)))
     return Type();
 
   {
     Type genType;

mlir/unittests/Support/CMakeLists.txt

 add_mlir_unittest(MLIRSupportTests
   DebugActionTest.cpp
   DebugCounterTest.cpp
   IndentedOstreamTest.cpp
   MathExtrasTest.cpp
+  StorageUniquerTest.cpp
 )
 
 target_link_libraries(MLIRSupportTests
   PRIVATE MLIRSupportIndentedOstream MLIRSupport)

mlir/unittests/Support/StorageUniquerTest.cpp

+//===- StorageUniquerTest.cpp - StorageUniquer Tests ----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/StorageUniquer.h"
+#include "gmock/gmock.h"
+
+using namespace mlir;
+
+namespace {
+/// Simple storage class used for testing.
+template <typename ConcreteT, typename... Args>
+struct SimpleStorage : public StorageUniquer::BaseStorage {
+  using Base = SimpleStorage<ConcreteT, Args...>;
+  using KeyTy = std::tuple<Args...>;
+
+  SimpleStorage(KeyTy key) : key(key) {}
+
+  /// Get an instance of this storage instance.
+  template <typename... ParamsT>
+  static ConcreteT *get(StorageUniquer &uniquer, ParamsT &&...params) {
+    return uniquer.get<ConcreteT>(
+        /*initFn=*/{}, std::make_tuple(std::forward<ParamsT>(params)...));
+  }
+
+  /// Construct an instance with the given storage allocator.
+  static ConcreteT *construct(StorageUniquer::StorageAllocator &alloc,
+                              KeyTy key) {
+    return new (alloc.allocate<ConcreteT>())
+        ConcreteT(std::forward<KeyTy>(key));
+  }
+  bool operator==(const KeyTy &key) const { return this->key == key; }
+
+  KeyTy key;
+};
+} // namespace
+
+TEST(StorageUniquerTest, NonTrivialDestructor) {
+  struct NonTrivialStorage : public SimpleStorage<NonTrivialStorage, bool *> {
+    using Base::Base;
+    ~NonTrivialStorage() {
+      bool *wasDestructed = std::get<0>(key);
+      *wasDestructed = true;
+    }
+  };
+
+  // Verify that the storage instance destructor was properly called.
+  bool wasDestructed = false;
+  {
+    StorageUniquer uniquer;
+    uniquer.registerParametricStorageType<NonTrivialStorage>();
+    NonTrivialStorage::get(uniquer, &wasDestructed);
+  }
+
+  EXPECT_TRUE(wasDestructed);
+}