blazingly fasterer

2026-07-21 15:43:35 +02:00
parent a893d23a74
commit 3e468b58c8
37 changed files with 1119 additions and 933 deletions
@@ -120,8 +120,8 @@ add_pim_library(OMPIMAccel
  OMSpatialToPim
  OMPimCommon
  OMPimBufferization
-  OMPimMemoryCoalescing
  OMPimHostConstantFolding
+  OMPimMemoryCoalescing
  OMPimVerification
  MLIRTensorInferTypeOpInterfaceImpl
 )
@@ -3,6 +3,7 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"

+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"

 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
@@ -12,8 +13,10 @@ namespace onnx_mlir {

 bool isCoreStaticAddressOp(mlir::Operation* op) {
  if (mlir::isa<mlir::affine::AffineApplyOp,
-                mlir::arith::ConstantOp, mlir::arith::AddIOp,
-                mlir::arith::SubIOp, mlir::arith::MulIOp,
+                mlir::arith::ConstantOp,
+                mlir::arith::AddIOp,
+                mlir::arith::SubIOp,
+                mlir::arith::MulIOp,
                mlir::arith::DivUIOp,
                mlir::arith::DivSIOp,
                mlir::arith::MinUIOp,
@@ -33,13 +36,13 @@ bool isCoreStaticAddressOp(mlir::Operation* op) {

 namespace {

-enum class CoreWalkMode { ExecuteAllIterations, StructuralExtremes };
-using CoreWalkCallback =
-  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)>;
+enum class CoreWalkMode {
+  ExecuteCommunication,
+  StructuralExtremes
+};
+using CoreWalkCallback = llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)>;

-static void propagateRegionResults(mlir::ValueRange results,
-                                   mlir::Region& region,
-                                   StaticValueKnowledge& knowledge) {
+static void propagateRegionResults(mlir::ValueRange results, mlir::Region& region, StaticValueKnowledge& knowledge) {
  if (region.empty())
    return;
  auto yield = mlir::cast<mlir::scf::YieldOp>(region.front().getTerminator());
@@ -50,23 +53,38 @@ static void propagateRegionResults(mlir::ValueRange results,
 static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block,
                                                const StaticValueKnowledge& initialKnowledge,
                                                CoreWalkMode mode,
+                                                const PimCoreCommunicationPlan* communicationPlan,
                                                CoreWalkCallback callback) {
  bool hasFailure = false;
  StaticValueKnowledge knowledge = initialKnowledge;
-  llvm::StringRef purpose = mode == CoreWalkMode::ExecuteAllIterations ? "codegen" : "verification";
-  for (mlir::Operation& op : block) {
+  llvm::StringRef purpose = mode == CoreWalkMode::ExecuteCommunication ? "communication verification" : "verification";
+  llvm::SmallVector<mlir::Operation*, 0> structuralOperations;
+  llvm::ArrayRef<mlir::Operation*> operations;
+  if (communicationPlan) {
+    auto it = communicationPlan->find(&block);
+    if (it != communicationPlan->end())
+      operations = it->second;
+  }
+  else {
+    structuralOperations.reserve(block.getOperations().size());
+    for (mlir::Operation& op : block)
+      structuralOperations.push_back(&op);
+    operations = structuralOperations;
+  }
+
+  for (mlir::Operation* operation : operations) {
+    mlir::Operation& op = *operation;
    if (mlir::isa<pim::PimHaltOp, mlir::scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
      continue;
    if (auto loadOp = mlir::dyn_cast<mlir::memref::LoadOp>(op);
        loadOp && succeeded(resolveIndexValue(loadOp.getResult(), knowledge)))
      continue;
-
    if (auto forOp = mlir::dyn_cast<mlir::scf::ForOp>(op)) {
      auto lower = resolveIndexValue(forOp.getLowerBound(), knowledge);
      auto upper = resolveIndexValue(forOp.getUpperBound(), knowledge);
      auto step = resolveIndexValue(forOp.getStep(), knowledge);
      if (failed(lower) || failed(upper) || failed(step)
-          || (mode == CoreWalkMode::ExecuteAllIterations && *step <= 0)) {
+          || (mode == CoreWalkMode::ExecuteCommunication && *step <= 0)) {
        forOp.emitOpError() << "requires statically evaluable scf.for bounds for PIM " << purpose;
        hasFailure = true;
        continue;
@@ -84,13 +102,13 @@ static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block,
        loopKnowledge.indexValues[forOp.getInductionVar()] = induction;
        for (auto [index, iterArg] : llvm::enumerate(forOp.getRegionIterArgs()))
          loopKnowledge.aliases[iterArg] = carryValues ? iterValues[index] : forOp.getInitArgs()[index];
-        hasFailure |= failed(walkPimCoreBlockImpl(body, loopKnowledge, mode, callback));
+        hasFailure |= failed(walkPimCoreBlockImpl(body, loopKnowledge, mode, communicationPlan, callback));
        auto yield = mlir::cast<mlir::scf::YieldOp>(body.getTerminator());
        for (auto [index, yielded] : llvm::enumerate(yield.getOperands()))
          iterValues[index] = resolveLoopCarriedAlias(yielded, loopKnowledge);
      };

-      if (mode == CoreWalkMode::ExecuteAllIterations) {
+      if (mode == CoreWalkMode::ExecuteCommunication) {
        for (int64_t induction = *lower; induction < *upper; induction += *step)
          visitIteration(induction, true);
      }
@@ -113,12 +131,14 @@ static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block,
        continue;
      }
      mlir::Region& selected = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion();
-      if (mode == CoreWalkMode::ExecuteAllIterations) {
-        hasFailure |= !selected.empty() && failed(walkPimCoreBlockImpl(selected.front(), knowledge, mode, callback));
+      if (mode == CoreWalkMode::ExecuteCommunication) {
+        hasFailure |= !selected.empty()
+                   && failed(walkPimCoreBlockImpl(selected.front(), knowledge, mode, communicationPlan, callback));
      }
      else {
        for (mlir::Region* region : {&ifOp.getThenRegion(), &ifOp.getElseRegion()})
-          hasFailure |= !region->empty() && failed(walkPimCoreBlockImpl(region->front(), knowledge, mode, callback));
+          hasFailure |= !region->empty()
+                     && failed(walkPimCoreBlockImpl(region->front(), knowledge, mode, communicationPlan, callback));
      }
      propagateRegionResults(ifOp.getResults(), selected, knowledge);
      continue;
@@ -138,33 +158,54 @@ static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block,
          break;
        }
      for (mlir::Region& region : switchOp->getRegions()) {
-        if (mode == CoreWalkMode::ExecuteAllIterations && &region != selected)
+        if (mode == CoreWalkMode::ExecuteCommunication && &region != selected)
          continue;
-        hasFailure |= failed(walkPimCoreBlockImpl(region.front(), knowledge, mode, callback));
+        hasFailure |= failed(walkPimCoreBlockImpl(region.front(), knowledge, mode, communicationPlan, callback));
      }
      propagateRegionResults(switchOp.getResults(), *selected, knowledge);
      continue;
    }

-    hasFailure |= failed(callback(op, knowledge));
+    if (mode != CoreWalkMode::ExecuteCommunication || mlir::isa<pim::PimSendOp, pim::PimReceiveOp>(op))
+      hasFailure |= failed(callback(op, knowledge));
  }
  return mlir::success(!hasFailure);
 }

 } // namespace

-mlir::LogicalResult
-walkPimCoreBlock(mlir::Block& block,
-                 const StaticValueKnowledge& knowledge,
-                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
-  return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::ExecuteAllIterations, callback);
+PimCoreCommunicationPlan buildPimCoreCommunicationPlan(mlir::Block& block) {
+  llvm::DenseSet<mlir::Operation*> communicationAncestors;
+  block.walk([&](mlir::Operation* op) {
+    if (!mlir::isa<pim::PimSendOp, pim::PimReceiveOp>(op))
+      return;
+    for (mlir::Operation* parent = op->getParentOp(); parent; parent = parent->getParentOp())
+      communicationAncestors.insert(parent);
+  });
+
+  PimCoreCommunicationPlan plan;
+  block.walk([&](mlir::Operation* op) {
+    bool isCommunication = mlir::isa<pim::PimSendOp, pim::PimReceiveOp>(op);
+    bool isControlFlow = mlir::isa<mlir::scf::ForOp, mlir::scf::IfOp, mlir::scf::IndexSwitchOp>(op);
+    if (isCommunication || (isControlFlow && (op->getNumResults() != 0 || communicationAncestors.contains(op))))
+      plan[op->getBlock()].push_back(op);
+  });
+  return plan;
+}
+
+mlir::LogicalResult walkPimCoreCommunicationBlock(
+  mlir::Block& block,
+  const PimCoreCommunicationPlan& plan,
+  const StaticValueKnowledge& knowledge,
+  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
+  return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::ExecuteCommunication, &plan, callback);
 }

 mlir::LogicalResult walkPimCoreBlockStructurally(
  mlir::Block& block,
  const StaticValueKnowledge& knowledge,
  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
-  return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::StructuralExtremes, callback);
+  return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::StructuralExtremes, nullptr, callback);
 }

 } // namespace onnx_mlir
@@ -3,23 +3,30 @@
 #include "mlir/IR/Block.h"
 #include "mlir/Support/LogicalResult.h"

+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/SmallVector.h"

 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"

 namespace onnx_mlir {

+using PimCoreCommunicationPlan = llvm::DenseMap<mlir::Block*, llvm::SmallVector<mlir::Operation*, 8>>;
+
 /// Returns true for ops in a `pim.core` body that only participate in static
 /// address or index computation and therefore do not emit PIM instructions.
 bool isCoreStaticAddressOp(mlir::Operation* op);

-/// Walks a `pim.core` body, statically unrolling nested `scf.for` loops when
-/// their bounds are known and invoking `callback` only on instruction-emitting
-/// operations.
-mlir::LogicalResult
-walkPimCoreBlock(mlir::Block& block,
-                 const StaticValueKnowledge& knowledge,
-                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback);
+/// Walks a `pim.core` body's communication stream, statically unrolling
+/// control flow that contains send/receive operations and invoking `callback`
+/// on those operations in execution order.
+PimCoreCommunicationPlan buildPimCoreCommunicationPlan(mlir::Block& block);
+
+mlir::LogicalResult walkPimCoreCommunicationBlock(
+  mlir::Block& block,
+  const PimCoreCommunicationPlan& plan,
+  const StaticValueKnowledge& knowledge,
+  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback);

 /// Walks a `pim.core`-like body structurally for verification without
 /// enumerating full loop trip counts. Loop bounds must still be statically
@@ -22,9 +22,19 @@
 #include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
 #include "src/Compiler/CompilerOptions.hpp"

+#include <cstdint>
+#include <limits>
+
 namespace onnx_mlir {

 inline constexpr llvm::StringLiteral kCoreIdAttrName = "coreId";
 inline constexpr llvm::StringLiteral kCoreIdsAttrName = "coreIds";
+inline constexpr llvm::StringLiteral kLocalMemoryAddressAttrName = "pim.local_memory_address";
+inline constexpr llvm::StringLiteral kLocalMemorySlotAttrName = "pim.local_memory_slot";
+inline constexpr llvm::StringLiteral kLocalMemorySlotSizeAttrName = "pim.local_memory_slot_size";
+inline constexpr llvm::StringLiteral kLocalMemoryFallbackCountAttrName = "pim.local_memory_fallback_count";
+inline constexpr llvm::StringLiteral kLocalMemoryNestedSingleUseCountAttrName =
+  "pim.local_memory_nested_single_use_count";
+inline constexpr size_t kPimLocalMemoryAddressLimit = static_cast<size_t>(std::numeric_limits<int32_t>::max());

 } // namespace onnx_mlir
@@ -18,7 +18,6 @@ add_pim_library(OMPimCompilerUtils
  PimArtifactWriter.cpp
  PimCodeGen.cpp
  PimCoreProgram.cpp
-  PimMemoryLiveness.cpp
  PimWeightEmitter.cpp

  EXCLUDE_FROM_OM_LIBS
@@ -31,8 +30,8 @@ add_pim_library(OMPimCompilerUtils
  OMPimCompilerOptions
  OMPimCommon
  OMPimBufferization
-  OMPimMemoryCoalescing
  OMPimHostConstantFolding
+  OMPimLocalMemoryPlanning
  OMPimVerification
  OMPimPasses
  OMONNXToSpatial
@@ -15,6 +15,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_ostream.h"

@@ -31,8 +32,8 @@

 #include "Common/IR/CompactAsmUtils.hpp"
 #include "Common/PimCommon.hpp"
-#include "Common/Support/Diagnostics.hpp"
 #include "Common/Support/CheckedArithmetic.hpp"
+#include "Common/Support/Diagnostics.hpp"
 #include "Common/Support/ReportUtils.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
@@ -43,7 +44,6 @@
 #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp"
-#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp"
 #include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

@@ -110,8 +110,6 @@ static Operation* getDiagnosticAnchor(mlir::Value value) {
 // PIM instruction immediates are serialized as signed int32_t fields today
 // (`sldi` goes through checkedI32OrCrash), so local addresses must stay within
 // the non-negative int32_t range.
-static constexpr size_t kPimAddressLimit = static_cast<size_t>(std::numeric_limits<int32_t>::max());
-
 static FailureOr<size_t> checkedAlignTo(size_t value, size_t alignment, Operation* anchor, StringRef fieldName) {
  if (alignment == 0)
    return value;
@@ -129,7 +127,7 @@ static void printMemoryOverflowDiagnostic(const MemoryValueKey& key,
  llvm::errs() << "Requested allocation size: " << requestedSize << " bytes\n";
  llvm::errs() << "Current firstAvailableAddress: " << currentFirstAvailableAddress << "\n";
  llvm::errs() << "Aligned end address: " << alignedEndAddress << "\n";
-  llvm::errs() << "Address limit: " << kPimAddressLimit << " (signed int32_t immediate range)\n";
+  llvm::errs() << "Address limit: " << kPimLocalMemoryAddressLimit << " (signed int32_t immediate range)\n";
  if (key.lane)
    llvm::errs() << "Lane: " << *key.lane << "\n";
  llvm::errs() << "Value: ";
@@ -152,10 +150,13 @@ size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) {
  FailureOr<size_t> checkedAlignedEnd = failure();
  if (succeeded(checkedEnd))
    checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
-  if (address > kPimAddressLimit || failed(checkedEnd) || *checkedEnd > kPimAddressLimit
-      || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimAddressLimit) {
+  if (address > kPimLocalMemoryAddressLimit || failed(checkedEnd) || *checkedEnd > kPimLocalMemoryAddressLimit
+      || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimLocalMemoryAddressLimit) {
    printMemoryOverflowDiagnostic(
-      key, size, firstAvailableAddress, succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
+      key,
+      size,
+      firstAvailableAddress,
+      succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimLocalMemoryAddressLimit);
    llvm_unreachable("PIM local memory allocation overflow");
  }
  firstAvailableAddress = *checkedAlignedEnd;
@@ -202,15 +203,6 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE
  }
 }

-PhysicalSlotInfo PimMemory::allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key) {
-  PhysicalSlotInfo slot;
-  slot.id = nextPhysicalSlotId++;
-  slot.address = allocateAddress(slotSize, key);
-  slot.size = slotSize;
-  localPhysicalSlots.push_back(slot);
-  return slot;
-}
-
 void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
  SmallDenseMap<memref::GlobalOp, mlir::Value, 8> globalConstants;
  SmallVector<std::pair<mlir::Value, mlir::Value>, 16> globalAliases;
@@ -249,71 +241,17 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
    globalMemEntriesMap[getMemoryValueKey(alias)] = getMemEntry(getMemoryValueKey(original));
 }

-void PimMemory::allocateCore(Operation* op, std::optional<unsigned> lane) {
-  auto intervals = buildLocalAllocIntervals(op, lane, pimMemoryReport == PimMemoryReportFull);
-  SmallVector<PlannedPhysicalSlot> plannedSlots = planPhysicalSlots(intervals);
-
-  SmallVector<size_t> slotOrder(plannedSlots.size());
-  std::iota(slotOrder.begin(), slotOrder.end(), 0);
-  llvm::stable_sort(slotOrder, [&](size_t lhsIndex, size_t rhsIndex) {
-    const PlannedPhysicalSlot& lhs = plannedSlots[lhsIndex];
-    const PlannedPhysicalSlot& rhs = plannedSlots[rhsIndex];
-    if (lhs.requiredSize != rhs.requiredSize)
-      return lhs.requiredSize > rhs.requiredSize;
-    return lhs.id < rhs.id;
-  });
-
-  SmallVector<bool, 16> usedExistingSlots(localPhysicalSlots.size(), false);
-  for (size_t slotIndex : slotOrder) {
-    PlannedPhysicalSlot& slot = plannedSlots[slotIndex];
-    size_t bestExistingIndex = std::numeric_limits<size_t>::max();
-    auto bestKey = std::tuple<size_t, size_t, size_t>(
-      std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max());
-
-    for (size_t existingIndex = 0; existingIndex < localPhysicalSlots.size(); ++existingIndex) {
-      if (usedExistingSlots[existingIndex])
-        continue;
-      const PhysicalSlotInfo& existingSlot = localPhysicalSlots[existingIndex];
-      if (existingSlot.size < slot.requiredSize)
-        continue;
-      auto candidateKey =
-        std::tuple<size_t, size_t, size_t>(existingSlot.size - slot.requiredSize, existingSlot.size, existingSlot.id);
-      if (candidateKey < bestKey) {
-        bestKey = candidateKey;
-        bestExistingIndex = existingIndex;
-      }
-    }
-
-    if (bestExistingIndex != std::numeric_limits<size_t>::max()) {
-      const PhysicalSlotInfo& existingSlot = localPhysicalSlots[bestExistingIndex];
-      slot.id = existingSlot.id;
-      slot.address = existingSlot.address;
-      slot.size = existingSlot.size;
-      usedExistingSlots[bestExistingIndex] = true;
-    }
-    else {
-      PhysicalSlotInfo newSlot = allocatePhysicalSlot(slot.requiredSize, intervals[slot.intervalIndices.front()].key);
-      slot.id = newSlot.id;
-      slot.address = newSlot.address;
-      slot.size = newSlot.size;
-      usedExistingSlots.push_back(true);
-    }
-
-    for (size_t intervalIndex : slot.intervalIndices) {
-      LocalAllocInterval& interval = intervals[intervalIndex];
-      interval.physicalSlotId = slot.id;
-      interval.assignedAddress = slot.address;
-      interval.physicalSlotSize = slot.size;
-      MemEntry memEntry {slot.address, interval.size};
-      ownedMemEntriesMap[interval.key] = memEntry;
-      globalMemEntriesMap[interval.key] = memEntry;
-    }
+void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane) {
+  if (localPhysicalSlots.empty()) {
+    llvm::append_range(localPhysicalSlots, plan.slots);
+    reportRow.logicalAllocaBytes = plan.logicalBytes;
+    reportRow.fallbackIntervals = plan.fallbackIntervals;
+    reportRow.nestedSingleUseIntervals = plan.nestedSingleUseIntervals;
  }
-
-  if (pimMemoryReport != PimMemoryReportNone) {
-    MemoryPlanArtifacts artifacts =
-      buildMemoryPlanArtifacts(op, lane, intervals, plannedSlots, kPimAddressLimit, pimMemoryReport);
-    livenessArtifacts += artifacts;
+  for (const CompiledLocalMemoryEntry& entry : plan.entries) {
+    MemoryValueKey key = getMemoryValueKey(entry.value, lane);
+    ownedMemEntriesMap[key] = entry.memory;
+    globalMemEntriesMap[key] = entry.memory;
  }
 }

@@ -475,13 +413,41 @@ void PimAcceleratorMemory::flushReport() {
  uint64_t totalGlobalMemory = hostReportRow.has_value() ? hostReportRow->sizeGlobal : 0;
  uint64_t totalWeightsMemory = totalWeightBytes;
  uint64_t totalCoresMemory = 0;
-  for (const MemoryReportEntry& entry : reportEntries)
+  uint64_t totalLogicalCoreMemory = 0;
+  uint64_t totalFallbackIntervals = 0;
+  uint64_t totalNestedSingleUseIntervals = 0;
+  uint64_t maximumCorePeak = 0;
+  std::string worstEntry = "<none>";
+  for (const MemoryReportEntry& entry : reportEntries) {
    totalCoresMemory += entry.totalAllocaBytes;
+    uint64_t factor = entry.kind == MemoryReportEntry::Kind::Batch ? entry.coreIds.size() : 1;
+    totalLogicalCoreMemory += entry.row.logicalAllocaBytes * factor;
+    totalFallbackIntervals += entry.row.fallbackIntervals * factor;
+    totalNestedSingleUseIntervals += entry.row.nestedSingleUseIntervals * factor;
+    if (entry.row.sizeAlloca <= maximumCorePeak)
+      continue;
+    maximumCorePeak = entry.row.sizeAlloca;
+    worstEntry = entry.kind == MemoryReportEntry::Kind::Batch ? "Batch " + std::to_string(entry.id)
+                                                              : "Core " + std::to_string(entry.coreIds.front());
+  }
+  uint64_t savedCoreMemory =
+    totalLogicalCoreMemory >= totalCoresMemory ? totalLogicalCoreMemory - totalCoresMemory : 0;
+  double savedPercent = totalLogicalCoreMemory == 0
+                        ? 0.0
+                        : 100.0 * static_cast<double>(savedCoreMemory) / static_cast<double>(totalLogicalCoreMemory);

-  llvm::SmallVector<ReportField, 3> totalFields = {
-    {"Global memory",  formatReportMemory(totalGlobalMemory) },
-    {"Weights memory", formatReportMemory(totalWeightsMemory)},
-    {"Cores memory",   formatReportMemory(totalCoresMemory)  }
+  llvm::SmallVector<ReportField, 10> totalFields = {
+    {"Global memory",               formatReportMemory(totalGlobalMemory)               },
+    {"Weights memory",              formatReportMemory(totalWeightsMemory)              },
+    {"Logical core allocations",    formatReportMemory(totalLogicalCoreMemory)           },
+    {"Physical core memory",        formatReportMemory(totalCoresMemory)                 },
+    {"Cores memory",                formatReportMemory(totalCoresMemory)                 },
+    {"Saved core memory",           formatReportMemory(savedCoreMemory)                  },
+    {"Saved core memory percent",   formatv("{0:F2}%", savedPercent).str()               },
+    {"Maximum core peak",           formatReportMemory(maximumCorePeak)                  },
+    {"Worst core/batch",            worstEntry                                           },
+    {"Fallback intervals",          std::to_string(totalFallbackIntervals)               },
+    {"Nested single-use intervals", std::to_string(totalNestedSingleUseIntervals)        }
  };
  printReportTotalsBlock(os, totalFields);

@@ -552,19 +518,14 @@ void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instructio

 void PimCodeGen::updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const {
  switch (instruction.opcode) {
-  case pim_binary::Opcode::sldi:
-    scalarRegisterValues[instruction.rd] = instruction.r2OrImm;
-    break;
+  case pim_binary::Opcode::sldi:  scalarRegisterValues[instruction.rd] = instruction.r2OrImm; break;
  case pim_binary::Opcode::sld:
  case pim_binary::Opcode::sadd:
  case pim_binary::Opcode::ssub:
  case pim_binary::Opcode::smul:
  case pim_binary::Opcode::saddi:
-  case pim_binary::Opcode::smuli:
-    scalarRegisterValues[instruction.rd].reset();
-    break;
-  default:
-    break;
+  case pim_binary::Opcode::smuli: scalarRegisterValues[instruction.rd].reset(); break;
+  default:                        break;
  }
 }

@@ -580,8 +541,8 @@ void PimCodeGen::genSetRegisterImmediate(uint8_t registerNumber, int32_t immedia
 }

 void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const {
-  genSetRegisterImmediate(
-    pim::checkedU8OrCrash(registerNumber, "register number"), pim::checkedI32OrCrash(immediate, "register immediate"));
+  genSetRegisterImmediate(pim::checkedU8OrCrash(registerNumber, "register number"),
+                          pim::checkedI32OrCrash(immediate, "register immediate"));
 }

 void PimCodeGen::setupRd(size_t rdAddress, size_t rdOffset) const {
@@ -729,8 +690,7 @@ void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKno
    for (size_t outerIndex = 0; outerIndex < outerCount; ++outerIndex) {
      size_t dstOffset = (outerIndex * outputConcatDim + concatOffset) * innerCount * elementSize;
      size_t srcOffset = outerIndex * inputConcatDim * innerCount * elementSize;
-      emitMemCopyOp(
-        pim_binary::Opcode::lmv, outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len");
+      emitMemCopyOp(pim_binary::Opcode::lmv, outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len");
    }

    concatOffset += inputConcatDim;
@@ -788,10 +748,11 @@ void PimCodeGen::codeGenTransposeOp(const CompiledTransposePlan& plan, const Sta
  size_t maxElementOffset = plan.totalBytes == 0 ? 0 : plan.totalBytes - plan.elementBytes;
  int32_t maxSourceAddress = pim::checkedI32OrCrash(
    pim::checkedAddOrCrash(srcAddr, maxElementOffset, "transpose source address"), "transpose source address");
-  int32_t maxDestinationAddress = pim::checkedI32OrCrash(
-    pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"), "transpose destination address");
-  (void)maxSourceAddress;
-  (void)maxDestinationAddress;
+  int32_t maxDestinationAddress =
+    pim::checkedI32OrCrash(pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"),
+                           "transpose destination address");
+  (void) maxSourceAddress;
+  (void) maxDestinationAddress;

  pim_binary::InstructionRecord copyInstruction;
  copyInstruction.opcode = pim_binary::Opcode::lmv;
@@ -871,6 +832,60 @@ static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
  return coreLikeOps;
 }

+static FailureOr<CompiledCoreMemoryPlan> compileCoreMemoryPlan(Operation* coreLikeOp) {
+  CompiledCoreMemoryPlan plan;
+  auto fallbackAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryFallbackCountAttrName);
+  auto nestedSingleUseAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryNestedSingleUseCountAttrName);
+  if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) {
+    coreLikeOp->emitError("requires complete PIM local-memory planning summary attributes before codegen");
+    return failure();
+  }
+  plan.fallbackIntervals = static_cast<uint64_t>(fallbackAttr.getInt());
+  plan.nestedSingleUseIntervals = static_cast<uint64_t>(nestedSingleUseAttr.getInt());
+  SmallDenseMap<size_t, size_t, 32> slotIndices;
+  bool hasFailure = false;
+  coreLikeOp->walk([&](memref::AllocOp allocOp) {
+    if (hasFailure)
+      return;
+    auto addressAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
+    auto slotAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotAttrName);
+    auto slotSizeAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotSizeAttrName);
+    if (!addressAttr || !slotAttr || !slotSizeAttr || addressAttr.getInt() < 0 || slotAttr.getInt() < 0
+        || slotSizeAttr.getInt() < 0) {
+      allocOp.emitOpError("requires a complete non-negative PIM local-memory plan before codegen");
+      hasFailure = true;
+      return;
+    }
+
+    auto checkedSize = pim::getCheckedShapedTypeSizeInBytes(
+      cast<ShapedType>(allocOp.getType()), allocOp, "planned local allocation byte size");
+    if (failed(checkedSize)) {
+      hasFailure = true;
+      return;
+    }
+    size_t address = static_cast<size_t>(addressAttr.getInt());
+    size_t slotId = static_cast<size_t>(slotAttr.getInt());
+    size_t slotSize = static_cast<size_t>(slotSizeAttr.getInt());
+    plan.entries.push_back({allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}});
+    plan.logicalBytes += *checkedSize;
+
+    auto [slotIt, inserted] = slotIndices.try_emplace(slotId, plan.slots.size());
+    if (inserted) {
+      plan.slots.push_back({slotId, 0, slotSize});
+      return;
+    }
+    const PhysicalSlotInfo& existing = plan.slots[slotIt->second];
+    if (existing.size != slotSize) {
+      allocOp.emitOpError("has a PIM local-memory arena inconsistent with another allocation");
+      hasFailure = true;
+    }
+  });
+  if (hasFailure)
+    return failure();
+  llvm::sort(plan.slots, [](const PhysicalSlotInfo& lhs, const PhysicalSlotInfo& rhs) { return lhs.id < rhs.id; });
+  return plan;
+}
+
 struct CoreEmissionResult {
  static constexpr size_t kMaxStoredCodegenDiagnostics = 8;

@@ -882,10 +897,8 @@ struct CoreEmissionResult {
  OnnxMlirCompilerErrorCodes status = CompilerSuccess;
  MemoryReportRow reportRow;
  llvm::SmallVector<ResolvedWeightView, 8> usedWeights;
-  MemoryPlanArtifacts livenessArtifacts;
  llvm::SmallVector<DiagnosticRecord, kMaxStoredCodegenDiagnostics> diagnostics;
  size_t diagnosticCount = 0;
-
  void recordDiagnostic(Operation* op, StringRef message) {
    ++diagnosticCount;
    if (diagnostics.size() < kMaxStoredCodegenDiagnostics)
@@ -1012,13 +1025,21 @@ static LogicalResult executeCompiledCorePlan(
    }

    auto emitBinary = [&](auto op, pim_binary::Opcode opcode) {
-      coreCodeGen.emitBinaryVectorOp(opcode, op.getOutputBuffer(), op.getLhs(), op.getRhs(),
-                                     getVectorByteSizeOrCrash(cast<ShapedType>(op.getLhs().getType())), knowledge);
+      coreCodeGen.emitBinaryVectorOp(opcode,
+                                     op.getOutputBuffer(),
+                                     op.getLhs(),
+                                     op.getRhs(),
+                                     getVectorByteSizeOrCrash(cast<ShapedType>(op.getLhs().getType())),
+                                     knowledge);
    };
    auto emitUnary = [&](auto op, pim_binary::Opcode opcode, int32_t r2OrImm, int32_t generic1) {
-      coreCodeGen.emitUnaryVectorOp(opcode, op.getOutputBuffer(), op.getInput(),
+      coreCodeGen.emitUnaryVectorOp(opcode,
+                                    op.getOutputBuffer(),
+                                    op.getInput(),
                                    getVectorByteSizeOrCrash(cast<ShapedType>(op.getInput().getType())),
-                                    knowledge, r2OrImm, generic1);
+                                    knowledge,
+                                    r2OrImm,
+                                    generic1);
    };

    switch (node.opKind) {
@@ -1038,20 +1059,19 @@ static LogicalResult executeCompiledCorePlan(
      else
        return failure();
      break;
-    case CompiledCoreOpKind::Transpose:
-      coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge);
-      break;
-    case CompiledCoreOpKind::VVAdd:  emitBinary(cast<pim::PimVVAddOp>(node.op), pim_binary::Opcode::vvadd); break;
-    case CompiledCoreOpKind::VVSub:  emitBinary(cast<pim::PimVVSubOp>(node.op), pim_binary::Opcode::vvsub); break;
-    case CompiledCoreOpKind::VVMul:  emitBinary(cast<pim::PimVVMulOp>(node.op), pim_binary::Opcode::vvmul); break;
-    case CompiledCoreOpKind::VVMax:  emitBinary(cast<pim::PimVVMaxOp>(node.op), pim_binary::Opcode::vvmax); break;
-    case CompiledCoreOpKind::VVDMul: emitBinary(cast<pim::PimVVDMulOp>(node.op), pim_binary::Opcode::vvdmul); break;
-    case CompiledCoreOpKind::VAvg:   emitUnary(cast<pim::PimVAvgOp>(node.op), pim_binary::Opcode::vavg, 1, 1); break;
-    case CompiledCoreOpKind::VRelu:  emitUnary(cast<pim::PimVReluOp>(node.op), pim_binary::Opcode::vrelu, 0, 0); break;
-    case CompiledCoreOpKind::VTanh:  emitUnary(cast<pim::PimVTanhOp>(node.op), pim_binary::Opcode::vtanh, 0, 0); break;
-    case CompiledCoreOpKind::VSigm:  emitUnary(cast<pim::PimVSigmOp>(node.op), pim_binary::Opcode::vsigm, 0, 0); break;
+    case CompiledCoreOpKind::Transpose: coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge); break;
+    case CompiledCoreOpKind::VVAdd:     emitBinary(cast<pim::PimVVAddOp>(node.op), pim_binary::Opcode::vvadd); break;
+    case CompiledCoreOpKind::VVSub:     emitBinary(cast<pim::PimVVSubOp>(node.op), pim_binary::Opcode::vvsub); break;
+    case CompiledCoreOpKind::VVMul:     emitBinary(cast<pim::PimVVMulOp>(node.op), pim_binary::Opcode::vvmul); break;
+    case CompiledCoreOpKind::VVMax:     emitBinary(cast<pim::PimVVMaxOp>(node.op), pim_binary::Opcode::vvmax); break;
+    case CompiledCoreOpKind::VVDMul:    emitBinary(cast<pim::PimVVDMulOp>(node.op), pim_binary::Opcode::vvdmul); break;
+    case CompiledCoreOpKind::VAvg:      emitUnary(cast<pim::PimVAvgOp>(node.op), pim_binary::Opcode::vavg, 1, 1); break;
+    case CompiledCoreOpKind::VRelu:     emitUnary(cast<pim::PimVReluOp>(node.op), pim_binary::Opcode::vrelu, 0, 0); break;
+    case CompiledCoreOpKind::VTanh:     emitUnary(cast<pim::PimVTanhOp>(node.op), pim_binary::Opcode::vtanh, 0, 0); break;
+    case CompiledCoreOpKind::VSigm:     emitUnary(cast<pim::PimVSigmOp>(node.op), pim_binary::Opcode::vsigm, 0, 0); break;
    case CompiledCoreOpKind::VSoftmax:
-      emitUnary(cast<pim::PimVSoftmaxOp>(node.op), pim_binary::Opcode::vsoftmax, 0, 0); break;
+      emitUnary(cast<pim::PimVSoftmaxOp>(node.op), pim_binary::Opcode::vsoftmax, 0, 0);
+      break;
    }
  }
  return success();
@@ -1103,8 +1123,7 @@ static void aliasMaterializedHostGlobals(CoreLikeOpTy coreLikeOp,
 }

 static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) {
-  std::string outputCorePath =
-    (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
+  std::string outputCorePath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
  std::error_code errorCode;
  raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
  if (errorCode) {
@@ -1128,8 +1147,7 @@ static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath
  if (!pimEmitJson.getValue())
    return CompilerSuccess;

-  std::string outputCoreJsonPath =
-    (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
+  std::string outputCoreJsonPath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
  errorCode = std::error_code();
  raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode);
  if (errorCode) {
@@ -1169,18 +1187,27 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
  SmallDenseMap<memref::GlobalOp, MemEntry, 16> materializedHostGlobals =
    collectMaterializedHostGlobals(moduleOp, funcOp, memory);
  llvm::DenseMap<Operation*, std::unique_ptr<CompiledCoreProgram>> compiledPrograms;
+  llvm::DenseMap<Operation*, std::unique_ptr<CompiledCoreMemoryPlan>> memoryPlans;
  for (Operation* op : coreLikeOps) {
    auto program = std::make_unique<CompiledCoreProgram>();
    if (failed(compileCoreProgram(op, *program)))
      return CompilerFailure;
    compiledPrograms.try_emplace(op, std::move(program));
+    auto memoryPlan = compileCoreMemoryPlan(op);
+    if (failed(memoryPlan))
+      return CompilerFailure;
+    memoryPlans.try_emplace(op, std::make_unique<CompiledCoreMemoryPlan>(std::move(*memoryPlan)));
  }
-
  auto getCompiledProgram = [&](Operation* op) {
    auto it = compiledPrograms.find(op);
    assert(it != compiledPrograms.end() && "missing compiled PIM core program");
    return it->second.get();
  };
+  auto getMemoryPlan = [&](Operation* op) {
+    auto it = memoryPlans.find(op);
+    assert(it != memoryPlans.end() && "missing PIM core memory plan");
+    return it->second.get();
+  };

  llvm::DenseMap<size_t, size_t> emittedCoreIds;
  size_t nextEmittedCoreId = 0;
@@ -1210,6 +1237,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
      CoreEmissionJob job;
      job.coreLikeOp = coreOp;
      job.program = getCompiledProgram(op);
+      job.memoryPlan = getMemoryPlan(op);
      job.emittedCoreId = emittedCoreIds.lookup(originalCoreId);
      jobs.push_back(std::move(job));
      continue;
@@ -1229,6 +1257,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
      CoreEmissionJob job;
      job.coreLikeOp = coreBatchOp;
      job.program = getCompiledProgram(op);
+      job.memoryPlan = getMemoryPlan(op);
      job.emittedCoreId = emittedCoreIds.lookup(originalCoreId);
      job.lanes = lanesByCoreId.lookup(originalCoreId);
      job.batchReportId = nextBatchReportId;
@@ -1332,17 +1361,16 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
    if (auto coreOp = dyn_cast<pim::PimCoreOp>(job.coreLikeOp)) {
      aliasMaterializedHostGlobals(coreOp, moduleOp, materializedHostGlobals, jobMemory);
      auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId);
-      deviceMemory.allocateCore(coreOp);
+      deviceMemory.allocateCore(*job.memoryPlan);

      StaticValueKnowledge knowledge = seedCoreCodegenKnowledge(coreOp);
      if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) {
-        (void)finalizeInstructions();
+        (void) finalizeInstructions();
        result.status = CompilerFailure;
        return result;
      }
      result.reportRow = deviceMemory.getReportRow();
      result.usedWeights = std::move(usedWeights);
-      result.livenessArtifacts = deviceMemory.getLivenessArtifacts();
    }
    else {
      auto coreBatchOp = cast<pim::PimCoreBatchOp>(job.coreLikeOp);
@@ -1352,10 +1380,10 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
      for (unsigned lane : job.lanes) {
        StaticValueKnowledge knowledge = seedCoreBatchCodegenKnowledge(coreBatchOp, lane);

-        deviceMemory.allocateCore(coreBatchOp, lane);
+        deviceMemory.allocateCore(*job.memoryPlan, lane);
        coreCodeGen.setBatchLane(lane);
        if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) {
-          (void)finalizeInstructions();
+          (void) finalizeInstructions();
          result.status = CompilerFailure;
          return result;
        }
@@ -1363,7 +1391,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::

      result.reportRow = deviceMemory.getReportRow();
      result.usedWeights = std::move(usedWeights);
-      result.livenessArtifacts = deviceMemory.getLivenessArtifacts();
    }

    if (!finalizeInstructions()) {
@@ -1388,9 +1415,8 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
  for (const CoreEmissionResult& result : jobResults) {
    if (!summaryAnchor && !result.diagnostics.empty())
      summaryAnchor = result.diagnostics.front().op;
-    for (const CoreEmissionResult::DiagnosticRecord& diagnostic : result.diagnostics) {
+    for (const CoreEmissionResult::DiagnosticRecord& diagnostic : result.diagnostics)
      diagnostics.report(diagnostic.op, [&](Operation* op) { op->emitError() << diagnostic.message; });
-    }
    size_t unreportedCount = result.diagnosticCount - result.diagnostics.size();
    diagnostics.noteFailures(static_cast<int64_t>(unreportedCount));
  }
@@ -1420,19 +1446,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
  auto weightEmission = createAndPopulateWeightFolder(weightRequests, outputDirPath);
  memory.setTotalWeightBytes(weightEmission.totalWeightBytes);
  auto& mapCoreWeightToFileName = weightEmission.mapCoreWeightToFileName;
-  if (std::string reportsRoot = getOutputDir(); !reportsRoot.empty()) {
-    std::string reportsDir = reportsRoot + "/reports";
-    sys::fs::remove(reportsDir + "/pim_memory_liveness_report.txt");
-    sys::fs::remove(reportsDir + "/pim_memory_liveness_report.json");
-    sys::fs::remove(reportsDir + "/pim_memory_liveness_timeline.dot");
-  }
-  std::fstream livenessReportFile;
-  std::unique_ptr<llvm::raw_os_ostream> livenessReportOs;
-  if (pimMemoryReport != PimMemoryReportNone) {
-    livenessReportFile = openReportFileWithExtension("pim_memory_liveness_report", "txt");
-    livenessReportOs = std::make_unique<llvm::raw_os_ostream>(livenessReportFile);
-  }
-
  for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
    const CoreEmissionJob& job = jobs[jobIndex];
    const CoreEmissionResult& result = jobResults[jobIndex];
@@ -1443,8 +1456,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
        return err;
      xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
      memory.recordCoreReport(job.emittedCoreId, result.reportRow);
-      if (livenessReportFile.is_open())
-        *livenessReportOs << "Core " << job.emittedCoreId << ":\n" << result.livenessArtifacts;
      continue;
    }
  }
@@ -1473,18 +1484,10 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
                             batchPerCoreRow.value_or(MemoryReportRow {}),
                             batchRow.numAlloca,
                             batchRow.sizeAlloca);
-    if (livenessReportFile.is_open())
-      for (size_t jobIndex : group)
-        *livenessReportOs << "Batch " << batchReportId << " core " << jobs[jobIndex].emittedCoreId << ":\n"
-                          << jobResults[jobIndex].livenessArtifacts;
  }

  maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1;

-  if (livenessReportFile.is_open()) {
-    livenessReportOs->flush();
-    livenessReportFile.close();
-  }
  memory.flushReport();
  return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath);
 }
@@ -39,8 +39,6 @@ struct PhysicalSlotInfo {
  size_t size = 0;
 };

-using MemoryPlanArtifacts = std::string;
-
 struct MemoryValueKey {
  mlir::Value value;
  std::optional<unsigned> lane;
@@ -48,15 +46,33 @@ struct MemoryValueKey {
  bool operator==(const MemoryValueKey& other) const { return value == other.value && lane == other.lane; }
 };

+struct CompiledLocalMemoryEntry {
+  mlir::Value value;
+  MemEntry memory;
+};
+
+struct CompiledCoreMemoryPlan {
+  llvm::SmallVector<CompiledLocalMemoryEntry, 32> entries;
+  llvm::SmallVector<PhysicalSlotInfo, 32> slots;
+  uint64_t logicalBytes = 0;
+  uint64_t fallbackIntervals = 0;
+  uint64_t nestedSingleUseIntervals = 0;
+};
+
 struct MemoryReportRow {
  uint64_t numAlloca = 0;
  uint64_t sizeAlloca = 0;
  uint64_t numGlobal = 0;
  uint64_t sizeGlobal = 0;
+  uint64_t logicalAllocaBytes = 0;
+  uint64_t fallbackIntervals = 0;
+  uint64_t nestedSingleUseIntervals = 0;

  bool operator==(const MemoryReportRow& other) const {
    return numAlloca == other.numAlloca && sizeAlloca == other.sizeAlloca && numGlobal == other.numGlobal
-        && sizeGlobal == other.sizeGlobal;
+        && sizeGlobal == other.sizeGlobal && logicalAllocaBytes == other.logicalAllocaBytes
+        && fallbackIntervals == other.fallbackIntervals
+        && nestedSingleUseIntervals == other.nestedSingleUseIntervals;
  }
 };

@@ -93,26 +109,22 @@ class PimMemory {
  llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap;
  llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32> ownedMemEntriesMap;
  MemoryReportRow reportRow;
-  MemoryPlanArtifacts livenessArtifacts;

  size_t minAlignment = 4;
  size_t firstAvailableAddress = 0;
-  size_t nextPhysicalSlotId = 0;

  MemEntry* gatherMemEntry(mlir::Value value, std::optional<unsigned> lane = std::nullopt);
  size_t allocateAddress(size_t size, const MemoryValueKey& key);
  void allocateGatheredMemory();
  void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind);
-  PhysicalSlotInfo allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key);

 public:
  PimMemory(llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap)
  : globalMemEntriesMap(globalMemEntriesMap) {}

  void allocateHost(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp);
-  void allocateCore(mlir::Operation* op, std::optional<unsigned> lane = std::nullopt);
+  void allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane = std::nullopt);
  MemoryReportRow getReportRow() const;
-  const MemoryPlanArtifacts& getLivenessArtifacts() const { return livenessArtifacts; }

  size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
  MemEntry getMemEntry(const MemoryValueKey& key) const;
@@ -160,6 +172,7 @@ public:
 struct CoreEmissionJob {
  mlir::Operation* coreLikeOp = nullptr;
  const CompiledCoreProgram* program = nullptr;
+  const CompiledCoreMemoryPlan* memoryPlan = nullptr;
  size_t emittedCoreId = 0;
  llvm::SmallVector<unsigned, 4> lanes;
  std::optional<uint64_t> batchReportId;
@@ -74,7 +74,7 @@ llvm::cl::opt<bool>

 llvm::cl::opt<bool>
  pimDisableMemoryCoalescing("pim-disable-memory-coalescing",
-                             llvm::cl::desc("Skip the PIM memory coalescing pass (developer diagnostic option)"),
+                             llvm::cl::desc("Skip the early PIM IR memory coalescing pass (developer diagnostic)"),
                             llvm::cl::init(false),
                             llvm::cl::cat(OnnxMlirOptions));

@@ -124,10 +124,10 @@ llvm::cl::opt<bool> pimTraceCommunicationMaterialization(
  llvm::cl::cat(OnnxMlirOptions));

 llvm::cl::opt<size_t>
-  crossbarSize("crossbar-size", llvm::cl::desc("Width and height of a single crossbar"), llvm::cl::init(2));
+  crossbarSize("crossbar-size", llvm::cl::desc("Width and height of a single crossbar"), llvm::cl::init(128));

 llvm::cl::opt<size_t>
-  crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(256));
+  crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(64));

 llvm::cl::opt<long> coresCount("core-count",
                               llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."),
@@ -53,8 +53,8 @@ extern llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport;
 extern llvm::cl::opt<PimConvLoweringType> pimConvLowering;
 extern llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow;

-extern llvm::cl::opt<bool> pimOnlyCodegen;
 extern llvm::cl::opt<bool> pimDisableMemoryCoalescing;
+extern llvm::cl::opt<bool> pimOnlyCodegen;
 extern llvm::cl::opt<bool> useExperimentalConvImpl;
 extern llvm::cl::opt<bool> pimEmitJson;
 extern llvm::cl::opt<bool> pimReportConvLowering;
@@ -21,6 +21,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
  verifyExplicitPimCoreCount();

  if (pimOnlyCodegen) {
+    pm.addPass(createPimLocalMemoryPlanningPass());
    pm.addPass(createEmitPimCodePass());
    return;
  }
@@ -53,6 +54,8 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
    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());
    pm.addPass(createMessagePass("Pim verified"));
    pm.addPass(createEmitPimCodePass());
@@ -0,0 +1,8 @@
+add_pim_library(OMPimLocalMemoryLifetimeAnalysis
+  LocalMemoryLifetimeAnalysis.cpp
+
+  EXCLUDE_FROM_OM_LIBS
+
+  INCLUDE_DIRS PUBLIC
+  ${PIM_PUBLIC_INCLUDE_DIRS}
+)
@@ -0,0 +1,71 @@
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
+
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace pim {
+
+bool isLocalMemoryAliasOp(Operation* op) {
+  return isa<memref::SubViewOp, memref::CastOp, memref::CollapseShapeOp, memref::ExpandShapeOp>(op);
+}
+
+LogicalResult walkLocalMemoryUses(Value root,
+                                  llvm::function_ref<LogicalResult(Value, Operation*)> visitUser) {
+  llvm::SmallPtrSet<Value, 16> visitedValues;
+  llvm::SmallPtrSet<Operation*, 32> visitedUsers;
+  llvm::SmallVector<Value> pendingValues {root};
+  auto addAlias = [&](Value value) { pendingValues.push_back(value); };
+
+  while (!pendingValues.empty()) {
+    Value value = pendingValues.pop_back_val();
+    if (!visitedValues.insert(value).second)
+      continue;
+    for (Operation* user : value.getUsers()) {
+      if (!visitedUsers.insert(user).second)
+        continue;
+      if (failed(visitUser(value, user)))
+        return failure();
+
+      if (isLocalMemoryAliasOp(user))
+        for (Value result : user->getResults())
+          addAlias(result);
+
+      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user))
+        for (OpResult result : user->getResults())
+          if (OpOperand* tied = dpsOp.getTiedOpOperand(result); tied && tied->get() == value)
+            addAlias(result);
+
+      if (auto forOp = dyn_cast<scf::ForOp>(user))
+        for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs()))
+          if (initArg == value) {
+            addAlias(forOp.getRegionIterArgs()[index]);
+            addAlias(forOp.getResult(index));
+          }
+
+      auto yieldOp = dyn_cast<scf::YieldOp>(user);
+      if (!yieldOp)
+        continue;
+      for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
+        if (operand != value)
+          continue;
+        if (auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp()))
+          addAlias(forOp.getResult(index));
+        else if (auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp()))
+          addAlias(ifOp.getResult(index));
+        else if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp()))
+          addAlias(switchOp.getResult(index));
+      }
+    }
+  }
+  return success();
+}
+
+} // namespace pim
+} // namespace onnx_mlir
@@ -0,0 +1,17 @@
+#pragma once
+
+#include "mlir/IR/Operation.h"
+
+#include "llvm/ADT/STLFunctionalExtras.h"
+
+namespace onnx_mlir {
+namespace pim {
+
+bool isLocalMemoryAliasOp(mlir::Operation* op);
+
+mlir::LogicalResult walkLocalMemoryUses(
+  mlir::Value root,
+  llvm::function_ref<mlir::LogicalResult(mlir::Value, mlir::Operation*)> visitUser);
+
+} // namespace pim
+} // namespace onnx_mlir
@@ -1,9 +1,11 @@
 add_onnx_mlir_dialect(Pim pim)
 add_onnx_mlir_dialect_doc(pim Pim.td)

