Add support for versioning properties in MLIR bytecode

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

Reviewed By: mehdi_amini

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

Author: Matteo Franciolini

mlir/include/mlir/IR/OpBase.td

@@ -2576,6 +2576,10 @@ class Op<Dialect dialect, string mnemonic, list<Trait> props = []> {
   // 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;

mlir/include/mlir/TableGen/Operator.h

@@ -352,6 +352,10 @@ class Operator {
 
   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();

mlir/lib/TableGen/Operator.cpp

@@ -854,3 +854,7 @@ std::string Operator::getRemoverName(StringRef name) const {
 }
 
 bool Operator::hasFolder() const { return def.getValueAsBit("hasFolder"); }
+
+bool Operator::useCustomPropertiesEncoding() const {
+  return def.getValueAsBit("useCustomPropertiesEncoding");
+}

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

@@ -4,19 +4,19 @@
 // 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-with-native-prop-2.0.mlirbc

@@ -10,7 +10,7 @@
 // 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}> : () -> ()
 
 //===--------------------------------------------------------------------===//
@@ -22,8 +22,8 @@
 // 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

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

@@ -0,0 +1,28 @@
+// 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/Bytecode/versioning/versioned-op-with-prop-2.0.mlirbc

@@ -114,6 +114,16 @@ 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>();
@@ -1148,6 +1158,54 @@ static ParseResult customParseProperties(OpAsmParser &parser,
   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]) ||
@@ -1220,6 +1278,69 @@ MutableOperandRange TestStoreWithARegionTerminator::getMutableSuccessorOperands(
   return MutableOperandRange(getOperation());
 }
 
+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("test");
+  if (succeeded(maybeVersion)) {
+    // 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 =
+        reinterpret_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("test");
+  bool needToParseAnotherInt = true;
+  if (succeeded(maybeVersion)) {
+    // 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 =
+        reinterpret_cast<const TestDialectVersion *>(*maybeVersion);
+    if ((version->major < 2))
+      needToParseAnotherInt = false;
+  }
+  if (needToParseAnotherInt && 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"

mlir/test/Bytecode/versioning/versioned_bytecode.mlir

@@ -98,6 +98,17 @@ class MyPropStruct {
     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. 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

mlir/test/Bytecode/versioning/versioned_op.mlir

@@ -122,11 +122,10 @@ struct TestBytecodeDialectInterface : public BytecodeDialectInterface {
     // 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();
   }

mlir/test/Bytecode/versioning/versioned_op_with_native_properties.mlir

@@ -785,7 +785,7 @@ def OpWithShapedTypeInferTypeInterfaceOp : TEST_Op<"op_with_shaped_type_infer_ty
   let results = (outs AnyTensor);
 }
 
-def OpWithShapedTypeInferTypeAdaptorInterfaceOp : 
+def OpWithShapedTypeInferTypeAdaptorInterfaceOp :
       TEST_Op<"op_with_shaped_type_infer_type_adaptor_if",
               [InferTensorTypeAdaptorWithReify]> {
   let arguments = (ins AnyTensor:$operand1, AnyTensor:$operand2);
@@ -2601,6 +2601,10 @@ def TestVersionedOpA : TEST_Op<"versionedA"> {
    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"> {
@@ -2742,6 +2746,51 @@ def TestOpWithNiceProperties : TEST_Op<"with_nice_properties"> {
   }];
 }
 
+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();
+    }
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Test Dataflow
 //===----------------------------------------------------------------------===//
@@ -2765,7 +2814,7 @@ def TestCallAndStoreOp : TEST_Op<"call_and_store",
 def TestStoreWithARegion : TEST_Op<"store_with_a_region",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
      SingleBlock]> {
-  let arguments = (ins 
+  let arguments = (ins
     Arg<AnyMemRef, "", [MemWrite]>:$address,
     BoolAttr:$store_before_region
   );