unify memory opimization passes

better reports
cleanups

src/PIM/CMakeLists.txt

@@ -121,7 +121,6 @@ add_pim_library(OMPIMAccel
   OMPimCommon
   OMPimBufferization
   OMPimHostConstantFolding
-  OMPimMemoryCoalescing
   OMPimVerification
   MLIRTensorInferTypeOpInterfaceImpl
 )

src/PIM/Common/PimCommon.hpp

@@ -23,6 +23,7 @@
 #include "src/Compiler/CompilerOptions.hpp"
 
 #include <cstdint>
+#include <array>
 #include <limits>
 
 namespace onnx_mlir {
@@ -30,11 +31,13 @@ namespace onnx_mlir {
 inline constexpr llvm::StringLiteral kCoreIdAttrName = "coreId";
 inline constexpr llvm::StringLiteral kCoreIdsAttrName = "coreIds";
 inline constexpr llvm::StringLiteral kLocalMemoryAddressAttrName = "pim.local_memory_address";
-inline constexpr llvm::StringLiteral kLocalMemorySlotAttrName = "pim.local_memory_slot";
-inline constexpr llvm::StringLiteral kLocalMemorySlotSizeAttrName = "pim.local_memory_slot_size";
-inline constexpr llvm::StringLiteral kLocalMemoryFallbackCountAttrName = "pim.local_memory_fallback_count";
-inline constexpr llvm::StringLiteral kLocalMemoryNestedSingleUseCountAttrName =
-  "pim.local_memory_nested_single_use_count";
+inline constexpr llvm::StringLiteral kLocalMemorySizeAttrName = "pim.local_memory_size";
+inline constexpr std::array<llvm::StringLiteral, 4> kRemovedLocalMemoryPlanAttrNames = {
+  "pim.local_memory_slot",
+  "pim.local_memory_slot_size",
+  "pim.local_memory_fallback_count",
+  "pim.local_memory_nested_single_use_count",
+};
 inline constexpr size_t kPimLocalMemoryAddressLimit = static_cast<size_t>(std::numeric_limits<int32_t>::max());
 
 } // namespace onnx_mlir

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;

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();
 };

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -19,10 +19,8 @@ llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport(
   "pim-memory-report",
   llvm::cl::desc("Emit a human-readable PIM memory planning report"),
   llvm::cl::values(clEnumValN(PimMemoryReportNone, "none", "Do not emit any PIM memory planning report")),
-  llvm::cl::values(
-    clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise slot reuse report with key offenders")),
-  llvm::cl::values(clEnumValN(PimMemoryReportFull, "full", "Emit the full detailed PIM memory planning report")),
-  llvm::cl::init(PimMemoryReportNone),
+  llvm::cl::values(clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise PIM memory summary")),
+  llvm::cl::init(PimMemoryReportSummary),
   llvm::cl::cat(OnnxMlirOptions));
 
 llvm::cl::opt<PimConvLoweringType> pimConvLowering(
@@ -72,12 +70,6 @@ llvm::cl::opt<bool>
                  llvm::cl::init(false),
                  llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool>
-  pimDisableMemoryCoalescing("pim-disable-memory-coalescing",
-                             llvm::cl::desc("Skip the early PIM IR memory coalescing pass (developer diagnostic)"),
-                             llvm::cl::init(false),
-                             llvm::cl::cat(OnnxMlirOptions));
-
 llvm::cl::opt<bool> useExperimentalConvImpl("use-experimental-conv-impl",
                                             llvm::cl::desc("Use experimental implementation for convolution"),
                                             llvm::cl::init(false),

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -23,7 +23,6 @@ typedef enum {
 typedef enum {
   PimMemoryReportNone = 0,
   PimMemoryReportSummary = 1,
-  PimMemoryReportFull = 2,
 } PimMemoryReportLevel;
 
 typedef enum {
@@ -53,7 +52,6 @@ extern llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport;
 extern llvm::cl::opt<PimConvLoweringType> pimConvLowering;
 extern llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow;
 
-extern llvm::cl::opt<bool> pimDisableMemoryCoalescing;
 extern llvm::cl::opt<bool> pimOnlyCodegen;
 extern llvm::cl::opt<bool> useExperimentalConvImpl;
 extern llvm::cl::opt<bool> pimEmitJson;

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -22,6 +22,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
 
   if (pimOnlyCodegen) {
     pm.addPass(createPimLocalMemoryPlanningPass());
+    pm.addPass(createPimVerificationPass());
     pm.addPass(createEmitPimCodePass());
     return;
   }
@@ -52,8 +53,6 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     pm.addPass(mlir::createLowerAffinePass());
     pm.addPass(createPimHostConstantFoldingPass());
     pm.addPass(createMessagePass("Pim host constants folded"));
-    if (!pimDisableMemoryCoalescing)
-      pm.addPass(createPimMemoryCoalescingPass());
     pm.addPass(createPimLocalMemoryPlanningPass());
     pm.addPass(createMessagePass("Pim local memory planned"));
     pm.addPass(createPimVerificationPass());

src/PIM/Dialect/Pim/Analysis/CMakeLists.txt

@@ -5,4 +5,8 @@ add_pim_library(OMPimLocalMemoryLifetimeAnalysis
 
   INCLUDE_DIRS PUBLIC
   ${PIM_PUBLIC_INCLUDE_DIRS}
+
+  LINK_LIBS PUBLIC
+  OMPimCommon
+  PimOps
 )

src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp

@@ -2,70 +2,204 @@
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace mlir;
 
-namespace onnx_mlir {
-namespace pim {
+namespace onnx_mlir::pim {
+namespace {
+
+struct OperationOrdering {
+  DenseMap<Operation*, uint64_t> position;
+  DenseMap<Operation*, uint64_t> subtreeEnd;
+  uint64_t next = 0;
+};
+
+static void orderOperation(Operation* op, OperationOrdering& ordering) {
+  uint64_t end = ordering.next;
+  ordering.position[op] = ordering.next++;
+  for (Region& region : op->getRegions())
+    for (Block& block : region)
+      for (Operation& nested : block) {
+        orderOperation(&nested, ordering);
+        end = std::max(end, ordering.subtreeEnd.lookup(&nested));
+      }
+  ordering.subtreeEnd[op] = end;
+}
+
+static OperationOrdering buildOrdering(Operation* coreLikeOp) {
+  OperationOrdering ordering;
+  if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
+    return ordering;
+  for (Operation& op : coreLikeOp->getRegion(0).front())
+    orderOperation(&op, ordering);
+  return ordering;
+}
+
+static bool isRuntimeMemoryTouch(Operation* op) {
+  return isa<PimMemCopyHostToDevOp,
+             PimMemCopyDevToHostOp,
+             PimMemCopyOp,
+             PimReceiveOp,
+             PimSendOp,
+             PimConcatOp,
+             PimVMMOp,
+             PimTransposeOp,
+             PimVVAddOp,
+             PimVVSubOp,
+             PimVVMulOp,
+             PimVVMaxOp,
+             PimVVDMulOp,
+             PimVAvgOp,
+             PimVReluOp,
+             PimVTanhOp,
+             PimVSigmOp,
+             PimVSoftmaxOp>(op);
+}
+
+static bool isIgnoredUser(Operation* op) {
+  return isLocalMemoryAliasOp(op) || isa<scf::ForOp, scf::YieldOp, memref::DeallocOp>(op)
+      || isCoreStaticAddressOp(op);
+}
+
+static bool isWithin(Value value, Region* region) {
+  if (auto argument = dyn_cast<BlockArgument>(value))
+    return argument.getOwner()->getParent() == region;
+  if (Operation* definingOp = value.getDefiningOp())
+    return definingOp->getParentRegion() == region || region->isAncestor(definingOp->getParentRegion());
+  return false;
+}
+
+static std::pair<uint64_t, uint64_t> getTouchRange(Value definingValue, Operation* user, const OperationOrdering& ordering) {
+  uint64_t start = ordering.position.lookup(user);
+  uint64_t end = start;
+  for (Operation* current = user; current; current = current->getParentOp())
+    if (auto forOp = dyn_cast<scf::ForOp>(current); forOp && !isWithin(definingValue, &forOp.getRegion())) {
+      start = std::min(start, ordering.position.lookup(forOp));
+      end = std::max(end, ordering.subtreeEnd.lookup(forOp));
+    }
+  return {start, end};
+}
+
+static FailureOr<size_t> getAllocationSize(memref::AllocOp allocation) {
+  auto type = dyn_cast<ShapedType>(allocation.getType());
+  if (!type)
+    return failure();
+  auto bytes = getCheckedShapedTypeSizeInBytes(type, allocation, "memory allocation byte size");
+  if (failed(bytes))
+    return failure();
+  return checkedSize(*bytes, allocation, "memory allocation byte size");
+}
+
+} // namespace
 
 bool isLocalMemoryAliasOp(Operation* op) {
   return isa<memref::SubViewOp, memref::CastOp, memref::CollapseShapeOp, memref::ExpandShapeOp>(op);
 }
 
-LogicalResult walkLocalMemoryUses(Value root,
-                                  llvm::function_ref<LogicalResult(Value, Operation*)> visitUser) {
+LogicalResult walkLocalMemoryUses(Value root, LocalMemoryUseCallback callback) {
   llvm::SmallPtrSet<Value, 16> visitedValues;
-  llvm::SmallPtrSet<Operation*, 32> visitedUsers;
-  llvm::SmallVector<Value> pendingValues {root};
-  auto addAlias = [&](Value value) { pendingValues.push_back(value); };
-
-  while (!pendingValues.empty()) {
-    Value value = pendingValues.pop_back_val();
+  llvm::SmallPtrSet<OpOperand*, 32> visitedUses;
+  SmallVector<Value> pending = {root};
+  while (!pending.empty()) {
+    Value value = pending.pop_back_val();
     if (!visitedValues.insert(value).second)
       continue;
-    for (Operation* user : value.getUsers()) {
-      if (!visitedUsers.insert(user).second)
-        continue;
-      if (failed(visitUser(value, user)))
+    for (OpOperand& use : value.getUses()) {
+      if (!visitedUses.insert(&use).second || failed(callback(use)))
         return failure();
+      Operation* user = use.getOwner();
 
-      if (isLocalMemoryAliasOp(user))
-        for (Value result : user->getResults())
-          addAlias(result);
+      if (isLocalMemoryAliasOp(user) && use.getOperandNumber() == 0) {
+        if (user->getNumResults() != 1)
+          return failure();
+        pending.push_back(user->getResult(0));
+      }
 
-      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user))
-        for (OpResult result : user->getResults())
-          if (OpOperand* tied = dpsOp.getTiedOpOperand(result); tied && tied->get() == value)
-            addAlias(result);
+      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user); dpsOp && dpsOp.isDpsInit(&use))
+        pending.push_back(dpsOp.getTiedOpResult(&use));
 
       if (auto forOp = dyn_cast<scf::ForOp>(user))
-        for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs()))
-          if (initArg == value) {
-            addAlias(forOp.getRegionIterArgs()[index]);
-            addAlias(forOp.getResult(index));
+        for (auto [index, init] : llvm::enumerate(forOp.getInitArgsMutable()))
+          if (&init == &use) {
+            pending.push_back(forOp.getRegionIterArgs()[index]);
+            pending.push_back(forOp.getResult(index));
           }
 
-      auto yieldOp = dyn_cast<scf::YieldOp>(user);
-      if (!yieldOp)
-        continue;
-      for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
-        if (operand != value)
-          continue;
-        if (auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp()))
-          addAlias(forOp.getResult(index));
-        else if (auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp()))
-          addAlias(ifOp.getResult(index));
-        else if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp()))
-          addAlias(switchOp.getResult(index));
+      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
+        unsigned index = use.getOperandNumber();
+        Operation* parent = yieldOp->getParentOp();
+        if (auto forOp = dyn_cast<scf::ForOp>(parent))
+          pending.push_back(forOp.getResult(index));
+        else if (auto ifOp = dyn_cast<scf::IfOp>(parent))
+          pending.push_back(ifOp.getResult(index));
+        else if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(parent))
+          pending.push_back(switchOp.getResult(index));
+        else
+          return failure();
       }
     }
   }
   return success();
 }
 
-} // namespace pim
-} // namespace onnx_mlir
+FailureOr<SmallVector<LocalMemoryInterval, 0>> analyzeLocalMemoryLifetimes(Operation* coreLikeOp) {
+  SmallVector<LocalMemoryInterval, 0> intervals;
+  OperationOrdering ordering = buildOrdering(coreLikeOp);
+  if (ordering.position.empty())
+    return intervals;
+  uint64_t fallbackEnd = ordering.next - 1;
+  bool failedSize = false;
+  coreLikeOp->walk([&](memref::AllocOp allocation) {
+    auto size = getAllocationSize(allocation);
+    if (failed(size)) {
+      failedSize = true;
+      return;
+    }
+
+    LocalMemoryInterval interval {allocation, ordering.position.lookup(allocation), ordering.position.lookup(allocation), *size};
+    bool hasRuntimeTouch = false;
+    bool needsFallback = false;
+    if (failed(walkLocalMemoryUses(allocation.getResult(), [&](OpOperand& use) {
+      Operation* user = use.getOwner();
+      if (isRuntimeMemoryTouch(user)) {
+        auto [start, end] = getTouchRange(allocation.getResult(), user, ordering);
+        if (!hasRuntimeTouch) {
+          interval.start = start;
+          interval.end = end;
+          hasRuntimeTouch = true;
+        }
+        else {
+          interval.start = std::min(interval.start, start);
+          interval.end = std::max(interval.end, end);
+        }
+      }
+      else if (!isIgnoredUser(user)) {
+        needsFallback = true;
+      }
+      return success();
+    })))
+      needsFallback = true;
+    if (!hasRuntimeTouch || needsFallback) {
+      interval.start = std::min(interval.start, ordering.position.lookup(allocation));
+      interval.end = fallbackEnd;
+    }
+    intervals.push_back(interval);
+  });
+  if (failedSize)
+    return failure();
+  return intervals;
+}
+
+bool localMemoryLifetimesOverlap(const LocalMemoryInterval& lhs, const LocalMemoryInterval& rhs) {
+  return !(lhs.end < rhs.start || rhs.end < lhs.start);
+}
+
+} // namespace onnx_mlir::pim

