From 6a226084018ad7088401d6f500ae40c81a06ba7e Mon Sep 17 00:00:00 2001 From: liuwenhao4 Date: Thu, 24 Sep 2020 16:14:19 +0800 Subject: [PATCH] Add the function that split size can be -1 at the end of split ops --- mindspore/lite/nnacl/split.c | 8 -------- mindspore/lite/src/runtime/kernel/arm/base/split_base.cc | 8 ++++++++ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/mindspore/lite/nnacl/split.c b/mindspore/lite/nnacl/split.c index dfa7ad8dd0..03e7f1eddf 100644 --- a/mindspore/lite/nnacl/split.c +++ b/mindspore/lite/nnacl/split.c @@ -50,14 +50,6 @@ int DoSplit(float *in_data, float **out_data, const int *input_shape, int offset split_which = i % num_split; split_times = i / num_split; int split_size = split_sizes[split_which]; - // support split size is -1 in the end. - if (split_size == -1) { - int split_dim_i = input_shape[split_dim]; - for (int j = 0; j < num_split - 1; ++j) { - split_dim_i -= split_sizes[j]; - } - split_size = split_dim_i; - } float *dst = out_data[split_which] + split_times * in_stride * split_size; (void)memcpy(dst, src, split_size * in_stride_bytes); src += split_size * in_stride; diff --git a/mindspore/lite/src/runtime/kernel/arm/base/split_base.cc b/mindspore/lite/src/runtime/kernel/arm/base/split_base.cc index 9c9d71c39c..89d9558bc8 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/split_base.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/split_base.cc @@ -58,6 +58,14 @@ int SplitBaseCPUKernel::ReSize() { } } + if (param->split_sizes_[param->num_split_ - 1] == -1) { + int split_shape_end = input_shape[param->split_dim_]; + for (int i = 0; i < param->num_split_ - 1; i++) { + split_shape_end -= param->split_sizes_[i]; + } + param->split_sizes_[param->num_split_ - 1] = split_shape_end; + } + num_unit_ = param->split_count_ * param->num_split_; thread_n_num_ = MSMIN(thread_count_, num_unit_); thread_n_stride_ = UP_DIV(num_unit_, thread_n_num_);