From b588f604a33b8f39e059b5e7b90c62b947d01b30 Mon Sep 17 00:00:00 2001
From: Maarten Arnst <maarten.arnst@uliege.be>
Date: Mon, 16 Dec 2024 12:53:04 +0100
Subject: [PATCH] Use rocprofiler-sdk-roctx in roctxconnector when available

---
 .github/workflows/build-with-kokkos.yml       |  4 ++--
 profiling/roctx-connector/CMakeLists.txt      | 19 ++++++++++++++-----
 .../roctx-connector/kp_roctx_connector.cpp    |  4 ++++
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/build-with-kokkos.yml b/.github/workflows/build-with-kokkos.yml
index ad56034b1..483ed7bbb 100644
--- a/.github/workflows/build-with-kokkos.yml
+++ b/.github/workflows/build-with-kokkos.yml
@@ -19,10 +19,10 @@ jobs:
           - image: nvidia/cuda:12.2.0-devel-ubuntu22.04
             preset: Cuda
             compiler: {cpp: g++-12, c: gcc-12}
-          - image: rocm/dev-ubuntu-22.04:5.4
+          - image: rocm/dev-ubuntu-22.04:5.7
             preset: ROCm
             compiler: default
-          - image: rocm/dev-ubuntu-22.04:5.7
+          - image: rocm/dev-ubuntu-22.04:6.3
             preset: ROCm
             compiler: default
     container:
diff --git a/profiling/roctx-connector/CMakeLists.txt b/profiling/roctx-connector/CMakeLists.txt
index 2c6857303..39f84adf6 100644
--- a/profiling/roctx-connector/CMakeLists.txt
+++ b/profiling/roctx-connector/CMakeLists.txt
@@ -1,7 +1,16 @@
-find_library(ROCM_ROCTX_LIB roctx64 REQUIRED HINTS $ENV{ROCM_PATH}/lib)
-find_path(ROCM_ROCTX_INCLUDE roctx.h REQUIRED HINTS $ENV{ROCM_PATH}/include/roctracer)
-
 kp_add_library(kp_roctx_connector kp_roctx_connector.cpp)
 
-target_include_directories(kp_roctx_connector PRIVATE ${ROCM_ROCTX_INCLUDE})
-target_link_libraries(kp_roctx_connector PRIVATE ${ROCM_ROCTX_LIB})
+# As of ROCm 6.2, it is recommended to use ROCtx provided by rocprofiler-sdk-roctx
+# instead of the "old" one provided by roctracer.
+#
+# See also: https://rocm.docs.amd.com/projects/rocprofiler-sdk/en/amd-mainline/how-to/using-rocprofv3.html
+find_package(rocprofiler-sdk-roctx CONFIG PATHS $ENV{ROCM_PATH})
+if(rocprofiler-sdk-roctx_FOUND)
+    target_link_libraries(kp_roctx_connector PRIVATE rocprofiler-sdk-roctx::rocprofiler-sdk-roctx)
+else()
+    find_library(ROCM_ROCTX_LIB roctx64 REQUIRED HINTS $ENV{ROCM_PATH}/lib)
+    find_path(ROCM_ROCTX_INCLUDE roctx.h REQUIRED HINTS $ENV{ROCM_PATH}/include/roctracer)
+
+    target_include_directories(kp_roctx_connector PRIVATE ${ROCM_ROCTX_INCLUDE})
+    target_link_libraries(kp_roctx_connector PRIVATE ${ROCM_ROCTX_LIB})
+endif()
diff --git a/profiling/roctx-connector/kp_roctx_connector.cpp b/profiling/roctx-connector/kp_roctx_connector.cpp
index 6c1ea182e..3a1d8b192 100644
--- a/profiling/roctx-connector/kp_roctx_connector.cpp
+++ b/profiling/roctx-connector/kp_roctx_connector.cpp
@@ -14,7 +14,11 @@
 //
 //@HEADER
 
+#if __has_include(<rocprofiler-sdk-roctx/roctx.h>)
+#include <rocprofiler-sdk-roctx/roctx.h>
+#else
 #include <roctx.h>
+#endif
 
 #include <cstdint>
 #include <iostream>