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. // Type alias for compatibility with the previous managed memory implementation.
using OwningDeviceMemory = ScopedDeviceMemory<uint8>; using OwningDeviceMemory = ScopedDeviceMemory<uint8>;
// Interface for device memory allocators used within the XLA service. An // Memory allocator interface for the device.
// allocator is responsible for allocating memory on all devices of a particular //
// platform. // Intended usage is through Allocate() functions which return an owning smart
// pointer.
class DeviceMemoryAllocator { class DeviceMemoryAllocator {
public: public:
// Parameter platform indicates which platform the allocator allocates memory // Parameter platform indicates which platform the allocator allocates memory
@ -186,7 +187,9 @@ class DeviceMemoryAllocator {
return Allocate(device_ordinal, size, retry_on_failure); 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; virtual port::Status Deallocate(int device_ordinal, DeviceMemoryBase mem) = 0;
// Return the platform that the allocator allocates memory on. // 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; class Stream;
// Interface that allows stream operations (e.g. // Interface for "scratch" allocator for device memory, which deallocates all
// Stream::ThenConvolveWithScratch) to optionally request scratch space be // buffers it has allocated at destruction. Returned memory pointers are not
// allocated in order to speed up the operation being enqueued. // owning.
// //
// Note that the caller is responsible for deallocating the scratch space at a // Used by stream operations (e.g. Stream::ThenConvolveWithScratch) to optonally
// known-safe point, when all scratch-memory-consuming kernels are known for // request scratch space to speed up the operation.
// 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.)
class ScratchAllocator { class ScratchAllocator {
public: public:
virtual ~ScratchAllocator(); virtual ~ScratchAllocator();