unify memory opimization passes

better reports
cleanups

src/PIM/Compiler/PimCodeGen.cpp

@@ -141,6 +141,19 @@ static void printMemoryOverflowDiagnostic(const MemoryValueKey& key,
   }
 }
 
+static std::fstream openMemoryReport(bool enabled) {
+  if (std::string outputRoot = getOutputDir(); !outputRoot.empty()) {
+    for (std::string name : {std::string("memory_") + "coalescing_report.txt",
+                             std::string("pim_memory_") + "liveness_report.txt",
+                             std::string("pim_memory_") + "liveness_report.json",
+                             std::string("pim_memory_") + "liveness_timeline.dot"})
+      sys::fs::remove(outputRoot + "/reports/" + name);
+    if (!enabled)
+      sys::fs::remove(outputRoot + "/reports/memory_report.txt");
+  }
+  return enabled ? openReportFile("memory_report") : std::fstream();
+}
+
 } // namespace
 
 size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) {
@@ -193,12 +206,11 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE
   globalMemEntriesMap[key] = memEntry;
 
   switch (reportKind) {
-  case MemoryReportKind::Alloca: break;
+  case MemoryReportKind::Alloca:
   case MemoryReportKind::Global:
-    ++reportRow.numGlobal;
-    reportRow.sizeGlobal += memEntry.size;
-    break;
   case MemoryReportKind::Input:
+    ++reportRow.hostObjectCount;
+    break;
   case MemoryReportKind::None:  break;
   }
 }
@@ -242,12 +254,15 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
 }
 
 void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane) {
-  if (localPhysicalSlots.empty()) {
-    llvm::append_range(localPhysicalSlots, plan.slots);
-    reportRow.logicalAllocaBytes = plan.logicalBytes;
-    reportRow.fallbackIntervals = plan.fallbackIntervals;
-    reportRow.nestedSingleUseIntervals = plan.nestedSingleUseIntervals;
+  if (!localArenaSize) {
+    localArenaSize = plan.arenaSize;
+    reportRow.logicalLocalAllocationCount = plan.logicalAllocationCount;
+    reportRow.logicalLocalBytes = plan.logicalBytes;
   }
+  else if (*localArenaSize != plan.arenaSize
+           || reportRow.logicalLocalAllocationCount != plan.logicalAllocationCount
+           || reportRow.logicalLocalBytes != plan.logicalBytes)
+    llvm_unreachable("inconsistent PIM local-memory plan across core-batch lanes");
   for (const CompiledLocalMemoryEntry& entry : plan.entries) {
     MemoryValueKey key = getMemoryValueKey(entry.value, lane);
     ownedMemEntriesMap[key] = entry.memory;
@@ -255,49 +270,10 @@ void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<u
   }
 }
 
-static void printHostMemoryReportRow(raw_ostream& os, const MemoryReportRow& row) {
-  llvm::SmallVector<ReportField, 2> fields = {
-    {"Number of globals", std::to_string(row.numGlobal)     },
-    {"Global memory",     formatReportMemory(row.sizeGlobal)}
-  };
-  printReportFlatFields(os, fields);
-}
-
-static void printCoreMemoryReportRow(raw_ostream& os, const MemoryReportEntry& entry) {
-  llvm::SmallVector<ReportField, 2> fields = {
-    {"Number of allocas", std::to_string(entry.row.numAlloca)     },
-    {"Allocated memory",  formatReportMemory(entry.row.sizeAlloca)}
-  };
-  printReportFlatFields(os, fields);
-}
-
-static void printBatchMemoryReportRow(raw_ostream& os, const MemoryReportEntry& entry) {
-  llvm::SmallVector<ReportField, 2> perCoreFields = {
-    {"Number of allocas", std::to_string(entry.row.numAlloca)     },
-    {"Allocated memory",  formatReportMemory(entry.row.sizeAlloca)}
-  };
-  llvm::SmallVector<ReportField, 2> totalFields = {
-    {"Number of allocas", std::to_string(entry.totalAllocaCount)    },
-    {"Batch memory",      formatReportMemory(entry.totalAllocaBytes)}
-  };
-  printReportPerCoreAndTotalFields(os, perCoreFields, totalFields);
-}
-
-static MemoryReportRow addMemoryReportRows(const MemoryReportRow& lhs, const MemoryReportRow& rhs) {
-  MemoryReportRow result = lhs;
-  result.numAlloca += rhs.numAlloca;
-  result.sizeAlloca += rhs.sizeAlloca;
-  result.numGlobal += rhs.numGlobal;
-  result.sizeGlobal += rhs.sizeGlobal;
-  return result;
-}
-
 MemoryReportRow PimMemory::getReportRow() const {
   MemoryReportRow row = reportRow;
-  row.numAlloca = localPhysicalSlots.size();
-  row.sizeAlloca = 0;
-  for (const PhysicalSlotInfo& slot : localPhysicalSlots)
-    row.sizeAlloca += slot.size;
+  row.hostBytes = firstAvailableAddress;
+  row.physicalLocalBytes = localArenaSize.value_or(0);
   return row;
 }
 
@@ -379,29 +355,32 @@ llvm::FailureOr<int64_t> PimAcceleratorMemory::getIndexValue(mlir::Value value,
   return compiledIt->second.evaluate(knowledge);
 }
 
+PimAcceleratorMemory::PimAcceleratorMemory()
+: hostMem(memEntriesMap), fileReport(openMemoryReport(pimMemoryReport == PimMemoryReportSummary)) {}
+
+PimAcceleratorMemory::PimAcceleratorMemory(
+  const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport)
+: memEntriesMap(initialMemEntries),
+  hostMem(memEntriesMap),
+  fileReport(enableReport ? openMemoryReport(true) : std::fstream()) {}
+
 void PimAcceleratorMemory::reportHost() { hostReportRow = hostMem.getReportRow(); }
 
 void PimAcceleratorMemory::recordCoreReport(size_t coreId, const MemoryReportRow& row) {
   reportEntries.push_back({MemoryReportEntry::Kind::Core,
                            coreId,
                            {pim::checkedI32OrCrash(coreId, "memory report core id")},
-                           row,
-                           row.numAlloca,
-                           row.sizeAlloca});
+                           row});
 }
 
 void PimAcceleratorMemory::recordBatchReport(uint64_t batchId,
                                              ArrayRef<int32_t> coreIds,
-                                             const MemoryReportRow& perCoreRow,
-                                             uint64_t totalAllocaCount,
-                                             uint64_t totalAllocaBytes) {
+                                             const MemoryReportRow& perCoreRow) {
   MemoryReportEntry entry;
   entry.kind = MemoryReportEntry::Kind::Batch;
   entry.id = batchId;
   llvm::append_range(entry.coreIds, coreIds);
   entry.row = perCoreRow;
-  entry.totalAllocaCount = totalAllocaCount;
-  entry.totalAllocaBytes = totalAllocaBytes;
   reportEntries.push_back(std::move(entry));
 }
 
