blazingly fasterer

src/PIM/Compiler/PimCodeGen.hpp

@@ -39,8 +39,6 @@ struct PhysicalSlotInfo {
   size_t size = 0;
 };
 
-using MemoryPlanArtifacts = std::string;
-
 struct MemoryValueKey {
   mlir::Value value;
   std::optional<unsigned> lane;
@@ -48,15 +46,33 @@ struct MemoryValueKey {
   bool operator==(const MemoryValueKey& other) const { return value == other.value && lane == other.lane; }
 };
 
+struct CompiledLocalMemoryEntry {
+  mlir::Value value;
+  MemEntry memory;
+};
+
+struct CompiledCoreMemoryPlan {
+  llvm::SmallVector<CompiledLocalMemoryEntry, 32> entries;
+  llvm::SmallVector<PhysicalSlotInfo, 32> slots;
+  uint64_t logicalBytes = 0;
+  uint64_t fallbackIntervals = 0;
+  uint64_t nestedSingleUseIntervals = 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;
 
   bool operator==(const MemoryReportRow& other) const {
     return numAlloca == other.numAlloca && sizeAlloca == other.sizeAlloca && numGlobal == other.numGlobal
-        && sizeGlobal == other.sizeGlobal;
+        && sizeGlobal == other.sizeGlobal && logicalAllocaBytes == other.logicalAllocaBytes
+        && fallbackIntervals == other.fallbackIntervals
+        && nestedSingleUseIntervals == other.nestedSingleUseIntervals;
   }
 };
 
@@ -93,26 +109,22 @@ class PimMemory {
   llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap;
   llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32> ownedMemEntriesMap;
   MemoryReportRow reportRow;
-  MemoryPlanArtifacts livenessArtifacts;
 
   size_t minAlignment = 4;
   size_t firstAvailableAddress = 0;
-  size_t nextPhysicalSlotId = 0;
 
   MemEntry* gatherMemEntry(mlir::Value value, std::optional<unsigned> lane = std::nullopt);
   size_t allocateAddress(size_t size, const MemoryValueKey& key);
   void allocateGatheredMemory();
   void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind);
-  PhysicalSlotInfo allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key);
 
 public:
   PimMemory(llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap)
   : globalMemEntriesMap(globalMemEntriesMap) {}
 
   void allocateHost(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp);
-  void allocateCore(mlir::Operation* op, std::optional<unsigned> lane = std::nullopt);
+  void allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane = std::nullopt);
   MemoryReportRow getReportRow() const;
-  const MemoryPlanArtifacts& getLivenessArtifacts() const { return livenessArtifacts; }
 
   size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
   MemEntry getMemEntry(const MemoryValueKey& key) const;
@@ -160,6 +172,7 @@ public:
 struct CoreEmissionJob {
   mlir::Operation* coreLikeOp = nullptr;
   const CompiledCoreProgram* program = nullptr;
+  const CompiledCoreMemoryPlan* memoryPlan = nullptr;
   size_t emittedCoreId = 0;
   llvm::SmallVector<unsigned, 4> lanes;
   std::optional<uint64_t> batchReportId;