[mlir] Add initial support for a binary serialization format

This commit adds a new bytecode serialization format for MLIR.
The actual serialization of MLIR to binary is relatively straightforward,
given the very very general structure of MLIR. The underlying basis for
this format is a variable-length encoding for integers, which gets heavily
used for nearly all aspects of the encoding (given that most of the encoding
is just indexing into lists).

The format currently does not provide support for custom attribute/type
serialization, and thus always uses an assembly format fallback. It also
doesn't provide support for resources. These will be added in followups,
the intention for this patch is to provide something that supports the
basic cases, and can be built on top of.

https://discourse.llvm.org/t/rfc-a-binary-serialization-format-for-mlir/63518

Differential Revision: https://reviews.llvm.org/D131747

mlir/docs/BytecodeFormat.md

@@ -0,0 +1,314 @@
+# MLIR Bytecode Format
+
+This documents describes the MLIR bytecode format and its encoding.
+
+[TOC]
+
+## Magic Number
+
+MLIR uses the following four-byte magic number to indicate bytecode files:
+
+'\[‘M’<sub>8</sub>, ‘L’<sub>8</sub>, ‘ï’<sub>8</sub>, ‘R’<sub>8</sub>\]'
+
+In hex:
+
+'\[‘4D’<sub>8</sub>, ‘4C’<sub>8</sub>, ‘EF’<sub>8</sub>, ‘52’<sub>8</sub>\]'
+
+## Format Overview
+
+An MLIR Bytecode file is comprised of a byte stream, with a few simple
+structural concepts layered on top.
+
+### Primitives
+
+#### Fixed-Width Integers
+
+```
+  byte                  ::= `0x00`...`0xFF`
+```
+
+Fixed width integers are unsigned integers of a known byte size. The values are
+stored in little-endian byte order.
+
+TODO: Add larger fixed width integers as necessary.
+
+#### Variable-Width Integers
+
+Variable width integers, or `VarInt`s, provide a compact representation for
+integers. Each encoded VarInt consists of one to nine bytes, which together
+represent a single 64-bit value. The MLIR bytecode utilizes the "PrefixVarInt"
+encoding for VarInts. This encoding is a variant of the
+[LEB128 ("Little-Endian Base 128")](https://en.wikipedia.org/wiki/LEB128)
+encoding, where each byte of the encoding provides up to 7 bits for the value,
+with the remaining bit used to store a tag indicating the number of bytes used
+for the encoding. This means that small unsigned integers (less than 2^7) may be
+stored in one byte, unsigned integers up to 2^14 may be stored in two bytes,
+etc.
+
+The first byte of the encoding includes a length prefix in the low bits. This
+prefix is a bit sequence of '0's followed by a terminal '1', or the end of the
+byte. The number of '0' bits indicate the number of _additional_ bytes, not
+including the prefix byte, used to encode the value. All of the remaining bits
+in the first byte, along with all of the bits in the additional bytes, provide
+the value of the integer. Below are the various possible encodings of the prefix
+byte:
+
+```
+xxxxxxx1:  7 value bits, the encoding uses 1 byte
+xxxxxx10: 14 value bits, the encoding uses 2 bytes
+xxxxx100: 21 value bits, the encoding uses 3 bytes
+xxxx1000: 28 value bits, the encoding uses 4 bytes
+xxx10000: 35 value bits, the encoding uses 5 bytes
+xx100000: 42 value bits, the encoding uses 6 bytes
+x1000000: 49 value bits, the encoding uses 7 bytes
+10000000: 56 value bits, the encoding uses 8 bytes
+00000000: 64 value bits, the encoding uses 9 bytes
+```
+
+#### Strings
+
+Strings are blobs of characters with an associated length.
+
+### Sections
+
+```
+section {
+  id: byte
+  length: varint
+}
+```
+
+Sections are a mechanism for grouping data within the bytecode. The 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.
+
+## 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
+
+The top-level structure of the bytecode contains the 4-byte "magic number", a
+version number, a null-terminated producer string, and a list of sections. Each
+section is currently only expected to appear once within a bytecode file.
+
+```
+bytecode {
+  magic: "MLïR",
+  version: varint,
+  producer: string,
+  sections: section[]
+}
+```
+
+### String Section
+
+```
+strings {
+  numStrings: varint,
+  reverseStringLengths: varint[],
+  stringData: byte[]
+}
+```
+
+The string section contains a table of strings referenced within the bytecode,
+more easily enabling string sharing. This section is encoded first with the
+total number of strings, followed by the sizes of each of the individual strings
+in reverse order. The remaining encoding contains a single blob containing all
+of the strings concatenated together.
+
+### Dialect Section
+
+The dialect section of the bytecode contains all of the dialects referenced
+within the encoded IR, and some information about the components of those
+dialects that were also referenced.
+
+```
+dialect_section {
+  numDialects: varint,
+  dialectNames: varint[],
+  opNames: op_name_group[]
+}
+
+op_name_group {
+  dialect: varint,
+  numOpNames: varint,
+  opNames: varint[]
+}
+```
+
+Dialects are encoded as indexes to the name string within the string section.
+Operation names are encoded in groups by dialect, with each group containing the
+dialect, the number of operation names, and the array of indexes to each name
+within the string section.
+
+### Attribute/Type Sections
+
+Attributes and types are encoded using two [sections](#sections), one section
+(`attr_type_section`) containing the actual encoded representation, and another
+section (`attr_type_offset_section`) containing the offsets of each encoded
+attribute/type into the previous section. This structure allows for attributes
+and types to always be lazily loaded on demand.
+
+```
+attr_type_section {
+  attrs: attribute[],
+  types: type[]  
+}
+attr_type_offset_section {
+  numAttrs: varint,
+  numTypes: varint,
+  offsets: attr_type_offset_group[]
+}
+
+attr_type_offset_group {
+  dialect: varint,
+  numElements: varint,
+  offsets: varint[] // (offset << 1) | (hasCustomEncoding) 
+}
+
+attribute {
+  encoding: ...
+}
+type {
+  encoding: ...
+}
+```
+
+Each `offset` in the `attr_type_offset_section` above is the size of the
+encoding for the attribute or type and a flag indicating if the encoding uses
+the textual assembly format, or a custom bytecode encoding. We avoid using the
+direct offset into the `attr_type_section`, as a smaller relative offsets
+provides more effective compression. Attributes and types are grouped by
+dialect, with each `attr_type_offset_group` in the offset section containing the
+corresponding parent dialect, number of elements, and offsets for each element
+within the group.
+
+#### Attribute/Type Encodings
+
+In the abstract, an attribute/type is encoded in one of two possible ways: via
+its assembly format, or via a custom dialect defined encoding.
+
+##### Assembly Format Fallback
+
+In the case where a dialect does not define a method for encoding the attribute
+or type, the textual assembly format of that attribute or type is used as a
+fallback. For example, a type of `!bytecode.type` would be encoded as the null
+terminated string "!bytecode.type". This ensures that every attribute and type
+may be encoded, even if the owning dialect has not yet opted in to a more
+efficient serialization.
+
+TODO: We shouldn't redundantly encode the dialect name here, we should use a
+reference to the parent dialect instead.
+
+##### Dialect Defined Encoding
+
+TODO: This is not yet supported.
+
+### IR Section
+
+The IR section contains the encoded form of operations within the bytecode.
+
+#### Operation Encoding
+
+```
+op {
+  name: varint,
+  encodingMask: byte,
+  location: varint,
+
+  attrDict: varint?,
+
+  numResults: varint?,
+  resultTypes: varint[],
+
+  numOperands: varint?,
+  operands: varint[],
+
+  numSuccessors: varint?,
+  successors: varint[],
+
+  regionEncoding: varint?, // (numRegions << 1) | (isIsolatedFromAbove) 
+  regions: region[]
+}
+```
+
+The encoding of an operation is important because this is generally the most
+commonly appearing structure in the bytecode. A single encoding is used for
+every type of operation. Given this prevelance, many of the fields of an
+operation are optional. The `encodingMask` field is a bitmask which indicates
+which of the components of the operation are present.
+
+##### Location
+
+The location is encoded as the index of the location within the attribute table.
+
+##### Attributes
+
+If the operation has attribues, the index of the operation attribute dictionary
+within the attribute table is encoded.
+
+##### Results
+
+If the operation has results, the number of results and the indexes of the
+result types within the type table are encoded.
+
+##### Operands
+
+If the operation has operands, the number of operands and the value index of
+each operand is encoded. This value index is the relative ordering of the
+definition of that value from the start of the first ancestor isolated region.
+
+##### Successors
+
+If the operation has successors, the number of successors and the indexes of the
+successor blocks within the parent region are encoded.
+
+##### Regions
+
+If the operation has regions, the number of regions and if the regions are
+isolated from above are encoded together in a single varint. Afterwards, each
+region is encoded inline.
+
+#### Region Encoding
+
+```
+region {
+  numBlocks: varint,
+
+  numValues: varint?,
+  blocks: block[]
+}
+```
+
+A region is encoded first with the number of blocks within. If the region is
+non-empty, the number of values defined directly within the region are encoded,
+followed by the blocks of the region.
+
+#### Block Encoding
+
+```
+block {
+  encoding: varint, // (numOps << 1) | (hasBlockArgs)
+  arguments: block_arguments?, // Optional based on encoding
+  ops : op[]
+}
+
+block_arguments {
+  numArgs: varint?,
+  args: block_argument[]
+}
+
+block_argument {
+  typeIndex: varint,
+  location: varint
+}
+```
+
+A block is encoded with an array of operations and block arguments. The first
+field is an encoding that combines the number of operations in the block, with a
+flag indicating if the block has arguments.

mlir/include/mlir/Bytecode/BytecodeReader.h

@@ -0,0 +1,34 @@
+//===- BytecodeReader.h - MLIR Bytecode Reader ------------------*- 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 read MLIR bytecode files/streams.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_BYTECODE_BYTECODEREADER_H
+#define MLIR_BYTECODE_BYTECODEREADER_H
+
+#include "mlir/IR/AsmState.h"
+#include "mlir/Support/LLVM.h"
+
+namespace llvm {
+class MemoryBufferRef;
+} // namespace llvm
+
+namespace mlir {
+/// Returns true if the given buffer starts with the magic bytes that signal
+/// MLIR bytecode.
+bool isBytecode(llvm::MemoryBufferRef buffer);
+
+/// Read the operations defined within the given memory buffer, containing MLIR
+/// bytecode, into the provided block.
+LogicalResult readBytecodeFile(llvm::MemoryBufferRef buffer, Block *block,
+                               const ParserConfig &config);
+} // namespace mlir
+
+#endif // MLIR_BYTECODE_BYTECODEREADER_H

mlir/include/mlir/Bytecode/BytecodeWriter.h

@@ -0,0 +1,36 @@
+//===- 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"
+
+namespace mlir {
+class Operation;
+
+//===----------------------------------------------------------------------===//
+// 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);
+
+} // namespace mlir
+
+#endif // MLIR_BYTECODE_BYTECODEWRITER_H

mlir/include/mlir/IR/OperationSupport.h

@@ -642,11 +642,11 @@ public:
   OperationState(Location location, OperationName name);
 
   OperationState(Location location, OperationName name, ValueRange operands,
-                 TypeRange types, ArrayRef<NamedAttribute> attributes,
+                 TypeRange types, ArrayRef<NamedAttribute> attributes = {},
                  BlockRange successors = {},
                  MutableArrayRef<std::unique_ptr<Region>> regions = {});
   OperationState(Location location, StringRef name, ValueRange operands,
-                 TypeRange types, ArrayRef<NamedAttribute> attributes,
+                 TypeRange types, ArrayRef<NamedAttribute> attributes = {},
                  BlockRange successors = {},
                  MutableArrayRef<std::unique_ptr<Region>> regions = {});
 

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

@@ -50,13 +50,15 @@ using PassPipelineFn = llvm::function_ref<LogicalResult(PassManager &pm)>;
 /// - preloadDialectsInContext will trigger the upfront loading of all
 ///   dialects from the global registry in the MLIRContext. This option is
 ///   deprecated and will be removed soon.
+/// - emitBytecode will generate bytecode output instead of text.
 LogicalResult MlirOptMain(llvm::raw_ostream &outputStream,
                           std::unique_ptr<llvm::MemoryBuffer> buffer,
                           const PassPipelineCLParser &passPipeline,
                           DialectRegistry &registry, bool splitInputFile,
                           bool verifyDiagnostics, bool verifyPasses,
                           bool allowUnregisteredDialects,
-                          bool preloadDialectsInContext = false);
+                          bool preloadDialectsInContext = false,
+                          bool emitBytecode = false);
 
 /// Support a callback to setup the pass manager.
 /// - passManagerSetupFn is the callback invoked to setup the pass manager to
@@ -67,7 +69,8 @@ LogicalResult MlirOptMain(llvm::raw_ostream &outputStream,
                           DialectRegistry &registry, bool splitInputFile,
                           bool verifyDiagnostics, bool verifyPasses,
                           bool allowUnregisteredDialects,
-                          bool preloadDialectsInContext = false);
+                          bool preloadDialectsInContext = false,
+                          bool emitBytecode = false);
 
 /// Implementation for tools like `mlir-opt`.
 /// - toolName is used for the header displayed by `--help`.

