228 lines · python
1"""GDB pretty printers for MLIR types."""2 3import gdb.printing4 5 6class StoragePrinter:7 """Prints bases of a struct and its fields."""8 9 def __init__(self, val):10 self.val = val11 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 23class TupleTypeStoragePrinter(StoragePrinter):24 def children(self):25 for child in StoragePrinter.children(self):26 yield child27 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 36class FusedLocationStoragePrinter(StoragePrinter):37 def children(self):38 for child in StoragePrinter.children(self):39 yield child40 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 = None57 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 return67 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_name73 )74 except gdb.error:75 # Skip when TypeID instance cannot be found in current context.76 continue77 if concrete_type and storage:78 self.map[int(storage)] = concrete_type79 80 def __getitem__(self, type_id):81 self._init_map()82 return self.map.get(int(type_id["storage"]))83 84 85storage_type_map = StorageTypeMap()86 87 88def 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 = string94 95 def to_string(self):96 return self.string97 98 concrete_type = storage_type_map[val]99 if not concrete_type:100 return None101 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 def __init__(self, type_id, impl):109 self.type_id = type_id110 self.impl = impl111 112 def children(self):113 yield "typeID", self.type_id114 yield "impl", self.impl115 116 def to_string(self):117 return "cast<%s>" % self.impl.type118 119 if not val["impl"]:120 return None121 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 None126 # 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 None130 return AttrOrTypePrinter(type_id, impl.cast(storage_type))131 132 133class ImplPrinter:134 """Printer for an instance with a single 'impl' member pointer."""135 136 def __init__(self, val):137 self.val = val138 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.name146 147 148# Printers of types deriving from Attribute::AttrBase or Type::TypeBase.149for name in [150 # mlir/IR/Attributes.h151 "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.h165 "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.h182 "CallSiteLoc",183 "FileLineColLoc",184 "FusedLoc",185 "NameLoc",186 "OpaqueLoc",187 "UnknownLoc",188]:189 storage_type_map.register_type("mlir::%s" % name) # Register for upcasting.190storage_type_map.register_type("void") # Register default.191 192 193pp = gdb.printing.RegexpCollectionPrettyPrinter("MLIRSupport")194 195pp.add_printer("mlir::OperationName", "^mlir::OperationName$", ImplPrinter)196pp.add_printer("mlir::Value", "^mlir::Value$", ImplPrinter)197 198# Printers for types deriving from AttributeStorage or TypeStorage.199pp.add_printer(200 "mlir::detail::FusedLocationStorage",201 "^mlir::detail::FusedLocationStorage",202 FusedLocationStoragePrinter,203)204pp.add_printer(205 "mlir::detail::TupleTypeStorage",206 "^mlir::detail::TupleTypeStorage$",207 TupleTypeStoragePrinter,208)209 210pp.add_printer("mlir::TypeID", "^mlir::TypeID$", get_type_id_printer)211 212 213def 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.224for name in ["Attribute", "Type"]:225 add_attr_or_type_printers(name)226 227gdb.printing.register_pretty_printer(gdb.current_objfile(), pp)228