[Mlir-commits] [mlir] [mlir] expose -debug-only equivalent to C and Python (PR #93175)

[Oleksandr Alex Zinenko]

https://github.com/ftynse created https://github.com/llvm/llvm-project/pull/93175

These are useful for finer-grain debugging and complement the already exposed global debug flag.

>From f7f8adc8fdb577eeb8a64af0dc6eb416fb6b9997 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <ftynse at gmail.com>
Date: Thu, 23 May 2024 13:22:06 +0200
Subject: [PATCH] [mlir] expose -debug-only equivalent to C and Python

These are useful for finer-grain debugging and complement the already
exposed global debug flag.
---
 mlir/include/mlir-c/Debug.h         | 13 +++++++++++++
 mlir/lib/Bindings/Python/IRCore.cpp | 15 ++++++++++++++-
 mlir/lib/CAPI/Debug/Debug.cpp       | 18 ++++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir-c/Debug.h b/mlir/include/mlir-c/Debug.h
index 2502f2fa23bf0..7dad73500858d 100644
--- a/mlir/include/mlir-c/Debug.h
+++ b/mlir/include/mlir-c/Debug.h
@@ -21,6 +21,19 @@ MLIR_CAPI_EXPORTED void mlirEnableGlobalDebug(bool enable);
 /// Retuns `true` if the global debugging flag is set, false otherwise.
 MLIR_CAPI_EXPORTED bool mlirIsGlobalDebugEnabled();
 
+/// Sets the current debug type, similarly to `-debug-only=type` in the
+/// command-line tools. Note that global debug should be enabled for any output
+/// to be produced.
+MLIR_CAPI_EXPORTED void mlirSetGlobalDebugType(const char *type);
+
+/// Sets multiple current debug types, similarly to `-debug-only=type1,type2" in
+/// the command-line tools. Note that global debug should be enabled for any
+/// output to be produced.
+MLIR_CAPI_EXPORTED void mlirSetGlobalDebugTypes(const char **types, intptr_t n);
+
+/// Checks if `type` is set as the current debug type.
+MLIR_CAPI_EXPORTED bool mlirIsCurrentDebugType(const char *type);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 01678a9719f90..2b2792ea6c776 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -240,7 +240,20 @@ struct PyGlobalDebugFlag {
     // Debug flags.
     py::class_<PyGlobalDebugFlag>(m, "_GlobalDebug", py::module_local())
         .def_property_static("flag", &PyGlobalDebugFlag::get,
-                             &PyGlobalDebugFlag::set, "LLVM-wide debug flag");
+                             &PyGlobalDebugFlag::set, "LLVM-wide debug flag")
+        .def_static(
+            "set_types",
+            [](const std::string &type) {
+              mlirSetGlobalDebugType(type.c_str());
+            },
+            "types"_a, "Sets specific debug types to be produced by LLVM")
+        .def_static("set_types", [](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());
+          mlirSetGlobalDebugTypes(pointers.data(), pointers.size());
+        });
   }
 };
 
diff --git a/mlir/lib/CAPI/Debug/Debug.cpp b/mlir/lib/CAPI/Debug/Debug.cpp
index 288ecd6012749..320ece4998e04 100644
--- a/mlir/lib/CAPI/Debug/Debug.cpp
+++ b/mlir/lib/CAPI/Debug/Debug.cpp
@@ -16,3 +16,21 @@
 void mlirEnableGlobalDebug(bool enable) { llvm::DebugFlag = enable; }
 
 bool mlirIsGlobalDebugEnabled() { return llvm::DebugFlag; }
+
+void mlirSetGlobalDebugType(const char *type) {
+  // Depending on the NDEBUG flag, this name can be either a function or a macro
+  // that expands to something that isn't a funciton call, so we cannot
+  // explicitly prefix it with `llvm::` or declare `using` it.
+  using namespace llvm;
+  setCurrentDebugType(type);
+}
+
+void mlirSetGlobalDebugTypes(const char **types, intptr_t n) {
+  using namespace llvm;
+  setCurrentDebugTypes(types, n);
+}
+
+bool mlirIsCurrentDebugType(const char *type) {
+  using namespace llvm;
+  return isCurrentDebugType(type);
+}