better throughput scheduling with steady state cost model (WIP)

This commit is contained in:
NiccoloN
2026-10-23 11:27:06 +02:00
parent 48e9fb3795
commit fd493ae6e8
51 changed files with 3581 additions and 247 deletions
+106 -29
View File
@@ -14,6 +14,9 @@
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
#include "PimPipelineSelection.hpp"
#include "PimInstructionTiming.hpp"
#include "PimPeriodicAnalysis.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
@@ -174,8 +177,9 @@ std::pair<size_t, size_t> getConfigPair(const llvm::json::Object& object,
}
void loadPimInterProcessorLatencies(
spatial::SchedulingTarget& target,
const llvm::json::Object& network) {
PimInstructionTiming& target,
const llvm::json::Object& network,
std::optional<int64_t> memorySwitch) {
std::optional<llvm::StringRef> filename =
network.getString("net_config_file_path");
if (!filename)
@@ -240,10 +244,29 @@ void loadPimInterProcessorLatencies(
? 0
: (latencySum + static_cast<Cost>(pairCount) - 1)
/ static_cast<Cost>(pairCount);
if (memorySwitch) {
const auto *memoryRow = latencies->getObject(std::to_string(*memorySwitch));
for (size_t core = 0; core < target.processorCount; ++core) {
const auto *row = latencies->getObject(std::to_string(core));
auto request = row ? row->getNumber(std::to_string(*memorySwitch)) : std::nullopt;
auto response = memoryRow ? memoryRow->getNumber(std::to_string(core)) : std::nullopt;
if (!request || !response || !std::isfinite(*request)
|| !std::isfinite(*response) || *request < 0 || *response < 0
|| *request >= static_cast<double>(std::numeric_limits<Cost>::max())
|| *response >= static_cast<double>(std::numeric_limits<Cost>::max())) {
target.sharedMemoryRequestLatencyNs.clear();
target.sharedMemoryResponseLatencyNs.clear();
break;
}
target.sharedMemoryRequestLatencyNs.push_back(std::ceil(*request));
target.sharedMemoryResponseLatencyNs.push_back(std::ceil(*response));
}
}
}
spatial::SchedulingTarget getPimSchedulingTarget() {
spatial::SchedulingTarget target = getDefaultPimSchedulingTarget();
PimInstructionTiming loadPimInstructionTiming() {
PimInstructionTiming target;
static_cast<spatial::SchedulingTarget &>(target) = getDefaultPimSchedulingTarget();
if (pimTargetConfig.empty())
return target;
@@ -298,12 +321,19 @@ spatial::SchedulingTarget getPimSchedulingTarget() {
|| target.matrixColumns != crossbarSize.getValue())
llvm::report_fatal_error("Pim target config resources do not match --core-count, "
"--crossbar-count, and --crossbar-size");
loadPimInterProcessorLatencies(target, network);
loadPimInterProcessorLatencies(target, network,
core.getInteger("global_memory_switch_id"));
target.processorPeriodNs =
getConfigCost(core, "period", target.processorPeriodNs);
target.signalIssueNs = target.processorPeriodNs;
target.controlMessageBytes = 1;
target.sharedMemoryReadReturnsPayload = false;
target.peerPermissionThenSourceRead = true;
target.localMemoryWidthBytes =
getConfigCost(localMemory, "data_width", target.localMemoryWidthBytes);
target.localMemoryPeriodNs =
getConfigCost(localMemory, "period", target.localMemoryPeriodNs);
target.localMemoryReadLatencyCycles =
getConfigCost(localMemory, "read_latency_cycle", target.localMemoryReadLatencyCycles);
target.localMemoryWriteLatencyCycles =
@@ -345,6 +375,65 @@ spatial::SchedulingTarget getPimSchedulingTarget() {
} // namespace
PimInstructionTiming getPimInstructionTiming() {
return loadPimInstructionTiming();
}
spatial::SchedulingTarget getPimSchedulingTarget() {
auto timing = getPimInstructionTiming();
auto target = static_cast<spatial::SchedulingTarget>(timing);
configurePimScheduleCosts(target, std::move(timing));
return target;
}
llvm::Expected<int64_t> getPimPeriodicTimingContract() {
auto buffer = llvm::MemoryBuffer::getFile(pimTargetConfig);
if (!buffer) return llvm::createStringError("queue selection requires --pim-target-config with instruction timing");
auto parsed = llvm::json::parse(buffer.get()->getBuffer());
if (!parsed) return parsed.takeError();
const auto *root = parsed->getAsObject();
const auto *chip = root ? root->getObject("chip_config") : nullptr;
const auto *core = chip ? chip->getObject("core_config") : nullptr;
const auto *matrix = core ? core->getObject("matrix_config") : nullptr;
if (!matrix || core->getInteger("rob_size").value_or(1) != 1)
return llvm::createStringError("queue selection requires a target with ROB capacity one and matrix timing");
auto bits = matrix->getInteger("cell_precision").value_or(2);
if (bits <= 0) return llvm::createStringError("matrix cell precision must be positive");
return bits;
}
void addPimLoweringPasses(PassManager &pm,
const spatial::SpatialTargetResources &targetResources, PimEmissionTargetType target,
bool reportProgress) {
auto message = [&](llvm::StringRef text) {
if (reportProgress) pm.addPass(createMessagePass(text.str()));
};
if (target >= EmitPim) {
pm.addPass(createSpatialToPimPass(targetResources));
message("Spatial lowered to Pim");
}
if (target >= EmitPimBufferized) {
pm.addPass(createPimBufferizationPreparationPass(pimVerifyBufferizationCopyFreedom.getValue()));
pm.addPass(createPimOneShotBufferizationPass());
pm.addPass(createPimMemoryNormalizationPass());
pm.addPass(createPimBufferizationVerificationPass());
message("Pim bufferized");
}
if (target >= EmitPimCodegen) {
pm.addPass(mlir::createLowerAffinePass());
pm.addPass(createPimHostConstantFoldingPass());
message("Pim host constants folded");
pm.addPass(createPimInstructionSelectionPass());
message("Pim instructions selected");
pm.addPass(createPimLocalMemoryPlanningPass());
message("Pim local memory planned");
pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue()));
message("Pim verified");
}
}
void addPassesPim(OwningOpRef<ModuleOp>& module,
PassManager& pm,
EmissionTargetType& emissionTarget,
@@ -374,34 +463,22 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
pm.addPass(createLowerSpatialPlansPass(targetResources, planningOptions, exportStage));
pm.addPass(createTrivialGraphComputeMergePass(
schedulingTarget.residentWeightCapacity, exportStage));
pm.addPass(spatial::createScheduleAndRealizeSpatialPass(
schedulingTarget, exportStage, pipelineStages.getValue()));
if (pipelineStages.getValue() > 1 && pimEmissionTarget >= EmitPimCodegen) {
pm.addPass(createPimPipelineSelectionPass(schedulingTarget, targetResources,
pipelineStages.getValue(), exportStage));
pm.addPass(createEmitPimCodePass());
return;
} else {
if (pipelineStages.getValue() > 1)
pm.addPass(createMessagePass("Queue-aware pipeline: intermediate emission retains the original schedule without backend evaluation"));
pm.addPass(spatial::createScheduleAndRealizeSpatialPass(
schedulingTarget, exportStage, pipelineStages.getValue(), pimPeriodicDiagnostics.getValue()));
}
pm.addPass(createMessagePass("Onnx lowered to Spatial"));
}
if (pimEmissionTarget >= EmitPim) {
pm.addPass(createSpatialToPimPass(targetResources));
pm.addPass(createMessagePass("Spatial lowered to Pim"));
}
if (pimEmissionTarget >= EmitPimBufferized) {
pm.addPass(createPimBufferizationPreparationPass(pimVerifyBufferizationCopyFreedom.getValue()));
pm.addPass(createPimOneShotBufferizationPass());
pm.addPass(createPimMemoryNormalizationPass());
pm.addPass(createPimBufferizationVerificationPass());
pm.addPass(createMessagePass("Pim bufferized"));
}
addPimLoweringPasses(pm, targetResources, pimEmissionTarget);
if (pimEmissionTarget >= EmitPimCodegen) {
pm.addPass(mlir::createLowerAffinePass());
pm.addPass(createPimHostConstantFoldingPass());
pm.addPass(createMessagePass("Pim host constants folded"));
pm.addPass(createPimInstructionSelectionPass());
pm.addPass(createMessagePass("Pim instructions selected"));
pm.addPass(createPimLocalMemoryPlanningPass());
pm.addPass(createMessagePass("Pim local memory planned"));
pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue()));
pm.addPass(createMessagePass("Pim verified"));
pm.addPass(createEmitPimCodePass());
pm.addPass(createMessagePass("Pim code emitted"));
}