• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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/squeeze_parameter.h"
18 using mindspore::schema::PrimitiveType_Squeeze;
19 
20 namespace mindspore {
21 namespace lite {
PopulateSqueezeParameter(const void * prim)22 OpParameter *PopulateSqueezeParameter(const void *prim) {
23   auto *primitive = static_cast<const schema::Primitive *>(prim);
24   MS_ASSERT(primitive != nullptr);
25   auto value = primitive->value_as_Squeeze();
26   if (value == nullptr) {
27     MS_LOG(ERROR) << "value is nullptr";
28     return nullptr;
29   }
30 
31   auto *param = reinterpret_cast<SqueezeParameter *>(malloc(sizeof(SqueezeParameter)));
32   if (param == nullptr) {
33     MS_LOG(ERROR) << "malloc SqueezeParameter failed.";
34     return nullptr;
35   }
36   memset(param, 0, sizeof(SqueezeParameter));
37 
38   param->op_parameter_.type_ = primitive->value_type();
39   auto axis = value->axis();
40   if (axis != nullptr) {
41     param->axis_size_ = axis->size();
42     if (param->axis_size_ > MAX_SHAPE_SIZE) {
43       MS_LOG(ERROR) << "Invalid axis size " << param->axis_size_;
44       free(param);
45       return nullptr;
46     }
47     for (size_t i = 0; i < param->axis_size_; i++) {
48       param->axis_[i] = *(axis->begin() + i);
49     }
50   } else {
51     param->axis_size_ = 0;
52   }
53   return reinterpret_cast<OpParameter *>(param);
54 }
55 
56 REG_POPULATE(PrimitiveType_Squeeze, PopulateSqueezeParameter, SCHEMA_CUR)
57 }  // namespace lite
58 }  // namespace mindspore
59