[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>

7 files changed, 158 insertions, 4 deletions

diff --git a/mlir/include/mlir-c/Debug.h b/mlir/include/mlir-c/Debug.h
index 7dad735..8f206e4 100644
--- a/mlir/include/mlir-c/Debug.h
+++ b/mlir/include/mlir-c/Debug.h
@@ -31,6 +31,24 @@ MLIR_CAPI_EXPORTED void mlirSetGlobalDebugType(const char *type);
 /// output to be produced.
 MLIR_CAPI_EXPORTED void mlirSetGlobalDebugTypes(const char **types, intptr_t n);
 
+/// Adds to the current debug type state, similarly to
+/// `-debug-only=prev_type,new_type` in the command-line tools. Note that global
+/// debug should be enabled for any output to be produced. A single append call
+/// can be reverted with mlirPopAppendedGlobalDebugTypes.
+MLIR_CAPI_EXPORTED void mlirAppendGlobalDebugType(const char *type);
+
+/// Adds to the current debug type state, similarly to
+/// `-debug-only=prev_type,new_type1,new_type2` in the command-line tools. Note
+/// that global debug should be enabled for any output to be produced. A single
+/// append call can be reverted with mlirPopAppendedGlobalDebugTypes.
+MLIR_CAPI_EXPORTED void mlirAppendGlobalDebugTypes(const char **types,
+                                                   intptr_t n);
+
+/// Restores the current debug type state to its state before the last append
+/// call. An appended state of `-debug-only=prev_type,new_type1,new_type2` would
+/// be `-debug-only=prev_type` after this call.
+MLIR_CAPI_EXPORTED void mlirPopAppendedGlobalDebugTypes();
+
 /// Checks if `type` is set as the current debug type.
 MLIR_CAPI_EXPORTED bool mlirIsCurrentDebugType(const char *type);
 
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();
+}
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index 56b9f17..ffba1e3 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -2818,3 +2818,5 @@ class VectorType(ShapedType):
 
 class _GlobalDebug:
     flag: ClassVar[bool] = False
+    def push_debug_only_flags(self, types: list[str]) -> None: ...
+    def pop_debug_only_flags(self) -> None: ...
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
index 0d2eaff..d23da98 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
@@ -12,9 +12,10 @@ __all__ = [
 ]
 
 class PassManager:
-    def __init__(self, context: _ir.Context | None = None) -> None: ...
+    def __init__(self, anchor_op: str | None = None) -> None: ...
     def _CAPICreate(self) -> object: ...
     def _testing_release(self) -> None: ...
