Provide both number from malloc and from vm to compare. Should most likely use max of these go best estimate model cost.

1 files changed, 38 insertions, 8 deletions

diff --git a/eval/src/apps/analyze_onnx_model/analyze_onnx_model.cpp b/eval/src/apps/analyze_onnx_model/analyze_onnx_model.cpp
index 1fdc6fe79f0..6e09fa2de0f 100644
--- a/eval/src/apps/analyze_onnx_model/analyze_onnx_model.cpp
+++ b/eval/src/apps/analyze_onnx_model/analyze_onnx_model.cpp
@@ -56,8 +56,10 @@ void extract(const vespalib::string &str, const vespalib::string &prefix, vespal
     }
 }
 struct MemoryUsage {
-    size_t size;
-    size_t rss;
+    size_t vm_size;
+    size_t rss_size;
+    size_t malloc_peak;
+    size_t malloc_current;
 };
 
 static const vespalib::string UNKNOWN = "unknown";
@@ -77,17 +79,43 @@ size_t convert(const vespalib::string & s) {
 }
 
 MemoryUsage extract_memory_usage() {
+    vespalib::string vm_size = UNKNOWN;
+    vespalib::string vm_rss = UNKNOWN;
+    FilePointer file(fopen("/proc/self/status", "r"));
+    if (file.valid()) {
+        vespalib::string line;
+        while (read_line(file, line)) {
+            extract(line, "VmSize:", vm_size);
+            extract(line, "VmRSS:", vm_rss);
+        }
+    }
+    MemoryUsage usage = {};
+    usage.vm_size = convert(vm_size);
+    usage.rss_size = convert(vm_rss);
+
+#if __GLIBC_PREREQ(2, 33)
+    struct mallinfo2 mallocInfo = mallinfo2();
+    usage.malloc_peak = size_t(info.usmblks);
+    usage.malloc_current = size_t(info.arena + info.hblkhd);
+#else
     struct mallinfo info = mallinfo();
+
     if (dlsym(RTLD_NEXT, "is_vespamalloc") != nullptr) {
-        return {size_t(info.usmblks) << 20, size_t(info.arena + info.hblkhd) << 20};
+        // Vespamalloc reports arena in 1M blocks as an 'int' is too small.
+        usage.malloc_peak = size_t(info.usmblks) * 1_Mi;
+        usage.malloc_current = size_t(info.arena + info.hblkhd) * 1_Mi;
     } else {
-        return {size_t(info.usmblks), size_t(info.arena + info.hblkhd)};
+        usage.malloc_peak = size_t(info.usmblks);
+        usage.malloc_current = size_t(info.arena + info.hblkhd);
     }
+#endif
+    return usage;
 }
 
 void report_memory_usage(const vespalib::string &desc) {
-    MemoryUsage vm = extract_memory_usage();
-    fprintf(stderr, "vm_size: %zu kB, vm_rss: %zu kB (%s)\n", vm.size/1024, vm.rss/1024, desc.c_str());
+    MemoryUsage m = extract_memory_usage();
+    fprintf(stderr, "vm_size: %zu kB, vm_rss: %zu kB, malloc_peak: %zu kb, malloc_curr: %zu (%s)\n",
+            m.vm_size/1_Ki, m.rss_size/1_Ki, m.malloc_peak/1_Ki, m.malloc_current/1_Ki, desc.c_str());
 }
 
 struct Options {
@@ -283,8 +311,10 @@ int probe_types() {
         types.setString(output.name, output_type.to_spec());
     }
     MemoryUsage vm_after = extract_memory_usage();
-    root.setLong("vm_size", vm_after.size - vm_before.size);
-    root.setLong("vm_rss", vm_after.rss - vm_before.rss);
+    root.setLong("vm_size", vm_after.vm_size - vm_before.vm_size);
+    root.setLong("vm_rss", vm_after.rss_size - vm_before.rss_size);
+    root.setLong("malloc_peak", vm_after.malloc_peak - vm_before.malloc_peak);
+    root.setLong("malloc_current", vm_after.malloc_current - vm_before.malloc_current);
     write_compact(result, std_out);
     return 0;
 }