mlir/lib/Bytecode/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_subdirectory(Reader)
+add_subdirectory(Writer)

mlir/lib/Bytecode/Encoding.h

@@ -0,0 +1,81 @@
+//===- Encoding.h - MLIR binary format encoding information -----*- 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 enum values describing the structure of MLIR bytecode
+// files.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LIB_MLIR_BYTECODE_ENCODING_H
+#define LIB_MLIR_BYTECODE_ENCODING_H
+
+#include <cstdint>
+
+namespace mlir {
+namespace bytecode {
+//===----------------------------------------------------------------------===//
+// General constants
+//===----------------------------------------------------------------------===//
+
+enum {
+  /// The current bytecode version.
+  kVersion = 0,
+};
+
+//===----------------------------------------------------------------------===//
+// Sections
+//===----------------------------------------------------------------------===//
+
+namespace Section {
+enum ID : uint8_t {
+  /// This section contains strings referenced within the bytecode.
+  kString = 0,
+
+  /// This section contains the dialects referenced within an IR module.
+  kDialect = 1,
+
+  /// This section contains the attributes and types referenced within an IR
+  /// module.
+  kAttrType = 2,
+
+  /// 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,
+
+  /// The total number of section types.
+  kNumSections = 5,
+};
+} // namespace Section
+
+//===----------------------------------------------------------------------===//
+// IR Section
+//===----------------------------------------------------------------------===//
+
+/// This enum represents a mask of all of the potential components of an
+/// operation. This mask is used when encoding an operation to indicate which
+/// components are present in the bytecode.
+namespace OpEncodingMask {
+enum : uint8_t {
+  // clang-format off
+  kHasAttrs         = 0b00000001,
+  kHasResults       = 0b00000010,
+  kHasOperands      = 0b00000100,
+  kHasSuccessors    = 0b00001000,
+  kHasInlineRegions = 0b00010000,
+  // clang-format on
+};
+} // namespace OpEncodingMask
+
+} // namespace bytecode
+} // namespace mlir
+
+#endif

