diff --git a/.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md b/.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md index da04e03..0f4187e 100644 --- a/.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md +++ b/.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md @@ -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 diff --git a/docs/compiler/options.md b/docs/compiler/options.md index 242c572..b64a367 100644 --- a/docs/compiler/options.md +++ b/docs/compiler/options.md @@ -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. | diff --git a/docs/compiler/scheduling-implementation.md b/docs/compiler/scheduling-implementation.md index 1dbcca6..3bb8ae9 100644 --- a/docs/compiler/scheduling-implementation.md +++ b/docs/compiler/scheduling-implementation.md @@ -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. diff --git a/docs/compiler/scheduling.md b/docs/compiler/scheduling.md index fd034b6..701a20e 100644 --- a/docs/compiler/scheduling.md +++ b/docs/compiler/scheduling.md @@ -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. diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp index cc8cf5c..d7655ef 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -132,12 +132,21 @@ llvm::cl::opt llvm::cl::desc("Number of crossbars in each core (required for Pim compilation)"), llvm::cl::init(0)); -llvm::cl::opt pipelineStages( +llvm::cl::opt 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="); + return stages; +} + llvm::cl::opt 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="); 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(coresCount.getValue()) < pipelineStages.getValue()) + size_t stages = getPimPipelineStages(); + if (static_cast(coresCount.getValue()) < stages) llvm::report_fatal_error("Pim compilation requires --pipeline not to exceed --core-count"); + size_t maximumStages = stages ? stages : static_cast(coresCount.getValue()); if (crossbarCountInCore.getValue() - > std::numeric_limits::max() / pipelineStages.getValue()) + > std::numeric_limits::max() / maximumStages) llvm::report_fatal_error("Pim compilation --crossbar-count * --pipeline overflows"); } diff --git a/src/PIM/Compiler/PimCompilerOptions.hpp b/src/PIM/Compiler/PimCompilerOptions.hpp index cbb36e2..8af9a65 100644 --- a/src/PIM/Compiler/PimCompilerOptions.hpp +++ b/src/PIM/Compiler/PimCompilerOptions.hpp @@ -65,7 +65,9 @@ extern llvm::cl::opt pimDisableSpatialPlanning; extern llvm::cl::opt crossbarSize; extern llvm::cl::opt crossbarCountInCore; -extern llvm::cl::opt pipelineStages; +extern llvm::cl::opt pipelineStages; +// Zero denotes automatic depth internally; numeric zero is never accepted. +size_t getPimPipelineStages(); extern llvm::cl::opt coresCount; extern llvm::cl::opt pimTargetConfig; extern llvm::cl::opt pimConvIm2colMaxElements; diff --git a/src/PIM/Compiler/PimCompilerUtils.cpp b/src/PIM/Compiler/PimCompilerUtils.cpp index 28b53c9..16e8d2a 100644 --- a/src/PIM/Compiler/PimCompilerUtils.cpp +++ b/src/PIM/Compiler/PimCompilerUtils.cpp @@ -463,16 +463,20 @@ void addPassesPim(OwningOpRef& 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")); } diff --git a/src/PIM/Compiler/PimPipelineSelection.cpp b/src/PIM/Compiler/PimPipelineSelection.cpp index 2632943..fb4ddb2 100644 --- a/src/PIM/Compiler/PimPipelineSelection.cpp +++ b/src/PIM/Compiler/PimPipelineSelection.cpp @@ -17,6 +17,7 @@ #include "llvm/Support/raw_os_ostream.h" #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" #include +#include using namespace mlir; using namespace onnx_mlir; @@ -87,7 +88,8 @@ struct PimPipelineSelectionPass final : PassWrapper(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 cores(stages); - std::vector 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 cores(depth); + std::vector pressure(depth); + std::vector active(depth); + for (const auto &[instance, core] : schedule.computeToCpuMap) + active[schedule.processorStages.empty() ? 0 : schedule.processorStages[core]] = true; + summary["active_stage_count"] = static_cast(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(cores[stage])); witness.push_back(pressure[stage]); } @@ -159,8 +166,8 @@ struct PimPipelineSelectionPass final : PassWrappercosts); 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(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(search->parents[i])}); + for (size_t i = 0; i < entries->size(); ++i) { + auto &entry = *(*entries)[i].getAsObject(); + entry["parent"] = static_cast(search->parents[i]); + entry["expansions"] = static_cast(search->expansions[i]); + entry["origin"] = i < search->initialSchedules ? "depth-initialization" : "refinement"; + } } report["generated_candidates"] = static_cast(prepared.candidates.size()); report["single_candidate_analysis_skipped"] = single; @@ -212,6 +224,44 @@ struct PimPipelineSelectionPass final : PassWrappertermination; report["queue_evaluation_budget"] = static_cast(kPipelineQueueEvaluationBudget); report["beam_width"] = static_cast(kPipelineBeamWidth); + report["maximum_depth"] = static_cast(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(search->initializationPackingQueries); + report["initialization_dp_transitions"] = static_cast(search->partitionTransitions); + report["initialization_ordered_threshold_scans"] = static_cast(search->orderedThresholdScans); + llvm::json::Array priorities; + for (auto [depth, priority] : search->depthPriorities) + priorities.push_back(llvm::json::Object{{"depth", static_cast(depth)}, + {"priority", static_cast(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 depths; + for (const auto &schedule : prepared.candidates) depths.insert(getPipelineStageCount(schedule)); + for (size_t depth : depths) explored.push_back(static_cast(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(depth)}, + {"evaluated_schedules", static_cast(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(prepared->candidates.size())}, {"selection_margin", kPipelineSelectionMargin}, {"prediction_error_bound_available", false}, {"temporary_disk_bytes", 0}}; + report["requested_depth"] = stages ? llvm::json::Value(static_cast(stages)) + : llvm::json::Value("auto"); if (pimDisableSynchronization) { report["synchronization_ablation"] = true; report["cost_model_program"] = "synchronized-reference"; @@ -308,7 +360,8 @@ struct PimPipelineSelectionPass final : PassWrappercandidates[*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 : PassWrappersetAttrs((*winner)->getAttrs()); checkpoint.getBodyRegion().takeBody(winner->getBodyRegion()); report["selection"] = static_cast(*selected); + report["selected_depth"] = static_cast(getPipelineStageCount(prepared->candidates[*selected])); report["status"] = "selected"; report["reason"] = reason; report["finalization_seconds"] = elapsed(finalize); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp index 8edeca1..0af3f67 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp @@ -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"); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp index 0a644c0..22b1833 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp @@ -84,8 +84,8 @@ private: FailureOr 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 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 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; diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp index d380a0e..c2ef4e9 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp @@ -9,6 +9,7 @@ namespace onnx_mlir::spatial { struct PreparedSpatialSchedules { ComputeGraph graph; std::vector candidates; + std::string latencyInitializationReason; }; // Success always includes at least one candidate, including an empty compute graph. diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.cpp index 3c1f37b..1f395b5 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.cpp @@ -1,5 +1,7 @@ #include "PeriodicDependencies.hpp" +#include "llvm/ADT/STLExtras.h" + #include namespace onnx_mlir::spatial { @@ -7,7 +9,7 @@ namespace onnx_mlir::spatial { mlir::FailureOr analyzePeriodicDependencies( size_t nodeCount, llvm::ArrayRef edges, std::string &error) { - std::vector> outgoing(nodeCount); + std::vector offsets(nodeCount + 1); std::vector indegree(nodeCount), recurrent; Cost upper = 0; for (auto [index, edge] : llvm::enumerate(edges)) { @@ -23,16 +25,26 @@ mlir::FailureOr 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 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 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 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 analyzePeriodicDependencies( if (recurrent.empty() || !upper) return result; + // Preserve the DAG's exact relaxation order without revisiting adjacency lists. + std::vector 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::max(); std::vector potential(nodeCount); std::vector predecessor(nodeCount, absent); @@ -92,9 +112,7 @@ mlir::FailureOr 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 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 diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelinePackingCostCache.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelinePackingCostCache.hpp new file mode 100644 index 0000000..ee7b5b5 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelinePackingCostCache.hpp @@ -0,0 +1,65 @@ +#pragma once + +#include "ResidentWeightPacking.hpp" + +#include +#include +#include +#include +#include + +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 a, llvm::ArrayRef b) const { + return std::lexicographical_compare(a.begin(), a.end(), b.begin(), b.end()); + } + }; + struct Entry { + size_t minimumCores, lastUse; + std::vector> costs; + }; + const IndexedResidentWeights &weights; + llvm::ArrayRef taskCosts; + size_t residentWeightCapacity, processorCount; + std::map, Entry, TaskOrder> entries; + size_t clock = 0; + + Entry &get(llvm::ArrayRef 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(tasks.begin(), tasks.end()), Entry{ + weights.getPackCountLowerBound(tasks, residentWeightCapacity), 0, + std::vector>(processorCount + 1)}).first; + } + found->second.lastUse = ++clock; + return found->second; + } + +public: + static constexpr size_t kMaximumEntries = 256; + PipelinePackingCostCache(const IndexedResidentWeights &weights, llvm::ArrayRef 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 tasks) { return get(tasks).minimumCores; } + Cost estimate(llvm::ArrayRef tasks, size_t cores) { + auto &value = get(tasks).costs[cores]; + if (!value) + value = weights.estimatePackCost(tasks, taskCosts, residentWeightCapacity, cores).value_or(std::numeric_limits::max()); + return *value; + } +}; + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp index 7f3c60a..08cf581 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp @@ -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 taskStages; std::vector stageSizes; std::vector constructivePacks; + std::vector constructiveCores; }; static bool consumesPipelineInput(const ComputeGraphNode &node) { @@ -225,161 +227,302 @@ struct PipelinePartitionCandidate { bool valid = false; }; -static FailureOr 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(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>> best(stages + 1, - std::vector>(count + 1, std::vector(units + 1))); - const Cost infinity = std::numeric_limits::max(); - std::vector> costs(count + 1, - std::vector(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 requestedCore; + std::vector spanAdjustments; + std::vector> costs; + bool sharedDepths; + std::vector> plans; + std::vector> priorities; + std::vector>> best; + static constexpr Cost infinity = std::numeric_limits::max(); + + void buildCosts() { + size_t count = groups.groups.size(); + costs.resize(count + 1, std::vector(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> 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> feasible; - if (!selected) { - feasible.resize(stages + 1, std::vector(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>(count + 1, + std::vector(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> minimumResourceStates(size_t minimumStages) { + size_t count = groups.groups.size(); + std::vector> feasible(stages + 1, std::vector(count + 1)); feasible[0][0].valid = true; std::vector> minimum(count + 1, std::vector(count + 1, std::numeric_limits::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(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(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(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(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 depths = {}) + : groups(groups), model(model), processors(processors), capacity(capacity), + stages(maximumStages), minimumAmount(depths.empty() ? 0 : 1), sharedDepths(!depths.empty()) { + budget = sharedDepths ? processors : std::min(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(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(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(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(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 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 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 priority(size_t depth) const { + if (sharedDepths) return priorities[depth]; + const auto &cell = best[depth][groups.groups.size()][budget]; + return cell.valid ? std::optional(cell.priority) : std::nullopt; + } + + FailureOr 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> 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(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 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 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 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 tasksByPhysicalCpu(physicalTarget.processorCount); std::vector 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 allocatePipelineStageCores( @@ -780,9 +939,11 @@ static std::optional 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 adjustable(stageCount); for (size_t stage : adjustableStages) adjustable[stage] = true; @@ -867,7 +1028,8 @@ static std::optional allocatePipelineStageCores( static void visitAdjacentPipelinePartitions(const PipelineGroupModel &groups, const PipelineStageAssignment &original, - const std::function &visit) { + const std::function &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 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 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 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; +using PipelineMove = std::tuple; static std::vector pipelineNeighborhood(const ComputeGraph &graph, const PipelineTaskModel &baseModel, const PipelineGroupModel &baseGroups, const PipelineSearchResult &search, ArrayRef beam, - const std::set &visited, size_t limit) { + PipelinePackingCostCache &packingCosts, + const std::set &visited, size_t limit, + size_t minimumDepth, size_t maximumDepth) { const auto &target = baseModel.target; std::vector proposals; for (size_t parent : beam) { + if (getPipelineStageCount(search.schedules[parent]) == 1) continue; std::vector 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 pressure(stages); if (!search.costs.empty()) { @@ -1105,26 +1341,19 @@ static std::vector pipelineNeighborhood(const ComputeGraph &gr pressure[search.schedules[parent].processorStages[core]], cost.processorTaskWorkNs[core]); } auto originalView = partitionProxy(graph, groups, model, original); - std::vector>> packCosts(stages, - std::vector>(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::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 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 pipelineNeighborhood(const ComputeGraph &gr })) return; Cost maximum = 0, total = 0; double affected = 0; - std::vector changed(stages); + size_t proposedStages = assignment.stageSizes.size(); + std::vector 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::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 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 active(stages); for (size_t stage : original.taskStages) active[stage] = true; @@ -1201,34 +1430,132 @@ static std::vector 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> 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 &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 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(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 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( + 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 visited; std::set placements{placementIdentity(graph, initial)}; - std::vector beam{0}; + std::set> attemptedAssignments; + for (const auto &schedule : result.schedules) placements.insert(placementIdentity(graph, schedule)); + std::vector 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 proposals; + if (pipelineStages) proposals = pipelineNeighborhood(graph, executionModel, *groups, result, beam, + packingCosts, visited, 2 * evaluationBudget, minimumDepth, result.maximumDepth); result.neighborhoodSeconds = std::chrono::duration( 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 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 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( 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 expandedParents; - std::vector 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( std::chrono::steady_clock::now() - packingStart).count(); @@ -1279,10 +1620,11 @@ mlir::FailureOr 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( std::chrono::steady_clock::now() - packingStart).count(); if (failed(packed)) continue; @@ -1292,24 +1634,15 @@ mlir::FailureOr 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); } diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp index cc4dcbe..a5eb8d6 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp @@ -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 schedules; std::vector costs; std::vector 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> depthPriorities; + size_t maximumDepth = 1; + std::vector> depthInitializationFailures; + size_t initialSchedules = 1; + std::vector 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; mlir::FailureOr optimizePipelineSchedule( diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ResidentWeightPacking.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ResidentWeightPacking.hpp index d5b743b..66bf4e1 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ResidentWeightPacking.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ResidentWeightPacking.hpp @@ -4,6 +4,7 @@ #include "llvm/ADT/SmallBitVector.h" #include "llvm/ADT/STLExtras.h" +#include #include #include #include @@ -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 costs; + const std::vector> &weights; + llvm::ArrayRef tasks; + size_t capacity, generation = 0; + std::vector seen; + Cost largest = 0, total = 0; + bool valid = true, cache; + std::map counts; + size_t scans = 0; + + size_t packCount(Cost maximum, std::vector *assignment = nullptr, + size_t stopAfter = std::numeric_limits::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 costs, + const std::vector> &weights, size_t weightCount, + llvm::ArrayRef 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 estimate(size_t packs, std::vector *assignment = nullptr, + Cost minimum = 0, std::optional 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> curve(size_t maximumPacks, + llvm::ArrayRef requested = {}) { + std::vector> result(maximumPacks + 1); + if (!valid || !maximumPacks) return result; + std::vector 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 estimateOrderedPackCost( llvm::ArrayRef taskCosts, const std::vector> &taskWeightIds, size_t weightCount, llvm::ArrayRef tasks, size_t residentWeightCapacity, size_t maximumPacks, std::vector *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 seen(weightCount); - size_t generation = 0; - auto pack = [&](Cost maximumCost, std::vector *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> estimatePackCostCurve(llvm::ArrayRef tasks, + llvm::ArrayRef costs, size_t capacity, size_t maximumPacks, size_t *thresholdScans = nullptr, + llvm::ArrayRef 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(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> pack( llvm::ArrayRef tasks, llvm::ArrayRef costs, diff --git a/test/PIM/CMakeLists.txt b/test/PIM/CMakeLists.txt index 5a12c1e..4d689e9 100644 --- a/test/PIM/CMakeLists.txt +++ b/test/PIM/CMakeLists.txt @@ -99,6 +99,11 @@ add_test(NAME PimQueueSelectionTest --compiler "$") set_tests_properties(PimQueueSelectionTest PROPERTIES LABELS pim-unittest) +add_test(NAME PimPipelineDepthTest + COMMAND "${PYTHON_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/test_PipelineDepth.py" + --compiler "$") +set_tests_properties(PimPipelineDepthTest PROPERTIES LABELS pim-unittest) + add_test(NAME PimEmissionTest COMMAND "${PYTHON_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/test_PimEmission.py" --compiler "$") diff --git a/test/PIM/PeriodicDependenciesTest.cpp b/test/PIM/PeriodicDependenciesTest.cpp index 4096b68..a6d0f5d 100644 --- a/test/PIM/PeriodicDependenciesTest.cpp +++ b/test/PIM/PeriodicDependenciesTest.cpp @@ -84,6 +84,15 @@ static void testFrozenDifferential() { static_cast(1 + component % 7)}); } differential(componentSize * components, large); + // Sparse adjacency includes empty rows, parallel edges, and shuffled sources. + std::vector 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}}); diff --git a/test/PIM/PimPipelineSelectionTest.cpp b/test/PIM/PimPipelineSelectionTest.cpp index f291c79..7fa7837 100644 --- a/test/PIM/PimPipelineSelectionTest.cpp +++ b/test/PIM/PimPipelineSelectionTest.cpp @@ -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; diff --git a/test/PIM/PipelineOptimizerTest.cpp b/test/PIM/PipelineOptimizerTest.cpp index a21e3c1..0664085 100644 --- a/test/PIM/PipelineOptimizerTest.cpp +++ b/test/PIM/PipelineOptimizerTest.cpp @@ -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 +#include #include #include +#include +#include #include #include @@ -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 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 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::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{0, 1, 2, 3, 4, 5}); + check(std::vector{0, 3, 1, 4, 2, 5}); + check(std::vector{0}); + check(std::vector{}); +} + 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 work; + std::vector 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 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 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(); @@ -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> 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 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(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(9), + TEST_CHECK(mlir::failed(packPipelineSchedule(graph, std::vector(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::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 stages, slots(3); + std::vector 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 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(std::chrono::steady_clock::now() - start).count(); + MergeScheduleResult initial; + std::string error; + TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, std::vector(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::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(std::chrono::steady_clock::now() - start).count() << '\n'; + } +} + +static void testNonunimodalPhysicalDepthCurve() { + TEST_CHECK(std::numeric_limits::is_iec559 && std::numeric_limits::digits == 53); + TEST_CHECK(std::numeric_limits::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 curve(5, std::numeric_limits::infinity()); + auto primary = curve, alternative = curve; + using Ratio = std::pair; + std::vector> certified(5); + auto certifiedPrimary = certified, certifiedAlternative = certified; + auto less = [](Ratio a, Ratio b) { return a.first * b.second < b.first * a.second; }; + std::vector> winnerCores(5), winnerStages(5); + std::vector evaluations(5); + for (size_t cuts = 0; cuts < 16; ++cuts) { + std::vector 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 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::infinity()); + long double high = std::nextafter(value, std::numeric_limits::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({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 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 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::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(); + 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 stages{0, 0, 0, 0, 1}; + double reference = std::numeric_limits::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 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 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 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(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::max(); + size_t evaluated = 0; + for (size_t mask = 0; mask < (size_t{1} << (work.size() - 1)); ++mask) { + std::vector 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> tasks(depth); + for (size_t task = 0; task < work.size(); ++task) tasks[stages[task]].push_back(task); + std::function 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 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::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::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 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 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::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({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); diff --git a/test/PIM/SpatialSchedulingTargetTest.cpp b/test/PIM/SpatialSchedulingTargetTest.cpp index 9423299..f7b5f0c 100644 --- a/test/PIM/SpatialSchedulingTargetTest.cpp +++ b/test/PIM/SpatialSchedulingTargetTest.cpp @@ -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(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(graph.nodes.size()), + {target.processorCount}, target, schedule, error))); std::vector 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(graph.nodes.size()), + {target.processorCount}, target, schedule, error))); assert(schedule.dominanceOrderCompute.size() == graph.nodes.size()); std::vector 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({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) diff --git a/test/PIM/test_HostPublication.py b/test/PIM/test_HostPublication.py index 2c4428a..25e53d5 100644 --- a/test/PIM/test_HostPublication.py +++ b/test/PIM/test_HostPublication.py @@ -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) diff --git a/test/PIM/test_PimEmission.py b/test/PIM/test_PimEmission.py index 41e93fd..c95b0fe 100644 --- a/test/PIM/test_PimEmission.py +++ b/test/PIM/test_PimEmission.py @@ -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() diff --git a/test/PIM/test_PimSpatialSelection.py b/test/PIM/test_PimSpatialSelection.py index 11926e5..fdea04a 100644 --- a/test/PIM/test_PimSpatialSelection.py +++ b/test/PIM/test_PimSpatialSelection.py @@ -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 diff --git a/test/PIM/test_PipelineDepth.py b/test/PIM/test_PipelineDepth.py new file mode 100644 index 0000000..0245828 --- /dev/null +++ b/test/PIM/test_PipelineDepth.py @@ -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() diff --git a/test/PIM/test_QueueSelection.py b/test/PIM/test_QueueSelection.py index 78e7a84..27cb3f4 100644 --- a/test/PIM/test_QueueSelection.py +++ b/test/PIM/test_QueueSelection.py @@ -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()} diff --git a/validation/README.md b/validation/README.md index 2ffbd45..df2e6c2 100644 --- a/validation/README.md +++ b/validation/README.md @@ -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. diff --git a/validation/raptor_validation/raptor.py b/validation/raptor_validation/raptor.py index 8c8e724..3f8e659 100644 --- a/validation/raptor_validation/raptor.py +++ b/validation/raptor_validation/raptor.py @@ -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) diff --git a/validation/tests/test_pim_experiments.py b/validation/tests/test_pim_experiments.py index 80c0731..27394f5 100644 --- a/validation/tests/test_pim_experiments.py +++ b/validation/tests/test_pim_experiments.py @@ -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" diff --git a/validation/tools/pim/ablation/README.md b/validation/tools/pim/ablation/README.md index 0024d93..5407482 100644 --- a/validation/tools/pim/ablation/README.md +++ b/validation/tools/pim/ablation/README.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. diff --git a/validation/tools/pim/ablation/run_ablation.py b/validation/tools/pim/ablation/run_ablation.py index bd53ec2..15cff0b 100755 --- a/validation/tools/pim/ablation/run_ablation.py +++ b/validation/tools/pim/ablation/run_ablation.py @@ -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"], ) ) diff --git a/validation/tools/pim/pimcomp/compare/README.md b/validation/tools/pim/pimcomp/compare/README.md index 4f850a5..34a1e8b 100644 --- a/validation/tools/pim/pimcomp/compare/README.md +++ b/validation/tools/pim/pimcomp/compare/README.md @@ -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. | diff --git a/validation/tools/pim/pimcomp/compare/compare_raptor_pimcomp_model.py b/validation/tools/pim/pimcomp/compare/compare_raptor_pimcomp_model.py index fb510e7..c3086a3 100755 --- a/validation/tools/pim/pimcomp/compare/compare_raptor_pimcomp_model.py +++ b/validation/tools/pim/pimcomp/compare/compare_raptor_pimcomp_model.py @@ -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, diff --git a/validation/tools/pim/pimcomp/compare/run_pimcomp_models.py b/validation/tools/pim/pimcomp/compare/run_pimcomp_models.py index 4d84331..00d60b0 100755 --- a/validation/tools/pim/pimcomp/compare/run_pimcomp_models.py +++ b/validation/tools/pim/pimcomp/compare/run_pimcomp_models.py @@ -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) ) diff --git a/validation/validate.py b/validation/validate.py index 9dbd97c..fa4ced2 100755 --- a/validation/validate.py +++ b/validation/validate.py @@ -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,