fix pir argminmax2 (#66622)

wanghuancoder · web-flow · commit ac6cbaf6eabe · 2024-07-31T09:31:51.000+08:00
* fix pir argmin argmax

* refine argmin

* refine

* refine

* refine
diff --git a/paddle/phi/kernels/cpu/arg_min_max_kernel.cc b/paddle/phi/kernels/cpu/arg_min_max_kernel.cc
@@ -201,8 +201,6 @@ PD_REGISTER_KERNEL(argmin,
                    CPU,
                    ALL_LAYOUT,
                    phi::ArgMinKernel,
-                   phi::dtype::float16,
-                   phi::dtype::bfloat16,
                    float,
                    double,
                    int32_t,
@@ -216,8 +214,6 @@ PD_REGISTER_KERNEL(argmax,
                    CPU,
                    ALL_LAYOUT,
                    phi::ArgMaxKernel,
-                   phi::dtype::float16,
-                   phi::dtype::bfloat16,
                    float,
                    double,
                    int32_t,
diff --git a/test/deprecated/legacy_test/test_arg_min_max_op.py b/test/deprecated/legacy_test/test_arg_min_max_op.py
@@ -146,6 +146,8 @@ def test_type_error(self):
 
     def test_bfp16(self):
         # in static mode
+        if not paddle.is_compiled_with_cuda():
+            return
         with program_guard(Program(), Program()):
             x = paddle.zeros(name='x', shape=[100, 10], dtype='uint16')
             t1 = paddle.argmin(x)
diff --git a/test/legacy_test/test_arg_min_max_v2_op.py b/test/legacy_test/test_arg_min_max_v2_op.py
@@ -375,11 +375,13 @@ def test_argmin_dtype_type():
 
 class TestArgMaxOpFp16(unittest.TestCase):
     def test_fp16(self):
-        x_np = np.random.random((10, 16)).astype('float16')
-        with paddle.static.program_guard(paddle.static.Program()):
-            x = paddle.static.data(shape=[10, 16], name='x', dtype='float16')
-            out = paddle.argmax(x)
-            if core.is_compiled_with_cuda():
+        if core.is_compiled_with_cuda():
+            x_np = np.random.random((10, 16)).astype('float16')
+            with paddle.static.program_guard(paddle.static.Program()):
+                x = paddle.static.data(
+                    shape=[10, 16], name='x', dtype='float16'
+                )
+                out = paddle.argmax(x)
                 place = paddle.CUDAPlace(0)
                 exe = paddle.static.Executor(place)
                 exe.run(paddle.static.default_startup_program())
@@ -388,11 +390,13 @@ def test_fp16(self):
 
 class TestArgMinOpFp16(unittest.TestCase):
     def test_fp16(self):
-        x_np = np.random.random((10, 16)).astype('float16')
-        with paddle.static.program_guard(paddle.static.Program()):
-            x = paddle.static.data(shape=[10, 16], name='x', dtype='float16')
-            out = paddle.argmin(x)
-            if core.is_compiled_with_cuda():
+        if core.is_compiled_with_cuda():
+            x_np = np.random.random((10, 16)).astype('float16')
+            with paddle.static.program_guard(paddle.static.Program()):
+                x = paddle.static.data(
+                    shape=[10, 16], name='x', dtype='float16'
+                )
+                out = paddle.argmin(x)
                 place = paddle.CUDAPlace(0)
                 exe = paddle.static.Executor(place)
                 exe.run(paddle.static.default_startup_program())
