blazingly fasterer

.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md

@@ -8,10 +8,10 @@ Every developer or coding agent modifying Spatial graph construction, graph
 verification, Blueprint handling, or `MergeComputeNodes` must read this file
 after `README.md` and `AGENTS.md`.
 
-`AGENTS.md` must contain this instruction:
+`AGENTS.md` must reference this invariant:
 
 ```text
-* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes.
+* `.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md`
 ```
 
 ## Scope

.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md

@@ -0,0 +1,67 @@
+# Compiler Performance Optimization Invariant
+
+## Scope
+
+This invariant applies to:
+
+- ONNX-to-Spatial lowering;
+- graph and trivial merge;
+- PEFT scheduling;
+- scheduled materialization;
+- deferred communication;
+- Spatial-to-PIM lowering;
+- bufferization;
+- memory coalescing;
+- liveness planning;
+- PIM code generation.
+
+## Invariant
+
+A compiler compile-time or memory optimization must not reduce available
+hardware parallelism or worsen theoretical execution time. In the absence of a
+precise runtime model, it is forbidden to serialize independent work, reduce
+the number of simultaneously schedulable compute instances, replace parallel
+graph or core batches with sequential loops, introduce recomputation, add
+runtime copies, increase communication or instruction count, or lengthen the
+schedule critical path merely to reduce compiler time or memory usage.
+
+An optimization is acceptable only when its static runtime proxies are equal or
+better.
+
+## Forbidden optimization trades
+
+Do not:
+
+- use fewer parallel lanes or cores to simplify IR;
+- scalarize a valid graph or core batch;
+- replace independent operations with one sequential loop;
+- recompute values to reduce storage;
+- add device/device or host/device copies to reduce peak memory;
+- change hardware parameters to make a benchmark easier;
+- reduce communication concurrency by imposing a global serial order;
+- increase emitted instruction count to reduce compiler analysis;
+- move PIM work to the host.
+
+## Required proof
+
+Every performance change must report before and after:
+
+- compiler wall time;
+- compiler peak RSS;
+- IR operation and value counts at the changed stage;
+- logical compute-instance count;
+- graph compute batch and scheduled batch lane counts;
+- PEFT class and core assignment;
+- maximum scheduled critical-path or step count;
+- MVM/VMM and vector instruction counts;
+- send and receive counts and bytes;
+- total emitted instruction count;
+- maximum instructions assigned to one core;
+- number and size of runtime memory copies;
+- final host, weights, logical-core, physical-core, and maximum-core memory;
+- numerical validation.
+
+## Stop rule
+
+If compile time or memory improves but a runtime proxy worsens, stop and reject
+the change.

AGENTS.md

@@ -1,5 +1,11 @@
 * Always read the full README.md before doing anything
-* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes.
+
+# Required project invariants
+
+Before modifying the relevant subsystem, read:
+
+* `.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md`
+* `.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md`
 * Build commands:
   * `cmake --build ./build_release`
   * `cmake --build ./build_debug`
@@ -165,6 +171,10 @@ Do not refactor when:
 * If a change adds a walk, cache, analysis, or structural traversal, justify why it is needed
 * For hot paths, prefer preserving existing asymptotic behavior unless a better structure is part of the requested change
 * If performance may change, mention the expected impact and suggest a targeted timing check
+* Compile-time and compiler-memory improvements must not trade away hardware
+  parallelism or theoretical runtime. Preserve or improve the static runtime
+  proxies defined in
+  `.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md`.
 
 # Goal-driven execution
 

README.md

@@ -67,20 +67,22 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    Converts tensor-semantics PIM IR into memref-semantics PIM IR using MLIR's
    bufferization interfaces.
 
-5. **Static memory coalescing**
+5. **Safe IR memory coalescing**
-   (`src/PIM/Dialect/Pim/Transforms/StaticMemoryCoalescing`).
+   (`src/PIM/Dialect/Pim/Transforms/MemoryCoalescing`).
-   Reuses compatible local memref allocations inside PIM cores before codegen.
+   Normalizes compatible block-local allocations before final planning.
-
+6. **PIM local-memory planning**
-6. **PIM code generation** (`src/PIM/Pass/PimCodegen` and
+   (`src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning`).
+   Computes whole-core lifetimes, coalesces non-overlapping allocations into
+   physical slots, assigns addresses, and records the explicit plan in PIM IR.
+7. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and
    `src/PIM/Compiler`).
-   Folds host constants, materializes remaining host constants, verifies PIM IR,
+   Verifies the memory plan and other PIM invariants, then emits `.pim` core
-   emits `.pim` core files, writes weights, and writes `memory.bin` /
+   files, weights, and `memory.bin` / `config.json` without rerunning liveness.
-   `config.json`.
 
 Supporting pieces:
 - `src/PIM/Common` - shared IR, filesystem, diagnostics, reports, and utility
   helpers.
-- `src/PIM/Compiler` - PIM compiler options, memory/address planning, binary
+- `src/PIM/Compiler` - PIM compiler options, planned-address materialization, binary
   instruction format, artifact writing, weight emission, and codegen entry
   points.
 - `src/PIM/Conversion/SpatialToGraphviz` - optional Spatial graphviz conversion
@@ -97,8 +99,8 @@ Pass these to `onnx-mlir` when compiling for PIM:
   `--EmitPimCodegen` - stop the PIM pipeline at the requested stage. The PIM
   default is `--EmitPimCodegen`.
 - `--core-count=<N>` - required positive core count for PIM compilation.
-- `--crossbar-size=<N>` - crossbar width/height. Default in code is `2`.
+- `--crossbar-size=<N>` - crossbar width/height. Default in code is `128`.
-- `--crossbar-count=<N>` - crossbars per core. Default in code is `256`.
+- `--crossbar-count=<N>` - crossbars per core. Default in code is `64`.
 - `--pim-only-codegen` - assume input is already bufferized PIM IR and only run
   the codegen tail.
 - `--pim-emit-json` - also emit `core_*.json` instruction files alongside
