From 15e8edb9c47e2e701c628de11d00ce9e9c015614 Mon Sep 17 00:00:00 2001
From: NiccoloN <niccolo.nicolosi@gmail.com>
Date: Sat, 25 Apr 2026 19:24:09 +0200
Subject: [PATCH] better spat computes merging

---
 README.md                                     | 154 +++++
 src/PIM/Compiler/PimCompilerOptions.cpp       |   4 +-
 .../ONNXToSpatial/ONNXToSpatialPass.cpp       | 104 +---
 src/PIM/Conversion/SpatialToPim/Common.cpp    |  19 +-
 .../DCPGraph/DCPAnalysis.cpp                  | 346 ++++++++---
 .../MergeComputeNodes/DCPGraph/Graph.cpp      | 514 ++++++++++++---
 .../MergeComputeNodes/DCPGraph/Graph.hpp      |  31 +-
 .../MergeComputeNodes/DCPGraph/GraphDebug.cpp |  41 +-
 .../MergeComputeNodes/DCPGraph/GraphDebug.hpp |   5 +-
 .../MergeComputeNodesPass.cpp                 | 585 +++++++++++++++++-
 test/PIM/DCPTest.cpp                          |  43 +-
 11 files changed, 1477 insertions(+), 369 deletions(-)

diff --git a/README.md b/README.md
index 3d26525..e973656 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,159 @@
 # Raptor
 
+Raptor is a domain-specific MLIR compiler for neural networks (ONNX format)
+targeting in-memory computing / processing-in-memory (PIM) architectures.
+It progressively lowers ONNX-MLIR through a set of MLIR dialects down to
+target-specific artifacts (currently JSON code for the `pimsim-nn` simulator).
+
+## Overview
+
+PIM architectures perform most of the computation directly in memory.
+Raptor's first supported target is `pimsim-nn`, which simulates a chip with:
+- a shared host memory,
+- a number of cores that do most of the computation directly in their memory
+  (vector ops, vmm/mvm on ReRAM crossbars),
+- no branching instructions (branchless architecture) and no hardware loop
+  support — any repeated work (e.g. convolutions) must be unrolled into
+  explicit per-iteration instructions.
+
+Because of this, the amount of emitted instructions explodes quickly and the
+compiler must optimize aggressively at every stage to keep compilation
+tractable.
+
+A second target, `PulPim`, is planned for an accelerator with RISC-V cores
+each carrying its own in-memory computing unit and crossbars. It will live in
+a dedicated dialect (future work).
+
+### Targets and simulators
+
+`pimsim-nn` (under `backend-simulators/pim/pimsim-nn`) is used for
+**performance** estimates (latency, energy), but does not functionally execute
+the JSON code it consumes. To validate the numerical correctness of the JSON
+code produced by Raptor (or, for comparison, by the `pimcomp` compiler), we use
+a Rust simulator we maintain in-tree at
+`backend-simulators/pim/pim-simulator`.
+
+## Compilation pipeline
+
+The PIM-related sources live under `src/PIM` and the tests under `test/PIM`.
+When working on this codebase, most changes should stay confined to those
+trees (you only need to look outside, e.g. at `onnx-mlir` or `llvm`, for
+framework-level details).
+
+High-level lowering flow:
+
+```
+ONNX-MLIR ──► Spatial ──► Pim (tensor) ──► Pim (bufferized) ──► PIM JSON
+```
+
+1. **ONNX → Spatial** (`src/PIM/Conversion/ONNXToSpatial`).
+   Lowers ONNX ops into the `spat` dialect (`src/PIM/Dialect/Spatial`).
+   Spatial models a high-level spatial in-memory accelerator: vmm/mvm
+   operations are accelerated by storing a constant RHS matrix into a
+   crossbar. Crossbars cannot be re-programmed during execution, have a
+   limited fixed size, and there is a limited number of them per core.
+   Conversion patterns are split by op family under
+   `Conversion/ONNXToSpatial/Patterns/{Math,NN,Tensor}` (Conv, Gemm, MatMul,
+   Elementwise, ReduceMean, Pool, Relu, Sigmoid, Softmax, Concat, Gather,
+   Reshape, Resize, Split).
+
+2. **Spatial → Pim** (`src/PIM/Conversion/SpatialToPim`).
+   Lowers Spatial to the `pim` dialect (`src/PIM/Dialect/Pim`), which
+   materializes PIM cores (`pim.core`), inter-core communication
+   (`pim.send` / `pim.receive`), halts, and crossbar-level operations.
+
+3. **Merge compute nodes** (`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`).
+   A DCP-inspired heuristic (Dynamic Critical Path — see the original
+   scheduling paper by Kwok & Ahmad,
+   [DCP-eScience2007](https://clouds.cis.unimelb.edu.au/papers/DCP-eScience2007.pdf))
+   that coarsens the virtual node graph and decides how to group compute
+   nodes onto cores. Our implementation is only DCP-*inspired*: it is a
+   heuristic with different assumptions from the paper (different cost
+   model, constraints from crossbar capacity / core resources, and a
+   windowed coarsening loop instead of full-graph reprioritization). The
+   `dcp-critical-window-size` option controls how many lowest-slack virtual
+   nodes each coarsening iteration considers (0 = legacy full-graph
+   analysis). Related sources: `DCPGraph/DCPAnalysis.cpp`, `Graph.cpp/.hpp`,
+   `MergeComputeNodesPass.cpp`.
+
+4. **Bufferization** (`src/PIM/Dialect/Pim/Transforms/Bufferization`).
+   Converts tensor-semantics PIM IR into memref-semantics PIM IR using the
+   standard MLIR `BufferizableOpInterface` machinery
+   (`OpBufferizationInterfaces.*`, `PimBufferization.td`).
+
+5. **PIM code generation** (`src/PIM/Pass/PimCodegen`):
+   - `HostConstantFolding` — folds host-side constants.
+   - `MaterializeHostConstantsPass` — materializes the remaining host
+     constants for emission.
+   - `VerificationPass` — checks invariants before emission.
+   - `EmitPimJsonPass` — emits the final PIM JSON consumed by `pimsim-nn`
+     and `pim-simulator`.
+
+Supporting pieces:
+- `src/PIM/Compiler` — PIM-specific compiler options (crossbar size/count,
+  core count, DCP window, experimental conv impl, concat error handling, …)
+  and `PimCodeGen` entry points.
+- `src/PIM/Common` — shared utilities (`PimCommon`, `LabeledList`).
+- `src/PIM/Pass` — auxiliary passes (`MessagePass`, `CountInstructionPass`)
+  and the `PIMPasses.h` registry used by `PimAccelerator`.
+- `src/PIM/PimAccelerator.{cpp,hpp}` — accelerator entry point: registers
+  dialects, passes, and plugs Raptor into the ONNX-MLIR driver.
+
+## Key compiler options
+
+Pass these on the `onnx-mlir` command line when compiling for PIM:
+
+- `--maccel=PIM` — select the PIM accelerator.
+- `--EmitSpatial` / `--EmitPim` / `--EmitPimBufferized` / `--EmitPimCodegen`
+  — stop the pipeline at the requested stage (default: `EmitPimCodegen`).
+- `--pim-only-codegen` — assume the input is already bufferized PIM IR and
+  run only the codegen tail.
+- `--crossbar-size=<N>` / `--crossbar-count=<N>` — crossbar dimensions and
+  per-core count.
+- `--core-count=<N>` — number of cores (`-1` picks the minimum).
+- `--dcp-critical-window-size=<N>` — DCP coarsening window (0 = legacy).
+- `--use-experimental-conv-impl` — alternative convolution lowering.
+- `--ignore-concat-error` — soft-fail corner case in `ConcatOp`.
+
+## Validation
+
+Functional validation lives in `validation/` and drives the Rust
+`pim-simulator` to compare Raptor's output against a reference.
+
+Per-operation validation (from `validation/`):
+
+```
+validate.py \
+    --raptor-path ../cmake-build-release/Release/bin/onnx-mlir \
+    --onnx-include-dir ../onnx-mlir/include
+```
+
+End-to-end network validation (example: first 4 layers of YOLOv11n):
+
+```
+validate.py \
+    --raptor-path ../cmake-build-release/Release/bin/onnx-mlir \
+    --onnx-include-dir ../onnx-mlir/include \
+    --operations-dir ./networks/yolo11n/depth_04 \
+    --crossbar-size 2048
+```
+
+Available networks under `validation/networks/`: `vgg16`, `yolo11n`.
+Available operations under `validation/operations/`: `add`, `conv`, `div`,
+`gather`, `gemm`, `gemv`, `mul`, `pool`, `reduce_mean`, `relu`, `resize`,
+`sigmoid`, `softmax`, `split`.
+
+## Rebuilding
+
+Release build (fast):
+
+```
+cmake --build /home/nico/raptor/raptor/cmake-build-release --target onnx-mlir -j 30
+```
+
+A slower debug build is also available — configure it the same way but with
+`-DCMAKE_BUILD_TYPE=Debug` (see installation instructions below).
+
 ## Build
 
 ### Protobuf
diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp
index 7f4b18c..1b3288a 100644
--- a/src/PIM/Compiler/PimCompilerOptions.cpp
+++ b/src/PIM/Compiler/PimCompilerOptions.cpp
@@ -41,7 +41,7 @@ llvm::cl::opt<size_t>
   crossbarSize("crossbar-size", llvm::cl::desc("Width and heigth of a single crossbar"), llvm::cl::init(2));
 
 llvm::cl::opt<size_t>
-  crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(2));
+  crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(256));
 
 llvm::cl::opt<long> coresCount("core-count",
                                llvm::cl::desc("Number of cores in the chip. `-1` to use the minimum amount of cores."),
@@ -51,7 +51,7 @@ llvm::cl::opt<size_t> dcpCriticalWindowSize(
   "dcp-critical-window-size",
   llvm::cl::desc("Number of lowest-slack virtual nodes considered by each DCP coarsening iteration. "
                  "Use 0 to run the legacy full-graph DCP analysis."),
-  llvm::cl::init(1024));
+  llvm::cl::init(4000));
 
 llvm::cl::opt<bool>
   ignoreConcatError("ignore-concat-error",
diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp
index 28e7078..6fc2240 100644
--- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp
+++ b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp
@@ -8,7 +8,6 @@
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
@@ -24,8 +23,6 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/DCPAnalysis.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -52,7 +49,6 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
 private:
   void annotateWeightsConstants(func::FuncOp funcOp) const;
   void encapsulateGlobalInstruction(func::FuncOp funcOp);
-  void mergeTriviallyConnectedComputes(func::FuncOp funcOp);
   LogicalResult promoteConstantInputsToWeights(func::FuncOp funcOp);
 };
 
@@ -155,8 +151,6 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
-  mergeTriviallyConnectedComputes(*entryFunc);
-
   // Dump to file for debug
   dumpModule(moduleOp, "spatial0");
 }
