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

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

Author: ftynse

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

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());
+        });
   }
 };
 

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