saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-06-06 00:03:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix id_to_allocator_map

Browse source
This commit is contained in:
Ryan Hill 2021-05-03 19:34:22 -07:00
parent bc7f8ad4f8
commit 1ff2002534
1 changed files with 3 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/python/onnxruntime_pybind_state.cc
View file

@ -484,10 +484,11 @@ static bool IsCudaDeviceIdValid(const onnxruntime::logging::Logger& logger, int
return true;
}
static std::unordered_map<OrtDevice::DeviceId, AllocatorPtr> id_to_allocator_map;
static AllocatorPtr GetCudaAllocator(OrtDevice::DeviceId id) {
// Current approach is not thread-safe, but there are some bigger infra pieces to put together in order to make
// multi-threaded CUDA allocation work we need to maintain a per-thread CUDA allocator
static std::unordered_map<OrtDevice::DeviceId, AllocatorPtr> id_to_allocator_map;
if (id_to_allocator_map.find(id) == id_to_allocator_map.end()) {
id_to_allocator_map.insert({id, GetProviderInfo_CUDA()->CreateCudaAllocator(id, gpu_mem_limit, arena_extend_strategy, external_allocator_info)});
@ -2098,6 +2099,7 @@ PYBIND11_MODULE(onnxruntime_pybind11_state, m) {
}
atexit([] {
id_to_allocator_map.clear();
UnloadSharedProviders();
});
#endif