better throughput scheduling with steady state cost model (WIP)
This commit is contained in:
NiccoloN
@@ -0,0 +1,225 @@
|
||||
#include "PimPipelineSelection.hpp"
|
||||
#include "PimCompilerOptions.hpp"
|
||||
#include "PimCompilerUtils.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"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
|
||||
#include "mlir/Pass/PassManager.h"
|
||||
#include "mlir/Pass/PassInstrumentation.h"
|
||||
#include "src/Compiler/CompilerOptions.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
|
||||
#include "llvm/Support/FormatVariadic.h"
|
||||
#include "llvm/Support/JSON.h"
|
||||
#include "llvm/Support/raw_os_ostream.h"
|
||||
#include <algorithm>
|
||||
#include <chrono>
|
||||
|
||||
using namespace mlir;
|
||||
using namespace onnx_mlir;
|
||||
using namespace onnx_mlir::spatial;
|
||||
|
||||
namespace {
|
||||
using Clock = std::chrono::steady_clock;
|
||||
double elapsed(Clock::time_point start) {
|
||||
return std::chrono::duration<double>(Clock::now() - start).count();
|
||||
}
|
||||
|
||||
llvm::json::Object countIR(Operation *root) {
|
||||
int64_t operations = 0, values = 0;
|
||||
root->walk([&](Operation *op) {
|
||||
++operations;
|
||||
values += op->getNumResults();
|
||||
for (Region ®ion : op->getRegions())
|
||||
for (Block &block : region) values += block.getNumArguments();
|
||||
});
|
||||
return llvm::json::Object{{"operations", operations}, {"values", values}};
|
||||
}
|
||||
|
||||
// Enabled with the existing compiler timing option. No instruction expansion.
|
||||
struct CandidateTiming final : PassInstrumentation {
|
||||
llvm::json::Object &summary;
|
||||
Clock::time_point start;
|
||||
explicit CandidateTiming(llvm::json::Object &summary) : summary(summary) {}
|
||||
void runBeforePass(Pass *pass, Operation *op) override {
|
||||
if (pass->getArgument() == "pim-one-shot-bufferization")
|
||||
summary["bufferization_input_ir"] = countIR(op);
|
||||
start = Clock::now();
|
||||
}
|
||||
void runAfterPass(Pass *pass, Operation *) override {
|
||||
auto *phases = summary.getObject("lowering_phases_seconds");
|
||||
if (!phases) {
|
||||
summary["lowering_phases_seconds"] = llvm::json::Object{};
|
||||
phases = summary.getObject("lowering_phases_seconds");
|
||||
}
|
||||
(*phases)[pass->getName()] = elapsed(start);
|
||||
}
|
||||
};
|
||||
|
||||
struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, OperationPass<ModuleOp>> {
|
||||
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimPipelineSelectionPass)
|
||||
PimPipelineSelectionPass(const SchedulingTarget &target, const SpatialTargetResources &resources,
|
||||
size_t stages, SpatialDataflowExportStage exportStage)
|
||||
: target(target), resources(resources), stages(stages), exportStage(exportStage) {}
|
||||
SchedulingTarget target;
|
||||
SpatialTargetResources resources;
|
||||
size_t stages;
|
||||
SpatialDataflowExportStage exportStage;
|
||||
|
||||
StringRef getArgument() const override { return "pim-pipeline-selection"; }
|
||||
StringRef getDescription() const override { return "Select a throughput pipeline using static instruction 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) {
|
||||
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);
|
||||
return success();
|
||||
}
|
||||
|
||||
void runOnOperation() override {
|
||||
ModuleOp module = 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());
|
||||
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));
|
||||
}
|
||||
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();
|
||||
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; }
|
||||
} else {
|
||||
auto file = openDialectDumpFileWithExtension("spatial3_scheduled_no_comm", "/reports", "txt");
|
||||
if (file) file << bestScheduleReport;
|
||||
}
|
||||
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()
|
||||
<< "; " << 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))); }
|
||||
}
|
||||
};
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<Pass> onnx_mlir::createPimPipelineSelectionPass(
|
||||
const SchedulingTarget &target, const SpatialTargetResources &resources,
|
||||
size_t stages, SpatialDataflowExportStage exportStage) {
|
||||
return std::make_unique<PimPipelineSelectionPass>(target, resources, stages, exportStage);
|
||||
}
|
||||
Reference in New Issue
Block a user