Correct the exact meaning of memory stats collected from the underlying system, and only output max RSS to avoid confusion.

PiperOrigin-RevId: 283023937
Change-Id: I00b76b6c0e8dc059c3cc22a92f76822d47c8e130

tensorflow/lite/profiling/memory_info.cc

@@ -35,20 +35,18 @@ MemoryUsage GetMemoryUsage() {
     result.max_rss_kb = res.ru_maxrss;
   }
   const auto mem = mallinfo();
-  result.total_allocated_bytes = mem.uordblks;
+  result.total_allocated_bytes = mem.arena;
+  result.in_use_allocated_bytes = mem.uordblks;
 #endif
   return result;
 }
 
-void MemoryUsage::SummaryToStream(std::ostream* stream) const {
-  *stream << "memory usage: max resident set size = " << max_rss_kb / 1024.0
+void MemoryUsage::AllStatsToStream(std::ostream* stream) const {
+  *stream << "max resident set size = " << max_rss_kb / 1024.0
           << " MB, total malloc-ed size = "
-          << total_allocated_bytes / 1024.0 / 1024.0 << " MB";
-}
-
-void MemoryUsage::ShortSummaryToStream(std::ostream* stream) const {
-  *stream << "max_rss_mb=" << max_rss_kb / 1024.0
-          << " total_malloced_mb=" << total_allocated_bytes / 1024.0 / 1024.0;
+          << total_allocated_bytes / 1024.0 / 1024.0
+          << " MB, in-use allocated/mmapped size = "
+          << in_use_allocated_bytes / 1024.0 / 1024.0 << " MB";
 }
 
 }  // namespace memory

tensorflow/lite/profiling/memory_info.h

@@ -26,21 +26,30 @@ struct MemoryUsage {
   static const int kValueNotSet;
 
   MemoryUsage()
-      : max_rss_kb(kValueNotSet), total_allocated_bytes(kValueNotSet) {}
+      : max_rss_kb(kValueNotSet),
+        total_allocated_bytes(kValueNotSet),
+        in_use_allocated_bytes(kValueNotSet) {}
 
   // The maximum memory size (in kilobytes) occupied by an OS process that is
   // held in main memory (RAM). Such memory usage information is generally
   // referred as resident set size (rss). This is an alias to rusage::ru_maxrss.
   int64_t max_rss_kb;
 
-  // Total allocated space in bytes. This is an alias to mallinfo::uordblks.
+  // Total non-mmapped space allocated from system in bytes. This is an alias to
+  // mallinfo::arena.
   int total_allocated_bytes;
 
+  // Total allocated (including mmapped) bytes that's in use (i.e. excluding
+  // those are freed). This is an alias to mallinfo::uordblks.
+  int in_use_allocated_bytes;
+
   MemoryUsage operator+(MemoryUsage const& obj) const {
     MemoryUsage res;
     res.max_rss_kb = max_rss_kb + obj.max_rss_kb;
     res.total_allocated_bytes =
         total_allocated_bytes + obj.total_allocated_bytes;
+    res.in_use_allocated_bytes =
+        in_use_allocated_bytes + obj.in_use_allocated_bytes;
     return res;
   }
 
@@ -49,15 +58,16 @@ struct MemoryUsage {
     res.max_rss_kb = max_rss_kb - obj.max_rss_kb;
     res.total_allocated_bytes =
         total_allocated_bytes - obj.total_allocated_bytes;
+    res.in_use_allocated_bytes =
+        in_use_allocated_bytes - obj.in_use_allocated_bytes;
     return res;
   }
 
-  void SummaryToStream(std::ostream* stream) const;
-  void ShortSummaryToStream(std::ostream* stream) const;
+  void AllStatsToStream(std::ostream* stream) const;
 
   friend std::ostream& operator<<(std::ostream& stream,
                                   const MemoryUsage& obj) {
-    obj.SummaryToStream(&stream);
+    obj.AllStatsToStream(&stream);
     return stream;
   }
 };

tensorflow/lite/profiling/memory_info_test.cc

@@ -25,23 +25,28 @@ TEST(MemoryUsage, AddAndSub) {
   MemoryUsage mem1, mem2;
   mem1.max_rss_kb = 5;
   mem1.total_allocated_bytes = 7000;
+  mem1.in_use_allocated_bytes = 2000;
 
   mem2.max_rss_kb = 3;
-  mem2.total_allocated_bytes = 5000;
+  mem2.total_allocated_bytes = 7000;
+  mem2.in_use_allocated_bytes = 4000;
 
   const auto add_mem = mem1 + mem2;
   EXPECT_EQ(8, add_mem.max_rss_kb);
-  EXPECT_EQ(12000, add_mem.total_allocated_bytes);
+  EXPECT_EQ(14000, add_mem.total_allocated_bytes);
+  EXPECT_EQ(6000, add_mem.in_use_allocated_bytes);
 
   const auto sub_mem = mem1 - mem2;
   EXPECT_EQ(2, sub_mem.max_rss_kb);
-  EXPECT_EQ(2000, sub_mem.total_allocated_bytes);
+  EXPECT_EQ(0, sub_mem.total_allocated_bytes);
+  EXPECT_EQ(-2000, sub_mem.in_use_allocated_bytes);
 }
 
 TEST(MemoryUsage, GetMemoryUsage) {
   MemoryUsage result;
   EXPECT_EQ(MemoryUsage::kValueNotSet, result.max_rss_kb);
   EXPECT_EQ(MemoryUsage::kValueNotSet, result.total_allocated_bytes);
+  EXPECT_EQ(MemoryUsage::kValueNotSet, result.in_use_allocated_bytes);
 
 #ifdef __linux__
   // Just allocate some space in heap so that we could meaningful memory usage

tensorflow/lite/tools/benchmark/benchmark_model.cc

@@ -192,8 +192,13 @@ TfLiteStatus BenchmarkModel::Run() {
                              inference_time_us, init_mem_usage,
                              overall_mem_usage});
 
-  TFLITE_LOG(INFO) << "Init " << init_mem_usage << std::endl
-                   << "Overall " << overall_mem_usage;
+  TFLITE_LOG(INFO)
+      << "Note: as the benchmark tool itself affects memory footprint, the "
+         "following is only APPROXIMATE to the actual memory footprint of the "
+         "model at runtime. Take the information at your discretion.";
+  TFLITE_LOG(INFO) << "Peak memory footprint (MB): init="
+                   << init_mem_usage.max_rss_kb / 1024.0
+                   << " overall=" << overall_mem_usage.max_rss_kb / 1024.0;
 
   return status;
 }