* updated cmake files for tensorrt
This commit is contained in:
stevenlix 2019-01-23 13:28:13 -08:00 committed by Changming Sun
parent 904d7c6ec8
commit 8ea7197b82
15 changed files with 77 additions and 2 deletions

3
.gitmodules vendored
View file

@ -22,3 +22,6 @@
[submodule "cmake/external/nsync"]
path = cmake/external/nsync
url = https://github.com/google/nsync
[submodule "cmake/external/onnx-tensorrt"]
path = cmake/external/onnx-tensorrt
url = https://github.com/stevenlix/onnx-tensorrt.git

View file

@ -64,6 +64,7 @@ option(onnxruntime_BUILD_SHARED_LIB "Build a shared library" OFF)
option(onnxruntime_ENABLE_MICROSOFT_INTERNAL "Use this option to enable/disable microsoft internal only code" OFF)
option(onnxruntime_USE_NUPHAR "Build with Nupha" OFF)
option(onnxruntime_USE_BRAINSLICE "Build with BrainSlice" OFF)
option(onnxruntime_USE_TRT "Build with TensorRT support" OFF)
set(protobuf_BUILD_TESTS OFF CACHE BOOL "Build protobuf tests" FORCE)
#nsync tests failed on Mac Build

1
cmake/external/onnx-tensorrt vendored Submodule

@ -0,0 +1 @@
Subproject commit 493487a7203d8a059a9b2288807cb700857fe5ca

View file

@ -23,5 +23,6 @@ constexpr const char* kCudaExecutionProvider = "CUDAExecutionProvider";
constexpr const char* kMklDnnExecutionProvider = "MKLDNNExecutionProvider";
constexpr const char* kNupharExecutionProvider = "NupharExecutionProvider";
constexpr const char* kBrainSliceExecutionProvider = "BrainSliceExecutionProvider";
constexpr const char* kTRTExecutionProvider = "TRTExecutionProvider";
} // namespace onnxruntime

View file

