1//===- DialectResourceBlobManager.h - Dialect Blob Management ---*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file defines utility classes for referencing and managing asm resource
10// blobs. These classes are intended to more easily facilitate the sharing of
11// large blobs, and their definition.
12//
13//===----------------------------------------------------------------------===//
14
15#ifndef MLIR_IR_DIALECTRESOURCEBLOBMANAGER_H
16#define MLIR_IR_DIALECTRESOURCEBLOBMANAGER_H
17
18#include "mlir/IR/AsmState.h"
19#include "mlir/IR/OpImplementation.h"
20#include "llvm/ADT/StringMap.h"
21#include "llvm/ADT/Twine.h"
22#include "llvm/Support/RWMutex.h"
23#include "llvm/Support/SMLoc.h"
24#include <optional>
25
26namespace mlir {
27//===----------------------------------------------------------------------===//
28// DialectResourceBlobManager
29//===---------------------------------------------------------------------===//
30
31/// This class defines a manager for dialect resource blobs. Blobs are uniqued
32/// by a given key, and represented using AsmResourceBlobs.
33class DialectResourceBlobManager {
34public:
35 /// The class represents an individual entry of a blob.
36 class BlobEntry {
37 public:
38 /// Return the key used to reference this blob.
39 StringRef getKey() const { return key; }
40
41 /// Return the blob owned by this entry if one has been initialized. Returns
42 /// nullptr otherwise.
43 const AsmResourceBlob *getBlob() const { return blob ? &*blob : nullptr; }
44 AsmResourceBlob *getBlob() { return blob ? &*blob : nullptr; }
45
46 /// Set the blob owned by this entry.
47 void setBlob(AsmResourceBlob &&newBlob) { blob = std::move(newBlob); }
48
49 private:
50 BlobEntry() = default;
51 BlobEntry(BlobEntry &&) = default;
52 BlobEntry &operator=(const BlobEntry &) = delete;
53 BlobEntry &operator=(BlobEntry &&) = delete;
54
55 /// Initialize this entry with the given key and blob.
56 void initialize(StringRef newKey, std::optional<AsmResourceBlob> newBlob) {
57 key = newKey;
58 blob = std::move(newBlob);
59 }
60
61 /// The key used for this blob.
62 StringRef key;
63
64 /// The blob that is referenced by this entry if it is valid.
65 std::optional<AsmResourceBlob> blob;
66
67 /// Allow access to the constructors.
68 friend DialectResourceBlobManager;
69 friend class llvm::StringMapEntryStorage<BlobEntry>;
70 };
71
72 /// Return the blob registered for the given name, or nullptr if no blob
73 /// is registered.
74 BlobEntry *lookup(StringRef name);
75 const BlobEntry *lookup(StringRef name) const {
76 return const_cast<DialectResourceBlobManager *>(this)->lookup(name);
77 }
78
79 /// Update the blob for the entry defined by the provided name. This method
80 /// asserts that an entry for the given name exists in the manager.
81 void update(StringRef name, AsmResourceBlob &&newBlob);
82
83 /// Insert a new entry with the provided name and optional blob data. The name
84 /// may be modified during insertion if another entry already exists with that
85 /// name. Returns the inserted entry.
86 BlobEntry &insert(StringRef name, std::optional<AsmResourceBlob> blob = {});
87 /// Insertion method that returns a dialect specific handle to the inserted
88 /// entry.
89 template <typename HandleT>
90 HandleT insert(typename HandleT::Dialect *dialect, StringRef name,
91 std::optional<AsmResourceBlob> blob = {}) {
92 BlobEntry &entry = insert(name, blob: std::move(blob));
93 return HandleT(&entry, dialect);
94 }
95
96private:
97 /// A mutex to protect access to the blob map.
98 llvm::sys::SmartRWMutex<true> blobMapLock;
99
100 /// The internal map of tracked blobs. StringMap stores entries in distinct
101 /// allocations, so we can freely take references to the data without fear of
102 /// invalidation during additional insertion/deletion.
103 llvm::StringMap<BlobEntry> blobMap;
104};
105
106//===----------------------------------------------------------------------===//
107// ResourceBlobManagerDialectInterface
108//===---------------------------------------------------------------------===//
109
110/// This class implements a dialect interface that provides common functionality
111/// for interacting with a resource blob manager.
112class ResourceBlobManagerDialectInterface
113 : public DialectInterface::Base<ResourceBlobManagerDialectInterface> {
114public:
115 ResourceBlobManagerDialectInterface(Dialect *dialect)
116 : Base(dialect),
117 blobManager(std::make_shared<DialectResourceBlobManager>()) {}
118
119 /// Return the blob manager held by this interface.
120 DialectResourceBlobManager &getBlobManager() { return *blobManager; }
121 const DialectResourceBlobManager &getBlobManager() const {
122 return *blobManager;
123 }
124
125 /// Set the blob manager held by this interface.
126 void
127 setBlobManager(std::shared_ptr<DialectResourceBlobManager> newBlobManager) {
128 blobManager = std::move(newBlobManager);
129 }
130
131private:
132 /// The blob manager owned by the dialect implementing this interface.
133 std::shared_ptr<DialectResourceBlobManager> blobManager;
134};
135
136/// This class provides a base class for dialects implementing the resource blob
137/// interface. It provides several additional dialect specific utilities on top
138/// of the generic interface. `HandleT` is the type of the handle used to
139/// reference a resource blob.
140template <typename HandleT>
141class ResourceBlobManagerDialectInterfaceBase
142 : public ResourceBlobManagerDialectInterface {
143public:
144 using ResourceBlobManagerDialectInterface::
145 ResourceBlobManagerDialectInterface;
146
147 /// Update the blob for the entry defined by the provided name. This method
148 /// asserts that an entry for the given name exists in the manager.
149 void update(StringRef name, AsmResourceBlob &&newBlob) {
150 getBlobManager().update(name, std::move(newBlob));
151 }
152
153 /// Insert a new resource blob entry with the provided name and optional blob
154 /// data. The name may be modified during insertion if another entry already
155 /// exists with that name. Returns a dialect specific handle to the inserted
156 /// entry.
157 HandleT insert(StringRef name, std::optional<AsmResourceBlob> blob = {}) {
158 return getBlobManager().template insert<HandleT>(
159 cast<typename HandleT::Dialect>(getDialect()), name, std::move(blob));
160 }
161
162 /// Build resources for each of the referenced blobs within this manager.
163 void buildResources(AsmResourceBuilder &provider,
164 ArrayRef<AsmDialectResourceHandle> referencedResources) {
165 for (const AsmDialectResourceHandle &handle : referencedResources) {
166 if (const auto *dialectHandle = dyn_cast<HandleT>(&handle)) {
167 if (auto *blob = dialectHandle->getBlob())
168 provider.buildBlob(dialectHandle->getKey(), *blob);
169 }
170 }
171 }
172};
173
174//===----------------------------------------------------------------------===//
175// DialectResourceBlobHandle
176//===----------------------------------------------------------------------===//
177
178/// This class defines a dialect specific handle to a resource blob. These
179/// handles utilize a StringRef for the internal key, and an AsmResourceBlob as
180/// the underlying data.
181template <typename DialectT>
182struct DialectResourceBlobHandle
183 : public AsmDialectResourceHandleBase<DialectResourceBlobHandle<DialectT>,
184 DialectResourceBlobManager::BlobEntry,
185 DialectT> {
186 using AsmDialectResourceHandleBase<DialectResourceBlobHandle<DialectT>,
187 DialectResourceBlobManager::BlobEntry,
188 DialectT>::AsmDialectResourceHandleBase;
189 using ManagerInterface = ResourceBlobManagerDialectInterfaceBase<
190 DialectResourceBlobHandle<DialectT>>;
191
192 /// Return the human readable string key for this handle.
193 StringRef getKey() const { return this->getResource()->getKey(); }
194
195 /// Return the blob referenced by this handle if the underlying resource has
196 /// been initialized. Returns nullptr otherwise.
197 AsmResourceBlob *getBlob() { return this->getResource()->getBlob(); }
198 const AsmResourceBlob *getBlob() const {
199 return this->getResource()->getBlob();
200 }
201
202 /// Get the interface for the dialect that owns handles of this type. Asserts
203 /// that the dialect is registered.
204 static ManagerInterface &getManagerInterface(MLIRContext *ctx) {
205 auto *dialect = ctx->getOrLoadDialect<DialectT>();
206 assert(dialect && "dialect not registered");
207
208 auto *iface = dialect->template getRegisteredInterface<ManagerInterface>();
209 assert(iface && "dialect doesn't provide the blob manager interface?");
210 return *iface;
211 }
212};
213
214} // namespace mlir
215
216#endif // MLIR_IR_DIALECTRESOURCEBLOBMANAGER_H
217

source code of mlir/include/mlir/IR/DialectResourceBlobManager.h