Unify TEST_CUDNN definition (pytorch#105594)

Signed-off-by: Edward Z. Yang <[email protected]>
Pull Request resolved: pytorch#105594
Approved by: https://github.com/larryliu0820, https://github.com/voznesenskym

Author: ezyang

test/jit/test_freezing.py

@@ -13,6 +13,7 @@
 from torch.testing._internal.common_quantization import skipIfNoFBGEMM
 from torch.testing._internal.common_quantized import override_quantized_engine
 from torch.testing._internal.common_utils import set_default_dtype, skipCUDAMemoryLeakCheckIf, TEST_WITH_ROCM
+from torch.testing._internal.common_cuda import TEST_CUDNN
 from torch.testing._internal.jit_utils import JitTestCase
 from torch.utils import mkldnn as mkldnn_utils
 
@@ -30,10 +31,6 @@
 
 TEST_CUDA = torch.cuda.is_available()
 TEST_ROCM = torch.cuda.is_available() and torch.version.hip is not None
-TEST_CUDNN = False
-if TEST_CUDA and not TEST_ROCM:  # Skip ROCM
-    torch.ones(1).cuda()  # initialize cuda context
-    TEST_CUDNN = TEST_CUDA and torch.backends.cudnn.is_acceptable(torch.tensor(1., device=torch.device('cuda:0')))
 
 def removeExceptions(graph):
     for n in graph.findAllNodes('prim::RaiseException'):

test/test_cpp_extensions_jit.py

@@ -11,6 +11,7 @@
 import glob
 
 import torch.testing._internal.common_utils as common
+from torch.testing._internal.common_cuda import TEST_CUDNN
 import torch
 import torch.backends.cudnn
 import torch.utils.cpp_extension
@@ -20,13 +21,7 @@
 
 
 TEST_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
-TEST_CUDNN = False
 TEST_ROCM = torch.cuda.is_available() and torch.version.hip is not None and ROCM_HOME is not None
-if TEST_CUDA and torch.version.cuda is not None:  # the skip CUDNN test for ROCm
-    CUDNN_HEADER_EXISTS = os.path.isfile(os.path.join(CUDA_HOME, "include/cudnn.h"))
-    TEST_CUDNN = (
-        TEST_CUDA and CUDNN_HEADER_EXISTS and torch.backends.cudnn.is_available()
-    )
 TEST_MPS = torch.backends.mps.is_available()
 IS_WINDOWS = sys.platform == "win32"
 

test/test_cpp_extensions_open_device_registration.py

@@ -15,13 +15,7 @@
 
 
 TEST_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
-TEST_CUDNN = False
 TEST_ROCM = torch.cuda.is_available() and torch.version.hip is not None and ROCM_HOME is not None
-if TEST_CUDA and torch.version.cuda is not None:  # the skip CUDNN test for ROCm
-    CUDNN_HEADER_EXISTS = os.path.isfile(os.path.join(CUDA_HOME, "include/cudnn.h"))
-    TEST_CUDNN = (
-        TEST_CUDA and CUDNN_HEADER_EXISTS and torch.backends.cudnn.is_available()
-    )
 
 
 def remove_build_path():