Function expand_nvtx_range_names

bench/python/benchmark.py:234–301

Converts a hierarchical NVTX range tree with parent-child relationship into a flat tree by adding the names of parent nodes in-front of all the child nodes. Hence, a tree like the following: root child_a sub_child_a child_b

(range_info)

def expand_nvtx_range_names(range_info):
    """
    Converts a hierarchical NVTX range tree with parent-child relationship into a flat
    tree by adding the names of parent nodes in-front of all the child nodes.
    Hence, a tree like the following:
        root
            child_a
                   sub_child_a
            child_b
            child_c
                   sub_child_c

    becomes:
        root
        root.child_a
        root.child_a.sub_child_a
        root.child_b
        root.child_c
        root.child_c.sub_child-c

    :param range_info: The range_info dictionary returned by the parsing functions.
    """
    final_dict = {}

    # Loop over all the process from the range info dictionary.
    for process_id in range_info:
        if process_id not in final_dict:
            final_dict[process_id] = {}

        # Loop over all the threads from the range info dictionary.
        for thread_id in range_info[process_id]:
            if thread_id not in final_dict[process_id]:
                final_dict[process_id][thread_id] = {}

            # Loop over all the ranges from the range info dictionary.
            for range_id in range_info[process_id][thread_id]:

                # Fetch the range information.
                nvtx_range = range_info[process_id][thread_id][range_id]

                # There are two cases to consider:
                # 1. This was a root node (i.e no parent)
                # 2. This is not a root node (i.e has a parent)
                #
                my_parent_id = nvtx_range.parent_range_id
                if my_parent_id and my_parent_id != "None":
                    # This is not a root node. Get the information of its parent.
                    parent_nvtx_range = range_info[process_id][thread_id][my_parent_id]
                    # prepend parent's name in the child's name
                    new_name = os.path.join(
                        parent_nvtx_range.flat_name, nvtx_range.flat_name
                    )

                    # Most important to update our existing range info dictionary
                    # so any nested children will end up using the new, fully
                    # qualified name of this range.
                    nvtx_range.flat_name = new_name
                    range_info[process_id][thread_id][range_id] = nvtx_range