NVSHMEM Initialization

This section describes some key NVSHMEM initialization considerations when developing applications using the NVSHMEM runtime.

Two-Stage Initialization

This method allows users to select a CUDA device after calling the NVSHMEM initialization APIs nvshmem_init, nvshmemx_init_attr, or nvshmemx_hostlib_init_attr. It is used when the application wants to assign a CUDA device based on the PE returned by nvshmem_my_pe. nvshmem_my_pe can only be called after nvshmem_init or nvshmemx_init_attr. Here is an example of the correct usage:

nvshmemx_init_attr(..., &attr);
// predefined_team can either NVSHMEM_TEAM_WORLD or NVSHMEMX_TEAM_NODE
int mype = [nvshmem_my_pe(), nvshmem_n_pes(), nvshmem_team_my_pe([predefined_team]];
cudaSetDevice(mype);
nvshmem(x)_*(...)

In a two-stage initialization, the actual initialization is deferred to the first call after the device is set. After the NVSHMEM initialization API returns, only a few APIs (listed in the example above) may be called without the library attempting to reinitialize. The library would attempt to reinitialize NVSHMEM in any of the host and on-stream collective APIs like nvshmem_malloc, nvshmem_calloc, nvshmem_align, nvshmem_barrier_all, nvshmem_sync_all, nvshmem_barrier_all_on_stream, or nvshmem_sync_all_on_stream. We recommend that you invoke any other API only after NVSHMEM initialization has completed. To know the status of NVSHMEM initialization, the application can use nvshmemx_init_status.