Add support for scratch reporting (#523)

* Add ToolsApiTable Add ToolsApiTable wrapping for scratch memory tracking * Add initial support for scratch memory tracking Buffering is implemented * cmake formatting (cmake-format) (#525) Co-authored-by: MythreyaK <MythreyaK@users.noreply.github.com> * source formatting (clang-format v11) (#524) Co-authored-by: MythreyaK <MythreyaK@users.noreply.github.com> * Add callback tracing for scratch Fixed the error where scratch tracking init was called irrespective of whether any client requested for it * Apply suggestions from code review Co-authored-by: Jonathan R. Madsen <jrmadsen@users.noreply.github.com> * Fix tools api copy/update Table were saved/updated incorrectly in previous commit. Also adds passing user data through the callback * Fix OpKind sequence for scratch tracking Previously scratch was using OpKind from rocprofiler-sdk, but templates were instantiated using API ID. These differ by 1 * Integration tests for scratch reporting Added buffer and callback integration tests for scratch reporting * source formatting (clang-format v11) (#550) Co-authored-by: MythreyaK <26112391+MythreyaK@users.noreply.github.com> * cmake formatting (cmake-format) (#551) Co-authored-by: MythreyaK <26112391+MythreyaK@users.noreply.github.com> * python formatting (black) (#549) Co-authored-by: MythreyaK <26112391+MythreyaK@users.noreply.github.com> * CI fixes * source formatting (clang-format v11) (#554) Co-authored-by: MythreyaK <26112391+MythreyaK@users.noreply.github.com> * Update api Rebase on main and updates based on PR feedback * Update scratch reporting and address PR comments - Added agent id to buffer records - Updated `test_internal_correlation_ids` - Is almost identical to one in async-copy - Updated scratch test to check for agent id - Updated queue id serialization in callback records (prints handle as nested key) - Remove `marker_api_traces` from scratch `test_internal_correlation_ids` validation test - Rename `amd_tools_api` to `scratch_memory` - Added doxygen comments - Remove scratch callback from `tool.cpp` - Replace assert with `LOF_IF` in `scratch_memory.cpp` * Update tools table Changed to match up with changes to hsa tables in main branch * Rework scratch memory structure * Update tests - Added suggestions from PR review, and updated tests accordingly * Misc cleanup * Update scratch test As of Apr 4th, `hsa_amd_agent_set_async_scratch_limit` is disabled. Note, > This API: `hsa_amd_agent_set_async_scratch_limit` is currently > disabled. We need some changes in CP firmware to be able to do this > and these changes are not ready yet. > With the current code, you will also not get notifications for > alternate-scratch allocations because this feature has been disabled > while CP firmware is making additional changes > We are hoping to have that feature enabled by ROCm-6.3 * Minor update to lib/rocprofiler-sdk/internal_threading.* - delay destruction of shared_ptrs of the tasks to prevent rare (but possible) data race on the destruction of the shared_ptr --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: MythreyaK <MythreyaK@users.noreply.github.com> Co-authored-by: Jonathan R. Madsen <jrmadsen@users.noreply.github.com> Co-authored-by: Jonathan R. Madsen <jonathanrmadsen@gmail.com>
2024-04-05 18:32:57 -07:00
@@ -49,6 +49,7 @@ add_subdirectory(bin)
 # validation tests
 add_subdirectory(kernel-tracing)
 add_subdirectory(async-copy-tracing)
+add_subdirectory(scratch-memory-tracing)
 add_subdirectory(c-tool)

 # rocprofv3 validation tests
@@ -12,6 +12,7 @@ add_subdirectory(simple-transpose)
 add_subdirectory(multistream)
 add_subdirectory(vector-operations)
 add_subdirectory(hip-in-libraries)
+add_subdirectory(scratch-memory)

 set(CMAKE_BUILD_RPATH
    "\$ORIGIN:\$ORIGIN/../lib:$<TARGET_FILE_DIR:rocprofiler-sdk-roctx::rocprofiler-sdk-roctx-shared-library>"
@@ -0,0 +1,47 @@
+#
+#
+#
+cmake_minimum_required(VERSION 3.21.0 FATAL_ERROR)
+
+if(NOT CMAKE_HIP_COMPILER)
+    find_program(
+        amdclangpp_EXECUTABLE
+        NAMES amdclang++
+        HINTS ${ROCM_PATH} ENV ROCM_PATH /opt/rocm
+        PATHS ${ROCM_PATH} ENV ROCM_PATH /opt/rocm
+        PATH_SUFFIXES bin llvm/bin NO_CACHE)
+    mark_as_advanced(amdclangpp_EXECUTABLE)
+
+    if(amdclangpp_EXECUTABLE)
+        set(CMAKE_HIP_COMPILER "${amdclangpp_EXECUTABLE}")
+    endif()
+endif()
+
+project(rocprofiler-tool-test-app-scratch-memory LANGUAGES CXX HIP)
+
+foreach(_TYPE DEBUG MINSIZEREL RELEASE RELWITHDEBINFO)
+    if("${CMAKE_HIP_FLAGS_${_TYPE}}" STREQUAL "")
+        set(CMAKE_HIP_FLAGS_${_TYPE} "${CMAKE_CXX_FLAGS_${_TYPE}}")
+    endif()
+endforeach()
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_EXTENSIONS OFF)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_HIP_STANDARD 17)
+set(CMAKE_HIP_EXTENSIONS OFF)
+set(CMAKE_HIP_STANDARD_REQUIRED ON)
+
+set_source_files_properties(scratch-memory.cpp PROPERTIES LANGUAGE HIP)
+add_executable(scratch-memory)
+target_sources(scratch-memory PRIVATE scratch-memory.cpp)
+target_compile_options(scratch-memory PRIVATE -W -Wall -Wextra -Wpedantic -Wshadow
+                                              -Werror)
+
+find_package(Threads REQUIRED)
+target_link_libraries(scratch-memory PRIVATE Threads::Threads hsa-runtime64)
+
+install(
+    TARGETS scratch-memory
+    DESTINATION bin
+    COMPONENT tests)
@@ -0,0 +1,239 @@
+// MIT License
+//
+// Copyright (c) 2023 Advanced Micro Devices, Inc. All rights reserved.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+#include <hip/hip_runtime.h>
+#include <hsa/hsa.h>
+#include <hsa/hsa_ext_amd.h>
+#include <cstdio>
+#include <iostream>
+#include <vector>
+
+#define hipCheckErr(errval)                                                                        \
+    do                                                                                             \
+    {                                                                                              \
+        hipCheckAndFail((errval), __FILE__, __LINE__);                                             \
+    } while(0)
+
+#define hipCheckLastError()                                                                        \
+    do                                                                                             \
+    {                                                                                              \
+        hipCheckErr(hipGetLastError());                                                            \
+    } while(0)
+
+#define HSA_CALL2(cmd)                                                                             \
+    do                                                                                             \
+    {                                                                                              \
+        hsa_status_t error = (cmd);                                                                \
+        if(error != HSA_STATUS_SUCCESS)                                                            \
+        {                                                                                          \
+            const char* errorStr;                                                                  \
+            hsa_status_string(error, &errorStr);                                                   \
+            std::cout << "Encountered HSA error (" << errorStr << ") at line " << __LINE__         \
+                      << " in file " << __FILE__ << "\n";                                          \
+            exit(-1);                                                                              \
+        }                                                                                          \
+    } while(0)
+
+namespace
+{
+inline void
+hipCheckAndFail(hipError_t errval, const char* file, int line)
+{
+    if(errval != hipSuccess)
+    {
+        std::cerr << "hip error: " << hipGetErrorString(errval) << std::endl;
+        std::cerr << "    Location: " << file << ":" << line << std::endl;
+        exit(errval);
+    }
+}
+
+hsa_status_t
+find_gpu_agents(hsa_agent_t agent, void* data)
+{
+    hsa_status_t      status;
+    hsa_device_type_t device_type;
+    status = hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, &device_type);
+    if(status == HSA_STATUS_SUCCESS && device_type == HSA_DEVICE_TYPE_GPU)
+    {
+        std::vector<hsa_agent_t>* agents = reinterpret_cast<std::vector<hsa_agent_t>*>(data);
+        agents->push_back(agent);
+    }
+    return HSA_STATUS_SUCCESS;
+}
+}  // namespace
+
+__global__ void
+test_kern_large(uint64_t* output)
+{
+    uint64_t result = 0;
+    int      test[4000];
+    memset(test, 5, 4000);
+    for(int& i : test)
+    {
+        i = i + 7;
+        *output += i;
+        result += i;
+    }
+    *output ^= result;
+    *output ^= result;
+}
+
+__global__ void
+test_kern_medium(uint64_t* output)
+{
+    uint64_t result = 0;
+    int      test[175];
+    memset(test, 5, 175);
+    for(int& i : test)
+    {
+        i = i + 7;
+        *output += i;
+        result += i;
+    }
+    *output ^= result;
+    *output ^= result;
+}
+
+__global__ void
+test_kern_small(uint64_t* output)
+{
+    uint64_t result = 0;
+    int      test[2];
+    for(int& i : test)
+    {
+        i = i + 7;
+        *output += i;
+        result += i;
+    }
+    *output ^= result;
+    *output ^= result;
+}
+
+// Checks whether we get a request-more-scratch when grid-x is incremented
+int
+test_gridx(uint64_t* data_ptr)
+{
+    *data_ptr = 0;
+    printf("Running Medium\n");
+    test_kern_medium<<<1000, 1>>>(data_ptr);
+    hipCheckErr(hipDeviceSynchronize());
+    printf("Running Medium - done\n");
+
+    printf("Running Medium-2 - should trigger more-scratch requests\n");
+    test_kern_medium<<<1500, 1>>>(data_ptr);
+    hipCheckErr(hipDeviceSynchronize());
+
+    printf("Running Medium-2 - done\n");
+    return 0;
+}
+
+// 1st allocation should go to primary, then large should still trigger a USO
+int
+test_primary_then_uso(uint64_t* data_ptr)
+{
+    printf("Running Medium - all slots\n");
+    test_kern_medium<<<10000, 1>>>(data_ptr);
+    hipCheckErr(hipDeviceSynchronize());
+    printf("Running Medium - done\n");
+
+    printf("Running Large - should trigger USO\n");
+    test_kern_large<<<1100, 1>>>(data_ptr);
+    hipCheckErr(hipDeviceSynchronize());
+    printf("Running Large - done\n");
+    return 0;
+}
+
+int
+test_scratch()
+{
+    uint64_t* data_ptr;
+    hipCheckErr(hipHostMalloc(&data_ptr, sizeof(uint64_t), 0));
+
+    std::vector<float> host_floats(1024);
+    float*             dev;
+
+    hipCheckErr(hipMalloc((void**) &dev, host_floats.size() * sizeof(float)));
+    hipCheckErr(hipMemcpy(
+        dev, host_floats.data(), host_floats.size() * sizeof(float), hipMemcpyHostToDevice));
+
+    *data_ptr = 0;
+
+    printf("Running test_primary_then_uso========================\n");
+    test_primary_then_uso(data_ptr);
+    printf("=====================================================\n");
+
+    printf("Running test_gridx===================================\n");
+    test_gridx(data_ptr);
+    printf("=====================================================\n");
+
+    printf("Running Small\n");
+    test_kern_small<<<1000, 1>>>(data_ptr);
+    hipCheckErr(hipDeviceSynchronize());
+    printf("Running Small - done\n");
+
+    printf("Running Medium\n");
+    test_kern_medium<<<1000, 1>>>(data_ptr);
+    hipCheckErr(hipDeviceSynchronize());
+    printf("Running Medium - done\n");
+
+    printf("Running Small\n");
+    test_kern_small<<<1000, 1>>>(data_ptr);
+    hipCheckErr(hipDeviceSynchronize());
+    printf("Running Small - done\n");
+
+    printf("Running Large\n");
+    test_kern_large<<<1100, 1>>>(data_ptr);
+    hipCheckErr(hipDeviceSynchronize());
+    printf("Running Large - done\n");
+
+    printf("Running Large\n");
+    test_kern_large<<<1000, 1>>>(data_ptr);
+    hipCheckErr(hipDeviceSynchronize());
+    printf("Running Large - done\n");
+
+    printf("Running Large\n");
+    test_kern_large<<<1000, 1>>>(data_ptr);
+    hipCheckErr(hipFree(dev));
+    hipCheckErr(hipDeviceSynchronize());
+    printf("Running Large - done\n");
+
+    return 0;
+}
+
+int
+main()
+{
+    hipCheckErr(hipInit(0));
+
+    std::vector<hsa_agent_t> agents;
+    HSA_CALL2(hsa_iterate_agents(find_gpu_agents, &agents));
+    size_t numAgents = agents.size();
+    printf("Detected %ld agents\n", numAgents);
+
+    for(size_t i = 0; i < agents.size(); ++i)
+    {
+        hipCheckErr(hipSetDevice(i));
+        test_scratch();
+    }
+
+    return 0;
+}
@@ -160,6 +160,62 @@ save(ArchiveT& ar, rocprofiler_hsa_api_retval_t data)
    SAVE_DATA_FIELD(uint64_t_retval);
 }

