[mlir] Add a generic DialectResourceBlobManager to simplify resource blob management

[rriddle]

Summary

The DialectResourceBlobManager class provides functionality for managing resource blobs
in a generic, dialect-agnostic fashion. In addition to this class, a dialect interface and custom
resource handle are provided to simplify referencing and interacting with the manager. These
classes intend to simplify the work required for dialects that want to manage resource blobs
during compilation, such as for large elements attrs. The old manager for the resource example
in the test dialect has been updated to use this, which provides and cleaner and more consistent API.

This commit also adds new HeapAsmResourceBlob and ImmortalAsmResourceBlob to simplify
creating resource blobs in common scenarios.

[jpienaar]

Nice! I just did a brief skim as I'm OOO this afternoon but priority tomorrow :)

mlir/include/mlir/IR/AsmState.h

86

Why is alignment needed for deleter?

95–98

I'm not a fan of unsigned as you know ;-)

186

/*dataIsMutable=*/true ?

So no option here for immutable data as there is no initializer?

205

Unmanged/Unowned ? The guaranteed to persist and immortal may make folks thing they do not need to worry about lifetimes, while this is opposite this class doesn't care ensure lifetimes.

210

It is unclear why this is always immutable, e.g., even if I don't own it and just want create a blob to handle memory I manage in a different way, I may be OK in mutating it. And this is probably also an invariant that has to be added to 205 if an external process is allowed to change this too. Say this was in a reference counted outer container of a framework, this is saying the framework is not allowed to release this memory or change the data while compilation is happening if I'm not mistaken (unless we completely treat the contents as opaque this could invalidate data dependent passes).

253

And up to implementation as to if owned or unowned?

[jpienaar]

Nice helpers/simple to use.

mlir/include/mlir/IR/DialectResourceBlobManager.h

132

Could we avoid shared_ptr? I'd expect the lifetime of the blob manager to be exceed that of the MLIRContext.

mlir/lib/IR/DialectResourceBlobManager.cpp

60

Could you add test that has multiple overlapping names with different data?

[rriddle]

mlir/include/mlir/IR/AsmState.h

86

Some aligned allocators want the alignment used when the data was allocated (e.g. llvm::deallocate_buffer).

95–98

Switched to size_t, which is what the standard uses for alignment values.

186

Added a flag for mutability just in case a user wants to allocate (but not allow mutation).

210

I switched to using "Unmanaged" and added a flag for mutability. I think that is cleaner overall, and given that these are mostly just utility for construction I don't think it's a big deal to make mutability configurable.

253

Exactly, whatever created the blob already made those decisions.

mlir/include/mlir/IR/DialectResourceBlobManager.h

132

I wanted a nice default for the places that don't care about sharing and just want an easily configurable default (e.g. LSP server, etc.). If we don't use shared_ptr, we'd need an optional backing storage in each instance just in case the user doesn't want to explicitly provide one. Using shared_ptr felt fine here given that: there are relatively few lifetime changes (which would actually incur the cost of shared_ptr), dereferencing is the same as a normal pointer, etc.

mlir/lib/IR/DialectResourceBlobManager.cpp

60

The original resource patch actually tested for this, which should now be covered by this functionality (given that it replaced the adhoc stuff previously added to the test dialect): https://github.com/llvm/llvm-project/blob/78015047b22dd64f7c647ab7ce04d42e761e7b8f/mlir/unittests/Parser/ResourceTest.cpp#L29

[jpienaar]

Looks good, thanks!

mlir/include/mlir/IR/AsmState.h

 ///     }
 ///   }
 /// #-}
 ///
 
 //===----------------------------------------------------------------------===//
 // Resource Entry
 
