Adding wrappers on HSA for executable load/unload and allowing multiple agents per context on ATT (#951)

* Codeobj wrappers around HSA calls for ATT * Formatting * Bookeeping * Tidy * Tidy * Update source/lib/rocprofiler-sdk/thread_trace/code_object.hpp Co-authored-by: Vladimir Indic <139573562+vlaindic@users.noreply.github.com> * Update source/lib/rocprofiler-sdk/thread_trace/att_core.hpp Co-authored-by: Vladimir Indic <139573562+vlaindic@users.noreply.github.com> * Variable naming --------- Co-authored-by: Vladimir Indic <139573562+vlaindic@users.noreply.github.com>
2024-06-25 13:46:13 -03:00
parent b62ba5f096
commit 8da0c35079
9 changed files with 314 additions and 155 deletions
@@ -0,0 +1,147 @@
+// MIT License
+//
+// Copyright (c) 2024 ROCm Developer Tools
+//
+// 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 "lib/rocprofiler-sdk/thread_trace/code_object.hpp"
+#include "lib/rocprofiler-sdk/code_object/code_object.hpp"
+
+namespace rocprofiler
+{
+namespace thread_trace
+{
+namespace code_object
+{
+std::mutex                         CodeobjCallbackRegistry::mut;
+std::set<CodeobjCallbackRegistry*> CodeobjCallbackRegistry::all_registries{};
+
+CodeobjCallbackRegistry::CodeobjCallbackRegistry(LoadCallback _ld, UnloadCallback _unld)
+: ld_fn(std::move(_ld))
+, unld_fn(std::move(_unld))
+{
+    std::unique_lock<std::mutex> lg(mut);
+    all_registries.insert(this);
+}
+
+CodeobjCallbackRegistry::~CodeobjCallbackRegistry()
+{
+    std::unique_lock<std::mutex> lg(mut);
+    all_registries.erase(this);
+}
+
+void
+CodeobjCallbackRegistry::Load(rocprofiler_agent_id_t agent,
+                              uint64_t               id,
+                              uint64_t               addr,
+                              uint64_t               size)
+{
+    std::unique_lock<std::mutex> lg(mut);
+    for(auto* reg : all_registries)
+        reg->ld_fn(agent, id, addr, size);
+}
+
+void
+CodeobjCallbackRegistry::Unload(uint64_t id)
+{
+    std::unique_lock<std::mutex> lg(mut);
+    for(auto* reg : all_registries)
+        reg->unld_fn(id);
+}
+
+void
+CodeobjCallbackRegistry::IterateLoaded() const
+{
+    std::unique_lock<std::mutex> lg(mut);
+
+    rocprofiler::code_object::iterate_loaded_code_objects(
+        [&](const rocprofiler::code_object::hsa::code_object& code_object) {
+            const auto& data = code_object.rocp_data;
+            ld_fn(data.rocp_agent, data.code_object_id, data.load_delta, data.load_size);
+        });
+}
+
+namespace
+{
+auto&
+get_freeze_function()
+{
+    static decltype(::hsa_executable_freeze)* _v = nullptr;
+    return _v;
+}
+
+auto&
+get_destroy_function()
+{
+    static decltype(::hsa_executable_destroy)* _v = nullptr;
+    return _v;
+}
+
+hsa_status_t
+executable_freeze(hsa_executable_t executable, const char* options)
+{
+    // Call underlying function
+    hsa_status_t status = CHECK_NOTNULL(get_freeze_function())(executable, options);
+    if(status != HSA_STATUS_SUCCESS) return status;
+
+    rocprofiler::code_object::iterate_loaded_code_objects(
+        [&](const rocprofiler::code_object::hsa::code_object& code_object) {
+            if(code_object.hsa_executable != executable) return;
+
+            const auto& data = code_object.rocp_data;
+            CodeobjCallbackRegistry::Load(
+                data.rocp_agent, data.code_object_id, data.load_delta, data.load_size);
+        });
+
+    return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t
+executable_destroy(hsa_executable_t executable)
+{
+    rocprofiler::code_object::iterate_loaded_code_objects(
+        [&](const rocprofiler::code_object::hsa::code_object& code_object) {
+            if(code_object.hsa_executable == executable)
+                CodeobjCallbackRegistry::Unload(code_object.rocp_data.code_object_id);
+        });
+
+    // Call underlying function
+    return CHECK_NOTNULL(get_destroy_function())(executable);
+}
+}  // namespace
+
+void
+initialize(HsaApiTable* table)
+{
+    (void) table;
+    auto& core_table = *table->core_;
+
+    get_freeze_function()                = CHECK_NOTNULL(core_table.hsa_executable_freeze_fn);
+    get_destroy_function()               = CHECK_NOTNULL(core_table.hsa_executable_destroy_fn);
+    core_table.hsa_executable_freeze_fn  = executable_freeze;
+    core_table.hsa_executable_destroy_fn = executable_destroy;
+    LOG_IF(FATAL, get_freeze_function() == core_table.hsa_executable_freeze_fn)
+        << "infinite recursion";
+    LOG_IF(FATAL, get_destroy_function() == core_table.hsa_executable_destroy_fn)
+        << "infinite recursion";
+}
+
+}  // namespace code_object
+}  // namespace thread_trace
+}  // namespace rocprofiler