!9462 add MapUniform dynamic kernel

From: @fangzehua
Reviewed-by: @kisnwang,@stsuteng
Signed-off-by: @stsuteng

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 <string>
+#include <memory>
+#include <vector>
+#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<kernel::AddressPtr> &inputs,
+                                 const std::vector<kernel::AddressPtr> & /*workspace*/,
+                                 const std::vector<kernel::AddressPtr> &outputs) {
+  if (dtype_ == kNumberTypeInt32) {
+    LaunchKernel<int>(inputs, outputs);
+  } else if (dtype_ == kNumberTypeInt64) {
+    LaunchKernel<int64_t>(inputs, outputs);
+  } else {
+    MS_LOG(ERROR) << "Only support int32, int64";
+    return false;
+  }
+  return true;
+}
+
+template <typename T>
+void MapUniformCPUKernel::LaunchKernel(const std::vector<AddressPtr> &inputs,
+                                       const std::vector<kernel::AddressPtr> &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<T *>(inputs[0]->addr);
+  auto per_group_size = *reinterpret_cast<T *>(inputs[1]->addr);
+  auto group_num = *reinterpret_cast<T *>(inputs[2]->addr);
+  auto output_x = reinterpret_cast<T *>(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

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 <math.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 MapUniformCPUKernel : public CPUKernel {
+ public:
+  MapUniformCPUKernel() = default;
+  ~MapUniformCPUKernel() 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 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_

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,

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

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

mindspore/core/base/core_ops.h

@@ -119,6 +119,7 @@ inline const PrimitivePtr kPrimDynamicGRUV2 = std::make_shared<Primitive>("Dynam
 inline const PrimitivePtr kPrimDynamicGRUV2Grad = std::make_shared<Primitive>("DynamicGRUV2Grad");
 inline const PrimitivePtr kPrimScatterAdd = std::make_shared<Primitive>("ScatterAdd");
 inline const PrimitivePtr kPrimScatterUpdate = std::make_shared<Primitive>("ScatterUpdate");
+inline const PrimitivePtr kPrimMapUniform = std::make_shared<Primitive>("MapUniform");
 inline const PrimitivePtr kPrimSplit = std::make_shared<Primitive>("Split");
 inline const PrimitivePtr kPrimSequenceMask = std::make_shared<Primitive>("SequenceMask");
 

mindspore/ops/operations/__init__.py

@@ -91,7 +91,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',

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.

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