Report page migration events as start/end (#793)

* Squashed commit of the following:

commit b76f2635f4b65599f03812a73d0cf410f5ada213
Author: Mythreya <mythreya.kuricheti@amd.com>
Date:   Fri Apr 26 00:29:09 2024 +0000

    Changed for PR feedback

commit bedb8ad566ff42fbf117b19202c26c507abcf8ac
Author: Jonathan R. Madsen <jonathanrmadsen@gmail.com>
Date:   Thu Apr 25 19:20:06 2024 -0500

    Fix installation

commit a98f8a69459a1450a1be9c98e20b3c1e7f2568c2
Author: Jonathan R. Madsen <jonathanrmadsen@gmail.com>
Date:   Thu Apr 25 19:16:35 2024 -0500

    Restructure the headers

commit 46489a020ffafdd5f4ce3f580469ff233ef67fe1
Author: Mythreya <mythreya.kuricheti@amd.com>
Date:   Tue Apr 23 23:31:10 2024 +0000

    Update hsa include

commit 8e795282cce348fc6aa736b7857b21aeb32aa20a
Author: Mythreya <mythreya.kuricheti@amd.com>
Date:   Tue Apr 23 23:02:32 2024 +0000

    Report page migration events as start/end

    * Updated tests accordingly
    * Page migration events are reported independently

commit 8784e5ad4895a626a2a8e4ac12f8021b34172bd4
Author: Mythreya <mythreya.kuricheti@amd.com>
Date:   Tue Apr 16 17:01:57 2024 +0000

    Update handling of dropped page migration events

    Previously, we dropped all locally buffered events when we detect that
    KFD has dropped some events. This may drop too many pending events too eagerly.

    When we receive an end event and cannot find the corresponding start,
    we can be sure that KFD has dropped some events in the immediate past.

    When this happens, we look through all locally buffered events and report
    the start events that are older than 10s as partial events --- they have
    no "end" information (we expect that the end events have been dropped).

    We also set the polling timeout to 10s to prevent the local buffer from
    getting too large with events waiting to be paired up.

    Updated tests

commit 2e8e0b07eeda9b5990e1ae8d28dcd3a035ce38e1
Author: Mythreya <mythreya.kuricheti@amd.com>
Date:   Tue Apr 16 17:01:31 2024 +0000

    Docs for triggers

* Fix page migration sample

* Fix hasher, kfd install

* Add hsa include
* Install KFD include dir

* Updates from code review

- single timestamp field
- node_id -> agent_id
- from_node -> from_agent
- to_node -> to_agent

* Misc revisions

* Remove page-migration install target

* Update page-migration pytest

* Tweak to serialization

* Address PR comments

* Update page-migration test

* Add cli args, update iterations

* Address PR comments

* Add abi.cpp for static_asserts
* Update page_migration gtest with only runtime tests
* Moved helpers into utils.hpp

---------

Co-authored-by: Jonathan R. Madsen <jonathanrmadsen@gmail.com>

tests/bin/page-migration/CMakeLists.txt

@@ -40,8 +40,3 @@ target_compile_options(page-migration PRIVATE -W -Wall -Wextra -Wpedantic -Wshad
 
 find_package(Threads REQUIRED)
 target_link_libraries(page-migration PRIVATE Threads::Threads)
-
-install(
-    TARGETS page-migration
-    DESTINATION bin
-    COMPONENT tests)

tests/bin/page-migration/page-migration.cpp

@@ -27,13 +27,14 @@
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
-#include <iomanip>
 #include <iostream>
 #include <mutex>
 #include <sstream>
 #include <stdexcept>
 #include <string>
 #include <string_view>
+#include <thread>
+#include <vector>
 
 #include <fcntl.h>
 #include <sys/mman.h>
@@ -117,14 +118,14 @@ private:
 };
 
 int
-main()
+run_test(int num_iter)
 {
     using namespace std::chrono_literals;
 
-    static constexpr auto NUM_PAGES       = 16;
-    const auto            PAGE_SIZE_BYTES = ::sysconf(_SC_PAGE_SIZE);
+    constexpr size_t NUM_PAGES       = 512;
+    const size_t     PAGE_SIZE_BYTES = ::sysconf(_SC_PAGE_SIZE);
 
-    size_t elem_count = (NUM_PAGES * PAGE_SIZE_BYTES) / sizeof(size_t);  // one page?
+    const size_t elem_count = (NUM_PAGES * PAGE_SIZE_BYTES) / sizeof(size_t);
 
     auto  alloc  = mmap_allocator(NUM_PAGES);
     void* data_v = alloc.get<void>();
@@ -142,8 +143,10 @@ main()
 
     HIP_API_CALL(hipHostRegister(data, elem_count * sizeof(size_t), hipHostRegisterDefault));
 
-    char maps[1024 * 1024];
-    std::memset(maps, '\0', 1024 * 1024);
+    constexpr size_t MAPS_BUFFER_SIZE = 1024 * 1024;
+
+    char maps[MAPS_BUFFER_SIZE];
+    std::memset(maps, '\0', MAPS_BUFFER_SIZE);
     auto fd = open("/proc/self/maps", O_RDONLY | O_CLOEXEC);
 
     if(fd == -1)
@@ -153,7 +156,7 @@ main()
         exit(-1);
     }
 
-    auto bytes = read(fd, maps, 1024 * 1024 - 1);
+    auto bytes = read(fd, maps, MAPS_BUFFER_SIZE - 1);
     if(bytes == -1)
     {
         auto ecode = errno;
@@ -190,33 +193,96 @@ main()
         }
     }
 
