(cherry picked from commit 94da229a7788d74685d1591a4e75a8341de64f41)
Этот коммит содержится в:
Wenkai Du
2023-07-21 07:31:27 -07:00
коммит произвёл GitHub
родитель 59db7b8722
Коммит a7fcd58a97
15 изменённых файлов: 159 добавлений и 56 удалений
+24 -6
Просмотреть файл
@@ -165,9 +165,15 @@ ncclResult_t ncclCudaMallocDebug(const char *filefunc, int line, T** ptr, size_t
cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
*ptr = nullptr;
CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
if (isFineGrain)
if (isFineGrain) {
#if defined(HIP_UNCACHED_MEMORY)
hipDeviceProp_t prop;
CUDACHECK(hipGetDeviceProperties(&prop, 0));
CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), prop.gcnArch/10 == 94 ? hipDeviceMallocUncached : hipDeviceMallocFinegrained), result, finish);
#else
CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
else
#endif
} else
CUDACHECKGOTO(cudaMalloc(ptr, nelem*sizeof(T)), result, finish);
finish:
CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
@@ -187,9 +193,15 @@ ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t
cudaStream_t stream = sideStream;
if (stream == nullptr)
CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
if (isFineGrain)
if (isFineGrain) {
#if defined(HIP_UNCACHED_MEMORY)
hipDeviceProp_t prop;
CUDACHECK(hipGetDeviceProperties(&prop, 0));
CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), prop.gcnArch/10 == 94 ? hipDeviceMallocUncached : hipDeviceMallocFinegrained), result, finish);
#else
CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
else
#endif
} else
CUDACHECKGOTO(cudaMalloc(ptr, nelem*sizeof(T)), result, finish);
CUDACHECKGOTO(cudaMemsetAsync(*ptr, 0, nelem*sizeof(T), stream), result, finish);
CUDACHECKGOTO(cudaStreamSynchronize(stream), result, finish);
@@ -215,9 +227,15 @@ ncclResult_t ncclCudaCallocAsyncDebug(const char *filefunc, int line, T** ptr, s
cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
*ptr = nullptr;
CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
if (isFineGrain)
if (isFineGrain) {
#if defined(HIP_UNCACHED_MEMORY)
hipDeviceProp_t prop;
CUDACHECK(hipGetDeviceProperties(&prop, 0));
CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), prop.gcnArch/10 == 94 ? hipDeviceMallocUncached : hipDeviceMallocFinegrained), result, finish);
#else
CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
else
#endif
} else
CUDACHECKGOTO(cudaMalloc(ptr, nelem*sizeof(T)), result, finish);
CUDACHECKGOTO(cudaMemsetAsync(*ptr, 0, nelem*sizeof(T), stream), result, finish);
int dev;