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 "constant_of_shape_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 int PARAM_MAX_NUM = 2;
24 static const int SCALAR_LENGTH = 1;
25 static const std::string OP_NAME = "ConstantOfShape";
26 
ConstantOfShapeBuilder()27 ConstantOfShapeBuilder::ConstantOfShapeBuilder() {}
28 
~ConstantOfShapeBuilder()29 ConstantOfShapeBuilder::~ConstantOfShapeBuilder() {}
30 
SetDataType(const std::shared_ptr<NNTensor> & tensor)31 OH_NN_ReturnCode ConstantOfShapeBuilder::SetDataType(const std::shared_ptr<NNTensor>& tensor)
32 {
33     if (tensor->GetDataType() != OH_NN_INT64) {
34         LOGE("[ConstantOfShape] The dataType should be type OH_NN_INT64.");
35         return OH_NN_INVALID_PARAMETER;
36     }
37 
38     if (tensor->GetElementCount() != SCALAR_LENGTH) {
39         LOGE("[ConstantOfShape] The dataType should be scalar.");
40         return OH_NN_INVALID_PARAMETER;
41     }
42 
43     void* buffer = tensor->GetBuffer();
44     if (buffer == nullptr) {
45         LOGE("[ConstantOfShape] Tensor buffer is nullptr.");
46         return OH_NN_INVALID_PARAMETER;
47     }
48     m_dataType = *static_cast<const int64_t*>(buffer);
49 
50     return OH_NN_SUCCESS;
51 }
52 
SetValue(const std::shared_ptr<NNTensor> & tensor)53 OH_NN_ReturnCode ConstantOfShapeBuilder::SetValue(const std::shared_ptr<NNTensor>& tensor)
54 {
55     if (tensor->GetDataType() != OH_NN_FLOAT32) {
56         LOGE("[ConstantOfShape] The value should be type OH_NN_FLOAT32.");
57         return OH_NN_INVALID_PARAMETER;
58     }
59 
60     m_value.clear();
61 
62     void* buffer = tensor->GetBuffer();
63     if (buffer == nullptr) {
64         LOGE("[ConstantOfShape] Tensor buffer is nullptr.");
65         return OH_NN_INVALID_PARAMETER;
66     }
67 
68     float* pValue = static_cast<float*>(buffer);
69 
70     uint32_t elementCount = tensor->GetElementCount();
71     for (uint32_t i = 0; i < elementCount; ++i) {
72         m_value.emplace_back(*pValue);
73         ++pValue;
74     }
75     return OH_NN_SUCCESS;
76 }
77 
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)78 OH_NN_ReturnCode ConstantOfShapeBuilder::Build(const std::vector<uint32_t>& paramsIndex,
79                                                const std::vector<uint32_t>& inputsIndex,
80                                                const std::vector<uint32_t>& outputsIndex,
81                                                const std::vector<std::shared_ptr<NNTensor>>& allTensors)
82 {
83     if (m_isBuild) {
84         LOGE("[ConstantOfShape] Build failed, the constantOfShape operation has been build. cannot build again.");
85         return OH_NN_OPERATION_FORBIDDEN;
86     }
87 
88     auto ret = CheckIOIndex(inputsIndex, outputsIndex, allTensors, INPUT_NUM, OUTPUT_NUM);
89     if (ret != OH_NN_SUCCESS) {
90         LOGE("[ConstantOfShape] Build failed, passed invalid input or output index.");
91         return ret;
92     }
93 
94     m_inputsIndex = inputsIndex;
95     m_outputsIndex = outputsIndex;
96 
97     ret = CheckParamIndex(paramsIndex, allTensors, PARAM_MAX_NUM);
98     if (ret != OH_NN_SUCCESS) {
99         LOGE("[ConstantOfShape] Build failed, passed invalid invalid index.");
100         return ret;
101     }
102 
103     for (int i : paramsIndex) {
104         std::shared_ptr<NNTensor> tensor = allTensors[i];
105         tensor->IdentifyOpParameter();
106         if (m_paramMap.find(tensor->GetType()) != m_paramMap.end()) {
107             ret = (this->*(m_paramMap[tensor->GetType()]))(tensor);
108         } else {
109             LOGE("[ConstantOfShape] Build failed, param invalid, type=%d", tensor->GetType());
110             return OH_NN_INVALID_PARAMETER;
111         }
112 
113         if (ret != OH_NN_SUCCESS) {
114             LOGE("[ConstantOfShape] Build failed, passed invalid param.");
115             return ret;
116         }
117     }
118 
119     m_name = OP_NAME;
120     m_isBuild = true;
121     return OH_NN_SUCCESS;
122 }
123 
GetPrimitive()124 LiteGraphPrimitvePtr ConstantOfShapeBuilder::GetPrimitive()
125 {
126     if (!m_isBuild) {
127         LOGE("[ConstantOfShape] GetPrimitive failed, cannot get primitive before call build.");
128         return {nullptr, DestroyLiteGraphPrimitive};
129     }
130 
131     void* primitive = mindspore::lite::MindIR_ConstantOfShape_CreatePrimitive(m_dataType, m_value);
132     LiteGraphPrimitvePtr graphPrimitivePtr(primitive, DestroyLiteGraphPrimitive) ;
133     return graphPrimitivePtr;
134 }
135 
136 REGISTER_OPS(ConstantOfShapeBuilder, OH_NN_OPS_CONSTANT_OF_SHAPE);
137 } // namespace Ops
138 } // namespace NeuralNetworkRuntime
139 } // namespace OHOS
140