nvshmem_gpu_symmetric_memory_put_operation_quickstart.py

1import cupy as cp
2import nvshmempy as nvshmem
3
4# Initialize NVSHMEM
5nvshmem.init()
6
7# Get the rank (PE ID) and the total number of PEs
8mype = nvshmem.my_pe()
9npes = nvshmem.n_pes()
10
11# Target PE for the data transfer (neighbor in a ring fashion)
12target_pe = (mype + 1) % npes
13
14# Allocate a symmetric array (accessible by all PEs)
15# Note: nvshmem.empty is used for symmetric allocation
16size = 1024
17data_send = cp.full(size, mype, dtype=cp.float32)
18data_recv = nvshmem.empty(size, dtype=cp.float32)
19
20# Ensure all PEs have finished allocation
21nvshmem.barrier_all()
22
23# Perform a remote 'put' operation: 
24# Copy data_send from the local PE to data_recv on the target PE
25nvshmem.put(data_recv, data_send, target_pe)
26
27# Quiet all pending NVSHMEM operations and synchronize
28nvshmem.barrier_all()
29
30# Verify the result: data_recv should contain the rank of the previous PE
31expected_val = (mype - 1 + npes) % npes
32if cp.all(data_recv == expected_val):
33    print(f"PE {mype}: Verification Successful!")
34else:
35    print(f"PE {mype}: Verification Failed!")
36
37# Finalize NVSHMEM
38nvshmem.finalize()