NVIDIA NVSHMEM (nvshmem4py) - CUDA 12

import nvshmem.core as nvshmem
import os

def main():
    # Initialize NVSHMEM. This is a collective operation.
    # In a real scenario, this script would be launched with `mpiexec` or `nvshmrun`.
    nvshmem.init()

    # Query PE information
    my_pe = nvshmem.my_pe()
    n_pes = nvshmem.n_pes()

    print(f"Hello from PE {my_pe} of {n_pes}")

    # Perform some simple collective (e.g., a barrier)
    # This ensures all PEs reach this point before proceeding
    nvshmem.barrier_all()

    # Finalize NVSHMEM. This is also a collective operation.
    nvshmem.finalize()

if __name__ == '__main__':
    # Note: This script needs to be run using an MPI launcher (e.g., mpiexec -n 2 python your_script.py)
    # or NVSHMEM's own launcher (nvshmrun). Running directly 'python your_script.py'
    # will result in an error or hang if NVSHMEM expects multiple processes.
    try:
        main()
    except Exception as e:
        # Catch potential errors if not launched collectively, for a more graceful exit
        print(f"Error: {e}")
        print("Please ensure the script is launched collectively, e.g., 'mpiexec -n 2 python quickstart.py'")