Function calc_mean_ranges

bench/python/benchmark.py:351–392

Calculates the mean of all NVTX ranges present in the all_range_info. Since NVTX ranges can be reported per process, per thread, we need to have a way to average those numbers. The mean here is computed by taking the average of all the numbers per process per thread.

(all_range_info)

def calc_mean_ranges(all_range_info):
    """
    Calculates the mean of all NVTX ranges present in the all_range_info. Since NVTX ranges
    can be reported per process, per thread, we need to have a way to average those numbers.
    The mean here is computed by taking the average of all the numbers per process per thread.
    """
    mean_range_info = {}

    # Aggregate all the values in a list keyed by the range names.
    for process_id in all_range_info:
        for thread_id in all_range_info[process_id]:
            for range_name in all_range_info[process_id][thread_id]:
                if range_name not in mean_range_info:
                    mean_range_info[range_name] = (
                        [],
                        [],
                    )  # Mean lists for CPU and GPU time info

                cpu_time = all_range_info[process_id][thread_id][
                    range_name
                ].cpu_time_info.duration_ms
                gpu_time = all_range_info[process_id][thread_id][
                    range_name
                ].gpu_time_info.duration_ms

                mean_range_info[range_name][0].append(cpu_time)
                mean_range_info[range_name][1].append(gpu_time)

    # Replace the list with the mean value.
    for range_name in mean_range_info:
        if len(mean_range_info[range_name]):
            cpu_ranges_list = mean_range_info[range_name][0]
            gpu_ranges_list = mean_range_info[range_name][1]

            avg_cpu_time = round(sum(cpu_ranges_list) / len(cpu_ranges_list), 4)
            avg_gpu_time = round(sum(gpu_ranges_list) / len(gpu_ranges_list), 4)

            mean_range_info[range_name] = (avg_cpu_time, avg_gpu_time)
        else:
            mean_range_info[range_name] = (0, 0)

    return mean_range_info


class NumpyValuesEncoder(json.JSONEncoder):