reading working

Author: nikalra

mlir/include/mlir/Bytecode/BytecodeOpInterface.td

@@ -49,17 +49,24 @@ def FallbackBytecodeOpInterface : OpInterface<"FallbackBytecodeOpInterface"> {
   let cppNamespace = "::mlir";
 
   let methods = [
+    StaticInterfaceMethod<[{
+      Set the original name for this operation from the bytecode.
+      }],
+      "void", "setOriginalOperationName", (ins
+          "::mlir::StringRef":$opName,
+          "::mlir::OperationState &":$state)
+    >,
     StaticInterfaceMethod<[{
       Read the properties blob for this operation from the bytecode and populate the state.
       }],
-      "LogicalResult", "readPropertiesBlob", (ins
-          "ArrayRef<char>":$blob,
+      "::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.
       }],
-      "ArrayRef<char>", "getPropertiesBlob", (ins)
+      "::mlir::ArrayRef<char>", "getPropertiesBlob", (ins)
     >,
   ];
 }

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

@@ -1415,8 +1415,8 @@ class mlir::BytecodeReader::Impl {
   /// Parse an operation name reference using the given reader, and set the
   /// `wasRegistered` flag that indicates if the bytecode was produced by a
   /// context where opName was registered.
-  FailureOr<OperationName> parseOpName(EncodingReader &reader,
-                                       std::optional<bool> &wasRegistered);
+  FailureOr<BytecodeOperationName *>
+  parseOpName(EncodingReader &reader, std::optional<bool> &wasRegistered);
 
   //===--------------------------------------------------------------------===//
   // Attribute/Type Section
@@ -1848,7 +1848,7 @@ BytecodeReader::Impl::parseDialectSection(ArrayRef<uint8_t> sectionData) {
   return success();
 }
 
-FailureOr<OperationName>
+FailureOr<BytecodeOperationName *>
 BytecodeReader::Impl::parseOpName(EncodingReader &reader,
                                   std::optional<bool> &wasRegistered) {
   BytecodeOperationName *opName = nullptr;
@@ -1868,21 +1868,28 @@ BytecodeReader::Impl::parseOpName(EncodingReader &reader,
       // Load the dialect and its version.
       DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
                                   dialectsMap, reader, version);
-      if (succeeded(opName->dialect->load(dialectReader, getContext()))) {
-        opName->opName.emplace(
-            (opName->dialect->name + "." + opName->name).str(), getContext());
-      } else if (auto fallbackOp =
-                     opName->dialect->interface->getFallbackOperationName();
-                 succeeded(fallbackOp)) {
-        // If the dialect's bytecode interface specifies a fallback op, we want
-        // to use that instead of an unregistered op.
-        opName->opName.emplace(*fallbackOp);
-      } else {
+      if (failed(opName->dialect->load(dialectReader, getContext())))
         return failure();
+
+      opName->opName.emplace((opName->dialect->name + "." + opName->name).str(),
+                             getContext());
+
+      // If the op is unregistered now, but was not marked as unregistered, try
+      // to parse it as a fallback op if the dialect's bytecode interface
+      // specifies one.
+      // We don't treat this condition as an error because we may still be able
+      // to parse the op as an unregistered op if it doesn't use custom
+      // properties encoding.
+      if (wasRegistered && !opName->opName->isRegistered()) {
+        if (auto fallbackOp =
+                opName->dialect->interface->getFallbackOperationName();
+            succeeded(fallbackOp)) {
+          opName->opName.emplace(*fallbackOp);
+        }
       }
     }
   }
-  return *opName->opName;
+  return opName;
 }
 
 //===----------------------------------------------------------------------===//
