[mlir:Bytecode] Add initial support for dialect defined attribute/type encodings

[rriddle]

Summary

Dialects can opt-in to providing custom encodings by implementing the
BytecodeDialectInterface. This interface provides hooks, namely
readAttribute/readType and writeAttribute/writeType, that will be used
by the bytecode reader and writer. These hooks are provided a reader and writer
implementation that can be used to encode various constructs in the underlying
bytecode format. A unique feature of this interface is that dialects may choose
to only encode a subset of their attributes and types in a custom bytecode
format, which can simplify adding new or experimental components that aren't
fully baked.

[zero9178]

mlir/include/mlir/Bytecode/BytecodeImplementation.h

44

What is the advantage of differentiating between default constructible and not here? Why not always use the FailureOr<T>() from?

mlir/lib/IR/BuiltinDialectBytecode.cpp

29–37

Should the varint for the length, written and read by (read|write)List be documented here as well?

[mehdi_amini]

Nice!

mlir/docs/BytecodeFormat.md

213

compatibility

Also we may over-promise at the moment here.

mlir/include/mlir/Bytecode/BytecodeImplementation.h

30

Please document a bit more who should subclass it (no one?) and how, where it get involved.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

345

Nit: "parent" is confusing to me here.

mlir/lib/IR/BuiltinDialectBytecode.cpp

29–37

I don't think so: the [] expresses that this is a list, the encoding of a list is an implementation detail of the list, not part of the schema here.

36

<StringAttr, Attribute> ?

I guess StringAttr may not be desirable because it is an ID here and it can be confusing?

[rriddle]

mlir/docs/BytecodeFormat.md

213

Yeah, just dropping this part of the sentence for now. I don't really want to promise anything at this point.

mlir/include/mlir/Bytecode/BytecodeImplementation.h

44

We have different APIs that use either LogicalResult+Reference or FailureOr, supporting both is convenient here. Switched to checking the callable type instead of basing this on default constructible.

mlir/lib/IR/BuiltinDialectBytecode.cpp

29–37

Yeah, I didn't want this schema to try to decipher the underlying implementation, just match what API the BytecodeReader/Writer interface provides. E.g. when we support IntegerAttr I want to just say APInt, not whatever specific implementation we use for encoding that. I suppose if we really wanted to make this explicitly differentiated from an array, we could do something like list<blah>.

36

I think StringAttr is fine to use here.

[akuegel]

It looks like this introduces a dependency cycle between MLIRBytecodeReader and MLIRIR, and MLIRBytecodeWriter and MLIRIR, because now MLIRIR uses the BytecodeImplementation.h header. We can work around that in Bazel by making the BytecodeImplementation.h header part of the "IR" target, but that is not nice.

mlir/docs/BytecodeFormat.md

 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.
