From dba44b338f2fc864c7895d7df86e72ebf1f5bf02 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Thu, 15 Oct 2026 14:57:33 -0700 Subject: [PATCH 01/14] Add FLOP estimation proposal document --- .../contrib_op_flop_estimation_rollout.md | 331 ++++++++++++++++++ 1 file changed, 331 insertions(+) create mode 100644 docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md 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..83809654f3936 --- /dev/null +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -0,0 +1,331 @@ +# Contrib Operator FLOP Estimation: Design and Rollout Plan + +**Status:** Proposal for team review; not a description of an implemented feature. + +## 1. Summary + +Introduce an opt-in diagnostic feature that estimates the mathematical work of a contrib operator invocation. +Start with a pure `com.microsoft::MatMulNBits` estimator and a coverage-aware static report. Then add concrete +runtime counts, Attention coverage, and measured effective throughput. 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. + +**Recommended first release:** counting conventions, `MatMulNBits` core dense-equivalent operations, and a versioned +developer report. Throughput and performance-policy changes are separate deliverables. + +## 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`. + +## 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 JSON is a proposed reporting format, not an +existing ORT API: + +```json +{ + "schema_version": 1, + "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", + "estimate_kind": "exact_under_convention", + "input_basis": "static_concrete", + "metric": "dense_equivalent_ops", + "scope": "matmul_core", + "operation_count": "24576", + "excluded_components": ["dequantization", "bias", "packing", "conversion", "padding"] + } + ], + "summary": { + "reported_node_count": 1, + "estimated_node_count": 1, + "unavailable_node_count": 0, + "known_dense_equivalent_ops": "24576", + "complete_for_declared_graph_and_component_scope": true, + "whole_operator_work_complete": false + } +} +``` + +Serialize counts as decimal strings to preserve integer precision in JSON consumers. An unavailable node omits +`operation_count` and includes a reason such as `unsupported_operator`, `unknown_shape`, `missing_semantic_value`, +`invalid_metadata`, `overflow`, or `runtime_execution_required`. + +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 counts | Opt-in invocation estimates using concrete shapes; event identity and bounded caching | Phase 2 | Execution/profiling owners | Dynamic-shape cases pass; no extra device synchronization; agreed overhead gate passes | +| 4. Attention coverage | Incremental packed Attention and GQA support; phase and active-length semantics | Phase 2; Phase 3 for runtime-only metadata | Attention owners | Prefill/decode/chunked, causal/windowed, and ragged reference cases pass | +| 5. Measured throughput | Matched CPU and eager-CUDA timing; units, scope, route provenance where available | Phase 3; Phase 4 only for Attention 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 | + +Phase 4 and Phase 5 can proceed independently after their prerequisites. Do not make pilot throughput wait for +complete Attention coverage. Phases 0-2 are the useful, independently shippable MVP. + +### Phase 2: static-report launch + +- Begin with an internal/developer reporting surface; the exact command or API is a Phase 0 decision. +- 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. + +### Phase 3: runtime refinement + +- 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. + +### Phase 4: Attention expansion + +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 + +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. + +## 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 | +| 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:** separately opt in to counts, then timing; record overhead and attribution gaps. +4. **Broader support:** add operators only 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 + +| 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 | + +## 10. Completion criteria + +The **MVP is complete** when a developer can request a persistent, versioned static report for an optimized graph, +see the correct `MatMulNBits` core count for supported concrete shapes, and understand every exclusion and missing +estimate without any change to inference behavior. + +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. From 6904645826ea7700eff867a90597aafd31edd142 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Mon, 19 Oct 2026 11:09:50 -0700 Subject: [PATCH 02/14] docs: outline contrib operator compute-estimation rollout --- .../contrib_op_flop_estimation_rollout.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md index 83809654f3936..1fa8d9254796d 100644 --- a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -25,6 +25,18 @@ The key distinction is: **Recommended first release:** counting conventions, `MatMulNBits` core dense-equivalent operations, and a versioned developer report. Throughput and performance-policy changes are separate deliverables. +### Initial implementation + +The internal [MatMulNBits compute estimator](../../onnxruntime/contrib_ops/cpu/quantization/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. + +This first increment does not add a graph report, runtime instrumentation, public API, or inference-policy changes. +The sources are included by the existing CPU contrib source glob, independently of CUDA, and the tests run in +`onnxruntime_provider_test` with `--gtest_filter="MatMulNBitsComputeEstimate*"` after rebuilding. + ## 2. Motivation and boundaries The feature should help users and ORT developers: From e4159acec0a7655eac51cec7df6f5db31b385aeb Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Mon, 19 Oct 2026 11:10:10 -0700 Subject: [PATCH 03/14] feat: add internal MatMulNBits dense-equivalent operation estimator --- .../matmul_nbits_compute_estimate.cc | 101 ++++++++ .../matmul_nbits_compute_estimate.h | 47 ++++ .../matmul_nbits_compute_estimate_test.cc | 235 ++++++++++++++++++ 3 files changed, 383 insertions(+) create mode 100644 onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc create mode 100644 onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h create mode 100644 onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc new file mode 100644 index 0000000000000..285e21403c4b2 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc @@ -0,0 +1,101 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h" + +#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 { + +MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute( + 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; + } + + 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}; +} + +#if !defined(ORT_MINIMAL_BUILD) +MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute(const Node& node) { + 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; + } + + 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); + } + } + + return EstimateMatMulNBitsCompute(dimensions, k_it->second.i(), n_it->second.i()); +} +#endif + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h new file mode 100644 index 0000000000000..46350aa0e131f --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h @@ -0,0 +1,47 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include + +namespace onnxruntime { +class Node; + +namespace contrib { + +enum class ComputeEstimateUnavailableReason { + UnsupportedOperator, + UnsupportedVersion, + UnknownShape, + InvalidMetadata, + Overflow, +}; + +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; + +// 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); + +#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); +#endif + +} // namespace contrib +} // namespace onnxruntime 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..abe7ac99846b8 --- /dev/null +++ b/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc @@ -0,0 +1,235 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "contrib_ops/cpu/quantization/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::EstimateMatMulNBitsCompute; +using contrib::MatMulNBitsComputeEstimate; +using contrib::MatMulNBitsComputeEstimateResult; + +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) { + 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) { + auto& missing = builder_.graph_.GetOrCreateNodeArg("", nullptr); + inputs.push_back(&missing); + inputs.push_back(&missing); + inputs.push_back(builder_.MakeInput(std::vector{32})); + } + 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, 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); +} + +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); +} + +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); +} + +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 From 2395592da173bdf9426eb06056b8e319f67311a3 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Mon, 19 Oct 2026 14:43:07 -0700 Subject: [PATCH 04/14] docs: expand contrib operator FLOP estimation rollout with ancillary components and counting categories --- .../contrib_op_flop_estimation_rollout.md | 87 +++++++++++++++++++ 1 file changed, 87 insertions(+) diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md index 1fa8d9254796d..bf483a0346d1b 100644 --- a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -25,6 +25,10 @@ The key distinction is: **Recommended first release:** counting conventions, `MatMulNBits` core dense-equivalent operations, and a versioned developer report. Throughput and performance-policy changes are separate deliverables. +Ancillary-work estimates are incremental follow-ups, not prerequisites for this first release. Define their +component contract during reporting design, add shape-based bias counts after static reporting, and add +implementation-dependent work only when its execution route and frequency can be established. + ### Initial implementation The internal [MatMulNBits compute estimator](../../onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h) @@ -116,6 +120,45 @@ dequantization, and an optional bias. The pilot counts **only core matrix multip 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. + ## 5. Proposed architecture and report ```text @@ -213,6 +256,19 @@ from particular contributors. Calendar estimates should follow contract review a Phase 4 and Phase 5 can proceed independently after their prerequisites. Do not make pilot throughput wait for complete Attention coverage. Phases 0-2 are the useful, independently shippable MVP. +The ancillary increments below extend these phases without renumbering or blocking the main rollout: + +| 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 | + +Agree on component semantics during Phase 2 design, but do not implement every component before shipping the MVP. +Increment 2a can proceed alongside runtime-count development. Increment 3a does not block Attention or core-only +throughput reporting; each new component is enabled only for routes with enough evidence. Full ancillary coverage +is not required before subsequent phases or the initial feature PR. + ### Phase 2: static-report launch - Begin with an internal/developer reporting surface; the exact command or API is a Phase 0 decision. @@ -221,6 +277,16 @@ complete Attention coverage. Phases 0-2 are the useful, independently shippable - 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 + +- 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 @@ -231,6 +297,22 @@ complete Attention coverage. Phases 0-2 are the useful, independently shippable - 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 + +For one supported route 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 First agree on useful unmasked work versus rectangular dense-equivalent work. For useful core work, define `P` as @@ -273,6 +355,8 @@ Label quantized results as dense-equivalent operations/s, not literal hardware T - 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 @@ -284,6 +368,8 @@ Use independent expected values rather than comparing two implementations of the | 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 | @@ -331,6 +417,7 @@ and rollback plan. | 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 | ## 10. Completion criteria From 58bbcc37be6792160aea6e058c4f20d11be168e9 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Mon, 19 Oct 2026 17:01:45 -0700 Subject: [PATCH 05/14] docs: prioritize MatMulNBits counting completion before expansion --- .../contrib_op_flop_estimation_rollout.md | 177 ++++++++++++++---- 1 file changed, 145 insertions(+), 32 deletions(-) diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md index bf483a0346d1b..0f873df75106c 100644 --- a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -1,13 +1,14 @@ # Contrib Operator FLOP Estimation: Design and Rollout Plan -**Status:** Proposal for team review; not a description of an implemented feature. +**Status:** Rollout proposal with an initial internal implementation; later phases remain proposed. ## 1. Summary Introduce an opt-in diagnostic feature that estimates the mathematical work of a contrib operator invocation. -Start with a pure `com.microsoft::MatMulNBits` estimator and a coverage-aware static report. Then add concrete -runtime counts, Attention coverage, and measured effective throughput. Keep dispatch, placement, memory admission, -and allocation behavior unchanged throughout the diagnostic rollout. +Start with a pure `com.microsoft::MatMulNBits` estimator and a coverage-aware static report. Then complete and +validate MatMulNBits counting, including bias, concrete runtime shapes, and supported route-specific ancillary +work, before expanding to Attention/GQA, measured throughput, or other operators. 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), @@ -25,9 +26,10 @@ The key distinction is: **Recommended first release:** counting conventions, `MatMulNBits` core dense-equivalent operations, and a versioned developer report. Throughput and performance-policy changes are separate deliverables. -Ancillary-work estimates are incremental follow-ups, not prerequisites for this first release. Define their -component contract during reporting design, add shape-based bias counts after static reporting, and add -implementation-dependent work only when its execution route and frequency can be established. +The core-only developer preview remains independently reviewable and shippable, but it is not the MatMulNBits +counting-completion milestone. Prioritize its ancillary and runtime-count follow-ups before starting broader +operator coverage or throughput work. Add implementation-dependent counts only when their route and frequency +can be established; otherwise retain explicit unavailable results rather than inventing a total. ### Initial implementation @@ -37,9 +39,51 @@ core dense-equivalent count under convention version 1 or an explicit unavailabl `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. -This first increment does not add a graph report, runtime instrumentation, public API, or inference-policy changes. -The sources are included by the existing CPU contrib source glob, independently of CUDA, and the tests run in -`onnxruntime_provider_test` with `--gtest_filter="MatMulNBitsComputeEstimate*"` after rebuilding. +The Phase-2 preview adds a static JSON report via the existing session configuration API. There is no runtime +instrumentation, new public API function, or inference-policy change. The estimator sources are included by the +existing CPU contrib source glob, independently of CUDA. + +### Generating a review artifact + +Set `session.compute_estimate_report_path` to a UTF-8 JSON file path before creating a session: + +```python +import onnxruntime as ort + +options = ort.SessionOptions() +options.add_session_config_entry("session.compute_estimate_report_path", "compute_estimates.json") +session = ort.InferenceSession("model.onnx", sess_options=options, providers=["CPUExecutionProvider"]) +``` + +The equivalent C++ configuration is: + +```cpp +Ort::SessionOptions options; +options.AddConfigEntry("session.compute_estimate_report_path", "compute_estimates.json"); +``` + +These examples require an ORT build containing this implementation, not an older installed wheel. +Session initialization generates the report from the final executable graph before removable attributes are +pruned. No `Run()` or profiling call is required. Relative paths use the process working directory; parents must +already exist. Existing files are overwritten, so choose a path distinct from model/optimized-model files and a +different path for each concurrent session. Reports contain node names and shape metadata; review them before sharing. + +Unknown/unsupported estimates are normal report entries. Generation/open/write failures log a warning and do not +fail session initialization. Minimal builds and compile-only sessions warn that export is unsupported. Full builds +with contrib ops disabled can export coverage reports but label the MatMulNBits estimator unavailable. +The unset/empty option performs no report traversal or file I/O. + +The developer-facing schema is version 1. `nodes` contains per-node components and positional input shapes +(unknown dimensions are `null`, omitted optional inputs have state `missing`). `graphs` contains separate summaries +for `main` and each nested graph. Nested identities use `main//`, recursively; +attribute segments escape `~` as `~0` and `/` as `~1`. The top-level summary covers only direct nodes in `main`. +Counts never aggregate alternative branches or unknown loop iteration counts into a per-run total. + +Run the targeted tests after rebuilding: + +```text +onnxruntime_provider_test --gtest_filter="ComputeEstimateReportTest.*:MatMulNBitsComputeEstimate*" +``` ## 2. Motivation and boundaries @@ -189,8 +233,8 @@ per-node result -> coverage-aware report -> optional timing correlation 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 JSON is a proposed reporting format, not an -existing ORT API: +estimate has a reason rather than a numeric placeholder. The following is an abbreviated version-1 report; +the implementation also emits positional input metadata, estimator support, and per-graph summaries: ```json { @@ -207,12 +251,24 @@ existing ORT API: "opset_version": 1, "provider": "CUDAExecutionProvider", "availability": "available", - "estimate_kind": "exact_under_convention", - "input_basis": "static_concrete", - "metric": "dense_equivalent_ops", - "scope": "matmul_core", - "operation_count": "24576", - "excluded_components": ["dequantization", "bias", "packing", "conversion", "padding"] + "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" + ] + } + ] } ], "summary": { @@ -220,15 +276,23 @@ existing ORT API: "estimated_node_count": 1, "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 + "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 node omits -`operation_count` and includes a reason such as `unsupported_operator`, `unknown_shape`, `missing_semantic_value`, -`invalid_metadata`, `overflow`, or `runtime_execution_required`. +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` is reserved for future runtime-dependent components. +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. @@ -249,14 +313,15 @@ from particular contributors. Calendar estimates should follow contract review a | 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 counts | Opt-in invocation estimates using concrete shapes; event identity and bounded caching | Phase 2 | Execution/profiling owners | Dynamic-shape cases pass; no extra device synchronization; agreed overhead gate passes | -| 4. Attention coverage | Incremental packed Attention and GQA support; phase and active-length semantics | Phase 2; Phase 3 for runtime-only metadata | Attention owners | Prefill/decode/chunked, causal/windowed, and ragged reference cases pass | -| 5. Measured throughput | Matched CPU and eager-CUDA timing; units, scope, route provenance where available | Phase 3; Phase 4 only for Attention throughput | Profiling + EP owners | Attribution, aggregation, invalid durations, and timing validation pass | +| 4. Attention coverage | Incremental packed Attention and GQA support; phase and active-length semantics | MatMulNBits counting-completion gate below; Phase 3 for runtime-only metadata | Attention owners | Prefill/decode/chunked, causal/windowed, and ragged reference cases pass | +| 5. Measured throughput | Matched CPU and eager-CUDA timing; units, scope, route provenance where available | MatMulNBits counting-completion gate and Phase 3; Phase 4 only for Attention 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 | -Phase 4 and Phase 5 can proceed independently after their prerequisites. Do not make pilot throughput wait for -complete Attention coverage. Phases 0-2 are the useful, independently shippable MVP. +Phases 0-2 remain the useful, independently shippable core-only MVP. The current priority after Phase 2 is +MatMulNBits counting completion, not Attention or timing. After that gate passes, Phases 4 and 5 can proceed +independently; MatMulNBits throughput need not wait for complete Attention coverage. -The ancillary increments below extend these phases without renumbering or blocking the main rollout: +The ancillary increments below retain their numbering but now participate in the MatMulNBits-first milestone: | Increment | Deliverables | Dependencies | Proposed owner | Exit gate | | --- | --- | --- | --- | --- | @@ -265,13 +330,46 @@ The ancillary increments below extend these phases without renumbering or blocki | 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 | Agree on component semantics during Phase 2 design, but do not implement every component before shipping the MVP. -Increment 2a can proceed alongside runtime-count development. Increment 3a does not block Attention or core-only -throughput reporting; each new component is enabled only for routes with enough evidence. Full ancillary coverage -is not required before subsequent phases or the initial feature PR. +Deliver Increment 2a next, then Phase 3 for MatMulNBits and Increment 3a for the agreed supported routes. +Logical reconstruction remains optional and cannot substitute for evidence of executed ancillary work. +Do not start Attention/GQA, other operator estimators, throughput, or cost-model implementation until the +MatMulNBits counting-completion gate passes. Each route-specific component is enabled only with sufficient evidence. + +### MatMulNBits-first order and completion gate + +Here, "MatMul counting" means the existing `com.microsoft::MatMulNBits` target, not new coverage of standard +`MatMul`, `Gemm`, or every quantized matrix operator. + +1. Finish Phase-2 report review and correctness validation, preserving the independent core-only preview. +2. Implement Increment 2a: separate bias counts and component-level coverage, including absent, unknown, empty, + and overflow cases. Keep the `2MNK` core convention unchanged. +3. Implement Phase 3 for MatMulNBits: concrete invocation geometry, execution identity, repeated/dynamic workloads, + and checked aggregation without stale cache reuse or extra device synchronization. +4. Implement Increment 3a for a bounded, owner-approved CPU/CUDA route matrix. Audit dequantization, unpacking, + conversion, layout, padding/recomputation, and prepacking; identify retained/eliminated work, reuse, repetition, + and initialization/first-use/per-invocation lifetimes. Record unvalidated routes and components as unavailable. +5. Review independent reference counts, end-to-end reports, inference compatibility, build coverage, and measured + count-only overhead against the agreed workload matrix before opening the next operator/timing workstream. + +The gate passes only when: + +- Core and bias counts match independent references for the declared static and runtime workload matrix. +- Every applicable ancillary component on each claimed supported route has a validated count and attribution; + components proven absent or eliminated are distinguished from components whose execution is unknown. +- Unsupported routes/components have explicit reasons and documented limitations. Reviewers approve the bounded + coverage and any deferred gaps; marking all ancillary work unavailable is not completion. +- Reports keep arithmetic categories and lifetimes separate, avoid overlap/shared-preparation double counting, + and do not claim literal total executed FLOPs or whole-operator completeness without supporting evidence. +- Dynamic/repeated/concurrent invocation tests, route/reuse changes, overflow, disabled behavior, and the + approved overhead gate pass. Published review artifacts show both covered and unavailable cases. + +This is completion of counting within a declared, validated support scope, not a promise to count every hardware +instruction or support every EP. Broader route coverage can follow, but silently treating unknown work as zero +or relaxing the gate to begin Attention is not the planned rollout. ### Phase 2: static-report launch -- Begin with an internal/developer reporting surface; the exact command or API is a Phase 0 decision. +- 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. @@ -315,6 +413,8 @@ first-use preparation from steady-state work. Do not attribute shared preparatio ### Phase 4: Attention expansion +Begin implementation only after the MatMulNBits counting-completion gate passes. + 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: @@ -338,6 +438,9 @@ belong to a separately named route-specific metric, not an undocumented change t ### Phase 5: throughput +Begin timing integration only after the MatMulNBits counting-completion gate passes. Count-only overhead +measurements needed to validate Phase 3 are not deferred by this restriction. + Only calculate a rate when the count and duration cover the same invocation and components: ```text @@ -393,7 +496,9 @@ comparing disabled, count-only, and timing-enabled modes. Record exact commands 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:** separately opt in to counts, then timing; record overhead and attribution gaps. -4. **Broader support:** add operators only with owner-reviewed conventions and reference cases. Stabilize a public +4. **MatMulNBits counting completion:** validate bias, runtime counts, and supported ancillary routes against the + gate above before starting Attention/GQA, other operators, or throughput implementation. +5. **Broader support:** add operators only 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 @@ -405,6 +510,10 @@ 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 | @@ -418,6 +527,7 @@ and rollback plan. | 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 bounded MatMulNBits counting-completion scope? | Agree on CPU/CUDA routes, workload matrix, applicable components, and explicit deferred gaps; require owner approval before broader work starts | Before Increment 3a; gate review before Phases 4-5 | ## 10. Completion criteria @@ -425,6 +535,9 @@ The **MVP is complete** when a developer can request a persistent, versioned sta see the correct `MatMulNBits` core count for supported concrete shapes, and understand every exclusion and missing estimate without any change to inference behavior. +The **MatMulNBits counting milestone is complete** only when the MatMulNBits-first gate in Section 6 passes. +Shipping the core-only MVP does not satisfy that gate or authorize starting broader operator/timing work. + 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. From 7b9373f638fb2b997c9a7cdc2e6b6c8c51e1aaa8 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Mon, 19 Oct 2026 17:02:30 -0700 Subject: [PATCH 06/14] feat: Add opt-in static compute estimate reporting for MatMulNBits --- .../onnxruntime_session_options_config_keys.h | 10 + .../core/framework/compute_estimate_report.cc | 256 ++++++++++++++ .../core/framework/compute_estimate_report.h | 21 ++ onnxruntime/core/session/inference_session.cc | 18 + .../compute_estimate_report_test.cc | 325 ++++++++++++++++++ 5 files changed, 630 insertions(+) create mode 100644 onnxruntime/core/framework/compute_estimate_report.cc create mode 100644 onnxruntime/core/framework/compute_estimate_report.h create mode 100644 onnxruntime/test/contrib_ops/compute_estimate_report_test.cc 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..b40cf7be7dee7 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,16 @@ static const char* const kOrtSessionOptionsMaxShapeOverride = "session.max_shape static const char* const kOrtSessionOptionsStrictWorkspaceVerification = "session.strict_workspace_verification"; +// Opt-in static compute-estimation JSON report, written after successful session initialization +// and before removable attributes are pruned. Value is a UTF-8 file path; empty/unset disables reporting. +// Existing files are overwritten. Parent directories are not created. Use distinct paths for concurrent sessions +// and do not use the model or optimized-model path. Reports include node names and canonical input shapes. +// Currently estimates only MatMulNBits core dense-equivalent operations. Subgraphs have separate summaries, +// not a combined per-run total. Unknown/unsupported counts are explicit report entries. +// Report failures log a warning without failing initialization. Minimal builds warn that export is unsupported. +// No profiling run, additional device synchronization, or new public API is required. +static const char* const kOrtSessionOptionsComputeEstimateReportPath = "session.compute_estimate_report_path"; + // 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/core/framework/compute_estimate_report.cc b/onnxruntime/core/framework/compute_estimate_report.cc new file mode 100644 index 0000000000000..8631fb595c36d --- /dev/null +++ b/onnxruntime/core/framework/compute_estimate_report.cc @@ -0,0 +1,256 @@ +// 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 "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/graph/constants.h" +#include "core/graph/graph.h" + +#if !defined(DISABLE_CONTRIB_OPS) +#include "contrib_ops/cpu/quantization/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"; + } + 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)}); + } + 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; + 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()) { + ++estimated_count; + if (!subtotal_overflow && !SafeAdd(known_ops, *operation_count, known_ops)) { + subtotal_overflow = true; + } + } else { + ++unavailable_count; + const std::string reason = entry["unavailable_reason"].get(); + reasons[reason] = 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 summary{{"reported_node_count", estimated_count + unavailable_count}, + {"estimated_node_count", estimated_count}, + {"unavailable_node_count", unavailable_count}, + {"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 && !subtotal_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", 1}, + {"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_only"; +#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; +} + +Status WriteComputeEstimateReport(const Graph& graph, const std::string& utf8_path) { + Status status; + ORT_TRY { + std::string report; + ORT_RETURN_IF_ERROR(BuildComputeEstimateReport(graph, report)); + 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; +} + +} // 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..21069bdc9f234 --- /dev/null +++ b/onnxruntime/core/framework/compute_estimate_report.h @@ -0,0 +1,21 @@ +// 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; + +// 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); + +} // namespace onnxruntime +#endif diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index e9015b3e40b09..7a4aea432bc01 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" @@ -3220,6 +3221,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) @@ -3296,6 +3300,20 @@ common::Status InferenceSession::Initialize() { ort_format_model_mapped_memory_.reset(); } + // Estimate before pruning attributes needed by operator metadata adapters. + const auto compute_report_path = session_options_.config_options.GetConfigOrDefault( + kOrtSessionOptionsComputeEstimateReportPath, ""); + if (!compute_report_path.empty()) { +#if !defined(ORT_MINIMAL_BUILD) + const auto report_status = WriteComputeEstimateReport(graph, compute_report_path); + if (!report_status.IsOK()) { + LOGS(*session_logger_, WARNING) << report_status.ErrorMessage(); + } +#else + LOGS(*session_logger_, WARNING) << "Compute estimate report export is unsupported in a minimal build."; +#endif + } + // once the model is saved, we may remove unnecessary attributes for inference session_state_->PruneRemovableAttributes(); 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..9f5ad30c5999b --- /dev/null +++ b/onnxruntime/test/contrib_ops/compute_estimate_report_test.cc @@ -0,0 +1,325 @@ +// 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 "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/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()); + } + + Node& AddMatMul(std::optional> shape = std::vector{2, 3, 64}, + bool add_identity = 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(); + auto& node = builder_.graph_.AddNode(builder_.graph_.GenerateNodeName("projection \"quoted\"\\\n"), + "MatMulNBits", "", {a, b, scales}, {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() { + std::ifstream stream(path_); + 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; + const auto utf8_path = ToUTF8String(path_.native()); + EXPECT_TRUE(options.config_options.AddConfigEntry(kOrtSessionOptionsComputeEstimateReportPath, + utf8_path.c_str()) + .IsOK()); + return options; + } + + 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"], 1); + EXPECT_EQ(report["estimator_support"], "matmul_nbits_core_only"); + 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); + EXPECT_EQ(report, Report()); +} + +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_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); + } + auto options = ReportOptions(); + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(Initialize(session)); + const auto exported = ReadReport(); + EXPECT_EQ(exported["graphs"].size(), 3u); + EXPECT_EQ(exported["summary"]["known_dense_equivalent_ops"], "0"); + RecordProperty("compute_estimate_report", exported.dump(2)); +} + +TEST_F(ComputeEstimateReportTest, SessionExportsFinalGraphAndRetainsInference) { + AddMatMul(std::vector{2, 3, 64}, 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)); + const auto report = ReadReport(); + EXPECT_EQ(report["nodes"].size(), session.GetGraph().NumberOfNodes()); + EXPECT_EQ(report["summary"]["reported_node_count"], 1); + EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "24576"); + EXPECT_EQ(report["nodes"][0]["provider"], kCpuExecutionProvider); + RecordProperty("compute_estimate_report", report.dump(2)); + + 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], 0.0f); + } +} + +TEST_F(ComputeEstimateReportTest, DisabledExportDoesNotWrite) { + AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + SessionOptions options; + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(Initialize(session)); + EXPECT_FALSE(std::filesystem::exists(path_)); +} + +TEST_F(ComputeEstimateReportTest, WriteFailureWarnsAndInitializationSucceeds) { + AddMatMul(); + ASSERT_STATUS_OK(model_.MainGraph().Resolve()); + auto options = ReportOptions(); + // Opening an existing directory as a file fails without requiring permission changes. + const auto invalid_path = ToUTF8String(std::filesystem::temp_directory_path().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_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, 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"); +} + +} // namespace +} // namespace test +} // namespace onnxruntime +#endif From 79c427d803fef30ed0058de21cc34691bf36f32e Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Fri, 23 Oct 2026 12:40:22 -0700 Subject: [PATCH 07/14] docs: align compute estimation rollout with MatMulNBits and GQA goals --- .../contrib_op_flop_estimation_rollout.md | 89 +++++++++++++++++-- 1 file changed, 84 insertions(+), 5 deletions(-) diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md index 0f873df75106c..84a0dc661958f 100644 --- a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -43,6 +43,14 @@ The Phase-2 preview adds a static JSON report via the existing session configura instrumentation, new public API function, or inference-policy change. The estimator sources are included by the existing CPU contrib source glob, independently of CUDA. +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. + ### Generating a review artifact Set `session.compute_estimate_report_path` to a UTF-8 JSON file path before creating a session: @@ -73,16 +81,32 @@ fail session initialization. Minimal builds and compile-only sessions warn that with contrib ops disabled can export coverage reports but label the MatMulNBits estimator unavailable. The unset/empty option performs no report traversal or file I/O. -The developer-facing schema is version 1. `nodes` contains per-node components and positional input shapes +The developer-facing schema is version 2. It replaces the core-only version-1 preview, adding bias and +component-level subtotals and introducing partial node availability. `estimator_support` is +`matmul_nbits_core_and_bias` when enabled. `nodes` contains per-node components and positional input shapes (unknown dimensions are `null`, omitted optional inputs have state `missing`). `graphs` contains separate summaries for `main` and each nested graph. Nested identities use `main//`, recursively; attribute segments escape `~` as `~0` and `/` as `~1`. The top-level summary covers only direct nodes in `main`. Counts never aggregate alternative branches or unknown loop iteration counts into a per-run total. +Node availability is `available` when both declared components are available, `partial` when only one is +available, and `unavailable` when neither is available (or the operator is unsupported). +The three node counts in each summary are disjoint: `estimated_node_count` means fully estimated within the +declared component scope, `partially_estimated_node_count` means partially estimated, and `unavailable_node_count` +means no component is available. Reason counts include partial nodes. +Each `component_summaries` entry reports its metric, convention, logical basis, per-invocation phase, +estimated/unavailable component counts, unavailable reasons, checked `known_operation_count`, and completeness. +An overflowing component subtotal omits its count and reports overflow. These subtotals cover declared components +of supported direct nodes, not unsupported nodes or complete operator work. +The existing `known_dense_equivalent_ops` and `subtotal_availability` remain core-only; bias is never added to that +subtotal. Overall declared-scope completeness also requires no partial nodes or component-subtotal overflow. +Core still lists bias among its exclusions; bias has its own `floating_point_ops` metric and `bias_add_v1` +convention. Neither component claims route-specific executed work or whole-operator completeness. + Run the targeted tests after rebuilding: ```text -onnxruntime_provider_test --gtest_filter="ComputeEstimateReportTest.*:MatMulNBitsComputeEstimate*" +onnxruntime_provider_test --gtest_filter="ComputeEstimateReportTest.*:MatMulNBitsComputeEstimate*:MatMulNBitsBiasEstimate*" ``` ## 2. Motivation and boundaries @@ -203,6 +227,41 @@ total executed FLOPs. Known-component subtotals must retain exclusions and incom 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 From logical workload to estimated executed FLOPs + +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 transition requires the following work within Increment 3a, before the MatMulNBits completion gate: + +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. Audit the route end to end, including library calls. 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 with independent reference cases and route-boundary, padding, reuse, and lifecycle tests. Use backend + documentation, generated-code inspection, and interpretable hardware counters where needed. 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. + ## 5. Proposed architecture and report ```text @@ -233,12 +292,13 @@ per-node result -> coverage-aware report -> optional timing correlation 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 an abbreviated version-1 report; +estimate has a reason rather than a numeric placeholder. The following is an abbreviated version-2 report +for a node without bias; the implementation also emits positional input metadata, estimator support, and per-graph summaries: ```json { - "schema_version": 1, + "schema_version": 2, "convention_version": "core_arithmetic_v1", "graph_scope": "optimized_executable", "nodes": [ @@ -267,6 +327,21 @@ the implementation also emits positional input metadata, estimator support, and "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" + ] } ] } @@ -274,6 +349,7 @@ the implementation also emits positional input metadata, estimator support, and "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", @@ -290,7 +366,8 @@ Serialize counts as decimal strings to preserve integer precision in JSON consum `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` is reserved for future runtime-dependent components. +`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. @@ -380,6 +457,8 @@ or relaxing the gate to begin Attention is not the planned rollout. ### 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. From f5ed43e130337a9402f7417d7af8bb820c950984 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Fri, 23 Oct 2026 12:41:11 -0700 Subject: [PATCH 08/14] feat: Add logical bias estimates to MatMulNBits compute reports --- .../matmul_nbits_compute_estimate.cc | 67 +++++++++- .../matmul_nbits_compute_estimate.h | 16 +++ .../core/framework/compute_estimate_report.cc | 113 ++++++++++++++++- .../compute_estimate_report_test.cc | 118 +++++++++++++++++- .../matmul_nbits_compute_estimate_test.cc | 118 +++++++++++++++++- 5 files changed, 414 insertions(+), 18 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc index 285e21403c4b2..3f2b9425b5e7f 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc @@ -4,6 +4,7 @@ #include "contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h" #include +#include #include "core/common/safeint.h" @@ -15,8 +16,9 @@ namespace onnxruntime { namespace contrib { +namespace { -MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute( +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; })) { @@ -31,6 +33,17 @@ MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute( 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}; } @@ -50,8 +63,36 @@ MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute( 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) -MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute(const Node& node) { +namespace { +template +auto EstimateFromNode(const Node& node) + -> std::conditional_t { if (node.Domain() != kMSDomain || node.OpType() != "MatMulNBits") { return ComputeEstimateUnavailableReason::UnsupportedOperator; } @@ -75,6 +116,13 @@ MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute(const Node& node) { 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; @@ -93,7 +141,20 @@ MatMulNBitsComputeEstimateResult EstimateMatMulNBitsCompute(const Node& node) { } } - return EstimateMatMulNBitsCompute(dimensions, k_it->second.i(), n_it->second.i()); + 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 diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h index 46350aa0e131f..6d4925a0c9cbc 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include #include @@ -19,6 +20,7 @@ enum class ComputeEstimateUnavailableReason { UnknownShape, InvalidMetadata, Overflow, + MissingSemanticValue, }; struct MatMulNBitsComputeEstimate { @@ -31,16 +33,30 @@ struct MatMulNBitsComputeEstimate { 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 diff --git a/onnxruntime/core/framework/compute_estimate_report.cc b/onnxruntime/core/framework/compute_estimate_report.cc index 8631fb595c36d..9ede8676a1fce 100644 --- a/onnxruntime/core/framework/compute_estimate_report.cc +++ b/onnxruntime/core/framework/compute_estimate_report.cc @@ -45,6 +45,8 @@ const char* ReasonName(contrib::ComputeEstimateUnavailableReason reason) { return "invalid_metadata"; case contrib::ComputeEstimateUnavailableReason::Overflow: return "overflow"; + case contrib::ComputeEstimateUnavailableReason::MissingSemanticValue: + return "missing_semantic_value"; } ORT_THROW("Invalid compute estimate reason."); } @@ -116,6 +118,42 @@ Json EstimateNode(const Node& node, std::optional& operation_count) { 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(); @@ -143,6 +181,18 @@ void AppendGraph(const Graph& graph, const std::string& path, Json& report) { 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(); @@ -160,15 +210,42 @@ void AppendGraph(const Graph& graph, const std::string& path, Json& report) { std::optional operation_count; entry.update(EstimateNode(node, operation_count)); if (operation_count.has_value()) { - ++estimated_count; if (!subtotal_overflow && !SafeAdd(known_ops, *operation_count, known_ops)) { subtotal_overflow = true; } + } + if (entry["availability"] == "available") { + ++estimated_count; } else { - ++unavailable_count; + 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(); @@ -185,13 +262,37 @@ void AppendGraph(const Graph& graph, const std::string& path, Json& report) { } } - Json summary{{"reported_node_count", estimated_count + unavailable_count}, + 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 && !subtotal_overflow}, + {"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}}; @@ -208,7 +309,7 @@ void AppendGraph(const Graph& graph, const std::string& path, Json& report) { Status BuildComputeEstimateReport(const Graph& graph, std::string& report) { Status status; ORT_TRY { - Json value{{"schema_version", 1}, + Json value{{"schema_version", 2}, {"convention_version", "core_arithmetic_v1"}, {"graph_scope", "optimized_executable"}, {"nodes", Json::array()}, @@ -216,7 +317,7 @@ Status BuildComputeEstimateReport(const Graph& graph, std::string& report) { #if defined(DISABLE_CONTRIB_OPS) value["estimator_support"] = "disabled"; #else - value["estimator_support"] = "matmul_nbits_core_only"; + value["estimator_support"] = "matmul_nbits_core_and_bias"; #endif AppendGraph(graph, "main", value); value["summary"] = value["graphs"][0]["summary"]; diff --git a/onnxruntime/test/contrib_ops/compute_estimate_report_test.cc b/onnxruntime/test/contrib_ops/compute_estimate_report_test.cc index 9f5ad30c5999b..10850d6a5a9e4 100644 --- a/onnxruntime/test/contrib_ops/compute_estimate_report_test.cc +++ b/onnxruntime/test/contrib_ops/compute_estimate_report_test.cc @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -51,13 +52,20 @@ class ComputeEstimateReportTest : public ::testing::Test { } Node& AddMatMul(std::optional> shape = std::vector{2, 3, 64}, - bool add_identity = false) { + 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", "", {a, b, scales}, {output}, nullptr, kMSDomain); + "MatMulNBits", "", inputs, {output}, nullptr, kMSDomain); node.AddAttribute("K", int64_t{64}); node.AddAttribute("N", int64_t{32}); node.AddAttribute("bits", int64_t{4}); @@ -112,8 +120,8 @@ TEST_F(ComputeEstimateReportTest, ComponentContractAndDeterminism) { auto& node = AddMatMul(); ASSERT_STATUS_OK(model_.MainGraph().Resolve()); const auto report = Report(); - EXPECT_EQ(report["schema_version"], 1); - EXPECT_EQ(report["estimator_support"], "matmul_nbits_core_only"); + 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()); @@ -131,15 +139,112 @@ TEST_F(ComputeEstimateReportTest, ComponentContractAndDeterminism) { 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); @@ -244,7 +349,7 @@ TEST_F(ComputeEstimateReportTest, NestedGraphsAreSeparateScopes) { } TEST_F(ComputeEstimateReportTest, SessionExportsFinalGraphAndRetainsInference) { - AddMatMul(std::vector{2, 3, 64}, true); + AddMatMul(std::vector{2, 3, 64}, true, true); ASSERT_STATUS_OK(model_.MainGraph().Resolve()); auto options = ReportOptions(); options.graph_optimization_level = TransformerLevel::Level1; @@ -255,6 +360,7 @@ TEST_F(ComputeEstimateReportTest, SessionExportsFinalGraphAndRetainsInference) { EXPECT_EQ(report["summary"]["reported_node_count"], 1); EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "24576"); EXPECT_EQ(report["nodes"][0]["provider"], kCpuExecutionProvider); + EXPECT_EQ(report["nodes"][0]["components"][1]["operation_count"], "192"); RecordProperty("compute_estimate_report", report.dump(2)); const auto& input_name = session.GetGraph().GetInputs()[0]->Name(); @@ -271,7 +377,7 @@ TEST_F(ComputeEstimateReportTest, SessionExportsFinalGraphAndRetainsInference) { 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], 0.0f); + EXPECT_EQ(result.Data()[i], 1.0f); } } diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc index abe7ac99846b8..ea49b5ada0044 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc @@ -25,10 +25,82 @@ 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); @@ -114,16 +186,16 @@ class MatMulNBitsComputeEstimateGraphTest : public ::testing::Test { template Node& AddMatMul(std::optional> shape = std::vector{2, 3, 64}, - int64_t bits = 4, bool bias = false) { + 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) { + if (bias || explicit_absent_bias) { auto& missing = builder_.graph_.GetOrCreateNodeArg("", nullptr); inputs.push_back(&missing); inputs.push_back(&missing); - inputs.push_back(builder_.MakeInput(std::vector{32})); + inputs.push_back(bias ? builder_.MakeInput(std::vector{32}) : &missing); } auto* output = builder_.MakeOutput(std::nullopt); auto& node = builder_.AddNode("MatMulNBits", inputs, {output}, kMSDomain); @@ -145,6 +217,43 @@ TEST_F(MatMulNBitsComputeEstimateGraphTest, ResolvedGraph) { 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()); @@ -186,6 +295,7 @@ TEST_F(MatMulNBitsComputeEstimateGraphTest, MissingAttribute) { ASSERT_STATUS_OK(model_.MainGraph().Resolve()); node.ClearAttribute("K"); ExpectUnavailable(EstimateMatMulNBitsCompute(node), ComputeEstimateUnavailableReason::InvalidMetadata); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(node), ComputeEstimateUnavailableReason::InvalidMetadata); } TEST_F(MatMulNBitsComputeEstimateGraphTest, WrongAttributeType) { @@ -205,6 +315,7 @@ TEST_F(MatMulNBitsComputeEstimateGraphTest, UnsupportedVersion) { ASSERT_STATUS_OK(model_.MainGraph().Resolve()); node.SetSinceVersion(2); ExpectUnavailable(EstimateMatMulNBitsCompute(node), ComputeEstimateUnavailableReason::UnsupportedVersion); + ExpectBiasUnavailable(EstimateMatMulNBitsBias(node), ComputeEstimateUnavailableReason::UnsupportedVersion); } TEST_F(MatMulNBitsComputeEstimateGraphTest, EmptyGraphWorkload) { @@ -218,6 +329,7 @@ TEST_F(MatMulNBitsComputeEstimateGraphTest, UnsupportedDomain) { 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) { From d07fe662ab3196a6f46445327dfc0b5929568fc9 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Mon, 26 Oct 2026 08:27:46 -0700 Subject: [PATCH 09/14] Add structured runtime workload collector foundation --- .../contrib_op_flop_estimation_rollout.md | 202 +++++++++----- .../framework/runtime_workload_collector.cc | 96 +++++++ .../framework/runtime_workload_collector.h | 115 ++++++++ .../runtime_workload_collector_test.cc | 261 ++++++++++++++++++ 4 files changed, 606 insertions(+), 68 deletions(-) create mode 100644 onnxruntime/core/framework/runtime_workload_collector.cc create mode 100644 onnxruntime/core/framework/runtime_workload_collector.h create mode 100644 onnxruntime/test/framework/runtime_workload_collector_test.cc diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md index 84a0dc661958f..2d341cc2443d2 100644 --- a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -5,9 +5,10 @@ ## 1. Summary Introduce an opt-in diagnostic feature that estimates the mathematical work of a contrib operator invocation. -Start with a pure `com.microsoft::MatMulNBits` estimator and a coverage-aware static report. Then complete and -validate MatMulNBits counting, including bias, concrete runtime shapes, and supported route-specific ancillary -work, before expanding to Attention/GQA, measured throughput, or other operators. Keep dispatch, placement, +The minimum operator scope is `com.microsoft::MatMulNBits` and GroupQueryAttention (GQA). Start with the existing +MatMulNBits estimator, bias component, and legacy static-report preview; next connect estimates to runtime workload metadata, +then 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 @@ -23,13 +24,16 @@ The key distinction is: - **FLOP/s or equivalent operations/s:** a count divided by a matched measured duration. A graph alone cannot establish this rate. -**Recommended first release:** counting conventions, `MatMulNBits` core dense-equivalent operations, and a versioned -developer report. Throughput and performance-policy changes are separate deliverables. +**Target deliverable:** runtime-derived workload estimates for MatMulNBits and GQA in a versioned developer report. +Initialization should configure runtime reporting, not generate the target workload report. The existing static +export is a legacy developer preview to replace in a later batch, not the desired reporting lifecycle. +Throughput and performance-policy changes are separate deliverables. -The core-only developer preview remains independently reviewable and shippable, but it is not the MatMulNBits -counting-completion milestone. Prioritize its ancillary and runtime-count follow-ups before starting broader -operator coverage or throughput work. Add implementation-dependent counts only when their route and frequency -can be established; otherwise retain explicit unavailable results rather than inventing a total. +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 @@ -39,8 +43,8 @@ core dense-equivalent count under convention version 1 or an explicit unavailabl `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 adds a static JSON report via the existing session configuration API. There is no runtime -instrumentation, new public API function, or inference-policy change. The estimator sources are included by the +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 CPU contrib source glob, independently of CUDA. Increment 2a adds a separate `bias` component: `MN` logical floating-point additions when the sixth schema input @@ -51,8 +55,46 @@ dimension `K` makes the core count zero but does not make bias additions zero if 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. +The existing initialization-time export below remains a legacy preview until replacement is implemented. +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 (proposed): execution integration.** Pass an optional recorder through execution context, capture metadata +at actual CPU MatMulNBits dispatch/backend call sites, and commit only established successful invocations. +Define run failure/incompleteness behavior and ensure collection errors cannot fail inference. Test disabled +behavior, actual routes, changing shapes, and concurrent/repeated runs. Do not use logs as data transport. + +**Batch 3 (proposed): runtime counts and export.** Apply existing checked core/bias formulas to invocation records, +export versioned runtime reports with partial/drop/run-status coverage, and replace initialization-time workload +export. Agree on configuration and export lifecycle before implementation; GQA extends this path afterward. + ### Generating a review artifact +The following describes the still-available legacy static preview, not the target runtime report. + Set `session.compute_estimate_report_path` to a UTF-8 JSON file path before creating a session: ```python @@ -227,7 +269,7 @@ total executed FLOPs. Known-component subtotals must retain exclusions and incom 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 From logical workload to estimated executed FLOPs +### 4.4 Baseline estimates and optional route-aware refinements Provide two distinct results rather than progressively relabeling the existing core count: @@ -239,28 +281,38 @@ Provide two distinct results rather than progressively relabeling the existing c 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 transition requires the following work within Increment 3a, before the MatMulNBits completion gate: +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. Audit the route end to end, including library calls. Establish floating-point versus integer core arithmetic, +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 with independent reference cases and route-boundary, padding, reuse, and lifecycle tests. Use backend - documentation, generated-code inspection, and interpretable hardware counters where needed. An opaque library - call is not sufficient evidence for its internal arithmetic count. +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. +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 @@ -389,16 +441,17 @@ from particular contributors. Calendar estimates should follow contract review a | 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 counts | Opt-in invocation estimates using concrete shapes; event identity and bounded caching | Phase 2 | Execution/profiling owners | Dynamic-shape cases pass; no extra device synchronization; agreed overhead gate passes | -| 4. Attention coverage | Incremental packed Attention and GQA support; phase and active-length semantics | MatMulNBits counting-completion gate below; Phase 3 for runtime-only metadata | Attention owners | Prefill/decode/chunked, causal/windowed, and ragged reference cases pass | -| 5. Measured throughput | Matched CPU and eager-CUDA timing; units, scope, route provenance where available | MatMulNBits counting-completion gate and Phase 3; Phase 4 only for Attention throughput | Profiling + EP owners | Attribution, aggregation, invalid durations, and timing validation pass | +| 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 the useful, independently shippable core-only MVP. The current priority after Phase 2 is -MatMulNBits counting completion, not Attention or timing. After that gate passes, Phases 4 and 5 can proceed -independently; MatMulNBits throughput need not wait for complete Attention coverage. +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 below retain their numbering but now participate in the MatMulNBits-first milestone: +The ancillary increments retain their numbering; route-specific work does not block the minimum operator scope: | Increment | Deliverables | Dependencies | Proposed owner | Exit gate | | --- | --- | --- | --- | --- | @@ -406,43 +459,44 @@ The ancillary increments below retain their numbering but now participate in the | 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 | -Agree on component semantics during Phase 2 design, but do not implement every component before shipping the MVP. -Deliver Increment 2a next, then Phase 3 for MatMulNBits and Increment 3a for the agreed supported routes. -Logical reconstruction remains optional and cannot substitute for evidence of executed ancillary work. -Do not start Attention/GQA, other operator estimators, throughput, or cost-model implementation until the -MatMulNBits counting-completion gate passes. Each route-specific component is enabled only with sufficient evidence. +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. -### MatMulNBits-first order and completion gate +### 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. Finish Phase-2 report review and correctness validation, preserving the independent core-only preview. -2. Implement Increment 2a: separate bias counts and component-level coverage, including absent, unknown, empty, - and overflow cases. Keep the `2MNK` core convention unchanged. -3. Implement Phase 3 for MatMulNBits: concrete invocation geometry, execution identity, repeated/dynamic workloads, - and checked aggregation without stale cache reuse or extra device synchronization. -4. Implement Increment 3a for a bounded, owner-approved CPU/CUDA route matrix. Audit dequantization, unpacking, - conversion, layout, padding/recomputation, and prepacking; identify retained/eliminated work, reuse, repetition, - and initialization/first-use/per-invocation lifetimes. Record unvalidated routes and components as unavailable. -5. Review independent reference counts, end-to-end reports, inference compatibility, build coverage, and measured - count-only overhead against the agreed workload matrix before opening the next operator/timing workstream. - -The gate passes only when: - -- Core and bias counts match independent references for the declared static and runtime workload matrix. -- Every applicable ancillary component on each claimed supported route has a validated count and attribution; - components proven absent or eliminated are distinguished from components whose execution is unknown. -- Unsupported routes/components have explicit reasons and documented limitations. Reviewers approve the bounded - coverage and any deferred gaps; marking all ancillary work unavailable is not completion. -- Reports keep arithmetic categories and lifetimes separate, avoid overlap/shared-preparation double counting, - and do not claim literal total executed FLOPs or whole-operator completeness without supporting evidence. -- Dynamic/repeated/concurrent invocation tests, route/reuse changes, overflow, disabled behavior, and the - approved overhead gate pass. Published review artifacts show both covered and unavailable cases. - -This is completion of counting within a declared, validated support scope, not a promise to count every hardware -instruction or support every EP. Broader route coverage can follow, but silently treating unknown work as zero -or relaxing the gate to begin Attention is not the planned rollout. +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 @@ -467,6 +521,13 @@ Initial bias implementation is available for review in the version-2 report desc ### 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. @@ -476,7 +537,7 @@ Initial bias implementation is available for review in the version-2 report desc ### Increment 3a: route-specific ancillary work -For one supported route at a time: +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. @@ -492,7 +553,10 @@ first-use preparation from steady-state work. Do not attribute shared preparatio ### Phase 4: Attention expansion -Begin implementation only after the MatMulNBits counting-completion gate passes. +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: @@ -517,8 +581,8 @@ belong to a separately named route-specific metric, not an undocumented change t ### Phase 5: throughput -Begin timing integration only after the MatMulNBits counting-completion gate passes. Count-only overhead -measurements needed to validate Phase 3 are not deferred by this restriction. +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: @@ -574,10 +638,11 @@ comparing disabled, count-only, and timing-enabled modes. Record exact commands 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:** separately opt in to counts, then timing; record overhead and attribution gaps. -4. **MatMulNBits counting completion:** validate bias, runtime counts, and supported ancillary routes against the - gate above before starting Attention/GQA, other operators, or throughput implementation. -5. **Broader support:** add operators only with owner-reviewed conventions and reference cases. Stabilize a public +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 @@ -606,7 +671,7 @@ additional build coverage, and later-phase semantics still require team agreemen | 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 bounded MatMulNBits counting-completion scope? | Agree on CPU/CUDA routes, workload matrix, applicable components, and explicit deferred gaps; require owner approval before broader work starts | Before Increment 3a; gate review before Phases 4-5 | +| 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 @@ -614,8 +679,9 @@ The **MVP is complete** when a developer can request a persistent, versioned sta see the correct `MatMulNBits` core count for supported concrete shapes, and understand every exclusion and missing estimate without any change to inference behavior. -The **MatMulNBits counting milestone is complete** only when the MatMulNBits-first gate in Section 6 passes. -Shipping the core-only MVP does not satisfy that gate or authorize starting broader operator/timing work. +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 diff --git a/onnxruntime/core/framework/runtime_workload_collector.cc b/onnxruntime/core/framework/runtime_workload_collector.cc new file mode 100644 index 0000000000000..8ad8ea0388b9f --- /dev/null +++ b/onnxruntime/core/framework/runtime_workload_collector.cc @@ -0,0 +1,96 @@ +// 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 "core/common/safeint.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_}; + 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; +} + +} // 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..d8745e2d1b3ee --- /dev/null +++ b/onnxruntime/core/framework/runtime_workload_collector.h @@ -0,0 +1,115 @@ +// 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 { + +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; + // 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; +}; + +// 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); + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RuntimeWorkloadCollector); + + // Copies metadata without retaining tensors, caller storage, or graph objects. + // Rejection is an explicit diagnostic error; future 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; + + private: + const uint64_t session_id_; + const uint64_t run_id_; + const RuntimeWorkloadLimits limits_; + mutable std::mutex mutex_; + InlinedVector records_; + RuntimeWorkloadCollectionStats stats_; +}; + +} // namespace onnxruntime + +#endif 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..2554e645d4b73 --- /dev/null +++ b/onnxruntime/test/framework/runtime_workload_collector_test.cc @@ -0,0 +1,261 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) + +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "core/framework/runtime_workload_collector.h" +#include "core/graph/constants.h" +#include "test/util/include/asserts.h" + +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); +} + +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); +} + +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, 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()); + } +} + +} // namespace +} // namespace test +} // namespace onnxruntime + +#endif From 99c5bbe120a89f40ae1bb7a259785a761d321d43 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Mon, 26 Oct 2026 10:55:40 -0700 Subject: [PATCH 10/14] Connect runtime workload collection to CPU MatMulNBits execution --- .../contrib_op_flop_estimation_rollout.md | 53 +- .../core/framework/op_kernel_context.h | 3 + .../onnxruntime/core/framework/run_options.h | 7 + .../cpu/quantization/matmul_nbits.cc | 52 +- .../framework/op_kernel_context_internal.h | 22 +- .../framework/runtime_workload_collector.cc | 74 ++- .../framework/runtime_workload_collector.h | 41 +- .../core/framework/sequential_executor.cc | 14 +- .../core/framework/sequential_executor.h | 3 +- onnxruntime/core/framework/utils.cc | 26 +- onnxruntime/core/framework/utils.h | 10 +- .../core/providers/cpu/controlflow/if.cc | 2 +- .../core/providers/cpu/controlflow/loop.cc | 2 +- .../providers/cpu/controlflow/scan_utils.cc | 2 +- onnxruntime/core/session/inference_session.cc | 28 +- onnxruntime/core/session/inference_session.h | 3 +- .../runtime_workload_collector_test.cc | 555 ++++++++++++++++++ 17 files changed, 857 insertions(+), 40 deletions(-) diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md index 2d341cc2443d2..bdd1fdffc02a0 100644 --- a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -82,10 +82,55 @@ handling, logger integration, or JSON/file export. No production inference behav 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 (proposed): execution integration.** Pass an optional recorder through execution context, capture metadata -at actual CPU MatMulNBits dispatch/backend call sites, and commit only established successful invocations. -Define run failure/incompleteness behavior and ensure collection errors cannot fail inference. Test disabled -behavior, actual routes, changing shapes, and concurrent/repeated runs. Do not use logs as data transport. +**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. 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 (proposed): runtime counts and export.** Apply existing checked core/bias formulas to invocation records, export versioned runtime reports with partial/drop/run-status coverage, and replace initialization-time workload 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/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/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 index 8ad8ea0388b9f..ab1ef73db13dc 100644 --- a/onnxruntime/core/framework/runtime_workload_collector.cc +++ b/onnxruntime/core/framework/runtime_workload_collector.cc @@ -5,11 +5,13 @@ #if !defined(ORT_MINIMAL_BUILD) +#include #include #include #include #include "core/common/safeint.h" +#include "core/graph/graph.h" namespace onnxruntime { namespace { @@ -79,7 +81,7 @@ Status RuntimeWorkloadCollector::GetSnapshot(RuntimeWorkloadSnapshot& output) co std::lock_guard lock(mutex_); Status status; ORT_TRY { - RuntimeWorkloadSnapshot snapshot{session_id_, run_id_, records_, stats_}; + RuntimeWorkloadSnapshot snapshot{session_id_, run_id_, records_, stats_, run_state_}; output = std::move(snapshot); return Status::OK(); } @@ -91,6 +93,76 @@ Status RuntimeWorkloadCollector::GetSnapshot(RuntimeWorkloadSnapshot& output) co 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 index d8745e2d1b3ee..20fa33c37cc1f 100644 --- a/onnxruntime/core/framework/runtime_workload_collector.h +++ b/onnxruntime/core/framework/runtime_workload_collector.h @@ -19,6 +19,16 @@ namespace onnxruntime { +class Graph; +class Node; + +enum class RuntimeWorkloadRunState { + NotStarted, + Running, + Succeeded, + Failed, +}; + enum class MatMulNBitsRuntimeRoute { EmptyOutput, PackedLut, @@ -76,6 +86,8 @@ struct RuntimeWorkloadCollectionStats { 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; }; @@ -85,6 +97,9 @@ struct RuntimeWorkloadSnapshot { 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; }; // One collector per run. Callers own its lifetime and must keep it alive until all submitters finish. @@ -92,22 +107,42 @@ struct RuntimeWorkloadSnapshot { 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; future execution integration must not fail inference on it. + // 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: - const uint64_t session_id_; - const uint64_t run_id_; + 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 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 7a4aea432bc01..48c8efdf45308 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -3724,15 +3724,33 @@ 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) + if (run_options.runtime_workload_collector && + run_options.runtime_workload_collector->BeginRun(session_id_)) { + workload_collector = run_options.runtime_workload_collector; + } + bool succeeded = false; + auto finish_collection = gsl::finally([&]() { + if (workload_collector) { + workload_collector->EndRun(succeeded); + } + }); +#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()) { @@ -3920,7 +3938,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 @@ -4077,7 +4095,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..f1ca7f3239630 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); diff --git a/onnxruntime/test/framework/runtime_workload_collector_test.cc b/onnxruntime/test/framework/runtime_workload_collector_test.cc index 2554e645d4b73..2a041e69f0901 100644 --- a/onnxruntime/test/framework/runtime_workload_collector_test.cc +++ b/onnxruntime/test/framework/runtime_workload_collector_test.cc @@ -4,6 +4,7 @@ #if !defined(ORT_MINIMAL_BUILD) #include +#include #include #include #include @@ -14,6 +15,14 @@ #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 { @@ -192,6 +201,27 @@ TEST(RuntimeWorkloadCollectorTest, ZeroLimitsExplicitlyRejectSubmissions) { 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}); @@ -254,6 +284,531 @@ TEST(RuntimeWorkloadCollectorTest, ConcurrentSubmissionsAndSnapshotsRespectLimit } } +#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; +} + +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); +} + +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); + 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"); + } +} + +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); + 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 From 350767053fa8f600ff36bfa30408a65172e24b57 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Mon, 26 Oct 2026 11:03:27 -0700 Subject: [PATCH 11/14] Fix workload collector propagation through transformer subgraphs Forward the actual workload collector through eight GPT beam/greedy and T5/Whisper calls. Preserve the T5 encoder/decoder MatMulNBits regression sequences while covering collector propagation. Validation: CPU Debug builds successful; 25 collector tests passed with 1 LUT skipped; 26 search regression tests passed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_op_flop_estimation_rollout.md | 5 +- .../cpu/transformers/beam_search_impl_gpt.h | 6 +- .../cpu/transformers/beam_search_impl_t5.h | 6 +- .../transformers/beam_search_impl_whisper.h | 6 +- .../cpu/transformers/greedy_search_impl_gpt.h | 6 +- .../runtime_workload_collector_test.cc | 94 +++++++++++++++++++ 6 files changed, 114 insertions(+), 9 deletions(-) diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md index bdd1fdffc02a0..43ef17fd17ea4 100644 --- a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -102,7 +102,10 @@ Collectors cannot be reused: subsequent/concurrent attempts disable capture for 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. The context owns the recorder; kernels do not store per-run mutable state. +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. 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/test/framework/runtime_workload_collector_test.cc b/onnxruntime/test/framework/runtime_workload_collector_test.cc index 2a041e69f0901..b8642980a4135 100644 --- a/onnxruntime/test/framework/runtime_workload_collector_test.cc +++ b/onnxruntime/test/framework/runtime_workload_collector_test.cc @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -411,6 +412,99 @@ ModelProto LoopWorkloadModel() { 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 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; From 907fa852704e3993a6b64dbd62691f747a365441 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Mon, 26 Oct 2026 13:39:11 -0700 Subject: [PATCH 12/14] Export runtime MatMulNBits compute estimates --- .../contrib_op_flop_estimation_rollout.md | 136 ++++--- .../onnxruntime_session_options_config_keys.h | 23 +- .../core/framework/compute_estimate_report.cc | 243 ++++++++++++- .../core/framework/compute_estimate_report.h | 11 + .../framework/runtime_workload_collector.cc | 2 +- .../framework/runtime_workload_collector.h | 1 + onnxruntime/core/session/inference_session.cc | 57 ++- onnxruntime/core/session/inference_session.h | 6 + .../compute_estimate_report_test.cc | 336 +++++++++++++++++- .../runtime_workload_collector_test.cc | 28 ++ 10 files changed, 745 insertions(+), 98 deletions(-) diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md index 43ef17fd17ea4..9587f9e11a0e4 100644 --- a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -1,13 +1,13 @@ # Contrib Operator FLOP Estimation: Design and Rollout Plan -**Status:** Rollout proposal with an initial internal implementation; later phases remain proposed. +**Status:** CPU MatMulNBits runtime count/report preview implemented; GQA and later phases remain proposed. ## 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 legacy static-report preview; next connect estimates to runtime workload metadata, -then add GQA baseline estimates. Route-aware adjustments are optional, incremental refinements, not an exact +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. @@ -25,8 +25,8 @@ The key distinction is: establish this rate. **Target deliverable:** runtime-derived workload estimates for MatMulNBits and GQA in a versioned developer report. -Initialization should configure runtime reporting, not generate the target workload report. The existing static -export is a legacy developer preview to replace in a later batch, not the desired reporting lifecycle. +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. @@ -58,7 +58,7 @@ 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. -The existing initialization-time export below remains a legacy preview until replacement is implemented. +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).** @@ -135,68 +135,103 @@ capability-gated LUT route, activation conversion, dynamic/repeated/concurrent r 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 (proposed): runtime counts and export.** Apply existing checked core/bias formulas to invocation records, -export versioned runtime reports with partial/drop/run-status coverage, and replace initialization-time workload -export. Agree on configuration and export lifecycle before implementation; GQA extends this path afterward. +**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 -The following describes the still-available legacy static preview, not the target runtime report. - -Set `session.compute_estimate_report_path` to a UTF-8 JSON file path before creating a session: +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.json") +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.json"); +options.AddConfigEntry("session.compute_estimate_report_path", "compute_estimates"); ``` These examples require an ORT build containing this implementation, not an older installed wheel. -Session initialization generates the report from the final executable graph before removable attributes are -pruned. No `Run()` or profiling call is required. Relative paths use the process working directory; parents must -already exist. Existing files are overwritten, so choose a path distinct from model/optimized-model files and a -different path for each concurrent session. Reports contain node names and shape metadata; review them before sharing. - -Unknown/unsupported estimates are normal report entries. Generation/open/write failures log a warning and do not -fail session initialization. Minimal builds and compile-only sessions warn that export is unsupported. Full builds -with contrib ops disabled can export coverage reports but label the MatMulNBits estimator unavailable. -The unset/empty option performs no report traversal or file I/O. - -The developer-facing schema is version 2. It replaces the core-only version-1 preview, adding bias and -component-level subtotals and introducing partial node availability. `estimator_support` is -`matmul_nbits_core_and_bias` when enabled. `nodes` contains per-node components and positional input shapes -(unknown dimensions are `null`, omitted optional inputs have state `missing`). `graphs` contains separate summaries -for `main` and each nested graph. Nested identities use `main//`, recursively; -attribute segments escape `~` as `~0` and `/` as `~1`. The top-level summary covers only direct nodes in `main`. -Counts never aggregate alternative branches or unknown loop iteration counts into a per-run total. - -Node availability is `available` when both declared components are available, `partial` when only one is -available, and `unavailable` when neither is available (or the operator is unsupported). -The three node counts in each summary are disjoint: `estimated_node_count` means fully estimated within the -declared component scope, `partially_estimated_node_count` means partially estimated, and `unavailable_node_count` -means no component is available. Reason counts include partial nodes. -Each `component_summaries` entry reports its metric, convention, logical basis, per-invocation phase, -estimated/unavailable component counts, unavailable reasons, checked `known_operation_count`, and completeness. -An overflowing component subtotal omits its count and reports overflow. These subtotals cover declared components -of supported direct nodes, not unsupported nodes or complete operator work. -The existing `known_dense_equivalent_ops` and `subtotal_availability` remain core-only; bias is never added to that -subtotal. Overall declared-scope completeness also requires no partial nodes or component-subtotal overflow. -Core still lists bias among its exclusions; bias has its own `floating_point_ops` metric and `bias_add_v1` -convention. Neither component claims route-specific executed work or whole-operator completeness. +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.*:MatMulNBitsComputeEstimate*:MatMulNBitsBiasEstimate*" +onnxruntime_provider_test --gtest_filter="ComputeEstimateReportTest.*:RuntimeComputeEstimateReportTest.*:MatMulNBitsComputeEstimate*:MatMulNBitsBiasEstimate*" +onnxruntime_test_all --gtest_filter="RuntimeWorkloadCollectorTest.*:RuntimeWorkloadExecutionTest.*" ``` ## 2. Motivation and boundaries @@ -392,7 +427,7 @@ per-node result -> coverage-aware report -> optional timing correlation 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 an abbreviated version-2 report +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: @@ -723,9 +758,10 @@ additional build coverage, and later-phase semantics still require team agreemen ## 10. Completion criteria -The **MVP is complete** when a developer can request a persistent, versioned static report for an optimized graph, -see the correct `MatMulNBits` core count for supported concrete shapes, and understand every exclusion and missing -estimate without any change to inference behavior. +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, 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 b40cf7be7dee7..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,16 +577,23 @@ static const char* const kOrtSessionOptionsMaxShapeOverride = "session.max_shape static const char* const kOrtSessionOptionsStrictWorkspaceVerification = "session.strict_workspace_verification"; -// Opt-in static compute-estimation JSON report, written after successful session initialization -// and before removable attributes are pruned. Value is a UTF-8 file path; empty/unset disables reporting. -// Existing files are overwritten. Parent directories are not created. Use distinct paths for concurrent sessions -// and do not use the model or optimized-model path. Reports include node names and canonical input shapes. -// Currently estimates only MatMulNBits core dense-equivalent operations. Subgraphs have separate summaries, -// not a combined per-run total. Unknown/unsupported counts are explicit report entries. -// Report failures log a warning without failing initialization. Minimal builds warn that export is unsupported. -// No profiling run, additional device synchronization, or new public API is required. +// 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/core/framework/compute_estimate_report.cc b/onnxruntime/core/framework/compute_estimate_report.cc index 9ede8676a1fce..ba51bce7c14d4 100644 --- a/onnxruntime/core/framework/compute_estimate_report.cc +++ b/onnxruntime/core/framework/compute_estimate_report.cc @@ -6,6 +6,8 @@ #if !defined(ORT_MINIMAL_BUILD) #include +#include +#include #include #include #include @@ -16,8 +18,10 @@ #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 "contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h" @@ -333,11 +337,10 @@ Status BuildComputeEstimateReport(const Graph& graph, std::string& report) { return status; } -Status WriteComputeEstimateReport(const Graph& graph, const std::string& utf8_path) { +namespace { +Status WriteReport(const std::string& report, const std::string& utf8_path) { Status status; ORT_TRY { - std::string report; - ORT_RETURN_IF_ERROR(BuildComputeEstimateReport(graph, report)); 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; @@ -353,5 +356,239 @@ Status WriteComputeEstimateReport(const Graph& graph, const std::string& utf8_pa 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 index 21069bdc9f234..756e2a73a5cd6 100644 --- a/onnxruntime/core/framework/compute_estimate_report.h +++ b/onnxruntime/core/framework/compute_estimate_report.h @@ -11,11 +11,22 @@ 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/runtime_workload_collector.cc b/onnxruntime/core/framework/runtime_workload_collector.cc index ab1ef73db13dc..be2c51ee37a04 100644 --- a/onnxruntime/core/framework/runtime_workload_collector.cc +++ b/onnxruntime/core/framework/runtime_workload_collector.cc @@ -81,7 +81,7 @@ Status RuntimeWorkloadCollector::GetSnapshot(RuntimeWorkloadSnapshot& output) co std::lock_guard lock(mutex_); Status status; ORT_TRY { - RuntimeWorkloadSnapshot snapshot{session_id_, run_id_, records_, stats_, run_state_}; + RuntimeWorkloadSnapshot snapshot{session_id_, run_id_, records_, stats_, run_state_, limits_}; output = std::move(snapshot); return Status::OK(); } diff --git a/onnxruntime/core/framework/runtime_workload_collector.h b/onnxruntime/core/framework/runtime_workload_collector.h index 20fa33c37cc1f..3372224b17bc5 100644 --- a/onnxruntime/core/framework/runtime_workload_collector.h +++ b/onnxruntime/core/framework/runtime_workload_collector.h @@ -100,6 +100,7 @@ struct RuntimeWorkloadSnapshot { // 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. diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 48c8efdf45308..f242b4c31f247 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -2777,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; @@ -3300,20 +3320,6 @@ common::Status InferenceSession::Initialize() { ort_format_model_mapped_memory_.reset(); } - // Estimate before pruning attributes needed by operator metadata adapters. - const auto compute_report_path = session_options_.config_options.GetConfigOrDefault( - kOrtSessionOptionsComputeEstimateReportPath, ""); - if (!compute_report_path.empty()) { -#if !defined(ORT_MINIMAL_BUILD) - const auto report_status = WriteComputeEstimateReport(graph, compute_report_path); - if (!report_status.IsOK()) { - LOGS(*session_logger_, WARNING) << report_status.ErrorMessage(); - } -#else - LOGS(*session_logger_, WARNING) << "Compute estimate report export is unsupported in a minimal build."; -#endif - } - // once the model is saved, we may remove unnecessary attributes for inference session_state_->PruneRemovableAttributes(); @@ -3726,14 +3732,31 @@ Status InferenceSession::Run(const RunOptions& run_options, RuntimeWorkloadCollector* workload_collector = nullptr; #if !defined(ORT_MINIMAL_BUILD) - if (run_options.runtime_workload_collector && - run_options.runtime_workload_collector->BeginRun(session_id_)) { - workload_collector = run_options.runtime_workload_collector; + 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 diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index f1ca7f3239630..7bdc2635d1132 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -1013,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 index 10850d6a5a9e4..44ebd2f132631 100644 --- a/onnxruntime/test/contrib_ops/compute_estimate_report_test.cc +++ b/onnxruntime/test/contrib_ops/compute_estimate_report_test.cc @@ -4,6 +4,7 @@ #if !defined(ORT_MINIMAL_BUILD) #include +#include #include #include #include @@ -12,6 +13,7 @@ #include #include #include +#include #include #include "gtest/gtest.h" @@ -20,6 +22,7 @@ #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" @@ -51,6 +54,25 @@ class ComputeEstimateReportTest : public ::testing::Test { 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); @@ -85,7 +107,9 @@ class ComputeEstimateReportTest : public ::testing::Test { } Json ReadReport() { - std::ifstream stream(path_); + 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); } @@ -103,6 +127,7 @@ class ComputeEstimateReportTest : public ::testing::Test { 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()) @@ -110,6 +135,21 @@ class ComputeEstimateReportTest : public ::testing::Test { 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_; @@ -339,29 +379,17 @@ TEST_F(ComputeEstimateReportTest, NestedGraphsAreSeparateScopes) { EXPECT_EQ(report["graphs"][i]["summary"]["known_dense_equivalent_ops"], "24576"); EXPECT_EQ(report["graphs"][i]["summary"]["per_run_total_available"], false); } - auto options = ReportOptions(); - InferenceSession session(options, GetEnvironment()); - ASSERT_STATUS_OK(Initialize(session)); - const auto exported = ReadReport(); - EXPECT_EQ(exported["graphs"].size(), 3u); - EXPECT_EQ(exported["summary"]["known_dense_equivalent_ops"], "0"); - RecordProperty("compute_estimate_report", exported.dump(2)); } -TEST_F(ComputeEstimateReportTest, SessionExportsFinalGraphAndRetainsInference) { +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)); - const auto report = ReadReport(); - EXPECT_EQ(report["nodes"].size(), session.GetGraph().NumberOfNodes()); - EXPECT_EQ(report["summary"]["reported_node_count"], 1); - EXPECT_EQ(report["summary"]["known_dense_equivalent_ops"], "24576"); - EXPECT_EQ(report["nodes"][0]["provider"], kCpuExecutionProvider); - EXPECT_EQ(report["nodes"][0]["components"][1]["operation_count"], "192"); - RecordProperty("compute_estimate_report", report.dump(2)); + 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(); @@ -379,6 +407,21 @@ TEST_F(ComputeEstimateReportTest, SessionExportsFinalGraphAndRetainsInference) { 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) { @@ -387,15 +430,20 @@ TEST_F(ComputeEstimateReportTest, DisabledExportDoesNotWrite) { 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, WriteFailureWarnsAndInitializationSucceeds) { +TEST_F(ComputeEstimateReportTest, WriteFailureWarnsAndInferenceSucceeds) { AddMatMul(); ASSERT_STATUS_OK(model_.MainGraph().Resolve()); auto options = ReportOptions(); - // Opening an existing directory as a file fails without requiring permission changes. - const auto invalid_path = ToUTF8String(std::filesystem::temp_directory_path().native()); + // 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(); @@ -406,11 +454,113 @@ TEST_F(ComputeEstimateReportTest, WriteFailureWarnsAndInitializationSucceeds) { 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()); @@ -423,6 +573,154 @@ TEST_F(ComputeEstimateReportTest, Utf8PathAndOverwrite) { 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 diff --git a/onnxruntime/test/framework/runtime_workload_collector_test.cc b/onnxruntime/test/framework/runtime_workload_collector_test.cc index b8642980a4135..82ce681fe35c1 100644 --- a/onnxruntime/test/framework/runtime_workload_collector_test.cc +++ b/onnxruntime/test/framework/runtime_workload_collector_test.cc @@ -13,6 +13,8 @@ #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" @@ -52,6 +54,14 @@ size_t PayloadBytes(const MatMulNBitsRuntimeMetadata& metadata) { (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; @@ -61,6 +71,8 @@ TEST(RuntimeWorkloadCollectorTest, EmptySnapshotHasRunIdentity) { 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) { @@ -496,6 +508,9 @@ TEST(RuntimeWorkloadExecutionTest, TransformerEncoderAndDecoderPropagateCollecto 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}; @@ -732,6 +747,9 @@ TEST(RuntimeWorkloadExecutionTest, FailedRunRetainsEarlierSuccessfulInvocation) 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) { @@ -747,6 +765,12 @@ TEST(RuntimeWorkloadExecutionTest, LoopInvocationsAndLimitsKeepSubgraphIdentity) 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"); @@ -829,6 +853,9 @@ TEST(RuntimeWorkloadExecutionTest, IfBranchesWithSameNodeIndexHaveDistinctGraphI 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"); } } @@ -859,6 +886,7 @@ TEST(RuntimeWorkloadExecutionTest, ScanInvocationsShareIdentityButNotInvocationI 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); From 71263abc9a81f4cfa92093fb71a90e0adf20f957 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Tue, 27 Oct 2026 08:28:08 -0700 Subject: [PATCH 13/14] docs: Update MatMulNBits rollout plan with counter completion details --- .../contrib_op_flop_estimation_rollout.md | 3 + .../matmul_nbits_counter_completion_plan.md | 456 ++++++++++++++++++ 2 files changed, 459 insertions(+) create mode 100644 docs/annotated_partitioning/matmul_nbits_counter_completion_plan.md diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md index 9587f9e11a0e4..ec0bb651fc0b0 100644 --- a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -2,6 +2,9 @@ **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. 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..4377780bb74ce --- /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/contrib_ops/cpu/quantization/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. From 640365fa7cc6385ff5a9a9dc0efcb7aa10c4ff38 Mon Sep 17 00:00:00 2001 From: Tommaso Adani Date: Tue, 27 Oct 2026 17:48:05 +0000 Subject: [PATCH 14/14] refactor: Update MatMulNBits compute estimation files and paths for consistency --- .../contrib_op_flop_estimation_rollout.md | 5 +++-- .../matmul_nbits_counter_completion_plan.md | 2 +- onnxruntime/core/framework/compute_estimate_report.cc | 2 +- .../framework}/matmul_nbits_compute_estimate.cc | 2 +- .../framework}/matmul_nbits_compute_estimate.h | 0 .../test/contrib_ops/matmul_nbits_compute_estimate_test.cc | 2 +- 6 files changed, 7 insertions(+), 6 deletions(-) rename onnxruntime/{contrib_ops/cpu/quantization => core/framework}/matmul_nbits_compute_estimate.cc (98%) rename onnxruntime/{contrib_ops/cpu/quantization => core/framework}/matmul_nbits_compute_estimate.h (100%) diff --git a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md index ec0bb651fc0b0..a38789c5cf3b7 100644 --- a/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md +++ b/docs/annotated_partitioning/contrib_op_flop_estimation_rollout.md @@ -40,7 +40,7 @@ or be treated as zero extra work. Timing and performance-policy changes remain s ### Initial implementation -The internal [MatMulNBits compute estimator](../../onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h) +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 @@ -48,7 +48,8 @@ shapes, not maximum-shape scenarios, and does not validate packed weights or rep 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 CPU contrib source glob, independently of CUDA. +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 diff --git a/docs/annotated_partitioning/matmul_nbits_counter_completion_plan.md b/docs/annotated_partitioning/matmul_nbits_counter_completion_plan.md index 4377780bb74ce..8e09513fdbe52 100644 --- a/docs/annotated_partitioning/matmul_nbits_counter_completion_plan.md +++ b/docs/annotated_partitioning/matmul_nbits_counter_completion_plan.md @@ -54,7 +54,7 @@ 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/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc). +- [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). diff --git a/onnxruntime/core/framework/compute_estimate_report.cc b/onnxruntime/core/framework/compute_estimate_report.cc index ba51bce7c14d4..7bbbf5b2c6177 100644 --- a/onnxruntime/core/framework/compute_estimate_report.cc +++ b/onnxruntime/core/framework/compute_estimate_report.cc @@ -24,7 +24,7 @@ #include "core/platform/env.h" #if !defined(DISABLE_CONTRIB_OPS) -#include "contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h" +#include "core/framework/matmul_nbits_compute_estimate.h" #endif namespace onnxruntime { diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc b/onnxruntime/core/framework/matmul_nbits_compute_estimate.cc similarity index 98% rename from onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc rename to onnxruntime/core/framework/matmul_nbits_compute_estimate.cc index 3f2b9425b5e7f..e3186b7c83a85 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.cc +++ b/onnxruntime/core/framework/matmul_nbits_compute_estimate.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h" +#include "core/framework/matmul_nbits_compute_estimate.h" #include #include diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h b/onnxruntime/core/framework/matmul_nbits_compute_estimate.h similarity index 100% rename from onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h rename to onnxruntime/core/framework/matmul_nbits_compute_estimate.h diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc index ea49b5ada0044..9bdb2ba7cf8fd 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_compute_estimate_test.cc @@ -10,7 +10,7 @@ #include "gtest/gtest.h" -#include "contrib_ops/cpu/quantization/matmul_nbits_compute_estimate.h" +#include "core/framework/matmul_nbits_compute_estimate.h" #if !defined(ORT_MINIMAL_BUILD) #include "core/graph/constants.h"