From 5729353e425fde561c49e34ee1f187611a2dce0b Mon Sep 17 00:00:00 2001 From: ashishfarmer Date: Fri, 13 Nov 2020 00:42:40 +0000 Subject: [PATCH 1/2] allow disabling atomics in rocblas --- aten/src/ATen/cuda/CublasHandlePool.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/aten/src/ATen/cuda/CublasHandlePool.cpp b/aten/src/ATen/cuda/CublasHandlePool.cpp index 0165c53ac60d..b034dd3d3c30 100644 --- a/aten/src/ATen/cuda/CublasHandlePool.cpp +++ b/aten/src/ATen/cuda/CublasHandlePool.cpp @@ -50,6 +50,11 @@ cublasHandle_t getCurrentCUDABlasHandle() { } else { TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH)); } +#endif +#if defined(__HIP_PLATFORM_HCC__) && HIP_VERSION >= 308 + if (at::globalContext().deterministic()) { + TORCH_CUDABLAS_CHECK(rocblas_set_atomics_mode(handle, rocblas_atomics_not_allowed)); + } #endif return handle; } From f2afa4ce5ba145f721752ed75b788cce326750c0 Mon Sep 17 00:00:00 2001 From: ashishfarmer Date: Tue, 1 Dec 2020 23:19:39 +0000 Subject: [PATCH 2/2] more robust deterministic mode checking --- aten/src/ATen/cuda/CublasHandlePool.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/aten/src/ATen/cuda/CublasHandlePool.cpp b/aten/src/ATen/cuda/CublasHandlePool.cpp index b034dd3d3c30..82421f49de1e 100644 --- a/aten/src/ATen/cuda/CublasHandlePool.cpp +++ b/aten/src/ATen/cuda/CublasHandlePool.cpp @@ -52,9 +52,13 @@ cublasHandle_t getCurrentCUDABlasHandle() { } #endif #if defined(__HIP_PLATFORM_HCC__) && HIP_VERSION >= 308 + rocblas_atomics_mode rocblas_mode; if (at::globalContext().deterministic()) { - TORCH_CUDABLAS_CHECK(rocblas_set_atomics_mode(handle, rocblas_atomics_not_allowed)); + rocblas_mode = rocblas_atomics_not_allowed; + } else { + rocblas_mode = rocblas_atomics_allowed; } + TORCH_CUDABLAS_CHECK(rocblas_set_atomics_mode(handle, rocblas_mode)); #endif return handle; }