BytecodeWriter.cpp

Go to the documentation of this file.

//===- 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 "IRNumbering.h"
#include "mlir/Bytecode/BytecodeImplementation.h"
#include "mlir/Bytecode/BytecodeOpInterface.h"
#include "mlir/Bytecode/Encoding.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/OpImplementation.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/CachedHashString.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/Endian.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Support/raw_ostream.h"
#include <optional>
 
#define DEBUG_TYPE "mlir-bytecode-writer"
 
using namespace mlir;
using namespace mlir::bytecode::detail;
 
//===----------------------------------------------------------------------===//
// BytecodeWriterConfig
//===----------------------------------------------------------------------===//
 
struct BytecodeWriterConfig::Impl {
  Impl(StringRef producer) : producer(producer) {}
 
  /// Version to use when writing.
  /// Note: This only differs from kVersion if a specific version is set.
  int64_t bytecodeVersion = bytecode::kVersion;
 
  /// A flag specifying whether to elide emission of resources into the bytecode
  /// file.
  bool shouldElideResourceData = false;
 
  /// A flag specifying whether to elide emission of locations.
  bool shouldElideLocations = false;
 
  /// A map containing dialect version information for each dialect to emit.
  llvm::StringMap<std::unique_ptr<DialectVersion>> dialectVersionMap;
 
  /// The producer of the bytecode.
  StringRef producer;
 
  /// Printer callbacks used to emit custom type and attribute encodings.
  llvm::SmallVector<std::unique_ptr<AttrTypeBytecodeWriter<Attribute>>>
      attributeWriterCallbacks;
  llvm::SmallVector<std::unique_ptr<AttrTypeBytecodeWriter<Type>>>
      typeWriterCallbacks;
 
  /// 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(FallbackAsmResourceMap &map,
                                           StringRef producer)
    : BytecodeWriterConfig(producer) {
  attachFallbackResourcePrinter(map);
}
BytecodeWriterConfig::BytecodeWriterConfig(BytecodeWriterConfig &&config)
    : impl(std::move(config.impl)) {}
 
BytecodeWriterConfig::~BytecodeWriterConfig() = default;
 
ArrayRef<std::unique_ptr<AttrTypeBytecodeWriter<Attribute>>>
BytecodeWriterConfig::getAttributeWriterCallbacks() const {
  return impl->attributeWriterCallbacks;
}
 
ArrayRef<std::unique_ptr<AttrTypeBytecodeWriter<Type>>>
BytecodeWriterConfig::getTypeWriterCallbacks() const {
  return impl->typeWriterCallbacks;
}
 
void BytecodeWriterConfig::attachAttributeCallback(
    std::unique_ptr<AttrTypeBytecodeWriter<Attribute>> callback) {
  impl->attributeWriterCallbacks.emplace_back(std::move(callback));
}
 
void BytecodeWriterConfig::attachTypeCallback(
    std::unique_ptr<AttrTypeBytecodeWriter<Type>> callback) {
  impl->typeWriterCallbacks.emplace_back(std::move(callback));
}
 
void BytecodeWriterConfig::attachResourcePrinter(
    std::unique_ptr<AsmResourcePrinter> printer) {
  impl->externalResourcePrinters.emplace_back(std::move(printer));
}
 
void BytecodeWriterConfig::setElideResourceDataFlag(
    bool shouldElideResourceData) {
  impl->shouldElideResourceData = shouldElideResourceData;
}
 
void BytecodeWriterConfig::setElideLocations(bool shouldElideLocations) {
  impl->shouldElideLocations = shouldElideLocations;
}
 
bool BytecodeWriterConfig::shouldElideLocations() const {
  return impl->shouldElideLocations;
}
 
void BytecodeWriterConfig::setDesiredBytecodeVersion(int64_t bytecodeVersion) {
  impl->bytecodeVersion = bytecodeVersion;
}
 
int64_t BytecodeWriterConfig::getDesiredBytecodeVersion() const {
  return impl->bytecodeVersion;
}
 
llvm::StringMap<std::unique_ptr<DialectVersion>> &
BytecodeWriterConfig::getDialectVersionMap() const {
  return impl->dialectVersionMap;
}
 
void BytecodeWriterConfig::setDialectVersion(
    llvm::StringRef dialectName,
    std::unique_ptr<DialectVersion> dialectVersion) const {
  assert(!impl->dialectVersionMap.contains(dialectName) &&
         "cannot override a previously set dialect version");
  impl->dialectVersionMap.insert({dialectName, std::move(dialectVersion)});
}
 
