diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/map_uniform_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/map_uniform_cpu_kernel.cc new file mode 100644 index 0000000000..c16dddbbc3 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/map_uniform_cpu_kernel.cc @@ -0,0 +1,67 @@ +/** + * 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. + */ + +#include "backend/kernel_compiler/cpu/map_uniform_cpu_kernel.h" +#include +#include +#include +#include "runtime/device/cpu/cpu_device_address.h" + +namespace mindspore { +namespace kernel { +void MapUniformCPUKernel::InitKernel(const CNodePtr &kernel_node) { + MS_EXCEPTION_IF_NULL(kernel_node); + node_ = kernel_node; + dtype_ = AnfAlgo::GetPrevNodeOutputInferDataType(kernel_node, 0); +} + +bool MapUniformCPUKernel::Launch(const std::vector &inputs, + const std::vector & /*workspace*/, + const std::vector &outputs) { + if (dtype_ == kNumberTypeInt32) { + LaunchKernel(inputs, outputs); + } else if (dtype_ == kNumberTypeInt64) { + LaunchKernel(inputs, outputs); + } else { + MS_LOG(ERROR) << "Only support int32, int64"; + return false; + } + return true; +} + +template +void MapUniformCPUKernel::LaunchKernel(const std::vector &inputs, + const std::vector &outputs) { + auto input_x_shape = AnfAlgo::GetPrevNodeOutputInferShape(node_, 0); + batch_size_ = 1; + for (size_t i = 0; i < input_x_shape.size(); ++i) { + batch_size_ *= input_x_shape[i]; + } + MS_LOG(INFO) << "Input size: " << batch_size_; + auto input_x = reinterpret_cast(inputs[0]->addr); + auto per_group_size = *reinterpret_cast(inputs[1]->addr); + auto group_num = *reinterpret_cast(inputs[2]->addr); + auto output_x = reinterpret_cast(outputs[0]->addr); + T max_num = group_num * per_group_size; + for (size_t i = 0; i < batch_size_; ++i) { + output_x[i] = input_x[i] % group_num * per_group_size + input_x[i] / group_num; + if (output_x[i] >= max_num) { + MS_LOG(EXCEPTION) << "Value can not >= " << max_num; + } + } +} +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/map_uniform_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/map_uniform_cpu_kernel.h new file mode 100644 index 0000000000..b69959a9c3 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/map_uniform_cpu_kernel.h @@ -0,0 +1,65 @@ +/** + * 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. + */ +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_MAP_UNIFORM_CPU_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_MAP_UNIFORM_CPU_KERNEL_H_ + +#include +#include +#include +#include +#include "backend/kernel_compiler/cpu/cpu_kernel.h" +#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h" + +namespace mindspore { +namespace kernel { +class MapUniformCPUKernel : public CPUKernel { + public: + MapUniformCPUKernel() = default; + ~MapUniformCPUKernel() override = default; + + void InitKernel(const CNodePtr &kernel_node) override; + + bool Launch(const std::vector &inputs, const std::vector &workspace, + const std::vector &outputs) override; + + template + void LaunchKernel(const std::vector &inputs, const std::vector &outputs); + + private: + size_t batch_size_{1}; + TypeId dtype_{kTypeUnknown}; + CNodePtr node_ = nullptr; +}; + +MS_REG_CPU_KERNEL(MapUniform, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeInt32), + MapUniformCPUKernel); + +MS_REG_CPU_KERNEL(MapUniform, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeInt64), + MapUniformCPUKernel); +} // namespace kernel +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_MAP_UNIFORM_CPU_KERNEL_H_ diff --git a/mindspore/core/abstract/infer_functions.h b/mindspore/core/abstract/infer_functions.h index 7d41b7bd7e..3ddbc46fb6 100644 --- a/mindspore/core/abstract/infer_functions.h +++ b/mindspore/core/abstract/infer_functions.h @@ -267,6 +267,8 @@ AbstractBasePtr InferImplGpuConvertToDynamicShape(const AnalysisEnginePtr &, con const AbstractBasePtrList &args_spec_list); AbstractBasePtr InferImplPad(const AnalysisEnginePtr &, const PrimitivePtr &primitive, const AbstractBasePtrList &args_spec_list); +AbstractBasePtr InferImplMapUniform(const AnalysisEnginePtr &, const PrimitivePtr &primitive, + const AbstractBasePtrList &args_spec_list); AbstractBasePtr InferImplSplit(const AnalysisEnginePtr &, const PrimitivePtr &primitive, const AbstractBasePtrList &args_spec_list); AbstractBasePtr InferImplSequenceMask(const AnalysisEnginePtr &, const PrimitivePtr &primitive, diff --git a/mindspore/core/abstract/prim_arrays.cc b/mindspore/core/abstract/prim_arrays.cc index 745890cdd6..66b188caf4 100644 --- a/mindspore/core/abstract/prim_arrays.cc +++ b/mindspore/core/abstract/prim_arrays.cc @@ -863,6 +863,14 @@ AbstractBasePtr InferImplReshape(const AnalysisEnginePtr &, const PrimitivePtr & return ret; } +AbstractBasePtr InferImplMapUniform(const AnalysisEnginePtr &, const PrimitivePtr &primitive, + const AbstractBasePtrList &args_spec_list) { + // Inputs: one tensor. + const std::string op_name = primitive->name(); + CheckArgsSize(op_name, args_spec_list, 3); + return args_spec_list[0]->Broaden(); +} + AbstractBasePtr InferImplSplit(const AnalysisEnginePtr &, const PrimitivePtr &primitive, const AbstractBasePtrList &args_spec_list) { const std::string op_name = primitive->name(); diff --git a/mindspore/core/abstract/primitive_infer_map.cc b/mindspore/core/abstract/primitive_infer_map.cc index 28078a9864..67d59bd6b3 100644 --- a/mindspore/core/abstract/primitive_infer_map.cc +++ b/mindspore/core/abstract/primitive_infer_map.cc @@ -74,6 +74,7 @@ PrimitiveEvalImplMap &GetPrimitiveToEvalImplMap() { {prim::kPrimDynamicShape, {InferImplDynamicShape, true}}, {prim::kPrimTranspose, {InferImplTranspose, true}}, {prim::kPrimReshape, {InferImplReshape, true}}, + {prim::kPrimMapUniform, {InferImplMapUniform, true}}, {prim::kPrimSplit, {InferImplSplit, true}}, {prim::kPrimSequenceMask, {InferImplSequenceMask, true}}, // Structure diff --git a/mindspore/core/base/core_ops.h b/mindspore/core/base/core_ops.h index ba24ff97b9..5d948e0a0e 100644 --- a/mindspore/core/base/core_ops.h +++ b/mindspore/core/base/core_ops.h @@ -119,6 +119,7 @@ inline const PrimitivePtr kPrimDynamicGRUV2 = std::make_shared("Dynam inline const PrimitivePtr kPrimDynamicGRUV2Grad = std::make_shared("DynamicGRUV2Grad"); inline const PrimitivePtr kPrimScatterAdd = std::make_shared("ScatterAdd"); inline const PrimitivePtr kPrimScatterUpdate = std::make_shared("ScatterUpdate"); +inline const PrimitivePtr kPrimMapUniform = std::make_shared("MapUniform"); inline const PrimitivePtr kPrimSplit = std::make_shared("Split"); inline const PrimitivePtr kPrimSequenceMask = std::make_shared("SequenceMask"); diff --git a/mindspore/ops/operations/__init__.py b/mindspore/ops/operations/__init__.py index 63bdf01b67..a238cade5b 100644 --- a/mindspore/ops/operations/__init__.py +++ b/mindspore/ops/operations/__init__.py @@ -90,7 +90,8 @@ from ._thor_ops import (CusBatchMatMul, CusCholeskyTrsm, CusFusedAbsMax1, CusImg CusMatMulCubeDenseRight, CusMatMulCubeFraczLeftCast, Im2Col, UpdateThorGradient, Cholesky, CholeskyTrsm, DetTriangle) from .sparse_ops import SparseToDense -from ._cache_ops import CacheSwapHashmap, SearchCacheIdx, CacheSwapTable, UpdateCache, MapCacheIdx, SubAndFilter +from ._cache_ops import (CacheSwapHashmap, SearchCacheIdx, CacheSwapTable, UpdateCache, MapCacheIdx, SubAndFilter, + MapUniform) __all__ = [ 'Unique', diff --git a/mindspore/ops/operations/_cache_ops.py b/mindspore/ops/operations/_cache_ops.py index 6617dae21d..c074656aa3 100644 --- a/mindspore/ops/operations/_cache_ops.py +++ b/mindspore/ops/operations/_cache_ops.py @@ -187,6 +187,46 @@ class SearchCacheIdx(PrimitiveWithInfer): return out_dtype +class MapUniform(PrimitiveWithCheck): + """ + Map a tensor by using fomula : value = key % `group_num` * `per_group_size` + key // `group_num`. + + Inputs: + - **input** (Tensor) - Input Tensor. + - **per_group_size** (int) - The size of each group. + - **group_num** (int) - The number of group. + + Outputs: + Tensor, has the same dtype and shape as the `input`. + + Supported Platforms: + `CPU` + + Examples: + >>> input_x = Tensor(np.array([0, 1, 2, 3, 4, 5, 6, 7])) + >>> per_group_size = 4 + >>> group_num = 2 + >>> map_uniform = ops.MapUniform() + >>> output = map_uniform(input_x, per_group_size, group_num) + >>> print(output) + [0, 4, 1, 5, 2, 6, 3, 7] + """ + + @prim_attr_register + def __init__(self): + """init MapUniform""" + self.init_prim_io_names(inputs=['input', 'per_group_size', 'group_num'], + outputs=['output']) + + def check_dtype(self, input_dtype, per_group_size_dtype, group_num_dtype): + validator.check_tensor_dtype_valid( + "input", input_dtype, mstype.int_type, self.name) + validator.check_value_type( + 'per_group_size', per_group_size_dtype, [mstype.Int], self.name) + validator.check_value_type( + 'group_num', group_num_dtype, [mstype.Int], self.name) + + class CacheSwapHashmap(PrimitiveWithInfer): """ Delete a hashmap entry,and insert a new key to hashmap, return the key and value of delete entry. diff --git a/tests/st/dynamic_shape/test_map_uniform_cpu.py b/tests/st/dynamic_shape/test_map_uniform_cpu.py new file mode 100644 index 0000000000..fc15ab9f22 --- /dev/null +++ b/tests/st/dynamic_shape/test_map_uniform_cpu.py @@ -0,0 +1,46 @@ +# 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. +# ============================================================================ +import numpy as np +import pytest +import mindspore.context as context +import mindspore.nn as nn +from mindspore import Tensor +import mindspore.common.dtype as mstype +from mindspore.ops import operations as P + +context.set_context(mode=context.GRAPH_MODE, device_target="CPU") + + +class Net(nn.Cell): + def __init__(self): + super(Net, self).__init__() + self.map_uniform = P.MapUniform() + self.per_group_size = 4 + self.group_num = 2 + + def construct(self, x): + return self.map_uniform(x, self.per_group_size, self.group_num) + + +@pytest.mark.level0 +@pytest.mark.platform_arm_ascend_training +@pytest.mark.platform_x86_ascend_training +@pytest.mark.env_onecard +def test_map_uniform(): + x = Tensor(np.array([0, 1, 2, 3, 4, 5, 6, 7]), mstype.int32) + net = Net() + output = net(x) + expect1 = np.array([0, 4, 1, 5, 2, 6, 3, 7]) + assert (output.asnumpy() == expect1).all()