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 | |
26 | namespace 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. |
33 | class DialectResourceBlobManager { |
34 | public: |
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 | |
96 | private: |
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. |
112 | class ResourceBlobManagerDialectInterface |
113 | : public DialectInterface::Base<ResourceBlobManagerDialectInterface> { |
114 | public: |
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 | |
131 | private: |
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. |
140 | template <typename HandleT> |
141 | class ResourceBlobManagerDialectInterfaceBase |
142 | : public ResourceBlobManagerDialectInterface { |
143 | public: |
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. |
181 | template <typename DialectT> |
182 | struct 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 | |