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 "unsqueeze_builder.h"
17
18 #include "mindir.h"
19
20 namespace OHOS {
21 namespace NeuralNetworkRuntime {
22 namespace Ops {
23 static const int INPUT_NUM = 1;
24 static const int OUTPUT_NUM = 1;
25 static const std::string OP_NAME = "Unsqueeze";
26
UnsqueezeBuilder()27 UnsqueezeBuilder::UnsqueezeBuilder() {}
28
~UnsqueezeBuilder()29 UnsqueezeBuilder::~UnsqueezeBuilder() {}
30
SetAxis(std::shared_ptr<NNTensor> tensor)31 OH_NN_ReturnCode UnsqueezeBuilder::SetAxis(std::shared_ptr<NNTensor> tensor)
32 {
33 // Set Axis
34 if (tensor->GetDataType() != OH_NN_INT64) {
35 LOGE("[UnsqueezeBuilder] The 2nd input axis should be type OH_NN_INT64.");
36 return OH_NN_INVALID_PARAMETER;
37 }
38
39 if (tensor->GetElementCount() != 1) {
40 LOGE("[UnsqueezeBuilder] The 2nd input axis should be scaler.");
41 return OH_NN_INVALID_PARAMETER;
42 }
43
44 m_axis.clear();
45
46 void* buffer = tensor->GetBuffer();
47 if (buffer == nullptr) {
48 LOGE("[UnsqueezeBuilder] Tensor buffer is nullptr.");
49 return OH_NN_INVALID_PARAMETER;
50 }
51 m_axis.emplace_back(*(static_cast<const int64_t*>(buffer)));
52
53 return OH_NN_SUCCESS;
54 }
55
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)56 OH_NN_ReturnCode UnsqueezeBuilder::Build(const std::vector<uint32_t>& paramsIndex,
57 const std::vector<uint32_t>& inputsIndex,
58 const std::vector<uint32_t>& outputsIndex,
59 const std::vector<std::shared_ptr<NNTensor>>& allTensors)
60 {
61 if (m_isBuild) {
62 LOGE("[UnsqueezeBuilder] Unsqueeze build operation has been build, cannot build again");
63 return OH_NN_OPERATION_FORBIDDEN;
64 }
65 OH_NN_ReturnCode returnCode = CheckIOIndex(inputsIndex, outputsIndex, allTensors, INPUT_NUM, OUTPUT_NUM);
66 if (returnCode != OH_NN_SUCCESS) {
67 LOGE("[UnsqueezeBuilder] Passed invalid input or output index.");
68 return returnCode;
69 }
70
71 m_inputsIndex = inputsIndex;
72 m_outputsIndex = outputsIndex;
73
74 for (int i : paramsIndex) {
75 std::shared_ptr<NNTensor> tensor = allTensors[i];
76 tensor->IdentifyOpParameter();
77 switch (tensor->GetType()) {
78 case OH_NN_UNSQUEEZE_AXIS:
79 returnCode = SetAxis(tensor);
80 break;
81 default:
82 LOGE("[UnsqueezeBuilder] Parameter Type is invalid. type=%d", tensor->GetType());
83 return OH_NN_INVALID_PARAMETER;
84 }
85
86 if (returnCode != OH_NN_SUCCESS) {
87 LOGE("[UnsqueezeBuilder] Passed invalid param.");
88 return returnCode;
89 }
90 }
91 // The quantization type of the first output determinies that of the operator.
92 SetQuantType(outputsIndex, allTensors);
93
94 m_isBuild = true;
95 m_name = OP_NAME;
96 return OH_NN_SUCCESS;
97 }
98
GetPrimitive()99 LiteGraphPrimitvePtr UnsqueezeBuilder::GetPrimitive()
100 {
101 if (!m_isBuild) {
102 LOGE("[UnsqueezeBuilder] Cannot get primitive before call build.");
103 return {nullptr, DestroyLiteGraphPrimitive};
104 }
105
106 auto primitive = mindspore::lite::MindIR_Unsqueeze_CreatePrimitive(m_axis);
107 if (primitive == nullptr) {
108 LOGE("[UnsqueezeBuilder] MindIR_Unsqueeze_CreatePrimitive failed.");
109 return {nullptr, DestroyLiteGraphPrimitive};
110 }
111
112 LiteGraphPrimitvePtr graphPrimitivePtr(primitive, DestroyLiteGraphPrimitive);
113 return graphPrimitivePtr;
114 }
115
116 REGISTER_OPS(UnsqueezeBuilder, OH_NN_OPS_UNSQUEEZE);
117 } // namespace Ops
118 } // namespace NeuralNetworkRuntime
119 } // namespace OHOS
120