//===----------------------------------------------------------------------===//
// 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, StringLiteral desc) {
    LDBG() << "patchByte(" << offset << ',' << uint64_t(value) << ")\t" << desc;
    assert(offset < size() && offset >= prevResultSize &&
           "cannot patch previously emitted data");
    currentResult[offset - prevResultSize] = value;
  }
 
  /// Emit the provided blob of data, which is owned by the caller and is
  /// guaranteed to not die before the end of the bytecode process.
  void emitOwnedBlob(ArrayRef<uint8_t> data, StringLiteral desc) {
    LDBG() << "emitOwnedBlob(" << data.size() << "b)\t" << desc;
    // Push the current buffer before adding the provided data.
    appendResult(std::move(currentResult));
    appendOwnedResult(data);
  }
 
  /// Emit the provided blob of data that has the given alignment, which is
  /// owned by the caller and is guaranteed to not die before the end of the
  /// bytecode process. The alignment value is also encoded, making it available
  /// on load.
  void emitOwnedBlobAndAlignment(ArrayRef<uint8_t> data, uint32_t alignment,
                                 StringLiteral desc) {
    emitVarInt(alignment, desc);
    emitVarInt(data.size(), desc);
 
    alignTo(alignment);
    emitOwnedBlob(data, desc);
  }
  void emitOwnedBlobAndAlignment(ArrayRef<char> data, uint32_t alignment,
                                 StringLiteral desc) {
    ArrayRef<uint8_t> castedData(reinterpret_cast<const uint8_t *>(data.data()),
                                 data.size());
    emitOwnedBlobAndAlignment(castedData, alignment, desc);
  }
 
  /// 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, "alignment byte");
 
    // 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, StringLiteral desc) {
    LDBG() << "emitByte(" << uint64_t(byte) << ")\t" << desc;
    currentResult.push_back(static_cast<uint8_t>(byte));
  }
 
  /// Emit a range of bytes.
  void emitBytes(ArrayRef<uint8_t> bytes, StringLiteral desc) {
    LDBG() << "emitBytes(" << bytes.size() << "b)\t" << desc;
    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, StringLiteral desc) {
    LDBG() << "emitVarInt(" << value << ")\t" << desc;
 
    // 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, desc);
    emitMultiByteVarInt(value, desc);
  }
 
  /// Emit a signed variable length integer. Signed varints are encoded using
  /// a varint with zigzag encoding, meaning that we use the low bit of the
  /// value to indicate the sign of the value. This allows for more efficient
  /// encoding of negative values by limiting the number of active bits
  void emitSignedVarInt(uint64_t value, StringLiteral desc) {
    emitVarInt((value << 1) ^ (uint64_t)((int64_t)value >> 63), desc);
  }
 
  /// 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, StringLiteral desc) {
    emitVarInt((value << 1) | (flag ? 1 : 0), desc);
  }
 
  //===--------------------------------------------------------------------===//
  // String Emission
 
  /// Emit the given string as a nul terminated string.
  void emitNulTerminatedString(StringRef str, StringLiteral desc) {
    emitString(str, desc);
    emitByte(0, "null terminator");
  }
 
  /// Emit the given string without a nul terminator.
  void emitString(StringRef str, StringLiteral desc) {
    emitBytes({reinterpret_cast<const uint8_t *>(str.data()), str.size()},
              desc);
  }
 
  //===--------------------------------------------------------------------===//
  // 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. 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, "section code");
    emitVarInt(emitter.size(), "section 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, "section alignment");
        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)
      prevResultStorage.push_back(std::move(result));
    llvm::append_range(prevResultList, emitter.prevResultList);
    prevResultSize += emitter.prevResultSize;
    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,
                                                   StringLiteral desc);
 
  /// Append a new result buffer to the current contents.
  void appendResult(std::vector<uint8_t> &&result) {
    if (result.empty())
      return;
    prevResultStorage.emplace_back(std::move(result));
    appendOwnedResult(prevResultStorage.back());
  }
  void appendOwnedResult(ArrayRef<uint8_t> result) {
    if (result.empty())
      return;
    prevResultSize += result.size();
    prevResultList.emplace_back(result);
  }
 
  /// 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;
 
  /// The highest required alignment for the start of this section.
  unsigned requiredAlignment = 1;
};
 
//===----------------------------------------------------------------------===//
// StringSectionBuilder
//===----------------------------------------------------------------------===//
 
namespace {
/// This class is used to simplify the process of emitting the string section.
class StringSectionBuilder {
public:
  /// Add the given string to the string section, and return the index of the
  /// string within the section.
  size_t insert(StringRef str) {
    auto it = strings.insert({llvm::CachedHashStringRef(str), strings.size()});
    return it.first->second;
  }
 
  /// Write the current set of strings to the given emitter.
  void write(EncodingEmitter &emitter) {
    emitter.emitVarInt(strings.size(), "string section 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))
      emitter.emitVarInt(it.first.size() + 1, "string size");
    // Emit the string data itself.
    for (const auto &it : strings)
      emitter.emitNulTerminatedString(it.first.val(), "string");
  }
 
private:
  /// A set of strings referenced within the bytecode. The value of the map is
  /// unused.
  llvm::MapVector<llvm::CachedHashStringRef, size_t> strings;
};
} // namespace
 
class DialectWriter : public DialectBytecodeWriter {
  using DialectVersionMapT = llvm::StringMap<std::unique_ptr<DialectVersion>>;
 
public:
  DialectWriter(int64_t bytecodeVersion, EncodingEmitter &emitter,
                IRNumberingState &numberingState,
                StringSectionBuilder &stringSection,
                const DialectVersionMapT &dialectVersionMap)
      : bytecodeVersion(bytecodeVersion), emitter(emitter),
        numberingState(numberingState), stringSection(stringSection),
        dialectVersionMap(dialectVersionMap) {}
 
  //===--------------------------------------------------------------------===//
  // IR
  //===--------------------------------------------------------------------===//
 
  void writeAttribute(Attribute attr) override {
    emitter.emitVarInt(numberingState.getNumber(attr), "dialect attr");
  }
  void writeOptionalAttribute(Attribute attr) override {
    if (!attr) {
      emitter.emitVarInt(0, "dialect optional attr none");
      return;
    }
    emitter.emitVarIntWithFlag(numberingState.getNumber(attr), true,
                               "dialect optional attr");
  }
 
