diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index aff5b9e..b50ea4b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -12,7 +12,7 @@ env: jobs: test: - name: Test (${{ matrix.os }}, minimal=${{ matrix.minimal }}) + name: MSRV (${{ matrix.os }}, minimal=${{ matrix.minimal }}) runs-on: ${{ matrix.os }} strategy: fail-fast: false @@ -27,7 +27,7 @@ jobs: uses: actions/checkout@v4 - name: Install Rust - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@1.97.0 with: components: rustfmt, clippy @@ -89,3 +89,59 @@ jobs: - name: Build release run: make release + + stable: + name: Current stable + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install Rust + uses: dtolnay/rust-toolchain@stable + with: + components: rustfmt, clippy + + - name: Install protoc + run: | + sudo apt-get update + sudo apt-get install -y protobuf-compiler + + - name: Cache cargo registry + uses: actions/cache@v4 + with: + path: ~/.cargo/registry + key: ubuntu-stable-cargo-registry-${{ hashFiles('**/Cargo.lock') }} + restore-keys: ubuntu-stable-cargo-registry- + + - name: Cache cargo index + uses: actions/cache@v4 + with: + path: ~/.cargo/git + key: ubuntu-stable-cargo-git-${{ hashFiles('**/Cargo.lock') }} + restore-keys: ubuntu-stable-cargo-git- + + - name: Cache target directory + uses: actions/cache@v4 + with: + path: target + key: ubuntu-stable-target-${{ hashFiles('**/Cargo.lock') }} + restore-keys: ubuntu-stable-target- + + - name: Cache ResNet-50 ONNX model + uses: actions/cache@v4 + with: + path: target/test-data + key: resnet50-onnx-opset16-v1 + + - name: Check formatting + run: make fmt-check + + - name: Run clippy + run: make lint + + - name: Run tests + run: make test + + - name: Build release + run: make release diff --git a/AGENTS.md b/AGENTS.md deleted file mode 100644 index 2a02ea8..0000000 --- a/AGENTS.md +++ /dev/null @@ -1,416 +0,0 @@ -# CLAUDE.md - -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. - -Please keep prose line lengths at or below 120 characters when editing markdown/text files. - -## Project Overview - -webnn-graph is a Rust implementation for a WebNN-oriented graph DSL. It provides a complete pipeline: -1. **Convert ONNX models** to WebNN format (NEW!) -2. Parse WebNN graph text files (.webnn) into JSON representation -3. Serialize JSON back to WebNN text format (full round-trip support) -4. Validate graph structure and weights manifests -5. Emit WebNN JavaScript builder code - -## Build and Test Commands - -Build the project: -```bash -make build -``` - -Run the CLI tool: -```bash -make run -``` - -Run tests: -```bash -make test -``` - -Run a specific test: -```bash -cargo test -``` - -Format code: -```bash -make fmt -``` - -Check code formatting: -```bash -make fmt-check -``` - -Run clippy linter: -```bash -make lint -``` - -Quick compile check: -```bash -make check -``` - -Clean build artifacts: -```bash -make clean -``` - -View all available commands: -```bash -make help -``` - -## CLI Usage - -The binary is named `webnn-graph` with ten subcommands. **Most commands accept both .webnn and .json -formats** (auto-detected). - -### Graph Operations - -Validate graph structure (accepts .webnn or .json): -```bash -make validate -# Or directly: -webnn-graph validate examples/resnet_head.webnn -webnn-graph validate graph.json --weights-manifest examples/weights.manifest.json -``` - -Emit JavaScript builder code (accepts .webnn or .json): -```bash -make emit-js -# Or directly: -webnn-graph emit-js examples/resnet_head.webnn > buildGraph.js -webnn-graph emit-js graph.json > buildGraph.js # Also works -``` - -Generate interactive HTML visualizer (accepts .webnn or .json): -```bash -make emit-html -# Or directly: -webnn-graph emit-html examples/resnet_head.webnn > visualizer.html -open visualizer.html - -# Features: -# - Interactive graph layout with zoom/pan -# - Node details sidebar on click -# - Export to PNG/SVG -# - Light/dark theme toggle -# - Works completely offline (no network requests) -``` - -Parse graph text to JSON (explicit conversion): -```bash -webnn-graph parse examples/resnet_head.webnn > graph.json -``` - -Serialize JSON back to WebNN text format (explicit conversion): -```bash -webnn-graph serialize graph.json > model.webnn - -# Complete round-trip: -webnn-graph parse model.webnn | webnn-graph serialize /dev/stdin > model_copy.webnn -``` - -### ONNX Conversion - -**Built-in Constant Folding** - -The converter includes built-in constant folding (enabled with `--optimize`) that automatically handles dynamic -shape patterns. WebNN does not support dynamic shapes at runtime, so the converter resolves all dynamic dimensions -at conversion time. - -**Why this is necessary**: WebNN's `reshape` operation requires the shape parameter to be a constant, not a -dynamically computed value. ONNX models with dynamic shapes use `Shape→Gather→Concat→Reshape` patterns that -must be resolved to static constants at conversion time. - -Convert ONNX models to WebNN format with constant folding: -```bash -# Basic conversion with optimization (recommended) -webnn-graph convert-onnx --input model.onnx --optimize \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 - -# Output: model.webnn + model.weights + model.manifest.json - -# Custom output paths -webnn-graph convert-onnx \ - --input model.onnx \ - --optimize \ - --output graph.webnn \ - --weights graph.weights \ - --manifest graph.manifest.json \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 - -# Inline weights for small models -webnn-graph convert-onnx --input model.onnx --optimize --inline-weights \ - --override-dim batch_size=1 - -# Output to JSON format -webnn-graph convert-onnx --input model.onnx --optimize --output model.json \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 -``` - -The `--optimize` flag performs constant folding, which: -- Evaluates Shape, Gather, Concat operations at conversion time -- Eliminates dynamic shape computation patterns -- Reduces model size by 40-50% for transformer models -- Makes all reshape operations use static constants - -**Supported operators** (NLP/Transformer focused): -- **MatMul, Gemm**: Matrix multiplication with options -- **Add, Sub, Mul, Div, Pow**: Element-wise operations -- **LayerNormalization, Softmax**: Normalization operations -- **Reshape, Transpose, Concat, Split**: Tensor manipulation -- **Constant folding**: Shape, Gather, Concat, Unsqueeze, Squeeze, Cast, Constant - -**Full pipeline example**: -```bash -# Step 1: Convert ONNX → WebNN with constant folding -webnn-graph convert-onnx --input bert-base.onnx --optimize \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 - -# Step 2: Generate JavaScript -webnn-graph emit-js bert-base.webnn > buildGraph.js - -# Example results for BERT models with --optimize: -# - Original: 637 nodes with Shape operations -# - After constant folding: 317 nodes (50% reduction), no Shape operations -# - All reshape operations use static constants -``` - -**See also**: [Dynamic Dimensions Guide](docs/dynamic-dimensions-guide.md) for help choosing dimension override values. - -### Weights Management - -Pack tensor files into binary .weights file: -```bash -webnn-graph pack-weights \ - --manifest weights.manifest.json \ - --input-dir ./tensors/ \ - --output model.weights -``` - -Unpack binary weights for inspection: -```bash -webnn-graph unpack-weights \ - --weights model.weights \ - --manifest weights.manifest.json \ - --output-dir ./extracted/ -``` - -Create manifest from tensor directory: -```bash -webnn-graph create-manifest \ - --input-dir ./tensors/ \ - --output weights.manifest.json \ - --endianness little -``` - -## Complete Workflow Example - -The DSL uses **.webnn as the primary format** (10x smaller than JSON). The DSL is designed for **complete -separation of concerns**: graph structure is reusable, weights are external, and input data is provided at -runtime. - -### 1. Define Graph (Data-Agnostic) - -Create `model.webnn` (primary format): -```webnn -webnn_graph "resnet_head" v1 { - inputs { - x: f32[1, 2048]; // Shape only, no actual data - } - consts { - W: f32[2048, 1000] @weights("W"); // Reference to external weights - b: f32[1000] @weights("b"); - } - nodes { - logits0 = matmul(x, W); - logits = add(logits0, b); - probs = softmax(logits, axis=1); - } - outputs { probs; } -} -``` - -### 2. Prepare Weights - -Pack your trained weights: -```bash -# Assuming you have W.bin and b.bin tensor files with W.meta.json and b.meta.json -webnn-graph pack-weights \ - --manifest weights.manifest.json \ - --input-dir ./trained_weights/ \ - --output model.weights -``` - -### 3. Generate JavaScript - -Build the runtime code directly from .webnn: -```bash -webnn-graph emit-js model.webnn > buildGraph.js -``` - -This generates both the `WeightsFile` helper class and the `buildGraph()` function. - -### 4. Runtime Usage (Browser/Node.js) - -Use the graph with different input data: -```javascript -import { WeightsFile, buildGraph } from './buildGraph.js'; - -// One-time setup: load graph structure + weights -const weights = await WeightsFile.load('model.weights', 'weights.manifest.json'); -const context = await navigator.ml.createContext(); -const graph = await buildGraph(context, weights); - -// Reuse graph with different input data (many times!) -const input1 = new Float32Array(2048).fill(1.0); -const result1 = await context.compute(graph, { x: input1 }); - -const input2 = new Float32Array(2048).fill(2.0); -const result2 = await context.compute(graph, { x: input2 }); - -// Same graph, different data - no rebuilding needed! -``` - -## Architecture - -### Module Structure - -- **ast.rs**: Core data structures for the graph JSON format - - `GraphJson`: Top-level structure containing inputs, consts, nodes, outputs, and optional quantized flag - - `OperandDesc`: Describes tensor shape and data type - - `DataType`: Enum for f32, f16, i4, u4, i32, u32, i64, u64, i8, u8 - - `ConstDecl`: Constant declarations with initialization (weights, scalar, or inline bytes) - - `Node`: Represents operations with inputs, options, and outputs - -- **parser.rs**: Pest-based parser for WebNN text format - - Uses `wg.pest` grammar file - - `parse_wg_text()`: Main entry point converting WebNN text to `GraphJson` - - Handles four main blocks: inputs, consts, nodes, outputs - - Extracts and stores graph name from header - -- **serialize.rs**: WebNN text format serializer - - `serialize_graph_to_wg_text()`: Main entry point converting `GraphJson` to WebNN text - - Generates properly formatted .webnn files with 2-space indentation - - Handles all const initializations (weights, scalar, inline bytes) - - Supports multi-output nodes and options serialization - - `SerializeError`: Error type for serialization failures - -- **validate.rs**: Graph validation logic - - `validate_graph()`: Checks format version, output presence, reference validity - - `validate_weights()`: Validates weights manifest against const declarations - -- **weights.rs**: Weights manifest handling - - `WeightsManifest`: External weights file structure - - `TensorEntry`: Individual tensor metadata (dataType, shape, byteOffset, byteLength) - - Helper functions: `dtype_size()`, `numel()` - -- **weights_io.rs**: Binary weights packing/unpacking - - `pack_weights()`: Combine tensor files into binary .weights format - - `unpack_weights()`: Extract tensors from binary weights file - - `create_manifest()`: Generate manifest from tensor directory - - Binary format: 4-byte magic "WGWT", 4-byte version, concatenated tensor data - -- **emit_js.rs**: JavaScript code generation - - `emit_weights_loader_js()`: Generates WeightsFile helper class - - `emit_builder_js()`: Generates WebNN MLGraphBuilder code from `GraphJson` - - Outputs complete module with loading, validation, and graph building - -- **main.rs**: CLI interface using clap - - Defines six subcommands: Parse, Validate, EmitJs, PackWeights, UnpackWeights, CreateManifest - -### WebNN Graph Language Format - -Graph files have a structured format with four main sections: - -``` -webnn_graph "name" v1 { - inputs { - x: f32[1, 2048]; - } - consts { - W: f32[2048, 1000] @weights("W"); - } - nodes { - result = operation(input1, input2, option=value); - } - outputs { result; } -} -``` - -**Note**: For quantized graphs, add the `@quantized` annotation in the header: -``` -webnn_graph "name" v1 @quantized { - ... -} -``` - -This sets the `quantized` field to `true` in the JSON representation, indicating the graph contains quantized -weights or operations using Int4/Uint4 data types. - -### Data Flow - -**Primary workflow** (.webnn is 10x smaller than JSON): -1. Author: Write `.webnn` text files (human-readable, compact) -2. Validate: `webnn-graph validate model.webnn` (parses internally) -3. Emit: `webnn-graph emit-js model.webnn` (generates JavaScript) - -**Optional conversions**: -- Parse: `.webnn` → `.json` (for programmatic manipulation) -- Serialize: `.json` → `.webnn` (for human editing) - -The tool supports full round-tripping: text ↔ JSON, preserving semantics. - -### Key Invariants - -- All node inputs must reference previously defined inputs, consts, or node outputs -- Graph outputs must reference valid node results -- Weights manifest entries must match const declarations in type and shape -- Node IDs must be unique -- Graph name is optional in JSON but will default to "graph" when serializing - -## Test Coverage - -The project has comprehensive test coverage across all modules (50 tests total): - -- **ast.rs**: Tests for DataType conversion, OperandDesc equality, ConstInit variants -- **weights.rs**: Tests for dtype_size, numel, and large value handling -- **weights_io.rs**: Tests for pack/unpack roundtrip, validation, and error handling -- **parser.rs**: Tests for parsing inputs, consts, nodes, outputs, and error handling -- **serialize.rs**: Tests for serialization, round-trip preservation, string escaping, all data types -- **validate.rs**: Tests for graph validation and weights validation -- **emit_js.rs**: Tests for JavaScript code generation with various graph configurations - -Run tests with: -```bash -make test -``` - -Generate coverage report (requires cargo-tarpaulin): -```bash -make coverage -``` - -Install development dependencies: -```bash -make dev-deps -``` - -## Development Notes - -- This is a reference scaffold - op semantics should be extended as needed -- The parser uses Pest grammar (src/wg.pest) for parsing WebNN text (.webnn files) -- JSON format uses camelCase for WebNN API compatibility (dataType, byteOffset, etc.) -- Const initializations support three modes: weights references, scalar values, inline bytes -- All code should be formatted with `make fmt` before committing -- Run `make lint` to check for common issues diff --git a/CLAUDE.md b/CLAUDE.md deleted file mode 100644 index 43c994c..0000000 --- a/CLAUDE.md +++ /dev/null @@ -1 +0,0 @@ -@AGENTS.md diff --git a/Cargo.toml b/Cargo.toml index c51859d..6c637ca 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,6 +2,7 @@ name = "webnn-graph" version = "0.3.0" edition = "2021" +rust-version = "1.97" license = "Apache-2.0" description = "Simple DSL for WebNN graphs" diff --git a/README.md b/README.md index b0e0a3b..1f1525b 100644 --- a/README.md +++ b/README.md @@ -1,448 +1,135 @@ # webnn-graph -`webnn-graph` is a small Rust library and CLI that defines a **WebNN-oriented -graph DSL**, parses it into a minimal AST, and enables multiple downstream uses -such as graph validation, serialization, and WebNN graph construction. +`webnn-graph` is a Rust library and command-line tool for a WebNN-oriented graph DSL. It parses and +serializes `.webnn` files, validates graph structure, manages external weights, emits JavaScript and +interactive HTML, and optionally converts ONNX models into the DSL. -The goal is to keep the language surface **very close to WebNN itself**, while -allowing graphs to be expressed declaratively and reused across tooling. +The browser-based graph visualizer is published at +[rustnn.github.io/webnn-graph](https://rustnn.github.io/webnn-graph/). -The project also implements a Netron-like **WebNN graph visualizer** that -allows for interactive exploration of graph structure. +## File model -Check it out at [https://rustnn.github.io/webnn-graph/](https://rustnn.github.io/webnn-graph/) +A model consists of a graph and, when constants are external, a weight sidecar: -## Development - -### Git Hooks (Recommended) - -Enable pre-commit hooks to run clippy locally: - -```bash -./scripts/setup-githooks.sh -``` - -## Conceptual Model - -A WebNN graph defined with this project is split across **three distinct files**, each with a single responsibility. - -### 1. Graph definition (`.webnn`) - -The `.webnn` file describes **only the structure of the graph**: - -- Inputs and their types -- Constants and their shapes -- Operator calls and their wiring -- Named outputs - -It contains **no actual tensor data**. - -This file is intended to be: -- Small -- Human-readable -- Easy to diff and review -- Stable across weight updates - -Its EBNF-like grammar: - -``` -File ::= Header Block* EOF - -Header ::= "webnn_graph" String "v" Int ("@quantized")? "{" - -Block ::= InputsBlock - | ConstsBlock - | NodesBlock - | OutputsBlock - | "}" (* closes the graph *) - -InputsBlock ::= "inputs" "{" InputDecl* "}" -ConstsBlock ::= "consts" "{" ConstDecl* "}" -NodesBlock ::= "nodes" "{" Stmt* "}" -OutputsBlock ::= "outputs" "{" OutputItem* "}" - -InputDecl ::= Ident ":" Type ";" -ConstDecl ::= Ident ":" Type ConstAnnot* ";" - -OutputItem ::= Ident ("," Ident)* ";"? (* optional semicolon *) - -Stmt ::= (MultiAssign | Assign) ";" -Assign ::= Ident "=" Expr -MultiAssign ::= "[" Ident ("," Ident)* "]" "=" Expr - -Expr ::= Call | Ident | Literal -Call ::= Ident "(" Args? ")" - -Args ::= Arg ("," Arg)* -Arg ::= Ident "=" Value | Value - -Value ::= Literal | Ident - -Literal ::= Array | String | Number | Boolean | Null -Array ::= "[" (Value ("," Value)*)? "]" - -Boolean ::= "true" | "false" -Null ::= "null" - -Type ::= DType Shape -DType ::= "f32" | "f16" | "i4" | "u4" | "i32" | "u32" | "i64" | "u64" | "i8" | "u8" -Shape ::= "[" (Int ("," Int)*)? "]" - -ConstAnnot ::= "@weights" "(" String ")" - | "@scalar" "(" Number ")" - -Ident ::= (ALPHA | "_") (ALNUM | "_")* -Int ::= DIGIT+ -Number ::= "-"? DIGIT+ ("." DIGIT+)? (("e"|"E") ("+"|"-")? DIGIT+)? -String ::= "\"" ( "\\\"" | "\\\\" | (ANY-but-quote) )* "\"" -``` +- `.webnn` is the compact, human-readable graph representation. +- `GraphJson` is the equivalent JSON AST used by the Rust API and tooling. +- `.safetensors` is a self-describing external-weight archive. +- `.weights` plus `.manifest.json` is the raw binary alternative used by the weight utilities and + ONNX converter. -### 2. Weights manifest (`.manifest.json`, optional) +The canonical contracts are documented in: -If the graph references external weights using `@weights("key")`, a manifest file can be provided to: +- [WebNN graph format](docs/webnn-format.md) +- [External weight format](docs/external-weights.md) -- Describe tensor shapes and data types -- Define offsets and sizes inside a binary weights file -- Validate that referenced weights are well-formed - -The manifest is metadata only. It does not contain raw tensor bytes. - -### 3. Binary weights file (`.weights`, optional) - -The `.weights` file is a simple concatenation of raw tensor data. - -It is: -- Compact -- Fast to load -- Independent from graph structure - -This separation allows the same graph definition to be reused with different trained weights. - - -## Core Idea - -The library parses the `.webnn` DSL into a **very small, intentionally simple AST**: - -- Inputs -- Constants -- Nodes (operator name, inputs, options) -- Outputs - -This AST is the **true internal representation** of a graph. - -Once parsed, the AST can be: -- Validated -- Serialized -- Transformed -- Used to construct a WebNN graph - -## Using the AST - -The AST is designed to be easy to consume from other tools. In particular, it can be used to: - -- load, save a build an WebNN graph and its weights using **rustnn** or **PyWebNN** -- Generate WebNN JavaScript `MLGraphBuilder` calls -- Perform lightweight graph analysis or transformations - -The library does not attempt to deeply re-specify WebNN semantics. Anything not -explicitly checked is passed through and left to the WebNN runtime to validate. - -## JSON Serialization (Secondary) - -In addition to the text DSL, the AST can be serialized to a **canonical JSON format**. - -Important points: - -- JSON is **not** the primary authoring format -- It exists as a convenience for programmatic manipulation -- It supports full round-trip conversion back to `.webnn` -- It can store optional metadata such as the graph name - -The JSON format is roughly **10x larger** than the `.webnn` DSL and is best suited for tooling, not manual editing. - -All CLI commands auto-detect and accept both formats. - - -## Features - -- **Convert ONNX models** to WebNN format with static lowering and dynamic input metadata -- Parse WebNN graph text (`.webnn`) into a simple AST -- Serialize the AST to canonical JSON -- Serialize JSON back to `.webnn` with full round-trip support -- Validate graph structure and optional weights manifest -- Emit WebNN JavaScript builder code (`MLGraphBuilder` calls) -- Pack and unpack binary weight files - -This is intended as a **small, hackable reference scaffold**, not a heavy framework. - -## Install - -### From source (local dev) +## Install and build ```bash -git clone https://github.com/tarekziade/webnn-graph -cd webnn-graph -make build -make run -# Or: -webnn-graph --help +cargo build +cargo test ``` -### Install the CLI with Cargo +ONNX conversion is enabled by the default `onnx` feature. Build only the parser, serializer, validators, +emitters, and weight utilities with: ```bash -cargo install webnn-graph +cargo build --no-default-features ``` -## Formats +## CLI -### Text format: .webnn +The CLI accepts `.webnn` or `GraphJson` where indicated. Run `webnn-graph --help` for the complete +option list. -The DSL is block-based and declarative: +| Command | Purpose | +| --- | --- | +| `parse` | Parse `.webnn` and print `GraphJson`. | +| `serialize` | Serialize `GraphJson` as `.webnn`. | +| `validate` | Validate a graph and, optionally, a raw-weight manifest. | +| `emit-js` | Emit WebNN builder JavaScript and the raw `.weights` loader. | +| `emit-html` | Emit a standalone interactive graph visualizer. | +| `pack-weights` | Pack tensor files into a `WGWT` `.weights` archive. | +| `unpack-weights` | Extract tensors from a `WGWT` `.weights` archive. | +| `create-manifest` | Create a raw-weight manifest from tensor files. | +| `extract-weights` | Move inline graph constants into a raw-weight archive. | +| `inline-weights` | Copy raw external weights into `GraphJson`. | +| `convert-onnx` | Convert ONNX to `.webnn` or `GraphJson` when the `onnx` feature is enabled. | -- inputs {} declares typed inputs -- consts {} declares typed constants -- nodes {} lists operator calls in order -- outputs {} declares named graph outputs - - -Types use: -``` -dtype[dim0, dim1, ...] -``` - -Supported dtypes: `f32`, `f16`, `i4`, `u4`, `i32`, `u32`, `i64`, `u64`, `i8`, `u8`. - -## ONNX to WebNN Conversion - -The CLI includes a powerful ONNX-to-WebNN converter that enables you to take existing ONNX -models and convert them to the WebNN format. - -### Prerequisites: Static Lowering for Shape Expressions - -**Important:** ONNX models may contain symbolic input dimensions. `webnn-graph` can preserve unresolved -symbolic input dimensions in graph metadata (v2), but operations such as `reshape` still require static -shape expressions for WebNN lowering. Dynamic input preservation is experimental and must be enabled with -`--experimental-dynamic-inputs`. Use `--optimize` and `--override-dim` as needed. - -#### Why is this necessary? - -WebNN's `reshape` operation requires the shape parameter to be a constant, not a dynamically -computed value. Many ONNX models (especially transformers/BERT) use dynamic shape patterns like: - -``` -Shape → Gather → Concat → Reshape -``` - -These patterns are typically resolved to static constants during conversion. - -#### Built-in Constant Folding - -The converter includes a constant folding engine that automatically: -- Evaluates `Shape` operations at conversion time -- Resolves `Gather` and `Concat` operations on constant data -- Eliminates dynamic shape computation patterns -- Reduces model size by 40-50% for transformer models - -Simply use the `--optimize` flag to enable constant folding: - -```bash -webnn-graph convert-onnx --input model.onnx --optimize \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 -``` - -**What constant folding does:** -- Identifies nodes with all-constant inputs -- Evaluates them at conversion time -- Replaces them with their computed results -- Removes the evaluated nodes from the graph - -**See also:** [Dynamic Dimensions Guide](docs/dynamic-dimensions-guide.md) for help choosing dimension values. - -### Converting ONNX Models - -Convert ONNX models to WebNN format with built-in constant folding: - -```bash -# Basic conversion with optimization (recommended) -webnn-graph convert-onnx --input model.onnx --optimize \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 - -# Output: model.webnn + model.weights + model.manifest.json - -# Custom output paths -webnn-graph convert-onnx \ - --input model.onnx \ - --optimize \ - --output graph.webnn \ - --weights graph.weights \ - --manifest graph.manifest.json \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 - -# Inline weights for small models (not recommended for large models) -webnn-graph convert-onnx --input model.onnx --optimize --inline-weights \ - --override-dim batch_size=1 - -# Output to JSON format instead of .webnn -webnn-graph convert-onnx --input model.onnx --optimize --output model.json \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 -``` - -### Example: Converting SmolLM-135M - -Download the ONNX model: +Examples: ```bash -curl -L "https://huggingface.co/HuggingFaceTB/SmolLM-135M/resolve/main/onnx/model.onnx?download=true" \ - -o /tmp/smol_hf.onnx -``` +# Parse and validate a graph. +cargo run -- parse examples/resnet_head.webnn > /tmp/resnet_head.json +cargo run -- validate /tmp/resnet_head.json -Then convert: +# Serialize GraphJson back to the text format. +cargo run -- serialize /tmp/resnet_head.json > /tmp/resnet_head.webnn -```bash -webnn-graph convert-onnx \ - --input /tmp/smol_hf.onnx \ - --optimize \ - --experimental-dynamic-inputs \ - --output /tmp/smol_hf.webnn \ - --weights /tmp/smol_hf.weights \ - --manifest /tmp/smol_hf.manifest.json +# Generate JavaScript or a standalone visualizer. +cargo run -- emit-js examples/resnet_head.webnn > /tmp/build_graph.js +cargo run -- emit-html examples/resnet_head.webnn > /tmp/graph.html ``` -Observed output from that run: -- Graph header: `webnn_graph "main_graph" v2` -- `/tmp/smol_hf.webnn`: ~694 KB -- `/tmp/smol_hf.weights`: ~513 MB -- `/tmp/smol_hf.manifest.json`: ~423 KB +See [examples/README.md](examples/README.md) for the raw-weight workflow. -### Example: Converting all-MiniLM-L6-v2 +## ONNX conversion -Download the ONNX model: +The converter accepts `ai.onnx` opsets 11 through 18. Static dimension overrides and optional constant +folding can resolve shape-critical ONNX expressions. Experimental bounded dynamic input metadata is available, +but operations whose arguments must be static still require concrete values. ```bash -curl -L "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/onnx/model.onnx?download=true" \ - -o /tmp/minilm.onnx -``` - -Convert with common sentence-embedding overrides: - -```bash -webnn-graph convert-onnx \ - --input /tmp/minilm.onnx \ - --optimize \ +cargo run -- convert-onnx \ + --input model.onnx \ + --output model.webnn \ + --weights model.weights \ + --manifest model.manifest.json \ --override-dim batch_size=1 \ --override-dim sequence_length=128 \ - --output /tmp/minilm.webnn \ - --weights /tmp/minilm.weights \ - --manifest /tmp/minilm.manifest.json + --optimize ``` -### Supported ONNX Operations +Without `--inline-weights`, conversion produces `.webnn`, `.weights`, and `.manifest.json` artifacts. Use +`--experimental-dynamic-inputs` to preserve unresolved input dimensions as bounded `dyn(...)` metadata where +the lowering can otherwise proceed. -The converter focuses on NLP/Transformer operations: +See: -- **Matrix operations**: MatMul, Gemm -- **Element-wise**: Add, Sub, Mul, Div, Pow -- **Normalization**: LayerNormalization, Softmax -- **Tensor manipulation**: Reshape, Transpose, Concat, Split, Squeeze, Unsqueeze -- **Activation**: Relu, Sigmoid, Tanh, Gelu, etc. -- **Reduction**: ReduceMean, ReduceSum, ReduceMax, ReduceMin -- **Utility**: Gather, Slice +- [ONNX to WebNN lowering](docs/onnx-lowering.md) +- [Dynamic dimensions](docs/dynamic-dimensions-guide.md) -### Complete ONNX Workflow Example +## Rust library -```bash -# Step 1: Convert ONNX → WebNN with constant folding -webnn-graph convert-onnx --input bert-base.onnx --optimize \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 \ - --override-dim token_type_ids=128 +The public library exposes the format AST, parser, serializer, validation helpers, external-weight resolver and +SafeTensors writer, JavaScript/HTML emitters, and the optional ONNX converter. -# Output: bert-base.webnn + bert-base.weights + bert-base.manifest.json +```rust +use std::error::Error; -# Step 2: Generate JavaScript for browser/runtime -webnn-graph emit-js bert-base.webnn > buildGraph.js +use webnn_graph::parser::parse_wg_text; +use webnn_graph::serialize::{serialize_graph_to_wg_text, SerializeOptions}; -# Step 3: (Optional) Create HTML visualizer -webnn-graph emit-html bert-base.webnn > visualizer.html -open visualizer.html +fn main() -> Result<(), Box> { +let graph = parse_wg_text(r#" +webnn_graph "identity" v1 { + inputs { x: f32[1]; } + nodes { y = identity(x); } + outputs { y; } +} +"#)?; -# The --optimize flag performs constant folding automatically: -# - Eliminates Shape/Gather/Concat patterns -# - Reduces model size by 40-50% -# - No external preprocessing needed! +let text = serialize_graph_to_wg_text(&graph, SerializeOptions::default())?; +# Ok::<(), Box>(()) ``` -**Example results for BERT models with `--optimize`:** -- **Original ONNX**: 637 nodes with Shape operations -- **After constant folding**: 317 nodes (50% reduction), no Shape operations -- All reshape shape parameters become static constants -- **WebNN output**: All reshape operations use static constants, fully compatible - -## Examples - -Below is the same graph expressed in webnn and JSON. - -### Text - -```webnn -webnn_graph "resnet_head" v1 { - inputs { - x: f32[1, 2048]; - } - - consts { - W: f32[2048, 1000] @weights("W"); - b: f32[1000] @weights("b"); - } - - nodes { - logits0 = matmul(x, W); - logits = add(logits0, b); - probs = softmax(logits, axis=1); - } - - outputs { probs; } -} -``` +## Development -### JSON - -```json -{ - "format": "webnn-graph-json", - "version": 1, - "inputs": { - "x": { "dataType": "float32", "shape": [1, 2048] } - }, - "consts": { - "W": { - "dataType": "float32", - "shape": [2048, 1000], - "init": { "kind": "weights", "ref": "W" } - }, - "b": { - "dataType": "float32", - "shape": [1000], - "init": { "kind": "weights", "ref": "b" } - } - }, - "nodes": [ - { "id": "logits0", "op": "matmul", "inputs": ["x", "W"], "options": {} }, - { "id": "logits", "op": "add", "inputs": ["logits0", "b"], "options": {} }, - { "id": "probs", "op": "softmax", "inputs": ["logits"], "options": { "axis": 1 } } - ], - "outputs": { "probs": "probs" } -} +```bash +make fmt-check +make lint +make test +cargo test --all-features +cargo test --no-default-features ``` -## Notes - -- Validation is intentionally lightweight and structural. -- Operator semantics are mostly pass-through. -- The design favors simplicity and reuse over completeness. -- The AST is stable and meant to be consumed by other WebNN tooling. +The repository keeps prose lines at or below 120 characters. diff --git a/docs/dynamic-dimensions-guide.md b/docs/dynamic-dimensions-guide.md index ed9568e..96e9583 100644 --- a/docs/dynamic-dimensions-guide.md +++ b/docs/dynamic-dimensions-guide.md @@ -1,83 +1,37 @@ -# Dynamic Dimensions Guide +# Dynamic dimensions -A practical guide for choosing dimension override values when converting ONNX models to WebNN format. +ONNX inputs may use symbolic or unknown dimensions. `webnn-graph` can either replace them with concrete values +or preserve them as bounded input metadata. Shape-critical operation arguments must still be resolvable during +conversion. -## Table of Contents +## Static overrides -- [Understanding Dynamic Dimensions](#understanding-dynamic-dimensions) -- [Inspection Methods](#inspection-methods) -- [Common Values by Model Type](#common-values-by-model-type) -- [Decision-Making Process](#decision-making-process) -- [Troubleshooting](#troubleshooting) +Provide one or more values directly: -## Understanding Dynamic Dimensions - -### What Are Dynamic Dimensions? - -ONNX models often use symbolic dimensions (like `batch_size`, `sequence_length`) instead of fixed -numbers. This allows the same model to handle different input sizes. - -Example ONNX input shape: -``` -input_ids: [batch_size, sequence_length] # Dynamic -``` - -In many models, shape-driving expressions must become static for conversion: -``` -input_ids: [1, 128] # Static -``` - -### Why Provide Overrides? - -WebNN executes in browsers and edge devices where: -- Memory must be allocated upfront -- Shape-driving expressions (for example, reshape targets) must be resolvable -- Performance is optimized for specific sizes - -`webnn-graph` can preserve unresolved symbolic **input metadata** in v2 graphs with -`--experimental-dynamic-inputs`, but conversion still needs concrete values when dynamic shape math -cannot be folded. - -## Inspection Methods - -### Method 1: Using Python + ONNX - -```python -pip install onnxslim -onnxslim --inspect model.onnx +```bash +webnn-graph convert-onnx \ + --input model.onnx \ + --override-dim batch_size=1 \ + --override-dim sequence_length=128 ``` -### Method 2: Using Netron - -1. Install Netron: `pip install netron` -2. Open model: `netron model.onnx` -3. Click on input nodes to see shape information -4. Look for dimension parameters like `batch_size`, `seq_len`, etc. - -### Method 3: Check Model Documentation - -Most models on Hugging Face include dimension information: +Or load an object from a file: ```bash -# Visit the model page -# https://huggingface.co// - -# Look for: -# - "Model Details" section -# - "max_seq_length" in config -# - Example usage code +webnn-graph convert-onnx \ + --input model.onnx \ + --override-dims-file dimensions.json ``` -### Method 4: Check the Sidecar File - -webnn-graph supports automatic dimension discovery via `.dims.json` files: +Both accepted JSON shapes are equivalent: -```bash -# If model.onnx has a model.dims.json file, check it: -cat model.dims.json +```json +{ + "batch_size": 1, + "sequence_length": 128 +} ``` -Example content: ```json { "freeDimensionOverrides": { @@ -87,436 +41,80 @@ Example content: } ``` -## Common Values by Model Type +Repeated `--override-dim` values are applied after `--override-dims-file` and replace the same key. -### Text / NLP Models (Transformers, BERT, GPT) +## Implicit sources and precedence -**Typical dimensions:** -```bash ---override-dim batch_size=1 ---override-dim sequence_length=<128|256|512> -``` +The converter resolves dimensions in this order: -**Choosing sequence_length:** - -| Model Type | Recommended | Max | Use Case | -|------------|-------------|-----|----------| -| Sentence embeddings (MiniLM, MPNet) | 128 | 256 | Sentences, titles, short text | -| Document classification (BERT) | 256 | 512 | Paragraphs, articles | -| Question answering (BERT, RoBERTa) | 384 | 512 | Q&A pairs with context | -| Text generation (GPT-2) | 512 | 1024 | Long-form generation | -| Long-document (Longformer) | 1024 | 4096 | Full documents | - -**How to determine:** -```python -# Check tokenizer max length -from transformers import AutoTokenizer -tokenizer = AutoTokenizer.from_pretrained("model-name") -print(f"Max length: {tokenizer.model_max_length}") -``` +1. Values supplied through `--override-dims-file` and `--override-dim`. +2. If no explicit values were supplied, `.dims.json` next to the ONNX file. +3. `freeDimensionOverrides` JSON stored in ONNX model metadata, filling names not already set. +4. Without experimental dynamic inputs, common batch names (`batch_size`, `batch`, `n`, and `b`, matched + case-insensitively) use an inference value of 1. -**Common dimension parameter names:** -- `batch_size`, `batch`, `N`, `B` -- `sequence_length`, `seq_len`, `max_len`, `T`, `L` -- `hidden_size`, `hidden_dim` (usually fixed, not dynamic) +Supplying any explicit override prevents automatic loading of the `.dims.json` sidecar. Put the complete desired +set in the explicit source when mixing would otherwise be required. -### Vision Models (CNNs, Vision Transformers) +An unresolved symbolic dimension produces an error naming the input and suggested `--override-dim` flag. +Non-positive and unnamed dimensions use a generated `_dim` hint. -**Typical dimensions:** -```bash ---override-dim batch_size=1 ---override-dim height=224 ---override-dim width=224 -``` +## Bounded dynamic input metadata -**Standard image sizes by architecture:** - -| Architecture | Size | Notes | -|--------------|------|-------| -| ResNet-50/101 | 224×224 | ImageNet standard | -| EfficientNet-B0 | 224×224 | Scales up with variants | -| EfficientNet-B7 | 600×600 | Higher accuracy, slower | -| Vision Transformer (ViT-B) | 224×224 or 384×384 | Two common variants | -| MobileNet V2/V3 | 224×224 | Mobile-optimized | -| YOLO (object detection) | 416×416 or 640×640 | Detection-specific | -| Semantic segmentation | 512×512 or 1024×1024 | Full-resolution | - -**How to determine:** -```python -# Check preprocessing configuration -from transformers import AutoImageProcessor -processor = AutoImageProcessor.from_pretrained("model-name") -print(f"Size: {processor.size}") -``` - -**Common dimension parameter names:** -- `batch_size`, `batch`, `N`, `B` -- `height`, `H`, `image_height` -- `width`, `W`, `image_width` -- `channels`, `C` (usually 3 for RGB, 1 for grayscale) - -### Audio Models (Speech Recognition, Audio Classification) +Enable experimental preservation with: -**Typical dimensions:** ```bash ---override-dim batch_size=1 ---override-dim sequence_length= # Based on audio duration -``` - -**Determining sequence_length for audio:** -```python -# Formula: sequence_length = sample_rate * duration_seconds / hop_length - -# Example for Wav2Vec2 (16kHz, 10 seconds) -sample_rate = 16000 -duration = 10 # seconds -hop_length = 320 # model-specific -sequence_length = (sample_rate * duration) // hop_length -# Result: ~500 for 10-second audio -``` - -**Common values:** -- Whisper: Processes 30-second chunks → sequence_length based on mel spectrogram frames -- Wav2Vec2: Variable based on audio duration -- Audio classification: Often 16000 samples (1 second at 16kHz) - -## Decision-Making Process - -### Step-by-Step Guide - -#### Step 1: Identify Dynamic Dimensions That Need Values - -Try converting without overrides first: - -```bash -./webnn-graph convert-onnx --input model.onnx -``` - -If conversion cannot resolve required dims, error output will indicate what to set: -``` -Error: unresolved dynamic dimension(s) require explicit overrides: - - input 'input_ids' dim 'batch_size': --override-dim batch_size= - - input 'input_ids' dim 'sequence_length': --override-dim sequence_length= -``` - -#### Step 2: Determine Model Type - -Look at the model filename or documentation: -- `*bert*`, `*roberta*`, `*gpt*` → Text model -- `*resnet*`, `*efficientnet*`, `*vit*` → Vision model -- `*wav2vec*`, `*whisper*` → Audio model - -#### Step 3: Start with Conservative Values - -**Always start with:** -- `batch_size=1` (single inference) -- Smallest reasonable size for other dimensions - -**Why start small?** -- Faster conversion and testing -- Less memory usage -- Easier to debug -- Can always increase later - -#### Step 4: Look Up Standard Values - -Use the tables in [Common Values by Model Type](#common-values-by-model-type) above. - -#### Step 5: Test and Iterate - -```bash -# Test with initial values -./webnn-graph convert-onnx \ +webnn-graph convert-onnx \ --input model.onnx \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 - -# If successful, test inference -# If shapes are wrong, adjust and retry + --experimental-dynamic-inputs ``` -### Example Workflows +Unresolved input dimensions become v2 descriptors such as: -#### Example 1: BERT Sentence Embeddings - -```bash -# Model: sentence-transformers/all-MiniLM-L12-v2 -# Task: Generate sentence embeddings - -# 1. Check documentation -# Hugging Face says: max_seq_length = 256 - -# 2. Choose conservative value -# Most sentences < 128 tokens, so start there - -# 3. Convert -./webnn-graph convert-onnx \ - --input all-MiniLM-L12-v2.onnx \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 \ - --optimize - -# 4. If you need longer sequences, increase -./webnn-graph convert-onnx \ - --input all-MiniLM-L12-v2.onnx \ - --override-dim batch_size=1 \ - --override-dim sequence_length=256 \ - --optimize -``` - -#### Example 2: ResNet Image Classification - -```bash -# Model: ResNet-50 -# Task: Image classification - -# 1. Standard size for ResNet is 224×224 -# 2. No need to check - this is well-known - -# 3. Convert -./webnn-graph convert-onnx \ - --input resnet50.onnx \ - --override-dim batch_size=1 \ - --override-dim height=224 \ - --override-dim width=224 -``` - -#### Example 3: Unknown Custom Model - -```bash -# 1. Inspect the model -python -c " -import onnx -model = onnx.load('custom_model.onnx') -for inp in model.graph.input: - print(inp.name, inp.type.tensor_type.shape) -" - -# Output shows: data [N, 3, H, W] -# This is an image model (3 channels) - -# 2. Try standard image sizes -./webnn-graph convert-onnx \ - --input custom_model.onnx \ - --override-dim N=1 \ - --override-dim H=224 \ - --override-dim W=224 - -# 3. If conversion fails, check error messages -# 4. Try other common sizes: 256, 299, 384, 512 -``` - -## Troubleshooting - -### Problem: "Dynamic dimensions require explicit overrides" - -**Solution:** Some model paths still need concrete values for symbolic dimensions. - -```bash -# Error shows which dimensions need values -Error: unresolved dynamic dimension(s) require explicit overrides: - - input 'input' dim 'height': --override-dim height= - -# Provide the missing dimensions -./webnn-graph convert-onnx \ - --input model.onnx \ - --override-dim height=224 \ - --override-dim width=224 -``` - -### Problem: "Shape inference failed" - -**Cause:** The dimension values you provided are incompatible with the model's operations. - -**Solution:** - -1. Check if the model has specific size requirements: -```python -# Some models only work with specific sizes -# e.g., YOLO expects multiples of 32 -``` - -2. Try multiples of common factors: -```bash -# For CNNs: try 224, 256, 288, 320, 384, 416, 512 -# For transformers: try 128, 256, 384, 512 -``` - -3. Enable optimization to help with shape inference: -```bash -./webnn-graph convert-onnx \ - --input model.onnx \ - --optimize \ - --override-dim ... +```webnn +inputs { + input_ids: i64[dyn("batch_size", 8), dyn("sequence_length", 4096)]; +} ``` -### Problem: "Constant folding failed" +The current default maximum sizes are: -**Cause:** The model uses dynamic operations that can't be resolved with the given dimensions. +- 4096 for names containing `past`, `seq`, or `length`, and for `s` or `t`; +- 8 for names containing `batch`, and for `b` or `n`; +- 65535 for other symbolic or generated names. -**Solution:** +Explicit overrides always produce static dimensions. When experimental preservation is enabled, unresolved batch +names remain bounded dynamic dimensions rather than receiving the static inference default of 1. -1. Make sure you've provided ALL dynamic dimensions: -```bash -# Check for missing dimensions -python -c " -import onnx -model = onnx.load('model.onnx') -for inp in model.graph.input: - for dim in inp.type.tensor_type.shape.dim: - if dim.dim_param: - print(f'Dynamic: {dim.dim_param}') -" -``` +If any input contains a dynamic dimension, the converter emits graph version 2. Otherwise it emits version 1. -2. Use the `--optimize` flag: -```bash -./webnn-graph convert-onnx \ - --input model.onnx \ - --optimize \ - --override-dim batch_size=1 \ - --override-dim sequence_length=128 -``` - -### Problem: Conversion succeeds but inference fails +## Static lowering constraints -**Cause:** The dimension values are too small for your actual input data. +Dynamic input metadata does not make every ONNX shape expression dynamic at runtime. Conversion still needs +concrete information for arguments such as reshape targets, slice bounds, axes, permutations, split sizes, and +other operation-specific attributes. -**Solution:** +`--optimize` runs the registered constant-folding evaluators before lowering. This can eliminate shape-producing +subgraphs when their inputs are constant: -1. Check your actual input size: -```javascript -// In JavaScript -console.log("Input shape:", inputData.shape); -``` - -2. Increase dimensions to match: ```bash -# If your inputs are 256 tokens but you used 128 -./webnn-graph convert-onnx \ +webnn-graph convert-onnx \ --input model.onnx \ - --override-dim sequence_length=256 # Increased from 128 -``` - -### Problem: Out of memory during conversion - -**Cause:** Dimension values are too large. - -**Solution:** - -1. Reduce to smaller values: -```bash -# Instead of 512, try 256 -# Instead of 256, try 128 -``` - -2. Use batch_size=1 (not higher): -```bash ---override-dim batch_size=1 # Always use 1 for inference -``` - -## Best Practices - -### 1. Always Start with batch_size=1 - -For inference in WebNN (browsers/edge devices): -```bash ---override-dim batch_size=1 -``` - -Only increase batch size if you're doing batch processing server-side. - -### 2. Match Your Use Case - -Choose dimensions based on your actual inputs: -- Short texts (tweets, titles): `sequence_length=128` -- Medium texts (articles): `sequence_length=256` -- Long texts (documents): `sequence_length=512+` - -### 3. Consider Memory Constraints - -Larger dimensions = more memory: -``` -Memory ∝ batch_size × sequence_length × hidden_size -``` - -For browser inference, prefer smaller dimensions. - -### 4. Use Standard Values When Possible - -Standard sizes are well-tested: -- Images: 224, 256, 384, 512 -- Text: 128, 256, 512, 1024 -- These are powers of 2 or common multiples - -### 5. Document Your Choices - -Create a `.dims.json` file alongside your model: - -```json -{ - "freeDimensionOverrides": { - "batch_size": 1, - "sequence_length": 128 - }, - "notes": "128 tokens handles 95% of our sentences. Max length is 256 if needed." -} -``` - -### 6. Test with Real Data - -After conversion, test with actual inputs: -```javascript -// Verify converted model works with your data -const result = await context.compute(graph, { - input_ids: actualInputIds, // Your real input - attention_mask: actualMask -}); -``` - -## Quick Reference - -### Text Models -```bash ---override-dim batch_size=1 --override-dim sequence_length=128 -``` - -### Vision Models -```bash ---override-dim batch_size=1 --override-dim height=224 --override-dim width=224 -``` - -### Unknown Model -```bash -# 1. Try without overrides (see error message) -# 2. Check model documentation -# 3. Start with smallest reasonable values -# 4. Iterate based on errors/results + --experimental-dynamic-inputs \ + --optimize ``` -## Additional Resources +If a required value remains unresolved, conversion fails instead of inventing a shape. Add a concrete override +for the relevant symbolic dimension or use a model whose shape-critical expression can be folded. -- [ONNX Model Zoo](https://github.com/onnx/models) - Standard models with known shapes -- [Hugging Face Model Hub](https://huggingface.co/models) - Model documentation -- [Netron](https://netron.app/) - Visual model inspector -- [WebNN Specification](https://www.w3.org/TR/webnn/) - WebNN requirements +## Choosing overrides -## Summary +Choose values from the model's tokenizer, processor, configuration, or deployment contract. Dimension names are +model-defined; inspect the ONNX input descriptors with a tool such as Netron or an ONNX library rather than +assuming a repository-wide list. -**The decision process:** -1. Try converting without overrides and capture unresolved dimensions from the error -2. Determine model type (text/vision/audio) -3. Look up standard values for that model type -4. Start with conservative (small) values -5. Test and iterate as needed - -**Most common pattern:** -```bash -./webnn-graph convert-onnx \ - --input model.onnx \ - --optimize \ - --override-dim batch_size=1 \ - --override-dim = -``` +Start with the actual production batch and sequence/image/audio sizes. Overrides specialize the converted graph, +so a different deployment shape normally requires another conversion or a supported bounded-dynamic input. -Remember: **batch_size=1** is almost always correct for inference! +See [ONNX to WebNN lowering](onnx-lowering.md) for the complete conversion flow and +[WebNN graph format](webnn-format.md) for the serialized dynamic-dimension syntax. diff --git a/docs/external-weights.md b/docs/external-weights.md new file mode 100644 index 0000000..6e930ac --- /dev/null +++ b/docs/external-weights.md @@ -0,0 +1,116 @@ +# External weight format + +A `GraphJson` constant with `ConstInit::Weights { ref }`, rendered in `.webnn` as `@weights("ref")`, obtains +its bytes from a sidecar. The graph declaration remains authoritative for the logical data type and shape. + +Two sidecar families are supported: + +- a self-describing SafeTensors archive; +- an opaque `.weights` byte blob addressed by a JSON manifest. + +## Discovery + +An explicit weights path is resolved relative to the graph file. Without one, discovery checks in order: + +1. `.safetensors` +2. `.weights` +3. `model.safetensors` +4. `model.weights` + +The first existing path wins. A SafeTensors path ignores the manifest argument. A `.weights` path requires an +explicit manifest or a discovered `.manifest.json` or `manifest.json`. + +Exact tensor names take precedence. As a compatibility fallback, archive and manifest names are normalized by +replacing `::` with `__` and `.` with `_`. Resolution fails when more than one stored name normalizes to the +same requested reference. + +## Loading and ownership + +SafeTensors and raw `.weights` files are read into memory while references are resolved. Every selected tensor +range is then copied into the owned `InlineBytes` representation used by `GraphJson` consumers. The format does +not provide decode-time weight streaming. + +Callers must keep a sidecar immutable while it is being resolved. Producers should complete a temporary file +and rename it into place rather than modifying an installed archive. + +## SafeTensors + +Ordinary logical types use the corresponding SafeTensors storage type: + +| Graph declaration | SafeTensors storage | +| --- | --- | +| `float32` | F32 | +| `float16` | F16 | +| `int32` | I32 | +| `uint32` | U32 | +| `int64` | I64 | +| `uint64` | U64 | +| `int8` | I8 | +| `uint8` | U8 | + +A BF16 tensor may satisfy a `float32` declaration. The loader converts its values to F32. Other dtype +mismatches are rejected, and ordinary tensor shapes must match the graph declaration exactly. + +### Packed Int4 and Uint4 + +SafeTensors has no native 4-bit storage type. Logical `int4` and `uint4` declarations use this versioned +extension: + +- archive metadata key: `rustnn.webnn.packed4` +- supported metadata value: `1` +- physical SafeTensors type: U8 +- physical shape: `[ceil(logical_element_count / 2)]` +- byte layout: the first logical element occupies the low nibble and the second occupies the high nibble + +The `.webnn` or `GraphJson` declaration retains the logical 4-bit dtype and original shape. The loader requires +the marker whenever an archive is used for an external 4-bit declaration. It rejects a missing or unknown +marker, non-U8 storage, an incorrect physical shape or byte count, overflow, and missing references. + +Archives containing only ordinary tensors need no packed-4-bit marker and remain compatible with earlier files. + +## Manifest-backed `.weights` + +For raw weights, the manifest contains a `tensors` object keyed by each `@weights` reference. Resolution uses +`byteOffset` and `byteLength` as an absolute half-open byte range from the beginning of the `.weights` file. + +The weight utilities use this version-1 manifest shape: + +```json +{ + "format": "wg-weights-manifest", + "version": 1, + "endianness": "little", + "tensors": { + "weight": { + "dataType": "float32", + "shape": [2, 2], + "byteOffset": 8, + "byteLength": 16, + "layout": "row-major" + } + } +} +``` + +The external resolver accepts related manifest layouts as long as `tensors..byteOffset` and +`byteLength` are present. It validates integer conversion, range addition, and file bounds. Graph dtype and +shape remain authoritative during resolution; the richer manifest fields are consumed by the pack/unpack tools. + +Two current producers use absolute offsets differently: + +- `pack-weights` and `extract-weights` write `WGWT`, followed by a little-endian U32 version (`1`), then tensor + bytes. Their first tensor offset is 8. `unpack-weights` and `inline-weights` validate this header. +- `convert-onnx` writes a headerless concatenation whose first tensor offset is 0. + +Because offsets are absolute, the shared external resolver can load both layouts without interpreting a header. + +## SafeTensors writer + +`write_external_weights_safetensors` accepts a `GraphJson`, bytes keyed by weight reference, and a destination +path. It rejects missing bytes, conflicting declarations, unsupported mappings, shape/length mismatches, and +packed element-count overflow. It adds the packed-4-bit marker only when the graph contains external Int4 or +Uint4 tensors. + +The writer calls `safetensors::serialize_to_file` for a uniquely named temporary file in the destination +directory and renames it only after successful serialization. Failed serialization or installation removes the +temporary file and does not install a partial final archive. diff --git a/docs/onnx-lowering.md b/docs/onnx-lowering.md index 0c23ddc..4af15cc 100644 --- a/docs/onnx-lowering.md +++ b/docs/onnx-lowering.md @@ -1,215 +1,88 @@ -# Reference: ONNX → WebNN Lowering - -This reference explains how `webnn-graph` lowers ONNX graphs into the WebNN DSL. - -## Key concepts up front -- **Shape inference**: collect and propagate concrete shapes for every value. Inputs and initializers seed - known shapes; integer constants feed shape math; ops like `MatMul`, `Transpose`, `Concat`, `Reduce*`, - `Gather`, `Reshape` (with known newShape), etc., derive output shapes. Some unresolved symbolic input - metadata may still be preserved in v2 graphs, but shape-critical paths must be static. -- **Const folding boundaries**: only small integer tensors used for shape/axes/newShape/starts/ends math are - folded. Real weight tensors are never folded—they stay as external weights or inline bytes for tiny - scalars. This keeps semantic fidelity while unlocking static shapes. -- **Serialization**: the `.webnn` output stores only structure (inputs, const declarations, nodes, outputs). - Actual tensor bytes live in `model.weights` with offsets/types in `model.manifest.json`. Inline bytes are - allowed for tiny scalars; everything else is referenced via `@weights("key")`. - -## Why simplification is mandatory -- **WebNN is static**: execution-relevant tensor shapes must be known at build time. -- **ONNX can be dynamic**: inputs often carry symbolic dims (e.g., `batch`, `seq_len`), and graphs may - manipulate shapes at runtime (`Shape → Gather → Concat → Reshape` pipelines, dynamic `Slice`, etc.). -- **Goal**: arrive at a graph where every tensor shape is concrete and every shape-producing expression is - either folded to a constant or rejected early with a clear error, while preserving input-dim metadata - when possible. - - -## Two-phase lowering -1) **Prepare the ONNX graph for static lowering** - - **Provide overrides when needed**: use `--override-dim name=value` to pin symbolic dims (e.g., - `batch=1`, `seq_len=128`). A sidecar `*.dims.json` can supply the same. Defaults may kick in for - common names (`batch`, `sequence_length`, etc.) when missing. - - **Enable constant folding**: use the `--optimize` flag to activate built-in constant folding that - eliminates dynamic-shape plumbing such as: - - `Shape` → `Gather` → `Concat` → `Reshape` - - Constant axes/starts/ends passed around as tensors - - **Result**: shape expressions required by WebNN become static; unresolved symbolic input dims can be - preserved as metadata in v2 graphs when `--experimental-dynamic-inputs` is enabled. - -2) **Lower the ONNX graph to WebNN DSL** - - **Opset guard**: only `ai.onnx` opset 11–18 is accepted. - - **Shape/type seeding**: inputs (minus initializers) become WebNN inputs; shapes can include dynamic - metadata for unresolved symbolic dims; - initializers become constants; integer constants are recorded for later shape math. - - **Static shape inference**: a conservative pass infers shapes for intermediates and folds small integer - tensors needed for axes/newShape/starts/ends. If unresolved dynamics block required inference, node - conversion fails with a targeted error. - - **Constant folding** (`--optimize`): evaluates Shape/Gather/Concat/Cast/Squeeze/Unsqueeze operations - at conversion time, replacing them with their computed constant values. Reduces graph size by 40-50% - for transformer models. - - **Node conversion**: each ONNX node is mapped to one or more WebNN ops via `OpRegistry`; purely - constant outputs are emitted as consts and skipped as nodes. - - **Serialization**: emit `.webnn` (structure only) plus `.weights` and `.manifest.json` (raw bytes + - offsets). Inline bytes are used only for tiny scalars; real tensors use `@weights("key")`. - -### Flow at a glance -```mermaid -flowchart LR - A["ONNX model (may have symbolic dims)"] - B["Shape overrides (--override-dim) + constant folding (--optimize)"] - C["Static shape inference"] - D["Constant folding: Shape/Gather/Concat/etc"] - E["Op mapping (OpRegistry)"] - F["WebNN DSL (.webnn)"] - G["Weights sidecars (.weights + .manifest.json)"] - - A --> B --> C --> D --> E --> F - D --> G +# ONNX to WebNN lowering + +The optional `onnx` feature converts ONNX models into the `GraphJson` AST and serializes them as `.webnn` or +JSON. The converter accepts `ai.onnx` opsets 11 through 18. Other domains are retained for operator-specific +handling rather than being checked by the `ai.onnx` opset guard. + +## Conversion flow + +1. Read and decode the ONNX protobuf. +2. Optionally run registered constant-folding evaluators with `--optimize`. +3. Collect explicit, sidecar, and model-metadata dimension overrides. +4. Convert graph inputs and initializers into WebNN inputs and constants. +5. Infer intermediate types and shapes required by each supported lowering. +6. Lower each ONNX node through the operator registry. +7. Optionally extract constants to `.weights` plus `.manifest.json`. +8. Serialize the resulting `GraphJson` as `.webnn` or JSON. + +Unsupported opsets, operators, dtypes, attributes, or unresolved shape-critical values fail conversion with an +error. The converter does not silently omit a node. + +## Dimensions and shape expressions + +Static overrides specialize symbolic ONNX inputs. Experimental bounded dynamic input metadata can preserve an +unresolved input dimension when downstream lowering can still determine every required operation argument. + +See [Dynamic dimensions](dynamic-dimensions-guide.md) for override precedence, `.dims.json`, metadata, inferred +batch defaults, and the `dyn(...)` representation. + +ONNX graphs often calculate operation arguments through small tensor subgraphs. With `--optimize`, registered +evaluators fold expressions whose inputs are known constants before lowering. Without a foldable value, an +operation that requires a static reshape target, axis, permutation, slice bound, or similar parameter is rejected. + +The exact supported behavior is operator- and opset-specific. Source tests are authoritative; this page does not +claim that every variant of a named ONNX operator is supported. + +## Constants and output artifacts + +By default, `convert-onnx` extracts initializers and large inline constants into a headerless `.weights` blob and +a manifest whose offsets start at zero. Constants larger than 1 KiB are moved out of the graph when extraction is +enabled. Smaller scalar and byte constants may remain inline. + +```bash +webnn-graph convert-onnx \ + --input model.onnx \ + --output model.webnn \ + --weights model.weights \ + --manifest model.manifest.json \ + --override-dim batch_size=1 \ + --optimize ``` +Omitting explicit output paths derives all three names from the ONNX filename. `--inline-weights` suppresses the +raw sidecars and retains constants in the graph representation. + +The ONNX converter currently emits manifest-backed raw weights, not SafeTensors. Consumers may subsequently save +the graph through a writer that produces `.webnn` plus `.safetensors`. -## How dynamic constructs are handled -- **Symbolic input dims**: may be preserved in v2 input metadata when unresolved if - `--experimental-dynamic-inputs` is enabled, but conversion still requires concrete values wherever - shape math must be static. -- **Shape-producing ops** (`Shape`, `Gather`, `Concat`, `Unsqueeze`, `Squeeze`, `Cast` of ints): - - If the inputs are compile-time constants, the converter folds them and records both the values and - shapes. - - If not foldable, the op is left as a WebNN `shape`/`gather`/`concat`/`unsqueeze`/`squeeze` node, but - only when its shape impact is already known and compatible with WebNN. -- **Reshape**: - - Requires `newShape` to be fully known. The converter pulls it from constant tensors or folded const - values; `-1` is resolved using the input element count. - - If `newShape` is not fully static, conversion fails with a clear “WebNN requires static newShape” error. -- **Slice**: - - Starts/ends/axes/steps must be constants. Steps other than 1 are rejected. Negative indices are - normalized using known input dims. Unknown dims lead to a failure. -- **Gather**: - - Axis is normalized; if both data and indices are constant, it is folded. Otherwise, shape inference - ensures the output shape is determined. -- **Transpose/Concat/Split/Unsqueeze/Squeeze**: - - Permutations/axes must be known. Concat requires all input shapes known; otherwise, conversion stops. -- **Constant folding scope**: - - Only small integer tensors used for shape math are folded. Real weight tensors are never folded; they - are carried through to the weights file. - -### Before/after examples for common tricky patterns -- **Dynamic reshape pipeline** - - Original ONNX: `X -> Shape -> Gather -> Concat -> Reshape(X, newShape=tensor)` with symbolic dims. - - With `--optimize`: `newShape` becomes a constant tensor (e.g., `[1,128,12,32]`); `Shape/Gather/Concat` - are removed. - - After WebNN lowering: a single `reshape(X, newShape=[1,128,12,32])` node; `newShape` is static and - embedded as an inline small const if needed. - ```mermaid - flowchart LR - A[X] --> B[Shape] - B --> C[Gather] - C --> D[Concat] - D --> E["Reshape X,newShape=tensor"] - subgraph constant_folding - F[X] --> G["Reshape X,newShape=[1,128,12,32]"] - end - subgraph webnn - H["reshape(X,[1,128,12,32])"] - end - ``` -- **Slice with computed bounds** - - Original ONNX: `starts`/`ends`/`axes`/`steps` produced by small subgraphs. - - With `--optimize`: those tensors become constants with normalized axes; negative indices are resolved. - - After WebNN lowering: one `slice` node with static starts/ends/axes/steps; rejected if any bound stays - dynamic or `step != 1`. - ```mermaid - flowchart LR - A[starts subgraph] --> B[Slice] - C[ends subgraph] --> B - D[axes subgraph] --> B - E[steps subgraph] --> B - subgraph constant_folding - F[starts const] --> G["Slice(static)"] - H[ends const] --> G - I[axes const] --> G - J[steps const] --> G - end - subgraph webnn - K[slice static bounds] - end - ``` -- **Axis/permute tensors** - - Original ONNX: axes for `Unsqueeze`/`Squeeze`/`Reduce*` or perm for `Transpose` are fed by tensors. - - With `--optimize`: axes/perm are folded into constant tensors. - - After WebNN lowering: ops carry inline static `axes` or `permutation` arrays; if still dynamic, the op - is rejected. - ```mermaid - flowchart LR - A[axes tensor] --> B[Unsqueeze] - subgraph constant_folding - C[axes const] --> D["Unsqueeze axes=[1,3]"] - end - subgraph webnn - E["unsqueeze axes=[1,3]"] - end - ``` -- **Gather for shape math** - - Original ONNX: `Gather(Shape(X), idx)` to pick a dim. - - With `--optimize`: `Shape` and `Gather` are removed and replaced by a scalar constant (e.g., sequence - length). - - After WebNN lowering: the scalar becomes an inline const; no `gather` node is emitted. - ```mermaid - flowchart LR - A[X] --> B[Shape] - B --> C[Gather idx] - subgraph constant_folding - D[const dim]:::const - end - subgraph webnn - E[inline scalar const]:::const - end - classDef const fill:#eef,stroke:#66f; - ``` -- **Broadcasted elementwise ops** - - Original ONNX: elementwise ops rely on runtime broadcasting. - - With `--optimize`: shapes are static and compatible; broadcasting stays implicit. - - After WebNN lowering: elementwise ops are emitted as-is; shapes are already known, so no extra shape - ops are needed. - ```mermaid - flowchart LR - A["X (static shape)"] --> C[Add] - B["Y (static shape)"] --> C - subgraph webnn - D["add(X,Y) with inferred broadcast"] - end - ``` - - - -## Operator mapping (ONNX → WebNN) -- **MatMul/Gemm**: `matmul` (plus optional transposes, alpha/beta scaling, and bias add for Gemm). - Transposes are emitted as separate `transpose` nodes when requested. -- **Elementwise**: `Add`, `Sub`, `Mul`, `Div`, `Pow` map directly to WebNN elementwise ops with - broadcasted shapes already inferred. -- **Activations**: `Relu`, `Gelu`, `Tanh`, `Sigmoid`, `Sqrt`, `Exp`, `Log`, `Abs`, `Neg`, `Erf` map - one-to-one. -- **Normalization**: `LayerNormalization` (with epsilon/axes) and `Softmax` (axis). -- **Reshape family**: `Reshape`, `Transpose`, `Concat`, `Split`, `Unsqueeze`, `Squeeze` with static - axes/newShape/permutation; failures if not static. -- **Utility**: `Shape`, `Gather`, `Slice` as described above; `Cast` with supported dtype mapping. -- **Reductions**: `ReduceMean`, `ReduceSum`, `ReduceMax`, `ReduceMin` with `axes` and `keepdims`. - -Unsupported ops (or ops with remaining dynamism) fail fast with an explicit “unsupported operator” or -“WebNN requires static …” message to keep the pipeline predictable. - - - -## End-to-end recipe (anchor example: `all-MiniLM-L6-v2-webnn`) -1) **Convert with overrides + folding**: - - `webnn-graph convert-onnx --input model.onnx --optimize \\ - --override-dim batch_size=1 \\ - --override-dim sequence_length=128 \\ - --output model.webnn --weights model.weights --manifest model.manifest.json` -2) **Lowering behavior**: - - Shape-driving expressions are folded to static constants when possible. - - Unresolved symbolic input dimensions can be preserved in v2 input metadata. -3) **Artifacts** - - `model.webnn`: structure-only graph; inputs pinned; consts reference `@weights("…")` or inline tiny - scalars; nodes are WebNN ops with sanitized IDs; outputs expose `last_hidden_state`. - - `model.weights`: raw little-endian tensor bytes, concatenated. - - `model.manifest.json`: dtype/shape/byte offsets for every `@weights` tensor. +See [External weight format](external-weights.md) for both raw and SafeTensors contracts. + +## Graph versions + +The converter emits graph version 2 when at least one input retains a bounded dynamic dimension; otherwise it +emits version 1. Both versions use the same nodes, constants, and output structures. + +The `@quantized` graph flag is metadata carried by `GraphJson` and serialization. It does not select a different +ONNX lowering pipeline by itself. + +## Diagnostics + +Use the global `--debug` flag before the subcommand to enable converter diagnostics: + +```bash +webnn-graph --debug convert-onnx --input model.onnx --optimize +``` + +When conversion fails, first check: + +- whether the model uses `ai.onnx` opset 11–18; +- whether every required symbolic dimension has an override or a usable bounded representation; +- whether `--optimize` can fold the shape-producing expression; +- whether the specific operator form and attributes have a registered lowering. + +After conversion, parse and validate the artifact independently: + +```bash +webnn-graph parse model.webnn > model.json +webnn-graph validate model.webnn --weights-manifest model.manifest.json +``` diff --git a/docs/t5-conversion.md b/docs/t5-conversion.md deleted file mode 100644 index 52d4df9..0000000 --- a/docs/t5-conversion.md +++ /dev/null @@ -1,125 +0,0 @@ -# T5 Model Conversion Guide - -This guide covers converting T5 (Text-To-Text Transfer Transformer) models from ONNX to WebNN format. - -## Prerequisites - -- T5 model exported to ONNX format (separate encoder and decoder files) -- webnn-graph built in release mode: `cargo build --release` - -## Supported Operations - -The converter supports all operations required for T5 models: - -**Constant Folding:** -- Range (positional encoding sequences) -- ConstantOfShape (attention masks) -- Shape, Gather, Concat, Unsqueeze, Squeeze, Cast - -**Computation:** -- MatMul, Gemm (attention and feed-forward layers) -- Add, Sub, Mul, Div, Pow, Min, Max (elementwise operations) -- LayerNormalization, Softmax (normalization) -- Greater, Less, Equal, GreaterOrEqual, LessOrEqual (comparisons) -- Where (conditional selection) -- Reshape, Transpose, Split, Concat (tensor manipulation) - -## T5 Encoder Conversion - -Convert the T5 encoder model: - -```bash -./target/release/webnn-graph convert-onnx \ - --input encoder_model.onnx \ - --optimize \ - --override-dim batch_size=1 \ - --override-dim encoder_sequence_length=128 -``` - -**Dimension Overrides:** -- `batch_size=1` - Typical for inference workloads -- `encoder_sequence_length=128` - Standard sequence length for T5-small (adjust for longer sequences) - -**Output Files:** -- `encoder_model.webnn` - Graph structure (human-readable) -- `encoder_model.weights` - Binary weights (model parameters) -- `encoder_model.manifest.json` - Weights metadata - -## T5 Decoder Conversion - -Convert the T5 decoder model: - -```bash -./target/release/webnn-graph convert-onnx \ - --input decoder_model.onnx \ - --optimize \ - --override-dim batch_size=1 \ - --override-dim decoder_sequence_length=128 \ - --override-dim encoder_sequence_length=128 -``` - -**Dimension Overrides:** -- `batch_size=1` - Typical for inference -- `decoder_sequence_length=128` - Target sequence length (adjust as needed) -- `encoder_sequence_length=128` - Must match encoder output length (for cross-attention) - -**Output Files:** -- `decoder_model.webnn` - Graph structure -- `decoder_model.weights` - Binary weights -- `decoder_model.manifest.json` - Weights metadata - -## Model Size Notes - -T5 models with constant folding enabled achieve ~40-50% size reduction: -- Dynamic shape operations are resolved at conversion time -- No Shape/Gather/Concat operations in final graph -- All reshape operations use static constants - -## Common Dimension Override Values - -**T5-small:** -- Encoder: `encoder_sequence_length=128` to `512` -- Decoder: `decoder_sequence_length=128` to `512` - -**T5-base:** -- Encoder: `encoder_sequence_length=128` to `512` -- Decoder: `decoder_sequence_length=128` to `512` - -**T5-large:** -- Encoder: `encoder_sequence_length=128` to `512` -- Decoder: `decoder_sequence_length=128` to `512` - -Choose sequence lengths based on your application: -- Short text (summaries, classification): 128-256 -- Medium text (translation, Q&A): 256-512 -- Long text (documents): 512+ - -## Validation - -After conversion, validate the graph structure: - -```bash -./target/release/webnn-graph validate encoder_model.webnn -./target/release/webnn-graph validate decoder_model.webnn -``` - -## JavaScript Generation - -Generate WebNN JavaScript code for browser/Node.js: - -```bash -./target/release/webnn-graph emit-js encoder_model.webnn > buildEncoder.js -./target/release/webnn-graph emit-js decoder_model.webnn > buildDecoder.js -``` - -## Troubleshooting - -**Missing operators:** If conversion fails with "unsupported operator" errors, please report the issue with: -- The operator name -- The model variant (T5-small, T5-base, etc.) -- The conversion command used - -**Dynamic dimensions:** All dynamic dimensions must be overridden at conversion time. WebNN requires static -shapes for all operations. - -See also: [Dynamic Dimensions Guide](dynamic-dimensions-guide.md) diff --git a/docs/webnn-format.md b/docs/webnn-format.md new file mode 100644 index 0000000..d05b90a --- /dev/null +++ b/docs/webnn-format.md @@ -0,0 +1,156 @@ +# WebNN graph format + +The `.webnn` format is a textual representation of the `GraphJson` AST exposed by this crate. It describes +graph structure and constant declarations; external tensor bytes live in sidecars documented in +[External weight format](external-weights.md). + +This DSL is a `webnn-graph` interchange format. It is not a file format defined by the W3C WebNN specification. + +## Document structure + +```webnn +webnn_graph "example" v2 @quantized { + inputs { + input: f32[dyn("batch", 8), 4]; + } + + consts { + bias: f32[] @scalar(0.5); + bytes: u8[4] @bytes([1, 2, 3, 4]); + weight: f32[4, 4] @weights("weight"); + } + + nodes { + sum = add(input, bias, metadata={kind: "example", axes: [1]}); + [left, right] = split(sum, splits=[2, 2], axis=1); + result = concat(left, right, axis=1); + } + + outputs { result; } +} +``` + +A document contains one header followed by any of these blocks: + +- `inputs`: named graph inputs and their descriptors. +- `consts`: named constants and their initialization method. +- `nodes`: operation calls and the operands they produce. +- `outputs`: names exported by the graph. + +Blocks may be omitted when empty. `#` starts a comment that continues to the end of the line. + +## Header and versions + +The header is `webnn_graph "name" vN {`. Supported serialized versions are `v1` and `v2`. + +- `v1` represents static input dimensions. +- `v2` can preserve bounded dynamic input dimensions as `dyn("name", max_size)`. + +The optional `@quantized` annotation records that the graph contains quantized representations. It is metadata +for consumers; it does not alter parsing or tensor bytes. + +## Data types and shapes + +The text data types map to `GraphJson::DataType` as follows: + +| Text | GraphJson | +| --- | --- | +| `f32` | `float32` | +| `f16` | `float16` | +| `i4` | `int4` | +| `u4` | `uint4` | +| `i32` | `int32` | +| `u32` | `uint32` | +| `i64` | `int64` | +| `u64` | `uint64` | +| `i8` | `int8` | +| `u8` | `uint8` | + +A shape is a comma-separated list in brackets. `[]` is a known rank-0 scalar, not an unknown shape. + +Input dimensions may be static integers or bounded dynamic dimensions. Constants must have fully static shapes. + +```webnn +scalar: f32[]; +static: f32[1, 128, 768]; +dynamic: f32[dyn("batch", 8), dyn("sequence", 4096), 768]; +``` + +Each dynamic dimension has a stable name and a maximum size. Execution consumers decide which concrete sizes +within those bounds they support. + +## Inputs and constants + +Inputs declare only a name, data type, and shape: + +```webnn +inputs { + input_ids: i64[1, 128]; +} +``` + +Constants use one initialization annotation: + +- `@weights("ref")` resolves bytes from an external sidecar. If no annotation is written, the parser uses + `@weights("constant_name")`. +- `@scalar(value)` stores a JSON number as a scalar initializer. +- `@bytes([0, 1, ...])` stores bytes directly in the graph. + +```webnn +consts { + matrix: f32[4, 4] @weights("encoder.matrix"); + epsilon: f32[] @scalar(0.00001); + mask: u8[4] @bytes([1, 1, 0, 0]); +} +``` + +The graph declaration is authoritative for the logical constant name, data type, and shape. See +[External weight format](external-weights.md) for sidecar lookup and storage rules. + +## Nodes and values + +A single-output operation assigns its result to one identifier: + +```webnn +sum = add(lhs, rhs); +``` + +A multi-output operation declares all output identifiers: + +```webnn +[first, second] = split(input, splits=[2, 2], axis=1); +``` + +Arguments without a name are operand references or literals. Named arguments become entries in the node's +`options` map. Values may be identifiers, strings, numbers, booleans, `null`, arrays, or JSON-like objects. + +Identifiers begin with an ASCII letter, `_`, `/`, or `$`. Remaining characters may also include digits and `.`. +Serialized graphs should prefer portable identifier names because downstream consumers may impose stricter rules. + +## Outputs + +The outputs block lists exported operand names: + +```webnn +outputs { logits; hidden_state; } +``` + +Each name is both the public output binding and the referenced operand name in the text format. `GraphJson` +stores outputs as a map from binding name to operand reference; serializing to `.webnn` emits the map keys. +Consumers that need distinct binding and operand names must normalize them before text serialization. + +## GraphJson relationship + +The parser maps `.webnn` into these `GraphJson` fields: + +- header name, version, and quantized flag; +- ordered maps of inputs and constants; +- an ordered list of nodes; +- an ordered map of output bindings. + +`parse_wg_text` parses the DSL. `serialize_graph_to_wg_text` serializes `GraphJson` versions 1 and 2. A graph +that uses constructs representable in the DSL is stable across parse → serialize → parse, modulo formatting and +map ordering. + +The maintained [format reference example](../examples/format_reference.webnn) is exercised by an automated +round-trip test. diff --git a/examples/README.md b/examples/README.md index 86f8b05..9cff589 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1,288 +1,105 @@ -# WebNN Graph Examples +# Examples -This directory contains a complete working example demonstrating the full workflow of the WebNN graph DSL. +The checked-in examples demonstrate the `.webnn` DSL, raw-weight utilities, JavaScript emitter, and interactive +HTML emitter. -## Contents +For format details, see: -- **resnet_head.webnn** - Graph definition (data-agnostic template) -- **weights.manifest.json** - Original manifest (for reference) -- **tensors/** - Example tensor files (W.bin, b.bin with metadata) -- **build_example.sh** - Automated build script -- **browser_example.html** - Interactive browser demo +- [WebNN graph format](../docs/webnn-format.md) +- [External weight format](../docs/external-weights.md) -## Generated Files (after running build script) +## Files -- **resnet_head.manifest.json** - Generated weights manifest -- **resnet_head.weights** - Binary weights file (7.8 MB) -- **buildGraph.js** - Generated JavaScript with WeightsFile helper +- `format_reference.webnn` exercises the maintained v2 grammar and is covered by a round-trip test. +- `resnet_head.webnn` is a small graph whose constants use external references. +- `tensors/` contains tensor metadata. The build script reconstructs untracked raw `.bin` inputs from the + checked-in reference archive. +- `build_example.sh` creates a manifest, packs a `WGWT` weights file, and emits JavaScript. +- `browser_example.html` is an example host for emitted WebNN builder code. -## Quick Start +Some generated artifacts are checked in for inspection. Rerun the script before relying on them after changing +the graph, manifest, tensors, or emitter. -### 1. Build the Example - -Run the automated build script: +## Parse, serialize, and validate ```bash -./examples/build_example.sh +cargo run -- parse examples/format_reference.webnn > /tmp/format_reference.json +cargo run -- serialize /tmp/format_reference.json > /tmp/format_reference.webnn +cargo run -- validate /tmp/format_reference.webnn ``` -This will: -1. Create a manifest from the tensor files -2. Pack the tensors into a binary `.weights` file -3. Generate the JavaScript module +The format-reference graph declares an external tensor for syntax coverage. Parsing and structural validation do +not resolve that sidecar. -### 2. Run in Browser +## Raw-weight workflow -Serve the examples directory with a web server: +Run the complete ResNet-head example from the repository root: ```bash -# Using Python -python3 -m http.server 8000 - -# Or using Node.js -npx http-server -p 8000 +./examples/build_example.sh ``` -Then open http://localhost:8000/browser_example.html - -### 3. Manual Workflow - -You can also run each step manually: +It performs the equivalent steps: ```bash -# Step 1: Create manifest from tensors cargo run -- create-manifest \ --input-dir examples/tensors \ - --output examples/my_manifest.json \ + --output examples/resnet_head.manifest.json \ --endianness little -# Step 2: Pack weights into binary cargo run -- pack-weights \ - --manifest examples/my_manifest.json \ + --manifest examples/resnet_head.manifest.json \ --input-dir examples/tensors \ - --output examples/my_model.weights - -# Step 3: Generate JavaScript -cargo run -- parse examples/resnet_head.webnn | \ - cargo run -- emit-js /dev/stdin > examples/my_graph.js + --output examples/resnet_head.weights -# Optional: Unpack weights for inspection -cargo run -- unpack-weights \ - --weights examples/my_model.weights \ - --manifest examples/my_manifest.json \ - --output-dir examples/unpacked/ -``` - -## Understanding the Example - -### Graph Definition (resnet_head.webnn) +cargo run -- validate examples/resnet_head.webnn \ + --weights-manifest examples/resnet_head.manifest.json -The `.webnn` file defines only the **structure** - no actual data: - -```webnn -webnn_graph "resnet_head" v1 { - inputs { - x: f32[1, 2048]; // Shape only! - } - consts { - W: f32[2048, 1000] @weights("W"); // External reference - b: f32[1000] @weights("b"); - } - nodes { - logits0 = matmul(x, W); - logits = add(logits0, b); - probs = softmax(logits, axis=1); - } - outputs { probs; } -} +cargo run -- emit-js examples/resnet_head.webnn > examples/buildGraph.js ``` -### Tensor Files (tensors/ directory) - -Each weight has: -- **W.bin** - Raw binary data (float32 values) -- **W.meta.json** - Metadata (shape, dataType, layout) - -### Generated JavaScript +The packed file begins with the `WGWT` version-1 header. Manifest byte offsets are absolute from the beginning +of that file, including its eight-byte header. -The `buildGraph.js` contains: +Inspect the archive by unpacking it into a temporary directory: -1. **WeightsFile class** - Loads and validates weights - ```javascript - const weights = await WeightsFile.load('model.weights', 'manifest.json'); - ``` - -2. **buildGraph function** - Constructs the WebNN graph - ```javascript - const graph = await buildGraph(context, weights); - ``` - -### Runtime Usage - -The key benefit: **graph reusability** - -```javascript -// One-time setup -const weights = await WeightsFile.load('resnet_head.weights', 'manifest.json'); -const context = await navigator.ml.createContext(); -const graph = await buildGraph(context, weights); - -// Run multiple times with different inputs -const result1 = await context.compute(graph, { x: input1 }); -const result2 = await context.compute(graph, { x: input2 }); -const result3 = await context.compute(graph, { x: input3 }); -// ... no rebuilding needed! +```bash +cargo run -- unpack-weights \ + --weights examples/resnet_head.weights \ + --manifest examples/resnet_head.manifest.json \ + --output-dir /tmp/resnet_head_tensors ``` -## File Formats +## JavaScript and HTML output -### Weights Binary Format (.weights) +Generate WebNN builder JavaScript: -``` -┌─────────────────────────────────────┐ -│ Magic: "WGWT" (4 bytes) │ -├─────────────────────────────────────┤ -│ Version: 1 (4 bytes, little-endian)│ -├─────────────────────────────────────┤ -│ W tensor data (8,192,000 bytes) │ -│ b tensor data (4,000 bytes) │ -└─────────────────────────────────────┘ +```bash +cargo run -- emit-js examples/resnet_head.webnn > /tmp/build_graph.js ``` -### Weights Manifest (.manifest.json) +Generate a standalone graph visualization: -Describes where each tensor lives in the binary file: - -```json -{ - "format": "wg-weights-manifest", - "version": 1, - "endianness": "little", - "tensors": { - "W": { - "dataType": "float32", - "shape": [2048, 1000], - "byteOffset": 8, - "byteLength": 8192000 - }, - "b": { ... } - } -} +```bash +cargo run -- emit-html examples/resnet_head.webnn > /tmp/resnet_head.html ``` -## Creating Your Own Tensors - -To add your own weights: - -1. Create binary files (e.g., using Python/NumPy): - ```python - import numpy as np - - # Create tensor - W = np.random.randn(2048, 1000).astype(np.float32) - W.tofile('W.bin') +Executing the generated builder requires a JavaScript environment that implements the WebNN APIs used by the +graph. The emitters do not provide a WebNN runtime. - # Create metadata - import json - with open('W.meta.json', 'w') as f: - json.dump({ - "dataType": "float32", - "shape": [2048, 1000], - "byteOffset": 8, # After header - "byteLength": W.nbytes, - "layout": "row-major" - }, f) - ``` +## ONNX conversion -2. Run the build workflow to pack and generate code - -## Troubleshooting - -**"WebNN API not supported"** -- Use a browser with WebNN support (Chrome/Edge with experimental features enabled) -- Visit `chrome://flags` and enable "Experimental Web Platform features" - -**"Failed to load weights"** -- Make sure you're serving the files via HTTP (not file://) -- Check that all generated files exist in the examples directory - -**"Invalid magic bytes"** -- Re-run `build_example.sh` to regenerate the weights file -- Ensure the manifest matches the weights file - -## ONNX Conversion Workflow - -If you have an existing ONNX model, you can convert it to WebNN format. This is especially useful for transformer models like BERT. - -### Built-in Constant Folding - -The converter includes built-in constant folding (enabled with `--optimize`) that automatically handles dynamic shape patterns. No external preprocessing needed! - -### Step-by-Step Example +ONNX conversion is available with the default `onnx` feature: ```bash -# Step 1: Convert ONNX to WebNN with constant folding -cargo run -- convert-onnx --input your-model.onnx --optimize \ +cargo run -- convert-onnx \ + --input model.onnx \ + --output model.webnn \ + --weights model.weights \ + --manifest model.manifest.json \ --override-dim batch_size=1 \ - --override-dim sequence_length=128 - -# This creates three files: -# - your-model.webnn (graph structure) -# - your-model.weights (binary weights) -# - your-model.manifest.json (weights metadata) - -# Step 2: Generate JavaScript -cargo run -- emit-js your-model.webnn > your-model.js - -# Step 3: Create an interactive visualizer -cargo run -- emit-html your-model.webnn > visualizer.html -open visualizer.html -``` - -### How Constant Folding Works - -WebNN doesn't support dynamic shapes. ONNX models (especially transformers) often use patterns like: - -``` -Shape → Gather → Concat → Reshape -``` - -The `--optimize` flag automatically evaluates these patterns at conversion time: - -- ✅ Before: `Shape` operation computes dimensions dynamically -- ✅ After: Reshape uses constant `[1, 128, 768]` directly - -**Results for BERT models with `--optimize`:** -- Original: ~637 nodes with Shape operations -- After constant folding: ~317 nodes (50% reduction) -- All reshape operations use static constants - -**See also:** [Dynamic Dimensions Guide](../docs/dynamic-dimensions-guide.md) for help choosing dimension override values. - -### Using the Converted Model - -The converted WebNN model can be used just like the manual examples: - -```javascript -// Load the converted model -const weights = await WeightsFile.load('your-model.weights', - 'your-model.manifest.json'); -const context = await navigator.ml.createContext(); -const graph = await buildGraph(context, weights); - -// Run inference with your input data -const inputIds = new Int32Array([101, 2023, 2003, ...]); -const attentionMask = new Int32Array([1, 1, 1, ...]); - -const result = await context.compute(graph, { - input_ids: inputIds, - attention_mask: attentionMask -}); + --optimize ``` -## Next Steps - -- Modify `resnet_head.webnn` to define your own graph -- Replace tensors with your trained model weights -- Convert existing ONNX models using the workflow above -- Use the generated JavaScript in your web application +See [ONNX to WebNN lowering](../docs/onnx-lowering.md) and +[Dynamic dimensions](../docs/dynamic-dimensions-guide.md) for current conversion behavior. diff --git a/examples/build_example.sh b/examples/build_example.sh index 76e9afb..4d66a74 100755 --- a/examples/build_example.sh +++ b/examples/build_example.sh @@ -6,13 +6,29 @@ set -e # Change to repository root cd "$(dirname "$0")/.." +TENSOR_DIR="examples/tensors" +TEMP_DIR="" +if [ ! -f "$TENSOR_DIR/W.bin" ] || [ ! -f "$TENSOR_DIR/b.bin" ]; then + TEMP_DIR=$(mktemp -d) + trap 'rm -rf "$TEMP_DIR"' EXIT + TENSOR_DIR="$TEMP_DIR/tensors" + mkdir -p "$TENSOR_DIR" + + echo "Raw tensor inputs are not checked in; unpacking the reference archive..." + cargo run --quiet -- unpack-weights \ + --weights examples/resnet_head.weights \ + --manifest examples/resnet_head.manifest.json \ + --output-dir "$TENSOR_DIR" + cp examples/tensors/*.meta.json "$TENSOR_DIR/" +fi + echo "=== WebNN Graph Complete Workflow Example ===" echo # Step 1: Create manifest from tensor directory echo "Step 1: Creating weights manifest from tensors..." cargo run --quiet -- create-manifest \ - --input-dir examples/tensors \ + --input-dir "$TENSOR_DIR" \ --output examples/resnet_head.manifest.json \ --endianness little echo @@ -21,14 +37,13 @@ echo echo "Step 2: Packing weights into binary format..." cargo run --quiet -- pack-weights \ --manifest examples/resnet_head.manifest.json \ - --input-dir examples/tensors \ + --input-dir "$TENSOR_DIR" \ --output examples/resnet_head.weights echo # Step 3: Parse graph and emit JavaScript echo "Step 3: Generating JavaScript code..." -cargo run --quiet -- parse examples/resnet_head.webnn | \ - cargo run --quiet -- emit-js /dev/stdin > examples/buildGraph.js +cargo run --quiet -- emit-js examples/resnet_head.webnn > examples/buildGraph.js echo "Generated examples/buildGraph.js" echo @@ -37,22 +52,4 @@ echo "=== Generated Files ===" ls -lh examples/resnet_head.manifest.json examples/resnet_head.weights examples/buildGraph.js | awk '{print $9, "-", $5}' echo -echo "=== Usage ===" -echo "The generated files can be used in a browser:" -echo " - buildGraph.js: Contains WeightsFile class and buildGraph() function" -echo " - resnet_head.weights: Binary weights file (8.0 MB)" -echo " - resnet_head.manifest.json: Weights metadata" -echo -echo "Example JavaScript:" -echo " import { WeightsFile, buildGraph } from './buildGraph.js';" -echo " const weights = await WeightsFile.load('resnet_head.weights', 'resnet_head.manifest.json');" -echo " const context = await navigator.ml.createContext();" -echo " const graph = await buildGraph(context, weights);" -echo " const result = await context.compute(graph, { x: inputData });" -echo - -echo "=== Optional: Unpack weights for inspection ===" -echo "cargo run -- unpack-weights \\" -echo " --weights examples/resnet_head.weights \\" -echo " --manifest examples/resnet_head.manifest.json \\" -echo " --output-dir examples/unpacked/" +echo "See examples/README.md for parse, serialize, validation, emitter, and unpack commands." diff --git a/examples/format_reference.webnn b/examples/format_reference.webnn new file mode 100644 index 0000000..34dd5f8 --- /dev/null +++ b/examples/format_reference.webnn @@ -0,0 +1,20 @@ +# Maintained syntax example for docs and parser/serializer round-trip coverage. +webnn_graph "format_reference" v2 @quantized { + inputs { + input: f32[dyn("batch", 8), 4]; + } + + consts { + bias: f32[] @scalar(0.5); + inline_mask: u8[4] @bytes([1, 1, 0, 0]); + external_weight: f32[4, 4] @weights("encoder.weight"); + } + + nodes { + sum = add(input, bias, metadata={kind: "reference", axes: [1], enabled: true}); + [left, right] = split(sum, splits=[2, 2], axis=1); + result = concat(left, right, axis=1); + } + + outputs { result; } +} diff --git a/src/external_weights.rs b/src/external_weights.rs index 58ff4d6..126def2 100644 --- a/src/external_weights.rs +++ b/src/external_weights.rs @@ -1,12 +1,13 @@ //! Resolve `@weights` / [`ConstInit::Weights`] using sidecar files //! next to a graph path (SafeTensors or manifest + raw weights blob). -use std::collections::HashMap; -use std::fs; +use std::collections::{BTreeMap, HashMap}; +use std::fs; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; use half::bf16; -use safetensors::tensor::Dtype as StDtype; +use safetensors::tensor::{Dtype as StDtype, TensorView}; use safetensors::SafeTensors; use serde::Deserialize; use thiserror::Error; @@ -22,6 +23,11 @@ pub const DEFAULT_PATH_SAFETENSORS: &str = "model.safetensors"; /// Default weights manifest basename when not using a stem-prefixed `*.manifest.json` file. pub const DEFAULT_PATH_MANIFEST: &str = "manifest.json"; +/// SafeTensors archive metadata key for packed logical Int4/Uint4 tensors. +pub const PACKED_4BIT_METADATA_KEY: &str = "rustnn.webnn.packed4"; +/// Current version of the packed logical Int4/Uint4 storage extension. +pub const PACKED_4BIT_METADATA_VERSION: &str = "1"; + /// Failure while resolving external weights for a [`GraphJson`]. #[derive(Debug, Error)] pub enum WeightResolveError { @@ -50,6 +56,30 @@ pub enum WeightResolveError { Missing(String), } +/// Failure while writing a SafeTensors archive for a [`GraphJson`]. +#[derive(Debug, Error)] +pub enum WeightWriteError { + /// The graph declarations or supplied tensor bytes are inconsistent. + #[error("[safetensors] {0}")] + Validation(String), + /// SafeTensors could not serialize the archive. + #[error("failed to write SafeTensors archive `{path}`: {source}")] + Serialize { + path: PathBuf, + #[source] + source: safetensors::SafeTensorError, + }, + /// The completed temporary archive could not be installed at its destination. + #[error("failed to install SafeTensors archive `{path}`: {source}")] + Install { + path: PathBuf, + #[source] + source: std::io::Error, + }, +} + +static TEMP_FILE_SEQUENCE: AtomicU64 = AtomicU64::new(0); + fn graph_has_external_weight_refs(graph_json: &GraphJson) -> bool { graph_json .consts @@ -77,6 +107,53 @@ fn safetensors_st_dtype_matches_ast(st: StDtype, ast: &AstDataType) -> bool { ) } +fn ast_dtype_to_safetensors(ast: &AstDataType) -> Option { + Some(match ast { + AstDataType::Float32 => StDtype::F32, + AstDataType::Float16 => StDtype::F16, + AstDataType::Int32 => StDtype::I32, + AstDataType::Uint32 => StDtype::U32, + AstDataType::Int64 => StDtype::I64, + AstDataType::Uint64 => StDtype::U64, + AstDataType::Int8 => StDtype::I8, + AstDataType::Uint8 => StDtype::U8, + AstDataType::Int4 | AstDataType::Uint4 => return None, + }) +} + +fn logical_element_count(name: &str, shape: &[u32]) -> Result { + shape.iter().try_fold(1usize, |count, &dimension| { + count + .checked_mul(dimension as usize) + .ok_or_else(|| format!("constant `{name}` element count overflows usize")) + }) +} + +fn packed_byte_length(name: &str, shape: &[u32]) -> Result { + Ok(logical_element_count(name, shape)?.div_ceil(2)) +} + +fn validate_packed4_marker(bytes: &[u8], path: &Path) -> Result<(), WeightResolveError> { + let (_, metadata) = SafeTensors::read_metadata(bytes).map_err(|error| { + WeightResolveError::Safetensors(format!("failed to read `{}`: {error}", path.display())) + })?; + let marker = metadata + .metadata() + .as_ref() + .and_then(|values| values.get(PACKED_4BIT_METADATA_KEY)); + match marker.map(String::as_str) { + Some(PACKED_4BIT_METADATA_VERSION) => Ok(()), + Some(version) => Err(WeightResolveError::Safetensors(format!( + "`{}` uses unsupported packed 4-bit format version `{version}`", + path.display() + ))), + None => Err(WeightResolveError::Safetensors(format!( + "`{}` contains external int4/uint4 weights but lacks the `{PACKED_4BIT_METADATA_KEY}` metadata marker", + path.display() + ))), + } +} + fn st_shape_matches_const(st_shape: &[usize], const_shape: &[u32]) -> bool { if st_shape.len() != const_shape.len() { return false; @@ -96,7 +173,7 @@ fn bf16_bytes_to_f32_le_bytes(data: &[u8]) -> Result, WeightResolveError ))); } let mut out = Vec::with_capacity(data.len() * 2); - for chunk in data.chunks_exact(2) { + for chunk in data.as_chunks::<2>().0 { let bits = u16::from_le_bytes([chunk[0], chunk[1]]); let v = bf16::from_bits(bits).to_f32(); out.extend_from_slice(&v.to_le_bytes()); @@ -136,6 +213,13 @@ fn resolve_tensor_view<'a>( .map_err(|e| WeightResolveError::Safetensors(format!("tensor `{ref}` (via `{orig}`): {e}"))) } +fn read_file(path: &Path) -> Result, WeightResolveError> { + fs::read(path).map_err(|source| WeightResolveError::ReadFile { + path: path.to_path_buf(), + source, + }) +} + fn inline_weights_from_safetensors( graph_json: &mut GraphJson, safetensors_path: &Path, @@ -151,14 +235,18 @@ fn inline_weights_from_safetensors( weight_ref_count ); - let bytes = fs::read(safetensors_path).map_err(|source| WeightResolveError::ReadFile { - path: safetensors_path.to_path_buf(), - source, - })?; - let st = SafeTensors::deserialize(&bytes).map_err(|e| { + let archive_bytes = read_file(safetensors_path)?; + let st = SafeTensors::deserialize(&archive_bytes).map_err(|e| { WeightResolveError::Safetensors(format!("`{}`: {e}", safetensors_path.display())) })?; let sanitized_map = safetensors_sanitized_name_map(&st)?; + let has_packed_refs = graph_json.consts.values().any(|decl| { + matches!(decl.data_type, AstDataType::Int4 | AstDataType::Uint4) + && matches!(decl.init, ConstInit::Weights { .. }) + }); + if has_packed_refs { + validate_packed4_marker(&archive_bytes, safetensors_path)?; + } for (const_name, const_decl) in graph_json.consts.iter_mut() { let ConstInit::Weights { r#ref } = &const_decl.init else { @@ -175,6 +263,33 @@ fn inline_weights_from_safetensors( return Err(e); } }; + + if matches!(const_decl.data_type, AstDataType::Int4 | AstDataType::Uint4) { + let expected = packed_byte_length(const_name, &const_decl.shape) + .map_err(WeightResolveError::Safetensors)?; + if view.dtype() != StDtype::U8 { + return Err(WeightResolveError::Safetensors(format!( + "packed 4-bit weight `{ref}` (constant `{const_name}`) must use SafeTensors U8 storage, found {:?}", + view.dtype() + ))); + } + if view.shape() != [expected] { + return Err(WeightResolveError::Safetensors(format!( + "packed 4-bit weight `{ref}` (constant `{const_name}`) storage shape {:?} does not match expected [{expected}] for logical shape {:?}", + view.shape(), const_decl.shape + ))); + } + if view.data().len() != expected { + return Err(WeightResolveError::Safetensors(format!( + "packed 4-bit weight `{ref}` (constant `{const_name}`) has {} bytes, expected {expected}", + view.data().len() + ))); + } + const_decl.init = ConstInit::InlineBytes { + bytes: view.data().to_vec(), + }; + continue; + } if !st_shape_matches_const(view.shape(), &const_decl.shape) { let msg = format!( "shape mismatch for weight `{ref}` (constant `{const_name}`): graph {:?} vs safetensors {:?}", @@ -272,10 +387,7 @@ fn inline_weights_from_manifest( path: manifest_path.to_path_buf(), source, })?; - let weights_bytes = fs::read(weights_path).map_err(|source| WeightResolveError::ReadFile { - path: weights_path.to_path_buf(), - source, - })?; + let weights_bytes = read_file(weights_path)?; let manifest: FlexibleManifest = serde_json::from_str(&manifest_text).map_err(|source| { WeightResolveError::ManifestJson { @@ -490,6 +602,114 @@ pub fn resolve_external_weights( inline_weights_from_manifest(graph_json, &mp, &wp) } +fn temporary_archive_path(path: &Path) -> PathBuf { + let sequence = TEMP_FILE_SEQUENCE.fetch_add(1, Ordering::Relaxed); + let name = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("weights.safetensors"); + path.with_file_name(format!(".{name}.{}.{}.tmp", std::process::id(), sequence)) +} + +/// Writes all external constants declared by `graph_json` to one SafeTensors archive. +/// +/// `external_weights` is keyed by each [`ConstInit::Weights`] reference. Logical Int4 and Uint4 +/// payloads remain packed low-nibble-first and are stored as one-dimensional U8 tensors. +pub fn write_external_weights_safetensors( + graph_json: &GraphJson, + external_weights: &HashMap, + destination: &Path, +) -> Result<(), WeightWriteError> { + let mut declarations: BTreeMap<&str, (&AstDataType, &[u32])> = BTreeMap::new(); + for (constant_name, declaration) in &graph_json.consts { + let ConstInit::Weights { r#ref } = &declaration.init else { + continue; + }; + if let Some((data_type, shape)) = declarations.get(r#ref.as_str()) { + if *data_type != &declaration.data_type || *shape != declaration.shape.as_slice() { + return Err(WeightWriteError::Validation(format!( + "weight reference `{ref}` has conflicting declarations (constant `{constant_name}`)" + ))); + } + } else { + declarations.insert(r#ref, (&declaration.data_type, &declaration.shape)); + } + } + + let mut views = Vec::with_capacity(declarations.len()); + let mut has_packed4 = false; + for (weight_ref, (data_type, logical_shape)) in declarations { + let bytes = external_weights.get(weight_ref).copied().ok_or_else(|| { + WeightWriteError::Validation(format!( + "weight reference `{weight_ref}` has no supplied tensor bytes" + )) + })?; + let (storage_dtype, storage_shape) = if matches!( + data_type, + AstDataType::Int4 | AstDataType::Uint4 + ) { + let packed_len = packed_byte_length(weight_ref, logical_shape) + .map_err(WeightWriteError::Validation)?; + if bytes.len() != packed_len { + return Err(WeightWriteError::Validation(format!( + "packed 4-bit weight `{weight_ref}` has {} bytes, expected {packed_len} for logical shape {logical_shape:?}", + bytes.len() + ))); + } + has_packed4 = true; + (StDtype::U8, vec![packed_len]) + } else { + let dtype = ast_dtype_to_safetensors(data_type).ok_or_else(|| { + WeightWriteError::Validation(format!( + "weight reference `{weight_ref}` has unsupported data type {data_type:?}" + )) + })?; + ( + dtype, + logical_shape + .iter() + .map(|&dimension| dimension as usize) + .collect(), + ) + }; + let view = TensorView::new(storage_dtype, storage_shape, bytes).map_err(|error| { + WeightWriteError::Validation(format!("weight reference `{weight_ref}`: {error}")) + })?; + views.push((weight_ref.to_string(), view)); + } + + let metadata = has_packed4.then(|| { + HashMap::from([( + PACKED_4BIT_METADATA_KEY.to_string(), + PACKED_4BIT_METADATA_VERSION.to_string(), + )]) + }); + let temporary = temporary_archive_path(destination); + if let Err(source) = safetensors::serialize_to_file(views, metadata, &temporary) { + let _ = fs::remove_file(&temporary); + return Err(WeightWriteError::Serialize { + path: destination.to_path_buf(), + source, + }); + } + + let install_result = match fs::rename(&temporary, destination) { + Ok(()) => Ok(()), + Err(source) if source.kind() == std::io::ErrorKind::AlreadyExists => { + fs::remove_file(destination).and_then(|()| fs::rename(&temporary, destination)) + } + Err(source) => Err(source), + }; + if let Err(source) = install_result { + let _ = fs::remove_file(&temporary); + return Err(WeightWriteError::Install { + path: destination.to_path_buf(), + source, + }); + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -842,4 +1062,236 @@ mod tests { other => panic!("expected inline bytes, got {:?}", other), } } + + fn external_graph(constants: Vec<(&str, AstDataType, Vec, &str)>) -> GraphJson { + let mut graph = crate::ast::new_graph_json(); + for (name, data_type, shape, weight_ref) in constants { + graph.consts.insert( + name.to_string(), + crate::ast::ConstDecl { + data_type, + shape, + init: ConstInit::Weights { + r#ref: weight_ref.to_string(), + }, + }, + ); + } + graph + } + + fn write_test_archive( + path: &Path, + tensors: Vec<(String, Dtype, Vec, Vec)>, + marker: Option<&str>, + ) { + let views = tensors + .iter() + .map(|(name, dtype, shape, bytes)| { + ( + name.clone(), + TensorView::new(*dtype, shape.clone(), bytes).unwrap(), + ) + }) + .collect::>(); + let metadata = marker.map(|version| { + HashMap::from([(PACKED_4BIT_METADATA_KEY.to_string(), version.to_string())]) + }); + safetensors::serialize_to_file(views, metadata, path).unwrap(); + } + + #[test] + fn mixed_ordinary_and_packed_weights_roundtrip_in_one_archive() { + let temp_dir = TempDir::new().unwrap(); + let graph_path = temp_dir.path().join("mixed.webnn"); + let archive_path = temp_dir.path().join("mixed.safetensors"); + let mut graph = external_graph(vec![ + ("even", AstDataType::Int4, vec![4], "even"), + ("odd", AstDataType::Uint4, vec![3], "odd"), + ("float", AstDataType::Float32, vec![2], "float"), + ]); + let even = [0x21, 0x43]; + let odd = [0x65, 0x07]; + let float = [0_u8, 0, 128, 63, 0, 0, 0, 64]; + let bytes = HashMap::from([ + ("even".to_string(), even.as_slice()), + ("odd".to_string(), odd.as_slice()), + ("float".to_string(), float.as_slice()), + ]); + + write_external_weights_safetensors(&graph, &bytes, &archive_path).unwrap(); + let archive_bytes = read_file(&archive_path).unwrap(); + let (_, metadata) = SafeTensors::read_metadata(&archive_bytes).unwrap(); + assert_eq!( + metadata.metadata().as_ref().unwrap()[PACKED_4BIT_METADATA_KEY], + PACKED_4BIT_METADATA_VERSION + ); + let archive = SafeTensors::deserialize(&archive_bytes).unwrap(); + assert_eq!(archive.tensor("even").unwrap().shape(), [2]); + assert_eq!(archive.tensor("odd").unwrap().shape(), [2]); + assert_eq!(archive.tensor("float").unwrap().shape(), [2]); + + resolve_external_weights(&mut graph, &graph_path, None, None).unwrap(); + assert_eq!( + graph.consts["even"].init, + ConstInit::InlineBytes { + bytes: even.to_vec() + } + ); + assert_eq!( + graph.consts["odd"].init, + ConstInit::InlineBytes { + bytes: odd.to_vec() + } + ); + assert_eq!( + graph.consts["float"].init, + ConstInit::InlineBytes { + bytes: float.to_vec() + } + ); + } + + #[test] + fn packed_weights_require_supported_marker() { + let temp_dir = TempDir::new().unwrap(); + let graph_path = temp_dir.path().join("packed.webnn"); + let archive_path = temp_dir.path().join("packed.safetensors"); + let graph = external_graph(vec![("q", AstDataType::Uint4, vec![3], "q")]); + + for marker in [None, Some("999")] { + write_test_archive( + &archive_path, + vec![("q".to_string(), Dtype::U8, vec![2], vec![0x21, 0x03])], + marker, + ); + let mut candidate = graph.clone(); + let error = resolve_external_weights(&mut candidate, &graph_path, None, None) + .unwrap_err() + .to_string(); + assert!(error.contains(if marker.is_some() { + "unsupported" + } else { + "lacks" + })); + } + } + + #[test] + fn packed_weights_validate_storage_dtype_and_shape() { + let temp_dir = TempDir::new().unwrap(); + let graph_path = temp_dir.path().join("packed.webnn"); + let archive_path = temp_dir.path().join("packed.safetensors"); + let graph = external_graph(vec![("q", AstDataType::Int4, vec![3], "q")]); + + write_test_archive( + &archive_path, + vec![("q".to_string(), Dtype::I8, vec![2], vec![0x21, 0x03])], + Some(PACKED_4BIT_METADATA_VERSION), + ); + let error = resolve_external_weights(&mut graph.clone(), &graph_path, None, None) + .unwrap_err() + .to_string(); + assert!(error.contains("must use SafeTensors U8")); + + write_test_archive( + &archive_path, + vec![("q".to_string(), Dtype::U8, vec![1, 2], vec![0x21, 0x03])], + Some(PACKED_4BIT_METADATA_VERSION), + ); + let error = resolve_external_weights(&mut graph.clone(), &graph_path, None, None) + .unwrap_err() + .to_string(); + assert!(error.contains("storage shape")); + } + + #[test] + fn writer_rejects_missing_malformed_and_overflowing_weights_without_partials() { + let temp_dir = TempDir::new().unwrap(); + let archive_path = temp_dir.path().join("packed.safetensors"); + let graph = external_graph(vec![("q", AstDataType::Uint4, vec![3], "q")]); + + let error = write_external_weights_safetensors(&graph, &HashMap::new(), &archive_path) + .unwrap_err() + .to_string(); + assert!(error.contains("no supplied tensor bytes")); + + let malformed = [0x21]; + let error = write_external_weights_safetensors( + &graph, + &HashMap::from([("q".to_string(), malformed.as_slice())]), + &archive_path, + ) + .unwrap_err() + .to_string(); + assert!(error.contains("expected 2")); + + let overflow = external_graph(vec![( + "q", + AstDataType::Int4, + vec![u32::MAX, u32::MAX, u32::MAX], + "q", + )]); + let empty = []; + let error = write_external_weights_safetensors( + &overflow, + &HashMap::from([("q".to_string(), empty.as_slice())]), + &archive_path, + ) + .unwrap_err() + .to_string(); + assert!(error.contains("element count overflows")); + assert!(!archive_path.exists()); + assert!(temp_dir.path().read_dir().unwrap().all(|entry| !entry + .unwrap() + .file_name() + .to_string_lossy() + .ends_with(".tmp"))); + } + + #[test] + fn packed_weights_share_sanitized_lookup_and_ambiguity_rules() { + let temp_dir = TempDir::new().unwrap(); + let graph_path = temp_dir.path().join("packed.webnn"); + let archive_path = temp_dir.path().join("packed.safetensors"); + let graph = external_graph(vec![("q", AstDataType::Uint4, vec![3], "a__b")]); + + write_test_archive( + &archive_path, + vec![("a::b".to_string(), Dtype::U8, vec![2], vec![0x21, 0x03])], + Some(PACKED_4BIT_METADATA_VERSION), + ); + let mut resolved = graph.clone(); + resolve_external_weights(&mut resolved, &graph_path, None, None).unwrap(); + assert_eq!( + resolved.consts["q"].init, + ConstInit::InlineBytes { + bytes: vec![0x21, 0x03] + } + ); + + write_test_archive( + &archive_path, + vec![ + ("a::b".to_string(), Dtype::U8, vec![2], vec![0x21, 0x03]), + ("a..b".to_string(), Dtype::U8, vec![2], vec![0x21, 0x03]), + ], + Some(PACKED_4BIT_METADATA_VERSION), + ); + let error = resolve_external_weights(&mut graph.clone(), &graph_path, None, None) + .unwrap_err() + .to_string(); + assert!(error.contains("ambiguous sanitized tensor name")); + + write_test_archive( + &archive_path, + vec![("other".to_string(), Dtype::U8, vec![2], vec![0x21, 0x03])], + Some(PACKED_4BIT_METADATA_VERSION), + ); + let missing = external_graph(vec![("q", AstDataType::Uint4, vec![3], "missing")]); + let error = resolve_external_weights(&mut missing.clone(), &graph_path, None, None) + .unwrap_err() + .to_string(); + assert!(error.contains("not found")); + } } diff --git a/src/lib.rs b/src/lib.rs index 596623d..386c687 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,10 @@ pub mod ast; pub mod debug; pub mod external_weights; -pub use external_weights::{resolve_external_weights, WeightResolveError}; +pub use external_weights::{ + resolve_external_weights, write_external_weights_safetensors, WeightResolveError, + WeightWriteError, PACKED_4BIT_METADATA_KEY, PACKED_4BIT_METADATA_VERSION, +}; pub mod emit_html; pub mod emit_js; diff --git a/src/onnx/constant_folding.rs b/src/onnx/constant_folding.rs index cfee43c..2a45788 100644 --- a/src/onnx/constant_folding.rs +++ b/src/onnx/constant_folding.rs @@ -70,7 +70,9 @@ impl TensorData { match data_type { x if x == TensorProto_DataType::Int64 as i32 => { let values = raw_data - .chunks_exact(8) + .as_chunks::<8>() + .0 + .iter() .map(|c| { i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) }) @@ -79,21 +81,27 @@ impl TensorData { } x if x == TensorProto_DataType::Int32 as i32 => { let values = raw_data - .chunks_exact(4) + .as_chunks::<4>() + .0 + .iter() .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]])) .collect(); Ok(TensorData::Int32(values)) } x if x == TensorProto_DataType::Float as i32 => { let values = raw_data - .chunks_exact(4) + .as_chunks::<4>() + .0 + .iter() .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) .collect(); Ok(TensorData::Float32(values)) } x if x == TensorProto_DataType::Double as i32 => { let values = raw_data - .chunks_exact(8) + .as_chunks::<8>() + .0 + .iter() .map(|c| { f64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) }) diff --git a/src/onnx/convert.rs b/src/onnx/convert.rs index f6e659b..add1777 100644 --- a/src/onnx/convert.rs +++ b/src/onnx/convert.rs @@ -390,14 +390,18 @@ fn infer_shape( shape_tensor .raw_data .as_slice() - .chunks_exact(4) + .as_chunks::<4>() + .0 + .iter() .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64) .collect() } else { shape_tensor .raw_data .as_slice() - .chunks_exact(8) + .as_chunks::<8>() + .0 + .iter() .map(|c| { i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) }) @@ -702,7 +706,9 @@ fn infer_shape( if !raw.is_empty() { if t.data_type == TensorProto_DataType::Int32 as i32 { return Some( - raw.chunks_exact(4) + raw.as_chunks::<4>() + .0 + .iter() .map(|c| { i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64 }) @@ -710,7 +716,9 @@ fn infer_shape( ); } else { return Some( - raw.chunks_exact(8) + raw.as_chunks::<8>() + .0 + .iter() .map(|c| { i64::from_le_bytes([ c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7], @@ -1850,11 +1858,15 @@ Provide --override-dim {}= or enable --experimental-dynamic-inputs.", let raw = initializer.raw_data.as_slice(); let values = if !raw.is_empty() { if initializer.data_type == TensorProto_DataType::Int32 as i32 { - raw.chunks_exact(4) + raw.as_chunks::<4>() + .0 + .iter() .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64) .collect() } else { - raw.chunks_exact(8) + raw.as_chunks::<8>() + .0 + .iter() .map(|c| { i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) }) @@ -1894,11 +1906,15 @@ Provide --override-dim {}= or enable --experimental-dynamic-inputs.", let raw = tensor.raw_data.as_slice(); let values = if !raw.is_empty() { if tensor.data_type == TensorProto_DataType::Int32 as i32 { - raw.chunks_exact(4) + raw.as_chunks::<4>() + .0 + .iter() .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64) .collect() } else { - raw.chunks_exact(8) + raw.as_chunks::<8>() + .0 + .iter() .map(|c| { i64::from_le_bytes([ c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7], diff --git a/src/onnx/ops/reshape.rs b/src/onnx/ops/reshape.rs index 459a509..717daf8 100644 --- a/src/onnx/ops/reshape.rs +++ b/src/onnx/ops/reshape.rs @@ -110,7 +110,9 @@ impl ReshapeHandler { let raw = t.raw_data.as_slice(); if !raw.is_empty() { let mut axes: Vec = raw - .chunks_exact(8) + .as_chunks::<8>() + .0 + .iter() .map(|c| { i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) }) @@ -179,13 +181,17 @@ impl ReshapeHandler { if !raw_data.is_empty() { match initializer.data_type { x if x == TensorProto_DataType::Int32 as i32 => raw_data - .chunks_exact(4) + .as_chunks::<4>() + .0 + .iter() .map(|chunk| { i32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]) as i64 }) .collect(), _ => raw_data - .chunks_exact(8) + .as_chunks::<8>() + .0 + .iter() .map(|chunk| { i64::from_le_bytes([ chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], @@ -629,13 +635,17 @@ impl ReshapeHandler { if !raw_data.is_empty() { match initializer.data_type { x if x == TensorProto_DataType::Int32 as i32 => raw_data - .chunks_exact(4) + .as_chunks::<4>() + .0 + .iter() .map(|chunk| { i32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]) as i64 }) .collect(), _ => raw_data - .chunks_exact(8) + .as_chunks::<8>() + .0 + .iter() .map(|chunk| { i64::from_le_bytes([ chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], @@ -1237,13 +1247,17 @@ impl ReshapeHandler { if !raw.is_empty() { match tensor.data_type { x if x == TensorProto_DataType::Int64 as i32 => raw - .chunks_exact(8) + .as_chunks::<8>() + .0 + .iter() .map(|c| { i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) }) .collect(), x if x == TensorProto_DataType::Int32 as i32 => raw - .chunks_exact(4) + .as_chunks::<4>() + .0 + .iter() .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64) .collect(), _ => { diff --git a/src/onnx/ops/utility.rs b/src/onnx/ops/utility.rs index cd9d6e1..51fb894 100644 --- a/src/onnx/ops/utility.rs +++ b/src/onnx/ops/utility.rs @@ -680,13 +680,17 @@ impl UtilityHandler { if !raw.is_empty() { if t.data_type == crate::protos::onnx::TensorProto_DataType::Int32 as i32 { return Some( - raw.chunks_exact(4) + raw.as_chunks::<4>() + .0 + .iter() .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64) .collect(), ); } return Some( - raw.chunks_exact(8) + raw.as_chunks::<8>() + .0 + .iter() .map(|c| { i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) }) diff --git a/src/onnx/shape_inference.rs b/src/onnx/shape_inference.rs index 711cdff..1fec275 100644 --- a/src/onnx/shape_inference.rs +++ b/src/onnx/shape_inference.rs @@ -358,11 +358,7 @@ fn infer_node_shape(node: &NodeProto, ctx: &InferenceResult) -> Option> "Concat" => { let mut shapes = Vec::new(); for inp in node.input.as_slice() { - if let Some(s) = ctx.value_shapes.get(inp.as_str()) { - shapes.push(s.clone()); - } else { - return None; - } + shapes.push(ctx.value_shapes.get(inp.as_str())?.clone()); } if shapes.is_empty() { return None; @@ -1379,11 +1375,15 @@ fn read_int_tensor(tensor: &TensorProto) -> Vec { if !raw.is_empty() { match tensor.data_type { x if x == TensorProto_DataType::Int32 as i32 => raw - .chunks_exact(4) + .as_chunks::<4>() + .0 + .iter() .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64) .collect(), _ => raw - .chunks_exact(8) + .as_chunks::<8>() + .0 + .iter() .map(|c| i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]])) .collect(), } diff --git a/tests/format_reference.rs b/tests/format_reference.rs new file mode 100644 index 0000000..8d23c6b --- /dev/null +++ b/tests/format_reference.rs @@ -0,0 +1,34 @@ +use std::fs; + +use webnn_graph::parser::parse_wg_text; +use webnn_graph::serialize::{serialize_graph_to_wg_text, SerializeOptions}; + +#[test] +fn format_reference_round_trips_through_graph_json() { + let path = format!( + "{}/examples/format_reference.webnn", + env!("CARGO_MANIFEST_DIR") + ); + let source = fs::read_to_string(path).expect("read format reference example"); + let parsed = parse_wg_text(&source).expect("parse format reference example"); + + assert_eq!(parsed.version, 2); + assert!(parsed.quantized); + assert!(parsed.inputs["input"].static_shape().is_none()); + assert!(parsed.consts["bias"].shape.is_empty()); + assert!(parsed.consts.contains_key("inline_mask")); + assert!(parsed.consts.contains_key("external_weight")); + assert_eq!( + parsed.nodes[1].outputs.as_deref(), + Some(&["left".into(), "right".into()][..]) + ); + + let serialized = serialize_graph_to_wg_text(&parsed, SerializeOptions::default()) + .expect("serialize format reference example"); + let reparsed = parse_wg_text(&serialized).expect("reparse serialized format reference example"); + + assert_eq!( + serde_json::to_value(parsed).expect("serialize original AST"), + serde_json::to_value(reparsed).expect("serialize round-tripped AST") + ); +}