mlir/lib/Bytecode/Reader/CMakeLists.txt

@@ -0,0 +1,11 @@
+add_mlir_library(MLIRBytecodeReader
+  BytecodeReader.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Bytecode
+
+  LINK_LIBS PUBLIC
+  MLIRAsmParser
+  MLIRIR
+  MLIRSupport
+  )

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

@@ -0,0 +1,520 @@
+//===- BytecodeWriter.cpp - MLIR Bytecode Writer --------------------------===//
+//
+// 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/Bytecode/BytecodeWriter.h"
+#include "../Encoding.h"
+#include "IRNumbering.h"
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/OpImplementation.h"
+#include "llvm/ADT/CachedHashString.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/Support/Debug.h"
+#include <random>
+
+#define DEBUG_TYPE "mlir-bytecode-writer"
+
+using namespace mlir;
+using namespace mlir::bytecode::detail;
+
+//===----------------------------------------------------------------------===//
+// 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
+/// emission.
+class EncodingEmitter {
+public:
+  EncodingEmitter() = default;
+  EncodingEmitter(const EncodingEmitter &) = delete;
+  EncodingEmitter &operator=(const EncodingEmitter &) = delete;
+
+  /// Write the current contents to the provided stream.
+  void writeTo(raw_ostream &os) const;
+
+  /// Return the current size of the encoded buffer.
+  size_t size() const { return prevResultSize + currentResult.size(); }
+
+  //===--------------------------------------------------------------------===//
+  // Emission
+  //===--------------------------------------------------------------------===//
+
+  /// 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;
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Integer Emission
+
+  /// Emit a single byte.
+  template <typename T>
+  void emitByte(T byte) {
+    currentResult.push_back(static_cast<uint8_t>(byte));
+  }
+
+  /// Emit a range of bytes.
+  void emitBytes(ArrayRef<uint8_t> bytes) {
+    llvm::append_range(currentResult, bytes);
+  }
+
+  /// Emit a variable length integer. 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.
+  void emitVarInt(uint64_t value) {
+    // In the most common case, the value can be represented in a single byte.
+    // Given how hot this case is, explicitly handle that here.
+    if ((value >> 7) == 0)
+      return emitByte((value << 1) | 0x1);
+    emitMultiByteVarInt(value);
+  }
+
+  /// Emit a variable length integer whose low bit is used to encode the
+  /// provided flag, i.e. encoded as: (value << 1) | (flag ? 1 : 0).
+  void emitVarIntWithFlag(uint64_t value, bool flag) {
+    emitVarInt((value << 1) | (flag ? 1 : 0));
+  }
+
+  //===--------------------------------------------------------------------===//
+  // String Emission
+
+  /// Emit the given string as a nul terminated string.
+  void emitNulTerminatedString(StringRef str) {
+    emitString(str);
+    emitByte(0);
+  }
+
+  /// Emit the given string without a nul terminator.
+  void emitString(StringRef str) {
+    emitBytes({reinterpret_cast<const uint8_t *>(str.data()), str.size()});
+  }
+
+  //===--------------------------------------------------------------------===//
+  // 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.
+    emitByte(code);
+    emitVarInt(emitter.size());
+
+    // 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:
+  /// Emit the given value using a variable width encoding. This method is a
+  /// fallback when the number of bytes needed to encode the value is greater
+  /// than 1. We mark it noinline here so that the single byte hot path isn't
+  /// pessimized.
+  LLVM_ATTRIBUTE_NOINLINE void emitMultiByteVarInt(uint64_t value);
+
+  /// Append a new result buffer to the current contents.
+  void appendResult(std::vector<uint8_t> &&result) {
+    prevResultSize += result.size();
+    prevResultStorage.emplace_back(std::move(result));
+    prevResultList.emplace_back(prevResultStorage.back());
+  }
+
+  /// The result of the emitter currently being built. We refrain from building
+  /// a single buffer to simplify emitting sections, large data, and more. The
+  /// result is thus represented using multiple distinct buffers, some of which
+  /// we own (via prevResultStorage), and some of which are just pointers into
+  /// 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;
+};
+
+/// 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 raw_emitter_ostream : public raw_ostream {
+public:
+  explicit raw_emitter_ostream(EncodingEmitter &emitter) : emitter(emitter) {
+    SetUnbuffered();
+  }
+
+private:
+  void write_impl(const char *ptr, size_t size) override {
+    emitter.emitBytes({reinterpret_cast<const uint8_t *>(ptr), size});
+  }
+  uint64_t current_pos() const override { return emitter.size(); }
+
+  /// The section being emitted to.
+  EncodingEmitter &emitter;
+};
+} // namespace
+
+void EncodingEmitter::writeTo(raw_ostream &os) const {
+  for (auto &prevResult : prevResultList)
+    os.write((const char *)prevResult.data(), prevResult.size());
+  os.write((const char *)currentResult.data(), currentResult.size());
+}
+
+void EncodingEmitter::emitMultiByteVarInt(uint64_t value) {
+  // Compute the number of bytes needed to encode the value. Each byte can hold
+  // up to 7-bits of data. We only check up to the number of bits we can encode
+  // in the first byte (8).
+  uint64_t it = value >> 7;
+  for (size_t numBytes = 2; numBytes < 9; ++numBytes) {
+    if (LLVM_LIKELY(it >>= 7) == 0) {
+      uint64_t encodedValue = (value << 1) | 0x1;
+      encodedValue <<= (numBytes - 1);
+      emitBytes({reinterpret_cast<uint8_t *>(&encodedValue), numBytes});
+      return;
+    }
+  }
+
+  // If the value is too large to encode in a single byte, emit a special all
+  // zero marker byte and splat the value directly.
+  emitByte(0);
+  emitBytes({reinterpret_cast<uint8_t *>(&value), sizeof(value)});
+}
+
+//===----------------------------------------------------------------------===//
+// Bytecode Writer
+//===----------------------------------------------------------------------===//
+
+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);
+
+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);
+
+  //===--------------------------------------------------------------------===//
+  // Strings
+
+  void writeStringSection(EncodingEmitter &emitter);
+
+  /// Get the number for the given shared string, that is contained within the
+  /// string section.
+  size_t getSharedStringNumber(StringRef str);
+
+  //===--------------------------------------------------------------------===//
+  // Fields
+
+  /// The IR numbering state generated for the root operation.
+  IRNumberingState numberingState;
+
+  /// A set of strings referenced within the bytecode. The value of the map is
+  /// unused.
+  llvm::MapVector<llvm::CachedHashStringRef, size_t> strings;
+};
+} // namespace
+
+void BytecodeWriter::write(Operation *rootOp, raw_ostream &os,
+                           StringRef producer) {
+  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);
+
+  // Emit the dialect section.
+  writeDialectSection(emitter);
+
+  // Emit the attributes and types section.
+  writeAttrTypeSection(emitter);
+
+  // Emit the IR section.
+  writeIRSection(emitter, rootOp);
+
+  // Emit the string section.
+  writeStringSection(emitter);
+
+  // Write the generated bytecode to the provided output stream.
+  emitter.writeTo(os);
+}
+
+//===----------------------------------------------------------------------===//
+// Dialects
+
+/// Write the given entries in contiguous groups with the same parent dialect.
+/// Each dialect sub-group is encoded with the parent dialect and number of
+/// elements, followed by the encoding for the entries. The given callback is
+/// invoked to encode each individual entry.
+template <typename EntriesT, typename EntryCallbackT>
+static void writeDialectGrouping(EncodingEmitter &emitter, EntriesT &&entries,
+                                 EntryCallbackT &&callback) {
+  for (auto it = entries.begin(), e = entries.end(); it != e;) {
+    auto groupStart = it++;
+
+    // Find the end of the group that shares the same parent dialect.
+    DialectNumbering *currentDialect = groupStart->dialect;
+    it = std::find_if(it, e, [&](const auto &entry) {
+      return entry.dialect != currentDialect;
+    });
+
+    // Emit the dialect and number of elements.
+    emitter.emitVarInt(currentDialect->number);
+    emitter.emitVarInt(std::distance(groupStart, it));
+
+    // Emit the entries within the group.
+    for (auto &entry : llvm::make_range(groupStart, it))
+      callback(entry);
+  }
+}
+
+void BytecodeWriter::writeDialectSection(EncodingEmitter &emitter) {
+  EncodingEmitter dialectEmitter;
+
+  // Emit the referenced dialects.
+  auto dialects = numberingState.getDialects();
+  dialectEmitter.emitVarInt(llvm::size(dialects));
+  for (DialectNumbering &dialect : dialects)
+    dialectEmitter.emitVarInt(getSharedStringNumber(dialect.name));
+
+  // Emit the referenced operation names grouped by dialect.
+  auto emitOpName = [&](OpNameNumbering &name) {
+    dialectEmitter.emitVarInt(getSharedStringNumber(name.name.stripDialect()));
+  };
+  writeDialectGrouping(dialectEmitter, numberingState.getOpNames(), emitOpName);
+
+  emitter.emitSection(bytecode::Section::kDialect, std::move(dialectEmitter));
+}
+
+//===----------------------------------------------------------------------===//
+// Attributes and Types
+
+void BytecodeWriter::writeAttrTypeSection(EncodingEmitter &emitter) {
+  EncodingEmitter attrTypeEmitter;
+  EncodingEmitter offsetEmitter;
+  offsetEmitter.emitVarInt(llvm::size(numberingState.getAttributes()));
+  offsetEmitter.emitVarInt(llvm::size(numberingState.getTypes()));
+
+  // A functor used to emit an attribute or type entry.
+  uint64_t prevOffset = 0;
+  auto emitAttrOrType = [&](auto &entry) {
+    // TODO: Allow dialects to provide more optimal implementations of attribute
+    // and type encodings.
+    bool hasCustomEncoding = false;
+
+    // Emit the entry using the textual format.
+    raw_emitter_ostream(attrTypeEmitter) << entry.getValue();
+    attrTypeEmitter.emitByte(0);
+
+    // Record the offset of this entry.
+    uint64_t curOffset = attrTypeEmitter.size();
+    offsetEmitter.emitVarIntWithFlag(curOffset - prevOffset, hasCustomEncoding);
+    prevOffset = curOffset;
+  };
+
+  // Emit the attribute and type entries for each dialect.
+  writeDialectGrouping(offsetEmitter, numberingState.getAttributes(),
+                       emitAttrOrType);
+  writeDialectGrouping(offsetEmitter, numberingState.getTypes(),
+                       emitAttrOrType);
+
+  // Emit the sections to the stream.
+  emitter.emitSection(bytecode::Section::kAttrTypeOffset,
+                      std::move(offsetEmitter));
+  emitter.emitSection(bytecode::Section::kAttrType, std::move(attrTypeEmitter));
+}
+
+//===----------------------------------------------------------------------===//
+// Operations
+
+void BytecodeWriter::writeBlock(EncodingEmitter &emitter, Block *block) {
+  ArrayRef<BlockArgument> args = block->getArguments();
+  bool hasArgs = !args.empty();
+
+  // Emit the number of operations in this block, and if it has arguments. We
+  // use the low bit of the operation count to indicate if the block has
+  // arguments.
+  unsigned numOps = numberingState.getOperationCount(block);
+  emitter.emitVarIntWithFlag(numOps, hasArgs);
+
+  // Emit the arguments of the block.
+  if (hasArgs) {
+    emitter.emitVarInt(args.size());
+    for (BlockArgument arg : args) {
+      emitter.emitVarInt(numberingState.getNumber(arg.getType()));
+      emitter.emitVarInt(numberingState.getNumber(arg.getLoc()));
+    }
+  }
+
+  // Emit the operations within the block.
+  for (Operation &op : *block)
+    writeOp(emitter, &op);
+}
+
+void BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
+  emitter.emitVarInt(numberingState.getNumber(op->getName()));
+
+  // Emit a mask for the operation components. We need to fill this in later
+  // (when we actually know what needs to be emitted), so emit a placeholder for
+  // now.
+  uint64_t maskOffset = emitter.size();
+  uint8_t opEncodingMask = 0;
+  emitter.emitByte(0);
+
+  // Emit the location for this operation.
+  emitter.emitVarInt(numberingState.getNumber(op->getLoc()));
+
+  // Emit the attributes of this operation.
+  DictionaryAttr attrs = op->getAttrDictionary();
+  if (!attrs.empty()) {
+    opEncodingMask |= bytecode::OpEncodingMask::kHasAttrs;
+    emitter.emitVarInt(numberingState.getNumber(op->getAttrDictionary()));
+  }
+
+  // Emit the result types of the operation.
+  if (unsigned numResults = op->getNumResults()) {
+    opEncodingMask |= bytecode::OpEncodingMask::kHasResults;
+    emitter.emitVarInt(numResults);
+    for (Type type : op->getResultTypes())
+      emitter.emitVarInt(numberingState.getNumber(type));
+  }
+
+  // Emit the operands of the operation.
+  if (unsigned numOperands = op->getNumOperands()) {
+    opEncodingMask |= bytecode::OpEncodingMask::kHasOperands;
+    emitter.emitVarInt(numOperands);
+    for (Value operand : op->getOperands())
+      emitter.emitVarInt(numberingState.getNumber(operand));
+  }
+
+  // Emit the successors of the operation.
+  if (unsigned numSuccessors = op->getNumSuccessors()) {
+    opEncodingMask |= bytecode::OpEncodingMask::kHasSuccessors;
+    emitter.emitVarInt(numSuccessors);
+    for (Block *successor : op->getSuccessors())
+      emitter.emitVarInt(numberingState.getNumber(successor));
+  }
+
+  // Check for regions.
+  unsigned numRegions = op->getNumRegions();
+  if (numRegions)
+    opEncodingMask |= bytecode::OpEncodingMask::kHasInlineRegions;
+
+  // Update the mask for the operation.
+  emitter.patchByte(maskOffset, opEncodingMask);
+
+  // With the mask emitted, we can now emit the regions of the operation. We do
+  // this after mask emission to avoid offset complications that may arise by
+  // emitting the regions first (e.g. if the regions are huge, backpatching the
+  // op encoding mask is more annoying).
+  if (numRegions) {
+    bool isIsolatedFromAbove = op->hasTrait<OpTrait::IsIsolatedFromAbove>();
+    emitter.emitVarIntWithFlag(numRegions, isIsolatedFromAbove);
+
+    for (Region &region : op->getRegions())
+      writeRegion(emitter, &region);
+  }
+}
+
+void BytecodeWriter::writeRegion(EncodingEmitter &emitter, Region *region) {
+  // If the region is empty, we only need to emit the number of blocks (which is
+  // zero).
+  if (region->empty())
+    return emitter.emitVarInt(/*numBlocks*/ 0);
+
+  // Emit the number of blocks and values within the region.
+  unsigned numBlocks, numValues;
+  std::tie(numBlocks, numValues) = numberingState.getBlockValueCount(region);
+  emitter.emitVarInt(numBlocks);
+  emitter.emitVarInt(numValues);
+
+  // Emit the blocks within the region.
+  for (Block &block : *region)
+    writeBlock(emitter, &block);
+}
+
+void BytecodeWriter::writeIRSection(EncodingEmitter &emitter, Operation *op) {
+  EncodingEmitter irEmitter;
+
+  // Write the IR section the same way as a block with no arguments. Note that
+  // the low-bit of the operation count for a block is used to indicate if the
+  // block has arguments, which in this case is always false.
+  irEmitter.emitVarIntWithFlag(/*numOps*/ 1, /*hasArgs*/ false);
+
+  // Emit the operations.
+  writeOp(irEmitter, op);
+
+  emitter.emitSection(bytecode::Section::kIR, std::move(irEmitter));
+}
+
+//===----------------------------------------------------------------------===//
+// Strings
+
+void BytecodeWriter::writeStringSection(EncodingEmitter &emitter) {
+  EncodingEmitter stringEmitter;
+  stringEmitter.emitVarInt(strings.size());
+
+  // Emit the sizes in reverse order, so that we don't need to backpatch an
+  // offset to the string data or have a separate section.
+  for (const auto &it : llvm::reverse(strings))
+    stringEmitter.emitVarInt(it.first.size() + 1);
+  // Emit the string data itself.
+  for (const auto &it : strings)
+    stringEmitter.emitNulTerminatedString(it.first.val());
+
+  emitter.emitSection(bytecode::Section::kString, std::move(stringEmitter));
+}
+
+size_t BytecodeWriter::getSharedStringNumber(StringRef str) {
+  auto it = strings.insert({llvm::CachedHashStringRef(str), strings.size()});
+  return it.first->second;
+}
+
+//===----------------------------------------------------------------------===//
+// Entry Points
+//===----------------------------------------------------------------------===//
+
+void mlir::writeBytecodeToFile(Operation *op, raw_ostream &os,
+                               StringRef producer) {
+  BytecodeWriter writer(op);
+  writer.write(op, os, producer);
+}

