llvm-project

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()
569