  void writeType(Type type) override {
    emitter.emitVarInt(numberingState.getNumber(type), "dialect type");
  }
 
  void writeResourceHandle(const AsmDialectResourceHandle &resource) override {
    emitter.emitVarInt(numberingState.getNumber(resource), "dialect resource");
  }
 
  //===--------------------------------------------------------------------===//
  // Primitives
  //===--------------------------------------------------------------------===//
 
  void writeVarInt(uint64_t value) override {
    emitter.emitVarInt(value, "dialect writer");
  }
 
  void writeSignedVarInt(int64_t value) override {
    emitter.emitSignedVarInt(value, "dialect writer");
  }
 
  void writeAPIntWithKnownWidth(const APInt &value) override {
    size_t bitWidth = value.getBitWidth();
 
    // If the value is a single byte, just emit it directly without going
    // through a varint.
    if (bitWidth <= 8)
      return emitter.emitByte(value.getLimitedValue(), "dialect APInt");
 
    // If the value fits within a single varint, emit it directly.
    if (bitWidth <= 64)
      return emitter.emitSignedVarInt(value.getLimitedValue(), "dialect APInt");
 
    // Otherwise, we need to encode a variable number of active words. We use
    // active words instead of the number of total words under the observation
    // that smaller values will be more common.
    unsigned numActiveWords = value.getActiveWords();
    emitter.emitVarInt(numActiveWords, "dialect APInt word count");
 
    const uint64_t *rawValueData = value.getRawData();
    for (unsigned i = 0; i < numActiveWords; ++i)
      emitter.emitSignedVarInt(rawValueData[i], "dialect APInt word");
  }
 
  void writeAPFloatWithKnownSemantics(const APFloat &value) override {
    writeAPIntWithKnownWidth(value.bitcastToAPInt());
  }
 
  void writeOwnedString(StringRef str) override {
    emitter.emitVarInt(stringSection.insert(str), "dialect string");
  }
 
  void writeOwnedBlob(ArrayRef<char> blob) override {
    emitter.emitVarInt(blob.size(), "dialect blob");
    emitter.emitOwnedBlob(
        ArrayRef<uint8_t>(reinterpret_cast<const uint8_t *>(blob.data()),
                          blob.size()),
        "dialect blob");
  }
 
  void writeUnownedBlob(ArrayRef<char> blob) override {
    emitter.emitVarInt(blob.size(), "dialect blob");
    emitter.emitBytes(
        ArrayRef<uint8_t>(reinterpret_cast<const uint8_t *>(blob.data()),
                          blob.size()),
        "dialect blob");
  }
 
  void writeOwnedBool(bool value) override {
    emitter.emitByte(value, "dialect bool");
  }
 
  int64_t getBytecodeVersion() const override { return bytecodeVersion; }
 
  FailureOr<const DialectVersion *>
  getDialectVersion(StringRef dialectName) const override {
    auto dialectEntry = dialectVersionMap.find(dialectName);
    if (dialectEntry == dialectVersionMap.end())
      return failure();
    return dialectEntry->getValue().get();
  }
 
private:
  int64_t bytecodeVersion;
  EncodingEmitter &emitter;
  IRNumberingState &numberingState;
  StringSectionBuilder &stringSection;
  const DialectVersionMapT &dialectVersionMap;
};
 
namespace {
class PropertiesSectionBuilder {
public:
  PropertiesSectionBuilder(IRNumberingState &numberingState,
                           StringSectionBuilder &stringSection,
                           const BytecodeWriterConfig::Impl &config)
      : numberingState(numberingState), stringSection(stringSection),
        config(config) {}
 
  /// Emit the op properties in the properties section and return the index of
  /// the properties within the section. Return -1 if no properties was emitted.
  std::optional<ssize_t> emit(Operation *op) {
    EncodingEmitter propertiesEmitter;
    if (!op->getPropertiesStorageSize())
      return std::nullopt;
    if (!op->isRegistered()) {
      // Unregistered op are storing properties as an optional attribute.
      Attribute prop = *op->getPropertiesStorage().as<Attribute *>();
      if (!prop)
        return std::nullopt;
      EncodingEmitter sizeEmitter;
      sizeEmitter.emitVarInt(numberingState.getNumber(prop), "properties size");
      scratch.clear();
      llvm::raw_svector_ostream os(scratch);
      sizeEmitter.writeTo(os);
      return emit(scratch);
    }
 
    EncodingEmitter emitter;
    DialectWriter propertiesWriter(config.bytecodeVersion, emitter,
                                   numberingState, stringSection,
                                   config.dialectVersionMap);
    auto iface = cast<BytecodeOpInterface>(op);
    iface.writeProperties(propertiesWriter);
    scratch.clear();
    llvm::raw_svector_ostream os(scratch);
    emitter.writeTo(os);
    return emit(scratch);
  }
 
  /// Write the current set of properties to the given emitter.
  void write(EncodingEmitter &emitter) {
    emitter.emitVarInt(propertiesStorage.size(), "properties size");
    if (propertiesStorage.empty())
      return;
    for (const auto &storage : propertiesStorage) {
      if (storage.empty()) {
        emitter.emitBytes(ArrayRef<uint8_t>(), "empty properties");
        continue;
      }
      emitter.emitBytes(ArrayRef(reinterpret_cast<const uint8_t *>(&storage[0]),
                                 storage.size()),
                        "property");
    }
  }
 
