Adds __str__ support to python mlir.ir.MlirModule.

[stellaraccident]

Summary

• Also raises an exception on parse error.
• Removes placeholder smoketest.
• Adds docstrings.

[ftynse]

mlir/lib/Bindings/Python/IRModules.cpp

19

Please fix linter warnings. MLIR uses camelCase for variables, I've seen kConstantName for constants.

50

What does the + do?

[stellaraccident]

Fix lint and remove lambda decay.

mlir/lib/Bindings/Python/IRModules.cpp

19

Thanks - I started doing that then second guessed thinking that was just Google style (too many style switches recently).

50

It is just the unary plus but the effect is subtle: It causes an explicit decay of a (non capturing) C++ closure to a raw function pointer.

Here is a good explanation: https://stackoverflow.com/questions/18889028/a-positive-lambda-what-sorcery-is-this

I just checked that it is not required in this case and removed it (there might be a compiler version dependency: I remember that on older compilers we had to use this idiom frequently when bridging C/C++ callbacks on prior projects).

mlir/lib/Bindings/Python/CMakeLists.txt

 # The actual extension library produces a shared-object or DLL and has
 # sources that must be compiled in accordance with pybind11 needs (RTTI and
 # exceptions).
 # TODO: Link the libraries separately once a helper function is available
 # to more generically add a pybind11 compliant library.
 add_library(MLIRBindingsPythonExtension ${PYEXT_LINK_MODE}
   MainModule.cpp
   IRModules.cpp
+  PybindUtils.cpp
 )
 
 target_include_directories(MLIRBindingsPythonExtension PRIVATE
   "${PYTHON_INCLUDE_DIRS}"
   "${pybind11_INCLUDE_DIRS}")
 
 # The extension itself must be compiled with RTTI and exceptions enabled.
 # Also, some warning classes triggered by pybind11 are disabled.

 if(NOT MSVC AND NOT APPLE)
   set_target_properties(MLIRBindingsPythonExtension
     PROPERTIES LINK_FLAGS
     "-Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/unix_version.lds")
 endif()
 
 target_link_libraries(MLIRBindingsPythonExtension
   PRIVATE
-  MLIRIR
   MLIRCAPIIR
   MLIRCAPIRegistration
   ${PYEXT_LIBADD}
 )

mlir/lib/Bindings/Python/IRModules.h

 //===- IRModules.h - IR Submodules of pybind module -----------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_BINDINGS_PYTHON_IRMODULES_H
 #define MLIR_BINDINGS_PYTHON_IRMODULES_H
 
 #include <pybind11/pybind11.h>
 
 #include "mlir-c/IR.h"
 
-namespace py = pybind11;
+namespace mlir {
+namespace python {
 
 class PyMlirContext;
 class PyMlirModule;
 
 /// Wrapper around MlirContext.
 class PyMlirContext {
 public:
   PyMlirContext() { context = mlirContextCreate(); }
if (module.ptr)
       mlirModuleDestroy(module);
   }
   /// Dumps the module.
   void dump();
 
   MlirModule module;
 };
 
-void populateIRSubmodule(py::module &m);
+void populateIRSubmodule(pybind11::module &m);
+
+} // namespace python
+} // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_IRMODULES_H

mlir/lib/Bindings/Python/IRModules.cpp

 //===- IRModules.cpp - IR Submodules of pybind module ---------------------===//
 //
 // 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 "IRModules.h"
