From 06528789d42cf7e744fdb391887dd6551517e35b Mon Sep 17 00:00:00 2001 From: danishnxt Date: Wed, 11 Nov 2020 14:21:40 -0500 Subject: [PATCH] DynamicShape_Support_SegmentSum --- .../arrays/unsorted_segment_sum_gpu_kernel.cc | 62 +++++++++++++++++++ .../arrays/unsorted_segment_sum_gpu_kernel.h | 7 +++ mindspore/core/abstract/prim_arrays.cc | 51 ++++++++++----- mindspore/ops/operations/array_ops.py | 3 +- 4 files changed, 105 insertions(+), 18 deletions(-) diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/unsorted_segment_sum_gpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/unsorted_segment_sum_gpu_kernel.cc index 4be887ec79..56c029b2d6 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/unsorted_segment_sum_gpu_kernel.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/unsorted_segment_sum_gpu_kernel.cc @@ -18,6 +18,7 @@ namespace mindspore { namespace kernel { + MS_REG_GPU_KERNEL_TWO( UnsortedSegmentSum, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat32), @@ -37,5 +38,66 @@ MS_REG_GPU_KERNEL_TWO( UnsortedSegmentSum, KernelAttr().AddInputAttr(kNumberTypeInt32).AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt32), UnsortedSegmentSumGpuKernel, int, int64_t) + +// Re-registration with 3 inputs - dynamic shape mode - sets of Int64/Int32 num segments types +MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum, + KernelAttr() + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeFloat32), + UnsortedSegmentSumGpuKernel, float, int) +MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum, + KernelAttr() + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeFloat32), + UnsortedSegmentSumGpuKernel, float, int) + +MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum, + KernelAttr() + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeFloat32), + UnsortedSegmentSumGpuKernel, float, int64_t) +MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum, + KernelAttr() + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeFloat32), + UnsortedSegmentSumGpuKernel, float, int64_t) + +MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeInt32), + UnsortedSegmentSumGpuKernel, int, int) +MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeInt32), + UnsortedSegmentSumGpuKernel, int, int) + +MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeInt32), + UnsortedSegmentSumGpuKernel, int, int64_t) +MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeInt32), + UnsortedSegmentSumGpuKernel, int, int64_t) } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/unsorted_segment_sum_gpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/unsorted_segment_sum_gpu_kernel.h index 8db9db50b8..5e1fbe0993 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/unsorted_segment_sum_gpu_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/unsorted_segment_sum_gpu_kernel.h @@ -62,6 +62,13 @@ class UnsortedSegmentSumGpuKernel : public GpuKernel { auto ids_shapes = AnfAlgo::GetInputRealDeviceShapeIfExist(kernel_node, 1); auto output_shapes = AnfAlgo::GetOutputRealDeviceShapeIfExist(kernel_node, 0); + size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node); + if (input_num == 3) { + MS_LOG(INFO) << "UnsortedSegmentSum Kernel Input count is 3 - dynamic mode"; + } else { + MS_LOG(INFO) << "UnsortedSegmentSum Kernel Input count is 2"; + } + auto axis = ids_shapes.size(); for (size_t i = 0; i < input_shapes.size(); i++) { if (i < axis) { diff --git a/mindspore/core/abstract/prim_arrays.cc b/mindspore/core/abstract/prim_arrays.cc index 566dd96e6d..720dcf8f79 100644 --- a/mindspore/core/abstract/prim_arrays.cc +++ b/mindspore/core/abstract/prim_arrays.cc @@ -203,35 +203,52 @@ AbstractBasePtr InferImplUnsortedSegmentSum(const AnalysisEnginePtr &, const Pri const AbstractBasePtrList &args_spec_list) { const std::string op_name = primitive->name(); CheckArgsSize(op_name, args_spec_list, 3); + // input x auto x = CheckArg(op_name, args_spec_list, 0); MS_EXCEPTION_IF_NULL(x); MS_EXCEPTION_IF_NULL(x->shape()); auto x_shape = x->shape()->shape(); - + // segment_ids auto segment_ids = CheckArg(op_name, args_spec_list, 1); MS_EXCEPTION_IF_NULL(segment_ids); MS_EXCEPTION_IF_NULL(segment_ids->shape()); auto segment_ids_shape = segment_ids->shape()->shape(); - auto num_segments = CheckArg(op_name, args_spec_list, 2); - + // checks on Tensors 0 and 1 types + (void)CheckTensorDType(x, {kFloat32, kInt32}, "Input 0 (x) for SequenceMask should be %s"); + (void)CheckTensorDType(segment_ids, {kInt32, kInt64}, "Input 1 (segment_ids) for SequenceMask should be %s"); ShapeVector shape; - auto num_segments_value = num_segments->BuildValue(); - MS_EXCEPTION_IF_NULL(num_segments_value); - if (!num_segments_value->isa()) { - MS_LOG(WARNING) << num_segments_value << "evaluator num_segments_value should be tensor, but got " - << num_segments_value->type_name(); - shape.emplace_back(-1); + ShapeVector max_shape; + ShapeVector min_shape; + int64_t num_segments_value; + if (args_spec_list[2]->isa()) { // Num segments is Tensor + auto num_segments = args_spec_list[2]->cast(); + MS_EXCEPTION_IF_NULL(num_segments); + auto num_segments_value_ptr = num_segments->BuildValue(); + MS_EXCEPTION_IF_NULL(num_segments_value_ptr); + auto num_segments_tensor = num_segments_value_ptr->cast(); + MS_EXCEPTION_IF_NULL(num_segments_tensor); + num_segments_value = *static_cast(num_segments_tensor->data_c()); + shape.emplace_back(num_segments_value); + } else if (args_spec_list[2]->isa()) { // Num segments is Scalar + auto num_segments = CheckArg(op_name, args_spec_list, 2); + num_segments_value = GetValue(num_segments->BuildValue()); + shape.emplace_back(num_segments_value); } else { - auto num_segments_tensor = num_segments_value->cast(); - int64_t value = *(static_cast(num_segments_tensor->data_c())); - MS_LOG(INFO) << "Infer UnsortedSegmentSum output shape:" << value; - shape.emplace_back(value); + MS_LOG(EXCEPTION) << "num_segments incorrect type in UnsortedSegmentSum"; } - shape.insert(shape.end(), x_shape.begin() + segment_ids_shape.size(), x_shape.end()); - - AbstractTensorPtr ret = std::make_shared(x->element(), std::make_shared(shape)); - return ret; + // calc max shape + if (!x->shape()->max_shape().empty()) { // copy max shape from x if present + std::copy(x->shape()->max_shape().begin(), x->shape()->max_shape().end(), std::back_inserter(max_shape)); + } else { // copy x shape directly if not present + std::copy(x->shape()->shape().begin(), x->shape()->shape().end(), std::back_inserter(max_shape)); + } + // calc min shape + min_shape.push_back(segment_ids_shape.size()); + std::copy(x->shape()->shape().begin() + segment_ids_shape.size(), x->shape()->shape().end(), + back_inserter(min_shape)); + // return shape, min shape, max shape + return std::make_shared(x->element(), std::make_shared(shape, min_shape, max_shape)); } AbstractBasePtr InferImplScatterAdd(const AnalysisEnginePtr &, const PrimitivePtr &primitive, diff --git a/mindspore/ops/operations/array_ops.py b/mindspore/ops/operations/array_ops.py index 093c1fa0f3..127369c9e0 100644 --- a/mindspore/ops/operations/array_ops.py +++ b/mindspore/ops/operations/array_ops.py @@ -1560,7 +1560,8 @@ class UnsortedSegmentSum(PrimitiveWithInfer): >>> input_x = Tensor([1, 2, 3, 4], mindspore.float32) >>> segment_ids = Tensor([0, 0, 1, 2], mindspore.int32) >>> num_segments = 4 - >>> P.UnsortedSegmentSum()(input_x, segment_ids, num_segments) + >>> output = P.UnsortedSegmentSum()(input_x, segment_ids, num_segments) + >>> print(output) [3, 3, 4, 0] """