better compile times for throughput scheduling

fix failing no-plan throughput ablation runs
update results

test/PIM/PipelineOptimizerTest.cpp

@@ -3,6 +3,8 @@
 #include "src/Accelerators/PIM/Compiler/PimInstructionTiming.hpp"
 #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/SpatialOps.hpp"
 #include <chrono>
 #include <functional>
 #include <map>
@@ -14,6 +16,76 @@ using namespace onnx_mlir::spatial;
 using E = PimExecutionEvent;
 static bool sawJointMove = false, sawSuccessiveMoves = false;
 
+static void testPackingFeasibility() {
+  ComputeGraph graph;
+  for (size_t task = 0; task < 4; ++task) {
+    ResidentWeight weight;
+    weight.opaqueLane = task % 2;
+    graph.nodes.push_back({{nullptr, static_cast<uint32_t>(task), 1}, 10, {weight}, task});
+  }
+  IndexedResidentWeights weights(graph);
+  const std::vector<size_t> tasks{0, 1, 2, 3};
+  const std::vector<Cost> work{10, 10, 10, 10};
+  std::vector<size_t> assignment{99};
+  TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
+    tasks, 1, 2, &assignment));
+  TEST_CHECK(assignment.empty());
+  TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
+    tasks, 1, 0));
+  TEST_CHECK(weights.estimatePackCost(tasks, work, 1, 2) == 20);
+  auto placement = weights.pack(tasks, work, 1, 2);
+  TEST_CHECK(placement && (*placement)[0] == (*placement)[2]
+    && (*placement)[1] == (*placement)[3] && (*placement)[0] != (*placement)[1]);
+  TEST_CHECK(!weights.estimatePackCost(tasks, work, 1, 1));
+  for (size_t cores = 4; cores <= 8; ++cores)
+    TEST_CHECK(estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
+      tasks, 1, cores) == 10);
+  weights.taskWeightIds[0] = {0, 1};
+  TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
+    tasks, 1, 4));
+  weights.taskWeightIds = {{0, 1}, {2}, {0, 1}, {2}};
+  weights.weightCount = 3;
+  TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
+    tasks, 2, 2));
+  TEST_CHECK(weights.estimatePackCost(tasks, work, 2, 2) == 20);
+}
+
+static void testQuantizedFeasibility(uint32_t first) {
+  mlir::MLIRContext context;
+  context.loadDialect<spatial::SpatialDialect>();
+  auto module = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  mlir::OpBuilder builder(module.getBodyRegion());
+  ComputeGraph graph;
+  graph.predecessors.resize(40); graph.successors.resize(40);
+  for (uint32_t group = 0; group < 2; ++group) {
+    const uint32_t lanes = group ? 40 - first : first;
+    auto batch = spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
+      mlir::TypeRange{}, builder.getI32IntegerAttr(lanes), mlir::ValueRange{}, mlir::ValueRange{});
+    for (uint32_t lane = 0; lane < lanes; ++lane) {
+      size_t task = graph.nodes.size();
+      ResidentWeight weight; weight.opaqueLane = task;
+      graph.nodes.push_back({{batch, lane, 1}, 10, {weight}, task});
+      if (group) {
+        graph.predecessors[task].push_back({0, {0, 0}});
+        graph.successors[0].push_back({task, {0, 0}});
+      }
+    }
+  }
+  SchedulingTarget target;
+  target.processorCount = 40; target.residentWeightCapacity = 1;
+  target.interProcessorLatencyNs.assign(1600, 0);
+  MergeScheduleResult schedule;
+  std::string error;
+  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
+  std::vector<size_t> budgets(2);
+  for (size_t stage : schedule.processorStages) ++budgets[stage];
+  TEST_CHECK(budgets == std::vector<size_t>({first, 40 - 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));
+  module.erase();
+}
+
 // Test-only workloads use stage-boundary host transfers and local direct links.
 // The same queue
 // engine and complete-state evidence rank reference and searched schedules.
@@ -112,6 +184,176 @@ static PimInstructionTiming hardware(size_t profile) {
   return target;
 }
 
