Add a GDB/LLDB interface for interactive debugging of MLIR Actions

[mehdi_amini]

Summary

This includes a small runtime acting as callback for the ExecutionEngine
and a C API that makes it possible to control from the debugger.

A python script for LLDB is included that hook a new mlir subcommand
and allows to set breakpoints and inspect the current action, the context
and the stack.

[zero9178]

mlir/lib/Debug/GdbDebugExecutionContextHook.cpp

301

__attribute__((noinline)) is not a thing when compiling with MSVC. This should use LLVM_ATTRIBUTE_NOINLINE from llvm/Support/Compilers.h

[mehdi_amini]

Use LLVM_ATTRIBUTE_NOINLINE

[mehdi_amini]

Thanks for catching this.

[mehdi_amini]

rebase

[Mogball]

mlir/include/mlir/Debug/GdbDebugExecutionContextHook.h

2

this only works for gdb?

mlir/lib/Debug/GdbDebugExecutionContextHook.cpp

306

isn't there an LLVM_EXPORT thing that can be used in place of this?

[mehdi_amini]

mlir/include/mlir/Debug/GdbDebugExecutionContextHook.h

2

If you saw the demo, you may have noticed it was done with lldb :)

mlir/lib/Debug/GdbDebugExecutionContextHook.cpp

306

Nothing portable that I know of?

[rriddle]

I think for this revision we definitely need some high level usage docs

mlir/include/mlir/Debug/GdbDebugExecutionContextHook.h

2

Can we avoid using the name GDB? It's somewhat confusing in the modern multi-debugger world.

23–24

Why lowercase names?

mlir/lib/Debug/GdbDebugExecutionContextHook.cpp

13

Please drop the commented out code. Also it feels like a lot of these includes could be dropped.

32–37

Can you document this a bit?

48

Please drop the newline here.

306

The LLVM dump methods (and what we use for the mlir visualizers) is: LLVM_ATTRIBUTE_NOINLINE LLVM_ATTRIBUTE_USED

[jpienaar]

mlir/include/mlir/Debug/GdbDebugExecutionContextHook.h

2

(i was even wondering about DAP)

[mehdi_amini]

Rebase and address comments

mlir/lib/Debug/GdbDebugExecutionContextHook.cpp

306

I am hesitant to rely on LLVM_ATTRIBUTE_USED which is non-standard (not supported by MSVC?) for something that isn't a "debugging" things like the dump methods but something core to the feature here.

[tschuett]

mlir/lib/Debug/DebuggerExecutionContextHook.cpp

45 ↗

