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/shuffle_tensorrt.h"
18 #include <vector>
19 #include <numeric>
20 #include <functional>
21
22 namespace mindspore::lite {
IsSupport(const schema::Primitive * primitive,const std::vector<mindspore::MSTensor> & in_tensors,const std::vector<mindspore::MSTensor> & out_tensors)23 int ShuffleTensorRT::IsSupport(const schema::Primitive *primitive, const std::vector<mindspore::MSTensor> &in_tensors,
24 const std::vector<mindspore::MSTensor> &out_tensors) {
25 if (!IsShapeKnown()) {
26 MS_LOG(ERROR) << "Unsupported input tensor unknown shape: " << op_name_;
27 return RET_ERROR;
28 }
29 switch (type_) {
30 case schema::PrimitiveType_Flatten:
31 case schema::PrimitiveType_Squeeze:
32 case schema::PrimitiveType_Unsqueeze: {
33 if (in_tensors.size() != 1) {
34 MS_LOG(ERROR) << "Unsupported in_tensors size " << in_tensors.size() << " of "
35 << schema::EnumNamePrimitiveType(type_);
36 return RET_ERROR;
37 }
38 break;
39 }
40 case schema::PrimitiveType_Reshape: {
41 if (in_tensors.size() != INPUT_SIZE2) {
42 MS_LOG(ERROR) << "PrimitiveType_Transpose Unsupported in_tensors size: " << in_tensors.size();
43 return RET_ERROR;
44 }
45 break;
46 }
47 case schema::PrimitiveType_Transpose: {
48 if (in_tensors.size() != INPUT_SIZE2) {
49 MS_LOG(ERROR) << "PrimitiveType_Transpose Unsupported in_tensors size: " << in_tensors.size();
50 return RET_ERROR;
51 }
52 if (in_tensors[1].Data() == nullptr) {
53 MS_LOG(ERROR) << "Unsupported shape tensor of " << schema::EnumNamePrimitiveType(type_);
54 return RET_ERROR;
55 }
56 break;
57 }
58 default: {
59 MS_LOG(ERROR) << "Unsupported op type:" << schema::EnumNamePrimitiveType(type_);
60 return RET_ERROR;
61 }
62 }
63 if (out_tensors.size() != 1) {
64 MS_LOG(ERROR) << "invalid output tensort size: " << out_tensors.size();
65 return RET_ERROR;
66 }
67 return RET_OK;
68 }
69
AddInnerOp(nvinfer1::INetworkDefinition * network)70 int ShuffleTensorRT::AddInnerOp(nvinfer1::INetworkDefinition *network) {
71 if (network == nullptr) {
72 MS_LOG(ERROR) << "network is invalid";
73 return RET_ERROR;
74 }
75 nvinfer1::ITensor *shuffler_input = tensorrt_in_tensors_[0].trt_tensor_;
76 MS_LOG(DEBUG) << "before transpose " << GetTensorFormat(shuffler_input, tensorrt_in_tensors_[0].format_);
77 if (tensorrt_in_tensors_[0].trt_tensor_->getDimensions().nbDims == DIMENSION_4D &&
78 !SameDims(tensorrt_in_tensors_[0].trt_tensor_->getDimensions(), in_tensors_[0].Shape())) {
79 // only valid for nchw or nhwc
80 if (tensorrt_in_tensors_[0].format_ == Format::NCHW) {
81 nvinfer1::IShuffleLayer *transpose_layer = NCHW2NHWC(network, *tensorrt_in_tensors_[0].trt_tensor_);
82 if (transpose_layer == nullptr) {
83 MS_LOG(ERROR) << "create transpose layer failed for " << op_name_;
84 }
85 transpose_layer->setName((op_name_ + "_transpose_in").c_str());
86 shuffler_input = transpose_layer->getOutput(0);
87 out_format_ = Format::NHWC;
88 } else if (tensorrt_in_tensors_[0].format_ == Format::NHWC) {
89 nvinfer1::IShuffleLayer *transpose_layer = NHWC2NCHW(network, *tensorrt_in_tensors_[0].trt_tensor_);
90 if (transpose_layer == nullptr) {
91 MS_LOG(ERROR) << "create transpose layer failed for " << op_name_;
92 }
93 transpose_layer->setName((op_name_ + "_transpose_in").c_str());
94 shuffler_input = transpose_layer->getOutput(0);
95 out_format_ = Format::NCHW;
96 } else {
97 MS_LOG(ERROR) << "invalid input format for " << op_name_;
98 return RET_ERROR;
99 }
100 } else {
101 out_format_ = tensorrt_in_tensors_[0].format_;
102 }
103 MS_LOG(DEBUG) << "after transpose " << GetTensorFormat(shuffler_input, out_format_);
104
105 nvinfer1::IShuffleLayer *shuffle_layer = network->addShuffle(*shuffler_input);
106 if (shuffle_layer == nullptr) {
107 MS_LOG(ERROR) << "add Shuffle op failed for TensorRT.";
108 return RET_ERROR;
109 }
110 shuffle_layer->setName(op_name_.c_str());
111
112 int ret = RET_OK;
113 switch (type_) {
114 case schema::PrimitiveType_Unsqueeze: {
115 ret = AddUnsqueezeOp(shuffle_layer);
116 break;
117 }
118 case schema::PrimitiveType_Squeeze: {
119 ret = AddSqueezeOp(shuffle_layer);
120 break;
121 }
122 case schema::PrimitiveType_Transpose: {
123 ret = AddTransposeOp(shuffle_layer);
124 break;
125 }
126 case schema::PrimitiveType_Reshape: {
127 ret = AddReshapeOp(shuffle_layer);
128 break;
129 }
130 case schema::PrimitiveType_Flatten: {
131 ret = AddFlattenOp(shuffle_layer);
132 break;
133 }
134 default:
135 MS_LOG(ERROR) << "Unsupported op type for " << op_name_;
136 return RET_ERROR;
137 }
138 if (ret != RET_OK) {
139 MS_LOG(ERROR) << "AddOp failed for " << op_name_;
140 return ret;
141 }
142
143 nvinfer1::ITensor *out_tensor = shuffle_layer->getOutput(0);
144 if (out_tensor == nullptr) {
145 MS_LOG(ERROR) << "output tensor create failed";
146 return RET_ERROR;
147 }
148 out_tensor->setName((op_name_ + "_output").c_str());
149 MS_LOG(DEBUG) << "output " << GetTensorFormat(out_tensor, out_format_);
150 this->AddInnerOutTensors(ITensorHelper{out_tensor, out_format_});
151 return RET_OK;
152 }
153
AddSqueezeOp(nvinfer1::IShuffleLayer * shuffle_layer)154 int ShuffleTensorRT::AddSqueezeOp(nvinfer1::IShuffleLayer *shuffle_layer) {
155 // squeeze
156 auto squeeze_op = this->op_primitive_->value_as_Squeeze();
157 if (squeeze_op == nullptr) {
158 MS_LOG(ERROR) << "SqueezeOp convert failed";
159 return RET_ERROR;
160 }
161
162 // axis
163 auto squeeze_shape = in_tensors_[0].Shape();
164 auto begin = std::begin(squeeze_shape);
165 auto axis = squeeze_op->axis();
166 if (axis == nullptr) {
167 MS_LOG(ERROR) << "AddSqueezeOp has invalid axis";
168 return RET_ERROR;
169 }
170
171 for (size_t i = 0; i < axis->size(); i++) {
172 if (squeeze_shape[axis->Get(i)] != 1) {
173 MS_LOG(WARNING) << "squeeze_shape value is not 1, need check";
174 }
175 squeeze_shape.erase(begin + axis->Get(i));
176 }
177
178 nvinfer1::Dims squeeze_dims = lite::ConvertCudaDims(squeeze_shape);
179 MS_LOG(DEBUG) << "AddSqueezeOp: " << op_name_ << " squeeze_dims.nbDims: " << squeeze_dims.nbDims;
180
181 shuffle_layer->setReshapeDimensions(squeeze_dims);
182 return shuffle_layer->getOutput(0) == nullptr ? RET_ERROR : RET_OK;
183 }
184
AddUnsqueezeOp(nvinfer1::IShuffleLayer * shuffle_layer)185 int ShuffleTensorRT::AddUnsqueezeOp(nvinfer1::IShuffleLayer *shuffle_layer) {
186 // Unsqueeze
187 auto unsqueeze_op = this->op_primitive_->value_as_Unsqueeze();
188 if (unsqueeze_op == nullptr) {
189 MS_LOG(ERROR) << "AddUnsqueezeOp convert failed";
190 return RET_ERROR;
191 }
192 if (in_tensors_.size() != 1) {
193 MS_LOG(WARNING) << "AddUnsqueezeOp size of in tensort needs check: " << in_tensors_.size();
194 }
195 // axis
196 auto unsqueeze_shape = tensorrt_in_tensors_[0].trt_tensor_->getDimensions();
197 std::vector<int64_t> new_shape(unsqueeze_shape.d, unsqueeze_shape.d + unsqueeze_shape.nbDims);
198 auto axis = unsqueeze_op->axis();
199
200 for (size_t i = 0; i < axis->size(); i++) {
201 new_shape.insert(new_shape.begin() + axis->Get(i), 1);
202 }
203
204 nvinfer1::Dims unsqueeze_dims = lite::ConvertCudaDims(new_shape);
205
206 shuffle_layer->setReshapeDimensions(unsqueeze_dims);
207 return shuffle_layer->getOutput(0) == nullptr ? RET_ERROR : RET_OK;
208 }
209
AddTransposeOp(nvinfer1::IShuffleLayer * shuffle_layer)210 int ShuffleTensorRT::AddTransposeOp(nvinfer1::IShuffleLayer *shuffle_layer) {
211 auto transpose_op = this->op_primitive_->value_as_Transpose();
212 if (transpose_op == nullptr) {
213 MS_LOG(ERROR) << "AddTransposeOp convert failed";
214 return RET_ERROR;
215 }
216 if (in_tensors_.size() != 2) {
217 MS_LOG(ERROR) << "AddTransposeOp size of in tensort needs check: " << in_tensors_.size();
218 return RET_ERROR;
219 }
220 // perm
221 mindspore::MSTensor perm_ternsor = in_tensors_[1];
222 if (perm_ternsor.Data() == nullptr) {
223 MS_LOG(ERROR) << "AddTransposeOp perm_ternsor data is invalid: " << op_name_;
224 return RET_ERROR;
225 }
226 int *perm_data = reinterpret_cast<int *>(perm_ternsor.MutableData());
227
228 nvinfer1::Permutation perm{};
229 for (int i = 0; i < perm_ternsor.ElementNum(); i++) {
230 perm.order[i] = *perm_data;
231 perm_data++;
232 }
233 shuffle_layer->setFirstTranspose(perm);
234 if (perm_ternsor.ElementNum() == DIMENSION_4D) {
235 if (perm.order[1] == 3 && perm.order[2] == 1 && perm.order[3] == 2) {
236 out_format_ = Format::NCHW;
237 } else if (perm.order[1] == 2 && perm.order[2] == 3 && perm.order[3] == 1) {
238 out_format_ = Format::NHWC;
239 } else {
240 MS_LOG(WARNING) << "input format and perm order is invalid: " << op_name_;
241 }
242 }
243 return RET_OK;
244 }
245
AddReshapeOp(nvinfer1::IShuffleLayer * shuffle_layer)246 int ShuffleTensorRT::AddReshapeOp(nvinfer1::IShuffleLayer *shuffle_layer) {
247 mindspore::MSTensor &shape_tensor = in_tensors_[1];
248 if (shape_tensor.Data() != nullptr) {
249 // static shuffle layer
250 nvinfer1::Dims reshape_dims = lite::ConvertCudaDims(shape_tensor.Data().get(), shape_tensor.ElementNum());
251 shuffle_layer->setReshapeDimensions(reshape_dims);
252 } else {
253 if (tensorrt_in_tensors_.size() != INPUT_SIZE2) {
254 MS_LOG(ERROR) << "invalid shape tensor for reshape " << op_name_;
255 return RET_ERROR;
256 }
257 shuffle_layer->setInput(1, *tensorrt_in_tensors_[1].trt_tensor_);
258 }
259 return RET_OK;
260 }
261
AddFlattenOp(nvinfer1::IShuffleLayer * shuffle_layer)262 int ShuffleTensorRT::AddFlattenOp(nvinfer1::IShuffleLayer *shuffle_layer) {
263 nvinfer1::Dims flatten_dims;
264 const std::vector<int64_t> &input_shape = in_tensors_[0].Shape();
265 flatten_dims.nbDims = DIMENSION_2D;
266 flatten_dims.d[0] = input_shape[0];
267 flatten_dims.d[1] = std::accumulate(input_shape.begin() + 1, input_shape.end(), 1, std::multiplies<int>());
268 shuffle_layer->setReshapeDimensions(flatten_dims);
269 return RET_OK;
270 }
271
InferReshapeDims(nvinfer1::Dims input_dims,nvinfer1::Dims * reshape_dims)272 int ShuffleTensorRT::InferReshapeDims(nvinfer1::Dims input_dims, nvinfer1::Dims *reshape_dims) {
273 // tensorrt support infer shape of 0 and -1
274 int infer_index = -1;
275 int known_cnt = 1;
276 for (int i = 0; i < reshape_dims->nbDims; i++) {
277 if (reshape_dims->d[i] == 0) {
278 reshape_dims->d[i] = input_dims.d[i];
279 known_cnt *= (input_dims.d[i] == -1 ? 1 : input_dims.d[i]);
280 } else if (reshape_dims->d[i] == -1) {
281 if (infer_index != -1) {
282 MS_LOG(ERROR) << "invalid dims (more than one infer dim) for reshape " << op_name_;
283 return RET_ERROR;
284 }
285 infer_index = i;
286 } else {
287 known_cnt *= input_dims.d[i];
288 }
289 }
290 if (infer_index != -1) {
291 size_t tot_cnt = 1;
292 for (int i = 0; i < input_dims.nbDims; i++) {
293 tot_cnt *= (input_dims.d[i] == -1 ? 1 : input_dims.d[i]);
294 }
295 if (known_cnt == 0) {
296 MS_LOG(ERROR) << "invalid known cnt for " << op_name_;
297 return RET_ERROR;
298 }
299 reshape_dims->d[infer_index] = tot_cnt / known_cnt;
300 MS_LOG(DEBUG) << "reshape infer_index: " << infer_index
301 << ", reshape infer value: " << reshape_dims->d[infer_index];
302 }
303 return RET_OK;
304 }
305 } // namespace mindspore::lite
306