[XPU] add argmin op and unit test (#68176)

skywalker2012 · web-flow · commit b718c6bdf4f6 · 2024-09-24T10:16:17.000+08:00
* fix

* rename ut file

* fix

* restore previous way

* fix typo

* add return check

* replace with right marco

* cast_v2 to cast

* fix
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -48,6 +48,8 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16})},
+      {"arg_min",
+       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"argsort_grad",
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
diff --git a/paddle/phi/kernels/xpu/arg_min_max_kernel.cc b/paddle/phi/kernels/xpu/arg_min_max_kernel.cc
@@ -22,6 +22,7 @@
 
 namespace phi {
 
+// ArgMax implementation
 template <typename T, typename Context>
 void ArgMaxKernel(const Context& dev_ctx,
                   const DenseTensor& x,
@@ -34,16 +35,15 @@ void ArgMaxKernel(const Context& dev_ctx,
       x.numel(),
       0,
       common::errors::InvalidArgument(
-          "argmin/argmax input numel must > 0, bug got %d", x.numel()));
+          "argmin/argmax input numel must > 0, but got %d", x.numel()));
   using XPUType = typename XPUTypeTrait<T>::Type;
   PADDLE_ENFORCE_EQ(
       (dtype == DataType::UNDEFINED || dtype == DataType::INT32 ||
        dtype == DataType::INT64),
       true,
       errors::InvalidArgument(
           "The attribute of dtype in xpu argmin/argmax must be [%s] or [%s], "
-          "but "
-          "received [%s]",
+          "but received [%s]",
           DataType::INT64,
           DataType::INT32,
           dtype));
@@ -59,64 +59,133 @@ void ArgMaxKernel(const Context& dev_ctx,
     if (axis_val < 0) axis_val += x_dims.size();
   }
   auto xdims_vec = common::vectorize<int>(x_dims);
-  int r = 0;
   if (dtype != DataType::INT32) {
     dev_ctx.template Alloc<int64_t>(out);
     if (x.dims().size() == 0) {
-      xpu::constant(dev_ctx.x_context(),
-                    out->data<int64_t>(),
-                    x.numel(),
-                    static_cast<int64_t>(0));
+      int r = xpu::constant(dev_ctx.x_context(),
+                            out->data<int64_t>(),
+                            x.numel(),
+                            static_cast<int64_t>(0));
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
       return;
     }
-    r = xpu::argmax(dev_ctx.x_context(),
-                    reinterpret_cast<const XPUType*>(x.data<T>()),
-                    out->data<int64_t>(),
-                    xdims_vec,
-                    axis_val);
-    PADDLE_ENFORCE_EQ(
-        r,
-        XPU_SUCCESS,
-        errors::External("XPU argmax kernel return wrong value[%d %s].",
-                         r,
-                         XPUAPIErrorMsg[r]));
+    int r = xpu::argmax(dev_ctx.x_context(),
+                        reinterpret_cast<const XPUType*>(x.data<T>()),
+                        out->data<int64_t>(),
+                        xdims_vec,
+                        axis_val);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "argmax");
   } else {
     DenseTensor out_int64;
     out_int64.Resize(out->dims());
     dev_ctx.template Alloc<int64_t>(&out_int64);
     if (x.dims().size() == 0) {
-      xpu::constant(dev_ctx.x_context(),
-                    out_int64.data<int64_t>(),
-                    x.numel(),
-                    static_cast<int64_t>(0));
+      int r = xpu::constant(dev_ctx.x_context(),
+                            out_int64.data<int64_t>(),
+                            x.numel(),
+                            static_cast<int64_t>(0));
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
     } else {
-      r = xpu::argmax(dev_ctx.x_context(),
-                      reinterpret_cast<const XPUType*>(x.data<T>()),
-                      out_int64.data<int64_t>(),
-                      xdims_vec,
-                      axis_val);
+      int r = xpu::argmax(dev_ctx.x_context(),
+                          reinterpret_cast<const XPUType*>(x.data<T>()),
+                          out_int64.data<int64_t>(),
+                          xdims_vec,
+                          axis_val);
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "argmax");
     }
 
-    PADDLE_ENFORCE_EQ(
-        r,
-        XPU_SUCCESS,
-        errors::External("XPU argmax kernel return wrong value[%d %s].",
-                         r,
-                         XPUAPIErrorMsg[r]));
     dev_ctx.template Alloc<int>(out);
