unify memory opimization passes

better reports
cleanups

.agents/invariants/PARALLELISM_AND_RUNTIME_INVARIANT.md

@@ -1,20 +0,0 @@
-# Parallelism and Runtime Invariant
-
-Compile-time optimization must not trade away execution parallelism or worsen
-theoretical runtime behavior. The absence of a precise runtime measurement does
-not make such a trade acceptable.
-
-In particular, a compile-time optimization must not:
-
-- serialize work that can execute independently;
-- coarsen lane, core, batch, or scheduling granularity in a way that reduces
-  available execution parallelism;
-- increase the theoretical runtime critical path, instruction count, memory
-  traffic, synchronization, or required copies merely to reduce compiler work;
-- merge independently scheduled work when doing so may reduce concurrency.
-
-Compiler representations may share analysis, planning, or code-generation
-work only when the emitted execution retains the same independent work,
-schedule semantics, and available parallelism. If runtime impact cannot be
-measured precisely, it must be established from the transformation's semantics;
-uncertainty is not permission to accept a possible runtime regression.

.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md

@@ -11,7 +11,6 @@ This invariant applies to:
 - deferred communication;
 - Spatial-to-PIM lowering;
 - bufferization;
-- memory coalescing;
 - liveness planning;
 - PIM code generation.
 
@@ -19,12 +18,17 @@ This invariant applies to:
 
 A compiler compile-time or memory optimization must not reduce available
 hardware parallelism or worsen theoretical execution time. In the absence of a
-precise runtime model, it is forbidden to serialize independent work, reduce
+precise runtime model, uncertainty is not permission to accept a possible
+runtime regression. It is forbidden to serialize independent work, reduce
 the number of simultaneously schedulable compute instances, replace parallel
 graph or core batches with sequential loops, introduce recomputation, add
 runtime copies, increase communication or instruction count, or lengthen the
 schedule critical path merely to reduce compiler time or memory usage.
 
+Compiler representations may share analysis, planning, or code-generation
+work only when the emitted execution retains the same independent work,
+schedule semantics, scheduling granularity, and available parallelism.
+
 An optimization is acceptable only when its static runtime proxies are equal or
 better.
 
@@ -33,6 +37,8 @@ better.
 Do not:
 
 - use fewer parallel lanes or cores to simplify IR;
+- coarsen lane, core, batch, or scheduling granularity when it may reduce
+  concurrency;
 - scalarize a valid graph or core batch;
 - replace independent operations with one sequential loop;
 - recompute values to reduce storage;

AGENTS.md

@@ -14,6 +14,7 @@ Before modifying the relevant subsystem, read:
 * Use the debug build only when it is useful to obtain a clear stack trace with symbols, inspect names, place breakpoints, or test a small case interactively
 * The debug build is very slow, so use it only on small fast tests such as operation validations, not on network validations
 * Always prepend rtk to shell commands if missing and if rtk is available
+* Always use the repository Python virtual environment (`.venv/bin/python` and its bundled tools) for Python commands
 
 # Core engineering philosophy
 

README.md

@@ -67,14 +67,11 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    Converts tensor-semantics PIM IR into memref-semantics PIM IR using MLIR's
    bufferization interfaces.
 
-5. **Safe IR memory coalescing**
+5. **PIM local-memory planning**
-   (`src/PIM/Dialect/Pim/Transforms/MemoryCoalescing`).
-   Normalizes compatible block-local allocations before final planning.
-6. **PIM local-memory planning**
    (`src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning`).
-   Computes whole-core lifetimes, coalesces non-overlapping allocations into
+   Computes whole-core lifetimes, reuses addresses for non-overlapping
-   physical slots, assigns addresses, and records the explicit plan in PIM IR.
+   allocations, and records the explicit plan in PIM IR.
-7. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and
+6. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and
    `src/PIM/Compiler`).
    Verifies the memory plan and other PIM invariants, then emits `.pim` core
    files, weights, and `memory.bin` / `config.json` without rerunning liveness.
@@ -90,9 +87,10 @@ Supporting pieces:
 - `src/PIM/Pass` - pass registration and auxiliary passes.
 - `src/PIM/PimAccelerator.{cpp,hpp}` - ONNX-MLIR accelerator entry point.
 
-## Key compiler options
+## PIM compiler options
 
-Pass these to `onnx-mlir` when compiling for PIM:
+Pass these to `onnx-mlir` when compiling for PIM. These are all Raptor/PIM-specific
+options; `onnx-mlir --help` lists the inherited ONNX-MLIR options.
 
 - `--maccel=PIM` - select the PIM accelerator.
 - `--EmitSpatial`, `--EmitPim`, `--EmitPimBufferized`,
@@ -101,14 +99,30 @@ Pass these to `onnx-mlir` when compiling for PIM:
 - `--core-count=<N>` - required positive core count for PIM compilation.
 - `--crossbar-size=<N>` - crossbar width/height. Default in code is `128`.
 - `--crossbar-count=<N>` - crossbars per core. Default in code is `64`.
+- `--pim-memory-report=<summary|none>` - emit the concise combined memory report
+  under `reports/memory_report.txt`, or disable it. Default is `summary`.
 - `--pim-only-codegen` - assume input is already bufferized PIM IR and only run
   the codegen tail.
 - `--pim-emit-json` - also emit `core_*.json` instruction files alongside
   `core_*.pim`.
