!40369 [feat][assistant][I4XJGO]add Polar operator

Merge pull request !40369 from 桂宁馨/Polar

mindspore/ccsrc/plugin/device/cpu/kernel/polar_cpu_kernel.cc

@@ -0,0 +1,90 @@
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * 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
+ *
+ * http://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.
+ */
+
+#include "plugin/device/cpu/kernel/polar_cpu_kernel.h"
+#include <algorithm>
+#include <complex>
+#include <functional>
+#include <cmath>
+#include <tuple>
+#include <type_traits>
+#include "utils/ms_utils.h"
+#include "plugin/device/cpu/hal/device/cpu_device_address.h"
+
+namespace {
+constexpr size_t kPolarInputsNum = 2;
+constexpr size_t kPolarOutputsNum = 1;
+
+#define POLAR_COMPUTE_CASE(DTYPE, TYPE)        \
+  case (DTYPE): {                              \
+    ret = LaunchKernel<TYPE>(inputs, outputs); \
+    break;                                     \
+  }
+}  // namespace
+
+namespace mindspore {
+namespace kernel {
+bool PolarCpuKernelMod::Init(const BaseOperatorPtr &base_operator, const std::vector<KernelTensorPtr> &inputs,
+                             const std::vector<KernelTensorPtr> &outputs) {
+  MS_EXCEPTION_IF_NULL(base_operator);
+  kernel_name_ = base_operator->name();
+  input1_dtype_ = inputs[0]->GetDtype();
+  return true;
+}
+
+bool PolarCpuKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &,
+                               const std::vector<AddressPtr> &outputs) {
+  bool ret = true;
+  CHECK_KERNEL_INPUTS_NUM(inputs.size(), kPolarInputsNum, kernel_name_);
+  CHECK_KERNEL_OUTPUTS_NUM(outputs.size(), kPolarOutputsNum, kernel_name_);
+  switch (input1_dtype_) {
+    POLAR_COMPUTE_CASE(kNumberTypeFloat32, float)
+    POLAR_COMPUTE_CASE(kNumberTypeFloat64, double)
+    default:
+      ret = false;
+      MS_EXCEPTION(TypeError) << "For Polar, unsupported input data type: " << TypeIdToString(input1_dtype_) << ".";
+  }
+  return ret;
+}
+
+template <typename T>
+bool PolarCpuKernelMod::LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &outputs) {
+  const auto abs = reinterpret_cast<T *>(inputs[0]->addr);
+  const auto angle = reinterpret_cast<T *>(inputs[1]->addr);
+  auto output_addr = reinterpret_cast<std::complex<T> *>(outputs[0]->addr);
+  size_t output_size = outputs[0]->size / sizeof(std::complex<T>);
+  auto task = [output_addr, abs, angle](size_t start, size_t end) {
+    for (size_t i = start; i < end; ++i) {
+      output_addr[i].real(abs[i] * cos(angle[i]));
+      output_addr[i].imag(abs[i] * sin(angle[i]));
+    }
+  };
+  ParallelLaunchAutoSearch(task, output_size, this, &parallel_search_info_);
+  return true;
+}
+
+std::vector<KernelAttr> PolarCpuKernelMod::GetOpSupport() {
+  static std::vector<KernelAttr> support_list = {
+    KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeComplex64),
+    KernelAttr()
+      .AddInputAttr(kNumberTypeFloat64)
+      .AddInputAttr(kNumberTypeFloat64)
+      .AddOutputAttr(kNumberTypeComplex128)};
+  return support_list;
+}
+MS_KERNEL_FACTORY_REG(NativeCpuKernelMod, Polar, PolarCpuKernelMod);
+}  // namespace kernel
+}  // namespace mindspore

mindspore/ccsrc/plugin/device/cpu/kernel/polar_cpu_kernel.h