-    r = xpu::cast_v2<int64_t, int>(dev_ctx.x_context(),
-                                   out_int64.data<int64_t>(),
-                                   out->data<int>(),
-                                   out_int64.numel());
-    PADDLE_ENFORCE_EQ(
-        r,
-        XPU_SUCCESS,
-        errors::External("XPU cast kernel return wrong value[%d %s].",
-                         r,
-                         XPUAPIErrorMsg[r]));
+    int r = xpu::cast<int64_t, int>(dev_ctx.x_context(),
+                                    out_int64.data<int64_t>(),
+                                    out->data<int>(),
+                                    out_int64.numel());
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
   }
 }
+
+// ArgMin implementation
+template <typename T, typename Context>
+void ArgMinKernel(const Context& dev_ctx,
+                  const DenseTensor& x,
+                  const Scalar& axis,
+                  bool keepdims,
+                  bool flatten,
+                  DataType dtype,
+                  DenseTensor* out) {
+  PADDLE_ENFORCE_GT(
+      x.numel(),
+      0,
+      common::errors::InvalidArgument(
+          "argmin/argmax input numel must > 0, but got %d", x.numel()));
+  using XPUType = typename XPUTypeTrait<T>::Type;
+  PADDLE_ENFORCE_EQ(
+      (dtype == DataType::UNDEFINED || dtype == DataType::INT32 ||
+       dtype == DataType::INT64),
+      true,
+      errors::InvalidArgument(
+          "The attribute of dtype in xpu argmin/argmax must be [%s] or [%s], "
+          "but received [%s]",
+          DataType::INT64,
+          DataType::INT32,
+          dtype));
+
+  DDim x_dims;
+  int axis_val = axis.to<int>();
+  if (flatten) {
+    x_dims = common::make_ddim({x.numel()});
+    // If flatten, the axis just as 0
+    axis_val = 0;
+  } else {
+    x_dims = x.dims();
+    if (axis_val < 0) axis_val += x_dims.size();
+  }
+  auto xdims_vec = common::vectorize<int>(x_dims);
+  if (dtype != DataType::INT32) {
+    dev_ctx.template Alloc<int64_t>(out);
+    if (x.dims().size() == 0) {
+      int r = xpu::constant(dev_ctx.x_context(),
+                            out->data<int64_t>(),
+                            x.numel(),
+                            static_cast<int64_t>(0));
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
+      return;
+    }
+    int r = xpu::argmin(dev_ctx.x_context(),
+                        reinterpret_cast<const XPUType*>(x.data<T>()),
+                        out->data<int64_t>(),
+                        xdims_vec,
+                        axis_val);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "argmin");
+  } else {
+    DenseTensor out_int64;
+    out_int64.Resize(out->dims());
+    dev_ctx.template Alloc<int64_t>(&out_int64);
+    if (x.dims().size() == 0) {
+      int r = xpu::constant(dev_ctx.x_context(),
+                            out_int64.data<int64_t>(),
+                            x.numel(),
+                            static_cast<int64_t>(0));
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
+    } else {
+      int r = xpu::argmin(dev_ctx.x_context(),
+                          reinterpret_cast<const XPUType*>(x.data<T>()),
+                          out_int64.data<int64_t>(),
+                          xdims_vec,
+                          axis_val);
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "argmin");
+    }
+
+    dev_ctx.template Alloc<int>(out);
+    int r = xpu::cast<int64_t, int>(dev_ctx.x_context(),
+                                    out_int64.data<int64_t>(),
+                                    out->data<int>(),
+                                    out_int64.numel());
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
+  }
+}
+
 }  // namespace phi
