67 lines
2.4 KiB
Plaintext
67 lines
2.4 KiB
Plaintext
/*
|
|
* Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
|
*
|
|
* NVIDIA CORPORATION and its licensors retain all intellectual property
|
|
* and proprietary rights in and to this software, related documentation
|
|
* and any modifications thereto. Any use, reproduction, disclosure or
|
|
* distribution of this software and related documentation without an express
|
|
* license agreement from NVIDIA CORPORATION is strictly prohibited.
|
|
*
|
|
* See COPYRIGHT.txt for license information
|
|
*/
|
|
|
|
#include <stdio.h>
|
|
#include "mpi.h"
|
|
#include "nvshmem.h"
|
|
#include "nvshmemx.h"
|
|
|
|
#define CUDA_CHECK(stmt) \
|
|
do { \
|
|
cudaError_t result = (stmt); \
|
|
if (cudaSuccess != result) { \
|
|
fprintf(stderr, "[%s:%d] CUDA failed with %s \n", __FILE__, __LINE__, \
|
|
cudaGetErrorString(result)); \
|
|
exit(-1); \
|
|
} \
|
|
} while (0)
|
|
|
|
__global__ void simple_shift(int *destination) {
|
|
int mype = nvshmem_my_pe();
|
|
int npes = nvshmem_n_pes();
|
|
int peer = (mype + 1) % npes;
|
|
|
|
nvshmem_int_p(destination, mype, peer);
|
|
}
|
|
|
|
int main(int argc, char *argv[]) {
|
|
int mype_node, msg;
|
|
cudaStream_t stream;
|
|
int rank, nranks;
|
|
MPI_Comm mpi_comm = MPI_COMM_WORLD;
|
|
nvshmemx_init_attr_t attr = NVSHMEMX_INIT_ATTR_INITIALIZER;
|
|
|
|
MPI_Init(&argc, &argv);
|
|
MPI_Comm_rank(MPI_COMM_WORLD, &rank);
|
|
MPI_Comm_size(MPI_COMM_WORLD, &nranks);
|
|
|
|
attr.mpi_comm = &mpi_comm;
|
|
nvshmemx_init_attr(NVSHMEMX_INIT_WITH_MPI_COMM, &attr);
|
|
mype_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE);
|
|
|
|
CUDA_CHECK(cudaSetDevice(mype_node));
|
|
CUDA_CHECK(cudaStreamCreate(&stream));
|
|
int *destination = (int *)nvshmem_malloc(sizeof(int));
|
|
|
|
simple_shift<<<1, 1, 0, stream>>>(destination);
|
|
nvshmemx_barrier_all_on_stream(stream);
|
|
CUDA_CHECK(cudaMemcpyAsync(&msg, destination, sizeof(int), cudaMemcpyDeviceToHost, stream));
|
|
|
|
CUDA_CHECK(cudaStreamSynchronize(stream));
|
|
printf("%d: received message %d\n", nvshmem_my_pe(), msg);
|
|
|
|
nvshmem_free(destination);
|
|
nvshmem_finalize();
|
|
MPI_Finalize();
|
|
return 0;
|
|
}
|