NVSHMEM aggregates the memory of multiple GPUs in a cluster into a Partitioned Global Address Space (PGAS) that enables fine-grained GPU-to-GPU data movement and synchronization from within a CUDA kernel.

Using NVSHMEM, developers can write long running kernels that include both communication and computation, reducing the need for synchronization with the CPU. These composite kernels also allow for fine-grain overlap of computation with communication as a result of thread warp scheduling on the GPU. NVSHMEM GPU-initiated communication can reduce overheads resulting from kernel launches, calls to CUDA API, and CPU-GPU synchronization. Reducing these overheads can enable significant gains in strong scaling of application workloads. When necessary, NVSHMEM also provides the flexibility of CPU-side calls for inter-GPU communication outside of CUDA kernels.

The Message Passing Interface (MPI) is one of the most commonly used communication libraries for scalable, high performance computing. A key design principle of NVSHMEM is support for fine-grain, highly-concurrent GPU-to-GPU communication from within CUDA kernels. Utilizing conventional MPI send and receive operations in such a communication regime can lead to significant efficiency challenges. Because of the need to match send and receive operations, MPI implementations can incur high locking (or atomics) overheads for shared data structures involved in messaging, serialization overheads that result from the MPI message ordering requirements, and protocol overheads that result from messages arriving at the receiver before they have posted the corresponding receive operation.

While there have been efforts to reduce the overhead of critical sections using fine-grained locking and multiple network end-points per process, matching between send and receive operations is inherent to the send-receive model of communication in MPI and is challenging to scale to highly threaded environments presented by GPUs. In contrast, one-sided communication primitives avoid these bottlenecks by enabling the initiating thread to specify all the information required to complete a data transfer. They can be directly translated to RDMA primitives exposed by the network hardware or to load/store operations on the NVLink fabric. Using asynchronous APIs, one-sided primitives also make it programmatically easier to interleave computation and communication, thereby having the potential for better overlap.

NVSHMEM support for GPU-initiated communication can significantly reduce communication and synchronization overheads, leading to improvements in strong scaling.

Strong scaling refers to how the solution time of a fixed problem changes as the number of processors is increased. In other words, strong scaling is the ability to solve a fixed problem faster by increasing the number of processors. This is a critical metric for many scientific, engineering, and data analytics applications as increasingly large clusters and more powerful GPUs become available.

Current state-of-the-art applications running on GPU clusters typically offload computation phases onto the GPU and rely on the CPU to manage communication between cluster nodes, using the Message Passing Interface (MPI) or OpenSHMEM. Dependency on the CPU for communication limits strong scalability, owing to the overhead of repeated kernel launches, CPU-GPU synchronization, underutilization of the GPU during communication phases, and underutilization of the network during compute phases. Some of these issues can be addressed by restructuring application code to overlap independent compute and communication phases using CUDA streams. These optimizations can lead to complex application code and their benefits usually diminish as the problem size per GPU becomes smaller.

As a problem is strong-scaled, the Amdahl’s fraction of the execution corresponding to CPU-GPU synchronization and communication increases. Thus, minimizing these overheads is critical for strong scaling of applications on GPU clusters. Furthermore, GPUs are designed for throughput and have enough state and parallelism to hide long latencies to global memory, which can allow them to be efficient at hiding data movement overheads. Following the CUDA programming model and best practices and utilizing NVSHMEM for GPU-initiated communication enables developers to take advantage of these latency hiding capabilities.

NVSHMEM extends the OpenSHMEM APIs to support clusters of NVIDIA GPUs. The following provides a brief summary of several key extensions:

The code snippet below shows a simple example of NVSHMEM usage within a CUDA kernel where PEs form a communication ring.

#include <stdio.h>
#include <cuda.h>
#include <nvshmem.h>
#include <nvshmemx.h>

__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(void) {
    int mype_node, msg;
    cudaStream_t stream;

    nvshmem_init();
    mype_node = nvshmemx_my_pe(NVSHMEMX_TEAM_NODE);
    cudaSetDevice(mype_node);
    cudaStreamCreate(&stream);

    int *destination = (int *) nvshmem_malloc(sizeof(int));

    simple_shift<<<1, 1, 0, stream>>>(destination);
    nvshmemx_barrier_all_on_stream(stream);
    cudaMemcpyAsync(&msg, destination, sizeof(int), cudaMemcpyDeviceToHost, stream);

    cudaStreamSynchronize(stream);
    printf("%d: received message %d\n", nvshmem_my_pe(), msg);

    nvshmem_free(destination);
    nvshmem_finalize();
    return 0;
}