@ -131,6 +131,12 @@ common::Status CopyOneInputAcrossDevices(const SessionState& session_state,
ORT_ENFORCE(p_input_provider);
}
//no copy for TRT
if (required_provider_type == onnxruntime::kTRTExecutionProvider) {
new_mlvalue = orig_mlvalue;
return Status::OK();
}
auto input_provider_type = p_input_provider->Type();
if (input_provider_type == required_provider_type && input_tensor_loc.mem_type == OrtMemTypeDefault) {
new_mlvalue = orig_mlvalue;

View file

@ -299,7 +299,8 @@ class InferenceSession::Impl {
for (auto& provider : providers) {
if (provider->Type() != onnxruntime::kCpuExecutionProvider &&
provider->Type() != onnxruntime::kMklDnnExecutionProvider &&
provider->Type() != onnxruntime::kNupharExecutionProvider) {
provider->Type() != onnxruntime::kNupharExecutionProvider &&
provider->Type() != onnxruntime::kTRTExecutionProvider) {
TransformerMemcpyImpl copy_impl(graph, provider->Type());
copy_impl.ModifyGraph(kernel_registry_manager);
}

View file

@ -17,5 +17,12 @@ IExecutionProvider* TestCudaExecutionProvider() {
return &cuda_provider;
}
#endif
#ifdef USE_TRT
IExecutionProvider* TestTRTExecutionProvider() {
static TRTExecutionProvider trt_provider;
return &trt_provider;
}
#endif
} // namespace test
} // namespace onnxruntime

View file

@ -9,6 +9,9 @@
#ifdef USE_CUDA
#include "core/providers/cuda/cuda_execution_provider.h"
#endif
#ifdef USE_TRT
#include "core/providers/trt/trt_execution_provider.h"
#endif
namespace onnxruntime {
namespace test {
@ -18,6 +21,10 @@ IExecutionProvider* TestCPUExecutionProvider();
IExecutionProvider* TestCudaExecutionProvider();
#endif
#ifdef USE_TRT
IExecutionProvider* TestTRTExecutionProvider();
#endif
template <typename T>
void CreateMLValue(AllocatorPtr alloc,
const std::vector<int64_t>& dims,

View file

@ -82,6 +82,7 @@ int real_main(int argc, char* argv[]) {
bool enable_cuda = false;
bool enable_mkl = false;
bool enable_nuphar = false;
bool enable_trt = false;
OrtLoggingLevel logging_level = ORT_LOGGING_LEVEL_WARNING;
{
int ch;
@ -131,6 +132,8 @@ int real_main(int argc, char* argv[]) {
enable_mkl = true;
} else if (!MyStrCmp(optarg, ORT_TSTR("nuphar"))) {
enable_nuphar = true;
} else if (!MyStrCmp(optarg, ORT_TSTR("trt"))) {
enable_trt = true;
} else {
usage();
return -1;
@ -229,6 +232,17 @@ int real_main(int argc, char* argv[]) {
#else
fprintf(stderr, "MKL-DNN is not supported in this build");
return -1;
#endif
}
if (enable_trt) {
#ifdef USE_TRT
OrtProviderFactoryInterface** f;
ORT_THROW_ON_ERROR(OrtCreateTRTExecutionProviderFactory(0, &f));
sf.AppendExecutionProvider(f);
OrtReleaseObject(f);
#else
fprintf(stderr, "TensorRT is not supported in this build");
return -1;
#endif
}
TestEnv args(tests, stat, sf);

View file

@ -62,6 +62,8 @@ namespace perftest {
test_config.machine_config.provider_type_name = onnxruntime::kMklDnnExecutionProvider;
} else if (!strcmp(optarg, "brainslice")) {
test_config.machine_config.provider_type_name = onnxruntime::kBrainSliceExecutionProvider;
} else if (!strcmp(optarg, "trt")) {
test_config.machine_config.provider_type_name = onnxruntime::kTRTExecutionProvider;
} else {
return false;
}

View file

@ -74,6 +74,15 @@ Status SessionFactory::create(std::shared_ptr<::onnxruntime::InferenceSession>&
FACTORY_PTR_HOLDER;
#else
ORT_THROW("This executable was not built with BrainSlice");
#endif
} else if (provider == onnxruntime::kTRTExecutionProvider) {
#if USE_TRT
OrtProviderFactoryInterface** f;
ORT_THROW_ON_ERROR(OrtCreateTRTExecutionProviderFactory(0, &f));
RegisterExecutionProvider(sess.get(), f);
FACTORY_PTR_HOLDER;
#else
ORT_THROW("TensorRT is not supported in this build");
#endif
}
//TODO: add more

View file

@ -71,5 +71,18 @@ std::unique_ptr<IExecutionProvider> DefaultBrainSliceExecutionProvider() {
#endif
}
std::unique_ptr<IExecutionProvider> DefaultTRTExecutionProvider() {
#ifdef USE_TRT
OrtProviderFactoryInterface** f;
ORT_THROW_ON_ERROR(OrtCreateTRTExecutionProviderFactory(0, &f));
FACTORY_PTR_HOLDER;
OrtProvider* out;
ORT_THROW_ON_ERROR((*f)->CreateProvider(f, &out));
return std::unique_ptr<IExecutionProvider>((IExecutionProvider*)out);
#else
return nullptr;
#endif
}
} // namespace test
} // namespace onnxruntime

View file

@ -12,6 +12,7 @@ std::unique_ptr<IExecutionProvider> DefaultCudaExecutionProvider();
std::unique_ptr<IExecutionProvider> DefaultMkldnnExecutionProvider(bool enable_arena = true);
std::unique_ptr<IExecutionProvider> DefaultNupharExecutionProvider();
std::unique_ptr<IExecutionProvider> DefaultBrainSliceExecutionProvider();
std::unique_ptr<IExecutionProvider> DefaultTRTExecutionProvider();
} // namespace test
} // namespace onnxruntime

View file

@ -15,4 +15,7 @@
#endif
#if USE_BRAINSLICE
#include "core/providers/brainslice/brainslice_provider_factory.h"
#endif
#endif
#if USE_TRT
#include "core/providers/trt/trt_provider_factory.h"
#endif

View file

@ -119,6 +119,8 @@ Use the individual flags to only run the specified stages.
parser.add_argument("--brain_slice_package_name", help="Name of brain slice packages")
parser.add_argument("--brain_slice_client_package_name", help="Name of brainslice client package")
parser.add_argument("--use_nuphar", action='store_true', help="Build with nuphar")
parser.add_argument("--use_trt", action='store_true', help="Build with trt")
parser.add_argument("--trt_path", action='store_true', help="Path to trt dir")
return parser.parse_args()
def resolve_executable_path(command_or_path):
@ -297,6 +299,7 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
"-Donnxruntime_USE_BRAINSLICE=" + ("ON" if args.use_brainslice else "OFF"),
"-Donnxruntime_USE_NUPHAR=" + ("ON" if args.use_nuphar else "OFF"),
"-Donnxruntime_USE_EIGEN_THREADPOOL=" + ("ON" if args.use_eigenthreadpool else "OFF"),
"-Donnxruntime_USE_TRT=" + ("ON" if args.use_trt else "OFF"),
]
if args.use_brainslice:
bs_pkg_name = args.brain_slice_package_name.split('.', 1)
@ -306,6 +309,9 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
"-Donnxruntime_BS_CLIENT_PACKAGE=%s/%s" % (args.brain_slice_package_path, args.brain_slice_client_package_name),
"-Donnxruntime_BRAINSLICE_dynamic_lib_PATH=%s/%s" % (args.brain_slice_package_path, bs_shared_lib_name)]
if args.use_trt:
cmake_args += ["-DTENSORRT_ROOT=%s" % args.trt_path]
if args.use_llvm:
cmake_args += ["-DLLVM_DIR=%s" % args.llvm_path]