+template <typename ArchiveT>
+void
+save(ArchiveT& ar, const hsa_queue_t& data)
+{
+    ar(make_nvp("queue_id", data.id));
+}
+
+template <typename ArchiveT>
+void
+save(ArchiveT& ar, hsa_amd_event_scratch_alloc_start_t data)
+{
+    ar(make_nvp("queue_id", *data.queue));
+    SAVE_DATA_FIELD(dispatch_id);
+}
+
+template <typename ArchiveT>
+void
+save(ArchiveT& ar, hsa_amd_event_scratch_alloc_end_t data)
+{
+    ar(make_nvp("queue_id", *data.queue));
+    SAVE_DATA_FIELD(dispatch_id);
+    SAVE_DATA_FIELD(size);
+    SAVE_DATA_FIELD(num_slots);
+    SAVE_DATA_FIELD(flags);
+}
+
+template <typename ArchiveT>
+void
+save(ArchiveT& ar, hsa_amd_event_scratch_free_start_t data)
+{
+    ar(make_nvp("queue_id", *data.queue));
+}
+
+template <typename ArchiveT>
+void
+save(ArchiveT& ar, hsa_amd_event_scratch_free_end_t data)
+{
+    ar(make_nvp("queue_id", *data.queue));
+    SAVE_DATA_FIELD(flags);
+}
+
+template <typename ArchiveT>
+void
+save(ArchiveT& ar, hsa_amd_event_scratch_async_reclaim_start_t data)
+{
+    ar(make_nvp("queue_id", *data.queue));
+}
+
+template <typename ArchiveT>
+void
+save(ArchiveT& ar, hsa_amd_event_scratch_async_reclaim_end_t data)
+{
+    ar(make_nvp("queue_id", *data.queue));
+    SAVE_DATA_FIELD(flags);
+}
+
 template <typename ArchiveT>
 void
 save(ArchiveT& ar, rocprofiler_marker_api_retval_t data)
