[mlir][python] Add simple debugging and printing helpers

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

Author: nicolasvasilache

mlir/include/mlir-c/IR.h

@@ -75,6 +75,13 @@ struct MlirNamedAttribute {
 };
 typedef struct MlirNamedAttribute MlirNamedAttribute;
 
+//===----------------------------------------------------------------------===//
+// Global API.
+//===----------------------------------------------------------------------===//
+
+/// Set the global debugging flag.
+MLIR_CAPI_EXPORTED void mlirEnableGlobalDebug(bool enable);
+
 //===----------------------------------------------------------------------===//
 // Context API.
 //===----------------------------------------------------------------------===//
@@ -119,6 +126,10 @@ mlirContextGetNumLoadedDialects(MlirContext context);
 MLIR_CAPI_EXPORTED MlirDialect mlirContextGetOrLoadDialect(MlirContext context,
                                                            MlirStringRef name);
 
+/// Set threading mode (must be set to false to print-ir-after-all).
+MLIR_CAPI_EXPORTED void mlirContextEnableMultithreading(MlirContext context,
+                                                        bool enable);
+
 /// Returns whether the given fully-qualified operation (i.e.
 /// 'dialect.operation') is registered with the context. This will return true
 /// if the dialect is loaded and the operation is registered within the

mlir/include/mlir-c/Pass.h

@@ -65,6 +65,14 @@ mlirPassManagerGetAsOpPassManager(MlirPassManager passManager);
 MLIR_CAPI_EXPORTED MlirLogicalResult
 mlirPassManagerRun(MlirPassManager passManager, MlirModule module);
 
+/// Enable print-ir-after-all.
+MLIR_CAPI_EXPORTED void
+mlirPassManagerEnableIRPrinting(MlirPassManager passManager);
+
+/// Enable / disable verify-each.
+MLIR_CAPI_EXPORTED void
+mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable);
+
 /// Nest an OpPassManager under the top-level PassManager, the nested
 /// passmanager will only run on operations matching the provided name.
 /// The returned OpPassManager will be destroyed when the parent is destroyed.

mlir/lib/Bindings/Python/IRCore.cpp

@@ -1712,6 +1712,11 @@ class PyOpAttributeMap {
 //------------------------------------------------------------------------------
 
 void mlir::python::populateIRCore(py::module &m) {
+  //----------------------------------------------------------------------------
+  // Mapping of Global functions
+  //----------------------------------------------------------------------------
+  m.def("_enable_debug", [](bool enable) { mlirEnableGlobalDebug(enable); });
+
   //----------------------------------------------------------------------------
   // Mapping of MlirContext
   //----------------------------------------------------------------------------
@@ -1766,6 +1771,10 @@ void mlir::python::populateIRCore(py::module &m) {
           [](PyMlirContext &self, bool value) {
             mlirContextSetAllowUnregisteredDialects(self.get(), value);
           })
+      .def("enable_multithreading",
+           [](PyMlirContext &self, bool enable) {
+             mlirContextEnableMultithreading(self.get(), enable);
+           })
       .def("is_registered_operation",
            [](PyMlirContext &self, std::string &name) {
              return mlirContextIsRegisteredOperation(

mlir/lib/Bindings/Python/Pass.cpp

@@ -68,6 +68,18 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyPassManager::createFromCapsule)
       .def("_testing_release", &PyPassManager::release,
            "Releases (leaks) the backing pass manager (testing)")
+      .def(
+          "enable_ir_printing",
+          [](PyPassManager &passManager) {
+            mlirPassManagerEnableIRPrinting(passManager.get());
+          },
+          "Enable print-ir-after-all.")
+      .def(
+          "enable_verifier",
+          [](PyPassManager &passManager, bool enable) {
+            mlirPassManagerEnableVerifier(passManager.get(), enable);
+          },
+          "Enable / disable verify-each.")
       .def_static(
           "parse",
           [](const std::string pipeline, DefaultingPyMlirContext context) {

mlir/lib/Bindings/Python/mlir/ir.py

@@ -6,3 +6,8 @@
 from ._cext_loader import _reexport_cext
 _reexport_cext("ir", __name__)
 del _reexport_cext
+
+# Extra functions that are not visible to _reexport_cext.
+# TODO: is this really necessary?
+from _mlir.ir import _enable_debug
+_enable_debug = _enable_debug

mlir/lib/CAPI/IR/IR.cpp

@@ -21,8 +21,16 @@
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Parser.h"
 
+#include "llvm/Support/Debug.h"
+
 using namespace mlir;
 
+//===----------------------------------------------------------------------===//
+// Global API.
+//===----------------------------------------------------------------------===//
+
+void mlirEnableGlobalDebug(bool enable) { ::llvm::DebugFlag = true; }
+
 //===----------------------------------------------------------------------===//
 // Context API.
 //===----------------------------------------------------------------------===//
@@ -64,6 +72,10 @@ bool mlirContextIsRegisteredOperation(MlirContext context, MlirStringRef name) {
   return unwrap(context)->isOperationRegistered(unwrap(name));
 }
 
+void mlirContextEnableMultithreading(MlirContext context, bool enable) {
+  return unwrap(context)->enableMultithreading(enable);
+}
+
 //===----------------------------------------------------------------------===//
 // Dialect API.
 //===----------------------------------------------------------------------===//

mlir/lib/CAPI/IR/Pass.cpp

@@ -38,6 +38,14 @@ MlirLogicalResult mlirPassManagerRun(MlirPassManager passManager,
   return wrap(unwrap(passManager)->run(unwrap(module)));
 }
 
+void mlirPassManagerEnableIRPrinting(MlirPassManager passManager) {
+  return unwrap(passManager)->enableIRPrinting();
+}
+
+void mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable) {
+  unwrap(passManager)->enableVerifier(enable);
+}
+
 MlirOpPassManager mlirPassManagerGetNestedUnder(MlirPassManager passManager,
                                                 MlirStringRef operationName) {
   return wrap(&unwrap(passManager)->nest(unwrap(operationName)));

mlir/test/Bindings/Python/context_managers.py

@@ -10,6 +10,13 @@ def run(f):
   assert Context._get_live_count() == 0
 
 
+# CHECK-LABEL: TEST: testExports
+def testExports():
+  from mlir.ir import _enable_debug
+
+run(testExports)
+
+
 # CHECK-LABEL: TEST: testContextEnterExit
 def testContextEnterExit():
   with Context() as ctx: