Move the storage of uniqued TypeStorage objects into TypeUniquer and give each context a unique TypeUniquer instance.

PiperOrigin-RevId: 229460053
This commit is contained in:
River Riddle 2019-01-15 16:10:34 -08:00 committed by jpienaar
parent 03e15e1b9f
commit 18fe1ffcd7
4 changed files with 100 additions and 108 deletions

View File

@ -29,6 +29,10 @@ class MLIRContextImpl;
class Location;
class Dialect;
namespace detail {
class TypeUniquer;
}
/// MLIRContext is the top-level object for a collection of MLIR modules. It
/// holds immortal uniqued objects like types, and the tables used to unique
/// them.
@ -93,6 +97,9 @@ public:
// MLIRContextImpl type.
MLIRContextImpl &getImpl() const { return *impl.get(); }
/// Get the type uniquer for this context.
detail::TypeUniquer &getTypeUniquer() const;
private:
const std::unique_ptr<MLIRContextImpl> impl;

View File

@ -22,8 +22,10 @@
#ifndef MLIR_IR_TYPE_SUPPORT_H
#define MLIR_IR_TYPE_SUPPORT_H
#include "mlir/IR/MLIRContext.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseMapInfo.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/StringRef.h"
#include <memory>
@ -141,7 +143,21 @@ namespace detail {
// MLIRContext. This class manages all creation and uniquing of types.
class TypeUniquer {
public:
/// Lookup key for storage types.
template <typename T, typename... Args>
static T get(MLIRContext *ctx, Args &&... args) {
TypeUniquer &instance = ctx->getTypeUniquer();
return instance.getImpl<T>(ctx, args...);
}
private:
/// A utility wrapper object representing a hashed storage object. This class
/// contains a storage object and an existing computed hash value.
struct HashedStorageType {
unsigned hashValue;
TypeStorage *storage;
};
/// A lookup key for derived instances of TypeStorage objects.
struct TypeLookupKey {
/// The known derived kind for the storage.
unsigned kind;
@ -153,14 +169,12 @@ public:
llvm::function_ref<bool(const TypeStorage *)> isEqual;
};
TypeUniquer(MLIRContext *ctx) : ctx(ctx) {}
/// Get an uniqued instance of a type T. This overload is used for derived
/// types that have complex storage or uniquing constraints.
template <typename T, typename... Args>
typename std::enable_if<
!std::is_same<typename T::ImplType, DefaultTypeStorage>::value, T>::type
get(unsigned kind, Args &&... args) {
getImpl(MLIRContext *ctx, unsigned kind, Args &&... args) {
using ImplType = typename T::ImplType;
using KeyTy = typename ImplType::KeyTy;
@ -177,22 +191,20 @@ public:
return static_cast<const ImplType &>(*existing) == derivedKey;
};
// Lookup an existing type with the given key.
TypeStorage *storage = lookup(TypeLookupKey{kind, hashValue, isEqual});
if (storage)
return T(storage);
// Look to see if the type has been created already.
auto existing =
storageTypes.insert_as({}, TypeLookupKey{kind, hashValue, isEqual});
// Get the dialect this type was registered to.
auto &dialect = lookupDialectForType<T>();
// If it has been created, return it.
if (!existing.second)
return T(existing.first->storage);
// Otherwise, construct and initialize the derived storage for this type
// instance.
TypeStorageAllocator allocator(ctx);
storage = ImplType::construct(allocator, derivedKey);
storage->initializeTypeInfo(dialect, kind);
// Insert the new type storage instance into the context.
insert(hashValue, storage);
TypeStorage *storage = ImplType::construct(allocator, derivedKey);
storage->initializeTypeInfo(lookupDialectForType<T>(ctx), kind);
*existing.first = HashedStorageType{hashValue, storage};
return T(storage);
}
@ -202,23 +214,31 @@ public:
template <typename T, typename... Args>
typename std::enable_if<
std::is_same<typename T::ImplType, DefaultTypeStorage>::value, T>::type
get(unsigned kind) {
auto &dialect = lookupDialectForType<T>();
return T(getSimple(dialect, kind));
getImpl(MLIRContext *ctx, unsigned kind) {
// Check for an existing instance with this kind.
auto *&result = simpleTypes[kind];
if (!result) {
// Otherwise, allocate and initialize one.
TypeStorageAllocator allocator(ctx);
result = new (allocator.allocate<DefaultTypeStorage>())
DefaultTypeStorage(lookupDialectForType<T>(ctx), kind);
}
return T(result);
}
private:
/// Get the dialect that the type 'T' was registered with.
template <typename T> const Dialect &lookupDialectForType() {
return lookupDialectForType(&T::typeID);
template <typename T>
static const Dialect &lookupDialectForType(MLIRContext *ctx) {
return lookupDialectForType(ctx, &T::typeID);
}
/// Get the dialect that registered the type with the provided typeid.
const Dialect &lookupDialectForType(const void *const typeID);
static const Dialect &lookupDialectForType(MLIRContext *ctx,
const void *const typeID);
/// Get or create a uniqued type by its kind. This overload is used for
/// simple types that are only uniqued by kind.
TypeStorage *getSimple(const Dialect &dialect, unsigned kind);
//===--------------------------------------------------------------------===//
// Key Construction
//===--------------------------------------------------------------------===//
/// Utilities for generating a derived storage key.
/// Overload for if the key can be directly constructed from the provided
@ -240,15 +260,45 @@ private:
return ImplTy::getKey(args...);
}
/// Look up a uniqued type with a lookup key. This is used if the type defines
/// a storage key.
TypeStorage *lookup(const TypeLookupKey &key);
//===--------------------------------------------------------------------===//
// Instance Storage
//===--------------------------------------------------------------------===//
/// Insert a new type storage into the context.
void insert(unsigned hashValue, TypeStorage *storage);
/// Storage info for derived TypeStorage objects.
struct StorageKeyInfo : DenseMapInfo<HashedStorageType> {
static HashedStorageType getEmptyKey() {
return HashedStorageType{0, DenseMapInfo<TypeStorage *>::getEmptyKey()};
}
static HashedStorageType getTombstoneKey() {
return HashedStorageType{0,
DenseMapInfo<TypeStorage *>::getTombstoneKey()};
}
/// The current context that a type is being uniqued from.
MLIRContext *ctx;
static unsigned getHashValue(const HashedStorageType &key) {
return key.hashValue;
}
static unsigned getHashValue(TypeLookupKey key) { return key.hashValue; }
static bool isEqual(const HashedStorageType &lhs,
const HashedStorageType &rhs) {
return lhs.storage == rhs.storage;
}
static bool isEqual(const TypeLookupKey &lhs,
const HashedStorageType &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 = llvm::DenseSet<HashedStorageType, StorageKeyInfo>;
StorageTypeSet storageTypes;
// Unique types with just the kind.
DenseMap<unsigned, TypeStorage *> simpleTypes;
};
} // namespace detail

View File

@ -147,7 +147,7 @@ public:
// Ensure that the invariants are correct for type construction.
assert(!ConcreteType::verifyConstructionInvariants(llvm::None, context,
args...));
return detail::TypeUniquer(context).get<ConcreteType>(kind, args...);
return detail::TypeUniquer::get<ConcreteType>(context, kind, args...);
}
/// Get or create a new ConcreteType instance within the context, defined at
@ -159,7 +159,7 @@ public:
// If the construction invariants fail then we return a null type.
if (ConcreteType::verifyConstructionInvariants(loc, context, args...))
return ConcreteType();
return detail::TypeUniquer(context).get<ConcreteType>(kind, args...);
return detail::TypeUniquer::get<ConcreteType>(context, kind, args...);
}
/// Default implementation that just returns false for success.

View File

@ -46,43 +46,6 @@ using namespace mlir::detail;
using namespace llvm;
namespace {
/// A utility wrapper object representing a hased storage type. This class
/// contains a type storage object and an existing computed hash value.
struct HashedStorageType {
unsigned hashValue;
TypeStorage *storage;
};
/// Storage info for storage types.
struct StorageTypeKeyInfo : DenseMapInfo<HashedStorageType> {
/// Storage types use the lookup key with the type uniquer.
using KeyTy = TypeUniquer::TypeLookupKey;
static HashedStorageType getEmptyKey() {
return HashedStorageType{0, DenseMapInfo<TypeStorage *>::getEmptyKey()};
}
static HashedStorageType getTombstoneKey() {
return HashedStorageType{0, DenseMapInfo<TypeStorage *>::getTombstoneKey()};
}
static unsigned getHashValue(const HashedStorageType &key) {
return key.hashValue;
}
static unsigned getHashValue(KeyTy key) { return key.hashValue; }
static bool isEqual(const HashedStorageType &lhs,
const HashedStorageType &rhs) {
return lhs.storage == rhs.storage;
}
static bool isEqual(const KeyTy &lhs, const HashedStorageType &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);
}
};
struct AffineMapKeyInfo : DenseMapInfo<AffineMap> {
// Affine maps are uniqued based on their dim/symbol counts and affine
// expressions.
@ -380,12 +343,7 @@ public:
DenseMap<int64_t, AffineConstantExprStorage *> constExprs;
/// Type uniquing.
// Unique types with specific hashing or storage constraints.
using StorageTypeSet = DenseSet<HashedStorageType, StorageTypeKeyInfo>;
StorageTypeSet storageTypes;
// Unique types with just the kind.
DenseMap<unsigned, TypeStorage *> simpleTypes;
TypeUniquer typeUniquer;
// Attribute uniquing.
BoolAttributeStorage *boolAttrs[2] = {nullptr};
@ -722,47 +680,24 @@ Location FusedLoc::get(ArrayRef<Location> locs, Attribute metadata,
// Type uniquing
//===----------------------------------------------------------------------===//
/// Get the type uniquer for this context.
TypeUniquer &MLIRContext::getTypeUniquer() const {
return getImpl().typeUniquer;
}
/// Get a reference to the internal allocator.
llvm::BumpPtrAllocator &TypeStorageAllocator::getAllocator() {
return ctx->getImpl().allocator;
}
/// Get or create a uniqued type by it's kind. This overload is used for
/// simple types that are only uniqued by kind.
TypeStorage *TypeUniquer::getSimple(const Dialect &dialect, unsigned kind) {
auto &impl = ctx->getImpl();
// Check for an existing instance with this kind.
auto *&result = impl.simpleTypes[kind];
if (result)
return result;
// Otherwise, create a new instance and return it.
result = impl.allocator.Allocate<DefaultTypeStorage>();
return new (result) DefaultTypeStorage{dialect, kind};
}
/// Get the dialect that registered the type with the provided typeid.
const Dialect &TypeUniquer::lookupDialectForType(const void *const typeID) {
const Dialect &TypeUniquer::lookupDialectForType(MLIRContext *ctx,
const void *const typeID) {
auto &impl = ctx->getImpl();
assert(impl.registeredTypes.count(typeID) && "typeID is not registered.");
return *impl.registeredTypes[typeID];
}
/// Look up a uniqued type with a lookup key. This is used if the type defines
/// a storage key.
TypeStorage *TypeUniquer::lookup(const TypeLookupKey &key) {
auto &impl = ctx->getImpl();
auto existing = impl.storageTypes.find_as(key);
return existing != impl.storageTypes.end() ? existing->storage : nullptr;
}
/// Insert a pre hashed storage type into the context.
void TypeUniquer::insert(unsigned hashValue, TypeStorage *storage) {
ctx->getImpl().storageTypes.insert(HashedStorageType{hashValue, storage});
}
//===----------------------------------------------------------------------===//
// Attribute uniquing
//===----------------------------------------------------------------------===//