!22149 [assistant][ops] Add New IsInf

Merge pull request !22149 from 孟权令/IsInf

mindspore/ccsrc/backend/kernel_compiler/cpu/isinf_cpu_kernel.cc

@@ -0,0 +1,82 @@
+/**
+ * Copyright 2021 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 "backend/kernel_compiler/cpu/isinf_cpu_kernel.h"
+#include <cmath>
+#include "abstract/utils.h"
+#include "runtime/device/cpu/cpu_device_address.h"
+
+namespace mindspore {
+namespace kernel {
+void IsInfCPUKernel::InitKernel(const CNodePtr &kernelNode) {
+  MS_EXCEPTION_IF_NULL(kernelNode);
+  size_t input_num = AnfAlgo::GetInputTensorNum(kernelNode);
+  if (input_num != 1) {
+    MS_LOG(EXCEPTION) << "Input number is " << input_num << ", but IsInfCPUKernel needs 1 inputs.";
+  }
+  size_t output_num = AnfAlgo::GetOutputTensorNum(kernelNode);
+  if (output_num != 1) {
+    MS_LOG(EXCEPTION) << "Output number is " << output_num << ", but IsInfCPUKernel needs 1 output.";
+  }
+
+  input_dtype_ = AnfAlgo::GetInputDeviceDataType(kernelNode, 0);
+  if (dtype_map_.find(input_dtype_) == dtype_map_.end()) {
+    MS_LOG(EXCEPTION) << "Unsupported input type found.";
+  }
+}
+
+bool IsInfCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs, const std::vector<kernel::AddressPtr> &,
+                            const std::vector<kernel::AddressPtr> &outputs) {
+  if (input_dtype_ == kNumberTypeFloat16) {
+    LaunchKernelFloat16(inputs, outputs);
+  } else if (input_dtype_ == kNumberTypeFloat32 || input_dtype_ == kNumberTypeFloat) {
+    LaunchKernelFloat<float>(inputs, outputs);
+  } else if (input_dtype_ == kNumberTypeFloat64) {
+    LaunchKernelFloat<double>(inputs, outputs);
+  } else {
+    MS_LOG(EXCEPTION) << "Only support float, but actual data type is " << TypeIdLabel(input_dtype_);
+  }
+  return true;
+}
+
+void IsInfCPUKernel::LaunchKernelFloat16(const std::vector<AddressPtr> &inputs,
+                                         const std::vector<kernel::AddressPtr> &outputs) {
+  float16 *input = reinterpret_cast<float16 *>(inputs[0]->addr);
+  bool *output = reinterpret_cast<bool *>(outputs[0]->addr);
+
+  size_t elem_num = inputs[0]->size / sizeof(float16);
+
+  for (size_t i = 0; i < elem_num; i++) {
+    float temp_num = static_cast<float>(input[i]);
+    output[i] = std::isinf(temp_num);
+  }
+}
+
+template <typename T>
+void IsInfCPUKernel::LaunchKernelFloat(const std::vector<AddressPtr> &inputs,
+                                       const std::vector<kernel::AddressPtr> &outputs) {
+  T *input = reinterpret_cast<T *>(inputs[0]->addr);
+  bool *output = reinterpret_cast<bool *>(outputs[0]->addr);
+
+  size_t elem_num = inputs[0]->size / sizeof(T);
+
+  for (size_t i = 0; i < elem_num; i++) {
+    output[i] = std::isinf(input[i]);
+  }
+}
+
+}  // namespace kernel
+}  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/cpu/isinf_cpu_kernel.h

@@ -0,0 +1,57 @@
+/**
+ * Copyright 2021 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_IsInf_CPU_KERNEL_H_
+#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_IsInf_CPU_KERNEL_H_
+
+#include <vector>
+#include <map>
+#include "backend/kernel_compiler/cpu/cpu_kernel.h"
+#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h"
+
+namespace mindspore {
+namespace kernel {
+class IsInfCPUKernel : public CPUKernel {
+ public:
+  IsInfCPUKernel() = default;
+  ~IsInfCPUKernel() override = default;
+
+  void InitKernel(const CNodePtr &kernelNode) override;
+
+  bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
+              const std::vector<AddressPtr> &outputs) override;
+
+ private:
+  template <typename T>
+  void LaunchKernelFloat(const std::vector<AddressPtr> &inputs, const std::vector<kernel::AddressPtr> &outputs);
+
+  void LaunchKernelFloat16(const std::vector<AddressPtr> &inputs, const std::vector<kernel::AddressPtr> &outputs);
+
+ private:
+  std::map<TypeId, size_t> dtype_map_ = {
+    {kNumberTypeFloat16, sizeof(float16)}, {kNumberTypeFloat32, sizeof(float)}, {kNumberTypeFloat64, sizeof(double)}};
+
+  TypeId input_dtype_{kTypeUnknown};
+};
+
+MS_REG_CPU_KERNEL(IsInf, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeBool), IsInfCPUKernel);
+
+MS_REG_CPU_KERNEL(IsInf, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeBool), IsInfCPUKernel);
+
+MS_REG_CPU_KERNEL(IsInf, KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeBool), IsInfCPUKernel);
+}  // namespace kernel
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_IsInf_CPU_KERNEL_H_

mindspore/core/ops/is_inf.cc

@@ -0,0 +1,57 @@
+/**
+ * Copyright 2021 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 <map>
+#include <string>
+
+#include "ops/is_inf.h"
+#include "ops/op_utils.h"
+#include "utils/check_convert_utils.h"
+#include "abstract/primitive_infer_map.h"
+
+namespace mindspore {
+namespace ops {
+namespace {
+abstract::ShapePtr InferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
+  MS_EXCEPTION_IF_NULL(primitive);
+  const int64_t input_num = 1;
+  CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, input_num, primitive->name());
+  for (const auto &item : input_args) {
+    MS_EXCEPTION_IF_NULL(item);
+  }
+  auto x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape())[kShape];
+  return std::make_shared<abstract::Shape>(x_shape);
+}
+
+TypePtr InferType(const PrimitivePtr &prim, const std::vector<AbstractBasePtr> &input_args) {
+  for (const auto &item : input_args) {
+    MS_EXCEPTION_IF_NULL(item);
+  }
+  const int64_t input_num = 1;
+  CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, input_num, prim->name());
+  CheckAndConvertUtils::CheckTensorTypeValid("x", input_args[0]->BuildType(), {kFloat16, kFloat32, kFloat64},
+                                             prim->name());
+  return kBool;
+}
+}  // namespace
+
+AbstractBasePtr IsInfInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
+                           const std::vector<AbstractBasePtr> &input_args) {
+  return abstract::MakeAbstract(InferShape(primitive, input_args), InferType(primitive, input_args));
+}
+REGISTER_PRIMITIVE_EVAL_IMPL(IsInf, prim::kPrimIsInf, IsInfInfer, nullptr, true);
+}  // namespace ops
+}  // namespace mindspore

mindspore/core/ops/is_inf.h

@@ -0,0 +1,42 @@
+/**
+ * Copyright 2021 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_IS_INF_H_
+#define MINDSPORE_CORE_OPS_IS_INF_H_
+#include <vector>
+#include <memory>
+
+#include "ops/primitive_c.h"
+#include "ops/op_utils.h"
+#include "abstract/abstract_value.h"
+#include "utils/check_convert_utils.h"
+
+namespace mindspore {
+namespace ops {
+constexpr auto kNameIsInf = "IsInf";
+class IsInf : public PrimitiveC {
+ public:
+  IsInf() : PrimitiveC(kNameIsInf) { InitIOName({"x"}, {"y"}); }
+  ~IsInf() = default;
+  MS_DECLARE_PARENT(IsInf, PrimitiveC);
+};
+
+AbstractBasePtr IsInfInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
+                           const std::vector<AbstractBasePtr> &input_args);
+using PrimIsInfPtr = std::shared_ptr<IsInf>;
+}  // namespace ops
+}  // namespace mindspore
+#endif  // MINDSPORE_CORE_OPS_IS_INF_H_

mindspore/ops/_op_impl/aicpu/__init__.py

@@ -34,6 +34,7 @@ from .get_next import _get_next_aicpu
 from .print_tensor import _print_aicpu
 from .topk import _top_k_aicpu
 from .is_finite import _is_finite_aicpu
+from .is_inf import _is_inf_aicpu
 from .reshape import _reshape_aicpu
 from .flatten import _flatten_aicpu
 from .squeeze import _squeeze_aicpu

mindspore/ops/_op_impl/aicpu/is_inf.py

@@ -0,0 +1,31 @@
+# Copyright 2021 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.
+# ============================================================================
+
+"""IsInf op"""
+from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
+
+is_inf_op_info = AiCPURegOp("IsInf") \
+    .fusion_type("OPAQUE") \
+    .input(0, "x", "required") \
+    .output(0, "y", "required") \
+    .dtype_format(DataType.F16_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.F32_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.F64_Default, DataType.BOOL_Default) \
+    .get_op_info()
+
+@op_info_register(is_inf_op_info)
+def _is_inf_aicpu():
+    """IsInf AiCPU register"""
+    return

mindspore/ops/operations/math_ops.py

@@ -4020,7 +4020,7 @@ class IsNan(PrimitiveWithInfer):
         return mstype.tensor_type(mstype.bool_)
 
 
-class IsInf(PrimitiveWithInfer):
+class IsInf(Primitive):
     r"""
     Determines which elements are inf or -inf for each position
 
@@ -4044,7 +4044,7 @@ class IsInf(PrimitiveWithInfer):
         TypeError: If `x` is not a Tensor.
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> is_inf = ops.IsInf()
@@ -4059,12 +4059,6 @@ class IsInf(PrimitiveWithInfer):
         """Initialize IsInf"""
         self.init_prim_io_names(inputs=['x'], outputs=['output'])
 
-    def infer_shape(self, x_shape):
-        return x_shape
-
-    def infer_dtype(self, x_dtype):
-        return mstype.tensor_type(mstype.bool_)
-
 
 class IsFinite(PrimitiveWithInfer):
     r"""

tests/ut/python/ops/test_ops.py

@@ -1193,6 +1193,11 @@ test_case_math_ops = [
         'desc_inputs': [Tensor(np.array([[1, 2], [3, 4], [5, 6]]).astype(np.float32)),
                         Tensor(np.array([[0.5, 1], [1, 1.5]]).astype(np.float32))],
         'skip': ['backward']}),
+    ('IsInf', {
+        'block': P.IsInf(),
+        'desc_inputs': [Tensor(np.array([np.log(-1), 1, np.log(0)]).astype(np.float32))],
+        'desc_bprop': [],
+        'skip': ['backward']}),
     ('ACos', {
         'block': P.ACos(),
         'desc_inputs': [Tensor(np.array([2., 3.]).astype(np.float32))],