diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md new file mode 100644 index 0000000000000..a38789c5cf3b7 --- /dev/null +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -0,0 +1,776 @@ +# Contrib Operator FLOP Estimation: Design and Rollout Plan + +**Status:** CPU MatMulNBits runtime count/report preview implemented; GQA and later phases remain proposed. + +For the post-Batch-3 execution sequence, validation gates, CUDA expansion and timing work, see the +[MatMulNBits counter completion plan](matmul_nbits_counter_completion_plan.md). + +## 1. Summary + +Introduce an opt-in diagnostic feature that estimates the mathematical work of a contrib operator invocation. +The minimum operator scope is `com.microsoft::MatMulNBits` and GroupQueryAttention (GQA). Start with the existing +MatMulNBits estimator, bias component, and runtime metadata/report integration; +next add GQA baseline estimates. Route-aware adjustments are optional, incremental refinements, not an exact +executed-work prerequisite for GQA. Keep dispatch, placement, +memory admission, and allocation behavior unchanged throughout the diagnostic rollout. + +This plan builds on the +[ORT Memory Estimates Copilot discussion](https://github.com/copilot/tasks/2cffab6d-eafa-4f8a-8595-590c601ee41c), +including its follow-up rollout and design proposals. The discussion was recovered from session history; its +recommendations are proposals, not evidence that the suggested APIs already exist. + +The key distinction is: + +- **FLOPs:** a count of floating-point mathematical operations for a defined workload. +- **Dense-equivalent operations:** comparable dense arithmetic for a quantized operator; not literal floating-point + instructions executed by that implementation. +- **FLOP/s or equivalent operations/s:** a count divided by a matched measured duration. A graph alone cannot + establish this rate. + +**Target deliverable:** runtime-derived workload estimates for MatMulNBits and GQA in a versioned developer report. +Initialization configures runtime reporting, without generating a workload report. Batch 3 replaces the legacy +initialization-time static export; graph-based helpers remain internal regression/reference utilities. +Throughput and performance-policy changes are separate deliverables. + +The initial core-only preview and the core-plus-bias increment remain independently reviewable and shippable. +Neither alone satisfies the minimum MatMulNBits-and-GQA deliverable. Prioritize useful, convention-based workload +estimates over exhaustive kernel auditing. Add implementation-dependent adjustments only when supported by +execution metadata and documented arithmetic models; an unknown route must not invalidate an available baseline +or be treated as zero extra work. Timing and performance-policy changes remain separate deliverables. + +### Initial implementation + +The internal [MatMulNBits compute estimator](../../onnxruntime/core/framework/matmul_nbits_compute_estimate.h) +provides a checked shape-based count and a static adapter for resolved in-tree graph nodes. Its result is either a +core dense-equivalent count under convention version 1 or an explicit unavailable reason. Rank-one inputs use +`M = 1`; unknown dimensions remain unavailable even alongside zero extents. The adapter reads canonical graph +shapes, not maximum-shape scenarios, and does not validate packed weights or replace normal operator validation. + +The Phase-2 preview added a static JSON report via the existing session configuration API. That preview does not +collect runtime counts or change inference policy. The estimator sources are included by the +existing framework source glob, independently of CPU contrib kernels and CUDA. Keeping the estimator +and report in the framework library avoids a reverse framework-to-provider static-library dependency. + +Increment 2a adds a separate `bias` component: `MN` logical floating-point additions when the sixth schema input +is present, or zero when it is omitted. Known absence does not require concrete input shapes; present bias with +unknown required geometry is unavailable. The pure helper accepts unknown presence and returns +`missing_semantic_value`; a resolved graph identifies presence from its optional input slot. A zero reduction +dimension `K` makes the core count zero but does not make bias additions zero if the output is nonempty. +Core arithmetic and its convention remain unchanged. Runtime invocation counts and route-specific executed FLOPs +are not implemented by this increment. + +### Structured runtime collector: review batches + +The target report uses metadata from actual runtime invocations, not initialization-time workload assumptions. +Batch 3 replaces the legacy initialization-time export with per-run reports. +Metadata transport must be typed and independent of logger severity or text parsing. + +**Batch 1: standalone storage primitive (implemented, not connected to inference).** +The internal [RuntimeWorkloadCollector](../../onnxruntime/core/framework/runtime_workload_collector.h) owns a +run's submitted MatMulNBits records. Each record copies graph/provider identity, node index, actual shapes, +matrix geometry/batch count, bias presence, activation element type, quantization settings, and optional +route/effective compute mode. Session/run identifiers are supplied by the caller; accepted records receive a +unique collector-local invocation ID in commit order, not execution order. Missing route fields remain missing. + +Limits are explicit constructor arguments, with no enabled-by-default collector. Record count and accounted +payload bytes are bounded. The byte budget includes record size, copied string bytes/terminators, and shape +elements; allocator capacity, bookkeeping, caller buffers, and snapshot copies are not a hard-memory guarantee. +Rejected submissions return a diagnostic Status and update saturating drop counters and collection completeness. +Allocation failures are converted to Status in exception-enabled builds; no-exception builds retain the existing +allocation-failure termination behavior. Thread-safe snapshots copy retained records and status without exposing +mutable collector state. Collection completeness describes submissions only, not run success or FLOP coverage. +Callers must keep the collector alive until all submitters finish. + +This batch has no kernel/context/session plumbing, run-ID generation, semantic validation, counts, run-outcome +handling, logger integration, or JSON/file export. No production inference behavior or public API changes. +The full-build-only implementation does not depend on contrib kernels or MLAS. Tests cover ownership, repeated +invocations, missing routes, empty-output metadata, isolation, limits, and concurrent submissions/snapshots. + +**Batch 2: execution integration (implemented).** An internal, typed +`RunOptions::runtime_workload_collector` pointer opts a canonical `InferenceSession::Run` into capture. +The caller owns a fresh collector, supplies explicit record/payload budgets, and keeps it alive until Run +returns (or asynchronous execution finishes). No C API setter, default budgets, file configuration, or +export lifecycle is introduced. For example, internal C++ callers can use: + +```cpp +RuntimeWorkloadCollector collector(RuntimeWorkloadLimits{1000, 1024 * 1024}); +RunOptions options; +options.runtime_workload_collector = &collector; +// session.Run(options, ...); +RuntimeWorkloadSnapshot snapshot; +// Check the inference Status and collector.GetSnapshot(snapshot) independently. +``` + +The canonical Run binds the session's identity and a process-local, monotonically assigned run ID. +Collectors cannot be reused: subsequent/concurrent attempts disable capture for the rejected run, preserve +the original identity/data, and increment `run_reuse_rejections`, marking collection incomplete without +failing inference. Separate collectors isolate concurrent runs of the same or different sessions. +The optional pointer follows the existing run-profiler execution/context plumbing, including If, Loop, +and Scan subgraphs and the GPT beam/greedy-search and T5/Whisper beam-search encoder/decoder subgraphs +(including GPT's initial decoder). A T5 regression test inserts zero-valued MatMulNBits probes into both +encoder and decoder, checks their runtime records, and verifies unchanged generated sequences. +The context owns the recorder; kernels do not store per-run mutable state. +Standalone and constant-folding contexts return no recorder. Transport uses typed metadata, never logs. +The C API and its versioned function table are unchanged; an appended full-build C++ context virtual provides +safe optional lookup rather than assuming every `OpKernelContext` is an internal execution context. + +CPU MatMulNBits submits only after its selected computation succeeds. Invalid inputs, thrown/failed +computation, and terminated executions do not submit successful-invocation records. Empty outputs submit +`EmptyOutput`/`NotExecuted` explicitly. Records include actual activation/output shapes, activation type, +bias presence, bits/block size, executable node index, provider, and an unambiguous graph path formed from +parent node indices and subgraph attributes. Repeated control-flow invocations retain the same node identity +but receive distinct commit-order invocation IDs. Matrix geometry/batch count describe the actual +`MatMulComputeHelper` backend geometry: shared-weight batches may be flattened into M, with batch count one. +Packed LUT, packed QNBit, and unpacked FP32 GEMM are distinguished. Compute mode describes the effective MLAS +interface, including the half-activation int8 path's conversion to SQNBit int8, not merely `accuracy_level`, +activation type, or a claim about every instruction executed inside a backend. + +Snapshots independently expose `NotStarted`, `Running`, `Succeeded`, or `Failed` run state and submission +completeness. A failed run retains successful earlier kernel records, but they are not a complete successful +workload. Limit rejection, allocation failure, and other metadata-capture exceptions are diagnostic only; +they update drop counters/completeness without replacing the inference result. Snapshot failures remain a +separate diagnostic Status. Allocation-failure recovery requires exception-enabled builds, as in Batch 1. +With capture disabled, there is no collector allocation, record/shape/string copying, run-ID assignment, +or collection locking. This is CPU MatMulNBits coverage, not all-provider/all-operator coverage; EP graph +replay does not synthesize records for kernels it does not invoke. + +Focused tests cover storage/lifecycle diagnostics, actual packed int8/FP32 and unpacked routes, a +capability-gated LUT route, activation conversion, dynamic/repeated/concurrent runs, node/session/run +isolation, If/Loop/Scan identity, empty outputs, rank-one/two-bit geometry, kernel/termination/downstream +failures, bounded collection, collector reuse, and unchanged output values. + +**Batch 3: runtime counts and export (implemented).** Apply the existing pure checked core/bias helpers to each +retained invocation's full activation shape, K, N and bias presence. The full shape already includes batching: +backend M/batch_count are provenance, not an additional multiplier. Missing route information does not invalidate +the baseline. No route-adjusted arithmetic, device timing, hardware calibration, GQA or CUDA capture is introduced. + +### Generating a review artifact + +Set `session.compute_estimate_report_path` to a UTF-8 **output prefix** before creating a session: + +```python +import onnxruntime as ort + +options = ort.SessionOptions() +options.add_session_config_entry("session.compute_estimate_report_path", "compute_estimates") +options.add_session_config_entry("session.compute_estimate_max_records", "10000") +options.add_session_config_entry("session.compute_estimate_max_payload_bytes", "16777216") +session = ort.InferenceSession("model.onnx", sess_options=options, providers=["CPUExecutionProvider"]) +# Supply the model's actual feeds; each run writes its own report. +outputs = session.run(None, feeds) +``` + +The equivalent C++ configuration is: + +```cpp +Ort::SessionOptions options; +options.AddConfigEntry("session.compute_estimate_report_path", "compute_estimates"); +``` + +These examples require an ORT build containing this implementation, not an older installed wheel. +Initialization validates configuration but writes no report. After each canonical `Run`, including failed runs, +the session closes collection, snapshots it and synchronously exports outside the kernel path. A report is available +when `Run` returns. Export adds CPU/file-I/O latency; it performs no additional device synchronization. +For asynchronous callers, wait for completion before inspecting output. + +Filenames have the form +`___session-_run-.json`. +The process token is initialized once at first export; session/run IDs distinguish concurrent sessions/runs. +Relative prefixes use the process working directory; parent directories must already exist. Choose a fresh prefix +for independent application launches. Reports contain actual tensor shapes and graph/node identities, but no tensor +contents; review metadata before sharing. + +The two optional budgets default to 10,000 records and 16 MiB accounted payload per run. They accept nonnegative +decimal integers representable as size_t; invalid values fail initialization. Zero intentionally retains no records. +These are collector accounting limits, not hard peak-memory limits: snapshots, JSON construction and serialization +consume additional memory. Every run creates a file; there is no disk-retention policy. Use a bounded diagnostic +window and manage report retention externally. + +If an internal caller also supplies `RunOptions::runtime_workload_collector`, the supplied collector is used for +both capture and export with its own limits. It must be fresh for every run. Reuse warns and skips capture/export +for that run rather than exporting a stale snapshot. Internal callers may also explicitly call the snapshot-based +`BuildComputeEstimateReport` or `WriteComputeEstimateReport` without enabling automatic export. + +Generation/open/write failures warn without changing inference status or outputs; they do not produce a replacement +success report. An interrupted or failed write can leave a partial file; only successfully parsed reports should be +consumed. Minimal builds and compile-only sessions warn that automatic export is unsupported. Full contrib-disabled +builds can export empty coverage reports with `estimator_support: disabled`. Unset/empty prefix performs no +automatic collection, report traversal or file I/O; an explicitly supplied internal collector still works. + +**Migration:** the old `session.compute_estimate_report_path` literal-file/static behavior is intentionally replaced. +Consumers must discover suffixed files and parse schema version 3 instead of expecting an initialization-time +version-2 `nodes`/`graphs` document. The graph-based version-2 builders remain internal only. + +### Runtime report contract (schema version 3) + +- `report_kind: runtime_workload`, session/run IDs and `run_state` identify the run. No timing/rate is implied. +- `invocations` contains only retained records from successful instrumented CPU MatMulNBits invocations, including + repeated executed subgraphs. Unexecuted branches and static nodes do not contribute. +- Each record includes graph/node/invocation/provider identity, actual shapes, backend geometry, bias presence, + activation type, quantization settings, and optional route/effective-mode provenance. Invocation IDs reflect commit + order, not execution order. No hardware instruction count or tiling metadata is claimed. +- Core uses `dense_equivalent_ops` / `core_arithmetic_v1`; bias uses `floating_point_ops` / `bias_add_v1`. + For A=[2,3,64], N=32 with bias: core=24576 and bias=192. Bias is never added to the core subtotal. +- Per-invocation availability is `available`, `partial` or `unavailable`. Missing estimates have explicit reasons and + omit `operation_count`; valid zero counts remain available. Counts are decimal strings to preserve uint64 precision. +- `component_summaries` aggregate retained executed records across graph scopes. Overflow omits the subtotal count. + `complete_for_captured_component_scope` requires successful Run, complete collection, available component estimates + and no subtotal overflow. A known subtotal on a failed run describes successful invocations before failure only. +- `collection` exposes the actual record/payload budgets, accounted retained bytes and + record/payload/allocation/metadata/reuse diagnostics independently + of run outcome. Missing records are never assigned zero work. +- `coverage_scope: instrumented_cpu_matmul_nbits_invocations_only`, `whole_run_work_complete: false` and + `whole_operator_work_complete: false` prevent captured subtotals from masquerading as all-model/all-provider counts. + An empty report is not evidence that a workload did no work. Unsupported operators, CUDA kernels, fused/opaque EP + partitions and graph replay are not instrumented, and their skipped invocation counts are not known. +- `route_adjustment_applied: false`; dequantization, packing, conversions, padding and route-specific extras are + explicitly excluded. Route names do not change the baseline formulas. + +CPU Debug correctness tests cover counts, partial/overflow/drop diagnostics, automatic per-run export, concurrent +runs, dynamic shapes, UTF-8 paths, file errors, disabled behavior, explicit collector coexistence, and executed +If/Loop/Scan/transformer records. This is a developer preview, not production overhead qualification or evidence of +CUDA/plugin/minimal runtime coverage; those need their own build/runtime validation and performance matrix. + +Run the targeted tests after rebuilding: + +```text +onnxruntime_provider_test --gtest_filter="ComputeEstimateReportTest.*:RuntimeComputeEstimateReportTest.*:MatMulNBitsComputeEstimate*:MatMulNBitsBiasEstimate*" +onnxruntime_test_all --gtest_filter="RuntimeWorkloadCollectorTest.*:RuntimeWorkloadExecutionTest.*" +``` + +## 2. Motivation and boundaries + +The feature should help users and ORT developers: + +1. Quantify workload growth across batch sizes, input lengths, and model variants. +2. Identify expensive supported operators while making missing coverage visible. +3. Compare latency and effective throughput for equivalent workloads. +4. Build evidence for later route-selection and CPU/GPU-placement experiments. + +It does **not** initially provide hardware instruction counts, predicted latency, total memory traffic, +universal contrib-op coverage, symbolic-expression algebra, training support, or automatic optimization decisions. +Do not interpret a high operation count as proof that an operator is compute-bound. + +## 3. Relationship to existing memory estimation + +The [constrained-environment roadmap](future_directions_constrained_env.md) and +[tracking issue #29775](https://github.com/microsoft/onnxruntime/issues/29775) provide useful architectural precedents: + +| Existing surface | Reusable lesson | Boundary for compute estimation | +| --- | --- | --- | +| [Level1MemoryEstimate](../../include/onnxruntime/core/framework/level1_memory_estimate.h) | Distinguish quantities with different meanings and lifetimes | Keep compute counts in a separate result; do not add them to byte budgets | +| [MatMulNBits memory adapter](../../onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_workspace_estimate.h) | Separate graph parsing from checked dimension arithmetic | Logical counts should not require CUDA device properties or CUDA build flags | +| [WorkspaceInputShape](../../include/onnxruntime/core/framework/workspace_input_shape.h) | Preserve positional inputs, missing optionals, partial shapes, and zero extents | It does not carry shape provenance; a maximum-shape hint is not a concrete invocation | +| [DeclareWorkspaceRequirements](../../include/onnxruntime/core/framework/op_kernel.h) | Use metadata without requiring live tensor allocations | Do not introduce a new kernel virtual method just to implement a reporting pilot | +| [Attention workspace roadmap](attention_workspace_estimation.md) | Validate workload geometry and distinguish runtime routes from AOT scenarios | Workspace bytes do not determine operation count, memory traffic, or latency | +| [Node profiling](../../onnxruntime/core/framework/sequential_executor.cc) | Existing events include node identity, provider, and shape information | An event named `kernel_time` is not automatically a device-execution interval | + +The eventual planner would consume separate inputs: + +```text +memory requirements + allocation lifetimes -> feasible candidates +measured/calibrated performance costs -> preference among feasible candidates +``` + +The count/reporting pilot must not depend on completion of memory-policy or preallocation work. + +## 4. Counting contract + +### 4.1 Required semantics + +- Count a multiply-add as **two operations**, including when implemented as one fused instruction. +- Version the counting convention independently from the report schema. +- Declare the metric and included components for every estimate. +- Interpret "exact" as exact under that convention and component scope, not exact hardware execution. +- Return unavailable with a machine-readable reason for unsupported operators, unknown required metadata, + invalid metadata, or arithmetic overflow. Never substitute zero for an unavailable estimate. +- Accept a legitimate empty workload as zero only when the required metadata establishes its validity. +- Keep static concrete shapes, runtime invocation shapes, and declared workload scenarios distinguishable. +- Do not call maximum-shape hints proven upper bounds unless the estimator establishes that guarantee. +- Keep floating-point counts and dense-equivalent counts in separate subtotals. + +Malformed model inputs remain subject to normal ORT validation. A diagnostic failure must not silently change +inference results or bypass model validation. Unsupported estimation is a reportable condition, not an inference +failure; allocation, serialization, and other operational errors must use standard error reporting. + +### 4.2 Pilot formula: MatMulNBits + +Use the operator's logical `K` and `N` attributes, not the packed weight storage dimensions. For a supported valid +input `A` with shape `[..., K]`: + +```text +M = product of A's leading dimensions +core dense-equivalent operations = 2 * M * N * K +``` + +The [MatMulNBits schema](../../onnxruntime/core/graph/contrib_ops/contrib_defs.cc) defines packed weights, +dequantization, and an optional bias. The pilot counts **only core matrix multiplication**: + +- Exclude dequantization, unpacking, prepacking, bias, conversions, padding, and implementation-specific extra work. +- Report those exclusions even when bias or other optional inputs are present. +- Do not scale the logical count down by the quantization bit width. +- Preserve the same core count across providers and kernel routes for the same logical invocation. +- Validate rank and agreement between the last input dimension and `K`; handle rank-one inputs according to the + validated schema/kernel contract rather than assuming they are supported. +- Use checked integer multiplication; reject invalid negative dimensions before unsigned conversion. + +Example: `A = [2, 3, 64]`, `K = 64`, `N = 32` gives `M = 6` and **24,576 dense-equivalent operations**. +This is not a complete count of all work performed by `MatMulNBits`. + +### 4.3 Ancillary components and counting categories + +Core versus ancillary describes an operation's role, not whether it is floating-point work. Floating-point bias +addition and dequantization arithmetic can contribute FLOPs; bitwise unpacking does not. Report named components +alongside the stable core metric rather than silently broadening the meaning of `2MNK`. + +| Component | Initial counting approach | Required basis and limitations | +| --- | --- | --- | +| Matrix-product core | `2MNK` dense-equivalent operations | Preserve the existing convention; not literal executed floating-point instructions | +| Optional bias | `MN` logical floating-point additions when bias is present | Known output geometry and optional-input presence; zero for absent bias, unavailable when presence or geometry is unknown | +| Logical weight reconstruction | Under an explicit scalar convention, `KN` subtractions and `KN` multiplications for `(q - zero_point) * scale` | One reconstruction of logical weights only; not a claim about per-invocation execution | +| Executed dequantization | Count arithmetic over the elements actually processed, including repetition | Validated route, dtype, padding, reuse, and preparation frequency; otherwise unavailable | +| Unpacking | Route-specific bitwise/integer operations where justified | Keep separate from FLOPs; do not assume a universal instruction count per quantized weight | +| Conversions | Number of elements converted, with source and destination types | Separate category, not an arithmetic FLOP count | +| Layout transforms | Logical elements or bytes moved | Not measured DRAM traffic; caches and reuse are outside this count | +| Padding and recomputation | Additional route-specific arithmetic, classified by arithmetic type | Count only work beyond the corresponding baseline; avoid overlap with other components | +| Prepacking | Separately scoped initialization or first-use work | Never charge automatically to every steady-state invocation | + +Logical reconstruction is an optional workload scenario, not an executed-work subtotal. A kernel may reconstruct +weights once and cache them, repeat reconstruction across tiles, or fold the arithmetic into another computation. +An integer subtraction or an optimized-away operation must not be reported as an executed floating-point operation. +Explicitly document the treatment of symmetric zero points, fused operations, and eliminated arithmetic. + +For every added component, retain: + +- metric/unit, included work, and counting-convention version; +- basis: logical semantics, declared scenario, or validated selected route; +- phase: initialization, first use, or per invocation; +- processed geometry and, where needed, repetition/reuse assumptions; +- available count or an explicit unavailable reason. + +Keep dense-equivalent operations, floating-point arithmetic, integer/bitwise operations, conversions, and logical +bytes moved in separate categories. Only combine non-overlapping components with compatible units, conventions, +bases, and phases. In particular, do not present core dense-equivalent operations plus ancillary FLOPs as literal +total executed FLOPs. Known-component subtotals must retain exclusions and incomplete-coverage indicators. + +Adding components improves coverage; validated route attribution improves fidelity to execution. Neither makes +the quantized core a hardware instruction count or predicts latency without measurement. + +### 4.4 Baseline estimates and optional route-aware refinements + +Provide two distinct results rather than progressively relabeling the existing core count: + +- **Logical FLOP-equivalent workload:** the dense-equivalent core plus explicitly declared, non-overlapping + logical floating-point components under a common workload convention. For the example above, core plus a + present bias is `24576 + 192 = 24768` FLOP-equivalent operations. This scope excludes reconstruction unless + a separate reconstruction scenario is declared; it is not a complete executed-work total. +- **Estimated executed FLOPs:** floating-point arithmetic attributed to a validated implementation route for a + specified workload and lifecycle phase. Integer/bitwise work, conversions, and bytes moved remain separate. + An FMA counts as two FLOPs; this is an arithmetic estimate, not a hardware instruction count. + +The primary deliverable is the logical baseline for MatMulNBits and GQA, with explicit component scope and +assumptions. It does not require instruction-exact counts. A route-adjusted arithmetic estimate is a separate, +optional result: record its model/convention, evidence, processed geometry, and limitations. Call it estimated +executed FLOPs only when the model accounts for the floating-point arithmetic attributed to that route and phase. +Do not sum a baseline component with its adjusted replacement. + +Route-family refinements in Increment 3a use focused implementation verification rather than requiring an +exhaustive audit of every kernel before baseline delivery: + +1. Select a bounded route, initially a CPU path exercised by the pilot tests, and confirm the actual kernel + selected. Record relevant algorithm, dtypes, quantization settings, and preparation state. Reuse actual route + decisions rather than introducing a parallel dispatch predictor. +2. Verify count-affecting behavior, including library calls where necessary. Establish floating-point versus integer core arithmetic, + retained/eliminated reconstruction, bias/epilogues, padded processing, repetitions, and additional reductions. + Do not treat `2MNK` as executed FLOPs when the selected route performs integer dot products. +3. Establish initialization, first-use, and per-invocation lifetimes, including cached/shared preparation. + Distinguish cold and warm invocations; do not charge shared or one-time work repeatedly. +4. Implement checked formulas from validated processed geometry and execution frequency. Count padding and + repetition within the affected component, without adding them again as overlapping extra work. +5. Validate the arithmetic model with independent reference cases and route-boundary, padding, reuse, and lifecycle + tests. Share formulas across variants only after confirming their assumptions. Use backend documentation and + deeper inspection where necessary; hardware counters are optional supporting evidence, not a baseline gate. + An opaque library call is not sufficient evidence for its internal arithmetic count. +6. Report logical and executed metrics side by side with route identity, phase, convention, precision, assumptions, + and coverage. Runtime shapes alone do not establish an executed-FLOP count. + +For one invocation, the executed estimate sums only validated floating-point core, reconstruction, epilogue, +and other arithmetic performed for that invocation. If any applicable component is unknown, expose a +known-component subtotal with explicit gaps, not a complete executed-FLOP total. Enable the executed metric only +for validated routes; adding ancillary logical estimates alone does not qualify a route as supported. If route +metadata or an adjustment model is unavailable, retain the logical baseline and mark the adjustment unavailable. +Fusion alone does not imply fewer FLOPs; it may reduce memory traffic without eliminating arithmetic. + +## 5. Proposed architecture and report + +```text +optimized executable graph OR runtime invocation + | + v +operator metadata adapter + | + v +validated, graph-free workload description + | + v +pure checked estimator + | + v +per-node result -> coverage-aware report -> optional timing correlation +``` + +### Implementation shape + +1. Start with an internal lookup keyed by domain, operator name, and supported schema-version range. +2. Keep formulas independent of `Node`, `GraphViewer`, CUDA, allocation, and execution. +3. Have static and runtime adapters call the same formula. +4. Reuse validated shape/attribute helpers where their contracts fit; do not copy a CUDA-only helper into a second + implementation or make core reporting depend on it. Extract only genuinely shared arithmetic when necessary. +5. Preserve input positions and distinguish an omitted optional input from a present input without shape metadata. +6. Cache only with all count-affecting dimensions, attributes, and semantic values in the key. +7. Defer public C/C++/Python APIs and plugin registration until the internal contract stabilizes. + +The result type should make contradictory states impossible: an available estimate has a count, and an unavailable +estimate has a reason rather than a numeric placeholder. The following is a historical, abbreviated static version-2 report +for a node without bias; +the implementation also emits positional input metadata, estimator support, and per-graph summaries: + +```json +{ + "schema_version": 2, + "convention_version": "core_arithmetic_v1", + "graph_scope": "optimized_executable", + "nodes": [ + { + "graph_path": "main", + "node_index": 12, + "node_name": "projection", + "domain": "com.microsoft", + "op_type": "MatMulNBits", + "opset_version": 1, + "provider": "CUDAExecutionProvider", + "availability": "available", + "whole_operator_work_complete": false, + "components": [ + { + "component": "matmul_core", + "estimate_kind": "exact_under_convention", + "input_basis": "static_concrete", + "basis": "logical_semantics", + "execution_phase": "per_invocation", + "convention_version": "core_arithmetic_v1", + "availability": "available", + "metric": "dense_equivalent_ops", + "scope": "matmul_core", + "operation_count": "24576", + "excluded_components": [ + "dequantization", "bias", "packing", "conversion", "padding", "route_specific_extra_work" + ] + }, + { + "component": "bias", + "scope": "bias", + "metric": "floating_point_ops", + "convention_version": "bias_add_v1", + "estimate_kind": "exact_under_convention", + "input_basis": "optional_input_absent", + "basis": "logical_semantics", + "execution_phase": "per_invocation", + "availability": "available", + "operation_count": "0", + "excluded_components": [ + "matmul_core", "dequantization", "packing", "conversion", "padding", "route_specific_extra_work" + ] + } + ] + } + ], + "summary": { + "reported_node_count": 1, + "estimated_node_count": 1, + "partially_estimated_node_count": 0, + "unavailable_node_count": 0, + "known_dense_equivalent_ops": "24576", + "subtotal_availability": "available", + "complete_for_declared_graph_and_component_scope": true, + "whole_operator_work_complete": false, + "count_scope": "direct_nodes_only", + "per_run_total_available": false, + "graph_path": "main" + } +} +``` + +Serialize counts as decimal strings to preserve integer precision in JSON consumers. An unavailable component omits +`operation_count` and includes a reason such as `unsupported_version`, `unknown_shape`, `invalid_metadata`, `overflow`, +or `estimators_disabled`. Unsupported operators and opaque partitions have unavailable node entries with empty +components. Control-flow nodes use `runtime_execution_required`; their nested graphs are reported separately. +`missing_semantic_value` indicates unknown semantic information such as bias presence in the pure helper. +The static resolved-graph adapter can determine bias presence without runtime data. +If a graph subtotal overflows, omit it and report `subtotal_unavailable_reason: overflow`; never wrap or silently +drop a node's count. A zero known subtotal with unavailable nodes is not a zero-work model. + +Report the optimized graph actually represented by ORT, with graph/subgraph identity and final node identity. +Estimate fused nodes directly or mark them unsupported; never count both a fusion and its removed constituents. +An EP-compiled opaque partition must be unavailable unless a valid attribution mechanism exists. + +Known-operation subtotals are not complete model totals. Node coverage is not compute-weighted coverage. +Do not mix declared-scenario estimates with concrete counts in the same total. Branch and loop bodies cannot be +summed into an exact per-run count without execution information. + +## 6. Phased rollout + +Each phase should be delivered through reviewable PRs. The named owners below are roles to assign, not commitments +from particular contributors. Calendar estimates should follow contract review and a pilot implementation. + +| Phase | Deliverables | Dependencies | Proposed owner | Exit gate | +| --- | --- | --- | --- | --- | +| 0. Contract | Metric conventions, exclusions, result states, report schema, supported build scope | None | Framework maintainer + operator owner | Resolve Q1-Q3 below; approve reference cases and unknown/error semantics | +| 1. Pure pilot | Internal result type and lookup; graph-free MatMulNBits estimator; unit tests | Phase 0 | Framework + quantization owners | Independent expected counts pass; no CUDA dependency; overflow and invalid metadata diagnosed | +| 2. Static reporting / MVP | Opt-in developer report over the optimized executable graph; coverage and provenance | Phase 1 | Framework + tooling owners | End-to-end report matches golden expectations; fusion/partial coverage are honest; inference unchanged | +| 3. Runtime workload metadata and counts | MatMulNBits invocation estimates, actual route/compute-type metadata where exposed, event identity and bounded caching; accommodate GQA metadata needs | Phase 2 and Increment 2a | Execution/profiling + quantization owners | Changing shapes and actual selected-route attribution pass; no extra device synchronization; agreed overhead gate passes | +| 4. Required GQA baseline / later Attention expansion | GQA core estimate, explicit masking convention, prefill/decode/chunked and active-length semantics; other Attention ops later | Phase 2 and approved GQA convention; Phase 3 integration for runtime-only metadata | Attention owners | Independent GQA references, missing-metadata handling, and static/runtime report integration pass for declared support scope | +| 5. Measured throughput | Matched CPU and eager-CUDA timing; units, scope, route provenance where available | Phase 3; Phase 4 for GQA throughput | Profiling + EP owners | Attribution, aggregation, invalid durations, and timing validation pass | +| 6. Cost-model experiments | Calibrated latency/transfer model; separate route/placement experiment | Phase 5 + applicable memory feasibility contracts | Performance + EP owners | Representative end-to-end benefit and fallback behavior proven before any policy adoption | + +Phases 0-2 remain an independently shippable initial MVP, extended by the implemented bias increment. +The current priority is a MatMulNBits runtime slice followed by required GQA baseline coverage, not exhaustive +executed-work accounting. GQA formula/convention work can proceed without waiting for all runtime infrastructure. +Timing does not block either baseline; MatMulNBits timing need not wait for broader Attention coverage. + +The ancillary increments retain their numbering; route-specific work does not block the minimum operator scope: + +| Increment | Deliverables | Dependencies | Proposed owner | Exit gate | +| --- | --- | --- | --- | --- | +| 2a. Shape-based ancillary work | Separate optional-bias count; component-level coverage and exclusions | Phase 2 and approved component contract | Framework + quantization owners | Present/absent bias and dynamic/empty geometry tests pass; core counts remain unchanged | +| Optional logical reconstruction | Separately labeled one-reconstruction scenario, if consumers need it | Phase 2 and approved scalar reconstruction convention | Quantization + tooling owners | Hand-calculated cases pass; never labeled per-invocation executed work | +| 3a. Route-specific ancillary work | Incremental dequantization, unpacking, conversion, layout, padding, and prepacking estimates | Phase 3 plus validated route and preparation-frequency metadata for each supported path | Quantization + EP owners | Route, reuse, padding, repetition, phase, and non-overlap tests pass; unknown paths remain unavailable | + +Increment 2a is implemented for review. Deliver Phase 3 first as a bounded MatMulNBits runtime slice, then +integrate GQA baseline estimates. Logical reconstruction and Increment 3a remain optional follow-ups. +Do not require complete ancillary coverage or exact executed counts before adding GQA. + +### Current implementation order and minimum-scope acceptance + +Here, "MatMul counting" means the existing `com.microsoft::MatMulNBits` target, not new coverage of standard +`MatMul`, `Gemm`, or every quantized matrix operator. + +1. Review and commit the existing MatMulNBits core-plus-bias feature, tests, and report documentation. +2. Implement the next coding task: capture actual MatMulNBits invocation dimensions, bias presence, selected + top-level route, and compute type where available; reuse the existing pure estimators and expose runtime results. + Design metadata/reporting with GQA needs in mind without speculative universal kernel abstractions. +3. Add GQA baseline formulas and adapters using query heads, head dimensions, query lengths, active KV lengths, + and an approved causal/windowed/ragged or rectangular convention. Support prefill, decode, and chunked workloads + within a declared scope; report missing host-visible semantic values explicitly. +4. Validate both required operators with independent references, changing workloads, report artifacts, compatibility, + and count-only overhead measurements. No route-specific refinement is required to pass baseline acceptance. +5. Add route-family adjustments incrementally using actual selected-route metadata and justified models of + reconstruction, padding, repetitions, or reductions. Keep preparation lifetime and arithmetic categories explicit. + +The minimum-scope acceptance gate requires: + +- MatMulNBits core and bias and GQA core estimates match independent references under documented conventions. +- Static concrete and runtime invocation results retain distinct provenance; repeated/dynamic/concurrent cases + have correct identity and no stale count reuse. Control-flow totals include only established executions. +- GQA uses query-head count and active KV lengths, not KV-head count or cache capacity as substitutes; the masking + convention, ancillary exclusions, and unsupported cases are visible. +- Available baseline estimates survive missing route metadata. Unsupported adjustments have explicit reasons; + neither unknown extra work nor missing baseline inputs are silently treated as zero. +- Reports separate logical baselines, adjusted replacements, categories, and phases, with no double counting or + unsubstantiated executed-FLOP/whole-operator claims. +- Supported build/inference compatibility and the approved overhead gate pass; review artifacts cover both + available and unavailable cases. + +This gate covers useful estimates for both required operators, not exact instruction counts or universal EP +coverage. Route-specific ancillary work, other Attention operators, and timing are independently reviewable +follow-ups, not prerequisites for the required GQA baseline. + +### Phase 2: static-report launch + +- Use an explicit-path session config entry for the developer preview, with warning-only report failures. +- Restrict initial support to full inference builds. Builds excluding contrib ops must retain unchanged behavior + and must not acquire a dependency on unavailable contrib implementations. +- Validate real optimized models, including cases where fusion removes all directly supported pilot nodes. +- Include feature availability, convention/schema versions, exclusions, and unavailable reasons in documentation. +- Publish reports without changing provider assignment, kernel selection, allocator behavior, or inference outputs. +- Establish component naming, units, provenance, execution phase, and completeness rules for later ancillary work. + Keep this initial contract small; the core-only report does not require speculative route-estimation machinery. + +### Increment 2a: shape-based ancillary work + +Initial bias implementation is available for review in the version-2 report described above. + +- Add optional bias as a separate `MN` logical floating-point-addition count, leaving `matmul_core` unchanged. +- Distinguish absent bias from unknown optional-input metadata; use checked arithmetic and existing unavailable states. +- Report component-level exclusions and coverage, not just whether a node has any estimate. +- If a logical reconstruction scenario is requested, report its subtraction and multiplication counts separately + under an approved convention. Do not add `2KN` to every invocation or imply reconstruction actually ran. + +### Phase 3: runtime refinement + +- Start with a MatMulNBits runtime slice. Read dimensions and bias presence from validated inputs, route/compute + type from actual implementation decisions, and preparation state where exposed. Reuse existing metadata; + never infer actual execution from a maximum-shape hint or a provider name alone. +- For GQA, collect validated head geometry, query/active KV lengths, masking/window settings, and per-sequence + lengths where host-visible. Missing device-only values remain unavailable; do not copy them solely for estimates. +- Add an opt-in runtime reporting hook: initialization-time static export cannot describe per-invocation routes. + Baseline counting must remain independent of whether route-specific metadata is available. +- Recompute counts for actual invocation geometry, rather than presenting static hints as observed work. +- Distinguish concurrent runs, repeated invocations, subgraphs, and node identity within one session. +- Aggregate only executed control-flow paths once runtime support exists; otherwise retain explicit incompleteness. +- Do not copy GPU scalar values to the host or synchronize solely to compute a count. +- Use already validated host-visible metadata, or report `missing_semantic_value`. +- Bound cache growth and test changing shapes and semantic values, not just repeated fixed-shape runs. + +### Increment 3a: route-specific ancillary work + +Optional refinement, not a MatMulNBits/GQA baseline gate. For one arithmetic family at a time: + +1. Reuse the kernel's validated geometry and actual route decision; do not duplicate its dispatch cascade. +2. Establish whether preparation runs during initialization, first use, or each invocation. +3. Determine processed elements, including padding, repeated tiles, split-K work, and any cached or shared preparation. +4. Classify floating-point arithmetic, integer/bitwise work, conversions, and logical data movement separately. +5. Add checked pure formulas fed by that metadata, with independent reference tests. +6. Report unavailable when route, reuse, or frequency is unknown; do not synchronize or copy device data solely for + estimation. Include all count-affecting route and preparation state in cache keys or invalidate cached estimates. + +Fusion does not imply that a component is absent or separately materialized. Verify which arithmetic is retained, +eliminated, or repeated before attributing it. Treat prepacking as its own lifetime-scoped component, and distinguish +first-use preparation from steady-state work. Do not attribute shared preparation in full to multiple consumers. + +### Phase 4: Attention expansion + +GQA is required baseline coverage, not an expansion blocked on MatMulNBits route audits. Implement it after the +first MatMulNBits runtime slice; convention and pure-formula work can begin earlier. Other Attention operators +follow separately. Core formulas must have exact tests under their declared workload convention, even though +their relationship to kernel execution is approximate. + +First agree on useful unmasked work versus rectangular dense-equivalent work. For useful core work, define `P` as +the sum of valid query-key pairs across batches and query heads: + +```text +QK-transpose + PV core operations = 2 * P * (D_qk + D_v) +``` + +For rectangular Attention with equal head dimensions `D`, this reduces to: + +```text +4 * batch * query_heads * query_length * active_KV_length * D +``` + +Use query heads for grouped-query Attention, not KV-head count. Separate active KV length from allocated cache +capacity. Count valid pairs explicitly for causal, sliding-window, and ragged sequences, including any past-length +offset. Packed token count alone does not determine the sum of per-sequence query-key pairs. + +State whether projections, softmax, bias, rotary transforms, and other auxiliary work are included. PackedAttention +and PackedMultiHeadAttention do not have identical projection work. Backend tiling, recomputation, and padding +belong to a separately named route-specific metric, not an undocumented change to the logical count. + +### Phase 5: throughput + +Timing is separate from the minimum count deliverable and does not require exhaustive route accounting. +Count-only overhead measurements remain part of runtime-count validation. + +Only calculate a rate when the count and duration cover the same invocation and components: + +```text +effective operations/s = operation_count / duration_seconds +effective TFLOP/s = logical_flops / duration_seconds / 10^12 +``` + +Label quantized results as dense-equivalent operations/s, not literal hardware TFLOP/s. + +- Identify timing scope: CPU invocation, host submission, correlated device execution, or end-to-end run. +- A multi-kernel op needs a defined correlated interval; do not divide its full count by one child kernel's duration. +- Separate warmup, initialization/prepack, compilation/tactic search, and steady-state execution. +- Omit rates for missing, zero, negative, or uncorrelated durations and report why. +- Aggregate repeated matched invocations using `sum(counts) / sum(durations)`, not the mean of their rates. +- Do not sum overlapping device intervals and call the result end-to-end latency. +- Keep CUDA Graph replay and unsupported EP timing explicitly unavailable until attribution is verified. +- Record device, dtype, workload, provider, and route when observable; unknown route identity must remain unknown. +- Name the numerator's component scope. Core-only throughput can ship without ancillary coverage; adding a component + must not silently redefine an existing rate or mix initialization work into a steady-state duration. + +## 7. Validation and acceptance + +Use independent expected values rather than comparing two implementations of the same formula. + +| Area | Required cases | Acceptance | +| --- | --- | --- | +| Pilot arithmetic | Small hand-calculated cases, batched leading dimensions, dtype/bit-width variants | Exact equality under the convention; example above equals 24,576 | +| Metadata | Unknown rank/dimensions, mismatched K, missing/invalid attributes, omitted optional inputs | Explicit correct reason; no zero fallback or unsigned wrap | +| Boundaries | Valid zero extent, large products, uint64 overflow, zero with other unknown dimensions | Documented validity rules and checked arithmetic; unavailable when validity is unproven | +| Semantics | Bias present/absent, prepacked weights, CPU/CUDA routes | Same core count; exclusions preserved | +| Logical ancillary work | Bias present/absent, empty/dynamic outputs, overflow, hand-calculated reconstruction cases | Exact component counts under the stated convention; logical reconstruction never claimed as observed execution | +| Route-specific ancillary work | Cached versus per-invocation preparation, dtype, padding, repeated tiles, split-K, fusion, shared prepacking, route changes | Independently verified counts and phase attribution; no overlap, stale reuse, or extra synchronization; unknown routes unavailable | +| Static graph | Unsupported ops, fusion, opaque partitions, subgraphs, shape scenarios | No double counting; exact subtotals and coverage state | +| Runtime | Changing batch/sequence sizes, repeated/concurrent runs, count-affecting scalar changes | Invocation-specific counts; no stale cache reuse or extra synchronization | +| Attention | Hand-enumerated pairs for causal/past/windowed/ragged cases; unequal QK/V dimensions | Exact pair-based counts for supported cases; explicit missing-metadata result otherwise | +| Timing | Multiple kernels, asynchronous GPU work, invalid durations, unequal repeated workloads | Correct attribution, units, and weighted aggregation | +| Compatibility | Feature disabled, contrib-disabled build, unsupported build/report request | Existing inference behavior preserved; feature unavailability surfaced | + +Proposed overhead gates, **subject to team approval before Phase 3**: + +- Disabled: no estimator invocation, report allocation, device event, or new synchronization. +- Static-only: no per-run work; measure report/initialization cost separately. +- Runtime-count-only: at most 1% median steady-state latency overhead on an agreed CPU/CUDA matrix, with uncertainty + reported. Measure tiny launch-bound workloads separately rather than hiding them in an average. +- Timing-enabled: publish overhead separately; it is not covered by the runtime-count-only budget. + +Use existing framework/operator test targets and profiling test infrastructure. Run CPU tests first; CUDA tests are +required when CUDA runtime correlation is added. Maintain a golden report/schema test and an independent benchmark +comparing disabled, count-only, and timing-enabled modes. Record exact commands and results in implementation PRs. + +## 8. Enablement, observability, and rollback + +1. **Internal validation:** maintainers exercise the pilot and golden reports on small deterministic graphs. +2. **Developer preview:** expose opt-in static reporting with documented coverage; gather model/report feedback. +3. **Runtime preview:** opt in to MatMulNBits runtime counts and exposed route metadata; record overhead and gaps. +4. **Required operator coverage:** add GQA baseline estimates and validate both operators against minimum-scope + acceptance. Do not block on executed-work completeness. +5. **Optional refinements and broader support:** add route adjustments, matched timing, and other operators with + owner-reviewed conventions and reference cases. Stabilize a public + API only after consumers validate the report and unknown-state semantics. + +Track unavailable-reason frequencies, supported-node coverage, report generation cost, and count/timing overhead +in local reports and benchmark results. This proposal does not add automatic telemetry. + +Rollback is to disable reporting/instrumentation, not to change inference policy. Keep count-only and timing modes +independently switchable. Any later placement/dispatch experiment needs its own flag, fallback, quality gates, +and rollback plan. + +## 9. Open questions for the team + +Phase-2 preview decisions: use `session.compute_estimate_report_path`, versioned component-oriented JSON, separate +nested graph scopes, and warning-only export failures. These choices are implemented for developer review; ownership, +additional build coverage, and later-phase semantics still require team agreement. + +| ID | Decision to resolve | Recommendation | Needed by | +| --- | --- | --- | --- | +| Q1 | Core arithmetic only or complete operator work? | Core arithmetic first; explicit exclusions and separate auxiliary components later | Phase 0 | +| Q2 | Is quantized dense-equivalent work an acceptable first pilot? | Yes; never label it literal executed FLOPs. If literal floating-point work is required first, choose a floating-point pilot before implementation | Phase 0 | +| Q3 | Who owns conventions, schema, and the first reporting surface? | Framework owns the contract; operator owners approve formulas; internal developer report before public APIs | Phase 0 | +| Q4 | Which full/minimal/training/plugin builds participate? | Full inference first; no public/plugin ABI changes in the MVP | Phase 2 | +| Q5 | How should scenarios, valid zeros, and incomplete control-flow totals appear? | Distinct provenance and unavailable reasons; no claimed per-run total without execution evidence | Phase 2 | +| Q6 | Is the proposed count-only overhead budget acceptable, and on which workloads? | Approve a representative CPU/CUDA matrix and a tiny-op cohort before instrumentation | Phase 3 | +| Q7 | Useful unmasked or rectangular dense-equivalent Attention work? | One explicitly named primary convention; optional secondary metric only if needed | Phase 4 | +| Q8 | Which timing paths reliably correlate full op execution, including capture? | CPU and eager CUDA first; defer capture and other EPs until validated | Phase 5 | +| Q9 | Should external EPs/custom ops register estimators? | Revisit after multiple built-in operators demonstrate the needed contract | Public API design | +| Q10 | When may estimates affect placement or dispatch? | Only after calibrated end-to-end experiments, including bandwidth and transfer costs | Phase 6 | +| Q11 | Which ancillary components and reconstruction conventions should be supported first? | Bias after static reporting; logical reconstruction only if needed; executed-work estimates require validated route and frequency metadata | Phase 2 design; per-route approval before Increment 3a | +| Q12 | What is the minimum supported workload/build scope for MatMulNBits and GQA? | Agree on baseline components, prefill/decode/chunked and masking cases, runtime metadata availability, and explicit gaps; route refinements are optional | Phase 3 design and Phase 4 baseline acceptance | + +## 10. Completion criteria + +The **CPU MatMulNBits MVP is complete** when a developer can opt into persistent, versioned per-run reports, +see correct core and bias estimates from actual executed shapes, and understand collection limits, failures, +coverage gaps and exclusions without changes to inference results or dispatch. Batch 3 provides this developer +preview; production rollout still requires overhead qualification and the declared platform/build matrix. + +The **minimum operator scope is complete** when MatMulNBits core-plus-bias and GQA core estimates satisfy the +Section 6 acceptance gate for the declared workload/build matrix, with documented conventions, runtime support, +and explicit gaps. Exact route-specific executed-FLOP counts and full ancillary coverage are not required. + +The **diagnostic feature is mature** when supported runtime invocations and matched throughput can be reported with +validated attribution, honest coverage, and measured overhead. Automatic performance planning remains a separate +feature, not a prerequisite for delivering useful counts. diff --git a/docs/annotated_partitioning/matmul_nbits_counter_completion_plan.md b/docs/annotated_partitioning/matmul_nbits_counter_completion_plan.md new file mode 100644 index 0000000000000..8e09513fdbe52 --- /dev/null +++ b/docs/annotated_partitioning/matmul_nbits_counter_completion_plan.md @@ -0,0 +1,456 @@ +# MatMulNBits Counter: Validation, Hardware Coverage, and Timing Completion Plan + +## 1. Objective and definition of done + +Finish a developer-facing diagnostic that reports useful, auditable MatMulNBits workload counts and measured +throughput without changing inference results, dispatch, or memory policy. + +This is the execution plan following Batches 1-3 of the +[contrib-op rollout](contrib_op_flop_estimation_rollout.md). It focuses on MatMulNBits; GQA remains a separate +required operator in the broader rollout, not a prerequisite for validating this counter. Exhaustive +hardware-instruction accounting must not delay either baseline. + +The deliverable has three distinct levels: + +| Level | Deliverable | Accuracy contract | +| --- | --- | --- | +| Logical workload | Runtime core and bias counts | Exact under the declared convention for supported captured invocations | +| Implementation work | Optional route-specific estimates | Approximate, with documented assumptions, evidence and supported configurations | +| Measured throughput | Count divided by matched execution duration | Measured rate with explicit numerator, timing scope and uncertainty | + +**FLOPs are a count; FLOP/s is a rate.** Adding timing does not turn a logical count into literal hardware work. +For quantized MatMulNBits, the primary rate must be named **dense-equivalent operations/s** or +**dense-equivalent TOP/s**, not actual executed floating-point TFLOP/s. +An estimated floating-point-only numerator can have its own estimated FLOP/s metric when justified. +Do not combine integer/bitwise operations, floating-point operations, conversions and bytes into one FLOP count. + +The feature is complete for a declared platform when: + +- The user's model has passed count and output-equivalence validation on that platform. +- Supported executed MatMulNBits invocations are accounted for, with missing coverage explicitly identified. +- Reference arithmetic checks pass exactly; optional implementation models meet their own evidence-based gates. +- Rates use correlated counts and durations, with no silent fallback to host launch timing. +- Count-only, timing-enabled and export overhead have been measured separately. +- Unsupported routes/platforms/replay modes remain explicit rather than returning misleading zero counts. + +## 2. Starting point: implemented versus unverified + +Implemented by Batches 1-3: + +- Bounded typed per-run collector, run identity/outcome, diagnostics and subgraph propagation. +- CPU MatMulNBits runtime shapes, bias presence, quantization metadata and selected route/effective-mode provenance. +- Checked core `2MNK` and bias `MN` formulas, with independent components. +- Opt-in schema-v3 runtime JSON export after canonical Run, with unique process/session/run filename suffixes. +- Configurable automatic limits: 10,000 records and 16 MiB accounted payload per run by default. +- Explicit overflow, unavailable estimates, drops and captured-scope completeness. + +Not implemented or not qualified: + +- CUDA capture, GPU architecture provenance, GPU completion-aware collection and device timing. +- Timing/rate metrics, route-adjusted work, complete model coverage or standard MatMul/Gemm instrumentation. +- Production overhead qualification, all platform/build combinations, and the user's own model validation. + +Existing implementation anchors: + +- [Collector contract](../../onnxruntime/core/framework/runtime_workload_collector.h). +- [Runtime report builders](../../onnxruntime/core/framework/compute_estimate_report.h). +- [Pure core/bias formulas](../../onnxruntime/core/framework/matmul_nbits_compute_estimate.cc). +- [CPU dispatch and recording](../../onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc). +- [CUDA implementation to extend](../../onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc). + +Prior CPU Debug results are baseline evidence, not proof for a new build: 88 estimator/report/transformer tests +passed, with 25 collector/execution passes and one capability-gated LUT skip. Re-run after implementation changes. + +## 3. Reviewable implementation sequence + +| Batch | Work | Prerequisite | Review evidence | +| --- | --- | --- | --- | +| 4 | Validate current CPU behavior on synthetic and user models | Existing Batches 1-3 | Reproducible runs, independent counts and coverage reconciliation | +| 5 | Fix demonstrated gaps and qualify CPU overhead | Batch 4 findings | Focused regressions and before/after measurements | +| 6 | Provider-neutral provenance and CUDA count capture | Stable logical contract; CPU fixes affecting shared transport | Cross-device count parity and explicit actual CUDA routing | +| 7 | CPU and CUDA timing with matched rates | Stable identity; CUDA capture for GPU timing | Correlated timelines, timing accuracy and overhead results | +| 8 | Optional implementation-work models and hardware qualification | Actual route evidence; timing where needed for profiling | Per-route formulas, held-out validation and supported-range table | +| 9 | Developer rollout and expansion to further EPs | Count/timing acceptance on each advertised platform | Usage recipe, known limitations and rollback instructions | + +Keep each batch independently reviewable. Capture numerical baseline improvements before adding refinements. +CPU timing can proceed before CUDA timing if the shared identity contract is settled. +Do not add GPU-specific arithmetic multipliers just to make the plan appear hardware-aware. + +## 4. Batch 4: test the existing work + +### 4.1 Verify the actual runtime artifact + +1. Build the branch containing Batch 3; record revision, build configuration and enabled EPs. +2. For Python, build/install that branch's wheel into an isolated virtual environment. + Record `onnxruntime.__file__`, version and available providers, plus wheel/build provenance. + A published wheel accepting a configuration key does not prove it contains the counter. +3. Keep the CPU Debug test build separate from the Python wheel: the existing + `build\Windows\FlopEstimate` configuration has Python bindings and CUDA disabled. +4. Run C++ tests from the build output directory. Confirm a nonzero test count and that the changed targets rebuilt. + +Existing targeted commands on Windows, after building the corresponding targets: + +```powershell +Set-Location C:\Users\tommasoadani\repos\onnxruntime\build\Windows\FlopEstimate\Debug + +.\onnxruntime_provider_test.exe --gtest_filter="ComputeEstimateReportTest.*:RuntimeComputeEstimateReportTest.*:MatMulNBitsComputeEstimate*:MatMulNBitsBiasEstimate*" +.\onnxruntime_test_all.exe --gtest_filter="RuntimeWorkloadCollectorTest.*:RuntimeWorkloadExecutionTest.*" +``` + +See the repository [build skill](../../.github/skills/ort-build/SKILL.md) and +[test skill](../../.github/skills/ort-test/SKILL.md) for wheel/build commands and test execution guidance. +Missing route capabilities must skip explicitly; skipped cases remain unvalidated on that machine. + +### 4.2 Run a small hand-checkable model + +Use the existing synthetic end-to-end report test first: A=[2,3,64], N=32 and bias. + +```text +M = product(A dimensions except the last) = 6 +core dense-equivalent operations = 2 * 6 * 32 * 64 = 24576 +bias logical floating-point additions = 6 * 32 = 192 +``` + +The local [MatMulNBits-plus-Add model](../../onnxruntime/test/testdata/transform/runtime_optimization/matmulnbits_add.onnx) +is another candidate. Its [generator](../../onnxruntime/test/testdata/transform/runtime_optimization/matmulnbits_add_gen.py) +defines A=[M,2], N=3, 4-bit weights and a separate Add. With optimizations disabled and M=4: +core=48, internal bias=0; the separate Add is outside this counter. +Execute and inspect the actual artifact before treating it as validated. +With optimizations enabled, inspect whether Add fused into MatMulNBits and update the component expectation. + +### 4.3 Test the user's own model + +1. Start with a CPU-compatible ONNX model containing `com.microsoft::MatMulNBits`. + Inspect operator domains/types and the optimized executable graph, not just the filename or quantization label. + A QDQ-plus-MatMul model or a fused opaque partition does not automatically qualify. +2. Save representative real input shapes/settings: batch size, prompt length, decode steps and cache usage as applicable. + Keep model/input contents local; record only non-sensitive metadata needed to reproduce the experiment. +3. Create a fresh report directory. Enable `session.compute_estimate_report_path` with an output prefix, + not a literal JSON filename. Leave explicit budgets large enough for the planned run. +4. Run without reporting and retain reference outputs. Run with reporting using identical inputs/settings. + Apply an appropriate predeclared dtype/backend tolerance; use exact equality where deterministic and appropriate. +5. Verify initialization writes no report and each Run produces one distinct schema-v3 report. +6. Require successful run state, no drops/reuse rejection, available core/bias components and captured-scope completeness. +7. Reconcile executed invocation identities/counts with an independent trace or controlled graph execution. + Static node counts alone do not account for branch selection, loop iterations, decoder subgraphs or optimization. +8. Recompute each recorded core/bias count independently and compare both per-invocation values and checked totals. +9. Repeat for dynamic batch/sequence sizes, multiple runs, concurrent callers and relevant subgraphs. +10. Repeat with zero record/payload budgets, disabled reporting, invalid output directory and intentional inference failure. + Inference behavior must remain unchanged; diagnostics and partial results must reflect what happened. + +**Do not accept an empty report as a successful count validation.** +`collection_complete` describes attempted submissions, not all-model coverage. +`whole_run_work_complete` and `whole_operator_work_complete` remain false for this scoped feature. + +Deliverable: a local validation bundle with build provenance, input shapes, configuration, output checks, generated +reports, independent reference totals, invocation reconciliation and unresolved gaps. Do not commit model weights, +tensor contents or routine logs. + +## 5. Counting strategy and accuracy evidence + +### 5.1 Maintain three independent checks + +| Check | Method | What it establishes | +| --- | --- | --- | +| Formula correctness | Hand examples and independent arbitrary-precision reference arithmetic | Exact counts under convention, including overflow behavior | +| Runtime coverage | Compare executed identities/multiplicity with independent execution evidence | Records are not missing, duplicated or attributed to the wrong run | +| Implementation work | Inspect actual selected backend path and corroborate with targeted profiling | Optional route estimates reflect the implementation within declared assumptions | + +Reference arithmetic should use actual A dimensions and logical K/N, never packed-weight storage dimensions. +Do not call the production helper from the reference oracle or multiply by backend batch_count after using full A shape. +Use arbitrary-precision integers in the oracle, then check whether the result fits the production uint64 contract. + +Cover aligned/unaligned K/N, vector inputs, batching, bias presence, supported bit widths/block sizes, empty outputs, +zero reduction where supported, repeated subgraphs and unsupported metadata. +Test helper-only semantic cases separately from configurations that an actual kernel accepts. +Property tests should check linear scaling of counts, route-independent baselines and absence of wraparound. + +### 5.2 Count actual implementation work only where justified + +Manually audit a bounded set of route families once, then automate formulas from runtime metadata: + +- Padded matrix/tile work only if those lanes actually perform arithmetic; padding storage alone is not proof. +- Dequantization and scale/zero-point arithmetic, distinguishing cached preparation from per-invocation reconstruction. +- Repeated reconstruction within tiles, split-K reductions, conversions and separate/fused bias work. +- Packing/prepacking and tactic search as initialization/first-use work, not automatically charged to every steady-state run. + +For each model, document formula, arithmetic categories, route/version applicability, lifetime, required metadata, +excluded work and evidence. Missing required details mean unavailable adjustment, not zero extra work. +Preserve the original logical count and convention; add separately named/versioned implementation components. + +Use backend source, controlled microbenchmarks and profiler observations together. +Profiler hardware counters are corroboration, not universal ground truth: instruction-to-operation mapping, +tensor-core tile dimensions, predication, precision, replay and counter availability vary by architecture. +Record metric definitions/tool versions and isolate the intended operation. Do not infer FLOPs from utilization +percentages or treat matching latency as proof of a correct operation count. + +### 5.3 Accuracy acceptance + +- Logical counts: exact agreement with all supported reference cases; tolerance is zero. +- Coverage: every expected supported invocation reconciles in controlled tests with no unexplained duplicates or drops. +- Implementation models: publish error/limitations per route/configuration where independent measurements exist. + If no defensible reference exists, label the result model-based/unvalidated rather than claiming a percentage. +- Performance prediction is optional and separate from measuring actual duration. + Any predicted latency requires held-out shapes and device/backend-specific calibration. +- Do not advertise a universal accuracy percentage or arbitrary uncertainty band before collecting evidence. + +### 5.4 Route order and the refinement loop + +Use the following order rather than attempting to model every CPU implementation simultaneously: + +1. **Unpacked floating-point route:** establish reconstruction frequency, logical GEMM work, actual padding behavior + and bias placement. Start with one CPU configuration and hand-checkable cases, then validate held-out shapes + and the user's model where that route actually executes. +2. **Packed QNBit:** choose one available effective compute mode, validate its arithmetic categories and reuse + assumptions, then expand to other modes/CPU capabilities individually. Do not equate quantized dot products + with floating-point GEMM just because the dense-equivalent count matches. +3. **LUT:** establish table construction, lookup/accumulation work and reuse frequency before publishing adjustments. + Missing LUT capability is an explicit coverage gap, not evidence that the route model passed. +4. **CUDA routes:** apply the same loop using actual device/route evidence and the completion rules in Section 7. + +Keep the baseline tests running for every route, even while only one route's refinement is under development. +Advance a refinement when its reference cases pass, actual route execution is proven, supported assumptions are +documented and held-out cases expose no unexplained discrepancies. Other routes may retain baseline-only support. + +The repeatable loop is: + +```text +Choose route and hypothesis + -> generate controlled factor matrix + -> run and prove actual route + -> compare with independent arithmetic/coverage references + -> investigate discrepancies using backend/profiler evidence + -> revise formula or collected metadata + -> add focused regression + -> rerun reference and held-out cases + -> validate on user's model + -> accept documented scope or repeat +``` + +Changing the formula is not always the fix: a mismatch may be missing capture, wrong route attribution, duplicated +batches, a different preparation lifetime or an invalid test configuration. Record the cause and evidence. +Do not tune arithmetic coefficients merely to make elapsed time agree. + +### 5.5 Planned automation: route-validation matrix runner + +**Yes, this workflow can be scripted.** Add a Python CLI as a follow-up implementation task, with proposed location +`tools\python\matmul_nbits_counter_matrix.py`. This script does **not** exist as part of this documentation change. +It should orchestrate model generation/execution, independent baseline checks and result aggregation; it cannot +automatically discover correct implementation-work formulas from timing alone. + +Implement a CPU unpacked pilot first, then extend the same runner to packed QNBit, LUT and CUDA. +Use existing model-generation/quantization helpers where suitable, but keep the arithmetic oracle independent +of the production estimator. Use an isolated environment with a wheel containing the feature. + +**Configurable factors** + +| Group | Factors to sweep | +| --- | --- | +| Geometry | M/N/K, leading batch dimensions, rank, aligned/irregular dimensions | +| Semantics | Bias absent/present, supported bit widths/block sizes, activation dtype | +| Route controls | Requested route/mode, constant versus dynamic weights where supported, relevant session settings | +| Lifecycle | Fresh session, first invocation, warmup and repeated steady-state invocations | +| Execution | EP/device, CPU thread count, serial/concurrent run cases | +| Diagnostics | Disabled, capture-only where accessible, export enabled; collection budgets | + +Use a small one-factor-at-a-time cohort first, followed by selected interactions and a separate held-out cohort. +Do not blindly run the full Cartesian product: preflight configuration validity, budget total cases/runtime, +and report unsupported combinations explicitly. Keep failed attempts visible rather than silently filtering them out. +CPU-only shared settings and CUDA-only settings must not be applied interchangeably. + +**Runner responsibilities** + +- Assign a stable case ID, matrix/cohort version and deterministic input seed. +- Record build/wheel provenance, model identity, device/backend configuration and requested settings. +- Run controlled cases in isolated output directories; isolate processes when session/global backend settings require it. +- Assert the **observed** route/effective mode from structured records, not logs or requested configuration. + Route controls are requests, not guarantees. A route mismatch must remain a visible non-pass outcome. +- Detect missing reports, failed runs, missing/duplicate invocations, drops and unavailable estimates. +- Compare outputs against a predeclared reference/tolerance and compute baseline counts using arbitrary-precision + reference arithmetic. Preserve decimal counts without floating-point precision loss in exports. +- Include route-adjusted expected/measured components only when a validated reference and required metadata exist. + Missing reference data is unavailable, not zero error. +- Separate arithmetic checks, coverage checks, numeric correctness and timing checks into distinct outcomes. +- Preserve requested repetitions, warmup policy and variability; initially leave unsupported timing fields unavailable. + CPU Run wall time includes synchronous export when enabled and must not be labeled operator compute time. +- Return a failing exit status for required-case failures, route mismatches or unexpectedly missing coverage. + Capability skips can be allowed by the declared matrix policy but must be counted and listed. + +**Decision-friendly output** + +Produce detailed machine-readable JSON, a flat CSV and a generated Markdown summary in the chosen results directory. +Keep per-invocation evidence available separately; aggregate by route/mode, geometry cohort and device. +Suggested columns: + +```text +case_id, cohort, device, EP, dtype, bits, block_size, A_shape, N, bias, +requested_route, observed_route, effective_mode, lifecycle, +expected_invocations, captured_invocations, collection_complete, +expected_core_ops, reported_core_ops, core_delta, +expected_bias_ops, reported_bias_ops, bias_delta, +refinement_version, implementation_reference_status, implementation_error, +output_check, timing_scope, duration, duration_unit, timing_status, repeats, variability, +status, reason, evidence_path +``` + +Use explicit statuses such as pass, arithmetic mismatch, output mismatch, route mismatch, incomplete capture, +unsupported configuration, capability skip and execution failure. Summary tables should show how many requested +cases actually exercised each route, all non-pass reasons, worst discrepancies and unvalidated configurations. +Link failures to raw reports and profiler/reference evidence. A dashboard is optional; CSV and Markdown suffice first. + +At each refinement iteration, compare matrices with the same case/cohort IDs and retain formula/build versions. +Review where counts or statuses changed and add minimal failing cases to the normal unit/regression suites. +Do not adjust models using the held-out cohort without replacing that cohort for subsequent independent validation. + +## 6. Batch 5: refine only demonstrated gaps + +Prioritize fixes from Batch 4 and retain a regression for each: + +1. Correct arithmetic, identity, missing-subgraph or export behavior before adding new metrics. +2. Reconcile graph-path encoding and schema semantics where cross-tool consumers need stable identities. +3. Surface coverage exclusions more usefully; distinguish known excluded nodes from unknown dynamic execution counts. +4. Qualify failed/async runs, collector lifetime and report publication semantics. + Consider atomic report publication if partial-file visibility causes consumer errors. +5. Measure collector storage, snapshot/serialization memory, export duration and growth in report files. + Add retention/sampling controls only if usage evidence requires them; preserve explicit partial-coverage semantics. +6. Retire legacy graph adapters only after their remaining tests/consumers are understood; retain useful pure math. + +Benchmark Release builds with reporting disabled, capture only, and capture plus export. +Include tiny operations, decode-like small M and larger prefill-like M, with repeated trials and variability. +Set an overhead budget with the team before production rollout; no numeric budget is claimed as achieved here. + +## 7. Batch 6: expand to CUDA, then other EPs + +### 7.1 Shared transport and schema + +- Separate logical operator metadata from provider-specific route details. + Current route/mode enums describe CPU paths; do not relabel CUDA routes as CPU LUT/QNBit modes. +- Make shared-provider/plugin boundaries explicit: use supported bridge mechanisms or a boundary-safe typed contract, + not direct in-tree framework dependencies from a provider DLL. +- Reuse pure arithmetic across providers through an appropriate shared layer; avoid linking CUDA to CPU kernel code. +- Generalize the report's current CPU-only provider check and coverage label as part of CUDA enablement. +- Capture device identity/architecture and backend/build versions once per session/device. + Correlate records to the actual device in multi-device sessions; GPU name alone is not a dispatch key. + +### 7.2 CUDA recording and lifecycle + +Audit actual dispatch branches in the CUDA implementation, including mixed-weight GEMM, architecture/layout-specific +paths, fallback dequantization/GEMM and tactic selection where applicable. These are inspection targets, not a claim +that every branch is available on every GPU or build. + +- Record actual host-visible dimensions, bias, precision, quantization settings and selected path. +- Record tile/tactic/reduction information only where exposed and needed by a supported refinement. +- Do not transfer tensors to the CPU or synchronize a device solely to calculate baseline counts. +- Distinguish host submission, device completion and failure. Successful launch is not confirmed GPU completion. +- Define deferred completion/event ownership before closing snapshots or reporting complete GPU results. + Keep pending/unverified records explicit when callers disable synchronization. +- Ensure profiling/autotuning trial launches do not become multiple logical invocations. +- Keep CUDA Graph capture/replay unavailable until replay-aware identities, counts and timing are validated. + +### 7.3 Hardware and build matrix + +Use available representatives of Ampere (for example A100), Hopper (H100) and Blackwell (B100/B200-class), +plus representative CPU configurations. Availability is not presumed: list untested hardware explicitly. +Record compute capability, driver/toolkit/library versions, build flags, dtype, prepacked layout and selected route. +Source support for one architecture or prepacked layout is not evidence that another supports it. + +Validate: + +- Same executed logical shapes produce identical baseline counts across hardware. +- EP-specific tests prove actual CUDA dispatch, not successful CPU fallback. +- Irregular dimensions, small/large M, bias, supported quantization settings and dynamic weights where allowed. +- Architecture-dependent route eligibility, fallback and resource/launch limits. +- Native/shared-provider, Windows/Linux and selected build guards for the advertised support matrix. +- Unsupported device/layout/build combinations give explicit diagnostics, not misleading count-only success. + +Further EPs follow the same contract. Opaque fused partitions require their own attribution strategy; +do not fabricate per-node records from pre-partition graphs. + +## 8. Batch 7: timing and measured throughput + +### 8.1 Define timing scope before implementing rates + +| Scope | Numerator | Duration | Label | +| --- | --- | --- | --- | +| CPU operator | Captured operator logical core count | CPU wall interval covering the operator computation, including its worker completion | Dense-equivalent operator throughput | +| CUDA operator | Captured operator logical core count | Correlated device interval covering all its constituent compute kernels | Dense-equivalent device throughput | +| End-to-end Run | Explicit supported-work subtotal | Completed run wall time, with transfer/export inclusion stated | Supported-work throughput over run latency, not whole-model FLOP/s | +| Optional floating-point implementation component | Validated estimated FP-only work | Matched component interval, or explicitly labeled enclosing-operator interval | Estimated FP work per matched/enclosing interval | + +A core-only logical numerator over the full operator duration is a useful **effective** throughput metric, +including bias/dequantization costs in the elapsed interval; it is not isolated core execution throughput. +Name this convention explicitly. Never divide the full operator count by only its fastest child kernel. + +### 8.2 Correlation and instrumentation + +1. Reuse existing ORT profiling infrastructure where it provides the required scope. +2. Introduce stable run/node/invocation correlation shared by records and timing events. + Current collector invocation IDs are assigned on successful submission; do not match events by list position, + node name, or commit order. Repeated/concurrent subgraph calls need distinct identities. +3. CPU: use a monotonic wall clock around the defined compute region, not CPU thread time that misses worker work. + State whether recorder overhead is included; keep JSON serialization/file I/O outside the compute interval. +4. CUDA: use correlated CUDA events on the actual stream or suitable provider profiling. + An operator spanning multiple streams requires dependency-aware attribution; a single-stream event pair is insufficient. +5. Reclaim bounded event/profiling resources after completion; surface exhaustion or unsupported timing as unavailable. + Count-only mode must not add event/synchronization overhead. Timing collection can be sampled explicitly. +6. Export finalized timing/rates only when completion is established. + Define an explicit finalize/deferred-export path where required, rather than forcing synchronization in count-only mode. +7. Separate warmup, first-use packing, tactic search and steady-state samples. +8. Version the report when adding timing fields. Include duration units, clock/scope, correlation identity, + completion state, availability reason and timing methodology. + +### 8.3 Calculation and aggregation + +```text +dense_equivalent_ops_per_second = matched_dense_equivalent_ops / duration_seconds +dense_equivalent_TOP_per_second = dense_equivalent_ops_per_second / 1e12 +estimated_FP_ops_per_second = matched_estimated_FP_ops / duration_seconds +``` + +- Omit a rate for missing, nonpositive, invalid or uncorrelated duration, failed/pending device work or unavailable count. +- Never infer time from GPU peak TFLOPS and present it as a measurement. +- Aggregate matched invocation samples with sum(counts)/sum(durations), not the arithmetic mean of their rates. +- Keep per-operator accumulated time separate from run wall time: overlapping intervals are not additive latency. +- Exclude warmup from steady-state aggregates and report median/tail latency, sample count and variability. +- Preserve incomplete/dropped sample counts; no full-run throughput claim from an unknown partial numerator. +- Keep bias logical additions separate. Whole-operator time is not an isolated bias duration. + +Validate event duration/correlation against independent timelines, with concurrent runs and multi-kernel routes. +Compare uninstrumented, count-only, timing-enabled and full-export runs to measure observer overhead. +Timing acceptance thresholds must account for timer resolution and very short operations and be approved before rollout. + +## 9. Batch 8: optional hardware-aware refinement + +Add a refinement only when it answers a consumer question that the baseline and measured duration cannot. +Start with the **CPU unpacked floating-point route**, then a bounded **packed QNBit** configuration, then **LUT**. +Validate one available CUDA route afterward, then broaden only after evidence supports each model. + +Each refinement must include: + +- A route/architecture/backend-version applicability key, not a GPU-model correction factor. +- Independent logical baseline, named modeled components and separate arithmetic/traffic categories. +- Source-based derivation plus controlled measurements and held-out workload validation where possible. +- Explicit unsupported/extrapolated cases and model/convention version. +- Preparation/reuse assumptions and whether work is amortized across invocations. + +For hardware performance comparisons, use measured rates on the actual devices. +If prediction for an unavailable device is needed, treat that as a separate calibrated model with error reporting. +Keep architecture-specific instruction auditing bounded; useful ballpark estimates do not require exhaustive assembly counting. + +## 10. Rollout gate and audit checklist + +- [ ] CPU reference arithmetic and existing regression suites pass on the final build. +- [ ] User-model inputs, output equivalence, invocation coverage and per-component totals are reconciled. +- [ ] Dynamic/repeated/concurrent/subgraph, zero, overflow, failure, drop and disabled cases pass. +- [ ] Report migration, coverage exclusions, filesystem failures and resource budgets are documented. +- [ ] Release-build collection/export memory and latency overhead are measured against agreed budgets. +- [ ] Every advertised EP/device/build combination has actual execution evidence; skips/untested cases are listed. +- [ ] GPU submission/completion and graph replay semantics are honest. +- [ ] Timing events correlate to the same invocations as counts; async/multi-kernel cases are verified. +- [ ] Rates distinguish logical dense-equivalent throughput from estimated executed floating-point work. +- [ ] Optional route refinements meet their own evidence gates without modifying the portable baseline. +- [ ] Usage examples and generated reports are available for reviewer sign-off. + +Ship count-only and timing as independently opt-in modes. Start with a limited developer preview on validated CPU +and CUDA configurations. Roll back by disabling diagnostics, not by changing inference dispatch. +Broader GQA support and other EPs can reuse the validated contract without waiting for exhaustive route refinement. diff --git a/include/onnxruntime/core/framework/op_kernel_context.h b/include/onnxruntime/core/framework/op_kernel_context.h index 4d110e128a434..ea998f4408373 100644 --- a/include/onnxruntime/core/framework/op_kernel_context.h +++ b/include/onnxruntime/core/framework/op_kernel_context.h @@ -5,6 +5,7 @@ namespace onnxruntime { class IExecutionFrame; #if !defined(ORT_MINIMAL_BUILD) class KernelPilot; +struct RuntimeWorkloadRecorder; #endif class Stream; namespace concurrency { @@ -216,6 +217,8 @@ class OpKernelContext { // Keep new virtuals after existing declarations to preserve their vtable slots. // Session-owned pilot for this kernel; nullptr when unavailable. virtual KernelPilot* GetKernelPilot() const { return nullptr; } + // Internal typed diagnostic transport. Standalone/constant-folding contexts do not collect. + virtual const RuntimeWorkloadRecorder* GetRuntimeWorkloadRecorder() const { return nullptr; } #endif private: diff --git a/include/onnxruntime/core/framework/run_options.h b/include/onnxruntime/core/framework/run_options.h index 54a0355537f8e..fcaea0e2a16c1 100644 --- a/include/onnxruntime/core/framework/run_options.h +++ b/include/onnxruntime/core/framework/run_options.h @@ -11,6 +11,7 @@ #include "core/framework/config_options.h" namespace onnxruntime { +class RuntimeWorkloadCollector; namespace lora { class LoraAdapter; } @@ -38,6 +39,12 @@ struct OrtRunOptions { // Set to 'true' to enable profiling for this run. bool enable_profiling = false; +#if !defined(ORT_MINIMAL_BUILD) + // Internal diagnostic hook. Caller owns a fresh collector through completion of this Run. + // No C API setter or report/export lifecycle is exposed. + onnxruntime::RuntimeWorkloadCollector* runtime_workload_collector = nullptr; +#endif + // File prefix for profiling result for this run. // The actual filename will be: _.json // Only used when enable_profiling is true. diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 02537bb643fe7..e1d450928b3fa 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -577,6 +577,23 @@ static const char* const kOrtSessionOptionsMaxShapeOverride = "session.max_shape static const char* const kOrtSessionOptionsStrictWorkspaceVerification = "session.strict_workspace_verification"; +// Opt-in runtime compute-estimation JSON reports. Value is a UTF-8 output prefix, not a literal filename. +// Empty/unset disables automatic capture/export. After each canonical Run, append a process token, session ID, +// run ID and ".json". Parent directories are not created. Use a new prefix for independent application launches. +// Replaces the legacy initialization-time static export; initialization no longer writes workload reports. +// Reports contain actual shapes, graph/node identities, CPU MatMulNBits core dense-equivalent and logical bias counts, +// route provenance, run outcome and collection diagnostics. Unsupported operators/providers are not instrumented; +// captured totals are not whole-model totals or hardware instruction counts. No additional device synchronization. +// File failures warn without failing inference. Minimal builds warn that export is unsupported. +// Internal callers supplying a RunOptions collector use that collector's limits instead of automatic limits. +static const char* const kOrtSessionOptionsComputeEstimateReportPath = "session.compute_estimate_report_path"; + +// Nonnegative decimal per-run automatic capture budgets; defaults: 10000 records and 16777216 bytes. +// Zero retains no records. Bytes are accounted payload, not a hard allocator/peak-memory cap. +// Invalid values fail initialization. Limits do not bound total report files/disk usage across runs. +static const char* const kOrtSessionOptionsComputeEstimateMaxRecords = "session.compute_estimate_max_records"; +static const char* const kOrtSessionOptionsComputeEstimateMaxPayloadBytes = "session.compute_estimate_max_payload_bytes"; + // Enable EP context feature to dump the partitioned graph which includes the EP context into Onnx file. // The dumped Onnx model with EP context can be used for future inference to avoid the EP graph partitioning/compile overhead. // "0": disable. (default) diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc index d42dbae635047..b514c06580bf2 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc @@ -12,6 +12,7 @@ #include "core/common/narrow.h" #include "core/common/safeint.h" #include "core/framework/op_kernel.h" +#include "core/framework/runtime_workload_collector.h" #include "core/mlas/inc/mlas_qnbit.h" #include "core/mlas/inc/mlas_q4.h" #include "core/providers/cpu/math/matmul_helper.h" @@ -1535,8 +1536,22 @@ Status MatMulNBits::Compute(OpKernelContext* ctx) const { Tensor* y = ctx->Output(0, helper.OutputShape()); +#if !defined(ORT_MINIMAL_BUILD) + const auto* recorder = ctx->GetRuntimeWorkloadRecorder(); + auto record = [&](MatMulNBitsRuntimeRoute route, MatMulNBitsRuntimeComputeMode mode) { + if (recorder) { + recorder->RecordMatMulNBits(a->Shape(), y->Shape(), helper.M(), helper.N(), helper.K(), + helper.OutputOffsets().size(), bias != nullptr, a->GetElementType(), + static_cast(nbits_), static_cast(block_size_), route, mode); + } + }; +#endif + // Bail out early if the output is going to be empty if (y->Shape().Size() == 0) { +#if !defined(ORT_MINIMAL_BUILD) + record(MatMulNBitsRuntimeRoute::EmptyOutput, MatMulNBitsRuntimeComputeMode::NotExecuted); +#endif return Status::OK(); } @@ -1556,11 +1571,38 @@ Status MatMulNBits::Compute(OpKernelContext* ctx) const { // MlasQNBitGemmPackQuantBDataSize() returns 0, we can consider calling MlasQNBitGemmBatch() // with B directly too. if (prefer_lut_gemm_) { - return ComputeBPackedLUT(a, bias, y, thread_pool, helper); + ORT_RETURN_IF_ERROR(ComputeBPackedLUT(a, bias, y, thread_pool, helper)); +#if !defined(ORT_MINIMAL_BUILD) + record(MatMulNBitsRuntimeRoute::PackedLut, MatMulNBitsRuntimeComputeMode::Lut); +#endif + return Status::OK(); } if (MlasIsQNBitGemmAvailable(nbits_, block_size_, compute_type_)) { - return ComputeBPacked(a, scales, zero_points, bias, y, allocator, thread_pool, helper); + ORT_RETURN_IF_ERROR(ComputeBPacked(a, scales, zero_points, bias, y, allocator, thread_pool, helper)); +#if !defined(ORT_MINIMAL_BUILD) + // HQNBIT_CompInt8 converts half activations and calls the SQNBIT backend on every platform. + std::optional mode; + switch (compute_type_) { + case SQNBIT_CompFp32: + mode = MatMulNBitsRuntimeComputeMode::SQNBitFp32; + break; + case HQNBIT_CompFp16: + mode = MatMulNBitsRuntimeComputeMode::HQNBitFp16; + break; + case BHQNBIT_CompBf16: + mode = MatMulNBitsRuntimeComputeMode::BHQNBitBf16; + break; + case SQNBIT_CompInt8: + case HQNBIT_CompInt8: + mode = MatMulNBitsRuntimeComputeMode::SQNBitInt8; + break; + } + if (mode) { + record(MatMulNBitsRuntimeRoute::PackedQNBit, *mode); + } +#endif + return Status::OK(); } } @@ -1571,7 +1613,11 @@ Status MatMulNBits::Compute(OpKernelContext* ctx) const { "This is because MLAS doesn't have an optimized quantized kernel " "for the requested compute configuration."; - return ComputeBUnpacked(a, b, scales, zero_points, reorder_idx, bias, y, allocator, thread_pool, helper); + ORT_RETURN_IF_ERROR(ComputeBUnpacked(a, b, scales, zero_points, reorder_idx, bias, y, allocator, thread_pool, helper)); +#if !defined(ORT_MINIMAL_BUILD) + record(MatMulNBitsRuntimeRoute::Unpacked, MatMulNBitsRuntimeComputeMode::Float32Gemm); +#endif + return Status::OK(); } #define REGISTER_MatMulNBits(T1) \ diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h index 9b1081c5f9c43..e658b76a8bdfd 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h @@ -304,7 +304,8 @@ Status BeamSearchGpt::Execute(const FeedsFetchesManager* init_run_feeds_fetch this->context_.Logger(), this->ort_stream_, /*sync_subgraph_fetches*/ false, - this->context_.GetRunProfiler()); + this->context_.GetRunProfiler(), + this->context_.GetRuntimeWorkloadCollector()); } else { #ifdef DEBUG_NODE_INPUTS_OUTPUTS const_cast(this->decoder_session_state_).IncrementGraphExecutionCounter(); @@ -319,7 +320,8 @@ Status BeamSearchGpt::Execute(const FeedsFetchesManager* init_run_feeds_fetch this->context_.Logger(), this->ort_stream_, /*sync_subgraph_fetches*/ false, - this->context_.GetRunProfiler()); + this->context_.GetRunProfiler(), + this->context_.GetRuntimeWorkloadCollector()); } ORT_RETURN_IF_ERROR(status); diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h index 70aeacc736107..7f0efd38b8b19 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h @@ -182,7 +182,8 @@ Status BeamSearchT5::Execute(const FeedsFetchesManager& encoder_feeds_fetches this->context_.Logger(), this->ort_stream_, /*sync_subgraph_fetches*/ false, - this->context_.GetRunProfiler()); + this->context_.GetRunProfiler(), + this->context_.GetRuntimeWorkloadCollector()); ORT_RETURN_IF_ERROR(status); #ifdef DEBUG_GENERATION @@ -362,7 +363,8 @@ Status BeamSearchT5::Execute(const FeedsFetchesManager& encoder_feeds_fetches this->context_.Logger(), this->ort_stream_, /*sync_subgraph_fetches*/ false, - this->context_.GetRunProfiler()); + this->context_.GetRunProfiler(), + this->context_.GetRuntimeWorkloadCollector()); ORT_RETURN_IF_ERROR(status); diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h index a9f783ee7cc0a..2242bcfd8af5d 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h @@ -188,7 +188,8 @@ Status BeamSearchWhisper::Execute(const FeedsFetchesManager& encoder_feeds_fe this->context_.Logger(), this->ort_stream_, /*sync_subgraph_fetches*/ false, - this->context_.GetRunProfiler()); + this->context_.GetRunProfiler(), + this->context_.GetRuntimeWorkloadCollector()); ORT_RETURN_IF_ERROR(status); #ifdef DEBUG_GENERATION @@ -388,7 +389,8 @@ Status BeamSearchWhisper::Execute(const FeedsFetchesManager& encoder_feeds_fe this->context_.Logger(), this->ort_stream_, /*sync_subgraph_fetches*/ false, - this->context_.GetRunProfiler()); + this->context_.GetRunProfiler(), + this->context_.GetRuntimeWorkloadCollector()); ORT_RETURN_IF_ERROR(status); diff --git a/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_gpt.h b/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_gpt.h index f5a416781d1e4..957182d28cae5 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_gpt.h +++ b/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_gpt.h @@ -296,7 +296,8 @@ Status GreedySearchGpt::Execute(const FeedsFetchesManager* init_ this->context_.Logger(), this->ort_stream_, /*sync_subgraph_fetches*/ false, - this->context_.GetRunProfiler()); + this->context_.GetRunProfiler(), + this->context_.GetRuntimeWorkloadCollector()); } else { #ifdef DEBUG_NODE_INPUTS_OUTPUTS const_cast(this->decoder_session_state_).IncrementGraphExecutionCounter(); @@ -311,7 +312,8 @@ Status GreedySearchGpt::Execute(const FeedsFetchesManager* init_ this->context_.Logger(), this->ort_stream_, /*sync_subgraph_fetches*/ false, - this->context_.GetRunProfiler()); + this->context_.GetRunProfiler(), + this->context_.GetRuntimeWorkloadCollector()); } ORT_RETURN_IF_ERROR(status); diff --git a/onnxruntime/core/framework/compute_estimate_report.cc b/onnxruntime/core/framework/compute_estimate_report.cc new file mode 100644 index 0000000000000..7bbbf5b2c6177 --- /dev/null +++ b/onnxruntime/core/framework/compute_estimate_report.cc @@ -0,0 +1,594 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/framework/compute_estimate_report.h" + +#if !defined(ORT_MINIMAL_BUILD) + +#include +#include +#include +#include +#include +#include +#include + +#include "nlohmann/json.hpp" +#include "core/common/common.h" +#include "core/common/inlined_containers.h" +#include "core/common/path_string.h" +#include "core/common/safeint.h" +#include "core/framework/runtime_workload_collector.h" +#include "core/graph/constants.h" +#include "core/graph/graph.h" +#include "core/platform/env.h" + +#if !defined(DISABLE_CONTRIB_OPS) +#include "core/framework/matmul_nbits_compute_estimate.h" +#endif + +namespace onnxruntime { +namespace { + +using Json = nlohmann::ordered_json; + +Json ExcludedComponents() { + return Json::array({"dequantization", "bias", "packing", "conversion", "padding", "route_specific_extra_work"}); +} + +#if !defined(DISABLE_CONTRIB_OPS) +const char* ReasonName(contrib::ComputeEstimateUnavailableReason reason) { + switch (reason) { + case contrib::ComputeEstimateUnavailableReason::UnsupportedOperator: + return "unsupported_operator"; + case contrib::ComputeEstimateUnavailableReason::UnsupportedVersion: + return "unsupported_version"; + case contrib::ComputeEstimateUnavailableReason::UnknownShape: + return "unknown_shape"; + case contrib::ComputeEstimateUnavailableReason::InvalidMetadata: + return "invalid_metadata"; + case contrib::ComputeEstimateUnavailableReason::Overflow: + return "overflow"; + case contrib::ComputeEstimateUnavailableReason::MissingSemanticValue: + return "missing_semantic_value"; + } + ORT_THROW("Invalid compute estimate reason."); +} +#endif + +Json InputShapes(const Node& node) { + Json result = Json::array(); + for (const auto* input : node.InputDefs()) { + if (input == nullptr || !input->Exists()) { + result.push_back({{"state", "missing"}}); + continue; + } + const auto* shape = input->Shape(); + Json value{{"name", input->Name()}, {"state", shape ? "present_with_shape" : "present_without_shape"}}; + if (shape) { + Json dimensions = Json::array(); + for (const auto& dim : shape->dim()) { + if (dim.has_dim_value()) { + dimensions.push_back(dim.dim_value()); + } else { + dimensions.push_back(nullptr); + } + } + value["dimensions"] = std::move(dimensions); + } + result.push_back(std::move(value)); + } + return result; +} + +// Lookup remains internal and deliberately narrow; route-dependent metrics do not belong here. +Json EstimateNode(const Node& node, std::optional& operation_count) { + operation_count.reset(); + Json result{{"availability", "unavailable"}}; + if (node.NodeType() == Node::Type::Fused) { + result["unavailable_reason"] = "opaque_partition"; + } else if (node.ContainsSubgraph()) { + result["unavailable_reason"] = "runtime_execution_required"; + } else if (node.Domain() != kMSDomain || node.OpType() != "MatMulNBits") { + result["unavailable_reason"] = "unsupported_operator"; + } else { + Json component{{"component", "matmul_core"}, + {"metric", "dense_equivalent_ops"}, + {"scope", "matmul_core"}, + {"convention_version", "core_arithmetic_v1"}, + {"input_basis", "canonical_graph_metadata"}, + {"basis", "logical_semantics"}, + {"execution_phase", "per_invocation"}, + {"excluded_components", ExcludedComponents()}}; +#if !defined(DISABLE_CONTRIB_OPS) + const auto estimate = contrib::EstimateMatMulNBitsCompute(node); + if (const auto* count = std::get_if(&estimate)) { + operation_count = count->dense_equivalent_ops; + component["availability"] = "available"; + component["estimate_kind"] = "exact_under_convention"; + component["input_basis"] = "static_concrete"; + component["operation_count"] = std::to_string(count->dense_equivalent_ops); + } else { + component["availability"] = "unavailable"; + component["unavailable_reason"] = + ReasonName(std::get(estimate)); + } +#else + component["availability"] = "unavailable"; + component["unavailable_reason"] = "estimators_disabled"; +#endif + result["availability"] = component["availability"]; + if (component.contains("unavailable_reason")) { + result["unavailable_reason"] = component["unavailable_reason"]; + } + result["components"] = Json::array({std::move(component)}); + Json bias{{"component", "bias"}, + {"metric", "floating_point_ops"}, + {"scope", "bias"}, + {"convention_version", "bias_add_v1"}, + {"input_basis", "canonical_graph_metadata"}, + {"basis", "logical_semantics"}, + {"execution_phase", "per_invocation"}, + {"excluded_components", Json::array({"matmul_core", "dequantization", "packing", "conversion", + "padding", "route_specific_extra_work"})}}; +#if !defined(DISABLE_CONTRIB_OPS) + const auto bias_estimate = contrib::EstimateMatMulNBitsBias(node); + if (const auto* count = std::get_if(&bias_estimate)) { + const auto& inputs = node.InputDefs(); + bias["input_basis"] = inputs.size() > 5 && inputs[5] != nullptr && inputs[5]->Exists() + ? "static_concrete" + : "optional_input_absent"; + bias["availability"] = "available"; + bias["estimate_kind"] = "exact_under_convention"; + bias["operation_count"] = std::to_string(count->floating_point_ops); + } else { + bias["availability"] = "unavailable"; + bias["unavailable_reason"] = ReasonName(std::get(bias_estimate)); + } +#else + bias["availability"] = "unavailable"; + bias["unavailable_reason"] = "estimators_disabled"; +#endif + const bool core_available = operation_count.has_value(); + const bool bias_available = bias["availability"] == "available"; + result["availability"] = core_available && bias_available ? "available" + : core_available || bias_available ? "partial" + : "unavailable"; + if (!bias_available && core_available) { + result["unavailable_reason"] = bias["unavailable_reason"]; + } + result["components"].push_back(std::move(bias)); + } + if (!result.contains("components")) { + result["components"] = Json::array(); + } + result["whole_operator_work_complete"] = false; + return result; +} + +std::string EscapePathSegment(std::string_view value) { + std::string result; + for (char c : value) { + if (c == '~') { + result += "~0"; + } else if (c == '/') { + result += "~1"; + } else { + result += c; + } + } + return result; +} + +void AppendGraph(const Graph& graph, const std::string& path, Json& report) { + uint64_t known_ops = 0; + bool subtotal_overflow = false; + size_t estimated_count = 0; + size_t unavailable_count = 0; + size_t partial_count = 0; + struct ComponentSubtotal { + std::string name; + std::string metric; + std::string convention; + uint64_t count = 0; + size_t estimated = 0; + size_t unavailable = 0; + bool overflow = false; + Json reasons = Json::object(); + }; + InlinedVector component_subtotals; + bool has_subgraphs = false; + Json reasons = Json::object(); + const auto graph_index = report["graphs"].size(); + report["graphs"].push_back({{"graph_path", path}, {"graph_name", graph.Name()}}); + + for (const auto& node : graph.Nodes()) { + Json entry{{"graph_path", path}, + {"node_index", node.Index()}, + {"node_name", node.Name()}, + {"domain", node.Domain()}, + {"op_type", node.OpType()}, + {"opset_version", node.SinceVersion()}, + {"provider", node.GetExecutionProviderType()}, + {"input_shapes", InputShapes(node)}}; + std::optional operation_count; + entry.update(EstimateNode(node, operation_count)); + if (operation_count.has_value()) { + if (!subtotal_overflow && !SafeAdd(known_ops, *operation_count, known_ops)) { + subtotal_overflow = true; + } + } + if (entry["availability"] == "available") { + ++estimated_count; + } else { + if (entry["availability"] == "partial") { + ++partial_count; + } else { + ++unavailable_count; + } + const std::string reason = entry["unavailable_reason"].get(); + reasons[reason] = reasons.value(reason, size_t{0}) + 1; + } + for (const auto& component : entry["components"]) { + const auto name = component["component"].get(); + auto it = std::find_if(component_subtotals.begin(), component_subtotals.end(), + [&](const auto& subtotal) { return subtotal.name == name; }); + if (it == component_subtotals.end()) { + component_subtotals.push_back({name, component["metric"].get(), + component["convention_version"].get()}); + it = component_subtotals.end() - 1; + } + if (component["availability"] == "available") { + ++it->estimated; + const auto count = std::stoull(component["operation_count"].get()); + if (!it->overflow && !SafeAdd(it->count, count, it->count)) { + it->overflow = true; + } + } else { + ++it->unavailable; + const auto reason = component["unavailable_reason"].get(); + it->reasons[reason] = it->reasons.value(reason, size_t{0}) + 1; + } + } + report["nodes"].push_back(std::move(entry)); + + const auto subgraphs = node.GetAttributeNameToSubgraphMap(); + InlinedVector attributes; + attributes.reserve(subgraphs.size()); + for (const auto& item : subgraphs) { + attributes.push_back(item.first); + } + std::sort(attributes.begin(), attributes.end()); + for (const auto& attribute : attributes) { + has_subgraphs = true; + AppendGraph(*subgraphs.at(attribute), + path + "/" + std::to_string(node.Index()) + "/" + EscapePathSegment(attribute), report); + } + } + + Json component_summaries = Json::array(); + bool component_overflow = false; + for (const auto& subtotal : component_subtotals) { + Json value{{"component", subtotal.name}, + {"metric", subtotal.metric}, + {"convention_version", subtotal.convention}, + {"basis", "logical_semantics"}, + {"execution_phase", "per_invocation"}, + {"estimated_component_count", subtotal.estimated}, + {"unavailable_component_count", subtotal.unavailable}, + {"unavailable_reason_counts", subtotal.reasons}, + {"subtotal_availability", subtotal.overflow ? "unavailable" : "available"}, + {"complete_for_declared_component_scope", subtotal.unavailable == 0 && !subtotal.overflow}}; + if (subtotal.overflow) { + component_overflow = true; + value["subtotal_unavailable_reason"] = "overflow"; + } else { + value["known_operation_count"] = std::to_string(subtotal.count); + } + component_summaries.push_back(std::move(value)); + } + Json summary{{"reported_node_count", estimated_count + partial_count + unavailable_count}, + {"estimated_node_count", estimated_count}, + {"partially_estimated_node_count", partial_count}, + {"unavailable_node_count", unavailable_count}, + {"component_summaries", std::move(component_summaries)}, + {"unavailable_reason_counts", std::move(reasons)}, + {"subtotal_availability", subtotal_overflow ? "unavailable" : "available"}, + {"contains_subgraphs", has_subgraphs}, + {"complete_for_declared_graph_and_component_scope", + unavailable_count == 0 && partial_count == 0 && !subtotal_overflow && !component_overflow}, + {"whole_operator_work_complete", false}, + {"count_scope", "direct_nodes_only"}, + {"per_run_total_available", false}}; + if (subtotal_overflow) { + summary["subtotal_unavailable_reason"] = "overflow"; + } else { + summary["known_dense_equivalent_ops"] = std::to_string(known_ops); + } + report["graphs"][graph_index]["summary"] = std::move(summary); +} + +} // namespace + +Status BuildComputeEstimateReport(const Graph& graph, std::string& report) { + Status status; + ORT_TRY { + Json value{{"schema_version", 2}, + {"convention_version", "core_arithmetic_v1"}, + {"graph_scope", "optimized_executable"}, + {"nodes", Json::array()}, + {"graphs", Json::array()}}; +#if defined(DISABLE_CONTRIB_OPS) + value["estimator_support"] = "disabled"; +#else + value["estimator_support"] = "matmul_nbits_core_and_bias"; +#endif + AppendGraph(graph, "main", value); + value["summary"] = value["graphs"][0]["summary"]; + value["summary"]["graph_path"] = "main"; + report = value.dump(2) + "\n"; + return Status::OK(); + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to generate compute estimate report: ", ex.what()); + }); + } + return status; +} + +namespace { +Status WriteReport(const std::string& report, const std::string& utf8_path) { + Status status; + ORT_TRY { + std::ofstream stream(ToPathString(utf8_path), std::ios::out | std::ios::trunc | std::ios::binary); + ORT_RETURN_IF_NOT(stream.is_open(), "Failed to open compute estimate report: ", utf8_path); + stream << report; + stream.close(); + ORT_RETURN_IF_NOT(stream.good(), "Failed to write compute estimate report: ", utf8_path); + return Status::OK(); + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to write compute estimate report: ", utf8_path, ": ", ex.what()); + }); + } + return status; +} + +const char* RouteName(MatMulNBitsRuntimeRoute route) { + switch (route) { + case MatMulNBitsRuntimeRoute::EmptyOutput: + return "empty_output"; + case MatMulNBitsRuntimeRoute::PackedLut: + return "packed_lut"; + case MatMulNBitsRuntimeRoute::PackedQNBit: + return "packed_qnbit"; + case MatMulNBitsRuntimeRoute::Unpacked: + return "unpacked"; + } + ORT_THROW("Invalid MatMulNBits runtime route."); +} + +const char* ModeName(MatMulNBitsRuntimeComputeMode mode) { + switch (mode) { + case MatMulNBitsRuntimeComputeMode::NotExecuted: + return "not_executed"; + case MatMulNBitsRuntimeComputeMode::Float32Gemm: + return "float32_gemm"; + case MatMulNBitsRuntimeComputeMode::Lut: + return "lut"; + case MatMulNBitsRuntimeComputeMode::SQNBitFp32: + return "sqnbit_fp32"; + case MatMulNBitsRuntimeComputeMode::HQNBitFp16: + return "hqnbit_fp16"; + case MatMulNBitsRuntimeComputeMode::BHQNBitBf16: + return "bhqnbit_bf16"; + case MatMulNBitsRuntimeComputeMode::SQNBitInt8: + return "sqnbit_int8"; + case MatMulNBitsRuntimeComputeMode::HQNBitInt8: + return "hqnbit_int8"; + } + ORT_THROW("Invalid MatMulNBits runtime compute mode."); +} + +const char* RunStateName(RuntimeWorkloadRunState state) { + switch (state) { + case RuntimeWorkloadRunState::NotStarted: + return "not_started"; + case RuntimeWorkloadRunState::Running: + return "running"; + case RuntimeWorkloadRunState::Succeeded: + return "succeeded"; + case RuntimeWorkloadRunState::Failed: + return "failed"; + } + ORT_THROW("Invalid runtime workload run state."); +} + +Json RuntimeComponent(bool bias) { + return {{"component", bias ? "bias" : "matmul_core"}, + {"metric", bias ? "floating_point_ops" : "dense_equivalent_ops"}, + {"convention_version", bias ? "bias_add_v1" : "core_arithmetic_v1"}, + {"input_basis", "runtime_metadata"}, + {"basis", "logical_semantics"}, + {"execution_phase", "per_invocation"}, + {"availability", "unavailable"}}; +} +} // namespace + +Status WriteComputeEstimateReport(const Graph& graph, const std::string& utf8_path) { + std::string report; + ORT_RETURN_IF_ERROR(BuildComputeEstimateReport(graph, report)); + return WriteReport(report, utf8_path); +} + +Status BuildComputeEstimateReport(const RuntimeWorkloadSnapshot& snapshot, std::string& report) { + Status status; + ORT_TRY { + const auto& stats = snapshot.stats; + Json value{{"schema_version", 3}, + {"report_kind", "runtime_workload"}, + {"session_id", snapshot.session_id}, + {"run_id", snapshot.run_id}, + {"run_state", RunStateName(snapshot.run_state)}, + {"coverage_scope", "instrumented_cpu_matmul_nbits_invocations_only"}, + {"whole_run_work_complete", false}, + {"whole_operator_work_complete", false}, + {"route_adjustment_applied", false}, + {"excluded_components", Json::array({"dequantization", "packing", "conversion", "padding", + "route_specific_extra_work"})}, + {"invocations", Json::array()}}; + value["collection"] = {{"collection_complete", stats.collection_complete}, + {"max_records", snapshot.limits.max_records}, + {"max_payload_bytes", snapshot.limits.max_payload_bytes}, + {"retained_record_count", snapshot.records.size()}, + {"retained_payload_bytes", stats.retained_payload_bytes}, + {"record_limit_drops", stats.record_limit_drops}, + {"payload_limit_drops", stats.payload_limit_drops}, + {"allocation_failure_drops", stats.allocation_failure_drops}, + {"metadata_failure_drops", stats.metadata_failure_drops}, + {"run_reuse_rejections", stats.run_reuse_rejections}}; +#if defined(DISABLE_CONTRIB_OPS) + value["estimator_support"] = "disabled"; +#else + value["estimator_support"] = "matmul_nbits_core_and_bias"; +#endif + struct Subtotal { + uint64_t count = 0; + size_t estimated = 0; + size_t unavailable = 0; + bool overflow = false; + Json reasons = Json::object(); + }; + std::array subtotals; + for (const auto& record : snapshot.records) { + const auto& metadata = record.matmul_nbits; + Json entry{{"invocation_id", record.invocation_id}, + {"graph_path", record.graph_path}, + {"node_index", record.node_index}, + {"provider", record.provider}, + {"domain", kMSDomain}, + {"op_type", "MatMulNBits"}, + {"input_shape", metadata.input_shape}, + {"output_shape", metadata.output_shape}, + {"backend_geometry", {{"m", metadata.m}, {"n", metadata.n}, {"k", metadata.k}, {"batch_count", metadata.batch_count}}}, + {"bias_present", metadata.bias_present}, + {"input_element_type", metadata.input_element_type}, + {"bits", metadata.bits}, + {"block_size", metadata.block_size}, + {"components", Json::array()}}; + if (metadata.route) entry["route"] = RouteName(*metadata.route); + if (metadata.compute_mode) entry["compute_mode"] = ModeName(*metadata.compute_mode); + size_t available = 0; + for (size_t i = 0; i < subtotals.size(); ++i) { + Json component = RuntimeComponent(i == 1); + std::optional count; + if (record.provider != kCpuExecutionProvider) { + component["unavailable_reason"] = "unsupported_provider"; + } else { +#if !defined(DISABLE_CONTRIB_OPS) + // Full activation dimensions already include all batches; do not multiply by backend batch_count again. + if (i == 0) { + const auto estimate = contrib::EstimateMatMulNBitsCompute(metadata.input_shape, metadata.k, metadata.n); + if (const auto* result = std::get_if(&estimate)) { + count = result->dense_equivalent_ops; + } else { + component["unavailable_reason"] = ReasonName(std::get(estimate)); + } + } else { + const auto estimate = contrib::EstimateMatMulNBitsBias( + metadata.input_shape, metadata.k, metadata.n, metadata.bias_present); + if (const auto* result = std::get_if(&estimate)) { + count = result->floating_point_ops; + } else { + component["unavailable_reason"] = ReasonName(std::get(estimate)); + } + } +#else + component["unavailable_reason"] = "estimators_disabled"; +#endif + } + auto& subtotal = subtotals[i]; + if (count) { + ++available; + ++subtotal.estimated; + component["availability"] = "available"; + component["estimate_kind"] = "exact_under_convention"; + component["operation_count"] = std::to_string(*count); + if (!subtotal.overflow && !SafeAdd(subtotal.count, *count, subtotal.count)) { + subtotal.overflow = true; + } + } else { + ++subtotal.unavailable; + const auto reason = component["unavailable_reason"].get(); + subtotal.reasons[reason] = subtotal.reasons.value(reason, size_t{0}) + 1; + } + entry["components"].push_back(std::move(component)); + } + entry["availability"] = available == 2 ? "available" : available == 1 ? "partial" + : "unavailable"; + value["invocations"].push_back(std::move(entry)); + } + value["component_summaries"] = Json::array(); + for (size_t i = 0; i < subtotals.size(); ++i) { + const auto& subtotal = subtotals[i]; + Json summary = RuntimeComponent(i == 1); + summary.erase("availability"); + summary["estimated_component_count"] = subtotal.estimated; + summary["unavailable_component_count"] = subtotal.unavailable; + summary["unavailable_reason_counts"] = subtotal.reasons; + summary["count_scope"] = "retained_successful_invocations"; + summary["subtotal_availability"] = subtotal.overflow ? "unavailable" : "available"; + summary["complete_for_captured_component_scope"] = + snapshot.run_state == RuntimeWorkloadRunState::Succeeded && stats.collection_complete && + subtotal.unavailable == 0 && !subtotal.overflow; + if (subtotal.overflow) { + summary["subtotal_unavailable_reason"] = "overflow"; + } else { + summary["known_operation_count"] = std::to_string(subtotal.count); + } + value["component_summaries"].push_back(std::move(summary)); + } + report = value.dump(2) + "\n"; + return Status::OK(); + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to generate runtime compute estimate report: ", ex.what()); + }); + } + return status; +} + +Status WriteComputeEstimateReport(const RuntimeWorkloadSnapshot& snapshot, const std::string& utf8_path) { + std::string report; + ORT_RETURN_IF_ERROR(BuildComputeEstimateReport(snapshot, report)); + return WriteReport(report, utf8_path); +} + +Status ExportRuntimeComputeEstimateReport(const RuntimeWorkloadCollector& collector, const std::string& utf8_prefix) { + Status status; + ORT_TRY { + RuntimeWorkloadSnapshot snapshot; + ORT_RETURN_IF_ERROR(collector.GetSnapshot(snapshot)); + ORT_RETURN_IF_NOT(snapshot.run_state == RuntimeWorkloadRunState::Succeeded || + snapshot.run_state == RuntimeWorkloadRunState::Failed, + "Runtime compute estimate export requires a completed run."); + static const auto process_token = + std::to_string(Env::Default().GetSelfPid()) + "_" + + std::to_string(std::chrono::system_clock::now().time_since_epoch().count()); + const auto path = utf8_prefix + "_" + process_token + "_session-" + std::to_string(snapshot.session_id) + + "_run-" + std::to_string(snapshot.run_id) + ".json"; + return WriteComputeEstimateReport(snapshot, path); + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to export runtime compute estimate report: ", ex.what()); + }); + } + return status; +} + +} // namespace onnxruntime +#endif diff --git a/onnxruntime/core/framework/compute_estimate_report.h b/onnxruntime/core/framework/compute_estimate_report.h new file mode 100644 index 0000000000000..756e2a73a5cd6 --- /dev/null +++ b/onnxruntime/core/framework/compute_estimate_report.h @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#if !defined(ORT_MINIMAL_BUILD) + +#include + +#include "core/common/status.h" + +namespace onnxruntime { +class Graph; +struct RuntimeWorkloadSnapshot; +class RuntimeWorkloadCollector; + +// Canonical shapes in the final executable graph only. Subgraphs have separate summaries; +// counts are never summed across control-flow paths into a per-run total. +Status BuildComputeEstimateReport(const Graph& graph, std::string& report); +Status WriteComputeEstimateReport(const Graph& graph, const std::string& utf8_path); + +// Counts only retained, successfully submitted invocations. Run outcome and collection/coverage +// completeness are separate; route metadata does not adjust the logical counting convention. +Status BuildComputeEstimateReport(const RuntimeWorkloadSnapshot& snapshot, std::string& report); +Status WriteComputeEstimateReport(const RuntimeWorkloadSnapshot& snapshot, const std::string& utf8_path); + +// Writes __session-_run-.json after the collector is closed. +// No directories are created. Failures are diagnostic and must not replace inference status. +Status ExportRuntimeComputeEstimateReport(const RuntimeWorkloadCollector& collector, const std::string& utf8_prefix); + +} // namespace onnxruntime +#endif diff --git a/onnxruntime/core/framework/matmul_nbits_compute_estimate.cc b/onnxruntime/core/framework/matmul_nbits_compute_estimate.cc new file mode 100644 index 0000000000000..e3186b7c83a85 --- /dev/null +++ b/onnxruntime/core/framework/matmul_nbits_compute_estimate.cc @@ -0,0 +1,162 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/framework/matmul_nbits_compute_estimate.h" + +#include +#include + +#include "core/common/safeint.h" + +#if !defined(ORT_MINIMAL_BUILD) +#include "core/common/inlined_containers.h" +#include "core/graph/constants.h" +#include "core/graph/graph.h" +#endif + +namespace onnxruntime { +namespace contrib { +namespace { + +std::optional ValidateGeometry( + gsl::span input_a_shape, int64_t k, int64_t n) { + if (input_a_shape.empty() || k < 0 || n < 0 || + std::any_of(input_a_shape.begin(), input_a_shape.end(), [](int64_t dim) { return dim < -1; })) { + return ComputeEstimateUnavailableReason::InvalidMetadata; + } + + if (input_a_shape.back() >= 0 && input_a_shape.back() != k) { + return ComputeEstimateUnavailableReason::InvalidMetadata; + } + + if (std::find(input_a_shape.begin(), input_a_shape.end(), -1) != input_a_shape.end()) { + return ComputeEstimateUnavailableReason::UnknownShape; + } + + return std::nullopt; +} + +} // namespace + +MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute( + gsl::span input_a_shape, int64_t k, int64_t n) { + if (const auto reason = ValidateGeometry(input_a_shape, k, n)) { + return *reason; + } + + if (n == 0 || std::find(input_a_shape.begin(), input_a_shape.end(), 0) != input_a_shape.end()) { + return MatMulNBitsComputeEstimate{0}; + } + + uint64_t count = 2; + if (!SafeMultiply(count, static_cast(n), count)) { + return ComputeEstimateUnavailableReason::Overflow; + } + + // Multiplying every A dimension incorporates both the leading-dimension product M and K. + for (int64_t dim : input_a_shape) { + if (!SafeMultiply(count, static_cast(dim), count)) { + return ComputeEstimateUnavailableReason::Overflow; + } + } + + return MatMulNBitsComputeEstimate{count}; +} + +MatMulNBitsBiasEstimateResult EstimateMatMulNBitsBias( + gsl::span input_a_shape, int64_t k, int64_t n, std::optional bias_present) { + if (!bias_present.has_value()) { + return ComputeEstimateUnavailableReason::MissingSemanticValue; + } + if (!*bias_present) { + return MatMulNBitsBiasEstimate{0}; + } + if (const auto reason = ValidateGeometry(input_a_shape, k, n)) { + return *reason; + } + + const auto leading_dimensions = input_a_shape.first(input_a_shape.size() - 1); + if (n == 0 || std::find(leading_dimensions.begin(), leading_dimensions.end(), 0) != leading_dimensions.end()) { + return MatMulNBitsBiasEstimate{0}; + } + uint64_t count = static_cast(n); + for (int64_t dim : leading_dimensions) { + if (!SafeMultiply(count, static_cast(dim), count)) { + return ComputeEstimateUnavailableReason::Overflow; + } + } + return MatMulNBitsBiasEstimate{count}; +} + +#if !defined(ORT_MINIMAL_BUILD) +namespace { +template +auto EstimateFromNode(const Node& node) + -> std::conditional_t { + if (node.Domain() != kMSDomain || node.OpType() != "MatMulNBits") { + return ComputeEstimateUnavailableReason::UnsupportedOperator; + } + + if (node.SinceVersion() != 1) { + return ComputeEstimateUnavailableReason::UnsupportedVersion; + } + + const auto& attributes = node.GetAttributes(); + const auto k_it = attributes.find("K"); + const auto n_it = attributes.find("N"); + if (k_it == attributes.end() || n_it == attributes.end() || + k_it->second.type() != ONNX_NAMESPACE::AttributeProto::INT || !k_it->second.has_i() || + n_it->second.type() != ONNX_NAMESPACE::AttributeProto::INT || !n_it->second.has_i() || + k_it->second.i() < 0 || n_it->second.i() < 0) { + return ComputeEstimateUnavailableReason::InvalidMetadata; + } + + const auto& inputs = node.InputDefs(); + if (inputs.empty() || inputs[0] == nullptr || !inputs[0]->Exists()) { + return ComputeEstimateUnavailableReason::InvalidMetadata; + } + + if constexpr (Bias) { + // Bias is the sixth schema input; trailing omitted inputs need not have placeholders. + if (inputs.size() <= 5 || inputs[5] == nullptr || !inputs[5]->Exists()) { + return MatMulNBitsBiasEstimate{0}; + } + } + + const auto* shape = inputs[0]->Shape(); + if (shape == nullptr) { + return ComputeEstimateUnavailableReason::UnknownShape; + } + + InlinedVector dimensions; + dimensions.reserve(shape->dim_size()); + for (const auto& dim : shape->dim()) { + if (dim.has_dim_value()) { + if (dim.dim_value() < 0) { + return ComputeEstimateUnavailableReason::InvalidMetadata; + } + dimensions.push_back(dim.dim_value()); + } else { + dimensions.push_back(-1); + } + } + + if constexpr (Bias) { + return EstimateMatMulNBitsBias(dimensions, k_it->second.i(), n_it->second.i(), true); + } else { + return EstimateMatMulNBitsCompute(dimensions, k_it->second.i(), n_it->second.i()); + } +} +} // namespace + +MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute(const Node& node) { + return EstimateFromNode(node); +} + +MatMulNBitsBiasEstimateResult EstimateMatMulNBitsBias(const Node& node) { + return EstimateFromNode(node); +} +#endif + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/core/framework/matmul_nbits_compute_estimate.h b/onnxruntime/core/framework/matmul_nbits_compute_estimate.h new file mode 100644 index 0000000000000..6d4925a0c9cbc --- /dev/null +++ b/onnxruntime/core/framework/matmul_nbits_compute_estimate.h @@ -0,0 +1,63 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include + +namespace onnxruntime { +class Node; + +namespace contrib { + +enum class ComputeEstimateUnavailableReason { + UnsupportedOperator, + UnsupportedVersion, + UnknownShape, + InvalidMetadata, + Overflow, + MissingSemanticValue, +}; + +struct MatMulNBitsComputeEstimate { + // Convention v1: two operations per multiply-add, counting only the logical matrix product. + // Excludes bias, dequantization, packing, conversions, padding, and route-specific extra work. + static constexpr uint32_t kConventionVersion = 1; + uint64_t dense_equivalent_ops; +}; + +using MatMulNBitsComputeEstimateResult = + std::variant; + +struct MatMulNBitsBiasEstimate { + uint64_t floating_point_ops; +}; + +using MatMulNBitsBiasEstimateResult = + std::variant; + +// A has shape [..., K]. Rank-one inputs have M = 1. -1 denotes an unknown dimension; +// unknown dimensions remain unavailable even when another dimension is zero. +// This is a metadata-only estimate, not validation of packed weights or other operator inputs. +MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute( + gsl::span input_a_shape, int64_t k, int64_t n); + +// Counts MN logical additions, not route-specific executed work. Known-absent bias is zero +// without requiring shape metadata; unknown presence is unavailable. Present bias uses +// the same shape-validity rules as the core, but K = 0 does not imply empty output. +MatMulNBitsBiasEstimateResult EstimateMatMulNBitsBias( + gsl::span input_a_shape, int64_t k, int64_t n, std::optional bias_present); + +#if !defined(ORT_MINIMAL_BUILD) +// Reads canonical metadata from a resolved in-tree graph, never maximum-shape hints. +// The same core count applies across providers, dtypes, and supported quantization bit widths. +MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute(const Node& node); +MatMulNBitsBiasEstimateResult EstimateMatMulNBitsBias(const Node& node); +#endif + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/core/framework/op_kernel_context_internal.h b/onnxruntime/core/framework/op_kernel_context_internal.h index 65259710cf15c..33454e3217e06 100644 --- a/onnxruntime/core/framework/op_kernel_context_internal.h +++ b/onnxruntime/core/framework/op_kernel_context_internal.h @@ -7,6 +7,7 @@ #include "core/framework/op_kernel.h" #include "core/framework/session_state.h" #include "core/session/onnxruntime_c_api.h" +#include "core/framework/runtime_workload_collector.h" // onnxruntime internal OpKernelContext derived class to provide additional // APIs that aren't desirable to add to the public OpKernelContext API @@ -14,6 +15,7 @@ namespace onnxruntime { class SessionState; class ExecutionFrame; +class RuntimeWorkloadCollector; class OpKernelContextInternal : public OpKernelContext { public: @@ -23,11 +25,17 @@ class OpKernelContextInternal : public OpKernelContext { const logging::Logger& logger, const bool& terminate_flag, Stream* stream, - profiling::Profiler* run_profiler = nullptr) + profiling::Profiler* run_profiler = nullptr, + RuntimeWorkloadCollector* workload_collector = nullptr) : OpKernelContext(&frame, &kernel, stream, session_state.GetThreadPool(), logger), session_state_(session_state), terminate_flag_(terminate_flag), run_profiler_(run_profiler) { +#if !defined(ORT_MINIMAL_BUILD) + workload_recorder_ = {workload_collector, &session_state.GetGraphViewer().GetGraph(), &kernel.Node()}; +#else + ORT_UNUSED_PARAMETER(workload_collector); +#endif const auto& implicit_inputs = kernel.Node().ImplicitInputDefs(); int num_implicit_inputs = static_cast(implicit_inputs.size()); implicit_input_values_.reserve(num_implicit_inputs); @@ -59,6 +67,17 @@ class OpKernelContextInternal : public OpKernelContext { return kernel_pilot_; } + const RuntimeWorkloadRecorder* GetRuntimeWorkloadRecorder() const override { + return workload_recorder_.collector ? &workload_recorder_ : nullptr; + } + RuntimeWorkloadCollector* GetRuntimeWorkloadCollector() const { + return workload_recorder_.collector; + } +#else + RuntimeWorkloadCollector* GetRuntimeWorkloadCollector() const { return nullptr; } +#endif + +#if !defined(ORT_MINIMAL_BUILD) // Called by the executor only after successful Compute(). To commit usage, Compute() must obtain // the pilot via GetKernelPilot(), then BeginInvocation() and collect its selection before returning. // Without a lookup this is a no-op; KernelPilot::RecordUsage() also checks for a pending invocation. @@ -132,6 +151,7 @@ class OpKernelContextInternal : public OpKernelContext { private: #if !defined(ORT_MINIMAL_BUILD) mutable KernelPilot* kernel_pilot_{nullptr}; + RuntimeWorkloadRecorder workload_recorder_{}; class AccountingAllocator : public IAllocator { public: diff --git a/onnxruntime/core/framework/runtime_workload_collector.cc b/onnxruntime/core/framework/runtime_workload_collector.cc new file mode 100644 index 0000000000000..be2c51ee37a04 --- /dev/null +++ b/onnxruntime/core/framework/runtime_workload_collector.cc @@ -0,0 +1,168 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/framework/runtime_workload_collector.h" + +#if !defined(ORT_MINIMAL_BUILD) + +#include +#include +#include +#include + +#include "core/common/safeint.h" +#include "core/graph/graph.h" + +namespace onnxruntime { +namespace { + +void IncrementDropCount(uint64_t& count) { + if (count != std::numeric_limits::max()) { + ++count; + } +} + +bool PayloadSize(std::string_view graph_path, std::string_view provider, + const MatMulNBitsRuntimeMetadata& metadata, size_t& bytes) { + size_t dimensions; + size_t shape_bytes; + bytes = sizeof(RuntimeWorkloadRecord); + return SafeAdd(metadata.input_shape.size(), metadata.output_shape.size(), dimensions) && + SafeMultiply(dimensions, sizeof(int64_t), shape_bytes) && + SafeAdd(bytes, shape_bytes, bytes) && + SafeAdd(bytes, graph_path.size(), bytes) && + SafeAdd(bytes, provider.size(), bytes) && + SafeAdd(bytes, size_t{2}, bytes); +} + +} // namespace + +RuntimeWorkloadCollector::RuntimeWorkloadCollector(uint64_t session_id, uint64_t run_id, + RuntimeWorkloadLimits limits) + : session_id_(session_id), run_id_(run_id), limits_(limits) {} + +Status RuntimeWorkloadCollector::RecordMatMulNBits(std::string_view graph_path, size_t node_index, + std::string_view provider, + const MatMulNBitsRuntimeMetadata& metadata) { + std::lock_guard lock(mutex_); + if (records_.size() >= limits_.max_records) { + stats_.collection_complete = false; + IncrementDropCount(stats_.record_limit_drops); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Runtime workload collector record limit reached."); + } + + size_t bytes; + if (!PayloadSize(graph_path, provider, metadata, bytes) || + bytes > limits_.max_payload_bytes - stats_.retained_payload_bytes) { + stats_.collection_complete = false; + IncrementDropCount(stats_.payload_limit_drops); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Runtime workload collector payload limit reached."); + } + + Status status; + ORT_TRY { + RuntimeWorkloadRecord record{records_.size(), std::string(graph_path), node_index, + std::string(provider), metadata}; + records_.push_back(std::move(record)); + stats_.retained_payload_bytes += bytes; + return Status::OK(); + } + ORT_CATCH(const std::bad_alloc&) { + ORT_HANDLE_EXCEPTION([&]() { + stats_.collection_complete = false; + IncrementDropCount(stats_.allocation_failure_drops); + status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Runtime workload collector allocation failed."); + }); + } + return status; +} + +Status RuntimeWorkloadCollector::GetSnapshot(RuntimeWorkloadSnapshot& output) const { + std::lock_guard lock(mutex_); + Status status; + ORT_TRY { + RuntimeWorkloadSnapshot snapshot{session_id_, run_id_, records_, stats_, run_state_, limits_}; + output = std::move(snapshot); + return Status::OK(); + } + ORT_CATCH(const std::bad_alloc&) { + ORT_HANDLE_EXCEPTION([&]() { + status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Runtime workload snapshot allocation failed."); + }); + } + return status; +} + +bool RuntimeWorkloadCollector::BeginRun(uint64_t session_id) { + static std::atomic next_run_id{1}; + std::lock_guard lock(mutex_); + if (run_state_ != RuntimeWorkloadRunState::NotStarted || !records_.empty()) { + stats_.collection_complete = false; + IncrementDropCount(stats_.run_reuse_rejections); + return false; + } + session_id_ = session_id; + run_id_ = next_run_id.fetch_add(1, std::memory_order_relaxed); + run_state_ = RuntimeWorkloadRunState::Running; + return true; +} + +void RuntimeWorkloadCollector::EndRun(bool succeeded) { + std::lock_guard lock(mutex_); + run_state_ = succeeded ? RuntimeWorkloadRunState::Succeeded : RuntimeWorkloadRunState::Failed; +} + +void RuntimeWorkloadCollector::RecordCollectionFailure(bool allocation_failure) { + std::lock_guard lock(mutex_); + stats_.collection_complete = false; + IncrementDropCount(allocation_failure ? stats_.allocation_failure_drops : stats_.metadata_failure_drops); +} + +namespace { +std::string GraphPath(const Graph& graph) { + const auto* parent = graph.ParentNode(); + if (parent == nullptr) { + return "main"; + } + for (const auto& [attribute, subgraph] : parent->GetAttributeNameToSubgraphMap()) { + if (subgraph == &graph) { + return GraphPath(*graph.ParentGraph()) + "/" + std::to_string(parent->Index()) + "/" + attribute; + } + } + ORT_THROW("Runtime workload subgraph identity is unavailable."); +} +} // namespace + +void RuntimeWorkloadRecorder::RecordMatMulNBits( + const TensorShape& input_shape, const TensorShape& output_shape, + int64_t m, int64_t n, int64_t k, size_t batch_count, bool bias_present, + int32_t input_element_type, uint32_t bits, uint32_t block_size, + MatMulNBitsRuntimeRoute route, MatMulNBitsRuntimeComputeMode mode) const { + ORT_TRY { + MatMulNBitsRuntimeMetadata metadata; + metadata.input_shape.assign(input_shape.GetDims().begin(), input_shape.GetDims().end()); + metadata.output_shape.assign(output_shape.GetDims().begin(), output_shape.GetDims().end()); + metadata.m = m; + metadata.n = n; + metadata.k = k; + metadata.batch_count = batch_count; + metadata.bias_present = bias_present; + metadata.input_element_type = input_element_type; + metadata.bits = bits; + metadata.block_size = block_size; + metadata.route = route; + metadata.compute_mode = mode; + ORT_IGNORE_RETURN_VALUE(collector->RecordMatMulNBits( + GraphPath(*graph), node->Index(), node->GetExecutionProviderType(), metadata)); + } + ORT_CATCH(const std::bad_alloc&) { + ORT_HANDLE_EXCEPTION([&]() { collector->RecordCollectionFailure(true); }); + } + ORT_CATCH(const std::exception&) { + ORT_HANDLE_EXCEPTION([&]() { collector->RecordCollectionFailure(false); }); + } +} + +} // namespace onnxruntime + +#endif diff --git a/onnxruntime/core/framework/runtime_workload_collector.h b/onnxruntime/core/framework/runtime_workload_collector.h new file mode 100644 index 0000000000000..3372224b17bc5 --- /dev/null +++ b/onnxruntime/core/framework/runtime_workload_collector.h @@ -0,0 +1,151 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#if !defined(ORT_MINIMAL_BUILD) + +#include +#include +#include +#include +#include +#include + +#include "core/common/common.h" +#include "core/common/inlined_containers.h" +#include "core/common/status.h" +#include "core/framework/tensor_shape.h" + +namespace onnxruntime { + +class Graph; +class Node; + +enum class RuntimeWorkloadRunState { + NotStarted, + Running, + Succeeded, + Failed, +}; + +enum class MatMulNBitsRuntimeRoute { + EmptyOutput, + PackedLut, + PackedQNBit, + Unpacked, +}; + +// These describe the effective backend interface, not the activation tensor's dtype +// or a claim that all arithmetic in that route uses this type. +enum class MatMulNBitsRuntimeComputeMode { + NotExecuted, + Float32Gemm, + Lut, + SQNBitFp32, + HQNBitFp16, + BHQNBitBf16, + SQNBitInt8, + HQNBitInt8, +}; + +struct MatMulNBitsRuntimeMetadata { + TensorShapeVector input_shape; + TensorShapeVector output_shape; + int64_t m = 0; + int64_t n = 0; + int64_t k = 0; + size_t batch_count = 0; + bool bias_present = false; + int32_t input_element_type = 0; + uint32_t bits = 0; + uint32_t block_size = 0; + std::optional route; + std::optional compute_mode; +}; + +struct RuntimeWorkloadLimits { + size_t max_records; + // Accounting budget: sizeof(record) plus copied string bytes (including terminators) + // and shape elements. Not a hard allocator/peak-memory limit; capacity and snapshots add overhead. + size_t max_payload_bytes; +}; + +struct RuntimeWorkloadRecord { + // Unique within the collector; commit order, not execution order. + size_t invocation_id; + std::string graph_path; + size_t node_index; + std::string provider; + MatMulNBitsRuntimeMetadata matmul_nbits; +}; + +struct RuntimeWorkloadCollectionStats { + size_t retained_payload_bytes = 0; + // Drop counters saturate rather than wrap. + uint64_t record_limit_drops = 0; + uint64_t payload_limit_drops = 0; + uint64_t allocation_failure_drops = 0; + uint64_t metadata_failure_drops = 0; + uint64_t run_reuse_rejections = 0; + // Completeness of submissions only, not operator/component coverage or run success. + bool collection_complete = true; +}; + +struct RuntimeWorkloadSnapshot { + uint64_t session_id = 0; + uint64_t run_id = 0; + InlinedVector records; + RuntimeWorkloadCollectionStats stats; + // Run outcome is independent of submission completeness. Retained records on Failed runs + // describe only successful kernels before failure, not a complete successful workload. + RuntimeWorkloadRunState run_state = RuntimeWorkloadRunState::NotStarted; + RuntimeWorkloadLimits limits{}; +}; + +// One collector per run. Callers own its lifetime and must keep it alive until all submitters finish. +// This internal storage primitive does not validate operator semantics, estimate counts, or export reports. +class RuntimeWorkloadCollector { + public: + RuntimeWorkloadCollector(uint64_t session_id, uint64_t run_id, RuntimeWorkloadLimits limits); + explicit RuntimeWorkloadCollector(RuntimeWorkloadLimits limits) + : RuntimeWorkloadCollector(0, 0, limits) {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RuntimeWorkloadCollector); + + // Copies metadata without retaining tensors, caller storage, or graph objects. + // Rejection is an explicit diagnostic error; execution integration must not fail inference on it. + [[nodiscard]] Status RecordMatMulNBits(std::string_view graph_path, size_t node_index, + std::string_view provider, const MatMulNBitsRuntimeMetadata& metadata); + // Copies a consistent snapshot. On allocation failure, returns an error without modifying output. + [[nodiscard]] Status GetSnapshot(RuntimeWorkloadSnapshot& output) const; + // Execution integration is single-use. A fresh collector is required for each Run. + // Reusing one (including concurrently) disables capture for the rejected run. + [[nodiscard]] bool BeginRun(uint64_t session_id); + void EndRun(bool succeeded); + void RecordCollectionFailure(bool allocation_failure); + + private: + uint64_t session_id_; + uint64_t run_id_; + const RuntimeWorkloadLimits limits_; + mutable std::mutex mutex_; + InlinedVector records_; + RuntimeWorkloadCollectionStats stats_; + RuntimeWorkloadRunState run_state_ = RuntimeWorkloadRunState::NotStarted; +}; + +// Owned by an execution context, never by a kernel. Shape/identity copies are made only when enabled. +struct RuntimeWorkloadRecorder { + RuntimeWorkloadCollector* collector; + const Graph* graph; + const Node* node; + + void RecordMatMulNBits(const TensorShape& input_shape, const TensorShape& output_shape, + int64_t m, int64_t n, int64_t k, size_t batch_count, bool bias_present, + int32_t input_element_type, uint32_t bits, uint32_t block_size, + MatMulNBitsRuntimeRoute route, MatMulNBitsRuntimeComputeMode mode) const; +}; + +} // namespace onnxruntime + +#endif diff --git a/onnxruntime/core/framework/sequential_executor.cc b/onnxruntime/core/framework/sequential_executor.cc index 68964bf891414..697049c563bb6 100644 --- a/onnxruntime/core/framework/sequential_executor.cc +++ b/onnxruntime/core/framework/sequential_executor.cc @@ -156,8 +156,10 @@ std::string ComposeSeriesName(const GraphViewer& graph_viewer) { class SessionScope { public: friend class KernelScope; - SessionScope(const SessionState& session_state, const ExecutionFrame& frame, profiling::Profiler* run_profiler) + SessionScope(const SessionState& session_state, const ExecutionFrame& frame, profiling::Profiler* run_profiler, + RuntimeWorkloadCollector* workload_collector = nullptr) : session_state_(session_state), + workload_collector_(workload_collector), run_profiler_(run_profiler) #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) , @@ -255,6 +257,7 @@ class SessionScope { } profiling::Profiler* GetRunProfiler() const { return run_profiler_; } + RuntimeWorkloadCollector* GetRuntimeWorkloadCollector() const { return workload_collector_; } void StopProfilingIfEnabled(profiling::EventCategory category, const std::string& event_name, @@ -290,6 +293,7 @@ class SessionScope { private: const SessionState& session_state_; + RuntimeWorkloadCollector* workload_collector_; profiling::Profiler* run_profiler_; TimePoint session_start_; #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) @@ -546,7 +550,8 @@ onnxruntime::Status ExecuteKernel(StreamExecutionContext& ctx, ctx.GetLogger(), terminate_flag, ctx.GetDeviceStream(stream_idx), - session_scope.GetRunProfiler()); + session_scope.GetRunProfiler(), + session_scope.GetRuntimeWorkloadCollector()); onnxruntime::Status status; auto& logger = ctx.GetLogger(); if (p_kernel->IsAsync()) { @@ -694,7 +699,8 @@ onnxruntime::Status ExecuteThePlan(const SessionState& session_state, gsl::span< const bool& terminate_flag, const bool only_execute_path_to_fetches, bool single_thread_mode, - profiling::Profiler* run_profiler) { + profiling::Profiler* run_profiler, + RuntimeWorkloadCollector* workload_collector) { auto* execution_plan = session_state.GetExecutionPlan(); VLOGS(logger, 0) << "Number of streams: " << execution_plan->execution_plan.size(); int32_t valid_streams = 0; @@ -737,7 +743,7 @@ onnxruntime::Status ExecuteThePlan(const SessionState& session_state, gsl::span< ORT_UNUSED_PARAMETER(only_execute_path_to_fetches); #endif - SessionScope session_scope(session_state, ctx.GetExecutionFrame(), run_profiler); + SessionScope session_scope(session_state, ctx.GetExecutionFrame(), run_profiler, workload_collector); auto* tp = single_thread_mode ? nullptr : session_state.GetInterOpThreadPool(); diff --git a/onnxruntime/core/framework/sequential_executor.h b/onnxruntime/core/framework/sequential_executor.h index f42ccc24dabd0..af80a21587c5f 100644 --- a/onnxruntime/core/framework/sequential_executor.h +++ b/onnxruntime/core/framework/sequential_executor.h @@ -47,7 +47,8 @@ onnxruntime::Status ExecuteThePlan(const SessionState& session_state, gsl::span< const bool& terminate_flag, const bool only_execute_path_to_fetches, bool single_thread_mode, - profiling::Profiler* run_profiler = nullptr); + profiling::Profiler* run_profiler = nullptr, + RuntimeWorkloadCollector* workload_collector = nullptr); #ifdef ENABLE_TRAINING onnxruntime::Status PartialExecuteThePlan(const SessionState& session_state, gsl::span feed_mlvalue_idxs, diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index 905cbc077af7c..975d15b703515 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -678,7 +678,8 @@ ExecuteGraphImpl(const SessionState& session_state, #endif const bool only_execute_path_to_fetches = false, Stream* parent_stream = nullptr, - profiling::Profiler* run_profiler = nullptr) { + profiling::Profiler* run_profiler = nullptr, + RuntimeWorkloadCollector* workload_collector = nullptr) { const auto& feeds_fetches_info = feeds_fetches_manager.GetFeedsFetchesInfo(); const auto& device_copy_checks = feeds_fetches_manager.GetDeviceCopyChecks(); #ifdef ORT_ENABLE_STREAM @@ -711,7 +712,7 @@ ExecuteGraphImpl(const SessionState& session_state, only_execute_path_to_fetches, // single thread mode single_thread_mode, - run_profiler)); + run_profiler, workload_collector)); ORT_RETURN_IF_ERROR(status); } else { auto feeds_to_use = feeds; @@ -747,7 +748,7 @@ ExecuteGraphImpl(const SessionState& session_state, terminate_flag, only_execute_path_to_fetches, single_thread_mode, - run_profiler)); + run_profiler, workload_collector)); ORT_RETURN_IF_ERROR(status); InlinedVector fetches_streams; fetches_streams.reserve(feeds_fetches_info.fetches_mlvalue_idxs.size()); @@ -786,7 +787,8 @@ common::Status ExecuteGraph(const SessionState& session_state, #endif bool only_execute_path_to_fetches, Stream* parent_stream, - profiling::Profiler* run_profiler) { + profiling::Profiler* run_profiler, + RuntimeWorkloadCollector* workload_collector) { ORT_RETURN_IF_ERROR(utils::InitializeFeedFetchCopyInfo(session_state, feeds_fetches_manager)); // finalize the copy info using the provided feeds and fetches. will update device_copy_checks in the background @@ -798,14 +800,14 @@ common::Status ExecuteGraph(const SessionState& session_state, device_stream_collection, only_execute_path_to_fetches, parent_stream, - run_profiler); + run_profiler, workload_collector); return retval; #else return ExecuteGraphImpl(session_state, feeds_fetches_manager, feeds, fetches, {}, execution_mode, terminate_flag, logger, only_execute_path_to_fetches, parent_stream, - run_profiler); + run_profiler, workload_collector); #endif } @@ -817,7 +819,8 @@ common::Status ExecuteGraph(const SessionState& session_state, DeviceStreamCollectionHolder& device_stream_collection_holder, #endif const logging::Logger& logger, - profiling::Profiler* run_profiler) { + profiling::Profiler* run_profiler, + RuntimeWorkloadCollector* workload_collector) { return ExecuteGraph(session_state, feeds_fetches_manager, feeds, fetches, execution_mode, run_options.terminate, logger, #ifdef ORT_ENABLE_STREAM @@ -825,7 +828,7 @@ common::Status ExecuteGraph(const SessionState& session_state, #endif run_options.only_execute_path_to_fetches, nullptr /* parent_stream */, - run_profiler); + run_profiler, workload_collector); } #ifdef ENABLE_TRAINING @@ -937,20 +940,21 @@ common::Status ExecuteSubgraph(const SessionState& session_state, const FeedsFet ExecutionMode execution_mode, const bool& terminate_flag, const logging::Logger& logger, Stream* parent_stream, bool sync_subgraph_fetches, - profiling::Profiler* run_profiler) { + profiling::Profiler* run_profiler, + RuntimeWorkloadCollector* workload_collector) { #ifdef ORT_ENABLE_STREAM DeviceStreamCollectionHolder device_stream_collection_holder(&session_state); DeviceStreamCollection* device_stream_collection = device_stream_collection_holder.p_.get(); auto retval = ExecuteGraphImpl(session_state, feeds_fetches_manager, feeds, fetches, fetch_allocators, execution_mode, terminate_flag, logger, device_stream_collection, false, parent_stream, - run_profiler); + run_profiler, workload_collector); if (device_stream_collection) ORT_CHECK_AND_SET_RETVAL(device_stream_collection->CleanUp(false)); #else auto retval = ExecuteGraphImpl(session_state, feeds_fetches_manager, feeds, fetches, fetch_allocators, execution_mode, terminate_flag, logger, false, parent_stream, - run_profiler); + run_profiler, workload_collector); #endif if (retval.IsOK() && sync_subgraph_fetches && parent_stream) { parent_stream->Flush(); diff --git a/onnxruntime/core/framework/utils.h b/onnxruntime/core/framework/utils.h index ba9d0e3c01c95..862ccf73b25b5 100644 --- a/onnxruntime/core/framework/utils.h +++ b/onnxruntime/core/framework/utils.h @@ -21,6 +21,7 @@ class KernelDef; class KernelRegistryManager; class IExecutionProvider; class Node; +class RuntimeWorkloadCollector; #if !defined(ORT_MINIMAL_BUILD) #endif class Tensor; @@ -91,7 +92,8 @@ common::Status ExecuteGraph(const SessionState& session_state, FeedsFetchesManag #endif bool only_execute_path_to_fetches = false, Stream* parent_stream = nullptr, - profiling::Profiler* run_profiler = nullptr); + profiling::Profiler* run_profiler = nullptr, + RuntimeWorkloadCollector* workload_collector = nullptr); common::Status ExecuteGraph(const SessionState& session_state, FeedsFetchesManager& feeds_fetches_manager, gsl::span feeds, std::vector& fetches, @@ -100,7 +102,8 @@ common::Status ExecuteGraph(const SessionState& session_state, FeedsFetchesManag DeviceStreamCollectionHolder& device_stream_collection_holder, #endif const logging::Logger& logger, - profiling::Profiler* run_profiler = nullptr); + profiling::Profiler* run_profiler = nullptr, + RuntimeWorkloadCollector* workload_collector = nullptr); #ifdef ENABLE_TRAINING common::Status ExecutePartialGraph(const SessionState& session_state, FeedsFetchesManager& feeds_fetches_manager, @@ -123,7 +126,8 @@ common::Status ExecuteSubgraph(const SessionState& session_state, const FeedsFet is complete. this is mainly used when the parent kernel depends on the CPU value of the subgraph fetches, i.e. the loop condition*/ bool sync_subgraph_fetches = false, - profiling::Profiler* run_profiler = nullptr); + profiling::Profiler* run_profiler = nullptr, + RuntimeWorkloadCollector* workload_collector = nullptr); bool IsInputOnCpu(const Node& node, const KernelCreateInfo* p_kci, size_t index); bool IsOutputOnCpu(const Node& node, const KernelCreateInfo* p_kci, size_t index); diff --git a/onnxruntime/core/providers/cpu/controlflow/if.cc b/onnxruntime/core/providers/cpu/controlflow/if.cc index 2de1aaa3d853c..8527cfbd9984a 100644 --- a/onnxruntime/core/providers/cpu/controlflow/if.cc +++ b/onnxruntime/core/providers/cpu/controlflow/if.cc @@ -470,7 +470,7 @@ Status IfImpl::Execute(const FeedsFetchesManager& ffm) { ExecutionMode::ORT_SEQUENTIAL, context_.GetTerminateFlag(), context_.Logger(), context_.GetComputeStream(), /*sync_subgraph_fetches*/ false, - context_.GetRunProfiler()); + context_.GetRunProfiler(), context_.GetRuntimeWorkloadCollector()); ORT_RETURN_IF_ERROR(status); diff --git a/onnxruntime/core/providers/cpu/controlflow/loop.cc b/onnxruntime/core/providers/cpu/controlflow/loop.cc index 41ab1bd0ff019..48f56bbe9b270 100644 --- a/onnxruntime/core/providers/cpu/controlflow/loop.cc +++ b/onnxruntime/core/providers/cpu/controlflow/loop.cc @@ -558,7 +558,7 @@ Status LoopImpl::Execute(const FeedsFetchesManager& ffm) { // because the fetch[0] is the loop condition which we need to access on CPU, // have to perofrm a stream sync to make sure the data arrived. true, - context_.GetRunProfiler()); + context_.GetRunProfiler(), context_.GetRuntimeWorkloadCollector()); ORT_RETURN_IF_ERROR(status); condition_mlvalue_ = fetches[0]; diff --git a/onnxruntime/core/providers/cpu/controlflow/scan_utils.cc b/onnxruntime/core/providers/cpu/controlflow/scan_utils.cc index 4a30351696520..1b46e0250ec2e 100644 --- a/onnxruntime/core/providers/cpu/controlflow/scan_utils.cc +++ b/onnxruntime/core/providers/cpu/controlflow/scan_utils.cc @@ -291,7 +291,7 @@ Status IterateSequence(OpKernelContextInternal& context, const SessionState& ses ExecutionMode::ORT_SEQUENTIAL, context.GetTerminateFlag(), context.Logger(), context.GetComputeStream(), /*sync_subgraph_fetches*/ false, - context.GetRunProfiler()); + context.GetRunProfiler(), context.GetRuntimeWorkloadCollector()); ORT_RETURN_IF_ERROR(status); diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index e9015b3e40b09..f242b4c31f247 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -20,6 +20,7 @@ #include "core/common/logging/logging.h" #if !defined(ORT_MINIMAL_BUILD) #include "core/common/json_utils.h" +#include "core/framework/compute_estimate_report.h" #endif #include "core/common/parse_string.h" #include "core/common/path_string.h" @@ -2776,6 +2777,26 @@ common::Status InferenceSession::Initialize() { ORT_TRY { ORT_TELEMETRY_CAPTURE_STATUS_BEGIN(status) LOGS(*session_logger_, INFO) << "Initializing session."; +#if !defined(ORT_MINIMAL_BUILD) + compute_estimate_report_prefix_ = session_options_.config_options.GetConfigOrDefault( + kOrtSessionOptionsComputeEstimateReportPath, ""); + const auto parse_compute_limit = [&](const char* key, const char* default_value, size_t& limit) -> Status { + const auto text = session_options_.config_options.GetConfigOrDefault(key, default_value); + ORT_RETURN_IF_NOT(!text.empty() && + std::all_of(text.begin(), text.end(), [](char c) { return c >= '0' && c <= '9'; }) && + TryParseStringWithClassicLocale(text, limit), + key, " must be a nonnegative decimal integer representable as size_t. Received: ", text); + return Status::OK(); + }; + ORT_RETURN_IF_ERROR(parse_compute_limit(kOrtSessionOptionsComputeEstimateMaxRecords, "10000", + compute_estimate_max_records_)); + ORT_RETURN_IF_ERROR(parse_compute_limit(kOrtSessionOptionsComputeEstimateMaxPayloadBytes, "16777216", + compute_estimate_max_payload_bytes_)); +#else + if (!session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsComputeEstimateReportPath, "").empty()) { + LOGS(*session_logger_, WARNING) << "Compute estimate report export is unsupported in a minimal build."; + } +#endif #if defined(ORT_MINIMAL_BUILD) for (const auto& [key, value] : session_options_.config_options.GetConfigOptionsMap()) { const std::string_view option = key; @@ -3220,6 +3241,9 @@ common::Status InferenceSession::Initialize() { // output via optimized_model_filepath. We still log the session-creation telemetry and record the // end-of-initialization bookkeeping that the normal path emits, so neither is lost by returning early. if (session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionCompileOnly, "0") == "1") { + if (!session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsComputeEstimateReportPath, "").empty()) { + LOGS(*session_logger_, WARNING) << "Compute estimate report export is unsupported in a compile-only session."; + } LOGS(*session_logger_, INFO) << "Compile-only session: skipping session-state finalization. The session is not runnable."; #if defined(ORT_USE_TELEMETRY) @@ -3706,15 +3730,50 @@ Status InferenceSession::Run(const RunOptions& run_options, sequential_run_lock.emplace(session_mutex_); } - return RunImpl(run_options, feed_names, feeds, output_names, p_fetches, p_fetches_device_info, - /*graph_capture_depth=*/0); + RuntimeWorkloadCollector* workload_collector = nullptr; +#if !defined(ORT_MINIMAL_BUILD) + std::optional automatic_collector; + auto* requested_collector = run_options.runtime_workload_collector; + if (!requested_collector && !compute_estimate_report_prefix_.empty()) { + automatic_collector.emplace(RuntimeWorkloadLimits{compute_estimate_max_records_, + compute_estimate_max_payload_bytes_}); + requested_collector = &*automatic_collector; + } + if (requested_collector) { + if (requested_collector->BeginRun(session_id_)) { + workload_collector = requested_collector; + } else { + LOGS(*session_logger_, WARNING) << "Runtime workload capture/export skipped: collector cannot be reused."; + } + } + bool succeeded = false; + auto finish_collection = gsl::finally([&]() { + if (workload_collector) { + workload_collector->EndRun(succeeded); + if (!compute_estimate_report_prefix_.empty()) { + const auto report_status = ExportRuntimeComputeEstimateReport(*workload_collector, + compute_estimate_report_prefix_); + if (!report_status.IsOK()) { + LOGS(*session_logger_, WARNING) << report_status.ErrorMessage(); + } + } + } + }); +#endif + auto status = RunImpl(run_options, feed_names, feeds, output_names, p_fetches, p_fetches_device_info, + /*graph_capture_depth=*/0, workload_collector); +#if !defined(ORT_MINIMAL_BUILD) + succeeded = status.IsOK(); +#endif + return status; } Status InferenceSession::RunImpl(const RunOptions& run_options, gsl::span feed_names, gsl::span feeds, gsl::span output_names, std::vector* p_fetches, const std::vector* p_fetches_device_info, - int graph_capture_depth) { + int graph_capture_depth, + RuntimeWorkloadCollector* workload_collector) { // Ignore run-level profiling request if session-level profiling is already enabled. std::optional run_profiler; if (run_options.enable_profiling && session_profiler_.IsEnabled()) { @@ -3902,7 +3961,7 @@ Status InferenceSession::RunImpl(const RunOptions& run_options, device_stream_collection_holder, #endif run_logger, - run_profiler ? &*run_profiler : nullptr); + run_profiler ? &*run_profiler : nullptr, workload_collector); } // info all execution providers InferenceSession:Run ended @@ -4059,7 +4118,7 @@ Status InferenceSession::RunImpl(const RunOptions& run_options, RunOptions recursive_run_options{run_options}; recursive_run_options.enable_profiling = false; ORT_RETURN_IF_ERROR(RunImpl(recursive_run_options, feed_names, feeds, output_names, p_fetches, - p_fetches_device_info, graph_capture_depth + 1)); + p_fetches_device_info, graph_capture_depth + 1, workload_collector)); } // Log runtime error telemetry if the return value is not OK diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index c5c9d0fab05c0..7bdc2635d1132 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -805,7 +805,8 @@ class InferenceSession { gsl::span feeds, gsl::span output_names, std::vector* p_fetches, const std::vector* p_fetches_device_info, - int graph_capture_depth); + int graph_capture_depth, + RuntimeWorkloadCollector* workload_collector = nullptr); void SetLoggingManager(const SessionOptions& session_options, const Environment& session_env); @@ -1012,6 +1013,12 @@ class InferenceSession { static std::atomic global_session_id_; // a monotonically increasing session id uint32_t session_id_; // the current session's id +#if !defined(ORT_MINIMAL_BUILD) + std::string compute_estimate_report_prefix_; + size_t compute_estimate_max_records_ = 10000; + size_t compute_estimate_max_payload_bytes_ = 16 * 1024 * 1024; +#endif + struct Telemetry { Telemetry(); uint32_t total_runs_since_last_ = 0; // the total number of Run() calls since the last report diff --git a/onnxruntime/test/contrib_ops/compute_estimate_report_test.cc b/onnxruntime/test/contrib_ops/compute_estimate_report_test.cc new file mode 100644 index 0000000000000..44ebd2f132631 --- /dev/null +++ b/onnxruntime/test/contrib_ops/compute_estimate_report_test.cc @@ -0,0 +1,729 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include +#include "nlohmann/json.hpp" +#include "core/common/path_string.h" +#include "core/framework/allocator.h" +#include "core/framework/compute_estimate_report.h" +#include "core/framework/runtime_workload_collector.h" +#include "core/graph/constants.h" +#include "core/graph/indexed_sub_graph.h" +#include "core/graph/model.h" +#include "core/session/environment.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/capturing_sink.h" +#include "test/test_environment.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" +#include "test/util/include/file_util.h" +#include "test/util/include/inference_session_wrapper.h" + +namespace onnxruntime { +namespace test { +namespace { + +using Json = nlohmann::json; + +class ComputeEstimateReportTest : public ::testing::Test { + protected: + ComputeEstimateReportTest() + : model_("compute_report", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 17}, {kMSDomain, 1}}, {}, DefaultLoggingManager().DefaultLogger()), + builder_(model_.MainGraph()) { + static std::atomic sequence{0}; + const auto stamp = std::chrono::steady_clock::now().time_since_epoch().count(); + path_ = std::filesystem::temp_directory_path() / + ("ort_compute_report_" + std::to_string(stamp) + "_" + std::to_string(sequence++) + ".json"); + cleanup_ = ScopedFileDeleter(path_.native()); + } + + ~ComputeEstimateReportTest() override { + for (const auto& path : RuntimeReportPaths()) { + std::error_code error; + std::filesystem::remove(path, error); + EXPECT_FALSE(error) << error.message(); + } + } + + std::vector RuntimeReportPaths() const { + std::vector paths; + const auto prefix = path_.filename().native() + ToPathString("_"); + for (const auto& entry : std::filesystem::directory_iterator(path_.parent_path())) { + if (entry.path().filename().native().find(prefix) == 0 && entry.path().extension() == ".json") { + paths.push_back(entry.path()); + } + } + return paths; + } + + Node& AddMatMul(std::optional> shape = std::vector{2, 3, 64}, + bool add_identity = false, bool bias = false) { + auto* a = builder_.MakeInput(shape); + auto* b = builder_.MakeInitializer({32, 2, 16}, std::vector(1024, 0x88)); + auto* scales = builder_.MakeInitializer({32, 2}, std::vector(64, 1.0f)); + auto* output = add_identity ? builder_.MakeIntermediate() : builder_.MakeOutput(); + std::vector inputs{a, b, scales}; + if (bias) { + auto& missing = builder_.graph_.GetOrCreateNodeArg("", nullptr); + inputs.push_back(&missing); + inputs.push_back(&missing); + inputs.push_back(builder_.MakeInitializer({32}, std::vector(32, 1.0f))); + } + auto& node = builder_.graph_.AddNode(builder_.graph_.GenerateNodeName("projection \"quoted\"\\\n"), + "MatMulNBits", "", inputs, {output}, nullptr, kMSDomain); + node.AddAttribute("K", int64_t{64}); + node.AddAttribute("N", int64_t{32}); + node.AddAttribute("bits", int64_t{4}); + node.AddAttribute("block_size", int64_t{32}); + if (add_identity) { + builder_.AddNode("Identity", {output}, {builder_.MakeOutput()}); + } + builder_.SetGraphOutputs(); + return node; + } + + Json Report() { + std::string value; + const auto status = BuildComputeEstimateReport(model_.MainGraph(), value); + EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); + return Json::parse(value); + } + + Json ReadReport() { + const auto paths = RuntimeReportPaths(); + ORT_ENFORCE(paths.size() == 1, "Expected exactly one runtime report, found ", paths.size()); + std::ifstream stream(paths[0]); + EXPECT_TRUE(stream.is_open()); + return Json::parse(stream); + } + + std::string ModelBytes() { + return model_.ToProto().SerializeAsString(); + } + + Status Initialize(InferenceSession& session) { + std::istringstream stream(ModelBytes()); + ORT_RETURN_IF_ERROR(session.Load(stream)); + return session.Initialize(); + } + + SessionOptions ReportOptions() { + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + options.intra_op_param.thread_pool_size = 1; + const auto utf8_path = ToUTF8String(path_.native()); + EXPECT_TRUE(options.config_options.AddConfigEntry(kOrtSessionOptionsComputeEstimateReportPath, + utf8_path.c_str()) + .IsOK()); + return options; + } + + Status Run(InferenceSession& session, const std::vector& shape = {2, 3, 64}, + bool terminate = false, RuntimeWorkloadCollector* collector = nullptr) { + OrtValue input; + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape(shape), std::make_shared(), input); + std::fill_n(input.GetMutable()->MutableData(), + static_cast(input.Get().Shape().Size()), 1.0f); + RunOptions options; + options.terminate = terminate; + options.runtime_workload_collector = collector; + std::vector outputs; + const std::vector output_names{model_.MainGraph().GetOutputs()[0]->Name()}; + return session.Run(options, {{model_.MainGraph().GetInputs()[0]->Name(), input}}, + output_names, &outputs); + } + + Model model_; + ModelTestBuilder builder_; + std::filesystem::path path_; + ScopedFileDeleter cleanup_; +}; + +TEST_F(ComputeEstimateReportTest, ComponentContractAndDeterminism) { + auto& node = AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + EXPECT_EQ(report["schema_version"], 2); + EXPECT_EQ(report["estimator_support"], "matmul_nbits_core_and_bias"); + ASSERT_EQ(report["nodes"].size(), 1u); + const auto& entry = report["nodes"][0]; + EXPECT_EQ(entry["node_name"], node.Name()); + EXPECT_EQ(entry["graph_path"], "main"); + EXPECT_EQ(entry["opset_version"], 1); + const auto& component = entry["components"][0]; + EXPECT_EQ(component["operation_count"], "24576"); + EXPECT_EQ(component["metric"], "dense_equivalent_ops"); + EXPECT_EQ(component["input_basis"], "static_concrete"); + EXPECT_EQ(component["execution_phase"], "per_invocation"); + EXPECT_EQ(component["basis"], "logical_semantics"); + EXPECT_EQ(component["convention_version"], "core_arithmetic_v1"); + EXPECT_EQ(component["excluded_components"], + Json::array({"dequantization", "bias", "packing", "conversion", "padding", "route_specific_extra_work"})); + EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "24576"); + EXPECT_EQ(report["summary"]["complete_for_declared_graph_and_component_scope"], true); + EXPECT_EQ(report["summary"]["whole_operator_work_complete"], false); + ASSERT_EQ(entry["components"].size(), 2u); + const auto& bias = entry["components"][1]; + EXPECT_EQ(bias["component"], "bias"); + EXPECT_EQ(bias["metric"], "floating_point_ops"); + EXPECT_EQ(bias["operation_count"], "0"); + EXPECT_EQ(bias["input_basis"], "optional_input_absent"); + EXPECT_EQ(bias["convention_version"], "bias_add_v1"); + EXPECT_EQ(report, Report()); +} + +TEST_F(ComputeEstimateReportTest, BiasIsSeparateFromCoreSubtotal) { + AddMatMul(std::vector{2, 3, 64}, false, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + const auto& bias = report["nodes"][0]["components"][1]; + EXPECT_EQ(bias["operation_count"], "192"); + EXPECT_EQ(bias["input_basis"], "static_concrete"); + EXPECT_EQ(bias["basis"], "logical_semantics"); + EXPECT_EQ(bias["execution_phase"], "per_invocation"); + EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "24576"); + const auto& summary = report["summary"]["component_summaries"][1]; + EXPECT_EQ(summary["component"], "bias"); + EXPECT_EQ(summary["known_operation_count"], "192"); + EXPECT_EQ(summary["estimated_component_count"], 1); + EXPECT_EQ(summary["unavailable_component_count"], 0); + EXPECT_EQ(summary["complete_for_declared_component_scope"], true); + EXPECT_EQ(report["summary"]["whole_operator_work_complete"], false); +} + +TEST_F(ComputeEstimateReportTest, UnknownBiasGeometryHasExplicitCoverageGap) { + AddMatMul(std::vector{-1, 64}, false, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + const auto& entry = report["nodes"][0]; + EXPECT_EQ(entry["availability"], "unavailable"); + EXPECT_EQ(entry["components"][1]["unavailable_reason"], "unknown_shape"); + EXPECT_FALSE(entry["components"][1].contains("operation_count")); + const auto& summary = report["summary"]["component_summaries"][1]; + EXPECT_EQ(summary["known_operation_count"], "0"); + EXPECT_EQ(summary["unavailable_component_count"], 1); + EXPECT_EQ(summary["unavailable_reason_counts"]["unknown_shape"], 1); + EXPECT_EQ(summary["complete_for_declared_component_scope"], false); +} + +TEST_F(ComputeEstimateReportTest, EmptyOutputHasZeroBias) { + AddMatMul(std::vector{0, 64}, false, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + EXPECT_EQ(report["nodes"][0]["components"][1]["operation_count"], "0"); + EXPECT_EQ(report["summary"]["complete_for_declared_graph_and_component_scope"], true); +} + +TEST_F(ComputeEstimateReportTest, CoreOverflowRetainsKnownBias) { + auto& node = AddMatMul(std::vector{1, std::numeric_limits::max()}, false, true); + node.AddAttribute("K", std::numeric_limits::max()); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + EXPECT_EQ(report["nodes"][0]["availability"], "partial"); + EXPECT_EQ(report["nodes"][0]["components"][0]["unavailable_reason"], "overflow"); + EXPECT_EQ(report["nodes"][0]["components"][1]["operation_count"], "32"); + EXPECT_EQ(report["summary"]["partially_estimated_node_count"], 1); + EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "0"); + EXPECT_EQ(report["summary"]["component_summaries"][1]["known_operation_count"], "32"); + EXPECT_EQ(report["summary"]["complete_for_declared_graph_and_component_scope"], false); +} + +TEST_F(ComputeEstimateReportTest, BiasSubtotalOverflowIsSeparateAndExplicit) { + for (int i = 0; i < 3; ++i) { + auto& node = AddMatMul(std::vector{std::numeric_limits::max() / 32, 0}, false, true); + node.AddAttribute("K", int64_t{0}); + } + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "0"); + const auto& summary = report["summary"]["component_summaries"][1]; + EXPECT_EQ(summary["subtotal_availability"], "unavailable"); + EXPECT_EQ(summary["subtotal_unavailable_reason"], "overflow"); + EXPECT_FALSE(summary.contains("known_operation_count")); + EXPECT_EQ(report["summary"]["complete_for_declared_graph_and_component_scope"], false); +} + +TEST_F(ComputeEstimateReportTest, KnownCoreDoesNotHideBiasOverflow) { + auto& node = AddMatMul(std::vector{std::numeric_limits::max(), 0}, false, true); + node.AddAttribute("K", int64_t{0}); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + const auto& entry = report["nodes"][0]; + EXPECT_EQ(entry["availability"], "partial"); + EXPECT_EQ(entry["unavailable_reason"], "overflow"); + EXPECT_EQ(entry["components"][0]["operation_count"], "0"); + EXPECT_EQ(entry["components"][1]["unavailable_reason"], "overflow"); + EXPECT_FALSE(entry["components"][1].contains("operation_count")); + EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "0"); + EXPECT_EQ(report["summary"]["component_summaries"][1]["unavailable_component_count"], 1); + EXPECT_EQ(report["summary"]["complete_for_declared_graph_and_component_scope"], false); +} + +TEST_F(ComputeEstimateReportTest, UnknownShapesHaveNoCount) { + AddMatMul(std::vector{-1, 64}); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + const auto& entry = report["nodes"][0]; + EXPECT_EQ(entry["unavailable_reason"], "unknown_shape"); + EXPECT_EQ(entry["availability"], "partial"); + EXPECT_EQ(entry["components"][1]["operation_count"], "0"); + EXPECT_EQ(report["summary"]["partially_estimated_node_count"], 1); + EXPECT_FALSE(entry["components"][0].contains("operation_count")); + EXPECT_EQ(entry["input_shapes"][0]["dimensions"], Json::array({nullptr, 64})); + EXPECT_EQ(report["summary"]["unavailable_reason_counts"]["unknown_shape"], 1); + EXPECT_EQ(report["summary"]["complete_for_declared_graph_and_component_scope"], false); +} + +TEST_F(ComputeEstimateReportTest, PartialCoverage) { + AddMatMul(std::vector{2, 3, 64}, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + EXPECT_EQ(report["summary"]["reported_node_count"], 2); + EXPECT_EQ(report["summary"]["estimated_node_count"], 1); + EXPECT_EQ(report["summary"]["unavailable_node_count"], 1); + EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "24576"); + EXPECT_EQ(report["summary"]["complete_for_declared_graph_and_component_scope"], false); +} + +TEST_F(ComputeEstimateReportTest, OpaquePartitionDoesNotCountUnderlyingCore) { + const auto& node = AddMatMul(std::vector{2, 3, 64}, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + IndexedSubGraph partition; + auto definition = std::make_unique(); + definition->name = "CompiledProjection"; + definition->domain = "test.partition"; + definition->since_version = 1; + for (const auto* input : node.InputDefs()) { + definition->inputs.push_back(input->Name()); + } + definition->outputs.push_back(model_.MainGraph().GetOutputs()[0]->Name()); + for (const auto& original : model_.MainGraph().Nodes()) { + partition.nodes.push_back(original.Index()); + } + partition.SetMetaDef(std::move(definition)); + model_.MainGraph().FuseSubGraph(partition, "compiled_projection"); + const auto report = Report(); + ASSERT_EQ(report["nodes"].size(), 1u); + EXPECT_EQ(report["nodes"][0]["unavailable_reason"], "opaque_partition"); + EXPECT_TRUE(report["nodes"][0]["components"].empty()); + EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "0"); +} + +TEST_F(ComputeEstimateReportTest, SubtotalOverflowIsExplicit) { + constexpr int64_t m = (int64_t{1} << 61) - 1; + // Each count is 2 * ((2^61) - 1); add enough nodes to exceed uint64 without any node overflowing. + for (int i = 0; i < 5; ++i) { + auto& node = AddMatMul(std::vector{m, 1}); + node.AddAttribute("K", int64_t{1}); + node.AddAttribute("N", int64_t{1}); + } + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + EXPECT_EQ(report["summary"]["estimated_node_count"], 5); + EXPECT_EQ(report["summary"]["subtotal_unavailable_reason"], "overflow"); + EXPECT_FALSE(report["summary"].contains("known_dense_equivalent_ops")); +} + +TEST_F(ComputeEstimateReportTest, NestedGraphsAreSeparateScopes) { + auto* condition = builder_.MakeInput(std::vector{}); + auto* captured = builder_.MakeInput(std::vector{2, 3, 64}); + model_.MainGraph().SetInputs({condition, captured}); + auto* output = builder_.MakeOutput(); + auto& node = builder_.AddNode("If", {condition}, {output}); + for (const auto* attribute : {"then_branch", "else_branch"}) { + Model branch("branch", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 17}, {kMSDomain, 1}}, {}, DefaultLoggingManager().DefaultLogger()); + ModelTestBuilder branch_builder(branch.MainGraph()); + auto* input = branch_builder.MakeInput(std::vector{2, 3, 64}, captured->Name()); + auto* weights = branch_builder.MakeInitializer({32, 2, 16}, std::vector(1024, 0x88)); + auto* scales = branch_builder.MakeInitializer({32, 2}, std::vector(64, 1.0f)); + auto& matmul = branch_builder.AddNode("MatMulNBits", {input, weights, scales}, + {branch_builder.MakeOutput()}, kMSDomain); + matmul.AddAttribute("K", int64_t{64}); + matmul.AddAttribute("N", int64_t{32}); + matmul.AddAttribute("bits", int64_t{4}); + matmul.AddAttribute("block_size", int64_t{32}); + branch_builder.SetGraphOutputs(); + ASSERT_STATUS_OK(branch.MainGraph().Resolve()); + auto proto = branch.ToProto().graph(); + proto.clear_input(); + node.AddAttribute(attribute, proto); + } + builder_.SetGraphOutputs(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto report = Report(); + ASSERT_EQ(report["graphs"].size(), 3u); + EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "0"); + EXPECT_EQ(report["summary"]["contains_subgraphs"], true); + EXPECT_EQ(report["nodes"][0]["unavailable_reason"], "runtime_execution_required"); + EXPECT_EQ(report["graphs"][1]["graph_path"], "main/0/else_branch"); + EXPECT_EQ(report["graphs"][2]["graph_path"], "main/0/then_branch"); + for (size_t i = 1; i < 3; ++i) { + EXPECT_EQ(report["graphs"][i]["summary"]["known_dense_equivalent_ops"], "24576"); + EXPECT_EQ(report["graphs"][i]["summary"]["per_run_total_available"], false); + } +} + +TEST_F(ComputeEstimateReportTest, SessionExportsExecutedInvocationsAndRetainsInference) { + AddMatMul(std::vector{2, 3, 64}, true, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + auto options = ReportOptions(); + options.graph_optimization_level = TransformerLevel::Level1; + InferenceSessionWrapper session(options, GetEnvironment()); + ASSERT_STATUS_OK(Initialize(session)); + EXPECT_FALSE(std::filesystem::exists(path_)); + EXPECT_TRUE(RuntimeReportPaths().empty()); + + const auto& input_name = session.GetGraph().GetInputs()[0]->Name(); + const auto& output_name = session.GetGraph().GetOutputs()[0]->Name(); + std::vector values(2 * 3 * 64, 1.0f); + OrtValue input; + auto allocator = std::make_shared(); + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape({2, 3, 64}), allocator, input); + std::copy(values.begin(), values.end(), input.GetMutable()->MutableData()); + std::vector outputs; + const std::vector output_names{output_name}; + ASSERT_STATUS_OK(session.Run(RunOptions{}, {{input_name, input}}, output_names, &outputs)); + ASSERT_EQ(outputs.size(), 1u); + const auto& result = outputs[0].Get(); + EXPECT_EQ(result.Shape(), TensorShape({2, 3, 32})); + for (size_t i = 0; i < 2 * 3 * 32; ++i) { + EXPECT_EQ(result.Data()[i], 1.0f); + } + const auto report = ReadReport(); + EXPECT_EQ(report["schema_version"], 3); + EXPECT_EQ(report["run_state"], "succeeded"); + EXPECT_EQ(report["session_id"], session.GetCurrentSessionId()); + EXPECT_EQ(report["collection"]["max_records"], 10000); + EXPECT_EQ(report["collection"]["max_payload_bytes"], 16777216); + ASSERT_EQ(report["invocations"].size(), 1u); + EXPECT_EQ(report["invocations"][0]["provider"], kCpuExecutionProvider); + EXPECT_EQ(report["invocations"][0]["input_shape"], Json::array({2, 3, 64})); + EXPECT_EQ(report["invocations"][0]["components"][0]["operation_count"], "24576"); + EXPECT_EQ(report["invocations"][0]["components"][1]["operation_count"], "192"); + EXPECT_TRUE(report["invocations"][0].contains("route")); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], true); + EXPECT_EQ(report["whole_run_work_complete"], false); + RecordProperty("runtime_compute_estimate_report", report.dump(2)); +} + +TEST_F(ComputeEstimateReportTest, DisabledExportDoesNotWrite) { + AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + SessionOptions options; + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(Initialize(session)); + ASSERT_STATUS_OK(Run(session)); + EXPECT_FALSE(std::filesystem::exists(path_)); + EXPECT_TRUE(RuntimeReportPaths().empty()); +} + +TEST_F(ComputeEstimateReportTest, WriteFailureWarnsAndInferenceSucceeds) { + AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + auto options = ReportOptions(); + // A file cannot be the parent directory of a report. + { + std::ofstream parent(path_); + } + const auto invalid_path = ToUTF8String((path_ / "report").native()); + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsComputeEstimateReportPath, + invalid_path.c_str())); + auto sink = std::make_unique(); + auto* captured = sink.get(); + auto logging_manager = std::make_unique( + std::move(sink), logging::Severity::kWARNING, false, logging::LoggingManager::InstanceType::Temporal); + std::unique_ptr environment; + ASSERT_STATUS_OK(Environment::Create(std::move(logging_manager), environment)); + InferenceSession session(options, *environment); + ASSERT_STATUS_OK(Initialize(session)); + EXPECT_FALSE(std::any_of(captured->Messages().begin(), captured->Messages().end(), [](const std::string& message) { + return message.find("Failed to open compute estimate report") != std::string::npos; + })); + ASSERT_STATUS_OK(Run(session)); + EXPECT_TRUE(std::any_of(captured->Messages().begin(), captured->Messages().end(), [](const std::string& message) { + return message.find("Failed to open compute estimate report") != std::string::npos; + })); +} + +TEST_F(ComputeEstimateReportTest, DynamicAndConcurrentRunsWriteDistinctReports) { + AddMatMul(std::vector{-1, 3, 64}, false, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + InferenceSession session(ReportOptions(), GetEnvironment()); + ASSERT_STATUS_OK(Initialize(session)); + std::atomic failures{0}; + std::array threads; + for (size_t i = 0; i < threads.size(); ++i) { + threads[i] = std::thread([&, i]() { + if (!Run(session, {static_cast(i + 1), 3, 64}).IsOK()) ++failures; + }); + } + for (auto& thread : threads) thread.join(); + EXPECT_EQ(failures, 0u); + const auto paths = RuntimeReportPaths(); + ASSERT_EQ(paths.size(), threads.size()); + std::vector ids; + for (const auto& path : paths) { + std::ifstream stream(path); + const auto report = Json::parse(stream); + const auto batches = report["invocations"][0]["input_shape"][0].get(); + EXPECT_EQ(report["component_summaries"][0]["known_operation_count"], std::to_string(batches * 12288)); + EXPECT_EQ(report["component_summaries"][1]["known_operation_count"], std::to_string(batches * 96)); + ids.push_back(report["run_id"].get()); + } + std::sort(ids.begin(), ids.end()); + EXPECT_EQ(std::adjacent_find(ids.begin(), ids.end()), ids.end()); + ASSERT_STATUS_OK(Run(session)); + EXPECT_EQ(RuntimeReportPaths().size(), threads.size() + 1); + InferenceSession another_session(ReportOptions(), GetEnvironment()); + ASSERT_STATUS_OK(Initialize(another_session)); + ASSERT_STATUS_OK(Run(another_session)); + EXPECT_EQ(RuntimeReportPaths().size(), threads.size() + 2); +} + +TEST_F(ComputeEstimateReportTest, AutomaticLimitsAndFailedRunsAreExplicit) { + AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + for (const auto* limit : {kOrtSessionOptionsComputeEstimateMaxRecords, + kOrtSessionOptionsComputeEstimateMaxPayloadBytes}) { + auto options = ReportOptions(); + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(limit, "0")); + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(Initialize(session)); + ASSERT_STATUS_OK(Run(session)); + auto report = ReadReport(); + EXPECT_TRUE(report["invocations"].empty()); + EXPECT_EQ(report["run_state"], "succeeded"); + EXPECT_EQ(report["collection"]["collection_complete"], false); + EXPECT_EQ(report["collection"][limit == kOrtSessionOptionsComputeEstimateMaxRecords + ? "record_limit_drops" + : "payload_limit_drops"], + 1); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], false); + for (const auto& path : RuntimeReportPaths()) ASSERT_TRUE(std::filesystem::remove(path)); + } + InferenceSession session(ReportOptions(), GetEnvironment()); + ASSERT_STATUS_OK(Initialize(session)); + EXPECT_FALSE(Run(session, {2, 3, 64}, true).IsOK()); + const auto report = ReadReport(); + EXPECT_EQ(report["run_state"], "failed"); + EXPECT_TRUE(report["invocations"].empty()); + EXPECT_EQ(report["collection"]["collection_complete"], true); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], false); +} + +TEST_F(ComputeEstimateReportTest, InvalidLimitsFailInitialization) { + AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + for (const auto* key : {kOrtSessionOptionsComputeEstimateMaxRecords, + kOrtSessionOptionsComputeEstimateMaxPayloadBytes}) { + for (const auto* value : {"", "-1", "1x", " 2", "18446744073709551616"}) { + auto options = ReportOptions(); + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(key, value)); + InferenceSession session(options, GetEnvironment()); + EXPECT_FALSE(Initialize(session).IsOK()) << key << "=" << value; + } + } + EXPECT_TRUE(RuntimeReportPaths().empty()); +} + +TEST_F(ComputeEstimateReportTest, ExplicitCollectorIsExportedWithoutChangingItsLimits) { + AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + auto options = ReportOptions(); + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsComputeEstimateMaxRecords, "0")); + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(Initialize(session)); + RuntimeWorkloadCollector collector({10, 8192}); + ASSERT_STATUS_OK(Run(session, {2, 3, 64}, false, &collector)); + const auto report = ReadReport(); + EXPECT_EQ(report["invocations"].size(), 1u); + EXPECT_EQ(report["collection"]["max_records"], 10); + EXPECT_EQ(report["collection"]["max_payload_bytes"], 8192); + ASSERT_STATUS_OK(Run(session, {2, 3, 64}, false, &collector)); + EXPECT_EQ(RuntimeReportPaths().size(), 1u); +} + +TEST_F(ComputeEstimateReportTest, Utf8PathAndOverwrite) { + AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + const auto unicode_path = path_.parent_path() / + ToPathString(path_.stem().string() + "_\xE6\xB5\x8B\xE8\xAF\x95.json"); + auto cleanup = gsl::finally([&]() { std::filesystem::remove(unicode_path); }); + const auto utf8_path = ToUTF8String(unicode_path.native()); + ASSERT_STATUS_OK(WriteComputeEstimateReport(model_.MainGraph(), utf8_path)); + ASSERT_STATUS_OK(WriteComputeEstimateReport(model_.MainGraph(), utf8_path)); + std::ifstream stream(unicode_path); + const auto report = Json::parse(stream); + EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "24576"); + path_ = unicode_path; + InferenceSession session(ReportOptions(), GetEnvironment()); + ASSERT_STATUS_OK(Initialize(session)); + EXPECT_TRUE(RuntimeReportPaths().empty()); + ASSERT_STATUS_OK(Run(session)); + const auto first = ReadReport(); + EXPECT_EQ(first["component_summaries"][0]["known_operation_count"], "24576"); + ASSERT_STATUS_OK(Run(session)); + EXPECT_EQ(RuntimeReportPaths().size(), 2u); +} + +RuntimeWorkloadSnapshot RuntimeSnapshot() { + RuntimeWorkloadSnapshot snapshot; + snapshot.session_id = 7; + snapshot.run_id = 11; + snapshot.run_state = RuntimeWorkloadRunState::Succeeded; + snapshot.limits = {10, 8192}; + MatMulNBitsRuntimeMetadata metadata; + metadata.input_shape = {2, 3, 64}; + metadata.output_shape = {2, 3, 32}; + metadata.m = 6; + metadata.n = 32; + metadata.k = 64; + metadata.batch_count = 1; + metadata.bias_present = true; + metadata.bits = 4; + metadata.block_size = 32; + metadata.input_element_type = ONNX_NAMESPACE::TensorProto::FLOAT; + snapshot.records.push_back({0, "main/0/body", 1, kCpuExecutionProvider, metadata}); + return snapshot; +} + +Json RuntimeReport(const RuntimeWorkloadSnapshot& snapshot) { + std::string text; + ORT_THROW_IF_ERROR(BuildComputeEstimateReport(snapshot, text)); + return Json::parse(text); +} + +TEST(RuntimeComputeEstimateReportTest, CountsAndRouteProvenanceAreIndependent) { + auto snapshot = RuntimeSnapshot(); + for (const auto route : {MatMulNBitsRuntimeRoute::PackedLut, MatMulNBitsRuntimeRoute::PackedQNBit, + MatMulNBitsRuntimeRoute::Unpacked, MatMulNBitsRuntimeRoute::EmptyOutput}) { + snapshot.records[0].matmul_nbits.route = route; + const auto report = RuntimeReport(snapshot); + EXPECT_EQ(report["invocations"][0]["components"][0]["operation_count"], "24576"); + EXPECT_EQ(report["invocations"][0]["components"][1]["operation_count"], "192"); + EXPECT_EQ(report["component_summaries"][0]["known_operation_count"], "24576"); + EXPECT_EQ(report["route_adjustment_applied"], false); + EXPECT_EQ(report, RuntimeReport(snapshot)); + } + snapshot.records[0].matmul_nbits.route.reset(); + const auto report = RuntimeReport(snapshot); + EXPECT_FALSE(report["invocations"][0].contains("route")); + EXPECT_FALSE(report["invocations"][0].contains("compute_mode")); + EXPECT_EQ(report["invocations"][0]["availability"], "available"); +} + +TEST(RuntimeComputeEstimateReportTest, SubgraphsAndRepeatedInvocationsSumOnlyExecutedRecords) { + auto snapshot = RuntimeSnapshot(); + auto record = snapshot.records[0]; + record.invocation_id = 1; + record.graph_path = "main/2/then_branch"; + snapshot.records.push_back(record); + record.invocation_id = 2; + snapshot.records.push_back(record); + const auto report = RuntimeReport(snapshot); + EXPECT_EQ(report["invocations"].size(), 3u); + EXPECT_EQ(report["component_summaries"][0]["known_operation_count"], "73728"); + EXPECT_EQ(report["component_summaries"][1]["known_operation_count"], "576"); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], true); +} + +TEST(RuntimeComputeEstimateReportTest, UnknownInvalidAndUnsupportedAreNotZero) { + auto snapshot = RuntimeSnapshot(); + snapshot.records[0].matmul_nbits.input_shape = {-1, 3, 64}; + auto report = RuntimeReport(snapshot); + EXPECT_EQ(report["invocations"][0]["availability"], "unavailable"); + EXPECT_EQ(report["invocations"][0]["components"][0]["unavailable_reason"], "unknown_shape"); + EXPECT_FALSE(report["invocations"][0]["components"][0].contains("operation_count")); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], false); + snapshot.records[0].matmul_nbits.input_shape = {2, 3, 63}; + report = RuntimeReport(snapshot); + EXPECT_EQ(report["invocations"][0]["components"][0]["unavailable_reason"], "invalid_metadata"); + snapshot.records[0].provider = kCudaExecutionProvider; + report = RuntimeReport(snapshot); + EXPECT_EQ(report["invocations"][0]["components"][0]["unavailable_reason"], "unsupported_provider"); +} + +TEST(RuntimeComputeEstimateReportTest, EmptyOutputAndZeroReductionKeepBiasSemantics) { + auto snapshot = RuntimeSnapshot(); + snapshot.records[0].matmul_nbits.input_shape = {2, 0, 64}; + auto report = RuntimeReport(snapshot); + EXPECT_EQ(report["invocations"][0]["components"][0]["operation_count"], "0"); + EXPECT_EQ(report["invocations"][0]["components"][1]["operation_count"], "0"); + snapshot.records[0].matmul_nbits.input_shape = {2, 3, 0}; + snapshot.records[0].matmul_nbits.k = 0; + report = RuntimeReport(snapshot); + EXPECT_EQ(report["invocations"][0]["components"][0]["operation_count"], "0"); + EXPECT_EQ(report["invocations"][0]["components"][1]["operation_count"], "192"); + snapshot.records[0].matmul_nbits.bias_present = false; + EXPECT_EQ(RuntimeReport(snapshot)["invocations"][0]["components"][1]["operation_count"], "0"); +} + +TEST(RuntimeComputeEstimateReportTest, InvocationAndSubtotalOverflowRemainDistinct) { + auto snapshot = RuntimeSnapshot(); + auto& metadata = snapshot.records[0].matmul_nbits; + metadata.input_shape = {std::numeric_limits::max(), 1}; + metadata.k = 1; + metadata.n = 2; + auto report = RuntimeReport(snapshot); + EXPECT_EQ(report["invocations"][0]["components"][0]["unavailable_reason"], "overflow"); + EXPECT_EQ(report["invocations"][0]["components"][1]["operation_count"], "18446744073709551614"); + EXPECT_EQ(report["invocations"][0]["availability"], "partial"); + metadata.n = 1; + auto record = snapshot.records[0]; + record.invocation_id = 1; + snapshot.records.push_back(record); + report = RuntimeReport(snapshot); + EXPECT_EQ(report["invocations"][0]["components"][0]["operation_count"], "18446744073709551614"); + EXPECT_EQ(report["component_summaries"][0]["subtotal_unavailable_reason"], "overflow"); + EXPECT_FALSE(report["component_summaries"][0].contains("known_operation_count")); + EXPECT_EQ(report["component_summaries"][1]["known_operation_count"], "18446744073709551614"); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], false); +} + +TEST(RuntimeComputeEstimateReportTest, FailedRunningAndIncompleteDoNotClaimCompleteTotals) { + auto snapshot = RuntimeSnapshot(); + for (const auto state : {RuntimeWorkloadRunState::Failed, RuntimeWorkloadRunState::Running, + RuntimeWorkloadRunState::NotStarted}) { + snapshot.run_state = state; + const auto report = RuntimeReport(snapshot); + EXPECT_EQ(report["component_summaries"][0]["known_operation_count"], "24576"); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], false); + } + snapshot.run_state = RuntimeWorkloadRunState::Succeeded; + snapshot.stats.collection_complete = false; + snapshot.stats.record_limit_drops = 3; + snapshot.stats.payload_limit_drops = 4; + snapshot.stats.metadata_failure_drops = 5; + snapshot.stats.allocation_failure_drops = 6; + snapshot.stats.run_reuse_rejections = 7; + const auto report = RuntimeReport(snapshot); + EXPECT_EQ(report["collection"]["record_limit_drops"], 3); + EXPECT_EQ(report["collection"]["payload_limit_drops"], 4); + EXPECT_EQ(report["collection"]["metadata_failure_drops"], 5); + EXPECT_EQ(report["collection"]["allocation_failure_drops"], 6); + EXPECT_EQ(report["collection"]["run_reuse_rejections"], 7); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], false); +} + +} // namespace +} // namespace test +} // namespace onnxruntime +#endif diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc new file mode 100644 index 0000000000000..9bdb2ba7cf8fd --- /dev/null +++ b/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc @@ -0,0 +1,347 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "core/framework/matmul_nbits_compute_estimate.h" + +#if !defined(ORT_MINIMAL_BUILD) +#include "core/graph/constants.h" +#include "core/graph/model.h" +#include "test/test_environment.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" +#endif + +namespace onnxruntime { +namespace test { +namespace { + +using contrib::ComputeEstimateUnavailableReason; +using contrib::EstimateMatMulNBitsBias; +using contrib::EstimateMatMulNBitsCompute; +using contrib::MatMulNBitsBiasEstimate; +using contrib::MatMulNBitsBiasEstimateResult; +using contrib::MatMulNBitsComputeEstimate; +using contrib::MatMulNBitsComputeEstimateResult; + +void ExpectBiasCount(const MatMulNBitsBiasEstimateResult& result, uint64_t expected) { + const auto* estimate = std::get_if(&result); + ASSERT_NE(estimate, nullptr); + EXPECT_EQ(estimate->floating_point_ops, expected); +} + +void ExpectBiasUnavailable(const MatMulNBitsBiasEstimateResult& result, ComputeEstimateUnavailableReason expected) { + const auto* reason = std::get_if(&result); + ASSERT_NE(reason, nullptr); + EXPECT_EQ(*reason, expected); +} + +TEST(MatMulNBitsBiasEstimateTest, MatrixBatchedAndRankOneCounts) { + ExpectBiasCount(EstimateMatMulNBitsBias(std::array{3, 64}, 64, 32, true), 96); + ExpectBiasCount(EstimateMatMulNBitsBias(std::array{2, 3, 64}, 64, 32, true), 192); + ExpectBiasCount(EstimateMatMulNBitsBias(std::array{64}, 64, 32, true), 32); +} + +TEST(MatMulNBitsBiasEstimateTest, AbsentBiasDoesNotRequireGeometry) { + ExpectBiasCount(EstimateMatMulNBitsBias({}, 64, 32, false), 0); + ExpectBiasCount(EstimateMatMulNBitsBias(std::array{-1, 64}, 64, 32, false), 0); +} + +TEST(MatMulNBitsBiasEstimateTest, UnknownPresenceIsNotZero) { + ExpectBiasUnavailable(EstimateMatMulNBitsBias(std::array{0, 64}, 64, 32, std::nullopt), + ComputeEstimateUnavailableReason::MissingSemanticValue); +} + +TEST(MatMulNBitsBiasEstimateTest, UnknownGeometryIsNotZero) { + ExpectBiasUnavailable(EstimateMatMulNBitsBias({}, 64, 32, true), + ComputeEstimateUnavailableReason::InvalidMetadata); + for (const auto& shape : {std::array{2, -1, 64}, std::array{0, -1, 64}, + std::array{2, 3, -1}}) { + ExpectBiasUnavailable(EstimateMatMulNBitsBias(shape, 64, 32, true), + ComputeEstimateUnavailableReason::UnknownShape); + } +} + +TEST(MatMulNBitsBiasEstimateTest, InvalidGeometry) { + ExpectBiasUnavailable(EstimateMatMulNBitsBias(std::array{-2, 64}, 64, 32, true), + ComputeEstimateUnavailableReason::InvalidMetadata); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(std::array{3, 63}, 64, 32, true), + ComputeEstimateUnavailableReason::InvalidMetadata); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(std::array{3, 64}, -1, 32, true), + ComputeEstimateUnavailableReason::InvalidMetadata); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(std::array{3, 64}, 64, -1, true), + ComputeEstimateUnavailableReason::InvalidMetadata); +} + +TEST(MatMulNBitsBiasEstimateTest, EmptyOutputAndZeroReduction) { + ExpectBiasCount(EstimateMatMulNBitsBias(std::array{0, 64}, 64, 32, true), 0); + ExpectBiasCount(EstimateMatMulNBitsBias(std::array{3, 64}, 64, 0, true), 0); + ExpectBiasCount(EstimateMatMulNBitsBias(std::array{3, 0}, 0, 32, true), 96); +} + +TEST(MatMulNBitsBiasEstimateTest, ZeroPrecedesOverflow) { + constexpr auto max = std::numeric_limits::max(); + ExpectBiasCount(EstimateMatMulNBitsBias(std::array{max, max, 0, 64}, 64, max, true), 0); +} + +TEST(MatMulNBitsBiasEstimateTest, OverflowAndIndependentCoreCount) { + constexpr auto max = std::numeric_limits::max(); + ExpectBiasCount(EstimateMatMulNBitsBias(std::array{max, 64}, 64, 2, true), + std::numeric_limits::max() - 1); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(std::array{max, 64}, 64, 3, true), + ComputeEstimateUnavailableReason::Overflow); + ExpectBiasCount(EstimateMatMulNBitsBias(std::array{max}, max, 32, true), 32); +} + +void ExpectCount(const MatMulNBitsComputeEstimateResult& result, uint64_t expected) { + const auto* estimate = std::get_if(&result); + ASSERT_NE(estimate, nullptr); + EXPECT_EQ(estimate->dense_equivalent_ops, expected); + EXPECT_EQ(estimate->kConventionVersion, 1u); +} + +void ExpectUnavailable(const MatMulNBitsComputeEstimateResult& result, + ComputeEstimateUnavailableReason expected) { + const auto* reason = std::get_if(&result); + ASSERT_NE(reason, nullptr); + EXPECT_EQ(*reason, expected); +} + +TEST(MatMulNBitsComputeEstimateTest, MatrixAndBatchedCounts) { + ExpectCount(EstimateMatMulNBitsCompute(std::array{3, 64}, 64, 32), 12288); + ExpectCount(EstimateMatMulNBitsCompute(std::array{2, 3, 64}, 64, 32), 24576); + ExpectCount(EstimateMatMulNBitsCompute(std::array{2, 3, 4, 64}, 64, 32), 98304); +} + +TEST(MatMulNBitsComputeEstimateTest, RankOneInput) { + ExpectCount(EstimateMatMulNBitsCompute(std::array{64}, 64, 32), 4096); +} + +TEST(MatMulNBitsComputeEstimateTest, EmptyWorkloads) { + ExpectCount(EstimateMatMulNBitsCompute(std::array{0, 64}, 64, 32), 0); + ExpectCount(EstimateMatMulNBitsCompute(std::array{3, 64}, 64, 0), 0); + ExpectCount(EstimateMatMulNBitsCompute(std::array{3, 0}, 0, 32), 0); +} + +TEST(MatMulNBitsComputeEstimateTest, ZeroPrecedesProductOverflow) { + constexpr auto max = std::numeric_limits::max(); + ExpectCount(EstimateMatMulNBitsCompute(std::array{max, max, 0, 64}, 64, max), 0); +} + +TEST(MatMulNBitsComputeEstimateTest, UnknownDimensionsAreNotZero) { + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{-1, 64}, 64, 32), + ComputeEstimateUnavailableReason::UnknownShape); + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{3, -1}, 64, 32), + ComputeEstimateUnavailableReason::UnknownShape); + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{0, -1, 64}, 64, 32), + ComputeEstimateUnavailableReason::UnknownShape); +} + +TEST(MatMulNBitsComputeEstimateTest, InvalidMetadata) { + ExpectUnavailable(EstimateMatMulNBitsCompute({}, 64, 32), ComputeEstimateUnavailableReason::InvalidMetadata); + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{3, 63}, 64, 32), + ComputeEstimateUnavailableReason::InvalidMetadata); + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{-2, 64}, 64, 32), + ComputeEstimateUnavailableReason::InvalidMetadata); + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{3, 64}, -1, 32), + ComputeEstimateUnavailableReason::InvalidMetadata); + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{3, 64}, 64, -1), + ComputeEstimateUnavailableReason::InvalidMetadata); + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{0, -2, 64}, 64, 32), + ComputeEstimateUnavailableReason::InvalidMetadata); +} + +TEST(MatMulNBitsComputeEstimateTest, CountsUseUnsigned64BitRange) { + constexpr auto max = std::numeric_limits::max(); + ExpectCount(EstimateMatMulNBitsCompute(std::array{max}, max, 1), + std::numeric_limits::max() - 1); +} + +TEST(MatMulNBitsComputeEstimateTest, Overflow) { + constexpr auto max = std::numeric_limits::max(); + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{2, max}, max, 1), + ComputeEstimateUnavailableReason::Overflow); + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{2}, 2, max), + ComputeEstimateUnavailableReason::Overflow); + ExpectUnavailable(EstimateMatMulNBitsCompute(std::array{max, max, 1}, 1, 1), + ComputeEstimateUnavailableReason::Overflow); +} + +#if !defined(ORT_MINIMAL_BUILD) +class MatMulNBitsComputeEstimateGraphTest : public ::testing::Test { + protected: + MatMulNBitsComputeEstimateGraphTest() + : model_("compute_estimate", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 17}, {kMSDomain, 1}}, {}, + DefaultLoggingManager().DefaultLogger()), + builder_(model_.MainGraph()) {} + + template + Node& AddMatMul(std::optional> shape = std::vector{2, 3, 64}, + int64_t bits = 4, bool bias = false, bool explicit_absent_bias = false) { + NodeArg* a = builder_.MakeInput(shape); + NodeArg* b = builder_.MakeInput(std::vector{32, 2, 32 * bits / 8}); + NodeArg* scales = builder_.MakeInput(std::vector{32, 2}); + std::vector inputs{a, b, scales}; + if (bias || explicit_absent_bias) { + auto& missing = builder_.graph_.GetOrCreateNodeArg("", nullptr); + inputs.push_back(&missing); + inputs.push_back(&missing); + inputs.push_back(bias ? builder_.MakeInput(std::vector{32}) : &missing); + } + auto* output = builder_.MakeOutput(std::nullopt); + auto& node = builder_.AddNode("MatMulNBits", inputs, {output}, kMSDomain); + node.AddAttribute("K", int64_t{64}); + node.AddAttribute("N", int64_t{32}); + node.AddAttribute("block_size", int64_t{32}); + node.AddAttribute("bits", bits); + builder_.SetGraphOutputs(); + return node; + } + + Model model_; + ModelTestBuilder builder_; +}; + +TEST_F(MatMulNBitsComputeEstimateGraphTest, ResolvedGraph) { + auto& node = AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectCount(EstimateMatMulNBitsCompute(node), 24576); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, BiasPresentAndCoreUnchanged) { + auto& node = AddMatMul(std::vector{2, 3, 64}, 4, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectBiasCount(EstimateMatMulNBitsBias(node), 192); + ExpectCount(EstimateMatMulNBitsCompute(node), 24576); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, BiasAbsentWithUnknownRank) { + auto& node = AddMatMul(std::nullopt); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectBiasCount(EstimateMatMulNBitsBias(node), 0); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, ExplicitlyOmittedBiasWithUnknownRank) { + auto& node = AddMatMul(std::nullopt, 4, false, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectBiasCount(EstimateMatMulNBitsBias(node), 0); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, BiasPresentWithUnknownRank) { + auto& node = AddMatMul(std::nullopt, 4, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(node), ComputeEstimateUnavailableReason::UnknownShape); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, BiasPresentWithSymbolicShape) { + auto& node = AddMatMul(std::vector{-1, 64}, 4, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(node), ComputeEstimateUnavailableReason::UnknownShape); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, BiasHalfPrecisionTwoBit) { + auto& node = AddMatMul(std::vector{2, 3, 64}, 2, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectBiasCount(EstimateMatMulNBitsBias(node), 192); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, RankOneGraph) { + auto& node = AddMatMul(std::vector{64}); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectCount(EstimateMatMulNBitsCompute(node), 4096); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, TwoBitWithBias) { + auto& node = AddMatMul(std::vector{2, 3, 64}, 2, true); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectCount(EstimateMatMulNBitsCompute(node), 24576); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, EightBitHalfPrecision) { + auto& node = AddMatMul(std::vector{2, 3, 64}, 8); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectCount(EstimateMatMulNBitsCompute(node), 24576); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, BFloat16) { + auto& node = AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectCount(EstimateMatMulNBitsCompute(node), 24576); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, UnknownRank) { + auto& node = AddMatMul(std::nullopt); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectUnavailable(EstimateMatMulNBitsCompute(node), ComputeEstimateUnavailableReason::UnknownShape); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, PartiallyKnownShape) { + auto& node = AddMatMul(std::vector{-1, 64}); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectUnavailable(EstimateMatMulNBitsCompute(node), ComputeEstimateUnavailableReason::UnknownShape); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, MissingAttribute) { + auto& node = AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + node.ClearAttribute("K"); + ExpectUnavailable(EstimateMatMulNBitsCompute(node), ComputeEstimateUnavailableReason::InvalidMetadata); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(node), ComputeEstimateUnavailableReason::InvalidMetadata); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, WrongAttributeType) { + auto& node = AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + node.AddAttribute("N", std::string("32")); + ExpectUnavailable(EstimateMatMulNBitsCompute(node), ComputeEstimateUnavailableReason::InvalidMetadata); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, UnresolvedVersion) { + auto& node = AddMatMul(); + ExpectUnavailable(EstimateMatMulNBitsCompute(node), ComputeEstimateUnavailableReason::UnsupportedVersion); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, UnsupportedVersion) { + auto& node = AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + node.SetSinceVersion(2); + ExpectUnavailable(EstimateMatMulNBitsCompute(node), ComputeEstimateUnavailableReason::UnsupportedVersion); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(node), ComputeEstimateUnavailableReason::UnsupportedVersion); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, EmptyGraphWorkload) { + auto& node = AddMatMul(std::vector{0, 64}); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectCount(EstimateMatMulNBitsCompute(node), 0); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, UnsupportedDomain) { + auto* input = builder_.MakeInput(std::vector{2, 64}); + auto* output = builder_.MakeOutput(std::nullopt); + auto& node = builder_.AddNode("MatMulNBits", {input}, {output}, "other.domain"); + ExpectUnavailable(EstimateMatMulNBitsCompute(node), ComputeEstimateUnavailableReason::UnsupportedOperator); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(node), ComputeEstimateUnavailableReason::UnsupportedOperator); +} + +TEST_F(MatMulNBitsComputeEstimateGraphTest, UnsupportedOperator) { + auto* input = builder_.MakeInput(std::vector{2, 3}); + auto* output = builder_.MakeOutput(std::nullopt); + auto& node = builder_.AddNode("Identity", {input}, {output}); + builder_.SetGraphOutputs(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + ExpectUnavailable(EstimateMatMulNBitsCompute(node), ComputeEstimateUnavailableReason::UnsupportedOperator); +} +#endif + +} // namespace +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/framework/runtime_workload_collector_test.cc b/onnxruntime/test/framework/runtime_workload_collector_test.cc new file mode 100644 index 0000000000000..82ce681fe35c1 --- /dev/null +++ b/onnxruntime/test/framework/runtime_workload_collector_test.cc @@ -0,0 +1,938 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "nlohmann/json.hpp" +#include "core/framework/compute_estimate_report.h" +#include "core/framework/runtime_workload_collector.h" +#include "core/graph/constants.h" +#include "test/util/include/asserts.h" +#if !defined(DISABLE_CONTRIB_OPS) +#include "core/framework/tensor.h" +#include "core/mlas/inc/mlas_qnbit.h" +#include "core/mlas/inc/mlas_q4.h" +#include "core/session/inference_session.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/test_environment.h" +#endif + +namespace onnxruntime { +namespace test { +namespace { + +MatMulNBitsRuntimeMetadata Metadata() { + MatMulNBitsRuntimeMetadata metadata; + metadata.input_shape = {2, 3, 64}; + metadata.output_shape = {2, 3, 32}; + metadata.m = 3; + metadata.n = 32; + metadata.k = 64; + metadata.batch_count = 2; + metadata.bias_present = true; + metadata.input_element_type = ONNX_NAMESPACE::TensorProto_DataType_FLOAT; + metadata.bits = 4; + metadata.block_size = 32; + metadata.route = MatMulNBitsRuntimeRoute::PackedQNBit; + metadata.compute_mode = MatMulNBitsRuntimeComputeMode::SQNBitInt8; + return metadata; +} + +size_t PayloadBytes(const MatMulNBitsRuntimeMetadata& metadata) { + return sizeof(RuntimeWorkloadRecord) + std::string("main").size() + + std::string(kCpuExecutionProvider).size() + 2 + + (metadata.input_shape.size() + metadata.output_shape.size()) * sizeof(int64_t); +} + +#if !defined(DISABLE_CONTRIB_OPS) +nlohmann::json Report(const RuntimeWorkloadSnapshot& snapshot) { + std::string text; + ORT_THROW_IF_ERROR(BuildComputeEstimateReport(snapshot, text)); + return nlohmann::json::parse(text); +} +#endif + +TEST(RuntimeWorkloadCollectorTest, EmptySnapshotHasRunIdentity) { + RuntimeWorkloadCollector collector(12, 34, {10, 8192}); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + EXPECT_EQ(snapshot.session_id, 12u); + EXPECT_EQ(snapshot.run_id, 34u); + EXPECT_TRUE(snapshot.records.empty()); + EXPECT_TRUE(snapshot.stats.collection_complete); + EXPECT_EQ(snapshot.stats.retained_payload_bytes, 0u); + EXPECT_EQ(snapshot.limits.max_records, 10u); + EXPECT_EQ(snapshot.limits.max_payload_bytes, 8192u); +} + +TEST(RuntimeWorkloadCollectorTest, OwnsSubmittedMetadataAndSnapshotCopies) { + RuntimeWorkloadCollector collector(12, 34, {10, 8192}); + auto metadata = Metadata(); + std::string graph = "main/2/then_branch"; + std::string provider = kCpuExecutionProvider; + ASSERT_STATUS_OK(collector.RecordMatMulNBits(graph, 7, provider, metadata)); + graph.clear(); + provider.clear(); + metadata.input_shape[0] = 99; + metadata.output_shape.clear(); + metadata.route.reset(); + metadata.compute_mode.reset(); + + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + ASSERT_EQ(snapshot.records.size(), 1u); + const auto& record = snapshot.records[0]; + EXPECT_EQ(record.invocation_id, 0u); + EXPECT_EQ(record.graph_path, "main/2/then_branch"); + EXPECT_EQ(record.node_index, 7u); + EXPECT_EQ(record.provider, kCpuExecutionProvider); + EXPECT_EQ(record.matmul_nbits.input_shape, (TensorShapeVector{2, 3, 64})); + EXPECT_EQ(record.matmul_nbits.output_shape, (TensorShapeVector{2, 3, 32})); + EXPECT_EQ(record.matmul_nbits.m, 3); + EXPECT_EQ(record.matmul_nbits.n, 32); + EXPECT_EQ(record.matmul_nbits.k, 64); + EXPECT_EQ(record.matmul_nbits.batch_count, 2u); + EXPECT_TRUE(record.matmul_nbits.bias_present); + EXPECT_EQ(record.matmul_nbits.input_element_type, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + EXPECT_EQ(record.matmul_nbits.bits, 4u); + EXPECT_EQ(record.matmul_nbits.block_size, 32u); + EXPECT_EQ(record.matmul_nbits.route, MatMulNBitsRuntimeRoute::PackedQNBit); + EXPECT_EQ(record.matmul_nbits.compute_mode, MatMulNBitsRuntimeComputeMode::SQNBitInt8); + + snapshot.records[0].matmul_nbits.input_shape.clear(); + snapshot.records.clear(); + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + ASSERT_EQ(snapshot.records.size(), 1u); + EXPECT_EQ(snapshot.records[0].matmul_nbits.input_shape, (TensorShapeVector{2, 3, 64})); +} + +TEST(RuntimeWorkloadCollectorTest, RepeatedNodeKeepsDistinctInvocationsAndChangingShapes) { + RuntimeWorkloadCollector collector(1, 2, {10, 8192}); + auto metadata = Metadata(); + ASSERT_STATUS_OK(collector.RecordMatMulNBits("main", 0, kCpuExecutionProvider, metadata)); + metadata.input_shape = {1, 5, 64}; + metadata.output_shape = {1, 5, 32}; + metadata.m = 5; + metadata.batch_count = 1; + ASSERT_STATUS_OK(collector.RecordMatMulNBits("main", 0, kCpuExecutionProvider, metadata)); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + ASSERT_EQ(snapshot.records.size(), 2u); + EXPECT_EQ(snapshot.records[0].invocation_id, 0u); + EXPECT_EQ(snapshot.records[1].invocation_id, 1u); + EXPECT_EQ(snapshot.records[0].matmul_nbits.input_shape, (TensorShapeVector{2, 3, 64})); + EXPECT_EQ(snapshot.records[1].matmul_nbits.input_shape, (TensorShapeVector{1, 5, 64})); + EXPECT_EQ(snapshot.records[1].matmul_nbits.m, 5); +} + +TEST(RuntimeWorkloadCollectorTest, MissingRouteMetadataRemainsMissing) { + RuntimeWorkloadCollector collector(1, 2, {10, 8192}); + auto metadata = Metadata(); + metadata.route.reset(); + metadata.compute_mode.reset(); + ASSERT_STATUS_OK(collector.RecordMatMulNBits("main", 0, kCpuExecutionProvider, metadata)); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + EXPECT_FALSE(snapshot.records[0].matmul_nbits.route.has_value()); + EXPECT_FALSE(snapshot.records[0].matmul_nbits.compute_mode.has_value()); + EXPECT_TRUE(snapshot.stats.collection_complete); +} + +TEST(RuntimeWorkloadCollectorTest, EmptyOutputRetainsExplicitNonExecutionMode) { + RuntimeWorkloadCollector collector(1, 2, {10, 8192}); + auto metadata = Metadata(); + metadata.input_shape = {2, 0, 64}; + metadata.output_shape = {2, 0, 32}; + metadata.m = 0; + metadata.route = MatMulNBitsRuntimeRoute::EmptyOutput; + metadata.compute_mode = MatMulNBitsRuntimeComputeMode::NotExecuted; + ASSERT_STATUS_OK(collector.RecordMatMulNBits("main", 0, kCpuExecutionProvider, metadata)); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + EXPECT_EQ(snapshot.records[0].matmul_nbits.route, MatMulNBitsRuntimeRoute::EmptyOutput); + EXPECT_EQ(snapshot.records[0].matmul_nbits.compute_mode, MatMulNBitsRuntimeComputeMode::NotExecuted); +} + +TEST(RuntimeWorkloadCollectorTest, RecordLimitReturnsErrorAndPreservesAcceptedData) { + RuntimeWorkloadCollector collector(1, 2, {1, 8192}); + const auto metadata = Metadata(); + ASSERT_STATUS_OK(collector.RecordMatMulNBits("main", 0, kCpuExecutionProvider, metadata)); + EXPECT_FALSE(collector.RecordMatMulNBits("main", 1, kCpuExecutionProvider, metadata).IsOK()); + EXPECT_FALSE(collector.RecordMatMulNBits("main", 2, kCpuExecutionProvider, metadata).IsOK()); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + ASSERT_EQ(snapshot.records.size(), 1u); + EXPECT_EQ(snapshot.records[0].node_index, 0u); + EXPECT_EQ(snapshot.stats.record_limit_drops, 2u); + EXPECT_EQ(snapshot.stats.payload_limit_drops, 0u); + EXPECT_FALSE(snapshot.stats.collection_complete); +} + +TEST(RuntimeWorkloadCollectorTest, PayloadBoundaryAndRejectedRecordDoNotConsumeBudget) { + const auto metadata = Metadata(); + const auto bytes = PayloadBytes(metadata); + RuntimeWorkloadCollector collector(1, 2, {10, bytes}); + EXPECT_FALSE(collector.RecordMatMulNBits("main/0/longer", 0, kCpuExecutionProvider, metadata).IsOK()); + ASSERT_STATUS_OK(collector.RecordMatMulNBits("main", 0, kCpuExecutionProvider, metadata)); + EXPECT_FALSE(collector.RecordMatMulNBits("main", 1, kCpuExecutionProvider, metadata).IsOK()); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + ASSERT_EQ(snapshot.records.size(), 1u); + EXPECT_EQ(snapshot.records[0].invocation_id, 0u); + EXPECT_EQ(snapshot.stats.retained_payload_bytes, bytes); + EXPECT_EQ(snapshot.stats.payload_limit_drops, 2u); + EXPECT_FALSE(snapshot.stats.collection_complete); + + RuntimeWorkloadCollector below_boundary(1, 3, {10, bytes - 1}); + EXPECT_FALSE(below_boundary.RecordMatMulNBits("main", 0, kCpuExecutionProvider, metadata).IsOK()); + ASSERT_STATUS_OK(below_boundary.GetSnapshot(snapshot)); + EXPECT_TRUE(snapshot.records.empty()); + EXPECT_EQ(snapshot.stats.retained_payload_bytes, 0u); +} + +TEST(RuntimeWorkloadCollectorTest, ZeroLimitsExplicitlyRejectSubmissions) { + RuntimeWorkloadCollector no_records(1, 2, {0, 8192}); + RuntimeWorkloadCollector no_bytes(1, 3, {10, 0}); + EXPECT_FALSE(no_records.RecordMatMulNBits("main", 0, kCpuExecutionProvider, Metadata()).IsOK()); + EXPECT_FALSE(no_bytes.RecordMatMulNBits("main", 0, kCpuExecutionProvider, Metadata()).IsOK()); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(no_records.GetSnapshot(snapshot)); + EXPECT_EQ(snapshot.stats.record_limit_drops, 1u); + EXPECT_FALSE(snapshot.stats.collection_complete); + ASSERT_STATUS_OK(no_bytes.GetSnapshot(snapshot)); + EXPECT_EQ(snapshot.stats.payload_limit_drops, 1u); + EXPECT_FALSE(snapshot.stats.collection_complete); +} + +TEST(RuntimeWorkloadCollectorTest, RunLifecycleAndCollectionFailuresAreIndependent) { + RuntimeWorkloadCollector collector(0, 0, {10, 8192}); + EXPECT_TRUE(collector.BeginRun(23)); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + EXPECT_EQ(snapshot.run_state, RuntimeWorkloadRunState::Running); + EXPECT_EQ(snapshot.session_id, 23u); + EXPECT_NE(snapshot.run_id, 0u); + EXPECT_FALSE(collector.BeginRun(24)); + collector.RecordCollectionFailure(true); + collector.RecordCollectionFailure(false); + collector.EndRun(false); + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + EXPECT_EQ(snapshot.run_state, RuntimeWorkloadRunState::Failed); + EXPECT_EQ(snapshot.session_id, 23u); + EXPECT_EQ(snapshot.stats.run_reuse_rejections, 1u); + EXPECT_EQ(snapshot.stats.allocation_failure_drops, 1u); + EXPECT_EQ(snapshot.stats.metadata_failure_drops, 1u); + EXPECT_FALSE(snapshot.stats.collection_complete); +} + +TEST(RuntimeWorkloadCollectorTest, SeparateRunsDoNotShareRecordsOrIds) { + RuntimeWorkloadCollector first(1, 2, {10, 8192}); + RuntimeWorkloadCollector second(1, 3, {10, 8192}); + ASSERT_STATUS_OK(first.RecordMatMulNBits("main", 7, kCpuExecutionProvider, Metadata())); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(second.GetSnapshot(snapshot)); + EXPECT_EQ(snapshot.run_id, 3u); + EXPECT_TRUE(snapshot.records.empty()); + ASSERT_STATUS_OK(second.RecordMatMulNBits("main", 8, kCpuExecutionProvider, Metadata())); + ASSERT_STATUS_OK(second.GetSnapshot(snapshot)); + EXPECT_EQ(snapshot.records[0].invocation_id, 0u); + EXPECT_EQ(snapshot.records[0].node_index, 8u); + ASSERT_STATUS_OK(first.GetSnapshot(snapshot)); + EXPECT_EQ(snapshot.run_id, 2u); + EXPECT_EQ(snapshot.records[0].node_index, 7u); +} + +TEST(RuntimeWorkloadCollectorTest, ConcurrentSubmissionsAndSnapshotsRespectLimits) { + constexpr size_t limit = 100; + constexpr size_t submissions_per_thread = 50; + RuntimeWorkloadCollector collector(1, 2, {limit, std::numeric_limits::max()}); + std::atomic accepted{0}; + std::atomic workers_done{false}; + std::atomic snapshot_failed{false}; + std::thread reader([&]() { + while (!workers_done.load()) { + RuntimeWorkloadSnapshot snapshot; + if (!collector.GetSnapshot(snapshot).IsOK() || snapshot.records.size() > limit) { + snapshot_failed = true; + } + } + }); + std::array workers; + for (size_t i = 0; i < workers.size(); ++i) { + workers[i] = std::thread([&, i]() { + const auto metadata = Metadata(); + for (size_t j = 0; j < submissions_per_thread; ++j) { + if (collector.RecordMatMulNBits("main", i, kCpuExecutionProvider, metadata).IsOK()) { + ++accepted; + } + } + }); + } + for (auto& worker : workers) { + worker.join(); + } + workers_done = true; + reader.join(); + EXPECT_FALSE(snapshot_failed); + EXPECT_EQ(accepted, limit); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + ASSERT_EQ(snapshot.records.size(), limit); + EXPECT_EQ(snapshot.stats.record_limit_drops, workers.size() * submissions_per_thread - limit); + EXPECT_EQ(snapshot.stats.retained_payload_bytes, limit * PayloadBytes(Metadata())); + EXPECT_FALSE(snapshot.stats.collection_complete); + for (size_t i = 0; i < limit; ++i) { + EXPECT_EQ(snapshot.records[i].invocation_id, i); + EXPECT_LT(snapshot.records[i].node_index, workers.size()); + } +} + +#if !defined(DISABLE_CONTRIB_OPS) + +using ONNX_NAMESPACE::GraphProto; +using ONNX_NAMESPACE::ModelProto; +using ONNX_NAMESPACE::TensorProto; +using ONNX_NAMESPACE::ValueInfoProto; + +void SetValueType(ValueInfoProto& value, const char* name, int type, int rank) { + value.set_name(name); + auto* tensor_type = value.mutable_type()->mutable_tensor_type(); + tensor_type->set_elem_type(type); + for (int i = 0; i < rank; ++i) { + tensor_type->mutable_shape()->add_dim()->set_dim_param(std::string(name) + std::to_string(i)); + } + if (rank == 0) { + tensor_type->mutable_shape(); + } +} + +void AddIntegerAttribute(ONNX_NAMESPACE::NodeProto& node, const char* name, int64_t value) { + auto* attribute = node.add_attribute(); + attribute->set_name(name); + attribute->set_type(ONNX_NAMESPACE::AttributeProto::INT); + attribute->set_i(value); +} + +ModelProto WorkloadModel(bool dynamic_weights = false, bool bias = true, int bits = 4, + bool half = false, int rank = 3, int64_t n = 32, int64_t k = 64) { + ModelProto model; + model.set_ir_version(8); + model.add_opset_import()->set_version(17); + auto* ms_opset = model.add_opset_import(); + ms_opset->set_domain(kMSDomain); + ms_opset->set_version(1); + auto& graph = *model.mutable_graph(); + graph.set_name("runtime_workload"); + const int type = half ? TensorProto::FLOAT16 : TensorProto::FLOAT; + SetValueType(*graph.add_input(), "A", type, rank); + SetValueType(*graph.add_output(), "Y", type, rank); + if (dynamic_weights) { + SetValueType(*graph.add_input(), "B", TensorProto::UINT8, 3); + } else { + auto* b = graph.add_initializer(); + b->set_name("B"); + b->set_data_type(TensorProto::UINT8); + b->add_dims(n); + b->add_dims(k / 32); + b->add_dims(32 * bits / 8); + b->set_raw_data(std::string(static_cast(n * k * bits / 8), bits == 2 ? '\xff' : '\x99')); + } + auto add_float_initializer = [&](const char* name, int64_t size, float value) { + auto* tensor = graph.add_initializer(); + tensor->set_name(name); + tensor->set_data_type(type); + tensor->add_dims(size); + for (int64_t i = 0; i < size; ++i) { + if (half) { + tensor->add_int32_data(MLFloat16(value).val); + } else { + tensor->add_float_data(value); + } + } + }; + add_float_initializer("scales", n * k / 32, 1.f); + if (bias) { + add_float_initializer("bias", n, 2.f); + } + auto* node = graph.add_node(); + node->set_op_type("MatMulNBits"); + node->set_domain(kMSDomain); + for (const char* input : {"A", "B", "scales"}) { + node->add_input(input); + } + if (bias) { + node->add_input(""); + node->add_input(""); + node->add_input("bias"); + } + node->add_output("Y"); + AddIntegerAttribute(*node, "K", k); + AddIntegerAttribute(*node, "N", n); + AddIntegerAttribute(*node, "bits", bits); + AddIntegerAttribute(*node, "block_size", 32); + AddIntegerAttribute(*node, "accuracy_level", 4); + return model; +} + +ModelProto LoopWorkloadModel() { + auto model = WorkloadModel(); + auto body = model.graph(); + body.set_name("workload_body"); + body.clear_input(); + SetValueType(*body.add_input(), "iteration", TensorProto::INT64, 0); + SetValueType(*body.add_input(), "condition", TensorProto::BOOL, 0); + body.clear_output(); + SetValueType(*body.add_output(), "condition_out", TensorProto::BOOL, 0); + SetValueType(*body.add_output(), "Y", TensorProto::FLOAT, 3); + auto* identity = body.add_node(); + identity->set_op_type("Identity"); + identity->add_input("condition"); + identity->add_output("condition_out"); + + auto& graph = *model.mutable_graph(); + graph.clear_initializer(); + graph.clear_node(); + graph.clear_output(); + SetValueType(*graph.add_output(), "Y", TensorProto::FLOAT, 4); + auto* iterations = graph.add_initializer(); + iterations->set_name("iterations"); + iterations->set_data_type(TensorProto::INT64); + iterations->add_int64_data(3); + auto* condition = graph.add_initializer(); + condition->set_name("condition"); + condition->set_data_type(TensorProto::BOOL); + condition->add_int32_data(1); + auto* loop = graph.add_node(); + loop->set_op_type("Loop"); + loop->add_input("iterations"); + loop->add_input("condition"); + loop->add_output("Y"); + auto* attribute = loop->add_attribute(); + attribute->set_name("body"); + attribute->set_type(ONNX_NAMESPACE::AttributeProto::GRAPH); + *attribute->mutable_g() = std::move(body); + return model; +} + +template +OrtValue FilledTensor(const std::vector& shape, T value); + +TEST(RuntimeWorkloadExecutionTest, TransformerEncoderAndDecoderPropagateCollector) { + ModelProto model; + std::ifstream stream("testdata/dummy_t5.onnx", std::ios::binary); + ASSERT_TRUE(stream.is_open()); + ASSERT_TRUE(model.ParseFromIstream(&stream)); + size_t instrumented = 0; + for (auto& node : *model.mutable_graph()->mutable_node()) { + for (auto& attribute : *node.mutable_attribute()) { + if (attribute.name() != "encoder" && attribute.name() != "decoder") { + continue; + } + auto& graph = *attribute.mutable_g(); + ASSERT_GT(graph.output_size(), 0); + const auto output_name = graph.output(0).name(); + const std::string prefix = "workload_probe_"; + for (auto& existing : *graph.mutable_node()) { + for (auto& output : *existing.mutable_output()) { + if (output == output_name) output = prefix + "original"; + } + for (auto& input : *existing.mutable_input()) { + if (input == output_name) input = prefix + "original"; + } + } + const auto probe = WorkloadModel(false, false); + for (auto initializer : probe.graph().initializer()) { + initializer.set_name(prefix + initializer.name()); + *graph.add_initializer() = std::move(initializer); + } + auto* a = graph.add_initializer(); + a->set_name(prefix + "A"); + a->set_data_type(TensorProto::FLOAT); + for (int64_t dim : {1, 1, 64}) a->add_dims(dim); + a->set_raw_data(std::string(64 * sizeof(float), '\0')); + for (auto probe_node : probe.graph().node()) { + for (auto& input : *probe_node.mutable_input()) { + if (!input.empty()) input = prefix + input; + } + for (auto& output : *probe_node.mutable_output()) output = prefix + output; + *graph.add_node() = std::move(probe_node); + } + auto* reduce = graph.add_node(); + reduce->set_op_type("ReduceSum"); + reduce->add_input(prefix + "Y"); + reduce->add_output(prefix + "zero"); + AddIntegerAttribute(*reduce, "keepdims", 0); + auto* add = graph.add_node(); + add->set_op_type("Add"); + add->add_input(prefix + "original"); + add->add_input(prefix + "zero"); + add->add_output(output_name); + ++instrumented; + } + } + ASSERT_EQ(instrumented, 2u); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + options.intra_op_param.thread_pool_size = 1; + InferenceSession session(options, GetEnvironment()); + const auto bytes = model.SerializeAsString(); + ASSERT_STATUS_OK(session.Load(bytes.data(), static_cast(bytes.size()))); + ASSERT_STATUS_OK(session.Initialize()); + OrtValue input = FilledTensor({1, 5}, 0); + const std::array tokens{14, 6, 13, 9, 7}; + std::copy(tokens.begin(), tokens.end(), input.GetMutable()->MutableData()); + RuntimeWorkloadCollector collector(0, 0, {100, 1024 * 1024}); + RunOptions run_options; + run_options.runtime_workload_collector = &collector; + const std::vector input_names{"encoder_input_ids"}, output_names{"sequences"}; + const std::vector feeds{input}; + std::vector outputs; + ASSERT_STATUS_OK(session.Run(run_options, input_names, feeds, output_names, &outputs)); + RuntimeWorkloadSnapshot snapshot; + ASSERT_STATUS_OK(collector.GetSnapshot(snapshot)); + size_t encoder_records = 0, decoder_records = 0; + for (const auto& record : snapshot.records) { + if (record.graph_path.find("/encoder") != std::string::npos) ++encoder_records; + if (record.graph_path.find("/decoder") != std::string::npos) ++decoder_records; + } + EXPECT_EQ(encoder_records, 1u); + EXPECT_GT(decoder_records, 1u); + EXPECT_TRUE(snapshot.stats.collection_complete); + const auto report = Report(snapshot); + EXPECT_EQ(report["invocations"].size(), snapshot.records.size()); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], true); + const std::vector expected{ + 2, 16, 6, 14, 1, 15, 6, 14, 1, 15, 2, 3, 4, 15, 6, 14, 1, 15, 6, 14, + 2, 16, 6, 14, 1, 15, 6, 14, 1, 14}; + ASSERT_EQ(outputs.size(), 1u); + const auto& output = outputs[0].Get(); + ASSERT_EQ(output.Shape().Size(), static_cast(expected.size())); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), output.Data())); +} + +std::unique_ptr WorkloadSession(const ModelProto& model, bool force_fp32 = false, + bool lut = false, bool parallel = false) { + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + options.intra_op_param.thread_pool_size = 1; + options.inter_op_param.thread_pool_size = 2; + options.execution_mode = parallel ? ORT_PARALLEL : ORT_SEQUENTIAL; + if (force_fp32) { + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry(kOrtSessionOptionsMlasQNBitForceFp32, "1")); + } + if (lut) { + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry(kOrtSessionOptionsMlasLutGemm, "1")); + } + auto session = std::make_unique(options, GetEnvironment()); + const auto bytes = model.SerializeAsString(); + ORT_THROW_IF_ERROR(session->Load(bytes.data(), static_cast(bytes.size()))); + ORT_THROW_IF_ERROR(session->Initialize()); + return session; +} + +template +OrtValue FilledTensor(const std::vector& shape, T value) { + OrtValue result; + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape(shape), std::make_shared(), result); + auto* tensor = result.GetMutable(); + std::fill_n(tensor->MutableData(), static_cast(tensor->Shape().Size()), value); + return result; +} + +Status RunWorkload(InferenceSession& session, RuntimeWorkloadCollector* collector, + std::vector& outputs, const std::vector& shape = {2, 3, 64}, + bool dynamic_weights = false, bool half = false, bool terminate = false) { + RunOptions options; + options.runtime_workload_collector = collector; + options.terminate = terminate; + std::vector names{"A"}; + std::vector feeds; + feeds.push_back(half ? FilledTensor(shape, MLFloat16(1.f)) : FilledTensor(shape, 1.f)); + if (dynamic_weights) { + names.push_back("B"); + feeds.push_back(FilledTensor({32, 2, 16}, 0x99)); + } + const std::vector output_names{"Y"}; + return session.Run(options, names, feeds, output_names, &outputs); +} + +void ExpectOutput(const std::vector& outputs, float expected, bool half = false) { + ASSERT_EQ(outputs.size(), 1u); + const auto& tensor = outputs[0].Get(); + for (int64_t i = 0; i < tensor.Shape().Size(); ++i) { + EXPECT_FLOAT_EQ(half ? tensor.Data()[i].ToFloat() : tensor.Data()[i], expected); + } +} + +RuntimeWorkloadSnapshot Snapshot(const RuntimeWorkloadCollector& collector) { + RuntimeWorkloadSnapshot snapshot; + ORT_THROW_IF_ERROR(collector.GetSnapshot(snapshot)); + return snapshot; +} + +TEST(RuntimeWorkloadExecutionTest, DynamicRepeatedRunsAndDisabledCapturePreserveOutputs) { + auto session = WorkloadSession(WorkloadModel(true)); + std::vector outputs; + ASSERT_STATUS_OK(RunWorkload(*session, nullptr, outputs, {2, 3, 64}, true)); + ExpectOutput(outputs, 66.f); + uint64_t previous_run = 0; + for (int64_t rows : {3, 5}) { + RuntimeWorkloadCollector collector(0, 0, {10, 8192}); + outputs.clear(); + ASSERT_STATUS_OK(RunWorkload(*session, &collector, outputs, {2, rows, 64}, true)); + ExpectOutput(outputs, 66.f); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), 1u); + EXPECT_EQ(snapshot.session_id, session->GetCurrentSessionId()); + EXPECT_NE(snapshot.run_id, previous_run); + previous_run = snapshot.run_id; + EXPECT_EQ(snapshot.run_state, RuntimeWorkloadRunState::Succeeded); + EXPECT_TRUE(snapshot.stats.collection_complete); + const auto& record = snapshot.records[0]; + EXPECT_EQ(record.invocation_id, 0u); + EXPECT_EQ(record.graph_path, "main"); + EXPECT_EQ(record.node_index, 0u); + EXPECT_EQ(record.provider, kCpuExecutionProvider); + const auto& metadata = record.matmul_nbits; + EXPECT_EQ(metadata.input_shape, (TensorShapeVector{2, rows, 64})); + EXPECT_EQ(metadata.output_shape, (TensorShapeVector{2, rows, 32})); + // MatMulComputeHelper flattens the shared-B batches for the actual backend call. + EXPECT_EQ(metadata.m, rows * 2); + EXPECT_EQ(metadata.batch_count, 1u); + EXPECT_EQ(metadata.n, 32); + EXPECT_EQ(metadata.k, 64); + EXPECT_TRUE(metadata.bias_present); + EXPECT_EQ(metadata.bits, 4u); + EXPECT_EQ(metadata.block_size, 32u); + EXPECT_EQ(metadata.route, MatMulNBitsRuntimeRoute::Unpacked); + EXPECT_EQ(metadata.compute_mode, MatMulNBitsRuntimeComputeMode::Float32Gemm); + } +} + +TEST(RuntimeWorkloadExecutionTest, PackedBackendModeIsEffectiveAndHalfTypeIsIndependent) { + if (!MlasIsQNBitGemmAvailable(4, 32, SQNBIT_CompInt8)) { + GTEST_SKIP() << "SQNBit int8 backend is unavailable."; + } + for (bool half : {false, true}) { + auto session = WorkloadSession(WorkloadModel(false, false, 4, half)); + RuntimeWorkloadCollector collector(0, 0, {10, 8192}); + std::vector outputs; + ASSERT_STATUS_OK(RunWorkload(*session, &collector, outputs, {2, 3, 64}, false, half)); + ExpectOutput(outputs, 64.f, half); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), 1u); + const auto& metadata = snapshot.records[0].matmul_nbits; + EXPECT_EQ(metadata.route, MatMulNBitsRuntimeRoute::PackedQNBit); + EXPECT_EQ(metadata.compute_mode, MatMulNBitsRuntimeComputeMode::SQNBitInt8); + EXPECT_EQ(metadata.input_element_type, half ? TensorProto::FLOAT16 : TensorProto::FLOAT); + EXPECT_FALSE(metadata.bias_present); + } +} + +TEST(RuntimeWorkloadExecutionTest, ForceFp32RecordsActualPackedMode) { +#if defined(MLAS_TARGET_AMD64_IX86) + if (!MlasIsQNBitGemmAvailable(4, 32, SQNBIT_CompFp32)) { + GTEST_SKIP() << "SQNBit fp32 backend is unavailable."; + } + auto session = WorkloadSession(WorkloadModel(), true); + RuntimeWorkloadCollector collector(0, 0, {10, 8192}); + std::vector outputs; + ASSERT_STATUS_OK(RunWorkload(*session, &collector, outputs)); + ExpectOutput(outputs, 66.f); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), 1u); + EXPECT_EQ(snapshot.records[0].matmul_nbits.route, MatMulNBitsRuntimeRoute::PackedQNBit); + EXPECT_EQ(snapshot.records[0].matmul_nbits.compute_mode, MatMulNBitsRuntimeComputeMode::SQNBitFp32); +#else + GTEST_SKIP() << "Force-fp32 session option is x86-only."; +#endif +} + +TEST(RuntimeWorkloadExecutionTest, EmptyOutputIsExplicitAndNoBackendIsClaimed) { + auto session = WorkloadSession(WorkloadModel()); + RuntimeWorkloadCollector collector(0, 0, {10, 8192}); + std::vector outputs; + ASSERT_STATUS_OK(RunWorkload(*session, &collector, outputs, {2, 0, 64})); + ASSERT_EQ(outputs.size(), 1u); + EXPECT_EQ(outputs[0].Get().Shape(), TensorShape({2, 0, 32})); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), 1u); + EXPECT_EQ(snapshot.records[0].matmul_nbits.m, 0); + EXPECT_EQ(snapshot.records[0].matmul_nbits.route, MatMulNBitsRuntimeRoute::EmptyOutput); + EXPECT_EQ(snapshot.records[0].matmul_nbits.compute_mode, MatMulNBitsRuntimeComputeMode::NotExecuted); + EXPECT_EQ(snapshot.run_state, RuntimeWorkloadRunState::Succeeded); +} + +TEST(RuntimeWorkloadExecutionTest, LutBackendIsRecordedOnlyWhenItActuallyRuns) { + if (!MlasIsLutGemmAvailable(64, 128, 4, 32)) { + GTEST_SKIP() << "LUT backend is unavailable."; + } + auto session = WorkloadSession(WorkloadModel(false, true, 4, false, 3, 64, 128), false, true); + RuntimeWorkloadCollector collector(0, 0, {10, 8192}); + std::vector outputs; + ASSERT_STATUS_OK(RunWorkload(*session, &collector, outputs, {2, 3, 128})); + ExpectOutput(outputs, 130.f); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), 1u); + EXPECT_EQ(snapshot.records[0].matmul_nbits.route, MatMulNBitsRuntimeRoute::PackedLut); + EXPECT_EQ(snapshot.records[0].matmul_nbits.compute_mode, MatMulNBitsRuntimeComputeMode::Lut); +} + +TEST(RuntimeWorkloadExecutionTest, TwoBitVectorInvocationKeepsActualGeometry) { + auto session = WorkloadSession(WorkloadModel(false, false, 2, false, 1)); + RuntimeWorkloadCollector collector(0, 0, {10, 8192}); + std::vector outputs; + ASSERT_STATUS_OK(RunWorkload(*session, &collector, outputs, {64})); + ExpectOutput(outputs, 64.f); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), 1u); + const auto& metadata = snapshot.records[0].matmul_nbits; + EXPECT_EQ(metadata.input_shape, TensorShapeVector{64}); + EXPECT_EQ(metadata.output_shape, TensorShapeVector{32}); + EXPECT_EQ(metadata.bits, 2u); + EXPECT_EQ(metadata.m, 1); + EXPECT_EQ(metadata.n, 32); + EXPECT_EQ(metadata.k, 64); + EXPECT_EQ(metadata.batch_count, 1u); +} + +TEST(RuntimeWorkloadExecutionTest, KernelFailureAndTerminationDoNotCommitInvocations) { + auto session = WorkloadSession(WorkloadModel()); + for (bool terminate : {false, true}) { + RuntimeWorkloadCollector collector(0, 0, {10, 8192}); + std::vector outputs; + EXPECT_FALSE(RunWorkload(*session, &collector, outputs, {2, 3, terminate ? 64 : 63}, + false, false, terminate) + .IsOK()); + auto snapshot = Snapshot(collector); + EXPECT_TRUE(snapshot.records.empty()); + EXPECT_EQ(snapshot.run_state, RuntimeWorkloadRunState::Failed); + EXPECT_TRUE(snapshot.stats.collection_complete); + } +} + +TEST(RuntimeWorkloadExecutionTest, FailedRunRetainsEarlierSuccessfulInvocation) { + auto model = WorkloadModel(); + auto* reshape = model.mutable_graph()->add_node(); + reshape->set_op_type("Reshape"); + reshape->add_input("Y"); + reshape->add_input("bad_shape"); + reshape->add_output("bad_output"); + auto* shape = model.mutable_graph()->add_initializer(); + shape->set_name("bad_shape"); + shape->set_data_type(TensorProto::INT64); + shape->add_dims(1); + shape->add_int64_data(5); + model.mutable_graph()->clear_output(); + SetValueType(*model.mutable_graph()->add_output(), "bad_output", TensorProto::FLOAT, 1); + auto session = WorkloadSession(model); + RuntimeWorkloadCollector collector(0, 0, {10, 8192}); + RunOptions options; + options.runtime_workload_collector = &collector; + const std::vector names{"A"}, output_names{"bad_output"}; + const std::vector feeds{FilledTensor({2, 3, 64}, 1.f)}; + std::vector outputs; + EXPECT_FALSE(session->Run(options, names, feeds, output_names, &outputs).IsOK()); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), 1u); + EXPECT_EQ(snapshot.records[0].matmul_nbits.route.has_value(), true); + EXPECT_EQ(snapshot.run_state, RuntimeWorkloadRunState::Failed); + const auto report = Report(snapshot); + EXPECT_EQ(report["component_summaries"][0]["known_operation_count"], "24576"); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], false); +} + +TEST(RuntimeWorkloadExecutionTest, LoopInvocationsAndLimitsKeepSubgraphIdentity) { + auto session = WorkloadSession(LoopWorkloadModel()); + for (size_t max_records : {size_t{10}, size_t{1}, size_t{0}}) { + RuntimeWorkloadCollector collector(0, 0, {max_records, 8192}); + std::vector outputs; + ASSERT_STATUS_OK(RunWorkload(*session, &collector, outputs)); + ExpectOutput(outputs, 66.f); + EXPECT_EQ(outputs[0].Get().Shape(), TensorShape({3, 2, 3, 32})); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), std::min(max_records, size_t{3})); + EXPECT_EQ(snapshot.stats.record_limit_drops, 3 - snapshot.records.size()); + EXPECT_EQ(snapshot.stats.collection_complete, max_records >= 3); + EXPECT_EQ(snapshot.run_state, RuntimeWorkloadRunState::Succeeded); + const auto report = Report(snapshot); + EXPECT_EQ(report["component_summaries"][0]["known_operation_count"], + std::to_string(snapshot.records.size() * 24576)); + EXPECT_EQ(report["component_summaries"][1]["known_operation_count"], + std::to_string(snapshot.records.size() * 192)); + EXPECT_EQ(report["component_summaries"][0]["complete_for_captured_component_scope"], max_records >= 3); + for (size_t i = 0; i < snapshot.records.size(); ++i) { + EXPECT_EQ(snapshot.records[i].invocation_id, i); + EXPECT_EQ(snapshot.records[i].graph_path, "main/0/body"); + EXPECT_EQ(snapshot.records[i].node_index, 0u); + } + } +} + +TEST(RuntimeWorkloadExecutionTest, PayloadFailureAndCollectorReuseDoNotFailInference) { + auto session = WorkloadSession(WorkloadModel(true)); + RuntimeWorkloadCollector collector(0, 0, {10, 0}); + for (int i = 0; i < 2; ++i) { + std::vector outputs; + ASSERT_STATUS_OK(RunWorkload(*session, &collector, outputs, {2, 3, 64}, true)); + ExpectOutput(outputs, 66.f); + } + auto snapshot = Snapshot(collector); + EXPECT_TRUE(snapshot.records.empty()); + EXPECT_EQ(snapshot.stats.payload_limit_drops, 1u); + EXPECT_EQ(snapshot.stats.run_reuse_rejections, 1u); + EXPECT_FALSE(snapshot.stats.collection_complete); + EXPECT_EQ(snapshot.run_state, RuntimeWorkloadRunState::Succeeded); +} + +TEST(RuntimeWorkloadExecutionTest, DistinctNodesKeepTheirExecutableNodeIndices) { + auto model = WorkloadModel(); + auto& graph = *model.mutable_graph(); + auto second = graph.node(0); + second.set_output(0, "second"); + *graph.add_node() = std::move(second); + auto output = graph.output(0); + output.set_name("second"); + *graph.add_output() = std::move(output); + auto session = WorkloadSession(model, false, false, true); + RuntimeWorkloadCollector collector({10, 8192}); + std::vector outputs; + ASSERT_STATUS_OK(RunWorkload(*session, &collector, outputs)); + ExpectOutput(outputs, 66.f); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), 2u); + EXPECT_NE(snapshot.records[0].node_index, snapshot.records[1].node_index); + EXPECT_LT(snapshot.records[0].node_index, 2u); + EXPECT_LT(snapshot.records[1].node_index, 2u); + EXPECT_EQ(snapshot.records[0].invocation_id, 0u); + EXPECT_EQ(snapshot.records[1].invocation_id, 1u); + EXPECT_EQ(snapshot.records[0].graph_path, "main"); + EXPECT_EQ(snapshot.records[1].graph_path, "main"); +} + +TEST(RuntimeWorkloadExecutionTest, IfBranchesWithSameNodeIndexHaveDistinctGraphIdentities) { + auto model = WorkloadModel(); + auto branch = model.graph(); + branch.clear_input(); + auto& graph = *model.mutable_graph(); + graph.clear_initializer(); + graph.clear_node(); + SetValueType(*graph.add_input(), "condition", TensorProto::BOOL, 0); + auto* node = graph.add_node(); + node->set_op_type("If"); + node->add_input("condition"); + node->add_output("Y"); + for (const char* name : {"then_branch", "else_branch"}) { + auto* attribute = node->add_attribute(); + attribute->set_name(name); + attribute->set_type(ONNX_NAMESPACE::AttributeProto::GRAPH); + *attribute->mutable_g() = branch; + attribute->mutable_g()->set_name(name); + } + auto session = WorkloadSession(model); + for (bool condition : {true, false}) { + RuntimeWorkloadCollector collector({10, 8192}); + RunOptions options; + options.runtime_workload_collector = &collector; + const std::vector names{"A", "condition"}, output_names{"Y"}; + const std::vector feeds{FilledTensor({2, 3, 64}, 1.f), FilledTensor({}, condition)}; + std::vector outputs; + ASSERT_STATUS_OK(session->Run(options, names, feeds, output_names, &outputs)); + ExpectOutput(outputs, 66.f); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), 1u); + EXPECT_EQ(snapshot.records[0].node_index, 0u); + EXPECT_EQ(snapshot.records[0].graph_path, condition ? "main/0/then_branch" : "main/0/else_branch"); + const auto report = Report(snapshot); + EXPECT_EQ(report["invocations"].size(), 1u); + EXPECT_EQ(report["component_summaries"][0]["known_operation_count"], "24576"); + } +} + +TEST(RuntimeWorkloadExecutionTest, ScanInvocationsShareIdentityButNotInvocationId) { + auto model = WorkloadModel(); + auto body = model.graph(); + body.set_name("scan_body"); + auto& graph = *model.mutable_graph(); + graph.clear_initializer(); + graph.clear_node(); + graph.clear_input(); + graph.clear_output(); + SetValueType(*graph.add_input(), "A", TensorProto::FLOAT, 4); + SetValueType(*graph.add_output(), "Y", TensorProto::FLOAT, 4); + auto* node = graph.add_node(); + node->set_op_type("Scan"); + node->add_input("A"); + node->add_output("Y"); + AddIntegerAttribute(*node, "num_scan_inputs", 1); + auto* attribute = node->add_attribute(); + attribute->set_name("body"); + attribute->set_type(ONNX_NAMESPACE::AttributeProto::GRAPH); + *attribute->mutable_g() = std::move(body); + auto session = WorkloadSession(model); + RuntimeWorkloadCollector collector({10, 8192}); + std::vector outputs; + ASSERT_STATUS_OK(RunWorkload(*session, &collector, outputs, {3, 2, 3, 64})); + ExpectOutput(outputs, 66.f); + auto snapshot = Snapshot(collector); + ASSERT_EQ(snapshot.records.size(), 3u); + EXPECT_EQ(Report(snapshot)["component_summaries"][0]["known_operation_count"], "73728"); + for (size_t i = 0; i < snapshot.records.size(); ++i) { + EXPECT_EQ(snapshot.records[i].graph_path, "main/0/body"); + EXPECT_EQ(snapshot.records[i].node_index, 0u); + EXPECT_EQ(snapshot.records[i].invocation_id, i); + EXPECT_EQ(snapshot.records[i].matmul_nbits.input_shape, (TensorShapeVector{2, 3, 64})); + } +} + +TEST(RuntimeWorkloadExecutionTest, ConcurrentRunsAndSessionsAreIsolated) { + auto first = WorkloadSession(WorkloadModel(true), false, false, true); + auto second = WorkloadSession(WorkloadModel(true)); + std::array snapshots; + std::array workers; + std::atomic failures{0}; + for (size_t i = 0; i < workers.size(); ++i) { + workers[i] = std::thread([&, i]() { + auto& session = i < 6 ? *first : *second; + RuntimeWorkloadCollector collector(0, 0, {10, 8192}); + std::vector outputs; + if (!RunWorkload(session, &collector, outputs, {2, static_cast(i + 1), 64}, true).IsOK() || + !collector.GetSnapshot(snapshots[i]).IsOK()) { + ++failures; + } + }); + } + for (auto& worker : workers) { + worker.join(); + } + ASSERT_EQ(failures, 0u); + for (size_t i = 0; i < snapshots.size(); ++i) { + const auto& snapshot = snapshots[i]; + EXPECT_EQ(snapshot.session_id, (i < 6 ? first : second)->GetCurrentSessionId()); + EXPECT_EQ(snapshot.run_state, RuntimeWorkloadRunState::Succeeded); + ASSERT_EQ(snapshot.records.size(), 1u); + EXPECT_EQ(snapshot.records[0].invocation_id, 0u); + EXPECT_EQ(snapshot.records[0].matmul_nbits.input_shape[1], static_cast(i + 1)); + for (size_t j = 0; j < i; ++j) { + EXPECT_NE(snapshot.run_id, snapshots[j].run_id); + } + } +} + +#endif // !defined(DISABLE_CONTRIB_OPS) + +} // namespace +} // namespace test +} // namespace onnxruntime + +#endif