+add_subdirectory(Analysis)
 add_subdirectory(Transforms/Bufferization)
-add_subdirectory(Transforms/MemoryCoalescing)
 add_subdirectory(Transforms/HostConstantFolding)
+add_subdirectory(Transforms/MemoryCoalescing)
+add_subdirectory(Transforms/LocalMemoryPlanning)
 add_subdirectory(Transforms/Verification)

 add_pim_library(PimOps
@@ -0,0 +1,14 @@
+add_pim_library(OMPimLocalMemoryPlanning
+  LocalMemoryPlanning.cpp
+
+  EXCLUDE_FROM_OM_LIBS
+
+  INCLUDE_DIRS PUBLIC
+  ${PIM_PUBLIC_INCLUDE_DIRS}
+
+  LINK_LIBS PUBLIC
+  OMPimCommon
+  OMPimCompilerOptions
+  OMPimLocalMemoryLifetimeAnalysis
+  PimOps
+)
@@ -1,14 +1,18 @@
 #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 "mlir/Interfaces/DestinationStyleOpInterface.h"

 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallPtrSet.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>
@@ -16,9 +20,13 @@

 #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/PimMemoryLiveness.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;
@@ -26,19 +34,6 @@ using namespace onnx_mlir;

 namespace {

-static std::optional<unsigned> getLaneForMemoryValue(mlir::Value value, std::optional<unsigned> lane) {
-  if (!lane)
-    return std::nullopt;
-  auto allocOp = value.getDefiningOp<memref::AllocOp>();
-  if (!allocOp || !allocOp->getParentOfType<pim::PimCoreBatchOp>())
-    return std::nullopt;
-  return lane;
-}
-
-static MemoryValueKey getMemoryValueKey(mlir::Value value, std::optional<unsigned> lane = std::nullopt) {
-  return {value, getLaneForMemoryValue(value, lane)};
-}
-
 struct MemoryTouchInterval {
  uint64_t start = 0;
  uint64_t end = 0;
@@ -51,7 +46,6 @@ struct MemoryTouchInterval {
  bool endUsedFallback = false;
  bool escapesLoop = false;
  std::string fallbackReason;
-  llvm::SmallVector<std::string, 8> aliasesFollowed;
 };

 struct OperationOrdering {
@@ -60,50 +54,6 @@ struct OperationOrdering {
  uint64_t nextPosition = 0;
 };

-static std::string printValueToString(mlir::Value value) {
-  std::string text;
-  llvm::raw_string_ostream os(text);
-  value.print(os);
-  os.flush();
-  return text;
-}
-
-static std::string printOperationToString(Operation* op) {
-  if (!op)
-    return "<none>";
-  std::string text;
-  llvm::raw_string_ostream os(text);
-  op->print(os);
-  os.flush();
-  return text;
-}
-
-static std::string printLocationToString(Location loc) {
-  std::string text;
-  llvm::raw_string_ostream os(text);
-  loc.print(os);
-  os.flush();
-  return text;
-}
-
-static std::string collapseWhitespace(StringRef text) {
-  std::string out;
-  out.reserve(text.size());
-  bool lastWasSpace = false;
-  for (char c : text) {
-    bool isSpace = c == ' ' || c == '\n' || c == '\t' || c == '\r';
-    if (isSpace) {
-      if (!lastWasSpace && !out.empty())
-        out.push_back(' ');
-      lastWasSpace = true;
-      continue;
-    }
-    out.push_back(c);
-    lastWasSpace = false;
-  }
-  return out;
-}
-
 static std::string abbreviate(StringRef text, size_t maxLen) {
  if (text.size() <= maxLen)
    return text.str();
@@ -111,21 +61,23 @@ static std::string abbreviate(StringRef text, size_t maxLen) {
 }

 static std::string summarizeValue(mlir::Value value, size_t maxLen = 72) {
-  return abbreviate(collapseWhitespace(printValueToString(value)), maxLen);
+  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>";
-  std::string prefix = op->getName().getStringRef().str();
-  std::string full = collapseWhitespace(printOperationToString(op));
-  if (full == prefix)
-    return prefix;
-  return abbreviate(prefix + " :: " + full, maxLen);
-}
-
-static std::string summarizeLocation(Location loc, size_t maxLen = 88) {
-  return abbreviate(collapseWhitespace(printLocationToString(loc)), maxLen);
+  return abbreviate(op->getName().getStringRef(), maxLen);
 }

 static void assignOperationOrdering(Operation* op, OperationOrdering& ordering) {
@@ -151,13 +103,6 @@ static OperationOrdering buildOperationOrdering(Operation* coreLikeOp) {
  return ordering;
 }

-static bool isSupportedAliasOp(Operation* op) {
-  return isa<memref::SubViewOp,
-             memref::CastOp,
-             memref::CollapseShapeOp,
-             memref::ExpandShapeOp>(op);
-}
-
 static bool isRuntimeMemoryTouchOp(Operation* op) {
  return isa<pim::PimMemCopyHostToDevOp,
             pim::PimMemCopyDevToHostOp,
@@ -180,7 +125,8 @@ static bool isRuntimeMemoryTouchOp(Operation* op) {
 }

 static bool isIgnoredLivenessUser(Operation* op) {
-  return isSupportedAliasOp(op) || isa<scf::ForOp, scf::YieldOp, memref::DeallocOp>(op) || isCoreStaticAddressOp(op);
+  return pim::isLocalMemoryAliasOp(op) || isa<scf::ForOp, scf::YieldOp, memref::DeallocOp>(op)
+      || isCoreStaticAddressOp(op);
 }

 static bool isWithin(mlir::Value value, Region* region) {
@@ -207,12 +153,6 @@ static void addFallbackReason(std::string& reason, StringRef newReason) {
  reason += newReason.str();
 }

-static void appendAliasDescription(llvm::SmallVectorImpl<std::string>& aliases, mlir::Value value) {
-  std::string text = printValueToString(value);
-  if (!llvm::is_contained(aliases, text))
-    aliases.push_back(std::move(text));
-}
-
 struct OrderedTouchRange {
  uint64_t start = 0;
  uint64_t end = 0;
@@ -233,97 +173,28 @@ getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const Operati
  return range;
 }

-static MemoryTouchInterval
-computeMemoryTouchInterval(memref::AllocOp allocOp,
-                           const OperationOrdering& ordering,
-                           uint64_t fallbackEnd,
-                           bool includeAliasDescriptions) {
+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 recordAlias = [&](mlir::Value value) {
-    if (includeAliasDescriptions)
-      appendAliasDescription(interval.aliasesFollowed, value);
-  };

-  SmallPtrSet<mlir::Value, 16> visitedValues;
-  SmallPtrSet<Operation*, 32> visitedUsers;
-  SmallVector<mlir::Value> pendingValues;
-  pendingValues.push_back(allocOp.getResult());
  auto parentLoop = allocOp->getParentOfType<scf::ForOp>();
-
-  while (!pendingValues.empty()) {
-    mlir::Value value = pendingValues.pop_back_val();
-    if (!visitedValues.insert(value).second)
-      continue;
-
-    for (Operation* user : value.getUsers()) {
-      if (!visitedUsers.insert(user).second)
-        continue;
-
-      if (isSupportedAliasOp(user)) {
-        for (mlir::Value result : user->getResults()) {
-          pendingValues.push_back(result);
-          recordAlias(result);
-        }
-      }
-
-      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user)) {
-        for (OpResult result : user->getResults()) {
-          OpOperand* tiedOperand = dpsOp.getTiedOpOperand(result);
-          if (!tiedOperand || tiedOperand->get() != value)
-            continue;
-          pendingValues.push_back(result);
-          recordAlias(result);
-        }
-      }
-
-      if (auto forOp = dyn_cast<scf::ForOp>(user)) {
-        for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs())) {
-          if (initArg != value)
-            continue;
-          pendingValues.push_back(forOp.getRegionIterArgs()[index]);
-          pendingValues.push_back(forOp.getResult(index));
-          recordAlias(forOp.getRegionIterArgs()[index]);
-          recordAlias(forOp.getResult(index));
-          if (parentLoop && forOp != parentLoop)
-            interval.escapesLoop = true;
-        }
-      }
-
+  (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 (ifOp) {
-          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
-            if (operand != value)
-              continue;
-            pendingValues.push_back(ifOp.getResult(index));
-            recordAlias(ifOp.getResult(index));
-          }
-        }
-        else if (indexSwitch) {
-          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
-            if (operand != value)
-              continue;
-            pendingValues.push_back(indexSwitch.getResult(index));
-            recordAlias(indexSwitch.getResult(index));
-          }
-        }
-        else if (!forOp) {
+        if (!forOp && !ifOp && !indexSwitch)
          addFallbackReason(interval.fallbackReason, "yield without scf.for parent");
-        }
-        else {
-          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
-            if (operand != value)
-              continue;
-            pendingValues.push_back(forOp.getResult(index));
-            recordAlias(forOp.getResult(index));
-            if (parentLoop && forOp == parentLoop)
-              interval.escapesLoop = true;
-          }
-        }
+        else if (forOp && parentLoop && forOp == parentLoop && llvm::is_contained(yieldOp.getOperands(), value))
+          interval.escapesLoop = true;
      }

      if (isRuntimeMemoryTouchOp(user)) {
@@ -345,23 +216,21 @@ computeMemoryTouchInterval(memref::AllocOp allocOp,
          interval.hasRuntimeUse = true;
        }
        else {
-          if (range.start < interval.start) {
+          if (range.start < interval.start)
            interval.start = range.start;
-          }
-          if (range.end > interval.end) {
+          if (range.end > interval.end)
            interval.end = range.end;
-          }
        }
-        continue;
+        return success();
      }

      if (isIgnoredLivenessUser(user))
-        continue;
+        return success();

      addFallbackReason(interval.fallbackReason, "unhandled user op");
      interval.endUsedFallback = true;
-    }
-  }
+      return success();
+    });

  if (!interval.hasRuntimeUse) {
    interval.startUsedAllocFallback = true;
@@ -374,9 +243,8 @@ computeMemoryTouchInterval(memref::AllocOp allocOp,
    return interval;
  }

-  if (interval.endUsedFallback) {
+  if (interval.endUsedFallback)
    interval.end = std::max(interval.end, fallbackEnd);
-  }

  return interval;
 }
