Unify gpu_support variable in python tests (pytorch#16748)

Summary:
Assign `has_gpu_support = has_cuda_support or has_hip_support` and make according changes in python tests.
Pull Request resolved: pytorch#16748

Differential Revision: D13983132

Pulled By: bddppq

fbshipit-source-id: ca496fd8c6ae3549b736bebd3ace7fa20a6dad7f

Author: rohithkrn

caffe2/contrib/nccl/nccl_ops_test.py

@@ -38,7 +38,7 @@ def benchmark(ws, net, warmups=5, iters=100):
     return after - before
 
 
-@unittest.skipIf(not workspace.has_gpu_support, "NCCL only on GPU")
+@unittest.skipIf(not workspace.has_cuda_support, "NCCL only on CUDA GPU")
 class NCCLOpsTest(hu.HypothesisTestCase):
     @given(n=st.integers(min_value=2, max_value=workspace.NumCudaDevices()),
            m=st.integers(min_value=1, max_value=1000),

caffe2/python/_import_c_extension.py

@@ -10,18 +10,19 @@
 # if that still fails, we will exit loud.
 with extension_loader.DlopenGuard():
     has_hip_support = False
+    has_cuda_support = False
     has_gpu_support = False
 
     try:
         from caffe2.python.caffe2_pybind11_state_gpu import *  # noqa
         if num_cuda_devices():  # noqa
-            has_gpu_support = True
+            has_gpu_support = has_cuda_support = True
     except ImportError as gpu_e:
         logging.info('Failed to import cuda module: {}'.format(gpu_e))
         try:
             from caffe2.python.caffe2_pybind11_state_hip import *  # noqa
             if num_hip_devices():
-                has_hip_support = True
+                has_gpu_support = has_hip_support = True
                 logging.info('This caffe2 python run has AMD GPU support!')
         except ImportError as hip_e:
             logging.info('Failed to import AMD hip module: {}'.format(hip_e))

caffe2/python/core_test.py

@@ -643,8 +643,7 @@ def test_inject_copy(self):
         self.assertEqual(op.input[2], "fc_b")
 
 
-@unittest.skipIf(not workspace.has_gpu_support
-                and not workspace.has_hip_support, 'No GPU support')
+@unittest.skipIf(not workspace.has_gpu_support, 'No GPU support')
 class TestInferDevice(test_util.TestCase):
 
     def setUp(self):

caffe2/python/gradient_check_test.py

@@ -23,7 +23,7 @@
 import unittest
 
 
-if (workspace.has_gpu_support or workspace.has_hip_support) and workspace.NumGpuDevices() > 0:
+if workspace.has_gpu_support and workspace.NumGpuDevices() > 0:
     gpu_device_option = caffe2_pb2.DeviceOption()
     gpu_device_option.device_type = workspace.GpuDeviceType
     cpu_device_option = caffe2_pb2.DeviceOption()

caffe2/python/hypothesis_test.py

@@ -320,7 +320,7 @@ def prod(xs):
         for param, _ in enumerate(inputs):
             self.assertGradientChecks(gc, op, inputs, param, [0])
 
-    @unittest.skipIf(not workspace.has_gpu_support and not workspace.has_hip_support,
+    @unittest.skipIf(not workspace.has_gpu_support,
                      "Skipping test due to no gpu present.")
     @given(hidden_size=st.integers(min_value=1, max_value=3),
            num_layers=st.integers(min_value=1, max_value=3),

caffe2/python/hypothesis_test_util.py

@@ -277,11 +277,12 @@ def tensors1d(n, min_len=1, max_len=64, dtype=np.float32, elements=None):
 
 
 cpu_do = caffe2_pb2.DeviceOption()
-gpu_do = caffe2_pb2.DeviceOption(device_type=caffe2_pb2.CUDA)
+cuda_do = caffe2_pb2.DeviceOption(device_type=caffe2_pb2.CUDA)
 hip_do = caffe2_pb2.DeviceOption(device_type=caffe2_pb2.HIP)
+gpu_do =  caffe2_pb2.DeviceOption(device_type=workspace.GpuDeviceType) # CUDA or ROCm
 # (bddppq) Do not rely on this no_hip option! It's just used to
 # temporarily skip some flaky tests on ROCM before it's getting more mature.
-_device_options_no_hip = [cpu_do] + ([gpu_do] if workspace.has_gpu_support else [])
+_device_options_no_hip = [cpu_do] + ([cuda_do] if workspace.has_cuda_support else [])
 device_options = _device_options_no_hip + ([hip_do] if workspace.has_hip_support else [])
 
 # Include device option for each GPU
@@ -304,7 +305,8 @@ def gradient_checker_device_option():
 )
 
 gcs_cpu_only = dict(gc=st.sampled_from([cpu_do]), dc=st.just([cpu_do]))
-gcs_gpu_only = dict(gc=st.sampled_from([gpu_do]), dc=st.just([gpu_do]))
+gcs_cuda_only = dict(gc=st.sampled_from([cuda_do]), dc=st.just([cuda_do]))
+gcs_gpu_only = dict(gc=st.sampled_from([gpu_do]), dc=st.just([gpu_do])) # CUDA or ROCm
 gcs_no_hip = dict(gc=st.sampled_from(_device_options_no_hip), dc=st.just(_device_options_no_hip))
 
 

caffe2/python/memonger_test.py

@@ -223,8 +223,7 @@ def test_gradient_optim(self, input_dim, output_dim, batch_size):
         np.testing.assert_almost_equal(loss, optimized_loss)
         np.testing.assert_almost_equal(grad, optimized_grad)
 
-    @unittest.skipIf(not workspace.has_gpu_support
-                    and not workspace.has_hip_support, "No gpu support.")
+    @unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
     def test_memonger_mix_cpu_gpu(self):
         '''
         Check that memonger does not make blobs cross CPU/GPU boundary

caffe2/python/model_device_test.py

@@ -136,8 +136,7 @@ def _testMiniAlexNet(self, order):
         )
         self.assertEqual(ret, True)
 
-    @unittest.skipIf(not workspace.has_gpu_support
-                    and not workspace.has_hip_support,
+    @unittest.skipIf(not workspace.has_gpu_support,
                      "No GPU support. Skipping test.")
     def testMiniAlexNetNCHW(self):
         self._testMiniAlexNet("NCHW")

caffe2/python/onnx/backend.py

@@ -919,7 +919,7 @@ def supports_device(cls, device_str):
         if device.type == DeviceType.CPU:
             return True
         elif core.IsGPUDeviceType(device.type):
-            return workspace.has_gpu_support or workspace.has_hip_support
+            return workspace.has_gpu_support
         return False
 
     @classmethod

caffe2/python/operator_test/activation_ops_test.py

@@ -41,8 +41,7 @@ def relu_ref(X):
         self.assertDeviceChecks(dc, op, [X], [0])
         self.assertGradientChecks(gc, op, [X], 0, [0])
 
-    @unittest.skipIf(not workspace.has_gpu_support and
-                    not workspace.has_hip_support,
+    @unittest.skipIf(not workspace.has_gpu_support,
                      "Relu for float16 can only run on GPU now.")
     @given(X=hu.tensor(dtype=np.float16), in_place=st.booleans(),
            engine=st.sampled_from(["", "CUDNN"]), **hu.gcs)

caffe2/python/operator_test/conv_test.py

@@ -748,7 +748,7 @@ def canonical(o):
             ["simple", "dag"]
             + (
                 ["async_dag"]
-                if workspace.has_gpu_support or workspace.has_hip_support
+                if workspace.has_gpu_support
                 else []
             )
         ),

caffe2/python/operator_test/cudnn_recurrent_test.py

@@ -11,7 +11,7 @@
 import unittest
 
 
-@unittest.skipIf((not workspace.has_gpu_support) and (not workspace.has_hip_support), "No gpu support.")
+@unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
 class TestLSTMs(unittest.TestCase):
 
     def testEqualToCudnn(self):

caffe2/python/operator_test/elementwise_op_broadcast_test.py

@@ -327,7 +327,7 @@ def test_sum_reduce(self, gc, dc):
         dc_cpu_only = [d for d in dc if d.device_type != caffe2_pb2.CUDA]
         self.assertDeviceChecks(dc_cpu_only, op, [X, Y], [0])
 
-    @unittest.skipIf(not workspace.has_gpu_support and not workspace.has_hip_support, "No gpu support")
+    @unittest.skipIf(not workspace.has_gpu_support, "No gpu support")
     @given(**hu.gcs)
     def test_sum_reduce_fp16(self, gc, dc):
         assume(core.IsGPUDeviceType(gc.device_type))

caffe2/python/operator_test/momentum_sgd_test.py

@@ -139,7 +139,7 @@ def sparse(grad, m, lr, param, i):
             [grad, m, lr, w, indices],
             sparse)
 
-    @unittest.skipIf(not workspace.has_gpu_support and not workspace.has_hip_support, "No gpu support.")
+    @unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
     @given(n=st.integers(4, 8), nesterov=st.booleans(), **hu.gcs)
     def test_fp16momentum_sgd(self, n, nesterov, gc, dc):
         assume(core.IsGPUDeviceType(gc.device_type))

caffe2/python/operator_test/pooling_test.py

@@ -63,7 +63,7 @@ def test_pooling_separate_stride_pad(self, stride_h, stride_w,
     @unittest.skipIf(not os.getenv('CAFFE2_DEBUG'),
                      "This is a test that reproduces a cudnn error. If you "
                      "want to run it, set env variable CAFFE2_DEBUG=1.")
-    @given(**hu.gcs_gpu_only)
+    @given(**hu.gcs_cuda_only)
     def test_pooling_big_batch(self, gc, dc):
         op = core.CreateOperator(
             "AveragePool",

caffe2/python/operator_test/softmax_ops_test.py

@@ -626,7 +626,7 @@ def test_compare_cpugpu(self):
                 "SpatialSoftmaxWithLoss",
                 ["X_gpu", "label_gpu"],
                 ["probs_gpu", "avgloss_gpu"],
-                device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
+                device_option=core.DeviceOption(workspace.GpuDeviceType, 0)
             )
 
             cpuop = core.CreateOperator(
@@ -649,7 +649,7 @@ def test_compare_cpugpu(self):
             # Initialize label. Some of the labels are (-1), i.e "DONT CARE"
             label = (np.random.rand(n, H, W) * (D + 1)).astype(np.int32) - 1
 
-            gpu0 = core.DeviceOption(caffe2_pb2.CUDA, 0)
+            gpu0 = core.DeviceOption(workspace.GpuDeviceType, 0)
             workspace.FeedBlob("X_cpu", X)
             workspace.FeedBlob("label_cpu", label)
             workspace.FeedBlob("X_gpu", X, device_option=gpu0)

caffe2/python/operator_test/spatial_bn_op_test.py

@@ -67,8 +67,7 @@ def reference_spatialbn_test(X, scale, bias, mean, var):
                                    reference_spatialbn_test)
         self.assertDeviceChecks(dc, op, [X, scale, bias, mean, var], [0])
 
-    @unittest.skipIf((not workspace.has_gpu_support) and (
-        not workspace.has_hip_support), "No gpu support")
+    @unittest.skipIf(not workspace.has_gpu_support, "No gpu support")
     @given(size=st.integers(7, 10),
            input_channels=st.integers(1, 10),
            batch_size=st.integers(0, 3),

caffe2/python/operator_test/transpose_op_test.py

@@ -38,9 +38,9 @@ def transpose_ref(X):
         self.assertDeviceChecks(dc, op, [X], [0])
         self.assertGradientChecks(gc, op, [X], 0, [0])
 
-    @unittest.skipIf(not workspace.has_gpu_support, "no gpu support")
+    @unittest.skipIf(not workspace.has_cuda_support, "no cuda support")
     @given(X=hu.tensor(dtype=np.float32), use_axes=st.booleans(),
-           **hu.gcs_gpu_only)
+           **hu.gcs_cuda_only)
     def test_transpose_cudnn(self, X, use_axes, gc, dc):
         ndim = len(X.shape)
         axes = np.arange(ndim)
@@ -49,11 +49,11 @@ def test_transpose_cudnn(self, X, use_axes, gc, dc):
         if (use_axes):
             op = core.CreateOperator(
                 "Transpose", ["X"], ["Y"], axes=axes, engine="CUDNN",
-                device_option=hu.gpu_do)
+                device_option=hu.cuda_do)
         else:
             op = core.CreateOperator(
                 "Transpose", ["X"], ["Y"], engine="CUDNN",
-                device_option=hu.gpu_do)
+                device_option=hu.cuda_do)
 
         def transpose_ref(X):
             if use_axes:

caffe2/python/optimizer_test.py

@@ -454,8 +454,7 @@ def test_caffe2_cpu_vs_numpy(self):
                 )
 
     @unittest.skip("Results might vary too much. Only for individual use.")
-    @unittest.skipIf(not workspace.has_gpu_support
-                    and not workspace.has_hip_support, "No gpu support")
+    @unittest.skipIf(not workspace.has_gpu_support, "No gpu support")
     def test_caffe2_gpu_vs_numpy(self):
         n_dim = 1000000
         n_iter = 50

caffe2/python/parallelize_bmuf_distributed_test.py

@@ -25,7 +25,7 @@ def bmuf_process(filestore_dir, process_id, shared_results,
     dyndep.InitOpsLibrary("@/caffe2/caffe2/distributed:file_store_handler_ops")
 
     if not cpu_device:
-        if not workspace.has_gpu_support and not workspace.has_hip_support:
+        if not workspace.has_gpu_support:
             log.info('No GPU support test is Ignored.')
             return
         if workspace.NumGpuDevices() < 4:

caffe2/python/test/executor_test.py

@@ -47,8 +47,7 @@ def run_model():
         )
 
 
-@unittest.skipIf(not workspace.has_gpu_support
-                and not workspace.has_hip_support, "no gpu")
+@unittest.skipIf(not workspace.has_gpu_support, "no gpu")
 class ExecutorGPUResNetTest(ExecutorTestBase):
     @given(executor=st.sampled_from(EXECUTORS),
            num_workers=st.sampled_from([8]))

caffe2/python/workspace.py

@@ -41,9 +41,10 @@
 operator_tracebacks = defaultdict(dict)
 
 is_asan = C.is_asan
-has_gpu_support = C.has_gpu_support
+has_cuda_support = C.has_cuda_support
 has_hip_support = C.has_hip_support
-if has_gpu_support:
+has_gpu_support = C.has_gpu_support
+if has_cuda_support:
     GpuDeviceType = caffe2_pb2.CUDA
     NumCudaDevices = C.num_cuda_devices
     # This is a duplicate of NumCudaDevices. Remove
@@ -69,7 +70,7 @@ def GetGpuPeerAccessPattern():
         return np.asarray(C.get_hip_peer_access_pattern())
     GetDeviceProperties = C.get_device_properties
 
-if not has_gpu_support and not has_hip_support:
+if not has_gpu_support:
     # setting cuda as the default GpuDeviceType as some tests
     # like core, scope tests use GpuDeviceType even without gpu support
     GpuDeviceType = caffe2_pb2.CUDA

caffe2/python/workspace_test.py

@@ -320,8 +320,7 @@ def testCreateWorkspace(self):
         self.assertTrue("test" in workspaces)
 
 
-@unittest.skipIf(not workspace.has_gpu_support
-                and not workspace.has_hip_support, "No gpu support.")
+@unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
 class TestWorkspaceGPU(test_util.TestCase):
 
     def setUp(self):