  /// Returns true if the section is empty.
  bool empty() { return propertiesStorage.empty(); }
 
private:
  /// Emit raw data and returns the offset in the internal buffer.
  /// Data are deduplicated and will be copied in the internal buffer only if
  /// they don't exist there already.
  ssize_t emit(ArrayRef<char> rawProperties) {
    // Populate a scratch buffer with the properties size.
    SmallVector<char> sizeScratch;
    {
      EncodingEmitter sizeEmitter;
      sizeEmitter.emitVarInt(rawProperties.size(), "properties");
      llvm::raw_svector_ostream os(sizeScratch);
      sizeEmitter.writeTo(os);
    }
    // Append a new storage to the table now.
    size_t index = propertiesStorage.size();
    propertiesStorage.emplace_back();
    std::vector<char> &newStorage = propertiesStorage.back();
    size_t propertiesSize = sizeScratch.size() + rawProperties.size();
    newStorage.reserve(propertiesSize);
    llvm::append_range(newStorage, sizeScratch);
    llvm::append_range(newStorage, rawProperties);
 
    // Try to de-duplicate the new serialized properties.
    // If the properties is a duplicate, pop it back from the storage.
    auto inserted = propertiesUniquing.insert(
        std::make_pair(ArrayRef<char>(newStorage), index));
    if (!inserted.second)
      propertiesStorage.pop_back();
    return inserted.first->getSecond();
  }
 
  /// Storage for properties.
  std::vector<std::vector<char>> propertiesStorage;
  SmallVector<char> scratch;
  DenseMap<ArrayRef<char>, int64_t> propertiesUniquing;
  IRNumberingState &numberingState;
  StringSectionBuilder &stringSection;
  const BytecodeWriterConfig::Impl &config;
};
} // namespace
 
/// 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) {
    SetUnbuffered();
  }
 
private:
  void write_impl(const char *ptr, size_t size) override {
    emitter.emitBytes({reinterpret_cast<const uint8_t *>(ptr), size},
                      "raw emitter");
  }
  uint64_t current_pos() const override { return emitter.size(); }
 
  /// The section being emitted to.
  EncodingEmitter &emitter;
};
} // namespace
 
void EncodingEmitter::writeTo(raw_ostream &os) const {
  // Reserve space in the ostream for the encoded contents.
  os.reserveExtraSpace(size());
 
  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, StringLiteral desc) {
  // 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);
      llvm::support::ulittle64_t encodedValueLE(encodedValue);
      emitBytes({reinterpret_cast<uint8_t *>(&encodedValueLE), numBytes}, desc);
      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, desc);
  llvm::support::ulittle64_t valueLE(value);
  emitBytes({reinterpret_cast<uint8_t *>(&valueLE), sizeof(valueLE)}, desc);
}
 
//===----------------------------------------------------------------------===//
// Bytecode Writer
//===----------------------------------------------------------------------===//
 
namespace {
class BytecodeWriter {
public:
  BytecodeWriter(Operation *op, const BytecodeWriterConfig &config)
      : numberingState(op, config), config(config.getImpl()),
        propertiesSection(numberingState, stringSection, config.getImpl()) {}
 
  /// Write the bytecode for the given root operation.
  LogicalResult write(Operation *rootOp, raw_ostream &os);
 
private:
  //===--------------------------------------------------------------------===//
  // Dialects
 
  void writeDialectSection(EncodingEmitter &emitter);
 
  //===--------------------------------------------------------------------===//
  // Attributes and Types
 
  void writeAttrTypeSection(EncodingEmitter &emitter);
 
  //===--------------------------------------------------------------------===//
  // Operations
 
  LogicalResult writeBlock(EncodingEmitter &emitter, Block *block);
  LogicalResult writeOp(EncodingEmitter &emitter, Operation *op);
  LogicalResult writeRegion(EncodingEmitter &emitter, Region *region);
  LogicalResult writeIRSection(EncodingEmitter &emitter, Operation *op);
 
  LogicalResult writeRegions(EncodingEmitter &emitter,
                             MutableArrayRef<Region> regions) {
    return success(llvm::all_of(regions, [&](Region &region) {
      return succeeded(writeRegion(emitter, &region));
    }));
  }
 
  //===--------------------------------------------------------------------===//
  // Resources
 
  void writeResourceSection(Operation *op, EncodingEmitter &emitter);
 
  //===--------------------------------------------------------------------===//
  // Strings
 
  void writeStringSection(EncodingEmitter &emitter);
 
  //===--------------------------------------------------------------------===//
  // Properties
 
  void writePropertiesSection(EncodingEmitter &emitter);
 
  //===--------------------------------------------------------------------===//
  // Helpers
 
  void writeUseListOrders(EncodingEmitter &emitter, uint8_t &opEncodingMask,
                          ValueRange range);
 
  //===--------------------------------------------------------------------===//
  // Fields
 
  /// The builder used for the string section.
  StringSectionBuilder stringSection;
 
  /// The IR numbering state generated for the root operation.
  IRNumberingState numberingState;
 
  /// Configuration dictating bytecode emission.
  const BytecodeWriterConfig::Impl &config;
 
