1 /**
2 * Copyright 2022 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/extendrt/delegate/tensorrt/tensorrt_context.h"
18 #include "src/extendrt/delegate/tensorrt/tensorrt_utils.h"
19
20 namespace mindspore::lite {
~TensorRTContext()21 TensorRTContext::~TensorRTContext() {
22 if (network_ != nullptr) {
23 network_->destroy();
24 network_ = nullptr;
25 }
26 for (auto ptr : owner_memorys_) {
27 free(ptr);
28 }
29 }
30
Init()31 bool TensorRTContext::Init() {
32 network_ = runtime_->GetBuilder()->createNetworkV2(
33 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
34 if (network_ == nullptr) {
35 MS_LOG(ERROR) << "New network init failed.";
36 return false;
37 }
38 return true;
39 }
40
SetRuntime(TensorRTRuntime * runtime)41 void TensorRTContext::SetRuntime(TensorRTRuntime *runtime) { runtime_ = runtime; }
42
network()43 nvinfer1::INetworkDefinition *TensorRTContext::network() { return network_; }
44
RegisterLayer(nvinfer1::ILayer * layer,const std::string & basename)45 void TensorRTContext::RegisterLayer(nvinfer1::ILayer *layer, const std::string &basename) {
46 if (layer == nullptr) {
47 MS_LOG(ERROR) << "Register null layer!";
48 return;
49 }
50 MS_LOG(DEBUG) << "ms_layer " << basename << " register";
51 layer->setName((basename + "_" + std::to_string(counter_++)).c_str());
52 }
53
RegisterTensor(ITensorHelper tensor,const std::string & basename)54 void TensorRTContext::RegisterTensor(ITensorHelper tensor, const std::string &basename) {
55 MS_LOG(DEBUG) << GetTensorFormat(tensor);
56 std::string trt_name = basename + "_" + std::to_string(counter_++);
57 tensor.trt_tensor_->setName(trt_name.c_str());
58 MS_LOG(DEBUG) << "ms_tensor " << basename << " register to " << trt_name;
59 ms_name2trt_tensor_[basename] = tensor;
60 }
61
RegisterTensorWithSameName(ITensorHelper tensor,const std::string & basename)62 void TensorRTContext::RegisterTensorWithSameName(ITensorHelper tensor, const std::string &basename) {
63 std::string trt_name = basename;
64 tensor.trt_tensor_->setName(trt_name.c_str());
65 MS_LOG(DEBUG) << "ms_tensor " << basename << " register to " << trt_name;
66 ms_name2trt_tensor_[basename] = tensor;
67 }
68
HasTensor(const std::string & name) const69 bool TensorRTContext::HasTensor(const std::string &name) const {
70 return ms_name2trt_tensor_.find(name) != ms_name2trt_tensor_.end();
71 }
72
MsName2Tensor(const std::string & ms_name)73 ITensorHelper TensorRTContext::MsName2Tensor(const std::string &ms_name) {
74 if (ms_name2trt_tensor_.find(ms_name) != ms_name2trt_tensor_.end()) {
75 return ms_name2trt_tensor_[ms_name];
76 }
77 MS_LOG(WARNING) << "Get Tensorrt tensor by ms_tensor: " << ms_name << " fail!";
78 return {};
79 }
80
81 template <typename T>
ConvertTo0DTensor(T value)82 nvinfer1::ITensor *TensorRTContext::ConvertTo0DTensor(T value) {
83 void *ptr = malloc(sizeof(T));
84 memcpy(ptr, reinterpret_cast<const void *>(&value), sizeof(T));
85 owner_memorys_.push_back(ptr);
86
87 nvinfer1::Weights weights{GetNvinferDataType<T>(), ptr, 1};
88 nvinfer1::Dims dims{};
89 nvinfer1::IConstantLayer *constant_tensor = network()->addConstant(dims, weights);
90 if (constant_tensor == nullptr) {
91 MS_LOG(ERROR) << "create constant_tensor failed.";
92 return nullptr;
93 }
94 return constant_tensor->getOutput(0);
95 }
96
97 template <typename T>
ConvertTo1DTensor(T value)98 nvinfer1::ITensor *TensorRTContext::ConvertTo1DTensor(T value) {
99 return ConvertTo1DTensor(std::vector<T>{value});
100 }
101
102 template <typename T>
ConvertTo1DTensor(const std::vector<T> & values)103 nvinfer1::ITensor *TensorRTContext::ConvertTo1DTensor(const std::vector<T> &values) {
104 void *ptr = malloc(values.size() * sizeof(T));
105 const T *begin = &values[0];
106 memcpy(ptr, reinterpret_cast<const void *>(begin), values.size() * sizeof(T));
107 owner_memorys_.push_back(ptr);
108
109 nvinfer1::Weights weights{GetNvinferDataType<T>(), ptr, static_cast<int64_t>(values.size())};
110 nvinfer1::Dims dims{1, {static_cast<int>(values.size())}};
111 nvinfer1::IConstantLayer *constant_tensor = network()->addConstant(dims, weights);
112 if (constant_tensor == nullptr) {
113 MS_LOG(ERROR) << "create constant_tensor failed.";
114 return nullptr;
115 }
116 return constant_tensor->getOutput(0);
117 }
118
119 template nvinfer1::ITensor *TensorRTContext::ConvertTo0DTensor(int value);
120 template nvinfer1::ITensor *TensorRTContext::ConvertTo0DTensor(float value);
121 template nvinfer1::ITensor *TensorRTContext::ConvertTo1DTensor(int value);
122 template nvinfer1::ITensor *TensorRTContext::ConvertTo1DTensor(float value);
123 template nvinfer1::ITensor *TensorRTContext::ConvertTo1DTensor(const std::vector<int> &values);
124 template nvinfer1::ITensor *TensorRTContext::ConvertTo1DTensor(const std::vector<float> &values);
125 } // namespace mindspore::lite
126