@@ -2227,8 +2234,12 @@ BytecodeReader::Impl::parseOpWithoutRegions(EncodingReader &reader,
                                             bool &isIsolatedFromAbove) {
   // Parse the name of the operation.
   std::optional<bool> wasRegistered;
-  FailureOr<OperationName> opName = parseOpName(reader, wasRegistered);
-  if (failed(opName))
+  FailureOr<BytecodeOperationName *> bytecodeOp =
+      parseOpName(reader, wasRegistered);
+  if (failed(bytecodeOp))
+    return failure();
+  auto opName = (*bytecodeOp)->opName;
+  if (!opName)
     return failure();
 
   // Parse the operation mask, which indicates which components of the operation
@@ -2245,6 +2256,9 @@ BytecodeReader::Impl::parseOpWithoutRegions(EncodingReader &reader,
   // With the location and name resolved, we can start building the operation
   // state.
   OperationState opState(opLoc, *opName);
+  // If this is a fallback op, provide the original name of the operation.
+  if (auto *iface = opName->getInterface<FallbackBytecodeOpInterface>())
+    iface->setOriginalOperationName((*bytecodeOp)->name, opState);
 
   // Parse the attributes of the operation.
   if (opMask & bytecode::OpEncodingMask::kHasAttrs) {

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

@@ -8,6 +8,7 @@
 
 #include "TestDialect.h"
 #include "TestOps.h"
+#include "mlir/IR/OperationSupport.h"
 #include "mlir/Interfaces/FoldInterfaces.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Transforms/InliningUtils.h"
@@ -92,6 +93,11 @@ struct TestBytecodeDialectInterface : public BytecodeDialectInterface {
     return Attribute();
   }
 
+  FailureOr<OperationName> getFallbackOperationName() const final {
+    return OperationName(TestBytecodeFallbackOp::getOperationName(),
+                         getContext());
+  }
+
   // Emit a specific version of the dialect.
   void writeVersion(DialectBytecodeWriter &writer) const final {
     // Construct the current dialect version.

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

@@ -9,9 +9,12 @@
 #include "TestDialect.h"
 #include "TestOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Interfaces/FunctionImplementation.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
+#include "llvm/Support/LogicalResult.h"
+#include <cstdint>
 
 using namespace mlir;
 using namespace test;
@@ -1230,6 +1233,53 @@ void TestVersionedOpA::writeProperties(mlir::DialectBytecodeWriter &writer) {
   writer.writeAttribute(prop.modifier);
 }
 
+//===----------------------------------------------------------------------===//
+// 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();
+
+//   return success();
+// }
+
+// void TestVersionedOpD::writeProperties(mlir::DialectBytecodeWriter &writer) {
+//   auto &prop = getProperties();
+//   writer.writeOwnedString("version 1");
+//   writer.writeAttribute(prop.attribute);
+// }
+
+//===----------------------------------------------------------------------===//
+// TestBytecodeFallbackOp
+//===----------------------------------------------------------------------===//
+
+void TestBytecodeFallbackOp::setOriginalOperationName(StringRef name,
+                                                      OperationState &state) {
+  state.getOrAddProperties<Properties>().setOpname(
+      StringAttr::get(state.getContext(), name));
+}
+
+LogicalResult
+TestBytecodeFallbackOp::readPropertiesBlob(ArrayRef<char> blob,
+                                           OperationState &state) {
+  state.getOrAddProperties<Properties>().bytecodeProperties =
+      DenseI8ArrayAttr::get(state.getContext(),
+                            ArrayRef((const int8_t *)blob.data(), blob.size()));
+  return success();
+}
+
+ArrayRef<char> TestBytecodeFallbackOp::getPropertiesBlob() {
+  return ArrayRef((const char *)getBytecodeProperties().data(),
+                  getBytecodeProperties().size());
+}
+
 //===----------------------------------------------------------------------===//
 // TestOpWithVersionedProperties
 //===----------------------------------------------------------------------===//

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

@@ -20,6 +20,7 @@ include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/PatternBase.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
+include "mlir/Bytecode/BytecodeOpInterface.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
@@ -31,7 +32,6 @@ include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/MemorySlotInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
-
 // Include the attribute definitions.
 include "TestAttrDefs.td"
 // Include the type definitions.
@@ -3030,6 +3030,25 @@ def TestVersionedOpC : TEST_Op<"versionedC"> {
   );
 }
 
+// def TestVersionedOpD : TEST_Op<"versionedD"> {
+//   let arguments = (ins AnyAttrOf<[TestAttrParams,
+//                                   I32ElementsAttr]>:$attribute
+//   );
+
+//   let useCustomPropertiesEncoding = 1;
+// }
+
+def TestBytecodeFallbackOp : TEST_Op<"bytecode.fallback", [
+  DeclareOpInterfaceMethods<FallbackBytecodeOpInterface, ["setOriginalOperationName", "readPropertiesBlob", "getPropertiesBlob"]>
+]> {
+  let arguments = (ins 
+    StrAttr:$opname,
+    DenseI8ArrayAttr:$bytecodeProperties,
+    Variadic<AnyType>:$operands);
+  let regions = (region VariadicRegion<AnyRegion>:$bodyRegions);
+  let results = (outs Variadic<AnyType>:$results);
+}
+
 //===----------------------------------------------------------------------===//
 // Test Properties
 //===----------------------------------------------------------------------===//