Record persistent memory in tracking allocator and cost model.

Change: 147123056

5 files changed, 34 insertions, 24 deletions

diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc
index 8eacbf82db..2261396819 100644
--- a/tensorflow/core/common_runtime/executor.cc
+++ b/tensorflow/core/common_runtime/executor.cc
@@ -161,10 +161,10 @@ void SetMemory(NodeExecStats* nt, OpKernelContext* ctx) {
     // be dereferenced again after this statement
     auto sizes = allocator_pair.second->GetSizesAndUnRef();
     memory->set_allocator_name(allocator_pair.first->Name());
-    int tb = sizes.first;
-    memory->set_total_bytes(tb);
+    memory->set_total_bytes(std::get<0>(sizes));
     if (allocator_pair.first->TracksAllocationSizes()) {
-      memory->set_peak_bytes(sizes.second);
+      memory->set_peak_bytes(std::get<1>(sizes));
+      memory->set_live_bytes(std::get<2>(sizes));
     }
   }
 }
diff --git a/tensorflow/core/framework/step_stats.proto b/tensorflow/core/framework/step_stats.proto
index 4488f985c7..ed43b7a585 100644
--- a/tensorflow/core/framework/step_stats.proto
+++ b/tensorflow/core/framework/step_stats.proto
@@ -17,6 +17,8 @@ message AllocatorMemoryUsed {
   string allocator_name = 1;
   int64 total_bytes = 2;
   int64 peak_bytes = 3;
+  // The bytes that are not deallocated.
+  int64 live_bytes = 4;
 }
 
 // Output sizes recorded for a single execution of a graph node.
diff --git a/tensorflow/core/framework/tracking_allocator.cc b/tensorflow/core/framework/tracking_allocator.cc
index 302bb1d981..1052ac0554 100644
--- a/tensorflow/core/framework/tracking_allocator.cc
+++ b/tensorflow/core/framework/tracking_allocator.cc
@@ -151,20 +151,22 @@ void TrackingAllocator::GetStats(AllocatorStats* stats) {
   allocator_->GetStats(stats);
 }
 