-/// This class is used to build resource entries for use by the printer. Each
-/// resource entry is represented using a key/value pair. The provided key must
-/// be unique within the current context, which allows for a client to provide
-/// resource entries without worrying about overlap with other clients.
-class AsmResourceBuilder {
-public:
-  virtual ~AsmResourceBuilder();
-
-  /// Build a resource entry represented by the given bool.
-  virtual void buildBool(StringRef key, bool data) = 0;
-
-  /// Build a resource entry represented by the given human-readable string
-  /// value.
-  virtual void buildString(StringRef key, StringRef data) = 0;
-
-  /// Build an resource entry represented by the given binary blob data.
-  virtual void buildBlob(StringRef key, ArrayRef<char> data,
-                         uint32_t dataAlignment) = 0;
-  /// Build an resource entry represented by the given binary blob data. This is
-  /// a useful overload if the data type is known. Note that this does not
-  /// support `char` element types to avoid accidentally not providing the
-  /// expected alignment of data in situations that treat blobs generically.
-  template <typename T>
-  std::enable_if_t<!std::is_same<T, char>::value> buildBlob(StringRef key,
-                                                            ArrayRef<T> data) {
-    buildBlob(
-        key, ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)),
-        alignof(T));
-  }
-};
-
 /// This class represents a processed binary blob of data. A resource blob is
 /// essentially a collection of data, potentially mutable, with an associated
 /// deleter function (used if the data needs to be destroyed).
 class AsmResourceBlob {
 public:
-  /// A deleter function that frees a blob given the data and allocation size.
-  using DeleterFn = llvm::unique_function<void(const void *data, size_t size)>;
+  /// A deleter function that frees a blob given the data, allocation size, and
+  /// allocation aligment.
+  using DeleterFn =
+      llvm::unique_function<void(void *data, size_t size, size_t align)>;
+
+  //===--------------------------------------------------------------------===//
+  // Construction
+  //===--------------------------------------------------------------------===//
 
   AsmResourceBlob() = default;
-  AsmResourceBlob(ArrayRef<char> data, DeleterFn deleter, bool dataIsMutable)
-      : data(data), deleter(std::move(deleter)), dataIsMutable(dataIsMutable) {}
+  AsmResourceBlob(ArrayRef<char> data, unsigned dataAlignment,
+                  DeleterFn deleter, bool dataIsMutable)
+      : data(data), dataAlignment(dataAlignment), deleter(std::move(deleter)),
+        dataIsMutable(dataIsMutable) {}
   /// Utility constructor that initializes a blob with a non-char type T.
   template <typename T, typename DelT>
   AsmResourceBlob(ArrayRef<T> data, DelT &&deleteFn, bool dataIsMutable)
       : data((const char *)data.data(), data.size() * sizeof(T)),
-        deleter([deleteFn = std::forward<DelT>(deleteFn)](const void *data,
-                                                          size_t size) {
-          return deleteFn((const T *)data, size);
+        dataAlignment(alignof(T)),
+        deleter([deleteFn = std::forward<DelT>(deleteFn)](
+                    void *data, size_t size, size_t align) {
+          return deleteFn((T *)data, size, align);
         }),
         dataIsMutable(dataIsMutable) {}
   AsmResourceBlob(AsmResourceBlob &&) = default;
-  AsmResourceBlob &operator=(AsmResourceBlob &&) = default;
+  AsmResourceBlob &operator=(AsmResourceBlob &&rhs) {
+    // Delete the current blob if necessary.
+    if (deleter)
+      deleter(const_cast<char *>(data.data()), data.size(), dataAlignment);
+
+    // Take the data entries from rhs.
+    data = rhs.data;
+    dataAlignment = rhs.dataAlignment;
+    deleter = std::move(rhs.deleter);
+    dataIsMutable = rhs.dataIsMutable;
+    return *this;
+  }
   AsmResourceBlob(const AsmResourceBlob &) = delete;
   AsmResourceBlob &operator=(const AsmResourceBlob &) = delete;
   ~AsmResourceBlob() {
     if (deleter)
-      deleter(data.data(), data.size());
+      deleter(const_cast<char *>(data.data()), data.size(), dataAlignment);
   }
 
+  //===--------------------------------------------------------------------===//
+  // Data Access
+  //===--------------------------------------------------------------------===//
+
+  /// Return the alignment of the underlying data.
+  unsigned getDataAlignment() const { return dataAlignment; }
+
   /// Return the raw underlying data of this blob.
   ArrayRef<char> getData() const { return data; }
 
+  /// Return the underlying data as an array of the given type. This is an
+  /// inherrently unsafe operation, and should only be used when the data is
+  /// known to be of the correct type.
+  template <typename T>
+  ArrayRef<T> getDataAs() const {
+    return llvm::makeArrayRef<T>((const T *)data.data(),
+                                 data.size() / sizeof(T));
+  }
+
   /// Return a mutable reference to the raw underlying data of this blob.
   /// Asserts that the blob `isMutable`.
   MutableArrayRef<char> getMutableData() {
     assert(isMutable() &&
            "cannot access mutable reference to non-mutable data");
     return MutableArrayRef<char>(const_cast<char *>(data.data()), data.size());
   }
 
   /// Return if the data of this blob is mutable.
   bool isMutable() const { return dataIsMutable; }
 
   /// Return the deleter function of this blob.
   DeleterFn &getDeleter() { return deleter; }
   const DeleterFn &getDeleter() const { return deleter; }
 
 private:
   /// The raw, properly aligned, blob data.
   ArrayRef<char> data;
 
+  /// The alignment of the data.
+  unsigned dataAlignment = 0;
+
   /// An optional deleter function used to deallocate the underlying data when
   /// necessary.
   DeleterFn deleter;
 
   /// Whether the data is mutable.
   bool dataIsMutable;
 };
 
+/// This class provides a simple utility wrapper for creating heap allocated
+/// AsmResourceBlobs.
+class HeapAsmResourceBlob : public AsmResourceBlob {
+public:
+  /// Create a new heap allocated blob with the given size and alignment.
+  static AsmResourceBlob allocate(size_t size, size_t align) {
+    return AsmResourceBlob(
+        ArrayRef<char>((char *)llvm::allocate_buffer(size, align), size), align,
+        llvm::deallocate_buffer, true);
+  }
+  /// Create a new heap allocated blob and copy the provided data into it.
+  static AsmResourceBlob allocateAndCopy(ArrayRef<char> data, size_t align) {
+    AsmResourceBlob blob = allocate(data.size(), align);
+    std::memcpy(blob.getMutableData().data(), data.data(), data.size());
+    return blob;
+  }
+  template <typename T>
+  static std::enable_if_t<!std::is_same<T, char>::value, AsmResourceBlob>
+  allocateAndCopy(ArrayRef<T> data) {
+    return allocateAndCopy(
+        ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)),
+        alignof(T));
+  }
+};
+/// This class provides a simple utility wrapper for creating "immortal"
+/// AsmResourceBlobs. The lifetime of the data provided to these blobs is
+/// guaranteed to persist beyond the lifetime of this reference.
+class ImmortalAsmResourceBlob : public AsmResourceBlob {
+public:
+  /// Create a new heap allocated blob and copy the provided data into it.
+  static AsmResourceBlob allocate(ArrayRef<char> data, size_t align) {
+    return AsmResourceBlob(data, align, /*deleter=*/{},
+                           /*dataIsMutable=*/false);
+  }
+  template <typename T>
+  static std::enable_if_t<!std::is_same<T, char>::value, AsmResourceBlob>
+  allocate(ArrayRef<T> data) {
+    return allocate(
+        ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)),
+        alignof(T));
+  }
+};
+
+/// This class is used to build resource entries for use by the printer. Each
+/// resource entry is represented using a key/value pair. The provided key must
+/// be unique within the current context, which allows for a client to provide
+/// resource entries without worrying about overlap with other clients.
+class AsmResourceBuilder {
+public:
+  virtual ~AsmResourceBuilder();
+
+  /// Build a resource entry represented by the given bool.
+  virtual void buildBool(StringRef key, bool data) = 0;
+
+  /// Build a resource entry represented by the given human-readable string
+  /// value.
+  virtual void buildString(StringRef key, StringRef data) = 0;
+
+  /// Build an resource entry represented by the given binary blob data.
+  virtual void buildBlob(StringRef key, ArrayRef<char> data,
+                         uint32_t dataAlignment) = 0;
+  /// Build an resource entry represented by the given binary blob data. This is
+  /// a useful overload if the data type is known. Note that this does not
+  /// support `char` element types to avoid accidentally not providing the
+  /// expected alignment of data in situations that treat blobs generically.
+  template <typename T>
+  std::enable_if_t<!std::is_same<T, char>::value> buildBlob(StringRef key,
+                                                            ArrayRef<T> data) {
+    buildBlob(
+        key, ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)),
+        alignof(T));
+  }
+  /// Build an resource entry represented by the given resource blob. This is
+  /// a useful overload if a blob already exists in-memory.
+  void buildBlob(StringRef key, const AsmResourceBlob &blob) {
+    buildBlob(key, blob.getData(), blob.getDataAlignment());
+  }
+};
+
 /// This class represents a single parsed resource entry.
 class AsmParsedResourceEntry {
 public:
   virtual ~AsmParsedResourceEntry();
 
   /// Return the key of the resource entry.
   virtual StringRef getKey() const = 0;
 
   /// Emit an error at the location of this entry.
   virtual InFlightDiagnostic emitError() const = 0;
 
   /// Parse the resource entry represented by a boolean. Returns failure if the
   /// entry does not correspond to a bool.
   virtual FailureOr<bool> parseAsBool() const = 0;
 
   /// Parse the resource entry represented by a human-readable string. Returns
   /// failure if the entry does not correspond to a string.
   virtual FailureOr<std::string> parseAsString() const = 0;
 
-  /// The type of an allocator function used to allocate memory for a blob when
-  /// required. The function is provided a size and alignment, and should return
-  /// an aligned allocation buffer.
+  /// An allocator function used to allocate memory for a blob when required.
+  /// The function is provided a size and alignment, and should return an
+  /// aligned allocation buffer.
   using BlobAllocatorFn =
-      function_ref<AsmResourceBlob(unsigned size, unsigned align)>;
+      function_ref<AsmResourceBlob(size_t size, size_t align)>;
 
   /// Parse the resource entry represented by a binary blob. Returns failure if
   /// the entry does not correspond to a blob. If the blob needed to be
   /// allocated, the given allocator function is invoked.
   virtual FailureOr<AsmResourceBlob>
   parseAsBlob(BlobAllocatorFn allocator) const = 0;
+  /// Parse the resource entry represented by a binary blob using heap
+  /// allocation.
+  FailureOr<AsmResourceBlob> parseAsBlob() const {
+    return parseAsBlob(HeapAsmResourceBlob::allocate);
+  }
 };
 
 //===----------------------------------------------------------------------===//
 // Resource Parser/Printer
 
 /// This class represents an instance of a resource parser. This class should be
 /// implemented by non-dialect clients that want to inject additional resources
 /// into MLIR assembly formats.

