• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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