MLIR  21.0.0git
DialectResourceBlobManager.h
Go to the documentation of this file.
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"
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.
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.
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, std::move(blob));
93  return HandleT(&entry, dialect);
94  }
95 
96  /// Provide access to all the registered blobs via a callable. During access
97  /// the blob map is guaranteed to remain unchanged.
98  void getBlobMap(llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)>
99  accessor) const;
100 
101 private:
102  /// A mutex to protect access to the blob map.
103  mutable llvm::sys::SmartRWMutex<true> blobMapLock;
104 
105  /// The internal map of tracked blobs. StringMap stores entries in distinct
106  /// allocations, so we can freely take references to the data without fear of
107  /// invalidation during additional insertion/deletion.
108  llvm::StringMap<BlobEntry> blobMap;
109 };
110 
111 //===----------------------------------------------------------------------===//
112 // ResourceBlobManagerDialectInterface
113 //===---------------------------------------------------------------------===//
114 
115 /// This class implements a dialect interface that provides common functionality
116 /// for interacting with a resource blob manager.
118  : public DialectInterface::Base<ResourceBlobManagerDialectInterface> {
119 public:
121  : Base(dialect),
122  blobManager(std::make_shared<DialectResourceBlobManager>()) {}
123 
124  /// Return the blob manager held by this interface.
125  DialectResourceBlobManager &getBlobManager() { return *blobManager; }
127  return *blobManager;
128  }
129 
130  /// Set the blob manager held by this interface.
131  void
132  setBlobManager(std::shared_ptr<DialectResourceBlobManager> newBlobManager) {
133  blobManager = std::move(newBlobManager);
134  }
135 
136 private:
137  /// The blob manager owned by the dialect implementing this interface.
138  std::shared_ptr<DialectResourceBlobManager> blobManager;
139 };
140 
141 /// This class provides a base class for dialects implementing the resource blob
142 /// interface. It provides several additional dialect specific utilities on top
143 /// of the generic interface. `HandleT` is the type of the handle used to
144 /// reference a resource blob.
145 template <typename HandleT>
148 public:
151 
152  /// Update the blob for the entry defined by the provided name. This method
153  /// asserts that an entry for the given name exists in the manager.
154  void update(StringRef name, AsmResourceBlob &&newBlob) {
155  getBlobManager().update(name, std::move(newBlob));
156  }
157 
158  /// Insert a new resource blob entry with the provided name and optional blob
159  /// data. The name may be modified during insertion if another entry already
160  /// exists with that name. Returns a dialect specific handle to the inserted
161  /// entry.
162  HandleT insert(StringRef name, std::optional<AsmResourceBlob> blob = {}) {
163  return getBlobManager().template insert<HandleT>(
164  cast<typename HandleT::Dialect>(getDialect()), name, std::move(blob));
165  }
166 
167  /// Build resources for each of the referenced blobs within this manager.
169  ArrayRef<AsmDialectResourceHandle> referencedResources) {
170  for (const AsmDialectResourceHandle &handle : referencedResources) {
171  if (const auto *dialectHandle = dyn_cast<HandleT>(&handle)) {
172  if (auto *blob = dialectHandle->getBlob())
173  provider.buildBlob(dialectHandle->getKey(), *blob);
174  }
175  }
176  }
177 };
178 
179 //===----------------------------------------------------------------------===//
180 // DialectResourceBlobHandle
181 //===----------------------------------------------------------------------===//
182 
183 /// This class defines a dialect specific handle to a resource blob. These
184 /// handles utilize a StringRef for the internal key, and an AsmResourceBlob as
185 /// the underlying data.
186 template <typename DialectT>
188  : public AsmDialectResourceHandleBase<DialectResourceBlobHandle<DialectT>,
189  DialectResourceBlobManager::BlobEntry,
190  DialectT> {
196 
197  /// Return the human readable string key for this handle.
198  StringRef getKey() const { return this->getResource()->getKey(); }
199 
200  /// Return the blob referenced by this handle if the underlying resource has
201  /// been initialized. Returns nullptr otherwise.
202  AsmResourceBlob *getBlob() { return this->getResource()->getBlob(); }
203  const AsmResourceBlob *getBlob() const {
204  return this->getResource()->getBlob();
205  }
206 
207  /// Get the interface for the dialect that owns handles of this type. Asserts
208  /// that the dialect is registered.
210  auto *dialect = ctx->getOrLoadDialect<DialectT>();
211  assert(dialect && "dialect not registered");
212 
213  auto *iface = dialect->template getRegisteredInterface<ManagerInterface>();
214  assert(iface && "dialect doesn't provide the blob manager interface?");
215  return *iface;
216  }
217 };
218 
219 } // namespace mlir
220 
221 #endif // MLIR_IR_DIALECTRESOURCEBLOBMANAGER_H
This class represents a CRTP base class for dialect resource handles.
DialectResourceBlobManager::BlobEntry * getResource()
Return the resource referenced by this handle.
AsmDialectResourceHandleBase(DialectResourceBlobManager::BlobEntry *resource, DialectT *dialect)
Construct a handle from a pointer to the resource.
This class represents an opaque handle to a dialect resource entry.
This class represents a processed binary blob of data.
Definition: AsmState.h:91
This class is used to build resource entries for use by the printer.
Definition: AsmState.h:247
virtual void buildBlob(StringRef key, ArrayRef< char > data, uint32_t dataAlignment)=0
Build an resource entry represented by the given binary blob data.
The class represents an individual entry of a blob.
void setBlob(AsmResourceBlob &&newBlob)
Set the blob owned by this entry.
const AsmResourceBlob * getBlob() const
Return the blob owned by this entry if one has been initialized.
StringRef getKey() const
Return the key used to reference this blob.
This class defines a manager for dialect resource blobs.
HandleT insert(typename HandleT::Dialect *dialect, StringRef name, std::optional< AsmResourceBlob > blob={})
Insertion method that returns a dialect specific handle to the inserted entry.
void update(StringRef name, AsmResourceBlob &&newBlob)
Update the blob for the entry defined by the provided name.
void getBlobMap(llvm::function_ref< void(const llvm::StringMap< BlobEntry > &)> accessor) const
Provide access to all the registered blobs via a callable.
BlobEntry & insert(StringRef name, std::optional< AsmResourceBlob > blob={})
Insert a new entry with the provided name and optional blob data.
BlobEntry * lookup(StringRef name)
Return the blob registered for the given name, or nullptr if no blob is registered.
const BlobEntry * lookup(StringRef name) const
Dialects are groups of MLIR operations, types and attributes, as well as behavior associated with the...
Definition: Dialect.h:38
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
T * getOrLoadDialect()
Get (or create) a dialect for the given derived dialect type.
Definition: MLIRContext.h:97
This class provides a base class for dialects implementing the resource blob interface.
void update(StringRef name, AsmResourceBlob &&newBlob)
Update the blob for the entry defined by the provided name.
void buildResources(AsmResourceBuilder &provider, ArrayRef< AsmDialectResourceHandle > referencedResources)
Build resources for each of the referenced blobs within this manager.
HandleT insert(StringRef name, std::optional< AsmResourceBlob > blob={})
Insert a new resource blob entry with the provided name and optional blob data.
This class implements a dialect interface that provides common functionality for interacting with a r...
DialectResourceBlobManager & getBlobManager()
Return the blob manager held by this interface.
void setBlobManager(std::shared_ptr< DialectResourceBlobManager > newBlobManager)
Set the blob manager held by this interface.
const DialectResourceBlobManager & getBlobManager() const
The base class used for all derived interface types.
Include the generated interface declarations.
This class defines a dialect specific handle to a resource blob.
StringRef getKey() const
Return the human readable string key for this handle.
const AsmResourceBlob * getBlob() const
static ManagerInterface & getManagerInterface(MLIRContext *ctx)
Get the interface for the dialect that owns handles of this type.
AsmResourceBlob * getBlob()
Return the blob referenced by this handle if the underlying resource has been initialized.