1 /*
2  * Copyright (c) 2023 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 "depth_to_space_builder.h"
17 
18 #include "transform.h"
19 #include "validation.h"
20 
21 namespace OHOS {
22 namespace NeuralNetworkRuntime {
23 namespace Ops {
24 static const int INPUT_NUM = 1;
25 static const int OUTPUT_NUM = 1;
26 static const int PARAM_MAX_NUM = 2;
27 static const int SCALAR_LENGTH = 1;
28 static const std::string OP_NAME = "DepthToSpace";
29 static const std::unordered_map<int, std::string> modeList = {{0, "DCR"}, {1, "CRD"}};
30 
DepthToSpaceBuilder()31 DepthToSpaceBuilder::DepthToSpaceBuilder() {}
32 
~DepthToSpaceBuilder()33 DepthToSpaceBuilder::~DepthToSpaceBuilder() {}
34 
SetBlockSize(const std::shared_ptr<NNTensor> & tensor)35 OH_NN_ReturnCode DepthToSpaceBuilder::SetBlockSize(const std::shared_ptr<NNTensor>& tensor)
36 {
37     if (tensor->GetDataType() != OH_NN_INT64) {
38         LOGE("[DepthToSpace] The blockSize should be type OH_NN_INT64.");
39         return OH_NN_INVALID_PARAMETER;
40     }
41 
42     if (tensor->GetElementCount() != SCALAR_LENGTH) {
43         LOGE("[DepthToSpace] The blockSize should be scalar.");
44         return OH_NN_INVALID_PARAMETER;
45     }
46 
47     void* buffer = tensor->GetBuffer();
48     if (buffer == nullptr) {
49         LOGE("[DepthToSpace] Tensor buffer is nullptr.");
50         return OH_NN_INVALID_PARAMETER;
51     }
52     m_blockSize = *(static_cast<const int64_t*>(buffer));
53 
54     return OH_NN_SUCCESS;
55 }
56 
SetMode(const std::shared_ptr<NNTensor> & tensor)57 OH_NN_ReturnCode DepthToSpaceBuilder::SetMode(const std::shared_ptr<NNTensor>& tensor)
58 {
59     if (tensor->GetDataType() != OH_NN_INT32) {
60         LOGE("[DepthToSpace] The mode should be type OH_NN_INT32.");
61         return OH_NN_INVALID_PARAMETER;
62     }
63 
64     void* buffer = tensor->GetBuffer();
65     if (buffer == nullptr) {
66         LOGE("[DepthToSpace] Tensor buffer is nullptr.");
67         return OH_NN_INVALID_PARAMETER;
68     }
69 
70     int modeKey = *(static_cast<int*>(buffer));
71     auto it = modeList.find(modeKey);
72     if (it != modeList.end()) {
73         m_mode = it->second;
74     } else {
75         LOGE("[DepthToSpace] The mode value should between [0, 1], but get %d.", modeKey);
76         LOGE("[DepthToSpace] mode value: 0-DCR, 1-CRD");
77         return OH_NN_INVALID_PARAMETER;
78     }
79 
80     return OH_NN_SUCCESS;
81 }
82 
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)83 OH_NN_ReturnCode DepthToSpaceBuilder::Build(const std::vector<uint32_t>& paramsIndex,
84                                             const std::vector<uint32_t>& inputsIndex,
85                                             const std::vector<uint32_t>& outputsIndex,
86                                             const std::vector<std::shared_ptr<NNTensor>>& allTensors)
87 {
88     if (m_isBuild) {
89         LOGE("[DepthToSpace] Build failed, the depthToSpace operation has been build. cannot build again.");
90         return OH_NN_OPERATION_FORBIDDEN;
91     }
92 
93     auto ret = CheckIOIndex(inputsIndex, outputsIndex, allTensors, INPUT_NUM, OUTPUT_NUM);
94     if (ret != OH_NN_SUCCESS) {
95         LOGE("[DepthToSpace] Build failed, passed invalid input or output index.");
96         return ret;
97     }
98 
99     m_inputsIndex = inputsIndex;
100     m_outputsIndex = outputsIndex;
101 
102     ret = CheckParamIndex(paramsIndex, allTensors, PARAM_MAX_NUM);
103     if (ret != OH_NN_SUCCESS) {
104         LOGE("[DepthToSpace] Build failed, passed invalid param index.");
105         return ret;
106     }
107 
108     for (int i : paramsIndex) {
109         std::shared_ptr<NNTensor> tensor = allTensors[i];
110         tensor->IdentifyOpParameter();
111         if (m_paramMap.find(tensor->GetType()) != m_paramMap.end()) {
112             ret = (this->*(m_paramMap[tensor->GetType()]))(tensor);
113         } else {
114             LOGE("[DepthToSpace] Build failed, param invalid, type=%d", tensor->GetType());
115             return OH_NN_INVALID_PARAMETER;
116         }
117 
118         if (ret != OH_NN_SUCCESS) {
119             LOGE("[DepthToSpace] Build failed, passed invalid param.");
120             return ret;
121         }
122     }
123 
124     m_name = OP_NAME;
125     m_isBuild = true;
126     return OH_NN_SUCCESS;
127 }
128 
GetPrimitive()129 LiteGraphPrimitvePtr DepthToSpaceBuilder::GetPrimitive()
130 {
131     if (!m_isBuild) {
132         LOGE("[DepthToSpace] GetPrimitive failed, cannot get primitive before call build.");
133         return {nullptr, DestroyLiteGraphPrimitive};
134     }
135 
136     mindspore::lite::Format format {mindspore::lite::FORMAT_NCHW};
137 
138     void* primitive = mindspore::lite::MindIR_DepthToSpace_CreatePrimitive(m_blockSize, format, m_mode);
139     LiteGraphPrimitvePtr graphPrimitivePtr(primitive, DestroyLiteGraphPrimitive) ;
140     return graphPrimitivePtr;
141 }
142 
143 REGISTER_OPS(DepthToSpaceBuilder, OH_NN_OPS_DEPTH_TO_SPACE);
144 } // namespace Ops
145 } // namespace NeuralNetworkRuntime
146 } // namespace OHOS