mlir/include/mlir/IR/Dialect.h

public:
   }
 
   //===--------------------------------------------------------------------===//
   // Interfaces
   //===--------------------------------------------------------------------===//
 
   /// Lookup an interface for the given ID if one is registered, otherwise
   /// nullptr.
-  const DialectInterface *getRegisteredInterface(TypeID interfaceID) {
+  DialectInterface *getRegisteredInterface(TypeID interfaceID) {
     auto it = registeredInterfaces.find(interfaceID);
     return it != registeredInterfaces.end() ? it->getSecond().get() : nullptr;
   }
   template <typename InterfaceT>
-  const InterfaceT *getRegisteredInterface() {
-    return static_cast<const InterfaceT *>(
+  InterfaceT *getRegisteredInterface() {
+    return static_cast<InterfaceT *>(
         getRegisteredInterface(InterfaceT::getInterfaceID()));
   }
 
   /// Lookup an op interface for the given ID if one is registered, otherwise
   /// nullptr.
   virtual void *getRegisteredInterfaceForOp(TypeID interfaceID,
                                             OperationName opName) {
     return nullptr;
public:
   void addInterface(std::unique_ptr<DialectInterface> interface);
 
   /// Register a set of dialect interfaces with this dialect instance.
   template <typename... Args>
   void addInterfaces() {
     (void)std::initializer_list<int>{
         0, (addInterface(std::make_unique<Args>(this)), 0)...};
   }
+  template <typename InterfaceT, typename... Args>
+  InterfaceT &addInterface(Args &&...args) {
+    InterfaceT *interface = new InterfaceT(this, std::forward<Args>(args)...);
+    addInterface(std::unique_ptr<DialectInterface>(interface));
+    return *interface;
+  }
 
 protected:
   /// The constructor takes a unique namespace for this dialect as well as the
   /// context to bind to.
   /// Note: The namespace must not contain '.' characters.
   /// Note: All operations belonging to this dialect must have names starting
   ///       with the namespace followed by '.'.
   /// Example:
struct isa_impl<
     T, ::mlir::Dialect,
     std::enable_if_t<std::is_base_of<::mlir::DialectInterface, T>::value>> {
   static inline bool doit(const ::mlir::Dialect &dialect) {
     return const_cast<::mlir::Dialect &>(dialect).getRegisteredInterface<T>();
   }
 };
 template <typename T>
 struct cast_retty_impl<T, ::mlir::Dialect *> {
-  using ret_type =
-      std::conditional_t<std::is_base_of<::mlir::Dialect, T>::value, T *,
-                         const T *>;
+  using ret_type = T *;
 };
 template <typename T>
 struct cast_retty_impl<T, ::mlir::Dialect> {
-  using ret_type =
-      std::conditional_t<std::is_base_of<::mlir::Dialect, T>::value, T &,
-                         const T &>;
+  using ret_type = T &;
 };
 
 template <typename T>
 struct cast_convert_val<T, ::mlir::Dialect, ::mlir::Dialect> {
   template <typename To>
   static std::enable_if_t<std::is_base_of<::mlir::Dialect, To>::value, To &>
   doitImpl(::mlir::Dialect &dialect) {
     return static_cast<To &>(dialect);
   }
   template <typename To>
   static std::enable_if_t<std::is_base_of<::mlir::DialectInterface, To>::value,
-                          const To &>
+                          To &>
   doitImpl(::mlir::Dialect &dialect) {
     return *dialect.getRegisteredInterface<To>();
   }
 
   static auto &doit(::mlir::Dialect &dialect) { return doitImpl<T>(dialect); }
 };
 template <class T>
 struct cast_convert_val<T, ::mlir::Dialect *, ::mlir::Dialect *> {

mlir/include/mlir/IR/DialectResourceBlobManager.h

+//===- 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.
+    const AsmResourceBlob &getBlob() const { return blob; }
+    AsmResourceBlob &getBlob() { return blob; }
+
+    /// 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, AsmResourceBlob &&newBlob) {
+      key = newKey;
+      blob = std::move(newBlob);
+    }
+
+    /// The key used for this blob.
+    StringRef key;
+
+    /// The blob that is referenced by this entry.
+    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, 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,
+                 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, 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))
+        provider.buildBlob(dialectHandle->getKey(), dialectHandle->getBlob());
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// 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;
+
+  /// Return the human readable string key for this handle.
+  StringRef getKey() const { return this->getResource()->getKey(); }
+
+  /// Return the blob referenced by this handle.
+  AsmResourceBlob &getBlob() { return this->getResource()->getBlob(); }
+  const AsmResourceBlob &getBlob() const {
+    return this->getResource()->getBlob();
+  }
+};
+
+} // namespace mlir
+
+#endif // MLIR_IR_DIALECTRESOURCEBLOBMANAGER_H

