Add support for versioning properties in MLIR bytecode

[mfrancio]

Summary

[mlir] Add support for custom readProperties/writeProperties methods.

Currently, operations that opt-in to adopt properties will see auto-generated readProperties/writeProperties methods to emit and parse bytecode. If a dialects opts in to use usePropertiesForAttributes, those definitions will be generated for the current definition of the op without the possibility to handle attribute versioning.

The patch adds the capability for an operation to define its own read/write methods for the encoding of properties so that versioned operations can handle upgrading properties encodings.

In addition to this, the patch adds an example showing versioning on NamedProperties through the dialect version API exposed by the reader.

[mehdi_amini]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1072

I am not fond of having this weird transactional way of changing the state of the reader temporarily to carry this information.

Here is a patch

map.patch7 KBDownload

that add a map that would allow clients to get a version for any dialect by name: can we use this kind of mechanism instead?

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1347

I think we should split the logic for adding bytecode in another method and keep the control flow simple.

Can you outline the code from line 1337 to the end?

(I think this method has outgrown what's reasonable anyway)

[mfrancio]

• Refactors the way the dialect version is exposed to the reader
• Outlines readProperties/writeProperties methods generation into a new method of OpEmitter

[mfrancio]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1072

Makes sense to do this - i leveraged it to refactor the dialect reader, which now holds a reference of the map since construction.

BTW, does it make sense anymore now to have readAttribute/readType taking a version?
https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bytecode/BytecodeImplementation.h#L301

I feel we should consolidate the two methods.

I will also use the same strategy for https://reviews.llvm.org/D153383, which simplifies things a bit.

[jpienaar]

BTW, does it make sense anymore now to have readAttribute/readType taking a version?

https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bytecode/BytecodeImplementation.h#L301

Not if already accessible on reader no.

I'm not sure what consolidating would mean though, I think we need a separate read for attribute and type given they only differ in return (c++) type

[mfrancio]

In D155340#4508225, @jpienaar wrote:

BTW, does it make sense anymore now to have readAttribute/readType taking a version?

https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bytecode/BytecodeImplementation.h#L301

Not if already accessible on reader no.

I'm not sure what consolidating would mean though, I think we need a separate read for attribute and type given they only differ in return (c++) type

Exactly, once the reader have access to the version there is no need of having a specific API with a version passed as an argument.

[mfrancio]

Rebase

[mfrancio]

Edit commit message

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;
 
+  /// Retrieve the dialect version for the dialect that is in processing. It is
+  /// responsibility of the caller to make sure that the version is actually
+  /// available.
+  std::shared_ptr<DialectVersion> getDialectVersion() const { return _version; }
+
+  /// Load a dialect version on the reader.
+  void loadVersion(std::shared_ptr<DialectVersion> version) {
+    _version = version;
+  };
+
   /// 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;
     if (failed(readVarInt(size)))
public:
   virtual LogicalResult readBlob(ArrayRef<char> &result) = 0;
 
   /// Read a bool from the bytecode.
   virtual LogicalResult readBool(bool &result) = 0;
 
 private:
   /// Read a handle to a dialect resource.
   virtual FailureOr<AsmDialectResourceHandle> readResourceHandle() = 0;
+
+  /// Reference to the dialect version we are currently reading.
+  std::shared_ptr<DialectVersion> _version;
 };
 
 //===----------------------------------------------------------------------===//
 // DialectBytecodeWriter
 //===----------------------------------------------------------------------===//
 
 /// This class defines a virtual interface for writing to a bytecode stream,
 /// providing hooks into the bytecode writer. As such, this class should only be
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;
virtual LogicalResult writeType(Type type,
                                   DialectBytecodeWriter &writer) const {
     return failure();
   }
 
   /// Write the version of this dialect to the given writer.
   virtual void writeVersion(DialectBytecodeWriter &writer) const {}
 
   // Read the version of this dialect from the provided reader and return it as
-  // a `unique_ptr` to a dialect version object.
-  virtual std::unique_ptr<DialectVersion>
+  // a `shared_ptr` to a dialect version object.
+  virtual std::shared_ptr<DialectVersion>
   readVersion(DialectBytecodeReader &reader) const {
     reader.emitError("Dialect does not support versioning");
     return nullptr;
   }
 
   /// Hook invoked after parsing completed, if a version directive was present
   /// and included an entry for the current dialect. This hook offers the
   /// opportunity to the dialect to visit the IR and upgrades constructs emitted

