Is this page helpful?

NVSHMEM Device Remote Memory Access (RMA) with Numba-CUDA DSL#

This section documents the NVSHMEM Device Remote Memory Access (RMA) operations with Numba-CUDA DSL.

Using put and get in a Numba-CUDA Kernel#

The following example demonstrates how to use the NVSHMEM put and get operations in a Numba-CUDA kernel. These allow threads to write to and read from memory on a remote PE (processing element) directly from device code.

import numpy as np
import cupy as cp
from numba import cuda
import nvshmem
import nvshmem.core.device.numba as nvshmem_numba
from mpi4py import MPI

@cuda.jit
def rma_kernel(src, dst, remote_buf, pe):
    tid = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    if tid == 0:
        # Put data from src to remote_buf on remote PE
        nvshmem_numba.put(remote_buf, src, pe)
        # Get data from remote_buf on remote PE to dst
        nvshmem_numba.get(dst, remote_buf, pe)

# Initialize NVSHMEM
dev = cudaDevice()
dev.set_current()
stream = dev.create_stream()
nvshmem.init(dev=dev, mpi_comm=MPI.COMM_WORLD, initializer_method="mpi", stream=stream)

# Get information about the current PE
me = nvshmem.my_pe()
n_pes = nvshmem.n_pes()

# Choose a remote PE (for example, next PE in a ring)
pe = (me + 1) % n_pes

# Allocate device buffers
src = nvshmem.array((1,), dtype=np.int32)
dst = nvshmem.array((1,), dtype=np.int32)
remote_buf = nvshmem.array((1,), dtype=np.int32)

# Launch kernel to perform put and get on remote PE's buffer
# Note, Numba-cuda does not accept a cuda.core Stream, so we need to pass the stream handle.
rma_kernel[1, 1](src, dst, remote_buf, pe, stream=int(stream.handle))

# Finalize NVSHMEM
nvshmem.finalize(dev=dev, stream=stream)

This example puts the value from src[0] to the remote_buf[0] on the next PE in a ring, and then gets the value back into dst[0]. Only thread 0 performs the RMA operations for demonstration purposes. In practice, you can have multiple threads performing RMA as needed.

For more details, refer to the Numba test suite and the NVSHMEM4Py documentation.

NVSHMEM4Py Memory Management with Numba-CUDA DSL#

This section documents the NVSHMEM4Py Memory Management with Numba-CUDA DSL.

NVSHMEM4Py provides functions to access remote symmetric and multicast buffers as CuPy arrays through the nvshmem.core.device.numba.mem module.