!25356 Add gpu kernel dynamicstitch

Merge pull request !25356 from VectorSL/dynmic_stitch

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/dynamic_stitch_impl.cu

@@ -0,0 +1,36 @@
+/**
+ * Copyright 2021 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 "dynamic_stitch_impl.cuh"
+#include "runtime/device/gpu/cuda_common.h"
+
+__global__ void StitchKernel(const int *index_addr, const unsigned char *data_addr, unsigned char *output_addr,
+                             const size_t index_num, const size_t data_size, int *max_index_dev) {
+  for (size_t i = 0; i < index_num; i++) {
+    int index = index_addr[i];
+    if (max_index_dev[0] < index) {
+      max_index_dev[0] = index;
+    }
+    for (size_t j = blockIdx.x * blockDim.x + threadIdx.x; j < data_size; j += gridDim.x * blockDim.x) {
+      output_addr[index * data_size + j] = data_addr[i * data_size + j];
+    }
+  }
+}
+
+void CallStitch(const int *index_addr, const unsigned char *data_addr, unsigned char *output_addr,
+                const size_t index_num, const size_t data_size, int *max_index_dev, cudaStream_t cuda_stream) {
+  StitchKernel<<<GET_BLOCKS(data_size), GET_THREADS, 0, cuda_stream>>>(index_addr, data_addr, output_addr, index_num,
+                                                                       data_size, max_index_dev);
+}

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/dynamic_stitch_impl.cuh

@@ -0,0 +1,24 @@
+/**
+ * Copyright 2021 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_GPU_CUDA_IMPL_DYNAMIC_STITCH_CUH_
+#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_CUDA_IMPL_DYNAMIC_STITCH_CUH_
+
+#include <cuda_runtime.h>
+
+void CallStitch(const int *index_addr, const unsigned char *data_addr, unsigned char *output_addr,
+                const size_t index_num, const size_t data_size, int *max_index_dev, cudaStream_t cuda_stream);
+#endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_CUDA_IMPL_DYNAMIC_STITCH_CUH_

mindspore/ccsrc/backend/kernel_compiler/gpu/other/dynamic_stitch_gpu_kernel.cc

@@ -0,0 +1,110 @@
+/**
+ * Copyright 2021 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/gpu/other/dynamic_stitch_gpu_kernel.h"
+#include <functional>
+#include "backend/kernel_compiler/common_utils.h"
+#include "backend/kernel_compiler/gpu/cuda_impl/dynamic_stitch_impl.cuh"
+#include "runtime/device/gpu/gpu_common.h"
+
+namespace mindspore {
+namespace kernel {
+DynamicStitchKernel::DynamicStitchKernel()
+    : n_(0), real_ele_num_(0), max_index_(0), one_data_ele_num_(0), data_type_size_(0) {
+  ResetResource();
+}
+
+DynamicStitchKernel::~DynamicStitchKernel() {}
+
+const std::vector<size_t> &DynamicStitchKernel::GetInputSizeList() const { return input_size_list_; }
+
+const std::vector<size_t> &DynamicStitchKernel::GetOutputSizeList() const { return output_size_list_; }
+
+const std::vector<size_t> &DynamicStitchKernel::GetWorkspaceSizeList() const { return workspace_size_list_; }
+
+bool DynamicStitchKernel::Init(const CNodePtr &kernel_node) {
+  kernel_node_ = kernel_node;
+  // Inputs: (indexlist, datalist)
+  size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node);
+  n_ = input_num / kDivNum2;
+
+  auto output_shape = AnfAlgo::GetOutputRealDeviceShapeIfExist(kernel_node, 0);
+  auto data_type = AnfAlgo::GetInputDeviceDataType(kernel_node, n_);
+  // Index type is restricted to int32 by kernel prim.
+  size_t index_type_size = sizeof(int);
+  data_type_size_ = GetDtypeNbyte(TypeIdToString(data_type, false));
+  auto first_data_shape = AnfAlgo::GetInputRealDeviceShapeIfExist(kernel_node, n_);
+  one_data_ele_num_ = first_data_shape[first_data_shape.size() - 1];
+  for (size_t i = 0; i < n_; i++) {
+    auto data_shape = AnfAlgo::GetInputRealDeviceShapeIfExist(kernel_node, n_ + i);
+    size_t data_size = std::accumulate(data_shape.begin(), data_shape.end(), 1, std::multiplies<size_t>());
+    //  Data size
+    input_size_list_.push_back(data_size * data_type_size_);
+    // Index size
+    input_size_list_.insert(input_size_list_.begin() + i, data_size / one_data_ele_num_ * index_type_size);
+  }
+  size_t output_size =
+    std::accumulate(output_shape.begin(), output_shape.end(), data_type_size_, std::multiplies<size_t>());
+
+  // For max_index
+  workspace_size_list_.push_back(index_type_size);
+  // One output
+  output_size_list_.push_back(output_size);
+  return true;
+}
+
+void DynamicStitchKernel::PostExecute() {
+  auto output_shape = AnfAlgo::GetOutputRealDeviceShapeIfExist(kernel_node_.lock(), 0);
+  output_shape[0] = max_index_ + 1;
+  auto data_type = AnfAlgo::GetInputDeviceDataType(kernel_node_.lock(), n_);
+  AnfAlgo::SetOutputInferTypeAndShape({data_type}, {output_shape}, kernel_node_.lock().get());
+  MS_LOG(DEBUG) << "Run PostExecute for dynamicstitch, real output shape is " << output_shape;
+}
+
+void DynamicStitchKernel::ResetResource() noexcept {
+  input_size_list_.clear();
+  output_size_list_.clear();
+  workspace_size_list_.clear();
+}
+
+void DynamicStitchKernel::InitSizeLists() { return; }
+
+bool DynamicStitchKernel::Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
+                                 const std::vector<AddressPtr> &outputs, void *stream) {
+  auto cuda_stream = reinterpret_cast<cudaStream_t>(stream);
+  auto max_index_dev = GetDeviceAddress<int>(workspace, 0);
+  auto output_addr = GetDeviceAddress<unsigned char>(outputs, 0);
+  // Init output  and max_index with 0
+  CHECK_CUDA_RET_WITH_EXCEPT(kernel_node_, cudaMemsetAsync(output_addr, 0, output_size_list_[0], cuda_stream),
+                             "Init output with cudamemset failed");
+  CHECK_CUDA_RET_WITH_EXCEPT(kernel_node_, cudaMemsetAsync(max_index_dev, 0, sizeof(int), cuda_stream),
+                             "Init max_index_dev with cudamemset failed");
+
+  for (size_t i = 0; i < n_; i++) {
+    auto index_addr = GetDeviceAddress<int>(inputs, i);
+    auto data_addr = GetDeviceAddress<unsigned char>(inputs, n_ + i);
+    size_t index_num = input_size_list_[i] / sizeof(int);
+    CallStitch(index_addr, data_addr, output_addr, index_num, one_data_ele_num_ * data_type_size_, max_index_dev,
+               cuda_stream);
+  }
+  int temp = 0;
+  CHECK_CUDA_RET_WITH_EXCEPT(kernel_node_,
+                             cudaMemcpyAsync(&temp, max_index_dev, sizeof(int), cudaMemcpyDeviceToHost, cuda_stream),
+                             "Copy max_index failed")
+  max_index_ = IntToSize(temp);
+  return true;
+}
+}  // namespace kernel
+}  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/gpu/other/dynamic_stitch_gpu_kernel.h

@@ -0,0 +1,61 @@
+/**
+ * Copyright 2021 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_GPU_OTHER_DYNAMIC_STITCH_GPU_KERNEL_H_
+#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_OTHER_DYNAMIC_STITCH_GPU_KERNEL_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+#include "backend/kernel_compiler/gpu/gpu_kernel.h"
+#include "backend/kernel_compiler/gpu/gpu_kernel_factory.h"
+
+namespace mindspore {
+namespace kernel {
+class DynamicStitchKernel : public GpuKernel {
+ public:
+  DynamicStitchKernel();
+  ~DynamicStitchKernel();
+
+  const std::vector<size_t> &GetInputSizeList() const override;
+  const std::vector<size_t> &GetOutputSizeList() const override;
+  const std::vector<size_t> &GetWorkspaceSizeList() const override;
+  bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
+              const std::vector<AddressPtr> &outputs, void *stream_ptr) override;
+  bool Init(const CNodePtr &kernel_node) override;
+  void ResetResource() noexcept override;
+  void PostExecute() override;
+  constexpr static size_t kDivNum2 = 2;
+
+ protected:
+  void InitSizeLists() override;
+
+ private:
+  size_t n_;
+  size_t real_ele_num_;
+  size_t max_index_;
+  size_t one_data_ele_num_;
+  size_t data_type_size_;
+  std::vector<size_t> input_size_list_;
+  std::vector<size_t> output_size_list_;
+  std::vector<size_t> workspace_size_list_;
+};
+
+MS_REG_GPU_KERNEL(DynamicStitch, DynamicStitchKernel)
+}  // namespace kernel
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_OTHER_DYNAMIC_STITCH_GPU_KERNEL_H_

tests/st/ops/gpu/test_dynamic_stitch_op.py

@@ -0,0 +1,57 @@
+# Copyright 2021 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
+import mindspore.context as context
+import mindspore.nn as nn
+from mindspore import Tensor
+from mindspore.ops.operations import _inner_ops as ops
+
+context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
+
+
+class Net(nn.Cell):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.stitch = ops.DynamicStitch()
+
+    def construct(self, indices, data):
+        return self.stitch(indices, data)
+
+@pytest.mark.level1
+@pytest.mark.platform_x86_gpu_training
+@pytest.mark.env_onecard
+def test_net_int32():
+    """
+    Feature: ALL TO ALL
+    Description: test cases for dynamicstitch.
+    Expectation: the result match expected array.
+    """
+    x1 = Tensor([6], mindspore.int32)
+    x2 = Tensor(np.array([4, 1]), mindspore.int32)
+    x3 = Tensor(np.array([[5, 2], [0, 3]]), mindspore.int32)
+    y1 = Tensor(np.array([[61, 62]]), mindspore.int32)
+    y2 = Tensor(np.array([[41, 42], [11, 12]]), mindspore.int32)
+    y3 = Tensor(np.array([[[51, 52], [21, 22]], [[1, 2], [31, 32]]]), mindspore.int32)
+    expected = np.array([[1, 2], [11, 12], [21, 22],
+                         [31, 32], [41, 42], [51, 52], [61, 62]]).astype(np.int32)
+
+    indices = [x1, x2, x3]
+    data = [y1, y2, y3]
+    net = Net()
+    output = net(indices, data)
+    assert np.array_equal(output.asnumpy(), expected)