diff --git a/mlir/include/mlir/IR/IntegerSet.h b/mlir/include/mlir/IR/IntegerSet.h --- a/mlir/include/mlir/IR/IntegerSet.h +++ b/mlir/include/mlir/IR/IntegerSet.h @@ -117,8 +117,6 @@ private: ImplType *set; - /// Sets with constraints fewer than kUniquingThreshold are uniqued. - constexpr static unsigned kUniquingThreshold = 4; }; // Make AffineExpr hashable. diff --git a/mlir/lib/IR/AffineMapDetail.h b/mlir/lib/IR/AffineMapDetail.h --- a/mlir/lib/IR/AffineMapDetail.h +++ b/mlir/lib/IR/AffineMapDetail.h @@ -15,12 +15,16 @@ #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" +#include "mlir/Support/StorageUniquer.h" #include "llvm/ADT/ArrayRef.h" namespace mlir { namespace detail { -struct AffineMapStorage { +struct AffineMapStorage : public StorageUniquer::BaseStorage { + /// The hash key used for uniquing. + using KeyTy = std::tuple>; + unsigned numDims; unsigned numSymbols; @@ -29,6 +33,22 @@ ArrayRef results; MLIRContext *context; + + bool operator==(const KeyTy &key) const { + return std::get<0>(key) == numDims && std::get<1>(key) == numSymbols && + std::get<2>(key) == results; + } + + // Constructs an AffineMapStorage from a key. The context must be set by the + // caller. + static AffineMapStorage * + construct(StorageUniquer::StorageAllocator &allocator, const KeyTy &key) { + auto *res = new (allocator.allocate()) AffineMapStorage(); + res->numDims = std::get<0>(key); + res->numSymbols = std::get<1>(key); + res->results = allocator.copyInto(std::get<2>(key)); + return res; + } }; } // end namespace detail diff --git a/mlir/lib/IR/IntegerSet.cpp b/mlir/lib/IR/IntegerSet.cpp --- a/mlir/lib/IR/IntegerSet.cpp +++ b/mlir/lib/IR/IntegerSet.cpp @@ -35,9 +35,6 @@ } bool IntegerSet::isEmptyIntegerSet() const { - // This will only work if uniquing is on. - static_assert(kUniquingThreshold >= 1, - "uniquing threshold should be at least one"); return *this == getEmptySet(set->dimCount, set->symbolCount, getContext()); } diff --git a/mlir/lib/IR/IntegerSetDetail.h b/mlir/lib/IR/IntegerSetDetail.h --- a/mlir/lib/IR/IntegerSetDetail.h +++ b/mlir/lib/IR/IntegerSetDetail.h @@ -14,12 +14,17 @@ #define INTEGERSETDETAIL_H_ #include "mlir/IR/AffineExpr.h" +#include "mlir/Support/StorageUniquer.h" #include "llvm/ADT/ArrayRef.h" namespace mlir { namespace detail { -struct IntegerSetStorage { +struct IntegerSetStorage : public StorageUniquer::BaseStorage { + /// The hash key used for uniquing. + using KeyTy = + std::tuple, ArrayRef>; + unsigned dimCount; unsigned symbolCount; @@ -29,6 +34,22 @@ // Bits to check whether a constraint is an equality or an inequality. ArrayRef eqFlags; + + bool operator==(const KeyTy &key) const { + return std::get<0>(key) == dimCount && std::get<1>(key) == symbolCount && + std::get<2>(key) == constraints && std::get<3>(key) == eqFlags; + } + + static IntegerSetStorage * + construct(StorageUniquer::StorageAllocator &allocator, const KeyTy &key) { + auto *res = + new (allocator.allocate()) IntegerSetStorage(); + res->dimCount = std::get<0>(key); + res->symbolCount = std::get<1>(key); + res->constraints = allocator.copyInto(std::get<2>(key)); + res->eqFlags = allocator.copyInto(std::get<3>(key)); + return res; + } }; } // end namespace detail diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp --- a/mlir/lib/IR/MLIRContext.cpp +++ b/mlir/lib/IR/MLIRContext.cpp @@ -112,91 +112,6 @@ }; } // end anonymous namespace. -//===----------------------------------------------------------------------===// -// AffineMap and IntegerSet hashing -//===----------------------------------------------------------------------===// - -/// A utility function to safely get or create a uniqued instance within the -/// given set container. -template -static ValueT safeGetOrCreate(DenseSet &container, - KeyT &&key, llvm::sys::SmartRWMutex &mutex, - bool threadingIsEnabled, - ConstructorFn &&constructorFn) { - // Check for an existing instance in read-only mode. - if (threadingIsEnabled) { - llvm::sys::SmartScopedReader instanceLock(mutex); - auto it = container.find_as(key); - if (it != container.end()) - return *it; - } - - // Acquire a writer-lock so that we can safely create the new instance. - ScopedWriterLock instanceLock(mutex, threadingIsEnabled); - - // Check for an existing instance again here, because another writer thread - // may have already created one. Otherwise, construct a new instance. - auto existing = container.insert_as(ValueT(), key); - if (existing.second) - return *existing.first = constructorFn(); - return *existing.first; -} - -namespace { -struct AffineMapKeyInfo : DenseMapInfo { - // Affine maps are uniqued based on their dim/symbol counts and affine - // expressions. - using KeyTy = std::tuple>; - using DenseMapInfo::isEqual; - - static unsigned getHashValue(const AffineMap &key) { - return getHashValue( - KeyTy(key.getNumDims(), key.getNumSymbols(), key.getResults())); - } - - static unsigned getHashValue(KeyTy key) { - return hash_combine( - std::get<0>(key), std::get<1>(key), - hash_combine_range(std::get<2>(key).begin(), std::get<2>(key).end())); - } - - static bool isEqual(const KeyTy &lhs, AffineMap rhs) { - if (rhs == getEmptyKey() || rhs == getTombstoneKey()) - return false; - return lhs == std::make_tuple(rhs.getNumDims(), rhs.getNumSymbols(), - rhs.getResults()); - } -}; - -struct IntegerSetKeyInfo : DenseMapInfo { - // Integer sets are uniqued based on their dim/symbol counts, affine - // expressions appearing in the LHS of constraints, and eqFlags. - using KeyTy = - std::tuple, ArrayRef>; - using DenseMapInfo::isEqual; - - static unsigned getHashValue(const IntegerSet &key) { - return getHashValue(KeyTy(key.getNumDims(), key.getNumSymbols(), - key.getConstraints(), key.getEqFlags())); - } - - static unsigned getHashValue(KeyTy key) { - return hash_combine( - std::get<0>(key), std::get<1>(key), - hash_combine_range(std::get<2>(key).begin(), std::get<2>(key).end()), - hash_combine_range(std::get<3>(key).begin(), std::get<3>(key).end())); - } - - static bool isEqual(const KeyTy &lhs, IntegerSet rhs) { - if (rhs == getEmptyKey() || rhs == getTombstoneKey()) - return false; - return lhs == std::make_tuple(rhs.getNumDims(), rhs.getNumSymbols(), - rhs.getConstraints(), rhs.getEqFlags()); - } -}; -} // end anonymous namespace. - //===----------------------------------------------------------------------===// // MLIRContextImpl //===----------------------------------------------------------------------===// @@ -279,19 +194,7 @@ // Affine uniquing //===--------------------------------------------------------------------===// - // Affine allocator and mutex for thread safety. - llvm::BumpPtrAllocator affineAllocator; - llvm::sys::SmartRWMutex affineMutex; - - // Affine map uniquing. - using AffineMapSet = DenseSet; - AffineMapSet affineMaps; - - // Integer set uniquing. - using IntegerSets = DenseSet; - IntegerSets integerSets; - - // Affine expression uniquing. + // Affine expression, map and integer set uniquing. StorageUniquer affineUniquer; //===--------------------------------------------------------------------===// @@ -415,6 +318,8 @@ impl->affineUniquer .registerParametricStorageType(); impl->affineUniquer.registerParametricStorageType(); + impl->affineUniquer.registerParametricStorageType(); + impl->affineUniquer.registerParametricStorageType(); } MLIRContext::~MLIRContext() {} @@ -995,21 +900,10 @@ ArrayRef results, MLIRContext *context) { auto &impl = context->getImpl(); - auto key = std::make_tuple(dimCount, symbolCount, results); - - // Safely get or create an AffineMap instance. - return safeGetOrCreate( - impl.affineMaps, key, impl.affineMutex, impl.threadingIsEnabled, [&] { - auto *res = impl.affineAllocator.Allocate(); - - // Copy the results into the bump pointer. - results = copyArrayRefInto(impl.affineAllocator, results); - - // Initialize the memory using placement new. - new (res) - detail::AffineMapStorage{dimCount, symbolCount, results, context}; - return AffineMap(res); - }); + auto *storage = impl.affineUniquer.get( + [&](AffineMapStorage *storage) { storage->context = context; }, dimCount, + symbolCount, results); + return AffineMap(storage); } /// Check whether the arguments passed to the AffineMap::get() are consistent. @@ -1069,33 +963,9 @@ assert(constraints.size() == eqFlags.size()); auto &impl = constraints[0].getContext()->getImpl(); - - // A utility function to construct a new IntegerSetStorage instance. - auto constructorFn = [&] { - auto *res = impl.affineAllocator.Allocate(); - - // Copy the results and equality flags into the bump pointer. - constraints = copyArrayRefInto(impl.affineAllocator, constraints); - eqFlags = copyArrayRefInto(impl.affineAllocator, eqFlags); - - // Initialize the memory using placement new. - new (res) - detail::IntegerSetStorage{dimCount, symbolCount, constraints, eqFlags}; - return IntegerSet(res); - }; - - // If this instance is uniqued, then we handle it separately so that multiple - // threads may simultaneously access existing instances. - if (constraints.size() < IntegerSet::kUniquingThreshold) { - auto key = std::make_tuple(dimCount, symbolCount, constraints, eqFlags); - return safeGetOrCreate(impl.integerSets, key, impl.affineMutex, - impl.threadingIsEnabled, constructorFn); - } - - // Otherwise, acquire a writer-lock so that we can safely create the new - // instance. - ScopedWriterLock affineLock(impl.affineMutex, impl.threadingIsEnabled); - return constructorFn(); + auto *storage = impl.affineUniquer.get( + [](IntegerSetStorage *) {}, dimCount, symbolCount, constraints, eqFlags); + return IntegerSet(storage); } //===----------------------------------------------------------------------===//