1 /*
2 * Copyright (c) 2022 Huawei Device Co., Ltd.
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 * http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15
16 #include "argmax_builder.h"
17
18 namespace OHOS {
19 namespace NeuralNetworkRuntime {
20 namespace Ops {
21 static const int INPUT_NUM = 1;
22 static const int OUTPUT_NUM = 1;
23 static const std::string OP_NAME = "ArgMax";
24
ArgMaxBuilder()25 ArgMaxBuilder::ArgMaxBuilder() {}
26
~ArgMaxBuilder()27 ArgMaxBuilder::~ArgMaxBuilder() {}
28
SetAxis(std::shared_ptr<NNTensor> tensor)29 OH_NN_ReturnCode ArgMaxBuilder::SetAxis(std::shared_ptr<NNTensor> tensor)
30 {
31 tensor->IdentifyOpParameter();
32
33 if (tensor->GetDataType() != OH_NN_INT64) {
34 LOGE("[ArgMax] SetAxis failed, the axis should be type HNN_INT64.");
35 return OH_NN_INVALID_PARAMETER;
36 }
37
38 void* buffer = tensor->GetBuffer();
39 if (buffer == nullptr) {
40 LOGE("[ArgMax] SetAxis GetBuffer return nullptr.");
41 return OH_NN_INVALID_PARAMETER;
42 }
43
44 m_axis = *(static_cast<int64_t*>(buffer));
45 return OH_NN_SUCCESS;
46 }
47
SetKeepdims(std::shared_ptr<NNTensor> tensor)48 OH_NN_ReturnCode ArgMaxBuilder::SetKeepdims(std::shared_ptr<NNTensor> tensor)
49 {
50 tensor->IdentifyOpParameter();
51
52 if (tensor->GetDataType() != OH_NN_BOOL) {
53 LOGE("[ArgMax] SetKeepdims failed, the keep_dims should be type HNN_BOOL.");
54 return OH_NN_INVALID_PARAMETER;
55 }
56
57 void* buffer = tensor->GetBuffer();
58 if (buffer == nullptr) {
59 LOGE("[ArgMax] SetKeepdims GetBuffer return nullptr.");
60 return OH_NN_INVALID_PARAMETER;
61 }
62 m_keepDims = *(static_cast<bool*>(buffer));
63
64 return OH_NN_SUCCESS;
65 }
66
67 /**
68 * Build method.
69 * 1.build primitive of ops.
70 * 2.build inputIndex of ops.
71 * 3.build outputIndex of ops.
72 */
Build(const std::vector<uint32_t> & paramsIndex,const std::vector<uint32_t> & inputsIndex,const std::vector<uint32_t> & outputsIndex,const std::vector<std::shared_ptr<NNTensor>> & allTensors)73 OH_NN_ReturnCode ArgMaxBuilder::Build(const std::vector<uint32_t>& paramsIndex,
74 const std::vector<uint32_t>& inputsIndex, const std::vector<uint32_t>& outputsIndex,
75 const std::vector<std::shared_ptr<NNTensor>>& allTensors)
76 {
77 if (m_isBuild) {
78 LOGE("[ArgMax] Build failed, build operation has been completed, cannot build again.");
79 return OH_NN_OPERATION_FORBIDDEN;
80 }
81
82 OH_NN_ReturnCode returnCode = CheckIOIndex(inputsIndex, outputsIndex, allTensors, INPUT_NUM, OUTPUT_NUM);
83 if (returnCode != OH_NN_SUCCESS) {
84 LOGE("[ArgMax] Build failed, passed invalid input or output index.");
85 return returnCode;
86 }
87 m_inputsIndex = inputsIndex;
88 m_outputsIndex = outputsIndex;
89
90 for (int i : paramsIndex) {
91 const std::shared_ptr<NNTensor> tensor = allTensors[i];
92 switch (tensor->GetType()) {
93 case OH_NN_ARG_MAX_AXIS:
94 returnCode = SetAxis(tensor);
95 break;
96 case OH_NN_ARG_MAX_KEEPDIMS:
97 returnCode = SetKeepdims(tensor);
98 break;
99 default:
100 LOGE("[ArgMax] Build failed, param invalid, type = %d.", tensor->GetType());
101 return OH_NN_INVALID_PARAMETER;
102 }
103 if (returnCode != OH_NN_SUCCESS) {
104 LOGE("[ArgMax] Build failed, passed invalid param.");
105 return returnCode;
106 }
107 }
108
109 // The quantization type of the first output determinies that of the operator.
110 SetQuantType(outputsIndex, allTensors);
111
112 m_name = OP_NAME;
113 m_isBuild = true;
114 return OH_NN_SUCCESS;
115 }
116
GetPrimitive()117 LiteGraphPrimitvePtr ArgMaxBuilder::GetPrimitive()
118 {
119 if (!m_isBuild) {
120 LOGE("[ArgMax] GetPrimitive failed, cannot get primitive before call build.");
121 return {nullptr, DestroyLiteGraphPrimitive};
122 }
123
124 void* primitive = mindspore::lite::MindIR_ArgMaxFusion_CreatePrimitive(m_axis, m_topK, m_keepDims, m_outMaxValue);
125 LiteGraphPrimitvePtr graphPrimitivePtr(primitive, DestroyLiteGraphPrimitive);
126 return graphPrimitivePtr;
127 }
128 REGISTER_OPS(ArgMaxBuilder, OH_NN_OPS_ARG_MAX);
129 } // namespace Ops
130 } // namespace NeuralNetworkRuntime
131 } // namespace OHOS
132