1"""Convert a PETSc XML file into a Flame Graph input file.""" 2 3import argparse 4import os 5import sys 6 7try: 8 from lxml import objectify 9except ImportError: 10 sys.exit("Import error: lxml must be installed. Try 'pip install lxml'.") 11 12 13def parse_time(event): 14 # The time can either be stored under 'value' or 'avgvalue' 15 if hasattr(event.time, "value"): 16 return event.time.value 17 elif hasattr(event.time, "avgvalue"): 18 return event.time.avgvalue 19 else: 20 raise AssertionError 21 22 23def make_line(callstack, time, total_time): 24 """The output time needs to be an integer for the file to be 25 accepted by speedscope (speedscope.app). Therefore we output it in 26 microseconds. It is originally a percentage of the total time 27 (given in seconds). 28 """ 29 event_str = ";".join(str(event.name) for event in callstack) 30 time_us = int(time / 100 * total_time * 1e6) 31 return f"{event_str} {time_us}" 32 33 34def traverse_children(parent, total_time, callstack=None): 35 if callstack == None: 36 callstack = [] 37 38 # Sort the events into 'self' and child events 39 self_events, child_events = [], [] 40 for event in parent.event: 41 if event.name == "self" or str(event.name).endswith("other-timed"): 42 self_events.append(event) 43 else: 44 child_events.append(event) 45 46 lines = [] 47 if self_events: 48 time = sum(parse_time(event) for event in self_events) 49 lines.append(make_line(callstack, time, total_time)) 50 51 for event in child_events: 52 # Check to see if event has any children. The latter check is for the 53 # case when the <events> tag is present but empty. 54 if hasattr(event, "events") and hasattr(event.events, "event"): 55 callstack.append(event) 56 lines.extend(traverse_children(event.events, total_time, callstack)) 57 callstack.pop() 58 else: 59 time = parse_time(event) 60 lines.append(make_line(callstack+[event], time, total_time)) 61 return lines 62 63 64def parse_args(): 65 parser = argparse.ArgumentParser() 66 parser.add_argument("infile", type=str, help="Input XML file") 67 parser.add_argument("outfile", type=str, help="Output file") 68 return parser.parse_args() 69 70 71def check_args(args): 72 if not args.infile.endswith(".xml"): 73 raise ValueError("Input file must be an XML file.") 74 if not os.path.exists(args.infile): 75 raise ValueError("The input file does not exist.") 76 77 78def main(): 79 args = parse_args() 80 check_args(args) 81 82 root = objectify.parse(args.infile).find("//timertree") 83 total_time = root.find("totaltime") 84 lines = traverse_children(root, total_time) 85 86 with open(args.outfile, "w") as f: 87 f.write("\n".join(lines)) 88 89 90if __name__ == "__main__": 91 main() 92