1 /**
2 * Copyright 2019-2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16 #include "src/ops/populate/populate_register.h"
17 #include "nnacl/split_parameter.h"
18 #include "nnacl/op_base.h"
19 using mindspore::schema::PrimitiveType_Split;
20
21 namespace mindspore {
22 namespace lite {
DestroySplitParameter(OpParameter * parameter)23 void DestroySplitParameter(OpParameter *parameter) {
24 MS_CHECK_PTR_IF_NULL(parameter);
25 auto param = reinterpret_cast<SplitParameter *>(parameter);
26 if (param->split_sizes_ != nullptr) {
27 free(param->split_sizes_);
28 param->split_sizes_ = nullptr;
29 }
30 }
31
PopulateSplitParameter(const void * prim)32 OpParameter *PopulateSplitParameter(const void *prim) {
33 MS_CHECK_TRUE_RET(prim != nullptr, nullptr);
34 auto primitive = static_cast<const schema::Primitive *>(prim);
35 auto value = primitive->value_as_Split();
36 if (value == nullptr) {
37 MS_LOG(ERROR) << "value is nullptr";
38 return nullptr;
39 }
40
41 auto *param = reinterpret_cast<SplitParameter *>(malloc(sizeof(SplitParameter)));
42 if (param == nullptr) {
43 MS_LOG(ERROR) << "malloc SplitParameter failed.";
44 return nullptr;
45 }
46 memset(param, 0, sizeof(SplitParameter));
47
48 param->op_parameter_.type_ = primitive->value_type();
49 param->num_split_ = value->output_num();
50 if (param->num_split_ > std::numeric_limits<int>::max() / static_cast<int>(sizeof(int)) || param->num_split_ <= 0) {
51 MS_LOG(ERROR) << "The value of param->num_split_ is not correct";
52 free(param);
53 return nullptr;
54 }
55
56 /* free split_sizes_ in split op base */
57 param->split_sizes_ = reinterpret_cast<int *>(malloc(static_cast<size_t>(param->num_split_) * sizeof(int)));
58 if (param->split_sizes_ == nullptr) {
59 MS_LOG(ERROR) << "malloc param split_sizes_ error";
60 free(param);
61 return nullptr;
62 }
63 param->op_parameter_.destroy_func_ = DestroySplitParameter;
64 memset(param->split_sizes_, 0, static_cast<size_t>(param->num_split_) * sizeof(int));
65 auto split_sizes_vector_ = value->size_splits();
66 if (split_sizes_vector_ != nullptr && split_sizes_vector_->size() <= static_cast<uint32_t>(param->num_split_)) {
67 int i = 0;
68 for (auto iter : *split_sizes_vector_) {
69 param->split_sizes_[i++] = iter;
70 }
71 param->split_count_ = param->num_split_;
72 } else {
73 param->split_count_ = 0;
74 }
75 param->split_dim_ = value->axis();
76 return reinterpret_cast<OpParameter *>(param);
77 }
78
79 REG_POPULATE(PrimitiveType_Split, PopulateSplitParameter, SCHEMA_CUR)
80 } // namespace lite
81 } // namespace mindspore
82