fix parsing & writing to be bitwise exact

Author: nikalra

mlir/include/mlir/Bytecode/BytecodeOpInterface.td

@@ -60,19 +60,7 @@ def FallbackBytecodeOpInterface : OpInterface<"FallbackBytecodeOpInterface"> {
       Get the original name for this operation from the bytecode.
       }],
       "::mlir::StringRef", "getOriginalOperationName", (ins)
-    >,
-    StaticInterfaceMethod<[{
-      Read the properties blob for this operation from the bytecode and populate the state.
-      }],
-      "::mlir::LogicalResult", "readPropertiesBlob", (ins
-          "::mlir::ArrayRef<char>":$blob,
-          "::mlir::OperationState &":$state)
-    >,
-    InterfaceMethod<[{
-      Get the properties blob for this operation to be emitted into the bytecode.
-      }],
-      "::mlir::ArrayRef<char>", "getPropertiesBlob", (ins)
-    >,
+    >
   ];
 }
 

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

@@ -1118,11 +1118,6 @@ class PropertiesSectionReader {
         return failure();
     }
 
-    // If the op is a "fallback" op, give it a handle to the raw properties
-    // buffer.
-    if (auto *iface = opName->getInterface<FallbackBytecodeOpInterface>())
-      return iface->readPropertiesBlob(rawProperties, opState);
-
     // Setup a new reader to read from the `rawProperties` sub-buffer.
     EncodingReader reader(
         StringRef(rawProperties.begin(), rawProperties.size()), fileLoc);
