better memory report

capped vector allocations at u32::MAX in rust simulator

src/PIM/Compiler/PimCodeGen.cpp

@@ -73,6 +73,22 @@ static MemoryValueKey getMemoryValueKey(mlir::Value value, std::optional<unsigne
   return {value, getLaneForMemoryValue(value, lane)};
 }
 
+static bool isInsidePimCoreLikeOp(memref::AllocOp allocOp) {
+  return allocOp->getParentOfType<pim::PimCoreOp>() || allocOp->getParentOfType<pim::PimCoreBatchOp>();
+}
+
+static MemoryReportKind classifyMemoryReportKind(mlir::Value value) {
+  if (isa<mlir::BlockArgument>(value))
+    return MemoryReportKind::Input;
+  if (auto* op = value.getDefiningOp()) {
+    if (isa<memref::AllocOp>(op))
+      return MemoryReportKind::Alloca;
+    if (isa<memref::GetGlobalOp>(op))
+      return MemoryReportKind::Global;
+  }
+  return MemoryReportKind::None;
+}
+
 static int32_t getVectorByteSizeOrCrash(ShapedType type) {
   auto byteSize = pim::getCheckedShapedTypeSizeInBytes(type, UnknownLoc::get(type.getContext()), "vector byte size");
   if (failed(byteSize))
@@ -89,19 +105,23 @@ MemEntry* PimMemory::gatherMemEntry(mlir::Value value, std::optional<unsigned> l
     pim::getCheckedShapedTypeSizeInBytes(type, UnknownLoc::get(type.getContext()), "memory allocation byte size");
   if (failed(checkedAllocSize))
     llvm_unreachable("Failed to compute checked allocation byte size");
-  size_t allocSize = static_cast<size_t>(*checkedAllocSize);
-  MemEntry memEntry = {0, allocSize};
-  return &memEntries.emplace_back(memEntry, getMemoryValueKey(value, lane)).first;
+  PendingMemEntry pending;
+  pending.memEntry = {0, *checkedAllocSize};
+  pending.key = getMemoryValueKey(value, lane);
+  pending.reportKind = classifyMemoryReportKind(value);
+  return &memEntries.emplace_back(std::move(pending)).memEntry;
 }
 
 void PimMemory::allocateGatheredMemory() {
-  llvm::sort(memEntries, [](auto a, auto b) -> bool { return a.first.size > b.first.size; });
-  for (auto& [memEntry, key] : memEntries)
-    allocateMemoryForValue(key, memEntry);
+  llvm::sort(memEntries, [](const PendingMemEntry& lhs, const PendingMemEntry& rhs) {
+    return lhs.memEntry.size > rhs.memEntry.size;
+  });
+  for (PendingMemEntry& pending : memEntries)
+    allocateMemoryForValue(pending.key, pending.memEntry, pending.reportKind);
   memEntries.clear();
 }
 
-void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry) {
+void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind) {
   memEntry.address = firstAvailableAddress;
   firstAvailableAddress += memEntry.size;
   // Alignment
@@ -110,6 +130,19 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE
 
   ownedMemEntriesMap[key] = memEntry;
   globalMemEntriesMap[key] = memEntry;
+
+  switch (reportKind) {
+  case MemoryReportKind::Alloca:
+    ++reportRow.numAlloca;
+    reportRow.sizeAlloca += memEntry.size;
+    break;
+  case MemoryReportKind::Global:
+    ++reportRow.numGlobal;
+    reportRow.sizeGlobal += memEntry.size;
+    break;
+  case MemoryReportKind::Input:
+  case MemoryReportKind::None:  break;
+  }
 }
 
 void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
@@ -140,7 +173,7 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
   });
 
   funcOp.walk([&](memref::AllocOp allocOp) {
-    if (!allocOp->getParentOfType<pim::PimCoreOp>() && !allocOp->getParentOfType<pim::PimCoreBatchOp>())
+    if (!isInsidePimCoreLikeOp(allocOp))
       gatherMemEntry(allocOp.getResult());
   });
 
@@ -193,23 +226,7 @@ static MemoryReportRow addMemoryReportRows(const MemoryReportRow& lhs, const Mem
   return result;
 }
 
-MemoryReportRow PimMemory::getReportRow() const {
-  MemoryReportRow row;
-  for (auto& [key, memEntry] : ownedMemEntriesMap) {
-    if (auto op = key.value.getDefiningOp()) {
-      if (isa<memref::AllocOp>(op)) {
-        row.numAlloca++;
-        row.sizeAlloca += memEntry.size;
-      }
-
-      if (isa<memref::GetGlobalOp>(op)) {
-        row.numGlobal++;
-        row.sizeGlobal += memEntry.size;
-      }
-    }
-  }
-  return row;
-}
+MemoryReportRow PimMemory::getReportRow() const { return reportRow; }
 
 void PimMemory::remove(mlir::Value val) {
   for (auto it = ownedMemEntriesMap.begin(); it != ownedMemEntriesMap.end();)
@@ -334,13 +351,15 @@ void PimAcceleratorMemory::flushReport() {
 
   llvm::raw_os_ostream os(fileReport);
   uint64_t totalGlobalMemory = hostReportRow.has_value() ? hostReportRow->sizeGlobal : 0;
+  uint64_t totalWeightsMemory = totalWeightBytes;
   uint64_t totalCoresMemory = 0;
   for (const MemoryReportEntry& entry : reportEntries)
     totalCoresMemory += entry.totalAllocaBytes;
 
-  llvm::SmallVector<ReportField, 2> totalFields = {
-    {"Global memory", formatReportMemory(totalGlobalMemory)},
-    {"Cores memory",  formatReportMemory(totalCoresMemory) }
+  llvm::SmallVector<ReportField, 3> totalFields = {
+    {"Global memory",  formatReportMemory(totalGlobalMemory) },
+    {"Weights memory", formatReportMemory(totalWeightsMemory)},
+    {"Cores memory",   formatReportMemory(totalCoresMemory)  }
   };
   printReportTotalsBlock(os, totalFields);
 
@@ -1358,7 +1377,9 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     request.weights = jobResults[jobIndex].usedWeights;
     weightRequests.push_back(std::move(request));
   }
-  auto mapCoreWeightToFileName = createAndPopulateWeightFolder(weightRequests, outputDirPath);
+  auto weightEmission = createAndPopulateWeightFolder(weightRequests, outputDirPath);
+  memory.setTotalWeightBytes(weightEmission.totalWeightBytes);
+  auto& mapCoreWeightToFileName = weightEmission.mapCoreWeightToFileName;
 
   for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
     const CoreEmissionJob& job = jobs[jobIndex];