experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[SE] Clarify ScratchAllocator and DeviceMemoryAllocator semantics

Browse Source 
PiperOrigin-RevId: 261809474
This commit is contained in:
George Karpenkov 2019-08-05 17:47:10 -07:00 committed by TensorFlower Gardener
parent 3c8582bf36
commit d30d01dba7
2 changed files with 12 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tensorflow/stream_executor/device_memory_allocator.h
View File

@ -147,9 +147,10 @@ class ScopedDeviceMemory {
// Type alias for compatibility with the previous managed memory implementation.
using OwningDeviceMemory = ScopedDeviceMemory<uint8>;
// Interface for device memory allocators used within the XLA service. An
// allocator is responsible for allocating memory on all devices of a particular
// platform.
// Memory allocator interface for the device.
//
// Intended usage is through Allocate() functions which return an owning smart
// pointer.
class DeviceMemoryAllocator {
public:
// Parameter platform indicates which platform the allocator allocates memory
@ -186,7 +187,9 @@ class DeviceMemoryAllocator {
return Allocate(device_ordinal, size, retry_on_failure);
}
// Must be a nop for null pointers.
// Must be a nop for null pointers. Should not be used.
//
// TODO(cheshire): Add deprecation notice.
virtual port::Status Deallocate(int device_ordinal, DeviceMemoryBase mem) = 0;
// Return the platform that the allocator allocates memory on.

14
tensorflow/stream_executor/scratch_allocator.h
View File

@ -27,16 +27,12 @@ namespace stream_executor {
class Stream;
// Interface that allows stream operations (e.g.
// Stream::ThenConvolveWithScratch) to optionally request scratch space be
// allocated in order to speed up the operation being enqueued.
// Interface for "scratch" allocator for device memory, which deallocates all
// buffers it has allocated at destruction. Returned memory pointers are not
// owning.
//
// Note that the caller is responsible for deallocating the scratch space at a
// known-safe point, when all scratch-memory-consuming kernels are known for
// sure to have finished; e.g. at stream synchronization time. This is different
// from a traditional C++ object allocator, where the client is responsible for
// releasing. (Conceptually, scratch memory is a form of "temporary" device
// memory allocation.)
// Used by stream operations (e.g. Stream::ThenConvolveWithScratch) to optonally
// request scratch space to speed up the operation.
class ScratchAllocator {
public:
virtual ~ScratchAllocator();