Memory Allocation Tracking (#1142)

* Initial commit: Need to implement wrapper function to collect data and test that wrapper function is correctly replacing core HSA functions

* Attempted to implement wrapper implementation for hsa memory allocation functions. Need to modify generate record files and test if implementation is working as expected

* Debugging and implementing generateCSV function

* Memory allocation size and starting address outputted to csv and json file formats

* Formatting

* Initial setup for OTF2 and Perfetto generation

* Collecting agent id for memory_allocation and formatting

* Modified memory_allocation.cpp to set up code for AMD_EXT commands

* Support for memory_pool_allocate added

* Removed accidently added file

* Made flag optional and added more OTF2 and Perfetto code. Needs testing to ensure perfetto and OTF2 works

* Formatting

* Fixed perfetto and otf2 output

* Fixed flag issue due to incorrect buffer use

* Updated documentation

* Small cleaning and comments

* Added test for HSA memory allocation tracing

* Fixed summary test validation errors due to allocation tracing. Added type to location_base to create unique event ids for allocation due to OTF2 trace error

* Decreased lower limit of hip calls for test

* Modified summary tests to vary number of allocate requests

* Minor fixes to address comments. Still need to address OTF2 comments

* Fix docs and changed OTF2 to use enum for type specified in location_base construction

* Fixed schema error

* Added vmem command tracking. Need to add test

* Updated test to work with vmem command and updated generateCSV to output int instead of hex string.

* OTF2 enum update and mispelling fix

* CI does not support Virtual Memory API. Removed vmem test. Will add back if CI is modifed to suport vmem API

* Update CMakeLists.txt for memory allocation test

* Updated summary test

* Minor fixes to address comments

* Moved domain_type.hpp enum to before LAST

* Fixed compile errors and formatting

* Fixed stats summary domain name error

* Added rocprofv3 test

* Page migration test fix

* Undo page migration test changes. Failures do not appear to have to do with memory allocation

tests/hsa-memory-allocation/CMakeLists.txt

@@ -0,0 +1,46 @@
+#
+#
+#
+cmake_minimum_required(VERSION 3.21.0 FATAL_ERROR)
+
+project(
+    rocprofiler-tests-memory-allocation-tracing
+    LANGUAGES CXX
+    VERSION 0.0.0)
+
+find_package(rocprofiler-sdk REQUIRED)
+
+if(ROCPROFILER_MEMCHECK_PRELOAD_ENV)
+    set(PRELOAD_ENV
+        "${ROCPROFILER_MEMCHECK_PRELOAD_ENV}:$<TARGET_FILE:rocprofiler-sdk-json-tool>")
+else()
+    set(PRELOAD_ENV "LD_PRELOAD=$<TARGET_FILE:rocprofiler-sdk-json-tool>")
+endif()
+
+add_test(NAME test-memory-allocation-tracing-execute
+         COMMAND $<TARGET_FILE:hsa-memory-allocation>)
+
+set(memory-allocation-tracing-env
+    "${PRELOAD_ENV}"
+    "ROCPROFILER_TOOL_OUTPUT_FILE=memory-allocation-tracing-test.json"
+    "LD_LIBRARY_PATH=$<TARGET_FILE_DIR:rocprofiler-sdk::rocprofiler-sdk-shared-library>:$ENV{LD_LIBRARY_PATH}"
+    )
+
+set_tests_properties(
+    test-memory-allocation-tracing-execute
+    PROPERTIES TIMEOUT 45 LABELS "integration-tests" ENVIRONMENT
+               "${memory-allocation-tracing-env}" FAIL_REGULAR_EXPRESSION
+               "${ROCPROFILER_DEFAULT_FAIL_REGEX}")
+
+# copy to binary directory
+rocprofiler_configure_pytest_files(COPY validate.py conftest.py CONFIG pytest.ini)
+
+add_test(NAME test-memory-allocation-tracing-validate
+         COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_BINARY_DIR}/validate.py --input
+                 ${CMAKE_CURRENT_BINARY_DIR}/memory-allocation-tracing-test.json)
+
+set_tests_properties(
+    test-memory-allocation-tracing-validate
+    PROPERTIES TIMEOUT 45 LABELS "integration-tests" DEPENDS
+               test-memory-allocation-tracing-execute FAIL_REGULAR_EXPRESSION
+               "${ROCPROFILER_DEFAULT_FAIL_REGEX}")

tests/hsa-memory-allocation/conftest.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+
+import json
+import pytest
+
+from rocprofiler_sdk.pytest_utils.dotdict import dotdict
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--input",
+        action="store",
+        default="memory-allocation-tracing-test.json",
+        help="Input JSON",
+    )
+
+
+@pytest.fixture
+def input_data(request):
+    filename = request.config.getoption("--input")
+    with open(filename, "r") as inp:
+        return dotdict(json.load(inp))

tests/hsa-memory-allocation/pytest.ini

@@ -0,0 +1,5 @@
+
+[pytest]
+addopts = --durations=20 -rA -s -vv
+testpaths = validate.py
+pythonpath = @ROCPROFILER_SDK_TESTS_BINARY_DIR@/pytest-packages

tests/hsa-memory-allocation/validate.py

@@ -0,0 +1,289 @@
+#!/usr/bin/env python3
+
+import sys
+import pytest
+
+
+# helper function
+def node_exists(name, data, min_len=1):
+    assert name in data
+    assert data[name] is not None
+    if isinstance(data[name], (list, tuple, dict, set)):
+        assert len(data[name]) >= min_len, f"{name}:\n{data}"
+
+
+def test_data_structure(input_data):
+    """verify minimum amount of expected data is present"""
+    data = input_data
+
+    node_exists("rocprofiler-sdk-json-tool", data)
+
+    sdk_data = data["rocprofiler-sdk-json-tool"]
+
+    node_exists("metadata", sdk_data)
+    node_exists("pid", sdk_data["metadata"])
+    node_exists("main_tid", sdk_data["metadata"])
+    node_exists("init_time", sdk_data["metadata"])
+    node_exists("fini_time", sdk_data["metadata"])
+
+    node_exists("agents", sdk_data)
+    node_exists("call_stack", sdk_data)
+    node_exists("callback_records", sdk_data)
+    node_exists("buffer_records", sdk_data)
+
+    node_exists("names", sdk_data["callback_records"])
+    node_exists("hsa_api_traces", sdk_data["callback_records"])
+    node_exists("memory_allocations", sdk_data["callback_records"])
+
+    node_exists("names", sdk_data["buffer_records"])
+    node_exists("hsa_api_traces", sdk_data["callback_records"])
+    node_exists("memory_allocations", sdk_data["buffer_records"])
+
+
+def test_size_entries(input_data):
+    # check that size fields are > 0 but account for function arguments
+    # which are named "size"
+    def check_size(data, bt):
+        if "size" in data.keys():
+            if isinstance(data["size"], str) and bt.endswith('["args"]'):
+                pass
+            else:
+                assert data["size"] > 0, f"origin: {bt}"
+
+    # recursively check the entire data structure
+    def iterate_data(data, bt):
+        if isinstance(data, (list, tuple)):
+            for i, itr in enumerate(data):
+                if isinstance(itr, dict):
+                    check_size(itr, f"{bt}[{i}]")
+                iterate_data(itr, f"{bt}[{i}]")
+        elif isinstance(data, dict):
+            check_size(data, f"{bt}")
+            for key, itr in data.items():
+                iterate_data(itr, f'{bt}["{key}"]')
+
+    # start recursive check over entire JSON dict
+    iterate_data(input_data, "input_data")
+
+
+def test_timestamps(input_data):
+    """Verify starting timestamps are less than ending timestamps"""
+    data = input_data
+    sdk_data = data["rocprofiler-sdk-json-tool"]
+
+    cb_start = {}
+    cb_end = {}
+    for titr in ["hsa_api_traces"]:
+        for itr in sdk_data["callback_records"][titr]:
+            cid = itr["correlation_id"]["internal"]
+            phase = itr["phase"]
+            if phase == 1:
+                cb_start[cid] = itr["timestamp"]
+            elif phase == 2:
+                cb_end[cid] = itr["timestamp"]
+                assert cb_start[cid] <= itr["timestamp"]
+            else:
+                assert phase == 1 or phase == 2
+
+        for itr in sdk_data["buffer_records"][titr]:
+            assert itr["start_timestamp"] <= itr["end_timestamp"]
+
+    for titr in ["memory_allocations"]:
+        for itr in sdk_data["buffer_records"][titr]:
+            assert itr["start_timestamp"] < itr["end_timestamp"], f"[{titr}] {itr}"
+            assert itr["correlation_id"]["internal"] > 0, f"[{titr}] {itr}"
+            assert itr["correlation_id"]["external"] > 0, f"[{titr}] {itr}"
+            assert (
+                sdk_data["metadata"]["init_time"] < itr["start_timestamp"]
+            ), f"[{titr}] {itr}"
+            assert (
+                sdk_data["metadata"]["init_time"] < itr["end_timestamp"]
+            ), f"[{titr}] {itr}"
+            assert (
+                sdk_data["metadata"]["fini_time"] > itr["start_timestamp"]
+            ), f"[{titr}] {itr}"
+            assert (
+                sdk_data["metadata"]["fini_time"] > itr["end_timestamp"]
+            ), f"[{titr}] {itr}"
+
+            api_start = cb_start[itr["correlation_id"]["internal"]]
+            # api_end = cb_end[itr["correlation_id"]["internal"]]
+            assert api_start < itr["start_timestamp"], f"[{titr}] {itr}"
+            # assert api_end <= itr["end_timestamp"], f"[{titr}] {itr}"
+
+
+def test_internal_correlation_ids(input_data):
+    """Assure correlation ids are unique"""
+    data = input_data
+    sdk_data = data["rocprofiler-sdk-json-tool"]
+
+    api_corr_ids = []
+    for titr in ["hsa_api_traces"]:
+        for itr in sdk_data["callback_records"][titr]:
+            api_corr_ids.append(itr["correlation_id"]["internal"])
+
+        for itr in sdk_data["buffer_records"][titr]:
+            api_corr_ids.append(itr["correlation_id"]["internal"])
+
+    api_corr_ids_sorted = sorted(api_corr_ids)
+    api_corr_ids_unique = list(set(api_corr_ids))
+
+    for itr in sdk_data["buffer_records"]["memory_allocations"]:
+        assert itr["correlation_id"]["internal"] in api_corr_ids_unique
+
+    len_corr_id_unq = len(api_corr_ids_unique)
+    assert len(api_corr_ids) != len_corr_id_unq
+    assert max(api_corr_ids_sorted) == len_corr_id_unq
+
+
+def test_external_correlation_ids(input_data):
+    data = input_data
+    sdk_data = data["rocprofiler-sdk-json-tool"]
+
+    extern_corr_ids = []
+    for titr in ["hsa_api_traces"]:
+        for itr in sdk_data["callback_records"][titr]:
+            assert itr["correlation_id"]["external"] > 0
+            assert itr["thread_id"] == itr["correlation_id"]["external"]
+            extern_corr_ids.append(itr["correlation_id"]["external"])
+
+    extern_corr_ids = list(set(sorted(extern_corr_ids)))
+    for titr in ["hsa_api_traces"]:
+        for itr in sdk_data["buffer_records"][titr]:
+            assert itr["correlation_id"]["external"] > 0, f"[{titr}] {itr}"
+            assert (
+                itr["thread_id"] == itr["correlation_id"]["external"]
+            ), f"[{titr}] {itr}"
+            assert itr["thread_id"] in extern_corr_ids, f"[{titr}] {itr}"
+            assert itr["correlation_id"]["external"] in extern_corr_ids, f"[{titr}] {itr}"
+
+    for titr in ["memory_allocations"]:
+        for itr in sdk_data["buffer_records"][titr]:
+            assert itr["correlation_id"]["external"] > 0, f"[{titr}] {itr}"
+            assert itr["correlation_id"]["external"] in extern_corr_ids, f"[{titr}] {itr}"
+
+        for itr in sdk_data["callback_records"][titr]:
+            assert itr["correlation_id"]["external"] > 0, f"[{titr}] {itr}"
+            assert itr["correlation_id"]["external"] in extern_corr_ids, f"[{titr}] {itr}"
+
+
+def test_memory_alloc_sizes(input_data):
+    """Ensure trace file memory allocation operations match up with the memory allocation operations performed in hsa-memory-allocation"""
+    data = input_data
+    sdk_data = data["rocprofiler-sdk-json-tool"]
+
+    # Op values:
+    #   0 == ??? (unknown)
+    #   1 == hsa_memory_allocate
+    #   2 == hsa_amd_ext_memory_allocate
+    #   3 == hsa_amd_vmem_handle_create
+    memory_alloc_cnt = dict(
+        [
+            (idx, {"agent": set(), "starting_addr": set(), "size": set(), "count": 0})
+            for idx in range(1, 4)
+        ]
+    )
+    for itr in sdk_data["buffer_records"]["memory_allocations"]:
+        op_id = itr["operation"]
+        assert op_id > 0 and op_id <= 3, f"{itr}"
+        memory_alloc_cnt[op_id]["count"] += 1
+        memory_alloc_cnt[op_id]["starting_addr"].add(itr.starting_address)
+        memory_alloc_cnt[op_id]["size"].add(itr.allocation_size)
+        memory_alloc_cnt[op_id]["agent"].add(itr.agent_id.handle)
+
+    for itr in sdk_data["callback_records"]["memory_copies"]:
+        op_id = itr.operation
+        assert op_id > 0 and op_id <= 3, f"{itr}"
+        memory_alloc_cnt[op_id]["count"] += 1
+
+        phase = itr.phase
+        pitr = itr.payload
+
+        assert phase is not None, f"{itr}"
+        assert pitr is not None, f"{itr}"
+
+        if phase == 1:
+            assert pitr.start_timestamp == 0, f"{itr}"
+            assert pitr.end_timestamp == 0, f"{itr}"
+        elif phase == 2:
+            assert pitr.start_timestamp > 0, f"{itr}"
+            assert pitr.end_timestamp > 0, f"{itr}"
+            assert pitr.end_timestamp >= pitr.start_timestamp, f"{itr}"
+
+            memory_alloc_cnt[op_id]["starting_addr"].add(pitr.starting_address)
+            memory_alloc_cnt[op_id]["size"].add(pitr.allocation_size)
+            memory_alloc_cnt[op_id]["agent"].add(pitr.agent_id.handle)
+        else:
+            assert phase == 1 or phase == 2, f"{itr}"
+
+    # In the memory allocation test which generates this file
+    # 6 hsa_memory_allocation calls with 1024 bytes were called
+    # and 9 hsa_amd_memory_pool_allocations with 512 bytes
+    # were called
+    assert memory_alloc_cnt[1]["count"] == 6
+    assert memory_alloc_cnt[2]["count"] == 9
+    # assert memory_alloc_cnt[3]["count"] == 3
+    assert len(memory_alloc_cnt[1]["starting_addr"]) == 6
+    assert len(memory_alloc_cnt[2]["starting_addr"]) == 9
+    # assert len(memory_alloc_cnt[3]["starting_addr"]) == 3
+    assert len(memory_alloc_cnt[1]["size"]) == 1
+    assert len(memory_alloc_cnt[2]["size"]) == 1
+    # assert len(memory_alloc_cnt[3]["size"]) == 1
+    assert 1024 in memory_alloc_cnt[1]["size"]
+    assert 512 in memory_alloc_cnt[2]["size"]
+    assert len(memory_alloc_cnt[1]["agent"]) == 1
+    assert len(memory_alloc_cnt[2]["agent"]) == 1
+    # assert len(memory_alloc_cnt[3]["agent"]) == 1
+    assert memory_alloc_cnt[1]["agent"] != memory_alloc_cnt[2]["agent"]
+    # assert memory_alloc_cnt[2]["agent"] == memory_alloc_cnt[3]["agent"]
+
+
+def test_retired_correlation_ids(input_data):
+    data = input_data
+    sdk_data = data["rocprofiler-sdk-json-tool"]
+
+    def _sort_dict(inp):
+        return dict(sorted(inp.items()))
+
+    api_corr_ids = {}
+    for titr in ["hsa_api_traces"]:
+        for itr in sdk_data["buffer_records"][titr]:
+            corr_id = itr["correlation_id"]["internal"]
+            assert corr_id not in api_corr_ids.keys()
+            api_corr_ids[corr_id] = itr
+
+    alloc_corr_ids = {}
+    for titr in ["memory_allocations"]:
+        for itr in sdk_data["buffer_records"][titr]:
+            corr_id = itr["correlation_id"]["internal"]
+            assert corr_id not in alloc_corr_ids.keys()
+            alloc_corr_ids[corr_id] = itr
+
+    retired_corr_ids = {}
+    for itr in sdk_data["buffer_records"]["retired_correlation_ids"]:
+        corr_id = itr["internal_correlation_id"]
+        assert corr_id not in retired_corr_ids.keys()
+        retired_corr_ids[corr_id] = itr
+
+    api_corr_ids = _sort_dict(api_corr_ids)
+    alloc_corr_ids = _sort_dict(alloc_corr_ids)
+    retired_corr_ids = _sort_dict(retired_corr_ids)
+
+    for cid, itr in alloc_corr_ids.items():
+        assert cid in retired_corr_ids.keys()
+        retired_ts = retired_corr_ids[cid]["timestamp"]
+        end_ts = itr["end_timestamp"]
+        assert (retired_ts - end_ts) > 0, f"correlation-id: {cid}, data: {itr}"
+
+    for cid, itr in api_corr_ids.items():
+        assert cid in retired_corr_ids.keys()
+        retired_ts = retired_corr_ids[cid]["timestamp"]
+        end_ts = itr["end_timestamp"]
+        assert (retired_ts - end_ts) > 0, f"correlation-id: {cid}, data: {itr}"
+
+    assert len(api_corr_ids.keys()) == (len(retired_corr_ids.keys()))
+
+
+if __name__ == "__main__":
+    exit_code = pytest.main(["-x", __file__] + sys.argv[1:])
+    sys.exit(exit_code)