From 481bbaf1fce5b7ef0162c6f9ecec73a80e263cc7 Mon Sep 17 00:00:00 2001
From: Peter Boyle <paboyle@ph.ed.ac.uk>
Date: Thu, 23 Mar 2023 12:55:31 -0400
Subject: [PATCH] Interface to query memory use

---
 Grid/allocator/MemoryManager.cc |  1 +
 Grid/allocator/MemoryManager.h  | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/Grid/allocator/MemoryManager.cc b/Grid/allocator/MemoryManager.cc
index e9097c75..a9e5c9b4 100644
--- a/Grid/allocator/MemoryManager.cc
+++ b/Grid/allocator/MemoryManager.cc
@@ -39,6 +39,7 @@ void MemoryManager::PrintBytes(void)
 }
 
 uint64_t MemoryManager::DeviceCacheBytes() { return CacheBytes[Acc] + CacheBytes[AccHuge] + CacheBytes[AccSmall]; }
+uint64_t MemoryManager::HostCacheBytes()   { return CacheBytes[Cpu] + CacheBytes[CpuHuge] + CacheBytes[CpuSmall]; }
 
 //////////////////////////////////////////////////////////////////////
 // Data tables for recently freed pooiniter caches
diff --git a/Grid/allocator/MemoryManager.h b/Grid/allocator/MemoryManager.h
index 7a5f978c..0dc78f04 100644
--- a/Grid/allocator/MemoryManager.h
+++ b/Grid/allocator/MemoryManager.h
@@ -71,6 +71,21 @@ enum ViewMode {
   CpuWriteDiscard = 0x10 // same for now
 };
 
+struct MemoryStatus {
+  uint64_t     DeviceBytes;
+  uint64_t     DeviceLRUBytes;
+  uint64_t     DeviceMaxBytes;
+  uint64_t     HostToDeviceBytes;
+  uint64_t     DeviceToHostBytes;
+  uint64_t     HostToDeviceXfer;
+  uint64_t     DeviceToHostXfer;
+  uint64_t     DeviceEvictions;
+  uint64_t     DeviceDestroy;
+  uint64_t     DeviceAllocCacheBytes;
+  uint64_t     HostAllocCacheBytes;
+};
+
+
 class MemoryManager {
 private:
 
@@ -124,7 +139,24 @@ private:
   static uint64_t     DeviceDestroy;
   
   static uint64_t     DeviceCacheBytes();
+  static uint64_t     HostCacheBytes();
 
+  static MemoryStatus GetFootprint(void) {
+    MemoryStatus stat;
+    stat.DeviceBytes       = DeviceBytes;
+    stat.DeviceLRUBytes    = DeviceLRUBytes;
+    stat.DeviceMaxBytes    = DeviceMaxBytes;
+    stat.HostToDeviceBytes = HostToDeviceBytes;
+    stat.DeviceToHostBytes = DeviceToHostBytes;
+    stat.HostToDeviceXfer  = HostToDeviceXfer;
+    stat.DeviceToHostXfer  = DeviceToHostXfer;
+    stat.DeviceEvictions   = DeviceEvictions;
+    stat.DeviceDestroy     = DeviceDestroy;
+    stat.DeviceAllocCacheBytes = DeviceCacheBytes();
+    stat.HostAllocCacheBytes   = HostCacheBytes();
+    return stat;
+  };
+  
  private:
 #ifndef GRID_UVM
   //////////////////////////////////////////////////////////////////////