mlir/include/mlir/IR/OpBase.td

class Op<Dialect dialect, string mnemonic, list<Trait> props = []> {
 
   // Whether this op has a static "canonicalize" method to perform "match and
   // rewrite patterns".
   bit hasCanonicalizeMethod = 0;
 
   // Whether this op has a folder.
   bit hasFolder = 0;
 
+  // Whether to let ops implement their custom `readProperties` and
+  // `writeProperties` methods to emit bytecode.
+  bit useCustomPropertiesEncoding = 0;
+
   // Op traits.
   // Note: The list of traits will be uniqued by ODS.
   list<Trait> traits = props;
 
   // Additional code that will be added to the public part of the generated
   // C++ code of the op declaration.
   code extraClassDeclaration = ?;
 

mlir/include/mlir/TableGen/Operator.h

public:
   /// Returns the setter name for the accessor of `name`.
   std::string getSetterName(StringRef name) const;
 
   /// Returns the remove name for the accessor of `name`.
   std::string getRemoverName(StringRef name) const;
 
   bool hasFolder() const;
 
+  /// Whether to generate the `readProperty`/`writeProperty` methods for
+  /// bytecode emission.
+  bool useCustomPropertiesEncoding() const;
+
 private:
   /// Populates the vectors containing operands, attributes, results and traits.
   void populateOpStructure();
 
   /// Populates type inference info (mostly equality) with input a mapping from
   /// names to indices for arguments and results.
   void populateTypeInferenceInfo(
       const llvm::StringMap<int> &argumentsAndResultsIndex);

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

struct BytecodeDialect {
 
   /// The name of the dialect.
   StringRef name;
 
   /// A buffer containing the encoding of the dialect version parsed.
   ArrayRef<uint8_t> versionBuffer;
 
   /// Lazy loaded dialect version from the handle above.
-  std::unique_ptr<DialectVersion> loadedVersion;
+  std::shared_ptr<DialectVersion> loadedVersion;
 };
 
 /// This struct represents an operation name entry within the bytecode.
 struct BytecodeOperationName {
   BytecodeOperationName(BytecodeDialect *dialect, StringRef name,
                         std::optional<bool> wasRegistered)
       : dialect(dialect), name(name), wasRegistered(wasRegistered) {}
 
LogicalResult initialize(Location fileLoc, ArrayRef<uint8_t> sectionData) {
     }
     if (!offsetsReader.empty())
       return offsetsReader.emitError()
              << "Broken properties section: didn't exhaust the offsets table";
     return success();
   }
 
   LogicalResult read(Location fileLoc, DialectReader &dialectReader,
-                     OperationName *opName, OperationState &opState) {
+                     OperationName *opName, OperationState &opState,
+                     const BytecodeDialect &dialect) {
     uint64_t propertiesIdx;
     if (failed(dialectReader.readVarInt(propertiesIdx)))
       return failure();
     if (propertiesIdx >= offsetTable.size())
       return dialectReader.emitError("Properties idx out-of-bound for ")
              << opName->getStringRef();
     size_t propertiesOffset = offsetTable[propertiesIdx];
     if (propertiesIdx >= propertiesBuffers.size())
ArrayRef<char> rawProperties;
               dialectReader.withEncodingReader(reader).readBlob(rawProperties)))
         return failure();
     }
     // Setup a new reader to read from the `rawProperties` sub-buffer.
     EncodingReader reader(
         StringRef(rawProperties.begin(), rawProperties.size()), fileLoc);
     DialectReader propReader = dialectReader.withEncodingReader(reader);
 
-    auto *iface = opName->getInterface<BytecodeOpInterface>();
-    if (iface)
+    // Load the version into the reader.
+    propReader.loadVersion(dialect.loadedVersion);
+
+    if (auto *iface = opName->getInterface<BytecodeOpInterface>())
       return iface->readProperties(propReader, opState);
     if (opName->isRegistered())
       return propReader.emitError(
                  "has properties but missing BytecodeOpInterface for ")
              << opName->getStringRef();
     // Unregistered op are storing properties as an attribute.
     return propReader.readAttribute(opState.propertiesAttr);
   }
private:
   //===--------------------------------------------------------------------===//
   // Dialect Section
 
   LogicalResult parseDialectSection(ArrayRef<uint8_t> sectionData);
 
   /// Parse an operation name reference using the given reader, and set the
   /// `wasRegistered` flag that indicates if the bytecode was produced by a
   /// context where opName was registered.
