From d30d01dba7c87e30506554a85b55cfea9e55f00c Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Mon, 5 Aug 2019 17:47:10 -0700
Subject: [PATCH] [SE] Clarify ScratchAllocator and DeviceMemoryAllocator
 semantics

PiperOrigin-RevId: 261809474
---
 .../stream_executor/device_memory_allocator.h      | 11 +++++++----
 tensorflow/stream_executor/scratch_allocator.h     | 14 +++++---------
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/tensorflow/stream_executor/device_memory_allocator.h b/tensorflow/stream_executor/device_memory_allocator.h
index 35b6b605a4e..9d309693720 100644
--- a/tensorflow/stream_executor/device_memory_allocator.h
+++ b/tensorflow/stream_executor/device_memory_allocator.h
@@ -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.
diff --git a/tensorflow/stream_executor/scratch_allocator.h b/tensorflow/stream_executor/scratch_allocator.h
index 31278937fe4..29b4e5aa012 100644
--- a/tensorflow/stream_executor/scratch_allocator.h
+++ b/tensorflow/stream_executor/scratch_allocator.h
@@ -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();