Plumb write_bytecode to the Python API

[mehdi_amini]

Summary

This adds a write_bytecode method to the Operation class.
The method takes a file handle and writes the binary blob to it.

[mehdi_amini]

Add missing CMake dependency

[ftynse]

mlir/include/mlir-c/IR.h

525–526

Nit: the format feels off.

[mehdi_amini]

Clang-format and rebase

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);
 
+/// Same as mlirOperationPrint but writing the bytecode format out.
+MLIR_CAPI_EXPORTED void mlirOperationWriteBytecode(MlirOperation op,
+                                 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/lib/Bindings/Python/IRCore.cpp

binary: Whether to return a bytes (True) or str (False) object. Defaults to
     False.
   ... others ...: See the print() method for common keyword arguments for
     configuring the printout.
 Returns:
   Either a bytes or str object, depending on the setting of the 'binary'
   argument.
 )";
 
+static const char kOperationPrintBytecodeDocstring[] =
+R"(Write the bytecode form of the operation to a file like object.
+
+Args:
+  file: The file like object to write to.
+)";
+
 static const char kOperationStrDunderDocstring[] =
     R"(Gets the assembly form of the operation with default options.
 
 If more advanced control over the assembly formatting or I/O options is needed,
 use the dedicated print or get_asm method, which supports keyword arguments to
 customize behavior.
 )";
 
if (useLocalScope)
     mlirOpPrintingFlagsUseLocalScope(flags);
 
   PyFileAccumulator accum(fileObject, binary);
   mlirOperationPrintWithFlags(operation, flags, accum.getCallback(),
                               accum.getUserData());
   mlirOpPrintingFlagsDestroy(flags);
 }
 
+void PyOperationBase::writeBytecode(py::object fileObject) {
+  PyOperation &operation = getOperation();
+  operation.checkValid();
+  PyFileAccumulator accum(fileObject, /*binary=*/true);
+  mlirOperationWriteBytecode(operation, accum.getCallback(),
+                              accum.getUserData());
+}
+
 py::object PyOperationBase::getAsm(bool binary,
                                    llvm::Optional<int64_t> largeElementsLimit,
                                    bool enableDebugInfo, bool prettyDebugInfo,
                                    bool printGenericOpForm, bool useLocalScope,
                                    bool assumeVerified) {
   py::object fileObject;
   if (binary) {
     fileObject = py::module::import("io").attr("BytesIO")();
py::class_<PyOperationBase>(m, "_OperationBase", py::module_local())
            // Careful: Lots of arguments must match up with print method.
            py::arg("file") = py::none(), py::arg("binary") = false,
            py::arg("large_elements_limit") = py::none(),
            py::arg("enable_debug_info") = false,
            py::arg("pretty_debug_info") = false,
            py::arg("print_generic_op_form") = false,
            py::arg("use_local_scope") = false,
            py::arg("assume_verified") = false, kOperationPrintDocstring)
+      .def("write_bytecode", &PyOperationBase::writeBytecode,
+           py::arg("file"),
+           kOperationPrintBytecodeDocstring)
       .def("get_asm", &PyOperationBase::getAsm,
            // Careful: Lots of arguments must match up with get_asm method.
            py::arg("binary") = false,
            py::arg("large_elements_limit") = py::none(),
            py::arg("enable_debug_info") = false,
            py::arg("pretty_debug_info") = false,
            py::arg("print_generic_op_form") = false,
            py::arg("use_local_scope") = false,

mlir/lib/Bindings/Python/IRModule.h

void print(pybind11::object fileObject, bool binary,
              bool prettyDebugInfo, bool printGenericOpForm, bool useLocalScope,
              bool assumeVerified);
   pybind11::object getAsm(bool binary,
                           llvm::Optional<int64_t> largeElementsLimit,
                           bool enableDebugInfo, bool prettyDebugInfo,
                           bool printGenericOpForm, bool useLocalScope,
                           bool assumeVerified);
 
+  // Implement the bound 'writeBytecode' method.
+  void writeBytecode(pybind11::object fileObject);
+
   /// Moves the operation before or after the other operation.
   void moveAfter(PyOperationBase &other);
   void moveBefore(PyOperationBase &other);
 
   /// Each must provide access to the raw Operation.
   virtual PyOperation &getOperation() = 0;
 };
 

mlir/lib/CAPI/IR/IR.cpp

 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Parser/Parser.h"
+#include "mlir/Bytecode/BytecodeWriter.h"
 
-#include "llvm/Support/Debug.h"
 #include <cstddef>
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Context API.
 //===----------------------------------------------------------------------===//
 

 }
 
 void mlirOperationPrintWithFlags(MlirOperation op, MlirOpPrintingFlags flags,
                                  MlirStringCallback callback, void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(op)->print(stream, *unwrap(flags));
 }
 
+void mlirOperationWriteBytecode(MlirOperation op,
+                                 MlirStringCallback callback, void *userData) {
+  detail::CallbackOstream stream(callback, userData);
+  writeBytecodeToFile(unwrap(op), stream);
+}
+
 void mlirOperationDump(MlirOperation op) { return unwrap(op)->dump(); }
 
 bool mlirOperationVerify(MlirOperation op) {
   return succeeded(verify(unwrap(op)));
 }
 
 void mlirOperationMoveAfter(MlirOperation op, MlirOperation other) {
   return unwrap(op)->moveAfter(unwrap(other));

mlir/test/python/ir/operation.py

def testOperationPrint():
   f = io.StringIO()
   # CHECK: <class 'str'>
   # 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()
+  module.operation.write_bytecode(bytecode_stream)
+  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"
+
+
   # Test print to binary file.
   f = io.BytesIO()
   # CHECK: <class 'bytes'>
   # CHECK: return %arg0 : i32
   module.operation.print(file=f, binary=True)
   bytes_value = f.getvalue()
   print(bytes_value.__class__)
   print(bytes_value)

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

header_deps = [
         ":IR",
         ":Pass",
         ":Support",
         "//llvm:Support",
     ],
     includes = ["include"],
     deps = [
         ":AsmParser",
+        ":BytecodeWriter",
         ":ConversionPassIncGen",
         ":FuncDialect",
         ":InferTypeOpInterface",
         ":Parser",
     ],
 )
 
 mlir_c_api_cc_library(