Expose callbacks for encoding of types/attributes

[mfrancio]

Summary

[mlir] Expose a mechanism to provide a callback for encoding types and attributes in MLIR bytecode.

Two callbacks are exposed, respectively, to the BytecodeWriterConfig and to the ParserConfig. At bytecode parsing/printing, clients have the ability to specify a callback to be used to optionally read/write the encoding. On failure, fallback path will execute the default parsers and printers for the dialect.

Testing shows how to leverage this functionality to support back-deployment and backward-compatibility usecases when roundtripping to bytecode a client dialect with type/attributes dependencies on upstream.

[rriddle]

I'm not sure I understand the rationale here. This looks like a very awkward side band, and quite invasive. I don't see why we'd want to open up the bytecode for arbitrary encodings, this should be driven solely by the dialect itself.

[mehdi_amini]

mlir/include/mlir/Bytecode/BytecodeWriter.h

63

Seems like not natural to me as an API to anchor this to the dialect, that seems a bit arbitrary to me.

When we talked about it I was thinking something more flat and simpler:

void addOverrideCallback(std::function<bool(Type)> callback) {
}
Smallvector<std::function<bool(Type)>> typeOverrideCallbacks;

And the writer would do:

// First try to process the given type with the provided override
for (auto &callback : typeOverrideCallbacks)
  if (callback(type)) return

// continue with normal emission

[mehdi_amini]

mlir/include/mlir/Bytecode/BytecodeWriter.h

63

I guess what I wrote is not enough: there is more than the exact encoding, there may be some remapping to another dialect as well.

[mfrancio]

In D153383#4436145, @rriddle wrote:

I'm not sure I understand the rationale here. This looks like a very awkward side band, and quite invasive. I don't see why we'd want to open up the bytecode for arbitrary encodings, this should be driven solely by the dialect itself.

Thanks for your feedback. I completely agree in the principle: a dialect should be the only driver of the encoding. However, this comes with a limitation - if a versioned client dialect wants to control its own encoding, there is no way to do it currently without re-defining and reimplementing all the types and attributes. The patch tries to address this problem by exposing a callback, which offers clients the chance to decouple the encoding of types and attributes that are defined as part of the upstream dialects from their upstream encoding, so that a versioned client dialect that uses unversioned upstream types/attributes can maintain forward/backward compatibility independently from the upstream development.

mlir/include/mlir/Bytecode/BytecodeWriter.h

63

I found it as a compelling and explicit way to override Type and Attributes printer/parser of any dialect with a specific encoding. But agreed, it is arbitrary. I can remove this anchor and simplify it a little bit.

[mfrancio]

Remove ability to specify the dialect associated with each callback.

[mfrancio]

mlir/include/mlir/Bytecode/BytecodeWriter.h

63

I think what implemented should be enough to be able to override the encoding of a type/attribute if its definition jumps from one dialect to another - this is true as long as the encoding is always owned by the callback, before and after the jump.

[mfrancio]

Adds a test exercising roundtrip to bytecode with a custom encoding of IntegerType.

[mfrancio]

mlir/test/lib/IR/TestBytecodeCallbacks.cpp

91

@mehdi_amini this gives an idea of what we would do on parsing - we would get the dialect version to parse from the version map. Right now though there is no system to specify a version on writing, and this is the closest I could go :).

We could post-fix this once the proper API exists if you are ok in leaving the TODO, or I can propose an implementation and finalize the work.

[mehdi_amini]

mlir/test/lib/IR/TestBytecodeCallbacks.cpp

52

Can you write Type entryValue and drop the if constexpr in the body?

66

I'd be interested to see an example where you actually take a !test.i32 as input and write it as a builtin IntegerType (and show that we can parse it as such), and vice-versa.

91

LG

[mfrancio]

Adds bytecode roundtrip tests with custom integer types.

[mfrancio]

mlir/test/lib/IR/TestBytecodeCallbacks.cpp

52

No, the callback needs to compile for both Type entryValue and Attribute entryValue. We could do this with two separate callbacks if you feel it's a bit cumbersome to force the use of auto in the callback signature.

66

Test added.

Note that it is not possible to parse !test.i32 as a native builtin integer type (hence, without adding a specific callback for it, or without a custom parser that falls back to the builtin integer type parser, which I did not implement for the sake of this test) because the the "owner" of such encoding is still the test dialect.

[mfrancio]

clang-format commit.

[mehdi_amini]

mlir/test/lib/IR/TestBytecodeCallbacks.cpp

52

What about taking a union of type/attr instead of templating this?
Right now this adds some cognitive complexity that does not seems necessarily justified to me.
Alternatively we could indeed split it in two callback registration, that can be fine as well.

66

I understand we need a callback, but can't you implement the callback in this file to show we can parse !test.i32 as a native builtin integer type?

[mfrancio]

• split callbacks between type and attributes
• expose builtin bytecode dialect interface
• parse explicitly with bytecode dialect interface when testing interoperability with bytecode types/attributes within a callback

[mfrancio]

mlir/test/lib/IR/TestBytecodeCallbacks.cpp

52

Switched to two callbacks - it is a bit more code, but the function signature is cleaner and more explicit.

66

I did export the integer type parser code within the callback - but effectively was not very clear. Right now I exported the bytecode dialect interface and used this explicitly to write/read, which helps showcasing the feature.

[mehdi_amini]

mlir/include/mlir/IR/AsmState.h

83

I would think we should distinguish between "no handling" and "error during parsing" here, that is the API should be able to fail the parsing.

[mfrancio]

• Handle failure in read callback API
• Add test showcasing the feature

[mfrancio]

mlir/include/mlir/IR/AsmState.h

83

Thanks for pointing this out, it is indeed a very useful feature. I changed the API and added a test for it.

[burmako]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1248

How do you envision the contract between customized serialization and deserialization? E.g. how does a consumer of bytecode payload know that a specific payload was generated via serializer callbacks and how do they know where to obtain the corresponding deserializer callbacks?

[mfrancio]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1248

Since the callbacks are driven by the client, it's up to the client to decide. Assuming you are serializing a module that contains a versioned dialect, I would envision handling such scenario through the dialect version.

I don't think upstream dialects should use those callbacks in any way.

[mfrancio]

Simplifies signature of the read callback by leveraging the dialect reader to retrieve context and dialect versions.

[jpienaar]

(forgot to add feedback earlier)

mlir/include/mlir/Bytecode/BytecodeWriter.h

57

Nit: let's keep these called writer rather than printer (could be emitter too as that matches some of the comments).

mlir/include/mlir/IR/AsmState.h

27

Let's keep these sorted

67

I don't think the _ naming convention is used anywhere else here, of you want to signal internal it could be in private section.

73

emitting -> parsing (or reading or ingesting)

562

So these are just flat arrays rather than grouped by (say) type it handles?

mlir/include/mlir/IR/BuiltinDialectBytecode.h

22

I think we should just make this now builtin::detail

24

