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 "runtime/device/gpu/trt_loader.h"
18
19 #include <dlfcn.h>
20 #include <memory>
21 #include <NvInferRuntimeCommon.h>
22 #include "backend/kernel_compiler/gpu/trt/trt_utils.h"
23
24 namespace mindspore {
25 namespace device {
26 namespace gpu {
TrtLoader()27 TrtLoader::TrtLoader()
28 : nvinfer_loaded_(false), nvinfer_handle_(nullptr), create_infer_builder_(nullptr), create_infer_runtime_(nullptr) {
29 nvinfer_handle_ = dlopen("libnvinfer.so.7", RTLD_NOW | RTLD_LOCAL);
30 if (nvinfer_handle_ == nullptr) {
31 MS_LOG(WARNING) << "Can not open libnvinfer.so.7 " << dlerror()
32 << ". Install Tensor-RT and export LD_LIBRARY_PATH=${TENSORRT_HOME}/lib:$LD_LIBRARY_PATH.";
33 return;
34 }
35
36 create_infer_builder_ = (CreateInferBuilder_t)dlsym(nvinfer_handle_, "createInferBuilder_INTERNAL");
37 if (create_infer_builder_ == nullptr) {
38 MS_LOG(WARNING) << "Failed to get createInferBuilder_INTERNAL symbol. " << dlerror();
39 return;
40 }
41
42 create_infer_runtime_ = (CreateInferRuntime_t)dlsym(nvinfer_handle_, "createInferRuntime_INTERNAL");
43 if (create_infer_runtime_ == nullptr) {
44 MS_LOG(WARNING) << "Failed to get createInferRuntime_INTERNAL symbol. " << dlerror();
45 return;
46 }
47
48 nvinfer_loaded_ = true;
49 }
50
~TrtLoader()51 TrtLoader::~TrtLoader() {
52 if (nvinfer_handle_ != nullptr) {
53 dlclose(nvinfer_handle_);
54 }
55 }
56
CreateInferBuilder(nvinfer1::ILogger * logger)57 std::shared_ptr<nvinfer1::IBuilder> TrtLoader::CreateInferBuilder(nvinfer1::ILogger *logger) {
58 return TrtPtr<nvinfer1::IBuilder>(create_infer_builder_(*logger, NV_TENSORRT_VERSION));
59 }
60
CreateInferRuntime(nvinfer1::ILogger * logger)61 std::shared_ptr<nvinfer1::IRuntime> TrtLoader::CreateInferRuntime(nvinfer1::ILogger *logger) {
62 return TrtPtr<nvinfer1::IRuntime>(create_infer_runtime_(*logger, NV_TENSORRT_VERSION));
63 }
64 } // namespace gpu
65 } // namespace device
66 } // namespace mindspore
67