|
|
Previously, xla::DeviceMemoryAllocator::Allocate returned a
stream_executor::DeviceMemoryBase. This is morally equivalent to a raw
pointer: It's on you the user to call Deallocate().
Unfortunately we ~never got this right. Essentially all users of
Allocate() call it in a loop, and TF_RETURN_IF_ERROR within the loop.
If any of these allocations fails (mostly commonly, due to OOM), we leak
everything we've allocated up until then.
This patch changes our API so that it returns an owning pointer. Now
things mostly Just Work.
Also worth calling out: The lambda in CpuExecutable::ExecuteOnStream
passed to ExecuteComputeFunction almost certainly had multithreaded
use-after-free bugs. This patch fixes them.
PiperOrigin-RevId: 196000535
|
Justin Lebar