@@ -410,93 +389,91 @@ void PimAcceleratorMemory::flushReport() {
     return;
 
   llvm::raw_os_ostream os(fileReport);
-  uint64_t totalGlobalMemory = hostReportRow.has_value() ? hostReportRow->sizeGlobal : 0;
-  uint64_t totalWeightsMemory = totalWeightBytes;
-  uint64_t totalCoresMemory = 0;
-  uint64_t totalLogicalCoreMemory = 0;
-  uint64_t totalFallbackIntervals = 0;
-  uint64_t totalNestedSingleUseIntervals = 0;
-  uint64_t maximumCorePeak = 0;
-  std::string worstEntry = "<none>";
-  for (const MemoryReportEntry& entry : reportEntries) {
-    totalCoresMemory += entry.totalAllocaBytes;
-    uint64_t factor = entry.kind == MemoryReportEntry::Kind::Batch ? entry.coreIds.size() : 1;
-    totalLogicalCoreMemory += entry.row.logicalAllocaBytes * factor;
-    totalFallbackIntervals += entry.row.fallbackIntervals * factor;
-    totalNestedSingleUseIntervals += entry.row.nestedSingleUseIntervals * factor;
-    if (entry.row.sizeAlloca <= maximumCorePeak)
-      continue;
-    maximumCorePeak = entry.row.sizeAlloca;
-    worstEntry = entry.kind == MemoryReportEntry::Kind::Batch ? "Batch " + std::to_string(entry.id)
-                                                              : "Core " + std::to_string(entry.coreIds.front());
-  }
-  uint64_t savedCoreMemory =
-    totalLogicalCoreMemory >= totalCoresMemory ? totalLogicalCoreMemory - totalCoresMemory : 0;
-  double savedPercent = totalLogicalCoreMemory == 0
-                        ? 0.0
-                        : 100.0 * static_cast<double>(savedCoreMemory) / static_cast<double>(totalLogicalCoreMemory);
-
-  llvm::SmallVector<ReportField, 10> totalFields = {
-    {"Global memory",               formatReportMemory(totalGlobalMemory)               },
-    {"Weights memory",              formatReportMemory(totalWeightsMemory)              },
-    {"Logical core allocations",    formatReportMemory(totalLogicalCoreMemory)           },
-    {"Physical core memory",        formatReportMemory(totalCoresMemory)                 },
-    {"Cores memory",                formatReportMemory(totalCoresMemory)                 },
-    {"Saved core memory",           formatReportMemory(savedCoreMemory)                  },
-    {"Saved core memory percent",   formatv("{0:F2}%", savedPercent).str()               },
-    {"Maximum core peak",           formatReportMemory(maximumCorePeak)                  },
-    {"Worst core/batch",            worstEntry                                           },
-    {"Fallback intervals",          std::to_string(totalFallbackIntervals)               },
-    {"Nested single-use intervals", std::to_string(totalNestedSingleUseIntervals)        }
+  struct ReportGroup {
+    MemoryReportRow row;
+    SmallVector<int32_t, 8> coreIds;
+    std::optional<uint64_t> batchId;
+  };
+  SmallVector<ReportGroup, 32> groups;
+  for (const MemoryReportEntry& entry : reportEntries) {
+    auto group = llvm::find_if(groups, [&](const ReportGroup& candidate) {
+      return candidate.row.logicalLocalBytes == entry.row.logicalLocalBytes
+          && candidate.row.physicalLocalBytes == entry.row.physicalLocalBytes;
+    });
+    if (group == groups.end()) {
+      ReportGroup newGroup;
+      newGroup.row = entry.row;
+      llvm::append_range(newGroup.coreIds, entry.coreIds);
+      if (entry.kind == MemoryReportEntry::Kind::Batch)
+        newGroup.batchId = entry.id;
+      groups.push_back(std::move(newGroup));
+      continue;
+    }
+    llvm::append_range(group->coreIds, entry.coreIds);
+    group->batchId.reset();
+  }
+  llvm::sort(groups, [](const ReportGroup& lhs, const ReportGroup& rhs) {
+    if (lhs.row.physicalLocalBytes != rhs.row.physicalLocalBytes)
+      return lhs.row.physicalLocalBytes > rhs.row.physicalLocalBytes;
+    return lhs.row.logicalLocalBytes > rhs.row.logicalLocalBytes;
+  });
+  uint64_t logicalBytes = 0;
+  uint64_t physicalBytes = 0;
+  for (const ReportGroup& group : groups) {
+    logicalBytes += group.row.logicalLocalBytes * group.coreIds.size();
+    physicalBytes += group.row.physicalLocalBytes * group.coreIds.size();
+  }
+  uint64_t savedBytes = logicalBytes - physicalBytes;
+  double savedPercent = logicalBytes == 0 ? 0.0 : 100.0 * savedBytes / logicalBytes;
+  uint64_t largest = groups.empty() ? 0 : groups.front().row.physicalLocalBytes;
+  auto printLabel = [&](const ReportGroup& group) {
+    if (group.batchId)
+      os << "Batch " << *group.batchId << " — cores ";
+    else if (group.coreIds.size() == 1)
+      os << "Core ";
+    else
+      os << "Cores ";
+    printCompressedIntegerEntries(os, ArrayRef<int32_t>(group.coreIds));
   };
-  printReportTotalsBlock(os, totalFields);
 
-  if (hostReportRow.has_value()) {
-    os << "\nHost:\n";
-    printHostMemoryReportRow(os, *hostReportRow);
+  os << "Summary\n";
+  os << "  Host memory: " << formatReportMemory(hostReportRow ? hostReportRow->hostBytes : 0) << "\n";
+  os << "  Weights memory: " << formatReportMemory(totalWeightBytes) << "\n";
+  os << "  Local memory before reuse: " << formatReportMemory(logicalBytes) << "\n";
+  os << "  Local memory after reuse: " << formatReportMemory(physicalBytes) << "\n";
+  os << "  Saved local memory: " << formatReportMemory(savedBytes) << " ("
+     << formatv("{0:F1}%", savedPercent) << ")\n";
+  os << "  Largest core local memory: " << formatReportMemory(largest) << "\n";
+  if (!groups.empty()) {
+    os << "    ";
+    printLabel(groups.front());
+    os << "\n";
   }
 
-  if (!reportEntries.empty()) {
-    if (hostReportRow.has_value())
-      os << "\n";
-    sortReportEntriesByFirstCore(reportEntries);
-
-    for (size_t index = 0; index < reportEntries.size();) {
-      size_t runEnd = index + 1;
-      while (runEnd < reportEntries.size() && reportEntries[runEnd].kind == reportEntries[index].kind
-             && reportEntries[runEnd].row == reportEntries[index].row
-             && reportEntries[runEnd].totalAllocaCount == reportEntries[index].totalAllocaCount
-             && reportEntries[runEnd].totalAllocaBytes == reportEntries[index].totalAllocaBytes) {
-        ++runEnd;
-      }
-
-      if (reportEntries[index].kind == MemoryReportEntry::Kind::Batch) {
-        os << "Batch ";
-        for (size_t batchIndex = index; batchIndex < runEnd; ++batchIndex) {
-          if (batchIndex != index)
-            os << ",\n      ";
-          os << reportEntries[batchIndex].id << " (cores ";
-          printCompressedIntegerEntries(os, ArrayRef<int32_t>(reportEntries[batchIndex].coreIds));
-          os << ")";
-        }
-      }
-      else {
-        llvm::SmallVector<int32_t, 8> coreIds;
-        for (size_t coreIndex = index; coreIndex < runEnd; ++coreIndex)
-          coreIds.push_back(reportEntries[coreIndex].coreIds.front());
-        os << "Core ";
-        printCompressedIntegerEntries(os, ArrayRef<int32_t>(coreIds));
-      }
-      os << ":\n";
-      if (reportEntries[index].kind == MemoryReportEntry::Kind::Batch)
-        printBatchMemoryReportRow(os, reportEntries[index]);
-      else
-        printCoreMemoryReportRow(os, reportEntries[index]);
-      printReportEntrySeparator(os, runEnd < reportEntries.size());
-
-      index = runEnd;
+  os << "\nLargest core groups\n";
+  constexpr size_t kGroupLimit = 8;
+  for (const ReportGroup& group : ArrayRef(groups).take_front(kGroupLimit)) {
+    printLabel(group);
+    os << "\n";
+    uint64_t groupSaved = group.row.logicalLocalBytes - group.row.physicalLocalBytes;
+    double groupPercent = group.row.logicalLocalBytes == 0
+                            ? 0.0
+                            : 100.0 * groupSaved / group.row.logicalLocalBytes;
+    if (group.coreIds.size() == 1) {
+      os << "  Local memory: " << formatReportMemory(group.row.logicalLocalBytes) << " → "
+         << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent)
+         << ")\n";
+    }
+    else {
+      os << "  Per core: " << formatReportMemory(group.row.logicalLocalBytes) << " → "
+         << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent)
+         << ")\n";
+      os << "  Total after reuse: " << formatReportMemory(group.row.physicalLocalBytes * group.coreIds.size())
+         << "\n";
     }
   }
+  if (groups.size() > kGroupLimit)
+    os << "  " << groups.size() - kGroupLimit << " smaller groups omitted\n";
 
   os.flush();
   fileReport.close();
@@ -834,25 +811,20 @@ static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
 
 static FailureOr<CompiledCoreMemoryPlan> compileCoreMemoryPlan(Operation* coreLikeOp) {
   CompiledCoreMemoryPlan plan;
-  auto fallbackAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryFallbackCountAttrName);
-  auto nestedSingleUseAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryNestedSingleUseCountAttrName);
-  if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) {
-    coreLikeOp->emitError("requires complete PIM local-memory planning summary attributes before codegen");
+  auto arenaAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemorySizeAttrName);
+  if (!arenaAttr || arenaAttr.getInt() < 0
+      || static_cast<uint64_t>(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) {
+    coreLikeOp->emitError("requires a valid pim.local_memory_size attribute before codegen");
     return failure();
   }
-  plan.fallbackIntervals = static_cast<uint64_t>(fallbackAttr.getInt());
-  plan.nestedSingleUseIntervals = static_cast<uint64_t>(nestedSingleUseAttr.getInt());
-  SmallDenseMap<size_t, size_t, 32> slotIndices;
+  plan.arenaSize = static_cast<size_t>(arenaAttr.getInt());
   bool hasFailure = false;
   coreLikeOp->walk([&](memref::AllocOp allocOp) {
     if (hasFailure)
       return;
     auto addressAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
-    auto slotAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotAttrName);
-    auto slotSizeAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotSizeAttrName);
-    if (!addressAttr || !slotAttr || !slotSizeAttr || addressAttr.getInt() < 0 || slotAttr.getInt() < 0
-        || slotSizeAttr.getInt() < 0) {
-      allocOp.emitOpError("requires a complete non-negative PIM local-memory plan before codegen");
+    if (!addressAttr || addressAttr.getInt() < 0) {
+      allocOp.emitOpError("requires a non-negative pim.local_memory_address attribute before codegen");
       hasFailure = true;
       return;
     }
@@ -864,25 +836,23 @@ static FailureOr<CompiledCoreMemoryPlan> compileCoreMemoryPlan(Operation* coreLi
       return;
     }
     size_t address = static_cast<size_t>(addressAttr.getInt());
-    size_t slotId = static_cast<size_t>(slotAttr.getInt());
-    size_t slotSize = static_cast<size_t>(slotSizeAttr.getInt());
-    plan.entries.push_back({allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}});
-    plan.logicalBytes += *checkedSize;
-
-    auto [slotIt, inserted] = slotIndices.try_emplace(slotId, plan.slots.size());
-    if (inserted) {
-      plan.slots.push_back({slotId, 0, slotSize});
+    if (address > plan.arenaSize || *checkedSize > plan.arenaSize - address) {
+      allocOp.emitOpError("has a planned local-memory range outside its core arena");
+      hasFailure = true;
       return;
     }
-    const PhysicalSlotInfo& existing = plan.slots[slotIt->second];
-    if (existing.size != slotSize) {
-      allocOp.emitOpError("has a PIM local-memory arena inconsistent with another allocation");
+    plan.entries.push_back({allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}});
+    auto logicalBytes = pim::checkedAdd(
+      static_cast<size_t>(plan.logicalBytes), static_cast<size_t>(*checkedSize), allocOp, "logical local bytes");
+    if (failed(logicalBytes)) {
       hasFailure = true;
+      return;
     }
+    plan.logicalBytes = *logicalBytes;
+    ++plan.logicalAllocationCount;
   });
   if (hasFailure)
     return failure();
-  llvm::sort(plan.slots, [](const PhysicalSlotInfo& lhs, const PhysicalSlotInfo& rhs) { return lhs.id < rhs.id; });
   return plan;
 }
 