@@ -201,6 +257,17 @@ save(ArchiveT& ar, rocprofiler_callback_tracing_hip_api_data_t data)
    SAVE_DATA_FIELD(retval);
 }

+template <typename ArchiveT>
+void
+save(ArchiveT& ar, rocprofiler_callback_tracing_scratch_memory_data_t data)
+{
+    SAVE_DATA_FIELD(size);
+    SAVE_DATA_FIELD(agent_id);
+    SAVE_DATA_FIELD(queue_id);
+    SAVE_DATA_FIELD(flags);
+    SAVE_DATA_FIELD(args_kind);
+}
+
 template <typename ArchiveT>
 void
 save(ArchiveT& ar, rocprofiler_callback_tracing_record_t data)
@@ -288,6 +355,22 @@ save(ArchiveT& ar, rocprofiler_buffer_tracing_memory_copy_record_t data)
    SAVE_DATA_FIELD(src_agent_id);
 }

+template <typename ArchiveT>
+void
+save(ArchiveT& ar, rocprofiler_buffer_tracing_scratch_memory_record_t data)
+{
+    SAVE_DATA_FIELD(size);
+    SAVE_DATA_FIELD(kind);
+    SAVE_DATA_FIELD(operation);
+    SAVE_DATA_FIELD(agent_id);
+    SAVE_DATA_FIELD(queue_id);
+    SAVE_DATA_FIELD(thread_id);
+    SAVE_DATA_FIELD(start_timestamp);
+    SAVE_DATA_FIELD(end_timestamp);
+    SAVE_DATA_FIELD(correlation_id);
+    SAVE_DATA_FIELD(flags);
+}
+
 template <typename ArchiveT>
 void
 save(ArchiveT& ar, rocprofiler_buffer_tracing_correlation_id_retirement_record_t data)