(On Diff #516196)

You need LLVM_THREAD_LOCAL.

[mehdi_amini]

Use LLVM_THREAD_LOCAL

mlir/lib/Debug/DebuggerExecutionContextHook.cpp

45 ↗

(On Diff #516196)

Interesting, I didn't know this one (or forgot about it, we don't use thread_local often).

That said I see one use of LLVM_THREAD_LOCAL in MLIR, vs >=15 used of thread_local as a keyword. Seems like right now we need a mass migration of MLIR!

[mehdi_amini]

One more LLVM_THREAD_LOCAL

[mehdi_amini]

rebase and address comments

[mehdi_amini]

In D144817#4195596, @rriddle wrote:

I think for this revision we definitely need some high level usage docs

This is now done, in a child revision: https://reviews.llvm.org/D149037

[rriddle]

mlir/lib/Debug/DebuggerExecutionContextHook.cpp

127 ↗

(On Diff #516230)

Extra newline here.

45 ↗

(On Diff #516196)

Is that advice actually up to date? I've not seen LLVM_THREAD_LOCAL recommended ever over thread_local, and only ever see thread_local in new code.

mlir/lib/Debug/ExecutionContext.cpp

61–65

llvm::interleave?

mlir/utils/lldb-scripts/action_debugging.py

52

[mehdi_amini]

mlir/lib/Debug/DebuggerExecutionContextHook.cpp

45 ↗

(On Diff #516196)

I would think that the main advantage may be supporting the LLVM config flag to disable threading entirely?

[tschuett]

There are actually two keywords for thread local storage and it is about with or without threads.

[rriddle]

In D144817#4291431, @tschuett wrote:

There are actually two keywords for thread local storage and it is about with or without threads.

Gotcha. I don't think we are consistent with using those at all.

[mehdi_amini]

Address comments

mlir/include/mlir/Debug/BreakpointManagers/TagBreakpointManager.h

private:
   friend class TagBreakpointManager;
 };
 
 class TagBreakpointManager
     : public BreakpointManagerBase<TagBreakpointManager> {
 public:
   Breakpoint *match(const Action &action) const override {
     auto it = breakpoints.find(action.getTag());
-    if (it != breakpoints.end() && it->second->isEnabled()) {
+    if (it != breakpoints.end() && it->second->isEnabled())
       return it->second.get();
-    }
     return {};
   }
 
   /// Add a breakpoint to the manager for the given tag and return it.
   /// If a breakpoint already exists for the given tag, return the existing
   /// instance.
   TagBreakpoint *addBreakpoint(StringRef tag) {
     auto result = breakpoints.insert({tag, nullptr});

mlir/include/mlir/Debug/ExecutionContext.h

 struct ActionActiveStack {
 public:
   ActionActiveStack(const ActionActiveStack *parent, const Action &action,
                     int depth)
       : parent(parent), action(action), depth(depth) {}
   const ActionActiveStack *getParent() const { return parent; }
   const Action &getAction() const { return action; }
   int getDepth() const { return depth; }
+  void print(raw_ostream &os, bool withContext) const;
+  void dump() const {
+    print(llvm::errs(), /*withContext=*/true);
+    llvm::errs() << "\n";
+  }
+  Breakpoint *getBreakpoint() const { return breakpoint; }
+  void setBreakpoint(Breakpoint *breakpoint) { this->breakpoint = breakpoint; }
 
 private:
+  Breakpoint *breakpoint = nullptr;
   const ActionActiveStack *parent;
   const Action &action;
   int depth;
 };
 
 /// The ExecutionContext is the main orchestration of the
 /// infrastructure, it acts as a handler in the MLIRContext for executing an
 /// Action. When an action is dispatched, it'll query its set of Breakpoints
public:
   using CallbackTy = llvm::function_ref<Control(const ActionActiveStack *)>;
 
   /// Create an ExecutionContext with a callback that is used to control the
   /// execution.
   ExecutionContext(CallbackTy callback) { setCallback(callback); }
   ExecutionContext() = default;
 
   /// Set the callback that is used to control the execution.
-  void setCallback(CallbackTy callback);
+  void setCallback(CallbackTy callback) {
+    onBreakpointControlExecutionCallback = callback;
+  }
 
   /// This abstract class defines the interface used to observe an Action
   /// execution. It allows to be notified before and after the callback is
   /// processed, but can't affect the execution.
   struct Observer {
     virtual ~Observer() = default;
     /// This method is called before the Action is executed
     /// If a breakpoint was hit, it is passed as an argument to the callback.
private:
   /// The list of managers that are queried for breakpoints.
   SmallVector<BreakpointManager *> breakpoints;
 };
 
 } // namespace tracing
 } // namespace mlir
 
 #endif // MLIR_TRACING_EXECUTIONCONTEXT_H
 No newline at end of file

mlir/include/mlir/Debug/GdbDebugExecutionContextHook.h

+//===- GdbDebugExecutionContextHook.h - GDB for Debugger Support *- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a set of C API functions that are used by the debugger to
+// interact with the ExecutionContext.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SUPPORT_GDBDEBUGEXECUTIONCONTEXTHOOK_H
+#define MLIR_SUPPORT_GDBDEBUGEXECUTIONCONTEXTHOOK_H
+
+#include "mlir-c/IR.h"
+#include "mlir/Debug/ExecutionContext.h"
+#include "llvm/Support/Compiler.h"
+
+extern "C" {
+struct breakpoint;
+struct irunit;
+typedef struct breakpoint *breakpointHandle;
+typedef struct irunit *irunitHandle;
+
+/// This is used by the debugger to control what to do after a breakpoint is
+/// hit. See tracing::ExecutionContext::Control for more information.
+void mlirDebuggerSetControl(int controlOption);
+
+/// Print the available context for the current Action.
+void mlirDebuggerPrintContext();
+
+/// Print the current action backtrace.
+void mlirDebuggerPrintActionBacktrace(bool withContext);
+
+//===----------------------------------------------------------------------===//
+// Cursor Management: The cursor is used to select an IRUnit from the context
+// and to navigate through the IRUnit hierarchy.
+//===----------------------------------------------------------------------===//
+
+/// Print the current IR unit cursor.
+void mlirDebuggerCursorPrint(bool withRegion);
+
+/// Select the IR unit from the current context by ID.
+void mlirDebuggerCursorSelectIRUnitFromContext(int index);
+
+/// Select the parent IR unit of the provided IR unit, or print an error if the
+/// IR unit has no parent.
+void mlirDebuggerCursorSelectParentIRUnit();
+
+/// Select the child IR unit at the provided index, print an error if the index
+/// is out of bound. For example if the irunit is an operation, the children IR
+/// units will be the operation's regions.
+void mlirDebuggerCursorSelectChildIRUnit(int index);
+
+/// Return the next IR unit logically in the IR. For example if the irunit is a
+/// Region the next IR unit will be the next region in the parent operation or
+/// nullptr if there is no next region.
+void mlirDebuggerCursorSelectPreviousIRUnit();
+
+/// Return the previous IR unit logically in the IR. For example if the irunit
+/// is a Region, the previous IR unit will be the previous region in the parent
+/// operation or nullptr if there is no previous region.
+void mlirDebuggerCursorSelectNextIRUnit();
+
+//===----------------------------------------------------------------------===//
+// Breakpoint Management
+//===----------------------------------------------------------------------===//
+
+/// Enable the provided breakpoint.
+void mlirDebuggerEnableBreakpoint(breakpointHandle breakpoint);
+
+/// Disable the provided breakpoint.
+void mlirDebuggerDisableBreakpoint(breakpointHandle breakpoint);
+
+/// Add a breakpoint matching exactly the provided tag.
+breakpointHandle mlirDebuggerAddTagBreakpoint(const char *tag);
+
+/// Add a breakpoint matching a pattern by name.
+void mlirDebuggerAddRewritePatternBreakpoint(const char *patternNameInfo);
+
+/// Add a breakpoint matching a file, line and column.
+void mlirDebuggerAddFileLineColLocBreakpoint(const char *file, int line,
+                                             int col);
+
+} // extern "C"
+
+namespace mlir {
+// Setup the GDB hooks as a callback on the provided ExecutionContext.
+void setupGdbDebugExecutionContextHook(
+    tracing::ExecutionContext &executionContext);
+
+} // namespace mlir
+
+#endif // MLIR_SUPPORT_GDBDEBUGEXECUTIONCONTEXTHOOK_H

mlir/include/mlir/IR/OperationSupport.h

public:
   /// readable 'pretty' form. Note: The IR generated with 'prettyForm' is not
   /// parsable.
   OpPrintingFlags &enableDebugInfo(bool enable = true, bool prettyForm = false);
 
   /// Always print operations in the generic form.
   OpPrintingFlags &printGenericOpForm();
 
   /// Skip printing regions.
-  OpPrintingFlags &skipRegions();
+  OpPrintingFlags &skipRegions(bool skip = true);
 
   /// Do not verify the operation when using custom operation printers.
   OpPrintingFlags &assumeVerified();
 
   /// Use local scope when printing the operation. This allows for using the
   /// printer in a more localized and thread-safe setting, but may not
   /// necessarily be identical to what the IR will look like when dumping
   /// the full module.

mlir/include/mlir/Rewrite/PatternApplicator.h

public:
   using Base = tracing::ActionImpl<ApplyPatternAction>;
   ApplyPatternAction(ArrayRef<IRUnit> irUnits, const Pattern &pattern)
       : Base(irUnits), pattern(pattern) {}
   static constexpr StringLiteral tag = "apply-pattern-action";
   static constexpr StringLiteral desc =
       "Encapsulate the application of rewrite patterns";
 
   void print(raw_ostream &os) const override {
-    os << "`" << tag << "`\n"
-       << " pattern: " << pattern.getDebugName() << '\n';
+    os << "`" << tag << " pattern: " << pattern.getDebugName();
   }
 
 private:
   const Pattern &pattern;
 };
 
 /// This class manages the application of a group of rewrite patterns, with a
 /// user-provided cost model.

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

public:
 
   /// Set the output format to bytecode instead of textual IR.
   MlirOptMainConfig &emitBytecode(bool emit) {
     emitBytecodeFlag = emit;
     return *this;
   }
   bool shouldEmitBytecode() const { return emitBytecodeFlag; }
 
+  /// Enable the gdb action hook: it makes GDB able to intercept MLIR Actions.
+  void enableGdbActionHook(bool enabled = true) {
+    enableGdbActionHookFlag = enabled;
+  }
+
+  /// Return true if the gdb action hook is enabled.
+  bool isGdbActionHookEnabled() const { return enableGdbActionHookFlag; }
+
   /// Set the filename to use for logging actions, use "-" for stdout.
   MlirOptMainConfig &logActionsTo(StringRef filename) {
     logActionsToFlag = filename;
     return *this;
   }
   /// Get the filename to use for logging actions.
   StringRef getLogActionsTo() const { return logActionsToFlag; }
 
protected:
   bool allowUnregisteredDialectsFlag = false;
 
   /// Print the pipeline that will be run.
   bool dumpPassPipelineFlag = false;
 
   /// Emit bytecode instead of textual assembly when generating output.
   bool emitBytecodeFlag = false;
 
+  /// Enable the gdb action hook: GDB can intercept MLIR Actions.
+  bool enableGdbActionHookFlag = false;
+
   /// Log action execution to the given file (or "-" for stdout)
   std::string logActionsToFlag;
 
   /// Location Breakpoints to filter the action logging.
   std::vector<tracing::BreakpointManager *> logActionLocationFilter;
 
   /// The callback to populate the pass manager.
   std::function<LogicalResult(PassManager &)> passPipelineCallback;

mlir/lib/Debug/CMakeLists.txt

 add_subdirectory(Observers)
 
 add_mlir_library(MLIRDebug
   DebugCounter.cpp
   ExecutionContext.cpp
   BreakpointManagers/FileLineColLocBreakpointManager.cpp
+  GdbDebugExecutionContextHook.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Debug
 
   LINK_LIBS PUBLIC
   ${LLVM_PTHREAD_LIB}
   MLIRIR
   MLIRSupport
 )
 

mlir/lib/Debug/ExecutionContext.cpp

 //===- ExecutionContext.cpp - Debug Execution Context Support -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Debug/ExecutionContext.h"
 
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/Support/FormatVariadic.h"
 
 #include <cstddef>
 
 using namespace mlir;
 using namespace mlir::tracing;
 
 //===----------------------------------------------------------------------===//
+// ActionActiveStack
+//===----------------------------------------------------------------------===//
+
+void ActionActiveStack::print(raw_ostream &os, bool withContext) const {
+  os << "ActionActiveStack depth " << getDepth() << "\n";
+  const ActionActiveStack *current = this;
+  int count = 0;
+  while (current) {
+    llvm::errs() << llvm::formatv("#{0,3}: ", count++);
+    current->action.print(llvm::errs());
+    llvm::errs() << "\n";
+    ArrayRef<IRUnit> context = current->action.getContextIRUnits();
+    if (withContext && !context.empty()) {
+      llvm::errs() << "Context:\n";
+      for (const IRUnit &unit : current->action.getContextIRUnits()) {
+        llvm::errs() << "  - ";
+        unit.print(llvm::errs());
+        llvm::errs() << "\n";
+      }
+      llvm::errs() << "\n";
+    }
+    current = current->parent;
+  }
+}
+
+//===----------------------------------------------------------------------===//
 // ExecutionContext
 //===----------------------------------------------------------------------===//
 
 static const thread_local ActionActiveStack *actionStack = nullptr;
 
-void ExecutionContext::setCallback(CallbackTy callback) {
-  onBreakpointControlExecutionCallback = callback;
-}
-
 void ExecutionContext::registerObserver(Observer *observer) {
   observers.push_back(observer);
 }
 
 void ExecutionContext::operator()(llvm::function_ref<void()> transform,
                                   const Action &action) {
   // Update the top of the stack with the current action.
   int depth = 0;
   if (actionStack)
     depth = actionStack->getDepth() + 1;
   ActionActiveStack info{actionStack, action, depth};
   actionStack = &info;
   auto raii = llvm::make_scope_exit([&]() { actionStack = info.getParent(); });
   Breakpoint *breakpoint = nullptr;
 
   // Invoke the callback here and handles control requests here.
   auto handleUserInput = [&]() -> bool {
     if (!onBreakpointControlExecutionCallback)
       return true;
     auto todoNext = onBreakpointControlExecutionCallback(actionStack);
     switch (todoNext) {
     case ExecutionContext::Apply:
       depthToBreak = std::nullopt;
void ExecutionContext::operator()(llvm::function_ref<void()> transform,
 
   // Try to find a breakpoint that would hit on this action.
   // Right now there is no way to collect them all, we stop at the first one.
   for (auto *breakpointManager : breakpoints) {
     breakpoint = breakpointManager->match(action);
     if (breakpoint)
       break;
   }
+  info.setBreakpoint(breakpoint);
 
   bool shouldExecuteAction = true;
   // If we have a breakpoint, or if `depthToBreak` was previously set and the
   // current depth matches, we invoke the user-provided callback.
   if (breakpoint || (depthToBreak && depth <= depthToBreak))
     shouldExecuteAction = handleUserInput();
 
   // Notify the observers about the current action.
   for (auto *observer : observers)
     observer->beforeExecute(actionStack, breakpoint, shouldExecuteAction);
 
   if (shouldExecuteAction) {
 
     // Execute the action here.
     transform();
 
     // Notify the observers about completion of the action.
     for (auto *observer : observers)
       observer->afterExecute(actionStack);
   }
 
   if (depthToBreak && depth <= depthToBreak)
     handleUserInput();
 }
 No newline at end of file

mlir/lib/Debug/GdbDebugExecutionContextHook.cpp

+//===- GdbDebugExecutionContextHook.cpp - GDB for Debugger Support --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Debug/GdbDebugExecutionContextHook.h"
+
+#include "mlir/Debug/BreakpointManagers/FileLineColLocBreakpointManager.h"
+// #include "mlir/Debug/BreakpointManagers/RewritePatternBreakpointManager.h"
+#include "mlir/Debug/BreakpointManagers/TagBreakpointManager.h"
+#include "mlir/Debug/ExecutionContext.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/Region.h"
+#include "mlir/IR/Unit.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/FormatVariadic.h"
+#include <cstddef>
+#include <signal.h>
+
+using namespace mlir;
+using namespace mlir::tracing;
+
+namespace {
+struct DebuggerState {
+  ExecutionContext::Control gdbControl = ExecutionContext::Apply;
+  TagBreakpointManager tagBreakpointManager;
+  FileLineColLocBreakpointManager fileLineColLocBreakpointManager;
+  DenseMap<unsigned, Breakpoint *> breakpointIdsMap;
+  const tracing::ActionActiveStack *actionActiveStack;
+  IRUnit cursor;
+};
+} // namespace
+
+static DebuggerState &getGlobalDebuggerState() {
+  static thread_local DebuggerState debuggerState;
+  return debuggerState;
+}
+
+extern "C" {
+void mlirDebuggerSetControl(int controlOption) {
+
+  getGlobalDebuggerState().gdbControl =
+      static_cast<ExecutionContext::Control>(controlOption);
+}
+
+void mlirDebuggerPrintContext() {
+  DebuggerState &state = getGlobalDebuggerState();
+  if (!state.actionActiveStack) {
+    llvm::outs() << "No active action.\n";
+    return;
+  }
+  const ArrayRef<IRUnit> &units =
+      state.actionActiveStack->getAction().getContextIRUnits();
+  llvm::outs() << units.size() << " available IRUnits:\n";
+  for (const IRUnit &unit : units) {
+    llvm::outs() << "  - ";
+    unit.print(
+        llvm::outs(),
+        OpPrintingFlags().useLocalScope().skipRegions().enableDebugInfo());
+    llvm::outs() << "\n";
+  }
+}
+
+void mlirDebuggerPrintActionBacktrace(bool withContext) {
+  DebuggerState &state = getGlobalDebuggerState();
+  if (!state.actionActiveStack) {
+    llvm::outs() << "No active action.\n";
+    return;
+  }
+  state.actionActiveStack->print(llvm::outs(), withContext);
+}
+
+//===----------------------------------------------------------------------===//
+// Cursor Management
+//===----------------------------------------------------------------------===//
+
+void mlirDebuggerCursorPrint(bool withRegion) {
+  auto &state = getGlobalDebuggerState();
+  if (!state.cursor) {
+    llvm::outs() << "No active MLIR cursor, select from the context first\n";
+    return;
+  }
+  state.cursor.print(llvm::outs(), OpPrintingFlags()
+                                       .skipRegions(!withRegion)
+                                       .useLocalScope()
+                                       .enableDebugInfo());
+  llvm::outs() << "\n";
+}
+
+void mlirDebuggerCursorSelectIRUnitFromContext(int index) {
+  auto &state = getGlobalDebuggerState();
+  if (!state.actionActiveStack) {
+    llvm::outs() << "No active MLIR Action stack\n";
+    return;
+  }
+  ArrayRef<IRUnit> units =
+      state.actionActiveStack->getAction().getContextIRUnits();
+  if (index < 0 || index >= static_cast<int>(units.size())) {
+    llvm::outs() << "Index invalid, bounds: [0, " << units.size()
+                 << "] but got " << index << "\n";
+    return;
+  }
+  state.cursor = units[index];
+  state.cursor.print(llvm::outs());
+  llvm::outs() << "\n";
+}
+
+void mlirDebuggerCursorSelectParentIRUnit() {
+  auto &state = getGlobalDebuggerState();
+  if (!state.cursor) {
+    llvm::outs() << "No active MLIR cursor, select from the context first\n";
+    return;
+  }
+  IRUnit *unit = &state.cursor;
+  if (auto *op = unit->dyn_cast<Operation *>()) {
+    state.cursor = op->getBlock();
+  } else if (auto *region = unit->dyn_cast<Region *>()) {
+
+    state.cursor = region->getParentOp();
+  } else if (auto *block = unit->dyn_cast<Block *>()) {
+    state.cursor = block->getParent();
+  } else {
+    llvm::outs() << "Current cursor is not a valid IRUnit";
+    return;
+  }
+  state.cursor.print(llvm::outs());
+  llvm::outs() << "\n";
+}
+
+void mlirDebuggerCursorSelectChildIRUnit(int index) {
+  auto &state = getGlobalDebuggerState();
+  if (!state.cursor) {
+    llvm::outs() << "No active MLIR cursor, select from the context first\n";
+    return;
+  }
+  IRUnit *unit = &state.cursor;
+  if (auto *op = unit->dyn_cast<Operation *>()) {
+    if (index < 0 || index >= static_cast<int>(op->getNumRegions())) {
+      llvm::outs() << "Index invalid, op has " << op->getNumRegions()
+                   << " but got " << index << "\n";
+      return;
+    }
+    state.cursor = &op->getRegion(index);
+  } else if (auto *region = unit->dyn_cast<Region *>()) {
+    auto block = region->begin();
+    int count = 0;
+    while (block != region->end() && count != index) {
+      ++block;
+      ++count;
+    }
+
+    if (block == region->end()) {
+      llvm::outs() << "Index invalid, region has " << count << " block but got "
+                   << index << "\n";
+      return;
+    }
+    state.cursor = &*block;
+  } else if (auto *block = unit->dyn_cast<Block *>()) {
+    auto op = block->begin();
+    int count = 0;
+    while (op != block->end() && count != index) {
+      ++op;
+      ++count;
+    }
+
+    if (op == block->end()) {
+      llvm::outs() << "Index invalid, block has " << count
+                   << "operations but got " << index << "\n";
+      return;
+    }
+    state.cursor = &*op;
+  } else {
+    llvm::outs() << "Current cursor is not a valid IRUnit";
+    return;
+  }
+  state.cursor.print(llvm::outs());
+  llvm::outs() << "\n";
+}
+
+void mlirDebuggerCursorSelectPreviousIRUnit() {
+  auto &state = getGlobalDebuggerState();
+  if (!state.cursor) {
+    llvm::outs() << "No active MLIR cursor, select from the context first\n";
+    return;
+  }
+  IRUnit *unit = &state.cursor;
+  if (auto *op = unit->dyn_cast<Operation *>()) {
+    Operation *previous = op->getPrevNode();
+    if (!previous) {
+      llvm::outs() << "No previous operation in the current block\n";
+      return;
+    }
+    state.cursor = previous;
+  } else if (auto *region = unit->dyn_cast<Region *>()) {
+    llvm::outs() << "Has region\n";
+    Operation *parent = region->getParentOp();
+    if (!parent) {
+      llvm::outs() << "No parent operation for the current region\n";
+      return;
+    }
+    if (region->getRegionNumber() == 0) {
+      llvm::outs() << "No previous region in the current operation\n";
+      return;
+    }
+    state.cursor =
+        &region->getParentOp()->getRegion(region->getRegionNumber() - 1);
+  } else if (auto *block = unit->dyn_cast<Block *>()) {
+    Block *previous = block->getPrevNode();
+    if (!previous) {
+      llvm::outs() << "No previous block in the current region\n";
+      return;
+    }
+    state.cursor = previous;
+  } else {
+    llvm::outs() << "Current cursor is not a valid IRUnit";
+    return;
+  }
+  state.cursor.print(llvm::outs());
+  llvm::outs() << "\n";
+}
+
+void mlirDebuggerCursorSelectNextIRUnit() {
+  auto &state = getGlobalDebuggerState();
+  if (!state.cursor) {
+    llvm::outs() << "No active MLIR cursor, select from the context first\n";
+    return;
+  }
+  IRUnit *unit = &state.cursor;
+  if (auto *op = unit->dyn_cast<Operation *>()) {
+    Operation *next = op->getNextNode();
+    if (!next) {
+      llvm::outs() << "No next operation in the current block\n";
+      return;
+    }
+    state.cursor = next;
+  } else if (auto *region = unit->dyn_cast<Region *>()) {
+    Operation *parent = region->getParentOp();
+    if (!parent) {
+      llvm::outs() << "No parent operation for the current region\n";
+      return;
+    }
+    if (region->getRegionNumber() == parent->getNumRegions() - 1) {
+      llvm::outs() << "No next region in the current operation\n";
+      return;
+    }
+    state.cursor =
+        &region->getParentOp()->getRegion(region->getRegionNumber() + 1);
+  } else if (auto *block = unit->dyn_cast<Block *>()) {
+    Block *next = block->getNextNode();
+    if (!next) {
+      llvm::outs() << "No next block in the current region\n";
+      return;
+    }
+    state.cursor = next;
+  } else {
+    llvm::outs() << "Current cursor is not a valid IRUnit";
+    return;
+  }
+  state.cursor.print(llvm::outs());
+  llvm::outs() << "\n";
+}
+
+//===----------------------------------------------------------------------===//
+// Breakpoint Management
+//===----------------------------------------------------------------------===//
+
+void mlirDebuggerEnableBreakpoint(breakpointHandle breakpoint) {
+  reinterpret_cast<Breakpoint *>(breakpoint)->enable();
+}
+
+void mlirDebuggerDisableBreakpoint(breakpointHandle breakpoint) {
+  reinterpret_cast<Breakpoint *>(breakpoint)->disable();
+}
+
+breakpointHandle mlirDebuggerAddTagBreakpoint(const char *tag) {
+  DebuggerState &state = getGlobalDebuggerState();
+  Breakpoint *breakpoint =
+      state.tagBreakpointManager.addBreakpoint(StringRef(tag, strlen(tag)));
+  int breakpointId = state.breakpointIdsMap.size() + 1;
+  state.breakpointIdsMap[breakpointId] = breakpoint;
+  return reinterpret_cast<breakpointHandle>(breakpoint);
+}
+
+void mlirDebuggerAddRewritePatternBreakpoint(const char *patternNameInfo) {}
+
+void mlirDebuggerAddFileLineColLocBreakpoint(const char *file, int line,
+                                             int col) {
+  getGlobalDebuggerState().fileLineColLocBreakpointManager.addBreakpoint(
+      StringRef(file, strlen(file)), line, col);
+}
+
+} // extern "C"
+
+LLVM_ATTRIBUTE_NOINLINE void mlirGdbBreakpointHook() {
+  static thread_local void *volatile sink;
+  sink = (void *)&sink;
+}
+
+static void preventLinkerDeadCodeElim() {
+  static void *volatile sink;
+  static bool initialized = [&]() {
+    sink = (void *)mlirDebuggerSetControl;
+    sink = (void *)mlirDebuggerEnableBreakpoint;
+    sink = (void *)mlirDebuggerDisableBreakpoint;
+    sink = (void *)mlirDebuggerPrintContext;
+    sink = (void *)mlirDebuggerPrintActionBacktrace;
+    sink = (void *)mlirDebuggerCursorPrint;
+    sink = (void *)mlirDebuggerCursorSelectIRUnitFromContext;
+    sink = (void *)mlirDebuggerCursorSelectParentIRUnit;
+    sink = (void *)mlirDebuggerCursorSelectChildIRUnit;
+    sink = (void *)mlirDebuggerCursorSelectPreviousIRUnit;
+    sink = (void *)mlirDebuggerCursorSelectNextIRUnit;
+    sink = (void *)mlirDebuggerAddTagBreakpoint;
+    sink = (void *)mlirDebuggerAddRewritePatternBreakpoint;
+    sink = (void *)mlirDebuggerAddFileLineColLocBreakpoint;
+    sink = (void *)&sink;
+    return true;
+  }();
+  (void)initialized;
+}
+
+static tracing::ExecutionContext::Control
+gdbCallBackFunction(const tracing::ActionActiveStack *actionStack) {
+  preventLinkerDeadCodeElim();
+  // Invoke the breakpoint hook, the debugger is supposed to trap this.
+  // The debugger controls the execution from there by invoking
+  // `mlirDebuggerSetControl()`.
+  auto &state = getGlobalDebuggerState();
+  state.actionActiveStack = actionStack;
+  getGlobalDebuggerState().gdbControl = ExecutionContext::Apply;
+  actionStack->getAction().print(llvm::outs());
+  llvm::outs() << "\n";
+  mlirGdbBreakpointHook();
+  return getGlobalDebuggerState().gdbControl;
+}
+
+namespace {
+/// Manage the stack of actions that are currently active.
+class DebuggerObserver : public ExecutionContext::Observer {
+  void beforeExecute(const ActionActiveStack *action, Breakpoint *breakpoint,
+                     bool willExecute) override {
+    auto &state = getGlobalDebuggerState();
+    state.actionActiveStack = action;
+  }
+  void afterExecute(const ActionActiveStack *action) override {
+    auto &state = getGlobalDebuggerState();
+    state.actionActiveStack = action->getParent();
+    state.cursor = nullptr;
+  }
+};
+} // namespace
+
+void mlir::setupGdbDebugExecutionContextHook(
+    tracing::ExecutionContext &executionContext) {
+  executionContext.setCallback(gdbCallBackFunction);
+  DebuggerState &state = getGlobalDebuggerState();
+  static DebuggerObserver observer;
+  executionContext.registerObserver(&observer);
+  executionContext.addBreakpointManager(&state.fileLineColLocBreakpointManager);
+  executionContext.addBreakpointManager(&state.tagBreakpointManager);
+}

mlir/lib/IR/AsmPrinter.cpp

 
 /// Always print operations in the generic form.
 OpPrintingFlags &OpPrintingFlags::printGenericOpForm() {
   printGenericOpFormFlag = true;
   return *this;
 }
 
 /// Always skip Regions.
-OpPrintingFlags &OpPrintingFlags::skipRegions() {
-  skipRegionsFlag = true;
+OpPrintingFlags &OpPrintingFlags::skipRegions(bool skip) {
+  skipRegionsFlag = skip;
   return *this;
 }
 
 /// Do not verify the operation when using custom operation printers.
 OpPrintingFlags &OpPrintingFlags::assumeVerified() {
   assumeVerifiedFlag = true;
   return *this;
 }

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

 // out. It is designed to support unit testing.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/Debug/Counter.h"
 #include "mlir/Debug/ExecutionContext.h"
+#include "mlir/Debug/GdbDebugExecutionContextHook.h"
 #include "mlir/Debug/Observers/ActionLogging.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
MlirOptMainConfigCLOptions() {
     static cl::opt<bool, /*ExternalStorage=*/true> dumpPassPipeline(
         "dump-pass-pipeline", cl::desc("Print the pipeline that will be run"),
         cl::location(dumpPassPipelineFlag), cl::init(false));
 
     static cl::opt<bool, /*ExternalStorage=*/true> emitBytecode(
         "emit-bytecode", cl::desc("Emit bytecode when generating output"),
         cl::location(emitBytecodeFlag), cl::init(false));
 
+    static cl::opt<bool, /*ExternalStorage=*/true> enableGdbHook(
+        "mlir-enable-gdb-hook",
+        cl::desc("Enable gdb hook for debugging MLIR Actions"),
+        cl::location(enableGdbActionHookFlag), cl::init(false));
+
     static cl::opt<bool, /*ExternalStorage=*/true> explicitModule(
         "no-implicit-module",
         cl::desc("Disable implicit addition of a top-level module op during "
                  "parsing"),
         cl::location(useExplicitModuleFlag), cl::init(false));
 
     static cl::opt<std::string, /*ExternalStorage=*/true> logActionsTo{
         "log-actions-to",
MlirOptMainConfig &MlirOptMainConfig::setPassPipelineParser(
   };
   return *this;
 }
 
 /// Set the ExecutionContext on the context and handle the observers.
 class InstallDebugHandler {
 public:
   InstallDebugHandler(MLIRContext &context, const MlirOptMainConfig &config) {
-    if (config.getLogActionsTo().empty()) {
+    if (config.getLogActionsTo().empty() && !config.isGdbActionHookEnabled()) {
       if (tracing::DebugCounter::isDebugCounterActivated())
         context.registerActionHandler(tracing::DebugCounter{});
       return;
     }
+    llvm::errs() << "ExecutionContext registered on the context";
     if (tracing::DebugCounter::isDebugCounterActivated())
       emitError(UnknownLoc::get(&context),
-                "Debug counters are incompatible with --log-actions-to option "
-                "and are disabled");
+                "Debug counters are incompatible with --log-actions-to and "
+                "--mlir-enable-gdb-hook options and are disabled");
+    if (!config.getLogActionsTo().empty()) {
       std::string errorMessage;
       logActionsFile = openOutputFile(config.getLogActionsTo(), &errorMessage);
       if (!logActionsFile) {
-      llvm::errs() << errorMessage << "\n";
+        llvm::errs() << "Error handling --log-actions-to: " << errorMessage
+                     << "\n";
         return;
       }
       logActionsFile->keep();
       raw_fd_ostream &logActionsStream = logActionsFile->os();
       actionLogger = std::make_unique<tracing::ActionLogger>(logActionsStream);
       for (const auto *locationBreakpoint : config.getLogActionsLocFilters())
         actionLogger->addBreakpointManager(locationBreakpoint);
-
       executionContext.registerObserver(actionLogger.get());
+    }
+    if (config.isGdbActionHookEnabled()) {
+      llvm::errs() << " (with GDB hook)";
+      setupGdbDebugExecutionContextHook(executionContext);
+    }
+    llvm::errs() << "\n";
     context.registerActionHandler(executionContext);
   }
 
 private:
   std::unique_ptr<llvm::ToolOutputFile> logActionsFile;
   tracing::ExecutionContext executionContext;
   std::unique_ptr<tracing::ActionLogger> actionLogger;
   std::vector<std::unique_ptr<tracing::FileLineColLocBreakpoint>>

mlir/utils/lldb-scripts/action_debugging.py

+#!/usr/bin/env python
+
+# ---------------------------------------------------------------------
+# Be sure to add the python path that points to the LLDB shared library.
+#
+# # To use this in the embedded python interpreter using "lldb" just
+# import it with the full path using the "command script import"
+# command
+#   (lldb) command script import /path/to/cmdtemplate.py
+# ---------------------------------------------------------------------
+
+import inspect
+import lldb
+import argparse
+import shlex
+import sys
+
+# Each new breakpoint gets a unique ID starting from 1.
+nextid = 1
+# List of breakpoint set from python, the key is the ID and the value the
+# actual breakpoint. These are NOT LLDB SBBreakpoint objects.
+breakpoints = dict()
+
+exprOptions = lldb.SBExpressionOptions()
+exprOptions.SetIgnoreBreakpoints()
+exprOptions.SetLanguage(lldb.eLanguageTypeC)
+
+
+class MlirDebug:
+    """MLIR debugger commands
+    This is the class that hooks into LLDB and registers the `mlir` command.
+    Other providers can register subcommands below this one.
+    """
+
+    lldb_command = "mlir"
+    parser = None
+
+    def __init__(self, debugger, unused):
+        super().__init__()
+        self.create_options()
+        self.help_string = MlirDebug.parser.format_help()
+
+    @classmethod
+    def create_options(cls):
+        if MlirDebug.parser:
+            return MlirDebug.parser
+        usage = "usage: %s [options]" % (cls.lldb_command)
+        description = "TODO."
+
+        # Pass add_help_option = False, since this keeps the command in line
+        #  with lldb commands, and we wire up "help command" to work by
+        # providing the long & short help methods below.
+        MlirDebug.parser = argparse.ArgumentParser(
+            prog=cls.lldb_command, usage=usage, description=description, add_help=False
+        )
+        MlirDebug.subparsers = MlirDebug.parser.add_subparsers(dest="command")
+        return MlirDebug.parser
+
+    def get_short_help(self):
+        return "MLIR debugger commands"
+
+    def get_long_help(self):
+        return self.help_string
+
+    def __call__(self, debugger, command, exe_ctx, result):
+        # Use the Shell Lexer to properly parse up command options just like a
+        # shell would
+        command_args = shlex.split(command)
+
+        try:
+            args = MlirDebug.parser.parse_args(command_args)
+        except:
+            result.SetError("option parsing failed")
+            raise
+        args.func(args, debugger, command, exe_ctx, result)
+
+    @classmethod
+    def on_process_start(frame, bp_loc, dict):
+        print("Process started")
+
+
+class SetControl:
+    # Define the subcommands that controls what to do when a breakpoint is hit.
+    # The key is the subcommand name, the value is a tuple of the command ID to
+    # pass to MLIR and the help string.
+    commands = {
+        "apply": (1, "Apply the current action and continue the execution"),
+        "skip": (2, "Skip the current action and continue the execution"),
+        "step": (3, "Step into the current action"),
+        "next": (4, "Step over the current action"),
+        "finish": (5, "Step out of the current action"),
+    }
+
+    @classmethod
+    def register_mlir_subparser(cls):
+        for cmd, (cmdInt, help) in cls.commands.items():
+            parser = MlirDebug.subparsers.add_parser(
+                cmd,
+                help=help,
+            )
+            parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        frame = exe_ctx.GetFrame()
+        if not frame.IsValid():
+            result.SetError("No valid frame (program not running?)")
+            return
+        cmdInt = cls.commands.get(options.command, None)
+        if not cmdInt:
+            result.SetError("Invalid command: %s" % (options.command))
+            return
+
+        result = frame.EvaluateExpression(
+            "((bool (*)(int))mlirDebuggerSetControl)(%d)" % (cmdInt[0]),
+            exprOptions,
+        )
+        if not result.error.Success():
+            print("Error setting up command: %s" % (result.error))
+            return
+        debugger.SetAsync(True)
+        result = exe_ctx.GetProcess().Continue()
+        debugger.SetAsync(False)
+
+
+class PrintContext:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "context", help="Print the current context"
+        )
+        cls.parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        frame = exe_ctx.GetFrame()
+        if not frame.IsValid():
+            result.SetError("Can't print context without a valid frame")
+            return
+        result = frame.EvaluateExpression(
+            "((bool (*)())&mlirDebuggerPrintContext)()", exprOptions
+        )
+        if not result.error.Success():
+            print("Error printing context: %s" % (result.error))
+            return
+
+
+class Backtrace:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "backtrace", aliases=["bt"], help="Print the current backtrace"
+        )
+        cls.parser.set_defaults(func=cls.process_options)
+        cls.parser.add_argument("--context", default=False, action="store_true")
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        frame = exe_ctx.GetFrame()
+        if not frame.IsValid():
+            result.SetError(
+                "Can't backtrace without a valid frame (program not running?)"
+            )
+        result = frame.EvaluateExpression(
+            "((bool(*)(bool))mlirDebuggerPrintActionBacktrace)(%d)" % (options.context),
+            exprOptions,
+        )
+        if not result.error.Success():
+            print("Error printing breakpoints: %s" % (result.error))
+            return
+
+
+###############################################################################
+# Cursor manipulation
+###############################################################################
+
+
+class PrintCursor:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "cursor-print", aliases=["cursor-p"], help="Print the current cursor"
+        )
+        cls.parser.add_argument(
+            "--print-region", "--regions", "-r", default=False, action="store_true"
+        )
+        cls.parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        frame = exe_ctx.GetFrame()
+        if not frame.IsValid():
+            result.SetError(
+                "Can't print cursor without a valid frame (program not running?)"
+            )
+        result = frame.EvaluateExpression(
+            "((bool(*)(bool))mlirDebuggerCursorPrint)(%d)" % (options.print_region),
+            exprOptions,
+        )
+        if not result.error.Success():
+            print("Error printing cursor: %s" % (result.error))
+            return
+
+
+class SelectCursorFromContext:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "cursor-select-from-context",
+            aliases=["cursor-s"],
+            help="Select the cursor from the current context",
+        )
+        cls.parser.add_argument("index", type=int, help="Index in the context")
+        cls.parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        frame = exe_ctx.GetFrame()
+        if not frame.IsValid():
+            result.SetError(
+                "Can't manipulate cursor without a valid frame (program not running?)"
+            )
+        result = frame.EvaluateExpression(
+            "((bool(*)(int))mlirDebuggerCursorSelectIRUnitFromContext)(%d)"
+            % options.index,
+            exprOptions,
+        )
+        if not result.error.Success():
+            print("Error manipulating cursor: %s" % (result.error))
+            return
+
+
+class CursorSelectParent:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "cursor-parent", aliases=["cursor-up"], help="Select the cursor parent"
+        )
+        cls.parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        frame = exe_ctx.GetFrame()
+        if not frame.IsValid():
+            result.SetError(
+                "Can't manipulate cursor without a valid frame (program not running?)"
+            )
+        result = frame.EvaluateExpression(
+            "((bool(*)())mlirDebuggerCursorSelectParentIRUnit)()",
+            exprOptions,
+        )
+        if not result.error.Success():
+            print("Error manipulating cursor: %s" % (result.error))
+            return
+
+
+class SelectCursorChild:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "cursor-child", aliases=["cursor-c"], help="Select the nth child"
+        )
+        cls.parser.add_argument("index", type=int, help="Index of the child to select")
+        cls.parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        frame = exe_ctx.GetFrame()
+        if not frame.IsValid():
+            result.SetError(
+                "Can't manipulate cursor without a valid frame (program not running?)"
+            )
+        result = frame.EvaluateExpression(
+            "((bool(*)(int))mlirDebuggerCursorSelectChildIRUnit)(%d)" % options.index,
+            exprOptions,
+        )
+        if not result.error.Success():
+            print("Error manipulating cursor: %s" % (result.error))
+            return
+
+
+class CursorSelecPrevious:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "cursor-previous",
+            aliases=["cursor-prev"],
+            help="Select the cursor previous element",
+        )
+        cls.parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        frame = exe_ctx.GetFrame()
+        if not frame.IsValid():
+            result.SetError(
+                "Can't manipulate cursor without a valid frame (program not running?)"
+            )
+        result = frame.EvaluateExpression(
+            "((bool(*)())mlirDebuggerCursorSelectPreviousIRUnit)()",
+            exprOptions,
+        )
+        if not result.error.Success():
+            print("Error manipulating cursor: %s" % (result.error))
+            return
+
+
+class CursorSelecNext:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "cursor-next", aliases=["cursor-n"], help="Select the cursor next element"
+        )
+        cls.parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        frame = exe_ctx.GetFrame()
+        if not frame.IsValid():
+            result.SetError(
+                "Can't manipulate cursor without a valid frame (program not running?)"
+            )
+        result = frame.EvaluateExpression(
+            "((bool(*)())mlirDebuggerCursorSelectNextIRUnit)()",
+            exprOptions,
+        )
+        if not result.error.Success():
+            print("Error manipulating cursor: %s" % (result.error))
+            return
+
+
+###############################################################################
+# Breakpoints
+###############################################################################
+
+
+class EnableBreakpoint:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "enable", help="Enable a single breakpoint (given its ID)"
+        )
+        cls.parser.add_argument("id", help="ID of the breakpoint to enable")
+        cls.parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        bp = breakpoints.get(int(options.id), None)
+        if not bp:
+            result.SetError("No breakpoint with ID %d" % int(options.id))
+            return
+        bp.enable(exe_ctx.GetFrame())
+
+
+class DisableBreakpoint:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "disable", help="Disable a single breakpoint (given its ID)"
+        )
+        cls.parser.add_argument("id", help="ID of the breakpoint to disable")
+        cls.parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        bp = breakpoints.get(int(options.id), None)
+        if not bp:
+            result.SetError("No breakpoint with ID %s" % options.id)
+            return
+        bp.disable(exe_ctx.GetFrame())
+
+
+class ListBreakpoints:
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            "list", help="List all current breakpoints"
+        )
+        cls.parser.set_defaults(func=cls.process_options)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        for id, bp in sorted(breakpoints.items()):
+            print(id, type(id), str(bp), "enabled" if bp.isEnabled else "disabled")
+
+
+class Breakpoint:
+    def __init__(self):
+        global nextid
+        self.id = nextid
+        nextid += 1
+        breakpoints[self.id] = self
+        self.isEnabled = True
+
+    def enable(self, frame=None):
+        self.isEnabled = True
+        if not frame or not frame.IsValid():
+            return
+        # use a C cast to force the type of the breakpoint handle to be void * so
+        # that we don't rely on DWARF. Also add a fake bool return value otherwise
+        # LLDB can't signal any error with the expression evaluation (at least I don't know how).
+        cmd = (
+            "((bool (*)(void *))mlirDebuggerEnableBreakpoint)((void *)%s)" % self.handle
+        )
+        result = frame.EvaluateExpression(cmd, exprOptions)
+        if not result.error.Success():
+            print("Error enabling breakpoint: %s" % (result.error))
+            return
+
+    def disable(self, frame=None):
+        self.isEnabled = False
+        if not frame or not frame.IsValid():
+            return
+        # use a C cast to force the type of the breakpoint handle to be void * so
+        # that we don't rely on DWARF. Also add a fake bool return value otherwise
+        # LLDB can't signal any error with the expression evaluation (at least I don't know how).
+        cmd = (
+            "((bool (*)(void *)) mlirDebuggerDisableBreakpoint)((void *)%s)"
+            % self.handle
+        )
+        result = frame.EvaluateExpression(cmd, exprOptions)
+        if not result.error.Success():
+            print("Error disabling breakpoint: %s" % (result.error))
+            return
+
+
+class TagBreakpoint(Breakpoint):
+    mlir_subcommand = "break-on-tag"
+
+    def __init__(self, tag):
+        super().__init__()
+        self.tag = tag
+
+    def __str__(self):
+        return "[%d] TagBreakpoint(%s)" % (self.id, self.tag)
+
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            cls.mlir_subcommand, help="add a breakpoint on actions' tag matching"
+        )
+        cls.parser.set_defaults(func=cls.process_options)
+        cls.parser.add_argument("tag", help="tag to match")
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        breakpoint = TagBreakpoint(options.tag)
+        print("Added breakpoint %s" % str(breakpoint))
+
+        frame = exe_ctx.GetFrame()
+        if frame.IsValid():
+            breakpoint.install(frame)
+
+    def install(self, frame):
+        result = frame.EvaluateExpression(
+            '((void *(*)(const char *))mlirDebuggerAddTagBreakpoint)("%s")'
+            % (self.tag),
+            exprOptions,
+        )
+        if not result.error.Success():
+            print("Error installing breakpoint: %s" % (result.error))
+            return
+        # Save the handle, this is necessary to implement enable/disable.
+        self.handle = result.GetValue()
+
+
+class FileLineBreakpoint(Breakpoint):
+    mlir_subcommand = "break-on-file"
+
+    def __init__(self, file, line, col):
+        super().__init__()
+        self.file = file
+        self.line = line
+        self.col = col
+
+    def __str__(self):
+        return "[%d] FileLineBreakpoint(%s, %d, %d)" % (
+            self.id,
+            self.file,
+            self.line,
+            self.col,
+        )
+
+    @classmethod
+    def register_mlir_subparser(cls):
+        cls.parser = MlirDebug.subparsers.add_parser(
+            cls.mlir_subcommand,
+            help="add a breakpoint that filters on location of the IR affected by an action. The syntax is file:line:col where file and col are optional",
+        )
+        cls.parser.set_defaults(func=cls.process_options)
+        cls.parser.add_argument("location", type=str)
+
+    @classmethod
+    def process_options(cls, options, debugger, command, exe_ctx, result):
+        split_loc = options.location.split(":")
+        file = split_loc[0]
+        line = int(split_loc[1]) if len(split_loc) > 1 else -1
+        col = int(split_loc[2]) if len(split_loc) > 2 else -1
+        breakpoint = FileLineBreakpoint(file, line, col)
+        print("Added breakpoint %s" % str(breakpoint))
+
+        frame = exe_ctx.GetFrame()
+        if frame.IsValid():
+            breakpoint.install(frame)
+
+    def install(self, frame):
+        result = frame.EvaluateExpression(
+            '((void *(*)(const char *, int, int))mlirDebuggerAddFileLineColLocBreakpoint)("%s", %d, %d)'
+            % (self.file, self.line, self.col),
+            exprOptions,
+        )
+        if not result.error.Success():
+            print("Error installing breakpoint: %s" % (result.error))
+            return
+        # Save the handle, this is necessary to implement enable/disable.
+        self.handle = result.GetValue()
+
+
+def on_start(frame, bpno, err):
+    print("MLIR debugger attaching...")
+    for _, bp in sorted(breakpoints.items()):
+        if bp.isEnabled:
+            print("Installing breakpoint %s" % (str(bp)))
+            bp.install(frame)
+        else:
+            print("Skipping disabled breakpoint %s" % (str(bp)))
+
+    return True
+
+
+def __lldb_init_module(debugger, dict):
+    target = debugger.GetTargetAtIndex(0)
+    debugger.SetAsync(False)
+    if not target:
+        print("No target is loaded, please load a target before loading this script.")
+        return
+    if debugger.GetNumTargets() > 1:
+        print(
+            "Multiple targets (%s) loaded, attaching MLIR debugging to %s"
+            % (debugger.GetNumTargets(), target)
+        )
+
+    # Register all classes that have a register_lldb_command method
+    module_name = __name__
+    parser = MlirDebug.create_options()
+    MlirDebug.__doc__ = parser.format_help()
+
+    # Add the MLIR entry point to LLDB as a command.
+    command = "command script add -o -c %s.%s %s" % (
+        module_name,
+        MlirDebug.__name__,
+        MlirDebug.lldb_command,
+    )
+    debugger.HandleCommand(command)
+
+    main_bp = target.BreakpointCreateByName("main")
+    main_bp.SetScriptCallbackFunction("action_debugging.on_start")
+    main_bp.SetAutoContinue(auto_continue=True)
+
+    on_breackpoint = target.BreakpointCreateByName("mlirGdbBreakpointHook")
+
+    print(
+        'The "{0}" command has been installed for target `{1}`, type "help {0}" or "{0} '
+        '--help" for detailed help.'.format(MlirDebug.lldb_command, target)
+    )
+    for _name, cls in inspect.getmembers(sys.modules[module_name]):
+        if inspect.isclass(cls) and getattr(cls, "register_mlir_subparser", None):
+            cls.register_mlir_subparser()