Wenkai Du
15cd2aff3c
Introducing 3 new APIs:
ncclResult_t ncclGather(const void* sendbuff, void* recvbuff, size_t sendcount,
ncclDataType_t datatype, int root, ncclComm_t comm, hipStream_t stream);
ncclResult_t ncclScatter(const void* sendbuff, void* recvbuff,
size_t recvcount, ncclDataType_t datatype, int root, ncclComm_t comm,
hipStream_t stream);
ncclResult_t ncclAllToAll(const void* sendbuff, void* recvbuff, size_t count,
ncclDataType_t datatype, ncclComm_t comm, hipStream_t stream);
Only out of place operation is supported.
Preprocessor symbol RCCL_GATHER_SCATTER=1 indicates API availibility.
By default the APIs launche RCCL kernel implementation, which can be disabled by
RCCL_ALLTOALL_KERNEL_DISABLE=1. Then the APIs use wrapper around ncclSend and ncclRecv.
[ROCm/rccl commit: e80e29573c]
79 lines
2.0 KiB
C++
79 lines
2.0 KiB
C++
/*************************************************************************
|
|
* Copyright (c) 2016-2020, NVIDIA CORPORATION. All rights reserved.
|
|
*
|
|
* See LICENSE.txt for license information
|
|
************************************************************************/
|
|
|
|
#ifndef NCCL_PROXY_H_
|
|
#define NCCL_PROXY_H_
|
|
|
|
#include <pthread.h>
|
|
|
|
enum ncclProxyOpState { ncclProxyOpNone, ncclProxyOpReady, ncclProxyOpProgress };
|
|
|
|
struct ncclProxyArgs;
|
|
typedef ncclResult_t (*proxyProgressFunc_t)(struct ncclProxyArgs*);
|
|
|
|
struct ncclProxyArgs {
|
|
proxyProgressFunc_t progress;
|
|
struct ncclChannel* channel;
|
|
struct ncclConnector* connector;
|
|
int sliceSteps;
|
|
int chunkSteps;
|
|
int nsteps;
|
|
uint64_t opCount;
|
|
int protocol;
|
|
ncclDataType_t dtype;
|
|
ncclRedOp_t redOp;
|
|
int state; // add component before this line -- it is left out during initialization
|
|
|
|
// Internal state
|
|
uint64_t head;
|
|
uint64_t tail;
|
|
uint64_t end;
|
|
void* requests[NCCL_STEPS];
|
|
int idle;
|
|
|
|
// Element linking
|
|
pthread_mutex_t mutex;
|
|
struct ncclProxyArgs* next;
|
|
struct ncclProxyArgs* nextPeer;
|
|
};
|
|
|
|
struct ncclProxyPool;
|
|
struct ncclProxyState {
|
|
pthread_cond_t cond;
|
|
pthread_mutex_t mutex;
|
|
bool stop;
|
|
struct ncclProxyArgs* ops;
|
|
struct ncclProxyArgs* pool;
|
|
struct ncclProxyPool* pools;
|
|
};
|
|
|
|
typedef ncclResult_t (*threadFunc_t)(struct ncclProxyArgs*);
|
|
|
|
enum proxyMode {
|
|
proxyRing = 0,
|
|
proxyFrom = 1,
|
|
proxyTo = 2
|
|
};
|
|
|
|
ncclResult_t ncclProxySaveColl(struct ncclProxyArgs* args, int pattern, int root, int nranks);
|
|
ncclResult_t ncclProxySaveP2p(struct ncclInfo* info, struct ncclChannel* channel);
|
|
ncclResult_t ncclProxySaveA2a(struct ncclProxyArgs* args, struct ncclInfo* info);
|
|
ncclResult_t ncclProxyStart(struct ncclComm* comm);
|
|
ncclResult_t ncclProxyCreate(struct ncclComm* comm);
|
|
ncclResult_t ncclProxyDestroy(struct ncclComm* comm);
|
|
|
|
#include <unistd.h>
|
|
|
|
// Spin wait until func evaluates to true
|
|
template<typename FUNC>
|
|
inline void transportProxyWait(const FUNC& func) {
|
|
while (!func()) {
|
|
sched_yield();
|
|
}
|
|
}
|
|
|
|
#endif
|