[MSLITE] modify compile for tensorrt delegate as gpu

This commit is contained in:
Liu_Xuu 2021-07-02 12:41:13 +08:00 committed by XuuLiu
parent 9055537f62
commit 28430bc1f5
5 changed files with 30 additions and 25 deletions

View File

@ -19,7 +19,8 @@ if(PLATFORM_ARM32 AND CMAKE_CXX_COMPILER_ID STREQUAL "Clang" AND CMAKE_CXX_COMPI
endif()
# Options that can be configured through environment variables or manually
set(MSLITE_GPU_BACKEND "opencl" CACHE STRING "enable gpu backend, only arm64 support, opencl/vulkan/cuda/off")
set(MSLITE_GPU_BACKEND "off" CACHE STRING "enable gpu backend, \
only arm64 support opencl, only x86_64 support tensorrt, opencl/vulkan/cuda/tensorrt/off")
option(MSLITE_ENABLE_NPU "enable npu, only arm64 or arm32 support" on)
option(MSLITE_ENABLE_TRAIN "enable train" on)
option(MSLITE_ENABLE_SSE "enable SSE instruction set, only x86_64 support" off)
@ -28,7 +29,6 @@ option(MSLITE_ENABLE_CONVERTER "enable converter, only x86_64 support" on)
option(MSLITE_ENABLE_TOOLS "enable tools" on)
option(MSLITE_ENABLE_TESTCASES "enable testcase" off)
option(MSLITE_ENABLE_NNIE "enable NNIE" off)
option(MSLITE_ENABLE_TENSORRT "enable TensorRT" off)
# Option that can be configured through manually
option(ENABLE_VERBOSE "" off)
@ -70,8 +70,9 @@ endif()
if(DEFINED ENV{MSLITE_ENABLE_NNIE})
set(MSLITE_ENABLE_NNIE $ENV{MSLITE_ENABLE_NNIE})
endif()
if(DEFINED ENV{MSLITE_ENABLE_TENSORRT})
set(MSLITE_ENABLE_TENSORRT $ENV{MSLITE_ENABLE_TENSORRT})
if(MSLITE_GPU_BACKEND STREQUAL off AND PLATFORM_ARM64)
set(MSLITE_GPU_BACKEND "opencl")
endif()
if(PLATFORM_ARM64 OR PLATFORM_ARM32)
@ -84,7 +85,6 @@ if(PLATFORM_ARM64 OR PLATFORM_ARM32)
set(CMAKE_FIND_ROOT_PATH_MODE_INCLUDE BOTH)
set(CMAKE_FIND_ROOT_PATH_MODE_PACKAGE BOTH)
else()
set(MSLITE_GPU_BACKEND off)
set(MSLITE_ENABLE_NPU off)
endif()
@ -105,20 +105,6 @@ if(MSLITE_ENABLE_NPU)
endif()
endif()
if(MSLITE_ENABLE_TENSORRT)
set(SUPPORT_TENSORRT on)
if(DEFINED ENV{TENSORRT_HOME})
message("TENSORRT_HOME = $ENV{TENSORRT_HOME}")
else()
message(FATAL_ERROR "please set TENSORRT_HOME, example: export TENSORRT_HOME=/root/usr/TensorRT-6.0.1.5/")
endif()
if(DEFINED ENV{CUDA_HOME})
message("CUDA_HOME = $ENV{CUDA_HOME}")
else()
message(FATAL_ERROR "please set CUDA_HOME, example: export CUDA_HOME=/usr/local/cuda-10.1/")
endif()
endif()
message(STATUS "************MindSpore Lite Build Option:************")
message(STATUS "\tMSLITE_GPU_BACKEND = \t${MSLITE_GPU_BACKEND}")
message(STATUS "\tMSLITE_ENABLE_NPU = \t${MSLITE_ENABLE_NPU}")
@ -128,7 +114,6 @@ message(STATUS "\tMSLITE_ENABLE_AVX = \t${MSLITE_ENABLE_AVX}")
message(STATUS "\tMSLITE_ENABLE_CONVERTER = \t${MSLITE_ENABLE_CONVERTER}")
message(STATUS "\tMSLITE_ENABLE_TOOLS = \t${MSLITE_ENABLE_TOOLS}")
message(STATUS "\tMSLITE_ENABLE_TESTCASES = \t${MSLITE_ENABLE_TESTCASES}")
message(STATUS "\tMSLITE_ENABLE_TENSORRT = \t${MSLITE_ENABLE_TENSORRT}")
if(ENABLE_ASAN)
add_definitions(-fsanitize=address -fno-omit-frame-pointer -mllvm -asan-use-private-alias=1)
@ -273,6 +258,20 @@ if(MSLITE_GPU_BACKEND STREQUAL cuda)
add_definitions(-DGPU_CUDA)
add_compile_definitions(SUPPORT_GPU)
endif()
if(MSLITE_GPU_BACKEND STREQUAL tensorrt)
add_compile_definitions(SUPPORT_GPU)
set(SUPPORT_TENSORRT on)
if(DEFINED ENV{TENSORRT_PATH})
message("TENSORRT_PATH = $ENV{TENSORRT_PATH}")
else()
message(FATAL_ERROR "please set TENSORRT_PATH, example: export TENSORRT_PATH=/root/usr/TensorRT-6.0.1.5/")
endif()
if(DEFINED ENV{CUDA_HOME})
message("CUDA_HOME = $ENV{CUDA_HOME}")
else()
message(FATAL_ERROR "please set CUDA_HOME, example: export CUDA_HOME=/usr/local/cuda-10.1/")
endif()
endif()
if(WIN32)
add_compile_definitions(BUILDING_DLL)

View File

@ -200,7 +200,7 @@ endif()
if(SUPPORT_TENSORRT)
add_compile_definitions(GPU_TENSORRT)
set(TENSORRT_PATH $ENV{TENSORRT_HOME})
set(TENSORRT_PATH $ENV{TENSORRT_PATH})
set(CUDA_PATH $ENV{CUDA_HOME})
set(TENSORRT_LIB_PATH ${TENSORRT_PATH}/lib)
set(CUDA_LIB_PATH ${CUDA_PATH}/lib64)

View File

@ -21,7 +21,7 @@
#ifdef SUPPORT_NPU
#include "include/HiAiModelManagerType.h"
#endif
#ifdef SUPPORT_GPU
#ifdef GPU_OPENCL
#include "src/runtime/gpu/opencl/opencl_runtime.h"
#endif
@ -170,7 +170,7 @@ bool InnerContext::IsCpuFloat16Enabled() const {
}
bool InnerContext::IsGpuFloat16Enabled() const {
#ifdef SUPPORT_GPU
#ifdef GPU_OPENCL
if (!IsGpuEnabled()) {
return false;
}

View File

@ -656,7 +656,13 @@ int LiteSession::Init(const Context *context) {
}
#endif
#if GPU_TENSORRT
delegate_ = std::shared_ptr<TensorRTDelegate>(new (std::nothrow) TensorRTDelegate());
if (delegate_ == nullptr && context_->IsGpuEnabled()) {
delegate_ = std::shared_ptr<TensorRTDelegate>(new (std::nothrow) TensorRTDelegate());
if (delegate_ == nullptr) {
MS_LOG(ERROR) << "New tensorrt delegate_ failed";
return RET_ERROR;
}
}
#endif
if (delegate_ != nullptr) {
auto delegate_ret = delegate_->Init();

View File

@ -544,7 +544,7 @@ kernel::LiteKernel *Scheduler::FindBackendKernel(const std::vector<Tensor *> &in
int kernel_thread_count = op_parameter->thread_num_;
op_parameter->is_train_session_ = is_train_session_;
kernel::KernelKey desc{kernel::KERNEL_ARCH::kCPU, data_type, static_cast<schema::PrimitiveType>(op_parameter->type_)};
#ifdef SUPPORT_GPU
#ifdef GPU_OPENCL
if (node->device_type_ == DT_GPU || node->device_type_ == kDefaultDeviceType) {
status = FindGpuKernel(in_tensors, out_tensors, op_parameter, desc, &kernel);
if (status == RET_OK) {