From ffb4b628269ca28fc39993d7758ffe1827f9d19f Mon Sep 17 00:00:00 2001 From: Hector Li Date: Thu, 31 Dec 2020 20:15:43 -0800 Subject: [PATCH] Fix allocator issue for TensorRT IOBinding (#6240) * Fix issue: https://github.com/microsoft/onnxruntime/issues/6094 Root cause: we didn't expose the OrtMemoryInfo for TRT, so it will cause issue if user want use IObinding for Tensorrt. Short term fix, add the OrtMemoryInfo for TRT. Long term should unify the allocator for CUDA and TRT --- include/onnxruntime/core/framework/allocator.h | 3 +++ onnxruntime/core/framework/allocator.cc | 8 ++++++++ .../tensorrt/tensorrt_execution_provider.cc | 7 ++----- onnxruntime/test/shared_lib/test_inference.cc | 15 ++++++++++++--- 4 files changed, 25 insertions(+), 8 deletions(-) diff --git a/include/onnxruntime/core/framework/allocator.h b/include/onnxruntime/core/framework/allocator.h index b66548fee6..1059c398f6 100644 --- a/include/onnxruntime/core/framework/allocator.h +++ b/include/onnxruntime/core/framework/allocator.h @@ -22,6 +22,9 @@ namespace onnxruntime { constexpr const char* CPU = "Cpu"; constexpr const char* CUDA = "Cuda"; constexpr const char* CUDA_PINNED = "CudaPinned"; +// TODO: Unify the allocator for CUDA and Tensorrt +constexpr const char* TRT = "Tensorrt"; +constexpr const char* TRT_PINNED = "TensorrtPinned"; constexpr const char* MIGRAPHX = "MIGraphX"; constexpr const char* MIGRAPHX_PINNED = "MIGraphXPinned"; diff --git a/onnxruntime/core/framework/allocator.cc b/onnxruntime/core/framework/allocator.cc index 6967d57f8d..6403e38960 100644 --- a/onnxruntime/core/framework/allocator.cc +++ b/onnxruntime/core/framework/allocator.cc @@ -69,6 +69,14 @@ ORT_API_STATUS_IMPL(OrtApis::CreateMemoryInfo, _In_ const char* name1, enum OrtA *out = new OrtMemoryInfo( onnxruntime::CUDA_PINNED, type, OrtDevice(OrtDevice::CPU, OrtDevice::MemType::CUDA_PINNED, static_cast(id1)), id1, mem_type1); + } else if (strcmp(name1, onnxruntime::TRT) == 0) { + *out = new OrtMemoryInfo( + onnxruntime::TRT, type, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast(id1)), id1, + mem_type1); + } else if (strcmp(name1, onnxruntime::TRT_PINNED) == 0) { + *out = new OrtMemoryInfo( + onnxruntime::TRT_PINNED, type, OrtDevice(OrtDevice::CPU, OrtDevice::MemType::CUDA_PINNED, static_cast(id1)), + id1, mem_type1); } else { return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Specified device is not supported."); } diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index c5cd8ee9bd..993530a0a3 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -303,9 +303,6 @@ bool CudaCall(cudaError retCode, const char* exprString, const return g_host->CudaCall_true(retCode, exprString, libName, successCode, msg); } -constexpr const char* TRT = "Tensorrt"; -constexpr const char* TRT_PINNED = "TensorrtPinned"; - class Memcpy final : public Provider_OpKernel { public: Memcpy(const OpKernelInfo&) {} @@ -387,12 +384,12 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv CUDA_CALL_THROW(cudaSetDevice(device_id_)); AllocatorCreationInfo default_memory_info( - [](int id) { return CreateCUDAAllocator(id, TRT); }, device_id_); + [](int id) { return CreateCUDAAllocator(id, onnxruntime::TRT); }, device_id_); allocator_ = CreateAllocator(default_memory_info); InsertAllocator(allocator_); AllocatorCreationInfo pinned_allocator_info( - [](int) { return CreateCUDAPinnedAllocator(0, TRT_PINNED); }, device_id_); + [](int) { return CreateCUDAPinnedAllocator(0, onnxruntime::TRT_PINNED); }, device_id_); InsertAllocator(CreateAllocator(pinned_allocator_info)); // Get environment variables diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index 737f5b6a0e..a36d66eae1 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -729,7 +729,7 @@ TEST(CApiTest, io_binding) { binding.ClearBoundOutputs(); } -#ifdef USE_CUDA +#if defined(USE_CUDA) || defined(USE_TENSORRT) TEST(CApiTest, io_binding_cuda) { struct CudaMemoryDeleter { explicit CudaMemoryDeleter(const Ort::Allocator* alloc) { @@ -743,10 +743,19 @@ TEST(CApiTest, io_binding_cuda) { }; Ort::SessionOptions session_options; - Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0)); + #ifdef USE_TENSORRT + Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_Tensorrt(session_options, 0)); + #else + Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0)); + #endif Ort::Session session(*ort_env, MODEL_URI, session_options); - Ort::MemoryInfo info_cuda("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemTypeDefault); + #ifdef USE_TENSORRT + Ort::MemoryInfo info_cuda("Tensorrt", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemTypeDefault); + #else + Ort::MemoryInfo info_cuda("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemTypeDefault); + #endif + Ort::Allocator cuda_allocator(session, info_cuda); auto allocator_info = cuda_allocator.GetInfo(); ASSERT_TRUE(info_cuda == allocator_info);