split get_places function, add is_custom_device in op_test (#74363)

* rebase latest

Author: datutu-L

test/legacy_test/op_test.py

@@ -389,29 +389,38 @@ def convert_uint16_to_float(in_list):
     return np.reshape(out, in_list.shape)
 
 
-def get_places(string_format=False):
+def get_places():
     places = []
-    if not string_format:
-        if (
-            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
-            in ['1', 'true', 'on']
-            or not core.is_compiled_with_cuda()
-        ):
-            places.append(base.CPUPlace())
-        if core.is_compiled_with_cuda():
-            places.append(base.CUDAPlace(0))
-    else:
-        if (
-            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
-            in ['1', 'true', 'on']
-            or not paddle.is_compiled_with_cuda()
-        ):
-            places.append('cpu')
-        if paddle.is_compiled_with_cuda():
-            places.append('gpu')
+    if (
+        os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+        in ['1', 'true', 'on']
+        or not core.is_compiled_with_cuda()
+    ):
+        places.append(base.CPUPlace())
+    if core.is_compiled_with_cuda():
+        places.append(base.CUDAPlace(0))
+    if is_custom_device():
+        dev_type = paddle.device.get_all_custom_device_type()[0]
+        places.append(base.CustomPlace(dev_type, 0))
     return places
 
 
+def get_devices():
+    devices = []
+    if (
+        os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+        in ['1', 'true', 'on']
+        or not paddle.is_compiled_with_cuda()
+    ):
+        devices.append('cpu')
+    if paddle.is_compiled_with_cuda():
+        devices.append('gpu')
+    if is_custom_device():
+        dev_type = paddle.device.get_all_custom_device_type()[0]
+        devices.append(f'{dev_type}:0')
+    return devices
+
+
 def get_device_place():
     if core.is_compiled_with_cuda():
         return base.CUDAPlace(0)
@@ -423,6 +432,15 @@ def get_device_place():
     return base.CPUPlace()
 
 
+def is_custom_device():
+    custom_dev_types = paddle.device.get_all_custom_device_type()
+    if custom_dev_types and paddle.device.is_compiled_with_custom_device(
+        custom_dev_types[0]
+    ):
+        return True
+    return False
+
+
 @contextmanager
 def auto_parallel_test_guard(test_info_path, generated_test_file_path):
     test_info_file, generated_test_file = None, None
@@ -2902,6 +2920,13 @@ def _get_places(self):
                     return [place]
                 else:
                     return []
+            elif is_custom_device():
+                dev_type = paddle.device.get_all_custom_device_type()[0]
+                place = core.CustomPlace(dev_type, 0)
+                if core.is_float16_supported(place):
+                    return [place]
+                else:
+                    return []
             else:
                 return []
         places = []
@@ -2931,6 +2956,9 @@ def _get_places(self):
             and not cpu_only
         ):
             places.append(core.CUDAPlace(0))
+        if is_custom_device():
+            dev_type = paddle.device.get_all_custom_device_type()[0]
+            places.append(core.CustomPlace(dev_type, 0))
         return places
 
     def check_output(

test/legacy_test/test_activation_op.py

@@ -23,6 +23,7 @@
     convert_float_to_uint16,
     get_device_place,
     get_places,
+    is_custom_device,
 )
 from scipy.special import erf, expit
 from utils import static_guard
@@ -497,7 +498,8 @@ def init_shape(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda() or core.is_compiled_with_rocm(),
+    not (core.is_compiled_with_cuda() or is_custom_device())
+    or core.is_compiled_with_rocm(),
     "core is not compiled with CUDA",
 )
 class TestSigmoidBF16(OpTest):
@@ -1765,7 +1767,8 @@ def init_dtype(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda() or core.is_compiled_with_rocm(),
+    not (core.is_compiled_with_cuda() or is_custom_device())
+    or core.is_compiled_with_rocm(),
     "core is not compiled with CUDA",
 )
 class TestSqrtBF16(OpTest):
@@ -2037,7 +2040,7 @@ def setUp(self):
         self.inputs = {'X': OpTest.np_dtype_to_base_dtype(x)}
         self.outputs = {'Out': out}
         self.convert_input_output()
-        if not core.is_compiled_with_cuda():
+        if not (core.is_compiled_with_cuda() or is_custom_device()):
             self.__class__.no_need_check_grad = True
 
     def init_shape(self):
@@ -2091,7 +2094,7 @@ def setUp(self):
         self.inputs = {'X': OpTest.np_dtype_to_base_dtype(x)}
         self.outputs = {'Out': out}
         self.convert_input_output()
-        if not core.is_compiled_with_cuda():
+        if not (core.is_compiled_with_cuda() or is_custom_device()):
             self.__class__.no_need_check_grad = True
 
     def init_shape(self):
@@ -4563,7 +4566,8 @@ def init_shape(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda() or core.is_compiled_with_rocm(),
+    not (core.is_compiled_with_cuda() or is_custom_device())
+    or core.is_compiled_with_rocm(),
     "core is not compiled with CUDA",
 )
 class TestSquareBF16(OpTest):
@@ -4917,7 +4921,8 @@ def init_shape(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda() or core.is_compiled_with_rocm(),
+    not (core.is_compiled_with_cuda() or is_custom_device())
+    or core.is_compiled_with_rocm(),
     "core is not compiled with CUDA",
 )
 class TestSoftplusBF16(OpTest):
