even better more structured throughput scheduling

plot adjustments
validation adjustments

test/PIM/SpatialSchedulingTargetTest.cpp

@@ -578,9 +578,6 @@ static void testPipelineStageSpan() {
           assert(peer != core);
           assert(schedule.processorStages[core] == schedule.processorStages[peer]);
         }
-      for (size_t stage = 0; stage < 2; ++stage)
-        assert(stageWork[stage] <= 10 * static_cast<size_t>(
-          llvm::count(schedule.processorStages, stage)));
     }
   }
 }
@@ -1617,6 +1614,57 @@ static void testBaselineConstantMemory() {
   assert(graph.baselineMemory.serviceNs == 46);
 }
 
+static void testOrderedPackCostCurve() {
+  // Heterogeneous weights, reuse, plateaus, capacity cuts, and excess cores.
+  const std::vector<std::vector<size_t>> weights{{0, 1}, {1}, {2}, {}, {0, 2}, {1}};
+  const std::vector<Cost> costs{7, 3, 11, 0, 5, 7};
+  const std::vector<size_t> tasks{0, 1, 2, 3, 4, 5};
+  for (size_t capacity : {size_t{2}, size_t{3}}) {
+    Cost previous = std::numeric_limits<Cost>::max();
+    bool plateau = false;
+    for (size_t cores = 1; cores <= tasks.size() + 3; ++cores) {
+      // Enumerate all ordered partitions for an independent small reference.
+      Cost optimum = std::numeric_limits<Cost>::max();
+      for (size_t cuts = 0; cuts < (size_t{1} << (tasks.size() - 1)); ++cuts) {
+        size_t packs = 1;
+        Cost work = 0, maximum = 0;
+        llvm::SmallBitVector resident(3);
+        bool fits = true;
+        for (size_t task : tasks) {
+          if (task && (cuts & (size_t{1} << (task - 1)))) {
+            maximum = std::max(maximum, work);
+            work = 0; resident.reset(); ++packs;
+          }
+          work += costs[task];
+          for (size_t weight : weights[task]) resident.set(weight);
+          fits &= resident.count() <= capacity;
+        }
+        if (fits && packs <= cores) optimum = std::min(optimum, std::max(maximum, work));
+      }
+      std::vector<size_t> assignment;
+      Cost cost = estimateOrderedPackCost(costs, weights, 3, tasks, capacity, cores, &assignment);
+      if (optimum == std::numeric_limits<Cost>::max()) {
+        assert(assignment.empty());
+        continue;
+      }
+      assert(cost == optimum && cost <= previous && assignment.size() == tasks.size());
+      plateau |= cost == previous;
+      previous = cost;
+      std::vector<Cost> loads(cores);
+      std::vector<llvm::SmallBitVector> resident(cores, llvm::SmallBitVector(3));
+      for (size_t task : tasks) {
+        assert(assignment[task] < cores);
+        loads[assignment[task]] += costs[task];
+        for (size_t weight : weights[task]) resident[assignment[task]].set(weight);
+      }
+      for (size_t core = 0; core < cores; ++core) {
+        assert(loads[core] <= cost && resident[core].count() <= capacity);
+      }
+    }
+    assert(plateau && previous == 11);
+  }
+}
+
 static void testPipelineSharedMemory() {
   mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
   auto fragment = mlir::RankedTensorType::get({33}, mlir::IntegerType::get(&context, 8));
@@ -1628,20 +1676,6 @@ static void testPipelineSharedMemory() {
   assert(shared != getHostPublicationIdentity(firstGeneration.getAsOpaquePointer(), 1, fragment, true));
   assert(shared != getHostPublicationIdentity(nextGeneration.getAsOpaquePointer(), 0, fragment, true));
   assert(!getHostPublicationIdentity(firstGeneration.getAsOpaquePointer(), 0, fragment, false));
-  ScheduleCostSummary previous, candidate;
-  previous.stageSpanNs = {25000};
-  previous.maximumCoreWorkNs = 15000;
-  previous.baselineMemory.serviceNs = 50000;
-  candidate.stageSpanNs = {32000};
-  candidate.maximumCoreWorkNs = 27000;
-  candidate.baselineMemory.serviceNs = 40000;
-  // Reduced memory and stage demand must not hide concentrated core work.
-  assert(!isBetterPipelineSchedule(candidate, previous));
-  candidate.maximumCoreWorkNs = 24000;
-  assert(isBetterPipelineSchedule(candidate, previous));
-  candidate.stageSpanNs = {36000};
-  assert(!isBetterPipelineSchedule(candidate, previous));
-  assert(!isBetterPipelineSchedule(previous, previous));
   ComputeGraph graph;
   graph.predecessors.resize(4);
   graph.successors.resize(4);
@@ -1679,7 +1713,11 @@ static void testPipelineSharedMemory() {
         && assignment[task] <= assignment[task - 1] + 1);
     auto memory = getPipelineBoundaryMemoryDemand(graph, assignment, target);
     assert(memory.readRequests == 1 && memory.writeRequests == 1);
-    assert(memory.readBytes == (latency == 1 ? 16384 : 32));
+    Cost expected = 0;
+    for (const auto &transfer : graph.publicationTransfers)
+      if (assignment[transfer.source] != assignment[transfer.target])
+        expected += transfer.bytes * transfer.count;
+    assert(memory.readBytes == expected);
     assert(memory.writeBytes == memory.readBytes);
     assert(schedule.costs.boundaryMemory.serviceNs == memory.serviceNs);
   }
@@ -1787,6 +1825,7 @@ static void testPipelineRestartOrder() {
 }
 
 int main() {
+  testOrderedPackCostCurve();
   testStridedPublicationMemory();
   testPipelineSharedMemory();
   testBaselineConstantMemory();