From fd493ae6e88197832d5d7685a47229287a391a58 Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Fri, 23 Oct 2026 11:27:06 +0200 Subject: [PATCH] better throughput scheduling with steady state cost model (WIP) --- docs/compiler/options.md | 4 +- docs/compiler/scheduling.md | 303 +++++++- src/PIM/Common/Support/DebugDump.cpp | 8 + src/PIM/Common/Support/DebugDump.hpp | 1 + src/PIM/Common/Support/FileSystemUtils.cpp | 7 + src/PIM/Common/Support/FileSystemUtils.hpp | 9 + src/PIM/Compiler/CMakeLists.txt | 8 + src/PIM/Compiler/PimCodeGen.cpp | 73 +- src/PIM/Compiler/PimCodeGen.hpp | 9 +- src/PIM/Compiler/PimCompilerOptions.cpp | 4 + src/PIM/Compiler/PimCompilerOptions.hpp | 1 + src/PIM/Compiler/PimCompilerUtils.cpp | 135 +++- src/PIM/Compiler/PimCompilerUtils.hpp | 12 + src/PIM/Compiler/PimInstructionTiming.hpp | 39 + src/PIM/Compiler/PimPeriodicAnalysis.cpp | 671 ++++++++++++++++++ src/PIM/Compiler/PimPeriodicAnalysis.hpp | 31 + src/PIM/Compiler/PimPeriodicDiagnostic.cpp | 76 ++ src/PIM/Compiler/PimPipelineSelection.cpp | 225 ++++++ src/PIM/Compiler/PimPipelineSelection.hpp | 9 + src/PIM/Compiler/PimScheduleDiagnostic.cpp | 467 ++++++++++++ src/PIM/Compiler/PimScheduleDiagnostic.hpp | 22 + src/PIM/Dialect/Spatial/CMakeLists.txt | 1 + .../DeferredCommunicationPlanning.cpp | 20 +- .../DeferredCommunicationPlanning.hpp | 3 +- .../DeferredCommunicationRealization.cpp | 5 +- .../DeferredCommunicationRealization.hpp | 4 +- .../ScheduledComputeMaterialization.cpp | 8 +- .../ScheduledComputeReport.cpp | 7 + .../ScheduledComputeReport.hpp | 6 + .../ScheduledSpatialPasses.cpp | 263 ++++--- .../ScheduledSpatialPasses.hpp | 23 + .../Scheduling/ComputeGraph.cpp | 24 +- .../Scheduling/PeriodicDependencies.cpp | 141 ++++ .../Scheduling/PeriodicDependencies.hpp | 34 + .../Scheduling/PipelineScheduling.cpp | 80 ++- .../Scheduling/PipelineScheduling.hpp | 3 +- .../Scheduling/ScheduleEvaluation.cpp | 30 + .../Scheduling/ScheduleEvaluation.hpp | 3 + .../Scheduling/SchedulingTarget.hpp | 73 ++ src/PIM/Passes/PIMPasses.h | 2 +- test/PIM/CMakeLists.txt | 28 + test/PIM/PeriodicDependenciesTest.cpp | 88 +++ test/PIM/PimPipelineSelectionTest.cpp | 80 +++ test/PIM/SpatialSchedulingTargetTest.cpp | 50 ++ test/PIM/StaticSchedulingAudit.c | 60 ++ test/PIM/test_HostPublication.py | 43 +- test/PIM/test_PimEmission.py | 89 +++ test/PIM/test_PimPeriodicDiagnostic.py | 334 +++++++++ test/PIM/test_PimTargetConfig.py | 19 + test/PIM/test_QueueSelection.py | 145 ++++ .../pimcomp/compare/results_comparison.csv | 48 +- 51 files changed, 3581 insertions(+), 247 deletions(-) create mode 100644 src/PIM/Compiler/PimInstructionTiming.hpp create mode 100644 src/PIM/Compiler/PimPeriodicAnalysis.cpp create mode 100644 src/PIM/Compiler/PimPeriodicAnalysis.hpp create mode 100644 src/PIM/Compiler/PimPeriodicDiagnostic.cpp create mode 100644 src/PIM/Compiler/PimPipelineSelection.cpp create mode 100644 src/PIM/Compiler/PimPipelineSelection.hpp create mode 100644 src/PIM/Compiler/PimScheduleDiagnostic.cpp create mode 100644 src/PIM/Compiler/PimScheduleDiagnostic.hpp create mode 100644 src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp create mode 100644 src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.cpp create mode 100644 src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.hpp create mode 100644 test/PIM/PeriodicDependenciesTest.cpp create mode 100644 test/PIM/PimPipelineSelectionTest.cpp create mode 100644 test/PIM/StaticSchedulingAudit.c create mode 100644 test/PIM/test_PimEmission.py create mode 100644 test/PIM/test_PimPeriodicDiagnostic.py create mode 100644 test/PIM/test_QueueSelection.py diff --git a/docs/compiler/options.md b/docs/compiler/options.md index 8e128ba..0e06ca1 100644 --- a/docs/compiler/options.md +++ b/docs/compiler/options.md @@ -11,11 +11,13 @@ 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` | Throughput stages; default: 1 (latency scheduling). | +| `--pipeline=N` | Static queue-aware throughput selection for final codegen with N > 1; intermediate targets retain the original schedule. Default: 1 (latency PEFT). | | `--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. | | `--pim-emit-json` | Also emit per-core JSON instructions. | +| `--pim-dump-ir` | Write diagnostic IR snapshots with elided large constants; default: off. | +| `--pim-periodic-diagnostics` | Write the optional pre-codegen recurrence report; default: off. | | `--pim-export-spatial-dataflow=none\|spatial1\|spatial2\|spatial3\|spatial4\|all` | Export graph, trivially merged, scheduled, and realized CSV snapshots under `reports/`; default: `none`. | | `--pim-conv-lowering=STRATEGY` | `auto` (default), `legacy`, `depthwise`, `packed-im2col`, `streamed-patch`, `streamed-packed`, `output-channel-tiled`, `input-k-tiled`, or `tiled-2d`. | | `--pim-conv-im2col-max-elements=N` | Global im2col limit before streaming; default: 1048576. | diff --git a/docs/compiler/scheduling.md b/docs/compiler/scheduling.md index 135d0e6..28b42eb 100644 --- a/docs/compiler/scheduling.md +++ b/docs/compiler/scheduling.md @@ -1,7 +1,8 @@ # Spatial scheduling -`ScheduleAndRealizeSpatial` constructs a shared `ComputeGraph`, then selects -latency PEFT or pipeline scheduling. Pipeline mode does not run PEFT. +`prepareSpatialSchedules` constructs the shared `ComputeGraph`. Latency mode +uses PEFT. For `--pipeline > 1`, the PIM compiler selects among bounded pipeline +candidates with static queue-aware instruction analysis. Pipeline mode does not run PEFT. See [implementation notes](scheduling-implementation.md) for fragment-view eligibility, receive forwarding, transfer rounding, cost-model boundaries, and recovery limits. @@ -57,19 +58,14 @@ paths exclude core-order constraints; completion includes them. Stage spans exclude cross-stage routing, restart/host synchronization, and contention. Simulator latency and measured throughput remain the performance gates. -Pipeline scheduling retains the original ordered-packing/dependency-tail -heuristic using nominal movement volumes. A second search uses execution -estimates and aggregate serialized shared-memory service, minimizing their -maximum. Its DP keeps at most eight compute/memory trade-offs per state; -four final alternatives and the initial partition undergo the existing core -allocation, topology placement, and capacity verification. The phase-1 winner -remains a fallback. A second bounded search accounts for shared publications. -Selection retains the fallback unless the modeled bottleneck improves. Selection also requires that -shared-memory service plus maximum stage span does not increase. The same -check applies to shared-memory service plus maximum per-core compute work. -These guards limit compute imbalance; their sums do not predict simulator -periods. Sharing-aware alternatives must also preserve maximum stage span, -because placement-dependent baseline loads remain approximate. PEFT is unchanged. +Pipeline candidate generation retains the original ordered-packing/dependency-tail +heuristic and the memory-aware and sharing-aware searches. The partition DP +keeps at most eight compute/memory trade-offs per state. Four final alternatives +and the initial partition undergo core allocation, topology placement, and +capacity verification. The previous compute/memory guards identify candidate +zero, the original fallback. They do not reject legal alternatives before queue +analysis. All candidates preserve candidate zero's non-empty pipeline stages. +PEFT is unchanged. `ComputeGraph` retains physical publication and consumer identities separately from aggregated dependency edges. Cross-stage demands charge one store per @@ -92,6 +88,283 @@ Placement ranks maximum per-core compute/incident-transfer demand, then stage completion and compute load; within 1% demand, stage span breaks the tie. These are placement heuristics, not a steady-state timing model. +## Periodic dependency diagnostics + +`--pim-periodic-diagnostics` writes `reports/pipeline_periodic.json` after deferred +boundary planning. This diagnostic does not change schedule selection or IR. +It consumes the cached compute-instance costs, materialized core order, physical +transfer families, publication store flags, and allocated READY/FREE registers. +Compute runs remain atomic; their lanes and loop iterations are not expanded. +Compact runs require one physical core; multi-core classes use per-lane steps. +One event sequence represents each physical core. The model retains one planned +load per physical consumer transfer and one store per shared publication. +Later load folding or strided transfer splitting can change request granularity. + +Edges express `t(v,k) >= t(u,k-d) + w`. Within-generation edges have distance +zero. Each core's completion-to-entry edge has distance one. Reader-entry FREE +signals therefore release the preceding generation through that edge. Signals +join their exact planned WAIT count. Dissemination barriers, restart permissions, +and the downstream release tree retain their core identities and order. +The diagnostic checks its event-register summary against the emitted program +in the synchronization tests; it does not replace deadlock verification. + +The injected target describes direct-transfer completion and link, local-memory, +shared-memory, and signal timing. The current adapter specifies peer permission +before a source read. SEND completion follows the read and payload enqueue; +RECV completion also requires payload delivery and a destination write. +Unsupported or incomplete target timing produces an `unavailable` report. + +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. 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. + +The reported interval estimate is the maximum of the modeled recurrence, +per-core compute work, and aggregate serialized shared-memory service. Memory +service is never assigned a fixed global request order. Taking the maximum +does not add the same service twice. The within-generation dependency span is +reported separately and is not an initiation interval. + +These are conditional estimates for the abstract execution model, not exact +simulator predictions or guaranteed physical lower bounds. Serial compute costs +can overestimate overlapping work. Baseline memory accesses retain their +aggregate service bound, but their positions are not modeled. Communication +introduced after deferred realization, including final output collection, is +also absent. Instruction issue, local contention, network injection queues, and +memory arbitration remain unmodeled. Missing effects are listed in the report; +`selection_eligible` remains false. This plan-level diagnostic does not rank +production candidates. Production uses the instruction-level queue model below. + +For `V` events, `E` edges, and `R` recurrence edges, construction requires +`O(T + F + S log S)` time and `O(V + E + S)` storage, for existing compute +instances `T`, physical transfers `F`, and synchronization events `S`. +Analysis requires `O(L (R+1) (V+E))` time and `O(V+E)` storage, where `L <= 64`. +Here `R` is at most the number of active physical cores. There is no dense +all-core dependency matrix or inference-generation unrolling. +For `C` analyzed candidates, time scales by `C`; scratch storage can be reused. +The current diagnostic analyzes only the selected schedule (`C = 1`). + +### Post-codegen instruction diagnostic + +`raptor-periodic` provides an explicit, read-only diagnostic over an emitted +binary artifact. It accepts native or oracle-exported artifacts. Use the +oracle-exported artifact when comparing against Pimsim performance results: + +```bash +build_release/Release/bin/raptor-periodic /tmp/model/pimsim_nn \ + --pim-target-config=validation/pimsim_configs/pimcomp/arch-a/throughput_config_1000ms.json \ + --core-count=168 --crossbar-count=96 --crossbar-size=128 > /tmp/periodic.json +``` + +The diagnostic uses physical instruction files and their array-group map. +Each file owns one physical core's work, including all lanes and static loop +iterations already emitted there. It does not expand compiler loops or batches. +It replaces Spatial compute costs rather than adding instruction costs to them. +The lightweight compiler report remains separate and does not invoke this tool. + +The supported execution contract requires ROB capacity one. Non-scalar +instructions therefore retain their serial execution order. Local compute and +blocking memory runs become weighted chunks between communication events. +Every emitted global request contributes its actual size, rounding, and count. +Inputs, constants, split transfers, folded loads, and output collection are +included through their emitted instructions. Their positions relative to +communication remain explicit. The global service bound remains separate; +the model does not impose an order on different cores' global requests. + +Timing queries reuse the injected memory, link, and matrix definitions. +PIM instruction semantics belong to the compiler target adapter. Vector +writeback follows the oracle's actual width rules. Matrix writeback uses the +physical array-group width. Local-memory timing uses its own configured period. +The oracle passes `vector_latency_cycle` directly to a nanosecond wait. +This parameter does not scale with the processor period; no separate vector +clock exists in the supported oracle configuration. +In the current oracle, global loads return zero network payload bytes. +Global stores retain their outgoing payload cost. + +The source switch has one unbounded FIFO for outgoing requests. Enqueue does +not wait for earlier requests. FIFO dispatch waits for their completion. +Direct SEND requires permission delivery before the source read. Its payload +enqueue event coincides with source-read completion. SEND completes after +enqueue, independently of FIFO dispatch or payload delivery. RECV requires delivery and +destination write. Permission responses join both peer readiness and the +receiver's outgoing FIFO. SYNC delivery can precede its local completion. +The model retains WAIT counts, restart events, and cross-generation FIFO state. +It does not replace the compiler's synchronization or deadlock verification. + +`--model-network-queues=false` and `--model-issue=false` support timing +comparisons. The latter omits scalar issue and retirement delays while retaining +SYNC's instruction-defined execution time. Completion notifications can join +the current edge before the clock process runs. The model uses minimum +retirement delays; it does not reproduce SystemC clock-phase ordering. +Scalar issue is charged serially, so these estimates are not guaranteed +physical lower bounds. The base recurrence estimate omits global arbitration, +frontend buffering, and initial bootstrap. Production does not rank candidates +with this recurrence-only estimate. +Reports with at most 128 events include `timing_graph` for focused timing checks. +Its dependencies distinguish payload enqueue from FIFO dispatch and delivery. +The instruction test accepts `--pimsim /path/to/ChipTest` for optional fast-mode +timing comparisons. These cover both hardware profiles and distinct processor, +local-memory, and vector latency settings. + +For `N` emitted instructions, `F` direct transfers, `S` synchronization events, +and `P` active cores, construction takes `O(N + (F+S) log(F+S))` time. +The indexed graph uses `O(F+S+P)` space and at most `2P` recurrence edges. +Global requests are counted during the same scan without per-request graph +nodes. One instruction file is memory-mapped at a time. Maximum-cycle-ratio +analysis uses the same bounded algorithm as the lightweight diagnostic. + +`--model-memory-controller` adds a separate `memory_controller` estimate. +Each global request retains its readiness, transport, service, and completion +dependencies. The controller serves requests in predicted arrival order. +Read and write requests share that FIFO. Zero-size requests also join it. +Source switches remain occupied until transport completion. The model does not +assign a fixed global request order or consume simulator trace order. + +The analysis executes the indexed periodic graph with timestamped tokens. +FIFO tokens retain dependencies across generations; the graph is not copied. +A request arrival reserves the controller after earlier arrivals, and its +completion enables dependent events. Independent cores and switches remain +concurrent. The maximum-cycle-ratio solver and its report remain unchanged. + +The queue estimate compares adjacent two-generation windows with the oracle's +one-percent convergence threshold. It reports the mean interval across both +windows. Per-core queue waits use the corresponding four completed body +generations. The maximum request wait includes bootstrap. Event-index order +resolves simultaneous readiness. A second run reverses that order and reports +`alternative_tie_interval_ns`. These two orders are sensitivity checks, not +bounds over every possible tie order. Clock phases and frontend overlap remain +approximate. The standalone report remains read-only. Production selection uses the same queue model through its C++ API. + +For `R` global requests, the queue mode adds `O(R)` events and dependencies. +For `V` resulting events, `E` dependencies, and generation bound `K`, execution +takes `O(K(E + V log V))` time. Storage is `O(V + E + B)`, where `B` counts +pending tokens. The event heap contains at most one event per node. FIFO token +storage retains pending generations, rather than instruction-stream copies. +The diagnostic limits each event to 64 generations and processed events to +`16(V + E)`. It reports `unavailable` if either bound prevents convergence. +The two tie-order runs reuse the same graph. Production uses this queue execution directly and omits maximum-cycle-ratio analysis. +The standalone frontend retains the recurrence witness. Neither analysis changes latency PEFT. + +## Default throughput selection + +Full PIM code generation performs these steps for `--pipeline > 1`. +`--pim-only-codegen` retains the physical schedule already present in its input. + +The normal compilation workflow is: + +1. Run the ONNX frontend, layout planning, common lowering, and graph preparation once. +2. Generate at most eleven verified physical schedules from the existing searches. +3. Remove duplicate placements, including their core order and stage identities. +4. Clone the prepared module for each candidate using MLIR's lossless `IRMapping`. +5. Remap each schedule's compute identities to the clone. +6. Realize communication and lower the clone through verification and instruction selection. +7. Generate binary instruction records in memory, without weights, exports, or candidate files. +8. Run both static controller-arbitration tie policies on that instruction program. +9. Select an eligible candidate and emit only its normal program. + +Immutable attributes, including constants and weights, remain shared in the MLIR +context. Candidate operations, values, communication plans, and synchronization +state remain independent. Trials never mutate the prepared checkpoint. Candidates +run sequentially. At most one best lowered module survives between trials. +The retained module follows the conservative selection rule throughout evaluation. +A close alternative does not evict candidate zero. Unsupported timing contracts +and synchronization ablation generate only the original schedule and lower it +once. The compiler checks this before collecting alternatives. The frontend is +not repeated. + +Intermediate targets (`--EmitSpatial`, `--EmitPim`, and `--EmitPimBufferized`) +retain the original schedule. They do not perform queue-aware selection or run +transformations beyond the requested stage. Their schedule can therefore differ +from final code generation. Full PIM code generation uses queue-aware selection. + +The PIM compiler owns instruction timing and controller behavior. Spatial owns +legal candidate generation and schedule identity remapping. `PimPeriodicAnalysis` +is the single instruction-level timing implementation. `raptor-periodic` is its +read-only diagnostic frontend. The compiler invokes neither that executable nor +an external selection script. + +`SchedulingTarget::prepareScheduleCost` is an optional whole-schedule cost query. +It accepts a verified target-lowered module and returns an opaque program identity +and a deferred static evaluation. Spatial compares the resulting interval +estimates and applies the target's reported uncertainty. Missing uncertainty +makes an estimate ineligible. The target owns instruction extraction, +protocol timing, contention, and unavailable-result reasons. Preparation and +evaluation do not change the candidate. Equal program identities avoid repeated +queue analysis. An absent query retains the original schedule. +If generation yields one distinct schedule, no comparison is required. The +compiler realizes and lowers it once, without a trial clone, duplicate +instruction generation, or queue analysis. The selection report records this +case separately from unavailable estimates. Final code generation is unchanged. + +PIM transport details reside in `PimInstructionTiming`. The optional plan-level +diagnostic also belongs to the PIM compiler. Spatial supplies its physical plans +through an injected reporting callback; it does not encode the PIM handshake. +The instruction emitter remains the source of register setup, address-dependent +instruction suppression, loop expansion, and physical array-group identity. +Candidate evaluation uses this emitter in memory after local-memory planning. +It does not reconstruct these rules from scalar Spatial operation costs. + +Production projects native fp32 transfer sizes and vector widths to the existing +oracle int8 timing contract. It preserves instruction counts, matrix precision, +vector element counts, synchronization, and transfer identities. This projection +does not change the emitted native program. Unsupported widths, targets without +complete timing, and unsupported instruction semantics produce explicit fallback +reasons. Synchronization ablation retains the original schedule. + +Both arbitration tie runs must converge within the existing generation and event +budgets. Their estimates are sensitivity checks, not formal bounds. Duplicate +instruction programs are analyzed once. The candidate with the lowest larger +estimate must separate from the original schedule. Close alternatives do not +cancel a gain that separates from the original. Comparison uses `[min(ties)/1.021, max(ties)/0.979]`. The fixed 2.1% prediction uncertainty +covers offline validation error with upward rounding. It is a selection margin, +not a fitted instruction latency or a guaranteed error bound. Overlapping ranges +with the original retain candidate zero. Missing original estimates also retain candidate zero. +Unavailable alternatives cannot displace an eligible original. This search finds +the best distinguishable candidate in its bounded set; it does not establish +global optimality. + +`reports/pipeline_selection.json` records candidate identities, estimates, +selection reasons, and generation, lowering, instruction-generation, queue-analysis, +and finalization times. The selected verified lowered module and its captured +schedule report are reused without repeating materialization or communication +realization. Memory reports come from normal final code generation. +With `--enable-timing`, candidate summaries also contain per-pass lowering times +and operation/value counts before bufferization and instruction generation. +Instruction-generation summaries separate preparation from per-core emission. +Explicit Spatial exports, IR snapshots, +and plan-level periodic diagnostics reconstruct the selected Spatial schedule +because those outputs require pre-lowering operations. This optional work does +not change the retained instruction program. Candidate instructions and rejected +IR are released immediately. No temporary compilation directories are created. +`--pim-dump-ir` explicitly enables diagnostic IR snapshots; +these snapshots elide large constants and cannot serve as lossless checkpoints. + +### Static analysis, diagnostics, and external validation + +Production decisions use only the input program, compiler options, candidate +instructions and plans, and configured target resources and timing. The static +model executes deterministic abstract events. It does not execute the generated +program. It reads no benchmark CSV, profiling trace, measured latency cache, or +reference artifact. Identical static inputs produce the same decision regardless +of historical measurements. Throughput deferred-demand ordering uses IR order, +rather than allocation addresses. Latency PEFT and its demand ordering are unchanged. + +Optional diagnostics expose recurrence witnesses, instruction timing, and resource +service. Offline Pimsim fast-mode measurements validate these predictions and +candidate rankings. Simulator results never enter production selection. Timing +corrections must follow verified hardware behavior or configured parameters. + +For `B <= 11` schedules, `T` compute instances, and `P` cores, captured schedules +use `O(B(T+P))` storage. Placement deduplication takes `O(B²T)` time. Each distinct +finalist requires downstream lowering and one instruction scan. Queue execution +uses the bounded complexity stated above. Peak storage includes the prepared +module, one best lowered module, one trial, and one queue graph. It does not scale +with retained complete candidate programs. Maximum-cycle-ratio diagnostics do not +run in the candidate evaluation loop. + ## Capacity and recovery Distinct-weight union / per-core capacity is a lower bound; an oversized task diff --git a/src/PIM/Common/Support/DebugDump.cpp b/src/PIM/Common/Support/DebugDump.cpp index 632a3a9..86cce37 100644 --- a/src/PIM/Common/Support/DebugDump.cpp +++ b/src/PIM/Common/Support/DebugDump.cpp @@ -1,4 +1,5 @@ #include "llvm/Support/raw_os_ostream.h" +#include "llvm/Support/CommandLine.h" #include @@ -6,6 +7,12 @@ #include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp" namespace onnx_mlir { +namespace { +llvm::cl::opt dumpIr("pim-dump-ir", llvm::cl::init(false), + llvm::cl::desc("Write PIM compiler IR snapshots (large constants are elided)")); +} + +bool isPimIrDumpEnabled() { return dumpIr; } std::fstream openDialectDumpFileWithExtension(const std::string& name, llvm::StringRef destination, llvm::StringRef extension) { std::string outputDir = getOutputDir(); @@ -18,6 +25,7 @@ std::fstream openDialectDumpFileWithExtension(const std::string& name, llvm::Str } void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified) { + if (!dumpIr) return; std::fstream file = openDialectDumpFileWithExtension(name, "/dialects", "mlir"); if (!file.is_open()) return; diff --git a/src/PIM/Common/Support/DebugDump.hpp b/src/PIM/Common/Support/DebugDump.hpp index 820fa67..cc3167a 100644 --- a/src/PIM/Common/Support/DebugDump.hpp +++ b/src/PIM/Common/Support/DebugDump.hpp @@ -11,6 +11,7 @@ namespace onnx_mlir { /// Emits a MLIR snapshot under the current compiler output /// directory for pass-level debugging. void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified = false); +bool isPimIrDumpEnabled(); /// Opens a file under the same dialect dump directory used by dumpModule. std::fstream openDialectDumpFileWithExtension(const std::string& name,llvm::StringRef destination = "/dialects", llvm::StringRef extension = "mlir"); diff --git a/src/PIM/Common/Support/FileSystemUtils.cpp b/src/PIM/Common/Support/FileSystemUtils.cpp index 84717ea..3727eb4 100644 --- a/src/PIM/Common/Support/FileSystemUtils.cpp +++ b/src/PIM/Common/Support/FileSystemUtils.cpp @@ -5,7 +5,14 @@ namespace onnx_mlir { +namespace { thread_local bool suppressOutput = false; } +ScopedCompilerOutputSuppression::ScopedCompilerOutputSuppression() : previous(suppressOutput) { + suppressOutput = true; +} +ScopedCompilerOutputSuppression::~ScopedCompilerOutputSuppression() { suppressOutput = previous; } + std::string getOutputDir() { + if (suppressOutput) return {}; if (outputBaseName.empty() || outputBaseName == "-") return {}; diff --git a/src/PIM/Common/Support/FileSystemUtils.hpp b/src/PIM/Common/Support/FileSystemUtils.hpp index 64f76db..13b691f 100644 --- a/src/PIM/Common/Support/FileSystemUtils.hpp +++ b/src/PIM/Common/Support/FileSystemUtils.hpp @@ -8,6 +8,15 @@ namespace onnx_mlir { /// current compiler invocation. std::string getOutputDir(); +// Candidate lowering must not emit reports or temporary artifacts. +class ScopedCompilerOutputSuppression { +public: + ScopedCompilerOutputSuppression(); + ~ScopedCompilerOutputSuppression(); +private: + bool previous; +}; + void createDirectory(const std::string& directory); } // namespace onnx_mlir diff --git a/src/PIM/Compiler/CMakeLists.txt b/src/PIM/Compiler/CMakeLists.txt index 7cf2f2d..c84fe34 100644 --- a/src/PIM/Compiler/CMakeLists.txt +++ b/src/PIM/Compiler/CMakeLists.txt @@ -15,6 +15,9 @@ add_pim_library(OMPimCompilerOptions add_pim_library(OMPimCompilerUtils PimCompilerUtils.cpp + PimScheduleDiagnostic.cpp + PimPeriodicAnalysis.cpp + PimPipelineSelection.cpp PimArtifactWriter.cpp PimCodeGen.cpp PimCoreProgram.cpp @@ -42,3 +45,8 @@ add_pim_library(OMPimCompilerUtils ACCEL_INCLUDE_DIRS PRIVATE ${PIM_ACCEL_INCLUDE_DIRS} ) + +add_onnx_mlir_executable(raptor-periodic NO_INSTALL + PimPeriodicDiagnostic.cpp + LINK_LIBS PRIVATE OMPimCompilerUtils +) diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index ec21da1..0539ac8 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -41,6 +42,7 @@ #include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp" #include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp" #include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp" +#include "src/Accelerators/PIM/Compiler/PimPeriodicAnalysis.hpp" #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp" #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" #include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp" @@ -893,6 +895,7 @@ struct CoreEmissionResult { OnnxMlirCompilerErrorCodes status = CompilerSuccess; MemoryReportRow reportRow; + std::string instructions; llvm::SmallVector usedWeights; llvm::SmallVector diagnostics; size_t diagnosticCount = 0; @@ -1147,8 +1150,10 @@ static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath return CompilerSuccess; } -OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::string& outputDirPath) { - if (!outputDirPath.empty()) { +OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::string& outputDirPath, + PimInstructionProgram *instructions, PimCodegenTiming *timing) { + auto start = std::chrono::steady_clock::now(); + if (!instructions && !outputDirPath.empty()) { if (auto error = sys::fs::create_directory(outputDirPath)) { errs() << "Error creating output directory: " << outputDirPath << ": " << error.message() << '\n'; return InvalidOutputFileAccess; @@ -1164,8 +1169,9 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: memory.hostMem.allocateHost(moduleOp, funcOp); memory.reportHost(); - if (auto err = writeMemoryBinary(moduleOp, funcOp, memory, outputDirPath)) - return err; + if (!instructions) + if (auto err = writeMemoryBinary(moduleOp, funcOp, memory, outputDirPath)) + return err; json::Object xbarsPerArrayGroup; size_t maxCoreId = 0; @@ -1295,15 +1301,22 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: std::error_code errorCode; auto outputCorePath = outputDirPath + "/core_" + std::to_string(job.physicalCoreId) + ".pim"; - raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None); - if (errorCode) { - errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n'; - result.status = InvalidOutputFileAccess; - return result; + llvm::SmallVector binary; + raw_svector_ostream memoryStream(binary); + std::unique_ptr fileStream; + if (!instructions) { + fileStream = std::make_unique(outputCorePath, errorCode, sys::fs::OF_None); + if (errorCode) { + errs() << "Error opening core file: " << outputCorePath << ": " << errorCode.message() << '\n'; + result.status = InvalidOutputFileAccess; + return result; + } } + raw_pwrite_stream &coreBinaryStream = instructions + ? static_cast(memoryStream) : *fileStream; std::unique_ptr coreJsonStream; - if (pimEmitJson.getValue()) { + if (!instructions && pimEmitJson.getValue()) { std::string outputCoreJsonPath = outputDirPath + "/core_" + std::to_string(job.physicalCoreId) + ".json"; errorCode = std::error_code(); coreJsonStream = std::make_unique(outputCoreJsonPath, errorCode); @@ -1319,10 +1332,11 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: PimCodeGen coreCodeGen(jobMemory, instructionWriter, coreJsonStream.get()); auto finalizeInstructions = [&]() { bool succeeded = mlir::succeeded(instructionWriter.finalize()); - coreBinaryStream.close(); - if (coreBinaryStream.has_error()) { - coreBinaryStream.clear_error(); - succeeded = false; + if (fileStream) { + fileStream->close(); + if (fileStream->has_error()) { fileStream->clear_error(); succeeded = false; } + } else { + result.instructions.assign(binary.begin(), binary.end()); } return succeeded; }; @@ -1376,8 +1390,14 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: }; std::vector jobResults(jobs.size()); + auto emissionStart = std::chrono::steady_clock::now(); + if (timing) + timing->preparationSeconds = std::chrono::duration(emissionStart - start).count(); mlir::parallelFor( moduleOp.getContext(), 0, jobs.size(), [&](size_t index) { jobResults[index] = emitJob(jobs[index]); }); + if (timing) + timing->emissionSeconds = std::chrono::duration( + std::chrono::steady_clock::now() - emissionStart).count(); pim::CappedDiagnosticReporter diagnostics; Operation* summaryAnchor = nullptr; @@ -1407,14 +1427,33 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: physicalCoreId < activePhysicalCores.size(); ++physicalCoreId) { if (activePhysicalCores[physicalCoreId]) continue; - if (auto err = - emitEmptyCoreArtifacts(outputDirPath, physicalCoreId)) - return err; + if (!instructions) + if (auto err = emitEmptyCoreArtifacts(outputDirPath, physicalCoreId)) + return err; xbarsPerArrayGroup["core" + std::to_string(physicalCoreId)] = json::Array {}; memory.recordCoreReport(physicalCoreId, MemoryReportRow {}); } + if (instructions) { + instructions->instructions.resize(maxPhysicalCoreId + 1); + instructions->arrayGroups.resize(maxPhysicalCoreId + 1); + for (size_t core = 0; core <= maxPhysicalCoreId; ++core) { + instructions->arrayGroups[core].clear(); + llvm::SmallVector empty; + raw_svector_ostream stream(empty); + pim_binary::writeHeader(stream); + instructions->instructions[core].assign(empty.begin(), empty.end()); + } + for (size_t i = 0; i < jobs.size(); ++i) { + size_t core = jobs[i].physicalCoreId; + instructions->instructions[core] = std::move(jobResults[i].instructions); + for (const auto &weight : jobResults[i].usedWeights) + instructions->arrayGroups[core].push_back(weight.shape[1] / crossbarSize.getValue()); + } + return CompilerSuccess; + } + llvm::SmallVector weightRequests; weightRequests.reserve(jobs.size()); for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) { diff --git a/src/PIM/Compiler/PimCodeGen.hpp b/src/PIM/Compiler/PimCodeGen.hpp index 2186a2f..85bf13a 100644 --- a/src/PIM/Compiler/PimCodeGen.hpp +++ b/src/PIM/Compiler/PimCodeGen.hpp @@ -25,6 +25,7 @@ namespace onnx_mlir { struct CompiledCoreProgram; +struct PimInstructionProgram; class PimInstructionWriter; struct MemEntry { @@ -226,7 +227,13 @@ public: }; -OnnxMlirCompilerErrorCodes compileToPimCode(mlir::ModuleOp& moduleOpRef, std::string& outputDirName); +struct PimCodegenTiming { + double preparationSeconds = 0; + double emissionSeconds = 0; +}; + +OnnxMlirCompilerErrorCodes compileToPimCode(mlir::ModuleOp& moduleOpRef, std::string& outputDirName, + PimInstructionProgram *instructions = nullptr, PimCodegenTiming *timing = nullptr); } // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp index 5650da9..cdcbdb1 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -89,6 +89,10 @@ llvm::cl::opt pimReportConvLowering("pim-report-conv-lowering", llvm::cl::init(true), llvm::cl::cat(OnnxMlirOptions)); +llvm::cl::opt pimPeriodicDiagnostics("pim-periodic-diagnostics", + llvm::cl::desc("Emit the optional pre-codegen recurrence timing report"), + llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions)); + llvm::cl::opt pimEmitJson("pim-emit-json", llvm::cl::desc("Also emit per-core JSON instruction files alongside binary .pim files"), llvm::cl::init(false), diff --git a/src/PIM/Compiler/PimCompilerOptions.hpp b/src/PIM/Compiler/PimCompilerOptions.hpp index 4b02afa..cbb36e2 100644 --- a/src/PIM/Compiler/PimCompilerOptions.hpp +++ b/src/PIM/Compiler/PimCompilerOptions.hpp @@ -56,6 +56,7 @@ extern llvm::cl::opt pimExportSpatialDataflow; extern llvm::cl::opt pimOnlyCodegen; extern llvm::cl::opt pimEmitJson; +extern llvm::cl::opt pimPeriodicDiagnostics; extern llvm::cl::opt pimReportConvLowering; extern llvm::cl::opt pimDetectCommunicationDeadlock; extern llvm::cl::opt pimVerifyBufferizationCopyFreedom; diff --git a/src/PIM/Compiler/PimCompilerUtils.cpp b/src/PIM/Compiler/PimCompilerUtils.cpp index 75a462b..45beba5 100644 --- a/src/PIM/Compiler/PimCompilerUtils.cpp +++ b/src/PIM/Compiler/PimCompilerUtils.cpp @@ -14,6 +14,9 @@ #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" +#include "PimPipelineSelection.hpp" +#include "PimInstructionTiming.hpp" +#include "PimPeriodicAnalysis.hpp" #include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" @@ -174,8 +177,9 @@ std::pair getConfigPair(const llvm::json::Object& object, } void loadPimInterProcessorLatencies( - spatial::SchedulingTarget& target, - const llvm::json::Object& network) { + PimInstructionTiming& target, + const llvm::json::Object& network, + std::optional memorySwitch) { std::optional filename = network.getString("net_config_file_path"); if (!filename) @@ -240,10 +244,29 @@ void loadPimInterProcessorLatencies( ? 0 : (latencySum + static_cast(pairCount) - 1) / static_cast(pairCount); + if (memorySwitch) { + const auto *memoryRow = latencies->getObject(std::to_string(*memorySwitch)); + for (size_t core = 0; core < target.processorCount; ++core) { + const auto *row = latencies->getObject(std::to_string(core)); + auto request = row ? row->getNumber(std::to_string(*memorySwitch)) : std::nullopt; + auto response = memoryRow ? memoryRow->getNumber(std::to_string(core)) : std::nullopt; + if (!request || !response || !std::isfinite(*request) + || !std::isfinite(*response) || *request < 0 || *response < 0 + || *request >= static_cast(std::numeric_limits::max()) + || *response >= static_cast(std::numeric_limits::max())) { + target.sharedMemoryRequestLatencyNs.clear(); + target.sharedMemoryResponseLatencyNs.clear(); + break; + } + target.sharedMemoryRequestLatencyNs.push_back(std::ceil(*request)); + target.sharedMemoryResponseLatencyNs.push_back(std::ceil(*response)); + } + } } -spatial::SchedulingTarget getPimSchedulingTarget() { - spatial::SchedulingTarget target = getDefaultPimSchedulingTarget(); +PimInstructionTiming loadPimInstructionTiming() { + PimInstructionTiming target; + static_cast(target) = getDefaultPimSchedulingTarget(); if (pimTargetConfig.empty()) return target; @@ -298,12 +321,19 @@ spatial::SchedulingTarget getPimSchedulingTarget() { || target.matrixColumns != crossbarSize.getValue()) llvm::report_fatal_error("Pim target config resources do not match --core-count, " "--crossbar-count, and --crossbar-size"); - loadPimInterProcessorLatencies(target, network); + loadPimInterProcessorLatencies(target, network, + core.getInteger("global_memory_switch_id")); target.processorPeriodNs = getConfigCost(core, "period", target.processorPeriodNs); + target.signalIssueNs = target.processorPeriodNs; + target.controlMessageBytes = 1; + target.sharedMemoryReadReturnsPayload = false; + target.peerPermissionThenSourceRead = true; target.localMemoryWidthBytes = getConfigCost(localMemory, "data_width", target.localMemoryWidthBytes); + target.localMemoryPeriodNs = + getConfigCost(localMemory, "period", target.localMemoryPeriodNs); target.localMemoryReadLatencyCycles = getConfigCost(localMemory, "read_latency_cycle", target.localMemoryReadLatencyCycles); target.localMemoryWriteLatencyCycles = @@ -345,6 +375,65 @@ spatial::SchedulingTarget getPimSchedulingTarget() { } // namespace +PimInstructionTiming getPimInstructionTiming() { + return loadPimInstructionTiming(); +} + +spatial::SchedulingTarget getPimSchedulingTarget() { + auto timing = getPimInstructionTiming(); + auto target = static_cast(timing); + configurePimScheduleCosts(target, std::move(timing)); + return target; +} + +llvm::Expected getPimPeriodicTimingContract() { + auto buffer = llvm::MemoryBuffer::getFile(pimTargetConfig); + if (!buffer) return llvm::createStringError("queue selection requires --pim-target-config with instruction timing"); + auto parsed = llvm::json::parse(buffer.get()->getBuffer()); + if (!parsed) return parsed.takeError(); + const auto *root = parsed->getAsObject(); + const auto *chip = root ? root->getObject("chip_config") : nullptr; + const auto *core = chip ? chip->getObject("core_config") : nullptr; + const auto *matrix = core ? core->getObject("matrix_config") : nullptr; + if (!matrix || core->getInteger("rob_size").value_or(1) != 1) + return llvm::createStringError("queue selection requires a target with ROB capacity one and matrix timing"); + auto bits = matrix->getInteger("cell_precision").value_or(2); + if (bits <= 0) return llvm::createStringError("matrix cell precision must be positive"); + return bits; +} + +void addPimLoweringPasses(PassManager &pm, + const spatial::SpatialTargetResources &targetResources, PimEmissionTargetType target, + bool reportProgress) { + auto message = [&](llvm::StringRef text) { + if (reportProgress) pm.addPass(createMessagePass(text.str())); + }; + if (target >= EmitPim) { + pm.addPass(createSpatialToPimPass(targetResources)); + message("Spatial lowered to Pim"); + } + + if (target >= EmitPimBufferized) { + pm.addPass(createPimBufferizationPreparationPass(pimVerifyBufferizationCopyFreedom.getValue())); + pm.addPass(createPimOneShotBufferizationPass()); + pm.addPass(createPimMemoryNormalizationPass()); + pm.addPass(createPimBufferizationVerificationPass()); + message("Pim bufferized"); + } + + if (target >= EmitPimCodegen) { + pm.addPass(mlir::createLowerAffinePass()); + pm.addPass(createPimHostConstantFoldingPass()); + message("Pim host constants folded"); + pm.addPass(createPimInstructionSelectionPass()); + message("Pim instructions selected"); + pm.addPass(createPimLocalMemoryPlanningPass()); + message("Pim local memory planned"); + pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue())); + message("Pim verified"); + } +} + void addPassesPim(OwningOpRef& module, PassManager& pm, EmissionTargetType& emissionTarget, @@ -374,34 +463,22 @@ void addPassesPim(OwningOpRef& module, pm.addPass(createLowerSpatialPlansPass(targetResources, planningOptions, exportStage)); pm.addPass(createTrivialGraphComputeMergePass( schedulingTarget.residentWeightCapacity, exportStage)); - pm.addPass(spatial::createScheduleAndRealizeSpatialPass( - schedulingTarget, exportStage, pipelineStages.getValue())); + if (pipelineStages.getValue() > 1 && pimEmissionTarget >= EmitPimCodegen) { + pm.addPass(createPimPipelineSelectionPass(schedulingTarget, targetResources, + pipelineStages.getValue(), exportStage)); + pm.addPass(createEmitPimCodePass()); + return; + } else { + if (pipelineStages.getValue() > 1) + pm.addPass(createMessagePass("Queue-aware pipeline: intermediate emission retains the original schedule without backend evaluation")); + pm.addPass(spatial::createScheduleAndRealizeSpatialPass( + schedulingTarget, exportStage, pipelineStages.getValue(), pimPeriodicDiagnostics.getValue())); + } pm.addPass(createMessagePass("Onnx lowered to Spatial")); } - if (pimEmissionTarget >= EmitPim) { - pm.addPass(createSpatialToPimPass(targetResources)); - pm.addPass(createMessagePass("Spatial lowered to Pim")); - } - - if (pimEmissionTarget >= EmitPimBufferized) { - pm.addPass(createPimBufferizationPreparationPass(pimVerifyBufferizationCopyFreedom.getValue())); - pm.addPass(createPimOneShotBufferizationPass()); - pm.addPass(createPimMemoryNormalizationPass()); - pm.addPass(createPimBufferizationVerificationPass()); - pm.addPass(createMessagePass("Pim bufferized")); - } - + addPimLoweringPasses(pm, targetResources, pimEmissionTarget); if (pimEmissionTarget >= EmitPimCodegen) { - pm.addPass(mlir::createLowerAffinePass()); - pm.addPass(createPimHostConstantFoldingPass()); - pm.addPass(createMessagePass("Pim host constants folded")); - pm.addPass(createPimInstructionSelectionPass()); - pm.addPass(createMessagePass("Pim instructions selected")); - pm.addPass(createPimLocalMemoryPlanningPass()); - pm.addPass(createMessagePass("Pim local memory planned")); - pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue())); - pm.addPass(createMessagePass("Pim verified")); pm.addPass(createEmitPimCodePass()); pm.addPass(createMessagePass("Pim code emitted")); } diff --git a/src/PIM/Compiler/PimCompilerUtils.hpp b/src/PIM/Compiler/PimCompilerUtils.hpp index d0e2ee2..395c947 100644 --- a/src/PIM/Compiler/PimCompilerUtils.hpp +++ b/src/PIM/Compiler/PimCompilerUtils.hpp @@ -1,6 +1,8 @@ #pragma once +#include "PimCompilerOptions.hpp" #include "mlir/IR/BuiltinOps.h" +#include "llvm/Support/Error.h" #include "mlir/Pass/PassManager.h" #include "mlir/Pass/PassRegistry.h" @@ -8,6 +10,16 @@ namespace onnx_mlir { +namespace spatial { struct SchedulingTarget; struct SpatialTargetResources; } +spatial::SchedulingTarget getPimSchedulingTarget(); +struct PimInstructionTiming; +PimInstructionTiming getPimInstructionTiming(); +llvm::Expected getPimPeriodicTimingContract(); + +void addPimLoweringPasses(mlir::PassManager &pm, + const spatial::SpatialTargetResources &resources, PimEmissionTargetType target, + bool reportProgress = true); + void addPassesPim(mlir::OwningOpRef& module, mlir::PassManager& pm, EmissionTargetType& emissionTarget, diff --git a/src/PIM/Compiler/PimInstructionTiming.hpp b/src/PIM/Compiler/PimInstructionTiming.hpp new file mode 100644 index 0000000..e6c40ad --- /dev/null +++ b/src/PIM/Compiler/PimInstructionTiming.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp" + +namespace onnx_mlir { + +// PIM transport and instruction timing. Spatial consumes only the base costs. +struct PimInstructionTiming : spatial::SchedulingTarget { + bool peerPermissionThenSourceRead = false; + Cost signalIssueNs = 0; + Cost controlMessageBytes = 0; + bool sharedMemoryReadReturnsPayload = true; + std::vector sharedMemoryRequestLatencyNs; + std::vector sharedMemoryResponseLatencyNs; + + bool hasPeriodicTransferTiming() const { + return processorCount && sharedMemoryWidthBytes && controlMessageBytes + && peerPermissionThenSourceRead + && sharedMemoryRequestLatencyNs.size() == processorCount + && sharedMemoryResponseLatencyNs.size() == processorCount + && interProcessorLatencyNs.size() == processorCount * processorCount + && localMemoryWidthBytes && transferWidthBytes; + } + + Cost getLinkServiceNs(Cost bytes, size_t source, size_t destination) const { + Cost flits = bytes / transferWidthBytes + (bytes % transferWidthBytes != 0); + return checkedMultiply(flits, + getInterProcessorLatencyNs(source, destination)); + } + + Cost getSharedMemoryLinkServiceNs(Cost bytes, size_t core, bool response) const { + Cost flits = bytes / transferWidthBytes + (bytes % transferWidthBytes != 0); + return checkedMultiply(flits, response ? sharedMemoryResponseLatencyNs[core] + : sharedMemoryRequestLatencyNs[core]); + } + +}; + +} // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimPeriodicAnalysis.cpp b/src/PIM/Compiler/PimPeriodicAnalysis.cpp new file mode 100644 index 0000000..7caa22a --- /dev/null +++ b/src/PIM/Compiler/PimPeriodicAnalysis.cpp @@ -0,0 +1,671 @@ +#include "PimPeriodicAnalysis.hpp" +#include "PimInstructionTiming.hpp" +#include "PimCodeGen.hpp" +#include "PimCompilerUtils.hpp" +#include "PimCompilerOptions.hpp" +#include "PimScheduleDiagnostic.hpp" +#include "llvm/Support/SHA256.h" +#include "llvm/ADT/StringExtras.h" +#include +#include "PimBinaryFormat.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp" + +#include "llvm/ADT/SmallVector.h" +#include +#include +#include +#include +#include + +using namespace onnx_mlir; +using namespace onnx_mlir::spatial; +using Opcode = pim_binary::Opcode; +namespace { +using Key = std::pair; +struct Event { int64_t core; uint32_t pc; llvm::StringRef kind; }; +struct Send { size_t enqueue, requestArrival, payloadArrival; Cost bytes; }; +struct Receive { size_t complete, permissionArrival, permissionDispatch; Cost bytes; }; +struct Request { size_t begin, arrival; }; +struct Notice { size_t arrival; int64_t source; }; +struct Wait { size_t end; int64_t expected; }; +struct Core { + size_t entry, tail, firstNetwork, lastNetwork; + Cost pending = 0, work = 0, instructions = 0; + std::array service{}; + bool blockingNetworkInChunk = false; +}; +constexpr size_t absent = std::numeric_limits::max(); + +llvm::json::Object estimateQueuedPeriod(const std::vector &events, + const std::vector &edges, + const std::vector> &services, const std::map &cores, + bool reverseTies) { + struct Node { + llvm::SmallVector incoming, outgoing; + size_t completed = 0, readyInputs = 0; + bool busy = false; + }; + struct Pending { + Cost time; + size_t node; + bool arrival; + }; + struct Tokens { + llvm::SmallVector values; + size_t head = 0; + }; + auto later = [reverseTies](const Pending &a, const Pending &b) { + if (a.time != b.time) return a.time > b.time; + return reverseTies ? a.node < b.node : a.node > b.node; + }; + std::vector nodes(events.size()); + std::vector tokens(edges.size()); + std::vector service(events.size(), 0); + std::vector memoryAccess(events.size(), false); + std::vector entries(events.size(), false); + for (auto [node, cost] : services) { + service[node] = cost; + memoryAccess[node] = true; + } + for (const auto &[core, state] : cores) entries[state.entry] = true; + for (auto [index, edge] : llvm::enumerate(edges)) { + nodes[edge.source].outgoing.push_back(index); + nodes[edge.target].incoming.push_back(index); + tokens[index].values.assign(edge.distance, 0); + nodes[edge.target].readyInputs += edge.distance != 0; + } + std::priority_queue, decltype(later)> pending(later); + auto schedule = [&](size_t node) { + auto &state = nodes[node]; + if (state.busy || state.readyInputs != state.incoming.size()) return; + Cost time = 0; + for (size_t edge : state.incoming) { + auto &queue = tokens[edge]; + time = std::max(time, queue.values[queue.head++]); + if (queue.head == queue.values.size()) { + queue.values.clear(); queue.head = 0; --state.readyInputs; + } else if (queue.head * 2 >= queue.values.size()) { + queue.values.erase(queue.values.begin(), queue.values.begin() + queue.head); + queue.head = 0; + } + } + state.busy = true; + pending.push({time - service[node], node, memoryAccess[node]}); + }; + for (size_t node = 0; node < nodes.size(); ++node) schedule(node); + constexpr size_t maximumRounds = 64; + std::array roundEnd{}; + std::array roundCount{}; + std::map> queueWait; + Cost available = 0; + Cost maximumQueueWait = 0; + size_t maximumQueueNode = 0; + size_t processed = 0; + const size_t budget = 16 * (events.size() + edges.size()); + llvm::json::Object result{{"status", "unavailable"}}; + while (!pending.empty() && processed++ < budget) { + auto next = pending.top(); pending.pop(); + auto &node = nodes[next.node]; + if (node.completed >= maximumRounds) { + result["reason"] = "queue diagnostic exceeded its per-event generation bound"; + return result; + } + if (next.arrival) { + Cost begin = std::max(available, next.time); + if (begin - next.time > maximumQueueWait) { + maximumQueueWait = begin - next.time; + maximumQueueNode = next.node; + } + queueWait[events[next.node].core][node.completed] += begin - next.time; + available = begin + service[next.node]; + pending.push({available, next.node, false}); + continue; + } + ++node.completed; + node.busy = false; + if (entries[next.node]) { + size_t round = node.completed - 1; + roundEnd[round] = std::max(roundEnd[round], next.time); + if (++roundCount[round] == cores.size() && round >= 4 && round % 2 == 0) { + Cost first = roundEnd[round - 2] - roundEnd[round - 4]; + Cost second = roundEnd[round] - roundEnd[round - 2]; + Cost maximum = std::max(first, second); + if (maximum && std::abs(static_cast(first) - second) <= maximum * 0.01) { + result["status"] = "estimated"; + result["interval_ns"] = (roundEnd[round] - roundEnd[round - 4]) / 4.0; + result["completed_generations"] = static_cast(round); + result["processed_events"] = static_cast(processed); + result["maximum_queue_wait_ns_including_bootstrap"] = static_cast(maximumQueueWait); + if (maximumQueueWait) { + result["maximum_queue_core"] = events[maximumQueueNode].core; + result["maximum_queue_pc"] = events[maximumQueueNode].pc; + } + llvm::json::Array waits; + for (const auto &[core, values] : queueWait) { + Cost wait = 0; + for (size_t generation = round - 4; generation < round; ++generation) + wait += values[generation]; + waits.push_back(llvm::json::Object{{"core", core}, {"queue_wait_ns", wait / 4.0}}); + } + result["per_core"] = std::move(waits); + result["tie_policy"] = reverseTies ? "reverse event-index order for simultaneous readiness" + : "event-index order for simultaneous readiness"; + return result; + } + } + } + for (size_t edge : node.outgoing) { + if (tokens[edge].values.empty()) ++nodes[edges[edge].target].readyInputs; + tokens[edge].values.push_back(next.time + edges[edge].delayNs); + schedule(edges[edge].target); + } + schedule(next.node); + } + result["reason"] = "queue diagnostic did not converge within its event budget"; + return result; +} + +class Analysis { +public: + PimInstructionTiming target; + PimPeriodicOptions options; + bool queues = true, issue = true, memoryQueue = false; + Cost cellBits; + std::vector events; + std::vector edges; + std::map cores; + SharedMemoryDemand memory; + std::vector> memoryServices; + std::array opcodeService{}, opcodeCounts{}; + std::string error; + + bool readCore(int64_t core, llvm::StringRef data, const llvm::json::Array &groups); + llvm::json::Object queueReport(); + bool connect(); + llvm::json::Object report(); + +private: + std::map> sends; + std::map> receives; + std::map> signals; + std::map waits; + size_t event(int64_t core, uint32_t pc, llvm::StringRef kind) { + events.push_back({core, pc, kind}); + return events.size() - 1; + } + size_t append(int64_t core, uint32_t pc, llvm::StringRef kind, Cost delay = 0) { + size_t next = event(core, pc, kind); + Core &state = cores.at(core); + edges.push_back({state.tail, next, delay, 0}); + state.tail = next; + return next; + } + void flush(int64_t core, uint32_t pc) { + Core &state = cores.at(core); + if (state.pending) { + append(core, pc, "local-work", state.pending); + state.work += state.pending; + state.pending = 0; + } + if (state.blockingNetworkInChunk) + state.lastNetwork = state.tail; + state.blockingNetworkInChunk = false; + } + Request network(int64_t core, uint32_t pc, size_t ready, + Cost latency, Cost hold = 0) { + Core &state = cores.at(core); + size_t begin = event(core, pc, "fifo-dispatch"), arrival = event(core, pc, "delivery"); + size_t end = hold ? event(core, pc, "switch-release") : arrival; + edges.push_back({ready, begin, 0, 0}); + if (queues && state.lastNetwork != absent) + edges.push_back({state.lastNetwork, begin, 0, 0}); + edges.push_back({begin, arrival, latency, 0}); + if (hold) + edges.push_back({arrival, end, hold, 0}); + if (state.firstNetwork == absent) + state.firstNetwork = begin; + state.lastNetwork = end; + return {begin, arrival}; + } + Cost local(Cost bytes, bool write) const { + return target.getLocalMemoryServiceNs(bytes, write); + } + Cost retirement(bool matrix = false) const { + // Completion can join the current clock edge before ClockDomain runs. + // ROB retirement needs another edge; matrix also has AG/commit registers. + return issue ? (matrix ? 3 : 1) * target.processorPeriodNs : 0; + } + bool communication(int64_t core, uint32_t pc, Opcode op, + int32_t peer, int32_t reg, int32_t count, Cost bytes); +}; + +bool Analysis::communication(int64_t core, uint32_t pc, Opcode op, + int32_t peer, int32_t reg, int32_t count, Cost bytes) { + flush(core, pc); + size_t begin = append(core, pc, "communication"); + if (op == Opcode::sync) { + if (peer < 0 || peer >= static_cast(target.processorCount) + || reg < 0 || reg >= static_cast(target.synchronizationRegisterCount)) { + error = "SYNC has an invalid core or event register"; + return false; + } + auto request = network(core, pc, begin, + target.getLinkServiceNs(target.controlMessageBytes, core, peer)); + signals[{peer, reg}].push_back({request.arrival, core}); + append(core, pc, "sync-complete", target.signalIssueNs + retirement()); + cores.at(core).work += target.signalIssueNs + retirement(); + cores.at(core).service[6] += target.signalIssueNs + retirement(); + opcodeService[static_cast(op)] += target.signalIssueNs + retirement(); + } else if (op == Opcode::wait) { + size_t end = append(core, pc, "wait-complete", retirement()); + if (reg < 0 || reg >= static_cast(target.synchronizationRegisterCount) + || count <= 0 || !waits.emplace(Key{core, reg}, Wait{end, count}).second) { + error = "WAIT requires one positive exact count per register per generation"; + return false; + } + cores.at(core).work += retirement(); + cores.at(core).service[6] += retirement(); + opcodeService[static_cast(op)] += retirement(); + } else { + if (peer < 0 || peer >= static_cast(target.processorCount)) { + error = "direct transfer has an invalid peer"; + return false; + } + size_t complete = event(core, pc, op == Opcode::send ? "send-complete" : "receive-complete"); + edges.push_back({begin, complete, 0, 0}); + cores.at(core).tail = complete; + if (op == Opcode::send) { + auto request = network(core, pc, begin, + target.getLinkServiceNs(target.controlMessageBytes, core, peer), 1); + // SwitchSocket::send enqueues immediately in an unbounded FIFO. + size_t enqueue = event(core, pc, "payload-enqueue"); + edges.push_back({enqueue, complete, retirement(), 0}); + auto payload = network(core, pc, enqueue, target.getLinkServiceNs(bytes, core, peer)); + sends[{core, peer}].push_back({enqueue, request.arrival, payload.arrival, bytes}); + } else { + auto permission = network(core, pc, begin, + target.getLinkServiceNs(target.controlMessageBytes, core, peer)); + receives[{peer, core}].push_back({complete, permission.arrival, permission.begin, bytes}); + } + cores.at(core).work += local(bytes, op == Opcode::recv) + retirement(); + cores.at(core).service[2] += local(bytes, op == Opcode::recv) + retirement(); + opcodeService[static_cast(op)] += local(bytes, op == Opcode::recv) + retirement(); + } + return true; +} + +bool Analysis::readCore(int64_t core, llvm::StringRef data, const llvm::json::Array &groups) { + if (data.size() < 12 || data.take_front(4) != "PIMB" + || llvm::support::endian::read32le(data.data() + 4) != 1) { + error = "invalid binary instruction header"; return false; + } + uint32_t count = llvm::support::endian::read32le(data.data() + 8); + if (data.size() != 12ull + 20ull * count) { + error = "binary instruction count does not match file size"; return false; + } + if (!count) return true; + size_t entry = event(core, 0, "body-entry"); + cores.emplace(core, Core{entry, entry, absent, absent}); + Core &state = cores.at(core); + Cost inputBits = 0, outputBits = 0; + for (uint32_t pc = 0; pc < count; ++pc) { + const char *record = data.data() + 12 + 20ull * pc; + unsigned rawOpcode = static_cast(record[0]); + if (rawOpcode > static_cast(Opcode::sync)) { error = "unsupported instruction opcode"; return false; } + Opcode op = static_cast(rawOpcode); + ++opcodeCounts[rawOpcode]; + int32_t arg = llvm::support::endian::read32le(record + 4); + int32_t a = llvm::support::endian::read32le(record + 8); + int32_t b = llvm::support::endian::read32le(record + 12); + int32_t length = llvm::support::endian::read32le(record + 16); + ++state.instructions; + if (op <= Opcode::setbw) { + if (op == Opcode::sld) { error = "scalar memory loads are unsupported by this diagnostic"; return false; } + if (op == Opcode::setbw) { + if (a <= 0 || b <= 0) { error = "invalid SETBW"; return false; } + if (options.oraclePrecision && (a != 32 || b != 32)) { + error = "oracle timing projection requires 32-bit native vector widths"; return false; + } + inputBits = options.oraclePrecision ? 8 : a; + outputBits = options.oraclePrecision ? 8 : b; + } + // shortcut: serialize scalar issue; model frontend buffering before selection. + Cost cost = issue ? target.processorPeriodNs : 0; + state.pending += cost; state.service[5] += cost; + opcodeService[static_cast(op)] += cost; + continue; + } + if (length < 0 || op > Opcode::sync || op == Opcode::vsoftmax) { + error = "unsupported instruction opcode or negative transfer length"; return false; + } + Cost bytes = length; + if (options.oraclePrecision && (op == Opcode::ld || op == Opcode::st + || op == Opcode::lldi || op == Opcode::lmv || op == Opcode::send || op == Opcode::recv)) { + if (length % 4) { error = "oracle timing projection requires element-aligned transfers"; return false; } + bytes /= 4; + } + Cost in = (inputBits + 7) / 8, out = (outputBits + 7) / 8; + if (op <= Opcode::vrsl && (!inputBits || !outputBits)) { + error = "compute instruction precedes SETBW"; return false; + } + if (op >= Opcode::send) { + if (!communication(core, pc, op, arg, a, b, bytes)) return false; + continue; + } + if (op == Opcode::ld || op == Opcode::st) { + bool write = op == Opcode::st; + auto demand = target.getSharedMemoryDemand(bytes, 1, write); + memory += demand; + Cost completion = (write ? 0 : local(bytes, true)) + retirement(); + Cost cost = (write ? local(bytes, false) : 0) + demand.serviceNs + completion + + target.getSharedMemoryLinkServiceNs(write ? bytes : target.controlMessageBytes, core, false) + + target.getSharedMemoryLinkServiceNs(!write && target.sharedMemoryReadReturnsPayload ? bytes : 0, core, true); + state.service[write ? 4 : 3] += cost; + opcodeService[static_cast(op)] += cost; + if (memoryQueue) { + flush(core, pc); + if (write) append(core, pc, "store-source-read", local(bytes, false)); + auto request = network(core, pc, state.tail, + target.getSharedMemoryLinkServiceNs(write ? bytes : target.controlMessageBytes, core, false)); + size_t done = event(core, pc, "memory-service-complete"); + edges.push_back({request.arrival, done, demand.serviceNs, 0}); + memoryServices.emplace_back(done, demand.serviceNs); + Cost reply = target.getSharedMemoryLinkServiceNs( + !write && target.sharedMemoryReadReturnsPayload ? bytes : 0, core, true); + size_t response = done; + if (reply) { + response = event(core, pc, "memory-response"); + edges.push_back({done, response, reply, 0}); + } + state.lastNetwork = response; + size_t complete = event(core, pc, "memory-complete"); + edges.push_back({response, complete, completion, 0}); + state.tail = complete; + state.work += cost; + continue; + } + if (!state.blockingNetworkInChunk) { + flush(core, pc); + if (write) { + Cost read = local(bytes, false); + append(core, pc, "store-source-read", read); + state.work += read; + cost -= read; + } + size_t begin = append(core, pc, "memory-queue-join"); + if (queues && state.lastNetwork != absent) + edges.push_back({state.lastNetwork, begin, 0, 0}); + if (state.firstNetwork == absent) + state.firstNetwork = begin; + state.blockingNetworkInChunk = true; + } + // ROB=1 serializes these blocking requests. No external dependency joins + // inside a chunk; retain their exact service sum without per-request nodes. + state.pending += cost; + continue; + } + Cost cost = 0; + unsigned resource = 2; + if (op == Opcode::mvmul) { + if (arg <= 0 || length < 0 || static_cast(length) >= groups.size() + || !groups[length].getAsInteger() || *groups[length].getAsInteger() <= 0) { + error = "matrix instruction references an invalid array group"; return false; + } + Cost output = (target.matrixColumns * cellBits / arg) * *groups[length].getAsInteger() * out; + cost = local(target.matrixRows * in * target.matrixInputReadMultiplier, false) + + target.getMatrixComputeLatencyNs(inputBits) + local(output, true); + cost += retirement(true); resource = 0; + } else if (op <= Opcode::vrsl) { + bool binary = op >= Opcode::vvadd && op <= Opcode::vvsra; + bool outputWidth = op == Opcode::vvmul || op == Opcode::vvsll || op == Opcode::vvsra || op == Opcode::vtanh || op == Opcode::vsigm || op == Opcode::vrsu || op == Opcode::vrsl; + // The oracle's VectorUnit uses input width for VAVG/VVDMUL writeback. + // It waits vector_latency_cycle nanoseconds, independently of core period. + cost = local(bytes * in, false) * (binary ? 2 : 1) + + ((bytes + target.vectorWidth - 1) / target.vectorWidth) * target.vectorLatencyCycles + + local(bytes * (outputWidth ? out : in), true) + retirement(); + resource = 1; + } else { + cost = local(bytes, true) + (op == Opcode::lmv ? local(bytes, false) : 0) + retirement(); + } + state.pending += cost; + state.service[resource] += cost; + opcodeService[static_cast(op)] += cost; + } + flush(core, count); + edges.push_back({state.tail, state.entry, 0, 1}); + if (queues && state.firstNetwork != absent) + edges.push_back({state.lastNetwork, state.firstNetwork, 0, 1}); + return true; +} + +bool Analysis::connect() { + if (sends.size() != receives.size() || signals.size() != waits.size()) { + error = "communication endpoint sets differ"; return false; + } + for (const auto &[key, producers] : sends) { + auto found = receives.find(key); + if (found == receives.end() || found->second.size() != producers.size()) { + error = "direct SEND/RECV FIFO counts differ"; return false; + } + for (auto [send, receive] : llvm::zip(producers, found->second)) { + if (send.bytes != receive.bytes) { error = "direct payload sizes differ"; return false; } + // The receiver responds after both its RECV and the sender request arrive. + // Its outgoing FIFO may still contain earlier SYNC or transport requests. + edges.push_back({send.requestArrival, receive.permissionDispatch, 0, 0}); + edges.push_back({receive.permissionArrival, send.enqueue, local(send.bytes, false), 0}); + edges.push_back({send.payloadArrival, receive.complete, local(send.bytes, true) + retirement(), 0}); + } + } + for (const auto &[key, pending] : waits) { + auto found = signals.find(key); + if (found == signals.end() || found->second.size() != static_cast(pending.expected)) { + error = "event register does not receive its exact static count"; return false; + } + for (const auto ¬ice : found->second) + edges.push_back({notice.arrival, pending.end, retirement(), 0}); + } + return true; +} + +llvm::json::Object Analysis::report() { + llvm::json::Object result; + if (cores.empty()) { + error = "artifact has no active repeating instruction stream"; + result["status"] = "unavailable"; result["reason"] = error; return result; + } + if (!connect()) { result["status"] = "unavailable"; result["reason"] = error; return result; } + result["global_memory_service_ns"] = static_cast(memory.serviceNs); + result["global_read_bytes"] = static_cast(memory.readBytes); + result["global_write_bytes"] = static_cast(memory.writeBytes); + result["global_read_requests"] = static_cast(memory.readRequests); + result["global_write_requests"] = static_cast(memory.writeRequests); + if (!options.detailed) { + result["status"] = "estimated"; + result["memory_controller"] = queueReport(); + return result; + } + auto estimate = analyzePeriodicDependencies(events.size(), edges, error); + if (mlir::failed(estimate)) { result["status"] = "unavailable"; result["reason"] = error; return result; } + result["status"] = "estimated"; + result["recurrence_lower_ns"] = estimate->recurrenceLowerNs; + result["recurrence_upper_ns"] = estimate->recurrenceUpperNs; + result["within_generation_span_ns"] = static_cast(estimate->withinGenerationSpanNs); + result["interval_estimate_ns"] = std::max(estimate->recurrenceUpperNs, static_cast(memory.serviceNs)); + result["limiting_bound"] = memory.serviceNs > estimate->recurrenceUpperNs ? "global-memory" : "recurrence"; + result["events"] = static_cast(events.size()); + result["edges"] = static_cast(edges.size()); + // Keep small timing fixtures inspectable without dumping network-size graphs. + if (events.size() <= 128) { + llvm::json::Array nodes, dependencies; + for (const auto &node : events) + nodes.push_back(llvm::json::Object{{"core", node.core}, {"pc", node.pc}, {"kind", node.kind}}); + for (const auto &edge : edges) + dependencies.push_back(llvm::json::Object{{"source", static_cast(edge.source)}, + {"target", static_cast(edge.target)}, {"delay_ns", static_cast(edge.delayNs)}, + {"distance", edge.distance}}); + result["timing_graph"] = llvm::json::Object{{"events", std::move(nodes)}, {"dependencies", std::move(dependencies)}}; + } + result["network_queues"] = queues; result["instruction_issue"] = issue; + if (memoryQueue) result["memory_controller"] = queueReport(); + result["resource_service_columns"] = llvm::json::Array{"matrix", "vector", "local_transfer", + "global_load", "global_store", "scalar_issue", "synchronization"}; + Cost directBytes = 0, directRequests = 0; + for (const auto &[key, transfers] : sends) + for (const auto &transfer : transfers) { directBytes += transfer.bytes; ++directRequests; } + result["direct_requests"] = static_cast(directRequests); + result["direct_bytes"] = static_cast(directBytes); + llvm::json::Object opcodeDemand, instructionCounts; + for (size_t op = 0; op < opcodeCounts.size(); ++op) + if (opcodeCounts[op]) { + opcodeDemand[pim_binary::kOpcodeNames[op]] = static_cast(opcodeService[static_cast(op)]); + instructionCounts[pim_binary::kOpcodeNames[op]] = static_cast(opcodeCounts[op]); + } + result["opcode_service_ns"] = std::move(opcodeDemand); + result["instruction_counts"] = std::move(instructionCounts); + llvm::json::Array perCore, cycle; + Cost maximumCoreWork = 0; + for (const auto &[core, state] : cores) { + maximumCoreWork = std::max(maximumCoreWork, state.work); + llvm::json::Array resource; + for (Cost value : state.service) resource.push_back(static_cast(value)); + perCore.push_back(llvm::json::Object{{"core", core}, {"instructions", static_cast(state.instructions)}, + {"serial_work_ns", static_cast(state.work)}, {"resource_service_ns", std::move(resource)}}); + } + std::map cycleCores; + Cost cycleDelay = 0, cycleDistance = 0; + for (size_t index : estimate->criticalCycle) { + const auto &edge = edges[index]; + const auto &from = events[edge.source], &to = events[edge.target]; + cycleCores[to.core] += edge.delayNs; + cycleDelay += edge.delayNs; cycleDistance += edge.distance; + if (cycle.size() < 128) + cycle.push_back(llvm::json::Object{{"source_core", from.core}, {"target_core", to.core}, + {"source_pc", from.pc}, {"target_pc", to.pc}, {"kind", to.kind}, + {"delay_ns", static_cast(edge.delayNs)}, {"distance", edge.distance}}); + } + llvm::json::Array critical; + for (auto [core, delay] : cycleCores) + critical.push_back(llvm::json::Object{{"core", core}, {"delay_ns", static_cast(delay)}}); + result["critical_cycle"] = std::move(cycle); result["critical_cores"] = std::move(critical); + result["critical_cycle_edges"] = static_cast(estimate->criticalCycle.size()); + result["critical_cycle_delay_ns"] = static_cast(cycleDelay); + result["critical_cycle_distance"] = static_cast(cycleDistance); + result["per_core"] = std::move(perCore); + result["maximum_core_serial_work_ns"] = static_cast(maximumCoreWork); + result["omitted_effects"] = llvm::json::Array{"global-memory arbitration", "clock-phase rounding and frontend buffering", + "scalar issue overlap (charged serially)", "initial inference bootstrap"}; + return result; +} + +llvm::json::Object Analysis::queueReport() { + auto queued = estimateQueuedPeriod(events, edges, memoryServices, cores, false); + auto alternative = estimateQueuedPeriod(events, edges, memoryServices, cores, true); + queued["alternative_tie_status"] = alternative.getString("status").value(); + if (auto reason = alternative.getString("reason")) + queued["alternative_tie_reason"] = reason->str(); + if (auto interval = alternative.getNumber("interval_ns")) + queued["alternative_tie_interval_ns"] = *interval; + if (auto interval = queued.getNumber("interval_ns")) + queued["memory_service_fraction"] = memory.serviceNs / *interval; + queued["arbitration"] = "FIFO by predicted controller arrival"; + return queued; +} + +} // namespace + +llvm::json::Object onnx_mlir::analyzePimInstructions(const PimInstructionProgram &program, + const PimInstructionTiming &target, int64_t cellBits, PimPeriodicOptions options) { + Analysis analysis; + analysis.target = target; + analysis.cellBits = cellBits; + analysis.options = options; + analysis.queues = options.networkQueues; + analysis.issue = options.instructionIssue; + analysis.memoryQueue = options.memoryController; + auto unavailable = [](llvm::StringRef reason) { + return llvm::json::Object{{"status", "unavailable"}, {"reason", reason.str()}}; + }; + if (!target.hasPeriodicTransferTiming() || cellBits <= 0) + return unavailable("target has no supported instruction timing contract"); + if (program.instructions.empty() || program.instructions.size() > target.processorCount + || program.instructions.size() != program.arrayGroups.size()) + return unavailable("instruction program core count or array groups do not fit the target"); + for (size_t id = 0; id < program.instructions.size(); ++id) { + llvm::json::Array groups; + for (int64_t count : program.arrayGroups[id]) groups.push_back(count); + if (!analysis.readCore(id, program.instructions[id], groups)) + return unavailable(analysis.error); + } + return analysis.report(); +} + +void onnx_mlir::configurePimScheduleCosts(spatial::SchedulingTarget &target, + PimInstructionTiming timing) { + auto hardware = std::make_shared(std::move(timing)); + target.reportPeriodicSchedule = [hardware](const spatial::ComputeGraph &graph, + const spatial::ScheduledComputeMaterializationResult &materialization, + const spatial::DeferredTransferPlan &transfers, + const spatial::DeferredBoundaryPlan &boundaries) { + dumpPeriodicScheduleEstimate(graph, materialization, transfers, boundaries, *hardware); + }; + auto contract = getPimPeriodicTimingContract(); + if (!contract) target.wholeScheduleCostUnavailable = llvm::toString(contract.takeError()); + if (!hardware->hasPeriodicTransferTiming()) + target.wholeScheduleCostUnavailable = "target has no supported instruction timing contract"; + if (pimDisableSynchronization) + target.wholeScheduleCostUnavailable = "synchronization ablation has no periodic queue timing contract"; + if (!target.wholeScheduleCostUnavailable.empty()) return; + target.prepareScheduleCost = [hardware, cellBits = *contract](mlir::ModuleOp module, + llvm::json::Object &summary) -> mlir::FailureOr { + auto start = std::chrono::steady_clock::now(); + PimInstructionProgram program; + PimCodegenTiming timing; + std::string unused; + if (compileToPimCode(module, unused, &program, &timing) != CompilerSuccess) + return mlir::failure(); + summary["codegen_preparation_seconds"] = timing.preparationSeconds; + summary["codegen_emission_seconds"] = timing.emissionSeconds; + summary["instruction_generation_seconds"] = std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + llvm::SHA256 hash; + int64_t instructionBytes = 0; + for (size_t core = 0; core < program.instructions.size(); ++core) { + instructionBytes += program.instructions[core].size(); + hash.update(program.instructions[core]); + for (int64_t group : program.arrayGroups[core]) hash.update(std::to_string(group) + ","); + hash.update(";"); + } + summary["instruction_bytes"] = instructionBytes; + summary["instruction_count"] = (instructionBytes + - pim_binary::kHeaderSize * program.instructions.size()) / pim_binary::kRecordSize; + spatial::ScheduleCostQuery query; + query.identity = llvm::toHex(hash.final()); + query.evaluate = [program = std::move(program), hardware, cellBits](llvm::json::Object &summary) { + PimPeriodicOptions options; + options.oraclePrecision = true; + auto report = analyzePimInstructions(program, *hardware, cellBits, options); + for (const char *field : {"global_memory_service_ns", "global_read_bytes", "global_write_bytes", + "global_read_requests", "global_write_requests"}) + if (auto value = report.getInteger(field)) summary[field] = *value; + spatial::WholeScheduleCost cost; + // Conservative selection margin; not a timing correction or formal bound. + cost.relativeUncertainty = 0.021; + if (auto *queue = report.getObject("memory_controller")) { + if (queue->getString("status") == "estimated" + && queue->getString("alternative_tie_status") == "estimated") { + cost.primary = queue->getNumber("interval_ns"); + cost.alternative = queue->getNumber("alternative_tie_interval_ns"); + } + if (auto reason = queue->getString("reason")) cost.reason = reason->str(); + if (!cost.eligible() && cost.reason.empty()) + if (auto reason = queue->getString("alternative_tie_reason")) + cost.reason = "alternative arbitration: " + reason->str(); + } else if (auto reason = report.getString("reason")) cost.reason = reason->str(); + if (!cost.eligible() && cost.reason.empty()) + cost.reason = "arbitration sensitivity estimate unavailable or nonconverged"; + return cost; + }; + return query; + }; +} diff --git a/src/PIM/Compiler/PimPeriodicAnalysis.hpp b/src/PIM/Compiler/PimPeriodicAnalysis.hpp new file mode 100644 index 0000000..118fcbd --- /dev/null +++ b/src/PIM/Compiler/PimPeriodicAnalysis.hpp @@ -0,0 +1,31 @@ +#pragma once + +#include "llvm/Support/JSON.h" +#include +#include +#include + +namespace onnx_mlir { +namespace spatial { struct SchedulingTarget; } +struct PimInstructionTiming; +void configurePimScheduleCosts(spatial::SchedulingTarget &target, PimInstructionTiming timing); + +// Binary instructions, indexed by physical core. No executable or measured data. +struct PimInstructionProgram { + std::vector instructions; + std::vector> arrayGroups; +}; + +struct PimPeriodicOptions { + bool networkQueues = true; + bool instructionIssue = true; + bool memoryController = true; + bool detailed = false; + // Apply the oracle's fp32-to-int8 timing projection without changing the program. + bool oraclePrecision = false; +}; + +llvm::json::Object analyzePimInstructions(const PimInstructionProgram &program, + const PimInstructionTiming &target, int64_t cellBits, + PimPeriodicOptions options = {}); +} // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimPeriodicDiagnostic.cpp b/src/PIM/Compiler/PimPeriodicDiagnostic.cpp new file mode 100644 index 0000000..d2d73c9 --- /dev/null +++ b/src/PIM/Compiler/PimPeriodicDiagnostic.cpp @@ -0,0 +1,76 @@ +#include "PimPeriodicAnalysis.hpp" +#include "PimInstructionTiming.hpp" +#include "PimCompilerOptions.hpp" +#include "PimCompilerUtils.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/MemoryBuffer.h" +#include + +using namespace onnx_mlir; +namespace { +llvm::cl::opt artifact(llvm::cl::Positional, llvm::cl::Required, + llvm::cl::desc("PIM binary artifact directory (native or oracle-exported)")); +llvm::cl::opt queues("model-network-queues", llvm::cl::init(true)); +llvm::cl::opt issue("model-issue", llvm::cl::init(true)); +llvm::cl::opt oraclePrecision("oracle-precision", llvm::cl::init(false), + llvm::cl::desc("Project native fp32 instructions to oracle int8 timing")); +llvm::cl::opt memoryQueue("model-memory-controller", llvm::cl::init(false)); +llvm::Expected readJson(llvm::StringRef path) { + auto buffer = llvm::MemoryBuffer::getFile(path); + if (!buffer) return llvm::errorCodeToError(buffer.getError()); + auto parsed = llvm::json::parse(buffer.get()->getBuffer()); + if (!parsed) return parsed.takeError(); + if (!parsed->getAsObject()) return llvm::createStringError("expected a JSON object"); + return std::move(*parsed->getAsObject()); +} +} // namespace + +int main(int argc, char **argv) { + llvm::cl::ParseCommandLineOptions(argc, argv, "Post-codegen periodic PIM diagnostics\n"); + auto start = std::chrono::steady_clock::now(); + auto program = readJson(artifact + "/config.json"); + auto cellBits = getPimPeriodicTimingContract(); + if (!program || !cellBits) { + llvm::errs() << (!program ? llvm::toString(program.takeError()) : llvm::toString(cellBits.takeError())) << '\n'; + return 1; + } + auto target = getPimInstructionTiming(); + PimInstructionProgram instructions; + auto *groups = program->getObject("array_group_map"); + if (!groups) { llvm::errs() << "missing physical array groups\n"; return 1; } + auto coreCount = program->getInteger("core_cnt"); + if (!coreCount || *coreCount <= 0 || *coreCount > static_cast(target.processorCount)) { + llvm::errs() << "artifact core count does not fit the target\n"; return 1; + } + for (int64_t id = 0; id < *coreCount; ++id) { + auto path = artifact + "/core_" + std::to_string(id) + ".pim"; + auto *array = groups->getArray("core" + std::to_string(id)); + if (!array) { llvm::errs() << "missing physical core array groups\n"; return 1; } + auto buffer = llvm::MemoryBuffer::getFile(path); + if (!buffer) { llvm::errs() << "cannot read instruction artifact\n"; return 1; } + instructions.instructions.push_back(buffer.get()->getBuffer().str()); + std::vector counts; + for (const auto &value : *array) { + auto count = value.getAsInteger(); + if (!count) { llvm::errs() << "invalid array group\n"; return 1; } + counts.push_back(*count); + } + instructions.arrayGroups.push_back(std::move(counts)); + } + PimPeriodicOptions options; + options.detailed = true; + options.oraclePrecision = oraclePrecision; + options.networkQueues = queues; + options.instructionIssue = issue; + options.memoryController = memoryQueue; + auto report = analyzePimInstructions(instructions, target, *cellBits, options); + report["diagnostic_only"] = true; + report["selection_eligible"] = false; + bool available = report.getString("status") == "estimated"; + report["analysis_wall_ms"] = std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + llvm::outs() << llvm::formatv("{0:2}\n", llvm::json::Value(std::move(report))); + return available ? 0 : 1; +} diff --git a/src/PIM/Compiler/PimPipelineSelection.cpp b/src/PIM/Compiler/PimPipelineSelection.cpp new file mode 100644 index 0000000..111651e --- /dev/null +++ b/src/PIM/Compiler/PimPipelineSelection.cpp @@ -0,0 +1,225 @@ +#include "PimPipelineSelection.hpp" +#include "PimCompilerOptions.hpp" +#include "PimCompilerUtils.hpp" +#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" +#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp" +#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp" +#include "mlir/Pass/PassManager.h" +#include "mlir/Pass/PassInstrumentation.h" +#include "src/Compiler/CompilerOptions.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/JSON.h" +#include "llvm/Support/raw_os_ostream.h" +#include +#include + +using namespace mlir; +using namespace onnx_mlir; +using namespace onnx_mlir::spatial; + +namespace { +using Clock = std::chrono::steady_clock; +double elapsed(Clock::time_point start) { + return std::chrono::duration(Clock::now() - start).count(); +} + +llvm::json::Object countIR(Operation *root) { + int64_t operations = 0, values = 0; + root->walk([&](Operation *op) { + ++operations; + values += op->getNumResults(); + for (Region ®ion : op->getRegions()) + for (Block &block : region) values += block.getNumArguments(); + }); + return llvm::json::Object{{"operations", operations}, {"values", values}}; +} + +// Enabled with the existing compiler timing option. No instruction expansion. +struct CandidateTiming final : PassInstrumentation { + llvm::json::Object &summary; + Clock::time_point start; + explicit CandidateTiming(llvm::json::Object &summary) : summary(summary) {} + void runBeforePass(Pass *pass, Operation *op) override { + if (pass->getArgument() == "pim-one-shot-bufferization") + summary["bufferization_input_ir"] = countIR(op); + start = Clock::now(); + } + void runAfterPass(Pass *pass, Operation *) override { + auto *phases = summary.getObject("lowering_phases_seconds"); + if (!phases) { + summary["lowering_phases_seconds"] = llvm::json::Object{}; + phases = summary.getObject("lowering_phases_seconds"); + } + (*phases)[pass->getName()] = elapsed(start); + } +}; + +struct PimPipelineSelectionPass final : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimPipelineSelectionPass) + PimPipelineSelectionPass(const SchedulingTarget &target, const SpatialTargetResources &resources, + size_t stages, SpatialDataflowExportStage exportStage) + : target(target), resources(resources), stages(stages), exportStage(exportStage) {} + SchedulingTarget target; + SpatialTargetResources resources; + size_t stages; + SpatialDataflowExportStage exportStage; + + StringRef getArgument() const override { return "pim-pipeline-selection"; } + StringRef getDescription() const override { return "Select a throughput pipeline using static instruction queue analysis"; } + + LogicalResult lower(ModuleOp module, llvm::json::Object *summary = nullptr) { + PassManager pm(module.getContext()); + addPimLoweringPasses(pm, resources, EmitPimCodegen, /*reportProgress=*/false); + if (enableTiming && summary) + pm.addInstrumentation(std::make_unique(*summary)); + return pm.run(module); + } + + LogicalResult evaluate(ModuleOp checkpoint, const MergeScheduleResult &schedule, + OwningOpRef &candidate, + std::string &scheduleReport, llvm::json::Object &summary) { + auto begin = Clock::now(); + IRMapping mapping; + candidate = cast(checkpoint->clone(mapping)); + auto remapped = remapSpatialSchedule(schedule, mapping); + if (failed(realizeSpatialSchedule(*candidate, remapped, target, stages, + SpatialDataflowExportStage::None, nullptr, &scheduleReport)) + || failed(lower(*candidate, &summary))) return failure(); + if (enableTiming) summary["codegen_input_ir"] = countIR(candidate->getOperation()); + summary["lowering_seconds"] = elapsed(begin); + return success(); + } + + void runOnOperation() override { + ModuleOp module = getOperation(); + auto start = Clock::now(); + std::string unsupported; + if (!target.prepareScheduleCost) + unsupported = target.wholeScheduleCostUnavailable.empty() + ? "target has no whole-schedule cost query" : target.wholeScheduleCostUnavailable; + auto prepared = prepareSpatialSchedules(module, target, stages, unsupported.empty()); + if (failed(prepared)) { signalPassFailure(); return; } + bool singleCandidate = unsupported.empty() && prepared->candidates.size() == 1; + double generationSeconds = elapsed(start); + std::vector scores(prepared->candidates.size()); + llvm::json::Array summaries; + OwningOpRef bestModule; + std::string bestScheduleReport; + size_t bestIndex = 0; + double loweringSeconds = 0, instructionSeconds = 0, queueSeconds = 0; + size_t lowered = 0, analyzed = 0, unavailable = 0, duplicates = 0; + for (size_t i = 0; i < prepared->candidates.size(); ++i) { + llvm::json::Object summary{{"index", static_cast(i)}}; + if (!unsupported.empty()) { + scores[i].reason = unsupported; + } else if (singleCandidate) { + scores[i].reason = "one distinct generated schedule; no cost comparison required"; + } else { + ScopedCompilerOutputSuppression silence; + OwningOpRef candidate; + std::string scheduleReport; + if (failed(evaluate(module, prepared->candidates[i], candidate, scheduleReport, summary))) { + if (!i) { signalPassFailure(); return; } + scores[i].reason = "candidate lowering failed verification"; + } else { + ++lowered; + loweringSeconds += *summary.getNumber("lowering_seconds"); + auto query = target.prepareScheduleCost(*candidate, summary); + if (failed(query)) { + if (!i) { signalPassFailure(); return; } + scores[i].reason = "candidate cost preparation failed verification"; + ++unavailable; + summary["eligible"] = false; + summary["reason"] = scores[i].reason; + summaries.push_back(std::move(summary)); + continue; + } + instructionSeconds += summary.getNumber("instruction_generation_seconds").value_or(0); + scores[i].identity = query->identity; + auto duplicate = std::find_if(scores.begin(), scores.begin() + i, + [&](const WholeScheduleCost &other) { return other.identity == scores[i].identity; }); + if (duplicate != scores.begin() + i) { + ++duplicates; + scores[i].reason = "duplicate emitted instruction program"; + summary["duplicate_of"] = static_cast(duplicate - scores.begin()); + } else { + auto begin = Clock::now(); + scores[i] = query->evaluate(summary); + scores[i].identity = query->identity; + double seconds = elapsed(begin); + queueSeconds += seconds; + ++analyzed; + summary["queue_seconds"] = seconds; + } + std::string partialReason; + size_t partialSelection = selectPipelineCandidate( + llvm::ArrayRef(scores.data(), i + 1), partialReason); + if (!bestModule || partialSelection == i) { + bestIndex = i; + bestModule = std::move(candidate); + bestScheduleReport = std::move(scheduleReport); + } + } + } + if (!singleCandidate && !scores[i].eligible() + && scores[i].reason != "duplicate emitted instruction program") ++unavailable; + summary["eligible"] = scores[i].eligible(); + summary["identity"] = scores[i].identity; + if (scores[i].primary) summary["interval_ns"] = *scores[i].primary; + if (scores[i].alternative) summary["alternative_tie_interval_ns"] = *scores[i].alternative; + if (scores[i].relativeUncertainty) + summary["relative_prediction_uncertainty"] = *scores[i].relativeUncertainty; + if (!scores[i].reason.empty()) summary["reason"] = scores[i].reason; + summaries.push_back(std::move(summary)); + } + std::string reason; + size_t selected = selectPipelineCandidate(scores, reason); + if (singleCandidate) reason = scores.front().reason; + else if (!unsupported.empty()) reason += ": " + unsupported; + else if (!scores.front().eligible() && !scores.front().reason.empty()) + reason += ": " + scores.front().reason; + auto finalizeStart = Clock::now(); + bool spatialArtifacts = exportStage != SpatialDataflowExportStage::None + || pimPeriodicDiagnostics || isPimIrDumpEnabled(); + // Explicit Spatial diagnostics require the pre-lowering representation. + if (!bestModule || spatialArtifacts) { + if (failed(realizeSpatialSchedule(module, prepared->candidates[selected], target, + stages, exportStage, pimPeriodicDiagnostics ? &prepared->graph : nullptr))) { signalPassFailure(); return; } + } else { + auto file = openDialectDumpFileWithExtension("spatial3_scheduled_no_comm", "/reports", "txt"); + if (file) file << bestScheduleReport; + } + if (bestModule) { + assert(selected == bestIndex && "retained module must be the selected candidate"); + module->setAttrs((*bestModule)->getAttrs()); + module.getBodyRegion().takeBody(bestModule->getBodyRegion()); + } else if (failed(lower(module))) { signalPassFailure(); return; } + llvm::errs() << "Queue-aware pipeline: candidate " << selected << " of " << scores.size() + << "; " << reason << "\n"; + llvm::json::Object report{{"selection", static_cast(selected)}, {"reason", reason}, + {"static_analysis", true}, {"strategy", "queue-aware"}, {"candidates", std::move(summaries)}, + {"single_candidate_analysis_skipped", singleCandidate}, + {"candidate_generation_seconds", generationSeconds}, {"candidate_lowering_seconds", loweringSeconds}, + {"instruction_generation_seconds", instructionSeconds}, + {"queue_seconds", queueSeconds}, {"finalization_seconds", elapsed(finalizeStart)}, {"total_seconds", elapsed(start)}, + {"generated_candidates", static_cast(scores.size())}, + {"lowered_candidates", static_cast(lowered)}, + {"reused_lowered_candidate", static_cast(bestModule)}, + {"reconstructed_spatial_artifacts", spatialArtifacts && static_cast(bestModule)}, + {"analyzed_candidates", static_cast(analyzed)}, + {"unavailable_estimates", static_cast(unavailable)}, + {"duplicate_programs", static_cast(duplicates)}, {"temporary_disk_bytes", 0}}; + auto file = openReportFileWithExtension("pipeline_selection", "json"); + if (file) { llvm::raw_os_ostream out(file); out << llvm::formatv("{0:2}\n", llvm::json::Value(std::move(report))); } + } +}; +} // namespace + +std::unique_ptr onnx_mlir::createPimPipelineSelectionPass( + const SchedulingTarget &target, const SpatialTargetResources &resources, + size_t stages, SpatialDataflowExportStage exportStage) { + return std::make_unique(target, resources, stages, exportStage); +} diff --git a/src/PIM/Compiler/PimPipelineSelection.hpp b/src/PIM/Compiler/PimPipelineSelection.hpp new file mode 100644 index 0000000..4c7d972 --- /dev/null +++ b/src/PIM/Compiler/PimPipelineSelection.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "mlir/Pass/Pass.h" + +namespace onnx_mlir { +namespace spatial { struct SchedulingTarget; struct SpatialTargetResources; enum class SpatialDataflowExportStage; } +std::unique_ptr createPimPipelineSelectionPass( + const spatial::SchedulingTarget &target, const spatial::SpatialTargetResources &resources, + size_t stages, spatial::SpatialDataflowExportStage exportStage); +} // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimScheduleDiagnostic.cpp b/src/PIM/Compiler/PimScheduleDiagnostic.cpp new file mode 100644 index 0000000..09c43ee --- /dev/null +++ b/src/PIM/Compiler/PimScheduleDiagnostic.cpp @@ -0,0 +1,467 @@ +#include "PimScheduleDiagnostic.hpp" +#include "PimInstructionTiming.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.hpp" +#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" + +#include "llvm/ADT/MapVector.h" +#include "llvm/Support/JSON.h" +#include "llvm/Support/raw_os_ostream.h" + +#include +#include + +namespace onnx_mlir { +using namespace spatial; +namespace { + +struct Event { + int64_t core; + llvm::StringRef kind; + int64_t identity; +}; +struct DirectEndpoint { + size_t begin; + size_t end; + int64_t core; + Cost bytes; +}; +struct Signal { + size_t event; + int64_t core; +}; +struct Wait { + size_t event; + int64_t count; +}; +using Register = std::pair; + +class PeriodicPlan { +public: + PeriodicPlan(const PimInstructionTiming &target, + const DeferredTransferPlan &transfers) + : target(target), transfers(transfers), computeWork(target.processorCount), + entry(target.processorCount), tail(target.processorCount), + active(target.processorCount) {} + + bool build(const ComputeGraph &graph, + const ScheduledComputeMaterializationResult &materialization, + const DeferredBoundaryPlan &boundaries); + + const PimInstructionTiming ⌖ + const DeferredTransferPlan &transfers; + std::vector events; + std::vector edges; + SharedMemoryDemand boundaryMemory; + std::vector computeWork; + std::string error; + void recordSynchronization(llvm::json::Object &report) const; + +private: + size_t event(int64_t core, llvm::StringRef kind, int64_t identity = -1) { + events.push_back({core, kind, identity}); + return events.size() - 1; + } + size_t append(int64_t core, llvm::StringRef kind, Cost delay = 0, + int64_t identity = -1) { + size_t next = event(core, kind, identity); + edges.push_back({tail[core], next, delay, 0}); + tail[core] = next; + return next; + } + void signal(int64_t source, int64_t targetCore, unsigned reg, + llvm::StringRef kind) { + size_t begin = append(source, kind, 0, targetCore); + signals[{targetCore, reg}].push_back({begin, source}); + append(source, "signal-complete", target.signalIssueNs); + } + void wait(int64_t core, unsigned reg, int64_t count, llvm::StringRef kind) { + if (!count) + return; + size_t next = append(core, kind, 0, reg); + if (!waits.try_emplace({core, reg}, Wait{next, count}).second) + error = "multiple waits consume one event register in a generation"; + } + void entrySynchronization(); + void completionSynchronization(); + bool boundary(const BoundaryProgram &program); + bool transfer(const ScheduledTransferSlice &slice, bool send); + bool connect(); + + std::vector entry, tail; + std::vector active; + std::map> signals; + std::map waits; + llvm::MapVector sends, receives; +}; + +void PeriodicPlan::entrySynchronization() { + for (const ScheduledInfo &scheduled : transfers.scheduled) + for (int64_t core : scheduled.cores) { + auto found = transfers.hostReleaseSignals.find(core); + if (found != transfers.hostReleaseSignals.end()) + for (HostReleaseSignal release : found->second) + signal(core, release.writerCore, release.eventRegister, "free-signal"); + } + unsigned rounds = 0; + for (size_t distance = 1; distance < transfers.stageZeroCores.size(); distance *= 2) + ++rounds; + unsigned last = transfers.synchronizationRegisterCount - 1; + for (auto [rank, core] : llvm::enumerate(transfers.downstreamCores)) { + int64_t parent = rank ? transfers.downstreamCores[(rank - 1) / 2] + : transfers.stageZeroCores.front(); + signal(core, parent, rank ? last - 1 : last - rounds, "restart-permission"); + } +} + +void PeriodicPlan::completionSynchronization() { + unsigned rounds = 0; + for (size_t distance = 1; distance < transfers.stageZeroCores.size(); distance *= 2) + ++rounds; + unsigned last = transfers.synchronizationRegisterCount - 1; + for (auto [rank, core] : llvm::enumerate(transfers.stageZeroCores)) + for (unsigned round = 0; round < rounds; ++round) { + int64_t peer = transfers.stageZeroCores[ + (rank + (size_t{1} << round)) % transfers.stageZeroCores.size()]; + signal(core, peer, last - round, "barrier-signal"); + wait(core, last - round, 1, "barrier-wait"); + } + if (!transfers.downstreamCores.empty()) { + int64_t leader = transfers.stageZeroCores.front(); + wait(leader, last - rounds, 1, "restart-permission-wait"); + signal(leader, transfers.downstreamCores.front(), last, "restart-signal"); + } + for (auto [rank, core] : llvm::enumerate(transfers.downstreamCores)) { + wait(core, last, 1, "restart-wait"); + size_t first = 2 * rank + 1; + size_t end = std::min(first + 2, transfers.downstreamCores.size()); + if (first < end) + wait(core, last - 1, end - first, "restart-permission-wait"); + for (size_t child = first; child < end; ++child) + signal(core, transfers.downstreamCores[child], last, "restart-signal"); + } +} + +bool PeriodicPlan::transfer(const ScheduledTransferSlice &slice, bool send) { + const ExternalTransferFamily &family = *slice.family; + auto type = mlir::dyn_cast(family.requirement->publicationFragmentType); + if (!type || !type.hasStaticShape() || !type.getElementType().isIntOrFloat()) { + error = "transfer payload has no static byte size"; + return false; + } + Cost bits = std::min(type.getElementType().getIntOrFloatBitWidth(), + target.computeBitwidth); + Cost payloadBits = checkedMultiply(static_cast(type.getNumElements()), bits); + Cost bytes = payloadBits / 8 + (payloadBits % 8 != 0); + for (size_t offset = 0; offset < slice.transferCount; ++offset) { + size_t index = slice.familyOffset + offset; + int64_t source = family.sourceCores.valueAt(index); + int64_t destination = family.targetCores.valueAt(index); + int64_t core = send ? source : destination; + int64_t channel = family.channelIds.valueAt(index); + if (!family.hostRouted) { + size_t begin = append(core, send ? "send-begin" : "receive-begin", 0, channel); + size_t end = append(core, send ? "send-complete" : "receive-complete", 0, channel); + if (!(send ? sends : receives).insert({channel, {begin, end, core, bytes}}).second) { + error = "direct channel has multiple physical endpoints"; + return false; + } + continue; + } + unsigned reg = family.eventRegisters.valueAt(index); + if (send) { + if (family.hostStoreFlags.valueAt(index)) { + SharedMemoryDemand demand = target.getSharedMemoryDemand(bytes, 1, true); + boundaryMemory += demand; + Cost service = checkedAdd(target.getLocalMemoryServiceNs(bytes, false), + checkedAdd(target.getSharedMemoryLinkServiceNs(bytes, core, false), demand.serviceNs)); + append(core, "host-store", service, family.hostOffsets.valueAt(index)); + } + signal(core, destination, reg, "ready-signal"); + } else { + wait(core, reg, family.waitValues.valueAt(index), "ready-wait"); + SharedMemoryDemand demand = target.getSharedMemoryDemand(bytes, 1, false); + boundaryMemory += demand; + Cost service = checkedAdd(target.getSharedMemoryLinkServiceNs( + target.controlMessageBytes, core, false), + checkedAdd(demand.serviceNs, checkedAdd( + target.getSharedMemoryLinkServiceNs( + target.sharedMemoryReadReturnsPayload ? bytes : 0, core, true), + target.getLocalMemoryServiceNs(bytes, true)))); + append(core, "host-load", service, family.hostOffsets.valueAt(index)); + } + } + return true; +} + +bool PeriodicPlan::boundary(const BoundaryProgram &program) { + for (const BoundaryInstruction &instruction : program.instructions) { + if (const auto *run = std::get_if(&instruction)) { + for (const auto &slice : run->slices) + if (!transfer(slice, true)) + return false; + } else if (const auto *run = std::get_if(&instruction)) { + for (const auto &slice : run->slices) + if (!transfer(slice, false)) + return false; + } else if (const auto *reuse = std::get_if(&instruction)) { + for (LaneInterval lanes : reuse->lanes.intervals()) + for (unsigned lane = lanes.begin; lane < lanes.end; ++lane) + wait(program.key.first->cores[lane], reuse->eventRegisters.valueAt(lane), + reuse->waitValues.valueAt(lane), "free-wait"); + } + } + return error.empty(); +} + +bool PeriodicPlan::connect() { + if (sends.size() != receives.size()) { + error = "direct channel endpoint counts differ"; + return false; + } + for (const auto &[channel, send] : sends) { + auto found = receives.find(channel); + if (found == receives.end() || found->second.bytes != send.bytes) { + error = "direct channel has no matching payload"; + return false; + } + const DirectEndpoint &receive = found->second; + size_t permission = event(receive.core, "peer-permission", channel); + edges.push_back({send.begin, permission, + target.getLinkServiceNs(target.controlMessageBytes, send.core, receive.core), 0}); + edges.push_back({receive.begin, permission, 0, 0}); + edges.push_back({permission, send.end, checkedAdd( + target.getLinkServiceNs(target.controlMessageBytes, receive.core, send.core), + target.getLocalMemoryServiceNs(send.bytes, false)), 0}); + edges.push_back({send.end, receive.end, checkedAdd( + target.getLinkServiceNs(send.bytes, send.core, receive.core), + target.getLocalMemoryServiceNs(send.bytes, true)), 0}); + } + if (signals.size() != waits.size()) { + error = "signal and wait register sets differ"; + return false; + } + for (const auto &[key, pending] : waits) { + auto found = signals.find(key); + if (found == signals.end() || pending.count < 0 + || found->second.size() != static_cast(pending.count)) { + error = "event register does not receive its exact static count"; + return false; + } + for (Signal signal : found->second) + edges.push_back({signal.event, pending.event, + target.getLinkServiceNs(target.controlMessageBytes, signal.core, key.first), 0}); + } + for (size_t core = 0; core < active.size(); ++core) + if (active[core]) + edges.push_back({tail[core], entry[core], 0, 1}); + return true; +} + +void PeriodicPlan::recordSynchronization(llvm::json::Object &report) const { + llvm::json::Array registers; + for (const auto &[key, pending] : waits) { + llvm::json::Array sources; + std::map counts; + auto found = signals.find(key); + if (found != signals.end()) + for (Signal signal : found->second) + ++counts[signal.core]; + for (auto [core, count] : counts) + sources.push_back(llvm::json::Object{{"core", core}, {"count", count}}); + registers.push_back(llvm::json::Object{{"core", key.first}, + {"register", key.second}, {"expected", pending.count}, + {"kind", events[pending.event].kind}, {"sources", std::move(sources)}}); + } + report["event_registers"] = std::move(registers); +} + +bool PeriodicPlan::build(const ComputeGraph &graph, + const ScheduledComputeMaterializationResult &materialization, + const DeferredBoundaryPlan &boundaries) { + if (!target.hasPeriodicTransferTiming()) { + error = "target has no complete periodic communication timing contract"; + return false; + } + if (transfers.scheduled.empty() || transfers.stageZeroCores.empty() + || !transfers.synchronizationRegisterCount) { + error = "periodic timing requires an active repeating pipeline body"; + return false; + } + for (auto [index, scheduled] : llvm::enumerate(transfers.scheduled)) { + if (!materialization.materializedSchedules[index].runs.empty() + && scheduled.cores.size() != 1) { + error = "compact compute runs require one physical core"; + return false; + } + for (int64_t core : scheduled.cores) { + if (active[core]) { + error = "physical core appears in multiple scheduled bodies"; + return false; + } + active[core] = true; + tail[core] = entry[core] = event(core, "body-entry"); + } + } + entrySynchronization(); + llvm::DenseMap, + const BoundaryProgram *> indexed; + for (const BoundaryProgram &program : boundaries.boundaries) + indexed[program.key] = &program; + for (auto [recordIndex, record] : llvm::enumerate(materialization.materializedSchedules)) { + const ScheduledInfo *scheduled = &transfers.scheduled[recordIndex]; + for (unsigned step = 0; step <= scheduled->stepCount; ++step) { + if (const auto *program = indexed.lookup({scheduled, step})) + if (!boundary(*program)) + return false; + if (step == scheduled->stepCount) + continue; + auto appendCompute = [&](const ComputeInstance &instance, int64_t core) { + auto found = graph.instanceToIndex.find(instance); + if (found == graph.instanceToIndex.end()) { + error = "scheduled compute instance is absent from the cost graph"; + return; + } + append(core, "compute", graph.nodes[found->second].cost, found->second); + computeWork[core] = checkedAdd(computeWork[core], graph.nodes[found->second].cost); + }; + if (!record.runs.empty()) { + Cost cost = 0; + for (const ComputeInstance &instance : record.runs[step].instances) { + auto found = graph.instanceToIndex.find(instance); + if (found == graph.instanceToIndex.end()) { + error = "scheduled run instance is absent from the cost graph"; + return false; + } + cost = checkedAdd(cost, graph.nodes[found->second].cost); + } + append(scheduled->cores.front(), "compute-run", cost, step); + int64_t core = scheduled->cores.front(); + computeWork[core] = checkedAdd(computeWork[core], cost); + } else { + for (auto [lane, instance] : llvm::enumerate(record.stepPlans[step].stepTuple.instances)) + appendCompute(instance, scheduled->cores[lane]); + } + } + } + completionSynchronization(); + return error.empty() && connect(); +} + +static mlir::FailureOr analyzePlan(PeriodicPlan &plan, + const ComputeGraph &graph, + const ScheduledComputeMaterializationResult &materialization, + const DeferredBoundaryPlan &boundaries) { + if (!plan.build(graph, materialization, boundaries)) + return mlir::failure(); + return analyzePeriodicDependencies(plan.events.size(), plan.edges, plan.error); +} + +} // namespace + +llvm::json::Object getPeriodicScheduleReport(const ComputeGraph &graph, + const ScheduledComputeMaterializationResult &materialization, + const DeferredTransferPlan &transfers, const DeferredBoundaryPlan &boundaries, + const PimInstructionTiming &target) { + auto start = std::chrono::steady_clock::now(); + PeriodicPlan plan(target, transfers); + llvm::json::Object report{{"diagnostic_only", true}, + {"selection_eligible", false}, + {"timing_model", "serial compute chunks; uncontended communication; no issue or local contention model"}}; + report["omitted_effects"] = llvm::json::Array{ + "placement and timing of baseline memory accesses", + "communication introduced after deferred realization, including final output collection", + "instruction issue, execution-unit overlap and local-memory contention", + "projection materialization, strided transfer splitting and load folding", + "network injection queues and global-memory arbitration"}; + report["native_host_buffer_bytes"] = static_cast(transfers.pipelineHostBufferBytes); + auto estimate = analyzePlan(plan, graph, materialization, boundaries); + if (mlir::succeeded(estimate)) { + plan.recordSynchronization(report); + SharedMemoryDemand memory = graph.baselineMemory; + memory += plan.boundaryMemory; + report["recurrence_lower_ns"] = estimate->recurrenceLowerNs; + report["recurrence_upper_ns"] = estimate->recurrenceUpperNs; + report["within_generation_span_ns"] = static_cast(estimate->withinGenerationSpanNs); + report["global_memory_service_ns"] = static_cast(memory.serviceNs); + report["baseline_memory_service_ns"] = static_cast(graph.baselineMemory.serviceNs); + report["boundary_memory_service_ns"] = static_cast(plan.boundaryMemory.serviceNs); + report["boundary_read_bytes"] = static_cast(plan.boundaryMemory.readBytes); + report["boundary_write_bytes"] = static_cast(plan.boundaryMemory.writeBytes); + report["boundary_read_requests"] = static_cast(plan.boundaryMemory.readRequests); + report["boundary_write_requests"] = static_cast(plan.boundaryMemory.writeRequests); + Cost coreWork = *std::max_element(plan.computeWork.begin(), plan.computeWork.end()); + report["maximum_core_compute_ns"] = static_cast(coreWork); + report["interval_estimate_ns"] = std::max({estimate->recurrenceUpperNs, + static_cast(memory.serviceNs), static_cast(coreWork)}); + report["relaxation_rounds"] = static_cast(estimate->relaxationRounds); + llvm::json::Array cycle; + std::map kinds; + std::map coreDelays; + std::map, Cost> linkDelays; + Cost delay = 0, distance = 0; + for (size_t index : estimate->criticalCycle) { + const auto &edge = plan.edges[index]; + const Event &source = plan.events[edge.source], &destination = plan.events[edge.target]; + delay = checkedAdd(delay, edge.delayNs); + distance = checkedAdd(distance, static_cast(edge.distance)); + ++kinds[destination.kind]; + coreDelays[destination.core] = checkedAdd(coreDelays[destination.core], edge.delayNs); + if (source.core != destination.core) + linkDelays[{source.core, destination.core}] = checkedAdd( + linkDelays[{source.core, destination.core}], edge.delayNs); + if (cycle.size() < 128) + cycle.push_back(llvm::json::Object{{"source_core", source.core}, + {"source_kind", source.kind}, {"target_core", destination.core}, + {"target_kind", destination.kind}, {"identity", destination.identity}, + {"delay_ns", static_cast(edge.delayNs)}, + {"distance", edge.distance}}); + } + llvm::json::Object counts; + for (auto [kind, count] : kinds) + counts[kind] = count; + report["critical_cycle"] = std::move(cycle); + report["critical_cycle_event_counts"] = std::move(counts); + report["critical_cycle_edges"] = static_cast(estimate->criticalCycle.size()); + report["critical_cycle_delay_ns"] = static_cast(delay); + report["critical_cycle_distance"] = static_cast(distance); + llvm::json::Array cores, links; + for (auto [core, ns] : coreDelays) + cores.push_back(llvm::json::Object{{"core", core}, {"delay_ns", static_cast(ns)}}); + for (auto [pair, ns] : linkDelays) + links.push_back(llvm::json::Object{{"source", pair.first}, {"target", pair.second}, + {"delay_ns", static_cast(ns)}}); + report["critical_cycle_cores"] = std::move(cores); + report["critical_cycle_links"] = std::move(links); + } + report["status"] = plan.error.empty() ? "estimated" : "unavailable"; + if (!plan.error.empty()) + report["reason"] = plan.error; + report["events"] = static_cast(plan.events.size()); + report["edges"] = static_cast(plan.edges.size()); + llvm::json::Object eventCounts; + std::map counts; + for (const Event &event : plan.events) + ++counts[event.kind]; + for (auto [kind, count] : counts) + eventCounts[kind] = count; + report["event_counts"] = std::move(eventCounts); + report["analysis_wall_ms"] = std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + return report; +} + +void dumpPeriodicScheduleEstimate(const ComputeGraph &graph, + const ScheduledComputeMaterializationResult &materialization, + const DeferredTransferPlan &transfers, const DeferredBoundaryPlan &boundaries, + const PimInstructionTiming &target) { + auto file = openDialectDumpFileWithExtension("pipeline_periodic", "/reports", "json"); + if (!file.is_open()) + return; + auto report = getPeriodicScheduleReport(graph, materialization, transfers, + boundaries, target); + llvm::raw_os_ostream os(file); + os << llvm::formatv("{0:2}\n", llvm::json::Value(std::move(report))); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimScheduleDiagnostic.hpp b/src/PIM/Compiler/PimScheduleDiagnostic.hpp new file mode 100644 index 0000000..392288a --- /dev/null +++ b/src/PIM/Compiler/PimScheduleDiagnostic.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp" +#include "llvm/Support/JSON.h" + +namespace onnx_mlir { +namespace spatial { struct ScheduledComputeMaterializationResult; } +struct PimInstructionTiming; + +llvm::json::Object getPeriodicScheduleReport(const spatial::ComputeGraph &graph, + const spatial::ScheduledComputeMaterializationResult &materialization, + const spatial::DeferredTransferPlan &transfers, const spatial::DeferredBoundaryPlan &boundaries, + const PimInstructionTiming &target); + +// Read-only diagnostics over the final physical transfer and synchronization plan. +void dumpPeriodicScheduleEstimate(const spatial::ComputeGraph &graph, + const spatial::ScheduledComputeMaterializationResult &materialization, + const spatial::DeferredTransferPlan &transfers, const spatial::DeferredBoundaryPlan &boundaries, + const PimInstructionTiming &target); + +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/CMakeLists.txt b/src/PIM/Dialect/Spatial/CMakeLists.txt index e1389e2..64f8431 100644 --- a/src/PIM/Dialect/Spatial/CMakeLists.txt +++ b/src/PIM/Dialect/Spatial/CMakeLists.txt @@ -39,6 +39,7 @@ add_pim_library(SpatialOps Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp Passes/Transforms/MergeComputeNodes/Scheduling/PipelineWorkloadPreparation.cpp Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.cpp + Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.cpp Passes/Transforms/TrivialGraphComputeMergePass.cpp EXCLUDE_FROM_OM_LIBS diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp index cee807c..46a8784 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp @@ -465,8 +465,20 @@ LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input, LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc, Block &body, ArrayRef plans, IRMapping &mapper, - llvm::SmallPtrSetImpl &absorbed) { + llvm::SmallPtrSetImpl &absorbed, + bool deterministicOrder) { + // Deferred demand order must not depend on allocator addresses. + DenseMap valueOrder; + if (deterministicOrder) { + for (BlockArgument argument : body.getArguments()) + valueOrder.try_emplace(argument, valueOrder.size()); + body.walk([&](Operation *op) { + for (Value result : op->getResults()) + valueOrder.try_emplace(result, valueOrder.size()); + }); + } for (const DeferredInputPlan &plan : plans) { + valueOrder.try_emplace(plan.graphInput, valueOrder.size()); if (plan.availableValue) { mapper.map(plan.graphInput, plan.availableValue); continue; } SmallVector roots; bool needsIdentity = false; @@ -498,7 +510,11 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc } } if (needsIdentity) roots.push_back(plan.graphInput); - llvm::sort(roots, [](Value a, Value b) { return a.getAsOpaquePointer() < b.getAsOpaquePointer(); }); + llvm::sort(roots, [&](Value a, Value b) { + if (!deterministicOrder) return a.getAsOpaquePointer() < b.getAsOpaquePointer(); + assert(valueOrder.count(a) && valueOrder.count(b)); + return valueOrder.lookup(a) < valueOrder.lookup(b); + }); roots.erase(std::unique(roots.begin(), roots.end()), roots.end()); Value sharedSource; if (!plan.blueprint) { diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp index 7f719d3..316a672 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp @@ -47,7 +47,8 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc Block &graphBody, ArrayRef plans, IRMapping &mapper, - llvm::SmallPtrSetImpl &absorbed); + llvm::SmallPtrSetImpl &absorbed, + bool deterministicOrder); } // namespace spatial } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp index 422f08d..59eafcc 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp @@ -148,7 +148,8 @@ static LogicalResult verifyDominance(func::FuncOp funcOp) { LogicalResult realizeDeferredCommunication(func::FuncOp funcOp, const ScheduledComputeMaterializationResult& materialization, const SchedulingTarget& target, - size_t pipelineStages) { + size_t pipelineStages, + const ComputeGraph *costGraph) { IRRewriter rewriter(funcOp.getContext()); eraseUnusedIdentityDeferredCommunications(funcOp, rewriter); @@ -172,6 +173,8 @@ LogicalResult realizeDeferredCommunication(func::FuncOp funcOp, *transfers, *schedule, target.synchronizationRegisterCount); if (failed(boundaries)) return funcOp.emitOpError("phase 2 failed to build sparse boundary programs"); + if (pipelineStages > 1 && costGraph && target.reportPeriodicSchedule) + target.reportPeriodicSchedule(*costGraph, materialization, *transfers, *boundaries); if (failed(retargetDeferredPublications(funcOp, *transfers)) || failed(replaceFinalGraphPublications(funcOp, *transfers))) diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp index 7a19be1..898886a 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp @@ -6,10 +6,12 @@ namespace onnx_mlir::spatial { struct ScheduledComputeMaterializationResult; struct SchedulingTarget; +struct ComputeGraph; mlir::LogicalResult realizeDeferredCommunication(mlir::func::FuncOp funcOp, const ScheduledComputeMaterializationResult& materialization, const SchedulingTarget& target, - size_t pipelineStages = 1); + size_t pipelineStages = 1, + const ComputeGraph *costGraph = nullptr); } // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp index 519f8af..c9727d1 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp @@ -239,7 +239,7 @@ static LogicalResult materializeResultfulBatchRun( Block &source = batch.getBody().front(); llvm::SmallPtrSet absorbed; - if (failed(materializeDeferredPayloadDemands(builder, bodyLoc, source, inputPlans, mapper, absorbed))) + if (failed(materializeDeferredPayloadDemands(builder, bodyLoc, source, inputPlans, mapper, absorbed, !schedule.processorStages.empty()))) return failure(); for (Operation &op : source.without_terminator()) if (!absorbed.contains(&op)) @@ -334,7 +334,7 @@ static LogicalResult materializeSingleCpuPeftClass( inputPlans.push_back(std::move(plan)); } llvm::SmallPtrSet absorbed; - if (failed(materializeDeferredPayloadDemands(rewriter, compute.getLoc(), compute.getBody().front(), inputPlans, mapper, absorbed))) + if (failed(materializeDeferredPayloadDemands(rewriter, compute.getLoc(), compute.getBody().front(), inputPlans, mapper, absorbed, !schedule.processorStages.empty()))) return failure(); cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues, absorbed); } else { @@ -589,7 +589,7 @@ static LogicalResult materializeMultiCpuPeftClass( } SmallVector yieldedValues; llvm::SmallPtrSet absorbed; - if (failed(materializeDeferredPayloadDemands(rewriter, compute.getLoc(), compute.getBody().front(), inputPlans, mapper, absorbed))) + if (failed(materializeDeferredPayloadDemands(rewriter, compute.getLoc(), compute.getBody().front(), inputPlans, mapper, absorbed, !schedule.processorStages.empty()))) return failure(); cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues, absorbed); for (Value yielded : yieldedValues) { @@ -668,7 +668,7 @@ static LogicalResult materializeMultiCpuPeftClass( for (unsigned index = 0; index < batch.getNumResults(); ++index) mapper.map(*batch.getOutputArgument(index), iterArgs[index]); llvm::SmallPtrSet absorbed; - if (failed(materializeDeferredPayloadDemands(builder, bodyLoc, batch.getBody().front(), inputPlans, mapper, absorbed))) + if (failed(materializeDeferredPayloadDemands(builder, bodyLoc, batch.getBody().front(), inputPlans, mapper, absorbed, !schedule.processorStages.empty()))) return failure(); for (Operation &op : batch.getBody().front().without_terminator()) if (!absorbed.contains(&op)) diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp index a1925a5..079a7c1 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp @@ -79,6 +79,13 @@ void dumpScheduledComputeReport(ModuleOp moduleOp, func::FuncOp funcOp, const Me std::fstream file = openDialectDumpFileWithExtension("spatial3_scheduled_no_comm", "/reports", "txt"); if (!file.is_open()) return; llvm::raw_os_ostream os(file); + writeScheduledComputeReport(os, moduleOp, funcOp, schedule, peftClassPlans, records); +} + +void writeScheduledComputeReport(raw_ostream &os, ModuleOp moduleOp, func::FuncOp funcOp, + const MergeScheduleResult &schedule, + const llvm::MapVector &peftClassPlans, + ArrayRef records) { OpPrintingFlags flags; flags.elideLargeElementsAttrs().enableDebugInfo(false, false).assumeVerified(); AsmState asmState(moduleOp.getOperation(), flags); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp index e17670a..9daee61 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp @@ -5,6 +5,12 @@ namespace onnx_mlir { namespace spatial { +void writeScheduledComputeReport(llvm::raw_ostream &os, ModuleOp moduleOp, + func::FuncOp funcOp, + const MergeScheduleResult &schedule, + const llvm::MapVector &peftClassPlans, + ArrayRef materializedSchedules); + void dumpScheduledComputeReport(ModuleOp moduleOp, func::FuncOp funcOp, const MergeScheduleResult &schedule, diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp index b0a633d..5ed6088 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp @@ -1,5 +1,6 @@ #include "mlir/Pass/Pass.h" +#include "ScheduledSpatialPasses.hpp" #include "DeferredCommunicationRealization.hpp" #include "ScheduledComputeReport.hpp" #include "ScheduledComputeVerification.hpp" @@ -40,9 +41,9 @@ struct ScheduleAndRealizeSpatialPass final ScheduleAndRealizeSpatialPass() = default; ScheduleAndRealizeSpatialPass(const SchedulingTarget& target, SpatialDataflowExportStage exportStage, - size_t pipelineStages) + size_t pipelineStages, bool periodicDiagnostics) : target(target), exportStage(exportStage), - pipelineStages(pipelineStages), hasTarget(true) {} + pipelineStages(pipelineStages), periodicDiagnostics(periodicDiagnostics), hasTarget(true) {} StringRef getArgument() const override { return "schedule-and-realize-spatial"; } StringRef getDescription() const override { @@ -64,113 +65,10 @@ struct ScheduleAndRealizeSpatialPass final signalPassFailure(); return; } - auto entry = requireEntry(moduleOp); - if (failed(entry)) { - signalPassFailure(); - return; - } - func::FuncOp entryFunc = *entry; - - ComputeGraph scheduledGraph; - MergeScheduleResult schedule; - for (;;) { - auto graph = buildComputeGraph(entryFunc, target); - if (failed(graph)) { - signalPassFailure(); - return; - } - scheduledGraph = std::move(*graph); - std::string pipelineError; - if (pipelineStages > 1) { - FailureOr preparation = - preparePipelineWorkload( - scheduledGraph, pipelineStages, target, pipelineError); - if (failed(preparation)) { - moduleOp.emitError() << pipelineError; - signalPassFailure(); - return; - } - if (*preparation == PipelineWorkloadPreparation::Changed) - continue; - } - if (pipelineStages == 1) { - MergeSchedulingAnalysis analysis( - scheduledGraph, target, entryFunc.getContext(), pipelineError); - if (failed(analysis.getResult())) { - moduleOp.emitError() << pipelineError; - signalPassFailure(); - return; - } - schedule = std::move(*analysis.getResult()); - break; - } - if (succeeded(runPipelineScheduler( - scheduledGraph, schedule, pipelineStages, target, pipelineError))) - break; - moduleOp.emitError() << pipelineError; - signalPassFailure(); - return; - } - PatternRewriter rewriter(moduleOp.getContext()); - FailureOr materialization = - materializeScheduledCompute(entryFunc, schedule, rewriter); - if (failed(materialization)) { - signalPassFailure(); - return; - } - ScheduledComputeMaterializationResult materializationResult = std::move(*materialization); - - if (failed(verifyMaterializedScheduleMapping( - entryFunc, schedule, materializationResult.peftClassPlans, - materializationResult.graphComputeToBlockMap, - materializationResult.materializedSchedules)) - || failed(verifyDeferredTransferPhase1Invariants(entryFunc)) - || failed(verifyScheduledMaterializationRecords(materializationResult.materializedSchedules)) - || failed(verifyScheduledSpatialInvariants(entryFunc))) { - moduleOp.emitError("scheduled Spatial phase verification failed"); - signalPassFailure(); - return; - } - - dumpModule(moduleOp, "spatial3_scheduled_no_comm", /*assumeVerified=*/true); - if (shouldExportSpatialDataflowStage(exportStage, SpatialDataflowExportStage::Spatial3) - && failed(exportSpatialDataflowCsvScheduled( - entryFunc, materializationResult.materializedSchedules, - "spatial3_scheduled_no_comm", "spatial3"))) { - signalPassFailure(); - return; - } - dumpScheduledComputeReport( - moduleOp, entryFunc, schedule, materializationResult.peftClassPlans, - materializationResult.materializedSchedules); - - if (failed(realizeDeferredCommunication( - entryFunc, materializationResult, target, pipelineStages))) { - moduleOp.emitError("Spatial communication realization failed"); - signalPassFailure(); - return; - } - - bool deferredRemains = false; - entryFunc.walk([&](SpatDeferredCommunicationOp deferred) { - if (deferredRemains) - return; - deferred.emitOpError("realized Spatial graph still contains deferred communication"); - deferredRemains = true; - }); - if (deferredRemains - || failed(verifyScheduledResultsLive(materializationResult.materializedSchedules)) - || failed(verifyScheduledSpatialInvariants(entryFunc))) { - moduleOp.emitError("realized Spatial communication verification failed"); - signalPassFailure(); - return; - } - - dumpModule(moduleOp, "spatial4_scheduled", /*assumeVerified=*/true); - if (shouldExportSpatialDataflowStage(exportStage, SpatialDataflowExportStage::Spatial4) - && failed(exportSpatialDataflowCsvScheduled( - entryFunc, materializationResult.materializedSchedules, - "spatial4_scheduled", "spatial4"))) + auto prepared = prepareSpatialSchedules(moduleOp, target, pipelineStages, false); + if (failed(prepared) || failed(realizeSpatialSchedule(moduleOp, + prepared->candidates.front(), target, pipelineStages, exportStage, + periodicDiagnostics ? &prepared->graph : nullptr))) signalPassFailure(); } @@ -178,11 +76,154 @@ private: SchedulingTarget target; SpatialDataflowExportStage exportStage = SpatialDataflowExportStage::None; size_t pipelineStages = 1; + bool periodicDiagnostics = false; bool hasTarget = false; }; } // namespace +FailureOr prepareSpatialSchedules(ModuleOp moduleOp, + const SchedulingTarget &target, size_t pipelineStages, bool collectCandidates) { + if (!hasValidTarget(target) || !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(); + auto entryFunc = *entry; + PreparedSpatialSchedules prepared; + auto &scheduledGraph = prepared.graph; + MergeScheduleResult schedule; + for (;;) { + auto graph = buildComputeGraph(entryFunc, target); + if (failed(graph)) { + return failure(); + } + scheduledGraph = std::move(*graph); + std::string pipelineError; + if (pipelineStages > 1) { + FailureOr preparation = + preparePipelineWorkload( + scheduledGraph, pipelineStages, target, pipelineError); + if (failed(preparation)) { + moduleOp.emitError() << pipelineError; + return failure(); + } + if (*preparation == PipelineWorkloadPreparation::Changed) + continue; + } + if (pipelineStages == 1) { + MergeSchedulingAnalysis analysis( + scheduledGraph, target, entryFunc.getContext(), pipelineError); + if (failed(analysis.getResult())) { + moduleOp.emitError() << pipelineError; + return failure(); + } + schedule = std::move(*analysis.getResult()); + break; + } + if (succeeded(runPipelineScheduler( + scheduledGraph, schedule, pipelineStages, target, pipelineError, collectCandidates ? &prepared.candidates : nullptr))) + break; + moduleOp.emitError() << pipelineError; + return failure(); + } + + if (prepared.candidates.empty()) prepared.candidates.push_back(std::move(schedule)); + return prepared; +} + +LogicalResult realizeSpatialSchedule(ModuleOp moduleOp, + const MergeScheduleResult &schedule, const SchedulingTarget &target, + size_t pipelineStages, SpatialDataflowExportStage exportStage, + const ComputeGraph *costGraph, std::string *scheduleReport) { + auto entry = requireEntry(moduleOp); + if (failed(entry)) return failure(); + auto entryFunc = *entry; + PatternRewriter rewriter(moduleOp.getContext()); + FailureOr materialization = + materializeScheduledCompute(entryFunc, schedule, rewriter); + if (failed(materialization)) { + return failure(); + } + ScheduledComputeMaterializationResult materializationResult = std::move(*materialization); + + if (failed(verifyMaterializedScheduleMapping( + entryFunc, schedule, materializationResult.peftClassPlans, + materializationResult.graphComputeToBlockMap, + materializationResult.materializedSchedules)) + || failed(verifyDeferredTransferPhase1Invariants(entryFunc)) + || failed(verifyScheduledMaterializationRecords(materializationResult.materializedSchedules)) + || failed(verifyScheduledSpatialInvariants(entryFunc))) { + moduleOp.emitError("scheduled Spatial phase verification failed"); + return failure(); + } + + dumpModule(moduleOp, "spatial3_scheduled_no_comm", /*assumeVerified=*/true); + if (shouldExportSpatialDataflowStage(exportStage, SpatialDataflowExportStage::Spatial3) + && failed(exportSpatialDataflowCsvScheduled( + entryFunc, materializationResult.materializedSchedules, + "spatial3_scheduled_no_comm", "spatial3"))) { + return failure(); + } + if (scheduleReport) { + llvm::raw_string_ostream os(*scheduleReport); + writeScheduledComputeReport(os, moduleOp, entryFunc, schedule, + materializationResult.peftClassPlans, materializationResult.materializedSchedules); + } else { + dumpScheduledComputeReport( + moduleOp, entryFunc, schedule, materializationResult.peftClassPlans, + materializationResult.materializedSchedules); + } + + if (failed(realizeDeferredCommunication( + entryFunc, materializationResult, target, pipelineStages, costGraph))) { + moduleOp.emitError("Spatial communication realization failed"); + return failure(); + } + + bool deferredRemains = false; + entryFunc.walk([&](SpatDeferredCommunicationOp deferred) { + if (deferredRemains) + return; + deferred.emitOpError("realized Spatial graph still contains deferred communication"); + deferredRemains = true; + }); + if (deferredRemains + || failed(verifyScheduledResultsLive(materializationResult.materializedSchedules)) + || failed(verifyScheduledSpatialInvariants(entryFunc))) { + moduleOp.emitError("realized Spatial communication verification failed"); + return failure(); + } + + dumpModule(moduleOp, "spatial4_scheduled", /*assumeVerified=*/true); + if (shouldExportSpatialDataflowStage(exportStage, SpatialDataflowExportStage::Spatial4) + && failed(exportSpatialDataflowCsvScheduled( + entryFunc, materializationResult.materializedSchedules, + "spatial4_scheduled", "spatial4"))) + return failure(); + return success(); +} + +MergeScheduleResult remapSpatialSchedule(const MergeScheduleResult &schedule, + const IRMapping &mapping) { + MergeScheduleResult result = schedule; + auto remap = [&](ComputeInstance instance) { + instance.op = mapping.lookup(instance.op); + return instance; + }; + for (auto &instance : result.dominanceOrderCompute) instance = remap(instance); + result.computeToCpuMap.clear(); + result.computeToCpuSlotMap.clear(); + result.computeToAestMap.clear(); + result.isLastComputeOfCpu.clear(); + for (auto [instance, cpu] : schedule.computeToCpuMap) result.computeToCpuMap[remap(instance)] = cpu; + for (auto [instance, slot] : schedule.computeToCpuSlotMap) result.computeToCpuSlotMap[remap(instance)] = slot; + for (auto [instance, start] : schedule.computeToAestMap) result.computeToAestMap[remap(instance)] = start; + for (auto instance : schedule.isLastComputeOfCpu) result.isLastComputeOfCpu.insert(remap(instance)); + for (auto &[cpu, instance] : result.cpuToLastComputeMap) instance = remap(instance); + return result; +} + std::unique_ptr createScheduleAndRealizeSpatialPass() { return std::make_unique(); } @@ -190,9 +231,9 @@ std::unique_ptr createScheduleAndRealizeSpatialPass() { std::unique_ptr createScheduleAndRealizeSpatialPass( const SchedulingTarget& target, SpatialDataflowExportStage exportStage, - size_t pipelineStages) { + size_t pipelineStages, bool periodicDiagnostics) { return std::make_unique( - target, exportStage, pipelineStages); + target, exportStage, pipelineStages, periodicDiagnostics); } } // namespace spatial diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp new file mode 100644 index 0000000..6626e2a --- /dev/null +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp @@ -0,0 +1,23 @@ +#pragma once +#include "Scheduling/ComputeGraph.hpp" +#include "Scheduling/MergeSchedule.hpp" +#include "SpatialDataflowCsvExporter.hpp" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/IRMapping.h" + +namespace onnx_mlir::spatial { +struct PreparedSpatialSchedules { + ComputeGraph graph; + std::vector candidates; +}; + +mlir::FailureOr prepareSpatialSchedules( + mlir::ModuleOp module, const SchedulingTarget &target, + size_t pipelineStages, bool collectCandidates); +mlir::LogicalResult realizeSpatialSchedule(mlir::ModuleOp module, + const MergeScheduleResult &schedule, const SchedulingTarget &target, + size_t pipelineStages, SpatialDataflowExportStage exportStage, + const ComputeGraph *costGraph = nullptr, std::string *scheduleReport = nullptr); +MergeScheduleResult remapSpatialSchedule(const MergeScheduleResult &schedule, + const mlir::IRMapping &mapping); +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp index bc6433d..5ba7f26 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp @@ -184,29 +184,7 @@ struct SchedulerCostModel { } Cost getMatrixComputeLatency(Cost inputBitwidth) const { - Cost inputTimes = ceilDiv(inputBitwidth, target.matrixInputResolutionBits); - Cost inputPasses = ceilDiv(target.matrixRows, target.matrixInputParallelism); - Cost outputPasses = ceilDiv(target.matrixColumns, target.matrixOutputParallelism); - Cost readCycles = ceilDiv(target.matrixReadLatencyNs, target.matrixPeriodNs); - Cost frontStage = target.matrixInputBufferLatencyCycles + target.matrixInputLatencyCycles - + readCycles + target.matrixSampleLatencyCycles; - Cost backPipe = std::max(target.matrixOutputLatencyCycles, - checkedAdd(target.matrixShiftLatencyCycles, - target.matrixBufferLatencyCycles)); - Cost backStage = checkedAdd( - checkedAdd(target.matrixOutputLatencyCycles, target.matrixShiftLatencyCycles), - target.matrixBufferLatencyCycles); - backStage = checkedAdd(backStage, checkedMultiply(outputPasses - 1, backPipe)); - Cost totalTimes = checkedMultiply(inputTimes, inputPasses); - if (!target.matrixPipeline) - return checkedMultiply( - checkedMultiply(checkedAdd(frontStage, backStage), totalTimes), - target.matrixPeriodNs); - Cost stagePipe = std::max(frontStage, backStage); - return checkedMultiply( - checkedAdd(checkedAdd(frontStage, backStage), - checkedMultiply(totalTimes - 1, stagePipe)), - target.matrixPeriodNs); + return target.getMatrixComputeLatencyNs(inputBitwidth); } Cost getWvmmCost(Type inputType, Type outputType) const { diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.cpp new file mode 100644 index 0000000..a583c03 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.cpp @@ -0,0 +1,141 @@ +#include "PeriodicDependencies.hpp" + +#include + +namespace onnx_mlir::spatial { + +mlir::FailureOr analyzePeriodicDependencies( + size_t nodeCount, llvm::ArrayRef edges, + std::string &error) { + std::vector> outgoing(nodeCount); + std::vector indegree(nodeCount), recurrent; + Cost upper = 0; + for (auto [index, edge] : llvm::enumerate(edges)) { + if (edge.source >= nodeCount || edge.target >= nodeCount) { + error = "periodic dependency references an invalid event"; + return mlir::failure(); + } + if (edge.delayNs > std::numeric_limits::max() - upper) { + error = "periodic delay sum exceeds the timing range"; + return mlir::failure(); + } + upper += edge.delayNs; + if (edge.distance) { + recurrent.push_back(index); + } else { + outgoing[edge.source].push_back(index); + ++indegree[edge.target]; + } + } + 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]]) + if (!--indegree[edges[index].target]) + order.push_back(edges[index].target); + if (order.size() != nodeCount) { + error = "periodic distance-zero dependencies contain a cycle"; + return mlir::failure(); + } + + PeriodicDependencyEstimate result; + std::vector span(nodeCount); + for (size_t node : order) { + result.withinGenerationSpanNs = + std::max(result.withinGenerationSpanNs, span[node]); + for (size_t index : outgoing[node]) { + const auto &edge = edges[index]; + span[edge.target] = std::max(span[edge.target], + checkedAdd(span[node], edge.delayNs)); + } + } + if (recurrent.empty() || !upper) + return result; + + const size_t absent = std::numeric_limits::max(); + std::vector potential(nodeCount); + std::vector predecessor(nodeCount, absent); + auto infeasible = [&](double interval) { + std::fill(potential.begin(), potential.end(), 0); + std::fill(predecessor.begin(), predecessor.end(), absent); + auto relax = [&](size_t index) { + const auto &edge = edges[index]; + double value = potential[edge.source] + edge.delayNs + - interval * edge.distance; + if (value <= potential[edge.target] + 1e-7) + return false; + potential[edge.target] = value; + predecessor[edge.target] = index; + return true; + }; + for (size_t node : order) + for (size_t index : outgoing[node]) + 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) { + ++result.relaxationRounds; + 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); + if (!changed) + return false; + } + return true; + }; + result.recurrenceUpperNs = upper; + for (unsigned iteration = 0; + iteration < 64 && result.recurrenceUpperNs - result.recurrenceLowerNs > 1; + ++iteration) { + double middle = (result.recurrenceLowerNs + result.recurrenceUpperNs) / 2; + if (middle == result.recurrenceLowerNs || middle == result.recurrenceUpperNs) + break; + if (infeasible(middle)) + result.recurrenceLowerNs = middle; + else + result.recurrenceUpperNs = middle; + } + if (result.recurrenceUpperNs - result.recurrenceLowerNs > 1) { + error = "periodic timing exceeds one-nanosecond floating-point precision"; + return mlir::failure(); + } + if (!infeasible(result.recurrenceLowerNs)) + return result; + + // Recover witness cycles from the last longest-path predecessor forest. + std::vector seen(nodeCount, absent); + double best = 0; + for (size_t start = 0; start < nodeCount; ++start) { + size_t node = start; + while (seen[node] == absent && predecessor[node] != absent) { + seen[node] = start; + node = edges[predecessor[node]].source; + } + if (seen[node] != start || predecessor[node] == absent) + continue; + size_t first = node; + Cost delay = 0, distance = 0; + std::vector cycle; + do { + size_t index = predecessor[node]; + cycle.push_back(index); + delay = checkedAdd(delay, edges[index].delayNs); + distance = checkedAdd(distance, static_cast(edges[index].distance)); + node = edges[index].source; + } while (node != first); + double ratio = distance ? static_cast(delay) / distance : 0; + if (ratio > best) { + best = ratio; + std::reverse(cycle.begin(), cycle.end()); + result.criticalCycle = std::move(cycle); + } + } + return result; +} + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.hpp new file mode 100644 index 0000000..0cb103f --- /dev/null +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include "Utils.hpp" +#include "mlir/Support/LogicalResult.h" +#include "llvm/ADT/ArrayRef.h" + +#include +#include + +namespace onnx_mlir::spatial { + +// t(target, k) >= t(source, k - distance) + delayNs. +struct PeriodicDependency { + size_t source = 0; + size_t target = 0; + Cost delayNs = 0; + unsigned distance = 0; +}; + +struct PeriodicDependencyEstimate { + double recurrenceLowerNs = 0; + double recurrenceUpperNs = 0; + Cost withinGenerationSpanNs = 0; + std::vector criticalCycle; + size_t relaxationRounds = 0; +}; + +// Requires an acyclic distance-zero graph. Brackets the maximum cycle ratio +// within one nanosecond; storage is linear, without generation unrolling. +mlir::FailureOr analyzePeriodicDependencies( + size_t nodeCount, llvm::ArrayRef edges, + std::string &error); + +} // 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 459133d..d7a2e86 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp @@ -1,6 +1,7 @@ #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallBitVector.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/Support/raw_ostream.h" #include #include @@ -1079,10 +1080,20 @@ static LogicalResult packPipelineAssignment(const ComputeGraph &graph, } +struct PipelineCandidateCapture { + bool enabled = false; + std::vector packed; + + void consider(const MergeScheduleResult &candidate) { + if (enabled) + packed.push_back(candidate); + } +}; + static LogicalResult assignPipelineCores(const ComputeGraph &graph, MergeScheduleResult &schedule, size_t pipelineStages, const SchedulingTarget &physicalTarget, std::string &error, bool memoryAware, - bool sharePublications = true) { + bool sharePublications = true, PipelineCandidateCapture *capture = nullptr) { PipelineCoreLayout balancedLayout(physicalTarget.processorCount, pipelineStages); if (!balancedLayout.isValid()) { error = "pipeline scheduling received an incompatible physical core layout"; @@ -1098,6 +1109,8 @@ static LogicalResult assignPipelineCores(const ComputeGraph &graph, if (failed(packPipelineAssignment(graph, taskModel, assignment, candidate, pipelineStages, physicalTarget, error))) continue; + if (capture) + capture->consider(candidate); auto maximumStageSpan = [](const MergeScheduleResult &result) { return result.costs.stageSpanNs.empty() ? Cost{0} : *llvm::max_element(result.costs.stageSpanNs); @@ -1162,27 +1175,68 @@ mlir::LogicalResult runPipelineScheduler(const ComputeGraph& graph, MergeScheduleResult& schedule, size_t pipelineStages, const SchedulingTarget& physicalTarget, - std::string& error) { + std::string& error, + std::vector *candidates) { PipelineCoreLayout pipelineLayout( physicalTarget.processorCount, pipelineStages); if (pipelineStages <= 1 || !pipelineLayout.isValid()) { error = "pipeline scheduling requires a valid balanced physical core layout"; return mlir::failure(); } - if (failed(assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, error, false, false))) + PipelineCandidateCapture capture{candidates != nullptr, {}}; + if (failed(assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, error, false, false, &capture))) return failure(); - if (!physicalTarget.sharedMemoryWidthBytes || graph.publicationTransfers.empty()) - return success(); std::string candidateError; // Retain the calibrated phase-1 winner before evaluating sharing-aware cuts. - (void)assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, candidateError, true, false); - std::vector taskStages; - for (const auto &node : graph.nodes) - taskStages.push_back(schedule.processorStages[ - schedule.computeToCpuMap.lookup(node.instance)]); - schedule.costs.boundaryMemory = getPipelineBoundaryMemoryDemand( - graph, taskStages, physicalTarget); - (void)assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, candidateError, true, true); + if (physicalTarget.sharedMemoryWidthBytes && !graph.publicationTransfers.empty()) { + (void)assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, candidateError, true, false, &capture); + std::vector taskStages; + for (const auto &node : graph.nodes) + taskStages.push_back(schedule.processorStages[ + schedule.computeToCpuMap.lookup(node.instance)]); + schedule.costs.boundaryMemory = getPipelineBoundaryMemoryDemand( + graph, taskStages, physicalTarget); + (void)assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, candidateError, true, true, &capture); + } + if (candidates) { + auto activeStages = [&](const MergeScheduleResult &result) { + llvm::SmallBitVector stages(pipelineStages); + for (const auto &node : graph.nodes) + stages.set(result.processorStages[result.computeToCpuMap.lookup(node.instance)]); + return stages; + }; + auto samePlacement = [&](const MergeScheduleResult &a, const MergeScheduleResult &b) { + for (const auto &node : graph.nodes) { + size_t cpu = a.computeToCpuMap.lookup(node.instance); + if (cpu != b.computeToCpuMap.lookup(node.instance) + || a.computeToCpuSlotMap.lookup(node.instance) != b.computeToCpuSlotMap.lookup(node.instance) + || a.processorStages[cpu] != b.processorStages[cpu]) + return false; + } + return true; + }; + auto originalStages = activeStages(schedule); + SmallVector alternatives; + for (auto [index, candidate] : llvm::enumerate(capture.packed)) { + if (activeStages(candidate) != originalStages || samePlacement(candidate, schedule)) + continue; + if (llvm::none_of(alternatives, [&](size_t other) { + return samePlacement(candidate, capture.packed[other]); + })) + alternatives.push_back(index); + } + candidates->push_back(schedule); + for (size_t index : alternatives) { + auto &candidate = capture.packed[index]; + std::vector taskStages; + for (const auto &node : graph.nodes) + taskStages.push_back(candidate.processorStages[ + candidate.computeToCpuMap.lookup(node.instance)]); + candidate.costs.boundaryMemory = getPipelineBoundaryMemoryDemand( + graph, taskStages, physicalTarget); + candidates->push_back(std::move(candidate)); + } + } return success(); } 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 314d3c1..ce90a0f 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp @@ -124,7 +124,8 @@ mlir::LogicalResult runPipelineScheduler(const ComputeGraph& graph, MergeScheduleResult& schedule, size_t pipelineStages, const SchedulingTarget& physicalTarget, - std::string& error); + std::string& error, + std::vector *candidates = nullptr); enum class PipelineWorkloadPreparation { Ready, diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.cpp index dee00d2..572b03f 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.cpp @@ -1,10 +1,40 @@ #include "ScheduleEvaluation.hpp" #include "MergeSchedule.hpp" +#include #include namespace onnx_mlir::spatial { +bool WholeScheduleCost::eligible() const { + return primary && alternative && std::isfinite(*primary) && std::isfinite(*alternative) + && *primary > 0 && *alternative > 0 && relativeUncertainty + && std::isfinite(*relativeUncertainty) && *relativeUncertainty >= 0 + && *relativeUncertainty < 1; +} + +size_t selectPipelineCandidate(llvm::ArrayRef scores, std::string &reason) { + if (scores.empty() || !scores.front().eligible()) { + reason = "original estimate unavailable"; + return 0; + } + size_t best = 0; + auto upper = [&](size_t i) { + return std::max(*scores[i].primary, *scores[i].alternative) + / (1 - *scores[i].relativeUncertainty); + }; + for (size_t i = 1; i < scores.size(); ++i) + if (scores[i].eligible() && upper(i) < upper(best)) best = i; + if (!best) { reason = "original has the lowest conservative estimate"; return 0; } + double originalLower = std::min(*scores.front().primary, *scores.front().alternative); + if (upper(best) >= originalLower / (1 + *scores.front().relativeUncertainty)) { + reason = "candidate and original overlap static prediction uncertainty"; + return 0; + } + reason = "candidate improves on original under both arbitration tie policies and prediction uncertainty"; + return best; +} + void recordScheduleCosts(const ScheduleEvaluation &evaluation, MergeScheduleResult &schedule) { schedule.costs.totalWorkNs = evaluation.totalWork; diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp index 393b088..26b4251 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp @@ -9,6 +9,9 @@ namespace onnx_mlir::spatial { struct MergeScheduleResult; +size_t selectPipelineCandidate(llvm::ArrayRef scores, + std::string &reason); + struct ScheduleEvaluation { std::vector