@@ -2182,6 +2177,9 @@ BytecodeReader::Impl::parseRegions(std::vector<RegionReadState> &regionStack,
 
         // Parse the bytecode.
         {
+          // If the op is registered (and serialized in a compatible manner), or
+          // unregistered but uses standard properties encoding, parsing without
+          // going through the fallback path should work.
           EncodingReader::Snapshot snapshot(reader);
           op = parseOpWithoutRegions(reader, readState, isIsolatedFromAbove,
                                      /*useDialectFallback=*/false);

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

@@ -518,13 +518,6 @@ class PropertiesSectionBuilder {
       return emit(scratch);
     }
 
-    if (auto iface = dyn_cast<FallbackBytecodeOpInterface>(op)) {
-      // Fallback ops should write the properties payload to the bytecode buffer
-      // directly.
-      ArrayRef<char> encodedProperties = iface.getPropertiesBlob();
-      return emit(encodedProperties);
-    }
-
     EncodingEmitter emitter;
     DialectWriter propertiesWriter(config.bytecodeVersion, emitter,
                                    numberingState, stringSection,

mlir/test/Bytecode/versioning/versioning-fallback.mlir

@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s --emit-bytecode > %T/versioning-fallback.mlirbc
+"test.versionedD"() <{attribute = #test.attr_params<42, 24>}> : () -> ()
+
+// COM: check that versionedD was parsed as a fallback op.
+// RUN: mlir-opt %T/versioning-fallback.mlirbc | FileCheck %s --check-prefix=CHECK-PARSE
+// CHECK-PARSE: test.bytecode.fallback 
+// CHECK-PARSE-SAME: opname = "test.versionedD"
+
+// COM: check that the bytecode roundtrip was successful
+// RUN: mlir-opt %T/versioning-fallback.mlirbc --verify-roundtrip
+
+// COM: check that the bytecode roundtrip is bitwise exact
+// RUN: mlir-opt %T/versioning-fallback.mlirbc --emit-bytecode | diff %T/versioning-fallback.mlirbc -

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

@@ -10,10 +10,13 @@
 #include "TestOps.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Interfaces/FunctionImplementation.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/LogicalResult.h"
 #include <cstdint>
 
@@ -1238,24 +1241,60 @@ void TestVersionedOpA::writeProperties(mlir::DialectBytecodeWriter &writer) {
 // TestVersionedOpD
 //===----------------------------------------------------------------------===//
 
-// LogicalResult
-// TestVersionedOpD::readProperties(mlir::DialectBytecodeReader &reader,
-//                                  mlir::OperationState &state) {
-//   auto &prop = state.getOrAddProperties<Properties>();
-//   StringRef res;
-//   if (failed(reader.readString(res)))
-//     return failure();
-//   if (failed(reader.readAttribute(prop.attribute)))
-//     return failure();
+LogicalResult
+TestVersionedOpD::readProperties(mlir::DialectBytecodeReader &reader,
+                                 mlir::OperationState &state) {
+  // Always fail so that this uses the fallback path.
+  return failure();
+}
+
+struct FallbackCompliantPropertiesEncoding {
+  int64_t version;
+  SmallVector<Attribute> requiredAttributes;
+  SmallVector<Attribute> optionalAttributes;
+
+  void writeProperties(DialectBytecodeWriter &writer) const {
+    // Write the op version.
+    writer.writeSignedVarInt(version);
+
+    // Write the required attributes.
+    writer.writeList(requiredAttributes,
+                     [&](Attribute attr) { writer.writeAttribute(attr); });
+
+    // Write the optional attributes.
+    writer.writeList(optionalAttributes, [&](Attribute attr) {
+      writer.writeOptionalAttribute(attr);
+    });
+  }
+
+  LogicalResult readProperties(DialectBytecodeReader &reader) {
+    // Read the op version.
+    if (failed(reader.readSignedVarInt(version)))
+      return failure();
 
-//   return success();
-// }
+    // Read the required attributes.
+    if (failed(reader.readList(requiredAttributes, [&](Attribute &attr) {
+          return reader.readAttribute(attr);
+        })))
+      return failure();
 
-// void TestVersionedOpD::writeProperties(mlir::DialectBytecodeWriter &writer) {
-//   auto &prop = getProperties();
-//   writer.writeOwnedString("version 1");
-//   writer.writeAttribute(prop.attribute);
-// }
+    // Read the optional attributes.
+    if (failed(reader.readList(optionalAttributes, [&](Attribute &attr) {
+          return reader.readOptionalAttribute(attr);
+        })))
+      return failure();
+
+    return success();
+  }
+};
+
+void TestVersionedOpD::writeProperties(mlir::DialectBytecodeWriter &writer) {
+  FallbackCompliantPropertiesEncoding encoding{
+      .version = 1,
+      .requiredAttributes = {getAttribute()},
+      .optionalAttributes = {}};
+  encoding.writeProperties(writer);
+}
 
 //===----------------------------------------------------------------------===//
 // TestBytecodeFallbackOp
@@ -1272,17 +1311,30 @@ StringRef TestBytecodeFallbackOp::getOriginalOperationName() {
 }
 
 LogicalResult
-TestBytecodeFallbackOp::readPropertiesBlob(ArrayRef<char> blob,
-                                           OperationState &state) {
-  state.getOrAddProperties<Properties>().bytecodeProperties =
-      DenseI8ArrayAttr::get(state.getContext(),
-                            ArrayRef((const int8_t *)blob.data(), blob.size()));
+TestBytecodeFallbackOp::readProperties(DialectBytecodeReader &reader,
+                                       OperationState &state) {
+  FallbackCompliantPropertiesEncoding encoding;
+  if (failed(encoding.readProperties(reader)))
+    return failure();
+
+  auto &props = state.getOrAddProperties<Properties>();
+  props.opversion = encoding.version;
+  props.encodedReqdAttributes =
+      ArrayAttr::get(state.getContext(), encoding.requiredAttributes);
+  props.encodedOptAttributes =
+      ArrayAttr::get(state.getContext(), encoding.optionalAttributes);
+
   return success();
 }
 
-ArrayRef<char> TestBytecodeFallbackOp::getPropertiesBlob() {
-  return ArrayRef((const char *)getBytecodeProperties().data(),
-                  getBytecodeProperties().size());
+void TestBytecodeFallbackOp::writeProperties(DialectBytecodeWriter &writer) {
+  FallbackCompliantPropertiesEncoding encoding{
+      .version = getOpversion(),
+      .requiredAttributes =
+          llvm::to_vector(getEncodedReqdAttributes().getValue()),
+      .optionalAttributes =
+          llvm::to_vector(getEncodedOptAttributes().getValue())};
+  encoding.writeProperties(writer);
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -3030,23 +3030,28 @@ def TestVersionedOpC : TEST_Op<"versionedC"> {
   );
 }
 
-// def TestVersionedOpD : TEST_Op<"versionedD"> {
-//   let arguments = (ins AnyAttrOf<[TestAttrParams,
-//                                   I32ElementsAttr]>:$attribute
-//   );
+// This op is used to generate tests for the bytecode dialect fallback path.
+def TestVersionedOpD : TEST_Op<"versionedD"> {
+  let arguments = (ins AnyAttrOf<[TestAttrParams,
+                                  I32ElementsAttr]>:$attribute
+  );
 
-//   let useCustomPropertiesEncoding = 1;
-// }
+  let useCustomPropertiesEncoding = 1;
+}
 
 def TestBytecodeFallbackOp : TEST_Op<"bytecode.fallback", [
   DeclareOpInterfaceMethods<FallbackBytecodeOpInterface, ["setOriginalOperationName", "readPropertiesBlob", "getPropertiesBlob"]>
 ]> {
   let arguments = (ins 
     StrAttr:$opname,
-    DenseI8ArrayAttr:$bytecodeProperties,
+    IntProp<"int64_t">:$opversion,
+    ArrayAttr:$encodedReqdAttributes,
+    ArrayAttr:$encodedOptAttributes,
     Variadic<AnyType>:$operands);
   let regions = (region VariadicRegion<AnyRegion>:$bodyRegions);
   let results = (outs Variadic<AnyType>:$results);
+
+  let useCustomPropertiesEncoding = 1;
 }
 
 //===----------------------------------------------------------------------===//