-- `--pim-export-spatial-dataflow=<none|spatial1|spatial2|spatial3|spatial4|all>` - control Spatial
+- `--pim-export-spatial-dataflow=<none|spatial1|spatial2|spatial3|spatial4|all>` -
-  dataflow CSV reports for the graph, trivially merged graph, scheduled, and
+  control Spatial dataflow CSV reports for the graph, trivially merged graph,
-  realized snapshots under `reports/`.
+  scheduled, and realized snapshots under `reports/`. Default is `none`.
+- `--pim-conv-lowering=<auto|legacy|depthwise|packed-im2col|streamed-patch|streamed-packed|output-channel-tiled|input-k-tiled|tiled-2d>` -
+  select the convolution lowering strategy. Default is `auto`.
+- `--pim-conv-im2col-max-elements=<N>` - maximum globally materialized im2col
+  elements per convolution before streaming. Default is `1048576`.
+- `--pim-conv-stream-chunk-positions=<N>` - maximum output positions per
+  streamed convolution chunk. Default is `1024`.
+- `--pim-report-conv-lowering=<true|false>` - emit the bounded convolution
+  lowering report. Default is `true`.
 - `--use-experimental-conv-impl` - use the alternate convolution lowering.
+- `--pim-detect-communication-deadlock` - statically simulate expanded
+  send/receive ordering and reject blocking deadlocks. Default is off.
+- `--pim-materialize-scalar-fanout-global-order` - use the experimental,
+  expensive globally ordered scalar-fanout materializer. Default is off.
+- `--pim-trace-communication-materialization` - emit verbose communication
+  materialization diagnostics and provenance attributes. Default is off.
 - `--ignore-concat-error` - soft-fail a ConcatOp corner case.
 
 ## Standard PIM hardware profile
@@ -149,7 +163,7 @@ Python dependencies used by the validation scripts are `numpy`, `onnx`, and
 Per-operation validation from the repository root:
 
 ```bash
-python validation/validate.py \
+.venv/bin/python validation/validate.py \
   --raptor-path build_release/Release/bin/onnx-mlir \
   --onnx-include-dir onnx-mlir/include \
   --operations-dir validation/operations \
@@ -165,7 +179,7 @@ Validate one network or a subset by pointing `--operations-dir` at any directory
 containing `.onnx` files:
 
 ```bash
-python validation/validate.py \
+.venv/bin/python validation/validate.py \
   --raptor-path build_release/Release/bin/onnx-mlir \
   --onnx-include-dir onnx-mlir/include \
   --operations-dir validation/networks/yolo11n/depth_04 \
@@ -173,7 +187,6 @@ python validation/validate.py \
   --crossbar-size 128 \
   --core-count 144 \
   --verbose \
-  --raptor-extra-arg=--pim-memory-report=summary \
   --raptor-extra-arg=--pim-detect-communication-deadlock \
   --raptor-extra-arg=--pim-export-spatial-dataflow=none
 ```
@@ -200,7 +213,7 @@ Each validation run writes artifacts in the model workspace, for example under
 The compiler currently dumps dialect snapshots such as `spatial0.mlir`,
 `spatial1_graph.mlir`, `spatial2_trivial_merged.mlir`,
 `spatial3_scheduled_no_comm.mlir`, `spatial4_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
-`pim2_folded.mlir`, `pim3_coalesced.mlir`, and `pim4_memory_planned.mlir` when an output directory is
+`pim2_folded.mlir`, and `pim3_memory_planned.mlir` when an output directory is
 available.
 
 To rerun the simulator manually with tracing after validation has produced a
@@ -229,7 +242,7 @@ Available operation suites under `validation/operations/`: `add`, `concat`,
 Generated operation tests can be regenerated with:
 
 ```bash
-python3 validation/operations/gen_tests.py
+.venv/bin/python validation/operations/gen_tests.py
 ```
 
 ## Build

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 kLocalMemorySizeAttrName = "pim.local_memory_size";
-inline constexpr llvm::StringLiteral kLocalMemorySlotSizeAttrName = "pim.local_memory_slot_size";
+inline constexpr std::array<llvm::StringLiteral, 4> kRemovedLocalMemoryPlanAttrNames = {
-inline constexpr llvm::StringLiteral kLocalMemoryFallbackCountAttrName = "pim.local_memory_fallback_count";
+  "pim.local_memory_slot",
-inline constexpr llvm::StringLiteral kLocalMemoryNestedSingleUseCountAttrName =
+  "pim.local_memory_slot_size",
-  "pim.local_memory_nested_single_use_count";
+  "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()) {
+  if (!localArenaSize) {
-    llvm::append_range(localPhysicalSlots, plan.slots);
+    localArenaSize = plan.arenaSize;
-    reportRow.logicalAllocaBytes = plan.logicalBytes;
+    reportRow.logicalLocalAllocationCount = plan.logicalAllocationCount;
-    reportRow.fallbackIntervals = plan.fallbackIntervals;
+    reportRow.logicalLocalBytes = plan.logicalBytes;
-    reportRow.nestedSingleUseIntervals = plan.nestedSingleUseIntervals;
   }
+  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.hostBytes = firstAvailableAddress;
-  row.sizeAlloca = 0;
+  row.physicalLocalBytes = localArenaSize.value_or(0);
-  for (const PhysicalSlotInfo& slot : localPhysicalSlots)
-    row.sizeAlloca += slot.size;
   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});
