better throughput in pipeline mode

src/PIM/Compiler/PimCodeGen.cpp

@@ -544,11 +544,23 @@ void PimCodeGen::setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address
   genSetRegisterImmediateUnsigned(1, pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address"));
 }
 
-void PimCodeGen::setupRdRs1Rs2(
+std::array<uint8_t, 3> PimCodeGen::setupRdRs1Rs2(
   size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const {
-  genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
-  genSetRegisterImmediateUnsigned(1, pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address"));
-  genSetRegisterImmediateUnsigned(2, pim::checkedAddOrCrash(rs2Address, rs2Offset, "rs2 address"));
+  size_t rd = pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address");
+  size_t rs1 = pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address");
+  size_t rs2 = pim::checkedAddOrCrash(rs2Address, rs2Offset, "rs2 address");
+  genSetRegisterImmediateUnsigned(0, rd);
+  uint8_t rs1Register = 0;
+  if (rd != rs1) {
+    genSetRegisterImmediateUnsigned(1, rs1);
+    rs1Register = 1;
+  }
+  if (rd == rs2)
+    return {0, rs1Register, 0};
+  if (rs1 == rs2)
+    return {0, rs1Register, rs1Register};
+  genSetRegisterImmediateUnsigned(2, rs2);
+  return {0, rs1Register, 2};
 }
 
 void PimCodeGen::emitMemCopyOp(pim_binary::Opcode opcode,
@@ -664,13 +676,13 @@ void PimCodeGen::codeGenVMVOp(pim::PimVMVOp vmvOp, const StaticValueKnowledge& k
   auto sourceType = cast<ShapedType>(vmvOp.getSource().getType());
   int32_t bitwidth = getVectorElementBitwidthOrCrash(sourceType);
   ensureVectorBitwidth(bitwidth, bitwidth);
-  setupRdRs1Rs2(addressOf(vmvOp.getTarget(), knowledge), *targetOffset,
-                addressOf(vmvOp.getSource(), knowledge), *sourceOffset, 0, *sourceStride);
+  auto registers = setupRdRs1Rs2(addressOf(vmvOp.getTarget(), knowledge), *targetOffset,
+                                 addressOf(vmvOp.getSource(), knowledge), *sourceOffset, 0, *sourceStride);
   pim_binary::InstructionRecord instruction;
   instruction.opcode = pim_binary::Opcode::vmv;
-  instruction.rd = 0;
-  instruction.r1 = 1;
-  instruction.r2OrImm = 2;
+  instruction.rd = registers[0];
+  instruction.r1 = registers[1];
+  instruction.r2OrImm = registers[2];
   instruction.generic3 = vmvOp.getLength();
   emitInstruction(instruction);
 }
@@ -780,12 +792,13 @@ void PimCodeGen::emitBinaryVectorOp(pim_binary::Opcode opcode,
   auto inputType = cast<ShapedType>(lhs.getType());
   ensureVectorBitwidth(getVectorElementBitwidthOrCrash(inputType),
                        getVectorElementBitwidthOrCrash(cast<ShapedType>(output.getType())));
-  setupRdRs1Rs2(addressOf(output, knowledge), 0, addressOf(lhs, knowledge), 0, addressOf(rhs, knowledge), 0);
+  auto registers = setupRdRs1Rs2(
+    addressOf(output, knowledge), 0, addressOf(lhs, knowledge), 0, addressOf(rhs, knowledge), 0);
   pim_binary::InstructionRecord instruction;
   instruction.opcode = opcode;
-  instruction.rd = 0;
-  instruction.r1 = 1;
-  instruction.r2OrImm = 2;
+  instruction.rd = registers[0];
+  instruction.r1 = registers[1];
+  instruction.r2OrImm = registers[2];
   instruction.generic3 = getVectorElementCountOrCrash(inputType);
   emitInstruction(instruction);
 }

src/PIM/Compiler/PimCodeGen.hpp

@@ -176,7 +176,7 @@ class PimCodeGen {
   void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const;
   void setupRd(size_t rdAddress, size_t rdOffset) const;
   void setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset) const;
-  void setupRdRs1Rs2(
+  std::array<uint8_t, 3> setupRdRs1Rs2(
     size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const;
 
   void emitMemCopyOp(pim_binary::Opcode opcode,

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -140,8 +140,8 @@ void verifyExplicitPimCoreCount() {
 void verifyPimPipelineStages() {
   if (pipelineStages.getValue() == 0)
     llvm::report_fatal_error("PIM compilation requires --pipeline to be positive");
-  if (static_cast<size_t>(coresCount.getValue()) % pipelineStages.getValue() != 0)
-    llvm::report_fatal_error("PIM compilation requires --core-count to be divisible by --pipeline");
+  if (static_cast<size_t>(coresCount.getValue()) < pipelineStages.getValue())
+    llvm::report_fatal_error("PIM compilation requires --pipeline not to exceed --core-count");
   if (crossbarCountInCore.getValue()
       > std::numeric_limits<size_t>::max() / pipelineStages.getValue())
     llvm::report_fatal_error("PIM compilation --crossbar-count * --pipeline overflows");

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp

@@ -370,6 +370,14 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {
     Location loc = reduceMeanOp.getLoc();
     RankedTensorType leafType = getAllOnesType(inputType, resultType.getElementType());
     RankedTensorType keepdimsType = getKeepdimsType(inputType, resultType.getElementType(), reducedAxes);
+    if (semantics->keepdims != 0 && inputType.getRank() == 4
+        && inputType.getDimSize(0) == 1 && semantics->axes == ArrayRef<int64_t>({2, 3})
+        && resultType == keepdimsType) {
+      auto plan = spatial::SpatGlobalAveragePoolPlanOp::create(
+        rewriter, loc, resultType, adaptor.getData(), spatial::getNCHWLayout(rewriter.getContext()));
+      rewriter.replaceOp(reduceMeanOp, plan.getResult());
+      return success();
+    }
     int64_t laneCount = 1;
     for (auto [dim, isReduced] : llvm::zip_equal(keepdimsType.getShape(), reducedAxes)) {
       if (isReduced)

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

@@ -34,7 +34,9 @@ static LogicalResult verifyNoEscapingRegionValues(Operation* owner, StringRef ph
                               << escapingUser->getName() << " at " << escapingUser->getLoc();
 }
 
-static LogicalResult placeLogicalProcessorsOnPhysicalCores(DeferredTransferPlan& plan, const SchedulingTarget& target) {
+static LogicalResult placeLogicalProcessorsOnPhysicalCores(
+    DeferredTransferPlan& plan, const SchedulingTarget& target,
+    size_t pipelineStages) {
   std::vector<Cost> logicalTrafficFlits(target.processorCount * target.processorCount, 0);
   for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges)
     for (const ExternalTransferFamily& transfer : exchange->external) {
@@ -55,8 +57,15 @@ static LogicalResult placeLogicalProcessorsOnPhysicalCores(DeferredTransferPlan&
       }
     }
 
+  std::vector<size_t> placementGroups;
+  if (pipelineStages > 1) {
+    if (plan.processorStages.size() != target.processorCount)
+      return failure();
+    placementGroups = plan.processorStages;
+  }
   std::vector<size_t> physicalCoreForLogicalProcessor =
-    mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, target);
+    mapLogicalProcessorsToPhysicalCores(
+      logicalTrafficFlits, target, placementGroups);
   auto getPhysicalCore = [&](int64_t logicalProcessor) {
     assert(logicalProcessor >= 0 && static_cast<size_t>(logicalProcessor) < physicalCoreForLogicalProcessor.size()
            && "logical processor is outside the scheduling target");
@@ -218,7 +227,8 @@ LogicalResult realizeDeferredCommunication(func::FuncOp funcOp,
     funcOp, materialization, pipelineStages, target.processorCount);
   if (failed(transfers))
     return funcOp.emitOpError("phase 2 failed to build symbolic transfer families");
-  if (failed(placeLogicalProcessorsOnPhysicalCores(*transfers, target)))
+  if (failed(placeLogicalProcessorsOnPhysicalCores(
+        *transfers, target, pipelineStages)))
     return failure();
   if (transfers->pipelineHostBufferBytes != 0) {
     auto bytes = pim::checkedCast<int64_t>(

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

@@ -32,9 +32,11 @@ static LogicalResult collectScheduledOperations(
     DeferredTransferPlan &plan,
     size_t pipelineStageCount,
     size_t processorCount) {
-  if (pipelineStageCount == 0 || processorCount % pipelineStageCount != 0)
+  if (pipelineStageCount == 0
+      || (pipelineStageCount > 1
+          && materialization.processorStages.size() != processorCount))
     return failure();
-  size_t stageSize = processorCount / pipelineStageCount;
+  plan.processorStages = materialization.processorStages;
   unsigned nextStream = 0;
   for (const ScheduledMaterializationRecord &record :
        materialization.materializedSchedules) {
@@ -56,8 +58,12 @@ static LogicalResult collectScheduledOperations(
       if (core >= processorCount)
         return op.emitOpError("phase 2 scheduled core is outside the target");
       info.cores.push_back(core);
-      if (pipelineStageCount > 1)
-        info.pipelineStages.push_back(core / stageSize);
+      if (pipelineStageCount > 1) {
+        size_t stage = materialization.processorStages[core];
+        if (stage >= pipelineStageCount)
+          return op.emitOpError("phase 2 scheduled core has an invalid pipeline stage");
+        info.pipelineStages.push_back(stage);
+      }
     }
     for (size_t lane = 0; lane < info.cores.size(); ++lane)
       info.streamIds.push_back(nextStream++);

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

@@ -8,6 +8,7 @@
 namespace onnx_mlir::spatial {
 
 struct DeferredTransferPlan {
+  std::vector<size_t> processorStages;
   llvm::SmallVector<ScheduledInfo, 0> scheduled;
   llvm::SmallVector<std::unique_ptr<ProducedValue>> producedStorage;
   llvm::DenseMap<int64_t, llvm::SmallVector<ProducedValue*>> producedByGraph;

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

@@ -808,7 +808,9 @@ materializeScheduledCompute(func::FuncOp funcOp,
     }
   }
 
-  return ScheduledComputeMaterializationResult {std::move(peftClassPlans), std::move(materializedSchedules), std::move(graphComputeToBlockMap)};
+  return ScheduledComputeMaterializationResult {
+    std::move(peftClassPlans), std::move(materializedSchedules),
+    std::move(graphComputeToBlockMap), schedule.processorStages};
 }
 
 

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

@@ -14,6 +14,7 @@ struct ScheduledComputeMaterializationResult {
   llvm::MapVector<size_t, PeftClassPlan> peftClassPlans;
   std::vector<ScheduledMaterializationRecord> materializedSchedules;
   DenseMap<GraphComputeBlockKey, Block *> graphComputeToBlockMap;
+  std::vector<size_t> processorStages;
 };
 
 FailureOr<BatchFragmentSpec>

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

@@ -39,8 +39,9 @@ static SchedulingTarget getPipelineSchedulingTarget(
   if (pipelineStages == 1)
     return physicalTarget;
 
+  PipelineCoreLayout layout(physicalTarget.processorCount, pipelineStages);
   SchedulingTarget schedulingTarget = physicalTarget;
-  schedulingTarget.processorCount = physicalTarget.processorCount / pipelineStages;
+  schedulingTarget.processorCount = layout.getLogicalProcessorCount();
   schedulingTarget.residentWeightCapacity = checkedMultiply(
     physicalTarget.residentWeightCapacity, pipelineStages);
   schedulingTarget.interProcessorLatencyNs.assign(
@@ -88,7 +89,8 @@ struct ScheduleAndRealizeSpatialPass final
       signalPassFailure();
       return;
     }
-    if (pipelineStages == 0 || target.processorCount % pipelineStages != 0
+    PipelineCoreLayout pipelineLayout(target.processorCount, pipelineStages);
+    if (!pipelineLayout.isValid()
         || (pipelineStages > 1
             && target.synchronizationRegisterCount == 0)
         || target.residentWeightCapacity
@@ -115,19 +117,24 @@ struct ScheduleAndRealizeSpatialPass final
       scheduledGraph = analysis.getGraph();
       schedule = std::move(analysis.getResult());
       std::string pipelineError;
+      if (pipelineStages > 1) {
+        FailureOr<PipelineWorkloadPreparation> preparation =
+          preparePipelineWorkload(
+            scheduledGraph, schedule, pipelineStages, target, pipelineError);
+        if (failed(preparation)) {
+          moduleOp.emitError() << pipelineError;
+          signalPassFailure();
+          return;
+        }
+        if (*preparation == PipelineWorkloadPreparation::Changed)
+          continue;
+      }
       if (succeeded(applyPipelineScheduling(
             scheduledGraph, schedule, pipelineStages, target, pipelineError)))
         break;
-      std::string splitError;
-      if (pipelineStages == 1
-          || failed(splitPipelineWorkload(
-                scheduledGraph, schedule, pipelineStages, target, splitError))) {
-        if (!splitError.empty())
-          pipelineError = splitError;
-        moduleOp.emitError() << pipelineError;
-        signalPassFailure();
-        return;
-      }
+      moduleOp.emitError() << pipelineError;
+      signalPassFailure();
+      return;
     }
     PatternRewriter rewriter(moduleOp.getContext());
     FailureOr<ScheduledComputeMaterializationResult> materialization =

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

@@ -14,6 +14,7 @@ namespace spatial {
 
 struct MergeScheduleResult {
   size_t processorCount = 0;
+  std::vector<size_t> processorStages;
   std::vector<ComputeInstance> dominanceOrderCompute;
   llvm::DenseMap<ComputeInstance, size_t> computeToCpuMap;
   llvm::DenseMap<ComputeInstance, size_t> computeToCpuSlotMap;

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

@@ -244,11 +244,13 @@ FailureOr<LanePublicationSignatures> buildLanePublicationSignatures(SpatComputeB
 } // namespace
 
 std::vector<size_t> mapLogicalProcessorsToPhysicalCores(ArrayRef<Cost> logicalTrafficFlits,
-                                                        const SchedulingTarget& target) {
+                                                        const SchedulingTarget& target,
+                                                        ArrayRef<size_t> placementGroups) {
   const size_t processorCount = target.processorCount;
   assert(logicalTrafficFlits.size() == processorCount * processorCount
          && "logical traffic matrix must cover every processor pair");
-
+  assert((placementGroups.empty() || placementGroups.size() == processorCount)
+         && "physical placement groups must cover every processor");
   std::vector<size_t> physicalCoreForLogicalProcessor(processorCount);
   std::iota(physicalCoreForLogicalProcessor.begin(), physicalCoreForLogicalProcessor.end(), 0);
 
@@ -266,6 +268,10 @@ std::vector<size_t> mapLogicalProcessorsToPhysicalCores(ArrayRef<Cost> logicalTr
     for (size_t peerLogicalProcessor = 0; peerLogicalProcessor < processorCount; ++peerLogicalProcessor) {
       if (peerLogicalProcessor == logicalProcessor)
         continue;
+      if (!placementGroups.empty()
+          && placementGroups[peerLogicalProcessor]
+               != placementGroups[logicalProcessor])
+        continue;
       size_t physicalCore = physicalCoreForLogicalProcessor[logicalProcessor];
       size_t peerPhysicalCore = physicalCoreForLogicalProcessor[peerLogicalProcessor];
       Cost currentCost = 0;

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

@@ -29,7 +29,8 @@ inline Time getPeftTransferTime(const TransferCost& transferCost,
 MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options);
 
 std::vector<size_t> mapLogicalProcessorsToPhysicalCores(llvm::ArrayRef<Cost> logicalTrafficFlits,
-                                                        const SchedulingTarget& target);
+                                                        const SchedulingTarget& target,
+                                                        llvm::ArrayRef<size_t> placementGroups = {});
 
 } // namespace spatial
 } // namespace onnx_mlir

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

@@ -1,7 +1,10 @@
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/STLExtras.h"
 
 #include <algorithm>
 #include <limits>
+#include <numeric>
 #include <optional>
 #include <queue>
 #include <tuple>
@@ -556,6 +559,10 @@ static LogicalResult splitBatchCompute(SpatGraphComputeBatch batch,
     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 =
@@ -601,11 +608,10 @@ static LogicalResult splitBatchCompute(SpatGraphComputeBatch batch,
   return success();
 }
 
-static LogicalResult splitPipelineWorkloadImpl(const ComputeGraph &graph,
-                                               const MergeScheduleResult &schedule,
-                                               size_t pipelineStages,
-                                               const SchedulingTarget &physicalTarget,
-                                               std::string &error) {
+static FailureOr<PipelineWorkloadPreparation> preparePipelineWorkloadImpl(
+    const ComputeGraph &graph, const MergeScheduleResult &schedule,
+    size_t pipelineStages, const SchedulingTarget &physicalTarget,
+    std::string &error) {
   size_t groupSize = schedule.processorCount;
   std::vector<TaskList> tasksByCpu(groupSize);
   for (size_t task = 0; task < graph.nodes.size(); ++task) {
@@ -654,7 +660,7 @@ static LogicalResult splitPipelineWorkloadImpl(const ComputeGraph &graph,
       std::string currentError;
       if (succeeded(splitBatchCompute(
             batch, pipelineStages, physicalTarget, currentError)))
-        return success();
+        return PipelineWorkloadPreparation::Changed;
       if (!currentError.empty())
         candidateError = currentError;
     }
@@ -663,7 +669,8 @@ static LogicalResult splitPipelineWorkloadImpl(const ComputeGraph &graph,
       : candidateError;
     return failure();
   }
-  return failure();
+
+  return PipelineWorkloadPreparation::Ready;
 }
 
 bool fits(const ComputeGraph& graph,
@@ -719,193 +726,585 @@ Cost findMaximumPackCost(const ComputeGraph& graph,
   return low;
 }
 
-static Cost getCoreCost(const TaskList &tasks, const TaskCosts &taskCosts) {
-  Cost cost = 0;
-  for (size_t task : tasks)
-    cost = checkedAdd(cost, taskCosts[task]);
-  return cost;
-}
-
-static Cost getStageMaximumAssemblyCost(
-    const std::vector<TaskList> &tasksByCpu,
-    const TaskCosts &assemblyCosts, size_t groupSize, size_t stage) {
-  Cost maximum = 0;
-  for (size_t cpu = stage * groupSize;
-      cpu < (stage + 1) * groupSize; ++cpu)
-    maximum = std::max(
-      maximum, getCoreCost(tasksByCpu[cpu], assemblyCosts));
-  return maximum;
-}
-
-static bool fitsResidentWeights(const ComputeGraph &graph,
-                                const TaskList &tasks, size_t candidate,
-                                size_t residentWeightCapacity) {
-  ResidentWeightSet weights;
-  for (size_t task : tasks)
-    insertResidentWeights(weights, graph.nodes[task].residentWeights);
-  return getResidentWeightUnionSize(
-           weights, graph.nodes[candidate].residentWeights)
-       <= residentWeightCapacity;
-}
-
-static void repackPipelineStage(
-    const ComputeGraph &graph, const TaskCosts &schedulingCosts,
-    const TaskCosts &assemblyCosts,
-    std::vector<TaskList> &tasksByCpu, size_t groupSize, size_t stage,
-    size_t residentWeightCapacity) {
-  TaskList tasks;
-  Cost originalMaximum = 0;
-  for (size_t cpu = stage * groupSize;
-       cpu < (stage + 1) * groupSize; ++cpu) {
-    llvm::append_range(tasks, tasksByCpu[cpu]);
-    originalMaximum = std::max(
-      originalMaximum,
-      getCoreCost(tasksByCpu[cpu], schedulingCosts));
+static Cost findMaximumIndexedPackCost(
+    const TaskCosts &taskCosts,
+    const std::vector<TaskList> &taskWeightIds, size_t weightCount,
+    const TaskList &tasks, size_t residentWeightCapacity,
+    size_t maximumPacks) {
+  Cost low = 0;
+  Cost high = 0;
+  for (size_t task : tasks) {
+    low = std::max(low, taskCosts[task]);
+    high = checkedAdd(high, taskCosts[task]);
   }
-  llvm::sort(tasks, [&](size_t lhs, size_t rhs) {
+  low = std::max(
+    low, high / maximumPacks + (high % maximumPacks != 0));
+
+  std::vector<size_t> seen(weightCount);
+  size_t generation = 0;
+  while (low < high) {
+    Cost middle = low + (high - low) / 2;
+    size_t packs = 1;
+    Cost cost = 0;
+    size_t packWeightCount = 0;
+    ++generation;
+    bool fits = true;
+    bool packEmpty = true;
+    for (size_t task : tasks) {
+      size_t addedWeights = 0;
+      for (size_t weight : taskWeightIds[task])
+        addedWeights += seen[weight] != generation;
+      Cost taskCost = taskCosts[task];
+      bool startsNewPack = !packEmpty
+        && (cost > middle - taskCost
+            || packWeightCount + addedWeights > residentWeightCapacity);
+      if (startsNewPack) {
+        if (++packs > maximumPacks) {
+          fits = false;
+          break;
+        }
+        cost = 0;
+        packWeightCount = 0;
+        ++generation;
+        packEmpty = true;
+      }
+      cost = checkedAdd(cost, taskCost);
+      for (size_t weight : taskWeightIds[task])
+        if (seen[weight] != generation) {
+          seen[weight] = generation;
+          ++packWeightCount;
+        }
+      packEmpty = false;
+    }
+    if (fits)
+      high = middle;
+    else
+      low = middle + 1;
+  }
+  return low;
+}
+
+static size_t findMinimumIndexedPackCount(
+    const std::vector<TaskList> &taskWeightIds, const TaskList &tasks,
+    size_t residentWeightCapacity) {
+  if (tasks.empty())
+    return 0;
+  size_t packs = 1;
+  TaskList weights;
+  for (size_t task : tasks) {
+    size_t unionSize = weights.size();
+    for (size_t weight : taskWeightIds[task])
+      unionSize += !llvm::is_contained(weights, weight);
+    if (!weights.empty() && unionSize > residentWeightCapacity) {
+      ++packs;
+      weights.clear();
+    }
+    for (size_t weight : taskWeightIds[task])
+      if (!llvm::is_contained(weights, weight))
+        weights.push_back(weight);
+  }
+  return packs;
+}
+
+struct PipelineGroup {
+  TaskList tasks;
+  TaskList successors;
+  size_t originalOrder = std::numeric_limits<size_t>::max();
+  TaskList weightIds;
+  bool consumesPipelineInput = false;
+};
+
+struct PipelineStageAssignment {
+  std::vector<size_t> taskStages;
+  std::vector<size_t> stageSizes;
+};
+
+static const TaskCosts &getPipelineBalanceCosts(
+    const ComputeGraph &graph, const PipelineTaskModel &model) {
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    if (graph.nodes[task].instance.op && model.assemblyCosts[task] > 1
+        && !model.predecessors[task].empty())
+      return model.assemblyCosts;
+  return model.schedulingCosts;
+}
+
+static bool consumesPipelineInput(const ComputeGraphNode &node) {
+  if (!node.instance.op)
+    return false;
+  return llvm::any_of(getComputeInstanceInputs(node.instance),
+                      [](Value input) { return isa<BlockArgument>(input); });
+}
+
+static FailureOr<PipelineStageAssignment> assignPipelineStages(
+    const ComputeGraph &graph, const PipelineTaskModel &model,
+    const PipelineCoreLayout &layout, size_t residentWeightCapacity,
+    std::string &error) {
+  if (graph.nodes.empty())
+    return PipelineStageAssignment {
+      {}, std::vector<size_t>(layout.getStageCount(), 1)};
+  std::vector<size_t> tasksByOrder(graph.nodes.size());
+  std::iota(tasksByOrder.begin(), tasksByOrder.end(), 0);
+  llvm::sort(tasksByOrder, [&](size_t lhs, size_t rhs) {
     return graph.nodes[lhs].originalOrder < graph.nodes[rhs].originalOrder;
   });
 
-  std::vector<TaskList> packed(groupSize);
-  std::vector<ResidentWeightSet> weights(groupSize);
-  TaskCosts loads(groupSize);
-  TaskCosts assemblyLoads(groupSize);
+  std::vector<PipelineGroup> groups;
+  std::vector<size_t> taskToGroup(graph.nodes.size());
+  ResidentWeightSet indexedWeights;
+  std::vector<TaskList> taskWeightIds(graph.nodes.size());
+  for (size_t task : tasksByOrder)
+    for (const ResidentWeight &weight : graph.nodes[task].residentWeights) {
+      auto indexed = llvm::find(indexedWeights, weight);
+      size_t id = indexed - indexedWeights.begin();
+      if (indexed == indexedWeights.end()) {
+        id = indexedWeights.size();
+        indexedWeights.push_back(weight);
+      }
+      taskWeightIds[task].push_back(id);
+    }
+  llvm::DenseMap<Operation *, size_t> operationToGroup;
+  llvm::DenseMap<Attribute, size_t> splitOperationToGroup;
+  for (size_t task : tasksByOrder) {
+    Operation *operation = graph.nodes[task].instance.op;
+    Attribute splitGroup = operation
+      ? operation->getAttr("pipeline.stage_group")
+      : Attribute();
+    size_t group;
+    auto existingSplit = splitGroup
+      ? splitOperationToGroup.find(splitGroup)
+      : splitOperationToGroup.end();
+    auto existingOperation = operation && !splitGroup
+      ? operationToGroup.find(operation)
+      : operationToGroup.end();
+    if (existingSplit != splitOperationToGroup.end()) {
+      group = existingSplit->second;
+    } else if (existingOperation != operationToGroup.end()) {
+      group = existingOperation->second;
+    } else {
+      group = groups.size();
+      groups.emplace_back();
+      if (splitGroup)
+        splitOperationToGroup[splitGroup] = group;
+      else if (operation)
+        operationToGroup[operation] = group;
+    }
+    taskToGroup[task] = group;
+    PipelineGroup &pipelineGroup = groups[group];
+    pipelineGroup.tasks.push_back(task);
+    pipelineGroup.originalOrder = std::min(
+      pipelineGroup.originalOrder, graph.nodes[task].originalOrder);
+    pipelineGroup.consumesPipelineInput |=
+      consumesPipelineInput(graph.nodes[task]);
+    for (size_t weight : taskWeightIds[task])
+      if (!llvm::is_contained(pipelineGroup.weightIds, weight))
+        pipelineGroup.weightIds.push_back(weight);
+  }
+
+  std::vector<size_t> indegree(groups.size());
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    for (size_t predecessor : model.predecessors[task]) {
+      size_t source = taskToGroup[predecessor];
+      size_t target = taskToGroup[task];
+      if (source == target
+          || llvm::is_contained(groups[source].successors, target))
+        continue;
+      groups[source].successors.push_back(target);
+      ++indegree[target];
+    }
+
+  auto laterOriginalOrder = [&](size_t lhs, size_t rhs) {
+    return groups[lhs].originalOrder > groups[rhs].originalOrder;
+  };
+  std::priority_queue<size_t, std::vector<size_t>, decltype(laterOriginalOrder)>
+    ready(laterOriginalOrder);
+  for (size_t group = 0; group < groups.size(); ++group)
+    if (indegree[group] == 0)
+      ready.push(group);
+
+  TaskList groupOrder;
+  while (!ready.empty()) {
+    size_t group = ready.top();
+    ready.pop();
+    groupOrder.push_back(group);
+    for (size_t successor : groups[group].successors)
+      if (--indegree[successor] == 0)
+        ready.push(successor);
+  }
+  if (groupOrder.size() != groups.size()) {
+    error = "pipeline scheduling cannot keep every operation in one stage "
+            "because the collapsed operation graph is cyclic";
+    return failure();
+  }
+
+  std::vector<size_t> position(groups.size());
+  for (auto [index, group] : llvm::enumerate(groupOrder))
+    position[group] = index;
+  size_t minimumStageZeroEnd = 0;
+  for (auto [index, group] : llvm::enumerate(groupOrder))
+    if (groups[group].consumesPipelineInput)
+      minimumStageZeroEnd = index + 1;
+  std::vector<size_t> furthestSuccessorBefore(groups.size() + 1, 0);
+  bool hasCrossingEdge = false;
+  size_t furthestSuccessor = 0;
+  for (size_t cut = 1; cut <= groups.size(); ++cut) {
+    size_t group = groupOrder[cut - 1];
+    for (size_t successor : groups[group].successors) {
+      if (position[successor] <= position[group]) {
+        error = "pipeline scheduling operation order is not topological";
+        return failure();
+      }
+      furthestSuccessor = std::max(furthestSuccessor, position[successor]);
+      hasCrossingEdge = true;
+    }
+    furthestSuccessorBefore[cut] = furthestSuccessor;
+  }
+
+  auto partitionGroups = [&](ArrayRef<size_t> coreCounts)
+      -> FailureOr<std::vector<size_t>> {
+    const Cost infinity = std::numeric_limits<Cost>::max();
+    const size_t noCut = std::numeric_limits<size_t>::max();
+    std::vector<std::vector<Cost>> best(
+      coreCounts.size() + 1,
+      std::vector<Cost>(groups.size() + 1, infinity));
+    std::vector<std::vector<size_t>> parent(
+      coreCounts.size() + 1,
+      std::vector<size_t>(groups.size() + 1, noCut));
+    std::vector<size_t> stageCostCache(coreCounts.size());
+    std::vector<size_t> cachedCoreCounts;
+    std::vector<std::vector<std::vector<Cost>>> segmentCostCaches;
+    for (auto [stage, coreCount] : llvm::enumerate(coreCounts)) {
+      auto cached = llvm::find(cachedCoreCounts, coreCount);
+      if (cached == cachedCoreCounts.end()) {
+        stageCostCache[stage] = segmentCostCaches.size();
+        cachedCoreCounts.push_back(coreCount);
+        segmentCostCaches.emplace_back(
+          groups.size() + 1,
+          std::vector<Cost>(groups.size() + 1, infinity));
+      } else {
+        stageCostCache[stage] = cached - cachedCoreCounts.begin();
+      }
+    }
+    best[0][0] = 0;
+
+    // ponytail: operation groups are small; replace this quadratic partition
+    // only if scheduling profiles show it matters.
+    for (size_t stage = 0; stage < coreCounts.size(); ++stage) {
+      size_t coreCount = coreCounts[stage];
+      size_t stageWeightCapacity = checkedMultiply(
+        coreCount, residentWeightCapacity);
+      for (size_t start = 0; start < groups.size(); ++start) {
+        if (best[stage][start] == infinity)
+          continue;
+        TaskList segmentTasks;
+        llvm::SmallBitVector segmentWeights(indexedWeights.size());
+        size_t segmentWeightCount = 0;
+        for (size_t end = start + 1; end <= groups.size(); ++end) {
+          const PipelineGroup &group = groups[groupOrder[end - 1]];
+          llvm::append_range(segmentTasks, group.tasks);
+          for (size_t weight : group.weightIds)
+            if (!segmentWeights.test(weight)) {
+              segmentWeights.set(weight);
+              ++segmentWeightCount;
+            }
+          if (segmentWeightCount > stageWeightCapacity)
+            break;
+          if (stage == 0 && end < minimumStageZeroEnd)
+            continue;
+          if (stage != 0 && hasCrossingEdge
+              && furthestSuccessorBefore[start] >= end)
+            continue;
+          Cost &segmentCost =
+            segmentCostCaches[stageCostCache[stage]][start][end];
+          if (segmentCost == infinity)
+            segmentCost = findMaximumIndexedPackCost(
+              model.schedulingCosts, taskWeightIds, indexedWeights.size(),
+              segmentTasks, residentWeightCapacity, coreCount);
+          Cost maximumLoad = std::max(best[stage][start], segmentCost);
+          if (maximumLoad < best[stage + 1][end]) {
+            best[stage + 1][end] = maximumLoad;
+            parent[stage + 1][end] = start;
+          }
+        }
+      }
+    }
+
+    size_t usedStages = 0;
+    Cost bestLoad = infinity;
+    for (size_t stages = 1; stages <= coreCounts.size(); ++stages)
+      if (best[stages][groups.size()] != infinity
+          && best[stages][groups.size()] <= bestLoad) {
+        bestLoad = best[stages][groups.size()];
+        usedStages = stages;
+      }
+    if (usedStages == 0) {
+      error = "pipeline scheduling cannot split operations into dependency-adjacent "
+              "stages within the physical crossbar limit";
+      return failure();
+    }
+
+    std::vector<size_t> groupStages(groups.size());
+    size_t end = groups.size();
+    for (size_t stage = usedStages; stage > 0; --stage) {
+      size_t start = parent[stage][end];
+      assert(start != noCut && "selected pipeline partition has no parent");
+      for (size_t position = start; position < end; ++position)
+        groupStages[groupOrder[position]] = stage - 1;
+      end = start;
+    }
+    return groupStages;
+  };
+
+  FailureOr<std::vector<size_t>> initialGroupStages =
+    partitionGroups(layout.getStageSizes());
+  if (failed(initialGroupStages))
+    return failure();
+
+  std::vector<TaskList> tasksByStage(layout.getStageCount());
+  for (size_t group : groupOrder)
+    llvm::append_range(
+      tasksByStage[(*initialGroupStages)[group]], groups[group].tasks);
+  std::vector<size_t> stageSizes(layout.getStageCount(), 1);
+  size_t assignedCores = stageSizes.size();
+  for (size_t stage = 0; stage < stageSizes.size(); ++stage) {
+    for (size_t task : tasksByStage[stage])
+      if (taskWeightIds[task].size() > residentWeightCapacity) {
+        error = "pipeline scheduling cannot fit one compute instance in a "
+                "physical core's crossbars";
+        return failure();
+      }
+    stageSizes[stage] = std::max(
+      stageSizes[stage], findMinimumIndexedPackCount(
+                           taskWeightIds, tasksByStage[stage],
+                           residentWeightCapacity));
+    assignedCores += stageSizes[stage] - 1;
+  }
+  if (assignedCores > layout.getProcessorCount()) {
+    error = "pipeline scheduling cannot fit dependency-adjacent stages "
+            "within the physical crossbar limit";
+    return failure();
+  }
+
+  const TaskCosts &balanceCosts = getPipelineBalanceCosts(graph, model);
+  auto getStageCost = [&](size_t stage, size_t coreCount) {
+    Cost schedulingCost = findMaximumIndexedPackCost(
+      model.schedulingCosts, taskWeightIds, indexedWeights.size(),
+      tasksByStage[stage], residentWeightCapacity, coreCount);
+    if (&balanceCosts == &model.schedulingCosts)
+      return schedulingCost;
+    Cost assemblyCost = findMaximumIndexedPackCost(
+      balanceCosts, taskWeightIds, indexedWeights.size(),
+      tasksByStage[stage], residentWeightCapacity, coreCount);
+    return std::max(schedulingCost, assemblyCost);
+  };
+  std::vector<Cost> stageCosts(stageSizes.size());
+  std::vector<Cost> nextStageCosts(stageSizes.size());
+  for (size_t stage = 0; stage < stageSizes.size(); ++stage)
+    stageCosts[stage] = getStageCost(stage, stageSizes[stage]);
+  for (size_t stage = 0; stage < stageSizes.size(); ++stage)
+    nextStageCosts[stage] = getStageCost(stage, stageSizes[stage] + 1);
+  while (assignedCores < layout.getProcessorCount()) {
+    size_t bestStage = 0;
+    Cost bestBenefit = 0;
+    for (size_t stage = 0; stage < stageSizes.size(); ++stage) {
+      Cost benefit = stageCosts[stage] - nextStageCosts[stage];
+      if (benefit > bestBenefit
+          || (benefit == bestBenefit
+              && (stageCosts[stage] > stageCosts[bestStage]
+                  || (stageCosts[stage] == stageCosts[bestStage]
+                      && stageSizes[stage] < stageSizes[bestStage])))) {
+        bestStage = stage;
+        bestBenefit = benefit;
+      }
+    }
+    ++stageSizes[bestStage];
+    stageCosts[bestStage] = nextStageCosts[bestStage];
+    nextStageCosts[bestStage] =
+      getStageCost(bestStage, stageSizes[bestStage] + 1);
+    ++assignedCores;
+  }
+  FailureOr<std::vector<size_t>> refinedGroupStages =
+    partitionGroups(stageSizes);
+  if (failed(refinedGroupStages))
+    return failure();
+  std::vector<size_t> taskStages(graph.nodes.size());
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    taskStages[task] = (*refinedGroupStages)[taskToGroup[task]];
+  return PipelineStageAssignment {
+    std::move(taskStages), std::move(stageSizes)};
+}
+
+static bool packPipelineStage(
+    const ComputeGraph &graph, const TaskCosts &schedulingCosts,
+    const TaskCosts &assemblyCosts,
+    std::vector<TaskList> &tasksByCpu, const PipelineCoreLayout &layout,
+    ArrayRef<size_t> topologicalPosition, size_t stage,
+    size_t residentWeightCapacity, const SchedulingTarget &target,
+    std::vector<size_t> &taskCpus) {
+  PipelineStageRange range = layout.getStageRange(stage);
+  TaskList tasks;
+  for (size_t cpu = range.begin; cpu < range.begin + range.size; ++cpu)
+    llvm::append_range(tasks, tasksByCpu[cpu]);
+  llvm::sort(tasks, [&](size_t lhs, size_t rhs) {
+    return topologicalPosition[lhs] < topologicalPosition[rhs];
+  });
+  std::vector<TaskList> packed(range.size);
+  std::vector<ResidentWeightSet> weights(range.size);
+  TaskCosts loads(range.size);
+  TaskCosts assemblyLoads(range.size);
   for (size_t task : tasks) {
     std::optional<size_t> bestCore;
-    std::optional<std::tuple<Cost, Cost, size_t, size_t>> bestScore;
-    for (size_t core = 0; core < groupSize; ++core) {
+    using PackScore = std::tuple<Cost, Time, Cost, Cost, size_t>;
+    std::optional<PackScore> bestScore;
+    for (size_t core = 0; core < range.size; ++core) {
       size_t unionSize = getResidentWeightUnionSize(
         weights[core], graph.nodes[task].residentWeights);
       if (unionSize > residentWeightCapacity)
         continue;
       size_t addedWeights = unionSize - weights[core].size();
-      auto score = std::make_tuple(
-        checkedAdd(assemblyLoads[core], assemblyCosts[task]),
-        checkedAdd(loads[core], schedulingCosts[task]), addedWeights, core);
+      Cost assemblyLoad = checkedAdd(
+        assemblyLoads[core], assemblyCosts[task]);
+      Cost schedulingLoad = checkedAdd(
+        loads[core], schedulingCosts[task]);
+      Time transferTime = 0;
+      size_t candidateCpu = range.begin + core;
+      for (const auto &[predecessor, transferCost] :
+           graph.predecessors[task])
+        if (taskCpus[predecessor] < target.processorCount)
+          transferTime = checkedAdd(
+            transferTime, getPeftTransferTime(
+                            transferCost, taskCpus[predecessor],
+                            candidateCpu, target));
+      PackScore score {
+        assemblyLoad, transferTime, schedulingLoad, addedWeights, core};
       if (!bestScore || score < *bestScore) {
         bestCore = core;
         bestScore = score;
       }
     }
     if (!bestCore)
-      return;
+      return false;
     packed[*bestCore].push_back(task);
     insertResidentWeights(
       weights[*bestCore], graph.nodes[task].residentWeights);
     loads[*bestCore] = checkedAdd(loads[*bestCore], schedulingCosts[task]);
     assemblyLoads[*bestCore] = checkedAdd(
       assemblyLoads[*bestCore], assemblyCosts[task]);
+    taskCpus[task] = range.begin + *bestCore;
   }
-  if (*std::max_element(loads.begin(), loads.end()) > originalMaximum)
-    return;
-  for (size_t core = 0; core < groupSize; ++core)
-    tasksByCpu[stage * groupSize + core] = std::move(packed[core]);
+  for (size_t core = 0; core < range.size; ++core)
+    tasksByCpu[range.begin + core] = std::move(packed[core]);
+  return true;
 }
 
-static void rebalancePipelineStages(
+static LogicalResult packPipelineStages(
     const ComputeGraph &graph, const PipelineTaskModel &model,
-    std::vector<TaskList> &tasksByCpu, size_t groupSize,
-    size_t pipelineStages, size_t residentWeightCapacity) {
-  size_t minimumAssemblyFanIn = std::numeric_limits<size_t>::max();
+    std::vector<TaskList> &tasksByCpu, const PipelineCoreLayout &layout,
+    size_t pipelineStages, size_t residentWeightCapacity,
+    const SchedulingTarget &target, size_t &failedStage,
+    std::string &error) {
+  std::vector<size_t> indegree(graph.nodes.size());
+  std::vector<TaskList> successors(graph.nodes.size());
   for (size_t task = 0; task < graph.nodes.size(); ++task)
-    if (graph.nodes[task].instance.op && model.assemblyCosts[task] > 1
-        && !model.predecessors[task].empty())
-      minimumAssemblyFanIn = std::min(
-        minimumAssemblyFanIn, model.predecessors[task].size());
-  bool hasAssembly = minimumAssemblyFanIn != std::numeric_limits<size_t>::max();
-  if (hasAssembly && groupSize < minimumAssemblyFanIn)
-    return;
-  const TaskCosts &balanceCosts =
-    hasAssembly ? model.assemblyCosts : model.schedulingCosts;
+    for (size_t predecessor : model.predecessors[task]) {
+      successors[predecessor].push_back(task);
+      ++indegree[task];
+    }
+  auto laterOriginalOrder = [&](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(laterOriginalOrder)>
+    ready(laterOriginalOrder);
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    if (indegree[task] == 0)
+      ready.push(task);
+  std::vector<size_t> topologicalPosition(graph.nodes.size());
+  size_t position = 0;
+  while (!ready.empty()) {
+    size_t task = ready.top();
+    ready.pop();
+    topologicalPosition[task] = position++;
+    for (size_t successor : successors[task])
+      if (--indegree[successor] == 0)
+        ready.push(successor);
+  }
+  if (position != graph.nodes.size()) {
+    error = "pipeline rebalancing received a cyclic task graph";
+    return failure();
+  }
 
-  Cost schedulingLimit = 0;
-  for (const TaskList &tasks : tasksByCpu)
-    schedulingLimit = std::max(
-      schedulingLimit, getCoreCost(tasks, model.schedulingCosts));
+  const TaskCosts &balanceCosts = getPipelineBalanceCosts(graph, model);
+  std::vector<size_t> taskCpus(graph.nodes.size(), target.processorCount);
   for (size_t stage = 0; stage < pipelineStages; ++stage)
-    repackPipelineStage(
-      graph, model.schedulingCosts, balanceCosts, tasksByCpu,
-      groupSize, stage,
-      residentWeightCapacity);
+    if (!packPipelineStage(
+          graph, model.schedulingCosts, balanceCosts, tasksByCpu,
+          layout, topologicalPosition, stage, residentWeightCapacity,
+          target, taskCpus)) {
+      failedStage = stage;
+      error = "pipeline scheduling cannot pack dependency-monotone stage "
+            + std::to_string(stage)
+            + " within the physical crossbar limit";
+      return failure();
+    }
+  return success();
+}
 
-  std::vector<size_t> taskToCpu(graph.nodes.size());
-  for (size_t cpu = 0; cpu < tasksByCpu.size(); ++cpu)
-    for (size_t task : tasksByCpu[cpu])
-      taskToCpu[task] = cpu;
-
-  bool changed;
-  do {
-    changed = false;
-    for (size_t sourceStage = pipelineStages; sourceStage-- > 1;) {
-      size_t targetStage = sourceStage - 1;
-      while (true) {
-        Cost sourceMaximum = getStageMaximumAssemblyCost(
-          tasksByCpu, balanceCosts, groupSize, sourceStage);
-        Cost targetMaximum = getStageMaximumAssemblyCost(
-          tasksByCpu, balanceCosts, groupSize, targetStage);
-        if (targetMaximum >= sourceMaximum)
-          break;
-
-        struct Move {
-          size_t sourceCpu;
-          size_t targetCpu;
-          size_t task;
-        };
-        std::optional<Move> best;
-        std::optional<std::tuple<size_t, Cost, Cost, size_t>> bestScore;
-        for (size_t sourceCpu = sourceStage * groupSize;
-             sourceCpu < (sourceStage + 1) * groupSize; ++sourceCpu) {
-          if (tasksByCpu[sourceCpu].empty())
-            continue;
-          size_t task = tasksByCpu[sourceCpu].front();
-          bool dependenciesReady = llvm::all_of(
-            model.predecessors[task], [&](size_t predecessor) {
-              return taskToCpu[predecessor] / groupSize <= targetStage;
-            });
-          if (!dependenciesReady)
-            continue;
-
-          Cost sourceAfter = getCoreCost(
-            tasksByCpu[sourceCpu], balanceCosts)
-                           - balanceCosts[task];
-          for (size_t targetCpu = targetStage * groupSize;
-               targetCpu < (targetStage + 1) * groupSize; ++targetCpu) {
-            const TaskList &targetTasks = tasksByCpu[targetCpu];
-            if (!fitsResidentWeights(
-                     graph, targetTasks, task, residentWeightCapacity))
-              continue;
-            Cost targetAfter = checkedAdd(
-              getCoreCost(targetTasks, balanceCosts), balanceCosts[task]);
-            Cost targetSchedulingAfter = checkedAdd(
-              getCoreCost(targetTasks, model.schedulingCosts),
-              model.schedulingCosts[task]);
-            if (targetAfter >= sourceMaximum
-                || targetSchedulingAfter > schedulingLimit)
-              continue;
-            auto score = std::make_tuple(
-              graph.nodes[task].originalOrder,
-              std::max(sourceAfter, targetAfter), targetAfter, targetCpu);
-            if (!bestScore || score < *bestScore) {
-              best = Move {sourceCpu, targetCpu, task};
-              bestScore = score;
-            }
-          }
-        }
-        if (!best)
-          break;
-        tasksByCpu[best->sourceCpu].erase(
-          tasksByCpu[best->sourceCpu].begin());
-        TaskList &targetTasks = tasksByCpu[best->targetCpu];
-        auto insertion = llvm::find_if(targetTasks, [&](size_t task) {
-          return graph.nodes[task].originalOrder
-               > graph.nodes[best->task].originalOrder;
-        });
-        targetTasks.insert(insertion, best->task);
-        taskToCpu[best->task] = best->targetCpu;
-        changed = true;
+static LogicalResult verifyPipelineStageAssignment(
+    const ComputeGraph &graph, const PipelineTaskModel &model,
+    const std::vector<TaskList> &tasksByCpu,
+    const PipelineCoreLayout &layout, std::string &error) {
+  const size_t noStage = std::numeric_limits<size_t>::max();
+  std::vector<size_t> taskStages(graph.nodes.size(), noStage);
+  llvm::DenseMap<Operation *, size_t> operationStages;
+  llvm::DenseMap<Attribute, size_t> splitOperationStages;
+  for (size_t cpu = 0; cpu < tasksByCpu.size(); ++cpu) {
+    std::optional<size_t> stage = layout.getStageForCore(cpu);
+    if (!stage) {
+      error = "pipeline scheduling assigned a task outside the stage layout";
+      return failure();
+    }
+    for (size_t task : tasksByCpu[cpu]) {
+      if (task >= graph.nodes.size() || taskStages[task] != noStage) {
+        error = "pipeline scheduling did not assign every task exactly once";
+        return failure();
+      }
+      taskStages[task] = *stage;
+      if (*stage != 0 && consumesPipelineInput(graph.nodes[task])) {
+        error = "pipeline scheduling assigned a direct function-input "
+                "consumer after stage zero";
+        return failure();
+      }
+      Operation *operation = graph.nodes[task].instance.op;
+      if (!operation)
+        continue;
+      Attribute splitGroup = operation->getAttr("pipeline.stage_group");
+      bool consistent;
+      if (splitGroup) {
+        auto [entry, inserted] =
+          splitOperationStages.try_emplace(splitGroup, *stage);
+        consistent = inserted || entry->second == *stage;
+      } else {
+        auto [entry, inserted] =
+          operationStages.try_emplace(operation, *stage);
+        consistent = inserted || entry->second == *stage;
+      }
+      if (!consistent) {
+        error = "pipeline scheduling split one operation across stages";
+        return failure();
       }
     }
-  } while (changed);
+  }
+  if (llvm::is_contained(taskStages, noStage)) {
+    error = "pipeline scheduling did not assign every task exactly once";
+    return failure();
+  }
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    for (size_t predecessor : model.predecessors[task])
+      if (taskStages[predecessor] > taskStages[task]
+          || taskStages[task] - taskStages[predecessor] > 1) {
+        error = "pipeline scheduling produced a backward or skipped-stage dependency";
+        return failure();
+      }
+  return success();
 }
 
 mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
@@ -914,7 +1313,13 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
                                         const SchedulingTarget& physicalTarget,
                                         std::string& error) {
   const size_t groupSize = schedule.processorCount;
-  std::vector<TaskList> tasksByCpu(groupSize);
+  PipelineCoreLayout balancedLayout(
+    physicalTarget.processorCount, pipelineStages);
+  if (!balancedLayout.isValid()
+      || groupSize != balancedLayout.getLogicalProcessorCount()) {
+    error = "pipeline scheduling received an incompatible physical core layout";
+    return mlir::failure();
+  }
   for (size_t task = 0; task < graph.nodes.size(); ++task) {
     const ComputeInstance& instance = graph.nodes[task].instance;
     auto cpu = schedule.computeToCpuMap.find(instance);
@@ -924,74 +1329,69 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
       error = "pipeline scheduling received an incomplete PEFT schedule";
       return mlir::failure();
     }
-    tasksByCpu[cpu->second].push_back(task);
   }
-  for (TaskList& tasks : tasksByCpu)
-    llvm::sort(tasks, [&](size_t lhs, size_t rhs) {
-      return schedule.computeToCpuSlotMap.lookup(graph.nodes[lhs].instance)
-           < schedule.computeToCpuSlotMap.lookup(graph.nodes[rhs].instance);
-    });
   PipelineTaskModel taskModel = getPipelineTaskModel(
     graph, schedule, physicalTarget);
-  const TaskCosts &taskCosts = taskModel.schedulingCosts;
-
+  FailureOr<PipelineStageAssignment> assignment = assignPipelineStages(
+    graph, taskModel, balancedLayout,
+    physicalTarget.residentWeightCapacity, error);
+  if (failed(assignment))
+    return failure();
   std::vector<TaskList> tasksByPhysicalCpu(physicalTarget.processorCount);
-  for (size_t sourceCpu = 0; sourceCpu < groupSize; ++sourceCpu) {
-    const TaskList& tasks = tasksByCpu[sourceCpu];
-    if (tasks.empty())
-      continue;
-    for (size_t task : tasks)
-      if (graph.nodes[task].residentWeights.size() > physicalTarget.residentWeightCapacity) {
-        error = "pipeline scheduling cannot fit one compute instance in a physical core's crossbars";
-        return mlir::failure();
-      }
-
-    Cost maximumCost = findMaximumPackCost(
-      graph, taskCosts, tasks, physicalTarget.residentWeightCapacity, pipelineStages);
-    if (!fits(graph, taskCosts, tasks, maximumCost,
-              physicalTarget.residentWeightCapacity, pipelineStages)) {
-      error = "pipeline scheduling cannot partition one PEFT core within the physical crossbar limit";
-      return mlir::failure();
+  std::vector<size_t> minimumPackableStageSizes(pipelineStages, 1);
+  std::string packingError;
+  bool packed = false;
+  for (size_t attempt = 0; attempt < physicalTarget.processorCount; ++attempt) {
+    PipelineCoreLayout candidateLayout(assignment->stageSizes);
+    for (TaskList &tasks : tasksByPhysicalCpu)
+      tasks.clear();
+    for (size_t task = 0; task < graph.nodes.size(); ++task) {
+      PipelineStageRange range =
+        candidateLayout.getStageRange(assignment->taskStages[task]);
+      tasksByPhysicalCpu[range.begin].push_back(task);
     }
-
-    const size_t desiredPacks = std::min(pipelineStages, tasks.size());
-    size_t stage = 0;
-    Cost packCost = 0;
-    ResidentWeightSet packWeights;
-    bool packEmpty = true;
-    for (size_t index = 0; index < tasks.size(); ++index) {
-      size_t task = tasks[index];
-      const ComputeGraphNode& node = graph.nodes[task];
-      Cost taskCost = taskCosts[task];
-      bool exceedsLimit =
-        !packEmpty
-        && (packCost > maximumCost - taskCost
-            || getResidentWeightUnionSize(packWeights, node.residentWeights) > physicalTarget.residentWeightCapacity);
-      bool reserveOneTaskPerPack = !packEmpty && tasks.size() - index == desiredPacks - stage - 1;
-      if (exceedsLimit || reserveOneTaskPerPack) {
-        ++stage;
-        packCost = 0;
-        packWeights.clear();
-        packEmpty = true;
-      }
-      if (stage >= pipelineStages) {
-        error = "pipeline scheduling produced too many packs";
-        return mlir::failure();
-      }
-      size_t physicalCpu = sourceCpu + stage * groupSize;
-      tasksByPhysicalCpu[physicalCpu].push_back(task);
-      packCost = checkedAdd(packCost, taskCost);
-      insertResidentWeights(packWeights, node.residentWeights);
-      packEmpty = false;
+    size_t failedStage = 0;
+    if (succeeded(packPipelineStages(
+          graph, taskModel, tasksByPhysicalCpu, candidateLayout,
+          pipelineStages, physicalTarget.residentWeightCapacity,
+          physicalTarget, failedStage, packingError))) {
+      packed = true;
+      break;
     }
+    minimumPackableStageSizes[failedStage] = std::max(
+      minimumPackableStageSizes[failedStage],
+      assignment->stageSizes[failedStage] + 1);
+    std::optional<size_t> donor;
+    for (size_t stage = 0; stage < pipelineStages; ++stage)
+      if (stage != failedStage
+          && assignment->stageSizes[stage]
+               > minimumPackableStageSizes[stage]
+          && (!donor
+              || assignment->stageSizes[stage]
+                   > assignment->stageSizes[*donor]))
+        donor = stage;
+    if (!donor)
+      break;
+    --assignment->stageSizes[*donor];
+    ++assignment->stageSizes[failedStage];
   }
-
-  rebalancePipelineStages(
-    graph, taskModel, tasksByPhysicalCpu, groupSize, pipelineStages,
-    physicalTarget.residentWeightCapacity);
+  if (!packed) {
+    error = packingError;
+    return failure();
+  }
+  PipelineCoreLayout pipelineLayout(assignment->stageSizes);
+  if (failed(verifyPipelineStageAssignment(
+        graph, taskModel, tasksByPhysicalCpu, pipelineLayout, error)))
+    return failure();
 
   schedule.computeToCpuMap.clear();
   schedule.processorCount = physicalTarget.processorCount;
+  schedule.processorStages.resize(physicalTarget.processorCount);
+  for (size_t stage = 0; stage < pipelineLayout.getStageCount(); ++stage) {
+    PipelineStageRange range = pipelineLayout.getStageRange(stage);
+    std::fill_n(
+      schedule.processorStages.begin() + range.begin, range.size, stage);
+  }
   schedule.computeToCpuSlotMap.clear();
   schedule.computeToAestMap.clear();
   schedule.isLastComputeOfCpu.clear();
@@ -1071,21 +1471,22 @@ mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
                                             std::string& error) {
   if (pipelineStages == 1)
     return mlir::success();
-  if (pipelineStages == 0 || schedule.processorCount == 0
-      || schedule.processorCount > std::numeric_limits<size_t>::max() / pipelineStages
-      || schedule.processorCount * pipelineStages != physicalTarget.processorCount) {
-    error = "pipeline scheduling requires physical cores = scheduled cores * pipeline stages";
+  PipelineCoreLayout pipelineLayout(
+    physicalTarget.processorCount, pipelineStages);
+  if (!pipelineLayout.isValid() || schedule.processorCount == 0
+      || schedule.processorCount
+           != pipelineLayout.getLogicalProcessorCount()) {
+    error = "pipeline scheduling requires a valid balanced physical core layout";
     return mlir::failure();
   }
   return assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, error);
 }
 
-mlir::LogicalResult splitPipelineWorkload(const ComputeGraph &graph,
-                                          const MergeScheduleResult &schedule,
-                                          size_t pipelineStages,
-                                          const SchedulingTarget &physicalTarget,
-                                          std::string &error) {
-  return splitPipelineWorkloadImpl(
+mlir::FailureOr<PipelineWorkloadPreparation> preparePipelineWorkload(
+    const ComputeGraph &graph, const MergeScheduleResult &schedule,
+    size_t pipelineStages, const SchedulingTarget &physicalTarget,
+    std::string &error) {
+  return preparePipelineWorkloadImpl(
     graph, schedule, pipelineStages, physicalTarget, error);
 }
 

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

@@ -2,8 +2,15 @@
 
 #include "mlir/Support/LogicalResult.h"
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include <algorithm>
 #include <cstddef>
+#include <numeric>
+#include <optional>
 #include <string>
+#include <vector>
 
 #include "ComputeGraph.hpp"
 #include "MergeSchedule.hpp"
@@ -11,16 +18,82 @@
 
 namespace onnx_mlir::spatial {
 
+struct PipelineStageRange {
+  size_t begin;
+  size_t size;
+};
+
+class PipelineCoreLayout {
+public:
+  PipelineCoreLayout(size_t processorCount, size_t stageCount)
+      : processorCount(processorCount), stageSizes(stageCount) {
+    if (stageCount == 0)
+      return;
+    size_t baseSize = processorCount / stageCount;
+    size_t largerStageCount = processorCount % stageCount;
+    for (size_t stage = 0; stage < stageCount; ++stage)
+      stageSizes[stage] = baseSize + (stage < largerStageCount);
+  }
+
+  explicit PipelineCoreLayout(llvm::ArrayRef<size_t> stageSizes)
+      : processorCount(std::accumulate(
+          stageSizes.begin(), stageSizes.end(), size_t {0})),
+        stageSizes(stageSizes.begin(), stageSizes.end()) {}
+
+  bool isValid() const {
+    return !stageSizes.empty()
+        && llvm::none_of(stageSizes, [](size_t size) { return size == 0; });
+  }
+
+  size_t getLogicalProcessorCount() const {
+    return isValid()
+      ? *std::min_element(stageSizes.begin(), stageSizes.end())
+      : 0;
+  }
+
+  size_t getStageCount() const { return stageSizes.size(); }
+
+  size_t getProcessorCount() const { return processorCount; }
+
+  llvm::ArrayRef<size_t> getStageSizes() const { return stageSizes; }
+
+  PipelineStageRange getStageRange(size_t stage) const {
+    return {std::accumulate(
+              stageSizes.begin(), stageSizes.begin() + stage, size_t {0}),
+            stageSizes[stage]};
+  }
+
+  std::optional<size_t> getStageForCore(size_t core) const {
+    if (!isValid() || core >= processorCount)
+      return std::nullopt;
+    size_t end = 0;
+    for (auto [stage, size] : llvm::enumerate(stageSizes)) {
+      end += size;
+      if (core < end)
+        return stage;
+    }
+    return std::nullopt;
+  }
+
+private:
+  size_t processorCount;
+  std::vector<size_t> stageSizes;
+};
+
 mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
                                             MergeScheduleResult& schedule,
                                             size_t pipelineStages,
                                             const SchedulingTarget& physicalTarget,
                                             std::string& error);
 
-mlir::LogicalResult splitPipelineWorkload(const ComputeGraph& graph,
-                                          const MergeScheduleResult& schedule,
-                                          size_t pipelineStages,
-                                          const SchedulingTarget& physicalTarget,
-                                          std::string& error);
+enum class PipelineWorkloadPreparation {
+  Ready,
+  Changed,
+};
+
+mlir::FailureOr<PipelineWorkloadPreparation> preparePipelineWorkload(
+    const ComputeGraph& graph, const MergeScheduleResult& schedule,
+    size_t pipelineStages, const SchedulingTarget& physicalTarget,
+    std::string& error);
 
 } // namespace onnx_mlir::spatial