mlir/lib/Bytecode/Writer/CMakeLists.txt

@@ -0,0 +1,11 @@
+add_mlir_library(MLIRBytecodeWriter
+  BytecodeWriter.cpp
+  IRNumbering.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Bytecode
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRSupport
+  )

mlir/lib/Bytecode/Writer/IRNumbering.cpp

@@ -0,0 +1,251 @@
+//===- 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/BytecodeWriter.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OpDefinition.h"
+
+using namespace mlir;
+using namespace mlir::bytecode::detail;
+
+//===----------------------------------------------------------------------===//
+// IR Numbering
+//===----------------------------------------------------------------------===//
+
+/// Group and sort the elements of the given range by their parent dialect. This
+/// grouping is applied to sub-sections of the ranged defined by how many bytes
+/// it takes to encode a varint index to that sub-section.
+template <typename T>
+static void groupByDialectPerByte(T range) {
+  if (range.empty())
+    return;
+
+  // A functor used to sort by a given dialect, with a desired dialect to be
+  // ordered first (to better enable sharing of dialects across byte groups).
+  auto sortByDialect = [](unsigned dialectToOrderFirst, const auto &lhs,
+                          const auto &rhs) {
+    if (lhs->dialect->number == dialectToOrderFirst)
+      return rhs->dialect->number != dialectToOrderFirst;
+    return lhs->dialect->number < rhs->dialect->number;
+  };
+
+  unsigned dialectToOrderFirst = 0;
+  size_t elementsInByteGroup = 0;
+  auto iterRange = range;
+  for (unsigned i = 1; i < 9; ++i) {
+    // Update the number of elements in the current byte grouping. Reminder
+    // that varint encodes 7-bits per byte, so that's how we compute the
+    // number of elements in each byte grouping.
+    elementsInByteGroup = (1 << (7 * i)) - elementsInByteGroup;
+
+    // Slice out the sub-set of elements that are in the current byte grouping
+    // to be sorted.
+    auto byteSubRange = iterRange.take_front(elementsInByteGroup);
+    iterRange = iterRange.drop_front(byteSubRange.size());
+
+    // Sort the sub range for this byte.
+    llvm::stable_sort(byteSubRange, [&](const auto &lhs, const auto &rhs) {
+      return sortByDialect(dialectToOrderFirst, lhs, rhs);
+    });
+
+    // Update the dialect to order first to be the dialect at the end of the
+    // current grouping. This seeks to allow larger dialect groupings across
+    // byte boundaries.
+    dialectToOrderFirst = byteSubRange.back()->dialect->number;
+
+    // If the data range is now empty, we are done.
+    if (iterRange.empty())
+      break;
+  }
+
+  // Assign the entry numbers based on the sort order.
+  for (auto &entry : llvm::enumerate(range))
+    entry.value()->number = entry.index();
+}
+
+IRNumberingState::IRNumberingState(Operation *op) {
+  // Number the root operation.
+  number(*op);
+
+  // Push all of the regions of the root operation onto the worklist.
+  SmallVector<std::pair<Region *, unsigned>, 8> numberContext;
+  for (Region &region : op->getRegions())
+    numberContext.emplace_back(&region, nextValueID);
+
+  // Iteratively process each of the nested regions.
+  while (!numberContext.empty()) {
+    Region *region;
+    std::tie(region, nextValueID) = numberContext.pop_back_val();
+    number(*region);
+
+    // Traverse into nested regions.
+    for (Operation &op : region->getOps()) {
+      // Isolated regions don't share value numbers with their parent, so we can
+      // start numbering these regions at zero.
+      unsigned opFirstValueID =
+          op.hasTrait<OpTrait::IsIsolatedFromAbove>() ? 0 : nextValueID;
+      for (Region &region : op.getRegions())
+        numberContext.emplace_back(&region, opFirstValueID);
+    }
+  }
+
+  // Number each of the dialects. For now this is just in the order they were
+  // found, given that the number of dialects on average is small enough to fit
+  // within a singly byte (128). If we ever have real world use cases that have
+  // a huge number of dialects, this could be made more intelligent.
+  for (auto &it : llvm::enumerate(dialects))
+    it.value().second->number = it.index();
+
+  // Number each of the recorded components within each dialect.
+
+  // First sort by ref count so that the most referenced elements are first. We
+  // try to bias more heavily used elements to the front. This allows for more
+  // frequently referenced things to be encoded using smaller varints.
+  auto sortByRefCountFn = [](const auto &lhs, const auto &rhs) {
+    return lhs->refCount > rhs->refCount;
+  };
+  llvm::stable_sort(orderedAttrs, sortByRefCountFn);
+  llvm::stable_sort(orderedOpNames, sortByRefCountFn);
+  llvm::stable_sort(orderedTypes, sortByRefCountFn);
+
+  // 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));
+}
+
+void IRNumberingState::number(Attribute attr) {
+  auto it = attrs.insert({attr, nullptr});
+  if (!it.second) {
+    ++it.first->second->refCount;
+    return;
+  }
+  auto *numbering = new (attrAllocator.Allocate()) AttributeNumbering(attr);
+  it.first->second = numbering;
+  orderedAttrs.push_back(numbering);
+
+  // Check for OpaqueAttr, which is a dialect-specific attribute that didn't
+  // have a registered dialect when it got created. We don't want to encode this
+  // as the builtin OpaqueAttr, we want to encode it as if the dialect was
+  // actually loaded.
+  if (OpaqueAttr opaqueAttr = attr.dyn_cast<OpaqueAttr>())
+    numbering->dialect = &numberDialect(opaqueAttr.getDialectNamespace());
+  else
+    numbering->dialect = &numberDialect(&attr.getDialect());
+}
+
+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->dialect = dialect;
+  }
+  return *numbering;
+}
+
+auto IRNumberingState::numberDialect(StringRef dialect) -> DialectNumbering & {
+  DialectNumbering *&numbering = dialects[dialect];
+  if (!numbering) {
+    numbering = new (dialectAllocator.Allocate())
+        DialectNumbering(dialect, dialects.size() - 1);
+  }
+  return *numbering;
+}
+
+void IRNumberingState::number(Region &region) {
+  if (region.empty())
+    return;
+  size_t firstValueID = nextValueID;
+
+  // Number the blocks within this region.
+  size_t blockCount = 0;
+  for (auto &it : llvm::enumerate(region)) {
+    blockIDs.try_emplace(&it.value(), it.index());
+    number(it.value());
+    ++blockCount;
+  }
+
+  // Remember the number of blocks and values in this region.
+  regionBlockValueCounts.try_emplace(&region, blockCount,
+                                     nextValueID - firstValueID);
+}
+
+void IRNumberingState::number(Operation &op) {
+  // Number the components of an operation that won't be numbered elsewhere
+  // (e.g. we don't number operands, regions, or successors here).
+  number(op.getName());
+  for (OpResult result : op.getResults()) {
+    valueIDs.try_emplace(result, nextValueID++);
+    number(result.getType());
+  }
+
+  // Only number the operation's dictionary if it isn't empty.
+  DictionaryAttr dictAttr = op.getAttrDictionary();
+  if (!dictAttr.empty())
+    number(dictAttr);
+
+  number(op.getLoc());
+}
+
+void IRNumberingState::number(OperationName opName) {
+  OpNameNumbering *&numbering = opNames[opName];
+  if (numbering) {
+    ++numbering->refCount;
+    return;
+  }
+  DialectNumbering *dialectNumber = nullptr;
+  if (Dialect *dialect = opName.getDialect())
+    dialectNumber = &numberDialect(dialect);
+  else
+    dialectNumber = &numberDialect(opName.getDialectNamespace());
+
+  numbering =
+      new (opNameAllocator.Allocate()) OpNameNumbering(dialectNumber, opName);
+  orderedOpNames.push_back(numbering);
+}
+
+void IRNumberingState::number(Type type) {
+  auto it = types.insert({type, nullptr});
+  if (!it.second) {
+    ++it.first->second->refCount;
+    return;
+  }
+  auto *numbering = new (typeAllocator.Allocate()) TypeNumbering(type);
+  it.first->second = numbering;
+  orderedTypes.push_back(numbering);
+
+  // Check for OpaqueType, which is a dialect-specific type that didn't have a
+  // registered dialect when it got created. We don't want to encode this as the
+  // builtin OpaqueType, we want to encode it as if the dialect was actually
+  // loaded.
+  if (OpaqueType opaqueType = type.dyn_cast<OpaqueType>())
+    numbering->dialect = &numberDialect(opaqueType.getDialectNamespace());
+  else
+    numbering->dialect = &numberDialect(&type.getDialect());
+}