@@ -402,11 +270,21 @@ static uint64_t getSlotLogicalBytes(const PlannedPhysicalSlot& slot, ArrayRef<Lo
  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,
-                                                                       std::optional<unsigned> lane,
-                                                                       bool includeAliasDescriptions) {
+SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp) {
  SmallVector<LocalAllocInterval, 0> intervals;
  OperationOrdering ordering = buildOperationOrdering(coreLikeOp);
  if (ordering.position.empty())
@@ -423,12 +301,10 @@ SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation
      llvm_unreachable("Failed to compute local allocation size");
    }

-    MemoryTouchInterval touchInterval =
-      computeMemoryTouchInterval(allocOp, ordering, fallbackEnd, includeAliasDescriptions);
+    MemoryTouchInterval touchInterval = computeMemoryTouchInterval(allocOp, ordering, fallbackEnd);
    LocalAllocInterval interval;
    interval.id = nextIntervalId++;
    interval.alloc = allocOp;
-    interval.key = getMemoryValueKey(allocOp.getResult(), lane);
    interval.start = touchInterval.start;
    interval.end = touchInterval.end;
    interval.size = *checkedSize;
@@ -442,7 +318,9 @@ SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation
    interval.insideNestedRegion = isNestedAllocation(coreLikeOp, allocOp);
    interval.escapesLoop = touchInterval.escapesLoop;
    interval.fallbackReason = std::move(touchInterval.fallbackReason);
