Enable the MLIR debug actions CL options in the compiler driver. (#18928

) Signed-off-by: Stella Laurenzo <stellaraccident@gmail.com>
iree-org · Oct 29, 2024 · 3b69679 · 3b69679
1 parent f4a5f13
commit 3b69679
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 0 deletions.
diff --git a/compiler/src/iree/compiler/API/Internal/BUILD.bazel b/compiler/src/iree/compiler/API/Internal/BUILD.bazel
@@ -38,6 +38,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:BuiltinToLLVMIRTranslation",
         "@llvm-project//mlir:BytecodeWriter",
         "@llvm-project//mlir:CAPIIR",
+        "@llvm-project//mlir:Debug",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Parser",
         "@llvm-project//mlir:Support",

diff --git a/compiler/src/iree/compiler/API/Internal/CMakeLists.txt b/compiler/src/iree/compiler/API/Internal/CMakeLists.txt
@@ -23,6 +23,7 @@ iree_cc_library(
     MLIRBuiltinToLLVMIRTranslation
     MLIRBytecodeWriter
     MLIRCAPIIR
+    MLIRDebug
     MLIRIR
     MLIRParser
     MLIRSupport

diff --git a/compiler/src/iree/compiler/API/Internal/CompilerDriver.cpp b/compiler/src/iree/compiler/API/Internal/CompilerDriver.cpp
@@ -67,6 +67,7 @@
 #include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/CAPI/IR.h"
 #include "mlir/CAPI/Wrap.h"
+#include "mlir/Debug/CLOptionsSetup.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
@@ -274,6 +275,7 @@ void GlobalInit::registerCommandLineOptions() {
   // Register pass manager command-line options like -mlir-print-ir-*.
   mlir::registerPassManagerCLOptions();
   mlir::registerDefaultTimingManagerCLOptions();
+  mlir::tracing::DebugConfig::registerCLOptions();
 
   // Bind session options to the command line environment.
   clPluginManagerOptions = &PluginManagerOptions::FromFlags::get();
@@ -366,6 +368,11 @@ struct Session {
   // All user access to the context is done via this reference.
   MLIRContext &context;
   OptionsBinder binder;
+
+  // Debug configuration.
+  mlir::tracing::DebugConfig debugConfig;
+  std::optional<mlir::tracing::InstallDebugHandler> debugHandlerInstall;
+
   // PluginManagerOptions must initialize first because the session depends on
   // it.
   PluginManagerOptions pluginManagerOptions;
@@ -402,6 +409,7 @@ Session::Session(GlobalInit &globalInit)
 
   // Bootstrap session options from the cl environment, if enabled.
   if (globalInit.usesCommandLine) {
+    debugConfig = mlir::tracing::DebugConfig::createFromCLOptions();
     pluginManagerOptions = *globalInit.clPluginManagerOptions;
     bindingOptions = *globalInit.clBindingOptions;
     inputOptions = *globalInit.clInputOptions;
@@ -417,6 +425,9 @@ Session::Session(GlobalInit &globalInit)
 #endif
   }
 
+  // Enable debug integration.
+  debugHandlerInstall.emplace(context, debugConfig);
+
   // Register each options struct with the binder so we can manipulate
   // mnemonically via the API.
   bindingOptions.bindOptions(binder);