@@ -0,0 +1,56 @@
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * 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
+ *
+ * http://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.
+ */
+
+#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_POLAR_CPU_KERNEL_H
+#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_POLAR_CPU_KERNEL_H
+
+#include <cmath>
+#include <vector>
+#include <tuple>
+#include <map>
+#include <memory>
+#include <complex>
+#include <string>
+#include "plugin/device/cpu/kernel/cpu_kernel.h"
+#include "plugin/factory/ms_factory.h"
+
+namespace mindspore {
+namespace kernel {
+class PolarCpuKernelMod : public NativeCpuKernelMod {
+ public:
+  PolarCpuKernelMod() = default;
+  ~PolarCpuKernelMod() override = default;
+
+  bool Init(const BaseOperatorPtr &base_operator, const std::vector<KernelTensorPtr> &inputs,
+            const std::vector<KernelTensorPtr> &outputs) override;
+
+  bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
+              const std::vector<AddressPtr> &outputs) override;
+
+ protected:
+  std::vector<KernelAttr> GetOpSupport() override;
+
+ private:
+  template <typename T>
+  bool LaunchKernel(const std::vector<kernel::AddressPtr> &inputs, const std::vector<kernel::AddressPtr> &outputs);
+  string kernel_name_;
+  TypeId input1_dtype_{kTypeUnknown};
+  TypeId input2_dtype_{kTypeUnknown};
+};
+}  // namespace kernel
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_POLAR_CPU_KERNEL_H

mindspore/core/ops/polar.h

@@ -1,44 +1,44 @@
-/**
- * Copyright 2022 Huawei Technologies Co., Ltd
- *
- * 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
- *
- * http://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.
- */
-
-#ifndef MINDSPORE_CORE_OPS_POLAR_H_
-#define MINDSPORE_CORE_OPS_POLAR_H_
-#include <map>
-#include <memory>
-#include <set>
-#include <string>
-#include <vector>
-
-#include "ops/base_operator.h"
-#include "mindapi/base/types.h"
-
-namespace mindspore {
-namespace ops {
-constexpr auto kNamePolar = "Polar";
-
-class MIND_API Polar : public BaseOperator {
- public:
-  MIND_API_BASE_MEMBER(Polar);
-  Polar() : BaseOperator(kNamePolar) { InitIOName({"abs", "angle"}, {"y"}); }
-};
-
-abstract::AbstractBasePtr PolarInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
-                                     const std::vector<abstract::AbstractBasePtr> &input_args);
-using PrimPolarPtr = std::shared_ptr<Polar>;
-}  // namespace ops
-}  // namespace mindspore
-
-#endif  // MINDSPORE_CORE_OPS_POLAR_H_
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * 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
+ *
+ * http://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.
+ */
+
+#ifndef MINDSPORE_CORE_OPS_POLAR_H_
+#define MINDSPORE_CORE_OPS_POLAR_H_
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "ops/base_operator.h"
+#include "mindapi/base/types.h"
+
+namespace mindspore {
+namespace ops {
+constexpr auto kNamePolar = "Polar";
+
+class MIND_API Polar : public BaseOperator {
+ public:
+  MIND_API_BASE_MEMBER(Polar);
+  Polar() : BaseOperator(kNamePolar) { InitIOName({"abs", "angle"}, {"y"}); }
+};
+
+abstract::AbstractBasePtr PolarInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
+                                     const std::vector<abstract::AbstractBasePtr> &input_args);
+using PrimPolarPtr = std::shared_ptr<Polar>;
+}  // namespace ops
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CORE_OPS_POLAR_H_

mindspore/python/mindspore/ops/_op_impl/aicpu/polar.py

@@ -0,0 +1,32 @@
+# Copyright 2022 Huawei Technologies Co., Ltd
+#
+# 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
+#
+# http://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.
+# ============================================================================
+
+"""Polar op"""
+from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
+
+polar_op_info = AiCPURegOp("Polar") \
+    .fusion_type("OPAQUE") \
+    .input(0, "abs", "required") \
+    .input(1, "angle", "required") \
+    .output(0, "y", "required") \
+    .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.C64_Default) \
+    .dtype_format(DataType.F64_Default, DataType.F64_Default, DataType.C128_Default) \
+    .get_op_info()
+
+
+@op_info_register(polar_op_info)
+def _polar_aicpu():
+    """Polar aicpu register"""
+    return

mindspore/python/mindspore/ops/function/math_func.py

@@ -63,6 +63,7 @@ from mindspore.ops.operations.math_ops import (
     InplaceUpdateV2,
     Igamma,
     Igammac,
+    Polar,
     Angle,
 )
 from mindspore.common.tensor import Tensor
