better scheduling

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

@@ -277,10 +277,10 @@ static LogicalResult assignPipelineSynchronization(
   llvm::sort(transfers.stageZeroCores);
   transfers.stageZeroCores.erase(
     llvm::unique(transfers.stageZeroCores), transfers.stageZeroCores.end());
-  llvm::sort(transfers.downstreamCores);
-  transfers.downstreamCores.erase(
-    llvm::unique(transfers.downstreamCores),
-    transfers.downstreamCores.end());
+  DenseSet<int64_t> downstreamSeen;
+  llvm::erase_if(transfers.downstreamCores, [&](int64_t core) {
+    return !downstreamSeen.insert(core).second;
+  });
   llvm::erase_if(transfers.downstreamCores, [&](int64_t core) {
     return llvm::is_contained(transfers.stageZeroCores, core);
   });

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

@@ -36,9 +36,9 @@ struct TransferGroup {
   SmallVector<StreamThreshold> dependencies;
   unsigned unsatisfied = 0;
   bool scheduled = false;
-  std::tuple<unsigned, unsigned, unsigned, uint64_t> originalPriority;
+  std::tuple<uint64_t, unsigned, unsigned, unsigned> originalPriority;
   std::tuple<unsigned, unsigned, unsigned,
-             std::tuple<unsigned, unsigned, unsigned, uint64_t>> staticPriority;
+             std::tuple<uint64_t, unsigned, unsigned, unsigned>> staticPriority;
 };
 
 struct StreamProgress {
@@ -189,8 +189,12 @@ static SmallVector<TransferGroup> buildGroups(DeferredTransferPlan& plan) {
       unsigned source = family.sourceStreams.valueAt(0);
       thresholdByStream[source] = std::max(thresholdByStream.lookup(source), family.requirement->producer->step + 1);
       minSource = std::min(minSource, source);
-      for (size_t index = 0; index < family.targetStreams.size(); ++index)
-        minTarget = std::min<unsigned>(minTarget, family.targetStreams.valueAt(index));
+      for (size_t index = 0; index < family.targetStreams.size(); ++index) {
+        unsigned target = family.targetStreams.valueAt(index);
+        thresholdByStream[target] = std::max(
+          thresholdByStream.lookup(target), exchange->consumerStep);
+        minTarget = std::min(minTarget, target);
+      }
     }
     for (LocalAvailabilityFamily& local : exchange->local)
       for (LaneInterval interval : local.targetLanes.intervals())
@@ -206,7 +210,7 @@ static SmallVector<TransferGroup> buildGroups(DeferredTransferPlan& plan) {
       group.dependencies.push_back({stream, thresholdByStream.lookup(stream)});
     group.unsatisfied =
       llvm::count_if(group.dependencies, [](StreamThreshold dependency) { return dependency.completedStep != 0; });
-    group.originalPriority = {exchange->consumerStep, minSource, minTarget, exchange->exchangeId};
+    group.originalPriority = {exchange->exchangeId, exchange->consumerStep, minSource, minTarget};
     orderGroupSlices(group);
     groups.push_back(std::move(group));
   }

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

@@ -24,6 +24,7 @@
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Support/TypeUtilities.hpp"
 
 namespace onnx_mlir {
@@ -233,8 +234,20 @@ Cost getOperationCost(Operation& op, const SchedulerCostModel& costModel) {
       costModel.getUnaryVectorCost(softmax.getInput().getType(), softmax.getOutput().getType());
     return checkedMultiply(unary, static_cast<Cost>(4));
   }
-  if (auto extract = dyn_cast<tensor::ExtractSliceOp>(&op))
+  if (auto extract = dyn_cast<tensor::ExtractSliceOp>(&op)) {
+    bool weightOnly = !extract.getResult().use_empty()
+      && llvm::all_of(extract.getResult().getUses(), [&](OpOperand &use) {
+        auto vmm = dyn_cast<SpatVMMOp>(use.getOwner());
+        return vmm && vmm.getWeight() == extract.getResult();
+      });
+    if (weightOnly)
+      return 0;
+    if (extract.getSourceType().hasStaticShape() && isContiguousSubviewWithDynamicOffsets(
+          extract.getSourceType().getShape(), extract.getMixedOffsets(),
+          extract.getStaticSizes(), extract.getStaticStrides()))
+      return 0;
     return costModel.getTensorMoveCost(extract.getResult().getType());
+  }
   if (auto insert = dyn_cast<tensor::InsertSliceOp>(&op))
     return costModel.getTensorMoveCost(insert.getSource().getType());
 

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

@@ -97,22 +97,24 @@ std::vector<ResidentWeightSet> planResidentWeightReservations(const ComputeGraph
                                                     size_t processorCount,
                                                               size_t residentWeightCapacity,
                                                               const TopologyModel& topology,
-                                                              bool preferWeightReuse) {
+                                                              bool preferWeightReuse,
+                                                              bool preserveTaskOrder = true) {
   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 (graph.nodes[lhs].residentWeights.size() != graph.nodes[rhs].residentWeights.size())
+    if ((!preserveTaskOrder || preferWeightReuse) && 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;
   });
 
   std::vector<ResidentWeightSet> reservations(processorCount);
   std::vector<Time> reservedLoad(processorCount, 0);
+  std::vector<size_t> reservedCpus(graph.nodes.size(), processorCount);
   for (size_t task : weightedTasks) {
     size_t bestProcessor = std::numeric_limits<size_t>::max();
-    using ReservationScore = std::tuple<Time, size_t, size_t, size_t>;
+    using ReservationScore = std::tuple<Time, size_t, Time, size_t, size_t>;
     std::optional<ReservationScore> bestScore;
     for (size_t processor = 0; processor < processorCount; ++processor) {
       size_t residentWeightUnion =
@@ -120,8 +122,19 @@ std::vector<ResidentWeightSet> planResidentWeightReservations(const ComputeGraph
       if (residentWeightUnion > residentWeightCapacity)
         continue;
       size_t addedWeights = residentWeightUnion - reservations[processor].size();
-      ReservationScore score {preferWeightReuse ? addedWeights : reservedLoad[processor],
-                              preferWeightReuse ? reservedLoad[processor] : addedWeights,
+      Time communication = 0;
+      for (const auto &[predecessor, cost] : graph.predecessors[task])
+        if (reservedCpus[predecessor] != processorCount)
+          communication = addOrMax(communication, getPeftTransferTime(
+            cost, reservedCpus[predecessor], processor, topology.target));
+      for (const auto &[successor, cost] : graph.successors[task])
+        if (reservedCpus[successor] != processorCount)
+          communication = addOrMax(communication, getPeftTransferTime(
+            cost, processor, reservedCpus[successor], topology.target));
+      Time load = addOrMax(reservedLoad[processor], communication);
+      ReservationScore score {preferWeightReuse ? addedWeights : load,
+                              preferWeightReuse ? load : addedWeights,
+                              communication,
                               topology.getCenterDistance(processor),
                               processor};
       if (!bestScore || score < *bestScore) {
@@ -130,6 +143,10 @@ 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)
+        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} "
                       "processors of capacity {3}",
@@ -141,6 +158,7 @@ std::vector<ResidentWeightSet> planResidentWeightReservations(const ComputeGraph
       llvm::report_fatal_error(llvm::StringRef(message));
     }
     insertResidentWeights(reservations[bestProcessor], graph.nodes[task].residentWeights);
+    reservedCpus[task] = bestProcessor;
     reservedLoad[bestProcessor] = addOrMax(reservedLoad[bestProcessor], graph.nodes[task].cost);
   }
   return reservations;
@@ -440,6 +458,13 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     if (scheduledCount != nodeCount)
       llvm::report_fatal_error("PEFT scheduler: failed to schedule every compute node");
 
+    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) {

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

@@ -47,6 +47,8 @@ struct PipelineTaskModel {
   TaskCosts schedulingCosts;
   TaskCosts assemblyCosts;
   std::vector<TaskList> predecessors;
+  std::vector<TaskList> weightIds;
+  size_t weightCount = 0;
 };
 
 static BlueprintAssemblyInfo getBlueprintAssemblyInfo(
@@ -114,6 +116,17 @@ static PipelineTaskModel getPipelineTaskModel(
   model.schedulingCosts.reserve(graph.nodes.size());
   model.assemblyCosts.reserve(graph.nodes.size());
   model.predecessors.resize(graph.nodes.size());
+  model.weightIds.resize(graph.nodes.size());
+  ResidentWeightSet indexedWeights;
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    for (const ResidentWeight &weight : graph.nodes[task].residentWeights) {
+      auto indexed = llvm::find(indexedWeights, weight);
+      size_t id = indexed - indexedWeights.begin();
+      if (indexed == indexedWeights.end())
+        indexedWeights.push_back(weight);
+      model.weightIds[task].push_back(id);
+    }
+  model.weightCount = indexedWeights.size();
   for (const ComputeGraphNode &node : graph.nodes)
     if (node.instance.op) {
       model.schedulingCosts.push_back(node.cost);
@@ -596,7 +609,7 @@ static Cost findMaximumIndexedPackCost(
     const TaskCosts &taskCosts,
     const std::vector<TaskList> &taskWeightIds, size_t weightCount,
     const TaskList &tasks, size_t residentWeightCapacity,
-    size_t maximumPacks) {
+    size_t maximumPacks, TaskList *taskPacks = nullptr) {
   Cost low = 0;
   Cost high = 0;
   for (size_t task : tasks) {
@@ -608,8 +621,7 @@ static Cost findMaximumIndexedPackCost(
 
   std::vector<size_t> seen(weightCount);
   size_t generation = 0;
-  while (low < high) {
-    Cost middle = low + (high - low) / 2;
+  auto pack = [&](Cost maximumCost, TaskList *assignment) {
     size_t packs = 1;
     Cost cost = 0;
     size_t packWeightCount = 0;
@@ -622,7 +634,7 @@ static Cost findMaximumIndexedPackCost(
         addedWeights += seen[weight] != generation;
       Cost taskCost = taskCosts[task];
       bool startsNewPack = !packEmpty
-        && (cost > middle - taskCost
+        && (cost > maximumCost - taskCost
             || packWeightCount + addedWeights > residentWeightCapacity);
       if (startsNewPack) {
         if (++packs > maximumPacks) {
@@ -640,13 +652,24 @@ static Cost findMaximumIndexedPackCost(
           seen[weight] = generation;
           ++packWeightCount;
         }
+      if (assignment)
+        assignment->push_back(packs - 1);
       packEmpty = false;
     }
-    if (fits)
+    return fits;
+  };
+  while (low < high) {
+    Cost middle = low + (high - low) / 2;
+    if (pack(middle, nullptr))
       high = middle;
     else
       low = middle + 1;
   }
+  if (taskPacks) {
+    taskPacks->clear();
+    if (!pack(low, taskPacks))
+      taskPacks->clear();
+  }
   return low;
 }
 
@@ -685,6 +708,30 @@ struct PipelineStageAssignment {
   std::vector<size_t> stageSizes;
 };
 
+static TaskCosts getPipelineStageTailCosts(
+    const std::vector<PipelineGroup> &groups, ArrayRef<size_t> groupOrder,
+    ArrayRef<size_t> groupStages, const TaskCosts &taskCosts, size_t stageCount) {
+  TaskCosts groupEnds(groups.size());
+  TaskCosts stageSpans(stageCount);
+  TaskCosts stageFloors(stageCount);
+  for (size_t group : groupOrder) {
+    size_t stage = groupStages[group];
+    Cost cost = 0;
+    for (size_t task : groups[group].tasks)
+      cost = std::max(cost, taskCosts[task]);
+    groupEnds[group] = checkedAdd(groupEnds[group], cost);
+    stageSpans[stage] = std::max(stageSpans[stage], groupEnds[group]);
+    stageFloors[stage] = std::max(stageFloors[stage], cost);
+    for (size_t successor : groups[group].successors)
+      if (groupStages[successor] == stage)
+        groupEnds[successor] = std::max(groupEnds[successor], groupEnds[group]);
+  }
+  // The largest task is already included in the per-core packing estimate.
+  for (size_t stage = 0; stage < stageCount; ++stage)
+    stageSpans[stage] -= stageFloors[stage];
+  return stageSpans;
+}
+
 static const TaskCosts &getPipelineBalanceCosts(
     const ComputeGraph &graph, const PipelineTaskModel &model) {
   for (size_t task = 0; task < graph.nodes.size(); ++task)
@@ -716,18 +763,7 @@ static FailureOr<PipelineStageAssignment> assignPipelineStages(
 
   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);
-    }
+  const auto &taskWeightIds = model.weightIds;
   llvm::DenseMap<Operation *, size_t> operationToGroup;
   llvm::DenseMap<Attribute, size_t> splitOperationToGroup;
   for (size_t task : tasksByOrder) {
@@ -802,6 +838,11 @@ static FailureOr<PipelineStageAssignment> assignPipelineStages(
     return failure();
   }
 
+  TaskCosts groupCosts(groups.size());
+  for (size_t group : groupOrder)
+    for (size_t task : groups[group].tasks)
+      groupCosts[group] = std::max(groupCosts[group], model.schedulingCosts[task]);
+
   std::vector<size_t> position(groups.size());
   for (auto [index, group] : llvm::enumerate(groupOrder))
     position[group] = index;
@@ -862,10 +903,21 @@ static FailureOr<PipelineStageAssignment> assignPipelineStages(
         if (best[stage][start] == infinity)
           continue;
         TaskList segmentTasks;
-        llvm::SmallBitVector segmentWeights(indexedWeights.size());
+        TaskCosts segmentEnds(groups.size());
+        Cost segmentSpan = 0;
+        Cost largestTask = 0;
+        llvm::SmallBitVector segmentWeights(model.weightCount);
         size_t segmentWeightCount = 0;
         for (size_t end = start + 1; end <= groups.size(); ++end) {
           const PipelineGroup &group = groups[groupOrder[end - 1]];
+          size_t groupIndex = groupOrder[end - 1];
+          segmentEnds[groupIndex] = checkedAdd(
+            segmentEnds[groupIndex], groupCosts[groupIndex]);
+          segmentSpan = std::max(segmentSpan, segmentEnds[groupIndex]);
+          largestTask = std::max(largestTask, groupCosts[groupIndex]);
+          for (size_t successor : group.successors)
+            segmentEnds[successor] = std::max(
+              segmentEnds[successor], segmentEnds[groupIndex]);
           llvm::append_range(segmentTasks, group.tasks);
           for (size_t weight : group.weightIds)
             if (!segmentWeights.test(weight)) {
@@ -881,10 +933,12 @@ static FailureOr<PipelineStageAssignment> assignPipelineStages(
             continue;
           Cost &segmentCost =
             segmentCostCaches[stageCostCache[stage]][start][end];
-          if (segmentCost == infinity)
+          if (segmentCost == infinity) {
             segmentCost = findMaximumIndexedPackCost(
-              model.schedulingCosts, taskWeightIds, indexedWeights.size(),
+              model.schedulingCosts, taskWeightIds, model.weightCount,
               segmentTasks, residentWeightCapacity, coreCount);
+            segmentCost = checkedAdd(segmentCost, segmentSpan - largestTask);
+          }
           Cost maximumLoad = std::max(best[stage][start], segmentCost);
           if (maximumLoad < best[stage + 1][end]) {
             best[stage + 1][end] = maximumLoad;
@@ -951,35 +1005,43 @@ static FailureOr<PipelineStageAssignment> assignPipelineStages(
   }
 
   const TaskCosts &balanceCosts = getPipelineBalanceCosts(graph, model);
+  TaskCosts stageTailCosts = getPipelineStageTailCosts(
+    groups, groupOrder, *initialGroupStages, model.schedulingCosts, stageSizes.size());
   auto getStageCost = [&](size_t stage, size_t coreCount) {
     Cost schedulingCost = findMaximumIndexedPackCost(
-      model.schedulingCosts, taskWeightIds, indexedWeights.size(),
+      model.schedulingCosts, taskWeightIds, model.weightCount,
       tasksByStage[stage], residentWeightCapacity, coreCount);
+    schedulingCost = checkedAdd(schedulingCost, stageTailCosts[stage]);
     if (&balanceCosts == &model.schedulingCosts)
       return schedulingCost;
     Cost assemblyCost = findMaximumIndexedPackCost(
-      balanceCosts, taskWeightIds, indexedWeights.size(),
+      balanceCosts, taskWeightIds, model.weightCount,
       tasksByStage[stage], residentWeightCapacity, coreCount);
     return std::max(schedulingCost, assemblyCost);
   };
   std::vector<Cost> stageCosts(stageSizes.size());
   std::vector<Cost> nextStageCosts(stageSizes.size());
+  TaskCosts stageWork(stageSizes.size());
+  for (size_t stage = 0; stage < stageSizes.size(); ++stage)
+    for (size_t task : tasksByStage[stage])
+      stageWork[stage] = checkedAdd(stageWork[stage], model.schedulingCosts[task]);
   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);
+  bool canPlaceTasksSeparately = graph.nodes.size() <= layout.getProcessorCount();
   while (assignedCores < layout.getProcessorCount()) {
-    size_t bestStage = 0;
-    Cost bestBenefit = 0;
+    size_t bestStage = stageSizes.size();
+    std::optional<std::tuple<Cost, Cost, Cost, size_t>> bestScore;
     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])))) {
+      Cost workPerCore = stageWork[stage] / stageSizes[stage];
+      auto score = std::make_tuple(benefit,
+        benefit == 0 && canPlaceTasksSeparately ? workPerCore : 0, stageCosts[stage],
+        layout.getProcessorCount() - stageSizes[stage]);
+      if (!bestScore || score > *bestScore) {
         bestStage = stage;
-        bestBenefit = benefit;
+        bestScore = score;
       }
     }
     ++stageSizes[bestStage];
@@ -1006,7 +1068,7 @@ static bool packPipelineStage(
     ArrayRef<size_t> topologicalPosition, size_t stage,
     size_t residentWeightCapacity, const SchedulingTarget &target,
     std::vector<size_t> &taskCpus, std::vector<Time> &taskEndTimes,
-    bool prioritizeWeightReuse = false) {
+    bool balanceWork = false, ArrayRef<size_t> taskPacks = {}) {
   PipelineStageRange range = layout.getStageRange(stage);
   TaskList tasks;
   for (size_t cpu = range.begin; cpu < range.begin + range.size; ++cpu)
@@ -1019,12 +1081,16 @@ static bool packPipelineStage(
   TaskCosts loads(range.size);
   TaskCosts assemblyLoads(range.size);
   std::vector<Time> coreReady(range.size);
-  for (size_t task : tasks) {
+  assert((taskPacks.empty() || taskPacks.size() == tasks.size())
+    && "pipeline partition must cover every stage task");
+  for (auto [taskIndex, task] : llvm::enumerate(tasks)) {
     std::optional<size_t> bestCore;
     Time bestEndTime = 0;
     using PackScore = std::tuple<Cost, Time, Cost, Cost, size_t>;
     std::optional<PackScore> bestScore;
     for (size_t core = 0; core < range.size; ++core) {
+      if (!taskPacks.empty() && taskPacks[taskIndex] != core)
+        continue;
       size_t unionSize = getResidentWeightUnionSize(
         weights[core], graph.nodes[task].residentWeights);
       if (unionSize > residentWeightCapacity)
@@ -1035,6 +1101,7 @@ static bool packPipelineStage(
       Cost schedulingLoad = checkedAdd(
         loads[core], schedulingCosts[task]);
       Time startTime = coreReady[core];
+      Time transferTime = 0;
       size_t candidateCpu = layout.getPhysicalCore(range.begin + core);
       // Only same-stage edges use direct sends/receives. Pipeline boundaries
       // use host buffers and belong to different samples in steady state.
@@ -1044,15 +1111,14 @@ static bool packPipelineStage(
             && layout.getStageForCore(taskCpus[predecessor]) == stage) {
           Time transfer = getPeftTransferTime(
             transferCost, taskCpus[predecessor], candidateCpu, target);
+          transferTime = addOrMax(transferTime, transfer);
           startTime = std::max(
             startTime, addOrMax(taskEndTimes[predecessor], transfer));
         }
       Time endTime = addOrMax(startTime, schedulingCosts[task]);
-      PackScore score = prioritizeWeightReuse
-        ? PackScore {
-            addedWeights, endTime, assemblyLoad, schedulingLoad, core}
-        : PackScore {
-            endTime, assemblyLoad, schedulingLoad, addedWeights, core};
+      PackScore score {
+        balanceWork ? addOrMax(schedulingLoad, transferTime) : endTime,
+        assemblyLoad, schedulingLoad, addedWeights, core};
       if (!bestScore || score < *bestScore) {
         bestCore = core;
         bestScore = score;
@@ -1115,21 +1181,59 @@ static LogicalResult packPipelineStages(
   const TaskCosts &balanceCosts = getPipelineBalanceCosts(graph, model);
   std::vector<size_t> taskCpus(graph.nodes.size(), target.processorCount);
   std::vector<Time> taskEndTimes(graph.nodes.size());
-  for (size_t stage = 0; stage < pipelineStages; ++stage)
-    if (!packPipelineStage(
-          graph, model.schedulingCosts, balanceCosts, tasksByCpu,
-          layout, topologicalPosition, stage, residentWeightCapacity,
-          target, taskCpus, taskEndTimes)
-        && !packPipelineStage(
-          graph, model.schedulingCosts, balanceCosts, tasksByCpu,
-          layout, topologicalPosition, stage, residentWeightCapacity,
-          target, taskCpus, taskEndTimes, /*prioritizeWeightReuse=*/true)) {
+  for (size_t stage = 0; stage < pipelineStages; ++stage) {
+    auto balancedTasks = tasksByCpu;
+    auto balancedCpus = taskCpus;
+    auto balancedEnds = taskEndTimes;
+    bool balanced = packPipelineStage(
+      graph, model.schedulingCosts, balanceCosts, balancedTasks,
+      layout, topologicalPosition, stage, residentWeightCapacity,
+      target, balancedCpus, balancedEnds, true);
+    bool packed = packPipelineStage(
+      graph, model.schedulingCosts, balanceCosts, tasksByCpu,
+      layout, topologicalPosition, stage, residentWeightCapacity,
+      target, taskCpus, taskEndTimes);
+    auto period = [&](const std::vector<TaskList> &assignment) {
+      Cost maximum = 0;
+      PipelineStageRange range = layout.getStageRange(stage);
+      for (size_t index = range.begin; index < range.begin + range.size; ++index) {
+        Cost work = 0;
+        for (size_t task : assignment[layout.getPhysicalCore(index)])
+          work = checkedAdd(work, balanceCosts[task]);
+        maximum = std::max(maximum, work);
+      }
+      return maximum;
+    };
+    if (balanced && (!packed || period(balancedTasks) < period(tasksByCpu))) {
+      tasksByCpu = std::move(balancedTasks);
+      taskCpus = std::move(balancedCpus);
+      taskEndTimes = std::move(balancedEnds);
+      packed = true;
+    }
+    TaskList taskPacks;
+    if (!packed) {
+      TaskList tasks;
+      PipelineStageRange range = layout.getStageRange(stage);
+      for (size_t index = range.begin; index < range.begin + range.size; ++index)
+        llvm::append_range(tasks, tasksByCpu[layout.getPhysicalCore(index)]);
+      llvm::sort(tasks, [&](size_t lhs, size_t rhs) {
+        return topologicalPosition[lhs] < topologicalPosition[rhs];
+      });
+      findMaximumIndexedPackCost(model.schedulingCosts, model.weightIds,
+        model.weightCount, tasks, residentWeightCapacity, range.size, &taskPacks);
+      packed = taskPacks.size() == tasks.size() && packPipelineStage(
+        graph, model.schedulingCosts, balanceCosts, tasksByCpu,
+        layout, topologicalPosition, stage, residentWeightCapacity,
+        target, taskCpus, taskEndTimes, false, taskPacks);
+    }
+    if (!packed) {
       failedStage = stage;
       error = "pipeline scheduling cannot pack dependency-monotone stage "
             + std::to_string(stage)
             + " within the physical crossbar limit";
       return failure();
     }
+  }
   return success();
 }