  /// Storage for the properties section
  PropertiesSectionBuilder propertiesSection;
};
} // namespace
 
LogicalResult BytecodeWriter::write(Operation *rootOp, raw_ostream &os) {
  EncodingEmitter emitter;
 
  // Emit the bytecode file header. This is how we identify the output as a
  // bytecode file.
  emitter.emitString("ML\xefR", "bytecode header");
 
  // Emit the bytecode version.
  if (config.bytecodeVersion < bytecode::kMinSupportedVersion ||
      config.bytecodeVersion > bytecode::kVersion)
    return rootOp->emitError()
           << "unsupported version requested " << config.bytecodeVersion
           << ", must be in range ["
           << static_cast<int64_t>(bytecode::kMinSupportedVersion) << ", "
           << static_cast<int64_t>(bytecode::kVersion) << ']';
  emitter.emitVarInt(config.bytecodeVersion, "bytecode version");
 
  // Emit the producer.
  emitter.emitNulTerminatedString(config.producer, "bytecode producer");
 
  // Emit the dialect section.
  writeDialectSection(emitter);
 
  // Emit the attributes and types section.
  writeAttrTypeSection(emitter);
 
  // Emit the IR section.
  if (failed(writeIRSection(emitter, rootOp)))
    return failure();
 
  // Emit the resources section.
  writeResourceSection(rootOp, emitter);
 
  // Emit the string section.
  writeStringSection(emitter);
 
  // Emit the properties section.
  if (config.bytecodeVersion >= bytecode::kNativePropertiesEncoding)
    writePropertiesSection(emitter);
  else if (!propertiesSection.empty())
    return rootOp->emitError(
        "unexpected properties emitted incompatible with bytecode <5");
 
  // Write the generated bytecode to the provided output stream.
  emitter.writeTo(os);
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// 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, "dialect number");
    emitter.emitVarInt(std::distance(groupStart, it), "dialect offset");
 
    // 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), "dialects count");
  for (DialectNumbering &dialect : dialects) {
    // Write the string section and get the ID.
    size_t nameID = stringSection.insert(dialect.name);
 
    if (config.bytecodeVersion < bytecode::kDialectVersioning) {
      dialectEmitter.emitVarInt(nameID, "dialect name ID");
      continue;
    }
 
    // Try writing the version to the versionEmitter.
    EncodingEmitter versionEmitter;
    if (dialect.interface) {
      // The writer used when emitting using a custom bytecode encoding.
      DialectWriter versionWriter(config.bytecodeVersion, versionEmitter,
                                  numberingState, stringSection,
                                  config.dialectVersionMap);
      dialect.interface->writeVersion(versionWriter);
    }
 
    // If the version emitter is empty, version is not available. We can encode
    // this in the dialect ID, so if there is no version, we don't write the
    // section.
    size_t versionAvailable = versionEmitter.size() > 0;
    dialectEmitter.emitVarIntWithFlag(nameID, versionAvailable,
                                      "dialect version");
    if (versionAvailable)
      dialectEmitter.emitSection(bytecode::Section::kDialectVersions,
                                 std::move(versionEmitter));
  }
 
  if (config.bytecodeVersion >= bytecode::kElideUnknownBlockArgLocation)
    dialectEmitter.emitVarInt(size(numberingState.getOpNames()),
                              "op names count");
 
  // Emit the referenced operation names grouped by dialect.
  auto emitOpName = [&](OpNameNumbering &name) {
    size_t stringId = stringSection.insert(name.name.stripDialect());
    if (config.bytecodeVersion < bytecode::kNativePropertiesEncoding)
      dialectEmitter.emitVarInt(stringId, "dialect op name");
    else
      dialectEmitter.emitVarIntWithFlag(stringId, name.name.isRegistered(),
                                        "dialect op name");
  };
  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()),
                           "attributes count");
  offsetEmitter.emitVarInt(llvm::size(numberingState.getTypes()),
                           "types count");
 
  // A functor used to emit an attribute or type entry.
  uint64_t prevOffset = 0;
  auto emitAttrOrType = [&](auto &entry) {
    auto entryValue = entry.getValue();
 
    auto emitAttrOrTypeRawImpl = [&]() -> void {
      RawEmitterOstream(attrTypeEmitter) << entryValue;
      attrTypeEmitter.emitByte(0, "attr/type separator");
    };
    auto emitAttrOrTypeImpl = [&]() -> bool {
      // TODO: We don't currently support custom encoded mutable types and
      // attributes.
      if (entryValue.template hasTrait<TypeTrait::IsMutable>() ||
          entryValue.template hasTrait<AttributeTrait::IsMutable>()) {
        emitAttrOrTypeRawImpl();
        return false;
      }
 
      DialectWriter dialectWriter(config.bytecodeVersion, attrTypeEmitter,
                                  numberingState, stringSection,
                                  config.dialectVersionMap);
      if constexpr (std::is_same_v<std::decay_t<decltype(entryValue)>, Type>) {
        for (const auto &callback : config.typeWriterCallbacks) {
          if (succeeded(callback->write(entryValue, dialectWriter)))
            return true;
        }
        if (const BytecodeDialectInterface *interface =
                entry.dialect->interface) {
          if (succeeded(interface->writeType(entryValue, dialectWriter)))
            return true;
        }
      } else {
        for (const auto &callback : config.attributeWriterCallbacks) {
          if (succeeded(callback->write(entryValue, dialectWriter)))
            return true;
        }
        if (const BytecodeDialectInterface *interface =
                entry.dialect->interface) {
          if (succeeded(interface->writeAttribute(entryValue, dialectWriter)))
            return true;
        }
      }
 
      // If the entry was not emitted using a callback or a dialect interface,
      // emit it using the textual format.
      emitAttrOrTypeRawImpl();
      return false;
    };
 
    bool hasCustomEncoding = emitAttrOrTypeImpl();
 
    // Record the offset of this entry.
    uint64_t curOffset = attrTypeEmitter.size();
    offsetEmitter.emitVarIntWithFlag(curOffset - prevOffset, hasCustomEncoding,
                                     "attr/type offset");
    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
//===----------------------------------------------------------------------===//
 
LogicalResult 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, "block num ops");
 
  // Emit the arguments of the block.
  if (hasArgs) {
    emitter.emitVarInt(args.size(), "block args count");
    for (BlockArgument arg : args) {
      Location argLoc = arg.getLoc();
      if (config.bytecodeVersion >= bytecode::kElideUnknownBlockArgLocation) {
        emitter.emitVarIntWithFlag(numberingState.getNumber(arg.getType()),
                                   !isa<UnknownLoc>(argLoc), "block arg type");
        if (!isa<UnknownLoc>(argLoc))
          emitter.emitVarInt(numberingState.getNumber(argLoc),
                             "block arg location");
      } else {
        emitter.emitVarInt(numberingState.getNumber(arg.getType()),
                           "block arg type");
        emitter.emitVarInt(numberingState.getNumber(argLoc),
                           "block arg location");
      }
    }
    if (config.bytecodeVersion >= bytecode::kUseListOrdering) {
      uint64_t maskOffset = emitter.size();
      uint8_t encodingMask = 0;
      emitter.emitByte(0, "use-list separator");
      writeUseListOrders(emitter, encodingMask, args);
      if (encodingMask)
        emitter.patchByte(maskOffset, encodingMask, "block patch encoding");
    }
  }
 
  // Emit the operations within the block.
  for (Operation &op : *block)
    if (failed(writeOp(emitter, &op)))
      return failure();
  return success();
}
 