@@ -127,6 +128,7 @@ sinc_ = Sinc()
 cos_ = P.Cos()
 tan_ = P.Tan()
 asin_ = P.Asin()
+polar_ = Polar()
 acos_ = P.ACos()
 atan_ = P.Atan()
 sinh_ = P.Sinh()
@@ -1673,6 +1675,45 @@ def arctan2(x, other):
     return _atan2(x, other)
 
 
+def polar(abs, angle):  # pylint: disable=redefined-outer-name
+    r"""
+    Returns the complex tensor at polar coordinates.
+
+    .. math::
+
+        y_{i} =  abs_{i} * cos(angle_{i}) + abs_{i} * sin(angle_{i}) * j
+
+    Args:
+        abs (Tensor): The shape of tensor is
+            :math:`(N,*)`, where :math:`*` means additional dimensions of size less than 8.
+            Must be one of the following types: float32, float64.
+
+        angle (Tensor): The shape of tensor is
+            :math:`(N,*)`, where :math:`*` means additional dimensions of size less than 8.
+            Must be one of the following types: float32, float64.
+
+    Outputs:
+        Tensor, has the same shape and data type as `abs`.
+
+    Raises:
+        TypeError: If neither `abs` nor `angle` is a Tensor.
+        TypeError: If the dtype of input is not one of: float32, float64.
+        TypeError: If the dtypes of two args are not the same.
+        ValueError: If the shape of `abs` is not the same as that of `angle`.
+
+    Supported Platforms:
+        ``GPU`` ``CPU``
+
+    Examples:
+        >>> abs = Tensor(np.array([1, 2]), mindspore.float64)
+        >>> angle = Tensor(np.array([np.pi / 2, 5 * np.pi / 4]), mindspore.float64)
+        >>> output = ops.polar(abs, angle)
+        >>> print(output)
+        [ 6.12323400e-17+1.j         -1.41421356e+00-1.41421356j]
+    """
+    return polar_(abs, angle)
+
+
 def asin(x):
     r"""
     Computes arcsine of input tensors element-wise.

mindspore/python/mindspore/ops/operations/math_ops.py

@@ -7241,7 +7241,7 @@ class Polar(Primitive):
         ValueError: If `abs`'s shape is not the same as `angle`.
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> polar = ops.Polar()

tests/ut/python/ops/test_math_ops.py

@@ -32,7 +32,7 @@ from mindspore.ops.operations.math_ops import Zeta, Igamma, Igammac, BatchMatMul
 from mindspore.ops.operations.math_ops import MatrixTriangularSolve
 from mindspore.ops.operations.sparse_ops import DenseToDenseSetOperation
 from mindspore.ops.operations.sparse_ops import DenseToSparseSetOperation
-from mindspore.ops.function.math_func import inplace_index_add
+from mindspore.ops.function.math_func import inplace_index_add, polar
 
 from mindspore.common.parameter import Parameter
 from mindspore.common.initializer import initializer
@@ -632,6 +632,15 @@ class KronFunc(nn.Cell):
         return self.kron(x, y)
 
 
+class PolarFunc(nn.Cell):
+    def __init__(self):
+        super(PolarFunc, self).__init__()
+        self.polar = polar
+
+    def construct(self, x, y):
+        return self.polar(x, y)
+
+
 class Rot90Func(nn.Cell):
     def __init__(self):
         super(Rot90Func, self).__init__()
@@ -849,6 +858,12 @@ test_case_math_ops = [
         'desc_inputs': [Tensor(np.array([[0, 1, 2], [3, 4, 5]]).astype(np.float32)),
                         Tensor(np.array([[-1, -2, -3], [-4, -6, -8]]).astype(np.float32))],
         'skip': ['backward']}),
+    ('Polar', {
+        'block': PolarFunc(),
+        'desc_inputs': [Tensor(np.array([[0, 1, 2], [3, 4, 5]]).astype(np.float32)),
+                        Tensor(np.array([[-1, -2, -3], [-4, -6, -8]]).astype(np.float32))],
+        'desc_bprop': [Tensor(np.array([1+2j, 2+3j, 3+4j], np.complex64))],
+    }),
     ('Rot90', {
         'block': Rot90Func(),
         'desc_inputs': [Tensor(np.array([[0, 1], [2, 3]]).astype(np.float32))],