@@ -5595,7 +5600,8 @@ def test_errors(self):
 # ------------------ Test Cudnn Activation----------------------
 def create_test_act_cudnn_class(parent, atol=1e-3, grad_atol=1e-3):
     @unittest.skipIf(
-        not core.is_compiled_with_cuda(), "core is not compiled with CUDA"
+        not (core.is_compiled_with_cuda() or is_custom_device()),
+        "core is not compiled with CUDA",
     )
     class TestActCudnn(parent):
         def init_kernel_type(self):

test/legacy_test/test_adadelta_op.py

@@ -15,7 +15,7 @@
 import unittest
 
 import numpy as np
-from op_test import OpTest, get_device_place, get_places
+from op_test import OpTest, get_device_place, get_devices
 
 import paddle
 from paddle import base
@@ -294,7 +294,7 @@ def _test_adadelta_op_dygraph_place_amp(self, place, use_amp=False):
         paddle.enable_static()
 
     def test_main(self):
-        for place in get_places(string_format=True):
+        for place in get_devices():
             use_amp_list = [True, False]
             for use_amp in use_amp_list:
                 self._test_adadelta_op_dygraph_place_amp(place, use_amp)

test/legacy_test/test_adagrad_op.py

@@ -17,7 +17,7 @@
 
 import numpy as np
 from op import Operator
-from op_test import OpTest, get_device_place, get_places
+from op_test import OpTest, get_device_place, get_devices, get_places
 
 import paddle
 from paddle.base import core
@@ -242,7 +242,7 @@ def _test_adagrad_op_dygraph_place_amp(self, place, use_amp=False):
         paddle.enable_static()
 
     def test_main(self):
-        for place in get_places(string_format=True):
+        for place in get_devices():
             use_amp_list = [True, False]
             for use_amp in use_amp_list:
                 self._test_adagrad_op_dygraph_place_amp(place, use_amp)

test/legacy_test/test_adam_op.py

@@ -16,7 +16,7 @@
 
 import numpy as np
 from op import Operator
-from op_test import OpTest, get_places
+from op_test import OpTest, get_devices, get_places
 
 import paddle
 from paddle import base
@@ -1296,7 +1296,7 @@ def _adam_optimize_static(
         return out
 
     def _get_places(self):
-        return get_places(string_format=True)
+        return get_devices()
 
     def _check_with_place_amp(self, place, use_amp):
         # test dygraph mode

test/legacy_test/test_adamax_op.py

@@ -15,7 +15,7 @@
 import unittest
 
 import numpy as np
-from op_test import OpTest, get_device_place, get_places
+from op_test import OpTest, get_device_place, get_devices
 
 import paddle
 
@@ -275,7 +275,7 @@ def _test_adamax_op_dygraph_place_amp(self, place, use_amp=False):
         paddle.enable_static()
 
     def _get_places(self):
-        return get_places(string_format=True)
+        return get_devices()
 
     def test_main(self):
         for place in self._get_places():

test/legacy_test/test_adamw_op.py

@@ -18,7 +18,7 @@
 from functools import partial
 
 import numpy as np
-from op_test import OpTest, get_places
+from op_test import OpTest, get_devices
 
 import paddle
 from paddle import base, nn
@@ -758,7 +758,7 @@ def _test_adamw_op_dygraph_place_amp(self, place, use_amp=False):
                 optimizer.clear_grad()
 
     def _get_places(self):
-        places = get_places(string_format=True)
+        places = get_devices()
         if paddle.is_compiled_with_xpu():
             places.append('xpu')
         return places

test/legacy_test/test_adaptive_log_softmax_with_loss.py

@@ -15,7 +15,7 @@
 import unittest
 
 import numpy as np
-from op_test import get_places
+from op_test import get_devices, get_places
 
 import paddle
 import paddle.optimizer as optim
@@ -58,7 +58,7 @@ def predict(self, input):
 class TestNNAdaptiveLogSoftmaxWithLossAPI(unittest.TestCase):
     def setUp(self):
         paddle.seed(2024)
-        self.place = get_places(string_format=True)
+        self.place = get_devices()
         self.log_np = np.random.randn(4, 8).astype('float32')
         self.predict_np = np.abs(np.random.randn(64, 8).astype('float32'))
 

test/legacy_test/test_blha_get_max_len_op.py

@@ -15,6 +15,7 @@
 import unittest
 
 import numpy as np
+from op_test import is_custom_device
 
 import paddle
 from paddle.base import core
@@ -109,7 +110,8 @@ def test_static_api(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    not (core.is_compiled_with_cuda() or is_custom_device())
+    and not core.is_compiled_with_xpu(),
     "Only support XPU or GPU in CUDA mode.",
 )
 class TestBlhaGetMaxLenOp_ZeroSize(unittest.TestCase):

test/legacy_test/test_cartesian_prod.py

@@ -16,7 +16,7 @@
 from itertools import product
 
 import numpy as np
-from op_test import get_places
+from op_test import get_devices
 
 import paddle
 from paddle.base import core
@@ -36,7 +36,7 @@ def setUp(self):
         self.c_np = np.random.random(self.c_shape).astype(self.dtype_np)
         self.d_np = np.empty(0, self.dtype_np)
 
-        self.place = get_places(string_format=True)
+        self.place = get_devices()
 
     def init_setting(self):
         self.dtype_np = 'float32'
@@ -119,7 +119,7 @@ def setUp(self):
         self.a_np = np.random.random(self.a_shape).astype(self.dtype_np)
         self.b_np = np.empty(0, self.dtype_np)
 
-        self.place = get_places(string_format=True)
+        self.place = get_devices()
 
     def init_setting(self):
         self.dtype_np = 'float32'