diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp index 1639001..9677454 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp @@ -638,7 +638,8 @@ static FailureOr 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 preparePipelineWorkloadImpl( }); std::string candidateError; for (size_t task : candidates) { - auto batch = dyn_cast( + auto batch = dyn_cast_or_null( graph.nodes[task].instance.op); if (!batch || batch->hasAttr("pipeline.split")) continue; @@ -664,10 +665,16 @@ static FailureOr 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 &tasksByCpu, const PipelineCoreLayout &layout, ArrayRef topologicalPosition, size_t stage, size_t residentWeightCapacity, const SchedulingTarget &target, - std::vector &taskCpus) { + std::vector &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) diff --git a/test/PIM/SpatialSchedulingTargetTest.cpp b/test/PIM/SpatialSchedulingTargetTest.cpp index 06f6166..16d121a 100644 --- a/test/PIM/SpatialSchedulingTargetTest.cpp +++ b/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;