src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp

@@ -1,17 +1,25 @@
 #pragma once
 
-#include "mlir/IR/Operation.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 
 #include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/SmallVector.h"
 
-namespace onnx_mlir {
-namespace pim {
+namespace onnx_mlir::pim {
+
+struct LocalMemoryInterval {
+  mlir::memref::AllocOp allocation;
+  uint64_t start = 0;
+  uint64_t end = 0;
+  size_t size = 0;
+};
+
+using LocalMemoryUseCallback = llvm::function_ref<mlir::LogicalResult(mlir::OpOperand&)>;
 
 bool isLocalMemoryAliasOp(mlir::Operation* op);
+mlir::LogicalResult walkLocalMemoryUses(mlir::Value root, LocalMemoryUseCallback callback);
+mlir::FailureOr<llvm::SmallVector<LocalMemoryInterval, 0>>
+analyzeLocalMemoryLifetimes(mlir::Operation* coreLikeOp);
+bool localMemoryLifetimesOverlap(const LocalMemoryInterval& lhs, const LocalMemoryInterval& rhs);
 
-mlir::LogicalResult walkLocalMemoryUses(
-  mlir::Value root,
-  llvm::function_ref<mlir::LogicalResult(mlir::Value, mlir::Operation*)> visitUser);
-
-} // namespace pim
-} // namespace onnx_mlir
+} // namespace onnx_mlir::pim

src/PIM/Dialect/Pim/CMakeLists.txt

@@ -4,7 +4,6 @@ add_onnx_mlir_dialect_doc(pim Pim.td)
 add_subdirectory(Analysis)
 add_subdirectory(Transforms/Bufferization)
 add_subdirectory(Transforms/HostConstantFolding)
-add_subdirectory(Transforms/MemoryCoalescing)
 add_subdirectory(Transforms/LocalMemoryPlanning)
 add_subdirectory(Transforms/Verification)
 

src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/CMakeLists.txt

@@ -8,7 +8,6 @@ add_pim_library(OMPimLocalMemoryPlanning
 
   LINK_LIBS PUBLIC
   OMPimCommon
-  OMPimCompilerOptions
   OMPimLocalMemoryLifetimeAnalysis
   PimOps
 )

src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp

@@ -1,668 +1,169 @@
-#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Pass/Pass.h"
-#include "mlir/IR/Value.h"
 
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/Support/MathExtras.h"
-#include "llvm/Support/FileSystem.h"
-#include "llvm/Support/raw_os_ostream.h"
-#include "llvm/Support/raw_ostream.h"
 
-#include <fstream>
-#include <memory>
 #include <numeric>
-#include <string>
 #include <tuple>
-#include <utility>
 
-#include "Common/Support/CheckedArithmetic.hpp"
-#include "Common/Support/ReportUtils.hpp"
-#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace llvm;
 using namespace mlir;