+
 PD_REGISTER_KERNEL(argmax,
                    XPU,
                    ALL_LAYOUT,
@@ -126,3 +195,8 @@ PD_REGISTER_KERNEL(argmax,
                    phi::dtype::float16) {
   kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
 }
+
+PD_REGISTER_KERNEL(
+    argmin, XPU, ALL_LAYOUT, phi::ArgMinKernel, float, phi::dtype::float16) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}
diff --git a/test/xpu/test_arg_min_op_xpu.py b/test/xpu/test_arg_min_op_xpu.py
@@ -0,0 +1,173 @@
+#  Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+from get_test_cover_info import (
+    XPUOpTestWrapper,
+    create_test_class,
+    get_xpu_op_support_types,
+)
+from op_test_xpu import XPUOpTest
+
+import paddle
+
+paddle.enable_static()
+
+
+class XPUTestArgMin(XPUOpTestWrapper):
+    def __init__(self):
+        self.op_name = 'arg_min'
+
+    class XPUBaseTestCase(XPUOpTest):
+        def initTestCase(self):
+            self.dims = (3, 4)
+            self.axis = 1
+
+        def setUp(self):
+            self.op_type = 'arg_min'
+            self.dtype = self.in_type
+            self.initTestCase()
+
+            self.x = (np.random.random(self.dims)).astype(self.dtype)
+            self.inputs = {'X': self.x}
+            self.attrs = {'axis': self.axis, 'use_xpu': True}
+            self.outputs = {'Out': np.argmin(self.x, axis=self.axis)}
+
+        def test_check_output(self):
+            if paddle.is_compiled_with_xpu():
+                place = paddle.XPUPlace(0)
+                self.check_output_with_place(place)
+
+    class TestArgMinCase1(XPUBaseTestCase):
+        def initTestCase(self):
+            self.dims = (3, 4, 5)
+            self.axis = -1
+
+    class TestArgMinCase2(XPUBaseTestCase):
+        def initTestCase(self):
+            self.dims = (3, 4, 5)
+            self.axis = 0
+
+    class TestArgMinCase3(XPUBaseTestCase):
+        def initTestCase(self):
+            self.dims = (3, 4, 5)
+            self.axis = 1
+
+    class TestArgMinCase4(XPUBaseTestCase):
+        def initTestCase(self):
+            self.dims = (3, 4, 5)
+            self.axis = 2
+
+    class TestArgMinCase5(XPUBaseTestCase):
+        def initTestCase(self):
+            self.dims = (3, 4)
+            self.axis = -1
+
+    class TestArgMinCase6(XPUBaseTestCase):
+        def initTestCase(self):
+            self.dims = (3, 4)
+            self.axis = 0
+
+    class TestArgMinCase7(XPUBaseTestCase):
+        def initTestCase(self):
+            self.dims = (3, 4)
+            self.axis = 1
+
+    class TestArgMinCase8(XPUBaseTestCase):
+        def initTestCase(self):
+            self.dims = (1,)
+            self.axis = 0
+
+    class TestArgMinCase9(XPUBaseTestCase):
+        def initTestCase(self):
+            self.dims = (2,)
+            self.axis = 0
+
+    class TestArgMinCase10(XPUBaseTestCase):
+        def initTestCase(self):
+            self.dims = (3,)
+            self.axis = 0
+
+
+support_types = get_xpu_op_support_types('arg_min')
+for stype in support_types:
+    create_test_class(globals(), XPUTestArgMin, stype)
+
+
+class TestArgMinAPI(unittest.TestCase):
+    def initTestCase(self):
+        self.dims = (3, 4, 5)
+        self.dtype = 'float32'
+        self.axis = 0
+
+    def setUp(self):
+        self.initTestCase()
+        self.__class__.use_xpu = True
+        self.place = [paddle.XPUPlace(0)]
+
+    def test_dygraph_api(self):
+        def run(place):
+            paddle.disable_static(place)
+            np.random.seed(2021)
+            numpy_input = (np.random.random(self.dims)).astype(self.dtype)
+            tensor_input = paddle.to_tensor(numpy_input)
+            numpy_output = np.argmin(numpy_input, axis=self.axis)
+            paddle_output = paddle.argmin(tensor_input, axis=self.axis)
+            np.testing.assert_allclose(
+                numpy_output, paddle_output.numpy(), rtol=1e-05
+            )
+            paddle.enable_static()
+
+        for place in self.place:
+            run(place)
+
+
+class TestArgMinAPI_2(unittest.TestCase):
+    def initTestCase(self):
+        self.dims = (3, 4, 5)
+        self.dtype = 'float32'
+        self.axis = 0
+        self.keep_dims = True
+
+    def setUp(self):
+        self.initTestCase()
+        self.__class__.use_xpu = True
+        self.place = [paddle.XPUPlace(0)]
+
+    def test_dygraph_api(self):
+        def run(place):
+            paddle.disable_static(place)
+            np.random.seed(2021)
+            numpy_input = (np.random.random(self.dims)).astype(self.dtype)
+            tensor_input = paddle.to_tensor(numpy_input)
+            numpy_output = np.argmin(numpy_input, axis=self.axis).reshape(
+                1, 4, 5
+            )
+            paddle_output = paddle.argmin(
+                tensor_input, axis=self.axis, keepdim=self.keep_dims
+            )
+            np.testing.assert_allclose(
+                numpy_output, paddle_output.numpy(), rtol=1e-05
+            )
+            self.assertEqual(numpy_output.shape, paddle_output.numpy().shape)
+            paddle.enable_static()
+
+        for place in self.place:
+            run(place)
+
+
+if __name__ == '__main__':
+    unittest.main()