-                           row.numAlloca,
-                           row.sizeAlloca});
 }
 
 void PimAcceleratorMemory::recordBatchReport(uint64_t batchId,
                                              ArrayRef<int32_t> coreIds,
-                                             const MemoryReportRow& perCoreRow,
+                                             const MemoryReportRow& perCoreRow) {
-                                             uint64_t totalAllocaCount,
-                                             uint64_t totalAllocaBytes) {
   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;
+  struct ReportGroup {
-  uint64_t totalWeightsMemory = totalWeightBytes;
+    MemoryReportRow row;
-  uint64_t totalCoresMemory = 0;
+    SmallVector<int32_t, 8> coreIds;
-  uint64_t totalLogicalCoreMemory = 0;
+    std::optional<uint64_t> batchId;
-  uint64_t totalFallbackIntervals = 0;
+  };
-  uint64_t totalNestedSingleUseIntervals = 0;
+  SmallVector<ReportGroup, 32> groups;
-  uint64_t maximumCorePeak = 0;
+  for (const MemoryReportEntry& entry : reportEntries) {
-  std::string worstEntry = "<none>";
+    auto group = llvm::find_if(groups, [&](const ReportGroup& candidate) {
-  for (const MemoryReportEntry& entry : reportEntries) {
+      return candidate.row.logicalLocalBytes == entry.row.logicalLocalBytes
-    totalCoresMemory += entry.totalAllocaBytes;
+          && candidate.row.physicalLocalBytes == entry.row.physicalLocalBytes;
-    uint64_t factor = entry.kind == MemoryReportEntry::Kind::Batch ? entry.coreIds.size() : 1;
+    });
-    totalLogicalCoreMemory += entry.row.logicalAllocaBytes * factor;
+    if (group == groups.end()) {
-    totalFallbackIntervals += entry.row.fallbackIntervals * factor;
+      ReportGroup newGroup;
-    totalNestedSingleUseIntervals += entry.row.nestedSingleUseIntervals * factor;
+      newGroup.row = entry.row;
-    if (entry.row.sizeAlloca <= maximumCorePeak)
+      llvm::append_range(newGroup.coreIds, entry.coreIds);
-      continue;
+      if (entry.kind == MemoryReportEntry::Kind::Batch)
-    maximumCorePeak = entry.row.sizeAlloca;
+        newGroup.batchId = entry.id;
-    worstEntry = entry.kind == MemoryReportEntry::Kind::Batch ? "Batch " + std::to_string(entry.id)
+      groups.push_back(std::move(newGroup));
-                                                              : "Core " + std::to_string(entry.coreIds.front());
+      continue;
-  }
+    }
-  uint64_t savedCoreMemory =
+    llvm::append_range(group->coreIds, entry.coreIds);
-    totalLogicalCoreMemory >= totalCoresMemory ? totalLogicalCoreMemory - totalCoresMemory : 0;
+    group->batchId.reset();
-  double savedPercent = totalLogicalCoreMemory == 0
+  }
-                        ? 0.0
+  llvm::sort(groups, [](const ReportGroup& lhs, const ReportGroup& rhs) {
-                        : 100.0 * static_cast<double>(savedCoreMemory) / static_cast<double>(totalLogicalCoreMemory);
+    if (lhs.row.physicalLocalBytes != rhs.row.physicalLocalBytes)
-
+      return lhs.row.physicalLocalBytes > rhs.row.physicalLocalBytes;
-  llvm::SmallVector<ReportField, 10> totalFields = {
+    return lhs.row.logicalLocalBytes > rhs.row.logicalLocalBytes;
-    {"Global memory",               formatReportMemory(totalGlobalMemory)               },
+  });
-    {"Weights memory",              formatReportMemory(totalWeightsMemory)              },
+  uint64_t logicalBytes = 0;
-    {"Logical core allocations",    formatReportMemory(totalLogicalCoreMemory)           },
+  uint64_t physicalBytes = 0;
-    {"Physical core memory",        formatReportMemory(totalCoresMemory)                 },
+  for (const ReportGroup& group : groups) {
-    {"Cores memory",                formatReportMemory(totalCoresMemory)                 },
+    logicalBytes += group.row.logicalLocalBytes * group.coreIds.size();
-    {"Saved core memory",           formatReportMemory(savedCoreMemory)                  },
+    physicalBytes += group.row.physicalLocalBytes * group.coreIds.size();
-    {"Saved core memory percent",   formatv("{0:F2}%", savedPercent).str()               },
+  }
-    {"Maximum core peak",           formatReportMemory(maximumCorePeak)                  },
+  uint64_t savedBytes = logicalBytes - physicalBytes;
-    {"Worst core/batch",            worstEntry                                           },
+  double savedPercent = logicalBytes == 0 ? 0.0 : 100.0 * savedBytes / logicalBytes;
-    {"Fallback intervals",          std::to_string(totalFallbackIntervals)               },
+  uint64_t largest = groups.empty() ? 0 : groups.front().row.physicalLocalBytes;
-    {"Nested single-use intervals", std::to_string(totalNestedSingleUseIntervals)        }
+  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 << "Summary\n";
-    os << "\nHost:\n";
+  os << "  Host memory: " << formatReportMemory(hostReportRow ? hostReportRow->hostBytes : 0) << "\n";
-    printHostMemoryReportRow(os, *hostReportRow);
+  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()) {
+  os << "\nLargest core groups\n";
-    if (hostReportRow.has_value())
+  constexpr size_t kGroupLimit = 8;
-      os << "\n";
+  for (const ReportGroup& group : ArrayRef(groups).take_front(kGroupLimit)) {
-    sortReportEntriesByFirstCore(reportEntries);
+    printLabel(group);
-
+    os << "\n";
-    for (size_t index = 0; index < reportEntries.size();) {
+    uint64_t groupSaved = group.row.logicalLocalBytes - group.row.physicalLocalBytes;
-      size_t runEnd = index + 1;
+    double groupPercent = group.row.logicalLocalBytes == 0
-      while (runEnd < reportEntries.size() && reportEntries[runEnd].kind == reportEntries[index].kind
+                            ? 0.0
-             && reportEntries[runEnd].row == reportEntries[index].row
+                            : 100.0 * groupSaved / group.row.logicalLocalBytes;
-             && reportEntries[runEnd].totalAllocaCount == reportEntries[index].totalAllocaCount
+    if (group.coreIds.size() == 1) {
-             && reportEntries[runEnd].totalAllocaBytes == reportEntries[index].totalAllocaBytes) {
+      os << "  Local memory: " << formatReportMemory(group.row.logicalLocalBytes) << " → "
-        ++runEnd;
+         << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent)
-      }
+         << ")\n";
-
+    }
-      if (reportEntries[index].kind == MemoryReportEntry::Kind::Batch) {
+    else {
-        os << "Batch ";
+      os << "  Per core: " << formatReportMemory(group.row.logicalLocalBytes) << " → "
-        for (size_t batchIndex = index; batchIndex < runEnd; ++batchIndex) {
+         << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent)
-          if (batchIndex != index)
+         << ")\n";
-            os << ",\n      ";
+      os << "  Total after reuse: " << formatReportMemory(group.row.physicalLocalBytes * group.coreIds.size())
-          os << reportEntries[batchIndex].id << " (cores ";
+         << "\n";
-          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;
     }
   }
+  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 arenaAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemorySizeAttrName);
-  auto nestedSingleUseAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryNestedSingleUseCountAttrName);
+  if (!arenaAttr || arenaAttr.getInt() < 0
-  if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) {
+      || static_cast<uint64_t>(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) {
-    coreLikeOp->emitError("requires complete PIM local-memory planning summary attributes before codegen");
+    coreLikeOp->emitError("requires a valid pim.local_memory_size attribute before codegen");
     return failure();
   }
-  plan.fallbackIntervals = static_cast<uint64_t>(fallbackAttr.getInt());
+  plan.arenaSize = static_cast<size_t>(arenaAttr.getInt());
-  plan.nestedSingleUseIntervals = static_cast<uint64_t>(nestedSingleUseAttr.getInt());
-  SmallDenseMap<size_t, size_t, 32> slotIndices;
   bool hasFailure = false;
   coreLikeOp->walk([&](memref::AllocOp allocOp) {
     if (hasFailure)
       return;
     auto addressAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
-    auto slotAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotAttrName);
+    if (!addressAttr || addressAttr.getInt() < 0) {
-    auto slotSizeAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotSizeAttrName);
+      allocOp.emitOpError("requires a non-negative pim.local_memory_address attribute before codegen");
-    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");
       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());
+    if (address > plan.arenaSize || *checkedSize > plan.arenaSize - address) {
-    size_t slotSize = static_cast<size_t>(slotSizeAttr.getInt());
+      allocOp.emitOpError("has a planned local-memory range outside its core arena");
-    plan.entries.push_back({allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}});
+      hasFailure = true;
-    plan.logicalBytes += *checkedSize;
-
-    auto [slotIt, inserted] = slotIndices.try_emplace(slotId, plan.slots.size());
-    if (inserted) {
-      plan.slots.push_back({slotId, 0, slotSize});
       return;
     }
-    const PhysicalSlotInfo& existing = plan.slots[slotIt->second];
+    plan.entries.push_back({allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}});
-    if (existing.size != slotSize) {
+    auto logicalBytes = pim::checkedAdd(
-      allocOp.emitOpError("has a PIM local-memory arena inconsistent with another allocation");
+      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,
+    memory.recordBatchReport(
-                             reportedCoreIds,
+      batchReportId, reportedCoreIds, batchPerCoreRow.value_or(MemoryReportRow {}));
-                             batchPerCoreRow.value_or(MemoryReportRow {}),
-                             batchRow.numAlloca,
-                             batchRow.sizeAlloca);
   }
 
   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;
+  size_t arenaSize = 0;
-  uint64_t nestedSingleUseIntervals = 0;
 };
 
 struct MemoryReportRow {
-  uint64_t numAlloca = 0;
+  uint64_t hostObjectCount = 0;
-  uint64_t sizeAlloca = 0;
+  uint64_t hostBytes = 0;
-  uint64_t numGlobal = 0;
+  uint64_t logicalLocalAllocationCount = 0;
-  uint64_t sizeGlobal = 0;
+  uint64_t logicalLocalBytes = 0;
-  uint64_t logicalAllocaBytes = 0;
+  uint64_t physicalLocalBytes = 0;
-  uint64_t fallbackIntervals = 0;
-  uint64_t nestedSingleUseIntervals = 0;
 
   bool operator==(const MemoryReportRow& other) const {
-    return numAlloca == other.numAlloca && sizeAlloca == other.sizeAlloca && numGlobal == other.numGlobal
+    return hostObjectCount == other.hostObjectCount && hostBytes == other.hostBytes
-        && sizeGlobal == other.sizeGlobal && logicalAllocaBytes == other.logicalAllocaBytes
+        && logicalLocalAllocationCount == other.logicalLocalAllocationCount
-        && fallbackIntervals == other.fallbackIntervals
+        && logicalLocalBytes == other.logicalLocalBytes && physicalLocalBytes == other.physicalLocalBytes;
-        && nestedSingleUseIntervals == other.nestedSingleUseIntervals;
   }
 };
 
