[mlir][bytecode] Return error instead of min version

[jpienaar]

Summary

Can't return a well-formed IR output while enabling version to be bumped
up during emission. Previously it would return min version but
potentially invalid IR which was confusing, instead make it return
error and abort immediately instead.

mlir/include/mlir-c/IR.h

 
 /// Same as mlirOperationPrint but accepts flags controlling the printing
 /// behavior.
 MLIR_CAPI_EXPORTED void mlirOperationPrintWithFlags(MlirOperation op,
                                                     MlirOpPrintingFlags flags,
                                                     MlirStringCallback callback,
                                                     void *userData);
 
-struct MlirBytecodeWriterResult {
-  int64_t minVersion;
-};
-typedef struct MlirBytecodeWriterResult MlirBytecodeWriterResult;
-
-inline static bool
-mlirBytecodeWriterResultGetMinVersion(MlirBytecodeWriterResult res) {
-  return res.minVersion;
-}
-
-/// Same as mlirOperationPrint but writing the bytecode format and returns the
-/// minimum bytecode version the consumer needs to support.
-MLIR_CAPI_EXPORTED MlirBytecodeWriterResult mlirOperationWriteBytecode(
-    MlirOperation op, MlirStringCallback callback, void *userData);
-
-/// Same as mlirOperationWriteBytecode but with writer config.
-MLIR_CAPI_EXPORTED MlirBytecodeWriterResult
-mlirOperationWriteBytecodeWithConfig(MlirOperation op,
-                                     MlirBytecodeWriterConfig config,
+/// Same as mlirOperationPrint but writing the bytecode format.
+MLIR_CAPI_EXPORTED void mlirOperationWriteBytecode(MlirOperation op,
                                                    MlirStringCallback callback,
                                                    void *userData);
 
+/// Same as mlirOperationWriteBytecode but with writer config and returns
+/// failure only if desired bytecode could not be honored.
+MLIR_CAPI_EXPORTED MlirLogicalResult mlirOperationWriteBytecodeWithConfig(
+    MlirOperation op, MlirBytecodeWriterConfig config,
+    MlirStringCallback callback, void *userData);
+
 /// Prints an operation to stderr.
 MLIR_CAPI_EXPORTED void mlirOperationDump(MlirOperation op);
 
 /// Verify the operation and return true if it passes, false if it fails.
 MLIR_CAPI_EXPORTED bool mlirOperationVerify(MlirOperation op);
 
 /// Moves the given operation immediately after the other operation in its
 /// parent block. The given operation may be owned by the caller or by its

mlir/include/mlir/Bytecode/BytecodeWriter.h

for (auto &printer : map.getPrinters())
       attachResourcePrinter(std::move(printer));
   }
 
 private:
   /// A pointer to allocated storage for the impl state.
   std::unique_ptr<Impl> impl;
 };
 
-/// Status of bytecode serialization.
-struct BytecodeWriterResult {
-  /// The minimum version of the reader required to read the serialized file.
-  int64_t minVersion;
-};
-
 //===----------------------------------------------------------------------===//
 // Entry Points
 //===----------------------------------------------------------------------===//
 
 /// Write the bytecode for the given operation to the provided output stream.
 /// For streams where it matters, the given stream should be in "binary" mode.
-BytecodeWriterResult
-writeBytecodeToFile(Operation *op, raw_ostream &os,
+/// It only ever fails if setDesiredByteCodeVersion can't be honored.
+LogicalResult writeBytecodeToFile(Operation *op, raw_ostream &os,
                                   const BytecodeWriterConfig &config = {});
 
 } // namespace mlir
 
 #endif // MLIR_BYTECODE_BYTECODEWRITER_H

mlir/lib/Bindings/Python/IRCore.cpp

 #include "Globals.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/Debug.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
+#include "mlir-c/Support.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <optional>
 #include <utility>
 
 namespace py = pybind11;
 using namespace mlir;
if (assumeVerified)
     mlirOpPrintingFlagsAssumeVerified(flags);
 
   PyFileAccumulator accum(fileObject, binary);
   mlirOperationPrintWithFlags(operation, flags, accum.getCallback(),
                               accum.getUserData());
   mlirOpPrintingFlagsDestroy(flags);
 }
 