-    for(int iter = 0; iter < 1000; ++iter)
+    hipStream_t stream{};
+    HIP_API_CALL(hipStreamCreate(&stream));
+
+    for(int iter = 0; iter < num_iter; ++iter)
     {
         for(size_t i = 0; i < elem_count; ++i)
+        {
             data[i] = i;
+        }
 
         // std::cout << "launching..." << std::endl;
-        hipLaunchKernelGGL(kernel, 128, 64, 0, 0, data, elem_count);
+        hipLaunchKernelGGL(kernel, 1024, 1024, 0, stream, data, elem_count);
 
-        // std::cout << "syncing..." << std::endl;
-        HIP_API_CALL(hipDeviceSynchronize());
+        HIP_API_CALL(hipStreamSynchronize(stream));
 
-        // std::cout << "checking..." << std::endl;
         for(size_t i = 0; i < elem_count; ++i)
         {
-            if(data[i] != (i * 2))
+            const auto data_i = data[i];
+            if(data_i != (i * 2))
             {
                 auto msg = std::stringstream{};
                 msg << "GPU computed value at " << i << " in iteration " << iter
-                    << " is incorrect. Expected " << (i * 2) << ", found " << data[i];
+                    << " is incorrect. Expected " << (i * 2) << ", found " << data_i;
                 throw std::runtime_error{msg.str()};
             }
         }
-
-        std::cout << "Iteration " << std::setw(2) << iter << ": correct\n" << std::flush;
     }
 
+    HIP_API_CALL(hipStreamDestroy(stream));
     HIP_API_CALL(hipDeviceSynchronize());
 
     return 0;
 }
+
+int
+main(int argc, const char** argv)
+{
+    using namespace std::chrono_literals;
+
+    const auto usage_msg = [](const char** _argv) {
+        fprintf(stderr, "usage: %s <NUMBER OF THREADS> <ITERATIONS PER THREAD>\n", _argv[0]);
+    };
+
+    if(argc != 3)
+    {
+        usage_msg(argv);
+        exit(EXIT_FAILURE);
+    }
+
+    for(int i = 1; i < argc; ++i)
+    {
+        auto _arg = std::string{argv[i]};
+        if(_arg == "?" || _arg == "-h" || _arg == "--help")
+        {
+            usage_msg(argv);
+            exit(EXIT_SUCCESS);
+        }
+    }
+
+    const auto num_threads = std::atoi(argv[1]);
+    if(num_threads < 1)
+    {
+        fprintf(stderr, "Error: Invalid value %d for num_threads (min 1)\n", num_threads);
+        exit(EXIT_FAILURE);
+    }
+
+    const auto num_iter = std::atoi(argv[2]);
+    if(num_iter < 1)
+    {
+        fprintf(stderr, "Error: Invalid value %d for num_iter (min 1)\n", num_iter);
+        exit(EXIT_FAILURE);
+    }
+
+    run_test(num_iter);
+
+    std::vector<std::thread> threads;
+    threads.reserve(num_threads);
+
+    std::cerr << "Running " << num_iter << " iterations/thread on " << num_threads << " threads\n";
+
+    for(auto i = 0; i < num_threads; ++i)
+    {
+        threads.emplace_back([_num_iter = num_iter]() { run_test(_num_iter); });
+    }
+
+    std::cerr << "Waiting for threads\n";
+    for(auto& t : threads)
+    {
+        t.join();
+    }
+
+    return 0;
+}

