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

90

@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

51

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

65

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.

90

LG

[mfrancio]

Adds bytecode roundtrip tests with custom integer types.

[mfrancio]

mlir/test/lib/IR/TestBytecodeCallbacks.cpp

51

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.

65

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

51

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.

65

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

51

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

65

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

1234

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

1234

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

28

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)

580

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

mlir/include/mlir/IR/BuiltinDialectBytecode.h

22 ↗

(On Diff #540749)

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

24 ↗

(On Diff #540749)

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 ↗

(On Diff #540749)

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

579

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 ↗

(On Diff #540749)

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

87 ↗

(On Diff #540749)

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

580

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 ↗

(On Diff #540749)

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 ↗

(On Diff #540749)

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 ↗

(On Diff #540749)

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 ↗

(On Diff #540749)

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 ↗

(On Diff #540749)

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 ↗

(On Diff #540749)

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

815

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 ↗

(On Diff #542263)

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 ↗

(On Diff #542263)

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 ↗

(On Diff #542263)

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 ↗

(On Diff #542263)

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

578

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

578

done! Thanks for the review.

[mfrancio]

Rebase

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<AsmAttrTypeBytecodePrinter>> &
+  getAttrTypePrinterCallbacks() const;
+
+  /// Attach a custom bytecode printer callback to the configuration for the
+  /// emission of custom type/attributes encodings.
+  void
+  attachAttrTypeCallback(std::unique_ptr<AsmAttrTypeBytecodePrinter> 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(
+                                           Type, DialectBytecodeWriter &)>> &&
+      std::is_convertible_v<
+          CallableT,
+          std::function<LogicalResult(Attribute, DialectBytecodeWriter &)>>>
+  attachAttrTypeCallback(CallableT &&emitFn) {
+    attachAttrTypeCallback(AsmAttrTypeBytecodePrinter::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 "llvm/ADT/StringMap.h"
 
 #include <memory>
 #include <variant>
 
 namespace mlir {
 class AsmResourcePrinter;
 class AsmDialectResourceHandle;
+class DialectBytecodeWriter;
+class DialectBytecodeReader;
+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.
+class AsmAttrTypeBytecodePrinter {
+public:
+  AsmAttrTypeBytecodePrinter() = default;
+  virtual ~AsmAttrTypeBytecodePrinter() = default;
+
+  virtual LogicalResult write(Type entry, DialectBytecodeWriter &writer) = 0;
+  virtual LogicalResult write(Attribute entry,
+                              DialectBytecodeWriter &writer) = 0;
+
+  /// Return an Attribute/Type printer implemented via the given callable, whose
+  /// form should match that of `write` functions above.
+  template <typename CallableT,
+            std::enable_if_t<
+                std::is_convertible_v<CallableT,
+                                      std::function<LogicalResult(
+                                          Type, DialectBytecodeWriter &)>> &&
+                    std::is_convertible_v<
+                        CallableT, std::function<LogicalResult(
+                                       Attribute, DialectBytecodeWriter &)>>,
+                bool> = true>
+  static std::unique_ptr<AsmAttrTypeBytecodePrinter>
+  fromCallable(CallableT &&writeFn) {
+    struct Processor : public AsmAttrTypeBytecodePrinter {
+      Processor(CallableT &&writeFn)
+          : AsmAttrTypeBytecodePrinter(), _writeFn(std::move(writeFn)) {}
+      LogicalResult write(Type entry, DialectBytecodeWriter &writer) override {
+        return _writeFn(entry, writer);
+      }
+      LogicalResult write(Attribute entry,
+                          DialectBytecodeWriter &writer) override {
+        return _writeFn(entry, 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 emitting MLIR
+/// bytecode.
+class AsmAttrTypeBytecodeParser {
+public:
+  AsmAttrTypeBytecodeParser() = default;
+  virtual ~AsmAttrTypeBytecodeParser() = default;
+
+  virtual void parse(MLIRContext *ctx, DialectBytecodeReader &reader,
+                     const llvm::StringMap<DialectVersion *> &versionMap,
+                     Type &entry) = 0;
+  virtual void parse(MLIRContext *ctx, DialectBytecodeReader &reader,
+                     const llvm::StringMap<DialectVersion *> &versionMap,
+                     Attribute &entry) = 0;
+
+  /// Return an Attribute/Type printer implemented via the given callable, whose
+  /// form should match that of `parse` functions above.
+  template <
+      typename CallableT,
+      std::enable_if_t<
+          std::is_convertible_v<
+              CallableT,
+              std::function<void(MLIRContext *, DialectBytecodeReader &,
+                                 const llvm::StringMap<DialectVersion *> &,
+                                 Type &)>> &&
+              std::is_convertible_v<
+                  CallableT,
+                  std::function<void(MLIRContext *, DialectBytecodeReader &,
+                                     const llvm::StringMap<DialectVersion *> &,
+                                     Attribute &)>>,
+          bool> = true>
+  static std::unique_ptr<AsmAttrTypeBytecodeParser>
+  fromCallable(CallableT &&parseFn) {
+    struct Processor : public AsmAttrTypeBytecodeParser {
+      Processor(CallableT &&parseFn)
+          : AsmAttrTypeBytecodeParser(), _parseFn(std::move(parseFn)) {}
+      void parse(MLIRContext *ctx, DialectBytecodeReader &reader,
+                 const llvm::StringMap<DialectVersion *> &versionMap,
+                 Type &entry) override {
+        return _parseFn(ctx, reader, versionMap, entry);
+      }
+      void parse(MLIRContext *ctx, DialectBytecodeReader &reader,
+                 const llvm::StringMap<DialectVersion *> &versionMap,
+                 Attribute &entry) override {
+        return _parseFn(ctx, reader, versionMap, 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.
+  auto &getAttrTypeBytecodeCallbacks() const { return attrTypeBytecodeParsers; }
+
+  /// Attach a custom bytecode parser callback to the configuration for parsing
+  /// of custom type/attributes encodings.
+  void attachAttrTypeBytecodeCallback(
+      std::unique_ptr<AsmAttrTypeBytecodeParser> parser) {
+    attrTypeBytecodeParsers.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<void(
+                         MLIRContext *, DialectBytecodeReader &,
+                         const llvm::StringMap<DialectVersion *> &, Type &)>> &&
+      std::is_convertible_v<
+          CallableT,
+          std::function<void(MLIRContext *, DialectBytecodeReader &,
+                             const llvm::StringMap<DialectVersion *> &,
+                             Attribute &)>>>
+  attachAttrTypeBytecodeCallback(CallableT &&parserFn) {
+    attachAttrTypeBytecodeCallback(AsmAttrTypeBytecodeParser::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>>
+      attrTypeBytecodeParsers;
 };
 
 //===----------------------------------------------------------------------===//
 // 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/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.
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, Location fileLoc,
+                 const ParserConfig &config)
       : stringReader(stringReader), resourceReader(resourceReader),
-        fileLoc(fileLoc) {}
+        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);
 
   /// Resolve the attribute or type at the given index. Returns nullptr on
   /// failure.
private:
   ResourceSectionReader &resourceReader;
 
   /// The set of attribute and type entries.
   SmallVector<AttrEntry> attributes;
   SmallVector<TypeEntry> types;
 
   /// A location used for error emission.
   Location fileLoc;
+
+  /// A map to retrieve parsed dialect versions associated to each dialect name.
+  llvm::StringMap<DialectVersion *> dialectsVersionMap;
+
+  /// Reference to the parser configuration.
+  const ParserConfig &parserConfig;
 };
 
 class DialectReader : public DialectBytecodeReader {
 public:
   DialectReader(AttrTypeReader &attrTypeReader,
                 StringSectionReader &stringReader,
                 ResourceSectionReader &resourceReader, EncodingReader &reader)
       : attrTypeReader(attrTypeReader), stringReader(stringReader),
AttrTypeReader::initialize(MutableArrayRef<BytecodeDialect> dialects,
   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");
   }
+
+  // Fill up the dialect to dialectVersion map for every dialect version
+  // available.
+  for (auto &dialect : dialects) {
+    EncodingReader encReader(dialect.versionBuffer, fileLoc);
+    DialectReader dialectReader(*this, stringReader, resourceReader, encReader);
+    if (failed(dialect.load(dialectReader, fileLoc.getContext())))
+      return failure();
+    if (dialect.loadedVersion.get())
+      dialectsVersionMap.insert({dialect.name, dialect.loadedVersion.get()});
+  }
+
   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;

 
 template <typename T>
 LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
                                                EncodingReader &reader,
                                                StringRef entryType) {
   DialectReader dialectReader(*this, stringReader, resourceReader, reader);
   if (failed(entry.dialect->load(dialectReader, fileLoc.getContext())))
     return failure();
+
+  // Try parsing with callbacks first if available.
+  for (const auto &callback : parserConfig.getAttrTypeBytecodeCallbacks()) {
+    callback->parse(fileLoc.getContext(), dialectReader,
+                    dialectsVersionMap, entry.entry);
+
+    // 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.
class mlir::BytecodeReader::Impl {
   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, 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;
+  llvm::StringMap<DialectVersion *> dialectMap;
   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;

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 <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<AsmAttrTypeBytecodePrinter>>
+      attrTypePrinterCallbacks;
+
   /// 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<AsmAttrTypeBytecodePrinter>> &
+BytecodeWriterConfig::getAttrTypePrinterCallbacks() const {
+  return impl->attrTypePrinterCallbacks;
+}
+
+void BytecodeWriterConfig::attachAttrTypeCallback(
+    std::unique_ptr<AsmAttrTypeBytecodePrinter> callback) {
+  impl->attrTypePrinterCallbacks.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) {
+    // TODO: We don't currently support custom encoded mutable types and
+    // attributes.
+    if (!entryValue.template hasTrait<TypeTrait::IsMutable>() &&
+        !entryValue.template hasTrait<AttributeTrait::IsMutable>()) {
       // The writer used when emitting using a custom bytecode encoding.
       DialectWriter dialectWriter(config.bytecodeVersion, attrTypeEmitter,
                                   numberingState, stringSection);
+      for (const auto &callback : config.attrTypePrinterCallbacks) {
+        if (succeeded(callback->write(entryValue, dialectWriter)))
+          hasCustomEncoding = true;
+      }
 
-      if constexpr (std::is_same_v<std::decay_t<decltype(entryValue)>, Type>) {
-        // TODO: We don't currently support custom encoded mutable types.
+      if (!hasCustomEncoding) {
+        if (const BytecodeDialectInterface *interface =
+                entry.dialect->interface) {
+          // The writer used when emitting using a custom bytecode encoding.
+          DialectWriter dialectWriter(config.bytecodeVersion, attrTypeEmitter,
+                                      numberingState, stringSection);
+          if constexpr (std::is_same_v<std::decay_t<decltype(entryValue)>,
+                                       Type>) {
             hasCustomEncoding =
-            !entryValue.template hasTrait<TypeTrait::IsMutable>() &&
                 succeeded(interface->writeType(entryValue, dialectWriter));
           } else {
-        // TODO: We don't currently support custom encoded mutable attributes.
             hasCustomEncoding =
-            !entryValue.template hasTrait<AttributeTrait::IsMutable>() &&
                 succeeded(interface->writeAttribute(entryValue, dialectWriter));
           }
         }
+      }
+    }
 
     // If the entry was not emitted using the dialect interface, emit it using
     // the textual format.
     if (!hasCustomEncoding) {
       RawEmitterOstream(attrTypeEmitter) << entryValue;
       attrTypeEmitter.emitByte(0);
     }
 
     // Record the offset of this entry.
     uint64_t curOffset = attrTypeEmitter.size();
     offsetEmitter.emitVarIntWithFlag(curOffset - prevOffset, hasCustomEncoding);
     prevOffset = curOffset;
   };
 
   // Emit the attribute and type entries for each dialect.

mlir/lib/Bytecode/Writer/IRNumbering.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.getAttrTypePrinterCallbacks()) {
+      NumberingDialectWriter writer(*this);
+      if (succeeded(callback->write(attr, writer)))
+        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.getAttrTypePrinterCallbacks()) {
+      NumberingDialectWriter writer(*this);
+      if (succeeded(callback->write(type, writer)))
+        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/test/Bytecode/bytecode_callback.mlir

+// RUN: mlir-opt %s --test-bytecode-callback | FileCheck %s
+
+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: func.func @base_test([[ARG0:%.+]]: i32) -> f32 {
+// CHECK:   [[VAR0:%.+]] = "test.addi"([[ARG0]], [[ARG0]]) : (i32, i32) -> i32
+// CHECK:   [[VAR1:%.+]] = "test.cast"([[VAR0]]) : (i32) -> f32
+// CHECK:   return [[VAR1]] : 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,

mlir/test/lib/IR/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
+  TestBytecodeCallbacks.cpp
   TestBuiltinAttributeInterfaces.cpp
   TestClone.cpp
   TestDiagnostics.cpp
   TestDominance.cpp
   TestFunc.cpp
   TestInterfaces.cpp
   TestMatchers.cpp
   TestLazyLoading.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/BuiltinOps.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/MemoryBufferRef.h"
+#include "llvm/Support/raw_ostream.h"
+#include <list>
+
+using namespace mlir;
+
+namespace {
+/// This is a test pass which uses callbacks to encode attributes and types in a
+/// custom fashion.
+struct TestBytecodeCallbackPass
+    : public PassWrapper<TestBytecodeCallbackPass, OperationPass<>> {
+  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";
+  }
+  TestBytecodeCallbackPass() = default;
+  TestBytecodeCallbackPass(const TestBytecodeCallbackPass &) {}
+
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    test::TestDialectVersion targetEmissionVersion(1, 2);
+    std::string bytecode;
+    {
+      // For testing purposes, 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 custom encoding.
+
+      // Note that the ID 999 does not correspond to a valid integer type in the
+      // builtin encoding.
+      BytecodeWriterConfig writeConfig;
+      writeConfig.attachAttrTypeCallback(
+          [&](auto entryValue, DialectBytecodeWriter &writer) -> LogicalResult {
+            // Do not override anything if version less than 2.0.
+            if (targetEmissionVersion.major >= 2)
+              return failure();
+            // We don't override any encoding, hence return failure.
+            if constexpr (std::is_same_v<std::decay_t<decltype(entryValue)>,
+                                         Type>) {
+              if (auto type = llvm::dyn_cast<IntegerType>(entryValue)) {
+                writer.writeVarInt(/* IntegerType */ 999);
+                writer.writeVarInt(type.getWidth() << 2 | type.getSignedness());
+                return success();
+              }
+            }
+            return failure();
+          });
+      llvm::raw_string_ostream os(bytecode);
+      if (failed(writeBytecodeToFile(op, os, writeConfig))) {
+        op->emitError() << "failed to write bytecode\n";
+        signalPassFailure();
+        return;
+      }
+    }
+    ParserConfig parseConfig(op->getContext(), /*verifyAfterParse=*/true);
+    parseConfig.attachAttrTypeBytecodeCallback(
+        [&](MLIRContext *ctx, DialectBytecodeReader &reader,
+            const llvm::StringMap<DialectVersion *> &versionMap,
+            auto &entry) -> void {
+          // Get test dialect version from the version map.
+          assert(
+              versionMap.contains("test") &&
+              "expected versionMap to contain all the available version info");
+          test::TestDialectVersion &version =
+              static_cast<test::TestDialectVersion &>(*versionMap.at("test"));
+
+          // 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;
+
+          if constexpr (std::is_same_v<std::decay_t<decltype(entry)>, Type>) {
+            uint64_t encoding;
+            if (failed(reader.readVarInt(encoding)) || encoding != 999)
+              return;
+            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(ctx, width, signedness);
+              return;
+            }
+            // Fall through and do not assign entry to fallback on the
+            // standard codepath for parsing types and attributes.
+          }
+          return;
+        });
+    auto newModuleOp = parseSourceString(StringRef(bytecode), parseConfig);
+    if (!newModuleOp.get()) {
+      op->emitError() << "failed to read bytecode\n";
+      signalPassFailure();
+    }
+    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();