fix failing pipeline runs

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

@@ -638,7 +638,8 @@ static FailureOr<PipelineWorkloadPreparation> preparePipelineWorkloadImpl(
     if (tasks.empty())
       continue;
     Cost maximumCost = findMaximumPackCost(
-      graph, taskCosts, tasks, physicalTarget.residentWeightCapacity, pipelineStages);
+      graph, taskCosts, tasks, physicalTarget.residentWeightCapacity,
+      pipelineStages);
     if (fits(graph, taskCosts, tasks, maximumCost,
              physicalTarget.residentWeightCapacity, pipelineStages))
       continue;
@@ -653,7 +654,7 @@ static FailureOr<PipelineWorkloadPreparation> preparePipelineWorkloadImpl(
     });
     std::string candidateError;
     for (size_t task : candidates) {
-      auto batch = dyn_cast<SpatGraphComputeBatch>(
+      auto batch = dyn_cast_or_null<SpatGraphComputeBatch>(
         graph.nodes[task].instance.op);
       if (!batch || batch->hasAttr("pipeline.split"))
         continue;
@@ -664,10 +665,16 @@ static FailureOr<PipelineWorkloadPreparation> preparePipelineWorkloadImpl(
       if (!currentError.empty())
         candidateError = currentError;
     }
-    error = candidateError.empty()
-      ? "pipeline scheduling could not find a splittable graph compute"
-      : candidateError;
-    return failure();
+    if (llvm::any_of(tasks, [&](size_t task) {
+          return graph.nodes[task].residentWeights.size()
+               > physicalTarget.residentWeightCapacity;
+        })) {
+      error = candidateError.empty()
+        ? "pipeline scheduling cannot split a compute instance to fit one "
+          "physical core's crossbars"
+        : candidateError;
+      return failure();
+    }
   }
 
   return PipelineWorkloadPreparation::Ready;
@@ -1139,7 +1146,7 @@ static bool packPipelineStage(
     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) {
+    std::vector<size_t> &taskCpus, bool prioritizeWeightReuse = false) {
   PipelineStageRange range = layout.getStageRange(stage);
   TaskList tasks;
   for (size_t cpu = range.begin; cpu < range.begin + range.size; ++cpu)
@@ -1174,8 +1181,11 @@ static bool packPipelineStage(
             transferTime, getPeftTransferTime(
                             transferCost, taskCpus[predecessor],
                             candidateCpu, target));
-      PackScore score {
-        assemblyLoad, transferTime, schedulingLoad, addedWeights, core};
+      PackScore score = prioritizeWeightReuse
+        ? PackScore {
+            addedWeights, assemblyLoad, transferTime, schedulingLoad, core}
+        : PackScore {
+            assemblyLoad, transferTime, schedulingLoad, addedWeights, core};
       if (!bestScore || score < *bestScore) {
         bestCore = core;
         bestScore = score;
@@ -1238,7 +1248,11 @@ static LogicalResult packPipelineStages(
     if (!packPipelineStage(
           graph, model.schedulingCosts, balanceCosts, tasksByCpu,
           layout, topologicalPosition, stage, residentWeightCapacity,
-          target, taskCpus)) {
+          target, taskCpus)
+        && !packPipelineStage(
+          graph, model.schedulingCosts, balanceCosts, tasksByCpu,
+          layout, topologicalPosition, stage, residentWeightCapacity,
+          target, taskCpus, /*prioritizeWeightReuse=*/true)) {
       failedStage = stage;
       error = "pipeline scheduling cannot pack dependency-monotone stage "
             + std::to_string(stage)

test/PIM/SpatialSchedulingTargetTest.cpp

@@ -130,6 +130,14 @@ int main() {
     3, 3, 3, 0,
   };
   std::string pipelineError;
+  ComputeGraph preparationGraph = graph;
+  ResidentWeight extraWeight;
+  extraWeight.opaqueLane = graph.nodes.size();
+  preparationGraph.nodes[2].residentWeights.push_back(extraWeight);
+  auto preparation = preparePipelineWorkload(
+    preparationGraph, logicalSchedule, 2, physical, pipelineError);
+  assert(mlir::succeeded(preparation));
+  assert(*preparation == PipelineWorkloadPreparation::Ready);
   ComputeGraph emptyGraph;
   MergeScheduleResult emptySchedule;
   emptySchedule.processorCount = 2;