@@ -0,0 +1,48 @@
+#
+#
+#
+cmake_minimum_required(VERSION 3.21.0 FATAL_ERROR)
+
+project(
+    rocprofiler-tests-scratch-memory-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-scratch-memory-tracing-execute COMMAND $<TARGET_FILE:scratch-memory>)
+
+set(scratch-memory-tracing-env
+    "${PRELOAD_ENV}"
+    "HSA_TOOLS_LIB=$<TARGET_FILE:rocprofiler::rocprofiler-shared-library>"
+    "ROCPROFILER_TOOL_OUTPUT_FILE=scratch-memory-tracing-test.json"
+    "LD_LIBRARY_PATH=$<TARGET_FILE_DIR:rocprofiler::rocprofiler-shared-library>:$ENV{LD_LIBRARY_PATH}"
+    )
+
+set_tests_properties(
+    test-scratch-memory-tracing-execute
+    PROPERTIES TIMEOUT 45 LABELS "integration-tests" ENVIRONMENT
+               "${scratch-memory-tracing-env}" FAIL_REGULAR_EXPRESSION
+               "threw an exception")
+
+foreach(FILENAME validate.py pytest.ini conftest.py)
+    configure_file(${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}
+                   ${CMAKE_CURRENT_BINARY_DIR}/${FILENAME} COPYONLY)
+endforeach()
+
+add_test(NAME test-scratch-memory-tracing-validate
+         COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_BINARY_DIR}/validate.py --input
+                 ${CMAKE_CURRENT_BINARY_DIR}/scratch-memory-tracing-test.json)
+
+set_tests_properties(
+    test-scratch-memory-tracing-validate
+    PROPERTIES TIMEOUT 45 LABELS "integration-tests" DEPENDS
+               test-scratch-memory-tracing-execute FAIL_REGULAR_EXPRESSION
+               "threw an exception")
@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+
+import json
+import pytest
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--input",
+        action="store",
+        default="scratch-memory-tracing-test.json",
+        help="Input JSON",
+    )
+
+
+@pytest.fixture
+def input_data(request):
+    filename = request.config.getoption("--input")
+    with open(filename, "r") as inp:
+        return json.load(inp)
@@ -0,0 +1,4 @@
+
+[pytest]
+addopts = --durations=20 -rA -s -vv
+testpaths = validate.py
@@ -0,0 +1,293 @@
+#!/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 test_data_structure(input_data):
+    """verify minimum amount of expected data is present"""
+    data = input_data
+    sdk_data = input_data["rocprofiler-sdk-json-tool"]
+
+    node_exists("rocprofiler-sdk-json-tool", data)
+
+    sdk_data = data["rocprofiler-sdk-json-tool"]
+
+    num_agents = len([agent for agent in sdk_data["agents"] if agent["type"] == 2])
+
+    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("code_objects", sdk_data["callback_records"])
+    node_exists("kernel_symbols", sdk_data["callback_records"])
+    node_exists("hsa_api_traces", sdk_data["callback_records"])
+    node_exists("hip_api_traces", sdk_data["callback_records"], 0)
+    node_exists("scratch_memory_traces", sdk_data["callback_records"], min_len=8)
+
+    node_exists("names", sdk_data["buffer_records"])
+    node_exists("kernel_dispatches", sdk_data["buffer_records"])
+    node_exists("memory_copies", sdk_data["buffer_records"], num_agents)
+    node_exists("hsa_api_traces", sdk_data["buffer_records"])
+    node_exists("hip_api_traces", sdk_data["buffer_records"], 0)
+    node_exists("retired_correlation_ids", sdk_data["buffer_records"])
+    node_exists("scratch_memory_traces", sdk_data["buffer_records"], min_len=8)
+
+
+def test_timestamps(input_data):
+    data = input_data
+    sdk_data = data["rocprofiler-sdk-json-tool"]
+
+    cb_start = {}
+    cb_end = {}
+    for titr in ["hsa_api_traces", "hip_api_traces"]:
+        for itr in sdk_data["callback_records"][titr]:
+            cid = itr["record"]["correlation_id"]["internal"]
+            phase = itr["record"]["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 ["kernel_dispatches", "memory_copies"]:
+        for itr in sdk_data["buffer_records"][titr]:
+            assert itr["start_timestamp"] < itr["end_timestamp"]
+            assert itr["correlation_id"]["internal"] > 0
+            assert itr["correlation_id"]["external"] > 0
+            assert sdk_data["metadata"]["init_time"] < itr["start_timestamp"]
+            assert sdk_data["metadata"]["init_time"] < itr["end_timestamp"]
+            assert sdk_data["metadata"]["fini_time"] > itr["start_timestamp"]
+            assert sdk_data["metadata"]["fini_time"] > itr["end_timestamp"]
+
+            # TODO(Is this check applicable for scratch, which doesn't use any correlation id?)
+            # api_start = cb_start[itr["correlation_id"]["internal"]]
+            # api_end = cb_end[itr["correlation_id"]["internal"]]
+            # assert api_start < itr["start_timestamp"]
+            # assert api_end <= itr["end_timestamp"]
+
+
+def test_internal_correlation_ids(input_data):
+    data = input_data
+    sdk_data = data["rocprofiler-sdk-json-tool"]
+
+    api_corr_ids = []
+    for titr in ["hsa_api_traces", "hip_api_traces"]:
+        for itr in sdk_data["callback_records"][titr]:
+            api_corr_ids.append(itr["record"]["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"]["kernel_dispatches"]:
+        assert itr["correlation_id"]["internal"] in api_corr_ids_unique
+
+    for itr in sdk_data["buffer_records"]["memory_copies"]:
+        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", "hip_api_traces"]:
+        for itr in sdk_data["callback_records"][titr]:
+            assert itr["record"]["correlation_id"]["external"] > 0
+            assert (
+                itr["record"]["thread_id"] == itr["record"]["correlation_id"]["external"]
+            )
+            extern_corr_ids.append(itr["record"]["correlation_id"]["external"])
+
+    extern_corr_ids = list(set(sorted(extern_corr_ids)))
+    for titr in ["hsa_api_traces", "hip_api_traces"]:
+        for itr in sdk_data["buffer_records"][titr]:
+            assert itr["correlation_id"]["external"] > 0
+            assert itr["thread_id"] == itr["correlation_id"]["external"]
+            assert itr["thread_id"] in extern_corr_ids
+            assert itr["correlation_id"]["external"] in extern_corr_ids
+
+    for itr in sdk_data["buffer_records"]["kernel_dispatches"]:
+        assert itr["correlation_id"]["external"] > 0
+        assert itr["correlation_id"]["external"] in extern_corr_ids
+
+    for itr in sdk_data["buffer_records"]["memory_copies"]:
+        assert itr["correlation_id"]["external"] > 0
+        assert itr["correlation_id"]["external"] in extern_corr_ids
+
+
+def op_name(op_name, record):
+    found_op = False
+    op_key = None
+
+    for kind_node in record["names"]["kind_names"]:
+        if kind_node["value"] == op_name:
+            op_key = kind_node["key"]
+
+    for op_node in record["names"]["operation_names"]:
+        if op_node["key"] == op_key:
+            return op_node
+
+
+# Tests above are identical to async-copy. Update as needed
+
+
+def test_scratch_memory_tracking(input_data):
+    sdk_data = input_data["rocprofiler-sdk-json-tool"]
+    callback_records = sdk_data["callback_records"]
+    buffer_records = sdk_data["buffer_records"]
+
+    scratch_callback_data = sdk_data["callback_records"]["scratch_memory_traces"]
+    scratch_buffer_data = sdk_data["buffer_records"]["scratch_memory_traces"]
+
+    cb_op_names = op_name("SCRATCH_MEMORY", callback_records)["value"]
+    bf_op_names = op_name("SCRATCH_MEMORY", buffer_records)["value"]
+
+    assert len(cb_op_names) == 4
+    assert len(bf_op_names) == 4
+
+    # op name -> enum value
+    scratch_cb_op_map = {node["value"]: node["key"] for node in cb_op_names}
+    scratch_bf_op_map = {node["value"]: node["key"] for node in bf_op_names}
+    assert scratch_cb_op_map == scratch_bf_op_map
+
+    scratch_reported_agent_ids = set()
+    detected_agents_ids = set(
+        agent["id"]["handle"] for agent in sdk_data["agents"] if agent["type"] == 2
+    )
+    # check buffering data
+    for node in scratch_buffer_data:
+        assert "size" in node
+        assert "kind" in node
+        assert "flags" in node
+        assert "thread_id" in node
+        assert "end_timestamp" in node
+        assert "start_timestamp" in node
+
+        assert "queue_id" in node
+        assert "agent_id" in node
+        assert "operation" in node
+        assert "handle" in node["queue_id"]
+
+        assert node["start_timestamp"] > 0
+        assert node["start_timestamp"] < node["end_timestamp"]
+
+        scratch_reported_agent_ids.add(node["agent_id"]["handle"])
+
+    assert 2**64 - 1 not in scratch_reported_agent_ids
+    assert scratch_reported_agent_ids == detected_agents_ids
+
+    # { thread-id -> [ events ],  ... }
+    cb_threads = defaultdict(list)
+    bf_threads = defaultdict(list)
+
+    # fetch node["payload"]
+    pl = lambda x: x["payload"]
+    # fetch node["record"]
+    rc = lambda x: x["record"]
+
+    for node in scratch_callback_data:
+        cb_threads[rc(node)["thread_id"]].append(node)
+
+    for node in scratch_buffer_data:
+        bf_threads[node["thread_id"]].append(node)
+
+    for thread_id, nodes in cb_threads.items():
+        assert thread_id > 0
+
+        # start must be followed by end
+        for inx in range(0, len(nodes), 2):
+            this_node = nodes[inx]
+            next_node = nodes[inx + 1]
+
+            assert rc(this_node)["phase"] + 1 == rc(next_node)["phase"]
+            assert rc(this_node)["thread_id"] == rc(next_node)["thread_id"]
+            assert this_node["timestamp"] < next_node["timestamp"]
+
+            # alloc has more data vs free and async reclaim
+            scratch_alloc_node = (
+                this_node["record"]["operation"]
+                == scratch_cb_op_map["SCRATCH_MEMORY_ALLOC"]
+            )
+            if scratch_alloc_node:
+                assert (
+                    pl(this_node)["queue_id"]["handle"]
+                    == pl(next_node)["queue_id"]["handle"]
+                )
+                assert (
+                    this_node["args"]["dispatch_id"] == next_node["args"]["dispatch_id"]
+                )
+                assert "size" in pl(next_node) and pl(next_node)["size"] > 0
+                assert (
+                    "num_slots" in next_node["args"]
+                    and next_node["args"]["num_slots"] > 0
+                )
+                assert "flags" in pl(next_node)
+
+    # callback data and buffer data must agree with each other
+    for bf_thr, bf_nodes in bf_threads.items():
+        cb_nodes = cb_threads[bf_thr]
+
+        for bf_node_inx in range(len(bf_nodes)):
+            # All these 3 should have same data
+            # timestamps are not same as callback records them at
+            # a different instant in time. Callback timestamp
+            # should be more than buffer timestamp
+            bf_node = bf_nodes[bf_node_inx]
+            cb_enter = cb_nodes[bf_node_inx * 2]
+            cb_exit = cb_nodes[bf_node_inx * 2 + 1]
+
+            assert (
+                bf_node["operation"]
+                == rc(cb_enter)["operation"]
+                == rc(cb_exit)["operation"]
+            )
+            assert (
+                bf_op_names[bf_node["operation"]]
+                == cb_op_names[rc(cb_enter)["operation"]]
+                == cb_op_names[rc(cb_exit)["operation"]]
+            )
+
+            assert bf_node["flags"] == pl(cb_exit)["flags"]
+
+            assert (
+                bf_node["thread_id"]
+                == rc(cb_enter)["thread_id"]
+                == rc(cb_exit)["thread_id"]
+            )
+
+
+if __name__ == "__main__":
+    exit_code = pytest.main(["-x", __file__] + sys.argv[1:])
+    sys.exit(exit_code)
@@ -239,6 +239,7 @@ get_callback_tracing_names()
        ROCPROFILER_CALLBACK_TRACING_MARKER_CORE_API,
        ROCPROFILER_CALLBACK_TRACING_MARKER_CONTROL_API,
        ROCPROFILER_CALLBACK_TRACING_MARKER_NAME_API,
+        ROCPROFILER_CALLBACK_TRACING_SCRATCH_MEMORY,
        ROCPROFILER_CALLBACK_TRACING_CODE_OBJECT,
    };

@@ -302,6 +303,7 @@ get_buffer_tracing_names()
        ROCPROFILER_BUFFER_TRACING_MARKER_CONTROL_API,
        ROCPROFILER_BUFFER_TRACING_MARKER_NAME_API,
        ROCPROFILER_BUFFER_TRACING_MEMORY_COPY,
+        ROCPROFILER_BUFFER_TRACING_SCRATCH_MEMORY,
    };

    auto cb_name_info = buffer_name_info{};
