Kathryn-cat commented on code in PR #18093:
URL: https://github.com/apache/tvm/pull/18093#discussion_r2167458235
##########
src/runtime/contrib/nvshmem/init.cc:
##########
@@ -106,12 +106,26 @@ void InitNVSHMEMWrapper(String args) {
InitNVSHMEM(uid_64, num_workers, worker_id_start);
}
+void NVSHMEMXCumoduleInit(void* cuModule) {
+ CUmodule mod = static_cast<CUmodule>(cuModule);
+ auto status = nvshmemx_init_status();
+ // The NVSHMEM library must have completed device initialization prior to
+ // nvshmemx_cumodule_init. If not, we skip the cumodule initialization.
Review Comment:
The design here is to enable NVSHMEM compilation and linking broadly for
every kernel, including those whose NVSHMEM context is not initialized and do
not use NVSHMEM in their kernels.
In such case, `nvshmemx_init_status()` is used to check whether we need to
call `nvshmemx_cumodule_init` or not. If not device initialized, we just skip
`nvshmemx_cumodule_init`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
