Files
Raptor/test/PIM/PipelineOptimizerTest.cpp
NiccoloN 4a1fe710b2
Validate Operations / validate-operations (push) Has been cancelled
force selected pipeline stages in throughput mode
increase throughput scheduler budget
add experimental throughput auto mode
2026-11-02 12:08:33 +01:00

1094 lines
55 KiB
C++

#include "TestCheck.hpp"
#include "src/Accelerators/PIM/Compiler/PimPeriodicAnalysis.hpp"
#include "src/Accelerators/PIM/Compiler/PimInstructionTiming.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ResidentWeightPacking.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelinePackingCostCache.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include <chrono>
#include <cmath>
#include <functional>
#include <map>
#include <numeric>
#include <random>
#include <set>
#include <iostream>
using namespace onnx_mlir;
using namespace onnx_mlir::spatial;
using E = PimExecutionEvent;
static bool sawJointMove = false, sawSuccessiveMoves = false;
static void testPackingCostCacheEviction() {
ComputeGraph graph;
const std::vector<Cost> work{0, 11, 23, 7, 19, 3};
for (uint32_t task = 0; task < work.size(); ++task) {
ResidentWeight weight;
weight.opaqueLane = task % 3;
graph.nodes.push_back({{nullptr, task, 1}, work[task], {weight}, task});
}
IndexedResidentWeights weights(graph);
PipelinePackingCostCache cache(weights, work, 1, 8);
static_assert(PipelinePackingCostCache::kMaximumEntries == 256);
std::vector<size_t> tasks{0, 1, 2, 3, 4, 5};
auto check = [&](const auto &order) {
TEST_CHECK(cache.minimumCores(order) == weights.getPackCountLowerBound(order, 1));
for (size_t cores : {1u, 2u, 3u, 5u, 8u}) {
Cost expected = weights.estimatePackCost(order, work, 1, cores)
.value_or(std::numeric_limits<Cost>::max());
TEST_CHECK(cache.estimate(order, cores) == expected);
TEST_CHECK(cache.estimate(order, cores) == expected);
}
TEST_CHECK(cache.size() <= PipelinePackingCostCache::kMaximumEntries);
};
size_t orders = 0;
do {
check(tasks);
++orders;
} while (std::next_permutation(tasks.begin(), tasks.end()));
TEST_CHECK(orders == 720 && cache.size() == PipelinePackingCostCache::kMaximumEntries);
// Revisit early entries after eviction. Ordered and constructive queries must agree.
check(std::vector<size_t>{0, 1, 2, 3, 4, 5});
check(std::vector<size_t>{0, 3, 1, 4, 2, 5});
check(std::vector<size_t>{0});
check(std::vector<size_t>{});
}
static void testPackingFeasibility() {
ComputeGraph graph;
for (size_t task = 0; task < 4; ++task) {
ResidentWeight weight;
weight.opaqueLane = task % 2;
graph.nodes.push_back({{nullptr, static_cast<uint32_t>(task), 1}, 10, {weight}, task});
}
IndexedResidentWeights weights(graph);
const std::vector<size_t> tasks{0, 1, 2, 3};
const std::vector<Cost> work{10, 10, 10, 10};
std::vector<size_t> assignment{99};
TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
tasks, 1, 2, &assignment));
TEST_CHECK(assignment.empty());
TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
tasks, 1, 0));
TEST_CHECK(weights.estimatePackCost(tasks, work, 1, 2) == 20);
auto placement = weights.pack(tasks, work, 1, 2);
TEST_CHECK(placement && (*placement)[0] == (*placement)[2]
&& (*placement)[1] == (*placement)[3] && (*placement)[0] != (*placement)[1]);
TEST_CHECK(!weights.estimatePackCost(tasks, work, 1, 1));
for (size_t cores = 4; cores <= 8; ++cores)
TEST_CHECK(estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
tasks, 1, cores) == 10);
weights.taskWeightIds[0] = {0, 1};
TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
tasks, 1, 4));
weights.taskWeightIds = {{0, 1}, {2}, {0, 1}, {2}};
weights.weightCount = 3;
TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
tasks, 2, 2));
TEST_CHECK(weights.estimatePackCost(tasks, work, 2, 2) == 20);
}
static void testPackingCurves() {
std::mt19937 random(1739);
for (size_t fixture = 0; fixture < 150; ++fixture) {
size_t count = fixture % 9, capacity = 1 + fixture % 3;
ComputeGraph graph;
std::vector<Cost> work;
std::vector<size_t> tasks;
for (uint32_t task = 0; task < count; ++task) {
ResidentWeightSet weights;
for (size_t id = 0; id < 4; ++id)
if (random() % 3 == 0) { ResidentWeight w; w.opaqueLane = id; weights.push_back(w); }
work.push_back(random() % 40);
tasks.push_back(task);
graph.nodes.push_back({{nullptr, task, 1}, work.back(), weights, task});
}
IndexedResidentWeights weights(graph);
auto curve = weights.estimatePackCostCurve(tasks, work, capacity, 8);
auto sparse = weights.estimatePackCostCurve(tasks, work, capacity, 8, nullptr, {1, 3, 6, 8});
for (size_t cores : {1u, 3u, 6u, 8u}) TEST_CHECK(sparse[cores] == curve[cores]);
for (size_t cores : {2u, 4u, 5u, 7u}) TEST_CHECK(!sparse[cores]);
for (size_t cores = 1; cores <= 8; ++cores) {
TEST_CHECK(curve[cores] == weights.estimatePackCost(tasks, work, capacity, cores));
std::optional<Cost> reference;
for (size_t mask = 0; mask < (size_t{1} << (count ? count - 1 : 0)); ++mask) {
size_t used = count ? 1 : 0;
Cost maximum = 0, current = 0;
std::set<size_t> residents;
bool fits = true;
for (size_t task = 0; task < count; ++task) {
if (task && ((mask >> (task - 1)) & 1)) {
maximum = std::max(maximum, current); current = 0; residents.clear(); ++used;
}
current += work[task];
for (size_t id : weights.taskWeightIds[task]) residents.insert(id);
fits &= residents.size() <= capacity;
}
maximum = std::max(maximum, current);
if (fits && used <= cores && (!reference || maximum < *reference)) reference = maximum;
}
TEST_CHECK(estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
tasks, capacity, cores) == reference);
}
}
}
static void testQuantizedFeasibility(uint32_t first) {
mlir::MLIRContext context;
context.loadDialect<spatial::SpatialDialect>();
auto module = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
mlir::OpBuilder builder(module.getBodyRegion());
ComputeGraph graph;
graph.predecessors.resize(40); graph.successors.resize(40);
for (uint32_t group = 0; group < 2; ++group) {
const uint32_t lanes = group ? 40 - first : first;
auto batch = spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
mlir::TypeRange{}, builder.getI32IntegerAttr(lanes), mlir::ValueRange{}, mlir::ValueRange{});
for (uint32_t lane = 0; lane < lanes; ++lane) {
size_t task = graph.nodes.size();
ResidentWeight weight; weight.opaqueLane = task;
graph.nodes.push_back({{batch, lane, 1}, 10, {weight}, task});
if (group) {
graph.predecessors[task].push_back({0, {0, 0}});
graph.successors[0].push_back({task, {0, 0}});
}
}
}
SchedulingTarget target;
target.processorCount = 40; target.residentWeightCapacity = 1;
target.interProcessorLatencyNs.assign(1600, 0);
MergeScheduleResult schedule;
std::string error;
TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
std::vector<size_t> budgets(2);
for (size_t stage : schedule.processorStages) ++budgets[stage];
TEST_CHECK(budgets == std::vector<size_t>({first, 40 - first}));
for (size_t task = 0; task < graph.nodes.size(); ++task)
TEST_CHECK(schedule.processorStages[schedule.computeToCpuMap.lookup(graph.nodes[task].instance)]
== (task < first ? 0 : 1));
auto query = [](const MergeScheduleResult &) {
WholeScheduleCost cost;
cost.primary = cost.alternative = 100;
cost.selectionMargin = kPipelineSelectionMargin;
cost.loweringSupport = StaticLoweringSupport::Supported;
for (auto &proof : cost.recurrence) proof = {ScheduleRecurrenceEvidence::Proved, 1, 100, {}};
return cost;
};
auto depthOne = schedule;
depthOne.processorStages.assign(target.processorCount, 0);
auto search = optimizePipelineSchedule(graph, depthOne, 0, target, query, error, 2);
TEST_CHECK(mlir::succeeded(search) && search->initialSchedules == 2);
TEST_CHECK(search->schedules[1].processorStages == schedule.processorStages);
TEST_CHECK(search->schedules[1].computeToCpuMap == schedule.computeToCpuMap);
module.erase();
}
// Test-only workloads use stage-boundary host transfers and local direct links.
// The same queue
// engine and complete-state evidence rank reference and searched schedules.
static WholeScheduleCost queueCost(const ComputeGraph &graph,
const MergeScheduleResult &schedule, const PimInstructionTiming &target,
bool restartProtocol = false) {
PimExecutionProgram program(target.processorCount);
std::vector<std::vector<size_t>> tasks(target.processorCount);
for (size_t task = 0; task < graph.nodes.size(); ++task)
tasks[schedule.computeToCpuMap.lookup(graph.nodes[task].instance)].push_back(task);
for (auto &core : tasks)
llvm::sort(core, [&](size_t a, size_t b) {
return schedule.computeToCpuSlotMap.lookup(graph.nodes[a].instance)
< schedule.computeToCpuSlotMap.lookup(graph.nodes[b].instance);
});
struct Transfer { size_t source, target, reader, writer, reg; bool host; };
std::vector<Transfer> transfers;
std::vector<std::set<size_t>> readers(target.processorCount), writers(target.processorCount);
// Branched synthetic workloads use host routing to avoid inventing a second
// blocking-transfer scheduler in this test-only objective.
bool branching = llvm::any_of(graph.successors, [](const auto &edges) { return edges.size() > 1; });
for (size_t task = 0; task < graph.nodes.size(); ++task)
for (const auto &[predecessor, dependency] : graph.predecessors[task]) {
size_t writer = schedule.computeToCpuMap.lookup(graph.nodes[predecessor].instance);
size_t reader = schedule.computeToCpuMap.lookup(graph.nodes[task].instance);
if (writer == reader) continue;
bool host = branching || schedule.processorStages[writer] != schedule.processorStages[reader];
transfers.push_back({predecessor, task, reader, writer, 1 + transfers.size(), host});
if (host) { readers[writer].insert(reader); writers[reader].insert(writer); }
}
std::vector<size_t> zero, downstream;
if (restartProtocol)
for (size_t core = 0; core < program.size(); ++core)
(schedule.processorStages[core] ? downstream : zero).push_back(core);
for (size_t core = 0; core < program.size(); ++core) {
auto &events = program[core];
for (size_t writer : writers[core]) events.emplace_back(E::Signal, 0, writer, 0, 1);
size_t rank = std::find(downstream.begin(), downstream.end(), core) - downstream.begin();
if (restartProtocol && rank < downstream.size())
events.emplace_back(E::Signal, 0, rank ? downstream[(rank - 1) / 2] : zero.front(), rank ? 29 : 30, 1);
if (!restartProtocol && schedule.processorStages[core] == 0) {
size_t peers = 0;
for (size_t peer = 0; peer < program.size(); ++peer)
if (peer != core && schedule.processorStages[peer] == 0) {
events.emplace_back(E::Signal, 0, peer, 100, 1); ++peers;
}
if (peers) events.emplace_back(E::Wait, 0, 0, 100, peers);
}
bool reuseWait = false;
for (size_t task : tasks[core]) {
for (const auto &transfer : transfers)
if (transfer.target == task) {
if (transfer.host) {
events.emplace_back(E::Wait, 0, 0, transfer.reg, 1);
events.emplace_back(E::Load, 64);
} else events.emplace_back(E::Receive, 64, transfer.writer);
}
E work(E::Work, graph.nodes[task].cost); work.workClass = E::Vector;
events.push_back(work);
for (const auto &transfer : transfers)
if (transfer.source == task) {
if (!transfer.host) {
events.emplace_back(E::Send, 64, transfer.reader); continue;
}
if (!reuseWait) {
events.emplace_back(E::Wait, 0, 0, 0, readers[core].size()); reuseWait = true;
}
events.emplace_back(E::Store, 64);
events.emplace_back(E::Signal, 0, transfer.reader, transfer.reg, 1);
}
}
if (restartProtocol && schedule.processorStages[core] == 0) {
size_t rankZero = std::find(zero.begin(), zero.end(), core) - zero.begin();
for (size_t step = 1, reg = 16; step < zero.size(); step *= 2, ++reg) {
events.emplace_back(E::Signal, 0, zero[(rankZero + step) % zero.size()], reg, 1);
events.emplace_back(E::Wait, 0, 0, reg, 1);
}
if (core == zero.front() && !downstream.empty()) {
events.emplace_back(E::Wait, 0, 0, 30, 1);
events.emplace_back(E::Signal, 0, downstream.front(), 28, 1);
}
} else if (restartProtocol) {
events.emplace_back(E::Wait, 0, 0, 28, 1);
size_t children = (2 * rank + 1 < downstream.size()) + (2 * rank + 2 < downstream.size());
if (children) events.emplace_back(E::Wait, 0, 0, 29, children);
for (size_t child : {2 * rank + 1, 2 * rank + 2})
if (child < downstream.size()) events.emplace_back(E::Signal, 0, downstream[child], 28, 1);
}
}
if (restartProtocol)
for (const auto &sequence : program)
for (const auto &event : sequence)
if (event.kind == E::Signal || event.kind == E::Wait)
TEST_CHECK(event.reg >= 0 && static_cast<size_t>(event.reg) < target.synchronizationRegisterCount);
PimPeriodicOptions options;
options.verifyPeriodicState = true; options.searchGuidance = !restartProtocol;
auto estimate = evaluatePimExecution(program, target, options);
auto cost = std::move(estimate.cost);
cost.loweringSupport = StaticLoweringSupport::Supported;
cost.selectionMargin = 0.021;
TEST_CHECK(cost.eligible());
return cost;
}
static PimInstructionTiming hardware(size_t profile) {
PimInstructionTiming target;
target.processorCount = 5;
target.residentWeightCapacity = 2;
target.estimatePipelineDependencySpan = [](llvm::ArrayRef<Cost> spans) { return spans.front(); };
target.peerPermissionThenSourceRead = true;
target.signalIssueNs = profile + 1;
target.controlMessageBytes = 1;
target.sharedMemoryWidthBytes = 64;
target.sharedMemoryReadLatencyNs = 17 + profile * 9;
target.sharedMemoryWriteLatencyNs = 23 + profile * 7;
target.synchronizationRegisterCount = 128;
target.sharedMemoryRequestLatencyNs.assign(5, 2);
target.sharedMemoryResponseLatencyNs.assign(5, 3);
target.interProcessorLatencyNs.resize(25);
for (size_t a = 0; a < 5; ++a)
for (size_t b = 0; b < 5; ++b) {
const size_t irregular[] = {0, 1, 5, 6, 20};
size_t x = profile ? irregular[a] : a, y = profile ? irregular[b] : b;
target.interProcessorLatencyNs[a * 5 + b] = x > y ? x - y : y - x;
}
return target;
}
static void testInitializerPackingWitness() {
mlir::MLIRContext context;
context.loadDialect<spatial::SpatialDialect>();
auto module = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
mlir::OpBuilder builder(module.getBodyRegion());
std::vector<mlir::Operation *> operations;
for (size_t i = 0; i < 3; ++i) operations.push_back(
spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
mlir::TypeRange{}, builder.getI32IntegerAttr(3), mlir::ValueRange{}, mlir::ValueRange{}));
const std::vector<std::vector<size_t>> resident{{2,1,5},{2,4},{0,2},{0},{4,6},{1},{2},{2},{1,4}};
const std::vector<Cost> costs{4,1,5,4,1,2,5,2,2};
ComputeGraph graph;
graph.predecessors.resize(9); graph.successors.resize(9);
for (uint32_t task = 0; task < 9; ++task) {
ResidentWeightSet weights;
for (size_t id : resident[task]) {
ResidentWeight w; w.opaqueLane = id; weights.push_back(w);
}
graph.nodes.push_back({{operations[task % 3], task / 3, 1}, costs[task], weights, task});
}
IndexedResidentWeights indexed(graph);
TEST_CHECK(!indexed.estimatePackCost({1,2,4,5,7,8}, costs, 3, 2));
std::vector<size_t> witness;
TEST_CHECK(indexed.estimatePackCost({1,4,7,2,5,8}, costs, 3, 2, &witness));
TEST_CHECK(witness.size() == 6);
auto target = hardware(1);
target.processorCount = 4; target.residentWeightCapacity = 3;
target.interProcessorLatencyNs.assign(16, 1);
MergeScheduleResult schedule;
std::string error;
TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
std::vector<ResidentWeightSet> actual(4);
for (const auto &node : graph.nodes) {
size_t core = schedule.computeToCpuMap.lookup(node.instance);
TEST_CHECK(core < 4);
insertResidentWeights(actual[core], node.residentWeights);
TEST_CHECK(actual[core].size() <= 3);
}
// Exhaust all core assignments under the same stage map. The initialized
// assignment must belong to this physical-capacity reference set.
size_t feasible = 0;
for (size_t encoded = 0; encoded < (size_t{1} << 18); ++encoded) {
size_t choices = encoded;
std::vector<ResidentWeightSet> weights(4);
std::vector<size_t> groupStages(3, 2);
bool valid = true;
for (size_t task = 0; task < 9 && valid; ++task) {
size_t core = choices % 4; choices /= 4;
size_t stage = schedule.processorStages[core], group = task % 3;
if (groupStages[group] == 2) groupStages[group] = stage;
valid &= groupStages[group] == stage;
insertResidentWeights(weights[core], graph.nodes[task].residentWeights);
valid &= weights[core].size() <= 3;
}
feasible += valid;
}
TEST_CHECK(feasible > 0);
MergeScheduleResult alternative;
TEST_CHECK(mlir::failed(packPipelineSchedule(graph, std::vector<size_t>(9),
{3, 1}, target, alternative, error)));
TEST_CHECK(error.find("nonempty compute stages") != std::string::npos);
std::cout << "initializer_physical_feasible_assignments=" << feasible << '\n';
module.erase();
}
static void testPhysicalReference() {
auto target = hardware(1);
target.processorCount = 3;
target.sharedMemoryRequestLatencyNs.resize(3);
target.sharedMemoryResponseLatencyNs.resize(3);
target.interProcessorLatencyNs = {0, 2, 11, 2, 0, 7, 11, 7, 0};
ComputeGraph graph;
graph.predecessors.resize(3); graph.successors.resize(3);
for (uint32_t task = 0; task < 3; ++task) {
ResidentWeight weight; weight.opaqueLane = task % 2;
graph.nodes.push_back({{nullptr, task, 1}, task == 1 ? 700u : 80u, {weight}, task});
if (task) {
graph.predecessors[task].push_back({task - 1, {64, 1}});
graph.successors[task - 1].push_back({task, {64, 1}});
}
}
double optimum = std::numeric_limits<double>::infinity();
size_t evaluations = 0;
for (size_t mask = 1; mask < 7; ++mask)
for (size_t placement = 0; placement < 27; ++placement) {
MergeScheduleResult schedule;
for (size_t core = 0; core < 3; ++core) schedule.processorStages.push_back((mask >> core) & 1);
size_t encoded = placement;
std::vector<size_t> stages, slots(3);
std::vector<ResidentWeightSet> resident(3);
bool valid = true;
for (const auto &node : graph.nodes) {
size_t core = encoded % 3; encoded /= 3;
schedule.computeToCpuMap[node.instance] = core;
schedule.computeToCpuSlotMap[node.instance] = slots[core]++;
stages.push_back(schedule.processorStages[core]);
insertResidentWeights(resident[core], node.residentWeights);
valid &= resident[core].size() <= target.residentWeightCapacity;
}
if (!valid || stages[0] != 0 || stages[1] < stages[0] || stages[2] < stages[1]
|| !llvm::is_contained(stages, 1)) continue;
auto cost = queueCost(graph, schedule, target);
optimum = std::min(optimum, std::max(*cost.primary, *cost.alternative));
++evaluations;
}
std::string error;
MergeScheduleResult initial;
TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, initial, 2, target, error)));
auto evaluate = [&](const MergeScheduleResult &schedule) { return queueCost(graph, schedule, target); };
auto search = optimizePipelineSchedule(graph, initial, 2, target, evaluate, error);
TEST_CHECK(mlir::succeeded(search));
size_t winner = search->costs.empty() ? 0 : selectPipelineCandidate(search->costs, error);
auto selected = evaluate(search->schedules[winner]);
double value = std::max(*selected.primary, *selected.alternative);
TEST_CHECK(value + 1e-6 >= optimum);
std::cout << "physical_reference_ns=" << optimum << " selected_ns=" << value
<< " gap_percent=" << 100 * (value / optimum - 1)
<< " physical_evaluations=" << evaluations
<< " search_evaluations=" << search->costs.size()
<< " packing_attempts=" << search->packingAttempts << '\n';
}
static void testAutoPhysicalReference() {
for (size_t profile = 0; profile < 2; ++profile) {
auto target = hardware(profile);
target.processorCount = 3;
target.sharedMemoryRequestLatencyNs.resize(3); target.sharedMemoryResponseLatencyNs.resize(3);
target.interProcessorLatencyNs = {0, 2, 11, 2, 0, 7, 11, 7, 0};
ComputeGraph graph;
graph.predecessors.resize(3); graph.successors.resize(3);
for (uint32_t task = 0; task < 3; ++task) {
ResidentWeight weight; weight.opaqueLane = task % 2;
graph.nodes.push_back({{nullptr, task, 1}, task == 1 ? 700u : 80u, {weight}, task});
if (task) {
graph.predecessors[task].push_back({task - 1, {64, 1}});
graph.successors[task - 1].push_back({task, {64, 1}});
}
}
double reference = std::numeric_limits<double>::infinity();
size_t referenceDepth = 0, evaluations = 0;
auto start = std::chrono::steady_clock::now();
for (size_t depth = 1; depth <= 3; ++depth) {
size_t maps = depth * depth * depth;
for (size_t encodedMap = 0; encodedMap < maps; ++encodedMap)
for (size_t placement = 0; placement < 27; ++placement) {
MergeScheduleResult schedule;
size_t map = encodedMap, encoded = placement;
schedule.processorCount = 3;
for (size_t core = 0; core < 3; ++core) { schedule.processorStages.push_back(map % depth); map /= depth; }
std::vector<size_t> stages, slots(3);
std::vector<ResidentWeightSet> resident(3);
bool valid = true;
for (const auto &node : graph.nodes) {
size_t core = encoded % 3; encoded /= 3;
schedule.computeToCpuMap[node.instance] = core;
schedule.computeToCpuSlotMap[node.instance] = slots[core]++;
stages.push_back(schedule.processorStages[core]);
insertResidentWeights(resident[core], node.residentWeights);
valid &= resident[core].size() <= target.residentWeightCapacity;
}
std::set<size_t> active(stages.begin(), stages.end());
if (!valid || active.size() != depth || stages[0] != 0 || stages[1] < stages[0]
|| stages[1] > stages[0] + 1 || stages[2] < stages[1] || stages[2] > stages[1] + 1) continue;
auto cost = queueCost(graph, schedule, target);
double value = std::max(*cost.primary, *cost.alternative);
if (value < reference) { reference = value; referenceDepth = depth; }
++evaluations;
}
}
double referenceSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
MergeScheduleResult initial;
std::string error;
TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, std::vector<size_t>(3), {3}, target, initial, error)));
auto query = [&](const MergeScheduleResult &schedule) { return queueCost(graph, schedule, target); };
start = std::chrono::steady_clock::now();
auto search = optimizePipelineSchedule(graph, initial, 0, target, query, error);
TEST_CHECK(mlir::succeeded(search));
auto repeat = optimizePipelineSchedule(graph, initial, 0, target, query, error);
TEST_CHECK(mlir::succeeded(repeat) && search->costs.size() == repeat->costs.size());
size_t winner = selectPipelineCandidate(search->costs, error);
TEST_CHECK(winner == selectPipelineCandidate(repeat->costs, error));
for (size_t i = 0; i < search->schedules.size(); ++i) {
TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
TEST_CHECK(search->costs[i].primary == repeat->costs[i].primary);
}
double best = std::numeric_limits<double>::infinity();
for (const auto &cost : search->costs) best = std::min(best, std::max(*cost.primary, *cost.alternative));
double selected = std::max(*search->costs[winner].primary, *search->costs[winner].alternative);
TEST_CHECK(best + 1e-6 >= reference);
std::cout << "auto_physical_profile=" << profile << " reference=" << reference
<< " reference_depth=" << referenceDepth << " best_evaluated=" << best << " selected=" << selected
<< " selected_depth=" << getPipelineStageCount(search->schedules[winner])
<< " gap_percent=" << 100 * (selected / reference - 1) << " reference_domain=all_physical_assignments_for_chain"
<< " reference_evaluations=" << evaluations << " evaluations=" << search->costs.size()
<< " packing_attempts=" << search->packingAttempts
<< " initialization_seconds=" << search->segmentCostSeconds + search->partitionStateSeconds
<< " reference_seconds=" << referenceSeconds
<< " search_repeat_seconds=" << std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count() << '\n';
}
}
static void testNonunimodalPhysicalDepthCurve() {
TEST_CHECK(std::numeric_limits<double>::is_iec559 && std::numeric_limits<double>::digits == 53);
TEST_CHECK(std::numeric_limits<long double>::digits >= 64);
auto target = hardware(1);
target.processorCount = 4;
target.synchronizationRegisterCount = 32;
target.sharedMemoryRequestLatencyNs.assign(4, 2);
target.sharedMemoryResponseLatencyNs.assign(4, 3);
target.sharedMemoryReadLatencyNs = target.sharedMemoryWriteLatencyNs = 14;
// Manhattan distances at (0,0), (0,11), (18,0), and (5,7).
target.interProcessorLatencyNs = {0, 11, 18, 12, 11, 0, 29, 9, 18, 29, 0, 20, 12, 9, 20, 0};
ComputeGraph graph;
graph.predecessors.resize(5); graph.successors.resize(5);
const Cost work[] = {163, 46, 69, 131, 79};
for (uint32_t task = 0; task < 5; ++task) {
ResidentWeight weight; weight.opaqueLane = task % 2;
graph.nodes.push_back({{nullptr, task, 1}, work[task], {weight}, task});
if (task) {
graph.predecessors[task].push_back({task - 1, {64, 1}});
graph.successors[task - 1].push_back({task, {64, 1}});
}
}
std::vector<double> curve(5, std::numeric_limits<double>::infinity());
auto primary = curve, alternative = curve;
using Ratio = std::pair<uint64_t, uint64_t>;
std::vector<std::optional<Ratio>> certified(5);
auto certifiedPrimary = certified, certifiedAlternative = certified;
auto less = [](Ratio a, Ratio b) { return a.first * b.second < b.first * a.second; };
std::vector<std::vector<size_t>> winnerCores(5), winnerStages(5);
std::vector<size_t> evaluations(5);
for (size_t cuts = 0; cuts < 16; ++cuts) {
std::vector<size_t> stages{0};
for (size_t task = 1; task < 5; ++task)
stages.push_back(stages.back() + ((cuts >> (task - 1)) & 1));
size_t depth = stages.back() + 1;
if (depth > 4) continue;
for (size_t placement = 0; placement < 1024; ++placement) {
MergeScheduleResult schedule;
schedule.processorCount = 4;
size_t encoded = placement;
std::vector<size_t> assigned(4, depth), slots(4), idle;
bool valid = true;
for (size_t task = 0; task < 5; ++task) {
size_t core = encoded % 4; encoded /= 4;
if (assigned[core] != depth && assigned[core] != stages[task]) { valid = false; break; }
assigned[core] = stages[task];
schedule.computeToCpuMap[graph.nodes[task].instance] = core;
schedule.computeToCpuSlotMap[graph.nodes[task].instance] = slots[core]++;
}
if (!valid) continue;
// There are only two resident weights; both fit on every core.
for (size_t core = 0; core < 4; ++core) if (assigned[core] == depth) idle.push_back(core);
size_t maps = 1;
for (size_t i = 0; i < idle.size(); ++i) maps *= depth;
for (size_t map = 0; map < maps; ++map) {
schedule.processorStages = assigned;
size_t digits = map;
for (size_t core : idle) { schedule.processorStages[core] = digits % depth; digits /= depth; }
auto cost = queueCost(graph, schedule, target, true);
Ratio worst{0, 1};
for (auto [policy, proof] : llvm::enumerate(cost.recurrence)) {
TEST_CHECK(proof.proved() && proof.periodGenerations > 0 && proof.periodGenerations <= 1024);
double value = *proof.intervalNs;
TEST_CHECK(value > 0 && value < (uint64_t{1} << 30));
// Queue timestamps are checked integers. Adjacent binary64 values
// enclose the final rounded duration/period division, independently of MCR.
long double low = std::nextafter(value, -std::numeric_limits<double>::infinity());
long double high = std::nextafter(value, std::numeric_limits<double>::infinity());
uint64_t lower = std::ceil(low * proof.periodGenerations);
uint64_t upper = std::floor(high * proof.periodGenerations);
TEST_CHECK(lower == upper && upper < (uint64_t{1} << 40));
uint64_t divisor = std::gcd(lower, proof.periodGenerations);
Ratio exact{lower / divisor, proof.periodGenerations / divisor};
auto &minimum = (policy ? certifiedAlternative : certifiedPrimary)[depth];
if (!minimum || less(exact, *minimum)) minimum = exact;
if (less(worst, exact)) worst = exact;
}
if (!certified[depth] || less(worst, *certified[depth])) certified[depth] = worst;
double value = std::max(*cost.primary, *cost.alternative);
if (value < curve[depth]) {
curve[depth] = value;
winnerStages[depth] = schedule.processorStages;
winnerCores[depth].clear();
for (const auto &node : graph.nodes)
winnerCores[depth].push_back(schedule.computeToCpuMap.lookup(node.instance));
}
primary[depth] = std::min(primary[depth], *cost.primary);
alternative[depth] = std::min(alternative[depth], *cost.alternative);
++evaluations[depth];
}
}
}
TEST_CHECK(curve[1] == 496 && curve[2] == 304 && curve[3] == 305 && curve[4] == 304.5);
TEST_CHECK(primary == curve && alternative == curve);
TEST_CHECK(certified[3] == Ratio(305, 1) && certified[4] == Ratio(609, 2));
TEST_CHECK(certifiedPrimary == certified && certifiedAlternative == certified);
TEST_CHECK(less(*certified[4], *certified[3]));
TEST_CHECK(evaluations == std::vector<size_t>({0, 1024, 1728, 792, 96}));
std::cout << "nonunimodal_physical_depth_reference";
for (size_t depth = 1; depth <= 4; ++depth)
std::cout << " p" << depth << "_ns=" << curve[depth] << " primary=" << primary[depth]
<< " alternative=" << alternative[depth] << " exact=" << certified[depth]->first
<< '/' << certified[depth]->second << " evaluations=" << evaluations[depth];
std::cout << " reference_domain=all_chain_partitions_and_physical_assignments_with_restart\n";
for (size_t depth = 1; depth <= 4; ++depth) {
std::cout << "physical_depth_winner=" << depth << " core_stages=";
for (size_t stage : winnerStages[depth]) std::cout << stage << ',';
std::cout << " task_cores=";
for (size_t core : winnerCores[depth]) std::cout << core << ',';
std::cout << '\n';
}
}
static void testWorseningIntermediate() {
// A controlled cost oracle isolates the search's non-monotone landscape.
// Hardware queue objectives are covered separately by runReference().
auto target = hardware(0);
target.processorCount = 8;
target.interProcessorLatencyNs.assign(64, 0);
ComputeGraph graph;
graph.predecessors.resize(5); graph.successors.resize(5);
for (uint32_t task = 0; task < 5; ++task)
graph.nodes.push_back({{nullptr, task, 1}, task == 4 ? 0u : 10u, {}, task});
auto evaluate = [&](const MergeScheduleResult &schedule) {
bool firstStage = true;
for (size_t i = 0; i < graph.nodes.size(); ++i)
firstStage &= schedule.processorStages[schedule.computeToCpuMap.lookup(graph.nodes[i].instance)]
== (i == 4 ? 1 : 0);
size_t cores = llvm::count(schedule.processorStages, 0);
double value = !firstStage ? 120 : cores == 1 ? 80 : cores >= 6 ? 100 : 110;
WholeScheduleCost cost;
cost.primary = cost.alternative = value;
cost.loweringSupport = StaticLoweringSupport::Supported;
cost.selectionMargin = kPipelineSelectionMargin;
for (auto &proof : cost.recurrence)
proof = {ScheduleRecurrenceEvidence::Proved, 1, value, {}};
return cost;
};
std::string error;
MergeScheduleResult initial, intermediate;
const std::vector<size_t> stages{0, 0, 0, 0, 1};
TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {7, 1}, target, initial, error)));
TEST_CHECK(evaluate(initial).primary == 100);
TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {3, 5}, target, intermediate, error)));
TEST_CHECK(evaluate(intermediate).primary == 110);
double reference = std::numeric_limits<double>::infinity();
for (size_t cores = 1; cores < 8; ++cores) {
MergeScheduleResult schedule;
TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {cores, 8 - cores}, target, schedule, error)));
reference = std::min(reference, *evaluate(schedule).primary);
}
TEST_CHECK(reference == 80);
auto search = optimizePipelineSchedule(graph, initial, 2, target, evaluate, error);
TEST_CHECK(mlir::succeeded(search));
size_t winner = selectPipelineCandidate(search->costs, error);
double selected = *search->costs[winner].primary;
size_t depth = 0;
for (size_t i = winner; i; i = search->parents[i]) ++depth;
TEST_CHECK(depth >= 2);
sawSuccessiveMoves = true;
TEST_CHECK(selected >= reference && selected <= 100);
TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget);
std::cout << "non_monotone_oracle_reference=" << reference << " selected=" << selected
<< " gap_percent=" << 100 * (selected / reference - 1)
<< " evaluations=" << search->costs.size()
<< " packing_attempts=" << search->packingAttempts << '\n';
}
static void testMultipleWorseningAllocations() {
mlir::MLIRContext context;
context.loadDialect<spatial::SpatialDialect>();
auto module = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
mlir::OpBuilder builder(module.getBodyRegion());
auto first = spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
mlir::TypeRange{}, builder.getI32IntegerAttr(4), mlir::ValueRange{}, mlir::ValueRange{});
auto last = spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
mlir::TypeRange{}, builder.getI32IntegerAttr(1), mlir::ValueRange{}, mlir::ValueRange{});
ComputeGraph graph;
graph.predecessors.resize(5); graph.successors.resize(5);
for (uint32_t task = 0; task < 5; ++task) {
graph.nodes.push_back({{task < 4 ? first.getOperation() : last.getOperation(),
task < 4 ? task : 0, 1}, task < 4 ? 10u : 0u, {}, task});
if (task < 4) {
graph.predecessors[4].push_back({task, {0, 0}});
graph.successors[task].push_back({4, {0, 0}});
}
}
auto target = hardware(0);
target.processorCount = 32; target.interProcessorLatencyNs.assign(1024, 0);
auto query = [&](const MergeScheduleResult &schedule) {
size_t cores = llvm::count(schedule.processorStages, 0);
double value = cores == 1 ? 80 : cores == 31 ? 100 : 110 + cores;
WholeScheduleCost cost;
cost.primary = cost.alternative = value;
cost.loweringSupport = StaticLoweringSupport::Supported;
cost.selectionMargin = kPipelineSelectionMargin;
for (auto &proof : cost.recurrence) proof = {ScheduleRecurrenceEvidence::Proved, 1, value, {}};
return cost;
};
std::string error;
MergeScheduleResult initial;
const std::vector<size_t> stages{0, 0, 0, 0, 1};
double reference = std::numeric_limits<double>::infinity();
for (size_t cores = 1; cores < 32; ++cores) {
MergeScheduleResult schedule;
TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {cores, 32 - cores}, target, schedule, error)));
reference = std::min(reference, *query(schedule).primary);
}
TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {31, 1}, target, initial, error)));
auto search = optimizePipelineSchedule(graph, initial, 2, target, query, error);
auto repeat = optimizePipelineSchedule(graph, initial, 2, target, query, error);
TEST_CHECK(mlir::succeeded(search) && mlir::succeeded(repeat));
TEST_CHECK(search->parents == repeat->parents && search->costs.size() == repeat->costs.size());
for (size_t i = 0; i < search->schedules.size(); ++i) {
TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
TEST_CHECK(search->costs[i].primary == repeat->costs[i].primary);
}
size_t winner = selectPipelineCandidate(search->costs, error), path = 0;
TEST_CHECK(reference == 80 && *search->costs[winner].primary == reference);
for (size_t i = winner; i; i = search->parents[i]) {
if (i != winner) TEST_CHECK(*search->costs[i].primary > 100);
++path;
}
TEST_CHECK(path >= 3);
std::cout << "multiple_worsening_reference=" << reference << " selected=" << *search->costs[winner].primary
<< " gap_percent=0 path_length=" << path << " evaluations=" << search->costs.size()
<< " packing_attempts=" << search->packingAttempts
<< " reference_domain=all_31_core_budgets_for_fixed_operation_groups\n";
module.erase();
}
static void runReference(size_t profile, size_t fixture) {
auto target = hardware(profile);
ComputeGraph graph;
const std::vector<std::vector<Cost>> work{{900, 70, 800, 90, 700},
{30, 700, 110, 600, 90}, {1000, 50, 40, 900, 60},
{90, 80, 100, 110, 120}, {100, 800, 100, 800, 100},
{900, 800, 40, 30, 20}};
size_t count = work[fixture].size();
graph.predecessors.resize(count); graph.successors.resize(count);
for (size_t task = 0; task < count; ++task) {
ResidentWeight weight; weight.opaqueLane = fixture == 4 ? task % 3 : task % 2;
graph.nodes.push_back({{nullptr, static_cast<uint32_t>(task), 1}, work[fixture][task],
{weight}, task});
auto edge = [&](size_t source) {
graph.predecessors[task].push_back({source, {64, 1}});
graph.successors[source].push_back({task, {64, 1}});
graph.publicationTransfers.push_back({source, task, {}, 0, {}, 64, 1});
};
if (fixture == 1) {
if (task == 1 || task == 2) edge(0);
if (task == 3) { edge(1); edge(2); }
if (task == 4) edge(3);
} else if (task) edge(task - 1);
if (fixture == 2 && task == 4) edge(1);
}
if (fixture == 3) {
target.sharedMemoryReadLatencyNs = 400;
target.sharedMemoryWriteLatencyNs = 600;
}
std::string error;
MergeScheduleResult initial;
TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, initial, 3, target, error)));
auto evaluate = [&](const MergeScheduleResult &schedule) { return queueCost(graph, schedule, target); };
auto begin = std::chrono::steady_clock::now();
auto search = optimizePipelineSchedule(graph, initial, 3, target, evaluate, error);
double searchSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - begin).count();
TEST_CHECK(mlir::succeeded(search));
auto repeat = optimizePipelineSchedule(graph, initial, 3, target, evaluate, error);
TEST_CHECK(mlir::succeeded(repeat) && search->schedules.size() == repeat->schedules.size());
TEST_CHECK(search->parents == repeat->parents && search->expansions == repeat->expansions);
TEST_CHECK(search->packingAttempts == repeat->packingAttempts
&& search->consideredAssignments == repeat->consideredAssignments);
for (size_t i = 0; i < search->schedules.size(); ++i) {
TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
TEST_CHECK(search->schedules[i].computeToCpuSlotMap == repeat->schedules[i].computeToCpuSlotMap);
TEST_CHECK(search->schedules[i].processorStages == repeat->schedules[i].processorStages);
TEST_CHECK(search->costs.empty() || search->costs[i].primary == repeat->costs[i].primary);
TEST_CHECK(search->costs.empty() || search->costs[i].alternative == repeat->costs[i].alternative);
}
double reference = std::numeric_limits<double>::infinity();
size_t referenceEvaluations = 0;
std::vector<size_t> stages(count), budgets(3, 1);
std::function<void(size_t)> partitions = [&](size_t task) {
if (task != count) {
for (size_t stage = 0; stage <= (task ? 2 : 0); ++stage) {
stages[task] = stage; partitions(task + 1);
}
return;
}
std::vector<bool> used(3);
for (size_t stage : stages) used[stage] = true;
if (!llvm::all_of(used, [](bool value) { return value; })) return;
for (size_t core0 = 1; core0 + 1 < target.processorCount; ++core0)
for (size_t core1 = 1; core0 + core1 < target.processorCount; ++core1) {
budgets = {core0, core1, target.processorCount - core0 - core1};
MergeScheduleResult schedule;
if (mlir::failed(packPipelineSchedule(graph, stages, budgets, target, schedule, error))) continue;
auto cost = evaluate(schedule); ++referenceEvaluations;
reference = std::min(reference, std::max(*cost.primary, *cost.alternative));
}
};
auto referenceBegin = std::chrono::steady_clock::now();
partitions(0);
double referenceSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - referenceBegin).count();
std::string reason;
size_t winner = search->costs.empty() ? 0 : selectPipelineCandidate(search->costs, reason);
auto selected = evaluate(search->schedules[winner]);
double cost = std::max(*selected.primary, *selected.alternative);
double bestEvaluated = cost;
for (const auto &estimate : search->costs)
bestEvaluated = std::min(bestEvaluated, std::max(*estimate.primary, *estimate.alternative));
TEST_CHECK(std::isfinite(reference) && cost + 1e-6 >= reference);
TEST_CHECK(bestEvaluated + 1e-6 >= reference);
TEST_CHECK(cost <= reference * (1 + kPipelineSelectionMargin) + 1e-6);
TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget
&& search->packingAttempts <= 4 * kPipelineQueueEvaluationBudget
&& search->consideredAssignments <= 4 * kPipelineQueueEvaluationBudget);
size_t depth = 0;
for (size_t i = winner; i; i = search->parents[i]) {
++depth;
const auto &current = search->schedules[i], &parent = search->schedules[search->parents[i]];
bool cutChanged = false;
for (const auto &node : graph.nodes)
cutChanged |= current.processorStages[current.computeToCpuMap.lookup(node.instance)]
!= parent.processorStages[parent.computeToCpuMap.lookup(node.instance)];
std::vector<size_t> currentBudgets(3), parentBudgets(3);
for (size_t core = 0; core < target.processorCount; ++core) {
++currentBudgets[current.processorStages[core]];
++parentBudgets[parent.processorStages[core]];
}
sawJointMove |= cutChanged && currentBudgets != parentBudgets;
}
sawSuccessiveMoves |= depth >= 2;
for (size_t i = 1; i < search->schedules.size(); ++i) {
const auto &child = search->schedules[i], &parent = search->schedules[search->parents[i]];
bool cut = false;
for (const auto &node : graph.nodes)
cut |= child.processorStages[child.computeToCpuMap.lookup(node.instance)]
!= parent.processorStages[parent.computeToCpuMap.lookup(node.instance)];
std::vector<size_t> a(3), b(3);
for (size_t stage : child.processorStages) ++a[stage];
for (size_t stage : parent.processorStages) ++b[stage];
sawJointMove |= cut && a != b;
}
std::cout << "profile=" << profile << " fixture=" << fixture
<< " reference_ns=" << reference << " selected_ns=" << cost
<< " gap_percent=" << 100 * (cost / reference - 1)
<< " best_evaluated_ns=" << bestEvaluated
<< " best_evaluated_gap_percent=" << 100 * (bestEvaluated / reference - 1)
<< " reference_scope=logical_with_deterministic_packing"
<< " reference_evaluations=" << referenceEvaluations
<< " search_evaluations=" << search->costs.size()
<< " packing_attempts=" << search->packingAttempts << " refinement_depth=" << depth
<< " search_seconds=" << searchSeconds << " reference_seconds=" << referenceSeconds
<< " total_seconds=" << std::chrono::duration<double>(std::chrono::steady_clock::now() - begin).count()
<< '\n';
auto unavailable = [&](const MergeScheduleResult &) { WholeScheduleCost value; value.reason = "test unavailable timing"; return value; };
auto unknown = optimizePipelineSchedule(graph, initial, 3, target, unavailable, error);
TEST_CHECK(search->costs.empty() || (mlir::failed(unknown) && error.find("test unavailable timing") != std::string::npos));
auto unproved = [&](const MergeScheduleResult &schedule) {
auto cost = evaluate(schedule);
cost.recurrence[0].status = ScheduleRecurrenceEvidence::Unproved;
cost.reason = "test unproved recurrence";
return cost;
};
auto invalid = optimizePipelineSchedule(graph, initial, 3, target, unproved, error);
TEST_CHECK(search->costs.empty() || (mlir::failed(invalid) && error.find("test unproved recurrence") != std::string::npos));
if (!profile && !fixture) {
auto badGuidance = [&](const MergeScheduleResult &schedule) {
auto cost = evaluate(schedule); cost.processorTaskWorkNs.resize(target.processorCount + 1);
return cost;
};
TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, badGuidance, error)));
TEST_CHECK(error.find("physical processor count") != std::string::npos);
TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, evaluate, error, 0)));
TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, evaluate, error, kPipelineQueueEvaluationBudget + 1)));
}
}
// Exhaustive small-grid DP oracle. Physical placement is checked separately.
static void testSharedPartitionStates() {
for (size_t profile = 0; profile < 2; ++profile) {
ComputeGraph graph;
const std::vector<Cost> work{10, 30, 5, 50, 20, 15};
graph.predecessors.resize(work.size()); graph.successors.resize(work.size());
for (uint32_t task = 0; task < work.size(); ++task) {
ResidentWeight weight; weight.opaqueLane = task % (profile + 2);
graph.nodes.push_back({{nullptr, task, 1}, work[task], {weight}, task});
if (task) {
graph.predecessors[task].push_back({task - 1, {}});
graph.successors[task - 1].push_back({task, {}});
}
}
if (profile) {
graph.predecessors[4].push_back({1, {}});
graph.successors[1].push_back({4, {}});
}
auto target = hardware(profile);
MergeScheduleResult peft;
std::string error;
TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, std::vector<size_t>(work.size()),
{target.processorCount}, target, peft, error)));
auto equalCost = [](const MergeScheduleResult &) {
WholeScheduleCost cost;
cost.primary = cost.alternative = 100;
cost.selectionMargin = kPipelineSelectionMargin;
cost.loweringSupport = StaticLoweringSupport::Supported;
for (auto &proof : cost.recurrence) proof = {ScheduleRecurrenceEvidence::Proved, 1, 100, {}};
return cost;
};
auto search = optimizePipelineSchedule(graph, peft, 0, target, equalCost, error);
if (mlir::failed(search)) std::cerr << error << '\n';
TEST_CHECK(mlir::succeeded(search));
for (size_t i = 1; i < search->initialSchedules; ++i) {
size_t depth = getPipelineStageCount(search->schedules[i]);
MergeScheduleResult independentlyInitialized;
TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, independentlyInitialized, depth, target, error)));
TEST_CHECK(search->schedules[i].processorStages == independentlyInitialized.processorStages);
TEST_CHECK(search->schedules[i].computeToCpuMap == independentlyInitialized.computeToCpuMap);
TEST_CHECK(search->schedules[i].computeToCpuSlotMap == independentlyInitialized.computeToCpuSlotMap);
}
IndexedResidentWeights weights(graph);
for (const auto &entry : search->depthPriorities) {
size_t depth = entry.first;
Cost actual = entry.second;
Cost reference = std::numeric_limits<Cost>::max();
size_t evaluated = 0;
for (size_t mask = 0; mask < (size_t{1} << (work.size() - 1)); ++mask) {
std::vector<size_t> stages(work.size());
for (size_t task = 1; task < work.size(); ++task)
stages[task] = stages[task - 1] + ((mask >> (task - 1)) & 1);
if (stages.back() + 1 != depth) continue;
bool legal = true;
for (size_t task = 0; task < work.size(); ++task)
for (auto [before, transfer] : graph.predecessors[task])
legal &= stages[task] <= stages[before] + 1;
if (!legal) continue;
std::vector<std::vector<size_t>> tasks(depth);
for (size_t task = 0; task < work.size(); ++task) tasks[stages[task]].push_back(task);
std::function<void(size_t, size_t, Cost)> allocate = [&](size_t stage, size_t remaining, Cost objective) {
if (stage == depth) {
if (!remaining) { reference = std::min(reference, objective); ++evaluated; }
return;
}
const auto &members = tasks[stage];
std::vector<Cost> ends(work.size());
Cost span = 0, largest = 0;
for (size_t task : members) {
ends[task] = checkedAdd(ends[task], work[task]);
span = std::max(span, ends[task]); largest = std::max(largest, work[task]);
for (auto [next, transfer] : graph.successors[task]) ends[next] = std::max(ends[next], ends[task]);
}
Cost packed = std::numeric_limits<Cost>::max();
for (size_t cores = 1; cores + depth - stage - 1 <= remaining; ++cores) {
if (auto cost = weights.estimatePackCost(members, work, target.residentWeightCapacity, cores))
packed = std::min(packed, *cost);
if (packed != std::numeric_limits<Cost>::max())
allocate(stage + 1, remaining - cores, std::max(objective, checkedAdd(span - largest, packed)));
}
};
allocate(0, target.processorCount, 0);
}
TEST_CHECK(reference == actual);
std::cout << "shared_dp_profile=" << profile << " depth=" << depth
<< " reference_priority=" << reference << " actual_priority=" << actual
<< " reference_allocations=" << evaluated << " dp_transitions=" << search->partitionTransitions << '\n';
}
}
}
static void testDepthDomain() {
ComputeGraph graph;
graph.predecessors.resize(12); graph.successors.resize(12);
for (uint32_t i = 0; i < 12; ++i) {
graph.nodes.push_back({{nullptr, i, 1}, 100, {}, i});
if (i) {
graph.predecessors[i].push_back({i - 1, {}});
graph.successors[i - 1].push_back({i, {}});
}
}
auto target = hardware(0);
target.processorCount = 16;
target.interProcessorLatencyNs.assign(256, 1);
std::string error;
for (size_t depth : {2, 4, 8}) {
MergeScheduleResult schedule;
TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, schedule, depth, target, error)));
TEST_CHECK(mlir::succeeded(verifyPipelineStageCount(schedule, depth, error)));
std::set<size_t> active;
for (const auto &node : graph.nodes)
active.insert(schedule.processorStages[schedule.computeToCpuMap.lookup(node.instance)]);
TEST_CHECK(active.size() == depth && *active.begin() == 0 && *active.rbegin() == depth - 1);
auto broken = schedule;
for (const auto &[instance, core] : broken.computeToCpuMap)
if (broken.processorStages[core] == depth - 1) broken.processorStages[core] = depth - 2;
TEST_CHECK(mlir::failed(verifyPipelineStageCount(broken, depth, error)));
}
MergeScheduleResult invalid;
TEST_CHECK(mlir::failed(runPipelineScheduler(graph, invalid, 13, target, error)));
TEST_CHECK(error.find("compute operation groups") != std::string::npos);
TEST_CHECK(mlir::failed(runPipelineScheduler(graph, invalid, 17, target, error)));
MergeScheduleResult peft;
peft.processorCount = target.processorCount;
for (size_t i = 0; i < graph.nodes.size(); ++i) {
peft.computeToCpuMap[graph.nodes[i].instance] = 0;
peft.computeToCpuSlotMap[graph.nodes[i].instance] = i;
}
auto incomplete = peft;
incomplete.processorStages = {0};
TEST_CHECK(mlir::failed(verifyPipelineStageCount(incomplete, 1, error)));
std::vector<size_t> winners;
for (double boundaryCost : {1000., 300., 50., 0.}) {
// A closed synthetic target contract isolates depth-domain selection.
// Complete queue objectives and physical placements are tested separately.
auto query = [&](const MergeScheduleResult &schedule) {
size_t depth = getPipelineStageCount(schedule);
double interval = 1200. / depth + boundaryCost * (depth - 1);
WholeScheduleCost cost;
cost.primary = cost.alternative = interval;
cost.loweringSupport = StaticLoweringSupport::Supported;
cost.selectionMargin = kPipelineSelectionMargin;
for (auto &proof : cost.recurrence)
proof = {ScheduleRecurrenceEvidence::Proved, 1, interval, {}};
return cost;
};
auto search = optimizePipelineSchedule(graph, peft, 0, target, query, error);
TEST_CHECK(mlir::succeeded(search));
auto repeat = optimizePipelineSchedule(graph, peft, 0, target, query, error);
TEST_CHECK(mlir::succeeded(repeat) && repeat->schedules.size() == search->schedules.size());
TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget);
size_t winner = selectPipelineCandidate(search->costs, error);
size_t second = selectPipelineCandidate(repeat->costs, error);
TEST_CHECK(winner == second);
for (size_t i = 0; i < search->schedules.size(); ++i) {
TEST_CHECK(search->schedules[i].processorStages == repeat->schedules[i].processorStages);
TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
TEST_CHECK(search->costs[i].primary == repeat->costs[i].primary);
TEST_CHECK(mlir::succeeded(verifyPipelineStageCount(search->schedules[i],
getPipelineStageCount(search->schedules[i]), error)));
}
winners.push_back(getPipelineStageCount(search->schedules[winner]));
double reference = std::numeric_limits<double>::infinity();
for (size_t depth = 1; depth <= 12; ++depth)
reference = std::min(reference, 1200. / depth + boundaryCost * (depth - 1));
std::cout << "auto_target_boundary_ns=" << boundaryCost << " depth=" << winners.back()
<< " reference_ns=" << reference << " selected_ns=" << *search->costs[winner].primary
<< " evaluations=" << search->costs.size() << '\n';
}
TEST_CHECK(winners == std::vector<size_t>({1, 2, 5, 12}));
graph.predecessors[11].push_back({0, {}});
graph.successors[0].push_back({11, {}});
TEST_CHECK(mlir::failed(runPipelineScheduler(graph, invalid, 8, target, error)));
auto unavailable = [](const MergeScheduleResult &) {
WholeScheduleCost cost; cost.reason = "unavailable auto test cost"; return cost;
};
TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, peft, 0, target, unavailable, error)));
TEST_CHECK(error.find("unavailable auto test cost") != std::string::npos);
}
int main() {
testPackingCostCacheEviction();
testPackingCurves();
testSharedPartitionStates();
testDepthDomain();
testInitializerPackingWitness();
testPackingFeasibility();
for (uint32_t first : {19u, 21u}) testQuantizedFeasibility(first);
testPhysicalReference();
testAutoPhysicalReference();
testNonunimodalPhysicalDepthCurve();
testWorseningIntermediate();
testMultipleWorseningAllocations();
for (size_t profile = 0; profile < 2; ++profile)
for (size_t fixture = 0; fixture < 6; ++fixture) runReference(profile, fixture);
TEST_CHECK(sawJointMove && sawSuccessiveMoves);
}