tests/page-migration/CMakeLists.txt

@@ -17,7 +17,7 @@ else()
     set(PRELOAD_ENV "LD_PRELOAD=$<TARGET_FILE:rocprofiler-sdk-json-tool>")
 endif()
 
-add_test(NAME test-page-migration-execute COMMAND $<TARGET_FILE:page-migration>)
+add_test(NAME test-page-migration-execute COMMAND $<TARGET_FILE:page-migration> 4 1024)
 
 set(page-migration-env
     "${PRELOAD_ENV}"

tests/page-migration/conftest.py

@@ -2,6 +2,7 @@
 
 import json
 import pytest
+from rocprofiler_sdk.pytest_utils.dotdict import dotdict
 
 
 def pytest_addoption(parser):
@@ -24,4 +25,4 @@ def input_data(request):
         return pytest.skip(
             "Skipping test because KFD does not support SVM event reporting"
         )
-    return data
+    return dotdict(data)

tests/page-migration/validate.py

@@ -271,71 +271,113 @@ def get_allocated_pages(callback_records):
             assert "hostPtr" in itr["args"].keys(), f"{itr}"
             host_register_record.append(itr)
 
-    assert len(host_register_record) == 1
-    alloc_size = int(host_register_record[0]["args"]["sizeBytes"], 10)
-    start_addr = int(host_register_record[0]["args"]["hostPtr"], 16)
-    end_addr = start_addr + alloc_size
+    num_host_register_calls = len(host_register_record)
+    assert num_host_register_calls == 5, "Expected 5 hipHostRegister calls in test"
 
-    return start_addr, end_addr, alloc_size
+    ret = []
+    for i in range(num_host_register_calls):
+        alloc_size = int(host_register_record[0]["args"]["sizeBytes"], 10)
+        start_addr = int(host_register_record[0]["args"]["hostPtr"], 16)
+        end_addr = start_addr + alloc_size
+        ret.append((start_addr, end_addr))
+    return ret
+
+
+def validate_node(id, nodes):
+    assert id.handle in nodes
 
 
 def test_page_migration_data(input_data):
     data = input_data
     sdk_data = data["rocprofiler-sdk-json-tool"]
-    buffer_records = sdk_data["buffer_records"]
-    callback_records = sdk_data["callback_records"]
-    page_migration_buffers = buffer_records["page_migration"]
+    buffer_records = sdk_data.buffer_records
+    callback_records = sdk_data.callback_records
+    page_migtation_buffers = buffer_records.page_migration
 
     _, bf_op_names = get_operation(buffer_records, "PAGE_MIGRATION")
     assert bf_op_names[0] == "PAGE_MIGRATION_NONE"
-    assert "PAGE_MIGRATION_PAGE_MIGRATE" in bf_op_names
-    assert len(bf_op_names) == 5
 
-    node_ids = set(x["gpu_id"] for x in sdk_data["agents"])
-    start_addr, end_addr, alloc_size = get_allocated_pages(callback_records)
+    for op_name in bf_op_names:
+        assert "PAGE_MIGRATION" in op_name
 
-    assert start_addr < end_addr and start_addr + alloc_size == end_addr
-    assert int(alloc_size) == 16 * 4096  # We allocated 16 pages in the test
+    assert len(bf_op_names) == 8
 
-    # PID must be same
-    assert len(set(r["pid"] for r in page_migration_buffers)) == 1
+    nodes = set(x.id.handle for x in sdk_data.agents)
+    allocations = get_allocated_pages(callback_records)
 
-    for r in page_migration_buffers:
-        op = r["operation"]
+    for start_addr, end_addr in allocations:
 
-        assert r["size"] == 136
-        assert op != 0 and bf_op_names[op] != "PAGE_MIGRATION_NONE"
-        assert bf_op_names[op].lower().replace("page_migration_", "") in r.keys()
+        assert (
+            start_addr < end_addr
+        ), "Expected start address less than end address for mmap range"
+        alloc_size = end_addr - start_addr
+        assert int(alloc_size) == 512 * 4096  # We allocated 512 pages in the test
 
-        if "page_migrate" in r:
-            assert r["page_migrate"]["from_node"] in node_ids
-            assert r["page_migrate"]["to_node"] in node_ids
-            assert r["page_migrate"]["prefetch_node"] in node_ids
-            assert r["page_migrate"]["preferred_node"] in node_ids
-            assert r["page_migrate"]["trigger"] >= 0
+        # PID must be same
+        assert len(set(r.pid for r in page_migtation_buffers)) == 1
 
-        if "queue_suspend" in r:
-            assert r["queue_suspend"]["trigger"] >= 0
-            assert r["queue_suspend"]["node_id"] in node_ids
+        for r in page_migtation_buffers:
+            op = r.operation
 
-        if "unmap_from_gpu" in r:
-            assert r["unmap_from_gpu"]["trigger"] >= 0
-            # unmap is "instantaneous"
-            assert 0 < r["start_timestamp"] == r["end_timestamp"]
-        else:
-            assert 0 < r["start_timestamp"] < r["end_timestamp"]
+            assert r.size == 160
+            assert op != 0 and bf_op_names[op] != "PAGE_MIGRATION_NONE"
+            assert bf_op_names[op].lower().replace("page_migration_", "") in r.keys()
 
-    # Check for events with our page
-    for r in page_migration_buffers:
+            if "page_fault_start" in r:
+                arg = r.page_fault_start
+                assert arg.read_fault < 2
+                validate_node(arg.agent_id, nodes)
+                assert arg.address > 0
 
-        if "page_migrate" in r and r["page_migrate"]["start_addr"] == start_addr:
-            assert end_addr == r["page_migrate"]["end_addr"]
+            if "page_fault_end" in r:
+                arg = r.page_fault_end
+                assert arg.migrated < 2
+                validate_node(arg.agent_id, nodes)
+                assert arg.address > 0
 
-        if "unmap_from_gpu" in r and r["unmap_from_gpu"]["start_addr"] == start_addr:
-            assert end_addr == r["unmap_from_gpu"]["end_addr"]
+            if "page_migrate_start" in r:
+                arg = r.page_migrate_start
+                assert (
+                    0 < arg.start_addr < arg.end_addr
+                ), "Expected start addr to be less than end addr"
+                if arg.start_addr == start_addr:
+                    assert arg.end_addr == end_addr
+                validate_node(arg.from_agent, nodes)
+                validate_node(arg.to_agent, nodes)
+                validate_node(arg.prefetch_agent, nodes)
+                validate_node(arg.preferred_agent, nodes)
+                assert 0 <= arg.trigger < 4
 
-    # TODO: Check if a migrate a->b is paired up with b->a
-    # It may not always be reported towards app finalization
+            if "page_migrate_end" in r:
+                arg = r.page_migrate_end
+                assert (
+                    0 < arg.start_addr < arg.end_addr
+                ), "Expected start addr to be less than end addr"
+                if arg.start_addr == start_addr:
+                    assert arg.end_addr == end_addr
+                validate_node(arg.from_agent, nodes)
+                validate_node(arg.to_agent, nodes)
+                assert 0 <= arg.trigger < 4
+
+            if "queue_eviction" in r:
+                arg = r.queue_eviction
+                validate_node(arg.agent_id, nodes)
+                assert 0 <= arg.trigger < 6
+
+            if "queue_restore" in r:
+                arg = r.queue_restore
+                assert arg.rescheduled < 2
+                validate_node(arg.agent_id, nodes)
+
+            if "unmap_from_gpu" in r:
+                arg = r.unmap_from_gpu
+                assert (
+                    0 < arg.start_addr < arg.end_addr
+                ), "Expected start addr to be less than end addr"
+                if arg.start_addr == start_addr:
+                    assert arg.end_addr == end_addr
+                validate_node(arg.agent_id, nodes)
+                assert 0 <= arg.trigger < 3
 
 
 if __name__ == "__main__":

tests/tools/json-tool.cpp

@@ -1277,7 +1277,7 @@ tool_init(rocprofiler_client_finalize_t fini_func, void* tool_data)
                                                          0,
                                                          page_migration_buffer);
 
-        constexpr auto message = "page migration service for memory copy configure";
+        constexpr auto message = "buffer tracing service for page migration configure";
         if(page_migration_status == ROCPROFILER_STATUS_ERROR_INCOMPATIBLE_KERNEL)
             std::cerr << message
                       << " failed: " << rocprofiler_get_status_string(page_migration_status)