p12346870
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!25356 Add gpu kernel dynamicstitch 

Merge pull request !25356 from VectorSL/dynmic_stitch
This commit is contained in:
i-robot 2021-10-29 09:38:12 +00:00 committed by Gitee
parent 0c93832a50 2d87da9a08
commit 5f9ae3ec34
5 changed files with 288 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/dynamic_stitch_impl.cu Normal file
View File

@ -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);
}

24
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/dynamic_stitch_impl.cuh Normal file
View File

@ -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_

110
mindspore/ccsrc/backend/kernel_compiler/gpu/other/dynamic_stitch_gpu_kernel.cc Normal file
View File

@ -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

61
mindspore/ccsrc/backend/kernel_compiler/gpu/other/dynamic_stitch_gpu_kernel.h Normal file
View File

@ -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_

57
tests/st/ops/gpu/test_dynamic_stitch_op.py Normal file
View File

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