[Mlir-commits] [mlir] 3449e7a - [mlir][bytecode] Fix dialect version parsing.

[Jacques Pienaar]

Author: Jacques Pienaar
Date: 2023-05-11T05:19:06-07:00
New Revision: 3449e7a832ba21f518e69c72985beea3b3c77a7e

URL: https://github.com/llvm/llvm-project/commit/3449e7a832ba21f518e69c72985beea3b3c77a7e
DIFF: https://github.com/llvm/llvm-project/commit/3449e7a832ba21f518e69c72985beea3b3c77a7e.diff

LOG: [mlir][bytecode] Fix dialect version parsing.

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.

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

Added: 
    

Modified: 
    mlir/lib/Bytecode/Reader/BytecodeReader.cpp
    mlir/test/Bytecode/general.mlir
    mlir/test/lib/Dialect/Test/TestDialect.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
index d6f1e18d35ae9..e39c56885c205 100644
--- a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
+++ b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
@@ -271,6 +271,8 @@ class EncodingReader {
     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
@@ -835,6 +837,13 @@ class DialectReader : public DialectBytecodeReader {
     return reader.emitError(msg);
   }
 
+  DialectReader withEncodingReader(EncodingReader &encReader) {
+    return DialectReader(attrTypeReader, stringReader, resourceReader,
+                         encReader);
+  }
+
+  Location getLoc() const { return reader.getLoc(); }
+
   //===--------------------------------------------------------------------===//
   // IR
   //===--------------------------------------------------------------------===//
@@ -1054,7 +1063,6 @@ LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
   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,
@@ -1378,7 +1386,9 @@ LogicalResult BytecodeDialect::load(DialectReader &reader, MLIRContext *ctx) {
              << 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();
   }
@@ -1448,9 +1458,8 @@ FailureOr<OperationName> BytecodeReader::parseOpName(EncodingReader &reader) {
   // 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(),

diff  --git a/mlir/test/Bytecode/general.mlir b/mlir/test/Bytecode/general.mlir
index 5d707adf0ae38..071b4c36526cd 100644
--- a/mlir/test/Bytecode/general.mlir
+++ b/mlir/test/Bytecode/general.mlir
@@ -4,6 +4,7 @@
 // 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.
@@ -21,6 +22,7 @@
 // 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>)

diff  --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 83226a5f3c245..4660d9abe6769 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -128,8 +128,7 @@ struct TestBytecodeDialectInterface : public BytecodeDialectInterface {
       writer.writeVarInt(concreteAttr.getV1());
       return success();
     }
-    writer.writeAttribute(attr);
-    return success();
+    return failure();
   }
 
   Attribute readAttribute(DialectBytecodeReader &reader,