unify memory opimization passes

better reports
cleanups

src/PIM/Compiler/PimCodeGen.hpp

@@ -33,12 +33,6 @@ struct MemEntry {
   size_t size;
 };
 
-struct PhysicalSlotInfo {
-  size_t id = 0;
-  size_t address = 0;
-  size_t size = 0;
-};
-
 struct MemoryValueKey {
   mlir::Value value;
   std::optional<unsigned> lane;
@@ -53,26 +47,22 @@ struct CompiledLocalMemoryEntry {
 
 struct CompiledCoreMemoryPlan {
   llvm::SmallVector<CompiledLocalMemoryEntry, 32> entries;
-  llvm::SmallVector<PhysicalSlotInfo, 32> slots;
+  uint64_t logicalAllocationCount = 0;
   uint64_t logicalBytes = 0;
-  uint64_t fallbackIntervals = 0;
-  uint64_t nestedSingleUseIntervals = 0;
+  size_t arenaSize = 0;
 };
 
 struct MemoryReportRow {
-  uint64_t numAlloca = 0;
-  uint64_t sizeAlloca = 0;
-  uint64_t numGlobal = 0;
-  uint64_t sizeGlobal = 0;
-  uint64_t logicalAllocaBytes = 0;
-  uint64_t fallbackIntervals = 0;
-  uint64_t nestedSingleUseIntervals = 0;
+  uint64_t hostObjectCount = 0;
+  uint64_t hostBytes = 0;
+  uint64_t logicalLocalAllocationCount = 0;
+  uint64_t logicalLocalBytes = 0;
+  uint64_t physicalLocalBytes = 0;
 
   bool operator==(const MemoryReportRow& other) const {
-    return numAlloca == other.numAlloca && sizeAlloca == other.sizeAlloca && numGlobal == other.numGlobal
-        && sizeGlobal == other.sizeGlobal && logicalAllocaBytes == other.logicalAllocaBytes
-        && fallbackIntervals == other.fallbackIntervals
-        && nestedSingleUseIntervals == other.nestedSingleUseIntervals;
+    return hostObjectCount == other.hostObjectCount && hostBytes == other.hostBytes
+        && logicalLocalAllocationCount == other.logicalLocalAllocationCount
+        && logicalLocalBytes == other.logicalLocalBytes && physicalLocalBytes == other.physicalLocalBytes;
   }
 };
 
@@ -99,16 +89,14 @@ struct MemoryReportEntry {
   uint64_t id = 0;
   llvm::SmallVector<int32_t, 8> coreIds;
   MemoryReportRow row;
-  uint64_t totalAllocaCount = 0;
-  uint64_t totalAllocaBytes = 0;
 };
 
 class PimMemory {
   llvm::SmallVector<PendingMemEntry, 32> memEntries;
-  llvm::SmallVector<PhysicalSlotInfo, 32> localPhysicalSlots;
   llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap;
   llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32> ownedMemEntriesMap;
   MemoryReportRow reportRow;
+  std::optional<size_t> localArenaSize;
 
   size_t minAlignment = 4;
   size_t firstAvailableAddress = 0;
@@ -145,12 +133,9 @@ private:
   mutable llvm::DenseMap<mlir::Value, CompiledAddressExpr> compiledAddressExprs;
 
 public:
-  PimAcceleratorMemory()
-  : hostMem(memEntriesMap), fileReport(openReportFile("memory_report")) {}
-  PimAcceleratorMemory(const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport)
-  : memEntriesMap(initialMemEntries),
-    hostMem(memEntriesMap),
-    fileReport(enableReport ? openReportFile("memory_report") : std::fstream()) {}
+  PimAcceleratorMemory();
+  PimAcceleratorMemory(
+    const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport);
 
   PimMemory& getOrCreateDeviceMem(size_t id);
 
@@ -160,11 +145,8 @@ public:
   llvm::FailureOr<int64_t> getIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge = {}) const;
   void reportHost();
   void recordCoreReport(size_t coreId, const MemoryReportRow& row);
-  void recordBatchReport(uint64_t batchId,
-                         llvm::ArrayRef<int32_t> coreIds,
-                         const MemoryReportRow& perCoreRow,
-                         uint64_t totalAllocaCount,
-                         uint64_t totalAllocaBytes);
+  void recordBatchReport(
+    uint64_t batchId, llvm::ArrayRef<int32_t> coreIds, const MemoryReportRow& perCoreRow);
   void setTotalWeightBytes(uint64_t bytes) { totalWeightBytes = bytes; }
   void flushReport();
 };