Fixing Code Object Data Race and Thread Safety & Adding validation test (#2014)

2025-11-26 20:28:52 -06:00
parent 5b75ec6a09
commit ed42157c31
11 changed files with 679 additions and 92 deletions
@@ -0,0 +1,202 @@
+// MIT License
+//
+// Copyright (c) 2023-2025 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.
+
+/**
+ * @file tests/bin/code-object-multi-threaded/client.cpp
+ *
+ * @brief ROCProfiler tool that tracks code object operations
+ *
+ * This tool uses rocprofiler-sdk callback tracing to monitor code object
+ * load/unload operations during multi-threaded execution, verifying that
+ * the thread-safety fixes work correctly.
+ */
+
+#include <rocprofiler-sdk/callback_tracing.h>
+#include <rocprofiler-sdk/fwd.h>
+#include <rocprofiler-sdk/registration.h>
+#include <rocprofiler-sdk/rocprofiler.h>
+
+#include <atomic>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <iostream>
+#include <mutex>
+#include <unordered_map>
+
+namespace
+{
+using code_obj_load_data_t = rocprofiler_callback_tracing_code_object_load_data_t;
+using kernel_symbol_data_t = rocprofiler_callback_tracing_code_object_kernel_symbol_register_data_t;
+
+// Thread-safe counters
+std::atomic<uint64_t> g_code_object_load_count{0};
+std::atomic<uint64_t> g_code_object_unload_count{0};
+std::atomic<uint64_t> g_kernel_symbol_load_count{0};
+std::atomic<uint64_t> g_kernel_symbol_unload_count{0};
+
+// Track which threads invoke callbacks
+std::mutex                             g_thread_mutex;
+std::unordered_map<uint64_t, uint64_t> g_thread_callback_counts;
+
+rocprofiler_client_id_t*      client_id        = nullptr;
+rocprofiler_client_finalize_t client_fini_func = nullptr;
+rocprofiler_context_id_t      client_ctx       = {};
+
+void
+codeobj_tracing_callback(rocprofiler_callback_tracing_record_t record,
+                         rocprofiler_user_data_t* /* user_data */,
+                         void* /* callback_data */)
+{
+    if(record.kind != ROCPROFILER_CALLBACK_TRACING_CODE_OBJECT) return;
+
+    // Track thread activity
+    {
+        std::lock_guard<std::mutex> lock(g_thread_mutex);
+        g_thread_callback_counts[record.thread_id]++;
+    }
+
+    if(record.operation == ROCPROFILER_CODE_OBJECT_LOAD)
+    {
+        if(record.phase == ROCPROFILER_CALLBACK_PHASE_LOAD)
+        {
+            g_code_object_load_count.fetch_add(1, std::memory_order_relaxed);
+        }
+        else if(record.phase == ROCPROFILER_CALLBACK_PHASE_UNLOAD)
+        {
+            g_code_object_unload_count.fetch_add(1, std::memory_order_relaxed);
+        }
+    }
+    else if(record.operation == ROCPROFILER_CODE_OBJECT_DEVICE_KERNEL_SYMBOL_REGISTER)
+    {
+        if(record.phase == ROCPROFILER_CALLBACK_PHASE_LOAD)
+        {
+            g_kernel_symbol_load_count.fetch_add(1, std::memory_order_relaxed);
+        }
+        else if(record.phase == ROCPROFILER_CALLBACK_PHASE_UNLOAD)
+        {
+            g_kernel_symbol_unload_count.fetch_add(1, std::memory_order_relaxed);
+        }
+    }
+}
+
+void
+tool_fini(void* /* tool_data */)
+{
+    std::cout << "\n=== ROCProfiler Multi-Threaded Code Object Test Results ===\n";
+    std::cout << "Code objects loaded:        " << g_code_object_load_count.load() << "\n";
+    std::cout << "Code objects unloaded:      " << g_code_object_unload_count.load() << "\n";
+    std::cout << "Kernel symbols loaded:      " << g_kernel_symbol_load_count.load() << "\n";
+    std::cout << "Kernel symbols unloaded:    " << g_kernel_symbol_unload_count.load() << "\n";
+
+    size_t num_threads_with_callbacks = 0;
+    {
+        std::lock_guard<std::mutex> lock(g_thread_mutex);
+        num_threads_with_callbacks = g_thread_callback_counts.size();
+        std::cout << "Threads that invoked callbacks: " << num_threads_with_callbacks << "\n";
+    }
+    std::cout << "===========================================================\n";
+
+    // Verify we actually traced something
+    if(g_code_object_load_count.load() == 0)
+    {
+        std::cerr << "ERROR: No code objects were traced!\n";
+        std::abort();
+    }
+
+    if(g_kernel_symbol_load_count.load() == 0)
+    {
+        std::cerr << "ERROR: No kernel symbols were traced!\n";
+        std::abort();
+    }
+
+    // Verify multi-threaded execution
+    if(num_threads_with_callbacks < 2)
+    {
+        std::cerr << "ERROR: Expected callbacks from multiple threads, got "
+                  << num_threads_with_callbacks << "\n";
+        std::abort();
+    }
+
+    std::cout << "Test PASSED: Successfully traced multi-threaded code object loading!\n";
+}
+
+int
+tool_init(rocprofiler_client_finalize_t fini_func, void* tool_data)
+{
+    client_fini_func = fini_func;
+
+    auto status = rocprofiler_create_context(&client_ctx);
+    if(status != ROCPROFILER_STATUS_SUCCESS)
+    {
+        std::cerr << "Failed to create context\n";
+        return -1;
+    }
+
+    status =
+        rocprofiler_configure_callback_tracing_service(client_ctx,
+                                                       ROCPROFILER_CALLBACK_TRACING_CODE_OBJECT,
+                                                       nullptr,
+                                                       0,
+                                                       codeobj_tracing_callback,
+                                                       tool_data);
+
+    if(status != ROCPROFILER_STATUS_SUCCESS)
+    {
+        std::cerr << "Failed to configure code object tracing\n";
+        return -1;
+    }
+
+    int valid_ctx = 0;
+    status        = rocprofiler_context_is_valid(client_ctx, &valid_ctx);
+    if(status != ROCPROFILER_STATUS_SUCCESS || valid_ctx == 0)
+    {
+        std::cerr << "Context is not valid\n";
+        return -1;
+    }
+
+    status = rocprofiler_start_context(client_ctx);
+    if(status != ROCPROFILER_STATUS_SUCCESS)
+    {
+        std::cerr << "Failed to start context\n";
+        return -1;
+    }
+
+    std::cout << "ROCProfiler multi-threaded code object tool initialized\n";
+    return 0;
+}
+}  // namespace
+
+extern "C" rocprofiler_tool_configure_result_t*
+rocprofiler_configure(uint32_t /* version */,
+                      const char* /* runtime_version */,
+                      uint32_t /* priority */,
+                      rocprofiler_client_id_t* id)
+{
+    id->name  = "CodeObjectMultiThreadedClient";
+    client_id = id;
+
+    static auto cfg = rocprofiler_tool_configure_result_t{
+        sizeof(rocprofiler_tool_configure_result_t), &tool_init, &tool_fini, nullptr};
+
+    return &cfg;
+}