Unify shape inference and enable executable WebNN reloads - #238
FelixKrall wants to merge 12 commits into
Conversation
## Feature additions - Serialize complete operation arguments and options into .webnn artifacts, including reshape dimensions, slice parameters, concat axes, permutations, and recurrent settings. - Serialize operand-valued options by stable exported name and resolve them back to graph-local operand IDs during loading. - Infer loaded Squeeze, Unsqueeze, Split, GRU, LSTM, and LSTMCell result shapes while tracking inferred state separately from valid scalar shapes. - Add executable save/reload/dispatch coverage for Slice/Concat/Reshape/Squeeze/Unsqueeze, Gemm bias references, and uneven multi-output Split graphs. ## Bugfixes - Preserve every output shape and data type for multi-output operations instead of updating and testing only the first result. - Reject unresolved required intermediate and output shapes rather than allowing placeholder empty shapes to be interpreted as scalars. - Correct Unsqueeze axis validation so an axis equal to the output rank is rejected. ## Refactors - Apply current rustfmt layout to the existing rustnn_build_graph entry point used to compile loaded GraphInfo artifacts. ## Behavioral impact and compatibility - Newly saved .webnn artifacts retain the method-level and operand-reference data required for executable reload; numeric operand options in existing artifacts remain accepted, while new artifacts use stable names. - Loading now returns a conversion error when a required operand shape cannot be inferred. Valid scalar operands remain supported through explicit inference-state tracking. - These changes provide the executable reload path consumed by onnx2webnn commit d721da0 for cache-backed Tiny RoFormer validation; no new RustNN public API is introduced by this commit. ## Validation - cargo fmt --all -- --check passed. - PROTOC=/home/fkrall/vscprojects/transformers-convert/tools/protoc/bin/protoc ORT_DYLIB_PATH=/home/fkrall/vscprojects/transformers-convert/tools/onnxruntime/onnxruntime-linux-x64-1.29.0/lib/libonnxruntime.so.1.29.0 cargo test --lib --features onnx-runtime passed 347 tests with 1 ignored. - The focused saved Split reload/dispatch regression passed. - The dependent onnx2webnn suite passed all 704 tests, including Tiny RoFormer cache export, reload, dispatch, and native ORT comparison.
## Feature additions - Add and publicly export run_onnx_path_with_inputs so ONNX Runtime can resolve external-data sidecars relative to the model file. ## Bugfixes - None. ## Refactors - Share input dispatch, descriptor validation, and output collection between memory-backed and path-backed ONNX Runtime sessions. ## Behavioral impact and compatibility - The new API is additive and available behind the existing onnx-runtime feature. - Existing in-memory execution APIs remain unchanged. - Path-backed execution uses the existing CPU execution provider and disabled graph optimizations. ## Validation - make fmt passed. - PROTOC=... ORT_DYLIB_PATH=... cargo test --lib --features onnx-runtime passed 347 tests with 1 ignored. - The dependent onnx2webnn suite passed, including path-based external-data execution. - make test did not complete because clippy -D warnings reports 56 existing warnings in untouched RustNN code, primarily coreml_mlprogram.rs. - git diff --cached --check passed.
## Feature additions - Document the required `protoc` compiler and optional backend-specific development dependencies. ## Bugfixes - None. ## Refactors - None. ## Behavioral impact and compatibility - Documentation only; no runtime or API behavior changes. ## Validation - `git diff --check` passed. - Tests not run because no code changed.
## Feature additions - Add an internal `GraphRecorder` that owns in-progress `GraphInfo` construction, stable operand allocation, constant storage, output registration, and final structural validation. - Add focused coverage for scalar and dynamic descriptors, atomic failures, forward references, output-count mismatches, all multi-output inference families, and builder/loader equivalence. ## Bugfixes - Infer operation outputs before inserting them, preventing failed inference from leaving partial operations or empty-shape placeholders. - Preserve serialized dynamic dimension names and bounds during `.webnn` loading. - Reject missing and forward operand references deterministically with operation context. ## Refactors - Route `MLGraphBuilder` inputs, constants, operations, and outputs through `GraphRecorder` without changing its public API. - Share one canonical output-descriptor inference entry point between `MLGraphBuilder` and GraphJSON loading. - Remove the loader-specific ten-pass inference loop, `known_shapes` tracking, default `Float32` intermediates, placeholder shapes, and shape-repair heuristics. - Update loader fixtures to construct structurally valid pass-through graphs using explicit identity operations. ## Behavioral impact and compatibility - Loading now processes GraphJSON nodes once in dependency order and either produces fully inferred operands or returns an error. - This intentionally breaks compatibility with experimental `.webnn`/JSON files that expose an input or constant directly as a graph output. Such files must insert an identity operation before the output. - Existing operation-produced `.webnn` outputs and onnx2webnn-generated models remain supported. - Public `MLGraphBuilder` signatures and onnx2webnn integration remain unchanged. - This removes empty-shape ambiguity from recorder and loader operation construction, but does not claim to remove every legacy interpretation of empty shapes elsewhere in RustNN. ## Validation - `cargo test --lib`: 361 passed. - `cargo test --lib --no-default-features`: passed. - `cargo check --no-default-features --features onnx-runtime`: passed. - `cargo test --lib --features cann-runtime-mock`: passed. - `cargo test --lib --features dynamic-inputs`: passed. - onnx2webnn `cargo test --release --locked`: passed. - Unchanged onnx2webnn skeleton sweep: 47 passed, 0 failed, 5 heavy entries skipped. - Tiny RoFormer numerical validation passed with both real and generated weights. - `cargo fmt --all -- --check` and `git diff --check`: passed. - `make test` remains blocked by 54 pre-existing Clippy errors in untouched RustNN code; no findings originated in this change. - WPT was unavailable because the installed Node runtime lacks `Float16Array`.
## Feature additions - None. ## Bugfixes - Treat empty operand shapes as rank-0 scalars during quantize/dequantize validation and GRU-cell inference instead of interpreting them as unresolved metadata. - Preserve scalar operands in ONNX converter shape tracking and return contextual errors when required operand descriptors are missing. - Use canonical Gather shape inference so scalar indices preserve data dimensions preceding the gather axis. - Always serialize required Reshape, Expand, and Tile shape arguments, including empty scalar targets. - Reject missing or malformed required Reshape, Expand, Tile, Pad, Slice, and Constant shape arguments instead of silently defaulting them to empty vectors. ## Refactors - Establish that every completed `OperandDescriptor` has a known shape: `[]` is scalar, nonempty shapes may contain bounded dynamic dimensions, and unresolved inference never enters `GraphInfo`. - Represent potentially absent shape-valued operation arguments with `Option`, distinguishing `None` from an explicitly supplied empty vector. - Seed ONNX conversion shape maps from every graph operand and replace missing-operand shape fallbacks with explicit failures. - Update historical loader documentation to remove obsolete empty-shape inference heuristics. - Add regression coverage for scalar descriptors, constants, GRU cells, quantization, JSON shape arguments, ONNX shape maps, and scalar Reshape conversion. ## Behavioral impact and compatibility - Serialized `OperandDescriptor` values must now include a `shape` field; `"shape": []` remains the valid scalar representation. - Experimental `.webnn` files that omitted required empty `newShape`, `repetitions`, padding, Slice, or Constant shape arguments are rejected and must be regenerated. - Scalar GRU hidden states are now correctly rejected as rank 0 rather than receiving an inferred rank-2 fallback. - Static shapes and bounded `Dimension::Dynamic` descriptors remain supported and retain their existing representation. - The public graph-builder API is unchanged. ## Validation - `cargo fmt --all` passed. - `cargo test --lib` passed: 369 tests. - `cargo test --lib --features dynamic-inputs` passed: 376 tests. - `cargo test --lib --features cann-runtime-mock` passed: 387 tests. - `cargo check --all-targets`, `cargo check --no-default-features`, and `cargo check --features onnx-runtime` passed. - `cargo clippy --all-targets -- -D warnings` reached Clippy but failed on the same 54 pre-existing warnings, primarily in the CoreML converter. - onnx2webnn `cargo test --release --locked` passed against the modified RustNN checkout. - The cached manifest skeleton sweep passed 47 of 47 selected models; five designated heavy cases were skipped. - Tiny RoFormer numerical validation passed with both real and generated weights after WebNN/Safetensors export and reload. - WPT execution was unavailable because the installed Node runtime does not provide `Float16Array`. - `git diff --check` passed.
## Feature additions - Add a versioned `rustnn.webnn.packed4=1` Safetensors extension that stores logical Int4 and Uint4 constants as their original packed bytes in U8 tensors. - Teach the `.webnn` loader to restore packed constants by stable weight reference before resolving ordinary external tensors. - Add mixed Int4/Uint4 save, reload, compilation, and execution coverage, including odd logical element counts. ## Bugfixes - Allow `.webnn` graphs containing packed 4-bit constants to be exported instead of rejecting types unsupported natively by Safetensors. - Validate packed element counts, byte lengths, storage dtype and shape, format version, missing weights, and ambiguous sanitized names with explicit errors. ## Refactors - None. ## Behavioral impact and compatibility - The `.webnn` declaration remains authoritative for each constant’s logical 4-bit dtype and shape; Safetensors contains one U8 tensor of `ceil(elements / 2)` packed bytes. - Ordinary Safetensors constants retain their existing representation and loading path. - External Int4/Uint4 archives must carry the supported RustNN metadata marker. Missing, unknown, or malformed packed-format metadata is now rejected explicitly. - Packed artifacts require a RustNN version containing this extension; the experimental cache format is intentionally versioned. ## Validation - `PROTOC=... LD_LIBRARY_PATH=... make build` passed using RustNN’s repository-local protoc installation. - `make fmt-check` and `git diff --check` passed. - Downstream onnx2webnn `make test` passed, including packed Uint4 conversion, export, reload, execution, and native-ORT comparison. - `make test` reached Clippy but stopped before RustNN’s test phase on 54 pre-existing `-D warnings` failures in unrelated CoreML and graph code; none were reported in the changed files.
## Feature additions - None. ## Bugfixes - Convert WebNN Slice sizes to ONNX ends as `start + extent`, independent of stride, instead of incorrectly multiplying the extent by the stride. - Apply the same extent semantics when generating LiteRT `STRIDED_SLICE` end tensors. ## Refactors - Extract LiteRT Slice end calculation into a focused helper shared by conversion and unit coverage. - Add ONNX and LiteRT regression tests proving that start 1 with extent 8 and stride 2 produces end 9. ## Behavioral impact and compatibility - Strided Slice graphs now select the WebNN-specified range on ONNX Runtime and LiteRT backends instead of extending the range by the stride twice. - This changes generated backend graphs for affected positive-stride Slice operations but does not change `GraphInfo`, the public builder API, or serialized `.webnn` schemas. - Existing `.webnn` graphs already store extents and require no format migration; they receive corrected execution after rebuilding with this RustNN version. ## Validation - `cargo test --lib` passed all 379 default RustNN library tests. - `cargo test strided_slice_uses_extent_for_onnx_end --lib -- --nocapture` passed. - `cargo check --no-default-features`, `cargo check --features dynamic-inputs`, and `cargo check --features onnx-runtime` passed using the repository-local protobuf toolchain. - The onnx2webnn positive-stride Slice export/reload/execution regression and complete 52-case real-weight sweep exercised the corrected ONNX backend successfully. - The LiteRT-specific unit test was added, but its feature build could not run because `flatc` is not installed in the environment. - `cargo clippy --all-targets -- -D warnings` was attempted and remains blocked by 54 pre-existing warnings, predominantly in `coreml_mlprogram.rs`; none are in this change. - `cargo fmt --all -- --check` and `git diff --check` passed.
| } | ||
|
|
||
| impl GraphRecorder { | ||
| pub(crate) fn new() -> Self { |
There was a problem hiding this comment.
to me it seems as if all those should be methods on GraphInfo and some of them slowly replacing the public access to fields of GraphInfo
There was a problem hiding this comment.
for me fine, to have this struct first. I can compare with the changes I made to GraphInfo. But long-term a manageable interface like this should replace direct data access GraphInfo
There was a problem hiding this comment.
Yes I agree, this is a first attempt to guide the project into this direction. Its motivation was rather to unify shape inference in .webnn loading and the Graph construction path that is used when using the builder API or the .webnn loader.
This does not make GraphInfo crate private yet, but it creates an implementation entry point to adopt when doing this in the future.
| pub data_type: String, | ||
| #[serde(default)] | ||
| pub shape: Vec<u32>, | ||
| pub shape: Option<Vec<u32>>, |
There was a problem hiding this comment.
what is the meaning of shape being None?
There was a problem hiding this comment.
None means the required constant shape was omitted during parsing or construction; it is not an unknown tensor shape. We need the distinction because Some(vec![]) is a valid rank-0 scalar. constant_shape() rejects None before the operation enters GraphInfo, so every completed operand still has a known shape.
mklimenko-nv
left a comment
There was a problem hiding this comment.
Aside from those comments, the upstream documentation has no page that describes the .webnn or weights format at all. Could you make a pass on the docs for your changes as well?
| name.replace("::", "__").replace('.', "_") | ||
| } | ||
|
|
||
| fn discover_external_weights(graph_path: &Path) -> Option<PathBuf> { |
There was a problem hiding this comment.
This is a copy-paste from webnn-graph. Can we avoid it and come to a single source of truth?
| return Ok(()); | ||
| } | ||
|
|
||
| let bytes = fs::read(&weights_path).map_err(|source| GraphError::io(&weights_path, source))?; |
There was a problem hiding this comment.
Here and below: the whole archive is read (potentially several gigabytes), the header is parsed, then everything is copied into InlineBytes at line 252 and the buffer is dropped.
| }; | ||
|
|
||
| resolve_packed_4bit_weights(&mut graph_json, path_ref)?; | ||
| webnn_graph::external_weights::resolve_external_weights(&mut graph_json, path_ref, None, None) |
There was a problem hiding this comment.
webnn_graph resolver discovers the file again and reads the full buffer again for the remaining tensors.
|
@mklimenko-nv thanks for catching the memory and code duplication issues. I was not aware of the webnn-graph ownership of this. I'll fix it ASAP and also include the docs you are requesting. |
## Feature additions - Add `from_graph_json_owned` so filesystem loading can move resolved inline constant buffers into `GraphInfo` without cloning them; retain the borrowed `from_graph_json` API as a compatible wrapper. - Document that `MLConstantOptions::shape = None` means the required shape was omitted, while `Some(vec![])` represents a scalar. ## Bugfixes - Pin RustNN to webnn-graph commit `32a685c0cc006e8b1d5c3e2a8499fc173cb5997c`, ensuring the shared packed-4-bit resolver, writer, and metadata constants are available without a local Cargo patch. ## Refactors - Replace RustNN’s duplicate SafeTensors discovery, packed-4-bit validation, and archive-writing implementation with the shared webnn-graph resolver and writer. - Resolve every external-weight archive once, move resolved buffers through the owned conversion path, and keep RustNN responsible only for graph semantics and backend execution. - Move RustNN’s direct SafeTensors dependency to dev-dependencies and retain one mixed ordinary/Int4/Uint4 export, reload, compile, and execution round trip. ## Behavioral impact and compatibility - Existing `.webnn` plus SafeTensors artifacts using `rustnn.webnn.packed4=1` remain supported with the same low-nibble-first layout. - The borrowed GraphJson conversion API remains available; filesystem loading now avoids a redundant copy of every resolved constant. - RustNN temporarily depends on the exact prerequisite commit from `FelixKrall/webnn-graph`; after that PR merges, this must be changed back to the merged `rustnn/webnn-graph` revision and the lockfile refreshed. ## Validation - `cargo test --lib` — 378 tests passed. - `cargo test --lib --no-default-features` — 378 tests passed. - `cargo test --lib loader::tests` — 26 focused loader tests passed against the pinned remote dependency. - `cargo test --lib webnn_json::tests::owned_graph_json_moves_inline_constant_bytes` — passed. - `cargo check --all-targets --no-default-features` — passed. - `cargo clippy --all-targets --no-default-features -- -D warnings` — passed. - The mixed packed-4-bit save, reload, compile, and ORT execution test passed with the ONNX Runtime feature enabled. - Formatting, dynamic-input checks, backend documentation drift checks, and `git diff --check` passed.
|
I have addressed your comments as such @mklimenko-nv.
I have also added docs on webnn-graph usage here and about its current conventions in the webnn-graph PR. I removed the stale agent like docs from the webnn-graph repo as well. |
## Feature additions - Add a user guide for saving, loading, compiling, executing, and distributing RustNN `.webnn` graphs with external weight sidecars. - Register the guide in the MkDocs User Guide navigation. - Link to `webnn-graph` as the canonical specification for graph and external-weight formats. ## Bugfixes - None. ## Refactors - None. ## Behavioral impact and compatibility - No RustNN API, graph-loading, serialization, backend, or execution behavior changes are intended. - The new page documents existing integration behavior without duplicating the canonical file-format specification. ## Validation - `mkdocs build --strict` passed using the dependencies declared in `docs/requirements.txt`. - Documentation line-length checks and `git diff --check` passed.
## Feature additions - Add a backend-independent Cargo integration test covering empty `expand`, `reshape`, `slice`, and `tile` arguments, nonempty operand-array routing, and missing required-argument errors. ## Bugfixes - Require operand arrays to be nonempty before routing them as positional operands, preserving explicit empty arrays for `OperationExtras`. - Upload WPT JSON and HTML artifacts from the paths produced by `test-wpt-report`, removing invalid backend-specific and CoreML-to-LiteRT path mappings. ## Refactors - None. ## Behavioral impact and compatibility - WPT graphs can now explicitly supply empty `newShape`, `starts`, `sizes`, and `repetitions` arrays while omitted required arguments remain errors. - Nonempty operand arrays retain their existing positional routing. - CI artifact names remain backend-specific, while isolated matrix jobs upload their common report paths. - No public API, serialized-format, backend-converter, or expected-failure-list changes. ## Validation - `cargo fmt --check` passed. - `cargo test --lib` passed all 378 tests. - `cargo test --test test_wpt_argument_routing` passed all 3 tests. - Focused ONNX WPT suites for expand, reshape, slice, and tile passed, including all five scalar regressions. - Full ONNX WPT passed all 2,521 cases. - Full LiteRT WPT reported 792 passed, 1,729 existing skips or expected failures, and 0 unexpected failures. - Workflow YAML parsing and `git diff --check` passed. - CoreML was not run locally because the verification host is Linux.
Summary
RustNN changes to enable full roundtrip numeric validation of onnx2webnn conversion. Tied to onnx2webnn #5 and dependend on webnn-graph #19. The reference in this PR needs to updated after the webnn-graph was merged.
Make serialized WebNN graphs independently reloadable and executable while unifying graph recording and shape inference between
MLGraphBuilderand the.webnnloader.Completed graphs now use an unambiguous shape model:
[]always means a known rank-zero scalar, bounded dynamic dimensions remain explicit, and unresolved descriptors exist only as temporary internal inference state.Feature additions
MLContext::rustnn_build_graphas a direct compilation entry point for deserializedGraphInfo. This avoids constructing a second, unusedGraphRecorderwhen compiling a graph reconstructed by the.webnnloader.run_onnx_path_with_inputsso native ONNX Runtime execution can resolve external-data sidecars relative to the model file. This supports reference validation of filesystem-backed models whose weights are not embedded in the ONNX protobuf.MLGraphBuilder.Int4andUint4dtype and shape remain in.webnn.U8.Bugfixes
ReshapeandExpandtargets asnewShape: [].start + extent; this fixes both the ONNX and LiteRT lowering paths.Refactors
GraphRecorderused by both graph construction and JSON loading.Optionor missing internal entries rather than publicOperandDescriptorvalues.Behavioral impact and compatibility
GraphRecorderremains crate-private and is not a new public API.MLContext::rustnn_build_graphandrun_onnx_path_with_inputsare additive public entry points.OperandDescriptorin completedGraphInfonow has a known shape:[]is scalar.Dimension::Dynamicentries..webnnformat has intentional compatibility breaks:Validation
cargo test --lib: 379 tests passed.dynamic-inputs, ONNX Runtime, and available backend-mock checks passed across the refactor.Int4andUint4constants.flatcwas unavailable.-D warningsfailures outside this change.