better scheduling

.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md

@@ -30,8 +30,9 @@ they must not infer it from a physical core number after placement.
 
 ## Ownership
 
-Physical PEFT supplies the initial placement. Stage partitioning is the first
-phase of pipeline scheduling and owns stage adjacency. It must construct a valid
+Latency PEFT and pipeline scheduling independently consume the shared compute
+graph and target cost model. Stage partitioning is the first phase of pipeline
+scheduling and owns stage adjacency. It must construct a valid
 operation-level partition before physical-core packing. Operations split for
 physical capacity retain one shared stage identity. Repacking may move work
 only within its assigned stage. Deferred-transfer planning and

src/PIM/Common/IR/ShapeUtils.cpp

@@ -15,6 +15,25 @@ llvm::SmallVector<int64_t> computeRowMajorStrides(llvm::ArrayRef<int64_t> shape)
   return strides;
 }
 
+TransposeRun getTransposeRun(llvm::ArrayRef<int64_t> sourceShape,
+                            llvm::ArrayRef<int64_t> permutation) {
+  auto sourceStrides = computeRowMajorStrides(sourceShape);
+  TransposeRun run{sourceShape.size(), 1, 1};
+  for (size_t dim = permutation.size(); dim-- > 0;) {
+    int64_t sourceDim = permutation[dim];
+    if (sourceShape[sourceDim] == 1)
+      continue;
+    int64_t stride = sourceStrides[sourceDim];
+    if (run.start != sourceShape.size() && stride != run.elements * run.sourceStride)
+      break;
+    run.start = dim;
+    if (run.elements == 1)
+      run.sourceStride = stride;
+    run.elements *= sourceShape[sourceDim];
+  }
+  return run;
+}
+
 llvm::SmallVector<int64_t>
 delinearizeIndex(int64_t linearIndex, llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> strides) {
   llvm::SmallVector<int64_t> indices(shape.size(), 0);

src/PIM/Common/IR/ShapeUtils.hpp

@@ -21,6 +21,15 @@ using CoreId = size_t;
 
 llvm::SmallVector<int64_t> computeRowMajorStrides(llvm::ArrayRef<int64_t> shape);
 
+struct TransposeRun {
+  size_t start;
+  int64_t elements;
+  int64_t sourceStride;
+};
+
+TransposeRun getTransposeRun(llvm::ArrayRef<int64_t> sourceShape,
+                            llvm::ArrayRef<int64_t> permutation);
+
 llvm::SmallVector<int64_t>
 delinearizeIndex(int64_t linearIndex, llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> strides);
 

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -78,6 +78,8 @@ spatial::SchedulingTarget getDefaultPimSchedulingTarget() {
   target.processorCount = static_cast<size_t>(coresCount.getValue());
   target.residentWeightCapacity = crossbarCountInCore.getValue();
   target.matrixRows = crossbarSize.getValue();
+  // The oracle charges eight input reads for each matrix operation.
+  target.matrixInputReadMultiplier = 8;
   target.matrixColumns = crossbarSize.getValue();
   target.synchronizationRegisterCount = kPimEventRegisterCount;
 

src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/PimBufferizationPass.cpp

@@ -275,9 +275,12 @@ static void forwardSingleConsumerReceiveCopies(func::FuncOp funcOp) {
       ? receive.getOutputBuffer().getDefiningOp<memref::AllocOp>()
       : memref::AllocOp();
     auto sourceOffset = getConstantIntValue(copy.getSourceOffset());
+    auto receiveOffset = receive
+      ? getConstantIntValue(receive.getOutputOffset()) : std::nullopt;
     if (!receive || !outputAlloc || !receive.getOutput().hasOneUse()
         || !outputAlloc.getResult().hasOneUse()
         || !sourceOffset || *sourceOffset != 0
+        || !receiveOffset || *receiveOffset != 0
         || copy.getSize() != receive.getSize()
         || receive->getBlock() != copy->getBlock())
       continue;
@@ -285,7 +288,9 @@ static void forwardSingleConsumerReceiveCopies(func::FuncOp funcOp) {
     bool canSinkReceive = true;
     for (Operation* between = receive->getNextNode(); between != copy;
          between = between->getNextNode()) {
-      if (!isMemoryEffectFree(between)) {
+      // The receive writes a private allocation; intervening loads cannot
+      // observe it, and the destination is not written until the copy.
+      if (!isMemoryEffectFree(between) && !isa<memref::LoadOp>(between)) {
         canSinkReceive = false;
         break;
       }
@@ -594,7 +599,6 @@ static LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) {
 }
 
 static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp) {
-  forwardSingleConsumerReceiveCopies(funcOp);
   forwardSingleConsumerContiguousInputCopies(funcOp);
   forwardSingleConsumerPimOutputCopies(funcOp);
 
@@ -676,6 +680,7 @@ static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp)
     moduleOp.emitError("failed to normalize Pim copy contiguity during bufferization");
     return failure();
   }
+  forwardSingleConsumerReceiveCopies(funcOp);
   annotateWeightsMemrefs(moduleOp, funcOp);
   dumpModule(moduleOp, "pim1_buff");
   return success();

src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/InstructionSelectionPass.cpp

@@ -84,20 +84,7 @@ struct LowerTransposePattern final : OpRewritePattern<pim::PimTransposeOp> {
     for (int64_t sourceDim : permutation)
       destinationShape.push_back(sourceShape[sourceDim]);
 
-    size_t runStart = rank;
-    int64_t runLength = 1;
-    int64_t sourceStride = 1;
-    for (size_t destinationDim = rank; destinationDim-- > 0;) {
-      if (destinationShape[destinationDim] == 1)
-        continue;
-      int64_t candidateStride = sourceStrides[permutation[destinationDim]];
-      if (runStart != rank && candidateStride != runLength * sourceStride)
-        break;
-      runStart = destinationDim;
-      if (runLength == 1)
-        sourceStride = candidateStride;
-      runLength *= destinationShape[destinationDim];
-    }
+    auto [runStart, runLength, sourceStride] = getTransposeRun(sourceShape, permutation);
     assert(runStart != rank && "non-storage-preserving transpose must have a non-unit dimension");
 
     int64_t outerCount = sourceType.getNumElements() / runLength;

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -34,8 +34,11 @@ add_pim_library(SpatialOps
   Passes/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
   Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
   Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.cpp
   Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
   Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/PipelineWorkloadPreparation.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.cpp
   Passes/Transforms/TrivialGraphComputeMergePass.cpp
 
   EXCLUDE_FROM_OM_LIBS

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp

@@ -417,7 +417,8 @@ LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, Bloc
                            && producer->instance.laneCount
                                 == static_cast<uint32_t>(batch.getLaneCount()));
   if (auto available = availableValues.find(key);
-      hasCompleteValue && available != availableValues.end()) {
+      hasCompleteValue && available != availableValues.end()
+      && (!batch || available->second.hasGraphPhysicalOrder)) {
     plan.availableValue = available->second.payload;
     return success();
   }

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp

@@ -387,8 +387,10 @@ FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp
       funcOp.emitOpError("global communication scheduler made no progress before IR mutation");
     unsigned reported = 0;
     for (const TransferGroup& group : groups) {
-      if (group.scheduled || reported++ == 8)
+      if (group.scheduled)
         continue;
+      if (reported++ == 8)
+        break;
       diagnostic << "; blocked parent " << group.exchange->exchangeId;
       auto dependency = llvm::find_if(group.dependencies, [&](StreamThreshold item) {
         return streams[item.stream].completedStep < item.completedStep;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp

@@ -1,5 +1,6 @@
 #include "ScheduledComputeMaterialization.hpp"
 #include "DeferredCommunicationPlanning.hpp"
+#include "DeferredProjectionAnalysis.hpp"
 
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -129,6 +130,36 @@ static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mappe
     yieldedValues.push_back(mapper.lookup(output));
 }
 
+} // namespace
+
+Value insertLocalFragment(OpBuilder &builder, Location loc, Value source,
+                                 Value destination, ArrayRef<OpFoldResult> offsets,
+                                 ArrayRef<OpFoldResult> sizes,
+                                 ArrayRef<OpFoldResult> strides) {
+  auto type = cast<RankedTensorType>(destination.getType());
+  auto sourceType = cast<RankedTensorType>(source.getType());
+  bool complete = type.hasStaticShape() && sourceType.hasStaticShape()
+    && type.getRank() == sourceType.getRank() + 1
+    && type.getDimSize(0) == 1
+    && type.getElementType() == sourceType.getElementType()
+    && type.getEncoding() == sourceType.getEncoding()
+    && llvm::equal(type.getShape().drop_front(), sourceType.getShape())
+    && offsets.size() == static_cast<size_t>(type.getRank())
+    && sizes.size() == offsets.size() && strides.size() == offsets.size();
+  if (complete)
+    for (auto [dimension, offset, size, stride] :
+         llvm::zip_equal(type.getShape(), offsets, sizes, strides))
+      complete &= getConstantIntValue(offset) == 0
+        && getConstantIntValue(size) == dimension
+        && getConstantIntValue(stride) == 1;
+  if (complete)
+    return *addLeadingUnitTensorDimension(builder, loc, source);
+  return tensor::InsertSliceOp::create(
+    builder, loc, source, destination, offsets, sizes, strides);
+}
+
+namespace {
+
 static LogicalResult materializeResultfulBatchRun(
     PatternRewriter &rewriter, SpatComputeBatch batch,
     const ScheduledInstanceRun &run, ValueRange scheduledWeights,
@@ -232,14 +263,11 @@ static LogicalResult materializeResultfulBatchRun(
         size_t resultIndex = outputIndexByArg.lookup(oldDest);
         SmallVector<OpFoldResult> offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper);
         offsets.front() = localLane;
-        current[resultIndex] = tensor::InsertSliceOp::create(builder,
-                                                             insert.getLoc(),
-                                                             mapper.lookup(insert.getSource()),
-                                                             current[resultIndex],
-                                                             offsets,
-                                                             remapMixedOffsets(insert.getMixedSizes(), mapper),
-                                                             remapMixedOffsets(insert.getMixedStrides(), mapper))
-                                 .getResult();
+        current[resultIndex] = insertLocalFragment(
+          builder, insert.getLoc(), mapper.lookup(insert.getSource()),
+          current[resultIndex], offsets,
+          remapMixedOffsets(insert.getMixedSizes(), mapper),
+          remapMixedOffsets(insert.getMixedStrides(), mapper));
       }
       llvm::append_range(yielded, current);
       return success();
@@ -264,6 +292,7 @@ static LogicalResult materializeSingleCpuPeftClass(
   Block *block = createScheduledComputeBlock(
     rewriter, scheduled, instances.front().op->getLoc());
   DenseMap<ProducerValueKey, MaterializedProducerRef> availableValues;
+  GraphBatchPublicationCache publicationCache;
   DenseMap<ProducerValueKey, unsigned> publicationIndices;
   for (const ScheduledPublication &publication : record.publications)
     publicationIndices[publication.producer] = publication.scheduledResultIndex;
@@ -326,13 +355,23 @@ static LogicalResult materializeSingleCpuPeftClass(
       int64_t runLaneStart = instance.laneStart;
       const ComputeInstance &last = run.instances.back();
       int64_t runLaneCount = last.laneStart + last.laneCount - runLaneStart;
+      bool hasGraphPhysicalOrder = false;
+      if (auto batch = dyn_cast<SpatComputeBatch>(instance.op);
+          batch && runLaneStart == 0 && runLaneCount == batch.getLaneCount()) {
+        auto publication = getGraphBatchPublicationMap(batch, resultIndex, publicationCache);
+        if (failed(publication))
+          return failure();
+        hasGraphPhysicalOrder = llvm::all_of(
+          llvm::enumerate((*publication)->graphLaneToPhysicalSlot),
+          [](auto entry) { return entry.index() == static_cast<size_t>(entry.value()); });
+      }
       availableValues[{{instance.op, static_cast<uint32_t>(runLaneStart),
                         static_cast<uint32_t>(runLaneCount)}, resultIndex}] =
-        {value, 0, runLaneCount};
+        {value, 0, runLaneCount, hasGraphPhysicalOrder};
       for (const ComputeInstance &member : run.instances) {
         ProducerValueKey key {member, resultIndex};
         int64_t payloadLaneStart = member.laneStart - runLaneStart;
-        availableValues[key] = {value, payloadLaneStart, member.laneCount};
+        availableValues[key] = {value, payloadLaneStart, member.laneCount, hasGraphPhysicalOrder};
         Value published;
         if (auto it = publicationIndices.find(key);
             it != publicationIndices.end()) {
@@ -654,14 +693,11 @@ static LogicalResult materializeMultiCpuPeftClass(
             size_t resultIndex = outputIndexByArg.lookup(oldDest);
             SmallVector<OpFoldResult> offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper);
             offsets.front() = innerLane;
-            current[resultIndex] = tensor::InsertSliceOp::create(builder,
-                                                                 insert.getLoc(),
-                                                                 mapper.lookup(insert.getSource()),
-                                                                 current[resultIndex],
-                                                                 offsets,
-                                                                 remapMixedOffsets(insert.getMixedSizes(), mapper),
-                                                                 remapMixedOffsets(insert.getMixedStrides(), mapper))
-                                     .getResult();
+            current[resultIndex] = insertLocalFragment(
+              builder, insert.getLoc(), mapper.lookup(insert.getSource()),
+              current[resultIndex], offsets,
+              remapMixedOffsets(insert.getMixedSizes(), mapper),
+              remapMixedOffsets(insert.getMixedStrides(), mapper));
           }
           llvm::append_range(yielded, current);
           return success();
@@ -706,11 +742,9 @@ FailureOr<ScheduledComputeMaterializationResult>
 materializeScheduledCompute(func::FuncOp funcOp,
                             const MergeScheduleResult &schedule,
                             PatternRewriter &rewriter) {
-  DenseMap<Operation *, int64_t> graphIds;
   int64_t nextGraphId = 0;
   for (Operation &op : funcOp.getOps())
     if (isa<SpatGraphCompute, SpatGraphComputeBatch>(op)) {
-      graphIds[&op] = nextGraphId;
       op.setAttr("scheduled.graph_id", rewriter.getI64IntegerAttr(nextGraphId++));
     }
 
@@ -726,13 +760,12 @@ materializeScheduledCompute(func::FuncOp funcOp,
   }
   for (auto &entry : peftClassPlans) {
     PeftClassPlan &peftClassPlan = entry.second;
-    llvm::sort(peftClassPlan.cpus);
+    // Keep class lanes in graph encounter order so physical placement does not
+    // turn regular source lanes into lookup tables. Core execution uses slots.
     for (size_t cpu : peftClassPlan.cpus)
       llvm::sort(peftClassPlan.instancesByCpu[cpu], [&](const ComputeInstance &lhs, const ComputeInstance &rhs) {
-        return std::tie(graphIds.find(lhs.op)->second,
-                        schedule.computeToCpuSlotMap.find(lhs)->second) <
-               std::tie(graphIds.find(rhs.op)->second,
-                        schedule.computeToCpuSlotMap.find(rhs)->second);
+        return schedule.computeToCpuSlotMap.lookup(lhs)
+             < schedule.computeToCpuSlotMap.lookup(rhs);
       });
     if (failed(verifyPeftClassPlan(funcOp.getOperation(), peftClassPlan, schedule)))
       return failure();

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp

@@ -17,6 +17,11 @@ struct ScheduledComputeMaterializationResult {
   std::vector<size_t> processorStages;
 };
 
+Value insertLocalFragment(OpBuilder &builder, Location loc, Value source,
+                          Value destination, ArrayRef<OpFoldResult> offsets,
+                          ArrayRef<OpFoldResult> sizes,
+                          ArrayRef<OpFoldResult> strides);
+
 FailureOr<BatchFragmentSpec>
 getBatchFragmentSpec(SpatComputeBatch batch,
                      unsigned resultIndex,

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp

@@ -94,6 +94,7 @@ struct MaterializedProducerRef {
   Value payload;
   int64_t payloadLaneStart = 0;
   int64_t payloadLaneCount = 1;
+  bool hasGraphPhysicalOrder = false;
 };
 
 struct MaterializedStepValue {

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp

@@ -108,9 +108,15 @@ LogicalResult verifyPeftClassPlan(Operation *diagnosticAnchor,
     if (it == peftClassPlan.instancesByCpu.end())
       return diagnosticAnchor->emitOpError("PEFT materialization class is missing a per-CPU schedule");
     schedules.push_back(&it->second);
-    for (const ComputeInstance &instance : it->second)
-      if (!schedule.computeToCpuSlotMap.count(instance))
-        return diagnosticAnchor->emitOpError("PEFT materialization class references a compute instance without a scheduler position");
+    for (auto [slot, instance] : llvm::enumerate(it->second)) {
+      auto position = schedule.computeToCpuSlotMap.find(instance);
+      auto processor = schedule.computeToCpuMap.find(instance);
+      if (position == schedule.computeToCpuSlotMap.end()
+          || processor == schedule.computeToCpuMap.end()
+          || processor->second != cpu || position->second != slot)
+        return diagnosticAnchor->emitOpError(
+          "scheduled materialization must preserve each core's scheduler slot order");
+    }
   }
 
   if (peftClassPlan.cpus.size() == 1)

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp

@@ -112,6 +112,18 @@ void dumpScheduledComputeReport(ModuleOp moduleOp, func::FuncOp funcOp, const Me
      << "  largest run: " << largestRun << "\n"
      << "  instances compacted: " << instancesCompacted << "\n"
      << "  compatible runs rejected: 0\n";
+  const ScheduleCostSummary &costs = schedule.costs;
+  os << "  estimated total serial work (ns): " << costs.totalWorkNs << "\n"
+     << "  estimated maximum core work (ns): " << costs.maximumCoreWorkNs << "\n"
+     << "  estimated dependency critical path (ns): " << costs.criticalPathNs << "\n"
+     << "  summed edge communication time (ns): " << costs.communicationNs << "\n"
+     << "  estimated single-inference completion (ns): " << costs.completionNs << "\n";
+  for (auto [stage, span] : llvm::enumerate(costs.stageSpanNs))
+    os << "  estimated stage " << stage << " span (ns, excludes boundary synchronization/contention): "
+       << span << "\n";
+  for (auto [stage, demand] : llvm::enumerate(costs.stageSerialDemandNs))
+    os << "  estimated stage " << stage << " maximum core compute/incident-transfer demand (ns): "
+       << demand << "\n";
   llvm::MapVector<size_t, SmallVector<size_t>> stageCores;
   for (auto [core, stage] : llvm::enumerate(schedule.processorStages))
     stageCores[stage].push_back(core);

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp

@@ -19,6 +19,7 @@ namespace {
 static bool hasValidTarget(const SchedulingTarget& target) {
   return target.processorCount != 0 && target.residentWeightCapacity != 0
       && target.transferWidthBytes != 0
+      && target.matrixInputReadMultiplier != 0
       && target.interProcessorLatencyNs.size() == target.processorCount * target.processorCount
       && (target.processorCount == 1 || target.averageInterProcessorLatencyNs != 0);
 }
@@ -73,7 +74,12 @@ struct ScheduleAndRealizeSpatialPass final
     ComputeGraph scheduledGraph;
     MergeScheduleResult schedule;
     for (;;) {
-      scheduledGraph = buildComputeGraph(entryFunc, target);
+      auto graph = buildComputeGraph(entryFunc, target);
+      if (failed(graph)) {
+        signalPassFailure();
+        return;
+      }
+      scheduledGraph = std::move(*graph);
       std::string pipelineError;
       if (pipelineStages > 1) {
         FailureOr<PipelineWorkloadPreparation> preparation =
@@ -87,10 +93,18 @@ struct ScheduleAndRealizeSpatialPass final
         if (*preparation == PipelineWorkloadPreparation::Changed)
           continue;
       }
-      MergeSchedulingAnalysis analysis(
-        scheduledGraph, target, entryFunc.getContext());
-      schedule = std::move(analysis.getResult());
-      if (succeeded(applyPipelineScheduling(
+      if (pipelineStages == 1) {
+        MergeSchedulingAnalysis analysis(
+          scheduledGraph, target, entryFunc.getContext(), pipelineError);
+        if (failed(analysis.getResult())) {
+          moduleOp.emitError() << pipelineError;
+          signalPassFailure();
+          return;
+        }
+        schedule = std::move(*analysis.getResult());
+        break;
+      }
+      if (succeeded(runPipelineScheduler(
             scheduledGraph, schedule, pipelineStages, target, pipelineError)))
         break;
       moduleOp.emitError() << pipelineError;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinAttributes.h"
@@ -7,8 +8,12 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Unit.h"
 #include "mlir/IR/Value.h"
+#include "mlir/Transforms/RegionUtils.h"
 
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallBitVector.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/Support/Casting.h"
 
 #include <algorithm>
@@ -21,6 +26,7 @@
 
 #include "ComputeGraph.hpp"
 #include "ComputeInstanceUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp"
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
@@ -43,18 +49,11 @@ TransferCost addTransferCosts(const TransferCost& lhs,
           checkedAdd(lhs.networkFlits, rhs.networkFlits)};
 }
 
-TransferCost scaleTransferCost(const TransferCost& cost,
-                               Cost numerator,
-                               Cost denominator = 1) {
-  assert(denominator > 0 && "transfer cost denominator must be positive");
-  return {checkedMultiply(cost.fixed, numerator) / denominator,
-          checkedMultiply(cost.networkFlits, numerator) / denominator};
-}
-
 struct SchedulerCostModel {
   static constexpr Cost kFallbackOperationCost = 1;
 
   const SchedulingTarget& target;
+  ComputeCostKind kind = ComputeCostKind::Execution;
 
   static Cost ceilDiv(Cost numerator, Cost denominator) {
     assert(denominator > 0 && "denominator must be positive");
@@ -110,7 +109,44 @@ struct SchedulerCostModel {
   }
 
   Cost getTensorMoveCost(Type type) const {
-    return getVectorReadWriteCost(getByteSize(type), 0);
+    Cost bytes = getByteSize(type, getComputeBitwidth(type));
+    return getVectorReadWriteCost(bytes, bytes);
+  }
+
+  Cost getTransposeCost(linalg::TransposeOp transpose) const {
+    if (kind == ComputeCostKind::VolumePriority)
+      return getTensorMoveCost(transpose.getInput().getType());
+    auto type = cast<ShapedType>(transpose.getInput().getType());
+    if (!type.hasStaticShape())
+      return getTensorMoveCost(type);
+    if (type.getNumElements() == 0)
+      return 0;
+    auto run = getTransposeRun(type.getShape(), transpose.getPermutation());
+    Cost bytes = ceilDiv(checkedMultiply(static_cast<Cost>(run.elements), getComputeBitwidth(type)), Cost{8});
+    Cost cost = getVectorReadWriteCost(bytes, bytes);
+    if (run.sourceStride != 1)
+      cost = checkedAdd(cost, getVectorComputeCost(run.elements));
+    return checkedMultiply(static_cast<Cost>(type.getNumElements()) / run.elements, cost);
+  }
+
+  Cost getConcatCost(ValueRange inputs, int64_t axis) const {
+    Cost cost = 0;
+    for (Value input : inputs) {
+      auto type = dyn_cast<ShapedType>(input.getType());
+      if (!type || !type.hasStaticShape())
+        return kFallbackOperationCost;
+      Cost blocks = 1;
+      Cost elements = 1;
+      for (auto [dim, size] : llvm::enumerate(type.getShape())) {
+        if (static_cast<int64_t>(dim) < axis)
+          blocks = checkedMultiply(blocks, static_cast<Cost>(size));
+        else
+          elements = checkedMultiply(elements, static_cast<Cost>(size));
+      }
+      Cost bytes = ceilDiv(checkedMultiply(elements, getComputeBitwidth(type)), Cost{8});
+      cost = checkedAdd(cost, checkedMultiply(blocks, getVectorReadWriteCost(bytes, bytes)));
+    }
+    return cost;
   }
 
   TransferCost getTransferCostFromBytes(Cost bytes) const {
@@ -176,7 +212,7 @@ struct SchedulerCostModel {
     Cost inputBitwidth = getComputeBitwidth(inputType);
     Cost inputBytes = checkedMultiply(target.matrixRows,
                                       ceilDiv(inputBitwidth, static_cast<Cost>(8)));
-    inputBytes = checkedMultiply(inputBytes, static_cast<Cost>(8));
+    inputBytes = checkedMultiply(inputBytes, target.matrixInputReadMultiplier);
     Cost outputBytes = getByteSize(outputType, getComputeBitwidth(outputType));
     return checkedAdd(getVectorReadWriteCost(inputBytes, outputBytes), getMatrixComputeLatency(inputBitwidth));
   }
@@ -250,6 +286,10 @@ Cost getOperationCost(Operation& op, const SchedulerCostModel& costModel) {
   }
   if (auto insert = dyn_cast<tensor::InsertSliceOp>(&op))
     return costModel.getTensorMoveCost(insert.getSource().getType());
+  if (auto transpose = dyn_cast<linalg::TransposeOp>(&op))
+    return costModel.getTransposeCost(transpose);
+  if (auto concat = dyn_cast<SpatConcatOp>(&op))
+    return costModel.getConcatCost(concat.getInputs(), concat.getAxis());
 
   Cost nestedCost = 0;
   for (Region& region : op.getRegions())
@@ -275,6 +315,30 @@ Cost getComputeBodyCost(Region& body, const SchedulerCostModel& costModel) {
   return getRegionCost(body, costModel);
 }
 
+Cost getBlueprintAssemblyCost(const ComputeInstance& instance,
+                               const SchedulerCostModel& costModel) {
+  Cost cost = 0;
+  for (Value input : getComputeInstanceInputs(instance)) {
+    auto blueprint = input.getDefiningOp<SpatBlueprintOp>();
+    if (!blueprint || !isFragmentAssembly(blueprint.getMode()))
+      continue;
+    auto type = cast<RankedTensorType>(input.getType());
+    int64_t rank = type.getRank();
+    ArrayRef<int64_t> sizes = blueprint.getFragmentSizes();
+    if (sizes.empty())
+      continue;
+    for (size_t fragment = 0; fragment < sizes.size() / rank; ++fragment) {
+      Cost elements = 1;
+      for (int64_t size : sizes.slice(fragment * rank, rank))
+        elements = checkedMultiply(elements, static_cast<Cost>(size));
+      Cost bytes = SchedulerCostModel::ceilDiv(
+        checkedMultiply(elements, costModel.getComputeBitwidth(type)), Cost{8});
+      cost = checkedAdd(cost, costModel.getVectorReadWriteCost(bytes, bytes));
+    }
+  }
+  return cost;
+}
+
 uint64_t countOperationInstances(Operation& op) {
   if (auto loop = dyn_cast<scf::ForOp>(&op)) {
     std::optional<uint64_t> tripCount = getStaticTripCount(loop);
@@ -320,12 +384,6 @@ static FailureOr<int64_t> evaluateIndexLike(OpFoldResult value,
                                             std::optional<uint32_t> lane,
                                             Value laneArg);
 
-bool isProjectedBatchOffset(OpFoldResult offset, Value laneArg) {
-  DenseMap<Value, int64_t> bindings;
-  return succeeded(evaluateIndexLike(offset, bindings, /*lane=*/0, laneArg))
-      && succeeded(evaluateIndexLike(offset, bindings, /*lane=*/1, laneArg));
-}
-
 std::optional<uint32_t> getConstantExtractLane(tensor::ExtractSliceOp extract) {
   if (extract.getMixedOffsets().empty())
     return std::nullopt;
@@ -348,65 +406,41 @@ std::optional<uint32_t> getConstantExtractLane(tensor::ExtractSliceOp extract) {
   return std::nullopt;
 }
 
-std::optional<TransferCost> getBatchProjectedInputTransferCost(
-    SpatComputeBatch batch, Value input,
-    const SchedulerCostModel& costModel) {
-  auto inputIt = llvm::find(batch.getInputs(), input);
-  if (inputIt == batch.getInputs().end())
-    return std::nullopt;
-
-  size_t inputIndex = std::distance(batch.getInputs().begin(), inputIt);
-  std::optional<BlockArgument> inputArg = batch.getInputArgument(inputIndex);
-  std::optional<BlockArgument> laneArg = batch.getLaneArgument();
-  if (!inputArg || !laneArg)
-    return std::nullopt;
-
-  TransferCost projectedCost;
-  for (Operation* user : inputArg->getUsers()) {
-    auto extract = dyn_cast<tensor::ExtractSliceOp>(user);
-    if (!extract || extract.getSource() != *inputArg)
-      return std::nullopt;
-    if (extract.getMixedOffsets().empty() || !isProjectedBatchOffset(extract.getMixedOffsets().front(), *laneArg))
-      return std::nullopt;
-
-    auto resultType = dyn_cast<ShapedType>(extract.getResult().getType());
-    if (!resultType || !resultType.hasStaticShape())
-      return std::nullopt;
-    projectedCost = addTransferCosts(
-      projectedCost,
-      costModel.getTransferCostFromBytes(
-        costModel.getByteSize(resultType, costModel.getComputeBitwidth(resultType))));
-  }
-
-  if (projectedCost.fixed == 0 && projectedCost.networkFlits == 0)
-    return std::nullopt;
-  return projectedCost;
-}
-
-static std::optional<SmallVector<ProducerValueRef, 4>>
-collectProjectedProducerValueRefs(SpatComputeBatch producer,
-                                  Value input,
-                                  const ComputeInstance& consumerInstance,
-                                  size_t processorCount) {
+static FailureOr<bool>
+collectProjectedEdges(SpatComputeBatch producer, Value input,
+                      const ComputeInstance& consumerInstance,
+                      size_t processorCount, size_t targetIndex,
+                      const ComputeGraph& graph,
+                      const SchedulerCostModel& costModel,
+                      GraphBatchPublicationCache& publicationCache,
+                      SmallVectorImpl<ComputeGraphEdge>& rawEdges) {
   auto consumer = dyn_cast<SpatComputeBatch>(consumerInstance.op);
   if (!consumer)
-    return std::nullopt;
+    return false;
   auto inputIt = llvm::find(consumer.getInputs(), input);
   if (inputIt == consumer.getInputs().end())
-    return std::nullopt;
+    return false;
   size_t inputIndex = std::distance(consumer.getInputs().begin(), inputIt);
   std::optional<BlockArgument> inputArg = consumer.getInputArgument(inputIndex);
   std::optional<BlockArgument> laneArg = consumer.getLaneArgument();
   if (!inputArg || !laneArg)
-    return std::nullopt;
+    return false;
+  auto publication = getGraphBatchPublicationMap(
+    producer, cast<OpResult>(input).getResultNumber(), publicationCache);
+  if (failed(publication))
+    return failure();
 
-  SmallVector<ProducerValueRef, 4> producers;
+  SmallVector<ComputeGraphEdge, 4> edges;
+  llvm::SmallBitVector selectedSlots(producer.getLaneCount());
+  TransferCost fragmentCost = costModel.getTransferCostFromBytes(
+    costModel.getByteSize((*publication)->publicationFragmentType,
+      costModel.getComputeBitwidth((*publication)->publicationFragmentType)));
   DenseMap<Value, int64_t> bindings;
   for (Operation* user : inputArg->getUsers()) {
     auto extract = dyn_cast<tensor::ExtractSliceOp>(user);
     if (!extract || extract.getSource() != *inputArg || extract.getMixedOffsets().empty()
         || extract.getMixedSizes().empty() || extract.getMixedStrides().empty())
-      return std::nullopt;
+      return false;
     for (uint32_t lane = consumerInstance.laneStart;
          lane < consumerInstance.laneStart + consumerInstance.laneCount;
          ++lane) {
@@ -417,142 +451,219 @@ collectProjectedProducerValueRefs(SpatComputeBatch producer,
       FailureOr<int64_t> stride =
         evaluateIndexLike(extract.getMixedStrides().front(), bindings, lane, *laneArg);
       if (failed(offset) || failed(size) || failed(stride) || *offset < 0 || *size <= 0 || *stride <= 0)
-        return std::nullopt;
+        return false;
       for (int64_t index = 0; index < *size; ++index) {
-        int64_t producerLane = *offset + index * *stride;
-        if (producerLane < 0 || producerLane >= producer.getLaneCount())
-          return std::nullopt;
+        int64_t slot = *offset + index * *stride;
+        if (slot < 0 || slot >= producer.getLaneCount())
+          return false;
+        if (selectedSlots.test(slot))
+          continue;
+        selectedSlots.set(slot);
+        int64_t producerLane = (*publication)->physicalSlotToGraphLane[slot];
         ComputeInstance instance = getBatchChunkForLane(
           producer, static_cast<uint32_t>(producerLane), processorCount);
-        if (llvm::none_of(producers, [&](const ProducerValueRef& ref) { return ref.instance == instance; }))
-          producers.push_back({instance, 0});
+        auto source = graph.instanceToIndex.find(instance);
+        if (source == graph.instanceToIndex.end())
+          return producer.emitOpError("projected producer has no compute instance"), failure();
+        edges.push_back({source->second, targetIndex, fragmentCost});
       }
     }
   }
-  return producers;
+  llvm::append_range(rawEdges, edges);
+  return true;
 }
 
-TransferCost getInputTransferCost(const ComputeInstance& consumerInstance,
-                                  Value input,
+std::optional<TransferCost> getProjectedInputTransferCost(
+    Value input, const ComputeInstance& instance,
+    const SchedulerCostModel& costModel) {
+  auto batch = dyn_cast<SpatComputeBatch>(instance.op);
+  if (!batch)
+    return std::nullopt;
+  auto found = llvm::find(batch.getInputs(), input);
+  if (found == batch.getInputs().end())
+    return std::nullopt;
+  auto argument = batch.getInputArgument(found - batch.getInputs().begin());
+  if (!argument)
+    return std::nullopt;
+  TransferCost cost;
+  for (Operation* user : argument->getUsers()) {
+    auto extract = dyn_cast<tensor::ExtractSliceOp>(user);
+    if (!extract || extract.getSource() != *argument)
+      return std::nullopt;
+    Cost count = instance.laneCount;
+    for (Operation* parent = extract->getParentOp(); parent != instance.op;
+         parent = parent->getParentOp()) {
+      auto loop = dyn_cast<scf::ForOp>(parent);
+      auto trips = loop ? getStaticTripCount(loop) : std::nullopt;
+      if (!trips)
+        return std::nullopt;
+      count = checkedMultiply(count, static_cast<Cost>(*trips));
+    }
+    auto type = extract.getResult().getType();
+    TransferCost fragment = costModel.getTransferCostFromBytes(
+      costModel.getByteSize(type, costModel.getComputeBitwidth(type)));
+    cost = addTransferCosts(cost,
+      {checkedMultiply(fragment.fixed, count),
+       checkedMultiply(fragment.networkFlits, count)});
+  }
+  return cost;
+}
+
+TransferCost getInputTransferCost(Value input,
                                   const SchedulerCostModel& costModel) {
   auto inputType = cast<ShapedType>(input.getType());
-  if (auto batch = dyn_cast<SpatComputeBatch>(consumerInstance.op))
-    if (std::optional<TransferCost> projectedCost =
-          getBatchProjectedInputTransferCost(batch, input, costModel))
-      return *projectedCost;
   return costModel.getTransferCostFromBytes(
     costModel.getByteSize(inputType, costModel.getComputeBitwidth(inputType)));
 }
 
-uint32_t getLaneOverlapCount(const ComputeInstance& lhs, const ComputeInstance& rhs) {
-  uint32_t lhsEnd = lhs.laneStart + lhs.laneCount;
-  uint32_t rhsEnd = rhs.laneStart + rhs.laneCount;
-  return std::max(lhs.laneStart, rhs.laneStart) < std::min(lhsEnd, rhsEnd)
-         ? std::min(lhsEnd, rhsEnd) - std::max(lhs.laneStart, rhs.laneStart)
-         : 0;
-}
-
-TransferCost scaleTransferCostByLaneCount(
-    const TransferCost& totalCost, uint32_t totalLaneCount,
-    uint32_t fragmentLaneCount) {
-  assert(totalLaneCount > 0 && "laneCount must be positive");
-  assert(fragmentLaneCount > 0 && "fragmentLaneCount must be positive");
-  if (fragmentLaneCount >= totalLaneCount)
-    return totalCost;
-  return scaleTransferCost(totalCost, static_cast<Cost>(fragmentLaneCount),
-                           static_cast<Cost>(totalLaneCount));
-}
-
-SmallVector<ProducerValueRef, 4> collectProducerValueRefs(
-    Value value, const ComputeInstance& consumerInstance,
-    size_t processorCount) {
+FailureOr<SmallVector<ProducerValueRef, 4>> collectProducerValueRefs(
+    Value value,
+    size_t processorCount, GraphBatchPublicationCache& publicationCache) {
   SmallVector<ProducerValueRef, 4> producers;
-  Operation* op = value.getDefiningOp();
-  if (!op)
-    return producers;
-
-  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
-    Value source = extract.getSource();
-    auto batch = dyn_cast_or_null<SpatComputeBatch>(source.getDefiningOp());
-    if (batch && batch.getNumResults() != 0) {
-      if (std::optional<uint32_t> lane = getConstantExtractLane(extract)) {
-        ComputeInstance instance =
-          getBatchChunkForLane(batch, *lane, processorCount);
-        producers.push_back({instance, 0});
-        return producers;
+  SmallVector<Value> pending{value};
+  llvm::SmallDenseSet<Value, 16> visited;
+  llvm::SmallDenseSet<std::pair<ComputeInstance, size_t>, 8> indexed;
+  auto add = [&](ComputeInstance instance, size_t result) {
+    if (indexed.insert({instance, result}).second)
+      producers.push_back({instance, result});
+  };
+  while (!pending.empty()) {
+    Value current = pending.pop_back_val();
+    Operation* op = current.getDefiningOp();
+    if (!op || !visited.insert(current).second)
+      continue;
+    if (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
+      Value source = extract.getSource();
+      auto batch = source.getDefiningOp<SpatComputeBatch>();
+      if (batch && batch.getNumResults() != 0) {
+        if (auto slot = getConstantExtractLane(extract);
+            slot && getConstantIntValue(extract.getMixedSizes().front()) == 1) {
+          unsigned result = cast<OpResult>(source).getResultNumber();
+          auto publication = getGraphBatchPublicationMap(batch, result, publicationCache);
+          if (failed(publication))
+            return failure();
+          if (*slot >= (*publication)->physicalSlotToGraphLane.size())
+            return extract.emitOpError("physical fragment slot is out of range"), failure();
+          add(getBatchChunkForLane(batch,
+                (*publication)->physicalSlotToGraphLane[*slot], processorCount), result);
+          continue;
+        }
       }
-
-      for (ComputeInstance instance :
-           getBatchChunksForRange(batch, 0,
-                                  static_cast<uint32_t>(batch.getLaneCount()),
-                                  processorCount))
-        producers.push_back({instance, 0});
-      return producers;
+      pending.push_back(source);
+      continue;
     }
-
-    value = source;
-    op = value.getDefiningOp();
-    if (!op)
-      return producers;
-  }
-
-  if (auto compute = dyn_cast<SpatCompute>(op)) {
-    producers.push_back({
-      ComputeInstance {compute.getOperation(), 0, 1},
-       static_cast<size_t>(cast<OpResult>(value).getResultNumber())
-    });
-    return producers;
-  }
-
-  if (auto batch = dyn_cast<SpatComputeBatch>(op)) {
-    if (batch.getNumResults() != 0) {
-      if (auto projected = collectProjectedProducerValueRefs(
-            batch, value, consumerInstance, processorCount))
-        return *projected;
-      std::optional<ProducerValueRef> producer =
-        getProducerValueRef(value, &consumerInstance, processorCount);
-      if (!producer)
-        return producers;
-      for (ComputeInstance instance :
-           getBatchChunksForRange(batch, producer->instance.laneStart,
-                                  producer->instance.laneCount,
-                                  processorCount))
-        producers.push_back({instance, 0});
-      return producers;
+    if (auto compute = dyn_cast<SpatCompute>(op)) {
+      add({compute, 0, 1}, cast<OpResult>(current).getResultNumber());
+      continue;
+    }
+    if (auto batch = dyn_cast<SpatComputeBatch>(op)) {
+      unsigned result = cast<OpResult>(current).getResultNumber();
+      for (ComputeInstance instance : getBatchChunksForRange(
+             batch, 0, batch.getLaneCount(), processorCount))
+        add(instance, result);
+      continue;
+    }
+    // Unproven shaping retains every producer, including captured region values.
+    llvm::append_range(pending, op->getOperands());
+    if (op->getNumRegions()) {
+      llvm::SetVector<Value> captures;
+      getUsedValuesDefinedAbove(op->getRegions(), captures);
+      llvm::append_range(pending, captures);
     }
-
-    uint32_t lane = cast<OpResult>(value).getResultNumber();
-    ComputeInstance instance =
-      getBatchChunkForLane(batch, lane, processorCount);
-    producers.push_back({instance, lane - instance.laneStart});
-    return producers;
   }
-
   return producers;
 }
 
+TransferCost getProducerTransferCost(Value input,
+    const ComputeInstance& consumerInstance, const ProducerValueRef& producerRef,
+    const SchedulerCostModel& costModel);
+
+LogicalResult collectBlueprintEdges(Value input, size_t targetIndex,
+                           const ComputeGraph& graph, size_t processorCount,
+                           const SchedulerCostModel& costModel,
+                           GraphBatchPublicationCache& publicationCache,
+                           SmallVectorImpl<ComputeGraphEdge>& edges) {
+  auto blueprint = input.getDefiningOp<SpatBlueprintOp>();
+  if (!blueprint || !isFragmentAssembly(blueprint.getMode()))
+    return success();
+  auto operands = blueprint.getFragmentOperandIndices();
+  auto slots = blueprint.getFragmentSourceSlots();
+  auto type = cast<RankedTensorType>(blueprint.getOutput().getType());
+  if (!operands || !slots)
+    return blueprint.emitOpError("compute graph requires complete fragment publication metadata");
+  SmallVector<Value> sources{blueprint.getInput()};
+  llvm::append_range(sources, blueprint.getFragments());
+  for (auto [fragment, operand] : llvm::enumerate(*operands)) {
+    Value source = sources[operand];
+    auto result = dyn_cast<OpResult>(source);
+    auto batch = result ? dyn_cast<SpatComputeBatch>(result.getOwner()) : SpatComputeBatch();
+    if (!batch) {
+      auto producers = collectProducerValueRefs(source, processorCount, publicationCache);
+      if (failed(producers))
+        return failure();
+      for (const auto &producer : *producers) {
+        auto indexed = graph.instanceToIndex.find(producer.instance);
+        if (indexed == graph.instanceToIndex.end())
+          return blueprint.emitOpError("fragment source has no compute instance");
+        edges.push_back({indexed->second, targetIndex,
+          getProducerTransferCost(source, graph.nodes[targetIndex].instance, producer, costModel)});
+      }
+      continue;
+    }
+    auto publication = getGraphBatchPublicationMap(
+      batch, result.getResultNumber(), publicationCache);
+    int64_t slot = (*slots)[fragment];
+    if (failed(publication))
+      return failure();
+    if (slot < 0
+        || slot >= static_cast<int64_t>((*publication)->physicalSlotToGraphLane.size()))
+      return blueprint.emitOpError("physical fragment slot is out of range");
+    int64_t lane = (*publication)->physicalSlotToGraphLane[slot];
+    ComputeInstance producer = getBatchChunkForLane(batch, lane, processorCount);
+    auto indexed = graph.instanceToIndex.find(producer);
+    if (indexed == graph.instanceToIndex.end())
+      return blueprint.emitOpError("fragment producer has no compute instance");
+    Cost elements = 1;
+    for (int64_t size : blueprint.getFragmentSizes().slice(
+           fragment * type.getRank(), type.getRank()))
+      elements = checkedMultiply(elements, static_cast<Cost>(size));
+    auto fragmentType = (*publication)->publicationFragmentType;
+    Cost bytes = costModel.getByteSize(fragmentType, costModel.getComputeBitwidth(fragmentType));
+    TransferCost transfer = costModel.getTransferCostFromBytes(bytes);
+    auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+    if (!sourceOffsets)
+      return blueprint.emitOpError("compute graph requires fragment source offsets");
+    // Full receives can write directly into the assembly destination. Its
+    // endpoint memory work is already included in the consumer's assembly cost.
+    if ((*sourceOffsets)[fragment] == 0
+        && elements == static_cast<Cost>(fragmentType.getNumElements()))
+      transfer.fixed = 0;
+    edges.push_back({indexed->second, targetIndex, transfer});
+  }
+  return success();
+}
+
 TransferCost getProducerTransferCost(
     Value input, const ComputeInstance& consumerInstance,
     const ProducerValueRef& producerRef,
     const SchedulerCostModel& costModel) {
   TransferCost transferCost =
-    getInputTransferCost(consumerInstance, input, costModel);
+    getInputTransferCost(input, costModel);
   auto producerBatch = dyn_cast<SpatComputeBatch>(producerRef.instance.op);
-  if (!producerBatch || producerBatch.getNumResults() == 0)
+  if (!producerBatch || producerBatch.getNumResults() == 0) {
+    if (auto projected = getProjectedInputTransferCost(input, consumerInstance, costModel))
+      return *projected;
     return transferCost;
-
-  if (auto consumerBatch = dyn_cast<SpatComputeBatch>(consumerInstance.op)) {
-    if (std::optional<TransferCost> projectedCost =
-          getBatchProjectedInputTransferCost(consumerBatch, input, costModel)) {
-      uint32_t overlapLaneCount = getLaneOverlapCount(consumerInstance, producerRef.instance);
-      return scaleTransferCost(
-        *projectedCost,
-        static_cast<Cost>(std::max<uint32_t>(1, overlapLaneCount)));
-    }
   }
 
-  return scaleTransferCostByLaneCount(
-    transferCost, static_cast<uint32_t>(producerBatch.getLaneCount()), producerRef.instance.laneCount);
+  auto type = cast<ShapedType>(producerBatch.getResult(producerRef.resultIndex).getType());
+  Cost elements = checkedMultiply(
+    static_cast<Cost>(type.getNumElements()) / producerBatch.getLaneCount(),
+    static_cast<Cost>(producerRef.instance.laneCount));
+  // Round the selected chunk once; dividing rounded whole-result bytes loses sub-byte payloads.
+  Cost bytes = SchedulerCostModel::ceilDiv(
+    checkedMultiply(elements, costModel.getComputeBitwidth(type)), Cost{8});
+  return costModel.getTransferCostFromBytes(bytes);
 }
 
 static ResidentWeight getOpaqueResidentWeight(Value value, std::optional<uint32_t> lane) {
@@ -819,13 +930,16 @@ ResidentWeightSet collectDistinctResidentWeights(Operation* owner, std::optional
   return usage;
 }
 
-Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target) {
-  SchedulerCostModel costModel {target};
+Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target,
+                            ComputeCostKind kind) {
+  SchedulerCostModel costModel {target, kind};
+  Cost assembly = getBlueprintAssemblyCost(instance, costModel);
   if (auto spatCompute = dyn_cast<SpatCompute>(instance.op))
-    return getComputeBodyCost(spatCompute.getBody(), costModel);
+    return checkedAdd(getComputeBodyCost(spatCompute.getBody(), costModel), assembly);
   auto batch = cast<SpatComputeBatch>(instance.op);
   return checkedMultiply(
-    getComputeBodyCost(batch.getBody(), costModel), static_cast<Cost>(instance.laneCount));
+    checkedAdd(getComputeBodyCost(batch.getBody(), costModel), assembly),
+    static_cast<Cost>(instance.laneCount));
 }
 
 bool containsResidentWeight(ArrayRef<ResidentWeight> usage, const ResidentWeight& weight) {
@@ -893,7 +1007,7 @@ ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& insta
   return tiled;
 }
 
-ComputeGraph buildComputeGraph(Operation* entryOp,
+FailureOr<ComputeGraph> buildComputeGraph(Operation* entryOp,
                                const SchedulingTarget& target,
                                size_t computePartitionCount) {
   ComputeGraph graph;
@@ -921,12 +1035,13 @@ ComputeGraph buildComputeGraph(Operation* entryOp,
             continue;
           size_t chunkCount =
             getBatchChunkTargetCount(batch, computePartitionCount);
+          Cost laneCost = getComputeInstanceCost({batch, 0, 1}, target);
           for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) {
             ComputeInstance instance = getBatchChunkForIndex(
               batch, chunkIndex, computePartitionCount);
             size_t index = graph.nodes.size();
             graph.nodes.push_back({instance,
-                                   getComputeInstanceCost(instance, target),
+                                   checkedMultiply(laneCost, static_cast<Cost>(instance.laneCount)),
                                    getComputeInstanceResidentWeights(instance, target),
                                    index});
             graph.instanceToIndex[instance] = index;
@@ -937,15 +1052,32 @@ ComputeGraph buildComputeGraph(Operation* entryOp,
   }
 
   llvm::SmallVector<ComputeGraphEdge, 16> rawEdges;
+  GraphBatchPublicationCache publicationCache;
   for (const auto& [targetIndex, node] : llvm::enumerate(graph.nodes)) {
     llvm::SmallVector<Value, 4> inputs = getComputeInstanceInputs(node.instance);
     for (Value input : inputs) {
-      for (const ProducerValueRef& producerRef :
-           collectProducerValueRefs(input, node.instance,
-                                    computePartitionCount)) {
+      if (auto producer = input.getDefiningOp<SpatComputeBatch>()) {
+        auto projected = collectProjectedEdges(producer, input, node.instance,
+          computePartitionCount, targetIndex, graph, costModel, publicationCache, rawEdges);
+        if (failed(projected))
+          return failure();
+        if (*projected)
+          continue;
+      }
+      if (auto blueprint = input.getDefiningOp<SpatBlueprintOp>();
+          blueprint && isFragmentAssembly(blueprint.getMode())) {
+        if (failed(collectBlueprintEdges(input, targetIndex, graph, computePartitionCount,
+              costModel, publicationCache, rawEdges)))
+          return failure();
+        continue;
+      }
+      auto producers = collectProducerValueRefs(input, computePartitionCount, publicationCache);
+      if (failed(producers))
+        return failure();
+      for (const ProducerValueRef& producerRef : *producers) {
         auto producerIt = graph.instanceToIndex.find(producerRef.instance);
         if (producerIt == graph.instanceToIndex.end())
-          continue;
+          return node.instance.op->emitOpError("input producer has no compute instance"), failure();
         rawEdges.push_back({producerIt->second,
                             targetIndex,
                             getProducerTransferCost(input, node.instance, producerRef, costModel)});

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp

@@ -37,16 +37,26 @@ namespace spatial {
 
 struct ComputeGraphNode {
   ComputeInstance instance;
-  Cost cost = 0;
+  Cost cost = 0; // Per-instance serial execution estimate, in nanoseconds.
   ResidentWeightSet residentWeights;
   size_t originalOrder = 0;
 };
 
 struct TransferCost {
-  Cost fixed = 0;
-  Cost networkFlits = 0;
+  Cost fixed = 0; // Endpoint memory time, in nanoseconds.
+  Cost networkFlits = 0; // Multiplied by the target's physical link time.
 };
 
+inline Time getTransferTime(const TransferCost& transferCost,
+                                size_t sourceProcessor,
+                                size_t targetProcessor,
+                                const SchedulingTarget& target) {
+  if (sourceProcessor == targetProcessor)
+    return 0;
+  return checkedAdd(transferCost.fixed,
+                    checkedMultiply(transferCost.networkFlits, target.getInterProcessorLatencyNs(sourceProcessor, targetProcessor)));
+}
+
 struct ComputeGraphEdge {
   size_t source = 0;
   size_t target = 0;
@@ -61,7 +71,7 @@ struct ComputeGraph {
   llvm::DenseMap<ComputeInstance, size_t> instanceToIndex;
 };
 
-ComputeGraph buildComputeGraph(mlir::Operation* entryOp,
+mlir::FailureOr<ComputeGraph> buildComputeGraph(mlir::Operation* entryOp,
                                const SchedulingTarget& target,
                                size_t computePartitionCount = 0);
 bool verifyAcyclic(const ComputeGraph& graph);
@@ -70,7 +80,14 @@ TransferCost getTransferCostFromBytes(Cost bytes,
                                       const SchedulingTarget& target);
 uint64_t countComputeBodyInstructions(mlir::Region& body);
 uint64_t countComputeBodyOperationInstances(mlir::Region& body);
-Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target);
+enum class ComputeCostKind {
+  Execution,
+  // A nominal volume-based scheduling weight, not predicted execution time.
+  VolumePriority,
+};
+
+Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target,
+                            ComputeCostKind kind = ComputeCostKind::Execution);
 ResidentWeightSet collectDistinctResidentWeights(mlir::Operation* owner,
                                                  std::optional<uint32_t> lane = std::nullopt);
 ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance);

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.cpp

@@ -0,0 +1,144 @@
+#include "MergeSchedule.hpp"
+#include "ComputeGraph.hpp"
+#include "llvm/ADT/STLExtras.h"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+using LanePublicationSignatures = llvm::SmallVector<llvm::SmallVector<int64_t, 8>, 8>;
+
+FailureOr<LanePublicationSignatures> buildLanePublicationSignatures(SpatComputeBatch batch,
+                                                                    GraphBatchPublicationCache& publicationCache) {
+  LanePublicationSignatures signatures(batch.getLaneCount());
+  for (auto [resultIndex, result] : llvm::enumerate(batch.getResults())) {
+    auto publicationMap = getGraphBatchPublicationMap(batch, resultIndex, publicationCache);
+    if (failed(publicationMap))
+      return failure();
+
+    for (auto [useIndex, use] : llvm::enumerate(result.getUses())) {
+      auto blueprint = dyn_cast<SpatBlueprintOp>(use.getOwner());
+      if (!blueprint || !isFragmentAssembly(blueprint.getMode()))
+        continue;
+      auto operandIndices = blueprint.getFragmentOperandIndices();
+      auto sourceSlots = blueprint.getFragmentSourceSlots();
+      auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+      auto fragmentStrides = blueprint.getFragmentStrides();
+      auto outputType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+      if (!operandIndices || !sourceSlots || !sourceOffsets || !fragmentStrides || !outputType
+          || !outputType.hasStaticShape())
+        return blueprint.emitOpError("PEFT publication compatibility requires complete static fragment metadata"),
+               failure();
+
+      llvm::ArrayRef<int64_t> fragmentOffsets = blueprint.getFragmentOffsets();
+      llvm::ArrayRef<int64_t> fragmentSizes = blueprint.getFragmentSizes();
+      int64_t rank = outputType.getRank();
+      if (rank <= 0 || fragmentOffsets.size() != fragmentSizes.size()
+          || fragmentOffsets.size() != fragmentStrides->size()
+          || fragmentOffsets.size() != operandIndices->size() * rank || sourceSlots->size() != operandIndices->size()
+          || sourceOffsets->size() != operandIndices->size())
+        return blueprint.emitOpError("PEFT publication compatibility found inconsistent fragment metadata"), failure();
+
+      llvm::SmallVector<llvm::SmallVector<size_t, 2>, 8> fragmentsByLane(batch.getLaneCount());
+      for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
+        if (operandIndex != static_cast<int64_t>(use.getOperandNumber()))
+          continue;
+        int64_t slot = (*sourceSlots)[fragmentIndex];
+        if (slot < 0 || slot >= static_cast<int64_t>((*publicationMap)->physicalSlotToGraphLane.size()))
+          return blueprint.emitOpError("PEFT publication fragment source slot is out of range"), failure();
+        int64_t graphLane = (*publicationMap)->physicalSlotToGraphLane[slot];
+        if (graphLane < 0 || graphLane >= batch.getLaneCount())
+          return blueprint.emitOpError("PEFT publication fragment has no graph lane owner"), failure();
+        fragmentsByLane[graphLane].push_back(fragmentIndex);
+      }
+
+      for (auto [lane, fragments] : llvm::enumerate(fragmentsByLane)) {
+        if (fragments.empty())
+          continue;
+        llvm::SmallVector<int64_t, 8>& signature = signatures[lane];
+        signature.push_back(resultIndex);
+        signature.push_back(useIndex);
+        signature.push_back(fragments.size());
+        signature.push_back(rank);
+        size_t firstFragment = fragments.front();
+        for (size_t fragmentIndex : fragments) {
+          signature.push_back((*sourceOffsets)[fragmentIndex] - (*sourceOffsets)[firstFragment]);
+          for (int64_t dim = 0; dim < rank; ++dim) {
+            size_t index = fragmentIndex * rank + dim;
+            size_t firstIndex = firstFragment * rank + dim;
+            signature.push_back(fragmentOffsets[index] - fragmentOffsets[firstIndex]);
+            signature.push_back(fragmentSizes[index]);
+            signature.push_back((*fragmentStrides)[index]);
+          }
+        }
+      }
+    }
+  }
+  return signatures;
+}
+
+} // namespace
+
+void populateScheduleEquivalenceClasses(
+    const ComputeGraph &graph,
+    const std::vector<std::vector<size_t>> &tasksByProcessor,
+    MergeScheduleResult &schedule) {
+  auto &equivalentClass = schedule.equivalentClass;
+  equivalentClass.clear();
+  size_t processorCount = tasksByProcessor.size();
+  GraphBatchPublicationCache publicationCache;
+  llvm::DenseMap<Operation *, LanePublicationSignatures> publicationSignatures;
+  llvm::DenseSet<Operation *> invalidPublicationSignatures;
+  auto haveCompatiblePublications = [&](const ComputeInstance& lhs, const ComputeInstance& rhs) {
+    auto batch = dyn_cast_or_null<SpatComputeBatch>(lhs.op);
+    if (!batch || batch.getNumResults() == 0)
+      return true;
+    auto signatures = publicationSignatures.find(lhs.op);
+    if (signatures == publicationSignatures.end() && !invalidPublicationSignatures.contains(lhs.op)) {
+      auto built = buildLanePublicationSignatures(batch, publicationCache);
+      if (failed(built))
+        invalidPublicationSignatures.insert(lhs.op);
+      else
+        signatures = publicationSignatures.try_emplace(lhs.op, std::move(*built)).first;
+    }
+    if (invalidPublicationSignatures.contains(lhs.op))
+      return false;
+    for (uint32_t lane = 0; lane < lhs.laneCount; ++lane)
+      if (signatures->second[lhs.laneStart + lane] != signatures->second[rhs.laneStart + lane])
+        return false;
+    return true;
+  };
+  for (size_t currentProcessor = 0; currentProcessor < processorCount; ++currentProcessor) {
+    if (tasksByProcessor[currentProcessor].empty())
+      continue;
+    for (size_t controlProcessor = currentProcessor + 1; controlProcessor < processorCount; ++controlProcessor) {
+      if (!schedule.processorStages.empty()
+          && schedule.processorStages[currentProcessor] != schedule.processorStages[controlProcessor])
+        continue;
+      if (tasksByProcessor[currentProcessor].size() != tasksByProcessor[controlProcessor].size())
+        continue;
+      auto& currentTasks = tasksByProcessor[currentProcessor];
+      auto& controlTasks = tasksByProcessor[controlProcessor];
+      bool equalSchedule = true;
+
+      for (auto [currentTask, controlTask] : llvm::zip(currentTasks, controlTasks)) {
+        const ComputeInstance currentComputeInstance = graph.nodes[currentTask].instance;
+        const ComputeInstance controlComputeInstance = graph.nodes[controlTask].instance;
+        if (!currentComputeInstance.op || currentComputeInstance.op != controlComputeInstance.op
+            || currentComputeInstance.laneCount != controlComputeInstance.laneCount
+            || !haveCompatiblePublications(currentComputeInstance, controlComputeInstance)) {
+          equalSchedule = false;
+          break;
+        }
+      }
+
+      if (equalSchedule) {
+        equivalentClass[currentProcessor].push_back(controlProcessor);
+        equivalentClass[controlProcessor].push_back(currentProcessor);
+      }
+    }
+  }
+
+}
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp

@@ -12,7 +12,18 @@
 namespace onnx_mlir {
 namespace spatial {
 
+struct ScheduleCostSummary {
+  uint64_t totalWorkNs = 0;
+  uint64_t maximumCoreWorkNs = 0;
+  uint64_t criticalPathNs = 0;
+  uint64_t communicationNs = 0;
+  uint64_t completionNs = 0;
+  std::vector<uint64_t> stageSpanNs;
+  std::vector<uint64_t> stageSerialDemandNs;
+};
+
 struct MergeScheduleResult {
+  ScheduleCostSummary costs;
   size_t processorCount = 0;
   std::vector<size_t> processorStages;
   std::vector<ComputeInstance> dominanceOrderCompute;
@@ -24,5 +35,11 @@ struct MergeScheduleResult {
   llvm::DenseMap<size_t, mlir::SmallVector<size_t, 5>> equivalentClass;
 };
 
+struct ComputeGraph;
+void populateScheduleEquivalenceClasses(
+    const ComputeGraph &graph,
+    const std::vector<std::vector<size_t>> &tasksByProcessor,
+    MergeScheduleResult &schedule);
+
 } // namespace spatial
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp

@@ -75,7 +75,7 @@ void verifySchedule(const ComputeGraph& graph,
     Time earliestTargetStart = addOrMax(sourceStart, graph.nodes[edge.source].cost);
     if (sourceCpu != targetCpu)
       earliestTargetStart = addOrMax(
-        earliestTargetStart, getPeftTransferTime(edge.transferCost, sourceCpu, targetCpu, target));
+        earliestTargetStart, getTransferTime(edge.transferCost, sourceCpu, targetCpu, target));
     if (targetStart < earliestTargetStart) {
       std::string message = llvm::formatv("merge scheduling: dependency legality failed between tasks {0} and {1}",
                                           graph.nodes[edge.source].originalOrder,
@@ -90,20 +90,24 @@ void verifySchedule(const ComputeGraph& graph,
 
 MergeSchedulingAnalysis::MergeSchedulingAnalysis(const ComputeGraph& computeGraph,
                                                  const SchedulingTarget& schedulingTarget,
-                                                 mlir::MLIRContext* context)
+                                                 mlir::MLIRContext* context, std::string& error)
 : context(context), target(schedulingTarget), graph(computeGraph) {
-  result = run();
+  result = run(error);
 }
 
-MergeScheduleResult MergeSchedulingAnalysis::run() {
-  if (!verifyAcyclic(graph))
-    llvm::report_fatal_error("merge scheduling: compute graph is cyclic");
+mlir::FailureOr<MergeScheduleResult> MergeSchedulingAnalysis::run(std::string& error) {
+  if (!verifyAcyclic(graph)) {
+    error = "merge scheduling: compute graph is cyclic";
+    return mlir::failure();
+  }
 
-  MergeScheduleResult schedule = runPeftScheduler(
+  auto schedule = runPeftScheduler(
     graph, PeftScheduleOptions {
       target,
-      context});
-  verifySchedule(graph, schedule, target);
+      context}, error);
+  if (mlir::failed(schedule))
+    return mlir::failure();
+  verifySchedule(graph, *schedule, target);
   return schedule;
 }
 

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <string>
+
 #include "mlir/IR/Operation.h"
 
 #include "MergeSchedule.hpp"
@@ -13,16 +15,16 @@ class MergeSchedulingAnalysis {
 public:
   MergeSchedulingAnalysis(const ComputeGraph& graph,
                           const SchedulingTarget& target,
-                          mlir::MLIRContext* context);
-  MergeScheduleResult& getResult() { return result; }
+                          mlir::MLIRContext* context, std::string& error);
+  mlir::FailureOr<MergeScheduleResult>& getResult() { return result; }
 
 private:
   mlir::MLIRContext* context = nullptr;
   const SchedulingTarget& target;
   const ComputeGraph& graph;
-  MergeScheduleResult result;
+  mlir::FailureOr<MergeScheduleResult> result;
 
-  MergeScheduleResult run();
+  mlir::FailureOr<MergeScheduleResult> run(std::string& error);
 };
 
 } // namespace spatial

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

@@ -1,6 +1,5 @@
 #include "mlir/IR/Threading.h"
 
-#include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 
@@ -11,7 +10,8 @@
 #include <vector>
 
 #include "PeftScheduler.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp"
+#include "ResidentWeightPacking.hpp"
+#include "ScheduleEvaluation.hpp"
 
 namespace onnx_mlir {
 namespace spatial {
@@ -32,12 +32,18 @@ struct ScheduledTask {
 
 struct TopologyModel {
   const SchedulingTarget& target;
+  std::vector<size_t> centerDistances;
+
+  explicit TopologyModel(const SchedulingTarget& target)
+      : target(target), centerDistances(target.processorCount) {
+    for (size_t processor = 0; processor < target.processorCount; ++processor)
+      for (size_t other = 0; other < target.processorCount; ++other)
+        centerDistances[processor] = checkedAdd(centerDistances[processor],
+          static_cast<size_t>(target.getInterProcessorLatencyNs(processor, other)));
+  }
 
   size_t getCenterDistance(size_t processor) const {
-    Cost total = 0;
-    for (size_t other = 0; other < target.processorCount; ++other)
-      total = checkedAdd(total, target.getInterProcessorLatencyNs(processor, other));
-    return static_cast<size_t>(total);
+    return centerDistances[processor];
   }
 };
 
@@ -93,18 +99,20 @@ bool hasHighResidentWeightPressure(const ComputeGraph& graph, size_t processorCo
   return false;
 }
 
-std::vector<ResidentWeightSet> planResidentWeightReservations(const ComputeGraph& graph,
+mlir::FailureOr<std::vector<ResidentWeightSet>> planResidentWeightReservations(const ComputeGraph& graph,
                                                     size_t processorCount,
                                                               size_t residentWeightCapacity,
                                                               const TopologyModel& topology,
                                                               bool preferWeightReuse,
-                                                              bool preserveTaskOrder = true) {
+                                                              std::string& error,
+                                                              bool largestWeightsFirst = false) {
   std::vector<size_t> weightedTasks;
   for (size_t task = 0; task < graph.nodes.size(); ++task)
     if (!graph.nodes[task].residentWeights.empty())
       weightedTasks.push_back(task);
   llvm::sort(weightedTasks, [&](size_t lhs, size_t rhs) {
-    if ((!preserveTaskOrder || preferWeightReuse) && graph.nodes[lhs].residentWeights.size() != graph.nodes[rhs].residentWeights.size())
+    if ((preferWeightReuse || largestWeightsFirst)
+        && graph.nodes[lhs].residentWeights.size() != graph.nodes[rhs].residentWeights.size())
       return graph.nodes[lhs].residentWeights.size() > graph.nodes[rhs].residentWeights.size();
     return graph.nodes[lhs].originalOrder < graph.nodes[rhs].originalOrder;
   });
@@ -125,11 +133,11 @@ std::vector<ResidentWeightSet> planResidentWeightReservations(const ComputeGraph
       Time communication = 0;
       for (const auto &[predecessor, cost] : graph.predecessors[task])
         if (reservedCpus[predecessor] != processorCount)
-          communication = addOrMax(communication, getPeftTransferTime(
+          communication = addOrMax(communication, getTransferTime(
             cost, reservedCpus[predecessor], processor, topology.target));
       for (const auto &[successor, cost] : graph.successors[task])
         if (reservedCpus[successor] != processorCount)
-          communication = addOrMax(communication, getPeftTransferTime(
+          communication = addOrMax(communication, getTransferTime(
             cost, processor, reservedCpus[successor], topology.target));
       Time load = addOrMax(reservedLoad[processor], communication);
       ReservationScore score {preferWeightReuse ? addedWeights : load,
@@ -143,19 +151,38 @@ std::vector<ResidentWeightSet> planResidentWeightReservations(const ComputeGraph
       }
     }
     if (bestProcessor == std::numeric_limits<size_t>::max()) {
-      // Task order improves locality but can fragment the available crossbars.
-      if (preserveTaskOrder && !preferWeightReuse)
+      // Preserve load-aware reservations before falling back to feasibility-only packing.
+      if (!preferWeightReuse && !largestWeightsFirst)
         return planResidentWeightReservations(graph, processorCount,
-          residentWeightCapacity, topology, preferWeightReuse, false);
-      std::string message =
-        llvm::formatv("PEFT reservation planner: cannot place task {0} with {1} distinct weights in {2} "
+          residentWeightCapacity, topology, preferWeightReuse, error, true);
+      IndexedResidentWeights indexed(graph);
+      std::vector<Cost> costs;
+      for (const ComputeGraphNode& node : graph.nodes)
+        costs.push_back(node.cost);
+      auto assignment = indexed.pack(
+        weightedTasks, costs, residentWeightCapacity, processorCount);
+      if (assignment) {
+        reservations.assign(processorCount, {});
+        for (auto [index, weightedTask] : llvm::enumerate(weightedTasks))
+          insertResidentWeights(reservations[(*assignment)[index]],
+                                graph.nodes[weightedTask].residentWeights);
+        return reservations;
+      }
+      size_t lowerBound = indexed.getPackCountLowerBound(
+        weightedTasks, residentWeightCapacity);
+      if (lowerBound > processorCount) {
+        error = "PEFT reservation planner: resident-weight lower bound exceeds physical capacity";
+        return mlir::failure();
+      }
+      error =
+        llvm::formatv("PEFT resident-weight heuristic search exhausted its bounded repair: could not place task {0} with {1} distinct weights in {2} "
                       "processors of capacity {3}",
                       graph.nodes[task].originalOrder,
                       graph.nodes[task].residentWeights.size(),
                       processorCount,
                       residentWeightCapacity)
           .str();
-      llvm::report_fatal_error(llvm::StringRef(message));
+      return mlir::failure();
     }
     insertResidentWeights(reservations[bestProcessor], graph.nodes[task].residentWeights);
     reservedCpus[task] = bestProcessor;
@@ -164,89 +191,24 @@ std::vector<ResidentWeightSet> planResidentWeightReservations(const ComputeGraph
   return reservations;
 }
 
-using LanePublicationSignatures = llvm::SmallVector<llvm::SmallVector<int64_t, 8>, 8>;
-
-FailureOr<LanePublicationSignatures> buildLanePublicationSignatures(SpatComputeBatch batch,
-                                                                    GraphBatchPublicationCache& publicationCache) {
-  LanePublicationSignatures signatures(batch.getLaneCount());
-  for (auto [resultIndex, result] : llvm::enumerate(batch.getResults())) {
-    auto publicationMap = getGraphBatchPublicationMap(batch, resultIndex, publicationCache);
-    if (failed(publicationMap))
-      return failure();
-
-    for (auto [useIndex, use] : llvm::enumerate(result.getUses())) {
-      auto blueprint = dyn_cast<SpatBlueprintOp>(use.getOwner());
-      if (!blueprint || !isFragmentAssembly(blueprint.getMode()))
-        continue;
-      auto operandIndices = blueprint.getFragmentOperandIndices();
-      auto sourceSlots = blueprint.getFragmentSourceSlots();
-      auto sourceOffsets = blueprint.getFragmentSourceOffsets();
-      auto fragmentStrides = blueprint.getFragmentStrides();
-      auto outputType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
-      if (!operandIndices || !sourceSlots || !sourceOffsets || !fragmentStrides || !outputType
-          || !outputType.hasStaticShape())
-        return blueprint.emitOpError("PEFT publication compatibility requires complete static fragment metadata"),
-               failure();
-
-      llvm::ArrayRef<int64_t> fragmentOffsets = blueprint.getFragmentOffsets();
-      llvm::ArrayRef<int64_t> fragmentSizes = blueprint.getFragmentSizes();
-      int64_t rank = outputType.getRank();
-      if (rank <= 0 || fragmentOffsets.size() != fragmentSizes.size()
-          || fragmentOffsets.size() != fragmentStrides->size()
-          || fragmentOffsets.size() != operandIndices->size() * rank || sourceSlots->size() != operandIndices->size()
-          || sourceOffsets->size() != operandIndices->size())
-        return blueprint.emitOpError("PEFT publication compatibility found inconsistent fragment metadata"), failure();
-
-      llvm::SmallVector<llvm::SmallVector<size_t, 2>, 8> fragmentsByLane(batch.getLaneCount());
-      for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
-        if (operandIndex != static_cast<int64_t>(use.getOperandNumber()))
-          continue;
-        int64_t slot = (*sourceSlots)[fragmentIndex];
-        if (slot < 0 || slot >= static_cast<int64_t>((*publicationMap)->physicalSlotToGraphLane.size()))
-          return blueprint.emitOpError("PEFT publication fragment source slot is out of range"), failure();
-        int64_t graphLane = (*publicationMap)->physicalSlotToGraphLane[slot];
-        if (graphLane < 0 || graphLane >= batch.getLaneCount())
-          return blueprint.emitOpError("PEFT publication fragment has no graph lane owner"), failure();
-        fragmentsByLane[graphLane].push_back(fragmentIndex);
-      }
-
-      for (auto [lane, fragments] : llvm::enumerate(fragmentsByLane)) {
-        if (fragments.empty())
-          continue;
-        llvm::SmallVector<int64_t, 8>& signature = signatures[lane];
-        signature.push_back(resultIndex);
-        signature.push_back(useIndex);
-        signature.push_back(fragments.size());
-        signature.push_back(rank);
-        size_t firstFragment = fragments.front();
-        for (size_t fragmentIndex : fragments) {
-          signature.push_back((*sourceOffsets)[fragmentIndex] - (*sourceOffsets)[firstFragment]);
-          for (int64_t dim = 0; dim < rank; ++dim) {
-            size_t index = fragmentIndex * rank + dim;
-            size_t firstIndex = firstFragment * rank + dim;
-            signature.push_back(fragmentOffsets[index] - fragmentOffsets[firstIndex]);
-            signature.push_back(fragmentSizes[index]);
-            signature.push_back((*fragmentStrides)[index]);
-          }
-        }
-      }
-    }
-  }
-  return signatures;
-}
-
 } // namespace
 
-MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options) {
+mlir::FailureOr<MergeScheduleResult> runPeftScheduler(const ComputeGraph& graph,
+    const PeftScheduleOptions& options, std::string& error) {
   const size_t nodeCount = graph.nodes.size();
   const size_t processorCount = options.target.processorCount;
-  if (processorCount == 0)
-    llvm::report_fatal_error("PEFT scheduler: processor count must be positive");
+  if (processorCount == 0 || options.target.residentWeightCapacity == 0) {
+    error = "PEFT scheduler: processor count and resident-weight capacity must be positive";
+    return mlir::failure();
+  }
   TopologyModel topology {options.target};
   const bool preferWeightReuse =
     hasHighResidentWeightPressure(graph, processorCount, options.target.residentWeightCapacity);
-  std::vector<ResidentWeightSet> capacityReservations = planResidentWeightReservations(
-    graph, processorCount, options.target.residentWeightCapacity, topology, preferWeightReuse);
+  auto reservations = planResidentWeightReservations(
+    graph, processorCount, options.target.residentWeightCapacity, topology, preferWeightReuse, error);
+  if (mlir::failed(reservations))
+    return mlir::failure();
+  std::vector<ResidentWeightSet> capacityReservations = std::move(*reservations);
 
   std::vector<std::vector<size_t>> reverseLevels = buildReverseLevels(graph);
 
@@ -268,30 +230,11 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
         computeNodeOct(i);
   }
 
-  struct RankEntry {
-    long double rank = 0.0L;
-    size_t node = 0;
-    size_t originalOrder = 0;
-  };
-  std::vector<RankEntry> ranks(nodeCount);
-  auto computeRank = [&](size_t node) {
-    long double rank = static_cast<long double>(oct[node]);
-    ranks[node] = {rank, node, graph.nodes[node].originalOrder};
-  };
-
-  if (options.context != nullptr)
-    mlir::parallelFor(options.context, 0, nodeCount, computeRank);
-  else
-    for (size_t node = 0; node < nodeCount; ++node)
-      computeRank(node);
-
   auto readyCompare = [&](size_t lhs, size_t rhs) {
-    const RankEntry& lhsRank = ranks[lhs];
-    const RankEntry& rhsRank = ranks[rhs];
-    if (lhsRank.rank != rhsRank.rank)
-      return lhsRank.rank < rhsRank.rank;
-    if (lhsRank.originalOrder != rhsRank.originalOrder)
-      return lhsRank.originalOrder > rhsRank.originalOrder;
+    if (oct[lhs] != oct[rhs])
+      return oct[lhs] < oct[rhs];
+    if (graph.nodes[lhs].originalOrder != graph.nodes[rhs].originalOrder)
+      return graph.nodes[lhs].originalOrder > graph.nodes[rhs].originalOrder;
     return lhs > rhs;
   };
 
@@ -309,8 +252,6 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     std::vector<ScheduledTask> schedules(nodeCount);
     std::vector<std::vector<size_t>> tasksByProcessor(processorCount);
     std::vector<std::vector<size_t>> timelineByProcessor(processorCount);
-    std::vector<size_t> taskOrder;
-    taskOrder.reserve(nodeCount);
 
     size_t scheduledCount = 0;
     while (!readyQueue.empty()) {
@@ -318,15 +259,12 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
       readyQueue.pop();
       if (scheduled[task])
         continue;
-      taskOrder.push_back(task);
 
       size_t bestProcessor = std::numeric_limits<size_t>::max();
       Time bestEst = 0;
       Time bestEft = 0;
-      Time bestOeft = std::numeric_limits<Time>::max();
-      unsigned int bestOverlapCount = 0;
-      size_t bestTaskCount = std::numeric_limits<size_t>::max();
-      size_t bestCenterDistance = std::numeric_limits<size_t>::max();
+      using PlacementScore = std::tuple<Time, Time, Time, size_t, size_t, unsigned, size_t>;
+      std::optional<PlacementScore> bestScore;
       size_t smallestResidentWeightUnion = std::numeric_limits<size_t>::max();
       bool residentWeightRejected = false;
 
@@ -343,7 +281,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
         Time dataReady = 0;
         for (const auto& [pred, comm] : graph.predecessors[task]) {
           const ScheduledTask& predSchedule = schedules[pred];
-          Time commPenalty = getPeftTransferTime(comm, predSchedule.processor, processor, options.target);
+          Time commPenalty = getTransferTime(comm, predSchedule.processor, processor, options.target);
           dataReady = std::max(dataReady, addOrMax(predSchedule.endTime, commPenalty));
         }
 
@@ -369,47 +307,14 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
         Time oeft = addOrMax(eft, oct[task]);
         size_t centerDistance = topology.getCenterDistance(processor);
         size_t taskCount = tasksByProcessor[processor].size();
-        bool betterResidentWeightChoice =
-          preferWeightReuse ? overlapCount > bestOverlapCount : overlapCount < bestOverlapCount;
-
-        if (oeft < bestOeft || (oeft == bestOeft && eft < bestEft)
-            || (oeft == bestOeft && eft == bestEft && est < bestEst)) {
+        PlacementScore score{oeft, eft, est, taskCount, centerDistance,
+          preferWeightReuse ? std::numeric_limits<unsigned>::max() - overlapCount : overlapCount,
+          processor};
+        if (!bestScore || score < *bestScore) {
           bestProcessor = processor;
           bestEst = est;
           bestEft = eft;
-          bestOeft = oeft;
-          bestOverlapCount = overlapCount;
-          bestTaskCount = taskCount;
-          bestCenterDistance = centerDistance;
-        }
-        else if (oeft == bestOeft && eft == bestEft && est == bestEst && taskCount < bestTaskCount) {
-          bestProcessor = processor;
-          bestEst = est;
-          bestEft = eft;
-          bestOeft = oeft;
-          bestOverlapCount = overlapCount;
-          bestTaskCount = taskCount;
-          bestCenterDistance = centerDistance;
-        }
-        else if (oeft == bestOeft && eft == bestEft && est == bestEst && taskCount == bestTaskCount
-                 && centerDistance < bestCenterDistance) {
-          bestProcessor = processor;
-          bestEst = est;
-          bestEft = eft;
-          bestOeft = oeft;
-          bestOverlapCount = overlapCount;
-          bestTaskCount = taskCount;
-          bestCenterDistance = centerDistance;
-        }
-        else if (oeft == bestOeft && eft == bestEft && est == bestEst && taskCount == bestTaskCount
-                 && centerDistance == bestCenterDistance && betterResidentWeightChoice) {
-          bestProcessor = processor;
-          bestEst = est;
-          bestEft = eft;
-          bestOeft = oeft;
-          bestOverlapCount = overlapCount;
-          bestTaskCount = taskCount;
-          bestCenterDistance = centerDistance;
+          bestScore = score;
         }
       }
 
@@ -460,35 +365,23 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
 
     if (chronologicalGaps) {
       tasksByProcessor = std::move(timelineByProcessor);
-      llvm::stable_sort(taskOrder, [&](size_t lhs, size_t rhs) {
-        return schedules[lhs].startTime < schedules[rhs].startTime;
-      });
     }
 
-    std::vector<Time> processorReady(processorCount, 0);
-    Time makespan = 0;
-    for (size_t task : taskOrder) {
-      size_t processor = schedules[task].processor;
-      Time dataReady = 0;
-      for (const auto& [pred, comm] : graph.predecessors[task]) {
-        const ScheduledTask& predecessor = schedules[pred];
-        Time transfer = getPeftTransferTime(comm, predecessor.processor, processor, options.target);
-        dataReady = std::max(dataReady, addOrMax(predecessor.endTime, transfer));
-      }
-      Time start = std::max(processorReady[processor], dataReady);
-      Time end = addOrMax(start, graph.nodes[task].cost);
-      schedules[task].startTime = start;
-      schedules[task].endTime = end;
-      processorReady[processor] = end;
-      makespan = std::max(makespan, end);
+    std::string error;
+    auto evaluated = evaluateSchedule(graph, tasksByProcessor, options.target, {}, error);
+    if (failed(evaluated))
+      llvm::report_fatal_error(llvm::Twine("PEFT scheduler: ") + error);
+    for (size_t task = 0; task < nodeCount; ++task) {
+      schedules[task].startTime = evaluated->taskStarts[task];
+      schedules[task].endTime = evaluated->taskEnds[task];
     }
-    return std::make_tuple(std::move(schedules), std::move(tasksByProcessor), makespan);
+    return std::make_tuple(std::move(schedules), std::move(tasksByProcessor), std::move(*evaluated));
   };
 
   auto chronological = placeTasks(true);
   auto materialization = placeTasks(false);
-  Time chronologicalMakespan = std::get<2>(chronological);
-  Time materializationMakespan = std::get<2>(materialization);
+  Time chronologicalMakespan = std::get<2>(chronological).completion;
+  Time materializationMakespan = std::get<2>(materialization).completion;
   Time chronologicalSaving =
     materializationMakespan > chronologicalMakespan ? materializationMakespan - chronologicalMakespan : 0;
   bool materiallyShorter = chronologicalSaving > materializationMakespan / kMakespanMaterialityDivisor;
@@ -505,58 +398,9 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     return graph.nodes[a].originalOrder < graph.nodes[b].originalOrder;
   });
 
-  // 5. Check if equal schedule in two level
-  llvm::DenseMap<size_t, mlir::SmallVector<size_t, 5>> equivalentClass;
-  GraphBatchPublicationCache publicationCache;
-  llvm::DenseMap<Operation *, LanePublicationSignatures> publicationSignatures;
-  llvm::DenseSet<Operation *> invalidPublicationSignatures;
-  auto haveCompatiblePublications = [&](const ComputeInstance& lhs, const ComputeInstance& rhs) {
-    auto batch = dyn_cast<SpatComputeBatch>(lhs.op);
-    if (!batch || batch.getNumResults() == 0)
-      return true;
-    auto signatures = publicationSignatures.find(lhs.op);
-    if (signatures == publicationSignatures.end() && !invalidPublicationSignatures.contains(lhs.op)) {
-      auto built = buildLanePublicationSignatures(batch, publicationCache);
-      if (failed(built))
-        invalidPublicationSignatures.insert(lhs.op);
-      else
-        signatures = publicationSignatures.try_emplace(lhs.op, std::move(*built)).first;
-    }
-    if (invalidPublicationSignatures.contains(lhs.op))
-      return false;
-    for (uint32_t lane = 0; lane < lhs.laneCount; ++lane)
-      if (signatures->second[lhs.laneStart + lane] != signatures->second[rhs.laneStart + lane])
-        return false;
-    return true;
-  };
-  for (size_t currentProcessor = 0; currentProcessor < processorCount - 1; ++currentProcessor) {
-    for (size_t controlProcessor = currentProcessor; controlProcessor < processorCount; ++controlProcessor) {
-      if (tasksByProcessor[currentProcessor].size() != tasksByProcessor[controlProcessor].size())
-        continue;
-      auto& currentTasks = tasksByProcessor[currentProcessor];
-      auto& controlTasks = tasksByProcessor[controlProcessor];
-      bool equalSchedule = true;
-
-      for (auto [currentTask, controlTask] : llvm::zip(currentTasks, controlTasks)) {
-        const ComputeInstance currentComputeInstance = graph.nodes[currentTask].instance;
-        const ComputeInstance controlComputeInstance = graph.nodes[controlTask].instance;
-        if (currentComputeInstance.op != controlComputeInstance.op
-            || currentComputeInstance.laneCount != controlComputeInstance.laneCount
-            || !haveCompatiblePublications(currentComputeInstance, controlComputeInstance)) {
-          equalSchedule = false;
-          break;
-        }
-      }
-
-      if (equalSchedule) {
-        equivalentClass[currentProcessor].push_back(controlProcessor);
-        equivalentClass[controlProcessor].push_back(currentProcessor);
-      }
-    }
-  }
-
   // 6. Populate Final Result
   MergeScheduleResult result;
+  recordScheduleCosts(std::get<2>(selected), result);
   result.processorCount = processorCount;
   result.dominanceOrderCompute.reserve(nodeCount);
 
@@ -578,7 +422,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     }
   }
 
-  result.equivalentClass = equivalentClass;
+  populateScheduleEquivalenceClasses(graph, tasksByProcessor, result);
 
   return result;
 }

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <string>
+
 #include "mlir/IR/MLIRContext.h"
 
 #include "ComputeGraph.hpp"
@@ -14,18 +16,9 @@ struct PeftScheduleOptions {
   mlir::MLIRContext* context = nullptr;
 };
 
-inline Time getPeftTransferTime(const TransferCost& transferCost,
-                                size_t sourceProcessor,
-                                size_t targetProcessor,
-                                const SchedulingTarget& target) {
-  if (sourceProcessor == targetProcessor)
-    return 0;
-  return checkedAdd(transferCost.fixed,
-                    checkedMultiply(transferCost.networkFlits, target.getInterProcessorLatencyNs(sourceProcessor, targetProcessor)));
-}
-
 // PEFT assigns physical cores using the injected target topology.
-MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options);
+mlir::FailureOr<MergeScheduleResult> runPeftScheduler(const ComputeGraph& graph,
+    const PeftScheduleOptions& options, std::string& error);
 
 } // namespace spatial
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp

@@ -10,6 +10,7 @@
 #include <numeric>
 #include <optional>
 #include <string>
+#include <tuple>
 #include <vector>
 
 #include "ComputeGraph.hpp"
@@ -18,6 +19,19 @@
 
 namespace onnx_mlir::spatial {
 
+using StagePlacementScore = std::tuple<Time, Time, Cost>;
+
+inline bool isBetterStagePlacement(const StagePlacementScore &candidate,
+                                   const StagePlacementScore &current) {
+  auto [demand, span, work] = candidate;
+  auto [oldDemand, oldSpan, oldWork] = current;
+  Time change = std::max(demand, oldDemand) - std::min(demand, oldDemand);
+  // Resolve near ties using latency, rather than magnifying proxy precision.
+  if (change <= std::max(demand, oldDemand) / 100)
+    return std::tie(span, demand, work) < std::tie(oldSpan, oldDemand, oldWork);
+  return candidate < current;
+}
+
 struct PipelineStageRange {
   size_t begin;
   size_t size;
@@ -85,7 +99,8 @@ private:
   std::vector<size_t> physicalCoreStages;
 };
 
-mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
+// Produces a complete physical schedule directly from the shared compute graph.
+mlir::LogicalResult runPipelineScheduler(const ComputeGraph& graph,
                                             MergeScheduleResult& schedule,
                                             size_t pipelineStages,
                                             const SchedulingTarget& physicalTarget,

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineWorkloadPreparation.cpp

@@ -0,0 +1,479 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
+
+#include "PipelineScheduling.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
+
+namespace onnx_mlir::spatial {
+namespace {
+using namespace mlir;
+
+struct SplitBody {
+  scf::ForOp loop;
+  SmallVector<SpatVMMOp, 32> vmms;
+};
+
+static void copyPipelineComputeAttrs(SpatGraphComputeBatch source,
+                                     SpatGraphComputeBatch destination) {
+  for (NamedAttribute attr : source->getAttrs()) {
+    StringRef name = attr.getName().getValue();
+    if (name == source.getOperandSegmentSizesAttrName()
+        || name == source.getLaneCountAttrName())
+      continue;
+    destination->setAttr(attr.getName(), attr.getValue());
+  }
+  destination->setAttr("pipeline.split", UnitAttr::get(source.getContext()));
+}
+
+static FailureOr<Value> createZeroTensor(OpBuilder &builder, Location loc, Type type) {
+  auto tensorType = dyn_cast<RankedTensorType>(type);
+  if (!tensorType || !tensorType.hasStaticShape())
+    return failure();
+  auto zero = DenseElementsAttr::get(
+    tensorType, builder.getZeroAttr(tensorType.getElementType()));
+  return arith::ConstantOp::create(builder, loc, zero).getResult();
+}
+
+static bool isInside(Operation *operation, Operation *ancestor) {
+  for (Operation *current = operation; current; current = current->getParentOp())
+    if (current == ancestor)
+      return true;
+  return false;
+}
+
+static FailureOr<SplitBody> findSplitBody(SpatGraphComputeBatch batch) {
+  if (batch.getNumResults() != 1 || batch.getBody().empty())
+    return failure();
+
+  SplitBody result;
+  Block &block = batch.getBody().front();
+  for (Operation &operation : block.without_terminator()) {
+    if (auto loop = dyn_cast<scf::ForOp>(&operation)) {
+      SmallVector<SpatVMMOp, 32> vmms;
+      loop.walk([&](SpatVMMOp vmm) { vmms.push_back(vmm); });
+      if (!vmms.empty()) {
+        if (result.loop)
+          return failure();
+        result.loop = loop;
+        result.vmms = std::move(vmms);
+      }
+      continue;
+    }
+    bool containsVmm = false;
+    operation.walk([&](SpatVMMOp) { containsVmm = true; });
+    if (containsVmm)
+      return failure();
+  }
+  if (!result.loop || result.vmms.size() < 2 || result.loop.getNumResults() != 1)
+    return failure();
+  if (!isa<SpatInParallelOp>(block.getTerminator()))
+    return failure();
+  if (!isa<RankedTensorType>(result.loop.getResult(0).getType()))
+    return failure();
+  return result;
+}
+
+static LogicalResult verifySplittableVmmUses(ArrayRef<SpatVMMOp> vmms,
+                                             std::string &error) {
+  for (SpatVMMOp vmm : vmms)
+    for (OpOperand &use : vmm.getResult().getUses())
+      if (!isa<SpatVAddOp>(use.getOwner())) {
+        error = "pipeline task split requires VMM results to feed an additive reduction";
+        return failure();
+      }
+  return success();
+}
+
+static SmallVector<BlockArgument, 4> collectAdditiveInputArguments(
+    SpatGraphComputeBatch batch, scf::ForOp loop) {
+  SmallVector<BlockArgument, 4> result;
+  Block &block = batch.getBody().front();
+  for (BlockArgument argument : block.getArguments()) {
+    if (argument == *batch.getLaneArgument())
+      continue;
+    bool hasUse = false;
+    bool additive = true;
+    for (OpOperand &use : argument.getUses()) {
+      if (!isInside(use.getOwner(), loop)) {
+        additive = false;
+        break;
+      }
+      hasUse = true;
+      if (!isa<SpatVAddOp>(use.getOwner())) {
+        additive = false;
+        break;
+      }
+    }
+    if (hasUse && additive && isa<RankedTensorType>(argument.getType()))
+      result.push_back(argument);
+  }
+  return result;
+}
+
+static void simplifyZeroAdds(Operation *root, ArrayRef<Value> zeroValues) {
+  bool changed = true;
+  while (changed) {
+    changed = false;
+    SmallVector<SpatVAddOp, 16> adds;
+    root->walk([&](SpatVAddOp add) { adds.push_back(add); });
+    for (SpatVAddOp add : adds) {
+      Value replacement;
+      if (llvm::is_contained(zeroValues, add.getLhs()))
+        replacement = add.getRhs();
+      else if (llvm::is_contained(zeroValues, add.getRhs()))
+        replacement = add.getLhs();
+      if (!replacement)
+        continue;
+      add.getResult().replaceAllUsesWith(replacement);
+      add.erase();
+      changed = true;
+    }
+  }
+}
+
+static FailureOr<Operation *> cloneSplitLoop(OpBuilder &builder,
+                                             scf::ForOp sourceLoop,
+                                             IRMapping &mapper,
+                                             ArrayRef<unsigned> selectedVmmIndices,
+                                             ArrayRef<SpatVMMOp> sourceVmms) {
+  Operation *cloned = builder.clone(*sourceLoop.getOperation(), mapper);
+  SmallVector<SpatVMMOp, 32> clonedVmms;
+  cloned->walk([&](SpatVMMOp vmm) { clonedVmms.push_back(vmm); });
+  if (clonedVmms.size() != sourceVmms.size()) {
+    cloned->erase();
+    return failure();
+  }
+
+  SmallVector<Value, 8> zeroValues;
+  for (auto [index, vmm] : llvm::enumerate(clonedVmms)) {
+    if (llvm::is_contained(selectedVmmIndices, static_cast<unsigned>(index)))
+      continue;
+    OpBuilder zeroBuilder(vmm.getOperation());
+    FailureOr<Value> zero = createZeroTensor(zeroBuilder, vmm.getLoc(), vmm.getResult().getType());
+    if (failed(zero)) {
+      cloned->erase();
+      return failure();
+    }
+    zeroValues.push_back(*zero);
+    vmm.getResult().replaceAllUsesWith(*zero);
+    vmm.erase();
+  }
+  simplifyZeroAdds(cloned, zeroValues);
+  return cloned;
+}
+
+static SmallVector<OpFoldResult> makeFragmentOffsets(OpBuilder &builder,
+                                                      Value lane,
+                                                      RankedTensorType fragmentType) {
+  SmallVector<OpFoldResult> offsets {lane};
+  offsets.append(fragmentType.getRank(), builder.getIndexAttr(0));
+  return offsets;
+}
+
+static SmallVector<OpFoldResult> makeFragmentSizes(OpBuilder &builder,
+                                                   RankedTensorType fragmentType) {
+  SmallVector<OpFoldResult> sizes {builder.getIndexAttr(1)};
+  for (int64_t dimension : fragmentType.getShape())
+    sizes.push_back(builder.getIndexAttr(dimension));
+  return sizes;
+}
+
+static SmallVector<OpFoldResult> makeFragmentStrides(OpBuilder &builder,
+                                                      RankedTensorType fragmentType) {
+  return SmallVector<OpFoldResult>(
+    fragmentType.getRank() + 1, builder.getIndexAttr(1));
+}
+
+static FailureOr<SpatGraphComputeBatch> createPartialBatch(
+    OpBuilder &builder,
+    SpatGraphComputeBatch source,
+    SplitBody &splitBody,
+    ArrayRef<unsigned> selectedVmmIndices,
+    ArrayRef<BlockArgument> additiveInputs) {
+  auto fragmentType = dyn_cast<RankedTensorType>(
+    splitBody.loop.getOperation()->getResult(0).getType());
+  if (!fragmentType || !fragmentType.hasStaticShape())
+    return failure();
+  RankedTensorType resultType =
+    getGraphBatchPhysicalResultType(source.getLaneCount(), fragmentType);
+  SmallVector<Value> partialInputs;
+  for (auto [index, input] : llvm::enumerate(source.getInputs()))
+    if (!llvm::is_contained(additiveInputs, *source.getInputArgument(index)))
+      partialInputs.push_back(input);
+  auto partial = createEmptySpatGraphComputeBatch(
+    builder, source.getLoc(), TypeRange {resultType}, source.getLaneCount(),
+    source.getWeights(), partialInputs);
+  if (failed(partial))
+    return failure();
+  copyPipelineComputeAttrs(source, *partial);
+
+  IRMapping mapper;
+  mapper.map(*source.getLaneArgument(), *(*partial).getLaneArgument());
+  for (size_t index = 0; index < source.getWeights().size(); ++index)
+    mapper.map(*source.getWeightArgument(index), *(*partial).getWeightArgument(index));
+  unsigned partialInputIndex = 0;
+  for (size_t index = 0; index < source.getInputs().size(); ++index) {
+    BlockArgument sourceArgument = *source.getInputArgument(index);
+    if (llvm::is_contained(additiveInputs, sourceArgument))
+      continue;
+    mapper.map(sourceArgument, *(*partial).getInputArgument(partialInputIndex++));
+  }
+  for (size_t index = 0; index < source.getOutputs().size(); ++index)
+    mapper.map(*source.getOutputArgument(index), *(*partial).getOutputArgument(index));
+  Block &sourceBlock = source.getBody().front();
+  Block &partialBlock = (*partial).getBody().front();
+  builder.setInsertionPointToEnd(&partialBlock);
+  for (Operation &operation : sourceBlock.without_terminator()) {
+    if (&operation == splitBody.loop)
+      break;
+    builder.clone(operation, mapper);
+  }
+
+  for (BlockArgument input : additiveInputs) {
+    FailureOr<Value> zero = createZeroTensor(
+      builder, input.getLoc(), input.getType());
+    if (failed(zero)) {
+      partial->erase();
+      return failure();
+    }
+    mapper.map(input, *zero);
+  }
+
+  if (failed(cloneSplitLoop(builder, splitBody.loop, mapper,
+                            selectedVmmIndices, splitBody.vmms))) {
+    partial->erase();
+    return failure();
+  }
+  auto partialLoop = cast<scf::ForOp>(
+    &partialBlock.back());
+  auto output = (*partial).getOutputArgument(0);
+  if (!output) {
+    partial->erase();
+    return failure();
+  }
+
+  builder.setInsertionPointToEnd(&partialBlock);
+  auto inParallel = SpatInParallelOp::create(builder, source.getLoc());
+  builder.setInsertionPointToEnd(&inParallel.getRegion().front());
+  tensor::ParallelInsertSliceOp::create(
+    builder, source.getLoc(), partialLoop.getResult(0), *output,
+    makeFragmentOffsets(builder, *(*partial).getLaneArgument(), fragmentType),
+    makeFragmentSizes(builder, fragmentType),
+    makeFragmentStrides(builder, fragmentType));
+  return *partial;
+}
+
+static FailureOr<Value> extractBatchFragment(OpBuilder &builder,
+                                             Value physicalValue,
+                                             Value lane,
+                                             RankedTensorType fragmentType,
+                                             Location loc) {
+  SmallVector<OpFoldResult> offsets = makeFragmentOffsets(builder, lane, fragmentType);
+  SmallVector<OpFoldResult> sizes = makeFragmentSizes(builder, fragmentType);
+  SmallVector<OpFoldResult> strides = makeFragmentStrides(builder, fragmentType);
+  return tensor::ExtractSliceOp::create(
+    builder, loc, fragmentType, physicalValue, offsets, sizes, strides).getResult();
+}
+
+static bool hasUnsupportedContinuationOperands(SpatGraphComputeBatch source,
+                                               scf::ForOp loop) {
+  Block &block = source.getBody().front();
+  bool afterLoop = false;
+  for (Operation &operation : block.without_terminator()) {
+    if (&operation == loop) {
+      afterLoop = true;
+      continue;
+    }
+    if (!afterLoop)
+      continue;
+    for (Value operand : operation.getOperands()) {
+      if (auto argument = dyn_cast<BlockArgument>(operand))
+        if (argument.getOwner() == &block)
+          return true;
+      Operation *definingOp = operand.getDefiningOp();
+      if (definingOp && definingOp->isBeforeInBlock(loop))
+        return true;
+    }
+  }
+  return false;
+}
+
+static FailureOr<SpatGraphComputeBatch> createContinuationBatch(
+    OpBuilder &builder,
+    SpatGraphComputeBatch source,
+    SplitBody &splitBody,
+    ArrayRef<Value> partialResults) {
+  if (partialResults.size() < 2
+      || hasUnsupportedContinuationOperands(source, splitBody.loop))
+    return failure();
+
+  SmallVector<Value> inputs(partialResults.begin(), partialResults.end());
+  auto continuation = createEmptySpatGraphComputeBatch(
+    builder, source.getLoc(), source.getResultTypes(), source.getLaneCount(),
+    ValueRange {}, inputs);
+  if (failed(continuation))
+    return failure();
+  copyPipelineComputeAttrs(source, *continuation);
+
+  IRMapping mapper;
+  mapper.map(*source.getLaneArgument(), *(*continuation).getLaneArgument());
+  for (size_t index = 0; index < source.getOutputs().size(); ++index)
+    mapper.map(*source.getOutputArgument(index), *(*continuation).getOutputArgument(index));
+  Block &block = (*continuation).getBody().front();
+  builder.setInsertionPointToEnd(&block);
+  auto fragmentType = dyn_cast<RankedTensorType>(
+    splitBody.loop.getOperation()->getResult(0).getType());
+  if (!fragmentType || !fragmentType.hasStaticShape()) {
+    continuation->erase();
+    return failure();
+  }
+
+  SmallVector<Value> partialFragments;
+  for (auto [index, partial] : llvm::enumerate(partialResults)) {
+    auto input = (*continuation).getInputArgument(index);
+    if (!input) {
+      continuation->erase();
+      return failure();
+    }
+    FailureOr<Value> fragment = extractBatchFragment(
+      builder, *input, *(*continuation).getLaneArgument(), fragmentType,
+      source.getLoc());
+    if (failed(fragment)) {
+      continuation->erase();
+      return failure();
+    }
+    partialFragments.push_back(*fragment);
+  }
+
+  Value combined = partialFragments.front();
+  for (Value fragment : llvm::drop_begin(partialFragments))
+    combined = SpatVAddOp::create(
+      builder, source.getLoc(), fragmentType, combined, fragment).getResult();
+  mapper.map(splitBody.loop.getOperation()->getResult(0), combined);
+
+  bool afterLoop = false;
+  for (Operation &operation : source.getBody().front().without_terminator()) {
+    if (&operation == splitBody.loop) {
+      afterLoop = true;
+      continue;
+    }
+    if (afterLoop)
+      builder.clone(operation, mapper);
+  }
+  builder.clone(*source.getBody().front().getTerminator(), mapper);
+  return *continuation;
+}
+
+static SmallVector<SmallVector<unsigned, 8>, 4>
+partitionVmmIndices(size_t vmmCount, size_t partCount) {
+  SmallVector<SmallVector<unsigned, 8>, 4> partitions;
+  size_t next = 0;
+  for (size_t part = 0; part < partCount; ++part) {
+    size_t remainingParts = partCount - part;
+    size_t remainingVmms = vmmCount - next;
+    size_t count = (remainingVmms + remainingParts - 1) / remainingParts;
+    SmallVector<unsigned, 8> selected;
+    for (size_t index = 0; index < count; ++index)
+      selected.push_back(static_cast<unsigned>(next + index));
+    next += count;
+    partitions.push_back(std::move(selected));
+  }
+  return partitions;
+}
+
+static void eraseCreatedSplits(ArrayRef<Operation *> created) {
+  for (Operation *operation : llvm::reverse(created))
+    operation->erase();
+}
+
+static LogicalResult splitBatchCompute(SpatGraphComputeBatch batch,
+                                       size_t pipelineStages,
+                                       const SchedulingTarget &physicalTarget,
+                                       std::string &error) {
+  if (batch->hasAttr("pipeline.split"))
+    return failure();
+  FailureOr<SplitBody> splitBody = findSplitBody(batch);
+  if (failed(splitBody))
+    return failure();
+  if (failed(verifySplittableVmmUses(splitBody->vmms, error)))
+    return failure();
+  if (!batch->hasAttr("pipeline.stage_group"))
+    batch->setAttr(
+      "pipeline.stage_group",
+      DistinctAttr::create(UnitAttr::get(batch.getContext())));
+
+  size_t partCount = std::min(pipelineStages, splitBody->vmms.size());
+  SmallVector<SmallVector<unsigned, 8>, 4> partitions =
+    partitionVmmIndices(splitBody->vmms.size(), partCount);
+  SmallVector<BlockArgument, 4> additiveInputs =
+    collectAdditiveInputArguments(batch, splitBody->loop);
+  SmallVector<Operation *, 4> created;
+  SmallVector<Value, 4> partialResults;
+  OpBuilder builder(batch.getOperation());
+  for (auto [partIndex, partition] : llvm::enumerate(partitions)) {
+    builder.setInsertionPoint(batch.getOperation());
+    ArrayRef<BlockArgument> zeroInputs =
+      partIndex == 0 ? ArrayRef<BlockArgument>() : additiveInputs;
+    FailureOr<SpatGraphComputeBatch> partial = createPartialBatch(
+      builder, batch, *splitBody, partition, zeroInputs);
+    if (failed(partial)) {
+      eraseCreatedSplits(created);
+      return failure();
+    }
+    ComputeInstance instance {
+      partial->getOperation(), 0,
+      static_cast<uint32_t>(partial->getLaneCount())};
+    if (getComputeInstanceResidentWeights(instance, physicalTarget).size()
+        > physicalTarget.residentWeightCapacity) {
+      eraseCreatedSplits(created);
+      partial->erase();
+      return failure();
+    }
+    created.push_back(partial->getOperation());
+    partialResults.push_back(partial->getResult(0));
+  }
+
+  builder.setInsertionPoint(batch.getOperation());
+  FailureOr<SpatGraphComputeBatch> continuation = createContinuationBatch(
+    builder, batch, *splitBody, partialResults);
+  if (failed(continuation)) {
+    eraseCreatedSplits(created);
+    return failure();
+  }
+  created.push_back(continuation->getOperation());
+  batch->replaceAllUsesWith(continuation->getOperation()->getResults());
+  batch->erase();
+  return success();
+}
+
+static FailureOr<PipelineWorkloadPreparation> preparePipelineWorkloadImpl(
+    const ComputeGraph &graph, size_t pipelineStages,
+    const SchedulingTarget &physicalTarget, std::string &error) {
+  for (const ComputeGraphNode &node : graph.nodes) {
+    if (node.residentWeights.size() <= physicalTarget.residentWeightCapacity)
+      continue;
+    auto batch = dyn_cast_or_null<SpatGraphComputeBatch>(node.instance.op);
+    if (batch && succeeded(splitBatchCompute(
+                   batch, pipelineStages, physicalTarget, error)))
+      return PipelineWorkloadPreparation::Changed;
+    if (error.empty())
+      error = "pipeline scheduling cannot split a compute instance to fit one "
+              "physical core's crossbars";
+    return failure();
+  }
+  return PipelineWorkloadPreparation::Ready;
+}
+
+} // namespace
+
+mlir::FailureOr<PipelineWorkloadPreparation> preparePipelineWorkload(
+    const ComputeGraph &graph, size_t pipelineStages,
+    const SchedulingTarget &physicalTarget, std::string &error) {
+  return preparePipelineWorkloadImpl(
+    graph, pipelineStages, physicalTarget, error);
+}
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/README.md

@@ -0,0 +1,177 @@
+# Spatial scheduling
+
+`ScheduleAndRealizeSpatial` selects one scheduling strategy after constructing
+the shared `ComputeGraph`. Pipeline scheduling does not run, require, or repair
+a latency PEFT placement.
+
+## Ownership and flow
+
+| Decision | Owner |
+| --- | --- |
+| Compute instances, dependency extraction, execution/transfer costs, resident-weight identities | `ComputeGraph` and `ComputeInstanceUtils` |
+| Target configuration and physical link costs | Compiler target adapter, exposed through `SchedulingTarget` |
+| Latency priorities, reservations, physical assignment, insertion gaps, per-core order | `PeftScheduler` |
+| Capacity-driven additive workload splitting | `PipelineWorkloadPreparation` |
+| Operation groups, stage adjacency and input constraints | `buildPipelineGroups` |
+| Stage partitioning | `partitionPipelineGroups` |
+| Stage core budgets | `assignPipelineStages` |
+| Physical topology clustering | `PipelineCoreLayout` |
+| Stage-local placement and resident-weight packing | `packPipelineStages` and `IndexedResidentWeights` |
+| Schedule timing evaluation | `evaluateSchedule` |
+| Publication-compatible core equivalence classes | `populateScheduleEquivalenceClasses` |
+| Emission of the selected core order and homogeneous runs | Scheduled compute planning/materialization |
+| Routing, local forwarding and transfer families | Deferred transfer planning |
+| Deadlock-free blocking communication order | Deferred communication scheduling |
+| READY/FREE counts, registers, overwrite boundaries and restart synchronization | Deferred boundary planning |
+| Explicit communication and synchronization operations | Deferred boundary realization |
+
+Both strategies produce `MergeScheduleResult`: physical core assignment,
+contiguous per-core execution slots, estimated starts, last-task markers,
+equivalence classes and cost summaries. Pipeline schedules also provide the
+stage of every physical core. Graph identity identifies publications; it does
+not order execution. Materialization sorts only by the verified scheduler slots
+and combines only consecutive compatible graph instances.
+Class lanes keep graph encounter order independently of physical core IDs,
+preserving regular source-lane mappings without changing any core's execution
+order. Complete singleton fragment publications use reshape views only with
+matching static shapes, zero offsets, complete sizes and unit strides, including
+scalar fragments. Other insertions retain their slice semantics. Receive forwarding
+runs after copy normalization and can cross lookup-table loads, while writes
+and synchronization remain barriers.
+
+Physical fragment slots are resolved through `getGraphBatchPublicationMap`.
+Direct statically evaluable batch projections can have many-to-many edges.
+Unproven projections retain the complete producer dependency. Blueprint
+dependencies belong to the shared graph, rather than a separate pipeline-only
+dependency graph.
+Sources reached through intermediate projections retain conservative producer
+dependencies, including region captures. Graph construction reports malformed
+publication metadata through normal pass failure.
+Projected communication costs count the selected physical fragments through
+the same publication map; graph-range overlap is not a transfer-volume estimate.
+Conservative costs scale payload bytes before rounding memory/packet costs, so
+small fragments retain endpoint and packet overhead. Conservative producer-chunk
+payloads scale elements/bits before byte rounding, preserving one aggregate
+chunk estimate; projected publication costs retain their per-fragment granularity.
+Spatial sub-byte types reach analysis, while the current PIM byte-addressed
+communication lowering rejects non-byte-sized element storage.
+
+## Cost dimensions
+
+Execution times and fixed transfer costs use integer nanoseconds. Hardware
+cycles are converted using the corresponding target period. Bytes, elements,
+flits, resident-weight slots and lane counts remain resource quantities.
+Blueprint reconstruction charges memory work separately for each fragment;
+small projected slices are not assumed to require a single copy. Full receives
+that can write directly into assembly storage do not charge endpoint memory work
+again on their edges. Later copy elimination can make this assembly estimate
+conservative. Blueprint reconstruction uses target memory read/write costs; element counts
+and cycles are not added to nanoseconds. Runtime-free views cost zero.
+Runtime transposes charge target memory read/write traffic rather than generic
+nested-operation fallback cost. Their shared lowering/run analysis charges rounded
+memory accesses per run and vector execution for strided runs. The matched oracle
+charges vector payload bytes, not the full strided address footprint. Instruction
+setup/issue time and contention remain approximation boundaries.
+Concatenation charges each input block separately
+for each outer row, including memory-width rounding. Axis, size and target-period
+scaling have focused coverage.
+
+The evaluator distinguishes total serial work, maximum per-core work,
+dependency critical path, summed edge transfer latency, single-inference
+completion and individual stage spans. The dependency path excludes core-order
+constraints; completion includes them. Summed communication is a work metric,
+not elapsed time. Stage spans exclude cross-stage routing, restart/host-buffer
+synchronization and network contention. They are scheduling proxies, not an
+authoritative prediction of steady-state throughput. Simulator latency and
+measured throughput period remain the performance gates.
+
+Partitioning retains its successful ordered-packing plus dependency-tail
+priority. This deliberately biases cuts away from long dependent groups and
+is not reported as predicted execution time. Pipeline partitioning and packing
+use nominal volume priorities: transpose traffic is weighted as a bulk movement,
+preserving the measured heuristic without pretending that individual runs are
+free. The per-operation priority adjustment is cached and scaled by chunk lanes.
+PEFT and final schedule evaluation always use physical execution estimates.
+Pipeline candidate placement compares
+maximum per-core serial compute/incident-transfer demand, then stage completion
+and compute load. Incident transfer latency is charged to each participating
+core as a placement proxy; it is not added twice to predicted elapsed time.
+This preserves useful parallelism that an arrival-time-only ranking can lose.
+The final reported compute/incident-transfer demand uses physical execution
+estimates, separately from the priority scores used to select the placement.
+When serial-demand estimates are within 1%, stage span breaks the near tie;
+this is a heuristic resolution rule, not a pipeline timing model. Core budgets
+come from stage allocation and capacity recovery. A proposed additional one-core
+load refinement was removed after a matched throughput regression; improving
+the placement proxy did not establish a better execution period.
+
+## Capacity and recovery
+
+The distinct-weight union divided by per-core capacity is a provable lower
+bound; an oversized task is individually infeasible. Ordered pack counts are
+heuristic core-budget estimates for unrestricted placement. Their failure does
+not prove hardware infeasibility. When these estimates exhaust the fleet, core
+allocation restarts from the lower bounds and actual packing establishes
+feasibility.
+
+Latency reservations prioritize physical communication/load, or reuse under
+high weight pressure. If graph-order reservations fail, a largest-weight-first
+retry retains load-aware placement before using the shared reuse-first best-fit
+constructor. The constructor establishes feasibility; its concentrated assignment
+does not replace a successful load-aware reservation. Pipeline packing compares arrival-aware and load-aware placements;
+ordered packing and the same shared constructor supply constructive recovery
+assignments. Failed search budgets guide bounded core reallocation, but are not
+resource lower bounds. Diagnostics distinguish a proved capacity limit from an
+unsuccessful heuristic search. No exact combinatorial solver is used.
+The shared constructor retries once in reverse order after greedy failure.
+Only if both orders fail, it replays at most16 single-choice deviations with a
+shared budget of65536 core-candidate checks, reusing its scratch buffers.
+Successful existing assignments are preserved; exhausted repair is not proof of
+infeasibility. PEFT propagates reservation failures to the scheduling pass for
+a bounded diagnostic and normal pass failure.
+
+## Complexity
+
+Let `T` be tasks, `E` dependencies, `P` processors, `G` operation groups, `S`
+stages, `R` resident-weight references, `W` distinct weights and `C` total modeled
+work.
+
+- Schedule evaluation takes `O((T + E) log T + P)` time with deterministic ready
+  ordering, and `O(T + E + P)` space.
+- Resident-weight indexing retains `O(R W)` worst-case equality comparisons.
+  Lower bounds take `O(R + W)` time and `O(W)` bits. Constructive best fit takes
+  `O(T log T + P R + P W)` time, with `O(P W)` bits and `O(T + P)` additional
+  storage. It runs on failure paths rather than replacing successful placement.
+  Recovery adds at most17 constructions with the same complexity (one reverse
+  traversal plus16 single-choice replays); repair candidate checks are also
+  capped at65536. Scratch storage is reused, with one extra `O(T)` failed-assignment
+  snapshot only after both original attempts fail.
+- The partition DP remains superlinear: up to
+  `O(S G² ((T + R) log C + W))` time for ordered segment estimates, plus group
+  dependency traversal, and `O(S G²)` cache space. This bounded existing search
+  is retained rather than replaced with global placement optimization.
+- Physical topology clustering and cached PEFT center distances take `O(P²)`
+  time. Center lookups are constant time instead of repeating a full physical
+  topology scan for each task/candidate processor.
+- Concatenation cost extraction takes `O(A D)` time and constant auxiliary space
+  for `A` inputs of rank `D`; it does not expand the outer rows or copy blocks.
+- Graph execution costs are scaled by chunk lane counts. Pipeline priority
+  analysis adds two cost traversals per operation template and caches the
+  adjustment, taking `O(T + D)` time and `O(T + G)` space for `D` total template
+  body operations. It does not revisit template bodies for every chunk.
+- Exact projection analysis visits each selected physical slot once per
+  consumer chunk and input use, with a producer-lane bitset for deduplication.
+  If `L` is producer lanes and `Q` is projected slot visits, it takes
+  `O(Q + L)` time and `O(L + Q)` temporary space; publication maps are cached
+  across all chunks. Cost extraction shares this traversal with dependencies.
+- Conservative producer tracing visits each value/projection once per input,
+  including captured region operands, with linear time and space in that input's
+  projection DAG. Direct physical-batch edges retain their existing fast path.
+- Pipeline mode removes a complete PEFT run, including its reservation and
+  insertion-gap searches. Latency mode retains those algorithms and their
+  existing potentially quadratic gap traversal.
+
+Remaining approximation boundaries include operation issue overhead, eliminated
+assembly copies, communication contention and pipeline synchronization. Weight
+indexing and ready-transfer signature searches remain possible compiler
+hotspots. Changes to these require matched execution measurements.

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ResidentWeightPacking.hpp

@@ -0,0 +1,126 @@
+#pragma once
+
+#include "ComputeGraph.hpp"
+#include "llvm/ADT/SmallBitVector.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include <numeric>
+#include <optional>
+#include <tuple>
+
+namespace onnx_mlir::spatial {
+
+struct IndexedResidentWeights {
+  std::vector<std::vector<size_t>> taskWeightIds;
+  size_t weightCount = 0;
+
+  explicit IndexedResidentWeights(const ComputeGraph &graph) {
+    taskWeightIds.resize(graph.nodes.size());
+    ResidentWeightSet indexed;
+    for (auto [task, node] : llvm::enumerate(graph.nodes))
+      for (const ResidentWeight &weight : node.residentWeights) {
+        auto found = llvm::find(indexed, weight);
+        size_t id = found - indexed.begin();
+        if (found == indexed.end())
+          indexed.push_back(weight);
+        taskWeightIds[task].push_back(id);
+      }
+    weightCount = indexed.size();
+  }
+
+  size_t getPackCountLowerBound(llvm::ArrayRef<size_t> tasks,
+                                size_t capacity) const {
+    assert(capacity > 0);
+    llvm::SmallBitVector distinct(weightCount);
+    for (size_t task : tasks) {
+      if (taskWeightIds[task].size() > capacity)
+        return std::numeric_limits<size_t>::max();
+      for (size_t weight : taskWeightIds[task])
+        distinct.set(weight);
+    }
+    size_t count = distinct.count();
+    return std::max<size_t>(!tasks.empty(), count / capacity + (count % capacity != 0));
+  }
+
+  // A constructive reuse-first best fit; failure is not a proof of infeasibility.
+  std::optional<std::vector<size_t>> pack(
+      llvm::ArrayRef<size_t> tasks, llvm::ArrayRef<Cost> costs,
+      size_t capacity, size_t packCount) const {
+    if (getPackCountLowerBound(tasks, capacity) > packCount)
+      return std::nullopt;
+    std::vector<size_t> order(tasks.size());
+    std::iota(order.begin(), order.end(), 0);
+    llvm::stable_sort(order, [&](size_t lhs, size_t rhs) {
+      size_t left = tasks[lhs], right = tasks[rhs];
+      if (taskWeightIds[left].size() != taskWeightIds[right].size())
+        return taskWeightIds[left].size() > taskWeightIds[right].size();
+      return costs[left] > costs[right];
+    });
+    std::vector<llvm::SmallBitVector> weights(
+      packCount, llvm::SmallBitVector(weightCount));
+    std::vector<size_t> weightCounts(packCount);
+    std::vector<Cost> loads(packCount);
+    std::vector<size_t> assignment(tasks.size());
+    constexpr size_t noDeviation = std::numeric_limits<size_t>::max();
+    size_t repairChecks = 65536;
+    auto construct = [&](size_t forbiddenPosition = std::numeric_limits<size_t>::max(),
+                         size_t forbiddenPack = 0) -> std::optional<std::vector<size_t>> {
+      for (auto &resident : weights)
+        resident.reset();
+      std::fill(weightCounts.begin(), weightCounts.end(), 0);
+      std::fill(loads.begin(), loads.end(), 0);
+      std::fill(assignment.begin(), assignment.end(), packCount);
+      for (size_t position : order) {
+        size_t task = tasks[position];
+        const auto &ids = taskWeightIds[task];
+        std::optional<std::tuple<size_t, size_t, Cost, size_t>> best;
+        for (size_t pack = 0; pack < packCount; ++pack) {
+          if (forbiddenPosition != noDeviation) {
+            if (repairChecks == 0)
+              return std::nullopt;
+            --repairChecks;
+            if (position == forbiddenPosition && pack == forbiddenPack)
+              continue;
+          }
+          size_t added = 0;
+          for (size_t id : ids)
+            added += !weights[pack].test(id);
+          size_t unionSize = weightCounts[pack] + added;
+          if (unionSize > capacity)
+            continue;
+          auto score = std::make_tuple(added, ids.empty() ? 0 : capacity - unionSize,
+                                       loads[pack], pack);
+          if (!best || score < *best)
+            best = score;
+        }
+        if (!best)
+          return std::nullopt;
+        size_t pack = std::get<3>(*best);
+        assignment[position] = pack;
+        weightCounts[pack] += std::get<0>(*best);
+        loads[pack] = checkedAdd(loads[pack], costs[task]);
+        for (size_t id : ids)
+          weights[pack].set(id);
+      }
+      return std::move(assignment);
+    };
+    if (auto assignment = construct())
+      return assignment;
+    // Reverse the greedy choices once, only after failure; neither order proves infeasibility.
+    std::reverse(order.begin(), order.end());
+    if (auto assignment = construct())
+      return assignment;
+    // Replay with one next-best choice: at most16 deviations and65536 core checks.
+    const auto greedyAssignment = assignment;
+    for (size_t index = 0; index < std::min<size_t>(16, order.size()) && repairChecks; ++index) {
+      size_t position = order[index];
+      if (greedyAssignment[position] == packCount)
+        break;
+      if (auto repaired = construct(position, greedyAssignment[position]))
+        return repaired;
+    }
+    return std::nullopt;
+  }
+};
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.cpp

@@ -0,0 +1,111 @@
+#include "ScheduleEvaluation.hpp"
+#include "MergeSchedule.hpp"
+
+#include <queue>
+
+namespace onnx_mlir::spatial {
+
+void recordScheduleCosts(const ScheduleEvaluation &evaluation,
+                         MergeScheduleResult &schedule) {
+  schedule.costs.totalWorkNs = evaluation.totalWork;
+  schedule.costs.maximumCoreWorkNs = evaluation.coreWork.empty() ? 0
+    : *std::max_element(evaluation.coreWork.begin(), evaluation.coreWork.end());
+  schedule.costs.criticalPathNs = evaluation.criticalPath;
+  schedule.costs.communicationNs = evaluation.communication;
+  schedule.costs.completionNs = evaluation.completion;
+  schedule.costs.stageSpanNs.assign(evaluation.stageSpans.begin(), evaluation.stageSpans.end());
+}
+
+mlir::FailureOr<ScheduleEvaluation> evaluateSchedule(
+    const ComputeGraph &graph,
+    const std::vector<std::vector<size_t>> &tasksByCore,
+    const SchedulingTarget &target, llvm::ArrayRef<size_t> processorStages,
+    std::string &error) {
+  size_t count = graph.nodes.size();
+  if (tasksByCore.size() != target.processorCount
+      || (!processorStages.empty() && processorStages.size() != target.processorCount)) {
+    error = "schedule evaluation received an incompatible core layout";
+    return mlir::failure();
+  }
+  std::vector<size_t> cores(count, target.processorCount);
+  std::vector<size_t> indegree(count);
+  std::vector<std::vector<size_t>> successors(count);
+  for (auto [core, tasks] : llvm::enumerate(tasksByCore))
+    for (auto [slot, task] : llvm::enumerate(tasks)) {
+      if (task >= count || cores[task] != target.processorCount) {
+        error = "schedule evaluation requires exactly one core per task";
+        return mlir::failure();
+      }
+      cores[task] = core;
+      if (slot) {
+        successors[tasks[slot - 1]].push_back(task);
+        ++indegree[task];
+      }
+    }
+  if (llvm::is_contained(cores, target.processorCount)) {
+    error = "schedule evaluation received an incomplete assignment";
+    return mlir::failure();
+  }
+  for (size_t task = 0; task < count; ++task)
+    for (const auto &[predecessor, transfer] : graph.predecessors[task]) {
+      successors[predecessor].push_back(task);
+      ++indegree[task];
+    }
+  auto laterOrder = [&](size_t lhs, size_t rhs) {
+    return graph.nodes[lhs].originalOrder > graph.nodes[rhs].originalOrder;
+  };
+  std::priority_queue<size_t, std::vector<size_t>, decltype(laterOrder)> ready(laterOrder);
+  for (size_t task = 0; task < count; ++task)
+    if (indegree[task] == 0)
+      ready.push(task);
+  ScheduleEvaluation result;
+  result.taskStarts.resize(count);
+  result.taskEnds.resize(count);
+  result.coreWork.resize(target.processorCount);
+  size_t stageCount = processorStages.empty() ? 1
+    : *std::max_element(processorStages.begin(), processorStages.end()) + 1;
+  result.stageSpans.resize(stageCount);
+  std::vector<Time> coreReady(target.processorCount);
+  std::vector<Time> stageCoreReady(target.processorCount);
+  std::vector<Time> stageEnds(count);
+  std::vector<Time> pathEnds(count);
+  size_t visited = 0;
+  while (!ready.empty()) {
+    size_t task = ready.top();
+    ready.pop();
+    size_t core = cores[task];
+    size_t stage = processorStages.empty() ? 0 : processorStages[core];
+    Time start = coreReady[core];
+    Time stageStart = stageCoreReady[core];
+    Time pathStart = 0;
+    for (const auto &[predecessor, transfer] : graph.predecessors[task]) {
+      Time communication = getTransferTime(transfer, cores[predecessor], core, target);
+      result.communication = checkedAdd(result.communication, communication);
+      start = std::max(start, checkedAdd(result.taskEnds[predecessor], communication));
+      pathStart = std::max(pathStart, checkedAdd(pathEnds[predecessor], communication));
+      if (processorStages.empty() || processorStages[cores[predecessor]] == stage)
+        stageStart = std::max(stageStart, checkedAdd(stageEnds[predecessor], communication));
+    }
+    Time cost = graph.nodes[task].cost;
+    result.totalWork = checkedAdd(result.totalWork, cost);
+    result.coreWork[core] = checkedAdd(result.coreWork[core], cost);
+    result.taskStarts[task] = start;
+    result.taskEnds[task] = coreReady[core] = checkedAdd(start, cost);
+    stageEnds[task] = stageCoreReady[core] = checkedAdd(stageStart, cost);
+    pathEnds[task] = checkedAdd(pathStart, cost);
+    result.stageSpans[stage] = std::max(result.stageSpans[stage], stageEnds[task]);
+    result.criticalPath = std::max(result.criticalPath, pathEnds[task]);
+    result.completion = std::max(result.completion, result.taskEnds[task]);
+    ++visited;
+    for (size_t successor : successors[task])
+      if (--indegree[successor] == 0)
+        ready.push(successor);
+  }
+  if (visited != count) {
+    error = "schedule evaluation found a cyclic graph or per-core order";
+    return mlir::failure();
+  }
+  return result;
+}
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp

@@ -0,0 +1,34 @@
+#pragma once
+
+#include "ComputeGraph.hpp"
+#include "mlir/Support/LogicalResult.h"
+
+#include <string>
+
+namespace onnx_mlir::spatial {
+
+struct MergeScheduleResult;
+
+struct ScheduleEvaluation {
+  std::vector<Time> taskStarts;
+  std::vector<Time> taskEnds;
+  std::vector<Time> coreWork;
+  std::vector<Time> stageSpans;
+  Time totalWork = 0;
+  Time criticalPath = 0;
+  Time completion = 0;
+  Time communication = 0;
+};
+
+// Nanosecond estimates with serial core execution and dependency arrival times.
+// Stage spans exclude cross-stage routing, synchronization and network contention.
+mlir::FailureOr<ScheduleEvaluation> evaluateSchedule(
+    const ComputeGraph &graph,
+    const std::vector<std::vector<size_t>> &tasksByCore,
+    const SchedulingTarget &target, llvm::ArrayRef<size_t> processorStages,
+    std::string &error);
+
+void recordScheduleCosts(const ScheduleEvaluation &evaluation,
+                         MergeScheduleResult &schedule);
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp

@@ -26,6 +26,7 @@ struct SchedulingTarget {
   size_t synchronizationRegisterCount = 0;
 
   Cost matrixRows = 128;
+  Cost matrixInputReadMultiplier = 1;
   Cost matrixColumns = 128;
   Cost matrixPeriodNs = 1;
   Cost matrixInputResolutionBits = 1;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/Utils.hpp

@@ -15,6 +15,8 @@
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using CPU = int;
+// Scheduling costs and times are integer nanoseconds, except explicitly named
+// resource quantities such as bytes, flits, weights, and hardware cycles.
 using Cost = unsigned long long;
 using Time = unsigned long long;