nvshmem_point_to_point_put_symmetric_memory_cupy.py

1import cupy as cp
2import nvshmempy as nvshmem
3
4# Initialize NVSHMEM
5nvshmem.init()
6
7# Get the rank (mype) and total number of processes (npes)
8mype = nvshmem.my_pe()
9npes = nvshmem.n_pes()
10
11# Check if we have at least 2 PEs for the P2P example
12if npes < 2:
13    if mype == 0:
14        print("This example requires at least 2 PEs.")
15    nvshmem.finalize()
16    exit()
17
18# Allocate a symmetric array (shared memory across PEs)
19# In NVSHMEM Python, we typically use nvshmem.empty or nvshmem.zeros
20# which returns a Cupy-like array that is registered for NVSHMEM operations.
21size = 10
22data = nvshmem.zeros((size,), dtype=cp.float32)
23
24# Source PE (Rank 0) prepares data and puts it into Destination PE (Rank 1)
25if mype == 0:
26    # Fill local data
27    data[:] = cp.arange(size, dtype=cp.float32)
28    print(f"PE {mype}: Sending data to PE 1")
29    
30    # Perform the 'put' operation: copy data to PE 1
31    # Syntax: nvshmem.put(dest_array, src_array, target_pe)
32    nvshmem.put(data, data, 1)
33    
34    # Ensure the operation is complete
35    nvshmem.quiet()
36
37# Synchronize all PEs
38nvshmem.barrier()
39
40# Destination PE (Rank 1) checks the received data
41if mype == 1:
42    print(f"PE {mype}: Data received from PE 0: {data}")
43    expected = cp.arange(size, dtype=cp.float32)
44    if cp.allclose(data, expected):
45        print("Success!")
46    else:
47        print("Failure!")
48
49# Finalize NVSHMEM
50nvshmem.finalize()