[Mlir-commits] [mlir] [mlir][bytecode] Fix D155919 and enable backward-compatibility and back-deployment between version 5 and version 6 of the bytecode encoding. (PR #70498)

[Matteo Franciolini]

https://github.com/mfrancio updated https://github.com/llvm/llvm-project/pull/70498

>From b9760c2ec5aec694e857b1e5fce7e3cada04e5c0 Mon Sep 17 00:00:00 2001
From: Matteo Franciolini <m_franciolini at apple.com>
Date: Fri, 27 Oct 2023 12:34:26 -0700
Subject: [PATCH] [mlir][bytecode] Fix D155919 and enable
 backward-compatibility and back-deployment between version 5 and version 6 of
 the bytecode encoding.

Before D155919, when a dialect was leveraging properties to store attributes with `usePropertiesForAttributes`, the operand segment sizes attribute was emitted in the property section in sorted order together with all the other attributes of an op. After D155919, version 5 of the bytecode was emitting and parsing operand segment sizes after all the properties of an op, breaking backward compatibility and back deployment. This patch fixes the emission ordering and allows to parse bytecode files emitted before (D155919) with version 5 of MLIR bytecode. The patch also enables to emit correctly version 5 of MLIR bytecode.
---
 mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp | 105 +++++++++++++++-----
 1 file changed, 81 insertions(+), 24 deletions(-)

diff --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index 7029c0eac15c392..a620265b3dd8809 100644
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
@@ -156,26 +156,36 @@ static const char *const valueRangeReturnCode = R"(
 )";
 
 /// Read operand/result segment_size from bytecode.
-static const char *const readBytecodeSegmentSize = R"(
-if ($_reader.getBytecodeVersion() < /*kNativePropertiesODSSegmentSize=*/6) {
-  ::mlir::DenseI32ArrayAttr attr;
-  if (::mlir::failed($_reader.readAttribute(attr))) return ::mlir::failure();
-  if (attr.size() > static_cast<int64_t>(sizeof($_storage) / sizeof(int32_t))) {
-    $_reader.emitError("size mismatch for operand/result_segment_size");
-    return ::mlir::failure();
+static const char *const readBytecodeSegmentSizeNative = R"(
+  if ($_reader.getBytecodeVersion() >= /*kNativePropertiesODSSegmentSize=*/6)
+    return $_reader.readSparseArray(::llvm::MutableArrayRef($_storage));
+)";
+
+static const char *const readBytecodeSegmentSizeLegacy = R"(
+  if ($_reader.getBytecodeVersion() < /*kNativePropertiesODSSegmentSize=*/6) {
+    auto &$_storage = prop.$_propName;
+    ::mlir::DenseI32ArrayAttr attr;
+    if (::mlir::failed($_reader.readAttribute(attr))) return ::mlir::failure();
+    if (attr.size() > static_cast<int64_t>(sizeof($_storage) / sizeof(int32_t))) {
+      $_reader.emitError("size mismatch for operand/result_segment_size");
+      return ::mlir::failure();
+    }
+    ::llvm::copy(::llvm::ArrayRef<int32_t>(attr), $_storage.begin());
   }
-  ::llvm::copy(::llvm::ArrayRef<int32_t>(attr), $_storage.begin());
-} else {
-  return $_reader.readSparseArray(::llvm::MutableArrayRef($_storage));
-}
 )";
 
 /// Write operand/result segment_size to bytecode.
-static const char *const writeBytecodeSegmentSize = R"(
-if ($_writer.getBytecodeVersion() < /*kNativePropertiesODSSegmentSize=*/6)
-  $_writer.writeAttribute(::mlir::DenseI32ArrayAttr::get(getContext(), $_storage));
-else
-  $_writer.writeSparseArray(::llvm::ArrayRef($_storage));
+static const char *const writeBytecodeSegmentSizeNative = R"(
+  if ($_writer.getBytecodeVersion() >= /*kNativePropertiesODSSegmentSize=*/6)
+    $_writer.writeSparseArray(::llvm::ArrayRef($_storage));
+)";
+
+/// Write operand/result segment_size to bytecode.
+static const char *const writeBytecodeSegmentSizeLegacy = R"(
+if ($_writer.getBytecodeVersion() < /*kNativePropertiesODSSegmentSize=*/6) {
+  auto &$_storage = prop.$_propName;
+  $_writer.writeAttribute(::mlir::DenseI32ArrayAttr::get($_ctxt, $_storage));
+}
 )";
 
 /// A header for indicating code sections.
@@ -389,6 +399,14 @@ class OpOrAdaptorHelper {
     return resultSegmentsSize;
   }
 
+  uint32_t getOperandSegmentSizesLegacyIndex() {
+    return operandSegmentSizesLegacyIndex;
+  }
+
+  uint32_t getResultSegmentSizesLegacyIndex() {
+    return resultSegmentSizesLegacyIndex;
+  }
+
 private:
   // Compute the attribute metadata.
   void computeAttrMetadata();
@@ -407,6 +425,12 @@ class OpOrAdaptorHelper {
   std::optional<NamedProperty> resultSegmentsSize;
   std::string resultSegmentsSizeStorage;
 
+  // Indices to store the position in the emission order of the operand/result
+  // segment sizes attribute if emitted as part of the properties for legacy
+  // bytecode encodings, i.e. versions less than 6.
+  uint32_t operandSegmentSizesLegacyIndex = 0;
+  uint32_t resultSegmentSizesLegacyIndex = 0;
+
   // The number of required attributes.
   unsigned numRequired;
 };
@@ -435,8 +459,8 @@ void OpOrAdaptorHelper::computeAttrMetadata() {
         "::mlir::DenseI32ArrayAttr::get($_ctxt, $_storage)",
         /*convertFromAttributeCall=*/
         "return convertFromAttribute($_storage, $_attr, $_diag);",
-        /*readFromMlirBytecodeCall=*/readBytecodeSegmentSize,
-        /*writeToMlirBytecodeCall=*/writeBytecodeSegmentSize,
+        /*readFromMlirBytecodeCall=*/readBytecodeSegmentSizeNative,
+        /*writeToMlirBytecodeCall=*/writeBytecodeSegmentSizeNative,
         /*hashPropertyCall=*/
         "::llvm::hash_combine_range(std::begin($_storage), "
         "std::end($_storage));",
@@ -478,6 +502,21 @@ void OpOrAdaptorHelper::computeAttrMetadata() {
                return lhs.attrName < rhs.attrName;
              });
 
+  // Store the position of the legacy operand_segment_sizes /
+  // result_segment_sizes so we can emit a backward compatible property readers
+  // and writers.
+  StringRef legacyOperandSegmentSizeName =
+      StringLiteral("operand_segment_sizes");
+  StringRef legacyResultSegmentSizeName = StringLiteral("result_segment_sizes");
+  operandSegmentSizesLegacyIndex = 0;
+  resultSegmentSizesLegacyIndex = 0;
+  for (auto item : sortedAttrMetadata) {
+    if (item.attrName < legacyOperandSegmentSizeName)
+      ++operandSegmentSizesLegacyIndex;
+    if (item.attrName < legacyResultSegmentSizeName)
+      ++resultSegmentSizesLegacyIndex;
+  }
+
   // Compute the subrange bounds for each attribute.
   numRequired = 0;
   for (AttributeMetadata &attr : sortedAttrMetadata) {
@@ -1580,15 +1619,33 @@ void OpEmitter::genPropertiesSupportForBytecode(
   readPropertiesMethod
       << "  auto &prop = state.getOrAddProperties<Properties>(); (void)prop;";
   writePropertiesMethod << "  auto &prop = getProperties(); (void)prop;\n";
-  for (const auto &attrOrProp : attrOrProperties) {
+  for (const auto &item : llvm::enumerate(attrOrProperties)) {
+    auto &attrOrProp = item.value();
+    FmtContext fctx;
+    fctx.addSubst("_reader", "reader")
+        .addSubst("_writer", "writer")
+        .addSubst("_storage", propertyStorage)
+        .addSubst("_ctxt", "this->getContext()");
+    // If the op emits operand/result segment sizes as a property, emit the
+    // legacy reader/writer in the appropriate order to allow backward
+    // compatibility and back deployment.
+    if (emitHelper.getOperandSegmentsSize().has_value() &&
+        item.index() == emitHelper.getOperandSegmentSizesLegacyIndex()) {
+      FmtContext fmtCtxt(fctx);
+      fmtCtxt.addSubst("_propName", operandSegmentAttrName);
+      readPropertiesMethod << tgfmt(readBytecodeSegmentSizeLegacy, &fmtCtxt);
+      writePropertiesMethod << tgfmt(writeBytecodeSegmentSizeLegacy, &fmtCtxt);
+    }
+    if (emitHelper.getResultSegmentsSize().has_value() &&
+        item.index() == emitHelper.getResultSegmentSizesLegacyIndex()) {
+      FmtContext fmtCtxt(fctx);
+      fmtCtxt.addSubst("_propName", resultSegmentAttrName);
+      readPropertiesMethod << tgfmt(readBytecodeSegmentSizeLegacy, &fmtCtxt);
+      writePropertiesMethod << tgfmt(writeBytecodeSegmentSizeLegacy, &fmtCtxt);
+    }
     if (const auto *namedProperty =
             attrOrProp.dyn_cast<const NamedProperty *>()) {
       StringRef name = namedProperty->name;
-      FmtContext fctx;
-      fctx.addSubst("_reader", "reader")
-          .addSubst("_writer", "writer")
-          .addSubst("_storage", propertyStorage)
-          .addSubst("_ctxt", "this->getContext()");
       readPropertiesMethod << formatv(
           R"(
   {{