[mlir] Refactor StorageUniquer to require registration of possible storage types

[rriddle]

Summary

This allows for bucketing the different possible storage types, with each bucket having its own allocator/mutex/instance map. This greatly reduces the amount of lock contention when multi-threading is enabled. On some non-trivial .mlir modules (>300K operations), this led to a compile time decrease of a single conversion pass by around half a second(>25%).

Depends On D82595

[jpienaar]

mlir/lib/Dialect/SDBM/SDBMDialect.cpp

18

Could we keep these sorted?

mlir/lib/Support/StorageUniquer.cpp

19

Mmm, this is one definition of simple, but simple is not very descriptive. KindUniquer ? OneOfKindUniquer ?

Or could go ParametricKindUniquer vs NonparametricKindUniquer as this is about whether the kind's uniquer has additional parameters to consider when uniquering.

Perhaps others would disagree here. This is also a local class ...

84

Can both be used at the same time? Would a given kind be able to use both?

[rriddle]

Resolve comments

[rriddle]

Run clang-format

[rriddle]

Rebase

[rriddle]

Rebase

[rriddle]

mlir/lib/Support/StorageUniquer.cpp

19

Used simple/complex because that is how it is described everywhere else, including user facing docs.

84

Not really, but I haven't been able to get a good merged representation that I liked. Having a union between the two always feels weird. Merged them together for now.

[rriddle]

Rebase

[mehdi_amini]

mlir/lib/Support/StorageUniquer.cpp

19

Can you document the struct?

[rriddle]

Resolve comments

[rriddle]

Rebase

mlir/include/mlir/Dialect/SDBM/SDBMDialect.h

 #include "mlir/IR/Dialect.h"
 #include "mlir/Support/StorageUniquer.h"
 
 namespace mlir {
 class MLIRContext;
 
 class SDBMDialect : public Dialect {
 public:
-  SDBMDialect(MLIRContext *context)
-      : Dialect(getDialectNamespace(), context, TypeID::get<SDBMDialect>()) {}
+  SDBMDialect(MLIRContext *context);
 
   /// Since there are no other virtual methods in this derived class, override
   /// the destructor so that key methods get defined in the corresponding
   /// module.
   ~SDBMDialect() override;
 
   static StringRef getDialectNamespace() { return "sdbm"; }
 

mlir/include/mlir/IR/AttributeSupport.h

 // A utility class to get, or create, unique instances of attributes within an
 // MLIRContext. This class manages all creation and uniquing of attributes.
 class AttributeUniquer {
 public:
   /// Get an uniqued instance of attribute T.
   template <typename T, typename... Args>
   static T get(MLIRContext *ctx, unsigned kind, Args &&... args) {
     return ctx->getAttributeUniquer().get<typename T::ImplType>(
+        T::getTypeID(),
         [ctx](AttributeStorage *storage) {
           initializeAttributeStorage(storage, ctx, T::getTypeID());
         },
         kind, std::forward<Args>(args)...);
   }
 
-  template <typename ImplType, typename... Args>
-  static LogicalResult mutate(MLIRContext *ctx, ImplType *impl,
+  template <typename T, typename... Args>
+  static LogicalResult mutate(MLIRContext *ctx, typename T::ImplType *impl,
                               Args &&...args) {
     assert(impl && "cannot mutate null attribute");
-    return ctx->getAttributeUniquer().mutate(impl, std::forward<Args>(args)...);
+    return ctx->getAttributeUniquer().mutate(T::getTypeID(), impl,
+                                             std::forward<Args>(args)...);
   }
 
 private:
   /// Initialize the given attribute storage instance.
   static void initializeAttributeStorage(AttributeStorage *storage,
                                          MLIRContext *ctx, TypeID attrID);
 };
 } // namespace detail
 
 } // end namespace mlir
 
 #endif

mlir/include/mlir/IR/StorageUniquerSupport.h

if (failed(ConcreteT::verifyConstructionInvariants(loc, args...)))
       return ConcreteT();
     return UniquerT::template get<ConcreteT>(loc.getContext(), kind, args...);
   }
 
   /// Mutate the current storage instance. This will not change the unique key.
   /// The arguments are forwarded to 'ConcreteT::mutate'.
   template <typename... Args>
   LogicalResult mutate(Args &&...args) {
-    return UniquerT::mutate(this->getContext(), getImpl(),
+    return UniquerT::template mutate<ConcreteT>(this->getContext(), getImpl(),
                                                 std::forward<Args>(args)...);
   }
 
   /// Default implementation that just returns success.
   template <typename... Args>
   static LogicalResult verifyConstructionInvariants(Args... args) {
     return success();
   }
 
   /// Utility for easy access to the storage instance.
   ImplType *getImpl() const { return static_cast<ImplType *>(this->impl); }
 };
 } // namespace detail
 } // namespace mlir
 
 #endif

mlir/include/mlir/IR/TypeSupport.h

 namespace detail {
 /// A utility class to get, or create, unique instances of types within an
 /// MLIRContext. This class manages all creation and uniquing of types.
 struct TypeUniquer {
   /// Get an uniqued instance of a type T.
   template <typename T, typename... Args>
   static T get(MLIRContext *ctx, unsigned kind, Args &&... args) {
     return ctx->getTypeUniquer().get<typename T::ImplType>(
+        T::getTypeID(),
         [&](TypeStorage *storage) {
           storage->initialize(AbstractType::lookup(T::getTypeID(), ctx));
         },
         kind, std::forward<Args>(args)...);
   }
 
   /// Change the mutable component of the given type instance in the provided
   /// context.
-  template <typename ImplType, typename... Args>
-  static LogicalResult mutate(MLIRContext *ctx, ImplType *impl,
+  template <typename T, typename... Args>
+  static LogicalResult mutate(MLIRContext *ctx, typename T::ImplType *impl,
                               Args &&...args) {
     assert(impl && "cannot mutate null type");
-    return ctx->getTypeUniquer().mutate(impl, std::forward<Args>(args)...);
+    return ctx->getTypeUniquer().mutate(T::getTypeID(), impl,
+                                        std::forward<Args>(args)...);
   }
 };
 } // namespace detail
 
 } // end namespace mlir
 
 #endif