@@ -488,12 +490,48 @@ struct marker_api_callback_record_t
    }
 };

+struct scratch_memory_callback_record_t
+{
+    uint64_t                                           timestamp = 0;
+    rocprofiler_callback_tracing_record_t              record    = {};
+    rocprofiler_callback_tracing_scratch_memory_data_t payload   = {};
+
+    template <typename ArchiveT>
+    void save(ArchiveT& ar) const
+    {
+        ar(cereal::make_nvp("timestamp", timestamp));
+        ar(cereal::make_nvp("record", record));
+        ar(cereal::make_nvp("payload", payload));
+
+        if constexpr(std::is_same<ArchiveT, cereal::BinaryOutputArchive>::value ||
+                     std::is_same<ArchiveT, cereal::PortableBinaryOutputArchive>::value)
+        {}
+        else
+        {
+            ar.setNextName("args");
+            ar.startNode();
+            if(payload.args_kind == HSA_AMD_TOOL_EVENT_SCRATCH_ALLOC_START)
+            {
+                ar(cereal::make_nvp("dispatch_id", payload.args.alloc_start.dispatch_id));
+            }
+            else if(payload.args_kind == HSA_AMD_TOOL_EVENT_SCRATCH_ALLOC_END)
+            {
+                ar(cereal::make_nvp("dispatch_id", payload.args.alloc_end.dispatch_id));
+                ar(cereal::make_nvp("size", payload.args.alloc_end.size));
+                ar(cereal::make_nvp("num_slots", payload.args.alloc_end.num_slots));
+            }
+            ar.finishNode();
+        }
+    }
+};
+
 auto code_object_records           = std::deque<code_object_callback_record_t>{};
 auto kernel_symbol_records         = std::deque<kernel_symbol_callback_record_t>{};
 auto hsa_api_cb_records            = std::deque<hsa_api_callback_record_t>{};
 auto marker_api_cb_records         = std::deque<marker_api_callback_record_t>{};
 auto counter_collection_bf_records = std::deque<rocprofiler_record_counter_t>{};
 auto hip_api_cb_records            = std::deque<hip_api_callback_record_t>{};
