[Mlir-commits] [mlir] e45705a - [MLIR] Use a shared uniquer for affine maps and integer sets.
Alex Zinenko
llvmlistbot at llvm.org
Thu Dec 2 14:49:39 PST 2021
Author: Ulysse Beaugnon
Date: 2021-12-02T23:49:32+01:00
New Revision: e45705ad5051247e1cb2355094ef77316ae24c37
URL: https://github.com/llvm/llvm-project/commit/e45705ad5051247e1cb2355094ef77316ae24c37
DIFF: https://github.com/llvm/llvm-project/commit/e45705ad5051247e1cb2355094ef77316ae24c37.diff
LOG: [MLIR] Use a shared uniquer for affine maps and integer sets.
Affine maps and integer sets previously relied on a single lock for creating unique instances. In a multi-threaded setting, this lock becomes a contention point. This commit updates AffineMap and IntegerSet to use StorageUniquer instead. StorageUniquer internally uses sharded locks and thread-local caches to reduce contention. It is already used for affine expressions, types and attributes. On my local machine, this gives me a 5X speedup for an application that manipulates a lot of affine maps and integer sets.
This commit also removes the integer set uniquer threshold. The threshold was used to avoid adding integer sets with a lot of constraints to the hash_map containing unique instances, but the constraints and the integer set were still allocated in the same allocator and never freed, thus not saving any space expect for the hash-map entry.
Reviewed By: rriddle
Differential Revision: https://reviews.llvm.org/D114942
Added:
Modified:
mlir/include/mlir/IR/IntegerSet.h
mlir/lib/IR/AffineMapDetail.h
mlir/lib/IR/IntegerSet.cpp
mlir/lib/IR/IntegerSetDetail.h
mlir/lib/IR/MLIRContext.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/IR/IntegerSet.h b/mlir/include/mlir/IR/IntegerSet.h
index 7a7457a213c33..77e1a0cd620cc 100644
--- a/mlir/include/mlir/IR/IntegerSet.h
+++ b/mlir/include/mlir/IR/IntegerSet.h
@@ -117,8 +117,6 @@ class IntegerSet {
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
index 2d059a131d217..d3090e3a59173 100644
--- 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, unsigned, ArrayRef<AffineExpr>>;
+
unsigned numDims;
unsigned numSymbols;
@@ -29,6 +33,22 @@ struct AffineMapStorage {
ArrayRef<AffineExpr> 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>()) 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
index bb1f2920a492f..706580aa838b0 100644
--- a/mlir/lib/IR/IntegerSet.cpp
+++ b/mlir/lib/IR/IntegerSet.cpp
@@ -35,9 +35,6 @@ unsigned IntegerSet::getNumInequalities() const {
}
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
index 3f049dccfc398..451c1594ca12e 100644
--- 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<unsigned, unsigned, ArrayRef<AffineExpr>, ArrayRef<bool>>;
+
unsigned dimCount;
unsigned symbolCount;
@@ -29,6 +34,22 @@ struct IntegerSetStorage {
// Bits to check whether a constraint is an equality or an inequality.
ArrayRef<bool> 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>()) 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
index 6f54703e9b625..9ccdd97d6fe7d 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -112,91 +112,6 @@ struct ScopedWriterLock {
};
} // end anonymous namespace.
-//===----------------------------------------------------------------------===//
-// AffineMap and IntegerSet hashing
-//===----------------------------------------------------------------------===//
-
-/// A utility function to safely get or create a uniqued instance within the
-/// given set container.
-template <typename ValueT, typename DenseInfoT, typename KeyT,
- typename ConstructorFn>
-static ValueT safeGetOrCreate(DenseSet<ValueT, DenseInfoT> &container,
- KeyT &&key, llvm::sys::SmartRWMutex<true> &mutex,
- bool threadingIsEnabled,
- ConstructorFn &&constructorFn) {
- // Check for an existing instance in read-only mode.
- if (threadingIsEnabled) {
- llvm::sys::SmartScopedReader<true> 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<AffineMap> {
- // Affine maps are uniqued based on their dim/symbol counts and affine
- // expressions.
- using KeyTy = std::tuple<unsigned, unsigned, ArrayRef<AffineExpr>>;
- using DenseMapInfo<AffineMap>::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<IntegerSet> {
- // Integer sets are uniqued based on their dim/symbol counts, affine
- // expressions appearing in the LHS of constraints, and eqFlags.
- using KeyTy =
- std::tuple<unsigned, unsigned, ArrayRef<AffineExpr>, ArrayRef<bool>>;
- using DenseMapInfo<IntegerSet>::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 @@ class MLIRContextImpl {
// Affine uniquing
//===--------------------------------------------------------------------===//
- // Affine allocator and mutex for thread safety.
- llvm::BumpPtrAllocator affineAllocator;
- llvm::sys::SmartRWMutex<true> affineMutex;
-
- // Affine map uniquing.
- using AffineMapSet = DenseSet<AffineMap, AffineMapKeyInfo>;
- AffineMapSet affineMaps;
-
- // Integer set uniquing.
- using IntegerSets = DenseSet<IntegerSet, IntegerSetKeyInfo>;
- IntegerSets integerSets;
-
- // Affine expression uniquing.
+ // Affine expression, map and integer set uniquing.
StorageUniquer affineUniquer;
//===--------------------------------------------------------------------===//
@@ -415,6 +318,8 @@ MLIRContext::MLIRContext(const DialectRegistry ®istry, Threading setting)
impl->affineUniquer
.registerParametricStorageType<AffineConstantExprStorage>();
impl->affineUniquer.registerParametricStorageType<AffineDimExprStorage>();
+ impl->affineUniquer.registerParametricStorageType<AffineMapStorage>();
+ impl->affineUniquer.registerParametricStorageType<IntegerSetStorage>();
}
MLIRContext::~MLIRContext() {}
@@ -995,21 +900,10 @@ AffineMap AffineMap::getImpl(unsigned dimCount, unsigned symbolCount,
ArrayRef<AffineExpr> 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<detail::AffineMapStorage>();
-
- // 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>(
+ [&](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 @@ IntegerSet IntegerSet::get(unsigned dimCount, unsigned symbolCount,
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<detail::IntegerSetStorage>();
-
- // 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>(
+ [](IntegerSetStorage *) {}, dimCount, symbolCount, constraints, eqFlags);
+ return IntegerSet(storage);
}
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list