1"""GDB pretty printers for MLIR types."""
2
3import gdb.printing
4
5
6class 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
22class TupleTypeStoragePrinter(StoragePrinter):
23
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
35class FusedLocationStoragePrinter(StoragePrinter):
36
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
49class 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      except gdb.error:
74        # Skip when TypeID instance cannot be found in current context.
75        continue
76      if concrete_type and storage:
77        self.map[int(storage)] = concrete_type
78
79  def __getitem__(self, type_id):
80    self._init_map()
81    return self.map.get(int(type_id['storage']))
82
83
84storage_type_map = StorageTypeMap()
85
86
87def get_type_id_printer(val):
88  """Returns a printer of the name of a mlir::TypeID."""
89
90  class TypeIdPrinter:
91
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
104def get_attr_or_type_printer(val, get_type_id):
105  """Returns a printer for mlir::Attribute or mlir::Type."""
106
107  class AttrOrTypePrinter:
108
109    def __init__(self, type_id, impl):
110      self.type_id = type_id
111      self.impl = impl
112
113    def children(self):
114      yield 'typeID', self.type_id
115      yield 'impl', self.impl
116
117    def to_string(self):
118      return 'cast<%s>' % self.impl.type
119
120  if not val['impl']:
121    return None
122  impl = val['impl'].dereference()
123  type_id = get_type_id(impl)
124  concrete_type = storage_type_map[type_id]
125  if not concrete_type:
126    return None
127  # 3rd template argument of StorageUserBase is the storage type.
128  storage_type = concrete_type.fields()[0].type.template_argument(2)
129  if not storage_type:
130    return None
131  return AttrOrTypePrinter(type_id, impl.cast(storage_type))
132
133
134class ImplPrinter:
135  """Printer for an instance with a single 'impl' member pointer."""
136
137  def __init__(self, val):
138    self.val = val
139    self.impl = val['impl']
140
141  def children(self):
142    if self.impl:
143      yield 'impl', self.impl.dereference()
144
145  def to_string(self):
146    return self.val.type.name
147
148
149# Printers of types deriving from Attribute::AttrBase or Type::TypeBase.
150for name in [
151    # mlir/IR/Attributes.h
152    'ArrayAttr',
153    'DictionaryAttr',
154    'FloatAttr',
155    'IntegerAttr',
156    'IntegerSetAttr',
157    'OpaqueAttr',
158    'StringAttr',
159    'SymbolRefAttr',
160    'TypeAttr',
161    'UnitAttr',
162    'DenseStringElementsAttr',
163    'DenseIntOrFPElementsAttr',
164    'OpaqueElementsAttr',
165    'SparseElementsAttr',
166    # mlir/IR/BuiltinTypes.h
167    'ComplexType',
168    'IndexType',
169    'IntegerType',
170    'Float16Type',
171    'Float32Type',
172    'Float64Type',
173    'Float80Type',
174    'Float128Type',
175    'NoneType',
176    'VectorType',
177    'RankedTensorType',
178    'UnrankedTensorType',
179    'MemRefType',
180    'UnrankedMemRefType',
181    'TupleType',
182    # mlir/IR/Location.h
183    'CallSiteLoc',
184    'FileLineColLoc',
185    'FusedLoc',
186    'NameLoc',
187    'OpaqueLoc',
188    'UnknownLoc'
189]:
190  storage_type_map.register_type('mlir::%s' % name)  # Register for upcasting.
191storage_type_map.register_type('void')  # Register default.
192
193
194pp = gdb.printing.RegexpCollectionPrettyPrinter('MLIRSupport')
195
196pp.add_printer('mlir::OperationName', '^mlir::OperationName$', ImplPrinter)
197pp.add_printer('mlir::Value', '^mlir::Value$', ImplPrinter)
198
199# Printers for types deriving from AttributeStorage or TypeStorage.
200pp.add_printer('mlir::detail::FusedLocationStorage',
201               '^mlir::detail::FusedLocationStorage',
202               FusedLocationStoragePrinter)
203pp.add_printer('mlir::detail::TupleTypeStorage',
204               '^mlir::detail::TupleTypeStorage$', TupleTypeStoragePrinter)
205
206pp.add_printer('mlir::TypeID', '^mlir::TypeID$', get_type_id_printer)
207
208
209def add_attr_or_type_printers(name):
210  """Adds printers for mlir::Attribute or mlir::Type and their Storage type."""
211  get_type_id = lambda val: val['abstract%s' % name]['typeID']
212  pp.add_printer('mlir::%s' % name, '^mlir::%s$' % name,
213                 lambda val: get_attr_or_type_printer(val, get_type_id))
214
215
216# Upcasting printers of mlir::Attribute and mlir::Type.
217for name in ['Attribute', 'Type']:
218  add_attr_or_type_printers(name)
219
220gdb.printing.register_pretty_printer(gdb.current_objfile(), pp)
221