@@ -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()
+  PimAcceleratorMemory();
-  : hostMem(memEntriesMap), fileReport(openReportFile("memory_report")) {}
+  PimAcceleratorMemory(
-  PimAcceleratorMemory(const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport)
+    const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport);
-  : memEntriesMap(initialMemEntries),
-    hostMem(memEntriesMap),
-    fileReport(enableReport ? openReportFile("memory_report") : std::fstream()) {}
 
   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,
+  void recordBatchReport(
-                         llvm::ArrayRef<int32_t> coreIds,
+    uint64_t batchId, llvm::ArrayRef<int32_t> coreIds, const MemoryReportRow& perCoreRow);
-                         const MemoryReportRow& perCoreRow,
-                         uint64_t totalAllocaCount,
-                         uint64_t totalAllocaBytes);
   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(
+  llvm::cl::values(clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise PIM memory summary")),
-    clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise slot reuse report with key offenders")),
+  llvm::cl::init(PimMemoryReportSummary),
-  llvm::cl::values(clEnumValN(PimMemoryReportFull, "full", "Emit the full detailed PIM memory planning report")),
-  llvm::cl::init(PimMemoryReportNone),
   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 onnx_mlir::pim {
-namespace 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,
+LogicalResult walkLocalMemoryUses(Value root, LocalMemoryUseCallback callback) {
-                                  llvm::function_ref<LogicalResult(Value, Operation*)> visitUser) {
   llvm::SmallPtrSet<Value, 16> visitedValues;
-  llvm::SmallPtrSet<Operation*, 32> visitedUsers;
+  llvm::SmallPtrSet<OpOperand*, 32> visitedUses;
-  llvm::SmallVector<Value> pendingValues {root};
+  SmallVector<Value> pending = {root};
-  auto addAlias = [&](Value value) { pendingValues.push_back(value); };
+  while (!pending.empty()) {
-
+    Value value = pending.pop_back_val();
-  while (!pendingValues.empty()) {
-    Value value = pendingValues.pop_back_val();
     if (!visitedValues.insert(value).second)
       continue;
-    for (Operation* user : value.getUsers()) {
+    for (OpOperand& use : value.getUses()) {
-      if (!visitedUsers.insert(user).second)
+      if (!visitedUses.insert(&use).second || failed(callback(use)))
-        continue;
-      if (failed(visitUser(value, user)))
         return failure();
+      Operation* user = use.getOwner();
 
-      if (isLocalMemoryAliasOp(user))
+      if (isLocalMemoryAliasOp(user) && use.getOperandNumber() == 0) {
-        for (Value result : user->getResults())
+        if (user->getNumResults() != 1)
-          addAlias(result);
+          return failure();
+        pending.push_back(user->getResult(0));
+      }
 
-      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user))
+      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user); dpsOp && dpsOp.isDpsInit(&use))
-        for (OpResult result : user->getResults())
+        pending.push_back(dpsOp.getTiedOpResult(&use));
-          if (OpOperand* tied = dpsOp.getTiedOpOperand(result); tied && tied->get() == value)
-            addAlias(result);
 
       if (auto forOp = dyn_cast<scf::ForOp>(user))
-        for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs()))
+        for (auto [index, init] : llvm::enumerate(forOp.getInitArgsMutable()))
-          if (initArg == value) {
+          if (&init == &use) {
-            addAlias(forOp.getRegionIterArgs()[index]);
+            pending.push_back(forOp.getRegionIterArgs()[index]);
-            addAlias(forOp.getResult(index));
+            pending.push_back(forOp.getResult(index));
           }
 
-      auto yieldOp = dyn_cast<scf::YieldOp>(user);
+      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
-      if (!yieldOp)
+        unsigned index = use.getOperandNumber();
-        continue;
+        Operation* parent = yieldOp->getParentOp();
-      for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
+        if (auto forOp = dyn_cast<scf::ForOp>(parent))
-        if (operand != value)
+          pending.push_back(forOp.getResult(index));
-          continue;
+        else if (auto ifOp = dyn_cast<scf::IfOp>(parent))
-        if (auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp()))
+          pending.push_back(ifOp.getResult(index));
-          addAlias(forOp.getResult(index));
+        else if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(parent))
-        else if (auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp()))
+          pending.push_back(switchOp.getResult(index));
-          addAlias(ifOp.getResult(index));
+        else
-        else if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp()))
+          return failure();
-          addAlias(switchOp.getResult(index));
       }
     }
   }
   return success();
 }
 
