[VM][PooledAllocator] try reallocation once when OOM

[ganler]

See TVM Discussion Topic.

This change aims to make TVM behaviour more robust when an OOM occurs and resolve a mysterious exception-uncaught bug.

Potential reviewers: @junrushao1994 @icemelon9 @jroesch

[YuchenJin]

LGTM! Thanks @ganler for the contribution!

[comaniac]

What to expect if it still failed here?

[ganler]

If it still fails, an InternalError will be thrown, causing a TVMError regarding OOM in the Python End.

[comaniac]

Hmm, would that be better if we let ReleaseAll return the size it released and check if it is larger than the requested size? So that we can directly throw a message including both sizes without calling alloc again.

[ganler]

Thanks for the suggestion. But IMHO this is not robust enough.

Say that we have 8 GB GPU memory, the PooledAllocator cached 4 GB and we want to allocate 6 GB.

• Applying your idea, ReleaseAll() returns "4GB" which is less than "6GB", thus resulting in a failed allocation.
• Instead, if we release unused memory and do re-allocation, "6GB" is very likely to be successfully allocated.

The big picture behind your idea is practical if we can have some APIs like "total_system_memory" and "available_system_memory", which may require introducing a series of runtime/driver libraries. e.g., cudaMemGetInfo by CudaRT (user space) or NVML (if some system privilege is allowed).

[comaniac]

Fair enough. I don't have other comments then.

[masahi]

This change doesn't solve the issue in #8233, because AllocDataSpace can be called from NDArray::Empty:

tvm/src/runtime/ndarray.cc

Lines 196 to 197 in 4d9bc9b

	DeviceAPI::Get(ret->device)
	->AllocDataSpace(ret->device, shape.size(), shape.data(), ret->dtype, mem_scope);

That call is not protected by try/catch, so if almost all memory are held by PooledAllocator and NDArray::Empty is called, the program crashes with the following error:

terminate called after throwing an instance of 'tvm::runtime::InternalError'
  what():  [19:12:54] /home/masa/projects/dev/tvm/src/runtime/vulkan/vulkan_stream.cc:123: 
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------
  Check failed: (__e == VK_SUCCESS) is false: Vulkan Error, code=-13: Unknown Vulkan error code
Stack trace:
  0: tvm::runtime::vulkan::VulkanStream::Synchronize()
  1: _ZN3tvm7runtime6vulkan15VulkanDeviceAPI13FreeDataSpac
  2: tvm::runtime::NDArray::Internal::DefaultDeleter(tvm::runtime::Object*)
  3: tvm::runtime::NDArray::CopyTo(DLDevice const&) const
  4: tvm::runtime::vm::CopyTo(tvm::runtime::ObjectRef, DLDevice const&)
  5: std::_Function_handler<void (tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*), tvm::runtime::vm::VirtualMachine::GetFunction(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, tvm::runtime::ObjectPtr<tvm::runtime::Object> const&)::$_6>::_M_invoke(std::_Any_data const&, tvm::runtime::TVMArgs&&, tvm::runtime::TVMRetValue*&&)
  6: TVMFuncCall

I think we need to revisit the memory release strategy of PooledAllocator.