Function parse_nvtx_gpu_proj_trace_json

bench/python/benchmark.py:146–231

Parses the nvtx_gpu_proj_trace JSON generated by NSYS and returns a dictionary keyed by process_id, thread_id and range_id. The values are various fields important to the benchmarking process. :param json_path: Full path to the nvtx_gpu_proj_trace.json file.

(json_path)

def parse_nvtx_gpu_proj_trace_json(json_path):
    """
    Parses the nvtx_gpu_proj_trace JSON generated by NSYS and returns a dictionary
    keyed by process_id, thread_id and range_id. The values are various fields
    important to the benchmarking process.
    :param json_path: Full path to the nvtx_gpu_proj_trace.json file.
    """

    #
    # The nvtx_gpu_proj_trace JSON has the following structure. It is a list of
    # dictionaries.
    # e.g.
    #  [ {
    #   "Projected Start (ns)": 2372801266,
    #   "Projected Duration (ns)"  : 13528369268,
    #       ...
    #   },
    # ...
    # ]
    #
    # We will store the parsed data in the range_info dictionary. The overall
    # structure of the dictionary is:
    # range_info = {
    #       process_id : {
    #           thread_id : {
    #               range_id : NvtxRange(flat_name, parent_range_id, cpu_duration_ms, gpu_duration_ms)
    #           }
    #       }
    # }
    #
    # NOTE: Even though this report returns the cpu_duration_ms and gpu_duration_ms, it will
    #       only do so for operations which had gpu_duration_ms > 0. For pure CPU operations,
    #       this report will not even return those ranges. That is the reason why we need to
    #       query the pushpop_trace report.
    #
    range_info = {}

    # Check if the file was empty or not. Empty file means no GPU ops were recorded.
    if os.stat(json_path).st_size == 0:
        return range_info

    # Read the JSON.
    with open(json_path, "r") as f:
        json_data = json.loads(f.read())

    for row in json_data:
        # Grab the necessary values from the JSON file.
        range_id = row["RangeId"]

        if not range_id or range_id == "None":
            continue

        flat_name = row["Name"]
        cpu_start_ns = float(row["Orig Start (ns)"])
        cpu_duration_ns = float(row["Orig Duration (ns)"])
        cpu_end_ns = cpu_start_ns + cpu_duration_ns

        gpu_start_ns = float(row["Projected Start (ns)"])