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/common/ops/populate/populate_register.h"
17 #include "nnacl/arg_min_max_parameter.h"
18 using mindspore::schema::PrimitiveType_ArgMaxFusion;
19
20 namespace mindspore {
21 namespace lite {
PopulateArgMaxParameter(const void * prim)22 OpParameter *PopulateArgMaxParameter(const void *prim) {
23 MS_CHECK_TRUE_RET(prim != nullptr, nullptr);
24 auto *arg_param = reinterpret_cast<ArgMinMaxParameter *>(malloc(sizeof(ArgMinMaxParameter)));
25 if (arg_param == nullptr) {
26 MS_LOG(ERROR) << "malloc ArgMinMaxParameter failed.";
27 return nullptr;
28 }
29 memset(arg_param, 0, sizeof(ArgMinMaxParameter));
30 auto *primitive = static_cast<const schema::Primitive *>(prim);
31 arg_param->op_parameter_.type_ = primitive->value_type();
32 auto param = primitive->value_as_ArgMaxFusion();
33 if (param == nullptr) {
34 MS_LOG(ERROR) << "param is nullptr";
35 free(arg_param);
36 return nullptr;
37 }
38 if (param->axis() > INT32_MAX) {
39 MS_LOG(ERROR) << "check le fail, value1:" << param->axis() << " value2: " << INT32_MAX;
40 free(arg_param);
41 return nullptr;
42 }
43
44 arg_param->axis_ = param->axis();
45 arg_param->topk_ = param->top_k();
46 arg_param->out_value_ = param->out_max_value();
47 arg_param->keep_dims_ = param->keep_dims();
48 return reinterpret_cast<OpParameter *>(arg_param);
49 }
50
51 REG_POPULATE(PrimitiveType_ArgMaxFusion, PopulateArgMaxParameter, SCHEMA_CUR)
52 } // namespace lite
53 } // namespace mindspore
54