[mlir][python] Add debug helpersusers/nico/python-2

This PR lets users provide --debug-only flags with python decorators.
This greatly simplifies the debugging experience in python.

Co-authored-by: Tres <tpopp@users.noreply.github.com>

2 files changed, 35 insertions, 1 deletions

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 002923b..6ce27d1 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -282,7 +282,26 @@ struct PyGlobalDebugFlag {
             pointers.push_back(str.c_str());
           nb::ft_lock_guard lock(mutex);
           mlirSetGlobalDebugTypes(pointers.data(), pointers.size());
-        });
+        })
+        .def_static(
+            "push_debug_only_flags",
+            [](const std::string &type) {
+              mlirAppendGlobalDebugType(type.c_str());
+            },
+            "flags"_a,
+            "Appends specific debug only flags which can be popped later.")
+        .def_static("push_debug_only_flags",
+                    [](const std::vector<std::string> &types) {
+                      std::vector<const char *> pointers;
+                      pointers.reserve(types.size());
+                      for (const std::string &str : types)
+                        pointers.push_back(str.c_str());
+                      mlirAppendGlobalDebugTypes(pointers.data(),
+                                                 pointers.size());
+                    })
+        .def_static(
+            "pop_debug_only_flags", []() { mlirPopAppendedGlobalDebugTypes(); },
+            "Removes the latest non-popped addition from push_debug_only_flags.");
   }
 
 private:
diff --git a/mlir/lib/CAPI/Debug/Debug.cpp b/mlir/lib/CAPI/Debug/Debug.cpp
index 320ece4..c76a109 100644
--- a/mlir/lib/CAPI/Debug/Debug.cpp
+++ b/mlir/lib/CAPI/Debug/Debug.cpp
@@ -34,3 +34,18 @@ bool mlirIsCurrentDebugType(const char *type) {
   using namespace llvm;
   return isCurrentDebugType(type);
 }
+
+void mlirAppendGlobalDebugType(const char *type) {
+  using namespace llvm;
+  appendDebugType(type);
+}
+
+void mlirAppendGlobalDebugTypes(const char **types, intptr_t n) {
+  using namespace llvm;
+  appendDebugTypes(types, n);
+}
+
+void mlirPopAppendedGlobalDebugTypes() {
+  using namespace llvm;
+  popAppendedDebugTypes();
+}