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

!41760 sparse_add, sparse_add_grad, conv_grad support dynamic rank 

Merge pull request !41760 from looop5/sparse_add_dyn_rank
This commit is contained in:
i-robot 2022-09-13 06:09:14 +00:00 committed by Gitee
parent 25c7f3a737 f79b4680d1
commit 691b7de54d
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
6 changed files with 48 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/conv_grad_filter_cpu_kernel.cc
View File

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

8
mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/conv_grad_input_cpu_kernel.cc
View File

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

4
mindspore/ccsrc/plugin/device/gpu/kernel/cuda_impl/cuda_class/sparse_add_grad_helper.h
View File

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

5
mindspore/ccsrc/plugin/device/gpu/kernel/sparse/sparse_add_gpu_kernel.cc
View File

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

24
mindspore/core/ops/grad/sparse_add_grad.cc
View File

@ -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<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);
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<AbstractTuple>(ret);

44
mindspore/core/ops/sparse_add.cc
View File

@ -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<AbstractTensor>(
a_indices->element()->BuildType(),
std::make_shared<mindspore::abstract::Shape>(out_indices_shape, min_indices_shape, max_indices_shape));
auto out_values = std::make_shared<AbstractTensor>(
a_value_type, std::make_shared<mindspore::abstract::Shape>(out_value_shape, min_value_shape, max_value_shape));
auto out_indices = std::make_shared<AbstractTensor>(a_indices->element()->BuildType(),
std::make_shared<mindspore::abstract::Shape>(out_indices_shape));
auto out_values =
std::make_shared<AbstractTensor>(a_value_type, std::make_shared<mindspore::abstract::Shape>(out_value_shape));
auto out_shape =
std::make_shared<AbstractTensor>(a_shape_type, std::make_shared<mindspore::abstract::Shape>(a_shape_shape));