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