saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-08 17:17:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix allocator issue for TensorRT IOBinding (#6240)

Browse source 
* 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
This commit is contained in:
Hector Li 2020-12-31 20:15:43 -08:00 committed by GitHub
parent 1685167e46
commit ffb4b62826
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 25 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
include/onnxruntime/core/framework/allocator.h
View file

@ -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";

8
onnxruntime/core/framework/allocator.cc
View file

@ -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<OrtDevice::DeviceId>(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<OrtDevice::DeviceId>(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<OrtDevice::DeviceId>(id1)),
id1, mem_type1);
} else {
return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Specified device is not supported.");
}

7
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
View file

@ -303,9 +303,6 @@ bool CudaCall<cudaError, true>(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

15
onnxruntime/test/shared_lib/test_inference.cc
View file

@ -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);