-} // namespace pim
+FailureOr<SmallVector<LocalMemoryInterval, 0>> analyzeLocalMemoryLifetimes(Operation* coreLikeOp) {
-} // namespace onnx_mlir
+  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 onnx_mlir::pim {
-namespace 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(
+} // namespace onnx_mlir::pim
-  mlir::Value root,
-  llvm::function_ref<mlir::LogicalResult(mlir::Value, mlir::Operation*)> visitUser);
-
-} // namespace pim
-} // namespace onnx_mlir

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 hasFailure = false;
-    bool failedPlanning = 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);
+      op->removeAttr(kLocalMemorySizeAttrName);
-      if (failed(assignPhysicalSlotAddresses(plan, kPimLocalMemoryAddressLimit))) {
+      for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
-        failedPlanning = true;
+        op->removeAttr(name);
-        return;
+      op->walk([&](memref::AllocOp allocation) {
-      }
+        allocation->removeAttr(kLocalMemoryAddressAttrName);
-
+        for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
-      size_t arenaSize = 0;
+          allocation->removeAttr(name);
-      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;
       });
-      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)
+      auto plan = buildCoreMemoryPlan(op);
-        return;
+      if (failed(plan)) {
-      std::string planReport = buildMemoryPlanReport(
+        hasFailure = true;
-        op, plan.intervals, plan.slots, kPimLocalMemoryAddressLimit, pimMemoryReport);
-      if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
-        *report << "Core " << coreOp.getCoreId() << ":\n";
-        *report << planReport;
         return;
       }
-
+      op->setAttr(kLocalMemorySizeAttrName, builder.getI64IntegerAttr(plan->arenaSize));
-      SmallVector<int32_t> coreIds = getBatchCoreIds(cast<pim::PimCoreBatchOp>(op));
+      for (const LocalMemoryPlacement& placement : plan->placements)
-      llvm::sort(coreIds);
+        for (size_t intervalIndex : placement.intervalIndices)
-      coreIds.erase(std::unique(coreIds.begin(), coreIds.end()), coreIds.end());
+          plan->intervals[intervalIndex].allocation->setAttr(
-      uint64_t batchId = nextBatchId++;
+            kLocalMemoryAddressAttrName, builder.getI64IntegerAttr(placement.address));
-      for (int32_t coreId : coreIds)
-        *report << "Batch " << batchId << " core " << coreId << ":\n" << planReport;
     });
-
+    if (hasFailure) {
-    if (report) {
-      report->flush();
-      reportFile.close();
-    }
-    if (failedPlanning) {
       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 "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <limits>
-#include <string>
-
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
 namespace onnx_mlir {
 
-struct LocalAllocInterval {
+struct LocalMemoryPlacement {
-  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;
   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<pim::LocalMemoryInterval, 0> intervals;
-  llvm::SmallVector<PlannedPhysicalSlot, 0> slots;
+  llvm::SmallVector<LocalMemoryPlacement, 0> placements;
+  size_t arenaSize = 0;
 };
 
-llvm::SmallVector<LocalAllocInterval, 0> buildLocalAllocIntervals(mlir::Operation* coreLikeOp);
+mlir::FailureOr<llvm::SmallVector<LocalMemoryPlacement, 0>>
-
+planLocalMemoryPlacements(llvm::ArrayRef<pim::LocalMemoryInterval> intervals, size_t addressLimit);
-llvm::SmallVector<PlannedPhysicalSlot, 0> planPhysicalSlots(llvm::MutableArrayRef<LocalAllocInterval> intervals);
+mlir::FailureOr<CoreMemoryPlan> buildCoreMemoryPlan(mlir::Operation* coreLikeOp);
-
-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);
 
 } // 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);
+  for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
-  auto nestedSingleUseAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryNestedSingleUseCountAttrName);
+    if (coreLikeOp->hasAttr(name)) {
-  if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) {
+      diagnostics.report(coreLikeOp, [name](Operation* op) {
-    diagnostics.report(coreLikeOp, [&](Operation* op) {
+        op->emitError() << "contains removed PIM local-memory planning attribute '" << name << "'";
-      op->emitError("requires complete non-negative PIM local-memory planning summary attributes");
+      });
+      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) {
+  if (hasFailure)
-    auto addressAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
+    return failure();
-    auto slotAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotAttrName);
+  uint64_t arenaSize = static_cast<uint64_t>(arenaAttr.getInt());
-    auto slotSizeAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotSizeAttrName);
+
-    if (!addressAttr || !slotAttr || !slotSizeAttr || addressAttr.getInt() < 0 || slotAttr.getInt() < 0
+  auto analyzed = pim::analyzeLocalMemoryLifetimes(coreLikeOp);
-        || slotSizeAttr.getInt() < 0) {
+  if (failed(analyzed)) {
-      diagnostics.report(allocOp, [&](Operation*) {
+    diagnostics.report(coreLikeOp, [](Operation* op) {
-        allocOp.emitOpError("requires a complete non-negative PIM local-memory plan");
+      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());
+    if (address % 4 != 0 || address > arenaSize || interval.size > arenaSize - address) {
-    uint64_t slotSize = static_cast<uint64_t>(slotSizeAttr.getInt());
+      diagnostics.report(allocation, [&](Operation* op) {
-    auto logicalSize = pim::getCheckedShapedTypeSizeInBytes(
+        op->emitOpError() << "has invalid PIM local-memory range [" << address << ", "
-      cast<ShapedType>(allocOp.getType()), allocOp, "planned local allocation byte size");
+                          << (address <= arenaSize && interval.size <= arenaSize - address
-    bool invalidRange = failed(logicalSize) || address > slotSize || *logicalSize > slotSize - address
+                                ? address + interval.size
-                     || slotSize > kPimLocalMemoryAddressLimit || address % 4 != 0 || slotId != 0;
+                                : arenaSize)
-    if (invalidRange) {
+                          << ") for arena size " << arenaSize;
-      diagnostics.report(allocOp, [&](Operation*) {
-        allocOp.emitOpError() << "has an invalid PIM local-memory range: address=" << address
-                              << ", logical size=" << (succeeded(logicalSize) ? *logicalSize : 0)
-                              << ", slot size=" << slotSize;
       });
       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});
+    const auto& interval = (*analyzed)[index];
-    if (!inserted && slotIt->second.size != slotSize) {
+    auto successor = active.lower_bound(address);
-      diagnostics.report(allocOp, [&](Operation*) {
+    auto conflict = [&](std::map<size_t, size_t>::iterator it) {
-        allocOp.emitOpError() << "has inconsistent size for PIM local-memory arena " << slotId;
+      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();
     }
-  });
+    active.emplace(address, index);
-  return success(!hasFailure);
+  }
+  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);

test/PIM/PimMemoryLivenessPlannerTest.cpp

@@ -4,103 +4,97 @@
 
 #include "src/Accelerators/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp"
 
-using onnx_mlir::LocalAllocInterval;
+using onnx_mlir::LocalMemoryPlacement;
-using onnx_mlir::CoreMemoryPlan;
+using onnx_mlir::pim::LocalMemoryInterval;
-using onnx_mlir::assignPhysicalSlotAddresses;
+using onnx_mlir::planLocalMemoryPlacements;
-using onnx_mlir::planPhysicalSlots;
 
 namespace {
 
-LocalAllocInterval makeInterval(size_t id, size_t size, uint64_t start, uint64_t end) {
+LocalMemoryInterval makeInterval(size_t size, uint64_t start, uint64_t end) {
-  LocalAllocInterval interval;
+  LocalMemoryInterval interval;
-  interval.id = id;
   interval.size = size;
   interval.start = start;
   interval.end = end;
   return interval;
 }
 
-size_t getPeak(llvm::ArrayRef<onnx_mlir::PlannedPhysicalSlot> slots) {
+size_t getPeak(llvm::ArrayRef<LocalMemoryPlacement> placements) {
   size_t peak = 0;
-  for (const auto& slot : slots)
+  for (const auto& placement : placements)
-    peak = std::max(peak, slot.address + slot.requiredSize);
+    peak = std::max(peak, placement.address + placement.size);
   return peak;
 }
 
-void assertSingleSlotCase(LocalAllocInterval a, LocalAllocInterval b, size_t expectedSlotSize) {
+void assertSinglePlacementCase(LocalMemoryInterval a, LocalMemoryInterval b, size_t expectedSize) {
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {a, b};
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {a, b};
-  auto slots = planPhysicalSlots(intervals);
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
-  assert(slots.size() == 1);
+  assert(mlir::succeeded(placements));
-  assert(slots.front().requiredSize == expectedSlotSize);
+  assert(placements->size() == 1);
-  assert(intervals[0].slotPlanIndex == intervals[1].slotPlanIndex);
+  assert(placements->front().size == expectedSize);
+  assert(placements->front().intervalIndices.size() == 2);
 }
 
 int testSameSizeNonOverlap() {
   std::cout << "testSameSizeNonOverlap:" << std::endl;
-  assertSingleSlotCase(makeInterval(0, 64, 0, 10), makeInterval(1, 64, 11, 20), 64);
+  assertSinglePlacementCase(makeInterval(64, 0, 10), makeInterval(64, 11, 20), 64);
   return 0;
 }
 
 int testLargerFirst() {
   std::cout << "testLargerFirst:" << std::endl;
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
-    makeInterval(0, 100, 0, 10), makeInterval(1, 40, 11, 20)};
+    makeInterval(100, 0, 10), makeInterval(40, 11, 20)};
-  auto slots = planPhysicalSlots(intervals);
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
-  assert(slots.size() == 2);
+  assert(mlir::succeeded(placements) && placements->size() == 2);
-  assert(getPeak(slots) == 100);
+  assert(getPeak(*placements) == 100);
   return 0;
 }
 
 int testSmallerFirst() {
   std::cout << "testSmallerFirst:" << std::endl;
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
-    makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20)};
+    makeInterval(40, 0, 10), makeInterval(100, 11, 20)};
-  auto slots = planPhysicalSlots(intervals);
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
-  assert(slots.size() == 2);
+  assert(mlir::succeeded(placements) && placements->size() == 2);
-  assert(getPeak(slots) == 100);
+  assert(getPeak(*placements) == 100);
   return 0;
 }
 
 int testOverlapNeedsTwoSlots() {
   std::cout << "testOverlapNeedsTwoSlots:" << std::endl;
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
-    makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)};
+    makeInterval(100, 0, 20), makeInterval(40, 10, 30)};
-  auto slots = planPhysicalSlots(intervals);
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
-  assert(slots.size() == 2);
+  assert(mlir::succeeded(placements) && placements->size() == 2);
-  assert(intervals[0].slotPlanIndex != intervals[1].slotPlanIndex);
+  assert((*placements)[0].address != (*placements)[1].address);
   return 0;
 }
 
 int testReuseChain() {
   std::cout << "testReuseChain:" << std::endl;
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
-    makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20), makeInterval(2, 20, 21, 30)};
+    makeInterval(40, 0, 10), makeInterval(100, 11, 20), makeInterval(20, 21, 30)};
-  auto slots = planPhysicalSlots(intervals);
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
-  assert(slots.size() == 3);
+  assert(mlir::succeeded(placements) && placements->size() == 3);
-  assert(getPeak(slots) == 100);
+  assert(getPeak(*placements) == 100);
   return 0;
 }
 
 int testPartialAddressReuse() {
   std::cout << "testPartialAddressReuse:" << std::endl;
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
-    makeInterval(0, 100, 0, 10), makeInterval(1, 60, 11, 20), makeInterval(2, 40, 11, 20)};
+    makeInterval(100, 0, 10), makeInterval(60, 11, 20), makeInterval(40, 11, 20)};
-  auto slots = planPhysicalSlots(intervals);
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
-  assert(getPeak(slots) == 100);
+  assert(mlir::succeeded(placements));
+  assert(getPeak(*placements) == 100);
   return 0;
 }
 
