From f79b4680d139dbd44111431ee38b6315ad579aee Mon Sep 17 00:00:00 2001 From: looop5 Date: Fri, 9 Sep 2022 15:33:03 +0800 Subject: [PATCH] sparse_add, sparse_add_grad, conv_grad support dynamic rank --- .../mkldnn/conv_grad_filter_cpu_kernel.cc | 7 +-- .../mkldnn/conv_grad_input_cpu_kernel.cc | 8 ++-- .../cuda_class/sparse_add_grad_helper.h | 4 -- .../kernel/sparse/sparse_add_gpu_kernel.cc | 5 ++- mindspore/core/ops/grad/sparse_add_grad.cc | 24 +++++----- mindspore/core/ops/sparse_add.cc | 44 +++++++++---------- 6 files changed, 48 insertions(+), 44 deletions(-) diff --git a/mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/conv_grad_filter_cpu_kernel.cc b/mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/conv_grad_filter_cpu_kernel.cc index e4c687224d2..f89a7e6c1d8 100644 --- a/mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/conv_grad_filter_cpu_kernel.cc +++ b/mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/conv_grad_filter_cpu_kernel.cc @@ -40,6 +40,10 @@ void ConvGradFilterCpuKernelMod::InitKernel(const CNodePtr &kernel_node) { std::vector dst_shape = AnfAlgo::GetInputDeviceShape(kernel_node, diff_dst_index_); std::vector weight_shape = AnfAlgo::GetOutputDeviceShape(kernel_node, 0); + if (AnfAlgo::IsShapesDynamic({src_shape, weight_shape, dst_shape})) { + return; + } + size_t src_dim = src_shape.size(); if (src_dim != SHAPE_4D && src_dim != SHAPE_5D) { MS_LOG(EXCEPTION) << "Conv Grad only supports 4D/5D input, but got " << src_dim << "D!"; @@ -52,9 +56,6 @@ void ConvGradFilterCpuKernelMod::InitKernel(const CNodePtr &kernel_node) { MS_LOG(EXCEPTION) << kernel_name_ << " only supports 5D input with NCDHW format, but got fornat " << format; } - if (AnfAlgo::IsShapesDynamic({src_shape, weight_shape, dst_shape})) { - return; - } dnnl::memory::dims kernel_size(weight_shape.begin() + NC_LEN, weight_shape.end()); const size_t group = LongToSize(common::AnfAlgo::GetNodeAttr(kernel_node, GROUP)); if (group > 1) { diff --git a/mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/conv_grad_input_cpu_kernel.cc b/mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/conv_grad_input_cpu_kernel.cc index 57e9b6e6ddf..ecf72da54d4 100644 --- a/mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/conv_grad_input_cpu_kernel.cc +++ b/mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/conv_grad_input_cpu_kernel.cc @@ -40,6 +40,11 @@ void ConvGradInputCpuKernelMod::InitKernel(const CNodePtr &kernel_node) { std::vector src_shape = AnfAlgo::GetOutputDeviceShape(kernel_node, 0); std::vector weight_shape = AnfAlgo::GetInputDeviceShape(kernel_node, weight_index_); std::vector dst_shape = AnfAlgo::GetInputDeviceShape(kernel_node, diff_dst_index_); + + if (AnfAlgo::IsShapesDynamic({src_shape, weight_shape, dst_shape})) { + return; + } + size_t src_dim = src_shape.size(); if (src_dim != SHAPE_4D && src_dim != SHAPE_5D) { MS_LOG(EXCEPTION) << "Conv grad only supports 4D/5D input, but got " << src_dim << "D!"; @@ -51,9 +56,6 @@ void ConvGradInputCpuKernelMod::InitKernel(const CNodePtr &kernel_node) { if (src_dim == SHAPE_5D && format != NCDHW) { MS_LOG(EXCEPTION) << kernel_name_ << " only supports 5D input with NCDHW format, but got format " << format; } - if (AnfAlgo::IsShapesDynamic({src_shape, weight_shape, dst_shape})) { - return; - } dnnl::memory::dims kernel_size(weight_shape.begin() + NC_LEN, weight_shape.end()); const size_t group = LongToSize(common::AnfAlgo::GetNodeAttr(kernel_node, GROUP)); if (group > 1) { diff --git a/mindspore/ccsrc/plugin/device/gpu/kernel/cuda_impl/cuda_class/sparse_add_grad_helper.h b/mindspore/ccsrc/plugin/device/gpu/kernel/cuda_impl/cuda_class/sparse_add_grad_helper.h index 82d28cdeb26..9681c56037c 100644 --- a/mindspore/ccsrc/plugin/device/gpu/kernel/cuda_impl/cuda_class/sparse_add_grad_helper.h +++ b/mindspore/ccsrc/plugin/device/gpu/kernel/cuda_impl/cuda_class/sparse_add_grad_helper.h @@ -152,8 +152,6 @@ class SparseAddGradHelperGpuKernel : public GpuKernelHelperBase { int CheckKernelParam() override { size_t dim = input_shapes_.at(0).size(); if (dim != 1) { - MS_LOG(ERROR) << "For '" << kernel_name_ << "', the " << 0 << "(st/nd/rd/th) input dim should be 1," - << " but got " << dim; return -1; } @@ -161,8 +159,6 @@ class SparseAddGradHelperGpuKernel : public GpuKernelHelperBase { for (size_t i = 1; i < size; i++) { size_t dim = input_shapes_.at(i).size(); if (dim != kSparseAddGradIndicesDim) { - MS_LOG(ERROR) << "For '" << kernel_name_ << "', the " << i << "(st/nd/rd/th) input dim should be 2," - << " but got " << dim; return -1; } } diff --git a/mindspore/ccsrc/plugin/device/gpu/kernel/sparse/sparse_add_gpu_kernel.cc b/mindspore/ccsrc/plugin/device/gpu/kernel/sparse/sparse_add_gpu_kernel.cc index 10690c14f2e..382c580078c 100644 --- a/mindspore/ccsrc/plugin/device/gpu/kernel/sparse/sparse_add_gpu_kernel.cc +++ b/mindspore/ccsrc/plugin/device/gpu/kernel/sparse/sparse_add_gpu_kernel.cc @@ -103,7 +103,9 @@ int SparseAddGpuKernelMod::Resize(const BaseOperatorPtr &base_operator, const st auto dense_shape = inputs.at(kSparseAddIndex2)->GetShapeVector(); auto b_indices_shape = inputs.at(kSparseAddIndex3)->GetShapeVector(); auto b_values_shape = inputs.at(kSparseAddIndex4)->GetShapeVector(); - rank_ = a_indices_shape.at(1); + if (a_indices_shape.size() >= kDim2 && a_indices_shape.at(1) >= 0) { + rank_ = LongToSize(a_indices_shape.at(1)); + } (void)std::transform(a_indices_shape.begin(), a_indices_shape.end(), std::back_inserter(a_indices_shape_), [](int64_t x) { return x < 0 ? 0 : LongToSize(x); }); (void)std::transform(a_values_shape.begin(), a_values_shape.end(), std::back_inserter(a_values_shape_), @@ -114,7 +116,6 @@ int SparseAddGpuKernelMod::Resize(const BaseOperatorPtr &base_operator, const st [](int64_t x) { return x < 0 ? 0 : LongToSize(x); }); (void)std::transform(b_values_shape.begin(), b_values_shape.end(), std::back_inserter(b_values_shape_), [](int64_t x) { return x < 0 ? 0 : LongToSize(x); }); - // rank_ = input_shape_.size(); a_indices_size_ = std::accumulate(a_indices_shape_.begin(), a_indices_shape_.end(), 1, std::multiplies{}); a_values_size_ = std::accumulate(a_values_shape_.begin(), a_values_shape_.end(), 1, std::multiplies{}); dense_shape_size_ = std::accumulate(dense_shape_.begin(), dense_shape_.end(), 1, std::multiplies{}); diff --git a/mindspore/core/ops/grad/sparse_add_grad.cc b/mindspore/core/ops/grad/sparse_add_grad.cc index 8a834e68fda..64b06f7c46b 100644 --- a/mindspore/core/ops/grad/sparse_add_grad.cc +++ b/mindspore/core/ops/grad/sparse_add_grad.cc @@ -37,8 +37,8 @@ constexpr size_t kSparseAddGradIndex2 = 2; constexpr size_t kSparseAddGradIndex3 = 3; inline void CheckSparseAddGradShape(const size_t sparse_shape_size, const size_t expected_dim, - const std::string &arg_name, const std::string &op_name) { - if (sparse_shape_size != expected_dim) { + const std::string &arg_name, const std::string &op_name, bool is_dyn_rank) { + if (!is_dyn_rank && sparse_shape_size != expected_dim) { MS_EXCEPTION(mindspore::ValueError) << "For '" << op_name << "', " << arg_name << " must be a " << expected_dim << "-dimensional tensor, but got a " << sparse_shape_size << "-dimensional tensor."; @@ -74,7 +74,7 @@ AbstractBasePtr SparseAddGradInfer(const abstract::AnalysisEnginePtr &, const Pr const size_t kInputNum = 4; constexpr size_t kIndicesShapeSize = 2; constexpr size_t kValuesShapeSize = 1; - (void)CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, kInputNum, name); + CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, kInputNum, name); auto bvg = input_args.at(kSparseAddGradIndex0); auto x1_indices = input_args.at(kSparseAddGradIndex1); auto x2_indices = input_args.at(kSparseAddGradIndex2); @@ -90,23 +90,27 @@ AbstractBasePtr SparseAddGradInfer(const abstract::AnalysisEnginePtr &, const Pr std::shared_ptr dx1 = nullptr; std::shared_ptr dx2 = nullptr; auto val_grad_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape()); - (void)CheckSparseAddGradShape(val_grad_shape[kShape].size(), kValuesShapeSize, "backprop_val_grad", name); + auto val_grad_shape_dyn_rank = IsDynamicRank(val_grad_shape[kShape]); + CheckSparseAddGradShape(val_grad_shape[kShape].size(), kValuesShapeSize, "backprop_val_grad", name, + val_grad_shape_dyn_rank); auto x1_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kSparseAddGradIndex1]->BuildShape()); auto dx1_shape = x1_shape[kShape]; - (void)CheckSparseAddGradShape(dx1_shape.size(), kIndicesShapeSize, "x1_indices", name); + auto dx1_shape_dyn_rank = IsDynamicRank(dx1_shape); + CheckSparseAddGradShape(dx1_shape.size(), kIndicesShapeSize, "x1_indices", name, dx1_shape_dyn_rank); auto type = SparseAddGradInferType(name, input_args, 0); ShapeVector shp = {dx1_shape.at(0)}; dx1 = std::make_shared(type, std::make_shared(shp)); auto x2_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kSparseAddGradIndex2]->BuildShape()); ShapeVector dx2_shape = x2_shape[kShape]; - (void)CheckSparseAddGradShape(dx2_shape.size(), kIndicesShapeSize, "x2_indices", name); + auto dx2_shape_dyn_rank = IsDynamicRank(dx2_shape); + CheckSparseAddGradShape(dx2_shape.size(), kIndicesShapeSize, "x2_indices", name, dx2_shape_dyn_rank); shp = {dx2_shape.at(0)}; dx2 = std::make_shared(type, std::make_shared(shp)); auto sum_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kSparseAddGradIndex3]->BuildShape()); - (void)CheckSparseAddGradShape(sum_shape[kShape].size(), kIndicesShapeSize, "sum_indices", name); - if (sum_shape[kShape][0] >= 0 && val_grad_shape[kShape][0] >= 0) { - (void)CheckSparseAddGradNNZ(sum_shape[kShape][0], val_grad_shape[kShape][0], "sum_indices", "backprop_val_grad", - name); + auto sum_shape_dyn_rank = IsDynamicRank(sum_shape[kShape]); + CheckSparseAddGradShape(sum_shape[kShape].size(), kIndicesShapeSize, "sum_indices", name, sum_shape_dyn_rank); + if (!sum_shape_dyn_rank && !val_grad_shape_dyn_rank && sum_shape[kShape][0] >= 0 && val_grad_shape[kShape][0] >= 0) { + CheckSparseAddGradNNZ(sum_shape[kShape][0], val_grad_shape[kShape][0], "sum_indices", "backprop_val_grad", name); } AbstractBasePtrList ret = {dx1, dx2}; return std::make_shared(ret); diff --git a/mindspore/core/ops/sparse_add.cc b/mindspore/core/ops/sparse_add.cc index e8c2ddc8f59..b796e27e713 100644 --- a/mindspore/core/ops/sparse_add.cc +++ b/mindspore/core/ops/sparse_add.cc @@ -31,8 +31,8 @@ using mindspore::abstract::AbstractTensor; using mindspore::abstract::AbstractTuple; namespace { inline void CheckSparseAddShape(const size_t sparse_shape_size, const size_t expected_dim, const std::string &arg_name, - const std::string &op_name) { - if (sparse_shape_size != expected_dim) { + const std::string &op_name, bool is_dyn_rank) { + if (!is_dyn_rank && sparse_shape_size != expected_dim) { MS_EXCEPTION(mindspore::ValueError) << "For " << op_name << ", " << arg_name << " must be a " << expected_dim << "-dimensional tensor, but got a " << sparse_shape_size << "-dimensional tensor."; @@ -97,21 +97,28 @@ AbstractBasePtr SparseAddInfer(const abstract::AnalysisEnginePtr &, const Primit // 2-D indices auto a_indices_shape = a_indices->shape()->shape(); auto b_indices_shape = b_indices->shape()->shape(); - CheckSparseAddShape(a_indices_shape.size(), kIndicesShape, "x1_indices", op_name); - CheckSparseAddShape(b_indices_shape.size(), kIndicesShape, "x2_indices", op_name); + auto a_indices_shape_dyn_rank = IsDynamicRank(a_indices_shape); + auto b_indices_shape_dyn_rank = IsDynamicRank(b_indices_shape); + CheckSparseAddShape(a_indices_shape.size(), kIndicesShape, "x1_indices", op_name, a_indices_shape_dyn_rank); + CheckSparseAddShape(b_indices_shape.size(), kIndicesShape, "x2_indices", op_name, b_indices_shape_dyn_rank); // 1-D values auto a_values_shape = a_values->shape()->shape(); auto b_values_shape = b_values->shape()->shape(); - CheckSparseAddShape(a_values_shape.size(), 1, "x1_values", op_name); - CheckSparseAddShape(b_values_shape.size(), 1, "x2_values", op_name); + auto a_values_shape_dyn_rank = IsDynamicRank(a_values_shape); + auto b_values_shape_dyn_rank = IsDynamicRank(b_values_shape); + CheckSparseAddShape(a_values_shape.size(), 1, "x1_values", op_name, a_values_shape_dyn_rank); + CheckSparseAddShape(b_values_shape.size(), 1, "x2_values", op_name, b_values_shape_dyn_rank); // 1-D shape auto a_shape_shape = a_shape->shape()->shape(); auto b_shape_shape = b_shape->shape()->shape(); - CheckSparseAddShape(a_shape_shape.size(), 1, "x1_shape", op_name); - CheckSparseAddShape(b_shape_shape.size(), 1, "x2_shape", op_name); + auto a_shape_shape_dyn_rank = IsDynamicRank(a_shape_shape); + auto b_shape_shape_dyn_rank = IsDynamicRank(b_shape_shape); + CheckSparseAddShape(a_shape_shape.size(), 1, "x1_shape", op_name, a_shape_shape_dyn_rank); + CheckSparseAddShape(b_shape_shape.size(), 1, "x2_shape", op_name, b_shape_shape_dyn_rank); // 0-D shape auto thresh_shape = thresh->shape()->shape(); - CheckSparseAddShape(thresh_shape.size(), 0, "thresh", op_name); + auto thresh_shape_dyn_rank = IsDynamicRank(thresh_shape); + CheckSparseAddShape(thresh_shape.size(), 0, "thresh", op_name, thresh_shape_dyn_rank); // Check dtype // a_indices and b_indices should be int64 @@ -138,10 +145,10 @@ AbstractBasePtr SparseAddInfer(const abstract::AnalysisEnginePtr &, const Primit (void)CheckAndConvertUtils::CheckTensorTypeValid("thresh", thresh->BuildType(), thresh_valid_types, op_name); // Check same nnz - if (a_indices_shape[0] >= 0 && a_values_shape[0] >= 0) { + if (!a_indices_shape_dyn_rank && !a_values_shape_dyn_rank && a_indices_shape[0] >= 0 && a_values_shape[0] >= 0) { CheckSparseAddNNZ(a_indices_shape[0], a_values_shape[0], "x1_indices", "x1_values", op_name); } - if (b_indices_shape[0] >= 0 && b_values_shape[0] >= 0) { + if (!b_indices_shape_dyn_rank && !b_values_shape_dyn_rank && b_indices_shape[0] >= 0 && b_values_shape[0] >= 0) { CheckSparseAddNNZ(b_indices_shape[0], b_values_shape[0], "x2_indices", "x2_values", op_name); } // Check same type @@ -157,20 +164,13 @@ AbstractBasePtr SparseAddInfer(const abstract::AnalysisEnginePtr &, const Primit CheckSparseAddSameDtype(a_value_type, thresh_type, "x1_values", "thresh", op_name); } - int64_t max_indices_shape_ = a_indices_shape[0] + b_indices_shape[0]; - int64_t min_indices_shape_ = std::max(a_indices_shape[0], b_indices_shape[0]); ShapeVector out_indices_shape{-1, 2}; ShapeVector out_value_shape{-1}; - ShapeVector min_value_shape{min_indices_shape_}; - ShapeVector max_value_shape{max_indices_shape_}; - ShapeVector min_indices_shape{min_indices_shape_, 2}; - ShapeVector max_indices_shape{max_indices_shape_, 2}; - auto out_indices = std::make_shared( - a_indices->element()->BuildType(), - std::make_shared(out_indices_shape, min_indices_shape, max_indices_shape)); - auto out_values = std::make_shared( - a_value_type, std::make_shared(out_value_shape, min_value_shape, max_value_shape)); + auto out_indices = std::make_shared(a_indices->element()->BuildType(), + std::make_shared(out_indices_shape)); + auto out_values = + std::make_shared(a_value_type, std::make_shared(out_value_shape)); auto out_shape = std::make_shared(a_shape_type, std::make_shared(a_shape_shape));