mlir/lib/Bytecode/Writer/IRNumbering.h

@@ -0,0 +1,193 @@
+//===- 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 "llvm/ADT/MapVector.h"
+
+namespace mlir {
+class BytecodeWriterConfig;
+
+namespace bytecode {
+namespace detail {
+struct DialectNumbering;
+
+//===----------------------------------------------------------------------===//
+// Attribute and Type Numbering
+//===----------------------------------------------------------------------===//
+
+/// This class represents a numbering entry for an Attribute or Type.
+struct AttrTypeNumbering {
+  AttrTypeNumbering(PointerUnion<Attribute, Type> value) : value(value) {}
+
+  /// The concrete value.
+  PointerUnion<Attribute, Type> value;
+
+  /// The number assigned to this value.
+  unsigned number = 0;
+
+  /// The number of references to this value.
+  unsigned refCount = 1;
+
+  /// The dialect of this value.
+  DialectNumbering *dialect = nullptr;
+};
+struct AttributeNumbering : public AttrTypeNumbering {
+  AttributeNumbering(Attribute value) : AttrTypeNumbering(value) {}
+  Attribute getValue() const { return value.get<Attribute>(); }
+};
+struct TypeNumbering : public AttrTypeNumbering {
+  TypeNumbering(Type value) : AttrTypeNumbering(value) {}
+  Type getValue() const { return value.get<Type>(); }
+};
+
+//===----------------------------------------------------------------------===//
+// OpName Numbering
+//===----------------------------------------------------------------------===//
+
+/// This class represents the numbering entry of an operation name.
+struct OpNameNumbering {
+  OpNameNumbering(DialectNumbering *dialect, OperationName name)
+      : dialect(dialect), name(name) {}
+
+  /// The dialect of this value.
+  DialectNumbering *dialect;
+
+  /// The concrete name.
+  OperationName name;
+
+  /// The number assigned to this name.
+  unsigned number = 0;
+
+  /// The number of references to this name.
+  unsigned refCount = 1;
+};
+
+//===----------------------------------------------------------------------===//
+// 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 loaded dialect, or nullptr if the dialect isn't loaded.
+  Dialect *dialect = nullptr;
+};
+
+//===----------------------------------------------------------------------===//
+// IRNumberingState
+//===----------------------------------------------------------------------===//
+
+/// This class manages numbering IR entities in preparation of bytecode
+/// emission.
+class IRNumberingState {
+public:
+  IRNumberingState(Operation *op);
+
+  /// Return the numbered dialects.
+  auto getDialects() {
+    return llvm::make_pointee_range(llvm::make_second_range(dialects));
+  }
+  auto getAttributes() { return llvm::make_pointee_range(orderedAttrs); }
+  auto getOpNames() { return llvm::make_pointee_range(orderedOpNames); }
+  auto getTypes() { return llvm::make_pointee_range(orderedTypes); }
+
+  /// Return the number for the given IR unit.
+  unsigned getNumber(Attribute attr) {
+    assert(attrs.count(attr) && "attribute not numbered");
+    return attrs[attr]->number;
+  }
+  unsigned getNumber(Block *block) {
+    assert(blockIDs.count(block) && "block not numbered");
+    return blockIDs[block];
+  }
+  unsigned getNumber(OperationName opName) {
+    assert(opNames.count(opName) && "opName not numbered");
+    return opNames[opName]->number;
+  }
+  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];
+  }
+
+  /// 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.
+  unsigned getOperationCount(Block *block) {
+    assert(blockOperationCounts.count(block) && "block not numbered");
+    return blockOperationCounts[block];
+  }
+
+private:
+  /// Number the given IR unit for bytecode emission.
+  void number(Attribute attr);
+  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);
+
+  /// 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;
+
+  /// Allocators used for the various numbering entries.
+  llvm::SpecificBumpPtrAllocator<AttributeNumbering> attrAllocator;
+  llvm::SpecificBumpPtrAllocator<DialectNumbering> dialectAllocator;
+  llvm::SpecificBumpPtrAllocator<OpNameNumbering> opNameAllocator;
+  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;
+
+  /// A map from region to the number of blocks and values within that region.
+  DenseMap<Region *, std::pair<unsigned, unsigned>> regionBlockValueCounts;
+
+  /// The next value ID to assign when numbering.
+  unsigned nextValueID = 0;
+};
+} // namespace detail
+} // namespace bytecode
+} // namespace mlir
+
+#endif

