[mlir:Bytecode] Add support for encoding resources

[rriddle]

Summary

Resources are encoded in two separate sections similarly to
attributes/types, one for the actual data and one for the data
offsets. Unlike other sections, the resource sections are optional
given that in many cases they won't be present. For testing,
bytecode serialization is added for DenseResourceElementsAttr.

[jpienaar]

(only partial scan)

mlir/docs/BytecodeFormat.md

105

allows ?

107

So the padding is already in file to allow for memory mapping? Or why padding in file? (Unless mistaken one could just had alignment considered when allocating)

272

And this is index into string pool?

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

92

I don't recall max alignment being specified above.

93

Where does this come into effect? E.g., if i had 12, where would it fail further? (I mean I can't think of when it would be useful)

117

An error without a message but just location?

[jpienaar]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

117

I missed that this was InflightDiagnostic so intended for streaming.

512

The newlines between functions or not here is a bit confusing (seems 2 together, newline, 2 together ...) and I can't make much sense of why. Lets have newline between all of the parse methods for consistency.

512

Where is this function used?

630

Is the expectations around these documented somewhere? E.g., they need to initialized already or some such.

So currently we'd emit a warning and skip over?

634

Is a continue needed here?

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

101

Where is this verified before here? E.g., if a file has an invalid alignment where would that be flagged before we get here.

237

Where is this used?

mlir/lib/Bytecode/Writer/IRNumbering.cpp

365

buildBlob before buildBool ?

mlir/lib/Bytecode/Writer/IRNumbering.h

95

Is this naming convention used for resources? (extern vs not, declaration vs definition)

122

SetVector here reminds me: do we have a test for consistent serialization? E.g., given same input you have same output (e.g., would we have detected if we used wrong set data structure here)

mlir/lib/IR/BuiltinDialectBytecode.cpp

406–411

Is this intended to be in the same order as AttributeCode ?

456

Would readResrouceHandle emit an error?

[rriddle]

mlir/docs/BytecodeFormat.md

107

Yeah, the padding is there to ensure that the data is already at the correct alignment in-file (e.g. to support mmaping).

272

Yep

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

92

Do you have any suggestions on how we should approach that? Realistically an upper limit on any system would be something like PAGE size. We want to allow "larger" alignments (at least larger than std::max_align_t), e.g. given that some places want larger alignments when processing the data.

93

We could technically allow it (by aligning to the next power of 2 and just padding to the weird alignment), but I don't have a use case and it's easier to just disallow for now.

512

It's a virtual function, it's called by users parsing resources.

630

Is the expectations around these documented somewhere? E.g., they need to initialized already or some such.

I need to finalize the docs for them and send that out.

So currently we'd emit a warning and skip over?

Yep.

634

No, we still need to skip over the entries, we just don't process them. There is a "continue" in parseGroup for the case of a null handler.

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

237

It's used, e.g., when emitting a section into another one (see emitSection above).

mlir/lib/Bytecode/Writer/IRNumbering.h

95

It is for dialect resources. We declare before defining, i.e. this is how the handles work (you can have a handle to a resource that isn't fully defined).

122

Added. It's also quite difficult to break this given that the API for resources requires passing in a SetVector, so we'd have to purposefully use something else and then construct a SetVector (which would give a werid code smell).

mlir/lib/IR/BuiltinDialectBytecode.cpp

406–411

No, given that we might have different AttributeCodes for the same attribute. It's cleaner to just use alphabetical order here, given that we shouldn't try to derive a connection between the two. I do group Locations separately, but that's mostly conventional with how we group builtin attributes in switches elsewhere in the codebase.

456

Yep.

[jpienaar]

Nice, like the structure containing the resource info.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

92

I don't really ... Perhaps a TODO for now and we can come back to it. PAGE seems appealing, but for now lets trust the user here.

mlir/lib/IR/BuiltinDialectBytecode.cpp

406–411

Ah so alphabetical but with with locations at end, could you add a // Locations or some such, I missed that.

[rriddle]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

92

Added a TODO, I'll add an error in a followup.

mlir/docs/BytecodeFormat.md

 #### Strings
 
 Strings are blobs of characters with an associated length.
 
 ### Sections
 
 ```
 section {
-  id: byte
-  length: varint
+  idAndIsAligned: byte // id | (hasAlign << 7)
+  length: varint,
+
+  alignment: varint?,
+  padding: byte[], // Padding bytes are always `0xCB`.
+
+  data: byte[]
 }
 ```
 
-Sections are a mechanism for grouping data within the bytecode. The enable
+Sections are a mechanism for grouping data within the bytecode. They enable
 delayed processing, which is useful for out-of-order processing of data,
-lazy-loading, and more. Each section contains a Section ID and a length (which
-allowing for skipping over the section).
-
-TODO: Sections should also carry an optional alignment. Add this when necessary.
+lazy-loading, and more. Each section contains a Section ID, whose high bit
+indicates if the section has alignment requirements, a length (which allows for
+skipping over the section), and an optional alignment. When an alignment is
+present, a variable number of padding bytes (0xCB) may appear before the section
+data. The alignment of a section must be a power of 2.
 
 ## MLIR Encoding
 
 Given the generic structure of MLIR, the bytecode encoding is actually fairly
 simplistic. It effectively maps to the core components of MLIR.
 
 ### Top Level Structure
 

 
 When implementing the bytecode interface, dialects are responsible for all
 aspects of the encoding. This includes the indicator for which kind of attribute
 or type is being encoded; the bytecode reader will only know that it has
 encountered an attribute or type of a given dialect, it doesn't encode any
 further information. As such, a common encoding idiom is to use a leading
 `varint` code to indicate how the attribute or type was encoded.
 
+### Resource Section
+
+Resources are encoded using two [sections](#sections), one section
+(`resource_section`) containing the actual encoded representation, and another
+section (`resource_offset_section`) containing the offsets of each encoded
+resource into the previous section.
+
+```
+resource_section {
+  resources: resource[]
+}
+resource {
+  value: resource_bool | resource_string | resource_blob
+}
+resource_bool {
+  value: byte
+}
+resource_string {
+  value: varint
+}
+resource_blob {
+  alignment: varint,
+  size: varint,
+  padding: byte[],
+  blob: byte[]
+}
+
+resource_offset_section {
+  numExternalResourceGroups: varint,
+  resourceGroups: resource_group[]
+}
+resource_group {
+  key: varint,
+  numResources: varint,
+  resources: resource_info[]
+}
+resource_info {
+  key: varint,
+  size: varint
+  kind: byte,
+}
+```
+
+Resources are grouped by the provider, either an external entity or a dialect,
+with each `resource_group` in the offset section containing the corresponding
+provider, number of elements, and info for each element within the group. For
+each element, we record the key, the value kind, and the encoded size. We avoid
+using the direct offset into the `resource_section`, as a smaller relative
+offsets provides more effective compression.
+
 ### IR Section
 
 The IR section contains the encoded form of operations within the bytecode.
 
 #### Operation Encoding
 
 ```
 op {

mlir/include/mlir/Bytecode/BytecodeImplementation.h

 
 #ifndef MLIR_BYTECODE_BYTECODEIMPLEMENTATION_H
 #define MLIR_BYTECODE_BYTECODEIMPLEMENTATION_H
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectInterface.h"
+#include "mlir/IR/OpImplementation.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/Twine.h"
 
 namespace mlir {
 //===----------------------------------------------------------------------===//
 // DialectBytecodeReader
 //===----------------------------------------------------------------------===//
 
LogicalResult readType(T &result) {
     if (failed(readType(baseResult)))
       return failure();
     if ((result = baseResult.dyn_cast<T>()))
       return success();
     return emitError() << "expected " << llvm::getTypeName<T>()
                        << ", but got: " << baseResult;
   }
 
+  /// Read a handle to a dialect resource.
+  template <typename ResourceT>
+  FailureOr<ResourceT> readResourceHandle() {
+    FailureOr<AsmDialectResourceHandle> handle = readResourceHandle();
+    if (failed(handle))
+      return failure();
+    if (auto *result = dyn_cast<ResourceT>(&*handle))
+      return std::move(*result);
+    return emitError() << "provided resource handle differs from the "
+                          "expected resource type";
+  }
+
   //===--------------------------------------------------------------------===//
   // Primitives
   //===--------------------------------------------------------------------===//
 
   /// Read a variable width integer.
   virtual LogicalResult readVarInt(uint64_t &result) = 0;
 
   /// Read a signed variable width integer.
   virtual LogicalResult readSignedVarInt(int64_t &result) = 0;
   LogicalResult readSignedVarInts(SmallVectorImpl<int64_t> &result) {
     return readList(result,
                     [this](int64_t &value) { return readSignedVarInt(value); });
   }
 
   /// Read an APInt that is known to have been encoded with the given width.
   virtual FailureOr<APInt> readAPIntWithKnownWidth(unsigned bitWidth) = 0;
 
   /// Read an APFloat that is known to have been encoded with the given
   /// semantics.
   virtual FailureOr<APFloat>
   readAPFloatWithKnownSemantics(const llvm::fltSemantics &semantics) = 0;
 
   /// Read a string from the bytecode.
   virtual LogicalResult readString(StringRef &result) = 0;
+
+private:
+  /// Read a handle to a dialect resource.
+  virtual FailureOr<AsmDialectResourceHandle> readResourceHandle() = 0;
 };
 
 //===----------------------------------------------------------------------===//
 // DialectBytecodeWriter
 //===----------------------------------------------------------------------===//
 
 /// This class defines a virtual interface for writing to a bytecode stream,
 /// providing hooks into the bytecode writer. As such, this class should only be
public:
 
   /// Write a reference to the given type.
   virtual void writeType(Type type) = 0;
   template <typename T>
   void writeTypes(ArrayRef<T> types) {
     writeList(types, [this](T type) { writeType(type); });
   }
 
+  /// Write the given handle to a dialect resource.
+  virtual void
+  writeResourceHandle(const AsmDialectResourceHandle &resource) = 0;
+
   //===--------------------------------------------------------------------===//
   // Primitives
   //===--------------------------------------------------------------------===//
 
   /// Write a variable width integer to the output stream. This should be the
   /// preferred method for emitting integers whenever possible.
   virtual void writeVarInt(uint64_t value) = 0;
 

mlir/include/mlir/Bytecode/BytecodeWriter.h

 //===- BytecodeWriter.h - MLIR Bytecode Writer ------------------*- 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 header defines interfaces to write MLIR bytecode files/streams.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_BYTECODE_BYTECODEWRITER_H
 #define MLIR_BYTECODE_BYTECODEWRITER_H
 
-#include "mlir/Support/LLVM.h"
-#include "llvm/ADT/StringRef.h"
+#include "mlir/IR/AsmState.h"
 
 namespace mlir {
 class Operation;
 
+/// This class contains the configuration used for the bytecode writer. It
+/// controls various aspects of bytecode generation, and contains all of the
+/// various bytecode writer hooks.
+class BytecodeWriterConfig {
+public:
+  /// `producer` is an optional string that can be used to identify the producer
+  /// of the bytecode when reading. It has no functional effect on the bytecode
+  /// serialization.
+  BytecodeWriterConfig(StringRef producer = "MLIR" LLVM_VERSION_STRING);
+  ~BytecodeWriterConfig();
+
+  /// An internal implementation class that contains the state of the
+  /// configuration.
+  struct Impl;
+
+  /// Return an instance of the internal implementation.
+  const Impl &getImpl() const { return *impl; }
+
+  //===--------------------------------------------------------------------===//
+  // Resources
+  //===--------------------------------------------------------------------===//
+
+  /// Attach the given resource printer to the writer configuration.
+  void attachResourcePrinter(std::unique_ptr<AsmResourcePrinter> printer);
+
+  /// Attach an resource printer, in the form of a callable, to the
+  /// configuration.
+  template <typename CallableT>
+  std::enable_if_t<std::is_convertible<
+      CallableT, function_ref<void(Operation *, AsmResourceBuilder &)>>::value>
+  attachResourcePrinter(StringRef name, CallableT &&printFn) {
+    attachResourcePrinter(AsmResourcePrinter::fromCallable(
+        name, std::forward<CallableT>(printFn)));
+  }
+
+private:
+  /// A pointer to allocated storage for the impl state.
+  std::unique_ptr<Impl> impl;
+};
+
 //===----------------------------------------------------------------------===//
 // Entry Points
 //===----------------------------------------------------------------------===//
 
 /// Write the bytecode for the given operation to the provided output stream.
 /// For streams where it matters, the given stream should be in "binary" mode.
-/// `producer` is an optional string that can be used to identify the producer
-/// of the bytecode when reading. It has no functional effect on the bytecode
-/// serialization.
 void writeBytecodeToFile(Operation *op, raw_ostream &os,
-                         StringRef producer = "MLIR" LLVM_VERSION_STRING);
+                         const BytecodeWriterConfig &config = {});
 
 } // namespace mlir
 
 #endif // MLIR_BYTECODE_BYTECODEWRITER_H

mlir/include/mlir/IR/AsmState.h

public:
   }
   /// 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 enum represents the different kinds of resource values.
+enum class AsmResourceEntryKind {
+  /// A blob of data with an accompanying alignment.
+  Blob,
+  /// A boolean value.
+  Bool,
+  /// A string value.
+  String,
+};
+StringRef toString(AsmResourceEntryKind kind);
+
 /// 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;
 
+  /// Return the kind of this value.
+  virtual AsmResourceEntryKind getKind() 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;
 

mlir/lib/AsmParser/Parser.cpp

public:
   ParsedResourceEntry(StringRef key, SMLoc keyLoc, Token value, Parser &p)
       : key(key), keyLoc(keyLoc), value(value), p(p) {}
   ~ParsedResourceEntry() override = default;
 
   StringRef getKey() const final { return key; }
 
   InFlightDiagnostic emitError() const final { return p.emitError(keyLoc); }
 
+  AsmResourceEntryKind getKind() const final {
+    if (value.isAny(Token::kw_true, Token::kw_false))
+      return AsmResourceEntryKind::Bool;
+    return value.getSpelling().startswith("\"0x")
+               ? AsmResourceEntryKind::Blob
+               : AsmResourceEntryKind::String;
+  }
+
   FailureOr<bool> parseAsBool() const final {
     if (value.is(Token::kw_true))
       return true;
     if (value.is(Token::kw_false))
       return false;
     return p.emitError(value.getLoc(),
                        "expected 'true' or 'false' value for key '" + key +
                            "'");

mlir/lib/Bytecode/Encoding.h

 namespace bytecode {
 //===----------------------------------------------------------------------===//
 // General constants
 //===----------------------------------------------------------------------===//
 
 enum {
   /// The current bytecode version.
   kVersion = 0,
+
+  /// An arbitrary value used to fill alignment padding.
+  kAlignmentByte = 0xCB,
 };
 
 //===----------------------------------------------------------------------===//
 // Sections
 //===----------------------------------------------------------------------===//
 
 namespace Section {
 enum ID : uint8_t {
enum ID : uint8_t {
   /// This section contains the offsets for the attribute and types within the
   /// AttrType section.
   kAttrTypeOffset = 3,
 
   /// This section contains the list of operations serialized into the bytecode,
   /// and their nested regions/operations.
   kIR = 4,
 
+  /// This section contains the resources of the bytecode.
+  kResource = 5,
+
+  /// This section contains the offsets of resources within the Resource
+  /// section.
+  kResourceOffset = 6,
+
   /// The total number of section types.
-  kNumSections = 5,
+  kNumSections = 7,
 };
 } // namespace Section
 
 //===----------------------------------------------------------------------===//
 // IR Section
 //===----------------------------------------------------------------------===//
 
 /// This enum represents a mask of all of the potential components of an

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Verifier.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/Support/SaveAndRestore.h"
 
 #define DEBUG_TYPE "mlir-bytecode-reader"
 
 using namespace mlir;
 
 /// Stringify the given section ID.
 static std::string toString(bytecode::Section::ID sectionID) {
   switch (sectionID) {
   case bytecode::Section::kString:
     return "String (0)";
   case bytecode::Section::kDialect:
     return "Dialect (1)";
   case bytecode::Section::kAttrType:
     return "AttrType (2)";
   case bytecode::Section::kAttrTypeOffset:
     return "AttrTypeOffset (3)";
   case bytecode::Section::kIR:
     return "IR (4)";
+  case bytecode::Section::kResource:
+    return "Resource (5)";
+  case bytecode::Section::kResourceOffset:
+    return "ResourceOffset (6)";
   default:
     return ("Unknown (" + Twine(static_cast<unsigned>(sectionID)) + ")").str();
   }
 }
 
+/// Returns true if the given top-level section ID is optional.
+static bool isSectionOptional(bytecode::Section::ID sectionID) {
+  switch (sectionID) {
+  case bytecode::Section::kString:
+  case bytecode::Section::kDialect:
+  case bytecode::Section::kAttrType:
+  case bytecode::Section::kAttrTypeOffset:
+  case bytecode::Section::kIR:
+    return false;
+  case bytecode::Section::kResource:
+  case bytecode::Section::kResourceOffset:
+    return true;
+  default:
+    llvm_unreachable("unknown section ID");
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // EncodingReader
 //===----------------------------------------------------------------------===//
 
 namespace {
 class EncodingReader {
 public:
   explicit EncodingReader(ArrayRef<uint8_t> contents, Location fileLoc)
       : dataIt(contents.data()), dataEnd(contents.end()), fileLoc(fileLoc) {}
   explicit EncodingReader(StringRef contents, Location fileLoc)
       : EncodingReader({reinterpret_cast<const uint8_t *>(contents.data()),
                         contents.size()},
                        fileLoc) {}
 
   /// Returns true if the entire section has been read.
   bool empty() const { return dataIt == dataEnd; }
 
   /// Returns the remaining size of the bytecode.
   size_t size() const { return dataEnd - dataIt; }
 
+  /// Align the current reader position to the specified alignment.
+  LogicalResult alignTo(unsigned alignment) {
+    if (!llvm::isPowerOf2_32(alignment))
+      return emitError("expected alignment to be a power-of-two");
+
+    // Shift the reader position to the next alignment boundary.
+    while (uintptr_t(dataIt) & (uintptr_t(alignment) - 1)) {
+      uint8_t padding;
+      if (failed(parseByte(padding)))
+        return failure();
+      if (padding != bytecode::kAlignmentByte) {
+        return emitError("expected alignment byte (0xCB), but got: '0x" +
+                         llvm::utohexstr(padding) + "'");
+      }
+    }
+
+    // TODO: Check that the current data pointer is actually at the expected
+    // alignment.
+
+    return success();
+  }
+
   /// Emit an error using the given arguments.
   template <typename... Args>
   InFlightDiagnostic emitError(Args &&...args) const {
     return ::emitError(fileLoc).append(std::forward<Args>(args)...);
   }
+  InFlightDiagnostic emitError() const { return ::emitError(fileLoc); }
 
   /// Parse a single byte from the stream.
   template <typename T>
   LogicalResult parseByte(T &value) {
     if (empty())
       return emitError("attempting to parse a byte at the end of the bytecode");
     value = static_cast<T>(*dataIt++);
     return success();
if (length > size()) {
       return emitError("attempting to parse ", length, " bytes when only ",
                        size(), " remain");
     }
     memcpy(result, dataIt, length);
     dataIt += length;
     return success();
   }
 
+  /// Parse an aligned blob of data, where the alignment was encoded alongside
+  /// the data.
+  LogicalResult parseBlobAndAlignment(ArrayRef<uint8_t> &data,
+                                      uint64_t &alignment) {
+    uint64_t dataSize;
+    if (failed(parseVarInt(alignment)) || failed(parseVarInt(dataSize)) ||
+        failed(alignTo(alignment)))
+      return failure();
+    return parseBytes(dataSize, data);
+  }
+
   /// Parse a variable length encoded integer from the byte stream. The first
   /// encoded byte contains a prefix in the low bits indicating the encoded
   /// length of the value. This length prefix is a bit sequence of '0's followed
   /// by a '1'. The number of '0' bits indicate the number of _additional_ bytes
   /// (not including the prefix byte). All remaining bits in the first byte,
   /// along with all of the bits in additional bytes, provide the value of the
   /// integer encoded in little-endian order.
   LogicalResult parseVarInt(uint64_t &result) {
LogicalResult parseNullTerminatedString(StringRef &result) {
     dataIt = (const uint8_t *)nulIt + 1;
     return success();
   }
 
   /// Parse a section header, placing the kind of section in `sectionID` and the
   /// contents of the section in `sectionData`.
   LogicalResult parseSection(bytecode::Section::ID &sectionID,
                              ArrayRef<uint8_t> &sectionData) {
+    uint8_t sectionIDAndHasAlignment;
     uint64_t length;
-    if (failed(parseByte(sectionID)) || failed(parseVarInt(length)))
+    if (failed(parseByte(sectionIDAndHasAlignment)) ||
+        failed(parseVarInt(length)))
       return failure();
+
+    // Extract the section ID and whether the section is aligned. The high bit
+    // of the ID is the alignment flag.
+    sectionID = static_cast<bytecode::Section::ID>(sectionIDAndHasAlignment &
+                                                   0b01111111);
+    bool hasAlignment = sectionIDAndHasAlignment & 0b10000000;
+
+    // Check that the section is actually valid before trying to process its
+    // data.
     if (sectionID >= bytecode::Section::kNumSections)
       return emitError("invalid section ID: ", unsigned(sectionID));
 
-    // Parse the actua section data now that we have its length.
+    // Process the section alignment if present.
+    if (hasAlignment) {
+      uint64_t alignment;
+      if (failed(parseVarInt(alignment)) || failed(alignTo(alignment)))
+        return failure();
+    }
+
+    // Parse the actual section data.
     return parseBytes(static_cast<size_t>(length), sectionData);
   }
 
 private:
   /// Parse a variable length encoded integer from the byte stream. This method
   /// is a fallback when the number of bytes used to encode the value is greater
   /// than 1, but less than the max (9). The provided `result` value can be
   /// assumed to already contain the first byte of the value.
LogicalResult load(EncodingReader &reader, MLIRContext *ctx) {
 
     // If the dialect was actually loaded, check to see if it has a bytecode
     // interface.
     if (loadedDialect)
       interface = dyn_cast<BytecodeDialectInterface>(loadedDialect);
     return success();
   }
 
+  /// Return the loaded dialect, or nullptr if the dialect is unknown. This can
+  /// only be called after `load`.
+  Dialect *getLoadedDialect() const {
+    assert(dialect &&
+           "expected `load` to be invoked before `getLoadedDialect`");
+    return *dialect;
+  }
+
   /// The loaded dialect entry. This field is None if we haven't attempted to
   /// load, nullptr if we failed to load, otherwise the loaded dialect.
   Optional<Dialect *> dialect;
 
   /// The bytecode interface of the dialect, or nullptr if the dialect does not
   /// implement the bytecode interface. This field should only be checked if the
   /// `dialect` field is non-None.
   const BytecodeDialectInterface *interface = nullptr;
static LogicalResult parseDialectGrouping(
 
   for (uint64_t i = 0; i < numEntries; ++i)
     if (failed(entryCallback(dialect)))
       return failure();
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// ResourceSectionReader
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class is used to read the resource section from the bytecode.
+class ResourceSectionReader {
+public:
+  /// Initialize the resource section reader with the given section data.
+  LogicalResult initialize(Location fileLoc, const ParserConfig &config,
+                           MutableArrayRef<BytecodeDialect> dialects,
+                           StringSectionReader &stringReader,
+                           ArrayRef<uint8_t> sectionData,
+                           ArrayRef<uint8_t> offsetSectionData);
+
+  /// Parse a dialect resource handle from the resource section.
+  LogicalResult parseResourceHandle(EncodingReader &reader,
+                                    AsmDialectResourceHandle &result) {
+    return parseEntry(reader, dialectResources, result, "resource handle");
+  }
+
+private:
+  /// The table of dialect resources within the bytecode file.
+  SmallVector<AsmDialectResourceHandle> dialectResources;
+};
+
+class ParsedResourceEntry : public AsmParsedResourceEntry {
+public:
+  ParsedResourceEntry(StringRef key, AsmResourceEntryKind kind,
+                      EncodingReader &reader, StringSectionReader &stringReader)
+      : key(key), kind(kind), reader(reader), stringReader(stringReader) {}
+  ~ParsedResourceEntry() override = default;
+
+  StringRef getKey() const final { return key; }
+
+  InFlightDiagnostic emitError() const final { return reader.emitError(); }
+
+  AsmResourceEntryKind getKind() const final { return kind; }
+
+  FailureOr<bool> parseAsBool() const final {
+    if (kind != AsmResourceEntryKind::Bool)
+      return emitError() << "expected a bool resource entry, but found a "
+                         << toString(kind) << " entry instead";
+
+    bool value;
+    if (failed(reader.parseByte(value)))
+      return failure();
+    return value;
+  }
+  FailureOr<std::string> parseAsString() const final {
+    if (kind != AsmResourceEntryKind::String)
+      return emitError() << "expected a string resource entry, but found a "
+                         << toString(kind) << " entry instead";
+
+    StringRef string;
+    if (failed(stringReader.parseString(reader, string)))
+      return failure();
+    return string.str();
+  }
+
+  FailureOr<AsmResourceBlob>
+  parseAsBlob(BlobAllocatorFn allocator) const final {
+    if (kind != AsmResourceEntryKind::Blob)
+      return emitError() << "expected a blob resource entry, but found a "
+                         << toString(kind) << " entry instead";
+
+    ArrayRef<uint8_t> data;
+    uint64_t alignment;
+    if (failed(reader.parseBlobAndAlignment(data, alignment)))
+      return failure();
+
+    // Allocate memory for the blob using the provided allocator and copy the
+    // data into it.
+    // FIXME: If the current holder of the bytecode can ensure its lifetime
+    // (e.g. when mmap'd), we should not copy the data. We should use the data
+    // from the bytecode directly.
+    AsmResourceBlob blob = allocator(data.size(), alignment);
+    assert(llvm::isAddrAligned(llvm::Align(alignment), blob.getData().data()) &&
+           blob.isMutable() &&
+           "blob allocator did not return a properly aligned address");
+    memcpy(blob.getMutableData().data(), data.data(), data.size());
+    return blob;
+  }
+
+private:
+  StringRef key;
+  AsmResourceEntryKind kind;
+  EncodingReader &reader;
+  StringSectionReader &stringReader;
+};
+} // namespace
+
+template <typename T>
+static LogicalResult
+parseResourceGroup(Location fileLoc, bool allowEmpty,
+                   EncodingReader &offsetReader, EncodingReader &resourceReader,
+                   StringSectionReader &stringReader, T *handler,
+                   function_ref<LogicalResult(StringRef)> processKeyFn = {}) {
+  uint64_t numResources;
+  if (failed(offsetReader.parseVarInt(numResources)))
+    return failure();
+
+  for (uint64_t i = 0; i < numResources; ++i) {
+    StringRef key;
+    AsmResourceEntryKind kind;
+    uint64_t resourceOffset;
+    ArrayRef<uint8_t> data;
+    if (failed(stringReader.parseString(offsetReader, key)) ||
+        failed(offsetReader.parseVarInt(resourceOffset)) ||
+        failed(offsetReader.parseByte(kind)) ||
+        failed(resourceReader.parseBytes(resourceOffset, data)))
+      return failure();
+
+    // Process the resource key.
+    if ((processKeyFn && failed(processKeyFn(key))))
+      return failure();
+
+    // If the resource data is empty and we allow it, don't error out when
+    // parsing below, just skip it.
+    if (allowEmpty && data.empty())
+      continue;
+
+    // Ignore the entry if we don't have a valid handler.
+    if (!handler)
+      continue;
+
+    // Otherwise, parse the resource value.
+    EncodingReader entryReader(data, fileLoc);
+    ParsedResourceEntry entry(key, kind, entryReader, stringReader);
+    if (failed(handler->parseResource(entry)))
+      return failure();
+    if (!entryReader.empty()) {
+      return entryReader.emitError(
+          "unexpected trailing bytes in resource entry '", key, "'");
+    }
+  }
+  return success();
+}
+
+LogicalResult
+ResourceSectionReader::initialize(Location fileLoc, const ParserConfig &config,
+                                  MutableArrayRef<BytecodeDialect> dialects,
+                                  StringSectionReader &stringReader,
+                                  ArrayRef<uint8_t> sectionData,
+                                  ArrayRef<uint8_t> offsetSectionData) {
+  EncodingReader resourceReader(sectionData, fileLoc);
+  EncodingReader offsetReader(offsetSectionData, fileLoc);
+
+  // Read the number of external resource providers.
+  uint64_t numExternalResourceGroups;
+  if (failed(offsetReader.parseVarInt(numExternalResourceGroups)))
+    return failure();
+
+  // Utility functor that dispatches to `parseResourceGroup`, but implicitly
+  // provides most of the arguments.
+  auto parseGroup = [&](auto *handler, bool allowEmpty = false,
+                        function_ref<LogicalResult(StringRef)> keyFn = {}) {
+    return parseResourceGroup(fileLoc, allowEmpty, offsetReader, resourceReader,
+                              stringReader, handler, keyFn);
+  };
+
+  // Read the external resources from the bytecode.
+  for (uint64_t i = 0; i < numExternalResourceGroups; ++i) {
+    StringRef key;
+    if (failed(stringReader.parseString(offsetReader, key)))
+      return failure();
+
+    // Get the handler for these resources.
+    // TODO: Should we require handling external resources in some scenarios?
+    AsmResourceParser *handler = config.getResourceParser(key);
+    if (!handler) {
+      emitWarning(fileLoc) << "ignoring unknown external resources for '" << key
+                           << "'";
+    }
+
+    if (failed(parseGroup(handler)))
+      return failure();
+  }
+
+  // Read the dialect resources from the bytecode.
+  MLIRContext *ctx = fileLoc->getContext();
+  while (!offsetReader.empty()) {
+    BytecodeDialect *dialect;
+    if (failed(parseEntry(offsetReader, dialects, dialect, "dialect")) ||
+        failed(dialect->load(resourceReader, ctx)))
+      return failure();
+    Dialect *loadedDialect = dialect->getLoadedDialect();
+    if (!loadedDialect) {
+      return resourceReader.emitError()
+             << "dialect '" << dialect->name << "' is unknown";
+    }
+    const auto *handler = dyn_cast<OpAsmDialectInterface>(loadedDialect);
+    if (!handler) {
+      return resourceReader.emitError()
+             << "unexpected resources for dialect '" << dialect->name << "'";
+    }
+
+    // Ensure that each resource is declared before being processed.
+    auto processResourceKeyFn = [&](StringRef key) -> LogicalResult {
+      FailureOr<AsmDialectResourceHandle> handle =
+          handler->declareResource(key);
+      if (failed(handle)) {
+        return resourceReader.emitError()
+               << "unknown 'resource' key '" << key << "' for dialect '"
+               << dialect->name << "'";
+      }
+      dialectResources.push_back(*handle);
+      return success();
+    };
+
+    // Parse the resources for this dialect. We allow empty resources because we
+    // just treat these as declarations.
+    if (failed(parseGroup(handler, /*allowEmpty=*/true, processResourceKeyFn)))
+      return failure();
+  }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // Attribute/Type Reader
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class provides support for reading attribute and type entries from the
 /// bytecode. Attribute and Type entries are read lazily on demand, so we use
 /// this reader to manage when to actually parse them from the bytecode.
 class AttrTypeReader {
struct Entry {
     bool hasCustomEncoding = false;
     /// The raw data of this entry in the bytecode.
     ArrayRef<uint8_t> data;
   };
   using AttrEntry = Entry<Attribute>;
   using TypeEntry = Entry<Type>;
 
 public:
-  AttrTypeReader(StringSectionReader &stringReader, Location fileLoc)
-      : stringReader(stringReader), fileLoc(fileLoc) {}
+  AttrTypeReader(StringSectionReader &stringReader,
+                 ResourceSectionReader &resourceReader, Location fileLoc)
+      : stringReader(stringReader), resourceReader(resourceReader),
+        fileLoc(fileLoc) {}
 
   /// Initialize the attribute and type information within the reader.
   LogicalResult initialize(MutableArrayRef<BytecodeDialect> dialects,
                            ArrayRef<uint8_t> sectionData,
                            ArrayRef<uint8_t> offsetSectionData);
 
   /// Resolve the attribute or type at the given index. Returns nullptr on
   /// failure.
private:
   template <typename T>
   LogicalResult parseCustomEntry(Entry<T> &entry, EncodingReader &reader,
                                  StringRef entryType);
 
   /// The string section reader used to resolve string references when parsing
   /// custom encoded attribute/type entries.
   StringSectionReader &stringReader;
 
+  /// The resource section reader used to resolve resource references when
+  /// parsing custom encoded attribute/type entries.
+  ResourceSectionReader &resourceReader;
+
   /// The set of attribute and type entries.
   SmallVector<AttrEntry> attributes;
   SmallVector<TypeEntry> types;
 
   /// A location used for error emission.
   Location fileLoc;
 };
 
 class DialectReader : public DialectBytecodeReader {
 public:
   DialectReader(AttrTypeReader &attrTypeReader,
-                StringSectionReader &stringReader, EncodingReader &reader)
+                StringSectionReader &stringReader,
+                ResourceSectionReader &resourceReader, EncodingReader &reader)
       : attrTypeReader(attrTypeReader), stringReader(stringReader),
-        reader(reader) {}
+        resourceReader(resourceReader), reader(reader) {}
 
   InFlightDiagnostic emitError(const Twine &msg) override {
     return reader.emitError(msg);
   }
 
   //===--------------------------------------------------------------------===//
   // IR
   //===--------------------------------------------------------------------===//
 
   LogicalResult readAttribute(Attribute &result) override {
     return attrTypeReader.parseAttribute(reader, result);
   }
 
   LogicalResult readType(Type &result) override {
     return attrTypeReader.parseType(reader, result);
   }
 
+  FailureOr<AsmDialectResourceHandle> readResourceHandle() override {
+    AsmDialectResourceHandle handle;
+    if (failed(resourceReader.parseResourceHandle(reader, handle)))
+      return failure();
+    return handle;
+  }
+
   //===--------------------------------------------------------------------===//
   // Primitives
   //===--------------------------------------------------------------------===//
 
   LogicalResult readVarInt(uint64_t &result) override {
     return reader.parseVarInt(result);
   }
 
public:
 
   LogicalResult readString(StringRef &result) override {
     return stringReader.parseString(reader, result);
   }
 
 private:
   AttrTypeReader &attrTypeReader;
   StringSectionReader &stringReader;
+  ResourceSectionReader &resourceReader;
   EncodingReader &reader;
 };
 } // namespace
 
 LogicalResult
 AttrTypeReader::initialize(MutableArrayRef<BytecodeDialect> dialects,
                            ArrayRef<uint8_t> sectionData,
                            ArrayRef<uint8_t> offsetSectionData) {
LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
 
   // Ensure that the dialect implements the bytecode interface.
   if (!entry.dialect->interface) {
     return reader.emitError("dialect '", entry.dialect->name,
                             "' does not implement the bytecode interface");
   }
 
   // Ask the dialect to parse the entry.
-  DialectReader dialectReader(*this, stringReader, reader);
+  DialectReader dialectReader(*this, stringReader, resourceReader, reader);
   if constexpr (std::is_same_v<T, Type>)
     entry.entry = entry.dialect->interface->readType(dialectReader);
   else
     entry.entry = entry.dialect->interface->readAttribute(dialectReader);
   return success(!!entry.entry);
 }
 
 //===----------------------------------------------------------------------===//
 // Bytecode Reader
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class is used to read a bytecode buffer and translate it into MLIR.
 class BytecodeReader {
 public:
   BytecodeReader(Location fileLoc, const ParserConfig &config)
-      : config(config), fileLoc(fileLoc), attrTypeReader(stringReader, fileLoc),
+      : config(config), fileLoc(fileLoc),
+        attrTypeReader(stringReader, resourceReader, fileLoc),
         // Use the builtin unrealized conversion cast operation to represent
         // forward references to values that aren't yet defined.
         forwardRefOpState(UnknownLoc::get(config.getContext()),
                           "builtin.unrealized_conversion_cast", ValueRange(),
                           NoneType::get(config.getContext())) {}
 
   /// Read the bytecode defined within `buffer` into the given block.
   LogicalResult read(llvm::MemoryBufferRef buffer, Block *block);
private:
   LogicalResult parseAttribute(EncodingReader &reader, T &result) {
     return attrTypeReader.parseAttribute(reader, result);
   }
   LogicalResult parseType(EncodingReader &reader, Type &result) {
     return attrTypeReader.parseType(reader, result);
   }
 
   //===--------------------------------------------------------------------===//
+  // Resource Section
+
+  LogicalResult
+  parseResourceSection(Optional<ArrayRef<uint8_t>> resourceData,
+                       Optional<ArrayRef<uint8_t>> resourceOffsetData);
+
+  //===--------------------------------------------------------------------===//
   // IR Section
 
   /// This struct represents the current read state of a range of regions. This
   /// struct is used to enable iterative parsing of regions.
   struct RegionReadState {
     RegionReadState(Operation *op, bool isIsolatedFromAbove)
         : RegionReadState(op->getRegions(), isIsolatedFromAbove) {}
     RegionReadState(MutableArrayRef<Region> regions, bool isIsolatedFromAbove)
private:
 
   /// The producer of the bytecode being read.
   StringRef producer;
 
   /// The table of IR units referenced within the bytecode file.
   SmallVector<BytecodeDialect> dialects;
   SmallVector<BytecodeOperationName> opNames;
 
+  /// The reader used to process resources within the bytecode.
+  ResourceSectionReader resourceReader;
+
   /// The table of strings referenced within the bytecode file.
   StringSectionReader stringReader;
 
   /// The current set of available IR value scopes.
   std::vector<ValueScope> valueScopes;
   /// A block containing the set of operations defined to create forward
   /// references.
   Block forwardRefOps;
while (!reader.empty()) {
 
     // Check for duplicate sections, we only expect one instance of each.
     if (sectionDatas[sectionID]) {
       return reader.emitError("duplicate top-level section: ",
                               toString(sectionID));
     }
     sectionDatas[sectionID] = sectionData;
   }
-  // Check that all of the sections were found.
+  // Check that all of the required sections were found.
   for (int i = 0; i < bytecode::Section::kNumSections; ++i) {
-    if (!sectionDatas[i]) {
+    bytecode::Section::ID sectionID = static_cast<bytecode::Section::ID>(i);
+    if (!sectionDatas[i] && !isSectionOptional(sectionID)) {
       return reader.emitError("missing data for top-level section: ",
-                              toString(bytecode::Section::ID(i)));
+                              toString(sectionID));
     }
   }
 
   // Process the string section first.
   if (failed(stringReader.initialize(
           fileLoc, *sectionDatas[bytecode::Section::kString])))
     return failure();
 
   // Process the dialect section.
   if (failed(parseDialectSection(*sectionDatas[bytecode::Section::kDialect])))
     return failure();
 
+  // Process the resource section if present.
+  if (failed(parseResourceSection(
+          sectionDatas[bytecode::Section::kResource],
+          sectionDatas[bytecode::Section::kResourceOffset])))
+    return failure();
+
   // Process the attribute and type section.
   if (failed(attrTypeReader.initialize(
           dialects, *sectionDatas[bytecode::Section::kAttrType],
           *sectionDatas[bytecode::Section::kAttrTypeOffset])))
     return failure();
 
   // Finally, process the IR section.
   return parseIRSection(*sectionDatas[bytecode::Section::kIR], block);
if (failed(opName->dialect->load(reader, getContext())))
       return failure();
     opName->opName.emplace((opName->dialect->name + "." + opName->name).str(),
                            getContext());
   }
   return *opName->opName;
 }
 
 //===----------------------------------------------------------------------===//
+// Resource Section
+
+LogicalResult BytecodeReader::parseResourceSection(
+    Optional<ArrayRef<uint8_t>> resourceData,
+    Optional<ArrayRef<uint8_t>> resourceOffsetData) {
+  // Ensure both sections are either present or not.
+  if (resourceData.has_value() != resourceOffsetData.has_value()) {
+    if (resourceOffsetData)
+      return emitError(fileLoc, "unexpected resource offset section when "
+                                "resource section is not present");
+    return emitError(
+        fileLoc,
+        "expected resource offset section when resource section is present");
+  }
+
+  // If the resource sections are absent, there is nothing to do.
+  if (!resourceData)
+    return success();
+
+  // Initialize the resource reader with the resource sections.
+  return resourceReader.initialize(fileLoc, config, dialects, stringReader,
+                                   *resourceData, *resourceOffsetData);
+}
+
+//===----------------------------------------------------------------------===//
 // IR Section
 
 LogicalResult BytecodeReader::parseIRSection(ArrayRef<uint8_t> sectionData,
                                              Block *block) {
   EncodingReader reader(sectionData, fileLoc);
 
   // A stack of operation regions currently being read from the bytecode.
   std::vector<RegionReadState> regionStack;

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

 #include <random>
 
 #define DEBUG_TYPE "mlir-bytecode-writer"
 
 using namespace mlir;
 using namespace mlir::bytecode::detail;
 
 //===----------------------------------------------------------------------===//
+// BytecodeWriterConfig
+//===----------------------------------------------------------------------===//
+
+struct BytecodeWriterConfig::Impl {
+  Impl(StringRef producer) : producer(producer) {}
+
+  /// The producer of the bytecode.
+  StringRef producer;
+
+  /// A collection of non-dialect resource printers.
+  SmallVector<std::unique_ptr<AsmResourcePrinter>> externalResourcePrinters;
+};
+
+BytecodeWriterConfig::BytecodeWriterConfig(StringRef producer)
+    : impl(std::make_unique<Impl>(producer)) {}
+BytecodeWriterConfig::~BytecodeWriterConfig() = default;
+
+void BytecodeWriterConfig::attachResourcePrinter(
+    std::unique_ptr<AsmResourcePrinter> printer) {
+  impl->externalResourcePrinters.emplace_back(std::move(printer));
+}
+
+//===----------------------------------------------------------------------===//
 // EncodingEmitter
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class functions as the underlying encoding emitter for the bytecode
 /// writer. This class is a bit different compared to other types of encoders;
 /// it does not use a single buffer, but instead may contain several buffers
 /// (some owned by the writer, and some not) that get concatted during the final
public:
 
   /// Backpatch a byte in the result buffer at the given offset.
   void patchByte(uint64_t offset, uint8_t value) {
     assert(offset < size() && offset >= prevResultSize &&
            "cannot patch previously emitted data");
     currentResult[offset - prevResultSize] = value;
   }
 
+  /// Emit the provided blob of data that has the given alignment. The alignment
+  /// value is also encoded, making it available on load.
+  void emitBlobAndAlignment(ArrayRef<uint8_t> data, uint32_t alignment) {
+    emitVarInt(alignment);
+    emitVarInt(data.size());
+
+    alignTo(alignment);
+    emitBytes(data);
+  }
+  void emitBlobAndAlignment(ArrayRef<char> data, uint32_t alignment) {
+    ArrayRef<uint8_t> castedData(reinterpret_cast<const uint8_t *>(data.data()),
+                                 data.size());
+    emitBlobAndAlignment(castedData, alignment);
+  }
+
+  /// Align the emitter to the given alignment.
+  void alignTo(unsigned alignment) {
+    if (alignment < 2)
+      return;
+    assert(llvm::isPowerOf2_32(alignment) && "expected valid alignment");
+
+    // Check to see if we need to emit any padding bytes to meet the desired
+    // alignment.
+    size_t curOffset = size();
+    size_t paddingSize = llvm::alignTo(curOffset, alignment) - curOffset;
+    while (paddingSize--)
+      emitByte(bytecode::kAlignmentByte);
+
+    // Keep track of the maximum required alignment.
+    requiredAlignment = std::max(requiredAlignment, alignment);
+  }
+
   //===--------------------------------------------------------------------===//
   // Integer Emission
 
   /// Emit a single byte.
   template <typename T>
   void emitByte(T byte) {
     currentResult.push_back(static_cast<uint8_t>(byte));
   }
public:
   }
 
   //===--------------------------------------------------------------------===//
   // Section Emission
 
   /// Emit a nested section of the given code, whose contents are encoded in the
   /// provided emitter.
   void emitSection(bytecode::Section::ID code, EncodingEmitter &&emitter) {
-    // Emit the section code and length.
+    // Emit the section code and length. The high bit of the code is used to
+    // indicate whether the section alignment is present, so save an offset to
+    // it.
+    uint64_t codeOffset = currentResult.size();
     emitByte(code);
     emitVarInt(emitter.size());
 
+    // Integrate the alignment of the section into this emitter if necessary.
+    unsigned emitterAlign = emitter.requiredAlignment;
+    if (emitterAlign > 1) {
+      if (size() & (emitterAlign - 1)) {
+        emitVarInt(emitterAlign);
+        alignTo(emitterAlign);
+
+        // Indicate that we needed to align the section, the high bit of the
+        // code field is used for this.
+        currentResult[codeOffset] |= 0b10000000;
+      } else {
+        // Otherwise, if we happen to be at a compatible offset, we just
+        // remember that we need this alignment.
+        requiredAlignment = std::max(requiredAlignment, emitterAlign);
+      }
+    }
+
     // Push our current buffer and then merge the provided section body into
     // ours.
     appendResult(std::move(currentResult));
     for (std::vector<uint8_t> &result : emitter.prevResultStorage)
       appendResult(std::move(result));
     appendResult(std::move(emitter.currentResult));
   }
 
private:
   /// externally owned buffers.
   std::vector<uint8_t> currentResult;
   std::vector<ArrayRef<uint8_t>> prevResultList;
   std::vector<std::vector<uint8_t>> prevResultStorage;
 
   /// An up-to-date total size of all of the buffers within `prevResultList`.
   /// This enables O(1) size checks of the current encoding.
   size_t prevResultSize = 0;
+
+  /// The highest required alignment for the start of this section.
+  unsigned requiredAlignment = 1;
 };
 
 /// A simple raw_ostream wrapper around a EncodingEmitter. This removes the need
 /// to go through an intermediate buffer when interacting with code that wants a
 /// raw_ostream.
 class RawEmitterOstream : public raw_ostream {
 public:
   explicit RawEmitterOstream(EncodingEmitter &emitter) : emitter(emitter) {

 //===----------------------------------------------------------------------===//
 
 namespace {
 class BytecodeWriter {
 public:
   BytecodeWriter(Operation *op) : numberingState(op) {}
 
   /// Write the bytecode for the given root operation.
-  void write(Operation *rootOp, raw_ostream &os, StringRef producer);
+  void write(Operation *rootOp, raw_ostream &os,
+             const BytecodeWriterConfig::Impl &config);
 
 private:
   //===--------------------------------------------------------------------===//
   // Dialects
 
   void writeDialectSection(EncodingEmitter &emitter);
 
   //===--------------------------------------------------------------------===//
   // Attributes and Types
 
   void writeAttrTypeSection(EncodingEmitter &emitter);
 
   //===--------------------------------------------------------------------===//
   // Operations
 
   void writeBlock(EncodingEmitter &emitter, Block *block);
   void writeOp(EncodingEmitter &emitter, Operation *op);
   void writeRegion(EncodingEmitter &emitter, Region *region);
   void writeIRSection(EncodingEmitter &emitter, Operation *op);
 
   //===--------------------------------------------------------------------===//
+  // Resources
+
+  void writeResourceSection(Operation *op, EncodingEmitter &emitter,
+                            const BytecodeWriterConfig::Impl &config);
+
+  //===--------------------------------------------------------------------===//
   // Strings
 
   void writeStringSection(EncodingEmitter &emitter);
 
   //===--------------------------------------------------------------------===//
   // Fields
 
   /// The builder used for the string section.
   StringSectionBuilder stringSection;
 
   /// The IR numbering state generated for the root operation.
   IRNumberingState numberingState;
 };
 } // namespace
 
 void BytecodeWriter::write(Operation *rootOp, raw_ostream &os,
-                           StringRef producer) {
+                           const BytecodeWriterConfig::Impl &config) {
   EncodingEmitter emitter;
 
   // Emit the bytecode file header. This is how we identify the output as a
   // bytecode file.
   emitter.emitString("ML\xefR");
 
   // Emit the bytecode version.
   emitter.emitVarInt(bytecode::kVersion);
 
   // Emit the producer.
-  emitter.emitNulTerminatedString(producer);
+  emitter.emitNulTerminatedString(config.producer);
 
   // Emit the dialect section.
   writeDialectSection(emitter);
 
   // Emit the attributes and types section.
   writeAttrTypeSection(emitter);
 
   // Emit the IR section.
   writeIRSection(emitter, rootOp);
 
+  // Emit the resources section.
+  writeResourceSection(rootOp, emitter, config);
+
   // Emit the string section.
   writeStringSection(emitter);
 
   // Write the generated bytecode to the provided output stream.
   emitter.writeTo(os);
 }
 
 //===----------------------------------------------------------------------===//
public:
 
   void writeAttribute(Attribute attr) override {
     emitter.emitVarInt(numberingState.getNumber(attr));
   }
   void writeType(Type type) override {
     emitter.emitVarInt(numberingState.getNumber(type));
   }
 
+  void writeResourceHandle(const AsmDialectResourceHandle &resource) override {
+    emitter.emitVarInt(numberingState.getNumber(resource));
+  }
+
   //===--------------------------------------------------------------------===//
   // Primitives
   //===--------------------------------------------------------------------===//
 
   void writeVarInt(uint64_t value) override { emitter.emitVarInt(value); }
 
   void writeSignedVarInt(int64_t value) override {
     emitter.emitSignedVarInt(value);
void BytecodeWriter::writeIRSection(EncodingEmitter &emitter, Operation *op) {
 
   // Emit the operations.
   writeOp(irEmitter, op);
 
   emitter.emitSection(bytecode::Section::kIR, std::move(irEmitter));
 }
 
 //===----------------------------------------------------------------------===//
+// Resources
+
+namespace {
+/// This class represents a resource builder implementation for the MLIR
+/// bytecode format.
+class ResourceBuilder : public AsmResourceBuilder {
+public:
+  using PostProcessFn = function_ref<void(StringRef, AsmResourceEntryKind)>;
+
+  ResourceBuilder(EncodingEmitter &emitter, StringSectionBuilder &stringSection,
+                  PostProcessFn postProcessFn)
+      : emitter(emitter), stringSection(stringSection),
+        postProcessFn(postProcessFn) {}
+  ~ResourceBuilder() override = default;
+
+  void buildBlob(StringRef key, ArrayRef<char> data,
+                 uint32_t dataAlignment) final {
+    emitter.emitBlobAndAlignment(data, dataAlignment);
+    postProcessFn(key, AsmResourceEntryKind::Blob);
+  }
+  void buildBool(StringRef key, bool data) final {
+    emitter.emitByte(data);
+    postProcessFn(key, AsmResourceEntryKind::Bool);
+  }
+  void buildString(StringRef key, StringRef data) final {
+    emitter.emitVarInt(stringSection.insert(data));
+    postProcessFn(key, AsmResourceEntryKind::String);
+  }
+
+private:
+  EncodingEmitter &emitter;
+  StringSectionBuilder &stringSection;
+  PostProcessFn postProcessFn;
+};
+} // namespace
+
+void BytecodeWriter::writeResourceSection(
+    Operation *op, EncodingEmitter &emitter,
+    const BytecodeWriterConfig::Impl &config) {
+  EncodingEmitter resourceEmitter;
+  EncodingEmitter resourceOffsetEmitter;
+  uint64_t prevOffset = 0;
+  SmallVector<std::tuple<StringRef, AsmResourceEntryKind, uint64_t>>
+      curResourceEntries;
+
+  // Functor used to process the offset for a resource of `kind` defined by
+  // 'key'.
+  auto appendResourceOffset = [&](StringRef key, AsmResourceEntryKind kind) {
+    uint64_t curOffset = resourceEmitter.size();
+    curResourceEntries.emplace_back(key, kind, curOffset - prevOffset);
+    prevOffset = curOffset;
+  };
+
+  // Functor used to emit a resource group defined by 'key'.
+  auto emitResourceGroup = [&](uint64_t key) {
+    resourceOffsetEmitter.emitVarInt(key);
+    resourceOffsetEmitter.emitVarInt(curResourceEntries.size());
+    for (auto [key, kind, size] : curResourceEntries) {
+      resourceOffsetEmitter.emitVarInt(stringSection.insert(key));
+      resourceOffsetEmitter.emitVarInt(size);
+      resourceOffsetEmitter.emitByte(kind);
+    }
+  };
+
+  // Builder used to emit resources.
+  ResourceBuilder entryBuilder(resourceEmitter, stringSection,
+                               appendResourceOffset);
+
+  // Emit the external resource entries.
+  resourceOffsetEmitter.emitVarInt(config.externalResourcePrinters.size());
+  for (const auto &printer : config.externalResourcePrinters) {
+    curResourceEntries.clear();
+    printer->buildResources(op, entryBuilder);
+    emitResourceGroup(stringSection.insert(printer->getName()));
+  }
+
+  // Emit the dialect resource entries.
+  for (DialectNumbering &dialect : numberingState.getDialects()) {
+    if (!dialect.asmInterface)
+      continue;
+    curResourceEntries.clear();
+    dialect.asmInterface->buildResources(op, dialect.resources, entryBuilder);
+
+    // Emit the declaration resources for this dialect, these didn't get emitted
+    // by the interface. These resources don't have data attached, so just use a
+    // "blob" kind as a placeholder.
+    for (const auto &resource : dialect.resourceMap)
+      if (resource.second->isDeclaration)
+        appendResourceOffset(resource.first, AsmResourceEntryKind::Blob);
+
+    // Emit the resource group for this dialect.
+    if (!curResourceEntries.empty())
+      emitResourceGroup(dialect.number);
+  }
+
+  // If we didn't emit any resource groups, elide the resource sections.
+  if (resourceOffsetEmitter.size() == 0)
+    return;
+
+  emitter.emitSection(bytecode::Section::kResourceOffset,
+                      std::move(resourceOffsetEmitter));
+  emitter.emitSection(bytecode::Section::kResource, std::move(resourceEmitter));
+}
+
+//===----------------------------------------------------------------------===//
 // Strings
 
 void BytecodeWriter::writeStringSection(EncodingEmitter &emitter) {
   EncodingEmitter stringEmitter;
   stringSection.write(stringEmitter);
   emitter.emitSection(bytecode::Section::kString, std::move(stringEmitter));
 }
 
 //===----------------------------------------------------------------------===//
 // Entry Points
 //===----------------------------------------------------------------------===//
 
 void mlir::writeBytecodeToFile(Operation *op, raw_ostream &os,
-                               StringRef producer) {
+                               const BytecodeWriterConfig &config) {
   BytecodeWriter writer(op);
-  writer.write(op, os, producer);
+  writer.write(op, os, config.getImpl());
 }

mlir/lib/Bytecode/Writer/IRNumbering.h

 //===- IRNumbering.h - MLIR bytecode IR numbering ---------------*- 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 contains various utilities that number IR structures in preparation
 // for bytecode emission.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LIB_MLIR_BYTECODE_WRITER_IRNUMBERING_H
 #define LIB_MLIR_BYTECODE_WRITER_IRNUMBERING_H
 
-#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/StringMap.h"
 
 namespace mlir {
 class BytecodeDialectInterface;
 class BytecodeWriterConfig;
 
 namespace bytecode {
 namespace detail {
 struct DialectNumbering;
struct OpNameNumbering {
   /// The number assigned to this name.
   unsigned number = 0;
 
   /// The number of references to this name.
   unsigned refCount = 1;
 };
 
 //===----------------------------------------------------------------------===//
+// Dialect Resource Numbering
+//===----------------------------------------------------------------------===//
+
+/// This class represents a numbering entry for a dialect resource.
+struct DialectResourceNumbering {
+  DialectResourceNumbering(std::string key) : key(std::move(key)) {}
+
+  /// The key used to reference this resource.
+  std::string key;
+
+  /// The number assigned to this resource.
+  unsigned number = 0;
+
+  /// A flag indicating if this resource is only a declaration, not a full
+  /// definition.
+  bool isDeclaration = true;
+};
+
+//===----------------------------------------------------------------------===//
 // Dialect Numbering
 //===----------------------------------------------------------------------===//
 
 /// This class represents a numbering entry for an Dialect.
 struct DialectNumbering {
   DialectNumbering(StringRef name, unsigned number)
       : name(name), number(number) {}
 
   /// The namespace of the dialect.
   StringRef name;
 
   /// The number assigned to the dialect.
   unsigned number;
 
   /// The bytecode dialect interface of the dialect if defined.
   const BytecodeDialectInterface *interface = nullptr;
+
+  /// The asm dialect interface of the dialect if defined.
+  const OpAsmDialectInterface *asmInterface = nullptr;
+
+  /// The referenced resources of this dialect.
+  SetVector<AsmDialectResourceHandle> resources;
+
+  /// A mapping from resource key to the corresponding resource numbering entry.
+  llvm::MapVector<StringRef, DialectResourceNumbering *> resourceMap;
 };
 
 //===----------------------------------------------------------------------===//
 // IRNumberingState
 //===----------------------------------------------------------------------===//
 
 /// This class manages numbering IR entities in preparation of bytecode
 /// emission.
public:
   unsigned getNumber(Type type) {
     assert(types.count(type) && "type not numbered");
     return types[type]->number;
   }
   unsigned getNumber(Value value) {
     assert(valueIDs.count(value) && "value not numbered");
     return valueIDs[value];
   }
+  unsigned getNumber(const AsmDialectResourceHandle &resource) {
+    assert(dialectResources.count(resource) && "resource not numbered");
+    return dialectResources[resource]->number;
+  }
 
   /// Return the block and value counts of the given region.
   std::pair<unsigned, unsigned> getBlockValueCount(Region *region) {
     assert(regionBlockValueCounts.count(region) && "value not numbered");
     return regionBlockValueCounts[region];
   }
 
   /// Return the number of operations in the given block.
private:
   void number(Block &block);
   DialectNumbering &numberDialect(Dialect *dialect);
   DialectNumbering &numberDialect(StringRef dialect);
   void number(Operation &op);
   void number(OperationName opName);
   void number(Region &region);
   void number(Type type);
 
+  /// Number the given dialect resources.
+  void number(Dialect *dialect, ArrayRef<AsmDialectResourceHandle> resources);
+
+  /// Finalize the numberings of any dialect resources.
+  void finalizeDialectResourceNumberings(Operation *rootOp);
+
   /// Mapping from IR to the respective numbering entries.
   DenseMap<Attribute, AttributeNumbering *> attrs;
   DenseMap<OperationName, OpNameNumbering *> opNames;
   DenseMap<Type, TypeNumbering *> types;
   DenseMap<Dialect *, DialectNumbering *> registeredDialects;
   llvm::MapVector<StringRef, DialectNumbering *> dialects;
   std::vector<AttributeNumbering *> orderedAttrs;
   std::vector<OpNameNumbering *> orderedOpNames;
   std::vector<TypeNumbering *> orderedTypes;
 
+  /// A mapping from dialect resource handle to the numbering for the referenced
+  /// resource.
+  llvm::DenseMap<AsmDialectResourceHandle, DialectResourceNumbering *>
+      dialectResources;
+
   /// Allocators used for the various numbering entries.
   llvm::SpecificBumpPtrAllocator<AttributeNumbering> attrAllocator;
   llvm::SpecificBumpPtrAllocator<DialectNumbering> dialectAllocator;
   llvm::SpecificBumpPtrAllocator<OpNameNumbering> opNameAllocator;
+  llvm::SpecificBumpPtrAllocator<DialectResourceNumbering> resourceAllocator;
   llvm::SpecificBumpPtrAllocator<TypeNumbering> typeAllocator;
 
   /// The value ID for each Block and Value.
   DenseMap<Block *, unsigned> blockIDs;
   DenseMap<Value, unsigned> valueIDs;
 
   /// The number of operations in each block.
   DenseMap<Block *, unsigned> blockOperationCounts;

mlir/lib/Bytecode/Writer/IRNumbering.cpp

 //===- IRNumbering.cpp - MLIR Bytecode IR numbering -----------------------===//
 //
 // 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 "IRNumbering.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Bytecode/BytecodeWriter.h"
+#include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 
 using namespace mlir;
 using namespace mlir::bytecode::detail;
 
 //===----------------------------------------------------------------------===//
 // NumberingDialectWriter
 //===----------------------------------------------------------------------===//
 
 struct IRNumberingState::NumberingDialectWriter : public DialectBytecodeWriter {
   NumberingDialectWriter(IRNumberingState &state) : state(state) {}
 
   void writeAttribute(Attribute attr) override { state.number(attr); }
   void writeType(Type type) override { state.number(type); }
+  void writeResourceHandle(const AsmDialectResourceHandle &resource) override {
+    state.number(resource.getDialect(), resource);
+  }
 
   /// Stubbed out methods that are not used for numbering.
   void writeVarInt(uint64_t) override {}
   void writeSignedVarInt(int64_t value) override {}
   void writeAPIntWithKnownWidth(const APInt &value) override {}
   void writeAPFloatWithKnownSemantics(const APFloat &value) override {}
   void writeOwnedString(StringRef) override {
     // TODO: It might be nice to prenumber strings and sort by the number of
IRNumberingState::IRNumberingState(Operation *op) {
   // After that, we apply a secondary ordering based on the parent dialect. This
   // ordering is applied to sub-sections of the element list defined by how many
   // bytes it takes to encode a varint index to that sub-section. This allows
   // for more efficiently encoding components of the same dialect (e.g. we only
   // have to encode the dialect reference once).
   groupByDialectPerByte(llvm::makeMutableArrayRef(orderedAttrs));
   groupByDialectPerByte(llvm::makeMutableArrayRef(orderedOpNames));
   groupByDialectPerByte(llvm::makeMutableArrayRef(orderedTypes));
+
+  // Finalize the numbering of the dialect resources.
+  finalizeDialectResourceNumberings(op);
 }
 
 void IRNumberingState::number(Attribute attr) {
   auto it = attrs.insert({attr, nullptr});
   if (!it.second) {
     ++it.first->second->refCount;
     return;
   }
if (OpaqueAttr opaqueAttr = attr.dyn_cast<OpaqueAttr>()) {
     return;
   }
   numbering->dialect = &numberDialect(&attr.getDialect());
 
   // If this attribute will be emitted using the bytecode format, perform a
   // dummy writing to number any nested components.
   if (const auto *interface = numbering->dialect->interface) {
     // TODO: We don't allow custom encodings for mutable attributes right now.
-    if (attr.hasTrait<AttributeTrait::IsMutable>())
-      return;
-
+    if (!attr.hasTrait<AttributeTrait::IsMutable>()) {
       NumberingDialectWriter writer(*this);
-    (void)interface->writeAttribute(attr, writer);
+      if (succeeded(interface->writeAttribute(attr, writer)))
+        return;
+    }
   }
+  // If this attribute will be emitted using the fallback, number the nested
+  // dialect resources. We don't number everything (e.g. no nested
+  // attributes/types), because we don't want to encode things we won't decode
+  // (the textual format can't really share much).
+  AsmState tempState(attr.getContext());
+  llvm::raw_null_ostream dummyOS;
+  attr.print(dummyOS, tempState);
+
+  // Number the used dialect resources.
+  for (const auto &it : tempState.getDialectResources())
+    number(it.getFirst(), it.getSecond().getArrayRef());
 }
 
 void IRNumberingState::number(Block &block) {
   // Number the arguments of the block.
   for (BlockArgument arg : block.getArguments()) {
     valueIDs.try_emplace(arg, nextValueID++);
     number(arg.getLoc());
     number(arg.getType());
   }
 
   // Number the operations in this block.
   unsigned &numOps = blockOperationCounts[&block];
   for (Operation &op : block) {
     number(op);
     ++numOps;
   }
 }
 
 auto IRNumberingState::numberDialect(Dialect *dialect) -> DialectNumbering & {
   DialectNumbering *&numbering = registeredDialects[dialect];
   if (!numbering) {
     numbering = &numberDialect(dialect->getNamespace());
     numbering->interface = dyn_cast<BytecodeDialectInterface>(dialect);
+    numbering->asmInterface = dyn_cast<OpAsmDialectInterface>(dialect);
   }
   return *numbering;
 }
 
 auto IRNumberingState::numberDialect(StringRef dialect) -> DialectNumbering & {
   DialectNumbering *&numbering = dialects[dialect];
   if (!numbering) {
     numbering = new (dialectAllocator.Allocate())
if (OpaqueType opaqueType = type.dyn_cast<OpaqueType>()) {
     return;
   }
   numbering->dialect = &numberDialect(&type.getDialect());
 
   // If this type will be emitted using the bytecode format, perform a dummy
   // writing to number any nested components.
   if (const auto *interface = numbering->dialect->interface) {
     // TODO: We don't allow custom encodings for mutable types right now.
-    if (type.hasTrait<TypeTrait::IsMutable>())
+    if (!type.hasTrait<TypeTrait::IsMutable>()) {
+      NumberingDialectWriter writer(*this);
+      if (succeeded(interface->writeType(type, writer)))
+        return;
+    }
+  }
+  // If this type will be emitted using the fallback, number the nested dialect
+  // resources. We don't number everything (e.g. no nested attributes/types),
+  // because we don't want to encode things we won't decode (the textual format
+  // can't really share much).
+  AsmState tempState(type.getContext());
+  llvm::raw_null_ostream dummyOS;
+  type.print(dummyOS, tempState);
+
+  // Number the used dialect resources.
+  for (const auto &it : tempState.getDialectResources())
+    number(it.getFirst(), it.getSecond().getArrayRef());
+}
+
+void IRNumberingState::number(Dialect *dialect,
+                              ArrayRef<AsmDialectResourceHandle> resources) {
+  DialectNumbering &dialectNumber = numberDialect(dialect);
+  assert(
+      dialectNumber.asmInterface &&
+      "expected dialect owning a resource to implement OpAsmDialectInterface");
+
+  for (const auto &resource : resources) {
+    // Check if this is a newly seen resource.
+    if (!dialectNumber.resources.insert(resource))
       return;
 
-    NumberingDialectWriter writer(*this);
-    (void)interface->writeType(type, writer);
+    auto *numbering =
+        new (resourceAllocator.Allocate()) DialectResourceNumbering(
+            dialectNumber.asmInterface->getResourceKey(resource));
+    dialectNumber.resourceMap.insert({numbering->key, numbering});
+    dialectResources.try_emplace(resource, numbering);
+  }
+}
+
+namespace {
+/// A dummy resource builder used to number dialect resources.
+struct NumberingResourceBuilder : public AsmResourceBuilder {
+  NumberingResourceBuilder(DialectNumbering *dialect, unsigned &nextResourceID)
+      : dialect(dialect), nextResourceID(nextResourceID) {}
+  ~NumberingResourceBuilder() override = default;
+
+  void buildBlob(StringRef key, ArrayRef<char>, uint32_t) final {
+    numberEntry(key);
+  }
+  void buildBool(StringRef key, bool) final { numberEntry(key); }
+  void buildString(StringRef key, StringRef) final {
+    // TODO: We could pre-number the value string here as well.
+    numberEntry(key);
+  }
+
+  /// Number the dialect entry for the given key.
+  void numberEntry(StringRef key) {
+    // TODO: We could pre-number resource key strings here as well.
+
+    auto it = dialect->resourceMap.find(key);
+    if (it != dialect->resourceMap.end()) {
+      it->second->number = nextResourceID++;
+      it->second->isDeclaration = false;
+    }
+  }
+
+  DialectNumbering *dialect;
+  unsigned &nextResourceID;
+};
+} // namespace
+
+void IRNumberingState::finalizeDialectResourceNumberings(Operation *rootOp) {
+  unsigned nextResourceID = 0;
+  for (DialectNumbering &dialect : getDialects()) {
+    if (!dialect.asmInterface)
+      continue;
+    NumberingResourceBuilder entryBuilder(&dialect, nextResourceID);
+    dialect.asmInterface->buildResources(rootOp, dialect.resources,
+                                         entryBuilder);
+
+    // Number any resources that weren't added by the dialect. This can happen
+    // if there was no backing data to the resource, but we still want these
+    // resource references to roundtrip, so we number them and indicate that the
+    // data is missing.
+    for (const auto &it : dialect.resourceMap)
+      if (it.second->isDeclaration)
+        it.second->number = nextResourceID++;
   }
 }

mlir/lib/IR/AsmPrinter.cpp

 // Resources
 //===----------------------------------------------------------------------===//
 
 AsmParsedResourceEntry::~AsmParsedResourceEntry() = default;
 AsmResourceBuilder::~AsmResourceBuilder() = default;
 AsmResourceParser::~AsmResourceParser() = default;
 AsmResourcePrinter::~AsmResourcePrinter() = default;
 
+StringRef mlir::toString(AsmResourceEntryKind kind) {
+  switch (kind) {
+  case AsmResourceEntryKind::Blob:
+    return "blob";
+  case AsmResourceEntryKind::Bool:
+    return "bool";
+  case AsmResourceEntryKind::String:
+    return "string";
+  }
+  llvm_unreachable("unknown AsmResourceEntryKind");
+}
+
 //===----------------------------------------------------------------------===//
 // AsmState
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace detail {
 class AsmStateImpl {
 public:

mlir/lib/IR/BuiltinDialectBytecode.cpp

 //===- BuiltinDialectBytecode.cpp - Builtin Bytecode Implementation -------===//
 //
 // 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 "BuiltinDialectBytecode.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/DialectResourceBlobManager.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Encoding
 //===----------------------------------------------------------------------===//
 
enum AttributeCode {
   ///     name: StringAttr,
   ///     childLoc: LocationAttr
   ///   }
   kNameLoc = 14,
 
   ///   UnknownLoc {
   ///   }
   kUnknownLoc = 15,
+
+  ///   DenseResourceElementsAttr {
+  ///     type: Type,
+  ///     handle: ResourceHandle
+  ///   }
+  kDenseResourceElementsAttr = 16,
 };
 
 /// This enum contains marker codes used to indicate which type is currently
 /// being decoded, and how it should be decoded. The order of these codes should
 /// generally be unchanged, as any changes will inevitably break compatibility
 /// with older bytecode.
 enum TypeCode {
   ///   IntegerType {
struct BuiltinDialectBytecodeInterface : public BytecodeDialectInterface {
   BuiltinDialectBytecodeInterface(Dialect *dialect)
       : BytecodeDialectInterface(dialect) {}
 
   //===--------------------------------------------------------------------===//
   // Attributes
 
   Attribute readAttribute(DialectBytecodeReader &reader) const override;
   ArrayAttr readArrayAttr(DialectBytecodeReader &reader) const;
+  DenseResourceElementsAttr
+  readDenseResourceElementsAttr(DialectBytecodeReader &reader) const;
   DictionaryAttr readDictionaryAttr(DialectBytecodeReader &reader) const;
   FloatAttr readFloatAttr(DialectBytecodeReader &reader) const;
   IntegerAttr readIntegerAttr(DialectBytecodeReader &reader) const;
   StringAttr readStringAttr(DialectBytecodeReader &reader, bool hasType) const;
   SymbolRefAttr readSymbolRefAttr(DialectBytecodeReader &reader,
                                   bool hasNestedRefs) const;
   TypeAttr readTypeAttr(DialectBytecodeReader &reader) const;
 
   LocationAttr readCallSiteLoc(DialectBytecodeReader &reader) const;
   LocationAttr readFileLineColLoc(DialectBytecodeReader &reader) const;
   LocationAttr readFusedLoc(DialectBytecodeReader &reader,
                             bool hasMetadata) const;
   LocationAttr readNameLoc(DialectBytecodeReader &reader) const;
 
   LogicalResult writeAttribute(Attribute attr,
                                DialectBytecodeWriter &writer) const override;
   void write(ArrayAttr attr, DialectBytecodeWriter &writer) const;
+  void write(DenseResourceElementsAttr attr,
+             DialectBytecodeWriter &writer) const;
   void write(DictionaryAttr attr, DialectBytecodeWriter &writer) const;
   void write(IntegerAttr attr, DialectBytecodeWriter &writer) const;
   void write(FloatAttr attr, DialectBytecodeWriter &writer) const;
   void write(StringAttr attr, DialectBytecodeWriter &writer) const;
   void write(SymbolRefAttr attr, DialectBytecodeWriter &writer) const;
   void write(TypeAttr attr, DialectBytecodeWriter &writer) const;
 
   void write(CallSiteLoc attr, DialectBytecodeWriter &writer) const;
Attribute BuiltinDialectBytecodeInterface::readAttribute(
   case builtin_encoding::kFusedLoc:
     return readFusedLoc(reader, /*hasMetadata=*/false);
   case builtin_encoding::kFusedLocWithMetadata:
     return readFusedLoc(reader, /*hasMetadata=*/true);
   case builtin_encoding::kNameLoc:
     return readNameLoc(reader);
   case builtin_encoding::kUnknownLoc:
     return UnknownLoc::get(getContext());
+  case builtin_encoding::kDenseResourceElementsAttr:
+    return readDenseResourceElementsAttr(reader);
   default:
     reader.emitError() << "unknown builtin attribute code: " << code;
     return Attribute();
   }
 }
 
 LogicalResult BuiltinDialectBytecodeInterface::writeAttribute(
     Attribute attr, DialectBytecodeWriter &writer) const {
   return TypeSwitch<Attribute, LogicalResult>(attr)
-      .Case<ArrayAttr, DictionaryAttr, FloatAttr, IntegerAttr, StringAttr,
-            SymbolRefAttr, TypeAttr, CallSiteLoc, FileLineColLoc, FusedLoc,
-            NameLoc>([&](auto attr) {
+      .Case<ArrayAttr, DenseResourceElementsAttr, DictionaryAttr, FloatAttr,
+            IntegerAttr, StringAttr, SymbolRefAttr, TypeAttr>([&](auto attr) {
+        write(attr, writer);
+        return success();
+      })
+      .Case<CallSiteLoc, FileLineColLoc, FusedLoc, NameLoc>([&](auto attr) {
         write(attr, writer);
         return success();
       })
       .Case([&](OpaqueLoc attr) { return write(attr, writer); })
       .Case([&](UnitAttr) {
         writer.writeVarInt(builtin_encoding::kUnitAttr);
         return success();
       })

 
 void BuiltinDialectBytecodeInterface::write(
     ArrayAttr attr, DialectBytecodeWriter &writer) const {
   writer.writeVarInt(builtin_encoding::kArrayAttr);
   writer.writeAttributes(attr.getValue());
 }
 
 //===----------------------------------------------------------------------===//
+// DenseResourceElementsAttr
+
+DenseResourceElementsAttr
+BuiltinDialectBytecodeInterface::readDenseResourceElementsAttr(
+    DialectBytecodeReader &reader) const {
+  ShapedType type;
+  if (failed(reader.readType(type)))
+    return DenseResourceElementsAttr();
+
+  FailureOr<DenseResourceElementsHandle> handle =
+      reader.readResourceHandle<DenseResourceElementsHandle>();
+  if (failed(handle))
+    return DenseResourceElementsAttr();
+
+  return DenseResourceElementsAttr::get(type, *handle);
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    DenseResourceElementsAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kDenseResourceElementsAttr);
+  writer.writeType(attr.getType());
+  writer.writeResourceHandle(attr.getRawHandle());
+}
+
+//===----------------------------------------------------------------------===//
 // DictionaryAttr
 
 DictionaryAttr BuiltinDialectBytecodeInterface::readDictionaryAttr(
     DialectBytecodeReader &reader) const {
   auto readNamedAttr = [&]() -> FailureOr<NamedAttribute> {
     StringAttr name;
     Attribute value;
     if (failed(reader.readAttribute(name)) ||

mlir/test/Bytecode/invalid/invalid-structure.mlir

 
 // RUN: not mlir-opt %S/invalid-structure-section-missing.mlirbc 2>&1 | FileCheck %s --check-prefix=SECTION_MISSING
 // SECTION_MISSING: missing data for top-level section: String (0)
 
 //===--------------------------------------------------------------------===//
 // ID
 
 // RUN: not mlir-opt %S/invalid-structure-section-id-unknown.mlirbc 2>&1 | FileCheck %s --check-prefix=SECTION_ID_UNKNOWN
-// SECTION_ID_UNKNOWN: invalid section ID: 255
+// SECTION_ID_UNKNOWN: invalid section ID: 127
 
 //===--------------------------------------------------------------------===//
 // Length
 
 // RUN: not mlir-opt %S/invalid-structure-section-length.mlirbc 2>&1 | FileCheck %s --check-prefix=SECTION_LENGTH
 // SECTION_LENGTH: attempting to parse a byte at the end of the bytecode
 
 //===--------------------------------------------------------------------===//
 // Duplicate
 
 // RUN: not mlir-opt %S/invalid-structure-section-duplicate.mlirbc 2>&1 | FileCheck %s --check-prefix=SECTION_DUPLICATE
 // SECTION_DUPLICATE: duplicate top-level section: String (0)

mlir/test/Bytecode/resources.mlir

+// RUN: mlir-opt -emit-bytecode %s | mlir-opt | FileCheck %s
+
+// Bytecode currently does not support big-endian platforms
+// UNSUPPORTED: s390x-
+
+// CHECK-LABEL: @TestDialectResources
+module @TestDialectResources attributes {
+  // CHECK: bytecode.test = dense_resource<decl_resource> : tensor<2xui32>
+  // CHECK: bytecode.test2 = dense_resource<resource> : tensor<4xf64>
+  // CHECK: bytecode.test3 = dense_resource<resource_2> : tensor<4xf64>
+  bytecode.test = dense_resource<decl_resource> : tensor<2xui32>,
+  bytecode.test2 = dense_resource<resource> : tensor<4xf64>,
+  bytecode.test3 = dense_resource<resource_2> : tensor<4xf64>
+} {}
+
+// CHECK: builtin: {
+// CHECK-NEXT: resource: "0x08000000010000000000000002000000000000000300000000000000"
+// CHECK-NEXT: resource_2: "0x08000000010000000000000002000000000000000300000000000000"
+
+{-#
+  dialect_resources: {
+    builtin: {
+      resource: "0x08000000010000000000000002000000000000000300000000000000",
+      resource_2: "0x08000000010000000000000002000000000000000300000000000000"
+    }
+  }
+#-}