Function parse_prof_blocks

profiler/parse_profiler_results.py:125–184

(
    prof_blocks: Dict[str, List[Tuple[List[ProfileData], List[MemAllocation]]]],
    allocator_dict: Dict[int, str],
    time_scale: TimeScale,
)

def parse_prof_blocks(
    prof_blocks: Dict[str, List[Tuple[List[ProfileData], List[MemAllocation]]]],
    allocator_dict: Dict[int, str],
    time_scale: TimeScale,
) -> Tuple[Dict[str, List[ProfileEvent]], Dict[str, List[MemEvent]]]:

    prof_data = OrderedDict()
    mem_prof_data = OrderedDict()

    # Iterate through all the profiling blocks data that have been grouped by name.
    for name, data_list in prof_blocks.items():
        prof_data_list: List[ProfileEvent] = []
        mem_prof_data_list: List[MemAllocation] = []
        # Each entry in data_list is a tuple in which the first entry is profiling data
        # and the second entry is memory allocation data, also each entry in data_list
        # represents one iteration of a code block.
        for i in range(len(data_list)):
            for idx, event in enumerate(data_list[i][0]):
                # If the event represented by the index idx already exists in the list
                # then just append the new time entry to the duration list present in
                # the event object. If it doesn't exist then create a new entry and add
                # it to the list.
                if idx < len(prof_data_list):
                    start_time, duration = adjust_time_scale(event, time_scale)
                    prof_data_list[idx].ts.append(start_time)
                    prof_data_list[idx].duration.append(duration)
                else:
                    start_time, duration = adjust_time_scale(event, time_scale)
                    prof_data_list.append(
                        ProfileEvent(
                            event.name,
                            [start_time],
                            [duration],
                            event.chain_idx,
                            event.instruction_idx,
                        )
                    )

            # Collect all the memory allocation events of this iteration of the code block
            for idx, event in enumerate(data_list[i][1]):
                if idx >= len(mem_prof_data_list):
                    mem_prof_data_list.append(event)

        # Group all the memory allocation events based on the allocator they were
        # allocated from.
        alloc_sum_dict: OrderedDict[int, int] = OrderedDict()
        for alloc in mem_prof_data_list:
            alloc_sum_dict[alloc.allocator_id] = (
                alloc_sum_dict.get(alloc.allocator_id, 0) + alloc.allocation_size
            )

        mem_prof_sum_list: List[MemEvent] = []
        for allocator_id, allocation_size in alloc_sum_dict.items():
            mem_prof_sum_list.append(
                MemEvent(allocator_dict[allocator_id], allocation_size)
            )
        prof_data[name] = prof_data_list
        mem_prof_data[name] = mem_prof_sum_list