-MlirBytecodeWriterResult
-PyOperationBase::writeBytecode(const py::object &fileObject,
+void PyOperationBase::writeBytecode(const py::object &fileObject,
                                     std::optional<int64_t> bytecodeVersion) {
   PyOperation &operation = getOperation();
   operation.checkValid();
   PyFileAccumulator accum(fileObject, /*binary=*/true);
 
   if (!bytecodeVersion.has_value())
     return mlirOperationWriteBytecode(operation, accum.getCallback(),
                                       accum.getUserData());
 
   MlirBytecodeWriterConfig config = mlirBytecodeWriterConfigCreate();
   mlirBytecodeWriterConfigDesiredEmitVersion(config, *bytecodeVersion);
-  return mlirOperationWriteBytecodeWithConfig(
+  MlirLogicalResult res = mlirOperationWriteBytecodeWithConfig(
       operation, config, accum.getCallback(), accum.getUserData());
+  if (mlirLogicalResultIsFailure(res))
+    throw py::value_error((Twine("Unable to honor desired bytecode version ") +
+                           Twine(*bytecodeVersion))
+                              .str());
 }
 
 py::object PyOperationBase::getAsm(bool binary,
                                    std::optional<int64_t> largeElementsLimit,
                                    bool enableDebugInfo, bool prettyDebugInfo,
                                    bool printGenericOpForm, bool useLocalScope,
                                    bool assumeVerified) {
   py::object fileObject;
py::class_<PySymbolTable>(m, "SymbolTable", py::module_local())
                   py::arg("symbol"), py::arg("visibility"))
       .def_static("replace_all_symbol_uses",
                   &PySymbolTable::replaceAllSymbolUses, py::arg("old_symbol"),
                   py::arg("new_symbol"), py::arg("from_op"))
       .def_static("walk_symbol_tables", &PySymbolTable::walkSymbolTables,
                   py::arg("from_op"), py::arg("all_sym_uses_visible"),
                   py::arg("callback"));
 
-  py::class_<MlirBytecodeWriterResult>(m, "BytecodeResult", py::module_local())
-      .def("min_version",
-           [](MlirBytecodeWriterResult &res) { return res.minVersion; });
-
   // Container bindings.
   PyBlockArgumentList::bind(m);
   PyBlockIterator::bind(m);
   PyBlockList::bind(m);
   PyOperationIterator::bind(m);
   PyOperationList::bind(m);
   PyOpAttributeMap::bind(m);
   PyOpOperandIterator::bind(m);

mlir/lib/Bindings/Python/IRModule.h

void print(pybind11::object fileObject, bool binary,
              bool assumeVerified);
   pybind11::object getAsm(bool binary,
                           std::optional<int64_t> largeElementsLimit,
                           bool enableDebugInfo, bool prettyDebugInfo,
                           bool printGenericOpForm, bool useLocalScope,
                           bool assumeVerified);
 
   // Implement the bound 'writeBytecode' method.
-  MlirBytecodeWriterResult
-  writeBytecode(const pybind11::object &fileObject,
+  void writeBytecode(const pybind11::object &fileObject,
                      std::optional<int64_t> bytecodeVersion);
 
   /// Moves the operation before or after the other operation.
   void moveAfter(PyOperationBase &other);
   void moveBefore(PyOperationBase &other);
 
   /// Verify the operation. Throws `MLIRError` if verification fails, and
   /// returns `true` otherwise.
   bool verify();

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

void BytecodeWriter::writeStringSection(EncodingEmitter &emitter) {
   stringSection.write(stringEmitter);
   emitter.emitSection(bytecode::Section::kString, std::move(stringEmitter));
 }
 
 //===----------------------------------------------------------------------===//
 // Entry Points
 //===----------------------------------------------------------------------===//
 
-BytecodeWriterResult
-mlir::writeBytecodeToFile(Operation *op, raw_ostream &os,
+LogicalResult 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};
+  // Currently there is no failure case.
+  return success();
 }

mlir/lib/CAPI/IR/IR.cpp

 }
 
 void mlirOperationPrintWithFlags(MlirOperation op, MlirOpPrintingFlags flags,
                                  MlirStringCallback callback, void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(op)->print(stream, *unwrap(flags));
 }
 
-MlirBytecodeWriterResult mlirOperationWriteBytecode(MlirOperation op,
-                                                    MlirStringCallback callback,
-                                                    void *userData) {
+void mlirOperationWriteBytecode(MlirOperation op, MlirStringCallback callback,
   detail::CallbackOstream stream(callback, userData);
-  MlirBytecodeWriterResult res;
-  BytecodeWriterResult r = writeBytecodeToFile(unwrap(op), stream);
-  res.minVersion = r.minVersion;
-  return res;
+  // As no desired version is set, no failure can occur.
+  (void)writeBytecodeToFile(unwrap(op), stream);
 }
 
-MlirBytecodeWriterResult mlirOperationWriteBytecodeWithConfig(
+MlirLogicalResult mlirOperationWriteBytecodeWithConfig(
     MlirOperation op, MlirBytecodeWriterConfig config,
     MlirStringCallback callback, void *userData) {
   detail::CallbackOstream stream(callback, userData);
-  BytecodeWriterResult r =
-      writeBytecodeToFile(unwrap(op), stream, *unwrap(config));
-  MlirBytecodeWriterResult res;
-  res.minVersion = r.minVersion;
-  return res;
+  return wrap(writeBytecodeToFile(unwrap(op), stream, *unwrap(config)));
 }
 
 void mlirOperationDump(MlirOperation op) { return unwrap(op)->dump(); }
 
 bool mlirOperationVerify(MlirOperation op) {
   return succeeded(verify(unwrap(op)));
 }
 

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

MLIRDocument::convertToBytecode() {
   }
 
   lsp::MLIRConvertBytecodeResult result;
   {
     BytecodeWriterConfig writerConfig(fallbackResourceMap);
 
     std::string rawBytecodeBuffer;
     llvm::raw_string_ostream os(rawBytecodeBuffer);
-    writeBytecodeToFile(&parsedIR.front(), os, writerConfig);
+    // No desired bytecode version set, so no need to check for error.
+    (void)writeBytecodeToFile(&parsedIR.front(), os, writerConfig);
     result.output = llvm::encodeBase64(rawBytecodeBuffer);
   }
   return result;
 }
 
 //===----------------------------------------------------------------------===//
 // MLIRTextFileChunk
 //===----------------------------------------------------------------------===//

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

performActions(raw_ostream &os,
   // Run the pipeline.
   if (failed(pm.run(*op)))
     return failure();
 
   // Print the output.
   TimingScope outputTiming = timing.nest("Output");
   if (config.shouldEmitBytecode()) {
     BytecodeWriterConfig writerConfig(fallbackResourceMap);
-    if (auto v = config.bytecodeVersionToEmit()) {
+    if (auto v = config.bytecodeVersionToEmit())
       writerConfig.setDesiredBytecodeVersion(*v);
-      // Returns failure if requested version couldn't be used for opt tools.
-      return success(
-          writeBytecodeToFile(op.get(), os, writerConfig).minVersion <= *v);
-    }
-    writeBytecodeToFile(op.get(), os, writerConfig);
-    return success();
+    return writeBytecodeToFile(op.get(), os, writerConfig);
   }
 
   if (config.bytecodeVersionToEmit().has_value())
     return emitError(UnknownLoc::get(pm.getContext()))
            << "bytecode version while not emitting bytecode";
   AsmState asmState(op.get(), OpPrintingFlags(), /*locationMap=*/nullptr,
                     &fallbackResourceMap);
   op.get()->print(os, asmState);

mlir/test/python/ir/operation.py

def testOperationPrint():
   # CHECK: return %arg0 : i32
   module.operation.print(file=f)
   str_value = f.getvalue()
   print(str_value.__class__)
   print(f.getvalue())
 
   # Test roundtrip to bytecode.
   bytecode_stream = io.BytesIO()
-  result = module.operation.write_bytecode(bytecode_stream, desired_version=1)
-  assert result.min_version() == 1, "Requested version not serialized to"
+  module.operation.write_bytecode(bytecode_stream, desired_version=1)
   bytecode = bytecode_stream.getvalue()
   assert bytecode.startswith(b'ML\xefR'), "Expected bytecode to start with MLïR"
   module_roundtrip = Module.parse(bytecode, ctx)
   f = io.StringIO()
   module_roundtrip.operation.print(file=f)
   roundtrip_value = f.getvalue()
   assert str_value == roundtrip_value, "Mismatch after roundtrip bytecode"