1 /**
2 * Copyright 2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "src/delegate/tensorrt/op/elementwise_tensorrt.h"
18 #include "src/delegate/tensorrt/tensorrt_utils.h"
19
20 namespace mindspore::lite {
IsSupport(const schema::Primitive * primitive,const std::vector<mindspore::MSTensor> & in_tensors,const std::vector<mindspore::MSTensor> & out_tensors)21 int ElementWiseTensorRT::IsSupport(const schema::Primitive *primitive,
22 const std::vector<mindspore::MSTensor> &in_tensors,
23 const std::vector<mindspore::MSTensor> &out_tensors) {
24 if (!IsShapeKnown()) {
25 MS_LOG(ERROR) << "Unsupported input tensor unknown shape: " << op_name_;
26 return RET_ERROR;
27 }
28 std::map<schema::PrimitiveType, nvinfer1::ElementWiseOperation> element_wise_ops = {
29 {schema::PrimitiveType_AddFusion, nvinfer1::ElementWiseOperation::kSUM},
30 {schema::PrimitiveType_PowFusion, nvinfer1::ElementWiseOperation::kPOW},
31 {schema::PrimitiveType_DivFusion, nvinfer1::ElementWiseOperation::kDIV},
32 {schema::PrimitiveType_SubFusion, nvinfer1::ElementWiseOperation::kSUB},
33 {schema::PrimitiveType_MulFusion, nvinfer1::ElementWiseOperation::kPROD},
34 };
35 auto iter_op = element_wise_ops.find(this->type_);
36 if (iter_op != element_wise_ops.end()) {
37 element_wise_op_ = iter_op->second;
38 } else {
39 // PrimitiveType_Eltwise
40 auto eltwise_op = op_primitive_->value_as_Eltwise();
41 if (eltwise_op == nullptr) {
42 MS_LOG(ERROR) << "convert to Eltwise failed: " << op_name_;
43 return RET_ERROR;
44 }
45 schema::EltwiseMode eltwiseMode = eltwise_op->mode();
46 std::map<schema::EltwiseMode, nvinfer1::ElementWiseOperation> eltwise_modes = {
47 {schema::EltwiseMode::EltwiseMode_SUM, nvinfer1::ElementWiseOperation::kSUM},
48 {schema::EltwiseMode::EltwiseMode_PROD, nvinfer1::ElementWiseOperation::kPROD},
49 {schema::EltwiseMode::EltwiseMode_MAXIMUM, nvinfer1::ElementWiseOperation::kMAX},
50 };
51 auto iter_mode = eltwise_modes.find(eltwiseMode);
52 if (iter_mode != eltwise_modes.end()) {
53 element_wise_op_ = iter_mode->second;
54 } else {
55 MS_LOG(ERROR) << "unsupported type for ElementWise op" << op_name_;
56 return RET_ERROR;
57 }
58 }
59
60 if (in_tensors.size() != INPUT_SIZE2) {
61 MS_LOG(ERROR) << "invalid input tensort size: " << in_tensors.size();
62 return RET_ERROR;
63 }
64 if (out_tensors.size() != 1) {
65 MS_LOG(ERROR) << "invalid output tensort size: " << out_tensors.size();
66 return RET_ERROR;
67 }
68
69 // if constant tensor is scalar, it needs to know another input tensor's shape to broadcast
70 if (in_tensors[0].Shape()[0] == -1 && in_tensors[1].Shape().size() == 0) {
71 MS_LOG(ERROR) << "invalid all input tensor shape unknown for: " << op_name_;
72 return RET_ERROR;
73 }
74
75 return RET_OK;
76 }
77
AddInnerOp(nvinfer1::INetworkDefinition * network)78 int ElementWiseTensorRT::AddInnerOp(nvinfer1::INetworkDefinition *network) {
79 if (network == nullptr) {
80 MS_LOG(ERROR) << "network or input tensor size is invalid";
81 return RET_ERROR;
82 }
83 first_in_tensor_index_ =
84 SameDims(tensorrt_in_tensors_[0].trt_tensor_->getDimensions(), in_tensors_[0].Shape()) ? 0 : 1;
85
86 if (this->tensorrt_in_tensors_.size() != INPUT_SIZE2) {
87 int ret = AddConstTensor(network);
88 if (ret != RET_OK) {
89 MS_LOG(ERROR) << "AddConstTensor failed for " << op_name_;
90 return ret;
91 }
92 }
93 MS_LOG(DEBUG) << "before transpose "
94 << GetTensorFormat(tensorrt_in_tensors_[first_in_tensor_index_].trt_tensor_,
95 tensorrt_in_tensors_[first_in_tensor_index_].format_);
96 MS_LOG(DEBUG) << "before transpose "
97 << GetTensorFormat(tensorrt_in_tensors_[1 - first_in_tensor_index_].trt_tensor_,
98 tensorrt_in_tensors_[1 - first_in_tensor_index_].format_);
99
100 if (tensorrt_in_tensors_[0].trt_tensor_->getDimensions().nbDims == DIMENSION_4D &&
101 tensorrt_in_tensors_[0].format_ != tensorrt_in_tensors_[1].format_) {
102 // when inputs format are different, change to NHWC
103 int transpose_input_tensor = tensorrt_in_tensors_[0].format_ == Format::NCHW ? 0 : 1;
104 nvinfer1::IShuffleLayer *transpose_layer =
105 NCHW2NHWC(network, *tensorrt_in_tensors_[transpose_input_tensor].trt_tensor_);
106 if (transpose_layer == nullptr) {
107 MS_LOG(ERROR) << "op action convert failed";
108 return RET_ERROR;
109 }
110 transpose_layer->setName((op_name_ + "_input_transpose2NHWC").c_str());
111 tensorrt_in_tensors_[transpose_input_tensor].trt_tensor_ = transpose_layer->getOutput(0);
112 tensorrt_in_tensors_[transpose_input_tensor].format_ = Format::NHWC;
113 } else if (tensorrt_in_tensors_[0].format_ != tensorrt_in_tensors_[1].format_) {
114 MS_LOG(ERROR) << "elementwise op inputs are in different format: " << op_name_;
115 return RET_ERROR;
116 }
117 MS_LOG(DEBUG) << "after transpose "
118 << GetTensorFormat(tensorrt_in_tensors_[first_in_tensor_index_].trt_tensor_,
119 tensorrt_in_tensors_[first_in_tensor_index_].format_);
120 MS_LOG(DEBUG) << "after transpose "
121 << GetTensorFormat(tensorrt_in_tensors_[1 - first_in_tensor_index_].trt_tensor_,
122 tensorrt_in_tensors_[1 - first_in_tensor_index_].format_);
123
124 nvinfer1::IElementWiseLayer *cal_layer =
125 network->addElementWise(*tensorrt_in_tensors_[first_in_tensor_index_].trt_tensor_,
126 *tensorrt_in_tensors_[1 - first_in_tensor_index_].trt_tensor_, element_wise_op_);
127
128 if (cal_layer == nullptr) {
129 MS_LOG(ERROR) << "addElementWise failed for TensorRT.";
130 return RET_ERROR;
131 }
132 cal_layer->setName(op_name_.c_str());
133
134 nvinfer1::ITensor *op_out_tensor = cal_layer->getOutput(0);
135 if (op_out_tensor == nullptr) {
136 MS_LOG(ERROR) << "addElementWise out tensor is nullptr.";
137 return RET_ERROR;
138 }
139 // add activation
140 nvinfer1::ITensor *activation_out_tensor = AddActivation(network, op_out_tensor);
141 op_out_tensor = (activation_out_tensor == nullptr) ? op_out_tensor : activation_out_tensor;
142
143 // scale and shift
144 if (type_ == schema::PrimitiveType_PowFusion) {
145 auto pow_op = op_primitive_->value_as_PowFusion();
146 if (pow_op == nullptr) {
147 MS_LOG(ERROR) << "PowFusion convert failed.";
148 return RET_ERROR;
149 }
150 float scale = pow_op->scale();
151 float shift = pow_op->shift();
152 if (abs(scale - 1) >= 1.0e-05 || abs(shift - 0) >= 1.0e-05) {
153 MS_LOG(WARNING) << "deal with scale and shift for pow op";
154 }
155 }
156 op_out_tensor->setName((op_name_ + "_output").c_str());
157 this->AddInnerOutTensors(ITensorHelper{op_out_tensor, tensorrt_in_tensors_[1].format_});
158 MS_LOG(DEBUG) << "output " << GetTensorFormat(op_out_tensor, tensorrt_in_tensors_[1].format_);
159 return RET_OK;
160 }
161
AddActivation(nvinfer1::INetworkDefinition * network,nvinfer1::ITensor * in_tensor)162 nvinfer1::ITensor *ElementWiseTensorRT::AddActivation(nvinfer1::INetworkDefinition *network,
163 nvinfer1::ITensor *in_tensor) {
164 schema::ActivationType activation = schema::ActivationType::ActivationType_NO_ACTIVATION;
165 switch (type_) {
166 case schema::PrimitiveType_AddFusion: {
167 auto sum_op = op_primitive_->value_as_AddFusion();
168 if (sum_op == nullptr) {
169 MS_LOG(ERROR) << "AddFusion convert failed.";
170 return nullptr;
171 }
172 activation = sum_op->activation_type();
173 break;
174 }
175 case schema::PrimitiveType_DivFusion: {
176 auto div_op = op_primitive_->value_as_DivFusion();
177 if (div_op == nullptr) {
178 MS_LOG(ERROR) << "DivFusion convert failed.";
179 return nullptr;
180 }
181 activation = div_op->activation_type();
182 break;
183 }
184 case schema::PrimitiveType_SubFusion: {
185 auto sub_op = op_primitive_->value_as_SubFusion();
186 if (sub_op == nullptr) {
187 MS_LOG(ERROR) << "SubFusion convert failed.";
188 return nullptr;
189 }
190 activation = sub_op->activation_type();
191 break;
192 }
193 case schema::PrimitiveType_MulFusion: {
194 auto mul_op = op_primitive_->value_as_MulFusion();
195 if (mul_op == nullptr) {
196 MS_LOG(ERROR) << "MulFusion convert failed.";
197 return nullptr;
198 }
199 activation = mul_op->activation_type();
200 break;
201 }
202 default:
203 MS_LOG(DEBUG) << "no activation need for: " << op_name_;
204 }
205 nvinfer1::ITensor *activation_out_tensor = nullptr;
206 if (activation != schema::ActivationType::ActivationType_NO_ACTIVATION) {
207 MS_LOG(WARNING) << "op: " << op_name_ << " has activation";
208 }
209 return activation_out_tensor;
210 }
AddConstTensor(nvinfer1::INetworkDefinition * network)211 int ElementWiseTensorRT::AddConstTensor(nvinfer1::INetworkDefinition *network) {
212 // create ITensor from MS constant tensor of index 1 - first_in_tensor_index_
213 nvinfer1::ITensor *constant_input = nullptr;
214 if (this->in_tensors_[1 - first_in_tensor_index_].Shape().size() == 0 ||
215 this->in_tensors_[1 - first_in_tensor_index_].ElementNum() == 1) {
216 constant_input = lite::ConvertScalarToITensor(network, this->in_tensors_[first_in_tensor_index_].Shape().size(),
217 in_tensors_[1 - first_in_tensor_index_].Data().get(),
218 in_tensors_[1 - first_in_tensor_index_].DataType());
219 if (constant_input == nullptr) {
220 MS_LOG(ERROR) << "create Itensor from scalar tensor failed: " << op_name_;
221 return RET_ERROR;
222 }
223 this->AddInnerInTensors(ITensorHelper{constant_input, tensorrt_in_tensors_[0].format_});
224 } else {
225 constant_input = lite::ConvertConstantTensor(network, in_tensors_[1 - first_in_tensor_index_]);
226 if (constant_input == nullptr) {
227 MS_LOG(ERROR) << "create Itensor from constant tensor failed: " << op_name_;
228 return RET_ERROR;
229 }
230 this->AddInnerInTensors(ITensorHelper{constant_input, Format::NHWC});
231 }
232 return RET_OK;
233 }
234 } // namespace mindspore::lite
235