better memory report

capped vector allocations at u32::MAX in rust simulator

src/PIM/Compiler/PimCodeGen.hpp

@@ -45,6 +45,19 @@ struct MemoryReportRow {
   }
 };
 
+enum class MemoryReportKind {
+  None,
+  Alloca,
+  Global,
+  Input
+};
+
+struct PendingMemEntry {
+  MemEntry memEntry;
+  MemoryValueKey key;
+  MemoryReportKind reportKind = MemoryReportKind::None;
+};
+
 struct MemoryReportEntry {
   enum class Kind {
     Core,
@@ -60,16 +73,17 @@ struct MemoryReportEntry {
 };
 
 class PimMemory {
-  llvm::SmallVector<std::pair<MemEntry, MemoryValueKey>, 32> memEntries;
+  llvm::SmallVector<PendingMemEntry, 32> memEntries;
   llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap;
   llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32> ownedMemEntriesMap;
+  MemoryReportRow reportRow;
 
   size_t minAlignment = 4;
   size_t firstAvailableAddress = 0;
 
   MemEntry* gatherMemEntry(mlir::Value value, std::optional<unsigned> lane = std::nullopt);
   void allocateGatheredMemory();
-  void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry);
+  void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind);
 
 public:
   PimMemory(llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap)
@@ -94,6 +108,7 @@ private:
   std::fstream fileReport;
   std::optional<MemoryReportRow> hostReportRow;
   llvm::SmallVector<MemoryReportEntry, 32> reportEntries;
+  uint64_t totalWeightBytes = 0;
   mutable llvm::DenseMap<mlir::Value, CompiledIndexExpr> compiledIndexExprs;
   mutable llvm::DenseMap<mlir::Value, CompiledAddressExpr> compiledAddressExprs;
 
@@ -118,6 +133,7 @@ public:
                          const MemoryReportRow& perCoreRow,
                          uint64_t totalAllocaCount,
                          uint64_t totalAllocaBytes);
+  void setTotalWeightBytes(uint64_t bytes) { totalWeightBytes = bytes; }
   void flushReport();
   void clean(mlir::Operation* op);
 };