#ifndef DISTRIBUTED_H
#define DISTRIBUTED_H
#include "tensor.h"

#define MPI_CHECK(cmd) do {                          \
  int e = cmd;                                      \
  if( e != MPI_SUCCESS ) {                          \
    printf("Failed: MPI error %s:%d '%d'\n",        \
        __FILE__,__LINE__, e);   \
    exit(EXIT_FAILURE);                             \
  }                                                 \
} while(0)

#define NCCL_CHECK(cmd) do {                         \
  ncclResult_t r = cmd;                             \
  if (r!= ncclSuccess) {                            \
    printf("Failed, NCCL error %s:%d '%s'\n",             \
        __FILE__,__LINE__,ncclGetErrorString(r));   \
    exit(EXIT_FAILURE);                             \
  }                                                 \
} while(0)

extern "C" {
  void init_process_group(int rank, int world_size);
  void broadcast_tensor(Tensor* tensor, int src);
  void allreduce_sum_tensor(Tensor* tensor);
  void allreduce_mean_tensor(Tensor* tensor);
}

#endif /* DISTRIBUTED_H */