mlir/lib/CMakeLists.txt

@@ -3,6 +3,7 @@ add_flag_if_supported("-Werror=global-constructors" WERROR_GLOBAL_CONSTRUCTOR)
 
 add_subdirectory(Analysis)
 add_subdirectory(AsmParser)
+add_subdirectory(Bytecode)
 add_subdirectory(Conversion)
 add_subdirectory(Dialect)
 add_subdirectory(IR)

mlir/lib/Parser/CMakeLists.txt

@@ -6,5 +6,6 @@ add_mlir_library(MLIRParser
 
   LINK_LIBS PUBLIC
   MLIRAsmParser
+  MLIRBytecodeReader
   MLIRIR
   )

mlir/lib/Parser/Parser.cpp

@@ -12,6 +12,7 @@
 
 #include "mlir/Parser/Parser.h"
 #include "mlir/AsmParser/AsmParser.h"
+#include "mlir/Bytecode/BytecodeReader.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
@@ -25,6 +26,8 @@ LogicalResult mlir::parseSourceFile(const llvm::SourceMgr &sourceMgr,
                                          sourceBuf->getBufferIdentifier(),
                                          /*line=*/0, /*column=*/0);
   }
+  if (isBytecode(*sourceBuf))
+    return readBytecodeFile(*sourceBuf, block, config);
   return parseAsmSourceFile(sourceMgr, block, config);
 }
 

