maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[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
This commit is contained in:
Ulysse Beaugnon 2021-12-02 23:49:02 +01:00 committed by Alex Zinenko
parent 36e6a259c8
commit e45705ad50
5 changed files with 53 additions and 147 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mlir/include/mlir/IR/IntegerSet.h
View File

@ -117,8 +117,6 @@ public:
private:
ImplType *set;
/// Sets with constraints fewer than kUniquingThreshold are uniqued.
constexpr static unsigned kUniquingThreshold = 4;
};
// Make AffineExpr hashable.

22
mlir/lib/IR/AffineMapDetail.h
View File

@ -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

3
mlir/lib/IR/IntegerSet.cpp
View File

@ -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());
}

23
mlir/lib/IR/IntegerSetDetail.h
View File

@ -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

150
mlir/lib/IR/MLIRContext.cpp
View File

@ -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 @@ public:
// 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 &registry, 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);
}
//===----------------------------------------------------------------------===//