Initialization and Finalization#

This section documents the initialization and finalization APIs in nvshmem.core.init_fini.

The nvshmem4py.core module provides initialization and finalization routines for the NVSHMEM runtime. These must be called before and after using any NVSHMEM features in Python.

Examples#

MPI-based initialization:

from mpi4py import MPI
from cuda.core.experimental import Device
import nvshmem.core as nvshmem

rank = MPI.COMM_WORLD.Get_rank()
dev = Device(rank % system.num_devices)
dev.set_current()

nvshmem.init(device=dev, mpi_comm=MPI.COMM_WORLD, initializer_method="mpi")

# ... use NVSHMEM ...

nvshmem.finalize()

UID-based initialization:

from mpi4py import MPI
from cuda.core.experimental import Device, system
import nvshmem.core as nvshmem
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
nranks = comm.Get_size()

dev = Device(rank % system.num_devices)
dev.set_current()

uid = nvshmem.get_unique_id(empty=(rank != 0))
comm.Bcast(uid._data.view(np.int8), root=0)

nvshmem.init(device=dev, uid=uid, rank=rank, nranks=nranks, initializer_method="uid")

# ... use NVSHMEM ...

nvshmem.finalize()

Emulated MPI initialization:

from mpi4py import MPI
from cuda.core.experimental import Device
import nvshmem.core as nvshmem

rank = MPI.COMM_WORLD.Get_rank()
dev = Device(rank % system.num_devices)
dev.set_current()

nvshmem.init(device=dev, mpi_comm=MPI.COMM_WORLD, initializer_method="emulated_mpi")

# ... use NVSHMEM ...

nvshmem.finalize()

API Reference#

Teams#

class nvshmem.core.Teams(IntEnum)#

NVSHMEM4Py uses an enumerator to refer to NVSHMEM Teams.

TEAM_WORLD#: The world team that contains all PEs in the NVSHMEM program.

TEAM_SHARED#: The team of PEs that share a memory domain. NVSHMEM_TEAM_SHARED refers to the team of all PEs that would mutually return a non-null address from a call to nvshmem_ptr for all symmetric heap objects. That is, nvshmem_ptr must return a non-null pointer to the local PE for all symmetric heap objects on all target PEs in the team. This means that symmetric heap objects on each PE are directly load/store accessible by all PEs in the team. Refer to Team Management for more details about its use.

TEAM_NODE#: The team of PEs that are on the same node

TEAM_SAME_MYPE_NODE#: The team of PEs that are the same PE within a node - that is to say, all PEs for which nvshmem.core.team_my_pe(Teams.TEAM_NODE) returns the same value.

TEAM_SAME_GPU#: The team of PEs that are on the same GPU

TEAM_GPU_LEADERS#: The team of PEs that are leaders of their respective GPUs

Initialization Methods#

NVSHMEM supports multiple bootstrap methods to initialize the runtime. You must explicitly specify one of the following using the initializer_method argument in nvshmem.core.init().

Supported methods:

"mpi": Initializes NVSHMEM using an MPI communicator (mpi4py is required).
"uid": Initializes NVSHMEM using a user-provided unique identifier and rank information.
"emulated_mpi": Uses MPI to broadcast a unique ID internally before doing UID-based init.

Querying Initialization Status#

class nvshmem.core.InitStatus(IntEnum)#

NVSHMEM4Py enumerator for initialization status.

STATUS_NOT_INITIALIZED#: The program is not initialized.

STATUS_IS_BOOTSTRAPPED#: The group of PEs is bootstrapped, but NVSHMEM is not initialized. This means processes can communicate with each other, but CUDA devices are not yet bound to a specific PE. After calling nvshmem.core.finalize() (after a successful initialization), the program will be in this state.

STATUS_IS_INITIALIZED#: The NVSHMEM runtime is initialized. After a succesful call to nvshmem.core.init(), the program will be in this state.

STATUS_LIMITED_MPG#: The NVSHMEM runtime is initialized with limited MPG support. Refer to MPG support for more details.

STATUS_FULL_MPG#: The NVSHMEM runtime is initialized with full MPG support. Refer to MPG support for more details.

STATUS_INVALID#: The program has an invalid state. This is typically due to an error in the initialization process.

Finalization#

When NVSHMEM operations are complete, call nvshmem.core.finalize() to clean up runtime resources.

Retrieving Version Information#

You can query NVSHMEM version details using:

Retrieving a Unique ID#

For UID-based initialization, use: