From 0afd607328fbfcccb698c26e85f723d9f6017fc6 Mon Sep 17 00:00:00 2001
From: Stanley Tsang <stanley.tsang@amd.com>
Date: Tue, 26 Oct 2021 10:10:04 -0700
Subject: [PATCH] Re-enable use of chrpath to manually set rpath for unit
 tests. (#448)

* Re-enable use of chrpath to manually set rpath for unit tests.

* Add check for chrpath

[ROCm/rccl commit: d23dfc12c1b23fcdc188ca8149497f7f054a12f7]
---
 projects/rccl/test/CMakeLists.txt | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/projects/rccl/test/CMakeLists.txt b/projects/rccl/test/CMakeLists.txt
index 4c8ebf73c5..7532d34c1f 100644
--- a/projects/rccl/test/CMakeLists.txt
+++ b/projects/rccl/test/CMakeLists.txt
@@ -5,6 +5,11 @@ if(BUILD_TESTS)
 
   message("Going to build unit tests (Installed in /test/UnitTests)")
 
+  find_program(CHRPATH chrpath)
+  if(NOT CHRPATH)
+      message(FATAL_ERROR "chrpath is required for UnitTests. Please install (e.g. sudo apt-get install chrpath)")
+  endif()
+
   include_directories(${GTEST_INCLUDE_DIRS})
 
   # Collect source files for tests
@@ -60,6 +65,17 @@ if(BUILD_TESTS)
     target_link_libraries(UnitTests PRIVATE rccl)
     target_link_libraries(UnitTestsMultiProcess PRIVATE rt rccl)
   endif()
+  # HIPCC adds /opt/rocm/lib as RPATH, even though the install process is supposed to
+  # remove RPATH.  It also occurs before any user-specified rpath, which effectively overrides the user rpath.
+  #  As a work-around, set the correct RPATH for the unit test executable as a post-install step
+  if (CMAKE_INSTALL_PREFIX STREQUAL "/opt/rocm")
+    # install_prefix/CMAKE_INSTALL_PREFIX was not explicitly specified, so look in build/release
+    add_custom_command( TARGET UnitTests POST_BUILD COMMAND chrpath ARGS -r ${CMAKE_BINARY_DIR}:/opt/rocm/lib ${CMAKE_BINARY_DIR}/test/UnitTests)
+    add_custom_command( TARGET UnitTestsMultiProcess POST_BUILD COMMAND chrpath ARGS -r ${CMAKE_BINARY_DIR}:/opt/rocm/lib ${CMAKE_BINARY_DIR}/test/UnitTestsMultiProcess)
+  else()
+    add_custom_command( TARGET UnitTests POST_BUILD COMMAND chrpath ARGS -r ${CMAKE_INSTALL_PREFIX}/lib:/opt/rocm/lib ${CMAKE_INSTALL_PREFIX}/test/UnitTests)
+    add_custom_command( TARGET UnitTestsMultiProcess POST_BUILD COMMAND chrpath ARGS -r ${CMAKE_INSTALL_PREFIX}/lib:/opt/rocm/lib ${CMAKE_INSTALL_PREFIX}/test/UnitTestsMultiProcess)
+  endif()
 else()
   message("Not building unit tests")
 endif()