/// Check legacy property reconstruction using scratch storage. Only inherent
/// attributes may restore properties; discardable attributes with native
/// property names must remain discardable.
static bool preservesLegacyProperties(Operation *op, DictionaryAttr attrs) {
  OperationName name = op->getName();
  // Operation property storage is aligned to eight bytes.
  SmallVector<uint64_t, 8> storage(
      llvm::divideCeil(op->getPropertiesStorageSize(), sizeof(uint64_t)));
  PropertyRef reconstructed(name.getOpPropertiesTypeID(), storage.data());
  name.initOpProperties(reconstructed, {});
  auto cleanup =
      llvm::scope_exit([&] { name.destroyOpProperties(reconstructed); });
 
  // Unregistered operations cannot restore their property attribute from the
  // legacy dictionary: their initialized property attribute is null.
  if (name.isRegistered()) {
    NamedAttrList legacyAttrs(attrs);
    name.populateDefaultAttrs(legacyAttrs);
    NamedAttrList inherentAttrs;
    for (NamedAttribute attr : legacyAttrs)
      if (name.getInherentAttr(op, attr.getName()))
        inherentAttrs.append(attr);
    if (!inherentAttrs.empty() &&
        failed(name.setOpPropertiesFromAttribute(
            name, reconstructed, inherentAttrs.getDictionary(op->getContext()),
            [&] {
              auto diag = mlir::emitError(op->getLoc());
              diag.abandon();
              return diag;
            })))
      return false;
  }
  return name.compareOpProperties(op->getPropertiesStorage(), reconstructed);
}
 