This example begins in main by initializing the NVSHMEM library, querying the PE’s ID in the on-node team, and using the on-node ID to set the CUDA device. The device must be set prior to allocating memory or launching any kernels. Next, a stream is created and a symmetric integer called destination is allocated on every PE. After this, the simple_shift kernel is launched on a single thread with a pointer to this symmetric object as its argument.

Figure 2. Illustration of the communication performed by the simple_shift kernel.

This kernel queries the global PE ID and the number of executing PEs. It then performs a single-element integer put operation to write the calling PE’s ID into destination at the PE with the next highest ID, or 0 in the case of the PE with the highest ID. The kernel is launched asynchronously on stream, followed by an NVSHMEM barrier on the stream to ensure that all updates have completed, and an asynchronous copy to copy the updated destination value to the host. The stream is then synchronized and the result is printed. Example output with 8 PEs is shown below.

0: received message 7
1: received message 0
2: received message 1
4: received message 3
6: received message 5
7: received message 6
3: received message 2
5: received message 4

Finally, the destination buffer is freed and the NVSHMEM library is finalized prior to the program exiting.

NVSHMEM can be used in conjunction with OpenSHMEM or MPI, making it easier for existing OpenSHMEM and MPI applications to be incrementally ported to use NVSHMEM. The code snippet below shows how NVSHMEM can be initialized within an MPI program. In this program, we assume that each MPI process is also an NVSHMEM PE -- that is, that each process has both an MPI rank as well as an NVSHMEM rank.

int main(int argc, char *argv[]) {
    int rank, ndevices;

    nvshmemx_init_attr_t attr;
    MPI_Comm comm = MPI_COMM_WORLD;
    attr.mpi_comm = &comm;

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(&rank, MPI_COMM_WORLD));	
    
    cudaGetDeviceCount(&ndevices);
    cudaSetDevice(rank % ndevices);
    nvshmemx_init_attr(NVSHMEMX_INIT_WITH_MPI_COMM, &attr);

    ...

    nvshmem_finalize();
    MPI_Finalize();
    return 0;
}

As shown in this example, the MPI (or OpenSHMEM) library should be initialized first. After MPI has been initialized, the MPI rank can be queried and used to set the CUDA device. An nvshmemx_init_attr_t structure is created and the mpi_comm field is assigned a reference to an MPI communicator handle. The nvshmemx_init_attr operation is used instead of nvshmem_init to enable MPI compatibility mode. For additional details, see .

NVSHMEM provides get and put APIs, which copy data from and to symmetric objects, respectively. Bulk transfer, scalar transfer, and interleaved versions of these APIs are provided. In addition, Atomic Memory Operations (AMOs) are also provided and can be used to perform atomic updates to symmetric variables. Using these APIs, NVSHMEM provides fine-grained and low-overhead to data stored in the partitioned global address space (PGAS) from inside CUDA kernels. By performing communication from within kernels, NVSHMEM also enables applications to benefit from the intrinsic latency-hiding capabilities of the GPU warp scheduling hardware.

In addition to put, get, and AMO library routines, applications can also use the nvshmem_ptr routine to query a direct pointer to data located in partitions of the PGAS located on other PEs. When the given memory at the specified PE is directly accessible, this function returns a valid pointer. Otherwise, it returns a null pointer. This allows applications to issue direct loads and stores to global memory. NVSHMEM APIs (and loads/stores when allowed by the hardware) can be used to access both local and remote data, allowing a single code path to handle both local and remote data.

NVSHMEM is a stateful library. It detects which GPU a PE is using when the PE calls into an NVSHMEM initialization routine. This information is stored inside the NVSHMEM runtime. All symmetric allocation calls made by the PE return device memory of the selected GPU. All NVSHMEM calls made by the PE are assumed to be made with respect to the selected GPU or from inside kernels launched on this GPU. This requires certain restrictions on PE-to-GPU mappings in applications using NVSHMEM.

NVSHMEM relies on data coalescing features in GPU hardware to achieve efficiency over the network when the data access API is used. It is important that application developers follow CUDA programming best practices that promote data coalescing when using fine-grained communication APIs in NVSHMEM.

CUDA kernels that do not use the NVSHMEM synchronization or collective APIs but use other NVSHMEM communication APIs can be launched with either the normal CUDA launch interfaces or the collective launch API. These kernels can still use other NVSHMEM device side APIs such as the one-sided data movement API.

NVSHMEM follows the OpenSHMEM memory model; however several important exceptions are made to adapt OpenSHMEM to the weakly consistent memory model provided by the GPU architecture, as noted in Modifications Of The OpenSHMEM Memory Model.