-    interval.aliasesFollowed = std::move(touchInterval.aliasesFollowed);
+    interval.valueSummary = summarizeValue(allocOp.getResult(), 88);
+    interval.firstTouchSummary = summarizeOperation(touchInterval.firstTouchOp);
+    interval.lastTouchSummary = summarizeOperation(touchInterval.lastTouchOp);
    intervals.push_back(std::move(interval));
  });

@@ -467,60 +345,94 @@ SmallVector<PlannedPhysicalSlot, 0> onnx_mlir::planPhysicalSlots(MutableArrayRef

  for (size_t intervalIndex : intervalOrder) {
    LocalAllocInterval& interval = intervals[intervalIndex];
-    PlannedPhysicalSlot* bestSlot = nullptr;
-    auto bestKey = std::tuple<size_t, size_t, size_t, size_t>(std::numeric_limits<size_t>::max(),
-                                                              std::numeric_limits<size_t>::max(),
-                                                              std::numeric_limits<size_t>::max(),
-                                                              std::numeric_limits<size_t>::max());
-
-    for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
-      PlannedPhysicalSlot& slot = slots[slotIndex];
-      bool compatible = true;
-      for (size_t otherIndex : slot.intervalIndices) {
-        if (intervalsOverlap(interval, intervals[otherIndex])) {
-          compatible = false;
-          break;
-        }
+    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);
      }
-      if (!compatible)
-        continue;
+    }

