nvshmem_cupy_symmetric_memory_allreduce_sum_quickstart.py

1import cupy as cp
2import nvshmem
3from nvshmem import libnvshmem as nvs
4
5# Initialize NVSHMEM
6# This will also initialize the underlying communication library (e.g., MPI or OpenSHMEM)
7nvshmem.init()
8
9# Get the process ID (mype) and total number of processes (npes)
10mype = nvshmem.my_pe()
11npes = nvshmem.n_pes()
12
13# Select the GPU device corresponding to the local rank
14# (Assuming 1 PE per GPU)
15device_id = mype % cp.cuda.runtime.getDeviceCount()
16cp.cuda.Device(device_id).use()
17
18# Allocate a symmetric array (accessible by all PEs)
19# The size is 10 integers
20size = 10
21data = nvshmem.alloc_array(size, dtype=cp.int32)
22
23# Fill the local array with the value of mype
24data[:] = cp.int32(mype)
25
26# Synchronize to ensure all PEs have initialized their data
27nvshmem.barrier_all()
28
29# Create a destination array for the reduction result
30recvbuf = cp.zeros(size, dtype=cp.int32)
31
32# Perform a collective sum reduction across all PEs
33# Every element in recvbuf will become sum(0, 1, ..., npes-1)
34nvshmem.allreduce(recvbuf, data, op=nvs.SUM)
35
36# Verify the result
37expected_sum = (npes * (npes - 1)) // 2
38print(f"PE {mype}: Result[0] = {recvbuf[0]}, Expected = {expected_sum}")
39
40# Finalize NVSHMEM
41nvshmem.finalize()