force selected pipeline stages in throughput mode

increase throughput scheduler budget
add experimental throughput auto mode

test/PIM/PipelineOptimizerTest.cpp

@@ -4,10 +4,14 @@
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ResidentWeightPacking.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelinePackingCostCache.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include <chrono>
+#include <cmath>
 #include <functional>
 #include <map>
+#include <numeric>
+#include <random>
 #include <set>
 #include <iostream>
 
@@ -16,6 +20,41 @@ using namespace onnx_mlir::spatial;
 using E = PimExecutionEvent;
 static bool sawJointMove = false, sawSuccessiveMoves = false;
 
+static void testPackingCostCacheEviction() {
+  ComputeGraph graph;
+  const std::vector<Cost> work{0, 11, 23, 7, 19, 3};
+  for (uint32_t task = 0; task < work.size(); ++task) {
+    ResidentWeight weight;
+    weight.opaqueLane = task % 3;
+    graph.nodes.push_back({{nullptr, task, 1}, work[task], {weight}, task});
+  }
+  IndexedResidentWeights weights(graph);
+  PipelinePackingCostCache cache(weights, work, 1, 8);
+  static_assert(PipelinePackingCostCache::kMaximumEntries == 256);
+  std::vector<size_t> tasks{0, 1, 2, 3, 4, 5};
+  auto check = [&](const auto &order) {
+    TEST_CHECK(cache.minimumCores(order) == weights.getPackCountLowerBound(order, 1));
+    for (size_t cores : {1u, 2u, 3u, 5u, 8u}) {
+      Cost expected = weights.estimatePackCost(order, work, 1, cores)
+        .value_or(std::numeric_limits<Cost>::max());
+      TEST_CHECK(cache.estimate(order, cores) == expected);
+      TEST_CHECK(cache.estimate(order, cores) == expected);
+    }
+    TEST_CHECK(cache.size() <= PipelinePackingCostCache::kMaximumEntries);
+  };
+  size_t orders = 0;
+  do {
+    check(tasks);
+    ++orders;
+  } while (std::next_permutation(tasks.begin(), tasks.end()));
+  TEST_CHECK(orders == 720 && cache.size() == PipelinePackingCostCache::kMaximumEntries);
+  // Revisit early entries after eviction. Ordered and constructive queries must agree.
+  check(std::vector<size_t>{0, 1, 2, 3, 4, 5});
+  check(std::vector<size_t>{0, 3, 1, 4, 2, 5});
+  check(std::vector<size_t>{0});
+  check(std::vector<size_t>{});
+}
+
 static void testPackingFeasibility() {
   ComputeGraph graph;
   for (size_t task = 0; task < 4; ++task) {
@@ -50,6 +89,51 @@ static void testPackingFeasibility() {
   TEST_CHECK(weights.estimatePackCost(tasks, work, 2, 2) == 20);
 }
 
+static void testPackingCurves() {
+  std::mt19937 random(1739);
+  for (size_t fixture = 0; fixture < 150; ++fixture) {
+    size_t count = fixture % 9, capacity = 1 + fixture % 3;
+    ComputeGraph graph;
+    std::vector<Cost> work;
+    std::vector<size_t> tasks;
+    for (uint32_t task = 0; task < count; ++task) {
+      ResidentWeightSet weights;
+      for (size_t id = 0; id < 4; ++id)
+        if (random() % 3 == 0) { ResidentWeight w; w.opaqueLane = id; weights.push_back(w); }
+      work.push_back(random() % 40);
+      tasks.push_back(task);
+      graph.nodes.push_back({{nullptr, task, 1}, work.back(), weights, task});
+    }
+    IndexedResidentWeights weights(graph);
+    auto curve = weights.estimatePackCostCurve(tasks, work, capacity, 8);
+    auto sparse = weights.estimatePackCostCurve(tasks, work, capacity, 8, nullptr, {1, 3, 6, 8});
+    for (size_t cores : {1u, 3u, 6u, 8u}) TEST_CHECK(sparse[cores] == curve[cores]);
+    for (size_t cores : {2u, 4u, 5u, 7u}) TEST_CHECK(!sparse[cores]);
+    for (size_t cores = 1; cores <= 8; ++cores) {
+      TEST_CHECK(curve[cores] == weights.estimatePackCost(tasks, work, capacity, cores));
+      std::optional<Cost> reference;
+      for (size_t mask = 0; mask < (size_t{1} << (count ? count - 1 : 0)); ++mask) {
+        size_t used = count ? 1 : 0;
+        Cost maximum = 0, current = 0;
+        std::set<size_t> residents;
+        bool fits = true;
+        for (size_t task = 0; task < count; ++task) {
+          if (task && ((mask >> (task - 1)) & 1)) {
+            maximum = std::max(maximum, current); current = 0; residents.clear(); ++used;
+          }
+          current += work[task];
+          for (size_t id : weights.taskWeightIds[task]) residents.insert(id);
+          fits &= residents.size() <= capacity;
+        }
+        maximum = std::max(maximum, current);
+        if (fits && used <= cores && (!reference || maximum < *reference)) reference = maximum;
+      }
+      TEST_CHECK(estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
+        tasks, capacity, cores) == reference);
+    }
+  }
+}
+
 static void testQuantizedFeasibility(uint32_t first) {
   mlir::MLIRContext context;
   context.loadDialect<spatial::SpatialDialect>();
@@ -83,6 +167,20 @@ static void testQuantizedFeasibility(uint32_t first) {
   for (size_t task = 0; task < graph.nodes.size(); ++task)
     TEST_CHECK(schedule.processorStages[schedule.computeToCpuMap.lookup(graph.nodes[task].instance)]
       == (task < first ? 0 : 1));
+  auto query = [](const MergeScheduleResult &) {
+    WholeScheduleCost cost;
+    cost.primary = cost.alternative = 100;
+    cost.selectionMargin = kPipelineSelectionMargin;
+    cost.loweringSupport = StaticLoweringSupport::Supported;
+    for (auto &proof : cost.recurrence) proof = {ScheduleRecurrenceEvidence::Proved, 1, 100, {}};
+    return cost;
+  };
+  auto depthOne = schedule;
+  depthOne.processorStages.assign(target.processorCount, 0);
+  auto search = optimizePipelineSchedule(graph, depthOne, 0, target, query, error, 2);
+  TEST_CHECK(mlir::succeeded(search) && search->initialSchedules == 2);
+  TEST_CHECK(search->schedules[1].processorStages == schedule.processorStages);
+  TEST_CHECK(search->schedules[1].computeToCpuMap == schedule.computeToCpuMap);
   module.erase();
 }
 
@@ -90,7 +188,8 @@ static void testQuantizedFeasibility(uint32_t first) {
 // The same queue
 // engine and complete-state evidence rank reference and searched schedules.
 static WholeScheduleCost queueCost(const ComputeGraph &graph,
-    const MergeScheduleResult &schedule, const PimInstructionTiming &target) {
+    const MergeScheduleResult &schedule, const PimInstructionTiming &target,
+    bool restartProtocol = false) {
   PimExecutionProgram program(target.processorCount);
   std::vector<std::vector<size_t>> tasks(target.processorCount);
   for (size_t task = 0; task < graph.nodes.size(); ++task)
@@ -115,10 +214,17 @@ static WholeScheduleCost queueCost(const ComputeGraph &graph,
       transfers.push_back({predecessor, task, reader, writer, 1 + transfers.size(), host});
       if (host) { readers[writer].insert(reader); writers[reader].insert(writer); }
     }
+  std::vector<size_t> zero, downstream;
+  if (restartProtocol)
+    for (size_t core = 0; core < program.size(); ++core)
+      (schedule.processorStages[core] ? downstream : zero).push_back(core);
   for (size_t core = 0; core < program.size(); ++core) {
     auto &events = program[core];
     for (size_t writer : writers[core]) events.emplace_back(E::Signal, 0, writer, 0, 1);
-    if (schedule.processorStages[core] == 0) {
+    size_t rank = std::find(downstream.begin(), downstream.end(), core) - downstream.begin();
+    if (restartProtocol && rank < downstream.size())
+      events.emplace_back(E::Signal, 0, rank ? downstream[(rank - 1) / 2] : zero.front(), rank ? 29 : 30, 1);
+    if (!restartProtocol && schedule.processorStages[core] == 0) {
       size_t peers = 0;
       for (size_t peer = 0; peer < program.size(); ++peer)
         if (peer != core && schedule.processorStages[peer] == 0) {
@@ -149,9 +255,31 @@ static WholeScheduleCost queueCost(const ComputeGraph &graph,
           events.emplace_back(E::Signal, 0, transfer.reader, transfer.reg, 1);
         }
     }
+    if (restartProtocol && schedule.processorStages[core] == 0) {
+      size_t rankZero = std::find(zero.begin(), zero.end(), core) - zero.begin();
+      for (size_t step = 1, reg = 16; step < zero.size(); step *= 2, ++reg) {
+        events.emplace_back(E::Signal, 0, zero[(rankZero + step) % zero.size()], reg, 1);
+        events.emplace_back(E::Wait, 0, 0, reg, 1);
+      }
+      if (core == zero.front() && !downstream.empty()) {
+        events.emplace_back(E::Wait, 0, 0, 30, 1);
+        events.emplace_back(E::Signal, 0, downstream.front(), 28, 1);
+      }
+    } else if (restartProtocol) {
+      events.emplace_back(E::Wait, 0, 0, 28, 1);
+      size_t children = (2 * rank + 1 < downstream.size()) + (2 * rank + 2 < downstream.size());
+      if (children) events.emplace_back(E::Wait, 0, 0, 29, children);
+      for (size_t child : {2 * rank + 1, 2 * rank + 2})
+        if (child < downstream.size()) events.emplace_back(E::Signal, 0, downstream[child], 28, 1);
+    }
   }
+  if (restartProtocol)
+    for (const auto &sequence : program)
+      for (const auto &event : sequence)
+        if (event.kind == E::Signal || event.kind == E::Wait)
+          TEST_CHECK(event.reg >= 0 && static_cast<size_t>(event.reg) < target.synchronizationRegisterCount);
   PimPeriodicOptions options;
-  options.verifyPeriodicState = true; options.searchGuidance = true;
+  options.verifyPeriodicState = true; options.searchGuidance = !restartProtocol;
   auto estimate = evaluatePimExecution(program, target, options);
   auto cost = std::move(estimate.cost);
   cost.loweringSupport = StaticLoweringSupport::Supported;
@@ -242,8 +370,9 @@ static void testInitializerPackingWitness() {
   }
   TEST_CHECK(feasible > 0);
   MergeScheduleResult alternative;
-  TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, std::vector<size_t>(9),
+  TEST_CHECK(mlir::failed(packPipelineSchedule(graph, std::vector<size_t>(9),
     {3, 1}, target, alternative, error)));
+  TEST_CHECK(error.find("nonempty compute stages") != std::string::npos);
   std::cout << "initializer_physical_feasible_assignments=" << feasible << '\n';
   module.erase();
 }
@@ -282,7 +411,8 @@ static void testPhysicalReference() {
         insertResidentWeights(resident[core], node.residentWeights);
         valid &= resident[core].size() <= target.residentWeightCapacity;
       }
-      if (!valid || stages[0] != 0 || stages[1] < stages[0] || stages[2] < stages[1]) continue;
+      if (!valid || stages[0] != 0 || stages[1] < stages[0] || stages[2] < stages[1]
+          || !llvm::is_contained(stages, 1)) continue;
       auto cost = queueCost(graph, schedule, target);
       optimum = std::min(optimum, std::max(*cost.primary, *cost.alternative));
       ++evaluations;
@@ -304,6 +434,198 @@ static void testPhysicalReference() {
     << " packing_attempts=" << search->packingAttempts << '\n';
 }
 
+static void testAutoPhysicalReference() {
+  for (size_t profile = 0; profile < 2; ++profile) {
+    auto target = hardware(profile);
+    target.processorCount = 3;
+    target.sharedMemoryRequestLatencyNs.resize(3); target.sharedMemoryResponseLatencyNs.resize(3);
+    target.interProcessorLatencyNs = {0, 2, 11, 2, 0, 7, 11, 7, 0};
+    ComputeGraph graph;
+    graph.predecessors.resize(3); graph.successors.resize(3);
+    for (uint32_t task = 0; task < 3; ++task) {
+      ResidentWeight weight; weight.opaqueLane = task % 2;
+      graph.nodes.push_back({{nullptr, task, 1}, task == 1 ? 700u : 80u, {weight}, task});
+      if (task) {
+        graph.predecessors[task].push_back({task - 1, {64, 1}});
+        graph.successors[task - 1].push_back({task, {64, 1}});
+      }
+    }
+    double reference = std::numeric_limits<double>::infinity();
+    size_t referenceDepth = 0, evaluations = 0;
+    auto start = std::chrono::steady_clock::now();
+    for (size_t depth = 1; depth <= 3; ++depth) {
+      size_t maps = depth * depth * depth;
+      for (size_t encodedMap = 0; encodedMap < maps; ++encodedMap)
+        for (size_t placement = 0; placement < 27; ++placement) {
+          MergeScheduleResult schedule;
+          size_t map = encodedMap, encoded = placement;
+          schedule.processorCount = 3;
+          for (size_t core = 0; core < 3; ++core) { schedule.processorStages.push_back(map % depth); map /= depth; }
+          std::vector<size_t> stages, slots(3);
+          std::vector<ResidentWeightSet> resident(3);
+          bool valid = true;
+          for (const auto &node : graph.nodes) {
+            size_t core = encoded % 3; encoded /= 3;
+            schedule.computeToCpuMap[node.instance] = core;
+            schedule.computeToCpuSlotMap[node.instance] = slots[core]++;
+            stages.push_back(schedule.processorStages[core]);
+            insertResidentWeights(resident[core], node.residentWeights);
+            valid &= resident[core].size() <= target.residentWeightCapacity;
+          }
+          std::set<size_t> active(stages.begin(), stages.end());
+          if (!valid || active.size() != depth || stages[0] != 0 || stages[1] < stages[0]
+              || stages[1] > stages[0] + 1 || stages[2] < stages[1] || stages[2] > stages[1] + 1) continue;
+          auto cost = queueCost(graph, schedule, target);
+          double value = std::max(*cost.primary, *cost.alternative);
+          if (value < reference) { reference = value; referenceDepth = depth; }
+          ++evaluations;
+        }
+    }
+    double referenceSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
+    MergeScheduleResult initial;
+    std::string error;
+    TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, std::vector<size_t>(3), {3}, target, initial, error)));
+    auto query = [&](const MergeScheduleResult &schedule) { return queueCost(graph, schedule, target); };
+    start = std::chrono::steady_clock::now();
+    auto search = optimizePipelineSchedule(graph, initial, 0, target, query, error);
+    TEST_CHECK(mlir::succeeded(search));
+    auto repeat = optimizePipelineSchedule(graph, initial, 0, target, query, error);
+    TEST_CHECK(mlir::succeeded(repeat) && search->costs.size() == repeat->costs.size());
+    size_t winner = selectPipelineCandidate(search->costs, error);
+    TEST_CHECK(winner == selectPipelineCandidate(repeat->costs, error));
+    for (size_t i = 0; i < search->schedules.size(); ++i) {
+      TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
+      TEST_CHECK(search->costs[i].primary == repeat->costs[i].primary);
+    }
+    double best = std::numeric_limits<double>::infinity();
+    for (const auto &cost : search->costs) best = std::min(best, std::max(*cost.primary, *cost.alternative));
+    double selected = std::max(*search->costs[winner].primary, *search->costs[winner].alternative);
+    TEST_CHECK(best + 1e-6 >= reference);
+    std::cout << "auto_physical_profile=" << profile << " reference=" << reference
+      << " reference_depth=" << referenceDepth << " best_evaluated=" << best << " selected=" << selected
+      << " selected_depth=" << getPipelineStageCount(search->schedules[winner])
+      << " gap_percent=" << 100 * (selected / reference - 1) << " reference_domain=all_physical_assignments_for_chain"
+      << " reference_evaluations=" << evaluations << " evaluations=" << search->costs.size()
+      << " packing_attempts=" << search->packingAttempts
+      << " initialization_seconds=" << search->segmentCostSeconds + search->partitionStateSeconds
+      << " reference_seconds=" << referenceSeconds
+      << " search_repeat_seconds=" << std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count() << '\n';
+  }
+}
+
+static void testNonunimodalPhysicalDepthCurve() {
+  TEST_CHECK(std::numeric_limits<double>::is_iec559 && std::numeric_limits<double>::digits == 53);
+  TEST_CHECK(std::numeric_limits<long double>::digits >= 64);
+  auto target = hardware(1);
+  target.processorCount = 4;
+  target.synchronizationRegisterCount = 32;
+  target.sharedMemoryRequestLatencyNs.assign(4, 2);
+  target.sharedMemoryResponseLatencyNs.assign(4, 3);
+  target.sharedMemoryReadLatencyNs = target.sharedMemoryWriteLatencyNs = 14;
+  // Manhattan distances at (0,0), (0,11), (18,0), and (5,7).
+  target.interProcessorLatencyNs = {0, 11, 18, 12, 11, 0, 29, 9, 18, 29, 0, 20, 12, 9, 20, 0};
+  ComputeGraph graph;
+  graph.predecessors.resize(5); graph.successors.resize(5);
+  const Cost work[] = {163, 46, 69, 131, 79};
+  for (uint32_t task = 0; task < 5; ++task) {
+    ResidentWeight weight; weight.opaqueLane = task % 2;
+    graph.nodes.push_back({{nullptr, task, 1}, work[task], {weight}, task});
+    if (task) {
+      graph.predecessors[task].push_back({task - 1, {64, 1}});
+      graph.successors[task - 1].push_back({task, {64, 1}});
+    }
+  }
+  std::vector<double> curve(5, std::numeric_limits<double>::infinity());
+  auto primary = curve, alternative = curve;
+  using Ratio = std::pair<uint64_t, uint64_t>;
+  std::vector<std::optional<Ratio>> certified(5);
+  auto certifiedPrimary = certified, certifiedAlternative = certified;
+  auto less = [](Ratio a, Ratio b) { return a.first * b.second < b.first * a.second; };
+  std::vector<std::vector<size_t>> winnerCores(5), winnerStages(5);
+  std::vector<size_t> evaluations(5);
+  for (size_t cuts = 0; cuts < 16; ++cuts) {
+    std::vector<size_t> stages{0};
+    for (size_t task = 1; task < 5; ++task)
+      stages.push_back(stages.back() + ((cuts >> (task - 1)) & 1));
+    size_t depth = stages.back() + 1;
+    if (depth > 4) continue;
+    for (size_t placement = 0; placement < 1024; ++placement) {
+      MergeScheduleResult schedule;
+      schedule.processorCount = 4;
+      size_t encoded = placement;
+      std::vector<size_t> assigned(4, depth), slots(4), idle;
+      bool valid = true;
+      for (size_t task = 0; task < 5; ++task) {
+        size_t core = encoded % 4; encoded /= 4;
+        if (assigned[core] != depth && assigned[core] != stages[task]) { valid = false; break; }
+        assigned[core] = stages[task];
+        schedule.computeToCpuMap[graph.nodes[task].instance] = core;
+        schedule.computeToCpuSlotMap[graph.nodes[task].instance] = slots[core]++;
+      }
+      if (!valid) continue;
+      // There are only two resident weights; both fit on every core.
+      for (size_t core = 0; core < 4; ++core) if (assigned[core] == depth) idle.push_back(core);
+      size_t maps = 1;
+      for (size_t i = 0; i < idle.size(); ++i) maps *= depth;
+      for (size_t map = 0; map < maps; ++map) {
+        schedule.processorStages = assigned;
+        size_t digits = map;
+        for (size_t core : idle) { schedule.processorStages[core] = digits % depth; digits /= depth; }
+        auto cost = queueCost(graph, schedule, target, true);
+        Ratio worst{0, 1};
+        for (auto [policy, proof] : llvm::enumerate(cost.recurrence)) {
+          TEST_CHECK(proof.proved() && proof.periodGenerations > 0 && proof.periodGenerations <= 1024);
+          double value = *proof.intervalNs;
+          TEST_CHECK(value > 0 && value < (uint64_t{1} << 30));
+          // Queue timestamps are checked integers. Adjacent binary64 values
+          // enclose the final rounded duration/period division, independently of MCR.
+          long double low = std::nextafter(value, -std::numeric_limits<double>::infinity());
+          long double high = std::nextafter(value, std::numeric_limits<double>::infinity());
+          uint64_t lower = std::ceil(low * proof.periodGenerations);
+          uint64_t upper = std::floor(high * proof.periodGenerations);
+          TEST_CHECK(lower == upper && upper < (uint64_t{1} << 40));
+          uint64_t divisor = std::gcd(lower, proof.periodGenerations);
+          Ratio exact{lower / divisor, proof.periodGenerations / divisor};
+          auto &minimum = (policy ? certifiedAlternative : certifiedPrimary)[depth];
+          if (!minimum || less(exact, *minimum)) minimum = exact;
+          if (less(worst, exact)) worst = exact;
+        }
+        if (!certified[depth] || less(worst, *certified[depth])) certified[depth] = worst;
+        double value = std::max(*cost.primary, *cost.alternative);
+        if (value < curve[depth]) {
+          curve[depth] = value;
+          winnerStages[depth] = schedule.processorStages;
+          winnerCores[depth].clear();
+          for (const auto &node : graph.nodes)
+            winnerCores[depth].push_back(schedule.computeToCpuMap.lookup(node.instance));
+        }
+        primary[depth] = std::min(primary[depth], *cost.primary);
+        alternative[depth] = std::min(alternative[depth], *cost.alternative);
+        ++evaluations[depth];
+      }
+    }
+  }
+  TEST_CHECK(curve[1] == 496 && curve[2] == 304 && curve[3] == 305 && curve[4] == 304.5);
+  TEST_CHECK(primary == curve && alternative == curve);
+  TEST_CHECK(certified[3] == Ratio(305, 1) && certified[4] == Ratio(609, 2));
+  TEST_CHECK(certifiedPrimary == certified && certifiedAlternative == certified);
+  TEST_CHECK(less(*certified[4], *certified[3]));
+  TEST_CHECK(evaluations == std::vector<size_t>({0, 1024, 1728, 792, 96}));
+  std::cout << "nonunimodal_physical_depth_reference";
+  for (size_t depth = 1; depth <= 4; ++depth)
+    std::cout << " p" << depth << "_ns=" << curve[depth] << " primary=" << primary[depth]
+      << " alternative=" << alternative[depth] << " exact=" << certified[depth]->first
+      << '/' << certified[depth]->second << " evaluations=" << evaluations[depth];
+  std::cout << " reference_domain=all_chain_partitions_and_physical_assignments_with_restart\n";
+  for (size_t depth = 1; depth <= 4; ++depth) {
+    std::cout << "physical_depth_winner=" << depth << " core_stages=";
+    for (size_t stage : winnerStages[depth]) std::cout << stage << ',';
+    std::cout << " task_cores=";
+    for (size_t core : winnerCores[depth]) std::cout << core << ',';
+    std::cout << '\n';
+  }
+}
+
 static void testWorseningIntermediate() {
   // A controlled cost oracle isolates the search's non-monotone landscape.
   // Hardware queue objectives are covered separately by runReference().
@@ -311,13 +633,14 @@ static void testWorseningIntermediate() {
   target.processorCount = 8;
   target.interProcessorLatencyNs.assign(64, 0);
   ComputeGraph graph;
-  graph.predecessors.resize(4); graph.successors.resize(4);
-  for (uint32_t task = 0; task < 4; ++task)
-    graph.nodes.push_back({{nullptr, task, 1}, 10, {}, task});
+  graph.predecessors.resize(5); graph.successors.resize(5);
+  for (uint32_t task = 0; task < 5; ++task)
+    graph.nodes.push_back({{nullptr, task, 1}, task == 4 ? 0u : 10u, {}, task});
   auto evaluate = [&](const MergeScheduleResult &schedule) {
-    bool firstStage = llvm::all_of(graph.nodes, [&](const auto &node) {
-      return schedule.processorStages[schedule.computeToCpuMap.lookup(node.instance)] == 0;
-    });
+    bool firstStage = true;
+    for (size_t i = 0; i < graph.nodes.size(); ++i)
+      firstStage &= schedule.processorStages[schedule.computeToCpuMap.lookup(graph.nodes[i].instance)]
+        == (i == 4 ? 1 : 0);
     size_t cores = llvm::count(schedule.processorStages, 0);
     double value = !firstStage ? 120 : cores == 1 ? 80 : cores >= 6 ? 100 : 110;
     WholeScheduleCost cost;
@@ -330,9 +653,9 @@ static void testWorseningIntermediate() {
   };
   std::string error;
   MergeScheduleResult initial, intermediate;
-  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, initial, 2, target, error)));
+  const std::vector<size_t> stages{0, 0, 0, 0, 1};
+  TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {7, 1}, target, initial, error)));
   TEST_CHECK(evaluate(initial).primary == 100);
