From 26317ea7d0fe0685df7a3e1cf0a3cd78cc776f4d Mon Sep 17 00:00:00 2001
From: ilgeco <gabriele.magnani@mail.polimi.it>
Date: Tue, 12 May 2026 10:38:35 +0200
Subject: [PATCH] Shorter Memory Reporty

---
 src/PIM/Compiler/PimCodeGen.cpp | 56 +++++++++++----------------------
 1 file changed, 18 insertions(+), 38 deletions(-)

diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp
index 288798c..792ba27 100644
--- a/src/PIM/Compiler/PimCodeGen.cpp
+++ b/src/PIM/Compiler/PimCodeGen.cpp
@@ -199,47 +199,27 @@ std::string formatMemory(uint64_t bytes) {
 }
 
 void PimMemory::report(llvm::raw_ostream& file) {
-  // Key: {OpName, Size}, Value: Vector of Addresses
-  // This groups all "memref.alloc" of "1KB" together
-  std::vector orderedList(globalMemEntriesMap.begin(), globalMemEntriesMap.end());
-  std::sort(
-    orderedList.begin(), orderedList.end(), [](auto lft, auto rgt) { return lft.second.address < rgt.second.address; });
-  auto newEnd = std::unique(orderedList.begin(), orderedList.end(), [](auto lft, auto rgt) {
-    return lft.second.address == rgt.second.address;
-  });
-  orderedList.erase(newEnd, orderedList.end());
-  std::sort(
-    orderedList.begin(), orderedList.end(), [](auto lft, auto rgt) { return lft.second.size < rgt.second.size; });
-  std::map<std::pair<std::string, uint64_t>, std::vector<uint64_t>> groupedStats;
+  uint64_t numAlloca = 0;
+  uint64_t sizeAlloca = 0;
+  uint64_t numGlobal = 0;
+  uint64_t sizeGlobal = 0;
+  for (auto& [val, memEntry] : globalMemEntriesMap) {
+    if (auto op = val.getDefiningOp()) {
+      if (auto allocaOp = dyn_cast<memref::AllocOp>(op)) {
+        numAlloca++;
+        sizeAlloca += memEntry.size;
+      }
 
-  for (auto& [value, memEntry] : orderedList) {
-    std::string opName = "Unknown/BlockArg";
-    if (auto op = value.getDefiningOp())
-      opName = op->getName().getStringRef().str();
-
-    groupedStats[{opName, memEntry.size}].push_back(memEntry.address);
-  }
-
-  file << "--- Memory Usage Report ---\n";
-
-  uint64_t totalMemory = 0;
-  for (auto const& [key, addresses] : groupedStats) {
-    const std::string& opName = key.first;
-    uint64_t size = key.second;
-
-    file.indent(4) << "Type: " << opName << " [" << formatMemory(size) << "]\n";
-    file.indent(6) << "Count:   " << addresses.size() << "\n";
-    file.indent(6) << "Total Memory:   " << formatMemory(size * addresses.size()) << "\n";
-    totalMemory += size * addresses.size();
-
-    // Optional: Print address range or first/last address to keep it concise
-    if (!addresses.empty()) {
-      auto [min, max] = std::minmax_element(addresses.begin(), addresses.end());
-      file.indent(6) << "Range:   " << llvm::format_hex(*min, 10) << " -> " << llvm::format_hex(*max, 10) << "\n";
+      if (auto allocaOp = dyn_cast<memref::GetGlobalOp>(op)) {
+        numGlobal++;
+        sizeGlobal += memEntry.size;
+      }
     }
     file << "\n";
     file << "Total Core Memory:  " << formatMemory(totalMemory) << "\n";
   }
+
+  file << numAlloca << " " << formatMemory(sizeAlloca) << " " << numGlobal << " " << formatMemory(sizeGlobal) << "\n";
 }
 
 void PimMemory::remove(mlir::Value val) {
@@ -289,14 +269,14 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value, const StaticValu
 
 void PimAcceleratorMemory::reportHost() {
   llvm::raw_os_ostream os(fileReport);
-  os << "Host Memory\n";
+  os << "Host Memory\t";
   hostMem.report(os);
   os.flush();
 }
 
 void PimAcceleratorMemory::reportCore(size_t coreId) {
   llvm::raw_os_ostream os(fileReport);
-  os << "Core " << coreId << " Memory\n";
+  os << "Core " << coreId << " Memory\t";
   deviceMem.at(coreId).report(os);
   os.flush();
 }