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

test/PIM/SpatialSchedulingTargetTest.cpp

@@ -4,14 +4,165 @@
 #include <vector>
 
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
 
 using namespace onnx_mlir::spatial;
 
+static void testResidentWeightViewCost() {
+  mlir::MLIRContext context;
+  context.loadDialect<SpatialDialect, mlir::tensor::TensorDialect>();
+  mlir::Block operands;
+  mlir::OwningOpRef<mlir::ModuleOp> owner =
+    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  mlir::OpBuilder builder(&context);
+  auto loc = owner->getLoc();
+  auto weightType = mlir::RankedTensorType::get({128, 128}, builder.getF32Type());
+  auto inputType = mlir::RankedTensorType::get({1, 128}, builder.getF32Type());
+  auto weight = operands.addArgument(weightType, loc);
+  auto input = operands.addArgument(inputType, loc);
+  auto wideType = mlir::RankedTensorType::get({1, 256}, builder.getF32Type());
+  auto wideInput = operands.addArgument(wideType, loc);
+  auto makeCompute = [&](bool weightView, unsigned inputView) {
+    builder.setInsertionPointToEnd(owner->getBody());
+    auto compute = SpatGraphCompute::create(builder, loc,
+      mlir::TypeRange{inputType}, mlir::ValueRange{weight},
+      mlir::ValueRange{inputView == 2 ? wideInput : input});
+    auto *block = new mlir::Block();
+    compute.getBody().push_back(block);
+    mlir::Value w = block->addArgument(weightType, loc);
+    mlir::Value x = block->addArgument(inputView == 2 ? wideType : inputType, loc);
+    builder.setInsertionPointToStart(block);
+    auto slice = [&](mlir::Value value, mlir::RankedTensorType type, bool strided = false) {
+      llvm::SmallVector<mlir::OpFoldResult> offsets(2, builder.getIndexAttr(0));
+      llvm::SmallVector<mlir::OpFoldResult> strides(2, builder.getIndexAttr(1));
+      if (strided)
+        strides.back() = builder.getIndexAttr(2);
+      llvm::SmallVector<mlir::OpFoldResult> sizes;
+      for (int64_t size : type.getShape())
+        sizes.push_back(builder.getIndexAttr(size));
+      return mlir::tensor::ExtractSliceOp::create(
+        builder, loc, type, value, offsets, sizes, strides).getResult();
+    };
+    if (weightView)
+      w = slice(w, weightType);
+    if (inputView)
+      x = slice(x, inputType, inputView == 2);
+    auto result = SpatVMMOp::create(builder, loc, inputType, w, x);
+    SpatYieldOp::create(builder, loc, mlir::ValueRange{result.getResult()});
+    return ComputeInstance{compute, 0, 1};
+  };
+  auto direct = makeCompute(false, false);
+  auto weightView = makeCompute(true, false);
+  auto inputView = makeCompute(false, true);
+  auto stridedInput = makeCompute(false, 2);
+  for (Cost readLatency : {1, 20}) {
+    SchedulingTarget target;
+    target.localMemoryReadLatencyCycles = readLatency;
+    target.localMemoryWriteLatencyCycles = 3;
+    target.localMemoryWidthBytes = 8;
+    target.processorPeriodNs = 1;
+    target.computeBitwidth = 8;
+    assert(getComputeInstanceCost(weightView, target) == getComputeInstanceCost(direct, target));
+    assert(getComputeInstanceCost(inputView, target) == getComputeInstanceCost(direct, target));
+    assert(getComputeInstanceCost(stridedInput, target)
+      == getComputeInstanceCost(direct, target) + 64 * readLatency);
+  }
+}
+
+static void testPipelineStageSpan() {
+  using namespace mlir;
+  MLIRContext context;
+  context.loadDialect<SpatialDialect>();
+  OwningOpRef<ModuleOp> owner = ModuleOp::create(UnknownLoc::get(&context));
+  OpBuilder builder(owner->getBodyRegion());
+  ComputeGraph graph;
+  graph.successors.resize(12);
+  graph.predecessors.resize(12);
+  for (size_t group = 0; group < 6; ++group) {
+    auto batch = SpatGraphComputeBatch::create(builder, owner->getLoc(),
+      TypeRange{}, builder.getI32IntegerAttr(2), ValueRange{}, ValueRange{});
+    for (uint32_t lane = 0; lane < 2; ++lane) {
+      size_t task = graph.nodes.size();
+      ComputeInstance instance{batch, lane, 1};
+      graph.nodes.push_back({instance, group == 0 ? 40ULL : 10ULL, {}, task});
+      graph.instanceToIndex[instance] = task;
+      if (group != 0)
+        for (size_t predecessor = task - lane - 2; predecessor < task - lane; ++predecessor) {
+          graph.predecessors[task].push_back({predecessor, {}});
+          graph.successors[predecessor].push_back({task, {}});
+        }
+    }
+  }
+  for (size_t coreCount : {6, 12, 20}) {
+    SchedulingTarget target;
+    target.processorCount = coreCount;
+    target.residentWeightCapacity = 1;
+    target.interProcessorLatencyNs.assign(coreCount * coreCount, 0);
+    auto schedule = runPeftScheduler(graph, {target, &context});
+    std::string error;
+    assert(succeeded(applyPipelineScheduling(graph, schedule, 2, target, error)));
+    Cost stageCosts[2] = {};
+    Cost stageWork[2] = {};
+    for (size_t task = 0; task < graph.nodes.size(); ++task) {
+      size_t core = schedule.computeToCpuMap.lookup(graph.nodes[task].instance);
+      size_t stage = schedule.processorStages[core];
+      stageWork[stage] += graph.nodes[task].cost;
+      if (task % 2 == 0)
+        stageCosts[stage] += graph.nodes[task].cost;
+    }
+    // Both 40/50 and 50/40 are optimal cuts for this 90-cycle chain.
+    assert(std::max(stageCosts[0], stageCosts[1]) == 50);
+    if (coreCount == 20)
+      for (size_t stage = 0; stage < 2; ++stage)
+        assert(stageWork[stage] <= 10 * static_cast<size_t>(
+          llvm::count(schedule.processorStages, stage)));
+  }
+}
+
+static void testPipelinePackingCapacity() {
+  using namespace mlir;
+  MLIRContext context;
+  context.loadDialect<SpatialDialect>();
+  OwningOpRef<ModuleOp> owner = ModuleOp::create(UnknownLoc::get(&context));
+  OpBuilder builder(owner->getBodyRegion());
+  auto batch = SpatGraphComputeBatch::create(builder, owner->getLoc(),
+    TypeRange{}, builder.getI32IntegerAttr(4), ValueRange{}, ValueRange{});
+  ComputeGraph graph;
+  graph.successors.resize(4);
+  graph.predecessors.resize(4);
+  for (uint32_t task = 0; task < 4; ++task) {
+    ResidentWeightSet weights;
+    for (uint32_t index = 0; index < (task == 3 ? 3u : 1u); ++index) {
+      ResidentWeight weight;
+      weight.opaqueLane = task == 3 ? index + 1 : 0;
+      weights.push_back(weight);
+    }
+    ComputeInstance instance{batch, task, 1};
+    graph.nodes.push_back({instance, 10, weights, task});
+    graph.instanceToIndex[instance] = task;
+  }
+  SchedulingTarget target;
+  target.processorCount = 4;
+  target.residentWeightCapacity = 3;
+  target.interProcessorLatencyNs.assign(16, 0);
+  auto schedule = runPeftScheduler(graph, {target, &context});
+  std::string error;
+  assert(succeeded(applyPipelineScheduling(graph, schedule, 2, target, error)));
+  std::vector<Cost> loads(4);
+  for (const auto &node : graph.nodes) {
+    size_t core = schedule.computeToCpuMap.lookup(node.instance);
+    assert(schedule.processorStages[core] == 0);
+    loads[core] += node.cost;
+  }
+  assert(*std::max_element(loads.begin(), loads.end()) == 20);
+}
+
 static void testPhysicalPeft() {
   TransferCost transfer {.fixed = 50, .networkFlits = 4};
   SchedulingTarget line;
@@ -51,7 +202,7 @@ static void testPhysicalPeft() {
       chain.predecessors[task].push_back({task - 1, transfer});
     }
   }