-  const std::vector<size_t> stages(4);
   TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {3, 5}, target, intermediate, error)));
   TEST_CHECK(evaluate(intermediate).primary == 110);
   double reference = std::numeric_limits<double>::infinity();
@@ -346,6 +669,10 @@ static void testWorseningIntermediate() {
   TEST_CHECK(mlir::succeeded(search));
   size_t winner = selectPipelineCandidate(search->costs, error);
   double selected = *search->costs[winner].primary;
+  size_t depth = 0;
+  for (size_t i = winner; i; i = search->parents[i]) ++depth;
+  TEST_CHECK(depth >= 2);
+  sawSuccessiveMoves = true;
   TEST_CHECK(selected >= reference && selected <= 100);
   TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget);
   std::cout << "non_monotone_oracle_reference=" << reference << " selected=" << selected
@@ -354,6 +681,69 @@ static void testWorseningIntermediate() {
     << " packing_attempts=" << search->packingAttempts << '\n';
 }
 
+static void testMultipleWorseningAllocations() {
+  mlir::MLIRContext context;
+  context.loadDialect<spatial::SpatialDialect>();
+  auto module = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  mlir::OpBuilder builder(module.getBodyRegion());
+  auto first = spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
+    mlir::TypeRange{}, builder.getI32IntegerAttr(4), mlir::ValueRange{}, mlir::ValueRange{});
+  auto last = spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
+    mlir::TypeRange{}, builder.getI32IntegerAttr(1), mlir::ValueRange{}, mlir::ValueRange{});
+  ComputeGraph graph;
+  graph.predecessors.resize(5); graph.successors.resize(5);
+  for (uint32_t task = 0; task < 5; ++task) {
+    graph.nodes.push_back({{task < 4 ? first.getOperation() : last.getOperation(),
+      task < 4 ? task : 0, 1}, task < 4 ? 10u : 0u, {}, task});
+    if (task < 4) {
+      graph.predecessors[4].push_back({task, {0, 0}});
+      graph.successors[task].push_back({4, {0, 0}});
+    }
+  }
+  auto target = hardware(0);
+  target.processorCount = 32; target.interProcessorLatencyNs.assign(1024, 0);
+  auto query = [&](const MergeScheduleResult &schedule) {
+    size_t cores = llvm::count(schedule.processorStages, 0);
+    double value = cores == 1 ? 80 : cores == 31 ? 100 : 110 + cores;
+    WholeScheduleCost cost;
+    cost.primary = cost.alternative = value;
+    cost.loweringSupport = StaticLoweringSupport::Supported;
+    cost.selectionMargin = kPipelineSelectionMargin;
+    for (auto &proof : cost.recurrence) proof = {ScheduleRecurrenceEvidence::Proved, 1, value, {}};
+    return cost;
+  };
+  std::string error;
+  MergeScheduleResult initial;
+  const std::vector<size_t> stages{0, 0, 0, 0, 1};
+  double reference = std::numeric_limits<double>::infinity();
+  for (size_t cores = 1; cores < 32; ++cores) {
+    MergeScheduleResult schedule;
+    TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {cores, 32 - cores}, target, schedule, error)));
+    reference = std::min(reference, *query(schedule).primary);
+  }
+  TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {31, 1}, target, initial, error)));
+  auto search = optimizePipelineSchedule(graph, initial, 2, target, query, error);
+  auto repeat = optimizePipelineSchedule(graph, initial, 2, target, query, error);
+  TEST_CHECK(mlir::succeeded(search) && mlir::succeeded(repeat));
+  TEST_CHECK(search->parents == repeat->parents && search->costs.size() == repeat->costs.size());
+  for (size_t i = 0; i < search->schedules.size(); ++i) {
+    TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
+    TEST_CHECK(search->costs[i].primary == repeat->costs[i].primary);
+  }
+  size_t winner = selectPipelineCandidate(search->costs, error), path = 0;
+  TEST_CHECK(reference == 80 && *search->costs[winner].primary == reference);
+  for (size_t i = winner; i; i = search->parents[i]) {
+    if (i != winner) TEST_CHECK(*search->costs[i].primary > 100);
+    ++path;
+  }
+  TEST_CHECK(path >= 3);
+  std::cout << "multiple_worsening_reference=" << reference << " selected=" << *search->costs[winner].primary
+    << " gap_percent=0 path_length=" << path << " evaluations=" << search->costs.size()
+    << " packing_attempts=" << search->packingAttempts
+    << " reference_domain=all_31_core_budgets_for_fixed_operation_groups\n";
+  module.erase();
+}
+
 static void runReference(size_t profile, size_t fixture) {
   auto target = hardware(profile);
   ComputeGraph graph;
@@ -393,6 +783,9 @@ static void runReference(size_t profile, size_t fixture) {
   TEST_CHECK(mlir::succeeded(search));
   auto repeat = optimizePipelineSchedule(graph, initial, 3, target, evaluate, error);
   TEST_CHECK(mlir::succeeded(repeat) && search->schedules.size() == repeat->schedules.size());
+  TEST_CHECK(search->parents == repeat->parents && search->expansions == repeat->expansions);
+  TEST_CHECK(search->packingAttempts == repeat->packingAttempts
+    && search->consideredAssignments == repeat->consideredAssignments);
   for (size_t i = 0; i < search->schedules.size(); ++i) {
     TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
     TEST_CHECK(search->schedules[i].computeToCpuSlotMap == repeat->schedules[i].computeToCpuSlotMap);
@@ -412,7 +805,7 @@ static void runReference(size_t profile, size_t fixture) {
     }
     std::vector<bool> used(3);
     for (size_t stage : stages) used[stage] = true;
-    if (used[2] && !used[1]) return;
+    if (!llvm::all_of(used, [](bool value) { return value; })) return;
     for (size_t core0 = 1; core0 + 1 < target.processorCount; ++core0)
       for (size_t core1 = 1; core0 + core1 < target.processorCount; ++core1) {
         budgets = {core0, core1, target.processorCount - core0 - core1};
@@ -454,6 +847,17 @@ static void runReference(size_t profile, size_t fixture) {
     sawJointMove |= cutChanged && currentBudgets != parentBudgets;
   }
   sawSuccessiveMoves |= depth >= 2;
+  for (size_t i = 1; i < search->schedules.size(); ++i) {
+    const auto &child = search->schedules[i], &parent = search->schedules[search->parents[i]];
+    bool cut = false;
+    for (const auto &node : graph.nodes)
+      cut |= child.processorStages[child.computeToCpuMap.lookup(node.instance)]
+        != parent.processorStages[parent.computeToCpuMap.lookup(node.instance)];
+    std::vector<size_t> a(3), b(3);
+    for (size_t stage : child.processorStages) ++a[stage];
+    for (size_t stage : parent.processorStages) ++b[stage];
+    sawJointMove |= cut && a != b;
+  }
   std::cout << "profile=" << profile << " fixture=" << fixture
     << " reference_ns=" << reference << " selected_ns=" << cost
     << " gap_percent=" << 100 * (cost / reference - 1)
@@ -489,12 +893,200 @@ static void runReference(size_t profile, size_t fixture) {
   }
 }
 
+// Exhaustive small-grid DP oracle. Physical placement is checked separately.
+static void testSharedPartitionStates() {
+  for (size_t profile = 0; profile < 2; ++profile) {
+    ComputeGraph graph;
+    const std::vector<Cost> work{10, 30, 5, 50, 20, 15};
+    graph.predecessors.resize(work.size()); graph.successors.resize(work.size());
+    for (uint32_t task = 0; task < work.size(); ++task) {
+      ResidentWeight weight; weight.opaqueLane = task % (profile + 2);
+      graph.nodes.push_back({{nullptr, task, 1}, work[task], {weight}, task});
+      if (task) {
+        graph.predecessors[task].push_back({task - 1, {}});
+        graph.successors[task - 1].push_back({task, {}});
+      }
+    }
+    if (profile) {
+      graph.predecessors[4].push_back({1, {}});
+      graph.successors[1].push_back({4, {}});
+    }
+    auto target = hardware(profile);
+    MergeScheduleResult peft;
+    std::string error;
+    TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, std::vector<size_t>(work.size()),
+      {target.processorCount}, target, peft, error)));
+    auto equalCost = [](const MergeScheduleResult &) {
+      WholeScheduleCost cost;
+      cost.primary = cost.alternative = 100;
+      cost.selectionMargin = kPipelineSelectionMargin;
+      cost.loweringSupport = StaticLoweringSupport::Supported;
+      for (auto &proof : cost.recurrence) proof = {ScheduleRecurrenceEvidence::Proved, 1, 100, {}};
+      return cost;
+    };
+    auto search = optimizePipelineSchedule(graph, peft, 0, target, equalCost, error);
+    if (mlir::failed(search)) std::cerr << error << '\n';
+    TEST_CHECK(mlir::succeeded(search));
+    for (size_t i = 1; i < search->initialSchedules; ++i) {
+      size_t depth = getPipelineStageCount(search->schedules[i]);
+      MergeScheduleResult independentlyInitialized;
+      TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, independentlyInitialized, depth, target, error)));
+      TEST_CHECK(search->schedules[i].processorStages == independentlyInitialized.processorStages);
+      TEST_CHECK(search->schedules[i].computeToCpuMap == independentlyInitialized.computeToCpuMap);
+      TEST_CHECK(search->schedules[i].computeToCpuSlotMap == independentlyInitialized.computeToCpuSlotMap);
+    }
+    IndexedResidentWeights weights(graph);
+    for (const auto &entry : search->depthPriorities) {
+      size_t depth = entry.first;
+      Cost actual = entry.second;
+      Cost reference = std::numeric_limits<Cost>::max();
+      size_t evaluated = 0;
+      for (size_t mask = 0; mask < (size_t{1} << (work.size() - 1)); ++mask) {
+        std::vector<size_t> stages(work.size());
+        for (size_t task = 1; task < work.size(); ++task)
+          stages[task] = stages[task - 1] + ((mask >> (task - 1)) & 1);
+        if (stages.back() + 1 != depth) continue;
+        bool legal = true;
+        for (size_t task = 0; task < work.size(); ++task)
+          for (auto [before, transfer] : graph.predecessors[task])
+            legal &= stages[task] <= stages[before] + 1;
+        if (!legal) continue;
+        std::vector<std::vector<size_t>> tasks(depth);
+        for (size_t task = 0; task < work.size(); ++task) tasks[stages[task]].push_back(task);
+        std::function<void(size_t, size_t, Cost)> allocate = [&](size_t stage, size_t remaining, Cost objective) {
+          if (stage == depth) {
+            if (!remaining) { reference = std::min(reference, objective); ++evaluated; }
+            return;
+          }
+          const auto &members = tasks[stage];
+          std::vector<Cost> ends(work.size());
+          Cost span = 0, largest = 0;
+          for (size_t task : members) {
+            ends[task] = checkedAdd(ends[task], work[task]);
+            span = std::max(span, ends[task]); largest = std::max(largest, work[task]);
+            for (auto [next, transfer] : graph.successors[task]) ends[next] = std::max(ends[next], ends[task]);
+          }
+          Cost packed = std::numeric_limits<Cost>::max();
+          for (size_t cores = 1; cores + depth - stage - 1 <= remaining; ++cores) {
+            if (auto cost = weights.estimatePackCost(members, work, target.residentWeightCapacity, cores))
+              packed = std::min(packed, *cost);
+            if (packed != std::numeric_limits<Cost>::max())
+              allocate(stage + 1, remaining - cores, std::max(objective, checkedAdd(span - largest, packed)));
+          }
+        };
+        allocate(0, target.processorCount, 0);
+      }
+      TEST_CHECK(reference == actual);
+      std::cout << "shared_dp_profile=" << profile << " depth=" << depth
+        << " reference_priority=" << reference << " actual_priority=" << actual
+        << " reference_allocations=" << evaluated << " dp_transitions=" << search->partitionTransitions << '\n';
+    }
+  }
+}
+
+static void testDepthDomain() {
+  ComputeGraph graph;
+  graph.predecessors.resize(12); graph.successors.resize(12);
+  for (uint32_t i = 0; i < 12; ++i) {
+    graph.nodes.push_back({{nullptr, i, 1}, 100, {}, i});
+    if (i) {
+      graph.predecessors[i].push_back({i - 1, {}});
+      graph.successors[i - 1].push_back({i, {}});
+    }
+  }
+  auto target = hardware(0);
+  target.processorCount = 16;
+  target.interProcessorLatencyNs.assign(256, 1);
+  std::string error;
+  for (size_t depth : {2, 4, 8}) {
+    MergeScheduleResult schedule;
+    TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, schedule, depth, target, error)));
+    TEST_CHECK(mlir::succeeded(verifyPipelineStageCount(schedule, depth, error)));
+    std::set<size_t> active;
+    for (const auto &node : graph.nodes)
+      active.insert(schedule.processorStages[schedule.computeToCpuMap.lookup(node.instance)]);
+    TEST_CHECK(active.size() == depth && *active.begin() == 0 && *active.rbegin() == depth - 1);
+    auto broken = schedule;
+    for (const auto &[instance, core] : broken.computeToCpuMap)
+      if (broken.processorStages[core] == depth - 1) broken.processorStages[core] = depth - 2;
+    TEST_CHECK(mlir::failed(verifyPipelineStageCount(broken, depth, error)));
+  }
+  MergeScheduleResult invalid;
+  TEST_CHECK(mlir::failed(runPipelineScheduler(graph, invalid, 13, target, error)));
+  TEST_CHECK(error.find("compute operation groups") != std::string::npos);
+  TEST_CHECK(mlir::failed(runPipelineScheduler(graph, invalid, 17, target, error)));
+
+  MergeScheduleResult peft;
+  peft.processorCount = target.processorCount;
+  for (size_t i = 0; i < graph.nodes.size(); ++i) {
+    peft.computeToCpuMap[graph.nodes[i].instance] = 0;
+    peft.computeToCpuSlotMap[graph.nodes[i].instance] = i;
+  }
+  auto incomplete = peft;
+  incomplete.processorStages = {0};
+  TEST_CHECK(mlir::failed(verifyPipelineStageCount(incomplete, 1, error)));
+  std::vector<size_t> winners;
+  for (double boundaryCost : {1000., 300., 50., 0.}) {
+    // A closed synthetic target contract isolates depth-domain selection.
+    // Complete queue objectives and physical placements are tested separately.
+    auto query = [&](const MergeScheduleResult &schedule) {
+      size_t depth = getPipelineStageCount(schedule);
+      double interval = 1200. / depth + boundaryCost * (depth - 1);
+      WholeScheduleCost cost;
+      cost.primary = cost.alternative = interval;
+      cost.loweringSupport = StaticLoweringSupport::Supported;
+      cost.selectionMargin = kPipelineSelectionMargin;
+      for (auto &proof : cost.recurrence)
+        proof = {ScheduleRecurrenceEvidence::Proved, 1, interval, {}};
+      return cost;
+    };
+    auto search = optimizePipelineSchedule(graph, peft, 0, target, query, error);
+    TEST_CHECK(mlir::succeeded(search));
+    auto repeat = optimizePipelineSchedule(graph, peft, 0, target, query, error);
+    TEST_CHECK(mlir::succeeded(repeat) && repeat->schedules.size() == search->schedules.size());
+    TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget);
+    size_t winner = selectPipelineCandidate(search->costs, error);
+    size_t second = selectPipelineCandidate(repeat->costs, error);
+    TEST_CHECK(winner == second);
+    for (size_t i = 0; i < search->schedules.size(); ++i) {
+      TEST_CHECK(search->schedules[i].processorStages == repeat->schedules[i].processorStages);
+      TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
+      TEST_CHECK(search->costs[i].primary == repeat->costs[i].primary);
+      TEST_CHECK(mlir::succeeded(verifyPipelineStageCount(search->schedules[i],
+        getPipelineStageCount(search->schedules[i]), error)));
+    }
+    winners.push_back(getPipelineStageCount(search->schedules[winner]));
+    double reference = std::numeric_limits<double>::infinity();
+    for (size_t depth = 1; depth <= 12; ++depth)
+      reference = std::min(reference, 1200. / depth + boundaryCost * (depth - 1));
+    std::cout << "auto_target_boundary_ns=" << boundaryCost << " depth=" << winners.back()
+      << " reference_ns=" << reference << " selected_ns=" << *search->costs[winner].primary
+      << " evaluations=" << search->costs.size() << '\n';
+  }
+  TEST_CHECK(winners == std::vector<size_t>({1, 2, 5, 12}));
+  graph.predecessors[11].push_back({0, {}});
+  graph.successors[0].push_back({11, {}});
+  TEST_CHECK(mlir::failed(runPipelineScheduler(graph, invalid, 8, target, error)));
+  auto unavailable = [](const MergeScheduleResult &) {
+    WholeScheduleCost cost; cost.reason = "unavailable auto test cost"; return cost;
+  };
+  TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, peft, 0, target, unavailable, error)));
+  TEST_CHECK(error.find("unavailable auto test cost") != std::string::npos);
+}
+
 int main() {
+  testPackingCostCacheEviction();
+  testPackingCurves();
+  testSharedPartitionStates();
+  testDepthDomain();
   testInitializerPackingWitness();
   testPackingFeasibility();
   for (uint32_t first : {19u, 21u}) testQuantizedFeasibility(first);
   testPhysicalReference();
+  testAutoPhysicalReference();
+  testNonunimodalPhysicalDepthCurve();
   testWorseningIntermediate();
+  testMultipleWorseningAllocations();
   for (size_t profile = 0; profile < 2; ++profile)
     for (size_t fixture = 0; fixture < 6; ++fixture) runReference(profile, fixture);
   TEST_CHECK(sawJointMove && sawSuccessiveMoves);