better throughput scheduling with steady state cost model, now limited to spatial

This commit is contained in:
NiccoloN
2026-10-26 12:07:01 +01:00
parent fd493ae6e8
commit 4b53f6cf10
62 changed files with 4293 additions and 871 deletions
+236 -128
View File
@@ -13,7 +13,7 @@
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/JSON.h"
#include "llvm/Support/raw_os_ostream.h"
#include <algorithm>
#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
#include <chrono>
using namespace mlir;
@@ -57,6 +57,14 @@ struct CandidateTiming final : PassInstrumentation {
}
};
void writeSelectionReport(llvm::json::Object report) {
auto file = openReportFileWithExtension("pipeline_selection", "json");
if (file) {
llvm::raw_os_ostream out(file);
out << llvm::formatv("{0:2}\n", llvm::json::Value(std::move(report)));
}
}
struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, OperationPass<ModuleOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimPipelineSelectionPass)
PimPipelineSelectionPass(const SchedulingTarget &target, const SpatialTargetResources &resources,
@@ -68,156 +76,256 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
SpatialDataflowExportStage exportStage;
StringRef getArgument() const override { return "pim-pipeline-selection"; }
StringRef getDescription() const override { return "Select a throughput pipeline using static instruction queue analysis"; }
StringRef getDescription() const override { return "Select a throughput pipeline using static Spatial queue analysis"; }
LogicalResult lower(ModuleOp module, llvm::json::Object *summary = nullptr) {
PassManager pm(module.getContext());
addPimLoweringPasses(pm, resources, EmitPimCodegen, /*reportProgress=*/false);
if (enableTiming && summary)
pm.addInstrumentation(std::make_unique<CandidateTiming>(*summary));
return pm.run(module);
}
LogicalResult evaluate(ModuleOp checkpoint, const MergeScheduleResult &schedule,
OwningOpRef<ModuleOp> &candidate,
std::string &scheduleReport, llvm::json::Object &summary) {
LogicalResult realizeCandidate(ModuleOp checkpoint, const MergeScheduleResult &schedule,
OwningOpRef<ModuleOp> &candidate, std::string &scheduleReport,
llvm::json::Object &summary) {
auto begin = Clock::now();
IRMapping mapping;
candidate = cast<ModuleOp>(checkpoint->clone(mapping));
auto remapped = remapSpatialSchedule(schedule, mapping);
if (failed(realizeSpatialSchedule(*candidate, remapped, target, stages,
SpatialDataflowExportStage::None, nullptr, &scheduleReport))
|| failed(lower(*candidate, &summary))) return failure();
if (enableTiming) summary["codegen_input_ir"] = countIR(candidate->getOperation());
summary["lowering_seconds"] = elapsed(begin);
SpatialDataflowExportStage::None, nullptr, &scheduleReport))) return failure();
summary["realization_seconds"] = elapsed(begin);
return success();
}
WholeScheduleCost spatialCost(ModuleOp candidate, llvm::json::Object &summary) {
WholeScheduleCost cost;
auto begin = Clock::now();
llvm::json::Object estimate;
if (target.estimateRealizedScheduleCost)
cost = target.estimateRealizedScheduleCost(candidate, estimate);
else cost.reason = target.wholeScheduleCostUnavailable.empty()
? "target has no realized-Spatial cost query" : target.wholeScheduleCostUnavailable;
cost.selectionMargin = kPipelineSelectionMargin;
estimate["safe_for_selection"] = cost.eligible();
estimate["diagnostic_only"] = false;
summary["spatial_cost"] = std::move(estimate);
summary["spatial_cost_seconds"] = elapsed(begin);
return cost;
}
LogicalResult evaluateSpatialCandidates(ModuleOp checkpoint,
const 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)}};
OwningOpRef<ModuleOp> candidate;
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);
}
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();
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;
}
summaries.push_back(std::move(summary));
}
report["candidates"] = std::move(summaries);
report["candidate_realization_seconds"] = realization;
report["event_extraction_seconds"] = extraction;
report["event_expansion_seconds"] = expansion;
report["queue_seconds"] = queue;
if (!errors.empty()) {
report["status"] = "incomplete-costs";
report["reason"] = errors;
checkpoint.emitError("Spatial pipeline cannot compare the complete candidate set:") << errors;
return failure();
}
return success();
}
void runOnOperation() override {
ModuleOp module = getOperation();
ModuleOp checkpoint = getOperation();
auto start = Clock::now();
std::string unsupported;
if (!target.prepareScheduleCost)
unsupported = target.wholeScheduleCostUnavailable.empty()
? "target has no whole-schedule cost query" : target.wholeScheduleCostUnavailable;
auto prepared = prepareSpatialSchedules(module, target, stages, unsupported.empty());
auto prepared = prepareSpatialSchedules(checkpoint, target, stages, /*collectCandidates=*/true);
if (failed(prepared)) { signalPassFailure(); return; }
bool singleCandidate = unsupported.empty() && prepared->candidates.size() == 1;
double generationSeconds = elapsed(start);
std::vector<WholeScheduleCost> scores(prepared->candidates.size());
llvm::json::Array summaries;
OwningOpRef<ModuleOp> bestModule;
std::string bestScheduleReport;
size_t bestIndex = 0;
double loweringSeconds = 0, instructionSeconds = 0, queueSeconds = 0;
size_t lowered = 0, analyzed = 0, unavailable = 0, duplicates = 0;
for (size_t i = 0; i < prepared->candidates.size(); ++i) {
llvm::json::Object summary{{"index", static_cast<int64_t>(i)}};
if (!unsupported.empty()) {
scores[i].reason = unsupported;
} else if (singleCandidate) {
scores[i].reason = "one distinct generated schedule; no cost comparison required";
} else {
ScopedCompilerOutputSuppression silence;
OwningOpRef<ModuleOp> candidate;
std::string scheduleReport;
if (failed(evaluate(module, prepared->candidates[i], candidate, scheduleReport, summary))) {
if (!i) { signalPassFailure(); return; }
scores[i].reason = "candidate lowering failed verification";
} else {
++lowered;
loweringSeconds += *summary.getNumber("lowering_seconds");
auto query = target.prepareScheduleCost(*candidate, summary);
if (failed(query)) {
if (!i) { signalPassFailure(); return; }
scores[i].reason = "candidate cost preparation failed verification";
++unavailable;
summary["eligible"] = false;
summary["reason"] = scores[i].reason;
summaries.push_back(std::move(summary));
continue;
}
instructionSeconds += summary.getNumber("instruction_generation_seconds").value_or(0);
scores[i].identity = query->identity;
auto duplicate = std::find_if(scores.begin(), scores.begin() + i,
[&](const WholeScheduleCost &other) { return other.identity == scores[i].identity; });
if (duplicate != scores.begin() + i) {
++duplicates;
scores[i].reason = "duplicate emitted instruction program";
summary["duplicate_of"] = static_cast<int64_t>(duplicate - scores.begin());
} else {
auto begin = Clock::now();
scores[i] = query->evaluate(summary);
scores[i].identity = query->identity;
double seconds = elapsed(begin);
queueSeconds += seconds;
++analyzed;
summary["queue_seconds"] = seconds;
}
std::string partialReason;
size_t partialSelection = selectPipelineCandidate(
llvm::ArrayRef<WholeScheduleCost>(scores.data(), i + 1), partialReason);
if (!bestModule || partialSelection == i) {
bestIndex = i;
bestModule = std::move(candidate);
bestScheduleReport = std::move(scheduleReport);
}
}
}
if (!singleCandidate && !scores[i].eligible()
&& scores[i].reason != "duplicate emitted instruction program") ++unavailable;
summary["eligible"] = scores[i].eligible();
summary["identity"] = scores[i].identity;
if (scores[i].primary) summary["interval_ns"] = *scores[i].primary;
if (scores[i].alternative) summary["alternative_tie_interval_ns"] = *scores[i].alternative;
if (scores[i].relativeUncertainty)
summary["relative_prediction_uncertainty"] = *scores[i].relativeUncertainty;
if (!scores[i].reason.empty()) summary["reason"] = scores[i].reason;
summaries.push_back(std::move(summary));
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},
{"prediction_error_bound_available", false}, {"temporary_disk_bytes", 0}};
std::vector<WholeScheduleCost> costs(prepared->candidates.size());
if (!single && failed(evaluateSpatialCandidates(checkpoint, *prepared, costs, report))) {
writeSelectionReport(std::move(report)); signalPassFailure(); return;
}
if (single) report["candidates"] = llvm::json::Array{llvm::json::Object{{"index", 0}}};
OwningOpRef<ModuleOp> winner;
std::string scheduleReport;
llvm::json::Array attempts;
double realizationSeconds = 0, backendSeconds = 0;
size_t lowered = 0;
auto attempt = [&](size_t index) {
llvm::json::Object summary{{"candidate", static_cast<int64_t>(index)}};
std::string capturedReport;
OwningOpRef<ModuleOp> candidate;
LogicalResult realized = failure();
{
ScopedCompilerOutputSuppression silence;
realized = realizeCandidate(checkpoint, prepared->candidates[index], candidate, capturedReport, summary);
}
PimCandidateResult result = PimCandidateResult::Failure;
if (failed(realized)) {
summary["reason"] = "selected Spatial schedule failed deterministic realization";
summary["status"] = "realization-failed";
} else {
++lowered;
result = lowerPimPipelineCandidate(*candidate, resources, summary);
}
realizationSeconds += summary.getNumber("realization_seconds").value_or(0);
backendSeconds += summary.getNumber("backend_seconds").value_or(0);
if (result == PimCandidateResult::Success) {
summary["status"] = "lowered";
winner = std::move(candidate);
scheduleReport = std::move(capturedReport);
} else if (result == PimCandidateResult::ResourceLimit)
llvm::errs() << "Spatial pipeline candidate " << index
<< " rejected by a backend resource limit; rejection recorded in pipeline_selection.json\n";
attempts.push_back(std::move(summary));
return result;
};
std::string reason;
size_t selected = selectPipelineCandidate(scores, reason);
if (singleCandidate) reason = scores.front().reason;
else if (!unsupported.empty()) reason += ": " + unsupported;
else if (!scores.front().eligible() && !scores.front().reason.empty())
reason += ": " + scores.front().reason;
auto finalizeStart = Clock::now();
std::optional<size_t> selected;
if (single) {
if (attempt(0) == PimCandidateResult::Success) selected = 0;
reason = selected ? "one distinct schedule; no cost comparison required"
: "the single schedule failed backend lowering; no alternatives are available";
} else selected = lowerSelectedPimCandidate(costs, attempt, reason);
report["backend_attempts"] = std::move(attempts);
report["lowered_candidates"] = static_cast<int64_t>(lowered);
report["costed_candidates"] = single ? 0 : static_cast<int64_t>(costs.size());
size_t attemptCount = report.getArray("backend_attempts")->size();
report["backend_retries"] = static_cast<int64_t>(attemptCount ? attemptCount - 1 : 0);
report["winner_realization_seconds"] = realizationSeconds;
report["backend_lowering_seconds"] = backendSeconds;
if (!selected) {
report["status"] = "backend-failure";
report["reason"] = reason;
std::string failures;
for (auto &entry : *report.getArray("backend_attempts")) {
auto *failed = entry.getAsObject();
failures += "\ncandidate " + std::to_string(*failed->getInteger("candidate"))
+ ": " + failed->getString("reason").value_or("backend failure").str();
}
checkpoint.emitError("Spatial pipeline backend selection failed; ") << reason << failures;
writeSelectionReport(std::move(report)); signalPassFailure(); return;
}
auto finalize = Clock::now();
bool spatialArtifacts = exportStage != SpatialDataflowExportStage::None
|| pimPeriodicDiagnostics || isPimIrDumpEnabled();
// Explicit Spatial diagnostics require the pre-lowering representation.
if (!bestModule || spatialArtifacts) {
if (failed(realizeSpatialSchedule(module, prepared->candidates[selected], target,
stages, exportStage, pimPeriodicDiagnostics ? &prepared->graph : nullptr))) { signalPassFailure(); return; }
|| pimPeriodicDiagnostics || isPimIrDumpEnabled();
if (spatialArtifacts) {
if (failed(realizeSpatialSchedule(checkpoint, prepared->candidates[*selected], target,
stages, exportStage, pimPeriodicDiagnostics ? &prepared->graph : nullptr))) {
signalPassFailure(); return;
}
} else {
auto file = openDialectDumpFileWithExtension("spatial3_scheduled_no_comm", "/reports", "txt");
if (file) file << bestScheduleReport;
if (file) file << scheduleReport;
}
if (bestModule) {
assert(selected == bestIndex && "retained module must be the selected candidate");
module->setAttrs((*bestModule)->getAttrs());
module.getBodyRegion().takeBody(bestModule->getBodyRegion());
} else if (failed(lower(module))) { signalPassFailure(); return; }
llvm::errs() << "Queue-aware pipeline: candidate " << selected << " of " << scores.size()
checkpoint->setAttrs((*winner)->getAttrs());
checkpoint.getBodyRegion().takeBody(winner->getBodyRegion());
report["selection"] = static_cast<int64_t>(*selected);
report["status"] = "selected";
report["reason"] = reason;
report["finalization_seconds"] = elapsed(finalize);
report["total_seconds"] = elapsed(start);
llvm::errs() << "Spatial queue-aware pipeline: candidate " << *selected << " of " << costs.size()
<< "; " << reason << "\n";
llvm::json::Object report{{"selection", static_cast<int64_t>(selected)}, {"reason", reason},
{"static_analysis", true}, {"strategy", "queue-aware"}, {"candidates", std::move(summaries)},
{"single_candidate_analysis_skipped", singleCandidate},
{"candidate_generation_seconds", generationSeconds}, {"candidate_lowering_seconds", loweringSeconds},
{"instruction_generation_seconds", instructionSeconds},
{"queue_seconds", queueSeconds}, {"finalization_seconds", elapsed(finalizeStart)}, {"total_seconds", elapsed(start)},
{"generated_candidates", static_cast<int64_t>(scores.size())},
{"lowered_candidates", static_cast<int64_t>(lowered)},
{"reused_lowered_candidate", static_cast<bool>(bestModule)},
{"reconstructed_spatial_artifacts", spatialArtifacts && static_cast<bool>(bestModule)},
{"analyzed_candidates", static_cast<int64_t>(analyzed)},
{"unavailable_estimates", static_cast<int64_t>(unavailable)},
{"duplicate_programs", static_cast<int64_t>(duplicates)}, {"temporary_disk_bytes", 0}};
auto file = openReportFileWithExtension("pipeline_selection", "json");
if (file) { llvm::raw_os_ostream out(file); out << llvm::formatv("{0:2}\n", llvm::json::Value(std::move(report))); }
writeSelectionReport(std::move(report));
}
};
} // namespace
PimCandidateResult onnx_mlir::lowerPimPipelineCandidate(ModuleOp module,
const SpatialTargetResources &resources, llvm::json::Object &summary) {
bool classified = true, anyError = false;
std::string diagnostic;
size_t count = 0;
ScopedDiagnosticHandler handler(module.getContext(), [&](Diagnostic &message) {
if (message.getSeverity() == DiagnosticSeverity::Error) {
anyError = true;
classified &= pim::isCandidateResourceError(message);
if (++count <= 8) diagnostic += message.str().substr(0, 2048) + "\n";
}
return success();
});
auto start = Clock::now();
PassManager pm(module.getContext());
addPimLoweringPasses(pm, resources, EmitPimCodegen, /*reportProgress=*/false);
if (enableTiming) pm.addInstrumentation(std::make_unique<CandidateTiming>(summary));
LogicalResult result = pm.run(module);
summary["backend_seconds"] = elapsed(start);
if (succeeded(result)) return PimCandidateResult::Success;
summary["reason"] = diagnostic.empty() ? "backend failed without a classified diagnostic" : diagnostic;
bool resource = anyError && classified;
summary["status"] = resource ? "resource-limit" : "systemic-or-unclassified-error";
return resource ? PimCandidateResult::ResourceLimit : PimCandidateResult::Failure;
}
std::optional<size_t> onnx_mlir::lowerSelectedPimCandidate(llvm::ArrayRef<WholeScheduleCost> costs,
const std::function<PimCandidateResult(size_t)> &lower, std::string &reason) {
if (costs.empty()) {
reason = "Spatial pipeline requires at least one candidate schedule";
return std::nullopt;
}
for (const auto &cost : costs)
if (!cost.eligible()) {
reason = "cannot compare the complete candidate set: unavailable cost or unproved eligibility";
return std::nullopt;
}
std::vector<WholeScheduleCost> remaining(costs.begin(), costs.end());
std::vector<size_t> indices;
for (size_t i = 0; i < costs.size(); ++i) indices.push_back(i);
while (!indices.empty()) {
size_t selected = selectPipelineCandidate(remaining, reason);
auto result = lower(indices[selected]);
if (result == PimCandidateResult::Success) return indices[selected];
if (result != PimCandidateResult::ResourceLimit) {
reason = "systemic or unclassified backend failure; recovery is not justified";
return std::nullopt;
}
remaining.erase(remaining.begin() + selected);
indices.erase(indices.begin() + selected);
}
reason = "all candidate schedules failed backend resource constraints";
return std::nullopt;
}
std::unique_ptr<Pass> onnx_mlir::createPimPipelineSelectionPass(
const SchedulingTarget &target, const SpatialTargetResources &resources,
size_t stages, SpatialDataflowExportStage exportStage) {