-  line.residentWeightCapacity = reversed.residentWeightCapacity = 3;
+  line.residentWeightCapacity = reversed.residentWeightCapacity = 1;
   for (const SchedulingTarget &target : {line, reversed}) {
     auto scheduled = runPeftScheduler(chain, {target, &context});
     size_t firstCore = scheduled.computeToCpuMap.lookup(chain.nodes[0].instance);
@@ -69,6 +220,58 @@ static void testPhysicalPeft() {
   }
 }
 
+static void testPeftReservationCapacity() {
+  mlir::MLIRContext context;
+  auto owner = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  ComputeGraph graph;
+  graph.successors.resize(4);
+  graph.predecessors.resize(4);
+  for (uint32_t task = 0; task < 4; ++task) {
+    ResidentWeightSet weights;
+    for (uint32_t index = 0; index < (task == 3 ? 4u : 1u); ++index) {
+      ResidentWeight weight;
+      weight.opaqueLane = task == 3 ? index + 1 : 0;
+      weights.push_back(weight);
+    }
+    graph.nodes.push_back({{owner, task, 1}, 10, weights, task});
+  }
+  SchedulingTarget target;
+  target.processorCount = 3;
+  target.residentWeightCapacity = 4;
+  target.interProcessorLatencyNs.assign(9, 0);
+  auto schedule = runPeftScheduler(graph, {target, &context});
+  size_t largeCore = schedule.computeToCpuMap.lookup(graph.nodes[3].instance);
+  for (size_t task : {0, 1, 2})
+    assert(schedule.computeToCpuMap.lookup(graph.nodes[task].instance) != largeCore);
+  owner.erase();
+}
+
+static void testPeftGapOrder() {
+  mlir::MLIRContext context;
+  mlir::OwningOpRef<mlir::ModuleOp> owner =
+    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  ComputeGraph graph;
+  graph.successors.resize(3);
+  graph.predecessors.resize(3);
+  for (uint32_t task = 0; task < 3; ++task) {
+    ResidentWeight weight;
+    weight.opaqueLane = task == 0 ? 0 : 1;
+    graph.nodes.push_back({{owner->getOperation(), task, 1},
+      task == 2 ? 20u : 10u, {weight}, task});
+  }
+  TransferCost transfer {.fixed = 100};
+  graph.successors[0].push_back({1, transfer});
+  graph.predecessors[1].push_back({0, transfer});
+  SchedulingTarget target;
+  target.processorCount = 2;
+  target.residentWeightCapacity = 1;
+  target.interProcessorLatencyNs.assign(4, 0);
+  auto schedule = runPeftScheduler(graph, {target, &context});
+  assert(schedule.computeToCpuSlotMap.lookup(graph.nodes[2].instance) == 0);
+  assert(schedule.computeToCpuSlotMap.lookup(graph.nodes[1].instance) == 1);
+  assert(schedule.computeToAestMap.lookup(graph.nodes[2].instance) == 0);
+}
+
 static void testPipelineInputArrival() {
   ComputeGraph graph;
   graph.successors.resize(4);
@@ -146,9 +349,39 @@ static void testReadyCommunication() {
   assert(scheduled->slices.size() == 2);
   for (const auto &slice : scheduled->slices)
     assert(slice.sourceInsertionStep == 1);
+
+  plan.stepCounts = {3, 3, 3};
+  for (auto &exchange : plan.exchanges)
+    exchange->consumerStep = 2;
+  scheduled = scheduleDeferredCommunication({}, plan);
+  assert(mlir::succeeded(scheduled));
+  for (const auto &slice : scheduled->slices)
+    assert(slice.targetInsertionStep == 2);
+}
+
+static void testPipelineRestartOrder() {
+  mlir::MLIRContext context;
+  auto owner = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  DeferredTransferPlan transfers;
+  for (int64_t core : {0, 7, 3, 5}) {
+    ScheduledInfo info;
+    info.op = owner;
+    info.cores.push_back(core);
+    info.pipelineStages.push_back(core == 0 ? 0 : 1);
+    transfers.scheduled.push_back(info);
+  }
+  auto boundaries = buildDeferredBoundaryPlan(transfers, {}, 16);
+  assert(mlir::succeeded(boundaries));
+  // Restarting the producer must not wait for a later consumer on core 3.
+  assert((transfers.downstreamCores == llvm::SmallVector<int64_t>{7, 3, 5}));
+  owner.erase();
 }
 
 int main() {
+  testResidentWeightViewCost();
+  testPipelineStageSpan();
+  testPipelinePackingCapacity();
+  testPipelineRestartOrder();
   testPipelineInputArrival();
   testReadyCommunication();
   PipelineCoreLayout unevenLayout(138, 4);
@@ -198,6 +431,8 @@ int main() {
   assert(slow.getInterProcessorLatencyNs(0, 1) == 10);
 
   testPhysicalPeft();
+  testPeftGapOrder();
+  testPeftReservationCapacity();
 
   ComputeGraph graph;
   graph.successors.resize(6);