+    def add(self, pipeline: str) -> None: ...
     def enable_ir_printing(
         self,
         print_before_all: bool = False,
diff --git a/mlir/python/mlir/utils.py b/mlir/python/mlir/utils.py
index c6e9b57..cf8ad220 100644
--- a/mlir/python/mlir/utils.py
+++ b/mlir/python/mlir/utils.py
@@ -209,3 +209,79 @@ def call_with_toplevel_context_create_module(
     decorated = with_toplevel_context_create_module(f)
     decorated()
     return decorated
+
+
+def _debug_flags_impl(flags: Sequence[str]) -> Iterator[None]:
+    from mlir.ir import _GlobalDebug
+
+    # Save the original debug state. The debug flags will be popped rather than
+    # manually copied and saved for later.
+    original_flag = _GlobalDebug.flag
+    _GlobalDebug.flag = True
+    _GlobalDebug.push_debug_only_flags(flags)
+
+    try:
+        yield
+    finally:
+        # Reset the global debug flag and remove the most recent flags that were
+        # appended. This assumes that nothing else popped when it should not have.
+        _GlobalDebug.flag = original_flag
+        _GlobalDebug.pop_debug_only_flags()
+
+
+@contextmanager
+def debug_flags_context(flags: Sequence[str]):
+    """Temporarily create a context that enables debugging with specified filters.
+
+    These would be the same as running with -debug-only=*flags. Where multiple contexts
+    will be joined together to create the full list if they are nested.
+
+    This requires that the core MLIR units were compiled without NDEBUG.
+    """
+    return _debug_flags_impl(flags)
+
+
+@contextmanager
+def debug_conversion(flags: Sequence[str] = []) -> Iterator[None]:
+    """Temporarily create a context that enables full conversion debugging,
+    potentially with additional specified filters.
+
+    These would be the same as running with -debug-only=*flags. Where multiple contexts
+    will be joined together to create the full list if they are nested.
+
+    This requires that the core MLIR units were compiled without NDEBUG.
+    """
+    return _debug_flags_impl(list(flags) + ["dialect-conversion"])
+
+
+@contextmanager
+def debug_greedy_rewriter(flags: Sequence[str] = []) -> Iterator[None]:
+    """Temporarily create a context that enables full conversion debugging,
+    potentially with additional specified filters.
+
+    These would be the same as running with -debug-only=*flags. Where multiple contexts
+    will be joined together to create the full list if they are nested.
+
+    This requires that the core MLIR units were compiled without NDEBUG.
+    """
+    return _debug_flags_impl(list(flags) + ["greedy_rewriter"])
+
+
+@contextmanager
+def debug_td(flags: Sequence[str] = [], *, full_debug: bool = False) -> Iterator[None]:
+    """Temporarily create a context that enables full transform dialect debugging,
+    potentially with additional specified filters.
+
+    These would be the same as running with -debug-only=*flags. Where multiple contexts
+    will be joined together to create the full list if they are nested.
+
+    This requires that the core MLIR units were compiled without NDEBUG.
+    """
+    return _debug_flags_impl(
+        list(flags)
+        + [
+            "transform-dialect",
+            "transform-dialect-print-top-level-after-all",
+        ]
+        + (["transform-dialect-full"] if full_debug else [])
+    )
diff --git a/mlir/test/python/utils.py b/mlir/test/python/utils.py
index 8435fdd..c700808 100644
--- a/mlir/test/python/utils.py
+++ b/mlir/test/python/utils.py
@@ -1,13 +1,15 @@
 # RUN: %python %s | FileCheck %s
+# RUN: %python %s 2>&1 | FileCheck %s --check-prefix=DEBUG_ONLY
 
 import unittest
 
 from mlir import ir
-from mlir.dialects import arith, builtin
-from mlir.extras import types as T
+from mlir.passmanager import PassManager
+from mlir.dialects import arith
 from mlir.utils import (
     call_with_toplevel_context_create_module,
     caller_mlir_context,
+    debug_conversion,
     using_mlir_context,
 )
 
@@ -31,6 +33,7 @@ class TestRequiredContext(unittest.TestCase):
             c = arith.ConstantOp(value=42.42, result=ir.F32Type.get()).result
             multiple_adds(c, c)
 
+            # CHECK-LABEL: module {
             # CHECK: constant
             # CHECK-NEXT: arith.addf
             # CHECK-NEXT: arith.addf
@@ -54,5 +57,25 @@ class TestRequiredContext(unittest.TestCase):
             pass
 
 
+class TestDebugOnlyFlags(unittest.TestCase):
+    def test_debug_types(self):
+        """Test checks --debug-only=xxx functionality is available in MLIR."""
+
+        @debug_conversion()
+        def lower(module) -> None:
+            pm = PassManager("builtin.module")
+            pm.add("convert-arith-to-llvm")
+            pm.run(module.operation)
+
+        @call_with_toplevel_context_create_module
+        def _(module) -> None:
+            c = arith.ConstantOp(value=42.42, result=ir.F32Type.get()).result
+            arith.AddFOp(c, c, fastmath=arith.FastMathFlags.nnan | arith.FastMathFlags.ninf)
+
+            # DEBUG_ONLY-LABEL: Legalizing operation : 'builtin.module'
+            #       DEBUG_ONLY: Legalizing operation : 'arith.addf'
+            lower(module)
+
+
 if __name__ == "__main__":
     unittest.main()