LogicalResult BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
  emitter.emitVarInt(numberingState.getNumber(op->getName()), "op name ID");
 
  // 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, "op separator");
 
  // Emit the location for this operation.
  emitter.emitVarInt(numberingState.getNumber(op->getLoc()), "op location");
 
  // Emit the attributes of this operation.
  DictionaryAttr attrs = op->getRawDictionaryAttrs();
  // Allow deployment to version <kNativePropertiesEncoding by merging inherent
  // attribute with the discardable ones. We should fail if there are any
  // conflicts. When properties are not used by the op, also store everything as
  // attributes.
  if (config.bytecodeVersion < bytecode::kNativePropertiesEncoding &&
      op->getPropertiesStorage()) {
    NamedAttrList allAttrs;
    op->getName().walkInherentAttrs(op, [&](StringRef name, Attribute &attr) {
      allAttrs.append(name, attr);
    });
    allAttrs.append(op->getDiscardableAttrDictionary().getValue());
    attrs = allAttrs.getDictionary(op->getContext());
 
    // Versions before kNativePropertiesEncoding can only reconstruct properties
    // from these attributes. Check that native properties would not silently
    // revert to their defaults. Later versions encode properties directly and
    // do not need this legacy reconstruction check.
    if (!preservesLegacyProperties(op, attrs))
      return op->emitOpError()
             << "cannot emit bytecode before version "
             << static_cast<int64_t>(bytecode::kNativePropertiesEncoding)
             << " without losing operation properties";
  }
  if (!attrs.empty()) {
    opEncodingMask |= bytecode::OpEncodingMask::kHasAttrs;
    emitter.emitVarInt(numberingState.getNumber(attrs), "op attrs count");
  }
 
  // Emit the properties of this operation, for now we still support deployment
  // to version <kNativePropertiesEncoding.
  if (config.bytecodeVersion >= bytecode::kNativePropertiesEncoding) {
    std::optional<ssize_t> propertiesId = propertiesSection.emit(op);
    if (propertiesId.has_value()) {
      opEncodingMask |= bytecode::OpEncodingMask::kHasProperties;
      emitter.emitVarInt(*propertiesId, "op properties ID");
    }
  }
 
  // Emit the result types of the operation.
  if (unsigned numResults = op->getNumResults()) {
    opEncodingMask |= bytecode::OpEncodingMask::kHasResults;
    emitter.emitVarInt(numResults, "op results count");
    for (Type type : op->getResultTypes())
      emitter.emitVarInt(numberingState.getNumber(type), "op result type");
  }
 
  // Emit the operands of the operation.
  if (unsigned numOperands = op->getNumOperands()) {
    opEncodingMask |= bytecode::OpEncodingMask::kHasOperands;
    emitter.emitVarInt(numOperands, "op operands count");
    for (Value operand : op->getOperands())
      emitter.emitVarInt(numberingState.getNumber(operand), "op operand types");
  }
 
  // Emit the successors of the operation.
  if (unsigned numSuccessors = op->getNumSuccessors()) {
    opEncodingMask |= bytecode::OpEncodingMask::kHasSuccessors;
    emitter.emitVarInt(numSuccessors, "op successors count");
    for (Block *successor : op->getSuccessors())
      emitter.emitVarInt(numberingState.getNumber(successor), "op successor");
  }
 
  // Emit the use-list orders to bytecode, so we can reconstruct the same order
  // at parsing.
  if (config.bytecodeVersion >= bytecode::kUseListOrdering)
    writeUseListOrders(emitter, opEncodingMask, ValueRange(op->getResults()));
 
  // Check for regions.
  unsigned numRegions = op->getNumRegions();
  if (numRegions)
    opEncodingMask |= bytecode::OpEncodingMask::kHasInlineRegions;
 
  // Update the mask for the operation.
  emitter.patchByte(maskOffset, opEncodingMask, "op encoding mask");
 
  // 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 = numberingState.isIsolatedFromAbove(op);
    emitter.emitVarIntWithFlag(numRegions, isIsolatedFromAbove,
                               "op regions count");
 
    // If the region is not isolated from above, or we are emitting bytecode
    // targeting version <kLazyLoading, we don't use a section.
    if (isIsolatedFromAbove &&
        config.bytecodeVersion >= bytecode::kLazyLoading) {
      EncodingEmitter regionEmitter;
      if (failed(writeRegions(regionEmitter, op->getRegions())))
        return failure();
      emitter.emitSection(bytecode::Section::kIR, std::move(regionEmitter));
 
    } else if (failed(writeRegions(emitter, op->getRegions()))) {
      return failure();
    }
  }
  return success();
}
 
void BytecodeWriter::writeUseListOrders(EncodingEmitter &emitter,
                                        uint8_t &opEncodingMask,
                                        ValueRange range) {
  // Loop over the results and store the use-list order per result index.
  llvm::MapVector<unsigned, llvm::SmallVector<unsigned>> map;
  for (auto item : llvm::enumerate(range)) {
    auto value = item.value();
    // No need to store a custom use-list order if the result does not have
    // multiple uses.
    if (value.use_empty() || value.hasOneUse())
      continue;
 
    // For each result, assemble the list of pairs (use-list-index,
    // global-value-index). While doing so, detect if the global-value-index is
    // already ordered with respect to the use-list-index.
    bool alreadyOrdered = true;
    auto &firstUse = *value.use_begin();
    uint64_t prevID = bytecode::getUseID(
        firstUse, numberingState.getNumber(firstUse.getOwner()));
    llvm::SmallVector<std::pair<unsigned, uint64_t>> useListPairs(
        {{0, prevID}});
 
    for (auto use : llvm::drop_begin(llvm::enumerate(value.getUses()))) {
      uint64_t currentID = bytecode::getUseID(
          use.value(), numberingState.getNumber(use.value().getOwner()));
      // The use-list order achieved when building the IR at parsing always
      // pushes new uses on front. Hence, if the order by unique ID is
      // monotonically decreasing, a roundtrip to bytecode preserves such order.
      alreadyOrdered &= (prevID > currentID);
      useListPairs.push_back({use.index(), currentID});
      prevID = currentID;
    }
 
    // Do not emit if the order is already sorted.
    if (alreadyOrdered)
      continue;
 
    // Sort the use indices by the unique ID indices in descending order.
    std::sort(
        useListPairs.begin(), useListPairs.end(),
        [](auto elem1, auto elem2) { return elem1.second > elem2.second; });
 
    map.try_emplace(item.index(), llvm::map_range(useListPairs, [](auto elem) {
                      return elem.first;
                    }));
  }
 
  if (map.empty())
    return;
 
  opEncodingMask |= bytecode::OpEncodingMask::kHasUseListOrders;
  // Emit the number of results that have a custom use-list order if the number
  // of results is greater than one.
  if (range.size() != 1) {
    emitter.emitVarInt(map.size(), "custom use-list size");
  }
 
  for (const auto &[resultIdx, useListOrder] : map) {
    // Compute the number of uses that are actually shuffled. If those are less
    // than half of the total uses, encoding the index pair `(src, dst)` is more
    // space efficient.
    size_t shuffledElements =
        llvm::count_if(llvm::enumerate(useListOrder),
                       [](auto item) { return item.index() != item.value(); });
    bool indexPairEncoding = shuffledElements < (useListOrder.size() / 2);
 
    // For single result, we don't need to store the result index.
    if (range.size() != 1)
      emitter.emitVarInt(resultIdx, "use-list result index");
 
    if (indexPairEncoding) {
      emitter.emitVarIntWithFlag(shuffledElements * 2, indexPairEncoding,
                                 "use-list index pair size");
      for (auto pair : llvm::enumerate(useListOrder)) {
        if (pair.index() != pair.value()) {
          emitter.emitVarInt(pair.value(), "use-list index pair first");
          emitter.emitVarInt(pair.index(), "use-list index pair second");
        }
      }
    } else {
      emitter.emitVarIntWithFlag(useListOrder.size(), indexPairEncoding,
                                 "use-list size");
      for (const auto &index : useListOrder)
        emitter.emitVarInt(index, "use-list order");
    }
  }
}
 