mlir/lib/Tools/mlir-opt/CMakeLists.txt

@@ -5,6 +5,7 @@ add_mlir_library(MLIROptLib
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/mlir-opt
 
   LINK_LIBS PUBLIC
+  MLIRBytecodeWriter
   MLIRPass
   MLIRParser
   MLIRSupport

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
+#include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -47,7 +48,8 @@ using namespace llvm;
 static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
                                     bool verifyPasses, SourceMgr &sourceMgr,
                                     MLIRContext *context,
-                                    PassPipelineFn passManagerSetupFn) {
+                                    PassPipelineFn passManagerSetupFn,
+                                    bool emitBytecode) {
   DefaultTimingManager tm;
   applyDefaultTimingManagerCLOptions(tm);
   TimingScope timing = tm.getRootScope();
@@ -86,8 +88,12 @@ static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
 
   // Print the output.
   TimingScope outputTiming = timing.nest("Output");
+  if (emitBytecode) {
+    writeBytecodeToFile(module->getOperation(), os);
+  } else {
     module->print(os);
     os << '\n';
+  }
   return success();
 }
 
@@ -97,8 +103,8 @@ static LogicalResult
 processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
               bool verifyDiagnostics, bool verifyPasses,
               bool allowUnregisteredDialects, bool preloadDialectsInContext,
-              PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
-              llvm::ThreadPool *threadPool) {
+              bool emitBytecode, PassPipelineFn passManagerSetupFn,
+              DialectRegistry &registry, llvm::ThreadPool *threadPool) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.
   SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
@@ -122,7 +128,7 @@ processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
   if (!verifyDiagnostics) {
     SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
     return performActions(os, verifyDiagnostics, verifyPasses, sourceMgr,
-                          &context, passManagerSetupFn);
+                          &context, passManagerSetupFn, emitBytecode);
   }
 
   SourceMgrDiagnosticVerifierHandler sourceMgrHandler(sourceMgr, &context);
@@ -131,7 +137,7 @@ processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
   // these actions succeed or fail, we only care what diagnostics they produce
   // and whether they match our expectations.
   (void)performActions(os, verifyDiagnostics, verifyPasses, sourceMgr, &context,
-                       passManagerSetupFn);
+                       passManagerSetupFn, emitBytecode);
 
   // Verify the diagnostic handler to make sure that each of the diagnostics
   // matched.
@@ -144,7 +150,8 @@ LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                                 DialectRegistry &registry, bool splitInputFile,
                                 bool verifyDiagnostics, bool verifyPasses,
                                 bool allowUnregisteredDialects,
-                                bool preloadDialectsInContext) {
+                                bool preloadDialectsInContext,
+                                bool emitBytecode) {
   // The split-input-file mode is a very specific mode that slices the file
   // up into small pieces and checks each independently.
   // We use an explicit threadpool to avoid creating and joining/destroying
@@ -163,8 +170,8 @@ LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                      raw_ostream &os) {
     return processBuffer(os, std::move(chunkBuffer), verifyDiagnostics,
                          verifyPasses, allowUnregisteredDialects,
-                         preloadDialectsInContext, passManagerSetupFn, registry,
-                         threadPool);
+                         preloadDialectsInContext, emitBytecode,
+                         passManagerSetupFn, registry, threadPool);
   };
   return splitAndProcessBuffer(std::move(buffer), chunkFn, outputStream,
                                splitInputFile, /*insertMarkerInOutput=*/true);
@@ -176,7 +183,8 @@ LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                                 DialectRegistry &registry, bool splitInputFile,
                                 bool verifyDiagnostics, bool verifyPasses,
                                 bool allowUnregisteredDialects,
-                                bool preloadDialectsInContext) {
+                                bool preloadDialectsInContext,
+                                bool emitBytecode) {
   auto passManagerSetupFn = [&](PassManager &pm) {
     auto errorHandler = [&](const Twine &msg) {
       emitError(UnknownLoc::get(pm.getContext())) << msg;
@@ -186,7 +194,8 @@ LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
   };
   return MlirOptMain(outputStream, std::move(buffer), passManagerSetupFn,
                      registry, splitInputFile, verifyDiagnostics, verifyPasses,
-                     allowUnregisteredDialects, preloadDialectsInContext);
+                     allowUnregisteredDialects, preloadDialectsInContext,
+                     emitBytecode);
 }
 
 LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
@@ -224,6 +233,10 @@ LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
       "show-dialects", cl::desc("Print the list of registered dialects"),
       cl::init(false));
 
+  static cl::opt<bool> emitBytecode(
+      "emit-bytecode", cl::desc("Emit bytecode when generating output"),
+      cl::init(false));
+
   InitLLVM y(argc, argv);
 
   // Register any command line options.
@@ -268,7 +281,8 @@ LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
 
   if (failed(MlirOptMain(output->os(), std::move(file), passPipeline, registry,
                          splitInputFile, verifyDiagnostics, verifyPasses,
-                         allowUnregisteredDialects, preloadDialectsInContext)))
+                         allowUnregisteredDialects, preloadDialectsInContext,
+                         emitBytecode)))
     return failure();
 
   // Keep the output file if the invocation of MlirOptMain was successful.

mlir/test/Bytecode/general.mlir

