[mlir][bytecode] Allow client to specify a desired version.

Add method to set a desired bytecode file format to generate. Change
write method to be able to return status including the minimum bytecode
version needed by reader. This enables generating an older version of
the bytecode (not dialect ops, attributes or types). But this does not
guarantee that an older version can always be generated, e.g., if a
dialect uses a new encoding only available at later bytecode version.
This clamps setting to at most current version.

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

2 files changed, 50 insertions, 21 deletions

diff --git a/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp b/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
index 8433f3d84b85..95729a2a4fa5 100644
--- a/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
+++ b/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
@@ -27,6 +27,10 @@ using namespace mlir::bytecode::detail;
 struct BytecodeWriterConfig::Impl {
   Impl(StringRef producer) : producer(producer) {}
 
+  /// Version to use when writing.
+  /// Note: This only differs from kVersion if a specific version is set.
+  int64_t bytecodeVersion = bytecode::kVersion;
+
   /// The producer of the bytecode.
   StringRef producer;
 
@@ -48,6 +52,12 @@ void BytecodeWriterConfig::attachResourcePrinter(
   impl->externalResourcePrinters.emplace_back(std::move(printer));
 }
 
+void BytecodeWriterConfig::setDesiredBytecodeVersion(int64_t bytecodeVersion) {
+  // Clamp to current version.
+  impl->bytecodeVersion =
+      std::min<int64_t>(bytecodeVersion, bytecode::kVersion);
+}
+
 //===----------------------------------------------------------------------===//
 // EncodingEmitter
 //===----------------------------------------------------------------------===//
@@ -295,7 +305,8 @@ private:
 
 class DialectWriter : public DialectBytecodeWriter {
 public:
-  DialectWriter(EncodingEmitter &emitter, IRNumberingState &numberingState,
+  DialectWriter(int64_t bytecodeVersion, EncodingEmitter &emitter,
+                IRNumberingState &numberingState,
                 StringSectionBuilder &stringSection)
       : emitter(emitter), numberingState(numberingState),
         stringSection(stringSection) {}
@@ -362,7 +373,10 @@ public:
         reinterpret_cast<const uint8_t *>(blob.data()), blob.size()));
   }
 
+  int64_t getBytecodeVersion() const override { return bytecodeVersion; }
+
 private:
+  int64_t bytecodeVersion;
   EncodingEmitter &emitter;
   IRNumberingState &numberingState;
   StringSectionBuilder &stringSection;
