nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

force selected pipeline stages in throughput mode
Validate Operations / validate-operations (push) Has been cancelled
Details

Browse Source 
increase throughput scheduler budget
add experimental throughput auto mode
This commit is contained in:
NiccoloN
2026-11-02 12:08:33 +01:00
parent 4168b588fa
commit 4a1fe710b2
35 changed files with 2047 additions and 398 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Compiler/PimPipelineSelection.cpp
+66 -12
View File
@@ -17,6 +17,7 @@
#include "llvm/Support/raw_os_ostream.h"
#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
#include <chrono>
#include <set>
using namespace mlir;
using namespace onnx_mlir;
@@ -87,7 +88,8 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
IRMapping mapping;
candidate = cast<ModuleOp>(checkpoint->clone(mapping));
auto remapped = remapSpatialSchedule(schedule, mapping);
if (failed(realizeSpatialSchedule(*candidate, remapped, target, stages,
if (failed(realizeSpatialSchedule(*candidate, remapped, target,
stages ? stages : getPipelineStageCount(schedule),
SpatialDataflowExportStage::None, nullptr, &scheduleReport))) return failure();
summary["realization_seconds"] = elapsed(begin);
return success();
@@ -142,15 +144,20 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
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];
size_t depth = getPipelineStageCount(schedule);
std::vector<size_t> cores(depth);
std::vector<double> pressure(depth);
std::vector<bool> active(depth);
for (const auto &[instance, core] : schedule.computeToCpuMap)
active[schedule.processorStages.empty() ? 0 : schedule.processorStages[core]] = true;
summary["active_stage_count"] = static_cast<int64_t>(llvm::count(active, true));
for (size_t core = 0; core < schedule.processorCount; ++core) {
size_t stage = schedule.processorStages.empty() ? 0 : schedule.processorStages[core];
++cores[stage];
if (core < cost.processorCriticalDelayNs.size())
pressure[stage] += cost.processorCriticalDelayNs[core];
}
for (size_t stage = 0; stage < stages; ++stage) {
for (size_t stage = 0; stage < depth; ++stage) {
budgets.push_back(static_cast<int64_t>(cores[stage]));
witness.push_back(pressure[stage]);
}
@@ -159,8 +166,8 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
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)];
return schedule.processorStages.empty() ? size_t{0}
: 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;
@@ -196,11 +203,16 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
costs = std::move(search->costs);
bool single = costs.empty();
if (single) {
report["candidates"] = llvm::json::Array{llvm::json::Object{{"index", 0}}};
report["candidates"] = llvm::json::Array{llvm::json::Object{{"index", 0},
{"active_stage_count", static_cast<int64_t>(getPipelineStageCount(prepared.candidates.front()))}}};
} 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])});
for (size_t i = 0; i < entries->size(); ++i) {
auto &entry = *(*entries)[i].getAsObject();
entry["parent"] = static_cast<int64_t>(search->parents[i]);
entry["expansions"] = static_cast<int64_t>(search->expansions[i]);
entry["origin"] = i < search->initialSchedules ? "depth-initialization" : "refinement";
}
}
report["generated_candidates"] = static_cast<int64_t>(prepared.candidates.size());
report["single_candidate_analysis_skipped"] = single;
@@ -212,6 +224,44 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
report["search_termination"] = search->termination;
report["queue_evaluation_budget"] = static_cast<int64_t>(kPipelineQueueEvaluationBudget);
report["beam_width"] = static_cast<int64_t>(kPipelineBeamWidth);
report["maximum_depth"] = static_cast<int64_t>(search->maximumDepth);
report["grouping_seconds"] = search->groupingSeconds;
report["initialization_segment_seconds"] = search->segmentCostSeconds;
report["initialization_dp_seconds"] = search->partitionStateSeconds;
report["initialization_placement_seconds"] = search->initializationPlacementSeconds;
report["initialization_packing_queries"] = static_cast<int64_t>(search->initializationPackingQueries);
report["initialization_dp_transitions"] = static_cast<int64_t>(search->partitionTransitions);
report["initialization_ordered_threshold_scans"] = static_cast<int64_t>(search->orderedThresholdScans);
llvm::json::Array priorities;
for (auto [depth, priority] : search->depthPriorities)
priorities.push_back(llvm::json::Object{{"depth", static_cast<int64_t>(depth)},
{"priority", static_cast<double>(priority)}});
report["depth_priorities"] = std::move(priorities);
report["frontier_policy"] = "incumbent-plus-bounded-lookahead-best-first";
if (!prepared.latencyInitializationReason.empty())
report["latency_initialization_reason"] = prepared.latencyInitializationReason;
llvm::json::Array explored;
std::set<size_t> depths;
for (const auto &schedule : prepared.candidates) depths.insert(getPipelineStageCount(schedule));
for (size_t depth : depths) explored.push_back(static_cast<int64_t>(depth));
report["explored_depths"] = std::move(explored);
llvm::json::Array domain;
for (size_t depth = stages ? stages : 1; depth <= search->maximumDepth; ++depth) {
size_t count = llvm::count_if(prepared.candidates, [&](const auto &schedule) {
return getPipelineStageCount(schedule) == depth;
});
llvm::json::Object entry{{"depth", static_cast<int64_t>(depth)},
{"evaluated_schedules", static_cast<int64_t>(single ? 0 : count)},
{"status", count ? (single ? "initialized-single-schedule" : "queue-validated") : "not-explored"}};
for (const auto &[failedDepth, reason] : search->depthInitializationFailures)
if (depth == failedDepth && !count) { entry["status"] = "initialization-failed"; entry["reason"] = reason; }
if (depth == 1 && !prepared.latencyInitializationReason.empty()) {
entry["status"] = "initialization-failed";
entry["reason"] = prepared.latencyInitializationReason;
}
domain.push_back(std::move(entry));
}
report["depth_search"] = std::move(domain);
return success();
}
@@ -229,6 +279,8 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
{"generated_candidates", static_cast<int64_t>(prepared->candidates.size())},
{"selection_margin", kPipelineSelectionMargin},
{"prediction_error_bound_available", false}, {"temporary_disk_bytes", 0}};
report["requested_depth"] = stages ? llvm::json::Value(static_cast<int64_t>(stages))
: llvm::json::Value("auto");
if (pimDisableSynchronization) {
report["synchronization_ablation"] = true;
report["cost_model_program"] = "synchronized-reference";
@@ -308,7 +360,8 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
|| pimPeriodicDiagnostics || isPimIrDumpEnabled();
if (spatialArtifacts) {
if (failed(realizeSpatialSchedule(checkpoint, prepared->candidates[*selected], target,
stages, exportStage, pimPeriodicDiagnostics ? &prepared->graph : nullptr))) {
stages ? stages : getPipelineStageCount(prepared->candidates[*selected]),
exportStage, pimPeriodicDiagnostics ? &prepared->graph : nullptr))) {
signalPassFailure(); return;
}
} else {
@@ -318,6 +371,7 @@ struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, Op
checkpoint->setAttrs((*winner)->getAttrs());
checkpoint.getBodyRegion().takeBody(winner->getBodyRegion());
report["selection"] = static_cast<int64_t>(*selected);
report["selected_depth"] = static_cast<int64_t>(getPipelineStageCount(prepared->candidates[*selected]));
report["status"] = "selected";
report["reason"] = reason;
report["finalization_seconds"] = elapsed(finalize);