+auto scratch_memory_cb_records     = std::deque<scratch_memory_callback_record_t>{};

 rocprofiler_thread_id_t
 push_external_correlation();
@@ -688,6 +726,15 @@ tool_tracing_callback(rocprofiler_callback_tracing_record_t record,
        marker_api_cb_records.emplace_back(
            marker_api_callback_record_t{ts, record, *data, std::move(args)});
    }
+    else if(record.kind == ROCPROFILER_CALLBACK_TRACING_SCRATCH_MEMORY)
+    {
+        auto* data =
+            static_cast<rocprofiler_callback_tracing_scratch_memory_data_t*>(record.payload);
+
+        static auto _mutex = std::mutex{};
+        auto        _lk    = std::unique_lock<std::mutex>{_mutex};
+        scratch_memory_cb_records.emplace_back(scratch_memory_callback_record_t{ts, record, *data});
+    }
    else
    {
        throw std::runtime_error{"unsupported callback kind"};
@@ -699,6 +746,7 @@ auto marker_api_bf_records   = std::deque<rocprofiler_buffer_tracing_marker_api_
 auto hip_api_bf_records      = std::deque<rocprofiler_buffer_tracing_hip_api_record_t>{};
 auto kernel_dispatch_records = std::deque<rocprofiler_buffer_tracing_kernel_dispatch_record_t>{};
 auto memory_copy_records     = std::deque<rocprofiler_buffer_tracing_memory_copy_record_t>{};
+auto scratch_memory_records  = std::deque<rocprofiler_buffer_tracing_scratch_memory_record_t>{};
 auto corr_id_retire_records =
    std::deque<rocprofiler_buffer_tracing_correlation_id_retirement_record_t>{};

@@ -783,6 +831,13 @@ tool_tracing_buffered(rocprofiler_context_id_t /*context*/,

                memory_copy_records.emplace_back(*record);
            }
+            else if(header->kind == ROCPROFILER_BUFFER_TRACING_SCRATCH_MEMORY)
+            {
+                auto* record = static_cast<rocprofiler_buffer_tracing_scratch_memory_record_t*>(
+                    header->payload);
+
+                scratch_memory_records.emplace_back(*record);
+            }
            else if(header->kind == ROCPROFILER_BUFFER_TRACING_CORRELATION_ID_RETIREMENT)
            {
                auto* record =
@@ -854,6 +909,7 @@ rocprofiler_context_id_t marker_api_buffered_ctx = {};
 rocprofiler_context_id_t kernel_dispatch_ctx     = {};
 rocprofiler_context_id_t memory_copy_ctx         = {};
 rocprofiler_context_id_t counter_collection_ctx  = {};
+rocprofiler_context_id_t scratch_memory_ctx      = {};
 rocprofiler_context_id_t corr_id_retire_ctx      = {};
 // buffers
 rocprofiler_buffer_id_t hsa_api_buffered_buffer    = {};
@@ -862,6 +918,7 @@ rocprofiler_buffer_id_t marker_api_buffered_buffer = {};
 rocprofiler_buffer_id_t kernel_dispatch_buffer     = {};
 rocprofiler_buffer_id_t memory_copy_buffer         = {};
 rocprofiler_buffer_id_t counter_collection_buffer  = {};
+rocprofiler_buffer_id_t scratch_memory_buffer      = {};
 rocprofiler_buffer_id_t corr_id_retire_buffer      = {};

 auto contexts = std::unordered_map<std::string_view, rocprofiler_context_id_t*>{
@@ -875,18 +932,18 @@ auto contexts = std::unordered_map<std::string_view, rocprofiler_context_id_t*>{
    {"KERNEL_DISPATCH", &kernel_dispatch_ctx},
    {"MEMORY_COPY", &memory_copy_ctx},
    {"COUNTER_COLLECTION", &counter_collection_ctx},
+    {"SCRATCH_MEMORY", &scratch_memory_ctx},
    {"CORRELATION_ID_RETIREMENT", &corr_id_retire_ctx},
 };

-auto buffers = std::array<rocprofiler_buffer_id_t*, 7>{
-    &hsa_api_buffered_buffer,
-    &hip_api_buffered_buffer,
-    &marker_api_buffered_buffer,
-    &kernel_dispatch_buffer,
-    &memory_copy_buffer,
-    &counter_collection_buffer,
-    &corr_id_retire_buffer,
-};
+auto buffers = std::array<rocprofiler_buffer_id_t*, 8>{&hsa_api_buffered_buffer,
+                                                       &hip_api_buffered_buffer,
+                                                       &marker_api_buffered_buffer,
+                                                       &kernel_dispatch_buffer,
+                                                       &memory_copy_buffer,
+                                                       &scratch_memory_buffer,
+                                                       &counter_collection_buffer,
+                                                       &corr_id_retire_buffer};

 auto agents = std::vector<rocprofiler_agent_t>{};

@@ -987,6 +1044,15 @@ tool_init(rocprofiler_client_finalize_t fini_func, void* tool_data)
                                                       nullptr),
        "hsa api tracing service configure");

+    ROCPROFILER_CALL(
+        rocprofiler_configure_callback_tracing_service(scratch_memory_ctx,
+                                                       ROCPROFILER_CALLBACK_TRACING_SCRATCH_MEMORY,
+                                                       nullptr,
+                                                       0,
+                                                       tool_tracing_callback,
+                                                       nullptr),
+        "scratch memory tracing service configure");
+
    constexpr auto buffer_size = 8192;
    constexpr auto watermark   = 7936;

@@ -1035,6 +1101,15 @@ tool_init(rocprofiler_client_finalize_t fini_func, void* tool_data)
                                               &memory_copy_buffer),
                     "buffer creation");

