Support Windows memory info

Memory info is a useful class for recording memory usage across
a variety of devices. Add Windows support to it.

PiperOrigin-RevId: 820360533

Author: tf-marissaw

tflite/profiling/memory_info.cc

@@ -25,6 +25,10 @@ limitations under the License.
 #elif defined(__APPLE__)
 #include <mach/mach.h>
 #include <malloc/malloc.h>
+#elif defined(_WIN32)
+#include <windows.h>
+// psapi must be included after windows.h.
+#include <psapi.h>
 #endif
 
 namespace tflite {
@@ -34,7 +38,7 @@ namespace memory {
 const size_t MemoryUsage::kValueNotSet = 0;
 
 bool MemoryUsage::IsSupported() {
-#if defined(__linux__) || defined(__APPLE__)
+#if defined(__linux__) || defined(__APPLE__) || defined(_WIN32)
   return true;
 #endif
   return false;
@@ -71,6 +75,30 @@ MemoryUsage GetMemoryUsage() {
   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;
+  HANDLE process_handle = GetCurrentProcess();
+  if (process_handle != nullptr &&
+      GetProcessMemoryInfo(process_handle, &process_memory_counters,
+                           sizeof(process_memory_counters))) {
+    result.mem_footprint_kb = process_memory_counters.WorkingSetSize / 1024;
+  } else {
+    result.mem_footprint_kb = -1;
+  }
+  CloseHandle(process_handle);
+  HANDLE process_heap = GetProcessHeap();
+  if (process_heap != nullptr && HeapLock(process_heap)) {
+    HEAP_SUMMARY heap_summary;
+    heap_summary.cb = sizeof(heap_summary);
+    if (HeapSummary(process_heap, 0, &heap_summary)) {
+      result.total_allocated_bytes = heap_summary.cbCommitted;
+      result.in_use_allocated_bytes = heap_summary.cbAllocated;
+    } else {
+      result.total_allocated_bytes = -1;
+      result.in_use_allocated_bytes = -1;
+    }
+    HeapUnlock(process_heap);
+  }
 #endif  // __linux__
   return result;
 }

tflite/profiling/memory_info.h

@@ -55,11 +55,19 @@ struct MemoryUsage {
   //    + purgeable_nonvolatile
   //    + purgeable_nonvolatile_compressed
   //    + page_table
+  //
+  // For Windows:
+  // This is the current memory size (in kilobytes) occupied by the OS process
+  // that is held in main memory (RAM). This is generally referred to as the
+  // working set size. This is an alias to
+  // PROCESS_MEMORY_COUNTERS::WorkingSetSize.
   int64_t mem_footprint_kb;
 
   // Total non-mmapped heap space allocated from system in bytes.
   // For Linux, this is an alias to mallinfo::arena.
   // For Mac, this is an alias to mstats::bytes_total
+  // For Windows, this is an alias to HEAP_SUMMARY::cbCommitted
+  // for the default process heap.
   //
   // This does not count mmapped heap space, nor does it count non-heap
   // uses of memory such as other mmapped space, thread stacks, globals,
@@ -70,6 +78,8 @@ struct MemoryUsage {
   // (i.e. excluding those have been freed).
   // For Linux, this is an alias to mallinfo::uordblks.
   // For Mac, this is an alias to mstats::bytes_used
+  // For Windows, this is an alias to HEAP_SUMMARY::cbAllocated
+  // for the default process heap.
   //
   // This does not count non-heap uses of mmap, nor does it count other
   // non-heap uses of memory such as thread stacks, globals, code, etc.

tflite/profiling/memory_info_test.cc

@@ -52,7 +52,7 @@ TEST(MemoryUsage, GetMemoryUsage) {
   EXPECT_EQ(MemoryUsage::kValueNotSet, result.total_allocated_bytes);
   EXPECT_EQ(MemoryUsage::kValueNotSet, result.in_use_allocated_bytes);
 
-#if defined(__linux__) || defined(__APPLE__)
+#if defined(__linux__) || defined(__APPLE__) || defined(_WIN32)
   // Just allocate some space in heap so that we have some meaningful
   // memory usage to report.
   constexpr int size = 10 * 1024 * 1024;
@@ -89,7 +89,7 @@ TEST(MemoryUsage, OutputMemoryUsageToStream) {
 }
 
 TEST(MemoryUsage, IsSupported) {
-#if defined(__linux__) || defined(__APPLE__)
+#if defined(__linux__) || defined(__APPLE__) || defined(_WIN32)
   EXPECT_TRUE(MemoryUsage::IsSupported());
 #else
   EXPECT_FALSE(MemoryUsage::IsSupported());