Expose callbacks for encoding of types/attributes

[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.

Reviewed By: rriddle

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

Author: Matteo Franciolini

mlir/include/mlir/Bytecode/BytecodeImplementation.h

@@ -24,6 +24,17 @@
 #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
 //===----------------------------------------------------------------------===//
@@ -38,7 +49,14 @@ class DialectBytecodeReader {
   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;
 
   /// Return the bytecode version being read.
   virtual uint64_t getBytecodeVersion() const = 0;
@@ -384,17 +402,6 @@ class DialectBytecodeWriter {
   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
 //===----------------------------------------------------------------------===//
@@ -409,47 +416,23 @@ class BytecodeDialectInterface
   //===--------------------------------------------------------------------===//
 
   /// 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
   //===--------------------------------------------------------------------===//

mlir/include/mlir/Bytecode/BytecodeReader.h

@@ -25,7 +25,6 @@ class SourceMgr;
 } // namespace llvm
 
 namespace mlir {
-
 /// The BytecodeReader allows to load MLIR bytecode files, while keeping the
 /// state explicitly available in order to support lazy loading.
 /// The `finalize` method must be called before destruction.

mlir/include/mlir/Bytecode/BytecodeReaderConfig.h

@@ -0,0 +1,120 @@
+//===- BytecodeReader.h - MLIR Bytecode Reader ------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header defines interfaces to read MLIR bytecode files/streams.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_BYTECODE_BYTECODEREADERCONFIG_H
+#define MLIR_BYTECODE_BYTECODEREADERCONFIG_H
+
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+
+namespace mlir {
+class Attribute;
+class DialectBytecodeReader;
+class Type;
+
+/// A class to interact with the attributes and types parser when parsing MLIR
+/// bytecode.
+template <class T>
+class AttrTypeBytecodeReader {
+public:
+  AttrTypeBytecodeReader() = default;
+  virtual ~AttrTypeBytecodeReader() = default;
+
+  virtual LogicalResult read(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<AttrTypeBytecodeReader<T>>
+  fromCallable(CallableT &&readFn) {
+    struct Processor : public AttrTypeBytecodeReader<T> {
+      Processor(CallableT &&readFn)
+          : AttrTypeBytecodeReader(), readFn(std::move(readFn)) {}
+      LogicalResult read(DialectBytecodeReader &reader, StringRef dialectName,
+                         T &entry) override {
+        return readFn(reader, dialectName, entry);
+      }
+
+      std::decay_t<CallableT> readFn;
+    };
+    return std::make_unique<Processor>(std::forward<CallableT>(readFn));
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// BytecodeReaderConfig
+//===----------------------------------------------------------------------===//
+
+/// A class containing bytecode-specific configurations of the `ParserConfig`.
+class BytecodeReaderConfig {
+public:
+  BytecodeReaderConfig() = default;
+
+  /// Returns the callbacks available to the parser.
+  ArrayRef<std::unique_ptr<AttrTypeBytecodeReader<Attribute>>>
+  getAttributeCallbacks() const {
+    return attributeBytecodeParsers;
+  }
+  ArrayRef<std::unique_ptr<AttrTypeBytecodeReader<Type>>>
+  getTypeCallbacks() const {
+    return typeBytecodeParsers;
+  }
+
+  /// Attach a custom bytecode parser callback to the configuration for parsing
+  /// of custom type/attributes encodings.
+  void attachAttributeCallback(
+      std::unique_ptr<AttrTypeBytecodeReader<Attribute>> parser) {
+    attributeBytecodeParsers.emplace_back(std::move(parser));
+  }
+  void
+  attachTypeCallback(std::unique_ptr<AttrTypeBytecodeReader<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 &)>>>
+  attachAttributeCallback(CallableT &&parserFn) {
+    attachAttributeCallback(AttrTypeBytecodeReader<Attribute>::fromCallable(
+        std::forward<CallableT>(parserFn)));
+  }
+  template <typename CallableT>
+  std::enable_if_t<std::is_convertible_v<
+      CallableT,
+      std::function<LogicalResult(DialectBytecodeReader &, StringRef, Type &)>>>
+  attachTypeCallback(CallableT &&parserFn) {
+    attachTypeCallback(AttrTypeBytecodeReader<Type>::fromCallable(
+        std::forward<CallableT>(parserFn)));
+  }
+
+private:
+  llvm::SmallVector<std::unique_ptr<AttrTypeBytecodeReader<Attribute>>>
+      attributeBytecodeParsers;
+  llvm::SmallVector<std::unique_ptr<AttrTypeBytecodeReader<Type>>>
+      typeBytecodeParsers;
+};
+
+} // namespace mlir
+
+#endif // MLIR_BYTECODE_BYTECODEREADERCONFIG_H

mlir/include/mlir/Bytecode/BytecodeWriter.h

@@ -17,6 +17,55 @@
 
 namespace mlir {
 class Operation;
+class DialectBytecodeWriter;
+
+/// A class to interact with the attributes and types printer when emitting MLIR
+/// bytecode.
+template <class T>
+class AttrTypeBytecodeWriter {
+public:
+  AttrTypeBytecodeWriter() = default;
+  virtual ~AttrTypeBytecodeWriter() = default;
+
+  /// Callback writer API used in IRNumbering, where groups are created and
+  /// type/attribute components are numbered. At this stage, writer is expected
+  /// to be a `NumberingDialectWriter`.
+  virtual LogicalResult write(T entry, std::optional<StringRef> &name,
+                              DialectBytecodeWriter &writer) = 0;
+
+  /// Callback writer API used in BytecodeWriter, where groups are created and
+  /// type/attribute components are numbered. Here, DialectBytecodeWriter is
+  /// expected to be an actual writer. The optional stringref specified by
+  /// the user is ignored, since the group was already specified when numbering
+  /// the IR.
+  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<AttrTypeBytecodeWriter<T>>
+  fromCallable(CallableT &&writeFn) {
+    struct Processor : public AttrTypeBytecodeWriter<T> {
+      Processor(CallableT &&writeFn)
+          : AttrTypeBytecodeWriter(), 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));
+  }
+};
 
 /// This class contains the configuration used for the bytecode writer. It
 /// controls various aspects of bytecode generation, and contains all of the
@@ -48,6 +97,43 @@ class BytecodeWriterConfig {
   /// Get the set desired bytecode version to emit.
   int64_t getDesiredBytecodeVersion() const;
 
+  //===--------------------------------------------------------------------===//
+  // Types and Attributes encoding
+  //===--------------------------------------------------------------------===//
+
+  /// Retrieve the callbacks.
+  ArrayRef<std::unique_ptr<AttrTypeBytecodeWriter<Attribute>>>
+  getAttributeWriterCallbacks() const;
+  ArrayRef<std::unique_ptr<AttrTypeBytecodeWriter<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<AttrTypeBytecodeWriter<Attribute>> callback);
+  void
+  attachTypeCallback(std::unique_ptr<AttrTypeBytecodeWriter<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(AttrTypeBytecodeWriter<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(AttrTypeBytecodeWriter<Type>::fromCallable(
+        std::forward<CallableT>(emitFn)));
+  }
+
   //===--------------------------------------------------------------------===//
   // Resources
   //===--------------------------------------------------------------------===//

mlir/include/mlir/IR/AsmState.h

@@ -14,6 +14,7 @@
 #ifndef MLIR_IR_ASMSTATE_H_
 #define MLIR_IR_ASMSTATE_H_
 
+#include "mlir/Bytecode/BytecodeReaderConfig.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/MapVector.h"
@@ -475,6 +476,11 @@ class ParserConfig {
   /// Returns if the parser should verify the IR after parsing.
   bool shouldVerifyAfterParse() const { return verifyAfterParse; }
 
+  /// Returns the parsing configurations associated to the bytecode read.
+  BytecodeReaderConfig &getBytecodeReaderConfig() const {
+    return const_cast<BytecodeReaderConfig &>(bytecodeReaderConfig);
+  }
+
   /// Return the resource parser registered to the given name, or nullptr if no
   /// parser with `name` is registered.
   AsmResourceParser *getResourceParser(StringRef name) const {
@@ -509,6 +515,7 @@ class ParserConfig {
   bool verifyAfterParse;
   DenseMap<StringRef, std::unique_ptr<AsmResourceParser>> resourceParsers;
   FallbackAsmResourceMap *fallbackResourceMap;
+  BytecodeReaderConfig bytecodeReaderConfig;
 };
 
 //===----------------------------------------------------------------------===//