[mlir][bytecode] Fix dialect version parsing.

[jpienaar]

Summary

We were querying the wrong EncReader along some paths that resulted in
failures depending on if one encountered an Attribute from an unloaded
dialect before encountering an operation from that dialect.

Also fix error where we were able to emit "custom" form for an attribute
without custom form in TestDialect.

[mehdi_amini]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1391

I am puzzled, versionBuffer wasn't used at all, how did it ever worked?

[mfrancio]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1391

It is used in parseOpName, before loading the dialect reader:

// Load the dialect and its version.
EncodingReader versionReader(opName->dialect->versionBuffer, fileLoc);

[mehdi_amini]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1391

OK so we should remove it from there now that it is here right ?

[jpienaar]

Remove unused encReader created

This was created but never used as it would be replaced in call.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1391

Yes, that's why this only happened when attribute got materialized first. And yes removed.

[jpienaar]

Op was being printed as dialect, change to dialect.op form.

[mehdi_amini]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

840

I'm fine with introducing this method, but it seems like if this is the pattern to create a new DialectReader for a specific encoding buffer, we should be consistent and use it everywhere we're doing it.

That is there are a couple of other places like:

// Initialize the resource reader with the resource sections.
DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
                            reader);

1390

Instead of `withE

[jpienaar]

(I think your other comment got cut off)

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

840

We don't have a DialectReader to start with there. This is effectively "take DialectReader where everything else remains the same but with different EncodingReader" (e.g., clone with new encoding reader). I don't see this pattern used elsewhere here.

[mehdi_amini]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

840

Ah I missed that the other case didn't have a DialectReader at hand!

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

if (hasAlignment) {
       if (failed(parseVarInt(alignment)) || failed(alignTo(alignment)))
         return failure();
     }
 
     // Parse the actual section data.
     return parseBytes(static_cast<size_t>(length), sectionData);
   }
 
+  Location getLoc() const { return fileLoc; }
+
 private:
   /// Parse a variable length encoded integer from the byte stream. This method
   /// is a fallback when the number of bytes used to encode the value is greater
   /// than 1, but less than the max (9). The provided `result` value can be
   /// assumed to already contain the first byte of the value.
   /// NOTE: This method is marked noinline to avoid pessimizing the common case
   /// of single byte encoding.
   LLVM_ATTRIBUTE_NOINLINE LogicalResult parseMultiByteVarInt(uint64_t &result) {
DialectReader(AttrTypeReader &attrTypeReader,
                 ResourceSectionReader &resourceReader, EncodingReader &reader)
       : attrTypeReader(attrTypeReader), stringReader(stringReader),
         resourceReader(resourceReader), reader(reader) {}
 
   InFlightDiagnostic emitError(const Twine &msg) override {
     return reader.emitError(msg);
   }
 
+  DialectReader withEncodingReader(EncodingReader &encReader) {
+    return DialectReader(attrTypeReader, stringReader, resourceReader,
+                         encReader);
+  }
+
+  Location getLoc() const { return reader.getLoc(); }
+
   //===--------------------------------------------------------------------===//
   // IR
   //===--------------------------------------------------------------------===//
 
   LogicalResult readAttribute(Attribute &result) override {
     return attrTypeReader.parseAttribute(reader, result);
   }
 

 
 template <typename T>
 LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
                                                EncodingReader &reader,
                                                StringRef entryType) {
   DialectReader dialectReader(*this, stringReader, resourceReader, reader);
   if (failed(entry.dialect->load(dialectReader, fileLoc.getContext())))
     return failure();
-
   // Ensure that the dialect implements the bytecode interface.
   if (!entry.dialect->interface) {
     return reader.emitError("dialect '", entry.dialect->name,
                             "' does not implement the bytecode interface");
   }
 
   // Ask the dialect to parse the entry. If the dialect is versioned, parse
   // using the versioned encoding readers.
LogicalResult BytecodeDialect::load(DialectReader &reader, MLIRContext *ctx) {
   if (loadedDialect)
     interface = dyn_cast<BytecodeDialectInterface>(loadedDialect);
   if (!versionBuffer.empty()) {
     if (!interface)
       return reader.emitError("dialect '")
              << name
              << "' does not implement the bytecode interface, "
                 "but found a version entry";
-    loadedVersion = interface->readVersion(reader);
+    EncodingReader encReader(versionBuffer, reader.getLoc());
+    DialectReader versionReader = reader.withEncodingReader(encReader);
+    loadedVersion = interface->readVersion(versionReader);
     if (!loadedVersion)
       return failure();
   }
   return success();
 }
 
 LogicalResult
 BytecodeReader::parseDialectSection(ArrayRef<uint8_t> sectionData) {
FailureOr<OperationName> BytecodeReader::parseOpName(EncodingReader &reader) {
   BytecodeOperationName *opName = nullptr;
   if (failed(parseEntry(reader, opNames, opName, "operation name")))
     return failure();
 
   // 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.
-    EncodingReader versionReader(opName->dialect->versionBuffer, fileLoc);
     DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
-                                versionReader);
+                                reader);
     if (failed(opName->dialect->load(dialectReader, getContext())))
       return failure();
     opName->opName.emplace((opName->dialect->name + "." + opName->name).str(),
                            getContext());
   }
   return *opName->opName;
 }
 

mlir/test/Bytecode/general.mlir

 // RUN: mlir-opt -allow-unregistered-dialect -emit-bytecode %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
 
 // Bytecode currently does not support big-endian platforms
 // UNSUPPORTED: target=s390x-{{.*}}
 
 // CHECK-LABEL: "bytecode.test1"
+// CHECK-NEXT:    "unregistered.op"() {test_attr = #test.dynamic_singleton} : () -> ()
 // CHECK-NEXT:    "bytecode.empty"() : () -> ()
 // CHECK-NEXT:    "bytecode.attributes"() {attra = 10 : i64, attrb = #bytecode.attr} : () -> ()
 // CHECK-NEXT{LITERAL}: "bytecode.sparse"() {value = sparse<[[2, 1], [1, 1], [1, 2]], [1.
 // CHECK-NEXT:    test.graph_region {
 // CHECK-NEXT:      "bytecode.operands"(%[[RESULTS:.*]]#0, %[[RESULTS]]#1, %[[RESULTS]]#2) : (i32, i64, i32) -> ()
 // CHECK-NEXT:      %[[RESULTS]]:3 = "bytecode.results"() : () -> (i32, i64, i32)
 // CHECK-NEXT:    }
 // CHECK-NEXT:    "bytecode.branch"()[^[[BLOCK:.*]]] : () -> ()
 // CHECK-NEXT:  ^[[BLOCK]](%[[ARG0:.*]]: i32, %[[ARG1:.*]]: !bytecode.int, %[[ARG2:.*]]: !pdl.operation):
 // CHECK-NEXT:    "bytecode.regions"() ({
 // CHECK-NEXT:      "bytecode.operands"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (i32, !bytecode.int, !pdl.operation) -> ()
 // CHECK-NEXT:      "bytecode.return"() : () -> ()
 // CHECK-NEXT:    }) : () -> ()
 // CHECK-NEXT:    "bytecode.return"() : () -> ()
 // CHECK-NEXT:  }) : () -> ()
 
 "bytecode.test1"() ({
+  "unregistered.op"() {test_attr = #test.dynamic_singleton} : () -> ()
   "bytecode.empty"() : () -> ()
   "bytecode.attributes"() {attra = 10, attrb = #bytecode.attr} : () -> ()
   %cst = "bytecode.sparse"() {value = sparse<[[2, 1], [1, 1], [1, 2]], [1.0, 5.0, 6.0]> : tensor<8x7xf32>} : () -> (tensor<8x7xf32>)
   test.graph_region {
     "bytecode.operands"(%results#0, %results#1, %results#2) : (i32, i64, i32) -> ()
     %results:3 = "bytecode.results"() : () -> (i32, i64, i32)
   }
   "bytecode.branch"()[^secondBlock] : () -> ()
 
 ^secondBlock(%arg1: i32, %arg2: !bytecode.int, %arg3: !pdl.operation):
   "bytecode.regions"() ({
     "bytecode.operands"(%arg1, %arg2, %arg3) : (i32, !bytecode.int, !pdl.operation) -> ()
     "bytecode.return"() : () -> ()
   }) : () -> ()
   "bytecode.return"() : () -> ()
 }) : () -> ()

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

struct TestBytecodeDialectInterface : public BytecodeDialectInterface {
   LogicalResult writeAttribute(Attribute attr,
                                DialectBytecodeWriter &writer) const final {
     if (auto concreteAttr = llvm::dyn_cast<TestAttrParamsAttr>(attr)) {
       writer.writeVarInt(test_encoding::k_attr_params);
       writer.writeVarInt(concreteAttr.getV0());
       writer.writeVarInt(concreteAttr.getV1());
       return success();
     }
-    writer.writeAttribute(attr);
-    return success();
+    return failure();
   }
 
   Attribute readAttribute(DialectBytecodeReader &reader,
                           const DialectVersion &version_) const final {
     const auto &version = static_cast<const TestDialectVersion &>(version_);
     if (version.major < 2)
       return readAttrOldEncoding(reader);
     if (version.major == 2 && version.minor == 0)