@@ -187,8 +181,8 @@ bool encapsulateConcat(IRRewriter& rewriter, Location loc, Operation* inst) {
   if (auto toRemoveOp = llvm::dyn_cast_if_present<tensor::ConcatOp>(inst)) {
     auto sources = toRemoveOp.getInputs();
     rewriter.setInsertionPointAfter(toRemoveOp);
-    if (llvm::any_of(
-          sources, [](auto source) { return isa_and_present<spatial::SpatCompute>(source.getDefiningOp()); })) {
+    if (llvm::any_of(sources,
+                     [](auto source) { return isa_and_present<spatial::SpatCompute>(source.getDefiningOp()); })) {
       auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), sources);
       SmallVector<Type> sourceTypes;
       SmallVector<Location> sourceLoc;
@@ -294,100 +288,6 @@ void ONNXToSpatialPass::encapsulateGlobalInstruction(func::FuncOp funcOp) {
   }
 }
 
-void ONNXToSpatialPass::mergeTriviallyConnectedComputes(func::FuncOp funcOp) {
-  Location loc = funcOp.getLoc();
-  IRRewriter rewriter(&getContext());
-  SmallVector<spatial::SpatCompute> trivialComputes;
-  llvm::SmallSet<spatial::SpatCompute, 8> toErase;
-
-  for (auto compute : funcOp.getOps<spatial::SpatCompute>())
-    if (compute->hasOneUse()) {
-      auto& use = *compute->getUses().begin();
-      auto user = dyn_cast<spatial::SpatCompute>(use.getOwner());
-
-      if (user && user.getInputs().size() == 1 && use.getOperandNumber() >= user.getWeights().size())
-        trivialComputes.push_back(compute);
-    }
-
-  while (!trivialComputes.empty()) {
-    auto compute = trivialComputes.front();
-
-    if (compute.use_empty()) {
-      std::swap(trivialComputes.front(), trivialComputes.back());
-      trivialComputes.pop_back();
-      continue;
-    }
-    auto& computeUse = *compute->getUses().begin();
-    auto child = cast<spatial::SpatCompute>(computeUse.getOwner());
-    auto usedResult = cast<OpResult>(computeUse.get()).getResultNumber();
-    auto childArgIndex = computeUse.getOperandNumber() - child.getWeights().size();
-
-    rewriter.setInsertionPointAfter(compute.getOperation());
-
-    auto newCompute =
-      spatial::SpatCompute::create(rewriter, loc, child.getResultTypes(), compute.getOperands());
-    newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(compute.getWeights().size()), static_cast<int>(compute.getInputs().size())});
-
-    IRMapping mapper;
-    auto weightMutableIter = newCompute.getWeightsMutable();
-    for (auto weight : child.getWeights()) {
-      auto founded = llvm::find(newCompute.getWeights(), weight);
-      if (founded == newCompute.getWeights().end()) {
-        weightMutableIter.append(weight);
-        auto last = weightMutableIter.end();
-        last = std::prev(last, 1);
-        mapper.map(weight, last->get());
-      }
-      else {
-        mapper.map(weight, *founded);
-      }
-    }
-
-    compute.getBodyRegion().cloneInto(&newCompute.getBodyRegion(), mapper);
-    auto newTerminator = newCompute.getBody().front().getTerminator();
-    mapper.map(child.getBody().front().getArgument(childArgIndex), newTerminator->getOperand(usedResult));
-    newTerminator->erase();
-    rewriter.setInsertionPoint(&newCompute.getBody().front(), newCompute.getBody().front().end());
-    for (auto& op : child.getBody().front()) {
-      auto newInst = rewriter.clone(op, mapper);
-
-      if (auto vmOp = llvm::dyn_cast<spatial::SpatWeightedMVMOp>(newInst)) {
-        auto oldIndex = vmOp.getWeightIndex();
-        auto newWeight = mapper.lookup(*std::next(child.getWeights().begin(), oldIndex));
-        auto newIndex = std::distance(newCompute.getWeights().begin(), llvm::find(newCompute.getWeights(), newWeight));
-        vmOp.setWeightIndex(newIndex);
-      }
-      if (auto vmOp = llvm::dyn_cast<spatial::SpatWeightedVMMOp>(newInst)) {
-        auto oldIndex = vmOp.getWeightIndex();
-        auto newWeight = mapper.lookup(*std::next(child.getWeights().begin(), oldIndex));
-        auto newIndex = std::distance(newCompute.getWeights().begin(), llvm::find(newCompute.getWeights(), newWeight));
-        vmOp.setWeightIndex(newIndex);
-      }
-    }
-
-    child.replaceAllUsesWith(newCompute);
-    toErase.insert(child);
-
-    std::swap(trivialComputes.front(), trivialComputes.back());
-    trivialComputes.pop_back();
-    toErase.insert(compute);
-
-    if (newCompute->hasOneUse()) {
-      auto& use = *newCompute->getUses().begin();
-      auto user = dyn_cast<spatial::SpatCompute>(use.getOwner());
-      if (user && user.getInputs().size() == 1 && use.getOperandNumber() >= user.getWeights().size())
-        trivialComputes.push_back(newCompute);
-    }
-  }
-
-  for (auto compute : toErase) {
-    for (Value result : compute->getResults())
-      result.dropAllUses();
-    compute.erase();
-  }
-}
-
 void ONNXToSpatialPass::annotateWeightsConstants(func::FuncOp funcOp) const {
   funcOp.walk([&](arith::ConstantOp constantOp) {
     if (hasOnlySpatialMvmVmmWeightUses(constantOp.getResult()))
diff --git a/src/PIM/Conversion/SpatialToPim/Common.cpp b/src/PIM/Conversion/SpatialToPim/Common.cpp
index 5d4af51..0f6859f 100644
--- a/src/PIM/Conversion/SpatialToPim/Common.cpp
+++ b/src/PIM/Conversion/SpatialToPim/Common.cpp
@@ -96,8 +96,8 @@ bool hasSpatialChannelTargetCoreIdAttr(mlir::Value channel) {
   return channelNewOp && channelNewOp->hasAttr(kChannelTargetCoreIdAttrName);
 }
 
-mlir::Value createPimReceiveFromSpatialChannel(
-  PatternRewriter& rewriter, Location loc, mlir::Value output, mlir::Value channel) {
+mlir::Value
+createPimReceiveFromSpatialChannel(PatternRewriter& rewriter, Location loc, mlir::Value output, mlir::Value channel) {
   mlir::Value outputBuffer = getBestOutputTensorFromOperandsOrAllocate(rewriter, output.getDefiningOp());
   auto sizeAttr = getTensorSizeInBytesAttr(rewriter, output);
   auto sourceCoreIdAttr = getSpatialChannelSourceCoreIdAttr(rewriter, channel);
@@ -127,6 +127,16 @@ SmallVector<mlir::Value> getOpOperandsSortedByUses(Operation* operation) {
   return map_to_vector(operandsAndUses, [](auto operandAndUse) { return operandAndUse.first; });
 }
 
+bool hasLaterUserInBlock(mlir::Value value, Operation* operation) {
+  for (Operation* user : value.getUsers()) {
+    if (user->getBlock() != operation->getBlock())
+      return true;
+    if (operation->isBeforeInBlock(user))
+      return true;
+  }
+  return false;
+}
+
 mlir::Value getBestOutputTensorFromOperandsOrAllocate(PatternRewriter& rewriter, Operation* operation) {
   assert("Only support operations with a single result" && operation->getNumResults() == 1);
   mlir::Value result = operation->getResult(0);
@@ -134,8 +144,9 @@ mlir::Value getBestOutputTensorFromOperandsOrAllocate(PatternRewriter& rewriter,
   assert("Only support result ShapedType as result type" && isa<ShapedType>(resultType));
 
   SmallVector<mlir::Value> operands = getOpOperandsSortedByUses(operation);
-  auto validOperands =
-    make_filter_range(operands, [resultType](mlir::Value operand) { return operand.getType() == resultType; });
+  auto validOperands = make_filter_range(operands, [operation, resultType](mlir::Value operand) {
+    return operand.getType() == resultType && !hasLaterUserInBlock(operand, operation);
+  });
   auto bestOperand = validOperands.begin();
 
   if (bestOperand != validOperands.end())
diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/DCPAnalysis.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/DCPAnalysis.cpp
index 781b5bc..a9e5e6c 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/DCPAnalysis.cpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/DCPAnalysis.cpp
@@ -3,15 +3,17 @@
 #include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/raw_ostream.h"
 
 #include <algorithm>
 #include <iterator>
-#include <map>
 #include <numeric>
 #include <optional>
-#include <set>
+#include <queue>
 #include <utility>
 #include <vector>
 
@@ -47,11 +49,13 @@ struct TimingInfo {
 
 struct WindowScheduleResult {
   std::vector<std::vector<size_t>> mergeGroups;
-  bool usedAllAvailableCpus = false;
+  CPU cpuCount = 0;
+  size_t mergedNodeCount = 0;
+  size_t maxMergeGroupSize = 0;
 };
 
 std::vector<IndexedEdge> aggregateEdges(ArrayRef<IndexedEdge> edges) {
-  std::map<std::pair<size_t, size_t>, Weight> edgeWeights;
+  llvm::DenseMap<std::pair<size_t, size_t>, Weight> edgeWeights;
   for (auto [start, end, weight] : edges) {
     size_t startIndex = static_cast<size_t>(start);
     size_t endIndex = static_cast<size_t>(end);
@@ -59,11 +63,9 @@ std::vector<IndexedEdge> aggregateEdges(ArrayRef<IndexedEdge> edges) {
       continue;
     auto key = std::make_pair(startIndex, endIndex);
     Weight edgeWeight = static_cast<Weight>(weight);
-    auto it = edgeWeights.find(key);
-    if (it == edgeWeights.end())
-      edgeWeights.insert({key, edgeWeight});
-    else
-      it->second = std::max(it->second, edgeWeight);
+    auto inserted = edgeWeights.try_emplace(key, edgeWeight);
+    if (!inserted.second)
+      inserted.first->second = std::max(inserted.first->second, edgeWeight);
   }
 
   std::vector<IndexedEdge> aggregatedEdges;
@@ -71,6 +73,11 @@ std::vector<IndexedEdge> aggregateEdges(ArrayRef<IndexedEdge> edges) {
   for (auto [key, weight] : edgeWeights)
     aggregatedEdges.push_back(
       {static_cast<int64_t>(key.first), static_cast<int64_t>(key.second), static_cast<int64_t>(weight)});
+  llvm::sort(aggregatedEdges, [](const IndexedEdge& lhs, const IndexedEdge& rhs) {
+    if (std::get<0>(lhs) != std::get<0>(rhs))
+      return std::get<0>(lhs) < std::get<0>(rhs);
+    return std::get<1>(lhs) < std::get<1>(rhs);
+  });
   return aggregatedEdges;
 }
 
@@ -157,10 +164,27 @@ TimingInfo computeTiming(const VirtualGraph& graph) {
   return timing;
 }
 
-std::vector<size_t> selectCriticalWindow(const TimingInfo& timing, size_t windowSize) {
-  std::vector<size_t> selected(timing.aest.size());
-  std::iota(selected.begin(), selected.end(), 0);
-  std::stable_sort(selected.begin(), selected.end(), [&](size_t lhs, size_t rhs) {
+std::vector<std::vector<size_t>> buildUndirectedAdjacency(const VirtualGraph& graph) {
+  std::vector<std::vector<size_t>> adjacency(graph.nodes.size());
+  for (auto [start, end, weight] : graph.edges) {
+    (void) weight;
+    size_t startIndex = static_cast<size_t>(start);
+    size_t endIndex = static_cast<size_t>(end);
+    assert(startIndex < graph.nodes.size() && endIndex < graph.nodes.size() && "virtual edge endpoint out of range");
+    adjacency[startIndex].push_back(endIndex);
+    adjacency[endIndex].push_back(startIndex);
+  }
+  for (auto& neighbours : adjacency) {
+    llvm::sort(neighbours);
+    neighbours.erase(std::unique(neighbours.begin(), neighbours.end()), neighbours.end());
+  }
+  return adjacency;
+}
+
+std::vector<size_t> selectCriticalWindow(const VirtualGraph& graph, const TimingInfo& timing, size_t windowSize) {
+  std::vector<size_t> ranked(timing.aest.size());
+  std::iota(ranked.begin(), ranked.end(), 0);
+  auto isHigherPriority = [&](size_t lhs, size_t rhs) {
     Time lhsSlack = slackOrZero(timing.aest[lhs], timing.alst[lhs]);
     Time rhsSlack = slackOrZero(timing.aest[rhs], timing.alst[rhs]);
     if (lhsSlack != rhsSlack)
@@ -168,19 +192,83 @@ std::vector<size_t> selectCriticalWindow(const TimingInfo& timing, size_t window
     if (timing.aest[lhs] != timing.aest[rhs])
       return timing.aest[lhs] < timing.aest[rhs];
     return lhs < rhs;
-  });
-  selected.resize(std::min(windowSize, selected.size()));
-  return selected;
-}
+  };
 
-std::vector<size_t> getOriginalSignature(const VirtualGraph& graph, ArrayRef<size_t> selectedNodes) {
-  std::vector<size_t> signature;
-  for (size_t nodeIndex : selectedNodes) {
-    const VirtualNode& node = graph.nodes[nodeIndex];
-    signature.insert(signature.end(), node.originalComputeIndices.begin(), node.originalComputeIndices.end());
+  windowSize = std::min(windowSize, ranked.size());
+  if (windowSize == 0)
+    return {};
+  if (windowSize == ranked.size()) {
+    llvm::sort(ranked, isHigherPriority);
+    return ranked;
   }
-  std::sort(signature.begin(), signature.end());
-  return signature;
+
+  size_t criticalPoolSize = std::min(ranked.size(), std::max(windowSize, windowSize * 2));
+  if (criticalPoolSize < ranked.size())
+    std::nth_element(
+      ranked.begin(), ranked.begin() + static_cast<std::ptrdiff_t>(criticalPoolSize), ranked.end(), isHigherPriority);
+
+  std::vector<char> inCriticalPool(ranked.size(), false);
+  for (size_t i = 0; i < criticalPoolSize; ++i)
+    inCriticalPool[ranked[i]] = true;
+
+  size_t seed = *std::min_element(ranked.begin(), ranked.end(), isHigherPriority);
+  std::vector<std::vector<size_t>> adjacency = buildUndirectedAdjacency(graph);
+  std::vector<size_t> selected;
+  std::vector<char> inWindow(ranked.size(), false);
+  selected.reserve(windowSize);
+
+  struct FrontierEntry {
+    size_t node;
+  };
+  auto frontierCompare = [&](FrontierEntry lhs, FrontierEntry rhs) { return isHigherPriority(rhs.node, lhs.node); };
+  std::priority_queue<FrontierEntry, std::vector<FrontierEntry>, decltype(frontierCompare)> frontier(frontierCompare);
+
+  auto addToWindow = [&](size_t node, const std::vector<char>& eligible) {
+    if (inWindow[node])
+      return;
+    inWindow[node] = true;
+    selected.push_back(node);
+    for (size_t neighbour : adjacency[node])
+      if (!inWindow[neighbour] && eligible[neighbour])
+        frontier.push({neighbour});
+  };
+
+  addToWindow(seed, inCriticalPool);
+  while (!frontier.empty() && selected.size() < windowSize) {
+    size_t node = frontier.top().node;
+    frontier.pop();
+    if (!inWindow[node])
+      addToWindow(node, inCriticalPool);
+  }
+
+  if (selected.size() < windowSize) {
+    std::vector<char> anyNode(ranked.size(), true);
+    for (size_t node : selected)
+      for (size_t neighbour : adjacency[node])
+        if (!inWindow[neighbour])
+          frontier.push({neighbour});
+    while (!frontier.empty() && selected.size() < windowSize) {
+      size_t node = frontier.top().node;
+      frontier.pop();
+      if (!inWindow[node])
+        addToWindow(node, anyNode);
+    }
+  }
+
+  if (selected.size() < windowSize) {
+    llvm::sort(ranked, isHigherPriority);
+    for (size_t node : ranked) {
+      if (selected.size() == windowSize)
+        break;
+      if (!inWindow[node]) {
+        inWindow[node] = true;
+        selected.push_back(node);
+      }
+    }
+  }
+
+  llvm::sort(selected, isHigherPriority);
+  return selected;
 }
 
 std::vector<IndexedEdge> buildWindowEdges(const VirtualGraph& graph, const std::vector<int64_t>& nodeToWindowIndex) {
@@ -216,25 +304,47 @@ WindowScheduleResult scheduleWindow(const VirtualGraph& graph, ArrayRef<size_t>
   windowGraph.runDcp();
 
   WindowScheduleResult result;
-  result.usedAllAvailableCpus = windowGraph.cpuCount() >= windowGraph.getMaxCpuCount();
+  result.cpuCount = windowGraph.cpuCount();
   for (CPU cpu = 0; cpu < windowGraph.cpuCount(); ++cpu) {
     auto scheduledTasks = windowGraph.getScheduledTasks(cpu);
     if (scheduledTasks.size() < 2)
       continue;
 
+    result.mergedNodeCount += scheduledTasks.size();
+    result.maxMergeGroupSize = std::max(result.maxMergeGroupSize, scheduledTasks.size());
     std::vector<size_t> mergeGroup;
     mergeGroup.reserve(scheduledTasks.size());
     for (const auto& task : scheduledTasks)
       mergeGroup.push_back(selectedNodes[task.nodeIndex]);
-    std::sort(mergeGroup.begin(), mergeGroup.end());
     result.mergeGroups.push_back(std::move(mergeGroup));
   }
   return result;
 }
 
-bool coarsenGraph(const VirtualGraph& graph, ArrayRef<std::vector<size_t>> mergeGroups, VirtualGraph& coarsenedGraph) {
+bool coarsenGraph(const VirtualGraph& graph,
+                  ArrayRef<std::vector<size_t>> mergeGroups,
+                  VirtualGraph& coarsenedGraph,
+                  std::vector<size_t>& oldToNewNode) {
+  TimingInfo timing = computeTiming(graph);
+  std::vector<size_t> topologicalRank(graph.nodes.size());
+  std::iota(topologicalRank.begin(), topologicalRank.end(), 0);
+  if (timing.valid)
+    for (auto [rank, nodeIndex] : llvm::enumerate(timing.topologicalOrder))
+      topologicalRank[nodeIndex] = rank;
+
+  std::vector<std::vector<size_t>> orderedMergeGroups;
+  orderedMergeGroups.reserve(mergeGroups.size());
+  for (const auto& mergeGroup : mergeGroups) {
+    orderedMergeGroups.emplace_back(mergeGroup.begin(), mergeGroup.end());
+    std::stable_sort(orderedMergeGroups.back().begin(), orderedMergeGroups.back().end(), [&](size_t lhs, size_t rhs) {
+      if (topologicalRank[lhs] != topologicalRank[rhs])
+        return topologicalRank[lhs] < topologicalRank[rhs];
+      return lhs < rhs;
+    });
+  }
+
   std::vector<int64_t> nodeToMergeGroup(graph.nodes.size(), -1);
-  for (auto [groupIndex, mergeGroup] : llvm::enumerate(mergeGroups)) {
+  for (auto [groupIndex, mergeGroup] : llvm::enumerate(orderedMergeGroups)) {
     if (mergeGroup.size() < 2)
       continue;
     for (size_t nodeIndex : mergeGroup) {
@@ -243,18 +353,21 @@ bool coarsenGraph(const VirtualGraph& graph, ArrayRef<std::vector<size_t>> merge
     }
   }
 
-  std::vector<std::optional<size_t>> mergeGroupToNewNode(mergeGroups.size());
-  std::vector<size_t> oldToNewNode(graph.nodes.size(), 0);
+  std::vector<std::optional<size_t>> mergeGroupToNewNode(orderedMergeGroups.size());
+  std::vector<size_t> newNodeRank;
+  oldToNewNode.assign(graph.nodes.size(), 0);
   bool mergedAny = false;
   coarsenedGraph.nodes.clear();
   coarsenedGraph.edges.clear();
   coarsenedGraph.nodes.reserve(graph.nodes.size());
+  newNodeRank.reserve(graph.nodes.size());
 
   for (size_t nodeIndex = 0; nodeIndex < graph.nodes.size(); ++nodeIndex) {
     int64_t mergeGroupIndex = nodeToMergeGroup[nodeIndex];
     if (mergeGroupIndex == -1) {
       oldToNewNode[nodeIndex] = coarsenedGraph.nodes.size();
       coarsenedGraph.nodes.push_back(graph.nodes[nodeIndex]);
+      newNodeRank.push_back(topologicalRank[nodeIndex]);
       continue;
     }
 
@@ -265,7 +378,7 @@ bool coarsenGraph(const VirtualGraph& graph, ArrayRef<std::vector<size_t>> merge
     }
 
     VirtualNode mergedNode;
-    for (size_t memberIndex : mergeGroups[static_cast<size_t>(mergeGroupIndex)]) {
+    for (size_t memberIndex : orderedMergeGroups[static_cast<size_t>(mergeGroupIndex)]) {
       const VirtualNode& memberNode = graph.nodes[memberIndex];
       mergedNode.originalComputeIndices.append(memberNode.originalComputeIndices.begin(),
                                                memberNode.originalComputeIndices.end());
@@ -276,8 +389,9 @@ bool coarsenGraph(const VirtualGraph& graph, ArrayRef<std::vector<size_t>> merge
 
     mergedAny = true;
     newNodeIndex = coarsenedGraph.nodes.size();
-    for (size_t memberIndex : mergeGroups[static_cast<size_t>(mergeGroupIndex)])
+    for (size_t memberIndex : orderedMergeGroups[static_cast<size_t>(mergeGroupIndex)])
       oldToNewNode[memberIndex] = *newNodeIndex;
+    newNodeRank.push_back(topologicalRank[orderedMergeGroups[static_cast<size_t>(mergeGroupIndex)].front()]);
     coarsenedGraph.nodes.push_back(std::move(mergedNode));
   }
 
@@ -291,75 +405,61 @@ bool coarsenGraph(const VirtualGraph& graph, ArrayRef<std::vector<size_t>> merge
     size_t newEnd = oldToNewNode[static_cast<size_t>(end)];
     if (newStart == newEnd)
       continue;
+    if (newNodeRank[newStart] >= newNodeRank[newEnd])
+      continue;
     remappedEdges.push_back({static_cast<int64_t>(newStart), static_cast<int64_t>(newEnd), weight});
   }
   coarsenedGraph.edges = aggregateEdges(remappedEdges);
 
-  return computeTiming(coarsenedGraph).valid;
+  return true;
 }
 
-bool coarsenGraphWithFallback(const VirtualGraph& graph,
-                              ArrayRef<std::vector<size_t>> mergeGroups,
-                              VirtualGraph& coarsenedGraph) {
-  if (coarsenGraph(graph, mergeGroups, coarsenedGraph))
-    return true;
+constexpr CPU kDefaultMaxCpuCount = 1000;
 
-  std::vector<size_t> orderedGroupIndices(mergeGroups.size());
-  std::iota(orderedGroupIndices.begin(), orderedGroupIndices.end(), 0);
-  std::stable_sort(orderedGroupIndices.begin(), orderedGroupIndices.end(), [&](size_t lhs, size_t rhs) {
-    return mergeGroups[lhs].size() > mergeGroups[rhs].size();
-  });
-
-  std::vector<std::vector<size_t>> acceptedMergeGroups;
-  acceptedMergeGroups.reserve(mergeGroups.size());
-  for (size_t groupIndex : orderedGroupIndices) {
-    std::vector<std::vector<size_t>> candidateMergeGroups = acceptedMergeGroups;
-    candidateMergeGroups.push_back(mergeGroups[groupIndex]);
-
-    VirtualGraph candidateGraph;
-    if (!coarsenGraph(graph, candidateMergeGroups, candidateGraph))
-      continue;
-
-    acceptedMergeGroups = std::move(candidateMergeGroups);
-    coarsenedGraph = std::move(candidateGraph);
-  }
-  return !acceptedMergeGroups.empty();
+CPU getVirtualGraphMaxCpuCount() {
+  if (coresCount.getValue() > 0)
+    return static_cast<CPU>(coresCount.getValue());
+  return kDefaultMaxCpuCount;
 }
 
-std::vector<size_t> computeOriginalTopologicalOrder(size_t computeCount, ArrayRef<IndexedEdge> edges) {
-  VirtualGraph graph;
-  graph.nodes.resize(computeCount);
-  graph.edges = aggregateEdges(edges);
-  TimingInfo timing = computeTiming(graph);
-  if (timing.valid)
-    return timing.topologicalOrder;
-
-  std::vector<size_t> fallbackOrder(computeCount);
-  std::iota(fallbackOrder.begin(), fallbackOrder.end(), 0);
-  return fallbackOrder;
+size_t getDcpCoarseningWindowSize(size_t nodeCount) {
+  size_t windowSize = std::min(dcpCriticalWindowSize.getValue(), nodeCount);
+  CPU maxCpuCount = std::max<CPU>(1, getVirtualGraphMaxCpuCount());
+  if (nodeCount > static_cast<size_t>(maxCpuCount))
+    windowSize = std::max(windowSize, std::min(nodeCount, static_cast<size_t>(maxCpuCount) + 1));
+  return windowSize;
 }
 
-DCPAnalysisResult buildResultFromVirtualGraph(const VirtualGraph& graph,
-                                              ArrayRef<SpatCompute> spatComputes,
-                                              ArrayRef<IndexedEdge> originalEdges) {
+DCPAnalysisResult buildResultFromVirtualGraph(const VirtualGraph& graph, ArrayRef<SpatCompute> spatComputes) {
   DCPAnalysisResult result;
-  std::vector<size_t> originalToVirtualNode(spatComputes.size(), 0);
-  for (auto [virtualNodeIndex, virtualNode] : llvm::enumerate(graph.nodes))
-    for (size_t originalIndex : virtualNode.originalComputeIndices)
-      originalToVirtualNode[originalIndex] = virtualNodeIndex;
 
-  auto dominanceOrder = computeOriginalTopologicalOrder(spatComputes.size(), originalEdges);
-  result.dominanceOrderCompute.reserve(dominanceOrder.size());
-  for (size_t originalIndex : dominanceOrder) {
-    SpatCompute spatCompute = spatComputes[originalIndex];
-    size_t cpu = originalToVirtualNode[originalIndex];
+  TimingInfo timing = computeTiming(graph);
+  std::vector<size_t> virtualNodeOrder;
+  if (timing.valid) {
+    virtualNodeOrder = std::move(timing.topologicalOrder);
+  }
+  else {
+    virtualNodeOrder.resize(graph.nodes.size());
+    std::iota(virtualNodeOrder.begin(), virtualNodeOrder.end(), 0);
+  }
+
+  std::vector<size_t> originalComputeToCpu(spatComputes.size(), 0);
+  for (auto [cpu, virtualNodeIndex] : llvm::enumerate(virtualNodeOrder)) {
+    const VirtualNode& virtualNode = graph.nodes[virtualNodeIndex];
+    for (size_t originalIndex : virtualNode.originalComputeIndices)
+      originalComputeToCpu[originalIndex] = cpu;
+  }
+
+  result.dominanceOrderCompute.reserve(spatComputes.size());
+  for (auto [originalIndex, spatCompute] : llvm::enumerate(spatComputes)) {
+    size_t cpu = originalComputeToCpu[originalIndex];
     result.dominanceOrderCompute.push_back(spatCompute);
     result.computeToCpuMap[spatCompute] = cpu;
     result.cpuToLastComputeMap[cpu] = spatCompute;
   }
-
-  for (auto [cpu, lastCompute] : result.cpuToLastComputeMap)
+  for (const auto& [cpu, lastCompute] : result.cpuToLastComputeMap)
     result.isLastComputeOfCpu.insert(lastCompute);
+
   return result;
 }
 
@@ -409,32 +509,74 @@ DCPAnalysisResult DCPAnalysis::run() {
     return runLegacyDcp(spatComputes, edges, entryOp->getContext());
 
   VirtualGraph virtualGraph = buildInitialVirtualGraph(spatComputes, edges);
-  std::set<std::vector<size_t>> seenCriticalWindows;
-  while (virtualGraph.nodes.size() > 1) {
-    TimingInfo timing = computeTiming(virtualGraph);
-    if (!timing.valid)
-      break;
-
-    auto selectedNodes = selectCriticalWindow(timing, dcpCriticalWindowSize.getValue());
-    if (selectedNodes.size() < 2)
-      break;
-
-    if (!seenCriticalWindows.insert(getOriginalSignature(virtualGraph, selectedNodes)).second)
-      break;
-
+  size_t iteration = 0;
+  auto tryCoarsenSelectedNodes = [&](ArrayRef<size_t> selectedNodes) {
+    size_t oldNodeCount = virtualGraph.nodes.size();
     WindowScheduleResult windowSchedule = scheduleWindow(virtualGraph, selectedNodes, entryOp->getContext());
-    if (windowSchedule.mergeGroups.empty())
-      break;
+    if (windowSchedule.mergeGroups.empty()) {
+      if (oldNodeCount >= 200)
+        llvm::errs() << llvm::formatv("[DCP-COARSEN] iter={0} old={1} selected={2} windowCpus={3} "
+                                      "groups=0 mergedNodes=0 maxGroup=0 new={1} changed=0\n",
+                                      iteration,
+                                      oldNodeCount,
+                                      selectedNodes.size(),
+                                      windowSchedule.cpuCount);
+      return false;
+    }
 
     VirtualGraph coarsenedGraph;
-    if (!coarsenGraphWithFallback(virtualGraph, windowSchedule.mergeGroups, coarsenedGraph))
-      break;
+    std::vector<size_t> oldToNewNode;
+    if (!coarsenGraph(virtualGraph, windowSchedule.mergeGroups, coarsenedGraph, oldToNewNode))
+      return false;
+    if (oldNodeCount >= 200 || coarsenedGraph.nodes.size() >= 200)
+      llvm::errs() << llvm::formatv("[DCP-COARSEN] iter={0} old={1} selected={2} windowCpus={3} "
+                                    "groups={4} mergedNodes={5} maxGroup={6} new={7} changed={8}\n",
+                                    iteration,
+                                    oldNodeCount,
+                                    selectedNodes.size(),
+                                    windowSchedule.cpuCount,
+                                    windowSchedule.mergeGroups.size(),
+                                    windowSchedule.mergedNodeCount,
+                                    windowSchedule.maxMergeGroupSize,
+                                    coarsenedGraph.nodes.size(),
+                                    oldNodeCount - coarsenedGraph.nodes.size());
     virtualGraph = std::move(coarsenedGraph);
-    if (windowSchedule.usedAllAvailableCpus)
+    return true;
+  };
+
+  while (virtualGraph.nodes.size() > 1) {
+    iteration++;
+    TimingInfo timing = computeTiming(virtualGraph);
+    if (!timing.valid) {
+      if (virtualGraph.nodes.size() >= 200)
+        llvm::errs() << llvm::formatv(
+          "[DCP-COARSEN] iter={0} old={1} invalid-timing\n", iteration, virtualGraph.nodes.size());
       break;
+    }
+
+    SmallVector<size_t> selectedNodes;
+    auto criticalWindow =
+      selectCriticalWindow(virtualGraph, timing, getDcpCoarseningWindowSize(virtualGraph.nodes.size()));
+    selectedNodes.append(criticalWindow.begin(), criticalWindow.end());
+
+    if (selectedNodes.size() < 2) {
+      if (virtualGraph.nodes.size() >= 200)
+        llvm::errs() << llvm::formatv("[DCP-COARSEN] iter={0} old={1} selected={2} stop=small-window\n",
+                                      iteration,
+                                      virtualGraph.nodes.size(),
+                                      selectedNodes.size());
+      break;
+    }
+
+    if (tryCoarsenSelectedNodes(selectedNodes))
+      continue;
+    if (virtualGraph.nodes.size() >= 200)
+      llvm::errs() << llvm::formatv(
+        "[DCP-COARSEN] iter={0} old={1} stop=no-merge\n", iteration, virtualGraph.nodes.size());
+    break;
   }
 
-  return buildResultFromVirtualGraph(virtualGraph, spatComputes, edges);
+  return buildResultFromVirtualGraph(virtualGraph, spatComputes);
 }
 
 } // namespace spatial
diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.cpp
index b24d3b1..6c48b1d 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.cpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.cpp
@@ -38,11 +38,14 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/ErrorHandling.h"
 
 #include <algorithm>
 #include <cassert>
 #include <chrono>
+#include <cstdint>
 #include <cstdio>
+#include <queue>
 #include <vector>
 
 #include "DCPAnalysis.hpp"
@@ -60,6 +63,7 @@ namespace {
 // Coarse-grained phase timers printed when DCP_SELECT_PROFILE is set.
 struct SelectTimers {
   double findSlot = 0.0;
+  double dedup = 0.0;
   double precheck = 0.0;
   double snapshotInsertUpdate = 0.0;
   double childSlot = 0.0;
@@ -70,9 +74,19 @@ struct SelectTimers {
   long tasksProcessed = 0;
   void dump(const char* label) const {
     std::fprintf(stderr,
-                 "[selectProfile:%s] tasks=%ld findSlot=%.2fs precheck=%.2fs snapUpd=%.2fs childSlot=%.2fs rollback=%.2fs iter=%ld precheckPass=%ld dcplPass=%ld\n",
-                 label, tasksProcessed, findSlot, precheck, snapshotInsertUpdate, childSlot,
-                 rollbackRestore, iterations, passedPrecheck, passedDcpl);
+                 "[selectProfile:%s] tasks=%ld dedup=%.2fs findSlot=%.2fs precheck=%.2fs snapUpd=%.2fs "
+                 "childSlot=%.2fs rollback=%.2fs iter=%ld precheckPass=%ld dcplPass=%ld\n",
+                 label,
+                 tasksProcessed,
+                 dedup,
+                 findSlot,
+                 precheck,
+                 snapshotInsertUpdate,
+                 childSlot,
+                 rollbackRestore,
+                 iterations,
+                 passedPrecheck,
+                 passedDcpl);
   }
   ~SelectTimers() {
     if (std::getenv("DCP_SELECT_PROFILE"))
@@ -83,6 +97,101 @@ static SelectTimers gSelectTimers;
 } // namespace
 #endif
 
+namespace {
+
+uint64_t mixHash(uint64_t seed, uint64_t value) {
+  seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2);
+  return seed;
+}
+
+uint64_t finishHash(uint64_t seed) {
+  seed ^= seed >> 33;
+  seed *= 0xff51afd7ed558ccdULL;
+  seed ^= seed >> 33;
+  seed *= 0xc4ceb9fe1a85ec53ULL;
+  seed ^= seed >> 33;
+  return seed;
+}
+
+uint64_t hashEdgeSignature(uint64_t neighborHash, Weight weight, uint64_t direction) {
+  uint64_t hash = mixHash(0x84222325cbf29ce4ULL, direction);
+  hash = mixHash(hash, neighborHash);
+  hash = mixHash(hash, static_cast<uint64_t>(weight));
+  return finishHash(hash);
+}
+
+struct CpuAestCache {
+  Time defaultAest = 0;
+  llvm::SmallDenseMap<CPU, Time, 8> colocatedParentAests;
+
+  Time get(CPU cpu) const {
+    auto it = colocatedParentAests.find(cpu);
+    if (it == colocatedParentAests.end())
+      return defaultAest;
+    return it->second;
+  }
+};
+
+struct CpuTimeMax {
+  CPU cpu = -1;
+  Time time = 0;
+};
+
+void updateCpuTimeMax(CpuTimeMax& first, CpuTimeMax& second, CPU cpu, Time time) {
+  if (first.cpu == cpu) {
+    first.time = std::max(first.time, time);
+    return;
+  }
+  if (second.cpu == cpu) {
+    second.time = std::max(second.time, time);
+    if (second.time > first.time)
+      std::swap(first, second);
+    return;
+  }
+  if (time >= first.time) {
+    second = first;
+    first = {cpu, time};
+    return;
+  }
+  if (time > second.time)
+    second = {cpu, time};
+}
+
+CpuAestCache computeCpuAestCache(TaskDCP* task) {
+  CpuAestCache cache;
+  llvm::SmallDenseMap<CPU, Time, 8> transferAestByCpu;
+  llvm::SmallDenseMap<CPU, Time, 8> localAestByCpu;
+  Time unscheduledTransferAest = 0;
+
+  for (const Edge& parentEdge : task->parents) {
+    Time parentFinish = addOrMax(parentEdge.first->getAest(), parentEdge.first->getWeight());
+    Time transferAest = addOrMax(parentFinish, getTransferCost(parentEdge.first, task));
+    if (std::optional<CPU> parentCpu = parentEdge.first->getCpu()) {
+      Time& cpuTransferAest = transferAestByCpu[*parentCpu];
+      cpuTransferAest = std::max(cpuTransferAest, transferAest);
+      Time& cpuLocalAest = localAestByCpu[*parentCpu];
+      cpuLocalAest = std::max(cpuLocalAest, parentFinish);
+      continue;
+    }
+    unscheduledTransferAest = std::max(unscheduledTransferAest, transferAest);
+  }
+
+  CpuTimeMax firstOther {-1, unscheduledTransferAest};
+  CpuTimeMax secondOther {-1, 0};
+  for (const auto& entry : transferAestByCpu)
+    updateCpuTimeMax(firstOther, secondOther, entry.first, entry.second);
+
+  cache.defaultAest = firstOther.time;
+  for (const auto& entry : localAestByCpu) {
+    CPU cpu = entry.first;
+    Time bestNonLocalParentAest = firstOther.cpu == cpu ? secondOther.time : firstOther.time;
+    cache.colocatedParentAests[cpu] = std::max(bestNonLocalParentAest, entry.second);
+  }
+  return cache;
+}
+
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // Edge manipulation
 //===----------------------------------------------------------------------===//
@@ -156,6 +265,49 @@ std::vector<TaskDCP*> GraphDCP::getRoots() {
   return tmp;
 }
 
+void GraphDCP::initTaskStructureHashes() {
+  taskStructureHashes.resize(nodes.size());
+  for (auto [index, task] : llvm::enumerate(nodes)) {
+    uint64_t hash = mixHash(0x7442b1129fd01363ULL, static_cast<uint64_t>(task.getWeight()));
+    hash = mixHash(hash, static_cast<uint64_t>(task.getCrossbarUsage()));
+    taskStructureHashes[index] = finishHash(hash);
+  }
+
+  std::vector<uint64_t> nextHashes(nodes.size());
+  std::vector<uint64_t> edgeHashes;
+  for (int iteration = 0; iteration < 4; ++iteration) {
+    for (auto [index, task] : llvm::enumerate(nodes)) {
+      uint64_t hash = mixHash(0x464dcab27ac82291ULL, taskStructureHashes[index]);
+      edgeHashes.clear();
+      edgeHashes.reserve(task.parents.size() + task.children.size());
+      for (const Edge& parent : task.parents)
+        if (!parent.isScheduling)
+          edgeHashes.push_back(
+            hashEdgeSignature(taskStructureHashes[getNodeIndex(parent.first)], parent.second, /*direction=*/0));
+      for (const Edge& child : task.children)
+        if (!child.isScheduling)
+          edgeHashes.push_back(
+            hashEdgeSignature(taskStructureHashes[getNodeIndex(child.first)], child.second, /*direction=*/1));
+      llvm::sort(edgeHashes);
+      hash = mixHash(hash, static_cast<uint64_t>(edgeHashes.size()));
+      for (uint64_t edgeHash : edgeHashes)
+        hash = mixHash(hash, edgeHash);
+      nextHashes[index] = finishHash(hash);
+    }
+    taskStructureHashes.swap(nextHashes);
+  }
+}
+
+// Compact dedup key for CPU `c` vs `candidate`: mixes candidateAest, crossbar
+// usage, and the incremental cpu structure hash. No heap allocation.
+uint64_t GraphDCP::computeCpuCandidateKey(Time candidateAest, CPU cpu) {
+  uint64_t hash = mixHash(0xd6e8feb86659fd93ULL, static_cast<uint64_t>(candidateAest));
+  hash = mixHash(hash, static_cast<uint64_t>(getCpuCrossbarUsage(cpu)));
+  auto it = cpuStructureHashes.find(cpu);
+  hash = mixHash(hash, it != cpuStructureHashes.end() ? it->second : 0ULL);
+  return finishHash(hash);
+}
+
 // Inserts `task` at `position` on `cpu`, wiring up scheduling edges with the
 // neighbouring tasks and keeping the global topological order consistent.
 TaskInsertion GraphDCP::insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position) {
@@ -164,6 +316,7 @@ TaskInsertion GraphDCP::insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position)
   task->setCpu(cpu);
   task->setWeight(scheduledWeight);
   reserveTaskCrossbars(cpu, task);
+  cpuStructureHashes[cpu] ^= taskStructureHashes[getNodeIndex(task)];
   auto& tasksInCpu = getOrCreateCpuTasks(cpu);
   unsigned int numCpuTasks = tasksInCpu.size();
   assert(position <= numCpuTasks && "Inserting in a not valid position");
@@ -201,6 +354,7 @@ TaskInsertion GraphDCP::insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position)
 
 void GraphDCP::removeTaskFromCPU(CPU cpu, TaskDCP* task) {
   releaseTaskCrossbars(cpu, task);
+  cpuStructureHashes[cpu] ^= taskStructureHashes[getNodeIndex(task)];
   task->resetCpu();
   task->resetWeight();
   auto& scheduledTasks = getOrCreateCpuTasks(cpu);
@@ -271,6 +425,21 @@ bool GraphDCP::wouldExhaustCrossbarCapacity(CPU cpu, const TaskDCP* task) const
   return nextUsage >= getCpuCrossbarCapacity();
 }
 
+size_t GraphDCP::crossbarsUsed() const {
+  CrossbarUsage crossbarEdge = static_cast<CrossbarUsage>(onnx_mlir::crossbarSize.getValue());
+  CrossbarUsage crossbarArea = crossbarEdge * crossbarEdge;
+  if (crossbarArea == 0)
+    return 0;
+  CrossbarUsage totalArea = 0;
+  for (const auto& [cpu, usage] : cpuCrossbarUsage)
+    totalArea = checkedAdd(totalArea, usage);
+  return static_cast<size_t>(totalArea / crossbarArea);
+}
+
+size_t GraphDCP::crossbarsAvailable() const {
+  return static_cast<size_t>(lastCpu) * onnx_mlir::crossbarCountInCore.getValue();
+}
+
 //===----------------------------------------------------------------------===//
 // AEST / ALST computation
 //===----------------------------------------------------------------------===//
@@ -456,9 +625,9 @@ void GraphDCP::updateAestFromTaskWithDescendants(TaskDCP* task, llvm::ArrayRef<T
   for (TaskDCP* descendant : descendantsTopoOrder)
     recomputeAest(descendant);
 
-  const bool oldMaxInvalidated = maxCompletionTask != nullptr
-                              && (maxCompletionTask == task
-                                  || llvm::is_contained(descendantsTopoOrder, maxCompletionTask));
+  const bool oldMaxInvalidated =
+    maxCompletionTask != nullptr
+    && (maxCompletionTask == task || llvm::is_contained(descendantsTopoOrder, maxCompletionTask));
   if (oldMaxInvalidated) {
     // The pre-update max came from a modified task; its completion has moved
     // upward, so modifiedMaxCompletion is an upper bound covering it. The
@@ -523,9 +692,9 @@ bool GraphDCP::tryUpdateAestWithinBudget(TaskDCP* task,
     if (!process(descendant))
       return false;
 
-  const bool oldMaxInvalidated = maxCompletionTask != nullptr
-                              && (maxCompletionTask == task
-                                  || llvm::is_contained(descendantsTopoOrder, maxCompletionTask));
+  const bool oldMaxInvalidated =
+    maxCompletionTask != nullptr
+    && (maxCompletionTask == task || llvm::is_contained(descendantsTopoOrder, maxCompletionTask));
   if (oldMaxInvalidated) {
     dcpl = modifiedMaxCompletion;
     maxCompletion = modifiedMaxCompletion;
@@ -546,6 +715,109 @@ bool GraphDCP::tryUpdateAestWithinBudget(TaskDCP* task,
   return true;
 }
 
+// Incrementally refreshes ALST after `task` was placed. The set of nodes whose
+// ALST is structurally affected by the insertion is exactly
+// `relations.ancestors ∪ {task}`: the task's outgoing transfer costs to
+// same-CPU real children become 0, and new scheduling edges create parent
+// relationships between `task` and its same-CPU neighbors. Every other node
+// keeps its relative distance to the sink boundary and only absorbs the
+// signed DCPL delta captured between `oldDcpl` and the now-updated `dcpl`.
+void GraphDCP::updateAlstFromScheduledTask(TaskDCP* task, const CandidateRelations& relations, Time oldDcpl) {
+  Time newDcpl = getDcpl();
+  // If the AEST update saturated dcpl (e.g. rescue placement on a
+  // crossbar-exhausted CPU sets task weight to UINT64_MAX), the shift delta
+  // would be meaningless. Fall back to a full recompute for this step only.
+  if (newDcpl == std::numeric_limits<Time>::max()) {
+    initAlst();
+    return;
+  }
+
+  if (newDcpl != oldDcpl) {
+    const bool increased = newDcpl > oldDcpl;
+    const Time delta = increased ? (newDcpl - oldDcpl) : (oldDcpl - newDcpl);
+    for (TaskDCP& node : topologicalOrder) {
+      if (&node == task || relations.ancestors.contains(&node))
+        continue;
+      Time alst = node.getAlst();
+      node.setAlst(increased ? addOrMax(alst, delta) : subtractOrZero(alst, delta));
+    }
+  }
+
+  auto recomputeAlst = [&](TaskDCP* node) {
+    Time minAlst = std::numeric_limits<Time>::max();
+    if (!node->hasChildren())
+      minAlst = subtractOrZero(newDcpl, node->getWeight());
+    for (const Edge& childEdge : node->children)
+      minAlst = std::min(minAlst,
+                         subtractOrZero(childEdge.first->getAlst(),
+                                        addOrMax(node->getWeight(), getTransferCost(node, childEdge.first))));
+    node->setAlst(minAlst);
+  };
+
+  // Walk the backward cone with a pending-children counter so that every
+  // ancestor is recomputed only after all of its affected children have
+  // been refreshed. This is resilient to staleness in the global
+  // `topologicalOrder` relative to freshly added scheduling edges.
+  llvm::DenseSet<TaskDCP*> affected = relations.ancestors;
+  affected.insert(task);
+  llvm::DenseMap<TaskDCP*, int> pendingAffectedChildren;
+  pendingAffectedChildren.reserve(affected.size());
+  std::vector<TaskDCP*> worklist;
+  worklist.reserve(affected.size());
+  for (TaskDCP* node : affected) {
+    int count = 0;
+    for (const Edge& childEdge : node->children)
+      if (affected.contains(childEdge.first))
+        count++;
+    pendingAffectedChildren[node] = count;
+    if (count == 0)
+      worklist.push_back(node);
+  }
+  while (!worklist.empty()) {
+    TaskDCP* node = worklist.back();
+    worklist.pop_back();
+    recomputeAlst(node);
+    for (const Edge& parentEdge : node->parents) {
+      if (!affected.contains(parentEdge.first))
+        continue;
+      auto it = pendingAffectedChildren.find(parentEdge.first);
+      assert(it != pendingAffectedChildren.end());
+      if (--it->second == 0)
+        worklist.push_back(parentEdge.first);
+    }
+  }
+
+// Opt-in consistency check: verifies the incremental ALST result against a
+// full initAlst() recomputation. Very expensive (O(V+E) per placement) - only
+// enable when investigating suspected drift.
+#ifdef DCP_DEBUG_CHECK_ALST
+  std::vector<Time> afterIncremental(nodes.size());
+  for (size_t i = 0; i < nodes.size(); ++i)
+    afterIncremental[i] = nodes[i].getAlst();
+  initAlst();
+  bool mismatched = false;
+  for (size_t i = 0; i < nodes.size(); ++i) {
+    if (afterIncremental[i] != nodes[i].getAlst()) {
+      if (!mismatched) {
+        llvm::errs() << "[alst-mismatch] placed=" << getNodeIndex(task) << " oldDcpl=" << oldDcpl
+                     << " newDcpl=" << newDcpl << " ancestors={";
+        for (TaskDCP* a : relations.ancestors)
+          llvm::errs() << getNodeIndex(a) << ",";
+        llvm::errs() << "}\n";
+        mismatched = true;
+      }
+      llvm::errs() << "  node=" << i << " incremental=" << afterIncremental[i] << " full=" << nodes[i].getAlst()
+                   << " weight=" << nodes[i].getWeight()
+                   << " cpu=" << (nodes[i].isScheduled() ? (int) *nodes[i].getCpu() : -1) << " children=[";
+      for (const Edge& e : nodes[i].children)
+        llvm::errs() << getNodeIndex(e.first) << (e.isScheduling ? "s" : "")
+                     << "(tc=" << getTransferCost(&nodes[i], e.first) << ",alst=" << e.first->getAlst() << "),";
+      llvm::errs() << "]\n";
+    }
+  }
+#endif
+}
+
 // Computes a localised ALST: only ancestors of the candidate (plus the
 // candidate itself) get recomputed, every other task keeps its current ALST.
 // Processes nodes in reverse dependency order using a pending-children
@@ -905,32 +1177,6 @@ GraphDCP::FindSlot GraphDCP::findSlotWithFixedFinalTime(
 // Candidate selection and processor assignment
 //===----------------------------------------------------------------------===//
 
-// Lowest slack wins; earliest AEST breaks ties. Critical-path tasks (zero
-// slack) naturally float to the front.
-TaskDCP* GraphDCP::findCandidate(const std::vector<TaskDCP*>& readyNodes) {
-  auto findBestNode = [](auto lft, auto rgt) {
-    Time leftSlack = slackOrZero((*lft)->getAest(), (*lft)->getAlst());
-    Time rightSlack = slackOrZero((*rgt)->getAest(), (*rgt)->getAlst());
-    if (leftSlack < rightSlack)
-      return lft;
-    if (rightSlack < leftSlack)
-      return rgt;
-    if ((*lft)->getAest() < (*rgt)->getAest())
-      return lft;
-    return rgt;
-  };
-
-  assert(!readyNodes.empty() && "expected at least one ready node");
-  auto validNode = readyNodes.begin();
-  auto bestNode = validNode;
-
-  while (validNode != readyNodes.end()) {
-    bestNode = findBestNode(validNode, bestNode);
-    std::advance(validNode, 1);
-  }
-  return *bestNode;
-}
-
 // Picks the best CPU + slot for `candidate`:
 //   * Phase 1 (parallel, read-only): call findSlot on every candidate CPU.
 //   * Phase 2 (sequential): process CPUs in ascending slot.aest order. For
@@ -939,7 +1185,7 @@ TaskDCP* GraphDCP::findCandidate(const std::vector<TaskDCP*>& readyNodes) {
 //     evaluate a slot for the smallest-slack child, then roll back.
 //   * Rescue (sequential): if nothing fit, grow the CPU count if allowed,
 //     otherwise pick the CPU that leads to the smallest DCPL increase.
-void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
+GraphDCP::CandidateRelations GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
   CandidateRelations relations = dcp_graph::computeCandidateRelations(candidate);
   relations.descendantsTopoOrder.reserve(relations.descendants.size());
   for (auto it = candidate->getTopologicalIterator(); it != topologicalOrder.end(); ++it) {
@@ -959,22 +1205,43 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
   const CrossbarUsage candidateFootprint = getTaskCrossbarFootprint(candidate);
   const bool candidateHasCrossbar = candidateFootprint != 0;
   const CrossbarUsage cpuCapacity = candidateHasCrossbar ? getCpuCrossbarCapacity() : 0;
+  DCP_DEBUG_IF(auto dedupStart = std::chrono::steady_clock::now();)
+  CpuAestCache cpuAests = computeCpuAestCache(candidate);
+  DCP_DEBUG_IF(const bool checkCpuAestCache = std::getenv("DCP_CHECK_CPU_AEST_CACHE") != nullptr;)
+  llvm::SmallDenseSet<uint64_t, 32> seenProcessorKeys;
+  seenProcessorKeys.reserve(static_cast<size_t>(topCpu + 1));
   for (CPU c = 0; c <= topCpu; c++) {
     if (candidateHasCrossbar && c != getLastCpu()) {
       CrossbarUsage nextUsage = checkedAdd(getCpuCrossbarUsage(c), candidateFootprint);
       if (nextUsage >= cpuCapacity)
         continue;
     }
+    Time candidateAest = cpuAests.get(c);
+    DCP_DEBUG_IF(if (checkCpuAestCache) {
+      Time recomputedAest = computeAestOnCpu(candidate, c);
+      if (candidateAest != recomputedAest) {
+        std::fprintf(stderr,
+                     "[DCP_CHECK_CPU_AEST_CACHE] mismatch candidate=%zu cpu=%d cached=%llu recomputed=%llu\n",
+                     getNodeIndex(candidate),
+                     c,
+                     static_cast<unsigned long long>(candidateAest),
+                     static_cast<unsigned long long>(recomputedAest));
+        llvm::report_fatal_error("DCP CPU AEST cache mismatch");
+      }
+    })
+    if (!seenProcessorKeys.insert(computeCpuCandidateKey(candidateAest, c)).second)
+      continue;
     processors.push_back(c);
   }
-
+  DCP_DEBUG_IF(gSelectTimers.dedup +=
+               std::chrono::duration<double>(std::chrono::steady_clock::now() - dedupStart).count();)
   if (processors.empty()) {
-    CPU bestCpu = canCreateNewCpu ? getLastCpu() : 0;
-    FindSlot bestSlot = {computeAestOnCpu(candidate, bestCpu), static_cast<int>(getOrCreateCpuTasks(bestCpu).size())};
-    if (canCreateNewCpu)
-      incrementLastCpu();
-    insertTaskInCPU(bestCpu, candidate, bestSlot.index);
-    return;
+    // processors.empty() implies !canCreateNewCpu: a fresh CPU always passes
+    // the crossbar filter and would have been added. Reaching here means every
+    // existing CPU is crossbar-exhausted and the task requires crossbar
+    // capacity — the placement is impossible.
+    llvm::report_fatal_error("DCP scheduler: crossbar capacity exhausted on all CPUs; "
+                             "cannot schedule task that requires crossbar allocation");
   }
 
   // Phase 1: parallel findSlot sweep (read-only over graph state).
@@ -1000,21 +1267,20 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
     for (size_t i = 0; i < processors.size(); ++i)
       sweep(i);
   DCP_DEBUG_IF(gSelectTimers.findSlot +=
-                 std::chrono::duration<double>(std::chrono::steady_clock::now() - sweepStart).count();)
+               std::chrono::duration<double>(std::chrono::steady_clock::now() - sweepStart).count();)
 
 #ifdef DCP_DEBUG_ENABLED
   {
     static bool reported = false;
     if (!reported) {
       reported = true;
-      std::fprintf(stderr,
-                   "[dcp] selectProcessor parallel sweep: context=%p mt=%d procs=%zu pool=%u\n",
-                   (void*) context,
-                   context != nullptr ? (int) context->isMultithreadingEnabled() : -1,
-                   processors.size(),
-                   context != nullptr && context->isMultithreadingEnabled()
-                     ? context->getThreadPool().getMaxConcurrency()
-                     : 0u);
+      std::fprintf(
+        stderr,
+        "[dcp] selectProcessor parallel sweep: context=%p mt=%d procs=%zu pool=%u\n",
+        (void*) context,
+        context != nullptr ? (int) context->isMultithreadingEnabled() : -1,
+        processors.size(),
+        context != nullptr && context->isMultithreadingEnabled() ? context->getThreadPool().getMaxConcurrency() : 0u);
     }
   }
 #endif
@@ -1055,9 +1321,10 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
       DCP_DEBUG_IF(auto t2 = std::chrono::steady_clock::now();)
       Weight candidateWeight = candidate->computeWeightOnCpu(this, currentCpu);
       Time candidateCompletion = addOrMax(slot.aest, candidateWeight);
-      bool skip = (!emptyCpu && candidateCompletion > currentDcpl)
-               || addOrMax(slot.aest, candidateCompletion) >= bestComposite;
-      DCP_DEBUG_IF(gSelectTimers.precheck += std::chrono::duration<double>(std::chrono::steady_clock::now() - t2).count();)
+      bool skip =
+        (!emptyCpu && candidateCompletion > currentDcpl) || addOrMax(slot.aest, candidateCompletion) >= bestComposite;
+      DCP_DEBUG_IF(gSelectTimers.precheck +=
+                   std::chrono::duration<double>(std::chrono::steady_clock::now() - t2).count();)
       if (skip)
         continue;
       DCP_DEBUG_IF(++gSelectTimers.passedPrecheck;)
@@ -1073,8 +1340,8 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         scheduleSnapshot = dcp_graph::captureLocalScheduleState(
           candidate, relations.descendants, dcpl, maxCompletion, secondMaxCompletion, maxCompletionTask);
         taskInsertion = insertTaskInCPU(currentCpu, candidate, slot.index);
-        bool withinBudget = tryUpdateAestWithinBudget(
-          candidate, llvm::ArrayRef<TaskDCP*>(relations.descendantsTopoOrder), currentDcpl);
+        bool withinBudget =
+          tryUpdateAestWithinBudget(candidate, llvm::ArrayRef<TaskDCP*>(relations.descendantsTopoOrder), currentDcpl);
         if (!withinBudget) {
           DCP_DEBUG_IF(auto t4 = std::chrono::steady_clock::now();)
           taskInsertion.rollBack();
@@ -1087,7 +1354,7 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         }
       }
       DCP_DEBUG_IF(gSelectTimers.snapshotInsertUpdate +=
-                     std::chrono::duration<double>(std::chrono::steady_clock::now() - t3).count();)
+                   std::chrono::duration<double>(std::chrono::steady_clock::now() - t3).count();)
       DCP_DEBUG_IF(++gSelectTimers.passedDcpl;)
 
       // Pick the tightest unscheduled child (smallest slack) and measure what
@@ -1135,7 +1402,7 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         dcp_graph::restoreLocalScheduleState(
           scheduleSnapshot, dcpl, maxCompletion, secondMaxCompletion, maxCompletionTask);
       DCP_DEBUG_IF(gSelectTimers.rollbackRestore +=
-                     std::chrono::duration<double>(std::chrono::steady_clock::now() - t6).count();)
+                   std::chrono::duration<double>(std::chrono::steady_clock::now() - t6).count();)
     }
   }
 
@@ -1150,7 +1417,9 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
     else {
       Time bestDcpl = std::numeric_limits<Time>::max();
       Time currentDcpl = getDcpl();
-      for (CPU c = 0; c < getLastCpu(); c++) {
+      for (CPU c : processors) {
+        if (c == getLastCpu())
+          continue;
         auto slot = findSlot(candidate, c, false, relations);
         if (slot.aest == std::numeric_limits<Time>::max())
           slot = findSlot(candidate, c, true, relations);
@@ -1159,8 +1428,7 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         // Cheap lower bound: post-insertion DCPL is at least max(currentDcpl,
         // candidate completion on this slot). Skip CPUs already worse than
         // the best seen.
-        Time lowerBound =
-          std::max(currentDcpl, addOrMax(slot.aest, candidate->computeWeightOnCpu(this, c)));
+        Time lowerBound = std::max(currentDcpl, addOrMax(slot.aest, candidate->computeWeightOnCpu(this, c)));
         if (lowerBound >= bestDcpl)
           continue;
         auto snapshot = dcp_graph::captureLocalScheduleState(
@@ -1169,23 +1437,37 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         updateAestFromTaskWithDescendants(candidate, llvm::ArrayRef<TaskDCP*>(relations.descendantsTopoOrder));
         Time candidateDcpl = getDcpl();
         taskInsertion.rollBack();
-        dcp_graph::restoreLocalScheduleState(
-          snapshot, dcpl, maxCompletion, secondMaxCompletion, maxCompletionTask);
+        dcp_graph::restoreLocalScheduleState(snapshot, dcpl, maxCompletion, secondMaxCompletion, maxCompletionTask);
         if (candidateDcpl < bestDcpl) {
           bestDcpl = candidateDcpl;
           bestCpu = c;
           bestSlot = slot;
         }
       }
-      if (bestCpu == -1) {
-        bestCpu = 0;
-        bestSlot = {computeAestOnCpu(candidate, bestCpu), static_cast<int>(getOrCreateCpuTasks(bestCpu).size())};
-      }
+      if (bestCpu == -1)
+        llvm::report_fatal_error("DCP scheduler: no valid slot found for task on any eligible CPU — "
+                                 "all slots are blocked by already-placed descendants");
     }
   }
   if (bestCpu == getLastCpu() && getLastCpu() < maxCpuCount)
     incrementLastCpu();
   insertTaskInCPU(bestCpu, candidate, bestSlot.index);
+
+  // Incremental AEST/ALST refresh replacing the full initAest/initAlst that
+  // used to run after every placement. Post-insertion relations pick up any
+  // new scheduling-edge ancestors/descendants introduced by the insertion.
+  Time oldDcpl = getDcpl();
+  CandidateRelations postRelations = dcp_graph::computeCandidateRelations(candidate);
+  llvm::SmallVector<TaskDCP*, 32> postDescendantsTopoOrder;
+  postDescendantsTopoOrder.reserve(postRelations.descendants.size());
+  for (auto it = candidate->getTopologicalIterator(); it != topologicalOrder.end(); ++it) {
+    TaskDCP* current = &*it;
+    if (current != candidate && postRelations.descendants.contains(current))
+      postDescendantsTopoOrder.push_back(current);
+  }
+  updateAestFromTaskWithDescendants(candidate, llvm::ArrayRef<TaskDCP*>(postDescendantsTopoOrder));
+  updateAlstFromScheduledTask(candidate, postRelations, oldDcpl);
+  return postRelations;
 }
 
 //===----------------------------------------------------------------------===//
@@ -1194,61 +1476,99 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
 
 void GraphDCP::runDcp() {
   initTopological();
+  initTaskStructureHashes();
   initAest();
   initAlst();
   dumpDot();
 
   dcp_graph::DcpProgressLogger progressLogger(nodes.size());
   llvm::DenseMap<TaskDCP*, int> unscheduledParents;
-  std::vector<TaskDCP*> readyNodes;
-  readyNodes.reserve(nodes.size());
+
+  // Min-heap over ready tasks: tightest slack first, earliest AEST as tiebreak.
+  // Lazy deletion: when AEST/ALST change after a placement, fresh entries are
+  // pushed for the affected tasks. Stale ones are detected on pop by comparing
+  // stored vs current (slack, aest) and re-pushed with the current values.
+  struct ReadyEntry {
+    Time slack;
+    Time aest;
+    TaskDCP* task;
+    bool operator>(const ReadyEntry& other) const {
+      if (slack != other.slack)
+        return slack > other.slack;
+      return aest > other.aest;
+    }
+  };
+  std::priority_queue<ReadyEntry, std::vector<ReadyEntry>, std::greater<ReadyEntry>> readyQueue;
+  size_t readyCount = 0;
+
+  auto pushReady = [&](TaskDCP* node) {
+    readyQueue.push({slackOrZero(node->getAest(), node->getAlst()), node->getAest(), node});
+  };
+
   for (auto& node : nodes) {
     int dependencyParents = dcp_graph::countDependencyParents(&node);
     unscheduledParents[&node] = dependencyParents;
-    if (dependencyParents == 0)
-      readyNodes.push_back(&node);
+    if (dependencyParents == 0) {
+      pushReady(&node);
+      ++readyCount;
+    }
   }
-  progressLogger.printStart(readyNodes.size());
+  size_t xbarsCapacity = static_cast<size_t>(maxCpuCount) * onnx_mlir::crossbarCountInCore.getValue();
+  progressLogger.printStart(readyCount, maxCpuCount, xbarsCapacity);
 
-  while (!readyNodes.empty()) {
-    DCP_DEBUG_IF(auto findStart = std::chrono::steady_clock::now();)
-    TaskDCP* candidate = findCandidate(readyNodes);
-    DCP_DEBUG_IF(progressLogger.recordFindDuration(
-      std::chrono::duration<double>(std::chrono::steady_clock::now() - findStart).count());)
-    fastRemove(readyNodes, candidate);
+  while (readyCount > 0) {
+    // Pop with lazy deletion: skip stale entries and re-push with current values.
+    TaskDCP* candidate = nullptr;
+    while (!readyQueue.empty()) {
+      auto entry = readyQueue.top();
+      readyQueue.pop();
+      Time curSlack = slackOrZero(entry.task->getAest(), entry.task->getAlst());
+      Time curAest = entry.task->getAest();
+      if (entry.slack == curSlack && entry.aest == curAest) {
+        candidate = entry.task;
+        break;
+      }
+      readyQueue.push({curSlack, curAest, entry.task});
+    }
+    assert(candidate != nullptr && "readyCount > 0 but heap exhausted");
+    --readyCount;
 
     DCP_DEBUG_IF(auto selectStart = std::chrono::steady_clock::now();)
-    selectProcessor(candidate, candidate->isCriticalPath());
+    CandidateRelations postRelations = selectProcessor(candidate, candidate->isCriticalPath());
     DCP_DEBUG_IF(
       double selectSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - selectStart).count();
       progressLogger.recordSelectDuration(selectSeconds);
-      progressLogger.maybePrintSlowCandidate(getNodeIndex(candidate), selectSeconds, readyNodes.size(), getLastCpu());
-    )
+      progressLogger.maybePrintSlowCandidate(getNodeIndex(candidate), selectSeconds, readyCount, getLastCpu());)
+
+    // Proactively refresh the heap priority for ready nodes whose AEST or ALST
+    // changed: ancestors had their ALST individually recomputed; descendants had
+    // their AEST bumped. Both may now sort differently than their stale entries.
+    for (TaskDCP* node : postRelations.ancestors)
+      if (!node->isScheduled() && unscheduledParents[node] == 0)
+        pushReady(node);
+    for (TaskDCP* node : postRelations.descendants)
+      if (!node->isScheduled() && unscheduledParents[node] == 0)
+        pushReady(node);
 
-    DCP_DEBUG_IF(auto updateStart = std::chrono::steady_clock::now();)
-    initAest();
-    initAlst();
-    DCP_DEBUG_IF(progressLogger.recordUpdateDuration(
-      std::chrono::duration<double>(std::chrono::steady_clock::now() - updateStart).count());)
     progressLogger.advanceCompleted();
-    progressLogger.printProgress(readyNodes.size(), getLastCpu(), "recompute", false);
+    progressLogger.printProgress(readyCount, getLastCpu(), maxCpuCount, crossbarsUsed(), crossbarsAvailable(), false);
 
     for (const auto& childEdge : candidate->children) {
       if (childEdge.isScheduling || childEdge.first->isScheduled())
         continue;
       int& dependencyParents = unscheduledParents[childEdge.first];
       assert(dependencyParents > 0 && "dependency parent count must stay positive");
-      dependencyParents--;
-      if (dependencyParents == 0)
-        readyNodes.push_back(childEdge.first);
+      --dependencyParents;
+      if (dependencyParents == 0) {
+        pushReady(childEdge.first);
+        ++readyCount;
+      }
     }
-    DCP_DEBUG_IF(
-      ++gSelectTimers.tasksProcessed;
-      if (std::getenv("DCP_SELECT_PROFILE") && (gSelectTimers.tasksProcessed % 100 == 0))
-        gSelectTimers.dump("tick");
-    )
+    DCP_DEBUG_IF(++gSelectTimers.tasksProcessed;
+                 if (std::getenv("DCP_SELECT_PROFILE") && (gSelectTimers.tasksProcessed % 100 == 0))
+                   gSelectTimers.dump("tick");)
   }
-  progressLogger.printProgress(readyNodes.size(), getLastCpu(), "done", true);
+  progressLogger.printProgress(0, getLastCpu(), maxCpuCount, crossbarsUsed(), crossbarsAvailable(), true);
   dumpDot();
 }
 
diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.hpp
index 7e9f4a5..e8303d6 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.hpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.hpp
@@ -4,6 +4,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 
+#include <cstdint>
 #include <list>
 #include <optional>
 #include <unordered_map>
@@ -48,8 +49,10 @@ private:
 
   std::vector<TaskDCP> nodes;
   onnx_mlir::LabeledList<TaskDCP> topologicalOrder;
+  std::vector<uint64_t> taskStructureHashes;
   std::vector<CpuTaskList> cpuTasks;
   std::unordered_map<CPU, CrossbarUsage> cpuCrossbarUsage;
+  llvm::DenseMap<CPU, uint64_t> cpuStructureHashes;
   CPU lastCpu = 0;
   long long flag = 1;
   Time dcpl = 0;
@@ -70,6 +73,7 @@ private:
 
   void initAest();
   void initAlst();
+  void initTaskStructureHashes();
 
   Time computeAestOnCpu(TaskDCP* task, CPU cpu);
   Time computeDcplOnCpu(TaskDCP* task, CPU cpu);
@@ -83,9 +87,15 @@ private:
   // `dcplBudget`, signalling that the new DCPL would exceed the budget.
   // Returns true iff the full propagation completed without exceeding the
   // budget. Uses the caller's snapshot to restore AEST on the aborted tail.
-  bool tryUpdateAestWithinBudget(TaskDCP* task,
-                                 llvm::ArrayRef<TaskDCP*> descendantsTopoOrder,
-                                 Time dcplBudget);
+  bool tryUpdateAestWithinBudget(TaskDCP* task, llvm::ArrayRef<TaskDCP*> descendantsTopoOrder, Time dcplBudget);
+
+  // Incrementally refreshes ALST after `task` has been scheduled. Nodes
+  // outside the backward cone (`relations.ancestors` plus `task`) retain
+  // their relative distance to the sink boundary and only absorb the signed
+  // DCPL delta (`newDcpl - oldDcpl`). `task` itself and its ancestors are
+  // recomputed in reverse topological order so that new same-CPU transfer
+  // costs (now zero) and scheduling-edge children are reflected.
+  void updateAlstFromScheduledTask(TaskDCP* task, const CandidateRelations& relations, Time oldDcpl);
 
   void initTopological();
   void topologicalMoveAfter(TaskDCP* task, TaskDCP* pivotPoint, TaskInsertion* insertion = nullptr);
@@ -94,8 +104,11 @@ private:
   llvm::DenseMap<TaskDCP*, Time> computeAlst(TaskDCP* task, CPU cpu, const CandidateRelations& relations);
   size_t getNodeIndex(const TaskDCP* task) const;
 
-  TaskDCP* findCandidate(const std::vector<TaskDCP*>& readyNodes);
-  void selectProcessor(TaskDCP* candidate, bool push);
+  // Returns a compact dedup key for CPU `c` when evaluating `candidate`:
+  // mixes candidateAest, crossbar usage, and the incremental cpu structure
+  // hash into a single uint64_t. Zero heap allocation.
+  uint64_t computeCpuCandidateKey(Time candidateAest, CPU cpu);
+  CandidateRelations selectProcessor(TaskDCP* candidate, bool push);
   CPU getLastCpu() const { return lastCpu; }
   void incrementLastCpu() { lastCpu++; }
   FindSlot findSlot(TaskDCP* candidate, CPU cpu, bool push, const CandidateRelations& relations);
@@ -115,8 +128,7 @@ private:
 
 public:
   void runDcp();
-  GraphDCP(llvm::ArrayRef<onnx_mlir::spatial::SpatCompute> spatComputes,
-           llvm::ArrayRef<IndexedEdge> edges)
+  GraphDCP(llvm::ArrayRef<onnx_mlir::spatial::SpatCompute> spatComputes, llvm::ArrayRef<IndexedEdge> edges)
   : nodes(), cpuTasks(), cpuCrossbarUsage() {
     for (auto spatCompute : spatComputes)
       nodes.emplace_back(spatCompute);
@@ -150,6 +162,11 @@ public:
   void setMaxCpuCount(int value) { maxCpuCount = value; }
   int getMaxCpuCount() const { return maxCpuCount; }
 
+  // Total crossbar units allocated across all active CPUs.
+  size_t crossbarsUsed() const;
+  // Maximum crossbar units available across all active CPUs (lastCpu * per-CPU capacity).
+  size_t crossbarsAvailable() const;
+
   // Optional MLIR context used to drive mlir::parallelFor inside runDcp. If
   // null the scheduler runs single-threaded (tests use this path).
   void setContext(mlir::MLIRContext* ctx) { context = ctx; }
diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/GraphDebug.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/GraphDebug.cpp
index 7f3a89d..1f907d1 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/GraphDebug.cpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/GraphDebug.cpp
@@ -35,10 +35,12 @@ void DcpProgressLogger::recordSelectDuration(double seconds) { selectProcessorSe
 void DcpProgressLogger::recordUpdateDuration(double seconds) { updateTimingSeconds += seconds; }
 void DcpProgressLogger::advanceCompleted(size_t taskCount) { completedTasks += taskCount; }
 
-void DcpProgressLogger::printStart(size_t readyCount) const {
+void DcpProgressLogger::printStart(size_t readyCount, int maxCpuCount, size_t xbarsCapacity) const {
   if (!logProgress)
     return;
-  llvm::errs() << llvm::formatv("[DCP] start: tasks={0} ready={1}\n", totalTasks, readyCount);
+  llvm::errs() << llvm::formatv(
+    "[DCP] start   tasks={0}   ready={1}   cpus=0/{2}   crossbars=0/{3}\n",
+    totalTasks, readyCount, maxCpuCount, xbarsCapacity);
 }
 
 void DcpProgressLogger::maybePrintSlowCandidate(size_t nodeIndex,
@@ -48,14 +50,15 @@ void DcpProgressLogger::maybePrintSlowCandidate(size_t nodeIndex,
   if (!logProgress || elapsedSeconds < 1.0)
     return;
 
-  llvm::errs() << llvm::formatv("[DCP] slow candidate node={0} elapsed={1} ready={2} cpus={3}\n",
+  llvm::errs() << llvm::formatv("[DCP] slow node={0}   elapsed={1}   ready={2}   cpus={3}\n",
                                 nodeIndex,
                                 formatDuration(elapsedSeconds),
                                 readyCount,
                                 cpuCount);
 }
 
-void DcpProgressLogger::printProgress(size_t readyCount, CPU cpuCount, llvm::StringRef stage, bool force) {
+void DcpProgressLogger::printProgress(
+  size_t readyCount, CPU cpuCount, int maxCpuCount, size_t xbarsUsed, size_t xbarsAvailable, bool force) {
   if (!logProgress)
     return;
 
@@ -68,19 +71,19 @@ void DcpProgressLogger::printProgress(size_t readyCount, CPU cpuCount, llvm::Str
   double etaSeconds = rate > 0.0 ? static_cast<double>(totalTasks - completedTasks) / rate : 0.0;
   double percent = totalTasks == 0 ? 100.0 : (100.0 * static_cast<double>(completedTasks) / totalTasks);
 
-  llvm::errs() << llvm::formatv("[DCP] {0}/{1} ({2:F1}%) ready={3} cpus={4} stage={5} elapsed={6} eta={7}\n",
-                                completedTasks,
-                                totalTasks,
-                                percent,
-                                readyCount,
-                                cpuCount,
-                                stage,
-                                formatDuration(elapsedSeconds),
-                                completedTasks == totalTasks ? "0:00" : formatDuration(etaSeconds));
-  llvm::errs() << llvm::formatv("        time(find={0}, select={1}, update={2})\n",
-                                formatDuration(findCandidateSeconds),
-                                formatDuration(selectProcessorSeconds),
-                                formatDuration(updateTimingSeconds));
+  bool done = completedTasks == totalTasks;
+  llvm::errs() << llvm::formatv(
+    "[DCP] {0}/{1} ({2:F0}%)   ready={3}   cpus={4}/{5}   crossbars={6}/{7}   {8}{9}\n",
+    completedTasks,
+    totalTasks,
+    percent,
+    readyCount,
+    cpuCount,
+    maxCpuCount,
+    xbarsUsed,
+    xbarsAvailable,
+    llvm::formatv("elapsed={0}", formatDuration(elapsedSeconds)).str(),
+    done ? "" : llvm::formatv("   eta={0}", formatDuration(etaSeconds)).str());
   lastProgressPrint = now;
 }
 
@@ -91,9 +94,9 @@ void DcpProgressLogger::recordFindDuration(double) {}
 void DcpProgressLogger::recordSelectDuration(double) {}
 void DcpProgressLogger::recordUpdateDuration(double) {}
 void DcpProgressLogger::advanceCompleted(size_t) {}
-void DcpProgressLogger::printStart(size_t) const {}
+void DcpProgressLogger::printStart(size_t, int, size_t) const {}
 void DcpProgressLogger::maybePrintSlowCandidate(size_t, double, size_t, CPU) const {}
-void DcpProgressLogger::printProgress(size_t, CPU, llvm::StringRef, bool) {}
+void DcpProgressLogger::printProgress(size_t, CPU, int, size_t, size_t, bool) {}
 
 #endif
 
diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/GraphDebug.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/GraphDebug.hpp
index 380f9df..14e3783 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/GraphDebug.hpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/GraphDebug.hpp
@@ -31,9 +31,10 @@ public:
   void recordUpdateDuration(double seconds);
   void advanceCompleted(size_t taskCount = 1);
 
-  void printStart(size_t readyCount) const;
+  void printStart(size_t readyCount, int maxCpuCount, size_t xbarsCapacity) const;
   void maybePrintSlowCandidate(size_t nodeIndex, double elapsedSeconds, size_t readyCount, CPU cpuCount) const;
-  void printProgress(size_t readyCount, CPU cpuCount, llvm::StringRef stage, bool force);
+  void printProgress(size_t readyCount, CPU cpuCount, int maxCpuCount,
+                     size_t xbarsUsed, size_t xbarsAvailable, bool force);
 
 #ifdef DCP_DEBUG_ENABLED
 private:
diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
index 976f418..89acae3 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
@@ -1,4 +1,5 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/PatternMatch.h"
@@ -10,23 +11,30 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_os_ostream.h"
+#include "llvm/Support/raw_ostream.h"
 
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
+#include <cstdlib>
 #include <fstream>
 #include <functional>
 #include <iterator>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <tuple>
+#include <utility>
 #include <vector>
 
 #include "DCPGraph/DCPAnalysis.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
 using namespace mlir;
 
@@ -34,6 +42,541 @@ namespace onnx_mlir {
 namespace {
 using SpatCompute = spatial::SpatCompute;
 
+SpatCompute getOriginalSpatCompute(Operation* op) {
+  while (auto extract = dyn_cast_if_present<tensor::ExtractSliceOp>(op))
+    op = extract.getSource().getDefiningOp();
+  return dyn_cast_if_present<SpatCompute>(op);
+}
+
+struct ComputeMotifInfo {
+  uint64_t instructionCount = 0;
+  uint64_t weightedMvmCount = 0;
+  uint64_t weightedVmmCount = 0;
+};
+
+void appendUnique(SmallVector<size_t>& values, size_t value) {
+  if (!llvm::is_contained(values, value))
+    values.push_back(value);
+}
+
+bool isTrivialSerialMergeCandidate(SpatCompute compute) {
+  if (!compute->hasOneUse())
+    return false;
+
+  auto& use = *compute->getUses().begin();
+  auto user = dyn_cast<SpatCompute>(use.getOwner());
+  return user && user.getInputs().size() == 1 && use.getOperandNumber() >= user.getWeights().size();
+}
+
+SmallVector<size_t> appendMissingWeightsAndBuildIndexMap(SpatCompute target, ValueRange sourceWeights) {
+  DenseMap<Value, SmallVector<size_t, 4>> targetWeightIndices;
+  for (auto [weightIndex, weight] : llvm::enumerate(target.getWeights()))
+    targetWeightIndices[weight].push_back(weightIndex);
+
+  DenseMap<Value, size_t> usedSourceWeightOccurrences;
+  SmallVector<size_t> sourceToTargetIndex;
+  sourceToTargetIndex.reserve(sourceWeights.size());
+  auto targetWeights = target.getWeightsMutable();
+  for (Value weight : sourceWeights) {
+    size_t occurrence = usedSourceWeightOccurrences[weight]++;
+    auto& matchingIndices = targetWeightIndices[weight];
+    if (occurrence >= matchingIndices.size()) {
+      size_t newIndex = target.getWeights().size();
+      targetWeights.append(weight);
+      matchingIndices.push_back(newIndex);
+      sourceToTargetIndex.push_back(newIndex);
+      continue;
+    }
+    sourceToTargetIndex.push_back(matchingIndices[occurrence]);
+  }
+  return sourceToTargetIndex;
+}
+
+void mergeTriviallyConnectedComputes(func::FuncOp funcOp) {
+  Location loc = funcOp.getLoc();
+  IRRewriter rewriter(funcOp->getContext());
+  SmallVector<SpatCompute> trivialComputes;
+  llvm::SmallSet<SpatCompute, 8> toErase;
+
+  for (auto compute : funcOp.getOps<SpatCompute>())
+    if (isTrivialSerialMergeCandidate(compute))
+      trivialComputes.push_back(compute);
+
+  while (!trivialComputes.empty()) {
+    auto compute = trivialComputes.front();
+
+    if (compute.use_empty()) {
+      std::swap(trivialComputes.front(), trivialComputes.back());
+      trivialComputes.pop_back();
+      continue;
+    }
+
+    auto& computeUse = *compute->getUses().begin();
+    auto child = cast<SpatCompute>(computeUse.getOwner());
+    auto usedResult = cast<OpResult>(computeUse.get()).getResultNumber();
+    auto childArgIndex = computeUse.getOperandNumber() - child.getWeights().size();
+
+    rewriter.setInsertionPointAfter(compute.getOperation());
+    auto newCompute = SpatCompute::create(rewriter, loc, child.getResultTypes(), compute.getOperands());
+    newCompute.getProperties().setOperandSegmentSizes(
+      {static_cast<int>(compute.getWeights().size()), static_cast<int>(compute.getInputs().size())});
+
+    IRMapping mapper;
+    SmallVector<size_t> childWeightToNewIndex = appendMissingWeightsAndBuildIndexMap(newCompute, child.getWeights());
+    for (auto [oldIndex, weight] : llvm::enumerate(child.getWeights()))
+      mapper.map(weight, *std::next(newCompute.getWeights().begin(), childWeightToNewIndex[oldIndex]));
+
+    compute.getBodyRegion().cloneInto(&newCompute.getBodyRegion(), mapper);
+    auto newTerminator = newCompute.getBody().front().getTerminator();
+    mapper.map(child.getBody().front().getArgument(childArgIndex), newTerminator->getOperand(usedResult));
+    newTerminator->erase();
+
+    rewriter.setInsertionPoint(&newCompute.getBody().front(), newCompute.getBody().front().end());
+    auto remapWeightIndex = [&](auto weightedOp) {
+      auto oldIndex = weightedOp.getWeightIndex();
+      assert(static_cast<size_t>(oldIndex) < childWeightToNewIndex.size() && "weight index out of range");
+      weightedOp.setWeightIndex(childWeightToNewIndex[oldIndex]);
+    };
+
+    for (auto& op : child.getBody().front()) {
+      auto newInst = rewriter.clone(op, mapper);
+      if (auto weightedMvmOp = dyn_cast<spatial::SpatWeightedMVMOp>(newInst))
+        remapWeightIndex(weightedMvmOp);
+      if (auto weightedVmmOp = dyn_cast<spatial::SpatWeightedVMMOp>(newInst))
+        remapWeightIndex(weightedVmmOp);
+    }
+
+    child.replaceAllUsesWith(newCompute);
+    toErase.insert(child);
+
+    std::swap(trivialComputes.front(), trivialComputes.back());
+    trivialComputes.pop_back();
+    toErase.insert(compute);
+
+    if (isTrivialSerialMergeCandidate(newCompute))
+      trivialComputes.push_back(newCompute);
+  }
+
+  for (auto compute : toErase) {
+    for (Value result : compute->getResults())
+      result.dropAllUses();
+    compute.erase();
+  }
+}
+
+struct WeightedVmmBandCandidate {
+  Operation* parent;
+  SpatCompute compute;
+};
+
+bool isSingleWeightedVmmCompute(SpatCompute compute) {
+  if (compute.getNumResults() != 1 || compute.getWeights().size() != 1 || compute.getInputs().size() != 1)
+    return false;
+
+  uint64_t weightedVmmCount = 0;
+  for (Operation& op : compute.getBody().front()) {
+    if (isa<spatial::SpatWeightedMVMOp>(&op))
+      return false;
+    if (isa<spatial::SpatWeightedVMMOp>(&op))
+      weightedVmmCount++;
+  }
+  return weightedVmmCount == 1;
+}
+
+std::optional<WeightedVmmBandCandidate> getWeightedVmmBandCandidate(SpatCompute compute) {
+  if (!isSingleWeightedVmmCompute(compute) || !compute->hasOneUse())
+    return std::nullopt;
+
+  auto& use = *compute->getUses().begin();
+  auto child = dyn_cast<SpatCompute>(use.getOwner());
+  if (!child || use.getOperandNumber() < child.getWeights().size())
+    return std::nullopt;
+
+  auto parent = getOriginalSpatCompute(compute.getInputs().front().getDefiningOp());
+  if (!parent || parent == child)
+    return std::nullopt;
+
+  return WeightedVmmBandCandidate {parent.getOperation(), compute};
+}
+
+bool haveSameWeightedVmmBandShape(SpatCompute lhs, SpatCompute rhs) {
+  return lhs.getWeights().front().getType() == rhs.getWeights().front().getType()
+      && lhs.getInputs().front().getType() == rhs.getInputs().front().getType()
+      && lhs.getResult(0).getType() == rhs.getResult(0).getType();
+}
+
+SpatCompute packWeightedVmmComputes(func::FuncOp funcOp, ArrayRef<SpatCompute> computes) {
+  assert(!computes.empty() && "expected at least one compute to pack");
+
+  IRRewriter rewriter(funcOp->getContext());
+  SpatCompute child = cast<SpatCompute>((*computes.front()->getUses().begin()).getOwner());
+  rewriter.setInsertionPoint(child);
+
+  SmallVector<Value> operands;
+  SmallVector<Type> inputTypes;
+  SmallVector<Location> inputLocs;
+  SmallVector<Type> resultTypes;
+  operands.reserve(computes.size() * 2);
+  inputTypes.reserve(computes.size());
+  inputLocs.reserve(computes.size());
+  resultTypes.reserve(computes.size());
+
+  for (SpatCompute compute : computes)
+    for (Value weight : compute.getWeights())
+      operands.push_back(weight);
+  for (SpatCompute compute : computes) {
+    for (Value input : compute.getInputs()) {
+      operands.push_back(input);
+      inputTypes.push_back(input.getType());
+      inputLocs.push_back(input.getLoc());
+    }
+    for (Type resultType : compute.getResultTypes())
+      resultTypes.push_back(resultType);
+  }
+
+  auto packedCompute = SpatCompute::create(rewriter, funcOp.getLoc(), resultTypes, operands);
+  packedCompute.getProperties().setOperandSegmentSizes(
+    {static_cast<int>(computes.size()), static_cast<int>(inputTypes.size())});
+  auto* block = rewriter.createBlock(&packedCompute.getBody(), packedCompute.getBody().end(), inputTypes, inputLocs);
+  rewriter.setInsertionPointToEnd(block);
+
+  SmallVector<Value> yieldValues;
+  yieldValues.reserve(resultTypes.size());
+  size_t inputBaseIndex = 0;
+  size_t weightBaseIndex = 0;
+  for (SpatCompute compute : computes) {
+    IRMapping mapper;
+    for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights()))
+      mapper.map(weight, *std::next(packedCompute.getWeights().begin(), weightBaseIndex + weightIndex));
+    for (auto [inputIndex, bbArg] : llvm::enumerate(compute.getBody().front().getArguments()))
+      mapper.map(bbArg, block->getArgument(inputBaseIndex + inputIndex));
+
+    auto remapWeightIndex = [&](auto weightedOp) {
+      weightedOp.setWeightIndex(weightBaseIndex + weightedOp.getWeightIndex());
+    };
+    for (Operation& op : compute.getBody().front()) {
+      if (auto yield = dyn_cast<spatial::SpatYieldOp>(&op)) {
+        for (Value yieldOperand : yield.getOperands())
+          yieldValues.push_back(mapper.lookup(yieldOperand));
+        continue;
+      }
+
+      Operation* cloned = rewriter.clone(op, mapper);
+      if (auto weightedMvmOp = dyn_cast<spatial::SpatWeightedMVMOp>(cloned))
+        remapWeightIndex(weightedMvmOp);
+      if (auto weightedVmmOp = dyn_cast<spatial::SpatWeightedVMMOp>(cloned))
+        remapWeightIndex(weightedVmmOp);
+    }
+
+    inputBaseIndex += compute.getInputs().size();
+    weightBaseIndex += compute.getWeights().size();
+  }
+  spatial::SpatYieldOp::create(rewriter, funcOp.getLoc(), yieldValues);
+
+  size_t resultIndex = 0;
+  for (SpatCompute compute : computes)
+    for (OpResult result : compute->getResults())
+      result.replaceAllUsesWith(packedCompute.getResult(resultIndex++));
+  for (SpatCompute compute : llvm::reverse(computes))
+    compute.erase();
+
+  return packedCompute;
+}
+
+size_t getFastPathCpuBudget() {
+  constexpr size_t kDefaultMaxCpuCount = 1000;
+  if (coresCount.getValue() > 0)
+    return static_cast<size_t>(coresCount.getValue());
+  return kDefaultMaxCpuCount;
+}
+
+size_t packWideWeightedVmmBands(func::FuncOp funcOp) {
+  constexpr size_t kMinBandSize = 64;
+  size_t cpuBudget = std::max<size_t>(1, getFastPathCpuBudget());
+  SmallVector<WeightedVmmBandCandidate> candidates;
+  for (SpatCompute compute : funcOp.getOps<SpatCompute>())
+    if (auto candidate = getWeightedVmmBandCandidate(compute))
+      candidates.push_back(*candidate);
+
+  llvm::stable_sort(candidates, [](const WeightedVmmBandCandidate& lhs, const WeightedVmmBandCandidate& rhs) {
+    if (lhs.parent != rhs.parent)
+      return lhs.parent < rhs.parent;
+    return lhs.compute->isBeforeInBlock(rhs.compute);
+  });
+
+  size_t packedBandCount = 0;
+  size_t packedNodeCount = 0;
+  size_t createdNodeCount = 0;
+  size_t minChunkSize = std::numeric_limits<size_t>::max();
+  size_t maxChunkSize = 0;
+  SmallVector<SmallVector<SpatCompute>> shapeGroups;
+  for (size_t parentBegin = 0; parentBegin < candidates.size();) {
+    size_t parentEnd = parentBegin + 1;
+    while (parentEnd < candidates.size() && candidates[parentEnd].parent == candidates[parentBegin].parent)
+      parentEnd++;
+
+    shapeGroups.clear();
+    for (size_t index = parentBegin; index < parentEnd; ++index) {
+      SpatCompute compute = candidates[index].compute;
+      auto groupIt = llvm::find_if(shapeGroups, [&](const SmallVector<SpatCompute>& group) {
+        return haveSameWeightedVmmBandShape(group.front(), compute);
+      });
+      if (groupIt == shapeGroups.end())
+        shapeGroups.push_back({compute});
+      else
+        groupIt->push_back(compute);
+    }
+
+    for (ArrayRef<SpatCompute> band : shapeGroups) {
+      size_t bandSize = band.size();
+      if (bandSize < kMinBandSize || bandSize <= cpuBudget)
+        continue;
+
+      size_t chunkSize = (bandSize + cpuBudget - 1) / cpuBudget;
+      packedBandCount++;
+      SmallVector<SpatCompute> chunk;
+      chunk.reserve(std::min(chunkSize, bandSize));
+      for (auto [index, compute] : llvm::enumerate(band)) {
+        chunk.push_back(compute);
+        if (chunk.size() == chunkSize || index + 1 == band.size()) {
+          minChunkSize = std::min(minChunkSize, chunk.size());
+          maxChunkSize = std::max(maxChunkSize, chunk.size());
+          packWeightedVmmComputes(funcOp, chunk);
+          packedNodeCount += chunk.size();
+          createdNodeCount++;
+          chunk.clear();
+        }
+      }
+    }
+    parentBegin = parentEnd;
+  }
+
+  if (packedNodeCount != 0)
+    llvm::errs() << llvm::formatv("[DCP-FASTPATH] wvmmBands={0} packedNodes={1} createdNodes={2} changed={3} "
+                                  "cpuBudget={4} chunkSizeRange={5}-{6}\n",
+                                  packedBandCount,
+                                  packedNodeCount,
+                                  createdNodeCount,
+                                  packedNodeCount - createdNodeCount,
+                                  cpuBudget,
+                                  minChunkSize,
+                                  maxChunkSize);
+  return packedNodeCount - createdNodeCount;
+}
+
+void emitMotifProfile(func::FuncOp funcOp) {
+  if (!std::getenv("DCP_MOTIF_PROFILE"))
+    return;
+
+  SmallVector<SpatCompute> computes(funcOp.getOps<SpatCompute>());
+  DenseMap<SpatCompute, size_t> computeToIndex;
+  computeToIndex.reserve(computes.size());
+  for (auto [index, compute] : llvm::enumerate(computes))
+    computeToIndex[compute] = index;
+
+  SmallVector<ComputeMotifInfo> computeInfos(computes.size());
+  SmallVector<SmallVector<size_t>> parents(computes.size());
+  SmallVector<SmallVector<size_t>> children(computes.size());
+
+  uint64_t weightedVmmNodeCount = 0;
+  uint64_t weightedVmmOpCount = 0;
+  uint64_t edgeCount = 0;
+
+  for (auto [index, compute] : llvm::enumerate(computes)) {
+    ComputeMotifInfo& info = computeInfos[index];
+    for (Operation& op : compute.getBody().front()) {
+      info.instructionCount++;
+      if (isa<spatial::SpatWeightedMVMOp>(&op))
+        info.weightedMvmCount++;
+      if (isa<spatial::SpatWeightedVMMOp>(&op))
+        info.weightedVmmCount++;
+    }
+    if (info.weightedVmmCount > 0) {
+      weightedVmmNodeCount++;
+      weightedVmmOpCount += info.weightedVmmCount;
+    }
+
+    for (Value input : compute.getInputs()) {
+      auto parent = getOriginalSpatCompute(input.getDefiningOp());
+      if (!parent || parent == compute)
+        continue;
+      auto parentIt = computeToIndex.find(parent);
+      if (parentIt == computeToIndex.end())
+        continue;
+
+      size_t parentIndex = parentIt->second;
+      size_t oldParentCount = parents[index].size();
+      appendUnique(parents[index], parentIndex);
+      if (parents[index].size() != oldParentCount) {
+        appendUnique(children[parentIndex], index);
+        edgeCount++;
+      }
+    }
+  }
+
+  uint64_t maxFanIn = 0;
+  uint64_t maxFanOut = 0;
+  uint64_t fanIn16 = 0;
+  uint64_t fanIn64 = 0;
+  uint64_t fanIn256 = 0;
+  uint64_t fanOut16 = 0;
+  uint64_t fanOut64 = 0;
+  uint64_t fanOut256 = 0;
+  for (size_t index = 0; index < computes.size(); ++index) {
+    uint64_t fanIn = parents[index].size();
+    uint64_t fanOut = children[index].size();
+    maxFanIn = std::max(maxFanIn, fanIn);
+    maxFanOut = std::max(maxFanOut, fanOut);
+    fanIn16 += fanIn >= 16;
+    fanIn64 += fanIn >= 64;
+    fanIn256 += fanIn >= 256;
+    fanOut16 += fanOut >= 16;
+    fanOut64 += fanOut >= 64;
+    fanOut256 += fanOut >= 256;
+  }
+
+  uint64_t serialChainCount = 0;
+  uint64_t serialChainNodeCount = 0;
+  uint64_t maxSerialChain = 0;
+  for (size_t index = 0; index < computes.size(); ++index) {
+    if (parents[index].size() == 1 && children[parents[index][0]].size() == 1)
+      continue;
+
+    uint64_t chainLength = 1;
+    size_t current = index;
+    while (children[current].size() == 1) {
+      size_t child = children[current][0];
+      if (parents[child].size() != 1)
+        break;
+      chainLength++;
+      current = child;
+    }
+
+    if (chainLength >= 2) {
+      serialChainCount++;
+      serialChainNodeCount += chainLength;
+      maxSerialChain = std::max(maxSerialChain, chainLength);
+    }
+  }
+
+  SmallVector<size_t> incomingEdgeCount;
+  incomingEdgeCount.reserve(parents.size());
+  for (ArrayRef<size_t> parentList : parents)
+    incomingEdgeCount.push_back(parentList.size());
+
+  SmallVector<uint64_t> level(computes.size(), 0);
+  SmallVector<size_t> readyNodes;
+  readyNodes.reserve(computes.size());
+  for (size_t index = 0; index < computes.size(); ++index)
+    if (incomingEdgeCount[index] == 0)
+      readyNodes.push_back(index);
+
+  size_t readyIndex = 0;
+  while (readyIndex != readyNodes.size()) {
+    size_t current = readyNodes[readyIndex++];
+    for (size_t child : children[current]) {
+      level[child] = std::max(level[child], level[current] + 1);
+      assert(incomingEdgeCount[child] > 0 && "incoming edge count underflow");
+      incomingEdgeCount[child]--;
+      if (incomingEdgeCount[child] == 0)
+        readyNodes.push_back(child);
+    }
+  }
+
+  SmallVector<uint64_t> weightedVmmNodesByLevel;
+  for (size_t index = 0; index < computes.size(); ++index) {
+    if (computeInfos[index].weightedVmmCount == 0)
+      continue;
+    if (weightedVmmNodesByLevel.size() <= level[index])
+      weightedVmmNodesByLevel.resize(level[index] + 1, 0);
+    weightedVmmNodesByLevel[level[index]]++;
+  }
+
+  uint64_t maxWeightedVmmLevel = 0;
+  uint64_t wideWeightedVmmLevels64 = 0;
+  uint64_t wideWeightedVmmLevels256 = 0;
+  for (uint64_t count : weightedVmmNodesByLevel) {
+    maxWeightedVmmLevel = std::max(maxWeightedVmmLevel, count);
+    wideWeightedVmmLevels64 += count >= 64;
+    wideWeightedVmmLevels256 += count >= 256;
+  }
+
+  using ShapeKey = std::tuple<uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t>;
+  SmallVector<ShapeKey> weightedVmmShapeKeys;
+  for (auto [index, compute] : llvm::enumerate(computes)) {
+    const ComputeMotifInfo& info = computeInfos[index];
+    if (info.weightedVmmCount == 0)
+      continue;
+    weightedVmmShapeKeys.push_back({info.instructionCount,
+                                    info.weightedVmmCount,
+                                    info.weightedMvmCount,
+                                    static_cast<uint64_t>(compute.getWeights().size()),
+                                    static_cast<uint64_t>(compute.getInputs().size()),
+                                    static_cast<uint64_t>(parents[index].size()),
+                                    static_cast<uint64_t>(children[index].size())});
+  }
+
+  llvm::sort(weightedVmmShapeKeys);
+  SmallVector<std::pair<uint64_t, ShapeKey>> weightedVmmShapeCounts;
+  for (size_t index = 0; index < weightedVmmShapeKeys.size();) {
+    size_t next = index + 1;
+    while (next < weightedVmmShapeKeys.size() && weightedVmmShapeKeys[next] == weightedVmmShapeKeys[index])
+      next++;
+    weightedVmmShapeCounts.push_back({next - index, weightedVmmShapeKeys[index]});
+    index = next;
+  }
+  llvm::sort(weightedVmmShapeCounts, [](const auto& lhs, const auto& rhs) {
+    if (lhs.first != rhs.first)
+      return lhs.first > rhs.first;
+    return lhs.second < rhs.second;
+  });
+
+  llvm::errs() << llvm::formatv("[DCP-MOTIF] computes={0} edges={1} wvmmNodes={2} wvmmOps={3} "
+                                "serialChains={4} serialChainNodes={5} maxSerialChain={6} "
+                                "maxFanIn={7} maxFanOut={8} fanIn>=16/64/256={9}/{10}/{11} "
+                                "fanOut>=16/64/256={12}/{13}/{14} topoVisited={15}\n",
+                                computes.size(),
+                                edgeCount,
+                                weightedVmmNodeCount,
+                                weightedVmmOpCount,
+                                serialChainCount,
+                                serialChainNodeCount,
+                                maxSerialChain,
+                                maxFanIn,
+                                maxFanOut,
+                                fanIn16,
+                                fanIn64,
+                                fanIn256,
+                                fanOut16,
+                                fanOut64,
+                                fanOut256,
+                                readyNodes.size());
+
+  llvm::errs() << llvm::formatv("[DCP-MOTIF] wvmmLevels={0} maxWvmmLevel={1} wideWvmmLevels>=64/256={2}/{3} "
+                                "shapeGroups={4}\n",
+                                weightedVmmNodesByLevel.size(),
+                                maxWeightedVmmLevel,
+                                wideWeightedVmmLevels64,
+                                wideWeightedVmmLevels256,
+                                weightedVmmShapeCounts.size());
+
+  for (size_t rank = 0, end = std::min<size_t>(weightedVmmShapeCounts.size(), 5); rank < end; ++rank) {
+    auto [count, shape] = weightedVmmShapeCounts[rank];
+    auto [insts, vmmOps, mvmOps, weights, inputs, fanIn, fanOut] = shape;
+    llvm::errs() << llvm::formatv("[DCP-MOTIF] wvmmShape rank={0} count={1} insts={2} vmmOps={3} "
+                                  "mvmOps={4} weights={5} inputs={6} fanIn={7} fanOut={8}\n",
+                                  rank,
+                                  count,
+                                  insts,
+                                  vmmOps,
+                                  mvmOps,
+                                  weights,
+                                  inputs,
+                                  fanIn,
+                                  fanOut);
+  }
+}
+
 void generateReport(func::FuncOp funcOp, const std::string& name) {
   std::string outputDir = getOutputDir();
   if (outputDir.empty())
@@ -95,23 +638,22 @@ void generateReport(func::FuncOp funcOp, const std::string& name) {
     auto expectedPrintedValue = currentComputeId + 1;
     bool rangePrinted = false;
     cI++;
-    for (; cI <= lastIndex; ++cI){
+    for (; cI <= lastIndex; ++cI) {
       auto candidateToPrint = std::get<0>(collectedData[cI]);
-      if (candidateToPrint == expectedPrintedValue){
+      if (candidateToPrint == expectedPrintedValue) {
         expectedPrintedValue = candidateToPrint + 1;
         rangePrinted = true;
-      } else {
-        if (rangePrinted) {
+      }
+      else {
+        if (rangePrinted)
           os << " - " << expectedPrintedValue - 1;
-        }
         os << " , " << candidateToPrint;
         rangePrinted = false;
         expectedPrintedValue = candidateToPrint + 1;
       }
     }
-    if (rangePrinted && currentComputeId != expectedPrintedValue - 1){
-          os << " - " << expectedPrintedValue - 1;
-    }
+    if (rangePrinted && currentComputeId != expectedPrintedValue - 1)
+      os << " - " << expectedPrintedValue - 1;
 
     os << " :\n";
     os << "\tNumber of instructions " << currentNumInst << "\n";
@@ -194,6 +736,10 @@ public:
   LogicalResult initialize(MLIRContext* context) override { return success(); }
 
   void runOnOperation() override {
+    mergeTriviallyConnectedComputes(getOperation());
+    packWideWeightedVmmBands(getOperation());
+    emitMotifProfile(getOperation());
+
     DCPAnalysisResult& analysisResult = getAnalysis<spatial::DCPAnalysis>().getResult();
     auto& lastComputeOfCpu = analysisResult.isLastComputeOfCpu;
     auto& cpuToLastComputeMap = analysisResult.cpuToLastComputeMap;
@@ -312,22 +858,34 @@ private:
     IRRewriter rewriter(&getContext());
     IRMapping mapper;
 
+    DenseMap<Value, SmallVector<size_t, 4>> toWeightIndices;
+    for (auto [weightIndex, weight] : llvm::enumerate(toCompute.getWeights()))
+      toWeightIndices[weight].push_back(weightIndex);
+    DenseMap<Value, size_t> usedFromWeightOccurrences;
+    SmallVector<size_t> fromWeightToNewIndex;
+    fromWeightToNewIndex.reserve(fromCompute.getWeights().size());
+
     auto weightMutableIter = toCompute.getWeightsMutable();
     for (auto weight : fromCompute.getWeights()) {
-      auto founded = llvm::find(toCompute.getWeights(), weight);
-      if (founded == toCompute.getWeights().end()) {
+      size_t occurrence = usedFromWeightOccurrences[weight]++;
+      auto& matchingIndices = toWeightIndices[weight];
+      if (occurrence >= matchingIndices.size()) {
         size_t sizeW = toCompute.getWeights().size();
         size_t sizeI = toCompute.getInputs().size();
         weightMutableIter.append(weight);
         auto last = weightMutableIter.end();
         last = std::prev(last, 1);
         mapper.map(weight, last->get());
+        matchingIndices.push_back(sizeW);
+        fromWeightToNewIndex.push_back(sizeW);
         assert(sizeW + 1 == toCompute.getWeights().size());
         assert(sizeI == toCompute.getInputs().size());
         assert(sizeW + sizeI + 1 == toCompute.getOperands().size());
       }
       else {
-        mapper.map(weight, *founded);
+        size_t newIndex = matchingIndices[occurrence];
+        mapper.map(weight, *std::next(toCompute.getWeights().begin(), newIndex));
+        fromWeightToNewIndex.push_back(newIndex);
       }
     }
 
@@ -379,9 +937,8 @@ private:
     ComputeValueResults computeValueResults;
     auto remapWeightIndex = [&](auto weightedOp) {
       auto oldIndex = weightedOp.getWeightIndex();
-      auto newWeight = mapper.lookup(*std::next(fromCompute.getWeights().begin(), oldIndex));
-      auto newIndex = std::distance(toCompute.getWeights().begin(), llvm::find(toCompute.getWeights(), newWeight));
-      weightedOp.setWeightIndex(newIndex);
+      assert(static_cast<size_t>(oldIndex) < fromWeightToNewIndex.size() && "weight index out of range");
+      weightedOp.setWeightIndex(fromWeightToNewIndex[oldIndex]);
     };
     for (auto& op : fromCompute.getOps()) {
       if (auto yield = dyn_cast<spatial::SpatYieldOp>(&op)) {
diff --git a/test/PIM/DCPTest.cpp b/test/PIM/DCPTest.cpp
index b813fb1..5687ac0 100644
--- a/test/PIM/DCPTest.cpp
+++ b/test/PIM/DCPTest.cpp
@@ -457,6 +457,10 @@ int testDCPGraphDiamondDependencies() {
   return 0;
 }
 
+// crossbarSize=4, crossbarCount=2 => capacity = 4*4*2 = 32.
+// Each task with crossbarUsage=1 needs footprint = 4*4 = 16, so at most 1 task
+// can fit per CPU (16+16 = 32 >= capacity). The scheduler must open a fresh CPU
+// for each task; all three end up on separate CPUs with their base weight.
 int testDCPGraphCrossbarExhaustion() {
   std::cout << "testDCPGraphCrossbarExhaustion:" << std::endl;
   configureDcpDotOutput();
@@ -474,36 +478,35 @@ int testDCPGraphCrossbarExhaustion() {
   const std::vector<Weight> nodeWeights = {10, 10, 10};
   const std::vector<CrossbarUsage> nodeCrossbarUsage = {1, 1, 1};
   GraphDCP graph(nodeWeights, {}, nodeCrossbarUsage);
-  graph.setMaxCpuCount(1);
+  graph.setMaxCpuCount(3);
   graph.runDcp();
 
-  if (graph.cpuCount() != 1) {
+  if (graph.cpuCount() != 3) {
     restoreCrossbarOptions();
-    std::cerr << "Expected exactly 1 CPU with maxCpuCount=1, got " << graph.cpuCount() << "\n";
+    std::cerr << "Expected 3 CPUs (one per task due to crossbar limit), got " << graph.cpuCount() << "\n";
     dumpDcpFailureArtifacts();
     return 1;
   }
 
-  auto scheduledTasks = graph.getScheduledTasks(0);
-  if (scheduledTasks.size() != 3) {
-    restoreCrossbarOptions();
-    std::cerr << "Expected all three tasks to be scheduled on CPU 0\n";
-    printCpuSchedule(graph, 0);
-    dumpDcpFailureArtifacts();
-    return 1;
-  }
-
-  if (scheduledTasks[0].weight != 10 || scheduledTasks[1].weight != std::numeric_limits<Weight>::max()
-      || scheduledTasks[2].weight != std::numeric_limits<Weight>::max()) {
-    restoreCrossbarOptions();
-    std::cerr << "Unexpected effective weights under crossbar exhaustion\n";
-    printCpuSchedule(graph, 0);
-    dumpDcpFailureArtifacts();
-    return 1;
+  int failures = 0;
+  for (CPU c = 0; c < 3; c++) {
+    auto scheduledTasks = graph.getScheduledTasks(c);
+    if (scheduledTasks.size() != 1) {
+      std::cerr << "Expected exactly 1 task on CPU " << c << ", got " << scheduledTasks.size() << "\n";
+      printCpuSchedule(graph, c);
+      failures++;
+      continue;
+    }
+    if (scheduledTasks[0].weight != 10) {
+      std::cerr << "Expected weight=10 on CPU " << c << ", got " << scheduledTasks[0].weight << "\n";
+      printCpuSchedule(graph, c);
+      failures++;
+    }
   }
 
   restoreCrossbarOptions();
-  return 0;
+  if (failures) dumpDcpFailureArtifacts();
+  return failures;
 }
 
 } // namespace