mlir/lib/IR/CMakeLists.txt

 add_mlir_library(MLIRIR
   AffineExpr.cpp
   AffineMap.cpp
   AsmPrinter.cpp
   Attributes.cpp
   Block.cpp
   Builders.cpp
   BuiltinAttributeInterfaces.cpp
   BuiltinAttributes.cpp
   BuiltinDialect.cpp
   BuiltinTypes.cpp
   BuiltinTypeInterfaces.cpp
   Diagnostics.cpp
   Dialect.cpp
+  DialectResourceBlobManager.cpp
   Dominance.cpp
   ExtensibleDialect.cpp
   FunctionImplementation.cpp
   FunctionInterfaces.cpp
   IntegerSet.cpp
   Location.cpp
   MLIRContext.cpp
   Operation.cpp

mlir/lib/IR/DialectResourceBlobManager.cpp

+//===- DialectResourceBlobManager.cpp - Dialect Blob Management -----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/DialectResourceBlobManager.h"
+#include "llvm/ADT/SmallString.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// DialectResourceBlobManager
+//===---------------------------------------------------------------------===//
+
+auto DialectResourceBlobManager::lookup(StringRef name) -> BlobEntry * {
+  llvm::sys::SmartScopedReader<true> reader(blobMapLock);
+
+  auto it = blobMap.find(name);
+  return it != blobMap.end() ? &it->second : nullptr;
+}
+
+void DialectResourceBlobManager::update(StringRef name,
+                                        AsmResourceBlob &&newBlob) {
+  BlobEntry *entry = lookup(name);
+  assert(entry && "`update` expects an existing entry for the provided name");
+  entry->setBlob(std::move(newBlob));
+}
+
+auto DialectResourceBlobManager::insert(StringRef name, AsmResourceBlob &&blob)
+    -> BlobEntry & {
+  llvm::sys::SmartScopedWriter<true> writer(blobMapLock);
+
+  // Functor used to attempt insertion with a given name.
+  auto tryInsertion = [&](StringRef name) -> BlobEntry * {
+    auto it = blobMap.try_emplace(name, BlobEntry());
+    if (it.second) {
+      it.first->second.initialize(it.first->getKey(), std::move(blob));
+      return &it.first->second;
+    }
+    return nullptr;
+  };
+
+  // Try inserting with the name provided by the user.
+  if (BlobEntry *entry = tryInsertion(name))
+    return *entry;
+
+  // If an entry already exists for the user provided name, tweak the name and
+  // re-attempt insertion until we find one that is unique.
+  llvm::SmallString<32> nameStorage(name);
+  nameStorage.push_back('_');
+  size_t nameCounter = 1;
+  do {
+    Twine(nameCounter++).toVector(nameStorage);
+
+    // Try inserting with the new name.
+    if (BlobEntry *entry = tryInsertion(name))
+      return *entry;
+    nameStorage.resize(name.size() + 1);
+  } while (true);
+}

mlir/test/lib/Dialect/Test/TestAttrDefs.td

 
 // Test simple extern 1D vector using ElementsAttrInterface.
 def TestExtern1DI64ElementsAttr : Test_Attr<"TestExtern1DI64Elements", [
     ElementsAttrInterface
   ]> {
   let mnemonic = "e1di64_elements";
   let parameters = (ins
     AttributeSelfTypeParameter<"", "::mlir::ShapedType">:$type,
-    ResourceHandleParameter<"TestExternalElementsDataHandle">:$handle
+    ResourceHandleParameter<"TestDialectResourceBlobHandle">:$handle
   );
   let extraClassDeclaration = [{
     /// Return the elements referenced by this attribute.
     llvm::ArrayRef<uint64_t> getElements() const;
 
     /// The set of data types that can be iterated by this attribute.
     using ContiguousIterableTypesT = std::tuple<uint64_t>;
 

mlir/test/lib/Dialect/Test/TestAttributes.h

 #include "TestTraits.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 
 #include "TestAttrInterfaces.h.inc"
 #include "TestOpEnums.h.inc"
+#include "mlir/IR/DialectResourceBlobManager.h"
 
 namespace test {
-struct TestExternalElementsDataHandle;
+class TestDialect;
+
+/// A handle used to reference external elements instances.
+using TestDialectResourceBlobHandle =
+    mlir::DialectResourceBlobHandle<TestDialect>;
 } // namespace test
 
 #define GET_ATTRDEF_CLASSES
 #include "TestAttrDefs.h.inc"
 
 #endif // MLIR_TESTATTRIBUTES_H

mlir/test/lib/Dialect/Test/TestAttributes.cpp

SubElementAttrInterface TestSubElementsAccessAttr::replaceImmediateSubAttribute(
   return get(getContext(), first, second, third);
 }
 
 //===----------------------------------------------------------------------===//
 // TestExtern1DI64ElementsAttr
 //===----------------------------------------------------------------------===//
 
 ArrayRef<uint64_t> TestExtern1DI64ElementsAttr::getElements() const {
-  return getHandle().getData()->getData();
+  return getHandle().getBlob().getDataAs<uint64_t>();
 }
 
 //===----------------------------------------------------------------------===//
 // Tablegen Generated Definitions
 //===----------------------------------------------------------------------===//
 
 #include "TestAttrInterfaces.cpp.inc"
 

mlir/test/lib/Dialect/Test/TestDialect.h

 #include "mlir/Dialect/DLTI/Traits.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/IR/ExtensibleDialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/RegionKindInterface.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
 #include "mlir/Interfaces/DerivedAttributeOpInterface.h"
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 namespace mlir {
 class DLTIDialect;
 class RewritePatternSet;
 } // namespace mlir
 
-namespace test {
-class TestDialect;
-
-//===----------------------------------------------------------------------===//
-// External Elements Data
-//===----------------------------------------------------------------------===//
-
-/// This class represents a single external elements instance. It keeps track of
-/// the data, and deallocates when destructed.
-class TestExternalElementsData : public mlir::AsmResourceBlob {
-public:
-  using mlir::AsmResourceBlob::AsmResourceBlob;
-  TestExternalElementsData(mlir::AsmResourceBlob &&blob)
-      : mlir::AsmResourceBlob(std::move(blob)) {}
-
-  /// Return the data of this external elements instance.
-  llvm::ArrayRef<uint64_t> getData() const;
-
-  /// Allocate a new external elements instance with the given number of
-  /// elements.
-  static TestExternalElementsData allocate(size_t numElements);
-};
-
-/// A handle used to reference external elements instances.
-struct TestExternalElementsDataHandle
-    : public mlir::AsmDialectResourceHandleBase<
-          TestExternalElementsDataHandle,
-          llvm::StringMapEntry<std::unique_ptr<TestExternalElementsData>>,
-          TestDialect> {
-  using AsmDialectResourceHandleBase::AsmDialectResourceHandleBase;
-
-  /// Return a key to use for this handle.
-  llvm::StringRef getKey() const { return getResource()->getKey(); }
-
-  /// Return the data referenced by this handle.
-  TestExternalElementsData *getData() const {
-    return getResource()->getValue().get();
-  }
-};
-
-/// This class acts as a manager for external elements data. It provides API
-/// for creating and accessing registered elements data.
-class TestExternalElementsDataManager {
-  using DataMap = llvm::StringMap<std::unique_ptr<TestExternalElementsData>>;
-
-public:
-  /// Return the data registered for the given name, or nullptr if no data is
-  /// registered.
-  const TestExternalElementsData *getData(llvm::StringRef name) const;
-
-  /// Register an entry with the provided name, which may be modified if another
-  /// entry was already inserted with that name. Returns the inserted entry.
-  std::pair<DataMap::iterator, bool> insert(llvm::StringRef name);
-
-  /// Set the data for the given entry, which is expected to exist.
-  void setData(llvm::StringRef name, TestExternalElementsData &&data);
-
-private:
-  llvm::StringMap<std::unique_ptr<TestExternalElementsData>> dataMap;
-};
-} // namespace test
-
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
 #include "TestOpInterfaces.h.inc"
 #include "TestOpsDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.h.inc"
 
 namespace test {
 void registerTestDialect(::mlir::DialectRegistry &registry);
 void populateTestReductionPatterns(::mlir::RewritePatternSet &patterns);
 } // namespace test
 
 #endif // MLIR_TESTDIALECT_H

mlir/test/lib/Dialect/Test/TestDialect.cpp

 using namespace mlir;
 using namespace test;
 
 void test::registerTestDialect(DialectRegistry &registry) {
   registry.insert<TestDialect>();
 }
 
 //===----------------------------------------------------------------------===//
-// External Elements Data
-//===----------------------------------------------------------------------===//
-
-ArrayRef<uint64_t> TestExternalElementsData::getData() const {
-  ArrayRef<char> data = AsmResourceBlob::getData();
-  return ArrayRef<uint64_t>((const uint64_t *)data.data(),
-                            data.size() / sizeof(uint64_t));
-}
-
-TestExternalElementsData
-TestExternalElementsData::allocate(size_t numElements) {
-  return TestExternalElementsData(
-      llvm::ArrayRef<uint64_t>(new uint64_t[numElements], numElements),
-      [](const uint64_t *data, size_t) { delete[] data; },
-      /*dataIsMutable=*/true);
-}
-
-const TestExternalElementsData *
-TestExternalElementsDataManager::getData(StringRef name) const {
-  auto it = dataMap.find(name);
-  return it != dataMap.end() ? &*it->second : nullptr;
-}
-
-std::pair<TestExternalElementsDataManager::DataMap::iterator, bool>
-TestExternalElementsDataManager::insert(StringRef name) {
-  auto it = dataMap.try_emplace(name, nullptr);
-  if (it.second)
-    return it;
-
-  llvm::SmallString<32> nameStorage(name);
-  nameStorage.push_back('_');
-  size_t nameCounter = 1;
-  do {
-    nameStorage += std::to_string(nameCounter++);
-    auto it = dataMap.try_emplace(nameStorage, nullptr);
-    if (it.second)
-      return it;
-    nameStorage.resize(name.size() + 1);
-  } while (true);
-}
-
-void TestExternalElementsDataManager::setData(StringRef name,
-                                              TestExternalElementsData &&data) {
-  auto it = dataMap.find(name);
-  assert(it != dataMap.end() && "data not registered");
-  it->second = std::make_unique<TestExternalElementsData>(std::move(data));
-}
-
-//===----------------------------------------------------------------------===//
 // TestDialect Interfaces
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// Testing the correctness of some traits.
 static_assert(
     llvm::is_detected<OpTrait::has_implicit_terminator_t,
                       SingleBlockImplicitTerminatorOp>::value,
     "has_implicit_terminator_t does not match SingleBlockImplicitTerminatorOp");
 static_assert(OpTrait::hasSingleBlockImplicitTerminator<
                   SingleBlockImplicitTerminatorOp>::value,
               "hasSingleBlockImplicitTerminator does not match "
               "SingleBlockImplicitTerminatorOp");
 
+struct TestResourceBlobManagerInterface
+    : public ResourceBlobManagerDialectInterfaceBase<
+          TestDialectResourceBlobHandle> {
+  using ResourceBlobManagerDialectInterfaceBase<
+      TestDialectResourceBlobHandle>::ResourceBlobManagerDialectInterfaceBase;
+};
+
 // Test support for interacting with the AsmPrinter.
 struct TestOpAsmInterface : public OpAsmDialectInterface {
   using OpAsmDialectInterface::OpAsmDialectInterface;
+  TestOpAsmInterface(Dialect *dialect, TestResourceBlobManagerInterface &mgr)
+      : OpAsmDialectInterface(dialect), blobManager(mgr) {}
 
   //===------------------------------------------------------------------===//
   // Aliases
   //===------------------------------------------------------------------===//
 
   AliasResult getAlias(Attribute attr, raw_ostream &os) const final {
     StringAttr strAttr = attr.dyn_cast<StringAttr>();
     if (!strAttr)
struct TestOpAsmInterface : public OpAsmDialectInterface {
   }
 
   //===------------------------------------------------------------------===//
   // Resources
   //===------------------------------------------------------------------===//
 
   std::string
   getResourceKey(const AsmDialectResourceHandle &handle) const override {
-    return cast<TestExternalElementsDataHandle>(handle).getKey().str();
+    return cast<TestDialectResourceBlobHandle>(handle).getKey().str();
   }
 
   FailureOr<AsmDialectResourceHandle>
   declareResource(StringRef key) const final {
-    TestDialect *dialect = cast<TestDialect>(getDialect());
-    TestExternalElementsDataManager &mgr = dialect->getExternalDataManager();
-
-    // Resolve the reference by inserting a new entry into the manager.
-    auto it = mgr.insert(key).first;
-    return TestExternalElementsDataHandle(&*it, dialect);
+    return blobManager.insert(key);
   }
 
   LogicalResult parseResource(AsmParsedResourceEntry &entry) const final {
-    TestDialect *dialect = cast<TestDialect>(getDialect());
-    TestExternalElementsDataManager &mgr = dialect->getExternalDataManager();
-
-    // The resource entries are external constant data.
-    auto blobAllocFn = [](unsigned size, unsigned align) {
-      assert(align == alignof(uint64_t) && "unexpected data alignment");
-      return TestExternalElementsData::allocate(size / sizeof(uint64_t));
-    };
-    FailureOr<AsmResourceBlob> blob = entry.parseAsBlob(blobAllocFn);
+    FailureOr<AsmResourceBlob> blob = entry.parseAsBlob();
     if (failed(blob))
       return failure();
 
-    mgr.setData(entry.getKey(), std::move(*blob));
+    // Update the blob for this entry.
+    blobManager.update(entry.getKey(), std::move(*blob));
     return success();
   }
 
   void
   buildResources(Operation *op,
                  const SetVector<AsmDialectResourceHandle> &referencedResources,
                  AsmResourceBuilder &provider) const final {
-    for (const AsmDialectResourceHandle &handle : referencedResources) {
-      const auto &testHandle = cast<TestExternalElementsDataHandle>(handle);
-      provider.buildBlob(testHandle.getKey(), testHandle.getData()->getData());
-    }
+    blobManager.buildResources(provider, referencedResources.getArrayRef());
   }
+
+private:
+  /// The blob manager for the dialect.
+  TestResourceBlobManagerInterface &blobManager;
 };
 
 struct TestDialectFoldInterface : public DialectFoldInterface {
   using DialectFoldInterface::DialectFoldInterface;
 
   /// Registered hook to check if the given region, which is attached to an
   /// operation that is *not* isolated from above, should be used when
   /// materializing constants.
void TestDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
   registerDynamicOp(getDynamicGenericOp(this));
   registerDynamicOp(getDynamicOneOperandTwoResultsOp(this));
   registerDynamicOp(getDynamicCustomParserPrinterOp(this));
 
-  addInterfaces<TestOpAsmInterface, TestDialectFoldInterface,
-                TestInlinerInterface, TestReductionPatternInterface>();
+  auto &blobInterface = addInterface<TestResourceBlobManagerInterface>();
+  addInterface<TestOpAsmInterface>(blobInterface);
+
+  addInterfaces<TestDialectFoldInterface, TestInlinerInterface,
+                TestReductionPatternInterface>();
   allowUnknownOperations();
 
   // Instantiate our fallback op interface that we'll use on specific
   // unregistered op.
   fallbackEffectOpInterfaces = new TestOpEffectInterfaceFallback;
 }
 TestDialect::~TestDialect() {
   delete static_cast<TestOpEffectInterfaceFallback *>(

mlir/test/lib/Dialect/Test/TestDialect.td

let extraClassDeclaration = [{
     ::llvm::unique_function<void(::mlir::Operation *,
                                  ::mlir::OpAsmPrinter &printer)>
      getOperationPrinter(::mlir::Operation *op) const override;
 
      ::mlir::Type parseType(::mlir::DialectAsmParser &parser) const override;
      void printType(::mlir::Type type,
                     ::mlir::DialectAsmPrinter &printer) const override;
 
-    /// Returns the external elements data manager for this dialect.
-    TestExternalElementsDataManager &getExternalDataManager() {
-      return externalDataManager;
-    }
-
   private:
     // Storage for a custom fallback interface.
     void *fallbackEffectOpInterfaces;
 
     ::mlir::Type parseTestType(::mlir::AsmParser &parser,
                                ::llvm::SetVector<::mlir::Type> &stack) const;
     void printTestType(::mlir::Type type, ::mlir::AsmPrinter &printer,
                        ::llvm::SetVector<::mlir::Type> &stack) const;
-
-    /// An external data manager used to test external elements data.
-    TestExternalElementsDataManager externalDataManager;
   }];
 }
 
 #endif // TEST_DIALECT