mlir/utils/lldb-scripts/action_debugging.py

   1 #!/usr/bin/env python
   2
   3 # ---------------------------------------------------------------------
   4 # Be sure to add the python path that points to the LLDB shared library.
   5 #
   6 # # To use this in the embedded python interpreter using "lldb" just
   7 # import it with the full path using the "command script import"
   8 # command
   9 #   (lldb) command script import /path/to/cmdtemplate.py
  10 # ---------------------------------------------------------------------
  11
  12 import inspect
  13 import lldb
  14 import argparse
  15 import shlex
  16 import sys
  17
  18 # Each new breakpoint gets a unique ID starting from 1.
  19 nextid = 1
  20 # List of breakpoint set from python, the key is the ID and the value the
  21 # actual breakpoint. These are NOT LLDB SBBreakpoint objects.
  22 breakpoints = dict()
  23
  24 exprOptions = lldb.SBExpressionOptions()
  25 exprOptions.SetIgnoreBreakpoints()
  26 exprOptions.SetLanguage(lldb.eLanguageTypeC)
  27
  28
  29 class MlirDebug:
  30     """MLIR debugger commands
  31     This is the class that hooks into LLDB and registers the `mlir` command.
  32     Other providers can register subcommands below this one.
  33     """
  34
  35     lldb_command = "mlir"
  36     parser = None
  37
  38     def __init__(self, debugger, unused):
  39         super().__init__()
  40         self.create_options()
  41         self.help_string = MlirDebug.parser.format_help()
  42
  43     @classmethod
  44     def create_options(cls):
  45         if MlirDebug.parser:
  46             return MlirDebug.parser
  47         usage = "usage: %s [options]" % (cls.lldb_command)
  48         description = "TODO."
  49
  50         # Pass add_help_option = False, since this keeps the command in line
  51         # with lldb commands, and we wire up "help command" to work by
  52         # providing the long & short help methods below.
  53         MlirDebug.parser = argparse.ArgumentParser(
  54             prog=cls.lldb_command, usage=usage, description=description, add_help=False
  55         )
  56         MlirDebug.subparsers = MlirDebug.parser.add_subparsers(dest="command")
  57         return MlirDebug.parser
  58
  59     def get_short_help(self):
  60         return "MLIR debugger commands"
  61
  62     def get_long_help(self):
  63         return self.help_string
  64
  65     def __call__(self, debugger, command, exe_ctx, result):
  66         # Use the Shell Lexer to properly parse up command options just like a
  67         # shell would
  68         command_args = shlex.split(command)
  69
  70         try:
  71             args = MlirDebug.parser.parse_args(command_args)
  72         except:
  73             result.SetError("option parsing failed")
  74             raise
  75         args.func(args, debugger, command, exe_ctx, result)
  76
  77     @classmethod
  78     def on_process_start(frame, bp_loc, dict):
  79         print("Process started")
  80
  81
  82 class SetControl:
  83     # Define the subcommands that controls what to do when a breakpoint is hit.
  84     # The key is the subcommand name, the value is a tuple of the command ID to
  85     # pass to MLIR and the help string.
  86     commands = {
  87         "apply": (1, "Apply the current action and continue the execution"),
  88         "skip": (2, "Skip the current action and continue the execution"),
  89         "step": (3, "Step into the current action"),
  90         "next": (4, "Step over the current action"),
  91         "finish": (5, "Step out of the current action"),
  92     }
  93
  94     @classmethod
  95     def register_mlir_subparser(cls):
  96         for cmd, (cmdInt, help) in cls.commands.items():
  97             parser = MlirDebug.subparsers.add_parser(
  98                 cmd,
  99                 help=help,
 100             )
 101             parser.set_defaults(func=cls.process_options)
 102
 103     @classmethod
 104     def process_options(cls, options, debugger, command, exe_ctx, result):
 105         frame = exe_ctx.GetFrame()
 106         if not frame.IsValid():
 107             result.SetError("No valid frame (program not running?)")
 108             return
 109         cmdInt = cls.commands.get(options.command, None)
 110         if not cmdInt:
 111             result.SetError("Invalid command: %s" % (options.command))
 112             return
 113
 114         result = frame.EvaluateExpression(
 115             "((bool (*)(int))mlirDebuggerSetControl)(%d)" % (cmdInt[0]),
 116             exprOptions,
 117         )
 118         if not result.error.Success():
 119             print("Error setting up command: %s" % (result.error))
 120             return
 121         debugger.SetAsync(True)
 122         result = exe_ctx.GetProcess().Continue()
 123         debugger.SetAsync(False)
 124
 125
 126 class PrintContext:
 127     @classmethod
 128     def register_mlir_subparser(cls):
 129         cls.parser = MlirDebug.subparsers.add_parser(
 130             "context", help="Print the current context"
 131         )
 132         cls.parser.set_defaults(func=cls.process_options)
 133
 134     @classmethod
 135     def process_options(cls, options, debugger, command, exe_ctx, result):
 136         frame = exe_ctx.GetFrame()
 137         if not frame.IsValid():
 138             result.SetError("Can't print context without a valid frame")
 139             return
 140         result = frame.EvaluateExpression(
 141             "((bool (*)())&mlirDebuggerPrintContext)()", exprOptions
 142         )
 143         if not result.error.Success():
 144             print("Error printing context: %s" % (result.error))
 145             return
 146
 147
 148 class Backtrace:
 149     @classmethod
 150     def register_mlir_subparser(cls):
 151         cls.parser = MlirDebug.subparsers.add_parser(
 152             "backtrace", aliases=["bt"], help="Print the current backtrace"
 153         )
 154         cls.parser.set_defaults(func=cls.process_options)
 155         cls.parser.add_argument("--context", default=False, action="store_true")
 156
 157     @classmethod
 158     def process_options(cls, options, debugger, command, exe_ctx, result):
 159         frame = exe_ctx.GetFrame()
 160         if not frame.IsValid():
 161             result.SetError(
 162                 "Can't backtrace without a valid frame (program not running?)"
 163             )
 164         result = frame.EvaluateExpression(
 165             "((bool(*)(bool))mlirDebuggerPrintActionBacktrace)(%d)" % (options.context),
 166             exprOptions,
 167         )
 168         if not result.error.Success():
 169             print("Error printing breakpoints: %s" % (result.error))
 170             return
 171
 172
 173 ###############################################################################
 174 # Cursor manipulation
 175 ###############################################################################
 176
 177
 178 class PrintCursor:
 179     @classmethod
 180     def register_mlir_subparser(cls):
 181         cls.parser = MlirDebug.subparsers.add_parser(
 182             "cursor-print", aliases=["cursor-p"], help="Print the current cursor"
 183         )
 184         cls.parser.add_argument(
 185             "--print-region", "--regions", "-r", default=False, action="store_true"
 186         )
 187         cls.parser.set_defaults(func=cls.process_options)
 188
 189     @classmethod
 190     def process_options(cls, options, debugger, command, exe_ctx, result):
 191         frame = exe_ctx.GetFrame()
 192         if not frame.IsValid():
 193             result.SetError(
 194                 "Can't print cursor without a valid frame (program not running?)"
 195             )
 196         result = frame.EvaluateExpression(
 197             "((bool(*)(bool))mlirDebuggerCursorPrint)(%d)" % (options.print_region),
 198             exprOptions,
 199         )
 200         if not result.error.Success():
 201             print("Error printing cursor: %s" % (result.error))
 202             return
 203
 204
 205 class SelectCursorFromContext:
 206     @classmethod
 207     def register_mlir_subparser(cls):
 208         cls.parser = MlirDebug.subparsers.add_parser(
 209             "cursor-select-from-context",
 210             aliases=["cursor-s"],
 211             help="Select the cursor from the current context",
 212         )
 213         cls.parser.add_argument("index", type=int, help="Index in the context")
 214         cls.parser.set_defaults(func=cls.process_options)
 215
 216     @classmethod
 217     def process_options(cls, options, debugger, command, exe_ctx, result):
 218         frame = exe_ctx.GetFrame()
 219         if not frame.IsValid():
 220             result.SetError(
 221                 "Can't manipulate cursor without a valid frame (program not running?)"
 222             )
 223         result = frame.EvaluateExpression(
 224             "((bool(*)(int))mlirDebuggerCursorSelectIRUnitFromContext)(%d)"
 225             % options.index,
 226             exprOptions,
 227         )
 228         if not result.error.Success():
 229             print("Error manipulating cursor: %s" % (result.error))
 230             return
 231
 232
 233 class CursorSelectParent:
 234     @classmethod
 235     def register_mlir_subparser(cls):
 236         cls.parser = MlirDebug.subparsers.add_parser(
 237             "cursor-parent", aliases=["cursor-up"], help="Select the cursor parent"
 238         )
 239         cls.parser.set_defaults(func=cls.process_options)
 240
 241     @classmethod
 242     def process_options(cls, options, debugger, command, exe_ctx, result):
 243         frame = exe_ctx.GetFrame()
 244         if not frame.IsValid():
 245             result.SetError(
 246                 "Can't manipulate cursor without a valid frame (program not running?)"
 247             )
 248         result = frame.EvaluateExpression(
 249             "((bool(*)())mlirDebuggerCursorSelectParentIRUnit)()",
 250             exprOptions,
 251         )
 252         if not result.error.Success():
 253             print("Error manipulating cursor: %s" % (result.error))
 254             return
 255
 256
 257 class SelectCursorChild:
 258     @classmethod
 259     def register_mlir_subparser(cls):
 260         cls.parser = MlirDebug.subparsers.add_parser(
 261             "cursor-child", aliases=["cursor-c"], help="Select the nth child"
 262         )
 263         cls.parser.add_argument("index", type=int, help="Index of the child to select")
 264         cls.parser.set_defaults(func=cls.process_options)
 265
 266     @classmethod
 267     def process_options(cls, options, debugger, command, exe_ctx, result):
 268         frame = exe_ctx.GetFrame()
 269         if not frame.IsValid():
 270             result.SetError(
 271                 "Can't manipulate cursor without a valid frame (program not running?)"
 272             )
 273         result = frame.EvaluateExpression(
 274             "((bool(*)(int))mlirDebuggerCursorSelectChildIRUnit)(%d)" % options.index,
 275             exprOptions,
 276         )
 277         if not result.error.Success():
 278             print("Error manipulating cursor: %s" % (result.error))
 279             return
 280
 281
 282 class CursorSelecPrevious:
 283     @classmethod
 284     def register_mlir_subparser(cls):
 285         cls.parser = MlirDebug.subparsers.add_parser(
 286             "cursor-previous",
 287             aliases=["cursor-prev"],
 288             help="Select the cursor previous element",
 289         )
 290         cls.parser.set_defaults(func=cls.process_options)
 291
 292     @classmethod
 293     def process_options(cls, options, debugger, command, exe_ctx, result):
 294         frame = exe_ctx.GetFrame()
 295         if not frame.IsValid():
 296             result.SetError(
 297                 "Can't manipulate cursor without a valid frame (program not running?)"
 298             )
 299         result = frame.EvaluateExpression(
 300             "((bool(*)())mlirDebuggerCursorSelectPreviousIRUnit)()",
 301             exprOptions,
 302         )
 303         if not result.error.Success():
 304             print("Error manipulating cursor: %s" % (result.error))
 305             return
 306
 307
 308 class CursorSelecNext:
 309     @classmethod
 310     def register_mlir_subparser(cls):
 311         cls.parser = MlirDebug.subparsers.add_parser(
 312             "cursor-next", aliases=["cursor-n"], help="Select the cursor next element"
 313         )
 314         cls.parser.set_defaults(func=cls.process_options)
 315
 316     @classmethod
 317     def process_options(cls, options, debugger, command, exe_ctx, result):
 318         frame = exe_ctx.GetFrame()
 319         if not frame.IsValid():
 320             result.SetError(
 321                 "Can't manipulate cursor without a valid frame (program not running?)"
 322             )
 323         result = frame.EvaluateExpression(
 324             "((bool(*)())mlirDebuggerCursorSelectNextIRUnit)()",
 325             exprOptions,
 326         )
 327         if not result.error.Success():
 328             print("Error manipulating cursor: %s" % (result.error))
 329             return
 330
 331
 332 ###############################################################################
 333 # Breakpoints
 334 ###############################################################################
 335
 336
 337 class EnableBreakpoint:
 338     @classmethod
 339     def register_mlir_subparser(cls):
 340         cls.parser = MlirDebug.subparsers.add_parser(
 341             "enable", help="Enable a single breakpoint (given its ID)"
 342         )
 343         cls.parser.add_argument("id", help="ID of the breakpoint to enable")
 344         cls.parser.set_defaults(func=cls.process_options)
 345
 346     @classmethod
 347     def process_options(cls, options, debugger, command, exe_ctx, result):
 348         bp = breakpoints.get(int(options.id), None)
 349         if not bp:
 350             result.SetError("No breakpoint with ID %d" % int(options.id))
 351             return
 352         bp.enable(exe_ctx.GetFrame())
 353
 354
 355 class DisableBreakpoint:
 356     @classmethod
 357     def register_mlir_subparser(cls):
 358         cls.parser = MlirDebug.subparsers.add_parser(
 359             "disable", help="Disable a single breakpoint (given its ID)"
 360         )
 361         cls.parser.add_argument("id", help="ID of the breakpoint to disable")
 362         cls.parser.set_defaults(func=cls.process_options)
 363
 364     @classmethod
 365     def process_options(cls, options, debugger, command, exe_ctx, result):
 366         bp = breakpoints.get(int(options.id), None)
 367         if not bp:
 368             result.SetError("No breakpoint with ID %s" % options.id)
 369             return
 370         bp.disable(exe_ctx.GetFrame())
 371
 372
 373 class ListBreakpoints:
 374     @classmethod
 375     def register_mlir_subparser(cls):
 376         cls.parser = MlirDebug.subparsers.add_parser(
 377             "list", help="List all current breakpoints"
 378         )
 379         cls.parser.set_defaults(func=cls.process_options)
 380
 381     @classmethod
 382     def process_options(cls, options, debugger, command, exe_ctx, result):
 383         for id, bp in sorted(breakpoints.items()):
 384             print(id, type(id), str(bp), "enabled" if bp.isEnabled else "disabled")
 385
 386
 387 class Breakpoint:
 388     def __init__(self):
 389         global nextid
 390         self.id = nextid
 391         nextid += 1
 392         breakpoints[self.id] = self
 393         self.isEnabled = True
 394
 395     def enable(self, frame=None):
 396         self.isEnabled = True
 397         if not frame or not frame.IsValid():
 398             return
 399         # use a C cast to force the type of the breakpoint handle to be void * so
 400         # that we don't rely on DWARF. Also add a fake bool return value otherwise
 401         # LLDB can't signal any error with the expression evaluation (at least I don't know how).
 402         cmd = (
 403             "((bool (*)(void *))mlirDebuggerEnableBreakpoint)((void *)%s)" % self.handle
 404         )
 405         result = frame.EvaluateExpression(cmd, exprOptions)
 406         if not result.error.Success():
 407             print("Error enabling breakpoint: %s" % (result.error))
 408             return
 409
 410     def disable(self, frame=None):
 411         self.isEnabled = False
 412         if not frame or not frame.IsValid():
 413             return
 414         # use a C cast to force the type of the breakpoint handle to be void * so
 415         # that we don't rely on DWARF. Also add a fake bool return value otherwise
 416         # LLDB can't signal any error with the expression evaluation (at least I don't know how).
 417         cmd = (
 418             "((bool (*)(void *)) mlirDebuggerDisableBreakpoint)((void *)%s)"
 419             % self.handle
 420         )
 421         result = frame.EvaluateExpression(cmd, exprOptions)
 422         if not result.error.Success():
 423             print("Error disabling breakpoint: %s" % (result.error))
 424             return
 425
 426
 427 class TagBreakpoint(Breakpoint):
 428     mlir_subcommand = "break-on-tag"
 429
 430     def __init__(self, tag):
 431         super().__init__()
 432         self.tag = tag
 433
 434     def __str__(self):
 435         return "[%d] TagBreakpoint(%s)" % (self.id, self.tag)
 436
 437     @classmethod
 438     def register_mlir_subparser(cls):
 439         cls.parser = MlirDebug.subparsers.add_parser(
 440             cls.mlir_subcommand, help="add a breakpoint on actions' tag matching"
 441         )
 442         cls.parser.set_defaults(func=cls.process_options)
 443         cls.parser.add_argument("tag", help="tag to match")
 444
 445     @classmethod
 446     def process_options(cls, options, debugger, command, exe_ctx, result):
 447         breakpoint = TagBreakpoint(options.tag)
 448         print("Added breakpoint %s" % str(breakpoint))
 449
 450         frame = exe_ctx.GetFrame()
 451         if frame.IsValid():
 452             breakpoint.install(frame)
 453
 454     def install(self, frame):
 455         result = frame.EvaluateExpression(
 456             '((void *(*)(const char *))mlirDebuggerAddTagBreakpoint)("%s")'
 457             % (self.tag),
 458             exprOptions,
 459         )
 460         if not result.error.Success():
 461             print("Error installing breakpoint: %s" % (result.error))
 462             return
 463         # Save the handle, this is necessary to implement enable/disable.
 464         self.handle = result.GetValue()
 465
 466
 467 class FileLineBreakpoint(Breakpoint):
 468     mlir_subcommand = "break-on-file"
 469
 470     def __init__(self, file, line, col):
 471         super().__init__()
 472         self.file = file
 473         self.line = line
 474         self.col = col
 475
 476     def __str__(self):
 477         return "[%d] FileLineBreakpoint(%s, %d, %d)" % (
 478             self.id,
 479             self.file,
 480             self.line,
 481             self.col,
 482         )
 483
 484     @classmethod
 485     def register_mlir_subparser(cls):
 486         cls.parser = MlirDebug.subparsers.add_parser(
 487             cls.mlir_subcommand,
 488             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",
 489         )
 490         cls.parser.set_defaults(func=cls.process_options)
 491         cls.parser.add_argument("location", type=str)
 492
 493     @classmethod
 494     def process_options(cls, options, debugger, command, exe_ctx, result):
 495         split_loc = options.location.split(":")
 496         file = split_loc[0]
 497         line = int(split_loc[1]) if len(split_loc) > 1 else -1
 498         col = int(split_loc[2]) if len(split_loc) > 2 else -1
 499         breakpoint = FileLineBreakpoint(file, line, col)
 500         print("Added breakpoint %s" % str(breakpoint))
 501
 502         frame = exe_ctx.GetFrame()
 503         if frame.IsValid():
 504             breakpoint.install(frame)
 505
 506     def install(self, frame):
 507         result = frame.EvaluateExpression(
 508             '((void *(*)(const char *, int, int))mlirDebuggerAddFileLineColLocBreakpoint)("%s", %d, %d)'
 509             % (self.file, self.line, self.col),
 510             exprOptions,
 511         )
 512         if not result.error.Success():
 513             print("Error installing breakpoint: %s" % (result.error))
 514             return
 515         # Save the handle, this is necessary to implement enable/disable.
 516         self.handle = result.GetValue()
 517
 518
 519 def on_start(frame, bpno, err):
 520     print("MLIR debugger attaching...")
 521     for _, bp in sorted(breakpoints.items()):
 522         if bp.isEnabled:
 523             print("Installing breakpoint %s" % (str(bp)))
 524             bp.install(frame)
 525         else:
 526             print("Skipping disabled breakpoint %s" % (str(bp)))
 527
 528     return True
 529
 530
 531 def __lldb_init_module(debugger, dict):
 532     target = debugger.GetTargetAtIndex(0)
 533     debugger.SetAsync(False)
 534     if not target:
 535         print("No target is loaded, please load a target before loading this script.")
 536         return
 537     if debugger.GetNumTargets() > 1:
 538         print(
 539             "Multiple targets (%s) loaded, attaching MLIR debugging to %s"
 540             % (debugger.GetNumTargets(), target)
 541         )
 542
 543     # Register all classes that have a register_lldb_command method
 544     module_name = __name__
 545     parser = MlirDebug.create_options()
 546     MlirDebug.__doc__ = parser.format_help()
 547
 548     # Add the MLIR entry point to LLDB as a command.
 549     command = "command script add -o -c %s.%s %s" % (
 550         module_name,
 551         MlirDebug.__name__,
 552         MlirDebug.lldb_command,
 553     )
 554     debugger.HandleCommand(command)
 555
 556     main_bp = target.BreakpointCreateByName("main")
 557     main_bp.SetScriptCallbackFunction("action_debugging.on_start")
 558     main_bp.SetAutoContinue(auto_continue=True)
 559
 560     on_breackpoint = target.BreakpointCreateByName("mlirDebuggerBreakpointHook")
 561
 562     print(
 563         'The "{0}" command has been installed for target `{1}`, type "help {0}" or "{0} '
 564         '--help" for detailed help.'.format(MlirDebug.lldb_command, target)
 565     )
 566     for _name, cls in inspect.getmembers(sys.modules[module_name]):
 567         if inspect.isclass(cls) and getattr(cls, "register_mlir_subparser", None):
 568             cls.register_mlir_subparser()