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

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
@@ -8,6 +8,19 @@ synchronization.
## Invariant
`--pipeline=1` preserves latency-oriented PEFT and remains the default.
An explicit `--pipeline=N`, with `N > 1`, requires exactly `N` nonempty
compute stages, identified by all integers from zero through `N - 1`.
Each stage owns at least one physical core and scheduled compute work.
Initialization, feasibility completion, refinement, allocation, placement,
and final verification must preserve this count. Fixed-depth refinement
must not remove stages or compact the schedule to a smaller depth.
An unsupported exact depth fails explicitly; empty stages are not a repair.
`--pipeline=auto` permits depth changes within the legal resource domain.
It compares PEFT p1 and pipeline schedules using steady-state queue costs.
Auto p1 preserves PEFT placement and uses the existing stage-zero restart
barrier for repeated inputs. Numeric p1 remains the latency program.
A scheduled compute operation and all of its lanes belong to exactly one
pipeline stage. An operation may consume results produced in its own stage or
the immediately preceding stage only. Therefore every compute-graph edge from
+1 -1
View File
@@ -11,7 +11,7 @@ Boolean switches default to off unless stated otherwise.
| `--maccel=PIM` | Select the PIM accelerator. |
| `--EmitSpatial`, `--EmitPim`, `--EmitPimBufferized`, `--EmitPimCodegen` | Stop at a pipeline stage; default: codegen. |
| `--core-count=N`, `--crossbar-size=N`, `--crossbar-count=N` | Required hardware resources; crossbars are square. |
| `--pipeline=N` | Static queue-aware throughput selection for final codegen with N > 1; intermediate targets retain the original schedule. Default: 1 (latency PEFT). |
| `--pipeline=N\|auto` | Default: 1 (latency PEFT). N > 1 requires exactly N nonempty compute stages. `auto` compares stage counts with static queue-aware throughput costs. Auto requires final codegen; numeric intermediate targets use the initialized schedule. |
| `--pim-target-config=PATH` | Target cost/topology configuration; resources must match hardware flags. Default: built-in model. |
| `--pim-memory-report=summary\|none` | Memory report at `reports/memory_report.txt`; default: `summary`. |
| `--pim-only-codegen` | Run the codegen tail on bufferized Pim IR. |
+22 -1
View File
@@ -138,7 +138,10 @@ remain approximations, not extra inferred memory traffic.
adjustment between physical execution cost and nominal volume priority, scaling
it by chunk lanes. Initialization uses nominal priorities; transposes count as
bulk moves there. One joint prefix/resource DP constructs the initial partition.
Bounded beam refinement changes partitions and their core allocations together.
Bounded incumbent/frontier refinement changes partitions and core allocations together.
Numeric depths require exactly that many nonempty compute stages throughout both DPs,
packing, refinement, and realization. Auto uses the same optimizer across a bounded
stage-count domain; all depths share the queue-evaluation budget.
Queue-aware Spatial evaluation scores every admitted physical placement. PEFT,
final schedule evaluation, and reported per-core demand use physical estimates.
These quantities must not be substituted for measured simulator periods.
@@ -255,3 +258,21 @@ Storage costs `O(P+R+G D)` for bounded generation count `G`, plus bounded recurr
Reports separate extraction, expansion, and queue-analysis times.
Fingerprints compare memory and blocking-operation order, without addresses or work between requests.
Equal fingerprints therefore do not prove equal request readiness or timing.
### Shared automatic-depth initialization
`PipelinePartitionTable` owns the ordered segment costs and joint prefix/resource
states. Automatic depth shares integer-core segment curves across its depth
initializations. Each depth samples its original bounded resource grid.
Its reserved-core count changes that grid, so identical DP states are not available across depths.
Both modes use the same recurrence, legality checks, and constructive packing
implementation.
Ordered cost curves share threshold scans. Scalar queries retain their original
cost and placement decisions. Compact plans retain the constructive core count
before spare cores are assigned. Minimum-resource completion uses the same
shared scalar-equivalent curves.
Automatic-depth initialization reports separate grouping, segment packing,
DP-state evaluation, and physical placement. Fixed-depth initialization remains
inside the candidate-generation timing. The complete queue model determines
eligibility and final selection. Initialization priorities are search costs, not timing bounds.
+117 -27
View File
@@ -1,7 +1,7 @@
# Spatial scheduling
`prepareSpatialSchedules` constructs the shared `ComputeGraph`. Latency mode
uses PEFT. For `--pipeline > 1`, the PIM compiler selects among bounded pipeline
uses PEFT. For `--pipeline > 1` or `--pipeline=auto`, the PIM compiler selects among bounded pipeline
candidates with static queue-aware Spatial execution analysis. Pipeline mode does not run PEFT.
See [implementation notes](scheduling-implementation.md) for fragment-view eligibility,
@@ -68,10 +68,12 @@ every physical core count for each segment. This preserves feasibility over
the segment assignments found by the constructive packers, independently of
resource quantization. Packing failure remains a bounded-search limitation,
not proved hardware infeasibility.
The seed minimizes maximum stage packing/dependency priority over the legal
active-stage counts within the configured depth. These priorities provide
a starting point, not a timing guarantee. Refinement can remove stage boundaries.
No empty stage is inserted between dependent operations.
For numeric depths greater than one, initialization requires exactly that many
nonempty compute stages. Both prefix DPs enforce this requirement. Insufficient
groups, resources, or legal stage cuts produce a diagnostic. A configured stage
array does not establish active compute stages; realization verifies actual
compute assignments and physical-core ownership. Packing priorities provide
a starting point, not a timing guarantee.
Every proposed partition receives a minimax core budget for its changed stages.
Unchanged stages retain their budgets. The allocator builds
@@ -87,14 +89,26 @@ retained assignment.
Task reordering cannot discard the initial partition's constructive capacity evidence.
This evidence does not prove final backend legality.
Refinement uses a beam of two schedules ranked by the larger queue estimate.
Its ordered neighborhood moves operation groups across adjacent stage boundaries,
inserts a boundary, or removes consecutive boundaries.
Refinement preserves contiguous groups in the initialization's topological order.
Coalescing compacts later stages. Each changed block redistributes its core pool.
Refinement retains the best evaluated schedule separately from expansion state.
Its ordered neighborhood moves operation groups across adjacent stage boundaries.
Only automatic depth search can insert or remove boundaries. Fixed-depth moves
preserve every required compute stage. Refinement preserves contiguous groups
in the initialization's topological order. Each changed block redistributes its
core pool. A depth change recomputes the allocation for its complete partition.
Core coordinates also transfer powers of two between stages, including one core.
These moves cross packing plateaus within a bounded number of evaluations.
The configured pipeline depth and total physical resources remain fixed.
Numeric pipeline depth and total physical resources remain fixed.
Automatic search changes depth within the same physical resource budget.
Exact-depth feasible schedule sets are not nested. Discrete placement, packing,
boundary transfers and restart synchronization can make the optimum output
interval non-unimodal in depth. A bounded search adds optimization error at each
depth. A poor initialized or refined schedule therefore cannot exclude nearby
depths. Numeric-depth results and automatic-search results do not establish a
global optimum over depths or placements.
Complete-state queue intervals use integer elapsed times and generation counts,
with a final floating-point division. Their precision is independent of the
one-nanosecond MCR diagnostic bracket.
A cheap proxy ranks the neighborhood before physical packing and queue analysis.
It uses ordered-packing curves, dependency spans and aggregate shared-memory demand.
@@ -104,15 +118,22 @@ Processor service and critical recurrence contributions guide ties toward affect
The PIM adapter supplies target-neutral processor guidance from the shared periodic
graph. PIM timing semantics remain outside Spatial.
The optimizer evaluates at most thirty-two complete Spatial queue costs and attempts
at most 128 constructive allocations and physical packings.
The optimizer evaluates at most sixty-four complete Spatial queue costs and attempts
at most 256 constructive allocations and physical packings.
The shortlist contains at most 128 proposals. Initialization retains at most sixteen
spare-core units. The packing-cost cache retains at most 256 ordered task lists.
Each round evaluates at most two new placements.
Each retained parent receives one queue-evaluation opportunity per round.
The first round can evaluate two neighbors of its single parent.
The two best evaluated schedules form the next beam.
The expansion frontier contains the best unexhausted schedule and one best-first
exploration schedule. Exploration uses waves of `ceil(sqrt(B))` expansions for
queue budget `B`, then ages that schedule behind less-expanded schedules. This
allows bounded lookahead through worse intermediate states while retaining the
best incumbent. At budget 64, the wave has eight expansions.
The two frontier slots do not change final selection.
All evaluated schedules remain available for final comparison.
Search stops at the evaluation or
packing budget, or when the beam's legal unvisited neighborhood is exhausted.
packing budget, or when all retained states have exhausted their legal neighborhood.
Canonical stage/budget vectors, physical core order and evaluation index break
ties deterministically. Actual physical placements, including inactive-core
stage identities, are deduplicated. The final 2.1% near-equivalence policy compares
@@ -120,8 +141,8 @@ the retained queue estimates. There are no independent legacy, sharing-aware,
allocator-only or boundary-shift candidate-generation policies.
Unavailable timing, unsupported static forms or unproved recurrence explicitly
stop a costed search. They are not silently discarded. One operation group or
one distinct neighborhood schedule skips optional queue comparison and lowers
stop a costed search. They are not silently discarded. One fixed-depth operation group or
one distinct fixed-depth neighborhood schedule skips optional queue comparison and lowers
once. Final backend resource retries remain restricted to already costed schedules.
`ComputeGraph` retains physical publication and consumer identities separately
@@ -175,8 +196,10 @@ The analysis brackets the maximum cycle ratio within one nanosecond. It tests
edge delays `w - II*d`, sweeping the distance-zero DAG between recurrence
relaxations. A positive predecessor cycle can terminate an infeasible probe early.
The final witness pass retains the full relaxation procedure and deterministic
witness selection. At most `R+1` sweeps detect a positive cycle, where `R` counts
positive-distance edges. Bisection uses at most 64 tests. Timing outside the
witness selection. Contiguous adjacency offsets preserve the original edge order.
Probes reuse one flattened DAG edge order. Recurrence arithmetic and witness
tie-breaking remain unchanged. At most `R+1` sweeps detect a positive cycle,
where `R` counts positive-distance edges. Bisection uses at most 64 tests. Timing outside the
supported numeric range produces an explicit unavailable result. A predecessor
cycle supplies a witness, including delay, iteration distance, cores, and links.
Detailed witness output is limited to 128 edges; totals cover the entire cycle.
@@ -311,10 +334,11 @@ The standalone frontend retains the recurrence witness. Neither analysis changes
## Default throughput selection
Full PIM code generation uses Spatial-only queue-aware selection for `--pipeline > 1`.
Full PIM code generation uses Spatial-only queue-aware selection for `--pipeline > 1` and `--pipeline=auto`.
Latency PEFT and `--pim-only-codegen` retain their existing behavior.
Intermediate emission retains the initialized schedule and stops at its requested stage.
Its schedule can differ from the final program.
Its schedule can differ from the final program. Automatic depth requires full
code generation and is rejected for intermediate emission.
The compiler performs these steps:
@@ -324,13 +348,14 @@ The compiler performs these steps:
4. Realize each admitted schedule in an independent clone of the prepared checkpoint.
5. Query the target adapter over realized Spatial IR, then release the clone.
6. Require supported static forms, available costs, and proved modeled recurrence.
7. Refine a two-schedule beam within thirty-two total queue evaluations.
7. Refine the incumbent and exploration frontier within sixty-four total queue evaluations.
8. Apply the existing 2.1% near-equivalence rule to the evaluated schedules.
9. Realize the selected schedule again from the unchanged checkpoint.
10. Lower, allocate, verify, and emit only the selected program.
A single candidate skips event extraction, queue analysis, recurrence proof, and comparison.
A single fixed-depth candidate skips event extraction, queue analysis, recurrence proof, and comparison.
It is realized and lowered once, even when optional timing is unavailable.
Auto still validates its p1 queue cost when no pipeline candidate can be initialized.
No rejected candidate enters PIM IR or instruction generation during cost evaluation.
Constants and weights remain shared immutable attributes in the MLIR context.
Candidate operations, communication plans, and synchronization state remain independent.
@@ -386,6 +411,44 @@ Final instruction generation uses the normal compiler timing report.
Explicit Spatial exports and IR dumps reconstruct pre-lowering artifacts for the selected schedule.
They do not change the emitted program. Diagnostic MLIR dumps elide large constants and are not lossless checkpoints.
### Fixed and automatic depth
`--pipeline=1` retains latency PEFT, including its existing instruction program.
`--pipeline=N`, for N greater than one, requires exactly N nonempty compute stages.
Every stage has physical cores. Boundary refinement cannot remove a required stage.
The compiler rejects infeasible exact depths instead of compacting stages.
`--pipeline=auto` compares steady-state output intervals with the same queue model
and final selection policy. Its p1 candidate uses the original PEFT schedule.
Auto adds the existing stage-zero restart barrier to the PEFT placement. This
protects changing input generations without changing numeric p1 latency compilation.
The queue model repeats that synchronized program; its cost is not single-inference latency.
If PEFT initialization fails, the report records that construction failure and
starts from a constructively feasible pipeline schedule. This is not proof that
all possible p1 placements are infeasible.
The legal ordered-prefix graph bounds the maximum depth, together with the
physical-core count. Auto shares segment packing curves and graph constraints
across depth initializations. Each depth retains its resource grid and
deterministic DP tie order. Reserved-core counts change those grids, so their
DP states cannot be reused unchanged across depths. Minimum-resource completion
reuses the shared integer-core cost curves.
Auto initialization covers depth 2 and the maximum depth, then bisects the
largest remaining depth gap. Ties prefer the shallower gap. At most
`min(1+ceil(log2(D-1)), D-1, floor(B/2))` pipeline seeds are constructed for
maximum depth D and queue budget B. Failed depth initializations are reported;
they do not invalidate feasible depths. All seeds and refinements share
sixty-four complete queue evaluations and the same incumbent/frontier policy.
Untested depths and unexplored neighbors remain explicit search limitations.
This bounded search does not prove a global optimum.
An unavailable cost or unproved recurrence for an admitted schedule stops comparison.
Only the final selected program enters PIM lowering.
`pipeline_selection.json` records requested and selected depths, explored depths,
per-depth initialization status, active stage identities, queue costs, search parents,
expansion counts, budgets, and the termination reason.
### Static analysis, diagnostics, and external validation
Production decisions use the input graph, compiler options, physical plans, and configured target parameters only.
@@ -398,18 +461,45 @@ Optional diagnostics expose resource work, recurrence, and ordered-request finge
Offline Pimsim fast-mode runs validate predictions and emitted programs.
Simulator measurements never enter production selection.
For `C <= 32` candidates, scheduling metadata uses `O(C(T+P))` storage for `T` compute instances and `P` cores.
Placement deduplication costs `O(C²(T+P))` time. The beam retains at most two schedules.
For `C <= 64` candidates, scheduling metadata uses `O(C(T+P))` storage for `T` compute instances and `P` cores.
Placement identities require `O(C(T+P))` storage; ordered-set deduplication compares
identities in `O((T+P) log C)` time per insertion. The active frontier has at most two schedules.
Auto adds `O(G²)` time and `O(G)` space for its legal depth upper bound.
For G operation groups, S numeric stages, P cores, and U spare-core units
(at most sixteen), the numeric quality DP has `O(S G² U²)` transitions and
`O(G² U + S G U)` storage. Minimum-resource completion can require `O(G² P)`
packing queries and `O(S G²)` transitions.
Auto shares segment curves across its depth initializations. Ordered minimax
queries reuse threshold scans and monotone neighboring-core bounds. Curves query only the
union of core counts used by the selected depth grids. Feasibility completion
computes missing integer-core costs lazily, with the same constructive packer.
Constructive unordered packing remains necessary when ordered packing fails. Illegal stage segments
have no packing queries. For maximum depth D, all depth-specific resource grids
have `O(L D G² U²)` DP transitions in total for L initialized depths.
Segment costs require at most `O(G² P Kpack)` constructive query work, with
reused ordered-threshold scans reducing repeated scalar searches.
Peak storage is `O(G² P + D G U + L T)` for T compute instances, plus packing
scratch. Uninitialized depths do not allocate DP states.
Each initialized depth retains a compact assignment and its constructive core
counts. The shared table is released before queue evaluation.
Depth seed reconstruction does not rerun packing-cost construction or the DP.
Profiles report segment work, DP transitions, threshold scans, and physical
initialization separately.
A round streams legal boundary and core moves into a shortlist of at most `2C`.
Each parent builds stage resource and dependency summaries in `O(T+E+P)` time and storage.
Unchanged-stage packing costs use an `O(SP)` cache. Changed task lists are not cached.
Exact packing costs and resource lower bounds share a search-local cache of at most
256 ordered task lists, independent of the queue budget. The immutable task costs,
weights and target fix each query. LRU eviction changes recomputation only.
For maximum list length T, the cache uses `O(256(T+P))` storage and
`O(T log 256)` lookup time; eviction scans at most 256 entries.
Each parent proposes `O(G + S² log P)` group, cut, and core moves.
For `M` proposed assignments, cheap allocation and scoring cost
`O(M (T log T + S P + E + F + S Kpack))`, plus bounded shortlist comparisons,
where `Kpack` is one existing ordered-packing query's cost.
Constructive allocation builds at most `S P` ordered-packing queries per admitted
partition. The graph and weight indices are reused. Search performs at most
128 physical packings; queue costs remain capped at thirty-two.
256 physical packings; queue costs remain capped at sixty-four.
Recurrence guidance uses the shared bounded maximum-cycle-ratio solver.
Candidate realization and extraction are sequential; complete candidate modules are not retained.
Normal successful compilation performs one candidate-specific backend lowering.
+16 -7
View File
@@ -132,12 +132,21 @@ llvm::cl::opt<size_t>
llvm::cl::desc("Number of crossbars in each core (required for Pim compilation)"),
llvm::cl::init(0));
llvm::cl::opt<size_t> pipelineStages(
llvm::cl::opt<std::string> pipelineStages(
"pipeline",
llvm::cl::desc("Number of throughput pipeline stages (1 preserves latency scheduling)"),
llvm::cl::init(1),
llvm::cl::desc("Exactly N nonempty stages, or auto for throughput depth selection (default: 1, latency PEFT)"),
llvm::cl::init("1"),
llvm::cl::cat(OnnxMlirOptions));
size_t getPimPipelineStages() {
llvm::StringRef value(pipelineStages.getValue());
if (value == "auto") return 0;
size_t stages;
if (value.getAsInteger(10, stages) || !stages)
llvm::report_fatal_error("Pim compilation requires --pipeline=<positive integer|auto>");
return stages;
}
llvm::cl::opt<long> coresCount("core-count",
llvm::cl::desc("Number of cores in the chip. Required for Pim compilation."),
llvm::cl::init(-1));
@@ -161,12 +170,12 @@ void verifyPimCompilerOptions() {
llvm::report_fatal_error("Pim compilation requires an explicit --crossbar-count=<positive integer>");
if (crossbarCountInCore.getValue() == 0)
llvm::report_fatal_error("Pim compilation requires --crossbar-count to be a positive integer");
if (pipelineStages.getValue() == 0)
llvm::report_fatal_error("Pim compilation requires --pipeline to be positive");
if (static_cast<size_t>(coresCount.getValue()) < pipelineStages.getValue())
size_t stages = getPimPipelineStages();
if (static_cast<size_t>(coresCount.getValue()) < stages)
llvm::report_fatal_error("Pim compilation requires --pipeline not to exceed --core-count");
size_t maximumStages = stages ? stages : static_cast<size_t>(coresCount.getValue());
if (crossbarCountInCore.getValue()
> std::numeric_limits<size_t>::max() / pipelineStages.getValue())
> std::numeric_limits<size_t>::max() / maximumStages)
llvm::report_fatal_error("Pim compilation --crossbar-count * --pipeline overflows");
}
+3 -1
View File
@@ -65,7 +65,9 @@ extern llvm::cl::opt<bool> pimDisableSpatialPlanning;
extern llvm::cl::opt<size_t> crossbarSize;
extern llvm::cl::opt<size_t> crossbarCountInCore;
extern llvm::cl::opt<size_t> pipelineStages;
extern llvm::cl::opt<std::string> pipelineStages;
// Zero denotes automatic depth internally; numeric zero is never accepted.
size_t getPimPipelineStages();
extern llvm::cl::opt<long> coresCount;
extern llvm::cl::opt<std::string> pimTargetConfig;
extern llvm::cl::opt<uint64_t> pimConvIm2colMaxElements;
+8 -4
View File
@@ -463,16 +463,20 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
pm.addPass(createLowerSpatialPlansPass(targetResources, planningOptions, exportStage));
pm.addPass(createTrivialGraphComputeMergePass(
schedulingTarget.residentWeightCapacity, exportStage));
if (pipelineStages.getValue() > 1 && pimEmissionTarget >= EmitPimCodegen) {
size_t stages = getPimPipelineStages();
if (stages != 1 && pimEmissionTarget >= EmitPimCodegen) {
pm.addPass(createPimPipelineSelectionPass(schedulingTarget, targetResources,
pipelineStages.getValue(), exportStage));
stages, exportStage));
pm.addPass(createEmitPimCodePass());
return;
} else {
if (pipelineStages.getValue() > 1)
if (stages == 0) {
llvm::report_fatal_error("--pipeline=auto requires full PIM code generation for queue-cost comparison");
}
if (stages > 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()));
schedulingTarget, exportStage, stages, pimPeriodicDiagnostics.getValue()));
}
pm.addPass(createMessagePass("Onnx lowered to Spatial"));
}
+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);
@@ -34,7 +34,7 @@ static LogicalResult collectScheduledOperations(
size_t pipelineStageCount,
size_t processorCount) {
if (pipelineStageCount == 0
|| (pipelineStageCount > 1
|| ((pipelineStageCount > 1 || !materialization.processorStages.empty())
&& materialization.processorStages.size() != processorCount))
return failure();
plan.processorStages = materialization.processorStages;
@@ -59,7 +59,7 @@ static LogicalResult collectScheduledOperations(
if (core >= processorCount)
return op.emitOpError("phase 2 scheduled core is outside the target");
info.cores.push_back(core);
if (pipelineStageCount > 1) {
if (!materialization.processorStages.empty()) {
size_t stage = materialization.processorStages[core];
if (stage >= pipelineStageCount)
return op.emitOpError("phase 2 scheduled core has an invalid pipeline stage");
@@ -84,8 +84,8 @@ private:
FailureOr<PreparedSpatialSchedules> prepareSpatialSchedules(ModuleOp moduleOp,
const SchedulingTarget &target, size_t pipelineStages) {
if (!hasValidTarget(target) || !PipelineCoreLayout(target.processorCount, pipelineStages).isValid()
|| (pipelineStages > 1 && !target.synchronizationRegisterCount))
if (!hasValidTarget(target) || (pipelineStages && !PipelineCoreLayout(target.processorCount, pipelineStages).isValid())
|| (pipelineStages != 1 && !target.synchronizationRegisterCount))
return moduleOp.emitError("Spatial scheduling requires an explicit valid target and pipeline layout"), failure();
auto entry = requireEntry(moduleOp);
if (failed(entry)) return failure();
@@ -100,10 +100,21 @@ FailureOr<PreparedSpatialSchedules> prepareSpatialSchedules(ModuleOp moduleOp,
}
scheduledGraph = std::move(*graph);
std::string pipelineError;
if (pipelineStages > 1) {
if (pipelineStages == 0) {
MergeSchedulingAnalysis latency(scheduledGraph, target, entryFunc.getContext(), pipelineError);
if (succeeded(latency.getResult())) {
schedule = std::move(*latency.getResult());
// Auto repeats the PEFT placement; explicit stage zero enables restart safety.
schedule.processorStages.assign(target.processorCount, 0);
break;
}
prepared.latencyInitializationReason = pipelineError;
pipelineError.clear();
}
if (pipelineStages != 1) {
FailureOr<PipelineWorkloadPreparation> preparation =
preparePipelineWorkload(
scheduledGraph, pipelineStages, target, pipelineError);
scheduledGraph, pipelineStages ? pipelineStages : target.processorCount, target, pipelineError);
if (failed(preparation)) {
moduleOp.emitError() << pipelineError;
return failure();
@@ -136,6 +147,9 @@ LogicalResult realizeSpatialSchedule(ModuleOp moduleOp,
const MergeScheduleResult &schedule, const SchedulingTarget &target,
size_t pipelineStages, SpatialDataflowExportStage exportStage,
const ComputeGraph *costGraph, std::string *scheduleReport) {
std::string depthError;
if (failed(verifyPipelineStageCount(schedule, pipelineStages, depthError)))
return moduleOp.emitError(depthError), failure();
auto entry = requireEntry(moduleOp);
if (failed(entry)) return failure();
auto entryFunc = *entry;
@@ -9,6 +9,7 @@ namespace onnx_mlir::spatial {
struct PreparedSpatialSchedules {
ComputeGraph graph;
std::vector<MergeScheduleResult> candidates;
std::string latencyInitializationReason;
};
// Success always includes at least one candidate, including an empty compute graph.
@@ -1,5 +1,7 @@
#include "PeriodicDependencies.hpp"
#include "llvm/ADT/STLExtras.h"
#include <limits>
namespace onnx_mlir::spatial {
@@ -7,7 +9,7 @@ namespace onnx_mlir::spatial {
mlir::FailureOr<PeriodicDependencyEstimate> analyzePeriodicDependencies(
size_t nodeCount, llvm::ArrayRef<PeriodicDependency> edges,
std::string &error) {
std::vector<std::vector<size_t>> outgoing(nodeCount);
std::vector<size_t> offsets(nodeCount + 1);
std::vector<size_t> indegree(nodeCount), recurrent;
Cost upper = 0;
for (auto [index, edge] : llvm::enumerate(edges)) {
@@ -23,16 +25,26 @@ mlir::FailureOr<PeriodicDependencyEstimate> analyzePeriodicDependencies(
if (edge.distance) {
recurrent.push_back(index);
} else {
outgoing[edge.source].push_back(index);
++offsets[edge.source + 1];
++indegree[edge.target];
}
}
for (size_t node = 1; node <= nodeCount; ++node) offsets[node] += offsets[node - 1];
std::vector<size_t> outgoing(offsets.back());
{
auto next = offsets;
for (auto [index, edge] : llvm::enumerate(edges))
if (!edge.distance) outgoing[next[edge.source]++] = index;
}
auto outgoingEdges = [&](size_t node) {
return llvm::ArrayRef(outgoing).slice(offsets[node], offsets[node + 1] - offsets[node]);
};
std::vector<size_t> order;
for (size_t node = 0; node < nodeCount; ++node)
if (!indegree[node])
order.push_back(node);
for (size_t cursor = 0; cursor < order.size(); ++cursor)
for (size_t index : outgoing[order[cursor]])
for (size_t index : outgoingEdges(order[cursor]))
if (!--indegree[edges[index].target])
order.push_back(edges[index].target);
if (order.size() != nodeCount) {
@@ -45,7 +57,7 @@ mlir::FailureOr<PeriodicDependencyEstimate> analyzePeriodicDependencies(
for (size_t node : order) {
result.withinGenerationSpanNs =
std::max(result.withinGenerationSpanNs, span[node]);
for (size_t index : outgoing[node]) {
for (size_t index : outgoingEdges(node)) {
const auto &edge = edges[index];
span[edge.target] = std::max(span[edge.target],
checkedAdd(span[node], edge.delayNs));
@@ -54,6 +66,14 @@ mlir::FailureOr<PeriodicDependencyEstimate> analyzePeriodicDependencies(
if (recurrent.empty() || !upper)
return result;
// Preserve the DAG's exact relaxation order without revisiting adjacency lists.
std::vector<size_t> zeroOrder;
zeroOrder.reserve(outgoing.size());
for (size_t node : order)
llvm::append_range(zeroOrder, outgoingEdges(node));
outgoing.clear(); outgoing.shrink_to_fit();
offsets.clear(); offsets.shrink_to_fit();
const size_t absent = std::numeric_limits<size_t>::max();
std::vector<double> potential(nodeCount);
std::vector<size_t> predecessor(nodeCount, absent);
@@ -92,9 +112,7 @@ mlir::FailureOr<PeriodicDependencyEstimate> analyzePeriodicDependencies(
predecessor[edge.target] = index;
return true;
};
for (size_t node : order)
for (size_t index : outgoing[node])
relax(index);
for (size_t index : zeroOrder) relax(index);
// Eliminate distance-zero paths by a DAG sweep. A simple residual path
// visits at most R recurrent edges, so R+1 rounds detect a positive cycle.
for (size_t round = 0; round <= recurrent.size(); ++round) {
@@ -102,9 +120,7 @@ mlir::FailureOr<PeriodicDependencyEstimate> analyzePeriodicDependencies(
bool changed = false;
for (size_t index : recurrent)
changed |= relax(index);
for (size_t node : order)
for (size_t index : outgoing[node])
changed |= relax(index);
for (size_t index : zeroOrder) changed |= relax(index);
if (!changed)
return false;
// A positive predecessor cycle already proves infeasibility. Keep the
@@ -0,0 +1,65 @@
#pragma once
#include "ResidentWeightPacking.hpp"
#include <algorithm>
#include <limits>
#include <map>
#include <optional>
#include <vector>
namespace onnx_mlir::spatial {
// Scalar proxy queries share immutable task costs and weights within one search.
class PipelinePackingCostCache {
struct TaskOrder {
using is_transparent = void;
bool operator()(llvm::ArrayRef<size_t> a, llvm::ArrayRef<size_t> b) const {
return std::lexicographical_compare(a.begin(), a.end(), b.begin(), b.end());
}
};
struct Entry {
size_t minimumCores, lastUse;
std::vector<std::optional<Cost>> costs;
};
const IndexedResidentWeights &weights;
llvm::ArrayRef<Cost> taskCosts;
size_t residentWeightCapacity, processorCount;
std::map<std::vector<size_t>, Entry, TaskOrder> entries;
size_t clock = 0;
Entry &get(llvm::ArrayRef<size_t> tasks) {
auto found = entries.find(tasks);
if (found == entries.end()) {
// Bound retained geometry independently of the number of legal moves.
if (size() == kMaximumEntries) {
auto oldest = std::min_element(entries.begin(), entries.end(), [](const auto &a, const auto &b) {
return a.second.lastUse < b.second.lastUse;
});
entries.erase(oldest);
}
found = entries.emplace(std::vector<size_t>(tasks.begin(), tasks.end()), Entry{
weights.getPackCountLowerBound(tasks, residentWeightCapacity), 0,
std::vector<std::optional<Cost>>(processorCount + 1)}).first;
}
found->second.lastUse = ++clock;
return found->second;
}
public:
static constexpr size_t kMaximumEntries = 256;
PipelinePackingCostCache(const IndexedResidentWeights &weights, llvm::ArrayRef<Cost> taskCosts,
size_t residentWeightCapacity, size_t processorCount)
: weights(weights), taskCosts(taskCosts), residentWeightCapacity(residentWeightCapacity),
processorCount(processorCount) {}
size_t size() const { return entries.size(); }
size_t minimumCores(llvm::ArrayRef<size_t> tasks) { return get(tasks).minimumCores; }
Cost estimate(llvm::ArrayRef<size_t> tasks, size_t cores) {
auto &value = get(tasks).costs[cores];
if (!value)
value = weights.estimatePackCost(tasks, taskCosts, residentWeightCapacity, cores).value_or(std::numeric_limits<Cost>::max());
return *value;
}
};
} // namespace onnx_mlir::spatial
@@ -19,6 +19,7 @@
#include "ComputeInstanceUtils.hpp"
#include "PipelineScheduling.hpp"
#include "PipelinePackingCostCache.hpp"
#include "ResidentWeightPacking.hpp"
#include "ScheduleEvaluation.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -73,6 +74,7 @@ struct PipelineStageAssignment {
std::vector<size_t> taskStages;
std::vector<size_t> stageSizes;
std::vector<size_t> constructivePacks;
std::vector<size_t> constructiveCores;
};
static bool consumesPipelineInput(const ComputeGraphNode &node) {
@@ -225,161 +227,302 @@ struct PipelinePartitionCandidate {
bool valid = false;
};
static FailureOr<PipelineStageAssignment> initialPipelinePartition(
const PipelineGroupModel &constraints, const PipelineTaskModel &model,
const PipelineCoreLayout &layout, size_t residentWeightCapacity,
std::string &error) {
size_t count = constraints.groups.size(), stages = layout.getStageCount();
if (!count) return PipelineStageAssignment{{},
{layout.getStageSizes().begin(), layout.getStageSizes().end()}, {}};
size_t extra = layout.getProcessorCount() - stages;
// Resource quantization bounds seed construction only. Refinement uses
// individual physical cores and final allocation uses every available core.
size_t units = std::min<size_t>(16, extra);
auto coreCount = [&](size_t amount) { return 1 + (units ? extra * amount / units : 0); };
// Share segment costs while retaining each depth's resource grid and DP tie order.
class PipelinePartitionTable {
using Cell = PipelinePartitionCandidate;
std::vector<std::vector<std::vector<Cell>>> best(stages + 1,
std::vector<std::vector<Cell>>(count + 1, std::vector<Cell>(units + 1)));
const Cost infinity = std::numeric_limits<Cost>::max();
std::vector<std::vector<TaskCosts>> costs(count + 1,
std::vector<TaskCosts>(count + 1, TaskCosts(units + 1, infinity)));
for (size_t start = 0; start < count; ++start) {
TaskList tasks;
TaskCosts ends(count);
Cost span = 0, largest = 0;
llvm::SmallBitVector weights(model.weights.weightCount);
for (size_t end = start + 1; end <= count; ++end) {
size_t index = constraints.groupOrder[end - 1];
const auto &group = constraints.groups[index];
ends[index] = checkedAdd(ends[index], constraints.groupPriorities[index]);
span = std::max(span, ends[index]);
largest = std::max(largest, constraints.groupPriorities[index]);
for (size_t successor : group.successors)
ends[successor] = std::max(ends[successor], ends[index]);
llvm::append_range(tasks, group.tasks);
for (size_t weight : group.weightIds) weights.set(weight);
for (size_t amount = 0; amount <= units; ++amount) {
size_t cores = coreCount(amount);
if (weights.count() > checkedMultiply(cores, residentWeightCapacity)) continue;
if (auto packing = model.weights.estimatePackCost(tasks,
model.schedulingCosts, residentWeightCapacity, cores))
costs[start][end][amount] = checkedAdd(span - largest, *packing);
}
const PipelineGroupModel &groups;
const PipelineTaskModel &model;
size_t processors, capacity, stages, minimumAmount, budget;
TaskList coreCounts, taskOrder, taskOffsets, requestedCores;
std::vector<bool> requestedCore;
std::vector<TaskCosts> spanAdjustments;
std::vector<std::vector<TaskCosts>> costs;
bool sharedDepths;
std::vector<std::optional<PipelineStageAssignment>> plans;
std::vector<std::optional<Cost>> priorities;
std::vector<std::vector<std::vector<Cell>>> best;
static constexpr Cost infinity = std::numeric_limits<Cost>::max();
void buildCosts() {
size_t count = groups.groups.size();
costs.resize(count + 1, std::vector<TaskCosts>(count + 1));
if (sharedDepths) spanAdjustments.resize(count + 1, TaskCosts(count + 1));
taskOffsets.push_back(0);
for (size_t group : groups.groupOrder) {
llvm::append_range(taskOrder, groups.groups[group].tasks);
taskOffsets.push_back(taskOrder.size());
}
}
best[0][0][0].valid = true;
for (size_t stage = 0; stage < stages; ++stage)
for (size_t start = 0; start < count; ++start)
for (size_t used = 0; used <= units; ++used) {
const auto &previous = best[stage][start][used];
if (!previous.valid) continue;
for (size_t end = start + 1; end <= count; ++end) {
if (stage == 0 && end < constraints.minimumStageZeroEnd) continue;
if (stage && constraints.hasCrossingEdge
&& constraints.furthestSuccessorBefore[start] >= end) continue;
for (size_t amount = 0; used + amount <= units; ++amount) {
Cost priority = costs[start][end][amount];
if (priority == infinity) continue;
Cell candidate{std::max(previous.priority, priority), start, amount, true};
auto &current = best[stage + 1][end][used + amount];
if (!current.valid || candidate.priority < current.priority) current = candidate;
for (size_t start = 0; start < count; ++start) {
TaskCosts ends(count);
Cost span = 0, largest = 0;
llvm::SmallBitVector weights(model.weights.weightCount);
for (size_t end = start + 1; end <= count; ++end) {
size_t index = groups.groupOrder[end - 1];
const auto &group = groups.groups[index];
ends[index] = checkedAdd(ends[index], groups.groupPriorities[index]);
span = std::max(span, ends[index]);
largest = std::max(largest, groups.groupPriorities[index]);
for (size_t successor : group.successors)
ends[successor] = std::max(ends[successor], ends[index]);
for (size_t weight : group.weightIds) weights.set(weight);
if (!legalSegment(start != 0, start, end)) continue;
if (sharedDepths) spanAdjustments[start][end] = span - largest;
auto &curve = costs[start][end];
curve.assign(budget + 1, infinity);
auto tasks = segmentTasks(start, end);
std::vector<std::optional<Cost>> shared;
if (minimumAmount) {
size_t scans = 0;
shared = model.weights.estimatePackCostCurve(tasks,
model.schedulingCosts, capacity, processors, &scans, requestedCores);
orderedThresholdScans += scans;
}
for (size_t amount = minimumAmount; amount <= budget; ++amount) {
size_t cores = coreCounts[amount];
if (sharedDepths && !requestedCore[cores]) continue;
if (weights.count() <= checkedMultiply(cores, capacity)) {
++packingQueries;
if (auto packed = minimumAmount ? shared[cores] : model.weights.estimatePackCost(tasks,
model.schedulingCosts, capacity, cores))
curve[amount] = checkedAdd(span - largest, *packed);
}
}
}
size_t selected = 0;
for (size_t stage = 1; stage <= stages; ++stage)
if (best[stage][count][units].valid
&& (!selected || best[stage][count][units].priority
< best[selected][count][units].priority))
selected = stage;
// Quantization controls quality, not feasibility. Complete the same prefix
// DP with minimum constructive core counts if no quantized state survives.
std::vector<std::vector<Cell>> feasible;
if (!selected) {
feasible.resize(stages + 1, std::vector<Cell>(count + 1));
}
}
bool legalSegment(size_t stage, size_t start, size_t end) const {
return !(stage == 0 && end < groups.minimumStageZeroEnd)
&& !(stage && groups.hasCrossingEdge && groups.furthestSuccessorBefore[start] >= end);
}
void buildStates() {
size_t count = groups.groups.size();
best.resize(stages + 1, std::vector<std::vector<Cell>>(count + 1,
std::vector<Cell>(budget + 1)));
best[0][0][0].valid = true;
for (size_t stage = 0; stage < stages; ++stage)
for (size_t start = 0; start < count; ++start) {
// Preserve numeric-depth grid states and their deterministic tie order.
for (size_t used = 0; used <= budget; ++used) {
const auto &previous = best[stage][start][used];
if (!previous.valid) continue;
for (size_t end = start + 1; end <= count; ++end) {
if (!legalSegment(stage, start, end)) continue;
for (size_t amount = 0; used + amount <= budget; ++amount) {
Cost priority = costs[start][end][sharedDepths ? coreCounts[amount] : amount];
if (priority == infinity) continue;
++transitions;
Cell candidate{std::max(previous.priority, priority), start, amount, true};
auto &current = best[stage + 1][end][used + amount];
if (!current.valid || candidate.priority < current.priority) current = candidate;
}
}
}
}
}
std::vector<std::vector<Cell>> minimumResourceStates(size_t minimumStages) {
size_t count = groups.groups.size();
std::vector<std::vector<Cell>> feasible(stages + 1, std::vector<Cell>(count + 1));
feasible[0][0].valid = true;
std::vector<std::vector<size_t>> minimum(count + 1,
std::vector<size_t>(count + 1, std::numeric_limits<size_t>::max()));
for (size_t start = 0; start < count; ++start) {
TaskList tasks;
for (size_t start = 0; start < count; ++start)
for (size_t end = start + 1; end <= count; ++end) {
llvm::append_range(tasks, constraints.groups[constraints.groupOrder[end - 1]].tasks);
const size_t lower = std::max<size_t>(1,
model.weights.getPackCountLowerBound(tasks, residentWeightCapacity));
const size_t maximum = layout.getProcessorCount() - stages + 1;
for (size_t cores = lower; cores <= maximum; ++cores)
if (model.weights.estimatePackCost(tasks, model.schedulingCosts,
residentWeightCapacity, cores)) {
auto tasks = segmentTasks(start, end);
size_t lower = std::max<size_t>(1, model.weights.getPackCountLowerBound(tasks, capacity));
for (size_t cores = lower; cores <= processors - minimumStages + 1; ++cores) {
if (!legalSegment(start != 0, start, end)) break;
if (sharedDepths && !requestedCore[cores] && costs[start][end][cores] == infinity) {
++packingQueries;
if (auto value = model.weights.estimatePackCost(tasks, model.schedulingCosts, capacity, cores))
costs[start][end][cores] = checkedAdd(spanAdjustments[start][end], *value);
}
if (!sharedDepths) ++packingQueries;
bool feasible = sharedDepths ? costs[start][end][cores] != infinity
: bool(model.weights.estimatePackCost(tasks, model.schedulingCosts, capacity, cores));
if (feasible) {
minimum[start][end] = cores;
break;
}
}
}
}
for (size_t stage = 0; stage < stages; ++stage)
for (size_t start = 0; start < count; ++start) {
if (!feasible[stage][start].valid) continue;
for (size_t end = start + 1; end <= count; ++end) {
if (stage == 0 && end < constraints.minimumStageZeroEnd) continue;
if (stage && constraints.hasCrossingEdge
&& constraints.furthestSuccessorBefore[start] >= end) continue;
const size_t cores = minimum[start][end];
if (cores > layout.getProcessorCount()) continue;
Cost used = checkedAdd(feasible[stage][start].priority, static_cast<Cost>(cores));
if (used > layout.getProcessorCount() - (stages - stage - 1)) continue;
if (!legalSegment(stage, start, end) || minimum[start][end] > processors) continue;
Cost used = checkedAdd(feasible[stage][start].priority, static_cast<Cost>(minimum[start][end]));
size_t reserved = minimumStages > stage + 1 ? minimumStages - stage - 1 : 0;
if (used > processors - reserved) continue;
auto &current = feasible[stage + 1][end];
if (!current.valid || used < current.priority)
current = {used, start, cores, true};
current = {used, start, minimum[start][end], true};
}
}
for (size_t stage = 1; stage <= stages; ++stage)
if (feasible[stage][count].valid
&& (!selected || feasible[stage][count].priority < feasible[selected][count].priority))
selected = stage;
return feasible;
}
if (!selected && count) {
error = "pipeline initialization found no constructively packable dependency-adjacent partition; "
"bounded packing failure does not prove physical infeasibility";
return failure();
public:
size_t packingQueries = 0, transitions = 0, orderedThresholdScans = 0;
double segmentSeconds = 0, stateSeconds = 0;
PipelinePartitionTable(const PipelineGroupModel &groups, const PipelineTaskModel &model,
size_t processors, size_t capacity, size_t maximumStages, ArrayRef<size_t> depths = {})
: groups(groups), model(model), processors(processors), capacity(capacity),
stages(maximumStages), minimumAmount(depths.empty() ? 0 : 1), sharedDepths(!depths.empty()) {
budget = sharedDepths ? processors : std::min<size_t>(16, processors - stages);
for (size_t amount = 0; amount <= budget; ++amount)
coreCounts.push_back(sharedDepths ? amount : 1 + (budget ? (processors - stages) * amount / budget : 0));
if (sharedDepths) {
requestedCore.assign(processors + 1, false);
for (size_t depth : depths) {
size_t extra = processors - depth, units = std::min<size_t>(16, extra);
for (size_t amount = 0; amount <= units; ++amount)
requestedCore[1 + (units ? extra * amount / units : 0)] = true;
}
for (size_t cores = 1; cores <= processors; ++cores)
if (requestedCore[cores]) requestedCores.push_back(cores);
}
auto start = std::chrono::steady_clock::now();
buildCosts();
segmentSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
start = std::chrono::steady_clock::now();
if (sharedDepths) {
size_t maximum = stages;
plans.resize(maximum + 1); priorities.resize(maximum + 1);
for (size_t depth : depths) {
stages = depth;
size_t extra = processors - stages;
budget = std::min<size_t>(16, extra);
coreCounts.clear();
for (size_t amount = 0; amount <= budget; ++amount)
coreCounts.push_back(1 + (budget ? extra * amount / budget : 0));
best.clear(); buildStates();
std::string error;
if (auto plan = reconstruct(stages, stages, error, false); succeeded(plan)) {
plans[stages] = std::move(*plan);
const auto &cell = best[stages][groups.groups.size()][budget];
Cost priority = 0;
if (cell.valid) priority = cell.priority;
else {
size_t begin = 0;
for (size_t stage = 0; stage < stages; ++stage) {
size_t end = begin;
while (end < groups.groups.size()
&& plans[stages]->taskStages[groups.groups[groups.groupOrder[end]].tasks.front()] == stage) ++end;
priority = std::max(priority, costs[begin][end][plans[stages]->constructiveCores[stage]]);
begin = end;
}
}
priorities[stages] = priority;
}
}
stages = maximum;
best.clear(); best.shrink_to_fit();
} else buildStates();
stateSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
}
PipelineStageAssignment result;
result.stageSizes.assign(stages, 1);
result.taskStages.resize(constraints.taskToGroup.size());
size_t end = count, remaining = units;
for (size_t stage = selected; stage; --stage) {
const auto &cell = feasible.empty() ? best[stage][end][remaining] : feasible[stage][end];
result.stageSizes[stage - 1] = feasible.empty() ? coreCount(cell.units) : cell.units;
for (size_t position = cell.start; position < end; ++position)
for (size_t task : constraints.groups[constraints.groupOrder[position]].tasks)
result.taskStages[task] = stage - 1;
end = cell.start;
if (feasible.empty()) remaining -= cell.units;
ArrayRef<size_t> segmentTasks(size_t start, size_t end) const {
return ArrayRef(taskOrder).slice(taskOffsets[start], taskOffsets[end] - taskOffsets[start]);
}
// Preserve the constructive witness before adding spare cores. Reordering
// tasks for topology-aware placement must not lose DP feasibility evidence.
result.constructivePacks.resize(model.schedulingCosts.size());
std::vector<TaskList> stageTasks(selected);
for (size_t group : constraints.groupOrder)
for (size_t task : constraints.groups[group].tasks)
stageTasks[result.taskStages[task]].push_back(task);
for (size_t stage = 0; stage < selected; ++stage) {
const auto &tasks = stageTasks[stage];
TaskList packs;
if (!model.weights.estimatePackCost(tasks, model.schedulingCosts,
residentWeightCapacity, result.stageSizes[stage], &packs)) {
error = "pipeline initializer could not reconstruct its constructive packing witness";
std::optional<Cost> priority(size_t depth) const {
if (sharedDepths) return priorities[depth];
const auto &cell = best[depth][groups.groups.size()][budget];
return cell.valid ? std::optional<Cost>(cell.priority) : std::nullopt;
}
FailureOr<PipelineStageAssignment> reconstruct(size_t minimumStages, size_t maximumStages,
std::string &error, bool constructPacks) {
size_t selected = 0, count = groups.groups.size();
for (size_t stage = minimumStages; stage <= maximumStages; ++stage)
if (best[stage][count][budget].valid && (!selected
|| best[stage][count][budget].priority < best[selected][count][budget].priority)) selected = stage;
std::vector<std::vector<Cell>> feasible;
if (!selected) {
feasible = minimumResourceStates(minimumStages);
for (size_t stage = minimumStages; stage <= maximumStages; ++stage)
if (feasible[stage][count].valid
&& (!selected || feasible[stage][count].priority < feasible[selected][count].priority)) selected = stage;
}
if (!selected) {
error = "pipeline initialization found no constructively packable nonempty "
"dependency-adjacent partition in the requested depth domain; "
"bounded packing failure does not prove physical infeasibility";
return failure();
}
for (auto [index, task] : llvm::enumerate(tasks))
result.constructivePacks[task] = packs[index];
PipelineStageAssignment result;
result.stageSizes.assign(selected, 1);
result.taskStages.resize(groups.taskToGroup.size());
size_t end = count, remaining = budget;
for (size_t stage = selected; stage; --stage) {
const auto &cell = feasible.empty() ? best[stage][end][remaining] : feasible[stage][end];
result.stageSizes[stage - 1] = feasible.empty() ? coreCounts[cell.units] : cell.units;
for (size_t task : segmentTasks(cell.start, end)) result.taskStages[task] = stage - 1;
end = cell.start;
if (feasible.empty()) remaining -= cell.units;
}
result.constructiveCores = result.stageSizes;
if (constructPacks && failed(constructWitness(result, error))) return failure();
size_t assigned = std::accumulate(result.stageSizes.begin(), result.stageSizes.end(), size_t{0});
for (; assigned < processors; ++assigned) {
size_t bestStage = 0;
for (size_t stage = 1; stage < selected; ++stage)
if (result.stageSizes[stage] < result.stageSizes[bestStage]) bestStage = stage;
++result.stageSizes[bestStage];
}
return result;
}
size_t assigned = std::accumulate(result.stageSizes.begin(), result.stageSizes.end(), size_t{0});
for (; assigned < layout.getProcessorCount(); ++assigned) {
size_t bestStage = 0;
for (size_t stage = 1; stage < std::max<size_t>(1, selected); ++stage)
if (result.stageSizes[stage] < result.stageSizes[bestStage]) bestStage = stage;
++result.stageSizes[bestStage];
LogicalResult constructWitness(PipelineStageAssignment &result, std::string &error) {
result.constructivePacks.resize(model.schedulingCosts.size());
size_t begin = 0, count = groups.groups.size();
for (size_t stage = 0; stage < result.stageSizes.size(); ++stage) {
size_t end = begin;
while (end < count && result.taskStages[groups.groups[groups.groupOrder[end]].tasks.front()] == stage) ++end;
auto tasks = segmentTasks(begin, end);
TaskList packs;
if (!model.weights.estimatePackCost(tasks, model.schedulingCosts, capacity,
result.constructiveCores[stage], &packs)) {
error = "pipeline initializer could not reconstruct its constructive packing witness";
return failure();
}
for (auto [index, task] : llvm::enumerate(tasks)) result.constructivePacks[task] = packs[index];
begin = end;
}
return success();
}
return result;
FailureOr<PipelineStageAssignment> partition(size_t minimumStages, size_t maximumStages,
std::string &error, bool constructPacks = true) {
if (!sharedDepths) return reconstruct(minimumStages, maximumStages, error, constructPacks);
size_t selected = 0;
for (size_t depth = minimumStages; depth <= maximumStages; ++depth)
if (plans[depth] && (!selected || *priorities[depth] < *priorities[selected])) selected = depth;
if (!selected) {
error = "shared initialization found no constructive partition at the requested depth";
return failure();
}
auto result = *plans[selected];
if (constructPacks && failed(constructWitness(result, error))) return failure();
return result;
}
};
static FailureOr<PipelineStageAssignment> initialPipelinePartition(
const PipelineGroupModel &groups, const PipelineTaskModel &model,
const PipelineCoreLayout &layout, size_t capacity,
std::string &error, size_t minimumStages) {
if (groups.groups.size() < minimumStages) {
error = "exact pipeline depth requires at least " + std::to_string(minimumStages)
+ " compute operation groups; found " + std::to_string(groups.groups.size());
return failure();
}
PipelinePartitionTable table(groups, model, layout.getProcessorCount(), capacity,
layout.getStageCount());
return table.partition(minimumStages, layout.getStageCount(), error);
}
static bool packPipelineStage(
@@ -679,6 +822,22 @@ static LogicalResult packPipelineAssignment(const ComputeGraph &graph,
const PipelineTaskModel &taskModel, PipelineStageAssignment assignment,
MergeScheduleResult &schedule, size_t pipelineStages,
const SchedulingTarget &physicalTarget, std::string &error) {
if (assignment.taskStages.size() != graph.nodes.size()
|| assignment.stageSizes.size() != pipelineStages || !pipelineStages
|| llvm::is_contained(assignment.stageSizes, 0)
|| std::accumulate(assignment.stageSizes.begin(), assignment.stageSizes.end(), size_t{0})
!= physicalTarget.processorCount
|| llvm::any_of(assignment.taskStages, [&](size_t stage) { return stage >= pipelineStages; })) {
error = "pipeline packing requires complete stage assignments and physical core budgets";
return failure();
}
std::vector<bool> active(pipelineStages);
for (size_t stage : assignment.taskStages) active[stage] = true;
if (llvm::is_contained(active, false)) {
error = "pipeline packing requires exactly " + std::to_string(pipelineStages)
+ " nonempty compute stages";
return failure();
}
std::vector<TaskList> tasksByPhysicalCpu(physicalTarget.processorCount);
std::vector<size_t> searchCoreFloors(pipelineStages, 1);
std::string packingError;
@@ -769,7 +928,7 @@ static LogicalResult packPipelineAssignment(const ComputeGraph &graph,
schedule.costs.boundaryMemory = getPipelineBoundaryMemoryDemand(
graph, assignment.taskStages, physicalTarget);
populateScheduleEquivalenceClasses(graph, tasksByPhysicalCpu, schedule);
return mlir::success();
return verifyPipelineStageCount(schedule, pipelineStages, error);
}
static std::optional<PipelineStageAssignment> allocatePipelineStageCores(
@@ -780,9 +939,11 @@ static std::optional<PipelineStageAssignment> allocatePipelineStageCores(
PipelineStageAssignment assignment;
for (size_t task = 0; task < graph.nodes.size(); ++task) {
size_t stage = taskStages[task];
if (stage >= stageCount) return std::nullopt;
assignment.taskStages.push_back(stage);
tasks[stage].push_back(task);
}
if (llvm::any_of(tasks, [](const auto &stage) { return stage.empty(); })) return std::nullopt;
assignment.stageSizes.resize(stageCount);
std::vector<bool> adjustable(stageCount);
for (size_t stage : adjustableStages) adjustable[stage] = true;
@@ -867,7 +1028,8 @@ static std::optional<PipelineStageAssignment> allocatePipelineStageCores(
static void visitAdjacentPipelinePartitions(const PipelineGroupModel &groups,
const PipelineStageAssignment &original,
const std::function<void(PipelineStageAssignment, TaskList)> &visit) {
const std::function<void(PipelineStageAssignment, TaskList)> &visit,
size_t minimumStages, size_t maximumStages) {
size_t count = groups.groups.size(), stages = original.stageSizes.size(), active = 0;
TaskList current(count);
for (size_t group = 0; group < count; ++group) {
@@ -875,11 +1037,12 @@ static void visitAdjacentPipelinePartitions(const PipelineGroupModel &groups,
active = std::max(active, current[group] + 1);
}
auto emit = [&](TaskList proposed) {
TaskList counts(stages), compact(stages);
TaskList counts(maximumStages), compact(maximumStages);
for (size_t stage : proposed) ++counts[stage];
size_t next = 0;
for (size_t stage = 0; stage < stages; ++stage)
for (size_t stage = 0; stage < maximumStages; ++stage)
if (counts[stage]) compact[stage] = next++;
if (next < minimumStages || next > maximumStages) return;
for (size_t &stage : proposed) stage = compact[stage];
if (proposed == current) return;
for (size_t group = 0; group < count; ++group) {
@@ -893,6 +1056,17 @@ static void visitAdjacentPipelinePartitions(const PipelineGroupModel &groups,
for (size_t group = 0; group < count; ++group)
for (size_t task : groups.groups[group].tasks)
candidate.taskStages[task] = proposed[group];
if (next != stages) {
candidate.stageSizes.assign(next, std::accumulate(original.stageSizes.begin(),
original.stageSizes.end(), size_t{0}) / next);
size_t remainder = std::accumulate(original.stageSizes.begin(),
original.stageSizes.end(), size_t{0}) % next;
for (size_t stage = 0; stage < remainder; ++stage) ++candidate.stageSizes[stage];
TaskList changed(next);
std::iota(changed.begin(), changed.end(), 0);
visit(std::move(candidate), std::move(changed));
return;
}
for (size_t group = 0; group < count; ++group) {
if (!llvm::is_contained(flows[current[group]], proposed[group]))
flows[current[group]].push_back(proposed[group]);
@@ -930,15 +1104,17 @@ static void visitAdjacentPipelinePartitions(const PipelineGroupModel &groups,
emit(std::move(proposed));
}
}
for (size_t begin = 0; begin + 1 < active; ++begin)
for (size_t end = begin + 1; end < active; ++end) {
auto proposed = current;
for (size_t &stage : proposed)
if (stage > end) stage -= end - begin;
else if (stage >= begin) stage = begin;
emit(std::move(proposed));
}
if (active < stages)
if (minimumStages < active) {
for (size_t begin = 0; begin + 1 < active; ++begin)
for (size_t end = begin + 1; end < active; ++end) {
auto proposed = current;
for (size_t &stage : proposed)
if (stage > end) stage -= end - begin;
else if (stage >= begin) stage = begin;
emit(std::move(proposed));
}
}
if (active < maximumStages)
for (size_t stage = 0; stage < active; ++stage) {
TaskList members;
for (size_t group : groups.groupOrder)
@@ -998,22 +1174,79 @@ PipelineCoreLayout::PipelineCoreLayout(ArrayRef<size_t> sizes,
}
}
size_t getPipelineStageCount(const MergeScheduleResult &schedule) {
return schedule.processorStages.empty() ? 1
: *std::max_element(schedule.processorStages.begin(), schedule.processorStages.end()) + 1;
}
mlir::LogicalResult verifyPipelineStageCount(
const MergeScheduleResult &schedule, size_t stages, std::string &error) {
if (!stages || (schedule.processorStages.empty() ? stages > 1
: schedule.processorStages.size() != schedule.processorCount)) {
error = "pipeline stage map does not cover every physical core";
return failure();
}
std::vector<bool> active(stages), assigned(stages);
for (size_t stage : schedule.processorStages) {
if (stage >= stages) { error = "physical core has an invalid pipeline stage"; return failure(); }
assigned[stage] = true;
}
if (stages == 1 && schedule.processorStages.empty()) assigned[0] = true;
for (const auto &[instance, core] : schedule.computeToCpuMap) {
if (core >= schedule.processorCount) { error = "compute has an invalid physical core"; return failure(); }
active[schedule.processorStages.empty() ? 0 : schedule.processorStages[core]] = true;
}
if (llvm::is_contained(assigned, false) || (stages > 1 && llvm::is_contained(active, false))) {
error = "requested pipeline depth " + std::to_string(stages)
+ " requires exactly that many nonempty compute stages with physical cores";
return failure();
}
return success();
}
static size_t maximumPipelineDepth(const PipelineGroupModel &groups, size_t processors) {
size_t count = groups.groups.size();
std::vector<size_t> stages(count + 1);
for (size_t start = 0; start < count; ++start) {
if (start && !stages[start]) continue;
for (size_t end = start + 1; end <= count; ++end) {
if (!start && end < groups.minimumStageZeroEnd) continue;
if (start && groups.hasCrossingEdge && groups.furthestSuccessorBefore[start] >= end) continue;
stages[end] = std::max(stages[end], stages[start] + 1);
}
}
return std::min(processors, stages[count]);
}
static LogicalResult initializePipelineSchedule(const ComputeGraph &graph,
const PipelineTaskModel &model, const PipelineGroupModel &groups,
size_t minimumStages, size_t maximumStages, const SchedulingTarget &target,
MergeScheduleResult &schedule, std::string &error) {
if (minimumStages > maximumStages) {
error = "requested pipeline depth exceeds the legal ordered compute-stage domain";
return failure();
}
PipelineCoreLayout layout(target.processorCount, maximumStages);
auto partition = initialPipelinePartition(groups, model, layout,
target.residentWeightCapacity, error, minimumStages);
if (failed(partition)) return failure();
return packPipelineAssignment(graph, model, *partition, schedule,
partition->stageSizes.size(), target, error);
}
mlir::LogicalResult runPipelineScheduler(const ComputeGraph &graph,
MergeScheduleResult &schedule, size_t pipelineStages,
const SchedulingTarget &target, std::string &error) {
PipelineCoreLayout layout(target.processorCount, pipelineStages);
if (pipelineStages <= 1 || !layout.isValid()) {
if (pipelineStages == 1 || pipelineStages > target.processorCount || !target.processorCount) {
error = "pipeline scheduling requires a valid physical core layout";
return failure();
}
PipelineTaskModel model(graph, target);
auto groups = buildPipelineGroups(graph, model, error);
if (failed(groups)) return failure();
auto partition = initialPipelinePartition(*groups, model, layout,
target.residentWeightCapacity, error);
if (failed(partition)) return failure();
return packPipelineAssignment(graph, model, *partition, schedule,
pipelineStages, target, error);
size_t maximum = pipelineStages ? pipelineStages : maximumPipelineDepth(*groups, target.processorCount);
return initializePipelineSchedule(graph, model, *groups,
pipelineStages ? pipelineStages : 2, maximum, target, schedule, error);
}
namespace {
@@ -1076,20 +1309,23 @@ struct PipelineProposal {
bool optimizeCores;
TaskList adjustableStages;
};
using PipelineMove = std::tuple<size_t, TaskList, TaskList, bool, TaskList>;
using PipelineMove = std::tuple<TaskList, TaskList, bool, TaskList>;
static std::vector<PipelineProposal> pipelineNeighborhood(const ComputeGraph &graph,
const PipelineTaskModel &baseModel, const PipelineGroupModel &baseGroups,
const PipelineSearchResult &search, ArrayRef<size_t> beam,
const std::set<PipelineMove> &visited, size_t limit) {
PipelinePackingCostCache &packingCosts,
const std::set<PipelineMove> &visited, size_t limit,
size_t minimumDepth, size_t maximumDepth) {
const auto &target = baseModel.target;
std::vector<PipelineProposal> proposals;
for (size_t parent : beam) {
if (getPipelineStageCount(search.schedules[parent]) == 1) continue;
std::vector<PipelineProposal> neighborhood;
const auto &model = baseModel;
const auto &groups = baseGroups;
auto original = scheduleAssignment(graph, search.schedules[parent],
search.schedules[parent].costs.stageSpanNs.size());
getPipelineStageCount(search.schedules[parent]));
size_t stages = original.stageSizes.size();
std::vector<double> pressure(stages);
if (!search.costs.empty()) {
@@ -1105,26 +1341,19 @@ static std::vector<PipelineProposal> pipelineNeighborhood(const ComputeGraph &gr
pressure[search.schedules[parent].processorStages[core]], cost.processorTaskWorkNs[core]);
}
auto originalView = partitionProxy(graph, groups, model, original);
std::vector<std::vector<std::optional<Cost>>> packCosts(stages,
std::vector<std::optional<Cost>>(target.processorCount + 1));
auto packingCost = [&](const TaskList &tasks, size_t cores, size_t stage) {
auto *cached = tasks == originalView.tasks[stage] ? &packCosts[stage][cores] : nullptr;
if (cached && *cached) return **cached;
Cost value = model.weights.estimatePackCost(tasks, model.schedulingCosts,
target.residentWeightCapacity, cores).value_or(std::numeric_limits<Cost>::max());
if (cached) *cached = value;
return value;
auto packingCost = [&](const TaskList &tasks, size_t cores) {
return packingCosts.estimate(tasks, cores);
};
Cost originalBound = originalView.memory;
if (target.estimatePipelineDependencySpan)
originalBound = std::max(originalBound, target.estimatePipelineDependencySpan(originalView.spans));
for (size_t stage = 0; stage < stages; ++stage)
originalBound = std::max(originalBound, packingCost(originalView.tasks[stage],
original.stageSizes[stage], stage));
original.stageSizes[stage]));
auto consider = [&](PipelineStageAssignment assignment, const PartitionProxy &view,
bool optimizeCores = false, ArrayRef<size_t> adjustable = {}) {
TaskList adjustableStages(adjustable.begin(), adjustable.end());
auto identity = std::make_tuple(parent, assignment.taskStages, assignment.stageSizes,
auto identity = std::make_tuple(assignment.taskStages, assignment.stageSizes,
optimizeCores, adjustableStages);
if (visited.count(identity) || llvm::any_of(neighborhood, [&](const auto &proposal) {
return proposal.parent == parent && proposal.optimizeCores == optimizeCores
@@ -1134,19 +1363,19 @@ static std::vector<PipelineProposal> pipelineNeighborhood(const ComputeGraph &gr
})) return;
Cost maximum = 0, total = 0;
double affected = 0;
std::vector<bool> changed(stages);
size_t proposedStages = assignment.stageSizes.size();
std::vector<bool> changed(std::max(stages, proposedStages));
for (size_t task = 0; task < graph.nodes.size(); ++task)
if (original.taskStages[task] != assignment.taskStages[task])
changed[original.taskStages[task]] = changed[assignment.taskStages[task]] = true;
for (size_t stage = 0; stage < stages; ++stage) {
if (model.weights.getPackCountLowerBound(view.tasks[stage],
target.residentWeightCapacity) > assignment.stageSizes[stage]) return;
Cost work = packingCost(view.tasks[stage], assignment.stageSizes[stage], stage);
for (size_t stage = 0; stage < proposedStages; ++stage) {
if (packingCosts.minimumCores(view.tasks[stage]) > assignment.stageSizes[stage]) return;
Cost work = packingCost(view.tasks[stage], assignment.stageSizes[stage]);
if (work == std::numeric_limits<Cost>::max()) return;
maximum = std::max(maximum, work);
// A within-generation span is latency, not an initiation-interval bound.
total = checkedAdd(total, view.spans[stage]);
if (changed[stage] || original.stageSizes[stage] != assignment.stageSizes[stage])
if (stage < stages && (changed[stage] || original.stageSizes[stage] != assignment.stageSizes[stage]))
affected += pressure[stage];
}
double priority = maximum;
@@ -1171,11 +1400,11 @@ static std::vector<PipelineProposal> pipelineNeighborhood(const ComputeGraph &gr
visitAdjacentPipelinePartitions(groups, original,
[&](PipelineStageAssignment partition, TaskList adjustable) {
if (auto allocation = allocatePipelineStageCores(graph, model,
partition.taskStages, stages, target, partition.stageSizes, adjustable, false)) {
partition.taskStages, partition.stageSizes.size(), target, partition.stageSizes, adjustable, false)) {
auto view = partitionProxy(graph, groups, model, *allocation);
consider(std::move(*allocation), view, true, adjustable);
}
});
}, minimumDepth, maximumDepth);
// Core transfers are coordinate moves in the same joint neighborhood.
std::vector<bool> active(stages);
for (size_t stage : original.taskStages) active[stage] = true;
@@ -1201,34 +1430,132 @@ static std::vector<PipelineProposal> pipelineNeighborhood(const ComputeGraph &gr
}
} // namespace
// Cover the depth-domain endpoints, then bisect the largest remaining gap.
// The logarithmic seed count reserves at least half the queue budget for refinement.
static TaskList pipelineDepthSeeds(size_t maximum, size_t budget) {
if (maximum < 2 || budget < 2) return {};
size_t count = 1;
for (size_t width = maximum - 1; width > 1; width = (width + 1) / 2) ++count;
count = std::min({count, maximum - 1, budget / 2});
TaskList result{2};
if (count == 1) return result;
result.push_back(maximum);
std::vector<std::pair<size_t, size_t>> intervals{{2, maximum}};
while (result.size() < count) {
auto best = std::max_element(intervals.begin(), intervals.end(), [](auto a, auto b) {
if (a.second - a.first != b.second - b.first)
return a.second - a.first < b.second - b.first;
return a.first > b.first;
});
if (best == intervals.end() || best->second - best->first <= 1) break;
auto [low, high] = *best;
intervals.erase(best);
size_t middle = low + (high - low) / 2;
result.push_back(middle);
intervals.emplace_back(low, middle);
intervals.emplace_back(middle, high);
}
return result;
}
static TaskList pipelineFrontier(const PipelineSearchResult &search,
const std::set<size_t> &exhausted, size_t evaluationBudget) {
size_t lookahead = 1;
while (lookahead * lookahead < evaluationBudget) ++lookahead;
TaskList ranked;
for (size_t i = 0; i < search.schedules.size(); ++i)
if (!exhausted.count(i) && getPipelineStageCount(search.schedules[i]) > 1) ranked.push_back(i);
auto score = [&](size_t i) {
const auto &cost = search.costs[i];
return std::make_tuple(std::max(*cost.primary, *cost.alternative), *cost.primary, *cost.alternative, i);
};
llvm::sort(ranked, [&](size_t a, size_t b) { return score(a) < score(b); });
if (ranked.size() <= kPipelineBeamWidth) return ranked;
size_t frontier = ranked[1];
for (size_t i : ranked)
if (i != ranked.front() && std::make_tuple(search.expansions[i] / lookahead, std::get<0>(score(i)), i)
< std::make_tuple(search.expansions[frontier] / lookahead, std::get<0>(score(frontier)), frontier)) frontier = i;
return {ranked.front(), frontier};
}
mlir::FailureOr<PipelineSearchResult> optimizePipelineSchedule(
const ComputeGraph &graph, const MergeScheduleResult &initial,
size_t pipelineStages, const SchedulingTarget &target,
const PipelineCostQuery &evaluate, std::string &error, size_t evaluationBudget) {
error.clear();
if (!evaluationBudget || evaluationBudget > kPipelineQueueEvaluationBudget) {
error = "pipeline queue-evaluation budget must be between one and thirty-two";
error = "pipeline queue-evaluation budget must be between one and "
+ std::to_string(kPipelineQueueEvaluationBudget);
return failure();
}
auto groupingStart = std::chrono::steady_clock::now();
PipelineTaskModel model(graph, target);
auto executionModel = model;
executionModel.schedulingCosts = model.executionCosts;
auto groups = buildPipelineGroups(graph, model, error);
if (failed(groups)) return failure();
PipelineSearchResult result;
result.groupingSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - groupingStart).count();
result.schedules.push_back(initial); result.parents.push_back(0);
if (groups->groups.size() <= 1) {
result.maximumDepth = pipelineStages ? pipelineStages : maximumPipelineDepth(*groups, target.processorCount);
if (failed(verifyPipelineStageCount(initial,
pipelineStages ? pipelineStages : getPipelineStageCount(initial), error))) return failure();
if (!pipelineStages && result.maximumDepth > 1 && evaluationBudget > 1) {
auto seedDepths = pipelineDepthSeeds(result.maximumDepth, evaluationBudget);
PipelinePartitionTable table(*groups, model, target.processorCount,
target.residentWeightCapacity, result.maximumDepth, seedDepths);
result.segmentCostSeconds = table.segmentSeconds;
result.partitionStateSeconds = table.stateSeconds;
result.initializationPackingQueries = table.packingQueries;
result.partitionTransitions = table.transitions;
result.orderedThresholdScans = table.orderedThresholdScans;
for (size_t depth : seedDepths)
if (auto priority = table.priority(depth)) result.depthPriorities.emplace_back(depth, *priority);
else result.depthInitializationFailures.emplace_back(depth,
"shared prefix/resource states have no constructive feasible partition; this does not prove physical infeasibility");
std::set<TaskList> seeds{placementIdentity(graph, initial)};
for (size_t depth : seedDepths) {
MergeScheduleResult seed;
std::string initializationError;
++result.packingAttempts;
auto start = std::chrono::steady_clock::now();
auto partition = table.partition(depth, depth, initializationError);
auto packed = failed(partition) ? failure() : packPipelineAssignment(graph,
model, *partition, seed, depth, target, initializationError);
result.initializationPlacementSeconds += std::chrono::duration<double>(
std::chrono::steady_clock::now() - start).count();
if (failed(packed)) {
result.depthInitializationFailures.emplace_back(depth, initializationError);
continue;
}
if (seeds.insert(placementIdentity(graph, seed)).second) {
result.schedules.push_back(std::move(seed));
result.parents.push_back(0);
}
}
}
result.initialSchedules = result.schedules.size();
if (pipelineStages && groups->groups.size() <= 1) {
result.termination = "one operation group; queue comparison skipped";
return result;
}
std::set<PipelineMove> visited;
std::set<TaskList> placements{placementIdentity(graph, initial)};
std::vector<size_t> beam{0};
std::set<std::pair<TaskList, TaskList>> attemptedAssignments;
for (const auto &schedule : result.schedules) placements.insert(placementIdentity(graph, schedule));
std::vector<size_t> beam;
PipelinePackingCostCache packingCosts(executionModel.weights, executionModel.schedulingCosts,
target.residentWeightCapacity, target.processorCount);
for (size_t i = 0; i < result.schedules.size(); ++i)
if (getPipelineStageCount(result.schedules[i]) > 1) beam.push_back(i);
size_t minimumDepth = pipelineStages ? pipelineStages : 2;
auto neighborhoodStart = std::chrono::steady_clock::now();
auto proposals = pipelineNeighborhood(graph, executionModel, *groups, result, beam, visited, 2 * evaluationBudget);
std::vector<PipelineProposal> proposals;
if (pipelineStages) proposals = pipelineNeighborhood(graph, executionModel, *groups, result, beam,
packingCosts, visited, 2 * evaluationBudget, minimumDepth, result.maximumDepth);
result.neighborhoodSeconds = std::chrono::duration<double>(
std::chrono::steady_clock::now() - neighborhoodStart).count();
if (proposals.empty() || evaluationBudget == 1) {
if (pipelineStages && (proposals.empty() || evaluationBudget == 1)) {
result.termination = "one neighborhood schedule; queue comparison skipped";
return result;
}
@@ -1243,35 +1570,49 @@ mlir::FailureOr<PipelineSearchResult> optimizePipelineSchedule(
})) error = "pipeline optimizer target guidance is unrepresentable";
}
if (error.empty() && !value.eligible()) error = "pipeline optimizer candidate "
+ std::to_string(result.costs.size()) + " cannot compare complete cost evidence: " + value.reason;
+ std::to_string(result.costs.size()) + " at depth " + std::to_string(getPipelineStageCount(schedule))
+ " cannot compare complete cost evidence: " + value.reason;
return value;
};
result.costs.push_back(cost(initial));
if (!error.empty()) return failure();
if (result.schedules.size() > evaluationBudget) result.schedules.resize(evaluationBudget);
result.parents.resize(result.schedules.size());
for (const auto &schedule : result.schedules) {
result.costs.push_back(cost(schedule));
if (!error.empty()) return failure();
}
result.expansions.resize(result.schedules.size());
std::set<size_t> exhausted;
beam = pipelineFrontier(result, exhausted, evaluationBudget);
constexpr size_t maximumPackingAttempts = 4 * kPipelineQueueEvaluationBudget;
while (result.costs.size() < evaluationBudget && result.packingAttempts < maximumPackingAttempts) {
neighborhoodStart = std::chrono::steady_clock::now();
proposals = pipelineNeighborhood(graph, executionModel, *groups, result, beam, visited, 2 * evaluationBudget);
llvm::erase_if(beam, [&](size_t i) { return i >= result.schedules.size(); });
proposals = pipelineNeighborhood(graph, executionModel, *groups, result, beam,
packingCosts, visited, 2 * evaluationBudget, minimumDepth, result.maximumDepth);
for (size_t parent : beam)
if (llvm::none_of(proposals, [&](const auto &p) { return p.parent == parent; })) exhausted.insert(parent);
result.neighborhoodSeconds += std::chrono::duration<double>(
std::chrono::steady_clock::now() - neighborhoodStart).count();
if (proposals.empty()) break;
size_t peftSeeds = llvm::count_if(result.schedules,
[](const auto &schedule) { return getPipelineStageCount(schedule) == 1; });
if (proposals.empty() && exhausted.size() + peftSeeds >= result.schedules.size()) break;
++result.rounds;
size_t evaluated = 0;
std::set<size_t> expandedParents;
std::vector<size_t> ranked;
for (const auto &proposal : proposals) {
if (result.costs.size() == evaluationBudget || evaluated == kPipelineBeamWidth
|| result.packingAttempts == maximumPackingAttempts) break;
if (beam.size() > 1 && expandedParents.count(proposal.parent)) continue;
auto move = std::make_tuple(proposal.assignment.taskStages, proposal.assignment.stageSizes,
proposal.optimizeCores, proposal.adjustableStages);
if (!visited.insert(std::move(move)).second) continue;
++result.consideredAssignments;
++result.packingAttempts;
visited.emplace(proposal.parent, proposal.assignment.taskStages,
proposal.assignment.stageSizes, proposal.optimizeCores, proposal.adjustableStages);
auto packingStart = std::chrono::steady_clock::now();
auto assignment = proposal.assignment;
if (proposal.optimizeCores) {
auto allocated = allocatePipelineStageCores(graph, executionModel, assignment.taskStages,
pipelineStages, target, assignment.stageSizes, proposal.adjustableStages);
assignment.stageSizes.size(), target, assignment.stageSizes, proposal.adjustableStages);
if (!allocated) {
result.packingSeconds += std::chrono::duration<double>(
std::chrono::steady_clock::now() - packingStart).count();
@@ -1279,10 +1620,11 @@ mlir::FailureOr<PipelineSearchResult> optimizePipelineSchedule(
}
assignment = std::move(*allocated);
}
if (!attemptedAssignments.emplace(assignment.taskStages, assignment.stageSizes).second) continue;
MergeScheduleResult schedule;
std::string packingError;
auto packed = packPipelineAssignment(graph, model, assignment, schedule,
pipelineStages, target, packingError);
assignment.stageSizes.size(), target, packingError);
result.packingSeconds += std::chrono::duration<double>(
std::chrono::steady_clock::now() - packingStart).count();
if (failed(packed)) continue;
@@ -1292,24 +1634,15 @@ mlir::FailureOr<PipelineSearchResult> optimizePipelineSchedule(
result.parents.push_back(proposal.parent);
result.schedules.push_back(std::move(schedule));
result.costs.push_back(std::move(estimate));
result.expansions.push_back(0);
++result.expansions[proposal.parent];
expandedParents.insert(proposal.parent);
++evaluated;
}
ranked.resize(result.costs.size());
std::iota(ranked.begin(), ranked.end(), 0);
llvm::sort(ranked, [&](size_t a, size_t b) {
auto score = [&](size_t i) {
const auto &value = result.costs[i];
return std::make_tuple(std::max(*value.primary, *value.alternative),
*value.primary, *value.alternative, i);
};
return score(a) < score(b);
});
ranked.resize(std::min(kPipelineBeamWidth, ranked.size()));
beam = std::move(ranked);
beam = pipelineFrontier(result, exhausted, evaluationBudget);
}
result.termination = result.costs.size() == evaluationBudget ? "queue-evaluation budget"
: result.packingAttempts == maximumPackingAttempts ? "physical-packing budget" : "beam neighborhood exhausted";
: result.packingAttempts == maximumPackingAttempts ? "physical-packing budget" : "frontier neighborhood exhausted";
return result;
}
@@ -1325,7 +1658,7 @@ mlir::LogicalResult packPipelineSchedule(const ComputeGraph &graph,
}
PipelineTaskModel model(graph, target);
return packPipelineAssignment(graph, model,
{{taskStages.begin(), taskStages.end()}, {stageSizes.begin(), stageSizes.end()}, {}},
{{taskStages.begin(), taskStages.end()}, {stageSizes.begin(), stageSizes.end()}, {}, {}},
schedule, stageSizes.size(), target, error);
}
@@ -106,18 +106,31 @@ mlir::LogicalResult runPipelineScheduler(const ComputeGraph &graph,
MergeScheduleResult &schedule, size_t pipelineStages,
const SchedulingTarget &target, std::string &error);
inline constexpr size_t kPipelineQueueEvaluationBudget = 32;
inline constexpr size_t kPipelineQueueEvaluationBudget = 64;
inline constexpr size_t kPipelineBeamWidth = 2;
// PEFT has no explicit processor-stage map and represents depth one.
size_t getPipelineStageCount(const MergeScheduleResult &schedule);
mlir::LogicalResult verifyPipelineStageCount(const MergeScheduleResult &schedule,
size_t stages, std::string &error);
struct PipelineSearchResult {
std::vector<MergeScheduleResult> schedules;
std::vector<WholeScheduleCost> costs;
std::vector<size_t> parents;
size_t rounds = 0, packingAttempts = 0, consideredAssignments = 0;
double neighborhoodSeconds = 0, packingSeconds = 0;
double groupingSeconds = 0, segmentCostSeconds = 0, partitionStateSeconds = 0, initializationPlacementSeconds = 0;
size_t initializationPackingQueries = 0, partitionTransitions = 0, orderedThresholdScans = 0;
std::vector<std::pair<size_t, Cost>> depthPriorities;
size_t maximumDepth = 1;
std::vector<std::pair<size_t, std::string>> depthInitializationFailures;
size_t initialSchedules = 1;
std::vector<size_t> expansions;
std::string termination;
};
// Depth zero permits automatic depth changes; depth one is an unmodified PEFT seed.
// The callback realizes an independent Spatial clone; it must not mutate graph IR.
using PipelineCostQuery = std::function<WholeScheduleCost(const MergeScheduleResult &)>;
mlir::FailureOr<PipelineSearchResult> optimizePipelineSchedule(
@@ -4,6 +4,7 @@
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/STLExtras.h"
#include <map>
#include <numeric>
#include <optional>
#include <tuple>
@@ -12,74 +13,111 @@ namespace onnx_mlir::spatial {
// Greedy maximal prefixes minimize ordered pack count: both constraints are
// closed under taking subranges. Feasible optimum costs cannot increase with packs.
class OrderedPackingCosts {
llvm::ArrayRef<Cost> costs;
const std::vector<std::vector<size_t>> &weights;
llvm::ArrayRef<size_t> tasks;
size_t capacity, generation = 0;
std::vector<size_t> seen;
Cost largest = 0, total = 0;
bool valid = true, cache;
std::map<Cost, size_t> counts;
size_t scans = 0;
size_t packCount(Cost maximum, std::vector<size_t> *assignment = nullptr,
size_t stopAfter = std::numeric_limits<size_t>::max()) {
if (cache && !assignment)
if (auto found = counts.find(maximum); found != counts.end()) return found->second;
++scans;
size_t packs = 1, weightCount = 0;
Cost work = 0;
bool empty = true;
++generation;
for (size_t task : tasks) {
size_t added = 0;
for (size_t weight : weights[task]) added += seen[weight] != generation;
if (!empty && (work > maximum - costs[task] || weightCount + added > capacity)) {
++packs; work = 0; weightCount = 0; ++generation;
if (!cache && !assignment && packs > stopAfter) return packs;
}
work = checkedAdd(work, costs[task]);
for (size_t weight : weights[task])
if (seen[weight] != generation) { seen[weight] = generation; ++weightCount; }
if (assignment) assignment->push_back(packs - 1);
empty = false;
}
if (cache && !assignment) counts.emplace(maximum, packs);
return packs;
}
public:
OrderedPackingCosts(llvm::ArrayRef<Cost> costs,
const std::vector<std::vector<size_t>> &weights, size_t weightCount,
llvm::ArrayRef<size_t> tasks, size_t capacity, bool cache = false)
: costs(costs), weights(weights), tasks(tasks), capacity(capacity), seen(weightCount), cache(cache) {
for (size_t task : tasks) {
valid &= weights[task].size() <= capacity;
largest = std::max(largest, costs[task]);
total = checkedAdd(total, costs[task]);
}
}
size_t getScanCount() const { return scans; }
std::optional<Cost> estimate(size_t packs, std::vector<size_t> *assignment = nullptr,
Cost minimum = 0, std::optional<Cost> maximum = std::nullopt) {
if (assignment) assignment->clear();
if (!valid || !packs) return std::nullopt;
Cost low = std::max({largest, total / packs + (total % packs != 0), minimum});
Cost high = maximum.value_or(total);
if (packCount(high, nullptr, packs) > packs) return std::nullopt;
while (low < high) {
Cost middle = low + (high - low) / 2;
if (packCount(middle, nullptr, packs) <= packs) high = middle;
else low = middle + 1;
}
if (assignment) packCount(low, assignment);
return low;
}
// Requested counts are sorted, unique, and in [1, maximumPacks].
std::vector<std::optional<Cost>> curve(size_t maximumPacks,
llvm::ArrayRef<size_t> requested = {}) {
std::vector<std::optional<Cost>> result(maximumPacks + 1);
if (!valid || !maximumPacks) return result;
std::vector<size_t> all;
if (requested.empty()) {
all.resize(maximumPacks); std::iota(all.begin(), all.end(), 1);
requested = all;
}
size_t first = std::lower_bound(requested.begin(), requested.end(), packCount(total)) - requested.begin();
// Neighboring resource states bound the same monotone minimax query.
auto fill = [&](auto &&fill, size_t first, size_t last, Cost low, Cost high) -> void {
if (first >= last) return;
if (low == high) {
for (size_t i = first; i < last; ++i) result[requested[i]] = low;
return;
}
size_t middle = first + (last - first) / 2;
auto value = estimate(requested[middle], nullptr, low, high);
if (!value) return;
result[requested[middle]] = value;
fill(fill, first, middle, *value, high);
fill(fill, middle + 1, last, low, *value);
};
fill(fill, first, requested.size(), largest, total);
return result;
}
};
inline std::optional<Cost> estimateOrderedPackCost(
llvm::ArrayRef<Cost> taskCosts,
const std::vector<std::vector<size_t>> &taskWeightIds, size_t weightCount,
llvm::ArrayRef<size_t> tasks, size_t residentWeightCapacity,
size_t maximumPacks, std::vector<size_t> *taskPacks = nullptr) {
if (taskPacks) taskPacks->clear();
if (!maximumPacks) return std::nullopt;
Cost low = 0;
Cost high = 0;
for (size_t task : tasks) {
if (taskWeightIds[task].size() > residentWeightCapacity) return std::nullopt;
low = std::max(low, taskCosts[task]);
high = checkedAdd(high, taskCosts[task]);
}
low = std::max(
low, high / maximumPacks + (high % maximumPacks != 0));
std::vector<size_t> seen(weightCount);
size_t generation = 0;
auto pack = [&](Cost maximumCost, std::vector<size_t> *assignment) {
size_t packs = 1;
Cost cost = 0;
size_t packWeightCount = 0;
++generation;
bool fits = true;
bool packEmpty = true;
for (size_t task : tasks) {
size_t addedWeights = 0;
for (size_t weight : taskWeightIds[task])
addedWeights += seen[weight] != generation;
Cost taskCost = taskCosts[task];
bool startsNewPack = !packEmpty
&& (cost > maximumCost - taskCost
|| packWeightCount + addedWeights > residentWeightCapacity);
if (startsNewPack) {
if (++packs > maximumPacks) {
fits = false;
break;
}
cost = 0;
packWeightCount = 0;
++generation;
packEmpty = true;
}
cost = checkedAdd(cost, taskCost);
for (size_t weight : taskWeightIds[task])
if (seen[weight] != generation) {
seen[weight] = generation;
++packWeightCount;
}
if (assignment)
assignment->push_back(packs - 1);
packEmpty = false;
}
return fits;
};
if (!pack(high, nullptr)) return std::nullopt;
while (low < high) {
Cost middle = low + (high - low) / 2;
if (pack(middle, nullptr))
high = middle;
else
low = middle + 1;
}
if (taskPacks) {
pack(low, taskPacks);
}
return low;
OrderedPackingCosts costs(taskCosts, taskWeightIds, weightCount, tasks, residentWeightCapacity);
return costs.estimate(maximumPacks, taskPacks);
}
@@ -131,6 +169,22 @@ struct IndexedResidentWeights {
return *llvm::max_element(loads);
}
// Exact scalar-query parity, with ordered threshold scans shared across cores.
std::vector<std::optional<Cost>> estimatePackCostCurve(llvm::ArrayRef<size_t> tasks,
llvm::ArrayRef<Cost> costs, size_t capacity, size_t maximumPacks, size_t *thresholdScans = nullptr,
llvm::ArrayRef<size_t> requested = {}) const {
OrderedPackingCosts ordered(costs, taskWeightIds, weightCount, tasks, capacity, true);
auto result = ordered.curve(maximumPacks, requested);
if (thresholdScans) *thresholdScans = ordered.getScanCount();
size_t minimum = std::max<size_t>(1, getPackCountLowerBound(tasks, capacity));
if (requested.empty()) {
for (size_t cores = minimum; cores <= maximumPacks; ++cores)
if (!result[cores]) result[cores] = estimatePackCost(tasks, costs, capacity, cores);
} else for (size_t cores : requested)
if (cores >= minimum && !result[cores]) result[cores] = estimatePackCost(tasks, costs, capacity, cores);
return result;
}
// A constructive reuse-first best fit; failure is not a proof of infeasibility.
std::optional<std::vector<size_t>> pack(
llvm::ArrayRef<size_t> tasks, llvm::ArrayRef<Cost> costs,
+5
View File
@@ -99,6 +99,11 @@ add_test(NAME PimQueueSelectionTest
--compiler "$<TARGET_FILE:onnx-mlir>")
set_tests_properties(PimQueueSelectionTest PROPERTIES LABELS pim-unittest)
add_test(NAME PimPipelineDepthTest
COMMAND "${PYTHON_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/test_PipelineDepth.py"
--compiler "$<TARGET_FILE:onnx-mlir>")
set_tests_properties(PimPipelineDepthTest PROPERTIES LABELS pim-unittest)
add_test(NAME PimEmissionTest
COMMAND "${PYTHON_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/test_PimEmission.py"
--compiler "$<TARGET_FILE:onnx-mlir>")
+9
View File
@@ -84,6 +84,15 @@ static void testFrozenDifferential() {
static_cast<unsigned>(1 + component % 7)});
}
differential(componentSize * components, large);
// Sparse adjacency includes empty rows, parallel edges, and shuffled sources.
std::vector<PeriodicDependency> sparse;
for (size_t node = 0; node < 16384; node += 4) {
sparse.push_back({node + 2, node, 11 + node % 7, 1});
sparse.push_back({node, node + 1, 3, 0});
sparse.push_back({node, node + 1, 3, 0});
sparse.push_back({node + 1, node + 2, 5, 0});
}
differential(16384, sparse);
differential(0, {});
differential(2, {{0, 1, 0, 0}, {1, 0, 0, 0}});
differential(2, {{0, 1, 10, 0}, {1, 0, 3, 0}});
+8 -1
View File
@@ -81,7 +81,14 @@ static void testSelectedLoweringFailure() {
Value result = SpatReluOp::create(builder, loc, type, input);
SpatYieldOp::create(builder, loc, ValueRange{result});
builder.setInsertionPointAfter(compute);
func::ReturnOp::create(builder, loc, compute.getResults());
auto next = SpatGraphCompute::create(builder, loc, TypeRange{type}, ValueRange{}, compute.getResults());
next.getBody().push_back(new Block);
Value nextInput = next.getBody().front().addArgument(type, loc);
builder.setInsertionPointToStart(&next.getBody().front());
Value nextResult = SpatReluOp::create(builder, loc, type, nextInput);
SpatYieldOp::create(builder, loc, ValueRange{nextResult});
builder.setInsertionPointAfter(next);
func::ReturnOp::create(builder, loc, next.getResults());
TEST_CHECK(succeeded(verify(*module)));
SchedulingTarget target;
target.processorCount = 2;
+606 -14
View File
@@ -4,10 +4,14 @@
#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>
@@ -16,6 +20,41 @@ 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) {
@@ -50,6 +89,51 @@ static void testPackingFeasibility() {
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>();
@@ -83,6 +167,20 @@ static void testQuantizedFeasibility(uint32_t 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();
}
@@ -90,7 +188,8 @@ static void testQuantizedFeasibility(uint32_t first) {
// 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) {
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)
@@ -115,10 +214,17 @@ static WholeScheduleCost queueCost(const ComputeGraph &graph,
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);
if (schedule.processorStages[core] == 0) {
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) {
@@ -149,9 +255,31 @@ static WholeScheduleCost queueCost(const ComputeGraph &graph,
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 = true;
options.verifyPeriodicState = true; options.searchGuidance = !restartProtocol;
auto estimate = evaluatePimExecution(program, target, options);
auto cost = std::move(estimate.cost);
cost.loweringSupport = StaticLoweringSupport::Supported;
@@ -242,8 +370,9 @@ static void testInitializerPackingWitness() {
}
TEST_CHECK(feasible > 0);
MergeScheduleResult alternative;
TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, std::vector<size_t>(9),
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();
}
@@ -282,7 +411,8 @@ static void testPhysicalReference() {
insertResidentWeights(resident[core], node.residentWeights);
valid &= resident[core].size() <= target.residentWeightCapacity;
}
if (!valid || stages[0] != 0 || stages[1] < stages[0] || stages[2] < stages[1]) continue;
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;
@@ -304,6 +434,198 @@ static void testPhysicalReference() {
<< " 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().
@@ -311,13 +633,14 @@ static void testWorseningIntermediate() {
target.processorCount = 8;
target.interProcessorLatencyNs.assign(64, 0);
ComputeGraph graph;
graph.predecessors.resize(4); graph.successors.resize(4);
for (uint32_t task = 0; task < 4; ++task)
graph.nodes.push_back({{nullptr, task, 1}, 10, {}, task});
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 = llvm::all_of(graph.nodes, [&](const auto &node) {
return schedule.processorStages[schedule.computeToCpuMap.lookup(node.instance)] == 0;
});
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;
@@ -330,9 +653,9 @@ static void testWorseningIntermediate() {
};
std::string error;
MergeScheduleResult initial, intermediate;
TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, initial, 2, target, error)));
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);
const std::vector<size_t> stages(4);
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();
@@ -346,6 +669,10 @@ static void testWorseningIntermediate() {
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
@@ -354,6 +681,69 @@ static void testWorseningIntermediate() {
<< " 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;
@@ -393,6 +783,9 @@ static void runReference(size_t profile, size_t fixture) {
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);
@@ -412,7 +805,7 @@ static void runReference(size_t profile, size_t fixture) {
}
std::vector<bool> used(3);
for (size_t stage : stages) used[stage] = true;
if (used[2] && !used[1]) return;
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};
@@ -454,6 +847,17 @@ static void runReference(size_t profile, size_t fixture) {
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)
@@ -489,12 +893,200 @@ static void runReference(size_t profile, size_t fixture) {
}
}
// 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);
+18 -7
View File
@@ -290,7 +290,8 @@ static void testTransposeCost() {
auto graph = buildCheckedComputeGraph(function, target);
MergeScheduleResult schedule;
std::string error;
assert(succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
assert(succeeded(packPipelineSchedule(graph, std::vector<size_t>(graph.nodes.size()),
{target.processorCount}, target, schedule, error)));
assert(schedule.costs.totalWorkNs == expected);
assert(schedule.costs.maximumCoreWorkNs == expected);
}
@@ -610,14 +611,15 @@ static void testPipelinePackingCapacity() {
target.interProcessorLatencyNs.assign(16, 0);
MergeScheduleResult schedule;
std::string error;
assert(succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
assert(succeeded(packPipelineSchedule(graph, std::vector<size_t>(graph.nodes.size()),
{target.processorCount}, target, schedule, error)));
std::vector<Cost> loads(4);
for (const auto &node : graph.nodes) {
size_t core = schedule.computeToCpuMap.lookup(node.instance);
assert(schedule.processorStages[core] == 0);
loads[core] += node.cost;
}
assert(*std::max_element(loads.begin(), loads.end()) == 20);
assert(*std::max_element(loads.begin(), loads.end()) == 10);
}
static void testRepeatedResidentWeightPacking() {
@@ -659,7 +661,10 @@ static void testRepeatedResidentWeightPacking() {
target.interProcessorLatencyNs[core * target.processorCount + core] = 0;
MergeScheduleResult schedule;
std::string error;
assert(succeeded(runPipelineScheduler(graph, schedule, stages, target, error)));
assert(failed(runPipelineScheduler(graph, schedule, stages, target, error)));
assert(error.find("compute operation groups") != std::string::npos);
assert(succeeded(packPipelineSchedule(graph, std::vector<size_t>(graph.nodes.size()),
{target.processorCount}, target, schedule, error)));
assert(schedule.dominanceOrderCompute.size() == graph.nodes.size());
std::vector<ResidentWeightSet> residents(target.processorCount);
for (const auto &node : graph.nodes) {
@@ -1747,7 +1752,9 @@ static void testPipelineSharedMemory() {
graph.publicationTransfers.push_back({0, 3, {}, 0, {}, 4096, 1});
MergeScheduleResult constrained;
std::string error;
assert(mlir::succeeded(runPipelineScheduler(graph, constrained, 4, target, error)));
assert(mlir::failed(runPipelineScheduler(graph, constrained, 4, target, error)));
assert(error.find("requested depth domain") != std::string::npos);
assert(mlir::succeeded(runPipelineScheduler(graph, constrained, 2, target, error)));
auto assignment = stages(constrained);
for (size_t task = 0; task < graph.nodes.size(); ++task)
for (const auto &[predecessor, transfer] : graph.predecessors[task])
@@ -1969,10 +1976,10 @@ int main() {
ComputeGraph emptyGraph;
MergeScheduleResult emptySchedule;
emptySchedule.processorCount = 4;
assert(mlir::succeeded(runPipelineScheduler(
assert(mlir::failed(runPipelineScheduler(
emptyGraph, emptySchedule, 2, physical, pipelineError)));
assert(emptySchedule.processorCount == physical.processorCount);
assert(emptySchedule.processorStages == std::vector<size_t>({0, 0, 1, 1}));
assert(pipelineError.find("compute operation groups") != std::string::npos);
MergeScheduleResult pipelineSchedule = physicalSchedule;
assert(mlir::succeeded(runPipelineScheduler(
@@ -2005,6 +2012,10 @@ int main() {
fourStagePhysical.interProcessorLatencyNs[core * 8 + core] = 0;
MergeScheduleResult fourStageSchedule = physicalSchedule;
fourStageSchedule.processorCount = 8;
assert(mlir::failed(runPipelineScheduler(
graph, fourStageSchedule, 4, fourStagePhysical, pipelineError)));
graph.successors[0].pop_back();
graph.predecessors[5].pop_back();
assert(mlir::succeeded(runPipelineScheduler(
graph, fourStageSchedule, 4, fourStagePhysical, pipelineError)));
for (size_t task = 0; task < graph.nodes.size(); ++task)
+2 -2
View File
@@ -161,14 +161,14 @@ def run_functional(artifact, inputs, references, outputs, out, policy, seed, rea
command = [str(sync.RUST_BINARY), "--folder", str(artifact),
"--output", str(out / "output.bin"), "--dump",
sync.compare.build_dump_ranges(artifact / "config.json", outputs),
"--mode", "throughput", "--batch-size", "8", "--input-dir", str(inputs[0].parent),
"--mode", "throughput", "--batch-size", str(len(inputs)), "--input-dir", str(inputs[0].parent),
"--batch-output-dir", str(out / "iterations"),
"--diagnostic-schedule-policy", policy, "--diagnostic-schedule-seed", str(seed)]
if reader:
command += ["--diagnostic-target-stall", f"{reader[0] + 1}:{reader[2]}:1:10000"]
result = subprocess.run(command, capture_output=True, text=True, timeout=120)
assert result.returncode == 0, result.stdout + result.stderr
for index, reference in enumerate(references[:8]):
for index, reference in enumerate(references[:len(inputs)]):
result = sync.compare.compare_simulator_outputs(
out / "iterations" / f"output_{index:06d}.bin", outputs, reference,
threshold=1e-3, rtol=1e-5)
+32 -12
View File
@@ -39,7 +39,7 @@ def main():
out = root / source.stem / target
out.mkdir(parents=True)
result = subprocess.run([args.compiler, str(source), "-o", str(out / "model"),
"--maccel=PIM", f"--{target}", "--pipeline=2", "--core-count=168",
"--maccel=PIM", f"--{target}", "--pipeline=1", "--core-count=168",
"--crossbar-count=96", "--crossbar-size=128", "--mlir-disable-threading", "--pim-dump-ir",
f"--pim-target-config={REPO}/validation/pimsim_configs/pimcomp/arch-a/throughput_config_1000ms.json"],
capture_output=True, text=True, timeout=30, env=dict(os.environ, OMP_NUM_THREADS="1"))
@@ -50,13 +50,8 @@ def main():
assert result.returncode == 0, output
if target == "EmitPimCodegen":
assert (out / "pim/config.json").exists()
report = json.loads((out / "reports/pipeline_selection.json").read_text())
assert report["generated_candidates"] == 1
assert report["single_candidate_analysis_skipped"]
assert report["lowered_candidates"] == 1
assert "no cost comparison required" in report["reason"]
assert not (out / "reports/pipeline_selection.json").exists()
else:
assert "intermediate emission retains the original schedule" in output
assert not (out / "reports/pipeline_selection.json").exists()
assert not (out / "reports/pipeline_periodic.json").exists()
assert not (out / "pim").exists()
@@ -71,13 +66,38 @@ def main():
assert not (out / "dialects/pim1_buff.mlir").exists()
else:
assert "pim.core" in ir and "memref<" in ir
weights = [helper.make_tensor(f"w{i}", TensorProto.FLOAT, [4, 4],
[float(i + 1) if row == column else 0. for row in range(4) for column in range(4)])
for i in range(2)]
graph = helper.make_graph([helper.make_node("MatMul", ["x", "w0"], ["a"]),
helper.make_node("MatMul", ["a", "w1"], ["z"])], "two-groups",
[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4])],
[helper.make_tensor_value_info("z", TensorProto.FLOAT, [1, 4])], initializer=weights)
chain = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
chain.ir_version = 10
chain_path = root / "two-groups.onnx"
onnx.save(chain, chain_path)
for target in ("EmitSpatial", "EmitPim", "EmitPimBufferized", "EmitPimCodegen"):
out = root / "two-groups" / target
out.mkdir(parents=True)
result = subprocess.run([args.compiler, str(chain_path), "-o", str(out / "model"),
"--maccel=PIM", f"--{target}", "--pipeline=2", "--core-count=2",
"--crossbar-count=1", "--crossbar-size=128", "--mlir-disable-threading"],
capture_output=True, text=True, timeout=30)
assert result.returncode == 0, result.stdout + result.stderr
if target == "EmitPimCodegen":
report = json.loads((out / "reports/pipeline_selection.json").read_text())
assert report["generated_candidates"] == report["lowered_candidates"] == 1
assert report["single_candidate_analysis_skipped"] and report["selected_depth"] == 2
else:
assert "intermediate emission retains the original schedule" in result.stdout + result.stderr
assert not (out / "pim").exists()
out = root / "explicit-diagnostic"
out.mkdir()
result = subprocess.run([args.compiler, str(small), "-o", str(out / "model"),
"--maccel=PIM", "--EmitSpatial", "--pipeline=2", "--core-count=168",
"--crossbar-count=96", "--crossbar-size=128", "--mlir-disable-threading",
"--pim-periodic-diagnostics",
f"--pim-target-config={REPO}/validation/pimsim_configs/pimcomp/arch-a/throughput_config_1000ms.json"],
result = subprocess.run([args.compiler, str(chain_path), "-o", str(out / "model"),
"--maccel=PIM", "--EmitSpatial", "--pipeline=2", "--core-count=2",
"--crossbar-count=1", "--crossbar-size=128", "--mlir-disable-threading",
"--pim-periodic-diagnostics"],
capture_output=True, text=True, timeout=30)
assert result.returncode == 0, result.stdout + result.stderr
assert (out / "reports/pipeline_periodic.json").exists()
-12
View File
@@ -69,18 +69,6 @@ def main():
"--pim-export-spatial-dataflow=all", "--pim-dump-ir", "--pim-periodic-diagnostics"))
assert diagnostic_program == program and diagnostic["selection"] == selected["selection"]
# A single schedule needs no optional timing contract.
source = root / "softmax.onnx"
tensor = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 8])
output = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 8])
model = helper.make_model(helper.make_graph([helper.make_node("Softmax", ["x"], ["y"], axis=1)],
"softmax", [tensor], [output]), opset_imports=[helper.make_opsetid("", 18)])
model.ir_version = 10
onnx.save(model, source)
selected, program, _ = compile(source)
assert selected["single_candidate_analysis_skipped"] and selected["lowered_candidates"] == 1
assert "spatial_cost" not in selected["candidates"][0]
# Softmax remains in the functional program but has no oracle timing work.
model = onnx.load(vgg)
output = model.graph.output[0].name
+203
View File
@@ -0,0 +1,203 @@
"""Strict stage identities and bounded automatic depth selection."""
import argparse
import hashlib
import json
import importlib.util
from pathlib import Path
import struct
import subprocess
import tempfile
import onnx
from onnx import TensorProto, helper
REPO = Path(__file__).resolve().parents[2]
def identity(folder):
return {str(p.relative_to(folder)): hashlib.sha256(p.read_bytes()).hexdigest()
for p in folder.rglob("*") if p.is_file()}
def validate_repeating(artifact, source, out, count):
spec = importlib.util.spec_from_file_location("publication", REPO / "test/PIM/test_HostPublication.py")
publication = importlib.util.module_from_spec(spec)
spec.loader.exec_module(publication)
from types import SimpleNamespace
options = SimpleNamespace(timeout=120, threshold=1e-3, rtol=1e-5)
samples, inputs, _, _ = publication.sync.make_inputs(source, count, 0, out)
references = publication.sync.make_references(source, samples, out, options)
_, outputs = publication.sync.compare.onnx_io(source)
streams = publication.read_events(artifact)
stalled = next((core, 0, event[0]) for core, events in streams.items()
for event in events if event[1] == 25)
for policy, seed in (("greedy", 0), ("randomized", 7), ("randomized", 11)):
publication.run_functional(artifact, inputs, references, outputs,
out / f"functional-{seed}", policy, seed, stalled if seed == 11 else None)
return streams
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--compiler", required=True)
args = parser.parse_args()
with tempfile.TemporaryDirectory(prefix="raptor-depth-") as temporary:
root = Path(temporary)
model = REPO / "validation/networks/pimcomp_models/vgg8/vgg8-mnist-reconstructed.onnx"
def compile_model(label, depth=None, source=model, extra=(), hardware=None):
out = root / label
out.mkdir()
cores, crossbars, config = hardware or (138, 128,
REPO / "validation/pimsim_configs/pimcomp/arch-b/throughput_config_1000ms.json")
command = [args.compiler, str(source), "-o", str(out / "model"),
"--maccel=PIM", "--EmitPimCodegen", f"--core-count={cores}", f"--crossbar-count={crossbars}",
"--crossbar-size=128", "--mlir-disable-threading", "--pim-detect-communication-deadlock",
*([f"--pim-target-config={config}"] if config is not None else []),
*([f"--pipeline={depth}"] if depth is not None else []), *extra]
result = subprocess.run(command, capture_output=True, text=True, timeout=120)
return out, result
latency, result = compile_model("default")
assert result.returncode == 0, result.stderr
explicit, result = compile_model("p1", 1)
assert result.returncode == 0, result.stderr
assert identity(latency / "pim") == identity(explicit / "pim")
assert not (explicit / "reports/pipeline_selection.json").exists()
for depth in (2, 4, 8):
out, result = compile_model(f"p{depth}", depth)
assert result.returncode == 0, result.stderr
report = json.loads((out / "reports/pipeline_selection.json").read_text())
assert report["requested_depth"] == report["selected_depth"] == depth
assert report["lowered_candidates"] == 1
for candidate in report["candidates"]:
assert candidate["active_stage_count"] == depth
assert {run[2] for run in candidate["task_stage_runs"]} == set(range(depth))
assert len(candidate["stage_core_budgets"]) == depth
assert all(candidate["stage_core_budgets"])
counts = {31: 0, 32: 0}
for path in (out / "pim").glob("core_*.pim"):
data = path.read_bytes()
magic, version, count = struct.unpack("<4sII", data[:12])
assert magic == b"PIMB" and version == 1 and len(data) == 12 + count * 20
for opcode in data[12::20]:
if opcode in counts:
counts[opcode] += 1
assert all(counts.values()), counts
previous = None
for repetition in range(2):
out, result = compile_model(f"auto-{repetition}", "auto")
assert result.returncode == 0, result.stderr
report = json.loads((out / "reports/pipeline_selection.json").read_text())
assert report["requested_depth"] == "auto"
assert 1 in report["explored_depths"]
assert report["queue_evaluation_budget"] == 64
assert report["costed_candidates"] <= report["queue_evaluation_budget"]
assert report["packing_attempts"] <= 4 * report["queue_evaluation_budget"]
assert report["beam_width"] == 2 and report["lowered_candidates"] == 1
assert all(c["eligible"] and c["convergence_proved"] for c in report["candidates"])
selected = report["candidates"][report["selection"]]
assert selected["active_stage_count"] == report["selected_depth"]
current = (report["selection"], report["selected_depth"],
[(c["interval_ns"], c["alternative_tie_interval_ns"]) for c in report["candidates"]],
identity(out / "pim"))
assert previous is None or current == previous
previous = current
# All dynamic-input consumers stay in stage zero. Auto p1 repeats the
# PEFT placement with the existing restart barrier, unlike latency p1.
weights = [helper.make_tensor(f"w{i}", TensorProto.FLOAT, [4, 4],
[float(i + 1) if r == c else 0. for r in range(4) for c in range(4)])
for i in range(2)]
fork = helper.make_model(helper.make_graph([
helper.make_node("MatMul", ["x", "w0"], ["y0"]),
helper.make_node("MatMul", ["x", "w1"], ["y1"])], "input-fork",
[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4])],
[helper.make_tensor_value_info(f"y{i}", TensorProto.FLOAT, [1, 4]) for i in range(2)],
initializer=weights), opset_imports=[helper.make_opsetid("", 18)])
fork.ir_version = 10
fork_path = root / "fork.onnx"
onnx.save(fork, fork_path)
out, result = compile_model("auto-p1", "auto", fork_path)
assert result.returncode == 0, result.stderr
report = json.loads((out / "reports/pipeline_selection.json").read_text())
assert report["selected_depth"] == 1 and report["explored_depths"] == [1]
assert report["candidates"][0]["eligible"] and report["costed_candidates"] == 1
opcodes = [opcode for path in (out / "pim").glob("core_*.pim")
for opcode in path.read_bytes()[12::20]]
assert 31 in opcodes and 32 in opcodes
validate_repeating(out / "pim", fork_path, out, 8)
_, result = compile_model("infeasible-fork", 2, fork_path)
assert result.returncode != 0 and "dependency-adjacent partition" in result.stderr
# Each weighted consumer also reads x, so every compute group stays in stage zero.
weights = [helper.make_tensor(f"w{i}", TensorProto.FLOAT, [128, 128],
[float(i + 1) if r == c else 0. for r in range(128) for c in range(128)])
for i in range(4)]
nodes = [helper.make_node("MatMul", ["x", "w0"], ["a"]),
helper.make_node("Gemm", ["a", "w1", "x"], ["b"]),
helper.make_node("Gemm", ["a", "w2", "x"], ["c"]),
helper.make_node("Gemm", ["b", "w3", "x"], ["tail0"])]
for i in range(24):
nodes.append(helper.make_node("Add", [f"tail{i}", "x"], [f"tail{i+1}"]))
tail = helper.make_model(helper.make_graph(nodes, "uneven-publications",
[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 128])],
[helper.make_tensor_value_info(name, TensorProto.FLOAT, [1, 128])
for name in ("a", "b", "c", "tail24")], initializer=weights),
opset_imports=[helper.make_opsetid("", 18)])
tail.ir_version = 10
tail_path = root / "long-tail.onnx"
onnx.save(tail, tail_path)
profile = REPO / "validation/pimsim_configs/pimcomp/arch-b/throughput_config_1000ms.json"
config = json.loads(profile.read_text())
config["chip_config"]["core_cnt"] = 4
config["chip_config"]["core_config"]["matrix_config"]["xbar_array_count"] = 1
network = config["chip_config"]["network_config"]
network["layout"] = [1, 4]
network["net_config_file_path"] = str(profile.parent / network["net_config_file_path"])
config_path = root / "four-core.json"
config_path.write_text(json.dumps(config))
out, result = compile_model("auto-long-tail", "auto", tail_path, hardware=(4, 1, config_path))
assert result.returncode == 0, result.stdout + result.stderr
report = json.loads((out / "reports/pipeline_selection.json").read_text())
assert report["selected_depth"] == 1 and report["explored_depths"] == [1]
candidate = report["candidates"][0]
assert candidate["convergence_proved"]
queue = candidate["spatial_cost"]["memory_controller"]
assert queue["periodic_state"]["complete_scheduler_state"]
assert queue["alternative_periodic_state"]["complete_scheduler_state"]
assert "complete-state recurrence" in candidate["spatial_cost"]["interval_source"]
latency, result = compile_model("latency-long-tail", 1, tail_path, hardware=(4, 1, config_path))
assert result.returncode == 0, result.stdout + result.stderr
headers = lambda folder: [line for line in (folder / "reports/spatial3_scheduled_no_comm.txt")
.read_text().splitlines() if line.startswith("C")]
assert headers(out) == headers(latency)
for opcode in range(9, 31):
count = lambda folder: sum(path.read_bytes()[12::20].count(opcode)
for path in (folder / "pim").glob("core_*.pim"))
assert count(out) == count(latency), opcode
streams = validate_repeating(out / "pim", tail_path, out, 24)
peers = {core: {event[2] for event in events if event[1] == 29}
for core, events in streams.items()}
assert any(len(receivers) >= 2 for receivers in peers.values()), peers
assert sum(event[1] == 26 for events in streams.values() for event in events) >= 4
lengths = [len(path.read_bytes()[12:]) // 20 for path in (out / "pim").glob("core_*.pim")]
assert max(lengths) > 2 * min(lengths), lengths
graph = helper.make_graph([helper.make_node("Relu", ["x"], ["y"])], "one-group",
[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 16])],
[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 16])])
tiny = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
tiny.ir_version = 10
source = root / "one-group.onnx"
onnx.save(tiny, source)
_, result = compile_model("infeasible", 2, source)
assert result.returncode != 0 and "compute operation groups" in result.stderr
for value in (0, -1, "unknown"):
_, result = compile_model(f"invalid-{value}", value, source)
assert result.returncode != 0 and "positive integer|auto" in result.stderr
_, result = compile_model("cores", 139, source)
assert result.returncode != 0 and "not to exceed --core-count" in result.stderr
_, result = compile_model("intermediate-auto", "auto", source, ("--EmitSpatial",))
assert result.returncode != 0 and "requires full PIM code generation" in result.stderr
if __name__ == "__main__":
main()
+2 -1
View File
@@ -133,7 +133,7 @@ def main():
out = root / f"vgg-{diagnostics}"
out.mkdir()
command = [compiler, str(vgg), "-o", str(out / "model"), "--maccel=PIM", "--EmitPimCodegen",
"--pipeline=2", "--core-count=138", "--crossbar-count=128", "--crossbar-size=128",
"--pipeline=auto", "--core-count=138", "--crossbar-count=128", "--crossbar-size=128",
"--mlir-disable-threading", "--pim-detect-communication-deadlock",
f"--pim-target-config={REPO}/validation/pimsim_configs/pimcomp/arch-b/throughput_config_1000ms.json"]
if diagnostics:
@@ -141,6 +141,7 @@ def main():
result = subprocess.run(command, cwd=root, env=env, capture_output=True, text=True, timeout=90)
assert result.returncode == 0, result.stdout + result.stderr
report = check_selected_program(out)
assert report["requested_depth"] == "auto" and report["selected_depth"] >= 1
assert bool((out / "dialects").exists()) == diagnostics
identity = {str(p.relative_to(out / "pim")): hashlib.sha256(p.read_bytes()).hexdigest()
for p in (out / "pim").rglob("*") if p.is_file()}
+6
View File
@@ -112,6 +112,12 @@ The regular validator uses these paths beside each model:
| `artifacts/raptor/throughput/pim/` | Pipeline-4 throughput artifacts. |
| `artifacts/simulation/{latency,throughput}/` | Functional simulator outputs. |
`--throughput-pipeline=N` selects an exact throughput depth; `auto` requests
automatic depth selection. The default remains 4. Compiler logs report the
selected depth when a selection report is available. Small operation graphs can
reject exact depth 4. Use `--throughput-pipeline=auto` to test their feasible
throughput depth without weakening explicit-depth constraints.
The [comparison runner](tools/pim/pimcomp/compare/README.md) has separate
architecture/pipeline directories and shared `artifacts/common/` references.
+32
View File
@@ -1,3 +1,5 @@
import argparse
import json
import re
import shlex
import shutil
@@ -20,6 +22,29 @@ PIM_PASS_LABEL_BY_SUFFIX = dict(PIM_PASS_LABELS)
TIMING_LINE_RE = re.compile(r"^\s*([0-9]+\.[0-9]+)\s+\(\s*[0-9.]+%\)\s+(.+?)\s*$")
def parse_pipeline_depth(value):
if value == "auto":
return value
try:
depth = int(value)
except ValueError:
raise argparse.ArgumentTypeError("pipeline must be a positive integer or auto") from None
if depth < 1:
raise argparse.ArgumentTypeError("pipeline must be a positive integer or auto")
return depth
def pipeline_sort_key(value):
return (str(value) == "auto", 0 if str(value) == "auto" else int(value))
def read_selected_pipeline_depth(pim_dir):
if pim_dir is None:
return None
report = Path(pim_dir).parent / "reports/pipeline_selection.json"
return json.loads(report.read_text()).get("selected_depth") if report.exists() else None
def _parse_pim_pass_timings(output_text):
pass_timings = {}
for line in output_text.splitlines():
@@ -79,6 +104,13 @@ def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
capture_output=True,
timeout_sec=timeout_sec,
)
selected_depth = read_selected_pipeline_depth(pim_dir)
if selected_depth is not None:
message = f"Selected pipeline depth: {selected_depth}"
if reporter is not None:
reporter.log(message)
else:
print(message)
if reporter is None:
print(Fore.GREEN + "Raptor execution successful" + Style.RESET_ALL)
return _parse_pim_pass_timings(output_text)
+71
View File
@@ -20,9 +20,80 @@ sys.path.insert(0, str(TOOLS / "pimcomp/compare"))
import run_ablation as ablation
import run_pimcomp_models as suite
import compare_raptor_pimcomp_model as comparator
from raptor_validation.raptor import parse_pipeline_depth, pipeline_sort_key, read_selected_pipeline_depth
from raptor_validation.validate_one import validate_outputs
def test_default_experiment_depths_exclude_auto():
expected = [("latency", 1), ("throughput", 2), ("throughput", 4), ("throughput", 8)]
assert [(mode, depth) for mode, depth, _ in suite.COMPARISONS] == expected
cases = ablation.runner_argument_sets(["--mode", "latency", "throughput"])
assert [ablation.case_key(arguments) for _, arguments in cases] == [
(mode, str(depth)) for mode, depth in expected]
assert all("auto" not in arguments for _, arguments in ablation.runner_argument_sets([]))
def test_operation_validation_forwards_requested_depth(tmp_path, monkeypatch):
import importlib.util
spec = importlib.util.spec_from_file_location("operation_validation", TOOLS.parents[1] / "validate.py")
operation_validation = importlib.util.module_from_spec(spec)
spec.loader.exec_module(operation_validation)
model = tmp_path / "model.onnx"
model.touch()
monkeypatch.setattr(operation_validation, "discover_onnx_files", lambda _: [model])
monkeypatch.setattr(sys, "argv", ["validate.py", "--operations-dir", str(tmp_path),
"--raptor-path", str(tmp_path / "compiler"), "--onnx-include-dir", str(tmp_path),
"--skip-non-functional-simulation", "--jobs", "1", "--throughput-pipeline", "8"])
class JobReached(Exception):
pass
def check_job(job):
assert job[1]["validation_kwargs"]["throughput_pipeline"] == 8
raise JobReached
monkeypatch.setattr(operation_validation, "run_validation_job", check_job)
with pytest.raises(JobReached):
operation_validation.main()
@pytest.mark.parametrize("depth", ["auto", "3"])
def test_comparison_and_ablation_accept_depth_domain(tmp_path, monkeypatch, depth):
commands = []
monkeypatch.setattr(suite, "run", lambda command, **kwargs: commands.append(command) or 0)
monkeypatch.setattr(sys, "argv", ["run_pimcomp_models.py", "--out-dir", str(tmp_path),
"--models", "vgg8", "--archs", "arch-b", "--mode", "throughput", "--pipeline", depth,
"--raptor-only", "--dry-run"])
assert suite.main() == 0
assert len(commands) == 2
assert f"--raptor-extra-arg=--pipeline={depth}" in commands[-1]
assert f"pipeline{depth}" in commands[-1][commands[-1].index("--out-dir") + 1]
cases = ablation.runner_argument_sets(["--mode", "throughput", "--pipeline", depth])
assert len(cases) == 1 and cases[0][0] == f"throughput/pipeline{depth}"
assert cases[0][1][-1] == depth
def test_pipeline_depth_parser():
for invalid in ["0", "-1", "unknown", "2.5"]:
with pytest.raises(argparse.ArgumentTypeError):
parse_pipeline_depth(invalid)
assert parse_pipeline_depth("1") == 1
assert parse_pipeline_depth("auto") == "auto"
assert sorted(["auto", "8", "1", "3"], key=pipeline_sort_key) == ["1", "3", "8", "auto"]
def test_selected_pipeline_report(tmp_path):
assert read_selected_pipeline_depth(None) is None
assert read_selected_pipeline_depth(tmp_path / "pim") is None
(tmp_path / "reports").mkdir()
report = tmp_path / "reports/pipeline_selection.json"
report.write_text(json.dumps({"requested_depth": "auto", "selected_depth": 3}))
assert read_selected_pipeline_depth(tmp_path / "pim") == 3
report.write_text(json.dumps({"requested_depth": 8}))
assert read_selected_pipeline_depth(tmp_path / "pim") is None
@pytest.mark.parametrize("mode", ["latency", "throughput"])
def test_comparison_report_energy_is_per_inference(tmp_path, mode):
report = tmp_path / "comparison_report.md"
+2
View File
@@ -31,6 +31,8 @@ Other [suite-runner options](../pimcomp/compare/README.md), including `--models`
`--archs`, and `--no-fast`, are forwarded. Supplying `--mode` or `--pipeline`
replaces the default latency/pipeline-4 case set. Mode-only selection runs all
suite pipelines for that mode. Pipeline-only selection uses its matching mode.
`--pipeline auto` requests automatic throughput depth selection; positive numeric
depths retain their exact-stage meaning.
`--only` is rejected because
the wrapper always runs Raptor only. Ctrl+C stops the active variant's workers.
@@ -25,6 +25,7 @@ from raptor_validation.pimcomp_models import (
add_models_argument,
)
from raptor_validation.artifacts import remove_lock_files
from raptor_validation.raptor import parse_pipeline_depth, pipeline_sort_key
@dataclass(frozen=True)
@@ -102,10 +103,12 @@ def runner_argument_sets(forwarded: list[str]) -> tuple[tuple[str, list[str]], .
)
parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
parser.add_argument("--mode", nargs="+", choices=("latency", "throughput"))
parser.add_argument("--pipeline", type=int, choices=(1, 2, 4, 8))
parser.add_argument("--pipeline", type=parse_pipeline_depth)
args, remaining = parser.parse_known_args(forwarded)
cases = [(mode, pipeline) for mode, pipeline in
(("latency", 1), ("throughput", 2), ("throughput", 4), ("throughput", 8))
requested = (("latency", 1), ("throughput", 2), ("throughput", 4), ("throughput", 8)) if args.pipeline is None else (
("latency", 1) if args.pipeline == 1 else ("throughput", args.pipeline),
)
cases = [(mode, pipeline) for mode, pipeline in requested
if (args.mode is None or mode in args.mode)
and (args.pipeline is None or pipeline == args.pipeline)]
if not cases:
@@ -268,7 +271,7 @@ def aggregate_results(
row["arch"],
row["model"],
row["mode"],
int(row["raptor_pipeline"]),
pipeline_sort_key(row["raptor_pipeline"]),
row["comparison"],
)
)
@@ -12,12 +12,15 @@ built release Raptor, Pimcomp, and `pimsim-nn`.
Defaults: all five models, Arch-A/B, latency pipeline 1, throughput pipelines
2/4/8, four workers, and a functional throughput batch of 64.
Use `--help` for the complete option reference.
Use `--help` for the complete option reference. Numeric depths greater than one
require exactly that many active compute stages. `--pipeline auto` requests
throughput depth selection. Reports include `raptor_selected_pipeline`; the
requested depth remains in the existing pipeline field.
| Option | Purpose |
|---|---|
| `--models MODEL [...]`, `--archs ARCH [...]` | Select models and hardware profiles. |
| `--mode {latency,throughput} [...]`, `--pipeline {1,2,4,8}` | Select experiments. |
| `--mode {latency,throughput} [...]`, `--pipeline N\|auto` | Select experiments. |
| `--out-dir PATH` | Suite root; default: `validation/networks/pimcomp_models`. |
| `--common-dir PATH` | Override the shared reference-artifact root. |
| `--only {raptor,pimcomp}` | Rerun one compiler and reuse the other's existing report. |
@@ -44,7 +44,7 @@ from raptor_validation.onnx_utils import (
write_input_batch_csv,
write_inputs_to_memory_bin,
)
from raptor_validation.raptor import compile_with_raptor
from raptor_validation.raptor import compile_with_raptor, read_selected_pipeline_depth
from raptor_validation.output_comparison import VALIDATION_POLICY, compare_output_arrays
from raptor_validation.pimsim_nn import (
export_raptor_pimsim_artifact,
@@ -2160,6 +2160,7 @@ def main():
"raptor_reuse_context": experiment_context(args, model_path, functional_model_path, hardware, "raptor"),
"pimcomp_config": str(args.pimcomp_config),
"raptor_extra_args": args.raptor_extra_arg,
"raptor_selected_pipeline": read_selected_pipeline_depth(raptor_pim_dir),
"reused_raptor_report": optional_path(args.reuse_raptor_report.resolve()) if reuse_raptor else None,
"reused_pimcomp_report": optional_path(reused_pimcomp_report_path) if reuse_pimcomp else None,
"failures": failures,
@@ -30,6 +30,7 @@ from raptor_validation.pimcomp_models import (
from raptor_validation.artifacts import artifacts_dir, remove_lock_files
from raptor_validation.pimsim_nn import parse_pimsim_nn_metrics
from raptor_validation.output_comparison import VALIDATION_POLICY
from raptor_validation.raptor import parse_pipeline_depth, pipeline_sort_key
from raptor_validation.validate_one import STAGE_COLORS
PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN"
@@ -56,7 +57,7 @@ class ComparisonSpec:
common_dir: Path
config: Path
mode: str
pipeline: int
pipeline: int | str
pimcomp_pipeline: str
shared_pimcomp_dir: Path
anchor: bool
@@ -71,7 +72,7 @@ def result_dir(
name: str,
arch: str,
mode: str,
pipeline: int,
pipeline: int | str,
ablation_variant: str | None = None,
) -> Path:
base = artifacts_dir(model_dir(root, name))
@@ -243,7 +244,7 @@ def write_results_csv(
row["arch"],
model_order[row["model"]],
mode_order[row["mode"]],
int(row["raptor_pipeline"]),
pipeline_sort_key(row["raptor_pipeline"]),
)
)
output.parent.mkdir(parents=True, exist_ok=True)
@@ -378,7 +379,7 @@ def comparison_command(
common_dir: Path,
config: Path,
mode: str,
pipeline: int,
pipeline: int | str,
pimcomp_pipeline: str,
pimsim_time_ms: int,
batch_size: int,
@@ -610,9 +611,8 @@ def main() -> int:
)
parser.add_argument(
"--pipeline",
type=int,
choices=(1, 2, 4, 8),
help="Run only this Raptor pipeline within the selected mode.",
type=parse_pipeline_depth,
help="Run this positive Raptor pipeline depth, or auto, within the selected mode.",
)
parser.add_argument(
"--pimsim-time-ms",
@@ -684,11 +684,15 @@ def main() -> int:
parser.error("--batch-size must be positive")
if args.timeout_seconds < 0:
parser.error("--timeout-seconds must be non-negative")
comparisons_by_arch: dict[str, tuple[tuple[str, int, str], ...]] = {}
comparisons_by_arch: dict[str, tuple[tuple[str, int | str, str], ...]] = {}
configs_by_arch: dict[str, dict[str, Path]] = {}
for arch in args.archs:
requested = COMPARISONS if args.pipeline is None else (
("latency", 1, "element") if args.pipeline == 1
else ("throughput", args.pipeline, "batch"),
)
comparisons = tuple(
comparison for comparison in COMPARISONS
comparison for comparison in requested
if comparison[0] in args.mode
and (args.pipeline is None or comparison[1] == args.pipeline)
)
+4 -2
View File
@@ -30,7 +30,7 @@ from raptor_validation.validate_one import (
pimcomp_compatibility_errors,
validate_network,
)
from raptor_validation.raptor import PIM_PASS_LABELS
from raptor_validation.raptor import PIM_PASS_LABELS, parse_pipeline_depth
DEFAULT_PIMCOMP_CONFIG = "arch-a"
PIMCOMP_CONFIG_CHOICES = ("arch-a", "arch-b", "arch-c")
@@ -222,6 +222,8 @@ def main():
help="Skip non-functional simulation.")
ap.add_argument("--no-fast", action="store_true",
help="Disable fast Pimsim throughput convergence for authoritative experiments.")
ap.add_argument("--throughput-pipeline", type=parse_pipeline_depth, default=THROUGHPUT_PIPELINE,
help="Throughput pipeline depth, or auto (default: 4).")
ap.add_argument("--threshold", type=float, default=1e-3,
help="Absolute tolerance for per-element output comparison.")
ap.add_argument("--relative-threshold", type=float, default=1e-5,
@@ -345,7 +347,7 @@ def main():
"raptor_extra_args": raptor_extra_args,
"pimsim_nn_build_dir": pimsim_nn_build_dir,
"pimsim_config_path": selected_pimsim_config,
"throughput_pipeline": THROUGHPUT_PIPELINE,
"throughput_pipeline": a.throughput_pipeline,
"throughput_batch_size": THROUGHPUT_BATCH_SIZE,
"throughput_pimsim_config_path": selected_throughput_pimsim_config,
"pimsim_fast": not a.no_fast,