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 #ifndef MINDSPORE_LITE_SRC_RUNTIME_DELEGATE_TENSORRT_UTILS_H_
17 #define MINDSPORE_LITE_SRC_RUNTIME_DELEGATE_TENSORRT_UTILS_H_
18 #include <vector>
19 #include <NvInfer.h>
20 #include <memory>
21 #include <string>
22 #include "src/delegate/tensorrt/op/tensorrt_op.h"
23 #include "mindspore/core/ir/dtype/type_id.h"
24 #include "schema/ops_generated.h"
25 #include "nnacl/pack.h"
26
27 #define kNCHW_N 0
28 #define kNCHW_C 1
29 #define kNCHW_H 2
30 #define kNCHW_W 3
31 namespace mindspore::lite {
32 struct ActivationParams {
33 nvinfer1::ActivationType activation_type;
34 bool has_alpha;
35 float alpha;
36 bool has_beta;
37 float beta;
38 };
39
40 typedef union float32_bits {
41 unsigned int u;
42 float f;
43 } float32_bits;
44
45 // Convert Tensor data to Cuda dims.
46 nvinfer1::Dims ConvertCudaDims(const void *data, int64_t size);
47
48 nvinfer1::Dims ConvertCudaDims(int data, size_t size);
49
50 bool SameDims(nvinfer1::Dims dims, const std::vector<int64_t> &shape);
51
52 std::vector<int64_t> ConvertMSShape(const nvinfer1::Dims dims);
53
54 nvinfer1::DataType ConvertDataType(DataType type_id);
55
56 nvinfer1::IShuffleLayer *NHWC2NCHW(nvinfer1::INetworkDefinition *network, const nvinfer1::ITensor &input);
57
58 nvinfer1::IShuffleLayer *NCHW2NHWC(nvinfer1::INetworkDefinition *network, const nvinfer1::ITensor &input);
59
60 ActivationParams ConvertActivationType(schema::ActivationType activation_type);
61
62 nvinfer1::ITensor *ConvertConstantTensor(nvinfer1::INetworkDefinition *network, const mindspore::MSTensor &ms_tensor);
63
64 nvinfer1::ITensor *ConvertTensorWithExpandDims(nvinfer1::INetworkDefinition *network,
65 const mindspore::MSTensor &ms_tensor, size_t expand_shape_size);
66
67 nvinfer1::ITensor *ConvertScalarToITensor(nvinfer1::INetworkDefinition *network, size_t shape_size, const void *value,
68 const DataType data_type);
69
70 nvinfer1::Weights TransposeWeight(const mindspore::MSTensor &ms_tensor, void **pack_weight);
71
72 nvinfer1::Weights TransposeWeightFP32(const mindspore::MSTensor &ms_tensor, void **pack_weight);
73
74 nvinfer1::Weights ConvertWeight(const mindspore::MSTensor &ms_tensor);
75
76 void SetCudaDevice(std::shared_ptr<GPUDeviceInfo> device_info_);
77
78 Format GetOutputFormat(Format input_format, nvinfer1::Permutation perm);
79
80 int ConvertAxisFromNHWC2NCHW(int nhwc_axis);
81
82 void PackNHWCToNCHWFp16(const void *src, void *dst, size_t batch, size_t plane, size_t channel, size_t task_id,
83 size_t thread_count);
84
85 std::string GetTensorFormat(nvinfer1::ITensor *trt_tensor, mindspore::Format format);
86
87 template <typename T1, typename T2>
SameDims(const std::vector<T1> & shape1,const std::vector<T2> & shape2)88 bool SameDims(const std::vector<T1> &shape1, const std::vector<T2> &shape2) {
89 if (shape1.size() != shape2.size()) {
90 return false;
91 }
92 for (size_t i = 0; i < shape1.size(); i++) {
93 if (std::abs(shape1[i] - shape2[i]) > 1e-6) {
94 return false;
95 }
96 }
97 return true;
98 }
99
100 template <typename T>
ConvertCudaDims(const std::vector<T> & shape)101 nvinfer1::Dims ConvertCudaDims(const std::vector<T> &shape) {
102 nvinfer1::Dims dims{};
103 if (!shape.empty() && shape.size() <= static_cast<size_t>(dims.MAX_DIMS)) {
104 dims.nbDims = shape.size();
105 for (int i = 0; i < dims.nbDims; i++) {
106 dims.d[i] = static_cast<int>(shape[i]);
107 }
108 } else {
109 MS_LOG(ERROR) << "invalid shape.";
110 }
111 return dims;
112 }
113 } // namespace mindspore::lite
114 #endif // MINDSPORE_LITE_SRC_RUNTIME_DELEGATE_TENSORRT_UTILS_H_
115