Do not use async stream for memory allocation and transport setup without graph (#460)

[ROCm/rccl commit: 33bdd557c8]
This commit is contained in:
Wenkai Du
2021-11-05 13:39:14 -07:00
committed by GitHub
parent fd98ee84b4
commit 680660fcf7
2 changed files with 18 additions and 1 deletions
+6
View File
@@ -56,16 +56,22 @@ extern struct allocationTracker allocTracker[];
template <typename T>
static ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, bool isFineGrain = false) {
#if CUDART_VERSION >= 11030
// Need async stream for P2P pre-connect + CUDA Graph
hipStream_t stream;
CUDACHECK(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
#endif
if (isFineGrain)
CUDACHECK(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained));
else
CUDACHECK(hipMalloc(ptr, nelem*sizeof(T)));
#if CUDART_VERSION >= 11030
CUDACHECK(hipMemsetAsync(*ptr, 0, nelem*sizeof(T), stream));
CUDACHECK(hipStreamSynchronize(stream));
CUDACHECK(hipStreamDestroy(stream));
#else
CUDACHECK(hipMemset(*ptr, 0, nelem*sizeof(T)));
#endif
INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
int dev;
CUDACHECK(hipGetDevice(&dev));