-  FailureOr<OperationName> parseOpName(EncodingReader &reader,
-                                       std::optional<bool> &wasRegistered);
+  FailureOr<BytecodeOperationName *> parseOpName(EncodingReader &reader);
 
   //===--------------------------------------------------------------------===//
   // Attribute/Type Section
 
   /// Parse an attribute or type using the given reader.
   template <typename T>
   LogicalResult parseAttribute(EncodingReader &reader, T &result) {
     return attrTypeReader.parseAttribute(reader, result);
if (version >= bytecode::kElideUnknownBlockArgLocation) {
     opNames.reserve(numOps);
   }
   while (!sectionReader.empty())
     if (failed(parseDialectGrouping(sectionReader, dialects, parseOpName)))
       return failure();
   return success();
 }
 
-FailureOr<OperationName>
-BytecodeReader::Impl::parseOpName(EncodingReader &reader,
-                                  std::optional<bool> &wasRegistered) {
+FailureOr<BytecodeOperationName *>
+BytecodeReader::Impl::parseOpName(EncodingReader &reader) {
   BytecodeOperationName *opName = nullptr;
   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);
     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())
         return emitError(fileLoc) << "has an empty opname for dialect '"
                                   << opName->dialect->name << "'\n";
 
       opName->opName.emplace(opName->dialect->name, getContext());
     } else {
       opName->opName.emplace((opName->dialect->name + "." + opName->name).str(),
                              getContext());
     }
   }
