Add peak private footprint memory measurement tracking

PiperOrigin-RevId: 820398210

Author: tf-marissaw

tflite/profiling/memory_info.cc

@@ -16,7 +16,12 @@ limitations under the License.
 
 #include <stddef.h>
 
+#include <cstdint>
+#include <fstream>
+#include <iostream>
 #include <ostream>
+#include <sstream>
+#include <string>
 
 #ifdef __linux__
 #include <malloc.h>
@@ -37,6 +42,38 @@ namespace memory {
 
 const size_t MemoryUsage::kValueNotSet = 0;
 
+namespace {
+
+#if defined(__linux__)
+// Returns the current VM swap in kilobytes on Linux.
+int64_t GetCurrentVmSwapKb() {
+  std::ifstream status_file("/proc/self/status");
+  if (!status_file.is_open()) {
+    return -1;
+  }
+  std::string line;
+  while (std::getline(status_file, line)) {
+    if (line.rfind("VmSwap:", 0) == 0) {
+      std::stringstream ss(line);
+      std::string key;
+      int64_t value_kb;
+      // The line format is "VmSwap:    1234 kB"
+      // We can extract the key ("VmSwap:") and the numeric value ("1234").
+      ss >> key >> value_kb;
+      if (!ss.fail()) {
+        return value_kb;
+      } else {
+        return -1;  // Indicate parsing error
+      }
+    }
+  }
+  // If the VmSwap line is not found, it means 0 swap is being used.
+  return 0;
+}
+#endif
+
+}  // namespace
+
 bool MemoryUsage::IsSupported() {
 #if defined(__linux__) || defined(__APPLE__) || defined(_WIN32)
   return true;
@@ -50,6 +87,10 @@ MemoryUsage GetMemoryUsage() {
   rusage res;
   if (getrusage(RUSAGE_SELF, &res) == 0) {
     result.mem_footprint_kb = res.ru_maxrss;
+    int64_t vm_swap_kb = GetCurrentVmSwapKb();
+    if (vm_swap_kb >= 0) {
+      result.private_footprint_bytes = (vm_swap_kb + res.ru_maxrss) * 1024;
+    }
   }
 #if defined(__NO_MALLINFO__)
   result.total_allocated_bytes = -1;
@@ -71,19 +112,24 @@ MemoryUsage GetMemoryUsage() {
   if (status == KERN_SUCCESS) {
     result.mem_footprint_kb =
         static_cast<int64_t>(vm_info.phys_footprint / 1024.0);
+    // TODO: b/421171145 - Consider subtracting shared_resident_kb.
+    result.private_footprint_bytes = vm_info.phys_footprint;
   }
   struct mstats stats = mstats();
   result.total_allocated_bytes = stats.bytes_total;
   result.in_use_allocated_bytes = stats.bytes_used;
 #elif defined(_WIN32)
-  PROCESS_MEMORY_COUNTERS process_memory_counters;
+  PROCESS_MEMORY_COUNTERS_EX process_memory_counters;
   HANDLE process_handle = GetCurrentProcess();
   if (process_handle != nullptr &&
-      GetProcessMemoryInfo(process_handle, &process_memory_counters,
+      GetProcessMemoryInfo(process_handle,
+                           (PROCESS_MEMORY_COUNTERS*)&process_memory_counters,
                            sizeof(process_memory_counters))) {
     result.mem_footprint_kb = process_memory_counters.WorkingSetSize / 1024;
+    result.private_footprint_bytes = process_memory_counters.PrivateUsage;
   } else {
     result.mem_footprint_kb = -1;
+    result.private_footprint_bytes = -1;
   }
   CloseHandle(process_handle);
   HANDLE process_heap = GetProcessHeap();
@@ -108,7 +154,9 @@ void MemoryUsage::AllStatsToStream(std::ostream* stream) const {
           << mem_footprint_kb / 1000.0 << " MB, total non-mmapped heap size = "
           << total_allocated_bytes / 1000.0 / 1000.0
           << " MB, in-use heap size = "
-          << in_use_allocated_bytes / 1000.0 / 1000.0 << " MB";
+          << in_use_allocated_bytes / 1000.0 / 1000.0
+          << " MB, private footprint = "
+          << private_footprint_bytes / 1000.0 / 1000.0 << " MB";
 }
 
 }  // namespace memory

tflite/profiling/memory_info.h

@@ -35,7 +35,8 @@ struct MemoryUsage {
   MemoryUsage()
       : mem_footprint_kb(kValueNotSet),
         total_allocated_bytes(kValueNotSet),
-        in_use_allocated_bytes(kValueNotSet) {}
+        in_use_allocated_bytes(kValueNotSet),
+        private_footprint_bytes(kValueNotSet) {}
 
   // The memory footprint (in kilobytes).
   //
@@ -85,13 +86,22 @@ struct MemoryUsage {
   // non-heap uses of memory such as thread stacks, globals, code, etc.
   size_t in_use_allocated_bytes;
 
+  // Private footprint (in kilobytes).
+  //
+  // For Linux this is the rusage::ru_maxrss + VmSwap.
+  // For Mac this is the task_vm_info::phys_footprint.
+  // For Windows this is the PrivateUsage.
+  size_t private_footprint_bytes;
+
   MemoryUsage operator+(MemoryUsage const& obj) const {
     MemoryUsage res;
     res.mem_footprint_kb = mem_footprint_kb + obj.mem_footprint_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;
+    res.private_footprint_bytes =
+        private_footprint_bytes + obj.private_footprint_bytes;
     return res;
   }
 
@@ -102,6 +112,8 @@ struct MemoryUsage {
         total_allocated_bytes - obj.total_allocated_bytes;
     res.in_use_allocated_bytes =
         in_use_allocated_bytes - obj.in_use_allocated_bytes;
+    res.private_footprint_bytes =
+        private_footprint_bytes - obj.private_footprint_bytes;
     return res;
   }
 

tflite/profiling/memory_info_test.cc

@@ -30,27 +30,32 @@ TEST(MemoryUsage, AddAndSub) {
   mem1.mem_footprint_kb = 5;
   mem1.total_allocated_bytes = 7000;
   mem1.in_use_allocated_bytes = 2000;
+  mem1.private_footprint_bytes = 1000;
 
   mem2.mem_footprint_kb = 3;
   mem2.total_allocated_bytes = 7000;
   mem2.in_use_allocated_bytes = 4000;
+  mem2.private_footprint_bytes = 500;
 
   const auto add_mem = mem1 + mem2;
   EXPECT_EQ(8, add_mem.mem_footprint_kb);
   EXPECT_EQ(14000, add_mem.total_allocated_bytes);
   EXPECT_EQ(6000, add_mem.in_use_allocated_bytes);
+  EXPECT_EQ(1500, add_mem.private_footprint_bytes);
 
   const auto sub_mem = mem1 - mem2;
   EXPECT_EQ(2, sub_mem.mem_footprint_kb);
   EXPECT_EQ(0, sub_mem.total_allocated_bytes);
   EXPECT_EQ(-2000, sub_mem.in_use_allocated_bytes);
+  EXPECT_EQ(500, sub_mem.private_footprint_bytes);
 }
 
 TEST(MemoryUsage, GetMemoryUsage) {
   MemoryUsage result;
   EXPECT_EQ(MemoryUsage::kValueNotSet, result.mem_footprint_kb);
   EXPECT_EQ(MemoryUsage::kValueNotSet, result.total_allocated_bytes);
   EXPECT_EQ(MemoryUsage::kValueNotSet, result.in_use_allocated_bytes);
+  EXPECT_EQ(MemoryUsage::kValueNotSet, result.private_footprint_bytes);
 
 #if defined(__linux__) || defined(__APPLE__) || defined(_WIN32)
   // Just allocate some space in heap so that we have some meaningful
@@ -72,6 +77,7 @@ TEST(MemoryUsage, GetMemoryUsage) {
   EXPECT_GE(result.mem_footprint_kb, size / 1024);
   EXPECT_GE(result.total_allocated_bytes, size);
   EXPECT_GE(result.in_use_allocated_bytes, size);
+  EXPECT_GE(result.private_footprint_bytes, size);
 #endif
 }
 

tflite/profiling/memory_latency_logger.cc

@@ -80,10 +80,16 @@ void MemoryLatencyLogger::Stop(absl::string_view log_message) {
                    << " MB,";
   }
   if (mem_monitor_->GetCurrentInUseMemoryInMB() < 0) {
-    message_stream << " current in-use: unknown";
+    message_stream << " current in-use: unknown,";
   } else {
     message_stream << " current in-use: "
-                   << mem_monitor_->GetCurrentInUseMemoryInMB() << " MB";
+                   << mem_monitor_->GetCurrentInUseMemoryInMB() << " MB,";
+  }
+  if (mem_monitor_->GetPeakPrivateFootprintInMB() < 0) {
+    message_stream << " peak private: unknown";
+  } else {
+    message_stream << " peak private: "
+                   << mem_monitor_->GetPeakPrivateFootprintInMB() << " MB";
   }
   TFLITE_LOG(INFO) << message_stream.str();
 }

tflite/profiling/memory_usage_monitor.cc

@@ -62,6 +62,11 @@ void MemoryUsageMonitor::Start() {
       if (current_in_use_bytes > peak_in_use_mem_bytes_) {
         peak_in_use_mem_bytes_ = current_in_use_bytes;
       }
+      int64_t current_private_footprint_bytes =
+          mem_info.private_footprint_bytes;
+      if (current_private_footprint_bytes > peak_private_footprint_bytes_) {
+        peak_private_footprint_bytes_ = current_private_footprint_bytes;
+      }
       if (stop_signal_->HasBeenNotified()) break;
       sampler_->SleepFor(sampling_interval_);
     }

tflite/profiling/memory_usage_monitor.h

@@ -87,6 +87,13 @@ class MemoryUsageMonitor {
     return BytesToMegabytes(peak_in_use_mem_bytes_);
   }
 
+  float GetPeakPrivateFootprintInMB() const {
+    if (!is_supported_ || check_memory_thd_ != nullptr) {
+      return kInvalidMemUsageMB;
+    }
+    return BytesToMegabytes(peak_private_footprint_bytes_);
+  }
+
   MemoryUsageMonitor(MemoryUsageMonitor&) = delete;
   MemoryUsageMonitor& operator=(const MemoryUsageMonitor&) = delete;
   MemoryUsageMonitor(MemoryUsageMonitor&&) = delete;
@@ -105,6 +112,7 @@ class MemoryUsageMonitor {
   std::unique_ptr<std::thread> check_memory_thd_ = nullptr;
   int64_t peak_mem_footprint_bytes_ = kInvalidMemUsageBytes;
   int64_t peak_in_use_mem_bytes_ = kInvalidMemUsageBytes;
+  int64_t peak_private_footprint_bytes_ = kInvalidMemUsageBytes;
 };
 
 }  // namespace memory