@@ -109,12 +111,28 @@ Pass these to `onnx-mlir` when compiling for PIM:
 - `--use-experimental-conv-impl` - use the alternate convolution lowering.
 - `--ignore-concat-error` - soft-fail a ConcatOp corner case.
 
+## Standard PIM hardware profile
+
+Raptor's standard development and YOLO validation profile is:
+
+| Parameter | Value |
+| --- | ---: |
+| Cores | 144 |
+| Crossbars per core | 64 |
+| Crossbar size | 128 × 128 |
+
+Canonical compiler flags:
+
+`--crossbar-count=64 --crossbar-size=128 --core-count=144`
+
+`--core-count` remains mandatory and must be passed explicitly to the compiler.
+
 Example:
 
 ```bash
 ./build_release/Release/bin/onnx-mlir model.onnx -o /tmp/raptor/model \
   --maccel=PIM --EmitPimCodegen \
-  --crossbar-size=2048 --crossbar-count=256 --core-count=1000
+  --crossbar-count=64 --crossbar-size=128 --core-count=144
 ```
 
 This writes PIM artifacts under `/tmp/raptor/pim/`.
@@ -131,21 +149,33 @@ Python dependencies used by the validation scripts are `numpy`, `onnx`, and
 Per-operation validation from the repository root:
 
 ```bash
-python3 validation/validate.py \
+python validation/validate.py \
   --raptor-path build_release/Release/bin/onnx-mlir \
   --onnx-include-dir onnx-mlir/include \
-  --core-count 1000
+  --operations-dir validation/operations \
+  --crossbar-count 64 \
+  --crossbar-size 128 \
+  --core-count 144 \
+  --verbose \
+  --raptor-extra-arg=--pim-detect-communication-deadlock \
+  --raptor-extra-arg=--pim-export-spatial-dataflow=none
 ```
 
 Validate one network or a subset by pointing `--operations-dir` at any directory
 containing `.onnx` files:
 
 ```bash
-python3 validation/validate.py \
+python validation/validate.py \
   --raptor-path build_release/Release/bin/onnx-mlir \
   --onnx-include-dir onnx-mlir/include \
   --operations-dir validation/networks/yolo11n/depth_04 \
-  --crossbar-size 2048 --crossbar-count 256 --core-count 1000
+  --crossbar-count 64 \
+  --crossbar-size 128 \
+  --core-count 144 \
+  --verbose \
+  --raptor-extra-arg=--pim-memory-report=summary \
+  --raptor-extra-arg=--pim-detect-communication-deadlock \
+  --raptor-extra-arg=--pim-export-spatial-dataflow=none
 ```
 
 Useful validation options:
@@ -170,7 +200,7 @@ Each validation run writes artifacts in the model workspace, for example under
 The compiler currently dumps dialect snapshots such as `spatial0.mlir`,
 `spatial1_graph.mlir`, `spatial2_trivial_merged.mlir`,
 `spatial3_scheduled_no_comm.mlir`, `spatial4_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
-`pim2_folded.mlir`, and `pim3_coalesced.mlir` when an output directory is
+`pim2_folded.mlir`, `pim3_coalesced.mlir`, and `pim4_memory_planned.mlir` when an output directory is
 available.
 
 To rerun the simulator manually with tracing after validation has produced a

src/PIM/CMakeLists.txt

@@ -120,8 +120,8 @@ add_pim_library(OMPIMAccel
   OMSpatialToPim
   OMPimCommon
   OMPimBufferization
-  OMPimMemoryCoalescing
   OMPimHostConstantFolding
+  OMPimMemoryCoalescing
   OMPimVerification
   MLIRTensorInferTypeOpInterfaceImpl
 )

src/PIM/Common/IR/CoreBlockUtils.cpp

@@ -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::ConstantOp,
-                mlir::arith::SubIOp, mlir::arith::MulIOp,
+                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 };
+enum class CoreWalkMode {
-using CoreWalkCallback =
+  ExecuteCommunication,
-  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)>;
+  StructuralExtremes
+};
+using CoreWalkCallback = llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)>;
 
-static void propagateRegionResults(mlir::ValueRange results,
+static void propagateRegionResults(mlir::ValueRange results, mlir::Region& region, StaticValueKnowledge& knowledge) {
-                                   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";
+  llvm::StringRef purpose = mode == CoreWalkMode::ExecuteCommunication ? "communication verification" : "verification";
-  for (mlir::Operation& op : block) {
+  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) {
+      if (mode == CoreWalkMode::ExecuteCommunication) {
-        hasFailure |= !selected.empty() && failed(walkPimCoreBlockImpl(selected.front(), knowledge, mode, callback));
+        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
+PimCoreCommunicationPlan buildPimCoreCommunicationPlan(mlir::Block& block) {
-walkPimCoreBlock(mlir::Block& block,
+  llvm::DenseSet<mlir::Operation*> communicationAncestors;
-                 const StaticValueKnowledge& knowledge,
+  block.walk([&](mlir::Operation* op) {
-                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
+    if (!mlir::isa<pim::PimSendOp, pim::PimReceiveOp>(op))
-  return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::ExecuteAllIterations, callback);
+      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

src/PIM/Common/IR/CoreBlockUtils.hpp

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

src/PIM/Common/PimCommon.hpp

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

src/PIM/Compiler/CMakeLists.txt

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

src/PIM/Compiler/PimCodeGen.cpp

@@ -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
+  if (address > kPimLocalMemoryAddressLimit || failed(checkedEnd) || *checkedEnd > kPimLocalMemoryAddressLimit
-      || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimAddressLimit) {
+      || 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) {
+void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane) {
-  auto intervals = buildLocalAllocIntervals(op, lane, pimMemoryReport == PimMemoryReportFull);
+  if (localPhysicalSlots.empty()) {
-  SmallVector<PlannedPhysicalSlot> plannedSlots = planPhysicalSlots(intervals);
+    llvm::append_range(localPhysicalSlots, plan.slots);
-
+    reportRow.logicalAllocaBytes = plan.logicalBytes;
-  SmallVector<size_t> slotOrder(plannedSlots.size());
+    reportRow.fallbackIntervals = plan.fallbackIntervals;
-  std::iota(slotOrder.begin(), slotOrder.end(), 0);
+    reportRow.nestedSingleUseIntervals = plan.nestedSingleUseIntervals;
-  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;
-    }
   }
-
+  for (const CompiledLocalMemoryEntry& entry : plan.entries) {
-  if (pimMemoryReport != PimMemoryReportNone) {
+    MemoryValueKey key = getMemoryValueKey(entry.value, lane);
-    MemoryPlanArtifacts artifacts =
+    ownedMemEntriesMap[key] = entry.memory;
-      buildMemoryPlanArtifacts(op, lane, intervals, plannedSlots, kPimAddressLimit, pimMemoryReport);
+    globalMemEntriesMap[key] = entry.memory;
-    livenessArtifacts += artifacts;
   }
 }
 
@@ -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 = {
+  llvm::SmallVector<ReportField, 10> totalFields = {
-    {"Global memory",  formatReportMemory(totalGlobalMemory) },
+    {"Global memory",               formatReportMemory(totalGlobalMemory)               },
-    {"Weights memory", formatReportMemory(totalWeightsMemory)},
+    {"Weights memory",              formatReportMemory(totalWeightsMemory)              },
-    {"Cores memory",   formatReportMemory(totalCoresMemory)  }
+    {"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:
+  case pim_binary::Opcode::sldi:  scalarRegisterValues[instruction.rd] = instruction.r2OrImm; break;
-    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:
+  case pim_binary::Opcode::smuli: scalarRegisterValues[instruction.rd].reset(); break;
-    scalarRegisterValues[instruction.rd].reset();
+  default:                        break;
-    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(
+  genSetRegisterImmediate(pim::checkedU8OrCrash(registerNumber, "register number"),
-    pim::checkedU8OrCrash(registerNumber, "register number"), pim::checkedI32OrCrash(immediate, "register immediate"));
+                          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(
+      emitMemCopyOp(pim_binary::Opcode::lmv, outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len");
-        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(
+  int32_t maxDestinationAddress =
-    pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"), "transpose destination address");
+    pim::checkedI32OrCrash(pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"),
-  (void)maxSourceAddress;
+                           "transpose destination address");
-  (void)maxDestinationAddress;
+  (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(),
+      coreCodeGen.emitBinaryVectorOp(opcode,
-                                     getVectorByteSizeOrCrash(cast<ShapedType>(op.getLhs().getType())), knowledge);
+                                     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:
+    case CompiledCoreOpKind::Transpose: coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge); break;
-      coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge);
+    case CompiledCoreOpKind::VVAdd:     emitBinary(cast<pim::PimVVAddOp>(node.op), pim_binary::Opcode::vvadd); break;
-      break;
+    case CompiledCoreOpKind::VVSub:     emitBinary(cast<pim::PimVVSubOp>(node.op), pim_binary::Opcode::vvsub); break;
-    case CompiledCoreOpKind::VVAdd:  emitBinary(cast<pim::PimVVAddOp>(node.op), pim_binary::Opcode::vvadd); break;
+    case CompiledCoreOpKind::VVMul:     emitBinary(cast<pim::PimVVMulOp>(node.op), pim_binary::Opcode::vvmul); break;
-    case CompiledCoreOpKind::VVSub:  emitBinary(cast<pim::PimVVSubOp>(node.op), pim_binary::Opcode::vvsub); break;
+    case CompiledCoreOpKind::VVMax:     emitBinary(cast<pim::PimVVMaxOp>(node.op), pim_binary::Opcode::vvmax); break;
-    case CompiledCoreOpKind::VVMul:  emitBinary(cast<pim::PimVVMulOp>(node.op), pim_binary::Opcode::vvmul); break;
+    case CompiledCoreOpKind::VVDMul:    emitBinary(cast<pim::PimVVDMulOp>(node.op), pim_binary::Opcode::vvdmul); break;
-    case CompiledCoreOpKind::VVMax:  emitBinary(cast<pim::PimVVMaxOp>(node.op), pim_binary::Opcode::vvmax); break;
+    case CompiledCoreOpKind::VAvg:      emitUnary(cast<pim::PimVAvgOp>(node.op), pim_binary::Opcode::vavg, 1, 1); break;
-    case CompiledCoreOpKind::VVDMul: emitBinary(cast<pim::PimVVDMulOp>(node.op), pim_binary::Opcode::vvdmul); break;
+    case CompiledCoreOpKind::VRelu:     emitUnary(cast<pim::PimVReluOp>(node.op), pim_binary::Opcode::vrelu, 0, 0); break;
-    case CompiledCoreOpKind::VAvg:   emitUnary(cast<pim::PimVAvgOp>(node.op), pim_binary::Opcode::vavg, 1, 1); break;
+    case CompiledCoreOpKind::VTanh:     emitUnary(cast<pim::PimVTanhOp>(node.op), pim_binary::Opcode::vtanh, 0, 0); break;
-    case CompiledCoreOpKind::VRelu:  emitUnary(cast<pim::PimVReluOp>(node.op), pim_binary::Opcode::vrelu, 0, 0); break;
+    case CompiledCoreOpKind::VSigm:     emitUnary(cast<pim::PimVSigmOp>(node.op), pim_binary::Opcode::vsigm, 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 =
+  std::string outputCorePath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
-    (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 =
+  std::string outputCoreJsonPath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
-    (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);
 }

src/PIM/Compiler/PimCodeGen.hpp

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

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -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."),

src/PIM/Compiler/PimCompilerOptions.hpp

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

src/PIM/Compiler/PimCompilerUtils.cpp

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

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

@@ -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();
+  return abbreviate(op->getName().getStringRef(), maxLen);
-  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);
 }
 
 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
+static MemoryTouchInterval computeMemoryTouchInterval(memref::AllocOp allocOp,
-computeMemoryTouchInterval(memref::AllocOp allocOp,
+                                                      const OperationOrdering& ordering,
-                           const OperationOrdering& ordering,
+                                                      uint64_t fallbackEnd) {
-                           uint64_t fallbackEnd,
-                           bool includeAliasDescriptions) {
   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>();
-
+  (void) pim::walkLocalMemoryUses(
-  while (!pendingValues.empty()) {
+    allocOp.getResult(),
-    mlir::Value value = pendingValues.pop_back_val();
+    [&](mlir::Value value, Operation* user) {
-    if (!visitedValues.insert(value).second)
+      if (auto forOp = dyn_cast<scf::ForOp>(user);
-      continue;
+          forOp && parentLoop && forOp != parentLoop && llvm::is_contained(forOp.getInitArgs(), value))
-
+        interval.escapesLoop = true;
-    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;
-        }
-      }
-
       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) {
+        if (!forOp && !ifOp && !indexSwitch)
-          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) {
           addFallbackReason(interval.fallbackReason, "yield without scf.for parent");
-        }
+        else if (forOp && parentLoop && forOp == parentLoop && llvm::is_contained(yieldOp.getOperands(), value))
-        else {
+          interval.escapesLoop = true;
-          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;
-          }
-        }
       }
 
       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,
+SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp) {
-                                                                       std::optional<unsigned> lane,
-                                                                       bool includeAliasDescriptions) {
   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 =
+    MemoryTouchInterval touchInterval = computeMemoryTouchInterval(allocOp, ordering, fallbackEnd);
-      computeMemoryTouchInterval(allocOp, ordering, fallbackEnd, includeAliasDescriptions);
     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;
+    SmallVector<const PlannedPhysicalSlot*, 16> conflictingSlots;
-    auto bestKey = std::tuple<size_t, size_t, size_t, size_t>(std::numeric_limits<size_t>::max(),
+    SmallVector<size_t, 16> candidateAddresses = {0};
-                                                              std::numeric_limits<size_t>::max(),
+    size_t currentPeak = 0;
-                                                              std::numeric_limits<size_t>::max(),
+    for (const PlannedPhysicalSlot& slot : slots) {
-                                                              std::numeric_limits<size_t>::max());
+      currentPeak = std::max(currentPeak, slot.address + slot.requiredSize);
-
+      if (llvm::any_of(slot.intervalIndices, [&](size_t otherIndex) {
-    for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
+            return intervalsOverlap(interval, intervals[otherIndex]);
-      PlannedPhysicalSlot& slot = slots[slotIndex];
+          })) {
-      bool compatible = true;
+        conflictingSlots.push_back(&slot);
-      for (size_t otherIndex : slot.intervalIndices) {
+        size_t candidate = llvm::alignTo(slot.address + slot.requiredSize, 4);
-        if (intervalsOverlap(interval, intervals[otherIndex])) {
+        if (!llvm::is_contained(candidateAddresses, candidate))
-          compatible = false;
+          candidateAddresses.push_back(candidate);
-          break;
-        }
       }
-      if (!compatible)
+    }
-        continue;
 
-      size_t resultingSize = std::max(slot.requiredSize, interval.size);
+    size_t bestAddress = std::numeric_limits<size_t>::max();
-      size_t growth = resultingSize - slot.requiredSize;
+    auto bestKey = std::tuple<size_t, size_t>(std::numeric_limits<size_t>::max(),
-      auto candidateKey =
+                                              std::numeric_limits<size_t>::max());
-        std::tuple<size_t, size_t, size_t, size_t>(growth, resultingSize, slot.intervalIndices.size(), slot.id);
+    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) {
+    auto reusable = llvm::find_if(slots, [&](const PlannedPhysicalSlot& slot) {
-      slots.push_back({slots.size(), interval.size, interval.size, 0, {intervalIndex}});
+      return slot.address == bestAddress && slot.requiredSize == interval.size;
-      interval.slotPlanIndex = slots.size() - 1;
+    });
-      interval.physicalSlotId = slots.back().id;
+    if (reusable != slots.end()) {
-      interval.physicalSlotSize = slots.back().requiredSize;
+      reusable->intervalIndices.push_back(intervalIndex);
-      continue;
+      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,
+CoreMemoryPlan onnx_mlir::buildCoreMemoryPlan(Operation* coreLikeOp) {
-                                                        std::optional<unsigned> lane,
+  CoreMemoryPlan plan;
-                                                        ArrayRef<LocalAllocInterval> intervals,
+  plan.coreLikeOp = coreLikeOp;
-                                                        ArrayRef<PlannedPhysicalSlot> slots,
+  plan.intervals = buildLocalAllocIntervals(coreLikeOp);
-                                                        size_t addressLimit,
+  plan.slots = planPhysicalSlots(plan.intervals);
-                                                        PimMemoryReportLevel reportLevel) {
+  return plan;
-  MemoryPlanArtifacts artifacts;
+}
+
+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) {
+  for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
-    totalPhysicalBytes += slot.size;
+    const PlannedPhysicalSlot& slot = slots[slotIndex];
-    largestPhysicalSlot = std::max(largestPhysicalSlot, slot.size);
+    largestPhysicalSlot = std::max(largestPhysicalSlot, slot.requiredSize);
-    if (slot.intervalIndices.size() > 1)
+    maximumAssignedAddress = std::max(maximumAssignedAddress, slot.address + slot.requiredSize);
-      reusedAllocations += slot.intervalIndices.size() - 1;
+    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 << "  address placements: " << slots.size() << "\n";
-  os << "  reused allocations: " << reusedAllocations << "\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)
+  for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
-    if (slot.intervalIndices.size() > 1)
+    if (hasAddressReuse(slotIndex, slots))
-      reusedSlots.push_back(&slot);
+      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)
+    if (lhs->requiredSize != rhs->requiredSize)
-      return lhs->size > rhs->size;
+      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)
+    if (lhs->requiredSize != rhs->requiredSize)
-      return lhs->size > rhs->size;
+      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
+       << "  size=" << formatReportMemory(slot->requiredSize) << "  interval=#" << interval.id
-       << "  value=" << summarizeValue(interval.key.value, 56) << "\n";
+       << "  value=" << abbreviate(interval.valueSummary, 56) << "\n";
-    os << "    first=" << summarizeOperation(interval.firstTouchOp, 40)
+    os << "    first=" << abbreviate(interval.firstTouchSummary, 40)
-       << "  last=" << summarizeOperation(interval.lastTouchOp, 40)
+       << "  last=" << abbreviate(interval.lastTouchSummary, 40)
        << "  nested=" << (interval.insideNestedRegion ? "yes" : "no")
        << "  escapes_loop=" << (interval.escapesLoop ? "yes" : "no") << "\n";
   }
-  size_t fallbackPrinted = 0;
+  if (singleUseSlots.empty())
-  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)
     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) << " ("
+      os << "  slot #" << slot.id << "  addr=" << slot.address << "  size=" << formatReportMemory(slot.requiredSize)
-         << slot.size << ")"
+         << " (" << 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)
+           << "  first=" << abbreviate(interval.firstTouchSummary, 48)
-           << "  last=" << summarizeOperation(interval.lastTouchOp, 48) << "\n";
+           << "  last=" << abbreviate(interval.lastTouchSummary, 48) << "\n";
-        os << "      value=" << summarizeValue(allocValue) << "\n";
+        os << "      value=" << abbreviate(interval.valueSummary, 72) << "\n";
-      }
+        if (!interval.fallbackReason.empty())
-    }
+          os << "      fallback_reason=" << interval.fallbackReason << "\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";
       }
     }
   }
@@ -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

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

@@ -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,
+struct CoreMemoryPlan {
-                                                                  std::optional<unsigned> lane,
+  mlir::Operation* coreLikeOp = nullptr;
-                                                                  bool includeAliasDescriptions = true);
+  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,
+CoreMemoryPlan buildCoreMemoryPlan(mlir::Operation* coreLikeOp);
-                                             std::optional<unsigned> lane,
+
-                                             llvm::ArrayRef<LocalAllocInterval> intervals,
+mlir::LogicalResult assignPhysicalSlotAddresses(CoreMemoryPlan& plan, size_t addressLimit);
-                                             llvm::ArrayRef<PlannedPhysicalSlot> slots,
+
-                                             size_t addressLimit,
+std::string buildMemoryPlanReport(mlir::Operation* coreLikeOp,
-                                             PimMemoryReportLevel reportLevel);
+                                  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;
+  if (failed(walkLocalMemoryUses(allocOp.getResult(), [&](Value, Operation* user) {
-  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 (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
-        auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
+        if (!isa<scf::ForOp>(yieldOp->getParentOp()))
-        if (!forOp)
           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 candidates = 0;
-  uint64_t numSkipped = 0;
+  uint64_t logicalBytes = 0;
-  uint64_t numRemoved = 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 {
+  std::string label;
-    Core,
+  uint64_t coreCount = 1;
-    Batch
-  };
-
-  Kind kind = Kind::Core;
-  uint64_t id = 0;
-  llvm::SmallVector<int32_t, 8> coreIds;
   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);
+    total.candidates += entry.row.candidates * entry.coreCount;
-    totalRow.numCandidates += entryTotal.numCandidates;
+    total.logicalBytes += entry.row.logicalBytes * entry.coreCount;
-    totalRow.numSkipped += entryTotal.numSkipped;
+    total.skipped += entry.row.skipped * entry.coreCount;
-    totalRow.numRemoved += entryTotal.numRemoved;
+    total.removed += entry.row.removed * entry.coreCount;
-    totalRow.savedBytes += entryTotal.savedBytes;
+    total.savedBytes += entry.row.savedBytes * entry.coreCount;
   }
-
+  printReportTotalsBlock(os,
-  llvm::SmallVector<ReportField, 4> totalFields = {
+                         {{"Number of candidates", std::to_string(total.candidates)},
-    {"Number of candidates", std::to_string(totalRow.numCandidates)},
+                          {"Logical candidate bytes", formatReportMemory(total.logicalBytes)},
-    {"Skipped allocations",  std::to_string(totalRow.numSkipped)   },
+                          {"Skipped allocations", std::to_string(total.skipped)},
-    {"Removed allocations",  std::to_string(totalRow.numRemoved)   },
+                          {"Removed allocations", std::to_string(total.removed)},
-    {"Saved memory",         formatMemory(totalRow.savedBytes)     }
+                          {"Saved memory", formatReportMemory(total.savedBytes)}});
-  };
+  for (const CoalescingReportEntry& entry : entries) {
-  printReportTotalsBlock(os, totalFields);
+    os << "\n" << entry.label << ":\n";
-  if (!entries.empty())
+    printReportFlatFields(os,
-    os << "\n";
+                          {{"Number of candidates", std::to_string(entry.row.candidates)},
-
+                           {"Logical candidate bytes", formatReportMemory(entry.row.logicalBytes)},
-  llvm::SmallVector<CoalescingReportEntry, 32> sortedEntries(entries.begin(), entries.end());
+                           {"Skipped allocations", std::to_string(entry.row.skipped)},
-  sortReportEntriesByFirstCore(sortedEntries);
+                           {"Removed allocations", std::to_string(entry.row.removed)},
-
+                           {"Saved memory", formatReportMemory(entry.row.savedBytes)}});
-  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;
   }
-
   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"; }
+  StringRef getDescription() const override { return "Safely coalesce compatible block-local PIM allocations"; }
-
-  PimMemoryCoalescingPass() = default;
-  PimMemoryCoalescingPass(const PimMemoryCoalescingPass& pass) {}
 
   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 {
+      CoalescingReportRow row {analysis.getCandidateCount(),
-        analysis.getCandidateCount(), stats.skippedAllocations, stats.removedAllocs, stats.savedBytes};
+                               analysis.getCandidateBytes(),
-
+                               stats.skippedAllocations,
+                               stats.removedAllocs,
+                               stats.savedBytes};
       if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
-        auto checkedCoreId =
+        reportEntries.push_back({"Core " + std::to_string(coreOp.getCoreId()), 1, row});
-          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});
         return;
       }
-
+      SmallVector<int32_t> coreIds = getBatchCoreIds(cast<pim::PimCoreBatchOp>(op));
-      auto coreIds = getBatchCoreIds(cast<pim::PimCoreBatchOp>(op));
+      reportEntries.push_back(
-      CoalescingReportEntry entry;
+        {"Batch " + std::to_string(nextBatchId++), std::max<uint64_t>(1, coreIds.size()), row});
-      entry.kind = CoalescingReportEntry::Kind::Batch;
-      entry.id = nextBatchId++;
-      llvm::append_range(entry.coreIds, coreIds);
-      entry.row = row;
-      reportEntries.push_back(std::move(entry));
     });
-
-    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 {
-enum class CommunicationEventKind { Send, Receive };
+  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
+     << (event.kind == CommunicationEventKind::Send ? "to" : "from") << " " << event.peerCoreId << " size "
-     << " size " << event.size << "B ordinal " << event.ordinal;
+     << event.size << "B ordinal " << event.ordinal;
-  if (event.minChannelId)
+  if (minChannelId)
-    os << " min_channel " << *event.minChannelId;
+    os << " min_channel " << *minChannelId;
-  if (event.commOrder)
+  if (commOrder)
-    os << " comm_order " << *event.commOrder;
+    os << " comm_order " << *commOrder;
-  if (!event.materializer.empty())
+  if (!materializer.empty())
-    os << " materializer " << event.materializer;
+    os << " materializer " << materializer;
-  if (event.traceId)
+  if (traceId)
-    os << " trace#" << *event.traceId;
+    os << " trace#" << *traceId;
-  if (!event.tracePhase.empty())
+  if (!tracePhase.empty())
-    os << " phase " << event.tracePhase;
+    os << " phase " << tracePhase;
-  if (event.traceClassId)
+  if (traceClassId)
-    os << " class " << event.traceClassKind << "#" << *event.traceClassId;
+    os << " class " << traceClassKind << "#" << *traceClassId;
-  if (event.traceBlockOrdinal)
+  if (traceBlockOrdinal)
-    os << " block_ordinal " << *event.traceBlockOrdinal;
+    os << " block_ordinal " << *traceBlockOrdinal;
-  if (!event.tracePayload.empty())
+  if (!tracePayload.empty())
-    os << " payload " << event.tracePayload;
+    os << " payload " << tracePayload;
-  if (!event.traceMessages.empty())
+  if (!traceMessages.empty())
-    os << " messages {" << event.traceMessages << "}";
+    os << " messages {" << traceMessages << "}";
-  if (!event.tracePrevOp.empty() || !event.traceNextOp.empty())
+  if (!tracePrevOp.empty() || !traceNextOp.empty())
-    os << " inserted_between [" << event.tracePrevOp << " | " << event.traceNextOp << "]";
+    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>
-static std::optional<size_t> findMatchingCounterpartIndex(const CommunicationEventVector& events,
+findMatchingCounterpartIndex(const CommunicationEventVector& events, const CommunicationEvent& event, size_t begin) {
-                                                           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()
+  os << "    peer core " << blockedEvent.peerCoreId << " current pc " << peerPc << " of " << peerEvents.size() << "\n";
-     << "\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,
+static CommunicationEvent makeCommunicationEvent(
-                                                 int64_t coreId,
+  CommunicationEventKind kind, int64_t coreId, int64_t peerCoreId, int64_t size, uint64_t ordinal, Operation* op) {
-                                                 int64_t peerCoreId,
+  return CommunicationEvent {kind, coreId, peerCoreId, size, ordinal, op};
-                                                 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 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) {
+  return walkPimCoreCommunicationBlock(
-    if (auto sendOp = dyn_cast<pim::PimSendOp>(&op)) {
+    block, plan, initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
-      auto targetCoreId = resolveIndexValue(sendOp.getTargetCoreId(), knowledge);
+      if (auto sendOp = dyn_cast<pim::PimSendOp>(&op)) {
-      if (failed(targetCoreId)) {
+        auto targetCoreId = resolveIndexValue(sendOp.getTargetCoreId(), knowledge);
-        diagnostics.report(&op, [](Operation* illegalOp) {
+        if (failed(targetCoreId)) {
-          illegalOp->emitOpError("cannot statically resolve send target core for PIM communication deadlock check");
+          diagnostics.report(&op, [](Operation* illegalOp) {
-        });
+            illegalOp->emitOpError("cannot statically resolve send target core for PIM communication deadlock check");
-        return failure();
+          });
+          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,
+      if (auto receiveOp = dyn_cast<pim::PimReceiveOp>(&op)) {
-                                              coreId,
+        auto sourceCoreId = resolveIndexValue(receiveOp.getSourceCoreId(), knowledge);
-                                              *targetCoreId,
+        if (failed(sourceCoreId)) {
-                                              sendOp.getSize(),
+          diagnostics.report(&op, [](Operation* illegalOp) {
-                                              static_cast<uint64_t>(events.size()),
+            illegalOp->emitOpError(
-                                              &op));
+              "cannot statically resolve receive source core for PIM communication deadlock check");
-      return success();
+          });
-    }
+          return failure();
+        }
 
-    if (auto receiveOp = dyn_cast<pim::PimReceiveOp>(&op)) {
+        events.push_back(makeCommunicationEvent(CommunicationEventKind::Receive,
-      auto sourceCoreId = resolveIndexValue(receiveOp.getSourceCoreId(), knowledge);
+                                                coreId,
-      if (failed(sourceCoreId)) {
+                                                *sourceCoreId,
-        diagnostics.report(&op, [](Operation* illegalOp) {
+                                                receiveOp.getSize(),
-          illegalOp->emitOpError("cannot statically resolve receive source core for PIM communication deadlock check");
+                                                static_cast<uint64_t>(events.size()),
-        });
+                                                &op));
-        return failure();
+        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()
+  auto diagnostic =
-      << "PIM communication deadlock check found a blocking send/receive cycle while statically simulating the "
+    moduleOp.emitError()
-         "expanded per-core communication streams; see the PIM static communication deadlock report above";
+    << "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())
+    std::string materializer = getNearestStringAttr(event.op, kRaptorMaterializerAttr);
-      note << " emitted by " << event.materializer;
+    if (!materializer.empty())
+      note << " emitted by " << materializer;
   }
 }
 
-static FailureOr<SmallVector<int64_t>> findCommunicationWaitCycle(
+static FailureOr<SmallVector<int64_t>>
-  const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
+findCommunicationWaitCycle(const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
-  const DenseMap<int64_t, size_t>& programCounters) {
+                           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(
+        PimCoreCommunicationPlan plan = buildPimCoreCommunicationPlan(coreOp.getBody().front());
-              coreOp.getBody().front(), coreId, StaticValueKnowledge {}, coreEvents[coreId], diagnostics)))
+        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) {
+          for (int32_t coreId : laneCoreIds)
-            if (failed(appendCoreCommunicationEvents(
+            if (failed(appendCoreCommunicationEvents(coreBatchOp.getBody().front(),
-                  coreBatchOp.getBody().front(), coreId, laneKnowledge, coreEvents[coreId], diagnostics)))
+                                                     plan,
+                                                     coreId,
+                                                     laneKnowledge,
+                                                     coreEvents[coreId],
+                                                     diagnostics)))
               hasFailure = true;
-          }
         }
       }
     }
@@ -607,9 +656,10 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
       return failure();
     }
 
-    auto diagnostic = moduleOp.emitError()
+    auto diagnostic =
-        << "PIM communication deadlock check stalled without finding a closed wait cycle; this usually means a "
+      moduleOp.emitError()
-           "send/receive peer is missing or ordered after a finished core";
+      << "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

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

src/PIM/Pass/PIMPasses.h

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

src/PIM/PimAccelerator.cpp

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

test/PIM/PimMemoryLivenessPlannerTest.cpp

@@ -2,9 +2,11 @@
 #include <cstdlib>
 #include <iostream>
 
-#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp"
 
 using onnx_mlir::LocalAllocInterval;
+using onnx_mlir::CoreMemoryPlan;
+using onnx_mlir::assignPhysicalSlotAddresses;
 using onnx_mlir::planPhysicalSlots;
 
 namespace {
@@ -18,12 +20,19 @@ LocalAllocInterval makeInterval(size_t id, size_t size, uint64_t start, uint64_t
   return interval;
 }
 
+size_t getPeak(llvm::ArrayRef<onnx_mlir::PlannedPhysicalSlot> slots) {
+  size_t peak = 0;
+  for (const auto& slot : slots)
+    peak = std::max(peak, slot.address + slot.requiredSize);
+  return peak;
+}
+
 void assertSingleSlotCase(LocalAllocInterval a, LocalAllocInterval b, size_t expectedSlotSize) {
   llvm::SmallVector<LocalAllocInterval, 4> intervals = {a, b};
   auto slots = planPhysicalSlots(intervals);
   assert(slots.size() == 1);
   assert(slots.front().requiredSize == expectedSlotSize);
-  assert(intervals[0].physicalSlotId == intervals[1].physicalSlotId);
+  assert(intervals[0].slotPlanIndex == intervals[1].slotPlanIndex);
 }
 
 int testSameSizeNonOverlap() {
@@ -34,13 +43,21 @@ int testSameSizeNonOverlap() {
 
 int testLargerFirst() {
   std::cout << "testLargerFirst:" << std::endl;
-  assertSingleSlotCase(makeInterval(0, 100, 0, 10), makeInterval(1, 40, 11, 20), 100);
+  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
+    makeInterval(0, 100, 0, 10), makeInterval(1, 40, 11, 20)};
+  auto slots = planPhysicalSlots(intervals);
+  assert(slots.size() == 2);
+  assert(getPeak(slots) == 100);
   return 0;
 }
 
 int testSmallerFirst() {
   std::cout << "testSmallerFirst:" << std::endl;
-  assertSingleSlotCase(makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20), 100);
+  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
+    makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20)};
+  auto slots = planPhysicalSlots(intervals);
+  assert(slots.size() == 2);
+  assert(getPeak(slots) == 100);
   return 0;
 }
 
@@ -50,7 +67,7 @@ int testOverlapNeedsTwoSlots() {
     makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)};
   auto slots = planPhysicalSlots(intervals);
   assert(slots.size() == 2);
-  assert(intervals[0].physicalSlotId != intervals[1].physicalSlotId);
+  assert(intervals[0].slotPlanIndex != intervals[1].slotPlanIndex);
   return 0;
 }
 
@@ -59,10 +76,31 @@ int testReuseChain() {
   llvm::SmallVector<LocalAllocInterval, 4> intervals = {
     makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20), makeInterval(2, 20, 21, 30)};
   auto slots = planPhysicalSlots(intervals);
-  assert(slots.size() == 1);
+  assert(slots.size() == 3);
-  assert(slots.front().requiredSize == 100);
+  assert(getPeak(slots) == 100);
-  assert(intervals[0].physicalSlotId == intervals[1].physicalSlotId);
+  return 0;
-  assert(intervals[1].physicalSlotId == intervals[2].physicalSlotId);
+}
+
+int testPartialAddressReuse() {
+  std::cout << "testPartialAddressReuse:" << std::endl;
+  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
+    makeInterval(0, 100, 0, 10), makeInterval(1, 60, 11, 20), makeInterval(2, 40, 11, 20)};
+  auto slots = planPhysicalSlots(intervals);
+  assert(getPeak(slots) == 100);
+  return 0;
+}
+
+int testAddressAssignment() {
+  std::cout << "testAddressAssignment:" << std::endl;
+  CoreMemoryPlan plan;
+  plan.intervals = {makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)};
+  plan.slots = planPhysicalSlots(plan.intervals);
+  assert(mlir::succeeded(assignPhysicalSlotAddresses(plan, 1024)));
+  assert(plan.slots.size() == 2);
+  const auto& firstSlot = plan.slots[plan.intervals[0].slotPlanIndex];
+  const auto& secondSlot = plan.slots[plan.intervals[1].slotPlanIndex];
+  assert(firstSlot.address == 0 && firstSlot.requiredSize == 100);
+  assert(secondSlot.address == 100 && secondSlot.requiredSize == 40);
   return 0;
 }
 
@@ -78,6 +116,8 @@ int main(int argc, char *argv[]) {
   failures += testSmallerFirst();
   failures += testOverlapNeedsTwoSlots();
   failures += testReuseChain();
+  failures += testPartialAddressReuse();
+  failures += testAddressAssignment();
   if (failures != 0) {
     std::cerr << failures << " test failures\n";
     return EXIT_FAILURE;

validation/raptor.py

@@ -11,6 +11,8 @@ PIM_PASS_LABELS = (
     ("SpatialToPimPass", "Spatial to PIM"),
     ("PimBufferizationPass", "Bufferize PIM"),
     ("HostConstantFoldingPass", "Fold Host Constants"),
+    ("PimMemoryCoalescingPass", "Coalesce Local Memory"),
+    ("PimLocalMemoryPlanningPass", "Plan Local Memory"),
     ("VerificationPass", "Verify PIM"),
     ("EmitPimCodePass", "Emit PIM Code"),
 )

validation/tools/classification_local_image_validation.py

@@ -250,9 +250,9 @@ def main():
     parser.add_argument("--raptor-path", type=Path, default=defaults["raptor_path"])
     parser.add_argument("--onnx-include-dir", type=Path, default=defaults["onnx_include_dir"])
     parser.add_argument("--simulator-dir", type=Path, default=defaults["simulator_dir"])
-    parser.add_argument("--crossbar-size", type=int, default=2048)
+    parser.add_argument("--crossbar-size", type=int, default=128)
-    parser.add_argument("--crossbar-count", type=int, default=256)
+    parser.add_argument("--crossbar-count", type=int, default=64)
-    parser.add_argument("--core-count", type=int, default=1000)
+    parser.add_argument("--core-count", type=int, default=144)
     parser.add_argument("--top-k", type=int, default=5)
     parser.add_argument("--command-timeout-seconds", type=float, default=7200.0)
     parser.add_argument("--skip-compile", action="store_true")

validation/tools/yolo_local_image_validation.py

@@ -198,9 +198,9 @@ def main():
     parser.add_argument("--raptor-path", type=Path, default=defaults["raptor_path"])
     parser.add_argument("--onnx-include-dir", type=Path, default=defaults["onnx_include_dir"])
     parser.add_argument("--simulator-dir", type=Path, default=defaults["simulator_dir"])
-    parser.add_argument("--crossbar-size", type=int, default=2048)
+    parser.add_argument("--crossbar-size", type=int, default=128)
-    parser.add_argument("--crossbar-count", type=int, default=256)
+    parser.add_argument("--crossbar-count", type=int, default=64)
-    parser.add_argument("--core-count", type=int, default=1000)
+    parser.add_argument("--core-count", type=int, default=144)
     parser.add_argument("--command-timeout-seconds", type=float, default=7200.0)
     parser.add_argument("--skip-compile", action="store_true")
     parser.add_argument("--verbose", action="store_true")

validation/tools/yolo_real_image_validation.py

@@ -399,9 +399,9 @@ def main():
     parser.add_argument("--remote-project", default="/home/gmagnani/Project/Raptor")
     parser.add_argument("--remote-python", default="/home/gmagnani/venv/bin/python")
     parser.add_argument("--network-dir", default="validation/networks/yolo11n/depth_51")
-    parser.add_argument("--crossbar-size", type=int, default=2048)
+    parser.add_argument("--crossbar-size", type=int, default=128)
-    parser.add_argument("--crossbar-count", type=int, default=256)
+    parser.add_argument("--crossbar-count", type=int, default=64)
-    parser.add_argument("--core-count", type=int, default=1000)
+    parser.add_argument("--core-count", type=int, default=144)
     parser.add_argument("--command-timeout-seconds", type=int, default=7200)
     parser.add_argument("--skip-compile", action="store_true")
     parser.add_argument("--annotated-dir", default="validation/networks/yolo11n/depth_51/real_image_validation/annotated")

validation/validate.py

@@ -72,10 +72,9 @@ def main():
     ap.add_argument("--relative-threshold", type=float, default=1e-5,
                     help="Relative tolerance for per-element output comparison.")
     ap.add_argument("--seed", type=int, default=0, help="RNG seed for generated validation inputs.")
-    ap.add_argument("--crossbar-size", type=int, default=64)
+    ap.add_argument("--crossbar-size", type=int, default=128)
-    ap.add_argument("--crossbar-count", type=int, default=8)
+    ap.add_argument("--crossbar-count", type=int, default=64)
-    ap.add_argument("--core-count", type=int, default=None,
+    ap.add_argument("--core-count", type=int, default=144)
-                    help="Core count to pass to Raptor. Required for PIM validation.")
     ap.add_argument("--pim-memory-report", choices=("none", "summary", "full"), default="none",
                     help="Emit a human-readable PIM memory planning report during codegen.")
     ap.add_argument("--raptor-extra-arg", action="append", default=[],
@@ -120,8 +119,6 @@ def main():
         missing_args.append("--raptor-path")
     if not a.onnx_include_dir:
         missing_args.append("--onnx-include-dir")
-    if a.core_count is None:
-        missing_args.append("--core-count")
     if missing_args:
         ap.error("the following arguments are required unless --clean is used: " + ", ".join(missing_args))
 

validation/validate_one.py

@@ -311,7 +311,7 @@ def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1
 
 
 def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
-                     simulator_dir, crossbar_size=64, crossbar_count=8, core_count=None,
+                     simulator_dir, crossbar_size=128, crossbar_count=64, core_count=144,
                      pim_memory_report="none", raptor_extra_args=None,
                      threshold=1e-3, rtol=1e-5,
                      seed=0, reporter=None, model_index=1, model_total=1, verbose=False,