diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/sub_and_filter_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/sub_and_filter_cpu_kernel.cc new file mode 100644 index 0000000000..f3090dbe71 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/sub_and_filter_cpu_kernel.cc @@ -0,0 +1,78 @@ +/** + * 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/sub_and_filter_cpu_kernel.h" +#include +#include "runtime/device/cpu/cpu_device_address.h" + +namespace mindspore { +namespace kernel { +void SubAndFilterCPUKernel::InitKernel(const CNodePtr &kernel_node) { + MS_EXCEPTION_IF_NULL(kernel_node); + node_ = kernel_node; + input_x_dtype_ = AnfAlgo::GetPrevNodeOutputInferDataType(kernel_node, 0); +} + +bool SubAndFilterCPUKernel::Launch(const std::vector &inputs, + const std::vector & /*workspace*/, + const std::vector &outputs) { + if (input_x_dtype_ == kNumberTypeInt32) { + LaunchKernel(inputs, outputs); + } else if (input_x_dtype_ == kNumberTypeInt64) { + LaunchKernel(inputs, outputs); + } else { + MS_LOG(ERROR) << "input x dtype only support int32, int64"; + return false; + } + return true; +} + +template +void SubAndFilterCPUKernel::LaunchKernel(const std::vector &inputs, + const std::vector &outputs) { + auto indices_shape = AnfAlgo::GetPrevNodeOutputInferShape(node_, 0); + + batch_size_ = 1; + for (size_t i = 0; i < indices_shape.size(); ++i) { + batch_size_ *= indices_shape[i]; + } + MS_LOG(INFO) << "SubAndFilter batch_size:" << batch_size_; + + T *input_x = reinterpret_cast(inputs[0]->addr); + T max_num = *reinterpret_cast(inputs[1]->addr); + T offset = *reinterpret_cast(inputs[2]->addr); + T *filter_res = reinterpret_cast(outputs[0]->addr); + T *filter_idx = reinterpret_cast(outputs[1]->addr); + + size_t count = 0; + for (size_t i = 0; i < batch_size_; ++i) { + T temp = input_x[i] - offset; + if (temp < 0 || temp >= max_num) continue; + filter_res[count] = temp; + filter_idx[count] = i; + count++; + } + MS_LOG(INFO) << "SubAndFilter output count is " << count; + std::vector out_shape; + out_shape.emplace_back(count); + std::vector dtypes; + for (size_t i = 0; i < AnfAlgo::GetOutputTensorNum(node_); i++) { + dtypes.push_back(AnfAlgo::GetOutputInferDataType(node_, i)); + } + AnfAlgo::SetOutputInferTypeAndShape(dtypes, {out_shape, out_shape}, node_.get()); +} +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/sub_and_filter_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/sub_and_filter_cpu_kernel.h new file mode 100644 index 0000000000..cdc0af1e7d --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/sub_and_filter_cpu_kernel.h @@ -0,0 +1,66 @@ +/** + * 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_SUB_AND_FILTER_CPU_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SUB_AND_FILTER_CPU_KERNEL_H_ + +#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 SubAndFilterCPUKernel : public CPUKernel { + public: + SubAndFilterCPUKernel() = default; + ~SubAndFilterCPUKernel() 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 input_x_dtype_{kTypeUnknown}; + CNodePtr node_ = nullptr; +}; + +MS_REG_CPU_KERNEL(SubAndFilter, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeInt32), + SubAndFilterCPUKernel); + +MS_REG_CPU_KERNEL(SubAndFilter, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeInt64), + SubAndFilterCPUKernel); +} // namespace kernel +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SUB_AND_FILTER_CPU_KERNEL_H_ diff --git a/mindspore/core/abstract/infer_functions.h b/mindspore/core/abstract/infer_functions.h index 61e6443f47..7853361a20 100644 --- a/mindspore/core/abstract/infer_functions.h +++ b/mindspore/core/abstract/infer_functions.h @@ -207,6 +207,8 @@ AbstractBasePtr InferImplDiv(const AnalysisEnginePtr &, const PrimitivePtr &prim const AbstractBasePtrList &args_spec_list); AbstractBasePtr InferImplRealDiv(const AnalysisEnginePtr &, const PrimitivePtr &primitive, const AbstractBasePtrList &args_spec_list); +AbstractBasePtr InferImplSubAndFilter(const AnalysisEnginePtr &, const PrimitivePtr &primitive, + const AbstractBasePtrList &args_spec_list); AbstractBasePtr InferImplMapCacheIdx(const AnalysisEnginePtr &, const PrimitivePtr &primitive, const AbstractBasePtrList &args_spec_list); AbstractBasePtr InferImplCacheSwapTable(const AnalysisEnginePtr &, const PrimitivePtr &primitive, diff --git a/mindspore/core/abstract/prim_arrays.cc b/mindspore/core/abstract/prim_arrays.cc index 7de00d5160..d12addfa63 100644 --- a/mindspore/core/abstract/prim_arrays.cc +++ b/mindspore/core/abstract/prim_arrays.cc @@ -462,6 +462,34 @@ AbstractBasePtr InferImplUpdateCache(const AnalysisEnginePtr &, const PrimitiveP return ret; } +AbstractBasePtr InferImplSubAndFilter(const AnalysisEnginePtr &, const PrimitivePtr &primitive, + const AbstractBasePtrList &args_spec_list) { + const std::string op_name = primitive->name(); + auto input_x = CheckArg(op_name, args_spec_list, 0); + auto input_x_shp = input_x->shape(); + MS_EXCEPTION_IF_NULL(input_x); + MS_EXCEPTION_IF_NULL(input_x_shp); + + ShapeVector shape; + ShapeVector min_shape; + ShapeVector max_shape; + if (!input_x_shp->max_shape().empty()) { + max_shape = input_x_shp->max_shape(); + } else { + max_shape = input_x_shp->shape(); + } + for (size_t i = 0; i < max_shape.size(); i++) { + shape.emplace_back(Shape::SHP_ANY); + min_shape.emplace_back(1); + } + auto filter_res = + std::make_shared(input_x->element(), std::make_shared(shape, min_shape, max_shape)); + auto filter_idx = + std::make_shared(input_x->element(), std::make_shared(shape, min_shape, max_shape)); + AbstractBasePtrList elements = {filter_res, filter_idx}; + return std::make_shared(elements); +} + AbstractBasePtr InferImplDiv(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 09d42586f7..64d81bea69 100644 --- a/mindspore/core/abstract/primitive_infer_map.cc +++ b/mindspore/core/abstract/primitive_infer_map.cc @@ -60,6 +60,7 @@ PrimitiveEvalImplMap &GetPrimitiveToEvalImplMap() { {prim::kPrimUnsortedSegmentSum, {InferImplUnsortedSegmentSum, true}}, {prim::kPrimUnsortedSegmentMax, {InferImplUnsortedSegmentMax, true}}, {prim::kPrimScatterAdd, {InferImplScatterAdd, true}}, + {prim::kPrimSubAndFilter, {InferImplSubAndFilter, true}}, {prim::kPrimScatterUpdate, {InferImplScatterUpdate, true}}, {prim::kPrimMapCacheIdx, {InferImplMapCacheIdx, true}}, {prim::kPrimCacheSwapTable, {InferImplCacheSwapTable, true}}, diff --git a/mindspore/core/base/core_ops.h b/mindspore/core/base/core_ops.h index 2ec7b17a4f..e964d7cee9 100644 --- a/mindspore/core/base/core_ops.h +++ b/mindspore/core/base/core_ops.h @@ -98,6 +98,7 @@ inline const PrimitivePtr kPrimUnsortedSegmentSum = std::make_shared( inline const PrimitivePtr kPrimUnsortedSegmentMin = std::make_shared("UnsortedSegmentMin"); inline const PrimitivePtr kPrimConcatOffset = std::make_shared("ConcatOffset"); inline const PrimitivePtr kPrimReshape = std::make_shared("Reshape"); +inline const PrimitivePtr kPrimSubAndFilter = std::make_shared("SubAndFilter"); inline const PrimitivePtr kPrimMapCacheIdx = std::make_shared("MapCacheIdx"); inline const PrimitivePtr kPrimUpdateCache = std::make_shared("UpdateCache"); inline const PrimitivePtr kPrimCacheSwapTable = std::make_shared("CacheSwapTable"); diff --git a/mindspore/ops/operations/__init__.py b/mindspore/ops/operations/__init__.py index b6f63f1146..2f359984ba 100644 --- a/mindspore/ops/operations/__init__.py +++ b/mindspore/ops/operations/__init__.py @@ -90,7 +90,7 @@ from ._thor_ops import (CusBatchMatMul, CusCholeskyTrsm, CusFusedAbsMax1, CusImg CusMatMulCubeDenseRight, CusMatMulCubeFraczLeftCast, Im2Col, UpdateThorGradient, CholeskyTrsm, DetTriangle) from .sparse_ops import SparseToDense -from ._cache_ops import CacheSwapHashmap, SearchCacheIdx, CacheSwapTable, UpdateCache, MapCacheIdx +from ._cache_ops import CacheSwapHashmap, SearchCacheIdx, CacheSwapTable, UpdateCache, MapCacheIdx, SubAndFilter __all__ = [ 'Unique', diff --git a/mindspore/ops/operations/_cache_ops.py b/mindspore/ops/operations/_cache_ops.py index 749924e877..6617dae21d 100644 --- a/mindspore/ops/operations/_cache_ops.py +++ b/mindspore/ops/operations/_cache_ops.py @@ -56,6 +56,51 @@ class UpdateCache(PrimitiveWithCheck): return input_x_dtype +class SubAndFilter(PrimitiveWithCheck): + """ + Dynamic kernel, sub an offset and + return the elements which in range [0, max_num). + + Inputs: + - **input_x** (Tensor) - Input tensor. + - **max_num** (Int) - The max value of element that after sub `offset`. + - **offset** (int) - Specifies the offset value of this `input_x`. + + Outputs: + tuple(Tensor), tuple of 2 tensors, filter_res and filter_idx. + - **filter_res** (Tensor) - The result that `input_x` minus `offset`, + and return which in the range [0, max_num). + - **filter_idx** (Tensor) - A tensor containing indices of elements in the input + coressponding to the output tensor. + + Supported Platforms: + `CPU` + + Examples: + >>> x = Tensor(np.array([1, 3, 5, 8, 9, 16]), mindspore.int32) + >>> max_num = 10 + >>> offset = 5 + >>> output = ops.SubAndFilter()(x, max_num, offset) + >>> print(output) + (Tensor(shape=[3], dtype=Int32, value= [0, 3, 4]), + Tensor(shape=[3], dtype=Int32, value= [2, 3, 4])) + """ + @prim_attr_register + def __init__(self): + """init SubAndFilter""" + + self.init_prim_io_names(inputs=['input_x', 'max_num', 'offset'], + outputs=['sub_res', 'sub_idx']) + + def check_shape(self, input_x_shape, max_num_shape, offset_shape): + return (-1, -1) + + def check_dtype(self, input_x_dtype, max_num_dtype, offset_dtype): + validator.check_tensor_dtype_valid( + "input_x", input_x_dtype, mstype.int_type, self.name) + return input_x_dtype + + class SearchCacheIdx(PrimitiveWithInfer): """ Search the keys of a hashmap, and return the values. @@ -254,7 +299,8 @@ class MapCacheIdx(PrimitiveWithCheck): hashmap_dtype = hashmap['dtype'] indices_dtype = indices['dtype'] args = {"hashmap": hashmap_dtype, "indices": indices_dtype} - validator.check_tensor_type_same(args, mstype.int_type, self.name) + validator.check_tensors_dtypes_same_and_valid( + args, mstype.int_type, self.name) out_dtype = (hashmap_dtype, hashmap_dtype, hashmap_dtype, hashmap_dtype) diff --git a/tests/st/dynamic_shape/test_sub_and_filter_cpu.py b/tests/st/dynamic_shape/test_sub_and_filter_cpu.py new file mode 100644 index 0000000000..542de78863 --- /dev/null +++ b/tests/st/dynamic_shape/test_sub_and_filter_cpu.py @@ -0,0 +1,48 @@ +# 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.sub_and_filter = P.SubAndFilter() + self.offset = 5 + self.max_num = 10 + + def construct(self, x): + return self.sub_and_filter(x, self.max_num, self.offset) + + +@pytest.mark.level0 +@pytest.mark.platform_arm_ascend_training +@pytest.mark.platform_x86_ascend_training +@pytest.mark.env_onecard +def test_sub_and_filter(): + x = Tensor(np.array([1, 3, 5, 9, 6, 15]), mstype.int32) + sub_and_filter = Net() + output = sub_and_filter(x) + expect1 = np.array([0, 4, 1]) + expect2 = np.array([2, 3, 4]) + assert (output[0].asnumpy() == expect1).all() + assert (output[1].asnumpy() == expect2).all()