[mlir] Improve debug flag management in Python bindings

[ftynse]

Summary

Expose the debug flag as a readable and assignable property of a
dedicated class instead of a write-only function. Actually test the fact
of setting the flag. Move test to a dedicated file, it has zero relation
to context_managers.py where it was added.

Arguably, it should be promoted from mlir.ir to mlir module, but we are
not re-exporting the latter and this functionality is purposefully
hidden so can stay in IR for now. Drop unnecessary export code.

Refactor C API and put Debug into a separate library, fix it to actually
set the flag to the given value.

[ftynse]

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

11

@nicolasvasilache, nope it was not necessary.

[nicolasvasilache]

Cool ,thanks!

mlir/include/mlir-c/Debug.h

+//===-- mlir-c/Debug.h - C API for MLIR/LLVM debugging functions --*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Support.h"
+
+#include <stdbool.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Sets the global debugging flag.
+MLIR_CAPI_EXPORTED void mlirEnableGlobalDebug(bool enable);
+
+/// Retuns `true` if the global debugging flag is set, false otherwise.
+MLIR_CAPI_EXPORTED bool mlirIsGlobalDebugEnabled();
+
+#ifdef __cplusplus
+}
+#endif
+
+#ifndef MLIR_C_DEBUG_H
+#define MLIR_C_DEBUG_H
+#endif // MLIR_C_DEBUG_H

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);

mlir/lib/Bindings/Python/IRCore.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 "IRModule.h"
 
 #include "Globals.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
+#include "mlir-c/Debug.h"
 #include "mlir-c/Registration.h"
 #include "llvm/ADT/SmallVector.h"
 #include <pybind11/stl.h>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 

 position in the argument list. If the value is an operation result, this is
 equivalent to printing the operation that produced it.
 )";
 
 //------------------------------------------------------------------------------
 // Utilities.
 //------------------------------------------------------------------------------
 
-// Helper for creating an @classmethod.
+/// Helper for creating an @classmethod.
 template <class Func, typename... Args>
 py::object classmethod(Func f, Args... args) {
   py::object cf = py::cpp_function(f, args...);
   return py::reinterpret_borrow<py::object>((PyClassMethod_New(cf.ptr())));
 }
 
 static py::object
 createCustomDialectWrapper(const std::string &dialectNamespace,
                            py::object dialectDescriptor) {
   auto dialectClass = PyGlobals::get().lookupDialectClass(dialectNamespace);
   if (!dialectClass) {
     // Use the base class.
     return py::cast(PyDialect(std::move(dialectDescriptor)));
   }
 
   // Create the custom implementation.
   return (*dialectClass)(std::move(dialectDescriptor));
 }
 
 static MlirStringRef toMlirStringRef(const std::string &s) {
   return mlirStringRefCreate(s.data(), s.size());
 }
 
+/// Wrapper for the global LLVM debugging flag.
+struct PyGlobalDebugFlag {
+  static void set(py::object &o, bool enable) { mlirEnableGlobalDebug(enable); }
+
+  static bool get(py::object) { return mlirIsGlobalDebugEnabled(); }
+
+  static void bind(py::module &m) {
+    // Debug flags.
+    py::class_<PyGlobalDebugFlag>(m, "_GlobalDebug")
+        .def_property_static("flag", &PyGlobalDebugFlag::get,
+                             &PyGlobalDebugFlag::set, "LLVM-wide debug flag");
+  }
+};
+
 //------------------------------------------------------------------------------
 // Collections.
 //------------------------------------------------------------------------------
 
 namespace {
 
 class PyRegionIterator {
 public:

 } // 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
+  // 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());
              return ref.releaseObject();
void mlir::python::populateIRCore(py::module &m) {
   PyBlockList::bind(m);
   PyOperationIterator::bind(m);
   PyOperationList::bind(m);
   PyOpAttributeMap::bind(m);
   PyOpOperandList::bind(m);
   PyOpResultList::bind(m);
   PyRegionIterator::bind(m);
   PyRegionList::bind(m);
+
+  // Debug bindings.
+  PyGlobalDebugFlag::bind(m);
 }

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
 

mlir/lib/CAPI/CMakeLists.txt

+add_subdirectory(Debug)
 add_subdirectory(Dialect)
 add_subdirectory(Conversion)
 add_subdirectory(ExecutionEngine)
 add_subdirectory(IR)
 add_subdirectory(Registration)
 add_subdirectory(Transforms)
 
 

mlir/lib/CAPI/Debug/CMakeLists.txt

+add_mlir_public_c_api_library(MLIRCAPIDebug
+  Debug.cpp
+
+  LINK_LIBS PUBLIC
+  MLIRSupport
+)

mlir/lib/CAPI/Debug/Debug.cpp

+//===- Debug.cpp - C Interface for MLIR/LLVM Debugging Functions ----------===//
+//
+// 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 "mlir-c/Debug.h"
+#include "mlir-c/Support.h"
+
+#include "mlir/CAPI/Support.h"
+
+#include "llvm/Support/Debug.h"
+
+void mlirEnableGlobalDebug(bool enable) { llvm::DebugFlag = enable; }
+
+bool mlirIsGlobalDebugEnabled() { return llvm::DebugFlag; }

mlir/lib/CAPI/IR/IR.cpp

 #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);
 }
 

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

mlir/test/Bindings/Python/debug.py

+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+
+
+# CHECK-LABEL: TEST: testNameIsPrivate
+def testNameIsPrivate():
+  # `import *` ignores private names starting with an understore, so the debug
+  # flag shouldn't be visible unless explicitly imported.
+  try:
+    _GlobalDebug.flag = True
+  except NameError:
+    pass
+  else:
+    assert False, "_GlobalDebug must not be available by default"
+
+run(testNameIsPrivate)
+
+
+# CHECK-LABEL: TEST: testDebugDlag
+def testDebugDlag():
+  # Private names must be imported expilcitly.
+  from mlir.ir import _GlobalDebug
+
+  # CHECK: False
+  print(_GlobalDebug.flag)
+  _GlobalDebug.flag = True
+  # CHECK: True
+  print(_GlobalDebug.flag)
+  _GlobalDebug.flag = False
+  # CHECK: False
+  print(_GlobalDebug.flag)
+
+run(testDebugDlag)
+