mlir/include/mlir/Support/StorageUniquer.h

 #define MLIR_SUPPORT_STORAGEUNIQUER_H
 
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/Allocator.h"
 
 namespace mlir {
+class TypeID;
+
 namespace detail {
 struct StorageUniquerImpl;
 
 /// Trait to check if ImplTy provides a 'getKey' method with types 'Args'.
 template <typename ImplTy, typename... Args>
 using has_impltype_getkey_t = decltype(ImplTy::getKey(std::declval<Args>()...));
 
 /// Trait to check if ImplTy provides a 'hashKey' method for 'T'.

 ///      that is called when mutating a storage instance. The first argument is
 ///      an allocator to store any mutable data, and the remaining arguments are
 ///      forwarded from the call site. The storage can be mutated at any time
 ///      after creation. Care must be taken to avoid excessive mutation since
 ///      the allocated storage can keep containing previous states. The return
 ///      value of the function is used to indicate whether the mutation was
 ///      successful, e.g., to limit the number of mutations or enable deferred
 ///      one-time assignment of the mutable component.
+///
+/// All storage classes must be registered with the uniquer via
+/// `registerStorageType` using an appropriate unique `TypeID` for the storage
+/// class.
 class StorageUniquer {
 public:
   StorageUniquer();
   ~StorageUniquer();
 
   /// Set the flag specifying if multi-threading is disabled within the uniquer.
   void disableMultithreading(bool disable = true);
 
+  /// Register a new storage object with this uniquer using the given unique
+  /// type id.
+  void registerStorageType(TypeID id);
+
   /// This class acts as the base storage that all storage classes must derived
   /// from.
   class BaseStorage {
   public:
     /// Get the kind classification of this storage.
     unsigned getKind() const { return kind; }
 
   protected:
private:
     llvm::BumpPtrAllocator allocator;
   };
 
   /// Gets a uniqued instance of 'Storage'. 'initFn' is an optional parameter
   /// that can be used to initialize a newly inserted storage instance. This
   /// function is used for derived types that have complex storage or uniquing
   /// constraints.
   template <typename Storage, typename Arg, typename... Args>
-  Storage *get(function_ref<void(Storage *)> initFn, unsigned kind, Arg &&arg,
-               Args &&... args) {
+  Storage *get(const TypeID &id, function_ref<void(Storage *)> initFn,
+               unsigned kind, Arg &&arg, Args &&...args) {
     // Construct a value of the derived key type.
     auto derivedKey =
         getKey<Storage>(std::forward<Arg>(arg), std::forward<Args>(args)...);
 
     // Create a hash of the kind and the derived key.
     unsigned hashValue = getHash<Storage>(kind, derivedKey);
 
     // Generate an equality function for the derived storage.
     auto isEqual = [&derivedKey](const BaseStorage *existing) {
       return static_cast<const Storage &>(*existing) == derivedKey;
     };
 
     // Generate a constructor function for the derived storage.
     auto ctorFn = [&](StorageAllocator &allocator) {
       auto *storage = Storage::construct(allocator, derivedKey);
       if (initFn)
         initFn(storage);
       return storage;
     };
 
     // Get an instance for the derived storage.
-    return static_cast<Storage *>(getImpl(kind, hashValue, isEqual, ctorFn));
+    return static_cast<Storage *>(
+        getImpl(id, kind, hashValue, isEqual, ctorFn));
   }
 
   /// Gets a uniqued instance of 'Storage'. 'initFn' is an optional parameter
   /// that can be used to initialize a newly inserted storage instance. This
   /// function is used for derived types that use no additional storage or
   /// uniquing outside of the kind.
   template <typename Storage>
-  Storage *get(function_ref<void(Storage *)> initFn, unsigned kind) {
+  Storage *get(const TypeID &id, function_ref<void(Storage *)> initFn,
+               unsigned kind) {
     auto ctorFn = [&](StorageAllocator &allocator) {
       auto *storage = new (allocator.allocate<Storage>()) Storage();
       if (initFn)
         initFn(storage);
       return storage;
     };
-    return static_cast<Storage *>(getImpl(kind, ctorFn));
+    return static_cast<Storage *>(getImpl(id, kind, ctorFn));
   }
 
   /// Changes the mutable component of 'storage' by forwarding the trailing
   /// arguments to the 'mutate' function of the derived class.
   template <typename Storage, typename... Args>
-  LogicalResult mutate(Storage *storage, Args &&...args) {
+  LogicalResult mutate(const TypeID &id, Storage *storage, Args &&...args) {
     auto mutationFn = [&](StorageAllocator &allocator) -> LogicalResult {
       return static_cast<Storage &>(*storage).mutate(
           allocator, std::forward<Args>(args)...);
     };
-    return mutateImpl(mutationFn);
+    return mutateImpl(id, mutationFn);
   }
 
   /// Erases a uniqued instance of 'Storage'. This function is used for derived
   /// types that have complex storage or uniquing constraints.
   template <typename Storage, typename Arg, typename... Args>
-  void erase(unsigned kind, Arg &&arg, Args &&... args) {
+  void erase(const TypeID &id, unsigned kind, Arg &&arg, Args &&...args) {
     // Construct a value of the derived key type.
     auto derivedKey =
         getKey<Storage>(std::forward<Arg>(arg), std::forward<Args>(args)...);
 
     // Create a hash of the kind and the derived key.
     unsigned hashValue = getHash<Storage>(kind, derivedKey);
 
     // Generate an equality function for the derived storage.
     auto isEqual = [&derivedKey](const BaseStorage *existing) {
       return static_cast<const Storage &>(*existing) == derivedKey;
     };
 
     // Attempt to erase the storage instance.
-    eraseImpl(kind, hashValue, isEqual, [](BaseStorage *storage) {
+    eraseImpl(id, kind, hashValue, isEqual, [](BaseStorage *storage) {
       static_cast<Storage *>(storage)->cleanup();
     });
   }
 
 private:
   /// Implementation for getting/creating an instance of a derived type with
   /// complex storage.
-  BaseStorage *getImpl(unsigned kind, unsigned hashValue,
+  BaseStorage *getImpl(const TypeID &id, unsigned kind, unsigned hashValue,
                        function_ref<bool(const BaseStorage *)> isEqual,
                        function_ref<BaseStorage *(StorageAllocator &)> ctorFn);
 
   /// Implementation for getting/creating an instance of a derived type with
   /// default storage.
-  BaseStorage *getImpl(unsigned kind,
+  BaseStorage *getImpl(const TypeID &id, unsigned kind,
                        function_ref<BaseStorage *(StorageAllocator &)> ctorFn);
 
   /// Implementation for erasing an instance of a derived type with complex
   /// storage.
-  void eraseImpl(unsigned kind, unsigned hashValue,
+  void eraseImpl(const TypeID &id, unsigned kind, unsigned hashValue,
                  function_ref<bool(const BaseStorage *)> isEqual,
                  function_ref<void(BaseStorage *)> cleanupFn);
 
   /// Implementation for mutating an instance of a derived storage.
   LogicalResult
-  mutateImpl(function_ref<LogicalResult(StorageAllocator &)> mutationFn);
+  mutateImpl(const TypeID &id,
+             function_ref<LogicalResult(StorageAllocator &)> mutationFn);
 
   /// The internal implementation class.
   std::unique_ptr<detail::StorageUniquerImpl> impl;
 
   //===--------------------------------------------------------------------===//
   // Key Construction
   //===--------------------------------------------------------------------===//
 
   /// Used to construct an instance of 'ImplTy::KeyTy' if there is an
   /// 'ImplTy::getKey' function for the provided arguments.
   template <typename ImplTy, typename... Args>
   static typename std::enable_if<
       llvm::is_detected<detail::has_impltype_getkey_t, ImplTy, Args...>::value,
       typename ImplTy::KeyTy>::type
-  getKey(Args &&... args) {
+  getKey(Args &&...args) {
     return ImplTy::getKey(args...);
   }
   /// If there is no 'ImplTy::getKey' method, then we try to directly construct
   /// the 'ImplTy::KeyTy' with the provided arguments.
   template <typename ImplTy, typename... Args>
   static typename std::enable_if<
       !llvm::is_detected<detail::has_impltype_getkey_t, ImplTy, Args...>::value,
       typename ImplTy::KeyTy>::type
-  getKey(Args &&... args) {
+  getKey(Args &&...args) {
     return typename ImplTy::KeyTy(args...);
   }
 
   //===--------------------------------------------------------------------===//
   // Key and Kind Hashing
   //===--------------------------------------------------------------------===//
 
   /// Used to generate a hash for the 'ImplTy::KeyTy' and kind of a storage

mlir/lib/Dialect/SDBM/SDBMDialect.cpp

 //===- SDBMDialect.cpp - MLIR SDBM Dialect --------------------------------===//
 //
 // 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/SDBM/SDBMDialect.h"
+#include "SDBMExprDetail.h"
 
 using namespace mlir;
 
+SDBMDialect::SDBMDialect(MLIRContext *context)
+    : Dialect(getDialectNamespace(), context, TypeID::get<SDBMDialect>()) {
+  uniquer.registerStorageType(TypeID::get<detail::SDBMBinaryExprStorage>());
+  uniquer.registerStorageType(TypeID::get<detail::SDBMConstantExprStorage>());
+  uniquer.registerStorageType(TypeID::get<detail::SDBMDiffExprStorage>());
+  uniquer.registerStorageType(TypeID::get<detail::SDBMNegExprStorage>());
+  uniquer.registerStorageType(TypeID::get<detail::SDBMTermExprStorage>());
+}
+
 SDBMDialect::~SDBMDialect() = default;

mlir/lib/Dialect/SDBM/SDBMExpr.cpp

SDBMSumExpr SDBMSumExpr::get(SDBMTermExpr lhs, SDBMConstantExpr rhs) {
   if (auto lhsSum = lhs.dyn_cast<SDBMSumExpr>()) {
     lhs = lhsSum.getLHS();
     rhs = SDBMConstantExpr::get(rhs.getDialect(),
                                 rhs.getValue() + lhsSum.getRHS().getValue());
   }
 
   StorageUniquer &uniquer = lhs.getDialect()->getUniquer();
   return uniquer.get<detail::SDBMBinaryExprStorage>(
+      TypeID::get<detail::SDBMBinaryExprStorage>(),
       /*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Add), lhs, rhs);
 }
 
 SDBMTermExpr SDBMSumExpr::getLHS() const {
   return static_cast<ImplType *>(impl)->lhs.cast<SDBMTermExpr>();
 }
 
 SDBMConstantExpr SDBMSumExpr::getRHS() const {

 //===----------------------------------------------------------------------===//
 
 SDBMDiffExpr SDBMDiffExpr::get(SDBMDirectExpr lhs, SDBMTermExpr rhs) {
   assert(lhs && "expected SDBM dimension");
   assert(rhs && "expected SDBM dimension");
 
   StorageUniquer &uniquer = lhs.getDialect()->getUniquer();
   return uniquer.get<detail::SDBMDiffExprStorage>(
+      TypeID::get<detail::SDBMDiffExprStorage>(),
       /*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Diff), lhs, rhs);
 }
 
 SDBMDirectExpr SDBMDiffExpr::getLHS() const {
   return static_cast<ImplType *>(impl)->lhs;
 }
 
 SDBMTermExpr SDBMDiffExpr::getRHS() const {
SDBMStripeExpr SDBMStripeExpr::get(SDBMDirectExpr var,
                                    SDBMConstantExpr stripeFactor) {
   assert(var && "expected SDBM variable expression");
   assert(stripeFactor && "expected non-null stripe factor");
   if (stripeFactor.getValue() <= 0)
     llvm::report_fatal_error("non-positive stripe factor");
 
   StorageUniquer &uniquer = var.getDialect()->getUniquer();
   return uniquer.get<detail::SDBMBinaryExprStorage>(
+      TypeID::get<detail::SDBMBinaryExprStorage>(),
       /*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Stripe), var,
       stripeFactor);
 }
 
 SDBMDirectExpr SDBMStripeExpr::getLHS() const {
   if (SDBMVaryingExpr lhs = static_cast<ImplType *>(impl)->lhs)
     return lhs.cast<SDBMDirectExpr>();
   return {};
SDBMDimExpr SDBMDimExpr::get(SDBMDialect *dialect, unsigned position) {
   assert(dialect && "expected non-null dialect");
 
   auto assignDialect = [dialect](detail::SDBMTermExprStorage *storage) {
     storage->dialect = dialect;
   };
 
   StorageUniquer &uniquer = dialect->getUniquer();
   return uniquer.get<detail::SDBMTermExprStorage>(
-      assignDialect, static_cast<unsigned>(SDBMExprKind::DimId), position);
+      TypeID::get<detail::SDBMTermExprStorage>(), assignDialect,
+      static_cast<unsigned>(SDBMExprKind::DimId), position);
 }
 
 //===----------------------------------------------------------------------===//
 // SDBMSymbolExpr
 //===----------------------------------------------------------------------===//
 
 SDBMSymbolExpr SDBMSymbolExpr::get(SDBMDialect *dialect, unsigned position) {
   assert(dialect && "expected non-null dialect");
 
   auto assignDialect = [dialect](detail::SDBMTermExprStorage *storage) {
     storage->dialect = dialect;
   };
 
   StorageUniquer &uniquer = dialect->getUniquer();
   return uniquer.get<detail::SDBMTermExprStorage>(
-      assignDialect, static_cast<unsigned>(SDBMExprKind::SymbolId), position);
+      TypeID::get<detail::SDBMTermExprStorage>(), assignDialect,
+      static_cast<unsigned>(SDBMExprKind::SymbolId), position);
 }
 
 //===----------------------------------------------------------------------===//
 // SDBMConstantExpr
 //===----------------------------------------------------------------------===//
 
 SDBMConstantExpr SDBMConstantExpr::get(SDBMDialect *dialect, int64_t value) {
   assert(dialect && "expected non-null dialect");
 
   auto assignCtx = [dialect](detail::SDBMConstantExprStorage *storage) {
     storage->dialect = dialect;
   };
 
   StorageUniquer &uniquer = dialect->getUniquer();
   return uniquer.get<detail::SDBMConstantExprStorage>(
-      assignCtx, static_cast<unsigned>(SDBMExprKind::Constant), value);
+      TypeID::get<detail::SDBMConstantExprStorage>(), assignCtx,
+      static_cast<unsigned>(SDBMExprKind::Constant), value);
 }
 
 int64_t SDBMConstantExpr::getValue() const {
   return static_cast<ImplType *>(impl)->constant;
 }
 
 //===----------------------------------------------------------------------===//
 // SDBMNegExpr
 //===----------------------------------------------------------------------===//
 
 SDBMNegExpr SDBMNegExpr::get(SDBMDirectExpr var) {
   assert(var && "expected non-null SDBM direct expression");
 
   StorageUniquer &uniquer = var.getDialect()->getUniquer();
   return uniquer.get<detail::SDBMNegExprStorage>(
+      TypeID::get<detail::SDBMNegExprStorage>(),
       /*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Neg), var);
 }
 
 SDBMDirectExpr SDBMNegExpr::getVar() const {
   return static_cast<ImplType *>(impl)->expr;
 }
 
 SDBMExpr mlir::ops_assertions::operator+(SDBMExpr lhs, SDBMExpr rhs) {

mlir/lib/IR/AffineExpr.cpp

 //===- AffineExpr.cpp - MLIR Affine Expr 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/AffineExpr.h"
 #include "AffineExprDetail.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Support/MathExtras.h"
+#include "mlir/Support/TypeID.h"
 #include "llvm/ADT/STLExtras.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 MLIRContext *AffineExpr::getContext() const { return expr->context; }
 
 AffineExprKind AffineExpr::getKind() const {

 static AffineExpr getAffineDimOrSymbol(AffineExprKind kind, unsigned position,
                                        MLIRContext *context) {
   auto assignCtx = [context](AffineDimExprStorage *storage) {
     storage->context = context;
   };
 
   StorageUniquer &uniquer = context->getAffineUniquer();
   return uniquer.get<AffineDimExprStorage>(
-      assignCtx, static_cast<unsigned>(kind), position);
+      TypeID::get<AffineDimExprStorage>(), assignCtx,
+      static_cast<unsigned>(kind), position);
 }
 
 AffineExpr mlir::getAffineDimExpr(unsigned position, MLIRContext *context) {
   return getAffineDimOrSymbol(AffineExprKind::DimId, position, context);
 }
 
 AffineSymbolExpr::AffineSymbolExpr(AffineExpr::ImplType *ptr)
     : AffineExpr(ptr) {}

 
 AffineExpr mlir::getAffineConstantExpr(int64_t constant, MLIRContext *context) {
   auto assignCtx = [context](AffineConstantExprStorage *storage) {
     storage->context = context;
   };
 
   StorageUniquer &uniquer = context->getAffineUniquer();
   return uniquer.get<AffineConstantExprStorage>(
-      assignCtx, static_cast<unsigned>(AffineExprKind::Constant), constant);
+      TypeID::get<AffineConstantExprStorage>(), assignCtx,
+      static_cast<unsigned>(AffineExprKind::Constant), constant);
 }
 
 /// Simplify add expression. Return nullptr if it can't be simplified.
 static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
   auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
   // Fold if both LHS, RHS are a constant.
   if (lhsConst && rhsConst)
AffineExpr AffineExpr::operator+(int64_t v) const {
   return *this + getAffineConstantExpr(v, getContext());
 }
 AffineExpr AffineExpr::operator+(AffineExpr other) const {
   if (auto simplified = simplifyAdd(*this, other))
     return simplified;
 
   StorageUniquer &uniquer = getContext()->getAffineUniquer();
   return uniquer.get<AffineBinaryOpExprStorage>(
+      TypeID::get<AffineBinaryOpExprStorage>(),
       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Add), *this, other);
 }
 
 /// Simplify a multiply expression. Return nullptr if it can't be simplified.
 static AffineExpr simplifyMul(AffineExpr lhs, AffineExpr rhs) {
   auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
 
AffineExpr AffineExpr::operator*(int64_t v) const {
   return *this * getAffineConstantExpr(v, getContext());
 }
 AffineExpr AffineExpr::operator*(AffineExpr other) const {
   if (auto simplified = simplifyMul(*this, other))
     return simplified;
 
   StorageUniquer &uniquer = getContext()->getAffineUniquer();
   return uniquer.get<AffineBinaryOpExprStorage>(
+      TypeID::get<AffineBinaryOpExprStorage>(),
       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mul), *this, other);
 }
 
 // Unary minus, delegate to operator*.
 AffineExpr AffineExpr::operator-() const {
   return *this * getAffineConstantExpr(-1, getContext());
 }
 
AffineExpr AffineExpr::floorDiv(uint64_t v) const {
   return floorDiv(getAffineConstantExpr(v, getContext()));
 }
 AffineExpr AffineExpr::floorDiv(AffineExpr other) const {
   if (auto simplified = simplifyFloorDiv(*this, other))
     return simplified;
 
   StorageUniquer &uniquer = getContext()->getAffineUniquer();
   return uniquer.get<AffineBinaryOpExprStorage>(
+      TypeID::get<AffineBinaryOpExprStorage>(),
       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::FloorDiv), *this,
       other);
 }
 
 static AffineExpr simplifyCeilDiv(AffineExpr lhs, AffineExpr rhs) {
   auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
 
AffineExpr AffineExpr::ceilDiv(uint64_t v) const {
   return ceilDiv(getAffineConstantExpr(v, getContext()));
 }
 AffineExpr AffineExpr::ceilDiv(AffineExpr other) const {
   if (auto simplified = simplifyCeilDiv(*this, other))
     return simplified;
 
   StorageUniquer &uniquer = getContext()->getAffineUniquer();
   return uniquer.get<AffineBinaryOpExprStorage>(
+      TypeID::get<AffineBinaryOpExprStorage>(),
       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::CeilDiv), *this,
       other);
 }
 
 static AffineExpr simplifyMod(AffineExpr lhs, AffineExpr rhs) {
   auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
 
AffineExpr AffineExpr::operator%(uint64_t v) const {
   return *this % getAffineConstantExpr(v, getContext());
 }
 AffineExpr AffineExpr::operator%(AffineExpr other) const {
   if (auto simplified = simplifyMod(*this, other))
     return simplified;
 
   StorageUniquer &uniquer = getContext()->getAffineUniquer();
   return uniquer.get<AffineBinaryOpExprStorage>(
+      TypeID::get<AffineBinaryOpExprStorage>(),
       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mod), *this, other);
 }
 
 AffineExpr AffineExpr::compose(AffineMap map) const {
   SmallVector<AffineExpr, 8> dimReplacements(map.getResults().begin(),
                                              map.getResults().end());
   return replaceDimsAndSymbols(dimReplacements, {});
 }

mlir/lib/IR/MLIRContext.cpp

MLIRContext::MLIRContext() : impl(new MLIRContextImpl()) {
   impl->unitAttr =
       AttributeUniquer::get<UnitAttr>(this, StandardAttributes::Unit);
   /// Unknown Location Attribute.
   impl->unknownLocAttr = AttributeUniquer::get<UnknownLoc>(
       this, StandardAttributes::UnknownLocation);
   /// The empty dictionary attribute.
   impl->emptyDictionaryAttr = AttributeUniquer::get<DictionaryAttr>(
       this, StandardAttributes::Dictionary, ArrayRef<NamedAttribute>());
+
+  // Register the affine storage objects with the uniquer.
+  impl->affineUniquer.registerStorageType(
+      TypeID::get<AffineBinaryOpExprStorage>());
+  impl->affineUniquer.registerStorageType(
+      TypeID::get<AffineConstantExprStorage>());
+  impl->affineUniquer.registerStorageType(TypeID::get<AffineDimExprStorage>());
 }
 
 MLIRContext::~MLIRContext() {}
 
 /// Copy the specified array of elements into memory managed by the provided
 /// bump pointer allocator.  This assumes the elements are all PODs.
 template <typename T>
 static ArrayRef<T> copyArrayRefInto(llvm::BumpPtrAllocator &allocator,

 
 void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
   auto &impl = context->getImpl();
   auto *newInfo =
       new (impl.abstractDialectSymbolAllocator.Allocate<AbstractType>())
           AbstractType(std::move(typeInfo));
   if (!impl.registeredTypes.insert({typeID, newInfo}).second)
     llvm::report_fatal_error("Dialect Type already registered.");
+  impl.typeUniquer.registerStorageType(typeID);
 }
 
 void Dialect::addAttribute(TypeID typeID, AbstractAttribute &&attrInfo) {
   auto &impl = context->getImpl();
   auto *newInfo =
       new (impl.abstractDialectSymbolAllocator.Allocate<AbstractAttribute>())
           AbstractAttribute(std::move(attrInfo));
   if (!impl.registeredAttributes.insert({typeID, newInfo}).second)
     llvm::report_fatal_error("Dialect Attribute already registered.");
+  impl.attributeUniquer.registerStorageType(typeID);
 }
 
 /// Get the dialect that registered the attribute with the provided typeid.
 const AbstractAttribute &AbstractAttribute::lookup(TypeID typeID,
                                                    MLIRContext *context) {
   auto &impl = context->getImpl();
   auto it = impl.registeredAttributes.find(typeID);
   if (it == impl.registeredAttributes.end())

mlir/lib/Support/StorageUniquer.cpp

 //===- StorageUniquer.cpp - Common Storage Class Uniquer ------------------===//
 //
 // 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 "mlir/Support/LLVM.h"
+#include "mlir/Support/TypeID.h"
 #include "llvm/Support/RWMutex.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
-namespace mlir {
-namespace detail {
-/// This is the implementation of the StorageUniquer class.
-struct StorageUniquerImpl {
+namespace {
+/// This class represents a uniquer for storage instances of a specific type. It
+/// contains all of the necessary data to unique storage instances in a thread
+/// safe way. This allows for the main uniquer to bucket each of the individual
+/// sub-types removing the need to lock the main uniquer itself.
+struct InstSpecificUniquer {
   using BaseStorage = StorageUniquer::BaseStorage;
   using StorageAllocator = StorageUniquer::StorageAllocator;
 
   /// A lookup key for derived instances of storage objects.
   struct LookupKey {
     /// The known derived kind for the storage.
     unsigned kind;
 
     /// The known hash value of the key.
     unsigned hashValue;
 
     /// An equality function for comparing with an existing storage instance.
     function_ref<bool(const BaseStorage *)> isEqual;
   };
 
   /// A utility wrapper object representing a hashed storage object. This class
   /// contains a storage object and an existing computed hash value.
   struct HashedStorage {
     unsigned hashValue;
     BaseStorage *storage;
   };
 
+  /// Storage info for derived TypeStorage objects.
+  struct StorageKeyInfo : DenseMapInfo<HashedStorage> {
+    static HashedStorage getEmptyKey() {
+      return HashedStorage{0, DenseMapInfo<BaseStorage *>::getEmptyKey()};
+    }
+    static HashedStorage getTombstoneKey() {
+      return HashedStorage{0, DenseMapInfo<BaseStorage *>::getTombstoneKey()};
+    }
+
+    static unsigned getHashValue(const HashedStorage &key) {
+      return key.hashValue;
+    }
+    static unsigned getHashValue(LookupKey key) { return key.hashValue; }
+
+    static bool isEqual(const HashedStorage &lhs, const HashedStorage &rhs) {
+      return lhs.storage == rhs.storage;
+    }
+    static bool isEqual(const LookupKey &lhs, const HashedStorage &rhs) {
+      if (isEqual(rhs, getEmptyKey()) || isEqual(rhs, getTombstoneKey()))
+        return false;
+      // If the lookup kind matches the kind of the storage, then invoke the
+      // equality function on the lookup key.
+      return lhs.kind == rhs.storage->getKind() && lhs.isEqual(rhs.storage);
+    }
+  };
+
+  /// Unique types with specific hashing or storage constraints.
+  using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
+  StorageTypeSet complexInstances;
+
+  /// Instances of this storage object.
+  llvm::SmallDenseMap<unsigned, BaseStorage *, 1> simpleInstances;
+
+  /// Allocator to use when constructing derived instances.
+  StorageAllocator allocator;
+
+  /// A mutex to keep type uniquing thread-safe.
+  llvm::sys::SmartRWMutex<true> mutex;
+};
+} // end anonymous namespace
+
+namespace mlir {
+namespace detail {
+/// This is the implementation of the StorageUniquer class.
+struct StorageUniquerImpl {
+  using BaseStorage = StorageUniquer::BaseStorage;
+  using StorageAllocator = StorageUniquer::StorageAllocator;
+
   /// Get or create an instance of a complex derived type.
   BaseStorage *
-  getOrCreate(unsigned kind, unsigned hashValue,
+  getOrCreate(TypeID id, unsigned kind, unsigned hashValue,
               function_ref<bool(const BaseStorage *)> isEqual,
               function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
-    LookupKey lookupKey{kind, hashValue, isEqual};
+    assert(instUniquers.count(id) && "creating unregistered storage instance");
+    InstSpecificUniquer::LookupKey lookupKey{kind, hashValue, isEqual};
+    InstSpecificUniquer &storageUniquer = *instUniquers[id];
     if (!threadingIsEnabled)
-      return getOrCreateUnsafe(kind, hashValue, lookupKey, ctorFn);
+      return getOrCreateUnsafe(storageUniquer, kind, lookupKey, ctorFn);
 
     // Check for an existing instance in read-only mode.
     {
-      llvm::sys::SmartScopedReader<true> typeLock(mutex);
-      auto it = storageTypes.find_as(lookupKey);
-      if (it != storageTypes.end())
+      llvm::sys::SmartScopedReader<true> typeLock(storageUniquer.mutex);
+      auto it = storageUniquer.complexInstances.find_as(lookupKey);
+      if (it != storageUniquer.complexInstances.end())
         return it->storage;
     }
 
     // Acquire a writer-lock so that we can safely create the new type instance.
-    llvm::sys::SmartScopedWriter<true> typeLock(mutex);
-    return getOrCreateUnsafe(kind, hashValue, lookupKey, ctorFn);
+    llvm::sys::SmartScopedWriter<true> typeLock(storageUniquer.mutex);
+    return getOrCreateUnsafe(storageUniquer, kind, lookupKey, ctorFn);
   }
-  /// Get or create an instance of a complex derived type in an unsafe fashion.
+  /// Get or create an instance of a complex derived type in an thread-unsafe
+  /// fashion.
   BaseStorage *
-  getOrCreateUnsafe(unsigned kind, unsigned hashValue, LookupKey &lookupKey,
+  getOrCreateUnsafe(InstSpecificUniquer &storageUniquer, unsigned kind,
+                    InstSpecificUniquer::LookupKey &lookupKey,
                     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
-    auto existing = storageTypes.insert_as({}, lookupKey);
+    auto existing = storageUniquer.complexInstances.insert_as({}, lookupKey);
     if (!existing.second)
       return existing.first->storage;
 
     // Otherwise, construct and initialize the derived storage for this type
     // instance.
-    BaseStorage *storage = initializeStorage(kind, ctorFn);
-    *existing.first = HashedStorage{hashValue, storage};
+    BaseStorage *storage =
+        initializeStorage(kind, storageUniquer.allocator, ctorFn);
+    *existing.first =
+        InstSpecificUniquer::HashedStorage{lookupKey.hashValue, storage};
     return storage;
   }
 
   /// Get or create an instance of a simple derived type.
   BaseStorage *
-  getOrCreate(unsigned kind,
+  getOrCreate(TypeID id, unsigned kind,
               function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
+    assert(instUniquers.count(id) && "creating unregistered storage instance");
+    InstSpecificUniquer &storageUniquer = *instUniquers[id];
     if (!threadingIsEnabled)
-      return getOrCreateUnsafe(kind, ctorFn);
+      return getOrCreateUnsafe(storageUniquer, kind, ctorFn);
 
     // Check for an existing instance in read-only mode.
     {
-      llvm::sys::SmartScopedReader<true> typeLock(mutex);
-      auto it = simpleTypes.find(kind);
-      if (it != simpleTypes.end())
+      llvm::sys::SmartScopedReader<true> typeLock(storageUniquer.mutex);
+      auto it = storageUniquer.simpleInstances.find(kind);
+      if (it != storageUniquer.simpleInstances.end())
         return it->second;
     }
 
     // Acquire a writer-lock so that we can safely create the new type instance.
-    llvm::sys::SmartScopedWriter<true> typeLock(mutex);
-    return getOrCreateUnsafe(kind, ctorFn);
+    llvm::sys::SmartScopedWriter<true> typeLock(storageUniquer.mutex);
+    return getOrCreateUnsafe(storageUniquer, kind, ctorFn);
   }
-  /// Get or create an instance of a simple derived type in an unsafe fashion.
+  /// Get or create an instance of a simple derived type in an thread-unsafe
+  /// fashion.
   BaseStorage *
-  getOrCreateUnsafe(unsigned kind,
+  getOrCreateUnsafe(InstSpecificUniquer &storageUniquer, unsigned kind,
                     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
-    auto &result = simpleTypes[kind];
+    auto &result = storageUniquer.simpleInstances[kind];
     if (result)
       return result;
 
     // Otherwise, create and return a new storage instance.
-    return result = initializeStorage(kind, ctorFn);
+    return result = initializeStorage(kind, storageUniquer.allocator, ctorFn);
   }
 
   /// Erase an instance of a complex derived type.
-  void erase(unsigned kind, unsigned hashValue,
+  void erase(TypeID id, unsigned kind, unsigned hashValue,
              function_ref<bool(const BaseStorage *)> isEqual,
              function_ref<void(BaseStorage *)> cleanupFn) {
-    LookupKey lookupKey{kind, hashValue, isEqual};
+    assert(instUniquers.count(id) && "erasing unregistered storage instance");
+    InstSpecificUniquer &storageUniquer = *instUniquers[id];
+    InstSpecificUniquer::LookupKey lookupKey{kind, hashValue, isEqual};
 
     // Acquire a writer-lock so that we can safely erase the type instance.
-    llvm::sys::SmartScopedWriter<true> typeLock(mutex);
-    auto existing = storageTypes.find_as(lookupKey);
-    if (existing == storageTypes.end())
+    llvm::sys::SmartScopedWriter<true> lock(storageUniquer.mutex);
+    auto existing = storageUniquer.complexInstances.find_as(lookupKey);
+    if (existing == storageUniquer.complexInstances.end())
       return;
 
     // Cleanup the storage and remove it from the map.
     cleanupFn(existing->storage);
-    storageTypes.erase(existing);
+    storageUniquer.complexInstances.erase(existing);
   }
 
   /// Mutates an instance of a derived storage in a thread-safe way.
   LogicalResult
-  mutate(function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
+  mutate(TypeID id,
+         function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
+    assert(instUniquers.count(id) && "mutating unregistered storage instance");
+    InstSpecificUniquer &storageUniquer = *instUniquers[id];
     if (!threadingIsEnabled)
-      return mutationFn(allocator);
+      return mutationFn(storageUniquer.allocator);
 
-    llvm::sys::SmartScopedWriter<true> lock(mutex);
-    return mutationFn(allocator);
+    llvm::sys::SmartScopedWriter<true> lock(storageUniquer.mutex);
+    return mutationFn(storageUniquer.allocator);
   }
 
   //===--------------------------------------------------------------------===//
   // Instance Storage
   //===--------------------------------------------------------------------===//
 
   /// Utility to create and initialize a storage instance.
   BaseStorage *
-  initializeStorage(unsigned kind,
+  initializeStorage(unsigned kind, StorageAllocator &allocator,
                     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
     BaseStorage *storage = ctorFn(allocator);
     storage->kind = kind;
     return storage;
   }
 
-  /// Storage info for derived TypeStorage objects.
-  struct StorageKeyInfo : DenseMapInfo<HashedStorage> {
-    static HashedStorage getEmptyKey() {
-      return HashedStorage{0, DenseMapInfo<BaseStorage *>::getEmptyKey()};
-    }
-    static HashedStorage getTombstoneKey() {
-      return HashedStorage{0, DenseMapInfo<BaseStorage *>::getTombstoneKey()};
-    }
-
-    static unsigned getHashValue(const HashedStorage &key) {
-      return key.hashValue;
-    }
-    static unsigned getHashValue(LookupKey key) { return key.hashValue; }
-
-    static bool isEqual(const HashedStorage &lhs, const HashedStorage &rhs) {
-      return lhs.storage == rhs.storage;
-    }
-    static bool isEqual(const LookupKey &lhs, const HashedStorage &rhs) {
-      if (isEqual(rhs, getEmptyKey()) || isEqual(rhs, getTombstoneKey()))
-        return false;
-      // If the lookup kind matches the kind of the storage, then invoke the
-      // equality function on the lookup key.
-      return lhs.kind == rhs.storage->getKind() && lhs.isEqual(rhs.storage);
-    }
-  };
-
-  /// Unique types with specific hashing or storage constraints.
-  using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
-  StorageTypeSet storageTypes;
-
-  /// Unique types with just the kind.
-  DenseMap<unsigned, BaseStorage *> simpleTypes;
-
-  /// Allocator to use when constructing derived type instances.
-  StorageUniquer::StorageAllocator allocator;
-
-  /// A mutex to keep type uniquing thread-safe.
-  llvm::sys::SmartRWMutex<true> mutex;
+  /// Map of type ids to the storage uniquer to use for registered objects.
+  DenseMap<TypeID, std::unique_ptr<InstSpecificUniquer>> instUniquers;
 
   /// Flag specifying if multi-threading is enabled within the uniquer.
   bool threadingIsEnabled = true;
 };
 } // end namespace detail
 } // namespace mlir
 
 StorageUniquer::StorageUniquer() : impl(new StorageUniquerImpl()) {}
 StorageUniquer::~StorageUniquer() {}
 
 /// Set the flag specifying if multi-threading is disabled within the uniquer.
 void StorageUniquer::disableMultithreading(bool disable) {
   impl->threadingIsEnabled = !disable;
 }
 
+/// Register a new storage object with this uniquer using the given unique type
+/// id.
+void StorageUniquer::registerStorageType(TypeID id) {
+  impl->instUniquers.try_emplace(id, std::make_unique<InstSpecificUniquer>());
+}
+
 /// Implementation for getting/creating an instance of a derived type with
 /// complex storage.
 auto StorageUniquer::getImpl(
-    unsigned kind, unsigned hashValue,
+    const TypeID &id, unsigned kind, unsigned hashValue,
     function_ref<bool(const BaseStorage *)> isEqual,
     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) -> BaseStorage * {
-  return impl->getOrCreate(kind, hashValue, isEqual, ctorFn);
+  return impl->getOrCreate(id, kind, hashValue, isEqual, ctorFn);
 }
 
 /// Implementation for getting/creating an instance of a derived type with
 /// default storage.
 auto StorageUniquer::getImpl(
-    unsigned kind, function_ref<BaseStorage *(StorageAllocator &)> ctorFn)
-    -> BaseStorage * {
-  return impl->getOrCreate(kind, ctorFn);
+    const TypeID &id, unsigned kind,
+    function_ref<BaseStorage *(StorageAllocator &)> ctorFn) -> BaseStorage * {
+  return impl->getOrCreate(id, kind, ctorFn);
 }
 
 /// Implementation for erasing an instance of a derived type with complex
 /// storage.
-void StorageUniquer::eraseImpl(unsigned kind, unsigned hashValue,
+void StorageUniquer::eraseImpl(const TypeID &id, unsigned kind,
+                               unsigned hashValue,
                                function_ref<bool(const BaseStorage *)> isEqual,
                                function_ref<void(BaseStorage *)> cleanupFn) {
-  impl->erase(kind, hashValue, isEqual, cleanupFn);
+  impl->erase(id, kind, hashValue, isEqual, cleanupFn);
 }
 
 /// Implementation for mutating an instance of a derived storage.
 LogicalResult StorageUniquer::mutateImpl(
+    const TypeID &id,
     function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
-  return impl->mutate(mutationFn);
+  return impl->mutate(id, mutationFn);
 }