diff --git a/.github/workflows/rocprofiler-sdk-rocm_release_compatibility.yml b/.github/workflows/rocprofiler-sdk-rocm_release_compatibility.yml index 0ab9053557a0..face781447d6 100644 --- a/.github/workflows/rocprofiler-sdk-rocm_release_compatibility.yml +++ b/.github/workflows/rocprofiler-sdk-rocm_release_compatibility.yml @@ -90,6 +90,14 @@ jobs: update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 10 --slave /usr/bin/g++ g++ /usr/bin/g++-11 --slave /usr/bin/gcov gcov /usr/bin/gcov-11 update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 20 --slave /usr/bin/g++ g++ /usr/bin/g++-12 --slave /usr/bin/gcov gcov /usr/bin/gcov-12 python3 -m pip install -U --user -r requirements.txt + # Held at 3.25 rather than the 4.x that requirements.txt resolves to: this + # branch's external submodules (json, pybind11, sqlite, yaml-cpp) still declare + # cmake_minimum_required below 3.5, which CMake 4 rejects outright. This must + # run after requirements.txt so the pin is not upgraded away. + python3 -m pip install -U --user cmake==3.25.0 + + # Add pip packages to path + echo "$HOME/.local/bin" >> "$GITHUB_PATH" # Add pip packages to path echo "$HOME/.local/bin" >> "$GITHUB_PATH" diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/context/context.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/context/context.cpp index bf8d0c92299f..6039d32851a0 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/context/context.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/context/context.cpp @@ -295,6 +295,13 @@ start_context(rocprofiler_context_id_t context_id) // conflicting context return ROCPROFILER_STATUS_ERROR_CONTEXT_CONFLICT; } + else if(cfg->dispatch_thread_trace && itr->dispatch_thread_trace && + cfg->dispatch_thread_trace->intersects(*itr->dispatch_thread_trace)) + { + // Two dispatch ATT contexts can run concurrently as long as they target disjoint + // sets of GPU agents. Overlapping agent sets would cross-talk in post_kernel_call. + return ROCPROFILER_STATUS_ERROR_CONTEXT_CONFLICT; + } } uint64_t rocp_tot_contexts = get_registered_contexts_impl()->size(); @@ -376,6 +383,19 @@ stop_context(rocprofiler_context_id_t idx) const context* _expected = itr.load(std::memory_order_acquire); if(_expected && _expected->context_idx == idx.handle) { + // Stop queue-interposed services before clearing the active slot so + // disable_serialization() always runs before dispatches stop being instrumented. + if(_expected->dispatch_counter_collection) + { + rocprofiler::counters::stop_context(const_cast(_expected)); + } + + if(_expected->dispatch_spm) + rocprofiler::spm::stop_context(const_cast(_expected)); + + if(_expected->device_thread_trace) _expected->device_thread_trace->stop_context(); + if(_expected->dispatch_thread_trace) _expected->dispatch_thread_trace->stop_context(); + bool success = itr.compare_exchange_strong(_expected, nullptr); if(success) @@ -383,18 +403,6 @@ stop_context(rocprofiler_context_id_t idx) auto nactive = get_num_active_contexts().load(std::memory_order_acquire); if(nactive > 0) get_num_active_contexts().fetch_sub(1, std::memory_order_release); - if(_expected->dispatch_counter_collection) - { - rocprofiler::counters::stop_context(const_cast(_expected)); - } - - if(_expected->dispatch_spm) - rocprofiler::spm::stop_context(const_cast(_expected)); - - if(_expected->device_thread_trace) _expected->device_thread_trace->stop_context(); - if(_expected->dispatch_thread_trace) - _expected->dispatch_thread_trace->stop_context(); - rocprofiler::hsa::queue_interposition:: notify_queue_interposition_consumer_context_stopped(_expected); diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/CMakeLists.txt b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/CMakeLists.txt index ab27514b1822..c085f2c7f7cd 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/CMakeLists.txt +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/CMakeLists.txt @@ -39,3 +39,4 @@ target_sources(rocprofiler-sdk-object-library PRIVATE ${ROCPROFILER_LIB_HSA_SOUR ${ROCPROFILER_LIB_HSA_HEADERS}) add_subdirectory(details) +add_subdirectory(queue_hooks) diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue.cpp index 36e67f1e642e..ae5bde818153 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue.cpp @@ -45,6 +45,7 @@ #include "lib/rocprofiler-sdk/pc_sampling/hsa_adapter.hpp" #include "lib/rocprofiler-sdk/pc_sampling/service.hpp" #include "lib/rocprofiler-sdk/registration.hpp" +#include "lib/rocprofiler-sdk/thread_trace/queue_hooks.hpp" #include "lib/rocprofiler-sdk/tracing/tracing.hpp" #include @@ -294,6 +295,15 @@ AsyncSignalHandler(hsa_signal_value_t /*signal_v*/, void* data) } }); + // Thread trace completion is migrated off the callback registry (see WriteInterceptor); + // invoke it explicitly here. + thread_trace::signal_completion_hook(queue_info_session.queue, + packet.kernel_packet, + _session, + packet, + packet.instrumentation_packets, + dispatch_time); + CHECK_NOTNULL(hsa::get_queue_controller()) ->serializer(&queue_info_session.queue) .wlock([&](auto& serializer) { @@ -428,8 +438,10 @@ WriteInterceptor(const void* packets, auto* gls = ::rocprofiler::hip::graph::current_launch_state(); const bool graph_launch_active = (gls != nullptr); + // Thread trace no longer registers a queue-controller callback, so it does not count toward + // get_notifiers(); detect it explicitly so an ATT-only run still enters the interceptor. const bool no_real_consumers = - (queue.get_notifiers() == 0 && + (queue.get_notifiers() == 0 && !thread_trace::is_any_active() && context::get_active_contexts(full_packet_instrumentation_context_filter).empty()); const bool has_kernel_replay = kernel_replay::has_active_replay_contexts(); @@ -827,6 +839,18 @@ WriteInterceptor(const void* packets, } }); + // Thread trace is migrated off the per-queue callback registry: call its hook + // explicitly (the other services still flow through signal_callback above). + thread_trace::write_hook(queue, + kernel_packet, + kernel_id, + dispatch_id, + &_packet_data.user_data, + _packet_data.tracing_data.external_correlation_ids, + corr_id, + _packet_data.instrumentation_packets, + _packet_data.is_serialized); + bool inserted_before = false; if(_packet_data.is_serialized) { @@ -1195,6 +1219,9 @@ WriteInterceptor(const void* packets, } }); + // Thread trace requires per-packet mode; it no longer participates in the registry above. + if(thread_trace::is_any_active()) should_batch_packets = false; + if(should_batch_packets) { ROCP_TRACE_IF(pkt_count > 1) << fmt::format( diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_controller.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_controller.cpp index 4b334e34fd1b..3120c9675658 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_controller.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_controller.cpp @@ -677,10 +677,19 @@ common::Synchronized& QueueController::serializer(const Queue* queue) { CHECK(queue); + const auto agent_id = queue->get_agent().get_rocp_agent()->id; + + // Read the refcount before taking the serializer lock: update_serialization() acquires + // them in that order, and reversing it here would deadlock. A concurrent enable/disable + // for this agent can therefore race with the creation below, exactly as the previous + // _serialized_enabled load could, and is resolved the same way -- the next transition + // through update_serialization() reaches the entry once it is in the map. + const bool should_serialize = is_serialization_enabled(agent_id); + common::Synchronized* ret = nullptr; _profiler_serializer.ulock( [&](const auto& m) { - if(auto ptr = m.find(queue->get_agent().get_rocp_agent()->id); ptr != m.end()) + if(auto ptr = m.find(agent_id); ptr != m.end()) { ret = ptr->second.get(); return true; @@ -688,10 +697,10 @@ QueueController::serializer(const Queue* queue) return false; }, [&](auto& m) { - ret = m.emplace(queue->get_agent().get_rocp_agent()->id, + ret = m.emplace(agent_id, std::make_shared>()) .first->second.get(); - if(_serialized_enabled.load() == true) + if(should_serialize) { ret->wlock([&](auto& serializer) { serializer.enable({}); }); } @@ -715,47 +724,120 @@ per_dev_map(const QueueController::queue_map_t& _queues_v) }; // namespace void -QueueController::disable_serialization() +QueueController::update_serialization(const agent_handle_set_t& agents, bool enable) { _queues.rlock([&](const queue_map_t& _queues_v) { - _serialized_enabled.store(false); auto pd_map = per_dev_map(_queues_v); - _profiler_serializer.wlock([&](auto& m) { - for(auto& [k, v] : m) + + // Agents whose effective state changed in this call; only those get their serializer + // toggled. Each entry stores the agent id together with its new effective state so + // that the serializer action is derived from the post-update refcount rather than + // from the caller's `enable` flag (which could be stale if another thread races + // between the refcount unlock and the serializer lock). + auto transitioned = std::vector>{}; + bool any_enabled = false; + + _serialization_refcount.wlock([&](auto& state) { + // An agent can be covered by `all` and by an explicit entry at the same time, so + // the transition has to be decided by comparing effective counts across the + // update rather than by watching a single counter hit zero. + auto was_enabled = std::unordered_map{}; + _profiler_serializer.rlock([&](const auto& serializers) { + for(const auto& [agent_id, _] : serializers) + was_enabled.emplace(agent_id, state.enabled(agent_id)); + }); + + if(agents.empty()) + { + if(enable) + ++state.all; + else if(state.all > 0) + --state.all; + } + else { - if(auto it = pd_map.find(k); it != pd_map.end()) + for(const auto& agent_id : agents) { - v->wlock([&](auto& serializer) { serializer.disable(it->second); }); + auto& count = state.per_agent[agent_id]; + if(enable) + ++count; + else if(count > 0) + --count; } - else + } + + _profiler_serializer.rlock([&](const auto& serializers) { + for(const auto& [agent_id, _] : serializers) { - v->wlock([&](auto& serializer) { serializer.disable({}); }); + auto itr = was_enabled.find(agent_id); + if(itr == was_enabled.end()) continue; + bool now_enabled = state.enabled(agent_id); + if(itr->second != now_enabled) transitioned.emplace_back(agent_id, now_enabled); } + }); + + any_enabled = state.any(); + }); + + _serialized_enabled.store(any_enabled); + + if(transitioned.empty()) return; + + _profiler_serializer.wlock([&](auto& m) { + // Deliberately not a structured binding: the serializer lambda below captures + // now_enabled, and capturing a structured binding is a C++20 extension. + for(const auto& transition : transitioned) + { + const auto agent_id = transition.first; + const bool now_enabled = transition.second; + + auto itr = m.find(agent_id); + if(itr == m.end() || !itr->second) continue; + + auto queues = hsa_barrier::queue_map_ptr_t{}; + if(auto it = pd_map.find(agent_id); it != pd_map.end()) queues = it->second; + + itr->second->wlock([&](auto& serializer) { + if(now_enabled) + serializer.enable(queues); + else + serializer.disable(queues); + }); } }); }); } +void +QueueController::disable_serialization() +{ + update_serialization({}, false); +} + void QueueController::enable_serialization() { - _queues.rlock([&](const queue_map_t& _queues_v) { - _serialized_enabled.store(true); - auto pd_map = per_dev_map(_queues_v); - _profiler_serializer.wlock([&](auto& m) { - for(auto& [k, v] : m) - { - if(auto it = pd_map.find(k); it != pd_map.end()) - { - v->wlock([&](auto& serializer) { serializer.enable(it->second); }); - } - else - { - v->wlock([&](auto& serializer) { serializer.enable({}); }); - } - } - }); - }); + update_serialization({}, true); +} + +void +QueueController::disable_serialization(const agent_handle_set_t& agents) +{ + update_serialization(agents, false); +} + +void +QueueController::enable_serialization(const agent_handle_set_t& agents) +{ + update_serialization(agents, true); +} + +bool +QueueController::is_serialization_enabled(rocprofiler_agent_id_t agent_id) const +{ + bool enabled = false; + _serialization_refcount.rlock([&](const auto& state) { enabled = state.enabled(agent_id); }); + return enabled; } void diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_controller.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_controller.hpp index b42ae02394f8..c1cbf4936480 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_controller.hpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_controller.hpp @@ -30,11 +30,13 @@ #include #include +#include #include #include #include #include +#include #include namespace rocprofiler @@ -100,13 +102,27 @@ class QueueController common::Synchronized& serializer(const Queue*); + // An empty set means every GPU agent. + using agent_handle_set_t = std::unordered_set; + /** - * Disable serialization for QueueController, has no effect if counter collection - * is not in use (which defaults to no serialization mechanism). Should only be used for - * testing. + * Enable/disable serialization for QueueController, has no effect if counter collection + * is not in use (which defaults to no serialization mechanism). + * + * Serialization is reference counted per agent. Counter collection, thread trace and SPM + * each enable it independently, so an agent stays serialized until every subsystem that + * asked for it has released it -- otherwise whichever service stops first would silently + * unserialize the ones still running. The no-argument overloads apply to every GPU agent, + * which is the behavior every caller had before agents could be scoped. */ void enable_serialization(); void disable_serialization(); + void enable_serialization(const agent_handle_set_t& agents); + void disable_serialization(const agent_handle_set_t& agents); + + // Whether the given agent currently has serialization enabled. Exposed for tests and for + // callers that need to reason about scope rather than trigger a transition. + bool is_serialization_enabled(rocprofiler_agent_id_t agent_id) const; // Prints current state of signals for queues, used for debugging. Only prints // serialization related signals if not compiled in debug mode. @@ -132,6 +148,39 @@ class QueueController std::unordered_map>>> _profiler_serializer; + // How many subsystems currently want serialization, per agent. `all` counts the callers + // that asked for every agent; it is kept separate from `per_agent` because an agent's + // serializer is created lazily on first use and an unscoped request has to cover agents + // that do not exist yet. + struct serialization_refcount + { + int64_t all = 0; + std::unordered_map per_agent = {}; + + int64_t count(rocprofiler_agent_id_t agent_id) const + { + auto itr = per_agent.find(agent_id); + return all + ((itr == per_agent.end()) ? 0 : itr->second); + } + + bool enabled(rocprofiler_agent_id_t agent_id) const { return count(agent_id) > 0; } + + bool any() const + { + if(all > 0) return true; + for(const auto& [agent_id, count] : per_agent) + { + if(count > 0) return true; + } + return false; + } + }; + + common::Synchronized _serialization_refcount; + + // Applies a +1/-1 to the refcount of each agent in `agents` (empty == all GPU agents) and + // toggles only the serializers whose effective count crossed zero. + void update_serialization(const agent_handle_set_t& agents, bool enable); }; QueueController* diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_hooks/CMakeLists.txt b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_hooks/CMakeLists.txt new file mode 100644 index 000000000000..5f5e8f8481f2 --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_hooks/CMakeLists.txt @@ -0,0 +1,9 @@ +# +# Shared, header-only client-id tags used by the per-service queue hooks. +# +set(ROCPROFILER_LIB_HSA_QUEUE_HOOKS_SOURCES) +set(ROCPROFILER_LIB_HSA_QUEUE_HOOKS_HEADERS client_ids.hpp) + +target_sources( + rocprofiler-sdk-object-library PRIVATE ${ROCPROFILER_LIB_HSA_QUEUE_HOOKS_SOURCES} + ${ROCPROFILER_LIB_HSA_QUEUE_HOOKS_HEADERS}) diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_hooks/client_ids.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_hooks/client_ids.hpp new file mode 100644 index 000000000000..b7a45802fef9 --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_hooks/client_ids.hpp @@ -0,0 +1,64 @@ +// MIT License +// +// Copyright (c) 2026 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. + +#pragma once + +#include + +namespace rocprofiler +{ +namespace hsa +{ +namespace queue_hooks +{ +// Tags identifying the producer subsystem of each inst_pkt_t entry, stored in the +// inst_pkt_t pair as an hsa::ClientID. +// +// The underlying type is fixed to int64_t to match hsa::ClientID. The enum is +// deliberately unscoped rather than an enum class: unscoped keeps the implicit +// conversion to ClientID, so the emplace sites need no cast, while still grouping +// the values under one type. +// +// The values are negative because they share the inst_pkt_t tag space with the +// per-queue callback registry, whose QueueController::add_callback hands out +// strictly positive ids counting up from 1. Migrated services (counters, thread +// trace, SPM) select their own packets by comparing against these tags, so a tag +// drawn from the registry's range could match a packet the service does not own. +// Services that have not migrated yet still receive registry ids, so the two +// ranges must stay disjoint until the registry is gone. Any tag added here must +// likewise be negative; -1 is left unused because QueueController::add_callback +// uses it as a local "unassigned" initializer. +enum client_id : int64_t +{ + COUNTERS_CLIENT_ID = -1001, + THREAD_TRACE_CLIENT_ID = -1002, + PC_SAMPLING_CLIENT_ID = -1003, + SPM_CLIENT_ID = -1004, +}; + +static_assert(COUNTERS_CLIENT_ID < 0 && THREAD_TRACE_CLIENT_ID < 0 && PC_SAMPLING_CLIENT_ID < 0 && + SPM_CLIENT_ID < 0, + "queue hook tags must stay disjoint from the positive ids handed out by " + "QueueController::add_callback"); +} // namespace queue_hooks +} // namespace hsa +} // namespace rocprofiler diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/CMakeLists.txt b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/CMakeLists.txt index 4ad32790258b..b08e867e48ee 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/CMakeLists.txt +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/CMakeLists.txt @@ -1,7 +1,8 @@ -set(ROCPROFILER_LIB_THREAD_TRACE_SOURCES core.cpp service.cpp code_object.cpp decode.cpp - dl.cpp threading.cpp hsa_util.cpp) +set(ROCPROFILER_LIB_THREAD_TRACE_SOURCES + core.cpp service.cpp code_object.cpp decode.cpp dl.cpp threading.cpp hsa_util.cpp + queue_hooks.cpp) set(ROCPROFILER_LIB_THREAD_TRACE_HEADERS core.hpp code_object.hpp dl.hpp threading.hpp - hsa_util.hpp trace_decoder_api.h) + hsa_util.hpp trace_decoder_api.h queue_hooks.hpp) target_sources( rocprofiler-sdk-object-library PRIVATE ${ROCPROFILER_LIB_THREAD_TRACE_SOURCES} ${ROCPROFILER_LIB_THREAD_TRACE_HEADERS}) diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/core.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/core.cpp index 8b2e3aede54f..05d9dc47a15f 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/core.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/core.cpp @@ -35,6 +35,7 @@ #include "lib/rocprofiler-sdk/context/context.hpp" #include "lib/rocprofiler-sdk/hsa/queue.hpp" #include "lib/rocprofiler-sdk/hsa/queue_controller.hpp" +#include "lib/rocprofiler-sdk/hsa/queue_hooks/client_ids.hpp" #include "lib/rocprofiler-sdk/internal_threading.hpp" #include "lib/rocprofiler-sdk/kernel_replay/local_context.hpp" #include "lib/rocprofiler-sdk/registration.hpp" @@ -80,10 +81,6 @@ struct cbdata_t uint64_t next_chunk = 0; }; -// Keeps track of a single client registering for serialized thread trace -// operations so we can gate new traces while one is active. -common::Synchronized> client; - // True once the HSA runtime is registered. Gates start_context() so pre-init // start requests are deferred and replayed by start_active_contexts(). std::atomic& @@ -462,16 +459,11 @@ DispatchThreadTracer::resource_init() void DispatchThreadTracer::resource_deinit() { - enabled.store(false, std::memory_order_release); - - if(auto* controller = hsa::get_queue_controller()) + if(enabled.exchange(false, std::memory_order_acq_rel)) { - client.wlock([&](auto& client_id) { - if(!client_id) return; - controller->remove_callback(*client_id); - client_id = std::nullopt; - }); - controller->disable_serialization(); + const auto serialization_agents = configured_agents(); + if(auto* controller = hsa::get_queue_controller()) + controller->disable_serialization(serialization_agents); } ROCP_TRACE << "Clearing agents"; @@ -543,69 +535,79 @@ DispatchThreadTracer::post_kernel_call(DispatchThreadTracer::inst_pkt_t& aql, for(auto& aql_pkt : aql) { + if(aql_pkt.second != hsa::queue_hooks::THREAD_TRACE_CLIENT_ID) continue; + auto* pkt = dynamic_cast(aql_pkt.first.get()); if(!pkt) continue; std::shared_lock lk(agents_map_mut); + + auto it = agents.find(pkt->GetAgent()); + if(it == agents.end() || it->second == nullptr) continue; + post_move_data.fetch_sub(1); if(pkt->after_krn_pkt.empty()) continue; - auto it = agents.find(pkt->GetAgent()); - if(it != agents.end() && it->second != nullptr) - it->second->iterate_data(pkt->GetHandle(), packet_data.user_data); + it->second->iterate_data(pkt->GetHandle(), packet_data.user_data); } } -void -DispatchThreadTracer::start_context() +std::unordered_set +DispatchThreadTracer::configured_agents() const { - using corr_id_map_t = hsa::queue_info_session_t::external_corr_id_map_t; + auto result = std::unordered_set{}; + std::shared_lock lk(agents_map_mut); + for(const auto& [agent_id, _] : params) + result.insert(agent_id); + return result; +} - // Only installs queue-controller callbacks (cached and applied to queues as - // they are created), so this is safe to call before hsa_init. - CHECK_NOTNULL(hsa::get_queue_controller())->enable_serialization(); - enabled.store(true, std::memory_order_release); +bool +DispatchThreadTracer::collects_on(rocprofiler_agent_id_t agent_id) const +{ + std::shared_lock lk(agents_map_mut); + return params.count(agent_id) > 0; +} + +bool +DispatchThreadTracer::intersects(const DispatchThreadTracer& rhs) const +{ + std::shared_lock lk(agents_map_mut); + std::shared_lock lk_rhs(rhs.agents_map_mut); + for(const auto& [agent_id, _] : params) + { + if(rhs.params.count(agent_id) > 0) return true; + } + return false; +} - // Only one thread should be attempting to enable/disable this context - client.wlock([&](auto& client_id) { - if(client_id) return; - - auto&& _callbacks = hsa::queue_callbacks_t{ - .batch_packets = []() { return false; }, - .write_interceptor = - [=](const hsa::Queue& q, - const hsa::rocprofiler_packet& /* kern_pkt */, - rocprofiler_kernel_id_t kernel_id, - rocprofiler_dispatch_id_t dispatch_id, - rocprofiler_user_data_t* user_data, - const corr_id_map_t& /* extern_corr_ids */, - const context::correlation_id* corr_id) { - return this->pre_kernel_call(q, kernel_id, dispatch_id, user_data, corr_id); - }, - .signal_completion = - [=](const hsa::Queue& /* q */, - hsa::rocprofiler_packet /* kern_pkt */, - std::shared_ptr& session, - hsa::packet_data_t& packet_data, - inst_pkt_t& aql, - kernel_dispatch::profiling_time) { - this->post_kernel_call(aql, *session, packet_data); - }}; - - client_id = CHECK_NOTNULL(hsa::get_queue_controller()) - ->add_callback(std::nullopt, std::move(_callbacks)); - }); +void +DispatchThreadTracer::start_context() const +{ + // Thread trace no longer registers a per-queue callback with the queue controller; the + // HSA write interceptor now calls thread_trace::write_hook / signal_completion_hook + // directly (see hsa/queue.cpp). Scope serialization to the agents configured on this + // context. An empty set still means every agent. + const auto serialization_agents = configured_agents(); + CHECK_NOTNULL(hsa::get_queue_controller())->enable_serialization(serialization_agents); + enabled.store(true, std::memory_order_release); } void -DispatchThreadTracer::stop_context() // NOLINT(readability-convert-member-functions-to-static) +DispatchThreadTracer::stop_context() const { - // Stop injecting ATT packets before transitioning serialization. Completion callbacks remain - // registered so packets already in the queues can drain through the serializer transition. + // Stop injecting ATT packets before transitioning serialization. Completion hooks remain + // reachable for packets already in the queues while the serializer transition drains. if(!enabled.exchange(false, std::memory_order_acq_rel)) return; - if(auto* controller = hsa::get_queue_controller()) controller->disable_serialization(); + // Drain in-flight dispatches before tearing down serialization so post_kernel_call can + // still run while the context remains registered. + hsa::queue_controller_sync(); + + const auto serialization_agents = configured_agents(); + if(auto* controller = hsa::get_queue_controller()) + controller->disable_serialization(serialization_agents); } DeviceThreadTracer::DeviceThreadTracer() diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/core.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/core.hpp index 70222c280ca8..adc788cf0213 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/core.hpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/core.hpp @@ -45,6 +45,7 @@ #include #include #include +#include #include namespace rocprofiler @@ -151,11 +152,14 @@ class DispatchThreadTracer ~DispatchThreadTracer() = default; /// Initializes shared resources needed by dispatch-based tracing. - void start_context(); - void stop_context(); + void start_context() const; + void stop_context() const; void resource_init(); void resource_deinit(); + bool collects_on(rocprofiler_agent_id_t agent_id) const; + bool intersects(const DispatchThreadTracer& rhs) const; + void add_agent(rocprofiler_agent_id_t agent, thread_trace_parameter_pack pack) { auto lk = std::unique_lock{agents_map_mut}; @@ -173,13 +177,18 @@ class DispatchThreadTracer const hsa::packet_data_t& packet_data); const auto& get_agents() const { return agents; } + /// Number of injected trace packets still awaiting post_kernel_call. + int pending_post_moves() const { return post_move_data.load(); } + + std::unordered_set configured_agents() const; + private: std::unordered_map> agents{}; std::unordered_map params{}; - std::shared_mutex agents_map_mut{}; - std::atomic post_move_data{0}; - std::atomic enabled{false}; + mutable std::shared_mutex agents_map_mut{}; + std::atomic post_move_data{0}; + mutable std::atomic enabled{false}; }; class DeviceThreadTracer diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/queue_hooks.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/queue_hooks.cpp new file mode 100644 index 000000000000..4a5d932e531d --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/queue_hooks.cpp @@ -0,0 +1,97 @@ +// MIT License +// +// Copyright (c) 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. + +#include "lib/rocprofiler-sdk/thread_trace/queue_hooks.hpp" +#include "lib/rocprofiler-sdk/context/context.hpp" +#include "lib/rocprofiler-sdk/hsa/agent_cache.hpp" +#include "lib/rocprofiler-sdk/hsa/queue.hpp" +#include "lib/rocprofiler-sdk/hsa/queue_hooks/client_ids.hpp" +#include "lib/rocprofiler-sdk/thread_trace/core.hpp" + +namespace rocprofiler +{ +namespace thread_trace +{ +namespace +{ +auto +thread_trace_contexts_filter() +{ + return [](const context::context* ctx) -> bool { + return ctx && ctx->dispatch_thread_trace != nullptr; + }; +} +} // namespace + +void +write_hook(const hsa::Queue& queue, + const hsa::rocprofiler_packet& /*kernel_packet*/, + rocprofiler_kernel_id_t kernel_id, + rocprofiler_dispatch_id_t dispatch_id, + rocprofiler_user_data_t* user_data, + const hsa::queue_info_session_t::external_corr_id_map_t& /*ext_corr_ids*/, + const context::correlation_id* correlation_id, + hsa::inst_pkt_t& inst_pkt, + bool& is_serialized) +{ + const auto agent_id = CHECK_NOTNULL(queue.get_agent().get_rocp_agent())->id; + + auto active = context::get_active_contexts(thread_trace_contexts_filter()); + for(const auto* ctx : active) + { + auto& tracer = *ctx->dispatch_thread_trace; + if(!tracer.collects_on(agent_id)) continue; + + auto [packet, bSerial] = + tracer.pre_kernel_call(queue, kernel_id, dispatch_id, user_data, correlation_id); + if(packet) + inst_pkt.emplace_back(std::move(packet), hsa::queue_hooks::THREAD_TRACE_CLIENT_ID); + is_serialized |= bSerial; + } +} + +void +signal_completion_hook(const hsa::Queue& /*queue*/, + const hsa::rocprofiler_packet& /*kernel_packet*/, + std::shared_ptr& session, + hsa::packet_data_t& packet_data, + hsa::inst_pkt_t& inst_pkt, + kernel_dispatch::profiling_time /*dispatch_time*/) +{ + // Route by packet provenance, not current activeness: post_kernel_call self-filters via + // THREAD_TRACE_CLIENT_ID and agent ownership, so in-flight dispatches still complete after + // stop_context removes the context from the active list. + auto contexts = context::get_registered_contexts(thread_trace_contexts_filter()); + for(const auto* ctx : contexts) + { + auto& tracer = *ctx->dispatch_thread_trace; + tracer.post_kernel_call(inst_pkt, *session, packet_data); + } +} + +bool +is_any_active() +{ + return !context::get_active_contexts(thread_trace_contexts_filter()).empty(); +} +} // namespace thread_trace +} // namespace rocprofiler diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/queue_hooks.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/queue_hooks.hpp new file mode 100644 index 000000000000..f1f2746b2a1a --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/queue_hooks.hpp @@ -0,0 +1,66 @@ +// MIT License +// +// Copyright (c) 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. + +#pragma once + +#include "lib/rocprofiler-sdk/context/context.hpp" +#include "lib/rocprofiler-sdk/hsa/queue.hpp" +#include "lib/rocprofiler-sdk/hsa/queue_info_session.hpp" +#include "lib/rocprofiler-sdk/hsa/rocprofiler_packet.hpp" +#include "lib/rocprofiler-sdk/kernel_dispatch/profiling_time.hpp" + +#include + +namespace rocprofiler +{ +namespace thread_trace +{ +// Explicit replacement for the thread-trace per-queue callback that used to be +// registered with the HSA queue controller. Iterates active dispatch_thread_trace +// contexts and calls each tracer's pre_kernel_call; appends produced packets to +// inst_pkt, OR-folding each tracer's serialize flag into is_serialized. +void +write_hook(const hsa::Queue& queue, + const hsa::rocprofiler_packet& kernel_packet, + rocprofiler_kernel_id_t kernel_id, + rocprofiler_dispatch_id_t dispatch_id, + rocprofiler_user_data_t* user_data, + const hsa::queue_info_session_t::external_corr_id_map_t& ext_corr_ids, + const context::correlation_id* correlation_id, + hsa::inst_pkt_t& inst_pkt, + bool& is_serialized); + +// Explicit replacement for the thread-trace completion callback. Iterates active +// dispatch_thread_trace contexts and calls each tracer's post_kernel_call. +void +signal_completion_hook(const hsa::Queue& queue, + const hsa::rocprofiler_packet& kernel_packet, + std::shared_ptr& session, + hsa::packet_data_t& packet, + hsa::inst_pkt_t& inst_pkt, + kernel_dispatch::profiling_time dispatch_time); + +// True if any context currently has dispatch thread trace active. +bool +is_any_active(); +} // namespace thread_trace +} // namespace rocprofiler diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/tests/CMakeLists.txt b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/tests/CMakeLists.txt index 0f4bc0b846fc..1edf3e34fbef 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/tests/CMakeLists.txt +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/tests/CMakeLists.txt @@ -71,3 +71,59 @@ rocprofiler_add_unit_test( "ROCPROFILER_METRICS_PATH=${_ROCPROFILER_SHARE_DIR}" "LD_LIBRARY_PATH=${CMAKE_INSTALL_PREFIX}/lib:${CMAKE_INSTALL_PREFIX}/llvm/lib:/opt/rocm/lib:/opt/rocm/llvm/lib:$ENV{LD_LIBRARY_PATH}" FAIL_REGULAR_EXPRESSION "${ROCPROFILER_DEFAULT_FAIL_REGEX}") + +set(ROCPROFILER_THREAD_TRACE_QUEUE_HOOKS_TEST_SOURCES "queue_hooks_test.cpp") + +add_executable(thread-trace-queue-hooks-test) +target_sources(thread-trace-queue-hooks-test + PRIVATE ${ROCPROFILER_THREAD_TRACE_QUEUE_HOOKS_TEST_SOURCES}) + +target_link_libraries( + thread-trace-queue-hooks-test + PRIVATE rocprofiler-sdk::rocprofiler-sdk-static-library + rocprofiler-sdk::rocprofiler-sdk-abseil + rocprofiler-sdk::rocprofiler-sdk-hsa-runtime + rocprofiler-sdk::rocprofiler-sdk-hip + rocprofiler-sdk::rocprofiler-sdk-common-library + GTest::gtest + GTest::gtest_main + rocprofiler-sdk::counter-test-constants) + +rocprofiler_add_unit_test( + TARGET thread-trace-queue-hooks-test + SOURCES ${ROCPROFILER_THREAD_TRACE_QUEUE_HOOKS_TEST_SOURCES} + TEST_LIST thread-trace-test_TESTS + TIMEOUT 10 + LABELS "unittests" + ENVIRONMENT + "ROCPROFILER_METRICS_PATH=${_ROCPROFILER_SHARE_DIR}" + "LD_LIBRARY_PATH=${PROJECT_BINARY_DIR}/${CMAKE_INSTALL_LIBDIR}:${CMAKE_INSTALL_PREFIX}/lib:${CMAKE_INSTALL_PREFIX}/llvm/lib:${ROCM_PATH}/lib:${ROCM_PATH}/llvm/lib:$ENV{LD_LIBRARY_PATH}" + FAIL_REGULAR_EXPRESSION "${ROCPROFILER_DEFAULT_FAIL_REGEX}") + +set(ROCPROFILER_THREAD_TRACE_PER_AGENT_TEST_SOURCES "per_agent_scoping_test.cpp") + +add_executable(thread-trace-per-agent-test) +target_sources(thread-trace-per-agent-test + PRIVATE ${ROCPROFILER_THREAD_TRACE_PER_AGENT_TEST_SOURCES}) + +target_link_libraries( + thread-trace-per-agent-test + PRIVATE rocprofiler-sdk::rocprofiler-sdk-static-library + rocprofiler-sdk::rocprofiler-sdk-abseil + rocprofiler-sdk::rocprofiler-sdk-hsa-runtime + rocprofiler-sdk::rocprofiler-sdk-hip + rocprofiler-sdk::rocprofiler-sdk-common-library + GTest::gtest + GTest::gtest_main + rocprofiler-sdk::counter-test-constants) + +rocprofiler_add_unit_test( + TARGET thread-trace-per-agent-test + SOURCES ${ROCPROFILER_THREAD_TRACE_PER_AGENT_TEST_SOURCES} + TEST_LIST thread-trace-test_TESTS + TIMEOUT 45 + LABELS "unittests" + ENVIRONMENT + "ROCPROFILER_METRICS_PATH=${_ROCPROFILER_SHARE_DIR}" + "LD_LIBRARY_PATH=${PROJECT_BINARY_DIR}/${CMAKE_INSTALL_LIBDIR}:${CMAKE_LIBRARY_OUTPUT_DIRECTORY}:${CMAKE_INSTALL_PREFIX}/lib:${CMAKE_INSTALL_PREFIX}/llvm/lib:${ROCM_PATH}/lib:${ROCM_PATH}/llvm/lib:$ENV{LD_LIBRARY_PATH}" + FAIL_REGULAR_EXPRESSION "${ROCPROFILER_DEFAULT_FAIL_REGEX}") diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/tests/per_agent_scoping_test.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/tests/per_agent_scoping_test.cpp new file mode 100644 index 000000000000..0f90ba2a849d --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/tests/per_agent_scoping_test.cpp @@ -0,0 +1,257 @@ +// MIT License +// +// Copyright (c) 2026 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 "lib/rocprofiler-sdk/agent.hpp" +#include "lib/rocprofiler-sdk/context/context.hpp" +#include "lib/rocprofiler-sdk/counters/tests/hsa_tables.hpp" +#include "lib/rocprofiler-sdk/hsa/agent_cache.hpp" +#include "lib/rocprofiler-sdk/hsa/queue.hpp" +#include "lib/rocprofiler-sdk/hsa/queue_controller.hpp" +#include "lib/rocprofiler-sdk/registration.hpp" +#include "lib/rocprofiler-sdk/thread_trace/core.hpp" +#include "lib/rocprofiler-sdk/thread_trace/queue_hooks.hpp" + +#include +#include + +#include +#include + +using namespace rocprofiler::counters::test_constants; +using namespace rocprofiler; + +namespace rocprofiler +{ +namespace hsa +{ +class PerAgentFakeQueue : public Queue +{ +public: + PerAgentFakeQueue(const AgentCache& a, rocprofiler_queue_id_t id) + : Queue(a, get_api_table()) + , _agent(a) + , _id(id) + {} + const AgentCache& get_agent() const final { return _agent; }; + rocprofiler_queue_id_t get_id() const final { return _id; }; + + ~PerAgentFakeQueue() override = default; + +private: + const AgentCache& _agent; + rocprofiler_queue_id_t _id = {}; +}; +} // namespace hsa +} // namespace rocprofiler + +namespace +{ +struct agent_pair +{ + const hsa::AgentCache* first = nullptr; + const hsa::AgentCache* second = nullptr; +}; + +void +test_init() +{ + static bool _initialized = false; + if(_initialized) return; + _initialized = true; + + HsaApiTable table; + table.amd_ext_ = &get_ext_table(); + table.core_ = &get_api_table(); + agent::construct_agent_cache(&table); + ASSERT_TRUE(hsa::get_queue_controller() != nullptr); + hsa::get_queue_controller()->init(get_api_table(), get_ext_table()); + registration::init_logging(); + registration::set_init_status(-1); + context::push_client(1); +} + +agent_pair +get_two_agents() +{ + auto result = agent_pair{}; + auto& agents = hsa::get_queue_controller()->get_supported_agents(); + for(const auto& [_, agent] : agents) + { + if(!agent.get_rocp_agent()) continue; + if(!result.first) + result.first = &agent; + else if(!result.second) + result.second = &agent; + else + break; + } + return result; +} + +rocprofiler_thread_trace_control_flags_t +noop_dispatch_cb(rocprofiler_agent_id_t, + rocprofiler_queue_id_t, + rocprofiler_async_correlation_id_t, + rocprofiler_kernel_id_t, + rocprofiler_dispatch_id_t, + void*, + rocprofiler_user_data_t*) +{ + return ROCPROFILER_THREAD_TRACE_CONTROL_NONE; +} + +void noop_shader_cb(rocprofiler_thread_trace_shader_data_t, rocprofiler_user_data_t) {} + +rocprofiler_context_id_t +make_att_context(rocprofiler_agent_id_t agent_id) +{ + auto ctx = rocprofiler_context_id_t{0}; + EXPECT_EQ(rocprofiler_create_context(&ctx), ROCPROFILER_STATUS_SUCCESS); + EXPECT_EQ(rocprofiler_configure_dispatch_thread_trace_service( + ctx, agent_id, nullptr, 0, noop_dispatch_cb, noop_shader_cb, nullptr), + ROCPROFILER_STATUS_SUCCESS); + return ctx; +} + +TEST(thread_trace_per_agent, collects_on_and_intersects_match_configured_agents) +{ + ASSERT_EQ(hsa_init(), HSA_STATUS_SUCCESS); + test_init(); + + auto agents = get_two_agents(); + if(!agents.second) GTEST_SKIP() << "fewer than two GPU agents available"; + + auto agent_a = agents.first->get_rocp_agent()->id; + auto agent_b = agents.second->get_rocp_agent()->id; + + auto left = thread_trace::DispatchThreadTracer{}; + auto right = thread_trace::DispatchThreadTracer{}; + left.add_agent(agent_a, {}); + right.add_agent(agent_b, {}); + + EXPECT_TRUE(left.collects_on(agent_a)); + EXPECT_FALSE(left.collects_on(agent_b)); + EXPECT_FALSE(left.intersects(right)); + + right.add_agent(agent_a, {}); + EXPECT_TRUE(left.intersects(right)); + + auto configured = left.configured_agents(); + EXPECT_EQ(configured.size(), 1U); + EXPECT_EQ(configured.count(agent_a), 1U); +} + +TEST(thread_trace_per_agent, disjoint_agent_contexts_can_be_active_together) +{ + ASSERT_EQ(hsa_init(), HSA_STATUS_SUCCESS); + test_init(); + + auto agents = get_two_agents(); + if(!agents.second) GTEST_SKIP() << "fewer than two GPU agents available"; + + auto ctx_a = make_att_context(agents.first->get_rocp_agent()->id); + auto ctx_b = make_att_context(agents.second->get_rocp_agent()->id); + + EXPECT_EQ(rocprofiler_start_context(ctx_a), ROCPROFILER_STATUS_SUCCESS); + EXPECT_EQ(rocprofiler_start_context(ctx_b), ROCPROFILER_STATUS_SUCCESS); + + EXPECT_EQ(rocprofiler_stop_context(ctx_b), ROCPROFILER_STATUS_SUCCESS); + EXPECT_EQ(rocprofiler_stop_context(ctx_a), ROCPROFILER_STATUS_SUCCESS); +} + +TEST(thread_trace_per_agent, overlapping_agent_contexts_still_conflict) +{ + ASSERT_EQ(hsa_init(), HSA_STATUS_SUCCESS); + test_init(); + + auto agents = get_two_agents(); + if(!agents.first) GTEST_SKIP() << "no GPU agent available"; + + auto agent_a = agents.first->get_rocp_agent()->id; + auto ctx_a = make_att_context(agent_a); + auto ctx_b = make_att_context(agent_a); + + ASSERT_EQ(rocprofiler_start_context(ctx_a), ROCPROFILER_STATUS_SUCCESS); + EXPECT_EQ(rocprofiler_start_context(ctx_b), ROCPROFILER_STATUS_ERROR_CONTEXT_CONFLICT); + ASSERT_EQ(rocprofiler_stop_context(ctx_a), ROCPROFILER_STATUS_SUCCESS); +} + +TEST(thread_trace_per_agent, serialization_is_scoped_to_the_contexts_agents) +{ + ASSERT_EQ(hsa_init(), HSA_STATUS_SUCCESS); + test_init(); + + auto agents = get_two_agents(); + if(!agents.second) GTEST_SKIP() << "fewer than two GPU agents available"; + + auto agent_a = agents.first->get_rocp_agent()->id; + auto agent_b = agents.second->get_rocp_agent()->id; + auto* controller = hsa::get_queue_controller(); + + auto queue_a = hsa::PerAgentFakeQueue{*agents.first, {.handle = 501}}; + auto queue_b = hsa::PerAgentFakeQueue{*agents.second, {.handle = 502}}; + controller->serializer(&queue_a); + controller->serializer(&queue_b); + + auto ctx = make_att_context(agent_b); + ASSERT_EQ(rocprofiler_start_context(ctx), ROCPROFILER_STATUS_SUCCESS); + + EXPECT_TRUE(controller->is_serialization_enabled(agent_b)); + EXPECT_FALSE(controller->is_serialization_enabled(agent_a)); + + ASSERT_EQ(rocprofiler_stop_context(ctx), ROCPROFILER_STATUS_SUCCESS); + EXPECT_FALSE(controller->is_serialization_enabled(agent_b)); +} + +TEST(thread_trace_per_agent, write_hook_ignores_dispatches_on_other_agents) +{ + ASSERT_EQ(hsa_init(), HSA_STATUS_SUCCESS); + test_init(); + + auto agents = get_two_agents(); + if(!agents.second) GTEST_SKIP() << "fewer than two GPU agents available"; + + auto ctx = make_att_context(agents.second->get_rocp_agent()->id); + ASSERT_EQ(rocprofiler_start_context(ctx), ROCPROFILER_STATUS_SUCCESS); + + auto packet = hsa::rocprofiler_packet{}; + auto corr_id = context::correlation_id{}; + corr_id.internal = 4242; + + { + auto queue_a = hsa::PerAgentFakeQueue{*agents.first, {.handle = 601}}; + auto inst_pkt = hsa::inst_pkt_t{}; + bool is_serialized = false; + auto user_data = rocprofiler_user_data_t{.value = corr_id.internal}; + + thread_trace::write_hook( + queue_a, packet, 42, 1, &user_data, {}, &corr_id, inst_pkt, is_serialized); + + EXPECT_FALSE(is_serialized) + << "a dispatch on an agent outside the context must not be serialized"; + EXPECT_TRUE(inst_pkt.empty()) + << "a dispatch on an agent outside the context must not be instrumented"; + } + + ASSERT_EQ(rocprofiler_stop_context(ctx), ROCPROFILER_STATUS_SUCCESS); +} +} // namespace diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/tests/queue_hooks_test.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/tests/queue_hooks_test.cpp new file mode 100644 index 000000000000..f80aca97a50f --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/tests/queue_hooks_test.cpp @@ -0,0 +1,178 @@ +// MIT License +// +// Copyright (c) 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. + +#include "lib/rocprofiler-sdk/thread_trace/queue_hooks.hpp" +#include "lib/rocprofiler-sdk/context/context.hpp" +#include "lib/rocprofiler-sdk/counters/tests/hsa_tables.hpp" +#include "lib/rocprofiler-sdk/hsa/agent_cache.hpp" +#include "lib/rocprofiler-sdk/hsa/hsa.hpp" +#include "lib/rocprofiler-sdk/hsa/queue.hpp" +#include "lib/rocprofiler-sdk/hsa/queue_controller.hpp" +#include "lib/rocprofiler-sdk/hsa/queue_hooks/client_ids.hpp" +#include "lib/rocprofiler-sdk/registration.hpp" +#include "lib/rocprofiler-sdk/thread_trace/core.hpp" + +#include +#include + +#include +#include + +using namespace rocprofiler::counters::test_constants; +using namespace rocprofiler; + +namespace rocprofiler +{ +namespace hsa +{ +class HookTestFakeQueue : public Queue +{ +public: + HookTestFakeQueue(const AgentCache& a, rocprofiler_queue_id_t id) + : Queue(a, get_api_table()) + , _agent(a) + , _id(id) + {} + + const AgentCache& get_agent() const final { return _agent; }; + rocprofiler_queue_id_t get_id() const final { return _id; }; + + ~HookTestFakeQueue() override = default; + +private: + const AgentCache& _agent; + rocprofiler_queue_id_t _id = {}; +}; +} // namespace hsa +} // namespace rocprofiler + +namespace +{ +void +test_init() +{ + static bool _initialized = false; + if(_initialized) return; + _initialized = true; + + HsaApiTable table; + table.amd_ext_ = &get_ext_table(); + table.core_ = &get_api_table(); + // ThreadTracerAgent reads these globals when resource_init() builds its agent map, so they + // have to be populated the same way att_packet_test does. + hsa::copy_table(table.core_, 0); + hsa::copy_table(table.amd_ext_, 0); + agent::construct_agent_cache(&table); + ASSERT_TRUE(hsa::get_queue_controller() != nullptr); + hsa::get_queue_controller()->init(get_api_table(), get_ext_table()); + registration::init_logging(); + registration::set_init_status(-1); + context::push_client(1); +} + +rocprofiler_thread_trace_control_flags_t +start_stop_dispatch_cb(rocprofiler_agent_id_t, + rocprofiler_queue_id_t, + rocprofiler_async_correlation_id_t, + rocprofiler_kernel_id_t, + rocprofiler_dispatch_id_t, + void*, + rocprofiler_user_data_t*) +{ + return ROCPROFILER_THREAD_TRACE_CONTROL_START_AND_STOP; +} + +TEST(ThreadTraceQueueHooks, IsAnyActiveReturnsFalseWhenNoContextActive) +{ + EXPECT_FALSE(rocprofiler::thread_trace::is_any_active()); +} + +// A dispatch instrumented while the context is active must still complete via +// signal_completion_hook after stop_context clears the active slot. +TEST(ThreadTraceQueueHooks, StopContextInFlightCompletionRoutesViaHookPath) +{ + ASSERT_EQ(hsa_init(), HSA_STATUS_SUCCESS); + test_init(); + + auto& agents = hsa::get_queue_controller()->get_supported_agents(); + const hsa::AgentCache* att_agent = nullptr; + for(const auto& [_, agent] : agents) + { + if(!agent.get_rocp_agent()) continue; + if(!agent.get_rocp_agent()->runtime_visibility.hsa) continue; + att_agent = &agent; + break; + } + if(!att_agent) GTEST_SKIP() << "no ATT-capable GPU agent available"; + + auto ctx = rocprofiler_context_id_t{0}; + ASSERT_EQ(rocprofiler_create_context(&ctx), ROCPROFILER_STATUS_SUCCESS); + ASSERT_EQ(rocprofiler_configure_dispatch_thread_trace_service( + ctx, + att_agent->get_rocp_agent()->id, + nullptr, + 0, + start_stop_dispatch_cb, + [](rocprofiler_thread_trace_shader_data_t, rocprofiler_user_data_t) {}, + nullptr), + ROCPROFILER_STATUS_SUCCESS); + ASSERT_EQ(rocprofiler_start_context(ctx), ROCPROFILER_STATUS_SUCCESS); + + auto* ctx_p = context::get_mutable_registered_context(ctx); + ASSERT_TRUE(ctx_p && ctx_p->dispatch_thread_trace); + auto& tracer = *ctx_p->dispatch_thread_trace; + + // thread_trace::initialize() only calls resource_init() on contexts that already exist when + // the HSA runtime registers. This context is created afterwards, so its agent map is still + // empty and pre_kernel_call would bail out before building a packet. + tracer.resource_init(); + + hsa::HookTestFakeQueue fq(*att_agent, {.handle = 901}); + hsa::rocprofiler_packet pkt{}; + context::correlation_id corr_id{}; + corr_id.internal = 42; + auto user_data = rocprofiler_user_data_t{.value = corr_id.internal}; + + hsa::inst_pkt_t inst_pkt; + bool is_serialized = false; + thread_trace::write_hook(fq, pkt, 1, 1, &user_data, {}, &corr_id, inst_pkt, is_serialized); + + ASSERT_FALSE(inst_pkt.empty()) << "write_hook must inject ATT control packet"; + EXPECT_GE(tracer.pending_post_moves(), 1); + + ASSERT_EQ(rocprofiler_stop_context(ctx), ROCPROFILER_STATUS_SUCCESS); + EXPECT_FALSE(thread_trace::is_any_active()); + + auto sess = std::make_shared(hsa::queue_info_session_t{.queue = fq}); + auto packet_data = hsa::packet_data_t{}; + packet_data.user_data = user_data; + + thread_trace::signal_completion_hook(fq, pkt, sess, packet_data, inst_pkt, {}); + + EXPECT_EQ(tracer.pending_post_moves(), 0) + << "post_move_data must drain via signal_completion_hook after stop_context"; + + registration::set_init_status(1); + registration::finalize(); + context::pop_client(1); +} +} // namespace