Wenkai Du
90 рядки
4.7 KiB
C
90 рядки
4.7 KiB
C
/*************************************************************************
|
|
* Copyright (c) 2016-2019, NVIDIA CORPORATION. All rights reserved.
|
|
* Modifications Copyright (c) 2019-2020 Advanced Micro Devices, Inc. All rights reserved.
|
|
*
|
|
* See LICENSE.txt for license information
|
|
************************************************************************/
|
|
|
|
#ifndef NCCL_INT_NET_H_
|
|
#define NCCL_INT_NET_H_
|
|
|
|
#include "nccl.h"
|
|
#include "nccl_net.h"
|
|
|
|
extern ncclNet_t* ncclNet;
|
|
typedef char ncclNetHandle_t[NCCL_NET_HANDLE_MAXSIZE];
|
|
|
|
// Translation to external API
|
|
static const char* ncclNetName() { return ncclNet->name; }
|
|
static ncclResult_t ncclNetDevices(int* ndev) { NCCLCHECK(ncclNet->devices(ndev)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetGetProperties(int dev, ncclNetProperties_t* props) { NCCLCHECK(ncclNet->getProperties(dev, props)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetListen(int dev, void* handle, void** listenComm) { NCCLCHECK(ncclNet->listen(dev, handle, listenComm)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetConnect(int dev, void* handle, void** sendComm) { NCCLCHECK(ncclNet->connect(dev, handle, sendComm)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetAccept(void* listenComm, void** recvComm) { NCCLCHECK(ncclNet->accept(listenComm, recvComm)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetRegMr(void* comm, void* data, int size, int type, void** mhandle) { NCCLCHECK(ncclNet->regMr(comm, data, size, type, mhandle)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetDeregMr(void* comm, void* mhandle) { NCCLCHECK(ncclNet->deregMr(comm, mhandle)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetIsend(void* sendComm, void* data, int size, void* mhandle, void** request) { NCCLCHECK(ncclNet->isend(sendComm, data, size, mhandle, request)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetIrecv(void* recvComm, void* data, int size, void* mhandle, void** request) { NCCLCHECK(ncclNet->irecv(recvComm, data, size, mhandle, request)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetFlush(void* recvComm, void* data, int size, void* mhandle) { NCCLCHECK(ncclNet->flush(recvComm, data, size, mhandle)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetTest(void* request, int* done, int* size) { NCCLCHECK(ncclNet->test(request, done, size)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetCloseSend(void* sendComm) { NCCLCHECK(ncclNet->closeSend(sendComm)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetCloseRecv(void* recvComm) { NCCLCHECK(ncclNet->closeRecv(recvComm)); return ncclSuccess; }
|
|
static ncclResult_t ncclNetCloseListen(void* listenComm) { NCCLCHECK(ncclNet->closeListen(listenComm)); return ncclSuccess; }
|
|
|
|
// Test whether the current GPU support GPU Direct RDMA.
|
|
#define GPU_BUF_SIZE (2*1024*1024)
|
|
static ncclResult_t ncclGpuGdrSupport(int* gdrSupport) {
|
|
int netDevs;
|
|
NCCLCHECK(ncclNetDevices(&netDevs));
|
|
pthread_mutex_t ncclParamMutexGpuGdrSupport = PTHREAD_MUTEX_INITIALIZER;
|
|
static int gdrSupportCached[16] = {-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1};
|
|
int cudaDev;
|
|
CUDACHECK(hipGetDevice(&cudaDev));
|
|
if (gdrSupportCached[cudaDev] != -1) {
|
|
*gdrSupport = gdrSupportCached[cudaDev];
|
|
return ncclSuccess;
|
|
}
|
|
*gdrSupport = 0;
|
|
for (int dev=0; dev<netDevs; dev++) {
|
|
// Find a net device which is GDR-capable
|
|
ncclNetProperties_t props;
|
|
NCCLCHECK(ncclNet->getProperties(dev, &props));
|
|
if ((props.ptrSupport & NCCL_PTR_CUDA) == 0) continue;
|
|
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
|
|
if (!hasFineGrainVramPcie()) continue;
|
|
#endif
|
|
|
|
// Allocate memory on the GPU and try to register it on the NIC.
|
|
void *lComm = NULL, *sComm = NULL, *rComm = NULL;
|
|
ncclNetHandle_t handle;
|
|
void* gpuPtr = NULL;
|
|
void* mHandle = NULL;
|
|
pthread_mutex_lock(&ncclParamMutexGpuGdrSupport);
|
|
NCCLCHECK(ncclNetListen(dev, &handle, &lComm));
|
|
NCCLCHECK(ncclNetConnect(dev, &handle, &sComm));
|
|
NCCLCHECK(ncclNetAccept(lComm, &rComm));
|
|
CUDACHECK(hipExtMallocWithFlags(&gpuPtr, GPU_BUF_SIZE, hipDeviceMallocFinegrained));
|
|
ncclDebugNoWarn = NCCL_NET;
|
|
if (ncclNetRegMr(sComm, gpuPtr, GPU_BUF_SIZE, NCCL_PTR_CUDA, &mHandle) == ncclSuccess) {
|
|
NCCLCHECK(ncclNetDeregMr(sComm, mHandle));
|
|
NCCLCHECK(ncclNetRegMr(rComm, gpuPtr, GPU_BUF_SIZE, NCCL_PTR_CUDA, &mHandle));
|
|
NCCLCHECK(ncclNetDeregMr(rComm, mHandle));
|
|
*gdrSupport = 1;
|
|
}
|
|
ncclDebugNoWarn = 0;
|
|
CUDACHECK(hipFree(gpuPtr));
|
|
NCCLCHECK(ncclNetCloseRecv(rComm));
|
|
NCCLCHECK(ncclNetCloseSend(sComm));
|
|
NCCLCHECK(ncclNetCloseListen(lComm));
|
|
pthread_mutex_unlock(&ncclParamMutexGpuGdrSupport);
|
|
break;
|
|
}
|
|
gdrSupportCached[cudaDev] = *gdrSupport;
|
|
return ncclSuccess;
|
|
}
|
|
|
|
extern ncclNet_t ncclNetIb;
|
|
extern ncclNet_t ncclNetSocket;
|
|
|
|
#endif
|