[rocshmem] Add rocshmem_hipmodule_init API for CUDA graph compatibility

projects/rocshmem/include/rocshmem/rocshmem.hpp

@@ -77,16 +77,41 @@ constexpr char VERSION[] = "3.2.1";
  */
 __host__ void rocshmem_init(void);
 
+
 /**
- * @brief Query rocSHMEM context from host API
+ * @brief Query rocSHMEM context from host API (DEPRECATED)
+ *
+ * @deprecated This API is deprecated. Use rocshmem_hipmodule_init() instead,
+ *             which provides better CUDA graph compatibility and automatic
+ *             module initialization.
  *
  * @param[out] ctx      Returns ROCSHMEM_CTX_DEFAULT device pointer that users
  *                      can query from one instance of rocshmem host library and
- *                      use use later for dynamic module initialization in
+ *                      use later for dynamic module initialization in
  *                      kernel bitcode device library in the same application
  */
+[[deprecated("Use rocshmem_hipmodule_init() instead")]]
 __host__ void * rocshmem_get_device_ctx();
 
+/**
+ * @brief Initialize rocSHMEM device context for a specific HIP module
+ *
+ * This function queries the ROCSHMEM_CTX_DEFAULT symbol from the provided
+ * HIP module and initializes it with the host-side context using stream-ordered
+ * memory operations. This is required for CUDA graph compatibility.
+ *
+ * @param[in] module    HIP module containing rocSHMEM device code
+ * @param[in] stream    HIP stream to use for context initialization (optional,
+ *                      uses hipStreamPerThread if nullptr)
+ *
+ * @return int          0 on success, non-zero on failure
+ *
+ * @note This function does not synchronize the stream, allowing it to be
+ *       used in CUDA graph capture contexts. The caller is responsible for
+ *       synchronization if needed.
+ */
+__host__ int rocshmem_hipmodule_init(hipModule_t module, hipStream_t stream = nullptr);
+
 /**
  * @brief Query rocSHMEM remote symmetric heap pointer
  *

projects/rocshmem/scripts/functional_tests/driver.sh

@@ -122,6 +122,7 @@ declare -A TEST_NUMBERS=(
   ["flood_get"]="86"
   ["flood_getnbi"]="87"
   ["flood_g"]="88"
+  ["hipmodule_init"]="89"
 )
 
 ExecTest() {
@@ -471,6 +472,7 @@ TestOther() {
   #       | Name             | Ranks | Workgroups | Threads | Max Message Size #
   ##############################################################################
   ExecTest  "init"             2       1            1
+ExecTest  "hipmodule_init    2       1            1
 
   ExecTest  "pingpong"         2       1            1
   ExecTest  "pingpong"         2       8            1
@@ -665,6 +667,7 @@ TestGDA() {
   #       | Name             | Ranks | Workgroups | Threads | Max Message Size #
   ##############################################################################
   ExecTest  "init"             2       1            1
+ExecTest  "hipmodule_init    2       1            1
 
   ExecTest  "pingpong"         2       1            1
   ExecTest  "pingpong"         2       8            1

projects/rocshmem/src/rocshmem_gpu.cpp

@@ -114,17 +114,78 @@ __device__ void rocshmem_wg_finalize() {}
 
 
 /******************************************************************************
-* These host APIs use Device side symbol - ROCSHMEM_CTX_DEFAULT so it needs
+* These host API use Device side symbol - ROCSHMEM_CTX_DEFAULT so it needs
 * to stay here to avoid getting pulled into other places in compilation
 ******************************************************************************/
 
 __host__ void * rocshmem_get_device_ctx() {
-  rocshmem_ctx_t ctx;
-
+  rocshmem_ctx_t ctx = {nullptr, nullptr};
   CHECK_HIP(hipMemcpyFromSymbol(&ctx, HIP_SYMBOL(ROCSHMEM_CTX_DEFAULT),
-                             sizeof(rocshmem_ctx_t)));
+                                sizeof(rocshmem_ctx_t)));
   return ctx.ctx_opaque;
+}
+
+__host__ int rocshmem_hipmodule_init(hipModule_t module, hipStream_t stream) {
+  // Step 1: Get device address of rocSHMEM's built-in ROCSHMEM_CTX_DEFAULT symbol
+  // Use hipGetSymbolAddress for graph-capture compatibility (device-to-device path)
+  void *source_ctx_device = nullptr;
+  hipError_t err = hipGetSymbolAddress(&source_ctx_device, HIP_SYMBOL(ROCSHMEM_CTX_DEFAULT));
+
+  if (err != hipSuccess) {
+    fprintf(stderr, "[rocSHMEM] Error: Failed to get address of built-in ROCSHMEM_CTX_DEFAULT: %s\n",
+            hipGetErrorString(err));
+    return ROCSHMEM_ERROR;
+  }
+
+  if (source_ctx_device == nullptr) {
+    fprintf(stderr, "[rocSHMEM] Error: Built-in ROCSHMEM_CTX_DEFAULT has null address\n");
+    return ROCSHMEM_ERROR;
+  }
+
+  // Step 2: Query the device symbol address from the user's HIP module
+  void *target_ctx_device = nullptr;
+  size_t symbol_size = 0;
+
+  err = hipModuleGetGlobal(
+      &target_ctx_device,
+      &symbol_size,
+      module,
+      "ROCSHMEM_CTX_DEFAULT"
+  );
+
+  if (err != hipSuccess) {
+    fprintf(stderr, "[rocSHMEM] Error: Failed to get ROCSHMEM_CTX_DEFAULT symbol from module: %s\n",
+            hipGetErrorString(err));
+    return ROCSHMEM_ERROR;
+  }
+
+  if (symbol_size != sizeof(rocshmem_ctx_t)) {
+    fprintf(stderr, "[rocSHMEM] Error: Symbol size mismatch. Expected %zu, got %zu\n",
+            sizeof(rocshmem_ctx_t), symbol_size);
+    return ROCSHMEM_ERROR;
+  }
+
+  // Step 3: Device-to-device copy using stream-ordered memcpy
+  // This is fully graph-capture compatible since both source and target are device memory
+  if (stream == nullptr) {
+    stream = hipStreamPerThread;
+  }
+
+  err = hipMemcpyAsync(
+      target_ctx_device,
+      source_ctx_device,
+      sizeof(rocshmem_ctx_t),
+      hipMemcpyDeviceToDevice,  // Device-to-device copy for graph capture compatibility
+      stream
+  );
+
+  if (err != hipSuccess) {
+    fprintf(stderr, "[rocSHMEM] Error: Failed to copy context to device: %s\n",
+            hipGetErrorString(err));
+    return ROCSHMEM_ERROR;
+  }
 
+  return 0;
 }
 
 /******************************************************************************

projects/rocshmem/tests/functional_tests/CMakeLists.txt

@@ -93,6 +93,8 @@ target_compile_definitions(
   ${PROJECT_NAME}
   PRIVATE
     $<$<TARGET_EXISTS:PMIx::pmix>:HAVE_PMIX=1>
+    CMAKE_SOURCE_DIR="${CMAKE_SOURCE_DIR}"
+    CMAKE_BINARY_DIR="${CMAKE_BINARY_DIR}"
 )
 
 configure_file(../../scripts/functional_tests/driver.sh rocshmem_functional_driver.sh COPYONLY)
@@ -105,4 +107,5 @@ target_link_libraries(
     roc::rocshmem
     $<TARGET_NAME_IF_EXISTS:MPI::MPI_CXX>
     $<TARGET_NAME_IF_EXISTS:PMIx::pmix>
+    hiprtc
 )