• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1"""GDB pretty printers for MLIR types."""
2
3import gdb.printing
4
5
6class IdentifierPrinter:
7  """Prints an mlir::Identifier instance."""
8
9  def __init__(self, val):
10    self.entry = val['entry']
11
12  def to_string(self):
13    ptr = (self.entry + 1).cast(gdb.lookup_type('char').pointer())
14    return ptr.string(length=self.entry['keyLength'])
15
16  def display_hint(self):
17    return 'string'
18
19
20class StoragePrinter:
21  """Prints bases of a struct and its fields."""
22
23  def __init__(self, val):
24    self.val = val
25
26  def children(self):
27    for field in self.val.type.fields():
28      if field.is_base_class:
29        yield ('<%s>' % field.name, self.val.cast(field.type))
30      else:
31        yield (field.name, self.val[field.name])
32
33
34class TupleTypeStoragePrinter(StoragePrinter):
35
36  def children(self):
37    for child in StoragePrinter.children(self):
38      yield child
39    pointer_type = gdb.lookup_type('mlir::Type').pointer()
40    elements = (self.val.address + 1).cast(pointer_type)
41    for i in range(self.val['numElements']):
42      yield 'elements[%u]' % i, elements[i]
43
44
45class RankedTypeStoragePrinter(StoragePrinter):
46
47  def children(self):
48    for child in StoragePrinter.children(self):
49      yield child
50    for i in range(self.val['shapeSize']):
51      yield 'shapeElements[%u]' % i, self.val['shapeElements'][i]
52
53
54class MemRefTypeStoragePrinter(RankedTypeStoragePrinter):
55
56  def children(self):
57    for child in RankedTypeStoragePrinter.children(self):
58      yield child
59    for i in range(self.val['numAffineMaps']):
60      yield 'affineMapsList[%u]' % i, self.val['affineMapsList'][i]
61
62
63class FusedLocationStoragePrinter(StoragePrinter):
64
65  def children(self):
66    for child in StoragePrinter.children(self):
67      yield child
68    pointer_type = gdb.lookup_type('mlir::Location').pointer()
69    elements = (self.val.address + 1).cast(pointer_type)
70    for i in range(self.val['numLocs']):
71      yield 'locs[%u]' % i, elements[i]
72
73
74class StorageUserBasePrinter:
75  """Printer for an mlir::detail::StorageUserBase instance."""
76
77  def __init__(self, val):
78    self.val = val
79
80  def children(self):
81    storage_type = self.val.type.template_argument(2)
82    yield 'impl', self.val['impl'].dereference().cast(storage_type)
83
84
85class StorageTypeMap:
86  """Maps a TypeID to the corresponding type derived from StorageUserBase.
87
88  Types need to be registered by name before the first lookup.
89  """
90
91  def __init__(self):
92    self.map = None
93    self.type_names = []
94
95  def register_type(self, type_name):
96    assert not self.map, 'register_type called after __getitem__'
97    self.type_names += [type_name]
98
99  def _init_map(self):
100    """Lazy initialization  of self.map."""
101    if self.map:
102      return
103    self.map = {}
104    for type_name in self.type_names:
105      concrete_type = gdb.lookup_type(type_name)
106      try:
107        storage = gdb.parse_and_eval(
108            "&'mlir::detail::TypeIDExported::get<%s>()::instance'" % type_name)
109      except gdb.error:
110        # Skip when TypeID instance cannot be found in current context.
111        continue
112      if concrete_type and storage:
113        self.map[int(storage)] = concrete_type
114
115  def __getitem__(self, type_id):
116    self._init_map()
117    return self.map.get(int(type_id['storage']))
118
119
120storage_type_map = StorageTypeMap()
121
122
123def get_type_id_printer(val):
124  """Returns a printer of the name of a mlir::TypeID."""
125
126  class StringPrinter:
127
128    def __init__(self, string):
129      self.string = string
130
131    def to_string(self):
132      return self.string
133
134  concrete_type = storage_type_map[val]
135  if not concrete_type:
136    return None
137  return StringPrinter('"%s"' % concrete_type.name)
138
139
140def get_attr_or_type_printer(val, get_type_id):
141  """Returns a printer for mlir::Attribute or mlir::Type."""
142
143  class UpcastPrinter:
144
145    def __init__(self, val, type):
146      self.val = val.cast(type)
147
148    def children(self):
149      yield 'cast<%s>' % self.val.type.name, self.val
150
151  if not val['impl']:
152    return None
153  type_id = get_type_id(val['impl'].dereference())
154  concrete_type = storage_type_map[type_id]
155  if not concrete_type:
156    return None
157  return UpcastPrinter(val, concrete_type)
158
159
160pp = gdb.printing.RegexpCollectionPrettyPrinter('MLIRSupport')
161
162pp.add_printer('mlir::Identifier', '^mlir::Identifier$', IdentifierPrinter)
163
164# Printers for types deriving from AttributeStorage or TypeStorage.
165pp.add_printer('mlir::detail::FusedLocationStorage',
166               '^mlir::detail::FusedLocationStorage',
167               FusedLocationStoragePrinter)
168pp.add_printer('mlir::detail::VectorTypeStorage',
169               '^mlir::detail::VectorTypeStorage', RankedTypeStoragePrinter)
170pp.add_printer('mlir::detail::RankedTensorTypeStorage',
171               '^mlir::detail::RankedTensorTypeStorage',
172               RankedTypeStoragePrinter)
173pp.add_printer('mlir::detail::MemRefTypeStorage',
174               '^mlir::detail::MemRefTypeStorage$', MemRefTypeStoragePrinter)
175pp.add_printer('mlir::detail::TupleTypeStorage',
176               '^mlir::detail::TupleTypeStorage$', TupleTypeStoragePrinter)
177
178# Printers for Attribute::AttrBase or Type::TypeBase typedefs.
179pp.add_printer('mlir::detail::StorageUserBase',
180               '^mlir::detail::StorageUserBase<.*>$', StorageUserBasePrinter)
181
182# Printers of types deriving from Attribute::AttrBase or Type::TypeBase.
183for name in [
184    # mlir/IR/Attributes.h
185    'ArrayAttr',
186    'DictionaryAttr',
187    'FloatAttr',
188    'IntegerAttr',
189    'IntegerSetAttr',
190    'OpaqueAttr',
191    'StringAttr',
192    'SymbolRefAttr',
193    'TypeAttr',
194    'UnitAttr',
195    'DenseStringElementsAttr',
196    'DenseIntOrFPElementsAttr',
197    'OpaqueElementsAttr',
198    'SparseElementsAttr',
199    # mlir/IR/BuiltinTypes.h
200    'ComplexType',
201    'IndexType',
202    'IntegerType',
203    'Float16Type',
204    'Float32Type',
205    'Float64Type',
206    'NoneType',
207    'VectorType',
208    'RankedTensorType',
209    'UnrankedTensorType',
210    'MemRefType',
211    'UnrankedMemRefType',
212    'TupleType',
213    # mlir/IR/Location.h
214    'CallSiteLoc',
215    'FileLineColLoc',
216    'FusedLoc',
217    'NameLoc',
218    'OpaqueLoc',
219    'UnknownLoc'
220]:
221  storage_type_map.register_type('mlir::%s' % name)  # Register for upcasting.
222
223pp.add_printer('mlir::TypeID', '^mlir::TypeID$', get_type_id_printer)
224
225
226def add_attr_or_type_printers(name):
227  """Adds printers for mlir::Attribute or mlir::Type and their Storage type."""
228  get_type_id = lambda val: val['abstract%s' % name]['typeID']
229  pp.add_printer('mlir::%s' % name, '^mlir::%s$' % name,
230                 lambda val: get_attr_or_type_printer(val, get_type_id))
231  pp.add_printer('mlir::%sStorage' % name, '^mlir::%sStorage$' % name,
232                 lambda val: get_type_id_printer(get_type_id(val)))
233
234
235# Upcasting printers of mlir::Attribute and mlir::Type.
236for name in ['Attribute', 'Type']:
237  add_attr_or_type_printers(name)
238
239gdb.printing.register_pretty_printer(gdb.current_objfile(), pp)
240