even better more structured throughput scheduling

plot adjustments
validation adjustments

test/PIM/PipelineOptimizerTest.cpp

@@ -0,0 +1,252 @@
+#include "TestCheck.hpp"
+#include "src/Accelerators/PIM/Compiler/PimPeriodicAnalysis.hpp"
+#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 <chrono>
+#include <functional>
+#include <map>
+#include <set>
+#include <iostream>
+
+using namespace onnx_mlir;
+using namespace onnx_mlir::spatial;
+using E = PimExecutionEvent;
+static bool sawJointMove = false, sawSuccessiveMoves = false;
+
+// 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.
+static WholeScheduleCost queueCost(const ComputeGraph &graph,
+    const MergeScheduleResult &schedule, const PimInstructionTiming &target) {
+  PimExecutionProgram program(target.processorCount);
+  std::vector<std::vector<size_t>> tasks(target.processorCount);
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    tasks[schedule.computeToCpuMap.lookup(graph.nodes[task].instance)].push_back(task);
+  for (auto &core : tasks)
+    llvm::sort(core, [&](size_t a, size_t b) {
+      return schedule.computeToCpuSlotMap.lookup(graph.nodes[a].instance)
+        < schedule.computeToCpuSlotMap.lookup(graph.nodes[b].instance);
+    });
+  struct Transfer { size_t source, target, reader, writer, reg; bool host; };
+  std::vector<Transfer> transfers;
+  std::vector<std::set<size_t>> readers(target.processorCount), writers(target.processorCount);
+  // Branched synthetic workloads use host routing to avoid inventing a second
+  // blocking-transfer scheduler in this test-only objective.
+  bool branching = llvm::any_of(graph.successors, [](const auto &edges) { return edges.size() > 1; });
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    for (const auto &[predecessor, dependency] : graph.predecessors[task]) {
+      size_t writer = schedule.computeToCpuMap.lookup(graph.nodes[predecessor].instance);
+      size_t reader = schedule.computeToCpuMap.lookup(graph.nodes[task].instance);
+      if (writer == reader) continue;
+      bool host = branching || schedule.processorStages[writer] != schedule.processorStages[reader];
+      transfers.push_back({predecessor, task, reader, writer, 1 + transfers.size(), host});
+      if (host) { readers[writer].insert(reader); writers[reader].insert(writer); }
+    }
+  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 peers = 0;
+      for (size_t peer = 0; peer < program.size(); ++peer)
+        if (peer != core && schedule.processorStages[peer] == 0) {
+          events.emplace_back(E::Signal, 0, peer, 100, 1); ++peers;
+        }
+      if (peers) events.emplace_back(E::Wait, 0, 0, 100, peers);
+    }
+    bool reuseWait = false;
+    for (size_t task : tasks[core]) {
+      for (const auto &transfer : transfers)
+        if (transfer.target == task) {
+          if (transfer.host) {
+            events.emplace_back(E::Wait, 0, 0, transfer.reg, 1);
+            events.emplace_back(E::Load, 64);
+          } else events.emplace_back(E::Receive, 64, transfer.writer);
+        }
+      E work(E::Work, graph.nodes[task].cost); work.workClass = E::Vector;
+      events.push_back(work);
+      for (const auto &transfer : transfers)
+        if (transfer.source == task) {
+          if (!transfer.host) {
+            events.emplace_back(E::Send, 64, transfer.reader); continue;
+          }
+          if (!reuseWait) {
+            events.emplace_back(E::Wait, 0, 0, 0, readers[core].size()); reuseWait = true;
+          }
+          events.emplace_back(E::Store, 64);
+          events.emplace_back(E::Signal, 0, transfer.reader, transfer.reg, 1);
+        }
+    }
+  }
+  PimPeriodicOptions options;
+  options.verifyPeriodicState = true; options.searchGuidance = true;
+  auto estimate = evaluatePimExecution(program, target, options);
+  auto cost = std::move(estimate.cost);
+  cost.loweringSupport = StaticLoweringSupport::Supported;
+  cost.selectionMargin = 0.021;
+  TEST_CHECK(cost.eligible());
+  return cost;
+}
+
+static PimInstructionTiming hardware(size_t profile) {
+  PimInstructionTiming target;
+  target.processorCount = 5;
+  target.residentWeightCapacity = 2;
+  target.estimatePipelineDependencySpan = [](llvm::ArrayRef<Cost> spans) { return spans.front(); };
+  target.peerPermissionThenSourceRead = true;
+  target.signalIssueNs = profile + 1;
+  target.controlMessageBytes = 1;
+  target.sharedMemoryWidthBytes = 64;
+  target.sharedMemoryReadLatencyNs = 17 + profile * 9;
+  target.sharedMemoryWriteLatencyNs = 23 + profile * 7;
+  target.synchronizationRegisterCount = 128;
+  target.sharedMemoryRequestLatencyNs.assign(5, 2);
+  target.sharedMemoryResponseLatencyNs.assign(5, 3);
+  target.interProcessorLatencyNs.resize(25);
+  for (size_t a = 0; a < 5; ++a)
+    for (size_t b = 0; b < 5; ++b) {
+      const size_t irregular[] = {0, 1, 5, 6, 20};
+      size_t x = profile ? irregular[a] : a, y = profile ? irregular[b] : b;
+      target.interProcessorLatencyNs[a * 5 + b] = x > y ? x - y : y - x;
+    }
+  return target;
+}
+
+static void runReference(size_t profile, size_t fixture) {
+  auto target = hardware(profile);
+  ComputeGraph graph;
+  const std::vector<std::vector<Cost>> work{{900, 70, 800, 90, 700},
+    {30, 700, 110, 600, 90}, {1000, 50, 40, 900, 60},
+    {90, 80, 100, 110, 120}, {100, 800, 100, 800, 100},
+    {900, 800, 40, 30, 20}};
+  size_t count = work[fixture].size();
+  graph.predecessors.resize(count); graph.successors.resize(count);
+  for (size_t task = 0; task < count; ++task) {
+    ResidentWeight weight; weight.opaqueLane = fixture == 4 ? task % 3 : task % 2;
+    graph.nodes.push_back({{nullptr, static_cast<uint32_t>(task), 1}, work[fixture][task],
+      {weight}, task});
+    auto edge = [&](size_t source) {
+      graph.predecessors[task].push_back({source, {64, 1}});
+      graph.successors[source].push_back({task, {64, 1}});
+      graph.publicationTransfers.push_back({source, task, {}, 0, {}, 64, 1});
+    };
+    if (fixture == 1) {
+      if (task == 1 || task == 2) edge(0);
+      if (task == 3) { edge(1); edge(2); }
+      if (task == 4) edge(3);
+    } else if (task) edge(task - 1);
+    if (fixture == 2 && task == 4) edge(1);
+  }
+  if (fixture == 3) {
+    target.sharedMemoryReadLatencyNs = 400;
+    target.sharedMemoryWriteLatencyNs = 600;
+  }
+  std::string error;
+  MergeScheduleResult initial;
+  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, initial, 3, target, error)));
+  auto evaluate = [&](const MergeScheduleResult &schedule) { return queueCost(graph, schedule, target); };
+  auto begin = std::chrono::steady_clock::now();
+  auto search = optimizePipelineSchedule(graph, initial, 3, target, evaluate, error);
+  double searchSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - begin).count();
+  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());
+  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);
+    TEST_CHECK(search->schedules[i].processorStages == repeat->schedules[i].processorStages);
+    TEST_CHECK(search->costs.empty() || search->costs[i].primary == repeat->costs[i].primary);
+    TEST_CHECK(search->costs.empty() || search->costs[i].alternative == repeat->costs[i].alternative);
+  }
+  double reference = std::numeric_limits<double>::infinity();
+  size_t referenceEvaluations = 0;
+  std::vector<size_t> stages(count), budgets(3, 1);
+  std::function<void(size_t)> partitions = [&](size_t task) {
+    if (task != count) {
+      for (size_t stage = 0; stage <= (task ? 2 : 0); ++stage) {
+        stages[task] = stage; partitions(task + 1);
+      }
+      return;
+    }
+    std::vector<bool> used(3);
+    for (size_t stage : stages) used[stage] = true;
+    if (used[2] && !used[1]) 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};
+        MergeScheduleResult schedule;
+        if (mlir::failed(packPipelineSchedule(graph, stages, budgets, target, schedule, error))) continue;
+        auto cost = evaluate(schedule); ++referenceEvaluations;
+        reference = std::min(reference, std::max(*cost.primary, *cost.alternative));
+      }
+  };
+  auto referenceBegin = std::chrono::steady_clock::now();
+  partitions(0);
+  double referenceSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - referenceBegin).count();
+  std::string reason;
+  size_t winner = search->costs.empty() ? 0 : selectPipelineCandidate(search->costs, reason);
+  auto selected = evaluate(search->schedules[winner]);
+  double cost = std::max(*selected.primary, *selected.alternative);
+  double bestEvaluated = cost;
+  for (const auto &estimate : search->costs)
+    bestEvaluated = std::min(bestEvaluated, std::max(*estimate.primary, *estimate.alternative));
+  TEST_CHECK(std::isfinite(reference) && cost + 1e-6 >= reference);
+  TEST_CHECK(bestEvaluated + 1e-6 >= reference);
+  TEST_CHECK(cost <= reference * (1 + kPipelineSelectionMargin) + 1e-6);
+  TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget
+    && search->packingAttempts <= 4 * kPipelineQueueEvaluationBudget
+    && search->consideredAssignments <= 4 * kPipelineQueueEvaluationBudget);
+  size_t depth = 0;
+  for (size_t i = winner; i; i = search->parents[i]) {
+    ++depth;
+    const auto &current = search->schedules[i], &parent = search->schedules[search->parents[i]];
+    bool cutChanged = false;
+    for (const auto &node : graph.nodes)
+      cutChanged |= current.processorStages[current.computeToCpuMap.lookup(node.instance)]
+        != parent.processorStages[parent.computeToCpuMap.lookup(node.instance)];
+    std::vector<size_t> currentBudgets(3), parentBudgets(3);
+    for (size_t core = 0; core < target.processorCount; ++core) {
+      ++currentBudgets[current.processorStages[core]];
+      ++parentBudgets[parent.processorStages[core]];
+    }
+    sawJointMove |= cutChanged && currentBudgets != parentBudgets;
+  }
+  sawSuccessiveMoves |= depth >= 2;
+  std::cout << "profile=" << profile << " fixture=" << fixture
+    << " reference_ns=" << reference << " selected_ns=" << cost
+    << " gap_percent=" << 100 * (cost / reference - 1)
+    << " best_evaluated_ns=" << bestEvaluated
+    << " best_evaluated_gap_percent=" << 100 * (bestEvaluated / reference - 1)
+    << " reference_evaluations=" << referenceEvaluations
+    << " search_evaluations=" << search->costs.size() << " refinement_depth=" << depth
+    << " search_seconds=" << searchSeconds << " reference_seconds=" << referenceSeconds
+    << " total_seconds=" << std::chrono::duration<double>(std::chrono::steady_clock::now() - begin).count()
+    << '\n';
+  auto unavailable = [&](const MergeScheduleResult &) { WholeScheduleCost value; value.reason = "test unavailable timing"; return value; };
+  auto unknown = optimizePipelineSchedule(graph, initial, 3, target, unavailable, error);
+  TEST_CHECK(search->costs.empty() || (mlir::failed(unknown) && error.find("test unavailable timing") != std::string::npos));
+  auto unproved = [&](const MergeScheduleResult &schedule) {
+    auto cost = evaluate(schedule);
+    cost.recurrence[0].status = ScheduleRecurrenceEvidence::Unproved;
+    cost.reason = "test unproved recurrence";
+    return cost;
+  };
+  auto invalid = optimizePipelineSchedule(graph, initial, 3, target, unproved, error);
+  TEST_CHECK(search->costs.empty() || (mlir::failed(invalid) && error.find("test unproved recurrence") != std::string::npos));
+  if (!profile && !fixture) {
+    auto badGuidance = [&](const MergeScheduleResult &schedule) {
+      auto cost = evaluate(schedule); cost.processorTaskWorkNs.resize(target.processorCount + 1);
+      return cost;
+    };
+    TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, badGuidance, error)));
+    TEST_CHECK(error.find("physical processor count") != std::string::npos);
+    TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, evaluate, error, 0)));
+    TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, evaluate, error, kPipelineQueueEvaluationBudget + 1)));
+  }
+}
+
+int main() {
+  for (size_t profile = 0; profile < 2; ++profile)
+    for (size_t fixture = 0; fixture < 6; ++fixture) runReference(profile, fixture);
+  TEST_CHECK(sawJointMove && sawSuccessiveMoves);
+}