+#include "PybindUtils.h"
+
+namespace py = pybind11;
+using namespace mlir::python;
+
+//------------------------------------------------------------------------------
+// Docstrings (trivial, non-duplicated docstrings are included inline).
+//------------------------------------------------------------------------------
+
+static const char kContextParseDocstring[] =
+    R"(Parses a module's assembly format from a string.
+
+Returns a new MlirModule or raises a ValueError if the parsing fails.
+)";
+
+static const char kOperationStrDunderDocstring[] =
+    R"(Prints 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 method, which supports keyword arguments to customize
+behavior.
+)";
+
+static const char kDumpDocstring[] =
+    R"(Dumps a debug representation of the object to stderr.)";
+
+//------------------------------------------------------------------------------
+// Conversion utilities.
+//------------------------------------------------------------------------------
+
+namespace {
+
+/// Accumulates into a python string from a method that accepts an
+/// MlirPrintCallback.
+struct PyPrintAccumulator {
+  py::list parts;
+
+  void *getUserData() { return this; }
+
+  MlirPrintCallback getCallback() {
+    return [](const char *part, intptr_t size, void *userData) {
+      PyPrintAccumulator *printAccum =
+          static_cast<PyPrintAccumulator *>(userData);
+      py::str pyPart(part, size); // Decodes as UTF-8 by default.
+      printAccum->parts.append(std::move(pyPart));
+    };
+  }
+
+  py::str join() {
+    py::str delim("", 0);
+    return delim.attr("join")(parts);
+  }
+};
+
+} // namespace
 
 //------------------------------------------------------------------------------
 // Context Wrapper Class.
 //------------------------------------------------------------------------------
 
 PyMlirModule PyMlirContext::parse(const std::string &module) {
   auto moduleRef = mlirModuleCreateParse(context, module.c_str());
+  if (!moduleRef.ptr) {
+    throw SetPyError(PyExc_ValueError,
+                     "Unable to parse module assembly (see diagnostics)");
+  }
   return PyMlirModule(moduleRef);
 }
 
 //------------------------------------------------------------------------------
 // Module Wrapper Class.
 //------------------------------------------------------------------------------
 
 void PyMlirModule::dump() { mlirOperationDump(mlirModuleGetOperation(module)); }
 
 //------------------------------------------------------------------------------
 // Populates the pybind11 IR submodule.
 //------------------------------------------------------------------------------
 
