diff --git a/include/rocshmem/rocshmem.hpp b/include/rocshmem/rocshmem.hpp index be4de8cf16..3541a93a0c 100644 --- a/include/rocshmem/rocshmem.hpp +++ b/include/rocshmem/rocshmem.hpp @@ -65,6 +65,17 @@ constexpr char VERSION[] = "3.0.0"; */ __host__ void rocshmem_init(MPI_Comm comm = MPI_COMM_WORLD); + +/** + * @brief Query rocSHMEM context from host API + * + * @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 + * kernel bitcode device library in the same application + */ +__host__ void * rocshmem_get_device_ctx(); + /** * @brief Initialize the rocSHMEM runtime and underlying transport layer * with an attempt to enable the requested thread support. diff --git a/src/rocshmem.cpp b/src/rocshmem.cpp index f992a70bbb..e7a8aec48d 100644 --- a/src/rocshmem.cpp +++ b/src/rocshmem.cpp @@ -263,6 +263,14 @@ rocshmem_ctx_t ROCSHMEM_HOST_CTX_DEFAULT; return ROCSHMEM_SUCCESS; } +__host__ void * rocshmem_get_device_ctx() { + void *ctx = nullptr; + + CHECK_HIP(hipMemcpyFromSymbol(&ctx, HIP_SYMBOL(ROCSHMEM_CTX_DEFAULT), + sizeof(rocshmem_ctx_t))); + return ctx; +} + [[maybe_unused]] __host__ int rocshmem_my_pe() { if (backend != nullptr) { return backend->getMyPE();