[rocprofiler-system]: Enable UCX Communication API tracing (#2306)
## Motivation Enable UCX communication tracing and communication metadata ## Technical Details Implement UCX API wrappers to trace transport-layer communication. This adds communication data tracking and exposes “UCX Comm Send/Recv” timelines, enabling detailed analysis of MPI, OpenSHMEM, and other UCX-based runtime communication patterns. - Implements function interception for UCX functions across multiple categories using gotcha component. - Extended comm_data component to track UCX send/recv operations - Added ucx_send and ucx_recv labels for Perfetto counter tracks. Integrated UCX data tracking with existing MPI/RCCL tracking infrastructure. - Added ROCPROFSYS_USE_UCX configuration option (enabled by default). - Created FindUCX.cmake module for UCX header detection. Falls back to internal UCX headers if system headers not found. - Updated all Dockerfiles to include UCX dependencies.
This commit is contained in:
Sajina PK
GitHub
+306
-1
@@ -85,6 +85,7 @@ metadata_initialize_comm_data_categories()
|
||||
trace_cache::get_metadata_registry().add_string(
|
||||
trait::name<category::comm_data>::value);
|
||||
trace_cache::get_metadata_registry().add_string(trait::name<category::mpi>::value);
|
||||
trace_cache::get_metadata_registry().add_string(trait::name<category::ucx>::value);
|
||||
|
||||
_is_initialized = true;
|
||||
}
|
||||
@@ -128,6 +129,16 @@ metadata_initialize_comm_data_pmc()
|
||||
trait::name<category::mpi>::description, LONG_DESCRIPTION, COMPONENT, MSG,
|
||||
rocprofsys::trace_cache::ABSOLUTE, BLOCK, EXPRESSION, 0, 0 });
|
||||
#endif
|
||||
trace_cache::get_metadata_registry().add_pmc_info(
|
||||
{ agent_type::CPU, DEVICE_ID, TARGET_ARCH, EVENT_CODE, INSTANCE_ID,
|
||||
comm_data::ucx_send::label, "Tracks UCX communication data sizes",
|
||||
trait::name<category::ucx>::description, LONG_DESCRIPTION, COMPONENT, MSG,
|
||||
rocprofsys::trace_cache::ABSOLUTE, BLOCK, EXPRESSION, 0, 0 });
|
||||
trace_cache::get_metadata_registry().add_pmc_info(
|
||||
{ agent_type::CPU, DEVICE_ID, TARGET_ARCH, EVENT_CODE, INSTANCE_ID,
|
||||
comm_data::ucx_recv::label, "Tracks UCX communication data sizes",
|
||||
trait::name<category::ucx>::description, LONG_DESCRIPTION, COMPONENT, MSG,
|
||||
rocprofsys::trace_cache::ABSOLUTE, BLOCK, EXPRESSION, 0, 0 });
|
||||
}
|
||||
|
||||
template <typename Track>
|
||||
@@ -172,6 +183,8 @@ comm_data::start()
|
||||
metadata_initialize_track<mpi_send>();
|
||||
metadata_initialize_track<mpi_recv>();
|
||||
#endif
|
||||
metadata_initialize_track<ucx_send>();
|
||||
metadata_initialize_track<ucx_recv>();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -195,7 +208,7 @@ comm_data::configure()
|
||||
_once = true;
|
||||
|
||||
comm_data_tracker_t::label() = "comm_data";
|
||||
comm_data_tracker_t::description() = "Tracks MPI/RCCL communication data sizes";
|
||||
comm_data_tracker_t::description() = "Tracks MPI/RCCL/UCX communication data sizes";
|
||||
comm_data_tracker_t::display_unit() = "MB";
|
||||
comm_data_tracker_t::unit() = units::megabyte;
|
||||
|
||||
@@ -471,6 +484,298 @@ comm_data::audit(const gotcha_data& _data, audit::incoming, const void*, int sen
|
||||
}
|
||||
#endif
|
||||
|
||||
// UCX communication tracking implementations
|
||||
|
||||
// ucp_tag_send_nbx: (void* ep, const void* buffer, size_t count, uint64_t tag, const
|
||||
// void* param)
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, const void*,
|
||||
size_t count, uint64_t tag, const void*)
|
||||
{
|
||||
if(count == 0) return;
|
||||
|
||||
if(get_use_perfetto()) write_perfetto_counter_track<ucx_send>(count);
|
||||
|
||||
{
|
||||
cache_comm_data_events<ucx_send>(0, count);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, count);
|
||||
add(JOIN('/', _name, JOIN('=', "tag", tag)), count);
|
||||
}
|
||||
}
|
||||
|
||||
// ucp_tag_recv_nbx: (void* worker, void* buffer, size_t count, uint64_t tag, uint64_t
|
||||
// tag_mask, const void* param)
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, void*, size_t count,
|
||||
uint64_t tag, uint64_t tag_mask, const void*)
|
||||
{
|
||||
if(count == 0) return;
|
||||
|
||||
if(get_use_perfetto()) write_perfetto_counter_track<ucx_recv>(count);
|
||||
|
||||
{
|
||||
cache_comm_data_events<ucx_recv>(0, count);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, count);
|
||||
add(JOIN('/', _name, JOIN('=', "tag", tag)), count);
|
||||
add(JOIN('/', _name, JOIN('=', "tag", tag), JOIN('=', "tag_mask", tag_mask)),
|
||||
count);
|
||||
}
|
||||
}
|
||||
|
||||
// ucp_put_nbx: (void* ep, const void* buffer, size_t count, uint64_t remote_addr, void*
|
||||
// rkey, const void* param)
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, const void*,
|
||||
size_t count, uint64_t remote_addr, void*, const void*)
|
||||
{
|
||||
if(count == 0) return;
|
||||
|
||||
if(get_use_perfetto()) write_perfetto_counter_track<ucx_send>(count);
|
||||
|
||||
{
|
||||
cache_comm_data_events<ucx_send>(0, count);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, count);
|
||||
add(JOIN('/', _name, JOIN('=', "remote_addr", remote_addr)), count);
|
||||
}
|
||||
}
|
||||
|
||||
// ucp_get_nbx: (void* ep, void* buffer, size_t count, uint64_t remote_addr, void* rkey,
|
||||
// const void* param)
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, void*, size_t count,
|
||||
uint64_t remote_addr, void*, const void*)
|
||||
{
|
||||
if(count == 0) return;
|
||||
|
||||
if(get_use_perfetto()) write_perfetto_counter_track<ucx_recv>(count);
|
||||
|
||||
{
|
||||
cache_comm_data_events<ucx_recv>(0, count);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, count);
|
||||
add(JOIN('/', _name, JOIN('=', "remote_addr", remote_addr)), count);
|
||||
}
|
||||
}
|
||||
|
||||
// ucp_am_send_nbx: (void* ep, unsigned id, const void* header, size_t header_length,
|
||||
// const void* buffer, size_t count, const void* param)
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, unsigned id,
|
||||
const void*, size_t header_length, const void*, size_t count,
|
||||
const void*)
|
||||
{
|
||||
if(count == 0 && header_length == 0) return;
|
||||
|
||||
size_t total_size = header_length + count;
|
||||
if(get_use_perfetto()) write_perfetto_counter_track<ucx_send>(total_size);
|
||||
|
||||
{
|
||||
cache_comm_data_events<ucx_send>(0, total_size);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, total_size);
|
||||
add(JOIN('/', _name, JOIN('=', "am_id", id)), total_size);
|
||||
}
|
||||
}
|
||||
|
||||
// ucp_stream_send_nbx: (void* ep, const void* buffer, size_t count, const void* param)
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, const void*,
|
||||
size_t count, const void*)
|
||||
{
|
||||
if(count == 0) return;
|
||||
|
||||
if(get_use_perfetto()) write_perfetto_counter_track<ucx_send>(count);
|
||||
|
||||
{
|
||||
cache_comm_data_events<ucx_send>(0, count);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, count);
|
||||
}
|
||||
}
|
||||
|
||||
// ucp_stream_recv_nbx: (void* ep, void* buffer, size_t count, size_t* length, const void*
|
||||
// param)
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, void*, size_t count,
|
||||
size_t*, const void*)
|
||||
{
|
||||
if(count == 0) return;
|
||||
|
||||
if(get_use_perfetto()) write_perfetto_counter_track<ucx_recv>(count);
|
||||
|
||||
{
|
||||
cache_comm_data_events<ucx_recv>(0, count);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, count);
|
||||
}
|
||||
}
|
||||
|
||||
// Legacy: ucp_tag_send_nb/nbx - send with tag matching (for old-style wrappers)
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, size_t count, void*,
|
||||
void*, void*)
|
||||
{
|
||||
if(count == 0) return;
|
||||
|
||||
if(get_use_perfetto()) write_perfetto_counter_track<ucx_send>(count);
|
||||
|
||||
{
|
||||
cache_comm_data_events<ucx_send>(0, count);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, count);
|
||||
}
|
||||
}
|
||||
|
||||
// Legacy: ucp_tag_recv_nb/nbx - receive with tag matching (for old-style wrappers)
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, size_t count, void*,
|
||||
void*, void*, void*, void*)
|
||||
{
|
||||
if(count == 0) return;
|
||||
|
||||
if(get_use_perfetto()) write_perfetto_counter_track<ucx_recv>(count);
|
||||
|
||||
{
|
||||
cache_comm_data_events<ucx_recv>(0, count);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, count);
|
||||
}
|
||||
}
|
||||
|
||||
// ucp_put/get operations - RMA
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, size_t length,
|
||||
uint64_t, void*, void*)
|
||||
{
|
||||
if(length == 0) return;
|
||||
|
||||
bool is_put = _data.tool_id.find("ucp_put") != std::string::npos;
|
||||
|
||||
if(get_use_perfetto())
|
||||
{
|
||||
if(is_put)
|
||||
write_perfetto_counter_track<ucx_send>(length);
|
||||
else
|
||||
write_perfetto_counter_track<ucx_recv>(length);
|
||||
}
|
||||
|
||||
{
|
||||
if(is_put)
|
||||
cache_comm_data_events<ucx_send>(0, length);
|
||||
else
|
||||
cache_comm_data_events<ucx_recv>(0, length);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, length);
|
||||
}
|
||||
}
|
||||
|
||||
// ucp_am_send_nb/nbx - active message send
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, unsigned, void*,
|
||||
size_t header_length, void*, size_t length, unsigned, void*)
|
||||
{
|
||||
size_t total_length = header_length + length;
|
||||
if(total_length == 0) return;
|
||||
|
||||
if(get_use_perfetto()) write_perfetto_counter_track<ucx_send>(total_length);
|
||||
|
||||
{
|
||||
cache_comm_data_events<ucx_send>(0, total_length);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, total_length);
|
||||
}
|
||||
}
|
||||
|
||||
// ucp_stream_send/recv operations
|
||||
void
|
||||
comm_data::audit(const gotcha_data& _data, audit::incoming, void*, void*, size_t count,
|
||||
void*, unsigned, void*)
|
||||
{
|
||||
if(count == 0) return;
|
||||
|
||||
bool is_send = _data.tool_id.find("send") != std::string::npos;
|
||||
|
||||
if(get_use_perfetto())
|
||||
{
|
||||
if(is_send)
|
||||
write_perfetto_counter_track<ucx_send>(count);
|
||||
else
|
||||
write_perfetto_counter_track<ucx_recv>(count);
|
||||
}
|
||||
|
||||
{
|
||||
if(is_send)
|
||||
cache_comm_data_events<ucx_send>(0, count);
|
||||
else
|
||||
cache_comm_data_events<ucx_recv>(0, count);
|
||||
}
|
||||
|
||||
if(rocprofsys::get_use_timemory())
|
||||
{
|
||||
auto _name = std::string_view{ _data.tool_id };
|
||||
tracker_t _t{ _name };
|
||||
add(_t, count);
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(ROCPROFSYS_USE_RCCL)
|
||||
// Kept for reference, but now gathered throught the SDK callbacks.
|
||||
|
||||
|
||||
Viittaa uudesa ongelmassa
Block a user