Two operations, either issued by the same PE or different PEs, that access the same memory location in parallel are in conflict when one or more of the operations performs a write. Such conflicts result in undefined behavior in the OpenSHMEM memory model. An exception is made when the operations are a combination of AMOs or AMOs and wait/test operations. A second exception is made when the operations are a combination of signal updates and wait/test operations.

Updates to globally accessible objects are unordered. A PE can enforce ordering of its updates with respect to accesses performed by the target PE using the nvshmem_fence operation. When updates performed by a PE must be ordered or made visible to PEs other than the target PE, the nvshmem_quiet operation should be used. While updates are unordered, updates made using NVSHMEM APIs are guaranteed to eventually complete without any additional actions performed by either the source or the target PE. As a result, NVSHMEM guarantees that updates will eventually become visible to other PEs through the NVSHMEM API. Updates are also stable in the sense once the update is visible to another API call, the update remains until replaced by another update. This guarantees that synchronization as described above completes in a finite amount of time.

The following OpenSHMEM 1.3 APIs are not supported in NVSHMEM:

• OpenSHMEM Fortran API
• shmem_global_exit
• shmem_pe_accessible
• shmem_addr_accessible
• shmem_realloc
• shmem_fcollect
• shmem_alltoalls
• shmem_lock
• _SHMEM_* constants (deprecated)
• CUDA supports only a subset of the atomic operations included in the OpenSHMEM specification. The long long, int64_t, ptrdiff_t types are currently unsupported for add, fetch_add, inc, and fetch_inc atomic operations in NVSHMEM.

The following OpenSHMEM 1.3 APIs are not supported over InfiniBand in NVSHMEM:

• shmem_iput
• shmem_iget
• shmem_atomic_<all operations>

For a complete list of API descriptions, refer to the OpenSHMEM 1.4 specification document, available at openshmem.org.

For a complete list of API descriptions, refer to the NVSHMEM API documentation.

NVSHMEM extends the remote memory access (get and put), memory ordering, collective communication, and collective synchronization APIs with support for CUDA streams. Each steam-based function performs the same operation as described by the OpenSHMEM specification. An additional argument of type cudaStream_t is added as the last argument to each function and indicates the stream on which the operation is enqueued.

Ordering APIs (fence, quiet and barrier) issued on the CPU and the GPU only order communication operations that were issued from the CPU and the GPU, respectively. To ensure completion of GPU-side operations from the CPU, the developer must ensure completion of the CUDA kernel from which the GPU-side operations were issued, using operations like cudaStreamSynchronize or cudaDeviceSynchronize. Stream-based quiet or barrier operations have the effect of a barrier being executed on the GPU in stream order, ensuring completion and ordering of only GPU-side operations.

The above table lists the NVSHMEM extension APIs that can be invoked by GPU threads. Each of the APIs corresponding to the entries in the description column have two variants each – one with the suffix _block and the other with the suffix _warp. For example, the OpenSHMEM API shmem_float_put has two extension APIs in NVSHMEM: nvshmemx_float_put_block and nvshmemx_float_put_warp.

These extension APIs are collective calls that must be called by every thread in the scope of the API and with exactly the same arguments. The scope of the *_block extension APIs is the block in which the thread resides. Similarly, the scope of the *_warp extension API is the warp in which the thread resides. For example, if thread 0 calls nvshmem_float_put_block, then every other thread that is in the same block as thread 0 must also call nvshmem_float_put_block with exactly the same arguments. Otherwise, the call will result in erroneous behavior or deadlock in the program. The NVSHMEM runtime may or may not leverage the multiple threads in the scope of the API to execute the API call.

The following code shows an MPI version of the simple shift program that was explained in The NVSHMEM Programming Model. It shows the use of the NVSHMEM attribute-based initialization API where the MPI communicator can be used to set up NVSHMEM.

#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;

    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);
    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;
}

The following code shows an example implementation of a single ring-based reduction where multiple iterations of the code, including computation, communication and synchronization are expressed as a single kernel.

This example also demonstrates the use of NVSHMEM collective launch, required when the NVSHMEM synchronization API is used from inside the CUDA kernel.

There is no MPI dependency for the example. NVSHMEM can be used to port existing MPI applications and develop new applications.

#include <stdio.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)
 