+static void testInitializerPackingWitness() {
+  mlir::MLIRContext context;
+  context.loadDialect<spatial::SpatialDialect>();
+  auto module = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  mlir::OpBuilder builder(module.getBodyRegion());
+  std::vector<mlir::Operation *> operations;
+  for (size_t i = 0; i < 3; ++i) operations.push_back(
+    spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
+      mlir::TypeRange{}, builder.getI32IntegerAttr(3), mlir::ValueRange{}, mlir::ValueRange{}));
+  const std::vector<std::vector<size_t>> resident{{2,1,5},{2,4},{0,2},{0},{4,6},{1},{2},{2},{1,4}};
+  const std::vector<Cost> costs{4,1,5,4,1,2,5,2,2};
+  ComputeGraph graph;
+  graph.predecessors.resize(9); graph.successors.resize(9);
+  for (uint32_t task = 0; task < 9; ++task) {
+    ResidentWeightSet weights;
+    for (size_t id : resident[task]) {
+      ResidentWeight w; w.opaqueLane = id; weights.push_back(w);
+    }
+    graph.nodes.push_back({{operations[task % 3], task / 3, 1}, costs[task], weights, task});
+  }
+  IndexedResidentWeights indexed(graph);
+  TEST_CHECK(!indexed.estimatePackCost({1,2,4,5,7,8}, costs, 3, 2));
+  std::vector<size_t> witness;
+  TEST_CHECK(indexed.estimatePackCost({1,4,7,2,5,8}, costs, 3, 2, &witness));
+  TEST_CHECK(witness.size() == 6);
+  auto target = hardware(1);
+  target.processorCount = 4; target.residentWeightCapacity = 3;
+  target.interProcessorLatencyNs.assign(16, 1);
+  MergeScheduleResult schedule;
+  std::string error;
+  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
+  std::vector<ResidentWeightSet> actual(4);
+  for (const auto &node : graph.nodes) {
+    size_t core = schedule.computeToCpuMap.lookup(node.instance);
+    TEST_CHECK(core < 4);
+    insertResidentWeights(actual[core], node.residentWeights);
+    TEST_CHECK(actual[core].size() <= 3);
+  }
+  // Exhaust all core assignments under the same stage map. The initialized
+  // assignment must belong to this physical-capacity reference set.
+  size_t feasible = 0;
+  for (size_t encoded = 0; encoded < (size_t{1} << 18); ++encoded) {
+    size_t choices = encoded;
+    std::vector<ResidentWeightSet> weights(4);
+    std::vector<size_t> groupStages(3, 2);
+    bool valid = true;
+    for (size_t task = 0; task < 9 && valid; ++task) {
+      size_t core = choices % 4; choices /= 4;
+      size_t stage = schedule.processorStages[core], group = task % 3;
+      if (groupStages[group] == 2) groupStages[group] = stage;
+      valid &= groupStages[group] == stage;
+      insertResidentWeights(weights[core], graph.nodes[task].residentWeights);
+      valid &= weights[core].size() <= 3;
+    }
+    feasible += valid;
+  }
+  TEST_CHECK(feasible > 0);
+  MergeScheduleResult alternative;
+  TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, std::vector<size_t>(9),
+    {3, 1}, target, alternative, error)));
+  std::cout << "initializer_physical_feasible_assignments=" << feasible << '\n';
+  module.erase();
+}
+
+static void testPhysicalReference() {
+  auto target = hardware(1);
+  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 optimum = std::numeric_limits<double>::infinity();
+  size_t evaluations = 0;
+  for (size_t mask = 1; mask < 7; ++mask)
+    for (size_t placement = 0; placement < 27; ++placement) {
+      MergeScheduleResult schedule;
+      for (size_t core = 0; core < 3; ++core) schedule.processorStages.push_back((mask >> core) & 1);
+      size_t encoded = placement;
+      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;
+      }
+      if (!valid || stages[0] != 0 || stages[1] < stages[0] || stages[2] < stages[1]) continue;
+      auto cost = queueCost(graph, schedule, target);
+      optimum = std::min(optimum, std::max(*cost.primary, *cost.alternative));
+      ++evaluations;
+    }
+  std::string error;
+  MergeScheduleResult initial;
+  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, initial, 2, target, error)));
+  auto evaluate = [&](const MergeScheduleResult &schedule) { return queueCost(graph, schedule, target); };
+  auto search = optimizePipelineSchedule(graph, initial, 2, target, evaluate, error);
+  TEST_CHECK(mlir::succeeded(search));
+  size_t winner = search->costs.empty() ? 0 : selectPipelineCandidate(search->costs, error);
+  auto selected = evaluate(search->schedules[winner]);
+  double value = std::max(*selected.primary, *selected.alternative);
+  TEST_CHECK(value + 1e-6 >= optimum);
+  std::cout << "physical_reference_ns=" << optimum << " selected_ns=" << value
+    << " gap_percent=" << 100 * (value / optimum - 1)
+    << " physical_evaluations=" << evaluations
+    << " search_evaluations=" << search->costs.size()
+    << " packing_attempts=" << search->packingAttempts << '\n';
+}
+
+static void testWorseningIntermediate() {
+  // A controlled cost oracle isolates the search's non-monotone landscape.
+  // Hardware queue objectives are covered separately by runReference().
+  auto target = hardware(0);
+  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});
+  auto evaluate = [&](const MergeScheduleResult &schedule) {
+    bool firstStage = llvm::all_of(graph.nodes, [&](const auto &node) {
+      return schedule.processorStages[schedule.computeToCpuMap.lookup(node.instance)] == 0;
+    });
+    size_t cores = llvm::count(schedule.processorStages, 0);
+    double value = !firstStage ? 120 : cores == 1 ? 80 : cores >= 6 ? 100 : 110;
+    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, intermediate;
+  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, initial, 2, target, 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();
+  for (size_t cores = 1; cores < 8; ++cores) {
+    MergeScheduleResult schedule;
+    TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {cores, 8 - cores}, target, schedule, error)));
+    reference = std::min(reference, *evaluate(schedule).primary);
+  }
+  TEST_CHECK(reference == 80);
+  auto search = optimizePipelineSchedule(graph, initial, 2, target, evaluate, error);
+  TEST_CHECK(mlir::succeeded(search));
+  size_t winner = selectPipelineCandidate(search->costs, error);
+  double selected = *search->costs[winner].primary;
+  TEST_CHECK(selected >= reference && selected <= 100);
+  TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget);
+  std::cout << "non_monotone_oracle_reference=" << reference << " selected=" << selected
+    << " gap_percent=" << 100 * (selected / reference - 1)
+    << " evaluations=" << search->costs.size()
+    << " packing_attempts=" << search->packingAttempts << '\n';
+}
+
 static void runReference(size_t profile, size_t fixture) {
   auto target = hardware(profile);
   ComputeGraph graph;
@@ -217,8 +459,10 @@ static void runReference(size_t profile, size_t fixture) {
     << " gap_percent=" << 100 * (cost / reference - 1)
     << " best_evaluated_ns=" << bestEvaluated
     << " best_evaluated_gap_percent=" << 100 * (bestEvaluated / reference - 1)
+    << " reference_scope=logical_with_deterministic_packing"
     << " reference_evaluations=" << referenceEvaluations
-    << " search_evaluations=" << search->costs.size() << " refinement_depth=" << depth
+    << " search_evaluations=" << search->costs.size()
+    << " packing_attempts=" << search->packingAttempts << " refinement_depth=" << depth
     << " search_seconds=" << searchSeconds << " reference_seconds=" << referenceSeconds
     << " total_seconds=" << std::chrono::duration<double>(std::chrono::steady_clock::now() - begin).count()
     << '\n';
@@ -246,6 +490,11 @@ static void runReference(size_t profile, size_t fixture) {
 }
 
 int main() {
+  testInitializerPackingWitness();
+  testPackingFeasibility();
+  for (uint32_t first : {19u, 21u}) testQuantizedFeasibility(first);
+  testPhysicalReference();
+  testWorseningIntermediate();
   for (size_t profile = 0; profile < 2; ++profile)
     for (size_t fixture = 0; fixture < 6; ++fixture) runReference(profile, fixture);
   TEST_CHECK(sawJointMove && sawSuccessiveMoves);