diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/sort_gpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/sort_gpu_kernel.cc new file mode 100644 index 0000000000..db45754a22 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/sort_gpu_kernel.cc @@ -0,0 +1,29 @@ +/** + * 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/arrays/sort_gpu_kernel.h" + +namespace mindspore { +namespace kernel { +MS_REG_GPU_KERNEL_ONE( + Sort, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeInt32), + SortGpuKernel, half) + +MS_REG_GPU_KERNEL_ONE( + Sort, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeInt32), + SortGpuKernel, float) +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/sort_gpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/sort_gpu_kernel.h new file mode 100644 index 0000000000..6e10550e43 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/sort_gpu_kernel.h @@ -0,0 +1,227 @@ +/** + * 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_SORT_GPU_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SORT_GPU_KERNEL_H_ + +#include +#include +#include +#include +#include + +#include "backend/kernel_compiler/gpu/gpu_kernel.h" +#include "backend/kernel_compiler/gpu/gpu_kernel_factory.h" +#include "backend/kernel_compiler/gpu/cuda_impl/topk_impl.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/transpose_impl.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/unary_op_impl.cuh" + +namespace mindspore { +namespace kernel { +template +class SortGpuKernel : public GpuKernel { + public: + SortGpuKernel() { ResetResource(); } + ~SortGpuKernel() = default; + + const std::vector &GetInputSizeList() const override { return input_size_list_; } + const std::vector &GetOutputSizeList() const override { return output_size_list_; } + const std::vector &GetWorkspaceSizeList() const override { return workspace_size_list_; } + + bool Launch(const std::vector &inputs, const std::vector &workspace, + const std::vector &outputs, void *stream_ptr) override { + T *input_device = GetDeviceAddress(inputs, 0); + + T *output_device = GetDeviceAddress(outputs, 0); + int32_t *indices_device = GetDeviceAddress(outputs, 1); + + T *temp_output_device = GetDeviceAddress(workspace, 0); + int32_t *temp_indices_device = GetDeviceAddress(workspace, 1); + size_t *input_shape_device = GetDeviceAddress(workspace, 2); + size_t *perm_device = GetDeviceAddress(workspace, 3); + size_t *transposed_shape_device = GetDeviceAddress(workspace, 4); + + CHECK_CUDA_RET_WITH_EXCEPT(kernel_node_, + cudaMemcpyAsync(input_shape_device, &input_shape_[0], workspace_size_list_[2], + cudaMemcpyHostToDevice, reinterpret_cast(stream_ptr)), + "cudaMemcpyAsync for input_shape_ failed"); + + CHECK_CUDA_RET_WITH_EXCEPT(kernel_node_, + cudaMemcpyAsync(perm_device, &perm_[0], workspace_size_list_[3], cudaMemcpyHostToDevice, + reinterpret_cast(stream_ptr)), + "cudaMemcpyAsync for perm_ failed"); + + // Sort is implemented using a combination of Neg, Transpose, and TopK. It's + // Not safe to treat Transpose and TopK as inplace operators, so we alternate + // between using temp_output_device and output_device for intermediate calculations, + // this way only a constant number of allocations is needed instead of needing to + // allocate once for each intermediate calculation. + T *intermediate_input_device = input_device; + T *intermediate_output_device = output_device; + + // if sort in descending order, negate input and negate back after sorting + if (!descending_) { + Negative(intermediate_input_device, intermediate_output_device, input_size_, + reinterpret_cast(stream_ptr)); + intermediate_input_device = output_device; + intermediate_output_device = temp_output_device; + } + + // transpose so that desired dimension to sort along becomes the last one + CalTranspose(input_size_, intermediate_input_device, input_shape_device, perm_device, input_rank_, + intermediate_output_device, reinterpret_cast(stream_ptr)); + intermediate_input_device = intermediate_output_device; + intermediate_output_device = intermediate_input_device == output_device ? temp_output_device : output_device; + + // topk sorts the input along the last dimension + FastTopK(outer_size_, inner_size_, intermediate_input_device, static_cast(input_shape_[axis_]), + intermediate_output_device, temp_indices_device, topk_init_, reinterpret_cast(stream_ptr)); + std::swap(intermediate_input_device, intermediate_output_device); + + CHECK_CUDA_RET_WITH_EXCEPT(kernel_node_, + cudaMemcpyAsync(transposed_shape_device, &transposed_shape_[0], workspace_size_list_[4], + cudaMemcpyHostToDevice, reinterpret_cast(stream_ptr)), + "cudaMemcpyAsync for transposed_shape_ failed"); + + // transpose the sorted output back to the original input shape + CalTranspose(input_size_, intermediate_input_device, transposed_shape_device, perm_device, input_rank_, + intermediate_output_device, reinterpret_cast(stream_ptr)); + + // transpose the indices back to the original input shape + CalTranspose(input_size_, temp_indices_device, transposed_shape_device, perm_device, input_rank_, indices_device, + reinterpret_cast(stream_ptr)); + + // negate back the sorted values if we negated prior to sorting + if (!descending_) { + std::swap(intermediate_input_device, intermediate_output_device); + Negative(intermediate_input_device, intermediate_output_device, input_size_, + reinterpret_cast(stream_ptr)); + } + + return true; + } + + bool Init(const CNodePtr &kernel_node) override { + size_t input_count = AnfAlgo::GetInputTensorNum(kernel_node); + if (input_count != 1) { + MS_LOG(ERROR) << input_count << " inputs were provided, but SortGpuKernel expects 2."; + return false; + } + + size_t output_count = AnfAlgo::GetOutputTensorNum(kernel_node); + if (output_count != 2) { + MS_LOG(ERROR) << "Number of outputs is " << output_count << ", but should be 2 for SortGpuKernel."; + return false; + } + + input_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, 0); + + input_rank_ = input_shape_.size(); + if (input_rank_ > TRANSPOSE_MAX_DIMENSION) { + MS_LOG(ERROR) << "Sort cannot support input that has more than " << TRANSPOSE_MAX_DIMENSION << " dimensions."; + } + + input_size_ = 1; + for (size_t i = 0; i < input_rank_; i++) { + input_size_ *= input_shape_[i]; + } + + descending_ = GetAttr(kernel_node, "descending"); + axis_ = GetAttr(kernel_node, "axis"); + + if (axis_ < 0) { + axis_ += input_rank_; + } + + perm_.resize(input_rank_); + std::iota(perm_.begin(), perm_.end(), 0); + std::swap(perm_[input_rank_ - 1], perm_[axis_]); + + transposed_shape_ = input_shape_; + std::swap(transposed_shape_[input_rank_ - 1], transposed_shape_[axis_]); + + inner_size_ = input_shape_[axis_]; + outer_size_ = input_size_ / inner_size_; + + if (std::is_same::value) { + // min value representable by float16, std::numeric_limits doesn't support half + topk_init_ = static_cast(-65504.); + } else { + topk_init_ = std::numeric_limits::lowest(); + } + + InitSizeLists(); + + return true; + } + + void ResetResource() noexcept override { + input_size_ = 0; + axis_ = 0; + descending_ = false; + input_shape_.clear(); + input_rank_ = 0; + transposed_shape_.clear(); + perm_.clear(); + outer_size_ = 0; + inner_size_ = 0; + topk_init_ = static_cast(0.); + input_size_list_.clear(); + output_size_list_.clear(); + workspace_size_list_.clear(); + } + + protected: + void InitSizeLists() override { + size_t input_bytes = input_size_ * sizeof(T); + size_t indices_bytes = input_size_ * sizeof(int32_t); + input_size_list_.push_back(input_bytes); + + // outputs: sorted values, indices + output_size_list_.push_back(input_bytes); + output_size_list_.push_back(indices_bytes); + + // workspace: temp output, temp indices, input shape, perm, transposed_shape + workspace_size_list_.push_back(input_bytes); + workspace_size_list_.push_back(indices_bytes); + workspace_size_list_.push_back(input_rank_ * sizeof(size_t)); + workspace_size_list_.push_back(input_rank_ * sizeof(size_t)); + workspace_size_list_.push_back(input_rank_ * sizeof(size_t)); + } + + private: + size_t input_size_; + int64_t axis_; + bool descending_; + std::vector input_shape_; + size_t input_rank_; + + // for transpose + std::vector transposed_shape_; + std::vector perm_; + + // for topk + size_t outer_size_; + size_t inner_size_; + T topk_init_; + + std::vector input_size_list_; + std::vector output_size_list_; + std::vector workspace_size_list_; +}; +} // namespace kernel +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SORT_GPU_KERNEL_H_ diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.cc index 59503128e9..3066c89ebd 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.cc @@ -1,5 +1,5 @@ /** - * Copyright 2020 Huawei Technologies Co., Ltd + * Copyright 2020-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. @@ -25,5 +25,13 @@ MS_REG_GPU_KERNEL_TWO(TopK, .AddOutputAttr(kNumberTypeFloat32) .AddOutputAttr(kNumberTypeInt32), TopKGpuKernel, float, int) -} + +MS_REG_GPU_KERNEL_TWO(TopK, + KernelAttr() + .AddInputAttr(kNumberTypeFloat16) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeFloat16) + .AddOutputAttr(kNumberTypeInt32), + TopKGpuKernel, half, int) +} // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cu index 5a57405260..2bd92288cc 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cu @@ -221,5 +221,7 @@ void FastTopK(const int outer_size, const int inner_size, const T *input, S k_cu } } +template void FastTopK(const int outer_size, const int inner_size, const half *input, int k_cut, half *output, + int *output_index, const half init_K, cudaStream_t stream); template void FastTopK(const int outer_size, const int inner_size, const float *input, int k_cut, float *output, int *output_index, const float init_K, cudaStream_t stream); diff --git a/mindspore/core/abstract/infer_functions.h b/mindspore/core/abstract/infer_functions.h index 6608fa6152..709a372944 100644 --- a/mindspore/core/abstract/infer_functions.h +++ b/mindspore/core/abstract/infer_functions.h @@ -292,6 +292,8 @@ AbstractBasePtr InferImplLoad(const AnalysisEnginePtr &, const PrimitivePtr &pri const AbstractBasePtrList &args_spec_list); AbstractBasePtr InferImplAssign(const AnalysisEnginePtr &, const PrimitivePtr &primitive, const AbstractBasePtrList &args_spec_list); +AbstractBasePtr InferImplSort(const AnalysisEnginePtr &, const PrimitivePtr &primitive, + const AbstractBasePtrList &args_spec_list); template AbstractBasePtr InferTupleOrListOrDictLen(const std::string &op_name, const AbstractBasePtrList &args_spec_list) { // Inputs: a tuple or list or dict. diff --git a/mindspore/core/abstract/prim_arrays.cc b/mindspore/core/abstract/prim_arrays.cc index cc5ff66887..bcfe1b5037 100644 --- a/mindspore/core/abstract/prim_arrays.cc +++ b/mindspore/core/abstract/prim_arrays.cc @@ -1132,5 +1132,29 @@ AbstractBasePtr InferImplArgMaxWithValue(const AnalysisEnginePtr &, const Primit AbstractBasePtrList result = {index, value}; return std::make_shared(result); } + +AbstractBasePtr InferImplSort(const AnalysisEnginePtr &, const PrimitivePtr &primitive, + const AbstractBasePtrList &args_spec_list) { + const std::string &op_name = primitive->name(); + CheckArgsSize(op_name, args_spec_list, 1); + AbstractTensorPtr input = CheckArg(op_name, args_spec_list, 0); + + TypePtrList supported_types = {kFloat16, kFloat32}; + (void)CheckTensorDType(input, supported_types, "input for Sort should be %s"); + + ValuePtr axis_ptr = primitive->GetAttr("axis"); + int64_t axis = GetValue(axis_ptr); + int64_t input_rank = input->shape()->shape().size(); + if (!(axis >= -input_rank && axis < input_rank)) { + MS_LOG(EXCEPTION) << "axis is not in the valid range [" << -input_rank << ", " << input_rank << ")."; + } + + auto sorted_values = std::make_shared(input->element(), input->shape()); + TypePtr idx_type = kInt32; + auto indices = std::make_shared(idx_type, input->shape()); + AbstractBasePtrList result = {sorted_values, indices}; + return std::make_shared(result); +} + } // namespace abstract } // namespace mindspore diff --git a/mindspore/core/abstract/primitive_infer_map.cc b/mindspore/core/abstract/primitive_infer_map.cc index 46aa76460e..93b7cdcbb6 100644 --- a/mindspore/core/abstract/primitive_infer_map.cc +++ b/mindspore/core/abstract/primitive_infer_map.cc @@ -94,6 +94,7 @@ PrimitiveEvalImplMap &GetPrimitiveToEvalImplMap() { {prim::kPrimMapUniform, {InferImplMapUniform, nullptr, true}}, {prim::kPrimSplit, {InferImplSplit, nullptr, true}}, {prim::kPrimSequenceMask, {InferImplSequenceMask, nullptr, true}}, + {prim::kPrimSort, {InferImplSort, nullptr, true}}, // Structure {prim::kPrimMakeTuple, {InferImplMakeTuple, nullptr, true}}, {prim::kPrimMakeList, {InferImplMakeList, nullptr, true}}, diff --git a/mindspore/core/base/core_ops.h b/mindspore/core/base/core_ops.h index 82c9bdc514..fcdc7eaaef 100644 --- a/mindspore/core/base/core_ops.h +++ b/mindspore/core/base/core_ops.h @@ -189,6 +189,7 @@ inline const PrimitivePtr kPrimReverseSequence = std::make_shared("Re inline const PrimitivePtr kPrimRank = std::make_shared("Rank"); inline const PrimitivePtr kPrimResizeBilinear = std::make_shared("ResizeBilinear"); inline const PrimitivePtr kPrimResizeGrad = std::make_shared("ResizeGrad"); +inline const PrimitivePtr kPrimSort = std::make_shared("Sort"); // NN inline const PrimitivePtr kPrimAdam = std::make_shared("Adam"); diff --git a/tests/st/ops/gpu/test_sort.py b/tests/st/ops/gpu/test_sort.py new file mode 100644 index 0000000000..30c1b47222 --- /dev/null +++ b/tests/st/ops/gpu/test_sort.py @@ -0,0 +1,160 @@ +# 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.context as context +import mindspore.nn as nn +from mindspore import Tensor +from mindspore.ops import operations as P + +class SortNet(nn.Cell): + def __init__(self, axis, descending): + super(SortNet, self).__init__() + self.sort = P.Sort(axis, descending) + + def construct(self, x): + return self.sort(x) + + +def sort_1d(descending, nptype): + context.set_context(mode=context.GRAPH_MODE, device_target="GPU") + + x_numpy = np.array([1, -2, 3, 4]).astype(np.float16) + x = Tensor(x_numpy) + sort_net = SortNet(0, descending) + output, indices = sort_net(x) + + expected_output = np.sort(x_numpy, 0) + expected_indices = np.array([1, 0, 2, 3]) + if descending: + expected_output = expected_output[::-1] + expected_indices = expected_indices[::-1] + + np.testing.assert_array_equal(output.asnumpy(), expected_output) + np.testing.assert_array_equal(indices.asnumpy(), expected_indices) + +def sort_3d(descending, nptype): + context.set_context(mode=context.GRAPH_MODE, device_target="GPU") + + x_numpy = np.array([[[1, 2, 3, 4], + [8, 7, 2, 0], + [9, 4, 1, 8]], + [[5, 4, 1, 8], + [2, 9, 0, 7], + [6, 1, 7, 4]]]).astype(nptype) + x = Tensor(x_numpy) + + axis = -1 + sort_net = SortNet(axis, descending) + output, indices = sort_net(x) + + expected_output = np.sort(x_numpy, axis) + expected_indices = np.array([[[0, 1, 2, 3], + [3, 2, 1, 0], + [2, 1, 3, 0]], + [[2, 1, 0, 3], + [2, 0, 3, 1], + [1, 3, 0, 2]]]) + if descending: + expected_output = expected_output[:, :, ::-1] + expected_indices = expected_indices[:, :, ::-1] + + np.testing.assert_array_equal(output.asnumpy(), expected_output) + np.testing.assert_array_equal(indices.asnumpy(), expected_indices) + + axis = 1 + sort_net = SortNet(axis, descending) + output, indices = sort_net(x) + + expected_output = np.sort(x_numpy, axis) + expected_indices = np.array([[[0, 0, 2, 1], + [1, 2, 1, 0], + [2, 1, 0, 2]], + [[1, 2, 1, 2], + [0, 0, 0, 1], + [2, 1, 2, 0]]]) + if descending: + expected_output = expected_output[:, ::-1, :] + expected_indices = expected_indices[:, ::-1, :] + + np.testing.assert_array_equal(output.asnumpy(), expected_output) + np.testing.assert_array_equal(indices.asnumpy(), expected_indices) + + axis = -3 + sort_net = SortNet(axis, descending) + output, indices = sort_net(x) + + expected_output = np.sort(x_numpy, axis) + expected_indices = np.array([[[0, 0, 1, 0], + [1, 0, 1, 0], + [1, 1, 0, 1]], + [[1, 1, 0, 1], + [0, 1, 0, 1], + [0, 0, 1, 0]]]) + if descending: + expected_output = expected_output[::-1, :, :] + expected_indices = expected_indices[::-1, :, :] + + np.testing.assert_array_equal(output.asnumpy(), expected_output) + np.testing.assert_array_equal(indices.asnumpy(), expected_indices) + +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_sort1d_float16(): + sort_1d(False, np.float16) + +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_sort1d_descending_float16(): + sort_1d(True, np.float16) + +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_sort1d_float32(): + sort_1d(False, np.float32) + +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_sort1d_descending_float32(): + sort_1d(True, np.float32) + +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_sort3d_float16(): + sort_3d(False, np.float16) + +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_sort3d_descending_float16(): + sort_3d(True, np.float16) + +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_sort3d_float32(): + sort_3d(False, np.float32) + +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_sort3d_descending_float32(): + sort_3d(True, np.float32)