@@ -421,11 +435,11 @@ void EncodingEmitter::emitMultiByteVarInt(uint64_t value) {
 namespace {
 class BytecodeWriter {
 public:
-  BytecodeWriter(Operation *op) : numberingState(op) {}
+  BytecodeWriter(Operation *op, const BytecodeWriterConfig::Impl &config)
+      : numberingState(op), config(config) {}
 
   /// Write the bytecode for the given root operation.
-  void write(Operation *rootOp, raw_ostream &os,
-             const BytecodeWriterConfig::Impl &config);
+  void write(Operation *rootOp, raw_ostream &os);
 
 private:
   //===--------------------------------------------------------------------===//
@@ -449,8 +463,7 @@ private:
   //===--------------------------------------------------------------------===//
   // Resources
 
-  void writeResourceSection(Operation *op, EncodingEmitter &emitter,
-                            const BytecodeWriterConfig::Impl &config);
+  void writeResourceSection(Operation *op, EncodingEmitter &emitter);
 
   //===--------------------------------------------------------------------===//
   // Strings
@@ -465,11 +478,13 @@ private:
 
   /// The IR numbering state generated for the root operation.
   IRNumberingState numberingState;
+
+  /// Configuration dictating bytecode emission.
+  const BytecodeWriterConfig::Impl &config;
 };
 } // namespace
 
-void BytecodeWriter::write(Operation *rootOp, raw_ostream &os,
-                           const BytecodeWriterConfig::Impl &config) {
+void BytecodeWriter::write(Operation *rootOp, raw_ostream &os) {
   EncodingEmitter emitter;
 
   // Emit the bytecode file header. This is how we identify the output as a
@@ -477,7 +492,7 @@ void BytecodeWriter::write(Operation *rootOp, raw_ostream &os,
   emitter.emitString("ML\xefR");
 
   // Emit the bytecode version.
-  emitter.emitVarInt(bytecode::kVersion);
+  emitter.emitVarInt(config.bytecodeVersion);
 
   // Emit the producer.
   emitter.emitNulTerminatedString(config.producer);
@@ -492,7 +507,7 @@ void BytecodeWriter::write(Operation *rootOp, raw_ostream &os,
   writeIRSection(emitter, rootOp);
 
   // Emit the resources section.
-  writeResourceSection(rootOp, emitter, config);
+  writeResourceSection(rootOp, emitter);
 
   // Emit the string section.
   writeStringSection(emitter);
@@ -540,12 +555,17 @@ void BytecodeWriter::writeDialectSection(EncodingEmitter &emitter) {
     // Write the string section and get the ID.
     size_t nameID = stringSection.insert(dialect.name);
 
+    if (config.bytecodeVersion == 0) {
+      dialectEmitter.emitVarInt(nameID);
+      continue;
+    }
+
     // Try writing the version to the versionEmitter.
     EncodingEmitter versionEmitter;
     if (dialect.interface) {
       // The writer used when emitting using a custom bytecode encoding.
-      DialectWriter versionWriter(versionEmitter, numberingState,
-                                  stringSection);
+      DialectWriter versionWriter(config.bytecodeVersion, versionEmitter,
+                                  numberingState, stringSection);
       dialect.interface->writeVersion(versionWriter);
     }
 
@@ -586,8 +606,8 @@ void BytecodeWriter::writeAttrTypeSection(EncodingEmitter &emitter) {
     bool hasCustomEncoding = false;
     if (const BytecodeDialectInterface *interface = entry.dialect->interface) {
       // The writer used when emitting using a custom bytecode encoding.
-      DialectWriter dialectWriter(attrTypeEmitter, numberingState,
-                                  stringSection);
+      DialectWriter dialectWriter(config.bytecodeVersion, attrTypeEmitter,
+                                  numberingState, stringSection);
 
       if constexpr (std::is_same_v<std::decay_t<decltype(entryValue)>, Type>) {
         // TODO: We don't currently support custom encoded mutable types.
@@ -787,9 +807,8 @@ private:
 };
 } // namespace
 
-void BytecodeWriter::writeResourceSection(
-    Operation *op, EncodingEmitter &emitter,
-    const BytecodeWriterConfig::Impl &config) {
+void BytecodeWriter::writeResourceSection(Operation *op,
+                                          EncodingEmitter &emitter) {
   EncodingEmitter resourceEmitter;
   EncodingEmitter resourceOffsetEmitter;
   uint64_t prevOffset = 0;
@@ -868,8 +887,12 @@ void BytecodeWriter::writeStringSection(EncodingEmitter &emitter) {
 // Entry Points
 //===----------------------------------------------------------------------===//
 
-void mlir::writeBytecodeToFile(Operation *op, raw_ostream &os,
-                               const BytecodeWriterConfig &config) {
-  BytecodeWriter writer(op);
-  writer.write(op, os, config.getImpl());
+BytecodeWriterResult
+mlir::writeBytecodeToFile(Operation *op, raw_ostream &os,
+                          const BytecodeWriterConfig &config) {
+  BytecodeWriter writer(op, config.getImpl());
+  writer.write(op, os);
+  // Return the bytecode version emitted - currently there is no additional
+  // feedback as to minimum beyond the requested one.
+  return {config.getImpl().bytecodeVersion};
 }
diff --git a/mlir/lib/Bytecode/Writer/IRNumbering.cpp b/mlir/lib/Bytecode/Writer/IRNumbering.cpp
index 06778b9c7c75..7f56e9a94d29 100644
--- a/mlir/lib/Bytecode/Writer/IRNumbering.cpp
+++ b/mlir/lib/Bytecode/Writer/IRNumbering.cpp
@@ -7,11 +7,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "IRNumbering.h"
+#include "../Encoding.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
+#include "llvm/Support/ErrorHandling.h"
 
 using namespace mlir;
 using namespace mlir::bytecode::detail;
@@ -41,6 +43,10 @@ struct IRNumberingState::NumberingDialectWriter : public DialectBytecodeWriter {
   }
   void writeOwnedBlob(ArrayRef<char> blob) override {}
 
+  int64_t getBytecodeVersion() const override {
+    llvm_unreachable("unexpected querying of version in IRNumbering");
+  }
+
   /// The parent numbering state that is populated by this writer.
   IRNumberingState &state;
 };