@@ -1462,7 +1432,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
 
   for (const SmallVector<size_t>& group : batchJobIndices) {
     SmallVector<int32_t> reportedCoreIds;
-    MemoryReportRow batchRow;
     std::optional<MemoryReportRow> batchPerCoreRow;
 
     for (size_t jobIndex : group) {
@@ -1475,15 +1444,13 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
       reportedCoreIds.push_back(pim::checkedI32OrCrash(job.emittedCoreId, "batch report core id"));
       if (!batchPerCoreRow)
         batchPerCoreRow = result.reportRow;
-      batchRow = addMemoryReportRows(batchRow, result.reportRow);
+      else if (!(*batchPerCoreRow == result.reportRow))
+        llvm_unreachable("one PIM core batch produced inconsistent per-core memory reports");
     }
 
     uint64_t batchReportId = jobs[group.front()].batchReportId.value_or(0);
-    memory.recordBatchReport(batchReportId,
-                             reportedCoreIds,
-                             batchPerCoreRow.value_or(MemoryReportRow {}),
-                             batchRow.numAlloca,
-                             batchRow.sizeAlloca);
+    memory.recordBatchReport(
+      batchReportId, reportedCoreIds, batchPerCoreRow.value_or(MemoryReportRow {}));
   }
 
   maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1;