-  return *opName->opName;
+  return opName;
 }
 
 //===----------------------------------------------------------------------===//
 // Resource Section
 
 LogicalResult BytecodeReader::Impl::parseResourceSection(
     EncodingReader &reader, std::optional<ArrayRef<uint8_t>> resourceData,
     std::optional<ArrayRef<uint8_t>> resourceOffsetData) {
BytecodeReader::Impl::parseRegions(std::vector<RegionReadState> &regionStack,
   return success();
 }
 
 FailureOr<Operation *>
 BytecodeReader::Impl::parseOpWithoutRegions(EncodingReader &reader,
                                             RegionReadState &readState,
                                             bool &isIsolatedFromAbove) {
   // Parse the name of the operation.
-  std::optional<bool> wasRegistered;
-  FailureOr<OperationName> opName = parseOpName(reader, wasRegistered);
-  if (failed(opName))
+  FailureOr<BytecodeOperationName *> bytecodeOpName = parseOpName(reader);
+  if (failed(bytecodeOpName))
     return failure();
 
+  OperationName opName = *(*bytecodeOpName)->opName;
+  std::optional<bool> wasRegistered = (*bytecodeOpName)->wasRegistered;
+  BytecodeDialect &bytecodeDialect = *(*bytecodeOpName)->dialect;
+
   // Parse the operation mask, which indicates which components of the operation
   // are present.
   uint8_t opMask;
   if (failed(reader.parseByte(opMask)))
     return failure();
 
   /// Parse the location.
   LocationAttr opLoc;
   if (failed(parseAttribute(reader, opLoc)))
     return failure();
 
   // With the location and name resolved, we can start building the operation
   // state.
-  OperationState opState(opLoc, *opName);
+  OperationState opState(opLoc, opName);
 
   // Parse the attributes of the operation.
   if (opMask & bytecode::OpEncodingMask::kHasAttrs) {
     DictionaryAttr dictAttr;
     if (failed(parseAttribute(reader, dictAttr)))
       return failure();
     opState.attributes = dictAttr;
   }
 
   if (opMask & bytecode::OpEncodingMask::kHasProperties) {
     // kHasProperties wasn't emitted in older bytecode, we should never get
     // there without also having the `wasRegistered` flag available.
-    if (!wasRegistered)
+    if (!wasRegistered.has_value())
       return emitError(fileLoc,
                        "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) {
+    if (*wasRegistered) {
       DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
                                   reader);
-      if (failed(
-              propertiesReader.read(fileLoc, dialectReader, &*opName, opState)))
+      if (failed(propertiesReader.read(fileLoc, dialectReader, &opName, opState,
+                                       bytecodeDialect)))
         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/TableGen/Operator.cpp

std::string Operator::getSetterName(StringRef name) const {
   return "set" + convertToCamelFromSnakeCase(name, /*capitalizeFirst=*/true);
 }
 
 std::string Operator::getRemoverName(StringRef name) const {
   return "remove" + convertToCamelFromSnakeCase(name, /*capitalizeFirst=*/true);
 }
 
 bool Operator::hasFolder() const { return def.getValueAsBit("hasFolder"); }
+
+bool Operator::useCustomPropertiesEncoding() const {
+  return def.getValueAsBit("useCustomPropertiesEncoding");
+}

mlir/test/Bytecode/versioning/versioned-op-with-prop-2.0.mlirbc

+module {
+  "test.versionedA"() <{dims = 123 : i64, modifier = false}> : () -> ()
+}
+

mlir/test/Bytecode/versioning/versioned_bytecode.mlir

 // This file contains test cases related to roundtripping.
 
 //===--------------------------------------------------------------------===//
 // Test roundtrip
 //===--------------------------------------------------------------------===//
 
-// RUN: mlir-opt %S/versioned-op-1.12.mlirbc -emit-bytecode \
+// RUN: mlir-opt %S/versioned-op-with-prop-1.12.mlirbc -emit-bytecode \
 // RUN:   -emit-bytecode-version=0 | mlir-opt -o %t.1 && \
-// RUN: mlir-opt %S/versioned-op-1.12.mlirbc -o %t.2 && \
+// RUN: mlir-opt %S/versioned-op-with-prop-1.12.mlirbc -o %t.2 && \
 // RUN: diff %t.1 %t.2
 
 //===--------------------------------------------------------------------===//
 // Test invalid versions
 //===--------------------------------------------------------------------===//
 
-// RUN: not mlir-opt %S/versioned-op-1.12.mlirbc -emit-bytecode \
+// RUN: not mlir-opt %S/versioned-op-with-prop-1.12.mlirbc -emit-bytecode \
 // RUN:   -emit-bytecode-version=-1 2>&1 | FileCheck %s --check-prefix=ERR_VERSION_NEGATIVE
 // ERR_VERSION_NEGATIVE: unsupported version requested -1, must be in range [{{[0-9]+}}, {{[0-9]+}}]
 
-// RUN: not mlir-opt %S/versioned-op-1.12.mlirbc -emit-bytecode \
+// RUN: not mlir-opt %S/versioned-op-with-prop-1.12.mlirbc -emit-bytecode \
 // RUN:   -emit-bytecode-version=999 2>&1 | FileCheck %s --check-prefix=ERR_VERSION_FUTURE
 // ERR_VERSION_FUTURE: unsupported version requested 999, must be in range [{{[0-9]+}}, {{[0-9]+}}]

mlir/test/Bytecode/versioning/versioned_op.mlir

 // This file contains test cases related to the dialect post-parsing upgrade
 // mechanism.
+// COM: those tests parse bytecode that was generated before test dialect
+//      adopted `usePropertiesFromAttributes`.
 
 //===--------------------------------------------------------------------===//
 // Test generic
 //===--------------------------------------------------------------------===//
 
 // COM: bytecode contains
 // COM: module {
 // COM:   version: 2.0
 // COM:   "test.versionedA"() <{dims = 123 : i64, modifier = false}> : () -> ()
 // COM: }
-// RUN: mlir-opt %S/versioned-op-2.0.mlirbc 2>&1 | FileCheck %s --check-prefix=CHECK1
+// RUN: mlir-opt %S/versioned-op-with-prop-2.0.mlirbc 2>&1 | FileCheck %s --check-prefix=CHECK1
 // CHECK1: "test.versionedA"() <{dims = 123 : i64, modifier = false}> : () -> ()
 
 //===--------------------------------------------------------------------===//
 // Test upgrade
 //===--------------------------------------------------------------------===//
 
 // COM: bytecode contains
 // COM: module {
 // COM:   version: 1.12
 // COM:   "test.versionedA"() <{dimensions = 123 : i64}> : () -> ()
 // COM: }
-// RUN: mlir-opt %S/versioned-op-1.12.mlirbc 2>&1 | FileCheck %s --check-prefix=CHECK2
-// CHECK2: "test.versionedA"() <{dims = 123 : i64, modifier = false}> : () -> ()
+// RUN: mlir-opt %S/versioned-op-with-prop-1.12.mlirbc 2>&1 | FileCheck %s --check-prefix=CHECK3
+// CHECK3: "test.versionedA"() <{dims = 123 : i64, modifier = false}> : () -> ()
 
 //===--------------------------------------------------------------------===//
 // Test forbidden downgrade
 //===--------------------------------------------------------------------===//
 
 // COM: bytecode contains
 // COM: module {
 // COM:   version: 2.2
 // COM:   "test.versionedA"() <{dims = 123 : i64, modifier = false}> : () -> ()
 // COM: }
 // RUN: not mlir-opt %S/versioned-op-2.2.mlirbc 2>&1 | FileCheck %s --check-prefix=ERR_NEW_VERSION
 // ERR_NEW_VERSION: current test dialect version is 2.0, can't parse version: 2.2

mlir/test/Bytecode/versioning/versioned_op_with_native_properties.mlir

+// This file contains test cases related to the dialect post-parsing upgrade
+// mechanism.
+// COM: those tests parse bytecode that was generated before test dialect
+//      adopted `usePropertiesFromAttributes`.
+
+//===--------------------------------------------------------------------===//
+// Test generic
+//===--------------------------------------------------------------------===//
+
+// COM: bytecode contains
+// COM: module {
+// COM:   version: 2.0
+// COM:   test.with_versioned_properties 1 | 2
+// COM: }
+// RUN: mlir-opt %S/versioned-op-with-native-prop-2.0.mlirbc 2>&1 | FileCheck %s --check-prefix=CHECK1
+// CHECK1: test.with_versioned_properties 1 | 2
+
+//===--------------------------------------------------------------------===//
+// Test upgrade
+//===--------------------------------------------------------------------===//
+
+// COM: bytecode contains
+// COM: module {
+// COM:   version: 1.12
+
+// COM: }
+// RUN: mlir-opt %S/versioned-op-with-native-prop-1.12.mlirbc 2>&1 | FileCheck %s --check-prefix=CHECK3
+// CHECK3: test.with_versioned_properties 1 | 0

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

public:
   static mlir::LogicalResult setFromAttr(MyPropStruct &prop,
                                          mlir::Attribute attr,
                                          mlir::InFlightDiagnostic *diag);
   llvm::hash_code hash() const;
   bool operator==(const MyPropStruct &rhs) const {
     return content == rhs.content;
   }
 };
+struct VersionedProperties {
+  // For the sake of testing, assume that this object was associated to
+  // version 1.2 of the test dialect when having only one int value. In the
+  // current version 2.0, the property has two values and on upgrade. We also
+  // assume that the class is upgrade-able if value2 = 0.
+  int value1;
+  int value2;
+  bool operator==(const VersionedProperties &rhs) const {
+    return value1 == rhs.value1 && value2 == rhs.value2;
+  }
+};
 } // namespace test
 
 #define GET_OP_CLASSES
 #include "TestOps.h.inc"
 
 namespace test {
 
 // Op deliberately defined in C++ code rather than ODS to test that C++

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

 static DictionaryAttr
 getPropertiesAsAttribute(MLIRContext *ctx,
                          const PropertiesWithCustomPrint &prop);
 static llvm::hash_code computeHash(const PropertiesWithCustomPrint &prop);
 static void customPrintProperties(OpAsmPrinter &p,
                                   const PropertiesWithCustomPrint &prop);
 static ParseResult customParseProperties(OpAsmParser &parser,
                                          PropertiesWithCustomPrint &prop);
+static LogicalResult setPropertiesFromAttribute(VersionedProperties &prop,
+                                                Attribute attr,
+                                                InFlightDiagnostic *diagnostic);
+static DictionaryAttr
+getPropertiesAsAttribute(MLIRContext *ctx,
+                         const VersionedProperties &prop);
+static llvm::hash_code computeHash(const VersionedProperties &prop);
+static void customPrintProperties(OpAsmPrinter &p,
+                                  const VersionedProperties &prop);
+static ParseResult customParseProperties(OpAsmParser &parser,
+                                         VersionedProperties &prop);
 
 void test::registerTestDialect(DialectRegistry &registry) {
   registry.insert<TestDialect>();
 }
 
 //===----------------------------------------------------------------------===//
 // TestDialect version utilities
 //===----------------------------------------------------------------------===//
struct TestBytecodeDialectInterface : public BytecodeDialectInterface {
 
   // Emit a specific version of the dialect.
   void writeVersion(DialectBytecodeWriter &writer) const final {
     auto version = TestDialectVersion();
     writer.writeVarInt(version.major); // major
     writer.writeVarInt(version.minor); // minor
   }
 
-  std::unique_ptr<DialectVersion>
+  std::shared_ptr<DialectVersion>
   readVersion(DialectBytecodeReader &reader) const final {
     uint64_t major, minor;
     if (failed(reader.readVarInt(major)) || failed(reader.readVarInt(minor)))
       return nullptr;
-    auto version = std::make_unique<TestDialectVersion>();
+    auto version = std::make_shared<TestDialectVersion>();
     version->major = major;
     version->minor = minor;
     return version;
   }
 
   LogicalResult upgradeFromVersion(Operation *topLevelOp,
                                    const DialectVersion &version_) const final {
     const auto &version = static_cast<const TestDialectVersion &>(version_);
     if ((version.major == 2) && (version.minor == 0))
       return success();
     if (version.major > 2 || (version.major == 2 && version.minor > 0)) {
       return topLevelOp->emitError()
              << "current test dialect version is 2.0, can't parse version: "
              << version.major << "." << version.minor;
     }
-    // Prior version 2.0, the old op supported only a single attribute called
-    // "dimensions". We can perform the upgrade.
     topLevelOp->walk([](TestVersionedOpA op) {
-      if (auto dims = op->getAttr("dimensions")) {
-        op->removeAttr("dimensions");
-        op->setAttr("dims", dims);
-      }
-      op->setAttr("modifier", BoolAttr::get(op->getContext(), false));
+      // Prior version 2.0, `readProperties` did not process the modifier
+      // attribute. Handle that according to the version here.
+      auto &prop = op.getProperties();
+      prop.modifier = BoolAttr::get(op->getContext(), false);
     });
     return success();
   }
 
 private:
   Attribute readAttrNewEncoding(DialectBytecodeReader &reader) const {
     uint64_t encoding;
     if (failed(reader.readVarInt(encoding)) ||
static ParseResult customParseProperties(OpAsmParser &parser,
                                          PropertiesWithCustomPrint &prop) {
   std::string label;
   if (parser.parseKeywordOrString(&label) || parser.parseKeyword("is") ||
       parser.parseInteger(prop.value))
     return failure();
   prop.label = std::make_shared<std::string>(std::move(label));
   return success();
 }
+static LogicalResult
+setPropertiesFromAttribute(VersionedProperties &prop, Attribute attr,
+                           InFlightDiagnostic *diagnostic) {
+  DictionaryAttr dict = dyn_cast<DictionaryAttr>(attr);
+  if (!dict) {
+    if (diagnostic)
+      *diagnostic << "expected DictionaryAttr to set VersionedProperties";
+    return failure();
+  }
+  auto value1Attr = dict.getAs<IntegerAttr>("value1");
+  if (!value1Attr) {
+    if (diagnostic)
+      *diagnostic << "expected IntegerAttr for key `value1`";
+    return failure();
+  }
+  auto value2Attr = dict.getAs<IntegerAttr>("value2");
+  if (!value2Attr) {
+    if (diagnostic)
+      *diagnostic << "expected IntegerAttr for key `value2`";
+    return failure();
+  }
+
+  prop.value1 = value1Attr.getValue().getSExtValue();
+  prop.value2 = value2Attr.getValue().getSExtValue();
+  return success();
+}
+static DictionaryAttr
+getPropertiesAsAttribute(MLIRContext *ctx,
+                         const VersionedProperties &prop) {
+  SmallVector<NamedAttribute> attrs;
+  Builder b{ctx};
+  attrs.push_back(b.getNamedAttr("value1", b.getI32IntegerAttr(prop.value1)));
+  attrs.push_back(b.getNamedAttr("value2", b.getI32IntegerAttr(prop.value2)));
+  return b.getDictionaryAttr(attrs);
+}
+static llvm::hash_code computeHash(const VersionedProperties &prop) {
+  return llvm::hash_combine(prop.value1, prop.value2);
+}
+static void customPrintProperties(OpAsmPrinter &p,
+                                  const VersionedProperties &prop) {
+  p << prop.value1 << " | " << prop.value2;
+}
+static ParseResult customParseProperties(OpAsmParser &parser,
+                                         VersionedProperties &prop) {
+  if (parser.parseInteger(prop.value1) || parser.parseVerticalBar() ||
+      parser.parseInteger(prop.value2))
+    return failure();
+  return success();
+}
 
 static bool parseUsingPropertyInCustom(OpAsmParser &parser, int64_t value[3]) {
   return parser.parseLSquare() || parser.parseInteger(value[0]) ||
          parser.parseComma() || parser.parseInteger(value[1]) ||
          parser.parseComma() || parser.parseInteger(value[2]) ||
          parser.parseRSquare();
 }
 
 static void printUsingPropertyInCustom(OpAsmPrinter &printer, Operation *op,
                                        ArrayRef<int64_t> value) {
   printer << '[' << value << ']';
 }
 
+LogicalResult
+TestVersionedOpA::readProperties(::mlir::DialectBytecodeReader &reader,
+                                 ::mlir::OperationState &state) {
+  auto &prop = state.getOrAddProperties<Properties>();
+  if (::mlir::failed(reader.readAttribute(prop.dims)))
+    return ::mlir::failure();
+
+  // Check if we have a version. If not, assume we are parsing the current
+  // version.
+  auto maybeVersion = reader.getDialectVersion();
+  if (maybeVersion.get()) {
+    // If version is less than 2.0, there is no additional attribute to parse.
+    // We can materialize missing properties post parsing before verification.
+    const auto &version =
+        static_cast<const TestDialectVersion &>(*maybeVersion);
+    if ((version.major < 2)) {
+      return success();
+    }
+  }
+
+  if (::mlir::failed(reader.readAttribute(prop.modifier)))
+    return ::mlir::failure();
+  return ::mlir::success();
+}
+
+void TestVersionedOpA::writeProperties(::mlir::DialectBytecodeWriter &writer) {
+  auto &prop = getProperties();
+  writer.writeAttribute(prop.dims);
+  writer.writeAttribute(prop.modifier);
+}
+
+::mlir::LogicalResult TestOpWithVersionedProperties::readFromMlirBytecode(
+    ::mlir::DialectBytecodeReader &reader, test::VersionedProperties &prop) {
+  uint64_t value1, value2 = 0;
+  if (failed(reader.readVarInt(value1)))
+    return failure();
+
+  // Check if we have a version. If not, assume we are parsing the current
+  // version.
+  auto maybeVersion = reader.getDialectVersion();
+  if (maybeVersion.get()) {
+    // If version is less than 2.0, there is no additional attribute to parse.
+    // We can materialize missing properties post parsing before verification.
+    const auto &version =
+        static_cast<const TestDialectVersion &>(*maybeVersion);
+    if ((version.major >= 2) && failed(reader.readVarInt(value2)))
+      return failure();
+  }
+
+  prop.value1 = value1;
+  prop.value2 = value2;
+  return success();
+}
+void TestOpWithVersionedProperties::writeToMlirBytecode(
+    ::mlir::DialectBytecodeWriter &writer,
+    const test::VersionedProperties &prop) {
+  writer.writeVarInt(prop.value1);
+  writer.writeVarInt(prop.value2);
+}
+
 #include "TestOpEnums.cpp.inc"
 #include "TestOpInterfaces.cpp.inc"
 #include "TestTypeInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

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

def TestVersionedOpA : TEST_Op<"versionedA"> {
   // In the current version (2.0) "dimensions" was renamed to "dims", and a new
   // boolean attribute "modifier" was added. The previous version of the op
   // corresponds to "modifier=false". We support loading old IR through
   // upgrading, see `upgradeFromVersion()` in `TestBytecodeDialectInterface`.
   let arguments = (ins
    AnyI64Attr:$dims,
    BoolAttr:$modifier
   );
+
+  // Since we use properties to store attributes, we need a custom encoding
+  // reader/writer to handle versioning.
+  let useCustomPropertiesEncoding = 1;
 }
 
 def TestVersionedOpB : TEST_Op<"versionedB"> {
   // A previous version of the dialect (let's say 1.*) we encoded TestAttrParams
   // with a custom encoding:
   //
   //    #test.attr_params<X, Y> -> { varInt: Y, varInt: X }
   //
::mlir::ParseResult TestOpWithNiceProperties::parseProperties(
       Properties &prop = result.getOrAddProperties<Properties>();
       if (customParseProperties(parser, prop.prop))
         return failure();
       return success();
     }
   }];
 }
 
+def VersionedProperties : Property<"VersionedProperties"> {
+  let convertToAttribute = [{
+      getPropertiesAsAttribute($_ctxt, $_storage)
+  }];
+  let convertFromAttribute = [{
+    return setPropertiesFromAttribute($_storage, $_attr, $_diag);
+  }];
+  let hashProperty = [{
+    computeHash($_storage);
+  }];
+}
+
+def TestOpWithVersionedProperties : TEST_Op<"with_versioned_properties"> {
+  let assemblyFormat = "prop-dict attr-dict";
+  let arguments = (ins
+     VersionedProperties:$prop
+  );
+  let extraClassDeclaration = [{
+    void printProperties(::mlir::MLIRContext *ctx, ::mlir::OpAsmPrinter &p,
+                         const Properties &prop);
+    static ::mlir::ParseResult  parseProperties(::mlir::OpAsmParser &parser,
+                                     ::mlir::OperationState &result);
+    static ::mlir::LogicalResult readFromMlirBytecode(
+        ::mlir::DialectBytecodeReader &,
+        test::VersionedProperties &prop);
+    static void writeToMlirBytecode(
+        ::mlir::DialectBytecodeWriter &,
+        const test::VersionedProperties &prop);
+  }];
+  let extraClassDefinition = [{
+    void TestOpWithVersionedProperties::printProperties(::mlir::MLIRContext *ctx,
+            ::mlir::OpAsmPrinter &p, const Properties &prop) {
+      customPrintProperties(p, prop.prop);
+    }
+    ::mlir::ParseResult TestOpWithVersionedProperties::parseProperties(
+        ::mlir::OpAsmParser &parser,
+        ::mlir::OperationState &result) {
+      Properties &prop = result.getOrAddProperties<Properties>();
+      if (customParseProperties(parser, prop.prop))
+        return failure();
+      return success();
+    }
+  }];
+}
+
 
 
 #endif // TEST_OPS

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

auto &verifyInherentAttrsMethod =
               "::mlir::LogicalResult", "verifyInherentAttrs",
               MethodParameter("::mlir::OperationName", "opName"),
               MethodParameter("::mlir::NamedAttrList &", "attrs"),
               MethodParameter(
                   "llvm::function_ref<::mlir::InFlightDiagnostic()>",
                   "getDiag"))
           ->body();
 
-  auto &readPropertiesMethod =
-      opClass
-          .addStaticMethod(
-              "::mlir::LogicalResult", "readProperties",
-              MethodParameter("::mlir::DialectBytecodeReader &", "reader"),
-              MethodParameter("::mlir::OperationState &", "state"))
-          ->body();
-
-  auto &writePropertiesMethod =
-      opClass
-          .addMethod(
-              "void", "writeProperties",
-              MethodParameter("::mlir::DialectBytecodeWriter &", "writer"))
-          ->body();
-
   opClass.declare<UsingDeclaration>("Properties", "FoldAdaptor::Properties");
 
   // Convert the property to the attribute form.
 
   setPropMethod << R"decl(
   ::mlir::DictionaryAttr dict = ::llvm::dyn_cast<::mlir::DictionaryAttr>(attr);
   if (!dict) {
     if (diag)
if (canEmitAttrVerifier(attr,
         return ::mlir::failure();
     }
 )",
                      name, constraintFn, namedAttr->attrName);
     }
   }
   verifyInherentAttrsMethod << "    return ::mlir::success();";
 
+  if (op.useCustomPropertiesEncoding()) {
+    opClass
+        .declareStaticMethod(
+            "::mlir::LogicalResult", "readProperties",
+            MethodParameter("::mlir::DialectBytecodeReader &", "reader"),
+            MethodParameter("::mlir::OperationState &", "state"));
+    opClass
+        .declareMethod(
+            "void", "writeProperties",
+            MethodParameter("::mlir::DialectBytecodeWriter &", "writer"));
+    return;
+  }
+
+  auto &readPropertiesMethod =
+      opClass
+          .addStaticMethod(
+              "::mlir::LogicalResult", "readProperties",
+              MethodParameter("::mlir::DialectBytecodeReader &", "reader"),
+              MethodParameter("::mlir::OperationState &", "state"))
+          ->body();
+
+  auto &writePropertiesMethod =
+      opClass
+          .addMethod(
+              "void", "writeProperties",
+              MethodParameter("::mlir::DialectBytecodeWriter &", "writer"))
+          ->body();
+
   // Populate bytecode serialization logic.
   readPropertiesMethod
       << "  auto &prop = state.getOrAddProperties<Properties>(); (void)prop;";
   writePropertiesMethod << "  auto &prop = getProperties(); (void)prop;\n";
   for (const auto &attrOrProp : attrOrProperties) {
     if (const auto *namedProperty =
             attrOrProp.dyn_cast<const NamedProperty *>()) {
       StringRef name = namedProperty->name;