I was semi in between exposing the interface vs just helper methods generated and then the add method below (the add method doesn't add much given full interface here). But I was thinking of different composition then.

47

Don't know Mehdi s opinion, but I'd probably put this just in builtin namespace along with dialect. This is rather "public" level API for me.

[mehdi_amini]

mlir/include/mlir/IR/AsmState.h

561

Nit: can you make the return type explicit (ArrayRef<...>) here?

I don't think the auto return is widely used in the codebase, and it goes a bit against the "use auto only when the type is obvious from the context" practice.

mlir/include/mlir/IR/BuiltinDialectBytecode.h

47

I agree: in general "detail" namespace aren't meant to contain things to be directly used by clients.
But also stepping back: I'm not convinced this entire file should be publicly exposed. We should be able to have a single entry point in MLIR for writing a type or an attribute.

mlir/lib/IR/BuiltinDialectBytecode.cpp

91

I believe that in general we prefer using namespace in implementation files, and fully qualify function definition?

[mfrancio]

• Makes BuiltinDialectBytecodeInterface private
• Exposes write/read bytecode functions for types and attributes using the builtin encoding
• Addresses few nits and comments

[mfrancio]

mlir/include/mlir/IR/AsmState.h

562

Well, the idea is that you would encode a function that handles the abstract type and not the concrete, then handling the concrete types using a type switch within the callback itself as needed.

I feel passing a callback per concrete type would be more cumbersome to use?

mlir/include/mlir/IR/BuiltinDialectBytecode.h

47

I avoided exposing the interface while adding an entry point to read/write types and attributes using the builtin encoding.

I used the builtin namespace even though I found out it wasn't used for the builtin dialect - it seems that this dialect lives directly at the mlir level.

However, I have the impression that read/write bytecode functions would be best under the builtin namespace to be explicit about the underlying encoding.

[mehdi_amini]

mlir/include/mlir/IR/BuiltinDialectBytecode.h

47

I really meant that we should be able to call « writeAttribute » without knowing the dialect to target: I don’t quite get what is special about the built in dialect here?

[mfrancio]

mlir/include/mlir/IR/BuiltinDialectBytecode.h

47

I got what you mean and yes, write functions for attribute and type could be made dialect agnostic - what made me stop is that it's now clear how that would extend to read functions, since the encoding does not contain dialect info, and there is nothing that prevents different dialects to use the same encoding for different things. I found this incongruence a bit confusing, to a point that it seemed suggesting that it would not fit in the current design?

For context, see for example Builtin and Quantization (those functions are generated through tablegen) - both dialect encode type 1 and there is no way to disambiguate in a "dialect agnostic" fashion.

Builtin:

static Type readType(MLIRContext* context, DialectBytecodeReader &reader) {
  uint64_t kind;
  if (failed(reader.readVarInt(kind)))
    return Type();
  switch (kind) {
    case 0:
      return readIntegerType(context, reader);
    case 1:
      return readIndexType(context, reader);

Quantization:

static Type readType(MLIRContext* context, DialectBytecodeReader &reader) {
  uint64_t kind;
  if (failed(reader.readVarInt(kind)))
    return Type();
  switch (kind) {
    case 1:
      return readAnyQuantizedType(context, reader);
    case 2:
      return readAnyQuantizedTypeWithExpressedType(context, reader);

Furthermore, it is reasonable to expect to have more of such conflicts as the MLIR bytecode gains popularity.

[mehdi_amini]

mlir/include/mlir/IR/BuiltinDialectBytecode.h

47

OK I get it: it is similar to what we do for print/parse textual assembly.

The convention there is that every type/attribute class expose a print/parse method, can we align on the same model for bytecode? We wouldn't need to emit the discriminant for which attribute it is when we know which one to emit!
(similar to textual ASM)

[mfrancio]

Removes readAttribute() and readType() APIs that take a version as argument since the dialect version is available through the dialect reader.

[mfrancio]

mlir/include/mlir/IR/BuiltinDialectBytecode.h

47

A non-breaking exension of the ASM approach to bytecode does not seem straightforward because of our current bytecode serialization approach.

To summarize the textual ASM approach: builtin types/attributes have known tokens that can be used to detect if a type or an attribute is owned by builtin. If it is not, textual ASM uses a special token (!), which triggers emission of an extended type. At parsing, the extended type token is processed first and this triggers the use of the specific dialect parser.

Now porting this approach into bytecode would mean changing how types or attributes are encoded. Types and attributes are grouped together when emitted to bytecode. At parsing, we parse first the dialect owning the group, which will determine the encoding (this is done per-group and not per-attribute, as opposed to the ASM case).

We could implement an API that works similarly to textual ASM, but unless i am missing some details, integrating it into the existing bytecode format would require some additional work if we want to maintain backwards compatibility, which is probably beyond the scope of the current patch. We could attempt doing this in the future and I would be happy to take the work.

[mehdi_amini]

mlir/include/mlir/IR/BuiltinDialectBytecode.h

47

I agree with you, the encoding difference makes it harder here. Seems reasonable.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

315

Formatting only?

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

828–834

Can you extract the above in a lambda "emitAttrOrTypeImpl(..)": you could do a lot of early return which would simplify the control flow and reduce the indentation.
The hasCustomEncoding boolean should disappear basically.

[mfrancio]

Extend callbacks to allow remapping from one dialect group to another when writing/reading bytecode.

[mfrancio]

Refactors emitAttrOrType lambda to simplify control flow.

[mfrancio]

improve few comments

[rriddle]

mlir/include/mlir/Bytecode/BytecodeWriter.h

56

Can you ArrayRef here and below instead?

mlir/include/mlir/IR/AsmState.h

41–44

It feels quite weird to have something bytecode related not in Bytecode/. Why does this need to be here instead of hooked into the BytecodeWriterConfig?

mlir/lib/IR/BuiltinDialectBytecode.cpp

115–130

Why do we need to expose these at all? As opposed to casting the builtin dialect to DialectBytecodeInterface, and going through the normal path?

[mfrancio]

mlir/include/mlir/Bytecode/BytecodeWriter.h

56

Absolutely, I'll push a revision.

mlir/include/mlir/IR/AsmState.h

41–44

I agree and I tried to search for a better location. The issue is that we currently have a unified entry point for the parser (ParserConfig) which works for both text and bytecode. Unless we want to refactor this, at least the reader side of this class need to stay in this header. Since the same logic is used for AsmResourcePrinter (defined here, but used in bytecode writer config), I thought this could work?

mlir/lib/IR/BuiltinDialectBytecode.cpp

115–130

That would require exposing the interface for builtin, which I had originally done in a previous version of the patch :). However, after few rounds of revisions the consensus was that it could have been better to leave the interface as an internal implementation detail and expose single hooks. In previous revisions, I was also asked to try to expose top level entry point for writing type/attributes in a dialect agnostic fashion, similarly to what the textual parser does, but we realized that it was not going to fit in the current design. Hence, exposing those hooks was the closest I could go to that idea, which preserves the original intent of @jpienaar (having the bytecode dialect interface as an internal implementation detail and not exposed).

Those functions here are only needed for the tests, but it is not unreasonable to expect clients to use them. Let me know if you are strongly against leaving as is, and I can revise as necessary.

[rriddle]

mlir/include/mlir/IR/AsmState.h

41–44

Hmrmrm, can we split out all of the bytecode config into a BytecodeReaderConfig and have that in the parser config? Would be good to keep all of the bytecode pieces isolated.

mlir/lib/IR/BuiltinDialectBytecode.cpp

115–130

DialectBytecodeInterface is an already exposed api, you can just do cast<DialectBytecodeInterface>(type.getDialect()) to get the virtual instance. I don't see which part of that requires exposing anything from the builtin dialect?

[mfrancio]

mlir/include/mlir/IR/AsmState.h

41–44

Sure, it seems like a good change. I'll do that and update the patch. Thanks for the feedback.

mlir/lib/IR/BuiltinDialectBytecode.cpp

115–130

I apologize for the confusion - it seems casting directly is not allowed, but I was able to retrieve the virtual instance through getRegisteredInterface<BytecodeDialectInterface>()! For some reason I had assumed at first that type id of the interface was going to be different. Thanks for pointing it out.

[mfrancio]

• Move bytecode related code into Bytecode/. with appropriate header
• Make naming convention uniform (parse<->read, print<->write)
• Rebase

[mfrancio]

mlir/include/mlir/IR/AsmState.h

41–44

I created a new header specific for the bytecode reader config to avoid a circular dependency (AsmState <-> BytecodeReader).

[rriddle]

Nice

mlir/include/mlir/IR/AsmState.h

560

Could we have just an accessor for the BytecodeReaderConfig? Would help really limit the bytecode parts of this API.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

315

Was the formatting off here before?

[mfrancio]

Expose the bytecode reader config directly to the parser config.

[mfrancio]

mlir/include/mlir/IR/AsmState.h

560

done! Thanks for the review.

[mfrancio]

Rebase

mlir/include/mlir/Bytecode/BytecodeImplementation.h

 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/Twine.h"
 
 namespace mlir {
+//===--------------------------------------------------------------------===//
+// Dialect Version Interface.
+//===--------------------------------------------------------------------===//
+
+/// This class is used to represent the version of a dialect, for the purpose
+/// of polymorphic destruction.
+class DialectVersion {
+public:
+  virtual ~DialectVersion() = default;
+};
+
 //===----------------------------------------------------------------------===//
 // DialectBytecodeReader
 //===----------------------------------------------------------------------===//
 
 /// This class defines a virtual interface for reading a bytecode stream,
 /// providing hooks into the bytecode reader. As such, this class should only be
 /// derived and defined by the main bytecode reader, users (i.e. dialects)
 /// should generally only interact with this class via the
 /// BytecodeDialectInterface below.
 class DialectBytecodeReader {
 public:
   virtual ~DialectBytecodeReader() = default;
 
   /// Emit an error to the reader.
-  virtual InFlightDiagnostic emitError(const Twine &msg = {}) = 0;
+  virtual InFlightDiagnostic emitError(const Twine &msg = {}) const = 0;
+
+  /// Retrieve the dialect version by name if available.
+  virtual FailureOr<const DialectVersion *>
+  getDialectVersion(StringRef dialectName) const = 0;
+
+  /// Retrieve the context associated to the reader.
+  virtual MLIRContext *getContext() const = 0;
 
   /// Read out a list of elements, invoking the provided callback for each
   /// element. The callback function may be in any of the following forms:
   ///   * LogicalResult(T &)
   ///   * FailureOr<T>()
   template <typename T, typename CallbackFn>
   LogicalResult readList(SmallVectorImpl<T> &result, CallbackFn &&callback) {
     uint64_t size;
public:
 
   /// Write a bool to the output stream.
   virtual void writeOwnedBool(bool value) = 0;
 
   /// Return the bytecode version being emitted for.
   virtual int64_t getBytecodeVersion() const = 0;
 };
 
-//===--------------------------------------------------------------------===//
-// Dialect Version Interface.
-//===--------------------------------------------------------------------===//
-
-/// This class is used to represent the version of a dialect, for the purpose
-/// of polymorphic destruction.
-class DialectVersion {
-public:
-  virtual ~DialectVersion() = default;
-};
-
 //===----------------------------------------------------------------------===//
 // 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.
+  /// method should return null in the case of failure. Optionally, the dialect
+  /// version can be accessed through the reader.
   virtual Attribute readAttribute(DialectBytecodeReader &reader) const {
     reader.emitError() << "dialect " << getDialect()->getNamespace()
                        << " does not support reading attributes from bytecode";
     return Attribute();
   }
 
-  /// Read a versioned attribute encoding belonging to this dialect from the
-  /// given reader. This method should return null in the case of failure, and
-  /// falls back to the non-versioned reader in case the dialect implements
-  /// versioning but it does not support versioned custom encodings for the
-  /// attributes.
-  virtual Attribute readAttribute(DialectBytecodeReader &reader,
-                                  const DialectVersion &version) const {
-    reader.emitError()
-        << "dialect " << getDialect()->getNamespace()
-        << " does not support reading versioned 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.
+  /// should return null in the case of failure. Optionally, the dialect version
+  /// can be accessed thorugh the reader.
   virtual Type readType(DialectBytecodeReader &reader) const {
     reader.emitError() << "dialect " << getDialect()->getNamespace()
                        << " does not support reading types from bytecode";
     return Type();
   }
 
-  /// Read a versioned type encoding belonging to this dialect from the given
-  /// reader. This method should return null in the case of failure, and
-  /// falls back to the non-versioned reader in case the dialect implements
-  /// versioning but it does not support versioned custom encodings for the
-  /// types.
-  virtual Type readType(DialectBytecodeReader &reader,
-                        const DialectVersion &version) const {
-    reader.emitError()
-        << "dialect " << getDialect()->getNamespace()
-        << " does not support reading versioned 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.

mlir/include/mlir/Bytecode/BytecodeWriter.h

public:
   /// the desired version. The bytecode writer entry point will return failure
   /// if it cannot emit the desired version.
   void setDesiredBytecodeVersion(int64_t bytecodeVersion);
 
   /// Get the set desired bytecode version to emit.
   int64_t getDesiredBytecodeVersion() const;
 
   //===--------------------------------------------------------------------===//
+  // Types and Attributes encoding
+  //===--------------------------------------------------------------------===//
+
+  /// Retrieve the callbacks.
+  llvm::SmallVector<std::unique_ptr<AsmAttrTypeBytecodeWriter<Attribute>>> &
+  getAttributeWriterCallbacks() const;
+  llvm::SmallVector<std::unique_ptr<AsmAttrTypeBytecodeWriter<Type>>> &
+  getTypeWriterCallbacks() const;
+
+  /// Attach a custom bytecode printer callback to the configuration for the
+  /// emission of custom type/attributes encodings.
+  void attachAttributeCallback(
+      std::unique_ptr<AsmAttrTypeBytecodeWriter<Attribute>> callback);
+  void
+  attachTypeCallback(std::unique_ptr<AsmAttrTypeBytecodeWriter<Type>> callback);
+
+  /// Attach a custom bytecode printer callback to the configuration for the
+  /// emission of custom type/attributes encodings.
+  template <typename CallableT>
+  std::enable_if_t<std::is_convertible_v<
+      CallableT,
+      std::function<LogicalResult(Attribute, std::optional<StringRef> &,
+                                  DialectBytecodeWriter &)>>>
+  attachAttributeCallback(CallableT &&emitFn) {
+    attachAttributeCallback(AsmAttrTypeBytecodeWriter<Attribute>::fromCallable(
+        std::forward<CallableT>(emitFn)));
+  }
+  template <typename CallableT>
+  std::enable_if_t<std::is_convertible_v<
+      CallableT, std::function<LogicalResult(Type, std::optional<StringRef> &,
+                                             DialectBytecodeWriter &)>>>
+  attachTypeCallback(CallableT &&emitFn) {
+    attachTypeCallback(AsmAttrTypeBytecodeWriter<Type>::fromCallable(
+        std::forward<CallableT>(emitFn)));
+  }
+
+  //===--------------------------------------------------------------------===//
   // Resources
   //===--------------------------------------------------------------------===//
 
   /// Attach the given resource printer to the writer configuration.
   void attachResourcePrinter(std::unique_ptr<AsmResourcePrinter> printer);
 
   /// Attach an resource printer, in the form of a callable, to the
   /// configuration.

mlir/include/mlir/IR/AsmState.h

 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/StringMap.h"
 
 #include <memory>
 #include <variant>
 
 namespace mlir {
-class AsmResourcePrinter;
 class AsmDialectResourceHandle;
+class AsmResourcePrinter;
+class DialectBytecodeReader;
+class DialectBytecodeWriter;
+class DialectVersion;
 class Operation;
 
 namespace detail {
 class AsmStateImpl;
 } // namespace detail
 
 //===----------------------------------------------------------------------===//
+// AsmAttrTypeBytecode Parser/Printer
+//===----------------------------------------------------------------------===//
+
+/// A class to interact with the attributes and types printer when emitting MLIR
+/// bytecode.
+template <class T>
+class AsmAttrTypeBytecodeWriter {
+public:
+  AsmAttrTypeBytecodeWriter() = default;
+  virtual ~AsmAttrTypeBytecodeWriter() = default;
+
+  virtual LogicalResult write(T entry, std::optional<StringRef> &name,
+                              DialectBytecodeWriter &writer) = 0;
+
+  LogicalResult write(T entry, DialectBytecodeWriter &writer) {
+    std::optional<StringRef> dummy;
+    return write(entry, dummy, writer);
+  }
+
+  /// Return an Attribute/Type printer implemented via the given callable, whose
+  /// form should match that of the `write` function above.
+  template <typename CallableT,
+            std::enable_if_t<std::is_convertible_v<
+                                 CallableT, std::function<LogicalResult(
+                                                T, std::optional<StringRef> &,
+                                                DialectBytecodeWriter &)>>,
+                             bool> = true>
+  static std::unique_ptr<AsmAttrTypeBytecodeWriter<T>>
+  fromCallable(CallableT &&writeFn) {
+    struct Processor : public AsmAttrTypeBytecodeWriter<T> {
+      Processor(CallableT &&writeFn)
+          : AsmAttrTypeBytecodeWriter(), writeFn(std::move(writeFn)) {}
+      LogicalResult write(T entry, std::optional<StringRef> &name,
+                          DialectBytecodeWriter &writer) override {
+        return writeFn(entry, name, writer);
+      }
+
+      std::decay_t<CallableT> writeFn;
+    };
+    return std::make_unique<Processor>(std::forward<CallableT>(writeFn));
+  }
+};
+
+/// A class to interact with the attributes and types parser when parsing MLIR
+/// bytecode.
+template <class T>
+class AsmAttrTypeBytecodeParser {
+public:
+  AsmAttrTypeBytecodeParser() = default;
+  virtual ~AsmAttrTypeBytecodeParser() = default;
+
+  virtual LogicalResult parse(DialectBytecodeReader &reader,
+                              StringRef dialectName, T &entry) = 0;
+
+  /// Return an Attribute/Type printer implemented via the given callable, whose
+  /// form should match that of the `parse` function above.
+  template <typename CallableT,
+            std::enable_if_t<
+                std::is_convertible_v<
+                    CallableT, std::function<LogicalResult(
+                                   DialectBytecodeReader &, StringRef, T &)>>,
+                bool> = true>
+  static std::unique_ptr<AsmAttrTypeBytecodeParser<T>>
+  fromCallable(CallableT &&parseFn) {
+    struct Processor : public AsmAttrTypeBytecodeParser<T> {
+      Processor(CallableT &&parseFn)
+          : AsmAttrTypeBytecodeParser(), parseFn(std::move(parseFn)) {}
+      LogicalResult parse(DialectBytecodeReader &reader, StringRef dialectName,
+                          T &entry) override {
+        return parseFn(reader, dialectName, entry);
+      }
+
+      std::decay_t<CallableT> parseFn;
+    };
+    return std::make_unique<Processor>(std::forward<CallableT>(parseFn));
+  }
+};
+
+//===----------------------------------------------------------------------===//
 // Resources
 //===----------------------------------------------------------------------===//
 
 /// The following classes enable support for parsing and printing resources
 /// within MLIR assembly formats. Resources are a mechanism by which dialects,
 /// and external clients, may attach additional information when parsing or
 /// printing IR without that information being encoded in the IR itself.
 /// Resources are not uniqued within the MLIR context, are not attached directly
public:
   }
 
   /// Return the MLIRContext to be used when parsing.
   MLIRContext *getContext() const { return context; }
 
   /// Returns if the parser should verify the IR after parsing.
   bool shouldVerifyAfterParse() const { return verifyAfterParse; }
 
+  /// Returns the callbacks available to the parser.
+  ArrayRef<std::unique_ptr<AsmAttrTypeBytecodeParser<Attribute>>>
+  getAttributeBytecodeCallbacks() const {
+    return attributeBytecodeParsers;
+  }
+  ArrayRef<std::unique_ptr<AsmAttrTypeBytecodeParser<Type>>>
+  getTypeBytecodeCallbacks() const {
+    return typeBytecodeParsers;
+  }
+
+  /// Attach a custom bytecode parser callback to the configuration for parsing
+  /// of custom type/attributes encodings.
+  void attachAttributeBytecodeCallback(
+      std::unique_ptr<AsmAttrTypeBytecodeParser<Attribute>> parser) {
+    attributeBytecodeParsers.emplace_back(std::move(parser));
+  }
+  void attachTypeBytecodeCallback(
+      std::unique_ptr<AsmAttrTypeBytecodeParser<Type>> parser) {
+    typeBytecodeParsers.emplace_back(std::move(parser));
+  }
+
+  /// Attach a custom bytecode parser callback to the configuration for parsing
+  /// of custom type/attributes encodings.
+  template <typename CallableT>
+  std::enable_if_t<std::is_convertible_v<
+      CallableT, std::function<LogicalResult(DialectBytecodeReader &, StringRef,
+                                             Attribute &)>>>
+  attachAttributeBytecodeCallback(CallableT &&parserFn) {
+    attachAttributeBytecodeCallback(
+        AsmAttrTypeBytecodeParser<Attribute>::fromCallable(
+            std::forward<CallableT>(parserFn)));
+  }
+  template <typename CallableT>
+  std::enable_if_t<std::is_convertible_v<
+      CallableT,
+      std::function<LogicalResult(DialectBytecodeReader &, StringRef, Type &)>>>
+  attachTypeBytecodeCallback(CallableT &&parserFn) {
+    attachTypeBytecodeCallback(AsmAttrTypeBytecodeParser<Type>::fromCallable(
+        std::forward<CallableT>(parserFn)));
+  }
+
   /// Return the resource parser registered to the given name, or nullptr if no
   /// parser with `name` is registered.
   AsmResourceParser *getResourceParser(StringRef name) const {
     auto it = resourceParsers.find(name);
     if (it != resourceParsers.end())
       return it->second.get();
     if (fallbackResourceMap)
       return &fallbackResourceMap->getParserFor(name);
attachResourceParser(AsmResourceParser::fromCallable(
         name, std::forward<CallableT>(parserFn)));
   }
 
 private:
   MLIRContext *context;
   bool verifyAfterParse;
   DenseMap<StringRef, std::unique_ptr<AsmResourceParser>> resourceParsers;
   FallbackAsmResourceMap *fallbackResourceMap;
+  llvm::SmallVector<std::unique_ptr<AsmAttrTypeBytecodeParser<Attribute>>>
+      attributeBytecodeParsers;
+  llvm::SmallVector<std::unique_ptr<AsmAttrTypeBytecodeParser<Type>>>
+      typeBytecodeParsers;
 };
 
 //===----------------------------------------------------------------------===//
 // AsmState
 //===----------------------------------------------------------------------===//
 
 /// This class provides management for the lifetime of the state used when
 /// printing the IR. It allows for alleviating the cost of recomputing the

mlir/include/mlir/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
+#ifndef MLIR_IR_BUILTINDIALECTBYTECODE_H
+#define MLIR_IR_BUILTINDIALECTBYTECODE_H
+
+#include "mlir/Bytecode/BytecodeImplementation.h"
 
 namespace mlir {
 class BuiltinDialect;
 
-namespace builtin_dialect_detail {
+namespace builtin {
+/// Utility read/write functions for types and attributes based on the builtin
+/// bytecode encoding.
+Attribute readAttribute(DialectBytecodeReader &reader);
+LogicalResult writeAttribute(Attribute attribute,
+                             DialectBytecodeWriter &writer);
+Type readType(DialectBytecodeReader &reader);
+LogicalResult writeType(Type type, DialectBytecodeWriter &writer);
+
 /// Add the interfaces necessary for encoding the builtin dialect components in
 /// bytecode.
 void addBytecodeInterface(BuiltinDialect *dialect);
-} // namespace builtin_dialect_detail
+} // namespace builtin
 } // namespace mlir
 
 #endif // LIB_MLIR_IR_BUILTINDIALECTBYTECODE_H

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

LLVM_ATTRIBUTE_NOINLINE LogicalResult parseMultiByteVarInt(uint64_t &result) {
     // implementation).
     uint32_t numBytes = llvm::countr_zero<uint32_t>(result);
     assert(numBytes > 0 && numBytes <= 7 &&
            "unexpected number of trailing zeros in varint encoding");
 
     // Parse in the remaining bytes of the value.
     llvm::support::ulittle64_t resultLE(result);
     if (failed(parseBytes(numBytes, reinterpret_cast<uint8_t *>(&resultLE) + 1)))
       return failure();
 
     // Shift out the low-order bits that were used to mark how the value was
     // encoded.
     result = resultLE >> (numBytes + 1);
     return success();
   }
 
   /// The current data iterator, and an iterator to the end of the buffer.

 class DialectReader;
 
 /// This struct represents a dialect entry within the bytecode.
 struct BytecodeDialect {
   /// Load the dialect into the provided context if it hasn't been loaded yet.
   /// Returns failure if the dialect couldn't be loaded *and* the provided
   /// context does not allow unregistered dialects. The provided reader is used
   /// for error emission if necessary.
-  LogicalResult load(DialectReader &reader, MLIRContext *ctx);
+  LogicalResult load(const DialectReader &reader, MLIRContext *ctx);
 
   /// Return the loaded dialect, or nullptr if the dialect is unknown. This can
   /// only be called after `load`.
   Dialect *getLoadedDialect() const {
     assert(dialect &&
            "expected `load` to be invoked before `getLoadedDialect`");
     return *dialect;
   }
struct BytecodeOperationName {
   /// Whether this operation was registered when the bytecode was produced.
   /// This flag is populated when bytecode version >=kNativePropertiesEncoding.
   std::optional<bool> wasRegistered;
 };
 } // namespace
 
 /// Parse a single dialect group encoded in the byte stream.
 static LogicalResult parseDialectGrouping(
-    EncodingReader &reader, MutableArrayRef<BytecodeDialect> dialects,
+    EncodingReader &reader,
+    MutableArrayRef<std::unique_ptr<BytecodeDialect>> dialects,
     function_ref<LogicalResult(BytecodeDialect *)> entryCallback) {
   // Parse the dialect and the number of entries in the group.
-  BytecodeDialect *dialect;
+  std::unique_ptr<BytecodeDialect> *dialect;
   if (failed(parseEntry(reader, dialects, dialect, "dialect")))
     return failure();
   uint64_t numEntries;
   if (failed(reader.parseVarInt(numEntries)))
     return failure();
 
   for (uint64_t i = 0; i < numEntries; ++i)
-    if (failed(entryCallback(dialect)))
+    if (failed(entryCallback(dialect->get())))
       return failure();
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ResourceSectionReader
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class is used to read the resource section from the bytecode.
 class ResourceSectionReader {
 public:
   /// Initialize the resource section reader with the given section data.
   LogicalResult
   initialize(Location fileLoc, const ParserConfig &config,
-             MutableArrayRef<BytecodeDialect> dialects,
+             MutableArrayRef<std::unique_ptr<BytecodeDialect>> dialects,
              StringSectionReader &stringReader, ArrayRef<uint8_t> sectionData,
              ArrayRef<uint8_t> offsetSectionData, DialectReader &dialectReader,
              const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef);
 
   /// Parse a dialect resource handle from the resource section.
   LogicalResult parseResourceHandle(EncodingReader &reader,
                                     AsmDialectResourceHandle &result) {
     return parseEntry(reader, dialectResources, result, "resource handle");
if (!entryReader.empty()) {
           "unexpected trailing bytes in resource entry '", key, "'");
     }
   }
   return success();
 }
 
 LogicalResult ResourceSectionReader::initialize(
     Location fileLoc, const ParserConfig &config,
-    MutableArrayRef<BytecodeDialect> dialects,
+    MutableArrayRef<std::unique_ptr<BytecodeDialect>> dialects,
     StringSectionReader &stringReader, ArrayRef<uint8_t> sectionData,
     ArrayRef<uint8_t> offsetSectionData, DialectReader &dialectReader,
     const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef) {
   EncodingReader resourceReader(sectionData, fileLoc);
   EncodingReader offsetReader(offsetSectionData, fileLoc);
 
   // Read the number of external resource providers.
   uint64_t numExternalResourceGroups;
for (uint64_t i = 0; i < numExternalResourceGroups; ++i) {
 
     if (failed(parseGroup(handler)))
       return failure();
   }
 
   // Read the dialect resources from the bytecode.
   MLIRContext *ctx = fileLoc->getContext();
   while (!offsetReader.empty()) {
-    BytecodeDialect *dialect;
+    std::unique_ptr<BytecodeDialect> *dialect;
     if (failed(parseEntry(offsetReader, dialects, dialect, "dialect")) ||
-        failed(dialect->load(dialectReader, ctx)))
+        failed((*dialect)->load(dialectReader, ctx)))
       return failure();
-    Dialect *loadedDialect = dialect->getLoadedDialect();
+    Dialect *loadedDialect = (*dialect)->getLoadedDialect();
     if (!loadedDialect) {
       return resourceReader.emitError()
-             << "dialect '" << dialect->name << "' is unknown";
+             << "dialect '" << (*dialect)->name << "' is unknown";
     }
     const auto *handler = dyn_cast<OpAsmDialectInterface>(loadedDialect);
     if (!handler) {
       return resourceReader.emitError()
-             << "unexpected resources for dialect '" << dialect->name << "'";
+             << "unexpected resources for dialect '" << (*dialect)->name << "'";
     }
 
     // Ensure that each resource is declared before being processed.
     auto processResourceKeyFn = [&](StringRef key) -> LogicalResult {
       FailureOr<AsmDialectResourceHandle> handle =
           handler->declareResource(key);
       if (failed(handle)) {
         return resourceReader.emitError()
                << "unknown 'resource' key '" << key << "' for dialect '"
-               << dialect->name << "'";
+               << (*dialect)->name << "'";
       }
       dialectResourceHandleRenamingMap[key] = handler->getResourceKey(*handle);
       dialectResources.push_back(*handle);
       return success();
     };
 
     // Parse the resources for this dialect. We allow empty resources because we
     // just treat these as declarations.
struct Entry {
     /// The raw data of this entry in the bytecode.
     ArrayRef<uint8_t> data;
   };
   using AttrEntry = Entry<Attribute>;
   using TypeEntry = Entry<Type>;
 
 public:
   AttrTypeReader(StringSectionReader &stringReader,
-                 ResourceSectionReader &resourceReader, Location fileLoc)
+                 ResourceSectionReader &resourceReader,
+                 const llvm::StringMap<BytecodeDialect *> &dialectsMap,
+                 Location fileLoc, const ParserConfig &config)
       : stringReader(stringReader), resourceReader(resourceReader),
-        fileLoc(fileLoc) {}
+        dialectsMap(dialectsMap), fileLoc(fileLoc), parserConfig(config) {}
 
   /// Initialize the attribute and type information within the reader.
-  LogicalResult initialize(MutableArrayRef<BytecodeDialect> dialects,
-                           ArrayRef<uint8_t> sectionData,
-                           ArrayRef<uint8_t> offsetSectionData);
+  LogicalResult
+  initialize(MutableArrayRef<std::unique_ptr<BytecodeDialect>> dialects,
 
   /// Resolve the attribute or type at the given index. Returns nullptr on
   /// failure.
   Attribute resolveAttribute(size_t index) {
     return resolveEntry(attributes, index, "Attribute");
   }
   Type resolveType(size_t index) { return resolveEntry(types, index, "Type"); }
 
private:
   /// The string section reader used to resolve string references when parsing
   /// custom encoded attribute/type entries.
   StringSectionReader &stringReader;
 
   /// The resource section reader used to resolve resource references when
   /// parsing custom encoded attribute/type entries.
   ResourceSectionReader &resourceReader;
 
+  /// The map of the loaded dialects used to retrieve dialect information, such
+  /// as the dialect version.
+  const llvm::StringMap<BytecodeDialect *> &dialectsMap;
+
   /// The set of attribute and type entries.
   SmallVector<AttrEntry> attributes;
   SmallVector<TypeEntry> types;
 
   /// A location used for error emission.
   Location fileLoc;
+
+  /// Reference to the parser configuration.
+  const ParserConfig &parserConfig;
 };
 
 class DialectReader : public DialectBytecodeReader {
 public:
   DialectReader(AttrTypeReader &attrTypeReader,
                 StringSectionReader &stringReader,
-                ResourceSectionReader &resourceReader, EncodingReader &reader)
+                ResourceSectionReader &resourceReader,
+                const llvm::StringMap<BytecodeDialect *> &dialectsMap,
+                EncodingReader &reader)
       : attrTypeReader(attrTypeReader), stringReader(stringReader),
-        resourceReader(resourceReader), reader(reader) {}
+        resourceReader(resourceReader), dialectsMap(dialectsMap),
+        reader(reader) {}
 
-  InFlightDiagnostic emitError(const Twine &msg) override {
+  InFlightDiagnostic emitError(const Twine &msg) const override {
     return reader.emitError(msg);
   }
 
-  DialectReader withEncodingReader(EncodingReader &encReader) {
+  FailureOr<const DialectVersion *>
+  getDialectVersion(StringRef dialectName) const override {
+    // First check if the dialect is available in the map.
+    auto dialectEntry = dialectsMap.find(dialectName);
+    if (dialectEntry == dialectsMap.end())
+      return failure();
+    // If the dialect was found, try to load it. This will trigger reading the
+    // bytecode version from the version buffer if it wasn't already processed.
+    // Return failure if either of those two actions could not be completed.
+    if (failed(dialectEntry->getValue()->load(*this, getLoc().getContext())) ||
+        dialectEntry->getValue()->loadedVersion.get() == nullptr)
+      return failure();
+    return dialectEntry->getValue()->loadedVersion.get();
+  }
+
+  MLIRContext *getContext() const override { return getLoc().getContext(); }
+
+  DialectReader withEncodingReader(EncodingReader &encReader) const {
     return DialectReader(attrTypeReader, stringReader, resourceReader,
-                         encReader);
+                         dialectsMap, encReader);
   }
 
   Location getLoc() const { return reader.getLoc(); }
 
   //===--------------------------------------------------------------------===//
   // IR
   //===--------------------------------------------------------------------===//
 
public:
   LogicalResult readBool(bool &result) override {
     return reader.parseByte(result);
   }
 
 private:
   AttrTypeReader &attrTypeReader;
   StringSectionReader &stringReader;
   ResourceSectionReader &resourceReader;
+  const llvm::StringMap<BytecodeDialect *> &dialectsMap;
   EncodingReader &reader;
 };
 
 /// Wraps the properties section and handles reading properties out of it.
 class PropertiesSectionReader {
 public:
   /// Initialize the properties section reader with the given section data.
   LogicalResult initialize(Location fileLoc, ArrayRef<uint8_t> sectionData) {
private:
   /// The properties buffer referenced within the bytecode file.
   ArrayRef<uint8_t> propertiesBuffers;
 
   /// Table of offset in the buffer above.
   SmallVector<int64_t> offsetTable;
 };
 } // namespace
 
-LogicalResult
-AttrTypeReader::initialize(MutableArrayRef<BytecodeDialect> dialects,
-                           ArrayRef<uint8_t> sectionData,
-                           ArrayRef<uint8_t> offsetSectionData) {
+LogicalResult AttrTypeReader::initialize(
+    MutableArrayRef<std::unique_ptr<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);
LogicalResult AttrTypeReader::initialize(
   if (failed(parseEntries(attributes)) || failed(parseEntries(types)))
     return failure();
 
   // Ensure that we read everything from the section.
   if (!offsetReader.empty()) {
     return offsetReader.emitError(
         "unexpected trailing data in the Attribute/Type offset section");
   }
+
   return success();
 }
 
 template <typename T>
 T AttrTypeReader::resolveEntry(SmallVectorImpl<Entry<T>> &entries, size_t index,
                                StringRef entryType) {
   if (index >= entries.size()) {
     emitError(fileLoc) << "invalid " << entryType << " index: " << index;
LogicalResult AttrTypeReader::parseAsmEntry(T &result, EncodingReader &reader,
   }
   return success();
 }
 
 template <typename T>
 LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
                                                EncodingReader &reader,
                                                StringRef entryType) {
-  DialectReader dialectReader(*this, stringReader, resourceReader, reader);
+  DialectReader dialectReader(*this, stringReader, resourceReader, dialectsMap,
+                              reader);
   if (failed(entry.dialect->load(dialectReader, fileLoc.getContext())))
     return failure();
+
+  if constexpr (std::is_same_v<T, Type>) {
+    // Try parsing with callbacks first if available.
+    for (const auto &callback : parserConfig.getTypeBytecodeCallbacks()) {
+      if (failed(
+              callback->parse(dialectReader, entry.dialect->name, entry.entry)))
+        return failure();
+      // Early return if parsing was successful.
+      if (!!entry.entry)
+        return success();
+
+      // Reset the reader if we failed to parse, so we can fall through the
+      // other parsing functions.
+      reader = EncodingReader(entry.data, reader.getLoc());
+    }
+  } else {
+    // Try parsing with callbacks first if available.
+    for (const auto &callback : parserConfig.getAttributeBytecodeCallbacks()) {
+      if (failed(
+              callback->parse(dialectReader, entry.dialect->name, entry.entry)))
+        return failure();
+      // Early return if parsing was successful.
+      if (!!entry.entry)
+        return success();
+
+      // Reset the reader if we failed to parse, so we can fall through the
+      // other parsing functions.
+      reader = EncodingReader(entry.data, reader.getLoc());
+    }
+  }
+
   // 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. If the dialect is versioned, parse
-  // using the versioned encoding readers.
-  if (entry.dialect->loadedVersion.get()) {
-    if constexpr (std::is_same_v<T, Type>)
-      entry.entry = entry.dialect->interface->readType(
-          dialectReader, *entry.dialect->loadedVersion);
-    else
-      entry.entry = entry.dialect->interface->readAttribute(
-          dialectReader, *entry.dialect->loadedVersion);
-
-  } else {
-    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
 //===----------------------------------------------------------------------===//
 
 /// This class is used to read a bytecode buffer and translate it into MLIR.
 class mlir::BytecodeReader::Impl {
   struct RegionReadState;
   using LazyLoadableOpsInfo =
       std::list<std::pair<Operation *, RegionReadState>>;
   using LazyLoadableOpsMap =
       DenseMap<Operation *, LazyLoadableOpsInfo::iterator>;
 
 public:
   Impl(Location fileLoc, const ParserConfig &config, bool lazyLoading,
        llvm::MemoryBufferRef buffer,
        const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef)
       : config(config), fileLoc(fileLoc), lazyLoading(lazyLoading),
-        attrTypeReader(stringReader, resourceReader, fileLoc),
+        attrTypeReader(stringReader, resourceReader, dialectsMap, fileLoc,
+                       config),
         // 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())),
         buffer(buffer), bufferOwnerRef(bufferOwnerRef) {}
 
   /// Read the bytecode defined within `buffer` into the given block.
private:
 
   /// The version of the bytecode being read.
   uint64_t version = 0;
 
   /// The producer of the bytecode being read.
   StringRef producer;
 
   /// The table of IR units referenced within the bytecode file.
-  SmallVector<BytecodeDialect> dialects;
+  SmallVector<std::unique_ptr<BytecodeDialect>> dialects;
+  llvm::StringMap<BytecodeDialect *> dialectsMap;
   SmallVector<BytecodeOperationName> opNames;
 
   /// The reader used to process resources within the bytecode.
   ResourceSectionReader resourceReader;
 
   /// Worklist of values with custom use-list orders to process before the end
   /// of the parsing.
   DenseMap<void *, UseListOrderStorage> valueToUseListMap;
LogicalResult BytecodeReader::Impl::parseVersion(EncodingReader &reader) {
   if (version < bytecode::kLazyLoading)
     lazyLoading = false;
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect Section
 
-LogicalResult BytecodeDialect::load(DialectReader &reader, MLIRContext *ctx) {
+LogicalResult BytecodeDialect::load(const DialectReader &reader,
+                                    MLIRContext *ctx) {
   if (dialect)
     return success();
   Dialect *loadedDialect = ctx->getOrLoadDialect(name);
   if (!loadedDialect && !ctx->allowsUnregisteredDialects()) {
     return reader.emitError("dialect '")
            << name
            << "' is unknown. If this is intended, please call "
               "allowUnregisteredDialects() on the MLIRContext, or use "
BytecodeReader::Impl::parseDialectSection(ArrayRef<uint8_t> sectionData) {
   // Parse the number of dialects in the section.
   uint64_t numDialects;
   if (failed(sectionReader.parseVarInt(numDialects)))
     return failure();
   dialects.resize(numDialects);
 
   // Parse each of the dialects.
   for (uint64_t i = 0; i < numDialects; ++i) {
+    dialects[i] = std::make_unique<BytecodeDialect>();
     /// Before version kDialectVersioning, there wasn't any versioning available
     /// for dialects, and the entryIdx represent the string itself.
     if (version < bytecode::kDialectVersioning) {
-      if (failed(stringReader.parseString(sectionReader, dialects[i].name)))
+      if (failed(stringReader.parseString(sectionReader, dialects[i]->name)))
         return failure();
       continue;
     }
+
     // Parse ID representing dialect and version.
     uint64_t dialectNameIdx;
     bool versionAvailable;
     if (failed(sectionReader.parseVarIntWithFlag(dialectNameIdx,
                                                  versionAvailable)))
       return failure();
     if (failed(stringReader.parseStringAtIndex(sectionReader, dialectNameIdx,
-                                               dialects[i].name)))
+                                               dialects[i]->name)))
       return failure();
     if (versionAvailable) {
       bytecode::Section::ID sectionID;
-      if (failed(
-              sectionReader.parseSection(sectionID, dialects[i].versionBuffer)))
+      if (failed(sectionReader.parseSection(sectionID,
+                                            dialects[i]->versionBuffer)))
         return failure();
       if (sectionID != bytecode::Section::kDialectVersions) {
         emitError(fileLoc, "expected dialect version section");
         return failure();
       }
     }
+    dialectsMap[dialects[i]->name] = dialects[i].get();
   }
 
   // Parse the operation names, which are grouped by dialect.
   auto parseOpName = [&](BytecodeDialect *dialect) {
     StringRef opName;
     std::optional<bool> wasRegistered;
     // Prior to version kNativePropertiesEncoding, the information about wheter
     // an op was registered or not wasn't encoded.
BytecodeReader::Impl::parseOpName(EncodingReader &reader,
   if (failed(parseEntry(reader, opNames, opName, "operation name")))
     return failure();
   wasRegistered = opName->wasRegistered;
   // Check to see if this operation name has already been resolved. If we
   // haven't, load the dialect and build the operation name.
   if (!opName->opName) {
     // Load the dialect and its version.
     DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
-                                reader);
+                                dialectsMap, reader);
     if (failed(opName->dialect->load(dialectReader, getContext())))
       return failure();
     // If the opName is empty, this is because we use to accept names such as
     // `foo` without any `.` separator. We shouldn't tolerate this in textual
     // format anymore but for now we'll be backward compatible. This can only
     // happen with unregistered dialects.
     if (opName->name.empty()) {
       if (opName->dialect->getLoadedDialect())
LogicalResult BytecodeReader::Impl::parseResourceSection(
   }
 
   // If the resource sections are absent, there is nothing to do.
   if (!resourceData)
     return success();
 
   // Initialize the resource reader with the resource sections.
   DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
-                              reader);
+                              dialectsMap, reader);
   return resourceReader.initialize(fileLoc, config, dialects, stringReader,
                                    *resourceData, *resourceOffsetData,
                                    dialectReader, bufferOwnerRef);
 }
 
 //===----------------------------------------------------------------------===//
 // UseListOrder Helpers
 
BytecodeReader::Impl::parseIRSection(ArrayRef<uint8_t> sectionData,
   }
 
   // Sort use-lists according to what specified in bytecode.
   if (failed(processUseLists(*moduleOp)))
     return reader.emitError(
         "parsed use-list orders were invalid and could not be applied");
 
   // Resolve dialect version.
-  for (const BytecodeDialect &byteCodeDialect : dialects) {
+  for (const std::unique_ptr<BytecodeDialect> &byteCodeDialect : dialects) {
     // Parsing is complete, give an opportunity to each dialect to visit the
     // IR and perform upgrades.
-    if (!byteCodeDialect.loadedVersion)
+    if (!byteCodeDialect->loadedVersion)
       continue;
-    if (byteCodeDialect.interface &&
-        failed(byteCodeDialect.interface->upgradeFromVersion(
-            *moduleOp, *byteCodeDialect.loadedVersion)))
+    if (byteCodeDialect->interface &&
+        failed(byteCodeDialect->interface->upgradeFromVersion(
+            *moduleOp, *byteCodeDialect->loadedVersion)))
       return failure();
   }
 
   // Verify that the parsed operations are valid.
   if (config.shouldVerifyAfterParse() && failed(verify(*moduleOp)))
     return failure();
 
   // Splice the parsed operations over to the provided top-level block.
if (!wasRegistered)
                        "Unexpected missing `wasRegistered` opname flag at "
                        "bytecode version ")
              << version << " with properties.";
     // When an operation is emitted without being registered, the properties are
     // stored as an attribute. Otherwise the op must implement the bytecode
     // interface and control the serialization.
     if (wasRegistered) {
       DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
-                                  reader);
+                                  dialectsMap, reader);
       if (failed(
               propertiesReader.read(fileLoc, dialectReader, &*opName, opState)))
         return failure();
     } else {
       // If the operation wasn't registered when it was emitted, the properties
       // was serialized as an attribute.
       if (failed(parseAttribute(reader, opState.propertiesAttr)))
         return failure();

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

 #include "mlir/Bytecode/Encoding.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/CachedHashString.h"
 #include "llvm/ADT/MapVector.h"
-#include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/raw_ostream.h"
 #include "llvm/Support/Endian.h"
-#include <cstddef>
-#include <cstdint>
-#include <cstring>
+#include "llvm/Support/raw_ostream.h"
 #include <optional>
-#include <sys/types.h>
 
 #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;
 
   /// The producer of the bytecode.
   StringRef producer;
 
+  /// Printer callbacks used to emit custom type and attribute encodings.
+  llvm::SmallVector<std::unique_ptr<AsmAttrTypeBytecodeWriter<Attribute>>>
+      attributeWriterCallbacks;
+  llvm::SmallVector<std::unique_ptr<AsmAttrTypeBytecodeWriter<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() = default;
 
+llvm::SmallVector<std::unique_ptr<AsmAttrTypeBytecodeWriter<Attribute>>> &
+BytecodeWriterConfig::getAttributeWriterCallbacks() const {
+  return impl->attributeWriterCallbacks;
+}
+
+llvm::SmallVector<std::unique_ptr<AsmAttrTypeBytecodeWriter<Type>>> &
+BytecodeWriterConfig::getTypeWriterCallbacks() const {
+  return impl->typeWriterCallbacks;
+}
+
+void BytecodeWriterConfig::attachAttributeCallback(
+    std::unique_ptr<AsmAttrTypeBytecodeWriter<Attribute>> callback) {
+  impl->attributeWriterCallbacks.emplace_back(std::move(callback));
+}
+
+void BytecodeWriterConfig::attachTypeCallback(
+    std::unique_ptr<AsmAttrTypeBytecodeWriter<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::setDesiredBytecodeVersion(int64_t bytecodeVersion) {
   impl->bytecodeVersion = bytecodeVersion;
 }
void BytecodeWriter::writeAttrTypeSection(EncodingEmitter &emitter) {
   offsetEmitter.emitVarInt(llvm::size(numberingState.getAttributes()));
   offsetEmitter.emitVarInt(llvm::size(numberingState.getTypes()));
 
   // A functor used to emit an attribute or type entry.
   uint64_t prevOffset = 0;
   auto emitAttrOrType = [&](auto &entry) {
     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) {
-      // The writer used when emitting using a custom bytecode encoding.
+    auto emitAttrOrTypeRawImpl = [&]() -> void {
+      RawEmitterOstream(attrTypeEmitter) << entryValue;
+      attrTypeEmitter.emitByte(0);
+    };
+    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);
-
       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));
+        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 {
-        // TODO: We don't currently support custom encoded mutable attributes.
-        hasCustomEncoding =
-            !entryValue.template hasTrait<AttributeTrait::IsMutable>() &&
-            succeeded(interface->writeAttribute(entryValue, dialectWriter));
+        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 the dialect interface, emit it using
-    // the textual format.
-    if (!hasCustomEncoding) {
-      RawEmitterOstream(attrTypeEmitter) << entryValue;
-      attrTypeEmitter.emitByte(0);
       }
 
+      // 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);
     prevOffset = curOffset;
   };
 
   // Emit the attribute and type entries for each dialect.
   writeDialectGrouping(offsetEmitter, numberingState.getAttributes(),

mlir/lib/Bytecode/Writer/IRNumbering.cpp

void IRNumberingState::number(Attribute attr) {
   if (OpaqueAttr opaqueAttr = dyn_cast<OpaqueAttr>(attr)) {
     numbering->dialect = &numberDialect(opaqueAttr.getDialectNamespace());
     return;
   }
   numbering->dialect = &numberDialect(&attr.getDialect());
 
   // If this attribute will be emitted using the bytecode format, perform a
   // dummy writing to number any nested components.
-  if (const auto *interface = numbering->dialect->interface) {
-    // TODO: We don't allow custom encodings for mutable attributes right now.
-    if (!attr.hasTrait<AttributeTrait::IsMutable>()) {
+    // Try overriding emission with callbacks.
+    for (const auto &callback : config.getAttributeWriterCallbacks()) {
+      NumberingDialectWriter writer(*this);
+      // The client has the ability to override the group name through the
+      // callback.
+      std::optional<StringRef> groupNameOverride;
+      if (succeeded(callback->write(attr, groupNameOverride, writer))) {
+        if (groupNameOverride.has_value())
+          numbering->dialect = &numberDialect(*groupNameOverride);
+        return;
+      }
+    }
+
+    if (const auto *interface = numbering->dialect->interface) {
       NumberingDialectWriter writer(*this);
       if (succeeded(interface->writeAttribute(attr, writer)))
         return;
     }
   }
   // If this attribute will be emitted using the fallback, number the nested
   // dialect resources. We don't number everything (e.g. no nested
   // attributes/types), because we don't want to encode things we won't decode
void IRNumberingState::number(Type type) {
   if (OpaqueType opaqueType = dyn_cast<OpaqueType>(type)) {
     numbering->dialect = &numberDialect(opaqueType.getDialectNamespace());
     return;
   }
   numbering->dialect = &numberDialect(&type.getDialect());
 
   // If this type will be emitted using the bytecode format, perform a dummy
   // writing to number any nested components.
-  if (const auto *interface = numbering->dialect->interface) {
-    // TODO: We don't allow custom encodings for mutable types right now.
-    if (!type.hasTrait<TypeTrait::IsMutable>()) {
+    // Try overriding emission with callbacks.
+    for (const auto &callback : config.getTypeWriterCallbacks()) {
+      NumberingDialectWriter writer(*this);
+      // The client has the ability to override the group name through the
+      // callback.
+      std::optional<StringRef> groupNameOverride;
+      if (succeeded(callback->write(type, groupNameOverride, writer))) {
+        if (groupNameOverride.has_value())
+          numbering->dialect = &numberDialect(*groupNameOverride);
+        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) {
       NumberingDialectWriter writer(*this);
       if (succeeded(interface->writeType(type, writer)))
         return;
     }
   }
   // If this type will be emitted using the fallback, number the nested dialect
   // resources. We don't number everything (e.g. no nested attributes/types),
   // because we don't want to encode things we won't decode (the textual format

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/Builders.h"
+#include "mlir/IR/BuiltinDialectBytecode.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeRange.h"
 
 using namespace mlir;
+using namespace builtin;
 
 //===----------------------------------------------------------------------===//
 // TableGen'erated dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/BuiltinDialect.cpp.inc"
 
 //===----------------------------------------------------------------------===//
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);
+  addBytecodeInterface(this);
 }
 
 //===----------------------------------------------------------------------===//
 // ModuleOp
 //===----------------------------------------------------------------------===//
 
 void ModuleOp::build(OpBuilder &builder, OperationState &state,
                      std::optional<StringRef> name) {

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/IR/BuiltinDialectBytecode.h"
 #include "AttributeDetail.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectResourceBlobManager.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 
-//===----------------------------------------------------------------------===//
-// BuiltinDialectBytecodeInterface
-//===----------------------------------------------------------------------===//
-
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Utility functions
 
 // TODO: Move these to separate file.
 
 // Returns the bitwidth if known, else return 0.
if (isSplat)
     return writer.writeOwnedString(attr.getRawStringData().front());
 
   for (StringRef str : attr.getRawStringData())
     writer.writeOwnedString(str);
 }
 
 #include "mlir/IR/BuiltinDialectBytecode.cpp.inc"
 
+//===----------------------------------------------------------------------===//
+// BuiltinDialectBytecodeInterface
+//===----------------------------------------------------------------------===//
+
 /// 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 {
     return ::readAttribute(getContext(), reader);
   }
 
   LogicalResult writeAttribute(Attribute attr,
                                DialectBytecodeWriter &writer) const override {
     return ::writeAttribute(attr, writer);
   }
 
   //===--------------------------------------------------------------------===//
   // Types
 
   Type readType(DialectBytecodeReader &reader) const override {
     return ::readType(getContext(), reader);
   }
 
   LogicalResult writeType(Type type,
                           DialectBytecodeWriter &writer) const override {
     return ::writeType(type, writer);
   }
 };
 } // namespace
 
-void builtin_dialect_detail::addBytecodeInterface(BuiltinDialect *dialect) {
+Attribute builtin::readAttribute(DialectBytecodeReader &reader) {
+  return ::readAttribute(reader.getContext(), reader);
+}
+
+LogicalResult builtin::writeAttribute(Attribute attribute,
+                                      DialectBytecodeWriter &writer) {
+  return ::writeAttribute(attribute, writer);
+}
+
+Type builtin::readType(DialectBytecodeReader &reader) {
+  return ::readType(reader.getContext(), reader);
+}
+
+LogicalResult builtin::writeType(Type type, DialectBytecodeWriter &writer) {
+  return ::writeType(type, writer);
+}
+
+void builtin::addBytecodeInterface(BuiltinDialect *dialect) {
   dialect->addInterfaces<BuiltinDialectBytecodeInterface>();
 }

mlir/test/Bytecode/bytecode_callback.mlir

+// RUN: mlir-opt %s --test-bytecode-callback="test-dialect-version=1.2" -verify-diagnostics | FileCheck %s --check-prefix=VERSION_1_2
+// RUN: mlir-opt %s --test-bytecode-callback="test-dialect-version=2.0" -verify-diagnostics | FileCheck %s --check-prefix=VERSION_2_0
+
+func.func @base_test(%arg0 : i32) -> f32 {
+  %0 = "test.addi"(%arg0, %arg0) : (i32, i32) -> i32
+  %1 = "test.cast"(%0) : (i32) -> f32
+  return %1 : f32
+}
+
+// VERSION_1_2: Overriding IntegerType encoding...
+// VERSION_1_2: Overriding parsing of IntegerType encoding...
+
+// VERSION_2_0-NOT: Overriding IntegerType encoding...
+// VERSION_2_0-NOT: Overriding parsing of IntegerType encoding...

mlir/test/Bytecode/bytecode_callback_full_override.mlir

+// RUN: not mlir-opt %s -split-input-file --test-bytecode-callback="callback-test=5" 2>&1 | FileCheck %s
+
+// CHECK-NOT: failed to read bytecode
+func.func @base_test(%arg0 : i32) -> f32 {
+  %0 = "test.addi"(%arg0, %arg0) : (i32, i32) -> i32
+  %1 = "test.cast"(%0) : (i32) -> f32
+  return %1 : f32
+}
+
+// -----
+
+// CHECK-LABEL: error: unknown attribute code: 99
+// CHECK: failed to read bytecode
+func.func @base_test(%arg0 : !test.i32) -> f32 {
+  %0 = "test.addi"(%arg0, %arg0) : (!test.i32, !test.i32) -> !test.i32
+  %1 = "test.cast"(%0) : (!test.i32) -> f32
+  return %1 : f32
+}

mlir/test/Bytecode/bytecode_callback_with_custom_attribute.mlir

+// RUN: mlir-opt %s -split-input-file --test-bytecode-callback="callback-test=3" | FileCheck %s --check-prefix=TEST_3
+// RUN: mlir-opt %s -split-input-file --test-bytecode-callback="callback-test=4" | FileCheck %s --check-prefix=TEST_4
+
+"test.versionedC"() <{attribute = #test.attr_params<42, 24>}> : () -> ()
+
+// TEST_3: Overriding TestAttrParamsAttr encoding...
+// TEST_3: "test.versionedC"() <{attribute = dense<[42, 24]> : tensor<2xi32>}> : () -> ()
+
+// -----
+
+"test.versionedC"() <{attribute = dense<[42, 24]> : tensor<2xi32>}> : () -> ()
+
+// TEST_4: Overriding parsing of TestAttrParamsAttr encoding...
+// TEST_4: "test.versionedC"() <{attribute = #test.attr_params<42, 24>}> : () -> ()

mlir/test/Bytecode/bytecode_callback_with_custom_type.mlir

+// RUN: mlir-opt %s -split-input-file --test-bytecode-callback="callback-test=1" | FileCheck %s --check-prefix=TEST_1
+// RUN: mlir-opt %s -split-input-file --test-bytecode-callback="callback-test=2" | FileCheck %s --check-prefix=TEST_2
+
+func.func @base_test(%arg0: !test.i32, %arg1: f32) {
+  return
+}
+
+// TEST_1: Overriding TestI32Type encoding...
+// TEST_1: func.func @base_test([[ARG0:%.+]]: i32, [[ARG1:%.+]]: f32) {
+
+// -----
+
+func.func @base_test(%arg0: i32, %arg1: f32) {
+  return
+}
+
+// TEST_2: Overriding parsing of TestI32Type encoding...
+// TEST_2: func.func @base_test([[ARG0:%.+]]: !test.i32, [[ARG1:%.+]]: f32) {

mlir/test/Bytecode/invalid/invalid_attr_type_section.mlir

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

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

 //===- TestDialect.h - MLIR Dialect for testing -----------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines a fake 'test' dialect that can be used for testing things
 // that do not have a respective counterpart in the main source directories.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TESTDIALECT_H
 #define MLIR_TESTDIALECT_H
 
-#include "TestTypes.h"
 #include "TestAttributes.h"
 #include "TestInterfaces.h"
+#include "TestTypes.h"
+#include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/DLTI/Traits.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"

 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
 #include "TestOpInterfaces.h.inc"
 #include "TestOpsDialect.h.inc"
 
 namespace test {
+
+//===----------------------------------------------------------------------===//
+// TestDialect version utilities
+//===----------------------------------------------------------------------===//
+
+struct TestDialectVersion : public mlir::DialectVersion {
+  TestDialectVersion() = default;
+  TestDialectVersion(uint32_t _major, uint32_t _minor)
+      : major(_major), minor(_minor){};
+  uint32_t major = 2;
+  uint32_t minor = 0;
+};
+
 // Define some classes to exercises the Properties feature.
 
 struct PropertiesWithCustomPrint {
   /// A shared_ptr to a const object is safe: it is equivalent to a value-based
   /// member. Here the label will be deallocated when the last operation
   /// refering to it is destroyed. However there is no pool-allocation: this is
   /// offloaded to the client.
   std::shared_ptr<const std::string> label;

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

 //===- TestDialect.cpp - MLIR Dialect for Testing -------------------------===//
 //
 // 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 "TestDialect.h"
 #include "TestAttributes.h"
 #include "TestInterfaces.h"
 #include "TestTypes.h"
-#include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/ExtensibleDialect.h"

 static ParseResult customParseProperties(OpAsmParser &parser,
                                          PropertiesWithCustomPrint &prop);
 
 void test::registerTestDialect(DialectRegistry &registry) {
   registry.insert<TestDialect>();
 }
 
 //===----------------------------------------------------------------------===//
-// TestDialect version utilities
-//===----------------------------------------------------------------------===//
-
-struct TestDialectVersion : public DialectVersion {
-  uint32_t major = 2;
-  uint32_t minor = 0;
-};
-
-//===----------------------------------------------------------------------===//
 // TestDialect Interfaces
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// Testing the correctness of some traits.
 static_assert(
     llvm::is_detected<OpTrait::has_implicit_terminator_t,
                       SingleBlockImplicitTerminatorOp>::value,
     "has_implicit_terminator_t does not match SingleBlockImplicitTerminatorOp");
 static_assert(OpTrait::hasSingleBlockImplicitTerminator<
                   SingleBlockImplicitTerminatorOp>::value,
               "hasSingleBlockImplicitTerminator does not match "
               "SingleBlockImplicitTerminatorOp");
 
 struct TestResourceBlobManagerInterface
     : public ResourceBlobManagerDialectInterfaceBase<
           TestDialectResourceBlobHandle> {
   using ResourceBlobManagerDialectInterfaceBase<
       TestDialectResourceBlobHandle>::ResourceBlobManagerDialectInterfaceBase;
 };
 
 namespace {
-enum test_encoding { k_attr_params = 0 };
+enum test_encoding { k_attr_params = 0, k_test_i32 = 99 };
 }
 
 // Test support for interacting with the Bytecode reader/writer.
 struct TestBytecodeDialectInterface : public BytecodeDialectInterface {
   using BytecodeDialectInterface::BytecodeDialectInterface;
   TestBytecodeDialectInterface(Dialect *dialect)
       : BytecodeDialectInterface(dialect) {}
 
+  LogicalResult writeType(Type type,
+                          DialectBytecodeWriter &writer) const final {
+    if (auto concreteType = llvm::dyn_cast<TestI32Type>(type)) {
+      writer.writeVarInt(test_encoding::k_test_i32);
+      return success();
+    }
+    return failure();
+  }
+
+  Type readType(DialectBytecodeReader &reader) const final {
+    uint64_t encoding;
+    if (failed(reader.readVarInt(encoding)))
+      return Type();
+    if (encoding == test_encoding::k_test_i32)
+      return TestI32Type::get(getContext());
+    return Type();
+  }
+
   LogicalResult writeAttribute(Attribute attr,
                                DialectBytecodeWriter &writer) const final {
     if (auto concreteAttr = llvm::dyn_cast<TestAttrParamsAttr>(attr)) {
       writer.writeVarInt(test_encoding::k_attr_params);
       writer.writeVarInt(concreteAttr.getV0());
       writer.writeVarInt(concreteAttr.getV1());
       return success();
     }
     return failure();
   }
 
-  Attribute readAttribute(DialectBytecodeReader &reader,
-                          const DialectVersion &version_) const final {
-    const auto &version = static_cast<const TestDialectVersion &>(version_);
+  Attribute readAttribute(DialectBytecodeReader &reader) const final {
+    auto versionOr = reader.getDialectVersion("test");
+    // Assume current version if not available through the reader.
+    const auto version =
+        (succeeded(versionOr))
+            ? *reinterpret_cast<const TestDialectVersion *>(*versionOr)
+            : TestDialectVersion();
     if (version.major < 2)
       return readAttrOldEncoding(reader);
     if (version.major == 2 && version.minor == 0)
       return readAttrNewEncoding(reader);
     // Forbid reading future versions by returning nullptr.
     return Attribute();
   }
 

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

 
 def TestOpWithVariadicResultsAndFolder: TEST_Op<"op_with_variadic_results_and_folder"> {
   let arguments = (ins Variadic<I32>);
   let results = (outs Variadic<I32>);
   let hasFolder = 1;
 }
 
 def TestAddIOp : TEST_Op<"addi"> {
-  let arguments = (ins I32:$op1, I32:$op2);
-  let results = (outs I32);
+  let arguments = (ins AnyTypeOf<[I32, TestI32]>:$op1,
+                       AnyTypeOf<[I32, TestI32]>:$op2);
+  let results = (outs AnyTypeOf<[I32, TestI32]>);
 }
 
 def TestCommutativeOp : TEST_Op<"op_commutative", [Commutative]> {
   let arguments = (ins I32:$op1, I32:$op2, I32:$op3, I32:$op4);
   let results = (outs I32);
 }
 
 def TestLargeCommutativeOp : TEST_Op<"op_large_commutative", [Commutative]> {
def TestVersionedOpB : TEST_Op<"versionedB"> {
   //
   // We support loading old IR through a custom readAttribute method, see
   // `readAttribute()` in `TestBytecodeDialectInterface`
   let arguments = (ins
     TestAttrParams:$attribute
   );
 }
 
+def TestVersionedOpC : TEST_Op<"versionedC"> {
+  let arguments = (ins AnyAttrOf<[TestAttrParams,
+                                  I32ElementsAttr]>:$attribute
+  );
+}
+
 //===----------------------------------------------------------------------===//
 // Test Properties
 //===----------------------------------------------------------------------===//
 
 
 // Op with a properties struct defined inline.
 def TestOpWithProperties : TEST_Op<"with_properties"> {
   let assemblyFormat = "prop-dict attr-dict";

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

 
 def TestTypeElseAnchorStruct : Test_Type<"TestTypeElseAnchorStruct"> {
   let parameters = (ins OptionalParameter<"std::optional<int>">:$a,
                         OptionalParameter<"std::optional<int>">:$b);
   let mnemonic = "else_anchor_struct";
   let assemblyFormat = "`<` (`?`) : (struct($a, $b)^)? `>`";
 }
 
+def TestI32 : Test_Type<"TestI32"> {
+  let mnemonic = "i32";
+}
+
 #endif // TEST_TYPEDEFS

mlir/test/lib/IR/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
+  TestBytecodeCallbacks.cpp
   TestBuiltinAttributeInterfaces.cpp
   TestBuiltinDistinctAttributes.cpp
   TestClone.cpp
   TestDiagnostics.cpp
   TestDominance.cpp
   TestFunc.cpp
   TestInterfaces.cpp
   TestMatchers.cpp

mlir/test/lib/IR/TestBytecodeCallbacks.cpp

+//===- TestBytecodeCallbacks.cpp - Pass to test bytecode callback hooks  --===//
+//
+// 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 "TestDialect.h"
+#include "mlir/Bytecode/BytecodeReader.h"
+#include "mlir/Bytecode/BytecodeWriter.h"
+#include "mlir/IR/BuiltinDialectBytecode.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/MemoryBufferRef.h"
+#include "llvm/Support/raw_ostream.h"
+#include <list>
+
+using namespace mlir;
+using namespace llvm;
+
+namespace {
+class TestDialectVersionParser : public cl::parser<test::TestDialectVersion> {
+public:
+  TestDialectVersionParser(cl::Option &O)
+      : cl::parser<test::TestDialectVersion>(O) {}
+
+  bool parse(cl::Option &O, StringRef /*argName*/, StringRef arg,
+             test::TestDialectVersion &v) {
+    long long major, minor;
+    if (getAsSignedInteger(arg.split(".").first, 10, major))
+      return O.error("Invalid argument '" + arg);
+    if (getAsSignedInteger(arg.split(".").second, 10, minor))
+      return O.error("Invalid argument '" + arg);
+    v = test::TestDialectVersion(major, minor);
+    // Returns true on error.
+    return false;
+  }
+  static void print(raw_ostream &os, const test::TestDialectVersion &v) {
+    os << v.major << "." << v.minor;
+  };
+};
+
+/// This is a test pass which uses callbacks to encode attributes and types in a
+/// custom fashion.
+struct TestBytecodeCallbackPass
+    : public PassWrapper<TestBytecodeCallbackPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestBytecodeCallbackPass)
+
+  StringRef getArgument() const final { return "test-bytecode-callback"; }
+  StringRef getDescription() const final {
+    return "Test encoding of a dialect type/attributes with a custom callback";
+  }
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<test::TestDialect>();
+  }
+  TestBytecodeCallbackPass() = default;
+  TestBytecodeCallbackPass(const TestBytecodeCallbackPass &) {}
+
+  void runOnOperation() override {
+    switch (testKind) {
+    case (0):
+      return runTest0(getOperation());
+    case (1):
+      return runTest1(getOperation());
+    case (2):
+      return runTest2(getOperation());
+    case (3):
+      return runTest3(getOperation());
+    case (4):
+      return runTest4(getOperation());
+    case (5):
+      return runTest5(getOperation());
+    default:
+      llvm_unreachable("unhandled test kind for TestBytecodeCallbacks pass");
+    }
+  }
+
+  mlir::Pass::Option<test::TestDialectVersion, TestDialectVersionParser>
+      targetVersion{*this, "test-dialect-version",
+                    llvm::cl::desc(
+                        "Specifies the test dialect version to emit and parse"),
+                    cl::init(test::TestDialectVersion())};
+
+  mlir::Pass::Option<int> testKind{
+      *this, "callback-test",
+      llvm::cl::desc("Specifies the test kind to execute"), cl::init(0)};
+
+private:
+  void doRoundtripWithConfigs(Operation *op,
+                              const BytecodeWriterConfig &writeConfig,
+                              const ParserConfig &parseConfig) {
+    std::string bytecode;
+    llvm::raw_string_ostream os(bytecode);
+    if (failed(writeBytecodeToFile(op, os, writeConfig))) {
+      op->emitError() << "failed to write bytecode\n";
+      signalPassFailure();
+      return;
+    }
+    auto newModuleOp = parseSourceString(StringRef(bytecode), parseConfig);
+    if (!newModuleOp.get()) {
+      op->emitError() << "failed to read bytecode\n";
+      signalPassFailure();
+      return;
+    }
+    // Print the module to the output stream, so that we can filecheck the
+    // result.
+    newModuleOp->print(llvm::outs());
+    return;
+  }
+
+  // Test0: let's assume that versions older than 2.0 were relying on a special
+  // integer attribute of the builtin dialect that is now deprecated. Assume
+  // that its encoding was made by two varInts, the first was the ID (999) and
+  // the second contained width and signedness info. We can emit it using a
+  // callback emitting a custom encoding, and parse it back with a custom parser
+  // reading the same encoding. Note that the ID 999 does not correspond to a
+  // valid integer type in the current encodings of builtin types.
+  void runTest0(Operation *op) {
+    auto newCtx = std::make_shared<MLIRContext>();
+    test::TestDialectVersion targetEmissionVersion = targetVersion;
+    BytecodeWriterConfig writeConfig;
+    writeConfig.attachTypeCallback(
+        [&](Type entryValue, std::optional<StringRef> &name,
+            DialectBytecodeWriter &writer) -> LogicalResult {
+          // Do not override anything if version less than 2.0.
+          if (targetEmissionVersion.major >= 2)
+            return failure();
+
+          // For version less than 2.0, override the encoding of IntegerType.
+          if (auto type = llvm::dyn_cast<IntegerType>(entryValue)) {
+            llvm::outs() << "Overriding IntegerType encoding...\n";
+            name = StringLiteral("funky");
+            writer.writeVarInt(/* IntegerType */ 999);
+            writer.writeVarInt(type.getWidth() << 2 | type.getSignedness());
+            return success();
+          }
+          return failure();
+        });
+    newCtx->appendDialectRegistry(op->getContext()->getDialectRegistry());
+    newCtx->allowUnregisteredDialects();
+    ParserConfig parseConfig(newCtx.get(), /*verifyAfterParse=*/true);
+    parseConfig.attachTypeBytecodeCallback([&](DialectBytecodeReader &reader,
+                                               StringRef dialectName,
+                                               Type &entry) -> LogicalResult {
+      // Get test dialect version from the version map.
+      auto versionOr = reader.getDialectVersion("test");
+      assert(
+          succeeded(versionOr) &&
+          "expected reader to be able to access the version for test dialect");
+      const auto *version =
+          reinterpret_cast<const test::TestDialectVersion *>(*versionOr);
+
+      // TODO: once back-deployment is formally supported,
+      // `targetEmissionVersion` will be encoded in the bytecode file, and
+      // exposed through the versionMap. Right now though this is not yet
+      // supported. For the purpose of the test, just use
+      // `targetEmissionVersion`.
+      (void)version;
+      if (targetEmissionVersion.major >= 2)
+        return success();
+
+      if (dialectName != StringLiteral("funky"))
+        return success();
+
+      uint64_t encoding;
+      if (failed(reader.readVarInt(encoding)) || encoding != 999)
+        return success();
+      llvm::outs() << "Overriding parsing of IntegerType encoding...\n";
+      uint64_t _widthAndSignedness, width;
+      IntegerType::SignednessSemantics signedness;
+      if (succeeded(reader.readVarInt(_widthAndSignedness)) &&
+          ((width = _widthAndSignedness >> 2), true) &&
+          ((signedness = static_cast<IntegerType::SignednessSemantics>(
+                _widthAndSignedness & 0x3)),
+           true))
+        entry = IntegerType::get(reader.getContext(), width, signedness);
+      // Return nullopt to fall through the rest of the parsing code path.
+      return success();
+    });
+    doRoundtripWithConfigs(op, writeConfig, parseConfig);
+    return;
+  }
+
+  // Test1: When writing bytecode, we override the encoding of TestI32Type with
+  // the encoding of builtin IntegerType. At parsing, we use such encoding to
+  // read the type and assemble the builtin IntegerType.
+  void runTest1(Operation *op) {
+    BytecodeWriterConfig writeConfig;
+    writeConfig.attachTypeCallback(
+        [&](Type entryValue, std::optional<StringRef> &name,
+            DialectBytecodeWriter &writer) -> LogicalResult {
+          // Emit TestIntegerType using the builtin dialect encoding.
+          if (llvm::isa<test::TestI32Type>(entryValue)) {
+            llvm::outs() << "Overriding TestI32Type encoding...\n";
+            auto builtinI32Type =
+                IntegerType::get(op->getContext(), 32,
+                                 IntegerType::SignednessSemantics::Signless);
+            name = StringLiteral("builtin");
+            if (succeeded(builtin::writeType(builtinI32Type, writer)))
+              return success();
+          }
+          return failure();
+        });
+    // We natively parse the attribute as a builtin, so no callback needed.
+    ParserConfig parseConfig(op->getContext(), /*verifyAfterParse=*/true);
+    doRoundtripWithConfigs(op, writeConfig, parseConfig);
+    return;
+  }
+
+  // Test2: When writing bytecode, we write standard builtin IntegerTypes. At
+  // parsing, we use the encoding of IntegerType to intercept all i32. Then,
+  // instead of creating i32s, we assemble TestI32Type and return it.
+  void runTest2(Operation *op) {
+    BytecodeWriterConfig writeConfig;
+    ParserConfig parseConfig(op->getContext(), /*verifyAfterParse=*/true);
+    parseConfig.attachTypeBytecodeCallback([&](DialectBytecodeReader &reader,
+                                               StringRef dialectName,
+                                               Type &entry) -> LogicalResult {
+      if (dialectName != StringLiteral("builtin"))
+        return success();
+      Type builtinAttr = builtin::readType(reader);
+      if (auto integerType = llvm::dyn_cast_or_null<IntegerType>(builtinAttr)) {
+        if (integerType.getWidth() == 32 && integerType.isSignless()) {
+          llvm::outs() << "Overriding parsing of TestI32Type encoding...\n";
+          entry = test::TestI32Type::get(reader.getContext());
+        }
+      }
+      return success();
+    });
+    doRoundtripWithConfigs(op, writeConfig, parseConfig);
+    return;
+  }
+
+  // Test3: When writing bytecode, we override the encoding of
+  // TestAttrParamsAttr with the encoding of builtin DenseIntElementsAttr. At
+  // parsing, we use such encoding to read the type and assemble the builtin
+  // DenseIntElementsAttr.
+  void runTest3(Operation *op) {
+    auto i32Type = IntegerType::get(op->getContext(), 32,
+                                    IntegerType::SignednessSemantics::Signless);
+    BytecodeWriterConfig writeConfig;
+    writeConfig.attachAttributeCallback(
+        [&](Attribute entryValue, std::optional<StringRef> &name,
+            DialectBytecodeWriter &writer) -> LogicalResult {
+          // Emit TestIntegerType using the builtin dialect encoding.
+          if (auto testParamAttrs =
+                  llvm::dyn_cast<test::TestAttrParamsAttr>(entryValue)) {
+            llvm::outs() << "Overriding TestAttrParamsAttr encoding...\n";
+            name = StringLiteral("builtin");
+            auto denseAttr = DenseIntElementsAttr::get(
+                RankedTensorType::get({2}, i32Type),
+                {testParamAttrs.getV0(), testParamAttrs.getV1()});
+            if (succeeded(builtin::writeAttribute(denseAttr, writer)))
+              return success();
+          }
+          return failure();
+        });
+    // We natively parse the attribute as a builtin, so no callback needed.
+    ParserConfig parseConfig(op->getContext(), /*verifyAfterParse=*/false);
+    doRoundtripWithConfigs(op, writeConfig, parseConfig);
+    return;
+  }
+
+  // Test4: When writing bytecode, we write standard builtin
+  // DenseIntElementsAttr. At parsing, we use the encoding of
+  // DenseIntElementsAttr to intercept all ElementsAttr that have shaped type of
+  // <2xi32>. Instead of assembling a DenseIntElementsAttr, we assemble
+  // TestAttrParamsAttr and return it.
+  void runTest4(Operation *op) {
+    auto i32Type = IntegerType::get(op->getContext(), 32,
+                                    IntegerType::SignednessSemantics::Signless);
+    BytecodeWriterConfig writeConfig;
+    ParserConfig parseConfig(op->getContext(), /*verifyAfterParse=*/false);
+    parseConfig.attachAttributeBytecodeCallback(
+        [&](DialectBytecodeReader &reader, StringRef dialectName,
+            Attribute &entry) -> LogicalResult {
+          // Override only the case where the return type of the builtin reader
+          // is an i32 and fall through on all the other cases, since we want to
+          // still use TestDialect normal codepath to parse the other types.
+          Attribute builtinAttr = builtin::readAttribute(reader);
+          if (auto denseAttr =
+                  llvm::dyn_cast_or_null<DenseIntElementsAttr>(builtinAttr)) {
+            if (denseAttr.getType().getShape() == ArrayRef<int64_t>(2) &&
+                denseAttr.getElementType() == i32Type) {
+              llvm::outs()
+                  << "Overriding parsing of TestAttrParamsAttr encoding...\n";
+              int v0 = denseAttr.getValues<IntegerAttr>()[0].getInt();
+              int v1 = denseAttr.getValues<IntegerAttr>()[1].getInt();
+              entry =
+                  test::TestAttrParamsAttr::get(reader.getContext(), v0, v1);
+            }
+          }
+          return success();
+        });
+    doRoundtripWithConfigs(op, writeConfig, parseConfig);
+    return;
+  }
+
+  // Test5: When writing bytecode, we want TestDialect to use nothing else than
+  // the builtin types and attributes and take full control of the encoding,
+  // returning failure if any type or attribute is not part of builtin.
+  void runTest5(Operation *op) {
+    BytecodeWriterConfig writeConfig;
+    writeConfig.attachAttributeCallback(
+        [&](Attribute attr, std::optional<StringRef> &name,
+            DialectBytecodeWriter &writer) -> LogicalResult {
+          return builtin::writeAttribute(attr, writer);
+        });
+    writeConfig.attachTypeCallback(
+        [&](Type type, std::optional<StringRef> &name,
+            DialectBytecodeWriter &writer) -> LogicalResult {
+          return builtin::writeType(type, writer);
+        });
+    ParserConfig parseConfig(op->getContext(), /*verifyAfterParse=*/false);
+    parseConfig.attachAttributeBytecodeCallback(
+        [&](DialectBytecodeReader &reader, StringRef dialectName,
+            Attribute &entry) -> LogicalResult {
+          Attribute builtinAttr = builtin::readAttribute(reader);
+          if (!builtinAttr)
+            return failure();
+          entry = builtinAttr;
+          return success();
+        });
+    parseConfig.attachTypeBytecodeCallback([&](DialectBytecodeReader &reader,
+                                               StringRef dialectName,
+                                               Type &entry) -> LogicalResult {
+      Type builtinType = builtin::readType(reader);
+      if (!builtinType) {
+        return failure();
+      }
+      entry = builtinType;
+      return success();
+    });
+    doRoundtripWithConfigs(op, writeConfig, parseConfig);
+    return;
+  }
+};
+} // namespace
+
+namespace mlir {
+void registerTestBytecodeCallbackPasses() {
+  PassRegistration<TestBytecodeCallbackPass>();
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerLoopLikeInterfaceTestPasses();
 void registerShapeFunctionTestPasses();
 void registerSideEffectTestPasses();
 void registerSliceAnalysisTestPass();
 void registerSymbolTestPasses();
 void registerRegionTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAffineReifyValueBoundsPass();
+void registerTestBytecodeCallbackPasses();
 void registerTestDecomposeAffineOpPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLoopPermutationPass();
 void registerTestMatchers();
 void registerTestOperationEqualPass();
void registerTestPasses() {
   registerSliceAnalysisTestPass();
   registerSymbolTestPasses();
   registerRegionTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAffineReifyValueBoundsPass();
   registerTestDecomposeAffineOpPass();
   registerTestAffineLoopUnswitchingPass();
   registerTestAllReduceLoweringPass();
+  registerTestBytecodeCallbackPasses();
   registerTestFunc();
   registerTestGpuMemoryPromotionPass();
   registerTestLoopPermutationPass();
   registerTestMatchers();
   registerTestOperationEqualPass();
   registerTestPrintDefUsePass();
   registerTestPrintInvalidPass();
   registerTestPrintNestingPass();