+    ROCPROFILER_CALL(rocprofiler_create_buffer(scratch_memory_ctx,
+                                               buffer_size,
+                                               watermark,
+                                               ROCPROFILER_BUFFER_POLICY_LOSSLESS,
+                                               tool_tracing_buffered,
+                                               tool_data,
+                                               &scratch_memory_buffer),
+                     "buffer creation");
+
    ROCPROFILER_CALL(rocprofiler_create_buffer(corr_id_retire_ctx,
                                               buffer_size,
                                               watermark,
@@ -1111,6 +1186,14 @@ tool_init(rocprofiler_client_finalize_t fini_func, void* tool_data)
                                                     memory_copy_buffer),
        "buffer tracing service for memory copy configure");

+    ROCPROFILER_CALL(
+        rocprofiler_configure_buffer_tracing_service(scratch_memory_ctx,
+                                                     ROCPROFILER_BUFFER_TRACING_SCRATCH_MEMORY,
+                                                     nullptr,
+                                                     0,
+                                                     scratch_memory_buffer),
+        "buffer tracing service for scratch memory configure");
+
    ROCPROFILER_CALL(rocprofiler_configure_buffer_tracing_service(
                         corr_id_retire_ctx,
                         ROCPROFILER_BUFFER_TRACING_CORRELATION_ID_RETIREMENT,
@@ -1247,6 +1330,7 @@ tool_fini(void* tool_data)
              << ", marker_api_callback_records=" << marker_api_cb_records.size()
              << ", kernel_dispatch_records=" << kernel_dispatch_records.size()
              << ", memory_copy_records=" << memory_copy_records.size()
+              << ", scratch_memory_records=" << scratch_memory_records.size()
              << ", hsa_api_bf_records=" << hsa_api_bf_records.size()
              << ", hip_api_bf_records=" << hip_api_bf_records.size()
              << ", marker_api_bf_records=" << marker_api_bf_records.size()
@@ -1334,6 +1418,7 @@ write_json(call_stack_t* _call_stack)
            json_ar(cereal::make_nvp("hsa_api_traces", hsa_api_cb_records));
            json_ar(cereal::make_nvp("hip_api_traces", hip_api_cb_records));
            json_ar(cereal::make_nvp("marker_api_traces", marker_api_cb_records));
+            json_ar(cereal::make_nvp("scratch_memory_traces", scratch_memory_cb_records));
        } catch(std::exception& e)
        {
            std::cerr << "[" << getpid() << "][" << __FUNCTION__
@@ -1349,6 +1434,7 @@ write_json(call_stack_t* _call_stack)
            json_ar(cereal::make_nvp("names", buffer_name_info));
            json_ar(cereal::make_nvp("kernel_dispatches", kernel_dispatch_records));
            json_ar(cereal::make_nvp("memory_copies", memory_copy_records));
+            json_ar(cereal::make_nvp("scratch_memory_traces", scratch_memory_records));
            json_ar(cereal::make_nvp("hsa_api_traces", hsa_api_bf_records));
            json_ar(cereal::make_nvp("hip_api_traces", hip_api_bf_records));
            json_ar(cereal::make_nvp("marker_api_traces", marker_api_bf_records));