[Rocprofiler-system]: Enable UCX Communication tracing

projects/rocprofiler-systems/cmake/Modules/FindUCX.cmake

@@ -0,0 +1,61 @@
+# ------------------------------------------------------------------------------#
+#
+# Finds headers for UCX (Unified Communication X)
+#
+# UCX is a high-performance communication framework used as a transport layer
+# for MPI and other communication libraries. This module locates UCX headers
+# (ucp.h, uct.h) for tracing and interception purposes.
+#
+# ------------------------------------------------------------------------------#
+
+include(FindPackageHandleStandardArgs)
+
+# ----------------------------------------------------------------------------------------#
+
+set(UCX_HEADERS_INCLUDE_DIR_INTERNAL
+    "${PROJECT_SOURCE_DIR}/source/lib/rocprof-sys/library/tpls/ucx"
+    CACHE PATH
+    "Path to internal UCX headers"
+)
+
+# ----------------------------------------------------------------------------------------#
+# Find UCX headers (ucp.h and uct.h are under ucx/ subdirectory)
+find_path(
+    UCX_HEADERS_INCLUDE_DIR
+    NAMES ucp/api/ucp.h
+    PATHS /usr/include /usr/local/include /opt/ucx/include
+)
+
+if(NOT EXISTS "${UCX_HEADERS_INCLUDE_DIR}")
+    rocprofiler_systems_message(
+        AUTHOR_WARNING
+        "UCX headers do not exist! Setting UCX_HEADERS_INCLUDE_DIR to internal directory: ${UCX_HEADERS_INCLUDE_DIR_INTERNAL}"
+    )
+    set(UCX_HEADERS_INCLUDE_DIR
+        "${UCX_HEADERS_INCLUDE_DIR_INTERNAL}"
+        CACHE PATH
+        "Path to UCX headers"
+        FORCE
+    )
+else()
+    rocprofiler_systems_message(STATUS "UCX headers found: ${UCX_HEADERS_INCLUDE_DIR}")
+endif()
+
+mark_as_advanced(UCX_HEADERS_INCLUDE_DIR)
+
+# ----------------------------------------------------------------------------------------#
+
+find_package_handle_standard_args(UCX DEFAULT_MSG UCX_HEADERS_INCLUDE_DIR)
+
+# ------------------------------------------------------------------------------#
+
+if(UCX_FOUND)
+    add_library(roc::ucx-headers INTERFACE IMPORTED)
+    target_include_directories(
+        roc::ucx-headers
+        SYSTEM
+        INTERFACE ${UCX_HEADERS_INCLUDE_DIR}
+    )
+endif()
+
+# ------------------------------------------------------------------------------#

projects/rocprofiler-systems/cmake/Packages.cmake

