rocm-systems/tests/rocprofv3/summary/validate.py

#!/usr/bin/env python3

import os
import re
import sys
import pytest


def test_hip_api_trace(json_data):
    data = json_data["rocprofiler-sdk-tool"]

    def get_operation_name(kind_id, op_id):
        return data["strings"]["buffer_records"][kind_id]["operations"][op_id]

    def get_kind_name(kind_id):
        return data["strings"]["buffer_records"][kind_id]["kind"]

    valid_domain_names = ("HIP_RUNTIME_API",)

    hip_api_data = data["buffer_records"]["hip_api"]

    functions = []
    for api in hip_api_data:
        kind = get_kind_name(api["kind"])
        assert kind in valid_domain_names
        assert api["end_timestamp"] >= api["start_timestamp"]
        functions.append(get_operation_name(api["kind"], api["operation"]))

    functions = list(set(functions))
    for itr in (
        "__hipPushCallConfiguration",
        "__hipPopCallConfiguration",
        "__hipRegisterFatBinary",
        "__hipRegisterFunction",
    ):
        assert itr not in functions, f"{itr}"

    for itr in (
        "hipMallocAsync",
        "hipMemcpyAsync",
        "hipMemsetAsync",
        "hipFreeAsync",
        "hipLaunchKernel",
    ):
        assert itr in functions, f"{itr}"


def test_kernel_trace(json_data):
    data = json_data["rocprofiler-sdk-tool"]

    def get_kernel_name(kernel_id):
        return data["kernel_symbols"][kernel_id]["formatted_kernel_name"]

    def get_kernel_rename(corr_id):
        for itr in data.strings.correlation_id.external:
            if itr.key == corr_id:
                return itr.value
        return None

    def get_kind_name(kind_id):
        return data["strings"]["buffer_records"][kind_id]["kind"]

    valid_kernel_names = "|".join(
        (
            "run",
            "run/iteration",
            "run/rank-([0-9]+)/thread-([0-9]+)/device-([0-9]+)/(begin|end)",
        )
    )
    valid_kernel_regex = re.compile("^({})$".format(valid_kernel_names))

    kernel_dispatch_data = data["buffer_records"]["kernel_dispatch"]
    for dispatch in kernel_dispatch_data:
        assert get_kind_name(dispatch["kind"]) == "KERNEL_DISPATCH"
        assert dispatch["correlation_id"]["internal"] > 0
        assert dispatch["correlation_id"]["external"] > 0

        dispatch_info = dispatch["dispatch_info"]
        assert dispatch_info["agent_id"]["handle"] > 0
        assert dispatch_info["queue_id"]["handle"] > 0
        assert dispatch_info["kernel_id"] > 0
        assert dispatch["end_timestamp"] >= dispatch["start_timestamp"]

        kernel_name = get_kernel_name(dispatch_info["kernel_id"])
        assert (
            re.search(valid_kernel_regex, kernel_name) is None
        ), f"kernel '{kernel_name}' matches regular expression '{valid_kernel_names}'"

        assert kernel_name not in valid_kernel_names

        external_corr_id = dispatch["correlation_id"]["external"]
        assert external_corr_id > 0

        kernel_rename = get_kernel_rename(external_corr_id)
        assert kernel_rename is not None, f"{dispatch}"
        assert kernel_rename != kernel_name, f"{dispatch}"
        assert (
            re.search(valid_kernel_regex, kernel_rename) is not None
        ), f"renamed kernel '{kernel_rename}' does not match regular expression '{valid_kernel_names}'"


def test_memory_copy_trace(json_data):
    data = json_data["rocprofiler-sdk-tool"]

    buffer_records = data["buffer_records"]
    agent_data = data["agents"]
    memory_copy_data = buffer_records["memory_copy"]

    def get_kind_name(kind_id):
        return data["strings"]["buffer_records"][kind_id]["kind"]

    def get_agent(node_id):
        for agent in agent_data:
            if agent["id"]["handle"] == node_id["handle"]:
                return agent
        return None

    assert len(memory_copy_data) == 12

    for row in memory_copy_data:
        src_agent = get_agent(row["src_agent_id"])
        dst_agent = get_agent(row["dst_agent_id"])
        assert get_kind_name(row["kind"]) == "MEMORY_COPY"
        assert src_agent is not None, f"{row}"
        assert dst_agent is not None, f"{row}"
        assert row["correlation_id"]["internal"] > 0
        assert row["end_timestamp"] >= row["start_timestamp"]


def test_metadata_data(json_data):
    data = json_data["rocprofiler-sdk-tool"]

    # patch summary groups for testing purposes
    # (it appears sometimes these contain single quotes and other times it doesn't)
    data.metadata.config.summary_groups = [
        f"{itr}".strip("'") for itr in data.metadata.config.summary_groups
    ]

    num_summary_grps = len(data.metadata.config.summary_groups)
    expected_summary_grps = (
        [
            "KERNEL_DISPATCH|MEMORY_COPY",
        ]
        if num_summary_grps == 1
        else [
            "KERNEL_DISPATCH|MEMORY_COPY",
            ".*_API",
        ]
    )

    assert data.metadata.config.summary is True
    assert data.metadata.config.summary_per_domain is True
    assert data.metadata.config.summary_unit == "nsec"
    assert data.metadata.config.summary_file == "summary"
    assert data.metadata.config.summary_groups == expected_summary_grps

    assert len(data.metadata.command) == 4

    # patch command to make it easier to test
    data.metadata.command[0] = os.path.basename(data.metadata.command[0])

    assert data.metadata.command == ["transpose", "2", "500", "10"]


def test_summary_data(json_data):
    data = json_data["rocprofiler-sdk-tool"]

    domains = []
    for itr in data.summary:
        domains.append(itr.domain)
        if itr.domain == "KERNEL_DISPATCH":
            assert itr.stats.count == 1004
            expected = dict([["run/iteration", 1000]])
            for oitr in itr.stats.operations:
                if oitr.key in expected.keys():
                    assert oitr.value.count == expected[oitr.key]
                else:
                    assert oitr.key.startswith("run/rank-")
                    assert (
                        re.match(
                            r"run/rank-([0-9]+)/thread-([0-9]+)/device-([0-9]+)/begin",
                            oitr.key,
                        )
                        is not None
                    )
                    assert oitr.value.count == 2
        elif itr.domain == "HIP_API":
            assert itr.stats.count == 2135
        elif itr.domain == "MEMORY_COPY":
            assert itr.stats.count == 12
        elif itr.domain == "MARKER_API":
            assert itr.stats.count == 1106
            expected = dict(
                [
                    ["run", 2],
                    ["run/iteration", 1000],
                    ["run/iteration/sync", 100],
                ]
            )
            for oitr in itr.stats.operations:
                if oitr.key in expected.keys():
                    assert oitr.value.count == expected[oitr.key]
                else:
                    assert oitr.key.startswith("run/rank-")
                    assert (
                        re.match(
                            r"run/rank-([0-9]+)/thread-([0-9]+)/device-([0-9]+)/(begin|end)",
                            oitr.key,
                        )
                        is not None
                    )
                    assert oitr.value.count == 1
        else:
            assert False, f"unhandled domain: {itr.domain}"

    assert len(list(set(domains))) == len(domains)


def test_summary_display_data(json_data, summary_data):
    data = json_data["rocprofiler-sdk-tool"]
    num_summary_grps = len(data.metadata.config.summary_groups)

    def get_df(domain):
        return summary_data[domain]

    def get_dims(df):
        # return rows x cols
        return [df.shape[0], df.shape[1]] if df is not None else [0, 0]

    hip = get_df("HIP_API")
    marker = get_df("MARKER_API")
    dispatch = get_df("KERNEL_DISPATCH")
    memcpy = get_df("MEMORY_COPY")
    dispatch_and_copy = get_df("KERNEL_DISPATCH + MEMORY_COPY")
    hip_and_marker = get_df("HIP_API + MARKER_API") if num_summary_grps > 1 else None
    total = get_df("SUMMARY")

    expected_hip_and_marker_dims = [20, 9] if hip_and_marker is not None else [0, 0]

    assert get_dims(hip) == [13, 9]
    assert get_dims(marker) == [7, 9]
    assert get_dims(dispatch) == [3, 9]
    assert get_dims(memcpy) == [2, 9]
    assert get_dims(dispatch_and_copy) == [5, 9]
    assert get_dims(hip_and_marker) == expected_hip_and_marker_dims
    assert get_dims(total) == [22, 9]


def test_perfetto_data(pftrace_data, json_data):
    import rocprofiler_sdk.tests.rocprofv3 as rocprofv3

    rocprofv3.test_perfetto_data(
        pftrace_data, json_data, ("hip", "marker", "kernel", "memory_copy")
    )


def test_otf2_data(otf2_data, json_data):
    import rocprofiler_sdk.tests.rocprofv3 as rocprofv3

    rocprofv3.test_otf2_data(
        otf2_data, json_data, ("hip", "marker", "kernel", "memory_copy")
    )


if __name__ == "__main__":
    exit_code = pytest.main(["-x", __file__] + sys.argv[1:])
    sys.exit(exit_code)