xref: /llvm-project/mlir/utils/lldb-scripts/action_debugging.py (revision 1020150e7a6f6d6f833c232125c5ab817c03c76b)
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
12import inspect
13import lldb
14import argparse
15import shlex
16import sys
17
18# Each new breakpoint gets a unique ID starting from 1.
19nextid = 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.
22breakpoints = dict()
23
24exprOptions = lldb.SBExpressionOptions()
25exprOptions.SetIgnoreBreakpoints()
26exprOptions.SetLanguage(lldb.eLanguageTypeC)
27
28
29class 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
82class 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
126class 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
148class 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
178class 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
205class 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
233class 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
257class 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
282class 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
308class 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
337class 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
355class 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
373class 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
387class 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
427class 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
467class 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
519def 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
531def __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()
569