@@ -0,0 +1,34 @@
+// RUN: mlir-opt -allow-unregistered-dialect -emit-bytecode %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
+
+// CHECK-LABEL: "bytecode.test1"
+// CHECK-NEXT:    "bytecode.empty"() : () -> ()
+// CHECK-NEXT:    "bytecode.attributes"() {attra = 10 : i64, attrb = #bytecode.attr} : () -> ()
+// CHECK-NEXT:    test.graph_region {
+// CHECK-NEXT:      "bytecode.operands"(%[[RESULTS:.*]]#0, %[[RESULTS]]#1, %[[RESULTS]]#2) : (i32, i64, i32) -> ()
+// CHECK-NEXT:      %[[RESULTS]]:3 = "bytecode.results"() : () -> (i32, i64, i32)
+// CHECK-NEXT:    }
+// CHECK-NEXT:    "bytecode.branch"()[^[[BLOCK:.*]]] : () -> ()
+// CHECK-NEXT:  ^[[BLOCK]](%[[ARG0:.*]]: i32, %[[ARG1:.*]]: !bytecode.int, %[[ARG2:.*]]: !pdl.operation):
+// CHECK-NEXT:    "bytecode.regions"() ({
+// CHECK-NEXT:      "bytecode.operands"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (i32, !bytecode.int, !pdl.operation) -> ()
+// CHECK-NEXT:      "bytecode.return"() : () -> ()
+// CHECK-NEXT:    }) : () -> ()
+// CHECK-NEXT:    "bytecode.return"() : () -> ()
+// CHECK-NEXT:  }) : () -> ()
+
+"bytecode.test1"() ({
+  "bytecode.empty"() : () -> ()
+  "bytecode.attributes"() {attra = 10, attrb = #bytecode.attr} : () -> ()
+  test.graph_region {
+    "bytecode.operands"(%results#0, %results#1, %results#2) : (i32, i64, i32) -> ()
+    %results:3 = "bytecode.results"() : () -> (i32, i64, i32)
+  }
+  "bytecode.branch"()[^secondBlock] : () -> ()
+
+^secondBlock(%arg1: i32, %arg2: !bytecode.int, %arg3: !pdl.operation):
+  "bytecode.regions"() ({
+    "bytecode.operands"(%arg1, %arg2, %arg3) : (i32, !bytecode.int, !pdl.operation) -> ()
+    "bytecode.return"() : () -> ()
+  }) : () -> ()
+  "bytecode.return"() : () -> ()
+}) : () -> ()

mlir/test/Bytecode/invalid/invalid-dialect_section.mlir

@@ -0,0 +1,19 @@
+// This file contains various failure test cases related to the structure of
+// the dialect section.
+
+//===--------------------------------------------------------------------===//
+// Dialect Name
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-dialect_section-dialect_string.mlirbc 2>&1 | FileCheck %s --check-prefix=DIALECT_STR
+// DIALECT_STR: invalid string index: 15
+
+//===--------------------------------------------------------------------===//
+// OpName
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-dialect_section-opname_dialect.mlirbc 2>&1 | FileCheck %s --check-prefix=OPNAME_DIALECT
+// OPNAME_DIALECT: invalid dialect index: 7
+
+// RUN: not mlir-opt %S/invalid-dialect_section-opname_string.mlirbc 2>&1 | FileCheck %s --check-prefix=OPNAME_STR
+// OPNAME_STR: invalid string index: 31

mlir/test/Bytecode/invalid/invalid-ir_section.mlir

@@ -0,0 +1,45 @@
+// This file contains various failure test cases related to the structure of
+// the IR section.
+
+//===--------------------------------------------------------------------===//
+// Operations
+//===--------------------------------------------------------------------===//
+
+//===--------------------------------------------------------------------===//
+// Name
+
+// RUN: not mlir-opt %S/invalid-ir_section-opname.mlirbc -allow-unregistered-dialect 2>&1 | FileCheck %s --check-prefix=OP_NAME
+// OP_NAME: invalid operation name index: 14
+
+//===--------------------------------------------------------------------===//
+// Loc
+
+// RUN: not mlir-opt %S/invalid-ir_section-loc.mlirbc -allow-unregistered-dialect 2>&1 | FileCheck %s --check-prefix=OP_LOC
+// OP_LOC: expected attribute of type: {{.*}}, but got: {attra = 10 : i64, attrb = #bytecode.attr}
+
+//===--------------------------------------------------------------------===//
+// Attr
+
+// RUN: not mlir-opt %S/invalid-ir_section-attr.mlirbc -allow-unregistered-dialect 2>&1 | FileCheck %s --check-prefix=OP_ATTR
+// OP_ATTR: expected attribute of type: {{.*}}, but got: loc(unknown)
+
+//===--------------------------------------------------------------------===//
+// Operands
+
+// RUN: not mlir-opt %S/invalid-ir_section-operands.mlirbc -allow-unregistered-dialect 2>&1 | FileCheck %s --check-prefix=OP_OPERANDS
+// OP_OPERANDS: invalid value index: 6
+
+// RUN: not mlir-opt %S/invalid-ir_section-forwardref.mlirbc -allow-unregistered-dialect 2>&1 | FileCheck %s --check-prefix=FORWARD_REF
+// FORWARD_REF: not all forward unresolved forward operand references
+
+//===--------------------------------------------------------------------===//
+// Results
+
+// RUN: not mlir-opt %S/invalid-ir_section-results.mlirbc -allow-unregistered-dialect 2>&1 | FileCheck %s --check-prefix=OP_RESULTS
+// OP_RESULTS: value index range was outside of the expected range for the parent region, got [3, 6), but the maximum index was 2
+
+//===--------------------------------------------------------------------===//
+// Successors
+
+// RUN: not mlir-opt %S/invalid-ir_section-successors.mlirbc -allow-unregistered-dialect 2>&1 | FileCheck %s --check-prefix=OP_SUCCESSORS
+// OP_SUCCESSORS: invalid successor index: 3

mlir/test/Bytecode/invalid/invalid-string_section.mlir

@@ -0,0 +1,26 @@
+// This file contains various failure test cases related to the structure of
+// the string section.
+
+//===--------------------------------------------------------------------===//
+// Count
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-string_section-count.mlirbc 2>&1 | FileCheck %s --check-prefix=COUNT
+// COUNT: attempting to parse a byte at the end of the bytecode
+
+//===--------------------------------------------------------------------===//
+// Invalid String
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-string_section-no_string.mlirbc 2>&1 | FileCheck %s --check-prefix=NO_STRING
+// NO_STRING: attempting to parse a byte at the end of the bytecode
+
+// RUN: not mlir-opt %S/invalid-string_section-large_string.mlirbc 2>&1 | FileCheck %s --check-prefix=LARGE_STRING
+// LARGE_STRING: string size exceeds the available data size
+
+//===--------------------------------------------------------------------===//
+// Trailing data
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-string_section-trailing_data.mlirbc 2>&1 | FileCheck %s --check-prefix=TRAILING_DATA
+// TRAILING_DATA: unexpected trailing data between the offsets for strings and their data

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

@@ -0,0 +1 @@
+ML颮
<01>

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

@@ -0,0 +1 @@
+ML颮<EFBFBD>

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

@@ -0,0 +1,44 @@
+// This file contains various failure test cases related to the structure of
+// a bytecode file.
+
+//===--------------------------------------------------------------------===//
+// Version
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-structure-version.mlirbc 2>&1 | FileCheck %s --check-prefix=VERSION
+// VERSION: bytecode version 127 is newer than the current version 0
+
+//===--------------------------------------------------------------------===//
+// Producer
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-structure-producer.mlirbc 2>&1 | FileCheck %s --check-prefix=PRODUCER
+// PRODUCER: malformed null-terminated string, no null character found
+
+//===--------------------------------------------------------------------===//
+// Section
+//===--------------------------------------------------------------------===//
+
+//===--------------------------------------------------------------------===//
+// Missing
+
+// 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
+
+//===--------------------------------------------------------------------===//
+// 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/invalid/invalid_attr_type_offset_section.mlir

@@ -0,0 +1,16 @@
+// This file contains various failure test cases related to the structure of
+// the attribute/type offset section.
+
+//===--------------------------------------------------------------------===//
+// Offset
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-attr_type_offset_section-large_offset.mlirbc 2>&1 | FileCheck %s --check-prefix=LARGE_OFFSET
+// LARGE_OFFSET: Attribute or Type entry offset points past the end of section
+
+//===--------------------------------------------------------------------===//
+// Trailing Data
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-attr_type_offset_section-trailing_data.mlirbc 2>&1 | FileCheck %s --check-prefix=TRAILING_DATA
+// TRAILING_DATA: unexpected trailing data in the Attribute/Type offset section

mlir/test/Bytecode/invalid/invalid_attr_type_section.mlir

@@ -0,0 +1,16 @@
+// This file contains various failure test cases related to the structure of
+// the attribute/type offset section.
+
+//===--------------------------------------------------------------------===//
+// Index
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-attr_type_section-index.mlirbc 2>&1 | FileCheck %s --check-prefix=INDEX
+// INDEX: invalid Attribute index: 3
+
+//===--------------------------------------------------------------------===//
+// Trailing Data
+//===--------------------------------------------------------------------===//
+
+// RUN: not mlir-opt %S/invalid-attr_type_section-trailing_data.mlirbc 2>&1 | FileCheck %s --check-prefix=TRAILING_DATA
+// TRAILING_DATA: trailing characters found after Attribute assembly format: trailing