#define NVSHMEM_CHECK(stmt)                               \
do {                                                    \
    int result = (stmt);                                \
    if (NVSHMEM_SUCCESS != result) {                      \
        fprintf(stderr, "[%s:%d] nvshmem failed with error %d \n",\
         __FILE__, __LINE__, result);                   \
        exit(-1);                                       \
    }                                                   \
} while (0)
 
 
__global__ void reduce_ring (int *target, int mype, int npes) {
  int peer = (mype + 1)%npes;
  int lvalue = mype;
 
  for (int i=1; i<npes; i++) {
      nvshmem_int_p(target, lvalue, peer);
      nvshmem_barrier_all();
      lvalue = *target + mype;
  }
}
 
int main (int c, char *v[])
{
  int mype, npes;
  nvshmem_init();
  mype = nvshmem_my_pe();
  npes = nvshmem_n_pes();
 
  //application picks the device each PE will use
  CUDA_CHECK(cudaSetDevice(mype));
  double *u = (double *) nvshmem_malloc(sizeof(double));
 
  void *args[] = {&u, &mype, &npes};
  dim3 dimBlock(1);
  dim3 dimGrid(1);
 
  NVSHMEM_CHECK(nvshmemx_collective_launch ((const void *)reduce_ring, dimGrid, dimBlock, args, 0 , 0));
  CUDA_CHECK(cudaDeviceSynchronize());
 
  printf("[%d of %d] run complete \n", mype, npes);
 
  nvshmem_free(u);
 
  nvshmem_finalize();
  return 0;
}

The following example shows how nvshmemx_*_on_stream functions can be used to enqueue a SHMEM operation onto a CUDA stream for execution in stream order. Specifically, the example shows the following:

• How a collective SHMEM reduction operation can be made to wait on a preceding kernel in the stream.
• How a kernel can be made to wait for a communication result from a previous collective SHMEM reduction operation.

The example shows one use case for relieving CPU control over GPU compute and communication.

#include <stdio.h>
#include "nvshmem.h"
#include "nvshmemx.h"

#define THRESHOLD 42
#define CORRECTION 7

#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 accumulate(int *input, int *partial_sum)
{
  int index = threadIdx.x;
  if (0 == index) *partial_sum = 0;
  __syncthreads();
  atomicAdd(partial_sum, input[index]);
}

__global__ void correct_accumulate(int *input, int *partial_sum, int *full_sum)
{
  int index = threadIdx.x;
  if (*full_sum > THRESHOLD) {
    input[index] = input[index] - CORRECTION;
  }
  if (0 == index) *partial_sum = 0;
  __syncthreads();
  atomicAdd(partial_sum, input[index]);
}

int main (int c, char *v[])
{
  int mype, npes;
  int *input;
  int *partial_sum;
  int *full_sum;
  int input_nelems = 512;
  int to_all_nelems = 1;
  int PE_start = 0;
  int PE_size = 0;
  int logPE_stride = 0;
  long *pSync;
  int *pWrk;
  cudaStream_t stream;

  nvshmem_init ();
  PE_size = nvshmem_n_pes();
  mype = nvshmem_my_pe();
  npes = nvshmem_n_pes();

  CUDA_CHECK(cudaSetDevice(mype));
  CUDA_CHECK(cudaStreamCreate(&stream));

  input = (int *) nvshmem_malloc(sizeof(int) * input_nelems);
  partial_sum = (int *) nvshmem_malloc(sizeof(int));
  full_sum = (int *) nvshmem_malloc(sizeof(int));
  pWrk = (int *) nvshmem_malloc(sizeof(int) * NVSHMEM_REDUCE_MIN_WRKDATA_SIZE);
  pSync = (long *) nvshmem_malloc(sizeof(long) * NVSHMEM_REDUCE_SYNC_SIZE);

  accumulate <<<1, input_nelems, 0, stream>>> (input, partial_sum);
  nvshmemx_int_sum_to_all_on_stream(full_sum, partial_sum, to_all_nelems, PE_start, logPE_stride, PE_size, pWrk, pSync, stream);
  correct_accumulate <<<1, input_nelems, 0, stream>>> (input, partial_sum, full_sum);
  CUDA_CHECK(cudaStreamSynchronize(stream));

  printf("[%d of %d] run complete \n", mype, npes);

  CUDA_CHECK(cudaStreamDestroy(stream));

  nvshmem_free(input);
  nvshmem_free(partial_sum);
  nvshmem_free(full_sum);
  nvshmem_free(pWrk);
  nvshmem_free(pSync);

  nvshmem_finalize();
  return 0;
}

The example in this section shows how nvshmemx_collect32_block can be used to leverage threads to accelerate a SHMEM collect operation when all threads in the block depend on the result of a preceding communication operation. For this instance, partial vector sums are computed across different PEs and have a SHMEM collect operation to obtain the complete sum across PEs.

