Move rocshmem_home logic to subproject

Author: alextmagro

transformer_engine/common/CMakeLists.txt

@@ -386,32 +386,13 @@ target_compile_definitions(transformer_engine PUBLIC NV_CUDNN_FRONTEND_USE_DYNAM
 else()
   option(NVTE_ENABLE_ROCSHMEM "Compile with ROCSHMEM library" OFF)
   if (NVTE_ENABLE_ROCSHMEM)
-    find_package(MPI REQUIRED)
-    target_link_libraries(transformer_engine PUBLIC MPI::MPI_CXX)
-    target_include_directories(transformer_engine PRIVATE ${MPI_CXX_INCLUDES})
     add_subdirectory(rocshmem_api)
-    if(DEFINED ENV{ROCSHMEM_HOME})
-      set(ROCSHMEM_HOME "$ENV{ROCSHMEM_HOME}" CACHE STRING "Location of ROCSHMEM installation")
-    else()
-      set(ROCSHMEM_HOME "/opt/rocm" CACHE STRING "Location of ROCSHMEM installation (default)")
-    endif()
-    set(ROCSHMEM_LIBRARY_PATH "${ROCSHMEM_HOME}/lib/librocshmem.a")
-    if (EXISTS ${ROCSHMEM_LIBRARY_PATH})
-      add_library(rocshmem STATIC IMPORTED)
-      set_target_properties(rocshmem PROPERTIES
-        IMPORTED_LOCATION "${ROCSHMEM_LIBRARY_PATH}"
-        IMPORTED_LINK_INTERFACE_LANGUAGES "CXX"
-      )
-    else()
-      message(FATAL_ERROR "ROCSHMEM library not found at ${ROCSHMEM_LIBRARY_PATH}. PLease set ROCSHMEM_HOME.")
-    endif()
     target_link_options(transformer_engine PRIVATE
         -fgpu-rdc
     )
     target_link_libraries(transformer_engine PUBLIC 
         -Wl,--whole-archive
         rocshmemapi
-        rocshmem
         -Wl,--no-whole-archive
     )
   endif()

transformer_engine/common/rocshmem_api/CMakeLists.txt

@@ -1,23 +1,12 @@
 # Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
 # License for AMD contributions = MIT. See LICENSE for more information
+
 cmake_minimum_required (VERSION 3.21)
 project(rocshmem LANGUAGES HIP)
 
 find_package(hip REQUIRED)
 find_package(MPI REQUIRED)
 
-
-if(DEFINED ENV{ROCSHMEM_HOME})
-  set(ROCSHMEM_HOME "$ENV{ROCSHMEM_HOME}" CACHE STRING "Location of ROCSHMEM installation")
-else()
-  set(ROCSHMEM_HOME "/opt/rocm" CACHE STRING "Location of ROCSHMEM installation (default)")
-endif()
-
-set(ROCSHMEM_INCLUDE_DIR "${ROCSHMEM_HOME}/include/rocshmem")
-if(NOT EXISTS "${ROCSHMEM_INCLUDE_DIR}")
-  set(ROCSHMEM_INCLUDE_DIR "${ROCSHMEM_HOME}/include")
-endif()
-
 add_library(rocshmemapi OBJECT rocshmem_waitkernel.hip)
 
 target_compile_options(rocshmemapi PRIVATE
@@ -30,14 +19,35 @@ set_target_properties(rocshmemapi PROPERTIES
     POSITION_INDEPENDENT_CODE ON
 )
 
+if(DEFINED ENV{ROCSHMEM_HOME})
+  set(ROCSHMEM_HOME "$ENV{ROCSHMEM_HOME}" CACHE STRING "Location of ROCSHMEM installation")
+else()
+  set(ROCSHMEM_HOME "/opt/rocm" CACHE STRING "Location of ROCSHMEM installation (default)")
+endif()
+set(ROCSHMEM_LIBRARY_PATH "${ROCSHMEM_HOME}/lib/librocshmem.a")
+if (EXISTS ${ROCSHMEM_LIBRARY_PATH})
+  add_library(rocshmem STATIC IMPORTED)
+  set_target_properties(rocshmem PROPERTIES
+    IMPORTED_LOCATION "${ROCSHMEM_LIBRARY_PATH}"
+    IMPORTED_LINK_INTERFACE_LANGUAGES "CXX"
+  )
+else()
+  message(FATAL_ERROR "ROCSHMEM library not found at ${ROCSHMEM_LIBRARY_PATH}. PLease set ROCSHMEM_HOME.")
+endif()
+
+set(ROCSHMEM_INCLUDE_DIR "${ROCSHMEM_HOME}/include/rocshmem")
+if(NOT EXISTS "${ROCSHMEM_INCLUDE_DIR}")
+  set(ROCSHMEM_INCLUDE_DIR "${ROCSHMEM_HOME}/include")
+endif()
+
 target_include_directories(rocshmemapi PUBLIC
         "${ROCSHMEM_INCLUDE_DIR}"
         "${CMAKE_CURRENT_SOURCE_DIR}"
         "${MPI_INCLUDE_PATH}"
 )
 
 target_link_libraries(rocshmemapi PUBLIC
-    "${ROCSHMEM_HOME}/lib/librocshmem.a"
+    rocshmem
     MPI::MPI_CXX
     hip::host
     hip::device