-using namespace onnx_mlir;
-
-namespace {
-
-struct MemoryTouchInterval {
-  uint64_t start = 0;
-  uint64_t end = 0;
-  Operation* firstTouchOp = nullptr;
-  Operation* lastTouchOp = nullptr;
-  uint64_t firstTouchPosition = 0;
-  uint64_t lastTouchPosition = 0;
-  bool hasRuntimeUse = false;
-  bool startUsedAllocFallback = false;
-  bool endUsedFallback = false;
-  bool escapesLoop = false;
-  std::string fallbackReason;
-};
-
-struct OperationOrdering {
-  llvm::DenseMap<Operation*, uint64_t> position;
-  llvm::DenseMap<Operation*, uint64_t> subtreeEnd;
-  uint64_t nextPosition = 0;
-};
-
-static std::string abbreviate(StringRef text, size_t maxLen) {
-  if (text.size() <= maxLen)
-    return text.str();
-  return (text.take_front(maxLen - 3) + "...").str();
-}
-
-static std::string summarizeValue(mlir::Value value, size_t maxLen = 72) {
-  std::string text;
-  llvm::raw_string_ostream os(text);
-  if (auto result = dyn_cast<OpResult>(value))
-    os << result.getOwner()->getName() << '#' << result.getResultNumber();
-  else if (auto blockArg = dyn_cast<BlockArgument>(value))
-    os << "block_arg#" << blockArg.getArgNumber();
-  else
-    os << "<unknown value>";
-  os << " : " << value.getType();
-  os.flush();
-  return abbreviate(text, maxLen);
-}
-
-static std::string summarizeOperation(Operation* op, size_t maxLen = 96) {
-  if (!op)
-    return "<none>";
-  return abbreviate(op->getName().getStringRef(), maxLen);
-}
-
-static void assignOperationOrdering(Operation* op, OperationOrdering& ordering) {
-  uint64_t position = ordering.nextPosition++;
-  ordering.position[op] = position;
-  uint64_t end = position;
-  for (Region& region : op->getRegions())
-    for (Block& block : region)
-      for (Operation& nestedOp : block) {
-        assignOperationOrdering(&nestedOp, ordering);
-        end = std::max(end, ordering.subtreeEnd.lookup(&nestedOp));
-      }
-  ordering.subtreeEnd[op] = end;
-}
-
-static OperationOrdering buildOperationOrdering(Operation* coreLikeOp) {
-  OperationOrdering ordering;
-  if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
-    return ordering;
-
-  for (Operation& op : coreLikeOp->getRegion(0).front())
-    assignOperationOrdering(&op, ordering);
-  return ordering;
-}
-
-static bool isRuntimeMemoryTouchOp(Operation* op) {
-  return isa<pim::PimMemCopyHostToDevOp,
-             pim::PimMemCopyDevToHostOp,
-             pim::PimMemCopyOp,
-             pim::PimReceiveOp,
-             pim::PimSendOp,
-             pim::PimConcatOp,
-             pim::PimVMMOp,
-             pim::PimTransposeOp,
-             pim::PimVVAddOp,
-             pim::PimVVSubOp,
-             pim::PimVVMulOp,
-             pim::PimVVMaxOp,
-             pim::PimVVDMulOp,
-             pim::PimVAvgOp,
-             pim::PimVReluOp,
-             pim::PimVTanhOp,
-             pim::PimVSigmOp,
-             pim::PimVSoftmaxOp>(op);
-}
-
-static bool isIgnoredLivenessUser(Operation* op) {
-  return pim::isLocalMemoryAliasOp(op) || isa<scf::ForOp, scf::YieldOp, memref::DeallocOp>(op)
-      || isCoreStaticAddressOp(op);
-}
-
-static bool isWithin(mlir::Value value, Region* region) {
-  if (!region)
-    return false;
-  if (auto blockArg = dyn_cast<BlockArgument>(value))
-    return blockArg.getOwner()->getParent() == region;
-  if (Operation* definingOp = value.getDefiningOp())
-    return definingOp->getParentRegion() == region || region->isAncestor(definingOp->getParentRegion());
-  return false;
-}
-
-static bool isNestedAllocation(Operation* coreLikeOp, memref::AllocOp allocOp) {
-  if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
-    return false;
-  return allocOp->getBlock() != &coreLikeOp->getRegion(0).front();
-}
-
-static void addFallbackReason(std::string& reason, StringRef newReason) {
-  if (newReason.empty())
-    return;
-  if (!reason.empty())
-    reason += "; ";
-  reason += newReason.str();
-}
-
-struct OrderedTouchRange {
-  uint64_t start = 0;
-  uint64_t end = 0;
-  bool escapedLoop = false;
-};
-
-static OrderedTouchRange
-getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const OperationOrdering& ordering) {
-  OrderedTouchRange range {ordering.position.lookup(user), ordering.position.lookup(user), false};
-  for (Operation* current = user; current; current = current->getParentOp()) {
-    auto forOp = dyn_cast<scf::ForOp>(current);
-    if (!forOp || isWithin(definingValue, &forOp.getRegion()))
-      continue;
-    range.start = std::min(range.start, ordering.position.lookup(forOp));
-    range.end = std::max(range.end, ordering.subtreeEnd.lookup(forOp));
-    range.escapedLoop = true;
-  }
-  return range;
-}
-
-static MemoryTouchInterval computeMemoryTouchInterval(memref::AllocOp allocOp,
-                                                      const OperationOrdering& ordering,
-                                                      uint64_t fallbackEnd) {
-  MemoryTouchInterval interval;
-  interval.start = ordering.position.lookup(allocOp);
-  interval.end = interval.start;
-
-  auto parentLoop = allocOp->getParentOfType<scf::ForOp>();
-  (void) pim::walkLocalMemoryUses(
-    allocOp.getResult(),
-    [&](mlir::Value value, Operation* user) {
-      if (auto forOp = dyn_cast<scf::ForOp>(user);
-          forOp && parentLoop && forOp != parentLoop && llvm::is_contained(forOp.getInitArgs(), value))
-        interval.escapesLoop = true;
-      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
-        auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
-        auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp());
-        auto indexSwitch = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp());
-        if (!forOp && !ifOp && !indexSwitch)
-          addFallbackReason(interval.fallbackReason, "yield without scf.for parent");
-        else if (forOp && parentLoop && forOp == parentLoop && llvm::is_contained(yieldOp.getOperands(), value))
-          interval.escapesLoop = true;
-      }
-
-      if (isRuntimeMemoryTouchOp(user)) {
-        uint64_t touchPosition = ordering.position.lookup(user);
-        if (!interval.hasRuntimeUse || touchPosition < interval.firstTouchPosition) {
-          interval.firstTouchPosition = touchPosition;
-          interval.firstTouchOp = user;
-        }
-        if (!interval.hasRuntimeUse || touchPosition > interval.lastTouchPosition) {
-          interval.lastTouchPosition = touchPosition;
-          interval.lastTouchOp = user;
-        }
-
-        OrderedTouchRange range = getEffectiveTouchRange(allocOp.getResult(), user, ordering);
-        interval.escapesLoop |= range.escapedLoop;
-        if (!interval.hasRuntimeUse) {
-          interval.start = range.start;
-          interval.end = range.end;
-          interval.hasRuntimeUse = true;
-        }
-        else {
-          if (range.start < interval.start)
-            interval.start = range.start;
-          if (range.end > interval.end)
-            interval.end = range.end;
-        }
-        return success();
-      }
-
-      if (isIgnoredLivenessUser(user))
-        return success();
-
-      addFallbackReason(interval.fallbackReason, "unhandled user op");
-      interval.endUsedFallback = true;
-      return success();
-    });
-
-  if (!interval.hasRuntimeUse) {
-    interval.startUsedAllocFallback = true;
-    interval.endUsedFallback = true;
-    interval.start = ordering.position.lookup(allocOp);
-    interval.end = fallbackEnd;
-    interval.firstTouchPosition = interval.start;
-    interval.lastTouchPosition = interval.end;
-    addFallbackReason(interval.fallbackReason, "no runtime memory touch");
-    return interval;
-  }
-
-  if (interval.endUsedFallback)
-    interval.end = std::max(interval.end, fallbackEnd);
-
-  return interval;
-}
-
-static FailureOr<size_t> getAllocSizeBytes(memref::AllocOp allocOp) {
-  auto type = dyn_cast<ShapedType>(allocOp.getType());
-  if (!type)
-    return failure();
-  auto checkedBytes = pim::getCheckedShapedTypeSizeInBytes(type, allocOp, "memory allocation byte size");
-  if (failed(checkedBytes))
-    return failure();
-  return pim::checkedSize(*checkedBytes, allocOp, "memory allocation byte size");
-}
-
-static bool intervalsOverlap(const LocalAllocInterval& lhs, const LocalAllocInterval& rhs) {
-  return !(lhs.end < rhs.start || rhs.end < lhs.start);
-}
-
-static uint64_t getSlotLogicalBytes(const PlannedPhysicalSlot& slot, ArrayRef<LocalAllocInterval> intervals) {
-  uint64_t slotLogicalBytes = 0;
-  for (size_t intervalIndex : slot.intervalIndices)
-    slotLogicalBytes += intervals[intervalIndex].size;
-  return slotLogicalBytes;
-}
-
-static bool placementsOverlap(const PlannedPhysicalSlot& lhs, const PlannedPhysicalSlot& rhs) {
-  return lhs.address < rhs.address + rhs.requiredSize && rhs.address < lhs.address + lhs.requiredSize;
-}
-
-static bool hasAddressReuse(size_t slotIndex, ArrayRef<PlannedPhysicalSlot> slots) {
-  if (slots[slotIndex].intervalIndices.size() > 1)
-    return true;
-  return llvm::any_of(llvm::enumerate(slots), [&](auto indexedSlot) {
-    return indexedSlot.index() != slotIndex && placementsOverlap(slots[slotIndex], indexedSlot.value());
-  });
-}
-
-} // namespace
-
-SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp) {
-  SmallVector<LocalAllocInterval, 0> intervals;
-  OperationOrdering ordering = buildOperationOrdering(coreLikeOp);
-  if (ordering.position.empty())
-    return intervals;
-
-  uint64_t fallbackEnd = ordering.nextPosition == 0 ? 0 : ordering.nextPosition - 1;
-  size_t nextIntervalId = 0;
-  coreLikeOp->walk([&](memref::AllocOp allocOp) {
-    auto checkedSize = getAllocSizeBytes(allocOp);
-    if (failed(checkedSize)) {
-      llvm::errs() << "Failed to compute local allocation size for value: ";
-      allocOp.getResult().print(llvm::errs());
-      llvm::errs() << "\n";
-      llvm_unreachable("Failed to compute local allocation size");
-    }
-
-    MemoryTouchInterval touchInterval = computeMemoryTouchInterval(allocOp, ordering, fallbackEnd);
-    LocalAllocInterval interval;
-    interval.id = nextIntervalId++;
-    interval.alloc = allocOp;
-    interval.start = touchInterval.start;
-    interval.end = touchInterval.end;
-    interval.size = *checkedSize;
-    interval.firstTouchOp = touchInterval.firstTouchOp;
-    interval.lastTouchOp = touchInterval.lastTouchOp;
-    interval.firstTouchPosition = touchInterval.firstTouchPosition;
-    interval.lastTouchPosition = touchInterval.lastTouchPosition;
-    interval.startUsedAllocFallback = touchInterval.startUsedAllocFallback;
-    interval.endUsedFallback = touchInterval.endUsedFallback;
-    interval.hasRuntimeUse = touchInterval.hasRuntimeUse;
-    interval.insideNestedRegion = isNestedAllocation(coreLikeOp, allocOp);
-    interval.escapesLoop = touchInterval.escapesLoop;
-    interval.fallbackReason = std::move(touchInterval.fallbackReason);
-    interval.valueSummary = summarizeValue(allocOp.getResult(), 88);
-    interval.firstTouchSummary = summarizeOperation(touchInterval.firstTouchOp);
-    interval.lastTouchSummary = summarizeOperation(touchInterval.lastTouchOp);
-    intervals.push_back(std::move(interval));
-  });
-
-  return intervals;
-}
-
-SmallVector<PlannedPhysicalSlot, 0> onnx_mlir::planPhysicalSlots(MutableArrayRef<LocalAllocInterval> intervals) {
-  SmallVector<PlannedPhysicalSlot, 0> slots;
-  SmallVector<size_t> intervalOrder(intervals.size());
-  std::iota(intervalOrder.begin(), intervalOrder.end(), 0);
-  llvm::stable_sort(intervalOrder, [&](size_t lhsIndex, size_t rhsIndex) {
-    const LocalAllocInterval& lhs = intervals[lhsIndex];
-    const LocalAllocInterval& rhs = intervals[rhsIndex];
-    if (lhs.size != rhs.size)
-      return lhs.size > rhs.size;
-    if (lhs.start != rhs.start)
-      return lhs.start < rhs.start;
-    if (lhs.end != rhs.end)
-      return lhs.end < rhs.end;
-    return lhs.id < rhs.id;
-  });
-
-  for (size_t intervalIndex : intervalOrder) {
-    LocalAllocInterval& interval = intervals[intervalIndex];
-    SmallVector<const PlannedPhysicalSlot*, 16> conflictingSlots;
-    SmallVector<size_t, 16> candidateAddresses = {0};
-    size_t currentPeak = 0;
-    for (const PlannedPhysicalSlot& slot : slots) {
-      currentPeak = std::max(currentPeak, slot.address + slot.requiredSize);
-      if (llvm::any_of(slot.intervalIndices, [&](size_t otherIndex) {
-            return intervalsOverlap(interval, intervals[otherIndex]);
-          })) {
-        conflictingSlots.push_back(&slot);
-        size_t candidate = llvm::alignTo(slot.address + slot.requiredSize, 4);
-        if (!llvm::is_contained(candidateAddresses, candidate))
-          candidateAddresses.push_back(candidate);
-      }
-    }
-
-    size_t bestAddress = std::numeric_limits<size_t>::max();
-    auto bestKey = std::tuple<size_t, size_t>(std::numeric_limits<size_t>::max(),
-                                              std::numeric_limits<size_t>::max());
-    for (size_t candidate : candidateAddresses) {
-      bool overlaps = llvm::any_of(conflictingSlots, [&](const PlannedPhysicalSlot* slot) {
-        return candidate < slot->address + slot->requiredSize && slot->address < candidate + interval.size;
-      });
-      if (overlaps)
-        continue;
-      auto candidateKey = std::tuple<size_t, size_t>(std::max(currentPeak, candidate + interval.size), candidate);
-      if (candidateKey < bestKey) {
-        bestKey = candidateKey;
-        bestAddress = candidate;
-      }
-    }
-    assert(bestAddress != std::numeric_limits<size_t>::max() && "address after all conflicts must be available");
-
-    auto reusable = llvm::find_if(slots, [&](const PlannedPhysicalSlot& slot) {
-      return slot.address == bestAddress && slot.requiredSize == interval.size;
-    });
-    if (reusable != slots.end()) {
-      reusable->intervalIndices.push_back(intervalIndex);
-      interval.slotPlanIndex = static_cast<size_t>(reusable - slots.begin());
-    }
-    else {
-      slots.push_back({slots.size(), interval.size, bestAddress, {intervalIndex}});
-      interval.slotPlanIndex = slots.size() - 1;
-    }
-  }
-
-  return slots;
-}
-
-CoreMemoryPlan onnx_mlir::buildCoreMemoryPlan(Operation* coreLikeOp) {
-  CoreMemoryPlan plan;
-  plan.coreLikeOp = coreLikeOp;
-  plan.intervals = buildLocalAllocIntervals(coreLikeOp);
-  plan.slots = planPhysicalSlots(plan.intervals);
-  return plan;
-}
-
-LogicalResult onnx_mlir::assignPhysicalSlotAddresses(CoreMemoryPlan& plan, size_t addressLimit) {
-  SmallVector<size_t> slotOrder(plan.slots.size());
-  std::iota(slotOrder.begin(), slotOrder.end(), 0);
-  llvm::stable_sort(slotOrder, [&](size_t lhsIndex, size_t rhsIndex) {
-    const PlannedPhysicalSlot& lhs = plan.slots[lhsIndex];
-    const PlannedPhysicalSlot& rhs = plan.slots[rhsIndex];
-    if (lhs.requiredSize != rhs.requiredSize)
-      return lhs.requiredSize > rhs.requiredSize;
-    return lhs.id < rhs.id;
-  });
-
-  size_t nextSlotId = 0;
-  for (size_t slotIndex : slotOrder) {
-    PlannedPhysicalSlot& slot = plan.slots[slotIndex];
-    if (slot.address > addressLimit || slot.requiredSize > addressLimit - slot.address) {
-      plan.coreLikeOp->emitError() << "PIM local memory plan exceeds the signed int32 address range while assigning "
-                                     "a "
-                                  << slot.requiredSize << " byte physical placement at address " << slot.address;
-      return failure();
-    }
-
-    slot.id = nextSlotId++;
-  }
-  return success();
-}
-
-std::string onnx_mlir::buildMemoryPlanReport(Operation* coreLikeOp,
-                                             ArrayRef<LocalAllocInterval> intervals,
-                                             ArrayRef<PlannedPhysicalSlot> slots,
-                                             size_t addressLimit,
-                                             PimMemoryReportLevel reportLevel) {
-  std::string artifacts;
-
-  uint64_t totalLogicalBytes = 0;
-  uint64_t totalPhysicalBytes = 0;
-  uint64_t fallbackIntervals = 0;
-  uint64_t noRuntimeTouchIntervals = 0;
-  uint64_t reusedAllocations = 0;
-  uint64_t nestedIntervals = 0;
-  uint64_t loopEscapingIntervals = 0;
-  size_t largestLogicalAllocation = 0;
-  size_t largestPhysicalSlot = 0;
-  size_t maximumAssignedAddress = 0;
-
-  for (const LocalAllocInterval& interval : intervals) {
-    totalLogicalBytes += interval.size;
-    largestLogicalAllocation = std::max(largestLogicalAllocation, interval.size);
-    if (interval.startUsedAllocFallback || interval.endUsedFallback)
-      ++fallbackIntervals;
-    if (!interval.hasRuntimeUse)
-      ++noRuntimeTouchIntervals;
-    if (interval.insideNestedRegion)
-      ++nestedIntervals;
-    if (interval.escapesLoop)
-      ++loopEscapingIntervals;
-  }
-  for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
-    const PlannedPhysicalSlot& slot = slots[slotIndex];
-    largestPhysicalSlot = std::max(largestPhysicalSlot, slot.requiredSize);
-    maximumAssignedAddress = std::max(maximumAssignedAddress, slot.address + slot.requiredSize);
-    if (hasAddressReuse(slotIndex, slots))
-      reusedAllocations += slot.intervalIndices.size();
-  }
-  totalPhysicalBytes = maximumAssignedAddress;
-
-  uint64_t savedBytes = totalLogicalBytes >= totalPhysicalBytes ? totalLogicalBytes - totalPhysicalBytes : 0;
-  double savedPercent =
-    totalLogicalBytes == 0 ? 0.0 : 100.0 * static_cast<double>(savedBytes) / static_cast<double>(totalLogicalBytes);
-
-  raw_string_ostream os(artifacts);
-  os << "=== PIM Memory Liveness Report ===\n";
-  os << "Op: " << coreLikeOp->getName() << "\n";
-  os << "Summary:\n";
-  os << "  logical allocation bytes:  " << formatReportMemory(totalLogicalBytes) << " (" << totalLogicalBytes << ")\n";
-  os << "  physical allocation bytes: " << formatReportMemory(totalPhysicalBytes) << " (" << totalPhysicalBytes
-     << ")\n";
-  os << "  saved bytes:               " << formatReportMemory(savedBytes) << " (" << savedBytes << ")\n";
-  os << "  saved percent: " << format("%.2f%%", savedPercent) << "\n";
-  os << "  intervals: " << intervals.size() << "\n";
-  os << "  address placements: " << slots.size() << "\n";
-  os << "  allocations sharing address space: " << reusedAllocations << "\n";
-  os << "  fallback intervals: " << fallbackIntervals << "\n";
-  os << "  intervals with no runtime memory touch: " << noRuntimeTouchIntervals << "\n";
-  os << "  nested allocations: " << nestedIntervals << "\n";
-  os << "  loop-escaping allocations: " << loopEscapingIntervals << "\n";
-  os << "  largest logical allocation: " << largestLogicalAllocation << "\n";
-  os << "  largest address placement: " << largestPhysicalSlot << "\n";
-  os << "  address limit: " << addressLimit << "\n";
-  os << "  peak physical memory: " << formatReportMemory(maximumAssignedAddress) << " (" << maximumAssignedAddress
-     << ")\n";
-  os << "  maximum assigned address: " << maximumAssignedAddress << "\n";
-
-  SmallVector<const PlannedPhysicalSlot*> reusedSlots;
-  SmallVector<const PlannedPhysicalSlot*> singleUseSlots;
-  for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
-    if (hasAddressReuse(slotIndex, slots))
-      reusedSlots.push_back(&slots[slotIndex]);
-    else
-      singleUseSlots.push_back(&slots[slotIndex]);
-  }
-
-  llvm::stable_sort(reusedSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) {
-    uint64_t lhsLogicalBytes = getSlotLogicalBytes(*lhs, intervals);
-    uint64_t rhsLogicalBytes = getSlotLogicalBytes(*rhs, intervals);
-    if (lhs->intervalIndices.size() != rhs->intervalIndices.size())
-      return lhs->intervalIndices.size() > rhs->intervalIndices.size();
-    if (lhsLogicalBytes != rhsLogicalBytes)
-      return lhsLogicalBytes > rhsLogicalBytes;
-    if (lhs->requiredSize != rhs->requiredSize)
-      return lhs->requiredSize > rhs->requiredSize;
-    return lhs->id < rhs->id;
-  });
-  llvm::stable_sort(singleUseSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) {
-    if (lhs->requiredSize != rhs->requiredSize)
-      return lhs->requiredSize > rhs->requiredSize;
-    return lhs->id < rhs->id;
-  });
-
-  constexpr size_t kSummaryReuseLimit = 6;
-  constexpr size_t kSummaryOffenderLimit = 10;
-
-  os << "\nBest Address Reuse:\n";
-  if (reusedSlots.empty()) {
-      os << "  no address ranges were reused\n";
-  }
-  else {
-    for (const PlannedPhysicalSlot* slot : ArrayRef(reusedSlots).take_front(kSummaryReuseLimit)) {
-      uint64_t slotLogicalBytes = getSlotLogicalBytes(*slot, intervals);
-      os << "  slot #" << slot->id << "  addr=" << slot->address << "  size=" << formatReportMemory(slot->requiredSize)
-         << "  intervals=" << slot->intervalIndices.size() << "  logical_sum=" << formatReportMemory(slotLogicalBytes)
-         << "\n";
-    }
-  }
-
-  os << "\nTop Offenders:\n";
-  for (const PlannedPhysicalSlot* slot : ArrayRef(singleUseSlots).take_front(kSummaryOffenderLimit)) {
-    const LocalAllocInterval& interval = intervals[slot->intervalIndices.front()];
-    os << "  slot #" << slot->id << " is single-use"
-       << "  size=" << formatReportMemory(slot->requiredSize) << "  interval=#" << interval.id
-       << "  value=" << abbreviate(interval.valueSummary, 56) << "\n";
-    os << "    first=" << abbreviate(interval.firstTouchSummary, 40)
-       << "  last=" << abbreviate(interval.lastTouchSummary, 40)
-       << "  nested=" << (interval.insideNestedRegion ? "yes" : "no")
-       << "  escapes_loop=" << (interval.escapesLoop ? "yes" : "no") << "\n";
-  }
-  if (singleUseSlots.empty())
-    os << "  no obvious blockers detected in this core\n";
-
-  if (reportLevel == PimMemoryReportFull) {
-    os << "\nSlot Reuse:\n";
-    for (const PlannedPhysicalSlot& slot : slots) {
-      uint64_t slotLogicalBytes = getSlotLogicalBytes(slot, intervals);
-      os << "  slot #" << slot.id << "  addr=" << slot.address << "  size=" << formatReportMemory(slot.requiredSize)
-         << " (" << slot.requiredSize << ")"
-         << "  intervals=" << slot.intervalIndices.size() << "  logical_sum=" << formatReportMemory(slotLogicalBytes)
-         << "\n";
-      for (size_t intervalIndex : slot.intervalIndices) {
-        const LocalAllocInterval& interval = intervals[intervalIndex];
-        os << "    [" << interval.start << "," << interval.end << "]"
-           << "  #" << interval.id << "  logical=" << formatReportMemory(interval.size) << " (" << interval.size
-           << ")"
-           << "  nested=" << (interval.insideNestedRegion ? "yes" : "no")
-           << "  escapes_loop=" << (interval.escapesLoop ? "yes" : "no")
-           << "  first=" << abbreviate(interval.firstTouchSummary, 48)
-           << "  last=" << abbreviate(interval.lastTouchSummary, 48) << "\n";
-        os << "      value=" << abbreviate(interval.valueSummary, 72) << "\n";
-        if (!interval.fallbackReason.empty())
-          os << "      fallback_reason=" << interval.fallbackReason << "\n";
-      }
-    }
-  }
-  os.flush();
-
-  return artifacts;
-}
 
 namespace onnx_mlir {
 namespace {
 
+static bool rangesOverlap(size_t lhsAddress, size_t lhsSize, size_t rhsAddress, size_t rhsSize) {
+  return lhsAddress < rhsAddress + rhsSize && rhsAddress < lhsAddress + lhsSize;
+}
+
+static FailureOr<size_t> alignedEnd(size_t address, size_t size, size_t limit) {
+  if (address > limit || size > limit - address)
+    return failure();
+  size_t end = address + size;
+  size_t padding = (4 - end % 4) % 4;
+  if (padding > limit - end)
+    return failure();
+  return end + padding;
+}
+
 struct PimLocalMemoryPlanningPass : PassWrapper<PimLocalMemoryPlanningPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimLocalMemoryPlanningPass)
 
   StringRef getArgument() const override { return "pim-local-memory-planning"; }
   StringRef getDescription() const override {
-    return "Plan liveness-based physical slots and addresses for PIM core-local memory";
+    return "Plan liveness-based addresses for PIM core-local memory";
   }
 
   void runOnOperation() override {
-    std::fstream reportFile;
-    std::unique_ptr<raw_os_ostream> report;
-    if (pimMemoryReport != PimMemoryReportNone) {
-      reportFile = openReportFileWithExtension("pim_memory_liveness_report", "txt");
-      if (reportFile.is_open())
-        report = std::make_unique<raw_os_ostream>(reportFile);
-    }
-    else if (std::string outputRoot = getOutputDir(); !outputRoot.empty()) {
-      sys::fs::remove(outputRoot + "/reports/pim_memory_liveness_report.txt");
-    }
-
     Builder builder(&getContext());
-    uint64_t nextBatchId = 0;
-    bool failedPlanning = false;
+    bool hasFailure = false;
     getOperation().walk([&](Operation* op) {
-      if (failedPlanning || !isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
+      if (hasFailure || !isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
         return;
 
-      CoreMemoryPlan plan = buildCoreMemoryPlan(op);
-      if (failed(assignPhysicalSlotAddresses(plan, kPimLocalMemoryAddressLimit))) {
-        failedPlanning = true;
-        return;
-      }
-
-      size_t arenaSize = 0;
-      for (const PlannedPhysicalSlot& placement : plan.slots)
-        arenaSize = std::max(arenaSize, placement.address + placement.requiredSize);
-      for (const LocalAllocInterval& interval : plan.intervals) {
-        const PlannedPhysicalSlot& slot = plan.slots[interval.slotPlanIndex];
-        interval.alloc->setAttr(kLocalMemoryAddressAttrName, builder.getI64IntegerAttr(slot.address));
-        interval.alloc->setAttr(kLocalMemorySlotAttrName, builder.getI64IntegerAttr(0));
-        interval.alloc->setAttr(kLocalMemorySlotSizeAttrName, builder.getI64IntegerAttr(arenaSize));
-      }
-      uint64_t fallbackIntervals = llvm::count_if(plan.intervals, [](const LocalAllocInterval& interval) {
-        return interval.startUsedAllocFallback || interval.endUsedFallback;
+      op->removeAttr(kLocalMemorySizeAttrName);
+      for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
+        op->removeAttr(name);
+      op->walk([&](memref::AllocOp allocation) {
+        allocation->removeAttr(kLocalMemoryAddressAttrName);
+        for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
+          allocation->removeAttr(name);
       });
-      uint64_t nestedSingleUseIntervals = llvm::count_if(plan.intervals, [&](const LocalAllocInterval& interval) {
-        return interval.insideNestedRegion && !hasAddressReuse(interval.slotPlanIndex, plan.slots);
-      });
-      op->setAttr(kLocalMemoryFallbackCountAttrName, builder.getI64IntegerAttr(fallbackIntervals));
-      op->setAttr(kLocalMemoryNestedSingleUseCountAttrName, builder.getI64IntegerAttr(nestedSingleUseIntervals));
 
-      if (!report)
-        return;
-      std::string planReport = buildMemoryPlanReport(
-        op, plan.intervals, plan.slots, kPimLocalMemoryAddressLimit, pimMemoryReport);
-      if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
-        *report << "Core " << coreOp.getCoreId() << ":\n";
-        *report << planReport;
+      auto plan = buildCoreMemoryPlan(op);
+      if (failed(plan)) {
+        hasFailure = true;
         return;
       }
-
-      SmallVector<int32_t> coreIds = getBatchCoreIds(cast<pim::PimCoreBatchOp>(op));
-      llvm::sort(coreIds);
-      coreIds.erase(std::unique(coreIds.begin(), coreIds.end()), coreIds.end());
-      uint64_t batchId = nextBatchId++;
-      for (int32_t coreId : coreIds)
-        *report << "Batch " << batchId << " core " << coreId << ":\n" << planReport;
+      op->setAttr(kLocalMemorySizeAttrName, builder.getI64IntegerAttr(plan->arenaSize));
+      for (const LocalMemoryPlacement& placement : plan->placements)
+        for (size_t intervalIndex : placement.intervalIndices)
+          plan->intervals[intervalIndex].allocation->setAttr(
+            kLocalMemoryAddressAttrName, builder.getI64IntegerAttr(placement.address));
     });
-
-    if (report) {
-      report->flush();
-      reportFile.close();
-    }
-    if (failedPlanning) {
+    if (hasFailure) {
       signalPassFailure();
       return;
     }
-    dumpModule(getOperation(), "pim4_memory_planned");
+    dumpModule(getOperation(), "pim3_memory_planned");
   }
 };
 
 } // namespace
 
+FailureOr<SmallVector<LocalMemoryPlacement, 0>>
+planLocalMemoryPlacements(ArrayRef<pim::LocalMemoryInterval> intervals, size_t addressLimit) {
+  SmallVector<LocalMemoryPlacement, 0> placements;
+  SmallVector<size_t> order(intervals.size());
+  std::iota(order.begin(), order.end(), 0);
+  llvm::stable_sort(order, [&](size_t lhsIndex, size_t rhsIndex) {
+    const auto& lhs = intervals[lhsIndex];
+    const auto& rhs = intervals[rhsIndex];
+    if (lhs.size != rhs.size)
+      return lhs.size > rhs.size;
+    if (lhs.start != rhs.start)
+      return lhs.start < rhs.start;
+    if (lhs.end != rhs.end)
+      return lhs.end < rhs.end;
+    return lhsIndex < rhsIndex;
+  });
+
+  for (size_t intervalIndex : order) {
+    const auto& interval = intervals[intervalIndex];
+    SmallVector<const LocalMemoryPlacement*, 16> conflicts;
+    SmallVector<size_t, 16> candidates = {0};
+    size_t currentPeak = 0;
+    for (const LocalMemoryPlacement& placement : placements) {
+      auto end = alignedEnd(placement.address, placement.size, addressLimit);
+      if (failed(end))
+        return failure();
+      currentPeak = std::max(currentPeak, *end);
+      if (!llvm::any_of(placement.intervalIndices, [&](size_t otherIndex) {
+            return pim::localMemoryLifetimesOverlap(interval, intervals[otherIndex]);
+          }))
+        continue;
+      conflicts.push_back(&placement);
+      if (!llvm::is_contained(candidates, *end))
+        candidates.push_back(*end);
+    }
+
+    size_t bestAddress = std::numeric_limits<size_t>::max();
+    auto bestKey = std::tuple(std::numeric_limits<size_t>::max(), bestAddress);
+    for (size_t candidate : candidates) {
+      auto end = alignedEnd(candidate, interval.size, addressLimit);
+      if (failed(end) || llvm::any_of(conflicts, [&](const LocalMemoryPlacement* placement) {
+            return rangesOverlap(candidate, interval.size, placement->address, placement->size);
+          }))
+        continue;
+      auto key = std::tuple(std::max(currentPeak, *end), candidate);
+      if (key < bestKey) {
+        bestKey = key;
+        bestAddress = candidate;
+      }
+    }
+    if (bestAddress == std::numeric_limits<size_t>::max())
+      return failure();
+
+    auto reusable = llvm::find_if(placements, [&](const LocalMemoryPlacement& placement) {
+      return placement.address == bestAddress && placement.size == interval.size;
+    });
+    if (reusable != placements.end())
+      reusable->intervalIndices.push_back(intervalIndex);
+    else
+      placements.push_back({bestAddress, interval.size, {intervalIndex}});
+  }
+  return placements;
+}
+
+FailureOr<CoreMemoryPlan> buildCoreMemoryPlan(Operation* coreLikeOp) {
+  CoreMemoryPlan plan;
+  plan.coreLikeOp = coreLikeOp;
+  auto intervals = pim::analyzeLocalMemoryLifetimes(coreLikeOp);
+  if (failed(intervals))
+    return failure();
+  plan.intervals = std::move(*intervals);
+  auto placements = planLocalMemoryPlacements(plan.intervals, kPimLocalMemoryAddressLimit);
+  if (failed(placements)) {
+    coreLikeOp->emitError("PIM local-memory plan exceeds the signed int32 address range");
+    return failure();
+  }
+  plan.placements = std::move(*placements);
+  for (const LocalMemoryPlacement& placement : plan.placements) {
+    auto end = alignedEnd(placement.address, placement.size, kPimLocalMemoryAddressLimit);
+    if (failed(end)) {
+      coreLikeOp->emitError("PIM local-memory plan has invalid address arithmetic");
+      return failure();
+    }
+    plan.arenaSize = std::max(plan.arenaSize, *end);
+  }
+  return plan;
+}
+
 std::unique_ptr<Pass> createPimLocalMemoryPlanningPass() {
   return std::make_unique<PimLocalMemoryPlanningPass>();
 }

src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp

@@ -1,64 +1,24 @@
 #pragma once
 
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <limits>
-#include <string>
-
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
 
 namespace onnx_mlir {
 
-struct LocalAllocInterval {
-  size_t id = 0;
-  mlir::memref::AllocOp alloc;
-  uint64_t start = 0;
-  uint64_t end = 0;
-  size_t size = 0;
-  mlir::Operation* firstTouchOp = nullptr;
-  mlir::Operation* lastTouchOp = nullptr;
-  uint64_t firstTouchPosition = 0;
-  uint64_t lastTouchPosition = 0;
-  bool startUsedAllocFallback = false;
-  bool endUsedFallback = false;
-  bool hasRuntimeUse = false;
-  bool insideNestedRegion = false;
-  bool escapesLoop = false;
-  std::string fallbackReason;
-  std::string valueSummary;
-  std::string firstTouchSummary;
-  std::string lastTouchSummary;
-  size_t slotPlanIndex = std::numeric_limits<size_t>::max();
-};
-
-struct PlannedPhysicalSlot {
-  size_t id = std::numeric_limits<size_t>::max();
-  size_t requiredSize = 0;
+struct LocalMemoryPlacement {
   size_t address = 0;
+  size_t size = 0;
   llvm::SmallVector<size_t, 8> intervalIndices;
 };
 
 struct CoreMemoryPlan {
   mlir::Operation* coreLikeOp = nullptr;
-  llvm::SmallVector<LocalAllocInterval, 0> intervals;
-  llvm::SmallVector<PlannedPhysicalSlot, 0> slots;
+  llvm::SmallVector<pim::LocalMemoryInterval, 0> intervals;
+  llvm::SmallVector<LocalMemoryPlacement, 0> placements;
+  size_t arenaSize = 0;
 };
 
-llvm::SmallVector<LocalAllocInterval, 0> buildLocalAllocIntervals(mlir::Operation* coreLikeOp);
-
-llvm::SmallVector<PlannedPhysicalSlot, 0> planPhysicalSlots(llvm::MutableArrayRef<LocalAllocInterval> intervals);
-
-CoreMemoryPlan buildCoreMemoryPlan(mlir::Operation* coreLikeOp);
-
-mlir::LogicalResult assignPhysicalSlotAddresses(CoreMemoryPlan& plan, size_t addressLimit);
-
-std::string buildMemoryPlanReport(mlir::Operation* coreLikeOp,
-                                  llvm::ArrayRef<LocalAllocInterval> intervals,
-                                  llvm::ArrayRef<PlannedPhysicalSlot> slots,
-                                  size_t addressLimit,
-                                  PimMemoryReportLevel reportLevel);
+mlir::FailureOr<llvm::SmallVector<LocalMemoryPlacement, 0>>
+planLocalMemoryPlacements(llvm::ArrayRef<pim::LocalMemoryInterval> intervals, size_t addressLimit);
+mlir::FailureOr<CoreMemoryPlan> buildCoreMemoryPlan(mlir::Operation* coreLikeOp);
 
 } // namespace onnx_mlir

src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/CMakeLists.txt

@@ -1,15 +0,0 @@
-add_pim_library(OMPimMemoryCoalescing
-  MemoryCoalescing.cpp
-  MemoryCoalescing.hpp
-  MemoryCoalescingPass.cpp
-
-  EXCLUDE_FROM_OM_LIBS
-
-  INCLUDE_DIRS PUBLIC
-  ${PIM_PUBLIC_INCLUDE_DIRS}
-
-  LINK_LIBS PUBLIC
-  OMPimCommon
-  OMPimLocalMemoryLifetimeAnalysis
-  PimOps
-)

src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.cpp

@@ -1,194 +0,0 @@
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <limits>
-
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace pim {
-
-namespace {
-
-static bool isCandidateAllocType(MemRefType type) {
-  return type && type.hasStaticShape() && type.getLayout().isIdentity()
-      && hasByteSizedElementType(type.getElementType());
-}
-
-static uint64_t getTypeSizeBytes(MemRefType type) {
-  return static_cast<uint64_t>(type.getNumElements() * getElementTypeSizeInBytes(type.getElementType()));
-}
-
-static Operation* getTopLevelAncestorInBlock(Operation* op, Block& block) {
-  Operation* current = op;
-  while (current && current->getBlock() != &block)
-    current = current->getParentOp();
-  return current;
-}
-
-static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis);
-
-static FailureOr<uint64_t>
-getLastUseInstruction(memref::AllocOp allocOp, Block& scopeBlock, const DenseMap<Operation*, uint64_t>& opOrder) {
-  uint64_t endInstruction = opOrder.lookup(allocOp);
-  if (failed(walkLocalMemoryUses(allocOp.getResult(), [&](Value, Operation* user) {
-      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
-        if (!isa<scf::ForOp>(yieldOp->getParentOp()))
-          return failure();
-      }
-      Operation* orderedUser = getTopLevelAncestorInBlock(user, scopeBlock);
-      if (!orderedUser)
-        return failure();
-      auto order = opOrder.find(orderedUser);
-      if (order == opOrder.end())
-        return failure();
-      endInstruction = std::max(endInstruction, order->second);
-      return success();
-    })))
-    return failure();
-
-  return endInstruction;
-}
-
-static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis) {
-  for (Operation& op : block)
-    for (Region& region : op.getRegions())
-      for (Block& nestedBlock : region)
-        analyzeBlock(nestedBlock, analysis);
-
-  DenseMap<Operation*, uint64_t> opOrder;
-  uint64_t nextInstruction = 0;
-  for (Operation& op : block)
-    opOrder.try_emplace(&op, nextInstruction++);
-
-  MemoryCoalescingBlockAnalysis blockAnalysis;
-  blockAnalysis.block = &block;
-
-  for (Operation& op : block) {
-    auto allocOp = dyn_cast<memref::AllocOp>(&op);
-    if (!allocOp)
-      continue;
-
-    auto allocType = dyn_cast<MemRefType>(allocOp.getType());
-    if (!isCandidateAllocType(allocType)) {
-      ++blockAnalysis.skippedAllocations;
-      continue;
-    }
-
-    auto endInstruction = getLastUseInstruction(allocOp, block, opOrder);
-    if (failed(endInstruction)) {
-      ++blockAnalysis.skippedAllocations;
-      continue;
-    }
-
-    blockAnalysis.candidates.push_back(
-      AllocationCandidate {allocOp, opOrder.lookup(allocOp), *endInstruction, getTypeSizeBytes(allocType)});
-  }
-
-  analysis.skippedAllocations += blockAnalysis.skippedAllocations;
-  if (!blockAnalysis.candidates.empty() || blockAnalysis.skippedAllocations != 0)
-    analysis.blocks.push_back(std::move(blockAnalysis));
-}
-
-} // namespace
-
-uint64_t MemoryCoalescingAnalysis::getCandidateCount() const {
-  uint64_t total = 0;
-  for (const MemoryCoalescingBlockAnalysis& block : blocks)
-    total += block.candidates.size();
-  return total;
-}
-
-uint64_t MemoryCoalescingAnalysis::getCandidateBytes() const {
-  uint64_t total = 0;
-  for (const MemoryCoalescingBlockAnalysis& block : blocks)
-    for (const AllocationCandidate& candidate : block.candidates)
-      total += candidate.sizeBytes;
-  return total;
-}
-
-MemoryCoalescingAnalysis analyzeMemoryCoalescingCandidates(Operation* coreLikeOp) {
-  MemoryCoalescingAnalysis analysis;
-  if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
-    return analysis;
-
-  analyzeBlock(coreLikeOp->getRegion(0).front(), analysis);
-  return analysis;
-}
-
-MemoryCoalescingStats
-coalesceMemory(Operation* coreLikeOp, const MemoryCoalescingAnalysis& analysis, RewriterBase& rewriter) {
-  (void) coreLikeOp;
-
-  MemoryCoalescingStats stats;
-  stats.skippedAllocations = analysis.skippedAllocations;
-
-  for (const MemoryCoalescingBlockAnalysis& blockAnalysis : analysis.blocks) {
-    auto candidates = blockAnalysis.candidates;
-    llvm::sort(candidates, [](const AllocationCandidate& lhs, const AllocationCandidate& rhs) {
-      if (lhs.startInstruction != rhs.startInstruction)
-        return lhs.startInstruction < rhs.startInstruction;
-      return lhs.endInstruction < rhs.endInstruction;
-    });
-
-    struct ActiveStorage {
-      memref::AllocOp root;
-      uint64_t endInstruction = 0;
-    };
-
-    SmallVector<ActiveStorage> active;
-    SmallVector<memref::AllocOp> freeList;
-
-    for (AllocationCandidate& candidate : candidates) {
-      for (auto it = active.begin(); it != active.end();) {
-        if (it->endInstruction < candidate.startInstruction) {
-          freeList.push_back(it->root);
-          it = active.erase(it);
-          continue;
-        }
-        ++it;
-      }
-
-      auto bestFit = freeList.end();
-      uint64_t bestFitBytes = std::numeric_limits<uint64_t>::max();
-      auto candidateType = cast<MemRefType>(candidate.alloc.getType());
-      for (auto it = freeList.begin(); it != freeList.end(); ++it) {
-        auto freeType = cast<MemRefType>((*it).getType());
-        if (freeType != candidateType)
-          continue;
-
-        uint64_t freeBytes = getTypeSizeBytes(freeType);
-        if (freeBytes < candidate.sizeBytes || freeBytes >= bestFitBytes)
-          continue;
-
-        bestFit = it;
-        bestFitBytes = freeBytes;
-      }
-
-      if (bestFit == freeList.end()) {
-        active.push_back(ActiveStorage {candidate.alloc, candidate.endInstruction});
-        continue;
-      }
-
-      memref::AllocOp root = *bestFit;
-      freeList.erase(bestFit);
-      candidate.alloc.getResult().replaceAllUsesWith(root.getResult());
-      rewriter.eraseOp(candidate.alloc);
-      active.push_back(ActiveStorage {root, candidate.endInstruction});
-      ++stats.removedAllocs;
-      stats.savedBytes += candidate.sizeBytes;
-    }
-  }
-
-  return stats;
-}
-
-} // namespace pim
-} // namespace onnx_mlir

src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp

@@ -1,44 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/IR/PatternMatch.h"
-
-#include "llvm/ADT/SmallVector.h"
-
-namespace onnx_mlir {
-namespace pim {
-
-struct AllocationCandidate {
-  mlir::memref::AllocOp alloc;
-  uint64_t startInstruction = 0;
-  uint64_t endInstruction = 0;
-  uint64_t sizeBytes = 0;
-};
-
-struct MemoryCoalescingBlockAnalysis {
-  mlir::Block* block = nullptr;
-  llvm::SmallVector<AllocationCandidate> candidates;
-  uint64_t skippedAllocations = 0;
-};
-
-struct MemoryCoalescingAnalysis {
-  llvm::SmallVector<MemoryCoalescingBlockAnalysis> blocks;
-  uint64_t skippedAllocations = 0;
-
-  uint64_t getCandidateCount() const;
-  uint64_t getCandidateBytes() const;
-};
-
-struct MemoryCoalescingStats {
-  uint64_t removedAllocs = 0;
-  uint64_t savedBytes = 0;
-  uint64_t skippedAllocations = 0;
-};
-
-MemoryCoalescingAnalysis analyzeMemoryCoalescingCandidates(mlir::Operation* coreLikeOp);
-
-MemoryCoalescingStats
-coalesceMemory(mlir::Operation* coreLikeOp, const MemoryCoalescingAnalysis& analysis, mlir::RewriterBase& rewriter);
-
-} // namespace pim
-} // namespace onnx_mlir

src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp

@@ -1,101 +0,0 @@
-#include "mlir/Pass/Pass.h"
-
-#include "llvm/Support/raw_os_ostream.h"
-
-#include <fstream>
-
-#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace {
-
-struct CoalescingReportRow {
-  uint64_t candidates = 0;
-  uint64_t logicalBytes = 0;
-  uint64_t skipped = 0;
-  uint64_t removed = 0;
-  uint64_t savedBytes = 0;
-};
-
-struct CoalescingReportEntry {
-  std::string label;
-  uint64_t coreCount = 1;
-  CoalescingReportRow row;
-};
-
-static void emitReport(ArrayRef<CoalescingReportEntry> entries) {
-  std::fstream file = openReportFile("memory_coalescing_report");
-  if (!file.is_open())
-    return;
-  llvm::raw_os_ostream os(file);
-  CoalescingReportRow total;
-  for (const CoalescingReportEntry& entry : entries) {
-    total.candidates += entry.row.candidates * entry.coreCount;
-    total.logicalBytes += entry.row.logicalBytes * entry.coreCount;
-    total.skipped += entry.row.skipped * entry.coreCount;
-    total.removed += entry.row.removed * entry.coreCount;
-    total.savedBytes += entry.row.savedBytes * entry.coreCount;
-  }
-  printReportTotalsBlock(os,
-                         {{"Number of candidates", std::to_string(total.candidates)},
-                          {"Logical candidate bytes", formatReportMemory(total.logicalBytes)},
-                          {"Skipped allocations", std::to_string(total.skipped)},
-                          {"Removed allocations", std::to_string(total.removed)},
-                          {"Saved memory", formatReportMemory(total.savedBytes)}});
-  for (const CoalescingReportEntry& entry : entries) {
-    os << "\n" << entry.label << ":\n";
-    printReportFlatFields(os,
-                          {{"Number of candidates", std::to_string(entry.row.candidates)},
-                           {"Logical candidate bytes", formatReportMemory(entry.row.logicalBytes)},
-                           {"Skipped allocations", std::to_string(entry.row.skipped)},
-                           {"Removed allocations", std::to_string(entry.row.removed)},
-                           {"Saved memory", formatReportMemory(entry.row.savedBytes)}});
-  }
-  os.flush();
-}
-
-struct PimMemoryCoalescingPass : PassWrapper<PimMemoryCoalescingPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimMemoryCoalescingPass)
-
-  StringRef getArgument() const override { return "pim-memory-coalescing"; }
-  StringRef getDescription() const override { return "Safely coalesce compatible block-local PIM allocations"; }
-
-  void runOnOperation() override {
-    IRRewriter rewriter(&getContext());
-    SmallVector<CoalescingReportEntry, 32> reportEntries;
-    uint64_t nextBatchId = 0;
-    getOperation().walk([&](Operation* op) {
-      if (!isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
-        return;
-      auto analysis = pim::analyzeMemoryCoalescingCandidates(op);
-      auto stats = pim::coalesceMemory(op, analysis, rewriter);
-      CoalescingReportRow row {analysis.getCandidateCount(),
-                               analysis.getCandidateBytes(),
-                               stats.skippedAllocations,
-                               stats.removedAllocs,
-                               stats.savedBytes};
-      if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
-        reportEntries.push_back({"Core " + std::to_string(coreOp.getCoreId()), 1, row});
-        return;
-      }
-      SmallVector<int32_t> coreIds = getBatchCoreIds(cast<pim::PimCoreBatchOp>(op));
-      reportEntries.push_back(
-        {"Batch " + std::to_string(nextBatchId++), std::max<uint64_t>(1, coreIds.size()), row});
-    });
-    emitReport(reportEntries);
-    dumpModule(getOperation(), "pim3_coalesced");
-  }
-};
-
-} // namespace
-
-std::unique_ptr<Pass> createPimMemoryCoalescingPass() { return std::make_unique<PimMemoryCoalescingPass>(); }
-
-} // namespace onnx_mlir

