forked from OSchip/llvm-project
[mlir] Refactor StorageUniquer to require registration of possible storage types
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%). Differential Revision: https://reviews.llvm.org/D82596
This commit is contained in:
parent
164a02d0fa
commit
86646be315
|
@ -17,8 +17,7 @@ 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
|
||||
|
|
|
@ -133,17 +133,19 @@ public:
|
|||
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:
|
||||
|
|
|
@ -109,7 +109,7 @@ protected:
|
|||
/// 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)...);
|
||||
}
|
||||
|
||||
|
|
|
@ -127,6 +127,7 @@ struct TypeUniquer {
|
|||
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));
|
||||
},
|
||||
|
@ -135,11 +136,12 @@ struct TypeUniquer {
|
|||
|
||||
/// 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
|
||||
|
|
|
@ -15,6 +15,8 @@
|
|||
#include "llvm/Support/Allocator.h"
|
||||
|
||||
namespace mlir {
|
||||
class TypeID;
|
||||
|
||||
namespace detail {
|
||||
struct StorageUniquerImpl;
|
||||
|
||||
|
@ -75,6 +77,10 @@ using has_impltype_hash_t = decltype(ImplTy::hashKey(std::declval<T>()));
|
|||
/// 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();
|
||||
|
@ -83,6 +89,10 @@ public:
|
|||
/// 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 {
|
||||
|
@ -140,8 +150,8 @@ public:
|
|||
/// 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)...);
|
||||
|
@ -163,7 +173,8 @@ public:
|
|||
};
|
||||
|
||||
// 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
|
||||
|
@ -171,31 +182,32 @@ public:
|
|||
/// 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)...);
|
||||
|
@ -209,7 +221,7 @@ public:
|
|||
};
|
||||
|
||||
// Attempt to erase the storage instance.
|
||||
eraseImpl(kind, hashValue, isEqual, [](BaseStorage *storage) {
|
||||
eraseImpl(id, kind, hashValue, isEqual, [](BaseStorage *storage) {
|
||||
static_cast<Storage *>(storage)->cleanup();
|
||||
});
|
||||
}
|
||||
|
@ -217,24 +229,25 @@ public:
|
|||
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;
|
||||
|
|
|
@ -7,7 +7,17 @@
|
|||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#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;
|
||||
|
|
|
@ -246,6 +246,7 @@ SDBMSumExpr SDBMSumExpr::get(SDBMTermExpr lhs, SDBMConstantExpr rhs) {
|
|||
|
||||
StorageUniquer &uniquer = lhs.getDialect()->getUniquer();
|
||||
return uniquer.get<detail::SDBMBinaryExprStorage>(
|
||||
TypeID::get<detail::SDBMBinaryExprStorage>(),
|
||||
/*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Add), lhs, rhs);
|
||||
}
|
||||
|
||||
|
@ -533,6 +534,7 @@ SDBMDiffExpr SDBMDiffExpr::get(SDBMDirectExpr lhs, SDBMTermExpr rhs) {
|
|||
|
||||
StorageUniquer &uniquer = lhs.getDialect()->getUniquer();
|
||||
return uniquer.get<detail::SDBMDiffExprStorage>(
|
||||
TypeID::get<detail::SDBMDiffExprStorage>(),
|
||||
/*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Diff), lhs, rhs);
|
||||
}
|
||||
|
||||
|
@ -573,6 +575,7 @@ SDBMStripeExpr SDBMStripeExpr::get(SDBMDirectExpr var,
|
|||
|
||||
StorageUniquer &uniquer = var.getDialect()->getUniquer();
|
||||
return uniquer.get<detail::SDBMBinaryExprStorage>(
|
||||
TypeID::get<detail::SDBMBinaryExprStorage>(),
|
||||
/*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Stripe), var,
|
||||
stripeFactor);
|
||||
}
|
||||
|
@ -608,7 +611,8 @@ SDBMDimExpr SDBMDimExpr::get(SDBMDialect *dialect, unsigned position) {
|
|||
|
||||
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);
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
@ -624,7 +628,8 @@ SDBMSymbolExpr SDBMSymbolExpr::get(SDBMDialect *dialect, unsigned position) {
|
|||
|
||||
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);
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
@ -640,7 +645,8 @@ SDBMConstantExpr SDBMConstantExpr::get(SDBMDialect *dialect, int64_t value) {
|
|||
|
||||
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 {
|
||||
|
@ -656,6 +662,7 @@ SDBMNegExpr SDBMNegExpr::get(SDBMDirectExpr var) {
|
|||
|
||||
StorageUniquer &uniquer = var.getDialect()->getUniquer();
|
||||
return uniquer.get<detail::SDBMNegExprStorage>(
|
||||
TypeID::get<detail::SDBMNegExprStorage>(),
|
||||
/*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Neg), var);
|
||||
}
|
||||
|
||||
|
|
|
@ -12,6 +12,7 @@
|
|||
#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;
|
||||
|
@ -448,7 +449,8 @@ static AffineExpr getAffineDimOrSymbol(AffineExprKind kind, unsigned position,
|
|||
|
||||
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) {
|
||||
|
@ -483,7 +485,8 @@ AffineExpr mlir::getAffineConstantExpr(int64_t constant, MLIRContext *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.
|
||||
|
@ -591,6 +594,7 @@ AffineExpr AffineExpr::operator+(AffineExpr other) const {
|
|||
|
||||
StorageUniquer &uniquer = getContext()->getAffineUniquer();
|
||||
return uniquer.get<AffineBinaryOpExprStorage>(
|
||||
TypeID::get<AffineBinaryOpExprStorage>(),
|
||||
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Add), *this, other);
|
||||
}
|
||||
|
||||
|
@ -651,6 +655,7 @@ AffineExpr AffineExpr::operator*(AffineExpr other) const {
|
|||
|
||||
StorageUniquer &uniquer = getContext()->getAffineUniquer();
|
||||
return uniquer.get<AffineBinaryOpExprStorage>(
|
||||
TypeID::get<AffineBinaryOpExprStorage>(),
|
||||
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mul), *this, other);
|
||||
}
|
||||
|
||||
|
@ -717,6 +722,7 @@ AffineExpr AffineExpr::floorDiv(AffineExpr other) const {
|
|||
|
||||
StorageUniquer &uniquer = getContext()->getAffineUniquer();
|
||||
return uniquer.get<AffineBinaryOpExprStorage>(
|
||||
TypeID::get<AffineBinaryOpExprStorage>(),
|
||||
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::FloorDiv), *this,
|
||||
other);
|
||||
}
|
||||
|
@ -760,6 +766,7 @@ AffineExpr AffineExpr::ceilDiv(AffineExpr other) const {
|
|||
|
||||
StorageUniquer &uniquer = getContext()->getAffineUniquer();
|
||||
return uniquer.get<AffineBinaryOpExprStorage>(
|
||||
TypeID::get<AffineBinaryOpExprStorage>(),
|
||||
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::CeilDiv), *this,
|
||||
other);
|
||||
}
|
||||
|
@ -807,6 +814,7 @@ AffineExpr AffineExpr::operator%(AffineExpr other) const {
|
|||
|
||||
StorageUniquer &uniquer = getContext()->getAffineUniquer();
|
||||
return uniquer.get<AffineBinaryOpExprStorage>(
|
||||
TypeID::get<AffineBinaryOpExprStorage>(),
|
||||
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mod), *this, other);
|
||||
}
|
||||
|
||||
|
|
|
@ -402,6 +402,13 @@ MLIRContext::MLIRContext() : impl(new MLIRContextImpl()) {
|
|||
/// 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() {}
|
||||
|
@ -571,6 +578,7 @@ void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
|
|||
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) {
|
||||
|
@ -580,6 +588,7 @@ void Dialect::addAttribute(TypeID typeID, AbstractAttribute &&attrInfo) {
|
|||
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.
|
||||
|
|
|
@ -9,15 +9,18 @@
|
|||
#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;
|
||||
|
||||
|
@ -40,113 +43,6 @@ struct StorageUniquerImpl {
|
|||
BaseStorage *storage;
|
||||
};
|
||||
|
||||
/// Get or create an instance of a complex derived type.
|
||||
BaseStorage *
|
||||
getOrCreate(unsigned kind, unsigned hashValue,
|
||||
function_ref<bool(const BaseStorage *)> isEqual,
|
||||
function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
|
||||
LookupKey lookupKey{kind, hashValue, isEqual};
|
||||
if (!threadingIsEnabled)
|
||||
return getOrCreateUnsafe(kind, hashValue, 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())
|
||||
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);
|
||||
}
|
||||
/// Get or create an instance of a complex derived type in an unsafe fashion.
|
||||
BaseStorage *
|
||||
getOrCreateUnsafe(unsigned kind, unsigned hashValue, LookupKey &lookupKey,
|
||||
function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
|
||||
auto existing = storageTypes.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};
|
||||
return storage;
|
||||
}
|
||||
|
||||
/// Get or create an instance of a simple derived type.
|
||||
BaseStorage *
|
||||
getOrCreate(unsigned kind,
|
||||
function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
|
||||
if (!threadingIsEnabled)
|
||||
return getOrCreateUnsafe(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())
|
||||
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);
|
||||
}
|
||||
/// Get or create an instance of a simple derived type in an unsafe fashion.
|
||||
BaseStorage *
|
||||
getOrCreateUnsafe(unsigned kind,
|
||||
function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
|
||||
auto &result = simpleTypes[kind];
|
||||
if (result)
|
||||
return result;
|
||||
|
||||
// Otherwise, create and return a new storage instance.
|
||||
return result = initializeStorage(kind, ctorFn);
|
||||
}
|
||||
|
||||
/// Erase an instance of a complex derived type.
|
||||
void erase(unsigned kind, unsigned hashValue,
|
||||
function_ref<bool(const BaseStorage *)> isEqual,
|
||||
function_ref<void(BaseStorage *)> cleanupFn) {
|
||||
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())
|
||||
return;
|
||||
|
||||
// Cleanup the storage and remove it from the map.
|
||||
cleanupFn(existing->storage);
|
||||
storageTypes.erase(existing);
|
||||
}
|
||||
|
||||
/// Mutates an instance of a derived storage in a thread-safe way.
|
||||
LogicalResult
|
||||
mutate(function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
|
||||
if (!threadingIsEnabled)
|
||||
return mutationFn(allocator);
|
||||
|
||||
llvm::sys::SmartScopedWriter<true> lock(mutex);
|
||||
return mutationFn(allocator);
|
||||
}
|
||||
|
||||
//===--------------------------------------------------------------------===//
|
||||
// Instance Storage
|
||||
//===--------------------------------------------------------------------===//
|
||||
|
||||
/// Utility to create and initialize a storage instance.
|
||||
BaseStorage *
|
||||
initializeStorage(unsigned kind,
|
||||
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() {
|
||||
|
@ -175,16 +71,149 @@ struct StorageUniquerImpl {
|
|||
|
||||
/// Unique types with specific hashing or storage constraints.
|
||||
using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
|
||||
StorageTypeSet storageTypes;
|
||||
StorageTypeSet complexInstances;
|
||||
|
||||
/// Unique types with just the kind.
|
||||
DenseMap<unsigned, BaseStorage *> simpleTypes;
|
||||
/// Instances of this storage object.
|
||||
llvm::SmallDenseMap<unsigned, BaseStorage *, 1> simpleInstances;
|
||||
|
||||
/// Allocator to use when constructing derived type instances.
|
||||
StorageUniquer::StorageAllocator allocator;
|
||||
/// 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(TypeID id, unsigned kind, unsigned hashValue,
|
||||
function_ref<bool(const BaseStorage *)> isEqual,
|
||||
function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
|
||||
assert(instUniquers.count(id) && "creating unregistered storage instance");
|
||||
InstSpecificUniquer::LookupKey lookupKey{kind, hashValue, isEqual};
|
||||
InstSpecificUniquer &storageUniquer = *instUniquers[id];
|
||||
if (!threadingIsEnabled)
|
||||
return getOrCreateUnsafe(storageUniquer, kind, lookupKey, ctorFn);
|
||||
|
||||
// Check for an existing instance in read-only mode.
|
||||
{
|
||||
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(storageUniquer.mutex);
|
||||
return getOrCreateUnsafe(storageUniquer, kind, lookupKey, ctorFn);
|
||||
}
|
||||
/// Get or create an instance of a complex derived type in an thread-unsafe
|
||||
/// fashion.
|
||||
BaseStorage *
|
||||
getOrCreateUnsafe(InstSpecificUniquer &storageUniquer, unsigned kind,
|
||||
InstSpecificUniquer::LookupKey &lookupKey,
|
||||
function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
|
||||
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, storageUniquer.allocator, ctorFn);
|
||||
*existing.first =
|
||||
InstSpecificUniquer::HashedStorage{lookupKey.hashValue, storage};
|
||||
return storage;
|
||||
}
|
||||
|
||||
/// Get or create an instance of a simple derived type.
|
||||
BaseStorage *
|
||||
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(storageUniquer, kind, ctorFn);
|
||||
|
||||
// Check for an existing instance in read-only mode.
|
||||
{
|
||||
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(storageUniquer.mutex);
|
||||
return getOrCreateUnsafe(storageUniquer, kind, ctorFn);
|
||||
}
|
||||
/// Get or create an instance of a simple derived type in an thread-unsafe
|
||||
/// fashion.
|
||||
BaseStorage *
|
||||
getOrCreateUnsafe(InstSpecificUniquer &storageUniquer, unsigned kind,
|
||||
function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
|
||||
auto &result = storageUniquer.simpleInstances[kind];
|
||||
if (result)
|
||||
return result;
|
||||
|
||||
// Otherwise, create and return a new storage instance.
|
||||
return result = initializeStorage(kind, storageUniquer.allocator, ctorFn);
|
||||
}
|
||||
|
||||
/// Erase an instance of a complex derived type.
|
||||
void erase(TypeID id, unsigned kind, unsigned hashValue,
|
||||
function_ref<bool(const BaseStorage *)> isEqual,
|
||||
function_ref<void(BaseStorage *)> cleanupFn) {
|
||||
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> 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);
|
||||
storageUniquer.complexInstances.erase(existing);
|
||||
}
|
||||
|
||||
/// Mutates an instance of a derived storage in a thread-safe way.
|
||||
LogicalResult
|
||||
mutate(TypeID id,
|
||||
function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
|
||||
assert(instUniquers.count(id) && "mutating unregistered storage instance");
|
||||
InstSpecificUniquer &storageUniquer = *instUniquers[id];
|
||||
if (!threadingIsEnabled)
|
||||
return mutationFn(storageUniquer.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, StorageAllocator &allocator,
|
||||
function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
|
||||
BaseStorage *storage = ctorFn(allocator);
|
||||
storage->kind = kind;
|
||||
return storage;
|
||||
}
|
||||
|
||||
/// 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;
|
||||
|
@ -200,33 +229,41 @@ 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);
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue