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/rocprofv3/CMakeLists.txt

@@ -29,6 +29,7 @@ add_subdirectory(tracing-hip-in-libraries)
 add_subdirectory(counter-collection)
 add_subdirectory(hsa-queue-dependency)
 add_subdirectory(kernel-rename)
+add_subdirectory(memory-allocation)
 add_subdirectory(aborted-app)
 add_subdirectory(summary)
 add_subdirectory(roctracer-roctx)

tests/rocprofv3/memory-allocation/CMakeLists.txt

@@ -0,0 +1,47 @@
+#
+#
+#
+cmake_minimum_required(VERSION 3.21.0 FATAL_ERROR)
+
+project(
+    rocprofiler-tests-rocprofv3-memory-allocation-tracing
+    LANGUAGES CXX
+    VERSION 0.0.0)
+
+find_package(rocprofiler-sdk REQUIRED)
+
+rocprofiler_configure_pytest_files(CONFIG pytest.ini COPY validate.py conftest.py)
+
+string(REPLACE "LD_PRELOAD=" "ROCPROF_PRELOAD=" PRELOAD_ENV
+               "${ROCPROFILER_MEMCHECK_PRELOAD_ENV}")
+
+set(memory-allocation-tracing-env "${PRELOAD_ENV}")
+
+add_test(
+    NAME rocprofv3-test-memory-allocation-tracing-execute
+    COMMAND
+        $<TARGET_FILE:rocprofiler-sdk::rocprofv3> --memory-allocation-trace -d
+        ${CMAKE_CURRENT_BINARY_DIR}/%tag%-trace -o out --output-format json pftrace otf2
+        --log-level env -- $<TARGET_FILE:hsa-memory-allocation>)
+
+set_tests_properties(
+    rocprofv3-test-memory-allocation-tracing-execute
+    PROPERTIES TIMEOUT 45 LABELS "integration-tests" ENVIRONMENT
+               "${memory-allocation-tracing-env}" FAIL_REGULAR_EXPRESSION
+               "threw an exception")
+
+add_test(
+    NAME rocprofv3-test-memory-allocation-tracing-validate
+    COMMAND
+        ${Python3_EXECUTABLE} ${CMAKE_CURRENT_BINARY_DIR}/validate.py --json-input
+        ${CMAKE_CURRENT_BINARY_DIR}/hsa-memory-allocation-trace/out_results.json
+        --pftrace-input
+        ${CMAKE_CURRENT_BINARY_DIR}/hsa-memory-allocation-trace/out_results.pftrace
+        --otf2-input
+        ${CMAKE_CURRENT_BINARY_DIR}/hsa-memory-allocation-trace/out_results.otf2)
+
+set_tests_properties(
+    rocprofv3-test-memory-allocation-tracing-validate
+    PROPERTIES TIMEOUT 45 LABELS "integration-tests" DEPENDS
+               rocprofv3-test-memory-allocation-tracing-execute FAIL_REGULAR_EXPRESSION
+               "AssertionError")

tests/rocprofv3/memory-allocation/conftest.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+
+import json
+import os
+import pytest
+
+from rocprofiler_sdk.pytest_utils.dotdict import dotdict
+from rocprofiler_sdk.pytest_utils import collapse_dict_list
+from rocprofiler_sdk.pytest_utils.perfetto_reader import PerfettoReader
+from rocprofiler_sdk.pytest_utils.otf2_reader import OTF2Reader
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--json-input",
+        action="store",
+        default="memory-allocation-tracing/out_results.json",
+        help="Input JSON",
+    )
+    parser.addoption(
+        "--pftrace-input",
+        action="store",
+        default="memory-allocation-tracing/out_results.pftrace",
+        help="Input JSON",
+    )
+    parser.addoption(
+        "--otf2-input",
+        action="store",
+        default="memory-allocation-tracing/out_results.otf2",
+        help="Input JSON",
+    )
+
+
+@pytest.fixture
+def json_data(request):
+    filename = request.config.getoption("--json-input")
+    with open(filename, "r") as inp:
+        return dotdict(collapse_dict_list(json.load(inp)))
+
+
+@pytest.fixture
+def pftrace_data(request):
+    filename = request.config.getoption("--pftrace-input")
+    return PerfettoReader(filename).read()[0]
+
+
+@pytest.fixture
+def otf2_data(request):
+    filename = request.config.getoption("--otf2-input")
+    if not os.path.exists(filename):
+        raise FileExistsError(f"{filename} does not exist")
+    return OTF2Reader(filename).read()[0]

tests/rocprofv3/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/rocprofv3/memory-allocation/validate.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+
+import sys
+import pytest
+import json
+
+from collections import defaultdict
+
+
+# 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
+
+
+def get_operation(record, kind_name, op_name=None):
+    for idx, itr in enumerate(record["strings"]["buffer_records"]):
+        if kind_name == itr["kind"]:
+            if op_name is None:
+                return idx, itr["operations"]
+            else:
+                for oidx, oname in enumerate(itr["operations"]):
+                    if op_name == oname:
+                        return oidx
+    return None
+
+
+def test_memory_allocation(json_data):
+    data = json_data["rocprofiler-sdk-tool"]
+    buffer_records = data["buffer_records"]
+
+    memory_allocation_data = buffer_records["memory_allocation"]
+
+    _, bf_op_names = get_operation(data, "MEMORY_ALLOCATION")
+
+    assert len(bf_op_names) == 4
+
+    allocation_reported_agent_ids = set()
+    # check buffering data
+    for node in memory_allocation_data:
+        assert "size" in node
+        assert "kind" in node
+        assert "operation" in node
+        assert "correlation_id" in node
+        assert "end_timestamp" in node
+        assert "start_timestamp" in node
+        assert "thread_id" in node
+
+        assert "agent_id" in node
+        assert "starting_address" in node
+        assert "allocation_size" in node
+
+        assert node.size > 0
+        assert node.allocation_size > 0
+        assert node.starting_address > 0
+        assert node.thread_id > 0
+        assert node.agent_id.handle > 0
+        assert node.start_timestamp > 0
+        assert node.end_timestamp > 0
+        assert node.start_timestamp < node.end_timestamp
+
+        assert data.strings.buffer_records[node.kind].kind == "MEMORY_ALLOCATION"
+        assert (
+            data.strings.buffer_records[node.kind].operations[node.operation]
+            in bf_op_names
+        )
+
+        allocation_reported_agent_ids.add(node["agent_id"]["handle"])
+
+    assert 2**64 - 1 not in allocation_reported_agent_ids
+
+
+def test_perfetto_data(pftrace_data, json_data):
+    import rocprofiler_sdk.tests.rocprofv3 as rocprofv3
+
+    rocprofv3.test_perfetto_data(pftrace_data, json_data, ("memory_allocation",))
+
+
+def test_otf2_data(otf2_data, json_data):
+    import rocprofiler_sdk.tests.rocprofv3 as rocprofv3
+
+    rocprofv3.test_otf2_data(otf2_data, json_data, ("memory_allocation",))
+
+
+if __name__ == "__main__":
+    exit_code = pytest.main(["-x", __file__] + sys.argv[1:])
+    sys.exit(exit_code)

tests/rocprofv3/summary/validate.py

@@ -188,6 +188,8 @@ def test_summary_data(json_data):
             assert itr.stats.count >= 2130 and itr.stats.count <= 2150
         elif itr.domain == "MEMORY_COPY":
             assert itr.stats.count == 12
+        elif itr.domain == "MEMORY_ALLOCATION":
+            assert itr.stats.count >= 10 and itr.stats.count <= 30
         elif itr.domain == "MARKER_API":
             assert itr.stats.count == 1106
             expected = dict(
@@ -231,6 +233,7 @@ def test_summary_display_data(json_data, summary_data):
     marker = get_df("MARKER_API")
     dispatch = get_df("KERNEL_DISPATCH")
     memcpy = get_df("MEMORY_COPY")
+    memalloc = get_df("MEMORY_ALLOCATION")
     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")
@@ -239,18 +242,21 @@ def test_summary_display_data(json_data, summary_data):
 
     assert get_dims(marker) == [7, 9], f"{marker}"
     assert get_dims(memcpy) == [2, 9], f"{memcpy}"
+    assert get_dims(memalloc) == [1, 9], f"{memalloc}"
     assert get_dims(dispatch) == [3, 9], f"{dispatch}"
     assert get_dims(dispatch_and_copy) == [5, 9], f"{dispatch_and_copy}"
     assert get_dims(hip) == [14, 9], f"{hip}"
     assert get_dims(hip_and_marker) == expected_hip_and_marker_dims, f"{hip_and_marker}"
-    assert get_dims(total) == [23, 9], f"{total}"
+    assert get_dims(total) == [24, 9], f"{total}"
 
 
 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")
+        pftrace_data,
+        json_data,
+        ("hip", "marker", "kernel", "memory_copy", "memory_allocation"),
     )
 
 
@@ -258,7 +264,9 @@ 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")
+        otf2_data,
+        json_data,
+        ("hip", "marker", "kernel", "memory_copy", "memory_allocation"),
     )