refactor / cleanup

Author: nikalra

mlir/include/mlir/Bytecode/BytecodeOpInterface.td

@@ -53,7 +53,7 @@ def FallbackBytecodeOpInterface : OpInterface<"FallbackBytecodeOpInterface"> {
       Set the original name for this operation from the bytecode.
       }],
       "void", "setOriginalOperationName", (ins
-          "::mlir::StringRef":$opName,
+          "const ::mlir::Twine&":$opName,
           "::mlir::OperationState &":$state)
     >,
     InterfaceMethod<[{

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

@@ -21,6 +21,7 @@
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/Twine.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/Support/SourceMgr.h"
@@ -2257,8 +2258,11 @@ BytecodeReader::Impl::parseOpWithoutRegions(EncodingReader &reader,
   // 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);
+  if (auto *iface = opName->getInterface<FallbackBytecodeOpInterface>()) {
+    const Twine originalName =
+        opName->getDialect()->getNamespace() + "." + (*bytecodeOp)->name;
+    iface->setOriginalOperationName(originalName, opState);
+  }
 
   // Parse the attributes of the operation.
   if (opMask & bytecode::OpEncodingMask::kHasAttrs) {

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

@@ -848,17 +848,12 @@ 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;
-
+    const bool isKnownOp = name.isOpaqueEntry || name.name.isRegistered();
     size_t stringId = stringSection.insert(name.name.stripDialect());
     if (config.bytecodeVersion < bytecode::kNativePropertiesEncoding)
       dialectEmitter.emitVarInt(stringId, "dialect op name");
     else
-      dialectEmitter.emitVarIntWithFlag(stringId, isRegistered,
-                                        "dialect op name");
+      dialectEmitter.emitVarIntWithFlag(stringId, isKnownOp, "dialect op name");
   };
   writeDialectGrouping(dialectEmitter, numberingState.getOpNames(), emitOpName);
 
@@ -1000,10 +995,8 @@ LogicalResult BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
   // 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());
+    opName =
+        OperationName(fallback.getOriginalOperationName(), op->getContext());
 
   emitter.emitVarInt(numberingState.getNumber(opName), "op name ID");
 

mlir/lib/Bytecode/Writer/IRNumbering.cpp

@@ -420,15 +420,14 @@ 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).
 
-  OperationName opName = op.getName();
-  // For fallback ops, create a new operation name referencing the original op
+  // For fallback ops, create a new OperationName referencing the original op
   // instead.
-  if (auto fallback = dyn_cast<FallbackBytecodeOpInterface>(op))
-    opName = OperationName((fallback->getDialect()->getNamespace() + "." +
-                            fallback.getOriginalOperationName())
-                               .str(),
-                           op.getContext());
-  number(opName);
+  if (auto fallback = dyn_cast<FallbackBytecodeOpInterface>(op)) {
+    OperationName opName(fallback.getOriginalOperationName(), op.getContext());
+    number(opName, /*isOpaque=*/true);
+  } else {
+    number(op.getName(), /*isOpaque=*/false);
+  }
 
   for (OpResult result : op.getResults()) {
     valueIDs.try_emplace(result, nextValueID++);
@@ -467,7 +466,7 @@ void IRNumberingState::number(Operation &op) {
   number(op.getLoc());
 }
 
-void IRNumberingState::number(OperationName opName) {
+void IRNumberingState::number(OperationName opName, bool isOpaque) {
   OpNameNumbering *&numbering = opNames[opName];
   if (numbering) {
     ++numbering->refCount;
@@ -479,8 +478,8 @@ void IRNumberingState::number(OperationName opName) {
   else
     dialectNumber = &numberDialect(opName.getDialectNamespace());
 
-  numbering =
-      new (opNameAllocator.Allocate()) OpNameNumbering(dialectNumber, opName);
+  numbering = new (opNameAllocator.Allocate())
+      OpNameNumbering(dialectNumber, opName, isOpaque);
   orderedOpNames.push_back(numbering);
 }
 

mlir/lib/Bytecode/Writer/IRNumbering.h

@@ -63,15 +63,18 @@ struct TypeNumbering : public AttrTypeNumbering {
 
 /// This class represents the numbering entry of an operation name.
 struct OpNameNumbering {
-  OpNameNumbering(DialectNumbering *dialect, OperationName name)
-      : dialect(dialect), name(name) {}
+  OpNameNumbering(DialectNumbering *dialect, OperationName name, bool isOpaque)
+      : dialect(dialect), name(name), isOpaqueEntry(isOpaque) {}
 
   /// The dialect of this value.
   DialectNumbering *dialect;
 
   /// The concrete name.
   OperationName name;
 
+  /// This entry represents an opaque operation entry.
+  bool isOpaqueEntry = false;
+
   /// The number assigned to this name.
   unsigned number = 0;
 
@@ -210,7 +213,7 @@ class IRNumberingState {
 
   /// Get the set desired bytecode version to emit.
   int64_t getDesiredBytecodeVersion() const;
-  
+
 private:
   /// This class is used to provide a fake dialect writer for numbering nested
   /// attributes and types.
@@ -225,7 +228,7 @@ class IRNumberingState {
   DialectNumbering &numberDialect(Dialect *dialect);
   DialectNumbering &numberDialect(StringRef dialect);
   void number(Operation &op);
-  void number(OperationName opName);
+  void number(OperationName opName, bool isOpaque);
   void number(Region &region);
   void number(Type type);
 

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

@@ -1261,7 +1261,7 @@ void TestVersionedOpA::writeProperties(mlir::DialectBytecodeWriter &writer) {
 // TestBytecodeFallbackOp
 //===----------------------------------------------------------------------===//
 
-void TestBytecodeFallbackOp::setOriginalOperationName(StringRef name,
+void TestBytecodeFallbackOp::setOriginalOperationName(const Twine &name,
                                                       OperationState &state) {
   state.getOrAddProperties<Properties>().setOpname(
       StringAttr::get(state.getContext(), name));