[mlir] Implements IR versioning capability (WIP)

[mehdi_amini]

Cleanup

[jpienaar]

mlir/docs/BytecodeFormat.md

411

So the number of instances here is not required as below one require a version for every dialect? (even if there is 1 dialect without a version, it is cheaper to have a varint for number of versioned dialects here up to 511 dialects in file it would seem)

mlir/include/mlir/IR/OpImplementation.h

472

Why should this be in OpImplementation.h instead of localized to Bytecode? (it seems like nothing actually queries it outside there)

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

437

Handle feels misleading, this is raw data from which version will be extracted.

1388

Should this be checked in the verify required sections function? Also, why is it required? E.g., what if I just want to use within same rev build?

mlir/test/Bytecode/versioning/versioned_op.mlir

1

Outdated.

41

Could you also add example of what .mlir input with version looks like?

[mehdi_amini]

Marked as abandoned, @mfrancio will merge this all into https://reviews.llvm.org/D143647

mlir/docs/BytecodeFormat.md

411

I haven't updated the doc from the revision I took as a basis, please ignore this file :)

mlir/include/mlir/IR/OpImplementation.h

472

Absolutely!

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1388

This is a nested section, we can't verify early I think.

It is not as much "required" as we don't expect any other section here, and I'm not sure how to express an optional section? (so we'll always write an empty one)

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

373

Code above is pure code motion because of DialectWriter used in BytecodeWriter::writeDialectSection() below.

mlir/test/Bytecode/versioning/versioned_op.mlir

41

We won't support any version in .mlir file right now (tests are outdated)

mlir/docs/BytecodeFormat.md

block_argument {
   typeIndex: varint,
   location: varint
 }
 ```
 
 A block is encoded with an array of operations and block arguments. The first
 field is an encoding that combines the number of operations in the block, with a
 flag indicating if the block has arguments.
+
+### Dialect Version Section
+
+```
+dialect_version_section {
+  dialects: dialect_info[]
+}
+
+dialect_info {
+  dialect: varint,
+  size: varint,
+  version: byte[]
+}
+```
+
+The dialect version section contains details about the version for each of the
+versioned dialects contained in the module. A dialect is considered versioned if
+implements the OpAsmDialectInterface hooks as described in the MLIR Language
+Reference. The dialect info section is optional and written only if the dialect 
+is registered in the context. When version information for a versioned dialect 
+are missing, a dialect will be parsed without performing upgrades.
 No newline at end of file

mlir/docs/LangRef.md

 See [here](DefiningDialects/AttributesAndTypes.md) on how to define dialect attribute values.
 
 ### Builtin Attribute Values
 
 The [builtin dialect](Dialects/Builtin.md) defines a set of attribute values
 that are directly usable by any other dialect in MLIR. These types cover a range
 from primitive integer and floating-point values, attribute dictionaries, dense
 multi-dimensional arrays, and more.
+
+### IR Versionning
+
+A dialect can opt-in to handle versioning in a custom way. Two hooks to expose a 
+dialect version to the printer and parser are available through the 
+`OpAsmDialectInterface`. First, the `getProducerVersion()` method allows to 
+inject a target dialect version to the printer, which will be written to the 
+output mlir file. Second, the `upgradeFromVersion()` method allows to retrieve 
+the version information while parsing the input IR, and gives an opportunity to 
+each dialect for which a version is present to perform IR upgrades.
+
+Version information are stored on a buffer through an `AsmDialectVersionHandle`. 
+There is no restriction on what kind of information a dialect is allowed to 
+encode to model its versioning. Currently, versioning is supported only for 
+bytecode formats.
 No newline at end of file

mlir/include/mlir/Bytecode/BytecodeImplementation.h

public:
   /// Write the given type, which belongs to this dialect, to the given writer.
   /// This method may return failure to indicate that the given type could not
   /// be encoded, in which case the textual format will be used to encode this
   /// type instead.
   virtual LogicalResult writeType(Type type,
                                   DialectBytecodeWriter &writer) const {
     return failure();
   }
+
+  /// Write the version of this dialect to the given writer.
+  /// The first emitted VarInt should be the size of the entry.
+  virtual void writeVersion(DialectBytecodeWriter &writer) const {}
+
+  virtual std::unique_ptr<DialectVersion>
+  parseVersion(DialectBytecodeReader &reader,
+               ArrayRef<uint8_t> versionHandle) const {
+    reader.emitError("Dialect does not support versioning");
+    return nullptr;
+  }
+
+  /// Hook invoked after parsing completed, if a version directive was present
+  /// and included an entry for the current dialect. This hook offers the
+  /// opportunity to the dialect to visit the IR and upgrades constructs emitted
+  /// by the version of the dialect corresponding to the provided version.
+  virtual LogicalResult
+  upgradeFromVersion(Operation *topLevelOp,
+                     const DialectVersion &version) const {
+    return success();
+  }
 };
 
 } // namespace mlir
 
 #endif // MLIR_BYTECODE_BYTECODEIMPLEMENTATION_H

mlir/include/mlir/IR/OpImplementation.h

inline OpAsmPrinter &operator<<(OpAsmPrinter &p, const T &values) {
   return p;
 }
 
 inline OpAsmPrinter &operator<<(OpAsmPrinter &p, Block *value) {
   p.printSuccessor(value);
   return p;
 }
 
+//===--------------------------------------------------------------------===//
+// Dialect Asm Version Interface.
+//===--------------------------------------------------------------------===//
+
+/// This class is used to represent the version of a dialect, for the purpose
+/// of polymorphic destruction.
+class DialectVersion {
+public:
+  virtual ~DialectVersion() = default;
+};
+
 //===----------------------------------------------------------------------===//
 // AsmParser
 //===----------------------------------------------------------------------===//
 
 /// This base class exposes generic asm parser hooks, usable across the various
 /// derived parsers.
 class AsmParser {
 public:
parseAffineExprOfSSAIds(SmallVectorImpl<UnresolvedOperand> &dimOperands,
                           SmallVectorImpl<UnresolvedOperand> &symbOperands,
                           AffineExpr &expr) = 0;
 
   //===--------------------------------------------------------------------===//
   // Argument Parsing
   //===--------------------------------------------------------------------===//
 
   struct Argument {
-    UnresolvedOperand ssaName;    // SourceLoc, SSA name, result #.
-    Type type;                    // Type.
-    DictionaryAttr attrs;         // Attributes if present.
+    UnresolvedOperand ssaName;         // SourceLoc, SSA name, result #.
+    Type type;                         // Type.
+    DictionaryAttr attrs;              // Attributes if present.
     std::optional<Location> sourceLoc; // Source location specifier if present.
   };
 
   /// Parse a single argument with the following syntax:
   ///
   ///   `%ssaName : !type { optionalAttrDict} loc(optionalSourceLoc)`
   ///
   /// If `allowType` is false or `allowAttrs` are false then the respective

mlir/lib/Bytecode/Encoding.h

 
 namespace mlir {
 namespace bytecode {
 //===----------------------------------------------------------------------===//
 // General constants
 //===----------------------------------------------------------------------===//
 
 enum {
+  /// The minimum supported version of the bytecode.
+  kMinSupportedVersion = 0,
+
   /// The current bytecode version.
-  kVersion = 0,
+  kVersion = 1,
 
   /// An arbitrary value used to fill alignment padding.
   kAlignmentByte = 0xCB,
 };
 
 //===----------------------------------------------------------------------===//
 // Sections
 //===----------------------------------------------------------------------===//
enum ID : uint8_t {
 
   /// This section contains the resources of the bytecode.
   kResource = 5,
 
   /// This section contains the offsets of resources within the Resource
   /// section.
   kResourceOffset = 6,
 
+  /// This section contains the versions of each dialect.
+  kDialectVersions = 7,
+
   /// The total number of section types.
-  kNumSections = 7,
+  kNumSections = 8,
 };
 } // namespace Section
 
 //===----------------------------------------------------------------------===//
 // IR Section
 //===----------------------------------------------------------------------===//
 
 /// This enum represents a mask of all of the potential components of an

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

static std::string toString(bytecode::Section::ID sectionID) {
   case bytecode::Section::kAttrTypeOffset:
     return "AttrTypeOffset (3)";
   case bytecode::Section::kIR:
     return "IR (4)";
   case bytecode::Section::kResource:
     return "Resource (5)";
   case bytecode::Section::kResourceOffset:
     return "ResourceOffset (6)";
+  case bytecode::Section::kDialectVersions:
+    return "DialectVersions (7)";
   default:
     return ("Unknown (" + Twine(static_cast<unsigned>(sectionID)) + ")").str();
   }
 }
 
 /// Returns true if the given top-level section ID is optional.
 static bool isSectionOptional(bytecode::Section::ID sectionID) {
   switch (sectionID) {
   case bytecode::Section::kString:
   case bytecode::Section::kDialect:
   case bytecode::Section::kAttrType:
   case bytecode::Section::kAttrTypeOffset:
   case bytecode::Section::kIR:
     return false;
   case bytecode::Section::kResource:
   case bytecode::Section::kResourceOffset:
+  case bytecode::Section::kDialectVersions:
     return true;
   default:
     llvm_unreachable("unknown section ID");
   }
 }
 
 //===----------------------------------------------------------------------===//
 // EncodingReader
LogicalResult StringSectionReader::initialize(Location fileLoc,
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // BytecodeDialect
 //===----------------------------------------------------------------------===//
 
 namespace {
+class DialectReader;
+
 /// This struct represents a dialect entry within the bytecode.
 struct BytecodeDialect {
   /// Load the dialect into the provided context if it hasn't been loaded yet.
   /// Returns failure if the dialect couldn't be loaded *and* the provided
   /// context does not allow unregistered dialects. The provided reader is used
   /// for error emission if necessary.
-  LogicalResult load(EncodingReader &reader, MLIRContext *ctx) {
-    if (dialect)
-      return success();
-    Dialect *loadedDialect = ctx->getOrLoadDialect(name);
-    if (!loadedDialect && !ctx->allowsUnregisteredDialects()) {
-      return reader.emitError(
-          "dialect '", name,
-          "' is unknown. If this is intended, please call "
-          "allowUnregisteredDialects() on the MLIRContext, or use "
-          "-allow-unregistered-dialect with the MLIR tool used.");
-    }
-    dialect = loadedDialect;
-
-    // If the dialect was actually loaded, check to see if it has a bytecode
-    // interface.
-    if (loadedDialect)
-      interface = dyn_cast<BytecodeDialectInterface>(loadedDialect);
-    return success();
-  }
+  LogicalResult load(DialectReader &reader, MLIRContext *ctx);
 
   /// Return the loaded dialect, or nullptr if the dialect is unknown. This can
   /// only be called after `load`.
   Dialect *getLoadedDialect() const {
     assert(dialect &&
            "expected `load` to be invoked before `getLoadedDialect`");
     return *dialect;
   }
 
   /// The loaded dialect entry. This field is std::nullopt if we haven't
   /// attempted to load, nullptr if we failed to load, otherwise the loaded
   /// dialect.
   std::optional<Dialect *> dialect;
 
   /// The bytecode interface of the dialect, or nullptr if the dialect does not
   /// implement the bytecode interface. This field should only be checked if the
   /// `dialect` field is not std::nullopt.
   const BytecodeDialectInterface *interface = nullptr;
 
   /// The name of the dialect.
   StringRef name;
+
+  /// Handle for the dialect version parsed.
+  ArrayRef<uint8_t> versionHandle;
+  /// Lazy loaded dialect version from the handle above.
+  std::unique_ptr<DialectVersion> loadedVersion;
 };
 
 /// This struct represents an operation name entry within the bytecode.
 struct BytecodeOperationName {
   BytecodeOperationName(BytecodeDialect *dialect, StringRef name)
       : dialect(dialect), name(name) {}
 
   /// The loaded operation name, or std::nullopt if it hasn't been processed

 /// This class is used to read the resource section from the bytecode.
 class ResourceSectionReader {
 public:
   /// Initialize the resource section reader with the given section data.
   LogicalResult
   initialize(Location fileLoc, const ParserConfig &config,
              MutableArrayRef<BytecodeDialect> dialects,
              StringSectionReader &stringReader, ArrayRef<uint8_t> sectionData,
-             ArrayRef<uint8_t> offsetSectionData,
+             ArrayRef<uint8_t> offsetSectionData, DialectReader &dialectReader,
              const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef);
 
   /// Parse a dialect resource handle from the resource section.
   LogicalResult parseResourceHandle(EncodingReader &reader,
                                     AsmDialectResourceHandle &result) {
     return parseEntry(reader, dialectResources, result, "resource handle");
   }
 
parseResourceGroup(Location fileLoc, bool allowEmpty,
   }
   return success();
 }
 
 LogicalResult ResourceSectionReader::initialize(
     Location fileLoc, const ParserConfig &config,
     MutableArrayRef<BytecodeDialect> dialects,
     StringSectionReader &stringReader, ArrayRef<uint8_t> sectionData,
-    ArrayRef<uint8_t> offsetSectionData,
+    ArrayRef<uint8_t> offsetSectionData, DialectReader &dialectReader,
     const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef) {
   EncodingReader resourceReader(sectionData, fileLoc);
   EncodingReader offsetReader(offsetSectionData, fileLoc);
 
   // Read the number of external resource providers.
   uint64_t numExternalResourceGroups;
   if (failed(offsetReader.parseVarInt(numExternalResourceGroups)))
     return failure();
if (failed(parseGroup(handler)))
       return failure();
   }
 
   // Read the dialect resources from the bytecode.
   MLIRContext *ctx = fileLoc->getContext();
   while (!offsetReader.empty()) {
     BytecodeDialect *dialect;
     if (failed(parseEntry(offsetReader, dialects, dialect, "dialect")) ||
-        failed(dialect->load(resourceReader, ctx)))
+        failed(dialect->load(dialectReader, ctx)))
       return failure();
     Dialect *loadedDialect = dialect->getLoadedDialect();
     if (!loadedDialect) {
       return resourceReader.emitError()
              << "dialect '" << dialect->name << "' is unknown";
     }
     const auto *handler = dyn_cast<OpAsmDialectInterface>(loadedDialect);
     if (!handler) {
LogicalResult AttrTypeReader::parseAsmEntry(T &result, EncodingReader &reader,
   }
   return success();
 }
 
 template <typename T>
 LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
                                                EncodingReader &reader,
                                                StringRef entryType) {
-  if (failed(entry.dialect->load(reader, fileLoc.getContext())))
+  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.
-  DialectReader dialectReader(*this, stringReader, resourceReader, reader);
   if constexpr (std::is_same_v<T, Type>)
     entry.entry = entry.dialect->interface->readType(dialectReader);
   else
     entry.entry = entry.dialect->interface->readAttribute(dialectReader);
   return success(!!entry.entry);
 }
 
 //===----------------------------------------------------------------------===//
private:
   LogicalResult parseType(EncodingReader &reader, Type &result) {
     return attrTypeReader.parseType(reader, result);
   }
 
   //===--------------------------------------------------------------------===//
   // Resource Section
 
   LogicalResult
-  parseResourceSection(std::optional<ArrayRef<uint8_t>> resourceData,
+  parseResourceSection(EncodingReader &reader,
+                       std::optional<ArrayRef<uint8_t>> resourceData,
                        std::optional<ArrayRef<uint8_t>> resourceOffsetData);
 
   //===--------------------------------------------------------------------===//
   // IR Section
 
   /// This struct represents the current read state of a range of regions. This
   /// struct is used to enable iterative parsing of regions.
   struct RegionReadState {
if (failed(stringReader.initialize(
     return failure();
 
   // Process the dialect section.
   if (failed(parseDialectSection(*sectionDatas[bytecode::Section::kDialect])))
     return failure();
 
   // Process the resource section if present.
   if (failed(parseResourceSection(
-          sectionDatas[bytecode::Section::kResource],
+          reader, sectionDatas[bytecode::Section::kResource],
           sectionDatas[bytecode::Section::kResourceOffset])))
     return failure();
 
   // Process the attribute and type section.
   if (failed(attrTypeReader.initialize(
           dialects, *sectionDatas[bytecode::Section::kAttrType],
           *sectionDatas[bytecode::Section::kAttrTypeOffset])))
     return failure();
 
   // Finally, process the IR section.
   return parseIRSection(*sectionDatas[bytecode::Section::kIR], block);
 }
 
 LogicalResult BytecodeReader::parseVersion(EncodingReader &reader) {
   if (failed(reader.parseVarInt(version)))
     return failure();
 
   // Validate the bytecode version.
   uint64_t currentVersion = bytecode::kVersion;
-  if (version < currentVersion) {
+  uint64_t minSupportedVersion = bytecode::kMinSupportedVersion;
+  if (version < minSupportedVersion) {
     return reader.emitError("bytecode version ", version,
                             " is older than the current version of ",
                             currentVersion, ", and upgrade is not supported");
   }
   if (version > currentVersion) {
     return reader.emitError("bytecode version ", version,
                             " is newer than the current version ",
                             currentVersion);
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect Section
 
+LogicalResult BytecodeDialect::load(DialectReader &reader, MLIRContext *ctx) {
+  if (dialect)
+    return success();
+  Dialect *loadedDialect = ctx->getOrLoadDialect(name);
+  if (!loadedDialect && !ctx->allowsUnregisteredDialects()) {
+    return reader.emitError("dialect '")
+           << name
+           << "' is unknown. If this is intended, please call "
+              "allowUnregisteredDialects() on the MLIRContext, or use "
+              "-allow-unregistered-dialect with the MLIR tool used.";
+  }
+  dialect = loadedDialect;
+
+  // If the dialect was actually loaded, check to see if it has a bytecode
+  // interface.
+  if (loadedDialect)
+    interface = dyn_cast<BytecodeDialectInterface>(loadedDialect);
+  if (!versionHandle.empty()) {
+    if (!interface)
+      return reader.emitError("dialect '")
+             << name
+             << "' does not implement the bytecode interface, "
+                "but found a version entry";
+    loadedVersion = interface->parseVersion(reader, versionHandle);
+    if (!loadedVersion)
+      return failure();
+  }
+  return success();
+}
+
 LogicalResult
 BytecodeReader::parseDialectSection(ArrayRef<uint8_t> sectionData) {
   EncodingReader sectionReader(sectionData, fileLoc);
 
   // Parse the number of dialects in the section.
   uint64_t numDialects;
   if (failed(sectionReader.parseVarInt(numDialects)))
     return failure();
   dialects.resize(numDialects);
 
   // Parse each of the dialects.
-  for (uint64_t i = 0; i < numDialects; ++i)
+  for (uint64_t i = 0; i < numDialects; ++i) {
     if (failed(stringReader.parseString(sectionReader, dialects[i].name)))
       return failure();
+    /// Before version 1, there wasn't any versioning available for dialects.
+    if (version == 0)
+      continue;
+    bytecode::Section::ID sectionID;
+    if (failed(
+            sectionReader.parseSection(sectionID, dialects[i].versionHandle)))
+      return failure();
+    if (sectionID != bytecode::Section::kDialectVersions) {
+      emitError(fileLoc, "expected dialect version section");
+      return failure();
+    }
+  }
 
   // Parse the operation names, which are grouped by dialect.
   auto parseOpName = [&](BytecodeDialect *dialect) {
     StringRef opName;
     if (failed(stringReader.parseString(sectionReader, opName)))
       return failure();
     opNames.emplace_back(dialect, opName);
     return success();
   };
   while (!sectionReader.empty())
     if (failed(parseDialectGrouping(sectionReader, dialects, parseOpName)))
       return failure();
   return success();
 }
 
 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.
+  DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
+                              reader);
   if (!opName->opName) {
-    if (failed(opName->dialect->load(reader, getContext())))
+    if (failed(opName->dialect->load(dialectReader, getContext())))
       return failure();
     opName->opName.emplace((opName->dialect->name + "." + opName->name).str(),
                            getContext());
   }
   return *opName->opName;
 }
 
 //===----------------------------------------------------------------------===//
 // Resource Section
 
 LogicalResult BytecodeReader::parseResourceSection(
-    std::optional<ArrayRef<uint8_t>> resourceData,
+    EncodingReader &reader, std::optional<ArrayRef<uint8_t>> resourceData,
     std::optional<ArrayRef<uint8_t>> resourceOffsetData) {
   // Ensure both sections are either present or not.
   if (resourceData.has_value() != resourceOffsetData.has_value()) {
     if (resourceOffsetData)
       return emitError(fileLoc, "unexpected resource offset section when "
                                 "resource section is not present");
     return emitError(
         fileLoc,
         "expected resource offset section when resource section is present");
   }
 
   // If the resource sections are absent, there is nothing to do.
   if (!resourceData)
     return success();
 
   // Initialize the resource reader with the resource sections.
+  DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
+                              reader);
   return resourceReader.initialize(fileLoc, config, dialects, stringReader,
                                    *resourceData, *resourceOffsetData,
-                                   bufferOwnerRef);
+                                   dialectReader, bufferOwnerRef);
 }
 
 //===----------------------------------------------------------------------===//
 // IR Section
 
 LogicalResult BytecodeReader::parseIRSection(ArrayRef<uint8_t> sectionData,
                                              Block *block) {
   EncodingReader reader(sectionData, fileLoc);
LogicalResult BytecodeReader::parseIRSection(ArrayRef<uint8_t> sectionData,
   while (!regionStack.empty())
     if (failed(parseRegions(reader, regionStack, regionStack.back())))
       return failure();
   if (!forwardRefOps.empty()) {
     return reader.emitError(
         "not all forward unresolved forward operand references");
   }
 
+  // Resolve dialect version.
+  for (const BytecodeDialect &byteCodeDialect : dialects) {
+    // Parsing is complete, give an opportunity to each dialect to visit the
+    // IR and perform upgrades.
+    if (!byteCodeDialect.loadedVersion)
+      continue;
+    if (byteCodeDialect.interface &&
+        failed(byteCodeDialect.interface->upgradeFromVersion(
+            *moduleOp, *byteCodeDialect.loadedVersion)))
+      return failure();
+  }
+
   // Verify that the parsed operations are valid.
   if (config.shouldVerifyAfterParse() && failed(verify(*moduleOp)))
     return failure();
 
   // Splice the parsed operations over to the provided top-level block.
   auto &parsedOps = moduleOp->getBody()->getOperations();
   auto &destOps = block->getOperations();
   destOps.splice(destOps.end(), parsedOps, parsedOps.begin(), parsedOps.end());

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

private:
   /// An up-to-date total size of all of the buffers within `prevResultList`.
   /// This enables O(1) size checks of the current encoding.
   size_t prevResultSize = 0;
 
   /// The highest required alignment for the start of this section.
   unsigned requiredAlignment = 1;
 };
 
+//===----------------------------------------------------------------------===//
+// StringSectionBuilder
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class is used to simplify the process of emitting the string section.
+class StringSectionBuilder {
+public:
+  /// Add the given string to the string section, and return the index of the
+  /// string within the section.
+  size_t insert(StringRef str) {
+    auto it = strings.insert({llvm::CachedHashStringRef(str), strings.size()});
+    return it.first->second;
+  }
+
+  /// Write the current set of strings to the given emitter.
+  void write(EncodingEmitter &emitter) {
+    emitter.emitVarInt(strings.size());
+
+    // Emit the sizes in reverse order, so that we don't need to backpatch an
+    // offset to the string data or have a separate section.
+    for (const auto &it : llvm::reverse(strings))
+      emitter.emitVarInt(it.first.size() + 1);
+    // Emit the string data itself.
+    for (const auto &it : strings)
+      emitter.emitNulTerminatedString(it.first.val());
+  }
+
+private:
+  /// A set of strings referenced within the bytecode. The value of the map is
+  /// unused.
+  llvm::MapVector<llvm::CachedHashStringRef, size_t> strings;
+};
+} // namespace
+
+class DialectWriter : public DialectBytecodeWriter {
+public:
+  DialectWriter(EncodingEmitter &emitter, IRNumberingState &numberingState,
+                StringSectionBuilder &stringSection)
+      : emitter(emitter), numberingState(numberingState),
+        stringSection(stringSection) {}
+
+  //===--------------------------------------------------------------------===//
+  // IR
+  //===--------------------------------------------------------------------===//
+
+  void writeAttribute(Attribute attr) override {
+    emitter.emitVarInt(numberingState.getNumber(attr));
+  }
+  void writeType(Type type) override {
+    emitter.emitVarInt(numberingState.getNumber(type));
+  }
+
+  void writeResourceHandle(const AsmDialectResourceHandle &resource) override {
+    emitter.emitVarInt(numberingState.getNumber(resource));
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Primitives
+  //===--------------------------------------------------------------------===//
+
+  void writeVarInt(uint64_t value) override { emitter.emitVarInt(value); }
+
+  void writeSignedVarInt(int64_t value) override {
+    emitter.emitSignedVarInt(value);
+  }
+
+  void writeAPIntWithKnownWidth(const APInt &value) override {
+    size_t bitWidth = value.getBitWidth();
+
+    // If the value is a single byte, just emit it directly without going
+    // through a varint.
+    if (bitWidth <= 8)
+      return emitter.emitByte(value.getLimitedValue());
+
+    // If the value fits within a single varint, emit it directly.
+    if (bitWidth <= 64)
+      return emitter.emitSignedVarInt(value.getLimitedValue());
+
+    // Otherwise, we need to encode a variable number of active words. We use
+    // active words instead of the number of total words under the observation
+    // that smaller values will be more common.
+    unsigned numActiveWords = value.getActiveWords();
+    emitter.emitVarInt(numActiveWords);
+
+    const uint64_t *rawValueData = value.getRawData();
+    for (unsigned i = 0; i < numActiveWords; ++i)
+      emitter.emitSignedVarInt(rawValueData[i]);
+  }
+
+  void writeAPFloatWithKnownSemantics(const APFloat &value) override {
+    writeAPIntWithKnownWidth(value.bitcastToAPInt());
+  }
+
+  void writeOwnedString(StringRef str) override {
+    emitter.emitVarInt(stringSection.insert(str));
+  }
+
+  void writeOwnedBlob(ArrayRef<char> blob) override {
+    emitter.emitVarInt(blob.size());
+    emitter.emitOwnedBlob(ArrayRef<uint8_t>(
+        reinterpret_cast<const uint8_t *>(blob.data()), blob.size()));
+  }
+
+private:
+  EncodingEmitter &emitter;
+  IRNumberingState &numberingState;
+  StringSectionBuilder &stringSection;
+};
+
 /// A simple raw_ostream wrapper around a EncodingEmitter. This removes the need
 /// to go through an intermediate buffer when interacting with code that wants a
 /// raw_ostream.
 class RawEmitterOstream : public raw_ostream {
 public:
   explicit RawEmitterOstream(EncodingEmitter &emitter) : emitter(emitter) {
     SetUnbuffered();
   }
void EncodingEmitter::emitMultiByteVarInt(uint64_t value) {
 
   // If the value is too large to encode in a single byte, emit a special all
   // zero marker byte and splat the value directly.
   emitByte(0);
   emitBytes({reinterpret_cast<uint8_t *>(&value), sizeof(value)});
 }
 
 //===----------------------------------------------------------------------===//
-// StringSectionBuilder
-//===----------------------------------------------------------------------===//
-
-namespace {
-/// This class is used to simplify the process of emitting the string section.
-class StringSectionBuilder {
-public:
-  /// Add the given string to the string section, and return the index of the
-  /// string within the section.
-  size_t insert(StringRef str) {
-    auto it = strings.insert({llvm::CachedHashStringRef(str), strings.size()});
-    return it.first->second;
-  }
-
-  /// Write the current set of strings to the given emitter.
-  void write(EncodingEmitter &emitter) {
-    emitter.emitVarInt(strings.size());
-
-    // Emit the sizes in reverse order, so that we don't need to backpatch an
-    // offset to the string data or have a separate section.
-    for (const auto &it : llvm::reverse(strings))
-      emitter.emitVarInt(it.first.size() + 1);
-    // Emit the string data itself.
-    for (const auto &it : strings)
-      emitter.emitNulTerminatedString(it.first.val());
-  }
-
-private:
-  /// A set of strings referenced within the bytecode. The value of the map is
-  /// unused.
-  llvm::MapVector<llvm::CachedHashStringRef, size_t> strings;
-};
-} // namespace
-
-//===----------------------------------------------------------------------===//
 // Bytecode Writer
 //===----------------------------------------------------------------------===//
 
 namespace {
 class BytecodeWriter {
 public:
   BytecodeWriter(Operation *op) : numberingState(op) {}
 
for (auto it = entries.begin(), e = entries.end(); it != e;) {
     // Emit the entries within the group.
     for (auto &entry : llvm::make_range(groupStart, it))
       callback(entry);
   }
 }
 
 void BytecodeWriter::writeDialectSection(EncodingEmitter &emitter) {
   EncodingEmitter dialectEmitter;
+  // The writer used when emitting using a custom bytecode encoding.
+  DialectWriter dialectWriter(dialectEmitter, numberingState, stringSection);
 
   // Emit the referenced dialects.
   auto dialects = numberingState.getDialects();
   dialectEmitter.emitVarInt(llvm::size(dialects));
-  for (DialectNumbering &dialect : dialects)
+  for (DialectNumbering &dialect : dialects) {
     dialectEmitter.emitVarInt(stringSection.insert(dialect.name));
+    EncodingEmitter versionEmitter;
+    if (dialect.interface) {
+      DialectWriter versionWriter(versionEmitter, numberingState,
+                                  stringSection);
+      dialect.interface->writeVersion(versionWriter);
+    }
+    dialectEmitter.emitSection(bytecode::Section::kDialectVersions,
+                               std::move(versionEmitter));
+  }
 
   // Emit the referenced operation names grouped by dialect.
   auto emitOpName = [&](OpNameNumbering &name) {
     dialectEmitter.emitVarInt(stringSection.insert(name.name.stripDialect()));
   };
   writeDialectGrouping(dialectEmitter, numberingState.getOpNames(), emitOpName);
 
   emitter.emitSection(bytecode::Section::kDialect, std::move(dialectEmitter));
 }
 
 //===----------------------------------------------------------------------===//
 // Attributes and Types
 
-namespace {
-class DialectWriter : public DialectBytecodeWriter {
-public:
-  DialectWriter(EncodingEmitter &emitter, IRNumberingState &numberingState,
-                StringSectionBuilder &stringSection)
-      : emitter(emitter), numberingState(numberingState),
-        stringSection(stringSection) {}
-
-  //===--------------------------------------------------------------------===//
-  // IR
-  //===--------------------------------------------------------------------===//
-
-  void writeAttribute(Attribute attr) override {
-    emitter.emitVarInt(numberingState.getNumber(attr));
-  }
-  void writeType(Type type) override {
-    emitter.emitVarInt(numberingState.getNumber(type));
-  }
-
-  void writeResourceHandle(const AsmDialectResourceHandle &resource) override {
-    emitter.emitVarInt(numberingState.getNumber(resource));
-  }
-
-  //===--------------------------------------------------------------------===//
-  // Primitives
-  //===--------------------------------------------------------------------===//
-
-  void writeVarInt(uint64_t value) override { emitter.emitVarInt(value); }
-
-  void writeSignedVarInt(int64_t value) override {
-    emitter.emitSignedVarInt(value);
-  }
-
-  void writeAPIntWithKnownWidth(const APInt &value) override {
-    size_t bitWidth = value.getBitWidth();
-
-    // If the value is a single byte, just emit it directly without going
-    // through a varint.
-    if (bitWidth <= 8)
-      return emitter.emitByte(value.getLimitedValue());
-
-    // If the value fits within a single varint, emit it directly.
-    if (bitWidth <= 64)
-      return emitter.emitSignedVarInt(value.getLimitedValue());
-
-    // Otherwise, we need to encode a variable number of active words. We use
-    // active words instead of the number of total words under the observation
-    // that smaller values will be more common.
-    unsigned numActiveWords = value.getActiveWords();
-    emitter.emitVarInt(numActiveWords);
-
-    const uint64_t *rawValueData = value.getRawData();
-    for (unsigned i = 0; i < numActiveWords; ++i)
-      emitter.emitSignedVarInt(rawValueData[i]);
-  }
-
-  void writeAPFloatWithKnownSemantics(const APFloat &value) override {
-    writeAPIntWithKnownWidth(value.bitcastToAPInt());
-  }
-
-  void writeOwnedString(StringRef str) override {
-    emitter.emitVarInt(stringSection.insert(str));
-  }
-
-  void writeOwnedBlob(ArrayRef<char> blob) override {
-    emitter.emitVarInt(blob.size());
-    emitter.emitOwnedBlob(ArrayRef<uint8_t>(
-        reinterpret_cast<const uint8_t *>(blob.data()), blob.size()));
-  }
-
-private:
-  EncodingEmitter &emitter;
-  IRNumberingState &numberingState;
-  StringSectionBuilder &stringSection;
-};
-} // namespace
-
 void BytecodeWriter::writeAttrTypeSection(EncodingEmitter &emitter) {
   EncodingEmitter attrTypeEmitter;
   EncodingEmitter offsetEmitter;
   offsetEmitter.emitVarInt(llvm::size(numberingState.getAttributes()));
   offsetEmitter.emitVarInt(llvm::size(numberingState.getTypes()));
 
   // A functor used to emit an attribute or type entry.
   uint64_t prevOffset = 0;

mlir/test/Bytecode/invalid/invalid-structure.mlir

 // This file contains various failure test cases related to the structure of
 // a bytecode file.
 
 // Bytecode currently does not support big-endian platforms
 // UNSUPPORTED: target=s390x-{{.*}}
 
 //===--------------------------------------------------------------------===//
 // Version
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-structure-version.mlirbc 2>&1 | FileCheck %s --check-prefix=VERSION
-// VERSION: bytecode version 127 is newer than the current version 0
+// VERSION: bytecode version 127 is newer than the current version
 
 //===--------------------------------------------------------------------===//
 // Producer
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-structure-producer.mlirbc 2>&1 | FileCheck %s --check-prefix=PRODUCER
 // PRODUCER: malformed null-terminated string, no null character found
 

mlir/test/Bytecode/versioning/versioned_op.mlir

+// This file contains various failure test cases related to the structure of
+// the dialect section.
+
+// Bytecode currently does not support big-endian platforms
+// UNSUPPORTED: target=s390x-{{.*}}
+
+//===--------------------------------------------------------------------===//
+// Test generic
+//===--------------------------------------------------------------------===//
+
+// COM: bytecode contains
+// COM: module {
+// COM:   version: 2.0
+// COM:   "test.versionedA"() {dims = 123 : i64, modifier = false} : () -> ()
+// COM: }
+// RUN: mlir-opt %S/versioned-op-2.0.mlirbc 2>&1 | FileCheck %s --check-prefix=CHECK1
+// CHECK1: "test.versionedA"() {dims = 123 : i64, modifier = false} : () -> ()
+
+//===--------------------------------------------------------------------===//
+// Test upgrade
+//===--------------------------------------------------------------------===//
+
+// COM: bytecode contains
+// COM: module {
+// COM:   version: 1.12
+// COM:   "test.versionedA"() {dimensions = 123 : i64} : () -> ()
+// COM: }
+// RUN: mlir-opt %S/versioned-op-1.12.mlirbc 2>&1 | FileCheck %s --check-prefix=CHECK2
+// CHECK2: "test.versionedA"() {dims = 123 : i64, modifier = false} : () -> ()
+
+//===--------------------------------------------------------------------===//
+// Test forbidden downgrade
+//===--------------------------------------------------------------------===//
+
+// COM: bytecode contains
+// COM: module {
+// COM:   version: 2.2
+// COM:   "test.versionedA"() {dims = 123 : i64, modifier = false} : () -> ()
+// COM: }
+// RUN: not mlir-opt %S/versioned-op-2.2.mlirbc 2>&1 | FileCheck %s --check-prefix=ERR_NEW_VERSION
+// ERR_NEW_VERSION: current test dialect version is 2.0, can't parse version: 2.2

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

 //===- TestDialect.cpp - MLIR Dialect for Testing -------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestDialect.h"
 #include "TestAttributes.h"
 #include "TestInterfaces.h"
 #include "TestTypes.h"
+#include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/ExtensibleDialect.h"
 #include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
-#include <optional>
+#include "llvm/Support/Endian.h"
 
+#include <_types/_uint64_t.h>
+#include <cstddef>
+#include <memory>
 #include <numeric>
+#include <optional>
 
 // Include this before the using namespace lines below to
 // test that we don't have namespace dependencies.
 #include "TestOpsDialect.cpp.inc"
 
 using namespace mlir;
 using namespace test;
 
 void test::registerTestDialect(DialectRegistry &registry) {
   registry.insert<TestDialect>();
 }
 
 //===----------------------------------------------------------------------===//
+// TestDialect version utilities
+//===----------------------------------------------------------------------===//
+
+struct TestDialectVersion : DialectVersion {
+  uint32_t major = 2;
+  uint32_t minor = 0;
+};
+
+//===----------------------------------------------------------------------===//
 // TestDialect Interfaces
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// Testing the correctness of some traits.
 static_assert(
     llvm::is_detected<OpTrait::has_implicit_terminator_t,
                       SingleBlockImplicitTerminatorOp>::value,
     "has_implicit_terminator_t does not match SingleBlockImplicitTerminatorOp");
 static_assert(OpTrait::hasSingleBlockImplicitTerminator<
                   SingleBlockImplicitTerminatorOp>::value,
               "hasSingleBlockImplicitTerminator does not match "
               "SingleBlockImplicitTerminatorOp");
 
 struct TestResourceBlobManagerInterface
     : public ResourceBlobManagerDialectInterfaceBase<
           TestDialectResourceBlobHandle> {
   using ResourceBlobManagerDialectInterfaceBase<
       TestDialectResourceBlobHandle>::ResourceBlobManagerDialectInterfaceBase;
 };
 
+// Test support for interacting with the Bytecode reader/writer.
+struct TestBytecodeDialectInterface : public BytecodeDialectInterface {
+  using BytecodeDialectInterface::BytecodeDialectInterface;
+  TestBytecodeDialectInterface(Dialect *dialect)
+      : BytecodeDialectInterface(dialect) {}
+
+  // Emit a specific version of the dialect.
+  void writeVersion(DialectBytecodeWriter &writer) const final {
+    llvm::SmallVector<uint8_t> bytes;
+    writer.writeVarInt(0); // major
+    writer.writeVarInt(2); // minor
+  }
+
+  std::unique_ptr<DialectVersion>
+  parseVersion(DialectBytecodeReader &reader,
+               ArrayRef<uint8_t> versionHandle) const final {
+    uint64_t major, minor;
+    if (failed(reader.readVarInt(major)) || failed(reader.readVarInt(minor)))
+      return nullptr;
+    auto version = std::make_unique<TestDialectVersion>();
+    version->major = major;
+    version->minor = minor;
+    return version;
+  }
+
+  LogicalResult upgradeFromVersion(Operation *topLevelOp,
+                                   const DialectVersion &version_) const final {
+    const auto &version = static_cast<const TestDialectVersion &>(version_);
+    if ((version.major == 2) && (version.minor == 0))
+      return success();
+    if (version.major > 2 || (version.major == 2 && version.minor > 0)) {
+      return topLevelOp->emitError()
+             << "current test dialect version is 2.0, can't parse version: "
+             << version.major << "." << version.minor;
+    }
+    // Prior version 2.0, the old op supported only a single attribute called
+    // "dimensions". We can perform the upgrade.
+    topLevelOp->walk([](TestVersionedOpA op) {
+      if (auto dims = op->getAttr("dimensions")) {
+        op->removeAttr("dimensions");
+        op->setAttr("dims", dims);
+      }
+      op->setAttr("modifier", BoolAttr::get(op->getContext(), false));
+    });
+
+    return success();
+  }
+};
+
 // Test support for interacting with the AsmPrinter.
 struct TestOpAsmInterface : public OpAsmDialectInterface {
   using OpAsmDialectInterface::OpAsmDialectInterface;
   TestOpAsmInterface(Dialect *dialect, TestResourceBlobManagerInterface &mgr)
       : OpAsmDialectInterface(dialect), blobManager(mgr) {}
 
   //===------------------------------------------------------------------===//
   // Aliases
#include "TestOps.cpp.inc"
   addOperations<ManualCppOpWithFold>();
   registerDynamicOp(getDynamicGenericOp(this));
   registerDynamicOp(getDynamicOneOperandTwoResultsOp(this));
   registerDynamicOp(getDynamicCustomParserPrinterOp(this));
 
   auto &blobInterface = addInterface<TestResourceBlobManagerInterface>();
   addInterface<TestOpAsmInterface>(blobInterface);
 
-  addInterfaces<TestDialectFoldInterface, TestInlinerInterface,
-                TestReductionPatternInterface>();
+  addInterfaces<TestBytecodeDialectInterface, TestDialectFoldInterface,
+                TestInlinerInterface, TestReductionPatternInterface>();
   allowUnknownOperations();
 
   // Instantiate our fallback op interface that we'll use on specific
   // unregistered op.
   fallbackEffectOpInterfaces = new TestOpEffectInterfaceFallback;
 }
 TestDialect::~TestDialect() {
   delete static_cast<TestOpEffectInterfaceFallback *>(
void TestOpWithRegionPattern::getCanonicalizationPatterns(
     RewritePatternSet &results, MLIRContext *context) {
   results.add<TestRemoveOpWithInnerOps>(context);
 }
 
 OpFoldResult TestOpWithRegionFold::fold(FoldAdaptor adaptor) {
   return getOperand();
 }
 
-OpFoldResult TestOpConstant::fold(FoldAdaptor adaptor) {
-  return getValue();
-}
+OpFoldResult TestOpConstant::fold(FoldAdaptor adaptor) { return getValue(); }
 
 LogicalResult TestOpWithVariadicResultsAndFolder::fold(
     FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) {
   for (Value input : this->getOperands()) {
     results.push_back(input);
   }
   return success();
 }

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

def TestCSEOfSingleBlockOp : TEST_Op<"cse_of_single_block_op",
   let results = (outs Variadic<AnyType>:$outputs);
   let regions = (region SizedRegion<1>:$region);
   let assemblyFormat = [{
     attr-dict `inputs` `(` $inputs `)`
     $region `:` type($inputs)  `->` type($outputs)
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Test Ops to upgrade base on the dialect versions
+//===----------------------------------------------------------------------===//
+
+def TestVersionedOpA : TEST_Op<"versionedA"> {
+  // A previous version of the dialect (let's say 1.*) supported an attribute
+  // named "dimensions":
+  // let arguments = (ins
+  //   AnyI64Attr:$dimensions
+  // );
+
+  // In the current version (2.0) "dimensions" was renamed to "dims", and a new
+  // boolean attribute "modifier" was added. The previous version of the op
+  // corresponds to "modifier=false". We support loading old IR through
+  // upgrading, see `upgradeFromVersion()` in `TestOpAsmInterface`.
+  let arguments = (ins
+   AnyI64Attr:$dims,
+   BoolAttr:$modifier
+  );
+}
+
 #endif // TEST_OPS