force selected pipeline stages in throughput mode
Validate Operations / validate-operations (push) Has been cancelled
Validate Operations / validate-operations (push) Has been cancelled
increase throughput scheduler budget add experimental throughput auto mode
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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. |
|
||||
|
||||
@@ -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
@@ -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.
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
+2
-2
@@ -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");
|
||||
|
||||
+18
-4
@@ -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.
|
||||
|
||||
+26
-10
@@ -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
|
||||
|
||||
+65
@@ -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
|
||||
+531
-198
@@ -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 ¤t = 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 ¤t = 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 ¤t = 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);
|
||||
}
|
||||
|
||||
|
||||
+14
-1
@@ -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(
|
||||
|
||||
+117
-63
@@ -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,
|
||||
|
||||
@@ -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>")
|
||||
|
||||
@@ -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}});
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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()}
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user