LogicalResult 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()) {
    emitter.emitVarInt(/*numBlocks*/ 0, "region block count empty");
    return success();
  }
 
  // Emit the number of blocks and values within the region.
  unsigned numBlocks, numValues;
  std::tie(numBlocks, numValues) = numberingState.getBlockValueCount(region);
  emitter.emitVarInt(numBlocks, "region block count");
  emitter.emitVarInt(numValues, "region value count");
 
  // Emit the blocks within the region.
  for (Block &block : *region)
    if (failed(writeBlock(emitter, &block)))
      return failure();
  return success();
}
 
LogicalResult 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, "ir section");
 
  // Emit the operations.
  if (failed(writeOp(irEmitter, op)))
    return failure();
 
  emitter.emitSection(bytecode::Section::kIR, std::move(irEmitter));
  return success();
}
 
//===----------------------------------------------------------------------===//
// 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, bool shouldElideData)
      : emitter(emitter), stringSection(stringSection),
        postProcessFn(postProcessFn), shouldElideData(shouldElideData) {}
  ~ResourceBuilder() override = default;
 
  void buildBlob(StringRef key, ArrayRef<char> data,
                 uint32_t dataAlignment) final {
    if (!shouldElideData)
      emitter.emitOwnedBlobAndAlignment(data, dataAlignment, "resource blob");
    postProcessFn(key, AsmResourceEntryKind::Blob);
  }
  void buildBool(StringRef key, bool data) final {
    if (!shouldElideData)
      emitter.emitByte(data, "resource bool");
    postProcessFn(key, AsmResourceEntryKind::Bool);
  }
  void buildString(StringRef key, StringRef data) final {
    if (!shouldElideData)
      emitter.emitVarInt(stringSection.insert(data), "resource string");
    postProcessFn(key, AsmResourceEntryKind::String);
  }
 
private:
  EncodingEmitter &emitter;
  StringSectionBuilder &stringSection;
  PostProcessFn postProcessFn;
  bool shouldElideData = false;
};
} // namespace
 
void BytecodeWriter::writeResourceSection(Operation *op,
                                          EncodingEmitter &emitter) {
  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, "resource group key");
    resourceOffsetEmitter.emitVarInt(curResourceEntries.size(),
                                     "resource group size");
    for (auto [key, kind, size] : curResourceEntries) {
      resourceOffsetEmitter.emitVarInt(stringSection.insert(key),
                                       "resource key");
      resourceOffsetEmitter.emitVarInt(size, "resource size");
      resourceOffsetEmitter.emitByte(kind, "resource kind");
    }
  };
 
  // Builder used to emit resources.
  ResourceBuilder entryBuilder(resourceEmitter, stringSection,
                               appendResourceOffset,
                               config.shouldElideResourceData);
 
  // Emit the external resource entries.
  resourceOffsetEmitter.emitVarInt(config.externalResourcePrinters.size(),
                                   "external resource printer count");
  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));
}
 
//===----------------------------------------------------------------------===//
// Properties
//===----------------------------------------------------------------------===//
 
void BytecodeWriter::writePropertiesSection(EncodingEmitter &emitter) {
  EncodingEmitter propertiesEmitter;
  propertiesSection.write(propertiesEmitter);
  emitter.emitSection(bytecode::Section::kProperties,
                      std::move(propertiesEmitter));
}
 
//===----------------------------------------------------------------------===//
// Entry Points
//===----------------------------------------------------------------------===//
 
LogicalResult mlir::writeBytecodeToFile(Operation *op, raw_ostream &os,
                                        const BytecodeWriterConfig &config) {
  // Validate before numbering the IR, which may invoke dialect writers.
  for (const auto &entry : config.getDialectVersionMap()) {
    Dialect *dialect = op->getContext()->getLoadedDialect(entry.getKey());
    if (!dialect)
      continue;
    auto *interface =
        dialect->getRegisteredInterface<BytecodeDialectInterface>();
    if (interface && failed(interface->verifyVersion(op, *entry.getValue())))
      return failure();
  }
  BytecodeWriter writer(op, config);
  return writer.write(op, os);
}