From ee7aef98dfe9ebc5a69a3d5675268517cea1f6a8 Mon Sep 17 00:00:00 2001 From: jiangjinsheng Date: Fri, 3 Jul 2020 11:42:40 +0800 Subject: [PATCH] vm for PopulationCount --- mindspore/ops/_op_impl/tbe/__init__.py | 1 + .../ops/_op_impl/tbe/population_count.py | 38 +++++++++++++++++++ mindspore/ops/operations/__init__.py | 5 ++- mindspore/ops/operations/other_ops.py | 31 +++++++++++++++ tests/ut/python/ops/test_ops.py | 5 ++- 5 files changed, 77 insertions(+), 3 deletions(-) create mode 100644 mindspore/ops/_op_impl/tbe/population_count.py diff --git a/mindspore/ops/_op_impl/tbe/__init__.py b/mindspore/ops/_op_impl/tbe/__init__.py index 35785a085c..76cea197ba 100644 --- a/mindspore/ops/_op_impl/tbe/__init__.py +++ b/mindspore/ops/_op_impl/tbe/__init__.py @@ -284,3 +284,4 @@ from .scatter_div import _scatter_div_tbe from .mod import _mod_tbe from .max_pool_grad_grad import _max_pool_grad_grad_tbe from .max_pool_grad_grad_with_argmax import _max_pool_grad_grad_with_argmax_tbe +from .population_count import _population_count_tbe diff --git a/mindspore/ops/_op_impl/tbe/population_count.py b/mindspore/ops/_op_impl/tbe/population_count.py new file mode 100644 index 0000000000..14feded367 --- /dev/null +++ b/mindspore/ops/_op_impl/tbe/population_count.py @@ -0,0 +1,38 @@ +# Copyright 2020 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. +# ============================================================================ + +"""PopulationCount op""" +from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType + +population_count_op_info = TBERegOp("PopulationCount") \ + .fusion_type("OPAQUE") \ + .async_flag(False) \ + .binfile_name("population_count.so") \ + .compute_cost(10) \ + .kernel_name("population_count") \ + .partial_flag(True) \ + .input(0, "x", False, "required", "all") \ + .output(0, "y", False, "required", "all") \ + .dtype_format(DataType.I16_5HD, DataType.U8_5HD) \ + .dtype_format(DataType.I16_Default, DataType.U8_Default) \ + .dtype_format(DataType.U16_5HD, DataType.U8_5HD) \ + .dtype_format(DataType.U16_Default, DataType.U8_Default) \ + .get_op_info() + + +@op_info_register(population_count_op_info) +def _population_count_tbe(): + """PopulationCount TBE register""" + return diff --git a/mindspore/ops/operations/__init__.py b/mindspore/ops/operations/__init__.py index b2d0fc7382..861a88ad08 100644 --- a/mindspore/ops/operations/__init__.py +++ b/mindspore/ops/operations/__init__.py @@ -76,7 +76,7 @@ from .nn_ops import (LSTM, SGD, Adam, SparseApplyAdam, SparseApplyLazyAdam, Appl ApplyAdaMax, ApplyAdadelta, ApplyAdagrad, ApplyAdagradV2, ApplyAddSign, ApplyPowerSign, ApplyGradientDescent, ApplyProximalGradientDescent, ApplyRMSProp, ApplyCenteredRMSProp, BasicLSTMCell, InTopK) -from .other_ops import (Assign, IOU, BoundingBoxDecode, BoundingBoxEncode, +from .other_ops import (Assign, IOU, BoundingBoxDecode, BoundingBoxEncode, PopulationCount, CheckValid, MakeRefKey, Partial, Depend, CheckBprop) from .thor_ops import * @@ -328,7 +328,8 @@ __all__ = [ "InplaceUpdate", "InTopK", "LRN", - "Mod" + "Mod", + "PopulationCount" ] __all__.sort() diff --git a/mindspore/ops/operations/other_ops.py b/mindspore/ops/operations/other_ops.py index b6b938d800..d72588b35f 100644 --- a/mindspore/ops/operations/other_ops.py +++ b/mindspore/ops/operations/other_ops.py @@ -51,6 +51,7 @@ class Assign(PrimitiveWithInfer): ('variable', sig_rw.RW_WRITE, sig_kind.KIND_POSITIONAL_KEYWORD, sig_kind.KIND_EMPTY_DEFAULT_VALUE, sig_dtype.T), ('value', sig_rw.RW_READ, sig_kind.KIND_POSITIONAL_KEYWORD, sig_kind.KIND_EMPTY_DEFAULT_VALUE, sig_dtype.T) ) + @prim_attr_register def __init__(self): self.init_prim_io_names(inputs=['ref', 'value'], outputs=['output']) @@ -324,6 +325,7 @@ class Partial(Primitive): partial_func = functools.partial(func, *args[1:]) return partial_func + class Depend(Primitive): """ Depend is used for process side-effect operations. @@ -457,3 +459,32 @@ class ConfusionMatrix(PrimitiveWithInfer): args = {"labels": labels, "predictions": predictions} validator.check_tensor_type_same(args, (mstype.number_type), self.name) return labels + + +class PopulationCount(PrimitiveWithInfer): + r""" + Calculate population count. + + Inputs: + - **input** (Tensor) - The data type should be int16 or uint16. + + Outputs: + Tensor, with shape same as the input. + + Examples: + >>> population_count = P.PopulationCount() + >>> x_input = Tensor([0, 1, 3], mindspore.int16) + >>> population_count(x_input) + """ + + @prim_attr_register + def __init__(self): + pass + + def infer_shape(self, x_shape): + return x_shape + + def infer_dtype(self, x_dtype): + args = {"x": x_dtype} + validator.check_tensor_type_same(args, (mstype.int16, mstype.uint16,), self.name) + return mstype.tensor_type(mstype.uint8) diff --git a/tests/ut/python/ops/test_ops.py b/tests/ut/python/ops/test_ops.py index 5262145c80..c222d9cc40 100755 --- a/tests/ut/python/ops/test_ops.py +++ b/tests/ut/python/ops/test_ops.py @@ -2133,7 +2133,10 @@ test_case_other_ops = [ 'desc_inputs': [Tensor(np.array([1.1]).astype(np.float32)), Tensor(np.array([1.2]).astype(np.float32))], 'skip': ['backward']}), - + ('PopulationCount', { + 'block': P.PopulationCount(), + 'desc_inputs': [Tensor(np.array([1, 2, 3]).astype(np.int16))], + 'skip': ['backward']}), ] test_case_quant_ops = [