add SubOffset cpu op

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 <string>
+#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<kernel::AddressPtr> &inputs,
+                                   const std::vector<kernel::AddressPtr> & /*workspace*/,
+                                   const std::vector<kernel::AddressPtr> &outputs) {
+  if (input_x_dtype_ == kNumberTypeInt32) {
+    LaunchKernel<int>(inputs, outputs);
+  } else if (input_x_dtype_ == kNumberTypeInt64) {
+    LaunchKernel<int64_t>(inputs, outputs);
+  } else {
+    MS_LOG(ERROR) << "input x dtype only support int32, int64";
+    return false;
+  }
+  return true;
+}
+
+template <typename T>
+void SubAndFilterCPUKernel::LaunchKernel(const std::vector<AddressPtr> &inputs,
+                                         const std::vector<kernel::AddressPtr> &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<T *>(inputs[0]->addr);
+  T max_num = *reinterpret_cast<T *>(inputs[1]->addr);
+  T offset = *reinterpret_cast<T *>(inputs[2]->addr);
+  T *filter_res = reinterpret_cast<T *>(outputs[0]->addr);
+  T *filter_idx = reinterpret_cast<T *>(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<size_t> out_shape;
+  out_shape.emplace_back(count);
+  std::vector<TypeId> 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

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 <vector>
+#include <memory>
+#include <unordered_map>
+#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<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
+              const std::vector<AddressPtr> &outputs) override;
+
+  template <typename T>
+  void LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<kernel::AddressPtr> &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_

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,

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<AbstractTensor>(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<AbstractTensor>(input_x->element(), std::make_shared<Shape>(shape, min_shape, max_shape));
+  auto filter_idx =
+    std::make_shared<AbstractTensor>(input_x->element(), std::make_shared<Shape>(shape, min_shape, max_shape));
+  AbstractBasePtrList elements = {filter_res, filter_idx};
+  return std::make_shared<AbstractTuple>(elements);
+}
+
 AbstractBasePtr InferImplDiv(const AnalysisEnginePtr &, const PrimitivePtr &primitive,
                              const AbstractBasePtrList &args_spec_list) {
   const std::string op_name = primitive->name();

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}},

mindspore/core/base/core_ops.h

@@ -98,6 +98,7 @@ inline const PrimitivePtr kPrimUnsortedSegmentSum = std::make_shared<Primitive>(
 inline const PrimitivePtr kPrimUnsortedSegmentMin = std::make_shared<Primitive>("UnsortedSegmentMin");
 inline const PrimitivePtr kPrimConcatOffset = std::make_shared<Primitive>("ConcatOffset");
 inline const PrimitivePtr kPrimReshape = std::make_shared<Primitive>("Reshape");
+inline const PrimitivePtr kPrimSubAndFilter = std::make_shared<Primitive>("SubAndFilter");
 inline const PrimitivePtr kPrimMapCacheIdx = std::make_shared<Primitive>("MapCacheIdx");
 inline const PrimitivePtr kPrimUpdateCache = std::make_shared<Primitive>("UpdateCache");
 inline const PrimitivePtr kPrimCacheSwapTable = std::make_shared<Primitive>("CacheSwapTable");

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',

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)
 

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()