!41760 sparse_add, sparse_add_grad, conv_grad support dynamic rank

Merge pull request !41760 from looop5/sparse_add_dyn_rank

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<int64_t> dst_shape = AnfAlgo::GetInputDeviceShape(kernel_node, diff_dst_index_);
   std::vector<int64_t> 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<int64_t>(kernel_node, GROUP));
   if (group > 1) {

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<int64_t> src_shape = AnfAlgo::GetOutputDeviceShape(kernel_node, 0);
   std::vector<int64_t> weight_shape = AnfAlgo::GetInputDeviceShape(kernel_node, weight_index_);
   std::vector<int64_t> 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<int64_t>(kernel_node, GROUP));
   if (group > 1) {

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

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

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) {
+                                    const std::string &arg_name, const std::string &op_name, bool is_dyn_rank) {
-  if (sparse_shape_size != expected_dim) {
+  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<AbstractTensor> dx1 = nullptr;
   std::shared_ptr<AbstractTensor> 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<AbstractTensor>(type, std::make_shared<mindspore::abstract::Shape>(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<AbstractTensor>(type, std::make_shared<mindspore::abstract::Shape>(shp));
   auto sum_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kSparseAddGradIndex3]->BuildShape());
-  (void)CheckSparseAddGradShape(sum_shape[kShape].size(), kIndicesShapeSize, "sum_indices", name);
+  auto sum_shape_dyn_rank = IsDynamicRank(sum_shape[kShape]);
-  if (sum_shape[kShape][0] >= 0 && val_grad_shape[kShape][0] >= 0) {
+  CheckSparseAddGradShape(sum_shape[kShape].size(), kIndicesShapeSize, "sum_indices", name, sum_shape_dyn_rank);
-    (void)CheckSparseAddGradNNZ(sum_shape[kShape][0], val_grad_shape[kShape][0], "sum_indices", "backprop_val_grad",
+  if (!sum_shape_dyn_rank && !val_grad_shape_dyn_rank && sum_shape[kShape][0] >= 0 && val_grad_shape[kShape][0] >= 0) {
-                                name);
+    CheckSparseAddGradNNZ(sum_shape[kShape][0], val_grad_shape[kShape][0], "sum_indices", "backprop_val_grad", name);
   }
   AbstractBasePtrList ret = {dx1, dx2};
   return std::make_shared<AbstractTuple>(ret);

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) {
+                                const std::string &op_name, bool is_dyn_rank) {
-  if (sparse_shape_size != expected_dim) {
+  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);
+  auto a_indices_shape_dyn_rank = IsDynamicRank(a_indices_shape);
-  CheckSparseAddShape(b_indices_shape.size(), kIndicesShape, "x2_indices", op_name);
+  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);
+  auto a_values_shape_dyn_rank = IsDynamicRank(a_values_shape);
-  CheckSparseAddShape(b_values_shape.size(), 1, "x2_values", op_name);
+  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);
+  auto a_shape_shape_dyn_rank = IsDynamicRank(a_shape_shape);
-  CheckSparseAddShape(b_shape_shape.size(), 1, "x2_shape", op_name);
+  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<AbstractTensor>(
+  auto out_indices = std::make_shared<AbstractTensor>(a_indices->element()->BuildType(),
-    a_indices->element()->BuildType(),
+                                                      std::make_shared<mindspore::abstract::Shape>(out_indices_shape));
-    std::make_shared<mindspore::abstract::Shape>(out_indices_shape, min_indices_shape, max_indices_shape));
+  auto out_values =
-  auto out_values = std::make_shared<AbstractTensor>(
+    std::make_shared<AbstractTensor>(a_value_type, std::make_shared<mindspore::abstract::Shape>(out_value_shape));
-    a_value_type, std::make_shared<mindspore::abstract::Shape>(out_value_shape, min_value_shape, max_value_shape));
   auto out_shape =
     std::make_shared<AbstractTensor>(a_shape_type, std::make_shared<mindspore::abstract::Shape>(a_shape_shape));