!22112 [MSLITE] fix allocator pool when load different network with same tensor name for tensorrt

Merge pull request !22112 from Liu_Xuu/trt_0820_bugs
This commit is contained in:
i-robot 2021-08-24 03:24:50 +00:00 committed by Gitee
commit 723109041a
4 changed files with 11 additions and 34 deletions

View File

@ -5,21 +5,15 @@ file(GLOB_RECURSE TENSORRT_RUNTIME_SRC
${CMAKE_CURRENT_SOURCE_DIR}/op/*.cc
${CMAKE_CURRENT_SOURCE_DIR}/../delegate_utils.cc
)
link_libraries(${CUDA_LIB_PATH}/libcudnn.so)
link_libraries(${CUDA_LIB_PATH}/libnvrtc.so)
link_libraries(${CUDA_LIB_PATH}/libcublas.so)
link_libraries(${CUDA_LIB_PATH}/libcublasLt.so)
add_library(libcudart SHARED IMPORTED)
set_target_properties(libcudart PROPERTIES IMPORTED_LOCATION ${CUDA_LIB_PATH}/libcudart.so)
add_library(libcudnn SHARED IMPORTED)
set_target_properties(libcudnn PROPERTIES IMPORTED_LOCATION ${CUDA_LIB_PATH}/libcudnn.so)
add_library(libnvrtc SHARED IMPORTED)
set_target_properties(libnvrtc PROPERTIES IMPORTED_LOCATION ${CUDA_LIB_PATH}/libnvrtc.so)
add_library(libcublas SHARED IMPORTED)
set_target_properties(libcublas PROPERTIES IMPORTED_LOCATION ${CUDA_LIB_PATH}/libcublas.so)
add_library(libcublasLt SHARED IMPORTED)
set_target_properties(libcublasLt PROPERTIES IMPORTED_LOCATION ${CUDA_LIB_PATH}/libcublasLt.so)
add_library(libnvinfer SHARED IMPORTED)
set_target_properties(libnvinfer PROPERTIES IMPORTED_LOCATION ${TENSORRT_LIB_PATH}/libnvinfer.so)
@ -28,9 +22,5 @@ add_dependencies(tensorrt_kernel_mid fbs_src)
target_link_libraries(
tensorrt_kernel_mid
libcudart
libcudnn
libnvrtc
libcublas
libcublasLt
libnvinfer
)

View File

@ -19,16 +19,6 @@
#include <vector>
namespace mindspore::lite {
static std::mutex g_mtx;
TensorRTRuntime *TensorRTRuntime::cuda_runtime_instance_ = nullptr;
TensorRTRuntime *TensorRTRuntime::GetInstance() {
std::unique_lock<std::mutex> lck(g_mtx);
static TensorRTRuntime cuda_runtime;
if (cuda_runtime_instance_ == nullptr) {
cuda_runtime_instance_ = &cuda_runtime;
}
return cuda_runtime_instance_;
}
int TensorRTRuntime::Init() {
if (is_init_) {
return RET_OK;
@ -39,7 +29,6 @@ int TensorRTRuntime::Init() {
return RET_ERROR;
}
builder_->setMaxBatchSize(MAX_BATCH_SIZE);
allocator_ = new (std::nothrow) TensorRTAllocator();
if (allocator_ == nullptr) {
MS_LOG(ERROR) << "Create allocator failed.";

View File

@ -45,8 +45,6 @@ class TensorRTRuntime {
~TensorRTRuntime();
static TensorRTRuntime *GetInstance();
int Init();
nvinfer1::IBuilder *GetBuilder() { return this->builder_; }
@ -58,7 +56,6 @@ class TensorRTRuntime {
TensorRTAllocator *GetAllocator() { return this->allocator_; }
private:
static TensorRTRuntime *cuda_runtime_instance_;
bool is_init_ = false;
nvinfer1::IBuilder *builder_{nullptr};
TensorRTLogger logger_;

View File

@ -54,7 +54,9 @@ int TensorRTSubGraph::Init() {
MS_LOG(ERROR) << "Get NPU subgraph input and output ops failed.";
return RET_ERROR;
}
runtime_ = TensorRTRuntime::GetInstance();
if (runtime_ == nullptr) {
runtime_ = new (std::nothrow) TensorRTRuntime();
}
ret = runtime_->Init();
if (ret != RET_OK) {
MS_LOG(ERROR) << "TensorRTRuntime init failed.";
@ -98,9 +100,8 @@ int TensorRTSubGraph::SetDeviceConfig() {
if (device_info_->GetEnableFP16() && SupportFP16()) {
config_->setFlag(nvinfer1::BuilderFlag::kFP16);
}
// config setMaxWorkspaceSize to 256 MB for max limit
config_->setMaxWorkspaceSize(256 * (1 << 20));
// config setMaxWorkspaceSize to 128 MB for max limit
config_->setMaxWorkspaceSize(32 * (1 << 20));
return RET_OK;
}