src/PIM/Dialect/Pim/Transforms/Verification/CMakeLists.txt

@@ -7,6 +7,7 @@ add_pim_library(OMPimVerification
   OMPimCommon
   OMPimCompilerOptions
   OMPimBufferization
+  OMPimLocalMemoryLifetimeAnalysis
   PimOps
   SpatialOps
 )

src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp

@@ -8,6 +8,7 @@
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include <map>
 #include <string>
 
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
@@ -18,6 +19,7 @@
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -109,61 +111,127 @@ static bool isCoreWeightBlockArgument(Value value) {
   return false;
 }
 
-struct VerifiedLocalMemorySlot {
-  uint64_t size = 0;
-};
-
 static LogicalResult
 verifyLocalMemoryPlan(Operation* coreLikeOp, pim::CappedDiagnosticReporter& diagnostics) {
-  DenseMap<uint64_t, VerifiedLocalMemorySlot> slots;
   bool hasFailure = false;
-  auto fallbackAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryFallbackCountAttrName);
-  auto nestedSingleUseAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryNestedSingleUseCountAttrName);
-  if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) {
-    diagnostics.report(coreLikeOp, [&](Operation* op) {
-      op->emitError("requires complete non-negative PIM local-memory planning summary attributes");
+  for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
+    if (coreLikeOp->hasAttr(name)) {
+      diagnostics.report(coreLikeOp, [name](Operation* op) {
+        op->emitError() << "contains removed PIM local-memory planning attribute '" << name << "'";
+      });
+      hasFailure = true;
+    }
+
+  auto arenaAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemorySizeAttrName);
+  if (!arenaAttr || arenaAttr.getInt() < 0
+      || static_cast<uint64_t>(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) {
+    diagnostics.report(coreLikeOp, [](Operation* op) {
+      op->emitError("requires one valid non-negative pim.local_memory_size attribute");
     });
     hasFailure = true;
   }
-  coreLikeOp->walk([&](memref::AllocOp allocOp) {
-    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) {
-      diagnostics.report(allocOp, [&](Operation*) {
-        allocOp.emitOpError("requires a complete non-negative PIM local-memory plan");
+  if (hasFailure)
+    return failure();
+  uint64_t arenaSize = static_cast<uint64_t>(arenaAttr.getInt());
+
+  auto analyzed = pim::analyzeLocalMemoryLifetimes(coreLikeOp);
+  if (failed(analyzed)) {
+    diagnostics.report(coreLikeOp, [](Operation* op) {
+      op->emitError("cannot analyze PIM local-memory lifetimes for plan verification");
+    });
+    return failure();
+  }
+
+  struct Event {
+    uint64_t position;
+    bool start;
+    size_t intervalIndex;
+  };
+  SmallVector<Event, 0> events;
+  SmallVector<size_t, 0> addresses(analyzed->size());
+  for (auto indexed : llvm::enumerate(*analyzed)) {
+    size_t index = indexed.index();
+    auto& interval = indexed.value();
+    memref::AllocOp allocation = interval.allocation;
+    for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
+      if (allocation->hasAttr(name)) {
+        diagnostics.report(allocation, [name](Operation* op) {
+          op->emitOpError() << "contains removed PIM local-memory planning attribute '" << name << "'";
+        });
+        hasFailure = true;
+      }
+    auto addressAttr = allocation->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
+    if (!addressAttr || addressAttr.getInt() < 0) {
+      diagnostics.report(allocation, [](Operation* op) {
+        op->emitOpError("requires one non-negative pim.local_memory_address attribute");
       });
       hasFailure = true;
-      return;
+      continue;
     }
-
     uint64_t address = static_cast<uint64_t>(addressAttr.getInt());
-    uint64_t slotId = static_cast<uint64_t>(slotAttr.getInt());
-    uint64_t slotSize = static_cast<uint64_t>(slotSizeAttr.getInt());
-    auto logicalSize = pim::getCheckedShapedTypeSizeInBytes(
-      cast<ShapedType>(allocOp.getType()), allocOp, "planned local allocation byte size");
-    bool invalidRange = failed(logicalSize) || address > slotSize || *logicalSize > slotSize - address
-                     || slotSize > kPimLocalMemoryAddressLimit || address % 4 != 0 || slotId != 0;
-    if (invalidRange) {
-      diagnostics.report(allocOp, [&](Operation*) {
-        allocOp.emitOpError() << "has an invalid PIM local-memory range: address=" << address
-                              << ", logical size=" << (succeeded(logicalSize) ? *logicalSize : 0)
-                              << ", slot size=" << slotSize;
+    if (address % 4 != 0 || address > arenaSize || interval.size > arenaSize - address) {
+      diagnostics.report(allocation, [&](Operation* op) {
+        op->emitOpError() << "has invalid PIM local-memory range [" << address << ", "
+                          << (address <= arenaSize && interval.size <= arenaSize - address
+                                ? address + interval.size
+                                : arenaSize)
+                          << ") for arena size " << arenaSize;
       });
       hasFailure = true;
-      return;
+      continue;
+    }
+    addresses[index] = static_cast<size_t>(address);
+    events.push_back({interval.start, true, index});
+    events.push_back({interval.end, false, index});
+  }
+  if (hasFailure)
+    return failure();
+
+  llvm::sort(events, [](const Event& lhs, const Event& rhs) {
+    return lhs.position != rhs.position ? lhs.position < rhs.position : lhs.start > rhs.start;
+  });
+  std::map<size_t, size_t> active;
+  for (const Event& event : events) {
+    size_t index = event.intervalIndex;
+    size_t address = addresses[index];
+    if (!event.start) {
+      auto it = active.find(address);
+      if (it != active.end() && it->second == index)
+        active.erase(it);
+      continue;
     }
 
-    auto [slotIt, inserted] = slots.try_emplace(slotId, VerifiedLocalMemorySlot {slotSize});
-    if (!inserted && slotIt->second.size != slotSize) {
-      diagnostics.report(allocOp, [&](Operation*) {
-        allocOp.emitOpError() << "has inconsistent size for PIM local-memory arena " << slotId;
+    const auto& interval = (*analyzed)[index];
+    auto successor = active.lower_bound(address);
+    auto conflict = [&](std::map<size_t, size_t>::iterator it) {
+      if (it == active.end())
+        return false;
+      const auto& other = (*analyzed)[it->second];
+      return address < it->first + other.size && it->first < address + interval.size;
+    };
+    auto predecessor = successor;
+    bool hasPredecessor = predecessor != active.begin();
+    if (hasPredecessor)
+      --predecessor;
+    auto conflicting = conflict(successor) ? successor
+                                           : hasPredecessor && conflict(predecessor) ? predecessor : active.end();
+    if (conflicting != active.end()) {
+      const auto& other = (*analyzed)[conflicting->second];
+      memref::AllocOp allocation = interval.allocation;
+      memref::AllocOp otherAllocation = other.allocation;
+      diagnostics.report(allocation, [&](Operation*) {
+        auto diagnostic = allocation.emitOpError()
+          << "PIM local-memory plan assigns simultaneously live allocations to overlapping ranges; first range ["
+          << conflicting->first << ", " << conflicting->first + other.size << "), second range [" << address
+          << ", " << address + interval.size << "), live positions overlap at ["
+          << std::max(interval.start, other.start) << ", " << std::min(interval.end, other.end) << "]";
+        diagnostic.attachNote(otherAllocation.getLoc()) << "first allocation";
       });
-      hasFailure = true;
+      return failure();
     }
-  });
-  return success(!hasFailure);
+    active.emplace(address, index);
+  }
+  return success();
 }
 
 static bool isSupportedCoreInstructionOp(Operation* op) {

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp

@@ -16,6 +16,11 @@ namespace spatial {
 
 namespace {
 
+// Pressure means distinct weights exceed half the fleet's one-copy capacity.
+// The reserved headroom conservatively avoids greedy capacity fragmentation;
+// makespan, communication, instruction-path, and local-memory proxies remain stable.
+constexpr size_t kHighCrossbarPressureCapacityDivisor = 2;
+
 struct ScheduledTask {
   size_t processor = std::numeric_limits<size_t>::max();
   Time startTime = 0;
@@ -136,7 +141,7 @@ void verifyOctTableSize(size_t nodeCount, size_t processorCount) {
 bool hasHighCrossbarPressure(const ComputeGraph& graph, size_t processorCount, size_t crossbarCapacity) {
   if (crossbarCapacity > std::numeric_limits<size_t>::max() / processorCount)
     return false;
-  const size_t threshold = processorCount * crossbarCapacity / 2;
+  const size_t threshold = processorCount * crossbarCapacity / kHighCrossbarPressureCapacityDivisor;
   CrossbarUsage distinctWeights;
   for (const ComputeGraphNode& node : graph.nodes) {
     for (const CrossbarWeight& weight : node.crossbarUsage) {

src/PIM/Pass/PIMPasses.h

@@ -15,8 +15,6 @@ std::unique_ptr<mlir::Pass> createSpatialToPimPass();
 
 std::unique_ptr<mlir::Pass> createPimBufferizationPass();
 
-std::unique_ptr<mlir::Pass> createPimMemoryCoalescingPass();
-
 std::unique_ptr<mlir::Pass> createMergeComputeNodesPass();
 
 std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass();

src/PIM/PimAccelerator.cpp

@@ -76,7 +76,6 @@ void PimAccelerator::registerPasses(int optLevel) const {
   registerPass(createLowerSpatialPlansPass);
   registerPass(createSpatialToPimPass);
   registerPass(createPimBufferizationPass);
-  registerPass(createPimMemoryCoalescingPass);
   registerPass(createTrivialGraphComputeMergePass);
   registerPass(createMergeComputeNodesPass);
   registerPass(createPimHostConstantFoldingPass);