llvm-project

1
"""GDB pretty printers for MLIR types."""
2

3
import gdb.printing
4

5

6
class StoragePrinter:
7
    """Prints bases of a struct and its fields."""
8

9
    def __init__(self, val):
10
        self.val = val
11

12
    def children(self):
13
        for field in self.val.type.fields():
14
            if field.is_base_class:
15
                yield "<%s>" % field.name, self.val.cast(field.type)
16
            else:
17
                yield field.name, self.val[field.name]
18

19
    def to_string(self):
20
        return "mlir::Storage"
21

22

23
class TupleTypeStoragePrinter(StoragePrinter):
24
    def children(self):
25
        for child in StoragePrinter.children(self):
26
            yield child
27
        pointer_type = gdb.lookup_type("mlir::Type").pointer()
28
        elements = (self.val.address + 1).cast(pointer_type)
29
        for i in range(self.val["numElements"]):
30
            yield "elements[%u]" % i, elements[i]
31

32
    def to_string(self):
33
        return "mlir::TupleTypeStorage of %u elements" % self.val["numElements"]
34

35

36
class FusedLocationStoragePrinter(StoragePrinter):
37
    def children(self):
38
        for child in StoragePrinter.children(self):
39
            yield child
40
        pointer_type = gdb.lookup_type("mlir::Location").pointer()
41
        elements = (self.val.address + 1).cast(pointer_type)
42
        for i in range(self.val["numLocs"]):
43
            yield "locs[%u]" % i, elements[i]
44

45
    def to_string(self):
46
        return "mlir::FusedLocationStorage of %u locs" % self.val["numLocs"]
47

48

49
class StorageTypeMap:
50
    """Maps a TypeID to the corresponding concrete type.
51

52
    Types need to be registered by name before the first lookup.
53
    """
54

55
    def __init__(self):
56
        self.map = None
57
        self.type_names = []
58

59
    def register_type(self, type_name):
60
        assert not self.map, "register_type called after __getitem__"
61
        self.type_names += [type_name]
62

63
    def _init_map(self):
64
        """Lazy initialization  of self.map."""
65
        if self.map:
66
            return
67
        self.map = {}
68
        for type_name in self.type_names:
69
            concrete_type = gdb.lookup_type(type_name)
70
            try:
71
                storage = gdb.parse_and_eval(
72
                    "&'mlir::detail::TypeIDExported::get<%s>()::instance'" % type_name
73
                )
74
            except gdb.error:
75
                # Skip when TypeID instance cannot be found in current context.
76
                continue
77
            if concrete_type and storage:
78
                self.map[int(storage)] = concrete_type
79

80
    def __getitem__(self, type_id):
81
        self._init_map()
82
        return self.map.get(int(type_id["storage"]))
83

84

85
storage_type_map = StorageTypeMap()
86

87

88
def get_type_id_printer(val):
89
    """Returns a printer of the name of a mlir::TypeID."""
90

91
    class TypeIdPrinter:
92
        def __init__(self, string):
93
            self.string = string
94

95
        def to_string(self):
96
            return self.string
97

98
    concrete_type = storage_type_map[val]
99
    if not concrete_type:
100
        return None
101
    return TypeIdPrinter("mlir::TypeID::get<%s>()" % concrete_type)
102

103

104
def get_attr_or_type_printer(val, get_type_id):
105
    """Returns a printer for mlir::Attribute or mlir::Type."""
106

107
    class AttrOrTypePrinter:
108
        def __init__(self, type_id, impl):
109
            self.type_id = type_id
110
            self.impl = impl
111

112
        def children(self):
113
            yield "typeID", self.type_id
114
            yield "impl", self.impl
115

116
        def to_string(self):
117
            return "cast<%s>" % self.impl.type
118

119
    if not val["impl"]:
120
        return None
121
    impl = val["impl"].dereference()
122
    type_id = get_type_id(impl)
123
    concrete_type = storage_type_map[type_id]
124
    if not concrete_type:
125
        return None
126
    # 3rd template argument of StorageUserBase is the storage type.
127
    storage_type = concrete_type.fields()[0].type.template_argument(2)
128
    if not storage_type:
129
        return None
130
    return AttrOrTypePrinter(type_id, impl.cast(storage_type))
131

132

133
class ImplPrinter:
134
    """Printer for an instance with a single 'impl' member pointer."""
135

136
    def __init__(self, val):
137
        self.val = val
138
        self.impl = val["impl"]
139

140
    def children(self):
141
        if self.impl:
142
            yield "impl", self.impl.dereference()
143

144
    def to_string(self):
145
        return self.val.type.name
146

147

148
# Printers of types deriving from Attribute::AttrBase or Type::TypeBase.
149
for name in [
150
    # mlir/IR/Attributes.h
151
    "ArrayAttr",
152
    "DictionaryAttr",
153
    "FloatAttr",
154
    "IntegerAttr",
155
    "IntegerSetAttr",
156
    "OpaqueAttr",
157
    "StringAttr",
158
    "SymbolRefAttr",
159
    "TypeAttr",
160
    "UnitAttr",
161
    "DenseStringElementsAttr",
162
    "DenseIntOrFPElementsAttr",
163
    "SparseElementsAttr",
164
    # mlir/IR/BuiltinTypes.h
165
    "ComplexType",
166
    "IndexType",
167
    "IntegerType",
168
    "Float16Type",
169
    "FloatTF32Type",
170
    "Float32Type",
171
    "Float64Type",
172
    "Float80Type",
173
    "Float128Type",
174
    "NoneType",
175
    "VectorType",
176
    "RankedTensorType",
177
    "UnrankedTensorType",
178
    "MemRefType",
179
    "UnrankedMemRefType",
180
    "TupleType",
181
    # mlir/IR/Location.h
182
    "CallSiteLoc",
183
    "FileLineColLoc",
184
    "FusedLoc",
185
    "NameLoc",
186
    "OpaqueLoc",
187
    "UnknownLoc",
188
]:
189
    storage_type_map.register_type("mlir::%s" % name)  # Register for upcasting.
190
storage_type_map.register_type("void")  # Register default.
191

192

193
pp = gdb.printing.RegexpCollectionPrettyPrinter("MLIRSupport")
194

195
pp.add_printer("mlir::OperationName", "^mlir::OperationName$", ImplPrinter)
196
pp.add_printer("mlir::Value", "^mlir::Value$", ImplPrinter)
197

198
# Printers for types deriving from AttributeStorage or TypeStorage.
199
pp.add_printer(
200
    "mlir::detail::FusedLocationStorage",
201
    "^mlir::detail::FusedLocationStorage",
202
    FusedLocationStoragePrinter,
203
)
204
pp.add_printer(
205
    "mlir::detail::TupleTypeStorage",
206
    "^mlir::detail::TupleTypeStorage$",
207
    TupleTypeStoragePrinter,
208
)
209

210
pp.add_printer("mlir::TypeID", "^mlir::TypeID$", get_type_id_printer)
211

212

213
def add_attr_or_type_printers(name):
214
    """Adds printers for mlir::Attribute or mlir::Type and their Storage type."""
215
    get_type_id = lambda val: val["abstract%s" % name]["typeID"]
216
    pp.add_printer(
217
        "mlir::%s" % name,
218
        "^mlir::%s$" % name,
219
        lambda val: get_attr_or_type_printer(val, get_type_id),
220
    )
221

222

223
# Upcasting printers of mlir::Attribute and mlir::Type.
224
for name in ["Attribute", "Type"]:
225
    add_attr_or_type_printers(name)
226

227
gdb.printing.register_pretty_printer(gdb.current_objfile(), pp)
228