[drprajap]

This commit adds a new API rocshmem_hipmodule_init() that allows
user-compiled HIP modules or device bitcode modules to be initialized with
rocSHMEM device contexts. This is essential for CUDA graph compatibility.

This API replaces manual rocshmem ctx management via rocshmem_get_device_ctx()
during JIT compilation.

Resolves: #2741

Technical Details

JIRA ID

Resolves AIROCSHMEM-145

Test Plan

Functional test is added to verify module loading after successfully setting device global symbol.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

Adds a new rocSHMEM host API intended to initialize a user-loaded HIP module’s ROCSHMEM_CTX_DEFAULT symbol so dynamically-loaded device code can use rocSHMEM in CUDA/HIP graph-capture-friendly workflows.

Changes:

• Introduces public API rocshmem_hipmodule_init(hipModule_t, hipStream_t) and implements it in rocshmem_gpu.cpp.
• Adds a new functional test (hipmodule_init) that HIPRTC-compiles a module defining ROCSHMEM_CTX_DEFAULT and attempts to initialize it.
• Wires the new test into the functional test harness (enum, factory, and driver mapping) and links hiprtc.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
projects/rocshmem/include/rocshmem/rocshmem.hpp	Adds the new public API declaration and documentation (removes prior API).
projects/rocshmem/src/rocshmem_gpu.cpp	Implements rocshmem_hipmodule_init() using HIP module symbol lookup and async memcpy.
projects/rocshmem/tests/functional_tests/hipmodule_init_tester.hpp	Declares new functional tester for HIP module init scenario.
projects/rocshmem/tests/functional_tests/hipmodule_init_tester.cpp	Implements HIPRTC compile/load + calls new API, then runs a simple module kernel.
projects/rocshmem/tests/functional_tests/tester.hpp	Adds new TestType enum entry for hipmodule init test.
projects/rocshmem/tests/functional_tests/tester.cpp	Registers and instantiates the new tester in the factory.
projects/rocshmem/tests/functional_tests/CMakeLists.txt	Adds build defines for source/binary include paths and links hiprtc.
projects/rocshmem/scripts/functional_tests/driver.sh	Adds CLI name → test number mapping for hipmodule_init.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[abouteiller]

Looks like the hiprtc is missing a dependency to compile the dynamic code in the CI @drprajap

[drprajap]

Looks like the hiprtc is missing a dependency to compile the dynamic code in the CI @drprajap

yes, because local testing has rocm path set in environment, this was not exposed in my testing, I have explicitly included it now

[omor1]

I think I'm somewhat confused as to why the compiled hiprtc kernel defines its own ROCSHMEM_CTX_DEFAULT symbol instead of using the one from the compiled library that it should link to. I'm not especially familiar with the hiprtc APIs and how to use them though, so I'm likely simply not understanding something about the workflow.

[drprajap]

I think I'm somewhat confused as to why the compiled hiprtc kernel defines its own ROCSHMEM_CTX_DEFAULT symbol instead of using the one from the compiled library that it should link to. I'm not especially familiar with the hiprtc APIs and how to use them though, so I'm likely simply not understanding something about the workflow.

Current test is not using existing rocshmem device library kernels for hiprtc compilation because that requires pulling up lots of dependencies for kernel compilation, I tried that and with existing rocshmem device API definitions, it is not minimal effort. So I chose to orchestrate ROCSHMEM_CTX_DEFAULT device side symbol definition in a minimal device kernel module compilation. The goal of this minimal unit level functional test is to verify Device module initialization path for given compiled module using existing host allocated 'ROCSHMEM_CTX_DEFAULT' symbol and see if new rocshmem_hipmodule_init API is able to get and set the device side symbols without any API compatibility issue.