@@ -42,6 +42,9 @@ rocprofiler_systems_add_interface_library(rocprofiler-systems-mpi
 rocprofiler_systems_add_interface_library(rocprofiler-systems-libva
     "Provides VA-API headers"
 )
+rocprofiler_systems_add_interface_library(rocprofiler-systems-ucx
+    "Provides UCX headers"
+)
 rocprofiler_systems_add_interface_library(rocprofiler-systems-bfd
     "Provides Binary File Descriptor (BFD)"
 )
@@ -968,6 +971,9 @@ target_include_directories(
     INTERFACE ${LIBVA_HEADERS_INCLUDE_DIR}
 )
 
+find_package(UCX ${rocprofiler_systems_FIND_QUIETLY} REQUIRED)
+target_include_directories(rocprofiler-systems-ucx INTERFACE ${UCX_HEADERS_INCLUDE_DIR})
+
 # ----------------------------------------------------------------------------------------#
 #
 # PTL (Parallel Tasking Library) submodule

projects/rocprofiler-systems/docker/Dockerfile.opensuse

@@ -26,7 +26,7 @@ RUN zypper --non-interactive update -y && \
     zypper --non-interactive install -y chrpath cmake curl dpkg-devel \
         gcc-c++ gcc-fortran git gmock gtest iproute2 libdrm-devel libnuma-devel \
         ninja nlohmann_json-devel openmpi3-devel python3-pip rpm-build \
-        sqlite3-devel wget && \
+        sqlite3-devel wget libucp-devel libuct-devel && \
     python3 -m pip install 'cmake==3.21'
 
 ARG ROCM_VERSION=0.0

projects/rocprofiler-systems/docker/Dockerfile.opensuse.ci

@@ -30,7 +30,7 @@ RUN zypper --non-interactive update -y && \
     zypper --non-interactive install -y chrpath cmake curl dpkg-devel \
         gcc-c++ gcc-fortran git gmock gtest iproute2 ninja nlohmann_json-devel \
         openmpi3-devel papi-devel python3-devel python3-pip rpm-build \
-        sqlite3-devel vim wget && \
+        sqlite3-devel vim wget libucp-devel libuct-devel && \
     zypper --non-interactive clean --all && \
     python3 -m pip install 'cmake==3.21' perfetto
 

projects/rocprofiler-systems/docker/Dockerfile.rhel

@@ -17,7 +17,7 @@ RUN yum groupinstall -y "Development Tools" && \
     yum install -y epel-release && crb enable && \
     yum install -y --allowerasing chrpath cmake curl dpkg-devel gmock-devel gtest-devel \
         iproute json-devel libdrm-devel ninja-build numactl-devel openmpi-devel \
-        papi-devel python3-pip sqlite-devel texinfo wget which zlib-devel && \
+        papi-devel python3-pip sqlite-devel texinfo wget which zlib-devel ucx-devel && \
     yum clean all && \
     python3 -m pip install 'cmake==3.21' && \
     python3 -m pip install 'perfetto'

projects/rocprofiler-systems/docker/Dockerfile.rhel.ci

@@ -21,7 +21,7 @@ RUN yum groupinstall -y "Development Tools" && \
     yum install -y epel-release && crb enable && \
     yum install -y --allowerasing chrpath cmake curl dpkg-devel gmock-devel gtest-devel \
         iproute json-devel ninja-build numactl-devel openmpi-devel papi-devel \
-        python3-devel python3-pip sqlite-devel texinfo wget which vim zlib-devel && \
+        python3-devel python3-pip sqlite-devel texinfo wget which vim zlib-devel ucx-devel && \
     yum clean all && \
     python3 -m pip install 'cmake==3.21' perfetto
 

projects/rocprofiler-systems/docker/Dockerfile.ubuntu

@@ -29,7 +29,8 @@ RUN apt-get update && \
         build-essential chrpath cmake curl flex gettext git-core gnupg2 iproute2 \
         libgmock-dev libgtest-dev libnuma1 libopenmpi-dev libpapi-dev libpfm4-dev \
         librpm-dev libsqlite3-dev libtool libudev1 lsb-release m4 ninja-build \
-        nlohmann-json3-dev python3-pip rpm texinfo wget && \
+        nlohmann-json3-dev python3-pip rpm texinfo wget \
+        libucx-dev ucx-utils && \
     OS_VERSION=$(grep '^VERSION_ID=' /etc/os-release | cut -d'=' -f2 | tr -d '"') && \
     OS_ID=$(grep '^ID=' /etc/os-release | cut -d'=' -f2 | tr -d '"') && \
     if [ "${OS_ID}" == "ubuntu" ] && [ "${OS_VERSION}" == "22.04" ]; then \

projects/rocprofiler-systems/docker/Dockerfile.ubuntu.ci

@@ -27,7 +27,8 @@ RUN apt-get update && \
         bzip2 chrpath cmake curl environment-modules flex gettext git-core gnupg2 \
         gzip iproute2 libgmock-dev libgtest-dev libiberty-dev libpapi-dev libpfm4-dev \
         libsqlite3-dev libtool locales lsb-release m4 ninja-build nlohmann-json3-dev \
-        python3-pip software-properties-common texinfo unzip wget vim zip zlib1g-dev && \
+        python3-pip software-properties-common texinfo unzip wget vim zip zlib1g-dev \
+        libucx-dev ucx-utils && \
     apt-get autoclean
 
 RUN OS_VERSION=$(grep '^VERSION_ID=' /etc/os-release | cut -d'=' -f2 | tr -d '"') && \

projects/rocprofiler-systems/source/lib/CMakeLists.txt

@@ -47,6 +47,7 @@ target_link_libraries(
         $<BUILD_INTERFACE:rocprofiler-systems::rocprofiler-systems-elfutils>
         $<BUILD_INTERFACE:rocprofiler-systems::rocprofiler-systems-bfd>
         $<BUILD_INTERFACE:rocprofiler-systems::rocprofiler-systems-mpi>
+        $<BUILD_INTERFACE:rocprofiler-systems::rocprofiler-systems-ucx>
         $<BUILD_INTERFACE:rocprofiler-systems::rocprofiler-systems-libva>
         $<BUILD_INTERFACE:rocprofiler-systems::rocprofiler-systems-ptl>
         $<BUILD_INTERFACE:rocprofiler-systems::rocprofiler-systems-rocm>

projects/rocprofiler-systems/source/lib/core/categories.hpp

@@ -127,8 +127,9 @@ ROCPROFSYS_DEFINE_CATEGORY(category, rocm_rccl, ROCPROFSYS_CATEGORY_ROCM_RCCL, "
 ROCPROFSYS_DEFINE_CATEGORY(category, pthread, ROCPROFSYS_CATEGORY_PTHREAD, "pthread", "POSIX threading functions")
 ROCPROFSYS_DEFINE_CATEGORY(category, kokkos, ROCPROFSYS_CATEGORY_KOKKOS, "kokkos", "KokkosTools regions")
 ROCPROFSYS_DEFINE_CATEGORY(category, mpi, ROCPROFSYS_CATEGORY_MPI, "mpi", "MPI regions")
+ROCPROFSYS_DEFINE_CATEGORY(category, ucx, ROCPROFSYS_CATEGORY_UCX, "ucx", "UCX regions")
 ROCPROFSYS_DEFINE_CATEGORY(category, process_sampling, ROCPROFSYS_CATEGORY_PROCESS_SAMPLING, "process_sampling", "Process-level data")
-ROCPROFSYS_DEFINE_CATEGORY(category, comm_data, ROCPROFSYS_CATEGORY_COMM_DATA, "comm_data", "MPI/RCCL counters for tracking amount of data sent or received")
+ROCPROFSYS_DEFINE_CATEGORY(category, comm_data, ROCPROFSYS_CATEGORY_COMM_DATA, "comm_data", "MPI/RCCL/UCX counters for tracking amount of data sent or received")
 ROCPROFSYS_DEFINE_CATEGORY(category, causal, ROCPROFSYS_CATEGORY_CAUSAL, "causal", "Causal profiling data")
 ROCPROFSYS_DEFINE_CATEGORY(category, cpu_freq, ROCPROFSYS_CATEGORY_CPU_FREQ, "cpu_frequency", "CPU frequency (collected in background thread)")
 ROCPROFSYS_DEFINE_CATEGORY(category, process_page, ROCPROFSYS_CATEGORY_PROCESS_PAGE, "process_physical_memory", "Physical memory usage (RSS) in process in MB (collected in background thread)")
@@ -207,6 +208,7 @@ using name = perfetto_category<Tp...>;
         ROCPROFSYS_PERFETTO_CATEGORY(category::pthread),                                 \
         ROCPROFSYS_PERFETTO_CATEGORY(category::kokkos),                                  \
         ROCPROFSYS_PERFETTO_CATEGORY(category::mpi),                                     \
+        ROCPROFSYS_PERFETTO_CATEGORY(category::ucx),                                     \
         ROCPROFSYS_PERFETTO_CATEGORY(category::sampling),                                \
         ROCPROFSYS_PERFETTO_CATEGORY(category::process_sampling),                        \
         ROCPROFSYS_PERFETTO_CATEGORY(category::comm_data),                               \

projects/rocprofiler-systems/source/lib/core/config.cpp

@@ -377,6 +377,10 @@ configure_settings(bool _init)
                               "Enable support for MPI functions", true, "mpi", "backend",
                               "parallelism");
 
+    ROCPROFSYS_CONFIG_SETTING(bool, "ROCPROFSYS_USE_UCX",
+                              "Enable support for UCX functions", true, "ucx", "backend",
+                              "parallelism");
+
     ROCPROFSYS_CONFIG_SETTING(
         bool, "ROCPROFSYS_USE_RCCLP",
         "Enable support for ROCm Communication Collectives Library (RCCL) Performance",
@@ -1938,6 +1942,13 @@ get_use_mpip()
     return static_cast<tim::tsettings<bool>&>(*_v->second).get();
 }
 
+bool&
+get_use_ucx()
+{
+    static auto _v = get_config()->find("ROCPROFSYS_USE_UCX");
+    return static_cast<tim::tsettings<bool>&>(*_v->second).get();
+}
+
 bool
 get_use_kokkosp()
 {

projects/rocprofiler-systems/source/lib/core/config.hpp

@@ -225,6 +225,9 @@ get_use_pid();
 bool&
 get_use_mpip();
 
+bool&
+get_use_ucx();
+
 bool
 get_use_kokkosp();
 

projects/rocprofiler-systems/source/lib/rocprof-sys-user/rocprofiler-systems/categories.h

@@ -80,6 +80,7 @@ extern "C"
         ROCPROFSYS_CATEGORY_PTHREAD,
         ROCPROFSYS_CATEGORY_KOKKOS,
         ROCPROFSYS_CATEGORY_MPI,
+        ROCPROFSYS_CATEGORY_UCX,
         ROCPROFSYS_CATEGORY_PROCESS_SAMPLING,
         ROCPROFSYS_CATEGORY_COMM_DATA,
         ROCPROFSYS_CATEGORY_CAUSAL,

projects/rocprofiler-systems/source/lib/rocprof-sys/library.cpp

@@ -55,6 +55,7 @@
 #include "library/components/mpi_gotcha.hpp"
 #include "library/components/numa_gotcha.hpp"
 #include "library/components/pthread_gotcha.hpp"
+#include "library/components/ucx_gotcha.hpp"
 #include "library/components/vaapi_gotcha.hpp"
 #include "library/coverage.hpp"
 #include "library/process_sampler.hpp"
@@ -609,6 +610,12 @@ rocprofsys_init_tooling_hidden(void)
     // start these gotchas once settings have been initialized
     if(get_init_bundle()) get_init_bundle()->start();
 
+    if(get_use_ucx())
+    {
+        LOG_DEBUG("Setting up UCX traces...\n");
+        component::ucx_gotcha::start();
+    }
+
     if(get_use_vaapi_tracing())
     {
         LOG_DEBUG("Setting up VA-API traces...");
@@ -900,6 +907,12 @@ rocprofsys_finalize_hidden(void)
     fini_bundle_t _finalization{};
     _finalization.start();
 
+    if(get_use_ucx())
+    {
+        LOG_DEBUG("Shutting down UCX tracing...\n");
+        component::ucx_gotcha::shutdown();
+    }
+
     if(get_use_vaapi_tracing())
     {
         LOG_DEBUG("Shutting down VA-API tracing...");

projects/rocprofiler-systems/source/lib/rocprof-sys/library/components/CMakeLists.txt

@@ -11,6 +11,7 @@ set(component_sources
     ${CMAKE_CURRENT_LIST_DIR}/kill_gotcha.cpp
     ${CMAKE_CURRENT_LIST_DIR}/mpi_gotcha.cpp
     ${CMAKE_CURRENT_LIST_DIR}/numa_gotcha.cpp
+    ${CMAKE_CURRENT_LIST_DIR}/ucx_gotcha.cpp
     ${CMAKE_CURRENT_LIST_DIR}/vaapi_gotcha.cpp
     ${CMAKE_CURRENT_LIST_DIR}/pthread_gotcha.cpp
     ${CMAKE_CURRENT_LIST_DIR}/pthread_create_gotcha.cpp
@@ -32,6 +33,7 @@ set(component_headers
     ${CMAKE_CURRENT_LIST_DIR}/mpip.hpp
     ${CMAKE_CURRENT_LIST_DIR}/mpi_gotcha.hpp
     ${CMAKE_CURRENT_LIST_DIR}/numa_gotcha.hpp
+    ${CMAKE_CURRENT_LIST_DIR}/ucx_gotcha.hpp
     ${CMAKE_CURRENT_LIST_DIR}/vaapi_gotcha.hpp
     ${CMAKE_CURRENT_LIST_DIR}/pthread_gotcha.hpp
     ${CMAKE_CURRENT_LIST_DIR}/pthread_create_gotcha.hpp