30 lines
1 KiB
C
30 lines
1 KiB
C
#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 */
|