#include <stdio.h>
#include "nvshmem.h"
#include "nvshmemx.h"

#define NTHREADS 512

#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 distributed_vector_sum(int *x, int *y, int *partial_sum, int *sum, long *pSync, int use_threadgroup, int mype, int npes)
{
  int index = threadIdx.x;
  int nelems = blockDim.x;
  int PE_start = 0;
  int logPE_stride = 0;
  partial_sum[index] = x[index] + y[index];

  if (use_threadgroup) {
    /* all threads realize the entire collect operation */
    nvshmemx_collect32_block(sum, partial_sum, nelems, PE_start, logPE_stride, npes, pSync);
  } else {
    /* thread 0 realizes the entire collect operation */
    if (0 == index) {
      nvshmem_collect32(sum, partial_sum, nelems, PE_start, logPE_stride, npes, pSync);
    }
  }
}

int main (int c, char *v[])
{
  int mype, npes;
  int *x;
  int *y;
  int *partial_sum;
  int *sum;
  int use_threadgroup = 1;
  long *pSync;
  int nthreads = NTHREADS;

  nvshmem_init ();
  npes = nvshmem_n_pes();
  mype = nvshmem_my_pe();

  CUDA_CHECK(cudaSetDevice(mype));

  x = (int *) nvshmem_malloc(sizeof(int) * nthreads);
  y = (int *) nvshmem_malloc(sizeof(int) * nthreads);
  partial_sum = (int *) nvshmem_malloc(sizeof(int) * nthreads);
  sum = (int *) nvshmem_malloc(sizeof(int) * nthreads * npes);
  pSync = (long *) nvshmem_malloc(sizeof(long) * NVSHMEM_COLLECT_SYNC_SIZE);

  void *args[] = {&x, &y, &partial_sum, &sum, &pSync, &use_threadgroup, &mype, &npes};
  dim3 dimBlock(nthreads);
  dim3 dimGrid(1);
  nvshmemx_collective_launch ((const void *)distributed_vector_sum, dimGrid, dimBlock, args, 0, 0);
  CUDA_CHECK(cudaDeviceSynchronize());

  printf("[%d of %d] run complete \n", mype, npes);

  nvshmem_free(x);
  nvshmem_free(y);
  nvshmem_free(partial_sum);
  nvshmem_free(sum);
  nvshmem_free(pSync);

  nvshmem_finalize();

  return 0;
}

In the example below, every thread in block 0 calls nvshmemx_float_put_block. Alternatively, every thread can call nvshmem_float_p, but nvshmem_float_p has a disadvantage that when the destination GPU is connected via InfiniBand, there is one RMA message for every single element, which can be detrimental to performance.

The disadvantage with using nvshmem_float_put in this case is that when the destination GPU is P2P-connected, a single thread will copy the entire data to the destination GPU. While nvshmemx_float_put_block can leverage all the threads in the block to copy the data in parallel to the destination GPU.

#include <stdio.h>
#include <assert.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)

#define THREADS_PER_BLOCK 1024

__global__ void set_and_shift_kernel (float *send_data, float *recv_data, int num_elems, int mype, int npes) {
  int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;

  /* set the corresponding element of send_data */
  if (thread_idx < num_elems)
    send_data[thread_idx] = mype;

  int peer = (mype + 1) % npes;

  int block_offset = blockIdx.x * blockDim.x;
  /* All threads in the block call API with same arguments */
    nvshmemx_float_put_block(recv_data + block_offset, send_data + block_offset, min(blockDim.x, num_elems, - block_offset), peer); 
}

int main (int c, char *v[])
{
  int mype, npes;
  float *send_data, *recv_data;
  int num_elems = 8192;
  int num_blocks;

  nvshmem_init ();
  mype = nvshmem_my_pe();
  npes = nvshmem_n_pes();

  //application picks the device each PE will use
  CUDA_CHECK(cudaSetDevice(mype));
  send_data = (float *) nvshmem_malloc(sizeof(float) * num_elems);
  recv_data = (float *) nvshmem_malloc(sizeof(float) * num_elems);

  assert(num_elems % THREADS_PER_BLOCK == 0); /* for simplicity */
  num_blocks = num_elems / THREADS_PER_BLOCK;

  set_and_shift_kernel<<<num_blocks, THREADS_PER_BLOCK>>> 
                     (send_data, recv_data, num_elems, mype, npes);
  CUDA_CHECK(cudaDeviceSynchronize());

  printf("[%d of %d] run complete \n", mype, npes);

  nvshmem_free(send_data);
  nvshmem_free(recv_data);

  nvshmem_finalize();
  return 0; 
}