even better more structured throughput scheduling

plot adjustments
validation adjustments

src/PIM/Compiler/PimPipelineSelection.cpp

@@ -1,6 +1,8 @@
 #include "PimPipelineSelection.hpp"
 #include "PimCompilerOptions.hpp"
 #include "PimCompilerUtils.hpp"
+#include "PimSpatialCost.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
 #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
 #include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
@@ -108,83 +110,133 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
   }
 
   LogicalResult evaluateSpatialCandidates(ModuleOp checkpoint,
-      const PreparedSpatialSchedules &prepared, std::vector<WholeScheduleCost> &costs,
+      PreparedSpatialSchedules &prepared, std::vector<WholeScheduleCost> &costs,
       llvm::json::Object &report) {
     llvm::json::Array summaries;
-    double realization = 0, extraction = 0, expansion = 0, queue = 0;
-    std::string errors;
-    for (size_t i = 0; i < costs.size(); ++i) {
-      llvm::json::Object summary{{"index", static_cast<int64_t>(i)}};
+    double realization = 0, extraction = 0, expansion = 0, queue = 0, evaluation = 0;
+    auto evaluate = [&](const MergeScheduleResult &schedule) {
+      size_t index = summaries.size();
+      llvm::json::Object summary{{"index", static_cast<int64_t>(index)}};
       OwningOpRef<ModuleOp> candidate;
+      WholeScheduleCost cost;
       std::string scheduleReport;
       {
         ScopedCompilerOutputSuppression silence;
-        if (failed(realizeCandidate(checkpoint, prepared.candidates[i], candidate, scheduleReport, summary)))
-          costs[i].reason = "Spatial schedule realization failed verification";
-        else costs[i] = spatialCost(*candidate, summary);
+        if (failed(realizeCandidate(checkpoint, schedule, candidate, scheduleReport, summary)))
+          cost.reason = "Spatial schedule realization failed verification";
+        else cost = spatialCost(*candidate, summary);
       }
       realization += summary.getNumber("realization_seconds").value_or(0);
-      extraction += costs[i].timing.extractionSeconds;
-      expansion += costs[i].timing.expansionSeconds;
-      queue += costs[i].timing.evaluationSeconds;
-      bool supported = costs[i].loweringSupport == StaticLoweringSupport::Supported;
-      auto convergence = costs[i].convergenceProved();
-      auto support = costs[i].loweringSupported();
-      bool converged = convergence == true;
-      bool safe = costs[i].eligible();
-      summary["cost_available"] = costs[i].available();
+      evaluation += summary.getNumber("spatial_cost_seconds").value_or(0);
+      extraction += cost.timing.extractionSeconds;
+      expansion += cost.timing.expansionSeconds;
+      queue += cost.timing.evaluationSeconds;
+      auto convergence = cost.convergenceProved(), support = cost.loweringSupported();
+      summary["cost_available"] = cost.available();
       summary["convergence_proved"] = convergence
         ? llvm::json::Value(*convergence) : llvm::json::Value(nullptr);
       summary["lowering_supported"] = support
         ? llvm::json::Value(*support) : llvm::json::Value(nullptr);
-      summary["eligible"] = safe;
-      if (costs[i].primary) summary["interval_ns"] = *costs[i].primary;
-      if (costs[i].alternative) summary["alternative_tie_interval_ns"] = *costs[i].alternative;
-      if (!safe) {
-        std::string property;
-        if (!costs[i].available()) property += "unavailable timing; ";
-        if (!supported) property += "unsupported or unknown static lowering; ";
-        if (!converged) property += "unproved modeled-state convergence; ";
-        if (property.empty()) property = "unavailable selection margin";
-        summary["reason"] = property + ": " + costs[i].reason;
-        errors += "\ncandidate " + std::to_string(i) + ": " + property + ": " + costs[i].reason;
+      summary["eligible"] = cost.eligible();
+      if (cost.primary) summary["interval_ns"] = *cost.primary;
+      if (cost.alternative) summary["alternative_tie_interval_ns"] = *cost.alternative;
+      llvm::json::Array budgets, witness;
+      std::vector<size_t> cores(stages);
+      std::vector<double> pressure(stages);
+      for (size_t core = 0; core < schedule.processorStages.size(); ++core) {
+        size_t stage = schedule.processorStages[core];
+        ++cores[stage];
+        if (core < cost.processorCriticalDelayNs.size())
+          pressure[stage] += cost.processorCriticalDelayNs[core];
       }
+      for (size_t stage = 0; stage < stages; ++stage) {
+        budgets.push_back(static_cast<int64_t>(cores[stage]));
+        witness.push_back(pressure[stage]);
+      }
+      summary["stage_core_budgets"] = std::move(budgets);
+      summary["recurrence_stage_contributions_ns"] = std::move(witness);
+      llvm::json::Array runs;
+      for (size_t begin = 0; begin < prepared.graph.nodes.size();) {
+        auto stageFor = [&](size_t task) {
+          return schedule.processorStages[schedule.computeToCpuMap.lookup(
+            prepared.graph.nodes[task].instance)];
+        };
+        size_t end = begin + 1, stage = stageFor(begin);
+        while (end < prepared.graph.nodes.size() && stageFor(end) == stage) ++end;
+        runs.push_back(llvm::json::Array{static_cast<int64_t>(begin),
+          static_cast<int64_t>(end - begin), static_cast<int64_t>(stage)});
+        begin = end;
+      }
+      summary["task_stage_runs"] = std::move(runs);
+      if (!cost.eligible()) summary["reason"] = cost.reason;
       summaries.push_back(std::move(summary));
-    }
-    report["candidates"] = std::move(summaries);
+      return cost;
+    };
+    std::string error;
+    auto start = Clock::now();
+    auto search = optimizePipelineSchedule(prepared.graph, prepared.candidates.front(),
+      stages, target, evaluate, error);
+    report["search_seconds"] = elapsed(start);
+    report["search_refinement_seconds"] = std::max(0.0, elapsed(start) - realization - evaluation);
+    report["cost_evaluation_seconds"] = evaluation;
     report["candidate_realization_seconds"] = realization;
     report["event_extraction_seconds"] = extraction;
     report["event_expansion_seconds"] = expansion;
     report["queue_seconds"] = queue;
-    if (!errors.empty()) {
+    report["candidates"] = std::move(summaries);
+    if (failed(search)) {
       report["status"] = "incomplete-costs";
-      report["reason"] = errors;
-      checkpoint.emitError("Spatial pipeline cannot compare the complete candidate set:") << errors;
+      report["reason"] = error;
+      checkpoint.emitError("Spatial pipeline cannot compare the complete candidate set: ") << error;
       return failure();
     }
+    prepared.candidates = std::move(search->schedules);
+    costs = std::move(search->costs);
+    bool single = costs.empty();
+    if (single) {
+      report["candidates"] = llvm::json::Array{llvm::json::Object{{"index", 0}}};
+    } else {
+      auto *entries = report.getArray("candidates");
+      for (size_t i = 0; i < entries->size(); ++i)
+        (*entries)[i].getAsObject()->insert({"parent", static_cast<int64_t>(search->parents[i])});
+    }
+    report["generated_candidates"] = static_cast<int64_t>(prepared.candidates.size());
+    report["single_candidate_analysis_skipped"] = single;
+    report["search_rounds"] = static_cast<int64_t>(search->rounds);
+    report["packing_attempts"] = static_cast<int64_t>(search->packingAttempts);
+    report["considered_assignments"] = static_cast<int64_t>(search->consideredAssignments);
+    report["search_termination"] = search->termination;
+    report["queue_evaluation_budget"] = static_cast<int64_t>(kPipelineQueueEvaluationBudget);
+    report["beam_width"] = static_cast<int64_t>(kPipelineBeamWidth);
     return success();
   }
 
   void runOnOperation() override {
     ModuleOp checkpoint = getOperation();
     auto start = Clock::now();
-    auto prepared = prepareSpatialSchedules(checkpoint, target, stages, /*collectCandidates=*/true);
+    auto prepared = prepareSpatialSchedules(checkpoint, target, stages);
     if (failed(prepared)) { signalPassFailure(); return; }
     if (prepared->candidates.empty()) {
       checkpoint.emitError("Spatial pipeline requires at least one candidate schedule");
       signalPassFailure(); return;
     }
-    bool single = prepared->candidates.size() == 1;
     llvm::json::Object report{{"strategy", "spatial-queue-aware"}, {"static_analysis", true},
       {"candidate_generation_seconds", elapsed(start)},
       {"generated_candidates", static_cast<int64_t>(prepared->candidates.size())},
-      {"single_candidate_analysis_skipped", single}, {"selection_margin", kPipelineSelectionMargin},
+      {"selection_margin", kPipelineSelectionMargin},
       {"prediction_error_bound_available", false}, {"temporary_disk_bytes", 0}};
-    std::vector<WholeScheduleCost> costs(prepared->candidates.size());
-    if (!single && failed(evaluateSpatialCandidates(checkpoint, *prepared, costs, report))) {
+    if (pimDisableSynchronization) {
+      report["synchronization_ablation"] = true;
+      report["cost_model_program"] = "synchronized-reference";
+      llvm::errs() << "Synchronization ablation: keep the normal selected schedule; "
+        "report zero-cost SYNC/WAIT separately and omit them in code generation. "
+        "Buffer reuse and restart correctness are not guaranteed.\n";
+    }
+    std::vector<WholeScheduleCost> costs;
+    if (failed(evaluateSpatialCandidates(checkpoint, *prepared, costs, report))) {
       writeSelectionReport(std::move(report)); signalPassFailure(); return;
     }
-    if (single) report["candidates"] = llvm::json::Array{llvm::json::Object{{"index", 0}}};
+    bool single = costs.empty();
     OwningOpRef<ModuleOp> winner;
     std::string scheduleReport;
     llvm::json::Array attempts;
@@ -204,6 +256,8 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
         summary["reason"] = "selected Spatial schedule failed deterministic realization";
         summary["status"] = "realization-failed";
       } else {
+        if (pimDisableSynchronization)
+          summary["synchronization_ablation_cost"] = analyzePimSpatialSynchronizationAblation(*candidate);
         ++lowered;
         result = lowerPimPipelineCandidate(*candidate, resources, summary);
       }
@@ -264,7 +318,7 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
     report["reason"] = reason;
     report["finalization_seconds"] = elapsed(finalize);
     report["total_seconds"] = elapsed(start);
-    llvm::errs() << "Spatial queue-aware pipeline: candidate " << *selected << " of " << costs.size()
+    llvm::errs() << "Spatial queue-aware pipeline: candidate " << *selected << " of " << prepared->candidates.size()
                  << "; " << reason << "\n";
     writeSelectionReport(std::move(report));
   }