Add new op BytesLimit.

This is similar to MaxBytesInUse.

Also add "jingyue" to contrib/memory_stats/OWNERS.
Change: 149785681

5 files changed, 58 insertions, 4 deletions

diff --git a/tensorflow/contrib/memory_stats/__init__.py b/tensorflow/contrib/memory_stats/__init__.py
index e310f9fd58..a2b2b65692 100644
--- a/tensorflow/contrib/memory_stats/__init__.py
+++ b/tensorflow/contrib/memory_stats/__init__.py
@@ -14,9 +14,11 @@
 # ==============================================================================
 """Ops for memory statistics.
 
+@@BytesLimit
 @@MaxBytesInUse
 """
 
+from tensorflow.contrib.memory_stats.python.ops.memory_stats_ops import BytesLimit
 from tensorflow.contrib.memory_stats.python.ops.memory_stats_ops import MaxBytesInUse
 
 from tensorflow.python.util.all_util import remove_undocumented
diff --git a/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc b/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc
index c5a4a2de75..23a682fb90 100644
--- a/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc
+++ b/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc
@@ -16,10 +16,10 @@ limitations under the License.
 
 namespace tensorflow {
 
-// Op that measures the peak memory in bytes.
-class MaxBytesInUseOp : public OpKernel {
+// Base class of ops that collects statistics of the allocator of a device.
+class MemoryStatsOp : public OpKernel {
  public:
-  explicit MaxBytesInUseOp(OpKernelConstruction* context) : OpKernel(context) {}
+  explicit MemoryStatsOp(OpKernelConstruction* context) : OpKernel(context) {}
 
   void Compute(OpKernelContext* context) override {
     Allocator* allocator =
@@ -30,7 +30,43 @@ class MaxBytesInUseOp : public OpKernel {
     Tensor* output_tensor = nullptr;
     OP_REQUIRES_OK(
         context, context->allocate_output(0, TensorShape({}), &output_tensor));
-    output_tensor->scalar<int64>()() = allocator_stats.max_bytes_in_use;
+    output_tensor->scalar<int64>()() = ExtractAllocatorStats(allocator_stats);
+  }
+
+ protected:
+  // Extracts a certain field (determined by subclasses) from an allocator
+  // stats.
+  virtual int64 ExtractAllocatorStats(
+      const AllocatorStats& allocator_stats) const = 0;
+};
+
+// Op that measures the total memory (in bytes) of a device.
+class BytesLimitOp : public MemoryStatsOp {
+ public:
+  explicit BytesLimitOp(OpKernelConstruction* context)
+      : MemoryStatsOp(context) {}
+
+ private:
+  int64 ExtractAllocatorStats(
+      const AllocatorStats& allocator_stats) const override {
+    return allocator_stats.bytes_limit;
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("BytesLimit").Device(DEVICE_CPU), BytesLimitOp);
+REGISTER_KERNEL_BUILDER(Name("BytesLimit").Device(DEVICE_GPU).HostMemory("out"),
+                        BytesLimitOp);
+
+// Op that measures the peak memory in bytes.
+class MaxBytesInUseOp : public MemoryStatsOp {
+ public:
+  explicit MaxBytesInUseOp(OpKernelConstruction* context)
+      : MemoryStatsOp(context) {}
+
+ private:
+  int64 ExtractAllocatorStats(
+      const AllocatorStats& allocator_stats) const override {
+    return allocator_stats.max_bytes_in_use;
   }
 };
 
diff --git a/tensorflow/contrib/memory_stats/ops/memory_stats_ops.cc b/tensorflow/contrib/memory_stats/ops/memory_stats_ops.cc
index cdca15388a..d5c6c4ea87 100644
--- a/tensorflow/contrib/memory_stats/ops/memory_stats_ops.cc
+++ b/tensorflow/contrib/memory_stats/ops/memory_stats_ops.cc
@@ -16,6 +16,7 @@ limitations under the License.
 
 namespace tensorflow {
 
+REGISTER_OP("BytesLimit").Output("out: int64").SetIsStateful();
 REGISTER_OP("MaxBytesInUse").Output("out: int64").SetIsStateful();
 
 }  // namespace tensorflow
diff --git a/tensorflow/contrib/memory_stats/python/kernel_tests/memory_stats_ops_test.py b/tensorflow/contrib/memory_stats/python/kernel_tests/memory_stats_ops_test.py
index b0d23b8155..0f3a5f1313 100644
--- a/tensorflow/contrib/memory_stats/python/kernel_tests/memory_stats_ops_test.py
+++ b/tensorflow/contrib/memory_stats/python/kernel_tests/memory_stats_ops_test.py
@@ -29,6 +29,16 @@ from tensorflow.python.platform import test
 
 class MemoryStatsOpsTest(test_util.TensorFlowTestCase):
 
+  def testBytesLimit(self):
+    # AllocatorStats.bytes_limit is set to zero for CPU allocators, so we skip
+    # the check.
+    if not test.is_gpu_available():
+      return
+
+    with self.test_session(use_gpu=True) as sess:
+      bytes_limit = sess.run(memory_stats_ops.BytesLimit())
+      self.assertLess(0, bytes_limit)
+
   # Tests the peak memory usage of the following computation.
   #   a   b
   #   | / |
diff --git a/tensorflow/contrib/memory_stats/python/ops/memory_stats_ops.py b/tensorflow/contrib/memory_stats/python/ops/memory_stats_ops.py
index 60472406a3..d35c6583ed 100644
--- a/tensorflow/contrib/memory_stats/python/ops/memory_stats_ops.py
+++ b/tensorflow/contrib/memory_stats/python/ops/memory_stats_ops.py
@@ -26,6 +26,11 @@ _memory_stats_ops_so = loader.load_op_library(
     resource_loader.get_path_to_datafile("_memory_stats_ops.so"))
 
 
+def BytesLimit():
+  """Generates an op that measures the total memory (in bytes) of a device."""
+  return gen_memory_stats_ops.bytes_limit()
+
+
 def MaxBytesInUse():
   """Generates an op that computes the peak memory of a device."""
   return gen_memory_stats_ops.max_bytes_in_use()