roundtrip working

Author: nikalra

mlir/include/mlir/Bytecode/BytecodeOpInterface.td

@@ -56,6 +56,11 @@ def FallbackBytecodeOpInterface : OpInterface<"FallbackBytecodeOpInterface"> {
           "::mlir::StringRef":$opName,
           "::mlir::OperationState &":$state)
     >,
+    InterfaceMethod<[{
+      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/lib/Bytecode/Writer/BytecodeWriter.cpp

@@ -848,11 +848,16 @@ void BytecodeWriter::writeDialectSection(EncodingEmitter &emitter) {
 
   // Emit the referenced operation names grouped by dialect.
   auto emitOpName = [&](OpNameNumbering &name) {
+    bool isRegistered = name.name.isRegistered();
+    // If we're writing a fallback op, write it as if it were a registered op.
+    if (name.name.hasInterface<FallbackBytecodeOpInterface>())
+      isRegistered = true;
+
     size_t stringId = stringSection.insert(name.name.stripDialect());
     if (config.bytecodeVersion < bytecode::kNativePropertiesEncoding)
       dialectEmitter.emitVarInt(stringId, "dialect op name");
     else
-      dialectEmitter.emitVarIntWithFlag(stringId, name.name.isRegistered(),
+      dialectEmitter.emitVarIntWithFlag(stringId, isRegistered,
                                         "dialect op name");
   };
   writeDialectGrouping(dialectEmitter, numberingState.getOpNames(), emitOpName);
@@ -991,7 +996,16 @@ LogicalResult BytecodeWriter::writeBlock(EncodingEmitter &emitter,
 }
 
 LogicalResult BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
-  emitter.emitVarInt(numberingState.getNumber(op->getName()), "op name ID");
+  OperationName opName = op->getName();
+  // For fallback ops, create a new operation name referencing the original op
+  // instead.
+  if (auto fallback = dyn_cast<FallbackBytecodeOpInterface>(op))
+    opName = OperationName((fallback->getDialect()->getNamespace() + "." +
+                            fallback.getOriginalOperationName())
+                               .str(),
+                           op->getContext());
+
+  emitter.emitVarInt(numberingState.getNumber(opName), "op name ID");
 
   // Emit a mask for the operation components. We need to fill this in later
   // (when we actually know what needs to be emitted), so emit a placeholder for

mlir/lib/Bytecode/Writer/IRNumbering.cpp

@@ -419,7 +419,17 @@ void IRNumberingState::number(Region &region) {
 void IRNumberingState::number(Operation &op) {
   // Number the components of an operation that won't be numbered elsewhere
   // (e.g. we don't number operands, regions, or successors here).
-  number(op.getName());
+
+  OperationName opName = op.getName();
+  // For fallback ops, create a new operation name referencing the original op
+  // instead.
+  if (auto fallback = dyn_cast<FallbackBytecodeOpInterface>(op))
+    opName = OperationName((fallback->getDialect()->getNamespace() + "." +
+                            fallback.getOriginalOperationName())
+                               .str(),
+                           op.getContext());
+  number(opName);
+
   for (OpResult result : op.getResults()) {
     valueIDs.try_emplace(result, nextValueID++);
     number(result.getType());

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

@@ -8,6 +8,7 @@
 
 #include "TestDialect.h"
 #include "TestOps.h"
+#include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Verifier.h"
@@ -1266,6 +1267,10 @@ void TestBytecodeFallbackOp::setOriginalOperationName(StringRef name,
       StringAttr::get(state.getContext(), name));
 }
 
+StringRef TestBytecodeFallbackOp::getOriginalOperationName() {
+  return getProperties().getOpname().getValue();
+}
+
 LogicalResult
 TestBytecodeFallbackOp::readPropertiesBlob(ArrayRef<char> blob,
                                            OperationState &state) {