-int testAddressAssignment() {
+int testAddressLimit() {
-  std::cout << "testAddressAssignment:" << std::endl;
+  std::cout << "testAddressLimit:" << std::endl;
-  CoreMemoryPlan plan;
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
-  plan.intervals = {makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)};
+    makeInterval(100, 0, 20), makeInterval(40, 10, 30)};
-  plan.slots = planPhysicalSlots(plan.intervals);
+  assert(mlir::failed(planLocalMemoryPlacements(intervals, 128)));
-  assert(mlir::succeeded(assignPhysicalSlotAddresses(plan, 1024)));
-  assert(plan.slots.size() == 2);
-  const auto& firstSlot = plan.slots[plan.intervals[0].slotPlanIndex];
-  const auto& secondSlot = plan.slots[plan.intervals[1].slotPlanIndex];
-  assert(firstSlot.address == 0 && firstSlot.requiredSize == 100);
-  assert(secondSlot.address == 100 && secondSlot.requiredSize == 40);
   return 0;
 }
 
@@ -117,7 +111,7 @@ int main(int argc, char *argv[]) {
   failures += testOverlapNeedsTwoSlots();
   failures += testReuseChain();
   failures += testPartialAddressReuse();
-  failures += testAddressAssignment();
+  failures += testAddressLimit();
   if (failures != 0) {
     std::cerr << failures << " test failures\n";
     return EXIT_FAILURE;

validation/raptor.py

@@ -11,7 +11,6 @@ PIM_PASS_LABELS = (
     ("SpatialToPimPass", "Spatial to PIM"),
     ("PimBufferizationPass", "Bufferize PIM"),
     ("HostConstantFoldingPass", "Fold Host Constants"),
-    ("PimMemoryCoalescingPass", "Coalesce Local Memory"),
     ("PimLocalMemoryPlanningPass", "Plan Local Memory"),
     ("VerificationPass", "Verify PIM"),
     ("EmitPimCodePass", "Emit PIM Code"),
@@ -43,7 +42,7 @@ def _format_command(cmd):
 
 def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
                         crossbar_size, crossbar_count, core_count=None,
-                        pim_memory_report="none", raptor_extra_args=None, cwd=None, verbose=False,
+                        raptor_extra_args=None, cwd=None, verbose=False,
                         reporter=None, timeout_sec=None):
     # Define the arguments, with the possibility to set crossbar size and count
     args = [
@@ -57,8 +56,6 @@ def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
     ]
     if core_count is not None:
         args.append(f"--core-count={core_count}")
-    if pim_memory_report != "none":
-        args.append(f"--pim-memory-report={pim_memory_report}")
     if raptor_extra_args:
         args.extend(str(arg) for arg in raptor_extra_args)
     if verbose:

validation/validate.py

@@ -75,8 +75,6 @@ def main():
     ap.add_argument("--crossbar-size", type=int, default=128)
     ap.add_argument("--crossbar-count", type=int, default=64)
     ap.add_argument("--core-count", type=int, default=144)
-    ap.add_argument("--pim-memory-report", choices=("none", "summary", "full"), default="none",
-                    help="Emit a human-readable PIM memory planning report during codegen.")
     ap.add_argument("--raptor-extra-arg", action="append", default=[],
                     help="Additional argument to pass through to the Raptor compiler. Repeat as needed.")
     ap.add_argument("--command-timeout-seconds", type=float, default=1000000.0,
@@ -144,7 +142,6 @@ def main():
             result = validate_network(
                 onnx_path, a.raptor_path, a.onnx_include_dir, simulator_dir,
                 crossbar_size=a.crossbar_size, crossbar_count=a.crossbar_count, core_count=a.core_count,
-                pim_memory_report=a.pim_memory_report,
                 raptor_extra_args=a.raptor_extra_arg,
                 command_timeout_seconds=a.command_timeout_seconds,
                 threshold=a.threshold,

validation/validate_one.py

@@ -312,7 +312,7 @@ def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1
 
 def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                      simulator_dir, crossbar_size=128, crossbar_count=64, core_count=144,
-                     pim_memory_report="none", raptor_extra_args=None,
+                     raptor_extra_args=None,
                      threshold=1e-3, rtol=1e-5,
                      seed=0, reporter=None, model_index=1, model_total=1, verbose=False,
                      command_timeout_seconds=60.0, mode=MODE_FULL):
@@ -367,7 +367,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                 pim_pass_timings = compile_with_raptor(
                     network_onnx_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
                     core_count=core_count,
-                    pim_memory_report=pim_memory_report, raptor_extra_args=raptor_extra_args,
+                    raptor_extra_args=raptor_extra_args,
                     cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
                 print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
                 reporter.advance()
@@ -407,7 +407,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             pim_pass_timings = compile_with_raptor(
                 network_onnx_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
                 core_count=core_count,
-                pim_memory_report=pim_memory_report, raptor_extra_args=raptor_extra_args,
+                raptor_extra_args=raptor_extra_args,
                 cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
             print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
             reporter.advance()