-std::pair<size_t, size_t> TrackingAllocator::GetSizesAndUnRef() {
+std::tuple<size_t, size_t, size_t> TrackingAllocator::GetSizesAndUnRef() {
   size_t high_watermark;
   size_t total_bytes;
+  size_t still_live_bytes;
   bool should_delete;
   {
     mutex_lock lock(mu_);
     high_watermark = high_watermark_;
     total_bytes = total_bytes_;
+    still_live_bytes = allocated_;
     should_delete = UnRef();
   }
   if (should_delete) {
     delete this;
   }
-  return std::make_pair(total_bytes, high_watermark);
+  return std::make_tuple(total_bytes, high_watermark, still_live_bytes);
 }
 
 bool TrackingAllocator::UnRef() {
diff --git a/tensorflow/core/framework/tracking_allocator.h b/tensorflow/core/framework/tracking_allocator.h
index f92ecc670f..92c89d30ac 100644
--- a/tensorflow/core/framework/tracking_allocator.h
+++ b/tensorflow/core/framework/tracking_allocator.h
@@ -59,19 +59,20 @@ class TrackingAllocator : public Allocator {
   void GetStats(AllocatorStats* stats) override;
 
   // If the underlying allocator tracks allocation sizes, this returns
-  // a pair where the first value is the total number of bytes
-  // allocated through this wrapper, and the second value is the high
-  // watermark of bytes allocated through this wrapper. If the
+  // a tuple where the first value is the total number of bytes
+  // allocated through this wrapper, the second value is the high
+  // watermark of bytes allocated through this wrapper and the third value is
+  // the allocated bytes through this wrapper that are still alive. If the
   // underlying allocator does not track allocation sizes the first
   // value is the total number of bytes requested through this wrapper
-  // and the second is 0.
+  // and the second and the third are 0.
   //
   // After GetSizesAndUnref is called, the only further calls allowed
   // on this wrapper are calls to DeallocateRaw with pointers that
   // were allocated by this wrapper and have not yet been
   // deallocated. After this call completes and all allocated pointers
   // have been deallocated the wrapper will delete itself.
-  std::pair<size_t, size_t> GetSizesAndUnRef();
+  std::tuple<size_t, size_t, size_t> GetSizesAndUnRef();
 
  protected:
   ~TrackingAllocator() override {}
diff --git a/tensorflow/core/framework/tracking_allocator_test.cc b/tensorflow/core/framework/tracking_allocator_test.cc
index 850cdc3909..ae440cc28b 100644
--- a/tensorflow/core/framework/tracking_allocator_test.cc
+++ b/tensorflow/core/framework/tracking_allocator_test.cc
@@ -75,10 +75,11 @@ TEST(TrackingAllocatorTest, SimpleNoTracking) {
   ta->DeallocateRaw(p1);
   void* p2 = ta->AllocateRaw(4, 12);
 
-  std::pair<size_t, size_t> sizes = ta->GetSizesAndUnRef();
+  std::tuple<size_t, size_t, size_t> sizes = ta->GetSizesAndUnRef();
 
-  EXPECT_EQ(16, sizes.first);
-  EXPECT_EQ(0, sizes.second);
+  EXPECT_EQ(16, std::get<0>(sizes));
+  EXPECT_EQ(0, std::get<1>(sizes));
+  EXPECT_EQ(0, std::get<2>(sizes));
 
   ta->DeallocateRaw(p2);
 
@@ -97,8 +98,9 @@ TEST(TrackingAllocatorTest, SimpleNoTracking) {
 
   sizes = ta->GetSizesAndUnRef();
 
-  EXPECT_LE(16, sizes.first);
-  EXPECT_LE(12, sizes.second);
+  EXPECT_LE(16, std::get<0>(sizes));
+  EXPECT_LE(12, std::get<1>(sizes));
+  EXPECT_LE(12, std::get<2>(sizes));
 
   ta->DeallocateRaw(p2);
 }
@@ -114,10 +116,11 @@ TEST(TrackingAllocatorTest, SimpleTracking) {
   ta->DeallocateRaw(p1);
   void* p2 = ta->AllocateRaw(4, 4);
 
-  std::pair<size_t, size_t> sizes = ta->GetSizesAndUnRef();
+  std::tuple<size_t, size_t, size_t> sizes = ta->GetSizesAndUnRef();
 
-  EXPECT_EQ(16, sizes.first);
-  EXPECT_EQ(12, sizes.second);
+  EXPECT_EQ(16, std::get<0>(sizes));
+  EXPECT_EQ(12, std::get<1>(sizes));
+  EXPECT_EQ(4, std::get<2>(sizes));
 
   ta->DeallocateRaw(p2);
 }
@@ -132,10 +135,11 @@ TEST(TrackingAllocatorTest, OutOfMemory) {
   void* p1 = ta->AllocateRaw(4, 12);
   EXPECT_EQ(nullptr, p1);
 
-  std::pair<size_t, size_t> sizes = ta->GetSizesAndUnRef();
+  std::tuple<size_t, size_t, size_t> sizes = ta->GetSizesAndUnRef();
 
-  EXPECT_EQ(0, sizes.first);
-  EXPECT_EQ(0, sizes.second);
+  EXPECT_EQ(0, std::get<0>(sizes));
+  EXPECT_EQ(0, std::get<1>(sizes));
+  EXPECT_EQ(0, std::get<2>(sizes));
 }
 
 TEST(TrackingAllocatorTest, FreeNullPtr) {
@@ -147,10 +151,11 @@ TEST(TrackingAllocatorTest, FreeNullPtr) {
 
   ta->DeallocateRaw(nullptr);
 
-  std::pair<size_t, size_t> sizes = ta->GetSizesAndUnRef();
+  std::tuple<size_t, size_t, size_t> sizes = ta->GetSizesAndUnRef();
 
-  EXPECT_EQ(0, sizes.first);
-  EXPECT_EQ(0, sizes.second);
+  EXPECT_EQ(0, std::get<0>(sizes));
+  EXPECT_EQ(0, std::get<1>(sizes));
+  EXPECT_EQ(0, std::get<2>(sizes));
 }
 
 }  // namespace tensorflow