-      size_t resultingSize = std::max(slot.requiredSize, interval.size);
-      size_t growth = resultingSize - slot.requiredSize;
-      auto candidateKey =
-        std::tuple<size_t, size_t, size_t, size_t>(growth, resultingSize, slot.intervalIndices.size(), slot.id);
+    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;
-        bestSlot = &slot;
+        bestAddress = candidate;
      }
    }
+    assert(bestAddress != std::numeric_limits<size_t>::max() && "address after all conflicts must be available");

-    if (!bestSlot) {
-      slots.push_back({slots.size(), interval.size, interval.size, 0, {intervalIndex}});
-      interval.slotPlanIndex = slots.size() - 1;
-      interval.physicalSlotId = slots.back().id;
-      interval.physicalSlotSize = slots.back().requiredSize;
-      continue;
+    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;
    }
-
-    bestSlot->requiredSize = std::max(bestSlot->requiredSize, interval.size);
-    bestSlot->size = bestSlot->requiredSize;
-    bestSlot->intervalIndices.push_back(intervalIndex);
-    interval.slotPlanIndex = static_cast<size_t>(bestSlot - slots.data());
-    interval.physicalSlotId = bestSlot->id;
-    interval.physicalSlotSize = bestSlot->requiredSize;
  }

  return slots;
 }

-MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp,
-                                                        std::optional<unsigned> lane,
-                                                        ArrayRef<LocalAllocInterval> intervals,
-                                                        ArrayRef<PlannedPhysicalSlot> slots,
-                                                        size_t addressLimit,
-                                                        PimMemoryReportLevel reportLevel) {
-  MemoryPlanArtifacts artifacts;
+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;
@@ -536,7 +448,6 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp,
  for (const LocalAllocInterval& interval : intervals) {
    totalLogicalBytes += interval.size;
    largestLogicalAllocation = std::max(largestLogicalAllocation, interval.size);
-    maximumAssignedAddress = std::max(maximumAssignedAddress, interval.assignedAddress + interval.physicalSlotSize);
    if (interval.startUsedAllocFallback || interval.endUsedFallback)
      ++fallbackIntervals;
    if (!interval.hasRuntimeUse)
@@ -546,12 +457,14 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp,
    if (interval.escapesLoop)
      ++loopEscapingIntervals;
  }
-  for (const PlannedPhysicalSlot& slot : slots) {
-    totalPhysicalBytes += slot.size;
-    largestPhysicalSlot = std::max(largestPhysicalSlot, slot.size);
-    if (slot.intervalIndices.size() > 1)
-      reusedAllocations += slot.intervalIndices.size() - 1;
+  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 =
@@ -560,8 +473,6 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp,
  raw_string_ostream os(artifacts);
  os << "=== PIM Memory Liveness Report ===\n";
  os << "Op: " << coreLikeOp->getName() << "\n";
-  if (lane)
-    os << "Lane: " << *lane << "\n";
  os << "Summary:\n";
  os << "  logical allocation bytes:  " << formatReportMemory(totalLogicalBytes) << " (" << totalLogicalBytes << ")\n";
  os << "  physical allocation bytes: " << formatReportMemory(totalPhysicalBytes) << " (" << totalPhysicalBytes
@@ -569,32 +480,27 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp,
  os << "  saved bytes:               " << formatReportMemory(savedBytes) << " (" << savedBytes << ")\n";
  os << "  saved percent: " << format("%.2f%%", savedPercent) << "\n";
  os << "  intervals: " << intervals.size() << "\n";
-  os << "  physical slots: " << slots.size() << "\n";
-  os << "  reused allocations: " << reusedAllocations << "\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 physical slot: " << largestPhysicalSlot << "\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";

-  os << "\nHow To Read:\n";
-  os << "  `summary` only shows the strongest reuse cases and the worst offenders.\n";
-  os << "  Use `--pim-memory-report=full` when you need the complete slot-by-slot and interval-by-interval dump.\n";
-  os << "  Large single-use slots, fallback intervals, and nested single-use allocations are the best places\n";
-  os << "  to inspect if allocations should be moved, sunk, or made easier to coalesce earlier in the pipeline.\n";
-
  SmallVector<const PlannedPhysicalSlot*> reusedSlots;
  SmallVector<const PlannedPhysicalSlot*> singleUseSlots;
-  for (const PlannedPhysicalSlot& slot : slots)
-    if (slot.intervalIndices.size() > 1)
-      reusedSlots.push_back(&slot);
+  for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
+    if (hasAddressReuse(slotIndex, slots))
+      reusedSlots.push_back(&slots[slotIndex]);
    else
-      singleUseSlots.push_back(&slot);
+      singleUseSlots.push_back(&slots[slotIndex]);
+  }

  llvm::stable_sort(reusedSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) {
    uint64_t lhsLogicalBytes = getSlotLogicalBytes(*lhs, intervals);
@@ -603,140 +509,66 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp,
      return lhs->intervalIndices.size() > rhs->intervalIndices.size();
    if (lhsLogicalBytes != rhsLogicalBytes)
      return lhsLogicalBytes > rhsLogicalBytes;
-    if (lhs->size != rhs->size)
-      return lhs->size > rhs->size;
+    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->size != rhs->size)
-      return lhs->size > rhs->size;
+    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 Reuse:\n";
+  os << "\nBest Address Reuse:\n";
  if (reusedSlots.empty()) {
-    os << "  no slots were shared by multiple intervals\n";
+      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->size)
+      os << "  slot #" << slot->id << "  addr=" << slot->address << "  size=" << formatReportMemory(slot->requiredSize)
         << "  intervals=" << slot->intervalIndices.size() << "  logical_sum=" << formatReportMemory(slotLogicalBytes)
         << "\n";
-      for (size_t intervalIndex : slot->intervalIndices) {
-        const LocalAllocInterval& interval = intervals[intervalIndex];
-        os << "    #" << interval.id << "  [" << interval.start << "," << interval.end << "]"
-           << "  logical=" << formatReportMemory(interval.size)
-           << "  first=" << summarizeOperation(interval.firstTouchOp, 40)
-           << "  last=" << summarizeOperation(interval.lastTouchOp, 40) << "\n";
-      }
    }
  }

  os << "\nTop Offenders:\n";
-  bool printedAttention = false;
  for (const PlannedPhysicalSlot* slot : ArrayRef(singleUseSlots).take_front(kSummaryOffenderLimit)) {
    const LocalAllocInterval& interval = intervals[slot->intervalIndices.front()];
-    printedAttention = true;
    os << "  slot #" << slot->id << " is single-use"
-       << "  size=" << formatReportMemory(slot->size) << "  interval=#" << interval.id
-       << "  value=" << summarizeValue(interval.key.value, 56) << "\n";
-    os << "    first=" << summarizeOperation(interval.firstTouchOp, 40)
-       << "  last=" << summarizeOperation(interval.lastTouchOp, 40)
+       << "  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";
  }
-  size_t fallbackPrinted = 0;
-  for (const LocalAllocInterval& interval : intervals) {
-    if (!(interval.startUsedAllocFallback || interval.endUsedFallback) || fallbackPrinted >= kSummaryOffenderLimit)
-      continue;
-    printedAttention = true;
-    ++fallbackPrinted;
-    os << "  fallback interval #" << interval.id << "  size=" << formatReportMemory(interval.size)
-       << "  value=" << summarizeValue(interval.key.value, 56) << "\n";
-    os << "    reason: " << (interval.fallbackReason.empty() ? "<none>" : interval.fallbackReason) << "\n";
-  }
-  size_t nestedPrinted = 0;
-  for (const LocalAllocInterval& interval : intervals) {
-    if (nestedPrinted >= kSummaryOffenderLimit)
-      break;
-    if (!(interval.insideNestedRegion && slots[interval.slotPlanIndex].intervalIndices.size() == 1))
-      continue;
-    printedAttention = true;
-    ++nestedPrinted;
-    os << "  nested single-use interval #" << interval.id << "  slot #" << interval.physicalSlotId
-       << "  size=" << formatReportMemory(interval.size) << "  value=" << summarizeValue(interval.key.value, 56)
-       << "\n";
-    os << "    hint: move or sink this alloc inside the nested region if the IR allows it.\n";
-  }
-  if (!printedAttention)
+  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.size) << " ("
-         << slot.size << ")"
+      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];
-        mlir::Value allocValue = interval.key.value;
        os << "    [" << interval.start << "," << interval.end << "]"
-           << "  #" << interval.id << "  logical=" << formatReportMemory(interval.size)
+           << "  #" << interval.id << "  logical=" << formatReportMemory(interval.size) << " (" << interval.size
+           << ")"
           << "  nested=" << (interval.insideNestedRegion ? "yes" : "no")
           << "  escapes_loop=" << (interval.escapesLoop ? "yes" : "no")
-           << "  first=" << summarizeOperation(interval.firstTouchOp, 48)
-           << "  last=" << summarizeOperation(interval.lastTouchOp, 48) << "\n";
-        os << "      value=" << summarizeValue(allocValue) << "\n";
-      }
-    }
-  }
-
-  if (reportLevel == PimMemoryReportFull) {
-    os << "\nInterval Details:\n";
-    for (const LocalAllocInterval& interval : intervals) {
-      const PlannedPhysicalSlot& slot = slots[interval.slotPlanIndex];
-      mlir::Value allocValue = interval.key.value;
-      Operation* definingOp = allocValue.getDefiningOp();
-      os << "  #" << interval.id << "  slot=" << slot.id << "  live=[" << interval.start << "," << interval.end << "]"
-         << "  logical=" << formatReportMemory(interval.size)
-         << "  slot_size=" << formatReportMemory(interval.physicalSlotSize) << "  addr=" << interval.assignedAddress
-         << "\n";
-      os << "    value=" << summarizeValue(allocValue, 88) << "\n";
-      os << "    type=" << allocValue.getType() << "\n";
-      os << "    loc="
-         << summarizeLocation(definingOp ? definingOp->getLoc() : UnknownLoc::get(coreLikeOp->getContext())) << "\n";
-      os << "    nested=" << (interval.insideNestedRegion ? "yes" : "no")
-         << "  escapes_loop=" << (interval.escapesLoop ? "yes" : "no")
-         << "  start_fallback=" << (interval.startUsedAllocFallback ? "yes" : "no")
-         << "  end_fallback=" << (interval.endUsedFallback ? "yes" : "no") << "\n";
-      os << "    first_use=" << summarizeOperation(interval.firstTouchOp) << " @" << interval.firstTouchPosition
-         << "\n";
-      os << "    last_use=" << summarizeOperation(interval.lastTouchOp) << " @" << interval.lastTouchPosition << "\n";
-      os << "    slot_peers=";
-      bool first = true;
-      for (size_t otherIndex : slot.intervalIndices) {
-        if (intervals[otherIndex].id == interval.id)
-          continue;
-        if (!first)
-          os << ", ";
-        os << "#" << intervals[otherIndex].id;
-        first = false;
-      }
-      if (first)
-        os << "<none>";
-      os << "\n";
-      if (!interval.fallbackReason.empty())
-        os << "    fallback_reason=" << interval.fallbackReason << "\n";
-      if (!interval.aliasesFollowed.empty()) {
-        os << "    aliases_followed=" << interval.aliasesFollowed.size() << "\n";
-        for (const std::string& alias : interval.aliasesFollowed)
-          os << "      - " << abbreviate(collapseWhitespace(alias), 108) << "\n";
+           << "  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";
      }
    }
  }
