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

[nicolasvasilache]

Export and test missing function

[nicolasvasilache]

Update commit message.

mlir/include/mlir-c/IR.h

 
 struct MlirNamedAttribute {
   MlirIdentifier name;
   MlirAttribute attribute;
 };
 typedef struct MlirNamedAttribute MlirNamedAttribute;
 
 //===----------------------------------------------------------------------===//
+// Global API.
+//===----------------------------------------------------------------------===//
+
+/// Set the global debugging flag.
+MLIR_CAPI_EXPORTED void mlirEnableGlobalDebug(bool enable);
+
+//===----------------------------------------------------------------------===//
 // Context API.
 //===----------------------------------------------------------------------===//
 
 /// Creates an MLIR context and transfers its ownership to the caller.
 MLIR_CAPI_EXPORTED MlirContext mlirContextCreate();
 
 /// Checks if two contexts are equal.
 MLIR_CAPI_EXPORTED bool mlirContextEqual(MlirContext ctx1, MlirContext ctx2);

 /// Gets the dialect instance owned by the given context using the dialect
 /// namespace to identify it, loads (i.e., constructs the instance of) the
 /// dialect if necessary. If the dialect is not registered with the context,
 /// returns null. Use mlirContextLoad<Name>Dialect to load an unregistered
 /// dialect.
 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
 /// dialect.
 MLIR_CAPI_EXPORTED bool mlirContextIsRegisteredOperation(MlirContext context,
                                                          MlirStringRef name);
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir-c/Pass.h

 /// Cast a top-level PassManager to a generic OpPassManager.
 MLIR_CAPI_EXPORTED MlirOpPassManager
 mlirPassManagerGetAsOpPassManager(MlirPassManager passManager);
 
 /// Run the provided `passManager` on the given `module`.
 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.
 /// To further nest more OpPassManager under the newly returned one, see
 /// `mlirOpPassManagerNest` below.
 MLIR_CAPI_EXPORTED MlirOpPassManager mlirPassManagerGetNestedUnder(
     MlirPassManager passManager, MlirStringRef operationName);
 

mlir/lib/Bindings/Python/IRCore.cpp

 } // end namespace
 
 //------------------------------------------------------------------------------
 // Populates the core exports of the 'ir' submodule.
 //------------------------------------------------------------------------------
 
 void mlir::python::populateIRCore(py::module &m) {
   //----------------------------------------------------------------------------
+  // Mapping of Global functions
+  //----------------------------------------------------------------------------
+  m.def("_enable_debug", [](bool enable) { mlirEnableGlobalDebug(enable); });
+
+  //----------------------------------------------------------------------------
   // Mapping of MlirContext
   //----------------------------------------------------------------------------
   py::class_<PyMlirContext>(m, "Context")
       .def(py::init<>(&PyMlirContext::createNewContextForInit))
       .def_static("_get_live_count", &PyMlirContext::getLiveCount)
       .def("_get_context_again",
            [](PyMlirContext &self) {
              PyMlirContextRef ref = PyMlirContext::forContext(self.get());
py::class_<PyMlirContext>(m, "Context")
       .def_property(
           "allow_unregistered_dialects",
           [](PyMlirContext &self) -> bool {
             return mlirContextGetAllowUnregisteredDialects(self.get());
           },
           [](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(
                  self.get(), MlirStringRef{name.data(), name.size()});
            });
 
   //----------------------------------------------------------------------------
   // Mapping of PyDialectDescriptor

mlir/lib/Bindings/Python/Pass.cpp

py::class_<PyPassManager>(m, "PassManager")
            }),
            py::arg("context") = py::none(),
            "Create a new PassManager for the current (or provided) Context.")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyPassManager::getCapsule)
       .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) {
             MlirPassManager passManager = mlirPassManagerCreate(context->get());
             MlirLogicalResult status = mlirParsePassPipeline(
                 mlirPassManagerGetAsOpPassManager(passManager),
                 mlirStringRefCreate(pipeline.data(), pipeline.size()));
             if (mlirLogicalResultIsFailure(status))

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

 #  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
 
 # Simply a wrapper around the extension module of the same name.
 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
 No newline at end of file

mlir/lib/CAPI/IR/IR.cpp

 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Verifier.h"
 #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.
 //===----------------------------------------------------------------------===//
 
 MlirContext mlirContextCreate() {
   auto *context = new MLIRContext;
   return wrap(context);
 }
 
MlirDialect mlirContextGetOrLoadDialect(MlirContext context,
                                         MlirStringRef name) {
   return wrap(unwrap(context)->getOrLoadDialect(unwrap(name)));
 }
 
 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.
 //===----------------------------------------------------------------------===//
 
 MlirContext mlirDialectGetContext(MlirDialect dialect) {
   return wrap(unwrap(dialect)->getContext());
 }
 

mlir/lib/CAPI/IR/Pass.cpp

mlirPassManagerGetAsOpPassManager(MlirPassManager passManager) {
   return wrap(static_cast<OpPassManager *>(unwrap(passManager)));
 }
 
 MlirLogicalResult mlirPassManagerRun(MlirPassManager passManager,
                                      MlirModule module) {
   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)));
 }
 
 MlirOpPassManager mlirOpPassManagerGetNestedUnder(MlirOpPassManager passManager,
                                                   MlirStringRef operationName) {
   return wrap(&unwrap(passManager)->nest(unwrap(operationName)));

mlir/test/Bindings/Python/context_managers.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   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:
     assert Context.current is ctx
   try:
     _ = Context.current
   except ValueError as e:
     # CHECK: No current Context