|
|
@@ -0,0 +1,215 @@ |
|
|
//===- DialectResourceBlobManager.h - Dialect Blob Management ---*- C++ -*-===// |
|
|
// |
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
|
|
// See https://llvm.org/LICENSE.txt for license information. |
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
|
|
// |
|
|
//===----------------------------------------------------------------------===// |
|
|
// |
|
|
// This file defines utility classes for referencing and managing asm resource |
|
|
// blobs. These classes are intended to more easily facilitate the sharing of |
|
|
// large blobs, and their definition. |
|
|
// |
|
|
//===----------------------------------------------------------------------===// |
|
|
|
|
|
#ifndef MLIR_IR_DIALECTRESOURCEBLOBMANAGER_H |
|
|
#define MLIR_IR_DIALECTRESOURCEBLOBMANAGER_H |
|
|
|
|
|
#include "mlir/IR/AsmState.h" |
|
|
#include "mlir/IR/OpImplementation.h" |
|
|
#include "llvm/ADT/StringMap.h" |
|
|
#include "llvm/ADT/Twine.h" |
|
|
#include "llvm/Support/RWMutex.h" |
|
|
#include "llvm/Support/SMLoc.h" |
|
|
|
|
|
namespace mlir { |
|
|
//===----------------------------------------------------------------------===// |
|
|
// DialectResourceBlobManager |
|
|
//===---------------------------------------------------------------------===// |
|
|
|
|
|
/// This class defines a manager for dialect resource blobs. Blobs are uniqued |
|
|
/// by a given key, and represented using AsmResourceBlobs. |
|
|
class DialectResourceBlobManager { |
|
|
public: |
|
|
/// The class represents an individual entry of a blob. |
|
|
class BlobEntry { |
|
|
public: |
|
|
/// Return the key used to reference this blob. |
|
|
StringRef getKey() const { return key; } |
|
|
|
|
|
/// Return the blob owned by this entry if one has been initialized. Returns |
|
|
/// nullptr otherwise. |
|
|
const AsmResourceBlob *getBlob() const { return blob ? &*blob : nullptr; } |
|
|
AsmResourceBlob *getBlob() { return blob ? &*blob : nullptr; } |
|
|
|
|
|
/// Set the blob owned by this entry. |
|
|
void setBlob(AsmResourceBlob &&newBlob) { blob = std::move(newBlob); } |
|
|
|
|
|
private: |
|
|
BlobEntry() = default; |
|
|
BlobEntry(BlobEntry &&) = default; |
|
|
BlobEntry &operator=(const BlobEntry &) = delete; |
|
|
BlobEntry &operator=(BlobEntry &&) = delete; |
|
|
|
|
|
/// Initialize this entry with the given key and blob. |
|
|
void initialize(StringRef newKey, Optional<AsmResourceBlob> newBlob) { |
|
|
key = newKey; |
|
|
blob = std::move(newBlob); |
|
|
} |
|
|
|
|
|
/// The key used for this blob. |
|
|
StringRef key; |
|
|
|
|
|
/// The blob that is referenced by this entry if it is valid. |
|
|
Optional<AsmResourceBlob> blob; |
|
|
|
|
|
/// Allow access to the constructors. |
|
|
friend DialectResourceBlobManager; |
|
|
friend class llvm::StringMapEntryStorage<BlobEntry>; |
|
|
}; |
|
|
|
|
|
/// Return the blob registered for the given name, or nullptr if no blob |
|
|
/// is registered. |
|
|
BlobEntry *lookup(StringRef name); |
|
|
const BlobEntry *lookup(StringRef name) const { |
|
|
return const_cast<DialectResourceBlobManager *>(this)->lookup(name); |
|
|
} |
|
|
|
|
|
/// Update the blob for the entry defined by the provided name. This method |
|
|
/// asserts that an entry for the given name exists in the manager. |
|
|
void update(StringRef name, AsmResourceBlob &&newBlob); |
|
|
|
|
|
/// Insert a new entry with the provided name and optional blob data. The name |
|
|
/// may be modified during insertion if another entry already exists with that |
|
|
/// name. Returns the inserted entry. |
|
|
BlobEntry &insert(StringRef name, Optional<AsmResourceBlob> blob = {}); |
|
|
/// Insertion method that returns a dialect specific handle to the inserted |
|
|
/// entry. |
|
|
template <typename HandleT> |
|
|
HandleT insert(typename HandleT::Dialect *dialect, StringRef name, |
|
|
Optional<AsmResourceBlob> blob = {}) { |
|
|
BlobEntry &entry = insert(name, std::move(blob)); |
|
|
return HandleT(&entry, dialect); |
|
|
} |
|
|
|
|
|
private: |
|
|
/// A mutex to protect access to the blob map. |
|
|
llvm::sys::SmartRWMutex<true> blobMapLock; |
|
|
|
|
|
/// The internal map of tracked blobs. StringMap stores entries in distinct |
|
|
/// allocations, so we can freely take references to the data without fear of |
|
|
/// invalidation during additional insertion/deletion. |
|
|
llvm::StringMap<BlobEntry> blobMap; |
|
|
}; |
|
|
|
|
|
//===----------------------------------------------------------------------===// |
|
|
// ResourceBlobManagerDialectInterface |
|
|
//===---------------------------------------------------------------------===// |
|
|
|
|
|
/// This class implements a dialect interface that provides common functionality |
|
|
/// for interacting with a resource blob manager. |
|
|
class ResourceBlobManagerDialectInterface |
|
|
: public DialectInterface::Base<ResourceBlobManagerDialectInterface> { |
|
|
public: |
|
|
ResourceBlobManagerDialectInterface(Dialect *dialect) |
|
|
: Base(dialect), |
|
|
blobManager(std::make_shared<DialectResourceBlobManager>()) {} |
|
|
|
|
|
/// Return the blob manager held by this interface. |
|
|
DialectResourceBlobManager &getBlobManager() { return *blobManager; } |
|
|
const DialectResourceBlobManager &getBlobManager() const { |
|
|
return *blobManager; |
|
|
} |
|
|
|
|
|
/// Set the blob manager held by this interface. |
|
|
void |
|
|
setBlobManager(std::shared_ptr<DialectResourceBlobManager> newBlobManager) { |
|
|
blobManager = std::move(newBlobManager); |
|
|
} |
|
|
|
|
|
private: |
|
|
/// The blob manager owned by the dialect implementing this interface. |
|
|
std::shared_ptr<DialectResourceBlobManager> blobManager; |
|
|
}; |
|
|
|
|
|
/// This class provides a base class for dialects implementing the resource blob |
|
|
/// interface. It provides several additional dialect specific utilities on top |
|
|
/// of the generic interface. `HandleT` is the type of the handle used to |
|
|
/// reference a resource blob. |
|
|
template <typename HandleT> |
|
|
class ResourceBlobManagerDialectInterfaceBase |
|
|
: public ResourceBlobManagerDialectInterface { |
|
|
public: |
|
|
using ResourceBlobManagerDialectInterface:: |
|
|
ResourceBlobManagerDialectInterface; |
|
|
|
|
|
/// Update the blob for the entry defined by the provided name. This method |
|
|
/// asserts that an entry for the given name exists in the manager. |
|
|
void update(StringRef name, AsmResourceBlob &&newBlob) { |
|
|
getBlobManager().update(name, std::move(newBlob)); |
|
|
} |
|
|
|
|
|
/// Insert a new resource blob entry with the provided name and optional blob |
|
|
/// data. The name may be modified during insertion if another entry already |
|
|
/// exists with that name. Returns a dialect specific handle to the inserted |
|
|
/// entry. |
|
|
HandleT insert(StringRef name, Optional<AsmResourceBlob> blob = {}) { |
|
|
return getBlobManager().template insert<HandleT>( |
|
|
cast<typename HandleT::Dialect>(getDialect()), name, std::move(blob)); |
|
|
} |
|
|
|
|
|
/// Build resources for each of the referenced blobs within this manager. |
|
|
void buildResources(AsmResourceBuilder &provider, |
|
|
ArrayRef<AsmDialectResourceHandle> referencedResources) { |
|
|
for (const AsmDialectResourceHandle &handle : referencedResources) { |
|
|
if (const auto *dialectHandle = dyn_cast<HandleT>(&handle)) { |
|
|
if (auto *blob = dialectHandle->getBlob()) |
|
|
provider.buildBlob(dialectHandle->getKey(), *blob); |
|
|
} |
|
|
} |
|
|
} |
|
|
}; |
|
|
|
|
|
//===----------------------------------------------------------------------===// |
|
|
// DialectResourceBlobHandle |
|
|
//===----------------------------------------------------------------------===// |
|
|
|
|
|
/// This class defines a dialect specific handle to a resource blob. These |
|
|
/// handles utilize a StringRef for the internal key, and an AsmResourceBlob as |
|
|
/// the underlying data. |
|
|
template <typename DialectT> |
|
|
struct DialectResourceBlobHandle |
|
|
: public AsmDialectResourceHandleBase<DialectResourceBlobHandle<DialectT>, |
|
|
DialectResourceBlobManager::BlobEntry, |
|
|
DialectT> { |
|
|
using AsmDialectResourceHandleBase<DialectResourceBlobHandle<DialectT>, |
|
|
DialectResourceBlobManager::BlobEntry, |
|
|
DialectT>::AsmDialectResourceHandleBase; |
|
|
using ManagerInterface = ResourceBlobManagerDialectInterfaceBase< |
|
|
DialectResourceBlobHandle<DialectT>>; |
|
|
|
|
|
/// Return the human readable string key for this handle. |
|
|
StringRef getKey() const { return this->getResource()->getKey(); } |
|
|
|
|
|
/// Return the blob referenced by this handle if the underlying resource has |
|
|
/// been initialized. Returns nullptr otherwise. |
|
|
AsmResourceBlob *getBlob() { return this->getResource()->getBlob(); } |
|
|
const AsmResourceBlob *getBlob() const { |
|
|
return this->getResource()->getBlob(); |
|
|
} |
|
|
|
|
|
/// Get the interface for the dialect that owns handles of this type. Asserts |
|
|
/// that the dialect is registered. |
|
|
static ManagerInterface &getManagerInterface(MLIRContext *ctx) { |
|
|
auto *dialect = ctx->getOrLoadDialect<DialectT>(); |
|
|
assert(dialect && "dialect not registered"); |
|
|
|
|
|
auto *iface = dialect->template getRegisteredInterface<ManagerInterface>(); |
|
|
assert(iface && "dialect doesn't provide the blob manager interface?"); |
|
|
return *iface; |
|
|
} |
|
|
}; |
|
|
|
|
|
} // namespace mlir |
|
|
|
|
|
#endif // MLIR_IR_DIALECTRESOURCEBLOBMANAGER_H |