+In addition to the assembly format fallback, dialects may also provide a custom
+encoding for their attributes and types. This is useful for both efficiency, as
+custom encoded constructs are significantly smaller and faster to read and
+write, as well as for compatability.
+
+Dialects can opt-in to providing custom encodings by implementing the
+`BytecodeDialectInterface`. This interface provides hooks, namely
+`readAttribute`/`readType` and `writeAttribute`/`writeType`, that will be used
+by the bytecode reader and writer. These hooks are provided a reader and writer
+implementation that can be used to encode various constructs in the underlying
+bytecode format. A unique feature of this interface is that dialects may choose
+to only encode a subset of their attributes and types in a custom bytecode
+format, which can simplify adding new or experimental components that aren't
+fully baked.
+
+When implementing the bytecode interface, dialects are responsible for all
+aspects of the encoding. This includes the indicator for which kind of attribute
+or type is being encoded, the bytecode reader will only know that it has
+encountered an attribute or type of a given dialect, it doesn't encode any
+further information. As such, a common encoding idiom is to use a leading
+`varint` code to indicate how the attribute or type was encoded.
 
 ### IR Section
 
 The IR section contains the encoded form of operations within the bytecode.
 
 #### Operation Encoding
 
 ```

mlir/include/mlir/Bytecode/BytecodeImplementation.h

+//===- BytecodeImplementation.h - MLIR Bytecode Implementation --*- 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 various interfaces and utilities necessary for dialects
+// to hook into bytecode serialization.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_BYTECODE_BYTECODEIMPLEMENTATION_H
+#define MLIR_BYTECODE_BYTECODEIMPLEMENTATION_H
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/DialectInterface.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/Twine.h"
+
+namespace mlir {
+//===----------------------------------------------------------------------===//
+// DialectBytecodeReader
+//===----------------------------------------------------------------------===//
+
+/// This class provides an interface for reading a dialect attribute or type
+/// from a bytecode stream.
+class DialectBytecodeReader {
+public:
+  virtual ~DialectBytecodeReader() = default;
+
+  /// Emit an error to the reader.
+  virtual InFlightDiagnostic emitError(const Twine &msg = {}) = 0;
+
+  //===--------------------------------------------------------------------===//
+  // IR
+  //===--------------------------------------------------------------------===//
+
+  /// Read out a list of elements, invoking the provided callback for each
+  /// element. If T is default constructible, the callback should be of the
+  /// form: `LogicalResult(T &)`. If T is not default constructible, the
+  /// callback should be of the form: `FailureOr<T>()`.
+  template <typename T, typename CallbackFn>
+  LogicalResult readList(SmallVectorImpl<T> &result, CallbackFn &&callback) {
+    uint64_t size;
+    if (failed(readVarInt(size)))
+      return failure();
+    result.reserve(size);
+
+    for (uint64_t i = 0; i < size; ++i) {
+      // If T is default constructible we can create a temporary instance, which
+      // simplifies the logic of the callback.
+      if constexpr (std::is_default_constructible_v<T>) {
+        T element = {};
+        if (failed(callback(element)))
+          return failure();
+        result.emplace_back(std::move(element));
+      } else {
+
+        // Otherwise, we need to route through FailureOr.
+        FailureOr<T> element = callback();
+        if (failed(element))
+          return failure();
+        result.emplace_back(std::move(*element));
+      }
+    }
+    return success();
+  }
+
+  /// Read a reference to the given attribute.
+  virtual LogicalResult readAttribute(Attribute &result) = 0;
+  template <typename T>
+  LogicalResult readAttributes(SmallVectorImpl<T> &attrs) {
+    return readList(attrs, [this](T &attr) { return readAttribute(attr); });
+  }
+
+  /// Read a reference to the given type.
+  virtual LogicalResult readType(Type &result) = 0;
+  template <typename T>
+  LogicalResult readTypes(SmallVectorImpl<T> &types) {
+    return readList(types, [this](T &type) { return readType(type); });
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Primitives
+  //===--------------------------------------------------------------------===//
+
+  /// Read a variable width integer.
+  // TODO: Add a signed variant when necessary.
+  virtual LogicalResult readVarInt(uint64_t &result) = 0;
+
+  /// Read a string from the bytecode.
+  virtual LogicalResult readString(StringRef &result) = 0;
+};
+
+//===----------------------------------------------------------------------===//
+// DialectBytecodeWriter
+//===----------------------------------------------------------------------===//
+
+/// This class provides an interface for writing a dialect attribute or type to
+/// a bytecode stream.
+class DialectBytecodeWriter {
+public:
+  virtual ~DialectBytecodeWriter() = default;
+
+  //===--------------------------------------------------------------------===//
+  // IR
+  //===--------------------------------------------------------------------===//
+
+  /// Write out a list of elements, invoking the provided callback for each
+  /// element.
+  template <typename RangeT, typename CallbackFn>
+  void writeList(RangeT &&range, CallbackFn &&callback) {
+    writeVarInt(llvm::size(range));
+    for (auto &element : range)
+      callback(element);
+  }
+
+  /// Write a reference to the given attribute.
+  virtual void writeAttribute(Attribute attr) = 0;
+  template <typename T>
+  void writeAttributes(ArrayRef<T> attrs) {
+    writeList(attrs, [this](T attr) { writeAttribute(attr); });
+  }
+
+  /// Write a reference to the given type.
+  virtual void writeType(Type type) = 0;
+  template <typename T>
+  void writeTypes(ArrayRef<T> types) {
+    writeList(types, [this](T type) { writeType(type); });
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Primitives
+  //===--------------------------------------------------------------------===//
+
+  /// Write a variable width integer to the output stream. This should be the
+  /// preferred method for emitting integers whenever possible.
+  // TODO: Add a signed variant when necessary.
+  virtual void writeVarInt(uint64_t value) = 0;
+
+  /// Write a string to the bytecode, which is owned by the caller and is
+  /// guaranteed to not die before the end of the bytecode process. This should
+  /// only be called if such a guarantee can be made, such as when the string is
+  /// owned by an attribute or type.
+  virtual void writeOwnedString(StringRef str) = 0;
+};
+
+//===----------------------------------------------------------------------===//
+// BytecodeDialectInterface
+//===----------------------------------------------------------------------===//
+
+class BytecodeDialectInterface
+    : public DialectInterface::Base<BytecodeDialectInterface> {
+public:
+  using Base::Base;
+
+  //===--------------------------------------------------------------------===//
+  // Reading
+  //===--------------------------------------------------------------------===//
+
+  /// Read an attribute belonging to this dialect from the given reader. This
+  /// method should return null in the case of failure.
+  virtual Attribute readAttribute(DialectBytecodeReader &reader) const {
+    reader.emitError() << "dialect " << getDialect()->getNamespace()
+                       << " does not support reading attributes from bytecode";
+    return Attribute();
+  }
+
+  /// Read a type belonging to this dialect from the given reader. This method
+  /// should return null in the case of failure.
+  virtual Type readType(DialectBytecodeReader &reader) const {
+    reader.emitError() << "dialect " << getDialect()->getNamespace()
+                       << " does not support reading types from bytecode";
+    return Type();
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Writing
+  //===--------------------------------------------------------------------===//
+
+  /// Write the given attribute, which belongs to this dialect, to the given
+  /// writer. This method may return failure to indicate that the given
+  /// attribute could not be encoded, in which case the textual format will be
+  /// used to encode this attribute instead.
+  virtual LogicalResult writeAttribute(Attribute attr,
+                                       DialectBytecodeWriter &writer) const {
+    return failure();
+  }
+
+  /// Write the given type, which belongs to this dialect, to the given writer.
+  /// This method may return failure to indicate that the given type could not
+  /// be encoded, in which case the textual format will be used to encode this
+  /// type instead.
+  virtual LogicalResult writeType(Type type,
+                                  DialectBytecodeWriter &writer) const {
+    return failure();
+  }
+};
+
+} // namespace mlir
+
+#endif // MLIR_BYTECODE_BYTECODEIMPLEMENTATION_H

mlir/include/mlir/IR/DialectInterface.h

public:
   /// The base class used for all derived interface types. This class provides
   /// utilities necessary for registration.
   template <typename ConcreteType>
   using Base = detail::DialectInterfaceBase<ConcreteType, DialectInterface>;
 
   /// Return the dialect that this interface represents.
   Dialect *getDialect() const { return dialect; }
 
+  /// Return the context that holds the parent dialect of this interface.
+  MLIRContext *getContext() const;
+
   /// Return the derived interface id.
   TypeID getID() const { return interfaceID; }
 
 protected:
   DialectInterface(Dialect *dialect, TypeID id)
       : dialect(dialect), interfaceID(id) {}
 
 private:

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

 //===- BytecodeReader.cpp - MLIR Bytecode Reader --------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 // TODO: Support for big-endian architectures.
 // TODO: Properly preserve use lists of values.
 
 #include "mlir/Bytecode/BytecodeReader.h"
 #include "../Encoding.h"
 #include "mlir/AsmParser/AsmParser.h"
+#include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Verifier.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Support/MemoryBufferRef.h"
public:
   /// Returns true if the entire section has been read.
   bool empty() const { return dataIt == dataEnd; }
 
   /// Returns the remaining size of the bytecode.
   size_t size() const { return dataEnd - dataIt; }
 
   /// Emit an error using the given arguments.
   template <typename... Args>
-  LogicalResult emitError(Args &&...args) const {
+  InFlightDiagnostic emitError(Args &&...args) const {
     return ::emitError(fileLoc).append(std::forward<Args>(args)...);
   }
 
   /// Parse a single byte from the stream.
   template <typename T>
   LogicalResult parseByte(T &value) {
     if (empty())
       return emitError("attempting to parse a byte at the end of the bytecode");
LogicalResult load(EncodingReader &reader, MLIRContext *ctx) {
     if (!loadedDialect && !ctx->allowsUnregisteredDialects()) {
       return reader.emitError(
           "dialect '", name,
           "' is unknown. If this is intended, please call "
           "allowUnregisteredDialects() on the MLIRContext, or use "
           "-allow-unregistered-dialect with the MLIR tool used.");
     }
     dialect = loadedDialect;
+
+    // If the dialect was actually loaded, check to see if it has a bytecode
+    // interface.
+    if (loadedDialect)
+      interface = dyn_cast<BytecodeDialectInterface>(loadedDialect);
     return success();
   }
 
   /// The loaded dialect entry. This field is None if we haven't attempted to
   /// load, nullptr if we failed to load, otherwise the loaded dialect.
   Optional<Dialect *> dialect;
 
+  /// The bytecode interface for the parent dialect, or nullptr if the parent
+  /// dialect does not implement the bytecode interface. This field should only
+  /// be checked if the `dialect` field is non-None.
+  const BytecodeDialectInterface *interface = nullptr;
+
   /// The name of the dialect.
   StringRef name;
 };
 
 /// This struct represents an operation name entry within the bytecode.
 struct BytecodeOperationName {
   BytecodeOperationName(BytecodeDialect *dialect, StringRef name)
       : dialect(dialect), name(name) {}
struct Entry {
     bool hasCustomEncoding = false;
     /// The raw data of this entry in the bytecode.
     ArrayRef<uint8_t> data;
   };
   using AttrEntry = Entry<Attribute>;
   using TypeEntry = Entry<Type>;
 
 public:
-  AttrTypeReader(Location fileLoc) : fileLoc(fileLoc) {}
+  AttrTypeReader(StringSectionReader &stringReader, Location fileLoc)
+      : stringReader(stringReader), fileLoc(fileLoc) {}
 
   /// Initialize the attribute and type information within the reader.
   LogicalResult initialize(MutableArrayRef<BytecodeDialect> dialects,
                            ArrayRef<uint8_t> sectionData,
                            ArrayRef<uint8_t> offsetSectionData);
 
   /// Resolve the attribute or type at the given index. Returns nullptr on
   /// failure.
LogicalResult parseAsmEntry(T &result, EncodingReader &reader,
                               StringRef entryType);
 
   /// Parse an entry using the given reader that was encoded using a custom
   /// bytecode format.
   template <typename T>
   LogicalResult parseCustomEntry(Entry<T> &entry, EncodingReader &reader,
                                  StringRef entryType);
 
+  /// The string section reader used to resolve string references when parsing
+  /// custom encoded attribute/type entries.
+  StringSectionReader &stringReader;
+
   /// The set of attribute and type entries.
   SmallVector<AttrEntry> attributes;
   SmallVector<TypeEntry> types;
 
   /// A location used for error emission.
   Location fileLoc;
 };
+
+class DialectReader : public DialectBytecodeReader {
+public:
+  DialectReader(AttrTypeReader &attrTypeReader,
+                StringSectionReader &stringReader, EncodingReader &reader)
+      : attrTypeReader(attrTypeReader), stringReader(stringReader),
+        reader(reader) {}
+
+  InFlightDiagnostic emitError(const Twine &msg) override {
+    return reader.emitError(msg);
+  }
+
+  //===--------------------------------------------------------------------===//
+  // IR
+  //===--------------------------------------------------------------------===//
+
+  LogicalResult readAttribute(Attribute &result) override {
+    return attrTypeReader.parseAttribute(reader, result);
+  }
+
+  LogicalResult readType(Type &result) override {
+    return attrTypeReader.parseType(reader, result);
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Primitives
+  //===--------------------------------------------------------------------===//
+
+  LogicalResult readVarInt(uint64_t &result) override {
+    return reader.parseVarInt(result);
+  }
+
+  LogicalResult readString(StringRef &result) override {
+    return stringReader.parseString(reader, result);
+  }
+
+private:
+  AttrTypeReader &attrTypeReader;
+  StringSectionReader &stringReader;
+  EncodingReader &reader;
+};
 } // namespace
 
 LogicalResult
 AttrTypeReader::initialize(MutableArrayRef<BytecodeDialect> dialects,
                            ArrayRef<uint8_t> sectionData,
                            ArrayRef<uint8_t> offsetSectionData) {
   EncodingReader offsetReader(offsetSectionData, fileLoc);
 
   // Parse the number of attribute and type entries.
   uint64_t numAttributes, numTypes;
   if (failed(offsetReader.parseVarInt(numAttributes)) ||
       failed(offsetReader.parseVarInt(numTypes)))
     return failure();
   attributes.resize(numAttributes);
   types.resize(numTypes);
 
   // A functor used to accumulate the offsets for the entries in the given
   // range.
   uint64_t currentOffset = 0;
   auto parseEntries = [&](auto &&range) {
     size_t currentIndex = 0, endIndex = range.size();
 
     // Parse an individual entry.
-    auto parseEntryFn = [&](BytecodeDialect *dialect) {
+    auto parseEntryFn = [&](BytecodeDialect *dialect) -> LogicalResult {
       auto &entry = range[currentIndex++];
 
       uint64_t entrySize;
       if (failed(offsetReader.parseVarIntWithFlag(entrySize,
                                                   entry.hasCustomEncoding)))
         return failure();
 
       // Verify that the offset is actually valid.
T AttrTypeReader::resolveEntry(SmallVectorImpl<Entry<T>> &entries, size_t index,
   if (entry.hasCustomEncoding) {
     if (failed(parseCustomEntry(entry, reader, entryType)))
       return T();
   } else if (failed(parseAsmEntry(entry.entry, reader, entryType))) {
     return T();
   }
 
   if (!reader.empty()) {
-    (void)reader.emitError("unexpected trailing bytes after " + entryType +
-                           " entry");
+    reader.emitError("unexpected trailing bytes after " + entryType + " entry");
     return T();
   }
   return entry.entry;
 }
 
 template <typename T>
 LogicalResult AttrTypeReader::parseAsmEntry(T &result, EncodingReader &reader,
                                             StringRef entryType) {
LogicalResult AttrTypeReader::parseAsmEntry(T &result, EncodingReader &reader,
   }
   return success();
 }
 
 template <typename T>
 LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
                                                EncodingReader &reader,
                                                StringRef entryType) {
-  // FIXME: Add support for reading custom attribute/type encodings.
-  return reader.emitError("unexpected Attribute encoding");
+  if (failed(entry.dialect->load(reader, fileLoc.getContext())))
+    return failure();
+
+  // Ensure that the dialect implements the bytecode interface.
+  if (!entry.dialect->interface) {
+    return reader.emitError("dialect '", entry.dialect->name,
+                            "' does not implement the bytecode interface");
+  }
+
+  // Ask the dialect to parse the entry.
+  DialectReader dialectReader(*this, stringReader, reader);
+  if constexpr (std::is_same_v<T, Type>)
+    entry.entry = entry.dialect->interface->readType(dialectReader);
+  else
+    entry.entry = entry.dialect->interface->readAttribute(dialectReader);
+  return success(!!entry.entry);
 }
 
 //===----------------------------------------------------------------------===//
 // Bytecode Reader
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class is used to read a bytecode buffer and translate it into MLIR.
 class BytecodeReader {
 public:
   BytecodeReader(Location fileLoc, const ParserConfig &config)
-      : config(config), fileLoc(fileLoc), attrTypeReader(fileLoc),
+      : config(config), fileLoc(fileLoc), attrTypeReader(stringReader, fileLoc),
         // Use the builtin unrealized conversion cast operation to represent
         // forward references to values that aren't yet defined.
         forwardRefOpState(UnknownLoc::get(config.getContext()),
                           "builtin.unrealized_conversion_cast", ValueRange(),
                           NoneType::get(config.getContext())) {}
 
   /// Read the bytecode defined within `buffer` into the given block.
   LogicalResult read(llvm::MemoryBufferRef buffer, Block *block);

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

 //===- 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/Bytecode/BytecodeImplementation.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>
 
void BytecodeWriter::writeDialectSection(EncodingEmitter &emitter) {
   writeDialectGrouping(dialectEmitter, numberingState.getOpNames(), emitOpName);
 
   emitter.emitSection(bytecode::Section::kDialect, std::move(dialectEmitter));
 }
 
 //===----------------------------------------------------------------------===//
 // Attributes and Types
 
+namespace {
+class DialectWriter : public DialectBytecodeWriter {
+public:
+  DialectWriter(EncodingEmitter &emitter, IRNumberingState &numberingState,
+                StringSectionBuilder &stringSection)
+      : emitter(emitter), numberingState(numberingState),
+        stringSection(stringSection) {}
+
+  //===--------------------------------------------------------------------===//
+  // IR
+  //===--------------------------------------------------------------------===//
+
+  void writeAttribute(Attribute attr) override {
+    emitter.emitVarInt(numberingState.getNumber(attr));
+  }
+  void writeType(Type type) override {
+    emitter.emitVarInt(numberingState.getNumber(type));
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Primitives
+  //===--------------------------------------------------------------------===//
+
+  void writeVarInt(uint64_t value) override { emitter.emitVarInt(value); }
+
+  void writeOwnedString(StringRef str) override {
+    emitter.emitVarInt(stringSection.insert(str));
+  }
+
+private:
+  EncodingEmitter &emitter;
+  IRNumberingState &numberingState;
+  StringSectionBuilder &stringSection;
+};
+} // namespace
+
 void BytecodeWriter::writeAttrTypeSection(EncodingEmitter &emitter) {
   EncodingEmitter attrTypeEmitter;
   EncodingEmitter offsetEmitter;
   offsetEmitter.emitVarInt(llvm::size(numberingState.getAttributes()));
   offsetEmitter.emitVarInt(llvm::size(numberingState.getTypes()));
 
+  // The writer used when emitting using a custom bytecode encoding.
+  DialectWriter dialectWriter(attrTypeEmitter, numberingState, stringSection);
+
   // 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.
+    auto entryValue = entry.getValue();
+
+    // First, try to emit this entry using the dialect bytecode interface.
     bool hasCustomEncoding = false;
+    if (const BytecodeDialectInterface *interface = entry.dialect->interface) {
+      if constexpr (std::is_same_v<std::decay_t<decltype(entryValue)>, Type>) {
+        // TODO: We don't currently support custom encoded mutable types.
+        hasCustomEncoding =
+            !entryValue.template hasTrait<TypeTrait::IsMutable>() &&
+            succeeded(interface->writeType(entryValue, dialectWriter));
+      } else {
+        // TODO: We don't currently support custom encoded mutable attributes.
+        hasCustomEncoding =
+            !entryValue.template hasTrait<AttributeTrait::IsMutable>() &&
+            succeeded(interface->writeAttribute(entryValue, dialectWriter));
+      }
+    }
 
-    // Emit the entry using the textual format.
-    raw_emitter_ostream(attrTypeEmitter) << entry.getValue();
+    // If the entry was not emitted using the dialect interface, emit it using
+    // the textual format.
+    if (!hasCustomEncoding) {
+      raw_emitter_ostream(attrTypeEmitter) << entryValue;
       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.

mlir/lib/Bytecode/Writer/IRNumbering.h

 
 #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 BytecodeDialectInterface;
 class BytecodeWriterConfig;
 
 namespace bytecode {
 namespace detail {
 struct DialectNumbering;
 
 //===----------------------------------------------------------------------===//
 // Attribute and Type Numbering
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;
+  /// The bytecode dialect interface of the dialect if defined.
+  const BytecodeDialectInterface *interface = nullptr;
 };
 
 //===----------------------------------------------------------------------===//
 // IRNumberingState
 //===----------------------------------------------------------------------===//
 
 /// This class manages numbering IR entities in preparation of bytecode
 /// emission.
public:
 
   /// Return the number of operations in the given block.
   unsigned getOperationCount(Block *block) {
     assert(blockOperationCounts.count(block) && "block not numbered");
     return blockOperationCounts[block];
   }
 
 private:
+  /// This class is used to provide a fake dialect writer for numbering nested
+  /// attributes and types.
+  struct NumberingDialectWriter;
+
   /// 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);

mlir/lib/Bytecode/Writer/IRNumbering.cpp

 //===- IRNumbering.cpp - MLIR Bytecode IR numbering -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "IRNumbering.h"
+#include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 
 using namespace mlir;
 using namespace mlir::bytecode::detail;
 
 //===----------------------------------------------------------------------===//
+// NumberingDialectWriter
+//===----------------------------------------------------------------------===//
+
+struct IRNumberingState::NumberingDialectWriter : public DialectBytecodeWriter {
+  NumberingDialectWriter(IRNumberingState &state) : state(state) {}
+
+  void writeAttribute(Attribute attr) override { state.number(attr); }
+  void writeType(Type type) override { state.number(type); }
+
+  /// Stubbed out methods that are not used for numbering.
+  void writeVarInt(uint64_t) override {}
+  void writeOwnedString(StringRef) override {
+    // TODO: It might be nice to prenumber strings and sort by the number of
+    // references. This could potentially be useful for optimizing things like
+    // file locations.
+  }
+
+  /// The parent numbering state that is populated by this writer.
+  IRNumberingState &state;
+};
+
+//===----------------------------------------------------------------------===//
 // 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) {
void IRNumberingState::number(Attribute attr) {
   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>())
+  if (OpaqueAttr opaqueAttr = attr.dyn_cast<OpaqueAttr>()) {
     numbering->dialect = &numberDialect(opaqueAttr.getDialectNamespace());
-  else
-    numbering->dialect = &numberDialect(&attr.getDialect());
+    return;
+  }
+
+  // If this attribute will be emitted using the bytecode format, perform a
+  // dummy writing to number any nested components.
+  if (const auto *interface = numbering->dialect->interface) {
+    // TODO: We don't allow custom encodings for mutable attributes right now.
+    if (attr.hasTrait<AttributeTrait::IsMutable>())
+      return;
+
+    NumberingDialectWriter writer(*this);
+    (void)interface->writeAttribute(attr, writer);
+  }
 }
 
 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;
+    numbering->interface = dyn_cast<BytecodeDialectInterface>(dialect);
   }
   return *numbering;
 }
 
 auto IRNumberingState::numberDialect(StringRef dialect) -> DialectNumbering & {
   DialectNumbering *&numbering = dialects[dialect];
   if (!numbering) {
     numbering = new (dialectAllocator.Allocate())
void IRNumberingState::number(Type type) {
   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>())
+  if (OpaqueType opaqueType = type.dyn_cast<OpaqueType>()) {
     numbering->dialect = &numberDialect(opaqueType.getDialectNamespace());
-  else
-    numbering->dialect = &numberDialect(&type.getDialect());
+    return;
+  }
+
+  // If this type will be emitted using the bytecode format, perform a dummy
+  // writing to number any nested components.
+  if (const auto *interface = numbering->dialect->interface) {
+    // TODO: We don't allow custom encodings for mutable types right now.
+    if (type.hasTrait<TypeTrait::IsMutable>())
+      return;
+
+    NumberingDialectWriter writer(*this);
+    (void)interface->writeType(type, writer);
+  }
 }

mlir/lib/IR/BuiltinDialect.cpp

 //===- BuiltinDialect.cpp - MLIR Builtin Dialect --------------------------===//
 //
 // 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 the Builtin dialect that contains all of the attributes,
 // operations, and types that are necessary for the validity of the IR.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/BuiltinDialect.h"
+#include "BuiltinDialectBytecode.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeRange.h"
void BuiltinDialect::initialize() {
   registerLocationAttributes();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/IR/BuiltinOps.cpp.inc"
       >();
 
   auto &blobInterface = addInterface<BuiltinBlobManagerInterface>();
   addInterface<BuiltinOpAsmDialectInterface>(blobInterface);
+  builtin_dialect_detail::addBytecodeInterface(this);
 }
 
 //===----------------------------------------------------------------------===//
 // ModuleOp
 //===----------------------------------------------------------------------===//
 
 void ModuleOp::build(OpBuilder &builder, OperationState &state,
                      Optional<StringRef> name) {

mlir/lib/IR/BuiltinDialectBytecode.h

+//===- BuiltinDialectBytecode.h - MLIR Bytecode Implementation --*- 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 hooks into the builtin dialect bytecode implementation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LIB_MLIR_IR_BUILTINDIALECTBYTECODE_H
+#define LIB_MLIR_IR_BUILTINDIALECTBYTECODE_H
+
+namespace mlir {
+class BuiltinDialect;
+
+namespace builtin_dialect_detail {
+/// Add the interfaces necessary for encoding the builtin dialect components in
+/// bytecode.
+void addBytecodeInterface(BuiltinDialect *dialect);
+} // namespace builtin_dialect_detail
+} // namespace mlir
+
+#endif // LIB_MLIR_IR_BUILTINDIALECTBYTECODE_H

mlir/lib/IR/BuiltinDialectBytecode.cpp

+//===- BuiltinDialectBytecode.cpp - Builtin Bytecode Implementation -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "BuiltinDialectBytecode.h"
+#include "mlir/Bytecode/BytecodeImplementation.h"
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Encoding
+//===----------------------------------------------------------------------===//
+
+namespace {
+namespace builtin_encoding {
+/// This enum contains marker codes used to indicate which attribute is
+/// currently being decoded, and how it should be decoded. The order of these
+/// codes should generally be unchanged, as any changes will inevitably break
+/// compatibility with older bytecode.
+enum AttributeCode {
+  ///   ArrayAttr {
+  ///     elements: Attribute[]
+  ///   }
+  ///
+  kArrayAttr = 0,
+
+  ///   DictionaryAttr {
+  ///     attrs: <Attribute, Attribute>[]
+  ///   }
+  kDictionaryAttr = 1,
+
+  ///   StringAttr {
+  ///     string
+  ///   }
+  kStringAttr = 2,
+};
+
+/// This enum contains marker codes used to indicate which type is currently
+/// being decoded, and how it should be decoded. The order of these codes should
+/// generally be unchanged, as any changes will inevitably break compatibility
+/// with older bytecode.
+enum TypeCode {
+  ///   IntegerType {
+  ///     widthAndSignedness: varint // (width << 2) | (signedness)
+  ///   }
+  ///
+  kIntegerType = 0,
+
+  ///   IndexType {
+  ///   }
+  ///
+  kIndexType = 1,
+
+  ///   FunctionType {
+  ///     inputs: Type[],
+  ///     results: Type[]
+  ///   }
+  ///
+  kFunctionType = 2,
+};
+
+} // namespace builtin_encoding
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// BuiltinDialectBytecodeInterface
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class implements the bytecode interface for the builtin dialect.
+struct BuiltinDialectBytecodeInterface : public BytecodeDialectInterface {
+  BuiltinDialectBytecodeInterface(Dialect *dialect)
+      : BytecodeDialectInterface(dialect) {}
+
+  //===--------------------------------------------------------------------===//
+  // Attributes
+
+  Attribute readAttribute(DialectBytecodeReader &reader) const override;
+  ArrayAttr readArrayAttr(DialectBytecodeReader &reader) const;
+  DictionaryAttr readDictionaryAttr(DialectBytecodeReader &reader) const;
+  StringAttr readStringAttr(DialectBytecodeReader &reader) const;
+
+  LogicalResult writeAttribute(Attribute attr,
+                               DialectBytecodeWriter &writer) const override;
+  void write(ArrayAttr attr, DialectBytecodeWriter &writer) const;
+  void write(DictionaryAttr attr, DialectBytecodeWriter &writer) const;
+  void write(StringAttr attr, DialectBytecodeWriter &writer) const;
+
+  //===--------------------------------------------------------------------===//
+  // Types
+
+  Type readType(DialectBytecodeReader &reader) const override;
+  IntegerType readIntegerType(DialectBytecodeReader &reader) const;
+  FunctionType readFunctionType(DialectBytecodeReader &reader) const;
+
+  LogicalResult writeType(Type type,
+                          DialectBytecodeWriter &writer) const override;
+  void write(IntegerType type, DialectBytecodeWriter &writer) const;
+  void write(FunctionType type, DialectBytecodeWriter &writer) const;
+};
+} // namespace
+
+void builtin_dialect_detail::addBytecodeInterface(BuiltinDialect *dialect) {
+  dialect->addInterfaces<BuiltinDialectBytecodeInterface>();
+}
+
+//===----------------------------------------------------------------------===//
+// Attributes: Reader
+
+Attribute BuiltinDialectBytecodeInterface::readAttribute(
+    DialectBytecodeReader &reader) const {
+  uint64_t code;
+  if (failed(reader.readVarInt(code)))
+    return Attribute();
+  switch (code) {
+  case builtin_encoding::kArrayAttr:
+    return readArrayAttr(reader);
+  case builtin_encoding::kDictionaryAttr:
+    return readDictionaryAttr(reader);
+  case builtin_encoding::kStringAttr:
+    return readStringAttr(reader);
+  default:
+    reader.emitError() << "unknown builtin attribute code: " << code;
+    return Attribute();
+  }
+}
+
+ArrayAttr BuiltinDialectBytecodeInterface::readArrayAttr(
+    DialectBytecodeReader &reader) const {
+  SmallVector<Attribute> elements;
+  if (failed(reader.readAttributes(elements)))
+    return ArrayAttr();
+  return ArrayAttr::get(getContext(), elements);
+}
+
+DictionaryAttr BuiltinDialectBytecodeInterface::readDictionaryAttr(
+    DialectBytecodeReader &reader) const {
+  auto readNamedAttr = [&]() -> FailureOr<NamedAttribute> {
+    StringAttr name;
+    Attribute value;
+    if (failed(reader.readAttribute(name)) ||
+        failed(reader.readAttribute(value)))
+      return failure();
+    return NamedAttribute(name, value);
+  };
+  SmallVector<NamedAttribute> attrs;
+  if (failed(reader.readList(attrs, readNamedAttr)))
+    return DictionaryAttr();
+  return DictionaryAttr::get(getContext(), attrs);
+}
+
+StringAttr BuiltinDialectBytecodeInterface::readStringAttr(
+    DialectBytecodeReader &reader) const {
+  StringRef string;
+  if (failed(reader.readString(string)))
+    return StringAttr();
+  return StringAttr::get(getContext(), string);
+}
+
+//===----------------------------------------------------------------------===//
+// Attributes: Writer
+
+LogicalResult BuiltinDialectBytecodeInterface::writeAttribute(
+    Attribute attr, DialectBytecodeWriter &writer) const {
+  return TypeSwitch<Attribute, LogicalResult>(attr)
+      .Case<ArrayAttr, DictionaryAttr, StringAttr>([&](auto attr) {
+        write(attr, writer);
+        return success();
+      })
+      .Default([&](Attribute) { return failure(); });
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    ArrayAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kArrayAttr);
+  writer.writeAttributes(attr.getValue());
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    DictionaryAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kDictionaryAttr);
+  writer.writeList(attr.getValue(), [&](NamedAttribute attr) {
+    writer.writeAttribute(attr.getName());
+    writer.writeAttribute(attr.getValue());
+  });
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    StringAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kStringAttr);
+  writer.writeOwnedString(attr.getValue());
+}
+
+//===----------------------------------------------------------------------===//
+// Types: Reader
+
+Type BuiltinDialectBytecodeInterface::readType(
+    DialectBytecodeReader &reader) const {
+  uint64_t code;
+  if (failed(reader.readVarInt(code)))
+    return Type();
+  switch (code) {
+  case builtin_encoding::kIntegerType:
+    return readIntegerType(reader);
+  case builtin_encoding::kIndexType:
+    return IndexType::get(getContext());
+
+  case builtin_encoding::kFunctionType:
+    return readFunctionType(reader);
+  default:
+    reader.emitError() << "unknown builtin type code: " << code;
+    return Type();
+  }
+}
+
+IntegerType BuiltinDialectBytecodeInterface::readIntegerType(
+    DialectBytecodeReader &reader) const {
+  uint64_t encoding;
+  if (failed(reader.readVarInt(encoding)))
+    return IntegerType();
+  return IntegerType::get(
+      getContext(), encoding >> 2,
+      static_cast<IntegerType::SignednessSemantics>(encoding & 0x3));
+}
+
+FunctionType BuiltinDialectBytecodeInterface::readFunctionType(
+    DialectBytecodeReader &reader) const {
+  SmallVector<Type> inputs, results;
+  if (failed(reader.readTypes(inputs)) || failed(reader.readTypes(results)))
+    return FunctionType();
+  return FunctionType::get(getContext(), inputs, results);
+}
+
+//===----------------------------------------------------------------------===//
+// Types: Writer
+
+LogicalResult BuiltinDialectBytecodeInterface::writeType(
+    Type type, DialectBytecodeWriter &writer) const {
+  return TypeSwitch<Type, LogicalResult>(type)
+      .Case<IntegerType, FunctionType>([&](auto type) {
+        write(type, writer);
+        return success();
+      })
+      .Case([&](IndexType) {
+        return writer.writeVarInt(builtin_encoding::kIndexType), success();
+      })
+      .Default([&](Type) { return failure(); });
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    IntegerType type, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kIntegerType);
+  writer.writeVarInt((type.getWidth() << 2) | type.getSignedness());
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    FunctionType type, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kFunctionType);
+  writer.writeTypes(type.getInputs());
+  writer.writeTypes(type.getResults());
+}

mlir/lib/IR/CMakeLists.txt

 add_mlir_library(MLIRIR
   AffineExpr.cpp
   AffineMap.cpp
   AsmPrinter.cpp
   Attributes.cpp
   Block.cpp
   Builders.cpp
   BuiltinAttributeInterfaces.cpp
   BuiltinAttributes.cpp
   BuiltinDialect.cpp
+  BuiltinDialectBytecode.cpp
   BuiltinTypes.cpp
   BuiltinTypeInterfaces.cpp
   Diagnostics.cpp
   Dialect.cpp
   DialectResourceBlobManager.cpp
   Dominance.cpp
   ExtensibleDialect.cpp
   FunctionImplementation.cpp

mlir/lib/IR/Dialect.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // Dialect Interface
 //===----------------------------------------------------------------------===//
 
 DialectInterface::~DialectInterface() = default;
 
+MLIRContext *DialectInterface::getContext() const {
+  return dialect->getContext();
+}
+
 DialectInterfaceCollectionBase::DialectInterfaceCollectionBase(
     MLIRContext *ctx, TypeID interfaceKind) {
   for (auto *dialect : ctx->getLoadedDialects()) {
     if (auto *interface = dialect->getRegisteredInterface(interfaceKind)) {
       interfaces.insert(interface);
       orderedInterfaces.push_back(interface);
     }
   }

mlir/test/Dialect/Builtin/Bytecode/attrs.mlir

+// RUN: mlir-opt -emit-bytecode %s | mlir-opt | FileCheck %s
+
+// Bytecode currently does not support big-endian platforms
+// UNSUPPORTED: s390x-
+
+// CHECK-LABEL: @TestArray
+module @TestArray attributes {
+  // CHECK: bytecode.array = [unit]
+  bytecode.array = [unit]
+} {}
+
+// CHECK-LABEL: @TestString
+module @TestString attributes {
+  // CHECK: bytecode.string = "hello"
+  bytecode.string = "hello"
+} {}

mlir/test/Dialect/Builtin/Bytecode/types.mlir

+// RUN: mlir-opt -emit-bytecode %s | mlir-opt | FileCheck %s
+
+// Bytecode currently does not support big-endian platforms
+// UNSUPPORTED: s390x-
+
+// CHECK-LABEL: @TestInteger
+module @TestInteger attributes {
+  // CHECK: bytecode.int = i1024,
+  // CHECK: bytecode.int1 = si32,
+  // CHECK: bytecode.int2 = ui512
+  bytecode.int = i1024,
+  bytecode.int1 = si32,
+  bytecode.int2 = ui512
+} {}
+
+// CHECK-LABEL: @TestIndex
+module @TestIndex attributes {
+  // CHECK: bytecode.index = index
+  bytecode.index = index
+} {}
+
+// CHECK-LABEL: @TestFunc
+module @TestFunc attributes {
+  // CHECK: bytecode.func = () -> (),
+  // CHECK: bytecode.func1 = (i1) -> i32
+  bytecode.func = () -> (),
+  bytecode.func1 = (i1) -> (i32)
+} {}