@@ -744,3 +576,95 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp,

  return artifacts;
 }
+
+namespace onnx_mlir {
+namespace {
+
+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";
+  }
+
+  void runOnOperation() override {
+    std::fstream reportFile;
+    std::unique_ptr<raw_os_ostream> report;
+    if (pimMemoryReport != PimMemoryReportNone) {
+      reportFile = openReportFileWithExtension("pim_memory_liveness_report", "txt");
+      if (reportFile.is_open())
+        report = std::make_unique<raw_os_ostream>(reportFile);
+    }
+    else if (std::string outputRoot = getOutputDir(); !outputRoot.empty()) {
+      sys::fs::remove(outputRoot + "/reports/pim_memory_liveness_report.txt");
+    }
+
+    Builder builder(&getContext());
+    uint64_t nextBatchId = 0;
+    bool failedPlanning = false;
+    getOperation().walk([&](Operation* op) {
+      if (failedPlanning || !isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
+        return;
+
+      CoreMemoryPlan plan = buildCoreMemoryPlan(op);
+      if (failed(assignPhysicalSlotAddresses(plan, kPimLocalMemoryAddressLimit))) {
+        failedPlanning = true;
+        return;
+      }
+
+      size_t arenaSize = 0;
+      for (const PlannedPhysicalSlot& placement : plan.slots)
+        arenaSize = std::max(arenaSize, placement.address + placement.requiredSize);
+      for (const LocalAllocInterval& interval : plan.intervals) {
+        const PlannedPhysicalSlot& slot = plan.slots[interval.slotPlanIndex];
+        interval.alloc->setAttr(kLocalMemoryAddressAttrName, builder.getI64IntegerAttr(slot.address));
+        interval.alloc->setAttr(kLocalMemorySlotAttrName, builder.getI64IntegerAttr(0));
+        interval.alloc->setAttr(kLocalMemorySlotSizeAttrName, builder.getI64IntegerAttr(arenaSize));
+      }
+      uint64_t fallbackIntervals = llvm::count_if(plan.intervals, [](const LocalAllocInterval& interval) {
+        return interval.startUsedAllocFallback || interval.endUsedFallback;
+      });
+      uint64_t nestedSingleUseIntervals = llvm::count_if(plan.intervals, [&](const LocalAllocInterval& interval) {
+        return interval.insideNestedRegion && !hasAddressReuse(interval.slotPlanIndex, plan.slots);
+      });
+      op->setAttr(kLocalMemoryFallbackCountAttrName, builder.getI64IntegerAttr(fallbackIntervals));
+      op->setAttr(kLocalMemoryNestedSingleUseCountAttrName, builder.getI64IntegerAttr(nestedSingleUseIntervals));
+
+      if (!report)
+        return;
+      std::string planReport = buildMemoryPlanReport(
+        op, plan.intervals, plan.slots, kPimLocalMemoryAddressLimit, pimMemoryReport);
+      if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
+        *report << "Core " << coreOp.getCoreId() << ":\n";
+        *report << planReport;
+        return;
+      }
+
+      SmallVector<int32_t> coreIds = getBatchCoreIds(cast<pim::PimCoreBatchOp>(op));
+      llvm::sort(coreIds);
+      coreIds.erase(std::unique(coreIds.begin(), coreIds.end()), coreIds.end());
+      uint64_t batchId = nextBatchId++;
+      for (int32_t coreId : coreIds)
+        *report << "Batch " << batchId << " core " << coreId << ":\n" << planReport;
+    });
+
+    if (report) {
+      report->flush();
+      reportFile.close();
+    }
+    if (failedPlanning) {
+      signalPassFailure();
+      return;
+    }
+    dumpModule(getOperation(), "pim4_memory_planned");
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createPimLocalMemoryPlanningPass() {
+  return std::make_unique<PimLocalMemoryPlanningPass>();
+}
+
+} // namespace onnx_mlir
@@ -6,10 +6,8 @@
 #include "llvm/ADT/SmallVector.h"

 #include <limits>
-#include <optional>
 #include <string>

-#include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"

 namespace onnx_mlir {
@@ -17,7 +15,6 @@ namespace onnx_mlir {
 struct LocalAllocInterval {
  size_t id = 0;
  mlir::memref::AllocOp alloc;
-  MemoryValueKey key;
  uint64_t start = 0;
  uint64_t end = 0;
  size_t size = 0;
@@ -31,32 +28,37 @@ struct LocalAllocInterval {
  bool insideNestedRegion = false;
  bool escapesLoop = false;
  std::string fallbackReason;
-  llvm::SmallVector<std::string, 8> aliasesFollowed;
+  std::string valueSummary;
+  std::string firstTouchSummary;
+  std::string lastTouchSummary;
  size_t slotPlanIndex = std::numeric_limits<size_t>::max();
-  size_t physicalSlotId = std::numeric_limits<size_t>::max();
-  size_t assignedAddress = 0;
-  size_t physicalSlotSize = 0;
 };

 struct PlannedPhysicalSlot {
  size_t id = std::numeric_limits<size_t>::max();
  size_t requiredSize = 0;
-  size_t size = 0;
  size_t address = 0;
  llvm::SmallVector<size_t, 8> intervalIndices;
 };

-llvm::SmallVector<LocalAllocInterval, 0> buildLocalAllocIntervals(mlir::Operation* coreLikeOp,
-                                                                  std::optional<unsigned> lane,
-                                                                  bool includeAliasDescriptions = true);
+struct CoreMemoryPlan {
+  mlir::Operation* coreLikeOp = nullptr;
+  llvm::SmallVector<LocalAllocInterval, 0> intervals;
+  llvm::SmallVector<PlannedPhysicalSlot, 0> slots;
+};
+
+llvm::SmallVector<LocalAllocInterval, 0> buildLocalAllocIntervals(mlir::Operation* coreLikeOp);

 llvm::SmallVector<PlannedPhysicalSlot, 0> planPhysicalSlots(llvm::MutableArrayRef<LocalAllocInterval> intervals);

-MemoryPlanArtifacts buildMemoryPlanArtifacts(mlir::Operation* coreLikeOp,
-                                             std::optional<unsigned> lane,
-                                             llvm::ArrayRef<LocalAllocInterval> intervals,
-                                             llvm::ArrayRef<PlannedPhysicalSlot> slots,
-                                             size_t addressLimit,
-                                             PimMemoryReportLevel reportLevel);
+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
@@ -10,5 +10,6 @@ add_pim_library(OMPimMemoryCoalescing

  LINK_LIBS PUBLIC
  OMPimCommon
+  OMPimLocalMemoryLifetimeAnalysis
  PimOps
 )
@@ -1,15 +1,13 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/Interfaces/DestinationStyleOpInterface.h"
-
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallPtrSet.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;
@@ -19,10 +17,6 @@ namespace pim {

 namespace {

-static bool isSupportedAliasOp(Operation* op) {
-  return isa<memref::SubViewOp, memref::CastOp, memref::CollapseShapeOp, memref::ExpandShapeOp>(op);
-}
-
 static bool isCandidateAllocType(MemRefType type) {
  return type && type.hasStaticShape() && type.getLayout().isIdentity()
      && hasByteSizedElementType(type.getElementType());
@@ -44,59 +38,21 @@ 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);
-  SmallPtrSet<Value, 16> visitedValues;
-  SmallPtrSet<Operation*, 16> visitedUsers;
-  SmallVector<Value> pendingValues;
-  pendingValues.push_back(allocOp.getResult());
-
-  while (!pendingValues.empty()) {
-    Value value = pendingValues.pop_back_val();
-    if (!visitedValues.insert(value).second)
-      continue;
-
-    for (Operation* user : value.getUsers()) {
-      if (!visitedUsers.insert(user).second)
-        continue;
-
-      if (isSupportedAliasOp(user))
-        llvm::append_range(pendingValues, user->getResults());
-
-      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user)) {
-        for (OpResult result : user->getResults()) {
-          OpOperand* tiedOperand = dpsOp.getTiedOpOperand(result);
-          if (tiedOperand && tiedOperand->get() == value)
-            pendingValues.push_back(result);
-        }
-      }
-
-      if (auto forOp = dyn_cast<scf::ForOp>(user)) {
-        for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs())) {
-          if (initArg != value)
-            continue;
-          pendingValues.push_back(forOp.getRegionIterArgs()[index]);
-          pendingValues.push_back(forOp.getResult(index));
-        }
-      }
-
+  if (failed(walkLocalMemoryUses(allocOp.getResult(), [&](Value, Operation* user) {
      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
-        auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
-        if (!forOp)
+        if (!isa<scf::ForOp>(yieldOp->getParentOp()))
          return failure();
-        for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands()))
-          if (operand == value)
-            pendingValues.push_back(forOp.getResult(index));
      }
-
      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;
 }
@@ -133,7 +89,7 @@ static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis) {
    }

    blockAnalysis.candidates.push_back(
-      AllocationCandidate {allocOp, &block, opOrder.lookup(allocOp), *endInstruction, getTypeSizeBytes(allocType)});
+      AllocationCandidate {allocOp, opOrder.lookup(allocOp), *endInstruction, getTypeSizeBytes(allocType)});
  }

  analysis.skippedAllocations += blockAnalysis.skippedAllocations;
@@ -150,6 +106,14 @@ uint64_t MemoryCoalescingAnalysis::getCandidateCount() const {
  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())
@@ -10,7 +10,6 @@ namespace pim {

 struct AllocationCandidate {
  mlir::memref::AllocOp alloc;
-  mlir::Block* scopeBlock = nullptr;
  uint64_t startInstruction = 0;
  uint64_t endInstruction = 0;
  uint64_t sizeBytes = 0;
@@ -27,6 +26,7 @@ struct MemoryCoalescingAnalysis {
  uint64_t skippedAllocations = 0;

  uint64_t getCandidateCount() const;
+  uint64_t getCandidateBytes() const;
 };

 struct MemoryCoalescingStats {
@@ -1,207 +1,94 @@
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"

-#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/raw_os_ostream.h"

 #include <fstream>

-#include "Common/IR/CompactAsmUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-#include "src/Accelerators/PIM/Common/Support/DebugDump.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;
-using namespace onnx_mlir::compact_asm;

 namespace onnx_mlir {
 namespace {

-// This pass is an IR cleanup step after bufferization. It only rewrites
-// obviously compatible local allocations with non-overlapping lifetimes inside
-// the same block and leaves the final physical memory planning to codegen.
-
 struct CoalescingReportRow {
-  uint64_t numCandidates = 0;
-  uint64_t numSkipped = 0;
-  uint64_t numRemoved = 0;
+  uint64_t candidates = 0;
+  uint64_t logicalBytes = 0;
+  uint64_t skipped = 0;
+  uint64_t removed = 0;
  uint64_t savedBytes = 0;
-
-  bool operator==(const CoalescingReportRow& other) const {
-    return numCandidates == other.numCandidates && numSkipped == other.numSkipped && numRemoved == other.numRemoved
-        && savedBytes == other.savedBytes;
-  }
 };

 struct CoalescingReportEntry {
-  enum class Kind {
-    Core,
-    Batch
-  };
-
-  Kind kind = Kind::Core;
-  uint64_t id = 0;
-  llvm::SmallVector<int32_t, 8> coreIds;
+  std::string label;
+  uint64_t coreCount = 1;
  CoalescingReportRow row;
 };

-static std::string formatMemory(uint64_t bytes) { return formatReportMemory(bytes); }
-
-static void printReportRow(raw_ostream& os, const CoalescingReportRow& row) {
-  llvm::SmallVector<ReportField, 4> fields = {
-    {"Number of candidates", std::to_string(row.numCandidates)},
-    {"Skipped allocations",  std::to_string(row.numSkipped)   },
-    {"Removed allocations",  std::to_string(row.numRemoved)   },
-    {"Saved memory",         formatMemory(row.savedBytes)     }
-  };
-  printReportFlatFields(os, fields);
-}
-
-static CoalescingReportRow getTotalRow(const CoalescingReportEntry& entry) {
-  uint64_t factor = std::max<uint64_t>(1, entry.coreIds.size());
-  return {entry.row.numCandidates * factor,
-          entry.row.numSkipped * factor,
-          entry.row.numRemoved * factor,
-          entry.row.savedBytes * factor};
-}
-
 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 totalRow;
+  CoalescingReportRow total;
  for (const CoalescingReportEntry& entry : entries) {
-    CoalescingReportRow entryTotal = getTotalRow(entry);
-    totalRow.numCandidates += entryTotal.numCandidates;
-    totalRow.numSkipped += entryTotal.numSkipped;
-    totalRow.numRemoved += entryTotal.numRemoved;
-    totalRow.savedBytes += entryTotal.savedBytes;
+    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;
  }
-
-  llvm::SmallVector<ReportField, 4> totalFields = {
-    {"Number of candidates", std::to_string(totalRow.numCandidates)},
-    {"Skipped allocations",  std::to_string(totalRow.numSkipped)   },
-    {"Removed allocations",  std::to_string(totalRow.numRemoved)   },
-    {"Saved memory",         formatMemory(totalRow.savedBytes)     }
-  };
-  printReportTotalsBlock(os, totalFields);
-  if (!entries.empty())
-    os << "\n";
-
-  llvm::SmallVector<CoalescingReportEntry, 32> sortedEntries(entries.begin(), entries.end());
-  sortReportEntriesByFirstCore(sortedEntries);
-
-  for (size_t index = 0; index < sortedEntries.size();) {
-    size_t runEnd = index + 1;
-    while (runEnd < sortedEntries.size() && sortedEntries[runEnd].kind == sortedEntries[index].kind
-           && sortedEntries[runEnd].row == sortedEntries[index].row) {
-      ++runEnd;
-    }
-
-    if (sortedEntries[index].kind == CoalescingReportEntry::Kind::Batch) {
-      os << "Batch ";
-      for (size_t batchIndex = index; batchIndex < runEnd; ++batchIndex) {
-        if (batchIndex != index)
-          os << ",\n      ";
-        os << sortedEntries[batchIndex].id << " (cores ";
-        printCompressedIntegerEntries(os, ArrayRef<int32_t>(sortedEntries[batchIndex].coreIds));
-        os << ")";
-      }
-    }
-    else {
-      llvm::SmallVector<int32_t, 8> coreIds;
-      for (size_t coreIndex = index; coreIndex < runEnd; ++coreIndex)
-        coreIds.push_back(sortedEntries[coreIndex].coreIds.front());
-      os << "Core ";
-      printCompressedIntegerEntries(os, ArrayRef<int32_t>(coreIds));
-    }
-
-    os << ":\n";
-    if (sortedEntries[index].kind == CoalescingReportEntry::Kind::Batch) {
-      llvm::SmallVector<ReportField, 4> perCoreFields = {
-        {"Number of candidates", std::to_string(sortedEntries[index].row.numCandidates)},
-        {"Skipped allocations",  std::to_string(sortedEntries[index].row.numSkipped)   },
-        {"Removed allocations",  std::to_string(sortedEntries[index].row.numRemoved)   },
-        {"Saved memory",         formatMemory(sortedEntries[index].row.savedBytes)     }
-      };
-      CoalescingReportRow totalRow = getTotalRow(sortedEntries[index]);
-      llvm::SmallVector<ReportField, 4> totalFields = {
-        {"Number of candidates", std::to_string(totalRow.numCandidates)},
-        {"Skipped allocations",  std::to_string(totalRow.numSkipped)   },
-        {"Removed allocations",  std::to_string(totalRow.numRemoved)   },
-        {"Saved memory",         formatMemory(totalRow.savedBytes)     }
-      };
-      printReportPerCoreAndTotalFields(os, perCoreFields, totalFields);
-    }
-    else {
-      printReportRow(os, sortedEntries[index].row);
-    }
-    printReportEntrySeparator(os, runEnd < sortedEntries.size());
-    index = runEnd;
+  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();
-  file.close();
 }

 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 "Analyze local PIM memory reuse opportunities"; }
-
-  PimMemoryCoalescingPass() = default;
-  PimMemoryCoalescingPass(const PimMemoryCoalescingPass& pass) {}
+  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;
-    bool hasFailure = false;
-
    getOperation().walk([&](Operation* op) {
-      if (hasFailure || !isa<pim::PimCoreOp, pim::PimCoreBatchOp>(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(), stats.skippedAllocations, stats.removedAllocs, stats.savedBytes};
-
+      CoalescingReportRow row {analysis.getCandidateCount(),
+                               analysis.getCandidateBytes(),
+                               stats.skippedAllocations,
+                               stats.removedAllocs,
+                               stats.savedBytes};
      if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
-        auto checkedCoreId =
-          pim::checkedI32(static_cast<uint64_t>(coreOp.getCoreId()), coreOp, "memory coalescing core id");
-        if (failed(checkedCoreId)) {
-          hasFailure = true;
-          return;
-        }
-        reportEntries.push_back(
-          {CoalescingReportEntry::Kind::Core, static_cast<uint64_t>(coreOp.getCoreId()), {*checkedCoreId}, row});
+        reportEntries.push_back({"Core " + std::to_string(coreOp.getCoreId()), 1, row});
        return;
      }
-
-      auto coreIds = getBatchCoreIds(cast<pim::PimCoreBatchOp>(op));
-      CoalescingReportEntry entry;
-      entry.kind = CoalescingReportEntry::Kind::Batch;
-      entry.id = nextBatchId++;
-      llvm::append_range(entry.coreIds, coreIds);
-      entry.row = row;
-      reportEntries.push_back(std::move(entry));
+      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});
    });
-
-    if (hasFailure) {
-      signalPassFailure();
-      return;
-    }
-
    emitReport(reportEntries);
    dumpModule(getOperation(), "pim3_coalesced");
  }
@@ -15,8 +15,9 @@
 #include "src/Accelerators/PIM/Common/IR/SubviewUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#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/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -108,6 +109,63 @@ static bool isCoreWeightBlockArgument(Value value) {
  return false;
 }

+struct VerifiedLocalMemorySlot {
+  uint64_t size = 0;
+};
+
+static LogicalResult
+verifyLocalMemoryPlan(Operation* coreLikeOp, pim::CappedDiagnosticReporter& diagnostics) {
+  DenseMap<uint64_t, VerifiedLocalMemorySlot> slots;
+  bool hasFailure = false;
+  auto fallbackAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryFallbackCountAttrName);
+  auto nestedSingleUseAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryNestedSingleUseCountAttrName);
+  if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) {
+    diagnostics.report(coreLikeOp, [&](Operation* op) {
+      op->emitError("requires complete non-negative PIM local-memory planning summary attributes");
+    });
+    hasFailure = true;
+  }
+  coreLikeOp->walk([&](memref::AllocOp allocOp) {
+    auto addressAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
+    auto slotAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotAttrName);
+    auto slotSizeAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotSizeAttrName);
+    if (!addressAttr || !slotAttr || !slotSizeAttr || addressAttr.getInt() < 0 || slotAttr.getInt() < 0
+        || slotSizeAttr.getInt() < 0) {
+      diagnostics.report(allocOp, [&](Operation*) {
+        allocOp.emitOpError("requires a complete non-negative PIM local-memory plan");
+      });
+      hasFailure = true;
+      return;
+    }
+
+    uint64_t address = static_cast<uint64_t>(addressAttr.getInt());
+    uint64_t slotId = static_cast<uint64_t>(slotAttr.getInt());
+    uint64_t slotSize = static_cast<uint64_t>(slotSizeAttr.getInt());
+    auto logicalSize = pim::getCheckedShapedTypeSizeInBytes(
+      cast<ShapedType>(allocOp.getType()), allocOp, "planned local allocation byte size");
+    bool invalidRange = failed(logicalSize) || address > slotSize || *logicalSize > slotSize - address
+                     || slotSize > kPimLocalMemoryAddressLimit || address % 4 != 0 || slotId != 0;
+    if (invalidRange) {
+      diagnostics.report(allocOp, [&](Operation*) {
+        allocOp.emitOpError() << "has an invalid PIM local-memory range: address=" << address
+                              << ", logical size=" << (succeeded(logicalSize) ? *logicalSize : 0)
+                              << ", slot size=" << slotSize;
+      });
+      hasFailure = true;
+      return;
+    }
+
+    auto [slotIt, inserted] = slots.try_emplace(slotId, VerifiedLocalMemorySlot {slotSize});
+    if (!inserted && slotIt->second.size != slotSize) {
+      diagnostics.report(allocOp, [&](Operation*) {
+        allocOp.emitOpError() << "has inconsistent size for PIM local-memory arena " << slotId;
+      });
+      hasFailure = true;
+    }
+  });
+  return success(!hasFailure);
+}
+
 static bool isSupportedCoreInstructionOp(Operation* op) {
  return isa<pim::PimMemCopyHostToDevOp,
             pim::PimMemCopyDevToHostOp,
@@ -148,8 +206,10 @@ static bool isHostAddressableValue(Value value, const StaticValueKnowledge& know
  return isa_and_nonnull<memref::GetGlobalOp>(base.getDefiningOp());
 }

-
-enum class CommunicationEventKind { Send, Receive };
+enum class CommunicationEventKind {
+  Send,
+  Receive
+};

 struct CommunicationEvent {
  CommunicationEventKind kind = CommunicationEventKind::Send;
@@ -157,19 +217,6 @@ struct CommunicationEvent {
  int64_t peerCoreId = 0;
  int64_t size = 0;
  uint64_t ordinal = 0;
-  std::optional<int64_t> minChannelId;
-  std::string materializer;
-  std::optional<int64_t> traceId;
-  std::optional<int64_t> commOrder;
-  std::optional<int64_t> traceClassId;
-  std::optional<int64_t> traceBlockOrdinal;
-  std::string traceKind;
-  std::string tracePhase;
-  std::string traceClassKind;
-  std::string tracePayload;
-  std::string traceMessages;
-  std::string tracePrevOp;
-  std::string traceNextOp;
  Operation* op = nullptr;
 };

@@ -183,7 +230,6 @@ constexpr StringLiteral kRaptorMinChannelIdAttr = "raptor.min_channel_id";
 constexpr StringLiteral kRaptorMaterializerAttr = "raptor.materializer";
 constexpr StringLiteral kRaptorCommOrderAttr = "raptor.comm_order";
 constexpr StringLiteral kRaptorCommTraceIdAttr = "raptor.comm_trace_id";
-constexpr StringLiteral kRaptorCommTraceKindAttr = "raptor.comm_trace_kind";
 constexpr StringLiteral kRaptorCommTracePhaseAttr = "raptor.comm_trace_phase";
 constexpr StringLiteral kRaptorCommTraceClassIdAttr = "raptor.comm_trace_class_id";
 constexpr StringLiteral kRaptorCommTraceClassKindAttr = "raptor.comm_trace_class_kind";
@@ -225,31 +271,44 @@ static std::string formatOperationSummary(Operation* op) {
 }

 static std::string formatCommunicationEvent(const CommunicationEvent& event) {
+  std::optional<int64_t> minChannelId = getNearestIntegerAttr(event.op, kRaptorMinChannelIdAttr);
+  std::optional<int64_t> commOrder = getNearestIntegerAttr(event.op, kRaptorCommOrderAttr);
+  std::optional<int64_t> traceId = getNearestIntegerAttr(event.op, kRaptorCommTraceIdAttr);
+  std::optional<int64_t> traceClassId = getNearestIntegerAttr(event.op, kRaptorCommTraceClassIdAttr);
+  std::optional<int64_t> traceBlockOrdinal = getNearestIntegerAttr(event.op, kRaptorCommTraceBlockOrdinalAttr);
+  std::string materializer = getNearestStringAttr(event.op, kRaptorMaterializerAttr);
+  std::string tracePhase = getNearestStringAttr(event.op, kRaptorCommTracePhaseAttr);
+  std::string traceClassKind = getNearestStringAttr(event.op, kRaptorCommTraceClassKindAttr);
+  std::string tracePayload = getNearestStringAttr(event.op, kRaptorCommTracePayloadAttr);
+  std::string traceMessages = getNearestStringAttr(event.op, kRaptorCommTraceMessagesAttr);
+  std::string tracePrevOp = getNearestStringAttr(event.op, kRaptorCommTracePrevOpAttr);
+  std::string traceNextOp = getNearestStringAttr(event.op, kRaptorCommTraceNextOpAttr);
+
  std::string text;
  llvm::raw_string_ostream os(text);
  os << "core " << event.coreId << " " << getCommunicationEventKindName(event.kind) << " "
-     << (event.kind == CommunicationEventKind::Send ? "to" : "from") << " " << event.peerCoreId
-     << " size " << event.size << "B ordinal " << event.ordinal;
-  if (event.minChannelId)
-    os << " min_channel " << *event.minChannelId;
-  if (event.commOrder)
-    os << " comm_order " << *event.commOrder;
-  if (!event.materializer.empty())
-    os << " materializer " << event.materializer;
-  if (event.traceId)
-    os << " trace#" << *event.traceId;
-  if (!event.tracePhase.empty())
-    os << " phase " << event.tracePhase;
-  if (event.traceClassId)
-    os << " class " << event.traceClassKind << "#" << *event.traceClassId;
-  if (event.traceBlockOrdinal)
-    os << " block_ordinal " << *event.traceBlockOrdinal;
-  if (!event.tracePayload.empty())
-    os << " payload " << event.tracePayload;
-  if (!event.traceMessages.empty())
-    os << " messages {" << event.traceMessages << "}";
-  if (!event.tracePrevOp.empty() || !event.traceNextOp.empty())
-    os << " inserted_between [" << event.tracePrevOp << " | " << event.traceNextOp << "]";
+     << (event.kind == CommunicationEventKind::Send ? "to" : "from") << " " << event.peerCoreId << " size "
+     << event.size << "B ordinal " << event.ordinal;
+  if (minChannelId)
+    os << " min_channel " << *minChannelId;
+  if (commOrder)
+    os << " comm_order " << *commOrder;
+  if (!materializer.empty())
+    os << " materializer " << materializer;
+  if (traceId)
+    os << " trace#" << *traceId;
+  if (!tracePhase.empty())
+    os << " phase " << tracePhase;
+  if (traceClassId)
+    os << " class " << traceClassKind << "#" << *traceClassId;
+  if (traceBlockOrdinal)
+    os << " block_ordinal " << *traceBlockOrdinal;
+  if (!tracePayload.empty())
+    os << " payload " << tracePayload;
+  if (!traceMessages.empty())
+    os << " messages {" << traceMessages << "}";
+  if (!tracePrevOp.empty() || !traceNextOp.empty())
+    os << " inserted_between [" << tracePrevOp << " | " << traceNextOp << "]";
  return os.str();
 }

@@ -261,10 +320,8 @@ static bool areMatchedCommunicationEvents(const CommunicationEvent& lhs, const C
      || (lhs.kind == CommunicationEventKind::Receive && rhs.kind == CommunicationEventKind::Send);
 }

-
-static std::optional<size_t> findMatchingCounterpartIndex(const CommunicationEventVector& events,
-                                                           const CommunicationEvent& event,
-                                                           size_t begin) {
+static std::optional<size_t>
+findMatchingCounterpartIndex(const CommunicationEventVector& events, const CommunicationEvent& event, size_t begin) {
  for (size_t index = begin; index < events.size(); ++index)
    if (areMatchedCommunicationEvents(event, events[index]))
      return index;
@@ -288,8 +345,7 @@ static void printCounterpartProbe(llvm::raw_ostream& os,
    peerPc = peerPcIt->second;

  os << "  counterpart probe for " << formatCommunicationEvent(blockedEvent) << "\n";
-  os << "    peer core " << blockedEvent.peerCoreId << " current pc " << peerPc << " of " << peerEvents.size()
-     << "\n";
+  os << "    peer core " << blockedEvent.peerCoreId << " current pc " << peerPc << " of " << peerEvents.size() << "\n";

  std::optional<size_t> nextMatch = findMatchingCounterpartIndex(peerEvents, blockedEvent, peerPc);
  std::optional<size_t> anyMatch = findMatchingCounterpartIndex(peerEvents, blockedEvent, 0);
@@ -317,7 +373,8 @@ static void printCounterpartProbe(llvm::raw_ostream& os,
    return;

  os << "    peer operations blocking before that counterpart:\n";
-  size_t end = std::min(peerEvents.size(), std::min(*nextMatch + static_cast<size_t>(1), peerPc + static_cast<size_t>(12)));
+  size_t end =
+    std::min(peerEvents.size(), std::min(*nextMatch + static_cast<size_t>(1), peerPc + static_cast<size_t>(12)));
  for (size_t index = peerPc; index < end; ++index) {
    os << (index == peerPc ? "      pc => " : "           ") << "#" << index << " "
       << formatCommunicationEvent(peerEvents[index]) << "\n";
@@ -327,77 +384,58 @@ static void printCounterpartProbe(llvm::raw_ostream& os,
    os << "      ... " << (*nextMatch - end + 1) << " more peer communication event(s) before the counterpart\n";
 }

-static CommunicationEvent makeCommunicationEvent(CommunicationEventKind kind,
-                                                 int64_t coreId,
-                                                 int64_t peerCoreId,
-                                                 int64_t size,
-                                                 uint64_t ordinal,
-                                                 Operation* op) {
-  return CommunicationEvent {kind,
-                             coreId,
-                             peerCoreId,
-                             size,
-                             ordinal,
-                             getNearestIntegerAttr(op, kRaptorMinChannelIdAttr),
-                             getNearestStringAttr(op, kRaptorMaterializerAttr),
-                             getNearestIntegerAttr(op, kRaptorCommTraceIdAttr),
-                             getNearestIntegerAttr(op, kRaptorCommOrderAttr),
-                             getNearestIntegerAttr(op, kRaptorCommTraceClassIdAttr),
-                             getNearestIntegerAttr(op, kRaptorCommTraceBlockOrdinalAttr),
-                             getNearestStringAttr(op, kRaptorCommTraceKindAttr),
-                             getNearestStringAttr(op, kRaptorCommTracePhaseAttr),
-                             getNearestStringAttr(op, kRaptorCommTraceClassKindAttr),
-                             getNearestStringAttr(op, kRaptorCommTracePayloadAttr),
-                             getNearestStringAttr(op, kRaptorCommTraceMessagesAttr),
-                             getNearestStringAttr(op, kRaptorCommTracePrevOpAttr),
-                             getNearestStringAttr(op, kRaptorCommTraceNextOpAttr),
-                             op};
+static CommunicationEvent makeCommunicationEvent(
+  CommunicationEventKind kind, int64_t coreId, int64_t peerCoreId, int64_t size, uint64_t ordinal, Operation* op) {
+  return CommunicationEvent {kind, coreId, peerCoreId, size, ordinal, op};
 }

 static LogicalResult appendCoreCommunicationEvents(Block& block,
+                                                   const PimCoreCommunicationPlan& plan,
                                                   int64_t coreId,
                                                   const StaticValueKnowledge& initialKnowledge,
                                                   SmallVectorImpl<CommunicationEvent>& events,
                                                   pim::CappedDiagnosticReporter& diagnostics) {
-  return walkPimCoreBlock(block, initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
-    if (auto sendOp = dyn_cast<pim::PimSendOp>(&op)) {
-      auto targetCoreId = resolveIndexValue(sendOp.getTargetCoreId(), knowledge);
-      if (failed(targetCoreId)) {
-        diagnostics.report(&op, [](Operation* illegalOp) {
-          illegalOp->emitOpError("cannot statically resolve send target core for PIM communication deadlock check");
-        });
-        return failure();
+  return walkPimCoreCommunicationBlock(
+    block, plan, initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
+      if (auto sendOp = dyn_cast<pim::PimSendOp>(&op)) {
+        auto targetCoreId = resolveIndexValue(sendOp.getTargetCoreId(), knowledge);
+        if (failed(targetCoreId)) {
+          diagnostics.report(&op, [](Operation* illegalOp) {
+            illegalOp->emitOpError("cannot statically resolve send target core for PIM communication deadlock check");
+          });
+          return failure();
+        }
+
+        events.push_back(makeCommunicationEvent(CommunicationEventKind::Send,
+                                                coreId,
+                                                *targetCoreId,
+                                                sendOp.getSize(),
+                                                static_cast<uint64_t>(events.size()),
+                                                &op));
+        return success();
      }

-      events.push_back(makeCommunicationEvent(CommunicationEventKind::Send,
-                                              coreId,
-                                              *targetCoreId,
-                                              sendOp.getSize(),
-                                              static_cast<uint64_t>(events.size()),
-                                              &op));
-      return success();
-    }
+      if (auto receiveOp = dyn_cast<pim::PimReceiveOp>(&op)) {
+        auto sourceCoreId = resolveIndexValue(receiveOp.getSourceCoreId(), knowledge);
+        if (failed(sourceCoreId)) {
+          diagnostics.report(&op, [](Operation* illegalOp) {
+            illegalOp->emitOpError(
+              "cannot statically resolve receive source core for PIM communication deadlock check");
+          });
+          return failure();
+        }

-    if (auto receiveOp = dyn_cast<pim::PimReceiveOp>(&op)) {
-      auto sourceCoreId = resolveIndexValue(receiveOp.getSourceCoreId(), knowledge);
-      if (failed(sourceCoreId)) {
-        diagnostics.report(&op, [](Operation* illegalOp) {
-          illegalOp->emitOpError("cannot statically resolve receive source core for PIM communication deadlock check");
-        });
-        return failure();
+        events.push_back(makeCommunicationEvent(CommunicationEventKind::Receive,
+                                                coreId,
+                                                *sourceCoreId,
+                                                receiveOp.getSize(),
+                                                static_cast<uint64_t>(events.size()),
+                                                &op));
+        return success();
      }

-      events.push_back(makeCommunicationEvent(CommunicationEventKind::Receive,
-                                              coreId,
-                                              *sourceCoreId,
-                                              receiveOp.getSize(),
-                                              static_cast<uint64_t>(events.size()),
-                                              &op));
      return success();
-    }
-
-    return success();
-  });
+    });
 }

 static void printCommunicationWindow(llvm::raw_ostream& os,
@@ -466,9 +504,10 @@ static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
                                           ArrayRef<int64_t> cycle) {
  printCommunicationDeadlockReport(coreEvents, programCounters, cycle);

-  auto diagnostic = moduleOp.emitError()
-      << "PIM communication deadlock check found a blocking send/receive cycle while statically simulating the "
-         "expanded per-core communication streams; see the PIM static communication deadlock report above";
+  auto diagnostic =
+    moduleOp.emitError()
+    << "PIM communication deadlock check found a blocking send/receive cycle while statically simulating the "
+       "expanded per-core communication streams; see the PIM static communication deadlock report above";

  for (int64_t coreId : cycle) {
    auto eventsIt = coreEvents.find(coreId);
@@ -479,14 +518,15 @@ static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
    const CommunicationEvent& event = eventsIt->second[pcIt->second];
    Diagnostic& note = diagnostic.attachNote(event.op->getLoc());
    note << formatCommunicationEvent(event);
-    if (!event.materializer.empty())
-      note << " emitted by " << event.materializer;
+    std::string materializer = getNearestStringAttr(event.op, kRaptorMaterializerAttr);
+    if (!materializer.empty())
+      note << " emitted by " << materializer;
  }
 }

-static FailureOr<SmallVector<int64_t>> findCommunicationWaitCycle(
-  const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
-  const DenseMap<int64_t, size_t>& programCounters) {
+static FailureOr<SmallVector<int64_t>>
+findCommunicationWaitCycle(const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
+                           const DenseMap<int64_t, size_t>& programCounters) {
  for (const auto& [startCoreId, events] : coreEvents) {
    auto startPcIt = programCounters.find(startCoreId);
    if (startPcIt == programCounters.end() || startPcIt->second >= events.size())
@@ -519,7 +559,7 @@ static FailureOr<SmallVector<int64_t>> findCommunicationWaitCycle(
 }

 static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
-                                                          pim::CappedDiagnosticReporter& diagnostics) {
+                                                         pim::CappedDiagnosticReporter& diagnostics) {
  DenseMap<int64_t, CommunicationEventVector> coreEvents;
  bool hasFailure = false;

@@ -530,14 +570,20 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
    for (Operation& op : funcOp.getBody().front().getOperations()) {
      if (auto coreOp = dyn_cast<pim::PimCoreOp>(&op)) {
        int64_t coreId = coreOp.getCoreId();
-        if (failed(appendCoreCommunicationEvents(
-              coreOp.getBody().front(), coreId, StaticValueKnowledge {}, coreEvents[coreId], diagnostics)))
+        PimCoreCommunicationPlan plan = buildPimCoreCommunicationPlan(coreOp.getBody().front());
+        if (failed(appendCoreCommunicationEvents(coreOp.getBody().front(),
+                                                 plan,
+                                                 coreId,
+                                                 StaticValueKnowledge {},
+                                                 coreEvents[coreId],
+                                                 diagnostics)))
          hasFailure = true;
        continue;
      }

      if (auto coreBatchOp = dyn_cast<pim::PimCoreBatchOp>(&op)) {
        SmallVector<int32_t> coreIds = getBatchCoreIds(coreBatchOp);
+        PimCoreCommunicationPlan plan = buildPimCoreCommunicationPlan(coreBatchOp.getBody().front());
        size_t laneCount = static_cast<size_t>(coreBatchOp.getLaneCount());
        for (size_t lane = 0; lane < laneCount; ++lane) {
          StaticValueKnowledge laneKnowledge;
@@ -546,11 +592,14 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
            laneKnowledge.aliases[coreBatchOp.getInputArgument(inputIndex)] = coreBatchOp.getInputs()[inputIndex];

          SmallVector<int32_t> laneCoreIds = getLaneChunkCoreIds(coreIds, laneCount, static_cast<unsigned>(lane));
-          for (int32_t coreId : laneCoreIds) {
-            if (failed(appendCoreCommunicationEvents(
-                  coreBatchOp.getBody().front(), coreId, laneKnowledge, coreEvents[coreId], diagnostics)))
+          for (int32_t coreId : laneCoreIds)
+            if (failed(appendCoreCommunicationEvents(coreBatchOp.getBody().front(),
+                                                     plan,
+                                                     coreId,
+                                                     laneKnowledge,
+                                                     coreEvents[coreId],
+                                                     diagnostics)))
              hasFailure = true;
-          }
        }
      }
    }
@@ -607,9 +656,10 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
      return failure();
    }

-    auto diagnostic = moduleOp.emitError()
-        << "PIM communication deadlock check stalled without finding a closed wait cycle; this usually means a "
-           "send/receive peer is missing or ordered after a finished core";
+    auto diagnostic =
+      moduleOp.emitError()
+      << "PIM communication deadlock check stalled without finding a closed wait cycle; this usually means a "
+         "send/receive peer is missing or ordered after a finished core";
    for (const auto& [coreId, events] : coreEvents) {
      size_t pc = programCounters[coreId];
      if (pc >= events.size())
@@ -652,6 +702,7 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
      for (Operation& op : funcOp.getBody().front().getOperations()) {
        if (auto coreOp = dyn_cast<pim::PimCoreOp>(&op)) {
          (void) verifyCoreWeights(moduleOp, coreOp, diagnostics);
+          (void) verifyLocalMemoryPlan(coreOp, diagnostics);
          StaticValueKnowledge knowledge;
          (void) verifyCoreLikeOperands(coreOp, knowledge, diagnostics);
          continue;
@@ -659,6 +710,7 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>

        if (auto coreBatchOp = dyn_cast<pim::PimCoreBatchOp>(&op)) {
          (void) verifyCoreWeights(moduleOp, coreBatchOp, diagnostics);
+          (void) verifyLocalMemoryPlan(coreBatchOp, diagnostics);
          llvm::SmallVector<unsigned, 2> lanes;
          lanes.push_back(0);
          if (coreBatchOp.getLaneCount() > 1)
@@ -133,6 +133,22 @@ 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;
+  CrossbarUsage distinctWeights;
+  for (const ComputeGraphNode& node : graph.nodes) {
+    for (const CrossbarWeight& weight : node.crossbarUsage) {
+      if (!containsCrossbarWeight(distinctWeights, weight))
+        distinctWeights.push_back(weight);
+      if (distinctWeights.size() > threshold)
+        return true;
+    }
+  }
+  return false;
+}
+
 } // namespace

 Time getPeftTransferTime(Time transferCost, size_t sourceProcessor, size_t targetProcessor, size_t processorCount) {
@@ -145,6 +161,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
  if (processorCount == 0)
    llvm::report_fatal_error("PEFT scheduler: processor count must be positive");
  MeshModel mesh = MeshModel::infer(processorCount);
+  const bool preferCrossbarReuse = hasHighCrossbarPressure(graph, processorCount, options.crossbarCapacity);

  verifyOctTableSize(nodeCount, processorCount);
  std::vector<std::vector<size_t>> reverseLevels = buildReverseLevels(graph);
@@ -282,6 +299,8 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
      Time eft = addOrMax(est, computeCost);
      Time oeft = addOrMax(eft, oct[task * processorCount + processor]);
      size_t centerDistance = mesh.getCenterDistance(processor);
+      bool betterCrossbarChoice =
+        preferCrossbarReuse ? overlapCount > bestOverlapCount : overlapCount < bestOverlapCount;

      if (oeft < bestOeft || (oeft == bestOeft && eft < bestEft)
          || (oeft == bestOeft && eft == bestEft && est < bestEst)) {
@@ -302,7 +321,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
        bestCenterDistance = centerDistance;
      }
      else if (oeft == bestOeft && eft == bestEft && est == bestEst
-               && centerDistance == bestCenterDistance && overlapCount < bestOverlapCount) {
+               && centerDistance == bestCenterDistance && betterCrossbarChoice) {
        bestProcessor = processor;
        bestEst = est;
        bestEft = eft;
@@ -23,6 +23,8 @@ std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass();

 std::unique_ptr<mlir::Pass> createPimHostConstantFoldingPass();

+std::unique_ptr<mlir::Pass> createPimLocalMemoryPlanningPass();
+
 std::unique_ptr<mlir::Pass> createPimVerificationPass();

 std::unique_ptr<mlir::Pass> createEmitPimCodePass();
@@ -80,6 +80,7 @@ void PimAccelerator::registerPasses(int optLevel) const {
  registerPass(createTrivialGraphComputeMergePass);
  registerPass(createMergeComputeNodesPass);
  registerPass(createPimHostConstantFoldingPass);
+  registerPass(createPimLocalMemoryPlanningPass);
  registerPass(createPimVerificationPass);
  registerPass(createEmitPimCodePass);
 }