-void populateIRSubmodule(py::module &m) {
+void mlir::python::populateIRSubmodule(py::module &m) {
   py::class_<PyMlirContext>(m, "MlirContext")
       .def(py::init<>())
-      .def("parse", &PyMlirContext::parse, py::keep_alive<0, 1>());
+      .def("parse", &PyMlirContext::parse, py::keep_alive<0, 1>(),
+           kContextParseDocstring);
 
-  py::class_<PyMlirModule>(m, "MlirModule").def("dump", &PyMlirModule::dump);
+  py::class_<PyMlirModule>(m, "MlirModule")
+      .def("dump", &PyMlirModule::dump, kDumpDocstring)
+      .def(
+          "__str__",
+          [](PyMlirModule &self) {
+            auto operation = mlirModuleGetOperation(self.module);
+            PyPrintAccumulator printAccum;
+            mlirOperationPrint(operation, printAccum.getCallback(),
+                               printAccum.getUserData());
+            return printAccum.join();
+          },
+          kOperationStrDunderDocstring);
 }

mlir/lib/Bindings/Python/MainModule.cpp

 //===- MainModule.cpp - Main pybind module --------------------------------===//
 //
 // 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 <tuple>
 
 #include <pybind11/pybind11.h>
 
 #include "IRModules.h"
-#include "mlir/IR/MLIRContext.h"
 
+namespace py = pybind11;
 using namespace mlir;
+using namespace mlir::python;
 
 PYBIND11_MODULE(_mlir, m) {
   m.doc() = "MLIR Python Native Extension";
 
-  m.def("get_test_value", []() {
-    // This is just calling a method on the MLIRContext as a smoketest
-    // for linkage.
-    MLIRContext context;
-    return std::make_tuple(std::string("From the native module"),
-                           context.isMultithreadingEnabled());
-  });
-
   // Define and populate IR submodule.
   auto irModule = m.def_submodule("ir", "MLIR IR Bindings");
   populateIRSubmodule(irModule);
 }

mlir/lib/Bindings/Python/PybindUtils.h

+//===- PybindUtils.h - Utilities for interop with pybind11 ------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_BINDINGS_PYTHON_PYBINDUTILS_H
+#define MLIR_BINDINGS_PYTHON_PYBINDUTILS_H
+
+#include <pybind11/pybind11.h>
+
+#include "llvm/ADT/Twine.h"
+
+namespace mlir {
+namespace python {
+
+// Sets a python error, ready to be thrown to return control back to the
+// python runtime.
+// Correct usage:
+//   throw SetPyError(PyExc_ValueError, "Foobar'd");
+pybind11::error_already_set SetPyError(PyObject *excClass, llvm::Twine message);
+
+} // namespace python
+} // namespace mlir
+
+#endif // MLIR_BINDINGS_PYTHON_PYBINDUTILS_H

mlir/lib/Bindings/Python/PybindUtils.cpp

+//===- PybindUtils.cpp - Utilities for interop with pybind11 --------------===//
+//
+// 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 "PybindUtils.h"
+
+namespace py = pybind11;
+
+pybind11::error_already_set mlir::python::SetPyError(PyObject *excClass,
+                                                     llvm::Twine message) {
+  auto messageStr = message.str();
+  PyErr_SetString(excClass, messageStr.c_str());
+  return pybind11::error_already_set();
+}

mlir/lib/CAPI/IR/IR.cpp

 /* ========================================================================== */
 
 MlirModule mlirModuleCreateEmpty(MlirLocation location) {
   return wrap(ModuleOp::create(unwrap(location)));
 }
 
 MlirModule mlirModuleCreateParse(MlirContext context, const char *module) {
   OwningModuleRef owning = parseSourceString(module, unwrap(context));
+  if (!owning)
+    return MlirModule{nullptr};
   return MlirModule{owning.release().getOperation()};
 }
 
 void mlirModuleDestroy(MlirModule module) {
   // Transfer ownership to an OwningModuleRef so that its destructor is called.
   OwningModuleRef(unwrap(module));
 }
 

mlir/test/Bindings/Python/ir_module_test.py

+# RUN: %PYTHON %s | FileCheck %s
+
+import mlir
+
+def run(f):
+  print("TEST:", f.__name__)
+  f()
+
+# Verify successful parse.
+# CHECK-LABEL: TEST: testParseSuccess
+# CHECK: module @successfulParse
+def testParseSuccess():
+  ctx = mlir.ir.MlirContext()
+  module = ctx.parse(r"""module @successfulParse {}""")
+  module.dump()  # Just outputs to stderr. Verifies that it functions.
+  print(str(module))
+
+run(testParseSuccess)
+
+
+# Verify parse error.
+# CHECK-LABEL: TEST: testParseError
+# CHECK: testParseError: Unable to parse module assembly (see diagnostics)
+def testParseError():
+  ctx = mlir.ir.MlirContext()
+  try:
+    module = ctx.parse(r"""}SYNTAX ERROR{""")
+  except ValueError as e:
+    print("testParseError:", e)
+  else:
+    print("Exception not produced")
+
+run(testParseError)
+
+
+# Verify round-trip of ASM that contains unicode.
+# Note that this does not test that the print path converts unicode properly
+# because MLIR asm always normalizes it to the hex encoding.
+# CHECK-LABEL: TEST: testRoundtripUnicode
+# CHECK: func @roundtripUnicode()
+# CHECK: foo = "\F0\9F\98\8A"
+def testRoundtripUnicode():
+  ctx = mlir.ir.MlirContext()
+  module = ctx.parse(r"""
+    func @roundtripUnicode() attributes { foo = "😊" }
+  """)
+  print(str(module))
+
+run(testRoundtripUnicode)

mlir/test/Bindings/Python/ir_test.py

-# RUN: %PYTHON %s | FileCheck %s
-
-import mlir
-
-TEST_MLIR_ASM = r"""
-module {
-}
-"""
-
-ctx = mlir.ir.MlirContext()
-module = ctx.parse(TEST_MLIR_ASM)
-module.dump()
-print(bool(module))
-# CHECK: True

mlir/test/Bindings/Python/smoke_test.py

-# RUN: %PYTHON %s | FileCheck %s
-
-import mlir
-
-# CHECK: From the native module
-print(mlir.get_test_value())