After merge

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

@@ -0,0 +1,8 @@
+add_pim_library(OMPimLocalMemoryLifetimeAnalysis
+  LocalMemoryLifetimeAnalysis.cpp
+
+  EXCLUDE_FROM_OM_LIBS
+
+  INCLUDE_DIRS PUBLIC
+  ${PIM_PUBLIC_INCLUDE_DIRS}
+)

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

@@ -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

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

@@ -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

src/PIM/Dialect/Pim/CMakeLists.txt

@@ -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

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

@@ -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
+)

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

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

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

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

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

@@ -10,5 +10,6 @@ add_pim_library(OMPimMemoryCoalescing
 
   LINK_LIBS PUBLIC
   OMPimCommon
+  OMPimLocalMemoryLifetimeAnalysis
   PimOps
 )

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

@@ -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())

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

@@ -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 {

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

@@ -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");
   }

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

@@ -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)

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

@@ -135,10 +135,27 @@ 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;
+}
+
 std::vector<CrossbarUsage> planCrossbarReservations(const ComputeGraph& graph,
                                                     size_t processorCount,
                                                     size_t crossbarCapacity,
-                                                    const MeshModel& mesh) {
+                                                    const MeshModel& mesh,
+                                                    bool preferCrossbarReuse) {
   std::vector<size_t> weightedTasks;
   for (size_t task = 0; task < graph.nodes.size(); ++task)
     if (!graph.nodes[task].crossbarUsage.empty())
@@ -161,8 +178,8 @@ std::vector<CrossbarUsage> planCrossbarReservations(const ComputeGraph& graph,
       if (crossbarUnion > crossbarCapacity)
         continue;
       size_t addedCrossbars = crossbarUnion - reservations[processor].size();
-      ReservationScore score {reservedLoad[processor],
-                              addedCrossbars,
+      ReservationScore score {preferCrossbarReuse ? addedCrossbars : reservedLoad[processor],
+                              preferCrossbarReuse ? reservedLoad[processor] : addedCrossbars,
                               mesh.getCenterDistance(processor),
                               processor};
       if (!bestScore || score < *bestScore) {
@@ -199,8 +216,9 @@ 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);
   std::vector<CrossbarUsage> capacityReservations =
-    planCrossbarReservations(graph, processorCount, options.crossbarCapacity, mesh);
+    planCrossbarReservations(graph, processorCount, options.crossbarCapacity, mesh, preferCrossbarReuse);
 
   verifyOctTableSize(nodeCount, processorCount);
   std::vector<std::vector<size_t>> reverseLevels = buildReverseLevels(graph);
@@ -339,6 +357,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)) {
@@ -359,7 +379,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;